From 8a3d695f428f991f55870e4088225fb871ec9f12 Mon Sep 17 00:00:00 2001 From: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> Date: Wed, 4 Nov 2026 20:51:37 +0800 Subject: [PATCH 01/20] fix(checkpoint): recover interrupted checkpoint operations Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> --- .../persistence-checkpoints-and-forks.md | 32 +- .../database/checkpoint_operations.go | 109 +++++ .../database/checkpoint_operations_test.go | 98 +++++ go/core/internal/database/checkpoints.go | 3 + .../internal/service/checkpoint/recovery.go | 135 ++++++ .../service/checkpoint/recovery_test.go | 383 ++++++++++++++++++ .../internal/service/checkpoint/service.go | 236 +++++++---- .../service/checkpoint/service_test.go | 194 ++++++--- go/core/pkg/app/app.go | 3 + .../core/000002_checkpoint_recovery.sql | 6 + go/core/pkg/migrations/runner_test.go | 3 + go/core/test/e2e/README.md | 6 + go/core/test/e2e/checkpoint_recovery_test.go | 96 +++++ go/core/test/e2e/connection_test.go | 36 ++ go/core/test/e2e/interaction_test.go | 15 +- go/core/test/e2e/scheduled_run_test.go | 30 +- 16 files changed, 1227 insertions(+), 158 deletions(-) create mode 100644 go/core/internal/database/checkpoint_operations.go create mode 100644 go/core/internal/database/checkpoint_operations_test.go create mode 100644 go/core/internal/service/checkpoint/recovery.go create mode 100644 go/core/internal/service/checkpoint/recovery_test.go create mode 100644 go/core/pkg/migrations/core/000002_checkpoint_recovery.sql create mode 100644 go/core/test/e2e/checkpoint_recovery_test.go diff --git a/docs/architecture/persistence-checkpoints-and-forks.md b/docs/architecture/persistence-checkpoints-and-forks.md index 45c1e1cce4..d8ffd3e9af 100644 --- a/docs/architecture/persistence-checkpoints-and-forks.md +++ b/docs/architecture/persistence-checkpoints-and-forks.md @@ -64,6 +64,27 @@ Creating a checkpoint does not suspend the Actor again: 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. +The RPC reserves durable work and waits for its outcome. Only the leader-elected +checkpoint worker creates and deletes snapshot tags; API replicas do not issue +competing copies. Canceling a request stops its wait, not the admitted operation. +Retrying with the same request ID observes the same checkpoint. Startup and +periodic scans resume `CREATING` and `DELETING` work without requiring the caller +to return. The worker polls once per second so admission on another API replica does not +wait for a caller-local signal; local API requests also wake it immediately. + +A pending tag is not evidence of abandonment. Substrate serializes `CreateTag` +against both the source Actor and the tag, including a copy still finishing after +an earlier controller disappeared. Recovery calls `CreateTag` under those leases +before interpreting `AlreadyExists`: a completed copy can be published, while a +partial copy can be compensated once the prior copy has ended. Transport errors +retain the reservation; checkpoint age never authorizes cleanup. + +Compensation first records the failure and exact tag UID while the checkpoint +remains `CREATING`. Source admission and all reference pins remain held. A restart +resumes deletion of that UID rather than creating or publishing another copy. +Only confirmed cleanup marks the checkpoint `FAILED`; a replaced or unverified +tag is never deleted. Database transactions do not span Substrate calls. + Kagent requests synchronous Substrate suspension before recording a checkpoint. If an earlier attempt already suspended the Actor but left a pending durable upload, the adapter waits within the request deadline for that same snapshot @@ -97,7 +118,7 @@ Failed checkpoints can also be deleted by their owner. Creation marks a checkpoi 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 +`DELETING` for the checkpoint worker to retry. Its history and runtime revision references remain retained until deletion completes. Checkpoints still in `CREATING` cannot be deleted through this path. @@ -165,11 +186,12 @@ snapshot forks. The Agent endpoint and globally unique task ID identify a task. Reads, writes, cancellation, subscriptions, authorization, and dispatch remain scoped to the resolved session. Forks have independent public identities. -## Unreleased schema +## Schema migrations -Until release, core schema changes are folded into `000001_initial.sql`. -Recreate development databases when that baseline changes; there is no upgrade -path from earlier development schemas. The Down migration removes the core schema. +Merged migration files are immutable. Core changes use new transactional Goose +migrations, including the partial index used to discover pending checkpoint work. +Development databases created from older, rewritten alpha baselines still need +to be recreated; new migrations do not bridge those historical schemas. Deploy the controller and clients together: clients must use `Session.context_id` or omit the context and let the routed gateway resolve it. diff --git a/go/core/internal/database/checkpoint_operations.go b/go/core/internal/database/checkpoint_operations.go new file mode 100644 index 0000000000..c3e36c3523 --- /dev/null +++ b/go/core/internal/database/checkpoint_operations.go @@ -0,0 +1,109 @@ +package database + +import ( + "context" + "fmt" + + "github.com/jackc/pgx/v5" + "google.golang.org/protobuf/proto" + + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" +) + +// CheckpointOperation contains the durable inputs for creation or deletion. +// Failure on a CREATING checkpoint means tag cleanup must finish before the +// reservation can become FAILED and release source-session admission. +type CheckpointOperation struct { + Checkpoint *apiv1.Checkpoint + Creator string + Snapshot *SessionTaskSnapshot + TagUID string +} + +// 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 := readCheckpoint(ctx, c.db, id, owner, nil) + 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 +// and deletion work across owners. Only the controller's 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') + AND ($1::uuid IS NULL OR id > $1::uuid) + ORDER BY id LIMIT $2 + `, pgx.RowToStructByName[sessionCheckpointRow], after, limit) + 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 +} + +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 +} + +// BeginCheckpointCleanup durably records creation compensation and its tag UID. +// The checkpoint remains CREATING, retaining history/revision references and +// blocking source mutations until external cleanup finishes. Identical retries +// succeed; a conflicting outcome or a terminal checkpoint returns ErrConflict. +func (c *Client) BeginCheckpointCleanup(ctx context.Context, id, tagUID, failure string) error { + if failure == "" { + return fmt.Errorf("checkpoint cleanup requires a failure") + } + return c.withTx(ctx, func(tx pgx.Tx) error { + row, err := lockCheckpoint(ctx, tx, id, true, "", nil) + if err != nil { + return notFoundOr(err) + } + checkpoint, err := toSessionCheckpoint(row) + if err != nil { + return err + } + if row.State != "CREATING" { + return ErrConflict + } + if checkpoint.Failure != nil { + if row.TagUID == tagUID && checkpoint.Failure.Message == failure { + return nil + } + return ErrConflict + } + checkpoint.Failure = &apiv1.Failure{Reason: "SnapshotTagFailed", Message: failure} + data, err := proto.Marshal(checkpoint) + if err != nil { + return fmt.Errorf("encode checkpoint cleanup: %w", err) + } + _, err = tx.Exec(ctx, `UPDATE session_checkpoint SET tag_uid = $2, data = $3 + WHERE id = $1 AND state = 'CREATING'`, row.ID, tagUID, data) + return err + }) +} 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..ee9f933e02 --- /dev/null +++ b/go/core/internal/database/checkpoint_operations_test.go @@ -0,0 +1,98 @@ +package database + +import ( + "testing" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/google/uuid" + "github.com/stretchr/testify/require" + + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" +) + +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 := client.ReserveSessionCheckpoint(ctx, + &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID)}, "alice", uuid.NewString()) + require.NoError(t, err) + return checkpoint, session +} + +func TestCheckpointCleanupSurvivesRestartAndFencesPublication(t *testing.T) { + pool := setupTestDB(t) + client := NewClient(pool, "public") + checkpoint, session := reserveRecoveryCheckpoint(t, client) + ctx := t.Context() + require.NoError(t, client.BeginCheckpointCleanup(ctx, checkpoint.Id, "tag-uid", "copy interrupted")) + client = NewClient(pool, "public") + op, err := client.GetCheckpointOperation(ctx, checkpoint.Id, "alice") + require.NoError(t, err) + require.Equal(t, "tag-uid", op.TagUID) + require.Equal(t, "copy interrupted", op.Checkpoint.Failure.Message) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_CREATING, op.Checkpoint.State) + _, err = client.GetCheckpointOperation(ctx, checkpoint.Id, "mallory") + require.ErrorIs(t, err, ErrNotFound) + _, err = client.GetSessionCheckpoint(ctx, checkpoint.Id, "alice") + require.ErrorIs(t, err, ErrNotFound) + require.ErrorIs(t, deleteSession(ctx, client, session.Id), ErrConflict) + _, _, err = client.ForkSession(ctx, checkpoint.Id, "alice", uuid.NewString(), uuid.NewString()) + require.ErrorIs(t, err, ErrNotFound) + _, err = client.FinalizeSessionCheckpoint(ctx, checkpoint.Id, "tag-uid", "s3://tags/copy", "") + require.ErrorIs(t, err, ErrConflict, "cleanup must not race a stale successful copy into READY") + require.ErrorIs(t, client.BeginCheckpointCleanup(ctx, checkpoint.Id, "replacement", "copy interrupted"), ErrConflict) + require.NoError(t, client.BeginCheckpointCleanup(ctx, checkpoint.Id, "tag-uid", "copy interrupted")) + _, err = client.FinalizeSessionCheckpoint(ctx, checkpoint.Id, "", "", "different failure") + require.ErrorIs(t, err, ErrConflict) + _, err = client.FinalizeSessionCheckpoint(ctx, checkpoint.Id, "", "", "copy interrupted") + require.NoError(t, err) + // Only completion of compensation releases source admission. + next := newSessionTask(uuid.NewString(), uuid.NewString()) + next.ContextID = session.ContextId + _, err = client.CreateRuntimeTask(ctx, session.Id, taskMutationHash("after-cleanup"), next, "") + require.NoError(t, err) +} + +func TestPendingCheckpointsPageBothPhasesAndRetainInputs(t *testing.T) { + client := NewClient(setupTestDB(t), "public") + creating, _ := reserveRecoveryCheckpoint(t, client) + deleting, _ := reserveRecoveryCheckpoint(t, client) + failed, _ := reserveRecoveryCheckpoint(t, client) + ctx := t.Context() + _, err := client.FinalizeSessionCheckpoint(ctx, deleting.Id, "retained-uid", "s3://tags/retained", "") + require.NoError(t, err) + _, _, err = client.BeginDeleteSessionCheckpoint(ctx, deleting.Id, "alice") + require.NoError(t, err) + _, err = client.FinalizeSessionCheckpoint(ctx, failed.Id, "", "", "failed") + require.NoError(t, err) + var after string + seen := map[string]*CheckpointOperation{} + for { + page, err := client.ListPendingCheckpoints(ctx, after, 1) + require.NoError(t, err) + if len(page) == 0 { + break + } + require.Len(t, page, 1) + op := page[0] + seen[op.Checkpoint.Id] = op + after = op.Checkpoint.Id + } + require.Len(t, seen, 2) + require.Equal(t, "s3://snapshots/source", seen[creating.Id].Snapshot.URI) + require.Equal(t, "alice", seen[creating.Id].Creator) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_DELETING, seen[deleting.Id].Checkpoint.State) + require.Equal(t, "s3://tags/retained", seen[deleting.Id].Snapshot.URI) + require.Equal(t, "retained-uid", seen[deleting.Id].TagUID) + require.NotContains(t, seen, failed.Id) +} diff --git a/go/core/internal/database/checkpoints.go b/go/core/internal/database/checkpoints.go index d2df87d782..466f244546 100644 --- a/go/core/internal/database/checkpoints.go +++ b/go/core/internal/database/checkpoints.go @@ -353,6 +353,9 @@ func (c *Client) FinalizeSessionCheckpoint(ctx context.Context, id, tagUID, snap } return ErrNotFound } + if result.Failure != nil && (failure == "" || result.Failure.Message != failure) { + return ErrConflict + } result.State = apiv1.CheckpointState_CHECKPOINT_STATE_READY if failure != "" { result.State = apiv1.CheckpointState_CHECKPOINT_STATE_FAILED diff --git a/go/core/internal/service/checkpoint/recovery.go b/go/core/internal/service/checkpoint/recovery.go new file mode 100644 index 0000000000..0ac14d3df0 --- /dev/null +++ b/go/core/internal/service/checkpoint/recovery.go @@ -0,0 +1,135 @@ +package checkpoint + +import ( + "context" + "errors" + "fmt" + "time" + + "golang.org/x/sync/errgroup" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + "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/substrate" + "github.com/kagent-dev/kagent/go/pkg/logging" +) + +// RecoveryWorker is the sole owner of checkpoint tag mutations. API replicas +// admit durable work and wait for its outcome; losing a caller does not cancel +// the copy. On restart or leader change the worker resumes persisted phases. +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: time.Second} +} + +func (*RecoveryWorker) NeedLeaderElection() bool { return true } + +func (w *RecoveryWorker) Start(ctx context.Context) error { + ticker := time.NewTicker(w.interval) + defer ticker.Stop() + for ctx.Err() == nil { + if err := w.sweep(ctx); err != nil && ctx.Err() == nil { + logging.FromContext(ctx).ErrorContext(ctx, "recover checkpoints", "error", err) + } + select { + case <-ctx.Done(): + case <-ticker.C: + case <-w.service.wake: + } + } + 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 + } + var group errgroup.Group + group.SetLimit(4) + for _, op := range operations { + group.Go(func() error { + attempt, cancel := context.WithTimeout(ctx, 2*time.Minute) + defer cancel() + if err := w.service.recoverOperation(attempt, op); err != nil && ctx.Err() == nil { + logging.FromContext(ctx).ErrorContext(ctx, "recover checkpoint operation", "checkpoint_id", op.Checkpoint.Id, "error", err) + } + return nil + }) + } + if err := group.Wait(); err != nil { + return err + } + if len(operations) < 100 { + return nil + } + after = operations[len(operations)-1].Checkpoint.Id + } + return ctx.Err() +} + +func (s *Service) recoverOperation(ctx context.Context, op *database.CheckpointOperation) error { + switch op.Checkpoint.State { + case apiv1.CheckpointState_CHECKPOINT_STATE_DELETING: + return s.finishDelete(ctx, op) + case apiv1.CheckpointState_CHECKPOINT_STATE_CREATING: + if op.Checkpoint.Failure != nil { + return s.finishCleanup(ctx, op) + } + tag, err := s.ensureTag(ctx, op.Checkpoint, op.Snapshot) + if err == nil { + _, err = s.store.FinalizeSessionCheckpoint(ctx, op.Checkpoint.Id, + tag.GetMetadata().GetUid(), substrate.SnapshotURI(tag.GetStatus().GetSnapshot()), "") + return err + } + if !errors.Is(err, errTagCopyFailed) { + return err + } + if tag != nil && !validTagSource(tag, op.Checkpoint, op.Snapshot) { + return fmt.Errorf("refuse cleanup of an unverified checkpoint tag: %w", err) + } + uid := tag.GetMetadata().GetUid() + if recordErr := s.store.BeginCheckpointCleanup(ctx, op.Checkpoint.Id, uid, err.Error()); recordErr != nil { + return recordErr + } + // Persist the exact phase before external cleanup. A restart must not + // mistake the same tag for a copy that can still become READY. + op, err = s.store.GetCheckpointOperation(ctx, op.Checkpoint.Id, op.Creator) + if err != nil { + return err + } + return s.finishCleanup(ctx, op) + default: + return nil + } +} + +func (s *Service) finishCleanup(ctx context.Context, op *database.CheckpointOperation) error { + name := tagName(op.Checkpoint.Id) + tag, err := s.tags.GetTag(ctx, op.Snapshot.Atespace, name) + if err != nil && status.Code(err) != codes.NotFound { + return err + } + if err == nil { + if op.TagUID == "" || tag.GetMetadata().GetUid() != op.TagUID || !validTagSource(tag, op.Checkpoint, op.Snapshot) { + return status.Error(codes.FailedPrecondition, "checkpoint cleanup tag identity changed") + } + if err := s.tags.DeleteTag(ctx, op.Snapshot.Atespace, name, op.TagUID); err != nil && status.Code(err) != codes.NotFound { + return err + } + } + _, err = s.store.FinalizeSessionCheckpoint(ctx, op.Checkpoint.Id, "", "", op.Checkpoint.Failure.Message) + return err +} 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..a605c511fb --- /dev/null +++ b/go/core/internal/service/checkpoint/recovery_test.go @@ -0,0 +1,383 @@ +package checkpoint + +import ( + "context" + "crypto/sha256" + "errors" + "sync" + "testing" + "time" + + "github.com/a2aproject/a2a-go/v2/a2a" + a2apb "github.com/a2aproject/a2a-go/v2/a2apb/v1" + "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgxpool" + "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/dbtest" + "github.com/kagent-dev/kagent/go/core/pkg/auth" +) + +// 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 startRecovery(t *testing.T, service *Service) func() { + t.Helper() + ctx, cancel := context.WithCancel(t.Context()) + done := make(chan error, 1) + worker := NewRecoveryWorker(service) + worker.interval = 10 * time.Millisecond + go func() { done <- worker.Start(ctx) }() + var once sync.Once + stop := func() { once.Do(func() { cancel(); require.NoError(t, <-done) }) } + t.Cleanup(stop) + return stop +} + +func recoverySource(t *testing.T, store *database.Client) (*apiv1.Session, string) { + t.Helper() + ctx := t.Context() + require.NoError(t, store.UpsertAgentDefinition(ctx, database.AgentDefinition{ + Namespace: "team-a", AgentName: "assistant", AgentUID: "agent-uid", DesiredRevision: "revision"})) + _, err := store.RecordRuntimeRevision(ctx, database.RuntimeRevision{ + Revision: "revision", Namespace: "team-a", AgentName: "assistant", AgentUID: "agent-uid", + SourceSnapshot: []byte("{}"), AgentCard: &a2apb.AgentCard{}, EgressDestinations: []string{}, + ActorTemplateAtespace: "team-a", ActorTemplateName: "runtime", ActorTemplateUID: "template-uid", + }, database.RecordRevisionOptions{Ready: true}) + require.NoError(t, err) + session, _, err := store.CreateSession(ctx, &apiv1.Session{Id: uuid.NewString(), Creator: "alice", + Agent: &apiv1.AgentProvenance{Ref: &apiv1.ResourceReference{Namespace: "team-a", Name: "assistant"}}}, uuid.NewString()) + require.NoError(t, err) + op, err := store.BeginSessionOperation(ctx, session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_CREATE) + 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, "runtime", "actor-uid", "") + 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: "team-a", URI: "s3://snapshots/source", Fidelity: "SNAPSHOT_FIDELITY_VOLUMES"})) + return session, string(task.ID) +} + +func waitCheckpointPhase(t *testing.T, store *database.Client, id string, phase apiv1.CheckpointState) *database.CheckpointOperation { + t.Helper() + var op *database.CheckpointOperation + require.Eventually(t, func() bool { + var err error + op, err = store.GetCheckpointOperation(t.Context(), id, "alice") + return err == nil && op.Checkpoint.State == phase + }, 5*time.Second, 10*time.Millisecond) + return op +} + +type failedCheckpointPublish struct { + *database.Client + entered chan struct{} + once sync.Once +} + +func (s *failedCheckpointPublish) FinalizeSessionCheckpoint(ctx context.Context, _, _, _, _ string) (*apiv1.Checkpoint, error) { + s.once.Do(func() { close(s.entered) }) + <-ctx.Done() + return nil, ctx.Err() +} + +func TestCheckpointRecoveryWithPostgres(t *testing.T) { + if testing.Short() { + t.Skip("requires PostgreSQL") + } + conn := dbtest.StartT(t.Context(), t) + require.NoError(t, dbtest.Migrate(conn, false)) + pool, err := pgxpool.New(t.Context(), conn) + require.NoError(t, err) + t.Cleanup(pool.Close) + newStore := func() *database.Client { return database.NewClient(pool, "public") } + ownerContext := func(ctx context.Context) context.Context { + return auth.AuthSessionTo(ctx, testSession{userID: "alice"}) + } + + t.Run("startup recovers a reservation whose caller disappeared", func(t *testing.T) { + store := newStore() + session, head := recoverySource(t, store) + tags := &lockedTags{base: &testTags{snapshotURI: "s3://snapshots/source"}} + service := NewService(store, testAuthorizer{}, tags, nil) + request := uuid.NewString() + ctx, cancel := context.WithCancel(ownerContext(t.Context())) + t.Cleanup(cancel) + done := make(chan error, 1) + go func() { _, err := service.Create(ctx, session.Id, request, head); done <- err }() + var pending []*database.CheckpointOperation + require.Eventually(t, func() bool { + var err error + pending, err = store.ListPendingCheckpoints(t.Context(), "", 10) + return err == nil && len(pending) == 1 + }, 5*time.Second, 10*time.Millisecond) + cancel() + require.ErrorIs(t, <-done, context.Canceled) + require.Zero(t, tags.base.createCalls, "API admission must not start its own copy") + replacement := NewService(newStore(), testAuthorizer{}, tags, nil) + stop := startRecovery(t, replacement) + replayed, err := replacement.Create(ownerContext(t.Context()), session.Id, request, head) + require.NoError(t, err) + require.Equal(t, pending[0].Checkpoint.Id, replayed.Id) + stop() + }) + + t.Run("canceled caller and concurrent replicas leave one live copy", func(t *testing.T) { + store := newStore() + session, head := recoverySource(t, store) + gate := make(chan struct{}) + tags := &lockedTags{base: &testTags{snapshotURI: "s3://snapshots/source"}, gate: gate, entered: make(chan struct{})} + leader := NewService(store, testAuthorizer{}, tags, nil) + stop := startRecovery(t, leader) + ctx, cancel := context.WithCancel(ownerContext(t.Context())) + result := make(chan error, 1) + request := uuid.NewString() + go func() { _, err := leader.Create(ctx, session.Id, request, head); result <- err }() + select { + case <-tags.entered: + case <-time.After(5 * time.Second): + t.Fatal("copy did not start") + } + cancel() + require.ErrorIs(t, <-result, context.Canceled) + // Another replica shares durable admission, not local singleflight. + follower := NewService(newStore(), testAuthorizer{}, tags, nil) + results := make(chan *apiv1.Checkpoint, 8) + errors := make(chan error, 8) + for range 8 { + go func() { + cp, err := follower.Create(ownerContext(t.Context()), session.Id, request, head) + results <- cp + errors <- err + }() + } + pending, err := store.ListPendingCheckpoints(t.Context(), "", 10) + require.NoError(t, err) + require.Len(t, pending, 1) + _, err = store.BeginSessionOperation(t.Context(), session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE) + require.ErrorIs(t, err, database.ErrConflict) + close(gate) + for range 8 { + require.NoError(t, <-errors) + require.Equal(t, pending[0].Checkpoint.Id, (<-results).Id) + } + stop() + tags.mu.Lock() + defer tags.mu.Unlock() + require.Equal(t, 1, tags.base.createCalls) + require.Zero(t, tags.base.deleteCalls) + }) + + t.Run("restart reuses completed tag after lost database publication", func(t *testing.T) { + store := newStore() + session, head := recoverySource(t, store) + tags := &lockedTags{base: &testTags{snapshotURI: "s3://snapshots/source"}} + fault := &failedCheckpointPublish{Client: store, entered: make(chan struct{})} + service := NewService(fault, testAuthorizer{}, tags, nil) + cp, _, err := service.beginCreate(ownerContext(t.Context()), session.Id, uuid.NewString(), head) + require.NoError(t, err) + stop := startRecovery(t, service) + select { + case <-fault.entered: + case <-time.After(5 * time.Second): + t.Fatal("publication did not start") + } + stop() + waitCheckpointPhase(t, store, cp.Id, apiv1.CheckpointState_CHECKPOINT_STATE_CREATING) + replacement := NewService(newStore(), testAuthorizer{}, tags, nil) + stop = startRecovery(t, replacement) + waitCheckpointPhase(t, store, cp.Id, apiv1.CheckpointState_CHECKPOINT_STATE_READY) + stop() + tags.mu.Lock() + defer tags.mu.Unlock() + require.Equal(t, "tag-uid", tags.base.created.GetMetadata().GetUid()) + require.Zero(t, tags.base.deleteCalls) + }) + + t.Run("restart completes compensation and releases source admission", func(t *testing.T) { + store := newStore() + session, head := recoverySource(t, store) + tags := &lockedTags{base: &testTags{snapshotURI: "s3://snapshots/source", deleteErr: errors.New("store offline"), + mutateTag: func(tag *ateapipb.Tag) { + tag.Status.Snapshot.DurableSnapshot.State = ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_PENDING + }}} + service := NewService(store, testAuthorizer{}, tags, nil) + cp, _, err := service.beginCreate(ownerContext(t.Context()), session.Id, uuid.NewString(), head) + require.NoError(t, err) + _, err = tags.base.CreateTag(t.Context(), "team-a", tagName(cp.Id), "session-"+session.Id) + require.NoError(t, err) + tags.base.createCalls = 0 + stop := startRecovery(t, service) + require.Eventually(t, func() bool { + op, err := store.GetCheckpointOperation(t.Context(), cp.Id, "alice") + return err == nil && op.Checkpoint.Failure != nil + }, 5*time.Second, 10*time.Millisecond) + stop() + tags.mu.Lock() + tags.base.deleteErr = nil + tags.mu.Unlock() + replacement := NewService(newStore(), testAuthorizer{}, tags, nil) + stop = startRecovery(t, replacement) + waitCheckpointPhase(t, store, cp.Id, apiv1.CheckpointState_CHECKPOINT_STATE_FAILED) + stop() + tags.mu.Lock() + require.Equal(t, 1, tags.base.createCalls) + require.Nil(t, tags.base.created) + tags.mu.Unlock() + message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("after recovery")) + message.ContextID = session.ContextId + task := a2a.NewSubmittedTask(message, message) + hash := sha256.Sum256([]byte(task.ID)) + _, err = store.CreateRuntimeTask(t.Context(), session.Id, hash[:], task, "") + require.NoError(t, err) + }) + + t.Run("restart finishes deletion after backend failure", func(t *testing.T) { + store := newStore() + session, head := recoverySource(t, store) + tags := &lockedTags{base: &testTags{snapshotURI: "s3://snapshots/source"}} + service := NewService(store, testAuthorizer{}, tags, nil) + stop := startRecovery(t, service) + cp, err := service.Create(ownerContext(t.Context()), session.Id, uuid.NewString(), head) + require.NoError(t, err) + stop() + tags.mu.Lock() + tags.base.deleteErr = status.Error(codes.Unavailable, "store offline") + tags.mu.Unlock() + _, err = service.beginDelete(ownerContext(t.Context()), cp.Id) + require.NoError(t, err) + stop = startRecovery(t, service) + waitCheckpointPhase(t, store, cp.Id, apiv1.CheckpointState_CHECKPOINT_STATE_DELETING) + stop() + tags.mu.Lock() + tags.base.deleteErr = nil + tags.mu.Unlock() + stop = startRecovery(t, NewService(newStore(), testAuthorizer{}, tags, nil)) + require.Eventually(t, func() bool { + _, err := store.GetCheckpointOperation(t.Context(), cp.Id, "alice") + return errors.Is(err, database.ErrNotFound) + }, 5*time.Second, 10*time.Millisecond) + stop() + tags.mu.Lock() + defer tags.mu.Unlock() + require.Nil(t, tags.base.created) + }) +} + +// 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) +} diff --git a/go/core/internal/service/checkpoint/service.go b/go/core/internal/service/checkpoint/service.go index de5cb6a390..ad4be9a968 100644 --- a/go/core/internal/service/checkpoint/service.go +++ b/go/core/internal/service/checkpoint/service.go @@ -6,10 +6,10 @@ import ( "errors" "fmt" "strings" + "time" "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" @@ -27,6 +27,9 @@ const ( ) type store interface { + GetCheckpointOperation(context.Context, string, string) (*database.CheckpointOperation, error) + ListPendingCheckpoints(context.Context, string, int) ([]*database.CheckpointOperation, error) + BeginCheckpointCleanup(context.Context, string, string, string) error ReserveSessionCheckpoint(context.Context, *apiv1.Checkpoint, string, string) (*apiv1.Checkpoint, *database.SessionTaskSnapshot, error) FinalizeSessionCheckpoint(context.Context, string, string, string, string) (*apiv1.Checkpoint, error) GetSessionCheckpoint(context.Context, string, string) (*apiv1.Checkpoint, error) @@ -49,15 +52,15 @@ 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 workflow workflow + wake chan struct{} } type ListRequest struct { @@ -72,40 +75,41 @@ type ListResult struct { } func NewService(store store, authorizer auth.Authorizer, tags tagClient, workflow workflow) *Service { - return &Service{store: store, authorizer: authorizer, tags: tags, workflow: workflow} + return &Service{store: store, authorizer: authorizer, tags: tags, workflow: workflow, wake: make(chan struct{}, 1)} } func (s *Service) Create(ctx context.Context, sessionID, requestID, expectedHeadTaskID string) (*apiv1.Checkpoint, error) { - sessionID, err := validateCreate(sessionID, requestID) + checkpoint, owner, err := s.beginCreate(ctx, sessionID, requestID, expectedHeadTaskID) if err != nil { return nil, err } + if checkpoint.State != apiv1.CheckpointState_CHECKPOINT_STATE_CREATING { + return checkpoint, nil + } + s.wakeRecovery() + return s.waitCreate(ctx, checkpoint.Id, owner) +} + +func (s *Service) beginCreate(ctx context.Context, sessionID, requestID, expectedHeadTaskID string) (*apiv1.Checkpoint, string, error) { + sessionID, err := validateCreate(sessionID, requestID) + if err != nil { + return nil, "", err + } if expectedHeadTaskID == "" { - return nil, status.Error(codes.InvalidArgument, "expected_head_task_id is required") + return nil, "", status.Error(codes.InvalidArgument, "expected_head_task_id is required") } userID, err := s.authorize(ctx, auth.VerbCreate, auth.Resource{Type: "Checkpoint"}) if err != nil { - return nil, err + return nil, "", err } 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 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) + return nil, "", serviceerrors.NewInternal("Failed to generate checkpoint identifier", err) } - checkpoint, snapshot, err := s.store.ReserveSessionCheckpoint(ctx, &apiv1.Checkpoint{Id: id.String(), SessionId: sessionID, HeadTaskId: expectedHeadTaskID}, userID, requestID) + checkpoint, _, 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) { @@ -113,80 +117,88 @@ func (s *Service) create(ctx context.Context, userID, sessionID, requestID, expe } failure, detailErr := status.New(codes.FailedPrecondition, err.Error()).WithDetails(&errdetails.ErrorInfo{Domain: "kagent.dev", Reason: reason}) if detailErr != nil { - return nil, detailErr + return nil, "", detailErr } - return nil, failure.Err() + return nil, "", failure.Err() } if errors.Is(err, database.ErrIdempotencyConflict) { - return nil, serviceerrors.NewAlreadyExists("request_id was already used for a different checkpoint", err) + return nil, "", serviceerrors.NewAlreadyExists("request_id was already used for a different checkpoint", err) } if errors.Is(err, database.ErrNotFound) { - return nil, serviceerrors.NewNotFound("Session not found", err) + return nil, "", serviceerrors.NewNotFound("Session not found", err) } if errors.Is(err, database.ErrConflict) || errors.Is(err, database.ErrFailedPrecondition) { - return nil, serviceerrors.NewFailedPrecondition(err.Error(), err) + return nil, "", serviceerrors.NewFailedPrecondition(err.Error(), err) } if err != nil { - return nil, serviceerrors.NewInternal("Failed to reserve checkpoint", err) + 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 checkpoint, 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. +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) (*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()) + 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 + } + name := tagName(checkpoint.Id) + // CreateTag holds Substrate's actor and tag leases. AlreadyExists is a + // synchronization barrier: a pending tag read after it is no longer being + // copied by an earlier call. A plain GetTag cannot establish that fact. 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 +208,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 } @@ -255,21 +267,36 @@ func (s *Service) List(ctx context.Context, request ListRequest) (ListResult, er } func (s *Service) Delete(ctx context.Context, checkpointID string) error { - checkpointID, err := validateIdentity(checkpointID) + owner, err := s.beginDelete(ctx, checkpointID) if err != nil { return err } + s.wakeRecovery() + return s.waitDeleted(ctx, checkpointID, owner) +} + +func (s *Service) beginDelete(ctx context.Context, checkpointID string) (string, error) { + checkpointID, err := validateIdentity(checkpointID) + if err != nil { + return "", err + } userID, err := s.authorize(ctx, auth.VerbDelete, auth.Resource{Type: "Checkpoint", Name: checkpointID}) if err != nil { - return err + return "", err } - snapshot, tagUID, err := s.store.BeginDeleteSessionCheckpoint(ctx, checkpointID, userID) + _, _, err = s.store.BeginDeleteSessionCheckpoint(ctx, checkpointID, userID) if errors.Is(err, database.ErrNotFound) { - return serviceerrors.NewNotFound("Checkpoint not found", err) + return "", serviceerrors.NewNotFound("Checkpoint not found", err) } if err != nil { - return serviceerrors.NewInternal("Failed to begin checkpoint deletion", err) + return "", serviceerrors.NewInternal("Failed to begin checkpoint deletion", err) } + return userID, nil +} + +func (s *Service) finishDelete(ctx context.Context, op *database.CheckpointOperation) error { + checkpointID, userID := op.Checkpoint.Id, op.Creator + 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) @@ -443,3 +470,62 @@ func decodePageToken(token string) (string, error) { } return validateIdentity(string(value)) } + +func (s *Service) wakeRecovery() { + select { + case s.wake <- struct{}{}: + default: + } +} + +func (s *Service) waitCreate(ctx context.Context, id, owner string) (*apiv1.Checkpoint, error) { + ticker := time.NewTicker(100 * time.Millisecond) + defer ticker.Stop() + for { + op, err := s.store.GetCheckpointOperation(ctx, id, owner) + if err != nil { + if ctx.Err() != nil { + return nil, ctx.Err() + } + if errors.Is(err, database.ErrNotFound) { + return nil, serviceerrors.NewNotFound("Checkpoint no longer exists", err) + } + return nil, serviceerrors.NewInternal("Failed to read checkpoint operation", err) + } + switch op.Checkpoint.State { + case apiv1.CheckpointState_CHECKPOINT_STATE_READY: + return op.Checkpoint, nil + case apiv1.CheckpointState_CHECKPOINT_STATE_DELETING: + return nil, serviceerrors.NewAborted("Checkpoint creation was superseded by deletion", nil) + case apiv1.CheckpointState_CHECKPOINT_STATE_FAILED: + return nil, serviceerrors.NewUnavailable("Failed to retain checkpoint snapshot", errors.New(op.Checkpoint.GetFailure().GetMessage())) + } + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-ticker.C: + } + } +} + +func (s *Service) waitDeleted(ctx context.Context, id, owner string) error { + ticker := time.NewTicker(100 * time.Millisecond) + defer ticker.Stop() + for { + _, err := s.store.GetCheckpointOperation(ctx, id, owner) + if errors.Is(err, database.ErrNotFound) { + return nil + } + if err != nil { + if ctx.Err() != nil { + return ctx.Err() + } + return serviceerrors.NewInternal("Failed to read checkpoint deletion", err) + } + select { + case <-ctx.Done(): + return ctx.Err() + case <-ticker.C: + } + } +} diff --git a/go/core/internal/service/checkpoint/service_test.go b/go/core/internal/service/checkpoint/service_test.go index f3c0b71f3b..e7e674a6e3 100644 --- a/go/core/internal/service/checkpoint/service_test.go +++ b/go/core/internal/service/checkpoint/service_test.go @@ -5,7 +5,6 @@ import ( "encoding/base64" "errors" "fmt" - "runtime" "strings" "sync/atomic" "testing" @@ -176,6 +175,7 @@ type testTags struct { getErr error createErr error deleteErr error + deleteResponseErr error actorErr error mutateTag func(*ateapipb.Tag) beforeDelete func() @@ -241,7 +241,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) { @@ -250,7 +250,7 @@ func TestCreatePreservesStoreConflictReason(t *testing.T) { storeErr := fmt.Errorf("Session cannot checkpoint in its current state: %w", cause) service := NewService(&testStore{reserveErr: storeErr}, testAuthorizer{}, nil, nil) ctx := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) - _, err := service.Create(ctx, "018f47a2-4efb-7c21-a848-123456789abc", "request-1", "task-1") + _, err := createAndRecover(service, ctx, "018f47a2-4efb-7c21-a848-123456789abc", "request-1", "task-1") require.Equal(t, serviceerrors.CodeFailedPrecondition, serviceerrors.CodeOf(err)) require.Equal(t, storeErr.Error(), serviceerrors.MessageOf(err)) require.ErrorIs(t, err, cause) @@ -268,7 +268,7 @@ func TestCreateDistinguishesPendingSnapshotFromAdvancedConversation(t *testing.T } { service := NewService(&testStore{reserveErr: test.err}, testAuthorizer{}, &testTags{}, &testWorkflow{}) ctx := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) - _, err := service.Create(ctx, "018f47a2-4efb-7c21-a848-123456789abc", "request", "task-a") + _, err := createAndRecover(service, ctx, "018f47a2-4efb-7c21-a848-123456789abc", "request", "task-a") require.Equal(t, codes.FailedPrecondition, status.Code(err)) details := status.Convert(err).Details() require.Len(t, details, 1) @@ -302,7 +302,7 @@ func TestCreateTagsRecordedSnapshotBoundary(t *testing.T) { service := NewService(store, &restrictiveAuthorizer{resource: "Session", verb: auth.VerbCreate}, tags, nil) ctx := auth.AuthSessionTo(context.Background(), testSession{userID: "alice"}) - checkpoint, err := service.Create(ctx, "018f47a2-4efb-7c21-a848-123456789abc", "request-1", "task-1") + checkpoint, err := createAndRecover(service, ctx, "018f47a2-4efb-7c21-a848-123456789abc", "request-1", "task-1") if err != nil { t.Fatal(err) } @@ -330,7 +330,7 @@ func TestCreateCleansTagBeforeFailing(t *testing.T) { 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") + _, err := createAndRecover(service, ctx, "018f47a2-4efb-7c21-a848-123456789abc", "request-1", "task-1") require.Error(t, err) require.Equal(t, 1, tags.deleteCalls) require.NotEmpty(t, store.failed) @@ -350,7 +350,7 @@ func TestDeleteHidesCheckpointBeforeDeletingTag(t *testing.T) { store.forkSourceErr = fmt.Errorf("deletion must use the snapshot returned by BeginDelete") tags.getErr = fmt.Errorf("tag lookup unavailable") - if err := service.Delete(ctx, checkpoint.GetId()); err == nil { + if err := deleteAndRecover(service, ctx, checkpoint.GetId()); err == nil { t.Fatal("Delete() succeeded without verifying the snapshot tag") } if checkpoint.State != apiv1.CheckpointState_CHECKPOINT_STATE_DELETING || tags.deleteCalls != 0 || store.deleted { @@ -358,11 +358,11 @@ 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, serviceerrors.CodeUnavailable, serviceerrors.CodeOf(deleteAndRecover(service, 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 - if err := service.Delete(ctx, checkpoint.GetId()); err != nil { + if err := deleteAndRecover(service, ctx, checkpoint.GetId()); err != nil { t.Fatal(err) } if checkpoint.State != apiv1.CheckpointState_CHECKPOINT_STATE_DELETING || tags.deleteCalls != 2 || !store.deleted { @@ -380,7 +380,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, serviceerrors.CodeUnavailable, serviceerrors.CodeOf(deleteAndRecover(service, ctx, checkpoint.GetId()))) require.False(t, store.deleted) require.Zero(t, tags.deleteCalls) require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_DELETING, checkpoint.State) @@ -388,7 +388,7 @@ func TestDeleteFailedCheckpointWithMissingTag(t *testing.T) { // Retry through a new service instance after the backend recovers. tags.getErr = nil service = NewService(store, testAuthorizer{}, tags, nil) - require.NoError(t, service.Delete(ctx, checkpoint.GetId())) + require.NoError(t, deleteAndRecover(service, ctx, checkpoint.GetId())) require.True(t, store.deleted) require.Zero(t, tags.deleteCalls) } @@ -415,7 +415,7 @@ func TestDeleteRejectsUnverifiedTag(t *testing.T) { }} service := NewService(store, testAuthorizer{}, tags, nil) ctx := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) - require.Equal(t, serviceerrors.CodeFailedPrecondition, serviceerrors.CodeOf(service.Delete(ctx, checkpoint.GetId()))) + require.Equal(t, serviceerrors.CodeFailedPrecondition, serviceerrors.CodeOf(deleteAndRecover(service, ctx, checkpoint.GetId()))) require.False(t, store.deleted) require.Zero(t, tags.deleteCalls) }) @@ -450,13 +450,13 @@ func TestCreateRetriesRetainedTagAfterPublishFailure(t *testing.T) { tags := &testTags{snapshotURI: "s3://snapshots/snapshot-1", createErr: status.Error(codes.DeadlineExceeded, "response lost")} 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") + _, err := createAndRecover(service, 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, 0, tags.deleteCalls) retained := tags.created store.finalizeErr = nil - checkpoint, err := service.Create(ctx, store.prepared.SessionId, "request-1", "task-1") + checkpoint, err := createAndRecover(service, ctx, store.prepared.SessionId, "request-1", "task-1") require.NoError(t, err) require.Same(t, retained, tags.created) require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_READY, checkpoint.State) @@ -472,7 +472,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) { @@ -485,11 +485,11 @@ func TestCreateRejectsInvalidTagAndKeepsCleanupRetryable(t *testing.T) { tags := &testTags{snapshotURI: "s3://snapshots/snapshot-1", mutateTag: tc.mutate, deleteErr: errors.New("cleanup unavailable")} 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") + _, err := createAndRecover(service, 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) tags.deleteErr = nil - _, err = service.Create(ctx, store.prepared.SessionId, "request-1", "task-1") + _, err = createAndRecover(service, ctx, store.prepared.SessionId, "request-1", "task-1") require.Error(t, err) require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_FAILED, store.prepared.State) require.Nil(t, tags.created) @@ -519,27 +519,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 := createAndRecover(service, 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) { @@ -553,17 +543,17 @@ func TestDeletePreservesTagReplacedAfterLookup(t *testing.T) { service := NewService(store, testAuthorizer{}, tags, nil) ctx := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) - require.Error(t, service.Delete(ctx, checkpoint.GetId())) + require.Error(t, deleteAndRecover(service, ctx, checkpoint.GetId())) require.Equal(t, 1, tags.deleteCalls) require.Equal(t, "replacement-uid", tags.created.GetMetadata().GetUid(), "replacement snapshot must survive") require.False(t, store.deleted, "retain the checkpoint reference after a rejected deletion") require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_DELETING, checkpoint.State) tags.beforeDelete = nil - require.Equal(t, serviceerrors.CodeFailedPrecondition, serviceerrors.CodeOf(service.Delete(ctx, checkpoint.GetId()))) + require.Equal(t, serviceerrors.CodeFailedPrecondition, serviceerrors.CodeOf(deleteAndRecover(service, ctx, checkpoint.GetId()))) require.Equal(t, 1, tags.deleteCalls, "retry must reject the replacement at lookup") tags.created = nil - require.NoError(t, service.Delete(ctx, checkpoint.GetId())) + require.NoError(t, deleteAndRecover(service, ctx, checkpoint.GetId())) require.True(t, store.deleted, "cleanup can finish once the tag is absent") } @@ -618,7 +608,7 @@ func TestCheckpointComposesPermissionsBeforeEffects(t *testing.T) { ctx := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) var err error if test.operation == "create" { - _, err = service.Create(ctx, sourceID, "request", "head") + _, err = createAndRecover(service, ctx, sourceID, "request", "head") } else { _, err = service.Fork(ctx, checkpointID, "request") } @@ -677,9 +667,9 @@ func TestCheckpointSharesAndCategoryDenialsProtectEveryOperation(t *testing.T) { return err }}, {"rename", auth.VerbUpdate, func(ctx context.Context, s *Service) error { _, err := s.Rename(ctx, id, "name"); return err }}, - {"delete", auth.VerbDelete, func(ctx context.Context, s *Service) error { return s.Delete(ctx, id) }}, + {"delete", auth.VerbDelete, func(ctx context.Context, s *Service) error { return deleteAndRecover(s, ctx, id) }}, {"create", auth.VerbCreate, func(ctx context.Context, s *Service) error { - _, err := s.Create(ctx, id, "request", "head") + _, err := createAndRecover(s, ctx, id, "request", "head") return err }}, {"fork", auth.VerbGet, func(ctx context.Context, s *Service) error { _, err := s.Fork(ctx, id, "request"); return err }}, @@ -703,10 +693,10 @@ func TestCreateRetriesReadyCheckpointAfterSourceDeletion(t *testing.T) { policy := &restrictiveAuthorizer{} svc := NewService(store, policy, tags, nil) ctx := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) - first, err := svc.Create(ctx, "018f47a2-4efb-7c21-a848-123456789abc", "request-1", "task-1") + first, err := createAndRecover(svc, ctx, "018f47a2-4efb-7c21-a848-123456789abc", "request-1", "task-1") require.NoError(t, err) store.sourceErr = database.ErrNotFound - retried, err := svc.Create(ctx, first.SessionId, "request-1", "task-1") + retried, err := createAndRecover(svc, ctx, first.SessionId, "request-1", "task-1") require.NoError(t, err) require.Equal(t, first.Id, retried.Id) require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_READY, retried.State) @@ -714,11 +704,11 @@ func TestCreateRetriesReadyCheckpointAfterSourceDeletion(t *testing.T) { require.Zero(t, store.sourceReads.Load()) before := store.reservations policy.resource, policy.verb = "Session", auth.VerbGet - _, err = svc.Create(ctx, first.SessionId, "request-1", "task-1") + _, err = createAndRecover(svc, ctx, first.SessionId, "request-1", "task-1") require.True(t, serviceerrors.IsCode(err, serviceerrors.CodePermissionDenied)) require.Equal(t, before, store.reservations) policy.resource = "" - _, err = svc.Create(auth.AuthSessionTo(t.Context(), testSession{userID: "bob"}), first.SessionId, "request-1", "task-1") + _, err = createAndRecover(svc, auth.AuthSessionTo(t.Context(), testSession{userID: "bob"}), first.SessionId, "request-1", "task-1") require.True(t, serviceerrors.IsCode(err, serviceerrors.CodeNotFound)) } @@ -782,13 +772,13 @@ func TestCheckpointCanonicalIdentityDenialsCannotBeBypassed(t *testing.T) { var err error switch op { case "create": - _, err = svc.Create(ctx, encoded, "request", "head") + _, err = createAndRecover(svc, ctx, encoded, "request", "head") case "get": _, err = svc.Get(ctx, encoded) case "rename": _, err = svc.Rename(ctx, encoded, "name") case "delete": - err = svc.Delete(ctx, encoded) + err = deleteAndRecover(svc, ctx, encoded) case "fork": _, err = svc.Fork(ctx, encoded, "request") } @@ -828,7 +818,7 @@ func TestCheckpointCanonicalIdentityReachesStorageAndRetries(t *testing.T) { svc := NewService(store, testAuthorizer{}, nil, nil) ctx := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) for _, id := range []string{canonical, strings.ToUpper(canonical)} { - got, err := svc.Create(ctx, id, "same-request", "head") + got, err := createAndRecover(svc, ctx, id, "same-request", "head") require.NoError(t, err) require.Equal(t, checkpointID, got.GetId()) require.Equal(t, canonical, store.sessionID) @@ -853,3 +843,115 @@ 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 (s *testStore) BeginCheckpointCleanup(_ context.Context, _, uid, failure string) error { + s.prepared.Failure = &apiv1.Failure{Reason: "SnapshotTagFailed", Message: failure} + s.tagUID = uid + return nil +} + +// Run one owner step without a background goroutine so fault injection stays +// deterministic. The PostgreSQL recovery tests exercise the public waiting RPCs +// concurrently with a real recovery worker and replacement service instances. +func createAndRecover(s *Service, ctx context.Context, id, requestID, head string) (*apiv1.Checkpoint, error) { + cp, owner, err := s.beginCreate(ctx, id, requestID, head) + if err != nil || cp.State != apiv1.CheckpointState_CHECKPOINT_STATE_CREATING { + return cp, err + } + op, err := s.store.GetCheckpointOperation(ctx, cp.Id, owner) + if err != nil { + return nil, err + } + if err := s.recoverOperation(ctx, op); err != nil { + return nil, err + } + return s.waitCreate(ctx, cp.Id, owner) +} + +func deleteAndRecover(s *Service, ctx context.Context, id string) error { + owner, err := s.beginDelete(ctx, id) + if err != nil { + return err + } + op, err := s.store.GetCheckpointOperation(ctx, id, owner) + if err != nil { + return err + } + return s.recoverOperation(ctx, op) +} + +func TestCleanupResumesAfterLostDeleteResponse(t *testing.T) { + for _, phase := range []apiv1.CheckpointState{apiv1.CheckpointState_CHECKPOINT_STATE_CREATING, apiv1.CheckpointState_CHECKPOINT_STATE_DELETING} { + t.Run(phase.String(), func(t *testing.T) { + cp := &apiv1.Checkpoint{Id: "11111111-1111-4111-8111-111111111111", SessionId: "22222222-2222-4222-8222-222222222222", State: phase} + if phase == apiv1.CheckpointState_CHECKPOINT_STATE_CREATING { + cp.Failure = &apiv1.Failure{Reason: "SnapshotTagFailed", Message: "copy interrupted"} + } + tags := &testTags{snapshotURI: "s3://snapshots/source", deleteResponseErr: status.Error(codes.Unavailable, "reply lost")} + _, err := tags.CreateTag(t.Context(), "team-a", tagName(cp.Id), substrate.ActorName(cp.SessionId)) + require.NoError(t, err) + snapshot := &database.SessionTaskSnapshot{Atespace: "team-a", URI: "s3://tags/checkpoint", Fidelity: "SNAPSHOT_FIDELITY_VOLUMES"} + store := &testStore{prepared: cp, snapshot: snapshot, tagUID: "tag-uid", checkpointOwner: "alice"} + service := NewService(store, testAuthorizer{}, tags, nil) + op, err := store.GetCheckpointOperation(t.Context(), cp.Id, "alice") + require.NoError(t, err) + require.Error(t, service.recoverOperation(t.Context(), op)) + require.Nil(t, tags.created) + // New owner sees the persisted phase and confirms absence rather than + // creating a replacement tag or discarding its reference prematurely. + service = NewService(store, testAuthorizer{}, tags, nil) + require.NoError(t, service.recoverOperation(t.Context(), op)) + if phase == apiv1.CheckpointState_CHECKPOINT_STATE_CREATING { + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_FAILED, cp.State) + } else { + require.True(t, store.deleted) + } + require.Equal(t, 1, tags.createCalls) + require.Equal(t, 1, tags.deleteCalls) + }) + } +} + +func TestCleanupRefusesReplacementTag(t *testing.T) { + cp := &apiv1.Checkpoint{Id: "11111111-1111-4111-8111-111111111111", SessionId: "22222222-2222-4222-8222-222222222222", State: apiv1.CheckpointState_CHECKPOINT_STATE_CREATING, + Failure: &apiv1.Failure{Reason: "SnapshotTagFailed", Message: "copy interrupted"}} + tags := &testTags{snapshotURI: "s3://snapshots/source"} + _, err := tags.CreateTag(t.Context(), "team-a", tagName(cp.Id), substrate.ActorName(cp.SessionId)) + require.NoError(t, err) + tags.created.Metadata.Uid = "replacement" + store := &testStore{prepared: cp, snapshot: &database.SessionTaskSnapshot{Atespace: "team-a"}, tagUID: "retained", checkpointOwner: "alice"} + service := NewService(store, testAuthorizer{}, tags, nil) + op, err := store.GetCheckpointOperation(t.Context(), cp.Id, "alice") + require.NoError(t, err) + require.Error(t, service.recoverOperation(t.Context(), op)) + require.Zero(t, tags.deleteCalls) + require.Equal(t, "replacement", tags.created.Metadata.Uid) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_CREATING, cp.State) +} + +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 := createAndRecover(service, 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) +} diff --git a/go/core/pkg/app/app.go b/go/core/pkg/app/app.go index 541109d3ec..8ba214b463 100644 --- a/go/core/pkg/app/app.go +++ b/go/core/pkg/app/app.go @@ -367,6 +367,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/000002_checkpoint_recovery.sql b/go/core/pkg/migrations/core/000002_checkpoint_recovery.sql new file mode 100644 index 0000000000..3af2993dfe --- /dev/null +++ b/go/core/pkg/migrations/core/000002_checkpoint_recovery.sql @@ -0,0 +1,6 @@ +-- +goose Up +CREATE INDEX session_checkpoint_pending_idx ON session_checkpoint (id) + WHERE state IN ('CREATING', 'DELETING'); + +-- +goose Down +DROP INDEX session_checkpoint_pending_idx; diff --git a/go/core/pkg/migrations/runner_test.go b/go/core/pkg/migrations/runner_test.go index 6cc9564c6d..8c15d7f17a 100644 --- a/go/core/pkg/migrations/runner_test.go +++ b/go/core/pkg/migrations/runner_test.go @@ -319,6 +319,9 @@ func TestBuiltinMigrationsRoundTrip(t *testing.T) { } for _, source := range sources { expected := []int64{0, 1} + if source.Name == "core" { + expected = append(expected, 2) + } 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/README.md b/go/core/test/e2e/README.md index bb20d5552e..9cab168475 100644 --- a/go/core/test/e2e/README.md +++ b/go/core/test/e2e/README.md @@ -126,6 +126,12 @@ checks also run in parallel. Trace cases share the receiver but select spans by their own trace IDs; the receiver is cleared once before starting those cases. Controller restart cases stay sequential with respect to the rest of the suite. +`TestCheckpointReplayAfterControllerRestart` loses a successful public checkpoint +response, replaces the controller once for the harness matrix, retries the same +request, and forks the retained snapshot to verify conversation continuity. +The response loss is caller-side after success; interrupted tag copying, +publication and compensation are exercised with PostgreSQL-backed workflow fault +tests. This test does not inject live Substrate packet loss. The scheduled-run restart test starts all selected harness executions before one shared controller restart, then releases their model responses and checks recovery in parallel. Each harness retains its own Session, task, and prompt-count assertions. 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..32384be423 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,36 @@ 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") + 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)) +} diff --git a/go/core/test/e2e/interaction_test.go b/go/core/test/e2e/interaction_test.go index bfa9173c96..cc8fdcd7f7 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 } From 2fad9c90444222965540720bb7a896439b6dafb1 Mon Sep 17 00:00:00 2001 From: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> Date: Thu, 5 Nov 2026 00:10:37 +0800 Subject: [PATCH 02/20] fix(checkpoint): fence recovery attempts across workers Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> --- .../persistence-checkpoints-and-forks.md | 9 + .../database/checkpoint_operations.go | 96 ++++++++++- .../database/checkpoint_operations_test.go | 118 ++++++++++++- go/core/internal/database/checkpoints.go | 50 +++++- .../internal/service/checkpoint/recovery.go | 46 +++-- .../service/checkpoint/recovery_test.go | 157 ++++++++++++++++-- .../internal/service/checkpoint/service.go | 12 +- .../service/checkpoint/service_test.go | 15 +- .../core/000002_checkpoint_recovery.sql | 11 ++ 9 files changed, 476 insertions(+), 38 deletions(-) diff --git a/docs/architecture/persistence-checkpoints-and-forks.md b/docs/architecture/persistence-checkpoints-and-forks.md index d8ffd3e9af..87a7e6b8a1 100644 --- a/docs/architecture/persistence-checkpoints-and-forks.md +++ b/docs/architecture/persistence-checkpoints-and-forks.md @@ -72,6 +72,15 @@ periodic scans resume `CREATING` and `DELETING` work without requiring the calle to return. The worker polls once per second so admission on another API replica does not wait for a caller-local signal; local API requests also wake it immediately. +Each attempt obtains a database execution claim before reading its current phase. +The claim and every backend call share a deadline of at most two minutes, starting +before claim acquisition. An overlapping worker cannot claim live work; after an +expired attempt, the replacement reads the persisted phase again. Cleanup, +publication, and deletion require the current executor identity under the row +lock, and a delayed release cannot revoke a replacement's claim. Leader election +reduces duplicate scans but is not the execution fence. A slow remote copy is +still synchronized through Substrate's tag lease after a client attempt expires. + A pending tag is not evidence of abandonment. Substrate serializes `CreateTag` against both the source Actor and the tag, including a copy still finishing after an earlier controller disappeared. Recovery calls `CreateTag` under those leases diff --git a/go/core/internal/database/checkpoint_operations.go b/go/core/internal/database/checkpoint_operations.go index c3e36c3523..c2a1caa4fe 100644 --- a/go/core/internal/database/checkpoint_operations.go +++ b/go/core/internal/database/checkpoint_operations.go @@ -2,8 +2,11 @@ package database import ( "context" + "errors" "fmt" + "time" + "github.com/google/uuid" "github.com/jackc/pgx/v5" "google.golang.org/protobuf/proto" @@ -20,6 +23,93 @@ type CheckpointOperation struct { TagUID string } +// ClaimCheckpointOperation grants one bounded attempt exclusive execution of an +// owner's pending checkpoint and returns its current phase, not a caller's earlier +// observation. The caller must establish its attempt deadline before claiming and +// use that context for every backend call. Missing, terminal, 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 + } + 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') + AND (executor_id IS NULL 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) + 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 { + // Lock by ID before testing the executor: an acquisition's COMMIT can + // still be in flight, with the old committed executor invisible to a + // conditional UPDATE. Wait for that transaction before deciding to no-op. + _, err := lockCheckpoint(ctx, tx, id, true, "", nil) + if errors.Is(err, pgx.ErrNoRows) { + return nil + } + if err != nil { + return err + } + _, err = tx.Exec(ctx, ` + UPDATE session_checkpoint SET executor_id = NULL, executor_expires_at = NULL + WHERE id = $1 AND executor_id = $2 + `, id, executor) + return err + }) +} + +// checkCheckpointExecutor requires the caller to hold the checkpoint row lock. +// Direct finalization is permitted only when no executor owns the row; +// claimed transitions reject expired and superseded attempts. +func checkCheckpointExecutor(ctx context.Context, tx pgx.Tx, id string, executor uuid.UUID) error { + var expected *uuid.UUID + if executor != uuid.Nil { + expected = &executor + } + matches, err := queryOne(ctx, tx, ` + SELECT executor_id IS NOT DISTINCT FROM $2::uuid + AND ($2::uuid IS NULL OR executor_expires_at > clock_timestamp()) + FROM session_checkpoint WHERE id = $1 + `, pgx.RowTo[bool], id, expected) + if err != nil { + return err + } + if !matches { + return ErrConflict + } + return nil +} + // 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. @@ -75,7 +165,8 @@ func checkpointOperation(row sessionCheckpointRow) (*CheckpointOperation, error) // The checkpoint remains CREATING, retaining history/revision references and // blocking source mutations until external cleanup finishes. Identical retries // succeed; a conflicting outcome or a terminal checkpoint returns ErrConflict. -func (c *Client) BeginCheckpointCleanup(ctx context.Context, id, tagUID, failure string) error { +// A non-nil executor must own the attempt; nil requires unclaimed work. +func (c *Client) BeginCheckpointCleanup(ctx context.Context, id string, executor uuid.UUID, tagUID, failure string) error { if failure == "" { return fmt.Errorf("checkpoint cleanup requires a failure") } @@ -84,6 +175,9 @@ func (c *Client) BeginCheckpointCleanup(ctx context.Context, id, tagUID, failure if err != nil { return notFoundOr(err) } + if err := checkCheckpointExecutor(ctx, tx, id, executor); err != nil { + return err + } checkpoint, err := toSessionCheckpoint(row) if err != nil { return err diff --git a/go/core/internal/database/checkpoint_operations_test.go b/go/core/internal/database/checkpoint_operations_test.go index ee9f933e02..34f411ee39 100644 --- a/go/core/internal/database/checkpoint_operations_test.go +++ b/go/core/internal/database/checkpoint_operations_test.go @@ -1,15 +1,127 @@ package database import ( + "context" + "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" 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 TestCheckpointExecutorFencesStaleTransitions(t *testing.T) { + client := NewClient(setupTestDB(t), "public") + checkpoint, _ := reserveRecoveryCheckpoint(t, client) + first := uuid.New() + ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond) + defer cancel() + _, err := client.ClaimCheckpointOperation(ctx, checkpoint.Id, "alice", first) + require.NoError(t, err) + otherCtx, cancelOther := context.WithTimeout(t.Context(), RuntimeOperationTimeout) + defer cancelOther() + _, err = client.ClaimCheckpointOperation(otherCtx, checkpoint.Id, "alice", uuid.New()) + require.ErrorIs(t, err, ErrConflict, "a live attempt must exclude another worker") + _, err = client.ClaimCheckpointOperation(otherCtx, checkpoint.Id, "mallory", uuid.New()) + require.ErrorIs(t, err, ErrConflict) + <-ctx.Done() + + second := uuid.New() + require.Eventually(t, func() bool { + _, err = client.ClaimCheckpointOperation(otherCtx, checkpoint.Id, "alice", second) + return err == nil + }, time.Second, 10*time.Millisecond) + // A late release and late writes must not revoke or complete the replacement. + require.NoError(t, client.ReleaseCheckpointOperation(t.Context(), checkpoint.Id, first)) + _, err = client.FinalizeClaimedSessionCheckpoint(t.Context(), checkpoint.Id, first, "old-tag", "s3://tags/old", "") + require.ErrorIs(t, err, ErrConflict) + require.ErrorIs(t, client.BeginCheckpointCleanup(t.Context(), checkpoint.Id, first, "old-tag", "old failure"), ErrConflict) + _, err = client.FinalizeSessionCheckpoint(t.Context(), checkpoint.Id, "unclaimed", "s3://tags/unclaimed", "") + require.ErrorIs(t, err, ErrConflict) + _, err = client.ClaimCheckpointOperation(otherCtx, checkpoint.Id, "alice", uuid.New()) + require.ErrorIs(t, err, ErrConflict, "the replacement must still own the attempt") + _, err = client.FinalizeClaimedSessionCheckpoint(otherCtx, checkpoint.Id, second, "new-tag", "s3://tags/new", "") + require.NoError(t, err) + require.NoError(t, client.ReleaseCheckpointOperation(t.Context(), checkpoint.Id, second)) + + _, _, err = client.BeginDeleteSessionCheckpoint(t.Context(), checkpoint.Id, "alice") + require.NoError(t, err) + deleting := uuid.New() + _, err = client.ClaimCheckpointOperation(otherCtx, checkpoint.Id, "alice", deleting) + require.NoError(t, err) + require.ErrorIs(t, client.DeleteClaimedSessionCheckpoint(t.Context(), checkpoint.Id, "alice", second), ErrConflict) + retained, err := client.GetCheckpointOperation(t.Context(), checkpoint.Id, "alice") + require.NoError(t, err) + require.Equal(t, "new-tag", retained.TagUID) + require.NoError(t, client.DeleteClaimedSessionCheckpoint(otherCtx, checkpoint.Id, "alice", deleting)) +} + func reserveRecoveryCheckpoint(t *testing.T, client *Client) (*apiv1.Checkpoint, *apiv1.Session) { t.Helper() ctx := t.Context() @@ -34,7 +146,7 @@ func TestCheckpointCleanupSurvivesRestartAndFencesPublication(t *testing.T) { client := NewClient(pool, "public") checkpoint, session := reserveRecoveryCheckpoint(t, client) ctx := t.Context() - require.NoError(t, client.BeginCheckpointCleanup(ctx, checkpoint.Id, "tag-uid", "copy interrupted")) + require.NoError(t, client.BeginCheckpointCleanup(ctx, checkpoint.Id, uuid.Nil, "tag-uid", "copy interrupted")) client = NewClient(pool, "public") op, err := client.GetCheckpointOperation(ctx, checkpoint.Id, "alice") require.NoError(t, err) @@ -50,8 +162,8 @@ func TestCheckpointCleanupSurvivesRestartAndFencesPublication(t *testing.T) { require.ErrorIs(t, err, ErrNotFound) _, err = client.FinalizeSessionCheckpoint(ctx, checkpoint.Id, "tag-uid", "s3://tags/copy", "") require.ErrorIs(t, err, ErrConflict, "cleanup must not race a stale successful copy into READY") - require.ErrorIs(t, client.BeginCheckpointCleanup(ctx, checkpoint.Id, "replacement", "copy interrupted"), ErrConflict) - require.NoError(t, client.BeginCheckpointCleanup(ctx, checkpoint.Id, "tag-uid", "copy interrupted")) + require.ErrorIs(t, client.BeginCheckpointCleanup(ctx, checkpoint.Id, uuid.Nil, "replacement", "copy interrupted"), ErrConflict) + require.NoError(t, client.BeginCheckpointCleanup(ctx, checkpoint.Id, uuid.Nil, "tag-uid", "copy interrupted")) _, err = client.FinalizeSessionCheckpoint(ctx, checkpoint.Id, "", "", "different failure") require.ErrorIs(t, err, ErrConflict) _, err = client.FinalizeSessionCheckpoint(ctx, checkpoint.Id, "", "", "copy interrupted") diff --git a/go/core/internal/database/checkpoints.go b/go/core/internal/database/checkpoints.go index 466f244546..23c5dc17c8 100644 --- a/go/core/internal/database/checkpoints.go +++ b/go/core/internal/database/checkpoints.go @@ -333,6 +333,19 @@ func (c *Client) ReserveSessionCheckpoint(ctx context.Context, checkpoint *apiv1 // 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) { + return c.finalizeSessionCheckpoint(ctx, id, uuid.Nil, tagUID, snapshotURI, failure) +} + +// FinalizeClaimedSessionCheckpoint publishes a creation outcome only while the +// given executor still owns its bounded attempt. Stale completion returns ErrConflict. +func (c *Client) FinalizeClaimedSessionCheckpoint(ctx context.Context, id string, executor uuid.UUID, tagUID, snapshotURI, failure string) (*apiv1.Checkpoint, error) { + if executor == uuid.Nil { + return nil, ErrFailedPrecondition + } + return c.finalizeSessionCheckpoint(ctx, id, executor, tagUID, snapshotURI, failure) +} + +func (c *Client) finalizeSessionCheckpoint(ctx context.Context, id string, executor uuid.UUID, 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") } @@ -342,6 +355,9 @@ func (c *Client) FinalizeSessionCheckpoint(ctx context.Context, id, tagUID, snap if err != nil { return notFoundOr(err) } + if err := checkCheckpointExecutor(ctx, tx, id, executor); err != nil { + return err + } result, err = toSessionCheckpoint(row) if err != nil { return err @@ -535,10 +551,36 @@ func (c *Client) BeginDeleteSessionCheckpoint(ctx context.Context, id, userID st // 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) + return c.deleteSessionCheckpoint(ctx, id, userID, uuid.Nil) +} + +// DeleteClaimedSessionCheckpoint removes an owned DELETING checkpoint only for +// its current executor, after external cleanup. A missing row is a successful retry. +func (c *Client) DeleteClaimedSessionCheckpoint(ctx context.Context, id, userID string, executor uuid.UUID) error { + if executor == uuid.Nil { + return ErrFailedPrecondition + } + return c.deleteSessionCheckpoint(ctx, id, userID, executor) +} + +func (c *Client) deleteSessionCheckpoint(ctx context.Context, id, userID string, executor uuid.UUID) error { + err := c.withTx(ctx, func(tx pgx.Tx) error { + _, err := lockCheckpoint(ctx, tx, id, false, userID, nil) + if errors.Is(err, pgx.ErrNoRows) { + return nil + } + if err != nil { + return err + } + if err := checkCheckpointExecutor(ctx, tx, id, executor); err != nil { + return err + } + _, err = tx.Exec(ctx, ` + DELETE FROM session_checkpoint + WHERE id = $1 AND user_id = $2 AND state = 'DELETING' + `, id, userID) + return err + }) if err != nil { return fmt.Errorf("delete Session checkpoint: %w", err) } diff --git a/go/core/internal/service/checkpoint/recovery.go b/go/core/internal/service/checkpoint/recovery.go index 0ac14d3df0..b577ba2ad0 100644 --- a/go/core/internal/service/checkpoint/recovery.go +++ b/go/core/internal/service/checkpoint/recovery.go @@ -6,6 +6,7 @@ import ( "fmt" "time" + "github.com/google/uuid" "golang.org/x/sync/errgroup" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" @@ -61,9 +62,7 @@ func (w *RecoveryWorker) sweep(ctx context.Context) error { group.SetLimit(4) for _, op := range operations { group.Go(func() error { - attempt, cancel := context.WithTimeout(ctx, 2*time.Minute) - defer cancel() - if err := w.service.recoverOperation(attempt, op); err != nil && ctx.Err() == nil { + if err := w.service.recoverOperation(ctx, op); err != nil && ctx.Err() == nil { logging.FromContext(ctx).ErrorContext(ctx, "recover checkpoint operation", "checkpoint_id", op.Checkpoint.Id, "error", err) } return nil @@ -80,17 +79,42 @@ func (w *RecoveryWorker) sweep(ctx context.Context) error { return ctx.Err() } -func (s *Service) recoverOperation(ctx context.Context, op *database.CheckpointOperation) error { +func (s *Service) recoverOperation(ctx context.Context, observed *database.CheckpointOperation) (err error) { + // Establish the attempt deadline before claiming. A paused executor cannot + // issue another RPC after its claim expires; unfinished copies synchronize + // through Substrate's tag lease when the replacement retries CreateTag. + ctx, cancel := context.WithTimeout(ctx, database.RuntimeOperationTimeout) + defer cancel() + executor := uuid.New() + id := observed.Checkpoint.Id + // The claim reply can be lost after the UPDATE committed. Release by the + // known executor even when acquisition returns an error; this cannot revoke + // another attempt and avoids waiting for expiry on graceful cancellation. + defer func() { + releaseCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + defer cancel() + err = errors.Join(err, s.store.ReleaseCheckpointOperation(releaseCtx, id, executor)) + }() + op, err := s.store.ClaimCheckpointOperation(ctx, id, observed.Creator, executor) + if errors.Is(err, database.ErrConflict) { + return nil // Another executor owns it, or the listed phase already ended. + } + if err != nil { + return err + } + if err := ctx.Err(); err != nil { + return err + } switch op.Checkpoint.State { case apiv1.CheckpointState_CHECKPOINT_STATE_DELETING: - return s.finishDelete(ctx, op) + return s.finishDelete(ctx, op, executor) case apiv1.CheckpointState_CHECKPOINT_STATE_CREATING: if op.Checkpoint.Failure != nil { - return s.finishCleanup(ctx, op) + return s.finishCleanup(ctx, op, executor) } tag, err := s.ensureTag(ctx, op.Checkpoint, op.Snapshot) if err == nil { - _, err = s.store.FinalizeSessionCheckpoint(ctx, op.Checkpoint.Id, + _, err = s.store.FinalizeClaimedSessionCheckpoint(ctx, op.Checkpoint.Id, executor, tag.GetMetadata().GetUid(), substrate.SnapshotURI(tag.GetStatus().GetSnapshot()), "") return err } @@ -101,7 +125,7 @@ func (s *Service) recoverOperation(ctx context.Context, op *database.CheckpointO return fmt.Errorf("refuse cleanup of an unverified checkpoint tag: %w", err) } uid := tag.GetMetadata().GetUid() - if recordErr := s.store.BeginCheckpointCleanup(ctx, op.Checkpoint.Id, uid, err.Error()); recordErr != nil { + if recordErr := s.store.BeginCheckpointCleanup(ctx, op.Checkpoint.Id, executor, uid, err.Error()); recordErr != nil { return recordErr } // Persist the exact phase before external cleanup. A restart must not @@ -110,13 +134,13 @@ func (s *Service) recoverOperation(ctx context.Context, op *database.CheckpointO if err != nil { return err } - return s.finishCleanup(ctx, op) + return s.finishCleanup(ctx, op, executor) default: return nil } } -func (s *Service) finishCleanup(ctx context.Context, op *database.CheckpointOperation) error { +func (s *Service) finishCleanup(ctx context.Context, op *database.CheckpointOperation, executor uuid.UUID) error { name := tagName(op.Checkpoint.Id) tag, err := s.tags.GetTag(ctx, op.Snapshot.Atespace, name) if err != nil && status.Code(err) != codes.NotFound { @@ -130,6 +154,6 @@ func (s *Service) finishCleanup(ctx context.Context, op *database.CheckpointOper return err } } - _, err = s.store.FinalizeSessionCheckpoint(ctx, op.Checkpoint.Id, "", "", op.Checkpoint.Failure.Message) + _, err = s.store.FinalizeClaimedSessionCheckpoint(ctx, op.Checkpoint.Id, executor, "", "", op.Checkpoint.Failure.Message) return err } diff --git a/go/core/internal/service/checkpoint/recovery_test.go b/go/core/internal/service/checkpoint/recovery_test.go index a605c511fb..9fb972977d 100644 --- a/go/core/internal/service/checkpoint/recovery_test.go +++ b/go/core/internal/service/checkpoint/recovery_test.go @@ -4,6 +4,7 @@ import ( "context" "crypto/sha256" "errors" + "strings" "sync" "testing" "time" @@ -11,7 +12,6 @@ import ( "github.com/a2aproject/a2a-go/v2/a2a" a2apb "github.com/a2aproject/a2a-go/v2/a2apb/v1" "github.com/google/uuid" - "github.com/jackc/pgx/v5/pgxpool" "github.com/stretchr/testify/require" ateapipb "go.ate.dev/substrate/api/ateapi/v1" "google.golang.org/grpc/codes" @@ -21,6 +21,7 @@ import ( 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/substrate" "github.com/kagent-dev/kagent/go/core/pkg/auth" ) @@ -87,18 +88,37 @@ func startRecovery(t *testing.T, service *Service) func() { } 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-" + actor.GetMetadata().GetUid() + sessionID := strings.TrimPrefix(actor.GetMetadata().GetName(), "session-") + _, err := uuid.Parse(sessionID) + require.NoError(t, err) require.NoError(t, store.UpsertAgentDefinition(ctx, database.AgentDefinition{ - Namespace: "team-a", AgentName: "assistant", AgentUID: "agent-uid", DesiredRevision: "revision"})) - _, err := store.RecordRuntimeRevision(ctx, database.RuntimeRevision{ - Revision: "revision", Namespace: "team-a", AgentName: "assistant", AgentUID: "agent-uid", + 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: "team-a", ActorTemplateName: "runtime", ActorTemplateUID: "template-uid", + 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: uuid.NewString(), Creator: "alice", - Agent: &apiv1.AgentProvenance{Ref: &apiv1.ResourceReference{Namespace: "team-a", Name: "assistant"}}}, uuid.NewString()) + 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) require.NoError(t, err) @@ -106,7 +126,7 @@ func recoverySource(t *testing.T, store *database.Client) (*apiv1.Session, strin 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, "runtime", "actor-uid", "") + 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 @@ -119,7 +139,7 @@ func recoverySource(t *testing.T, store *database.Client) (*apiv1.Session, strin boundary, err := store.ClaimSessionQuiescence(ctx) require.NoError(t, err) require.NoError(t, store.FinishSessionQuiescence(ctx, boundary, - &database.SessionTaskSnapshot{Atespace: "team-a", URI: "s3://snapshots/source", Fidelity: "SNAPSHOT_FIDELITY_VOLUMES"})) + &database.SessionTaskSnapshot{Atespace: space, URI: substrate.SnapshotURI(snapshot), Fidelity: snapshot.GetFidelity().String()})) return session, string(task.ID) } @@ -140,19 +160,46 @@ type failedCheckpointPublish struct { once sync.Once } -func (s *failedCheckpointPublish) FinalizeSessionCheckpoint(ctx context.Context, _, _, _, _ string) (*apiv1.Checkpoint, error) { +func (s *failedCheckpointPublish) FinalizeClaimedSessionCheckpoint(ctx context.Context, _ string, _ uuid.UUID, _, _, _ string) (*apiv1.Checkpoint, error) { s.once.Do(func() { close(s.entered) }) <-ctx.Done() return nil, ctx.Err() } +// 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 +} + func TestCheckpointRecoveryWithPostgres(t *testing.T) { if testing.Short() { t.Skip("requires PostgreSQL") } conn := dbtest.StartT(t.Context(), t) require.NoError(t, dbtest.Migrate(conn, false)) - pool, err := pgxpool.New(t.Context(), conn) + pool, err := database.Connect(t.Context(), &database.PostgresConfig{Pool: dbtest.PoolConfig(conn), Schema: "public"}) require.NoError(t, err) t.Cleanup(pool.Close) newStore := func() *database.Client { return database.NewClient(pool, "public") } @@ -207,6 +254,7 @@ func TestCheckpointRecoveryWithPostgres(t *testing.T) { require.ErrorIs(t, <-result, context.Canceled) // Another replica shares durable admission, not local singleflight. follower := NewService(newStore(), testAuthorizer{}, tags, nil) + stopFollower := startRecovery(t, follower) results := make(chan *apiv1.Checkpoint, 8) errors := make(chan error, 8) for range 8 { @@ -226,6 +274,7 @@ func TestCheckpointRecoveryWithPostgres(t *testing.T) { require.NoError(t, <-errors) require.Equal(t, pending[0].Checkpoint.Id, (<-results).Id) } + stopFollower() stop() tags.mu.Lock() defer tags.mu.Unlock() @@ -297,6 +346,92 @@ func TestCheckpointRecoveryWithPostgres(t *testing.T) { require.NoError(t, err) }) + t.Run("Given compensation already deleted the tag, when stale creation resumes, assert that source admission reopens without another copy", func(t *testing.T) { + store := newStore() + session, head := recoverySource(t, store) + tags := &testTags{snapshotURI: "s3://snapshots/source"} + service := NewService(store, testAuthorizer{}, tags, nil) + cp, _, err := service.beginCreate(ownerContext(t.Context()), session.Id, uuid.NewString(), head) + require.NoError(t, err) + stale, err := store.GetCheckpointOperation(t.Context(), cp.Id, "alice") + require.NoError(t, err) + tag, err := tags.CreateTag(t.Context(), "team-a", tagName(cp.Id), substrate.ActorName(session.Id)) + require.NoError(t, err) + ctx, cancel := context.WithTimeout(t.Context(), database.RuntimeOperationTimeout) + defer cancel() + executor := uuid.New() + _, err = store.ClaimCheckpointOperation(ctx, cp.Id, "alice", executor) + require.NoError(t, err) + require.NoError(t, store.BeginCheckpointCleanup(ctx, cp.Id, executor, tag.Metadata.Uid, "copy interrupted")) + require.NoError(t, tags.DeleteTag(ctx, "team-a", tagName(cp.Id), tag.Metadata.Uid)) + require.NoError(t, store.ReleaseCheckpointOperation(ctx, cp.Id, executor)) + + require.NoError(t, service.recoverOperation(t.Context(), stale)) + require.Nil(t, tags.created, "a stale discovery result must not recreate a compensated tag") + require.Equal(t, 1, tags.createCalls) + waitCheckpointPhase(t, store, cp.Id, apiv1.CheckpointState_CHECKPOINT_STATE_FAILED) + _, err = store.BeginSessionOperation(t.Context(), session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE) + require.NoError(t, err) + _, err = service.beginDelete(ownerContext(t.Context()), cp.Id) + require.NoError(t, err) + require.NoError(t, service.recoverOperation(t.Context(), stale)) + require.NoError(t, service.recoverOperation(t.Context(), stale), "deleted work must not recreate an orphan tag") + _, err = store.GetCheckpointOperation(t.Context(), cp.Id, "alice") + require.ErrorIs(t, err, database.ErrNotFound) + require.Nil(t, tags.created) + require.Equal(t, 1, tags.createCalls) + }) + + t.Run("Given a paused executor, when its attempt expires and recovery succeeds, assert that its late continuation cannot create a tag", func(t *testing.T) { + store := newStore() + session, head := recoverySource(t, store) + tags := &testTags{snapshotURI: "s3://snapshots/source"} + resume := make(chan struct{}) + paused := &pausedCheckpointClaim{Client: store, claimed: make(chan struct{}), resume: resume} + old := NewService(paused, testAuthorizer{}, tags, nil) + cp, _, err := old.beginCreate(ownerContext(t.Context()), session.Id, uuid.NewString(), head) + require.NoError(t, err) + op, err := store.GetCheckpointOperation(t.Context(), cp.Id, "alice") + require.NoError(t, err) + ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond) + defer cancel() + done := make(chan error, 1) + go func() { done <- old.recoverOperation(ctx, op) }() + select { + case <-paused.claimed: + case <-time.After(time.Second): + t.Fatal("old worker never claimed the operation") + } + <-ctx.Done() + replacement := NewService(newStore(), testAuthorizer{}, tags, nil) + require.Eventually(t, func() bool { + require.NoError(t, replacement.recoverOperation(t.Context(), op)) + current, err := store.GetCheckpointOperation(t.Context(), cp.Id, "alice") + return err == nil && current.Checkpoint.State == apiv1.CheckpointState_CHECKPOINT_STATE_READY + }, time.Second, 10*time.Millisecond) + close(resume) + require.ErrorIs(t, <-done, context.DeadlineExceeded) + require.Equal(t, 1, tags.createCalls, "the expired worker must issue no backend mutation") + require.Zero(t, tags.deleteCalls) + }) + + t.Run("Given a lost committed claim reply, when a replacement retries, assert that it can immediately resume", func(t *testing.T) { + store := newStore() + session, head := recoverySource(t, store) + tags := &testTags{snapshotURI: "s3://snapshots/source"} + service := NewService(&lostCheckpointClaim{Client: store}, testAuthorizer{}, tags, nil) + cp, _, err := service.beginCreate(ownerContext(t.Context()), session.Id, uuid.NewString(), head) + require.NoError(t, err) + op, err := store.GetCheckpointOperation(t.Context(), cp.Id, "alice") + require.NoError(t, err) + require.Error(t, service.recoverOperation(t.Context(), op)) + require.Zero(t, tags.createCalls) + replacement := NewService(newStore(), testAuthorizer{}, tags, nil) + require.NoError(t, replacement.recoverOperation(t.Context(), op)) + waitCheckpointPhase(t, store, cp.Id, apiv1.CheckpointState_CHECKPOINT_STATE_READY) + require.Equal(t, 1, tags.createCalls) + }) + t.Run("restart finishes deletion after backend failure", func(t *testing.T) { store := newStore() session, head := recoverySource(t, store) diff --git a/go/core/internal/service/checkpoint/service.go b/go/core/internal/service/checkpoint/service.go index ad4be9a968..c777b34445 100644 --- a/go/core/internal/service/checkpoint/service.go +++ b/go/core/internal/service/checkpoint/service.go @@ -29,14 +29,16 @@ const ( type store interface { GetCheckpointOperation(context.Context, string, string) (*database.CheckpointOperation, error) ListPendingCheckpoints(context.Context, string, int) ([]*database.CheckpointOperation, error) - BeginCheckpointCleanup(context.Context, string, string, string) error + ClaimCheckpointOperation(context.Context, string, string, uuid.UUID) (*database.CheckpointOperation, error) + ReleaseCheckpointOperation(context.Context, string, uuid.UUID) error + BeginCheckpointCleanup(context.Context, string, uuid.UUID, string, string) error ReserveSessionCheckpoint(context.Context, *apiv1.Checkpoint, string, string) (*apiv1.Checkpoint, *database.SessionTaskSnapshot, error) - FinalizeSessionCheckpoint(context.Context, string, string, string, string) (*apiv1.Checkpoint, error) + FinalizeClaimedSessionCheckpoint(context.Context, string, 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 + 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) } @@ -294,7 +296,7 @@ func (s *Service) beginDelete(ctx context.Context, checkpointID string) (string, return userID, nil } -func (s *Service) finishDelete(ctx context.Context, op *database.CheckpointOperation) error { +func (s *Service) finishDelete(ctx context.Context, op *database.CheckpointOperation, executor uuid.UUID) error { checkpointID, userID := op.Checkpoint.Id, op.Creator snapshot, tagUID := op.Snapshot, op.TagUID tag, err := s.tags.GetTag(ctx, snapshot.Atespace, tagName(checkpointID)) @@ -313,7 +315,7 @@ func (s *Service) finishDelete(ctx context.Context, op *database.CheckpointOpera return serviceerrors.NewUnavailable("Failed to delete checkpoint snapshot tag", err) } } - if err := s.store.DeleteSessionCheckpoint(ctx, checkpointID, userID); err != nil { + if err := s.store.DeleteClaimedSessionCheckpoint(ctx, checkpointID, userID, executor); err != nil { return serviceerrors.NewInternal("Failed to delete checkpoint", err) } return nil diff --git a/go/core/internal/service/checkpoint/service_test.go b/go/core/internal/service/checkpoint/service_test.go index e7e674a6e3..b89d1f88e6 100644 --- a/go/core/internal/service/checkpoint/service_test.go +++ b/go/core/internal/service/checkpoint/service_test.go @@ -9,6 +9,7 @@ import ( "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" @@ -82,7 +83,7 @@ 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, _ string, _ uuid.UUID, tagUID, snapshotURI, failure string) (*apiv1.Checkpoint, error) { if s.finalizeErr != nil { return nil, s.finalizeErr } @@ -130,7 +131,7 @@ func (s *testStore) BeginDeleteSessionCheckpoint(context.Context, string, string return s.snapshot, s.tagUID, nil } -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 } @@ -859,7 +860,7 @@ func (*testStore) ListPendingCheckpoints(context.Context, string, int) ([]*datab return nil, nil } -func (s *testStore) BeginCheckpointCleanup(_ context.Context, _, uid, failure string) error { +func (s *testStore) BeginCheckpointCleanup(_ context.Context, _ string, _ uuid.UUID, uid, failure string) error { s.prepared.Failure = &apiv1.Failure{Reason: "SnapshotTagFailed", Message: failure} s.tagUID = uid return nil @@ -955,3 +956,11 @@ func TestCreationRefusesUnrelatedTag(t *testing.T) { 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 +} diff --git a/go/core/pkg/migrations/core/000002_checkpoint_recovery.sql b/go/core/pkg/migrations/core/000002_checkpoint_recovery.sql index 3af2993dfe..b65e29882d 100644 --- a/go/core/pkg/migrations/core/000002_checkpoint_recovery.sql +++ b/go/core/pkg/migrations/core/000002_checkpoint_recovery.sql @@ -1,6 +1,17 @@ -- +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) + ); + CREATE INDEX session_checkpoint_pending_idx ON session_checkpoint (id) WHERE state IN ('CREATING', 'DELETING'); -- +goose Down DROP INDEX session_checkpoint_pending_idx; +ALTER TABLE session_checkpoint + DROP CONSTRAINT session_checkpoint_executor_pair, + DROP COLUMN executor_expires_at, + DROP COLUMN executor_id; From 329d7014ac7f3130599ffb32c054f3ad0235fc5a Mon Sep 17 00:00:00 2001 From: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> Date: Thu, 5 Nov 2026 00:19:49 +0800 Subject: [PATCH 03/20] test(checkpoint): exercise recovery with real Substrate Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> --- .github/workflows/ci.yaml | 3 + .../checkpoint/recovery_substrate_test.go | 325 ++++++++++++++++++ .../service/checkpoint/recovery_test.go | 4 +- go/core/test/e2e/README.md | 16 + go/core/test/e2e/checkpoint-recovery.sh | 45 +++ 5 files changed, 391 insertions(+), 2 deletions(-) create mode 100644 go/core/internal/service/checkpoint/recovery_substrate_test.go create mode 100644 go/core/test/e2e/checkpoint-recovery.sh diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index e6ea2cecea..54d38946b8 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -259,6 +259,9 @@ jobs: sleep 5 done test "$READY" = True + # Exercise recovery with real Substrate and PostgreSQL, including lost + # tag responses and compensation interrupted before deletion. + bash core/test/e2e/checkpoint-recovery.sh export KAGENT_E2E_API_URL="http://$(kubectl get svc -n kagent kagent-controller -o jsonpath='{.status.loadBalancer.ingress[0].ip}'):8083" echo "KAGENT_E2E_API_URL: $KAGENT_E2E_API_URL" make core/bin/kagent-local diff --git a/go/core/internal/service/checkpoint/recovery_substrate_test.go b/go/core/internal/service/checkpoint/recovery_substrate_test.go new file mode 100644 index 0000000000..2d9eafd611 --- /dev/null +++ b/go/core/internal/service/checkpoint/recovery_substrate_test.go @@ -0,0 +1,325 @@ +package checkpoint + +import ( + "context" + "crypto/tls" + "crypto/x509" + "encoding/json" + "os" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + ateapipb "go.ate.dev/substrate/api/ateapi/v1" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/credentials" + "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/substrate" + "github.com/kagent-dev/kagent/go/core/pkg/auth" +) + +// This suite uses real Actors, snapshot copies, UID-conditional deletion and +// PostgreSQL. Faults discard actual RPC results or prevent one RPC from being +// delivered; no backend response or persisted lifecycle state is fabricated. +// The configuration is test-only; CI runs this after installing Substrate. +func TestCheckpointRecoveryWithSubstrate(t *testing.T) { + path := os.Getenv("KAGENT_E2E_CHECKPOINT_SUBSTRATE_CONFIG") + if path == "" { + t.Skip("requires real Substrate; run core/test/e2e/checkpoint-recovery.sh") + } + var config struct { + Endpoint string `json:"endpoint"` + ServerName string `json:"serverName"` + CAFile string `json:"caFile"` + ClientCertFile string `json:"clientCertFile"` + Atespace string `json:"atespace"` + Image string `json:"image"` + } + data, err := os.ReadFile(path) + require.NoError(t, err) + require.NoError(t, json.Unmarshal(data, &config)) + require.NotEmpty(t, config.Endpoint) + require.NotEmpty(t, config.Atespace) + require.NotEmpty(t, config.Image) + certificate, err := tls.LoadX509KeyPair(config.ClientCertFile, config.ClientCertFile) + require.NoError(t, err) + ca, err := os.ReadFile(config.CAFile) + require.NoError(t, err) + roots := x509.NewCertPool() + require.True(t, roots.AppendCertsFromPEM(ca)) + conn, err := grpc.NewClient(config.Endpoint, grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{ + MinVersion: tls.VersionTLS12, RootCAs: roots, Certificates: []tls.Certificate{certificate}, ServerName: config.ServerName, + }))) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, conn.Close()) }) + client := &substrate.Client{ControlClient: ateapipb.NewControlClient(conn)} + ctx, cancel := context.WithTimeout(t.Context(), 8*time.Minute) + t.Cleanup(cancel) + template, err := client.CreateActorTemplate(ctx, &ateapipb.ActorTemplate{ + Metadata: &ateapipb.ResourceMetadata{Atespace: config.Atespace, Name: "checkpoint-recovery-" + uuid.NewString()}, + Resources: &ateapipb.ResourceRequirements{Limits: &ateapipb.Resources{CpuMillicores: 500, MemoryBytes: 256 << 20}}, + Containers: []*ateapipb.Container{{Name: "source", Image: config.Image, + Command: []string{"/bin/sh", "-c", "echo checkpoint-recovery > /data/proof; exec sleep 3600"}, + VolumeMounts: []*ateapipb.VolumeMount{{Name: "data", MountPath: "/data"}}, + }}, + Volumes: []*ateapipb.Volume{{Name: "data", DurableDir: &ateapipb.DurableDirVolumeSource{}}}, + SnapshotConfig: &ateapipb.SnapshotConfig{PreferredFidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES, + GoldenSnapshotConfig: &ateapipb.GoldenSnapshotConfig{Mode: ateapipb.GoldenSnapshotMode_GOLDEN_SNAPSHOT_MODE_DISABLED}}, + }) + require.NoError(t, err) + t.Cleanup(func() { + cleanup, stop := context.WithTimeout(context.Background(), time.Minute) + defer stop() + require.NoError(t, client.DeleteActorTemplate(cleanup, config.Atespace, template.Metadata.Name, template.Metadata.Uid)) + }) + 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) + newStore := func() *database.Client { return database.NewClient(pool, "public") } + owner := auth.AuthSessionTo(ctx, testSession{userID: "alice"}) + newSource := func(t *testing.T) (*database.Client, *apiv1.Session, string) { + t.Helper() + name := substrate.ActorName(uuid.NewString()) + _, err := client.CreateActor(ctx, config.Atespace, name, config.Atespace, template.Metadata.Name) + require.NoError(t, err) + t.Cleanup(func() { + cleanup, stop := context.WithTimeout(context.Background(), time.Minute) + defer stop() + err := client.DeleteActor(cleanup, config.Atespace, name) + if status.Code(err) != codes.NotFound { + require.NoError(t, err) + } + }) + _, err = client.ResumeActor(ctx, config.Atespace, name) + require.NoError(t, err) + actor, err := client.SuspendActor(ctx, config.Atespace, name) + require.NoError(t, err) + require.NotEmpty(t, substrate.SnapshotURI(substrate.LatestActorSnapshot(actor))) + store := newStore() + session, head := recoverySourceForActor(t, store, actor) + return store, session, head + } + deleteCheckpoint := func(t *testing.T, service *Service, id string) { + t.Helper() + stop := startRecovery(t, service) + require.NoError(t, service.Delete(owner, id)) + stop() + _, err := service.store.GetCheckpointOperation(ctx, id, "alice") + require.ErrorIs(t, err, database.ErrNotFound) + _, err = client.GetTag(ctx, config.Atespace, tagName(id)) + require.Equal(t, codes.NotFound, status.Code(err)) + } + cleanupTag := func(t *testing.T, id string) { + t.Helper() + t.Cleanup(func() { + cleanup, stop := context.WithTimeout(context.Background(), time.Minute) + defer stop() + tag, err := client.GetTag(cleanup, config.Atespace, tagName(id)) + if status.Code(err) == codes.NotFound { + return + } + require.NoError(t, err) + require.NoError(t, client.DeleteTag(cleanup, config.Atespace, tagName(id), tag.GetMetadata().GetUid())) + }) + } + + t.Run("Given an abandoned reservation, when a replacement starts, assert that it publishes a real retained copy", func(t *testing.T) { + store, source, head := newSource(t) + service := NewService(store, testAuthorizer{}, client, nil) + request := uuid.NewString() + cp, _, err := service.beginCreate(owner, source.Id, request, head) + require.NoError(t, err) + cleanupTag(t, cp.Id) + replacement := NewService(newStore(), testAuthorizer{}, client, nil) + stop := startRecovery(t, replacement) + ready, err := replacement.Create(owner, source.Id, request, head) + require.NoError(t, err) + stop() + require.Equal(t, cp.Id, ready.Id) + tag, err := client.GetTag(ctx, config.Atespace, tagName(cp.Id)) + require.NoError(t, err) + require.NotEmpty(t, substrate.SnapshotURI(tag.GetStatus().GetSnapshot())) + clone := substrate.ActorName(uuid.NewString()) + _, err = client.CreateActorFromTag(ctx, config.Atespace, clone, config.Atespace, template.Metadata.Name, config.Atespace, tagName(cp.Id)) + require.NoError(t, err) + t.Cleanup(func() { + cleanup, stop := context.WithTimeout(context.Background(), time.Minute) + defer stop() + if err := client.DeleteActor(cleanup, config.Atespace, clone); status.Code(err) != codes.NotFound { + require.NoError(t, err) + } + }) + _, err = client.ResumeActor(ctx, config.Atespace, clone) + require.NoError(t, err, "the recovered tag must contain a usable snapshot") + require.NoError(t, client.DeleteActor(ctx, config.Atespace, clone)) + deleteCheckpoint(t, replacement, cp.Id) + }) + + t.Run("Given a lost successful tag response, when recovery restarts, assert that it adopts the same backend UID", func(t *testing.T) { + store, source, head := newSource(t) + lost := &faultedSubstrateTags{Client: client, loseCreateResponse: true} + service := NewService(store, testAuthorizer{}, lost, nil) + request := uuid.NewString() + cp, _, err := service.beginCreate(owner, source.Id, request, head) + require.NoError(t, err) + cleanupTag(t, cp.Id) + op, err := store.GetCheckpointOperation(ctx, cp.Id, "alice") + require.NoError(t, err) + require.Error(t, service.recoverOperation(ctx, op)) + tag, err := client.GetTag(ctx, config.Atespace, tagName(cp.Id)) + require.NoError(t, err) + pending, err := store.GetCheckpointOperation(ctx, cp.Id, "alice") + require.NoError(t, err) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_CREATING, pending.Checkpoint.State) + replacement := NewService(newStore(), testAuthorizer{}, client, nil) + require.NoError(t, replacement.recoverOperation(ctx, pending)) + published, err := store.GetCheckpointOperation(ctx, cp.Id, "alice") + require.NoError(t, err) + require.Equal(t, tag.Metadata.Uid, published.TagUID) + deleteCheckpoint(t, replacement, cp.Id) + }) + + t.Run("Given a completed copy without database publication, when the worker restarts, assert that the original tag is published", func(t *testing.T) { + store, source, head := newSource(t) + fault := &failedCheckpointPublish{Client: store, entered: make(chan struct{})} + old := NewService(fault, testAuthorizer{}, client, nil) + request := uuid.NewString() + cp, _, err := old.beginCreate(owner, source.Id, request, head) + require.NoError(t, err) + cleanupTag(t, cp.Id) + op, err := store.GetCheckpointOperation(ctx, cp.Id, "alice") + require.NoError(t, err) + attempt, cancelAttempt := context.WithCancel(ctx) + defer cancelAttempt() + done := make(chan error, 1) + go func() { done <- old.recoverOperation(attempt, op) }() + select { + case <-fault.entered: + case err := <-done: + t.Fatalf("copy failed before reaching database publication: %v", err) + case <-ctx.Done(): + t.Fatal(ctx.Err()) + } + tag, err := client.GetTag(ctx, config.Atespace, tagName(cp.Id)) + require.NoError(t, err) + cancelAttempt() + require.ErrorIs(t, <-done, context.Canceled) + replacement := NewService(newStore(), testAuthorizer{}, client, nil) + require.NoError(t, replacement.recoverOperation(ctx, op)) + ready, err := store.GetCheckpointOperation(ctx, cp.Id, "alice") + require.NoError(t, err) + require.Equal(t, tag.Metadata.Uid, ready.TagUID) + deleteCheckpoint(t, replacement, cp.Id) + }) + + t.Run("Given a source deleted after reservation, when tag creation fails, assert that failed checkpoint deletion releases its row", func(t *testing.T) { + store, source, head := newSource(t) + service := NewService(store, testAuthorizer{}, client, nil) + cp, _, err := service.beginCreate(owner, source.Id, uuid.NewString(), head) + require.NoError(t, err) + cleanupTag(t, cp.Id) + require.NoError(t, client.DeleteActor(ctx, config.Atespace, substrate.ActorName(source.Id))) + op, err := store.GetCheckpointOperation(ctx, cp.Id, "alice") + require.NoError(t, err) + require.NoError(t, service.recoverOperation(ctx, op)) + waitCheckpointPhase(t, store, cp.Id, apiv1.CheckpointState_CHECKPOINT_STATE_FAILED) + _, err = store.BeginSessionOperation(ctx, source.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE) + require.NoError(t, err, "failed creation must stop blocking the source") + deleteCheckpoint(t, NewService(newStore(), testAuthorizer{}, client, nil), cp.Id) + }) + + t.Run("Given failed compensation, when recovery restarts, assert that cleanup precedes reopening the source", func(t *testing.T) { + store, source, head := newSource(t) + fault := &faultedSubstrateTags{Client: client, failDelete: true, afterCreate: func() error { + _, err := client.ResumeActor(ctx, config.Atespace, substrate.ActorName(source.Id)) + return err + }} + service := NewService(store, testAuthorizer{}, fault, nil) + cp, _, err := service.beginCreate(owner, source.Id, uuid.NewString(), head) + require.NoError(t, err) + cleanupTag(t, cp.Id) + op, err := store.GetCheckpointOperation(ctx, cp.Id, "alice") + require.NoError(t, err) + require.Error(t, service.recoverOperation(ctx, op)) + pending, err := store.GetCheckpointOperation(ctx, cp.Id, "alice") + require.NoError(t, err) + require.NotNil(t, pending.Checkpoint.Failure) + require.NotEmpty(t, pending.TagUID) + _, err = store.BeginSessionOperation(ctx, source.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE) + require.ErrorIs(t, err, database.ErrConflict) + replacement := NewService(newStore(), testAuthorizer{}, client, nil) + require.NoError(t, replacement.recoverOperation(ctx, op), "a stale discovery must resume persisted compensation") + waitCheckpointPhase(t, store, cp.Id, apiv1.CheckpointState_CHECKPOINT_STATE_FAILED) + _, err = client.GetTag(ctx, config.Atespace, tagName(cp.Id)) + require.Equal(t, codes.NotFound, status.Code(err)) + _, err = store.BeginSessionOperation(ctx, source.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE) + require.NoError(t, err) + deleteCheckpoint(t, replacement, cp.Id) + }) + + t.Run("Given a lost delete response, when recovery restarts, assert that an absent real tag completes deletion", func(t *testing.T) { + store, source, head := newSource(t) + service := NewService(store, testAuthorizer{}, client, nil) + stop := startRecovery(t, service) + cp, err := service.Create(owner, source.Id, uuid.NewString(), head) + require.NoError(t, err) + cleanupTag(t, cp.Id) + stop() + _, err = service.beginDelete(owner, cp.Id) + require.NoError(t, err) + op, err := store.GetCheckpointOperation(ctx, cp.Id, "alice") + require.NoError(t, err) + lost := NewService(store, testAuthorizer{}, &faultedSubstrateTags{Client: client, loseDeleteResponse: true}, nil) + require.Error(t, lost.recoverOperation(ctx, op)) + _, err = client.GetTag(ctx, config.Atespace, tagName(cp.Id)) + require.Equal(t, codes.NotFound, status.Code(err)) + deleteCheckpoint(t, NewService(newStore(), testAuthorizer{}, client, nil), cp.Id) + }) +} + +type faultedSubstrateTags struct { + *substrate.Client + loseCreateResponse bool + loseDeleteResponse bool + failDelete bool + afterCreate func() error +} + +func (f *faultedSubstrateTags) CreateTag(ctx context.Context, space, name, actor string) (*ateapipb.Tag, error) { + tag, err := f.Client.CreateTag(ctx, space, name, actor) + if err != nil { + return nil, err + } + if f.afterCreate != nil { + if err := f.afterCreate(); err != nil { + return nil, err + } + } + if f.loseCreateResponse { + return nil, status.Error(codes.Unavailable, "injected loss of successful CreateTag response") + } + return tag, nil +} + +func (f *faultedSubstrateTags) DeleteTag(ctx context.Context, space, name, uid string) error { + if f.failDelete { + return status.Error(codes.Unavailable, "injected failure to deliver DeleteTag") + } + if err := f.Client.DeleteTag(ctx, space, name, uid); err != nil { + return err + } + if f.loseDeleteResponse { + return status.Error(codes.Unavailable, "injected loss of successful DeleteTag response") + } + return nil +} + +var _ tagClient = (*faultedSubstrateTags)(nil) diff --git a/go/core/internal/service/checkpoint/recovery_test.go b/go/core/internal/service/checkpoint/recovery_test.go index 9fb972977d..8c79ddf8e5 100644 --- a/go/core/internal/service/checkpoint/recovery_test.go +++ b/go/core/internal/service/checkpoint/recovery_test.go @@ -105,7 +105,7 @@ func recoverySourceForActor(t *testing.T, store *database.Client, actor *ateapip ctx := t.Context() space := actor.GetMetadata().GetAtespace() snapshot := substrate.LatestActorSnapshot(actor) - revision := "revision-" + actor.GetMetadata().GetUid() + revision := "revision-" + snapshot.GetActorTemplateUid() sessionID := strings.TrimPrefix(actor.GetMetadata().GetName(), "session-") _, err := uuid.Parse(sessionID) require.NoError(t, err) @@ -197,7 +197,7 @@ func TestCheckpointRecoveryWithPostgres(t *testing.T) { if testing.Short() { t.Skip("requires PostgreSQL") } - conn := dbtest.StartT(t.Context(), t) + conn := dbtest.StartT(context.WithoutCancel(t.Context()), t) require.NoError(t, dbtest.Migrate(conn, false)) pool, err := database.Connect(t.Context(), &database.PostgresConfig{Pool: dbtest.PoolConfig(conn), Schema: "public"}) require.NoError(t, err) diff --git a/go/core/test/e2e/README.md b/go/core/test/e2e/README.md index 9cab168475..db94876051 100644 --- a/go/core/test/e2e/README.md +++ b/go/core/test/e2e/README.md @@ -132,6 +132,22 @@ request, and forks the retained snapshot to verify conversation continuity. The response loss is caller-side after success; interrupted tag copying, publication and compensation are exercised with PostgreSQL-backed workflow fault tests. This test does not inject live Substrate packet loss. + +`bash core/test/e2e/checkpoint-recovery.sh` (from `go/`) runs the checkpoint +lifecycle service against real PostgreSQL and the installed Substrate. Set +`KAGENT_E2E_RUNTIME_IMAGE` to the digest-pinned Go ADK image. The script reads the +test controller's projected identity into temporary private files and forwards +the Substrate service; it leaves controller configuration unchanged. Both CI +runtime jobs run it before the public API E2E suite, and a configured environment +must fail rather than skip on missing prerequisites. It owns its ActorTemplate, +Actors and Tags. Cases cover abandoned reservations, lost successful tag replies, +copy completion before database publication, +a source deleted before tag creation, failed compensation after a real source +change, and lost successful deletion replies. Transport faults are injected at +the service's client boundary around real calls; backend responses and database +phases are not fabricated. The deterministic PostgreSQL tests separately cover +pending-copy leases and overlapping/expired executors. These are not claims of +live object-store outage or mid-copy packet-loss coverage. The scheduled-run restart test starts all selected harness executions before one shared controller restart, then releases their model responses and checks recovery in parallel. Each harness retains its own Session, task, and prompt-count assertions. diff --git a/go/core/test/e2e/checkpoint-recovery.sh b/go/core/test/e2e/checkpoint-recovery.sh new file mode 100644 index 0000000000..2606b6a16b --- /dev/null +++ b/go/core/test/e2e/checkpoint-recovery.sh @@ -0,0 +1,45 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Run from any directory against the disposable E2E installation. Reuse its +# controller identity; private certificate material never enters test output. +cd "$(dirname "${BASH_SOURCE[0]}")/../../.." +: "${KAGENT_E2E_RUNTIME_IMAGE:?set the digest-pinned Go ADK image}" +checkpoint_tmp=$(mktemp -d) +chmod 700 "$checkpoint_tmp" +forward_pid= +cleanup() { + if [[ -n "$forward_pid" ]]; then + kill "$forward_pid" 2>/dev/null || true + wait "$forward_pid" 2>/dev/null || true + fi + rm -rf "$checkpoint_tmp" +} +trap cleanup EXIT +kubectl -n kagent exec deployment/kagent-controller -- \ + cat /run/substrate-servicedns/trust-bundle.pem > "$checkpoint_tmp/ca.pem" +kubectl -n kagent exec deployment/kagent-controller -- \ + cat /run/substrate-podidentity/credential-bundle.pem > "$checkpoint_tmp/client.pem" +chmod 600 "$checkpoint_tmp/client.pem" +kubectl -n ate-system port-forward service/api :443 \ + > "$checkpoint_tmp/forward.log" 2>&1 & +forward_pid=$! +port= +for _ in $(seq 1 100); do + port=$(awk '$1 == "Forwarding" && $3 ~ /^127[.]0[.]0[.]1:/ { split($3, address, ":"); print address[2]; exit }' "$checkpoint_tmp/forward.log") + if [[ -n "$port" ]]; then break; fi + if ! kill -0 "$forward_pid" 2>/dev/null; then + cat "$checkpoint_tmp/forward.log" >&2 + exit 1 + fi + sleep 0.1 +done +test -n "$port" +jq -n --arg endpoint "127.0.0.1:$port" --arg ca "$checkpoint_tmp/ca.pem" \ + --arg certificate "$checkpoint_tmp/client.pem" --arg image "$KAGENT_E2E_RUNTIME_IMAGE" \ + '{endpoint: $endpoint, serverName: "api.ate-system.svc", caFile: $ca, + clientCertFile: $certificate, atespace: "kagent", image: $image}' \ + > "$checkpoint_tmp/config.json" +KAGENT_E2E_CHECKPOINT_SUBSTRATE_CONFIG="$checkpoint_tmp/config.json" \ + go test -race ./core/internal/service/checkpoint \ + -run '^TestCheckpointRecoveryWithSubstrate$' -v -count=1 -timeout=10m From 75e80dbbe3e54ee37579d65c55282c379908e82d Mon Sep 17 00:00:00 2001 From: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> Date: Thu, 5 Nov 2026 00:31:22 +0800 Subject: [PATCH 04/20] docs(checkpoint): clarify recovery upgrade ordering Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> --- docs/architecture/persistence-checkpoints-and-forks.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/architecture/persistence-checkpoints-and-forks.md b/docs/architecture/persistence-checkpoints-and-forks.md index 87a7e6b8a1..f81d54e995 100644 --- a/docs/architecture/persistence-checkpoints-and-forks.md +++ b/docs/architecture/persistence-checkpoints-and-forks.md @@ -81,6 +81,10 @@ lock, and a delayed release cannot revoke a replacement's claim. Leader election reduces duplicate scans but is not the execution fence. A slow remote copy is still synchronized through Substrate's tag lease after a client attempt expires. +Upgrading from per-request tag writers requires stopping the old controller +replicas before starting this version. Those writers do not acquire checkpoint +execution claims; mixed-version rolling operation is not fenced by this worker. + A pending tag is not evidence of abandonment. Substrate serializes `CreateTag` against both the source Actor and the tag, including a copy still finishing after an earlier controller disappeared. Recovery calls `CreateTag` under those leases From ed88d1aa3c26b113bb01dcc807dc725b6522baec Mon Sep 17 00:00:00 2001 From: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> Date: Thu, 5 Nov 2026 16:42:18 +0800 Subject: [PATCH 05/20] fix(checkpoint): stop legacy writers before recovery upgrades Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> --- .github/workflows/ci.yaml | 2 +- Makefile | 2 +- .../persistence-checkpoints-and-forks.md | 11 +- .../test/upgrade/controller_handoff_test.go | 46 +++++++ go/core/test/upgrade/rolling_upgrade_test.go | 128 +++++++++++------- go/core/test/upgrade/upgrade_test.go | 6 + helm/README.md | 11 ++ .../templates/controller-deployment.yaml | 5 + .../tests/controller-configmap_test.yaml | 10 ++ helm/kagent/values.yaml | 3 + 10 files changed, 171 insertions(+), 53 deletions(-) create mode 100644 go/core/test/upgrade/controller_handoff_test.go diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 54d38946b8..8364ac4d9e 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -425,7 +425,7 @@ jobs: uses: ./.github/actions/upgrade-test-setup with: upgrade-from: ${{ matrix.upgrade-from }} - - name: Run rolling upgrade tests + - name: Verify controller upgrade handoff and schema compatibility if: steps.prep.outputs.skip != 'true' env: OPENAI_API_KEY: fake diff --git a/Makefile b/Makefile index 9fc3512c77..f2a5fd5328 100644 --- a/Makefile +++ b/Makefile @@ -579,7 +579,7 @@ run-upgrade-tests: announce-upgrade-from build install-previous-release ## Test # controller replicas (needed to observe the old-code/new-schema rollout window). .PHONY: run-rolling-upgrade-tests run-rolling-upgrade-tests: UPGRADE_PREV_EXTRA_ARGS = --set controller.replicas=2 -run-rolling-upgrade-tests: announce-upgrade-from build install-previous-release ## Install the previous release with 2 controller replicas, build the current images, and run the rolling upgrade e2e test +run-rolling-upgrade-tests: announce-upgrade-from build install-previous-release ## Verify stop-before-start controller upgrade from 2 previous-release replicas and schema compatibility @echo "=== Rolling upgrade test: $(UPGRADE_FROM_VERSION) -> $(VERSION) (registry=$(DOCKER_REGISTRY)) ===" cd go && \ KAGENT_E2E_RUN_ROLLING_UPGRADE_TESTS=true \ diff --git a/docs/architecture/persistence-checkpoints-and-forks.md b/docs/architecture/persistence-checkpoints-and-forks.md index f81d54e995..a04c920c66 100644 --- a/docs/architecture/persistence-checkpoints-and-forks.md +++ b/docs/architecture/persistence-checkpoints-and-forks.md @@ -81,9 +81,14 @@ lock, and a delayed release cannot revoke a replacement's claim. Leader election reduces duplicate scans but is not the execution fence. A slow remote copy is still synchronized through Substrate's tag lease after a client attempt expires. -Upgrading from per-request tag writers requires stopping the old controller -replicas before starting this version. Those writers do not acquire checkpoint -execution claims; mixed-version rolling operation is not fenced by this worker. +The controller Helm Deployment uses `Recreate`: all previous replicas stop +before a new controller starts. This introduces a controller API availability +gap during upgrades; existing Substrate Actors are not restarted. Previous +per-request tag writers do not acquire execution claims, so mixed-version +checkpoint execution is unsafe. Custom deployment tooling must preserve this +stop-before-start ordering, without force-deleting terminating writers. Before +rolling back to a chart that uses `RollingUpdate`, stop all current controllers +and wait for them to exit; the old chart does not enforce this handoff. A pending tag is not evidence of abandonment. Substrate serializes `CreateTag` against both the source Actor and the tag, including a copy still finishing after diff --git a/go/core/test/upgrade/controller_handoff_test.go b/go/core/test/upgrade/controller_handoff_test.go new file mode 100644 index 0000000000..39fbffdf49 --- /dev/null +++ b/go/core/test/upgrade/controller_handoff_test.go @@ -0,0 +1,46 @@ +package upgrade + +import ( + "testing" + + "github.com/stretchr/testify/require" + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/watch" +) + +func TestControllerHandoff(t *testing.T) { + old := corev1.Pod{ObjectMeta: metav1.ObjectMeta{Name: "old", UID: "old-uid"}} + replacement := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{Name: "old", UID: "new-uid"}} + t.Run("Given a terminating previous writer, when a replacement is created, assert that the overlap is rejected", func(t *testing.T) { + handoff := newControllerHandoff([]corev1.Pod{old}) + terminating := old.DeepCopy() + now := metav1.Now() + terminating.DeletionTimestamp = &now + require.NoError(t, handoff.observe(watch.Event{Type: watch.Modified, Object: terminating})) + require.ErrorContains(t, handoff.observe(watch.Event{Type: watch.Added, Object: replacement}), "previous writers still exist") + require.False(t, handoff.sawReplacement) + }) + t.Run("Given the previous writer has exited, when a replacement reuses its name, assert that the new UID completes the handoff", func(t *testing.T) { + handoff := newControllerHandoff([]corev1.Pod{old}) + require.NoError(t, handoff.observe(watch.Event{Type: watch.Deleted, Object: &old})) + require.NoError(t, handoff.observe(watch.Event{Type: watch.Added, Object: replacement})) + require.Empty(t, handoff.old) + require.True(t, handoff.sawReplacement) + }) + t.Run("Given an active handoff watch, when Kubernetes reports a watch error, assert that verification fails closed", func(t *testing.T) { + handoff := newControllerHandoff([]corev1.Pod{old}) + require.ErrorContains(t, handoff.observe(watch.Event{Type: watch.Error, Object: &metav1.Status{Reason: metav1.StatusReasonExpired}}), "watch failed") + require.False(t, handoff.sawReplacement) + }) + for _, phase := range []corev1.PodPhase{corev1.PodSucceeded, corev1.PodFailed} { + t.Run("Given a previous writer, when Kubernetes reports "+string(phase)+", assert that a replacement can start before object deletion", func(t *testing.T) { + handoff := newControllerHandoff([]corev1.Pod{old}) + exited := old.DeepCopy() + exited.Status.Phase = phase + require.NoError(t, handoff.observe(watch.Event{Type: watch.Modified, Object: exited})) + require.NoError(t, handoff.observe(watch.Event{Type: watch.Added, Object: replacement})) + require.True(t, handoff.sawReplacement) + }) + } +} diff --git a/go/core/test/upgrade/rolling_upgrade_test.go b/go/core/test/upgrade/rolling_upgrade_test.go index a0172faa3a..14a03d3689 100644 --- a/go/core/test/upgrade/rolling_upgrade_test.go +++ b/go/core/test/upgrade/rolling_upgrade_test.go @@ -8,6 +8,12 @@ import ( "time" "github.com/stretchr/testify/require" + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/types" + "k8s.io/apimachinery/pkg/watch" + "k8s.io/client-go/kubernetes" + "k8s.io/client-go/tools/clientcmd" kagentenv "github.com/kagent-dev/kagent/go/core/pkg/env" ) @@ -31,18 +37,16 @@ func TestRollingUpgradeCompatibility(t *testing.T) { waitForReadyPods(t, env, postgresSelector, 3*time.Minute) waitForPostgresSchema(t, env, 3*time.Minute) - // Run even when there is no migration delta: a rolling upgrade rolls the new - // image regardless of migrations, so the deploy can still break for - // non-schema reasons (a crashing new image, readiness, old pods against - // new-code-created resources). When the target build does add migrations, the - // same flow additionally exercises the old-code/new-schema window below. - // Keep multiple old controller pods around during the rollout. With a single - // replica the old-code/new-schema window can be too small to observe reliably. + // Multiple previous-generation API replicas must all exit before checkpoint + // recovery starts. Keep a terminating writer alive long enough that an unsafe + // RollingUpdate also has a visible overlap window. kubectl(t, env, 2*time.Minute, "scale", "deployment/kagent-controller", "-n", env.namespace, "--replicas=2", ) + kubectl(t, env, time.Minute, "patch", "deployment/kagent-controller", "-n", env.namespace, + "--type=strategic", "-p", `{"spec":{"template":{"spec":{"terminationGracePeriodSeconds":30,"containers":[{"name":"controller","lifecycle":{"preStop":{"sleep":{"seconds":5}}}}]}}}}`) kubectl(t, env, 3*time.Minute, "rollout", "status", "deployment/kagent-controller", "-n", env.namespace, @@ -51,15 +55,29 @@ func TestRollingUpgradeCompatibility(t *testing.T) { oldPods := podNamesForSelector(t, env, controllerSelector) require.NotEmpty(t, oldPods, "expected old controller pods before rolling upgrade") - // Seed with the baseline schema before the target controller applies new - // migrations. The compatibility canary below verifies this row is still - // readable while old pods are alive against the target schema. + // Preserve old-shape SQL read/write canaries across the upgrade. The + // controller handoff itself is deliberately stop-before-start, not a claim + // that old per-request checkpoint writers support mixed-version execution. pgExec(t, env, fmt.Sprintf( "INSERT INTO tool (id, server_name, group_kind, description) VALUES (%s, %s, %s, 'rolling baseline tool')", pgQuote(baselineToolID), pgQuote(toolServerName), pgQuote(groupKind))) ctx, cancel := context.WithTimeout(t.Context(), 10*time.Minute) defer cancel() + config, err := clientcmd.NewNonInteractiveDeferredLoadingClientConfig( + clientcmd.NewDefaultClientConfigLoadingRules(), &clientcmd.ConfigOverrides{CurrentContext: env.kubeContext}).ClientConfig() + require.NoError(t, err) + client, err := kubernetes.NewForConfig(config) + require.NoError(t, err) + pods, err := client.CoreV1().Pods(env.namespace).List(ctx, metav1.ListOptions{LabelSelector: controllerSelector}) + require.NoError(t, err) + require.NotEmpty(t, pods.Items, "expected previous controller pod UIDs before upgrade") + observer := newControllerHandoff(pods.Items) + events, err := client.CoreV1().Pods(env.namespace).Watch(ctx, metav1.ListOptions{ + LabelSelector: controllerSelector, ResourceVersion: pods.ResourceVersion, + }) + require.NoError(t, err) + defer events.Stop() done := make(chan upgradeResult, 1) cmd := upgradeCommand(ctx, env) go func() { @@ -67,36 +85,28 @@ func TestRollingUpgradeCompatibility(t *testing.T) { done <- upgradeResult{out: string(out), err: err} }() - // Wait until the target schema has landed while at least one previous-release - // controller pod is still ready. That is the rolling deploy hazard: old code - // can keep serving briefly after a new pod has applied migrations. Surface an - // upgrade failure immediately instead of letting it look like a - // schema-observation timeout. - var upgradeErr error - var upgradeOut string - require.Eventually(t, func() bool { + var result upgradeResult + finished := false + for !finished || !observer.sawReplacement { select { - case r := <-done: - if r.err != nil { - upgradeErr, upgradeOut = r.err, r.out - return true - } - // The upgrade finished successfully before we caught the window; put the - // result back so the final read below still sees it, then keep polling. - done <- r - default: - } - state, err := pgMigrationStateE(t, env) - if err != nil { - return false + case event, ok := <-events.ResultChan(): + require.True(t, ok, "pod watch ended before the controller handoff was verified") + require.NoError(t, observer.observe(event)) + case result = <-done: + require.NoError(t, result.err, "upgrade failed:\n%s", result.out) + finished = true + case <-ctx.Done(): + t.Fatal(ctx.Err()) } - return state.version == targetCoreVersion && anyPodsReady(t, env, oldPods) - }, 6*time.Minute, 500*time.Millisecond, "target schema was not observed while old controller pods were still ready") - require.NoError(t, upgradeErr, "upgrade failed before target schema was observed:\n%s", upgradeOut) + } + require.Empty(t, observer.old, "all previous writers must have exited") + strategy := kubectl(t, env, time.Minute, "get", "deployment/kagent-controller", "-n", env.namespace, + "-o", "jsonpath={.spec.strategy.type}") + require.Equal(t, "Recreate", strings.TrimSpace(strategy)) // These SQL canaries intentionally use the pre-upgrade column shape. They do // not prove every previous-release code path works, but they catch migrations - // that break basic old read/write assumptions during a rolling deployment. + // that break basic old read/write assumptions after the controller handoff. require.Equal(t, 1, pgQueryInt(t, env, fmt.Sprintf("SELECT count(*) FROM tool WHERE id = %s AND server_name = %s AND group_kind = %s", pgQuote(baselineToolID), pgQuote(toolServerName), pgQuote(groupKind))), @@ -111,9 +121,6 @@ func TestRollingUpgradeCompatibility(t *testing.T) { "old-shape tool write did not survive against target schema", ) - result := <-done - require.NoError(t, result.err, "upgrade failed:\n%s", result.out) - kubectl(t, env, 3*time.Minute, "rollout", "status", "deployment/kagent-controller", "-n", env.namespace, @@ -159,18 +166,43 @@ func podNamesForSelectorE(t *testing.T, env upgradeEnv, selector string) ([]stri return pods, nil } -func anyPodsReady(t *testing.T, env upgradeEnv, pods []string) bool { - t.Helper() +type controllerHandoff struct { + old map[types.UID]bool + sawReplacement bool +} +func newControllerHandoff(pods []corev1.Pod) *controllerHandoff { + old := make(map[types.UID]bool, len(pods)) for _, pod := range pods { - out, err := kubectlOutput(t, env, 10*time.Second, - "get", "pod", pod, - "-n", env.namespace, - "-o", fmt.Sprintf("jsonpath={.status.containerStatuses[?(@.name==%q)].ready}", controllerContainer), - ) - if err == nil && strings.TrimSpace(out) == "true" { - return true + if pod.Status.Phase != corev1.PodSucceeded && pod.Status.Phase != corev1.PodFailed { + old[pod.UID] = true + } + } + return &controllerHandoff{old: old} +} + +// Terminating and unready old pods can still issue checkpoint writes. Only +// deletion or a terminal phase (all processes exited) releases the barrier, +// matching the Deployment controller's Recreate condition. +func (h *controllerHandoff) observe(event watch.Event) error { + if event.Type == watch.Error { + return fmt.Errorf("controller pod watch failed: %v", event.Object) + } + pod, ok := event.Object.(*corev1.Pod) + if !ok { + return fmt.Errorf("unexpected controller watch object %T", event.Object) + } + if h.old[pod.UID] { + if event.Type == watch.Deleted || pod.Status.Phase == corev1.PodSucceeded || pod.Status.Phase == corev1.PodFailed { + delete(h.old, pod.UID) + } + return nil + } + if event.Type == watch.Added || event.Type == watch.Modified { + if len(h.old) != 0 { + return fmt.Errorf("new controller %s started while %d previous writers still exist", pod.Name, len(h.old)) } + h.sawReplacement = true } - return false + return nil } diff --git a/go/core/test/upgrade/upgrade_test.go b/go/core/test/upgrade/upgrade_test.go index df6a40c03c..1fa5022eae 100644 --- a/go/core/test/upgrade/upgrade_test.go +++ b/go/core/test/upgrade/upgrade_test.go @@ -237,6 +237,12 @@ func TestUpgrade(t *testing.T) { ctx, cancel := context.WithTimeout(t.Context(), 10*time.Minute) defer cancel() + // The previous chart uses RollingUpdate and its checkpoint writers do + // not acquire execution claims. Finish the current generation before + // restoring that chart; schema compatibility does not fence writers. + kubectl(t, env, time.Minute, "scale", "deployment/kagent-controller", "-n", env.namespace, "--replicas=0") + kubectl(t, env, 3*time.Minute, "wait", "pod", "-n", env.namespace, + "-l", controllerSelector, "--for=delete", "--timeout=3m") out, err := installPreviousReleaseCommand(ctx, env).CombinedOutput() require.NoError(t, err, "rolling back to the previous release failed:\n%s", string(out)) kubectl(t, env, 3*time.Minute, diff --git a/helm/README.md b/helm/README.md index 8149ba7fcc..212e9ec332 100644 --- a/helm/README.md +++ b/helm/README.md @@ -304,6 +304,17 @@ validation commands, defaults, and runtime configuration. ## Upgrading +The controller Deployment uses `Recreate`: previous controller processes exit +before the next generation starts. Plan for a controller API availability gap; +existing Substrate Actors are not restarted. This ordering prevents legacy +checkpoint writers from racing the new recovery worker. Custom deployment +tooling must preserve it without force-deleting terminating writers. + +Before rolling back to an older chart that uses `RollingUpdate`, scale the +current controller Deployment to zero and wait for its Pods to exit. Only then +restore the old chart, which does not enforce this handoff. See the +[checkpoint recovery contract](../docs/architecture/persistence-checkpoints-and-forks.md). + When upgrading, make sure to upgrade both charts: ```bash diff --git a/helm/kagent/templates/controller-deployment.yaml b/helm/kagent/templates/controller-deployment.yaml index 824f44efcd..4caf31211c 100644 --- a/helm/kagent/templates/controller-deployment.yaml +++ b/helm/kagent/templates/controller-deployment.yaml @@ -21,6 +21,11 @@ metadata: {{- include "kagent.controller.labels" . | nindent 4 }} spec: replicas: {{ .Values.controller.replicas }} + # Old checkpoint writers do not participate in recovery execution claims. + # Wait for their processes to exit before a new worker can start. + strategy: + type: Recreate + rollingUpdate: null selector: matchLabels: {{- include "kagent.controller.selectorLabels" . | nindent 6 }} diff --git a/helm/kagent/tests/controller-configmap_test.yaml b/helm/kagent/tests/controller-configmap_test.yaml index 15a6c077a6..0fa06e7b3a 100644 --- a/helm/kagent/tests/controller-configmap_test.yaml +++ b/helm/kagent/tests/controller-configmap_test.yaml @@ -3,6 +3,16 @@ templates: - controller-configmap.yaml - controller-deployment.yaml tests: + - it: stops previous controller writers before starting the new generation + template: controller-deployment.yaml + asserts: + - equal: + path: spec.strategy.type + value: Recreate + - equal: + path: spec.strategy.rollingUpdate + value: null + - it: configures Actor JWT discovery with the Substrate trust bundle template: controller-configmap.yaml set: diff --git a/helm/kagent/values.yaml b/helm/kagent/values.yaml index 69c5b77417..abf45cacde 100644 --- a/helm/kagent/values.yaml +++ b/helm/kagent/values.yaml @@ -224,6 +224,9 @@ config: # ============================================================================== controller: + # Controller upgrades use Recreate so legacy checkpoint writers exit before + # recovery workers start. The control-plane API is briefly unavailable; runtime + # Actors are not restarted by this controller replacement. # Standard OTEL overrides and external SDK credentials for this process only. env: [] envFrom: [] From a26dc54e8436324f007ab192f27a513d373dc836 Mon Sep 17 00:00:00 2001 From: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> Date: Thu, 5 Nov 2026 16:55:15 +0800 Subject: [PATCH 06/20] test(checkpoint): recover partially copied tags after storage denial Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> --- .../checkpoint/recovery_objectstore_test.go | 55 +++++++++ .../checkpoint/recovery_substrate_test.go | 110 ++++++++++++++++-- go/core/test/e2e/checkpoint-recovery.sh | 62 ++++++---- 3 files changed, 198 insertions(+), 29 deletions(-) create mode 100644 go/core/internal/service/checkpoint/recovery_objectstore_test.go diff --git a/go/core/internal/service/checkpoint/recovery_objectstore_test.go b/go/core/internal/service/checkpoint/recovery_objectstore_test.go new file mode 100644 index 0000000000..bf43545299 --- /dev/null +++ b/go/core/internal/service/checkpoint/recovery_objectstore_test.go @@ -0,0 +1,55 @@ +package checkpoint + +import ( + "context" + "testing" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/credentials" + "github.com/aws/aws-sdk-go-v2/service/s3" + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +// Fault policies and injected objects belong to a bucket created exclusively +// for this scenario; the installation's shared snapshot bucket is untouched. +func checkpointObjectStore(t *testing.T, endpoint, accessKey, secretKey string) (*s3.Client, string) { + t.Helper() + require.NotEmpty(t, endpoint, "real object-store endpoint is required") + require.NotEmpty(t, accessKey) + require.NotEmpty(t, secretKey) + client := s3.NewFromConfig(aws.Config{ + Region: "us-east-1", Credentials: credentials.NewStaticCredentialsProvider(accessKey, secretKey, ""), RetryMaxAttempts: 1, + }, func(o *s3.Options) { o.BaseEndpoint = aws.String(endpoint); o.UsePathStyle = true }) + bucket := "checkpoint-recovery-" + uuid.NewString() + _, err := client.CreateBucket(t.Context(), &s3.CreateBucketInput{Bucket: aws.String(bucket)}) + require.NoError(t, err) + t.Cleanup(func() { + ctx, cancel := context.WithTimeout(context.Background(), time.Minute) + defer cancel() + _, err := client.DeleteBucketPolicy(ctx, &s3.DeleteBucketPolicyInput{Bucket: aws.String(bucket)}) + require.NoError(t, err) + for _, key := range checkpointObjectKeys(t, ctx, client, bucket, "") { + _, err := client.DeleteObject(ctx, &s3.DeleteObjectInput{Bucket: aws.String(bucket), Key: aws.String(key)}) + require.NoError(t, err) + } + _, err = client.DeleteBucket(ctx, &s3.DeleteBucketInput{Bucket: aws.String(bucket)}) + require.NoError(t, err) + }) + return client, bucket +} + +func checkpointObjectKeys(t *testing.T, ctx context.Context, client *s3.Client, bucket, prefix string) []string { + t.Helper() + pages := s3.NewListObjectsV2Paginator(client, &s3.ListObjectsV2Input{Bucket: aws.String(bucket), Prefix: aws.String(prefix)}) + var keys []string + for pages.HasMorePages() { + page, err := pages.NextPage(ctx) + require.NoError(t, err) + for _, object := range page.Contents { + keys = append(keys, aws.ToString(object.Key)) + } + } + return keys +} diff --git a/go/core/internal/service/checkpoint/recovery_substrate_test.go b/go/core/internal/service/checkpoint/recovery_substrate_test.go index 2d9eafd611..8b1270bc6f 100644 --- a/go/core/internal/service/checkpoint/recovery_substrate_test.go +++ b/go/core/internal/service/checkpoint/recovery_substrate_test.go @@ -5,10 +5,15 @@ import ( "crypto/tls" "crypto/x509" "encoding/json" + "fmt" + "net/url" "os" + "strings" "testing" "time" + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/s3" "github.com/google/uuid" "github.com/stretchr/testify/require" ateapipb "go.ate.dev/substrate/api/ateapi/v1" @@ -16,6 +21,7 @@ import ( "google.golang.org/grpc/codes" "google.golang.org/grpc/credentials" "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" @@ -26,7 +32,8 @@ import ( // This suite uses real Actors, snapshot copies, UID-conditional deletion and // PostgreSQL. Faults discard actual RPC results or prevent one RPC from being -// delivered; no backend response or persisted lifecycle state is fabricated. +// delivered, or deny reads/deletes in an exclusively owned object-store bucket; +// no backend response or persisted lifecycle state is fabricated. // The configuration is test-only; CI runs this after installing Substrate. func TestCheckpointRecoveryWithSubstrate(t *testing.T) { path := os.Getenv("KAGENT_E2E_CHECKPOINT_SUBSTRATE_CONFIG") @@ -34,12 +41,15 @@ func TestCheckpointRecoveryWithSubstrate(t *testing.T) { t.Skip("requires real Substrate; run core/test/e2e/checkpoint-recovery.sh") } var config struct { - Endpoint string `json:"endpoint"` - ServerName string `json:"serverName"` - CAFile string `json:"caFile"` - ClientCertFile string `json:"clientCertFile"` - Atespace string `json:"atespace"` - Image string `json:"image"` + Endpoint string `json:"endpoint"` + ServerName string `json:"serverName"` + CAFile string `json:"caFile"` + ClientCertFile string `json:"clientCertFile"` + Atespace string `json:"atespace"` + Image string `json:"image"` + ObjectStoreEndpoint string `json:"objectStoreEndpoint"` + ObjectStoreAccessKey string `json:"objectStoreAccessKey"` + ObjectStoreSecretKey string `json:"objectStoreSecretKey"` } data, err := os.ReadFile(path) require.NoError(t, err) @@ -85,7 +95,7 @@ func TestCheckpointRecoveryWithSubstrate(t *testing.T) { t.Cleanup(pool.Close) newStore := func() *database.Client { return database.NewClient(pool, "public") } owner := auth.AuthSessionTo(ctx, testSession{userID: "alice"}) - newSource := func(t *testing.T) (*database.Client, *apiv1.Session, string) { + newSourceForTemplate := func(t *testing.T, template *ateapipb.ActorTemplate) (*database.Client, *apiv1.Session, string) { t.Helper() name := substrate.ActorName(uuid.NewString()) _, err := client.CreateActor(ctx, config.Atespace, name, config.Atespace, template.Metadata.Name) @@ -107,6 +117,9 @@ func TestCheckpointRecoveryWithSubstrate(t *testing.T) { session, head := recoverySourceForActor(t, store, actor) return store, session, head } + newSource := func(t *testing.T) (*database.Client, *apiv1.Session, string) { + return newSourceForTemplate(t, template) + } deleteCheckpoint := func(t *testing.T, service *Service, id string) { t.Helper() stop := startRecovery(t, service) @@ -265,6 +278,87 @@ func TestCheckpointRecoveryWithSubstrate(t *testing.T) { deleteCheckpoint(t, replacement, cp.Id) }) + t.Run("Given a partially copied tag and denied object cleanup, when recovery restarts after storage repair, assert that cleanup precedes releasing source references", func(t *testing.T) { + objects, bucket := checkpointObjectStore(t, config.ObjectStoreEndpoint, config.ObjectStoreAccessKey, config.ObjectStoreSecretKey) + privateTemplate := proto.Clone(template).(*ateapipb.ActorTemplate) + privateTemplate.Metadata = &ateapipb.ResourceMetadata{Atespace: config.Atespace, Name: "checkpoint-partial-" + uuid.NewString()} + privateTemplate.Status = nil + privateTemplate.SnapshotConfig.StorageLocation = "gs://" + bucket + privateTemplate, err := client.CreateActorTemplate(ctx, privateTemplate) + require.NoError(t, err) + t.Cleanup(func() { + cleanup, stop := context.WithTimeout(context.Background(), time.Minute) + defer stop() + require.NoError(t, client.DeleteActorTemplate(cleanup, config.Atespace, privateTemplate.Metadata.Name, privateTemplate.Metadata.Uid)) + }) + store, source, head := newSourceForTemplate(t, privateTemplate) + actor, err := client.GetActor(ctx, config.Atespace, substrate.ActorName(source.Id)) + require.NoError(t, err) + snapshot, err := url.Parse(substrate.SnapshotURI(substrate.LatestActorSnapshot(actor))) + require.NoError(t, err) + require.Equal(t, bucket, snapshot.Host) + prefix := strings.Trim(snapshot.Path, "/") + "/" + // Substrate copies in sorted order with eight workers. More than eight + // preceding objects make the final denied read occur after real copies. + for i := range 64 { + _, err := objects.PutObject(ctx, &s3.PutObjectInput{Bucket: aws.String(bucket), Key: aws.String(fmt.Sprintf("%szz-copy-probe/%04d", prefix, i)), Body: strings.NewReader("partial-copy-proof")}) + require.NoError(t, err) + } + deniedKey := prefix + "zz-copy-probe/9999-denied" + _, err = objects.PutObject(ctx, &s3.PutObjectInput{Bucket: aws.String(bucket), Key: aws.String(deniedKey), Body: strings.NewReader("deny this read")}) + require.NoError(t, err) + policy, err := json.Marshal(map[string]any{ + "Version": "2012-10-17", "Statement": []any{ + map[string]any{"Effect": "Deny", "Principal": "*", "Action": "s3:GetObject", "Resource": "arn:aws:s3:::" + bucket + "/" + deniedKey}, + map[string]any{"Effect": "Deny", "Principal": "*", "Action": "s3:DeleteObject", "Resource": "arn:aws:s3:::" + bucket + "/atespaces/" + config.Atespace + "/tags/*"}, + }, + }) + require.NoError(t, err) + service := NewService(store, testAuthorizer{}, client, nil) + cp, _, err := service.beginCreate(owner, source.Id, uuid.NewString(), head) + require.NoError(t, err) + cleanupTag(t, cp.Id) + _, err = objects.PutBucketPolicy(ctx, &s3.PutBucketPolicyInput{Bucket: aws.String(bucket), Policy: aws.String(string(policy))}) + require.NoError(t, err) + // Remove the denial before any Actor/tag/template cleanup, even on failure. + t.Cleanup(func() { + cleanup, stop := context.WithTimeout(context.Background(), time.Minute) + defer stop() + _, err := objects.DeleteBucketPolicy(cleanup, &s3.DeleteBucketPolicyInput{Bucket: aws.String(bucket)}) + require.NoError(t, err) + }) + op, err := store.GetCheckpointOperation(ctx, cp.Id, "alice") + require.NoError(t, err) + copyErr := service.recoverOperation(ctx, op) + require.Error(t, copyErr, "the actual object read must fail") + t.Logf("expected object-store copy failure: %v", copyErr) + tag, err := client.GetTag(ctx, config.Atespace, tagName(cp.Id)) + require.NoError(t, err) + require.Empty(t, substrate.SnapshotURI(tag.GetStatus().GetSnapshot())) + tagPrefix := "atespaces/" + config.Atespace + "/tags/" + tag.Metadata.Uid + "/" + require.NotEmpty(t, checkpointObjectKeys(t, ctx, objects, bucket, tagPrefix), "copy failure must leave actual destination objects") + replacement := NewService(newStore(), testAuthorizer{}, client, nil) + cleanupErr := replacement.recoverOperation(ctx, op) + require.Error(t, cleanupErr, "the real object deletion must also fail") + t.Logf("expected object-store cleanup failure: %v", cleanupErr) + pending, err := store.GetCheckpointOperation(ctx, cp.Id, "alice") + require.NoError(t, err) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_CREATING, pending.Checkpoint.State) + require.NotNil(t, pending.Checkpoint.Failure) + require.Equal(t, tag.Metadata.Uid, pending.TagUID) + _, err = store.BeginSessionOperation(ctx, source.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE) + require.ErrorIs(t, err, database.ErrConflict, "partial tag cleanup must retain source admission") + _, err = objects.DeleteBucketPolicy(ctx, &s3.DeleteBucketPolicyInput{Bucket: aws.String(bucket)}) + require.NoError(t, err) + replacement = NewService(newStore(), testAuthorizer{}, client, nil) + require.NoError(t, replacement.recoverOperation(ctx, pending)) + waitCheckpointPhase(t, store, cp.Id, apiv1.CheckpointState_CHECKPOINT_STATE_FAILED) + require.Empty(t, checkpointObjectKeys(t, ctx, objects, bucket, tagPrefix), "all partially copied objects must be collected") + _, err = store.BeginSessionOperation(ctx, source.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE) + require.NoError(t, err, "storage repair and cleanup must reopen the source") + deleteCheckpoint(t, replacement, cp.Id) + }) + t.Run("Given a lost delete response, when recovery restarts, assert that an absent real tag completes deletion", func(t *testing.T) { store, source, head := newSource(t) service := NewService(store, testAuthorizer{}, client, nil) diff --git a/go/core/test/e2e/checkpoint-recovery.sh b/go/core/test/e2e/checkpoint-recovery.sh index 2606b6a16b..a94f0b0099 100644 --- a/go/core/test/e2e/checkpoint-recovery.sh +++ b/go/core/test/e2e/checkpoint-recovery.sh @@ -7,12 +7,14 @@ cd "$(dirname "${BASH_SOURCE[0]}")/../../.." : "${KAGENT_E2E_RUNTIME_IMAGE:?set the digest-pinned Go ADK image}" checkpoint_tmp=$(mktemp -d) chmod 700 "$checkpoint_tmp" -forward_pid= +forward_pids=() cleanup() { - if [[ -n "$forward_pid" ]]; then - kill "$forward_pid" 2>/dev/null || true - wait "$forward_pid" 2>/dev/null || true - fi + for pid in "${forward_pids[@]:-}"; do + if [[ -n "$pid" ]]; then + kill "$pid" 2>/dev/null || true + wait "$pid" 2>/dev/null || true + fi + done rm -rf "$checkpoint_tmp" } trap cleanup EXIT @@ -21,25 +23,43 @@ kubectl -n kagent exec deployment/kagent-controller -- \ kubectl -n kagent exec deployment/kagent-controller -- \ cat /run/substrate-podidentity/credential-bundle.pem > "$checkpoint_tmp/client.pem" chmod 600 "$checkpoint_tmp/client.pem" -kubectl -n ate-system port-forward service/api :443 \ - > "$checkpoint_tmp/forward.log" 2>&1 & -forward_pid=$! -port= -for _ in $(seq 1 100); do - port=$(awk '$1 == "Forwarding" && $3 ~ /^127[.]0[.]0[.]1:/ { split($3, address, ":"); print address[2]; exit }' "$checkpoint_tmp/forward.log") - if [[ -n "$port" ]]; then break; fi - if ! kill -0 "$forward_pid" 2>/dev/null; then - cat "$checkpoint_tmp/forward.log" >&2 - exit 1 - fi - sleep 0.1 -done -test -n "$port" -jq -n --arg endpoint "127.0.0.1:$port" --arg ca "$checkpoint_tmp/ca.pem" \ +forward_port() { + local service=$1 remote_port=$2 forward_pid log + log="$checkpoint_tmp/$service-forward.log" + kubectl -n ate-system port-forward "service/$service" ":$remote_port" > "$log" 2>&1 & + forward_pid=$! + forward_pids+=("$forward_pid") + checkpoint_forwarded_port= + for _ in $(seq 1 100); do + checkpoint_forwarded_port=$(awk '$1 == "Forwarding" && $3 ~ /^127[.]0[.]0[.]1:/ { split($3, address, ":"); print address[2]; exit }' "$log") + if [[ -n "$checkpoint_forwarded_port" ]]; then return; fi + if ! kill -0 "$forward_pid" 2>/dev/null; then + cat "$log" >&2 + return 1 + fi + sleep 0.1 + done + echo "Timed out forwarding $service" >&2 + return 1 +} +forward_port api 443 +api_port=$checkpoint_forwarded_port +forward_port rustfs 9000 +object_store_port=$checkpoint_forwarded_port +# The disposable Substrate installation provides literal RustFS credentials. +# Keep them in private files rather than command arguments or test output. +kubectl -n ate-system get deployment/rustfs -o json > "$checkpoint_tmp/object-store.json" +jq -n --slurpfile store "$checkpoint_tmp/object-store.json" \ + --arg endpoint "127.0.0.1:$api_port" --arg ca "$checkpoint_tmp/ca.pem" \ + --arg objectStoreEndpoint "http://127.0.0.1:$object_store_port" \ --arg certificate "$checkpoint_tmp/client.pem" --arg image "$KAGENT_E2E_RUNTIME_IMAGE" \ '{endpoint: $endpoint, serverName: "api.ate-system.svc", caFile: $ca, - clientCertFile: $certificate, atespace: "kagent", image: $image}' \ + clientCertFile: $certificate, atespace: "kagent", image: $image, + objectStoreEndpoint: $objectStoreEndpoint, + objectStoreAccessKey: ([$store[0].spec.template.spec.containers[].env[] | select(.name == "RUSTFS_ACCESS_KEY").value][0] // error("missing RustFS access key")), + objectStoreSecretKey: ([$store[0].spec.template.spec.containers[].env[] | select(.name == "RUSTFS_SECRET_KEY").value][0] // error("missing RustFS secret key"))}' \ > "$checkpoint_tmp/config.json" +chmod 600 "$checkpoint_tmp/config.json" KAGENT_E2E_CHECKPOINT_SUBSTRATE_CONFIG="$checkpoint_tmp/config.json" \ go test -race ./core/internal/service/checkpoint \ -run '^TestCheckpointRecoveryWithSubstrate$' -v -count=1 -timeout=10m From c935cc3cb13243b3175de0d2a87a2770a76f0a59 Mon Sep 17 00:00:00 2001 From: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> Date: Fri, 6 Nov 2026 18:24:50 +0800 Subject: [PATCH 07/20] fix(checkpoint): address recovery and rolling upgrade review Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> --- Makefile | 2 +- .../persistence-checkpoints-and-forks.md | 43 ++- .../database/checkpoint_operations.go | 68 +++- .../database/checkpoint_operations_test.go | 35 ++ go/core/internal/database/checkpoints.go | 2 +- .../internal/service/checkpoint/recovery.go | 110 +++++- go/core/internal/service/checkpoint/retry.go | 68 ++++ .../service/checkpoint/review_test.go | 330 ++++++++++++++++++ .../internal/service/checkpoint/service.go | 61 +++- .../service/checkpoint/service_test.go | 5 + go/core/pkg/app/app.go | 8 + go/core/pkg/app/checkpoint_activation.go | 111 ++++++ go/core/pkg/app/checkpoint_activation_test.go | 144 ++++++++ go/core/pkg/env/controller.go | 6 + .../core/000002_checkpoint_recovery.sql | 4 +- go/core/test/e2e/README.md | 10 + .../test/e2e/checkpoint_activation_test.go | 90 +++++ go/core/test/e2e/connection_test.go | 8 +- .../test/upgrade/controller_handoff_test.go | 46 --- go/core/test/upgrade/rolling_upgrade_test.go | 128 +++---- helm/README.md | 20 +- .../templates/controller-deployment.yaml | 8 +- .../templates/rbac/leader-election-role.yaml | 7 + .../tests/controller-configmap_test.yaml | 18 +- helm/kagent/tests/leader-election_test.yaml | 14 + helm/kagent/values.yaml | 5 +- 26 files changed, 1151 insertions(+), 200 deletions(-) create mode 100644 go/core/internal/service/checkpoint/retry.go create mode 100644 go/core/internal/service/checkpoint/review_test.go create mode 100644 go/core/pkg/app/checkpoint_activation.go create mode 100644 go/core/pkg/app/checkpoint_activation_test.go create mode 100644 go/core/test/e2e/checkpoint_activation_test.go delete mode 100644 go/core/test/upgrade/controller_handoff_test.go diff --git a/Makefile b/Makefile index f2a5fd5328..b27f6815dd 100644 --- a/Makefile +++ b/Makefile @@ -579,7 +579,7 @@ run-upgrade-tests: announce-upgrade-from build install-previous-release ## Test # controller replicas (needed to observe the old-code/new-schema rollout window). .PHONY: run-rolling-upgrade-tests run-rolling-upgrade-tests: UPGRADE_PREV_EXTRA_ARGS = --set controller.replicas=2 -run-rolling-upgrade-tests: announce-upgrade-from build install-previous-release ## Verify stop-before-start controller upgrade from 2 previous-release replicas and schema compatibility +run-rolling-upgrade-tests: announce-upgrade-from build install-previous-release ## Verify rolling controller upgrade from 2 previous-release replicas and schema compatibility @echo "=== Rolling upgrade test: $(UPGRADE_FROM_VERSION) -> $(VERSION) (registry=$(DOCKER_REGISTRY)) ===" cd go && \ KAGENT_E2E_RUN_ROLLING_UPGRADE_TESTS=true \ diff --git a/docs/architecture/persistence-checkpoints-and-forks.md b/docs/architecture/persistence-checkpoints-and-forks.md index a04c920c66..91e2218ebf 100644 --- a/docs/architecture/persistence-checkpoints-and-forks.md +++ b/docs/architecture/persistence-checkpoints-and-forks.md @@ -81,14 +81,41 @@ lock, and a delayed release cannot revoke a replacement's claim. Leader election reduces duplicate scans but is not the execution fence. A slow remote copy is still synchronized through Substrate's tag lease after a client attempt expires. -The controller Helm Deployment uses `Recreate`: all previous replicas stop -before a new controller starts. This introduces a controller API availability -gap during upgrades; existing Substrate Actors are not restarted. Previous -per-request tag writers do not acquire execution claims, so mixed-version -checkpoint execution is unsafe. Custom deployment tooling must preserve this -stop-before-start ordering, without force-deleting terminating writers. Before -rolling back to a chart that uses `RollingUpdate`, stop all current controllers -and wait for them to exit; the old chart does not enforce this handoff. +Discovery enqueues checkpoint IDs independently of four execution workers, so a +slow copy cannot hold up unrelated new work. Failed attempts retry with capped +exponential backoff from one to thirty seconds and jitter, including at the cap. +Scans and immediate dirty-key requeues respect each checkpoint's retry deadline. +The retry schedule is local to a worker; reservations and execution claims remain +durable across restarts. + +Permanent operation errors are recorded under the current execution claim, +without deleting the checkpoint or releasing its references. Public waiters on +any API replica return the recorded status, such as `FailedPrecondition`, rather +than waiting for their deadline. Backend configuration is excluded from these +public error messages. Confirmed progress clears the error. Both creation and +deletion waiters use exponential polling backoff from 100 ms to one second with +jitter; canceling a caller still leaves the admitted work available for recovery. + +Controller upgrades remain `RollingUpdate`. New Pods advertise the private +`kagent.dev/checkpoint-execution-claims: v1` annotation. Each API replica uses the +live Kubernetes API to wait for all unmarked Pods in its Deployment's stable +selector to exit and their ReplicaSets to acknowledge scale-down to zero. +Unready and terminating legacy Pods still block activation, as do old ReplicaSets +that could recreate them. +This barrier runs independently of leader election and Pod readiness, so new +replicas can serve other traffic while Kubernetes drains the previous generation. +Checkpoint creation and deletion return retryable `Unavailable` before admitting +new work until activation; the recovery worker waits on the same barrier. +Already admitted reservations and reference pins survive the transition. + +Compatible future rollouts activate immediately without a controller API outage. +Custom deployments must preserve the protocol annotation and controller Pod +identity. A standalone development controller conservatively checks controller +Pods in its namespace; stop legacy Pods before running it. Kubernetes read errors +keep recovery closed. Do not force-delete terminating legacy writers. Before +rolling back to an older chart without the protocol, quiesce checkpoint requests +and finish pending recovery, then stop all current controllers +and wait for them to exit: the startup barrier does not fence a later downgrade. A pending tag is not evidence of abandonment. Substrate serializes `CreateTag` against both the source Actor and the tag, including a copy still finishing after diff --git a/go/core/internal/database/checkpoint_operations.go b/go/core/internal/database/checkpoint_operations.go index c2a1caa4fe..3b2fd93650 100644 --- a/go/core/internal/database/checkpoint_operations.go +++ b/go/core/internal/database/checkpoint_operations.go @@ -8,6 +8,7 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" + statuspb "google.golang.org/genproto/googleapis/rpc/status" "google.golang.org/protobuf/proto" apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" @@ -21,6 +22,13 @@ type CheckpointOperation struct { Creator string Snapshot *SessionTaskSnapshot TagUID string + // A permanent attempt error does not change the phase or authorize cleanup. + OperationError *statuspb.Status +} + +type checkpointOperationRow struct { + sessionCheckpointRow + OperationError []byte } // ClaimCheckpointOperation grants one bounded attempt exclusive execution of an @@ -48,8 +56,8 @@ func (c *Client) ClaimCheckpointOperation(ctx context.Context, id, owner string, AND (executor_id IS NULL 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) + source_history_id, prepared_revision, operation_error + `, pgx.RowToStructByName[checkpointOperationRow], id, owner, executor, time.Until(deadline).Milliseconds()+1) if errors.Is(err, pgx.ErrNoRows) { return ErrConflict } @@ -114,7 +122,12 @@ func checkCheckpointExecutor(ctx context.Context, tx pgx.Tx, id string, executor // 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 := readCheckpoint(ctx, c.db, id, owner, nil) + 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, operation_error + FROM session_checkpoint WHERE id = $1 AND user_id = $2 + `, pgx.RowToStructByName[checkpointOperationRow], id, owner) if err != nil { return nil, fmt.Errorf("get checkpoint operation: %w", notFoundOr(err)) } @@ -132,12 +145,12 @@ func (c *Client) ListPendingCheckpoints(ctx context.Context, afterID string, lim 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 + source_history_id, prepared_revision, operation_error FROM session_checkpoint WHERE state IN ('CREATING', 'DELETING') AND ($1::uuid IS NULL OR id > $1::uuid) ORDER BY id LIMIT $2 - `, pgx.RowToStructByName[sessionCheckpointRow], after, limit) + `, pgx.RowToStructByName[checkpointOperationRow], after, limit) if err != nil { return nil, fmt.Errorf("list pending checkpoints: %w", err) } @@ -152,13 +165,23 @@ func (c *Client) ListPendingCheckpoints(ctx context.Context, afterID string, lim return result, nil } -func checkpointOperation(row sessionCheckpointRow) (*CheckpointOperation, error) { - checkpoint, err := toSessionCheckpoint(row) +func checkpointOperation(row checkpointOperationRow) (*CheckpointOperation, error) { + checkpoint, err := toSessionCheckpoint(row.sessionCheckpointRow) if err != nil { return nil, err } + var operationError *statuspb.Status + if len(row.OperationError) != 0 { + operationError = &statuspb.Status{} + if err := proto.Unmarshal(row.OperationError, operationError); err != nil { + return nil, fmt.Errorf("decode checkpoint operation error: %w", err) + } + if operationError.Code <= 0 || operationError.Code > 16 { + return nil, fmt.Errorf("invalid checkpoint operation error code %d", operationError.Code) + } + } return &CheckpointOperation{Checkpoint: checkpoint, Creator: row.UserID, - Snapshot: checkpointSnapshot(row), TagUID: row.TagUID}, nil + Snapshot: checkpointSnapshot(row.sessionCheckpointRow), TagUID: row.TagUID, OperationError: operationError}, nil } // BeginCheckpointCleanup durably records creation compensation and its tag UID. @@ -196,8 +219,35 @@ func (c *Client) BeginCheckpointCleanup(ctx context.Context, id string, executor if err != nil { return fmt.Errorf("encode checkpoint cleanup: %w", err) } - _, err = tx.Exec(ctx, `UPDATE session_checkpoint SET tag_uid = $2, data = $3 + _, err = tx.Exec(ctx, `UPDATE session_checkpoint SET tag_uid = $2, data = $3, operation_error = NULL WHERE id = $1 AND state = 'CREATING'`, row.ID, tagUID, data) return err }) } + +// RecordCheckpointOperationError preserves an attempt's permanent failure for +// waiters on any API replica without releasing its reservation or references. +// The current unexpired executor must own pending work; stale writes conflict. +func (c *Client) RecordCheckpointOperationError(ctx context.Context, id string, executor uuid.UUID, failure *statuspb.Status) error { + if failure == nil || failure.Code <= 0 || failure.Code > 16 { + return ErrFailedPrecondition + } + data, err := proto.Marshal(failure) + if err != nil { + return err + } + return c.withTx(ctx, func(tx pgx.Tx) error { + row, err := lockCheckpoint(ctx, tx, id, true, "", nil) + if err != nil { + return notFoundOr(err) + } + if err := checkCheckpointExecutor(ctx, tx, id, executor); err != nil { + return err + } + if row.State != "CREATING" && row.State != "DELETING" { + return ErrConflict + } + _, err = tx.Exec(ctx, `UPDATE session_checkpoint SET operation_error = $2 WHERE id = $1`, id, data) + return err + }) +} diff --git a/go/core/internal/database/checkpoint_operations_test.go b/go/core/internal/database/checkpoint_operations_test.go index 34f411ee39..b429eecf94 100644 --- a/go/core/internal/database/checkpoint_operations_test.go +++ b/go/core/internal/database/checkpoint_operations_test.go @@ -11,6 +11,8 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" "github.com/stretchr/testify/require" + statuspb "google.golang.org/genproto/googleapis/rpc/status" + "google.golang.org/grpc/codes" apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" ) @@ -102,6 +104,7 @@ func TestCheckpointExecutorFencesStaleTransitions(t *testing.T) { _, err = client.FinalizeClaimedSessionCheckpoint(t.Context(), checkpoint.Id, first, "old-tag", "s3://tags/old", "") require.ErrorIs(t, err, ErrConflict) require.ErrorIs(t, client.BeginCheckpointCleanup(t.Context(), checkpoint.Id, first, "old-tag", "old failure"), ErrConflict) + require.ErrorIs(t, client.RecordCheckpointOperationError(t.Context(), checkpoint.Id, first, &statuspb.Status{Code: int32(codes.FailedPrecondition), Message: "stale failure"}), ErrConflict) _, err = client.FinalizeSessionCheckpoint(t.Context(), checkpoint.Id, "unclaimed", "s3://tags/unclaimed", "") require.ErrorIs(t, err, ErrConflict) _, err = client.ClaimCheckpointOperation(otherCtx, checkpoint.Id, "alice", uuid.New()) @@ -208,3 +211,35 @@ func TestPendingCheckpointsPageBothPhasesAndRetainInputs(t *testing.T) { require.Equal(t, "retained-uid", seen[deleting.Id].TagUID) require.NotContains(t, seen, failed.Id) } + +func TestCheckpointOperationErrorRetainsPendingState(t *testing.T) { + client := NewClient(setupTestDB(t), "public") + checkpoint, _ := reserveRecoveryCheckpoint(t, client) + ctx, cancel := context.WithTimeout(t.Context(), RuntimeOperationTimeout) + defer cancel() + executor := uuid.New() + _, err := client.ClaimCheckpointOperation(ctx, checkpoint.Id, "alice", executor) + require.NoError(t, err) + failure := &statuspb.Status{Code: int32(codes.FailedPrecondition), Message: "tag was replaced"} + require.NoError(t, client.RecordCheckpointOperationError(ctx, checkpoint.Id, executor, failure)) + require.NoError(t, client.ReleaseCheckpointOperation(ctx, checkpoint.Id, executor)) + // Given a failed attempt, when another client reads the operation, assert + // that the failure and the pending reservation both survive release. + pending, err := client.GetCheckpointOperation(ctx, checkpoint.Id, "alice") + require.NoError(t, err) + require.Equal(t, failure.Code, pending.OperationError.Code) + require.Equal(t, failure.Message, pending.OperationError.Message) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_CREATING, pending.Checkpoint.State) + _, err = client.GetCheckpointOperation(ctx, checkpoint.Id, "bob") + require.ErrorIs(t, err, ErrNotFound) + next := uuid.New() + _, err = client.ClaimCheckpointOperation(ctx, checkpoint.Id, "alice", next) + require.NoError(t, err) + require.ErrorIs(t, client.RecordCheckpointOperationError(ctx, checkpoint.Id, executor, failure), ErrConflict) + _, err = client.FinalizeClaimedSessionCheckpoint(ctx, checkpoint.Id, next, "repaired-tag", "s3://snapshots/repaired", "") + require.NoError(t, err) + ready, err := client.GetCheckpointOperation(ctx, checkpoint.Id, "alice") + require.NoError(t, err) + require.Nil(t, ready.OperationError, "confirmed progress must clear the previous error") + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_READY, ready.Checkpoint.State) +} diff --git a/go/core/internal/database/checkpoints.go b/go/core/internal/database/checkpoints.go index 23c5dc17c8..b7bf736408 100644 --- a/go/core/internal/database/checkpoints.go +++ b/go/core/internal/database/checkpoints.go @@ -386,7 +386,7 @@ func (c *Client) finalizeSessionCheckpoint(ctx context.Context, id string, execu 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 + data = $4, operation_error = NULL WHERE id = $1 AND state = 'CREATING' `, row.ID, tagUID, snapshotURI, data) diff --git a/go/core/internal/service/checkpoint/recovery.go b/go/core/internal/service/checkpoint/recovery.go index b577ba2ad0..283767167d 100644 --- a/go/core/internal/service/checkpoint/recovery.go +++ b/go/core/internal/service/checkpoint/recovery.go @@ -4,12 +4,13 @@ import ( "context" "errors" "fmt" + "sync" "time" "github.com/google/uuid" - "golang.org/x/sync/errgroup" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" + "k8s.io/client-go/util/workqueue" "sigs.k8s.io/controller-runtime/pkg/manager" apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" @@ -18,39 +19,68 @@ import ( "github.com/kagent-dev/kagent/go/pkg/logging" ) +type checkpointKey struct{ id, owner string } + +type checkpointRetry struct { + backoff retryBackoff + readyAt time.Time +} + // RecoveryWorker is the sole owner of checkpoint tag mutations. API replicas // admit durable work and wait for its outcome; losing a caller does not cancel // the copy. On restart or leader change the worker resumes persisted phases. type RecoveryWorker struct { service *Service interval time.Duration + queue workqueue.TypedDelayingInterface[checkpointKey] + mu sync.Mutex + retries map[checkpointKey]*checkpointRetry } var _ manager.Runnable = (*RecoveryWorker)(nil) var _ manager.LeaderElectionRunnable = (*RecoveryWorker)(nil) func NewRecoveryWorker(service *Service) *RecoveryWorker { - return &RecoveryWorker{service: service, interval: time.Second} + return &RecoveryWorker{service: service, interval: time.Second, retries: make(map[checkpointKey]*checkpointRetry)} } func (*RecoveryWorker) NeedLeaderElection() bool { return true } func (w *RecoveryWorker) Start(ctx context.Context) error { - ticker := time.NewTicker(w.interval) + if w.service.activation != nil { + select { + case <-ctx.Done(): + return nil + case <-w.service.activation: + } + } + w.queue = workqueue.NewTypedDelayingQueueWithConfig(workqueue.TypedDelayingQueueConfig[checkpointKey]{Clock: w.service.clock}) + var workers sync.WaitGroup + for range 4 { + workers.Go(func() { + for w.process(ctx) { + } + }) + } + ticker := w.service.clock.NewTicker(w.interval) defer ticker.Stop() for ctx.Err() == nil { if err := w.sweep(ctx); err != nil && ctx.Err() == nil { - logging.FromContext(ctx).ErrorContext(ctx, "recover checkpoints", "error", err) + logging.FromContext(ctx).ErrorContext(ctx, "discover checkpoints", "error", err) } select { case <-ctx.Done(): - case <-ticker.C: + case <-ticker.C(): case <-w.service.wake: } } + w.queue.ShutDown() + workers.Wait() return nil } +// Discovery never waits for a copy. The keyed queue deduplicates IDs and keeps +// spare workers available to requests arriving behind a slow operation. func (w *RecoveryWorker) sweep(ctx context.Context) error { var after string for ctx.Err() == nil { @@ -58,18 +88,8 @@ func (w *RecoveryWorker) sweep(ctx context.Context) error { if err != nil { return err } - var group errgroup.Group - group.SetLimit(4) for _, op := range operations { - group.Go(func() error { - if err := w.service.recoverOperation(ctx, op); err != nil && ctx.Err() == nil { - logging.FromContext(ctx).ErrorContext(ctx, "recover checkpoint operation", "checkpoint_id", op.Checkpoint.Id, "error", err) - } - return nil - }) - } - if err := group.Wait(); err != nil { - return err + w.enqueue(checkpointKey{id: op.Checkpoint.Id, owner: op.Creator}) } if len(operations) < 100 { return nil @@ -79,6 +99,57 @@ func (w *RecoveryWorker) sweep(ctx context.Context) error { return ctx.Err() } +func (w *RecoveryWorker) enqueue(key checkpointKey) { + w.mu.Lock() + defer w.mu.Unlock() + if retry := w.retries[key]; retry != nil { + if delay := retry.readyAt.Sub(w.service.clock.Now()); delay > 0 { + w.queue.AddAfter(key, delay) + return + } + } + w.queue.Add(key) +} + +func (w *RecoveryWorker) process(ctx context.Context) bool { + key, shutdown := w.queue.Get() + if shutdown { + return false + } + defer w.queue.Done(key) + // A scan can mark a running key dirty just before it fails. Recheck here so + // Done's immediate requeue cannot bypass the newly established retry delay. + w.mu.Lock() + if retry := w.retries[key]; retry != nil { + if delay := retry.readyAt.Sub(w.service.clock.Now()); delay > 0 { + w.queue.AddAfter(key, delay) + w.mu.Unlock() + return true + } + } + w.mu.Unlock() + err := w.service.recoverOperation(ctx, &database.CheckpointOperation{Checkpoint: &apiv1.Checkpoint{Id: key.id}, Creator: key.owner}) + w.mu.Lock() + defer w.mu.Unlock() + if err == nil { + delete(w.retries, key) + return true + } + if ctx.Err() != nil { + return false + } + retry := w.retries[key] + if retry == nil { + retry = &checkpointRetry{backoff: retryBackoff{delay: time.Second, cap: 30 * time.Second}} + w.retries[key] = retry + } + delay := retry.backoff.next() + retry.readyAt = w.service.clock.Now().Add(delay) + w.queue.AddAfter(key, delay) + logging.FromContext(ctx).ErrorContext(ctx, "recover checkpoint operation", "checkpoint_id", key.id, "retry_after", delay, "error", err) + return true +} + func (s *Service) recoverOperation(ctx context.Context, observed *database.CheckpointOperation) (err error) { // Establish the attempt deadline before claiming. A paused executor cannot // issue another RPC after its claim expires; unfinished copies synchronize @@ -102,6 +173,13 @@ func (s *Service) recoverOperation(ctx context.Context, observed *database.Check if err != nil { return err } + defer func() { + if failure := permanentOperationError(err); failure != nil { + recordCtx, stop := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + defer stop() + err = errors.Join(err, s.store.RecordCheckpointOperationError(recordCtx, id, executor, failure.Proto())) + } + }() if err := ctx.Err(); err != nil { return err } diff --git a/go/core/internal/service/checkpoint/retry.go b/go/core/internal/service/checkpoint/retry.go new file mode 100644 index 0000000000..1b9542bb13 --- /dev/null +++ b/go/core/internal/service/checkpoint/retry.go @@ -0,0 +1,68 @@ +package checkpoint + +import ( + "context" + "errors" + "math/rand/v2" + "time" + + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/kagent-dev/kagent/go/core/internal/service/serviceerrors" +) + +type retryBackoff struct { + delay time.Duration + cap time.Duration +} + +// Keep jitter at the cap too: synchronized callers must not converge onto the +// same interval after a long wait. Every delay is within 90-100% of its bound. +func (b *retryBackoff) next() time.Duration { + delay := b.delay + b.delay = min(2*b.delay, b.cap) + return delay - time.Duration(rand.Int64N(int64(delay/10)+1)) +} + +func (s *Service) waitRetry(ctx context.Context, delay time.Duration) error { + timer := s.clock.NewTimer(delay) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C(): + return nil + } +} + +func permanentOperationError(err error) *status.Status { + var failure *status.Status + message := "Checkpoint recovery cannot continue safely" + var public *serviceerrors.Error + if errors.As(err, &public) { + message = public.Message() + } + switch serviceerrors.CodeOf(err) { + case serviceerrors.CodeFailedPrecondition: + failure = status.New(codes.FailedPrecondition, message) + case serviceerrors.CodeInvalidArgument: + failure = status.New(codes.InvalidArgument, message) + case serviceerrors.CodePermissionDenied: + failure = status.New(codes.PermissionDenied, message) + case serviceerrors.CodeUnauthenticated: + failure = status.New(codes.Unauthenticated, message) + case "": + failure = status.Convert(err) + default: + return nil + } + switch failure.Code() { + case codes.FailedPrecondition, codes.InvalidArgument, codes.PermissionDenied, codes.Unauthenticated: + // Backend status text and details can quote private configuration. Only + // service-owned public messages cross this durable API error boundary. + return status.New(failure.Code(), message) + default: + return nil + } +} diff --git a/go/core/internal/service/checkpoint/review_test.go b/go/core/internal/service/checkpoint/review_test.go new file mode 100644 index 0000000000..af9be23ef0 --- /dev/null +++ b/go/core/internal/service/checkpoint/review_test.go @@ -0,0 +1,330 @@ +package checkpoint + +import ( + "context" + "errors" + "sync" + "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" + "google.golang.org/protobuf/proto" + clocktesting "k8s.io/utils/clock/testing" + + 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/internal/substrate" + "github.com/kagent-dev/kagent/go/core/pkg/auth" +) + +func TestPublicDeleteRetainsPermanentFailure(t *testing.T) { + store := reviewStore(t) + source, head := recoverySource(t, store) + tags := &lockedTags{base: &testTags{snapshotURI: "s3://snapshots/source"}} + service := NewService(store, testAuthorizer{}, tags, nil) + owner := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) + cp, err := createAndRecover(service, owner, source.Id, uuid.NewString(), head) + require.NoError(t, err) + tags.mu.Lock() + tags.base.created.Metadata.Uid = "replacement-uid" + tags.mu.Unlock() + stop := startRecovery(t, service) + // Given a replaced tag, when a separate API replica calls public Delete, + // assert that it receives the worker's failure without waiting for timeout. + follower := NewService(store, testAuthorizer{}, tags, nil) + ctx, cancel := context.WithTimeout(owner, 3*time.Second) + defer cancel() + require.Equal(t, codes.FailedPrecondition, status.Code(follower.Delete(ctx, cp.Id))) + stop() + pending, err := store.GetCheckpointOperation(t.Context(), cp.Id, "alice") + require.NoError(t, err) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_DELETING, pending.Checkpoint.State) + require.Equal(t, int32(codes.FailedPrecondition), pending.OperationError.GetCode()) + require.Equal(t, "tag-uid", pending.TagUID) + tags.mu.Lock() + require.Zero(t, tags.base.deleteCalls) + tags.mu.Unlock() + // A fresh service with no worker or local error cache must observe the same + // durable result, while another owner cannot read the operation error. + follower = NewService(store, testAuthorizer{}, tags, nil) + require.Equal(t, codes.FailedPrecondition, status.Code(follower.Delete(ctx, cp.Id))) + _, err = store.GetCheckpointOperation(t.Context(), cp.Id, "bob") + require.ErrorIs(t, err, database.ErrNotFound) +} + +func TestPermanentRecoveryErrorsKeepBackendConfigurationPrivate(t *testing.T) { + const private = "private-credential-value" + for _, scenario := range []struct { + name string + failure error + message string + }{ + {name: "backend status", failure: status.Error(codes.FailedPrecondition, private), message: "Checkpoint recovery cannot continue safely"}, + {name: "service error with private cause", failure: serviceerrors.NewFailedPrecondition("Snapshot identity changed", errors.New(private)), message: "Snapshot identity changed"}, + } { + t.Run("Given a "+scenario.name+", when recovery persists a permanent error, assert that callers see only the safe message", func(t *testing.T) { + failure := permanentOperationError(scenario.failure) + require.Equal(t, codes.FailedPrecondition, failure.Code()) + require.Equal(t, scenario.message, failure.Message()) + require.NotContains(t, failure.Message(), private) + }) + } +} + +func TestRecoveryDiscoversDuringSlowCopy(t *testing.T) { + store := reviewStore(t) + slow, slowHead := recoverySource(t, store) + fast, fastHead := recoverySource(t, store) + gate := make(chan struct{}) + var release sync.Once + defer release.Do(func() { close(gate) }) + slowTags := &lockedTags{base: &testTags{snapshotURI: "s3://snapshots/source"}, gate: gate, entered: make(chan struct{})} + tags := &indexedTags{actors: map[string]*lockedTags{ + substrate.ActorName(slow.Id): slowTags, + substrate.ActorName(fast.Id): {base: &testTags{snapshotURI: "s3://snapshots/source"}}, + }} + service := NewService(store, testAuthorizer{}, tags, nil) + stop := startRecovery(t, service) + defer stop() + owner := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) + slowDone := make(chan error, 1) + go func() { _, err := service.Create(owner, slow.Id, uuid.NewString(), slowHead); slowDone <- err }() + select { + case <-slowTags.entered: + case <-time.After(5 * time.Second): + t.Fatal("slow copy did not start") + } + // Given a blocked copy, when a later request arrives, assert that a free + // worker publishes it while the original copy is still blocked. + ctx, cancel := context.WithTimeout(owner, 2*time.Second) + defer cancel() + cp, err := service.Create(ctx, fast.Id, uuid.NewString(), fastHead) + require.NoError(t, err) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_READY, cp.State) + select { + case err := <-slowDone: + t.Fatalf("slow copy unexpectedly finished: %v", err) + default: + } + release.Do(func() { close(gate) }) + require.NoError(t, <-slowDone) +} + +func TestRecoveryScansRespectRetryDelay(t *testing.T) { + store := reviewStore(t) + source, head := recoverySource(t, store) + clock := clocktesting.NewFakeClock(time.Now()) + firstFailure := make(chan struct{}) + tags := &countedFailingTags{first: firstFailure, entered: make(chan struct{}), lockedTags: &lockedTags{base: &testTags{snapshotURI: "s3://snapshots/source", createErr: status.Error(codes.PermissionDenied, "storage denied")}}} + service := NewService(store, testAuthorizer{}, tags, nil) + service.clock = clock + owner := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) + cp, _, err := service.beginCreate(owner, source.Id, uuid.NewString(), head) + require.NoError(t, err) + worker := NewRecoveryWorker(service) + worker.interval = 10 * time.Millisecond + ctx, cancel := context.WithCancel(t.Context()) + done := make(chan error, 1) + go func() { done <- worker.Start(ctx) }() + t.Cleanup(func() { cancel(); require.NoError(t, <-done) }) + waitRetry := func(attempt int64) { + require.Eventually(t, func() bool { + worker.mu.Lock() + defer worker.mu.Unlock() + retry := worker.retries[checkpointKey{id: cp.Id, owner: "alice"}] + return tags.calls.Load() >= attempt && retry != nil && retry.readyAt.After(clock.Now()) + }, 3*time.Second, time.Millisecond) + } + select { + case <-tags.entered: + case <-time.After(3 * time.Second): + t.Fatal("first attempt did not start") + } + // A discovery while the first attempt runs marks the queue key dirty. The + // post-failure immediate requeue must still respect the new retry deadline. + require.NoError(t, worker.sweep(t.Context())) + close(firstFailure) + waitRetry(1) + // Given a persistent failure, when discovery is repeatedly woken before + // retry eligibility, assert that it makes no extra backend call or unpins. + for range 8 { + clock.Step(100 * time.Millisecond) + service.wakeRecovery() + } + require.Never(t, func() bool { return tags.calls.Load() > 1 }, 100*time.Millisecond, time.Millisecond) + _, err = store.BeginSessionOperation(t.Context(), source.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE) + require.ErrorIs(t, err, database.ErrConflict) + clock.Step(300 * time.Millisecond) + service.wakeRecovery() + waitRetry(2) + // The second retry has a 1.8-2 second delay, not the first one's interval. + clock.Step(1500 * time.Millisecond) + service.wakeRecovery() + require.Equal(t, int64(2), tags.calls.Load()) + clock.Step(600 * time.Millisecond) + waitRetry(3) + // Repairing the backend eventually completes the same admitted operation. + tags.mu.Lock() + tags.base.createErr = nil + tags.mu.Unlock() + clock.Step(5 * time.Second) + service.wakeRecovery() + waitCheckpointPhase(t, store, cp.Id, apiv1.CheckpointState_CHECKPOINT_STATE_READY) +} + +func TestPublicCheckpointWaitBackoff(t *testing.T) { + for _, operation := range []string{"Create", "Delete"} { + t.Run("Given pending work, when public "+operation+" waits for a minute, assert that reads stay bounded and cancellation returns promptly", func(t *testing.T) { + store := &countedReadStore{testStore: &testStore{ + prepared: &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: uuid.NewString(), State: apiv1.CheckpointState_CHECKPOINT_STATE_CREATING}, + snapshot: &database.SessionTaskSnapshot{}, + }} + clock := clocktesting.NewFakeClock(time.Now()) + service := NewService(store, testAuthorizer{}, nil, nil) + service.clock = clock + ctx, cancel := context.WithCancel(auth.AuthSessionTo(t.Context(), testSession{userID: "alice"})) + defer cancel() + done := make(chan error, 1) + go func() { + if operation == "Create" { + _, err := service.Create(ctx, store.prepared.SessionId, "request", "head") + done <- err + } else { + done <- service.Delete(ctx, store.prepared.Id) + } + }() + for range 6000 { + require.Eventually(t, clock.HasWaiters, time.Second, time.Millisecond) + clock.Step(10 * time.Millisecond) + } + cancel() + select { + case err := <-done: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(time.Second): + t.Fatal("canceled caller remained blocked") + } + t.Logf("60 seconds of waiting performed %d reads", store.reads.Load()) + require.Less(t, store.reads.Load(), int64(80)) + require.Greater(t, store.reads.Load(), int64(55), "waiters must keep checking durable progress") + }) + } +} + +func TestCheckpointActivationPreservesAdmittedWork(t *testing.T) { + store := reviewStore(t) + source, head := recoverySource(t, store) + tags := &lockedTags{base: &testTags{snapshotURI: "s3://snapshots/source"}, entered: make(chan struct{})} + service := NewService(store, 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) + ready := make(chan struct{}) + service.GateRecovery(ready) + stop := startRecovery(t, service) + defer stop() + // Given admitted legacy work, when replacements start before writers drain, + // assert that public mutations fail promptly and no worker starts a copy. + _, err = service.Create(owner, source.Id, request, head) + require.Equal(t, serviceerrors.CodeUnavailable, serviceerrors.CodeOf(err)) + require.Equal(t, serviceerrors.CodeUnavailable, serviceerrors.CodeOf(service.Delete(owner, cp.Id))) + select { + case <-tags.entered: + t.Fatal("recovery executed before activation") + default: + } + pending, err := store.GetCheckpointOperation(t.Context(), cp.Id, "alice") + require.NoError(t, err) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_CREATING, pending.Checkpoint.State) + _, err = store.BeginSessionOperation(t.Context(), source.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE) + require.ErrorIs(t, err, database.ErrConflict) + close(ready) + published, err := service.Create(owner, source.Id, request, head) + require.NoError(t, err) + require.Equal(t, cp.Id, published.Id) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_READY, published.State) +} + +func reviewStore(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") +} + +type countedReadStore struct { + *testStore + reads atomic.Int64 +} + +func (s *countedReadStore) GetCheckpointOperation(_ context.Context, _ string, owner string) (*database.CheckpointOperation, error) { + s.reads.Add(1) + return &database.CheckpointOperation{Checkpoint: proto.CloneOf(s.prepared), Creator: owner}, nil +} + +type countedFailingTags struct { + *lockedTags + calls atomic.Int64 + first <-chan struct{} + entered chan struct{} +} + +func (t *countedFailingTags) CreateTag(ctx context.Context, space, name, actor string) (*ateapipb.Tag, error) { + call := t.calls.Add(1) + if call == 1 && t.first != nil { + close(t.entered) + select { + case <-t.first: + case <-ctx.Done(): + return nil, ctx.Err() + } + } + t.mu.Lock() + failure := t.base.createErr + t.mu.Unlock() + if failure != nil { + return nil, failure + } + return t.lockedTags.CreateTag(ctx, space, name, actor) +} + +type indexedTags struct { + actors map[string]*lockedTags + tags sync.Map +} + +func (t *indexedTags) GetActor(ctx context.Context, space, name string) (*ateapipb.Actor, error) { + return t.actors[name].GetActor(ctx, space, name) +} + +func (t *indexedTags) CreateTag(ctx context.Context, space, name, actor string) (*ateapipb.Tag, error) { + t.tags.Store(name, t.actors[actor]) + return t.actors[actor].CreateTag(ctx, space, name, actor) +} + +func (t *indexedTags) GetTag(ctx context.Context, space, name string) (*ateapipb.Tag, error) { + client, ok := t.tags.Load(name) + if !ok { + return nil, status.Error(codes.NotFound, "tag absent") + } + return client.(*lockedTags).GetTag(ctx, space, name) +} + +func (t *indexedTags) DeleteTag(ctx context.Context, space, name, uid string) error { + client, ok := t.tags.Load(name) + if !ok { + return status.Error(codes.NotFound, "tag absent") + } + return client.(*lockedTags).DeleteTag(ctx, space, name, uid) +} diff --git a/go/core/internal/service/checkpoint/service.go b/go/core/internal/service/checkpoint/service.go index c777b34445..5da85fd156 100644 --- a/go/core/internal/service/checkpoint/service.go +++ b/go/core/internal/service/checkpoint/service.go @@ -11,8 +11,10 @@ import ( "github.com/google/uuid" ateapipb "go.ate.dev/substrate/api/ateapi/v1" "google.golang.org/genproto/googleapis/rpc/errdetails" + statuspb "google.golang.org/genproto/googleapis/rpc/status" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" + "k8s.io/utils/clock" apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" "github.com/kagent-dev/kagent/go/core/internal/database" @@ -22,8 +24,9 @@ import ( ) const ( - defaultPageSize = 50 - maxPageSize = 100 + defaultPageSize = 50 + maxPageSize = 100 + recoveryWaitingMessage = "Checkpoint recovery is waiting for previous controller writers to exit" ) type store interface { @@ -31,6 +34,7 @@ type store interface { 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 + RecordCheckpointOperationError(context.Context, string, uuid.UUID, *statuspb.Status) error BeginCheckpointCleanup(context.Context, string, uuid.UUID, string, string) error ReserveSessionCheckpoint(context.Context, *apiv1.Checkpoint, string, string) (*apiv1.Checkpoint, *database.SessionTaskSnapshot, error) FinalizeClaimedSessionCheckpoint(context.Context, string, uuid.UUID, string, string, string) (*apiv1.Checkpoint, error) @@ -63,6 +67,8 @@ type Service struct { tags tagClient workflow workflow wake chan struct{} + clock clock.WithTicker + activation <-chan struct{} } type ListRequest struct { @@ -77,7 +83,24 @@ type ListResult struct { } func NewService(store store, authorizer auth.Authorizer, tags tagClient, workflow workflow) *Service { - return &Service{store: store, authorizer: authorizer, tags: tags, workflow: workflow, wake: make(chan struct{}, 1)} + return &Service{store: store, authorizer: authorizer, tags: tags, workflow: workflow, wake: make(chan struct{}, 1), clock: clock.RealClock{}} +} + +// GateRecovery installs the controller startup barrier before this service is +// used. All API replicas reject new mutations until legacy writers have exited; +// the worker waits on the same barrier. Other checkpoint reads remain available. +func (s *Service) GateRecovery(ready <-chan struct{}) { s.activation = ready } + +func (s *Service) recoveryActive() bool { + if s.activation == nil { + return true + } + select { + case <-s.activation: + return true + default: + return false + } } func (s *Service) Create(ctx context.Context, sessionID, requestID, expectedHeadTaskID string) (*apiv1.Checkpoint, error) { @@ -107,6 +130,9 @@ func (s *Service) beginCreate(ctx context.Context, sessionID, requestID, expecte if _, err := s.authorize(ctx, auth.VerbGet, auth.Resource{Type: "Session", Name: sessionID}); err != nil { return nil, "", err } + if !s.recoveryActive() { + return nil, "", serviceerrors.NewUnavailable(recoveryWaitingMessage, nil) + } id, err := uuid.NewV7() if err != nil { return nil, "", serviceerrors.NewInternal("Failed to generate checkpoint identifier", err) @@ -286,6 +312,9 @@ func (s *Service) beginDelete(ctx context.Context, checkpointID string) (string, if err != nil { return "", err } + if !s.recoveryActive() { + return "", serviceerrors.NewUnavailable(recoveryWaitingMessage, nil) + } _, _, err = s.store.BeginDeleteSessionCheckpoint(ctx, checkpointID, userID) if errors.Is(err, database.ErrNotFound) { return "", serviceerrors.NewNotFound("Checkpoint not found", err) @@ -481,8 +510,7 @@ func (s *Service) wakeRecovery() { } func (s *Service) waitCreate(ctx context.Context, id, owner string) (*apiv1.Checkpoint, error) { - ticker := time.NewTicker(100 * time.Millisecond) - defer ticker.Stop() + backoff := retryBackoff{delay: 100 * time.Millisecond, cap: time.Second} for { op, err := s.store.GetCheckpointOperation(ctx, id, owner) if err != nil { @@ -502,19 +530,19 @@ func (s *Service) waitCreate(ctx context.Context, id, owner string) (*apiv1.Chec case apiv1.CheckpointState_CHECKPOINT_STATE_FAILED: return nil, serviceerrors.NewUnavailable("Failed to retain checkpoint snapshot", errors.New(op.Checkpoint.GetFailure().GetMessage())) } - select { - case <-ctx.Done(): - return nil, ctx.Err() - case <-ticker.C: + if op.OperationError != nil { + return nil, status.FromProto(op.OperationError).Err() + } + if err := s.waitRetry(ctx, backoff.next()); err != nil { + return nil, err } } } func (s *Service) waitDeleted(ctx context.Context, id, owner string) error { - ticker := time.NewTicker(100 * time.Millisecond) - defer ticker.Stop() + backoff := retryBackoff{delay: 100 * time.Millisecond, cap: time.Second} for { - _, err := s.store.GetCheckpointOperation(ctx, id, owner) + op, err := s.store.GetCheckpointOperation(ctx, id, owner) if errors.Is(err, database.ErrNotFound) { return nil } @@ -524,10 +552,11 @@ func (s *Service) waitDeleted(ctx context.Context, id, owner string) error { } return serviceerrors.NewInternal("Failed to read checkpoint deletion", err) } - select { - case <-ctx.Done(): - return ctx.Err() - case <-ticker.C: + if op.OperationError != nil { + return status.FromProto(op.OperationError).Err() + } + if err := s.waitRetry(ctx, backoff.next()); err != nil { + return err } } } diff --git a/go/core/internal/service/checkpoint/service_test.go b/go/core/internal/service/checkpoint/service_test.go index b89d1f88e6..151ad894a9 100644 --- a/go/core/internal/service/checkpoint/service_test.go +++ b/go/core/internal/service/checkpoint/service_test.go @@ -13,6 +13,7 @@ import ( "github.com/stretchr/testify/require" ateapipb "go.ate.dev/substrate/api/ateapi/v1" "google.golang.org/genproto/googleapis/rpc/errdetails" + statuspb "google.golang.org/genproto/googleapis/rpc/status" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" "google.golang.org/protobuf/types/known/timestamppb" @@ -964,3 +965,7 @@ func (s *testStore) ClaimCheckpointOperation(ctx context.Context, id, owner stri func (*testStore) ReleaseCheckpointOperation(context.Context, string, uuid.UUID) error { return nil } + +func (s *testStore) RecordCheckpointOperationError(context.Context, string, uuid.UUID, *statuspb.Status) error { + return nil +} diff --git a/go/core/pkg/app/app.go b/go/core/pkg/app/app.go index 8ba214b463..52b0788698 100644 --- a/go/core/pkg/app/app.go +++ b/go/core/pkg/app/app.go @@ -33,6 +33,7 @@ import ( "google.golang.org/protobuf/proto" corev1 "k8s.io/api/core/v1" k8sruntime "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/types" utilruntime "k8s.io/apimachinery/pkg/util/runtime" clientgoscheme "k8s.io/client-go/kubernetes/scheme" typedauth "k8s.io/client-go/kubernetes/typed/authentication/v1" @@ -73,6 +74,7 @@ import ( "github.com/kagent-dev/kagent/go/core/pkg/auth" "github.com/kagent-dev/kagent/go/core/pkg/config" "github.com/kagent-dev/kagent/go/core/pkg/egresspolicy" + kagentenv "github.com/kagent-dev/kagent/go/core/pkg/env" "github.com/kagent-dev/kagent/go/core/pkg/migrations" "github.com/kagent-dev/kagent/go/core/pkg/postgres" "github.com/kagent-dev/kagent/go/pkg/logging" @@ -367,6 +369,12 @@ 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) + activation := &checkpointActivation{reader: manager.GetAPIReader(), namespace: resourceNamespace, + podName: kagentenv.ControllerPodName.Get(), podUID: types.UID(kagentenv.ControllerPodUID.Get()), ready: make(chan struct{})} + checkpoints.GateRecovery(activation.ready) + if err := manager.Add(activation); err != nil { + return fmt.Errorf("register checkpoint activation barrier: %w", err) + } if err := manager.Add(checkpoint.NewRecoveryWorker(checkpoints)); err != nil { return fmt.Errorf("register checkpoint recovery worker: %w", err) } diff --git a/go/core/pkg/app/checkpoint_activation.go b/go/core/pkg/app/checkpoint_activation.go new file mode 100644 index 0000000000..a5d2e45ff9 --- /dev/null +++ b/go/core/pkg/app/checkpoint_activation.go @@ -0,0 +1,111 @@ +package app + +import ( + "context" + "fmt" + "time" + + appsv1 "k8s.io/api/apps/v1" + corev1 "k8s.io/api/core/v1" + "k8s.io/apimachinery/pkg/types" + "sigs.k8s.io/controller-runtime/pkg/client" + + "github.com/kagent-dev/kagent/go/pkg/logging" +) + +const ( + checkpointClaimsAnnotation = "kagent.dev/checkpoint-execution-claims" + checkpointClaimsVersion = "v1" + controllerComponentLabel = "app.kubernetes.io/component" +) + +// This barrier runs on every API replica, independently of leader election and +// readiness. A replacement can serve other API traffic and become Ready while +// the Deployment finishes rolling its old replicas out. Only then can recovery +// or new checkpoint mutations start. Use the live API reader, never the cache. +type checkpointActivation struct { + reader client.Reader + namespace, podName string + podUID types.UID + ready chan struct{} +} + +func (*checkpointActivation) NeedLeaderElection() bool { return false } + +func (g *checkpointActivation) Start(ctx context.Context) error { + ticker := time.NewTicker(time.Second) + defer ticker.Stop() + for { + drained, err := g.drained(ctx) + if err != nil && ctx.Err() == nil { + logging.FromContext(ctx).ErrorContext(ctx, "check previous checkpoint writers", "error", err) + } + if err == nil && drained { + close(g.ready) + <-ctx.Done() + return nil + } + select { + case <-ctx.Done(): + return nil + case <-ticker.C: + } + } +} + +func (g *checkpointActivation) drained(ctx context.Context) (bool, error) { + selector := client.MatchingLabels{controllerComponentLabel: "controller"} + if g.podName != "" { + var self corev1.Pod + if err := g.reader.Get(ctx, types.NamespacedName{Namespace: g.namespace, Name: g.podName}, &self); err != nil { + return false, err + } + if self.UID != g.podUID || g.podUID == "" { + return false, fmt.Errorf("controller pod identity changed") + } + if self.Labels[controllerComponentLabel] != "controller" { + return false, fmt.Errorf("controller pod lacks component identity") + } + // Never include version or pod-template-hash: those would hide the old + // generation. These are the chart's stable Deployment selector labels. + for _, key := range []string{"app.kubernetes.io/name", "app.kubernetes.io/instance", controllerComponentLabel} { + if self.Labels[key] == "" { + return false, fmt.Errorf("controller pod lacks selector label %s", key) + } + selector[key] = self.Labels[key] + } + } + // A failed old Pod disappearing is insufficient while its ReplicaSet can + // recreate it. Confirm scale-down before the final Pod read; status must + // acknowledge that generation and report no remaining replicas. + var replicas appsv1.ReplicaSetList + if err := g.reader.List(ctx, &replicas, client.InNamespace(g.namespace), selector); err != nil { + return false, err + } + for _, rs := range replicas.Items { + if rs.Spec.Template.Annotations[checkpointClaimsAnnotation] == checkpointClaimsVersion { + continue + } + if rs.Spec.Replicas == nil || *rs.Spec.Replicas != 0 || rs.Status.ObservedGeneration < rs.Generation || rs.Status.Replicas != 0 { + return false, nil + } + } + var pods corev1.PodList + if err := g.reader.List(ctx, &pods, client.InNamespace(g.namespace), selector); err != nil { + return false, err + } + for _, pod := range pods.Items { + if pod.UID == g.podUID && g.podUID != "" { + continue + } + if pod.Status.Phase == corev1.PodSucceeded || pod.Status.Phase == corev1.PodFailed { + continue + } + // Unready and terminating legacy pods can still issue writes. Only + // process exit or a compatible protocol marker makes them safe. + if pod.Annotations[checkpointClaimsAnnotation] != checkpointClaimsVersion { + return false, nil + } + } + return true, nil +} diff --git a/go/core/pkg/app/checkpoint_activation_test.go b/go/core/pkg/app/checkpoint_activation_test.go new file mode 100644 index 0000000000..e61303b3c5 --- /dev/null +++ b/go/core/pkg/app/checkpoint_activation_test.go @@ -0,0 +1,144 @@ +package app + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/require" + appsv1 "k8s.io/api/apps/v1" + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/runtime" + "sigs.k8s.io/controller-runtime/pkg/client/fake" +) + +func TestCheckpointActivationDrainsLegacyWriters(t *testing.T) { + for _, scenario := range []struct { + name string + phase corev1.PodPhase + terminating, compatible, otherRelease bool + ready bool + }{ + {name: "running legacy writer", phase: corev1.PodRunning}, + {name: "unready legacy writer", phase: corev1.PodPending}, + {name: "terminating legacy writer", phase: corev1.PodRunning, terminating: true}, + {name: "exited legacy writer", phase: corev1.PodSucceeded, ready: true}, + {name: "failed legacy writer", phase: corev1.PodFailed, ready: true}, + {name: "compatible writer", phase: corev1.PodRunning, compatible: true, ready: true}, + {name: "another release", phase: corev1.PodRunning, otherRelease: true, ready: true}, + } { + t.Run("Given a "+scenario.name+", when recovery checks activation, assert that only drained or compatible writers permit execution", func(t *testing.T) { + scheme := runtime.NewScheme() + require.NoError(t, corev1.AddToScheme(scheme)) + require.NoError(t, appsv1.AddToScheme(scheme)) + self := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{Namespace: "kagent", Name: "new", UID: "new-uid", Labels: map[string]string{ + "app.kubernetes.io/name": "custom-name", "app.kubernetes.io/instance": "release-a", controllerComponentLabel: "controller", + "pod-template-hash": "new-generation", + }}} + old := self.DeepCopy() + old.Name = "previous" + old.UID = "old-uid" + old.Labels["pod-template-hash"] = "old-generation" + old.Status.Phase = scenario.phase + if scenario.compatible { + old.Annotations = map[string]string{checkpointClaimsAnnotation: checkpointClaimsVersion} + } + if scenario.otherRelease { + old.Labels["app.kubernetes.io/instance"] = "release-b" + } + if scenario.terminating { + now := metav1.Now() + old.DeletionTimestamp = &now + old.Finalizers = []string{"test/retain"} + } + reader := fake.NewClientBuilder().WithScheme(scheme).WithObjects(self, old).Build() + gate := &checkpointActivation{reader: reader, namespace: self.Namespace, podName: self.Name, podUID: self.UID} + ready, err := gate.drained(t.Context()) + require.NoError(t, err) + require.Equal(t, scenario.ready, ready) + require.False(t, gate.NeedLeaderElection(), "followers must activate their API admission too") + }) + } +} + +func TestCheckpointActivationFailsClosedOnMissingIdentity(t *testing.T) { + scheme := runtime.NewScheme() + require.NoError(t, corev1.AddToScheme(scheme)) + require.NoError(t, appsv1.AddToScheme(scheme)) + reader := fake.NewClientBuilder().WithScheme(scheme).Build() + gate := &checkpointActivation{reader: reader, namespace: "kagent", podName: "missing", podUID: "missing-uid"} + ready, err := gate.drained(t.Context()) + require.Error(t, err) + require.False(t, ready) + // Local controller development has no Pod identity and is safe only after + // any old controller Pods in the namespace have been stopped. + gate.podName = "" + gate.podUID = "" + ready, err = gate.drained(t.Context()) + require.NoError(t, err) + require.True(t, ready) +} + +func TestCheckpointActivationWaitsForLegacyReplicaSetScaleDown(t *testing.T) { + for _, scenario := range []struct { + name string + desired int32 + observed int64 + remaining int32 + ready bool + }{ + {name: "old Pod can be recreated", desired: 1, observed: 2}, + {name: "scale-down not yet observed", observed: 1}, + {name: "old replica still reported", observed: 2, remaining: 1}, + {name: "scale-down completed", observed: 2, ready: true}, + } { + t.Run("Given no live legacy Pods, when "+scenario.name+", assert that activation requires confirmed retirement", func(t *testing.T) { + scheme := runtime.NewScheme() + require.NoError(t, corev1.AddToScheme(scheme)) + require.NoError(t, appsv1.AddToScheme(scheme)) + previous := &appsv1.ReplicaSet{ObjectMeta: metav1.ObjectMeta{Namespace: "kagent", Name: "previous", Generation: 2, Labels: map[string]string{controllerComponentLabel: "controller"}}, + Spec: appsv1.ReplicaSetSpec{Replicas: new(scenario.desired)}, Status: appsv1.ReplicaSetStatus{ObservedGeneration: scenario.observed, Replicas: scenario.remaining}, + } + reader := fake.NewClientBuilder().WithScheme(scheme).WithObjects(previous).Build() + gate := &checkpointActivation{reader: reader, namespace: "kagent"} + ready, err := gate.drained(t.Context()) + require.NoError(t, err) + require.Equal(t, scenario.ready, ready) + }) + } +} + +func TestCheckpointActivationOpensAfterLegacyExit(t *testing.T) { + scheme := runtime.NewScheme() + require.NoError(t, corev1.AddToScheme(scheme)) + require.NoError(t, appsv1.AddToScheme(scheme)) + legacy := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{Namespace: "kagent", Name: "legacy", UID: "legacy-uid", Labels: map[string]string{controllerComponentLabel: "controller"}}} + reader := fake.NewClientBuilder().WithScheme(scheme).WithObjects(legacy).Build() + gate := &checkpointActivation{reader: reader, namespace: "kagent", ready: make(chan struct{})} + ctx, cancel := context.WithCancel(t.Context()) + done := make(chan error, 1) + go func() { done <- gate.Start(ctx) }() + t.Cleanup(func() { cancel(); require.NoError(t, <-done) }) + // Given a non-leader API's closed startup gate, when the legacy Pod exits, + // assert that the running gate opens and continues until shutdown. + require.Never(t, func() bool { + select { + case <-gate.ready: + return true + default: + return false + } + }, 20*time.Millisecond, time.Millisecond) + require.NoError(t, reader.Delete(t.Context(), legacy)) + select { + case <-gate.ready: + case <-time.After(2 * time.Second): + t.Fatal("activation remained closed after legacy retirement") + } + select { + case err := <-done: + t.Fatalf("gate stopped before controller shutdown: %v", err) + default: + } +} diff --git a/go/core/pkg/env/controller.go b/go/core/pkg/env/controller.go index d18fe0f682..c5acaa2e7c 100644 --- a/go/core/pkg/env/controller.go +++ b/go/core/pkg/env/controller.go @@ -2,3 +2,9 @@ package env // KUBECONFIG remains an external client-go interface. var Kubeconfig = RegisterStringVar("KUBECONFIG", "", "Kubernetes client configuration file list. When unset, client-go uses its normal in-cluster or user kubeconfig discovery.", ComponentController, ComponentCLI, ComponentTesting) + +// The chart projects these identities through the Kubernetes Downward API. +var ( + ControllerPodName = RegisterStringVar("K8S_POD_NAME", "", "Controller Pod name projected by Kubernetes for startup coordination and telemetry.", ComponentController) + ControllerPodUID = RegisterStringVar("K8S_POD_UID", "", "Controller Pod UID projected by Kubernetes for startup coordination and telemetry.", ComponentController) +) diff --git a/go/core/pkg/migrations/core/000002_checkpoint_recovery.sql b/go/core/pkg/migrations/core/000002_checkpoint_recovery.sql index b65e29882d..740c56f4d5 100644 --- a/go/core/pkg/migrations/core/000002_checkpoint_recovery.sql +++ b/go/core/pkg/migrations/core/000002_checkpoint_recovery.sql @@ -1,6 +1,7 @@ -- +goose Up ALTER TABLE session_checkpoint ADD COLUMN executor_id uuid, + ADD COLUMN operation_error bytea, ADD COLUMN executor_expires_at timestamptz, ADD CONSTRAINT session_checkpoint_executor_pair CHECK ( (executor_id IS NULL) = (executor_expires_at IS NULL) @@ -14,4 +15,5 @@ DROP INDEX session_checkpoint_pending_idx; ALTER TABLE session_checkpoint DROP CONSTRAINT session_checkpoint_executor_pair, DROP COLUMN executor_expires_at, - DROP COLUMN executor_id; + DROP COLUMN executor_id, + DROP COLUMN operation_error; diff --git a/go/core/test/e2e/README.md b/go/core/test/e2e/README.md index db94876051..5cb0a1e97f 100644 --- a/go/core/test/e2e/README.md +++ b/go/core/test/e2e/README.md @@ -329,3 +329,13 @@ missing worker capacity fails the test rather than skipping continuation coverag There are no retries around list calls, so a Substrate token-format change fails on continuation. Run with `KAGENT_E2E_API_URL` and the normal kagent harness fixture configured, using `go test ./core/test/e2e -run '^TestSubstrateInventoryPagination$'`. + +`TestCheckpointActivationDuringControllerReplacement` keeps an unmarked, unready +Pod alive while replacing the actual controller. The replacement must become +Ready and serve its public health check while checkpoint mutations return +retryable `Unavailable`. A terminating legacy-shaped Pod still holds the barrier; +only its exit permits mutation admission. This controller activation scenario is +independent of harness execution and runs in both runtime lanes. The placeholder +models legacy Pod lifetime; it does not claim an old-binary release upgrade. +Published-release upgrade legs may still skip when no compatible database-managed +baseline exists, so their successful job status alone is not upgrade evidence. diff --git a/go/core/test/e2e/checkpoint_activation_test.go b/go/core/test/e2e/checkpoint_activation_test.go new file mode 100644 index 0000000000..3694e9bd90 --- /dev/null +++ b/go/core/test/e2e/checkpoint_activation_test.go @@ -0,0 +1,90 @@ +package e2e_test + +import ( + "context" + "maps" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" + corev1 "k8s.io/api/core/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" + metav1 "k8s.io/apimachinery/pkg/apis/meta/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" + kagentenv "github.com/kagent-dev/kagent/go/core/pkg/env" +) + +// This is a controller activation test, independent of harness execution. The +// unmarked, deliberately unready Pod represents a legacy process still alive; +// it does not fabricate a checkpoint, backend phase, or old-binary upgrade. +func TestCheckpointActivationDuringControllerReplacement(t *testing.T) { + target := interactionTarget(t) + image := kagentenv.E2ERuntimeImage.Get() + require.NotEmpty(t, image, "requires the digest-pinned E2E shell fixture") + kube := interactionKubeClient(t) + ctx := metadata.AppendToOutgoingContext(t.Context(), "x-user-id", "e2e") + var 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) + old := pods.Items[0] + labels := maps.Clone(old.Labels) + delete(labels, "pod-template-hash") + legacy := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{Namespace: old.Namespace, GenerateName: "checkpoint-legacy-writer-", Labels: labels}, Spec: corev1.PodSpec{ + TerminationGracePeriodSeconds: new(int64(30)), + Containers: []corev1.Container{{Name: "writer", Image: image, Command: []string{"/bin/sh", "-c", "exec sleep 3600"}, + ReadinessProbe: &corev1.Probe{ProbeHandler: corev1.ProbeHandler{Exec: &corev1.ExecAction{Command: []string{"/bin/sh", "-c", "exit 1"}}}}, + Lifecycle: &corev1.Lifecycle{PreStop: &corev1.LifecycleHandler{Sleep: &corev1.SleepAction{Seconds: 5}}}, + }}, + }} + require.NoError(t, kube.Create(ctx, legacy)) + t.Cleanup(func() { + cleanup, cancel := context.WithTimeout(context.Background(), time.Minute) + defer cancel() + if err := kube.Delete(cleanup, legacy); !apierrors.IsNotFound(err) { + require.NoError(t, err) + } + }) + require.NoError(t, wait.PollUntilContextTimeout(ctx, 100*time.Millisecond, time.Minute, true, func(ctx context.Context) (bool, error) { + if err := kube.Get(ctx, ctrlclient.ObjectKeyFromObject(legacy), legacy); err != nil { + return false, err + } + return legacy.Status.Phase == corev1.PodRunning, nil + })) + // Given an unready legacy replica, when the actual controller is replaced, + // assert that its public health and readiness recover without activating + // checkpoint mutations. The fake replica cannot receive Service traffic. + restartControllerPod(t, kube, old) + conn := newControllerConn(t, target) + checkpoints := apiv1.NewCheckpointServiceClient(conn) + create := &apiv1.CreateCheckpointRequest{SessionId: uuid.NewString(), RequestId: uuid.NewString(), ExpectedHeadTaskId: "activation-probe"} + remove := &apiv1.DeleteCheckpointRequest{CheckpointId: uuid.NewString()} + _, err := checkpoints.CreateCheckpoint(ctx, create) + require.Equal(t, codes.Unavailable, status.Code(err)) + _, err = checkpoints.DeleteCheckpoint(ctx, remove) + require.Equal(t, codes.Unavailable, status.Code(err)) + require.NoError(t, kube.Delete(ctx, legacy)) + // A deletion timestamp alone must not activate recovery during preStop. + _, err = checkpoints.CreateCheckpoint(ctx, create) + require.Equal(t, codes.Unavailable, status.Code(err)) + require.NoError(t, wait.PollUntilContextTimeout(ctx, 100*time.Millisecond, time.Minute, true, func(ctx context.Context) (bool, error) { + _, err := checkpoints.CreateCheckpoint(ctx, create) + if status.Code(err) == codes.Unavailable { + return false, nil + } + // NotFound proves admission reached the real store after activation; + // the probe intentionally never admits synthetic checkpoint work. + if status.Code(err) == codes.NotFound { + return true, nil + } + return false, err + })) + _, err = checkpoints.DeleteCheckpoint(ctx, remove) + require.Equal(t, codes.NotFound, status.Code(err)) +} diff --git a/go/core/test/e2e/connection_test.go b/go/core/test/e2e/connection_test.go index 32384be423..c73d561c0d 100644 --- a/go/core/test/e2e/connection_test.go +++ b/go/core/test/e2e/connection_test.go @@ -136,7 +136,13 @@ func restartTestController(t *testing.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] + 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{} diff --git a/go/core/test/upgrade/controller_handoff_test.go b/go/core/test/upgrade/controller_handoff_test.go deleted file mode 100644 index 39fbffdf49..0000000000 --- a/go/core/test/upgrade/controller_handoff_test.go +++ /dev/null @@ -1,46 +0,0 @@ -package upgrade - -import ( - "testing" - - "github.com/stretchr/testify/require" - corev1 "k8s.io/api/core/v1" - metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" - "k8s.io/apimachinery/pkg/watch" -) - -func TestControllerHandoff(t *testing.T) { - old := corev1.Pod{ObjectMeta: metav1.ObjectMeta{Name: "old", UID: "old-uid"}} - replacement := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{Name: "old", UID: "new-uid"}} - t.Run("Given a terminating previous writer, when a replacement is created, assert that the overlap is rejected", func(t *testing.T) { - handoff := newControllerHandoff([]corev1.Pod{old}) - terminating := old.DeepCopy() - now := metav1.Now() - terminating.DeletionTimestamp = &now - require.NoError(t, handoff.observe(watch.Event{Type: watch.Modified, Object: terminating})) - require.ErrorContains(t, handoff.observe(watch.Event{Type: watch.Added, Object: replacement}), "previous writers still exist") - require.False(t, handoff.sawReplacement) - }) - t.Run("Given the previous writer has exited, when a replacement reuses its name, assert that the new UID completes the handoff", func(t *testing.T) { - handoff := newControllerHandoff([]corev1.Pod{old}) - require.NoError(t, handoff.observe(watch.Event{Type: watch.Deleted, Object: &old})) - require.NoError(t, handoff.observe(watch.Event{Type: watch.Added, Object: replacement})) - require.Empty(t, handoff.old) - require.True(t, handoff.sawReplacement) - }) - t.Run("Given an active handoff watch, when Kubernetes reports a watch error, assert that verification fails closed", func(t *testing.T) { - handoff := newControllerHandoff([]corev1.Pod{old}) - require.ErrorContains(t, handoff.observe(watch.Event{Type: watch.Error, Object: &metav1.Status{Reason: metav1.StatusReasonExpired}}), "watch failed") - require.False(t, handoff.sawReplacement) - }) - for _, phase := range []corev1.PodPhase{corev1.PodSucceeded, corev1.PodFailed} { - t.Run("Given a previous writer, when Kubernetes reports "+string(phase)+", assert that a replacement can start before object deletion", func(t *testing.T) { - handoff := newControllerHandoff([]corev1.Pod{old}) - exited := old.DeepCopy() - exited.Status.Phase = phase - require.NoError(t, handoff.observe(watch.Event{Type: watch.Modified, Object: exited})) - require.NoError(t, handoff.observe(watch.Event{Type: watch.Added, Object: replacement})) - require.True(t, handoff.sawReplacement) - }) - } -} diff --git a/go/core/test/upgrade/rolling_upgrade_test.go b/go/core/test/upgrade/rolling_upgrade_test.go index 14a03d3689..a0172faa3a 100644 --- a/go/core/test/upgrade/rolling_upgrade_test.go +++ b/go/core/test/upgrade/rolling_upgrade_test.go @@ -8,12 +8,6 @@ import ( "time" "github.com/stretchr/testify/require" - corev1 "k8s.io/api/core/v1" - metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" - "k8s.io/apimachinery/pkg/types" - "k8s.io/apimachinery/pkg/watch" - "k8s.io/client-go/kubernetes" - "k8s.io/client-go/tools/clientcmd" kagentenv "github.com/kagent-dev/kagent/go/core/pkg/env" ) @@ -37,16 +31,18 @@ func TestRollingUpgradeCompatibility(t *testing.T) { waitForReadyPods(t, env, postgresSelector, 3*time.Minute) waitForPostgresSchema(t, env, 3*time.Minute) - // Multiple previous-generation API replicas must all exit before checkpoint - // recovery starts. Keep a terminating writer alive long enough that an unsafe - // RollingUpdate also has a visible overlap window. + // Run even when there is no migration delta: a rolling upgrade rolls the new + // image regardless of migrations, so the deploy can still break for + // non-schema reasons (a crashing new image, readiness, old pods against + // new-code-created resources). When the target build does add migrations, the + // same flow additionally exercises the old-code/new-schema window below. + // Keep multiple old controller pods around during the rollout. With a single + // replica the old-code/new-schema window can be too small to observe reliably. kubectl(t, env, 2*time.Minute, "scale", "deployment/kagent-controller", "-n", env.namespace, "--replicas=2", ) - kubectl(t, env, time.Minute, "patch", "deployment/kagent-controller", "-n", env.namespace, - "--type=strategic", "-p", `{"spec":{"template":{"spec":{"terminationGracePeriodSeconds":30,"containers":[{"name":"controller","lifecycle":{"preStop":{"sleep":{"seconds":5}}}}]}}}}`) kubectl(t, env, 3*time.Minute, "rollout", "status", "deployment/kagent-controller", "-n", env.namespace, @@ -55,29 +51,15 @@ func TestRollingUpgradeCompatibility(t *testing.T) { oldPods := podNamesForSelector(t, env, controllerSelector) require.NotEmpty(t, oldPods, "expected old controller pods before rolling upgrade") - // Preserve old-shape SQL read/write canaries across the upgrade. The - // controller handoff itself is deliberately stop-before-start, not a claim - // that old per-request checkpoint writers support mixed-version execution. + // Seed with the baseline schema before the target controller applies new + // migrations. The compatibility canary below verifies this row is still + // readable while old pods are alive against the target schema. pgExec(t, env, fmt.Sprintf( "INSERT INTO tool (id, server_name, group_kind, description) VALUES (%s, %s, %s, 'rolling baseline tool')", pgQuote(baselineToolID), pgQuote(toolServerName), pgQuote(groupKind))) ctx, cancel := context.WithTimeout(t.Context(), 10*time.Minute) defer cancel() - config, err := clientcmd.NewNonInteractiveDeferredLoadingClientConfig( - clientcmd.NewDefaultClientConfigLoadingRules(), &clientcmd.ConfigOverrides{CurrentContext: env.kubeContext}).ClientConfig() - require.NoError(t, err) - client, err := kubernetes.NewForConfig(config) - require.NoError(t, err) - pods, err := client.CoreV1().Pods(env.namespace).List(ctx, metav1.ListOptions{LabelSelector: controllerSelector}) - require.NoError(t, err) - require.NotEmpty(t, pods.Items, "expected previous controller pod UIDs before upgrade") - observer := newControllerHandoff(pods.Items) - events, err := client.CoreV1().Pods(env.namespace).Watch(ctx, metav1.ListOptions{ - LabelSelector: controllerSelector, ResourceVersion: pods.ResourceVersion, - }) - require.NoError(t, err) - defer events.Stop() done := make(chan upgradeResult, 1) cmd := upgradeCommand(ctx, env) go func() { @@ -85,28 +67,36 @@ func TestRollingUpgradeCompatibility(t *testing.T) { done <- upgradeResult{out: string(out), err: err} }() - var result upgradeResult - finished := false - for !finished || !observer.sawReplacement { + // Wait until the target schema has landed while at least one previous-release + // controller pod is still ready. That is the rolling deploy hazard: old code + // can keep serving briefly after a new pod has applied migrations. Surface an + // upgrade failure immediately instead of letting it look like a + // schema-observation timeout. + var upgradeErr error + var upgradeOut string + require.Eventually(t, func() bool { select { - case event, ok := <-events.ResultChan(): - require.True(t, ok, "pod watch ended before the controller handoff was verified") - require.NoError(t, observer.observe(event)) - case result = <-done: - require.NoError(t, result.err, "upgrade failed:\n%s", result.out) - finished = true - case <-ctx.Done(): - t.Fatal(ctx.Err()) + case r := <-done: + if r.err != nil { + upgradeErr, upgradeOut = r.err, r.out + return true + } + // The upgrade finished successfully before we caught the window; put the + // result back so the final read below still sees it, then keep polling. + done <- r + default: } - } - require.Empty(t, observer.old, "all previous writers must have exited") - strategy := kubectl(t, env, time.Minute, "get", "deployment/kagent-controller", "-n", env.namespace, - "-o", "jsonpath={.spec.strategy.type}") - require.Equal(t, "Recreate", strings.TrimSpace(strategy)) + state, err := pgMigrationStateE(t, env) + if err != nil { + return false + } + return state.version == targetCoreVersion && anyPodsReady(t, env, oldPods) + }, 6*time.Minute, 500*time.Millisecond, "target schema was not observed while old controller pods were still ready") + require.NoError(t, upgradeErr, "upgrade failed before target schema was observed:\n%s", upgradeOut) // These SQL canaries intentionally use the pre-upgrade column shape. They do // not prove every previous-release code path works, but they catch migrations - // that break basic old read/write assumptions after the controller handoff. + // that break basic old read/write assumptions during a rolling deployment. require.Equal(t, 1, pgQueryInt(t, env, fmt.Sprintf("SELECT count(*) FROM tool WHERE id = %s AND server_name = %s AND group_kind = %s", pgQuote(baselineToolID), pgQuote(toolServerName), pgQuote(groupKind))), @@ -121,6 +111,9 @@ func TestRollingUpgradeCompatibility(t *testing.T) { "old-shape tool write did not survive against target schema", ) + result := <-done + require.NoError(t, result.err, "upgrade failed:\n%s", result.out) + kubectl(t, env, 3*time.Minute, "rollout", "status", "deployment/kagent-controller", "-n", env.namespace, @@ -166,43 +159,18 @@ func podNamesForSelectorE(t *testing.T, env upgradeEnv, selector string) ([]stri return pods, nil } -type controllerHandoff struct { - old map[types.UID]bool - sawReplacement bool -} +func anyPodsReady(t *testing.T, env upgradeEnv, pods []string) bool { + t.Helper() -func newControllerHandoff(pods []corev1.Pod) *controllerHandoff { - old := make(map[types.UID]bool, len(pods)) for _, pod := range pods { - if pod.Status.Phase != corev1.PodSucceeded && pod.Status.Phase != corev1.PodFailed { - old[pod.UID] = true - } - } - return &controllerHandoff{old: old} -} - -// Terminating and unready old pods can still issue checkpoint writes. Only -// deletion or a terminal phase (all processes exited) releases the barrier, -// matching the Deployment controller's Recreate condition. -func (h *controllerHandoff) observe(event watch.Event) error { - if event.Type == watch.Error { - return fmt.Errorf("controller pod watch failed: %v", event.Object) - } - pod, ok := event.Object.(*corev1.Pod) - if !ok { - return fmt.Errorf("unexpected controller watch object %T", event.Object) - } - if h.old[pod.UID] { - if event.Type == watch.Deleted || pod.Status.Phase == corev1.PodSucceeded || pod.Status.Phase == corev1.PodFailed { - delete(h.old, pod.UID) - } - return nil - } - if event.Type == watch.Added || event.Type == watch.Modified { - if len(h.old) != 0 { - return fmt.Errorf("new controller %s started while %d previous writers still exist", pod.Name, len(h.old)) + out, err := kubectlOutput(t, env, 10*time.Second, + "get", "pod", pod, + "-n", env.namespace, + "-o", fmt.Sprintf("jsonpath={.status.containerStatuses[?(@.name==%q)].ready}", controllerContainer), + ) + if err == nil && strings.TrimSpace(out) == "true" { + return true } - h.sawReplacement = true } - return nil + return false } diff --git a/helm/README.md b/helm/README.md index 212e9ec332..beb5d238b6 100644 --- a/helm/README.md +++ b/helm/README.md @@ -304,16 +304,18 @@ validation commands, defaults, and runtime configuration. ## Upgrading -The controller Deployment uses `Recreate`: previous controller processes exit -before the next generation starts. Plan for a controller API availability gap; -existing Substrate Actors are not restarted. This ordering prevents legacy -checkpoint writers from racing the new recovery worker. Custom deployment -tooling must preserve it without force-deleting terminating writers. - -Before rolling back to an older chart that uses `RollingUpdate`, scale the +The controller keeps `RollingUpdate`. New replicas become Ready and serve API +traffic while checkpoint recovery waits for legacy writer Pods to exit. +Checkpoint creation and deletion return retryable `Unavailable` during this +handoff, without admitting new work; other APIs remain available. Keep the chart's +private execution-protocol annotation and Downward API Pod identity when using +custom deployment templates. Do not force-delete terminating legacy writers. + +Before rolling back to an older chart without the execution protocol, quiesce +checkpoint requests and finish pending recovery. Then scale the current controller Deployment to zero and wait for its Pods to exit. Only then -restore the old chart, which does not enforce this handoff. See the -[checkpoint recovery contract](../docs/architecture/persistence-checkpoints-and-forks.md). +restore the old chart; startup activation does not protect a later downgrade. +See the [checkpoint recovery contract](../docs/architecture/persistence-checkpoints-and-forks.md). When upgrading, make sure to upgrade both charts: diff --git a/helm/kagent/templates/controller-deployment.yaml b/helm/kagent/templates/controller-deployment.yaml index 4caf31211c..2f97bbad1a 100644 --- a/helm/kagent/templates/controller-deployment.yaml +++ b/helm/kagent/templates/controller-deployment.yaml @@ -21,11 +21,8 @@ metadata: {{- include "kagent.controller.labels" . | nindent 4 }} spec: replicas: {{ .Values.controller.replicas }} - # Old checkpoint writers do not participate in recovery execution claims. - # Wait for their processes to exit before a new worker can start. strategy: - type: Recreate - rollingUpdate: null + type: RollingUpdate selector: matchLabels: {{- include "kagent.controller.selectorLabels" . | nindent 6 }} @@ -33,9 +30,10 @@ spec: metadata: annotations: checksum/configmap: {{ include (print $.Template.BasePath "/controller-configmap.yaml") . | sha256sum }} - {{- with .Values.controller.podAnnotations | default .Values.podAnnotations }} + {{- with omit (.Values.controller.podAnnotations | default .Values.podAnnotations | default dict) "kagent.dev/checkpoint-execution-claims" }} {{- toYaml . | nindent 8 }} {{- end }} + kagent.dev/checkpoint-execution-claims: v1 labels: {{- /* extra pod labels: component merged over global, selector labels applied last so they can never be overridden */}} diff --git a/helm/kagent/templates/rbac/leader-election-role.yaml b/helm/kagent/templates/rbac/leader-election-role.yaml index 1e8f7384ba..389403b4ff 100644 --- a/helm/kagent/templates/rbac/leader-election-role.yaml +++ b/helm/kagent/templates/rbac/leader-election-role.yaml @@ -6,6 +6,13 @@ metadata: labels: {{- include "kagent.controller.labels" . | nindent 4 }} rules: +# Controller startup coordinates checkpoint recovery with legacy replicas. +- apiGroups: [""] + resources: [pods] + verbs: [get, list] +- apiGroups: [apps] + resources: [replicasets] + verbs: [list] - apiGroups: - coordination.k8s.io resources: diff --git a/helm/kagent/tests/controller-configmap_test.yaml b/helm/kagent/tests/controller-configmap_test.yaml index 0fa06e7b3a..9b13d29f6b 100644 --- a/helm/kagent/tests/controller-configmap_test.yaml +++ b/helm/kagent/tests/controller-configmap_test.yaml @@ -3,15 +3,25 @@ templates: - controller-configmap.yaml - controller-deployment.yaml tests: - - it: stops previous controller writers before starting the new generation + - it: preserves rolling availability and advertises the checkpoint execution protocol template: controller-deployment.yaml asserts: - equal: path: spec.strategy.type - value: Recreate + value: RollingUpdate - equal: - path: spec.strategy.rollingUpdate - value: null + path: spec.template.metadata.annotations["kagent.dev/checkpoint-execution-claims"] + value: v1 + + - it: keeps user annotations from overriding the checkpoint execution protocol + template: controller-deployment.yaml + set: + controller.podAnnotations: + kagent.dev/checkpoint-execution-claims: incompatible + asserts: + - equal: + path: spec.template.metadata.annotations["kagent.dev/checkpoint-execution-claims"] + value: v1 - it: configures Actor JWT discovery with the Substrate trust bundle template: controller-configmap.yaml diff --git a/helm/kagent/tests/leader-election_test.yaml b/helm/kagent/tests/leader-election_test.yaml index f595c80269..cc47732248 100644 --- a/helm/kagent/tests/leader-election_test.yaml +++ b/helm/kagent/tests/leader-election_test.yaml @@ -7,6 +7,20 @@ tests: set: controller.replicas: 1 asserts: + - contains: + path: rules + content: + apiGroups: [apps] + resources: [replicasets] + verbs: [list] + template: rbac/leader-election-role.yaml + - contains: + path: rules + content: + apiGroups: [""] + resources: [pods] + verbs: [get, list] + template: rbac/leader-election-role.yaml - contains: path: rules content: diff --git a/helm/kagent/values.yaml b/helm/kagent/values.yaml index abf45cacde..c37f9d986f 100644 --- a/helm/kagent/values.yaml +++ b/helm/kagent/values.yaml @@ -224,9 +224,8 @@ config: # ============================================================================== controller: - # Controller upgrades use Recreate so legacy checkpoint writers exit before - # recovery workers start. The control-plane API is briefly unavailable; runtime - # Actors are not restarted by this controller replacement. + # Controller upgrades remain rolling. Checkpoint mutations and recovery wait + # until legacy writer Pods exit; other API traffic remains available. # Standard OTEL overrides and external SDK credentials for this process only. env: [] envFrom: [] From 27eb9ea1f254512136342ac6c06e85cfb262ac51 Mon Sep 17 00:00:00 2001 From: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> Date: Fri, 6 Nov 2026 18:35:48 +0800 Subject: [PATCH 08/20] test(checkpoint): strengthen activation isolation and cleanup Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> --- .../service/checkpoint/review_test.go | 21 +++++++----- go/core/pkg/app/checkpoint_activation_test.go | 2 +- .../test/e2e/checkpoint_activation_test.go | 33 ++++++++++++++++--- 3 files changed, 42 insertions(+), 14 deletions(-) diff --git a/go/core/internal/service/checkpoint/review_test.go b/go/core/internal/service/checkpoint/review_test.go index af9be23ef0..3d6f46f39a 100644 --- a/go/core/internal/service/checkpoint/review_test.go +++ b/go/core/internal/service/checkpoint/review_test.go @@ -69,7 +69,7 @@ func TestPermanentRecoveryErrorsKeepBackendConfigurationPrivate(t *testing.T) { {name: "backend status", failure: status.Error(codes.FailedPrecondition, private), message: "Checkpoint recovery cannot continue safely"}, {name: "service error with private cause", failure: serviceerrors.NewFailedPrecondition("Snapshot identity changed", errors.New(private)), message: "Snapshot identity changed"}, } { - t.Run("Given a "+scenario.name+", when recovery persists a permanent error, assert that callers see only the safe message", func(t *testing.T) { + t.Run("Given a "+scenario.name+", when recovery classifies a permanent failure, assert that callers see only the safe message", func(t *testing.T) { failure := permanentOperationError(scenario.failure) require.Equal(t, codes.FailedPrecondition, failure.Code()) require.Equal(t, scenario.message, failure.Message()) @@ -233,14 +233,19 @@ func TestCheckpointActivationPreservesAdmittedWork(t *testing.T) { defer stop() // Given admitted legacy work, when replacements start before writers drain, // assert that public mutations fail promptly and no worker starts a copy. - _, err = service.Create(owner, source.Id, request, head) + blocked, cancel := context.WithTimeout(owner, time.Second) + defer cancel() + _, err = service.Create(blocked, source.Id, request, head) require.Equal(t, serviceerrors.CodeUnavailable, serviceerrors.CodeOf(err)) - require.Equal(t, serviceerrors.CodeUnavailable, serviceerrors.CodeOf(service.Delete(owner, cp.Id))) - select { - case <-tags.entered: - t.Fatal("recovery executed before activation") - default: - } + require.Equal(t, serviceerrors.CodeUnavailable, serviceerrors.CodeOf(service.Delete(blocked, cp.Id))) + require.Never(t, func() bool { + select { + case <-tags.entered: + return true + default: + return false + } + }, 100*time.Millisecond, time.Millisecond, "recovery executed before activation") pending, err := store.GetCheckpointOperation(t.Context(), cp.Id, "alice") require.NoError(t, err) require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_CREATING, pending.Checkpoint.State) diff --git a/go/core/pkg/app/checkpoint_activation_test.go b/go/core/pkg/app/checkpoint_activation_test.go index e61303b3c5..030e5371bd 100644 --- a/go/core/pkg/app/checkpoint_activation_test.go +++ b/go/core/pkg/app/checkpoint_activation_test.go @@ -93,7 +93,7 @@ func TestCheckpointActivationWaitsForLegacyReplicaSetScaleDown(t *testing.T) { {name: "old replica still reported", observed: 2, remaining: 1}, {name: "scale-down completed", observed: 2, ready: true}, } { - t.Run("Given no live legacy Pods, when "+scenario.name+", assert that activation requires confirmed retirement", func(t *testing.T) { + t.Run("Given "+scenario.name+", when recovery checks activation, assert that only confirmed retirement permits execution", func(t *testing.T) { scheme := runtime.NewScheme() require.NoError(t, corev1.AddToScheme(scheme)) require.NoError(t, appsv1.AddToScheme(scheme)) diff --git a/go/core/test/e2e/checkpoint_activation_test.go b/go/core/test/e2e/checkpoint_activation_test.go index 3694e9bd90..92a1f542e8 100644 --- a/go/core/test/e2e/checkpoint_activation_test.go +++ b/go/core/test/e2e/checkpoint_activation_test.go @@ -2,13 +2,16 @@ package e2e_test import ( "context" + "fmt" "maps" "testing" "time" "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" corev1 "k8s.io/api/core/v1" @@ -45,11 +48,24 @@ func TestCheckpointActivationDuringControllerReplacement(t *testing.T) { }} require.NoError(t, kube.Create(ctx, legacy)) t.Cleanup(func() { - cleanup, cancel := context.WithTimeout(context.Background(), time.Minute) + cleanup, cancel := context.WithTimeout(metadata.AppendToOutgoingContext(context.Background(), "x-user-id", "e2e"), time.Minute) defer cancel() if err := kube.Delete(cleanup, legacy); !apierrors.IsNotFound(err) { require.NoError(t, err) } + // Keep a failed scenario from leaving the following tests behind a + // terminating placeholder or an activation check not yet completed. + require.NoError(t, wait.PollUntilContextTimeout(cleanup, 100*time.Millisecond, time.Minute, true, func(ctx context.Context) (bool, error) { + err := kube.Get(ctx, ctrlclient.ObjectKeyFromObject(legacy), &corev1.Pod{}) + return apierrors.IsNotFound(err), ctrlclient.IgnoreNotFound(err) + })) + conn, err := grpc.NewClient(target, grpc.WithTransportCredentials(insecure.NewCredentials())) + require.NoError(t, err) + defer func() { require.NoError(t, conn.Close()) }() + require.NoError(t, waitForControllerAPI(cleanup, conn)) + require.NoError(t, waitForCheckpointAdmission(cleanup, apiv1.NewCheckpointServiceClient(conn), &apiv1.CreateCheckpointRequest{ + SessionId: uuid.NewString(), RequestId: uuid.NewString(), ExpectedHeadTaskId: "activation-cleanup-probe", + })) }) require.NoError(t, wait.PollUntilContextTimeout(ctx, 100*time.Millisecond, time.Minute, true, func(ctx context.Context) (bool, error) { if err := kube.Get(ctx, ctrlclient.ObjectKeyFromObject(legacy), legacy); err != nil { @@ -73,7 +89,13 @@ func TestCheckpointActivationDuringControllerReplacement(t *testing.T) { // A deletion timestamp alone must not activate recovery during preStop. _, err = checkpoints.CreateCheckpoint(ctx, create) require.Equal(t, codes.Unavailable, status.Code(err)) - require.NoError(t, wait.PollUntilContextTimeout(ctx, 100*time.Millisecond, time.Minute, true, func(ctx context.Context) (bool, error) { + require.NoError(t, waitForCheckpointAdmission(ctx, checkpoints, create)) + _, err = checkpoints.DeleteCheckpoint(ctx, remove) + require.Equal(t, codes.NotFound, status.Code(err)) +} + +func waitForCheckpointAdmission(ctx context.Context, checkpoints apiv1.CheckpointServiceClient, create *apiv1.CreateCheckpointRequest) error { + return wait.PollUntilContextTimeout(ctx, 100*time.Millisecond, time.Minute, true, func(ctx context.Context) (bool, error) { _, err := checkpoints.CreateCheckpoint(ctx, create) if status.Code(err) == codes.Unavailable { return false, nil @@ -83,8 +105,9 @@ func TestCheckpointActivationDuringControllerReplacement(t *testing.T) { if status.Code(err) == codes.NotFound { return true, nil } + if err == nil { + return false, fmt.Errorf("unknown probe Session was admitted") + } return false, err - })) - _, err = checkpoints.DeleteCheckpoint(ctx, remove) - require.Equal(t, codes.NotFound, status.Code(err)) + }) } From e4019822aa4561c029a70f1227654d10272dec4b Mon Sep 17 00:00:00 2001 From: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> Date: Sat, 7 Nov 2026 16:06:23 +0800 Subject: [PATCH 09/20] fix(checkpoint): surface permanent errors across recovery paths Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> --- .../internal/service/checkpoint/recovery.go | 4 +- .../service/checkpoint/review_test.go | 103 ++++++++++++++++++ .../internal/service/checkpoint/service.go | 4 +- .../service/checkpoint/service_test.go | 4 +- 4 files changed, 109 insertions(+), 6 deletions(-) diff --git a/go/core/internal/service/checkpoint/recovery.go b/go/core/internal/service/checkpoint/recovery.go index 283767167d..05c3b3ab0b 100644 --- a/go/core/internal/service/checkpoint/recovery.go +++ b/go/core/internal/service/checkpoint/recovery.go @@ -3,7 +3,6 @@ package checkpoint import ( "context" "errors" - "fmt" "sync" "time" @@ -15,6 +14,7 @@ import ( 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" ) @@ -200,7 +200,7 @@ func (s *Service) recoverOperation(ctx context.Context, observed *database.Check return err } if tag != nil && !validTagSource(tag, op.Checkpoint, op.Snapshot) { - return fmt.Errorf("refuse cleanup of an unverified checkpoint tag: %w", err) + return serviceerrors.NewFailedPrecondition("Checkpoint snapshot tag identity changed", err) } uid := tag.GetMetadata().GetUid() if recordErr := s.store.BeginCheckpointCleanup(ctx, op.Checkpoint.Id, executor, uid, err.Error()); recordErr != nil { diff --git a/go/core/internal/service/checkpoint/review_test.go b/go/core/internal/service/checkpoint/review_test.go index 3d6f46f39a..9f65c45f74 100644 --- a/go/core/internal/service/checkpoint/review_test.go +++ b/go/core/internal/service/checkpoint/review_test.go @@ -59,6 +59,109 @@ func TestPublicDeleteRetainsPermanentFailure(t *testing.T) { require.ErrorIs(t, err, database.ErrNotFound) } +func TestPublicDeleteReportsPermanentBackendErrors(t *testing.T) { + for _, method := range []string{"GetTag", "DeleteTag"} { + for _, code := range []codes.Code{codes.FailedPrecondition, codes.InvalidArgument, codes.PermissionDenied, codes.Unauthenticated} { + t.Run("Given "+method+" returns "+code.String()+", when public Delete runs, assert that its safe error survives worker replacement and recovery", func(t *testing.T) { + store := reviewStore(t) + source, head := recoverySource(t, store) + tags := &lockedTags{base: &testTags{snapshotURI: "s3://snapshots/source"}} + service := NewService(store, testAuthorizer{}, tags, nil) + owner := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) + cp, err := createAndRecover(service, owner, source.Id, uuid.NewString(), head) + require.NoError(t, err) + const private = "private-backend-configuration" + tags.mu.Lock() + if method == "GetTag" { + tags.base.getErr = status.Error(code, private) + } else { + tags.base.deleteErr = status.Error(code, private) + } + tags.mu.Unlock() + stop := startRecovery(t, service) + follower := NewService(store, testAuthorizer{}, tags, nil) + ctx, cancel := context.WithTimeout(owner, 3*time.Second) + defer cancel() + err = follower.Delete(ctx, cp.Id) + require.Equal(t, code, status.Code(err)) + require.NotContains(t, err.Error(), private) + stop() + pending, err := store.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) + require.Equal(t, int32(code), pending.OperationError.GetCode()) + tags.mu.Lock() + require.NotNil(t, tags.base.created, "a failed deletion must retain the tag") + tags.mu.Unlock() + follower = NewService(store, testAuthorizer{}, tags, nil) + require.Equal(t, code, status.Code(follower.Delete(ctx, cp.Id))) + // Repair and replace the worker: the reported error must not prevent + // completion or release the reference before deletion succeeds. + tags.mu.Lock() + tags.base.getErr, tags.base.deleteErr = nil, nil + tags.mu.Unlock() + stop = startRecovery(t, follower) + require.Eventually(t, func() bool { + _, err := store.GetCheckpointOperation(t.Context(), cp.Id, "alice") + return errors.Is(err, database.ErrNotFound) + }, 3*time.Second, 10*time.Millisecond) + stop() + tags.mu.Lock() + require.Nil(t, tags.base.created) + tags.mu.Unlock() + }) + } + } +} + +func TestPublicCreateRejectsUnverifiedExistingTag(t *testing.T) { + store := reviewStore(t) + source, head := recoverySource(t, store) + tags := &lockedTags{base: &testTags{snapshotURI: "s3://snapshots/source"}} + service := NewService(store, 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), "another-actor") + require.NoError(t, err) + stop := startRecovery(t, service) + // Given a pending reservation with a tag belonging to another Actor, when + // a fresh API receives public Create, assert that it reports the identity + // error without deleting the tag or releasing source admission. + follower := NewService(store, testAuthorizer{}, tags, nil) + ctx, cancel := context.WithTimeout(owner, 3*time.Second) + defer cancel() + _, err = follower.Create(ctx, source.Id, request, head) + require.Equal(t, codes.FailedPrecondition, status.Code(err)) + stop() + pending, err := store.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, "an unverified tag must not authorize compensation") + require.Empty(t, pending.TagUID) + require.Equal(t, int32(codes.FailedPrecondition), pending.OperationError.GetCode()) + _, err = store.BeginSessionOperation(t.Context(), source.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE) + require.ErrorIs(t, err, database.ErrConflict) + tags.mu.Lock() + require.Zero(t, tags.base.deleteCalls) + require.Equal(t, "another-actor", tags.base.created.SourceActor.Name) + // Once the backend identity is repaired, a new worker may adopt the same + // checkpoint and clear the durable error without a new reservation. + tags.base.created.SourceActor.Name = substrate.ActorName(source.Id) + tags.mu.Unlock() + stop = startRecovery(t, follower) + waitCheckpointPhase(t, store, cp.Id, apiv1.CheckpointState_CHECKPOINT_STATE_READY) + stop() + published, err := follower.Create(ctx, source.Id, request, head) + require.NoError(t, err) + require.Equal(t, cp.Id, published.Id) + pending, err = store.GetCheckpointOperation(t.Context(), cp.Id, "alice") + require.NoError(t, err) + require.Nil(t, pending.OperationError) +} + func TestPermanentRecoveryErrorsKeepBackendConfigurationPrivate(t *testing.T) { const private = "private-credential-value" for _, scenario := range []struct { diff --git a/go/core/internal/service/checkpoint/service.go b/go/core/internal/service/checkpoint/service.go index 5da85fd156..8e9102e868 100644 --- a/go/core/internal/service/checkpoint/service.go +++ b/go/core/internal/service/checkpoint/service.go @@ -330,7 +330,7 @@ func (s *Service) finishDelete(ctx context.Context, op *database.CheckpointOpera 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 fmt.Errorf("get checkpoint snapshot tag: %w", err) } if err == nil { // Failed creation records no tag UID after compensation. An unexpected tag @@ -341,7 +341,7 @@ func (s *Service) finishDelete(ctx context.Context, op *database.CheckpointOpera 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 fmt.Errorf("delete checkpoint snapshot tag: %w", err) } } if err := s.store.DeleteClaimedSessionCheckpoint(ctx, checkpointID, userID, executor); err != nil { diff --git a/go/core/internal/service/checkpoint/service_test.go b/go/core/internal/service/checkpoint/service_test.go index 151ad894a9..2c514fcebe 100644 --- a/go/core/internal/service/checkpoint/service_test.go +++ b/go/core/internal/service/checkpoint/service_test.go @@ -360,7 +360,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(deleteAndRecover(service, ctx, checkpoint.GetId()))) + require.Equal(t, codes.Unavailable, status.Code(deleteAndRecover(service, 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 @@ -382,7 +382,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(deleteAndRecover(service, ctx, checkpoint.GetId()))) + require.Equal(t, codes.Unavailable, status.Code(deleteAndRecover(service, ctx, checkpoint.GetId()))) require.False(t, store.deleted) require.Zero(t, tags.deleteCalls) require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_DELETING, checkpoint.State) From 71fc58cd9a22ccf59e5cc0e57e802c7f65f18a8f Mon Sep 17 00:00:00 2001 From: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> Date: Mon, 9 Nov 2026 00:59:39 +0800 Subject: [PATCH 10/20] fix(checkpoint): hide private claim release errors Preserve the current backend status when release also fails. Log the private database cause and join a service-owned message. Add public Create and Delete regressions that reproduce the leak before the fix. Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> --- .../internal/service/checkpoint/recovery.go | 8 ++++- .../service/checkpoint/review_test.go | 35 +++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/go/core/internal/service/checkpoint/recovery.go b/go/core/internal/service/checkpoint/recovery.go index 57db855cb9..b11ac8b68e 100644 --- a/go/core/internal/service/checkpoint/recovery.go +++ b/go/core/internal/service/checkpoint/recovery.go @@ -106,7 +106,13 @@ func (s *Service) claimAndExecute(ctx context.Context, observed *database.Checkp func (s *Service) release(ctx context.Context, id string, executor uuid.UUID) error { releaseCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) defer cancel() - return s.store.ReleaseCheckpointOperation(releaseCtx, id, executor) + if err := s.store.ReleaseCheckpointOperation(releaseCtx, id, executor); 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) (result *apiv1.Checkpoint, err error) { diff --git a/go/core/internal/service/checkpoint/review_test.go b/go/core/internal/service/checkpoint/review_test.go index aaff3e4e10..a9d1c49eb1 100644 --- a/go/core/internal/service/checkpoint/review_test.go +++ b/go/core/internal/service/checkpoint/review_test.go @@ -401,3 +401,38 @@ func (s *lostCheckpointFinalization) FinalizeClaimedSessionCheckpoint(ctx contex } return nil, status.Error(codes.Unavailable, "finalization reply lost") } + +func TestPublicCheckpointErrorsHideReleaseFailures(t *testing.T) { + persistence := reviewStore(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") +} From b322b72672fcde9c2893d66e2dc35da6f88903b8 Mon Sep 17 00:00:00 2001 From: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> Date: Mon, 9 Nov 2026 01:40:25 +0800 Subject: [PATCH 11/20] fix(checkpoint): pace recovery and verify takeover boundaries Wait a full interval after recovery sweeps. Cover a real busy Substrate copy after caller cancellation and drain legacy writers before upgrade or rollback. Verify that terminating Pods cannot open the upgrade boundary. Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> --- .../persistence-checkpoints-and-forks.md | 26 ++- .../internal/service/checkpoint/recovery.go | 9 +- .../checkpoint/recovery_substrate_test.go | 162 ++++++++++++++++++ .../service/checkpoint/recovery_test.go | 55 ++++++ go/core/test/upgrade/checkpoint_drain_test.go | 116 +++++++++++++ go/core/test/upgrade/rolling_upgrade_test.go | 2 + go/core/test/upgrade/upgrade_test.go | 25 ++- 7 files changed, 384 insertions(+), 11 deletions(-) create mode 100644 go/core/test/upgrade/checkpoint_drain_test.go diff --git a/docs/architecture/persistence-checkpoints-and-forks.md b/docs/architecture/persistence-checkpoints-and-forks.md index 43a44893d6..2eb9b5f261 100644 --- a/docs/architecture/persistence-checkpoints-and-forks.md +++ b/docs/architecture/persistence-checkpoints-and-forks.md @@ -79,17 +79,22 @@ later repaired request could accidentally replay. Public errors exclude backend configuration. A canceled or uncertain attempt retains `CREATING` and expires its claim immediately while retaining the executor identity. -A small leader-elected sweep runs at startup and every thirty seconds. It reads -only expired `CREATING` and `DELETING` work and uses the same execution path as -requests. Claims fence publication and deletion across replicas, restarts, and -paused executors; a delayed release cannot revoke a replacement claim. The sweep -interval bounds background retries. Live requests execute independently of the -sweep, without a queue, retry map, or wakeup protocol. +A small leader-elected sweep runs at startup and waits thirty seconds after each +completed sweep before starting another. It reads only expired `CREATING` and +`DELETING` work and uses the same execution path as requests. Claims fence publication and deletion across replicas, restarts, and +paused executors; a delayed release cannot revoke a replacement claim. Slow +operations can delay later background work; each attempt remains bounded by two +minutes. The post-sweep interval paces background retries. Live requests execute +independently of the sweep, without a queue, retry map, or wakeup protocol. A pending tag is not evidence of abandonment. Substrate holds its Actor and Tag leases through `CreateTag`, including copying and finalization. Recovery calls `CreateTag` before interpreting `AlreadyExists`: that reply establishes that the previous serialized attempt stopped, but does not establish copy completion. +A busy backend lease returns `Aborted`; that leaves creation pending without +inspecting or cleaning the still-live tag. A later retry can adopt the completed +copy under the same request ID. + Kagent must still verify the tag's identity, template lineage, completed storage, and the source Actor's snapshot and UID. Transport errors retain admission until a subsequent serialized attempt establishes an outcome. @@ -110,7 +115,14 @@ age exceeds the two-minute operation timeout. This grace is an additional delay, not proof that an old writer exited: old code has no bounded overall operation and can wait in a reservation transaction. Do not mix old checkpoint writers with new recovery. Future rollouts between claim-aware controllers can remain rolling. -Before downgrading, finish pending checkpoint work and drain all current writers. +Before downgrading to a controller without execution claims, finish pending +checkpoint operations, delete failed +checkpoints that retain tag UIDs, and drain all current writers. Legacy deletion +requires a completed URI and cannot collect a retained partial failed tag. +The upgrade suite keeps its old-code/new-schema checks, drains every legacy Pod +before starting the new controller, and repeats that boundary before application +rollback. A separate process-lifetime test holds a terminating legacy Pod to +verify that scaling to zero or losing readiness alone does not open the boundary. Kagent requests synchronous Substrate suspension before recording a checkpoint. If an earlier attempt already suspended the Actor but left a pending durable diff --git a/go/core/internal/service/checkpoint/recovery.go b/go/core/internal/service/checkpoint/recovery.go index b11ac8b68e..382fe13e2e 100644 --- a/go/core/internal/service/checkpoint/recovery.go +++ b/go/core/internal/service/checkpoint/recovery.go @@ -32,15 +32,18 @@ func NewRecoveryWorker(service *Service) *RecoveryWorker { func (*RecoveryWorker) NeedLeaderElection() bool { return true } func (w *RecoveryWorker) Start(ctx context.Context) error { - ticker := time.NewTicker(w.interval) - defer ticker.Stop() 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(): - case <-ticker.C: + timer.Stop() + return nil + case <-timer.C: } } return nil diff --git a/go/core/internal/service/checkpoint/recovery_substrate_test.go b/go/core/internal/service/checkpoint/recovery_substrate_test.go index 8a9ef6aa3f..7b44085898 100644 --- a/go/core/internal/service/checkpoint/recovery_substrate_test.go +++ b/go/core/internal/service/checkpoint/recovery_substrate_test.go @@ -9,6 +9,7 @@ import ( "net/url" "os" "strings" + "sync/atomic" "testing" "time" @@ -17,6 +18,7 @@ import ( "github.com/google/uuid" "github.com/stretchr/testify/require" ateapipb "go.ate.dev/substrate/api/ateapi/v1" + "golang.org/x/sync/errgroup" "google.golang.org/grpc" "google.golang.org/grpc/codes" "google.golang.org/grpc/credentials" @@ -344,6 +346,124 @@ func TestCheckpointRecoveryWithSubstrate(t *testing.T) { require.Empty(t, checkpointObjectKeys(t, ctx, objects, bucket, tagPrefix), "all partially copied objects must be collected") }) + t.Run("Given a live remote copy after caller cancellation, when another replica retries Create, assert that busy work stays pending and a later retry adopts its UID", func(t *testing.T) { + objects, bucket := checkpointObjectStore(t, config.ObjectStoreEndpoint, config.ObjectStoreAccessKey, config.ObjectStoreSecretKey) + privateTemplate := proto.Clone(template).(*ateapipb.ActorTemplate) + privateTemplate.Metadata = &ateapipb.ResourceMetadata{Atespace: config.Atespace, Name: "checkpoint-concurrent-" + uuid.NewString()} + privateTemplate.Status = nil + privateTemplate.SnapshotConfig.StorageLocation = "gs://" + bucket + privateTemplate, err := client.CreateActorTemplate(ctx, privateTemplate) + require.NoError(t, err) + t.Cleanup(func() { + cleanup, stop := context.WithTimeout(context.Background(), time.Minute) + defer stop() + require.NoError(t, client.DeleteActorTemplate(cleanup, config.Atespace, privateTemplate.Metadata.Name, privateTemplate.Metadata.Uid)) + }) + store, source, head := newSourceForTemplate(t, privateTemplate) + actor, err := client.GetActor(ctx, config.Atespace, substrate.ActorName(source.Id)) + require.NoError(t, err) + snapshot, err := url.Parse(substrate.SnapshotURI(substrate.LatestActorSnapshot(actor))) + require.NoError(t, err) + require.Equal(t, bucket, snapshot.Host) + // Many real objects keep copying observable after the first destination + // object appears. A completed setup copy must never silently pass this test. + var uploads errgroup.Group + uploads.SetLimit(8) + for i := range 4096 { + uploads.Go(func() error { + _, err := objects.PutObject(ctx, &s3.PutObjectInput{Bucket: aws.String(bucket), + Key: aws.String(fmt.Sprintf("%s/concurrent-copy/%04d", strings.Trim(snapshot.Path, "/"), i)), + Body: strings.NewReader(strings.Repeat("copy-proof", 1024))}) + return err + }) + } + require.NoError(t, uploads.Wait()) + request := uuid.NewString() + service := NewService(store, testAuthorizer{}, client, nil) + cp, _, err := service.beginCreate(owner, source.Id, request, head) + require.NoError(t, err) + cleanupTag(t, cp.Id) + remoteCtx, stopRemote := context.WithTimeout(context.WithoutCancel(ctx), 2*time.Minute) + defer stopRemote() + remote := &disconnectedSubstrateCopy{Client: client, remote: remoteCtx, done: make(chan struct{})} + callerCtx, cancelCaller := context.WithCancel(owner) + defer cancelCaller() + callerDone := make(chan error, 1) + go func() { + _, err := NewService(store, testAuthorizer{}, remote, nil).Create(callerCtx, source.Id, request, head) + callerDone <- err + }() + var copying *ateapipb.Tag + require.Eventually(t, func() bool { + var err error + copying, err = client.GetTag(ctx, config.Atespace, tagName(cp.Id)) + if err != nil || substrate.SnapshotURI(copying.GetStatus().GetSnapshot()) != "" { + return false + } + prefix := "atespaces/" + config.Atespace + "/tags/" + copying.Metadata.Uid + "/" + page, err := objects.ListObjectsV2(ctx, &s3.ListObjectsV2Input{Bucket: aws.String(bucket), Prefix: aws.String(prefix), MaxKeys: aws.Int32(1)}) + return err == nil && len(page.Contents) > 0 + }, 30*time.Second, 10*time.Millisecond, "a real pending copy with destination objects was not observed") + cancelCaller() + require.ErrorIs(t, <-callerDone, context.Canceled) + observed := &observedSubstrateCopy{Client: client, entered: make(chan struct{})} + type createResult struct { + checkpoint *apiv1.Checkpoint + err error + } + result := make(chan createResult, 1) + go func() { + ready, err := NewService(newStore(), testAuthorizer{}, observed, nil).Create(owner, source.Id, request, head) + result <- createResult{checkpoint: ready, err: err} + }() + select { + case <-observed.entered: + case <-ctx.Done(): + t.Fatal(ctx.Err()) + } + select { + case <-remote.done: + t.Fatal("the remote copy finished before takeover; concurrency was not exercised") + default: + } + pending, err := store.GetCheckpointOperation(ctx, cp.Id, "alice") + require.NoError(t, err) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_CREATING, pending.Checkpoint.State) + select { + case outcome := <-result: + require.Equal(t, codes.Aborted, status.Code(outcome.err), "a busy backend lease is retryable, not proof of a failed copy") + require.Nil(t, outcome.checkpoint) + case <-ctx.Done(): + t.Fatal(ctx.Err()) + } + require.Zero(t, observed.tagReads.Load(), "a busy copy must not be inspected as an abandoned pending tag") + pending, err = store.GetCheckpointOperation(ctx, cp.Id, "alice") + require.NoError(t, err) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_CREATING, pending.Checkpoint.State) + require.Nil(t, pending.Checkpoint.Failure) + _, err = store.BeginSessionOperation(ctx, source.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE, false) + require.ErrorIs(t, err, database.ErrConflict, "live copying must retain source admission") + select { + case <-remote.done: + case <-ctx.Done(): + t.Fatal(ctx.Err()) + } + require.NoError(t, remote.err) + ready, err := NewService(newStore(), testAuthorizer{}, client, nil).Create(owner, source.Id, request, head) + require.NoError(t, err) + require.Equal(t, cp.Id, ready.Id) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_READY, ready.State) + published, err := store.GetCheckpointOperation(ctx, cp.Id, "alice") + require.NoError(t, err) + require.Equal(t, copying.Metadata.Uid, published.TagUID) + require.Equal(t, remote.tag.Metadata.Uid, published.TagUID) + tag, err := client.GetTag(ctx, config.Atespace, tagName(cp.Id)) + require.NoError(t, err) + require.NotEmpty(t, substrate.SnapshotURI(tag.GetStatus().GetSnapshot())) + t.Log("replacement received Aborted while the real copy was running; a later retry adopted the same completed tag UID") + deleteCheckpoint(t, NewService(newStore(), testAuthorizer{}, client, nil), cp.Id) + }) + t.Run("Given a lost delete response, when recovery restarts, assert that an absent real tag completes deletion", func(t *testing.T) { store, source, head := newSource(t) service := NewService(store, testAuthorizer{}, client, nil) @@ -400,3 +520,45 @@ func (f *faultedSubstrateTags) DeleteTag(ctx context.Context, space, name, uid s } var _ tagClient = (*faultedSubstrateTags)(nil) + +// The client loses cancellation delivery to the server: the caller exits while +// the actual Substrate RPC keeps copying. No backend state or reply is invented. +type disconnectedSubstrateCopy struct { + *substrate.Client + remote context.Context + done chan struct{} + tag *ateapipb.Tag + err error +} + +func (c *disconnectedSubstrateCopy) CreateTag(ctx context.Context, space, name, actor string) (*ateapipb.Tag, error) { + go func() { + c.tag, c.err = c.Client.CreateTag(c.remote, space, name, actor) + close(c.done) + }() + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-c.done: + return c.tag, c.err + } +} + +type observedSubstrateCopy struct { + *substrate.Client + entered chan struct{} + tagReads atomic.Int32 +} + +func (c *observedSubstrateCopy) CreateTag(ctx context.Context, space, name, actor string) (*ateapipb.Tag, error) { + close(c.entered) + return c.Client.CreateTag(ctx, space, name, actor) +} + +func (c *observedSubstrateCopy) GetTag(ctx context.Context, space, name string) (*ateapipb.Tag, error) { + c.tagReads.Add(1) + return c.Client.GetTag(ctx, space, name) +} + +var _ tagClient = (*disconnectedSubstrateCopy)(nil) +var _ tagClient = (*observedSubstrateCopy)(nil) diff --git a/go/core/internal/service/checkpoint/recovery_test.go b/go/core/internal/service/checkpoint/recovery_test.go index 25dd96a99e..623fc4418d 100644 --- a/go/core/internal/service/checkpoint/recovery_test.go +++ b/go/core/internal/service/checkpoint/recovery_test.go @@ -4,9 +4,12 @@ import ( "context" "crypto/sha256" "errors" + "fmt" "strings" "sync" + "sync/atomic" "testing" + "testing/synctest" "time" "github.com/a2aproject/a2a-go/v2/a2a" @@ -259,3 +262,55 @@ func (s *Service) beginDelete(ctx context.Context, id string) (string, error) { _, err = s.store.BeginDeleteSessionCheckpoint(ctx, id, owner, executor) return owner, errors.Join(err, s.release(ctx, id, executor)) } + +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/test/upgrade/checkpoint_drain_test.go b/go/core/test/upgrade/checkpoint_drain_test.go new file mode 100644 index 0000000000..d8c25ba927 --- /dev/null +++ b/go/core/test/upgrade/checkpoint_drain_test.go @@ -0,0 +1,116 @@ +package upgrade + +import ( + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + appsv1 "k8s.io/api/apps/v1" + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/util/wait" + "k8s.io/utils/ptr" + + kagentenv "github.com/kagent-dev/kagent/go/core/pkg/env" +) + +// Exercise the process-drain boundary in an isolated namespace. This models +// legacy process lifetime; TestUpgrade uses the same drain with real old binaries. +func TestUpgradeWaitsForTerminatingCheckpointWriters(t *testing.T) { + if kagentenv.E2ERunUpgradeTests.Get() != "true" { + t.Skip("requires a disposable upgrade cluster") + } + env := loadUpgradeEnv(t) + env.namespace = "checkpoint-upgrade-" + uuid.NewString() + kubectl(t, env, time.Minute, "create", "namespace", env.namespace) + t.Cleanup(func() { + _, err := kubectlOutputContext(context.WithoutCancel(t.Context()), env, time.Minute, + "delete", "namespace", env.namespace, "--wait=false") + require.NoError(t, err) + }) + labels := map[string]string{"app.kubernetes.io/name": "kagent", "app.kubernetes.io/component": "controller"} + image := kagentenv.E2ERuntimeImage.Get() + if image == "" { + image = "busybox:1.36" + } + deployment := &appsv1.Deployment{ + TypeMeta: metav1.TypeMeta{APIVersion: "apps/v1", Kind: "Deployment"}, + ObjectMeta: metav1.ObjectMeta{Name: "kagent-controller", Namespace: env.namespace}, + Spec: appsv1.DeploymentSpec{Replicas: ptr.To[int32](2), Selector: &metav1.LabelSelector{MatchLabels: labels}, + Template: corev1.PodTemplateSpec{ObjectMeta: metav1.ObjectMeta{Labels: labels}, Spec: corev1.PodSpec{ + TerminationGracePeriodSeconds: ptr.To[int64](1), + Containers: []corev1.Container{{Name: controllerContainer, Image: image, + Command: []string{"/bin/sh", "-c", "trap 'exit 0' TERM; while :; do sleep 1; done"}}}, + }}}, + } + data, err := json.Marshal(deployment) + require.NoError(t, err) + manifest := filepath.Join(t.TempDir(), "legacy-writers.json") + require.NoError(t, os.WriteFile(manifest, data, 0o600)) + kubectl(t, env, time.Minute, "apply", "-f", manifest) + kubectl(t, env, 2*time.Minute, "rollout", "status", "deployment/kagent-controller", "-n", env.namespace, "--timeout=2m") + pods := podNamesForSelector(t, env, controllerSelector) + require.Len(t, pods, 2) + held := pods[0] + // Hold the Pod object after its process becomes unready. A scale request + // or readiness transition alone must never authorize starting new writers. + kubectl(t, env, time.Minute, "patch", "pod", held, "-n", env.namespace, + "--type=merge", "-p", `{"metadata":{"finalizers":["checkpoint.kagent.dev/upgrade-test"]}}`) + released := false + t.Cleanup(func() { + if !released { + _, _ = kubectlOutputContext(context.WithoutCancel(t.Context()), env, time.Minute, "patch", "pod", held, "-n", env.namespace, + "--type=merge", "-p", `{"metadata":{"finalizers":[]}}`) + } + }) + drained := make(chan error, 1) + go func() { drained <- drainLegacyCheckpointWriters(t, env) }() + require.Eventually(t, func() bool { + out, err := kubectlOutput(t, env, 10*time.Second, "get", "pod", held, "-n", env.namespace, + "-o", `jsonpath={.metadata.deletionTimestamp}{"|"}{.status.containerStatuses[?(@.name=="controller")].ready}`) + deleting, readiness, ok := strings.Cut(out, "|") + return err == nil && ok && deleting != "" && readiness == "false" + }, time.Minute, 100*time.Millisecond) + select { + case err := <-drained: + t.Fatalf("upgrade boundary passed a terminating legacy writer: %v", err) + case <-time.After(200 * time.Millisecond): + } + kubectl(t, env, time.Minute, "patch", "pod", held, "-n", env.namespace, + "--type=merge", "-p", `{"metadata":{"finalizers":[]}}`) + released = true + select { + case err := <-drained: + require.NoError(t, err) + case <-time.After(time.Minute): + t.Fatal("upgrade boundary did not open after all legacy writers exited") + } + require.Empty(t, podNamesForSelector(t, env, controllerSelector)) + scaleController(t, env, 1) + newPods := podNamesForSelector(t, env, controllerSelector) + require.Len(t, newPods, 1) + require.NotContains(t, pods, newPods[0], "new writers must start only after old Pods are gone") +} + +func drainLegacyCheckpointWriters(t *testing.T, env upgradeEnv) error { + t.Helper() + if _, err := kubectlOutput(t, env, time.Minute, "scale", "deployment/kagent-controller", + "-n", env.namespace, "--replicas=0"); err != nil { + return fmt.Errorf("stop legacy checkpoint writers: %w", err) + } + if _, err := kubectlOutput(t, env, 2*time.Minute, "rollout", "status", "deployment/kagent-controller", + "-n", env.namespace, "--timeout=2m"); err != nil { + return fmt.Errorf("observe legacy controller scale-down: %w", err) + } + return wait.PollUntilContextTimeout(t.Context(), 500*time.Millisecond, 2*time.Minute, true, func(context.Context) (bool, error) { + pods, err := podNamesForSelectorE(t, env, controllerSelector) + return err == nil && len(pods) == 0, err + }) +} diff --git a/go/core/test/upgrade/rolling_upgrade_test.go b/go/core/test/upgrade/rolling_upgrade_test.go index a0172faa3a..6891508b48 100644 --- a/go/core/test/upgrade/rolling_upgrade_test.go +++ b/go/core/test/upgrade/rolling_upgrade_test.go @@ -31,6 +31,8 @@ func TestRollingUpgradeCompatibility(t *testing.T) { waitForReadyPods(t, env, postgresSelector, 3*time.Minute) waitForPostgresSchema(t, env, 3*time.Minute) + // Checkpoint mutations remain quiescent in this SQL compatibility lane. + // TestUpgrade separately drains legacy writers at the first checkpoint upgrade. // Run even when there is no migration delta: a rolling upgrade rolls the new // image regardless of migrations, so the deploy can still break for // non-schema reasons (a crashing new image, readiness, old pods against diff --git a/go/core/test/upgrade/upgrade_test.go b/go/core/test/upgrade/upgrade_test.go index 4cfd1c7864..54a1af336b 100644 --- a/go/core/test/upgrade/upgrade_test.go +++ b/go/core/test/upgrade/upgrade_test.go @@ -81,9 +81,13 @@ func TestUpgrade(t *testing.T) { var pgBaselineState postgresMigrationState var baselineVectorVersion int var cleanPreviousSchema string + var legacyCheckpointWriters bool if !t.Run("seed baseline data before upgrade", func(t *testing.T) { pgBaselineState = pgMigrationState(t, env) + legacyCheckpointWriters = pgQuery(t, env, `SELECT + EXISTS (SELECT 1 FROM information_schema.tables WHERE table_schema = current_schema() AND table_name = 'session_checkpoint') + AND NOT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_schema = current_schema() AND table_name = 'session_checkpoint' AND column_name = 'executor_id')`) == "t" baselineVectorVersion = pgTrackVersion(t, env, "vector_schema_migrations") cleanPreviousSchema = pgSchemaDump(t, env, kagentSchema) t.Logf("baseline Postgres schema_migrations version: %d vector=%d (target=%d)", @@ -142,6 +146,17 @@ func TestUpgrade(t *testing.T) { return } + if legacyCheckpointWriters && !t.Run("drain legacy checkpoint writers before starting the new controller", func(t *testing.T) { + // Old code can still mutate checkpoints on the expanded schema. Keep + // the earlier compatibility checks, then establish the first-upgrade + // boundary before any claim-aware controller is started. + require.NotEmpty(t, podNamesForSelector(t, env, controllerSelector)) + require.NoError(t, drainLegacyCheckpointWriters(t, env)) + require.Empty(t, podNamesForSelector(t, env, controllerSelector)) + }) { + return + } + if !t.Run("upgrade to current build", func(t *testing.T) { ctx, cancel := context.WithTimeout(t.Context(), 15*time.Minute) defer cancel() @@ -234,6 +249,11 @@ func TestUpgrade(t *testing.T) { } if !t.Run("roll back application to previous release", func(t *testing.T) { + if legacyCheckpointWriters { + require.NoError(t, drainLegacyCheckpointWriters(t, env)) + require.Zero(t, pgQueryInt(t, env, "SELECT count(*) FROM session_checkpoint WHERE state IN ('CREATING', 'DELETING') OR (state = 'FAILED' AND tag_uid <> '')"), + "finish pending operations and delete retained failed checkpoints before restoring a legacy writer") + } ctx, cancel := context.WithTimeout(t.Context(), 10*time.Minute) defer cancel() @@ -524,8 +544,11 @@ func kubectl(t *testing.T, env upgradeEnv, timeout time.Duration, args ...string func kubectlOutput(t *testing.T, env upgradeEnv, timeout time.Duration, args ...string) (string, error) { t.Helper() + return kubectlOutputContext(t.Context(), env, timeout, args...) +} - ctx, cancel := context.WithTimeout(t.Context(), timeout) +func kubectlOutputContext(ctx context.Context, env upgradeEnv, timeout time.Duration, args ...string) (string, error) { + ctx, cancel := context.WithTimeout(ctx, timeout) defer cancel() fullArgs := append([]string{"--context", env.kubeContext}, args...) From 4d4386bf06d5786f6a94cdb0a12f9523d76edc1d Mon Sep 17 00:00:00 2001 From: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> Date: Mon, 9 Nov 2026 02:10:18 +0800 Subject: [PATCH 12/20] fix(checkpoint): observe superseded completion outcomes Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> --- .../persistence-checkpoints-and-forks.md | 56 +++++-- .../service/checkpoint/completion_test.go | 148 ++++++++++++++++++ .../service/checkpoint/migration_test.go | 101 ++++++++++++ .../internal/service/checkpoint/recovery.go | 30 ++-- .../internal/service/checkpoint/service.go | 29 +++- go/core/test/upgrade/checkpoint_drain_test.go | 29 ++++ .../test/upgrade/checkpoint_rollback_test.go | 102 ++++++++++++ go/core/test/upgrade/upgrade_test.go | 3 +- 8 files changed, 469 insertions(+), 29 deletions(-) create mode 100644 go/core/internal/service/checkpoint/completion_test.go create mode 100644 go/core/internal/service/checkpoint/migration_test.go create mode 100644 go/core/test/upgrade/checkpoint_rollback_test.go diff --git a/docs/architecture/persistence-checkpoints-and-forks.md b/docs/architecture/persistence-checkpoints-and-forks.md index 2eb9b5f261..73764acd5b 100644 --- a/docs/architecture/persistence-checkpoints-and-forks.md +++ b/docs/architecture/persistence-checkpoints-and-forks.md @@ -74,6 +74,10 @@ Each attempt shares a request deadline of at most two minutes with its claim. Retries with the same owner/request ID reuse the existing reservation. They acquire expired work and execute the same function, or return the current outcome when another request owns it. A live operation returns retryable `Aborted`. +If completion loses execution authority, the caller only observes the current +outcome: READY is returned, pending work returns `Aborted`, and a completed +deletion succeeds. Create returns `NotFound` if the checkpoint was deleted. +The observation never repeats a backend mutation or overwrites the new executor. Backend errors describe the current attempt, without persisting an error that a later repaired request could accidentally replay. Public errors exclude backend configuration. A canceled or uncertain attempt retains `CREATING` and expires its @@ -87,10 +91,11 @@ operations can delay later background work; each attempt remains bounded by two minutes. The post-sweep interval paces background retries. Live requests execute independently of the sweep, without a queue, retry map, or wakeup protocol. -A pending tag is not evidence of abandonment. Substrate holds its Actor and Tag -leases through `CreateTag`, including copying and finalization. Recovery calls -`CreateTag` before interpreting `AlreadyExists`: that reply establishes that the -previous serialized attempt stopped, but does not establish copy completion. +A pending tag is not evidence of abandonment. Substrate v1.0.0 acquires its Actor +and Tag leases before copying and finalizing in `CreateTag`. Recovery calls +`CreateTag` before interpreting `AlreadyExists`, which means the name is taken, +not that its copy completed. Interpreting a retained partial copy as stopped +depends on the backend serialization contract below. A busy backend lease returns `Aborted`; that leaves creation pending without inspecting or cleaning the still-live tag. A later retry can adopt the completed copy under the same request ID. @@ -99,6 +104,25 @@ Kagent must still verify the tag's identity, template lineage, completed storage and the source Actor's snapshot and UID. Transport errors retain admission until a subsequent serialized attempt establishes an outcome. +### Required backend contract: pending Substrate confirmation + +The recovery path requires previous external writes to have stopped before a +new serialized attempt can classify a retained tag as failed or delete it. +Besides `AlreadyExists`, the service interprets `CreateTag` replies of +`FailedPrecondition` and `NotFound` as definitive enough to inspect the tag. +Those paths, and UID-conditional deletion, require the same safety guarantee. + +The live Substrate regression verifies contention while the original RPC still +holds valid leases: the caller disconnects without delivering cancellation to +that RPC, another replica receives `Aborted`, and a later retry adopts its UID. +It does not verify cancellation delivered to Substrate, failed lease renewal, +process failure, or delayed object-store writes after lease loss. Lease tokens +protect lease ownership; neither a kagent claim timeout nor a successful lease +acquisition alone proves that earlier external writes stopped. The stronger +contract needs confirmation and fault coverage in Substrate. If it does not +hold, fix the owning backend or retain admission until the outcome is safe; +do not infer safety from elapsed time. This remains a release acceptance item. + After a copy has stopped, failed verification marks the checkpoint `FAILED` and releases source admission. Its exact tag UID remains retained for explicit UID-conditional deletion, including when tag storage is still `PENDING` and has @@ -106,15 +130,19 @@ no completed URI. Failed checkpoints cannot be forked. Their history and revisio references remain pinned until deletion succeeds. An unrelated or replaced tag is never adopted or deleted. -The first upgrade from a controller without checkpoint execution claims requires -an explicit boundary: quiesce checkpoint creation/deletion, stop every old -controller, and wait for those processes to exit before starting the new version -and resuming checkpoint requests. Merely marking a Pod unready is insufficient. +### Proposed initial upgrade boundary: pending maintainer acceptance + +The proposed first upgrade from a controller without checkpoint execution claims +quiesces checkpoint creation/deletion, stops every old controller, and waits for +those processes to exit before starting the new version and resuming requests. +This interrupts the controller API. Merely marking a Pod unready is insufficient. Legacy reservations have NULL claims; recovery waits until their UUIDv7 creation age exceeds the two-minute operation timeout. This grace is an additional delay, not proof that an old writer exited: old code has no bounded overall operation and can wait in a reservation transaction. Do not mix old checkpoint writers with -new recovery. Future rollouts between claim-aware controllers can remain rolling. +new recovery. Claim-aware controllers can coordinate through the database during +later rolling upgrades; this does not establish compatibility of every other +release change. Before downgrading to a controller without execution claims, finish pending checkpoint operations, delete failed checkpoints that retain tag UIDs, and drain all current writers. Legacy deletion @@ -123,6 +151,16 @@ The upgrade suite keeps its old-code/new-schema checks, drains every legacy Pod before starting the new controller, and repeats that boundary before application rollback. A separate process-lifetime test holds a terminating legacy Pod to verify that scaling to zero or losing readiness alone does not open the boundary. +The PostgreSQL migration regression seeds a populated version 3 checkpoint before +claim columns exist, upgrades it, and verifies recovery preserves its request and +history boundary. The rollback guard is also exercised with a retained failed +partial-tag fixture. These are storage fixtures, not previous-release binaries. + +The production installer still performs normal Helm upgrades; only the upgrade +test explicitly drains checkpoint writers. The proposed maintenance window is +not an accepted release policy or an enforced installer sequence. Before release, +maintainers must choose that boundary with public upgrade/rollback instructions, +or a staged compatible rollout if API interruption is unacceptable. Kagent requests synchronous Substrate suspension before recording a checkpoint. If an earlier attempt already suspended the Actor but left a pending durable 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..84583fc412 --- /dev/null +++ b/go/core/internal/service/checkpoint/completion_test.go @@ -0,0 +1,148 @@ +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 := reviewStore(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)) + } + } + service := NewService(gate, 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": + require.True(t, serviceerrors.IsCode(err, serviceerrors.CodeUnavailable), "%v", 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)) + } + }) + } +} + +func TestDeleteObservesReplacementCompletion(t *testing.T) { + persistence := reviewStore(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)) + } + } + err = NewService(gate, 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") + _, 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 + beforeCreateFinish func(context.Context, *apiv1.Checkpoint, uuid.UUID, string, string) + beforeDeleteFinish func(context.Context, string, string, uuid.UUID) + conflicts int +} + +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 +} 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..db5eb89e1d --- /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. +// Backend serialization is tested separately against real Substrate. +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/recovery.go b/go/core/internal/service/checkpoint/recovery.go index 382fe13e2e..65d06aeecd 100644 --- a/go/core/internal/service/checkpoint/recovery.go +++ b/go/core/internal/service/checkpoint/recovery.go @@ -90,14 +90,7 @@ func (s *Service) claimAndExecute(ctx context.Context, observed *database.Checkp }() op, err := s.store.ClaimCheckpointOperation(ctx, id, owner, executor) if errors.Is(err, database.ErrConflict) { - current, readErr := s.store.GetCheckpointOperation(ctx, id, owner) - if errors.Is(readErr, database.ErrNotFound) { - return nil, nil - } - if readErr != nil { - return nil, readErr - } - return createOutcome(current.Checkpoint) + return s.currentOutcome(ctx, id, owner) } if err != nil { return nil, err @@ -106,6 +99,19 @@ func (s *Service) claimAndExecute(ctx context.Context, observed *database.Checkp return s.execute(ctx, op, executor) } +// 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 { releaseCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) defer cancel() @@ -144,12 +150,16 @@ func (s *Service) execute(ctx context.Context, op *database.CheckpointOperation, if err != nil { failure = "Checkpoint snapshot copy could not be verified" } - // CreateTag's serialized reply established that copying stopped. FAILED can - // release admission now and retain the UID for a later explicit Delete. + // Under the required backend serialization contract, the definitive reply + // permits FAILED to release admission and retain its UID for explicit Delete. + // Expiry of the kagent claim alone never establishes that copying stopped. 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) { + return s.currentOutcome(finishCtx, op.Checkpoint.Id, op.Creator) + } if err != nil { return nil, err } diff --git a/go/core/internal/service/checkpoint/service.go b/go/core/internal/service/checkpoint/service.go index e697d9f7b6..514c0c86ba 100644 --- a/go/core/internal/service/checkpoint/service.go +++ b/go/core/internal/service/checkpoint/service.go @@ -96,12 +96,13 @@ func (s *Service) Create(ctx context.Context, sessionID, requestID, expectedHead } if op.Checkpoint.Id != id.String() { result, err = s.claimAndExecute(ctx, op) - if err == nil && result == nil { - return nil, serviceerrors.NewNotFound("Checkpoint no longer exists", nil) - } - return result, err + } else { + result, err = s.execute(ctx, op, executor) + } + if err == nil && result == nil { + return nil, serviceerrors.NewNotFound("Checkpoint no longer exists", nil) } - return s.execute(ctx, op, executor) + return result, err } func (s *Service) reserveCreate(ctx context.Context, sessionID, requestID, expectedHeadTaskID, id string, executor uuid.UUID) (*database.CheckpointOperation, error) { @@ -163,9 +164,11 @@ func (s *Service) ensureTag(ctx context.Context, checkpoint *apiv1.Checkpoint, r return nil, beforeErr } name := tagName(checkpoint.Id) - // CreateTag holds Substrate's actor and tag leases. AlreadyExists is a - // synchronization barrier: a pending tag read after it is no longer being - // copied by an earlier call. A plain GetTag cannot establish that fact. + // v1.0.0 takes Actor and Tag leases before its workflow checks. Interpreting + // terminal replies requires the backend guarantee that earlier external + // writes cannot outlive serialization, including cancellation/lease loss; + // see the architecture document's required backend contract. GetTag alone + // cannot establish that a pending copy stopped. tag, err := s.tags.CreateTag(ctx, reference.Atespace, name, actorName) if err != nil { switch status.Code(err) { @@ -328,6 +331,16 @@ func (s *Service) finishDelete(ctx context.Context, op *database.CheckpointOpera } } 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 diff --git a/go/core/test/upgrade/checkpoint_drain_test.go b/go/core/test/upgrade/checkpoint_drain_test.go index d8c25ba927..e24fdd0e09 100644 --- a/go/core/test/upgrade/checkpoint_drain_test.go +++ b/go/core/test/upgrade/checkpoint_drain_test.go @@ -11,6 +11,7 @@ import ( "time" "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgxpool" "github.com/stretchr/testify/require" appsv1 "k8s.io/api/apps/v1" corev1 "k8s.io/api/core/v1" @@ -19,6 +20,7 @@ import ( "k8s.io/utils/ptr" kagentenv "github.com/kagent-dev/kagent/go/core/pkg/env" + "github.com/kagent-dev/kagent/go/core/pkg/postgres" ) // Exercise the process-drain boundary in an isolated namespace. This models @@ -114,3 +116,30 @@ func drainLegacyCheckpointWriters(t *testing.T, env upgradeEnv) error { return err == nil && len(pods) == 0, err }) } + +func requireLegacyCheckpointRollbackReady(t *testing.T, env upgradeEnv) { + t.Helper() + port, stop := startPortForward(t, env, postgresServiceName, 5432) + defer stop() + config, err := postgres.PoolConfig(postgresURL(t, env, port), postgresRole) + require.NoError(t, err) + config.ConnConfig.RuntimeParams["search_path"] = kagentSchema + pool, err := pgxpool.NewWithConfig(t.Context(), config) + require.NoError(t, err) + defer pool.Close() + require.NoError(t, checkLegacyCheckpointRollback(t.Context(), pool)) +} + +// Legacy deletion requires a completed URI and cannot collect retained partial +// tags. Keep this guard shared with the PostgreSQL rollback regression. +func checkLegacyCheckpointRollback(ctx context.Context, pool *pgxpool.Pool) error { + var blockers int + if err := pool.QueryRow(ctx, `SELECT count(*) FROM session_checkpoint + WHERE state IN ('CREATING', 'DELETING') OR (state = 'FAILED' AND tag_uid <> '')`).Scan(&blockers); err != nil { + return fmt.Errorf("check checkpoint rollback compatibility: %w", err) + } + if blockers != 0 { + return fmt.Errorf("finish pending operations and delete retained failed checkpoints before restoring a legacy writer: %d incompatible checkpoints", blockers) + } + return nil +} diff --git a/go/core/test/upgrade/checkpoint_rollback_test.go b/go/core/test/upgrade/checkpoint_rollback_test.go new file mode 100644 index 0000000000..b601ebcd5a --- /dev/null +++ b/go/core/test/upgrade/checkpoint_rollback_test.go @@ -0,0 +1,102 @@ +package upgrade + +import ( + "context" + "crypto/sha256" + "testing" + + "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" + + 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" +) + +// Exercise the actual rollback guard with valid persisted checkpoints. The +// partial tag is a storage fixture; real object cleanup has its own Substrate test. +func TestUpgradeRejectsIncompatibleCheckpointRollback(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) + store := database.NewClient(pool, "public") + source, head := checkpointRollbackSource(t, store) + ctx, cancel := context.WithTimeout(t.Context(), database.RuntimeOperationTimeout) + defer cancel() + require.NoError(t, checkLegacyCheckpointRollback(ctx, pool)) + executor := uuid.New() + checkpoint, _, err := store.ReserveSessionCheckpoint(ctx, + &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: source.Id, HeadTaskId: head}, "alice", uuid.NewString(), executor) + require.NoError(t, err) + require.Error(t, checkLegacyCheckpointRollback(ctx, pool), "pending creation must block rollback") + failed, err := store.FinalizeClaimedSessionCheckpoint(ctx, checkpoint, executor, "partial-tag-uid", "", "copy failed") + require.NoError(t, err) + retained, err := store.GetCheckpointOperation(ctx, checkpoint.Id, "alice") + require.NoError(t, err) + require.Equal(t, "partial-tag-uid", retained.TagUID) + require.Equal(t, "s3://snapshots/source", retained.Snapshot.URI, "the failed copy has no completed tag URI") + require.Error(t, checkLegacyCheckpointRollback(ctx, pool), "a retained FAILED tag must block rollback") + deletion := uuid.New() + _, err = store.BeginDeleteSessionCheckpoint(ctx, failed.Id, "alice", deletion) + require.NoError(t, err) + require.Error(t, checkLegacyCheckpointRollback(ctx, pool), "pending cleanup must block rollback") + require.NoError(t, store.DeleteClaimedSessionCheckpoint(ctx, failed.Id, "alice", deletion)) + require.NoError(t, checkLegacyCheckpointRollback(ctx, pool), "completed cleanup permits rollback") + + for _, outcome := range []string{"ready", "failed without a tag"} { + t.Run("Given a checkpoint "+outcome+", when rollback is checked, assert that no incompatible work remains", func(t *testing.T) { + executor := uuid.New() + checkpoint, _, err := store.ReserveSessionCheckpoint(ctx, + &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: source.Id, HeadTaskId: head}, "alice", uuid.NewString(), executor) + require.NoError(t, err) + uid, uri, failure := "complete-tag-uid", "s3://tags/complete", "" + if outcome != "ready" { + uid, uri, failure = "", "", "source was deleted before copying" + } + _, err = store.FinalizeClaimedSessionCheckpoint(ctx, checkpoint, executor, uid, uri, failure) + require.NoError(t, err) + require.NoError(t, checkLegacyCheckpointRollback(ctx, pool)) + }) + } +} + +func checkpointRollbackSource(t *testing.T, store *database.Client) (*apiv1.Session, string) { + t.Helper() + ctx := t.Context() + require.NoError(t, store.UpsertAgentDefinition(ctx, database.AgentDefinition{ + Namespace: "team-a", AgentName: "assistant", AgentUID: "agent-uid", DesiredRevision: "revision"})) + _, err := store.RecordRuntimeRevision(ctx, database.RuntimeRevision{ + Revision: "revision", Namespace: "team-a", AgentName: "assistant", AgentUID: "agent-uid", + SourceSnapshot: []byte("{}"), AgentCard: &a2apb.AgentCard{}, EgressDestinations: []string{}, + ActorTemplateAtespace: "team-a", ActorTemplateName: "runtime", ActorTemplateUID: "template-uid", + }, database.RecordRevisionOptions{Ready: true}) + require.NoError(t, err) + session, _, err := store.CreateSession(ctx, &apiv1.Session{Id: uuid.NewString(), Creator: "alice", + Agent: &apiv1.AgentProvenance{Ref: &apiv1.ResourceReference{Namespace: "team-a", 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, "runtime", "actor-uid", "") + 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: "team-a", URI: "s3://snapshots/source", Fidelity: "SNAPSHOT_FIDELITY_VOLUMES"})) + return session, string(task.ID) +} diff --git a/go/core/test/upgrade/upgrade_test.go b/go/core/test/upgrade/upgrade_test.go index 54a1af336b..13d99ec1aa 100644 --- a/go/core/test/upgrade/upgrade_test.go +++ b/go/core/test/upgrade/upgrade_test.go @@ -251,8 +251,7 @@ func TestUpgrade(t *testing.T) { if !t.Run("roll back application to previous release", func(t *testing.T) { if legacyCheckpointWriters { require.NoError(t, drainLegacyCheckpointWriters(t, env)) - require.Zero(t, pgQueryInt(t, env, "SELECT count(*) FROM session_checkpoint WHERE state IN ('CREATING', 'DELETING') OR (state = 'FAILED' AND tag_uid <> '')"), - "finish pending operations and delete retained failed checkpoints before restoring a legacy writer") + requireLegacyCheckpointRollbackReady(t, env) } ctx, cancel := context.WithTimeout(t.Context(), 10*time.Minute) defer cancel() From b6aeda441808bc0eab9d50fbf72ce3a35da4b332 Mon Sep 17 00:00:00 2001 From: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> Date: Mon, 9 Nov 2026 02:24:05 +0800 Subject: [PATCH 13/20] fix(checkpoint): skip obsolete claim releases Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> --- .../persistence-checkpoints-and-forks.md | 4 ++ .../service/checkpoint/completion_test.go | 58 ++++++++++++++++++- .../internal/service/checkpoint/recovery.go | 26 ++++----- 3 files changed, 73 insertions(+), 15 deletions(-) diff --git a/docs/architecture/persistence-checkpoints-and-forks.md b/docs/architecture/persistence-checkpoints-and-forks.md index 73764acd5b..61f1a4eb15 100644 --- a/docs/architecture/persistence-checkpoints-and-forks.md +++ b/docs/architecture/persistence-checkpoints-and-forks.md @@ -78,6 +78,10 @@ If completion loses execution authority, the caller only observes the current outcome: READY is returned, pending work returns `Aborted`, and a completed deletion succeeds. Create returns `NotFound` if the checkpoint was deleted. The observation never repeats a backend mutation or overwrites the new executor. +Create observers skip obsolete claim releases, so a release failure cannot mask +READY or the translation of absence to `NotFound`. A definitive acquisition +conflict also needs no release. Other acquisition errors may hide a committed +claim; those attempts still release their specific executor before returning. Backend errors describe the current attempt, without persisting an error that a later repaired request could accidentally replay. Public errors exclude backend configuration. A canceled or uncertain attempt retains `CREATING` and expires its diff --git a/go/core/internal/service/checkpoint/completion_test.go b/go/core/internal/service/checkpoint/completion_test.go index 84583fc412..026b0a0d50 100644 --- a/go/core/internal/service/checkpoint/completion_test.go +++ b/go/core/internal/service/checkpoint/completion_test.go @@ -47,7 +47,8 @@ func TestCreateObservesReplacementCompletion(t *testing.T) { require.NoError(t, persistence.DeleteClaimedSessionCheckpoint(ctx, checkpoint.Id, "alice", deletion)) } } - service := NewService(gate, testAuthorizer{}, tags, nil) + 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") @@ -79,10 +80,37 @@ func TestCreateObservesReplacementCompletion(t *testing.T) { 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 := reviewStore(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 := reviewStore(t) for _, outcome := range []string{"pending", "deleted"} { @@ -104,11 +132,13 @@ func TestDeleteObservesReplacementCompletion(t *testing.T) { require.NoError(t, persistence.DeleteClaimedSessionCheckpoint(ctx, id, owner, replacement)) } } - err = NewService(gate, testAuthorizer{}, tags, nil).Delete(owner, checkpoint.Id) + 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 { @@ -124,11 +154,23 @@ func TestDeleteObservesReplacementCompletion(t *testing.T) { // 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) 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) @@ -146,3 +188,15 @@ func (s *checkpointCompletionGate) DeleteClaimedSessionCheckpoint(ctx context.Co } 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") +} diff --git a/go/core/internal/service/checkpoint/recovery.go b/go/core/internal/service/checkpoint/recovery.go index 65d06aeecd..d5fea229a5 100644 --- a/go/core/internal/service/checkpoint/recovery.go +++ b/go/core/internal/service/checkpoint/recovery.go @@ -79,23 +79,20 @@ func (s *Service) recoverOperation(ctx context.Context, observed *database.Check // 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) (result *apiv1.Checkpoint, err error) { +func (s *Service) claimAndExecute(ctx context.Context, observed *database.CheckpointOperation) (*apiv1.Checkpoint, error) { executor := uuid.New() id, owner := observed.Checkpoint.Id, observed.Creator - acquired := false - defer func() { - if !acquired { - err = errors.Join(err, s.release(ctx, id, executor)) - } - }() 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 { - return nil, err + // A lost COMMIT reply can hide successful acquisition. Retire this + // executor before returning the uncertain error, as on reservation. + return nil, errors.Join(err, s.release(ctx, id, executor)) } - acquired = true return s.execute(ctx, op, executor) } @@ -125,9 +122,9 @@ func (s *Service) release(ctx context.Context, id string, executor uuid.UUID) er } func (s *Service) execute(ctx context.Context, op *database.CheckpointOperation, executor uuid.UUID) (result *apiv1.Checkpoint, err error) { - completed := false + releaseNeeded := true defer func() { - if !completed { + if releaseNeeded { err = errors.Join(err, s.release(ctx, op.Checkpoint.Id, executor)) } }() @@ -136,7 +133,7 @@ func (s *Service) execute(ctx context.Context, op *database.CheckpointOperation, } if op.Checkpoint.State == apiv1.CheckpointState_CHECKPOINT_STATE_DELETING { err = s.finishDelete(ctx, op, executor) - completed = err == nil + releaseNeeded = err != nil return nil, err } tag, err := s.ensureTag(ctx, op.Checkpoint, op.Snapshot) @@ -158,11 +155,14 @@ func (s *Service) execute(ctx context.Context, op *database.CheckpointOperation, 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, err } - completed = true + releaseNeeded = false return createOutcome(checkpoint) } From 040f27b67925130d4b879877b082475c405eb5ee Mon Sep 17 00:00:00 2001 From: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> Date: Mon, 9 Nov 2026 14:15:07 +0800 Subject: [PATCH 14/20] fix(checkpoint): expose failed records for owner cleanup Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> --- .../persistence-checkpoints-and-forks.md | 10 ++ go/core/internal/database/checkpoints.go | 20 ++- .../internal/database/client_session_test.go | 10 +- .../grpcserver/checkpoint_failure_test.go | 119 ++++++++++++++++ .../service/checkpoint/completion_test.go | 2 +- .../service/checkpoint/discovery_test.go | 131 ++++++++++++++++++ .../internal/service/checkpoint/recovery.go | 2 +- .../service/checkpoint/review_test.go | 2 +- .../internal/service/checkpoint/service.go | 32 ++++- 9 files changed, 313 insertions(+), 15 deletions(-) create mode 100644 go/core/internal/grpcserver/checkpoint_failure_test.go create mode 100644 go/core/internal/service/checkpoint/discovery_test.go diff --git a/docs/architecture/persistence-checkpoints-and-forks.md b/docs/architecture/persistence-checkpoints-and-forks.md index 61f1a4eb15..f4d485cda4 100644 --- a/docs/architecture/persistence-checkpoints-and-forks.md +++ b/docs/architecture/persistence-checkpoints-and-forks.md @@ -134,6 +134,16 @@ no completed URI. Failed checkpoints cannot be forked. Their history and revisio references remain pinned until deletion succeeds. An unrelated or replaced tag is never adopted or deleted. +Owner-scoped Get and List expose both `READY` and `FAILED` checkpoints, including +after source deletion. `CREATING` and `DELETING` remain hidden, and fork reads and +transactions still require `READY`. A terminal failed Create returns +`FailedPrecondition` with `ErrorInfo.reason = KAGENT_CHECKPOINT_FAILED` and the +checkpoint UUID in `ResourceInfo.resource_name`. Repeating that request returns +the same terminal identity without copying again. The owner can discover and +delete the failed checkpoint before creating a new request. Public failed +checkpoint responses use a generic service-owned failure message, including for +legacy rows that retained raw backend errors; the stored operation is not rewritten. + ### Proposed initial upgrade boundary: pending maintainer acceptance The proposed first upgrade from a controller without checkpoint execution claims diff --git a/go/core/internal/database/checkpoints.go b/go/core/internal/database/checkpoints.go index a5e76df81f..3176de5a51 100644 --- a/go/core/internal/database/checkpoints.go +++ b/go/core/internal/database/checkpoints.go @@ -372,10 +372,16 @@ func (c *Client) FinalizeClaimedSessionCheckpoint(ctx context.Context, checkpoin return result, nil } -// GetSessionCheckpoint returns an owned READY checkpoint. Missing checkpoints, other -// owners, and other lifecycle states return ErrNotFound. +// GetSessionCheckpoint returns an owned READY or FAILED checkpoint so the owner +// can discover a failed copy for cleanup. Missing checkpoints, other owners, +// CREATING, and DELETING return ErrNotFound. Fork readers remain READY-only. 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 IN ('READY', 'FAILED') + `, pgx.RowToStructByName[sessionCheckpointRow], id, userID) if err != nil { return nil, fmt.Errorf("get Session checkpoint: %w", notFoundOr(err)) } @@ -440,9 +446,9 @@ 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. +// ListSessionCheckpoints returns an owner's READY and FAILED checkpoints for the +// source session in ascending ID order after afterID, up to limit. CREATING and +// DELETING stay hidden. Retained checkpoints remain listable after 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, @@ -450,7 +456,7 @@ func (c *Client) ListSessionCheckpoints(ctx context.Context, sessionID, userID, FROM session_checkpoint WHERE source_session_id = $1 AND user_id = $2 - AND state = 'READY' + AND state IN ('READY', 'FAILED') AND (NULLIF($3::text, '') IS NULL OR id > NULLIF($3::text, '')::uuid) ORDER BY id LIMIT $4 diff --git a/go/core/internal/database/client_session_test.go b/go/core/internal/database/client_session_test.go index bf80f4d3f4..1e2d769a75 100644 --- a/go/core/internal/database/client_session_test.go +++ b/go/core/internal/database/client_session_test.go @@ -343,9 +343,17 @@ func TestSessionCheckpointRetainsRecordedBoundary(t *testing.T) { t.Fatalf("checkpoint replay after source deletion = %+v, error %v", replayed, err) } listed, err := client.ListSessionCheckpoints(ctx, sessionID, "alice", "", 10) - if err != nil || len(listed) != 1 || listed[0].Id != checkpoint.Id { + if err != nil || len(listed) != 2 || listed[0].Id != checkpoint.Id || listed[1].Id != failed.Id { t.Fatalf("listed checkpoints = %+v, error %v", listed, err) } + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_FAILED, listed[1].State) + page, err := client.ListSessionCheckpoints(ctx, sessionID, "alice", checkpoint.Id, 1) + require.NoError(t, err) + require.Len(t, page, 1) + require.Equal(t, failed.Id, page[0].Id, "pagination must include failed cleanup candidates") + 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) } diff --git a/go/core/internal/grpcserver/checkpoint_failure_test.go b/go/core/internal/grpcserver/checkpoint_failure_test.go new file mode 100644 index 0000000000..f9d7a1fa5a --- /dev/null +++ b/go/core/internal/grpcserver/checkpoint_failure_test.go @@ -0,0 +1,119 @@ +package grpcserver + +import ( + "context" + "crypto/sha256" + "testing" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/google/uuid" + "github.com/stretchr/testify/require" + ateapipb "go.ate.dev/substrate/api/ateapi/v1" + "google.golang.org/genproto/googleapis/rpc/errdetails" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/metadata" + "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/checkpoint" + "github.com/kagent-dev/kagent/go/core/internal/substrate" + "github.com/kagent-dev/kagent/go/core/pkg/auth" +) + +func TestFailedCheckpointIdentitySurvivesGRPCMapping(t *testing.T) { + dsn := dbtest.StartT(context.WithoutCancel(t.Context()), t) + dbtest.MigrateT(t, dsn, false) + pool, err := database.Connect(t.Context(), &database.PostgresConfig{Pool: dbtest.PoolConfig(dsn)}) + require.NoError(t, err) + t.Cleanup(pool.Close) + store := database.NewClient(pool, "public") + session := createTaskStoreSession(t, store) + message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("checkpoint boundary")) + message.ContextID = session.ContextId + task := a2a.NewSubmittedTask(message, message) + task.Status.State = a2a.TaskStateCompleted + digest := sha256.Sum256([]byte(task.ID)) + version, err := store.CreateRuntimeTask(t.Context(), session.Id, digest[:], task, "") + require.NoError(t, err) + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + boundary, err := store.ClaimSessionQuiescence(t.Context()) + 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"})) + ctx, cancel := context.WithTimeout(t.Context(), database.RuntimeOperationTimeout) + defer cancel() + executor, requestID := uuid.New(), uuid.NewString() + failed, _, err := store.ReserveSessionCheckpoint(ctx, &apiv1.Checkpoint{ + Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID)}, "alice", requestID, executor) + require.NoError(t, err) + failed, err = store.FinalizeClaimedSessionCheckpoint(ctx, failed, executor, "partial-tag-uid", "", "private legacy backend detail") + require.NoError(t, err) + tags := &checkpointFailureTags{tag: &ateapipb.Tag{Metadata: &ateapipb.ResourceMetadata{ + Atespace: "team-a", Name: "checkpoint-" + failed.Id, Uid: "partial-tag-uid"}}} + connection := newConfigurationConnection(t, nil, func(config *Config) { + config.CheckpointService = checkpoint.NewService(store, &auth.NoopAuthorizer{}, tags, nil) + }) + client := apiv1.NewCheckpointServiceClient(connection) + ctx = metadata.AppendToOutgoingContext(ctx, "x-user-id", "alice") + _, err = client.CreateCheckpoint(ctx, &apiv1.CreateCheckpointRequest{ + SessionId: failed.SessionId, RequestId: requestID, ExpectedHeadTaskId: failed.HeadTaskId}) + require.Equal(t, codes.FailedPrecondition, status.Code(err)) + var reason, id string + for _, detail := range status.Convert(err).Details() { + switch info := detail.(type) { + case *errdetails.ErrorInfo: + require.Equal(t, "kagent.dev", info.Domain) + reason = info.Reason + case *errdetails.ResourceInfo: + id = info.ResourceName + } + } + require.Equal(t, "KAGENT_CHECKPOINT_FAILED", reason) + require.Equal(t, failed.Id, id) + require.NotContains(t, err.Error(), "private") + read, err := client.GetCheckpoint(ctx, &apiv1.GetCheckpointRequest{CheckpointId: id}) + require.NoError(t, err) + require.Equal(t, failed.State, read.Checkpoint.State) + require.NotContains(t, read.String(), "private") + listed, err := client.ListCheckpoints(ctx, &apiv1.ListCheckpointsRequest{SessionId: failed.SessionId}) + require.NoError(t, err) + require.Len(t, listed.Checkpoints, 1) + require.Equal(t, id, listed.Checkpoints[0].Id) + require.NotContains(t, listed.String(), "private") + other := metadata.AppendToOutgoingContext(t.Context(), "x-user-id", "bob") + _, err = client.GetCheckpoint(other, &apiv1.GetCheckpointRequest{CheckpointId: id}) + require.Equal(t, codes.NotFound, status.Code(err)) + _, err = client.DeleteCheckpoint(other, &apiv1.DeleteCheckpointRequest{CheckpointId: id}) + require.Equal(t, codes.NotFound, status.Code(err)) + _, err = client.ForkSession(ctx, &apiv1.ForkSessionRequest{CheckpointId: id, RequestId: uuid.NewString()}) + require.Equal(t, codes.NotFound, status.Code(err)) + _, err = client.DeleteCheckpoint(ctx, &apiv1.DeleteCheckpointRequest{CheckpointId: id}) + require.NoError(t, err) + require.Nil(t, tags.tag) + _, err = store.GetCheckpointOperation(t.Context(), id, "alice") + require.ErrorIs(t, err, database.ErrNotFound) + listed, err = client.ListCheckpoints(ctx, &apiv1.ListCheckpointsRequest{SessionId: failed.SessionId}) + require.NoError(t, err) + require.Empty(t, listed.Checkpoints) +} + +// Only the external Tag API is replaced: the generated client, interceptors, +// service and PostgreSQL use their production implementations. +type checkpointFailureTags struct { + *substrate.Client + tag *ateapipb.Tag +} + +func (s *checkpointFailureTags) GetTag(context.Context, string, string) (*ateapipb.Tag, error) { + return s.tag, nil +} + +func (s *checkpointFailureTags) DeleteTag(_ context.Context, atespace, name, uid string) error { + if s.tag.GetMetadata().GetAtespace() != atespace || s.tag.GetMetadata().GetName() != name || s.tag.GetMetadata().GetUid() != uid { + return status.Error(codes.FailedPrecondition, "tag identity changed") + } + s.tag = nil + return nil +} diff --git a/go/core/internal/service/checkpoint/completion_test.go b/go/core/internal/service/checkpoint/completion_test.go index 026b0a0d50..16e9046c9c 100644 --- a/go/core/internal/service/checkpoint/completion_test.go +++ b/go/core/internal/service/checkpoint/completion_test.go @@ -69,7 +69,7 @@ func TestCreateObservesReplacementCompletion(t *testing.T) { _, err = NewService(persistence, testAuthorizer{}, tags, nil).Create(owner, source.Id, request, head) require.NoError(t, err) case "failed": - require.True(t, serviceerrors.IsCode(err, serviceerrors.CodeUnavailable), "%v", err) + assertFailedCheckpointIdentity(t, err, checkpointID) op, readErr := persistence.GetCheckpointOperation(owner, checkpointID, "alice") require.NoError(t, readErr) require.Equal(t, "replacement rejected the copy", op.Checkpoint.Failure.Message) diff --git a/go/core/internal/service/checkpoint/discovery_test.go b/go/core/internal/service/checkpoint/discovery_test.go new file mode 100644 index 0000000000..e5f86c583c --- /dev/null +++ b/go/core/internal/service/checkpoint/discovery_test.go @@ -0,0 +1,131 @@ +package checkpoint + +import ( + "context" + "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" + "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/service/serviceerrors" + "github.com/kagent-dev/kagent/go/core/pkg/auth" +) + +func TestOwnerDiscoversFailedCheckpointAndDeletesPartialTag(t *testing.T) { + persistence := reviewStore(t) + owner := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) + other := auth.AuthSessionTo(t.Context(), testSession{userID: "bob"}) + 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() + created, err := service.Create(owner, source.Id, request, head) + require.Equal(t, codes.Unavailable, status.Code(err), "an uncertain attempt stays retryable") + require.Nil(t, created) + pending, err := service.List(owner, ListRequest{SessionID: source.Id}) + require.NoError(t, err) + require.Empty(t, pending.Checkpoints, "CREATING remains hidden") + + _, terminalErr := service.Create(owner, source.Id, request, head) + require.Error(t, terminalErr) + // Discover the identity through public methods only. Neither the backend + // fixture nor an internal lifecycle read supplies the ID used for cleanup. + listed, err := service.List(owner, ListRequest{SessionID: source.Id, PageSize: 1}) + require.NoError(t, err) + require.Len(t, listed.Checkpoints, 1) + failed := listed.Checkpoints[0] + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_FAILED, failed.State) + require.Empty(t, listed.NextPageToken) + assertFailedCheckpointIdentity(t, terminalErr, failed.Id) + read, err := service.Get(owner, failed.Id) + require.NoError(t, err) + require.True(t, proto.Equal(failed, read)) + _, err = service.Create(owner, source.Id, request, head) + assertFailedCheckpointIdentity(t, err, failed.Id) + require.Equal(t, 2, tags.createCalls, "terminal retries must not repeat copying") + + _, err = service.Get(other, failed.Id) + require.True(t, serviceerrors.IsCode(err, serviceerrors.CodeNotFound)) + foreign, err := service.List(other, ListRequest{SessionID: source.Id}) + require.NoError(t, err) + require.Empty(t, foreign.Checkpoints) + require.True(t, serviceerrors.IsCode(service.Delete(other, failed.Id), serviceerrors.CodeNotFound)) + _, err = service.Fork(owner, failed.Id, uuid.NewString()) + require.True(t, serviceerrors.IsCode(err, serviceerrors.CodeNotFound), "FAILED must never become a fork source") + _, _, err = persistence.ForkSession(owner, failed.Id, "alice", uuid.NewString(), uuid.NewString()) + require.ErrorIs(t, err, database.ErrNotFound, "the store must independently reject FAILED forks") + + tags.deleteErr = status.Error(codes.PermissionDenied, "private object-store configuration") + require.Equal(t, codes.PermissionDenied, status.Code(service.Delete(owner, failed.Id))) + retained, err := persistence.GetCheckpointOperation(owner, failed.Id, "alice") + require.NoError(t, err) + require.Equal(t, "tag-uid", retained.TagUID, "failed cleanup must retain its deletion handle") + _, err = service.Get(owner, failed.Id) + require.True(t, serviceerrors.IsCode(err, serviceerrors.CodeNotFound), "DELETING remains hidden") + tags.deleteErr = nil + require.NoError(t, NewService(persistence, testAuthorizer{}, tags, nil).Delete(owner, failed.Id)) + require.Nil(t, tags.created) + _, err = persistence.GetCheckpointOperation(owner, failed.Id, "alice") + require.ErrorIs(t, err, database.ErrNotFound, "successful cleanup must release the retained row and references") + listed, err = service.List(owner, ListRequest{SessionID: source.Id}) + require.NoError(t, err) + require.Empty(t, listed.Checkpoints) +} + +func TestFailedCheckpointReadsHideLegacyBackendDetails(t *testing.T) { + persistence := reviewStore(t) + owner := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) + source, head := recoverySource(t, persistence) + ctx, cancel := context.WithTimeout(owner, database.RuntimeOperationTimeout) + defer cancel() + executor := uuid.New() + checkpoint, _, err := persistence.ReserveSessionCheckpoint(ctx, + &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: source.Id, HeadTaskId: head}, "alice", uuid.NewString(), executor) + require.NoError(t, err) + const privateFailure = "private bucket and credentials in a legacy failure" + _, err = persistence.FinalizeClaimedSessionCheckpoint(ctx, checkpoint, executor, "legacy-tag", "", privateFailure) + require.NoError(t, err) + service := NewService(persistence, testAuthorizer{}, &testTags{}, nil) + read, err := service.Get(owner, checkpoint.Id) + require.NoError(t, err) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_FAILED, read.State) + require.NotEmpty(t, read.GetFailure().GetMessage()) + require.NotContains(t, read.String(), "private") + listed, err := service.List(owner, ListRequest{SessionID: source.Id}) + require.NoError(t, err) + require.Len(t, listed.Checkpoints, 1) + require.True(t, proto.Equal(read, listed.Checkpoints[0])) + stored, err := persistence.GetCheckpointOperation(owner, checkpoint.Id, "alice") + require.NoError(t, err) + require.Equal(t, privateFailure, stored.Checkpoint.Failure.Message, "public projection must not rewrite durable data") + require.NoError(t, service.Delete(owner, checkpoint.Id)) +} + +func assertFailedCheckpointIdentity(t *testing.T, err error, id string) { + t.Helper() + require.Equal(t, codes.FailedPrecondition, status.Code(err), "a terminal failure must not invite transient retries") + var reason, resource string + for _, detail := range status.Convert(err).Details() { + switch info := detail.(type) { + case *errdetails.ErrorInfo: + require.Equal(t, "kagent.dev", info.Domain) + reason = info.Reason + case *errdetails.ResourceInfo: + require.Equal(t, "kagent.api.v1.Checkpoint", info.ResourceType) + resource = info.ResourceName + } + } + require.Equal(t, "KAGENT_CHECKPOINT_FAILED", reason) + require.Equal(t, id, resource) + require.NotContains(t, err.Error(), "private") +} diff --git a/go/core/internal/service/checkpoint/recovery.go b/go/core/internal/service/checkpoint/recovery.go index d5fea229a5..b23e522bc9 100644 --- a/go/core/internal/service/checkpoint/recovery.go +++ b/go/core/internal/service/checkpoint/recovery.go @@ -145,7 +145,7 @@ func (s *Service) execute(ctx context.Context, op *database.CheckpointOperation, } failure := "" if err != nil { - failure = "Checkpoint snapshot copy could not be verified" + failure = checkpointFailureMessage } // Under the required backend serialization contract, the definitive reply // permits FAILED to release admission and retain its UID for explicit Delete. diff --git a/go/core/internal/service/checkpoint/review_test.go b/go/core/internal/service/checkpoint/review_test.go index a9d1c49eb1..1d995e2b1f 100644 --- a/go/core/internal/service/checkpoint/review_test.go +++ b/go/core/internal/service/checkpoint/review_test.go @@ -125,7 +125,7 @@ func TestCheckpointRequestsAndRecovery(t *testing.T) { service := NewService(store, testAuthorizer{}, tags, nil) request := uuid.NewString() _, err := service.Create(owner, source.Id, request, head) - require.True(t, serviceerrors.IsCode(err, serviceerrors.CodeUnavailable)) + 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) diff --git a/go/core/internal/service/checkpoint/service.go b/go/core/internal/service/checkpoint/service.go index 514c0c86ba..7da4acabbd 100644 --- a/go/core/internal/service/checkpoint/service.go +++ b/go/core/internal/service/checkpoint/service.go @@ -12,6 +12,7 @@ import ( "google.golang.org/genproto/googleapis/rpc/errdetails" "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" @@ -21,8 +22,9 @@ import ( ) const ( - defaultPageSize = 50 - maxPageSize = 100 + defaultPageSize = 50 + maxPageSize = 100 + checkpointFailureMessage = "Checkpoint snapshot copy could not be verified" ) type store interface { @@ -249,7 +251,7 @@ func (s *Service) Get(ctx context.Context, checkpointID string) (*apiv1.Checkpoi if err != nil { return nil, serviceerrors.NewInternal("Failed to get checkpoint", err) } - return checkpoint, nil + return publicCheckpoint(checkpoint), nil } func (s *Service) List(ctx context.Context, request ListRequest) (ListResult, error) { @@ -277,6 +279,9 @@ func (s *Service) List(ctx context.Context, request ListRequest) (ListResult, er return ListResult{}, serviceerrors.NewInternal("Failed to list checkpoints", err) } result := ListResult{Checkpoints: rows[:min(len(rows), pageSize)]} + for i, checkpoint := range result.Checkpoints { + result.Checkpoints[i] = publicCheckpoint(checkpoint) + } if len(rows) > pageSize { result.NextPageToken = encodePageToken(rows[pageSize-1].GetId()) } @@ -503,12 +508,31 @@ func createOutcome(checkpoint *apiv1.Checkpoint) (*apiv1.Checkpoint, error) { case apiv1.CheckpointState_CHECKPOINT_STATE_READY: return checkpoint, nil case apiv1.CheckpointState_CHECKPOINT_STATE_FAILED: - return nil, serviceerrors.NewUnavailable("Failed to retain checkpoint snapshot", nil) + failure, err := status.New(codes.FailedPrecondition, + "Checkpoint creation failed; delete the failed checkpoint before creating a new request").WithDetails( + &errdetails.ErrorInfo{Domain: "kagent.dev", Reason: "KAGENT_CHECKPOINT_FAILED"}, + &errdetails.ResourceInfo{ResourceType: "kagent.api.v1.Checkpoint", ResourceName: checkpoint.Id}, + ) + if err != nil { + return nil, serviceerrors.NewInternal("Failed to report checkpoint failure", err) + } + return nil, failure.Err() default: return nil, serviceerrors.NewAborted("Checkpoint operation is still pending; retry the same request", nil) } } +// Failed records from older controllers can contain raw backend errors. Public +// reads expose a generic failure without rewriting the retained operation data. +func publicCheckpoint(checkpoint *apiv1.Checkpoint) *apiv1.Checkpoint { + if checkpoint.State != apiv1.CheckpointState_CHECKPOINT_STATE_FAILED { + return checkpoint + } + result := proto.CloneOf(checkpoint) + result.Failure = &apiv1.Failure{Reason: "SnapshotTagFailed", Message: checkpointFailureMessage} + return result +} + // 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 { From 5620b6fe816ffa207309786902587e06f9a85e98 Mon Sep 17 00:00:00 2001 From: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> Date: Mon, 9 Nov 2026 15:12:55 +0800 Subject: [PATCH 15/20] fix(checkpoint): complete public failure cleanup paths Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> --- .../persistence-checkpoints-and-forks.md | 9 + .../internal/mcp/checkpoint_errors_test.go | 169 ++++++++++++++++++ go/core/internal/mcp/checkpoints.go | 52 ++++-- .../service/checkpoint/public_errors_test.go | 38 ++++ .../internal/service/checkpoint/recovery.go | 4 +- ui/playwright/tests/chat/checkpoints.spec.ts | 49 +++++ ui/src/api/domain/checkpoints.ts | 6 +- ui/src/api/grpc/checkpoints.test.ts | 36 +++- ui/src/api/grpc/operations.ts | 37 ++-- .../components/chat/SnapshotDetailsModal.tsx | 15 +- ui/src/mocks/mockBackend.test.ts | 17 ++ ui/src/mocks/state.ts | 2 + ui/src/mocks/transport.ts | 18 +- ui/src/pages/AgentChatPage.tsx | 55 +++++- 14 files changed, 467 insertions(+), 40 deletions(-) create mode 100644 go/core/internal/mcp/checkpoint_errors_test.go create mode 100644 go/core/internal/service/checkpoint/public_errors_test.go diff --git a/docs/architecture/persistence-checkpoints-and-forks.md b/docs/architecture/persistence-checkpoints-and-forks.md index f4d485cda4..7b9c48cd1f 100644 --- a/docs/architecture/persistence-checkpoints-and-forks.md +++ b/docs/architecture/persistence-checkpoints-and-forks.md @@ -144,6 +144,15 @@ delete the failed checkpoint before creating a new request. Public failed checkpoint responses use a generic service-owned failure message, including for legacy rows that retained raw backend errors; the stored operation is not rewritten. +MCP checkpoint tools use the same classified public errors, excluding storage +causes and wrapped cancellation details. Backend status is retained as +`kagent.dev/grpc-code`; terminal creation failures also retain +`kagent.dev/error-reason` and `kagent.dev/checkpoint-id` for cleanup. The Web UI +lists failed snapshots separately from saved transcript boundaries, refreshes +after a failed creation reply, and reads subsequent list pages. Failed records +can be inspected and deleted, but cannot be forked or renamed. A failed deletion +keeps its details open for retry. + ### Proposed initial upgrade boundary: pending maintainer acceptance The proposed first upgrade from a controller without checkpoint execution claims 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..d4f0a88f4c --- /dev/null +++ b/go/core/internal/mcp/checkpoint_errors_test.go @@ -0,0 +1,169 @@ +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 TestCheckpointToolsRetainTerminalIdentity(t *testing.T) { + store := &checkpointErrorStore{phase: "failed"} + h := &Handler{checkpoints: checkpoint.NewService(store, &auth.NoopAuthorizer{}, nil, nil)} + result := checkpointToolResponse(t, h.checkpoints, createCheckpointToolName, map[string]any{ + "session_id": testSessionID, "expected_head_task_id": testTaskID, "request_id": "failed-request"}) + require.Equal(t, true, result["isError"]) + meta := result["_meta"].(map[string]any) + require.Equal(t, "FailedPrecondition", meta["kagent.dev/grpc-code"]) + require.Equal(t, "KAGENT_CHECKPOINT_FAILED", meta["kagent.dev/error-reason"]) + id := meta["kagent.dev/checkpoint-id"].(string) + require.Equal(t, "33333333-3333-4333-8333-333333333333", id) + read, output, err := h.getCheckpoint(authContext(), nil, CheckpointIDInput{CheckpointID: id}) + require.NoError(t, err) + require.False(t, read.IsError) + require.Equal(t, id, output.Checkpoint.ID) + require.NotContains(t, output.Checkpoint.Failure.Message, "private") +} + +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")} + h := &Handler{checkpoints: checkpoint.NewService(&checkpointErrorStore{phase: "backend"}, &auth.NoopAuthorizer{}, backend, nil)} + result, _, err := h.deleteCheckpoint(authContext(), nil, CheckpointIDInput{CheckpointID: "33333333-3333-4333-8333-333333333333"}) + require.NoError(t, err) + require.True(t, result.IsError) + require.Equal(t, 1, backend.calls) + require.Equal(t, code.String(), result.Meta["kagent.dev/grpc-code"]) + require.NotContains(t, fmt.Sprint(result.Content), "private") + }) + } +} + +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)} + h := &Handler{checkpoints: checkpoint.NewService(store, &auth.NoopAuthorizer{}, nil, nil)} + result, _, err := h.createCheckpoint(authContext(), nil, CreateCheckpointInput{ + SessionID: testSessionID, ExpectedHeadTaskID: testTaskID, RequestID: "canceled-request"}) + require.NoError(t, err) + require.True(t, result.IsError) + require.Equal(t, status.FromContextError(cause).Code().String(), result.Meta["kagent.dev/grpc-code"]) + require.NotContains(t, fmt.Sprint(result.Content), "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 +} + +func (s *checkpointErrorStore) ReserveSessionCheckpoint(_ context.Context, cp *apiv1.Checkpoint, _ string, _ string, _ uuid.UUID) (*apiv1.Checkpoint, *database.SessionTaskSnapshot, error) { + 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 + } + if s.phase == "failed" { + cp.State = apiv1.CheckpointState_CHECKPOINT_STATE_FAILED + cp.Failure = &apiv1.Failure{Message: "private legacy failure"} + } + 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) 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 +} + +func (s *checkpointErrorStore) GetSessionCheckpoint(context.Context, string, string) (*apiv1.Checkpoint, error) { + return &apiv1.Checkpoint{Id: "33333333-3333-4333-8333-333333333333", State: apiv1.CheckpointState_CHECKPOINT_STATE_FAILED, Failure: &apiv1.Failure{Message: "private legacy failure"}}, nil +} + +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..f8b0b72e97 100644 --- a/go/core/internal/mcp/checkpoints.go +++ b/go/core/internal/mcp/checkpoints.go @@ -2,16 +2,19 @@ package mcp import ( "context" + "errors" "fmt" "time" "github.com/google/uuid" "github.com/modelcontextprotocol/go-sdk/mcp" "google.golang.org/genproto/googleapis/rpc/errdetails" + "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/service/checkpoint" + "github.com/kagent-dev/kagent/go/core/internal/service/serviceerrors" ) const ( @@ -87,7 +90,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 +98,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,13 +106,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) - 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} - } - } - return result, CreateCheckpointOutput{}, nil + return checkpointToolError(err), CreateCheckpointOutput{}, nil } output := CreateCheckpointOutput{Checkpoint: checkpointSummary(created)} return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: fmt.Sprintf("Created checkpoint %s", created.GetId())}}}, output, nil @@ -121,7 +118,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 +130,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 +143,39 @@ func stableRequestID(id string) string { return uuid.NewString() } +// Checkpoint services classify storage failures and sanitize backend statuses. +// Project only their public message, retaining status and cleanup identity for +// MCP clients that do not pass through the gRPC error-mapping interceptor. +func checkpointToolError(err error) *mcp.CallToolResult { + if errors.Is(err, context.Canceled) { + err = status.Error(codes.Canceled, "Checkpoint operation canceled") + } else if errors.Is(err, context.DeadlineExceeded) { + err = status.Error(codes.DeadlineExceeded, "Checkpoint operation deadline exceeded") + } + if serviceerrors.CodeOf(err) != "" { + return sessionToolError(err) + } + value, ok := status.FromError(err) + if !ok { + return sessionToolError(serviceerrors.NewInternal("Checkpoint operation could not complete", err)) + } + result := toolError(errors.New(value.Message())) + result.Meta = mcp.Meta{"kagent.dev/grpc-code": value.Code().String()} + for _, detail := range value.Details() { + switch info := detail.(type) { + case *errdetails.ErrorInfo: + if info.Domain == "kagent.dev" { + result.Meta["kagent.dev/error-reason"] = info.Reason + } + case *errdetails.ResourceInfo: + if info.ResourceType == "kagent.api.v1.Checkpoint" { + result.Meta["kagent.dev/checkpoint-id"] = info.ResourceName + } + } + } + return result +} + 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/public_errors_test.go b/go/core/internal/service/checkpoint/public_errors_test.go new file mode 100644 index 0000000000..218aa19c33 --- /dev/null +++ b/go/core/internal/service/checkpoint/public_errors_test.go @@ -0,0 +1,38 @@ +package checkpoint + +import ( + "errors" + "testing" + + "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 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") + }) + } +} diff --git a/go/core/internal/service/checkpoint/recovery.go b/go/core/internal/service/checkpoint/recovery.go index b23e522bc9..6b887d3c89 100644 --- a/go/core/internal/service/checkpoint/recovery.go +++ b/go/core/internal/service/checkpoint/recovery.go @@ -91,7 +91,7 @@ func (s *Service) claimAndExecute(ctx context.Context, observed *database.Checkp if err != nil { // A lost COMMIT reply can hide successful acquisition. Retire this // executor before returning the uncertain error, as on reservation. - return nil, errors.Join(err, s.release(ctx, id, executor)) + return nil, errors.Join(publicOperationError(err), s.release(ctx, id, executor)) } return s.execute(ctx, op, executor) } @@ -161,7 +161,7 @@ func (s *Service) execute(ctx context.Context, op *database.CheckpointOperation, return s.currentOutcome(finishCtx, op.Checkpoint.Id, op.Creator) } if err != nil { - return nil, err + return nil, publicOperationError(err) } releaseNeeded = false return createOutcome(checkpoint) diff --git a/ui/playwright/tests/chat/checkpoints.spec.ts b/ui/playwright/tests/chat/checkpoints.spec.ts index ec5318a116..5cfe0e24bf 100644 --- a/ui/playwright/tests/chat/checkpoints.spec.ts +++ b/ui/playwright/tests/chat/checkpoints.spec.ts @@ -13,6 +13,55 @@ const dividers = (page: import("@playwright/test").Page) => /** The seeded boundary's generated name: its conversation, joined to its turn. */ const SEEDED_NAME = `${instances.ready}-seed-task-1`; +test.describe("chat: failed snapshots", () => { + test.use({ expectedNoise: [/Could not (?:checkpoint the conversation|delete the checkpoint):/] }); + + test("chat: a failed creation can be found, inspected and deleted after reload", async ({ page }) => { + await page.goto(`${agentChat(instances.ready)}?checkpoint=failed`); + await expect(dividers(page)).toHaveCount(1); + await sendAndAwaitTurn(page, "A new turn for a failed snapshot."); + await page.getByTestId("chat-checkpoint").click(); + + const notice = page.getByTestId("chat-failed-snapshots"); + await expect(notice).toBeVisible(); + // A failure is a cleanup record, not a saved history boundary. + await expect(dividers(page)).toHaveCount(1); + await expect(page.getByTestId("chat-checkpoint")).toBeEnabled(); + const failed = notice.locator('[data-testid^="chat-failed-snapshot-"]'); + const id = (await failed.getAttribute("data-testid"))!.replace("chat-failed-snapshot-", ""); + + await page.reload(); + await expect(page.getByTestId(`chat-failed-snapshot-${id}`)).toBeVisible(); + await page.getByTestId(`chat-failed-snapshot-${id}`).click(); + await expect(page.getByTestId("snapshot-details-state")).toHaveAttribute("data-state", "failed"); + await expect(page.getByTestId("snapshot-details-fields")).toContainText(id); + await expect(page.getByTestId("snapshot-details-fork")).toBeDisabled(); + await expect(page.getByTestId("snapshot-details-rename")).toBeDisabled(); + await page.evaluate(() => { + const url = new URL(window.location.href); + url.searchParams.set("checkpointDelete", "denied"); + window.history.replaceState({}, "", url); + }); + await page.getByTestId("snapshot-details-delete").click(); + await pressOnce(page.getByTestId("snapshot-details-delete-confirm")); + await expect(page.getByTestId("snapshot-action-error")).toBeVisible(); + await expect(page.getByTestId("snapshot-details-body")).toBeVisible(); + await expect(notice).toBeVisible(); + await page.evaluate(() => { + const url = new URL(window.location.href); + url.searchParams.delete("checkpointDelete"); + window.history.replaceState({}, "", url); + }); + await page.getByTestId("snapshot-details-delete").click(); + await pressOnce(page.getByTestId("snapshot-details-delete-confirm")); + await expect(notice).toHaveCount(0); + await expect(page.getByTestId("snapshot-details-body")).toHaveCount(0); + await page.reload(); + await expect(page.getByTestId("chat-failed-snapshots")).toHaveCount(0); + await expect(dividers(page)).toHaveCount(1); + }); +}); + /** Opens one boundary's record, which is what pressing its line does. */ async function openSnapshot( page: import("@playwright/test").Page, diff --git a/ui/src/api/domain/checkpoints.ts b/ui/src/api/domain/checkpoints.ts index 51f556a1e5..62c12d4a8e 100644 --- a/ui/src/api/domain/checkpoints.ts +++ b/ui/src/api/domain/checkpoints.ts @@ -14,9 +14,9 @@ * How far the controller has got with retaining the snapshot. * * `CheckpointState` in the proto, spelled as words for the reasons given in - * `agentInstances.ts`. Creation answers `ready` or `failed` and never `creating` — - * the RPC copies the snapshot before it returns — so a caller that wants to fork can - * check the state it was handed rather than poll. + * `agentInstances.ts`. Successful creation answers `ready`; a terminal failure is + * an error. Owner queries also return `failed` records so they can be deleted. + * Only `ready` records can be forked or renamed. */ export type CheckpointState = | "unspecified" diff --git a/ui/src/api/grpc/checkpoints.test.ts b/ui/src/api/grpc/checkpoints.test.ts index 71ed7f9488..6a473a6d40 100644 --- a/ui/src/api/grpc/checkpoints.test.ts +++ b/ui/src/api/grpc/checkpoints.test.ts @@ -41,7 +41,7 @@ it("forks through Session RPCs and reads the new conversation identity", async ( expect(result).toMatchObject({ id: "forked-session", contextId: "forked-session", name: "Saved turn" }); }); -it.each(["KAGENT_CHECKPOINT_CONVERSATION_ADVANCED", undefined])("does not retry %s", async (reason) => { +it.each(["KAGENT_CHECKPOINT_CONVERSATION_ADVANCED", "KAGENT_CHECKPOINT_FAILED", undefined])("does not retry %s", async (reason) => { let calls = 0; setApiTransport(createRouterTransport((router) => { router.service(CheckpointService, { @@ -56,3 +56,37 @@ it.each(["KAGENT_CHECKPOINT_CONVERSATION_ADVANCED", undefined])("does not retry await expect(defaultOperations["agentInstances.checkpoints.create"]({ id: "instance", requestId: "request", expectedHeadTaskId: "turn-a" }, {})).rejects.toMatchObject({ reason }); expect(calls).toBe(1); }); + +it("reads failed cleanup candidates beyond the first page", async () => { + const tokens: string[] = []; + setApiTransport(createRouterTransport((router) => { + router.service(CheckpointService, { + listCheckpoints(request) { + tokens.push(request.page?.pageToken ?? ""); + return request.page?.pageToken + ? { checkpoints: [{ id: "failed", sessionId: request.sessionId, state: CheckpointState.FAILED, failure: { message: "Copy could not be verified" } }], page: {} } + : { checkpoints: [{ id: "ready", sessionId: request.sessionId, state: CheckpointState.READY }], page: { nextPageToken: "next" } }; + }, + }); + })); + const result = await defaultOperations["agentInstances.checkpoints.list"]({ id: "instance" }, {}); + expect(tokens).toEqual(["", "next"]); + expect(result).toEqual(expect.arrayContaining([ + expect.objectContaining({ id: "ready", state: "ready" }), + expect.objectContaining({ id: "failed", state: "failed", failure: "Copy could not be verified" }), + ])); +}); + +it("rejects a repeated snapshot page rather than returning an incomplete list", async () => { + let calls = 0; + setApiTransport(createRouterTransport((router) => { + router.service(CheckpointService, { + listCheckpoints() { + calls++; + return { checkpoints: [], page: { nextPageToken: "same" } }; + }, + }); + })); + await expect(defaultOperations["agentInstances.checkpoints.list"]({ id: "instance" }, {})).rejects.toThrow("repeated a page"); + expect(calls).toBe(2); +}); diff --git a/ui/src/api/grpc/operations.ts b/ui/src/api/grpc/operations.ts index 2ae8b2d8bc..230a6f93e3 100644 --- a/ui/src/api/grpc/operations.ts +++ b/ui/src/api/grpc/operations.ts @@ -625,6 +625,7 @@ function toAgentInstance(instance: PbSession): AgentInstance { * worth reporting rather than a list worth truncating. */ const INSTANCE_PAGE_LIMIT = 50; +const CHECKPOINT_PAGE_LIMIT = 50; function toAgentInstanceShare(share: PbSessionShare): AgentInstanceShare { return { @@ -803,22 +804,30 @@ const agentInstances: Pick< return checkpoint; }, - /* - * Newest first, because that is the order the chat needs them in and the controller - * does not promise one. One page: a conversation's boundaries are counted in - * handfuls, and paging a list this short would be machinery with nothing to do. - */ + // FAILED cleanup candidates can be on later pages. Read the complete list + // before sorting it newest-first, just as for conversation discovery above. "agentInstances.checkpoints.list": async (input, options) => { const name = "CheckpointService/ListCheckpoints"; - const response = await rpc(name, options.signal, () => - serviceClient(CheckpointService).listCheckpoints( - { sessionId: input.id }, - call("agentInstances.checkpoints.list", options), - ), - ); - return list(response.checkpoints) - .map(toCheckpoint) - .sort((left, right) => (right.createdAt ?? "").localeCompare(left.createdAt ?? "")); + const rows: Checkpoint[] = []; + let pageToken = ""; + const seen = new Set(); + for (let page = 0; page < CHECKPOINT_PAGE_LIMIT; page += 1) { + const response = await rpc(name, options.signal, () => + serviceClient(CheckpointService).listCheckpoints( + { sessionId: input.id, page: { pageToken } }, + call("agentInstances.checkpoints.list", options), + ), + ); + rows.push(...list(response.checkpoints).map(toCheckpoint)); + const next = response.page?.nextPageToken ?? ""; + if (!next) return rows.sort((left, right) => (right.createdAt ?? "").localeCompare(left.createdAt ?? "")); + if (seen.has(next)) { + throw new ApiError("The API repeated a page of snapshots instead of advancing.", { kind: "parse", url: name }); + } + seen.add(next); + pageToken = next; + } + throw new ApiError("The snapshot list could not be read to the end.", { kind: "parse", url: name }); }, /* diff --git a/ui/src/components/chat/SnapshotDetailsModal.tsx b/ui/src/components/chat/SnapshotDetailsModal.tsx index 249fbe39f5..a0ece8a0fc 100644 --- a/ui/src/components/chat/SnapshotDetailsModal.tsx +++ b/ui/src/components/chat/SnapshotDetailsModal.tsx @@ -1,5 +1,5 @@ import { useState } from "react"; -import { Button, Descriptions, Modal, Popconfirm, Space, Tag, Tooltip, Typography } from "antd"; +import { Alert, Button, Descriptions, Modal, Popconfirm, Space, Tag, Tooltip, Typography } from "antd"; import { Eraser, GitFork, Pencil } from "lucide-react"; import { useTheme } from "@emotion/react"; import { canForkFrom, type Checkpoint, type CheckpointState } from "@/api"; @@ -67,15 +67,19 @@ export function SnapshotDetailsModal({ const theme = useTheme(); const [isRenaming, setRenaming] = useState(false); const [isBusy, setBusy] = useState(false); + const [actionError, setActionError] = useState(); const label = snapshotLabel(checkpoint); const forkable = canForkFrom(checkpoint); const state = STATE_APPEARANCE[checkpoint.state]; async function act(run: () => void | Promise) { setBusy(true); + setActionError(undefined); try { await run(); onClose(); + } catch (cause: unknown) { + setActionError(cause instanceof Error ? cause.message : String(cause)); } finally { setBusy(false); } @@ -106,11 +110,12 @@ export function SnapshotDetailsModal({ children: ( {label} - + + + ))} + + + } + /> + ) : null} + {checkpoints.error ? ( + void checkpoints.refresh().catch((cause: unknown) => { + console.error("Could not re-read the saved boundaries:", cause); + })}> + Retry + + } + /> + ) : null} Date: Mon, 9 Nov 2026 21:25:48 +0800 Subject: [PATCH 16/20] fix(checkpoint): reconcile snapshot deletion retries in the UI Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> --- ui/playwright/tests/chat/checkpoints.spec.ts | 71 ++++++++++++++++++++ ui/src/api/grpc/checkpoints.test.ts | 55 +++++++++++++++ ui/src/api/grpc/operations.ts | 37 +++++++--- ui/src/components/chat/ChatTranscript.tsx | 2 +- ui/src/components/chat/CheckpointDivider.tsx | 11 ++- ui/src/mocks/transport.ts | 13 +++- 6 files changed, 177 insertions(+), 12 deletions(-) diff --git a/ui/playwright/tests/chat/checkpoints.spec.ts b/ui/playwright/tests/chat/checkpoints.spec.ts index 5cfe0e24bf..4a1015a953 100644 --- a/ui/playwright/tests/chat/checkpoints.spec.ts +++ b/ui/playwright/tests/chat/checkpoints.spec.ts @@ -13,6 +13,77 @@ const dividers = (page: import("@playwright/test").Page) => /** The seeded boundary's generated name: its conversation, joined to its turn. */ const SEEDED_NAME = `${instances.ready}-seed-task-1`; +async function deleteOutcome(page: import("@playwright/test").Page, outcome?: string) { + await page.evaluate((value) => { + const url = new URL(window.location.href); + if (value) url.searchParams.set("checkpointDelete", value); + else url.searchParams.delete("checkpointDelete"); + window.history.replaceState({}, "", url); + }, outcome); +} + +test.describe("chat: snapshot delete failures", () => { + test.use({ expectedNoise: [/Could not (?:checkpoint the conversation|delete the checkpoint):/] }); + + test("chat: a refused transcript delete reports failure and can be retried", async ({ page }) => { + await page.goto(`${agentChat(instances.ready)}?checkpointDelete=denied`); + const line = dividers(page).first(); + await expect(line).toBeVisible(); + const id = (await line.getAttribute("data-testid"))!.replace("chat-checkpoint-mark-", ""); + await line.getByTestId(`chat-checkpoint-delete-${id}`).click(); + await pressOnce(page.getByTestId(`chat-checkpoint-delete-confirm-${id}`)); + await expect(page.getByText("Could not delete: Checkpoint operation could not complete", { exact: true })).toBeVisible(); + await expect(line).toBeVisible(); + await expect(page.getByText("Delete this snapshot?", { exact: true })).toBeHidden(); + await deleteOutcome(page); + await line.getByTestId(`chat-checkpoint-delete-${id}`).click(); + await pressOnce(page.getByTestId(`chat-checkpoint-delete-confirm-${id}`)); + await expect(dividers(page)).toHaveCount(0); + await page.reload(); + await expect(page.getByTestId("chat-checkpoint")).toBeVisible(); + await expect(dividers(page)).toHaveCount(0); + }); + + test("chat: a completed cleanup closes stale failed details on retry", async ({ page }) => { + await page.goto(`${agentChat(instances.ready)}?checkpoint=failed`); + await expect(dividers(page)).toHaveCount(1); + await sendAndAwaitTurn(page, "A new turn for cleanup after a lost deletion reply."); + await page.getByTestId("chat-checkpoint").click(); + const notice = page.getByTestId("chat-failed-snapshots"); + await expect(notice).toBeVisible(); + await notice.locator('[data-testid^="chat-failed-snapshot-"]').click(); + await deleteOutcome(page, "lost"); + await page.getByTestId("snapshot-details-delete").click(); + await pressOnce(page.getByTestId("snapshot-details-delete-confirm")); + await expect(page.getByTestId("snapshot-action-error")).toBeVisible(); + await expect(notice).toBeVisible(); + await expect(page.getByText("Delete this snapshot?", { exact: true })).toBeHidden(); + await page.getByTestId("snapshot-details-delete").click(); + await pressOnce(page.getByTestId("snapshot-details-delete-confirm")); + await expect(page.getByTestId("snapshot-details-body")).toHaveCount(0); + await expect(notice).toHaveCount(0); + await page.reload(); + await expect(dividers(page)).toHaveCount(1); + await expect(notice).toHaveCount(0); + }); + + test("chat: a referenced snapshot is retained when deletion returns not found", async ({ page }) => { + await page.goto(`${agentChat(instances.ready)}?checkpointDelete=referenced`); + const line = dividers(page).first(); + await expect(line).toBeVisible(); + await openSnapshot(page, line); + await page.getByTestId("snapshot-details-delete").click(); + await pressOnce(page.getByTestId("snapshot-details-delete-confirm")); + await expect(page.getByTestId("snapshot-action-error")).toBeVisible(); + await expect(page.getByTestId("snapshot-details-body")).toBeVisible(); + await expect(line).toBeVisible(); + await page.reload(); + await expect(dividers(page)).toHaveCount(1); + await openSnapshot(page, dividers(page).first()); + await expect(page.getByTestId("snapshot-details-state")).toHaveAttribute("data-state", "ready"); + }); +}); + test.describe("chat: failed snapshots", () => { test.use({ expectedNoise: [/Could not (?:checkpoint the conversation|delete the checkpoint):/] }); diff --git a/ui/src/api/grpc/checkpoints.test.ts b/ui/src/api/grpc/checkpoints.test.ts index 6a473a6d40..29835e2e29 100644 --- a/ui/src/api/grpc/checkpoints.test.ts +++ b/ui/src/api/grpc/checkpoints.test.ts @@ -90,3 +90,58 @@ it("rejects a repeated snapshot page rather than returning an incomplete list", await expect(defaultOperations["agentInstances.checkpoints.list"]({ id: "instance" }, {})).rejects.toThrow("repeated a page"); expect(calls).toBe(2); }); + +it("confirms a missing snapshot after a delete retry without masking a retained record", async () => { + let retained = true; + const calls: string[] = []; + setApiTransport(createRouterTransport((router) => { + router.service(CheckpointService, { + deleteCheckpoint(request) { + calls.push(`delete:${request.checkpointId}`); + throw new ConnectError("checkpoint still referenced or already deleted", Code.NotFound); + }, + getCheckpoint(request) { + calls.push(`get:${request.checkpointId}`); + if (!retained) throw new ConnectError("checkpoint absent", Code.NotFound); + return { checkpoint: { id: request.checkpointId, state: CheckpointState.READY } }; + }, + }); + })); + const remove = () => defaultOperations["agentInstances.checkpoints.delete"]({ checkpointId: "checkpoint" }, {}); + await expect(remove()).rejects.toMatchObject({ code: "NotFound", message: "checkpoint still referenced or already deleted" }); + retained = false; + await expect(remove()).resolves.toBeUndefined(); + expect(calls).toEqual(["delete:checkpoint", "get:checkpoint", "delete:checkpoint", "get:checkpoint"]); +}); + +it.each([Code.PermissionDenied, Code.Unavailable, "empty" as const])("preserves a failed absence check (%s)", async (failure) => { + setApiTransport(createRouterTransport((router) => { + router.service(CheckpointService, { + deleteCheckpoint() { throw new ConnectError("checkpoint absent or referenced", Code.NotFound); }, + getCheckpoint() { + if (failure === "empty") return {}; + throw new ConnectError("lookup failed", failure); + }, + }); + })); + await expect(defaultOperations["agentInstances.checkpoints.delete"]({ checkpointId: "checkpoint" }, {})).rejects.toMatchObject( + failure === "empty" ? { kind: "parse" } : { code: Code[failure], message: "lookup failed" }, + ); +}); + +it.each([undefined, Code.PermissionDenied, Code.Unavailable])("does not look up a snapshot after delete outcome %s", async (failure) => { + let lookups = 0; + setApiTransport(createRouterTransport((router) => { + router.service(CheckpointService, { + deleteCheckpoint() { + if (failure !== undefined) throw new ConnectError("delete failed", failure); + return {}; + }, + getCheckpoint() { lookups++; return {}; }, + }); + })); + const result = defaultOperations["agentInstances.checkpoints.delete"]({ checkpointId: "checkpoint" }, {}); + if (failure === undefined) await expect(result).resolves.toBeUndefined(); + else await expect(result).rejects.toMatchObject({ code: Code[failure] }); + expect(lookups).toBe(0); +}); diff --git a/ui/src/api/grpc/operations.ts b/ui/src/api/grpc/operations.ts index 230a6f93e3..0a44d73dae 100644 --- a/ui/src/api/grpc/operations.ts +++ b/ui/src/api/grpc/operations.ts @@ -37,8 +37,8 @@ import { ScheduledRunService } from "@/generated/kagent/api/v1/scheduled_runs_pb * * `CheckpointService` is reached through the `agentInstances.checkpoints.*` ids, and * **`agentInstances.fork`** composes two of them — a checkpoint of the conversation - * as it stands, then a fork of it. `GetCheckpoint` and `DeleteCheckpoint` have no id: - * the chat reads boundaries from the list and nothing removes one yet. + * as it stands, then a fork of it. Deletion confirms absence through `GetCheckpoint` + * after NotFound, since a retained checkpoint with dependent forks returns that too. */ import { ModelService } from "@/generated/kagent/api/v1/models_pb"; @@ -870,12 +870,33 @@ const agentInstances: Pick< }, "agentInstances.checkpoints.delete": async (input, options) => { - await rpc("CheckpointService/DeleteCheckpoint", options.signal, () => - serviceClient(CheckpointService).deleteCheckpoint( - { checkpointId: input.checkpointId }, - call("agentInstances.checkpoints.delete", options), - ), - ); + try { + await rpc("CheckpointService/DeleteCheckpoint", options.signal, () => + serviceClient(CheckpointService).deleteCheckpoint( + { checkpointId: input.checkpointId }, + call("agentInstances.checkpoints.delete", options), + ), + ); + } catch (error) { + if (!(error instanceof ApiError) || error.code !== "NotFound") throw error; + // Recovery may have completed deletion after an earlier failed reply. A + // fork-pinned snapshot also rejects Delete with NotFound, so confirm absence + // before letting callers remove their local record and close its details. + const name = "CheckpointService/GetCheckpoint"; + try { + const response = await rpc(name, options.signal, () => + serviceClient(CheckpointService).getCheckpoint( + { checkpointId: input.checkpointId }, + call("agentInstances.checkpoints.delete", options), + ), + ); + required(response.checkpoint, name, "checkpoint"); + } catch (lookupError) { + if (lookupError instanceof ApiError && lookupError.code === "NotFound") return; + throw lookupError; + } + throw error; + } }, /* diff --git a/ui/src/components/chat/ChatTranscript.tsx b/ui/src/components/chat/ChatTranscript.tsx index 0ac76403b5..9ddaecb013 100644 --- a/ui/src/components/chat/ChatTranscript.tsx +++ b/ui/src/components/chat/ChatTranscript.tsx @@ -80,7 +80,7 @@ export function ChatTranscript({ /** Starts a new conversation from a boundary. */ onFork?: (checkpointId: string) => void; /** Drops a boundary and the runtime stored with it. */ - onDeleteCheckpoint?: (checkpointId: string) => void; + onDeleteCheckpoint?: (checkpointId: string) => void | Promise; /** Opens active callbacks for one A2A task. */ onManageTaskPush?: (taskId: string) => void; /** The boundaries the controller has described, for the line to name itself by. */ diff --git a/ui/src/components/chat/CheckpointDivider.tsx b/ui/src/components/chat/CheckpointDivider.tsx index 632d85ca53..007566d787 100644 --- a/ui/src/components/chat/CheckpointDivider.tsx +++ b/ui/src/components/chat/CheckpointDivider.tsx @@ -45,7 +45,7 @@ export function CheckpointDivider({ /** Names this boundary, and with it the forks taken from it. */ onRename?: () => void; /** Removes this boundary. */ - onDelete?: () => void; + onDelete?: () => void | Promise; }) { const theme = useTheme(); const { mode } = useThemeMode(); @@ -286,7 +286,14 @@ export function CheckpointDivider({ cancelButtonProps={{ "data-testid": `chat-checkpoint-delete-cancel-${checkpointId}`, }} - onConfirm={onDelete} + onConfirm={async () => { + try { + await onDelete(); + } catch { + // The page reports the failure and retains the row for retry. + // Popconfirm rethrows rejected callbacks as uncaught errors. + } + }} > {/* No tooltip, unlike its neighbours. The pointer that opens the confirmation is still resting on this button, so the tooltip stays up diff --git a/ui/src/mocks/transport.ts b/ui/src/mocks/transport.ts index d7aca725df..cc413202e7 100644 --- a/ui/src/mocks/transport.ts +++ b/ui/src/mocks/transport.ts @@ -951,11 +951,22 @@ on(CheckpointService.method.updateCheckpointName, (input) => { return { checkpoint: checkpointMessage(renamed) }; }); +on(CheckpointService.method.getCheckpoint, (input) => { + const checkpoint = checkpointById(input.checkpointId); + if (!checkpoint) throw notFound(`Checkpoint ${input.checkpointId}`); + return { checkpoint: checkpointMessage(checkpoint) }; +}); + on(CheckpointService.method.deleteCheckpoint, (input) => { - if (new URLSearchParams(window.location.search).get("checkpointDelete") === "denied") { + const outcome = new URLSearchParams(window.location.search).get("checkpointDelete"); + if (outcome === "denied") { throw new ConnectError("Checkpoint operation could not complete", Code.PermissionDenied); } + // The store also refuses deletion with NotFound while forks reference a snapshot. + if (outcome === "referenced") throw notFound(`Checkpoint ${input.checkpointId}`); if (!deleteCheckpoint(input.checkpointId)) throw notFound(`Checkpoint ${input.checkpointId}`); + // The next retry sees NotFound: deletion committed but its reply never arrived. + if (outcome === "lost") throw new ConnectError("Checkpoint operation could not complete", Code.Unavailable); return {}; }); From 15556c9b53c82cd0df0bd1f809d09deaec84715d Mon Sep 17 00:00:00 2001 From: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> Date: Tue, 10 Nov 2026 11:00:54 +0800 Subject: [PATCH 17/20] fix(checkpoint): retire uncertain reservations and reconcile refreshes Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> --- .../persistence-checkpoints-and-forks.md | 8 ++ .../database/checkpoint_operations.go | 41 ++++-- .../database/checkpoint_reservation_test.go | 71 +++++++++++ .../internal/mcp/checkpoint_errors_test.go | 3 + .../service/checkpoint/completion_test.go | 5 + .../service/checkpoint/public_errors_test.go | 29 +++++ .../internal/service/checkpoint/recovery.go | 10 +- .../checkpoint/reservation_commit_test.go | 120 ++++++++++++++++++ .../internal/service/checkpoint/service.go | 8 +- .../service/checkpoint/service_test.go | 4 + ui/playwright/tests/chat/checkpoints.spec.ts | 52 ++++++++ ui/src/api/hooks/useApiResource.test.tsx | 65 ++++++++++ ui/src/api/hooks/useApiResource.ts | 22 ++-- ui/src/api/hooks/useCheckpoints.test.tsx | 44 +++++++ ui/src/api/hooks/useCheckpoints.ts | 21 ++- ui/src/mocks/transport.ts | 3 + ui/src/pages/AgentChatPage.tsx | 1 + ui/src/pages/ScheduledRunsPage.test.tsx | 12 +- 18 files changed, 486 insertions(+), 33 deletions(-) create mode 100644 go/core/internal/database/checkpoint_reservation_test.go create mode 100644 go/core/internal/service/checkpoint/reservation_commit_test.go create mode 100644 ui/src/api/hooks/useApiResource.test.tsx create mode 100644 ui/src/api/hooks/useCheckpoints.test.tsx diff --git a/docs/architecture/persistence-checkpoints-and-forks.md b/docs/architecture/persistence-checkpoints-and-forks.md index 7b9c48cd1f..ea08c97a4c 100644 --- a/docs/architecture/persistence-checkpoints-and-forks.md +++ b/docs/architecture/persistence-checkpoints-and-forks.md @@ -82,6 +82,11 @@ Create observers skip obsolete claim releases, so a release failure cannot mask READY or the translation of absence to `NotFound`. A definitive acquisition conflict also needs no release. Other acquisition errors may hide a committed claim; those attempts still release their specific executor before returning. +An uncertain initial reservation also waits for its source Session transaction +before looking for the inserted row: locking the checkpoint alone cannot see an +uncommitted INSERT. This coordination is confined to reservation errors and keeps +the normal two-write path unchanged. Retirement remains bounded by five seconds; +if storage prevents it, claim expiry still permits later recovery. Backend errors describe the current attempt, without persisting an error that a later repaired request could accidentally replay. Public errors exclude backend configuration. A canceled or uncertain attempt retains `CREATING` and expires its @@ -152,6 +157,9 @@ lists failed snapshots separately from saved transcript boundaries, refreshes after a failed creation reply, and reads subsequent list pages. Failed records can be inspected and deleted, but cannot be forked or renamed. A failed deletion keeps its details open for retry. +Confirmed deletion removes the cached record before rereading the list, so a +subsequent read outage cannot leave a deleted FAILED entry on screen. Failed +refreshes retain the last data and expose an error with a retry action. ### Proposed initial upgrade boundary: pending maintainer acceptance diff --git a/go/core/internal/database/checkpoint_operations.go b/go/core/internal/database/checkpoint_operations.go index 98fb2b2822..c217b0c5d8 100644 --- a/go/core/internal/database/checkpoint_operations.go +++ b/go/core/internal/database/checkpoint_operations.go @@ -73,22 +73,39 @@ func (c *Client) ClaimCheckpointOperation(ctx context.Context, id, owner string, // 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 { - // Lock by ID before testing the executor: an acquisition's COMMIT can - // still be in flight, with the old committed executor invisible to a - // conditional UPDATE. Wait for that transaction before deciding to no-op. - _, err := lockCheckpoint(ctx, tx, id, true, "", nil) + 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 } - if err != 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 - }) + } + _, 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. 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/mcp/checkpoint_errors_test.go b/go/core/internal/mcp/checkpoint_errors_test.go index d4f0a88f4c..f2f2e2cc1d 100644 --- a/go/core/internal/mcp/checkpoint_errors_test.go +++ b/go/core/internal/mcp/checkpoint_errors_test.go @@ -138,6 +138,9 @@ func (s *checkpointErrorStore) ClaimCheckpointOperation(context.Context, string, 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") diff --git a/go/core/internal/service/checkpoint/completion_test.go b/go/core/internal/service/checkpoint/completion_test.go index 16e9046c9c..19a16d6dd2 100644 --- a/go/core/internal/service/checkpoint/completion_test.go +++ b/go/core/internal/service/checkpoint/completion_test.go @@ -200,3 +200,8 @@ func (s *checkpointReleaseFailure) ReleaseCheckpointOperation(context.Context, s 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/public_errors_test.go b/go/core/internal/service/checkpoint/public_errors_test.go index 218aa19c33..57f37d0abe 100644 --- a/go/core/internal/service/checkpoint/public_errors_test.go +++ b/go/core/internal/service/checkpoint/public_errors_test.go @@ -1,8 +1,10 @@ package checkpoint import ( + "context" "errors" "testing" + "time" "github.com/google/uuid" "github.com/stretchr/testify/require" @@ -13,6 +15,18 @@ import ( "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 @@ -36,3 +50,18 @@ func TestCreateClassifiesPublicationFailures(t *testing.T) { }) } } + +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 index 6b887d3c89..2023c06ce4 100644 --- a/go/core/internal/service/checkpoint/recovery.go +++ b/go/core/internal/service/checkpoint/recovery.go @@ -110,9 +110,17 @@ func (s *Service) currentOutcome(ctx context.Context, id, owner string) (*apiv1. } 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 := s.store.ReleaseCheckpointOperation(releaseCtx, id, executor); err != nil { + 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. 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 7da4acabbd..156d36bc86 100644 --- a/go/core/internal/service/checkpoint/service.go +++ b/go/core/internal/service/checkpoint/service.go @@ -32,6 +32,7 @@ type store interface { 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 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) @@ -124,8 +125,11 @@ func (s *Service) reserveCreate(ctx context.Context, sessionID, requestID, expec } checkpoint, snapshot, err := s.store.ReserveSessionCheckpoint(ctx, &apiv1.Checkpoint{Id: id, SessionId: sessionID, HeadTaskId: expectedHeadTaskID}, userID, requestID, executor) if err != nil { - // The INSERT may have committed before its reply was lost. - err = errors.Join(err, s.release(ctx, id, executor)) + // 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) + })) } if errors.Is(err, database.ErrCheckpointAdvanced) || errors.Is(err, database.ErrSnapshotPending) { reason := "KAGENT_CHECKPOINT_CONVERSATION_ADVANCED" diff --git a/go/core/internal/service/checkpoint/service_test.go b/go/core/internal/service/checkpoint/service_test.go index 2d5b4347e5..97fb38ed51 100644 --- a/go/core/internal/service/checkpoint/service_test.go +++ b/go/core/internal/service/checkpoint/service_test.go @@ -890,3 +890,7 @@ func (s *testStore) ClaimCheckpointOperation(ctx context.Context, id, owner stri func (*testStore) ReleaseCheckpointOperation(context.Context, string, uuid.UUID) error { return nil } + +func (*testStore) ReleaseCheckpointReservation(context.Context, string, string, uuid.UUID) error { + return nil +} diff --git a/ui/playwright/tests/chat/checkpoints.spec.ts b/ui/playwright/tests/chat/checkpoints.spec.ts index 4a1015a953..d79eba3fda 100644 --- a/ui/playwright/tests/chat/checkpoints.spec.ts +++ b/ui/playwright/tests/chat/checkpoints.spec.ts @@ -22,6 +22,58 @@ async function deleteOutcome(page: import("@playwright/test").Page, outcome?: st }, outcome); } +async function listUnavailable(page: import("@playwright/test").Page, unavailable: boolean) { + await page.evaluate((value) => { + const url = new URL(window.location.href); + if (value) url.searchParams.set("checkpointList", "unavailable"); + else url.searchParams.delete("checkpointList"); + window.history.replaceState({}, "", url); + }, unavailable); +} + +test.describe("chat: snapshot refresh failures", () => { + test.use({ expectedNoise: [/Could not (?:checkpoint the conversation|re-read the saved boundaries):/] }); + + test("chat: a failed creation followed by a list outage offers a repaired retry", async ({ page }) => { + await page.goto(`${agentChat(instances.ready)}?checkpoint=failed`); + await expect(dividers(page)).toHaveCount(1); + await sendAndAwaitTurn(page, "A retained failure whose first discovery read fails."); + await listUnavailable(page, true); + await page.getByTestId("chat-checkpoint").click(); + await expect(page.getByText("Could not load snapshots", { exact: true })).toBeVisible(); + await expect(page.getByTestId("chat-failed-snapshots")).toHaveCount(0); + await listUnavailable(page, false); + await page.getByRole("button", { name: "Retry", exact: true }).click(); + await expect(page.getByTestId("chat-failed-snapshots")).toBeVisible(); + await expect(page.getByText("Could not load snapshots", { exact: true })).toHaveCount(0); + await page.reload(); + await expect(page.getByTestId("chat-failed-snapshots")).toBeVisible(); + }); + + test("chat: deletion removes a failed entry even when the following list read fails", async ({ page }) => { + await page.goto(`${agentChat(instances.ready)}?checkpoint=failed`); + await expect(dividers(page)).toHaveCount(1); + await sendAndAwaitTurn(page, "Delete committed before discovery becomes unavailable."); + await page.getByTestId("chat-checkpoint").click(); + const notice = page.getByTestId("chat-failed-snapshots"); + await expect(notice).toBeVisible(); + await notice.locator('[data-testid^="chat-failed-snapshot-"]').click(); + await listUnavailable(page, true); + await page.getByTestId("snapshot-details-delete").click(); + await pressOnce(page.getByTestId("snapshot-details-delete-confirm")); + await expect(page.getByTestId("snapshot-details-body")).toHaveCount(0); + await expect(notice).toHaveCount(0); + await expect(page.getByText("Could not load snapshots", { exact: true })).toBeVisible(); + await listUnavailable(page, false); + await page.getByRole("button", { name: "Retry", exact: true }).click(); + await expect(page.getByText("Could not load snapshots", { exact: true })).toHaveCount(0); + await expect(notice).toHaveCount(0); + await page.reload(); + await expect(dividers(page)).toHaveCount(1); + await expect(notice).toHaveCount(0); + }); +}); + test.describe("chat: snapshot delete failures", () => { test.use({ expectedNoise: [/Could not (?:checkpoint the conversation|delete the checkpoint):/] }); diff --git a/ui/src/api/hooks/useApiResource.test.tsx b/ui/src/api/hooks/useApiResource.test.tsx new file mode 100644 index 0000000000..270a9ddddb --- /dev/null +++ b/ui/src/api/hooks/useApiResource.test.tsx @@ -0,0 +1,65 @@ +import { act, renderHook, waitFor } from "@testing-library/react"; +import type { ReactNode } from "react"; +import { SWRConfig } from "swr"; +import { expect, it, vi } from "vitest"; +import { ApiError } from "../ApiError"; +import { useApiResource } from "./useApiResource"; + +function wrapper({ children }: { children: ReactNode }) { + return new Map(), shouldRetryOnError: false, dedupingInterval: 0 }}>{children}; +} + +it("retains loaded data and exposes a failed refresh until a repaired retry succeeds", async () => { + const failure = new ApiError("List unavailable", { kind: "network", url: "test/list" }); + const fetcher = vi.fn<() => Promise>() + .mockResolvedValueOnce(["retained"]) + .mockRejectedValueOnce(failure) + .mockResolvedValueOnce([]); + const { result } = renderHook(() => useApiResource(["refresh-test"], fetcher), { wrapper }); + await waitFor(() => expect(result.current.data).toEqual(["retained"])); + await act(async () => { await expect(result.current.refresh()).rejects.toBe(failure); }); + expect(result.current.data).toEqual(["retained"]); + expect(result.current.error).toBe(failure); + expect(result.current.isEmpty).toBe(false); + await act(async () => { await result.current.refresh(); }); + expect(result.current.data).toEqual([]); + expect(result.current.error).toBeUndefined(); + expect(result.current.isEmpty).toBe(true); + expect(fetcher).toHaveBeenCalledTimes(3); +}); + +it("keeps refresh failures scoped to their resource key", async () => { + const failure = new ApiError("Old list unavailable", { kind: "network", url: "test/list" }); + const fetcher = vi.fn<() => Promise>() + .mockResolvedValueOnce(["first"]) + .mockRejectedValueOnce(failure) + .mockResolvedValueOnce(["second"]); + const { result, rerender } = renderHook(({ id }) => useApiResource(["scoped-refresh", id], fetcher), { initialProps: { id: "first" }, wrapper }); + await waitFor(() => expect(result.current.data).toEqual(["first"])); + await act(async () => { await expect(result.current.refresh()).rejects.toBe(failure); }); + rerender({ id: "second" }); + await waitFor(() => expect(result.current.data).toEqual(["second"])); + expect(result.current.error).toBeUndefined(); +}); + +it("does not read an idle resource when refresh is invoked", async () => { + const fetcher = vi.fn(async () => ["unexpected"]); + const { result } = renderHook(() => useApiResource(null, fetcher), { wrapper }); + await act(async () => { await result.current.refresh(); }); + expect(fetcher).not.toHaveBeenCalled(); + expect(result.current.data).toBeUndefined(); + expect(result.current.isLoading).toBe(false); +}); + +it("keeps an earlier refresh callback bound to its original resource after navigation", async () => { + const fetcher = vi.fn(async (id: string) => [id]); + const { result, rerender } = renderHook(({ id }) => useApiResource(["navigation-refresh", id], () => fetcher(id)), { initialProps: { id: "first" }, wrapper }); + await waitFor(() => expect(result.current.data).toEqual(["first"])); + const firstRefresh = result.current.refresh; + rerender({ id: "second" }); + await waitFor(() => expect(result.current.data).toEqual(["second"])); + await act(async () => { await firstRefresh(); }); + expect(fetcher.mock.calls).toEqual([["first"], ["second"]]); + expect(result.current.data).toEqual(["second"]); + expect(result.current.error).toBeUndefined(); +}); diff --git a/ui/src/api/hooks/useApiResource.ts b/ui/src/api/hooks/useApiResource.ts index d0aa2052bb..e737e15bf4 100644 --- a/ui/src/api/hooks/useApiResource.ts +++ b/ui/src/api/hooks/useApiResource.ts @@ -7,7 +7,7 @@ * empty state. This wraps it once so all four states are named. */ -import useSWR, { type SWRConfiguration } from "swr"; +import useSWR, { unstable_serialize, useSWRConfig, type SWRConfiguration } from "swr"; import { ApiError } from "../ApiError"; export interface ApiResource { @@ -34,7 +34,8 @@ export function useApiResource( fetcher: () => Promise, config?: SWRConfiguration, ): ApiResource { - const { data, error, isLoading, isValidating, mutate } = useSWR( + const { cache, mutate } = useSWRConfig(); + const { data, error, isLoading, isValidating } = useSWR( key, fetcher, config, @@ -51,18 +52,15 @@ export function useApiResource( /** * Refetches, and rejects when the refetch failed. * - * The bare `mutate()` this used to be does not reject: SWR captures a failed - * revalidation into `error` and resolves anyway, so a caller awaiting it cannot - * tell a refresh that worked from one that did not. That is fine for a - * background revalidation and wrong for a button someone pressed — it made the - * Refresh control report success over a page showing a load error. - * - * Handing `mutate` the promise instead makes it the new value: SWR records the - * rejection in `error` as before, so the page still shows what went wrong, and - * rethrows it here so the caller knows too. + * Revalidation updates SWR's shared error/loading state; mutating with a + * rejected fetch promise does not. SWR resolves even a failed revalidation, + * so read the completed result from this provider's cache and rethrow it. + * The render's `error` would still describe the previous read here. */ refresh: async () => { - await mutate(fetcher(), { revalidate: false }); + await mutate(key); + const failure = cache.get(unstable_serialize(key))?.error; + if (failure) throw failure; }, }; } diff --git a/ui/src/api/hooks/useCheckpoints.test.tsx b/ui/src/api/hooks/useCheckpoints.test.tsx new file mode 100644 index 0000000000..b31786a606 --- /dev/null +++ b/ui/src/api/hooks/useCheckpoints.test.tsx @@ -0,0 +1,44 @@ +import { act, renderHook, waitFor } from "@testing-library/react"; +import type { ReactNode } from "react"; +import { SWRConfig } from "swr"; +import { afterEach, expect, it } from "vitest"; +import { Code, ConnectError, createRouterTransport } from "@connectrpc/connect"; +import { CheckpointService, CheckpointState } from "@/generated/kagent/api/v1/checkpoints_pb"; +import { setApiTransport } from "../transport"; +import { useCheckpoints } from "./useCheckpoints"; + +function wrapper({ children }: { children: ReactNode }) { + return new Map(), shouldRetryOnError: false, dedupingInterval: 0 }}>{children}; +} + +afterEach(() => setApiTransport(undefined)); + +it("shares confirmed deletion with readers while keeping other conversations and read failures intact", async () => { + let unavailable = false; + const requests: string[] = []; + setApiTransport(createRouterTransport(({ service }) => service(CheckpointService, { + listCheckpoints(request) { + requests.push(request.sessionId); + if (unavailable) throw new ConnectError("List unavailable", Code.Unavailable); + return { checkpoints: [ + { id: `${request.sessionId}-ready`, state: CheckpointState.READY }, + { id: `${request.sessionId}-failed`, state: CheckpointState.FAILED }, + ] }; + }, + }))); + const { result } = renderHook(() => ({ first: useCheckpoints("first"), shared: useCheckpoints("first"), other: useCheckpoints("other") }), { wrapper }); + await waitFor(() => expect(result.current.first.data).toHaveLength(2)); + await waitFor(() => expect(result.current.other.data).toHaveLength(2)); + const reads = requests.length; + await act(async () => { await result.current.first.forget("first-failed"); }); + expect(result.current.first.data?.map(row => row.id)).toEqual(["first-ready"]); + expect(result.current.shared.data?.map(row => row.id)).toEqual(["first-ready"]); + expect(result.current.other.data?.map(row => row.id)).toEqual(["other-ready", "other-failed"]); + expect(requests).toHaveLength(reads); + unavailable = true; + await act(async () => { await expect(result.current.first.refresh()).rejects.toMatchObject({ code: "Unavailable" }); }); + expect(result.current.first.error?.code).toBe("Unavailable"); + expect(result.current.shared.error?.code).toBe("Unavailable"); + expect(result.current.first.data?.map(row => row.id)).toEqual(["first-ready"]); + expect(result.current.other.error).toBeUndefined(); +}); diff --git a/ui/src/api/hooks/useCheckpoints.ts b/ui/src/api/hooks/useCheckpoints.ts index dc0c652509..d6fd610951 100644 --- a/ui/src/api/hooks/useCheckpoints.ts +++ b/ui/src/api/hooks/useCheckpoints.ts @@ -1,6 +1,7 @@ import { apiClient } from "../client"; import type { Checkpoint } from "../domain/checkpoints"; import { type ApiResource, useApiResource } from "./useApiResource"; +import { useSWRConfig } from "swr"; /** * The turn boundaries saved against one conversation. @@ -8,9 +9,23 @@ import { type ApiResource, useApiResource } from "./useApiResource"; * Held back until the conversation is known, which `useApiResource` reports as idle * rather than as loading. */ -export function useCheckpoints(id: string | undefined): ApiResource { - return useApiResource( - id ? ["agentInstances.checkpoints.list", id] : null, +export function useCheckpoints(id: string | undefined): ApiResource & { + forget: (checkpointId: string) => Promise; +} { + const { mutate } = useSWRConfig(); + const key = id ? ["agentInstances.checkpoints.list", id] : null; + const resource = useApiResource( + key, () => apiClient.agentInstances.checkpoints.list(id ?? ""), ); + return { + ...resource, + // Confirmed deletion is already authoritative even if the next list fails. + // Update the shared record independently of that subsequent read. + forget: async (checkpointId) => { + if (key) { + await mutate(key, (rows) => rows?.filter((row) => row.id !== checkpointId), { revalidate: false }); + } + }, + }; } diff --git a/ui/src/mocks/transport.ts b/ui/src/mocks/transport.ts index cc413202e7..6da1b33087 100644 --- a/ui/src/mocks/transport.ts +++ b/ui/src/mocks/transport.ts @@ -972,6 +972,9 @@ on(CheckpointService.method.deleteCheckpoint, (input) => { on(CheckpointService.method.listCheckpoints, (input, call) => { const instance = instanceFor(requireInstanceId(input.sessionId), call); + if (new URLSearchParams(window.location.search).get("checkpointList") === "unavailable") { + throw new ConnectError("Snapshot list unavailable", Code.Unavailable); + } return { checkpoints: readCheckpoints(instance.id).map(checkpointMessage), page: {} }; }); diff --git a/ui/src/pages/AgentChatPage.tsx b/ui/src/pages/AgentChatPage.tsx index b0fec3ef5b..6f805f7dad 100644 --- a/ui/src/pages/AgentChatPage.tsx +++ b/ui/src/pages/AgentChatPage.tsx @@ -259,6 +259,7 @@ export function AgentChatPage() { async (checkpointId: string) => { try { await apiClient.agentInstances.checkpoints.remove(checkpointId); + await checkpoints.forget(checkpointId); setSavedHere((current) => ({ conversation: current.conversation, marks: new Map([...current.marks].filter(([, saved]) => saved !== checkpointId)), diff --git a/ui/src/pages/ScheduledRunsPage.test.tsx b/ui/src/pages/ScheduledRunsPage.test.tsx index 0442091dcf..44150ee6e3 100644 --- a/ui/src/pages/ScheduledRunsPage.test.tsx +++ b/ui/src/pages/ScheduledRunsPage.test.tsx @@ -54,18 +54,24 @@ it("reuses a manual run request ID after a lost response, then gives the next fi }))); renderSchedule(); const run = await screen.findByRole("button", { name: "Run" }); + // antd can clear disabled before its loading state stops ignoring clicks. + const waitForRun = () => waitFor(() => { + expect(run).toBeEnabled(); + expect(run).not.toHaveClass("ant-btn-loading"); + }); + await waitForRun(); fireEvent.click(run); await screen.findByText(/Response lost/); - await waitFor(() => expect(run).toBeEnabled()); + await waitForRun(); fireEvent.click(run); await screen.findByText(/Execution queued/); - await waitFor(() => expect(run).toBeEnabled()); + await waitForRun(); fireEvent.click(run); await waitFor(() => expect(requestIds).toHaveLength(3)); expect(requestIds[0]).toBeTruthy(); expect(requestIds[1]).toBe(requestIds[0]); expect(requestIds[2]).not.toBe(requestIds[0]); - await waitFor(() => expect(run).toBeEnabled()); + await waitForRun(); }); it("keeps unsaved edits after a conflict and preserves the exact timeout on the wire", async () => { From cf917206d1a8982dbedc7a522eba6a2b19ff3cc9 Mon Sep 17 00:00:00 2001 From: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> Date: Tue, 10 Nov 2026 11:29:56 +0800 Subject: [PATCH 18/20] fix(checkpoint): discard obsolete refresh failures Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> --- .../persistence-checkpoints-and-forks.md | 2 + .../database/checkpoint_operations_test.go | 34 +++++++- ui/playwright/tests/chat/checkpoints.spec.ts | 43 ++++++++++ ui/src/api/hooks/useApiResource.test.tsx | 83 ++++++++++++++++++- ui/src/api/hooks/useApiResource.ts | 46 +++++++++- ui/src/api/hooks/useCheckpoints.test.tsx | 23 +++++ ui/src/testUtils/deferred.ts | 7 ++ 7 files changed, 230 insertions(+), 8 deletions(-) create mode 100644 ui/src/testUtils/deferred.ts diff --git a/docs/architecture/persistence-checkpoints-and-forks.md b/docs/architecture/persistence-checkpoints-and-forks.md index ea08c97a4c..c16b6682af 100644 --- a/docs/architecture/persistence-checkpoints-and-forks.md +++ b/docs/architecture/persistence-checkpoints-and-forks.md @@ -160,6 +160,8 @@ keeps its details open for retry. Confirmed deletion removes the cached record before rereading the list, so a subsequent read outage cannot leave a deleted FAILED entry on screen. Failed refreshes retain the last data and expose an error with a retry action. +An older read failure cannot replace a newer read's outcome or undo a confirmed +cache update; only the current read may publish an error or schedule its retry. ### Proposed initial upgrade boundary: pending maintainer acceptance diff --git a/go/core/internal/database/checkpoint_operations_test.go b/go/core/internal/database/checkpoint_operations_test.go index 17971b7f50..cf06fcef7f 100644 --- a/go/core/internal/database/checkpoint_operations_test.go +++ b/go/core/internal/database/checkpoint_operations_test.go @@ -3,6 +3,7 @@ package database import ( "context" "encoding/binary" + "errors" "strings" "sync" "testing" @@ -98,12 +99,18 @@ func reserveRecoveryCheckpoint(t *testing.T, client *Client) (*apiv1.Checkpoint, } func TestCheckpointClaimsFenceCompletionAndDeletion(t *testing.T) { - client := NewClient(setupTestDB(t), "public") + 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(), 80*time.Millisecond) + ctx, cancel := context.WithTimeout(t.Context(), 500*time.Millisecond) defer cancel() - _, err := client.ClaimCheckpointOperation(ctx, checkpoint.Id, "alice", first) + _, err = client.ClaimCheckpointOperation(ctx, checkpoint.Id, "alice", first) require.NoError(t, err) nextCtx, stop := context.WithTimeout(t.Context(), RuntimeOperationTimeout) defer stop() @@ -116,7 +123,13 @@ func TestCheckpointClaimsFenceCompletionAndDeletion(t *testing.T) { require.Empty(t, page, "live claims are not sweep work") <-ctx.Done() second := uuid.New() - op, err := client.ClaimCheckpointOperation(nextCtx, checkpoint.Id, "alice", second) + // 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", "") @@ -136,6 +149,19 @@ func TestCheckpointClaimsFenceCompletionAndDeletion(t *testing.T) { 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) diff --git a/ui/playwright/tests/chat/checkpoints.spec.ts b/ui/playwright/tests/chat/checkpoints.spec.ts index d79eba3fda..f7c02c7867 100644 --- a/ui/playwright/tests/chat/checkpoints.spec.ts +++ b/ui/playwright/tests/chat/checkpoints.spec.ts @@ -34,6 +34,49 @@ async function listUnavailable(page: import("@playwright/test").Page, unavailabl test.describe("chat: snapshot refresh failures", () => { test.use({ expectedNoise: [/Could not (?:checkpoint the conversation|re-read the saved boundaries):/] }); + test("chat: a repaired snapshot list ignores an older failed refresh", async ({ page }) => { + await page.goto(`${agentChat(instances.ready)}?checkpoint=failed`); + await expect(dividers(page)).toHaveCount(1); + await sendAndAwaitTurn(page, "Overlapping discovery requests after storage repair."); + await listUnavailable(page, true); + await page.getByTestId("chat-checkpoint").click(); + await expect(page.getByText("Could not load snapshots", { exact: true })).toBeVisible(); + await page.evaluate(() => { + const url = new URL(window.location.href); + url.searchParams.delete("checkpointList"); + url.searchParams.set("mock", "error"); + window.history.replaceState({}, "", url); + // Delay just this next fixture RPC, whose error scenario is captured + // before sleeping. The repaired retry must finish ahead of its failure. + const original = window.setTimeout; + window.setTimeout = Object.assign((handler: TimerHandler, delay?: number, ...args: unknown[]): number => { + if (delay === 450 && typeof handler === "function") { + window.setTimeout = original; + document.body.dataset.delayedSnapshotRead = "started"; + return original(() => { + handler(...args); + document.body.dataset.delayedSnapshotRead = "completed"; + }, 1800); + } + return original(handler, delay, ...args); + }, original); + }); + await page.getByRole("button", { name: "Retry", exact: true }).click(); + await expect(page.locator("body")).toHaveAttribute("data-delayed-snapshot-read", "started"); + await page.evaluate(() => { + const url = new URL(window.location.href); + url.searchParams.set("mock", "ok"); + window.history.replaceState({}, "", url); + }); + await page.getByRole("button", { name: "Retry", exact: true }).click(); + await expect(page.getByTestId("chat-failed-snapshots")).toBeVisible(); + await expect(page.getByText("Could not load snapshots", { exact: true })).toHaveCount(0); + await expect(page.locator("body")).toHaveAttribute("data-delayed-snapshot-read", "completed"); + // Let the old RPC rejection and React's error render reach the page. + await page.evaluate(() => new Promise(resolve => requestAnimationFrame(() => requestAnimationFrame(() => resolve())))); + await expect(page.getByText("Could not load snapshots", { exact: true })).toHaveCount(0); + }); + test("chat: a failed creation followed by a list outage offers a repaired retry", async ({ page }) => { await page.goto(`${agentChat(instances.ready)}?checkpoint=failed`); await expect(dividers(page)).toHaveCount(1); diff --git a/ui/src/api/hooks/useApiResource.test.tsx b/ui/src/api/hooks/useApiResource.test.tsx index 270a9ddddb..c8167fc87c 100644 --- a/ui/src/api/hooks/useApiResource.test.tsx +++ b/ui/src/api/hooks/useApiResource.test.tsx @@ -1,9 +1,10 @@ import { act, renderHook, waitFor } from "@testing-library/react"; import type { ReactNode } from "react"; -import { SWRConfig } from "swr"; +import { SWRConfig, useSWRConfig } from "swr"; import { expect, it, vi } from "vitest"; import { ApiError } from "../ApiError"; import { useApiResource } from "./useApiResource"; +import { deferred } from "@/testUtils/deferred"; function wrapper({ children }: { children: ReactNode }) { return new Map(), shouldRetryOnError: false, dedupingInterval: 0 }}>{children}; @@ -63,3 +64,83 @@ it("keeps an earlier refresh callback bound to its original resource after navig expect(result.current.data).toEqual(["second"]); expect(result.current.error).toBeUndefined(); }); + +it.each(["refresh", "background"] as const)("ignores an older failure after a newer %s succeeds for all readers", async (mode) => { + const older = deferred(); + const failure = new ApiError("Earlier read failed", { kind: "network", url: "test/list" }); + const fetcher = vi.fn<() => Promise>() + .mockResolvedValueOnce(["initial"]) + .mockReturnValueOnce(older.promise) + .mockResolvedValueOnce(["repaired"]); + const { result } = renderHook(() => ({ + first: useApiResource(["ordered", mode], fetcher), + shared: useApiResource(["ordered", mode], fetcher), + revalidate: useSWRConfig().mutate, + }), { wrapper }); + await waitFor(() => expect(result.current.first.data).toEqual(["initial"])); + let earlier!: Promise; + act(() => { earlier = result.current.first.refresh().catch(error => error); }); + await waitFor(() => expect(fetcher).toHaveBeenCalledTimes(2)); + await act(async () => { + if (mode === "refresh") await result.current.first.refresh(); + else await result.current.revalidate(["ordered", mode]); + }); + expect(result.current.first.data).toEqual(["repaired"]); + await act(async () => { older.reject(failure); await earlier; }); + expect(result.current.first.error).toBeUndefined(); + expect(result.current.shared.error).toBeUndefined(); + expect(result.current.first.data).toEqual(["repaired"]); +}); + +it("preserves the newer failure when an older read fails afterward", async () => { + const older = deferred(); + const failure = new ApiError("Current read failed", { kind: "network", url: "test/list" }); + const fetcher = vi.fn<() => Promise>() + .mockResolvedValueOnce(["retained"]) + .mockReturnValueOnce(older.promise) + .mockRejectedValueOnce(failure); + const { result } = renderHook(() => useApiResource(["ordered-failure"], fetcher), { wrapper }); + await waitFor(() => expect(result.current.data).toEqual(["retained"])); + let earlier!: Promise; + act(() => { earlier = result.current.refresh().catch(error => error); }); + await waitFor(() => expect(fetcher).toHaveBeenCalledTimes(2)); + await act(async () => { await expect(result.current.refresh()).rejects.toBe(failure); }); + await act(async () => { older.reject(new ApiError("Obsolete failure", { kind: "network", url: "test/list" })); await earlier; }); + expect(result.current.error).toBe(failure); + expect(result.current.data).toEqual(["retained"]); +}); + +it("keeps concurrent reads with the same key isolated across cache providers", async () => { + const pending = deferred(); + const failure = new ApiError("First provider failed", { kind: "network", url: "test/list" }); + const first = renderHook(() => useApiResource(["provider-isolation"], () => pending.promise), { wrapper }); + const second = renderHook(() => useApiResource(["provider-isolation"], async () => ["healthy"]), { wrapper }); + await waitFor(() => expect(second.result.current.data).toEqual(["healthy"])); + await act(async () => { pending.reject(failure); }); + await waitFor(() => expect(first.result.current.error).toBe(failure)); + expect(second.result.current.data).toEqual(["healthy"]); + expect(second.result.current.error).toBeUndefined(); +}); + +it("reports a later current failure normally after overlapping reads have settled", async () => { + const older = deferred(); + const failure = new ApiError("New current failure", { kind: "network", url: "test/list" }); + const onError = vi.fn(); + const fetcher = vi.fn<() => Promise>() + .mockResolvedValueOnce(["initial"]) + .mockReturnValueOnce(older.promise) + .mockResolvedValueOnce(["repaired"]) + .mockRejectedValueOnce(failure); + const { result } = renderHook(() => useApiResource(["after-overlap"], fetcher, { onError }), { wrapper }); + await waitFor(() => expect(result.current.data).toEqual(["initial"])); + let earlier!: Promise; + act(() => { earlier = result.current.refresh().catch(error => error); }); + await waitFor(() => expect(fetcher).toHaveBeenCalledTimes(2)); + await act(async () => { await result.current.refresh(); }); + await act(async () => { older.reject(new ApiError("Obsolete failure", { kind: "network", url: "test/list" })); await earlier; }); + expect(onError).not.toHaveBeenCalled(); + await act(async () => { await expect(result.current.refresh()).rejects.toBe(failure); }); + expect(result.current.error).toBe(failure); + expect(onError).toHaveBeenCalledTimes(1); + expect(onError.mock.calls[0][0]).toBe(failure); +}); diff --git a/ui/src/api/hooks/useApiResource.ts b/ui/src/api/hooks/useApiResource.ts index e737e15bf4..f60d7df8c6 100644 --- a/ui/src/api/hooks/useApiResource.ts +++ b/ui/src/api/hooks/useApiResource.ts @@ -7,7 +7,7 @@ * empty state. This wraps it once so all four states are named. */ -import useSWR, { unstable_serialize, useSWRConfig, type SWRConfiguration } from "swr"; +import useSWR, { unstable_serialize, useSWRConfig, type Cache, type SWRConfiguration } from "swr"; import { ApiError } from "../ApiError"; export interface ApiResource { @@ -35,9 +35,10 @@ export function useApiResource( config?: SWRConfiguration, ): ApiResource { const { cache, mutate } = useSWRConfig(); + const resourceKey = unstable_serialize(key); const { data, error, isLoading, isValidating } = useSWR( key, - fetcher, + () => readResource(cache, resourceKey, fetcher), config, ); @@ -59,12 +60,51 @@ export function useApiResource( */ refresh: async () => { await mutate(key); - const failure = cache.get(unstable_serialize(key))?.error; + const failure = cache.get(resourceKey)?.error; if (failure) throw failure; }, }; } +interface ResourceReads { + latest: Promise; + pending: number; +} + +// Share ordering across readers and automatic revalidation on the same provider. +// Entries exist only while reads are in flight, including superseded requests. +const readsByCache = new WeakMap>(); + +function readResource(cache: Cache, key: string, fetcher: () => Promise): Promise { + let reads = readsByCache.get(cache); + if (!reads) { + reads = new Map(); + readsByCache.set(cache, reads); + } + const active = reads; + const dataBefore = cache.get(key)?.data; + const request: Promise = Promise.resolve().then(fetcher).catch((cause: unknown) => { + if (active.get(key)?.latest !== request || cache.get(key)?.data !== dataBefore) { + // SWR fences obsolete successes, but can still publish obsolete errors. + // Resolve through the cache so its request/mutation fence discards this + // result too. Only the current read may publish a failure or retry it. + return cache.get(key)?.data as T; + } + throw cause; + }); + const previous = active.get(key); + active.set(key, { latest: request, pending: (previous?.pending ?? 0) + 1 }); + const finish = () => { + const current = active.get(key); + if (current && --current.pending === 0) { + active.delete(key); + if (active.size === 0) readsByCache.delete(cache); + } + }; + void request.then(finish, finish); + return request; +} + function isEmptyResult(data: unknown): boolean { if (data === undefined || data === null) return false; // Not loaded is not empty. if (Array.isArray(data)) return data.length === 0; diff --git a/ui/src/api/hooks/useCheckpoints.test.tsx b/ui/src/api/hooks/useCheckpoints.test.tsx index b31786a606..fe0937f07b 100644 --- a/ui/src/api/hooks/useCheckpoints.test.tsx +++ b/ui/src/api/hooks/useCheckpoints.test.tsx @@ -6,6 +6,7 @@ import { Code, ConnectError, createRouterTransport } from "@connectrpc/connect"; import { CheckpointService, CheckpointState } from "@/generated/kagent/api/v1/checkpoints_pb"; import { setApiTransport } from "../transport"; import { useCheckpoints } from "./useCheckpoints"; +import { deferred } from "@/testUtils/deferred"; function wrapper({ children }: { children: ReactNode }) { return new Map(), shouldRetryOnError: false, dedupingInterval: 0 }}>{children}; @@ -42,3 +43,25 @@ it("shares confirmed deletion with readers while keeping other conversations and expect(result.current.first.data?.map(row => row.id)).toEqual(["first-ready"]); expect(result.current.other.error).toBeUndefined(); }); + +it.each(["success", "failure"] as const)("does not restore a deleted record or obsolete error from an earlier %s", async (outcome) => { + const pending = deferred<{ checkpoints: Array<{ id: string; state: CheckpointState }> }>(); + const records = [{ id: "ready", state: CheckpointState.READY }, { id: "failed", state: CheckpointState.FAILED }]; + let reads = 0; + setApiTransport(createRouterTransport(({ service }) => service(CheckpointService, { + listCheckpoints() { return ++reads === 2 ? pending.promise : { checkpoints: records }; }, + }))); + const { result } = renderHook(() => useCheckpoints("conversation"), { wrapper }); + await waitFor(() => expect(result.current.data).toHaveLength(2)); + let earlier!: Promise; + act(() => { earlier = result.current.refresh().catch(error => error); }); + await waitFor(() => expect(reads).toBe(2)); + await act(async () => { await result.current.forget("failed"); }); + await act(async () => { + if (outcome === "success") pending.resolve({ checkpoints: records }); + else pending.reject(new ConnectError("Obsolete read failure", Code.Unavailable)); + await earlier; + }); + expect(result.current.data?.map(row => row.id)).toEqual(["ready"]); + expect(result.current.error).toBeUndefined(); +}); diff --git a/ui/src/testUtils/deferred.ts b/ui/src/testUtils/deferred.ts new file mode 100644 index 0000000000..d09f3e7110 --- /dev/null +++ b/ui/src/testUtils/deferred.ts @@ -0,0 +1,7 @@ +/** Control completion order at an asynchronous test boundary. */ +export function deferred() { + let resolve!: (value: T) => void; + let reject!: (error: unknown) => void; + const promise = new Promise((yes, no) => { resolve = yes; reject = no; }); + return { promise, resolve, reject }; +} From 80f1bb797c869f48e628e987c803394992ed3e93 Mon Sep 17 00:00:00 2001 From: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> Date: Tue, 10 Nov 2026 16:01:36 +0800 Subject: [PATCH 19/20] fix(checkpoint): collect failed tags internally and narrow recovery scope Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> --- .github/workflows/ci.yaml | 3 - .../persistence-checkpoints-and-forks.md | 188 ++---- .../database/checkpoint_operations.go | 44 +- .../database/checkpoint_operations_test.go | 98 +++ go/core/internal/database/checkpoints.go | 14 +- .../internal/database/client_session_test.go | 8 +- .../grpcserver/checkpoint_failure_test.go | 119 ---- .../internal/mcp/checkpoint_errors_test.go | 76 ++- go/core/internal/mcp/checkpoints.go | 42 +- .../service/checkpoint/completion_test.go | 8 +- .../service/checkpoint/discovery_test.go | 131 ---- .../service/checkpoint/failed_cleanup_test.go | 134 +++++ .../service/checkpoint/migration_test.go | 2 +- ...iew_test.go => operation_recovery_test.go} | 66 +- .../internal/service/checkpoint/recovery.go | 39 +- .../checkpoint/recovery_objectstore_test.go | 55 -- .../checkpoint/recovery_substrate_test.go | 564 ------------------ .../service/checkpoint/recovery_test.go | 24 - .../internal/service/checkpoint/service.go | 91 ++- .../service/checkpoint/service_test.go | 5 + go/core/test/e2e/README.md | 23 - go/core/test/e2e/checkpoint-recovery.sh | 65 -- go/core/test/upgrade/checkpoint_drain_test.go | 145 ----- .../test/upgrade/checkpoint_rollback_test.go | 102 ---- go/core/test/upgrade/rolling_upgrade_test.go | 2 - go/core/test/upgrade/upgrade_test.go | 24 +- ui/playwright/tests/chat/checkpoints.spec.ts | 215 ------- ui/src/api/domain/checkpoints.ts | 6 +- ui/src/api/grpc/checkpoints.test.ts | 91 +-- ui/src/api/grpc/operations.ts | 74 +-- ui/src/api/hooks/useApiResource.test.tsx | 146 ----- ui/src/api/hooks/useApiResource.ts | 64 +- ui/src/api/hooks/useCheckpoints.test.tsx | 67 --- ui/src/api/hooks/useCheckpoints.ts | 21 +- ui/src/components/chat/ChatTranscript.tsx | 2 +- ui/src/components/chat/CheckpointDivider.tsx | 11 +- .../components/chat/SnapshotDetailsModal.tsx | 15 +- ui/src/mocks/mockBackend.test.ts | 17 - ui/src/mocks/state.ts | 2 - ui/src/mocks/transport.ts | 32 +- ui/src/pages/AgentChatPage.tsx | 56 +- ui/src/pages/ScheduledRunsPage.test.tsx | 12 +- ui/src/testUtils/deferred.ts | 7 - 43 files changed, 566 insertions(+), 2344 deletions(-) delete mode 100644 go/core/internal/grpcserver/checkpoint_failure_test.go delete mode 100644 go/core/internal/service/checkpoint/discovery_test.go create mode 100644 go/core/internal/service/checkpoint/failed_cleanup_test.go rename go/core/internal/service/checkpoint/{review_test.go => operation_recovery_test.go} (88%) delete mode 100644 go/core/internal/service/checkpoint/recovery_objectstore_test.go delete mode 100644 go/core/internal/service/checkpoint/recovery_substrate_test.go delete mode 100644 go/core/test/e2e/checkpoint-recovery.sh delete mode 100644 go/core/test/upgrade/checkpoint_drain_test.go delete mode 100644 go/core/test/upgrade/checkpoint_rollback_test.go delete mode 100644 ui/src/api/hooks/useApiResource.test.tsx delete mode 100644 ui/src/api/hooks/useCheckpoints.test.tsx delete mode 100644 ui/src/testUtils/deferred.ts diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 277651807c..a7f7d861bb 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -265,9 +265,6 @@ jobs: sleep 5 done test "$READY" = True - # Exercise recovery with real Substrate and PostgreSQL, including lost - # tag responses and deletion of failed partial copies. - bash core/test/e2e/checkpoint-recovery.sh export KAGENT_E2E_API_URL="http://$(kubectl get svc -n kagent kagent-controller -o jsonpath='{.status.loadBalancer.ingress[0].ip}'):8083" echo "KAGENT_E2E_API_URL: $KAGENT_E2E_API_URL" make core/bin/kagent-local diff --git a/docs/architecture/persistence-checkpoints-and-forks.md b/docs/architecture/persistence-checkpoints-and-forks.md index c16b6682af..5954c30104 100644 --- a/docs/architecture/persistence-checkpoints-and-forks.md +++ b/docs/architecture/persistence-checkpoints-and-forks.md @@ -59,141 +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. - -The request inserts the `CREATING` reservation and its bounded execution claim -in one transaction, copies the tag inline, then conditionally publishes `READY` -or `FAILED` while clearing the claim in the same update. Database transactions -never span Substrate calls. The normal create path writes only the reservation -and final outcome; it does not acquire a second claim or poll for a worker. - -Each attempt shares a request deadline of at most two minutes with its claim. -Retries with the same owner/request ID reuse the existing reservation. They -acquire expired work and execute the same function, or return the current outcome -when another request owns it. A live operation returns retryable `Aborted`. -If completion loses execution authority, the caller only observes the current -outcome: READY is returned, pending work returns `Aborted`, and a completed -deletion succeeds. Create returns `NotFound` if the checkpoint was deleted. -The observation never repeats a backend mutation or overwrites the new executor. -Create observers skip obsolete claim releases, so a release failure cannot mask -READY or the translation of absence to `NotFound`. A definitive acquisition -conflict also needs no release. Other acquisition errors may hide a committed -claim; those attempts still release their specific executor before returning. -An uncertain initial reservation also waits for its source Session transaction -before looking for the inserted row: locking the checkpoint alone cannot see an -uncommitted INSERT. This coordination is confined to reservation errors and keeps -the normal two-write path unchanged. Retirement remains bounded by five seconds; -if storage prevents it, claim expiry still permits later recovery. -Backend errors describe the current attempt, without persisting an error that a -later repaired request could accidentally replay. Public errors exclude backend -configuration. A canceled or uncertain attempt retains `CREATING` and expires its -claim immediately while retaining the executor identity. - -A small leader-elected sweep runs at startup and waits thirty seconds after each -completed sweep before starting another. It reads only expired `CREATING` and -`DELETING` work and uses the same execution path as requests. Claims fence publication and deletion across replicas, restarts, and -paused executors; a delayed release cannot revoke a replacement claim. Slow -operations can delay later background work; each attempt remains bounded by two -minutes. The post-sweep interval paces background retries. Live requests execute -independently of the sweep, without a queue, retry map, or wakeup protocol. - -A pending tag is not evidence of abandonment. Substrate v1.0.0 acquires its Actor -and Tag leases before copying and finalizing in `CreateTag`. Recovery calls -`CreateTag` before interpreting `AlreadyExists`, which means the name is taken, -not that its copy completed. Interpreting a retained partial copy as stopped -depends on the backend serialization contract below. -A busy backend lease returns `Aborted`; that leaves creation pending without -inspecting or cleaning the still-live tag. A later retry can adopt the completed -copy under the same request ID. - -Kagent must still verify the tag's identity, template lineage, completed storage, -and the source Actor's snapshot and UID. Transport errors retain admission until -a subsequent serialized attempt establishes an outcome. - -### Required backend contract: pending Substrate confirmation - -The recovery path requires previous external writes to have stopped before a -new serialized attempt can classify a retained tag as failed or delete it. -Besides `AlreadyExists`, the service interprets `CreateTag` replies of -`FailedPrecondition` and `NotFound` as definitive enough to inspect the tag. -Those paths, and UID-conditional deletion, require the same safety guarantee. - -The live Substrate regression verifies contention while the original RPC still -holds valid leases: the caller disconnects without delivering cancellation to -that RPC, another replica receives `Aborted`, and a later retry adopts its UID. -It does not verify cancellation delivered to Substrate, failed lease renewal, -process failure, or delayed object-store writes after lease loss. Lease tokens -protect lease ownership; neither a kagent claim timeout nor a successful lease -acquisition alone proves that earlier external writes stopped. The stronger -contract needs confirmation and fault coverage in Substrate. If it does not -hold, fix the owning backend or retain admission until the outcome is safe; -do not infer safety from elapsed time. This remains a release acceptance item. - -After a copy has stopped, failed verification marks the checkpoint `FAILED` and -releases source admission. Its exact tag UID remains retained for explicit -UID-conditional deletion, including when tag storage is still `PENDING` and has -no completed URI. Failed checkpoints cannot be forked. Their history and revision -references remain pinned until deletion succeeds. An unrelated or replaced tag -is never adopted or deleted. - -Owner-scoped Get and List expose both `READY` and `FAILED` checkpoints, including -after source deletion. `CREATING` and `DELETING` remain hidden, and fork reads and -transactions still require `READY`. A terminal failed Create returns -`FailedPrecondition` with `ErrorInfo.reason = KAGENT_CHECKPOINT_FAILED` and the -checkpoint UUID in `ResourceInfo.resource_name`. Repeating that request returns -the same terminal identity without copying again. The owner can discover and -delete the failed checkpoint before creating a new request. Public failed -checkpoint responses use a generic service-owned failure message, including for -legacy rows that retained raw backend errors; the stored operation is not rewritten. - -MCP checkpoint tools use the same classified public errors, excluding storage -causes and wrapped cancellation details. Backend status is retained as -`kagent.dev/grpc-code`; terminal creation failures also retain -`kagent.dev/error-reason` and `kagent.dev/checkpoint-id` for cleanup. The Web UI -lists failed snapshots separately from saved transcript boundaries, refreshes -after a failed creation reply, and reads subsequent list pages. Failed records -can be inspected and deleted, but cannot be forked or renamed. A failed deletion -keeps its details open for retry. -Confirmed deletion removes the cached record before rereading the list, so a -subsequent read outage cannot leave a deleted FAILED entry on screen. Failed -refreshes retain the last data and expose an error with a retry action. -An older read failure cannot replace a newer read's outcome or undo a confirmed -cache update; only the current read may publish an error or schedule its retry. - -### Proposed initial upgrade boundary: pending maintainer acceptance - -The proposed first upgrade from a controller without checkpoint execution claims -quiesces checkpoint creation/deletion, stops every old controller, and waits for -those processes to exit before starting the new version and resuming requests. -This interrupts the controller API. Merely marking a Pod unready is insufficient. -Legacy reservations have NULL claims; recovery waits until their UUIDv7 creation -age exceeds the two-minute operation timeout. This grace is an additional delay, -not proof that an old writer exited: old code has no bounded overall operation -and can wait in a reservation transaction. Do not mix old checkpoint writers with -new recovery. Claim-aware controllers can coordinate through the database during -later rolling upgrades; this does not establish compatibility of every other -release change. -Before downgrading to a controller without execution claims, finish pending -checkpoint operations, delete failed -checkpoints that retain tag UIDs, and drain all current writers. Legacy deletion -requires a completed URI and cannot collect a retained partial failed tag. -The upgrade suite keeps its old-code/new-schema checks, drains every legacy Pod -before starting the new controller, and repeats that boundary before application -rollback. A separate process-lifetime test holds a terminating legacy Pod to -verify that scaling to zero or losing readiness alone does not open the boundary. -The PostgreSQL migration regression seeds a populated version 3 checkpoint before -claim columns exist, upgrades it, and verifies recovery preserves its request and -history boundary. The rollback guard is also exercised with a retained failed -partial-tag fixture. These are storage fixtures, not previous-release binaries. - -The production installer still performs normal Helm upgrades; only the upgrade -test explicitly drains checkpoint writers. The proposed maintenance window is -not an accepted release policy or an enforced installer sequence. Before release, -maintainers must choose that boundary with public upgrade/rollback instructions, -or a staged compatible rollout if API interruption is unacceptable. +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 @@ -207,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 is definitively verified or failed. A lost response can reuse a completed Tag; -a stopped incomplete copy becomes FAILED and retains its Tag UID for deletion. 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. @@ -224,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 after the serialized copy has stopped and verification has failed. -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 the checkpoint worker 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 @@ -296,12 +191,11 @@ snapshot forks. The Agent endpoint and globally unique task ID identify a task. Reads, writes, cancellation, subscriptions, authorization, and dispatch remain scoped to the resolved session. Forks have independent public identities. -## Schema migrations +## Unreleased schema -Merged migration files are immutable. Core changes use new transactional Goose -migrations, including the two nullable checkpoint claim columns. -Development databases created from older, rewritten alpha baselines still need -to be recreated; new migrations do not bridge those historical schemas. +Until release, core schema changes are folded into `000001_initial.sql`. +Recreate development databases when that baseline changes; there is no upgrade +path from earlier development schemas. The Down migration removes the core schema. Deploy the controller and clients together: clients must use `Session.context_id` or omit the context and let the routed gateway resolve it. diff --git a/go/core/internal/database/checkpoint_operations.go b/go/core/internal/database/checkpoint_operations.go index c217b0c5d8..b6ae54de99 100644 --- a/go/core/internal/database/checkpoint_operations.go +++ b/go/core/internal/database/checkpoint_operations.go @@ -22,9 +22,9 @@ type CheckpointOperation struct { } // ClaimCheckpointOperation grants one bounded attempt exclusive execution of an -// owner's pending checkpoint and returns its current phase, not a caller's earlier -// observation. The caller must establish its attempt deadline before claiming and -// use that context for every backend call. Missing, terminal, or actively claimed +// 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() @@ -47,8 +47,9 @@ func (c *Client) ClaimCheckpointOperation(ctx context.Context, id, owner string, 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') - AND ((executor_id IS NULL AND $5::boolean) OR executor_expires_at <= clock_timestamp()) + 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 @@ -124,8 +125,8 @@ func (c *Client) GetCheckpointOperation(ctx context.Context, id, owner string) ( return checkpointOperation(row) } -// ListPendingCheckpoints returns a bounded ID-ordered page of durable creation -// and deletion work across owners. Only the controller's checkpoint worker uses +// 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 @@ -137,8 +138,8 @@ func (c *Client) ListPendingCheckpoints(ctx context.Context, afterID string, lim snapshot_atespace, snapshot_uri, snapshot_fidelity, tag_uid, state, data, source_history_id, prepared_revision FROM session_checkpoint - WHERE state IN ('CREATING', 'DELETING') - AND (executor_expires_at <= statement_timestamp() OR (executor_id IS NULL AND id < $3::uuid)) + 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())) @@ -156,6 +157,27 @@ func (c *Client) ListPendingCheckpoints(ctx context.Context, afterID string, lim 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 { @@ -165,8 +187,8 @@ func checkpointOperation(row sessionCheckpointRow) (*CheckpointOperation, error) Snapshot: checkpointSnapshot(row), TagUID: row.TagUID}, nil } -// Legacy rows have no claim. The UUID grace avoids taking freshly reserved work -// during rollout, but is not a fence: the first upgrade must drain old writers. +// 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) diff --git a/go/core/internal/database/checkpoint_operations_test.go b/go/core/internal/database/checkpoint_operations_test.go index cf06fcef7f..6186ea6971 100644 --- a/go/core/internal/database/checkpoint_operations_test.go +++ b/go/core/internal/database/checkpoint_operations_test.go @@ -79,6 +79,104 @@ func TestCheckpointReleaseWaitsForClaimCommit(t *testing.T) { 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() diff --git a/go/core/internal/database/checkpoints.go b/go/core/internal/database/checkpoints.go index 3176de5a51..a9077fbce3 100644 --- a/go/core/internal/database/checkpoints.go +++ b/go/core/internal/database/checkpoints.go @@ -372,15 +372,14 @@ func (c *Client) FinalizeClaimedSessionCheckpoint(ctx context.Context, checkpoin return result, nil } -// GetSessionCheckpoint returns an owned READY or FAILED checkpoint so the owner -// can discover a failed copy for cleanup. Missing checkpoints, other owners, -// CREATING, and DELETING return ErrNotFound. Fork readers remain READY-only. +// 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 := 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 IN ('READY', 'FAILED') + 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)) @@ -446,9 +445,8 @@ func checkpointSnapshot(row sessionCheckpointRow) *SessionTaskSnapshot { } } -// ListSessionCheckpoints returns an owner's READY and FAILED checkpoints for the -// source session in ascending ID order after afterID, up to limit. CREATING and -// DELETING stay hidden. Retained checkpoints remain listable after source deletion. +// ListSessionCheckpoints returns an owner's READY checkpoints for the source +// 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, @@ -456,7 +454,7 @@ func (c *Client) ListSessionCheckpoints(ctx context.Context, sessionID, userID, FROM session_checkpoint WHERE source_session_id = $1 AND user_id = $2 - AND state IN ('READY', 'FAILED') + AND state = 'READY' AND (NULLIF($3::text, '') IS NULL OR id > NULLIF($3::text, '')::uuid) ORDER BY id LIMIT $4 diff --git a/go/core/internal/database/client_session_test.go b/go/core/internal/database/client_session_test.go index 1e2d769a75..ac905ba220 100644 --- a/go/core/internal/database/client_session_test.go +++ b/go/core/internal/database/client_session_test.go @@ -343,14 +343,14 @@ func TestSessionCheckpointRetainsRecordedBoundary(t *testing.T) { t.Fatalf("checkpoint replay after source deletion = %+v, error %v", replayed, err) } listed, err := client.ListSessionCheckpoints(ctx, sessionID, "alice", "", 10) - if err != nil || len(listed) != 2 || listed[0].Id != checkpoint.Id || listed[1].Id != failed.Id { + if err != nil || len(listed) != 1 || listed[0].Id != checkpoint.Id { t.Fatalf("listed checkpoints = %+v, error %v", listed, err) } - require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_FAILED, listed[1].State) page, err := client.ListSessionCheckpoints(ctx, sessionID, "alice", checkpoint.Id, 1) require.NoError(t, err) - require.Len(t, page, 1) - require.Equal(t, failed.Id, page[0].Id, "pagination must include failed cleanup candidates") + 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) diff --git a/go/core/internal/grpcserver/checkpoint_failure_test.go b/go/core/internal/grpcserver/checkpoint_failure_test.go deleted file mode 100644 index f9d7a1fa5a..0000000000 --- a/go/core/internal/grpcserver/checkpoint_failure_test.go +++ /dev/null @@ -1,119 +0,0 @@ -package grpcserver - -import ( - "context" - "crypto/sha256" - "testing" - - "github.com/a2aproject/a2a-go/v2/a2a" - "github.com/google/uuid" - "github.com/stretchr/testify/require" - ateapipb "go.ate.dev/substrate/api/ateapi/v1" - "google.golang.org/genproto/googleapis/rpc/errdetails" - "google.golang.org/grpc/codes" - "google.golang.org/grpc/metadata" - "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/checkpoint" - "github.com/kagent-dev/kagent/go/core/internal/substrate" - "github.com/kagent-dev/kagent/go/core/pkg/auth" -) - -func TestFailedCheckpointIdentitySurvivesGRPCMapping(t *testing.T) { - dsn := dbtest.StartT(context.WithoutCancel(t.Context()), t) - dbtest.MigrateT(t, dsn, false) - pool, err := database.Connect(t.Context(), &database.PostgresConfig{Pool: dbtest.PoolConfig(dsn)}) - require.NoError(t, err) - t.Cleanup(pool.Close) - store := database.NewClient(pool, "public") - session := createTaskStoreSession(t, store) - message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("checkpoint boundary")) - message.ContextID = session.ContextId - task := a2a.NewSubmittedTask(message, message) - task.Status.State = a2a.TaskStateCompleted - digest := sha256.Sum256([]byte(task.ID)) - version, err := store.CreateRuntimeTask(t.Context(), session.Id, digest[:], task, "") - require.NoError(t, err) - require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) - boundary, err := store.ClaimSessionQuiescence(t.Context()) - 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"})) - ctx, cancel := context.WithTimeout(t.Context(), database.RuntimeOperationTimeout) - defer cancel() - executor, requestID := uuid.New(), uuid.NewString() - failed, _, err := store.ReserveSessionCheckpoint(ctx, &apiv1.Checkpoint{ - Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID)}, "alice", requestID, executor) - require.NoError(t, err) - failed, err = store.FinalizeClaimedSessionCheckpoint(ctx, failed, executor, "partial-tag-uid", "", "private legacy backend detail") - require.NoError(t, err) - tags := &checkpointFailureTags{tag: &ateapipb.Tag{Metadata: &ateapipb.ResourceMetadata{ - Atespace: "team-a", Name: "checkpoint-" + failed.Id, Uid: "partial-tag-uid"}}} - connection := newConfigurationConnection(t, nil, func(config *Config) { - config.CheckpointService = checkpoint.NewService(store, &auth.NoopAuthorizer{}, tags, nil) - }) - client := apiv1.NewCheckpointServiceClient(connection) - ctx = metadata.AppendToOutgoingContext(ctx, "x-user-id", "alice") - _, err = client.CreateCheckpoint(ctx, &apiv1.CreateCheckpointRequest{ - SessionId: failed.SessionId, RequestId: requestID, ExpectedHeadTaskId: failed.HeadTaskId}) - require.Equal(t, codes.FailedPrecondition, status.Code(err)) - var reason, id string - for _, detail := range status.Convert(err).Details() { - switch info := detail.(type) { - case *errdetails.ErrorInfo: - require.Equal(t, "kagent.dev", info.Domain) - reason = info.Reason - case *errdetails.ResourceInfo: - id = info.ResourceName - } - } - require.Equal(t, "KAGENT_CHECKPOINT_FAILED", reason) - require.Equal(t, failed.Id, id) - require.NotContains(t, err.Error(), "private") - read, err := client.GetCheckpoint(ctx, &apiv1.GetCheckpointRequest{CheckpointId: id}) - require.NoError(t, err) - require.Equal(t, failed.State, read.Checkpoint.State) - require.NotContains(t, read.String(), "private") - listed, err := client.ListCheckpoints(ctx, &apiv1.ListCheckpointsRequest{SessionId: failed.SessionId}) - require.NoError(t, err) - require.Len(t, listed.Checkpoints, 1) - require.Equal(t, id, listed.Checkpoints[0].Id) - require.NotContains(t, listed.String(), "private") - other := metadata.AppendToOutgoingContext(t.Context(), "x-user-id", "bob") - _, err = client.GetCheckpoint(other, &apiv1.GetCheckpointRequest{CheckpointId: id}) - require.Equal(t, codes.NotFound, status.Code(err)) - _, err = client.DeleteCheckpoint(other, &apiv1.DeleteCheckpointRequest{CheckpointId: id}) - require.Equal(t, codes.NotFound, status.Code(err)) - _, err = client.ForkSession(ctx, &apiv1.ForkSessionRequest{CheckpointId: id, RequestId: uuid.NewString()}) - require.Equal(t, codes.NotFound, status.Code(err)) - _, err = client.DeleteCheckpoint(ctx, &apiv1.DeleteCheckpointRequest{CheckpointId: id}) - require.NoError(t, err) - require.Nil(t, tags.tag) - _, err = store.GetCheckpointOperation(t.Context(), id, "alice") - require.ErrorIs(t, err, database.ErrNotFound) - listed, err = client.ListCheckpoints(ctx, &apiv1.ListCheckpointsRequest{SessionId: failed.SessionId}) - require.NoError(t, err) - require.Empty(t, listed.Checkpoints) -} - -// Only the external Tag API is replaced: the generated client, interceptors, -// service and PostgreSQL use their production implementations. -type checkpointFailureTags struct { - *substrate.Client - tag *ateapipb.Tag -} - -func (s *checkpointFailureTags) GetTag(context.Context, string, string) (*ateapipb.Tag, error) { - return s.tag, nil -} - -func (s *checkpointFailureTags) DeleteTag(_ context.Context, atespace, name, uid string) error { - if s.tag.GetMetadata().GetAtespace() != atespace || s.tag.GetMetadata().GetName() != name || s.tag.GetMetadata().GetUid() != uid { - return status.Error(codes.FailedPrecondition, "tag identity changed") - } - s.tag = nil - return nil -} diff --git a/go/core/internal/mcp/checkpoint_errors_test.go b/go/core/internal/mcp/checkpoint_errors_test.go index f2f2e2cc1d..78fd0177f7 100644 --- a/go/core/internal/mcp/checkpoint_errors_test.go +++ b/go/core/internal/mcp/checkpoint_errors_test.go @@ -49,35 +49,32 @@ func TestCheckpointToolsHideStorageFailures(t *testing.T) { } } -func TestCheckpointToolsRetainTerminalIdentity(t *testing.T) { - store := &checkpointErrorStore{phase: "failed"} - h := &Handler{checkpoints: checkpoint.NewService(store, &auth.NoopAuthorizer{}, nil, nil)} - result := checkpointToolResponse(t, h.checkpoints, createCheckpointToolName, map[string]any{ - "session_id": testSessionID, "expected_head_task_id": testTaskID, "request_id": "failed-request"}) - require.Equal(t, true, result["isError"]) - meta := result["_meta"].(map[string]any) - require.Equal(t, "FailedPrecondition", meta["kagent.dev/grpc-code"]) - require.Equal(t, "KAGENT_CHECKPOINT_FAILED", meta["kagent.dev/error-reason"]) - id := meta["kagent.dev/checkpoint-id"].(string) - require.Equal(t, "33333333-3333-4333-8333-333333333333", id) - read, output, err := h.getCheckpoint(authContext(), nil, CheckpointIDInput{CheckpointID: id}) - require.NoError(t, err) - require.False(t, read.IsError) - require.Equal(t, id, output.Checkpoint.ID) - require.NotContains(t, output.Checkpoint.Failure.Message, "private") -} - 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")} - h := &Handler{checkpoints: checkpoint.NewService(&checkpointErrorStore{phase: "backend"}, &auth.NoopAuthorizer{}, backend, nil)} - result, _, err := h.deleteCheckpoint(authContext(), nil, CheckpointIDInput{CheckpointID: "33333333-3333-4333-8333-333333333333"}) - require.NoError(t, err) - require.True(t, result.IsError) + 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.Equal(t, code.String(), result.Meta["kagent.dev/grpc-code"]) - require.NotContains(t, fmt.Sprint(result.Content), "private") + 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"]) }) } } @@ -85,13 +82,12 @@ func TestCheckpointToolsRetainCurrentBackendStatus(t *testing.T) { 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)} - h := &Handler{checkpoints: checkpoint.NewService(store, &auth.NoopAuthorizer{}, nil, nil)} - result, _, err := h.createCheckpoint(authContext(), nil, CreateCheckpointInput{ - SessionID: testSessionID, ExpectedHeadTaskID: testTaskID, RequestID: "canceled-request"}) - require.NoError(t, err) - require.True(t, result.IsError) - require.Equal(t, status.FromContextError(cause).Code().String(), result.Meta["kagent.dev/grpc-code"]) - require.NotContains(t, fmt.Sprint(result.Content), "private") + 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") } } @@ -107,12 +103,16 @@ func checkpointToolResponse(t *testing.T, service *checkpoint.Service, tool stri type checkpointErrorStore struct { *database.Client - phase string - failures int - claimErr error + 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") @@ -122,10 +122,6 @@ func (s *checkpointErrorStore) ReserveSessionCheckpoint(_ context.Context, cp *a if s.phase == "claim failure" { cp.State = apiv1.CheckpointState_CHECKPOINT_STATE_CREATING } - if s.phase == "failed" { - cp.State = apiv1.CheckpointState_CHECKPOINT_STATE_FAILED - cp.Failure = &apiv1.Failure{Message: "private legacy failure"} - } return cp, &database.SessionTaskSnapshot{}, nil } func (s *checkpointErrorStore) ClaimCheckpointOperation(context.Context, string, string, uuid.UUID) (*database.CheckpointOperation, error) { @@ -156,10 +152,6 @@ func (s *checkpointErrorStore) BeginDeleteSessionCheckpoint(context.Context, str return nil, database.ErrConflict } -func (s *checkpointErrorStore) GetSessionCheckpoint(context.Context, string, string) (*apiv1.Checkpoint, error) { - return &apiv1.Checkpoint{Id: "33333333-3333-4333-8333-333333333333", State: apiv1.CheckpointState_CHECKPOINT_STATE_FAILED, Failure: &apiv1.Failure{Message: "private legacy failure"}}, nil -} - type checkpointErrorBackend struct { *substrate.Client err error diff --git a/go/core/internal/mcp/checkpoints.go b/go/core/internal/mcp/checkpoints.go index f8b0b72e97..80f7704f84 100644 --- a/go/core/internal/mcp/checkpoints.go +++ b/go/core/internal/mcp/checkpoints.go @@ -9,12 +9,10 @@ import ( "github.com/google/uuid" "github.com/modelcontextprotocol/go-sdk/mcp" "google.golang.org/genproto/googleapis/rpc/errdetails" - "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/service/checkpoint" - "github.com/kagent-dev/kagent/go/core/internal/service/serviceerrors" ) const ( @@ -106,7 +104,13 @@ 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 { - return checkpointToolError(err), CreateCheckpointOutput{}, nil + 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} + } + } + return result, CreateCheckpointOutput{}, nil } output := CreateCheckpointOutput{Checkpoint: checkpointSummary(created)} return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: fmt.Sprintf("Created checkpoint %s", created.GetId())}}}, output, nil @@ -143,37 +147,15 @@ func stableRequestID(id string) string { return uuid.NewString() } -// Checkpoint services classify storage failures and sanitize backend statuses. -// Project only their public message, retaining status and cleanup identity for -// MCP clients that do not pass through the gRPC error-mapping interceptor. +// 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 = status.Error(codes.Canceled, "Checkpoint operation canceled") + err = errors.New("checkpoint operation canceled") } else if errors.Is(err, context.DeadlineExceeded) { - err = status.Error(codes.DeadlineExceeded, "Checkpoint operation deadline exceeded") - } - if serviceerrors.CodeOf(err) != "" { - return sessionToolError(err) - } - value, ok := status.FromError(err) - if !ok { - return sessionToolError(serviceerrors.NewInternal("Checkpoint operation could not complete", err)) - } - result := toolError(errors.New(value.Message())) - result.Meta = mcp.Meta{"kagent.dev/grpc-code": value.Code().String()} - for _, detail := range value.Details() { - switch info := detail.(type) { - case *errdetails.ErrorInfo: - if info.Domain == "kagent.dev" { - result.Meta["kagent.dev/error-reason"] = info.Reason - } - case *errdetails.ResourceInfo: - if info.ResourceType == "kagent.api.v1.Checkpoint" { - result.Meta["kagent.dev/checkpoint-id"] = info.ResourceName - } - } + err = errors.New("checkpoint operation deadline exceeded") } - return result + return sessionToolError(err) } func checkpointSummary(value *apiv1.Checkpoint) CheckpointSummary { diff --git a/go/core/internal/service/checkpoint/completion_test.go b/go/core/internal/service/checkpoint/completion_test.go index 19a16d6dd2..717256500f 100644 --- a/go/core/internal/service/checkpoint/completion_test.go +++ b/go/core/internal/service/checkpoint/completion_test.go @@ -15,7 +15,7 @@ import ( ) func TestCreateObservesReplacementCompletion(t *testing.T) { - persistence := reviewStore(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"}) @@ -69,7 +69,7 @@ func TestCreateObservesReplacementCompletion(t *testing.T) { _, err = NewService(persistence, testAuthorizer{}, tags, nil).Create(owner, source.Id, request, head) require.NoError(t, err) case "failed": - assertFailedCheckpointIdentity(t, err, checkpointID) + 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) @@ -86,7 +86,7 @@ func TestCreateObservesReplacementCompletion(t *testing.T) { } func TestClaimConflictObservesCompletionWithoutRelease(t *testing.T) { - persistence := reviewStore(t) + persistence := recoveryStore(t) owner := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) source, head := recoverySource(t, persistence) tags := &testTags{snapshotURI: "s3://snapshots/source"} @@ -112,7 +112,7 @@ func TestClaimConflictObservesCompletionWithoutRelease(t *testing.T) { } func TestDeleteObservesReplacementCompletion(t *testing.T) { - persistence := reviewStore(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"}) diff --git a/go/core/internal/service/checkpoint/discovery_test.go b/go/core/internal/service/checkpoint/discovery_test.go deleted file mode 100644 index e5f86c583c..0000000000 --- a/go/core/internal/service/checkpoint/discovery_test.go +++ /dev/null @@ -1,131 +0,0 @@ -package checkpoint - -import ( - "context" - "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" - "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/service/serviceerrors" - "github.com/kagent-dev/kagent/go/core/pkg/auth" -) - -func TestOwnerDiscoversFailedCheckpointAndDeletesPartialTag(t *testing.T) { - persistence := reviewStore(t) - owner := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) - other := auth.AuthSessionTo(t.Context(), testSession{userID: "bob"}) - 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() - created, err := service.Create(owner, source.Id, request, head) - require.Equal(t, codes.Unavailable, status.Code(err), "an uncertain attempt stays retryable") - require.Nil(t, created) - pending, err := service.List(owner, ListRequest{SessionID: source.Id}) - require.NoError(t, err) - require.Empty(t, pending.Checkpoints, "CREATING remains hidden") - - _, terminalErr := service.Create(owner, source.Id, request, head) - require.Error(t, terminalErr) - // Discover the identity through public methods only. Neither the backend - // fixture nor an internal lifecycle read supplies the ID used for cleanup. - listed, err := service.List(owner, ListRequest{SessionID: source.Id, PageSize: 1}) - require.NoError(t, err) - require.Len(t, listed.Checkpoints, 1) - failed := listed.Checkpoints[0] - require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_FAILED, failed.State) - require.Empty(t, listed.NextPageToken) - assertFailedCheckpointIdentity(t, terminalErr, failed.Id) - read, err := service.Get(owner, failed.Id) - require.NoError(t, err) - require.True(t, proto.Equal(failed, read)) - _, err = service.Create(owner, source.Id, request, head) - assertFailedCheckpointIdentity(t, err, failed.Id) - require.Equal(t, 2, tags.createCalls, "terminal retries must not repeat copying") - - _, err = service.Get(other, failed.Id) - require.True(t, serviceerrors.IsCode(err, serviceerrors.CodeNotFound)) - foreign, err := service.List(other, ListRequest{SessionID: source.Id}) - require.NoError(t, err) - require.Empty(t, foreign.Checkpoints) - require.True(t, serviceerrors.IsCode(service.Delete(other, failed.Id), serviceerrors.CodeNotFound)) - _, err = service.Fork(owner, failed.Id, uuid.NewString()) - require.True(t, serviceerrors.IsCode(err, serviceerrors.CodeNotFound), "FAILED must never become a fork source") - _, _, err = persistence.ForkSession(owner, failed.Id, "alice", uuid.NewString(), uuid.NewString()) - require.ErrorIs(t, err, database.ErrNotFound, "the store must independently reject FAILED forks") - - tags.deleteErr = status.Error(codes.PermissionDenied, "private object-store configuration") - require.Equal(t, codes.PermissionDenied, status.Code(service.Delete(owner, failed.Id))) - retained, err := persistence.GetCheckpointOperation(owner, failed.Id, "alice") - require.NoError(t, err) - require.Equal(t, "tag-uid", retained.TagUID, "failed cleanup must retain its deletion handle") - _, err = service.Get(owner, failed.Id) - require.True(t, serviceerrors.IsCode(err, serviceerrors.CodeNotFound), "DELETING remains hidden") - tags.deleteErr = nil - require.NoError(t, NewService(persistence, testAuthorizer{}, tags, nil).Delete(owner, failed.Id)) - require.Nil(t, tags.created) - _, err = persistence.GetCheckpointOperation(owner, failed.Id, "alice") - require.ErrorIs(t, err, database.ErrNotFound, "successful cleanup must release the retained row and references") - listed, err = service.List(owner, ListRequest{SessionID: source.Id}) - require.NoError(t, err) - require.Empty(t, listed.Checkpoints) -} - -func TestFailedCheckpointReadsHideLegacyBackendDetails(t *testing.T) { - persistence := reviewStore(t) - owner := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) - source, head := recoverySource(t, persistence) - ctx, cancel := context.WithTimeout(owner, database.RuntimeOperationTimeout) - defer cancel() - executor := uuid.New() - checkpoint, _, err := persistence.ReserveSessionCheckpoint(ctx, - &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: source.Id, HeadTaskId: head}, "alice", uuid.NewString(), executor) - require.NoError(t, err) - const privateFailure = "private bucket and credentials in a legacy failure" - _, err = persistence.FinalizeClaimedSessionCheckpoint(ctx, checkpoint, executor, "legacy-tag", "", privateFailure) - require.NoError(t, err) - service := NewService(persistence, testAuthorizer{}, &testTags{}, nil) - read, err := service.Get(owner, checkpoint.Id) - require.NoError(t, err) - require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_FAILED, read.State) - require.NotEmpty(t, read.GetFailure().GetMessage()) - require.NotContains(t, read.String(), "private") - listed, err := service.List(owner, ListRequest{SessionID: source.Id}) - require.NoError(t, err) - require.Len(t, listed.Checkpoints, 1) - require.True(t, proto.Equal(read, listed.Checkpoints[0])) - stored, err := persistence.GetCheckpointOperation(owner, checkpoint.Id, "alice") - require.NoError(t, err) - require.Equal(t, privateFailure, stored.Checkpoint.Failure.Message, "public projection must not rewrite durable data") - require.NoError(t, service.Delete(owner, checkpoint.Id)) -} - -func assertFailedCheckpointIdentity(t *testing.T, err error, id string) { - t.Helper() - require.Equal(t, codes.FailedPrecondition, status.Code(err), "a terminal failure must not invite transient retries") - var reason, resource string - for _, detail := range status.Convert(err).Details() { - switch info := detail.(type) { - case *errdetails.ErrorInfo: - require.Equal(t, "kagent.dev", info.Domain) - reason = info.Reason - case *errdetails.ResourceInfo: - require.Equal(t, "kagent.api.v1.Checkpoint", info.ResourceType) - resource = info.ResourceName - } - } - require.Equal(t, "KAGENT_CHECKPOINT_FAILED", reason) - require.Equal(t, id, resource) - require.NotContains(t, err.Error(), "private") -} 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..e6aa55d702 --- /dev/null +++ b/go/core/internal/service/checkpoint/failed_cleanup_test.go @@ -0,0 +1,134 @@ +package checkpoint + +import ( + "context" + "errors" + "testing" + + "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 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 index db5eb89e1d..1f86a66e4d 100644 --- a/go/core/internal/service/checkpoint/migration_test.go +++ b/go/core/internal/service/checkpoint/migration_test.go @@ -21,7 +21,7 @@ import ( ) // This is a populated legacy-schema fixture, not a previous-release binary. -// Backend serialization is tested separately against real Substrate. +// 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) diff --git a/go/core/internal/service/checkpoint/review_test.go b/go/core/internal/service/checkpoint/operation_recovery_test.go similarity index 88% rename from go/core/internal/service/checkpoint/review_test.go rename to go/core/internal/service/checkpoint/operation_recovery_test.go index 1d995e2b1f..10f6a8e216 100644 --- a/go/core/internal/service/checkpoint/review_test.go +++ b/go/core/internal/service/checkpoint/operation_recovery_test.go @@ -21,7 +21,7 @@ import ( ) func TestCheckpointRequestsAndRecovery(t *testing.T) { - store := reviewStore(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) { @@ -260,6 +260,60 @@ func TestCheckpointRequestsAndRecovery(t *testing.T) { }) } +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 @@ -273,7 +327,7 @@ func (s *countedClaims) ClaimCheckpointOperation(ctx context.Context, id, owner return s.Client.ClaimCheckpointOperation(ctx, id, owner, executor) } -func reviewStore(t *testing.T) *database.Client { +func recoveryStore(t *testing.T) *database.Client { t.Helper() dsn := dbtest.StartT(context.WithoutCancel(t.Context()), t) require.NoError(t, dbtest.Migrate(dsn, false)) @@ -293,7 +347,7 @@ func TestPublicBackendErrorsHideConfiguration(t *testing.T) { } func TestPublicDeleteReportsCurrentBackendStatus(t *testing.T) { - persistence := reviewStore(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} { @@ -324,7 +378,7 @@ func TestPublicDeleteReportsCurrentBackendStatus(t *testing.T) { } func TestPublicCreateRetainsUnverifiedTag(t *testing.T) { - persistence := reviewStore(t) + persistence := recoveryStore(t) source, head := recoverySource(t, persistence) tags := &testTags{snapshotURI: "s3://snapshots/source"} service := NewService(persistence, testAuthorizer{}, tags, nil) @@ -355,7 +409,7 @@ func TestPublicCreateRetainsUnverifiedTag(t *testing.T) { } func TestSlowRecoveryDoesNotBlockRequests(t *testing.T) { - persistence := reviewStore(t) + persistence := recoveryStore(t) slowSource, slowHead := recoverySource(t, persistence) fastSource, fastHead := recoverySource(t, persistence) owner := auth.AuthSessionTo(t.Context(), testSession{userID: "alice"}) @@ -403,7 +457,7 @@ func (s *lostCheckpointFinalization) FinalizeClaimedSessionCheckpoint(ctx contex } func TestPublicCheckpointErrorsHideReleaseFailures(t *testing.T) { - persistence := reviewStore(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) { diff --git a/go/core/internal/service/checkpoint/recovery.go b/go/core/internal/service/checkpoint/recovery.go index 2023c06ce4..113c5e593a 100644 --- a/go/core/internal/service/checkpoint/recovery.go +++ b/go/core/internal/service/checkpoint/recovery.go @@ -89,11 +89,14 @@ func (s *Service) claimAndExecute(ctx context.Context, observed *database.Checkp 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) + return s.execute(ctx, op, executor, false) } // Losing execution authority permits observation, never another backend call. @@ -129,7 +132,7 @@ func releaseClaim(ctx context.Context, id string, release func(context.Context) return nil } -func (s *Service) execute(ctx context.Context, op *database.CheckpointOperation, executor uuid.UUID) (result *apiv1.Checkpoint, err error) { +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 { @@ -144,7 +147,12 @@ func (s *Service) execute(ctx context.Context, op *database.CheckpointOperation, releaseNeeded = err != nil return nil, err } - tag, err := s.ensureTag(ctx, op.Checkpoint, op.Snapshot) + 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) } @@ -155,9 +163,7 @@ func (s *Service) execute(ctx context.Context, op *database.CheckpointOperation, if err != nil { failure = checkpointFailureMessage } - // Under the required backend serialization contract, the definitive reply - // permits FAILED to release admission and retain its UID for explicit Delete. - // Expiry of the kagent claim alone never establishes that copying stopped. + // 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, @@ -174,3 +180,24 @@ func (s *Service) execute(ctx context.Context, op *database.CheckpointOperation, 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_objectstore_test.go b/go/core/internal/service/checkpoint/recovery_objectstore_test.go deleted file mode 100644 index bf43545299..0000000000 --- a/go/core/internal/service/checkpoint/recovery_objectstore_test.go +++ /dev/null @@ -1,55 +0,0 @@ -package checkpoint - -import ( - "context" - "testing" - "time" - - "github.com/aws/aws-sdk-go-v2/aws" - "github.com/aws/aws-sdk-go-v2/credentials" - "github.com/aws/aws-sdk-go-v2/service/s3" - "github.com/google/uuid" - "github.com/stretchr/testify/require" -) - -// Fault policies and injected objects belong to a bucket created exclusively -// for this scenario; the installation's shared snapshot bucket is untouched. -func checkpointObjectStore(t *testing.T, endpoint, accessKey, secretKey string) (*s3.Client, string) { - t.Helper() - require.NotEmpty(t, endpoint, "real object-store endpoint is required") - require.NotEmpty(t, accessKey) - require.NotEmpty(t, secretKey) - client := s3.NewFromConfig(aws.Config{ - Region: "us-east-1", Credentials: credentials.NewStaticCredentialsProvider(accessKey, secretKey, ""), RetryMaxAttempts: 1, - }, func(o *s3.Options) { o.BaseEndpoint = aws.String(endpoint); o.UsePathStyle = true }) - bucket := "checkpoint-recovery-" + uuid.NewString() - _, err := client.CreateBucket(t.Context(), &s3.CreateBucketInput{Bucket: aws.String(bucket)}) - require.NoError(t, err) - t.Cleanup(func() { - ctx, cancel := context.WithTimeout(context.Background(), time.Minute) - defer cancel() - _, err := client.DeleteBucketPolicy(ctx, &s3.DeleteBucketPolicyInput{Bucket: aws.String(bucket)}) - require.NoError(t, err) - for _, key := range checkpointObjectKeys(t, ctx, client, bucket, "") { - _, err := client.DeleteObject(ctx, &s3.DeleteObjectInput{Bucket: aws.String(bucket), Key: aws.String(key)}) - require.NoError(t, err) - } - _, err = client.DeleteBucket(ctx, &s3.DeleteBucketInput{Bucket: aws.String(bucket)}) - require.NoError(t, err) - }) - return client, bucket -} - -func checkpointObjectKeys(t *testing.T, ctx context.Context, client *s3.Client, bucket, prefix string) []string { - t.Helper() - pages := s3.NewListObjectsV2Paginator(client, &s3.ListObjectsV2Input{Bucket: aws.String(bucket), Prefix: aws.String(prefix)}) - var keys []string - for pages.HasMorePages() { - page, err := pages.NextPage(ctx) - require.NoError(t, err) - for _, object := range page.Contents { - keys = append(keys, aws.ToString(object.Key)) - } - } - return keys -} diff --git a/go/core/internal/service/checkpoint/recovery_substrate_test.go b/go/core/internal/service/checkpoint/recovery_substrate_test.go deleted file mode 100644 index 7b44085898..0000000000 --- a/go/core/internal/service/checkpoint/recovery_substrate_test.go +++ /dev/null @@ -1,564 +0,0 @@ -package checkpoint - -import ( - "context" - "crypto/tls" - "crypto/x509" - "encoding/json" - "fmt" - "net/url" - "os" - "strings" - "sync/atomic" - "testing" - "time" - - "github.com/aws/aws-sdk-go-v2/aws" - "github.com/aws/aws-sdk-go-v2/service/s3" - "github.com/google/uuid" - "github.com/stretchr/testify/require" - ateapipb "go.ate.dev/substrate/api/ateapi/v1" - "golang.org/x/sync/errgroup" - "google.golang.org/grpc" - "google.golang.org/grpc/codes" - "google.golang.org/grpc/credentials" - "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/dbtest" - "github.com/kagent-dev/kagent/go/core/internal/substrate" - "github.com/kagent-dev/kagent/go/core/pkg/auth" -) - -// This suite uses real Actors, snapshot copies, UID-conditional deletion and -// PostgreSQL. Faults discard actual RPC results or prevent one RPC from being -// delivered, or deny reads/deletes in an exclusively owned object-store bucket; -// no backend response or persisted lifecycle state is fabricated. -// The configuration is test-only; CI runs this after installing Substrate. -func TestCheckpointRecoveryWithSubstrate(t *testing.T) { - path := os.Getenv("KAGENT_E2E_CHECKPOINT_SUBSTRATE_CONFIG") - if path == "" { - t.Skip("requires real Substrate; run core/test/e2e/checkpoint-recovery.sh") - } - var config struct { - Endpoint string `json:"endpoint"` - ServerName string `json:"serverName"` - CAFile string `json:"caFile"` - ClientCertFile string `json:"clientCertFile"` - Atespace string `json:"atespace"` - Image string `json:"image"` - ObjectStoreEndpoint string `json:"objectStoreEndpoint"` - ObjectStoreAccessKey string `json:"objectStoreAccessKey"` - ObjectStoreSecretKey string `json:"objectStoreSecretKey"` - } - data, err := os.ReadFile(path) - require.NoError(t, err) - require.NoError(t, json.Unmarshal(data, &config)) - require.NotEmpty(t, config.Endpoint) - require.NotEmpty(t, config.Atespace) - require.NotEmpty(t, config.Image) - certificate, err := tls.LoadX509KeyPair(config.ClientCertFile, config.ClientCertFile) - require.NoError(t, err) - ca, err := os.ReadFile(config.CAFile) - require.NoError(t, err) - roots := x509.NewCertPool() - require.True(t, roots.AppendCertsFromPEM(ca)) - conn, err := grpc.NewClient(config.Endpoint, grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{ - MinVersion: tls.VersionTLS12, RootCAs: roots, Certificates: []tls.Certificate{certificate}, ServerName: config.ServerName, - }))) - require.NoError(t, err) - t.Cleanup(func() { require.NoError(t, conn.Close()) }) - client := &substrate.Client{ControlClient: ateapipb.NewControlClient(conn)} - ctx, cancel := context.WithTimeout(t.Context(), 8*time.Minute) - t.Cleanup(cancel) - template, err := client.CreateActorTemplate(ctx, &ateapipb.ActorTemplate{ - Metadata: &ateapipb.ResourceMetadata{Atespace: config.Atespace, Name: "checkpoint-recovery-" + uuid.NewString()}, - Resources: &ateapipb.ResourceRequirements{Limits: &ateapipb.Resources{CpuMillicores: 500, MemoryBytes: 256 << 20}}, - Containers: []*ateapipb.Container{{Name: "source", Image: config.Image, - Command: []string{"/bin/sh", "-c", "echo checkpoint-recovery > /data/proof; exec sleep 3600"}, - VolumeMounts: []*ateapipb.VolumeMount{{Name: "data", MountPath: "/data"}}, - }}, - Volumes: []*ateapipb.Volume{{Name: "data", DurableDir: &ateapipb.DurableDirVolumeSource{}}}, - SnapshotConfig: &ateapipb.SnapshotConfig{PreferredFidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES, - GoldenSnapshotConfig: &ateapipb.GoldenSnapshotConfig{Mode: ateapipb.GoldenSnapshotMode_GOLDEN_SNAPSHOT_MODE_DISABLED}}, - }) - require.NoError(t, err) - t.Cleanup(func() { - cleanup, stop := context.WithTimeout(context.Background(), time.Minute) - defer stop() - require.NoError(t, client.DeleteActorTemplate(cleanup, config.Atespace, template.Metadata.Name, template.Metadata.Uid)) - }) - 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) - newStore := func() *database.Client { return database.NewClient(pool, "public") } - owner := auth.AuthSessionTo(ctx, testSession{userID: "alice"}) - newSourceForTemplate := func(t *testing.T, template *ateapipb.ActorTemplate) (*database.Client, *apiv1.Session, string) { - t.Helper() - name := substrate.ActorName(uuid.NewString()) - created, err := client.CreateActor(ctx, config.Atespace, name, config.Atespace, template.Metadata.Name) - require.NoError(t, err) - t.Cleanup(func() { - cleanup, stop := context.WithTimeout(context.Background(), time.Minute) - defer stop() - err := client.DeleteActor(cleanup, config.Atespace, name, created.GetMetadata().GetUid()) - if status.Code(err) != codes.NotFound { - require.NoError(t, err) - } - }) - _, err = client.ResumeActor(ctx, config.Atespace, name) - require.NoError(t, err) - actor, err := client.SuspendActor(ctx, config.Atespace, name) - require.NoError(t, err) - require.NotEmpty(t, substrate.SnapshotURI(substrate.LatestActorSnapshot(actor))) - store := newStore() - session, head := recoverySourceForActor(t, store, actor) - return store, session, head - } - newSource := func(t *testing.T) (*database.Client, *apiv1.Session, string) { - return newSourceForTemplate(t, template) - } - deleteCheckpoint := func(t *testing.T, service *Service, id string) { - t.Helper() - require.NoError(t, service.Delete(owner, id)) - _, err := service.store.GetCheckpointOperation(ctx, id, "alice") - require.ErrorIs(t, err, database.ErrNotFound) - _, err = client.GetTag(ctx, config.Atespace, tagName(id)) - require.Equal(t, codes.NotFound, status.Code(err)) - } - cleanupTag := func(t *testing.T, id string) { - t.Helper() - t.Cleanup(func() { - cleanup, stop := context.WithTimeout(context.Background(), time.Minute) - defer stop() - tag, err := client.GetTag(cleanup, config.Atespace, tagName(id)) - if status.Code(err) == codes.NotFound { - return - } - require.NoError(t, err) - require.NoError(t, client.DeleteTag(cleanup, config.Atespace, tagName(id), tag.GetMetadata().GetUid())) - }) - } - - t.Run("Given an abandoned reservation, when a replacement starts, assert that it publishes a real retained copy", func(t *testing.T) { - store, source, head := newSource(t) - service := NewService(store, testAuthorizer{}, client, nil) - request := uuid.NewString() - cp, _, err := service.beginCreate(owner, source.Id, request, head) - require.NoError(t, err) - cleanupTag(t, cp.Id) - replacement := NewService(newStore(), testAuthorizer{}, client, nil) - ready, err := replacement.Create(owner, source.Id, request, head) - require.NoError(t, err) - require.Equal(t, cp.Id, ready.Id) - tag, err := client.GetTag(ctx, config.Atespace, tagName(cp.Id)) - require.NoError(t, err) - require.NotEmpty(t, substrate.SnapshotURI(tag.GetStatus().GetSnapshot())) - clone := substrate.ActorName(uuid.NewString()) - cloned, err := client.CreateActorFromTag(ctx, config.Atespace, clone, config.Atespace, template.Metadata.Name, config.Atespace, tagName(cp.Id)) - require.NoError(t, err) - t.Cleanup(func() { - cleanup, stop := context.WithTimeout(context.Background(), time.Minute) - defer stop() - if err := client.DeleteActor(cleanup, config.Atespace, clone, cloned.GetMetadata().GetUid()); status.Code(err) != codes.NotFound { - require.NoError(t, err) - } - }) - _, err = client.ResumeActor(ctx, config.Atespace, clone) - require.NoError(t, err, "the recovered tag must contain a usable snapshot") - require.NoError(t, client.DeleteActor(ctx, config.Atespace, clone, cloned.GetMetadata().GetUid())) - deleteCheckpoint(t, replacement, cp.Id) - }) - - t.Run("Given a lost successful tag response, when recovery restarts, assert that it adopts the same backend UID", func(t *testing.T) { - store, source, head := newSource(t) - lost := &faultedSubstrateTags{Client: client, loseCreateResponse: true} - service := NewService(store, testAuthorizer{}, lost, nil) - request := uuid.NewString() - cp, _, err := service.beginCreate(owner, source.Id, request, head) - require.NoError(t, err) - cleanupTag(t, cp.Id) - op, err := store.GetCheckpointOperation(ctx, cp.Id, "alice") - require.NoError(t, err) - require.Error(t, service.recoverOperation(ctx, op)) - tag, err := client.GetTag(ctx, config.Atespace, tagName(cp.Id)) - require.NoError(t, err) - pending, err := store.GetCheckpointOperation(ctx, cp.Id, "alice") - require.NoError(t, err) - require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_CREATING, pending.Checkpoint.State) - replacement := NewService(newStore(), testAuthorizer{}, client, nil) - require.NoError(t, replacement.recoverOperation(ctx, pending)) - published, err := store.GetCheckpointOperation(ctx, cp.Id, "alice") - require.NoError(t, err) - require.Equal(t, tag.Metadata.Uid, published.TagUID) - deleteCheckpoint(t, replacement, cp.Id) - }) - - t.Run("Given a completed copy without database publication, when the worker restarts, assert that the original tag is published", func(t *testing.T) { - store, source, head := newSource(t) - fault := &failedCheckpointPublish{Client: store, entered: make(chan struct{})} - old := NewService(fault, testAuthorizer{}, client, nil) - request := uuid.NewString() - cp, _, err := old.beginCreate(owner, source.Id, request, head) - require.NoError(t, err) - cleanupTag(t, cp.Id) - op, err := store.GetCheckpointOperation(ctx, cp.Id, "alice") - require.NoError(t, err) - require.Error(t, old.recoverOperation(ctx, op)) - tag, err := client.GetTag(ctx, config.Atespace, tagName(cp.Id)) - require.NoError(t, err) - replacement := NewService(newStore(), testAuthorizer{}, client, nil) - require.NoError(t, replacement.recoverOperation(ctx, op)) - ready, err := store.GetCheckpointOperation(ctx, cp.Id, "alice") - require.NoError(t, err) - require.Equal(t, tag.Metadata.Uid, ready.TagUID) - deleteCheckpoint(t, replacement, cp.Id) - }) - - t.Run("Given a source deleted after reservation, when tag creation fails, assert that failed checkpoint deletion releases its row", func(t *testing.T) { - store, source, head := newSource(t) - service := NewService(store, testAuthorizer{}, client, nil) - cp, _, err := service.beginCreate(owner, source.Id, uuid.NewString(), head) - require.NoError(t, err) - cleanupTag(t, cp.Id) - actor, err := client.GetActor(ctx, config.Atespace, substrate.ActorName(source.Id)) - require.NoError(t, err) - require.NoError(t, client.DeleteActor(ctx, config.Atespace, substrate.ActorName(source.Id), actor.GetMetadata().GetUid())) - op, err := store.GetCheckpointOperation(ctx, cp.Id, "alice") - require.NoError(t, err) - require.Error(t, service.recoverOperation(ctx, op)) - waitCheckpointPhase(t, store, cp.Id, apiv1.CheckpointState_CHECKPOINT_STATE_FAILED) - _, err = store.BeginSessionOperation(ctx, source.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE, false) - require.NoError(t, err, "failed creation must stop blocking the source") - deleteCheckpoint(t, NewService(newStore(), testAuthorizer{}, client, nil), cp.Id) - }) - - t.Run("Given a source changed after copying, when deletion is interrupted, assert that its failed tag remains removable by a replacement", func(t *testing.T) { - store, source, head := newSource(t) - fault := &faultedSubstrateTags{Client: client, failDelete: true, afterCreate: func() error { - _, err := client.ResumeActor(ctx, config.Atespace, substrate.ActorName(source.Id)) - return err - }} - service := NewService(store, testAuthorizer{}, fault, nil) - cp, _, err := service.beginCreate(owner, source.Id, uuid.NewString(), head) - require.NoError(t, err) - cleanupTag(t, cp.Id) - op, err := store.GetCheckpointOperation(ctx, cp.Id, "alice") - require.NoError(t, err) - require.Error(t, service.recoverOperation(ctx, op)) - pending, err := store.GetCheckpointOperation(ctx, cp.Id, "alice") - require.NoError(t, err) - require.NotNil(t, pending.Checkpoint.Failure) - require.NotEmpty(t, pending.TagUID) - require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_FAILED, pending.Checkpoint.State) - _, err = store.BeginSessionOperation(ctx, source.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE, false) - require.NoError(t, err, "verification failure releases source admission after copying stopped") - require.Error(t, service.Delete(owner, cp.Id), "an unavailable delete must retain the tag reference") - deleting, err := store.GetCheckpointOperation(ctx, cp.Id, "alice") - require.NoError(t, err) - require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_DELETING, deleting.Checkpoint.State) - require.Equal(t, pending.TagUID, deleting.TagUID) - deleteCheckpoint(t, NewService(newStore(), testAuthorizer{}, client, nil), cp.Id) - }) - - t.Run("Given a partially copied tag and denied object cleanup, when the owner retries after storage repair, assert that retained objects and references are removed", func(t *testing.T) { - objects, bucket := checkpointObjectStore(t, config.ObjectStoreEndpoint, config.ObjectStoreAccessKey, config.ObjectStoreSecretKey) - privateTemplate := proto.Clone(template).(*ateapipb.ActorTemplate) - privateTemplate.Metadata = &ateapipb.ResourceMetadata{Atespace: config.Atespace, Name: "checkpoint-partial-" + uuid.NewString()} - privateTemplate.Status = nil - privateTemplate.SnapshotConfig.StorageLocation = "gs://" + bucket - privateTemplate, err := client.CreateActorTemplate(ctx, privateTemplate) - require.NoError(t, err) - t.Cleanup(func() { - cleanup, stop := context.WithTimeout(context.Background(), time.Minute) - defer stop() - require.NoError(t, client.DeleteActorTemplate(cleanup, config.Atespace, privateTemplate.Metadata.Name, privateTemplate.Metadata.Uid)) - }) - store, source, head := newSourceForTemplate(t, privateTemplate) - actor, err := client.GetActor(ctx, config.Atespace, substrate.ActorName(source.Id)) - require.NoError(t, err) - snapshot, err := url.Parse(substrate.SnapshotURI(substrate.LatestActorSnapshot(actor))) - require.NoError(t, err) - require.Equal(t, bucket, snapshot.Host) - prefix := strings.Trim(snapshot.Path, "/") + "/" - // Substrate copies in sorted order with eight workers. More than eight - // preceding objects make the final denied read occur after real copies. - for i := range 64 { - _, err := objects.PutObject(ctx, &s3.PutObjectInput{Bucket: aws.String(bucket), Key: aws.String(fmt.Sprintf("%szz-copy-probe/%04d", prefix, i)), Body: strings.NewReader("partial-copy-proof")}) - require.NoError(t, err) - } - deniedKey := prefix + "zz-copy-probe/9999-denied" - _, err = objects.PutObject(ctx, &s3.PutObjectInput{Bucket: aws.String(bucket), Key: aws.String(deniedKey), Body: strings.NewReader("deny this read")}) - require.NoError(t, err) - policy, err := json.Marshal(map[string]any{ - "Version": "2012-10-17", "Statement": []any{ - map[string]any{"Effect": "Deny", "Principal": "*", "Action": "s3:GetObject", "Resource": "arn:aws:s3:::" + bucket + "/" + deniedKey}, - map[string]any{"Effect": "Deny", "Principal": "*", "Action": "s3:DeleteObject", "Resource": "arn:aws:s3:::" + bucket + "/atespaces/" + config.Atespace + "/tags/*"}, - }, - }) - require.NoError(t, err) - service := NewService(store, testAuthorizer{}, client, nil) - cp, _, err := service.beginCreate(owner, source.Id, uuid.NewString(), head) - require.NoError(t, err) - cleanupTag(t, cp.Id) - _, err = objects.PutBucketPolicy(ctx, &s3.PutBucketPolicyInput{Bucket: aws.String(bucket), Policy: aws.String(string(policy))}) - require.NoError(t, err) - // Remove the denial before any Actor/tag/template cleanup, even on failure. - t.Cleanup(func() { - cleanup, stop := context.WithTimeout(context.Background(), time.Minute) - defer stop() - _, err := objects.DeleteBucketPolicy(cleanup, &s3.DeleteBucketPolicyInput{Bucket: aws.String(bucket)}) - require.NoError(t, err) - }) - op, err := store.GetCheckpointOperation(ctx, cp.Id, "alice") - require.NoError(t, err) - copyErr := service.recoverOperation(ctx, op) - require.Error(t, copyErr, "the actual object read must fail") - t.Logf("expected object-store copy failure: %v", copyErr) - tag, err := client.GetTag(ctx, config.Atespace, tagName(cp.Id)) - require.NoError(t, err) - require.Empty(t, substrate.SnapshotURI(tag.GetStatus().GetSnapshot())) - tagPrefix := "atespaces/" + config.Atespace + "/tags/" + tag.Metadata.Uid + "/" - require.NotEmpty(t, checkpointObjectKeys(t, ctx, objects, bucket, tagPrefix), "copy failure must leave actual destination objects") - replacement := NewService(newStore(), testAuthorizer{}, client, nil) - verificationErr := replacement.recoverOperation(ctx, op) - require.Error(t, verificationErr, "serialized AlreadyExists must verify and fail the partial copy") - pending, err := store.GetCheckpointOperation(ctx, cp.Id, "alice") - require.NoError(t, err) - require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_FAILED, pending.Checkpoint.State) - require.NotNil(t, pending.Checkpoint.Failure) - require.Equal(t, tag.Metadata.Uid, pending.TagUID) - _, err = store.BeginSessionOperation(ctx, source.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE, false) - require.NoError(t, err, "copying has stopped, so failed creation must release source admission") - require.Error(t, replacement.Delete(owner, cp.Id), "the real object deletion must fail while storage denies it") - deleting, err := store.GetCheckpointOperation(ctx, cp.Id, "alice") - require.NoError(t, err) - require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_DELETING, deleting.Checkpoint.State) - require.Equal(t, tag.Metadata.Uid, deleting.TagUID) - _, err = objects.DeleteBucketPolicy(ctx, &s3.DeleteBucketPolicyInput{Bucket: aws.String(bucket)}) - require.NoError(t, err) - replacement = NewService(newStore(), testAuthorizer{}, client, nil) - deleteCheckpoint(t, replacement, cp.Id) - require.Empty(t, checkpointObjectKeys(t, ctx, objects, bucket, tagPrefix), "all partially copied objects must be collected") - }) - - t.Run("Given a live remote copy after caller cancellation, when another replica retries Create, assert that busy work stays pending and a later retry adopts its UID", func(t *testing.T) { - objects, bucket := checkpointObjectStore(t, config.ObjectStoreEndpoint, config.ObjectStoreAccessKey, config.ObjectStoreSecretKey) - privateTemplate := proto.Clone(template).(*ateapipb.ActorTemplate) - privateTemplate.Metadata = &ateapipb.ResourceMetadata{Atespace: config.Atespace, Name: "checkpoint-concurrent-" + uuid.NewString()} - privateTemplate.Status = nil - privateTemplate.SnapshotConfig.StorageLocation = "gs://" + bucket - privateTemplate, err := client.CreateActorTemplate(ctx, privateTemplate) - require.NoError(t, err) - t.Cleanup(func() { - cleanup, stop := context.WithTimeout(context.Background(), time.Minute) - defer stop() - require.NoError(t, client.DeleteActorTemplate(cleanup, config.Atespace, privateTemplate.Metadata.Name, privateTemplate.Metadata.Uid)) - }) - store, source, head := newSourceForTemplate(t, privateTemplate) - actor, err := client.GetActor(ctx, config.Atespace, substrate.ActorName(source.Id)) - require.NoError(t, err) - snapshot, err := url.Parse(substrate.SnapshotURI(substrate.LatestActorSnapshot(actor))) - require.NoError(t, err) - require.Equal(t, bucket, snapshot.Host) - // Many real objects keep copying observable after the first destination - // object appears. A completed setup copy must never silently pass this test. - var uploads errgroup.Group - uploads.SetLimit(8) - for i := range 4096 { - uploads.Go(func() error { - _, err := objects.PutObject(ctx, &s3.PutObjectInput{Bucket: aws.String(bucket), - Key: aws.String(fmt.Sprintf("%s/concurrent-copy/%04d", strings.Trim(snapshot.Path, "/"), i)), - Body: strings.NewReader(strings.Repeat("copy-proof", 1024))}) - return err - }) - } - require.NoError(t, uploads.Wait()) - request := uuid.NewString() - service := NewService(store, testAuthorizer{}, client, nil) - cp, _, err := service.beginCreate(owner, source.Id, request, head) - require.NoError(t, err) - cleanupTag(t, cp.Id) - remoteCtx, stopRemote := context.WithTimeout(context.WithoutCancel(ctx), 2*time.Minute) - defer stopRemote() - remote := &disconnectedSubstrateCopy{Client: client, remote: remoteCtx, done: make(chan struct{})} - callerCtx, cancelCaller := context.WithCancel(owner) - defer cancelCaller() - callerDone := make(chan error, 1) - go func() { - _, err := NewService(store, testAuthorizer{}, remote, nil).Create(callerCtx, source.Id, request, head) - callerDone <- err - }() - var copying *ateapipb.Tag - require.Eventually(t, func() bool { - var err error - copying, err = client.GetTag(ctx, config.Atespace, tagName(cp.Id)) - if err != nil || substrate.SnapshotURI(copying.GetStatus().GetSnapshot()) != "" { - return false - } - prefix := "atespaces/" + config.Atespace + "/tags/" + copying.Metadata.Uid + "/" - page, err := objects.ListObjectsV2(ctx, &s3.ListObjectsV2Input{Bucket: aws.String(bucket), Prefix: aws.String(prefix), MaxKeys: aws.Int32(1)}) - return err == nil && len(page.Contents) > 0 - }, 30*time.Second, 10*time.Millisecond, "a real pending copy with destination objects was not observed") - cancelCaller() - require.ErrorIs(t, <-callerDone, context.Canceled) - observed := &observedSubstrateCopy{Client: client, entered: make(chan struct{})} - type createResult struct { - checkpoint *apiv1.Checkpoint - err error - } - result := make(chan createResult, 1) - go func() { - ready, err := NewService(newStore(), testAuthorizer{}, observed, nil).Create(owner, source.Id, request, head) - result <- createResult{checkpoint: ready, err: err} - }() - select { - case <-observed.entered: - case <-ctx.Done(): - t.Fatal(ctx.Err()) - } - select { - case <-remote.done: - t.Fatal("the remote copy finished before takeover; concurrency was not exercised") - default: - } - pending, err := store.GetCheckpointOperation(ctx, cp.Id, "alice") - require.NoError(t, err) - require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_CREATING, pending.Checkpoint.State) - select { - case outcome := <-result: - require.Equal(t, codes.Aborted, status.Code(outcome.err), "a busy backend lease is retryable, not proof of a failed copy") - require.Nil(t, outcome.checkpoint) - case <-ctx.Done(): - t.Fatal(ctx.Err()) - } - require.Zero(t, observed.tagReads.Load(), "a busy copy must not be inspected as an abandoned pending tag") - pending, err = store.GetCheckpointOperation(ctx, cp.Id, "alice") - require.NoError(t, err) - require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_CREATING, pending.Checkpoint.State) - require.Nil(t, pending.Checkpoint.Failure) - _, err = store.BeginSessionOperation(ctx, source.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE, false) - require.ErrorIs(t, err, database.ErrConflict, "live copying must retain source admission") - select { - case <-remote.done: - case <-ctx.Done(): - t.Fatal(ctx.Err()) - } - require.NoError(t, remote.err) - ready, err := NewService(newStore(), testAuthorizer{}, client, nil).Create(owner, source.Id, request, head) - require.NoError(t, err) - require.Equal(t, cp.Id, ready.Id) - require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_READY, ready.State) - published, err := store.GetCheckpointOperation(ctx, cp.Id, "alice") - require.NoError(t, err) - require.Equal(t, copying.Metadata.Uid, published.TagUID) - require.Equal(t, remote.tag.Metadata.Uid, published.TagUID) - tag, err := client.GetTag(ctx, config.Atespace, tagName(cp.Id)) - require.NoError(t, err) - require.NotEmpty(t, substrate.SnapshotURI(tag.GetStatus().GetSnapshot())) - t.Log("replacement received Aborted while the real copy was running; a later retry adopted the same completed tag UID") - deleteCheckpoint(t, NewService(newStore(), testAuthorizer{}, client, nil), cp.Id) - }) - - t.Run("Given a lost delete response, when recovery restarts, assert that an absent real tag completes deletion", func(t *testing.T) { - store, source, head := newSource(t) - service := NewService(store, testAuthorizer{}, client, nil) - cp, err := service.Create(owner, source.Id, uuid.NewString(), head) - require.NoError(t, err) - cleanupTag(t, cp.Id) - _, err = service.beginDelete(owner, cp.Id) - require.NoError(t, err) - op, err := store.GetCheckpointOperation(ctx, cp.Id, "alice") - require.NoError(t, err) - lost := NewService(store, testAuthorizer{}, &faultedSubstrateTags{Client: client, loseDeleteResponse: true}, nil) - require.Error(t, lost.recoverOperation(ctx, op)) - _, err = client.GetTag(ctx, config.Atespace, tagName(cp.Id)) - require.Equal(t, codes.NotFound, status.Code(err)) - deleteCheckpoint(t, NewService(newStore(), testAuthorizer{}, client, nil), cp.Id) - }) -} - -type faultedSubstrateTags struct { - *substrate.Client - loseCreateResponse bool - loseDeleteResponse bool - failDelete bool - afterCreate func() error -} - -func (f *faultedSubstrateTags) CreateTag(ctx context.Context, space, name, actor string) (*ateapipb.Tag, error) { - tag, err := f.Client.CreateTag(ctx, space, name, actor) - if err != nil { - return nil, err - } - if f.afterCreate != nil { - if err := f.afterCreate(); err != nil { - return nil, err - } - } - if f.loseCreateResponse { - return nil, status.Error(codes.Unavailable, "injected loss of successful CreateTag response") - } - return tag, nil -} - -func (f *faultedSubstrateTags) DeleteTag(ctx context.Context, space, name, uid string) error { - if f.failDelete { - return status.Error(codes.Unavailable, "injected failure to deliver DeleteTag") - } - if err := f.Client.DeleteTag(ctx, space, name, uid); err != nil { - return err - } - if f.loseDeleteResponse { - return status.Error(codes.Unavailable, "injected loss of successful DeleteTag response") - } - return nil -} - -var _ tagClient = (*faultedSubstrateTags)(nil) - -// The client loses cancellation delivery to the server: the caller exits while -// the actual Substrate RPC keeps copying. No backend state or reply is invented. -type disconnectedSubstrateCopy struct { - *substrate.Client - remote context.Context - done chan struct{} - tag *ateapipb.Tag - err error -} - -func (c *disconnectedSubstrateCopy) CreateTag(ctx context.Context, space, name, actor string) (*ateapipb.Tag, error) { - go func() { - c.tag, c.err = c.Client.CreateTag(c.remote, space, name, actor) - close(c.done) - }() - select { - case <-ctx.Done(): - return nil, ctx.Err() - case <-c.done: - return c.tag, c.err - } -} - -type observedSubstrateCopy struct { - *substrate.Client - entered chan struct{} - tagReads atomic.Int32 -} - -func (c *observedSubstrateCopy) CreateTag(ctx context.Context, space, name, actor string) (*ateapipb.Tag, error) { - close(c.entered) - return c.Client.CreateTag(ctx, space, name, actor) -} - -func (c *observedSubstrateCopy) GetTag(ctx context.Context, space, name string) (*ateapipb.Tag, error) { - c.tagReads.Add(1) - return c.Client.GetTag(ctx, space, name) -} - -var _ tagClient = (*disconnectedSubstrateCopy)(nil) -var _ tagClient = (*observedSubstrateCopy)(nil) diff --git a/go/core/internal/service/checkpoint/recovery_test.go b/go/core/internal/service/checkpoint/recovery_test.go index 623fc4418d..023edc9917 100644 --- a/go/core/internal/service/checkpoint/recovery_test.go +++ b/go/core/internal/service/checkpoint/recovery_test.go @@ -24,7 +24,6 @@ import ( 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" - "github.com/kagent-dev/kagent/go/core/pkg/auth" ) // lockedTags models Substrate's tag lease independently of a request or worker. @@ -132,17 +131,6 @@ func recoverySourceForActor(t *testing.T, store *database.Client, actor *ateapip return session, string(task.ID) } -func waitCheckpointPhase(t *testing.T, store *database.Client, id string, phase apiv1.CheckpointState) *database.CheckpointOperation { - t.Helper() - var op *database.CheckpointOperation - require.Eventually(t, func() bool { - var err error - op, err = store.GetCheckpointOperation(t.Context(), id, "alice") - return err == nil && op.Checkpoint.State == phase - }, 5*time.Second, 10*time.Millisecond) - return op -} - type failedCheckpointPublish struct { *database.Client entered chan struct{} @@ -251,18 +239,6 @@ func (s *Service) beginCreate(ctx context.Context, source, request, head string) return op.Checkpoint, op.Creator, nil } -func (s *Service) beginDelete(ctx context.Context, id string) (string, error) { - ctx, cancel := context.WithTimeout(ctx, database.RuntimeOperationTimeout) - defer cancel() - owner, err := s.authorize(ctx, auth.VerbDelete, auth.Resource{Type: "Checkpoint", Name: id}) - if err != nil { - return "", err - } - executor := uuid.New() - _, err = s.store.BeginDeleteSessionCheckpoint(ctx, id, owner, executor) - return owner, errors.Join(err, s.release(ctx, id, executor)) -} - 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) { diff --git a/go/core/internal/service/checkpoint/service.go b/go/core/internal/service/checkpoint/service.go index 156d36bc86..7a1b67ebb0 100644 --- a/go/core/internal/service/checkpoint/service.go +++ b/go/core/internal/service/checkpoint/service.go @@ -12,7 +12,6 @@ import ( "google.golang.org/genproto/googleapis/rpc/errdetails" "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" @@ -33,6 +32,7 @@ type store interface { 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) @@ -100,7 +100,7 @@ func (s *Service) Create(ctx context.Context, sessionID, requestID, expectedHead if op.Checkpoint.Id != id.String() { result, err = s.claimAndExecute(ctx, op) } else { - result, err = s.execute(ctx, op, executor) + result, err = s.execute(ctx, op, executor, true) } if err == nil && result == nil { return nil, serviceerrors.NewNotFound("Checkpoint no longer exists", nil) @@ -124,7 +124,7 @@ func (s *Service) reserveCreate(ctx context.Context, sessionID, requestID, expec return nil, err } checkpoint, snapshot, err := s.store.ReserveSessionCheckpoint(ctx, &apiv1.Checkpoint{Id: id, SessionId: sessionID, HeadTaskId: expectedHeadTaskID}, userID, requestID, executor) - if err != nil { + 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 { @@ -162,19 +162,20 @@ 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) (*ateapipb.Tag, error) { +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) - // v1.0.0 takes Actor and Tag leases before its workflow checks. Interpreting - // terminal replies requires the backend guarantee that earlier external - // writes cannot outlive serialization, including cancellation/lease loss; - // see the architecture document's required backend contract. GetTag alone - // cannot establish that a pending copy stopped. + // 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 { switch status.Code(err) { @@ -255,7 +256,7 @@ func (s *Service) Get(ctx context.Context, checkpointID string) (*apiv1.Checkpoi if err != nil { return nil, serviceerrors.NewInternal("Failed to get checkpoint", err) } - return publicCheckpoint(checkpoint), nil + return checkpoint, nil } func (s *Service) List(ctx context.Context, request ListRequest) (ListResult, error) { @@ -283,9 +284,6 @@ func (s *Service) List(ctx context.Context, request ListRequest) (ListResult, er return ListResult{}, serviceerrors.NewInternal("Failed to list checkpoints", err) } result := ListResult{Checkpoints: rows[:min(len(rows), pageSize)]} - for i, checkpoint := range result.Checkpoints { - result.Checkpoints[i] = publicCheckpoint(checkpoint) - } if len(rows) > pageSize { result.NextPageToken = encodePageToken(rows[pageSize-1].GetId()) } @@ -306,7 +304,6 @@ func (s *Service) Delete(ctx context.Context, checkpointID string) (err error) { executor := uuid.New() op, err := s.store.BeginDeleteSessionCheckpoint(ctx, checkpointID, userID, executor) if err != nil { - err = errors.Join(err, s.release(ctx, checkpointID, executor)) if errors.Is(err, database.ErrConflict) { _, err = s.claimAndExecute(ctx, &database.CheckpointOperation{Checkpoint: &apiv1.Checkpoint{Id: checkpointID}, Creator: userID}) return err @@ -314,14 +311,38 @@ func (s *Service) Delete(ctx context.Context, checkpointID string) (err error) { 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) + _, 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 { @@ -332,26 +353,13 @@ func (s *Service) finishDelete(ctx context.Context, op *database.CheckpointOpera // 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) || - (op.Checkpoint.Failure == nil && 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 publicOperationError(fmt.Errorf("delete checkpoint snapshot tag: %w", 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 } @@ -512,29 +520,16 @@ func createOutcome(checkpoint *apiv1.Checkpoint) (*apiv1.Checkpoint, error) { case apiv1.CheckpointState_CHECKPOINT_STATE_READY: return checkpoint, nil case apiv1.CheckpointState_CHECKPOINT_STATE_FAILED: - failure, err := status.New(codes.FailedPrecondition, - "Checkpoint creation failed; delete the failed checkpoint before creating a new request").WithDetails( - &errdetails.ErrorInfo{Domain: "kagent.dev", Reason: "KAGENT_CHECKPOINT_FAILED"}, - &errdetails.ResourceInfo{ResourceType: "kagent.api.v1.Checkpoint", ResourceName: checkpoint.Id}, - ) - if err != nil { - return nil, serviceerrors.NewInternal("Failed to report checkpoint failure", err) - } - return nil, failure.Err() + return nil, status.Error(codes.FailedPrecondition, "Checkpoint creation failed") default: return nil, serviceerrors.NewAborted("Checkpoint operation is still pending; retry the same request", nil) } } -// Failed records from older controllers can contain raw backend errors. Public -// reads expose a generic failure without rewriting the retained operation data. -func publicCheckpoint(checkpoint *apiv1.Checkpoint) *apiv1.Checkpoint { - if checkpoint.State != apiv1.CheckpointState_CHECKPOINT_STATE_FAILED { - return checkpoint - } - result := proto.CloneOf(checkpoint) - result.Failure = &apiv1.Failure{Reason: "SnapshotTagFailed", Message: checkpointFailureMessage} - return result +// 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 diff --git a/go/core/internal/service/checkpoint/service_test.go b/go/core/internal/service/checkpoint/service_test.go index 97fb38ed51..f7bd8f1f52 100644 --- a/go/core/internal/service/checkpoint/service_test.go +++ b/go/core/internal/service/checkpoint/service_test.go @@ -894,3 +894,8 @@ func (*testStore) ReleaseCheckpointOperation(context.Context, string, uuid.UUID) 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/test/e2e/README.md b/go/core/test/e2e/README.md index 3d394ad8e5..34e3c4d8d2 100644 --- a/go/core/test/e2e/README.md +++ b/go/core/test/e2e/README.md @@ -136,29 +136,6 @@ checks also run in parallel. Trace cases share the receiver but select spans by their own trace IDs; the receiver is cleared once before starting those cases. Controller restart cases stay sequential with respect to the rest of the suite. -`TestCheckpointReplayAfterControllerRestart` loses a successful public checkpoint -response, replaces the controller once for the harness matrix, retries the same -request, and forks the retained snapshot to verify conversation continuity. -The response loss is caller-side after success; interrupted tag copying, -publication and deletion are exercised with PostgreSQL-backed workflow fault -tests. This test does not inject live Substrate packet loss. - -`bash core/test/e2e/checkpoint-recovery.sh` (from `go/`) runs the checkpoint -lifecycle service against real PostgreSQL and the installed Substrate. Set -`KAGENT_E2E_RUNTIME_IMAGE` to the digest-pinned Go ADK image. The script reads the -test controller's projected identity into temporary private files and forwards -the Substrate service; it leaves controller configuration unchanged. Both CI -runtime jobs run it before the public API E2E suite, and a configured environment -must fail rather than skip on missing prerequisites. It owns its ActorTemplate, -Actors and Tags. Cases cover abandoned reservations, lost successful tag replies, -copy completion before database publication, -a source deleted before tag creation, a failed retained copy after a real source -change, denied deletion of actual partially copied objects followed by storage -repair, and lost successful deletion replies. Transport faults are injected at -the service's client boundary around real calls; backend responses and database -phases are not fabricated. The deterministic PostgreSQL tests separately cover -pending-copy leases and overlapping/expired executors. The partial-copy case applies a real object-store policy denial; transport loss -remains injected at the client boundary. The scheduled-run restart test starts all selected harness executions before one shared controller restart, then releases their model responses and checks recovery in parallel. Each harness retains its own Session, task, and prompt-count assertions. diff --git a/go/core/test/e2e/checkpoint-recovery.sh b/go/core/test/e2e/checkpoint-recovery.sh deleted file mode 100644 index a94f0b0099..0000000000 --- a/go/core/test/e2e/checkpoint-recovery.sh +++ /dev/null @@ -1,65 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -# Run from any directory against the disposable E2E installation. Reuse its -# controller identity; private certificate material never enters test output. -cd "$(dirname "${BASH_SOURCE[0]}")/../../.." -: "${KAGENT_E2E_RUNTIME_IMAGE:?set the digest-pinned Go ADK image}" -checkpoint_tmp=$(mktemp -d) -chmod 700 "$checkpoint_tmp" -forward_pids=() -cleanup() { - for pid in "${forward_pids[@]:-}"; do - if [[ -n "$pid" ]]; then - kill "$pid" 2>/dev/null || true - wait "$pid" 2>/dev/null || true - fi - done - rm -rf "$checkpoint_tmp" -} -trap cleanup EXIT -kubectl -n kagent exec deployment/kagent-controller -- \ - cat /run/substrate-servicedns/trust-bundle.pem > "$checkpoint_tmp/ca.pem" -kubectl -n kagent exec deployment/kagent-controller -- \ - cat /run/substrate-podidentity/credential-bundle.pem > "$checkpoint_tmp/client.pem" -chmod 600 "$checkpoint_tmp/client.pem" -forward_port() { - local service=$1 remote_port=$2 forward_pid log - log="$checkpoint_tmp/$service-forward.log" - kubectl -n ate-system port-forward "service/$service" ":$remote_port" > "$log" 2>&1 & - forward_pid=$! - forward_pids+=("$forward_pid") - checkpoint_forwarded_port= - for _ in $(seq 1 100); do - checkpoint_forwarded_port=$(awk '$1 == "Forwarding" && $3 ~ /^127[.]0[.]0[.]1:/ { split($3, address, ":"); print address[2]; exit }' "$log") - if [[ -n "$checkpoint_forwarded_port" ]]; then return; fi - if ! kill -0 "$forward_pid" 2>/dev/null; then - cat "$log" >&2 - return 1 - fi - sleep 0.1 - done - echo "Timed out forwarding $service" >&2 - return 1 -} -forward_port api 443 -api_port=$checkpoint_forwarded_port -forward_port rustfs 9000 -object_store_port=$checkpoint_forwarded_port -# The disposable Substrate installation provides literal RustFS credentials. -# Keep them in private files rather than command arguments or test output. -kubectl -n ate-system get deployment/rustfs -o json > "$checkpoint_tmp/object-store.json" -jq -n --slurpfile store "$checkpoint_tmp/object-store.json" \ - --arg endpoint "127.0.0.1:$api_port" --arg ca "$checkpoint_tmp/ca.pem" \ - --arg objectStoreEndpoint "http://127.0.0.1:$object_store_port" \ - --arg certificate "$checkpoint_tmp/client.pem" --arg image "$KAGENT_E2E_RUNTIME_IMAGE" \ - '{endpoint: $endpoint, serverName: "api.ate-system.svc", caFile: $ca, - clientCertFile: $certificate, atespace: "kagent", image: $image, - objectStoreEndpoint: $objectStoreEndpoint, - objectStoreAccessKey: ([$store[0].spec.template.spec.containers[].env[] | select(.name == "RUSTFS_ACCESS_KEY").value][0] // error("missing RustFS access key")), - objectStoreSecretKey: ([$store[0].spec.template.spec.containers[].env[] | select(.name == "RUSTFS_SECRET_KEY").value][0] // error("missing RustFS secret key"))}' \ - > "$checkpoint_tmp/config.json" -chmod 600 "$checkpoint_tmp/config.json" -KAGENT_E2E_CHECKPOINT_SUBSTRATE_CONFIG="$checkpoint_tmp/config.json" \ - go test -race ./core/internal/service/checkpoint \ - -run '^TestCheckpointRecoveryWithSubstrate$' -v -count=1 -timeout=10m diff --git a/go/core/test/upgrade/checkpoint_drain_test.go b/go/core/test/upgrade/checkpoint_drain_test.go deleted file mode 100644 index e24fdd0e09..0000000000 --- a/go/core/test/upgrade/checkpoint_drain_test.go +++ /dev/null @@ -1,145 +0,0 @@ -package upgrade - -import ( - "context" - "encoding/json" - "fmt" - "os" - "path/filepath" - "strings" - "testing" - "time" - - "github.com/google/uuid" - "github.com/jackc/pgx/v5/pgxpool" - "github.com/stretchr/testify/require" - appsv1 "k8s.io/api/apps/v1" - corev1 "k8s.io/api/core/v1" - metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" - "k8s.io/apimachinery/pkg/util/wait" - "k8s.io/utils/ptr" - - kagentenv "github.com/kagent-dev/kagent/go/core/pkg/env" - "github.com/kagent-dev/kagent/go/core/pkg/postgres" -) - -// Exercise the process-drain boundary in an isolated namespace. This models -// legacy process lifetime; TestUpgrade uses the same drain with real old binaries. -func TestUpgradeWaitsForTerminatingCheckpointWriters(t *testing.T) { - if kagentenv.E2ERunUpgradeTests.Get() != "true" { - t.Skip("requires a disposable upgrade cluster") - } - env := loadUpgradeEnv(t) - env.namespace = "checkpoint-upgrade-" + uuid.NewString() - kubectl(t, env, time.Minute, "create", "namespace", env.namespace) - t.Cleanup(func() { - _, err := kubectlOutputContext(context.WithoutCancel(t.Context()), env, time.Minute, - "delete", "namespace", env.namespace, "--wait=false") - require.NoError(t, err) - }) - labels := map[string]string{"app.kubernetes.io/name": "kagent", "app.kubernetes.io/component": "controller"} - image := kagentenv.E2ERuntimeImage.Get() - if image == "" { - image = "busybox:1.36" - } - deployment := &appsv1.Deployment{ - TypeMeta: metav1.TypeMeta{APIVersion: "apps/v1", Kind: "Deployment"}, - ObjectMeta: metav1.ObjectMeta{Name: "kagent-controller", Namespace: env.namespace}, - Spec: appsv1.DeploymentSpec{Replicas: ptr.To[int32](2), Selector: &metav1.LabelSelector{MatchLabels: labels}, - Template: corev1.PodTemplateSpec{ObjectMeta: metav1.ObjectMeta{Labels: labels}, Spec: corev1.PodSpec{ - TerminationGracePeriodSeconds: ptr.To[int64](1), - Containers: []corev1.Container{{Name: controllerContainer, Image: image, - Command: []string{"/bin/sh", "-c", "trap 'exit 0' TERM; while :; do sleep 1; done"}}}, - }}}, - } - data, err := json.Marshal(deployment) - require.NoError(t, err) - manifest := filepath.Join(t.TempDir(), "legacy-writers.json") - require.NoError(t, os.WriteFile(manifest, data, 0o600)) - kubectl(t, env, time.Minute, "apply", "-f", manifest) - kubectl(t, env, 2*time.Minute, "rollout", "status", "deployment/kagent-controller", "-n", env.namespace, "--timeout=2m") - pods := podNamesForSelector(t, env, controllerSelector) - require.Len(t, pods, 2) - held := pods[0] - // Hold the Pod object after its process becomes unready. A scale request - // or readiness transition alone must never authorize starting new writers. - kubectl(t, env, time.Minute, "patch", "pod", held, "-n", env.namespace, - "--type=merge", "-p", `{"metadata":{"finalizers":["checkpoint.kagent.dev/upgrade-test"]}}`) - released := false - t.Cleanup(func() { - if !released { - _, _ = kubectlOutputContext(context.WithoutCancel(t.Context()), env, time.Minute, "patch", "pod", held, "-n", env.namespace, - "--type=merge", "-p", `{"metadata":{"finalizers":[]}}`) - } - }) - drained := make(chan error, 1) - go func() { drained <- drainLegacyCheckpointWriters(t, env) }() - require.Eventually(t, func() bool { - out, err := kubectlOutput(t, env, 10*time.Second, "get", "pod", held, "-n", env.namespace, - "-o", `jsonpath={.metadata.deletionTimestamp}{"|"}{.status.containerStatuses[?(@.name=="controller")].ready}`) - deleting, readiness, ok := strings.Cut(out, "|") - return err == nil && ok && deleting != "" && readiness == "false" - }, time.Minute, 100*time.Millisecond) - select { - case err := <-drained: - t.Fatalf("upgrade boundary passed a terminating legacy writer: %v", err) - case <-time.After(200 * time.Millisecond): - } - kubectl(t, env, time.Minute, "patch", "pod", held, "-n", env.namespace, - "--type=merge", "-p", `{"metadata":{"finalizers":[]}}`) - released = true - select { - case err := <-drained: - require.NoError(t, err) - case <-time.After(time.Minute): - t.Fatal("upgrade boundary did not open after all legacy writers exited") - } - require.Empty(t, podNamesForSelector(t, env, controllerSelector)) - scaleController(t, env, 1) - newPods := podNamesForSelector(t, env, controllerSelector) - require.Len(t, newPods, 1) - require.NotContains(t, pods, newPods[0], "new writers must start only after old Pods are gone") -} - -func drainLegacyCheckpointWriters(t *testing.T, env upgradeEnv) error { - t.Helper() - if _, err := kubectlOutput(t, env, time.Minute, "scale", "deployment/kagent-controller", - "-n", env.namespace, "--replicas=0"); err != nil { - return fmt.Errorf("stop legacy checkpoint writers: %w", err) - } - if _, err := kubectlOutput(t, env, 2*time.Minute, "rollout", "status", "deployment/kagent-controller", - "-n", env.namespace, "--timeout=2m"); err != nil { - return fmt.Errorf("observe legacy controller scale-down: %w", err) - } - return wait.PollUntilContextTimeout(t.Context(), 500*time.Millisecond, 2*time.Minute, true, func(context.Context) (bool, error) { - pods, err := podNamesForSelectorE(t, env, controllerSelector) - return err == nil && len(pods) == 0, err - }) -} - -func requireLegacyCheckpointRollbackReady(t *testing.T, env upgradeEnv) { - t.Helper() - port, stop := startPortForward(t, env, postgresServiceName, 5432) - defer stop() - config, err := postgres.PoolConfig(postgresURL(t, env, port), postgresRole) - require.NoError(t, err) - config.ConnConfig.RuntimeParams["search_path"] = kagentSchema - pool, err := pgxpool.NewWithConfig(t.Context(), config) - require.NoError(t, err) - defer pool.Close() - require.NoError(t, checkLegacyCheckpointRollback(t.Context(), pool)) -} - -// Legacy deletion requires a completed URI and cannot collect retained partial -// tags. Keep this guard shared with the PostgreSQL rollback regression. -func checkLegacyCheckpointRollback(ctx context.Context, pool *pgxpool.Pool) error { - var blockers int - if err := pool.QueryRow(ctx, `SELECT count(*) FROM session_checkpoint - WHERE state IN ('CREATING', 'DELETING') OR (state = 'FAILED' AND tag_uid <> '')`).Scan(&blockers); err != nil { - return fmt.Errorf("check checkpoint rollback compatibility: %w", err) - } - if blockers != 0 { - return fmt.Errorf("finish pending operations and delete retained failed checkpoints before restoring a legacy writer: %d incompatible checkpoints", blockers) - } - return nil -} diff --git a/go/core/test/upgrade/checkpoint_rollback_test.go b/go/core/test/upgrade/checkpoint_rollback_test.go deleted file mode 100644 index b601ebcd5a..0000000000 --- a/go/core/test/upgrade/checkpoint_rollback_test.go +++ /dev/null @@ -1,102 +0,0 @@ -package upgrade - -import ( - "context" - "crypto/sha256" - "testing" - - "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" - - 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" -) - -// Exercise the actual rollback guard with valid persisted checkpoints. The -// partial tag is a storage fixture; real object cleanup has its own Substrate test. -func TestUpgradeRejectsIncompatibleCheckpointRollback(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) - store := database.NewClient(pool, "public") - source, head := checkpointRollbackSource(t, store) - ctx, cancel := context.WithTimeout(t.Context(), database.RuntimeOperationTimeout) - defer cancel() - require.NoError(t, checkLegacyCheckpointRollback(ctx, pool)) - executor := uuid.New() - checkpoint, _, err := store.ReserveSessionCheckpoint(ctx, - &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: source.Id, HeadTaskId: head}, "alice", uuid.NewString(), executor) - require.NoError(t, err) - require.Error(t, checkLegacyCheckpointRollback(ctx, pool), "pending creation must block rollback") - failed, err := store.FinalizeClaimedSessionCheckpoint(ctx, checkpoint, executor, "partial-tag-uid", "", "copy failed") - require.NoError(t, err) - retained, err := store.GetCheckpointOperation(ctx, checkpoint.Id, "alice") - require.NoError(t, err) - require.Equal(t, "partial-tag-uid", retained.TagUID) - require.Equal(t, "s3://snapshots/source", retained.Snapshot.URI, "the failed copy has no completed tag URI") - require.Error(t, checkLegacyCheckpointRollback(ctx, pool), "a retained FAILED tag must block rollback") - deletion := uuid.New() - _, err = store.BeginDeleteSessionCheckpoint(ctx, failed.Id, "alice", deletion) - require.NoError(t, err) - require.Error(t, checkLegacyCheckpointRollback(ctx, pool), "pending cleanup must block rollback") - require.NoError(t, store.DeleteClaimedSessionCheckpoint(ctx, failed.Id, "alice", deletion)) - require.NoError(t, checkLegacyCheckpointRollback(ctx, pool), "completed cleanup permits rollback") - - for _, outcome := range []string{"ready", "failed without a tag"} { - t.Run("Given a checkpoint "+outcome+", when rollback is checked, assert that no incompatible work remains", func(t *testing.T) { - executor := uuid.New() - checkpoint, _, err := store.ReserveSessionCheckpoint(ctx, - &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: source.Id, HeadTaskId: head}, "alice", uuid.NewString(), executor) - require.NoError(t, err) - uid, uri, failure := "complete-tag-uid", "s3://tags/complete", "" - if outcome != "ready" { - uid, uri, failure = "", "", "source was deleted before copying" - } - _, err = store.FinalizeClaimedSessionCheckpoint(ctx, checkpoint, executor, uid, uri, failure) - require.NoError(t, err) - require.NoError(t, checkLegacyCheckpointRollback(ctx, pool)) - }) - } -} - -func checkpointRollbackSource(t *testing.T, store *database.Client) (*apiv1.Session, string) { - t.Helper() - ctx := t.Context() - require.NoError(t, store.UpsertAgentDefinition(ctx, database.AgentDefinition{ - Namespace: "team-a", AgentName: "assistant", AgentUID: "agent-uid", DesiredRevision: "revision"})) - _, err := store.RecordRuntimeRevision(ctx, database.RuntimeRevision{ - Revision: "revision", Namespace: "team-a", AgentName: "assistant", AgentUID: "agent-uid", - SourceSnapshot: []byte("{}"), AgentCard: &a2apb.AgentCard{}, EgressDestinations: []string{}, - ActorTemplateAtespace: "team-a", ActorTemplateName: "runtime", ActorTemplateUID: "template-uid", - }, database.RecordRevisionOptions{Ready: true}) - require.NoError(t, err) - session, _, err := store.CreateSession(ctx, &apiv1.Session{Id: uuid.NewString(), Creator: "alice", - Agent: &apiv1.AgentProvenance{Ref: &apiv1.ResourceReference{Namespace: "team-a", 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, "runtime", "actor-uid", "") - 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: "team-a", URI: "s3://snapshots/source", Fidelity: "SNAPSHOT_FIDELITY_VOLUMES"})) - return session, string(task.ID) -} diff --git a/go/core/test/upgrade/rolling_upgrade_test.go b/go/core/test/upgrade/rolling_upgrade_test.go index 6891508b48..a0172faa3a 100644 --- a/go/core/test/upgrade/rolling_upgrade_test.go +++ b/go/core/test/upgrade/rolling_upgrade_test.go @@ -31,8 +31,6 @@ func TestRollingUpgradeCompatibility(t *testing.T) { waitForReadyPods(t, env, postgresSelector, 3*time.Minute) waitForPostgresSchema(t, env, 3*time.Minute) - // Checkpoint mutations remain quiescent in this SQL compatibility lane. - // TestUpgrade separately drains legacy writers at the first checkpoint upgrade. // Run even when there is no migration delta: a rolling upgrade rolls the new // image regardless of migrations, so the deploy can still break for // non-schema reasons (a crashing new image, readiness, old pods against diff --git a/go/core/test/upgrade/upgrade_test.go b/go/core/test/upgrade/upgrade_test.go index 13d99ec1aa..4cfd1c7864 100644 --- a/go/core/test/upgrade/upgrade_test.go +++ b/go/core/test/upgrade/upgrade_test.go @@ -81,13 +81,9 @@ func TestUpgrade(t *testing.T) { var pgBaselineState postgresMigrationState var baselineVectorVersion int var cleanPreviousSchema string - var legacyCheckpointWriters bool if !t.Run("seed baseline data before upgrade", func(t *testing.T) { pgBaselineState = pgMigrationState(t, env) - legacyCheckpointWriters = pgQuery(t, env, `SELECT - EXISTS (SELECT 1 FROM information_schema.tables WHERE table_schema = current_schema() AND table_name = 'session_checkpoint') - AND NOT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_schema = current_schema() AND table_name = 'session_checkpoint' AND column_name = 'executor_id')`) == "t" baselineVectorVersion = pgTrackVersion(t, env, "vector_schema_migrations") cleanPreviousSchema = pgSchemaDump(t, env, kagentSchema) t.Logf("baseline Postgres schema_migrations version: %d vector=%d (target=%d)", @@ -146,17 +142,6 @@ func TestUpgrade(t *testing.T) { return } - if legacyCheckpointWriters && !t.Run("drain legacy checkpoint writers before starting the new controller", func(t *testing.T) { - // Old code can still mutate checkpoints on the expanded schema. Keep - // the earlier compatibility checks, then establish the first-upgrade - // boundary before any claim-aware controller is started. - require.NotEmpty(t, podNamesForSelector(t, env, controllerSelector)) - require.NoError(t, drainLegacyCheckpointWriters(t, env)) - require.Empty(t, podNamesForSelector(t, env, controllerSelector)) - }) { - return - } - if !t.Run("upgrade to current build", func(t *testing.T) { ctx, cancel := context.WithTimeout(t.Context(), 15*time.Minute) defer cancel() @@ -249,10 +234,6 @@ func TestUpgrade(t *testing.T) { } if !t.Run("roll back application to previous release", func(t *testing.T) { - if legacyCheckpointWriters { - require.NoError(t, drainLegacyCheckpointWriters(t, env)) - requireLegacyCheckpointRollbackReady(t, env) - } ctx, cancel := context.WithTimeout(t.Context(), 10*time.Minute) defer cancel() @@ -543,11 +524,8 @@ func kubectl(t *testing.T, env upgradeEnv, timeout time.Duration, args ...string func kubectlOutput(t *testing.T, env upgradeEnv, timeout time.Duration, args ...string) (string, error) { t.Helper() - return kubectlOutputContext(t.Context(), env, timeout, args...) -} -func kubectlOutputContext(ctx context.Context, env upgradeEnv, timeout time.Duration, args ...string) (string, error) { - ctx, cancel := context.WithTimeout(ctx, timeout) + ctx, cancel := context.WithTimeout(t.Context(), timeout) defer cancel() fullArgs := append([]string{"--context", env.kubeContext}, args...) diff --git a/ui/playwright/tests/chat/checkpoints.spec.ts b/ui/playwright/tests/chat/checkpoints.spec.ts index f7c02c7867..ec5318a116 100644 --- a/ui/playwright/tests/chat/checkpoints.spec.ts +++ b/ui/playwright/tests/chat/checkpoints.spec.ts @@ -13,221 +13,6 @@ const dividers = (page: import("@playwright/test").Page) => /** The seeded boundary's generated name: its conversation, joined to its turn. */ const SEEDED_NAME = `${instances.ready}-seed-task-1`; -async function deleteOutcome(page: import("@playwright/test").Page, outcome?: string) { - await page.evaluate((value) => { - const url = new URL(window.location.href); - if (value) url.searchParams.set("checkpointDelete", value); - else url.searchParams.delete("checkpointDelete"); - window.history.replaceState({}, "", url); - }, outcome); -} - -async function listUnavailable(page: import("@playwright/test").Page, unavailable: boolean) { - await page.evaluate((value) => { - const url = new URL(window.location.href); - if (value) url.searchParams.set("checkpointList", "unavailable"); - else url.searchParams.delete("checkpointList"); - window.history.replaceState({}, "", url); - }, unavailable); -} - -test.describe("chat: snapshot refresh failures", () => { - test.use({ expectedNoise: [/Could not (?:checkpoint the conversation|re-read the saved boundaries):/] }); - - test("chat: a repaired snapshot list ignores an older failed refresh", async ({ page }) => { - await page.goto(`${agentChat(instances.ready)}?checkpoint=failed`); - await expect(dividers(page)).toHaveCount(1); - await sendAndAwaitTurn(page, "Overlapping discovery requests after storage repair."); - await listUnavailable(page, true); - await page.getByTestId("chat-checkpoint").click(); - await expect(page.getByText("Could not load snapshots", { exact: true })).toBeVisible(); - await page.evaluate(() => { - const url = new URL(window.location.href); - url.searchParams.delete("checkpointList"); - url.searchParams.set("mock", "error"); - window.history.replaceState({}, "", url); - // Delay just this next fixture RPC, whose error scenario is captured - // before sleeping. The repaired retry must finish ahead of its failure. - const original = window.setTimeout; - window.setTimeout = Object.assign((handler: TimerHandler, delay?: number, ...args: unknown[]): number => { - if (delay === 450 && typeof handler === "function") { - window.setTimeout = original; - document.body.dataset.delayedSnapshotRead = "started"; - return original(() => { - handler(...args); - document.body.dataset.delayedSnapshotRead = "completed"; - }, 1800); - } - return original(handler, delay, ...args); - }, original); - }); - await page.getByRole("button", { name: "Retry", exact: true }).click(); - await expect(page.locator("body")).toHaveAttribute("data-delayed-snapshot-read", "started"); - await page.evaluate(() => { - const url = new URL(window.location.href); - url.searchParams.set("mock", "ok"); - window.history.replaceState({}, "", url); - }); - await page.getByRole("button", { name: "Retry", exact: true }).click(); - await expect(page.getByTestId("chat-failed-snapshots")).toBeVisible(); - await expect(page.getByText("Could not load snapshots", { exact: true })).toHaveCount(0); - await expect(page.locator("body")).toHaveAttribute("data-delayed-snapshot-read", "completed"); - // Let the old RPC rejection and React's error render reach the page. - await page.evaluate(() => new Promise(resolve => requestAnimationFrame(() => requestAnimationFrame(() => resolve())))); - await expect(page.getByText("Could not load snapshots", { exact: true })).toHaveCount(0); - }); - - test("chat: a failed creation followed by a list outage offers a repaired retry", async ({ page }) => { - await page.goto(`${agentChat(instances.ready)}?checkpoint=failed`); - await expect(dividers(page)).toHaveCount(1); - await sendAndAwaitTurn(page, "A retained failure whose first discovery read fails."); - await listUnavailable(page, true); - await page.getByTestId("chat-checkpoint").click(); - await expect(page.getByText("Could not load snapshots", { exact: true })).toBeVisible(); - await expect(page.getByTestId("chat-failed-snapshots")).toHaveCount(0); - await listUnavailable(page, false); - await page.getByRole("button", { name: "Retry", exact: true }).click(); - await expect(page.getByTestId("chat-failed-snapshots")).toBeVisible(); - await expect(page.getByText("Could not load snapshots", { exact: true })).toHaveCount(0); - await page.reload(); - await expect(page.getByTestId("chat-failed-snapshots")).toBeVisible(); - }); - - test("chat: deletion removes a failed entry even when the following list read fails", async ({ page }) => { - await page.goto(`${agentChat(instances.ready)}?checkpoint=failed`); - await expect(dividers(page)).toHaveCount(1); - await sendAndAwaitTurn(page, "Delete committed before discovery becomes unavailable."); - await page.getByTestId("chat-checkpoint").click(); - const notice = page.getByTestId("chat-failed-snapshots"); - await expect(notice).toBeVisible(); - await notice.locator('[data-testid^="chat-failed-snapshot-"]').click(); - await listUnavailable(page, true); - await page.getByTestId("snapshot-details-delete").click(); - await pressOnce(page.getByTestId("snapshot-details-delete-confirm")); - await expect(page.getByTestId("snapshot-details-body")).toHaveCount(0); - await expect(notice).toHaveCount(0); - await expect(page.getByText("Could not load snapshots", { exact: true })).toBeVisible(); - await listUnavailable(page, false); - await page.getByRole("button", { name: "Retry", exact: true }).click(); - await expect(page.getByText("Could not load snapshots", { exact: true })).toHaveCount(0); - await expect(notice).toHaveCount(0); - await page.reload(); - await expect(dividers(page)).toHaveCount(1); - await expect(notice).toHaveCount(0); - }); -}); - -test.describe("chat: snapshot delete failures", () => { - test.use({ expectedNoise: [/Could not (?:checkpoint the conversation|delete the checkpoint):/] }); - - test("chat: a refused transcript delete reports failure and can be retried", async ({ page }) => { - await page.goto(`${agentChat(instances.ready)}?checkpointDelete=denied`); - const line = dividers(page).first(); - await expect(line).toBeVisible(); - const id = (await line.getAttribute("data-testid"))!.replace("chat-checkpoint-mark-", ""); - await line.getByTestId(`chat-checkpoint-delete-${id}`).click(); - await pressOnce(page.getByTestId(`chat-checkpoint-delete-confirm-${id}`)); - await expect(page.getByText("Could not delete: Checkpoint operation could not complete", { exact: true })).toBeVisible(); - await expect(line).toBeVisible(); - await expect(page.getByText("Delete this snapshot?", { exact: true })).toBeHidden(); - await deleteOutcome(page); - await line.getByTestId(`chat-checkpoint-delete-${id}`).click(); - await pressOnce(page.getByTestId(`chat-checkpoint-delete-confirm-${id}`)); - await expect(dividers(page)).toHaveCount(0); - await page.reload(); - await expect(page.getByTestId("chat-checkpoint")).toBeVisible(); - await expect(dividers(page)).toHaveCount(0); - }); - - test("chat: a completed cleanup closes stale failed details on retry", async ({ page }) => { - await page.goto(`${agentChat(instances.ready)}?checkpoint=failed`); - await expect(dividers(page)).toHaveCount(1); - await sendAndAwaitTurn(page, "A new turn for cleanup after a lost deletion reply."); - await page.getByTestId("chat-checkpoint").click(); - const notice = page.getByTestId("chat-failed-snapshots"); - await expect(notice).toBeVisible(); - await notice.locator('[data-testid^="chat-failed-snapshot-"]').click(); - await deleteOutcome(page, "lost"); - await page.getByTestId("snapshot-details-delete").click(); - await pressOnce(page.getByTestId("snapshot-details-delete-confirm")); - await expect(page.getByTestId("snapshot-action-error")).toBeVisible(); - await expect(notice).toBeVisible(); - await expect(page.getByText("Delete this snapshot?", { exact: true })).toBeHidden(); - await page.getByTestId("snapshot-details-delete").click(); - await pressOnce(page.getByTestId("snapshot-details-delete-confirm")); - await expect(page.getByTestId("snapshot-details-body")).toHaveCount(0); - await expect(notice).toHaveCount(0); - await page.reload(); - await expect(dividers(page)).toHaveCount(1); - await expect(notice).toHaveCount(0); - }); - - test("chat: a referenced snapshot is retained when deletion returns not found", async ({ page }) => { - await page.goto(`${agentChat(instances.ready)}?checkpointDelete=referenced`); - const line = dividers(page).first(); - await expect(line).toBeVisible(); - await openSnapshot(page, line); - await page.getByTestId("snapshot-details-delete").click(); - await pressOnce(page.getByTestId("snapshot-details-delete-confirm")); - await expect(page.getByTestId("snapshot-action-error")).toBeVisible(); - await expect(page.getByTestId("snapshot-details-body")).toBeVisible(); - await expect(line).toBeVisible(); - await page.reload(); - await expect(dividers(page)).toHaveCount(1); - await openSnapshot(page, dividers(page).first()); - await expect(page.getByTestId("snapshot-details-state")).toHaveAttribute("data-state", "ready"); - }); -}); - -test.describe("chat: failed snapshots", () => { - test.use({ expectedNoise: [/Could not (?:checkpoint the conversation|delete the checkpoint):/] }); - - test("chat: a failed creation can be found, inspected and deleted after reload", async ({ page }) => { - await page.goto(`${agentChat(instances.ready)}?checkpoint=failed`); - await expect(dividers(page)).toHaveCount(1); - await sendAndAwaitTurn(page, "A new turn for a failed snapshot."); - await page.getByTestId("chat-checkpoint").click(); - - const notice = page.getByTestId("chat-failed-snapshots"); - await expect(notice).toBeVisible(); - // A failure is a cleanup record, not a saved history boundary. - await expect(dividers(page)).toHaveCount(1); - await expect(page.getByTestId("chat-checkpoint")).toBeEnabled(); - const failed = notice.locator('[data-testid^="chat-failed-snapshot-"]'); - const id = (await failed.getAttribute("data-testid"))!.replace("chat-failed-snapshot-", ""); - - await page.reload(); - await expect(page.getByTestId(`chat-failed-snapshot-${id}`)).toBeVisible(); - await page.getByTestId(`chat-failed-snapshot-${id}`).click(); - await expect(page.getByTestId("snapshot-details-state")).toHaveAttribute("data-state", "failed"); - await expect(page.getByTestId("snapshot-details-fields")).toContainText(id); - await expect(page.getByTestId("snapshot-details-fork")).toBeDisabled(); - await expect(page.getByTestId("snapshot-details-rename")).toBeDisabled(); - await page.evaluate(() => { - const url = new URL(window.location.href); - url.searchParams.set("checkpointDelete", "denied"); - window.history.replaceState({}, "", url); - }); - await page.getByTestId("snapshot-details-delete").click(); - await pressOnce(page.getByTestId("snapshot-details-delete-confirm")); - await expect(page.getByTestId("snapshot-action-error")).toBeVisible(); - await expect(page.getByTestId("snapshot-details-body")).toBeVisible(); - await expect(notice).toBeVisible(); - await page.evaluate(() => { - const url = new URL(window.location.href); - url.searchParams.delete("checkpointDelete"); - window.history.replaceState({}, "", url); - }); - await page.getByTestId("snapshot-details-delete").click(); - await pressOnce(page.getByTestId("snapshot-details-delete-confirm")); - await expect(notice).toHaveCount(0); - await expect(page.getByTestId("snapshot-details-body")).toHaveCount(0); - await page.reload(); - await expect(page.getByTestId("chat-failed-snapshots")).toHaveCount(0); - await expect(dividers(page)).toHaveCount(1); - }); -}); - /** Opens one boundary's record, which is what pressing its line does. */ async function openSnapshot( page: import("@playwright/test").Page, diff --git a/ui/src/api/domain/checkpoints.ts b/ui/src/api/domain/checkpoints.ts index 62c12d4a8e..51f556a1e5 100644 --- a/ui/src/api/domain/checkpoints.ts +++ b/ui/src/api/domain/checkpoints.ts @@ -14,9 +14,9 @@ * How far the controller has got with retaining the snapshot. * * `CheckpointState` in the proto, spelled as words for the reasons given in - * `agentInstances.ts`. Successful creation answers `ready`; a terminal failure is - * an error. Owner queries also return `failed` records so they can be deleted. - * Only `ready` records can be forked or renamed. + * `agentInstances.ts`. Creation answers `ready` or `failed` and never `creating` — + * the RPC copies the snapshot before it returns — so a caller that wants to fork can + * check the state it was handed rather than poll. */ export type CheckpointState = | "unspecified" diff --git a/ui/src/api/grpc/checkpoints.test.ts b/ui/src/api/grpc/checkpoints.test.ts index 29835e2e29..71ed7f9488 100644 --- a/ui/src/api/grpc/checkpoints.test.ts +++ b/ui/src/api/grpc/checkpoints.test.ts @@ -41,7 +41,7 @@ it("forks through Session RPCs and reads the new conversation identity", async ( expect(result).toMatchObject({ id: "forked-session", contextId: "forked-session", name: "Saved turn" }); }); -it.each(["KAGENT_CHECKPOINT_CONVERSATION_ADVANCED", "KAGENT_CHECKPOINT_FAILED", undefined])("does not retry %s", async (reason) => { +it.each(["KAGENT_CHECKPOINT_CONVERSATION_ADVANCED", undefined])("does not retry %s", async (reason) => { let calls = 0; setApiTransport(createRouterTransport((router) => { router.service(CheckpointService, { @@ -56,92 +56,3 @@ it.each(["KAGENT_CHECKPOINT_CONVERSATION_ADVANCED", "KAGENT_CHECKPOINT_FAILED", await expect(defaultOperations["agentInstances.checkpoints.create"]({ id: "instance", requestId: "request", expectedHeadTaskId: "turn-a" }, {})).rejects.toMatchObject({ reason }); expect(calls).toBe(1); }); - -it("reads failed cleanup candidates beyond the first page", async () => { - const tokens: string[] = []; - setApiTransport(createRouterTransport((router) => { - router.service(CheckpointService, { - listCheckpoints(request) { - tokens.push(request.page?.pageToken ?? ""); - return request.page?.pageToken - ? { checkpoints: [{ id: "failed", sessionId: request.sessionId, state: CheckpointState.FAILED, failure: { message: "Copy could not be verified" } }], page: {} } - : { checkpoints: [{ id: "ready", sessionId: request.sessionId, state: CheckpointState.READY }], page: { nextPageToken: "next" } }; - }, - }); - })); - const result = await defaultOperations["agentInstances.checkpoints.list"]({ id: "instance" }, {}); - expect(tokens).toEqual(["", "next"]); - expect(result).toEqual(expect.arrayContaining([ - expect.objectContaining({ id: "ready", state: "ready" }), - expect.objectContaining({ id: "failed", state: "failed", failure: "Copy could not be verified" }), - ])); -}); - -it("rejects a repeated snapshot page rather than returning an incomplete list", async () => { - let calls = 0; - setApiTransport(createRouterTransport((router) => { - router.service(CheckpointService, { - listCheckpoints() { - calls++; - return { checkpoints: [], page: { nextPageToken: "same" } }; - }, - }); - })); - await expect(defaultOperations["agentInstances.checkpoints.list"]({ id: "instance" }, {})).rejects.toThrow("repeated a page"); - expect(calls).toBe(2); -}); - -it("confirms a missing snapshot after a delete retry without masking a retained record", async () => { - let retained = true; - const calls: string[] = []; - setApiTransport(createRouterTransport((router) => { - router.service(CheckpointService, { - deleteCheckpoint(request) { - calls.push(`delete:${request.checkpointId}`); - throw new ConnectError("checkpoint still referenced or already deleted", Code.NotFound); - }, - getCheckpoint(request) { - calls.push(`get:${request.checkpointId}`); - if (!retained) throw new ConnectError("checkpoint absent", Code.NotFound); - return { checkpoint: { id: request.checkpointId, state: CheckpointState.READY } }; - }, - }); - })); - const remove = () => defaultOperations["agentInstances.checkpoints.delete"]({ checkpointId: "checkpoint" }, {}); - await expect(remove()).rejects.toMatchObject({ code: "NotFound", message: "checkpoint still referenced or already deleted" }); - retained = false; - await expect(remove()).resolves.toBeUndefined(); - expect(calls).toEqual(["delete:checkpoint", "get:checkpoint", "delete:checkpoint", "get:checkpoint"]); -}); - -it.each([Code.PermissionDenied, Code.Unavailable, "empty" as const])("preserves a failed absence check (%s)", async (failure) => { - setApiTransport(createRouterTransport((router) => { - router.service(CheckpointService, { - deleteCheckpoint() { throw new ConnectError("checkpoint absent or referenced", Code.NotFound); }, - getCheckpoint() { - if (failure === "empty") return {}; - throw new ConnectError("lookup failed", failure); - }, - }); - })); - await expect(defaultOperations["agentInstances.checkpoints.delete"]({ checkpointId: "checkpoint" }, {})).rejects.toMatchObject( - failure === "empty" ? { kind: "parse" } : { code: Code[failure], message: "lookup failed" }, - ); -}); - -it.each([undefined, Code.PermissionDenied, Code.Unavailable])("does not look up a snapshot after delete outcome %s", async (failure) => { - let lookups = 0; - setApiTransport(createRouterTransport((router) => { - router.service(CheckpointService, { - deleteCheckpoint() { - if (failure !== undefined) throw new ConnectError("delete failed", failure); - return {}; - }, - getCheckpoint() { lookups++; return {}; }, - }); - })); - const result = defaultOperations["agentInstances.checkpoints.delete"]({ checkpointId: "checkpoint" }, {}); - if (failure === undefined) await expect(result).resolves.toBeUndefined(); - else await expect(result).rejects.toMatchObject({ code: Code[failure] }); - expect(lookups).toBe(0); -}); diff --git a/ui/src/api/grpc/operations.ts b/ui/src/api/grpc/operations.ts index 0a44d73dae..2ae8b2d8bc 100644 --- a/ui/src/api/grpc/operations.ts +++ b/ui/src/api/grpc/operations.ts @@ -37,8 +37,8 @@ import { ScheduledRunService } from "@/generated/kagent/api/v1/scheduled_runs_pb * * `CheckpointService` is reached through the `agentInstances.checkpoints.*` ids, and * **`agentInstances.fork`** composes two of them — a checkpoint of the conversation - * as it stands, then a fork of it. Deletion confirms absence through `GetCheckpoint` - * after NotFound, since a retained checkpoint with dependent forks returns that too. + * as it stands, then a fork of it. `GetCheckpoint` and `DeleteCheckpoint` have no id: + * the chat reads boundaries from the list and nothing removes one yet. */ import { ModelService } from "@/generated/kagent/api/v1/models_pb"; @@ -625,7 +625,6 @@ function toAgentInstance(instance: PbSession): AgentInstance { * worth reporting rather than a list worth truncating. */ const INSTANCE_PAGE_LIMIT = 50; -const CHECKPOINT_PAGE_LIMIT = 50; function toAgentInstanceShare(share: PbSessionShare): AgentInstanceShare { return { @@ -804,30 +803,22 @@ const agentInstances: Pick< return checkpoint; }, - // FAILED cleanup candidates can be on later pages. Read the complete list - // before sorting it newest-first, just as for conversation discovery above. + /* + * Newest first, because that is the order the chat needs them in and the controller + * does not promise one. One page: a conversation's boundaries are counted in + * handfuls, and paging a list this short would be machinery with nothing to do. + */ "agentInstances.checkpoints.list": async (input, options) => { const name = "CheckpointService/ListCheckpoints"; - const rows: Checkpoint[] = []; - let pageToken = ""; - const seen = new Set(); - for (let page = 0; page < CHECKPOINT_PAGE_LIMIT; page += 1) { - const response = await rpc(name, options.signal, () => - serviceClient(CheckpointService).listCheckpoints( - { sessionId: input.id, page: { pageToken } }, - call("agentInstances.checkpoints.list", options), - ), - ); - rows.push(...list(response.checkpoints).map(toCheckpoint)); - const next = response.page?.nextPageToken ?? ""; - if (!next) return rows.sort((left, right) => (right.createdAt ?? "").localeCompare(left.createdAt ?? "")); - if (seen.has(next)) { - throw new ApiError("The API repeated a page of snapshots instead of advancing.", { kind: "parse", url: name }); - } - seen.add(next); - pageToken = next; - } - throw new ApiError("The snapshot list could not be read to the end.", { kind: "parse", url: name }); + const response = await rpc(name, options.signal, () => + serviceClient(CheckpointService).listCheckpoints( + { sessionId: input.id }, + call("agentInstances.checkpoints.list", options), + ), + ); + return list(response.checkpoints) + .map(toCheckpoint) + .sort((left, right) => (right.createdAt ?? "").localeCompare(left.createdAt ?? "")); }, /* @@ -870,33 +861,12 @@ const agentInstances: Pick< }, "agentInstances.checkpoints.delete": async (input, options) => { - try { - await rpc("CheckpointService/DeleteCheckpoint", options.signal, () => - serviceClient(CheckpointService).deleteCheckpoint( - { checkpointId: input.checkpointId }, - call("agentInstances.checkpoints.delete", options), - ), - ); - } catch (error) { - if (!(error instanceof ApiError) || error.code !== "NotFound") throw error; - // Recovery may have completed deletion after an earlier failed reply. A - // fork-pinned snapshot also rejects Delete with NotFound, so confirm absence - // before letting callers remove their local record and close its details. - const name = "CheckpointService/GetCheckpoint"; - try { - const response = await rpc(name, options.signal, () => - serviceClient(CheckpointService).getCheckpoint( - { checkpointId: input.checkpointId }, - call("agentInstances.checkpoints.delete", options), - ), - ); - required(response.checkpoint, name, "checkpoint"); - } catch (lookupError) { - if (lookupError instanceof ApiError && lookupError.code === "NotFound") return; - throw lookupError; - } - throw error; - } + await rpc("CheckpointService/DeleteCheckpoint", options.signal, () => + serviceClient(CheckpointService).deleteCheckpoint( + { checkpointId: input.checkpointId }, + call("agentInstances.checkpoints.delete", options), + ), + ); }, /* diff --git a/ui/src/api/hooks/useApiResource.test.tsx b/ui/src/api/hooks/useApiResource.test.tsx deleted file mode 100644 index c8167fc87c..0000000000 --- a/ui/src/api/hooks/useApiResource.test.tsx +++ /dev/null @@ -1,146 +0,0 @@ -import { act, renderHook, waitFor } from "@testing-library/react"; -import type { ReactNode } from "react"; -import { SWRConfig, useSWRConfig } from "swr"; -import { expect, it, vi } from "vitest"; -import { ApiError } from "../ApiError"; -import { useApiResource } from "./useApiResource"; -import { deferred } from "@/testUtils/deferred"; - -function wrapper({ children }: { children: ReactNode }) { - return new Map(), shouldRetryOnError: false, dedupingInterval: 0 }}>{children}; -} - -it("retains loaded data and exposes a failed refresh until a repaired retry succeeds", async () => { - const failure = new ApiError("List unavailable", { kind: "network", url: "test/list" }); - const fetcher = vi.fn<() => Promise>() - .mockResolvedValueOnce(["retained"]) - .mockRejectedValueOnce(failure) - .mockResolvedValueOnce([]); - const { result } = renderHook(() => useApiResource(["refresh-test"], fetcher), { wrapper }); - await waitFor(() => expect(result.current.data).toEqual(["retained"])); - await act(async () => { await expect(result.current.refresh()).rejects.toBe(failure); }); - expect(result.current.data).toEqual(["retained"]); - expect(result.current.error).toBe(failure); - expect(result.current.isEmpty).toBe(false); - await act(async () => { await result.current.refresh(); }); - expect(result.current.data).toEqual([]); - expect(result.current.error).toBeUndefined(); - expect(result.current.isEmpty).toBe(true); - expect(fetcher).toHaveBeenCalledTimes(3); -}); - -it("keeps refresh failures scoped to their resource key", async () => { - const failure = new ApiError("Old list unavailable", { kind: "network", url: "test/list" }); - const fetcher = vi.fn<() => Promise>() - .mockResolvedValueOnce(["first"]) - .mockRejectedValueOnce(failure) - .mockResolvedValueOnce(["second"]); - const { result, rerender } = renderHook(({ id }) => useApiResource(["scoped-refresh", id], fetcher), { initialProps: { id: "first" }, wrapper }); - await waitFor(() => expect(result.current.data).toEqual(["first"])); - await act(async () => { await expect(result.current.refresh()).rejects.toBe(failure); }); - rerender({ id: "second" }); - await waitFor(() => expect(result.current.data).toEqual(["second"])); - expect(result.current.error).toBeUndefined(); -}); - -it("does not read an idle resource when refresh is invoked", async () => { - const fetcher = vi.fn(async () => ["unexpected"]); - const { result } = renderHook(() => useApiResource(null, fetcher), { wrapper }); - await act(async () => { await result.current.refresh(); }); - expect(fetcher).not.toHaveBeenCalled(); - expect(result.current.data).toBeUndefined(); - expect(result.current.isLoading).toBe(false); -}); - -it("keeps an earlier refresh callback bound to its original resource after navigation", async () => { - const fetcher = vi.fn(async (id: string) => [id]); - const { result, rerender } = renderHook(({ id }) => useApiResource(["navigation-refresh", id], () => fetcher(id)), { initialProps: { id: "first" }, wrapper }); - await waitFor(() => expect(result.current.data).toEqual(["first"])); - const firstRefresh = result.current.refresh; - rerender({ id: "second" }); - await waitFor(() => expect(result.current.data).toEqual(["second"])); - await act(async () => { await firstRefresh(); }); - expect(fetcher.mock.calls).toEqual([["first"], ["second"]]); - expect(result.current.data).toEqual(["second"]); - expect(result.current.error).toBeUndefined(); -}); - -it.each(["refresh", "background"] as const)("ignores an older failure after a newer %s succeeds for all readers", async (mode) => { - const older = deferred(); - const failure = new ApiError("Earlier read failed", { kind: "network", url: "test/list" }); - const fetcher = vi.fn<() => Promise>() - .mockResolvedValueOnce(["initial"]) - .mockReturnValueOnce(older.promise) - .mockResolvedValueOnce(["repaired"]); - const { result } = renderHook(() => ({ - first: useApiResource(["ordered", mode], fetcher), - shared: useApiResource(["ordered", mode], fetcher), - revalidate: useSWRConfig().mutate, - }), { wrapper }); - await waitFor(() => expect(result.current.first.data).toEqual(["initial"])); - let earlier!: Promise; - act(() => { earlier = result.current.first.refresh().catch(error => error); }); - await waitFor(() => expect(fetcher).toHaveBeenCalledTimes(2)); - await act(async () => { - if (mode === "refresh") await result.current.first.refresh(); - else await result.current.revalidate(["ordered", mode]); - }); - expect(result.current.first.data).toEqual(["repaired"]); - await act(async () => { older.reject(failure); await earlier; }); - expect(result.current.first.error).toBeUndefined(); - expect(result.current.shared.error).toBeUndefined(); - expect(result.current.first.data).toEqual(["repaired"]); -}); - -it("preserves the newer failure when an older read fails afterward", async () => { - const older = deferred(); - const failure = new ApiError("Current read failed", { kind: "network", url: "test/list" }); - const fetcher = vi.fn<() => Promise>() - .mockResolvedValueOnce(["retained"]) - .mockReturnValueOnce(older.promise) - .mockRejectedValueOnce(failure); - const { result } = renderHook(() => useApiResource(["ordered-failure"], fetcher), { wrapper }); - await waitFor(() => expect(result.current.data).toEqual(["retained"])); - let earlier!: Promise; - act(() => { earlier = result.current.refresh().catch(error => error); }); - await waitFor(() => expect(fetcher).toHaveBeenCalledTimes(2)); - await act(async () => { await expect(result.current.refresh()).rejects.toBe(failure); }); - await act(async () => { older.reject(new ApiError("Obsolete failure", { kind: "network", url: "test/list" })); await earlier; }); - expect(result.current.error).toBe(failure); - expect(result.current.data).toEqual(["retained"]); -}); - -it("keeps concurrent reads with the same key isolated across cache providers", async () => { - const pending = deferred(); - const failure = new ApiError("First provider failed", { kind: "network", url: "test/list" }); - const first = renderHook(() => useApiResource(["provider-isolation"], () => pending.promise), { wrapper }); - const second = renderHook(() => useApiResource(["provider-isolation"], async () => ["healthy"]), { wrapper }); - await waitFor(() => expect(second.result.current.data).toEqual(["healthy"])); - await act(async () => { pending.reject(failure); }); - await waitFor(() => expect(first.result.current.error).toBe(failure)); - expect(second.result.current.data).toEqual(["healthy"]); - expect(second.result.current.error).toBeUndefined(); -}); - -it("reports a later current failure normally after overlapping reads have settled", async () => { - const older = deferred(); - const failure = new ApiError("New current failure", { kind: "network", url: "test/list" }); - const onError = vi.fn(); - const fetcher = vi.fn<() => Promise>() - .mockResolvedValueOnce(["initial"]) - .mockReturnValueOnce(older.promise) - .mockResolvedValueOnce(["repaired"]) - .mockRejectedValueOnce(failure); - const { result } = renderHook(() => useApiResource(["after-overlap"], fetcher, { onError }), { wrapper }); - await waitFor(() => expect(result.current.data).toEqual(["initial"])); - let earlier!: Promise; - act(() => { earlier = result.current.refresh().catch(error => error); }); - await waitFor(() => expect(fetcher).toHaveBeenCalledTimes(2)); - await act(async () => { await result.current.refresh(); }); - await act(async () => { older.reject(new ApiError("Obsolete failure", { kind: "network", url: "test/list" })); await earlier; }); - expect(onError).not.toHaveBeenCalled(); - await act(async () => { await expect(result.current.refresh()).rejects.toBe(failure); }); - expect(result.current.error).toBe(failure); - expect(onError).toHaveBeenCalledTimes(1); - expect(onError.mock.calls[0][0]).toBe(failure); -}); diff --git a/ui/src/api/hooks/useApiResource.ts b/ui/src/api/hooks/useApiResource.ts index f60d7df8c6..d0aa2052bb 100644 --- a/ui/src/api/hooks/useApiResource.ts +++ b/ui/src/api/hooks/useApiResource.ts @@ -7,7 +7,7 @@ * empty state. This wraps it once so all four states are named. */ -import useSWR, { unstable_serialize, useSWRConfig, type Cache, type SWRConfiguration } from "swr"; +import useSWR, { type SWRConfiguration } from "swr"; import { ApiError } from "../ApiError"; export interface ApiResource { @@ -34,11 +34,9 @@ export function useApiResource( fetcher: () => Promise, config?: SWRConfiguration, ): ApiResource { - const { cache, mutate } = useSWRConfig(); - const resourceKey = unstable_serialize(key); - const { data, error, isLoading, isValidating } = useSWR( + const { data, error, isLoading, isValidating, mutate } = useSWR( key, - () => readResource(cache, resourceKey, fetcher), + fetcher, config, ); @@ -53,58 +51,22 @@ export function useApiResource( /** * Refetches, and rejects when the refetch failed. * - * Revalidation updates SWR's shared error/loading state; mutating with a - * rejected fetch promise does not. SWR resolves even a failed revalidation, - * so read the completed result from this provider's cache and rethrow it. - * The render's `error` would still describe the previous read here. + * The bare `mutate()` this used to be does not reject: SWR captures a failed + * revalidation into `error` and resolves anyway, so a caller awaiting it cannot + * tell a refresh that worked from one that did not. That is fine for a + * background revalidation and wrong for a button someone pressed — it made the + * Refresh control report success over a page showing a load error. + * + * Handing `mutate` the promise instead makes it the new value: SWR records the + * rejection in `error` as before, so the page still shows what went wrong, and + * rethrows it here so the caller knows too. */ refresh: async () => { - await mutate(key); - const failure = cache.get(resourceKey)?.error; - if (failure) throw failure; + await mutate(fetcher(), { revalidate: false }); }, }; } -interface ResourceReads { - latest: Promise; - pending: number; -} - -// Share ordering across readers and automatic revalidation on the same provider. -// Entries exist only while reads are in flight, including superseded requests. -const readsByCache = new WeakMap>(); - -function readResource(cache: Cache, key: string, fetcher: () => Promise): Promise { - let reads = readsByCache.get(cache); - if (!reads) { - reads = new Map(); - readsByCache.set(cache, reads); - } - const active = reads; - const dataBefore = cache.get(key)?.data; - const request: Promise = Promise.resolve().then(fetcher).catch((cause: unknown) => { - if (active.get(key)?.latest !== request || cache.get(key)?.data !== dataBefore) { - // SWR fences obsolete successes, but can still publish obsolete errors. - // Resolve through the cache so its request/mutation fence discards this - // result too. Only the current read may publish a failure or retry it. - return cache.get(key)?.data as T; - } - throw cause; - }); - const previous = active.get(key); - active.set(key, { latest: request, pending: (previous?.pending ?? 0) + 1 }); - const finish = () => { - const current = active.get(key); - if (current && --current.pending === 0) { - active.delete(key); - if (active.size === 0) readsByCache.delete(cache); - } - }; - void request.then(finish, finish); - return request; -} - function isEmptyResult(data: unknown): boolean { if (data === undefined || data === null) return false; // Not loaded is not empty. if (Array.isArray(data)) return data.length === 0; diff --git a/ui/src/api/hooks/useCheckpoints.test.tsx b/ui/src/api/hooks/useCheckpoints.test.tsx deleted file mode 100644 index fe0937f07b..0000000000 --- a/ui/src/api/hooks/useCheckpoints.test.tsx +++ /dev/null @@ -1,67 +0,0 @@ -import { act, renderHook, waitFor } from "@testing-library/react"; -import type { ReactNode } from "react"; -import { SWRConfig } from "swr"; -import { afterEach, expect, it } from "vitest"; -import { Code, ConnectError, createRouterTransport } from "@connectrpc/connect"; -import { CheckpointService, CheckpointState } from "@/generated/kagent/api/v1/checkpoints_pb"; -import { setApiTransport } from "../transport"; -import { useCheckpoints } from "./useCheckpoints"; -import { deferred } from "@/testUtils/deferred"; - -function wrapper({ children }: { children: ReactNode }) { - return new Map(), shouldRetryOnError: false, dedupingInterval: 0 }}>{children}; -} - -afterEach(() => setApiTransport(undefined)); - -it("shares confirmed deletion with readers while keeping other conversations and read failures intact", async () => { - let unavailable = false; - const requests: string[] = []; - setApiTransport(createRouterTransport(({ service }) => service(CheckpointService, { - listCheckpoints(request) { - requests.push(request.sessionId); - if (unavailable) throw new ConnectError("List unavailable", Code.Unavailable); - return { checkpoints: [ - { id: `${request.sessionId}-ready`, state: CheckpointState.READY }, - { id: `${request.sessionId}-failed`, state: CheckpointState.FAILED }, - ] }; - }, - }))); - const { result } = renderHook(() => ({ first: useCheckpoints("first"), shared: useCheckpoints("first"), other: useCheckpoints("other") }), { wrapper }); - await waitFor(() => expect(result.current.first.data).toHaveLength(2)); - await waitFor(() => expect(result.current.other.data).toHaveLength(2)); - const reads = requests.length; - await act(async () => { await result.current.first.forget("first-failed"); }); - expect(result.current.first.data?.map(row => row.id)).toEqual(["first-ready"]); - expect(result.current.shared.data?.map(row => row.id)).toEqual(["first-ready"]); - expect(result.current.other.data?.map(row => row.id)).toEqual(["other-ready", "other-failed"]); - expect(requests).toHaveLength(reads); - unavailable = true; - await act(async () => { await expect(result.current.first.refresh()).rejects.toMatchObject({ code: "Unavailable" }); }); - expect(result.current.first.error?.code).toBe("Unavailable"); - expect(result.current.shared.error?.code).toBe("Unavailable"); - expect(result.current.first.data?.map(row => row.id)).toEqual(["first-ready"]); - expect(result.current.other.error).toBeUndefined(); -}); - -it.each(["success", "failure"] as const)("does not restore a deleted record or obsolete error from an earlier %s", async (outcome) => { - const pending = deferred<{ checkpoints: Array<{ id: string; state: CheckpointState }> }>(); - const records = [{ id: "ready", state: CheckpointState.READY }, { id: "failed", state: CheckpointState.FAILED }]; - let reads = 0; - setApiTransport(createRouterTransport(({ service }) => service(CheckpointService, { - listCheckpoints() { return ++reads === 2 ? pending.promise : { checkpoints: records }; }, - }))); - const { result } = renderHook(() => useCheckpoints("conversation"), { wrapper }); - await waitFor(() => expect(result.current.data).toHaveLength(2)); - let earlier!: Promise; - act(() => { earlier = result.current.refresh().catch(error => error); }); - await waitFor(() => expect(reads).toBe(2)); - await act(async () => { await result.current.forget("failed"); }); - await act(async () => { - if (outcome === "success") pending.resolve({ checkpoints: records }); - else pending.reject(new ConnectError("Obsolete read failure", Code.Unavailable)); - await earlier; - }); - expect(result.current.data?.map(row => row.id)).toEqual(["ready"]); - expect(result.current.error).toBeUndefined(); -}); diff --git a/ui/src/api/hooks/useCheckpoints.ts b/ui/src/api/hooks/useCheckpoints.ts index d6fd610951..dc0c652509 100644 --- a/ui/src/api/hooks/useCheckpoints.ts +++ b/ui/src/api/hooks/useCheckpoints.ts @@ -1,7 +1,6 @@ import { apiClient } from "../client"; import type { Checkpoint } from "../domain/checkpoints"; import { type ApiResource, useApiResource } from "./useApiResource"; -import { useSWRConfig } from "swr"; /** * The turn boundaries saved against one conversation. @@ -9,23 +8,9 @@ import { useSWRConfig } from "swr"; * Held back until the conversation is known, which `useApiResource` reports as idle * rather than as loading. */ -export function useCheckpoints(id: string | undefined): ApiResource & { - forget: (checkpointId: string) => Promise; -} { - const { mutate } = useSWRConfig(); - const key = id ? ["agentInstances.checkpoints.list", id] : null; - const resource = useApiResource( - key, +export function useCheckpoints(id: string | undefined): ApiResource { + return useApiResource( + id ? ["agentInstances.checkpoints.list", id] : null, () => apiClient.agentInstances.checkpoints.list(id ?? ""), ); - return { - ...resource, - // Confirmed deletion is already authoritative even if the next list fails. - // Update the shared record independently of that subsequent read. - forget: async (checkpointId) => { - if (key) { - await mutate(key, (rows) => rows?.filter((row) => row.id !== checkpointId), { revalidate: false }); - } - }, - }; } diff --git a/ui/src/components/chat/ChatTranscript.tsx b/ui/src/components/chat/ChatTranscript.tsx index 9ddaecb013..0ac76403b5 100644 --- a/ui/src/components/chat/ChatTranscript.tsx +++ b/ui/src/components/chat/ChatTranscript.tsx @@ -80,7 +80,7 @@ export function ChatTranscript({ /** Starts a new conversation from a boundary. */ onFork?: (checkpointId: string) => void; /** Drops a boundary and the runtime stored with it. */ - onDeleteCheckpoint?: (checkpointId: string) => void | Promise; + onDeleteCheckpoint?: (checkpointId: string) => void; /** Opens active callbacks for one A2A task. */ onManageTaskPush?: (taskId: string) => void; /** The boundaries the controller has described, for the line to name itself by. */ diff --git a/ui/src/components/chat/CheckpointDivider.tsx b/ui/src/components/chat/CheckpointDivider.tsx index 007566d787..632d85ca53 100644 --- a/ui/src/components/chat/CheckpointDivider.tsx +++ b/ui/src/components/chat/CheckpointDivider.tsx @@ -45,7 +45,7 @@ export function CheckpointDivider({ /** Names this boundary, and with it the forks taken from it. */ onRename?: () => void; /** Removes this boundary. */ - onDelete?: () => void | Promise; + onDelete?: () => void; }) { const theme = useTheme(); const { mode } = useThemeMode(); @@ -286,14 +286,7 @@ export function CheckpointDivider({ cancelButtonProps={{ "data-testid": `chat-checkpoint-delete-cancel-${checkpointId}`, }} - onConfirm={async () => { - try { - await onDelete(); - } catch { - // The page reports the failure and retains the row for retry. - // Popconfirm rethrows rejected callbacks as uncaught errors. - } - }} + onConfirm={onDelete} > {/* No tooltip, unlike its neighbours. The pointer that opens the confirmation is still resting on this button, so the tooltip stays up diff --git a/ui/src/components/chat/SnapshotDetailsModal.tsx b/ui/src/components/chat/SnapshotDetailsModal.tsx index a0ece8a0fc..249fbe39f5 100644 --- a/ui/src/components/chat/SnapshotDetailsModal.tsx +++ b/ui/src/components/chat/SnapshotDetailsModal.tsx @@ -1,5 +1,5 @@ import { useState } from "react"; -import { Alert, Button, Descriptions, Modal, Popconfirm, Space, Tag, Tooltip, Typography } from "antd"; +import { Button, Descriptions, Modal, Popconfirm, Space, Tag, Tooltip, Typography } from "antd"; import { Eraser, GitFork, Pencil } from "lucide-react"; import { useTheme } from "@emotion/react"; import { canForkFrom, type Checkpoint, type CheckpointState } from "@/api"; @@ -67,19 +67,15 @@ export function SnapshotDetailsModal({ const theme = useTheme(); const [isRenaming, setRenaming] = useState(false); const [isBusy, setBusy] = useState(false); - const [actionError, setActionError] = useState(); const label = snapshotLabel(checkpoint); const forkable = canForkFrom(checkpoint); const state = STATE_APPEARANCE[checkpoint.state]; async function act(run: () => void | Promise) { setBusy(true); - setActionError(undefined); try { await run(); onClose(); - } catch (cause: unknown) { - setActionError(cause instanceof Error ? cause.message : String(cause)); } finally { setBusy(false); } @@ -110,12 +106,11 @@ export function SnapshotDetailsModal({ children: ( {label} - + - - ))} - - - } - /> - ) : null} - {checkpoints.error ? ( - void checkpoints.refresh().catch((cause: unknown) => { - console.error("Could not re-read the saved boundaries:", cause); - })}> - Retry - - } - /> - ) : null} waitFor(() => { - expect(run).toBeEnabled(); - expect(run).not.toHaveClass("ant-btn-loading"); - }); - await waitForRun(); fireEvent.click(run); await screen.findByText(/Response lost/); - await waitForRun(); + await waitFor(() => expect(run).toBeEnabled()); fireEvent.click(run); await screen.findByText(/Execution queued/); - await waitForRun(); + await waitFor(() => expect(run).toBeEnabled()); fireEvent.click(run); await waitFor(() => expect(requestIds).toHaveLength(3)); expect(requestIds[0]).toBeTruthy(); expect(requestIds[1]).toBe(requestIds[0]); expect(requestIds[2]).not.toBe(requestIds[0]); - await waitForRun(); + await waitFor(() => expect(run).toBeEnabled()); }); it("keeps unsaved edits after a conflict and preserves the exact timeout on the wire", async () => { diff --git a/ui/src/testUtils/deferred.ts b/ui/src/testUtils/deferred.ts deleted file mode 100644 index d09f3e7110..0000000000 --- a/ui/src/testUtils/deferred.ts +++ /dev/null @@ -1,7 +0,0 @@ -/** Control completion order at an asynchronous test boundary. */ -export function deferred() { - let resolve!: (value: T) => void; - let reject!: (error: unknown) => void; - const promise = new Promise((yes, no) => { resolve = yes; reject = no; }); - return { promise, resolve, reject }; -} From 238ef3d246d677ed7e95e24babccef4707e5127e Mon Sep 17 00:00:00 2001 From: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> Date: Tue, 10 Nov 2026 17:15:45 +0800 Subject: [PATCH 20/20] test(checkpoint): cover failed tag collection after deletion and takeover Signed-off-by: Junfeng Yu <64075563+Jacob953@users.noreply.github.com> --- .../service/checkpoint/completion_test.go | 10 ++ .../service/checkpoint/failed_cleanup_test.go | 107 ++++++++++++++++++ 2 files changed, 117 insertions(+) diff --git a/go/core/internal/service/checkpoint/completion_test.go b/go/core/internal/service/checkpoint/completion_test.go index 717256500f..07b2cdc982 100644 --- a/go/core/internal/service/checkpoint/completion_test.go +++ b/go/core/internal/service/checkpoint/completion_test.go @@ -157,6 +157,7 @@ type checkpointCompletionGate struct { 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 } @@ -189,6 +190,15 @@ func (s *checkpointCompletionGate) DeleteClaimedSessionCheckpoint(ctx context.Co 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 { diff --git a/go/core/internal/service/checkpoint/failed_cleanup_test.go b/go/core/internal/service/checkpoint/failed_cleanup_test.go index e6aa55d702..6b3d03b31a 100644 --- a/go/core/internal/service/checkpoint/failed_cleanup_test.go +++ b/go/core/internal/service/checkpoint/failed_cleanup_test.go @@ -4,6 +4,7 @@ import ( "context" "errors" "testing" + "time" "github.com/google/uuid" "github.com/stretchr/testify/require" @@ -127,6 +128,112 @@ func TestRecoveryRetainsPriorTagWhenSourceChanged(t *testing.T) { 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))