fix(checkpoint): recover interrupted checkpoint operations

docs/architecture/persistence-checkpoints-and-forks.md

@@ -59,10 +59,39 @@ Input-required and auth-required tasks are durably suspended with VOLUMES snapsh
 but are not checkpointable or forkable; callers must resolve the interaction first.
 Creating a checkpoint does not suspend the Actor again:
 
-1. Reserve the checkpoint in PostgreSQL.
+1. Reserve the checkpoint and a bounded execution claim in PostgreSQL.
 2. Verify that the suspended Actor still holds the external snapshot URI and fidelity recorded on the boundary.
 3. Create a Substrate `Tag`, which copies the Actor's current snapshot into independent storage, and verify the source did not change during the copy.
-4. Atomically persist the Tag UID and copied snapshot URI and mark the checkpoint ready.
+4. Conditionally persist the Tag UID and copied snapshot URI and mark the checkpoint ready, only while this executor's claim is live.
+
+Create and Delete execute inline; they do not wait for a background worker.
+The reservation owns the first attempt's claim. Retries acquire an expired claim
+and read the current phase before making backend calls. Claims last at most two
+minutes, and finalization clears the claim atomically. Uncertain outcomes retain
+durable work for retry. A leader-elected sweep resumes expired creation/deletion
+and collects failed tags, waiting 30 seconds after each completed sweep.
+
+Only READY checkpoints are public and forkable. Failed verification records an
+internal FAILED result and releases source admission. The sweep deletes any
+retained tag by UID, then conditionally clears that UID and its execution claim.
+Cleanup errors preserve the handle for retry. The FAILED row retains request
+identity, history and revision references, so replay returns the same terminal
+failure without copying again. A fresh invalid source fails before starting a
+copy; recovery still synchronizes with a possible earlier copy through Substrate.
+
+Ordinary rolling upgrades are supported. The claim migration adds nullable columns;
+legacy NULL-claim creation/deletion waits two minutes by UUIDv7 age before takeover.
+Overlapping tag attempts serialize through Substrate's Actor and Tag leases.
+During the mixed-version window, old controllers still have their existing
+unconditional failure-cleanup race; UID-conditional cleanup applies once the
+rollout completes. Rollback leaves the additive schema in place.
+
+Substrate v1.0.0 acquires both leases before returning AlreadyExists, and deletion
+uses the Tag lease. Object-store writes can outlive lease loss, but each copied
+snapshot uses a tag-UID-specific destination prefix that is never reused. Late
+writes can leak storage or finish a FAILED tag; they cannot make it forkable.
+READY still requires COMPLETED storage and matching source/tag identity. Stronger
+object-store fencing belongs in Substrate and does not block this recovery path.
 
 Kagent requests synchronous Substrate suspension before recording a checkpoint.
 If an earlier attempt already suspended the Actor but left a pending durable
@@ -76,8 +105,8 @@ and requires a COMPLETED durable copy. The Tag's template UID must
 match that lineage, and the Actor UID must remain unchanged across the copy.
 
 The `CREATING` reservation blocks task writes and idle or explicit lifecycle changes until the
-copy completes or cleanup finishes. A lost response can reuse a completed Tag;
-an incomplete copy is deleted before the failed reservation is released. A
+copy is verified or marked FAILED. A lost response can reuse a completed Tag;
+an incomplete copy becomes internal FAILED work for background collection. A
 failed database finalization leaves the reservation and completed Tag for retry.
 Idle lifecycle work records snapshots without retaining every turn: only explicit
 checkpoints survive subsequent suspends or source deletion.
@@ -93,13 +122,10 @@ deleted while its context and checkpoint remain.
 Deletion first hides the checkpoint, then deletes its snapshot tag, then removes
 the row. A checkpoint referenced by a fork cannot be deleted. Substrate deletes the Tag's copied snapshot with the Tag.
 
