Skip to content

fix(checkpoint): recover interrupted checkpoint operations - #3519

Open
Jacob953 wants to merge 22 commits into
kagent-dev:mainfrom
Jacob953:fix/checkpoint-operation-recovery
Open

Jacob953 wants to merge 22 commits into
kagent-dev:mainfrom
Jacob953:fix/checkpoint-operation-recovery

Conversation

@Jacob953

@Jacob953 Jacob953 commented Nov 4, 2026 •

Copy link
Copy Markdown

A canceled request or controller restart can leave a CREATING checkpoint blocking its source Session indefinitely. Checkpoint requests now execute a bounded attempt inline, with a small sweep to recover work when the caller never returns. This addresses the remaining lifecycle recovery scope of #2746 after #3199.

Behavior

  • Reservation inserts CREATING and its execution claim atomically. The normal Create path has two writes: reservation and conditional finalization to READY or FAILED, clearing the claim. No database transaction spans a Substrate call.
  • Retries reuse the owner/request identity and acquire expired work. Live claims return Aborted. An executor that loses authority observes the current outcome without repeating backend work or overwriting the replacement executor. Only uncertain acquisition outcomes need claim retirement.
  • A leader-elected sweep resumes expired creation/deletion and collects FAILED rows retaining a tag UID. Collection deletes by UID, then conditionally clears that UID and the live claim. Errors retain the handle for retry; the terminal request result remains stored, so replay never starts another copy.
  • FAILED stays internal: Get/List and forks accept only READY. The public FAILED discovery metadata and associated UI changes have been removed. Existing MCP snapshot-pending/conversation-advanced retry reasons remain intact; private backend and storage causes stay out of public errors.
  • A fresh source that already fails snapshot verification does not start a copy. Recovery still synchronizes with a possible prior copy through Substrate's leases. READY requires completed storage, matching tag/source identity and lineage, and an unchanged source Actor.

The explicit transaction around claim acquisition is retained for cancellation safety. A new PostgreSQL regression delays the real UPDATE until after caller cancellation and retirement, suppressing query cancellation to model lost delivery. With autocommit, the delayed UPDATE leaves a live claim; with the transaction, the canceled COMMIT prevents it. Existing coverage also verifies that retirement waits behind an already acquired row lock.

Upgrade and scope

Ordinary rolling upgrades are supported. Migration 000004 adds nullable claim columns; legacy NULL-claim creation/deletion gets the two-minute UUIDv7 grace. During the mixed-version window, old controllers retain their existing unconditional failure-cleanup race; the new code uses UID-conditional cleanup. Rollback leaves the additive schema in place. The checkpoint-specific drain/rollback tests and upgrade hooks have been removed.

Substrate v1.0.0 serializes create/delete with Actor/Tag leases and uses a unique, never-reused tag-UID destination prefix. Object-store writes can outlive lease loss and leak storage or finish an internal FAILED tag, but cannot make it forkable. Stronger object-store fencing is a Substrate follow-up, not a kagent release blocker.

The direct Substrate/object-store fault suites and their credential-exporting CI script are out of this PR. The regular API E2E remains: lose a successful Create response, restart the controller, replay the same request, and fork through the API. The architecture doc now describes the lifecycle and rolling-upgrade boundary without the implementation history.

