Cancelling an AI task and undoing what it already saved

01 Background

An AI task runs across several stores and cannot be wrapped in one long transaction, so I had to decide what a cancellation should undo and how far it should reach.

02 Decisions and implementation

1. Deciding that stopping execution was not enough

   Problem An AI task saves documents, folders, wiki pages and skills several times between model calls, so stopping execution left part of a cancelled task in place. Yet execution spanning the core DB, the AI DB and object storage could not be wrapped in one long transaction either.

   Decision I chose to undo already-committed changes too. Asking the model to perform the inverse cannot guarantee original values, approval scope or idempotency, and locking the whole service in one transaction cannot absorb long-running model, HTTP and object-storage work or its failures, so both were rejected.

   Implementation The AI side derives reverse actions from the run record while the backend performs atomic recovery of business data. Agent tools store a pre-change snapshot with reverse-action arguments and an idempotency key tied to the original approval, and ordinary tasks record before and after values through each database’s triggers inside the same transaction. Object storage persists the original and the intended result before writing.

   Verification A test that runs the real worker makes model creation fail during recovery and confirms the cancellation finishes without any model or gateway call. Cancelling a two-folder, six-document organisation task at each of 9 points returned the workspace to its original state, and after a lost response a new worker resumed recovery with the same idempotency key.

2. Blocking later writes and undoing children first

   Problem A handler still running after a cancellation would overwrite recovered values, and undoing a parent first left the changes its children had made behind. Recovery could fail, leaving users unable to return to the version they wanted.

   Decision Persist the cancellation first to block later writes, and once the running handler finishes, undo children first in reverse order.

   Implementation The agent turn and recovery share one PostgreSQL execution lock. The lock releases when a worker connection ends, so a leftover in-progress record cannot block recovery. Body recovery is saved as a new revision, so a late event cannot overwrite the source again.

   Verification Tests confirm that writes after cancellation are refused, recovery waits for the running handler, and children are restored before their parent. Even after the running process was force-killed, recovery waited on the same lock and finished, and a late event from an earlier revision could not overwrite the restored body. If another user later changed the data, recovery did not overwrite it and stayed retryable; cancelling again after the conflict cleared completed normally.

3. Recovering without erasing other tasks’ contributions

   Problem A wiki concept page merges contributions from several documents, so overwriting it with an older snapshot also removed valid contributions added later by other documents.

   Decision Removing just the relevant part from the finished result was hard because which sentence came from which document cannot be told apart reliably. Restoring an entire older snapshot was rejected too, because information is lost if a new document contributed in between. Instead the history is append-only and recovery excludes only the target contribution before reassembling.

   Implementation Operations, page revisions and per-document contributions are recorded append-only, and recovery restores or reassembles the page with the target contribution excluded. Pages and history are soft-deleted rather than removed, which preserves the audit trail.

   Verification Right before recovery, the contribution list (operation, revision, active flag) is re-read under a lock, and recovery is refused if it differs from the preview. Tests confirm that pages are rebuilt from the selected contributions only and that if any page fails, no current state changes. Recovery itself is recorded as a new operation and revision.

03 Improved processing flow

flowchart TD
 A[User cancels] --> B[Persist cancellation first, block later writes]
 B --> C[Wait for the running handler to finish]
 C --> D[Derive reverse actions from the run record, no model call]
 D --> E[Undo children first, in reverse]
 E --> F{Current values and references unchanged}
 F -->|Yes| G[Atomic recovery, body saved as a new revision]
 F -->|No| H[Leave as rollback_failed for retry]
 G --> I{Full recovery confirmed}
 I -->|Yes| J[cancelled]
 I -->|No| H

flowchart TD
 A[User cancels] --> B[Persist cancellation first, block later writes]
 B --> C[Wait for the running handler to finish]
 C --> D[Derive reverse actions from the run record, no model call]
 D --> E[Undo children first, in reverse]
 E --> F{Current values and references unchanged}
 F -->|Yes| G[Atomic recovery, body saved as a new revision]
 F -->|No| H[Leave as rollback_failed for retry]
 G --> I{Full recovery confirmed}
 I -->|Yes| J[cancelled]
 I -->|No| H

04 Results and lessons

A cancelled task’s changes are undone while other tasks’ contributions survive

Content, names, parents and ordering that users see are restored, while version, audit and idempotency records are kept. Recovery runs from the run record alone, with no model settings or calls, and it checks current values and downstream references so it never overwrites another task’s work. Cancellation is final only once full recovery is confirmed; conflicts and unclear outcomes stay retryable. For a feature where AI edits a user’s documents, I settled what could be undone and what had to be recorded before building it. In a manual run on real Gemini, Kafka, PostgreSQL and Redis, nothing changed before approval, and cancelling during the last document move restored 2 folders and 6 documents with 8 inverse operations.