-Failed checkpoints can also be deleted by their owner. Creation marks a checkpoint
-failed only after snapshot-tag cleanup succeeds or confirms that no tag exists.
-Deletion verifies that absence before removing the row; it never deletes an
-unexpected tag without a retained UID. Backend errors leave the checkpoint in
-`DELETING` for a later delete request to retry. Its history and runtime revision
-references remain retained until deletion completes. Checkpoints still in
-`CREATING` cannot be deleted through this path.
+Deletion never removes an unexpected tag without its retained UID. Backend errors
+leave the checkpoint in `DELETING` for a request or the sweep to retry. History and
+runtime revision references remain retained until deletion completes. Checkpoints
+still in `CREATING` cannot be deleted through this path.
 
 Checkpoint deletion passes the retained tag UID as a Substrate delete precondition,
 so a tag replaced after lookup cannot have its snapshot deleted. A rejected delete

go/core/internal/database/active_task_test.go

@@ -24,7 +24,7 @@ func TestCheckpointCreationBlocksSessionTaskWrites(t *testing.T) {
 	task.Status.State = a2a.TaskStateCompleted
 	require.NoError(t, saveRuntimeTask(t, client, session.GetId(), task, task,
 		&SessionTaskSnapshot{Atespace: "team-a", URI: "snapshot", Fidelity: "SNAPSHOT_FIDELITY_VOLUMES"}))
-	checkpoint, _, err := client.ReserveSessionCheckpoint(ctx,
+	checkpoint, _, err := reserveCheckpoint(ctx, client,
 		&apiv1.Checkpoint{Id: uuid.NewString(), SessionId: session.GetId(), HeadTaskId: string(task.ID)}, "alice", uuid.NewString())
 	require.NoError(t, err)
 
@@ -33,13 +33,13 @@ func TestCheckpointCreationBlocksSessionTaskWrites(t *testing.T) {
 	_, err = client.CreateRuntimeTask(ctx, session.GetId(), taskMutationHash("next-request"), next, "")
 	require.ErrorIs(t, err, ErrConflict)
 	require.ErrorIs(t, saveRuntimeTask(t, client, session.GetId(), task, task, nil), ErrFailedPrecondition)
-	_, _, err = client.BeginDeleteSessionCheckpoint(ctx, checkpoint.GetId(), "alice")
+	_, _, err = beginDeleteCheckpoint(ctx, client, checkpoint.GetId(), "alice")
 	require.ErrorIs(t, err, ErrNotFound)
 
-	_, err = client.FinalizeSessionCheckpoint(ctx, checkpoint.GetId(), "", "", "snapshot failed")
+	_, err = finalizeCheckpoint(ctx, client, checkpoint.GetId(), "", "", "snapshot failed")
 	require.NoError(t, err)
-	_, err = client.FinalizeSessionCheckpoint(ctx, checkpoint.GetId(), "tag", "retained", "")
-	require.ErrorIs(t, err, ErrNotFound)
+	_, err = finalizeCheckpoint(ctx, client, checkpoint.GetId(), "tag", "retained", "")
+	require.ErrorIs(t, err, ErrConflict)
 	version, err := client.CreateRuntimeTask(ctx, session.GetId(), taskMutationHash("next-request"), next, "")
 	require.NoError(t, err)
 	require.Positive(t, version)

go/core/internal/database/checkpoint_helpers_test.go

@@ -0,0 +1,64 @@
+package database
+
+import (
+	"context"
+	"errors"
+
+	"github.com/google/uuid"
+
+	apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1"
+)
+
+// Fixtures admit a bounded attempt and abandon it explicitly. Production callers
+// retain this claim while performing the external copy.
+func reserveCheckpoint(ctx context.Context, client *Client, checkpoint *apiv1.Checkpoint, owner, request string) (*apiv1.Checkpoint, *SessionTaskSnapshot, error) {
+	ctx, cancel := context.WithTimeout(ctx, RuntimeOperationTimeout)
+	defer cancel()
+	executor := uuid.New()
+	result, snapshot, err := client.ReserveSessionCheckpoint(ctx, checkpoint, owner, request, executor)
+	err = errors.Join(err, client.ReleaseCheckpointOperation(context.WithoutCancel(ctx), checkpoint.Id, executor))
+	return result, snapshot, err
+}
+
+func finalizeCheckpoint(ctx context.Context, client *Client, id, uid, uri, failure string) (*apiv1.Checkpoint, error) {
+	ctx, cancel := context.WithTimeout(ctx, RuntimeOperationTimeout)
+	defer cancel()
+	op, err := client.GetCheckpointOperation(ctx, id, "alice")
+	if err != nil {
+		return nil, err
+	}
+	executor := uuid.New()
+	op, err = client.ClaimCheckpointOperation(ctx, id, op.Creator, executor)
+	if err != nil {
+		return nil, err
+	}
+	return client.FinalizeClaimedSessionCheckpoint(ctx, op.Checkpoint, executor, uid, uri, failure)
+}
+
+func deleteCheckpoint(ctx context.Context, client *Client, id, owner string) error {
+	ctx, cancel := context.WithTimeout(ctx, RuntimeOperationTimeout)
+	defer cancel()
+	executor := uuid.New()
+	_, err := client.ClaimCheckpointOperation(ctx, id, owner, executor)
+	if err != nil {
+		return err
+	}
+	return client.DeleteClaimedSessionCheckpoint(ctx, id, owner, executor)
+}
+
+func beginDeleteCheckpoint(ctx context.Context, client *Client, id, owner string) (*SessionTaskSnapshot, string, error) {
+	ctx, cancel := context.WithTimeout(ctx, RuntimeOperationTimeout)
+	defer cancel()
+	executor := uuid.New()
+	op, err := client.BeginDeleteSessionCheckpoint(ctx, id, owner, executor)
+	if errors.Is(err, ErrConflict) {
+		op, err = client.GetCheckpointOperation(ctx, id, owner)
+	}
+	if err != nil {
+		return nil, "", err
+	}
+	if err = client.ReleaseCheckpointOperation(context.WithoutCancel(ctx), id, executor); err != nil {
+		return nil, "", err
+	}
+	return op.Snapshot, op.TagUID, nil
+}

go/core/internal/database/checkpoint_operations.go

@@ -0,0 +1,196 @@
+package database
+
+import (
+	"context"
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/jackc/pgx/v5"
+
+	apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1"
+)
+
+// CheckpointOperation contains the durable inputs for creation or deletion.
+type CheckpointOperation struct {
+	Checkpoint *apiv1.Checkpoint
+	Creator    string
+	Snapshot   *SessionTaskSnapshot
+	TagUID     string
+}
+
+// ClaimCheckpointOperation grants one bounded attempt exclusive execution of an
+// owner's pending work, including FAILED tag cleanup, and returns its current
+// phase. The caller must establish its attempt deadline before claiming and
+// use that context for every backend call. Missing, completed, or actively claimed
+// work returns ErrConflict. No database lock survives this call.
+func (c *Client) ClaimCheckpointOperation(ctx context.Context, id, owner string, executor uuid.UUID) (*CheckpointOperation, error) {
+	deadline, ok := ctx.Deadline()
+	if !ok || executor == uuid.Nil || time.Until(deadline) > RuntimeOperationTimeout {
+		return nil, ErrFailedPrecondition
+	}
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+	legacyEligible := false
+	if identifier, parseErr := uuid.Parse(id); parseErr == nil && identifier.Version() == 7 {
+		seconds, nanos := identifier.Time().UnixTime()
+		legacyEligible = time.Since(time.Unix(seconds, nanos)) > RuntimeOperationTimeout
+	}
+	var operation *CheckpointOperation
+	// Do not let a canceled autocommit UPDATE acquire a claim after its caller
+	// has already tried to release it. A lost COMMIT reply still holds the row
+	// lock until the transaction ends, so release serializes behind that outcome.
+	err := c.withTx(ctx, func(tx pgx.Tx) error {
+		row, err := queryOne(ctx, tx, `
+		UPDATE session_checkpoint SET executor_id = $3,
+		    executor_expires_at = clock_timestamp() + $4 * interval '1 millisecond'
+		WHERE id = $1 AND user_id = $2
+		    AND (state IN ('CREATING', 'DELETING') OR (state = 'FAILED' AND tag_uid <> ''))
+		    AND ((executor_id IS NULL AND (state = 'FAILED' OR $5::boolean)) OR executor_expires_at <= clock_timestamp())
+		RETURNING id, source_session_id, user_id, request_id, head_task_id, history_sequence,
+		    snapshot_atespace, snapshot_uri, snapshot_fidelity, tag_uid, state, data,
+		    source_history_id, prepared_revision
+		`, pgx.RowToStructByName[sessionCheckpointRow], id, owner, executor, time.Until(deadline).Milliseconds()+1, legacyEligible)
+		if errors.Is(err, pgx.ErrNoRows) {
+			return ErrConflict
+		}
+		if err != nil {
+			return err
+		}
+		operation, err = checkpointOperation(row)
+		return err
+	})
+	if err != nil {
+		return nil, fmt.Errorf("claim checkpoint operation: %w", err)
+	}
+	return operation, nil
+}
+
+// ReleaseCheckpointOperation releases only this executor's attempt. It does not
+// abandon durable work, and a delayed release cannot revoke a replacement claim.
+// Missing rows and other executors are successful no-ops.
+func (c *Client) ReleaseCheckpointOperation(ctx context.Context, id string, executor uuid.UUID) error {
+	return c.withTx(ctx, func(tx pgx.Tx) error {
+		return releaseCheckpointOperation(ctx, tx, id, executor)
+	})
+}
+
+// ReleaseCheckpointReservation retires an uncertain initial INSERT's executor.
+// Reservation locks its source Session before inserting; wait for that transaction
+// before looking for a row that may not yet be visible. Missing Sessions and
+// checkpoints are harmless; another executor is never released. The lock order
+// remains Session then checkpoint, and no lock survives this transaction.
+func (c *Client) ReleaseCheckpointReservation(ctx context.Context, sessionID, id string, executor uuid.UUID) error {
+	return c.withTx(ctx, func(tx pgx.Tx) error {
+		if _, err := lockSession(ctx, tx, sessionID); err != nil && !errors.Is(err, pgx.ErrNoRows) {
+			return err
+		}
+		return releaseCheckpointOperation(ctx, tx, id, executor)
+	})
+}
+
+// releaseCheckpointOperation waits behind an existing row's acquisition before
+// testing its executor. Initial INSERTs also need their source Session lock.
+// Missing rows and other executors are successful no-ops.
+func releaseCheckpointOperation(ctx context.Context, tx pgx.Tx, id string, executor uuid.UUID) error {
+	if _, err := lockCheckpoint(ctx, tx, id, true, "", nil); err != nil {
+		if errors.Is(err, pgx.ErrNoRows) {
+			return nil
+		}
+		return err
+	}
+	_, err := tx.Exec(ctx, `
+		UPDATE session_checkpoint SET executor_expires_at = clock_timestamp()
+		WHERE id = $1 AND executor_id = $2
+	`, id, executor)
+	return err
+}
+
+// GetCheckpointOperation returns an owner's checkpoint in any lifecycle phase.
+// Missing and other-owner records return ErrNotFound. It does not authorize
+// callers; services compose permissions before using this internal lifecycle read.
+func (c *Client) GetCheckpointOperation(ctx context.Context, id, owner string) (*CheckpointOperation, error) {
+	row, err := queryOne(ctx, c.db, `
+        SELECT id, source_session_id, user_id, request_id, head_task_id, history_sequence,
+            snapshot_atespace, snapshot_uri, snapshot_fidelity, tag_uid, state, data,
+            source_history_id, prepared_revision
+        FROM session_checkpoint WHERE id = $1 AND user_id = $2
+    `, pgx.RowToStructByName[sessionCheckpointRow], id, owner)
+	if err != nil {
+		return nil, fmt.Errorf("get checkpoint operation: %w", notFoundOr(err))
+	}
+	return checkpointOperation(row)
+}
+
+// ListPendingCheckpoints returns a bounded ID-ordered page of durable creation,
+// deletion, and FAILED tag cleanup across owners. Only the checkpoint worker uses
+// this unscoped query; age never establishes whether a tag copy is still live.
+func (c *Client) ListPendingCheckpoints(ctx context.Context, afterID string, limit int) ([]*CheckpointOperation, error) {
+	var after *string
+	if afterID != "" {
+		after = &afterID
+	}
+	rows, err := queryMany(ctx, c.db, `
+		SELECT id, source_session_id, user_id, request_id, head_task_id, history_sequence,
+		    snapshot_atespace, snapshot_uri, snapshot_fidelity, tag_uid, state, data,
+		    source_history_id, prepared_revision
+		FROM session_checkpoint
+		WHERE (state IN ('CREATING', 'DELETING') OR (state = 'FAILED' AND tag_uid <> ''))
+          AND (executor_expires_at <= statement_timestamp() OR (executor_id IS NULL AND (state = 'FAILED' OR id < $3::uuid)))
+		  AND ($1::uuid IS NULL OR id > $1::uuid)
+		ORDER BY id LIMIT $2
+	`, pgx.RowToStructByName[sessionCheckpointRow], after, limit, legacyCheckpointCutoff(time.Now()))
+	if err != nil {
+		return nil, fmt.Errorf("list pending checkpoints: %w", err)
+	}
+	result := make([]*CheckpointOperation, 0, len(rows))
+	for _, row := range rows {
+		op, err := checkpointOperation(row)
+		if err != nil {
+			return nil, err
+		}
+		result = append(result, op)
+	}
+	return result, nil
+}
+
+// ClearFailedCheckpointTag records successful UID-conditional cleanup without
+// discarding the owner's terminal request result. Only the current executor can
+// clear the retained UID and claim atomically; stale completions return ErrConflict.
+func (c *Client) ClearFailedCheckpointTag(ctx context.Context, id, owner, tagUID string, executor uuid.UUID) error {
+	if tagUID == "" || executor == uuid.Nil {
+		return ErrFailedPrecondition
+	}
+	tag, err := c.db.Exec(ctx, `
+		UPDATE session_checkpoint SET tag_uid = '', executor_id = NULL, executor_expires_at = NULL
+		WHERE id = $1 AND user_id = $2 AND state = 'FAILED' AND tag_uid = $3
+		    AND executor_id = $4 AND executor_expires_at > clock_timestamp()
+	`, id, owner, tagUID, executor)
+	if err != nil {
+		return fmt.Errorf("clear failed checkpoint tag: %w", err)
+	}
+	if tag.RowsAffected() != 1 {
+		return ErrConflict
+	}
+	return nil
+}
+
+func checkpointOperation(row sessionCheckpointRow) (*CheckpointOperation, error) {
+	checkpoint, err := toSessionCheckpoint(row)
+	if err != nil {
+		return nil, err
+	}
+	return &CheckpointOperation{Checkpoint: checkpoint, Creator: row.UserID,
+		Snapshot: checkpointSnapshot(row), TagUID: row.TagUID}, nil
+}
+
+// Legacy rows have no claim. The UUID grace delays freshly reserved work during
+// rolling upgrades; overlapping tag attempts still serialize through Substrate.
+func legacyCheckpointCutoff(now time.Time) uuid.UUID {
+	var id uuid.UUID
+	binary.BigEndian.PutUint64(id[:8], uint64(now.Add(-RuntimeOperationTimeout).UnixMilli())<<16)
+	return id
+}