Validation

  • Race-enabled checkpoint, Session, PostgreSQL store, migration, gRPC and MCP tests pass; store coverage includes inline SQL preparation and migration Up/Down. Kubernetes-backed checks use envtest 1.37.0.
  • Permanent PostgreSQL regressions cover FAILED collection after source deletion and replacement ownership (already collected, deleted, or still pending). Removing completed-cleanup observation makes both completion regressions fail, confirming that they detect obsolete-release errors masking success.
  • Full public API E2E passes on both gVisor and Cloud Hypervisor MicroVM in independent Linux validation: each leg has 93 passing top-level tests and 425 passing entries including subtests (1290.022s / 1586.059s). Each has the same 7 explicit harness-capability skips: BYO built-in default images; Claude/Codex ADK ask-user-disabled behavior; Codex shared-agent fixture; BYO/Claude/Codex structured output. No checkpoint case is skipped: all 16 checkpoint harness cases pass in each sandbox, including lost successful Create response/controller restart/replay/fork, lifecycle, paused-task rejection and MCP checkpoint/fork. MicroVM runs with a verified KVM_CREATE_VM, not a fallback or skipped leg.
  • Published v1.0.0 upgrade and rollback passes all 13 subtests (510.77s); real rolling upgrade passes (184.41s). The baseline actually runs rather than taking the skip path. Core migration 3→4 is verified with previous-release invocation/restart against the target schema, data retention, clean-install schema equality, current invocation, application rollback and schema rollback. The rolling test observes the target schema while old controller pods are still ready and verifies old-shape reads/writes.
  • These four successful jobs use the production sources of PR head 238ef3d2 unchanged; a validation-only workflow lives on the fork's temporary branch, outside this PR. They use ordinary API E2E and the repository's existing upgrade targets, without exporting controller credentials or object-store keys. Upstream Actions still require maintainer approval; these are independent fork results.
  • The additional local ARM64 check passes all 16 checkpoint harness cases in 46.544s on disposable Kind 1.37.0, real Substrate v1.0.0 and gVisor v7. The Linux E2E legs use Kind 1.37.0, four workers and E2E parallelism 2.
  • Go lint: 0 issues. Previously completed UI validation remains 792 unit tests, typecheck and 6 checkpoint browser scenarios on Chrome with the mock backend; UI lint has 0 errors and 6 existing warnings. There are no net UI changes.
  • Telemetry live-check remains a diagnostic limitation: each Linux E2E leg reports 546 semantic-convention violations, mainly ADK generate-content provider/span-kind conventions, compaction attributes and the session-expiration metric. The corresponding ADK, registry and expiration instrumentation is unchanged by this PR. The API tests pass, but this does not establish clean telemetry conformance.

Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com>
@github-actions github-actions Bot added bug Something isn't working and removed bug Something isn't working labels Nov 4, 2026
Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com>
Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com>
@Jacob953
Jacob953 marked this pull request as ready for review November 4, 2026 16:29
@Jacob953
Jacob953 requested a review from a team as a code owner November 4, 2026 16:29
@github-actions github-actions Bot added bug Something isn't working and removed bug Something isn't working labels Nov 4, 2026
Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com>
@github-actions github-actions Bot added bug Something isn't working and removed bug Something isn't working labels Nov 4, 2026
Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com>
Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com>
@github-actions github-actions Bot added bug Something isn't working and removed bug Something isn't working labels Nov 5, 2026
Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com>
@github-actions github-actions Bot added bug Something isn't working and removed bug Something isn't working labels Nov 6, 2026
Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com>
@github-actions github-actions Bot added bug Something isn't working and removed bug Something isn't working labels Nov 6, 2026
Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com>
Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com>
@github-actions github-actions Bot added bug Something isn't working and removed bug Something isn't working labels Nov 7, 2026

@EItanya EItanya left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 AI-generated review, left by Claude on behalf of @EItanya.

Thanks for digging into this, the stranded CREATING reservation is a real problem worth fixing. Before going line by line, though, I think the overall approach should change. Keeping DB operations simple and transactional matters a lot in this system, and the repo already has a pattern for exactly this kind of problem. This PR builds a separate system next to it.

The core problem

A CREATING row blocks its source Session in about seven places: task admission, quiescence, recovery, deletion and lifecycle operations. If the inline creator dies or its caller cancels, nothing ever finishes that row. Interrupted deletion is secondary: it only keeps references alive until Delete is retried.

The existing pattern

runtime_instance handles abandoned durable operations like this (runtime_instances.go claimRuntimeOperation / ReleaseRuntimeOperation / completeRuntimeOperation, and session/workflow.go):

  • Work runs inline in the request.
  • An execution claim (executor_id, executor_expires_at) is bounded by RuntimeOperationTimeout.
  • The claim is a single conditional UPDATE.
  • Completion clears the claim in the same transaction that records the outcome.
  • On error the claim is released; a caller that fails to claim gets the current outcome back.

