diff --git a/docs/architecture/persistence-checkpoints-and-forks.md b/docs/architecture/persistence-checkpoints-and-forks.md index 8d4d16c0a7..5954c30104 100644 --- a/docs/architecture/persistence-checkpoints-and-forks.md +++ b/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 diff --git a/go/core/internal/database/active_task_test.go b/go/core/internal/database/active_task_test.go index fe7b29480d..be191eecd4 100644 --- a/go/core/internal/database/active_task_test.go +++ b/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) diff --git a/go/core/internal/database/checkpoint_helpers_test.go b/go/core/internal/database/checkpoint_helpers_test.go new file mode 100644 index 0000000000..590260d837 --- /dev/null +++ b/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 +} diff --git a/go/core/internal/database/checkpoint_operations.go b/go/core/internal/database/checkpoint_operations.go new file mode 100644 index 0000000000..b6ae54de99 --- /dev/null +++ b/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 +} diff --git a/go/core/internal/database/checkpoint_operations_test.go b/go/core/internal/database/checkpoint_operations_test.go new file mode 100644 index 0000000000..6186ea6971 --- /dev/null +++ b/go/core/internal/database/checkpoint_operations_test.go @@ -0,0 +1,314 @@ +package database + +import ( + "context" + "encoding/binary" + "errors" + "strings" + "sync" + "testing" + "time" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/stretchr/testify/require" + "google.golang.org/protobuf/proto" + + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" +) + +type checkpointCommitGate struct { + entered chan struct{} + resume chan struct{} + once sync.Once +} + +func (g *checkpointCommitGate) TraceQueryStart(ctx context.Context, _ *pgx.Conn, data pgx.TraceQueryStartData) context.Context { + if strings.EqualFold(data.SQL, "commit") { + g.once.Do(func() { + close(g.entered) + select { + case <-g.resume: + case <-ctx.Done(): + } + }) + } + return ctx +} + +func (*checkpointCommitGate) TraceQueryEnd(context.Context, *pgx.Conn, pgx.TraceQueryEndData) {} + +func TestCheckpointReleaseWaitsForClaimCommit(t *testing.T) { + pool := setupTestDB(t) + checkpoint, _ := reserveRecoveryCheckpoint(t, NewClient(pool, "public")) + gate := &checkpointCommitGate{entered: make(chan struct{}), resume: make(chan struct{})} + config := pool.Config().Copy() + config.ConnConfig.Tracer = gate + claimedPool, err := Connect(t.Context(), &PostgresConfig{Pool: config, Schema: "public"}) + require.NoError(t, err) + t.Cleanup(claimedPool.Close) + var unblock sync.Once + resume := func() { unblock.Do(func() { close(gate.resume) }) } + t.Cleanup(resume) + client := NewClient(claimedPool, "public") + ctx, cancel := context.WithTimeout(t.Context(), RuntimeOperationTimeout) + defer cancel() + executor := uuid.New() + claimDone := make(chan error, 1) + go func() { + _, err := client.ClaimCheckpointOperation(ctx, checkpoint.Id, "alice", executor) + claimDone <- err + }() + select { + case <-gate.entered: + case <-ctx.Done(): + t.Fatal(ctx.Err()) + } + releaseDone := make(chan error, 1) + go func() { releaseDone <- client.ReleaseCheckpointOperation(ctx, checkpoint.Id, executor) }() + select { + case err := <-releaseDone: + t.Fatalf("release passed an uncommitted claim instead of waiting: %v", err) + case <-time.After(100 * time.Millisecond): + } + resume() + require.NoError(t, <-claimDone) + require.NoError(t, <-releaseDone) + _, err = client.ClaimCheckpointOperation(ctx, checkpoint.Id, "alice", uuid.New()) + require.NoError(t, err, "release must retire the claim after its delayed commit") +} + +func TestCanceledCheckpointClaimCannotCommitAfterRelease(t *testing.T) { + pool := setupTestDB(t) + client := NewClient(pool, "public") + checkpoint, _ := reserveRecoveryCheckpoint(t, client) + gate := &checkpointUpdateGate{entered: make(chan struct{}), resume: make(chan struct{})} + config := pool.Config().Copy() + config.ConnConfig.Tracer = gate + latePool, err := Connect(t.Context(), &PostgresConfig{Pool: config, Schema: "public"}) + require.NoError(t, err) + t.Cleanup(latePool.Close) + var once sync.Once + resume := func() { once.Do(func() { close(gate.resume) }) } + t.Cleanup(resume) + ctx, cancel := context.WithTimeout(t.Context(), RuntimeOperationTimeout) + defer cancel() + executor := uuid.New() + done := make(chan error, 1) + go func() { + _, err := NewClient(latePool, "public").ClaimCheckpointOperation(ctx, checkpoint.Id, "alice", executor) + done <- err + }() + select { + case <-gate.entered: + case <-ctx.Done(): + t.Fatal(ctx.Err()) + } + cancel() + // No UPDATE has locked the row yet: retirement can complete first. + require.NoError(t, client.ReleaseCheckpointOperation(t.Context(), checkpoint.Id, executor)) + resume() + claimErr := <-done + replacement, stop := context.WithTimeout(t.Context(), RuntimeOperationTimeout) + defer stop() + _, err = client.ClaimCheckpointOperation(replacement, checkpoint.Id, "alice", uuid.New()) + require.NoError(t, err, "a retired caller must not leave a late claim behind") + require.ErrorIs(t, claimErr, context.Canceled) +} + +// Deliver the UPDATE after cancellation and retirement, modeling lost query +// cancellation. PostgreSQL executes the real write; COMMIT keeps the caller's +// canceled context. No persisted lifecycle state or timestamp is fabricated. +type checkpointUpdateGate struct { + entered, resume chan struct{} + once sync.Once +} + +func (g *checkpointUpdateGate) TraceQueryStart(ctx context.Context, _ *pgx.Conn, data pgx.TraceQueryStartData) context.Context { + if strings.Contains(data.SQL, "UPDATE session_checkpoint SET executor_id") { + g.once.Do(func() { + close(g.entered) + <-g.resume + }) + return context.WithoutCancel(ctx) + } + return ctx +} + +func (*checkpointUpdateGate) TraceQueryEnd(context.Context, *pgx.Conn, pgx.TraceQueryEndData) {} + +func TestFailedCheckpointCollectionPreservesRequestAndFencesCleanup(t *testing.T) { + client := NewClient(setupTestDB(t), "public") + checkpoint, _ := reserveRecoveryCheckpoint(t, client) + ctx, cancel := context.WithTimeout(t.Context(), RuntimeOperationTimeout) + defer cancel() + creator := uuid.New() + _, err := client.ClaimCheckpointOperation(ctx, checkpoint.Id, "alice", creator) + require.NoError(t, err) + failed, err := client.FinalizeClaimedSessionCheckpoint(ctx, checkpoint, creator, "partial-tag", "", "copy failed") + require.NoError(t, err) + page, err := client.ListPendingCheckpoints(ctx, "", 100) + require.NoError(t, err) + require.Len(t, page, 1, "FAILED cleanup has no legacy-age grace") + executor := uuid.New() + _, err = client.ClaimCheckpointOperation(ctx, checkpoint.Id, "alice", executor) + require.NoError(t, err) + require.ErrorIs(t, client.ClearFailedCheckpointTag(ctx, checkpoint.Id, "mallory", "partial-tag", executor), ErrConflict) + require.ErrorIs(t, client.ClearFailedCheckpointTag(ctx, checkpoint.Id, "alice", "replacement", executor), ErrConflict) + require.ErrorIs(t, client.ClearFailedCheckpointTag(ctx, checkpoint.Id, "alice", "partial-tag", creator), ErrConflict) + require.NoError(t, client.ReleaseCheckpointOperation(ctx, checkpoint.Id, executor)) + require.ErrorIs(t, client.ClearFailedCheckpointTag(ctx, checkpoint.Id, "alice", "partial-tag", executor), ErrConflict, "an expired executor cannot record cleanup") + executor = uuid.New() + _, err = client.ClaimCheckpointOperation(ctx, checkpoint.Id, "alice", executor) + require.NoError(t, err) + require.NoError(t, client.ClearFailedCheckpointTag(ctx, checkpoint.Id, "alice", "partial-tag", executor)) + retained, err := client.GetCheckpointOperation(ctx, checkpoint.Id, "alice") + require.NoError(t, err) + require.Empty(t, retained.TagUID) + require.True(t, proto.Equal(failed, retained.Checkpoint), "cleanup keeps the terminal request result") + cleared, err := queryOne(ctx, client.db, `SELECT executor_id IS NULL AND executor_expires_at IS NULL FROM session_checkpoint WHERE id = $1`, pgx.RowTo[bool], checkpoint.Id) + require.NoError(t, err) + require.True(t, cleared) + page, err = client.ListPendingCheckpoints(ctx, "", 100) + require.NoError(t, err) + require.Empty(t, page) + _, err = client.ClaimCheckpointOperation(ctx, checkpoint.Id, "alice", uuid.New()) + require.ErrorIs(t, err, ErrConflict) +} + +func reserveRecoveryCheckpoint(t *testing.T, client *Client) (*apiv1.Checkpoint, *apiv1.Session) { + t.Helper() + ctx := t.Context() + sessionFixture(t, client, ctx, "team-a", "revision", "assistant", "kagent") + session, _, err := client.CreateSession(ctx, newSessionRequest(uuid.NewString(), "assistant", "kagent", ""), uuid.NewString()) + require.NoError(t, err) + _, err = markSessionReady(ctx, client, session.Id, "runtime") + require.NoError(t, err) + task := newSessionTask(uuid.NewString(), uuid.NewString()) + task.ContextID = session.ContextId + task.Status.State = a2a.TaskStateCompleted + require.NoError(t, saveRuntimeTask(t, client, session.Id, task, task, + &SessionTaskSnapshot{Atespace: "team-a", URI: "s3://snapshots/source", Fidelity: "SNAPSHOT_FIDELITY_VOLUMES"})) + checkpoint, _, err := reserveCheckpoint(ctx, client, + &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID)}, "alice", uuid.NewString()) + require.NoError(t, err) + return checkpoint, session +} + +func TestCheckpointClaimsFenceCompletionAndDeletion(t *testing.T) { + pool := setupTestDB(t) + config := pool.Config().Copy() + config.ConnConfig.Tracer = &checkpointClaimTransit{} + claimPool, err := Connect(t.Context(), &PostgresConfig{Pool: config, Schema: "public"}) + require.NoError(t, err) + t.Cleanup(claimPool.Close) + client := NewClient(claimPool, "public") + checkpoint, _ := reserveRecoveryCheckpoint(t, client) + first := uuid.New() + ctx, cancel := context.WithTimeout(t.Context(), 500*time.Millisecond) + defer cancel() + _, err = client.ClaimCheckpointOperation(ctx, checkpoint.Id, "alice", first) + require.NoError(t, err) + nextCtx, stop := context.WithTimeout(t.Context(), RuntimeOperationTimeout) + defer stop() + _, err = client.ClaimCheckpointOperation(nextCtx, checkpoint.Id, "alice", uuid.New()) + require.ErrorIs(t, err, ErrConflict) + _, err = client.ClaimCheckpointOperation(nextCtx, checkpoint.Id, "mallory", uuid.New()) + require.ErrorIs(t, err, ErrConflict) + page, err := client.ListPendingCheckpoints(t.Context(), "", 100) + require.NoError(t, err) + require.Empty(t, page, "live claims are not sweep work") + <-ctx.Done() + second := uuid.New() + // A client deadline does not include transit before PostgreSQL starts its + // lease clock. Observe actual takeover eligibility instead of guessing it. + var op *CheckpointOperation + require.Eventually(t, func() bool { + op, err = client.ClaimCheckpointOperation(nextCtx, checkpoint.Id, "alice", second) + return !errors.Is(err, ErrConflict) + }, 5*time.Second, 10*time.Millisecond, "expired claim must become eligible for takeover") + require.NoError(t, err) + require.NoError(t, client.ReleaseCheckpointOperation(t.Context(), checkpoint.Id, first)) + _, err = client.FinalizeClaimedSessionCheckpoint(nextCtx, op.Checkpoint, first, "old", "s3://tags/old", "") + require.ErrorIs(t, err, ErrConflict) + ready, err := client.FinalizeClaimedSessionCheckpoint(nextCtx, op.Checkpoint, second, "retained", "s3://tags/retained", "") + require.NoError(t, err) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_READY, ready.State) + // Completion must clear both columns in the same write, before any release. + claim, err := queryOne(t.Context(), client.db, `SELECT executor_id IS NULL AND executor_expires_at IS NULL FROM session_checkpoint WHERE id = $1`, pgx.RowTo[bool], checkpoint.Id) + require.NoError(t, err) + require.True(t, claim) + deletion := uuid.New() + deleting, err := client.BeginDeleteSessionCheckpoint(nextCtx, checkpoint.Id, "alice", deletion) + require.NoError(t, err) + require.Equal(t, "retained", deleting.TagUID) + require.ErrorIs(t, client.DeleteClaimedSessionCheckpoint(nextCtx, checkpoint.Id, "alice", second), ErrConflict) + require.NoError(t, client.DeleteClaimedSessionCheckpoint(nextCtx, checkpoint.Id, "alice", deletion)) +} + +// Delay the first claim after its duration was calculated, so the database lease +// deterministically outlives the client timer. Stored timestamps stay untouched. +type checkpointClaimTransit struct{ once sync.Once } + +func (g *checkpointClaimTransit) TraceQueryStart(ctx context.Context, _ *pgx.Conn, data pgx.TraceQueryStartData) context.Context { + if strings.Contains(data.SQL, "UPDATE session_checkpoint SET executor_id") { + g.once.Do(func() { time.Sleep(100 * time.Millisecond) }) + } + return ctx +} + +func (*checkpointClaimTransit) TraceQueryEnd(context.Context, *pgx.Conn, pgx.TraceQueryEndData) {} + +func TestLegacyUnversionedCheckpointCannotBeClaimed(t *testing.T) { + client := NewClient(setupTestDB(t), "public") + checkpoint, _ := reserveRecoveryCheckpoint(t, client) + // Prepare a legacy fixture with no claim: the migration preserves such + // rows, but new Reserve must always insert a non-NULL bounded executor. + _, err := client.db.Exec(t.Context(), `UPDATE session_checkpoint SET executor_id = NULL, executor_expires_at = NULL WHERE id = $1`, checkpoint.Id) + require.NoError(t, err) + ctx, cancel := context.WithTimeout(t.Context(), RuntimeOperationTimeout) + defer cancel() + _, err = client.ClaimCheckpointOperation(ctx, checkpoint.Id, "alice", uuid.New()) + require.ErrorIs(t, err, ErrConflict, "unknown legacy v4 rows cannot be age-fenced") + _, err = client.GetCheckpointOperation(ctx, checkpoint.Id, "mallory") + require.ErrorIs(t, err, ErrNotFound) +} + +func TestLegacyCheckpointGraceAndExpiredClaims(t *testing.T) { + client := NewClient(setupTestDB(t), "public") + for _, age := range []time.Duration{0, 3 * time.Minute} { + checkpoint, _ := reserveRecoveryCheckpoint(t, client) + id, err := uuid.NewV7() + require.NoError(t, err) + binary.BigEndian.PutUint64(id[:8], uint64(time.Now().Add(-age).UnixMilli())<<16) + id[6] = 0x70 + value := proto.CloneOf(checkpoint) + value.Id = id.String() + data, err := proto.Marshal(value) + require.NoError(t, err) + // A legacy reservation is intentionally constructed as migration input. + _, err = client.db.Exec(t.Context(), `UPDATE session_checkpoint SET id = $2, data = $3, executor_id = NULL, executor_expires_at = NULL WHERE id = $1`, checkpoint.Id, id, data) + require.NoError(t, err) + ctx, cancel := context.WithTimeout(t.Context(), RuntimeOperationTimeout) + executor := uuid.New() + page, err := client.ListPendingCheckpoints(ctx, "", 100) + require.NoError(t, err) + op, err := client.ClaimCheckpointOperation(ctx, id.String(), "alice", executor) + if age == 0 { + require.ErrorIs(t, err, ErrConflict) + require.Empty(t, page) + } else { + require.NoError(t, err) + require.Len(t, page, 1) + failed, err := client.FinalizeClaimedSessionCheckpoint(ctx, op.Checkpoint, executor, "partial-tag", "", "copy failed") + require.NoError(t, err) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_FAILED, failed.State) + retained, err := client.GetCheckpointOperation(ctx, id.String(), "alice") + require.NoError(t, err) + require.Equal(t, "partial-tag", retained.TagUID) + require.Equal(t, "s3://snapshots/source", retained.Snapshot.URI) + } + cancel() + } +} diff --git a/go/core/internal/database/checkpoint_reservation_test.go b/go/core/internal/database/checkpoint_reservation_test.go new file mode 100644 index 0000000000..a0d5874e76 --- /dev/null +++ b/go/core/internal/database/checkpoint_reservation_test.go @@ -0,0 +1,71 @@ +package database + +import ( + "context" + "sync" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" +) + +// Same delayed-COMMIT boundary as TestCheckpointReleaseWaitsForClaimCommit, +// applied to the initial reservation INSERT instead of an existing row UPDATE. +func TestCheckpointReleaseWaitsForReservationCommit(t *testing.T) { + pool := setupTestDB(t) + client := NewClient(pool, "public") + seed, source := reserveRecoveryCheckpoint(t, client) + _, err := finalizeCheckpoint(t.Context(), client, seed.Id, "seed", "s3://tags/seed", "") + require.NoError(t, err) + + gate := &checkpointCommitGate{entered: make(chan struct{}), resume: make(chan struct{})} + config := pool.Config().Copy() + config.ConnConfig.Tracer = gate + reservedPool, err := Connect(t.Context(), &PostgresConfig{Pool: config, Schema: "public"}) + require.NoError(t, err) + t.Cleanup(reservedPool.Close) + var once sync.Once + resume := func() { once.Do(func() { close(gate.resume) }) } + t.Cleanup(resume) + reserving := NewClient(reservedPool, "public") + ctx, cancel := context.WithTimeout(t.Context(), RuntimeOperationTimeout) + defer cancel() + executor, id, request := uuid.New(), uuid.NewString(), uuid.NewString() + reserveDone := make(chan error, 1) + go func() { + _, _, err := reserving.ReserveSessionCheckpoint(ctx, &apiv1.Checkpoint{Id: id, SessionId: source.Id, HeadTaskId: seed.HeadTaskId}, "alice", request, executor) + reserveDone <- err + }() + select { + case <-gate.entered: + case <-ctx.Done(): + t.Fatal(ctx.Err()) + } + releaseDone := make(chan error, 1) + go func() { releaseDone <- client.ReleaseCheckpointReservation(ctx, source.Id, id, executor) }() + early := false + select { + case err := <-releaseDone: + require.NoError(t, err) + early = true + case <-time.After(100 * time.Millisecond): + } + resume() + require.NoError(t, <-reserveDone) + if !early { + require.NoError(t, <-releaseDone) + } + replacement := uuid.New() + op, err := client.ClaimCheckpointOperation(ctx, id, "alice", replacement) + require.False(t, early, "release passed an invisible INSERT instead of waiting for its commit") + require.NoError(t, err, "a late reservation commit must not restore a live claim after release returned") + require.NoError(t, client.ReleaseCheckpointReservation(ctx, source.Id, id, executor)) + _, err = client.ClaimCheckpointOperation(ctx, id, "alice", uuid.New()) + require.ErrorIs(t, err, ErrConflict, "an obsolete reservation release must not revoke the replacement") + _, err = client.FinalizeClaimedSessionCheckpoint(ctx, op.Checkpoint, replacement, "retained", "s3://tags/retained", "") + require.NoError(t, err) + require.NoError(t, client.ReleaseCheckpointReservation(ctx, uuid.NewString(), uuid.NewString(), executor), "a missing source and reservation are harmless") +} diff --git a/go/core/internal/database/checkpoints.go b/go/core/internal/database/checkpoints.go index d2df87d782..a9077fbce3 100644 --- a/go/core/internal/database/checkpoints.go +++ b/go/core/internal/database/checkpoints.go @@ -6,6 +6,7 @@ import ( "fmt" "slices" "strings" + "time" "github.com/a2aproject/a2a-go/v2/a2a" "github.com/google/uuid" @@ -181,11 +182,16 @@ func (c *Client) ForkSession(ctx context.Context, checkpointID, userID, requestI // ReserveSessionCheckpoint captures the terminal task named by HeadTaskId, // provided it is still the latest boundary. ErrSnapshotPending permits retry; // ErrCheckpointAdvanced requires a new selection. The session must be READY. -// A repeated owner/requestID returns the same reservation for the same source -// and task, or ErrIdempotencyConflict otherwise. A +// The INSERT takes executor ownership under the caller deadline. A repeated +// owner/requestID returns the same reservation for the same source and task +// without taking a new claim, or ErrIdempotencyConflict otherwise. A // CREATING reservation blocks new task writes and lifecycle operations while the caller // retains the external snapshot. -func (c *Client) ReserveSessionCheckpoint(ctx context.Context, checkpoint *apiv1.Checkpoint, userID, requestID string) (*apiv1.Checkpoint, *SessionTaskSnapshot, error) { +func (c *Client) ReserveSessionCheckpoint(ctx context.Context, checkpoint *apiv1.Checkpoint, userID, requestID string, executor uuid.UUID) (*apiv1.Checkpoint, *SessionTaskSnapshot, error) { + deadline, ok := ctx.Deadline() + if !ok || executor == uuid.Nil || time.Until(deadline) > RuntimeOperationTimeout { + return nil, nil, ErrFailedPrecondition + } if checkpoint == nil { return nil, nil, fmt.Errorf("missing checkpoint") } @@ -289,8 +295,8 @@ func (c *Client) ReserveSessionCheckpoint(ctx context.Context, checkpoint *apiv1 row, err := queryOne(ctx, tx, ` INSERT INTO session_checkpoint (id, source_session_id, user_id, request_id, head_task_id, history_sequence, snapshot_atespace, snapshot_uri, snapshot_fidelity, source_history_id, - prepared_revision, data, state) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, - $10, $11, $12, 'CREATING') + prepared_revision, data, state, executor_id, executor_expires_at) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, + $10, $11, $12, 'CREATING', $13, clock_timestamp() + $14 * interval '1 millisecond') ON CONFLICT DO NOTHING 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 @@ -298,12 +304,12 @@ func (c *Client) ReserveSessionCheckpoint(ctx context.Context, checkpoint *apiv1 pgx.RowToStructByName[sessionCheckpointRow], checkpoint.GetId(), checkpoint.GetSessionId(), userID, requestID, boundary.ID, *boundary.HistorySequence, *boundary.SnapshotAtespace, *boundary.SnapshotURI, *boundary.SnapshotFidelity, session.HistoryID, - session.PreparedRevision, data, + session.PreparedRevision, data, executor, time.Until(deadline).Milliseconds()+1, ) if errors.Is(err, pgx.ErrNoRows) { existing, existingErr := readCheckpointRequest(ctx, tx, userID, requestID) if existingErr == nil { - if existing.SourceSessionID != sourceID { + if existing.SourceSessionID != sourceID || existing.HeadTaskID != checkpoint.HeadTaskId { return ErrIdempotencyConflict } snapshot = checkpointSnapshot(existing) @@ -328,67 +334,53 @@ func (c *Client) ReserveSessionCheckpoint(ctx context.Context, checkpoint *apiv1 return result, snapshot, nil } -// FinalizeSessionCheckpoint records either a retained tag UID and snapshot URI as -// READY, or a failure as FAILED. An identical terminal retry returns the existing -// checkpoint; a different terminal result returns ErrNotFound. This internal operation -// does not check ownership or create the external snapshot. -func (c *Client) FinalizeSessionCheckpoint(ctx context.Context, id, tagUID, snapshotURI, failure string) (*apiv1.Checkpoint, error) { - if (tagUID == "") == (failure == "") || (tagUID == "") != (snapshotURI == "") { - return nil, fmt.Errorf("finalize Session checkpoint requires tag UID and snapshot URI, or failure") +// FinalizeClaimedSessionCheckpoint atomically publishes READY or FAILED and clears +// the claim. Only a current, unexpired executor can complete CREATING work. +// A failed copy retains its tag UID for UID-conditional deletion, even if its +// storage is incomplete. Its saved URI remains the source snapshot in that case. +func (c *Client) FinalizeClaimedSessionCheckpoint(ctx context.Context, checkpoint *apiv1.Checkpoint, executor uuid.UUID, tagUID, snapshotURI, failure string) (*apiv1.Checkpoint, error) { + if checkpoint == nil || executor == uuid.Nil || (failure == "" && (tagUID == "" || snapshotURI == "")) { + return nil, ErrFailedPrecondition } - var result *apiv1.Checkpoint - err := c.withTx(ctx, func(tx pgx.Tx) error { - row, err := lockCheckpoint(ctx, tx, id, true, "", nil) - if err != nil { - return notFoundOr(err) - } - result, err = toSessionCheckpoint(row) - if err != nil { - return err - } - if row.State != "CREATING" { - if (row.State == "READY" && row.TagUID == tagUID && row.SnapshotURI == snapshotURI && failure == "") || - (row.State == "FAILED" && tagUID == "" && result.GetFailure().GetMessage() == failure) { - return nil - } - return ErrNotFound - } - result.State = apiv1.CheckpointState_CHECKPOINT_STATE_READY - if failure != "" { - result.State = apiv1.CheckpointState_CHECKPOINT_STATE_FAILED - result.Failure = &apiv1.Failure{Reason: "SnapshotTagFailed", Message: failure} - } - data, err := proto.Marshal(result) - if err != nil { - return fmt.Errorf("encode checkpoint: %w", err) - } - tag, err := tx.Exec(ctx, ` - UPDATE session_checkpoint - SET state = CASE WHEN $2::text <> '' THEN 'READY' ELSE 'FAILED' END, - tag_uid = $2, - snapshot_uri = CASE WHEN $2::text <> '' THEN $3::text ELSE snapshot_uri END, - data = $4 - WHERE id = $1 - AND state = 'CREATING' - `, row.ID, tagUID, snapshotURI, data) - if err != nil { - return err - } - if tag.RowsAffected() != 1 { - return ErrNotFound - } - return nil - }) + result := proto.Clone(checkpoint).(*apiv1.Checkpoint) + result.State = apiv1.CheckpointState_CHECKPOINT_STATE_READY + result.Failure = nil + state := "READY" + if failure != "" { + state = "FAILED" + result.State = apiv1.CheckpointState_CHECKPOINT_STATE_FAILED + result.Failure = &apiv1.Failure{Reason: "SnapshotTagFailed", Message: failure} + } + data, err := proto.Marshal(result) + if err != nil { + return nil, fmt.Errorf("encode checkpoint: %w", err) + } + tag, err := c.db.Exec(ctx, ` + UPDATE session_checkpoint + SET state = $2, tag_uid = $3, + snapshot_uri = CASE WHEN $4::text <> '' THEN $4::text ELSE snapshot_uri END, + data = $5, executor_id = NULL, executor_expires_at = NULL + WHERE id = $1 AND state = 'CREATING' + AND executor_id = $6 AND executor_expires_at > clock_timestamp() + `, checkpoint.Id, state, tagUID, snapshotURI, data, executor) if err != nil { return nil, fmt.Errorf("finalize Session checkpoint: %w", err) } + if tag.RowsAffected() != 1 { + return nil, ErrConflict + } return result, nil } -// GetSessionCheckpoint returns an owned READY checkpoint. Missing checkpoints, other -// owners, and other lifecycle states return ErrNotFound. +// GetSessionCheckpoint returns an owned READY checkpoint. All other lifecycle +// states and other owners return ErrNotFound. func (c *Client) GetSessionCheckpoint(ctx context.Context, id, userID string) (*apiv1.Checkpoint, error) { - row, err := readCheckpoint(ctx, c.db, id, userID, new("READY")) + 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 AND state = 'READY' + `, pgx.RowToStructByName[sessionCheckpointRow], id, userID) if err != nil { return nil, fmt.Errorf("get Session checkpoint: %w", notFoundOr(err)) } @@ -454,8 +446,7 @@ func checkpointSnapshot(row sessionCheckpointRow) *SessionTaskSnapshot { } // ListSessionCheckpoints returns an owner's READY checkpoints for the source -// session in ascending ID order after afterID, up to limit. Retained checkpoints remain -// listable after the source session is deleted. +// session in ascending ID order. Retained checkpoints outlive source deletion. func (c *Client) ListSessionCheckpoints(ctx context.Context, sessionID, userID, afterID string, limit int) ([]*apiv1.Checkpoint, error) { rows, err := queryMany(ctx, c.db, ` SELECT id, source_session_id, user_id, request_id, head_task_id, history_sequence, snapshot_atespace, @@ -484,18 +475,24 @@ func (c *Client) ListSessionCheckpoints(ctx context.Context, sessionID, userID, return result, nil } -// BeginDeleteSessionCheckpoint marks an owned READY or FAILED checkpoint DELETING and returns -// the snapshot and tag identity for external cleanup. Retrying a DELETING checkpoint is -// allowed. A fork reference, another lifecycle state, or a missing/unowned checkpoint -// returns ErrNotFound. -func (c *Client) BeginDeleteSessionCheckpoint(ctx context.Context, id, userID string) (*SessionTaskSnapshot, string, error) { - var snapshot *SessionTaskSnapshot - var tagUID string +// BeginDeleteSessionCheckpoint atomically marks an owned READY or FAILED +// checkpoint DELETING and takes its initial bounded claim. DELETING returns +// ErrConflict so a retry can acquire expired work. A checkpoint with dependent +// forks cannot be deleted; missing, unowned or ineligible rows return ErrNotFound. +func (c *Client) BeginDeleteSessionCheckpoint(ctx context.Context, id, userID string, executor uuid.UUID) (*CheckpointOperation, error) { + deadline, ok := ctx.Deadline() + if !ok || executor == uuid.Nil || time.Until(deadline) > RuntimeOperationTimeout { + return nil, ErrFailedPrecondition + } + var operation *CheckpointOperation err := c.withTx(ctx, func(tx pgx.Tx) error { row, err := lockCheckpoint(ctx, tx, id, false, userID, nil) if err != nil { return notFoundOr(err) } + if row.State == "DELETING" { + return ErrConflict + } checkpoint, err := toSessionCheckpoint(row) if err != nil { return err @@ -506,39 +503,48 @@ func (c *Client) BeginDeleteSessionCheckpoint(ctx context.Context, id, userID st return fmt.Errorf("encode checkpoint: %w", err) } tag, err := tx.Exec(ctx, ` - UPDATE session_checkpoint - SET state = 'DELETING', data = $3 - WHERE session_checkpoint.id = $1 AND session_checkpoint.user_id = $2 - AND session_checkpoint.state IN ('READY', 'FAILED', 'DELETING') - AND NOT EXISTS ( - SELECT 1 FROM session i WHERE i.pinned_checkpoint_id = session_checkpoint.id - ) - `, row.ID, userID, data) + UPDATE session_checkpoint + SET state = 'DELETING', data = $3, executor_id = $4, + executor_expires_at = clock_timestamp() + $5 * interval '1 millisecond' + WHERE session_checkpoint.id = $1 AND session_checkpoint.user_id = $2 + AND session_checkpoint.state IN ('READY', 'FAILED') + AND NOT EXISTS ( + SELECT 1 FROM session i WHERE i.pinned_checkpoint_id = session_checkpoint.id + ) + `, row.ID, userID, data, executor, time.Until(deadline).Milliseconds()+1) if err != nil { return err } if tag.RowsAffected() != 1 { return ErrNotFound } - snapshot, tagUID = checkpointSnapshot(row), row.TagUID + operation = &CheckpointOperation{Checkpoint: checkpoint, Snapshot: checkpointSnapshot(row), TagUID: row.TagUID, Creator: userID} return nil }) if err != nil { - return nil, "", fmt.Errorf("begin delete Session checkpoint: %w", err) + return nil, fmt.Errorf("begin delete Session checkpoint: %w", err) } - return snapshot, tagUID, nil + return operation, nil } -// DeleteSessionCheckpoint removes an owned DELETING checkpoint after the caller -// cleans up its external tag. No matching row is a successful no-op. -func (c *Client) DeleteSessionCheckpoint(ctx context.Context, id, userID string) error { - _, err := c.db.Exec(ctx, ` - DELETE FROM session_checkpoint - WHERE id = $1 AND user_id = $2 AND state = 'DELETING' - `, id, userID) +// DeleteClaimedSessionCheckpoint removes an owned DELETING checkpoint after +// external cleanup. Only its current, unexpired executor can remove the row; +// a missing or superseded operation returns ErrConflict for the caller to reread. +func (c *Client) DeleteClaimedSessionCheckpoint(ctx context.Context, id, userID string, executor uuid.UUID) error { + if executor == uuid.Nil { + return ErrFailedPrecondition + } + tag, err := c.db.Exec(ctx, ` + DELETE FROM session_checkpoint + WHERE id = $1 AND user_id = $2 AND state = 'DELETING' + AND executor_id = $3 AND executor_expires_at > clock_timestamp() + `, id, userID, executor) if err != nil { return fmt.Errorf("delete Session checkpoint: %w", err) } + if tag.RowsAffected() != 1 { + return ErrConflict + } return nil } diff --git a/go/core/internal/database/client_session_test.go b/go/core/internal/database/client_session_test.go index 21973bd144..ac905ba220 100644 --- a/go/core/internal/database/client_session_test.go +++ b/go/core/internal/database/client_session_test.go @@ -29,7 +29,7 @@ func TestMalformedDatabaseIDsReturnErrors(t *testing.T) { {"delete session", func() error { return deleteSession(ctx, client, "invalid") }}, {"get checkpoint", func() error { _, err := client.GetSessionCheckpoint(ctx, "invalid", "alice"); return err }}, {"reserve checkpoint", func() error { - _, _, err := client.ReserveSessionCheckpoint(ctx, &apiv1.Checkpoint{SessionId: "invalid", HeadTaskId: "task-1"}, "alice", "request") + _, _, err := reserveCheckpoint(ctx, client, &apiv1.Checkpoint{SessionId: "invalid", HeadTaskId: "task-1"}, "alice", "request") return err }}, {"fork", func() error { @@ -228,11 +228,11 @@ func TestSessionCheckpointRetainsRootFSFidelity(t *testing.T) { task.Status.State = a2a.TaskStateCompleted want := &SessionTaskSnapshot{Atespace: "team-a", URI: "s3://snapshots/rootfs", Fidelity: "SNAPSHOT_FIDELITY_ROOTFS"} require.NoError(t, saveRuntimeTask(t, client, session.Id, task, task, want)) - checkpoint, snapshot, err := client.ReserveSessionCheckpoint(ctx, + checkpoint, snapshot, err := reserveCheckpoint(ctx, client, &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID)}, "alice", uuid.NewString()) require.NoError(t, err) require.Equal(t, want, snapshot) - _, err = client.FinalizeSessionCheckpoint(ctx, checkpoint.Id, "tag-uid", "s3://tags/rootfs", "") + _, err = finalizeCheckpoint(ctx, client, checkpoint.Id, "tag-uid", "s3://tags/rootfs", "") require.NoError(t, err) retained, _, err := client.GetSessionCheckpointSnapshot(ctx, checkpoint.Id, "alice") require.NoError(t, err) @@ -263,7 +263,7 @@ func TestSessionCheckpointRetainsRecordedBoundary(t *testing.T) { t.Fatal(err) } - checkpoint, snapshot, err := client.ReserveSessionCheckpoint(ctx, &apiv1.Checkpoint{Id: "22222222-2222-4222-8222-222222222222", SessionId: sessionID, HeadTaskId: "task-1"}, "alice", "checkpoint-request") + checkpoint, snapshot, err := reserveCheckpoint(ctx, client, &apiv1.Checkpoint{Id: "22222222-2222-4222-8222-222222222222", SessionId: sessionID, HeadTaskId: "task-1"}, "alice", "checkpoint-request") if err != nil { t.Fatalf("ReserveSessionCheckpoint() = %+v, error %v", checkpoint, err) } @@ -294,11 +294,11 @@ func TestSessionCheckpointRetainsRecordedBoundary(t *testing.T) { } _, err = client.BeginSessionOperation(ctx, sessionID, apiv1.RuntimeOperation_RUNTIME_OPERATION_SUSPEND, false) require.ErrorIs(t, err, ErrConflict) - replayed, replayedSnapshot, err := client.ReserveSessionCheckpoint(ctx, &apiv1.Checkpoint{Id: "33333333-3333-4333-8333-333333333333", SessionId: sessionID, HeadTaskId: "task-1"}, "alice", "checkpoint-request") + replayed, replayedSnapshot, err := reserveCheckpoint(ctx, client, &apiv1.Checkpoint{Id: "33333333-3333-4333-8333-333333333333", SessionId: sessionID, HeadTaskId: "task-1"}, "alice", "checkpoint-request") if err != nil || replayed.Id != checkpoint.Id || replayedSnapshot == nil || *replayedSnapshot != *snapshot { t.Fatalf("replayed checkpoint = %+v, error %v", replayed, err) } - ready, err := client.FinalizeSessionCheckpoint(ctx, checkpoint.GetId(), "tag-uid", "s3://tags/checkpoint", "") + ready, err := finalizeCheckpoint(ctx, client, checkpoint.GetId(), "tag-uid", "s3://tags/checkpoint", "") if err != nil || ready.State != apiv1.CheckpointState_CHECKPOINT_STATE_READY { t.Fatalf("ready checkpoint = %+v, error %v", ready, err) } @@ -321,14 +321,14 @@ func TestSessionCheckpointRetainsRecordedBoundary(t *testing.T) { if err != nil || tagUID != "tag-uid" || retained.URI != "s3://tags/checkpoint" || retained.Fidelity != snapshot.Fidelity { t.Fatalf("checkpoint tag = %q, error %v", tagUID, err) } - if replayed, err := client.FinalizeSessionCheckpoint(ctx, checkpoint.GetId(), "tag-uid", "s3://tags/checkpoint", ""); err != nil || replayed.State != apiv1.CheckpointState_CHECKPOINT_STATE_READY { + if replayed, _, err := reserveCheckpoint(ctx, client, &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: sessionID, HeadTaskId: "task-1"}, "alice", "checkpoint-request"); err != nil || replayed.State != apiv1.CheckpointState_CHECKPOINT_STATE_READY { t.Fatalf("replayed ready checkpoint = %+v, error %v", replayed, err) } - failed, _, err := client.ReserveSessionCheckpoint(ctx, &apiv1.Checkpoint{Id: "44444444-4444-4444-8444-444444444444", SessionId: sessionID, HeadTaskId: "task-1"}, "alice", "failed-checkpoint-request") + failed, _, err := reserveCheckpoint(ctx, client, &apiv1.Checkpoint{Id: "44444444-4444-4444-8444-444444444444", SessionId: sessionID, HeadTaskId: "task-1"}, "alice", "failed-checkpoint-request") if err != nil { t.Fatal(err) } - failed, err = client.FinalizeSessionCheckpoint(ctx, failed.GetId(), "", "", "tag creation failed") + failed, err = finalizeCheckpoint(ctx, client, failed.GetId(), "", "", "tag creation failed") if err != nil || failed.State != apiv1.CheckpointState_CHECKPOINT_STATE_FAILED || failed.GetFailure().GetMessage() != "tag creation failed" { t.Fatalf("failed checkpoint = %+v, error %v", failed, err) } @@ -338,7 +338,7 @@ func TestSessionCheckpointRetainsRecordedBoundary(t *testing.T) { if err := deleteSession(ctx, client, sessionID); err != nil { t.Fatal(err) } - replayed, replayedSnapshot, err = client.ReserveSessionCheckpoint(ctx, &apiv1.Checkpoint{Id: "55555555-5555-4555-8555-555555555555", SessionId: sessionID, HeadTaskId: "task-1"}, "alice", "checkpoint-request") + replayed, replayedSnapshot, err = reserveCheckpoint(ctx, client, &apiv1.Checkpoint{Id: "55555555-5555-4555-8555-555555555555", SessionId: sessionID, HeadTaskId: "task-1"}, "alice", "checkpoint-request") if err != nil || replayed.Id != checkpoint.Id || replayedSnapshot == nil || *replayedSnapshot != *retained { t.Fatalf("checkpoint replay after source deletion = %+v, error %v", replayed, err) } @@ -346,10 +346,18 @@ func TestSessionCheckpointRetainsRecordedBoundary(t *testing.T) { if err != nil || len(listed) != 1 || listed[0].Id != checkpoint.Id { t.Fatalf("listed checkpoints = %+v, error %v", listed, err) } - if ref, tag, err := client.BeginDeleteSessionCheckpoint(ctx, checkpoint.GetId(), "mallory"); !errors.Is(err, ErrNotFound) || ref != nil || tag != "" { + page, err := client.ListSessionCheckpoints(ctx, sessionID, "alice", checkpoint.Id, 1) + require.NoError(t, err) + require.Empty(t, page, "pagination must not expose internal FAILED rows") + _, err = client.GetSessionCheckpoint(ctx, failed.Id, "alice") + require.ErrorIs(t, err, ErrNotFound) + foreign, err := client.ListSessionCheckpoints(ctx, sessionID, "mallory", "", 10) + require.NoError(t, err) + require.Empty(t, foreign) + if ref, tag, err := beginDeleteCheckpoint(ctx, client, checkpoint.GetId(), "mallory"); !errors.Is(err, ErrNotFound) || ref != nil || tag != "" { t.Fatalf("unauthorized deletion = %+v, %q, %v", ref, tag, err) } - deletingSnapshot, deletingTagUID, err := client.BeginDeleteSessionCheckpoint(ctx, checkpoint.GetId(), "alice") + deletingSnapshot, deletingTagUID, err := beginDeleteCheckpoint(ctx, client, checkpoint.GetId(), "alice") if err != nil || deletingSnapshot == nil || *deletingSnapshot != *retained || deletingTagUID != "tag-uid" { t.Fatalf("deleting snapshot = %+v, tag = %q, error %v", deletingSnapshot, deletingTagUID, err) } @@ -359,11 +367,11 @@ func TestSessionCheckpointRetainsRecordedBoundary(t *testing.T) { if _, err := client.UpdateCheckpointName(ctx, checkpoint.GetId(), "alice", "too late"); !errors.Is(err, ErrNotFound) { t.Fatalf("rename of a deleting checkpoint = %v, want not found", err) } - deletingSnapshot, deletingTagUID, err = client.BeginDeleteSessionCheckpoint(ctx, checkpoint.GetId(), "alice") + deletingSnapshot, deletingTagUID, err = beginDeleteCheckpoint(ctx, client, checkpoint.GetId(), "alice") if err != nil || deletingSnapshot == nil || *deletingSnapshot != *retained || deletingTagUID != "tag-uid" { t.Fatalf("deleting snapshot = %+v, tag = %q, error %v", deletingSnapshot, deletingTagUID, err) } - if err := client.DeleteSessionCheckpoint(ctx, checkpoint.GetId(), "alice"); err != nil { + if err := deleteCheckpoint(ctx, client, checkpoint.GetId(), "alice"); err != nil { t.Fatal(err) } } @@ -374,7 +382,7 @@ func TestReserveSessionCheckpointRejectsCorruptSource(t *testing.T) { client := NewClient(db, "public") sessionID := "99999999-9999-4999-8999-999999999999" insertReadySessionFixture(t, db, sessionID, "session-request", []byte{0xff}) - _, _, err := client.ReserveSessionCheckpoint(ctx, &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: sessionID}, "alice", "checkpoint-request") + _, _, err := reserveCheckpoint(ctx, client, &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: sessionID}, "alice", "checkpoint-request") require.ErrorContains(t, err, "decode Session") } @@ -429,7 +437,7 @@ func TestForkSessionCopiesBoundedHistory(t *testing.T) { &SessionTaskSnapshot{Atespace: "team-a", URI: "s3://snapshots/snapshot-1", Fidelity: "SNAPSHOT_FIDELITY_VOLUMES"}); err != nil { t.Fatal(err) } - _, _, err = client.ReserveSessionCheckpoint(ctx, &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: source.GetId(), HeadTaskId: "task-1"}, "alice", "hitl-checkpoint-request") + _, _, err = reserveCheckpoint(ctx, client, &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: source.GetId(), HeadTaskId: "task-1"}, "alice", "hitl-checkpoint-request") require.ErrorIs(t, err, ErrFailedPrecondition) first.Status.State = a2a.TaskStateCompleted @@ -437,13 +445,13 @@ func TestForkSessionCopiesBoundedHistory(t *testing.T) { &SessionTaskSnapshot{Atespace: "team-a", URI: "s3://snapshots/snapshot-1", Fidelity: "SNAPSHOT_FIDELITY_VOLUMES"}); err != nil { t.Fatal(err) } - checkpoint, _, err := client.ReserveSessionCheckpoint(ctx, &apiv1.Checkpoint{Id: "99999999-9999-4999-8999-999999999999", SessionId: source.GetId(), HeadTaskId: "task-1"}, "alice", "checkpoint-request-1") + checkpoint, _, err := reserveCheckpoint(ctx, client, &apiv1.Checkpoint{Id: "99999999-9999-4999-8999-999999999999", SessionId: source.GetId(), HeadTaskId: "task-1"}, "alice", "checkpoint-request-1") if err != nil { t.Fatal(err) } _, err = client.GetSessionCheckpointForkSource(ctx, checkpoint.GetId(), "alice") require.ErrorIs(t, err, ErrNotFound, "creating checkpoints cannot supply fork provenance") - if _, err := client.FinalizeSessionCheckpoint(ctx, checkpoint.GetId(), "tag-uid-1", "s3://tags/checkpoint", ""); err != nil { + if _, err := finalizeCheckpoint(ctx, client, checkpoint.GetId(), "tag-uid-1", "s3://tags/checkpoint", ""); err != nil { t.Fatal(err) } _, err = client.UpdateSessionName(ctx, source.GetId(), "alice", "Renamed after checkpoint") @@ -525,18 +533,18 @@ func TestForkSessionCopiesBoundedHistory(t *testing.T) { if _, _, err := client.ForkSession(ctx, "aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa", "alice", "fork-request-1", "ignored"); !errors.Is(err, ErrIdempotencyConflict) { t.Fatalf("conflicting fork request error = %v", err) } - if _, _, err := client.BeginDeleteSessionCheckpoint(ctx, checkpoint.GetId(), "alice"); !errors.Is(err, ErrNotFound) { + if _, _, err := beginDeleteCheckpoint(ctx, client, checkpoint.GetId(), "alice"); !errors.Is(err, ErrNotFound) { t.Fatalf("delete referenced checkpoint error = %v", err) } if _, err := markSessionReady(ctx, client, fork.GetId(), "fork.example"); err != nil { t.Fatal(err) } - checkpoint2, _, err := client.ReserveSessionCheckpoint(ctx, &apiv1.Checkpoint{Id: "bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb", SessionId: fork.GetId(), HeadTaskId: string(copied.ID)}, "alice", "checkpoint-request-2") + checkpoint2, _, err := reserveCheckpoint(ctx, client, &apiv1.Checkpoint{Id: "bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb", SessionId: fork.GetId(), HeadTaskId: string(copied.ID)}, "alice", "checkpoint-request-2") if err != nil { t.Fatal(err) } - if _, err := client.FinalizeSessionCheckpoint(ctx, checkpoint2.GetId(), "tag-uid-2", "s3://tags/checkpoint", ""); err != nil { + if _, err := finalizeCheckpoint(ctx, client, checkpoint2.GetId(), "tag-uid-2", "s3://tags/checkpoint", ""); err != nil { t.Fatal(err) } if err := deleteSession(ctx, client, fork.GetId()); err != nil { @@ -549,9 +557,9 @@ func TestForkSessionCopiesBoundedHistory(t *testing.T) { } // Deleting a fork releases its checkpoint pin but preserves its original // request identity, even after that checkpoint itself has been removed. - _, _, err = client.BeginDeleteSessionCheckpoint(ctx, checkpoint.GetId(), "alice") + _, _, err = beginDeleteCheckpoint(ctx, client, checkpoint.GetId(), "alice") require.NoError(t, err) - require.NoError(t, client.DeleteSessionCheckpoint(ctx, checkpoint.GetId(), "alice")) + require.NoError(t, deleteCheckpoint(ctx, client, checkpoint.GetId(), "alice")) _, _, err = client.ForkSession(ctx, checkpoint.GetId(), "alice", "fork-request-1", uuid.NewString()) require.ErrorIs(t, err, ErrFailedPrecondition) _, _, err = client.ForkSession(ctx, checkpoint2.GetId(), "alice", "fork-request-1", uuid.NewString()) @@ -826,14 +834,14 @@ func TestForkTaskOrderAndAuthorityIsolation(t *testing.T) { require.NoError(t, saveRuntimeTask(t, client, source.GetId(), task, task, &SessionTaskSnapshot{Atespace: "team-a", URI: "snapshot-" + id, Fidelity: "SNAPSHOT_FIDELITY_VOLUMES"})) } - checkpoint, _, err := client.ReserveSessionCheckpoint(ctx, &apiv1.Checkpoint{ + checkpoint, _, err := reserveCheckpoint(ctx, client, &apiv1.Checkpoint{ Id: uuid.NewString(), SessionId: source.GetId(), HeadTaskId: "a-second", }, "alice", uuid.NewString()) require.NoError(t, err) waiting, err := client.GetSessionTask(ctx, source.GetId(), "z-first", nil) require.NoError(t, err) require.ErrorIs(t, saveRuntimeTask(t, client, source.GetId(), waiting, waiting, nil), ErrFailedPrecondition) - _, err = client.FinalizeSessionCheckpoint(ctx, checkpoint.GetId(), "tag", "tag-snapshot", "") + _, err = finalizeCheckpoint(ctx, client, checkpoint.GetId(), "tag", "tag-snapshot", "") require.NoError(t, err) _, _, err = client.ForkSession(ctx, checkpoint.GetId(), "bob", uuid.NewString(), uuid.NewString()) require.ErrorIs(t, err, ErrNotFound) @@ -863,12 +871,12 @@ func TestForkTaskOrderAndAuthorityIsolation(t *testing.T) { } } // A fork can itself be checkpointed without losing task chronology. - nested, snapshot, err := client.ReserveSessionCheckpoint(ctx, &apiv1.Checkpoint{ + nested, snapshot, err := reserveCheckpoint(ctx, client, &apiv1.Checkpoint{ Id: uuid.NewString(), SessionId: fork.GetId(), HeadTaskId: forkHead, }, "alice", uuid.NewString()) require.NoError(t, err) require.Equal(t, "tag-snapshot", snapshot.URI) - _, err = client.FinalizeSessionCheckpoint(ctx, nested.GetId(), "nested-tag", "nested-snapshot", "") + _, err = finalizeCheckpoint(ctx, client, nested.GetId(), "nested-tag", "nested-snapshot", "") require.NoError(t, err) fork2, _, err := client.ForkSession(ctx, nested.GetId(), "alice", uuid.NewString(), uuid.NewString()) require.NoError(t, err) diff --git a/go/core/internal/database/protobuf_test.go b/go/core/internal/database/protobuf_test.go index bd94210782..b9f01ce4ee 100644 --- a/go/core/internal/database/protobuf_test.go +++ b/go/core/internal/database/protobuf_test.go @@ -177,9 +177,9 @@ func TestProtobufPersistenceLifecycle(t *testing.T) { require.NoError(t, saveRuntimeTask(t, client, session.Id, task, &a2a.TaskStatusUpdateEvent{TaskID: task.ID, ContextID: task.ContextID, Status: task.Status}, &SessionTaskSnapshot{Atespace: "team-a", URI: "s3://snapshots/snapshot", Fidelity: "SNAPSHOT_FIDELITY_VOLUMES"})) checkpointRequest := &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID)} addUnknown(checkpointRequest) - checkpoint, _, err := client.ReserveSessionCheckpoint(ctx, checkpointRequest, "alice", "checkpoint") + checkpoint, _, err := reserveCheckpoint(ctx, client, checkpointRequest, "alice", "checkpoint") require.NoError(t, err) - checkpoint, err = client.FinalizeSessionCheckpoint(ctx, checkpoint.Id, "tag-uid", "s3://tags/checkpoint", "") + checkpoint, err = finalizeCheckpoint(ctx, client, checkpoint.Id, "tag-uid", "s3://tags/checkpoint", "") require.NoError(t, err) checkpointRow, err := readCheckpoint(ctx, q, checkpoint.Id, "alice", new("READY")) require.NoError(t, err) @@ -214,7 +214,7 @@ func TestProtobufPersistenceLifecycle(t *testing.T) { require.Equal(t, futureTask.Status.ProtoReflect().GetUnknown(), forkTask.Status.ProtoReflect().GetUnknown()) require.Equal(t, a2apb.TaskState_TASK_STATE_COMPLETED, forkTask.Status.State) require.NoError(t, deleteSession(ctx, client, fork.Id)) - _, _, err = client.BeginDeleteSessionCheckpoint(ctx, checkpoint.Id, "alice") + _, _, err = beginDeleteCheckpoint(ctx, client, checkpoint.Id, "alice") require.NoError(t, err) checkpointRow, err = readCheckpoint(ctx, q, checkpoint.Id, "alice", nil) require.NoError(t, err) @@ -222,7 +222,7 @@ func TestProtobufPersistenceLifecycle(t *testing.T) { require.NoError(t, proto.Unmarshal(checkpointRow.Data, deleting)) require.Equal(t, checkpointRequest.ProtoReflect().GetUnknown(), deleting.ProtoReflect().GetUnknown()) require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_DELETING, deleting.State) - require.NoError(t, client.DeleteSessionCheckpoint(ctx, checkpoint.Id, "alice")) + require.NoError(t, deleteCheckpoint(ctx, client, checkpoint.Id, "alice")) _, err = client.GetSessionCheckpoint(ctx, checkpoint.Id, "alice") require.ErrorIs(t, err, ErrNotFound) } diff --git a/go/core/internal/database/runtime_revisions_test.go b/go/core/internal/database/runtime_revisions_test.go index 5d4fe8b3ee..6022e663eb 100644 --- a/go/core/internal/database/runtime_revisions_test.go +++ b/go/core/internal/database/runtime_revisions_test.go @@ -101,6 +101,7 @@ func TestRuntimeRevisionCollectionPreservesSessionAndCheckpoint(t *testing.T) { }{ {name: "ready", tagUID: "tag-uid", snapshotURI: "s3://tags/checkpoint"}, {name: "failed", failure: "tag creation failed"}, + {name: "partial tag", tagUID: "partial-uid", failure: "copy failed"}, } { t.Run(tc.name, func(t *testing.T) { client := NewClient(setupTestDB(t), "public") @@ -133,7 +134,7 @@ func TestRuntimeRevisionCollectionPreservesSessionAndCheckpoint(t *testing.T) { task.Status.State = a2a.TaskStateCompleted require.NoError(t, saveRuntimeTask(t, client, session.GetId(), task, task, &SessionTaskSnapshot{Atespace: "team-a", URI: "s3://snapshots/task", Fidelity: "SNAPSHOT_FIDELITY_MEMORY"})) - checkpoint, _, err := client.ReserveSessionCheckpoint(ctx, &apiv1.Checkpoint{ + checkpoint, _, err := reserveCheckpoint(ctx, client, &apiv1.Checkpoint{ Id: uuid.NewString(), SessionId: session.GetId(), HeadTaskId: string(task.ID), }, "alice", "checkpoint") require.NoError(t, err) @@ -144,24 +145,24 @@ func TestRuntimeRevisionCollectionPreservesSessionAndCheckpoint(t *testing.T) { // A checkpoint retains the runtime throughout creation, use, and deletion, // even after the source session and active template pair are gone. assertRetained() - _, _, err = client.BeginDeleteSessionCheckpoint(ctx, checkpoint.GetId(), "alice") + _, _, err = beginDeleteCheckpoint(ctx, client, checkpoint.GetId(), "alice") require.ErrorIs(t, err, ErrNotFound, "an in-flight checkpoint must not be deleted") - _, err = client.FinalizeSessionCheckpoint(ctx, checkpoint.GetId(), tc.tagUID, tc.snapshotURI, tc.failure) + _, err = finalizeCheckpoint(ctx, client, checkpoint.GetId(), tc.tagUID, tc.snapshotURI, tc.failure) require.NoError(t, err) require.NoError(t, deleteSession(ctx, client, session.GetId())) assertRetained() - _, _, err = client.BeginDeleteSessionCheckpoint(ctx, checkpoint.GetId(), "mallory") + _, _, err = beginDeleteCheckpoint(ctx, client, checkpoint.GetId(), "mallory") require.ErrorIs(t, err, ErrNotFound, "only the owner can delete the checkpoint") - _, _, err = client.BeginDeleteSessionCheckpoint(ctx, checkpoint.GetId(), "alice") + _, _, err = beginDeleteCheckpoint(ctx, client, checkpoint.GetId(), "alice") require.NoError(t, err) assertRetained() // A new store instance must resume a deletion without losing its references. restarted := NewClient(sharedDB, "public") - _, tagUID, err := restarted.BeginDeleteSessionCheckpoint(ctx, checkpoint.GetId(), "alice") + _, tagUID, err := beginDeleteCheckpoint(ctx, restarted, checkpoint.GetId(), "alice") require.NoError(t, err) require.Equal(t, tc.tagUID, tagUID) assertRetained() - require.NoError(t, client.DeleteSessionCheckpoint(ctx, checkpoint.GetId(), "alice")) + require.NoError(t, deleteCheckpoint(ctx, client, checkpoint.GetId(), "alice")) revisions, err := client.ListUnreferencedRuntimeRevisions(ctx) require.NoError(t, err) require.Len(t, revisions, 1) diff --git a/go/core/internal/database/session_expiration_test.go b/go/core/internal/database/session_expiration_test.go index 763ce2bda9..cafea4d20e 100644 --- a/go/core/internal/database/session_expiration_test.go +++ b/go/core/internal/database/session_expiration_test.go @@ -230,13 +230,13 @@ func TestSessionExpirationCheckpointAdmission(t *testing.T) { task := newSessionTask(uuid.NewString(), "message") task.ContextID, task.Status.State = session.ContextId, a2a.TaskStateCompleted require.NoError(t, saveRuntimeTask(t, client, session.Id, task, task, &SessionTaskSnapshot{Atespace: "team-a", URI: "s3://snapshot", Fidelity: "SNAPSHOT_FIDELITY_VOLUMES"})) - checkpoint, _, err := client.ReserveSessionCheckpoint(t.Context(), &apiv1.Checkpoint{ + checkpoint, _, err := reserveCheckpoint(t.Context(), client, &apiv1.Checkpoint{ Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID), }, session.Creator, "checkpoint") require.NoError(t, err) _, err = client.BeginIdleSessionDeletion(t.Context(), session.Id, time.Now()) require.ErrorIs(t, err, ErrConflict) - _, err = client.FinalizeSessionCheckpoint(t.Context(), checkpoint.Id, "tag-uid", "s3://snapshot", "") + _, err = finalizeCheckpoint(t.Context(), client, checkpoint.Id, "tag-uid", "s3://snapshot", "") require.NoError(t, err) beforeFork := time.Now() fork, _, err := client.ForkSession(t.Context(), checkpoint.Id, session.Creator, "fork", uuid.NewString()) diff --git a/go/core/internal/database/session_recovery_test.go b/go/core/internal/database/session_recovery_test.go index 0412f02c07..3270fb9221 100644 --- a/go/core/internal/database/session_recovery_test.go +++ b/go/core/internal/database/session_recovery_test.go @@ -80,7 +80,7 @@ func TestCrashedRecoverySettlesTasksAtomically(t *testing.T) { require.ErrorIs(t, err, ErrConflict) } require.ErrorIs(t, client.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "next"), ErrConflict) - _, _, err = client.ReserveSessionCheckpoint(t.Context(), &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID)}, session.Creator, uuid.NewString()) + _, _, err = reserveCheckpoint(t.Context(), client, &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID)}, session.Creator, uuid.NewString()) require.ErrorIs(t, err, ErrConflict) }) } @@ -108,7 +108,7 @@ func TestCrashedRecoveryAdmissionRespectsFences(t *testing.T) { require.NoError(t, saveRuntimeTask(t, client, session.Id, task, task, &SessionTaskSnapshot{ Atespace: "team-a", URI: "s3://snapshots/completed", Fidelity: "SNAPSHOT_FIDELITY_VOLUMES", })) - _, _, err := client.ReserveSessionCheckpoint(t.Context(), &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID)}, session.Creator, uuid.NewString()) + _, _, err := reserveCheckpoint(t.Context(), client, &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID)}, session.Creator, uuid.NewString()) require.NoError(t, err) case test.superseded: // Another revert finishes on the same Actor UID, and a new turn diff --git a/go/core/internal/database/task_events_test.go b/go/core/internal/database/task_events_test.go index fa382e6362..5499a0e961 100644 --- a/go/core/internal/database/task_events_test.go +++ b/go/core/internal/database/task_events_test.go @@ -79,7 +79,7 @@ func TestTaskViewsRebuildFromEvents(t *testing.T) { task.Status = a2a.TaskStatus{State: a2a.TaskStateInputRequired, Message: question} require.NoError(t, saveRuntimeTask(t, client, session.Id, task, task, nil)) assertReplay() - _, _, err = client.ReserveSessionCheckpoint(ctx, &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID)}, "alice", "hitl-checkpoint") + _, _, err = reserveCheckpoint(ctx, client, &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID)}, "alice", "hitl-checkpoint") require.ErrorIs(t, err, ErrFailedPrecondition) // Both a user reply and an immediate message result must persist their status. @@ -97,9 +97,9 @@ func TestTaskViewsRebuildFromEvents(t *testing.T) { require.NoError(t, saveRuntimeTask(t, client, session.Id, task, message, snapshot)) assertReplay() } - checkpoint, _, err := client.ReserveSessionCheckpoint(ctx, &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID)}, "alice", "checkpoint") + checkpoint, _, err := reserveCheckpoint(ctx, client, &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID)}, "alice", "checkpoint") require.NoError(t, err) - _, err = client.FinalizeSessionCheckpoint(ctx, checkpoint.Id, "tag", "retained", "") + _, err = finalizeCheckpoint(ctx, client, checkpoint.Id, "tag", "retained", "") require.NoError(t, err) boundaryEvents, err := readCheckpointEvents(ctx, q, uuid.MustParse(checkpoint.Id)) require.NoError(t, err) diff --git a/go/core/internal/database/task_store_test.go b/go/core/internal/database/task_store_test.go index 1ba79bae67..6fde139bca 100644 --- a/go/core/internal/database/task_store_test.go +++ b/go/core/internal/database/task_store_test.go @@ -229,15 +229,15 @@ func TestRuntimeForkRetainsOnlyTheCheckpointBoundary(t *testing.T) { hash := sha256.Sum256([]byte("checkpoint boundary")) version, err := client.UpdateSessionTask(t.Context(), source.Id, initialVersion, hash[:], completed, completed, "") require.NoError(t, err) - _, _, err = client.ReserveSessionCheckpoint(t.Context(), &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: source.Id, HeadTaskId: string(completed.ID)}, "alice", uuid.NewString()) + _, _, err = reserveCheckpoint(t.Context(), client, &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: source.Id, HeadTaskId: string(completed.ID)}, "alice", uuid.NewString()) require.ErrorIs(t, err, ErrFailedPrecondition) require.NoError(t, client.SettleSessionTask(t.Context(), source.Id, string(completed.ID), version)) boundary, err := client.ClaimSessionQuiescence(t.Context()) require.NoError(t, err) require.NoError(t, client.FinishSessionQuiescence(t.Context(), boundary, &SessionTaskSnapshot{Atespace: "team-a", URI: "turn-N", Fidelity: "SNAPSHOT_FIDELITY_VOLUMES"})) - checkpoint, _, err := client.ReserveSessionCheckpoint(t.Context(), &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: source.Id, HeadTaskId: string(completed.ID)}, "alice", uuid.NewString()) + checkpoint, _, err := reserveCheckpoint(t.Context(), client, &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: source.Id, HeadTaskId: string(completed.ID)}, "alice", uuid.NewString()) require.NoError(t, err) - _, err = client.FinalizeSessionCheckpoint(t.Context(), checkpoint.Id, "tag-N", "retained-N", "") + _, err = finalizeCheckpoint(t.Context(), client, checkpoint.Id, "tag-N", "retained-N", "") require.NoError(t, err) newInput := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("source after checkpoint")) newInput.ContextID = source.ContextId @@ -469,26 +469,26 @@ func TestCheckpointPinsExpectedTaskWhileSnapshotIsPending(t *testing.T) { require.NoError(t, err) require.NoError(t, client.SettleSessionTask(ctx, session.Id, string(task.ID), version)) request := &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID)} - _, _, err = client.ReserveSessionCheckpoint(ctx, request, "alice", "checkpoint") + _, _, err = reserveCheckpoint(ctx, client, request, "alice", "checkpoint") require.ErrorIs(t, err, ErrSnapshotPending) work, err := client.ClaimSessionQuiescence(ctx) require.NoError(t, err) require.NoError(t, client.FinishSessionQuiescence(ctx, work, &SessionTaskSnapshot{Atespace: "team-a", URI: "snapshot", Fidelity: "SNAPSHOT_FIDELITY_VOLUMES"})) - saved, _, err := client.ReserveSessionCheckpoint(ctx, request, "alice", "saved") + saved, _, err := reserveCheckpoint(ctx, client, request, "alice", "saved") require.NoError(t, err) - _, err = client.FinalizeSessionCheckpoint(ctx, saved.Id, "tag", "retained-snapshot", "") + _, err = finalizeCheckpoint(ctx, client, saved.Id, "tag", "retained-snapshot", "") require.NoError(t, err) next := newSessionTask(uuid.NewString(), "next-message") next.ContextID = session.ContextId _, err = client.CreateRuntimeTask(ctx, session.Id, taskMutationHash("next"), next, "") require.NoError(t, err) - _, _, err = client.ReserveSessionCheckpoint(ctx, request, "alice", "checkpoint") + _, _, err = reserveCheckpoint(ctx, client, request, "alice", "checkpoint") require.ErrorIs(t, err, ErrCheckpointAdvanced) - replayed, _, err := client.ReserveSessionCheckpoint(ctx, request, "alice", "saved") + replayed, _, err := reserveCheckpoint(ctx, client, request, "alice", "saved") require.NoError(t, err) require.Equal(t, saved.Id, replayed.Id) request.HeadTaskId = string(next.ID) - _, _, err = client.ReserveSessionCheckpoint(ctx, request, "alice", "saved") + _, _, err = reserveCheckpoint(ctx, client, request, "alice", "saved") require.ErrorIs(t, err, ErrIdempotencyConflict) } @@ -534,7 +534,7 @@ func TestCheckpointRejectsOlderPendingTask(t *testing.T) { require.NoError(t, saveRuntimeTask(t, client, session.Id, task, task, nil)) task.Status.State = a2a.TaskStateCompleted require.NoError(t, saveRuntimeTask(t, client, session.Id, task, task, &SessionTaskSnapshot{Atespace: "team-a", URI: "snapshot", Fidelity: "SNAPSHOT_FIDELITY_VOLUMES"})) - _, _, err := client.ReserveSessionCheckpoint(ctx, &apiv1.Checkpoint{ + _, _, err := reserveCheckpoint(ctx, client, &apiv1.Checkpoint{ Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID), }, "alice", uuid.NewString()) require.ErrorIs(t, err, ErrFailedPrecondition) diff --git a/go/core/internal/mcp/checkpoint_errors_test.go b/go/core/internal/mcp/checkpoint_errors_test.go new file mode 100644 index 0000000000..78fd0177f7 --- /dev/null +++ b/go/core/internal/mcp/checkpoint_errors_test.go @@ -0,0 +1,164 @@ +package mcp + +import ( + "context" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "testing" + + "github.com/a2aproject/a2a-go/v2/a2asrv" + "github.com/google/uuid" + "github.com/stretchr/testify/require" + ateapipb "go.ate.dev/substrate/api/ateapi/v1" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" + "github.com/kagent-dev/kagent/go/core/internal/database" + "github.com/kagent-dev/kagent/go/core/internal/service/checkpoint" + "github.com/kagent-dev/kagent/go/core/internal/substrate" + "github.com/kagent-dev/kagent/go/core/pkg/auth" +) + +func TestCheckpointToolsHideStorageFailures(t *testing.T) { + for _, phase := range []string{"reserve failure", "claim failure", "list failure", "fork failure", "delete claim failure"} { + t.Run("Given "+phase+", when the checkpoint tool is called, assert that internal details stay private", func(t *testing.T) { + store := &checkpointErrorStore{phase: phase} + tool := createCheckpointToolName + args := map[string]any{"session_id": testSessionID, "expected_head_task_id": testTaskID, "request_id": "checkpoint-request"} + if phase == "list failure" { + tool = listCheckpointsToolName + args = map[string]any{"session_id": testSessionID} + } + if phase == "fork failure" { + tool = forkSessionToolName + args = map[string]any{"checkpoint_id": "33333333-3333-4333-8333-333333333333", "request_id": "checkpoint-fork"} + } + if phase == "delete claim failure" { + tool = deleteCheckpointToolName + args = map[string]any{"checkpoint_id": "33333333-3333-4333-8333-333333333333"} + } + result := checkpointToolResponse(t, checkpoint.NewService(store, &auth.NoopAuthorizer{}, nil, nil), tool, args) + text := result["content"].([]any)[0].(map[string]any)["text"].(string) + require.Equal(t, true, result["isError"]) + require.Equal(t, 1, store.failures, "the intended failure must be reached") + require.NotContains(t, text, "private-db.internal", "MCP must not expose the storage cause") + }) + } +} + +func TestCheckpointToolsRetainCurrentBackendStatus(t *testing.T) { + for _, code := range []codes.Code{codes.PermissionDenied, codes.FailedPrecondition, codes.InvalidArgument, codes.Unauthenticated, codes.Unavailable, codes.Aborted} { + t.Run("Given backend "+code.String()+", when Delete is called, assert that the current code survives without private details", func(t *testing.T) { + backend := &checkpointErrorBackend{err: status.Error(code, "private backend endpoint")} + result := checkpointToolResponse(t, checkpoint.NewService(&checkpointErrorStore{phase: "backend"}, &auth.NoopAuthorizer{}, backend, nil), deleteCheckpointToolName, + map[string]any{"checkpoint_id": "33333333-3333-4333-8333-333333333333"}) + text := result["content"].([]any)[0].(map[string]any)["text"].(string) + require.Equal(t, true, result["isError"]) + require.Equal(t, 1, backend.calls) + require.Contains(t, text, code.String()) + require.NotContains(t, text, "private") + }) + } +} + +func TestCheckpointToolsPreserveSnapshotRetryReasons(t *testing.T) { + for reason, failure := range map[string]error{ + "KAGENT_CHECKPOINT_SNAPSHOT_PENDING": database.ErrSnapshotPending, + "KAGENT_CHECKPOINT_CONVERSATION_ADVANCED": database.ErrCheckpointAdvanced, + } { + t.Run("Given "+reason+", when Create is called, assert that existing retry metadata is retained", func(t *testing.T) { + store := &checkpointErrorStore{reserveErr: failure} + result := checkpointToolResponse(t, checkpoint.NewService(store, &auth.NoopAuthorizer{}, nil, nil), createCheckpointToolName, + map[string]any{"session_id": testSessionID, "expected_head_task_id": testTaskID, "request_id": "retry-reason"}) + require.Equal(t, true, result["isError"]) + require.Equal(t, reason, result["_meta"].(map[string]any)["kagent.dev/error-reason"]) + }) + } +} + +func TestCheckpointToolsSanitizeWrappedCancellation(t *testing.T) { + for _, cause := range []error{context.Canceled, context.DeadlineExceeded} { + store := &checkpointErrorStore{phase: "claim failure", claimErr: fmt.Errorf("private database endpoint: %w", cause)} + result := checkpointToolResponse(t, checkpoint.NewService(store, &auth.NoopAuthorizer{}, nil, nil), createCheckpointToolName, + map[string]any{"session_id": testSessionID, "expected_head_task_id": testTaskID, "request_id": "canceled-request"}) + text := result["content"].([]any)[0].(map[string]any)["text"].(string) + require.Equal(t, true, result["isError"]) + require.NotEmpty(t, text) + require.NotContains(t, text, "private") + } +} + +func checkpointToolResponse(t *testing.T, service *checkpoint.Service, tool string, args map[string]any) map[string]any { + t.Helper() + h, err := New(Config{Agents: testAgentService(t), Sessions: testSessionService(), Checkpoints: service, Gateway: &a2asrv.InterceptedHandler{Handler: &fakeGateway{}}}) + require.NoError(t, err) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { h.ServeHTTP(w, r.WithContext(authContext())) })) + t.Cleanup(server.Close) + response := rawMCPCall(t, server.URL, "tools/call", map[string]any{"name": tool, "arguments": args}, false) + return response["result"].(map[string]any) +} + +type checkpointErrorStore struct { + *database.Client + phase string + failures int + claimErr error + reserveErr error +} + +func (s *checkpointErrorStore) ReserveSessionCheckpoint(_ context.Context, cp *apiv1.Checkpoint, _ string, _ string, _ uuid.UUID) (*apiv1.Checkpoint, *database.SessionTaskSnapshot, error) { + if s.reserveErr != nil { + return nil, nil, s.reserveErr + } + if s.phase == "reserve failure" { + s.failures++ + return nil, nil, errors.New("connect to private-db.internal user=controller database=kagent: connection refused") + } + cp.Id = "33333333-3333-4333-8333-333333333333" + cp.State = apiv1.CheckpointState_CHECKPOINT_STATE_READY + if s.phase == "claim failure" { + cp.State = apiv1.CheckpointState_CHECKPOINT_STATE_CREATING + } + return cp, &database.SessionTaskSnapshot{}, nil +} +func (s *checkpointErrorStore) ClaimCheckpointOperation(context.Context, string, string, uuid.UUID) (*database.CheckpointOperation, error) { + s.failures++ + if s.claimErr != nil { + return nil, s.claimErr + } + return nil, errors.New("claim checkpoint operation: private-db.internal user=controller database=kagent: connection refused") +} +func (s *checkpointErrorStore) ReleaseCheckpointOperation(context.Context, string, uuid.UUID) error { + return nil +} +func (s *checkpointErrorStore) ReleaseCheckpointReservation(context.Context, string, string, uuid.UUID) error { + return nil +} +func (s *checkpointErrorStore) ListSessionCheckpoints(context.Context, string, string, string, int) ([]*apiv1.Checkpoint, error) { + s.failures++ + return nil, errors.New("query private-db.internal user=controller database=kagent: connection refused") +} +func (s *checkpointErrorStore) GetSessionCheckpointForkSource(context.Context, string, string) (*database.SessionCheckpointForkSource, error) { + s.failures++ + return nil, errors.New("read fork source from private-db.internal user=controller database=kagent: connection refused") +} +func (s *checkpointErrorStore) BeginDeleteSessionCheckpoint(context.Context, string, string, uuid.UUID) (*database.CheckpointOperation, error) { + if s.phase == "backend" { + return &database.CheckpointOperation{Checkpoint: &apiv1.Checkpoint{Id: "33333333-3333-4333-8333-333333333333", State: apiv1.CheckpointState_CHECKPOINT_STATE_DELETING}, Creator: "alice", Snapshot: &database.SessionTaskSnapshot{Atespace: "team-a"}}, nil + } + return nil, database.ErrConflict +} + +type checkpointErrorBackend struct { + *substrate.Client + err error + calls int +} + +func (b *checkpointErrorBackend) GetTag(context.Context, string, string) (*ateapipb.Tag, error) { + b.calls++ + return nil, b.err +} diff --git a/go/core/internal/mcp/checkpoints.go b/go/core/internal/mcp/checkpoints.go index 874e3378a2..80f7704f84 100644 --- a/go/core/internal/mcp/checkpoints.go +++ b/go/core/internal/mcp/checkpoints.go @@ -2,6 +2,7 @@ package mcp import ( "context" + "errors" "fmt" "time" @@ -87,7 +88,7 @@ func (h *Handler) registerCheckpointTools(server *mcp.Server) { func (h *Handler) getCheckpoint(ctx context.Context, _ *mcp.CallToolRequest, input CheckpointIDInput) (*mcp.CallToolResult, CreateCheckpointOutput, error) { value, err := h.checkpoints.Get(ctx, input.CheckpointID) if err != nil { - return sessionToolError(err), CreateCheckpointOutput{}, nil + return checkpointToolError(err), CreateCheckpointOutput{}, nil } output := CreateCheckpointOutput{Checkpoint: checkpointSummary(value)} return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: fmt.Sprintf("Checkpoint %s is %s", value.GetId(), value.GetState())}}}, output, nil @@ -95,7 +96,7 @@ func (h *Handler) getCheckpoint(ctx context.Context, _ *mcp.CallToolRequest, inp func (h *Handler) deleteCheckpoint(ctx context.Context, _ *mcp.CallToolRequest, input CheckpointIDInput) (*mcp.CallToolResult, DeleteCheckpointOutput, error) { if err := h.checkpoints.Delete(ctx, input.CheckpointID); err != nil { - return sessionToolError(err), DeleteCheckpointOutput{}, nil + return checkpointToolError(err), DeleteCheckpointOutput{}, nil } return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: fmt.Sprintf("Deleted checkpoint %s", input.CheckpointID)}}}, DeleteCheckpointOutput(input), nil } @@ -103,7 +104,7 @@ func (h *Handler) deleteCheckpoint(ctx context.Context, _ *mcp.CallToolRequest, func (h *Handler) createCheckpoint(ctx context.Context, _ *mcp.CallToolRequest, input CreateCheckpointInput) (*mcp.CallToolResult, CreateCheckpointOutput, error) { created, err := h.checkpoints.Create(ctx, input.SessionID, stableRequestID(input.RequestID), input.ExpectedHeadTaskID) if err != nil { - result := toolError(err) + result := checkpointToolError(err) for _, detail := range status.Convert(err).Details() { if info, ok := detail.(*errdetails.ErrorInfo); ok && info.Domain == "kagent.dev" { result.Meta = mcp.Meta{"kagent.dev/error-reason": info.Reason} @@ -121,7 +122,7 @@ func (h *Handler) listCheckpoints(ctx context.Context, _ *mcp.CallToolRequest, i PageSize: input.PageSize, PageToken: input.PageToken, }) if err != nil { - return toolError(err), ListCheckpointsOutput{}, nil + return checkpointToolError(err), ListCheckpointsOutput{}, nil } output := ListCheckpointsOutput{Checkpoints: make([]CheckpointSummary, len(listed.Checkpoints)), NextPageToken: listed.NextPageToken} for i, item := range listed.Checkpoints { @@ -133,7 +134,7 @@ func (h *Handler) listCheckpoints(ctx context.Context, _ *mcp.CallToolRequest, i func (h *Handler) forkSession(ctx context.Context, _ *mcp.CallToolRequest, input ForkSessionInput) (*mcp.CallToolResult, ForkSessionOutput, error) { session, err := h.checkpoints.Fork(ctx, input.CheckpointID, stableRequestID(input.RequestID)) if err != nil { - return toolError(err), ForkSessionOutput{}, nil + return checkpointToolError(err), ForkSessionOutput{}, nil } output := ForkSessionOutput{Session: sessionSummary(session)} return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: fmt.Sprintf("Created Session %s", session.GetId())}}}, output, nil @@ -146,6 +147,17 @@ func stableRequestID(id string) string { return uuid.NewString() } +// MCP bypasses the gRPC error interceptor. Use classified public messages and +// strip private causes from wrapped cancellation errors at this boundary too. +func checkpointToolError(err error) *mcp.CallToolResult { + if errors.Is(err, context.Canceled) { + err = errors.New("checkpoint operation canceled") + } else if errors.Is(err, context.DeadlineExceeded) { + err = errors.New("checkpoint operation deadline exceeded") + } + return sessionToolError(err) +} + func checkpointSummary(value *apiv1.Checkpoint) CheckpointSummary { result := CheckpointSummary{ ID: value.GetId(), SessionID: value.GetSessionId(), Name: value.GetName(), diff --git a/go/core/internal/service/checkpoint/completion_test.go b/go/core/internal/service/checkpoint/completion_test.go new file mode 100644 index 0000000000..07b2cdc982 --- /dev/null +++ b/go/core/internal/service/checkpoint/completion_test.go @@ -0,0 +1,217 @@ +package checkpoint + +import ( + "context" + "errors" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" + "github.com/kagent-dev/kagent/go/core/internal/database" + "github.com/kagent-dev/kagent/go/core/internal/service/serviceerrors" + "github.com/kagent-dev/kagent/go/core/pkg/auth" +) + +func TestCreateObservesReplacementCompletion(t *testing.T) { + persistence := recoveryStore(t) + for _, outcome := range []string{"pending", "ready", "failed", "deleted"} { + t.Run("Given replacement ownership with outcome "+outcome+", when the original Create finishes, assert that it observes the current result", func(t *testing.T) { + owner := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) + source, head := recoverySource(t, persistence) + tags := &testTags{snapshotURI: "s3://snapshots/source"} + var replacement uuid.UUID + var checkpointID string + gate := &checkpointCompletionGate{Client: persistence} + gate.beforeCreateFinish = func(ctx context.Context, checkpoint *apiv1.Checkpoint, executor uuid.UUID, uid, uri string) { + checkpointID = checkpoint.Id + require.NoError(t, persistence.ReleaseCheckpointOperation(ctx, checkpoint.Id, executor)) + replacement = uuid.New() + op, err := persistence.ClaimCheckpointOperation(ctx, checkpoint.Id, "alice", replacement) + require.NoError(t, err) + if outcome == "pending" { + return + } + failure := "" + if outcome == "failed" { + failure = "replacement rejected the copy" + } + _, err = persistence.FinalizeClaimedSessionCheckpoint(ctx, op.Checkpoint, replacement, uid, uri, failure) + require.NoError(t, err) + if outcome == "deleted" { + deletion := uuid.New() + _, err := persistence.BeginDeleteSessionCheckpoint(ctx, checkpoint.Id, "alice", deletion) + require.NoError(t, err) + require.NoError(t, tags.DeleteTag(ctx, "team-a", tagName(checkpoint.Id), uid)) + require.NoError(t, persistence.DeleteClaimedSessionCheckpoint(ctx, checkpoint.Id, "alice", deletion)) + } + } + releaseFailure := &checkpointReleaseFailure{store: gate} + service := NewService(releaseFailure, testAuthorizer{}, tags, nil) + request := uuid.NewString() + checkpoint, err := service.Create(owner, source.Id, request, head) + require.Equal(t, 1, gate.conflicts, "the real store must reject the original executor") + require.Equal(t, 1, tags.createCalls, "observing completion must not repeat backend work") + switch outcome { + case "ready": + require.NoError(t, err) + require.Equal(t, checkpointID, checkpoint.Id) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_READY, checkpoint.State) + case "pending": + require.True(t, serviceerrors.IsCode(err, serviceerrors.CodeAborted), "%v", err) + require.Nil(t, checkpoint) + ctx, cancel := context.WithTimeout(owner, database.RuntimeOperationTimeout) + defer cancel() + _, err := persistence.ClaimCheckpointOperation(ctx, checkpointID, "alice", uuid.New()) + require.ErrorIs(t, err, database.ErrConflict, "the stale release must preserve replacement ownership") + require.NoError(t, persistence.ReleaseCheckpointOperation(ctx, checkpointID, replacement)) + _, err = NewService(persistence, testAuthorizer{}, tags, nil).Create(owner, source.Id, request, head) + require.NoError(t, err) + case "failed": + assertFailedCheckpointOutcome(t, err) + op, readErr := persistence.GetCheckpointOperation(owner, checkpointID, "alice") + require.NoError(t, readErr) + require.Equal(t, "replacement rejected the copy", op.Checkpoint.Failure.Message) + case "deleted": + require.True(t, serviceerrors.IsCode(err, serviceerrors.CodeNotFound), "%v", err) + require.Nil(t, checkpoint, "Create must not return an empty successful response") + } + if outcome != "deleted" { + require.NoError(t, NewService(persistence, testAuthorizer{}, tags, nil).Delete(owner, checkpointID)) + } + require.Zero(t, releaseFailure.calls, "an observer must not issue a stale release") + }) + } +} + +func TestClaimConflictObservesCompletionWithoutRelease(t *testing.T) { + persistence := recoveryStore(t) + owner := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) + source, head := recoverySource(t, persistence) + tags := &testTags{snapshotURI: "s3://snapshots/source"} + service := NewService(persistence, testAuthorizer{}, tags, nil) + request := uuid.NewString() + pending, _, err := service.beginCreate(owner, source.Id, request, head) + require.NoError(t, err) + gate := &checkpointCompletionGate{Client: persistence, beforeClaim: func(ctx context.Context) { + // A replacement finishes between the retry's reservation read and its + // conditional claim. The original retry never acquires execution rights. + _, err := service.Create(ctx, source.Id, request, head) + require.NoError(t, err) + }} + releaseFailure := &checkpointReleaseFailure{store: gate} + checkpoint, err := NewService(releaseFailure, testAuthorizer{}, tags, nil).Create(owner, source.Id, request, head) + require.NoError(t, err, "a failed stale release must not replace observed success") + require.Equal(t, pending.Id, checkpoint.Id) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_READY, checkpoint.State) + require.Equal(t, 1, gate.conflicts, "the real store must reject the original claim") + require.Equal(t, 1, tags.createCalls, "only the replacement may copy") + require.Zero(t, releaseFailure.calls) + require.NoError(t, service.Delete(owner, checkpoint.Id)) +} + +func TestDeleteObservesReplacementCompletion(t *testing.T) { + persistence := recoveryStore(t) + for _, outcome := range []string{"pending", "deleted"} { + t.Run("Given replacement ownership with outcome "+outcome+", when the original Delete finishes, assert that it observes the current result", func(t *testing.T) { + owner := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) + source, head := recoverySource(t, persistence) + tags := &testTags{snapshotURI: "s3://snapshots/source"} + service := NewService(persistence, testAuthorizer{}, tags, nil) + checkpoint, err := service.Create(owner, source.Id, uuid.NewString(), head) + require.NoError(t, err) + var replacement uuid.UUID + gate := &checkpointCompletionGate{Client: persistence} + gate.beforeDeleteFinish = func(ctx context.Context, id, owner string, executor uuid.UUID) { + require.NoError(t, persistence.ReleaseCheckpointOperation(ctx, id, executor)) + replacement = uuid.New() + _, err := persistence.ClaimCheckpointOperation(ctx, id, owner, replacement) + require.NoError(t, err) + if outcome == "deleted" { + require.NoError(t, persistence.DeleteClaimedSessionCheckpoint(ctx, id, owner, replacement)) + } + } + releaseFailure := &checkpointReleaseFailure{store: gate} + err = NewService(releaseFailure, testAuthorizer{}, tags, nil).Delete(owner, checkpoint.Id) + require.Equal(t, 1, gate.conflicts) + require.Equal(t, 1, tags.deleteCalls, "observing completion must not repeat backend cleanup") + if outcome == "deleted" { + require.NoError(t, err, "a replacement already completed deletion") + require.Zero(t, releaseFailure.calls, "confirmed deletion needs no further release") + _, err = persistence.GetCheckpointOperation(owner, checkpoint.Id, "alice") + require.ErrorIs(t, err, database.ErrNotFound) + } else { + require.True(t, serviceerrors.IsCode(err, serviceerrors.CodeAborted), "%v", err) + require.NoError(t, persistence.ReleaseCheckpointOperation(owner, checkpoint.Id, replacement)) + require.NoError(t, service.Delete(owner, checkpoint.Id)) + } + }) + } +} + +// Replace ownership after backend work, just before the original conditional +// completion. Store methods perform every transition and reject the stale write. +type checkpointCompletionGate struct { + *database.Client + beforeClaim func(context.Context) + beforeCreateFinish func(context.Context, *apiv1.Checkpoint, uuid.UUID, string, string) + beforeDeleteFinish func(context.Context, string, string, uuid.UUID) + beforeFailedClear func(context.Context, string, string, string, uuid.UUID) + conflicts int +} + +func (s *checkpointCompletionGate) ClaimCheckpointOperation(ctx context.Context, id, owner string, executor uuid.UUID) (*database.CheckpointOperation, error) { + if s.beforeClaim != nil { + s.beforeClaim(ctx) + } + op, err := s.Client.ClaimCheckpointOperation(ctx, id, owner, executor) + if errors.Is(err, database.ErrConflict) { + s.conflicts++ + } + return op, err +} + +func (s *checkpointCompletionGate) FinalizeClaimedSessionCheckpoint(ctx context.Context, checkpoint *apiv1.Checkpoint, executor uuid.UUID, uid, uri, failure string) (*apiv1.Checkpoint, error) { + s.beforeCreateFinish(ctx, checkpoint, executor, uid, uri) + result, err := s.Client.FinalizeClaimedSessionCheckpoint(ctx, checkpoint, executor, uid, uri, failure) + if errors.Is(err, database.ErrConflict) { + s.conflicts++ + } + return result, err +} + +func (s *checkpointCompletionGate) DeleteClaimedSessionCheckpoint(ctx context.Context, id, owner string, executor uuid.UUID) error { + s.beforeDeleteFinish(ctx, id, owner, executor) + err := s.Client.DeleteClaimedSessionCheckpoint(ctx, id, owner, executor) + if errors.Is(err, database.ErrConflict) { + s.conflicts++ + } + return err +} + +func (s *checkpointCompletionGate) ClearFailedCheckpointTag(ctx context.Context, id, owner, uid string, executor uuid.UUID) error { + s.beforeFailedClear(ctx, id, owner, uid, executor) + err := s.Client.ClearFailedCheckpointTag(ctx, id, owner, uid, executor) + if errors.Is(err, database.ErrConflict) { + s.conflicts++ + } + return err +} + +// Any unnecessary release fails, making it observable if cleanup masks a result +// already established by the real store. Ownership transfers use the Client. +type checkpointReleaseFailure struct { + store + calls int +} + +func (s *checkpointReleaseFailure) ReleaseCheckpointOperation(context.Context, string, uuid.UUID) error { + s.calls++ + return errors.New("private database release failure") +} + +func (s *checkpointReleaseFailure) ReleaseCheckpointReservation(context.Context, string, string, uuid.UUID) error { + s.calls++ + return errors.New("private database reservation release failure") +} diff --git a/go/core/internal/service/checkpoint/failed_cleanup_test.go b/go/core/internal/service/checkpoint/failed_cleanup_test.go new file mode 100644 index 0000000000..6b3d03b31a --- /dev/null +++ b/go/core/internal/service/checkpoint/failed_cleanup_test.go @@ -0,0 +1,241 @@ +package checkpoint + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + ateapipb "go.ate.dev/substrate/api/ateapi/v1" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" + "github.com/kagent-dev/kagent/go/core/internal/database" + "github.com/kagent-dev/kagent/go/core/internal/service/serviceerrors" + "github.com/kagent-dev/kagent/go/core/pkg/auth" +) + +func TestFailedTagsAreCollectedWithoutExposingFailedCheckpoints(t *testing.T) { + persistence := recoveryStore(t) + owner := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) + source, head := recoverySource(t, persistence) + tags := &testTags{snapshotURI: "s3://snapshots/source", createErr: status.Error(codes.Unavailable, "copy interrupted"), + mutateTag: func(tag *ateapipb.Tag) { + tag.Status.Snapshot.DurableSnapshot.State = ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_PENDING + }} + service := NewService(persistence, testAuthorizer{}, tags, nil) + request := uuid.NewString() + _, err := service.Create(owner, source.Id, request, head) + require.Equal(t, codes.Unavailable, status.Code(err)) + _, err = service.Create(owner, source.Id, request, head) + assertFailedCheckpointOutcome(t, err) + pending, err := persistence.ListPendingCheckpoints(owner, "", 100) + require.NoError(t, err) + require.Len(t, pending, 1, "a new FAILED tag is immediately eligible for collection") + id := pending[0].Checkpoint.Id + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_FAILED, pending[0].Checkpoint.State) + _, err = service.Get(owner, id) + require.True(t, serviceerrors.IsCode(err, serviceerrors.CodeNotFound)) + listed, err := service.List(owner, ListRequest{SessionID: source.Id}) + require.NoError(t, err) + require.Empty(t, listed.Checkpoints) + _, err = service.Fork(owner, id, uuid.NewString()) + require.True(t, serviceerrors.IsCode(err, serviceerrors.CodeNotFound)) + _, _, err = persistence.ForkSession(owner, id, "alice", uuid.NewString(), uuid.NewString()) + require.ErrorIs(t, err, database.ErrNotFound) + tags.deleteErr = status.Error(codes.PermissionDenied, "storage denied") + require.NoError(t, NewRecoveryWorker(service).sweep(owner)) + retained, err := persistence.GetCheckpointOperation(owner, id, "alice") + require.NoError(t, err) + require.Equal(t, "tag-uid", retained.TagUID) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_FAILED, retained.Checkpoint.State) + tags.deleteErr = nil + lost := &lostFailedCollection{Client: persistence} + require.NoError(t, NewRecoveryWorker(NewService(lost, testAuthorizer{}, tags, nil)).sweep(owner)) + require.Equal(t, 1, lost.calls, "the UID-clear reply is lost after the real write commits") + collected, err := persistence.GetCheckpointOperation(owner, id, "alice") + require.NoError(t, err) + require.Empty(t, collected.TagUID) + require.Equal(t, retained.Checkpoint, collected.Checkpoint, "collection preserves the terminal request result") + require.Nil(t, tags.created) + _, err = service.Create(owner, source.Id, request, head) + assertFailedCheckpointOutcome(t, err) + require.Equal(t, 2, tags.createCalls, "terminal retries must not copy again") + pending, err = persistence.ListPendingCheckpoints(owner, "", 100) + require.NoError(t, err) + require.Empty(t, pending, "collected FAILED rows must not be retried") + require.NoError(t, NewRecoveryWorker(service).sweep(owner)) + require.Equal(t, 2, tags.createCalls) +} + +type lostFailedCollection struct { + *database.Client + calls int +} + +func (s *lostFailedCollection) ClearFailedCheckpointTag(ctx context.Context, id, owner, uid string, executor uuid.UUID) error { + if err := s.Client.ClearFailedCheckpointTag(ctx, id, owner, uid, executor); err != nil { + return err + } + s.calls++ + return errors.New("cleanup reply lost") +} + +func TestFreshInvalidSnapshotDoesNotStartCopy(t *testing.T) { + persistence := recoveryStore(t) + owner := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) + source, head := recoverySource(t, persistence) + tags := &testTags{snapshotURI: "s3://snapshots/changed"} + service := NewService(persistence, testAuthorizer{}, tags, nil) + _, err := service.Create(owner, source.Id, uuid.NewString(), head) + assertFailedCheckpointOutcome(t, err) + require.Zero(t, tags.createCalls) + pending, err := persistence.ListPendingCheckpoints(owner, "", 100) + require.NoError(t, err) + require.Empty(t, pending, "there is no tag to collect") + _, err = persistence.BeginSessionOperation(owner, source.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE, false) + require.NoError(t, err, "failed verification releases source admission") +} + +func TestRecoveryRetainsPriorTagWhenSourceChanged(t *testing.T) { + persistence := recoveryStore(t) + owner := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) + source, head := recoverySource(t, persistence) + tags := &testTags{snapshotURI: "s3://snapshots/source", createErr: status.Error(codes.Unavailable, "copy reply lost"), + mutateTag: func(tag *ateapipb.Tag) { + tag.Status.Snapshot.DurableSnapshot.State = ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_PENDING + }} + service := NewService(persistence, testAuthorizer{}, tags, nil) + request := uuid.NewString() + _, err := service.Create(owner, source.Id, request, head) + require.Equal(t, codes.Unavailable, status.Code(err)) + tags.snapshotURI = "s3://snapshots/changed" + _, err = service.Create(owner, source.Id, request, head) + assertFailedCheckpointOutcome(t, err) + require.Equal(t, 2, tags.createCalls, "recovery must serialize with the prior copy") + pending, err := persistence.ListPendingCheckpoints(owner, "", 100) + require.NoError(t, err) + require.Len(t, pending, 1) + require.Equal(t, "tag-uid", pending[0].TagUID) + tags.created.Metadata.Uid = "replacement-uid" + require.NoError(t, NewRecoveryWorker(service).sweep(owner)) + require.Zero(t, tags.deleteCalls, "collection must never delete a replacement tag") + retained, err := persistence.GetCheckpointOperation(owner, pending[0].Checkpoint.Id, "alice") + require.NoError(t, err) + require.Equal(t, "tag-uid", retained.TagUID) +} + +func TestFailedCollectionSurvivesSourceDeletion(t *testing.T) { + f := newFailedCollectionFixture(t) + source := f.operation.Checkpoint.SessionId + work, err := f.persistence.BeginSessionOperation(f.owner, source, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE, false) + require.NoError(t, err) + ctx, cancel := context.WithTimeout(f.owner, time.Minute) + defer cancel() + executor := uuid.New() + acquired, err := f.persistence.ClaimSessionOperation(ctx, source, work.ID, executor) + require.NoError(t, err) + require.True(t, acquired) + deleted, err := f.persistence.FinishSessionOperation(ctx, source, work.ID, executor, "", "", "") + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_DELETED, deleted.State) + f.tags.actorErr = status.Error(codes.NotFound, "source no longer exists") + require.NoError(t, NewRecoveryWorker(f.service).sweep(f.owner)) + collected, err := f.persistence.GetCheckpointOperation(f.owner, f.operation.Checkpoint.Id, "alice") + require.NoError(t, err) + require.Empty(t, collected.TagUID) + require.Equal(t, f.operation.Checkpoint, collected.Checkpoint) + _, err = f.service.Create(f.owner, source, f.request, f.operation.Checkpoint.HeadTaskId) + assertFailedCheckpointOutcome(t, err) + require.Equal(t, 1, f.tags.createCalls, "replay after source deletion must not copy again") +} + +func TestFailedCollectionObservesReplacementCompletion(t *testing.T) { + for _, outcome := range []string{"collected", "deleted"} { + t.Run("Given replacement "+outcome+", when stale cleanup finishes, assert that completed work is observed without an obsolete release", func(t *testing.T) { + f := newFailedCollectionFixture(t) + gate := &checkpointCompletionGate{Client: f.persistence} + gate.beforeFailedClear = func(ctx context.Context, id, owner, uid string, executor uuid.UUID) { + require.NoError(t, f.persistence.ReleaseCheckpointOperation(ctx, id, executor)) + replacement := uuid.New() + if outcome == "deleted" { + _, err := f.persistence.BeginDeleteSessionCheckpoint(ctx, id, owner, replacement) + require.NoError(t, err) + require.NoError(t, f.persistence.DeleteClaimedSessionCheckpoint(ctx, id, owner, replacement)) + return + } + _, err := f.persistence.ClaimCheckpointOperation(ctx, id, owner, replacement) + require.NoError(t, err) + require.NoError(t, f.persistence.ClearFailedCheckpointTag(ctx, id, owner, uid, replacement)) + } + releases := &checkpointReleaseFailure{store: gate} + service := NewService(releases, testAuthorizer{}, f.tags, nil) + require.NoError(t, service.recoverOperation(f.owner, f.operation), "an obsolete release failure must not mask completed cleanup") + require.Equal(t, 1, gate.conflicts, "the real store must reject the original executor") + require.Zero(t, releases.calls) + require.Equal(t, 1, f.tags.createCalls) + }) + } +} + +func TestFailedCollectionPreservesReplacementClaim(t *testing.T) { + f := newFailedCollectionFixture(t) + replacement := uuid.New() + gate := &checkpointCompletionGate{Client: f.persistence} + gate.beforeFailedClear = func(ctx context.Context, id, owner, _ string, executor uuid.UUID) { + require.NoError(t, f.persistence.ReleaseCheckpointOperation(ctx, id, executor)) + _, err := f.persistence.ClaimCheckpointOperation(ctx, id, owner, replacement) + require.NoError(t, err) + } + service := NewService(gate, testAuthorizer{}, f.tags, nil) + err := service.recoverOperation(f.owner, f.operation) + require.True(t, serviceerrors.IsCode(err, serviceerrors.CodeAborted), "%v", err) + require.Equal(t, 1, gate.conflicts) + ctx, cancel := context.WithTimeout(f.owner, time.Minute) + defer cancel() + _, err = f.persistence.ClaimCheckpointOperation(ctx, f.operation.Checkpoint.Id, "alice", uuid.New()) + require.ErrorIs(t, err, database.ErrConflict, "the stale executor's real release must preserve replacement ownership") + require.NoError(t, f.persistence.ReleaseCheckpointOperation(ctx, f.operation.Checkpoint.Id, replacement)) + require.NoError(t, NewRecoveryWorker(f.service).sweep(f.owner)) + collected, err := f.persistence.GetCheckpointOperation(f.owner, f.operation.Checkpoint.Id, "alice") + require.NoError(t, err) + require.Empty(t, collected.TagUID, "a subsequent sweep must finish after the replacement retires") + require.Equal(t, 1, f.tags.createCalls) +} + +type failedCollectionFixture struct { + persistence *database.Client + owner context.Context + service *Service + tags *testTags + operation *database.CheckpointOperation + request string +} + +func newFailedCollectionFixture(t *testing.T) *failedCollectionFixture { + t.Helper() + persistence := recoveryStore(t) + owner := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) + source, head := recoverySource(t, persistence) + tags := &testTags{snapshotURI: "s3://snapshots/source", mutateTag: func(tag *ateapipb.Tag) { + tag.Status.Snapshot.DurableSnapshot.State = ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_PENDING + }} + service := NewService(persistence, testAuthorizer{}, tags, nil) + request := uuid.NewString() + _, err := service.Create(owner, source.Id, request, head) + assertFailedCheckpointOutcome(t, err) + pending, err := persistence.ListPendingCheckpoints(owner, "", 100) + require.NoError(t, err) + require.Len(t, pending, 1) + return &failedCollectionFixture{persistence: persistence, owner: owner, service: service, + tags: tags, operation: pending[0], request: request} +} + +func assertFailedCheckpointOutcome(t *testing.T, err error) { + t.Helper() + require.Equal(t, codes.FailedPrecondition, status.Code(err)) + require.Empty(t, status.Convert(err).Details(), "internal failure must not add public discovery identity") +} diff --git a/go/core/internal/service/checkpoint/migration_test.go b/go/core/internal/service/checkpoint/migration_test.go new file mode 100644 index 0000000000..1f86a66e4d --- /dev/null +++ b/go/core/internal/service/checkpoint/migration_test.go @@ -0,0 +1,101 @@ +package checkpoint + +import ( + "context" + "encoding/binary" + "testing" + "time" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/pressly/goose/v3" + "github.com/stretchr/testify/require" + "google.golang.org/protobuf/proto" + "google.golang.org/protobuf/types/known/timestamppb" + + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" + "github.com/kagent-dev/kagent/go/core/internal/database" + "github.com/kagent-dev/kagent/go/core/internal/dbtest" + "github.com/kagent-dev/kagent/go/core/pkg/auth" + "github.com/kagent-dev/kagent/go/core/pkg/migrations" +) + +// This is a populated legacy-schema fixture, not a previous-release binary. +// The normal API E2E covers controller restart and request replay separately. +func TestCheckpointRecoveryAcrossClaimMigration(t *testing.T) { + dsn := dbtest.StartT(context.WithoutCancel(t.Context()), t) + config := dbtest.PoolConfig(dsn) + sources := migrations.BuiltinSourcesInSchema(false, "public", "public") + require.NoError(t, migrations.WithProvider(t.Context(), config, sources[0], func(provider *goose.Provider) error { + _, err := provider.UpTo(t.Context(), 3) + return err + })) + pool, err := database.Connect(t.Context(), &database.PostgresConfig{Pool: config, Schema: "public"}) + require.NoError(t, err) + t.Cleanup(pool.Close) + persistence := database.NewClient(pool, "public") + source, head := recoverySource(t, persistence) + var historyID uuid.UUID + var sequence int64 + require.NoError(t, pool.QueryRow(t.Context(), ` + SELECT s.history_id, t.history_sequence FROM session s + JOIN session_task t ON t.history_id = s.history_id + WHERE s.id = $1 AND t.id = $2`, source.Id, head).Scan(&historyID, &sequence)) + created := time.Now().Add(-3 * time.Minute) + id, err := uuid.NewV7() + require.NoError(t, err) + binary.BigEndian.PutUint64(id[:8], uint64(created.UnixMilli())<<16) + id[6] = 0x70 + request := uuid.NewString() + legacy := &apiv1.Checkpoint{Id: id.String(), SessionId: source.Id, HeadTaskId: head, + HistorySequence: uint64(sequence), State: apiv1.CheckpointState_CHECKPOINT_STATE_CREATING, + CreatedAt: timestamppb.New(created), Name: "legacy checkpoint"} + data, err := proto.Marshal(legacy) + require.NoError(t, err) + // Seed the old table before claim columns exist. All history and revision + // references come from a settled Session, rather than invented parent rows. + _, err = pool.CopyFrom(t.Context(), pgx.Identifier{"session_checkpoint"}, []string{ + "id", "source_session_id", "user_id", "request_id", "head_task_id", "history_sequence", + "snapshot_atespace", "snapshot_uri", "snapshot_fidelity", "source_history_id", "prepared_revision", "state", "data", + }, pgx.CopyFromRows([][]any{{id, source.Id, "alice", request, head, sequence, + "team-a", "s3://snapshots/source", "SNAPSHOT_FIDELITY_VOLUMES", historyID, source.PreparedRevision, "CREATING", data}})) + require.NoError(t, err) + var claimColumns int + require.NoError(t, pool.QueryRow(t.Context(), `SELECT count(*) FROM information_schema.columns + WHERE table_schema = 'public' AND table_name = 'session_checkpoint' + AND column_name IN ('executor_id', 'executor_expires_at')`).Scan(&claimColumns)) + require.Zero(t, claimColumns) + + require.NoError(t, migrations.RunUp(t.Context(), config, sources)) + op, err := persistence.GetCheckpointOperation(t.Context(), id.String(), "alice") + require.NoError(t, err) + require.True(t, proto.Equal(legacy, op.Checkpoint), "migration must preserve the complete legacy checkpoint") + var unclaimed bool + require.NoError(t, pool.QueryRow(t.Context(), `SELECT executor_id IS NULL AND executor_expires_at IS NULL + FROM session_checkpoint WHERE id = $1`, id).Scan(&unclaimed)) + require.True(t, unclaimed) + _, err = persistence.BeginSessionOperation(t.Context(), source.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE, false) + require.ErrorIs(t, err, database.ErrConflict, "migration must retain source admission") + + tags := &testTags{snapshotURI: "s3://snapshots/source"} + service := NewService(persistence, testAuthorizer{}, tags, nil) + require.NoError(t, NewRecoveryWorker(service).sweep(t.Context())) + recovered, err := persistence.GetSessionCheckpoint(t.Context(), legacy.Id, "alice") + require.NoError(t, err, "the sweep must finish legacy work before any client retries") + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_READY, recovered.State) + owner := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) + replayed, err := service.Create(owner, source.Id, request, head) + require.NoError(t, err) + require.Equal(t, legacy.Id, replayed.Id) + require.Equal(t, legacy.Name, replayed.Name) + require.Equal(t, legacy.HistorySequence, replayed.HistorySequence) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_READY, replayed.State) + require.Equal(t, 1, tags.createCalls, "replaying the old request must not create another copy") + retained, err := persistence.GetSessionCheckpointForkSource(owner, legacy.Id, "alice") + require.NoError(t, err) + require.Equal(t, "tag-uid", retained.TagUID) + require.Equal(t, "s3://tags/checkpoint", retained.Snapshot.URI) + require.NoError(t, service.Delete(owner, legacy.Id)) + _, err = persistence.BeginSessionOperation(owner, source.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE, false) + require.NoError(t, err, "recovered work must stop blocking source deletion") +} diff --git a/go/core/internal/service/checkpoint/operation_recovery_test.go b/go/core/internal/service/checkpoint/operation_recovery_test.go new file mode 100644 index 0000000000..10f6a8e216 --- /dev/null +++ b/go/core/internal/service/checkpoint/operation_recovery_test.go @@ -0,0 +1,492 @@ +package checkpoint + +import ( + "context" + "errors" + "sync/atomic" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + ateapipb "go.ate.dev/substrate/api/ateapi/v1" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" + "github.com/kagent-dev/kagent/go/core/internal/database" + "github.com/kagent-dev/kagent/go/core/internal/dbtest" + "github.com/kagent-dev/kagent/go/core/internal/service/serviceerrors" + "github.com/kagent-dev/kagent/go/core/pkg/auth" +) + +func TestCheckpointRequestsAndRecovery(t *testing.T) { + store := recoveryStore(t) + owner := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) + + t.Run("Given a ready source, when Create is called without a worker, assert that it returns a retained checkpoint inline", func(t *testing.T) { + source, head := recoverySource(t, store) + counted := &countedClaims{Client: store} + tags := &testTags{snapshotURI: "s3://snapshots/source"} + service := NewService(counted, testAuthorizer{}, tags, nil) + request := uuid.NewString() + cp, err := service.Create(owner, source.Id, request, head) + require.NoError(t, err) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_READY, cp.State) + require.Zero(t, counted.claims.Load(), "the reservation INSERT must own the first attempt") + retained, err := store.GetCheckpointOperation(t.Context(), cp.Id, "alice") + require.NoError(t, err) + require.Equal(t, "tag-uid", retained.TagUID) + require.Equal(t, "s3://tags/checkpoint", retained.Snapshot.URI) + repeated, err := service.Create(owner, source.Id, request, head) + require.NoError(t, err) + require.Equal(t, cp.Id, repeated.Id) + require.Equal(t, 1, tags.createCalls) + require.NoError(t, service.Delete(owner, cp.Id)) + }) + + for _, later := range []codes.Code{codes.OK, codes.Unavailable} { + t.Run("Given a denied deletion, when its backend changes to "+later.String()+", assert that a public retry executes the current attempt", func(t *testing.T) { + source, head := recoverySource(t, store) + tags := &testTags{snapshotURI: "s3://snapshots/source"} + service := NewService(store, testAuthorizer{}, tags, nil) + cp, err := service.Create(owner, source.Id, uuid.NewString(), head) + require.NoError(t, err) + tags.deleteErr = status.Error(codes.PermissionDenied, "private bucket and credentials") + require.Equal(t, codes.PermissionDenied, status.Code(service.Delete(owner, cp.Id))) + calls := tags.deleteCalls + tags.deleteErr = nil + if later != codes.OK { + tags.deleteErr = status.Error(later, "temporary outage") + } + replacement := NewService(store, testAuthorizer{}, tags, nil) + err = replacement.Delete(owner, cp.Id) + require.Equal(t, later, status.Code(err)) + require.Equal(t, calls+1, tags.deleteCalls, "retry must issue a new backend attempt") + if later != codes.OK { + pending, readErr := store.GetCheckpointOperation(t.Context(), cp.Id, "alice") + require.NoError(t, readErr) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_DELETING, pending.Checkpoint.State) + require.Equal(t, "tag-uid", pending.TagUID) + tags.deleteErr = nil + require.NoError(t, replacement.Delete(owner, cp.Id)) + } + _, err = store.GetCheckpointOperation(t.Context(), cp.Id, "alice") + require.ErrorIs(t, err, database.ErrNotFound) + }) + } + + t.Run("Given a canceled caller with a live claim, when a replica retries and later takes over, assert that only the owning request copies", func(t *testing.T) { + source, head := recoverySource(t, store) + gate := make(chan struct{}) + tags := &lockedTags{base: &testTags{snapshotURI: "s3://snapshots/source"}, gate: gate, entered: make(chan struct{})} + service := NewService(store, testAuthorizer{}, tags, nil) + request := uuid.NewString() + ctx, cancel := context.WithCancel(owner) + defer cancel() + done := make(chan error, 1) + go func() { _, err := service.Create(ctx, source.Id, request, head); done <- err }() + select { + case <-tags.entered: + case <-time.After(5 * time.Second): + t.Fatal("inline copy was never reached") + } + follower := NewService(store, testAuthorizer{}, tags, nil) + _, err := follower.Create(owner, source.Id, request, head) + require.True(t, serviceerrors.IsCode(err, serviceerrors.CodeAborted), "%v", err) + cancel() + require.ErrorIs(t, <-done, context.Canceled) + close(gate) + cp, err := follower.Create(owner, source.Id, request, head) + require.NoError(t, err) + require.Equal(t, 1, tags.base.createCalls) + require.NoError(t, follower.Delete(owner, cp.Id)) + }) + + t.Run("Given abandoned work, when a replacement sweeps, assert that expired creation is published", func(t *testing.T) { + source, head := recoverySource(t, store) + tags := &testTags{snapshotURI: "s3://snapshots/source"} + service := NewService(store, testAuthorizer{}, tags, nil) + cp, _, err := service.beginCreate(owner, source.Id, uuid.NewString(), head) + require.NoError(t, err) + replacement := NewService(store, testAuthorizer{}, tags, nil) + require.NoError(t, NewRecoveryWorker(replacement).sweep(t.Context())) + ready, err := store.GetSessionCheckpoint(t.Context(), cp.Id, "alice") + require.NoError(t, err) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_READY, ready.State) + require.NoError(t, replacement.Delete(owner, cp.Id)) + }) + + t.Run("Given a stopped partial copy, when it fails verification, assert that admission opens while its UID remains deletable after restart", func(t *testing.T) { + source, head := recoverySource(t, store) + tags := &testTags{snapshotURI: "s3://snapshots/source", mutateTag: func(tag *ateapipb.Tag) { + tag.Status.Snapshot.DurableSnapshot.State = ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_PENDING + }} + service := NewService(store, testAuthorizer{}, tags, nil) + request := uuid.NewString() + _, err := service.Create(owner, source.Id, request, head) + require.Equal(t, codes.FailedPrecondition, status.Code(err)) + cp, _, err := service.beginCreate(owner, source.Id, request, head) + require.NoError(t, err) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_FAILED, cp.State) + pending, err := store.GetCheckpointOperation(t.Context(), cp.Id, "alice") + require.NoError(t, err) + require.Equal(t, "tag-uid", pending.TagUID) + require.Equal(t, "s3://snapshots/source", pending.Snapshot.URI, "pending copy has no completed tag URI") + _, err = store.GetSessionCheckpointForkSource(t.Context(), cp.Id, "alice") + require.ErrorIs(t, err, database.ErrNotFound) + _, err = store.BeginSessionOperation(t.Context(), source.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE, false) + require.NoError(t, err, "the serialized copy has stopped; source admission must open") + tags.deleteErr = status.Error(codes.PermissionDenied, "storage denied") + require.Equal(t, codes.PermissionDenied, status.Code(service.Delete(owner, cp.Id))) + retained, err := store.GetCheckpointOperation(t.Context(), cp.Id, "alice") + require.NoError(t, err) + require.Equal(t, "tag-uid", retained.TagUID) + tags.deleteErr = nil + require.NoError(t, NewService(store, testAuthorizer{}, tags, nil).Delete(owner, cp.Id)) + require.Nil(t, tags.created) + }) + + t.Run("Given lost create and publication replies, when the same request retries, assert that it adopts the retained tag", func(t *testing.T) { + for _, publishFailure := range []bool{false, true} { + source, head := recoverySource(t, store) + tags := &testTags{snapshotURI: "s3://snapshots/source"} + var persistence storeInterface = store + if publishFailure { + persistence = &failedCheckpointPublish{Client: store, entered: make(chan struct{})} + } else { + tags.createErr = status.Error(codes.Unavailable, "copy reply lost") + } + service := NewService(persistence, testAuthorizer{}, tags, nil) + request := uuid.NewString() + _, err := service.Create(owner, source.Id, request, head) + require.Error(t, err) + require.NotNil(t, tags.created) + uid := tags.created.Metadata.Uid + replacement := NewService(store, testAuthorizer{}, tags, nil) + cp, err := replacement.Create(owner, source.Id, request, head) + require.NoError(t, err) + op, err := store.GetCheckpointOperation(t.Context(), cp.Id, "alice") + require.NoError(t, err) + require.Equal(t, uid, op.TagUID) + require.NoError(t, replacement.Delete(owner, cp.Id)) + } + }) + + t.Run("Given a lost delete reply, when a replacement retries, assert that backend absence completes deletion", func(t *testing.T) { + source, head := recoverySource(t, store) + tags := &testTags{snapshotURI: "s3://snapshots/source", deleteResponseErr: status.Error(codes.Unavailable, "delete reply lost")} + service := NewService(store, testAuthorizer{}, tags, nil) + cp, err := service.Create(owner, source.Id, uuid.NewString(), head) + require.NoError(t, err) + require.Equal(t, codes.Unavailable, status.Code(service.Delete(owner, cp.Id))) + require.Nil(t, tags.created) + require.NoError(t, NewService(store, testAuthorizer{}, tags, nil).Delete(owner, cp.Id)) + require.Equal(t, 1, tags.deleteCalls) + }) + + t.Run("Given a lost claim reply, when the request retries, assert that release makes it immediately reclaimable", func(t *testing.T) { + source, head := recoverySource(t, store) + tags := &testTags{snapshotURI: "s3://snapshots/source"} + service := NewService(store, testAuthorizer{}, tags, nil) + request := uuid.NewString() + cp, _, err := service.beginCreate(owner, source.Id, request, head) + require.NoError(t, err) + lost := NewService(&lostCheckpointClaim{Client: store}, testAuthorizer{}, tags, nil) + _, err = lost.Create(owner, source.Id, request, head) + require.Equal(t, codes.Unavailable, status.Code(err)) + require.Zero(t, tags.createCalls) + ready, err := service.Create(owner, source.Id, request, head) + require.NoError(t, err) + require.Equal(t, cp.Id, ready.Id) + require.NoError(t, service.Delete(owner, cp.Id)) + }) + t.Run("Given an executor paused after claiming, when its deadline passes and a replacement succeeds, assert that the stale executor cannot copy", func(t *testing.T) { + source, head := recoverySource(t, store) + tags := &testTags{snapshotURI: "s3://snapshots/source"} + service := NewService(store, testAuthorizer{}, tags, nil) + request := uuid.NewString() + cp, _, err := service.beginCreate(owner, source.Id, request, head) + require.NoError(t, err) + resume := make(chan struct{}) + paused := &pausedCheckpointClaim{Client: store, claimed: make(chan struct{}), resume: resume} + old := NewService(paused, testAuthorizer{}, tags, nil) + ctx, cancel := context.WithTimeout(owner, 80*time.Millisecond) + defer cancel() + done := make(chan error, 1) + go func() { _, err := old.Create(ctx, source.Id, request, head); done <- err }() + select { + case <-paused.claimed: + case <-time.After(5 * time.Second): + t.Fatal("claim never reached") + } + <-ctx.Done() + ready, err := service.Create(owner, source.Id, request, head) + require.NoError(t, err) + require.Equal(t, cp.Id, ready.Id) + close(resume) + require.ErrorIs(t, <-done, context.DeadlineExceeded) + require.Equal(t, 1, tags.createCalls) + require.NoError(t, service.Delete(owner, cp.Id)) + }) + t.Run("Given a lost reservation reply, when the same request retries, assert that it reclaims the committed reservation immediately", func(t *testing.T) { + source, head := recoverySource(t, store) + tags := &testTags{snapshotURI: "s3://snapshots/source"} + request := uuid.NewString() + lost := NewService(&lostCheckpointReservation{Client: store}, testAuthorizer{}, tags, nil) + _, err := lost.Create(owner, source.Id, request, head) + require.Error(t, err) + require.Zero(t, tags.createCalls) + replacement := NewService(store, testAuthorizer{}, tags, nil) + cp, err := replacement.Create(owner, source.Id, request, head) + require.NoError(t, err) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_READY, cp.State) + require.NoError(t, replacement.Delete(owner, cp.Id)) + }) + + t.Run("Given a lost successful finalization reply, when Create retries, assert that the published checkpoint is returned without copying again", func(t *testing.T) { + source, head := recoverySource(t, store) + tags := &testTags{snapshotURI: "s3://snapshots/source"} + request := uuid.NewString() + lost := NewService(&lostCheckpointFinalization{Client: store}, testAuthorizer{}, tags, nil) + _, err := lost.Create(owner, source.Id, request, head) + require.Error(t, err) + replacement := NewService(store, testAuthorizer{}, tags, nil) + cp, err := replacement.Create(owner, source.Id, request, head) + require.NoError(t, err) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_READY, cp.State) + require.Equal(t, 1, tags.createCalls) + require.NoError(t, replacement.Delete(owner, cp.Id)) + }) +} + +func TestDefiniteAcquisitionFailuresDoNotRetireClaims(t *testing.T) { + for _, failure := range []error{database.ErrCheckpointAdvanced, database.ErrSnapshotPending, + database.ErrIdempotencyConflict, database.ErrNotFound, database.ErrConflict, database.ErrFailedPrecondition} { + t.Run("Given a definite reserve rejection, when Create returns "+failure.Error()+", assert that no retirement is attempted", func(t *testing.T) { + persistence := &acquisitionFailureStore{testStore: &testStore{reserveErr: failure}} + tags := &testTags{} + owner := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) + _, err := NewService(persistence, testAuthorizer{}, tags, nil).Create(owner, uuid.NewString(), uuid.NewString(), "head") + require.Error(t, err) + require.Zero(t, persistence.releases) + require.Zero(t, tags.createCalls) + }) + } + for _, operation := range []string{"Delete", "claim"} { + t.Run("Given a definite precondition rejection, when "+operation+" is attempted, assert that no retirement is attempted", func(t *testing.T) { + persistence := &acquisitionFailureStore{testStore: &testStore{}, failure: database.ErrFailedPrecondition} + service := NewService(persistence, testAuthorizer{}, nil, nil) + owner := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) + var err error + if operation == "Delete" { + err = service.Delete(owner, uuid.NewString()) + } else { + _, err = service.claimAndExecute(owner, &database.CheckpointOperation{Checkpoint: &apiv1.Checkpoint{Id: uuid.NewString()}, Creator: "alice"}) + } + require.Error(t, err) + require.Zero(t, persistence.releases) + }) + } +} + +type acquisitionFailureStore struct { + *testStore + failure error + releases int +} + +func (s *acquisitionFailureStore) BeginDeleteSessionCheckpoint(context.Context, string, string, uuid.UUID) (*database.CheckpointOperation, error) { + return nil, s.failure +} + +func (s *acquisitionFailureStore) ClaimCheckpointOperation(context.Context, string, string, uuid.UUID) (*database.CheckpointOperation, error) { + return nil, s.failure +} + +func (s *acquisitionFailureStore) ReleaseCheckpointOperation(context.Context, string, uuid.UUID) error { + s.releases++ + return errors.New("unexpected claim retirement") +} + +func (s *acquisitionFailureStore) ReleaseCheckpointReservation(context.Context, string, string, uuid.UUID) error { + s.releases++ + return errors.New("unexpected reservation retirement") +} + +// Alias prevents the fixture's local store variable from shadowing the interface. +type storeInterface = store + +type countedClaims struct { + *database.Client + claims atomic.Int64 +} + +func (s *countedClaims) ClaimCheckpointOperation(ctx context.Context, id, owner string, executor uuid.UUID) (*database.CheckpointOperation, error) { + s.claims.Add(1) + return s.Client.ClaimCheckpointOperation(ctx, id, owner, executor) +} + +func recoveryStore(t *testing.T) *database.Client { + t.Helper() + dsn := dbtest.StartT(context.WithoutCancel(t.Context()), t) + require.NoError(t, dbtest.Migrate(dsn, false)) + pool, err := database.Connect(t.Context(), &database.PostgresConfig{Pool: dbtest.PoolConfig(dsn), Schema: "public"}) + require.NoError(t, err) + t.Cleanup(pool.Close) + return database.NewClient(pool, "public") +} + +func TestPublicBackendErrorsHideConfiguration(t *testing.T) { + tags := &testTags{snapshotURI: "s3://snapshots/snapshot-1", createErr: status.Error(codes.PermissionDenied, "secret endpoint and bucket")} + service := NewService(&testStore{}, testAuthorizer{}, tags, nil) + _, err := service.Create(auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}), uuid.NewString(), "request", "head") + require.Equal(t, codes.PermissionDenied, status.Code(err)) + require.NotContains(t, err.Error(), "secret endpoint") + require.False(t, errors.Is(err, database.ErrConflict)) +} + +func TestPublicDeleteReportsCurrentBackendStatus(t *testing.T) { + persistence := recoveryStore(t) + owner := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) + for _, operation := range []string{"lookup", "delete"} { + for _, code := range []codes.Code{codes.FailedPrecondition, codes.InvalidArgument, codes.PermissionDenied, codes.Unauthenticated} { + t.Run(operation+"/"+code.String(), func(t *testing.T) { + source, head := recoverySource(t, persistence) + tags := &testTags{snapshotURI: "s3://snapshots/source"} + service := NewService(persistence, testAuthorizer{}, tags, nil) + cp, err := service.Create(owner, source.Id, uuid.NewString(), head) + require.NoError(t, err) + failure := status.Error(code, "private storage endpoint") + if operation == "lookup" { + tags.getErr = failure + } else { + tags.deleteErr = failure + } + err = service.Delete(owner, cp.Id) + require.Equal(t, code, status.Code(err)) + require.NotContains(t, err.Error(), "private storage") + pending, err := persistence.GetCheckpointOperation(t.Context(), cp.Id, "alice") + require.NoError(t, err) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_DELETING, pending.Checkpoint.State) + require.Equal(t, "tag-uid", pending.TagUID) + tags.getErr, tags.deleteErr = nil, nil + require.NoError(t, NewService(persistence, testAuthorizer{}, tags, nil).Delete(owner, cp.Id)) + }) + } + } +} + +func TestPublicCreateRetainsUnverifiedTag(t *testing.T) { + persistence := recoveryStore(t) + source, head := recoverySource(t, persistence) + tags := &testTags{snapshotURI: "s3://snapshots/source"} + service := NewService(persistence, testAuthorizer{}, tags, nil) + owner := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) + request := uuid.NewString() + cp, _, err := service.beginCreate(owner, source.Id, request, head) + require.NoError(t, err) + _, err = tags.CreateTag(t.Context(), "team-a", tagName(cp.Id), "unrelated") + require.NoError(t, err) + _, err = service.Create(owner, source.Id, request, head) + require.True(t, serviceerrors.IsCode(err, serviceerrors.CodeFailedPrecondition)) + pending, err := persistence.GetCheckpointOperation(t.Context(), cp.Id, "alice") + require.NoError(t, err) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_CREATING, pending.Checkpoint.State) + require.Nil(t, pending.Checkpoint.Failure) + require.Zero(t, tags.deleteCalls) + for _, force := range []bool{false, true} { + _, err = persistence.BeginSessionOperation(t.Context(), source.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE, force) + require.ErrorIs(t, err, database.ErrConflict) + } + // Substrate source references are immutable. Model removal of the unrelated + // tag by its owner instead of pretending it can be edited into our tag. + require.NoError(t, tags.DeleteTag(t.Context(), "team-a", tagName(cp.Id), tags.created.Metadata.Uid)) + ready, err := NewService(persistence, testAuthorizer{}, tags, nil).Create(owner, source.Id, request, head) + require.NoError(t, err) + require.Equal(t, cp.Id, ready.Id) + require.NoError(t, service.Delete(owner, cp.Id)) +} + +func TestSlowRecoveryDoesNotBlockRequests(t *testing.T) { + persistence := recoveryStore(t) + slowSource, slowHead := recoverySource(t, persistence) + fastSource, fastHead := recoverySource(t, persistence) + owner := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) + gate := make(chan struct{}) + slowTags := &lockedTags{base: &testTags{snapshotURI: "s3://snapshots/source"}, gate: gate, entered: make(chan struct{})} + slow := NewService(persistence, testAuthorizer{}, slowTags, nil) + cp, _, err := slow.beginCreate(owner, slowSource.Id, uuid.NewString(), slowHead) + require.NoError(t, err) + done := make(chan error, 1) + go func() { done <- NewRecoveryWorker(slow).sweep(t.Context()) }() + select { + case <-slowTags.entered: + case <-time.After(5 * time.Second): + t.Fatal("sweep never began") + } + fast := NewService(persistence, testAuthorizer{}, &testTags{snapshotURI: "s3://snapshots/source"}, nil) + ctx, cancel := context.WithTimeout(owner, time.Second) + defer cancel() + ready, err := fast.Create(ctx, fastSource.Id, uuid.NewString(), fastHead) + require.NoError(t, err, "a request must execute without waiting behind the recovery sweep") + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_READY, ready.State) + close(gate) + require.NoError(t, <-done) + require.NoError(t, slow.Delete(owner, cp.Id)) + require.NoError(t, fast.Delete(owner, ready.Id)) +} + +// The reservation committed; only its successful reply was lost. +type lostCheckpointReservation struct{ *database.Client } + +func (s *lostCheckpointReservation) ReserveSessionCheckpoint(ctx context.Context, checkpoint *apiv1.Checkpoint, owner, request string, executor uuid.UUID) (*apiv1.Checkpoint, *database.SessionTaskSnapshot, error) { + if _, _, err := s.Client.ReserveSessionCheckpoint(ctx, checkpoint, owner, request, executor); err != nil { + return nil, nil, err + } + return nil, nil, status.Error(codes.Unavailable, "reservation reply lost") +} + +type lostCheckpointFinalization struct{ *database.Client } + +func (s *lostCheckpointFinalization) FinalizeClaimedSessionCheckpoint(ctx context.Context, checkpoint *apiv1.Checkpoint, executor uuid.UUID, uid, uri, failure string) (*apiv1.Checkpoint, error) { + if _, err := s.Client.FinalizeClaimedSessionCheckpoint(ctx, checkpoint, executor, uid, uri, failure); err != nil { + return nil, err + } + return nil, status.Error(codes.Unavailable, "finalization reply lost") +} + +func TestPublicCheckpointErrorsHideReleaseFailures(t *testing.T) { + persistence := recoveryStore(t) + owner := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) + for _, operation := range []string{"Create", "Delete"} { + t.Run("Given a private claim release failure, when "+operation+" also fails at the backend, assert that its current status excludes database details", func(t *testing.T) { + source, head := recoverySource(t, persistence) + tags := &testTags{snapshotURI: "s3://snapshots/source"} + service := NewService(&lostCheckpointRelease{Client: persistence}, testAuthorizer{}, tags, nil) + var err error + if operation == "Create" { + tags.createErr = status.Error(codes.PermissionDenied, "private backend endpoint") + _, err = service.Create(owner, source.Id, uuid.NewString(), head) + } else { + cp, createErr := service.Create(owner, source.Id, uuid.NewString(), head) + require.NoError(t, createErr) + tags.deleteErr = status.Error(codes.PermissionDenied, "private backend endpoint") + err = service.Delete(owner, cp.Id) + } + require.Equal(t, codes.PermissionDenied, status.Code(err)) + require.NotContains(t, err.Error(), "private") + }) + } +} + +type lostCheckpointRelease struct { + *database.Client +} + +func (s *lostCheckpointRelease) ReleaseCheckpointOperation(ctx context.Context, id string, executor uuid.UUID) error { + if err := s.Client.ReleaseCheckpointOperation(ctx, id, executor); err != nil { + return err + } + return errors.New("private database endpoint in lost release reply") +} diff --git a/go/core/internal/service/checkpoint/public_errors_test.go b/go/core/internal/service/checkpoint/public_errors_test.go new file mode 100644 index 0000000000..57f37d0abe --- /dev/null +++ b/go/core/internal/service/checkpoint/public_errors_test.go @@ -0,0 +1,67 @@ +package checkpoint + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/kagent-dev/kagent/go/core/internal/service/serviceerrors" + "github.com/kagent-dev/kagent/go/core/pkg/auth" +) + +func TestCreateHidesReservationReleaseFailuresAfterCancellation(t *testing.T) { + store := &failedReservationRelease{testStore: &testStore{reserveErr: errors.New("private reservation connection detail")}} + owner, cancel := context.WithCancel(auth.AuthSessionTo(t.Context(), testSession{userID: "alice"})) + cancel() + _, err := NewService(store, testAuthorizer{}, &testTags{}, nil).Create(owner, uuid.NewString(), uuid.NewString(), "task-1") + require.Equal(t, 1, store.calls) + require.NoError(t, store.contextErr, "reservation retirement must outlive caller cancellation") + require.True(t, store.bounded, "retirement must keep the five-second bound") + require.Equal(t, serviceerrors.CodeInternal, serviceerrors.CodeOf(err)) + require.NotContains(t, serviceerrors.MessageOf(err), "private") +} + +func TestCreateClassifiesPublicationFailures(t *testing.T) { + for _, scenario := range []struct { + name string + err error + serviceCode serviceerrors.Code + grpcCode codes.Code + }{ + {"storage failure", errors.New("private PostgreSQL connection detail"), serviceerrors.CodeInternal, codes.Unknown}, + {"lost reply", status.Error(codes.Unavailable, "private finalization reply"), "", codes.Unavailable}, + } { + t.Run("Given "+scenario.name+", when checkpoint publication fails, assert that only a classified public error is returned", func(t *testing.T) { + store := &testStore{finalizeErr: scenario.err} + tags := &testTags{snapshotURI: "s3://snapshots/snapshot-1"} + owner := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) + _, err := NewService(store, testAuthorizer{}, tags, nil).Create(owner, uuid.NewString(), uuid.NewString(), "task-1") + require.Error(t, err) + require.Equal(t, 1, tags.createCalls) + require.Equal(t, scenario.serviceCode, serviceerrors.CodeOf(err)) + require.Equal(t, scenario.grpcCode, status.Code(err)) + require.NotContains(t, serviceerrors.MessageOf(err), "private") + }) + } +} + +type failedReservationRelease struct { + *testStore + calls int + contextErr error + bounded bool +} + +func (s *failedReservationRelease) ReleaseCheckpointReservation(ctx context.Context, _, _ string, _ uuid.UUID) error { + s.calls++ + s.contextErr = ctx.Err() + deadline, ok := ctx.Deadline() + s.bounded = ok && time.Until(deadline) > 0 && time.Until(deadline) <= 5*time.Second + return errors.New("private retirement connection detail") +} diff --git a/go/core/internal/service/checkpoint/recovery.go b/go/core/internal/service/checkpoint/recovery.go new file mode 100644 index 0000000000..113c5e593a --- /dev/null +++ b/go/core/internal/service/checkpoint/recovery.go @@ -0,0 +1,203 @@ +package checkpoint + +import ( + "context" + "errors" + "time" + + "github.com/google/uuid" + "sigs.k8s.io/controller-runtime/pkg/manager" + + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" + "github.com/kagent-dev/kagent/go/core/internal/database" + "github.com/kagent-dev/kagent/go/core/internal/service/serviceerrors" + "github.com/kagent-dev/kagent/go/core/internal/substrate" + "github.com/kagent-dev/kagent/go/pkg/logging" +) + +// RecoveryWorker resumes expired durable work with the same execution path as +// requests. The sweep interval bounds retries; live requests are never queued. +type RecoveryWorker struct { + service *Service + interval time.Duration +} + +var _ manager.Runnable = (*RecoveryWorker)(nil) +var _ manager.LeaderElectionRunnable = (*RecoveryWorker)(nil) + +func NewRecoveryWorker(service *Service) *RecoveryWorker { + return &RecoveryWorker{service: service, interval: 30 * time.Second} +} + +func (*RecoveryWorker) NeedLeaderElection() bool { return true } + +func (w *RecoveryWorker) Start(ctx context.Context) error { + for ctx.Err() == nil { + if err := w.sweep(ctx); err != nil && ctx.Err() == nil { + logging.FromContext(ctx).ErrorContext(ctx, "recover checkpoints", "error", err) + } + // Start the delay after the sweep: a ticker can retain a stale tick + // during slow recovery and immediately retry the last failed attempt. + timer := time.NewTimer(w.interval) + select { + case <-ctx.Done(): + timer.Stop() + return nil + case <-timer.C: + } + } + return nil +} + +func (w *RecoveryWorker) sweep(ctx context.Context) error { + var after string + for ctx.Err() == nil { + operations, err := w.service.store.ListPendingCheckpoints(ctx, after, 100) + if err != nil { + return err + } + for _, op := range operations { + if err := w.service.recoverOperation(ctx, op); err != nil && ctx.Err() == nil { + logging.FromContext(ctx).ErrorContext(ctx, "recover checkpoint", "checkpoint_id", op.Checkpoint.Id, "error", err) + } + } + if len(operations) < 100 { + return nil + } + after = operations[len(operations)-1].Checkpoint.Id + } + return ctx.Err() +} + +func (s *Service) recoverOperation(ctx context.Context, observed *database.CheckpointOperation) error { + ctx, cancel := context.WithTimeout(ctx, database.RuntimeOperationTimeout) + defer cancel() + _, err := s.claimAndExecute(ctx, observed) + return err +} + +// Claim rereads the current phase: a sweep's observation can already be stale. +// A failed acquisition returns the current outcome rather than polling or +// replaying an earlier attempt's error. +func (s *Service) claimAndExecute(ctx context.Context, observed *database.CheckpointOperation) (*apiv1.Checkpoint, error) { + executor := uuid.New() + id, owner := observed.Checkpoint.Id, observed.Creator + op, err := s.store.ClaimCheckpointOperation(ctx, id, owner, executor) + if errors.Is(err, database.ErrConflict) { + // A definite conflict acquired nothing. Do not let an unrelated release + // failure replace the outcome this caller can only observe. + return s.currentOutcome(ctx, id, owner) + } + if err != nil { + if definiteAcquisitionError(err) { + return nil, publicOperationError(err) + } + // A lost COMMIT reply can hide successful acquisition. Retire this + // executor before returning the uncertain error, as on reservation. + return nil, errors.Join(publicOperationError(err), s.release(ctx, id, executor)) + } + return s.execute(ctx, op, executor, false) +} + +// Losing execution authority permits observation, never another backend call. +// Absence lets Delete succeed; Create translates it to NotFound at its boundary. +func (s *Service) currentOutcome(ctx context.Context, id, owner string) (*apiv1.Checkpoint, error) { + current, err := s.store.GetCheckpointOperation(ctx, id, owner) + if errors.Is(err, database.ErrNotFound) { + return nil, nil + } + if err != nil { + return nil, serviceerrors.NewInternal("Failed to observe checkpoint operation", err) + } + return createOutcome(current.Checkpoint) +} + +func (s *Service) release(ctx context.Context, id string, executor uuid.UUID) error { + return releaseClaim(ctx, id, func(releaseCtx context.Context) error { + return s.store.ReleaseCheckpointOperation(releaseCtx, id, executor) + }) +} + +// Release stays bounded after caller cancellation, and neither release path may +// expose storage details when its error is joined to the public attempt result. +func releaseClaim(ctx context.Context, id string, release func(context.Context) error) error { + releaseCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + defer cancel() + if err := release(releaseCtx); err != nil { + logging.FromContext(ctx).ErrorContext(ctx, "release checkpoint execution claim", "checkpoint_id", id, "error", err) + // This error is joined to the current attempt's public status. Keep + // database details in logs without replacing the original status code. + return errors.New("checkpoint execution claim could not be released") + } + return nil +} + +func (s *Service) execute(ctx context.Context, op *database.CheckpointOperation, executor uuid.UUID, firstAttempt bool) (result *apiv1.Checkpoint, err error) { + releaseNeeded := true + defer func() { + if releaseNeeded { + err = errors.Join(err, s.release(ctx, op.Checkpoint.Id, executor)) + } + }() + if err := ctx.Err(); err != nil { + return nil, err + } + if op.Checkpoint.State == apiv1.CheckpointState_CHECKPOINT_STATE_DELETING { + err = s.finishDelete(ctx, op, executor) + releaseNeeded = err != nil + return nil, err + } + if op.Checkpoint.State == apiv1.CheckpointState_CHECKPOINT_STATE_FAILED { + err = s.collectFailedTag(ctx, op, executor) + releaseNeeded = err != nil + return nil, err + } + tag, err := s.ensureTag(ctx, op.Checkpoint, op.Snapshot, firstAttempt) + if err != nil && !errors.Is(err, errTagCopyFailed) { + return nil, publicOperationError(err) + } + if err != nil && tag != nil && !validTagSource(tag, op.Checkpoint, op.Snapshot) { + return nil, serviceerrors.NewFailedPrecondition("Checkpoint snapshot tag identity changed", err) + } + failure := "" + if err != nil { + failure = checkpointFailureMessage + } + // FAILED releases source admission; the sweep collects its retained tag by UID. + finishCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + defer cancel() + checkpoint, err := s.store.FinalizeClaimedSessionCheckpoint(finishCtx, op.Checkpoint, executor, + tag.GetMetadata().GetUid(), substrate.SnapshotURI(tag.GetStatus().GetSnapshot()), failure) + if errors.Is(err, database.ErrConflict) { + // The conditional write proves authority was lost. This is now an + // observer; a stale release must not override the result it reads. + releaseNeeded = false + return s.currentOutcome(finishCtx, op.Checkpoint.Id, op.Creator) + } + if err != nil { + return nil, publicOperationError(err) + } + releaseNeeded = false + return createOutcome(checkpoint) +} + +func (s *Service) collectFailedTag(ctx context.Context, op *database.CheckpointOperation, executor uuid.UUID) error { + if err := s.deleteTag(ctx, op); err != nil { + return err + } + err := s.store.ClearFailedCheckpointTag(ctx, op.Checkpoint.Id, op.Creator, op.TagUID, executor) + if errors.Is(err, database.ErrConflict) { + current, readErr := s.store.GetCheckpointOperation(ctx, op.Checkpoint.Id, op.Creator) + if errors.Is(readErr, database.ErrNotFound) || (readErr == nil && current.TagUID == "") { + return nil + } + if readErr != nil { + return serviceerrors.NewInternal("Failed to observe checkpoint cleanup", readErr) + } + return serviceerrors.NewAborted("Checkpoint cleanup is still pending", nil) + } + if err != nil { + return serviceerrors.NewInternal("Failed to collect checkpoint tag", err) + } + return nil +} diff --git a/go/core/internal/service/checkpoint/recovery_test.go b/go/core/internal/service/checkpoint/recovery_test.go new file mode 100644 index 0000000000..023edc9917 --- /dev/null +++ b/go/core/internal/service/checkpoint/recovery_test.go @@ -0,0 +1,292 @@ +package checkpoint + +import ( + "context" + "crypto/sha256" + "errors" + "fmt" + "strings" + "sync" + "sync/atomic" + "testing" + "testing/synctest" + "time" + + "github.com/a2aproject/a2a-go/v2/a2a" + a2apb "github.com/a2aproject/a2a-go/v2/a2apb/v1" + "github.com/google/uuid" + "github.com/stretchr/testify/require" + ateapipb "go.ate.dev/substrate/api/ateapi/v1" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + "google.golang.org/protobuf/proto" + + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" + "github.com/kagent-dev/kagent/go/core/internal/database" + "github.com/kagent-dev/kagent/go/core/internal/substrate" +) + +// lockedTags models Substrate's tag lease independently of a request or worker. +// The copy can outlive its client; another CreateTag waits for the lease before +// returning AlreadyExists. GetTag alone can observe a still-live pending copy. +type lockedTags struct { + mu sync.Mutex + base *testTags + gate <-chan struct{} + entered chan struct{} + once sync.Once +} + +var _ tagClient = (*lockedTags)(nil) + +func (t *lockedTags) GetActor(ctx context.Context, space, name string) (*ateapipb.Actor, error) { + t.mu.Lock() + defer t.mu.Unlock() + return t.base.GetActor(ctx, space, name) +} + +func (t *lockedTags) GetTag(ctx context.Context, space, name string) (*ateapipb.Tag, error) { + t.mu.Lock() + defer t.mu.Unlock() + tag, err := t.base.GetTag(ctx, space, name) + return proto.CloneOf(tag), err +} + +func (t *lockedTags) CreateTag(ctx context.Context, space, name, actor string) (*ateapipb.Tag, error) { + if t.entered != nil { + t.once.Do(func() { close(t.entered) }) + } + if t.gate != nil { + select { + case <-t.gate: + case <-ctx.Done(): + return nil, ctx.Err() + } + } + t.mu.Lock() + defer t.mu.Unlock() + tag, err := t.base.CreateTag(ctx, space, name, actor) + return proto.CloneOf(tag), err +} + +func (t *lockedTags) DeleteTag(ctx context.Context, space, name, uid string) error { + t.mu.Lock() + defer t.mu.Unlock() + return t.base.DeleteTag(ctx, space, name, uid) +} + +func recoverySource(t *testing.T, store *database.Client) (*apiv1.Session, string) { + t.Helper() + return recoverySourceForActor(t, store, &ateapipb.Actor{ + Metadata: &ateapipb.ResourceMetadata{Atespace: "team-a", Name: substrate.ActorName(uuid.NewString()), Uid: "actor-uid"}, + ActorTemplate: &ateapipb.ObjectRef{Atespace: "team-a", Name: "runtime"}, + Status: &ateapipb.ActorStatus{Snapshots: []*ateapipb.Snapshot{{ + Generation: 1, ActorTemplateUid: "template-uid", Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES, + DurableSnapshot: &ateapipb.SnapshotStorage{State: ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_COMPLETED, + Object: &ateapipb.ObjectSnapshot{SnapshotUri: "s3://snapshots/source"}}, + }}}, + }) +} + +func recoverySourceForActor(t *testing.T, store *database.Client, actor *ateapipb.Actor) (*apiv1.Session, string) { + t.Helper() + ctx := t.Context() + space := actor.GetMetadata().GetAtespace() + snapshot := substrate.LatestActorSnapshot(actor) + revision := "revision-" + snapshot.GetActorTemplateUid() + sessionID := strings.TrimPrefix(actor.GetMetadata().GetName(), "session-") + _, err := uuid.Parse(sessionID) + require.NoError(t, err) + require.NoError(t, store.UpsertAgentDefinition(ctx, database.AgentDefinition{ + Namespace: space, AgentName: "assistant", AgentUID: "agent-uid", DesiredRevision: revision})) + _, err = store.RecordRuntimeRevision(ctx, database.RuntimeRevision{ + Revision: revision, Namespace: space, AgentName: "assistant", AgentUID: "agent-uid", + SourceSnapshot: []byte("{}"), AgentCard: &a2apb.AgentCard{}, EgressDestinations: []string{}, + ActorTemplateAtespace: space, ActorTemplateName: actor.GetActorTemplate().GetName(), ActorTemplateUID: snapshot.GetActorTemplateUid(), + }, database.RecordRevisionOptions{Ready: true}) + require.NoError(t, err) + session, _, err := store.CreateSession(ctx, &apiv1.Session{Id: sessionID, Creator: "alice", + Agent: &apiv1.AgentProvenance{Ref: &apiv1.ResourceReference{Namespace: space, Name: "assistant"}}}, uuid.NewString()) + require.NoError(t, err) + op, err := store.BeginSessionOperation(ctx, session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_CREATE, false) + require.NoError(t, err) + executor := uuid.New() + claimed, err := store.ClaimSessionOperation(ctx, session.Id, op.ID, executor) + require.NoError(t, err) + require.True(t, claimed) + session, err = store.FinishSessionOperation(ctx, session.Id, op.ID, executor, actor.GetActorTemplate().GetName(), actor.GetMetadata().GetUid(), "") + require.NoError(t, err) + message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("hello")) + message.ContextID = session.ContextId + task := a2a.NewSubmittedTask(message, message) + task.Status.State = a2a.TaskStateCompleted + hash := sha256.Sum256([]byte(task.ID)) + version, err := store.CreateRuntimeTask(ctx, session.Id, hash[:], task, "") + require.NoError(t, err) + require.NoError(t, store.SettleSessionTask(ctx, session.Id, string(task.ID), version)) + boundary, err := store.ClaimSessionQuiescence(ctx) + require.NoError(t, err) + require.NoError(t, store.FinishSessionQuiescence(ctx, boundary, + &database.SessionTaskSnapshot{Atespace: space, URI: substrate.SnapshotURI(snapshot), Fidelity: snapshot.GetFidelity().String()})) + return session, string(task.ID) +} + +type failedCheckpointPublish struct { + *database.Client + entered chan struct{} + once sync.Once +} + +func (s *failedCheckpointPublish) FinalizeClaimedSessionCheckpoint(ctx context.Context, _ *apiv1.Checkpoint, _ uuid.UUID, _, _, _ string) (*apiv1.Checkpoint, error) { + s.once.Do(func() { close(s.entered) }) + return nil, errors.New("publication reply lost") +} + +// pausedCheckpointClaim models an executor paused after its successful database +// reply. Its attempt deadline still elapses while a replacement makes progress. +type pausedCheckpointClaim struct { + *database.Client + claimed chan struct{} + resume <-chan struct{} +} + +type lostCheckpointClaim struct{ *database.Client } + +func (s *lostCheckpointClaim) ClaimCheckpointOperation(ctx context.Context, id, owner string, executor uuid.UUID) (*database.CheckpointOperation, error) { + if _, err := s.Client.ClaimCheckpointOperation(ctx, id, owner, executor); err != nil { + return nil, err + } + return nil, status.Error(codes.Unavailable, "committed claim response lost") +} + +func (s *pausedCheckpointClaim) ClaimCheckpointOperation(ctx context.Context, id, owner string, executor uuid.UUID) (*database.CheckpointOperation, error) { + op, err := s.Client.ClaimCheckpointOperation(ctx, id, owner, executor) + if err != nil { + return nil, err + } + close(s.claimed) + <-s.resume + return op, nil +} + +// An old controller's remote copy can still be running after failover. Model +// the backend lease directly: the new CreateTag cannot answer AlreadyExists +// until that copy ends, even though GetTag would already return a pending row. +type serializedCopyTags struct { + *testTags + entered chan struct{} + release <-chan struct{} +} + +func (s *serializedCopyTags) CreateTag(ctx context.Context, _, _, _ string) (*ateapipb.Tag, error) { + close(s.entered) + select { + case <-s.release: + s.created.Status.Snapshot.DurableSnapshot.State = ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_COMPLETED + return nil, status.Error(codes.AlreadyExists, "copy finished") + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +func TestRecoveryWaitsForLiveBackendCopyBeforeInspectingPendingTag(t *testing.T) { + cp := &apiv1.Checkpoint{Id: "11111111-1111-4111-8111-111111111111", SessionId: "22222222-2222-4222-8222-222222222222", State: apiv1.CheckpointState_CHECKPOINT_STATE_CREATING} + base := &testTags{snapshotURI: "s3://snapshots/source"} + _, err := base.CreateTag(t.Context(), "team-a", "checkpoint-"+cp.Id, "session-"+cp.SessionId) + require.NoError(t, err) + base.created.Status.Snapshot.DurableSnapshot.State = ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_PENDING + release := make(chan struct{}) + tags := &serializedCopyTags{testTags: base, entered: make(chan struct{}), release: release} + store := &testStore{prepared: cp, snapshot: &database.SessionTaskSnapshot{Atespace: "team-a", URI: "s3://snapshots/source", Fidelity: "SNAPSHOT_FIDELITY_VOLUMES"}, checkpointOwner: "alice"} + service := NewService(store, testAuthorizer{}, tags, nil) + op, err := store.GetCheckpointOperation(t.Context(), cp.Id, "alice") + require.NoError(t, err) + result := make(chan error, 1) + go func() { result <- service.recoverOperation(t.Context(), op) }() + select { + case <-tags.entered: + case <-time.After(time.Second): + t.Fatal("backend synchronization was skipped") + } + select { + case err := <-result: + t.Fatalf("recovered before the live copy ended: %v", err) + case <-time.After(50 * time.Millisecond): + } + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_CREATING, cp.State) + require.Zero(t, base.deleteCalls) + require.Zero(t, base.getCalls, "a pending tag must not be interpreted before the lease barrier") + close(release) + require.NoError(t, <-result) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_READY, cp.State) + require.Zero(t, base.deleteCalls) +} + +// Fixture for a process that reserved work but exited before calling Substrate. +func (s *Service) beginCreate(ctx context.Context, source, request, head string) (*apiv1.Checkpoint, string, error) { + ctx, cancel := context.WithTimeout(ctx, database.RuntimeOperationTimeout) + defer cancel() + executor := uuid.New() + id, err := uuid.NewV7() + if err != nil { + return nil, "", err + } + op, err := s.reserveCreate(ctx, source, request, head, id.String(), executor) + err = errors.Join(err, s.release(ctx, id.String(), executor)) + if err != nil { + return nil, "", err + } + return op.Checkpoint, op.Creator, nil +} + +func TestRecoveryWaitsAfterSlowSweep(t *testing.T) { + for _, failed := range []bool{false, true} { + t.Run(fmt.Sprintf("Given a slow sweep with failure=%t, when it finishes, assert that another sweep waits a full interval", failed), func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + gate := make(chan struct{}) + persistence := &blockedCheckpointSweep{testStore: &testStore{}, gate: gate, failed: failed} + worker := NewRecoveryWorker(NewService(persistence, testAuthorizer{}, nil, nil)) + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + done := make(chan error, 1) + go func() { done <- worker.Start(ctx) }() + synctest.Wait() + require.Equal(t, 1, int(persistence.calls.Load())) + time.Sleep(3 * worker.interval) + close(gate) + synctest.Wait() + require.Equal(t, 1, int(persistence.calls.Load()), "a stale tick must not immediately start another sweep") + time.Sleep(worker.interval - time.Millisecond) + synctest.Wait() + require.Equal(t, 1, int(persistence.calls.Load())) + time.Sleep(time.Millisecond) + synctest.Wait() + require.Equal(t, 2, int(persistence.calls.Load())) + cancel() + synctest.Wait() + require.NoError(t, <-done) + }) + }) + } +} + +type blockedCheckpointSweep struct { + *testStore + gate <-chan struct{} + calls atomic.Int32 + failed bool +} + +func (s *blockedCheckpointSweep) ListPendingCheckpoints(ctx context.Context, _ string, _ int) ([]*database.CheckpointOperation, error) { + if s.calls.Add(1) == 1 { + select { + case <-s.gate: + case <-ctx.Done(): + return nil, ctx.Err() + } + if s.failed { + return nil, errors.New("sweep unavailable") + } + } + return nil, nil +} diff --git a/go/core/internal/service/checkpoint/reservation_commit_test.go b/go/core/internal/service/checkpoint/reservation_commit_test.go new file mode 100644 index 0000000000..125f4945f0 --- /dev/null +++ b/go/core/internal/service/checkpoint/reservation_commit_test.go @@ -0,0 +1,120 @@ +package checkpoint + +import ( + "context" + "strings" + "sync" + "testing" + "time" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/stretchr/testify/require" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" + "github.com/kagent-dev/kagent/go/core/internal/database" + "github.com/kagent-dev/kagent/go/core/internal/dbtest" + "github.com/kagent-dev/kagent/go/core/pkg/auth" +) + +func TestCreateRetiresReservationThatCommitsAfterReplyLoss(t *testing.T) { + dsn := dbtest.StartT(context.WithoutCancel(t.Context()), t) + require.NoError(t, dbtest.Migrate(dsn, false)) + pool, err := database.Connect(t.Context(), &database.PostgresConfig{Pool: dbtest.PoolConfig(dsn), Schema: "public"}) + require.NoError(t, err) + t.Cleanup(pool.Close) + persistence := database.NewClient(pool, "public") + source, head := recoverySource(t, persistence) + gate := &reservationCommitGate{entered: make(chan struct{}), resume: make(chan struct{})} + config := dbtest.PoolConfig(dsn) + config.ConnConfig.Tracer = gate + latePool, err := database.Connect(t.Context(), &database.PostgresConfig{Pool: config, Schema: "public"}) + require.NoError(t, err) + t.Cleanup(latePool.Close) + var once sync.Once + resume := func() { once.Do(func() { close(gate.resume) }) } + t.Cleanup(resume) + lost := &lostUncommittedReservation{Client: persistence, reserving: database.NewClient(latePool, "public"), gate: gate, + committed: make(chan error, 1), releasing: make(chan struct{})} + owner := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) + tags := &testTags{snapshotURI: "s3://snapshots/source"} + request := uuid.NewString() + first := make(chan error, 1) + go func() { + _, err := NewService(lost, testAuthorizer{}, tags, nil).Create(owner, source.Id, request, head) + first <- err + }() + select { + case <-lost.releasing: + case <-time.After(5 * time.Second): + t.Fatal("reservation cleanup was not attempted") + } + select { + case err := <-first: + t.Fatalf("Create returned before retiring its uncommitted reservation: %v", err) + case <-time.After(100 * time.Millisecond): + } + resume() + require.NoError(t, <-lost.committed, "the reservation must actually commit") + require.Error(t, <-first, "reply loss remains an error after its claim is retired") + require.Zero(t, tags.createCalls) + service := NewService(persistence, testAuthorizer{}, tags, nil) + result, err := service.Create(owner, source.Id, request, head) + require.NoError(t, err, "the same request must reclaim immediately without waiting for expiry") + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_READY, result.GetState()) + require.Equal(t, 1, tags.createCalls) + require.NoError(t, service.Delete(owner, result.Id)) +} + +type reservationCommitGate struct { + entered, resume chan struct{} + once sync.Once +} + +func (g *reservationCommitGate) TraceQueryStart(ctx context.Context, _ *pgx.Conn, data pgx.TraceQueryStartData) context.Context { + if strings.EqualFold(data.SQL, "commit") { + g.once.Do(func() { + close(g.entered) + select { + case <-g.resume: + case <-ctx.Done(): + } + }) + } + return ctx +} + +func (*reservationCommitGate) TraceQueryEnd(context.Context, *pgx.Conn, pgx.TraceQueryEndData) {} + +// PostgreSQL still commits the real reservation; only its transport reply is +// interrupted before completion becomes visible. No stored phase is fabricated. +type lostUncommittedReservation struct { + *database.Client + reserving *database.Client + gate *reservationCommitGate + committed chan error + releasing chan struct{} +} + +func (s *lostUncommittedReservation) ReserveSessionCheckpoint(ctx context.Context, checkpoint *apiv1.Checkpoint, owner, request string, executor uuid.UUID) (*apiv1.Checkpoint, *database.SessionTaskSnapshot, error) { + attempt, cancel := context.WithTimeout(context.WithoutCancel(ctx), time.Minute) + go func() { + defer cancel() + _, _, err := s.reserving.ReserveSessionCheckpoint(attempt, checkpoint, owner, request, executor) + s.committed <- err + }() + <-s.gate.entered + return nil, nil, status.Error(codes.Unavailable, "reservation reply interrupted") +} + +func (s *lostUncommittedReservation) ReleaseCheckpointReservation(ctx context.Context, sessionID, id string, executor uuid.UUID) error { + close(s.releasing) + return s.Client.ReleaseCheckpointReservation(ctx, sessionID, id, executor) +} + +func (s *lostUncommittedReservation) ReleaseCheckpointOperation(ctx context.Context, id string, executor uuid.UUID) error { + close(s.releasing) + return s.Client.ReleaseCheckpointOperation(ctx, id, executor) +} diff --git a/go/core/internal/service/checkpoint/service.go b/go/core/internal/service/checkpoint/service.go index de5cb6a390..7a1b67ebb0 100644 --- a/go/core/internal/service/checkpoint/service.go +++ b/go/core/internal/service/checkpoint/service.go @@ -9,7 +9,6 @@ import ( "github.com/google/uuid" ateapipb "go.ate.dev/substrate/api/ateapi/v1" - "golang.org/x/sync/singleflight" "google.golang.org/genproto/googleapis/rpc/errdetails" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" @@ -22,18 +21,25 @@ import ( ) const ( - defaultPageSize = 50 - maxPageSize = 100 + defaultPageSize = 50 + maxPageSize = 100 + checkpointFailureMessage = "Checkpoint snapshot copy could not be verified" ) type store interface { - ReserveSessionCheckpoint(context.Context, *apiv1.Checkpoint, string, string) (*apiv1.Checkpoint, *database.SessionTaskSnapshot, error) - FinalizeSessionCheckpoint(context.Context, string, string, string, string) (*apiv1.Checkpoint, error) + GetCheckpointOperation(context.Context, string, string) (*database.CheckpointOperation, error) + ListPendingCheckpoints(context.Context, string, int) ([]*database.CheckpointOperation, error) + ClaimCheckpointOperation(context.Context, string, string, uuid.UUID) (*database.CheckpointOperation, error) + ReleaseCheckpointOperation(context.Context, string, uuid.UUID) error + ReleaseCheckpointReservation(context.Context, string, string, uuid.UUID) error + ClearFailedCheckpointTag(context.Context, string, string, string, uuid.UUID) error + ReserveSessionCheckpoint(context.Context, *apiv1.Checkpoint, string, string, uuid.UUID) (*apiv1.Checkpoint, *database.SessionTaskSnapshot, error) + FinalizeClaimedSessionCheckpoint(context.Context, *apiv1.Checkpoint, uuid.UUID, string, string, string) (*apiv1.Checkpoint, error) GetSessionCheckpoint(context.Context, string, string) (*apiv1.Checkpoint, error) GetSessionCheckpointForkSource(context.Context, string, string) (*database.SessionCheckpointForkSource, error) ListSessionCheckpoints(context.Context, string, string, string, int) ([]*apiv1.Checkpoint, error) - BeginDeleteSessionCheckpoint(context.Context, string, string) (*database.SessionTaskSnapshot, string, error) - DeleteSessionCheckpoint(context.Context, string, string) error + BeginDeleteSessionCheckpoint(context.Context, string, string, uuid.UUID) (*database.CheckpointOperation, error) + DeleteClaimedSessionCheckpoint(context.Context, string, string, uuid.UUID) error ForkSession(context.Context, string, string, string, string) (*apiv1.Session, bool, error) UpdateCheckpointName(context.Context, string, string, string) (*apiv1.Checkpoint, error) } @@ -49,11 +55,10 @@ type tagClient interface { DeleteTag(context.Context, string, string, string) error } +var _ store = (*database.Client)(nil) +var _ tagClient = (*substrate.Client)(nil) + type Service struct { - // creates coalesces identical requests within this service session. - // TODO: route tag creation and cleanup through durable session ownership - // so retries on different replicas cannot race. - creates singleflight.Group store store authorizer auth.Authorizer tags tagClient @@ -75,7 +80,35 @@ func NewService(store store, authorizer auth.Authorizer, tags tagClient, workflo return &Service{store: store, authorizer: authorizer, tags: tags, workflow: workflow} } -func (s *Service) Create(ctx context.Context, sessionID, requestID, expectedHeadTaskID string) (*apiv1.Checkpoint, error) { +// Create executes new work in the request. The reservation owns its bounded +// claim from INSERT, so the normal path needs no subsequent acquisition query. +func (s *Service) Create(ctx context.Context, sessionID, requestID, expectedHeadTaskID string) (result *apiv1.Checkpoint, err error) { + ctx, cancel := context.WithTimeout(ctx, database.RuntimeOperationTimeout) + defer cancel() + executor := uuid.New() + id, err := uuid.NewV7() + if err != nil { + return nil, serviceerrors.NewInternal("Failed to generate checkpoint identifier", err) + } + op, err := s.reserveCreate(ctx, sessionID, requestID, expectedHeadTaskID, id.String(), executor) + if err != nil { + return nil, err + } + if op.Checkpoint.State != apiv1.CheckpointState_CHECKPOINT_STATE_CREATING { + return createOutcome(op.Checkpoint) + } + if op.Checkpoint.Id != id.String() { + result, err = s.claimAndExecute(ctx, op) + } else { + result, err = s.execute(ctx, op, executor, true) + } + if err == nil && result == nil { + return nil, serviceerrors.NewNotFound("Checkpoint no longer exists", nil) + } + return result, err +} + +func (s *Service) reserveCreate(ctx context.Context, sessionID, requestID, expectedHeadTaskID, id string, executor uuid.UUID) (*database.CheckpointOperation, error) { sessionID, err := validateCreate(sessionID, requestID) if err != nil { return nil, err @@ -90,22 +123,14 @@ func (s *Service) Create(ctx context.Context, sessionID, requestID, expectedHead if _, err := s.authorize(ctx, auth.VerbGet, auth.Resource{Type: "Session", Name: sessionID}); err != nil { return nil, err } - key := fmt.Sprintf("%q/%q/%q/%q", userID, sessionID, requestID, expectedHeadTaskID) - result, err, _ := s.creates.Do(key, func() (any, error) { - return s.create(ctx, userID, sessionID, requestID, expectedHeadTaskID) - }) - if err != nil { - return nil, err - } - return result.(*apiv1.Checkpoint), nil -} - -func (s *Service) create(ctx context.Context, userID, sessionID, requestID, expectedHeadTaskID string) (*apiv1.Checkpoint, error) { - id, err := uuid.NewV7() - if err != nil { - return nil, serviceerrors.NewInternal("Failed to generate checkpoint identifier", err) + checkpoint, snapshot, err := s.store.ReserveSessionCheckpoint(ctx, &apiv1.Checkpoint{Id: id, SessionId: sessionID, HeadTaskId: expectedHeadTaskID}, userID, requestID, executor) + if err != nil && !definiteAcquisitionError(err) { + // The INSERT may still be committing after its reply was lost. Its + // source Session lock makes the new row visible before retiring the claim. + err = errors.Join(err, releaseClaim(ctx, id, func(releaseCtx context.Context) error { + return s.store.ReleaseCheckpointReservation(releaseCtx, sessionID, id, executor) + })) } - checkpoint, snapshot, err := s.store.ReserveSessionCheckpoint(ctx, &apiv1.Checkpoint{Id: id.String(), SessionId: sessionID, HeadTaskId: expectedHeadTaskID}, userID, requestID) if errors.Is(err, database.ErrCheckpointAdvanced) || errors.Is(err, database.ErrSnapshotPending) { reason := "KAGENT_CHECKPOINT_CONVERSATION_ADVANCED" if errors.Is(err, database.ErrSnapshotPending) { @@ -129,64 +154,75 @@ func (s *Service) create(ctx context.Context, userID, sessionID, requestID, expe if err != nil { return nil, serviceerrors.NewInternal("Failed to reserve checkpoint", err) } - if checkpoint.State != apiv1.CheckpointState_CHECKPOINT_STATE_CREATING { - return checkpoint, nil - } - - tag, err := s.ensureTag(ctx, checkpoint, snapshot) - if err != nil { - cleanupErr := s.tags.DeleteTag(ctx, snapshot.Atespace, tagName(checkpoint.GetId()), "") - if cleanupErr == nil || status.Code(cleanupErr) == codes.NotFound { - _, finalizeErr := s.store.FinalizeSessionCheckpoint(ctx, checkpoint.GetId(), "", "", err.Error()) - err = errors.Join(err, finalizeErr) - } else { - err = errors.Join(err, fmt.Errorf("cleanup checkpoint tag: %w", cleanupErr)) - } - return nil, serviceerrors.NewUnavailable("Failed to retain checkpoint snapshot", err) - } - checkpoint, err = s.store.FinalizeSessionCheckpoint(ctx, checkpoint.GetId(), tag.GetMetadata().GetUid(), substrate.SnapshotURI(tag.GetStatus().GetSnapshot()), "") - if err != nil { - return nil, serviceerrors.NewInternal("Failed to publish checkpoint", err) - } - return checkpoint, nil + return &database.CheckpointOperation{Checkpoint: checkpoint, Snapshot: snapshot, Creator: userID}, nil } -// The CREATING reservation blocks task admission and lifecycle changes while -// CreateTag copies the Actor's current snapshot. Verify both sides of the copy -// because ate-api also permits Actors to be changed outside kagent. -func (s *Service) ensureTag(ctx context.Context, checkpoint *apiv1.Checkpoint, reference *database.SessionTaskSnapshot) (*ateapipb.Tag, error) { - actorName := substrate.ActorName(checkpoint.GetSessionId()) - actor, err := s.verifySnapshot(ctx, actorName, reference) - if err != nil { - return nil, err - } - name := tagName(checkpoint.GetId()) +var errTagCopyFailed = errors.New("snapshot tag copy cannot be published") + +// ensureTag synchronizes with any prior backend copy and verifies both sides +// of the copy because ate-api also permits Actors to change outside kagent. +// The CREATING reservation keeps source admission closed until finalization. +func (s *Service) ensureTag(ctx context.Context, checkpoint *apiv1.Checkpoint, reference *database.SessionTaskSnapshot, firstAttempt bool) (*ateapipb.Tag, error) { + actorName := substrate.ActorName(checkpoint.SessionId) + actor, beforeErr := s.verifySnapshot(ctx, actorName, reference) + if beforeErr != nil && status.Code(beforeErr) != codes.FailedPrecondition && status.Code(beforeErr) != codes.NotFound { + return nil, beforeErr + } + if beforeErr != nil && firstAttempt { + // No copy has been requested for this newly reserved checkpoint. + return nil, errors.Join(errTagCopyFailed, beforeErr) + } + name := tagName(checkpoint.Id) + // Recovery may have an earlier copy: synchronize through the Actor/Tag leases + // before interpreting an existing pending tag. A fresh invalid source never + // starts a copy; completed tags still require full identity verification. tag, err := s.tags.CreateTag(ctx, reference.Atespace, name, actorName) if err != nil { - tag, err = s.tags.GetTag(ctx, reference.Atespace, name) - if err != nil { + switch status.Code(err) { + case codes.AlreadyExists, codes.FailedPrecondition, codes.NotFound: + createErr := err + tag, err = s.tags.GetTag(ctx, reference.Atespace, name) + if status.Code(err) == codes.NotFound && status.Code(createErr) != codes.AlreadyExists { + return nil, errors.Join(errTagCopyFailed, createErr) + } + if err != nil { + return nil, fmt.Errorf("read snapshot tag after create: %w", err) + } + default: + // A canceled or unavailable RPC may still own the copy. Keep its + // reservation and retry through the same backend serialization barrier. return nil, fmt.Errorf("create snapshot tag: %w", err) } } - metadata, source, snapshot := tag.GetMetadata(), tag.GetSourceActor(), tag.GetStatus().GetSnapshot() - if metadata.GetAtespace() != reference.Atespace || metadata.GetName() != name || metadata.GetUid() == "" || - source.GetAtespace() != reference.Atespace || source.GetName() != actorName || + if beforeErr != nil { + return tag, errors.Join(errTagCopyFailed, beforeErr) + } + snapshot := tag.GetStatus().GetSnapshot() + if !validTagSource(tag, checkpoint, reference) || tag.GetStatus().GetActorTemplateUid() != substrate.LatestActorSnapshot(actor).GetActorTemplateUid() || - substrate.SnapshotURI(snapshot) == "" || - snapshot.GetFidelity().String() != reference.Fidelity || - tag.GetScope() != ateapipb.TagScope_TAG_SCOPE_ATESPACE { - return nil, fmt.Errorf("snapshot tag %s/%s returned invalid identity", reference.Atespace, name) + substrate.SnapshotURI(snapshot) == "" || snapshot.GetFidelity().String() != reference.Fidelity { + return tag, fmt.Errorf("snapshot tag returned invalid identity or incomplete copy: %w", errTagCopyFailed) } verified, err := s.verifySnapshot(ctx, actorName, reference) if err != nil { + if status.Code(err) == codes.FailedPrecondition || status.Code(err) == codes.NotFound { + return tag, errors.Join(errTagCopyFailed, err) + } return nil, err } if verified.GetMetadata().GetUid() != actor.GetMetadata().GetUid() { - return nil, fmt.Errorf("checkpoint Actor identity changed while copying snapshot") + return tag, fmt.Errorf("checkpoint Actor changed while copying: %w", errTagCopyFailed) } return tag, nil } +func validTagSource(tag *ateapipb.Tag, checkpoint *apiv1.Checkpoint, snapshot *database.SessionTaskSnapshot) bool { + metadata, source := tag.GetMetadata(), tag.GetSourceActor() + return metadata.GetAtespace() == snapshot.Atespace && metadata.GetName() == tagName(checkpoint.Id) && metadata.GetUid() != "" && + source.GetAtespace() == snapshot.Atespace && source.GetName() == substrate.ActorName(checkpoint.SessionId) && + tag.GetScope() == ateapipb.TagScope_TAG_SCOPE_ATESPACE +} + func (s *Service) verifySnapshot(ctx context.Context, actorName string, reference *database.SessionTaskSnapshot) (*ateapipb.Actor, error) { actor, err := s.tags.GetActor(ctx, reference.Atespace, actorName) if err != nil { @@ -196,10 +232,10 @@ func (s *Service) verifySnapshot(ctx context.Context, actorName string, referenc if metadata.GetAtespace() != reference.Atespace || metadata.GetName() != actorName || metadata.GetUid() == "" || actor.GetStatus().GetState() != ateapipb.ActorState_ACTOR_STATE_SUSPENDED || reference.URI == "" || substrate.SnapshotURI(snapshot) != reference.URI || snapshot.GetActorTemplateUid() == "" { - return nil, fmt.Errorf("checkpoint Actor %s/%s snapshot changed", reference.Atespace, actorName) + return nil, status.Errorf(codes.FailedPrecondition, "checkpoint Actor %s/%s snapshot changed", reference.Atespace, actorName) } if fidelity := snapshot.GetFidelity().String(); fidelity != reference.Fidelity { - return nil, fmt.Errorf("checkpoint Actor %s/%s snapshot fidelity changed", reference.Atespace, actorName) + return nil, status.Errorf(codes.FailedPrecondition, "checkpoint Actor %s/%s snapshot fidelity changed", reference.Atespace, actorName) } return actor, nil } @@ -254,8 +290,10 @@ func (s *Service) List(ctx context.Context, request ListRequest) (ListResult, er return result, nil } -func (s *Service) Delete(ctx context.Context, checkpointID string) error { - checkpointID, err := validateIdentity(checkpointID) +func (s *Service) Delete(ctx context.Context, checkpointID string) (err error) { + ctx, cancel := context.WithTimeout(ctx, database.RuntimeOperationTimeout) + defer cancel() + checkpointID, err = validateIdentity(checkpointID) if err != nil { return err } @@ -263,32 +301,65 @@ func (s *Service) Delete(ctx context.Context, checkpointID string) error { if err != nil { return err } - snapshot, tagUID, err := s.store.BeginDeleteSessionCheckpoint(ctx, checkpointID, userID) - if errors.Is(err, database.ErrNotFound) { - return serviceerrors.NewNotFound("Checkpoint not found", err) - } + executor := uuid.New() + op, err := s.store.BeginDeleteSessionCheckpoint(ctx, checkpointID, userID, executor) if err != nil { + if errors.Is(err, database.ErrConflict) { + _, err = s.claimAndExecute(ctx, &database.CheckpointOperation{Checkpoint: &apiv1.Checkpoint{Id: checkpointID}, Creator: userID}) + return err + } + if errors.Is(err, database.ErrNotFound) { + return serviceerrors.NewNotFound("Checkpoint not found", err) + } + if !definiteAcquisitionError(err) { + err = errors.Join(err, s.release(ctx, checkpointID, executor)) + } return serviceerrors.NewInternal("Failed to begin checkpoint deletion", err) } + _, err = s.execute(ctx, op, executor, false) + return err +} + +func (s *Service) finishDelete(ctx context.Context, op *database.CheckpointOperation, executor uuid.UUID) error { + checkpointID, userID := op.Checkpoint.Id, op.Creator + if err := s.deleteTag(ctx, op); err != nil { + return err + } + if err := s.store.DeleteClaimedSessionCheckpoint(ctx, checkpointID, userID, executor); err != nil { + if errors.Is(err, database.ErrConflict) { + _, readErr := s.store.GetCheckpointOperation(ctx, checkpointID, userID) + if errors.Is(readErr, database.ErrNotFound) { + return nil + } + if readErr != nil { + return serviceerrors.NewInternal("Failed to observe checkpoint deletion", readErr) + } + return serviceerrors.NewAborted("Checkpoint deletion is still pending; retry the same request", nil) + } + return serviceerrors.NewInternal("Failed to delete checkpoint", err) + } + return nil +} + +func (s *Service) deleteTag(ctx context.Context, op *database.CheckpointOperation) error { + checkpointID := op.Checkpoint.Id + snapshot, tagUID := op.Snapshot, op.TagUID tag, err := s.tags.GetTag(ctx, snapshot.Atespace, tagName(checkpointID)) if err != nil && status.Code(err) != codes.NotFound { - return serviceerrors.NewUnavailable("Failed to get checkpoint snapshot tag", err) + return publicOperationError(fmt.Errorf("get checkpoint snapshot tag: %w", err)) } if err == nil { - // Failed creation records no tag UID after compensation. An unexpected tag - // at that name is not ours to delete, even if its snapshot URI matches. + // The retained UID is authoritative for failed partial copies too: + // their storage may still be PENDING and has no completed snapshot URI. if tagUID == "" || tag.GetMetadata().GetUid() != tagUID || tag.GetMetadata().GetAtespace() != snapshot.Atespace || tag.GetMetadata().GetName() != tagName(checkpointID) || - substrate.SnapshotURI(tag.GetStatus().GetSnapshot()) != snapshot.URI { + (op.Checkpoint.State != apiv1.CheckpointState_CHECKPOINT_STATE_FAILED && op.Checkpoint.Failure == nil && substrate.SnapshotURI(tag.GetStatus().GetSnapshot()) != snapshot.URI) { return serviceerrors.NewFailedPrecondition("Checkpoint snapshot tag identity changed", nil) } if err := s.tags.DeleteTag(ctx, snapshot.Atespace, tagName(checkpointID), tagUID); err != nil && status.Code(err) != codes.NotFound { - return serviceerrors.NewUnavailable("Failed to delete checkpoint snapshot tag", err) + return publicOperationError(fmt.Errorf("delete checkpoint snapshot tag: %w", err)) } } - if err := s.store.DeleteSessionCheckpoint(ctx, checkpointID, userID); err != nil { - return serviceerrors.NewInternal("Failed to delete checkpoint", err) - } return nil } @@ -443,3 +514,32 @@ func decodePageToken(token string) (string, error) { } return validateIdentity(string(value)) } + +func createOutcome(checkpoint *apiv1.Checkpoint) (*apiv1.Checkpoint, error) { + switch checkpoint.State { + case apiv1.CheckpointState_CHECKPOINT_STATE_READY: + return checkpoint, nil + case apiv1.CheckpointState_CHECKPOINT_STATE_FAILED: + return nil, status.Error(codes.FailedPrecondition, "Checkpoint creation failed") + default: + return nil, serviceerrors.NewAborted("Checkpoint operation is still pending; retry the same request", nil) + } +} + +// Semantic rejections cannot have committed a reservation or execution claim. +func definiteAcquisitionError(err error) bool { + return errors.Is(err, database.ErrFailedPrecondition) || errors.Is(err, database.ErrConflict) || + errors.Is(err, database.ErrNotFound) || errors.Is(err, database.ErrIdempotencyConflict) +} + +// Backend status text/details can contain private storage configuration. Keep +// current status codes while exposing only a service-owned public message. +func publicOperationError(err error) error { + if err == nil || serviceerrors.CodeOf(err) != "" || errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return err + } + if backend, ok := status.FromError(err); ok { + return status.Error(backend.Code(), "Checkpoint operation could not complete") + } + return serviceerrors.NewInternal("Checkpoint operation could not complete", err) +} diff --git a/go/core/internal/service/checkpoint/service_test.go b/go/core/internal/service/checkpoint/service_test.go index f3c0b71f3b..f7bd8f1f52 100644 --- a/go/core/internal/service/checkpoint/service_test.go +++ b/go/core/internal/service/checkpoint/service_test.go @@ -5,11 +5,11 @@ import ( "encoding/base64" "errors" "fmt" - "runtime" "strings" "sync/atomic" "testing" + "github.com/google/uuid" "github.com/stretchr/testify/require" ateapipb "go.ate.dev/substrate/api/ateapi/v1" "google.golang.org/genproto/googleapis/rpc/errdetails" @@ -59,7 +59,7 @@ func (s *testStore) GetSession(_ context.Context, id, owner string) (*apiv1.Sess return &apiv1.Session{Id: id, Creator: owner, Agent: &apiv1.AgentProvenance{Ref: &apiv1.ResourceReference{Namespace: "team-a", Name: "assistant"}}}, nil } -func (s *testStore) ReserveSessionCheckpoint(_ context.Context, checkpoint *apiv1.Checkpoint, owner, _ string) (*apiv1.Checkpoint, *database.SessionTaskSnapshot, error) { +func (s *testStore) ReserveSessionCheckpoint(_ context.Context, checkpoint *apiv1.Checkpoint, owner, _ string, _ uuid.UUID) (*apiv1.Checkpoint, *database.SessionTaskSnapshot, error) { s.reservations++ if s.reserveErr != nil { return nil, nil, s.reserveErr @@ -83,12 +83,17 @@ func (s *testStore) ReserveSessionCheckpoint(_ context.Context, checkpoint *apiv return checkpoint, s.snapshot, nil } -func (s *testStore) FinalizeSessionCheckpoint(_ context.Context, _ string, tagUID, snapshotURI, failure string) (*apiv1.Checkpoint, error) { +func (s *testStore) FinalizeClaimedSessionCheckpoint(_ context.Context, _ *apiv1.Checkpoint, _ uuid.UUID, tagUID, snapshotURI, failure string) (*apiv1.Checkpoint, error) { if s.finalizeErr != nil { return nil, s.finalizeErr } if failure != "" { s.prepared.State, s.failed = apiv1.CheckpointState_CHECKPOINT_STATE_FAILED, failure + s.prepared.Failure = &apiv1.Failure{Reason: "SnapshotTagFailed", Message: failure} + s.tagUID = tagUID + if snapshotURI != "" { + s.snapshot.URI = snapshotURI + } } else { s.prepared.State, s.tagUID = apiv1.CheckpointState_CHECKPOINT_STATE_READY, tagUID s.snapshot.URI = snapshotURI @@ -123,15 +128,15 @@ func (*testStore) ListSessionCheckpoints(context.Context, string, string, string return nil, nil } -func (s *testStore) BeginDeleteSessionCheckpoint(context.Context, string, string) (*database.SessionTaskSnapshot, string, error) { +func (s *testStore) BeginDeleteSessionCheckpoint(ctx context.Context, id, owner string, _ uuid.UUID) (*database.CheckpointOperation, error) { if s.prepared == nil { - return nil, "", database.ErrNotFound + return nil, database.ErrNotFound } s.prepared.State = apiv1.CheckpointState_CHECKPOINT_STATE_DELETING - return s.snapshot, s.tagUID, nil + return s.GetCheckpointOperation(ctx, id, owner) } -func (s *testStore) DeleteSessionCheckpoint(context.Context, string, string) error { +func (s *testStore) DeleteClaimedSessionCheckpoint(context.Context, string, string, uuid.UUID) error { s.deleted = true return nil } @@ -176,6 +181,7 @@ type testTags struct { getErr error createErr error deleteErr error + deleteResponseErr error actorErr error mutateTag func(*ateapipb.Tag) beforeDelete func() @@ -241,7 +247,7 @@ func (t *testTags) DeleteTag(_ context.Context, _, _, uid string) error { return status.Error(codes.Aborted, "tag UID changed") } t.created = nil - return nil + return t.deleteResponseErr } func TestCreatePreservesStoreConflictReason(t *testing.T) { @@ -315,7 +321,7 @@ func TestCreateTagsRecordedSnapshotBoundary(t *testing.T) { } } -func TestCreateCleansTagBeforeFailing(t *testing.T) { +func TestCreateRetainsFailedTagUntilExplicitDelete(t *testing.T) { for _, test := range []struct { name string snapshotURIAfterCreate string @@ -332,8 +338,11 @@ func TestCreateCleansTagBeforeFailing(t *testing.T) { _, err := service.Create(ctx, "018f47a2-4efb-7c21-a848-123456789abc", "request-1", "task-1") require.Error(t, err) - require.Equal(t, 1, tags.deleteCalls) + require.Zero(t, tags.deleteCalls) + require.Equal(t, "tag-uid", store.tagUID) require.NotEmpty(t, store.failed) + require.NoError(t, service.Delete(ctx, store.prepared.Id)) + require.Equal(t, 1, tags.deleteCalls) }) } } @@ -358,7 +367,7 @@ func TestDeleteHidesCheckpointBeforeDeletingTag(t *testing.T) { } tags.getErr = nil tags.deleteErr = status.Error(codes.Unavailable, "tag delete unavailable") - require.Equal(t, serviceerrors.CodeUnavailable, serviceerrors.CodeOf(service.Delete(ctx, checkpoint.GetId()))) + require.Equal(t, codes.Unavailable, status.Code(service.Delete(ctx, checkpoint.GetId()))) require.False(t, store.deleted, "a failed tag delete must retain the database reference") require.NotNil(t, tags.created) tags.deleteErr = nil @@ -380,7 +389,7 @@ func TestDeleteFailedCheckpointWithMissingTag(t *testing.T) { ctx := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) tags.getErr = status.Error(codes.Unavailable, "tag lookup unavailable") - require.Equal(t, serviceerrors.CodeUnavailable, serviceerrors.CodeOf(service.Delete(ctx, checkpoint.GetId()))) + require.Equal(t, codes.Unavailable, status.Code(service.Delete(ctx, checkpoint.GetId()))) require.False(t, store.deleted) require.Zero(t, tags.deleteCalls) require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_DELETING, checkpoint.State) @@ -447,7 +456,7 @@ func TestForkCreatesSessionFromRetainedCheckpointAfterSourceDeletion(t *testing. func TestCreateRetriesRetainedTagAfterPublishFailure(t *testing.T) { store := &testStore{finalizeErr: errors.New("database unavailable")} - tags := &testTags{snapshotURI: "s3://snapshots/snapshot-1", createErr: status.Error(codes.DeadlineExceeded, "response lost")} + tags := &testTags{snapshotURI: "s3://snapshots/snapshot-1"} service := NewService(store, testAuthorizer{}, tags, nil) ctx := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) _, err := service.Create(ctx, "018f47a2-4efb-7c21-a848-123456789abc", "request-1", "task-1") @@ -463,7 +472,7 @@ func TestCreateRetriesRetainedTagAfterPublishFailure(t *testing.T) { require.Equal(t, "s3://tags/checkpoint", store.snapshot.URI) } -func TestCreateRejectsInvalidTagAndKeepsCleanupRetryable(t *testing.T) { +func TestFailedCheckpointKeepsTagUntilDeletionSucceeds(t *testing.T) { cases := []struct { name string mutate func(*ateapipb.Tag) @@ -472,7 +481,7 @@ func TestCreateRejectsInvalidTagAndKeepsCleanupRetryable(t *testing.T) { {"upload pending", func(tag *ateapipb.Tag) { tag.Status.Snapshot.DurableSnapshot.State = ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_PENDING }}, - {"wrong source", func(tag *ateapipb.Tag) { tag.SourceActor.Name = "other" }}, + {"wrong template", func(tag *ateapipb.Tag) { tag.Status.ActorTemplateUid = "other" }}, {"missing template", func(tag *ateapipb.Tag) { tag.Status.ActorTemplateUid = "" }}, {"wrong scope", func(tag *ateapipb.Tag) { @@ -487,11 +496,14 @@ func TestCreateRejectsInvalidTagAndKeepsCleanupRetryable(t *testing.T) { ctx := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) _, err := service.Create(ctx, "018f47a2-4efb-7c21-a848-123456789abc", "request-1", "task-1") require.Error(t, err) - require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_CREATING, store.prepared.State) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_FAILED, store.prepared.State) + require.Equal(t, "tag-uid", store.tagUID) + require.Error(t, service.Delete(ctx, store.prepared.Id)) tags.deleteErr = nil _, err = service.Create(ctx, store.prepared.SessionId, "request-1", "task-1") require.Error(t, err) - require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_FAILED, store.prepared.State) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_DELETING, store.prepared.State) + require.NoError(t, service.Delete(ctx, store.prepared.Id)) require.Nil(t, tags.created) }) } @@ -519,27 +531,17 @@ func TestForkRejectsReplacedTag(t *testing.T) { } } -func TestConcurrentCreateRetainsOneTag(t *testing.T) { +func TestRepeatedCreateDoesNotCopyReadyTag(t *testing.T) { store := &testStore{} - tags := &testTags{snapshotURI: "s3://snapshots/snapshot-1", mutateTag: func(*ateapipb.Tag) { runtime.Gosched() }} + tags := &testTags{snapshotURI: "s3://snapshots/snapshot-1"} service := NewService(store, testAuthorizer{}, tags, nil) ctx := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) - start := make(chan struct{}) - results := make(chan error, 32) - for range cap(results) { - go func() { - <-start - _, err := service.Create(ctx, "018f47a2-4efb-7c21-a848-123456789abc", "request-1", "task-1") - results <- err - }() - } - close(start) - for range cap(results) { - require.NoError(t, <-results) + for range 32 { + _, err := service.Create(ctx, "018f47a2-4efb-7c21-a848-123456789abc", "request-1", "task-1") + require.NoError(t, err) } require.Equal(t, 1, tags.createCalls) - require.Equal(t, 0, tags.deleteCalls) - require.Equal(t, "s3://tags/checkpoint", store.snapshot.URI) + require.Zero(t, tags.deleteCalls) } func TestDeletePreservesTagReplacedAfterLookup(t *testing.T) { @@ -806,9 +808,9 @@ type canonicalCheckpointStore struct { afterID string } -func (c *canonicalCheckpointStore) ReserveSessionCheckpoint(ctx context.Context, value *apiv1.Checkpoint, owner, requestID string) (*apiv1.Checkpoint, *database.SessionTaskSnapshot, error) { +func (c *canonicalCheckpointStore) ReserveSessionCheckpoint(ctx context.Context, value *apiv1.Checkpoint, owner, requestID string, executor uuid.UUID) (*apiv1.Checkpoint, *database.SessionTaskSnapshot, error) { c.sessionID = value.GetSessionId() - return c.testStore.ReserveSessionCheckpoint(ctx, value, owner, requestID) + return c.testStore.ReserveSessionCheckpoint(ctx, value, owner, requestID, executor) } func (c *canonicalCheckpointStore) GetSessionCheckpoint(ctx context.Context, id, owner string) (*apiv1.Checkpoint, error) { @@ -853,3 +855,47 @@ func TestForkCanonicalizesSavedSourceIdentity(t *testing.T) { require.Equal(t, 1, store.forkSourceReads) require.Nil(t, store.forked) } + +func (s *testStore) GetCheckpointOperation(ctx context.Context, id, owner string) (*database.CheckpointOperation, error) { + if s.deleted { + return nil, database.ErrNotFound + } + cp, err := s.GetSessionCheckpoint(ctx, id, owner) + if err != nil { + return nil, err + } + return &database.CheckpointOperation{Checkpoint: cp, Creator: owner, Snapshot: s.snapshot, TagUID: s.tagUID}, nil +} + +func (*testStore) ListPendingCheckpoints(context.Context, string, int) ([]*database.CheckpointOperation, error) { + return nil, nil +} + +func TestCreationRefusesUnrelatedTag(t *testing.T) { + store := &testStore{} + tags := &testTags{snapshotURI: "s3://snapshots/snapshot-1", mutateTag: func(tag *ateapipb.Tag) { tag.SourceActor.Name = "unrelated" }} + service := NewService(store, testAuthorizer{}, tags, nil) + _, err := service.Create(auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}), "018f47a2-4efb-7c21-a848-123456789abc", "request", "task-1") + require.Error(t, err) + require.Zero(t, tags.deleteCalls) + require.NotNil(t, tags.created) + require.Nil(t, store.prepared.Failure) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_CREATING, store.prepared.State) +} + +func (s *testStore) ClaimCheckpointOperation(ctx context.Context, id, owner string, _ uuid.UUID) (*database.CheckpointOperation, error) { + return s.GetCheckpointOperation(ctx, id, owner) +} + +func (*testStore) ReleaseCheckpointOperation(context.Context, string, uuid.UUID) error { + return nil +} + +func (*testStore) ReleaseCheckpointReservation(context.Context, string, string, uuid.UUID) error { + return nil +} + +func (s *testStore) ClearFailedCheckpointTag(context.Context, string, string, string, uuid.UUID) error { + s.tagUID = "" + return nil +} diff --git a/go/core/internal/service/session/checkpoint_helpers_test.go b/go/core/internal/service/session/checkpoint_helpers_test.go new file mode 100644 index 0000000000..8b20688dfa --- /dev/null +++ b/go/core/internal/service/session/checkpoint_helpers_test.go @@ -0,0 +1,48 @@ +package session + +import ( + "context" + "errors" + + "github.com/google/uuid" + + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" + "github.com/kagent-dev/kagent/go/core/internal/database" +) + +func reserveCheckpoint(ctx context.Context, client *lifecycleTestStore, checkpoint *apiv1.Checkpoint, owner, request string) (*apiv1.Checkpoint, *database.SessionTaskSnapshot, error) { + ctx, cancel := context.WithTimeout(ctx, database.RuntimeOperationTimeout) + defer cancel() + executor := uuid.New() + cp, snapshot, err := client.ReserveSessionCheckpoint(ctx, checkpoint, owner, request, executor) + err = errors.Join(err, client.ReleaseCheckpointOperation(context.WithoutCancel(ctx), checkpoint.Id, executor)) + return cp, snapshot, err +} + +func finalizeCheckpoint(ctx context.Context, client *lifecycleTestStore, id, uid, uri, failure string) (*apiv1.Checkpoint, error) { + ctx, cancel := context.WithTimeout(ctx, database.RuntimeOperationTimeout) + defer cancel() + executor := uuid.New() + op, err := client.ClaimCheckpointOperation(ctx, id, "alice", executor) + if err != nil { + return nil, err + } + return client.FinalizeClaimedSessionCheckpoint(ctx, op.Checkpoint, executor, uid, uri, failure) +} + +func beginDeleteCheckpoint(ctx context.Context, client *lifecycleTestStore, id, owner string) (*database.SessionTaskSnapshot, string, error) { + ctx, cancel := context.WithTimeout(ctx, database.RuntimeOperationTimeout) + defer cancel() + executor := uuid.New() + op, err := client.BeginDeleteSessionCheckpoint(ctx, id, owner, executor) + if errors.Is(err, database.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 +} diff --git a/go/core/internal/service/session/quiescence_test.go b/go/core/internal/service/session/quiescence_test.go index 99e8852d72..bc7a7562e6 100644 --- a/go/core/internal/service/session/quiescence_test.go +++ b/go/core/internal/service/session/quiescence_test.go @@ -182,7 +182,7 @@ func TestIdleLifecycleDoesNotOwnTaskPublication(t *testing.T) { _, err = store.CreateRuntimeTask(t.Context(), session.Id, hash[:], next, "") require.ErrorIs(t, err, database.ErrFailedPrecondition) checkpoint := &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID)} - _, _, err = store.ReserveSessionCheckpoint(t.Context(), checkpoint, "alice", "checkpoint") + _, _, err = reserveCheckpoint(t.Context(), store, checkpoint, "alice", "checkpoint") require.ErrorIs(t, err, database.ErrFailedPrecondition) close(release) @@ -198,7 +198,7 @@ func TestIdleLifecycleDoesNotOwnTaskPublication(t *testing.T) { require.False(t, pending, "uncertain claims must not keep recovery retries active") } else { require.Eventually(t, func() bool { - _, snapshot, err := store.ReserveSessionCheckpoint(t.Context(), checkpoint, "alice", "checkpoint") + _, snapshot, err := reserveCheckpoint(t.Context(), store, checkpoint, "alice", "checkpoint") return err == nil && snapshot.URI == "s3://snapshots/snapshot-1" }, 5*time.Second, 10*time.Millisecond) require.EqualValues(t, 1, actors.mutations.Load(), "database retries must not suspend the actor again") diff --git a/go/core/internal/service/session/quiescence_wake_test.go b/go/core/internal/service/session/quiescence_wake_test.go index 691235c5d1..9fb347d6b0 100644 --- a/go/core/internal/service/session/quiescence_wake_test.go +++ b/go/core/internal/service/session/quiescence_wake_test.go @@ -191,7 +191,7 @@ func TestQuiescenceDrainsConcurrentSettlementsAcrossReplicas(t *testing.T) { require.Equal(t, ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES, substrate.LatestActorSnapshot(actor).GetFidelity()) require.NotEmpty(t, substrate.SnapshotURI(substrate.LatestActorSnapshot(actor))) if states[i%len(states)].Terminal() { - _, snapshot, err := store.ReserveSessionCheckpoint(t.Context(), &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: requests[i].TaskId}, "alice", uuid.NewString()) + _, snapshot, err := reserveCheckpoint(t.Context(), store, &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: requests[i].TaskId}, "alice", uuid.NewString()) require.NoError(t, err) require.Equal(t, "s3://snapshots/snapshot-1", snapshot.URI) } diff --git a/go/core/internal/service/session/recovery_test.go b/go/core/internal/service/session/recovery_test.go index 2f0d05dcc5..1b166b0614 100644 --- a/go/core/internal/service/session/recovery_test.go +++ b/go/core/internal/service/session/recovery_test.go @@ -333,7 +333,7 @@ func TestExplicitRevertRespectsCheckpointCreation(t *testing.T) { require.NoError(t, store.FinishSessionQuiescence(t.Context(), work, &database.SessionTaskSnapshot{ Atespace: "team-a", URI: "s3://snapshots/completed", Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES.String(), })) - _, _, err = store.ReserveSessionCheckpoint(t.Context(), &apiv1.Checkpoint{ + _, _, err = reserveCheckpoint(t.Context(), store, &apiv1.Checkpoint{ Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID), }, session.Creator, uuid.NewString()) require.NoError(t, err) @@ -421,7 +421,7 @@ func TestExplicitRevertCannotCheckpointDiscardedTerminalState(t *testing.T) { require.NoError(t, err) _, err = store.ClaimSessionQuiescence(t.Context()) require.ErrorIs(t, err, database.ErrNotFound, "the reverted snapshot cannot represent discarded terminal state") - _, _, err = store.ReserveSessionCheckpoint(t.Context(), &apiv1.Checkpoint{ + _, _, err = reserveCheckpoint(t.Context(), store, &apiv1.Checkpoint{ Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID), }, session.Creator, uuid.NewString()) require.ErrorIs(t, err, database.ErrSnapshotPending) diff --git a/go/core/internal/service/session/workflow_retry_test.go b/go/core/internal/service/session/workflow_retry_test.go index f3c105e9b0..a315000080 100644 --- a/go/core/internal/service/session/workflow_retry_test.go +++ b/go/core/internal/service/session/workflow_retry_test.go @@ -124,7 +124,7 @@ func TestLifecycleClientRetriesAmbiguousMutation(t *testing.T) { require.ErrorIs(t, err, database.ErrConflict, "Delete must not erase uncertain work") } if checkpointID != "" { - _, _, err = store.BeginDeleteSessionCheckpoint(t.Context(), checkpointID, session.Creator) + _, _, err = beginDeleteCheckpoint(t.Context(), store, checkpointID, session.Creator) require.ErrorIs(t, err, database.ErrNotFound, "uncertain fork must retain its checkpoint pin") checkpoint, err := store.GetSessionCheckpoint(t.Context(), checkpointID, session.Creator) require.NoError(t, err) diff --git a/go/core/internal/service/session/workflow_test.go b/go/core/internal/service/session/workflow_test.go index af6651bf92..aae17aa85d 100644 --- a/go/core/internal/service/session/workflow_test.go +++ b/go/core/internal/service/session/workflow_test.go @@ -347,9 +347,9 @@ func lifecycleForkFixture(t *testing.T, store *lifecycleTestStore, actors *lifec require.NoError(t, err) require.NoError(t, store.FinishSessionQuiescence(t.Context(), boundary, &database.SessionTaskSnapshot{Atespace: "team-a", URI: "s3://snapshots/source", Fidelity: "SNAPSHOT_FIDELITY_VOLUMES"})) - checkpoint, _, err := store.ReserveSessionCheckpoint(t.Context(), &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: source.Id, HeadTaskId: string(task.ID)}, source.Creator, uuid.NewString()) + checkpoint, _, err := reserveCheckpoint(t.Context(), store, &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: source.Id, HeadTaskId: string(task.ID)}, source.Creator, uuid.NewString()) require.NoError(t, err) - _, err = store.FinalizeSessionCheckpoint(t.Context(), checkpoint.Id, "tag-uid", "s3://snapshots/snapshot-1", "") + _, err = finalizeCheckpoint(t.Context(), store, checkpoint.Id, "tag-uid", "s3://snapshots/snapshot-1", "") require.NoError(t, err) requestID := uuid.NewString() fork, _, err := store.ForkSession(t.Context(), checkpoint.Id, source.Creator, requestID, uuid.NewString()) diff --git a/go/core/pkg/app/app.go b/go/core/pkg/app/app.go index d10cf3ba53..f9b829dc85 100644 --- a/go/core/pkg/app/app.go +++ b/go/core/pkg/app/app.go @@ -376,6 +376,9 @@ func Run(ctx context.Context, input *configv1.ControllerConfig, opts Options) er shareMaxTTL := cfg.GetSessions().GetShareMaxTtl().AsDuration() sessions := sessionsvc.NewService(store, authorizer, sessionWorkflow, sessionsvc.WithShareMaxTTL(shareMaxTTL)) checkpoints := checkpoint.NewService(store, authorizer, actors, sessionWorkflow) + if err := manager.Add(checkpoint.NewRecoveryWorker(checkpoints)); err != nil { + return fmt.Errorf("register checkpoint recovery worker: %w", err) + } gatewayDialer, err := a2agateway.NewRuntimeDialer( cfg.GetSubstrate(), authenticator, diff --git a/go/core/pkg/migrations/core/000004_checkpoint_recovery.sql b/go/core/pkg/migrations/core/000004_checkpoint_recovery.sql new file mode 100644 index 0000000000..aac6862cd7 --- /dev/null +++ b/go/core/pkg/migrations/core/000004_checkpoint_recovery.sql @@ -0,0 +1,13 @@ +-- +goose Up +ALTER TABLE session_checkpoint + ADD COLUMN executor_id uuid, + ADD COLUMN executor_expires_at timestamptz, + ADD CONSTRAINT session_checkpoint_executor_pair CHECK ( + (executor_id IS NULL) = (executor_expires_at IS NULL) + ); + +-- +goose Down +ALTER TABLE session_checkpoint + DROP CONSTRAINT session_checkpoint_executor_pair, + DROP COLUMN executor_expires_at, + DROP COLUMN executor_id; diff --git a/go/core/pkg/migrations/runner_test.go b/go/core/pkg/migrations/runner_test.go index 10fae8fbd3..2a082b2a02 100644 --- a/go/core/pkg/migrations/runner_test.go +++ b/go/core/pkg/migrations/runner_test.go @@ -320,7 +320,7 @@ func TestBuiltinMigrationsRoundTrip(t *testing.T) { for _, source := range sources { expected := []int64{0, 1} if source.Name == "core" { - expected = append(expected, 2, 3) + expected = append(expected, 2, 3, 4) } if versions := testVersions(t, dsn, source.TrackingTable); !slices.Equal(versions, expected) { t.Fatalf("%s versions = %v, want %v", source.Name, versions, expected) diff --git a/go/core/test/e2e/checkpoint_recovery_test.go b/go/core/test/e2e/checkpoint_recovery_test.go new file mode 100644 index 0000000000..dc57832b94 --- /dev/null +++ b/go/core/test/e2e/checkpoint_recovery_test.go @@ -0,0 +1,96 @@ +package e2e_test + +import ( + "context" + "strings" + "testing" + "time" + + a2apb "github.com/a2aproject/a2a-go/v2/a2apb/v1" + "github.com/google/uuid" + "github.com/stretchr/testify/require" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" + "google.golang.org/protobuf/proto" + "k8s.io/apimachinery/pkg/util/wait" + + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" +) + +// Sequential: pod replacement must not disrupt unrelated live scenarios. The +// interceptor drops a successful public response; Substrate still performs real +// snapshot copying. Interrupted backend phases are covered by the PostgreSQL +// recovery tests, not claimed as live packet-loss injection here. +func TestCheckpointReplayAfterControllerRestart(t *testing.T) { + prepared := 0 + forEachHarness(t, func(t *testing.T, harness testHarness) { + fixture := newInteractionFixture(t, harness, interactionTarget(t), startForkMemoryMock(t)) + _, _, task := fixture.send(t, "What is 2+2?") + request := &apiv1.CreateCheckpointRequest{SessionId: fixture.sessionID, RequestId: uuid.NewString(), ExpectedHeadTaskId: string(task.ID)} + var lost *apiv1.Checkpoint + conn, err := grpc.NewClient(interactionTarget(t), grpc.WithTransportCredentials(insecure.NewCredentials()), + grpc.WithUnaryInterceptor(func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoke grpc.UnaryInvoker, options ...grpc.CallOption) error { + err := invoke(ctx, method, req, reply, cc, options...) + if err == nil && method == apiv1.CheckpointService_CreateCheckpoint_FullMethodName { + lost = proto.CloneOf(reply.(*apiv1.CreateCheckpointResponse).Checkpoint) + return status.Error(codes.Unavailable, "successful checkpoint reply lost") + } + return err + })) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, conn.Close()) }) + require.NoError(t, waitForControllerAPI(t.Context(), conn)) + client := apiv1.NewCheckpointServiceClient(conn) + require.NoError(t, wait.PollUntilContextTimeout(fixture.ctx, 100*time.Millisecond, time.Minute, true, func(ctx context.Context) (bool, error) { + _, err := client.CreateCheckpoint(ctx, request) + if checkpointSnapshotPending(err) { + return false, nil + } + if lost != nil && status.Code(err) == codes.Unavailable { + return true, nil + } + return false, err + })) + require.NotNil(t, lost) + t.Cleanup(func() { + ctx, cancel := context.WithTimeout(metadata.AppendToOutgoingContext(context.Background(), "x-user-id", "e2e"), time.Minute) + defer cancel() + _, err := fixture.checkpoints.DeleteCheckpoint(ctx, &apiv1.DeleteCheckpointRequest{CheckpointId: lost.Id}) + if status.Code(err) != codes.NotFound { + require.NoError(t, err) + } + }) + prepared++ + // The parent restarts one controller for all harnesses before these + // continuations resume in parallel. + t.Parallel() + replacement := newControllerConn(t, interactionTarget(t)) + checkpoints := apiv1.NewCheckpointServiceClient(replacement) + replayed, err := checkpoints.CreateCheckpoint(fixture.ctx, request) + require.NoError(t, err) + require.Equal(t, lost.Id, replayed.Checkpoint.Id) + require.Equal(t, lost.HistorySequence, replayed.Checkpoint.HistorySequence) + listed, err := checkpoints.ListCheckpoints(fixture.ctx, &apiv1.ListCheckpointsRequest{SessionId: fixture.sessionID}) + require.NoError(t, err) + require.Len(t, listed.Checkpoints, 1, "retry must retain the original checkpoint") + forked, err := checkpoints.ForkSession(fixture.ctx, &apiv1.ForkSessionRequest{CheckpointId: lost.Id, RequestId: uuid.NewString()}) + require.NoError(t, err) + fork := forked.Session + t.Cleanup(func() { + ctx, cancel := context.WithTimeout(metadata.AppendToOutgoingContext(context.Background(), "x-user-id", "e2e"), time.Minute) + defer cancel() + require.NoError(t, deleteIdleSession(ctx, apiv1.NewSessionServiceClient(replacement), fork.Id)) + }) + forkFixture := *fixture + forkFixture.client = a2apb.NewA2AServiceClient(replacement) + forkFixture.sessionID, forkFixture.contextID = fork.Id, fork.ContextId + _, _, answer := forkFixture.send(t, "What was the answer before the checkpoint?") + require.True(t, strings.Contains(taskText(answer), "The answer is 4."), "fork did not restore the retained conversation: %s", taskText(answer)) + }) + if prepared > 0 && !t.Failed() { + restartTestController(t) + } +} diff --git a/go/core/test/e2e/connection_test.go b/go/core/test/e2e/connection_test.go index ad4a4fb7e2..c73d561c0d 100644 --- a/go/core/test/e2e/connection_test.go +++ b/go/core/test/e2e/connection_test.go @@ -16,6 +16,9 @@ import ( grpc_health_v1 "google.golang.org/grpc/health/grpc_health_v1" "google.golang.org/grpc/status" "google.golang.org/grpc/test/bufconn" + corev1 "k8s.io/api/core/v1" + "k8s.io/apimachinery/pkg/util/wait" + ctrlclient "sigs.k8s.io/controller-runtime/pkg/client" ) func newControllerConn(t *testing.T, target string) *grpc.ClientConn { @@ -121,3 +124,42 @@ func (s *readinessHealthServer) Check(context.Context, *grpc_health_v1.HealthChe s.calls.Add(1) return &grpc_health_v1.HealthCheckResponse{Status: s.health}, status.Error(s.code, "controller unavailable") } + +// restartTestController replaces the single test controller and waits for +// its public API. Call only from sequential tests with an endpoint that survives +// pod replacement. +func restartTestController(t *testing.T) { + t.Helper() + ctx, cancel := context.WithTimeout(t.Context(), 2*time.Minute) + defer cancel() + kube := interactionKubeClient(t) + pods := &corev1.PodList{} + require.NoError(t, kube.List(ctx, pods, ctrlclient.InNamespace("kagent"), ctrlclient.MatchingLabels{"app.kubernetes.io/component": "controller"})) + require.Len(t, pods.Items, 1, "restart test requires a single controller replica") + restartControllerPod(t, kube, pods.Items[0]) +} + +func restartControllerPod(t *testing.T, kube ctrlclient.Client, old corev1.Pod) { + t.Helper() + ctx, cancel := context.WithTimeout(t.Context(), 2*time.Minute) + defer cancel() + require.NoError(t, kube.Delete(ctx, &old)) + require.NoError(t, wait.PollUntilContextTimeout(ctx, time.Second, 90*time.Second, true, func(ctx context.Context) (bool, error) { + current := &corev1.PodList{} + if err := kube.List(ctx, current, ctrlclient.InNamespace("kagent"), ctrlclient.MatchingLabels{"app.kubernetes.io/component": "controller"}); err != nil { + return false, err + } + for _, pod := range current.Items { + if pod.UID == old.UID { + continue + } + for _, condition := range pod.Status.Conditions { + if condition.Type == corev1.PodReady && condition.Status == corev1.ConditionTrue { + return true, nil + } + } + } + return false, nil + }), "wait for replacement controller") + newControllerConn(t, interactionTarget(t)) +} diff --git a/go/core/test/e2e/interaction_test.go b/go/core/test/e2e/interaction_test.go index 8a77a0d16c..d259c6a97b 100644 --- a/go/core/test/e2e/interaction_test.go +++ b/go/core/test/e2e/interaction_test.go @@ -301,10 +301,8 @@ func createCheckpoint(t *testing.T, ctx context.Context, client apiv1.Checkpoint err := wait.PollUntilContextTimeout(ctx, 100*time.Millisecond, 2*time.Minute, true, func(ctx context.Context) (bool, error) { var err error result, err = client.CreateCheckpoint(ctx, request) - for _, detail := range status.Convert(err).Details() { - if info, ok := detail.(*errdetails.ErrorInfo); ok && info.Reason == "KAGENT_CHECKPOINT_SNAPSHOT_PENDING" { - return false, nil - } + if checkpointSnapshotPending(err) { + return false, nil } return err == nil, err }) @@ -314,6 +312,15 @@ func createCheckpoint(t *testing.T, ctx context.Context, client apiv1.Checkpoint return result } +func checkpointSnapshotPending(err error) bool { + for _, detail := range status.Convert(err).Details() { + if info, ok := detail.(*errdetails.ErrorInfo); ok && info.Reason == "KAGENT_CHECKPOINT_SNAPSHOT_PENDING" { + return true + } + } + return false +} + func TestSessionCheckpoint(t *testing.T) { t.Parallel() forEachHarness(t, func(t *testing.T, harness testHarness) { diff --git a/go/core/test/e2e/scheduled_run_test.go b/go/core/test/e2e/scheduled_run_test.go index f5c9062f82..2372156332 100644 --- a/go/core/test/e2e/scheduled_run_test.go +++ b/go/core/test/e2e/scheduled_run_test.go @@ -27,9 +27,7 @@ import ( "google.golang.org/grpc/status" "google.golang.org/protobuf/proto" "google.golang.org/protobuf/types/known/durationpb" - corev1 "k8s.io/api/core/v1" "k8s.io/apimachinery/pkg/util/wait" - ctrlclient "sigs.k8s.io/controller-runtime/pkg/client" apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" ) @@ -152,32 +150,8 @@ func TestScheduledRunControllerRestart(t *testing.T) { if len(releases) == 0 || t.Failed() { return } - ctx, cancel := context.WithTimeout(t.Context(), 2*time.Minute) - defer cancel() - kube := interactionKubeClient(t) - pods := &corev1.PodList{} - require.NoError(t, kube.List(ctx, pods, ctrlclient.InNamespace("kagent"), ctrlclient.MatchingLabels{"app.kubernetes.io/component": "controller"})) - require.Len(t, pods.Items, 1, "restart test requires a single controller replica") - old := pods.Items[0] - require.NoError(t, kube.Delete(ctx, &old)) - require.NoError(t, wait.PollUntilContextTimeout(ctx, time.Second, 90*time.Second, true, func(ctx context.Context) (bool, error) { - current := &corev1.PodList{} - if err := kube.List(ctx, current, ctrlclient.InNamespace("kagent"), ctrlclient.MatchingLabels{"app.kubernetes.io/component": "controller"}); err != nil { - return false, err - } - for _, pod := range current.Items { - if pod.UID == old.UID { - continue - } - for _, condition := range pod.Status.Conditions { - if condition.Type == corev1.PodReady && condition.Status == corev1.ConditionTrue { - return true, nil - } - } - } - return false, nil - }), "wait for replacement controller") - newControllerConn(t, interactionTarget(t)) + restartTestController(t) + restarted = true }