The PR adds the same claim columns, but moves all work into a leader-elected worker while API replicas poll for the result. Most of the added complexity follows from that one choice:

  • Durable operation_error column, RecordCheckpointOperationError, and polling waiters: needed only because the caller no longer runs the attempt itself.
  • Work queue, in-memory retry map, wake channel, worker goroutines, a 1 s ListPendingCheckpoints poll, and a new index: the machinery for scheduling and discovering that work.
  • BeginCheckpointCleanup / finishCleanup intermediate phase: a failure recorded while still in CREATING.
  • Activation barrier with a pod annotation, pod/ReplicaSet RBAC and Helm wiring: keeps old inline binaries from racing the new worker.

There's also a concrete bug in the durable-error path, which I reproduced on 6f1dbe8. operation_error is not cleared when a new Delete starts (BeginDeleteSessionCheckpoint) or when a later attempt fails transiently. After the backend is repaired, a retried Delete returns the stale PermissionDenied within about 2 ms. It keeps doing so until the worker's next backoff retry happens to succeed. The design below removes this class of bug entirely.

Suggested design

Make checkpoint creation follow the runtime-operation pattern, and add the smallest possible recovery path:

  1. Reserve takes the claim. ReserveSessionCheckpoint sets executor_id / executor_expires_at in the same INSERT that creates the CREATING row, so the fast path needs no extra query.
  2. Do the work inline, as today: ensureTag under the request deadline.
  3. Finish with one atomic update. A single UPDATE ... WHERE state = 'CREATING' AND executor_id = $me AND executor_expires_at > now() sets READY or FAILED and clears the claim. The fast path stays at two writes: reserve, then finalize.
  4. Drop the cleanup phase. When the copy has finished but fails verification, go to FAILED and keep tag_uid. Delete already removes a tag by its UID, and fork only accepts READY. This removes BeginCheckpointCleanup / finishCleanup and the "failure while still CREATING" state. If holding admission until cleanup is a hard requirement, keep a compensation step, but as one more conditional UPDATE rather than a separate phase.
  5. Leave it pending on uncertain errors. For a transport error or cancellation, leave the row CREATING and release the claim (executor_expires_at = now()), like ReleaseRuntimeOperation. CreateTag's serialization on the Actor and tag already covers a copy that is still finishing.
  6. Retrying takes over. A Create retried with the same request_id claims the existing row and drives it inline, or returns the current outcome if someone else holds the claim.
  7. One small sweep for callers that never come back. Periodically, one query selects CREATING (and optionally DELETING) rows with an expired claim and runs the same function as step 2. Claims already exclude live work, and claim expiry plus the sweep interval act as the retry backoff. No queue, retry map, wake channel or durable error is needed.
  8. No activation barrier. New code never leaves a CREATING row with a NULL claim, so a NULL claim can only come from an old binary. The sweep can skip those until the UUIDv7 id is older than RuntimeOperationTimeout. That bounds the overlap with old binaries without Kubernetes reads, RBAC or annotations. Pre-1.0, documenting it would also be reasonable.

Compared with the current PR, this means:

  • Columns: 2 new (the same claim pair as runtime_instance) instead of 3.
  • Index: none, unless the sweep needs one.
  • Store methods: about 4 new or changed instead of 8.
  • Fast path: reserve and finalize, with no polling.
  • Kubernetes and Helm: no changes.

It also removes the dead unclaimed FinalizeSessionCheckpoint / DeleteSessionCheckpoint paths, since finalization would always require the executor.

One open question for both designs: they rely on CreateTag serializing on the Actor and tag, so that AlreadyExists means any earlier copy has finished. It would be good to confirm that with the Substrate side.

Happy to discuss. If there's a requirement I'm missing that makes the worker model necessary, please point it out.

@github-actions github-actions Bot added bug Something isn't working and removed bug Something isn't working labels Nov 8, 2026
Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com>
@github-actions github-actions Bot added bug Something isn't working and removed bug Something isn't working labels Nov 8, 2026
Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com>
@github-actions github-actions Bot added bug Something isn't working and removed bug Something isn't working labels Nov 9, 2026
@Jacob953

Jacob953 commented Nov 9, 2026

Copy link
Copy Markdown
Author

@EItanya @shashankram — could we settle these two release decisions following the architectural review? Implementation and validation details remain in the PR description.

1. First upgrade from controllers without checkpoint claims

Is a documented, one-time controller API maintenance window acceptable for this transition in 1.0?

My recommendation is to follow the release-wide upgrade policy: quiesce checkpoint mutations, stop every legacy controller, wait for its process to exit, then start claim-aware controllers and resume requests. Legacy writers have no overall operation deadline; UUIDv7 age plus the two-minute grace does not fence a paused writer. Losing readiness or requesting scale-to-zero is also insufficient until the old processes exit.

Before rollback to a claimless binary, pending checkpoint work must finish, FAILED checkpoints retaining tags must be deleted, and current writers must drain. The old deletion path cannot collect partial tags without a completed URI.

The upgrade tests exercise this drain boundary, but the production installer still performs ordinary Helm upgrades. If maintainers accept the maintenance window, the agreed procedure needs public upgrade/rollback instructions and any required installer enforcement before release. If uninterrupted API availability is required, I suggest agreeing on a staged compatibility rollout at the release level before enabling recovery, rather than adding a checkpoint-specific Kubernetes activation protocol. Later claim-aware rollouts can coordinate through database claims, subject to the rest of the release's compatibility requirements.

2. Substrate's serialization guarantee for external writes

Can the Substrate side confirm that, after a serialized CreateTag attempt returns AlreadyExists, FailedPrecondition, or NotFound, earlier attempts for that Actor/tag can no longer write to the copied snapshot? This needs to cover delivered cancellation, lease renewal failure, process failure, and delayed object-store writes. Does UID-conditional DeleteTag provide the corresponding coordination with outstanding copy writes?

The live concurrency regression covers a copy still holding valid leases after cancellation delivery is lost: a competing attempt gets Aborted, then a later retry adopts the completed tag. It does not establish the cancellation/lease-loss guarantee above. Draining kagent controllers does not establish that guarantee either.

My recommendation is to document and fault-test this contract in Substrate. If it does not hold, the owning backend needs to supply a definitive safe outcome; until then, kagent must retain CREATING admission for an uncertain copy rather than infer safety from timeout or tag-name existence. The guarantee is needed before marking a retained partial copy FAILED and permitting cleanup.

Could we confirm the accepted upgrade policy and who can validate the Substrate contract, so the remaining release work has a clear owner?

Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com>
@github-actions github-actions Bot added bug Something isn't working and removed bug Something isn't working labels Nov 9, 2026
Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com>
@github-actions github-actions Bot added bug Something isn't working and removed bug Something isn't working labels Nov 9, 2026
Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com>
@github-actions github-actions Bot added bug Something isn't working and removed bug Something isn't working labels Nov 10, 2026
Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com>
@github-actions github-actions Bot added bug Something isn't working and removed bug Something isn't working labels Nov 10, 2026
@EItanya

EItanya commented Nov 10, 2026

Copy link
Copy Markdown
Contributor

🤖 Comment left by Claude on behalf of @EItanya.

Thanks for reworking this so quickly, @Jacob953. The inline claim design looks much better: reserve takes the claim, finalize is a single conditional update, and the worker, the durable error and the activation barrier are all gone. Answers to your two questions are below, followed by a few remaining concerns about the code.

1. First upgrade

I'd rather not have a maintenance window, and from what I can tell we don't need one. I went through each case where an old controller and a new one run at the same time:

  • Schema. 000004 only adds nullable columns. The migration runner already tolerates a database that is newer than the binary (it only rejects unknown versions <= latest), so old pods keep running and restarting. Rollback doesn't need a down migration.
  • A row an old controller is still creating. The sweep leaves NULL-claim rows alone for 2 minutes. The old writer's request can't outlive its pod, and a client cancellation cancels Substrate's workflow context. If the two do overlap, both go through CreateTag's leases, and each finalize only applies to a CREATING row, so whichever finishes second just reads the outcome.
  • Deletes. Deleting by UID is idempotent. An old controller can't delete a FAILED checkpoint whose tag is partial, so that row waits in DELETING until a new controller picks it up.
  • The old failure path's DeleteTag(name, ""). On failure, an old controller can delete a tag that a new one has already finalized. But two old replicas can already do this to each other on main today (that's the singleflight TODO), and this PR fixes it once the rollout finishes. The upgrade doesn't make it any worse.

So I think an ordinary rolling upgrade is fine. Can we replace the quiesce/drain procedure in the doc with a short note about the mixed-version window, and drop checkpoint_drain_test.go and checkpoint_rollback_test.go?

2. Substrate's guarantee

I read through the v1.0.0 server code (controlapi/workflow_tag.go, workflow.go, store/atepg/lease.go):

  • TagActorSnapshot holds the Actor and Tag leases for the whole workflow, under a context that combines the request, a 5-minute workflow deadline and lease loss.
  • AlreadyExists is only returned after both leases have been acquired. A concurrent live attempt gets Aborted instead.
  • A pending tag is never resumed. The doc comment says a create that dies leaves the name taken, and to retry you delete and create again. That matches how this PR treats AlreadyExists plus PENDING.
  • Substrate's own DeleteTag uses the same tag lease to keep out a live copy.

You're right that this isn't airtight. The copy runs in the object-store plugin, and no fencing token reaches the object store, so an in-flight write can land after the lease is gone. However, I don't think kagent depends on that guarantee. The destination prefix is keyed by the tag UID and never reused. A late write can only leak objects after cleanup (a Substrate storage concern that affects every client) or finish a tag we've already marked FAILED. Fork only accepts READY, and Delete still works by UID. We also never publish READY without verifying that the tag is COMPLETED with the right identity.

So I'd trim "Required backend contract: pending Substrate confirmation" down to a short note rather than block on it. A Substrate issue about fencing object-store writes would still be nice to have, but it isn't a release blocker for us.

Remaining concerns

FAILED checkpoints are now public. Because a FAILED row keeps its tag, Get and List now return FAILED rows, and the PR adds UI, MCP and Playwright changes so users can find and delete them. That's a public API change, and it's much harder to take back than to add later. Can we keep FAILED internal and have the system collect these tags itself? For example, the sweep could also select state = 'FAILED' AND tag_uid <> '', delete the tag by UID, and then run UPDATE ... SET tag_uid = '' WHERE id = $1 AND state = 'FAILED' AND tag_uid = $2. That's one more conditional update with no new column, and the Get/List/UI/MCP changes could all be reverted.

Size. The PR is currently +4.2k lines across 58 files. If FAILED stays internal and the upgrade tests go, a lot of that drops out. I'd also prefer not to wire checkpoint-recovery.sh into ci.yaml. Copying the controller's credential bundle and the RustFS keys out of the cluster to run a go test from CI isn't something we want to normalize. Could the real-Substrate scenarios run as regular E2E tests against the API instead, or move to a follow-up?

Architecture doc. The new section reads more like an implementation log than a design doc. For example, "Create observers skip obsolete claim releases…" is better as a code comment. I think a few paragraphs are enough: the claim lifecycle, the sweep, and how FAILED is collected.

Smaller things:

  • ensureTag still calls CreateTag after verifySnapshot has already reported that the Actor's snapshot changed. If no earlier attempt exists, that creates a new copy of the wrong snapshot only to mark it FAILED. Any particular reason to use CreateTag as the barrier there instead of failing directly?
  • reserveCreate releases the claim after every reserve error, including definite ones such as ErrCheckpointAdvanced. Can we only release when the outcome is uncertain?
  • Do we need the transaction around the single UPDATE in ClaimCheckpointOperation? ReleaseCheckpointOperation already locks the row first, so I'm not sure the extra round trips buy anything.
  • review_test.go should be named after the behavior it tests rather than the review round.

…cope

Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com>
@github-actions github-actions Bot added bug Something isn't working and removed bug Something isn't working labels Nov 10, 2026
…over

Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com>
@github-actions github-actions Bot added bug Something isn't working and removed bug Something isn't working labels Nov 10, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants