From 1dbf687cc773d5cec393c7ca9628c950b2f89870 Mon Sep 17 00:00:00 2001 From: Ada Lundhe Date: Fri, 23 Oct 2026 12:31:02 -0700 Subject: [PATCH 01/12] fix(session): add durable recovery primitives Signed-off-by: Ada Lundhe --- go/core/internal/database/session_recovery.go | 197 ++++++++++++++ .../database/session_recovery_test.go | 255 ++++++++++++++++++ .../internal/database/task_finalization.go | 96 +++---- go/core/internal/substrate/client.go | 10 + go/core/internal/substrate/snapshot.go | 11 + go/core/internal/substrate/snapshot_test.go | 20 ++ 6 files changed, 544 insertions(+), 45 deletions(-) create mode 100644 go/core/internal/database/session_recovery.go create mode 100644 go/core/internal/database/session_recovery_test.go diff --git a/go/core/internal/database/session_recovery.go b/go/core/internal/database/session_recovery.go new file mode 100644 index 0000000000..a803907ab9 --- /dev/null +++ b/go/core/internal/database/session_recovery.go @@ -0,0 +1,197 @@ +package database + +import ( + "context" + "fmt" + "time" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "google.golang.org/protobuf/types/known/timestamppb" + + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" +) + +// SessionRuntimeCrashed identifies a recovery admitted after Substrate confirmed +// that the recorded Actor crashed. Transport failures cannot admit this work. +const SessionRuntimeCrashed = "RuntimeCrashed" + +// ListSessionRecoveryCandidates finds recoveries admitted by an explicit request. +// Retrying durable work cannot admit recovery for a READY Session. +func (c *Client) ListSessionRecoveryCandidates(ctx context.Context, afterID string, limit int) ([]*apiv1.Session, error) { + rows, err := queryMany(ctx, c.db, ` + SELECT id, user_id, prepared_revision, state, data, operation, context_id, + source_checkpoint_id, history_id, operation_id, executor_id + FROM session_record i + WHERE ($1::text = '' OR i.id > $1::uuid) + AND i.state = 'RUNTIME_STATE_FAILED' AND i.operation = 'RUNTIME_OPERATION_RESUME' + ORDER BY i.id LIMIT $2 + `, pgx.RowToStructByName[sessionRow], afterID, limit) + if err != nil { + return nil, err + } + result := make([]*apiv1.Session, 0, len(rows)) + for _, row := range rows { + session, err := toSession(row) + if err != nil { + return nil, err + } + result = append(result, session) + } + return result, nil +} + +// SessionLifecycleGeneration returns the ID of the Session's latest lifecycle +// operation, or uuid.Nil when none is recorded. Every operation replaces it, and +// completion keeps it. Recovery reads it before observing Substrate. +func (c *Client) SessionLifecycleGeneration(ctx context.Context, id string) (uuid.UUID, error) { + generation, err := queryOne(ctx, c.db, `SELECT operation_id FROM session_record WHERE id = $1`, pgx.RowTo[*uuid.UUID], id) + if err != nil { + return uuid.Nil, notFoundOr(err) + } + if generation == nil { + return uuid.Nil, nil + } + return *generation, nil +} + +// BeginCrashedSessionRecovery fences confirmed-crash recovery under the Session lock. +// It preserves durable results; obsolete Actors and conflicting operations fail. +// generation is the caller's SessionLifecycleGeneration read from before it +// observed the crash. A lifecycle operation begun since then makes the observation +// stale, because a concurrent revert keeps the Actor UID and returns to READY. +func (c *Client) BeginCrashedSessionRecovery(ctx context.Context, id string, generation uuid.UUID, actorUID string, snapshot SessionTaskSnapshot) (*SessionOperation, error) { + var operation *SessionOperation + err := c.withTx(ctx, func(tx pgx.Tx) error { + row, err := lockSession(ctx, tx, id) + if err != nil { + return notFoundOr(err) + } + matches, err := queryOne(ctx, tx, `SELECT actor_uid = $2 FROM session WHERE id = $1`, pgx.RowTo[bool], row.ID, actorUID) + if err != nil { + return err + } + if !matches || actorUID == "" { + return ErrConflict + } + session, err := toSession(row) + if err != nil { + return err + } + if session.State == apiv1.RuntimeState_RUNTIME_STATE_FAILED && session.Operation == apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME && session.Failure.GetReason() == SessionRuntimeCrashed { + operation, err = toSessionOperation(row) + return err + } + if session.State != apiv1.RuntimeState_RUNTIME_STATE_READY || session.Operation != apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE { + return ErrConflict + } + current := uuid.Nil + if row.OperationID != nil { + current = *row.OperationID + } + if current != generation { + return ErrConflict + } + // An issued idle suspension does not exclude recovery here. It cannot finish + // against the crashed Actor, and its worker keeps the claim when suspension + // fails, so waiting for it would block recovery permanently. Settlement below + // retires the claim in this transaction; a late finish for it changes nothing. + // The generation check keeps the crash observation current: without another + // lifecycle operation, nothing returns the crashed Actor to service. + busy, err := queryOne(ctx, tx, ` + SELECT EXISTS (SELECT 1 FROM session_checkpoint WHERE source_session_id = $1 AND state = 'CREATING') + `, pgx.RowTo[bool], row.ID) + if err != nil { + return err + } + if busy { + return ErrConflict + } + row.State = apiv1.RuntimeState_RUNTIME_STATE_FAILED.String() + row.runtimeInstanceRow, err = beginRuntimeOperation(ctx, tx, row.runtimeInstanceRow, runtimeKindAgent, apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME) + if err != nil { + return err + } + session.State, session.Operation = apiv1.RuntimeState_RUNTIME_STATE_FAILED, apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME + session.Failure = &apiv1.Failure{Reason: SessionRuntimeCrashed, Message: "Runtime crashed; recovery to its last snapshot is pending"} + session.UpdatedAt = timestamppb.Now() + row.Data, err = marshalSession(session) + if err != nil { + return err + } + if err := execSQL(ctx, tx, `UPDATE session SET data = $2, dispatch_id = NULL, dispatch_expires_at = NULL WHERE id = $1`, row.ID, row.Data); err != nil { + return err + } + if err := settleCrashedSessionTask(ctx, tx, row, snapshot); err != nil { + return err + } + operation, err = toSessionOperation(row) + return err + }) + return operation, err +} + +// settleCrashedSessionTask retains captured waits and fails discarded nonterminal work. +// Callers hold the Session lock; recovery cannot checkpoint discarded state. +func settleCrashedSessionTask(ctx context.Context, tx pgx.Tx, session sessionRow, snapshot SessionTaskSnapshot) error { + if err := execSQL(ctx, tx, `UPDATE session_task_event SET quiescence_pending = FALSE WHERE history_id = $1 AND quiescence_pending`, session.HistoryID); err != nil { + return err + } + taskIDs, err := queryMany(ctx, tx, ` + SELECT t.id FROM session_task t WHERE t.history_id = $1 + AND (t.state IN ('TASK_STATE_SUBMITTED', 'TASK_STATE_WORKING') + OR (t.state IN ('TASK_STATE_INPUT_REQUIRED', 'TASK_STATE_AUTH_REQUIRED') + AND NOT EXISTS (SELECT 1 FROM session_task_event captured + WHERE captured.history_id = t.history_id AND $2::text <> '' + AND captured.snapshot_uri = $2 AND captured.snapshot_fidelity = $3 + AND captured.snapshot_atespace = $4 + AND captured.sequence >= (SELECT MAX(sequence) FROM session_task_event latest + WHERE latest.history_id = t.history_id AND latest.task_id = t.id)))) + `, pgx.RowTo[string], session.HistoryID, snapshot.URI, snapshot.Fidelity, snapshot.Atespace) + if err != nil { + return fmt.Errorf("read interrupted Session tasks: %w", err) + } + for _, taskID := range taskIDs { + version, err := taskVersion(ctx, tx, session.HistoryID, taskID) + if err != nil { + return err + } + pending, err := queryOne(ctx, tx, ` + SELECT EXISTS (SELECT 1 FROM session_task_event + WHERE sequence = $1 AND NOT published AND quiescence_pending IS NOT NULL) + `, pgx.RowTo[bool], version) + if err != nil { + return err + } + if pending { + if err := settleSessionTask(ctx, tx, session, taskID, version); err != nil { + return err + } + } + stored, err := readSessionTask(ctx, tx, session.HistoryID, taskID) + if err != nil { + return err + } + task, err := unmarshalSessionTask(stored.Data) + if err != nil { + return err + } + if !task.Status.State.Terminal() { + task.Status = a2a.TaskStatus{State: a2a.TaskStateFailed, Timestamp: new(time.Now().UTC()), + Message: a2a.NewMessage(a2a.MessageRoleAgent, a2a.NewTextPart("The runtime crashed before this turn completed. Send a new message to continue."))} + event := &a2a.TaskStatusUpdateEvent{TaskID: task.ID, ContextID: task.ContextID, Status: task.Status} + if err := storeSessionTaskEvent(ctx, tx, session, task, event, false); err != nil { + return err + } + version, err = taskVersion(ctx, tx, session.HistoryID, taskID) + if err != nil { + return err + } + if err := publishPushBoundary(ctx, tx, session.HistoryID, taskID, version, task); err != nil { + return err + } + } + } + return nil +} diff --git a/go/core/internal/database/session_recovery_test.go b/go/core/internal/database/session_recovery_test.go new file mode 100644 index 0000000000..68827b938d --- /dev/null +++ b/go/core/internal/database/session_recovery_test.go @@ -0,0 +1,255 @@ +package database + +import ( + "testing" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/stretchr/testify/require" + "google.golang.org/protobuf/proto" + + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" +) + +func TestCrashedRecoverySettlesTasksAtomically(t *testing.T) { + for _, test := range []struct { + name string + state, want a2a.TaskState + published, captured bool + }{ + {name: "submitted", state: a2a.TaskStateSubmitted, want: a2a.TaskStateFailed}, + {name: "working", state: a2a.TaskStateWorking, want: a2a.TaskStateFailed}, + {name: "completed before publication", state: a2a.TaskStateCompleted, want: a2a.TaskStateCompleted}, + {name: "unpublished input wait", state: a2a.TaskStateInputRequired, want: a2a.TaskStateFailed}, + {name: "unpublished auth wait", state: a2a.TaskStateAuthRequired, want: a2a.TaskStateFailed}, + {name: "published input wait", state: a2a.TaskStateInputRequired, want: a2a.TaskStateFailed, published: true}, + {name: "published auth wait", state: a2a.TaskStateAuthRequired, want: a2a.TaskStateFailed, published: true}, + {name: "captured input wait", state: a2a.TaskStateInputRequired, want: a2a.TaskStateInputRequired, published: true, captured: true}, + {name: "captured auth wait", state: a2a.TaskStateAuthRequired, want: a2a.TaskStateAuthRequired, published: true, captured: true}, + } { + t.Run(test.name, func(t *testing.T) { + client := NewClient(setupTestDB(t), "public") + sessionFixture(t, client, t.Context(), "team-a", "revision", "assistant", "kagent") + session, _, err := client.CreateSession(t.Context(), newSessionRequest(uuid.NewString(), "assistant", "kagent", ""), uuid.NewString()) + require.NoError(t, err) + session, err = markSessionReady(t.Context(), client, session.Id, "runtime.example") + require.NoError(t, err) + task := newSessionTask(uuid.NewString(), "original input") + task.ContextID, task.Status.State = session.ContextId, test.state + version, err := client.CreateRuntimeTask(t.Context(), session.Id, taskMutationHash("create"), task, "") + require.NoError(t, err) + if test.published { + require.NoError(t, client.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + } + snapshot := SessionTaskSnapshot{} + if test.captured { + work, err := client.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + snapshot = SessionTaskSnapshot{Atespace: "team-a", URI: "s3://snapshots/captured", Fidelity: "SNAPSHOT_FIDELITY_VOLUMES"} + require.NoError(t, client.FinishSessionQuiescence(t.Context(), work, &snapshot)) + } + before, _, err := client.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + page, err := client.ListSessionRecoveryCandidates(t.Context(), "", 100) + require.NoError(t, err) + require.Empty(t, page) + + operation, err := client.BeginCrashedSessionRecovery(t.Context(), session.Id, lifecycleGeneration(t, client, session.Id), "actor-"+session.Id, snapshot) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_FAILED, operation.Instance.State) + require.Equal(t, apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME, operation.Instance.Operation) + require.Equal(t, SessionRuntimeCrashed, operation.Instance.Failure.Reason) + stored, err := client.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + require.Equal(t, test.want, stored.Status.State) + require.Equal(t, before.History, stored.History) + require.Equal(t, before.Artifacts, stored.Artifacts) + if test.want == test.state { + require.Equal(t, before.Status, stored.Status) + } + _, settledVersion, err := client.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + replayed, err := client.BeginCrashedSessionRecovery(t.Context(), session.Id, lifecycleGeneration(t, client, session.Id), "actor-"+session.Id, snapshot) + require.NoError(t, err) + require.Equal(t, operation.ID, replayed.ID) + _, replayedVersion, err := client.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + require.Equal(t, settledVersion, replayedVersion) + _, err = client.UpdateSessionTask(t.Context(), session.Id, version, taskMutationHash("late"), task, task, "") + if test.state.Terminal() { + require.ErrorIs(t, err, ErrFailedPrecondition) + } else { + require.ErrorIs(t, err, ErrConflict) + } + require.ErrorIs(t, client.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "next"), ErrConflict) + _, _, err = client.ReserveSessionCheckpoint(t.Context(), &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID)}, session.Creator, uuid.NewString()) + require.ErrorIs(t, err, ErrConflict) + page, err = client.ListSessionRecoveryCandidates(t.Context(), "", 100) + require.NoError(t, err) + require.Len(t, page, 1) + require.Equal(t, session.Id, page[0].Id) + }) + } +} + +func TestCrashedRecoveryAdmissionRespectsFences(t *testing.T) { + for _, test := range []struct { + name string + actorUID string + checkpoint, superseded, lifecycle bool + }{ + {name: "replacement actor", actorUID: "replacement"}, + {name: "checkpoint", checkpoint: true}, + {name: "revert after the crash was observed", superseded: true}, + {name: "lifecycle", lifecycle: true}, + } { + t.Run(test.name, func(t *testing.T) { + client := NewClient(setupTestDB(t), "public") + sessionFixture(t, client, t.Context(), "team-a", "revision", "assistant", "kagent") + session, task := waitingTaskFixture(t, client) + uid := "actor-" + session.Id + // Callers read the generation before they observe the crash. + generation := lifecycleGeneration(t, client, session.Id) + switch { + case test.actorUID != "": + uid = test.actorUID + case test.checkpoint: + task.Status.State = a2a.TaskStateCompleted + require.NoError(t, saveRuntimeTask(t, client, session.Id, task, task, &SessionTaskSnapshot{ + Atespace: "team-a", URI: "s3://snapshots/completed", Fidelity: "SNAPSHOT_FIDELITY_VOLUMES", + })) + _, _, err := client.ReserveSessionCheckpoint(t.Context(), &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID)}, session.Creator, uuid.NewString()) + require.NoError(t, err) + case test.superseded: + // Another revert finishes on the same Actor UID, and a new turn + // starts once ingress wakes the restored Actor. + completeCrashedRecovery(t, client, session.Id, uid) + task = newSessionTask(uuid.NewString(), "next input") + task.ContextID, task.Status.State = session.ContextId, a2a.TaskStateWorking + _, err := client.CreateRuntimeTask(t.Context(), session.Id, taskMutationHash("next"), task, "") + require.NoError(t, err) + case test.lifecycle: + _, err := client.BeginSessionOperation(t.Context(), session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_SUSPEND) + require.NoError(t, err) + } + before, err := client.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + _, err = client.BeginCrashedSessionRecovery(t.Context(), session.Id, generation, uid, SessionTaskSnapshot{}) + require.ErrorIs(t, err, ErrConflict) + after, err := client.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + require.True(t, proto.Equal(before, after)) + if test.superseded { + running, _, err := client.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateWorking, running.Status.State, "a stale observation must not fail the new turn") + } + }) + } +} + +// A crash during post-turn suspension leaves the worker's claim issued, because +// the worker cannot tell whether its suspension took effect. The confirmed crash +// settles that: admission retires the claim instead of waiting for it. +func TestCrashedRecoveryRetiresIssuedIdleSuspension(t *testing.T) { + client := NewClient(setupTestDB(t), "public") + sessionFixture(t, client, t.Context(), "team-a", "revision", "assistant", "kagent") + session, task := waitingTaskFixture(t, client) + _, version, err := client.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + for _, state := range []a2a.TaskState{a2a.TaskStateWorking, a2a.TaskStateInputRequired} { + task.Status.State = state + version, err = client.UpdateSessionTask(t.Context(), session.Id, version, taskMutationHash(string(state)), task, task, "") + require.NoError(t, err) + } + require.NoError(t, client.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + work, err := client.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + uid := "actor-" + session.Id + + operation, err := client.BeginCrashedSessionRecovery(t.Context(), session.Id, lifecycleGeneration(t, client, session.Id), uid, SessionTaskSnapshot{}) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME, operation.Instance.Operation) + pending, err := queryOne(t.Context(), client.db, `SELECT quiescence_pending FROM session_task_event WHERE sequence = $1`, pgx.RowTo[bool], work.Version) + require.NoError(t, err) + require.False(t, pending, "admission retires the claim in its own transaction") + finishCrashedRecovery(t, client, session.Id, operation, uid) + + late := SessionTaskSnapshot{Atespace: "team-a", URI: "s3://snapshots/late", Fidelity: "SNAPSHOT_FIDELITY_VOLUMES"} + require.NoError(t, client.FinishSessionQuiescence(t.Context(), work, &late)) + row, err := readSession(t.Context(), client.db, session.Id) + require.NoError(t, err) + recorded, err := queryOne(t.Context(), client.db, ` + SELECT EXISTS (SELECT 1 FROM session_task WHERE history_id = $1 AND snapshot_uri = $2) + OR EXISTS (SELECT 1 FROM session_task_event WHERE history_id = $1 AND snapshot_uri = $2) + `, pgx.RowTo[bool], row.HistoryID, late.URI) + require.NoError(t, err) + require.False(t, recorded, "a late finish of the retired claim records nothing") + _, err = client.ClaimSessionQuiescence(t.Context()) + require.ErrorIs(t, err, ErrNotFound) + next := newSessionTask(uuid.NewString(), "next input") + next.ContextID, next.Status.State = session.ContextId, a2a.TaskStateWorking + _, err = client.CreateRuntimeTask(t.Context(), session.Id, taskMutationHash("next"), next, "") + require.NoError(t, err, "the retired claim no longer blocks a new turn") +} + +func lifecycleGeneration(t *testing.T, client *Client, sessionID string) uuid.UUID { + t.Helper() + generation, err := client.SessionLifecycleGeneration(t.Context(), sessionID) + require.NoError(t, err) + return generation +} + +// completeCrashedRecovery admits and finishes a revert as a successful concurrent +// caller would, leaving the Session READY on the same Actor UID. +func completeCrashedRecovery(t *testing.T, client *Client, sessionID, actorUID string) { + t.Helper() + operation, err := client.BeginCrashedSessionRecovery(t.Context(), sessionID, lifecycleGeneration(t, client, sessionID), actorUID, SessionTaskSnapshot{}) + require.NoError(t, err) + finishCrashedRecovery(t, client, sessionID, operation, actorUID) +} + +func finishCrashedRecovery(t *testing.T, client *Client, sessionID string, operation *SessionOperation, actorUID string) { + t.Helper() + executor := uuid.New() + claimed, err := client.ClaimSessionOperation(t.Context(), sessionID, operation.ID, executor) + require.NoError(t, err) + require.True(t, claimed) + recovered, err := client.FinishSessionOperation(t.Context(), sessionID, operation.ID, executor, "", actorUID, "") + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, recovered.State) +} + +func TestCrashedRecoveryAdmissionRollsBackOnMalformedTask(t *testing.T) { + client := NewClient(setupTestDB(t), "public") + sessionFixture(t, client, t.Context(), "team-a", "revision", "assistant", "kagent") + session, task := waitingTaskFixture(t, client) + _, version, err := client.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + for _, state := range []a2a.TaskState{a2a.TaskStateWorking, a2a.TaskStateInputRequired} { + task.Status.State = state + version, err = client.UpdateSessionTask(t.Context(), session.Id, version, taskMutationHash(string(state)), task, task, "") + require.NoError(t, err) + } + require.NoError(t, client.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + row, err := readSession(t.Context(), client.db, session.Id) + require.NoError(t, err) + before, err := readSessionTask(t.Context(), client.db, row.HistoryID, string(task.ID)) + require.NoError(t, err) + // A corrupt projection makes settlement fail after the lifecycle write. + _, err = client.db.Exec(t.Context(), `UPDATE session_task SET data = $2 WHERE history_id = $1 AND id = $3`, row.HistoryID, []byte{0xff}, string(task.ID)) + require.NoError(t, err) + _, err = client.BeginCrashedSessionRecovery(t.Context(), session.Id, lifecycleGeneration(t, client, session.Id), "actor-"+session.Id, SessionTaskSnapshot{}) + require.Error(t, err) + current, err := client.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, current.State) + require.Equal(t, apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE, current.Operation) + require.Nil(t, current.Failure) + _, err = client.db.Exec(t.Context(), `UPDATE session_task SET data = $2 WHERE history_id = $1 AND id = $3`, row.HistoryID, before.Data, string(task.ID)) + require.NoError(t, err) + work, err := client.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + require.Equal(t, string(task.ID), work.TaskID) +} diff --git a/go/core/internal/database/task_finalization.go b/go/core/internal/database/task_finalization.go index 9907257b21..ac209c8dfe 100644 --- a/go/core/internal/database/task_finalization.go +++ b/go/core/internal/database/task_finalization.go @@ -23,54 +23,60 @@ func (c *Client) SettleSessionTask(ctx context.Context, sessionID, taskID string if session.State == "RUNTIME_STATE_DELETED" { return ErrNotFound } - type boundary struct { - Data []byte - ExpectedVersion int64 - Published bool - } - row, err := queryOne(ctx, tx, ` - SELECT data, expected_version, published FROM session_task_event - WHERE history_id = $1 AND task_id = $2 AND sequence = $3 - AND expected_version IS NOT NULL AND quiescence_pending IS NOT NULL - `, pgx.RowToStructByName[boundary], session.HistoryID, taskID, version) - if errors.Is(err, pgx.ErrNoRows) { - return ErrConflict - } - if err != nil { - return err - } - if row.Published { - return nil - } - stored, err := readSessionTask(ctx, tx, session.HistoryID, taskID) - if err != nil { - return err - } - wire, err := applyStoredBoundary(stored.Data, row.Data) - if err != nil { - return err - } - task, err := pbconv.FromProtoTask(wire) - if err != nil { - return err - } - data, err := proto.Marshal(wire) - if err != nil { - return err - } - if _, err := saveTaskProjection(ctx, tx, session.HistoryID, taskID, string(task.Status.State), task.Status.Timestamp, data); err != nil { - return err - } - if err := execSQL(ctx, tx, ` - UPDATE session_task_event SET published = TRUE - WHERE history_id = $1 AND task_id = $2 AND sequence > $3 AND sequence <= $4 - `, session.HistoryID, taskID, row.ExpectedVersion, version); err != nil { - return err - } - return publishPushBoundary(ctx, tx, session.HistoryID, taskID, version, task) + return settleSessionTask(ctx, tx, session, taskID, version) }) } +// settleSessionTask publishes an already saved boundary under the Session row lock. +// Recovery uses the same publication after confirming the runtime crashed. +func settleSessionTask(ctx context.Context, tx pgx.Tx, session sessionRow, taskID string, version int64) error { + type boundary struct { + Data []byte + ExpectedVersion int64 + Published bool + } + row, err := queryOne(ctx, tx, ` + SELECT data, expected_version, published FROM session_task_event + WHERE history_id = $1 AND task_id = $2 AND sequence = $3 + AND expected_version IS NOT NULL AND quiescence_pending IS NOT NULL + `, pgx.RowToStructByName[boundary], session.HistoryID, taskID, version) + if errors.Is(err, pgx.ErrNoRows) { + return ErrConflict + } + if err != nil { + return err + } + if row.Published { + return nil + } + stored, err := readSessionTask(ctx, tx, session.HistoryID, taskID) + if err != nil { + return err + } + wire, err := applyStoredBoundary(stored.Data, row.Data) + if err != nil { + return err + } + task, err := pbconv.FromProtoTask(wire) + if err != nil { + return err + } + data, err := proto.Marshal(wire) + if err != nil { + return err + } + if _, err := saveTaskProjection(ctx, tx, session.HistoryID, taskID, string(task.Status.State), task.Status.Timestamp, data); err != nil { + return err + } + if err := execSQL(ctx, tx, ` + UPDATE session_task_event SET published = TRUE + WHERE history_id = $1 AND task_id = $2 AND sequence > $3 AND sequence <= $4 + `, session.HistoryID, taskID, row.ExpectedVersion, version); err != nil { + return err + } + return publishPushBoundary(ctx, tx, session.HistoryID, taskID, version, task) +} + func applyStoredBoundary(data, eventData []byte) (*a2apb.Task, error) { task, event := &a2apb.Task{}, &a2apb.StreamResponse{} if err := proto.Unmarshal(data, task); err != nil { diff --git a/go/core/internal/substrate/client.go b/go/core/internal/substrate/client.go index 4494b87c6e..798e6f72f1 100644 --- a/go/core/internal/substrate/client.go +++ b/go/core/internal/substrate/client.go @@ -200,6 +200,16 @@ func (c *Client) ResumeActor(ctx context.Context, atespace, actorID string) (*at return resp.GetActor(), nil } +func (c *Client) RevertActor(ctx context.Context, atespace, actorID string) (*ateapipb.Actor, error) { + ctx, cancel := c.callCtx(ctx) + defer cancel() + resp, err := c.ControlClient.RevertActor(ctx, &ateapipb.RevertActorRequest{Actor: actorRef(atespace, actorID)}) + if err != nil { + return nil, fmt.Errorf("failed to revert Actor %s/%s: %w", atespace, actorID, err) + } + return resp.GetActor(), nil +} + // SuspendActor waits for durable storage, including an upload left by an earlier // attempt. Substrate's synchronous mode returns already-suspended Actors as-is. func (c *Client) SuspendActor(ctx context.Context, atespace, actorID string) (*ateapipb.Actor, error) { diff --git a/go/core/internal/substrate/snapshot.go b/go/core/internal/substrate/snapshot.go index 02434ab83d..de2597052b 100644 --- a/go/core/internal/substrate/snapshot.go +++ b/go/core/internal/substrate/snapshot.go @@ -22,3 +22,14 @@ func SnapshotURI(snapshot *ateapipb.Snapshot) string { } return storage.GetObject().GetSnapshotUri() } + +// Checkpoints require the newest generation; revert retains the newest durable copy. +func LatestDurableActorSnapshot(actor *ateapipb.Actor) *ateapipb.Snapshot { + var latest *ateapipb.Snapshot + for _, snapshot := range actor.GetStatus().GetSnapshots() { + if SnapshotURI(snapshot) != "" && snapshot.GetGeneration() > latest.GetGeneration() { + latest = snapshot + } + } + return latest +} diff --git a/go/core/internal/substrate/snapshot_test.go b/go/core/internal/substrate/snapshot_test.go index 0d29f25f65..647e4aa903 100644 --- a/go/core/internal/substrate/snapshot_test.go +++ b/go/core/internal/substrate/snapshot_test.go @@ -40,3 +40,23 @@ func TestActorSnapshotGenerationAndDurability(t *testing.T) { require.Nil(t, LatestActorSnapshot(nil)) require.Empty(t, SnapshotURI(nil)) } + +func TestRevertSnapshotSkipsIncompleteGenerations(t *testing.T) { + durable := &ateapipb.Snapshot{Generation: 2, DurableSnapshot: &ateapipb.SnapshotStorage{ + Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_COMPLETED, + Object: &ateapipb.ObjectSnapshot{SnapshotUri: "s3://snapshots/durable"}, + }} + for _, latest := range []*ateapipb.Snapshot{ + {Generation: 3, Locality: "node-1"}, + {Generation: 3, DurableSnapshot: &ateapipb.SnapshotStorage{Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_PENDING}}, + {Generation: 3, DurableSnapshot: &ateapipb.SnapshotStorage{Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_COMPLETED}}, + } { + actor := &ateapipb.Actor{Status: &ateapipb.ActorStatus{Snapshots: []*ateapipb.Snapshot{latest, durable}}} + require.Same(t, durable, LatestDurableActorSnapshot(actor)) + require.Same(t, latest, LatestActorSnapshot(actor)) + } + require.Nil(t, LatestDurableActorSnapshot(nil)) + require.Nil(t, LatestDurableActorSnapshot(&ateapipb.Actor{Status: &ateapipb.ActorStatus{ + Snapshots: []*ateapipb.Snapshot{{Generation: 1, Locality: "node-1"}}, + }})) +} From 0d206ce74fa43734030c1cd9c39f59935b2f860c Mon Sep 17 00:00:00 2001 From: Ada Lundhe Date: Fri, 23 Oct 2026 12:31:55 -0700 Subject: [PATCH 02/12] fix(session): implement manual crashed-runtime recovery Signed-off-by: Ada Lundhe --- .../internal/service/session/quiescence.go | 21 +- .../service/session/quiescence_test.go | 4 + go/core/internal/service/session/recovery.go | 154 +++++ .../internal/service/session/recovery_test.go | 541 ++++++++++++++++++ go/core/internal/service/session/workflow.go | 7 + 5 files changed, 723 insertions(+), 4 deletions(-) create mode 100644 go/core/internal/service/session/recovery.go create mode 100644 go/core/internal/service/session/recovery_test.go diff --git a/go/core/internal/service/session/quiescence.go b/go/core/internal/service/session/quiescence.go index 283e70c7cc..aba8c40ffb 100644 --- a/go/core/internal/service/session/quiescence.go +++ b/go/core/internal/service/session/quiescence.go @@ -28,16 +28,29 @@ func (w *ActorWorkflow) wakeQuiescence() { } } -// Start suspends idle sessions to VOLUMES independently of task publication. -// Startup and periodic recovery scans discover work missed by local signals. -// Workers drain claims with bounded concurrency; task reads never wait for them. +// Start suspends idle sessions to VOLUMES independently of task publication and +// retries crash recoveries admitted by explicit requests. +// Startup and periodic scans discover idle work missed by local signals; crash +// checks use the same recovery interval. Workers drain claims with bounded concurrency. func (w *ActorWorkflow) Start(ctx context.Context) error { var workers sync.WaitGroup - for range 4 { + for worker := range 4 { workers.Go(func() { timer := time.NewTimer(w.recoveryInterval) defer timer.Stop() + var recoveryAfter string + nextRecovery := time.Now().Add(w.recoveryInterval) for ctx.Err() == nil { + if worker == 0 && !time.Now().Before(nextRecovery) { + scanCtx, cancel := context.WithTimeout(ctx, database.RuntimeOperationTimeout) + after, err := w.recoverSessionPage(scanCtx, recoveryAfter) + cancel() + if err != nil && ctx.Err() == nil { + logging.FromContext(ctx).ErrorContext(ctx, "scan interrupted sessions", "error", err) + } + recoveryAfter = after + nextRecovery = time.Now().Add(w.recoveryInterval) + } work, err := w.store.ClaimSessionQuiescence(ctx) if err == nil { // Hand off a hint before runtime I/O so a single settlement diff --git a/go/core/internal/service/session/quiescence_test.go b/go/core/internal/service/session/quiescence_test.go index 99e8852d72..d6b9af4991 100644 --- a/go/core/internal/service/session/quiescence_test.go +++ b/go/core/internal/service/session/quiescence_test.go @@ -32,6 +32,10 @@ type idleQuiescenceStore struct { var _ workflowStore = (*idleQuiescenceStore)(nil) +func (s *idleQuiescenceStore) ListSessionRecoveryCandidates(context.Context, string, int) ([]*apiv1.Session, error) { + return nil, nil +} + func (s *idleQuiescenceStore) ClaimSessionQuiescence(context.Context) (*database.SessionQuiescence, error) { s.claims.Add(1) if s.claimErr != nil { diff --git a/go/core/internal/service/session/recovery.go b/go/core/internal/service/session/recovery.go new file mode 100644 index 0000000000..27922bc3ed --- /dev/null +++ b/go/core/internal/service/session/recovery.go @@ -0,0 +1,154 @@ +package session + +import ( + "context" + "errors" + "fmt" + "time" + + "github.com/google/uuid" + ateapipb "go.ate.dev/substrate/api/ateapi/v1" + "google.golang.org/protobuf/proto" + + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" + "github.com/kagent-dev/kagent/go/core/internal/database" + "github.com/kagent-dev/kagent/go/core/internal/substrate" + "github.com/kagent-dev/kagent/go/pkg/logging" +) + +// recoverSessionPage retries admitted recoveries in ID order. +// READY Sessions are excluded, so a background retry cannot admit new work. +func (w *ActorWorkflow) recoverSessionPage(ctx context.Context, afterID string) (string, error) { + const pageSize = 100 + sessions, err := w.store.ListSessionRecoveryCandidates(ctx, afterID, pageSize) + if err != nil { + return afterID, err + } + for _, session := range sessions { + attemptCtx, cancel := context.WithTimeout(ctx, database.RuntimeOperationTimeout) + err := w.recoverSessionCandidate(attemptCtx, session) + cancel() + if err != nil && !errors.Is(err, database.ErrConflict) && ctx.Err() == nil { + logging.FromContext(ctx).ErrorContext(ctx, "recover crashed session", "session_id", session.Id, "error", err) + } + afterID = session.Id + if ctx.Err() != nil { + return afterID, ctx.Err() + } + } + if len(sessions) < pageSize { + afterID = "" + } + return afterID, nil +} + +func (w *ActorWorkflow) recoverSessionCandidate(ctx context.Context, session *apiv1.Session) error { + _, err := w.Revert(ctx, session) + if errors.Is(err, database.ErrFailedPrecondition) { + return nil + } + return err +} + +// Revert restores a confirmed crashed runtime without replaying its interrupted turn. +// Retries join the same durable operation and preserve the Actor identity. +func (w *ActorWorkflow) Revert(ctx context.Context, session *apiv1.Session) (*apiv1.Session, error) { + ctx, cancel := context.WithTimeout(ctx, database.RuntimeOperationTimeout) + defer cancel() + if session.State == apiv1.RuntimeState_RUNTIME_STATE_FAILED && session.Failure.GetReason() == database.SessionRuntimeCrashed { + return w.run(ctx, session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME) + } + if session.State != apiv1.RuntimeState_RUNTIME_STATE_READY || session.Operation != apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE { + return nil, database.ErrConflict + } + revision, err := w.store.GetRuntimeRevision(ctx, session.PreparedRevision) + if err != nil { + return nil, err + } + // Read before observing Substrate. A concurrent revert keeps the Actor UID and + // returns to READY, so only the generation shows that this crash observation + // predates it. + generation, err := w.store.SessionLifecycleGeneration(ctx, session.Id) + if err != nil { + return nil, err + } + actor, err := w.actors.GetActor(ctx, revision.ActorTemplateAtespace, substrate.ActorName(session.Id)) + if err != nil { + return nil, err + } + if err := w.verifyActor(ctx, session, revision, actor); err != nil { + return nil, err + } + if actor.GetStatus().GetState() != ateapipb.ActorState_ACTOR_STATE_CRASHED { + return nil, fmt.Errorf("runtime is not crashed: %w", database.ErrFailedPrecondition) + } + snapshot := substrate.LatestDurableActorSnapshot(actor) + operation, err := w.store.BeginCrashedSessionRecovery(ctx, session.Id, generation, actor.GetMetadata().GetUid(), database.SessionTaskSnapshot{ + Atespace: actor.GetMetadata().GetAtespace(), URI: substrate.SnapshotURI(snapshot), + Fidelity: snapshot.GetFidelity().String(), + }) + if err != nil { + return nil, err + } + return w.execute(ctx, operation) +} + +// recoverCrashedSession reconciles lost revert responses without replacing the Actor. +func (w *ActorWorkflow) recoverCrashedSession(ctx context.Context, operation *database.SessionOperation) (_ *apiv1.Session, err error) { + session := operation.Instance + if session.Operation != apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME { + return nil, database.ErrConflict + } + executorID := uuid.New() + claimed, err := w.store.ClaimSessionOperation(ctx, session.Id, operation.ID, executorID) + if err != nil { + return nil, err + } + if !claimed { + current, err := w.store.GetSessionOperation(ctx, session.Id, operation.ID) + if err != nil { + return nil, err + } + return operationOutcome(current) + } + defer func() { + if err != nil { + releaseCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + defer cancel() + err = errors.Join(err, w.store.ReleaseRuntimeOperation(releaseCtx, session.Id, operation.ID, executorID)) + } + }() + revision, err := w.store.GetRuntimeRevision(ctx, session.PreparedRevision) + if err != nil { + return nil, err + } + atespace, name := revision.ActorTemplateAtespace, substrate.ActorName(session.Id) + actor, err := w.actors.GetActor(ctx, atespace, name) + if err != nil { + return nil, err + } + if err := w.verifyActor(ctx, session, revision, actor); err != nil { + return nil, err + } + before := actor + beforeSnapshot := proto.CloneOf(substrate.LatestDurableActorSnapshot(actor)) + if beforeSnapshot != nil { + beforeSnapshot.Locality = "" + } + switch actor.GetStatus().GetState() { + case ateapipb.ActorState_ACTOR_STATE_CRASHED, ateapipb.ActorState_ACTOR_STATE_REVERTING: + actor, err = w.actors.RevertActor(ctx, atespace, name) + if err != nil { + return nil, fmt.Errorf("revert crashed Actor %s/%s: %w", atespace, name, err) + } + case ateapipb.ActorState_ACTOR_STATE_SUSPENDED: + default: + return nil, fmt.Errorf("crashed Actor %s/%s has unexpected state %s", atespace, name, actor.GetStatus().GetState()) + } + if !validActorIdentity(actor, revision, name) || actor.GetMetadata().GetUid() != before.GetMetadata().GetUid() || actor.GetStatus().GetState() != ateapipb.ActorState_ACTOR_STATE_SUSPENDED || !proto.Equal(beforeSnapshot, substrate.LatestActorSnapshot(actor)) { + return nil, fmt.Errorf("revert Actor %s/%s returned unexpected identity or state", atespace, name) + } + finishCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + defer cancel() + return w.store.FinishSessionOperation(finishCtx, session.Id, operation.ID, executorID, "", actor.GetMetadata().GetUid(), "") +} diff --git a/go/core/internal/service/session/recovery_test.go b/go/core/internal/service/session/recovery_test.go new file mode 100644 index 0000000000..d7ce338a83 --- /dev/null +++ b/go/core/internal/service/session/recovery_test.go @@ -0,0 +1,541 @@ +package session + +import ( + "context" + "crypto/sha256" + "fmt" + "slices" + "testing" + "time" + + "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/grpc/codes" + "google.golang.org/grpc/status" + "google.golang.org/protobuf/proto" + + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" + "github.com/kagent-dev/kagent/go/core/internal/database" + "github.com/kagent-dev/kagent/go/core/internal/substrate" +) + +func (a *lifecycleTestActors) RevertActor(_ context.Context, atespace, name string) (*ateapipb.Actor, error) { + a.mu.Lock() + defer a.mu.Unlock() + actor := a.actors[actorKey(atespace, name)] + if actor.Status.State != ateapipb.ActorState_ACTOR_STATE_CRASHED && actor.Status.State != ateapipb.ActorState_ACTOR_STATE_REVERTING { + return nil, status.Error(codes.FailedPrecondition, "actor is not crashed or reverting") + } + actor.Status.State = ateapipb.ActorState_ACTOR_STATE_SUSPENDED + actor.Status.Crash = nil + snapshot := proto.CloneOf(substrate.LatestDurableActorSnapshot(actor)) + actor.Status.Snapshots = nil + if snapshot != nil { + snapshot.Locality = "" + actor.Status.Snapshots = []*ateapipb.Snapshot{snapshot} + } + return proto.CloneOf(actor), nil +} + +type recoveryTestActors struct { + *lifecycleTestActors + revert func(context.Context, string, string) (*ateapipb.Actor, error) + readErr error + // afterRead runs once, after the first successful read returns its copy. + afterRead func(context.Context) +} + +func (a *recoveryTestActors) RevertActor(ctx context.Context, atespace, name string) (*ateapipb.Actor, error) { + return a.revert(ctx, atespace, name) +} + +func (a *recoveryTestActors) GetActor(ctx context.Context, atespace, name string) (*ateapipb.Actor, error) { + if a.readErr != nil { + return nil, a.readErr + } + actor, err := a.lifecycleTestActors.GetActor(ctx, atespace, name) + if hook := a.afterRead; hook != nil && err == nil { + a.afterRead = nil + hook(ctx) + } + return actor, err +} + +func revertGeneration(t *testing.T, store *lifecycleTestStore, sessionID string) uuid.UUID { + t.Helper() + generation, err := store.SessionLifecycleGeneration(t.Context(), sessionID) + require.NoError(t, err) + return generation +} + +func crashedTurnFixture(t *testing.T) (*lifecycleTestStore, *lifecycleTestActors, *apiv1.Session, *a2a.Task, int64) { + t.Helper() + store, session := lifecycleFixture(t) + actors := &lifecycleTestActors{actors: map[string]*ateapipb.Actor{}} + var err error + session, err = NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Second).Create(t.Context(), session) + require.NoError(t, err) + message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("original input")) + message.ContextID = session.ContextId + task := a2a.NewSubmittedTask(message, message) + task.Status.State = a2a.TaskStateWorking + hash := sha256.Sum256([]byte("working")) + version, err := store.CreateRuntimeTask(t.Context(), session.Id, hash[:], task, "") + require.NoError(t, err) + actors.actors[actorKey("team-a", substrate.ActorName(session.Id))].Status.State = ateapipb.ActorState_ACTOR_STATE_CRASHED + return store, actors, session, task, version +} + +func TestCrashedSessionRecoveryRetriesWithoutLosingState(t *testing.T) { + for _, test := range []struct { + name string + lostResponse bool + wantReverts int + }{ + {name: "revert failed", wantReverts: 2}, + {name: "revert response lost", lostResponse: true, wantReverts: 1}, + } { + t.Run(test.name, func(t *testing.T) { + store, base, session, task, _ := crashedTurnFixture(t) + calls := 0 + actors := &recoveryTestActors{lifecycleTestActors: base} + actors.revert = func(ctx context.Context, atespace, name string) (*ateapipb.Actor, error) { + calls++ + if calls == 1 { + if test.lostResponse { + _, err := base.RevertActor(ctx, atespace, name) + require.NoError(t, err) + } + return nil, status.Error(codes.Unavailable, "response unavailable") + } + return base.RevertActor(ctx, atespace, name) + } + workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Second) + _, err := workflow.Revert(t.Context(), session) + require.Error(t, err) + failed, err := store.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_FAILED, failed.State) + require.Equal(t, apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME, failed.Operation) + require.Equal(t, database.SessionRuntimeCrashed, failed.Failure.Reason) + interrupted, interruptedVersion, err := store.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateFailed, interrupted.Status.State) + published, err := store.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateFailed, published.Status.State) + require.Len(t, published.History, 1) + require.ErrorIs(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), uuid.NewString()), database.ErrConflict) + // A new workflow instance discovers the admitted operation from PostgreSQL. + restarted := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Second) + _, err = restarted.recoverSessionPage(t.Context(), "") + require.NoError(t, err) + ready, err := store.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, ready.State) + require.Nil(t, ready.Failure) + stored, err := store.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateFailed, stored.Status.State) + require.Equal(t, test.wantReverts, calls) + require.Len(t, stored.History, 1, "the original input remains archived once") + _, recoveredVersion, err := store.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + require.Equal(t, interruptedVersion, recoveredVersion, "retries do not append another terminal event") + _, err = restarted.Revert(t.Context(), ready) + require.ErrorIs(t, err, database.ErrFailedPrecondition) + require.Equal(t, test.wantReverts, calls, "completed recovery cannot discard another execution") + }) + } +} + +func TestCrashedSessionRecoveryFencesConcurrentWork(t *testing.T) { + store, base, session, task, version := crashedTurnFixture(t) + entered, release := make(chan struct{}), make(chan struct{}) + actors := &recoveryTestActors{lifecycleTestActors: base} + actors.revert = func(ctx context.Context, atespace, name string) (*ateapipb.Actor, error) { + close(entered) + select { + case <-release: + return base.RevertActor(ctx, atespace, name) + case <-ctx.Done(): + return nil, ctx.Err() + } + } + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Second) + done := make(chan error, 1) + go func() { _, err := workflow.Revert(ctx, session); done <- err }() + select { + case <-entered: + case <-time.After(5 * time.Second): + t.Fatal("recovery was not admitted") + } + failed, err := store.GetSessionByID(ctx, session.Id) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_FAILED, failed.State) + _, err = workflow.Revert(ctx, failed) + require.ErrorIs(t, err, database.ErrConflict) + hash := sha256.Sum256([]byte("late write")) + _, err = store.UpdateSessionTask(ctx, session.Id, version, hash[:], task, task, "") + require.ErrorIs(t, err, database.ErrConflict) + close(release) + require.NoError(t, <-done) + _, err = store.UpdateSessionTask(ctx, session.Id, version, hash[:], task, task, "") + require.ErrorIs(t, err, database.ErrConflict) +} + +func TestSessionRecoveryRequiresConfirmedRecordedActorCrash(t *testing.T) { + for _, test := range []struct { + name string + state ateapipb.ActorState + uid string + readErr error + wantErr bool + }{ + {name: "quiet running actor", state: ateapipb.ActorState_ACTOR_STATE_RUNNING, wantErr: true}, + {name: "transient read failure", readErr: status.Error(codes.Unavailable, "unreachable"), wantErr: true}, + {name: "same-name replacement", state: ateapipb.ActorState_ACTOR_STATE_CRASHED, uid: "another-actor", wantErr: true}, + } { + t.Run(test.name, func(t *testing.T) { + store, base, session, task, _ := crashedTurnFixture(t) + actor := base.actors[actorKey("team-a", substrate.ActorName(session.Id))] + actor.Status.State = test.state + if test.uid != "" { + actor.Metadata.Uid = test.uid + } + actors := &recoveryTestActors{lifecycleTestActors: base, readErr: test.readErr} + actors.revert = func(context.Context, string, string) (*ateapipb.Actor, error) { + t.Fatal("unconfirmed crash issued a revert") + return nil, nil + } + _, err := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Second).Revert(t.Context(), session) + if test.wantErr { + require.Error(t, err) + } else { + require.NoError(t, err) + } + current, err := store.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, current.State) + require.Nil(t, current.Failure) + stored, err := store.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateWorking, stored.Status.State) + }) + } +} + +func TestSessionRecoveryCandidatePolicy(t *testing.T) { + store, actors, active, _, _ := crashedTurnFixture(t) + workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Second) + for _, state := range []a2a.TaskState{"", a2a.TaskStateInputRequired, a2a.TaskStateAuthRequired, a2a.TaskStateWorking} { + session, _, err := store.CreateSession(t.Context(), &apiv1.Session{Id: uuid.NewString(), Creator: "alice", Agent: active.Agent}, uuid.NewString()) + require.NoError(t, err) + session, err = workflow.Create(t.Context(), session) + require.NoError(t, err) + if state == "" { + continue + } + message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("input")) + message.ContextID = session.ContextId + task := a2a.NewSubmittedTask(message, message) + task.Status.State = state + hash := sha256.Sum256([]byte(task.ID)) + version, err := store.CreateRuntimeTask(t.Context(), session.Id, hash[:], task, "") + require.NoError(t, err) + if state != a2a.TaskStateWorking { + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + } + } + page, err := store.ListSessionRecoveryCandidates(t.Context(), "", 100) + require.NoError(t, err) + require.Empty(t, page, "manual recovery does not inspect idle or active runtimes") + actor, err := actors.GetActor(t.Context(), "team-a", substrate.ActorName(active.Id)) + require.NoError(t, err) + _, err = store.BeginCrashedSessionRecovery(t.Context(), active.Id, revertGeneration(t, store, active.Id), actor.GetMetadata().GetUid(), database.SessionTaskSnapshot{}) + require.NoError(t, err) + page, err = store.ListSessionRecoveryCandidates(t.Context(), "", 100) + require.NoError(t, err) + require.Len(t, page, 1, "explicitly admitted recovery remains eligible for retry") + require.Equal(t, active.Id, page[0].Id) +} + +func TestCrashedSessionRecoversAnOlderResumedTask(t *testing.T) { + store, actors, session, older, version := crashedTurnFixture(t) + older.Status.State = a2a.TaskStateInputRequired + hash := sha256.Sum256([]byte("waiting")) + waitingVersion, err := store.UpdateSessionTask(t.Context(), session.Id, version, hash[:], older, older, "") + require.NoError(t, err) + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(older.ID), waitingVersion)) + work, err := store.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + require.NoError(t, store.FinishSessionQuiescence(t.Context(), work, &database.SessionTaskSnapshot{ + Atespace: "team-a", URI: "s3://snapshots/waiting-turn", Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES.String(), + })) + + message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("another turn")) + message.ContextID = session.ContextId + newer := a2a.NewSubmittedTask(message, message) + newer.Status.State = a2a.TaskStateCompleted + hash = sha256.Sum256([]byte("newer completed task")) + newerVersion, err := store.CreateRuntimeTask(t.Context(), session.Id, hash[:], newer, "") + require.NoError(t, err) + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(newer.ID), newerVersion)) + work, err = store.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + require.NoError(t, store.FinishSessionQuiescence(t.Context(), work, &database.SessionTaskSnapshot{ + Atespace: "team-a", URI: "s3://snapshots/newer-completed-turn", Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES.String(), + })) + + older.Status.State = a2a.TaskStateWorking + hash = sha256.Sum256([]byte("resume older task")) + _, err = store.UpdateSessionTask(t.Context(), session.Id, waitingVersion, hash[:], older, older, "") + require.NoError(t, err) + workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Second) + _, err = workflow.Revert(t.Context(), session) + require.NoError(t, err) + failed, err := store.GetSettledSessionTask(t.Context(), session.Id, string(older.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateFailed, failed.Status.State) + retained, retainedVersion, err := store.GetVersionedSessionTask(t.Context(), session.Id, string(newer.ID)) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateCompleted, retained.Status.State) + require.Equal(t, newerVersion, retainedVersion, "recovery cannot rewrite a later completed task") + require.NoError(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), uuid.NewString())) +} + +func TestExplicitRevertRespectsCheckpointCreation(t *testing.T) { + store, actors, session, task, version := crashedTurnFixture(t) + task.Status.State = a2a.TaskStateCompleted + hash := sha256.Sum256([]byte("completed")) + version, err := store.UpdateSessionTask(t.Context(), session.Id, version, hash[:], task, task, "") + require.NoError(t, err) + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + work, err := store.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + require.NoError(t, store.FinishSessionQuiescence(t.Context(), work, &database.SessionTaskSnapshot{ + Atespace: "team-a", URI: "s3://snapshots/completed", Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES.String(), + })) + _, _, err = store.ReserveSessionCheckpoint(t.Context(), &apiv1.Checkpoint{ + Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID), + }, session.Creator, uuid.NewString()) + require.NoError(t, err) + workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute) + _, err = workflow.Revert(t.Context(), session) + require.ErrorIs(t, err, database.ErrConflict) + actor, err := actors.GetActor(t.Context(), "team-a", substrate.ActorName(session.Id)) + require.NoError(t, err) + require.Equal(t, ateapipb.ActorState_ACTOR_STATE_CRASHED, actor.Status.State) + stored, err := store.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateCompleted, stored.Status.State) +} + +// The Actor crashed during post-turn suspension, and the worker kept its claim +// because it cannot tell whether the suspension took effect. +func TestExplicitRevertRetiresIdleSuspensionOfCrashedActor(t *testing.T) { + store, actors, session, task, version := crashedTurnFixture(t) + task.Status.State = a2a.TaskStateCompleted + hash := sha256.Sum256([]byte("completed")) + version, err := store.UpdateSessionTask(t.Context(), session.Id, version, hash[:], task, task, "") + require.NoError(t, err) + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + work, err := store.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute) + recovered, err := workflow.Revert(t.Context(), session) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, recovered.State) + actor, err := actors.GetActor(t.Context(), "team-a", substrate.ActorName(session.Id)) + require.NoError(t, err) + require.Equal(t, ateapipb.ActorState_ACTOR_STATE_SUSPENDED, actor.Status.State) + stored, err := store.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateCompleted, stored.Status.State) + require.NoError(t, store.FinishSessionQuiescence(t.Context(), work, &database.SessionTaskSnapshot{ + Atespace: "team-a", URI: "s3://snapshots/late", Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES.String(), + }), "a late finish of the retired claim is a no-op") + _, err = store.ClaimSessionQuiescence(t.Context()) + require.ErrorIs(t, err, database.ErrNotFound) +} + +// Revert A observes the crash, then pauses while Revert B restores the same Actor +// and ingress wakes it for a new turn. A then sees the same Actor UID and a READY +// Session, so only the lifecycle generation shows that its observation is stale. +func TestExplicitRevertRejectsCrashObservedBeforeAnotherRevert(t *testing.T) { + store, base, session, _, _ := crashedTurnFixture(t) + key := actorKey("team-a", substrate.ActorName(session.Id)) + var next *a2a.Task + paused := &recoveryTestActors{lifecycleTestActors: base, revert: base.RevertActor, afterRead: func(ctx context.Context) { + recovered, err := NewActorWorkflow(store, base, make(chan struct{}, 1), time.Minute).Revert(ctx, session) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, recovered.State) + base.mu.Lock() + base.actors[key].Status.State = ateapipb.ActorState_ACTOR_STATE_RUNNING + base.mu.Unlock() + message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("next input")) + message.ContextID = session.ContextId + next = a2a.NewSubmittedTask(message, message) + next.Status.State = a2a.TaskStateWorking + hash := sha256.Sum256([]byte("next")) + _, err = store.CreateRuntimeTask(ctx, session.Id, hash[:], next, "") + require.NoError(t, err) + }} + _, err := NewActorWorkflow(store, paused, make(chan struct{}, 1), time.Minute).Revert(t.Context(), session) + require.ErrorIs(t, err, database.ErrConflict) + current, err := store.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, current.State) + require.Equal(t, apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE, current.Operation) + running, _, err := store.GetVersionedSessionTask(t.Context(), session.Id, string(next.ID)) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateWorking, running.Status.State, "the stale revert must not fail the new turn") +} + +func TestExplicitRevertCannotCheckpointDiscardedTerminalState(t *testing.T) { + store, actors, session, task, version := crashedTurnFixture(t) + task.Status.State = a2a.TaskStateCompleted + hash := sha256.Sum256([]byte("completed before snapshot")) + version, err := store.UpdateSessionTask(t.Context(), session.Id, version, hash[:], task, task, "") + require.NoError(t, err) + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute) + _, err = workflow.Revert(t.Context(), session) + require.NoError(t, err) + _, err = store.ClaimSessionQuiescence(t.Context()) + require.ErrorIs(t, err, database.ErrNotFound, "the reverted snapshot cannot represent discarded terminal state") + _, _, err = store.ReserveSessionCheckpoint(t.Context(), &apiv1.Checkpoint{ + Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID), + }, session.Creator, uuid.NewString()) + require.ErrorIs(t, err, database.ErrSnapshotPending) + stored, err := store.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateCompleted, stored.Status.State) +} + +func TestExplicitRevertPreservesOnlyCapturedWaits(t *testing.T) { + for _, state := range []a2a.TaskState{a2a.TaskStateInputRequired, a2a.TaskStateAuthRequired} { + for _, captured := range []bool{false, true} { + t.Run(fmt.Sprint(state, "/captured=", captured), func(t *testing.T) { + store, actors, session, task, version := crashedTurnFixture(t) + task.Status.State = state + hash := sha256.Sum256([]byte(state)) + version, err := store.UpdateSessionTask(t.Context(), session.Id, version, hash[:], task, task, "") + require.NoError(t, err) + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + if captured { + work, err := store.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + require.NoError(t, store.FinishSessionQuiescence(t.Context(), work, &database.SessionTaskSnapshot{ + Atespace: "team-a", URI: "s3://snapshots/waiting", Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES.String(), + })) + actors.actors[actorKey("team-a", substrate.ActorName(session.Id))].Status.Snapshots = []*ateapipb.Snapshot{{Generation: 1, Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES, DurableSnapshot: &ateapipb.SnapshotStorage{Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_COMPLETED, Object: &ateapipb.ObjectSnapshot{SnapshotUri: "s3://snapshots/waiting"}}}} + actor := actors.actors[actorKey("team-a", substrate.ActorName(session.Id))] + actor.Status.Snapshots[0].Locality = "lost-node" + actor.Status.Snapshots = append(actor.Status.Snapshots, &ateapipb.Snapshot{ + Generation: 2, DurableSnapshot: &ateapipb.SnapshotStorage{Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_PENDING}, + }) + } + _, err = NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute).Revert(t.Context(), session) + require.NoError(t, err) + stored, err := store.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + want := a2a.TaskStateFailed + if captured { + want = state + } + require.Equal(t, want, stored.Status.State) + }) + } + } +} + +func TestExplicitRevertRetainsWaitCapturedByLaterTurn(t *testing.T) { + store, actors, session, waiting, version := crashedTurnFixture(t) + waiting.Status.State = a2a.TaskStateInputRequired + hash := sha256.Sum256([]byte("waiting")) + version, err := store.UpdateSessionTask(t.Context(), session.Id, version, hash[:], waiting, waiting, "") + require.NoError(t, err) + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(waiting.ID), version)) + message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("another turn")) + message.ContextID = session.ContextId + later := a2a.NewSubmittedTask(message, message) + later.Status.State = a2a.TaskStateCompleted + hash = sha256.Sum256([]byte("later completed turn")) + version, err = store.CreateRuntimeTask(t.Context(), session.Id, hash[:], later, "") + require.NoError(t, err) + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(later.ID), version)) + work, err := store.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + require.NoError(t, store.FinishSessionQuiescence(t.Context(), work, &database.SessionTaskSnapshot{ + Atespace: "team-a", URI: "s3://snapshots/later-turn", Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES.String(), + })) + actors.actors[actorKey("team-a", substrate.ActorName(session.Id))].Status.Snapshots = []*ateapipb.Snapshot{{Generation: 1, Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES, DurableSnapshot: &ateapipb.SnapshotStorage{Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_COMPLETED, Object: &ateapipb.ObjectSnapshot{SnapshotUri: "s3://snapshots/later-turn"}}}} + _, err = NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute).Revert(t.Context(), session) + require.NoError(t, err) + retained, err := store.GetSettledSessionTask(t.Context(), session.Id, string(waiting.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateInputRequired, retained.Status.State) + retained, err = store.GetSettledSessionTask(t.Context(), session.Id, string(later.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateCompleted, retained.Status.State) +} + +func TestSessionRecoveryPagesPastUnavailableActors(t *testing.T) { + store, base, first, _, _ := crashedTurnFixture(t) + actors := &recoveryTestActors{lifecycleTestActors: base, revert: base.RevertActor, readErr: status.Error(codes.Unavailable, "control plane unavailable")} + workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute) + ids := []string{first.Id} + for range 100 { + session, _, err := store.CreateSession(t.Context(), &apiv1.Session{ + Id: uuid.NewString(), Creator: first.Creator, Agent: first.Agent, + }, uuid.NewString()) + require.NoError(t, err) + session, err = NewActorWorkflow(store, base, make(chan struct{}, 1), time.Minute).Create(t.Context(), session) + require.NoError(t, err) + message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("accepted turn")) + message.ContextID = session.ContextId + task := a2a.NewSubmittedTask(message, message) + task.Status.State = a2a.TaskStateWorking + hash := sha256.Sum256([]byte(task.ID)) + _, err = store.CreateRuntimeTask(t.Context(), session.Id, hash[:], task, "") + require.NoError(t, err) + base.actors[actorKey("team-a", substrate.ActorName(session.Id))].Status.State = ateapipb.ActorState_ACTOR_STATE_CRASHED + ids = append(ids, session.Id) + } + for _, id := range ids { + actor, err := base.GetActor(t.Context(), "team-a", substrate.ActorName(id)) + require.NoError(t, err) + _, err = store.BeginCrashedSessionRecovery(t.Context(), id, revertGeneration(t, store, id), actor.GetMetadata().GetUid(), database.SessionTaskSnapshot{}) + require.NoError(t, err) + } + slices.Sort(ids) + after, err := workflow.recoverSessionPage(t.Context(), "") + require.NoError(t, err) + require.Equal(t, ids[99], after, "failed Actor reads still advance the bounded page") + actors.readErr = nil + after, err = workflow.recoverSessionPage(t.Context(), after) + require.NoError(t, err) + require.Empty(t, after, "the short second page resets the cursor") + last, err := base.GetActor(t.Context(), "team-a", substrate.ActorName(ids[100])) + require.NoError(t, err) + require.Equal(t, ateapipb.ActorState_ACTOR_STATE_SUSPENDED, last.GetStatus().GetState(), "the second page recovers the 101st candidate") + firstActor, err := base.GetActor(t.Context(), "team-a", substrate.ActorName(ids[0])) + require.NoError(t, err) + require.Equal(t, ateapipb.ActorState_ACTOR_STATE_CRASHED, firstActor.GetStatus().GetState(), "unavailable first-page Actors still need recovery") + after, err = workflow.recoverSessionPage(t.Context(), after) + require.NoError(t, err) + require.Equal(t, ids[99], after, "the wrapped scan revisits previously unavailable Actors") + after, err = workflow.recoverSessionPage(t.Context(), after) + require.NoError(t, err) + require.Empty(t, after, "an empty page resets the cursor") + for _, id := range ids { + actor, err := base.GetActor(t.Context(), "team-a", substrate.ActorName(id)) + require.NoError(t, err) + require.Equal(t, ateapipb.ActorState_ACTOR_STATE_SUSPENDED, actor.GetStatus().GetState()) + } +} diff --git a/go/core/internal/service/session/workflow.go b/go/core/internal/service/session/workflow.go index 741954f8c5..ea99fa4eb9 100644 --- a/go/core/internal/service/session/workflow.go +++ b/go/core/internal/service/session/workflow.go @@ -15,6 +15,9 @@ import ( ) type workflowStore interface { + ListSessionRecoveryCandidates(context.Context, string, int) ([]*apiv1.Session, error) + SessionLifecycleGeneration(context.Context, string) (uuid.UUID, error) + BeginCrashedSessionRecovery(context.Context, string, uuid.UUID, string, database.SessionTaskSnapshot) (*database.SessionOperation, error) ClaimSessionQuiescence(context.Context) (*database.SessionQuiescence, error) HasPendingSessionQuiescence(context.Context) (bool, error) FinishSessionQuiescence(context.Context, *database.SessionQuiescence, *database.SessionTaskSnapshot) error @@ -30,6 +33,7 @@ type workflowStore interface { type actorClient interface { substrate.LifecycleClient + RevertActor(context.Context, string, string) (*ateapipb.Actor, error) } // ActorWorkflow runs the imperative Substrate operations behind Session @@ -129,6 +133,9 @@ func (w *ActorWorkflow) execute(ctx context.Context, operation *database.Session if operation.Instance.Operation == apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE { return operationOutcome(operation) } + if operation.Instance.State == apiv1.RuntimeState_RUNTIME_STATE_FAILED && operation.Instance.Failure.GetReason() == database.SessionRuntimeCrashed { + return w.recoverCrashedSession(ctx, operation) + } kind := operation.Instance.Operation session := operation.Instance revision, err := w.store.GetRuntimeRevision(ctx, session.GetPreparedRevision()) From ccb7dc1911a4402d6148474c5a6262fc4404eb4a Mon Sep 17 00:00:00 2001 From: Ada Lundhe Date: Fri, 23 Oct 2026 12:33:56 -0700 Subject: [PATCH 03/12] fix(api): expose manual crashed-session revert Signed-off-by: Ada Lundhe --- docs/architecture/runtime-and-lifecycle.md | 34 +++ docs/lifecycle-retries.md | 19 +- go/api/gen/kagent/api/v1/sessions.pb.go | 267 ++++++++++++------ go/api/gen/kagent/api/v1/sessions_grpc.pb.go | 40 +++ go/core/internal/grpcserver/policy.go | 1 + go/core/internal/grpcserver/policy_test.go | 1 + .../internal/grpcserver/protovalidate_test.go | 1 + go/core/internal/grpcserver/session.go | 8 + .../grpcserver/session_recovery_test.go | 200 +++++++++++++ go/core/internal/mcp/tasks_test.go | 4 + go/core/internal/service/session/service.go | 19 ++ .../internal/service/session/service_test.go | 45 ++- go/core/test/e2e/session_revert_test.go | 151 ++++++++++ proto/kagent/api/v1/sessions.proto | 10 + ui/src/generated/kagent/api/v1/sessions_pb.ts | 64 ++++- 15 files changed, 753 insertions(+), 111 deletions(-) create mode 100644 go/core/internal/grpcserver/session_recovery_test.go create mode 100644 go/core/test/e2e/session_revert_test.go diff --git a/docs/architecture/runtime-and-lifecycle.md b/docs/architecture/runtime-and-lifecycle.md index 9616656216..32b7e651cb 100644 --- a/docs/architecture/runtime-and-lifecycle.md +++ b/docs/architecture/runtime-and-lifecycle.md @@ -205,6 +205,40 @@ claims still block new work, but completed results remain readable. The recorded actor UID is checked before lifecycle calls; a same-name replacement cannot be adopted implicitly. +`SessionService.RevertSession` lets an authorized caller restore a confirmed +crashed runtime. It accepts active and idle Sessions, but verifies Substrate's +CRASHED state for the recorded Actor UID and pinned template before admitting +recovery. Quiet turns, transport failures, missing Actors and same-name +replacements cannot authorize a revert. + +Admission records FAILED/RESUME with a RuntimeCrashed failure under the Session +lock. The same transaction publishes a durable terminal result saved before the +crash, or fails nonterminal work absent from the retained snapshot. Waiting tasks +captured by that snapshot remain available. Checkpoint creation excludes competing +recovery. + +The caller reads the Session's lifecycle generation before it observes Substrate. +Admission rejects the crash when another lifecycle operation began after that +read. A concurrent revert keeps the Actor UID and returns the Session to READY, +so only the generation exposes a stale observation. + +An issued idle suspension does not exclude recovery. It cannot finish against the +crashed Actor, and its worker keeps the claim when suspension fails. Admission +retires the claim in its transaction, so a late finish records nothing. + +The worker reverts the same Actor to its last durable snapshot without replaying +input. A runtime without a durable snapshot returns to its initial state. +An interrupted revert continues from REVERTING; a lost response can be reconciled +from SUSPENDED. Already admitted recoveries survive API restarts and retry on the +existing recovery interval in ID-ordered pages of at most 100. The worker does +not inspect READY Sessions or initiate a new recovery. + +A successful revert returns the Session to READY with the same Actor suspended. +The next authorized interaction resumes it through Substrate ingress. Runtime +changes after the retained snapshot are lost; the discarded turn receives no +checkpoint at that older snapshot. Another revert against the now-healthy Actor +returns `FAILED_PRECONDITION`. + ```mermaid sequenceDiagram participant Client diff --git a/docs/lifecycle-retries.md b/docs/lifecycle-retries.md index cf404a5835..78b61a129b 100644 --- a/docs/lifecycle-retries.md +++ b/docs/lifecycle-retries.md @@ -2,8 +2,8 @@ Session and Sandbox lifecycle mutations run within the API request. The server records the intended operation before calling the runtime and records completion -before returning success. Sessions expose Create, Fork, and Delete; explicit Suspend -and Resume are available only for Sandboxes. Session runtime suspension and resumption +before returning success. Sessions expose Create, Fork, Revert, and Delete; explicit +Suspend and Resume are available only for Sandboxes. Session runtime suspension and resumption are platform-managed. A failed request can have changed the runtime: a timeout does not mean that nothing happened, and canceling a request does not undo it. @@ -39,6 +39,21 @@ for subsequent invocations. Omitting it generates a new ID on each invocation, so rerunning that command is a new creation rather than a retry. MCP callers should likewise retain the `request_id` passed to `create_sandbox`. +## Recovering a crashed Session + +Call `SessionService.RevertSession` with the Session ID to restore its recorded runtime to the last durable snapshot. +The caller needs Session update access. The operation verifies Substrate's CRASHED state and the recorded Actor identity. +Recovery fails interrupted turns and waiting tasks absent from the retained snapshot. +It preserves durable terminal results and returns the Session to READY with its Actor suspended. +It discards runtime changes since that snapshot and does not replay input. +A runtime without a snapshot restarts from its initial state. + +Retry an interrupted revert against the same Session ID. Pending recovery retains its operation across API restarts. +The lifecycle worker retries already admitted recoveries; it does not initiate recovery for other crashed Sessions. +After recovery succeeds, another revert returns `FAILED_PRECONDITION` because the Actor is healthy. +Checkpoint creation returns `ABORTED` until it finishes. +A revert that observed the crash before another lifecycle operation began also returns `ABORTED`. + ## Errors | gRPC status | Client action | diff --git a/go/api/gen/kagent/api/v1/sessions.pb.go b/go/api/gen/kagent/api/v1/sessions.pb.go index 5ce6e04b22..ba4dbc92d3 100644 --- a/go/api/gen/kagent/api/v1/sessions.pb.go +++ b/go/api/gen/kagent/api/v1/sessions.pb.go @@ -617,6 +617,94 @@ func (x *UpdateSessionNameResponse) GetSession() *Session { return nil } +type RevertSessionRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + SessionId string `protobuf:"bytes,1,opt,name=session_id,json=sessionId,proto3" json:"session_id,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *RevertSessionRequest) Reset() { + *x = RevertSessionRequest{} + mi := &file_kagent_api_v1_sessions_proto_msgTypes[10] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *RevertSessionRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RevertSessionRequest) ProtoMessage() {} + +func (x *RevertSessionRequest) ProtoReflect() protoreflect.Message { + mi := &file_kagent_api_v1_sessions_proto_msgTypes[10] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RevertSessionRequest.ProtoReflect.Descriptor instead. +func (*RevertSessionRequest) Descriptor() ([]byte, []int) { + return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{10} +} + +func (x *RevertSessionRequest) GetSessionId() string { + if x != nil { + return x.SessionId + } + return "" +} + +type RevertSessionResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Session *Session `protobuf:"bytes,1,opt,name=session,proto3" json:"session,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *RevertSessionResponse) Reset() { + *x = RevertSessionResponse{} + mi := &file_kagent_api_v1_sessions_proto_msgTypes[11] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *RevertSessionResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RevertSessionResponse) ProtoMessage() {} + +func (x *RevertSessionResponse) ProtoReflect() protoreflect.Message { + mi := &file_kagent_api_v1_sessions_proto_msgTypes[11] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RevertSessionResponse.ProtoReflect.Descriptor instead. +func (*RevertSessionResponse) Descriptor() ([]byte, []int) { + return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{11} +} + +func (x *RevertSessionResponse) GetSession() *Session { + if x != nil { + return x.Session + } + return nil +} + type DeleteSessionRequest struct { state protoimpl.MessageState `protogen:"open.v1"` SessionId string `protobuf:"bytes,1,opt,name=session_id,json=sessionId,proto3" json:"session_id,omitempty"` @@ -626,7 +714,7 @@ type DeleteSessionRequest struct { func (x *DeleteSessionRequest) Reset() { *x = DeleteSessionRequest{} - mi := &file_kagent_api_v1_sessions_proto_msgTypes[10] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[12] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -638,7 +726,7 @@ func (x *DeleteSessionRequest) String() string { func (*DeleteSessionRequest) ProtoMessage() {} func (x *DeleteSessionRequest) ProtoReflect() protoreflect.Message { - mi := &file_kagent_api_v1_sessions_proto_msgTypes[10] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[12] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -651,7 +739,7 @@ func (x *DeleteSessionRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use DeleteSessionRequest.ProtoReflect.Descriptor instead. func (*DeleteSessionRequest) Descriptor() ([]byte, []int) { - return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{10} + return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{12} } func (x *DeleteSessionRequest) GetSessionId() string { @@ -670,7 +758,7 @@ type DeleteSessionResponse struct { func (x *DeleteSessionResponse) Reset() { *x = DeleteSessionResponse{} - mi := &file_kagent_api_v1_sessions_proto_msgTypes[11] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[13] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -682,7 +770,7 @@ func (x *DeleteSessionResponse) String() string { func (*DeleteSessionResponse) ProtoMessage() {} func (x *DeleteSessionResponse) ProtoReflect() protoreflect.Message { - mi := &file_kagent_api_v1_sessions_proto_msgTypes[11] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[13] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -695,7 +783,7 @@ func (x *DeleteSessionResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use DeleteSessionResponse.ProtoReflect.Descriptor instead. func (*DeleteSessionResponse) Descriptor() ([]byte, []int) { - return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{11} + return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{13} } func (x *DeleteSessionResponse) GetSession() *Session { @@ -719,7 +807,7 @@ type SessionShare struct { func (x *SessionShare) Reset() { *x = SessionShare{} - mi := &file_kagent_api_v1_sessions_proto_msgTypes[12] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[14] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -731,7 +819,7 @@ func (x *SessionShare) String() string { func (*SessionShare) ProtoMessage() {} func (x *SessionShare) ProtoReflect() protoreflect.Message { - mi := &file_kagent_api_v1_sessions_proto_msgTypes[12] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[14] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -744,7 +832,7 @@ func (x *SessionShare) ProtoReflect() protoreflect.Message { // Deprecated: Use SessionShare.ProtoReflect.Descriptor instead. func (*SessionShare) Descriptor() ([]byte, []int) { - return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{12} + return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{14} } func (x *SessionShare) GetId() string { @@ -788,7 +876,7 @@ type CreateSessionShareRequest struct { func (x *CreateSessionShareRequest) Reset() { *x = CreateSessionShareRequest{} - mi := &file_kagent_api_v1_sessions_proto_msgTypes[13] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[15] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -800,7 +888,7 @@ func (x *CreateSessionShareRequest) String() string { func (*CreateSessionShareRequest) ProtoMessage() {} func (x *CreateSessionShareRequest) ProtoReflect() protoreflect.Message { - mi := &file_kagent_api_v1_sessions_proto_msgTypes[13] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[15] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -813,7 +901,7 @@ func (x *CreateSessionShareRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use CreateSessionShareRequest.ProtoReflect.Descriptor instead. func (*CreateSessionShareRequest) Descriptor() ([]byte, []int) { - return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{13} + return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{15} } func (x *CreateSessionShareRequest) GetSessionId() string { @@ -841,7 +929,7 @@ type CreateSessionShareResponse struct { func (x *CreateSessionShareResponse) Reset() { *x = CreateSessionShareResponse{} - mi := &file_kagent_api_v1_sessions_proto_msgTypes[14] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[16] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -853,7 +941,7 @@ func (x *CreateSessionShareResponse) String() string { func (*CreateSessionShareResponse) ProtoMessage() {} func (x *CreateSessionShareResponse) ProtoReflect() protoreflect.Message { - mi := &file_kagent_api_v1_sessions_proto_msgTypes[14] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[16] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -866,7 +954,7 @@ func (x *CreateSessionShareResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use CreateSessionShareResponse.ProtoReflect.Descriptor instead. func (*CreateSessionShareResponse) Descriptor() ([]byte, []int) { - return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{14} + return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{16} } func (x *CreateSessionShareResponse) GetShare() *SessionShare { @@ -893,7 +981,7 @@ type ListSessionSharesRequest struct { func (x *ListSessionSharesRequest) Reset() { *x = ListSessionSharesRequest{} - mi := &file_kagent_api_v1_sessions_proto_msgTypes[15] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[17] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -905,7 +993,7 @@ func (x *ListSessionSharesRequest) String() string { func (*ListSessionSharesRequest) ProtoMessage() {} func (x *ListSessionSharesRequest) ProtoReflect() protoreflect.Message { - mi := &file_kagent_api_v1_sessions_proto_msgTypes[15] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[17] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -918,7 +1006,7 @@ func (x *ListSessionSharesRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use ListSessionSharesRequest.ProtoReflect.Descriptor instead. func (*ListSessionSharesRequest) Descriptor() ([]byte, []int) { - return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{15} + return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{17} } func (x *ListSessionSharesRequest) GetSessionId() string { @@ -945,7 +1033,7 @@ type ListSessionSharesResponse struct { func (x *ListSessionSharesResponse) Reset() { *x = ListSessionSharesResponse{} - mi := &file_kagent_api_v1_sessions_proto_msgTypes[16] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[18] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -957,7 +1045,7 @@ func (x *ListSessionSharesResponse) String() string { func (*ListSessionSharesResponse) ProtoMessage() {} func (x *ListSessionSharesResponse) ProtoReflect() protoreflect.Message { - mi := &file_kagent_api_v1_sessions_proto_msgTypes[16] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[18] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -970,7 +1058,7 @@ func (x *ListSessionSharesResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use ListSessionSharesResponse.ProtoReflect.Descriptor instead. func (*ListSessionSharesResponse) Descriptor() ([]byte, []int) { - return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{16} + return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{18} } func (x *ListSessionSharesResponse) GetShares() []*SessionShare { @@ -996,7 +1084,7 @@ type RevokeSessionShareRequest struct { func (x *RevokeSessionShareRequest) Reset() { *x = RevokeSessionShareRequest{} - mi := &file_kagent_api_v1_sessions_proto_msgTypes[17] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[19] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1008,7 +1096,7 @@ func (x *RevokeSessionShareRequest) String() string { func (*RevokeSessionShareRequest) ProtoMessage() {} func (x *RevokeSessionShareRequest) ProtoReflect() protoreflect.Message { - mi := &file_kagent_api_v1_sessions_proto_msgTypes[17] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[19] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1021,7 +1109,7 @@ func (x *RevokeSessionShareRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use RevokeSessionShareRequest.ProtoReflect.Descriptor instead. func (*RevokeSessionShareRequest) Descriptor() ([]byte, []int) { - return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{17} + return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{19} } func (x *RevokeSessionShareRequest) GetShareId() string { @@ -1039,7 +1127,7 @@ type RevokeSessionShareResponse struct { func (x *RevokeSessionShareResponse) Reset() { *x = RevokeSessionShareResponse{} - mi := &file_kagent_api_v1_sessions_proto_msgTypes[18] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[20] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1051,7 +1139,7 @@ func (x *RevokeSessionShareResponse) String() string { func (*RevokeSessionShareResponse) ProtoMessage() {} func (x *RevokeSessionShareResponse) ProtoReflect() protoreflect.Message { - mi := &file_kagent_api_v1_sessions_proto_msgTypes[18] + mi := &file_kagent_api_v1_sessions_proto_msgTypes[20] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1064,7 +1152,7 @@ func (x *RevokeSessionShareResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use RevokeSessionShareResponse.ProtoReflect.Descriptor instead. func (*RevokeSessionShareResponse) Descriptor() ([]byte, []int) { - return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{18} + return file_kagent_api_v1_sessions_proto_rawDescGZIP(), []int{20} } var File_kagent_api_v1_sessions_proto protoreflect.FileDescriptor @@ -1118,6 +1206,11 @@ const file_kagent_api_v1_sessions_proto_rawDesc = "" + "\x04name\x18\x02 \x01(\tB=\xbaH:r8\x18\xc8\x0123^(?:$|[^\\p{Z}\\p{Cc}](?:[^\\p{Cc}]*[^\\p{Z}\\p{Cc}])?)$R\x04name\"M\n" + "\x19UpdateSessionNameResponse\x120\n" + "\asession\x18\x01 \x01(\v2\x16.kagent.api.v1.SessionR\asession\"?\n" + + "\x14RevertSessionRequest\x12'\n" + + "\n" + + "session_id\x18\x01 \x01(\tB\b\xbaH\x05r\x03\xb0\x01\x01R\tsessionId\"I\n" + + "\x15RevertSessionResponse\x120\n" + + "\asession\x18\x01 \x01(\v2\x16.kagent.api.v1.SessionR\asession\"?\n" + "\x14DeleteSessionRequest\x12'\n" + "\n" + "session_id\x18\x01 \x01(\tB\b\xbaH\x05r\x03\xb0\x01\x01R\tsessionId\"I\n" + @@ -1147,13 +1240,14 @@ const file_kagent_api_v1_sessions_proto_rawDesc = "" + "\x04page\x18\x02 \x01(\v2\x1b.kagent.api.v1.PageResponseR\x04page\"@\n" + "\x19RevokeSessionShareRequest\x12#\n" + "\bshare_id\x18\x01 \x01(\tB\b\xbaH\x05r\x03\xb0\x01\x01R\ashareId\"\x1c\n" + - "\x1aRevokeSessionShareResponse2\x9a\x06\n" + + "\x1aRevokeSessionShareResponse2\xf6\x06\n" + "\x0eSessionService\x12Z\n" + "\rCreateSession\x12#.kagent.api.v1.CreateSessionRequest\x1a$.kagent.api.v1.CreateSessionResponse\x12Q\n" + "\n" + "GetSession\x12 .kagent.api.v1.GetSessionRequest\x1a!.kagent.api.v1.GetSessionResponse\x12W\n" + "\fListSessions\x12\".kagent.api.v1.ListSessionsRequest\x1a#.kagent.api.v1.ListSessionsResponse\x12f\n" + "\x11UpdateSessionName\x12'.kagent.api.v1.UpdateSessionNameRequest\x1a(.kagent.api.v1.UpdateSessionNameResponse\x12Z\n" + + "\rRevertSession\x12#.kagent.api.v1.RevertSessionRequest\x1a$.kagent.api.v1.RevertSessionResponse\x12Z\n" + "\rDeleteSession\x12#.kagent.api.v1.DeleteSessionRequest\x1a$.kagent.api.v1.DeleteSessionResponse\x12i\n" + "\x12CreateSessionShare\x12(.kagent.api.v1.CreateSessionShareRequest\x1a).kagent.api.v1.CreateSessionShareResponse\x12f\n" + "\x11ListSessionShares\x12'.kagent.api.v1.ListSessionSharesRequest\x1a(.kagent.api.v1.ListSessionSharesResponse\x12i\n" + @@ -1171,7 +1265,7 @@ func file_kagent_api_v1_sessions_proto_rawDescGZIP() []byte { return file_kagent_api_v1_sessions_proto_rawDescData } -var file_kagent_api_v1_sessions_proto_msgTypes = make([]protoimpl.MessageInfo, 19) +var file_kagent_api_v1_sessions_proto_msgTypes = make([]protoimpl.MessageInfo, 21) var file_kagent_api_v1_sessions_proto_goTypes = []any{ (*Failure)(nil), // 0: kagent.api.v1.Failure (*Session)(nil), // 1: kagent.api.v1.Session @@ -1183,67 +1277,72 @@ var file_kagent_api_v1_sessions_proto_goTypes = []any{ (*ListSessionsResponse)(nil), // 7: kagent.api.v1.ListSessionsResponse (*UpdateSessionNameRequest)(nil), // 8: kagent.api.v1.UpdateSessionNameRequest (*UpdateSessionNameResponse)(nil), // 9: kagent.api.v1.UpdateSessionNameResponse - (*DeleteSessionRequest)(nil), // 10: kagent.api.v1.DeleteSessionRequest - (*DeleteSessionResponse)(nil), // 11: kagent.api.v1.DeleteSessionResponse - (*SessionShare)(nil), // 12: kagent.api.v1.SessionShare - (*CreateSessionShareRequest)(nil), // 13: kagent.api.v1.CreateSessionShareRequest - (*CreateSessionShareResponse)(nil), // 14: kagent.api.v1.CreateSessionShareResponse - (*ListSessionSharesRequest)(nil), // 15: kagent.api.v1.ListSessionSharesRequest - (*ListSessionSharesResponse)(nil), // 16: kagent.api.v1.ListSessionSharesResponse - (*RevokeSessionShareRequest)(nil), // 17: kagent.api.v1.RevokeSessionShareRequest - (*RevokeSessionShareResponse)(nil), // 18: kagent.api.v1.RevokeSessionShareResponse - (*ResourceReference)(nil), // 19: kagent.api.v1.ResourceReference - (RuntimeState)(0), // 20: kagent.api.v1.RuntimeState - (RuntimeOperation)(0), // 21: kagent.api.v1.RuntimeOperation - (*timestamppb.Timestamp)(nil), // 22: google.protobuf.Timestamp - (*PageRequest)(nil), // 23: kagent.api.v1.PageRequest - (*PageResponse)(nil), // 24: kagent.api.v1.PageResponse - (*durationpb.Duration)(nil), // 25: google.protobuf.Duration + (*RevertSessionRequest)(nil), // 10: kagent.api.v1.RevertSessionRequest + (*RevertSessionResponse)(nil), // 11: kagent.api.v1.RevertSessionResponse + (*DeleteSessionRequest)(nil), // 12: kagent.api.v1.DeleteSessionRequest + (*DeleteSessionResponse)(nil), // 13: kagent.api.v1.DeleteSessionResponse + (*SessionShare)(nil), // 14: kagent.api.v1.SessionShare + (*CreateSessionShareRequest)(nil), // 15: kagent.api.v1.CreateSessionShareRequest + (*CreateSessionShareResponse)(nil), // 16: kagent.api.v1.CreateSessionShareResponse + (*ListSessionSharesRequest)(nil), // 17: kagent.api.v1.ListSessionSharesRequest + (*ListSessionSharesResponse)(nil), // 18: kagent.api.v1.ListSessionSharesResponse + (*RevokeSessionShareRequest)(nil), // 19: kagent.api.v1.RevokeSessionShareRequest + (*RevokeSessionShareResponse)(nil), // 20: kagent.api.v1.RevokeSessionShareResponse + (*ResourceReference)(nil), // 21: kagent.api.v1.ResourceReference + (RuntimeState)(0), // 22: kagent.api.v1.RuntimeState + (RuntimeOperation)(0), // 23: kagent.api.v1.RuntimeOperation + (*timestamppb.Timestamp)(nil), // 24: google.protobuf.Timestamp + (*PageRequest)(nil), // 25: kagent.api.v1.PageRequest + (*PageResponse)(nil), // 26: kagent.api.v1.PageResponse + (*durationpb.Duration)(nil), // 27: google.protobuf.Duration } var file_kagent_api_v1_sessions_proto_depIdxs = []int32{ - 19, // 0: kagent.api.v1.Session.agent:type_name -> kagent.api.v1.ResourceReference - 20, // 1: kagent.api.v1.Session.state:type_name -> kagent.api.v1.RuntimeState - 21, // 2: kagent.api.v1.Session.operation:type_name -> kagent.api.v1.RuntimeOperation + 21, // 0: kagent.api.v1.Session.agent:type_name -> kagent.api.v1.ResourceReference + 22, // 1: kagent.api.v1.Session.state:type_name -> kagent.api.v1.RuntimeState + 23, // 2: kagent.api.v1.Session.operation:type_name -> kagent.api.v1.RuntimeOperation 0, // 3: kagent.api.v1.Session.failure:type_name -> kagent.api.v1.Failure - 22, // 4: kagent.api.v1.Session.created_at:type_name -> google.protobuf.Timestamp - 22, // 5: kagent.api.v1.Session.updated_at:type_name -> google.protobuf.Timestamp - 19, // 6: kagent.api.v1.CreateSessionRequest.agent:type_name -> kagent.api.v1.ResourceReference + 24, // 4: kagent.api.v1.Session.created_at:type_name -> google.protobuf.Timestamp + 24, // 5: kagent.api.v1.Session.updated_at:type_name -> google.protobuf.Timestamp + 21, // 6: kagent.api.v1.CreateSessionRequest.agent:type_name -> kagent.api.v1.ResourceReference 1, // 7: kagent.api.v1.CreateSessionResponse.session:type_name -> kagent.api.v1.Session 1, // 8: kagent.api.v1.GetSessionResponse.session:type_name -> kagent.api.v1.Session - 23, // 9: kagent.api.v1.ListSessionsRequest.page:type_name -> kagent.api.v1.PageRequest - 19, // 10: kagent.api.v1.ListSessionsRequest.agent:type_name -> kagent.api.v1.ResourceReference + 25, // 9: kagent.api.v1.ListSessionsRequest.page:type_name -> kagent.api.v1.PageRequest + 21, // 10: kagent.api.v1.ListSessionsRequest.agent:type_name -> kagent.api.v1.ResourceReference 1, // 11: kagent.api.v1.ListSessionsResponse.sessions:type_name -> kagent.api.v1.Session - 24, // 12: kagent.api.v1.ListSessionsResponse.page:type_name -> kagent.api.v1.PageResponse + 26, // 12: kagent.api.v1.ListSessionsResponse.page:type_name -> kagent.api.v1.PageResponse 1, // 13: kagent.api.v1.UpdateSessionNameResponse.session:type_name -> kagent.api.v1.Session - 1, // 14: kagent.api.v1.DeleteSessionResponse.session:type_name -> kagent.api.v1.Session - 22, // 15: kagent.api.v1.SessionShare.created_at:type_name -> google.protobuf.Timestamp - 22, // 16: kagent.api.v1.SessionShare.expires_at:type_name -> google.protobuf.Timestamp - 25, // 17: kagent.api.v1.CreateSessionShareRequest.ttl:type_name -> google.protobuf.Duration - 12, // 18: kagent.api.v1.CreateSessionShareResponse.share:type_name -> kagent.api.v1.SessionShare - 23, // 19: kagent.api.v1.ListSessionSharesRequest.page:type_name -> kagent.api.v1.PageRequest - 12, // 20: kagent.api.v1.ListSessionSharesResponse.shares:type_name -> kagent.api.v1.SessionShare - 24, // 21: kagent.api.v1.ListSessionSharesResponse.page:type_name -> kagent.api.v1.PageResponse - 2, // 22: kagent.api.v1.SessionService.CreateSession:input_type -> kagent.api.v1.CreateSessionRequest - 4, // 23: kagent.api.v1.SessionService.GetSession:input_type -> kagent.api.v1.GetSessionRequest - 6, // 24: kagent.api.v1.SessionService.ListSessions:input_type -> kagent.api.v1.ListSessionsRequest - 8, // 25: kagent.api.v1.SessionService.UpdateSessionName:input_type -> kagent.api.v1.UpdateSessionNameRequest - 10, // 26: kagent.api.v1.SessionService.DeleteSession:input_type -> kagent.api.v1.DeleteSessionRequest - 13, // 27: kagent.api.v1.SessionService.CreateSessionShare:input_type -> kagent.api.v1.CreateSessionShareRequest - 15, // 28: kagent.api.v1.SessionService.ListSessionShares:input_type -> kagent.api.v1.ListSessionSharesRequest - 17, // 29: kagent.api.v1.SessionService.RevokeSessionShare:input_type -> kagent.api.v1.RevokeSessionShareRequest - 3, // 30: kagent.api.v1.SessionService.CreateSession:output_type -> kagent.api.v1.CreateSessionResponse - 5, // 31: kagent.api.v1.SessionService.GetSession:output_type -> kagent.api.v1.GetSessionResponse - 7, // 32: kagent.api.v1.SessionService.ListSessions:output_type -> kagent.api.v1.ListSessionsResponse - 9, // 33: kagent.api.v1.SessionService.UpdateSessionName:output_type -> kagent.api.v1.UpdateSessionNameResponse - 11, // 34: kagent.api.v1.SessionService.DeleteSession:output_type -> kagent.api.v1.DeleteSessionResponse - 14, // 35: kagent.api.v1.SessionService.CreateSessionShare:output_type -> kagent.api.v1.CreateSessionShareResponse - 16, // 36: kagent.api.v1.SessionService.ListSessionShares:output_type -> kagent.api.v1.ListSessionSharesResponse - 18, // 37: kagent.api.v1.SessionService.RevokeSessionShare:output_type -> kagent.api.v1.RevokeSessionShareResponse - 30, // [30:38] is the sub-list for method output_type - 22, // [22:30] is the sub-list for method input_type - 22, // [22:22] is the sub-list for extension type_name - 22, // [22:22] is the sub-list for extension extendee - 0, // [0:22] is the sub-list for field type_name + 1, // 14: kagent.api.v1.RevertSessionResponse.session:type_name -> kagent.api.v1.Session + 1, // 15: kagent.api.v1.DeleteSessionResponse.session:type_name -> kagent.api.v1.Session + 24, // 16: kagent.api.v1.SessionShare.created_at:type_name -> google.protobuf.Timestamp + 24, // 17: kagent.api.v1.SessionShare.expires_at:type_name -> google.protobuf.Timestamp + 27, // 18: kagent.api.v1.CreateSessionShareRequest.ttl:type_name -> google.protobuf.Duration + 14, // 19: kagent.api.v1.CreateSessionShareResponse.share:type_name -> kagent.api.v1.SessionShare + 25, // 20: kagent.api.v1.ListSessionSharesRequest.page:type_name -> kagent.api.v1.PageRequest + 14, // 21: kagent.api.v1.ListSessionSharesResponse.shares:type_name -> kagent.api.v1.SessionShare + 26, // 22: kagent.api.v1.ListSessionSharesResponse.page:type_name -> kagent.api.v1.PageResponse + 2, // 23: kagent.api.v1.SessionService.CreateSession:input_type -> kagent.api.v1.CreateSessionRequest + 4, // 24: kagent.api.v1.SessionService.GetSession:input_type -> kagent.api.v1.GetSessionRequest + 6, // 25: kagent.api.v1.SessionService.ListSessions:input_type -> kagent.api.v1.ListSessionsRequest + 8, // 26: kagent.api.v1.SessionService.UpdateSessionName:input_type -> kagent.api.v1.UpdateSessionNameRequest + 10, // 27: kagent.api.v1.SessionService.RevertSession:input_type -> kagent.api.v1.RevertSessionRequest + 12, // 28: kagent.api.v1.SessionService.DeleteSession:input_type -> kagent.api.v1.DeleteSessionRequest + 15, // 29: kagent.api.v1.SessionService.CreateSessionShare:input_type -> kagent.api.v1.CreateSessionShareRequest + 17, // 30: kagent.api.v1.SessionService.ListSessionShares:input_type -> kagent.api.v1.ListSessionSharesRequest + 19, // 31: kagent.api.v1.SessionService.RevokeSessionShare:input_type -> kagent.api.v1.RevokeSessionShareRequest + 3, // 32: kagent.api.v1.SessionService.CreateSession:output_type -> kagent.api.v1.CreateSessionResponse + 5, // 33: kagent.api.v1.SessionService.GetSession:output_type -> kagent.api.v1.GetSessionResponse + 7, // 34: kagent.api.v1.SessionService.ListSessions:output_type -> kagent.api.v1.ListSessionsResponse + 9, // 35: kagent.api.v1.SessionService.UpdateSessionName:output_type -> kagent.api.v1.UpdateSessionNameResponse + 11, // 36: kagent.api.v1.SessionService.RevertSession:output_type -> kagent.api.v1.RevertSessionResponse + 13, // 37: kagent.api.v1.SessionService.DeleteSession:output_type -> kagent.api.v1.DeleteSessionResponse + 16, // 38: kagent.api.v1.SessionService.CreateSessionShare:output_type -> kagent.api.v1.CreateSessionShareResponse + 18, // 39: kagent.api.v1.SessionService.ListSessionShares:output_type -> kagent.api.v1.ListSessionSharesResponse + 20, // 40: kagent.api.v1.SessionService.RevokeSessionShare:output_type -> kagent.api.v1.RevokeSessionShareResponse + 32, // [32:41] is the sub-list for method output_type + 23, // [23:32] is the sub-list for method input_type + 23, // [23:23] is the sub-list for extension type_name + 23, // [23:23] is the sub-list for extension extendee + 0, // [0:23] is the sub-list for field type_name } func init() { file_kagent_api_v1_sessions_proto_init() } @@ -1259,7 +1358,7 @@ func file_kagent_api_v1_sessions_proto_init() { GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_kagent_api_v1_sessions_proto_rawDesc), len(file_kagent_api_v1_sessions_proto_rawDesc)), NumEnums: 0, - NumMessages: 19, + NumMessages: 21, NumExtensions: 0, NumServices: 1, }, diff --git a/go/api/gen/kagent/api/v1/sessions_grpc.pb.go b/go/api/gen/kagent/api/v1/sessions_grpc.pb.go index f5bfe3112d..4c262b9426 100644 --- a/go/api/gen/kagent/api/v1/sessions_grpc.pb.go +++ b/go/api/gen/kagent/api/v1/sessions_grpc.pb.go @@ -23,6 +23,7 @@ const ( SessionService_GetSession_FullMethodName = "/kagent.api.v1.SessionService/GetSession" SessionService_ListSessions_FullMethodName = "/kagent.api.v1.SessionService/ListSessions" SessionService_UpdateSessionName_FullMethodName = "/kagent.api.v1.SessionService/UpdateSessionName" + SessionService_RevertSession_FullMethodName = "/kagent.api.v1.SessionService/RevertSession" SessionService_DeleteSession_FullMethodName = "/kagent.api.v1.SessionService/DeleteSession" SessionService_CreateSessionShare_FullMethodName = "/kagent.api.v1.SessionService/CreateSessionShare" SessionService_ListSessionShares_FullMethodName = "/kagent.api.v1.SessionService/ListSessionShares" @@ -43,6 +44,8 @@ type SessionServiceClient interface { GetSession(ctx context.Context, in *GetSessionRequest, opts ...grpc.CallOption) (*GetSessionResponse, error) ListSessions(ctx context.Context, in *ListSessionsRequest, opts ...grpc.CallOption) (*ListSessionsResponse, error) UpdateSessionName(ctx context.Context, in *UpdateSessionNameRequest, opts ...grpc.CallOption) (*UpdateSessionNameResponse, error) + // Restores a crashed runtime to its last snapshot without replaying its turn. + RevertSession(ctx context.Context, in *RevertSessionRequest, opts ...grpc.CallOption) (*RevertSessionResponse, error) DeleteSession(ctx context.Context, in *DeleteSessionRequest, opts ...grpc.CallOption) (*DeleteSessionResponse, error) CreateSessionShare(ctx context.Context, in *CreateSessionShareRequest, opts ...grpc.CallOption) (*CreateSessionShareResponse, error) ListSessionShares(ctx context.Context, in *ListSessionSharesRequest, opts ...grpc.CallOption) (*ListSessionSharesResponse, error) @@ -97,6 +100,16 @@ func (c *sessionServiceClient) UpdateSessionName(ctx context.Context, in *Update return out, nil } +func (c *sessionServiceClient) RevertSession(ctx context.Context, in *RevertSessionRequest, opts ...grpc.CallOption) (*RevertSessionResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(RevertSessionResponse) + err := c.cc.Invoke(ctx, SessionService_RevertSession_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + func (c *sessionServiceClient) DeleteSession(ctx context.Context, in *DeleteSessionRequest, opts ...grpc.CallOption) (*DeleteSessionResponse, error) { cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) out := new(DeleteSessionResponse) @@ -151,6 +164,8 @@ type SessionServiceServer interface { GetSession(context.Context, *GetSessionRequest) (*GetSessionResponse, error) ListSessions(context.Context, *ListSessionsRequest) (*ListSessionsResponse, error) UpdateSessionName(context.Context, *UpdateSessionNameRequest) (*UpdateSessionNameResponse, error) + // Restores a crashed runtime to its last snapshot without replaying its turn. + RevertSession(context.Context, *RevertSessionRequest) (*RevertSessionResponse, error) DeleteSession(context.Context, *DeleteSessionRequest) (*DeleteSessionResponse, error) CreateSessionShare(context.Context, *CreateSessionShareRequest) (*CreateSessionShareResponse, error) ListSessionShares(context.Context, *ListSessionSharesRequest) (*ListSessionSharesResponse, error) @@ -177,6 +192,9 @@ func (UnimplementedSessionServiceServer) ListSessions(context.Context, *ListSess func (UnimplementedSessionServiceServer) UpdateSessionName(context.Context, *UpdateSessionNameRequest) (*UpdateSessionNameResponse, error) { return nil, status.Error(codes.Unimplemented, "method UpdateSessionName not implemented") } +func (UnimplementedSessionServiceServer) RevertSession(context.Context, *RevertSessionRequest) (*RevertSessionResponse, error) { + return nil, status.Error(codes.Unimplemented, "method RevertSession not implemented") +} func (UnimplementedSessionServiceServer) DeleteSession(context.Context, *DeleteSessionRequest) (*DeleteSessionResponse, error) { return nil, status.Error(codes.Unimplemented, "method DeleteSession not implemented") } @@ -282,6 +300,24 @@ func _SessionService_UpdateSessionName_Handler(srv interface{}, ctx context.Cont return interceptor(ctx, in, info, handler) } +func _SessionService_RevertSession_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(RevertSessionRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(SessionServiceServer).RevertSession(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: SessionService_RevertSession_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(SessionServiceServer).RevertSession(ctx, req.(*RevertSessionRequest)) + } + return interceptor(ctx, in, info, handler) +} + func _SessionService_DeleteSession_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { in := new(DeleteSessionRequest) if err := dec(in); err != nil { @@ -377,6 +413,10 @@ var SessionService_ServiceDesc = grpc.ServiceDesc{ MethodName: "UpdateSessionName", Handler: _SessionService_UpdateSessionName_Handler, }, + { + MethodName: "RevertSession", + Handler: _SessionService_RevertSession_Handler, + }, { MethodName: "DeleteSession", Handler: _SessionService_DeleteSession_Handler, diff --git a/go/core/internal/grpcserver/policy.go b/go/core/internal/grpcserver/policy.go index 4b734fe890..3bcc71b972 100644 --- a/go/core/internal/grpcserver/policy.go +++ b/go/core/internal/grpcserver/policy.go @@ -73,6 +73,7 @@ func DefaultMethodPolicies() MethodPolicies { // A rename is the only write on this service that is not a lifecycle // operation, and it must not inherit the read mode its neighbours carry. policies[apiv1.SessionService_UpdateSessionName_FullMethodName] = auth.AccessUpdate + policies[apiv1.SessionService_RevertSession_FullMethodName] = auth.AccessUpdate policies[apiv1.SessionService_DeleteSession_FullMethodName] = auth.AccessDelete policies[apiv1.SessionService_CreateSessionShare_FullMethodName] = auth.AccessCreate policies[apiv1.SessionService_ListSessionShares_FullMethodName] = auth.AccessRead diff --git a/go/core/internal/grpcserver/policy_test.go b/go/core/internal/grpcserver/policy_test.go index ca936c7340..d4fbd3f335 100644 --- a/go/core/internal/grpcserver/policy_test.go +++ b/go/core/internal/grpcserver/policy_test.go @@ -38,6 +38,7 @@ func TestSessionServicePoliciesMatchTheirEffect(t *testing.T) { {name: "get", method: apiv1.SessionService_GetSession_FullMethodName, want: pkgauth.AccessRead}, {name: "list", method: apiv1.SessionService_ListSessions_FullMethodName, want: pkgauth.AccessRead}, {name: "rename is a write", method: apiv1.SessionService_UpdateSessionName_FullMethodName, want: pkgauth.AccessUpdate}, + {name: "revert", method: apiv1.SessionService_RevertSession_FullMethodName, want: pkgauth.AccessUpdate}, {name: "delete", method: apiv1.SessionService_DeleteSession_FullMethodName, want: pkgauth.AccessDelete}, {name: "create share", method: apiv1.SessionService_CreateSessionShare_FullMethodName, want: pkgauth.AccessCreate}, {name: "list shares", method: apiv1.SessionService_ListSessionShares_FullMethodName, want: pkgauth.AccessRead}, diff --git a/go/core/internal/grpcserver/protovalidate_test.go b/go/core/internal/grpcserver/protovalidate_test.go index 740ce014fe..64e1877b20 100644 --- a/go/core/internal/grpcserver/protovalidate_test.go +++ b/go/core/internal/grpcserver/protovalidate_test.go @@ -88,6 +88,7 @@ func TestInvalidSessionAndCheckpointIDsNeverReachHandlers(t *testing.T) { requests := []proto.Message{ &apiv1.GetSessionRequest{SessionId: "invalid"}, &apiv1.UpdateSessionNameRequest{SessionId: "invalid"}, + &apiv1.RevertSessionRequest{SessionId: "invalid"}, &apiv1.DeleteSessionRequest{SessionId: "invalid"}, &apiv1.CreateSessionShareRequest{SessionId: "invalid"}, &apiv1.ListSessionSharesRequest{SessionId: "invalid"}, diff --git a/go/core/internal/grpcserver/session.go b/go/core/internal/grpcserver/session.go index bcb55b8905..ddd925a139 100644 --- a/go/core/internal/grpcserver/session.go +++ b/go/core/internal/grpcserver/session.go @@ -51,6 +51,14 @@ func (s *sessionServer) UpdateSessionName(ctx context.Context, request *apiv1.Up return &apiv1.UpdateSessionNameResponse{Session: session}, nil } +func (s *sessionServer) RevertSession(ctx context.Context, request *apiv1.RevertSessionRequest) (*apiv1.RevertSessionResponse, error) { + session, err := s.service.Revert(ctx, request.GetSessionId()) + if err != nil { + return nil, err + } + return &apiv1.RevertSessionResponse{Session: session}, nil +} + func (s *sessionServer) DeleteSession(ctx context.Context, request *apiv1.DeleteSessionRequest) (*apiv1.DeleteSessionResponse, error) { session, err := s.service.Delete(ctx, request.GetSessionId()) if err != nil { diff --git a/go/core/internal/grpcserver/session_recovery_test.go b/go/core/internal/grpcserver/session_recovery_test.go new file mode 100644 index 0000000000..7515035aab --- /dev/null +++ b/go/core/internal/grpcserver/session_recovery_test.go @@ -0,0 +1,200 @@ +package grpcserver + +import ( + "context" + "crypto/sha256" + "net" + "sync" + "testing" + "time" + + "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/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/grpc/test/bufconn" + "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" + authimpl "github.com/kagent-dev/kagent/go/core/internal/httpserver/auth" + sessionsvc "github.com/kagent-dev/kagent/go/core/internal/service/session" + "github.com/kagent-dev/kagent/go/core/internal/substrate" + "github.com/kagent-dev/kagent/go/core/pkg/auth" +) + +type revertTestControl struct { + ateapipb.ControlClient + mu sync.Mutex + actors map[string]*ateapipb.Actor + revert func(*ateapipb.Actor) error +} + +func (c *revertTestControl) GetActor(_ context.Context, request *ateapipb.GetActorRequest, _ ...grpc.CallOption) (*ateapipb.Actor, error) { + c.mu.Lock() + defer c.mu.Unlock() + actor := c.actors[request.GetActor().GetName()] + if actor == nil { + return nil, status.Error(codes.NotFound, "actor not found") + } + return proto.CloneOf(actor), nil +} + +func (c *revertTestControl) RevertActor(_ context.Context, request *ateapipb.RevertActorRequest, _ ...grpc.CallOption) (*ateapipb.RevertActorResponse, error) { + c.mu.Lock() + defer c.mu.Unlock() + actor := c.actors[request.GetActor().GetName()] + if c.revert != nil { + if err := c.revert(actor); err != nil { + return nil, err + } + } + actor.Status.State = ateapipb.ActorState_ACTOR_STATE_SUSPENDED + return &ateapipb.RevertActorResponse{Actor: proto.CloneOf(actor)}, nil +} + +func TestRevertSessionThroughGRPC(t *testing.T) { + if testing.Short() { + t.Skip("requires PostgreSQL") + } + dsn := dbtest.StartT(context.WithoutCancel(t.Context()), t) + dbtest.MigrateT(t, dsn, false) + db, err := database.Connect(t.Context(), &database.PostgresConfig{Pool: dbtest.PoolConfig(dsn)}) + require.NoError(t, err) + t.Cleanup(db.Close) + store := database.NewClient(db, "public") + control := &revertTestControl{actors: map[string]*ateapipb.Actor{}} + actors := &substrate.Client{ControlClient: control} + workflow := sessionsvc.NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute) + listener := bufconn.Listen(DefaultMaxMessageSize) + server, err := New(Config{ + Listener: listener, Authenticator: &authimpl.InsecureAuthenticator{}, + SystemService: testSystemService(), ShareStore: store, + SessionService: sessionsvc.NewService(store, &auth.NoopAuthorizer{}, workflow), + }) + require.NoError(t, err) + ctx, stop := context.WithCancel(t.Context()) + done := make(chan error, 1) + go func() { done <- server.Start(ctx) }() + t.Cleanup(func() { stop(); require.NoError(t, <-done) }) + connection, err := grpc.NewClient("passthrough:///bufnet", grpc.WithTransportCredentials(insecure.NewCredentials()), + grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) { return listener.Dial() })) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, connection.Close()) }) + client := apiv1.NewSessionServiceClient(connection) + owner := metadata.NewOutgoingContext(t.Context(), metadata.Pairs("x-user-id", "alice")) + visitor := metadata.NewOutgoingContext(t.Context(), metadata.Pairs("x-user-id", "bob")) + _, err = client.RevertSession(owner, &apiv1.RevertSessionRequest{SessionId: "invalid"}) + require.Equal(t, codes.InvalidArgument, status.Code(err)) + + for _, test := range []struct { + name string + task a2a.TaskState + noSnapshot bool + lostResponse bool + }{ + {name: "interrupted turn", task: a2a.TaskStateWorking}, + {name: "completed turn", task: a2a.TaskStateCompleted}, + {name: "idle session"}, + {name: "published input wait", task: a2a.TaskStateInputRequired}, + {name: "published auth wait", task: a2a.TaskStateAuthRequired}, + {name: "first-turn crash", task: a2a.TaskStateWorking, noSnapshot: true}, + {name: "revert response lost", task: a2a.TaskStateWorking, lostResponse: true}, + } { + t.Run(test.name, func(t *testing.T) { + session := createTaskStoreSession(t, store) + name := substrate.ActorName(session.Id) + control.actors[name] = &ateapipb.Actor{ + Metadata: &ateapipb.ResourceMetadata{Atespace: "team-a", Name: name, Uid: "actor-uid"}, + ActorTemplate: &ateapipb.ObjectRef{Atespace: "team-a", Name: "assistant-kagent-revision"}, + Status: &ateapipb.ActorStatus{State: ateapipb.ActorState_ACTOR_STATE_RUNNING, + Snapshots: []*ateapipb.Snapshot{{Generation: 1, Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES, DurableSnapshot: &ateapipb.SnapshotStorage{Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_COMPLETED, Object: &ateapipb.ObjectSnapshot{SnapshotUri: "s3://snapshots/last-good"}}}}}, + } + var task *a2a.Task + if test.task != "" { + message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("original input")) + message.ContextID = session.ContextId + task = a2a.NewSubmittedTask(message, message) + task.Status.State = test.task + hash := sha256.Sum256([]byte(task.ID)) + version, err := store.CreateRuntimeTask(t.Context(), session.Id, hash[:], task, "") + require.NoError(t, err) + if test.task == a2a.TaskStateInputRequired || test.task == a2a.TaskStateAuthRequired { + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + } + } + request := &apiv1.RevertSessionRequest{SessionId: session.Id} + if test.noSnapshot { + control.actors[name].Status.Snapshots = nil + } + share, err := client.CreateSessionShare(owner, &apiv1.CreateSessionShareRequest{ + SessionId: session.Id, + }) + require.NoError(t, err) + readOnly := metadata.AppendToOutgoingContext(visitor, "x-share-token", share.Token) + _, err = client.RevertSession(readOnly, request) + require.Equal(t, codes.PermissionDenied, status.Code(err)) + _, err = client.RevertSession(visitor, request) + require.Equal(t, codes.NotFound, status.Code(err)) + _, err = client.RevertSession(owner, request) + require.Equal(t, codes.FailedPrecondition, status.Code(err), "healthy runtimes cannot revert") + control.actors[name].Status.State = ateapipb.ActorState_ACTOR_STATE_CRASHED + before := proto.CloneOf(control.actors[name]) + control.actors[name].Metadata.Uid = "replacement-uid" + _, err = client.RevertSession(owner, request) + require.Equal(t, codes.Unavailable, status.Code(err), "a same-name replacement is not this runtime") + control.actors[name].Metadata.Uid = before.Metadata.Uid + delete(control.actors, name) + _, err = client.RevertSession(owner, request) + require.Equal(t, codes.Unavailable, status.Code(err), "a missing runtime cannot authorize recovery") + control.actors[name] = proto.CloneOf(before) + calls := 0 + failOnce := true + control.revert = func(actor *ateapipb.Actor) error { + calls++ + if failOnce { + failOnce = false + if test.lostResponse { + actor.Status.State = ateapipb.ActorState_ACTOR_STATE_SUSPENDED + } + return status.Error(codes.Unavailable, "temporary revert outage") + } + return nil + } + _, err = client.RevertSession(owner, request) + require.Equal(t, codes.Unavailable, status.Code(err)) + failed, err := client.GetSession(owner, &apiv1.GetSessionRequest{SessionId: session.Id}) + require.NoError(t, err) + require.Equal(t, database.SessionRuntimeCrashed, failed.Session.Failure.Reason) + ready, err := client.RevertSession(owner, request) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, ready.Session.State) + require.Nil(t, ready.Session.Failure) + if test.lostResponse { + require.Equal(t, 1, calls, "a lost response cannot repeat the revert") + } else { + require.Equal(t, 2, calls) + } + after := control.actors[name] + require.Equal(t, before.Metadata.Uid, after.Metadata.Uid) + require.True(t, proto.Equal(substrate.LatestDurableActorSnapshot(before), substrate.LatestDurableActorSnapshot(after))) + require.Equal(t, ateapipb.ActorState_ACTOR_STATE_SUSPENDED, after.Status.State) + if task != nil { + stored, err := store.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + want := test.task + if want == a2a.TaskStateWorking || want == a2a.TaskStateInputRequired || want == a2a.TaskStateAuthRequired { + want = a2a.TaskStateFailed + } + require.Equal(t, want, stored.Status.State) + } + require.NoError(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), uuid.NewString())) + }) + } +} diff --git a/go/core/internal/mcp/tasks_test.go b/go/core/internal/mcp/tasks_test.go index 4d1774cc8c..ff87ed8623 100644 --- a/go/core/internal/mcp/tasks_test.go +++ b/go/core/internal/mcp/tasks_test.go @@ -566,6 +566,10 @@ func (*fakeSessionWorkflow) Create(_ context.Context, session *apiv1.Session) (* return session, nil } +func (*fakeSessionWorkflow) Revert(_ context.Context, session *apiv1.Session) (*apiv1.Session, error) { + return session, nil +} + func (*fakeSessionWorkflow) Delete(_ context.Context, session *apiv1.Session) (*apiv1.Session, error) { return session, nil } diff --git a/go/core/internal/service/session/service.go b/go/core/internal/service/session/service.go index aba2b28cac..bc2d22d32f 100644 --- a/go/core/internal/service/session/service.go +++ b/go/core/internal/service/session/service.go @@ -38,6 +38,7 @@ type store interface { type sessionWorkflow interface { Create(context.Context, *apiv1.Session) (*apiv1.Session, error) + Revert(context.Context, *apiv1.Session) (*apiv1.Session, error) Delete(context.Context, *apiv1.Session) (*apiv1.Session, error) } @@ -281,6 +282,24 @@ func (s *Service) Delete(ctx context.Context, id string) (*apiv1.Session, error) return session, nil } +func (s *Service) Revert(ctx context.Context, id string) (*apiv1.Session, error) { + session, err := s.getAuthorized(ctx, id, auth.VerbUpdate) + if err != nil { + return nil, err + } + session, err = s.workflow.Revert(ctx, session) + if errors.Is(err, database.ErrFailedPrecondition) { + return nil, serviceerrors.NewFailedPrecondition(err.Error(), err) + } + if errors.Is(err, database.ErrConflict) { + return nil, serviceerrors.NewAborted(err.Error(), err) + } + if err != nil { + return nil, serviceerrors.NewUnavailable("Failed to revert Session", err) + } + return session, nil +} + // CreateShare mints a read-only share of the caller's session. A positive ttl bounds how long its // token grants access and must not exceed the configured maximum; zero takes the maximum, // or leaves the share valid until it is revoked or the session deleted when none is set. diff --git a/go/core/internal/service/session/service_test.go b/go/core/internal/service/session/service_test.go index 894f152569..a63207de0b 100644 --- a/go/core/internal/service/session/service_test.go +++ b/go/core/internal/service/session/service_test.go @@ -129,6 +129,10 @@ func (w serviceTestWorkflow) Create(_ context.Context, session *apiv1.Session) ( return session, w.err } +func (w serviceTestWorkflow) Revert(_ context.Context, session *apiv1.Session) (*apiv1.Session, error) { + return session, w.err +} + func (w serviceTestWorkflow) Delete(_ context.Context, session *apiv1.Session) (*apiv1.Session, error) { return session, w.err } @@ -201,23 +205,33 @@ func TestServiceCreateRejectsInvalidOrUnauthorizedRequests(t *testing.T) { } } -func TestServiceDeletePreservesContentionErrors(t *testing.T) { - for _, failure := range []struct { - cause error - code serviceerrors.Code +func TestServiceLifecycleMethodsPreserveContentionErrors(t *testing.T) { + for _, test := range []struct { + name string + call func(*Service, context.Context, string) (*apiv1.Session, error) }{ - {database.ErrConflict, serviceerrors.CodeAborted}, - {database.ErrFailedPrecondition, serviceerrors.CodeFailedPrecondition}, + {name: "revert", call: (*Service).Revert}, + {name: "delete", call: (*Service).Delete}, } { - conflict := fmt.Errorf("Session has an operation in progress: %w", failure.cause) - service := NewService(&serviceTestStore{}, serviceTestAuthorizer{}, serviceTestWorkflow{err: conflict}) - _, err := service.Delete(serviceTestContext("alice"), "8bd650a8-9775-488f-8bc1-0d52bf7bdcab") - if !serviceerrors.IsCode(err, failure.code) { - t.Fatalf("error = %v, want code %s", err, failure.code) - } - if serviceerrors.MessageOf(err) != conflict.Error() || !errors.Is(err, failure.cause) { - t.Fatalf("error = %v, want preserved contention reason", err) - } + t.Run(test.name, func(t *testing.T) { + for _, failure := range []struct { + cause error + code serviceerrors.Code + }{ + {database.ErrConflict, serviceerrors.CodeAborted}, + {database.ErrFailedPrecondition, serviceerrors.CodeFailedPrecondition}, + } { + conflict := fmt.Errorf("Session has an operation in progress: %w", failure.cause) + service := NewService(&serviceTestStore{}, serviceTestAuthorizer{}, serviceTestWorkflow{err: conflict}) + _, err := test.call(service, serviceTestContext("alice"), "8bd650a8-9775-488f-8bc1-0d52bf7bdcab") + if !serviceerrors.IsCode(err, failure.code) { + t.Fatalf("error = %v, want code %s", err, failure.code) + } + if serviceerrors.MessageOf(err) != conflict.Error() || !errors.Is(err, failure.cause) { + t.Fatalf("error = %v, want preserved contention reason", err) + } + } + }) } } @@ -553,6 +567,7 @@ func TestServiceRejectsReadOnlyShareMutationsWithoutTransport(t *testing.T) { name string call func() error }{ + {name: "revert", call: func() error { _, err := service.Revert(ctx, id); return err }}, {name: "delete", call: func() error { _, err := service.Delete(ctx, id); return err }}, {name: "rename", call: func() error { _, err := service.Rename(ctx, id, "title"); return err }}, {name: "share listing", call: func() error { _, err := service.ListShares(ctx, id, 0, ""); return err }}, diff --git a/go/core/test/e2e/session_revert_test.go b/go/core/test/e2e/session_revert_test.go new file mode 100644 index 0000000000..83a655b33f --- /dev/null +++ b/go/core/test/e2e/session_revert_test.go @@ -0,0 +1,151 @@ +package e2e_test + +import ( + "context" + "net" + "net/http" + "net/http/httptest" + "net/http/httputil" + "net/url" + "sync/atomic" + "testing" + "time" + + "github.com/a2aproject/a2a-go/v2/a2a" + a2apb "github.com/a2aproject/a2a-go/v2/a2apb/v1" + "github.com/kagent-dev/mockllm" + "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" + corev1 "k8s.io/api/core/v1" + 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" + "github.com/kagent-dev/kagent/go/core/internal/substrate" +) + +func TestSessionRevertCrashedRuntime(t *testing.T) { + target := interactionTarget(t) + forEachHarness(t, func(t *testing.T, harness testHarness) { + cfg, err := mockllm.LoadConfigFromFile("mocks/invoke_agent.json", interactionMocks) + require.NoError(t, err) + upstream, err := url.Parse(startMockLLMConfig(t, cfg)) + require.NoError(t, err) + proxy := httputil.NewSingleHostReverseProxy(upstream) + // Only an interaction runs a Session's Actor; quiescence suspends it once + // the turn settles. The gate holds the second turn's model call open so + // the Actor is still running when its worker is deleted. + var blockNext atomic.Bool + blocked := make(chan struct{}, 1) + release := make(chan struct{}) + gate := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodPost && blockNext.CompareAndSwap(true, false) { + blocked <- struct{}{} + select { + case <-r.Context().Done(): + case <-release: + } + return + } + proxy.ServeHTTP(w, r) + })) + require.NoError(t, gate.Listener.Close()) + gate.Listener, err = net.Listen("tcp", "0.0.0.0:0") + require.NoError(t, err) + gate.Start() + t.Cleanup(gate.Close) + fixture := newInteractionFixture(t, harness, target, reachableModelURL(t, gate.URL)) + defer func() { + close(release) + gate.CloseClientConnections() + }() + _, _, task := fixture.send(t, "What is 2+2?") + require.Equal(t, a2a.TaskStateCompleted, task.Status.State) + assertActorSuspended(t, fixture) + getTask := &a2apb.GetTaskRequest{Tenant: fixture.tenant, Id: string(task.ID)} + before, err := fixture.client.GetTask(fixture.ctx, getTask) + require.NoError(t, err) + session, err := fixture.sessions.GetSession(fixture.ctx, &apiv1.GetSessionRequest{SessionId: fixture.sessionID}) + require.NoError(t, err) + namespace, name := session.Session.Agent.Namespace, substrate.ActorName(fixture.sessionID) + blockNext.Store(true) + _, turn := newMessageRequest(t, "What is 2+2?") + turn.Tenant, turn.Message.ContextId = fixture.tenant, fixture.sessionID + turnCtx, stopTurn := context.WithCancel(fixture.ctx) + defer stopTurn() + sendDone := make(chan error, 1) + go func() { _, err := sendMessageWithRetry(turnCtx, fixture.client, turn); sendDone <- err }() + select { + case <-blocked: + case <-time.After(time.Minute): + t.Fatal("second turn never reached the blocking model") + } + running, err := findSubstrateActor(fixture.ctx, fixture.system, namespace, name) + require.NoError(t, err) + require.NotNil(t, running) + require.Equal(t, ateapipb.ActorState_ACTOR_STATE_RUNNING, running.GetStatus().GetState()) + require.NotEmpty(t, substrate.SnapshotURI(substrate.LatestDurableActorSnapshot(running))) + request := &apiv1.RevertSessionRequest{SessionId: fixture.sessionID} + _, err = fixture.sessions.RevertSession(fixture.ctx, request) + require.Equal(t, codes.FailedPrecondition, status.Code(err), "healthy runtimes cannot revert") + assignment := running.GetStatus().GetWorkerAssignment() + require.Equal(t, namespace, assignment.GetWorkerNamespace()) + require.NotEmpty(t, assignment.GetWorkerPod()) + for token := ""; ; { + page, err := fixture.system.ListSubstrateActors(fixture.ctx, &apiv1.ListSubstrateActorsRequest{ + Page: &apiv1.PageRequest{Limit: 100, PageToken: token}, + }) + require.NoError(t, err) + require.Empty(t, page.GetAteApiError()) + for _, actor := range page.GetActors() { + other := actor.GetStatus().GetWorkerAssignment() + if actor.GetMetadata().GetUid() != running.GetMetadata().GetUid() { + require.False(t, other.GetWorkerNamespace() == assignment.GetWorkerNamespace() && other.GetWorkerPod() == assignment.GetWorkerPod(), + "refuse to delete a worker shared with another Actor") + } + } + token = page.GetPage().GetNextPageToken() + if token == "" { + break + } + } + kube := interactionKubeClient(t) + pod := &corev1.Pod{} + require.NoError(t, kube.Get(fixture.ctx, ctrlclient.ObjectKey{ + Namespace: assignment.GetWorkerNamespace(), Name: assignment.GetWorkerPod(), + }, pod)) + require.NotEmpty(t, pod.UID) + require.NoError(t, kube.Delete(fixture.ctx, pod, &ctrlclient.DeleteOptions{ + GracePeriodSeconds: new(int64(0)), Preconditions: &metav1.Preconditions{UID: &pod.UID}, + })) + err = wait.PollUntilContextTimeout(fixture.ctx, 100*time.Millisecond, time.Minute, true, func(ctx context.Context) (bool, error) { + actor, err := findSubstrateActor(ctx, fixture.system, namespace, name) + return actor.GetStatus().GetState() == ateapipb.ActorState_ACTOR_STATE_CRASHED, err + }) + require.NoError(t, err) + reverted, err := fixture.sessions.RevertSession(fixture.ctx, request) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, reverted.Session.State) + recovered, err := findSubstrateActor(fixture.ctx, fixture.system, namespace, name) + require.NoError(t, err) + require.Equal(t, running.GetMetadata().GetUid(), recovered.GetMetadata().GetUid()) + require.Equal(t, ateapipb.ActorState_ACTOR_STATE_SUSPENDED, recovered.GetStatus().GetState()) + require.Equal(t, substrate.SnapshotURI(substrate.LatestDurableActorSnapshot(running)), substrate.SnapshotURI(substrate.LatestActorSnapshot(recovered))) + retained, err := fixture.client.GetTask(fixture.ctx, getTask) + require.NoError(t, err) + require.True(t, proto.Equal(before, retained), "recovery preserves the completed task and its history") + stopTurn() + select { + case <-sendDone: + case <-time.After(10 * time.Second): + t.Fatal("interrupted caller did not stop") + } + _, _, next := fixture.send(t, "What is 2+2?") + require.Equal(t, a2a.TaskStateCompleted, next.Status.State) + require.Contains(t, taskText(next), "4") + }) +} diff --git a/proto/kagent/api/v1/sessions.proto b/proto/kagent/api/v1/sessions.proto index cb6c9c730f..b5c6dcb820 100644 --- a/proto/kagent/api/v1/sessions.proto +++ b/proto/kagent/api/v1/sessions.proto @@ -20,6 +20,8 @@ service SessionService { rpc GetSession(GetSessionRequest) returns (GetSessionResponse); rpc ListSessions(ListSessionsRequest) returns (ListSessionsResponse); rpc UpdateSessionName(UpdateSessionNameRequest) returns (UpdateSessionNameResponse); + // Restores a crashed runtime to its last snapshot without replaying its turn. + rpc RevertSession(RevertSessionRequest) returns (RevertSessionResponse); rpc DeleteSession(DeleteSessionRequest) returns (DeleteSessionResponse); rpc CreateSessionShare(CreateSessionShareRequest) returns (CreateSessionShareResponse); rpc ListSessionShares(ListSessionSharesRequest) returns (ListSessionSharesResponse); @@ -112,6 +114,14 @@ message UpdateSessionNameResponse { Session session = 1; } +message RevertSessionRequest { + string session_id = 1 [(buf.validate.field).string.uuid = true]; +} + +message RevertSessionResponse { + Session session = 1; +} + message DeleteSessionRequest { string session_id = 1 [(buf.validate.field).string.uuid = true]; } diff --git a/ui/src/generated/kagent/api/v1/sessions_pb.ts b/ui/src/generated/kagent/api/v1/sessions_pb.ts index 8d8b034588..6c2400615d 100644 --- a/ui/src/generated/kagent/api/v1/sessions_pb.ts +++ b/ui/src/generated/kagent/api/v1/sessions_pb.ts @@ -17,7 +17,7 @@ import type { Message } from "@bufbuild/protobuf"; * Describes the file kagent/api/v1/sessions.proto. */ export const file_kagent_api_v1_sessions: GenFile = /*@__PURE__*/ - fileDesc("ChxrYWdlbnQvYXBpL3YxL3Nlc3Npb25zLnByb3RvEg1rYWdlbnQuYXBpLnYxIioKB0ZhaWx1cmUSDgoGcmVhc29uGAEgASgJEg8KB21lc3NhZ2UYAiABKAkixwMKB1Nlc3Npb24SCgoCaWQYASABKAkSDwoHY3JlYXRvchgCIAEoCRIvCgVhZ2VudBgPIAEoCzIgLmthZ2VudC5hcGkudjEuUmVzb3VyY2VSZWZlcmVuY2USGQoRcHJlcGFyZWRfcmV2aXNpb24YBSABKAkSFQoNYTJhX2F1dGhvcml0eRgGIAEoCRIqCgVzdGF0ZRgHIAEoDjIbLmthZ2VudC5hcGkudjEuUnVudGltZVN0YXRlEjIKCW9wZXJhdGlvbhgIIAEoDjIfLmthZ2VudC5hcGkudjEuUnVudGltZU9wZXJhdGlvbhInCgdmYWlsdXJlGAkgASgLMhYua2FnZW50LmFwaS52MS5GYWlsdXJlEi4KCmNyZWF0ZWRfYXQYCiABKAsyGi5nb29nbGUucHJvdG9idWYuVGltZXN0YW1wEi4KCnVwZGF0ZWRfYXQYCyABKAsyGi5nb29nbGUucHJvdG9idWYuVGltZXN0YW1wEgwKBG5hbWUYDSABKAkSEgoKY29udGV4dF9pZBgOIAEoCUoECAMQBEoECAQQBUoECAwQDVIHaGFybmVzc1IOYWdlbnRfdGVtcGxhdGVSBmxhYmVscyLhAQoUQ3JlYXRlU2Vzc2lvblJlcXVlc3QSNwoFYWdlbnQYBSABKAsyIC5rYWdlbnQuYXBpLnYxLlJlc291cmNlUmVmZXJlbmNlQga6SAPIAQESHgoKcmVxdWVzdF9pZBgDIAEoCUIKukgHcgUQARiAARJLCgRuYW1lGAQgASgJQj26SDpyOBjIATIzXig/OiR8W15ccHtafVxwe0NjfV0oPzpbXlxwe0NjfV0qW15ccHtafVxwe0NjfV0pPykkSgQIARACSgQIAhADUgdoYXJuZXNzUg5hZ2VudF90ZW1wbGF0ZSJAChVDcmVhdGVTZXNzaW9uUmVzcG9uc2USJwoHc2Vzc2lvbhgBIAEoCzIWLmthZ2VudC5hcGkudjEuU2Vzc2lvbiIxChFHZXRTZXNzaW9uUmVxdWVzdBIcCgpzZXNzaW9uX2lkGAEgASgJQgi6SAVyA7ABASI9ChJHZXRTZXNzaW9uUmVzcG9uc2USJwoHc2Vzc2lvbhgBIAEoCzIWLmthZ2VudC5hcGkudjEuU2Vzc2lvbiK/AQoTTGlzdFNlc3Npb25zUmVxdWVzdBIUCgxhbGxfY3JlYXRvcnMYAiABKAgSKAoEcGFnZRgDIAEoCzIaLmthZ2VudC5hcGkudjEuUGFnZVJlcXVlc3QSLwoFYWdlbnQYBiABKAsyIC5rYWdlbnQuYXBpLnYxLlJlc291cmNlUmVmZXJlbmNlSgQIARACSgQIBBAFSgQIBRAGUgxtYXRjaF9sYWJlbHNSDmFnZW50X3RlbXBsYXRlUgdoYXJuZXNzImsKFExpc3RTZXNzaW9uc1Jlc3BvbnNlEigKCHNlc3Npb25zGAEgAygLMhYua2FnZW50LmFwaS52MS5TZXNzaW9uEikKBHBhZ2UYAiABKAsyGy5rYWdlbnQuYXBpLnYxLlBhZ2VSZXNwb25zZSKFAQoYVXBkYXRlU2Vzc2lvbk5hbWVSZXF1ZXN0EhwKCnNlc3Npb25faWQYASABKAlCCLpIBXIDsAEBEksKBG5hbWUYAiABKAlCPbpIOnI4GMgBMjNeKD86JHxbXlxwe1p9XHB7Q2N9XSg/OlteXHB7Q2N9XSpbXlxwe1p9XHB7Q2N9XSk/KSQiRAoZVXBkYXRlU2Vzc2lvbk5hbWVSZXNwb25zZRInCgdzZXNzaW9uGAEgASgLMhYua2FnZW50LmFwaS52MS5TZXNzaW9uIjQKFERlbGV0ZVNlc3Npb25SZXF1ZXN0EhwKCnNlc3Npb25faWQYASABKAlCCLpIBXIDsAEBIkAKFURlbGV0ZVNlc3Npb25SZXNwb25zZRInCgdzZXNzaW9uGAEgASgLMhYua2FnZW50LmFwaS52MS5TZXNzaW9uIo4BCgxTZXNzaW9uU2hhcmUSCgoCaWQYASABKAkSEgoKc2Vzc2lvbl9pZBgCIAEoCRIuCgpjcmVhdGVkX2F0GAMgASgLMhouZ29vZ2xlLnByb3RvYnVmLlRpbWVzdGFtcBIuCgpleHBpcmVzX2F0GAQgASgLMhouZ29vZ2xlLnByb3RvYnVmLlRpbWVzdGFtcCJrChlDcmVhdGVTZXNzaW9uU2hhcmVSZXF1ZXN0EhwKCnNlc3Npb25faWQYASABKAlCCLpIBXIDsAEBEjAKA3R0bBgCIAEoCzIZLmdvb2dsZS5wcm90b2J1Zi5EdXJhdGlvbkIIukgFqgECKgAiVwoaQ3JlYXRlU2Vzc2lvblNoYXJlUmVzcG9uc2USKgoFc2hhcmUYASABKAsyGy5rYWdlbnQuYXBpLnYxLlNlc3Npb25TaGFyZRINCgV0b2tlbhgCIAEoCSJiChhMaXN0U2Vzc2lvblNoYXJlc1JlcXVlc3QSHAoKc2Vzc2lvbl9pZBgBIAEoCUIIukgFcgOwAQESKAoEcGFnZRgCIAEoCzIaLmthZ2VudC5hcGkudjEuUGFnZVJlcXVlc3QicwoZTGlzdFNlc3Npb25TaGFyZXNSZXNwb25zZRIrCgZzaGFyZXMYASADKAsyGy5rYWdlbnQuYXBpLnYxLlNlc3Npb25TaGFyZRIpCgRwYWdlGAIgASgLMhsua2FnZW50LmFwaS52MS5QYWdlUmVzcG9uc2UiNwoZUmV2b2tlU2Vzc2lvblNoYXJlUmVxdWVzdBIaCghzaGFyZV9pZBgBIAEoCUIIukgFcgOwAQEiHAoaUmV2b2tlU2Vzc2lvblNoYXJlUmVzcG9uc2UymgYKDlNlc3Npb25TZXJ2aWNlEloKDUNyZWF0ZVNlc3Npb24SIy5rYWdlbnQuYXBpLnYxLkNyZWF0ZVNlc3Npb25SZXF1ZXN0GiQua2FnZW50LmFwaS52MS5DcmVhdGVTZXNzaW9uUmVzcG9uc2USUQoKR2V0U2Vzc2lvbhIgLmthZ2VudC5hcGkudjEuR2V0U2Vzc2lvblJlcXVlc3QaIS5rYWdlbnQuYXBpLnYxLkdldFNlc3Npb25SZXNwb25zZRJXCgxMaXN0U2Vzc2lvbnMSIi5rYWdlbnQuYXBpLnYxLkxpc3RTZXNzaW9uc1JlcXVlc3QaIy5rYWdlbnQuYXBpLnYxLkxpc3RTZXNzaW9uc1Jlc3BvbnNlEmYKEVVwZGF0ZVNlc3Npb25OYW1lEicua2FnZW50LmFwaS52MS5VcGRhdGVTZXNzaW9uTmFtZVJlcXVlc3QaKC5rYWdlbnQuYXBpLnYxLlVwZGF0ZVNlc3Npb25OYW1lUmVzcG9uc2USWgoNRGVsZXRlU2Vzc2lvbhIjLmthZ2VudC5hcGkudjEuRGVsZXRlU2Vzc2lvblJlcXVlc3QaJC5rYWdlbnQuYXBpLnYxLkRlbGV0ZVNlc3Npb25SZXNwb25zZRJpChJDcmVhdGVTZXNzaW9uU2hhcmUSKC5rYWdlbnQuYXBpLnYxLkNyZWF0ZVNlc3Npb25TaGFyZVJlcXVlc3QaKS5rYWdlbnQuYXBpLnYxLkNyZWF0ZVNlc3Npb25TaGFyZVJlc3BvbnNlEmYKEUxpc3RTZXNzaW9uU2hhcmVzEicua2FnZW50LmFwaS52MS5MaXN0U2Vzc2lvblNoYXJlc1JlcXVlc3QaKC5rYWdlbnQuYXBpLnYxLkxpc3RTZXNzaW9uU2hhcmVzUmVzcG9uc2USaQoSUmV2b2tlU2Vzc2lvblNoYXJlEigua2FnZW50LmFwaS52MS5SZXZva2VTZXNzaW9uU2hhcmVSZXF1ZXN0Gikua2FnZW50LmFwaS52MS5SZXZva2VTZXNzaW9uU2hhcmVSZXNwb25zZUI9WjtnaXRodWIuY29tL2thZ2VudC1kZXYva2FnZW50L2dvL2FwaS9nZW4va2FnZW50L2FwaS92MTthcGl2MWIGcHJvdG8z", [file_buf_validate_validate, file_google_protobuf_duration, file_google_protobuf_timestamp, file_kagent_api_v1_common, file_kagent_api_v1_runtime]); + fileDesc("ChxrYWdlbnQvYXBpL3YxL3Nlc3Npb25zLnByb3RvEg1rYWdlbnQuYXBpLnYxIioKB0ZhaWx1cmUSDgoGcmVhc29uGAEgASgJEg8KB21lc3NhZ2UYAiABKAkixwMKB1Nlc3Npb24SCgoCaWQYASABKAkSDwoHY3JlYXRvchgCIAEoCRIvCgVhZ2VudBgPIAEoCzIgLmthZ2VudC5hcGkudjEuUmVzb3VyY2VSZWZlcmVuY2USGQoRcHJlcGFyZWRfcmV2aXNpb24YBSABKAkSFQoNYTJhX2F1dGhvcml0eRgGIAEoCRIqCgVzdGF0ZRgHIAEoDjIbLmthZ2VudC5hcGkudjEuUnVudGltZVN0YXRlEjIKCW9wZXJhdGlvbhgIIAEoDjIfLmthZ2VudC5hcGkudjEuUnVudGltZU9wZXJhdGlvbhInCgdmYWlsdXJlGAkgASgLMhYua2FnZW50LmFwaS52MS5GYWlsdXJlEi4KCmNyZWF0ZWRfYXQYCiABKAsyGi5nb29nbGUucHJvdG9idWYuVGltZXN0YW1wEi4KCnVwZGF0ZWRfYXQYCyABKAsyGi5nb29nbGUucHJvdG9idWYuVGltZXN0YW1wEgwKBG5hbWUYDSABKAkSEgoKY29udGV4dF9pZBgOIAEoCUoECAMQBEoECAQQBUoECAwQDVIHaGFybmVzc1IOYWdlbnRfdGVtcGxhdGVSBmxhYmVscyLhAQoUQ3JlYXRlU2Vzc2lvblJlcXVlc3QSNwoFYWdlbnQYBSABKAsyIC5rYWdlbnQuYXBpLnYxLlJlc291cmNlUmVmZXJlbmNlQga6SAPIAQESHgoKcmVxdWVzdF9pZBgDIAEoCUIKukgHcgUQARiAARJLCgRuYW1lGAQgASgJQj26SDpyOBjIATIzXig/OiR8W15ccHtafVxwe0NjfV0oPzpbXlxwe0NjfV0qW15ccHtafVxwe0NjfV0pPykkSgQIARACSgQIAhADUgdoYXJuZXNzUg5hZ2VudF90ZW1wbGF0ZSJAChVDcmVhdGVTZXNzaW9uUmVzcG9uc2USJwoHc2Vzc2lvbhgBIAEoCzIWLmthZ2VudC5hcGkudjEuU2Vzc2lvbiIxChFHZXRTZXNzaW9uUmVxdWVzdBIcCgpzZXNzaW9uX2lkGAEgASgJQgi6SAVyA7ABASI9ChJHZXRTZXNzaW9uUmVzcG9uc2USJwoHc2Vzc2lvbhgBIAEoCzIWLmthZ2VudC5hcGkudjEuU2Vzc2lvbiK/AQoTTGlzdFNlc3Npb25zUmVxdWVzdBIUCgxhbGxfY3JlYXRvcnMYAiABKAgSKAoEcGFnZRgDIAEoCzIaLmthZ2VudC5hcGkudjEuUGFnZVJlcXVlc3QSLwoFYWdlbnQYBiABKAsyIC5rYWdlbnQuYXBpLnYxLlJlc291cmNlUmVmZXJlbmNlSgQIARACSgQIBBAFSgQIBRAGUgxtYXRjaF9sYWJlbHNSDmFnZW50X3RlbXBsYXRlUgdoYXJuZXNzImsKFExpc3RTZXNzaW9uc1Jlc3BvbnNlEigKCHNlc3Npb25zGAEgAygLMhYua2FnZW50LmFwaS52MS5TZXNzaW9uEikKBHBhZ2UYAiABKAsyGy5rYWdlbnQuYXBpLnYxLlBhZ2VSZXNwb25zZSKFAQoYVXBkYXRlU2Vzc2lvbk5hbWVSZXF1ZXN0EhwKCnNlc3Npb25faWQYASABKAlCCLpIBXIDsAEBEksKBG5hbWUYAiABKAlCPbpIOnI4GMgBMjNeKD86JHxbXlxwe1p9XHB7Q2N9XSg/OlteXHB7Q2N9XSpbXlxwe1p9XHB7Q2N9XSk/KSQiRAoZVXBkYXRlU2Vzc2lvbk5hbWVSZXNwb25zZRInCgdzZXNzaW9uGAEgASgLMhYua2FnZW50LmFwaS52MS5TZXNzaW9uIjQKFFJldmVydFNlc3Npb25SZXF1ZXN0EhwKCnNlc3Npb25faWQYASABKAlCCLpIBXIDsAEBIkAKFVJldmVydFNlc3Npb25SZXNwb25zZRInCgdzZXNzaW9uGAEgASgLMhYua2FnZW50LmFwaS52MS5TZXNzaW9uIjQKFERlbGV0ZVNlc3Npb25SZXF1ZXN0EhwKCnNlc3Npb25faWQYASABKAlCCLpIBXIDsAEBIkAKFURlbGV0ZVNlc3Npb25SZXNwb25zZRInCgdzZXNzaW9uGAEgASgLMhYua2FnZW50LmFwaS52MS5TZXNzaW9uIo4BCgxTZXNzaW9uU2hhcmUSCgoCaWQYASABKAkSEgoKc2Vzc2lvbl9pZBgCIAEoCRIuCgpjcmVhdGVkX2F0GAMgASgLMhouZ29vZ2xlLnByb3RvYnVmLlRpbWVzdGFtcBIuCgpleHBpcmVzX2F0GAQgASgLMhouZ29vZ2xlLnByb3RvYnVmLlRpbWVzdGFtcCJrChlDcmVhdGVTZXNzaW9uU2hhcmVSZXF1ZXN0EhwKCnNlc3Npb25faWQYASABKAlCCLpIBXIDsAEBEjAKA3R0bBgCIAEoCzIZLmdvb2dsZS5wcm90b2J1Zi5EdXJhdGlvbkIIukgFqgECKgAiVwoaQ3JlYXRlU2Vzc2lvblNoYXJlUmVzcG9uc2USKgoFc2hhcmUYASABKAsyGy5rYWdlbnQuYXBpLnYxLlNlc3Npb25TaGFyZRINCgV0b2tlbhgCIAEoCSJiChhMaXN0U2Vzc2lvblNoYXJlc1JlcXVlc3QSHAoKc2Vzc2lvbl9pZBgBIAEoCUIIukgFcgOwAQESKAoEcGFnZRgCIAEoCzIaLmthZ2VudC5hcGkudjEuUGFnZVJlcXVlc3QicwoZTGlzdFNlc3Npb25TaGFyZXNSZXNwb25zZRIrCgZzaGFyZXMYASADKAsyGy5rYWdlbnQuYXBpLnYxLlNlc3Npb25TaGFyZRIpCgRwYWdlGAIgASgLMhsua2FnZW50LmFwaS52MS5QYWdlUmVzcG9uc2UiNwoZUmV2b2tlU2Vzc2lvblNoYXJlUmVxdWVzdBIaCghzaGFyZV9pZBgBIAEoCUIIukgFcgOwAQEiHAoaUmV2b2tlU2Vzc2lvblNoYXJlUmVzcG9uc2Uy9gYKDlNlc3Npb25TZXJ2aWNlEloKDUNyZWF0ZVNlc3Npb24SIy5rYWdlbnQuYXBpLnYxLkNyZWF0ZVNlc3Npb25SZXF1ZXN0GiQua2FnZW50LmFwaS52MS5DcmVhdGVTZXNzaW9uUmVzcG9uc2USUQoKR2V0U2Vzc2lvbhIgLmthZ2VudC5hcGkudjEuR2V0U2Vzc2lvblJlcXVlc3QaIS5rYWdlbnQuYXBpLnYxLkdldFNlc3Npb25SZXNwb25zZRJXCgxMaXN0U2Vzc2lvbnMSIi5rYWdlbnQuYXBpLnYxLkxpc3RTZXNzaW9uc1JlcXVlc3QaIy5rYWdlbnQuYXBpLnYxLkxpc3RTZXNzaW9uc1Jlc3BvbnNlEmYKEVVwZGF0ZVNlc3Npb25OYW1lEicua2FnZW50LmFwaS52MS5VcGRhdGVTZXNzaW9uTmFtZVJlcXVlc3QaKC5rYWdlbnQuYXBpLnYxLlVwZGF0ZVNlc3Npb25OYW1lUmVzcG9uc2USWgoNUmV2ZXJ0U2Vzc2lvbhIjLmthZ2VudC5hcGkudjEuUmV2ZXJ0U2Vzc2lvblJlcXVlc3QaJC5rYWdlbnQuYXBpLnYxLlJldmVydFNlc3Npb25SZXNwb25zZRJaCg1EZWxldGVTZXNzaW9uEiMua2FnZW50LmFwaS52MS5EZWxldGVTZXNzaW9uUmVxdWVzdBokLmthZ2VudC5hcGkudjEuRGVsZXRlU2Vzc2lvblJlc3BvbnNlEmkKEkNyZWF0ZVNlc3Npb25TaGFyZRIoLmthZ2VudC5hcGkudjEuQ3JlYXRlU2Vzc2lvblNoYXJlUmVxdWVzdBopLmthZ2VudC5hcGkudjEuQ3JlYXRlU2Vzc2lvblNoYXJlUmVzcG9uc2USZgoRTGlzdFNlc3Npb25TaGFyZXMSJy5rYWdlbnQuYXBpLnYxLkxpc3RTZXNzaW9uU2hhcmVzUmVxdWVzdBooLmthZ2VudC5hcGkudjEuTGlzdFNlc3Npb25TaGFyZXNSZXNwb25zZRJpChJSZXZva2VTZXNzaW9uU2hhcmUSKC5rYWdlbnQuYXBpLnYxLlJldm9rZVNlc3Npb25TaGFyZVJlcXVlc3QaKS5rYWdlbnQuYXBpLnYxLlJldm9rZVNlc3Npb25TaGFyZVJlc3BvbnNlQj1aO2dpdGh1Yi5jb20va2FnZW50LWRldi9rYWdlbnQvZ28vYXBpL2dlbi9rYWdlbnQvYXBpL3YxO2FwaXYxYgZwcm90bzM", [file_buf_validate_validate, file_google_protobuf_duration, file_google_protobuf_timestamp, file_kagent_api_v1_common, file_kagent_api_v1_runtime]); /** * @generated from message kagent.api.v1.Failure @@ -294,6 +294,40 @@ export type UpdateSessionNameResponse = Message<"kagent.api.v1.UpdateSessionName export const UpdateSessionNameResponseSchema: GenMessage = /*@__PURE__*/ messageDesc(file_kagent_api_v1_sessions, 9); +/** + * @generated from message kagent.api.v1.RevertSessionRequest + */ +export type RevertSessionRequest = Message<"kagent.api.v1.RevertSessionRequest"> & { + /** + * @generated from field: string session_id = 1; + */ + sessionId: string; +}; + +/** + * Describes the message kagent.api.v1.RevertSessionRequest. + * Use `create(RevertSessionRequestSchema)` to create a new message. + */ +export const RevertSessionRequestSchema: GenMessage = /*@__PURE__*/ + messageDesc(file_kagent_api_v1_sessions, 10); + +/** + * @generated from message kagent.api.v1.RevertSessionResponse + */ +export type RevertSessionResponse = Message<"kagent.api.v1.RevertSessionResponse"> & { + /** + * @generated from field: kagent.api.v1.Session session = 1; + */ + session?: Session | undefined; +}; + +/** + * Describes the message kagent.api.v1.RevertSessionResponse. + * Use `create(RevertSessionResponseSchema)` to create a new message. + */ +export const RevertSessionResponseSchema: GenMessage = /*@__PURE__*/ + messageDesc(file_kagent_api_v1_sessions, 11); + /** * @generated from message kagent.api.v1.DeleteSessionRequest */ @@ -309,7 +343,7 @@ export type DeleteSessionRequest = Message<"kagent.api.v1.DeleteSessionRequest"> * Use `create(DeleteSessionRequestSchema)` to create a new message. */ export const DeleteSessionRequestSchema: GenMessage = /*@__PURE__*/ - messageDesc(file_kagent_api_v1_sessions, 10); + messageDesc(file_kagent_api_v1_sessions, 12); /** * @generated from message kagent.api.v1.DeleteSessionResponse @@ -326,7 +360,7 @@ export type DeleteSessionResponse = Message<"kagent.api.v1.DeleteSessionResponse * Use `create(DeleteSessionResponseSchema)` to create a new message. */ export const DeleteSessionResponseSchema: GenMessage = /*@__PURE__*/ - messageDesc(file_kagent_api_v1_sessions, 11); + messageDesc(file_kagent_api_v1_sessions, 13); /** * Grants read-only access to a session, including A2A get, list, and subscribe. @@ -362,7 +396,7 @@ export type SessionShare = Message<"kagent.api.v1.SessionShare"> & { * Use `create(SessionShareSchema)` to create a new message. */ export const SessionShareSchema: GenMessage = /*@__PURE__*/ - messageDesc(file_kagent_api_v1_sessions, 12); + messageDesc(file_kagent_api_v1_sessions, 14); /** * @generated from message kagent.api.v1.CreateSessionShareRequest @@ -388,7 +422,7 @@ export type CreateSessionShareRequest = Message<"kagent.api.v1.CreateSessionShar * Use `create(CreateSessionShareRequestSchema)` to create a new message. */ export const CreateSessionShareRequestSchema: GenMessage = /*@__PURE__*/ - messageDesc(file_kagent_api_v1_sessions, 13); + messageDesc(file_kagent_api_v1_sessions, 15); /** * @generated from message kagent.api.v1.CreateSessionShareResponse @@ -412,7 +446,7 @@ export type CreateSessionShareResponse = Message<"kagent.api.v1.CreateSessionSha * Use `create(CreateSessionShareResponseSchema)` to create a new message. */ export const CreateSessionShareResponseSchema: GenMessage = /*@__PURE__*/ - messageDesc(file_kagent_api_v1_sessions, 14); + messageDesc(file_kagent_api_v1_sessions, 16); /** * @generated from message kagent.api.v1.ListSessionSharesRequest @@ -434,7 +468,7 @@ export type ListSessionSharesRequest = Message<"kagent.api.v1.ListSessionSharesR * Use `create(ListSessionSharesRequestSchema)` to create a new message. */ export const ListSessionSharesRequestSchema: GenMessage = /*@__PURE__*/ - messageDesc(file_kagent_api_v1_sessions, 15); + messageDesc(file_kagent_api_v1_sessions, 17); /** * @generated from message kagent.api.v1.ListSessionSharesResponse @@ -456,7 +490,7 @@ export type ListSessionSharesResponse = Message<"kagent.api.v1.ListSessionShares * Use `create(ListSessionSharesResponseSchema)` to create a new message. */ export const ListSessionSharesResponseSchema: GenMessage = /*@__PURE__*/ - messageDesc(file_kagent_api_v1_sessions, 16); + messageDesc(file_kagent_api_v1_sessions, 18); /** * @generated from message kagent.api.v1.RevokeSessionShareRequest @@ -473,7 +507,7 @@ export type RevokeSessionShareRequest = Message<"kagent.api.v1.RevokeSessionShar * Use `create(RevokeSessionShareRequestSchema)` to create a new message. */ export const RevokeSessionShareRequestSchema: GenMessage = /*@__PURE__*/ - messageDesc(file_kagent_api_v1_sessions, 17); + messageDesc(file_kagent_api_v1_sessions, 19); /** * @generated from message kagent.api.v1.RevokeSessionShareResponse @@ -486,7 +520,7 @@ export type RevokeSessionShareResponse = Message<"kagent.api.v1.RevokeSessionSha * Use `create(RevokeSessionShareResponseSchema)` to create a new message. */ export const RevokeSessionShareResponseSchema: GenMessage = /*@__PURE__*/ - messageDesc(file_kagent_api_v1_sessions, 18); + messageDesc(file_kagent_api_v1_sessions, 20); /** * Lifecycle mutations execute inline. After a retryable error, repeat the same @@ -530,6 +564,16 @@ export const SessionService: GenService<{ input: typeof UpdateSessionNameRequestSchema; output: typeof UpdateSessionNameResponseSchema; }, + /** + * Restores a crashed runtime to its last snapshot without replaying its turn. + * + * @generated from rpc kagent.api.v1.SessionService.RevertSession + */ + revertSession: { + methodKind: "unary"; + input: typeof RevertSessionRequestSchema; + output: typeof RevertSessionResponseSchema; + }, /** * @generated from rpc kagent.api.v1.SessionService.DeleteSession */ From b885681c2eef407deee16e8715adb635890d62fa Mon Sep 17 00:00:00 2001 From: Ada Lundhe Date: Tue, 3 Nov 2026 09:56:56 -0800 Subject: [PATCH 04/12] fix(session): scan crash recoveries on their own ticker, not quiescence worker 0 Start ran the crash-recovery scan inline on quiescence worker 0, between idle claims. A slow idle suspension on that worker held every recovery retry for as long as the suspension took, and a slow revert held up idle work in turn. No test exercised the scan through Start: the existing tests called recoverSessionPage directly, and the Start-level store stub returned no candidates. The scan now runs in its own goroutine with its own ticker, at the same recovery interval. The four quiescence workers return to their shape before recovery existed. The page cursor still carries across ticks, so a backlog larger than one page cannot starve the Sessions past it. TestStartRetriesAdmittedRecoveryWhileIdleWorkBlocks admits a recovery whose first revert fails, then starts the workflow with all four quiescence workers blocked in idle suspensions. The Session must reach READY while they stay blocked. Verified: the new test fails on the previous Start ("Condition never satisfied") and passes with this change. go test for ./internal/service/..., ./internal/substrate/... and ./internal/database/... passes against Postgres; golangci-lint reports 0 issues. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Ada Lundhe --- .../internal/service/session/quiescence.go | 45 +++++++++++------ .../internal/service/session/recovery_test.go | 48 +++++++++++++++++++ 2 files changed, 78 insertions(+), 15 deletions(-) diff --git a/go/core/internal/service/session/quiescence.go b/go/core/internal/service/session/quiescence.go index aba8c40ffb..8d741149e8 100644 --- a/go/core/internal/service/session/quiescence.go +++ b/go/core/internal/service/session/quiescence.go @@ -30,27 +30,18 @@ func (w *ActorWorkflow) wakeQuiescence() { // Start suspends idle sessions to VOLUMES independently of task publication and // retries crash recoveries admitted by explicit requests. -// Startup and periodic scans discover idle work missed by local signals; crash -// checks use the same recovery interval. Workers drain claims with bounded concurrency. +// Startup and periodic scans discover idle work missed by local signals. Workers +// drain claims with bounded concurrency. Crash recovery scans on its own ticker, so +// a slow idle suspension cannot delay a recovery retry and a slow revert cannot +// hold up idle work. func (w *ActorWorkflow) Start(ctx context.Context) error { var workers sync.WaitGroup - for worker := range 4 { + workers.Go(func() { w.recoverCrashedSessions(ctx) }) + for range 4 { workers.Go(func() { timer := time.NewTimer(w.recoveryInterval) defer timer.Stop() - var recoveryAfter string - nextRecovery := time.Now().Add(w.recoveryInterval) for ctx.Err() == nil { - if worker == 0 && !time.Now().Before(nextRecovery) { - scanCtx, cancel := context.WithTimeout(ctx, database.RuntimeOperationTimeout) - after, err := w.recoverSessionPage(scanCtx, recoveryAfter) - cancel() - if err != nil && ctx.Err() == nil { - logging.FromContext(ctx).ErrorContext(ctx, "scan interrupted sessions", "error", err) - } - recoveryAfter = after - nextRecovery = time.Now().Add(w.recoveryInterval) - } work, err := w.store.ClaimSessionQuiescence(ctx) if err == nil { // Hand off a hint before runtime I/O so a single settlement @@ -85,6 +76,30 @@ func (w *ActorWorkflow) Start(ctx context.Context) error { return nil } +// recoverCrashedSessions retries admitted crash recoveries one page per tick. +// The cursor carries across ticks, so a large backlog cannot starve the Sessions +// after its first page, and an Actor that stays unreadable is revisited when the +// scan wraps. +func (w *ActorWorkflow) recoverCrashedSessions(ctx context.Context) { + ticker := time.NewTicker(w.recoveryInterval) + defer ticker.Stop() + var after string + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + scanCtx, cancel := context.WithTimeout(ctx, database.RuntimeOperationTimeout) + next, err := w.recoverSessionPage(scanCtx, after) + cancel() + if err != nil && ctx.Err() == nil { + logging.FromContext(ctx).ErrorContext(ctx, "scan interrupted sessions", "error", err) + } + after = next + } +} + func (w *ActorWorkflow) quiesceIdleSession(ctx context.Context, work *database.SessionQuiescence) { runtimeCtx, cancel := context.WithTimeout(ctx, 2*time.Minute) snapshot, err := w.Quiesce(runtimeCtx, work.Session) diff --git a/go/core/internal/service/session/recovery_test.go b/go/core/internal/service/session/recovery_test.go index d7ce338a83..0a7ef5a51c 100644 --- a/go/core/internal/service/session/recovery_test.go +++ b/go/core/internal/service/session/recovery_test.go @@ -5,6 +5,7 @@ import ( "crypto/sha256" "fmt" "slices" + "sync/atomic" "testing" "time" @@ -539,3 +540,50 @@ func TestSessionRecoveryPagesPastUnavailableActors(t *testing.T) { require.Equal(t, ateapipb.ActorState_ACTOR_STATE_SUSPENDED, actor.GetStatus().GetState()) } } + +// blockedSuspendActors holds every idle suspension in runtime I/O until the +// test ends, the way a slow Substrate upload would. +type blockedSuspendActors struct { + *lifecycleTestActors + entered atomic.Int32 +} + +func (a *blockedSuspendActors) SuspendActor(ctx context.Context, _, _ string) (*ateapipb.Actor, error) { + a.entered.Add(1) + <-ctx.Done() + return nil, ctx.Err() +} + +// Start retries an admitted recovery on its own ticker. Every quiescence worker +// is blocked in an idle suspension for the whole test, so a scan that shared +// those workers would never run. +func TestStartRetriesAdmittedRecoveryWhileIdleWorkBlocks(t *testing.T) { + store, base, crashed, _, _ := crashedTurnFixture(t) + unavailable := &recoveryTestActors{lifecycleTestActors: base, revert: func(context.Context, string, string) (*ateapipb.Actor, error) { + return nil, status.Error(codes.Unavailable, "control plane unavailable") + }} + _, err := NewActorWorkflow(store, unavailable, make(chan struct{}, 1), time.Hour).Revert(t.Context(), crashed) + require.Error(t, err, "the first attempt leaves an admitted recovery for Start to retry") + + const idleSessions = 4 // one per quiescence worker + for range idleSessions { + idle, _, err := store.CreateSession(t.Context(), &apiv1.Session{Id: uuid.NewString(), Creator: crashed.Creator, Agent: crashed.Agent}, uuid.NewString()) + require.NoError(t, err) + idle, err = NewActorWorkflow(store, base, make(chan struct{}, 1), time.Hour).Create(t.Context(), idle) + require.NoError(t, err) + request := stageQuiescenceTask(t, store.Client, idle, a2a.TaskStateCompleted) + require.NoError(t, store.SettleSessionTask(t.Context(), idle.Id, request.TaskId, request.Version)) + } + + actors := &blockedSuspendActors{lifecycleTestActors: base} + startQuiescenceWorker(t, NewActorWorkflow(store, actors, make(chan struct{}, 1), 20*time.Millisecond)) + require.Eventually(t, func() bool { return actors.entered.Load() == idleSessions }, 5*time.Second, 10*time.Millisecond) + require.Eventually(t, func() bool { + current, err := store.GetSessionByID(t.Context(), crashed.Id) + return err == nil && current.State == apiv1.RuntimeState_RUNTIME_STATE_READY + }, 5*time.Second, 10*time.Millisecond) + actor, err := base.GetActor(t.Context(), "team-a", substrate.ActorName(crashed.Id)) + require.NoError(t, err) + require.Equal(t, ateapipb.ActorState_ACTOR_STATE_SUSPENDED, actor.GetStatus().GetState()) + require.EqualValues(t, idleSessions, actors.entered.Load(), "idle work stayed blocked while recovery completed") +} From cfb50f6654a73bb18313483f1e9d82523fe9f938 Mon Sep 17 00:00:00 2001 From: Ada Lundhe Date: Tue, 3 Nov 2026 09:59:12 -0800 Subject: [PATCH 05/12] fix(session): revert crashed runtimes through execute, reading before the claim Crash recovery had its own copy of execute: claim, GetActor, verify, revert, finish. It claimed before its reads, and a released claim keeps executor_id. One failed GetActor therefore made Delete refuse the Session with "begin Session operation: resource conflict" until recovery succeeded. A deleted Actor made that permanent. Execute also sent every operation on a crashed Session into the recovery copy, so even a Delete that superseded the recovery would have returned a conflict. Recovery now runs through execute like suspend, resume and delete. substrate.PrepareActorRevert does the read-only checks before the claim, and ApplyActorTransition issues RevertActor, which moves onto LifecycleClient. The checks are unchanged: the Actor keeps its UID, ends SUSPENDED, and restores the durable snapshot observed before the revert. A SUSPENDED Actor still completes a recovery whose revert response was lost. A failed preparation keeps the RESUME intent instead of clearing it, because admission already failed the interrupted turn. Nothing was issued, so Delete can supersede it. Rejected: keeping the separate path and only moving its reads above the claim. That fixes the stuck Delete, but fixes to execute would still have to land twice. Verified: TestDeleteSupersedesUnissuedCrashRecovery covers an unreadable and a missing Actor; both fail on the previous code with "resource conflict" and pass here. TestActorRevertRestoresTheObservedSnapshotInPlace covers the revert transition's states and its UID and snapshot checks. go test for ./internal/service/..., ./internal/substrate/... and ./internal/database/... passes against Postgres; golangci-lint reports 0 issues. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Ada Lundhe --- go/core/internal/service/session/recovery.go | 63 --------------- .../internal/service/session/recovery_test.go | 48 ++++++++++++ go/core/internal/service/session/workflow.go | 31 ++++++-- go/core/internal/substrate/lifecycle.go | 51 ++++++++++++ go/core/internal/substrate/lifecycle_test.go | 77 +++++++++++++++++++ 5 files changed, 199 insertions(+), 71 deletions(-) diff --git a/go/core/internal/service/session/recovery.go b/go/core/internal/service/session/recovery.go index 27922bc3ed..cecde11d45 100644 --- a/go/core/internal/service/session/recovery.go +++ b/go/core/internal/service/session/recovery.go @@ -4,11 +4,8 @@ import ( "context" "errors" "fmt" - "time" - "github.com/google/uuid" ateapipb "go.ate.dev/substrate/api/ateapi/v1" - "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" @@ -92,63 +89,3 @@ func (w *ActorWorkflow) Revert(ctx context.Context, session *apiv1.Session) (*ap } return w.execute(ctx, operation) } - -// recoverCrashedSession reconciles lost revert responses without replacing the Actor. -func (w *ActorWorkflow) recoverCrashedSession(ctx context.Context, operation *database.SessionOperation) (_ *apiv1.Session, err error) { - session := operation.Instance - if session.Operation != apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME { - return nil, database.ErrConflict - } - executorID := uuid.New() - claimed, err := w.store.ClaimSessionOperation(ctx, session.Id, operation.ID, executorID) - if err != nil { - return nil, err - } - if !claimed { - current, err := w.store.GetSessionOperation(ctx, session.Id, operation.ID) - if err != nil { - return nil, err - } - return operationOutcome(current) - } - defer func() { - if err != nil { - releaseCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) - defer cancel() - err = errors.Join(err, w.store.ReleaseRuntimeOperation(releaseCtx, session.Id, operation.ID, executorID)) - } - }() - revision, err := w.store.GetRuntimeRevision(ctx, session.PreparedRevision) - if err != nil { - return nil, err - } - atespace, name := revision.ActorTemplateAtespace, substrate.ActorName(session.Id) - actor, err := w.actors.GetActor(ctx, atespace, name) - if err != nil { - return nil, err - } - if err := w.verifyActor(ctx, session, revision, actor); err != nil { - return nil, err - } - before := actor - beforeSnapshot := proto.CloneOf(substrate.LatestDurableActorSnapshot(actor)) - if beforeSnapshot != nil { - beforeSnapshot.Locality = "" - } - switch actor.GetStatus().GetState() { - case ateapipb.ActorState_ACTOR_STATE_CRASHED, ateapipb.ActorState_ACTOR_STATE_REVERTING: - actor, err = w.actors.RevertActor(ctx, atespace, name) - if err != nil { - return nil, fmt.Errorf("revert crashed Actor %s/%s: %w", atespace, name, err) - } - case ateapipb.ActorState_ACTOR_STATE_SUSPENDED: - default: - return nil, fmt.Errorf("crashed Actor %s/%s has unexpected state %s", atespace, name, actor.GetStatus().GetState()) - } - if !validActorIdentity(actor, revision, name) || actor.GetMetadata().GetUid() != before.GetMetadata().GetUid() || actor.GetStatus().GetState() != ateapipb.ActorState_ACTOR_STATE_SUSPENDED || !proto.Equal(beforeSnapshot, substrate.LatestActorSnapshot(actor)) { - return nil, fmt.Errorf("revert Actor %s/%s returned unexpected identity or state", atespace, name) - } - finishCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) - defer cancel() - return w.store.FinishSessionOperation(finishCtx, session.Id, operation.ID, executorID, "", actor.GetMetadata().GetUid(), "") -} diff --git a/go/core/internal/service/session/recovery_test.go b/go/core/internal/service/session/recovery_test.go index 0a7ef5a51c..3769f82e22 100644 --- a/go/core/internal/service/session/recovery_test.go +++ b/go/core/internal/service/session/recovery_test.go @@ -587,3 +587,51 @@ func TestStartRetriesAdmittedRecoveryWhileIdleWorkBlocks(t *testing.T) { require.Equal(t, ateapipb.ActorState_ACTOR_STATE_SUSPENDED, actor.GetStatus().GetState()) require.EqualValues(t, idleSessions, actors.entered.Load(), "idle work stayed blocked while recovery completed") } + +// A recovery whose Actor cannot be read has issued nothing, so it must not take +// a claim that would make Delete refuse the Session. +func TestDeleteSupersedesUnissuedCrashRecovery(t *testing.T) { + for _, test := range []struct { + name string + gone bool + readErr error + }{ + {name: "actor read fails", readErr: status.Error(codes.Unavailable, "control plane unavailable")}, + {name: "actor is gone", gone: true}, + } { + t.Run(test.name, func(t *testing.T) { + store, base, session, _, _ := crashedTurnFixture(t) + key := actorKey("team-a", substrate.ActorName(session.Id)) + _, err := store.BeginCrashedSessionRecovery(t.Context(), session.Id, revertGeneration(t, store, session.Id), base.actors[key].GetMetadata().GetUid(), database.SessionTaskSnapshot{}) + require.NoError(t, err) + if test.gone { + delete(base.actors, key) + } + actors := &recoveryTestActors{lifecycleTestActors: base, readErr: test.readErr} + actors.revert = func(context.Context, string, string) (*ateapipb.Actor, error) { + t.Fatal("an unread Actor was reverted") + return nil, nil + } + workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute) + admitted, err := store.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + _, err = workflow.Revert(t.Context(), admitted) + require.Error(t, err) + _, err = workflow.recoverSessionPage(t.Context(), "") + require.NoError(t, err) + + pending, err := store.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_FAILED, pending.State) + require.Equal(t, apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME, pending.Operation, "a failed read keeps the recovery for the next scan") + deleted, err := NewActorWorkflow(store, base, make(chan struct{}, 1), time.Minute).Delete(t.Context(), pending) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_DELETED, deleted.State) + _, err = base.GetActor(t.Context(), "team-a", substrate.ActorName(session.Id)) + require.Equal(t, codes.NotFound, status.Code(err)) + page, err := store.ListSessionRecoveryCandidates(t.Context(), "", 100) + require.NoError(t, err) + require.Empty(t, page) + }) + } +} diff --git a/go/core/internal/service/session/workflow.go b/go/core/internal/service/session/workflow.go index ea99fa4eb9..8698bf9adb 100644 --- a/go/core/internal/service/session/workflow.go +++ b/go/core/internal/service/session/workflow.go @@ -33,7 +33,6 @@ type workflowStore interface { type actorClient interface { substrate.LifecycleClient - RevertActor(context.Context, string, string) (*ateapipb.Actor, error) } // ActorWorkflow runs the imperative Substrate operations behind Session @@ -133,9 +132,6 @@ func (w *ActorWorkflow) execute(ctx context.Context, operation *database.Session if operation.Instance.Operation == apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE { return operationOutcome(operation) } - if operation.Instance.State == apiv1.RuntimeState_RUNTIME_STATE_FAILED && operation.Instance.Failure.GetReason() == database.SessionRuntimeCrashed { - return w.recoverCrashedSession(ctx, operation) - } kind := operation.Instance.Operation session := operation.Instance revision, err := w.store.GetRuntimeRevision(ctx, session.GetPreparedRevision()) @@ -164,11 +160,15 @@ func (w *ActorWorkflow) execute(ctx context.Context, operation *database.Session creation.Snapshot = &substrate.ActorSnapshot{Tag: &ateapipb.ObjectRef{Atespace: snapshot.Atespace, Name: tagName}, URI: snapshot.URI, Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES} } } - prepare := substrate.PrepareActorTransition - if operation.ExecutorID != uuid.Nil { - prepare = substrate.PrepareActorRetry + var transition *substrate.ActorTransition + switch { + case crashRecovery(operation): + transition, err = substrate.PrepareActorRevert(ctx, w.actors, binding) + case operation.ExecutorID != uuid.Nil: + transition, err = substrate.PrepareActorRetry(ctx, w.actors, binding, kind, creation) + default: + transition, err = substrate.PrepareActorTransition(ctx, w.actors, binding, kind, creation) } - transition, err := prepare(ctx, w.actors, binding, kind, creation) if err != nil { return w.failPreparation(ctx, operation, err) } @@ -221,6 +221,12 @@ func (w *ActorWorkflow) failPreparation(ctx context.Context, admitted *database. if admitted.ExecutorID != uuid.Nil { return nil, cause // An earlier attempt may have issued work; retain its intent. } + if crashRecovery(admitted) { + // Admission already failed the interrupted turn, so clearing the operation + // would leave a FAILED Session that nothing retries. Keep the intent for the + // next scan. No claim was taken, so Delete can still supersede it. + return nil, cause + } finishCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) defer cancel() _, err := w.store.FinishSessionOperation(finishCtx, admitted.Instance.Id, admitted.ID, uuid.Nil, "", "", "lifecycle preparation failed") @@ -234,6 +240,15 @@ func (w *ActorWorkflow) failPreparation(ctx context.Context, admitted *database. return nil, errors.Join(cause, err) } +// crashRecovery reports a RESUME admitted after Substrate confirmed a crash. +// It reverts the recorded Actor in place instead of resuming it. +func crashRecovery(operation *database.SessionOperation) bool { + session := operation.Instance + return session.State == apiv1.RuntimeState_RUNTIME_STATE_FAILED && + session.Operation == apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME && + session.Failure.GetReason() == database.SessionRuntimeCrashed +} + func operationOutcome(operation *database.SessionOperation) (*apiv1.Session, error) { if operation.Instance.Operation == apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE { return operation.Instance, nil diff --git a/go/core/internal/substrate/lifecycle.go b/go/core/internal/substrate/lifecycle.go index 49450c2830..773de26c6b 100644 --- a/go/core/internal/substrate/lifecycle.go +++ b/go/core/internal/substrate/lifecycle.go @@ -19,6 +19,7 @@ type LifecycleClient interface { CheckActorEgressPolicy(context.Context, string, string) error EnsureActorEgressPolicy(context.Context, string, string, string) error ResumeActor(context.Context, string, string) (*ateapipb.Actor, error) + RevertActor(context.Context, string, string) (*ateapipb.Actor, error) SuspendActor(context.Context, string, string) (*ateapipb.Actor, error) DeleteActor(context.Context, string, string) error } @@ -51,6 +52,9 @@ type ActorTransition struct { operation apiv1.RuntimeOperation actor *ateapipb.Actor creation *ActorCreation + // revert restores a crashed Actor to its last durable snapshot. The + // operation stays RESUME because that is what the Session records. + revert bool } // ActorUID returns the observed runtime identity, including after creation. @@ -118,6 +122,26 @@ func prepareActorTransition(ctx context.Context, actors LifecycleClient, binding return &ActorTransition{binding: binding, operation: operation, actor: actor, creation: creation}, nil } +// PrepareActorRevert performs the read-only checks for crash recovery before +// the store's durable claim. Reading first keeps a failed or missing Actor from +// leaving a claim behind, so Delete can still supersede the unissued recovery. +// A SUSPENDED Actor means an earlier revert completed and its response was lost. +func PrepareActorRevert(ctx context.Context, actors LifecycleClient, binding ActorBinding) (*ActorTransition, error) { + actor, err := actors.GetActor(ctx, binding.Atespace, binding.Name) + if err != nil { + return nil, fmt.Errorf("get Actor %s/%s: %w", binding.Atespace, binding.Name, err) + } + if !binding.Matches(actor) { + return nil, fmt.Errorf("actor %s/%s identity or template changed", binding.Atespace, binding.Name) + } + switch actor.GetStatus().GetState() { + case ateapipb.ActorState_ACTOR_STATE_CRASHED, ateapipb.ActorState_ACTOR_STATE_REVERTING, ateapipb.ActorState_ACTOR_STATE_SUSPENDED: + default: + return nil, fmt.Errorf("crashed Actor %s/%s has unexpected state %s", binding.Atespace, binding.Name, actor.GetStatus().GetState()) + } + return &ActorTransition{binding: binding, operation: apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME, actor: actor, revert: true}, nil +} + // ApplyActorTransition requires durable execution authority. Errors retain the // claim's durable intent and pins. Another bounded attempt can retry that intent. func ApplyActorTransition(ctx context.Context, actors LifecycleClient, transition *ActorTransition) error { @@ -125,6 +149,9 @@ func ApplyActorTransition(ctx context.Context, actors LifecycleClient, transitio return err } binding, actor := transition.binding, transition.actor + if transition.revert { + return transition.applyRevert(ctx, actors) + } var err error switch transition.operation { case apiv1.RuntimeOperation_RUNTIME_OPERATION_CREATE: @@ -202,6 +229,30 @@ func ApplyActorTransition(ctx context.Context, actors LifecycleClient, transitio return nil } +// applyRevert must not replace the Actor or roll it past the snapshot observed +// before the claim. A changed UID or snapshot means something other than this +// recovery acted on the runtime, so the Session must not adopt the result. +func (t *ActorTransition) applyRevert(ctx context.Context, actors LifecycleClient) error { + binding, actor := t.binding, t.actor + before := proto.CloneOf(LatestDurableActorSnapshot(actor)) + if before != nil { + // Revert restores onto any node, so the locality of the source copy is + // expected to change. + before.Locality = "" + } + if state := actor.GetStatus().GetState(); state == ateapipb.ActorState_ACTOR_STATE_CRASHED || state == ateapipb.ActorState_ACTOR_STATE_REVERTING { + reverted, err := actors.RevertActor(ctx, binding.Atespace, binding.Name) + if err != nil { + return fmt.Errorf("revert crashed Actor %s/%s: %w", binding.Atespace, binding.Name, err) + } + actor = reverted + } + if t.checkResult(actor, ateapipb.ActorState_ACTOR_STATE_SUSPENDED) != nil || !proto.Equal(before, LatestActorSnapshot(actor)) { + return fmt.Errorf("revert Actor %s/%s returned unexpected identity or state", binding.Atespace, binding.Name) + } + return nil +} + // Matches checks that the actor belongs to the expected runtime and template. func (b ActorBinding) Matches(actor *ateapipb.Actor) bool { return actor.GetMetadata().GetName() == b.Name && actor.GetMetadata().GetAtespace() == b.Atespace && actor.GetMetadata().GetUid() != "" && diff --git a/go/core/internal/substrate/lifecycle_test.go b/go/core/internal/substrate/lifecycle_test.go index 45e12c8ad4..b462c3ebc7 100644 --- a/go/core/internal/substrate/lifecycle_test.go +++ b/go/core/internal/substrate/lifecycle_test.go @@ -23,6 +23,8 @@ type transitionActors struct { checkErr error resumeUID string suspendSnapshot *ateapipb.Snapshot + revertUID string + revertSnapshot *ateapipb.Snapshot } func (a *transitionActors) GetActor(context.Context, string, string) (*ateapipb.Actor, error) { @@ -70,6 +72,26 @@ func (a *transitionActors) SuspendActor(context.Context, string, string) (*ateap return proto.CloneOf(a.actor), nil } +// RevertActor restores the latest durable snapshot unless the test overrides +// the identity or snapshot it returns. +func (a *transitionActors) RevertActor(context.Context, string, string) (*ateapipb.Actor, error) { + a.calls = append(a.calls, "revert") + a.actor.Status.State = ateapipb.ActorState_ACTOR_STATE_SUSPENDED + snapshot := proto.CloneOf(LatestDurableActorSnapshot(a.actor)) + if a.revertSnapshot != nil { + snapshot = proto.CloneOf(a.revertSnapshot) + } + a.actor.Status.Snapshots = nil + if snapshot != nil { + snapshot.Locality = "" + a.actor.Status.Snapshots = []*ateapipb.Snapshot{snapshot} + } + if a.revertUID != "" { + a.actor.Metadata.Uid = a.revertUID + } + return proto.CloneOf(a.actor), nil +} + func (a *transitionActors) DeleteActor(context.Context, string, string) error { a.calls = append(a.calls, "delete") a.actor = nil @@ -230,3 +252,58 @@ func TestActorSuspendRequiresDurableVolumesSnapshot(t *testing.T) { }) } } + +func TestActorRevertRestoresTheObservedSnapshotInPlace(t *testing.T) { + durable := func(generation int32, uri string) *ateapipb.Snapshot { + return &ateapipb.Snapshot{Generation: generation, Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES, Locality: "lost-node", + DurableSnapshot: &ateapipb.SnapshotStorage{Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_COMPLETED, Object: &ateapipb.ObjectSnapshot{SnapshotUri: uri}}} + } + for _, test := range []struct { + name string + missing bool + state ateapipb.ActorState + revertUID string + revertSnapshot *ateapipb.Snapshot + wantPrepareErr bool + wantApplyErr bool + wantCalls []string + }{ + {name: "crashed", state: ateapipb.ActorState_ACTOR_STATE_CRASHED, wantCalls: []string{"revert"}}, + {name: "revert in progress", state: ateapipb.ActorState_ACTOR_STATE_REVERTING, wantCalls: []string{"revert"}}, + {name: "earlier revert response lost", state: ateapipb.ActorState_ACTOR_STATE_SUSPENDED}, + {name: "running actor", state: ateapipb.ActorState_ACTOR_STATE_RUNNING, wantPrepareErr: true}, + {name: "actor gone", missing: true, wantPrepareErr: true}, + {name: "replaced actor", state: ateapipb.ActorState_ACTOR_STATE_CRASHED, revertUID: "another-actor", wantApplyErr: true, wantCalls: []string{"revert"}}, + {name: "different snapshot", state: ateapipb.ActorState_ACTOR_STATE_CRASHED, revertSnapshot: durable(1, "s3://snapshots/older"), wantApplyErr: true, wantCalls: []string{"revert"}}, + } { + t.Run(test.name, func(t *testing.T) { + actors := &transitionActors{revertUID: test.revertUID, revertSnapshot: test.revertSnapshot} + if !test.missing { + actors.actor = &ateapipb.Actor{ + Metadata: &ateapipb.ResourceMetadata{Atespace: "team", Name: "instance", Uid: "actor-uid"}, + ActorTemplate: &ateapipb.ObjectRef{Atespace: "team", Name: "revision"}, + Status: &ateapipb.ActorStatus{State: test.state, Snapshots: []*ateapipb.Snapshot{durable(2, "s3://snapshots/latest")}}, + } + if test.state == ateapipb.ActorState_ACTOR_STATE_SUSPENDED { + actors.actor.Status.Snapshots[0].Locality = "" + } + } + binding := ActorBinding{Revision: "revision", Atespace: "team", Name: "instance", TemplateAtespace: "team", TemplateName: "revision"} + transition, err := PrepareActorRevert(t.Context(), actors, binding) + if test.wantPrepareErr { + require.Error(t, err) + require.Empty(t, actors.calls, "a failed read issues nothing") + return + } + require.NoError(t, err) + require.Equal(t, "actor-uid", transition.ActorUID()) + err = ApplyActorTransition(t.Context(), actors, transition) + if test.wantApplyErr { + require.Error(t, err) + } else { + require.NoError(t, err) + } + require.Equal(t, test.wantCalls, actors.calls) + }) + } +} From 17022a21ae3574415ab3ea9e9a332a76f82684d2 Mon Sep 17 00:00:00 2001 From: Ada Lundhe Date: Tue, 3 Nov 2026 11:23:31 -0800 Subject: [PATCH 06/12] fix(session): refuse unrecoverable reverts as FAILED_PRECONDITION, not UNAVAILABLE RevertSession reported a missing or same-name-replaced Actor as `UNAVAILABLE`. The retry guide tells clients to retry that code, so a client that followed it retried a runtime no revert can restore until its deadline. These refusals now return `FAILED_PRECONDITION`. Transient Substrate and database failures still return `UNAVAILABLE`. `FAILED_PRECONDITION` then meant two opposite things: the runtime is already healthy, which includes a retry after a lost success response, or it can never be recovered. A `google.rpc.ErrorInfo` detail in domain `kagent.dev` now names `KAGENT_SESSION_RUNTIME_NOT_CRASHED` or `KAGENT_SESSION_RUNTIME_UNRECOVERABLE`, and the retry guide gives the client action for each. Telling clients to read `GetSession` lost: a Session whose Actor was replaced is also READY with no failure, so the Session alone cannot separate the two. A revert of a Session at rest outside READY returned `ABORTED`. A Session suspended by deadline cleanup stays suspended until an interaction resumes it, so the retry `ABORTED` invites never converged. Only an active lifecycle operation is now contention; a resting Session returns `KAGENT_SESSION_RUNTIME_NOT_CRASHED`. The admitted-recovery retry path is unchanged: it enters through the FAILED/RESUME branch, which none of these checks reach. Verified: the session and grpcserver suites pass against PostgreSQL, apart from TestSandboxTemplateCatalog, which needs envtest binaries. That includes TestRevertSessionThroughGRPC, which now pins the code and reason for a healthy, replaced, missing and already-recovered runtime. golangci-lint reports 0 issues and gofmt is clean. On a kind cluster with gVisor workers, the worker-loss revert E2E passes for all four harnesses. On the same cluster, a revert after the Actor was deleted, or replaced under the same name, returned `UNAVAILABLE` before this change and now returns `FAILED_PRECONDITION` with `KAGENT_SESSION_RUNTIME_UNRECOVERABLE`. A revert retried after recovery returns `KAGENT_SESSION_RUNTIME_NOT_CRASHED`. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Ada Lundhe --- docs/architecture/runtime-and-lifecycle.md | 10 +++- docs/lifecycle-retries.md | 13 +++- .../grpcserver/session_recovery_test.go | 27 ++++++++- go/core/internal/service/session/recovery.go | 36 ++++++++++- .../internal/service/session/recovery_test.go | 60 +++++++++++++++---- go/core/internal/service/session/service.go | 26 ++++++++ .../internal/service/session/service_test.go | 25 ++++++++ go/core/test/e2e/session_revert_test.go | 22 ++++++- 8 files changed, 198 insertions(+), 21 deletions(-) diff --git a/docs/architecture/runtime-and-lifecycle.md b/docs/architecture/runtime-and-lifecycle.md index 32b7e651cb..c37095eb0e 100644 --- a/docs/architecture/runtime-and-lifecycle.md +++ b/docs/architecture/runtime-and-lifecycle.md @@ -209,7 +209,13 @@ adopted implicitly. crashed runtime. It accepts active and idle Sessions, but verifies Substrate's CRASHED state for the recorded Actor UID and pinned template before admitting recovery. Quiet turns, transport failures, missing Actors and same-name -replacements cannot authorize a revert. +replacements cannot authorize a revert. A quiet turn returns `FAILED_PRECONDITION` +with reason `KAGENT_SESSION_RUNTIME_NOT_CRASHED`. A missing or replaced Actor +returns `FAILED_PRECONDITION` with `KAGENT_SESSION_RUNTIME_UNRECOVERABLE`, +because no retry restores it. A transport failure returns `UNAVAILABLE`. +Only an active lifecycle operation returns `ABORTED`. A Session at rest outside +READY returns `KAGENT_SESSION_RUNTIME_NOT_CRASHED`, so a client does not retry +it until its deadline. Admission records FAILED/RESUME with a RuntimeCrashed failure under the Session lock. The same transaction publishes a durable terminal result saved before the @@ -237,7 +243,7 @@ A successful revert returns the Session to READY with the same Actor suspended. The next authorized interaction resumes it through Substrate ingress. Runtime changes after the retained snapshot are lost; the discarded turn receives no checkpoint at that older snapshot. Another revert against the now-healthy Actor -returns `FAILED_PRECONDITION`. +returns `FAILED_PRECONDITION` with `KAGENT_SESSION_RUNTIME_NOT_CRASHED`. ```mermaid sequenceDiagram diff --git a/docs/lifecycle-retries.md b/docs/lifecycle-retries.md index 78b61a129b..fb59f20355 100644 --- a/docs/lifecycle-retries.md +++ b/docs/lifecycle-retries.md @@ -50,9 +50,20 @@ A runtime without a snapshot restarts from its initial state. Retry an interrupted revert against the same Session ID. Pending recovery retains its operation across API restarts. The lifecycle worker retries already admitted recoveries; it does not initiate recovery for other crashed Sessions. -After recovery succeeds, another revert returns `FAILED_PRECONDITION` because the Actor is healthy. Checkpoint creation returns `ABORTED` until it finishes. A revert that observed the crash before another lifecycle operation began also returns `ABORTED`. +So does a revert while another lifecycle operation is active. + +A revert that can never succeed against the Session as it stands returns `FAILED_PRECONDITION`. +Its `google.rpc.ErrorInfo` detail, in domain `kagent.dev`, names one of two reasons: + +| Reason | Meaning | Client action | +| --- | --- | --- | +| `KAGENT_SESSION_RUNTIME_NOT_CRASHED` | The runtime is healthy or suspended. A retry after a revert whose response was lost lands here. | Treat as recovered and continue the conversation. | +| `KAGENT_SESSION_RUNTIME_UNRECOVERABLE` | The recorded Actor is missing, or a same-name Actor replaced it. | Stop retrying. Start another Session, or fork one from a checkpoint. | + +A `FAILED_PRECONDITION` without either reason follows the general table below. +Read the reason rather than the message text, which is not stable. ## Errors diff --git a/go/core/internal/grpcserver/session_recovery_test.go b/go/core/internal/grpcserver/session_recovery_test.go index 7515035aab..ec4138ad71 100644 --- a/go/core/internal/grpcserver/session_recovery_test.go +++ b/go/core/internal/grpcserver/session_recovery_test.go @@ -12,6 +12,7 @@ import ( "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" "google.golang.org/grpc/codes" "google.golang.org/grpc/credentials/insecure" @@ -143,16 +144,16 @@ func TestRevertSessionThroughGRPC(t *testing.T) { _, err = client.RevertSession(visitor, request) require.Equal(t, codes.NotFound, status.Code(err)) _, err = client.RevertSession(owner, request) - require.Equal(t, codes.FailedPrecondition, status.Code(err), "healthy runtimes cannot revert") + requireRevertRefusal(t, err, "KAGENT_SESSION_RUNTIME_NOT_CRASHED", "healthy runtimes cannot revert") control.actors[name].Status.State = ateapipb.ActorState_ACTOR_STATE_CRASHED before := proto.CloneOf(control.actors[name]) control.actors[name].Metadata.Uid = "replacement-uid" _, err = client.RevertSession(owner, request) - require.Equal(t, codes.Unavailable, status.Code(err), "a same-name replacement is not this runtime") + requireRevertRefusal(t, err, "KAGENT_SESSION_RUNTIME_UNRECOVERABLE", "a same-name replacement is not this runtime") control.actors[name].Metadata.Uid = before.Metadata.Uid delete(control.actors, name) _, err = client.RevertSession(owner, request) - require.Equal(t, codes.Unavailable, status.Code(err), "a missing runtime cannot authorize recovery") + requireRevertRefusal(t, err, "KAGENT_SESSION_RUNTIME_UNRECOVERABLE", "a missing runtime cannot authorize recovery") control.actors[name] = proto.CloneOf(before) calls := 0 failOnce := true @@ -176,6 +177,10 @@ func TestRevertSessionThroughGRPC(t *testing.T) { require.NoError(t, err) require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, ready.Session.State) require.Nil(t, ready.Session.Failure) + // A client whose success response was lost retries and learns that + // the runtime is already healthy, not that it cannot be recovered. + _, err = client.RevertSession(owner, request) + requireRevertRefusal(t, err, "KAGENT_SESSION_RUNTIME_NOT_CRASHED", "a recovered runtime has nothing to revert") if test.lostResponse { require.Equal(t, 1, calls, "a lost response cannot repeat the revert") } else { @@ -198,3 +203,19 @@ func TestRevertSessionThroughGRPC(t *testing.T) { }) } } + +// requireRevertRefusal pins both the status code and the ErrorInfo reason, +// because the two FAILED_PRECONDITION refusals ask opposite things of a client. +func requireRevertRefusal(t *testing.T, err error, reason, message string) { + t.Helper() + refusal, ok := status.FromError(err) + require.True(t, ok, message) + require.Equal(t, codes.FailedPrecondition, refusal.Code(), message) + var reasons []string + for _, detail := range refusal.Details() { + if info, ok := detail.(*errdetails.ErrorInfo); ok && info.GetDomain() == "kagent.dev" { + reasons = append(reasons, info.GetReason()) + } + } + require.Equal(t, []string{reason}, reasons, message) +} diff --git a/go/core/internal/service/session/recovery.go b/go/core/internal/service/session/recovery.go index 27922bc3ed..ceab783ccf 100644 --- a/go/core/internal/service/session/recovery.go +++ b/go/core/internal/service/session/recovery.go @@ -8,6 +8,8 @@ import ( "github.com/google/uuid" 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" @@ -16,6 +18,19 @@ import ( "github.com/kagent-dev/kagent/go/pkg/logging" ) +// Revert refuses a Session with nothing to act on for one of two reasons. They +// ask opposite things of a client, so the API reports them with distinct reasons. +// errRuntimeNotCrashed means nothing needs recovering: the runtime is healthy or +// suspended. A retry after a revert whose response was lost lands here, and the +// next interaction proceeds. +// errRuntimeUnrecoverable means the recorded Actor is gone or was replaced. No +// retry restores it, and the client has to start another Session. Transient +// Substrate and database failures wrap neither, so they stay retryable. +var ( + errRuntimeNotCrashed = fmt.Errorf("runtime is not crashed: %w", database.ErrFailedPrecondition) + errRuntimeUnrecoverable = fmt.Errorf("runtime cannot be recovered: %w", database.ErrFailedPrecondition) +) + // recoverSessionPage retries admitted recoveries in ID order. // READY Sessions are excluded, so a background retry cannot admit new work. func (w *ActorWorkflow) recoverSessionPage(ctx context.Context, afterID string) (string, error) { @@ -58,9 +73,16 @@ func (w *ActorWorkflow) Revert(ctx context.Context, session *apiv1.Session) (*ap if session.State == apiv1.RuntimeState_RUNTIME_STATE_FAILED && session.Failure.GetReason() == database.SessionRuntimeCrashed { return w.run(ctx, session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME) } - if session.State != apiv1.RuntimeState_RUNTIME_STATE_READY || session.Operation != apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE { + // Only an active lifecycle operation is contention: a retry converges once it + // finishes. A Session at rest in another state, such as one suspended by + // deadline cleanup, stays there until an interaction resumes it, so a retry + // would spin until the caller's deadline. + if session.Operation != apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE { return nil, database.ErrConflict } + if session.State != apiv1.RuntimeState_RUNTIME_STATE_READY { + return nil, fmt.Errorf("runtime state is %s: %w", session.State, errRuntimeNotCrashed) + } revision, err := w.store.GetRuntimeRevision(ctx, session.PreparedRevision) if err != nil { return nil, err @@ -73,14 +95,22 @@ func (w *ActorWorkflow) Revert(ctx context.Context, session *apiv1.Session) (*ap return nil, err } actor, err := w.actors.GetActor(ctx, revision.ActorTemplateAtespace, substrate.ActorName(session.Id)) + if status.Code(err) == codes.NotFound { + return nil, fmt.Errorf("runtime Actor is missing: %w", errRuntimeUnrecoverable) + } if err != nil { return nil, err } - if err := w.verifyActor(ctx, session, revision, actor); err != nil { + if !validActorIdentity(actor, revision, substrate.ActorName(session.Id)) { + return nil, fmt.Errorf("runtime Actor identity or template changed: %w", errRuntimeUnrecoverable) + } + if err := w.verifyActor(ctx, session, revision, actor); errors.Is(err, database.ErrNotFound) { + return nil, fmt.Errorf("runtime Actor was replaced: %w", errRuntimeUnrecoverable) + } else if err != nil { return nil, err } if actor.GetStatus().GetState() != ateapipb.ActorState_ACTOR_STATE_CRASHED { - return nil, fmt.Errorf("runtime is not crashed: %w", database.ErrFailedPrecondition) + return nil, errRuntimeNotCrashed } snapshot := substrate.LatestDurableActorSnapshot(actor) operation, err := w.store.BeginCrashedSessionRecovery(ctx, session.Id, generation, actor.GetMetadata().GetUid(), database.SessionTaskSnapshot{ diff --git a/go/core/internal/service/session/recovery_test.go b/go/core/internal/service/session/recovery_test.go index d7ce338a83..3918fc1ae0 100644 --- a/go/core/internal/service/session/recovery_test.go +++ b/go/core/internal/service/session/recovery_test.go @@ -190,15 +190,20 @@ func TestCrashedSessionRecoveryFencesConcurrentWork(t *testing.T) { func TestSessionRecoveryRequiresConfirmedRecordedActorCrash(t *testing.T) { for _, test := range []struct { - name string - state ateapipb.ActorState - uid string - readErr error - wantErr bool + name string + state ateapipb.ActorState + uid string + template string + readErr error + // want is the refusal a client sees; nil means the error must stay + // retryable, so it may match neither refusal. + want error }{ - {name: "quiet running actor", state: ateapipb.ActorState_ACTOR_STATE_RUNNING, wantErr: true}, - {name: "transient read failure", readErr: status.Error(codes.Unavailable, "unreachable"), wantErr: true}, - {name: "same-name replacement", state: ateapipb.ActorState_ACTOR_STATE_CRASHED, uid: "another-actor", wantErr: true}, + {name: "quiet running actor", state: ateapipb.ActorState_ACTOR_STATE_RUNNING, want: errRuntimeNotCrashed}, + {name: "transient read failure", readErr: status.Error(codes.Unavailable, "unreachable")}, + {name: "missing actor", readErr: status.Error(codes.NotFound, "actor not found"), want: errRuntimeUnrecoverable}, + {name: "same-name replacement", state: ateapipb.ActorState_ACTOR_STATE_CRASHED, uid: "another-actor", want: errRuntimeUnrecoverable}, + {name: "template changed", state: ateapipb.ActorState_ACTOR_STATE_CRASHED, template: "another-template", want: errRuntimeUnrecoverable}, } { t.Run(test.name, func(t *testing.T) { store, base, session, task, _ := crashedTurnFixture(t) @@ -207,16 +212,21 @@ func TestSessionRecoveryRequiresConfirmedRecordedActorCrash(t *testing.T) { if test.uid != "" { actor.Metadata.Uid = test.uid } + if test.template != "" { + actor.ActorTemplate.Name = test.template + } actors := &recoveryTestActors{lifecycleTestActors: base, readErr: test.readErr} actors.revert = func(context.Context, string, string) (*ateapipb.Actor, error) { t.Fatal("unconfirmed crash issued a revert") return nil, nil } _, err := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Second).Revert(t.Context(), session) - if test.wantErr { - require.Error(t, err) + require.Error(t, err) + if test.want != nil { + require.ErrorIs(t, err, test.want) + require.ErrorIs(t, err, database.ErrFailedPrecondition) } else { - require.NoError(t, err) + require.NotErrorIs(t, err, database.ErrFailedPrecondition, "a transient failure must stay retryable") } current, err := store.GetSessionByID(t.Context(), session.Id) require.NoError(t, err) @@ -229,6 +239,34 @@ func TestSessionRecoveryRequiresConfirmedRecordedActorCrash(t *testing.T) { } } +// Only an active lifecycle operation is contention a retry waits out. A Session +// at rest outside READY, such as one suspended by deadline cleanup, would +// return ABORTED forever if it were reported as contention. +func TestExplicitRevertSeparatesContentionFromRestingState(t *testing.T) { + for _, test := range []struct { + name string + state apiv1.RuntimeState + operation apiv1.RuntimeOperation + want error + }{ + {name: "suspended by deadline cleanup", state: apiv1.RuntimeState_RUNTIME_STATE_SUSPENDED, operation: apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE, want: errRuntimeNotCrashed}, + {name: "suspension in progress", state: apiv1.RuntimeState_RUNTIME_STATE_READY, operation: apiv1.RuntimeOperation_RUNTIME_OPERATION_SUSPEND, want: database.ErrConflict}, + {name: "resume in progress", state: apiv1.RuntimeState_RUNTIME_STATE_SUSPENDED, operation: apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME, want: database.ErrConflict}, + {name: "creation in progress", state: apiv1.RuntimeState_RUNTIME_STATE_CREATING, operation: apiv1.RuntimeOperation_RUNTIME_OPERATION_CREATE, want: database.ErrConflict}, + } { + t.Run(test.name, func(t *testing.T) { + // Both checks run on the Session the caller read, before any store + // or Substrate access, so neither dependency is needed. + workflow := NewActorWorkflow(nil, nil, make(chan struct{}, 1), time.Second) + _, err := workflow.Revert(t.Context(), &apiv1.Session{Id: uuid.NewString(), State: test.state, Operation: test.operation}) + require.ErrorIs(t, err, test.want) + if test.want == database.ErrConflict { + require.NotErrorIs(t, err, database.ErrFailedPrecondition) + } + }) + } +} + func TestSessionRecoveryCandidatePolicy(t *testing.T) { store, actors, active, _, _ := crashedTurnFixture(t) workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Second) diff --git a/go/core/internal/service/session/service.go b/go/core/internal/service/session/service.go index bc2d22d32f..4de99281cd 100644 --- a/go/core/internal/service/session/service.go +++ b/go/core/internal/service/session/service.go @@ -11,6 +11,9 @@ import ( "time" "github.com/google/uuid" + "google.golang.org/genproto/googleapis/rpc/errdetails" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" "google.golang.org/protobuf/types/known/timestamppb" utilvalidation "k8s.io/apimachinery/pkg/util/validation" @@ -288,6 +291,16 @@ func (s *Service) Revert(ctx context.Context, id string) (*apiv1.Session, error) return nil, err } session, err = s.workflow.Revert(ctx, session) + if reason := revertRefusalReason(err); reason != "" { + // Both refusals are FAILED_PRECONDITION, but one means the runtime is + // already healthy and the other that no retry can restore it. The + // ErrorInfo reason lets a client tell them apart without parsing text. + refusal, detailErr := status.New(codes.FailedPrecondition, err.Error()).WithDetails(&errdetails.ErrorInfo{Domain: "kagent.dev", Reason: reason}) + if detailErr != nil { + return nil, detailErr + } + return nil, refusal.Err() + } if errors.Is(err, database.ErrFailedPrecondition) { return nil, serviceerrors.NewFailedPrecondition(err.Error(), err) } @@ -300,6 +313,19 @@ func (s *Service) Revert(ctx context.Context, id string) (*apiv1.Session, error) return session, nil } +// revertRefusalReason names the ErrorInfo reason for a revert that can never +// succeed against the Session as it stands, or returns "" for any other error. +func revertRefusalReason(err error) string { + switch { + case errors.Is(err, errRuntimeNotCrashed): + return "KAGENT_SESSION_RUNTIME_NOT_CRASHED" + case errors.Is(err, errRuntimeUnrecoverable): + return "KAGENT_SESSION_RUNTIME_UNRECOVERABLE" + default: + return "" + } +} + // CreateShare mints a read-only share of the caller's session. A positive ttl bounds how long its // token grants access and must not exceed the configured maximum; zero takes the maximum, // or leaves the share valid until it is revoked or the session deleted when none is set. diff --git a/go/core/internal/service/session/service_test.go b/go/core/internal/service/session/service_test.go index a63207de0b..124aaef5e2 100644 --- a/go/core/internal/service/session/service_test.go +++ b/go/core/internal/service/session/service_test.go @@ -13,6 +13,9 @@ import ( "github.com/google/uuid" "github.com/stretchr/testify/require" + "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/database" @@ -205,6 +208,28 @@ func TestServiceCreateRejectsInvalidOrUnauthorizedRequests(t *testing.T) { } } +func TestServiceRevertNamesRefusalReason(t *testing.T) { + for _, test := range []struct { + cause error + reason string + }{ + {errRuntimeNotCrashed, "KAGENT_SESSION_RUNTIME_NOT_CRASHED"}, + {fmt.Errorf("runtime Actor was replaced: %w", errRuntimeUnrecoverable), "KAGENT_SESSION_RUNTIME_UNRECOVERABLE"}, + } { + service := NewService(&serviceTestStore{}, serviceTestAuthorizer{}, serviceTestWorkflow{err: test.cause}) + _, err := service.Revert(serviceTestContext("alice"), "8bd650a8-9775-488f-8bc1-0d52bf7bdcab") + refusal, ok := status.FromError(err) + require.True(t, ok) + require.Equal(t, codes.FailedPrecondition, refusal.Code()) + require.Equal(t, test.cause.Error(), refusal.Message()) + require.Len(t, refusal.Details(), 1) + info, ok := refusal.Details()[0].(*errdetails.ErrorInfo) + require.True(t, ok) + require.Equal(t, "kagent.dev", info.GetDomain()) + require.Equal(t, test.reason, info.GetReason()) + } +} + func TestServiceLifecycleMethodsPreserveContentionErrors(t *testing.T) { for _, test := range []struct { name string diff --git a/go/core/test/e2e/session_revert_test.go b/go/core/test/e2e/session_revert_test.go index 83a655b33f..2e27cd322a 100644 --- a/go/core/test/e2e/session_revert_test.go +++ b/go/core/test/e2e/session_revert_test.go @@ -16,6 +16,7 @@ import ( "github.com/kagent-dev/mockllm" "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" @@ -91,7 +92,7 @@ func TestSessionRevertCrashedRuntime(t *testing.T) { require.NotEmpty(t, substrate.SnapshotURI(substrate.LatestDurableActorSnapshot(running))) request := &apiv1.RevertSessionRequest{SessionId: fixture.sessionID} _, err = fixture.sessions.RevertSession(fixture.ctx, request) - require.Equal(t, codes.FailedPrecondition, status.Code(err), "healthy runtimes cannot revert") + require.Equal(t, "KAGENT_SESSION_RUNTIME_NOT_CRASHED", revertRefusalReason(t, err), "healthy runtimes cannot revert") assignment := running.GetStatus().GetWorkerAssignment() require.Equal(t, namespace, assignment.GetWorkerNamespace()) require.NotEmpty(t, assignment.GetWorkerPod()) @@ -130,6 +131,10 @@ func TestSessionRevertCrashedRuntime(t *testing.T) { reverted, err := fixture.sessions.RevertSession(fixture.ctx, request) require.NoError(t, err) require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, reverted.Session.State) + // A client whose success response was lost retries and must learn that + // the runtime is healthy, not that it cannot be recovered. + _, err = fixture.sessions.RevertSession(fixture.ctx, request) + require.Equal(t, "KAGENT_SESSION_RUNTIME_NOT_CRASHED", revertRefusalReason(t, err), "a recovered runtime has nothing to revert") recovered, err := findSubstrateActor(fixture.ctx, fixture.system, namespace, name) require.NoError(t, err) require.Equal(t, running.GetMetadata().GetUid(), recovered.GetMetadata().GetUid()) @@ -149,3 +154,18 @@ func TestSessionRevertCrashedRuntime(t *testing.T) { require.Contains(t, taskText(next), "4") }) } + +// revertRefusalReason returns the kagent.dev ErrorInfo reason of a +// FAILED_PRECONDITION revert refusal. +func revertRefusalReason(t *testing.T, err error) string { + t.Helper() + refusal, ok := status.FromError(err) + require.True(t, ok, "revert error %v is not a gRPC status", err) + require.Equal(t, codes.FailedPrecondition, refusal.Code(), "revert error: %v", err) + for _, detail := range refusal.Details() { + if info, ok := detail.(*errdetails.ErrorInfo); ok && info.GetDomain() == "kagent.dev" { + return info.GetReason() + } + } + return "" +} From 04b839bb01f197054111132b624065af02b33d5a Mon Sep 17 00:00:00 2001 From: Ada Lundhe Date: Wed, 4 Nov 2026 08:39:22 -0800 Subject: [PATCH 07/12] test(session): use SnapshotStorage state in recovery tests Main renamed SnapshotStorage.status to state (SnapshotStorageState). Update the recovery and substrate snapshot tests to the new field. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Ada Lundhe --- go/core/internal/service/session/recovery_test.go | 6 +++--- go/core/internal/substrate/lifecycle_test.go | 2 +- go/core/internal/substrate/snapshot_test.go | 6 +++--- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/go/core/internal/service/session/recovery_test.go b/go/core/internal/service/session/recovery_test.go index 3769f82e22..842c4a3b98 100644 --- a/go/core/internal/service/session/recovery_test.go +++ b/go/core/internal/service/session/recovery_test.go @@ -434,11 +434,11 @@ func TestExplicitRevertPreservesOnlyCapturedWaits(t *testing.T) { require.NoError(t, store.FinishSessionQuiescence(t.Context(), work, &database.SessionTaskSnapshot{ Atespace: "team-a", URI: "s3://snapshots/waiting", Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES.String(), })) - actors.actors[actorKey("team-a", substrate.ActorName(session.Id))].Status.Snapshots = []*ateapipb.Snapshot{{Generation: 1, Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES, DurableSnapshot: &ateapipb.SnapshotStorage{Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_COMPLETED, Object: &ateapipb.ObjectSnapshot{SnapshotUri: "s3://snapshots/waiting"}}}} + actors.actors[actorKey("team-a", substrate.ActorName(session.Id))].Status.Snapshots = []*ateapipb.Snapshot{{Generation: 1, Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES, DurableSnapshot: &ateapipb.SnapshotStorage{State: ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_COMPLETED, Object: &ateapipb.ObjectSnapshot{SnapshotUri: "s3://snapshots/waiting"}}}} actor := actors.actors[actorKey("team-a", substrate.ActorName(session.Id))] actor.Status.Snapshots[0].Locality = "lost-node" actor.Status.Snapshots = append(actor.Status.Snapshots, &ateapipb.Snapshot{ - Generation: 2, DurableSnapshot: &ateapipb.SnapshotStorage{Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_PENDING}, + Generation: 2, DurableSnapshot: &ateapipb.SnapshotStorage{State: ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_PENDING}, }) } _, err = NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute).Revert(t.Context(), session) @@ -475,7 +475,7 @@ func TestExplicitRevertRetainsWaitCapturedByLaterTurn(t *testing.T) { require.NoError(t, store.FinishSessionQuiescence(t.Context(), work, &database.SessionTaskSnapshot{ Atespace: "team-a", URI: "s3://snapshots/later-turn", Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES.String(), })) - actors.actors[actorKey("team-a", substrate.ActorName(session.Id))].Status.Snapshots = []*ateapipb.Snapshot{{Generation: 1, Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES, DurableSnapshot: &ateapipb.SnapshotStorage{Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_COMPLETED, Object: &ateapipb.ObjectSnapshot{SnapshotUri: "s3://snapshots/later-turn"}}}} + actors.actors[actorKey("team-a", substrate.ActorName(session.Id))].Status.Snapshots = []*ateapipb.Snapshot{{Generation: 1, Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES, DurableSnapshot: &ateapipb.SnapshotStorage{State: ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_COMPLETED, Object: &ateapipb.ObjectSnapshot{SnapshotUri: "s3://snapshots/later-turn"}}}} _, err = NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute).Revert(t.Context(), session) require.NoError(t, err) retained, err := store.GetSettledSessionTask(t.Context(), session.Id, string(waiting.ID), nil) diff --git a/go/core/internal/substrate/lifecycle_test.go b/go/core/internal/substrate/lifecycle_test.go index 70efdc88cc..21e9eaa3dd 100644 --- a/go/core/internal/substrate/lifecycle_test.go +++ b/go/core/internal/substrate/lifecycle_test.go @@ -256,7 +256,7 @@ func TestActorSuspendRequiresDurableVolumesSnapshot(t *testing.T) { func TestActorRevertRestoresTheObservedSnapshotInPlace(t *testing.T) { durable := func(generation int32, uri string) *ateapipb.Snapshot { return &ateapipb.Snapshot{Generation: generation, Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES, Locality: "lost-node", - DurableSnapshot: &ateapipb.SnapshotStorage{Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_COMPLETED, Object: &ateapipb.ObjectSnapshot{SnapshotUri: uri}}} + DurableSnapshot: &ateapipb.SnapshotStorage{State: ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_COMPLETED, Object: &ateapipb.ObjectSnapshot{SnapshotUri: uri}}} } for _, test := range []struct { name string diff --git a/go/core/internal/substrate/snapshot_test.go b/go/core/internal/substrate/snapshot_test.go index dfda60e0cf..4e923131a8 100644 --- a/go/core/internal/substrate/snapshot_test.go +++ b/go/core/internal/substrate/snapshot_test.go @@ -43,13 +43,13 @@ func TestActorSnapshotGenerationAndDurability(t *testing.T) { func TestRevertSnapshotSkipsIncompleteGenerations(t *testing.T) { durable := &ateapipb.Snapshot{Generation: 2, DurableSnapshot: &ateapipb.SnapshotStorage{ - Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_COMPLETED, + State: ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_COMPLETED, Object: &ateapipb.ObjectSnapshot{SnapshotUri: "s3://snapshots/durable"}, }} for _, latest := range []*ateapipb.Snapshot{ {Generation: 3, Locality: "node-1"}, - {Generation: 3, DurableSnapshot: &ateapipb.SnapshotStorage{Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_PENDING}}, - {Generation: 3, DurableSnapshot: &ateapipb.SnapshotStorage{Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_COMPLETED}}, + {Generation: 3, DurableSnapshot: &ateapipb.SnapshotStorage{State: ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_PENDING}}, + {Generation: 3, DurableSnapshot: &ateapipb.SnapshotStorage{State: ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_COMPLETED}}, } { actor := &ateapipb.Actor{Status: &ateapipb.ActorStatus{Snapshots: []*ateapipb.Snapshot{latest, durable}}} require.Same(t, durable, LatestDurableActorSnapshot(actor)) From 683c3462d2ef66c75da677e4d6fbe88996b4a90d Mon Sep 17 00:00:00 2001 From: Ada Lundhe Date: Fri, 23 Oct 2026 12:31:02 -0700 Subject: [PATCH 08/12] fix(session): add durable recovery primitives Signed-off-by: Ada Lundhe --- go/core/internal/database/session_recovery.go | 197 ++++++++++++++ .../database/session_recovery_test.go | 255 ++++++++++++++++++ .../internal/database/task_finalization.go | 96 +++---- go/core/internal/substrate/client.go | 10 + go/core/internal/substrate/snapshot.go | 11 + go/core/internal/substrate/snapshot_test.go | 20 ++ 6 files changed, 544 insertions(+), 45 deletions(-) create mode 100644 go/core/internal/database/session_recovery.go create mode 100644 go/core/internal/database/session_recovery_test.go diff --git a/go/core/internal/database/session_recovery.go b/go/core/internal/database/session_recovery.go new file mode 100644 index 0000000000..a803907ab9 --- /dev/null +++ b/go/core/internal/database/session_recovery.go @@ -0,0 +1,197 @@ +package database + +import ( + "context" + "fmt" + "time" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "google.golang.org/protobuf/types/known/timestamppb" + + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" +) + +// SessionRuntimeCrashed identifies a recovery admitted after Substrate confirmed +// that the recorded Actor crashed. Transport failures cannot admit this work. +const SessionRuntimeCrashed = "RuntimeCrashed" + +// ListSessionRecoveryCandidates finds recoveries admitted by an explicit request. +// Retrying durable work cannot admit recovery for a READY Session. +func (c *Client) ListSessionRecoveryCandidates(ctx context.Context, afterID string, limit int) ([]*apiv1.Session, error) { + rows, err := queryMany(ctx, c.db, ` + SELECT id, user_id, prepared_revision, state, data, operation, context_id, + source_checkpoint_id, history_id, operation_id, executor_id + FROM session_record i + WHERE ($1::text = '' OR i.id > $1::uuid) + AND i.state = 'RUNTIME_STATE_FAILED' AND i.operation = 'RUNTIME_OPERATION_RESUME' + ORDER BY i.id LIMIT $2 + `, pgx.RowToStructByName[sessionRow], afterID, limit) + if err != nil { + return nil, err + } + result := make([]*apiv1.Session, 0, len(rows)) + for _, row := range rows { + session, err := toSession(row) + if err != nil { + return nil, err + } + result = append(result, session) + } + return result, nil +} + +// SessionLifecycleGeneration returns the ID of the Session's latest lifecycle +// operation, or uuid.Nil when none is recorded. Every operation replaces it, and +// completion keeps it. Recovery reads it before observing Substrate. +func (c *Client) SessionLifecycleGeneration(ctx context.Context, id string) (uuid.UUID, error) { + generation, err := queryOne(ctx, c.db, `SELECT operation_id FROM session_record WHERE id = $1`, pgx.RowTo[*uuid.UUID], id) + if err != nil { + return uuid.Nil, notFoundOr(err) + } + if generation == nil { + return uuid.Nil, nil + } + return *generation, nil +} + +// BeginCrashedSessionRecovery fences confirmed-crash recovery under the Session lock. +// It preserves durable results; obsolete Actors and conflicting operations fail. +// generation is the caller's SessionLifecycleGeneration read from before it +// observed the crash. A lifecycle operation begun since then makes the observation +// stale, because a concurrent revert keeps the Actor UID and returns to READY. +func (c *Client) BeginCrashedSessionRecovery(ctx context.Context, id string, generation uuid.UUID, actorUID string, snapshot SessionTaskSnapshot) (*SessionOperation, error) { + var operation *SessionOperation + err := c.withTx(ctx, func(tx pgx.Tx) error { + row, err := lockSession(ctx, tx, id) + if err != nil { + return notFoundOr(err) + } + matches, err := queryOne(ctx, tx, `SELECT actor_uid = $2 FROM session WHERE id = $1`, pgx.RowTo[bool], row.ID, actorUID) + if err != nil { + return err + } + if !matches || actorUID == "" { + return ErrConflict + } + session, err := toSession(row) + if err != nil { + return err + } + if session.State == apiv1.RuntimeState_RUNTIME_STATE_FAILED && session.Operation == apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME && session.Failure.GetReason() == SessionRuntimeCrashed { + operation, err = toSessionOperation(row) + return err + } + if session.State != apiv1.RuntimeState_RUNTIME_STATE_READY || session.Operation != apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE { + return ErrConflict + } + current := uuid.Nil + if row.OperationID != nil { + current = *row.OperationID + } + if current != generation { + return ErrConflict + } + // An issued idle suspension does not exclude recovery here. It cannot finish + // against the crashed Actor, and its worker keeps the claim when suspension + // fails, so waiting for it would block recovery permanently. Settlement below + // retires the claim in this transaction; a late finish for it changes nothing. + // The generation check keeps the crash observation current: without another + // lifecycle operation, nothing returns the crashed Actor to service. + busy, err := queryOne(ctx, tx, ` + SELECT EXISTS (SELECT 1 FROM session_checkpoint WHERE source_session_id = $1 AND state = 'CREATING') + `, pgx.RowTo[bool], row.ID) + if err != nil { + return err + } + if busy { + return ErrConflict + } + row.State = apiv1.RuntimeState_RUNTIME_STATE_FAILED.String() + row.runtimeInstanceRow, err = beginRuntimeOperation(ctx, tx, row.runtimeInstanceRow, runtimeKindAgent, apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME) + if err != nil { + return err + } + session.State, session.Operation = apiv1.RuntimeState_RUNTIME_STATE_FAILED, apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME + session.Failure = &apiv1.Failure{Reason: SessionRuntimeCrashed, Message: "Runtime crashed; recovery to its last snapshot is pending"} + session.UpdatedAt = timestamppb.Now() + row.Data, err = marshalSession(session) + if err != nil { + return err + } + if err := execSQL(ctx, tx, `UPDATE session SET data = $2, dispatch_id = NULL, dispatch_expires_at = NULL WHERE id = $1`, row.ID, row.Data); err != nil { + return err + } + if err := settleCrashedSessionTask(ctx, tx, row, snapshot); err != nil { + return err + } + operation, err = toSessionOperation(row) + return err + }) + return operation, err +} + +// settleCrashedSessionTask retains captured waits and fails discarded nonterminal work. +// Callers hold the Session lock; recovery cannot checkpoint discarded state. +func settleCrashedSessionTask(ctx context.Context, tx pgx.Tx, session sessionRow, snapshot SessionTaskSnapshot) error { + if err := execSQL(ctx, tx, `UPDATE session_task_event SET quiescence_pending = FALSE WHERE history_id = $1 AND quiescence_pending`, session.HistoryID); err != nil { + return err + } + taskIDs, err := queryMany(ctx, tx, ` + SELECT t.id FROM session_task t WHERE t.history_id = $1 + AND (t.state IN ('TASK_STATE_SUBMITTED', 'TASK_STATE_WORKING') + OR (t.state IN ('TASK_STATE_INPUT_REQUIRED', 'TASK_STATE_AUTH_REQUIRED') + AND NOT EXISTS (SELECT 1 FROM session_task_event captured + WHERE captured.history_id = t.history_id AND $2::text <> '' + AND captured.snapshot_uri = $2 AND captured.snapshot_fidelity = $3 + AND captured.snapshot_atespace = $4 + AND captured.sequence >= (SELECT MAX(sequence) FROM session_task_event latest + WHERE latest.history_id = t.history_id AND latest.task_id = t.id)))) + `, pgx.RowTo[string], session.HistoryID, snapshot.URI, snapshot.Fidelity, snapshot.Atespace) + if err != nil { + return fmt.Errorf("read interrupted Session tasks: %w", err) + } + for _, taskID := range taskIDs { + version, err := taskVersion(ctx, tx, session.HistoryID, taskID) + if err != nil { + return err + } + pending, err := queryOne(ctx, tx, ` + SELECT EXISTS (SELECT 1 FROM session_task_event + WHERE sequence = $1 AND NOT published AND quiescence_pending IS NOT NULL) + `, pgx.RowTo[bool], version) + if err != nil { + return err + } + if pending { + if err := settleSessionTask(ctx, tx, session, taskID, version); err != nil { + return err + } + } + stored, err := readSessionTask(ctx, tx, session.HistoryID, taskID) + if err != nil { + return err + } + task, err := unmarshalSessionTask(stored.Data) + if err != nil { + return err + } + if !task.Status.State.Terminal() { + task.Status = a2a.TaskStatus{State: a2a.TaskStateFailed, Timestamp: new(time.Now().UTC()), + Message: a2a.NewMessage(a2a.MessageRoleAgent, a2a.NewTextPart("The runtime crashed before this turn completed. Send a new message to continue."))} + event := &a2a.TaskStatusUpdateEvent{TaskID: task.ID, ContextID: task.ContextID, Status: task.Status} + if err := storeSessionTaskEvent(ctx, tx, session, task, event, false); err != nil { + return err + } + version, err = taskVersion(ctx, tx, session.HistoryID, taskID) + if err != nil { + return err + } + if err := publishPushBoundary(ctx, tx, session.HistoryID, taskID, version, task); err != nil { + return err + } + } + } + return nil +} diff --git a/go/core/internal/database/session_recovery_test.go b/go/core/internal/database/session_recovery_test.go new file mode 100644 index 0000000000..68827b938d --- /dev/null +++ b/go/core/internal/database/session_recovery_test.go @@ -0,0 +1,255 @@ +package database + +import ( + "testing" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/stretchr/testify/require" + "google.golang.org/protobuf/proto" + + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" +) + +func TestCrashedRecoverySettlesTasksAtomically(t *testing.T) { + for _, test := range []struct { + name string + state, want a2a.TaskState + published, captured bool + }{ + {name: "submitted", state: a2a.TaskStateSubmitted, want: a2a.TaskStateFailed}, + {name: "working", state: a2a.TaskStateWorking, want: a2a.TaskStateFailed}, + {name: "completed before publication", state: a2a.TaskStateCompleted, want: a2a.TaskStateCompleted}, + {name: "unpublished input wait", state: a2a.TaskStateInputRequired, want: a2a.TaskStateFailed}, + {name: "unpublished auth wait", state: a2a.TaskStateAuthRequired, want: a2a.TaskStateFailed}, + {name: "published input wait", state: a2a.TaskStateInputRequired, want: a2a.TaskStateFailed, published: true}, + {name: "published auth wait", state: a2a.TaskStateAuthRequired, want: a2a.TaskStateFailed, published: true}, + {name: "captured input wait", state: a2a.TaskStateInputRequired, want: a2a.TaskStateInputRequired, published: true, captured: true}, + {name: "captured auth wait", state: a2a.TaskStateAuthRequired, want: a2a.TaskStateAuthRequired, published: true, captured: true}, + } { + t.Run(test.name, func(t *testing.T) { + client := NewClient(setupTestDB(t), "public") + sessionFixture(t, client, t.Context(), "team-a", "revision", "assistant", "kagent") + session, _, err := client.CreateSession(t.Context(), newSessionRequest(uuid.NewString(), "assistant", "kagent", ""), uuid.NewString()) + require.NoError(t, err) + session, err = markSessionReady(t.Context(), client, session.Id, "runtime.example") + require.NoError(t, err) + task := newSessionTask(uuid.NewString(), "original input") + task.ContextID, task.Status.State = session.ContextId, test.state + version, err := client.CreateRuntimeTask(t.Context(), session.Id, taskMutationHash("create"), task, "") + require.NoError(t, err) + if test.published { + require.NoError(t, client.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + } + snapshot := SessionTaskSnapshot{} + if test.captured { + work, err := client.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + snapshot = SessionTaskSnapshot{Atespace: "team-a", URI: "s3://snapshots/captured", Fidelity: "SNAPSHOT_FIDELITY_VOLUMES"} + require.NoError(t, client.FinishSessionQuiescence(t.Context(), work, &snapshot)) + } + before, _, err := client.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + page, err := client.ListSessionRecoveryCandidates(t.Context(), "", 100) + require.NoError(t, err) + require.Empty(t, page) + + operation, err := client.BeginCrashedSessionRecovery(t.Context(), session.Id, lifecycleGeneration(t, client, session.Id), "actor-"+session.Id, snapshot) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_FAILED, operation.Instance.State) + require.Equal(t, apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME, operation.Instance.Operation) + require.Equal(t, SessionRuntimeCrashed, operation.Instance.Failure.Reason) + stored, err := client.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + require.Equal(t, test.want, stored.Status.State) + require.Equal(t, before.History, stored.History) + require.Equal(t, before.Artifacts, stored.Artifacts) + if test.want == test.state { + require.Equal(t, before.Status, stored.Status) + } + _, settledVersion, err := client.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + replayed, err := client.BeginCrashedSessionRecovery(t.Context(), session.Id, lifecycleGeneration(t, client, session.Id), "actor-"+session.Id, snapshot) + require.NoError(t, err) + require.Equal(t, operation.ID, replayed.ID) + _, replayedVersion, err := client.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + require.Equal(t, settledVersion, replayedVersion) + _, err = client.UpdateSessionTask(t.Context(), session.Id, version, taskMutationHash("late"), task, task, "") + if test.state.Terminal() { + require.ErrorIs(t, err, ErrFailedPrecondition) + } else { + require.ErrorIs(t, err, ErrConflict) + } + require.ErrorIs(t, client.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "next"), ErrConflict) + _, _, err = client.ReserveSessionCheckpoint(t.Context(), &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID)}, session.Creator, uuid.NewString()) + require.ErrorIs(t, err, ErrConflict) + page, err = client.ListSessionRecoveryCandidates(t.Context(), "", 100) + require.NoError(t, err) + require.Len(t, page, 1) + require.Equal(t, session.Id, page[0].Id) + }) + } +} + +func TestCrashedRecoveryAdmissionRespectsFences(t *testing.T) { + for _, test := range []struct { + name string + actorUID string + checkpoint, superseded, lifecycle bool + }{ + {name: "replacement actor", actorUID: "replacement"}, + {name: "checkpoint", checkpoint: true}, + {name: "revert after the crash was observed", superseded: true}, + {name: "lifecycle", lifecycle: true}, + } { + t.Run(test.name, func(t *testing.T) { + client := NewClient(setupTestDB(t), "public") + sessionFixture(t, client, t.Context(), "team-a", "revision", "assistant", "kagent") + session, task := waitingTaskFixture(t, client) + uid := "actor-" + session.Id + // Callers read the generation before they observe the crash. + generation := lifecycleGeneration(t, client, session.Id) + switch { + case test.actorUID != "": + uid = test.actorUID + case test.checkpoint: + task.Status.State = a2a.TaskStateCompleted + require.NoError(t, saveRuntimeTask(t, client, session.Id, task, task, &SessionTaskSnapshot{ + Atespace: "team-a", URI: "s3://snapshots/completed", Fidelity: "SNAPSHOT_FIDELITY_VOLUMES", + })) + _, _, err := client.ReserveSessionCheckpoint(t.Context(), &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID)}, session.Creator, uuid.NewString()) + require.NoError(t, err) + case test.superseded: + // Another revert finishes on the same Actor UID, and a new turn + // starts once ingress wakes the restored Actor. + completeCrashedRecovery(t, client, session.Id, uid) + task = newSessionTask(uuid.NewString(), "next input") + task.ContextID, task.Status.State = session.ContextId, a2a.TaskStateWorking + _, err := client.CreateRuntimeTask(t.Context(), session.Id, taskMutationHash("next"), task, "") + require.NoError(t, err) + case test.lifecycle: + _, err := client.BeginSessionOperation(t.Context(), session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_SUSPEND) + require.NoError(t, err) + } + before, err := client.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + _, err = client.BeginCrashedSessionRecovery(t.Context(), session.Id, generation, uid, SessionTaskSnapshot{}) + require.ErrorIs(t, err, ErrConflict) + after, err := client.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + require.True(t, proto.Equal(before, after)) + if test.superseded { + running, _, err := client.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateWorking, running.Status.State, "a stale observation must not fail the new turn") + } + }) + } +} + +// A crash during post-turn suspension leaves the worker's claim issued, because +// the worker cannot tell whether its suspension took effect. The confirmed crash +// settles that: admission retires the claim instead of waiting for it. +func TestCrashedRecoveryRetiresIssuedIdleSuspension(t *testing.T) { + client := NewClient(setupTestDB(t), "public") + sessionFixture(t, client, t.Context(), "team-a", "revision", "assistant", "kagent") + session, task := waitingTaskFixture(t, client) + _, version, err := client.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + for _, state := range []a2a.TaskState{a2a.TaskStateWorking, a2a.TaskStateInputRequired} { + task.Status.State = state + version, err = client.UpdateSessionTask(t.Context(), session.Id, version, taskMutationHash(string(state)), task, task, "") + require.NoError(t, err) + } + require.NoError(t, client.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + work, err := client.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + uid := "actor-" + session.Id + + operation, err := client.BeginCrashedSessionRecovery(t.Context(), session.Id, lifecycleGeneration(t, client, session.Id), uid, SessionTaskSnapshot{}) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME, operation.Instance.Operation) + pending, err := queryOne(t.Context(), client.db, `SELECT quiescence_pending FROM session_task_event WHERE sequence = $1`, pgx.RowTo[bool], work.Version) + require.NoError(t, err) + require.False(t, pending, "admission retires the claim in its own transaction") + finishCrashedRecovery(t, client, session.Id, operation, uid) + + late := SessionTaskSnapshot{Atespace: "team-a", URI: "s3://snapshots/late", Fidelity: "SNAPSHOT_FIDELITY_VOLUMES"} + require.NoError(t, client.FinishSessionQuiescence(t.Context(), work, &late)) + row, err := readSession(t.Context(), client.db, session.Id) + require.NoError(t, err) + recorded, err := queryOne(t.Context(), client.db, ` + SELECT EXISTS (SELECT 1 FROM session_task WHERE history_id = $1 AND snapshot_uri = $2) + OR EXISTS (SELECT 1 FROM session_task_event WHERE history_id = $1 AND snapshot_uri = $2) + `, pgx.RowTo[bool], row.HistoryID, late.URI) + require.NoError(t, err) + require.False(t, recorded, "a late finish of the retired claim records nothing") + _, err = client.ClaimSessionQuiescence(t.Context()) + require.ErrorIs(t, err, ErrNotFound) + next := newSessionTask(uuid.NewString(), "next input") + next.ContextID, next.Status.State = session.ContextId, a2a.TaskStateWorking + _, err = client.CreateRuntimeTask(t.Context(), session.Id, taskMutationHash("next"), next, "") + require.NoError(t, err, "the retired claim no longer blocks a new turn") +} + +func lifecycleGeneration(t *testing.T, client *Client, sessionID string) uuid.UUID { + t.Helper() + generation, err := client.SessionLifecycleGeneration(t.Context(), sessionID) + require.NoError(t, err) + return generation +} + +// completeCrashedRecovery admits and finishes a revert as a successful concurrent +// caller would, leaving the Session READY on the same Actor UID. +func completeCrashedRecovery(t *testing.T, client *Client, sessionID, actorUID string) { + t.Helper() + operation, err := client.BeginCrashedSessionRecovery(t.Context(), sessionID, lifecycleGeneration(t, client, sessionID), actorUID, SessionTaskSnapshot{}) + require.NoError(t, err) + finishCrashedRecovery(t, client, sessionID, operation, actorUID) +} + +func finishCrashedRecovery(t *testing.T, client *Client, sessionID string, operation *SessionOperation, actorUID string) { + t.Helper() + executor := uuid.New() + claimed, err := client.ClaimSessionOperation(t.Context(), sessionID, operation.ID, executor) + require.NoError(t, err) + require.True(t, claimed) + recovered, err := client.FinishSessionOperation(t.Context(), sessionID, operation.ID, executor, "", actorUID, "") + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, recovered.State) +} + +func TestCrashedRecoveryAdmissionRollsBackOnMalformedTask(t *testing.T) { + client := NewClient(setupTestDB(t), "public") + sessionFixture(t, client, t.Context(), "team-a", "revision", "assistant", "kagent") + session, task := waitingTaskFixture(t, client) + _, version, err := client.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + for _, state := range []a2a.TaskState{a2a.TaskStateWorking, a2a.TaskStateInputRequired} { + task.Status.State = state + version, err = client.UpdateSessionTask(t.Context(), session.Id, version, taskMutationHash(string(state)), task, task, "") + require.NoError(t, err) + } + require.NoError(t, client.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + row, err := readSession(t.Context(), client.db, session.Id) + require.NoError(t, err) + before, err := readSessionTask(t.Context(), client.db, row.HistoryID, string(task.ID)) + require.NoError(t, err) + // A corrupt projection makes settlement fail after the lifecycle write. + _, err = client.db.Exec(t.Context(), `UPDATE session_task SET data = $2 WHERE history_id = $1 AND id = $3`, row.HistoryID, []byte{0xff}, string(task.ID)) + require.NoError(t, err) + _, err = client.BeginCrashedSessionRecovery(t.Context(), session.Id, lifecycleGeneration(t, client, session.Id), "actor-"+session.Id, SessionTaskSnapshot{}) + require.Error(t, err) + current, err := client.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, current.State) + require.Equal(t, apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE, current.Operation) + require.Nil(t, current.Failure) + _, err = client.db.Exec(t.Context(), `UPDATE session_task SET data = $2 WHERE history_id = $1 AND id = $3`, row.HistoryID, before.Data, string(task.ID)) + require.NoError(t, err) + work, err := client.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + require.Equal(t, string(task.ID), work.TaskID) +} diff --git a/go/core/internal/database/task_finalization.go b/go/core/internal/database/task_finalization.go index 9907257b21..ac209c8dfe 100644 --- a/go/core/internal/database/task_finalization.go +++ b/go/core/internal/database/task_finalization.go @@ -23,54 +23,60 @@ func (c *Client) SettleSessionTask(ctx context.Context, sessionID, taskID string if session.State == "RUNTIME_STATE_DELETED" { return ErrNotFound } - type boundary struct { - Data []byte - ExpectedVersion int64 - Published bool - } - row, err := queryOne(ctx, tx, ` - SELECT data, expected_version, published FROM session_task_event - WHERE history_id = $1 AND task_id = $2 AND sequence = $3 - AND expected_version IS NOT NULL AND quiescence_pending IS NOT NULL - `, pgx.RowToStructByName[boundary], session.HistoryID, taskID, version) - if errors.Is(err, pgx.ErrNoRows) { - return ErrConflict - } - if err != nil { - return err - } - if row.Published { - return nil - } - stored, err := readSessionTask(ctx, tx, session.HistoryID, taskID) - if err != nil { - return err - } - wire, err := applyStoredBoundary(stored.Data, row.Data) - if err != nil { - return err - } - task, err := pbconv.FromProtoTask(wire) - if err != nil { - return err - } - data, err := proto.Marshal(wire) - if err != nil { - return err - } - if _, err := saveTaskProjection(ctx, tx, session.HistoryID, taskID, string(task.Status.State), task.Status.Timestamp, data); err != nil { - return err - } - if err := execSQL(ctx, tx, ` - UPDATE session_task_event SET published = TRUE - WHERE history_id = $1 AND task_id = $2 AND sequence > $3 AND sequence <= $4 - `, session.HistoryID, taskID, row.ExpectedVersion, version); err != nil { - return err - } - return publishPushBoundary(ctx, tx, session.HistoryID, taskID, version, task) + return settleSessionTask(ctx, tx, session, taskID, version) }) } +// settleSessionTask publishes an already saved boundary under the Session row lock. +// Recovery uses the same publication after confirming the runtime crashed. +func settleSessionTask(ctx context.Context, tx pgx.Tx, session sessionRow, taskID string, version int64) error { + type boundary struct { + Data []byte + ExpectedVersion int64 + Published bool + } + row, err := queryOne(ctx, tx, ` + SELECT data, expected_version, published FROM session_task_event + WHERE history_id = $1 AND task_id = $2 AND sequence = $3 + AND expected_version IS NOT NULL AND quiescence_pending IS NOT NULL + `, pgx.RowToStructByName[boundary], session.HistoryID, taskID, version) + if errors.Is(err, pgx.ErrNoRows) { + return ErrConflict + } + if err != nil { + return err + } + if row.Published { + return nil + } + stored, err := readSessionTask(ctx, tx, session.HistoryID, taskID) + if err != nil { + return err + } + wire, err := applyStoredBoundary(stored.Data, row.Data) + if err != nil { + return err + } + task, err := pbconv.FromProtoTask(wire) + if err != nil { + return err + } + data, err := proto.Marshal(wire) + if err != nil { + return err + } + if _, err := saveTaskProjection(ctx, tx, session.HistoryID, taskID, string(task.Status.State), task.Status.Timestamp, data); err != nil { + return err + } + if err := execSQL(ctx, tx, ` + UPDATE session_task_event SET published = TRUE + WHERE history_id = $1 AND task_id = $2 AND sequence > $3 AND sequence <= $4 + `, session.HistoryID, taskID, row.ExpectedVersion, version); err != nil { + return err + } + return publishPushBoundary(ctx, tx, session.HistoryID, taskID, version, task) +} + func applyStoredBoundary(data, eventData []byte) (*a2apb.Task, error) { task, event := &a2apb.Task{}, &a2apb.StreamResponse{} if err := proto.Unmarshal(data, task); err != nil { diff --git a/go/core/internal/substrate/client.go b/go/core/internal/substrate/client.go index c6eef1345b..6ac07cdd2c 100644 --- a/go/core/internal/substrate/client.go +++ b/go/core/internal/substrate/client.go @@ -203,6 +203,16 @@ func (c *Client) ResumeActor(ctx context.Context, atespace, actorID string) (*at return resp.GetActor(), nil } +func (c *Client) RevertActor(ctx context.Context, atespace, actorID string) (*ateapipb.Actor, error) { + ctx, cancel := c.callCtx(ctx) + defer cancel() + resp, err := c.ControlClient.RevertActor(ctx, &ateapipb.RevertActorRequest{Actor: actorRef(atespace, actorID)}) + if err != nil { + return nil, fmt.Errorf("failed to revert Actor %s/%s: %w", atespace, actorID, err) + } + return resp.GetActor(), nil +} + // SuspendActor waits for durable storage, including an upload left by an earlier // attempt. Substrate's synchronous mode returns already-suspended Actors as-is. func (c *Client) SuspendActor(ctx context.Context, atespace, actorID string) (*ateapipb.Actor, error) { diff --git a/go/core/internal/substrate/snapshot.go b/go/core/internal/substrate/snapshot.go index 3d1163c28e..302c07e975 100644 --- a/go/core/internal/substrate/snapshot.go +++ b/go/core/internal/substrate/snapshot.go @@ -22,3 +22,14 @@ func SnapshotURI(snapshot *ateapipb.Snapshot) string { } return storage.GetObject().GetSnapshotUri() } + +// Checkpoints require the newest generation; revert retains the newest durable copy. +func LatestDurableActorSnapshot(actor *ateapipb.Actor) *ateapipb.Snapshot { + var latest *ateapipb.Snapshot + for _, snapshot := range actor.GetStatus().GetSnapshots() { + if SnapshotURI(snapshot) != "" && snapshot.GetGeneration() > latest.GetGeneration() { + latest = snapshot + } + } + return latest +} diff --git a/go/core/internal/substrate/snapshot_test.go b/go/core/internal/substrate/snapshot_test.go index faee8a936c..dfda60e0cf 100644 --- a/go/core/internal/substrate/snapshot_test.go +++ b/go/core/internal/substrate/snapshot_test.go @@ -40,3 +40,23 @@ func TestActorSnapshotGenerationAndDurability(t *testing.T) { require.Nil(t, LatestActorSnapshot(nil)) require.Empty(t, SnapshotURI(nil)) } + +func TestRevertSnapshotSkipsIncompleteGenerations(t *testing.T) { + durable := &ateapipb.Snapshot{Generation: 2, DurableSnapshot: &ateapipb.SnapshotStorage{ + Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_COMPLETED, + Object: &ateapipb.ObjectSnapshot{SnapshotUri: "s3://snapshots/durable"}, + }} + for _, latest := range []*ateapipb.Snapshot{ + {Generation: 3, Locality: "node-1"}, + {Generation: 3, DurableSnapshot: &ateapipb.SnapshotStorage{Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_PENDING}}, + {Generation: 3, DurableSnapshot: &ateapipb.SnapshotStorage{Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_COMPLETED}}, + } { + actor := &ateapipb.Actor{Status: &ateapipb.ActorStatus{Snapshots: []*ateapipb.Snapshot{latest, durable}}} + require.Same(t, durable, LatestDurableActorSnapshot(actor)) + require.Same(t, latest, LatestActorSnapshot(actor)) + } + require.Nil(t, LatestDurableActorSnapshot(nil)) + require.Nil(t, LatestDurableActorSnapshot(&ateapipb.Actor{Status: &ateapipb.ActorStatus{ + Snapshots: []*ateapipb.Snapshot{{Generation: 1, Locality: "node-1"}}, + }})) +} From 7fb50c6514735c5f00de0c88db80e89150caea83 Mon Sep 17 00:00:00 2001 From: Ada Lundhe Date: Fri, 23 Oct 2026 12:31:55 -0700 Subject: [PATCH 09/12] fix(session): implement manual crashed-runtime recovery Signed-off-by: Ada Lundhe --- .../internal/service/session/quiescence.go | 21 +- .../service/session/quiescence_test.go | 4 + go/core/internal/service/session/recovery.go | 154 +++++ .../internal/service/session/recovery_test.go | 541 ++++++++++++++++++ go/core/internal/service/session/workflow.go | 7 + 5 files changed, 723 insertions(+), 4 deletions(-) create mode 100644 go/core/internal/service/session/recovery.go create mode 100644 go/core/internal/service/session/recovery_test.go diff --git a/go/core/internal/service/session/quiescence.go b/go/core/internal/service/session/quiescence.go index 283e70c7cc..aba8c40ffb 100644 --- a/go/core/internal/service/session/quiescence.go +++ b/go/core/internal/service/session/quiescence.go @@ -28,16 +28,29 @@ func (w *ActorWorkflow) wakeQuiescence() { } } -// Start suspends idle sessions to VOLUMES independently of task publication. -// Startup and periodic recovery scans discover work missed by local signals. -// Workers drain claims with bounded concurrency; task reads never wait for them. +// Start suspends idle sessions to VOLUMES independently of task publication and +// retries crash recoveries admitted by explicit requests. +// Startup and periodic scans discover idle work missed by local signals; crash +// checks use the same recovery interval. Workers drain claims with bounded concurrency. func (w *ActorWorkflow) Start(ctx context.Context) error { var workers sync.WaitGroup - for range 4 { + for worker := range 4 { workers.Go(func() { timer := time.NewTimer(w.recoveryInterval) defer timer.Stop() + var recoveryAfter string + nextRecovery := time.Now().Add(w.recoveryInterval) for ctx.Err() == nil { + if worker == 0 && !time.Now().Before(nextRecovery) { + scanCtx, cancel := context.WithTimeout(ctx, database.RuntimeOperationTimeout) + after, err := w.recoverSessionPage(scanCtx, recoveryAfter) + cancel() + if err != nil && ctx.Err() == nil { + logging.FromContext(ctx).ErrorContext(ctx, "scan interrupted sessions", "error", err) + } + recoveryAfter = after + nextRecovery = time.Now().Add(w.recoveryInterval) + } work, err := w.store.ClaimSessionQuiescence(ctx) if err == nil { // Hand off a hint before runtime I/O so a single settlement diff --git a/go/core/internal/service/session/quiescence_test.go b/go/core/internal/service/session/quiescence_test.go index 99e8852d72..d6b9af4991 100644 --- a/go/core/internal/service/session/quiescence_test.go +++ b/go/core/internal/service/session/quiescence_test.go @@ -32,6 +32,10 @@ type idleQuiescenceStore struct { var _ workflowStore = (*idleQuiescenceStore)(nil) +func (s *idleQuiescenceStore) ListSessionRecoveryCandidates(context.Context, string, int) ([]*apiv1.Session, error) { + return nil, nil +} + func (s *idleQuiescenceStore) ClaimSessionQuiescence(context.Context) (*database.SessionQuiescence, error) { s.claims.Add(1) if s.claimErr != nil { diff --git a/go/core/internal/service/session/recovery.go b/go/core/internal/service/session/recovery.go new file mode 100644 index 0000000000..27922bc3ed --- /dev/null +++ b/go/core/internal/service/session/recovery.go @@ -0,0 +1,154 @@ +package session + +import ( + "context" + "errors" + "fmt" + "time" + + "github.com/google/uuid" + ateapipb "go.ate.dev/substrate/api/ateapi/v1" + "google.golang.org/protobuf/proto" + + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" + "github.com/kagent-dev/kagent/go/core/internal/database" + "github.com/kagent-dev/kagent/go/core/internal/substrate" + "github.com/kagent-dev/kagent/go/pkg/logging" +) + +// recoverSessionPage retries admitted recoveries in ID order. +// READY Sessions are excluded, so a background retry cannot admit new work. +func (w *ActorWorkflow) recoverSessionPage(ctx context.Context, afterID string) (string, error) { + const pageSize = 100 + sessions, err := w.store.ListSessionRecoveryCandidates(ctx, afterID, pageSize) + if err != nil { + return afterID, err + } + for _, session := range sessions { + attemptCtx, cancel := context.WithTimeout(ctx, database.RuntimeOperationTimeout) + err := w.recoverSessionCandidate(attemptCtx, session) + cancel() + if err != nil && !errors.Is(err, database.ErrConflict) && ctx.Err() == nil { + logging.FromContext(ctx).ErrorContext(ctx, "recover crashed session", "session_id", session.Id, "error", err) + } + afterID = session.Id + if ctx.Err() != nil { + return afterID, ctx.Err() + } + } + if len(sessions) < pageSize { + afterID = "" + } + return afterID, nil +} + +func (w *ActorWorkflow) recoverSessionCandidate(ctx context.Context, session *apiv1.Session) error { + _, err := w.Revert(ctx, session) + if errors.Is(err, database.ErrFailedPrecondition) { + return nil + } + return err +} + +// Revert restores a confirmed crashed runtime without replaying its interrupted turn. +// Retries join the same durable operation and preserve the Actor identity. +func (w *ActorWorkflow) Revert(ctx context.Context, session *apiv1.Session) (*apiv1.Session, error) { + ctx, cancel := context.WithTimeout(ctx, database.RuntimeOperationTimeout) + defer cancel() + if session.State == apiv1.RuntimeState_RUNTIME_STATE_FAILED && session.Failure.GetReason() == database.SessionRuntimeCrashed { + return w.run(ctx, session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME) + } + if session.State != apiv1.RuntimeState_RUNTIME_STATE_READY || session.Operation != apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE { + return nil, database.ErrConflict + } + revision, err := w.store.GetRuntimeRevision(ctx, session.PreparedRevision) + if err != nil { + return nil, err + } + // Read before observing Substrate. A concurrent revert keeps the Actor UID and + // returns to READY, so only the generation shows that this crash observation + // predates it. + generation, err := w.store.SessionLifecycleGeneration(ctx, session.Id) + if err != nil { + return nil, err + } + actor, err := w.actors.GetActor(ctx, revision.ActorTemplateAtespace, substrate.ActorName(session.Id)) + if err != nil { + return nil, err + } + if err := w.verifyActor(ctx, session, revision, actor); err != nil { + return nil, err + } + if actor.GetStatus().GetState() != ateapipb.ActorState_ACTOR_STATE_CRASHED { + return nil, fmt.Errorf("runtime is not crashed: %w", database.ErrFailedPrecondition) + } + snapshot := substrate.LatestDurableActorSnapshot(actor) + operation, err := w.store.BeginCrashedSessionRecovery(ctx, session.Id, generation, actor.GetMetadata().GetUid(), database.SessionTaskSnapshot{ + Atespace: actor.GetMetadata().GetAtespace(), URI: substrate.SnapshotURI(snapshot), + Fidelity: snapshot.GetFidelity().String(), + }) + if err != nil { + return nil, err + } + return w.execute(ctx, operation) +} + +// recoverCrashedSession reconciles lost revert responses without replacing the Actor. +func (w *ActorWorkflow) recoverCrashedSession(ctx context.Context, operation *database.SessionOperation) (_ *apiv1.Session, err error) { + session := operation.Instance + if session.Operation != apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME { + return nil, database.ErrConflict + } + executorID := uuid.New() + claimed, err := w.store.ClaimSessionOperation(ctx, session.Id, operation.ID, executorID) + if err != nil { + return nil, err + } + if !claimed { + current, err := w.store.GetSessionOperation(ctx, session.Id, operation.ID) + if err != nil { + return nil, err + } + return operationOutcome(current) + } + defer func() { + if err != nil { + releaseCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + defer cancel() + err = errors.Join(err, w.store.ReleaseRuntimeOperation(releaseCtx, session.Id, operation.ID, executorID)) + } + }() + revision, err := w.store.GetRuntimeRevision(ctx, session.PreparedRevision) + if err != nil { + return nil, err + } + atespace, name := revision.ActorTemplateAtespace, substrate.ActorName(session.Id) + actor, err := w.actors.GetActor(ctx, atespace, name) + if err != nil { + return nil, err + } + if err := w.verifyActor(ctx, session, revision, actor); err != nil { + return nil, err + } + before := actor + beforeSnapshot := proto.CloneOf(substrate.LatestDurableActorSnapshot(actor)) + if beforeSnapshot != nil { + beforeSnapshot.Locality = "" + } + switch actor.GetStatus().GetState() { + case ateapipb.ActorState_ACTOR_STATE_CRASHED, ateapipb.ActorState_ACTOR_STATE_REVERTING: + actor, err = w.actors.RevertActor(ctx, atespace, name) + if err != nil { + return nil, fmt.Errorf("revert crashed Actor %s/%s: %w", atespace, name, err) + } + case ateapipb.ActorState_ACTOR_STATE_SUSPENDED: + default: + return nil, fmt.Errorf("crashed Actor %s/%s has unexpected state %s", atespace, name, actor.GetStatus().GetState()) + } + if !validActorIdentity(actor, revision, name) || actor.GetMetadata().GetUid() != before.GetMetadata().GetUid() || actor.GetStatus().GetState() != ateapipb.ActorState_ACTOR_STATE_SUSPENDED || !proto.Equal(beforeSnapshot, substrate.LatestActorSnapshot(actor)) { + return nil, fmt.Errorf("revert Actor %s/%s returned unexpected identity or state", atespace, name) + } + finishCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + defer cancel() + return w.store.FinishSessionOperation(finishCtx, session.Id, operation.ID, executorID, "", actor.GetMetadata().GetUid(), "") +} diff --git a/go/core/internal/service/session/recovery_test.go b/go/core/internal/service/session/recovery_test.go new file mode 100644 index 0000000000..d7ce338a83 --- /dev/null +++ b/go/core/internal/service/session/recovery_test.go @@ -0,0 +1,541 @@ +package session + +import ( + "context" + "crypto/sha256" + "fmt" + "slices" + "testing" + "time" + + "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/grpc/codes" + "google.golang.org/grpc/status" + "google.golang.org/protobuf/proto" + + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" + "github.com/kagent-dev/kagent/go/core/internal/database" + "github.com/kagent-dev/kagent/go/core/internal/substrate" +) + +func (a *lifecycleTestActors) RevertActor(_ context.Context, atespace, name string) (*ateapipb.Actor, error) { + a.mu.Lock() + defer a.mu.Unlock() + actor := a.actors[actorKey(atespace, name)] + if actor.Status.State != ateapipb.ActorState_ACTOR_STATE_CRASHED && actor.Status.State != ateapipb.ActorState_ACTOR_STATE_REVERTING { + return nil, status.Error(codes.FailedPrecondition, "actor is not crashed or reverting") + } + actor.Status.State = ateapipb.ActorState_ACTOR_STATE_SUSPENDED + actor.Status.Crash = nil + snapshot := proto.CloneOf(substrate.LatestDurableActorSnapshot(actor)) + actor.Status.Snapshots = nil + if snapshot != nil { + snapshot.Locality = "" + actor.Status.Snapshots = []*ateapipb.Snapshot{snapshot} + } + return proto.CloneOf(actor), nil +} + +type recoveryTestActors struct { + *lifecycleTestActors + revert func(context.Context, string, string) (*ateapipb.Actor, error) + readErr error + // afterRead runs once, after the first successful read returns its copy. + afterRead func(context.Context) +} + +func (a *recoveryTestActors) RevertActor(ctx context.Context, atespace, name string) (*ateapipb.Actor, error) { + return a.revert(ctx, atespace, name) +} + +func (a *recoveryTestActors) GetActor(ctx context.Context, atespace, name string) (*ateapipb.Actor, error) { + if a.readErr != nil { + return nil, a.readErr + } + actor, err := a.lifecycleTestActors.GetActor(ctx, atespace, name) + if hook := a.afterRead; hook != nil && err == nil { + a.afterRead = nil + hook(ctx) + } + return actor, err +} + +func revertGeneration(t *testing.T, store *lifecycleTestStore, sessionID string) uuid.UUID { + t.Helper() + generation, err := store.SessionLifecycleGeneration(t.Context(), sessionID) + require.NoError(t, err) + return generation +} + +func crashedTurnFixture(t *testing.T) (*lifecycleTestStore, *lifecycleTestActors, *apiv1.Session, *a2a.Task, int64) { + t.Helper() + store, session := lifecycleFixture(t) + actors := &lifecycleTestActors{actors: map[string]*ateapipb.Actor{}} + var err error + session, err = NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Second).Create(t.Context(), session) + require.NoError(t, err) + message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("original input")) + message.ContextID = session.ContextId + task := a2a.NewSubmittedTask(message, message) + task.Status.State = a2a.TaskStateWorking + hash := sha256.Sum256([]byte("working")) + version, err := store.CreateRuntimeTask(t.Context(), session.Id, hash[:], task, "") + require.NoError(t, err) + actors.actors[actorKey("team-a", substrate.ActorName(session.Id))].Status.State = ateapipb.ActorState_ACTOR_STATE_CRASHED + return store, actors, session, task, version +} + +func TestCrashedSessionRecoveryRetriesWithoutLosingState(t *testing.T) { + for _, test := range []struct { + name string + lostResponse bool + wantReverts int + }{ + {name: "revert failed", wantReverts: 2}, + {name: "revert response lost", lostResponse: true, wantReverts: 1}, + } { + t.Run(test.name, func(t *testing.T) { + store, base, session, task, _ := crashedTurnFixture(t) + calls := 0 + actors := &recoveryTestActors{lifecycleTestActors: base} + actors.revert = func(ctx context.Context, atespace, name string) (*ateapipb.Actor, error) { + calls++ + if calls == 1 { + if test.lostResponse { + _, err := base.RevertActor(ctx, atespace, name) + require.NoError(t, err) + } + return nil, status.Error(codes.Unavailable, "response unavailable") + } + return base.RevertActor(ctx, atespace, name) + } + workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Second) + _, err := workflow.Revert(t.Context(), session) + require.Error(t, err) + failed, err := store.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_FAILED, failed.State) + require.Equal(t, apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME, failed.Operation) + require.Equal(t, database.SessionRuntimeCrashed, failed.Failure.Reason) + interrupted, interruptedVersion, err := store.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateFailed, interrupted.Status.State) + published, err := store.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateFailed, published.Status.State) + require.Len(t, published.History, 1) + require.ErrorIs(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), uuid.NewString()), database.ErrConflict) + // A new workflow instance discovers the admitted operation from PostgreSQL. + restarted := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Second) + _, err = restarted.recoverSessionPage(t.Context(), "") + require.NoError(t, err) + ready, err := store.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, ready.State) + require.Nil(t, ready.Failure) + stored, err := store.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateFailed, stored.Status.State) + require.Equal(t, test.wantReverts, calls) + require.Len(t, stored.History, 1, "the original input remains archived once") + _, recoveredVersion, err := store.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + require.Equal(t, interruptedVersion, recoveredVersion, "retries do not append another terminal event") + _, err = restarted.Revert(t.Context(), ready) + require.ErrorIs(t, err, database.ErrFailedPrecondition) + require.Equal(t, test.wantReverts, calls, "completed recovery cannot discard another execution") + }) + } +} + +func TestCrashedSessionRecoveryFencesConcurrentWork(t *testing.T) { + store, base, session, task, version := crashedTurnFixture(t) + entered, release := make(chan struct{}), make(chan struct{}) + actors := &recoveryTestActors{lifecycleTestActors: base} + actors.revert = func(ctx context.Context, atespace, name string) (*ateapipb.Actor, error) { + close(entered) + select { + case <-release: + return base.RevertActor(ctx, atespace, name) + case <-ctx.Done(): + return nil, ctx.Err() + } + } + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Second) + done := make(chan error, 1) + go func() { _, err := workflow.Revert(ctx, session); done <- err }() + select { + case <-entered: + case <-time.After(5 * time.Second): + t.Fatal("recovery was not admitted") + } + failed, err := store.GetSessionByID(ctx, session.Id) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_FAILED, failed.State) + _, err = workflow.Revert(ctx, failed) + require.ErrorIs(t, err, database.ErrConflict) + hash := sha256.Sum256([]byte("late write")) + _, err = store.UpdateSessionTask(ctx, session.Id, version, hash[:], task, task, "") + require.ErrorIs(t, err, database.ErrConflict) + close(release) + require.NoError(t, <-done) + _, err = store.UpdateSessionTask(ctx, session.Id, version, hash[:], task, task, "") + require.ErrorIs(t, err, database.ErrConflict) +} + +func TestSessionRecoveryRequiresConfirmedRecordedActorCrash(t *testing.T) { + for _, test := range []struct { + name string + state ateapipb.ActorState + uid string + readErr error + wantErr bool + }{ + {name: "quiet running actor", state: ateapipb.ActorState_ACTOR_STATE_RUNNING, wantErr: true}, + {name: "transient read failure", readErr: status.Error(codes.Unavailable, "unreachable"), wantErr: true}, + {name: "same-name replacement", state: ateapipb.ActorState_ACTOR_STATE_CRASHED, uid: "another-actor", wantErr: true}, + } { + t.Run(test.name, func(t *testing.T) { + store, base, session, task, _ := crashedTurnFixture(t) + actor := base.actors[actorKey("team-a", substrate.ActorName(session.Id))] + actor.Status.State = test.state + if test.uid != "" { + actor.Metadata.Uid = test.uid + } + actors := &recoveryTestActors{lifecycleTestActors: base, readErr: test.readErr} + actors.revert = func(context.Context, string, string) (*ateapipb.Actor, error) { + t.Fatal("unconfirmed crash issued a revert") + return nil, nil + } + _, err := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Second).Revert(t.Context(), session) + if test.wantErr { + require.Error(t, err) + } else { + require.NoError(t, err) + } + current, err := store.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, current.State) + require.Nil(t, current.Failure) + stored, err := store.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateWorking, stored.Status.State) + }) + } +} + +func TestSessionRecoveryCandidatePolicy(t *testing.T) { + store, actors, active, _, _ := crashedTurnFixture(t) + workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Second) + for _, state := range []a2a.TaskState{"", a2a.TaskStateInputRequired, a2a.TaskStateAuthRequired, a2a.TaskStateWorking} { + session, _, err := store.CreateSession(t.Context(), &apiv1.Session{Id: uuid.NewString(), Creator: "alice", Agent: active.Agent}, uuid.NewString()) + require.NoError(t, err) + session, err = workflow.Create(t.Context(), session) + require.NoError(t, err) + if state == "" { + continue + } + message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("input")) + message.ContextID = session.ContextId + task := a2a.NewSubmittedTask(message, message) + task.Status.State = state + hash := sha256.Sum256([]byte(task.ID)) + version, err := store.CreateRuntimeTask(t.Context(), session.Id, hash[:], task, "") + require.NoError(t, err) + if state != a2a.TaskStateWorking { + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + } + } + page, err := store.ListSessionRecoveryCandidates(t.Context(), "", 100) + require.NoError(t, err) + require.Empty(t, page, "manual recovery does not inspect idle or active runtimes") + actor, err := actors.GetActor(t.Context(), "team-a", substrate.ActorName(active.Id)) + require.NoError(t, err) + _, err = store.BeginCrashedSessionRecovery(t.Context(), active.Id, revertGeneration(t, store, active.Id), actor.GetMetadata().GetUid(), database.SessionTaskSnapshot{}) + require.NoError(t, err) + page, err = store.ListSessionRecoveryCandidates(t.Context(), "", 100) + require.NoError(t, err) + require.Len(t, page, 1, "explicitly admitted recovery remains eligible for retry") + require.Equal(t, active.Id, page[0].Id) +} + +func TestCrashedSessionRecoversAnOlderResumedTask(t *testing.T) { + store, actors, session, older, version := crashedTurnFixture(t) + older.Status.State = a2a.TaskStateInputRequired + hash := sha256.Sum256([]byte("waiting")) + waitingVersion, err := store.UpdateSessionTask(t.Context(), session.Id, version, hash[:], older, older, "") + require.NoError(t, err) + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(older.ID), waitingVersion)) + work, err := store.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + require.NoError(t, store.FinishSessionQuiescence(t.Context(), work, &database.SessionTaskSnapshot{ + Atespace: "team-a", URI: "s3://snapshots/waiting-turn", Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES.String(), + })) + + message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("another turn")) + message.ContextID = session.ContextId + newer := a2a.NewSubmittedTask(message, message) + newer.Status.State = a2a.TaskStateCompleted + hash = sha256.Sum256([]byte("newer completed task")) + newerVersion, err := store.CreateRuntimeTask(t.Context(), session.Id, hash[:], newer, "") + require.NoError(t, err) + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(newer.ID), newerVersion)) + work, err = store.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + require.NoError(t, store.FinishSessionQuiescence(t.Context(), work, &database.SessionTaskSnapshot{ + Atespace: "team-a", URI: "s3://snapshots/newer-completed-turn", Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES.String(), + })) + + older.Status.State = a2a.TaskStateWorking + hash = sha256.Sum256([]byte("resume older task")) + _, err = store.UpdateSessionTask(t.Context(), session.Id, waitingVersion, hash[:], older, older, "") + require.NoError(t, err) + workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Second) + _, err = workflow.Revert(t.Context(), session) + require.NoError(t, err) + failed, err := store.GetSettledSessionTask(t.Context(), session.Id, string(older.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateFailed, failed.Status.State) + retained, retainedVersion, err := store.GetVersionedSessionTask(t.Context(), session.Id, string(newer.ID)) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateCompleted, retained.Status.State) + require.Equal(t, newerVersion, retainedVersion, "recovery cannot rewrite a later completed task") + require.NoError(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), uuid.NewString())) +} + +func TestExplicitRevertRespectsCheckpointCreation(t *testing.T) { + store, actors, session, task, version := crashedTurnFixture(t) + task.Status.State = a2a.TaskStateCompleted + hash := sha256.Sum256([]byte("completed")) + version, err := store.UpdateSessionTask(t.Context(), session.Id, version, hash[:], task, task, "") + require.NoError(t, err) + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + work, err := store.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + require.NoError(t, store.FinishSessionQuiescence(t.Context(), work, &database.SessionTaskSnapshot{ + Atespace: "team-a", URI: "s3://snapshots/completed", Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES.String(), + })) + _, _, err = store.ReserveSessionCheckpoint(t.Context(), &apiv1.Checkpoint{ + Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID), + }, session.Creator, uuid.NewString()) + require.NoError(t, err) + workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute) + _, err = workflow.Revert(t.Context(), session) + require.ErrorIs(t, err, database.ErrConflict) + actor, err := actors.GetActor(t.Context(), "team-a", substrate.ActorName(session.Id)) + require.NoError(t, err) + require.Equal(t, ateapipb.ActorState_ACTOR_STATE_CRASHED, actor.Status.State) + stored, err := store.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateCompleted, stored.Status.State) +} + +// The Actor crashed during post-turn suspension, and the worker kept its claim +// because it cannot tell whether the suspension took effect. +func TestExplicitRevertRetiresIdleSuspensionOfCrashedActor(t *testing.T) { + store, actors, session, task, version := crashedTurnFixture(t) + task.Status.State = a2a.TaskStateCompleted + hash := sha256.Sum256([]byte("completed")) + version, err := store.UpdateSessionTask(t.Context(), session.Id, version, hash[:], task, task, "") + require.NoError(t, err) + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + work, err := store.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute) + recovered, err := workflow.Revert(t.Context(), session) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, recovered.State) + actor, err := actors.GetActor(t.Context(), "team-a", substrate.ActorName(session.Id)) + require.NoError(t, err) + require.Equal(t, ateapipb.ActorState_ACTOR_STATE_SUSPENDED, actor.Status.State) + stored, err := store.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateCompleted, stored.Status.State) + require.NoError(t, store.FinishSessionQuiescence(t.Context(), work, &database.SessionTaskSnapshot{ + Atespace: "team-a", URI: "s3://snapshots/late", Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES.String(), + }), "a late finish of the retired claim is a no-op") + _, err = store.ClaimSessionQuiescence(t.Context()) + require.ErrorIs(t, err, database.ErrNotFound) +} + +// Revert A observes the crash, then pauses while Revert B restores the same Actor +// and ingress wakes it for a new turn. A then sees the same Actor UID and a READY +// Session, so only the lifecycle generation shows that its observation is stale. +func TestExplicitRevertRejectsCrashObservedBeforeAnotherRevert(t *testing.T) { + store, base, session, _, _ := crashedTurnFixture(t) + key := actorKey("team-a", substrate.ActorName(session.Id)) + var next *a2a.Task + paused := &recoveryTestActors{lifecycleTestActors: base, revert: base.RevertActor, afterRead: func(ctx context.Context) { + recovered, err := NewActorWorkflow(store, base, make(chan struct{}, 1), time.Minute).Revert(ctx, session) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, recovered.State) + base.mu.Lock() + base.actors[key].Status.State = ateapipb.ActorState_ACTOR_STATE_RUNNING + base.mu.Unlock() + message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("next input")) + message.ContextID = session.ContextId + next = a2a.NewSubmittedTask(message, message) + next.Status.State = a2a.TaskStateWorking + hash := sha256.Sum256([]byte("next")) + _, err = store.CreateRuntimeTask(ctx, session.Id, hash[:], next, "") + require.NoError(t, err) + }} + _, err := NewActorWorkflow(store, paused, make(chan struct{}, 1), time.Minute).Revert(t.Context(), session) + require.ErrorIs(t, err, database.ErrConflict) + current, err := store.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, current.State) + require.Equal(t, apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE, current.Operation) + running, _, err := store.GetVersionedSessionTask(t.Context(), session.Id, string(next.ID)) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateWorking, running.Status.State, "the stale revert must not fail the new turn") +} + +func TestExplicitRevertCannotCheckpointDiscardedTerminalState(t *testing.T) { + store, actors, session, task, version := crashedTurnFixture(t) + task.Status.State = a2a.TaskStateCompleted + hash := sha256.Sum256([]byte("completed before snapshot")) + version, err := store.UpdateSessionTask(t.Context(), session.Id, version, hash[:], task, task, "") + require.NoError(t, err) + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute) + _, err = workflow.Revert(t.Context(), session) + require.NoError(t, err) + _, err = store.ClaimSessionQuiescence(t.Context()) + require.ErrorIs(t, err, database.ErrNotFound, "the reverted snapshot cannot represent discarded terminal state") + _, _, err = store.ReserveSessionCheckpoint(t.Context(), &apiv1.Checkpoint{ + Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID), + }, session.Creator, uuid.NewString()) + require.ErrorIs(t, err, database.ErrSnapshotPending) + stored, err := store.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateCompleted, stored.Status.State) +} + +func TestExplicitRevertPreservesOnlyCapturedWaits(t *testing.T) { + for _, state := range []a2a.TaskState{a2a.TaskStateInputRequired, a2a.TaskStateAuthRequired} { + for _, captured := range []bool{false, true} { + t.Run(fmt.Sprint(state, "/captured=", captured), func(t *testing.T) { + store, actors, session, task, version := crashedTurnFixture(t) + task.Status.State = state + hash := sha256.Sum256([]byte(state)) + version, err := store.UpdateSessionTask(t.Context(), session.Id, version, hash[:], task, task, "") + require.NoError(t, err) + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + if captured { + work, err := store.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + require.NoError(t, store.FinishSessionQuiescence(t.Context(), work, &database.SessionTaskSnapshot{ + Atespace: "team-a", URI: "s3://snapshots/waiting", Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES.String(), + })) + actors.actors[actorKey("team-a", substrate.ActorName(session.Id))].Status.Snapshots = []*ateapipb.Snapshot{{Generation: 1, Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES, DurableSnapshot: &ateapipb.SnapshotStorage{Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_COMPLETED, Object: &ateapipb.ObjectSnapshot{SnapshotUri: "s3://snapshots/waiting"}}}} + actor := actors.actors[actorKey("team-a", substrate.ActorName(session.Id))] + actor.Status.Snapshots[0].Locality = "lost-node" + actor.Status.Snapshots = append(actor.Status.Snapshots, &ateapipb.Snapshot{ + Generation: 2, DurableSnapshot: &ateapipb.SnapshotStorage{Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_PENDING}, + }) + } + _, err = NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute).Revert(t.Context(), session) + require.NoError(t, err) + stored, err := store.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + want := a2a.TaskStateFailed + if captured { + want = state + } + require.Equal(t, want, stored.Status.State) + }) + } + } +} + +func TestExplicitRevertRetainsWaitCapturedByLaterTurn(t *testing.T) { + store, actors, session, waiting, version := crashedTurnFixture(t) + waiting.Status.State = a2a.TaskStateInputRequired + hash := sha256.Sum256([]byte("waiting")) + version, err := store.UpdateSessionTask(t.Context(), session.Id, version, hash[:], waiting, waiting, "") + require.NoError(t, err) + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(waiting.ID), version)) + message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("another turn")) + message.ContextID = session.ContextId + later := a2a.NewSubmittedTask(message, message) + later.Status.State = a2a.TaskStateCompleted + hash = sha256.Sum256([]byte("later completed turn")) + version, err = store.CreateRuntimeTask(t.Context(), session.Id, hash[:], later, "") + require.NoError(t, err) + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, string(later.ID), version)) + work, err := store.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + require.NoError(t, store.FinishSessionQuiescence(t.Context(), work, &database.SessionTaskSnapshot{ + Atespace: "team-a", URI: "s3://snapshots/later-turn", Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES.String(), + })) + actors.actors[actorKey("team-a", substrate.ActorName(session.Id))].Status.Snapshots = []*ateapipb.Snapshot{{Generation: 1, Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES, DurableSnapshot: &ateapipb.SnapshotStorage{Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_COMPLETED, Object: &ateapipb.ObjectSnapshot{SnapshotUri: "s3://snapshots/later-turn"}}}} + _, err = NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute).Revert(t.Context(), session) + require.NoError(t, err) + retained, err := store.GetSettledSessionTask(t.Context(), session.Id, string(waiting.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateInputRequired, retained.Status.State) + retained, err = store.GetSettledSessionTask(t.Context(), session.Id, string(later.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateCompleted, retained.Status.State) +} + +func TestSessionRecoveryPagesPastUnavailableActors(t *testing.T) { + store, base, first, _, _ := crashedTurnFixture(t) + actors := &recoveryTestActors{lifecycleTestActors: base, revert: base.RevertActor, readErr: status.Error(codes.Unavailable, "control plane unavailable")} + workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute) + ids := []string{first.Id} + for range 100 { + session, _, err := store.CreateSession(t.Context(), &apiv1.Session{ + Id: uuid.NewString(), Creator: first.Creator, Agent: first.Agent, + }, uuid.NewString()) + require.NoError(t, err) + session, err = NewActorWorkflow(store, base, make(chan struct{}, 1), time.Minute).Create(t.Context(), session) + require.NoError(t, err) + message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("accepted turn")) + message.ContextID = session.ContextId + task := a2a.NewSubmittedTask(message, message) + task.Status.State = a2a.TaskStateWorking + hash := sha256.Sum256([]byte(task.ID)) + _, err = store.CreateRuntimeTask(t.Context(), session.Id, hash[:], task, "") + require.NoError(t, err) + base.actors[actorKey("team-a", substrate.ActorName(session.Id))].Status.State = ateapipb.ActorState_ACTOR_STATE_CRASHED + ids = append(ids, session.Id) + } + for _, id := range ids { + actor, err := base.GetActor(t.Context(), "team-a", substrate.ActorName(id)) + require.NoError(t, err) + _, err = store.BeginCrashedSessionRecovery(t.Context(), id, revertGeneration(t, store, id), actor.GetMetadata().GetUid(), database.SessionTaskSnapshot{}) + require.NoError(t, err) + } + slices.Sort(ids) + after, err := workflow.recoverSessionPage(t.Context(), "") + require.NoError(t, err) + require.Equal(t, ids[99], after, "failed Actor reads still advance the bounded page") + actors.readErr = nil + after, err = workflow.recoverSessionPage(t.Context(), after) + require.NoError(t, err) + require.Empty(t, after, "the short second page resets the cursor") + last, err := base.GetActor(t.Context(), "team-a", substrate.ActorName(ids[100])) + require.NoError(t, err) + require.Equal(t, ateapipb.ActorState_ACTOR_STATE_SUSPENDED, last.GetStatus().GetState(), "the second page recovers the 101st candidate") + firstActor, err := base.GetActor(t.Context(), "team-a", substrate.ActorName(ids[0])) + require.NoError(t, err) + require.Equal(t, ateapipb.ActorState_ACTOR_STATE_CRASHED, firstActor.GetStatus().GetState(), "unavailable first-page Actors still need recovery") + after, err = workflow.recoverSessionPage(t.Context(), after) + require.NoError(t, err) + require.Equal(t, ids[99], after, "the wrapped scan revisits previously unavailable Actors") + after, err = workflow.recoverSessionPage(t.Context(), after) + require.NoError(t, err) + require.Empty(t, after, "an empty page resets the cursor") + for _, id := range ids { + actor, err := base.GetActor(t.Context(), "team-a", substrate.ActorName(id)) + require.NoError(t, err) + require.Equal(t, ateapipb.ActorState_ACTOR_STATE_SUSPENDED, actor.GetStatus().GetState()) + } +} diff --git a/go/core/internal/service/session/workflow.go b/go/core/internal/service/session/workflow.go index 741954f8c5..ea99fa4eb9 100644 --- a/go/core/internal/service/session/workflow.go +++ b/go/core/internal/service/session/workflow.go @@ -15,6 +15,9 @@ import ( ) type workflowStore interface { + ListSessionRecoveryCandidates(context.Context, string, int) ([]*apiv1.Session, error) + SessionLifecycleGeneration(context.Context, string) (uuid.UUID, error) + BeginCrashedSessionRecovery(context.Context, string, uuid.UUID, string, database.SessionTaskSnapshot) (*database.SessionOperation, error) ClaimSessionQuiescence(context.Context) (*database.SessionQuiescence, error) HasPendingSessionQuiescence(context.Context) (bool, error) FinishSessionQuiescence(context.Context, *database.SessionQuiescence, *database.SessionTaskSnapshot) error @@ -30,6 +33,7 @@ type workflowStore interface { type actorClient interface { substrate.LifecycleClient + RevertActor(context.Context, string, string) (*ateapipb.Actor, error) } // ActorWorkflow runs the imperative Substrate operations behind Session @@ -129,6 +133,9 @@ func (w *ActorWorkflow) execute(ctx context.Context, operation *database.Session if operation.Instance.Operation == apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE { return operationOutcome(operation) } + if operation.Instance.State == apiv1.RuntimeState_RUNTIME_STATE_FAILED && operation.Instance.Failure.GetReason() == database.SessionRuntimeCrashed { + return w.recoverCrashedSession(ctx, operation) + } kind := operation.Instance.Operation session := operation.Instance revision, err := w.store.GetRuntimeRevision(ctx, session.GetPreparedRevision()) From 7615ecb0a95ad4f5f50f59841d456d30c92684b0 Mon Sep 17 00:00:00 2001 From: Ada Lundhe Date: Tue, 3 Nov 2026 09:56:56 -0800 Subject: [PATCH 10/12] fix(session): scan crash recoveries on their own ticker, not quiescence worker 0 Start ran the crash-recovery scan inline on quiescence worker 0, between idle claims. A slow idle suspension on that worker held every recovery retry for as long as the suspension took, and a slow revert held up idle work in turn. No test exercised the scan through Start: the existing tests called recoverSessionPage directly, and the Start-level store stub returned no candidates. The scan now runs in its own goroutine with its own ticker, at the same recovery interval. The four quiescence workers return to their shape before recovery existed. The page cursor still carries across ticks, so a backlog larger than one page cannot starve the Sessions past it. TestStartRetriesAdmittedRecoveryWhileIdleWorkBlocks admits a recovery whose first revert fails, then starts the workflow with all four quiescence workers blocked in idle suspensions. The Session must reach READY while they stay blocked. Verified: the new test fails on the previous Start ("Condition never satisfied") and passes with this change. go test for ./internal/service/..., ./internal/substrate/... and ./internal/database/... passes against Postgres; golangci-lint reports 0 issues. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Ada Lundhe --- .../internal/service/session/quiescence.go | 45 +++++++++++------ .../internal/service/session/recovery_test.go | 48 +++++++++++++++++++ 2 files changed, 78 insertions(+), 15 deletions(-) diff --git a/go/core/internal/service/session/quiescence.go b/go/core/internal/service/session/quiescence.go index aba8c40ffb..8d741149e8 100644 --- a/go/core/internal/service/session/quiescence.go +++ b/go/core/internal/service/session/quiescence.go @@ -30,27 +30,18 @@ func (w *ActorWorkflow) wakeQuiescence() { // Start suspends idle sessions to VOLUMES independently of task publication and // retries crash recoveries admitted by explicit requests. -// Startup and periodic scans discover idle work missed by local signals; crash -// checks use the same recovery interval. Workers drain claims with bounded concurrency. +// Startup and periodic scans discover idle work missed by local signals. Workers +// drain claims with bounded concurrency. Crash recovery scans on its own ticker, so +// a slow idle suspension cannot delay a recovery retry and a slow revert cannot +// hold up idle work. func (w *ActorWorkflow) Start(ctx context.Context) error { var workers sync.WaitGroup - for worker := range 4 { + workers.Go(func() { w.recoverCrashedSessions(ctx) }) + for range 4 { workers.Go(func() { timer := time.NewTimer(w.recoveryInterval) defer timer.Stop() - var recoveryAfter string - nextRecovery := time.Now().Add(w.recoveryInterval) for ctx.Err() == nil { - if worker == 0 && !time.Now().Before(nextRecovery) { - scanCtx, cancel := context.WithTimeout(ctx, database.RuntimeOperationTimeout) - after, err := w.recoverSessionPage(scanCtx, recoveryAfter) - cancel() - if err != nil && ctx.Err() == nil { - logging.FromContext(ctx).ErrorContext(ctx, "scan interrupted sessions", "error", err) - } - recoveryAfter = after - nextRecovery = time.Now().Add(w.recoveryInterval) - } work, err := w.store.ClaimSessionQuiescence(ctx) if err == nil { // Hand off a hint before runtime I/O so a single settlement @@ -85,6 +76,30 @@ func (w *ActorWorkflow) Start(ctx context.Context) error { return nil } +// recoverCrashedSessions retries admitted crash recoveries one page per tick. +// The cursor carries across ticks, so a large backlog cannot starve the Sessions +// after its first page, and an Actor that stays unreadable is revisited when the +// scan wraps. +func (w *ActorWorkflow) recoverCrashedSessions(ctx context.Context) { + ticker := time.NewTicker(w.recoveryInterval) + defer ticker.Stop() + var after string + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + scanCtx, cancel := context.WithTimeout(ctx, database.RuntimeOperationTimeout) + next, err := w.recoverSessionPage(scanCtx, after) + cancel() + if err != nil && ctx.Err() == nil { + logging.FromContext(ctx).ErrorContext(ctx, "scan interrupted sessions", "error", err) + } + after = next + } +} + func (w *ActorWorkflow) quiesceIdleSession(ctx context.Context, work *database.SessionQuiescence) { runtimeCtx, cancel := context.WithTimeout(ctx, 2*time.Minute) snapshot, err := w.Quiesce(runtimeCtx, work.Session) diff --git a/go/core/internal/service/session/recovery_test.go b/go/core/internal/service/session/recovery_test.go index d7ce338a83..0a7ef5a51c 100644 --- a/go/core/internal/service/session/recovery_test.go +++ b/go/core/internal/service/session/recovery_test.go @@ -5,6 +5,7 @@ import ( "crypto/sha256" "fmt" "slices" + "sync/atomic" "testing" "time" @@ -539,3 +540,50 @@ func TestSessionRecoveryPagesPastUnavailableActors(t *testing.T) { require.Equal(t, ateapipb.ActorState_ACTOR_STATE_SUSPENDED, actor.GetStatus().GetState()) } } + +// blockedSuspendActors holds every idle suspension in runtime I/O until the +// test ends, the way a slow Substrate upload would. +type blockedSuspendActors struct { + *lifecycleTestActors + entered atomic.Int32 +} + +func (a *blockedSuspendActors) SuspendActor(ctx context.Context, _, _ string) (*ateapipb.Actor, error) { + a.entered.Add(1) + <-ctx.Done() + return nil, ctx.Err() +} + +// Start retries an admitted recovery on its own ticker. Every quiescence worker +// is blocked in an idle suspension for the whole test, so a scan that shared +// those workers would never run. +func TestStartRetriesAdmittedRecoveryWhileIdleWorkBlocks(t *testing.T) { + store, base, crashed, _, _ := crashedTurnFixture(t) + unavailable := &recoveryTestActors{lifecycleTestActors: base, revert: func(context.Context, string, string) (*ateapipb.Actor, error) { + return nil, status.Error(codes.Unavailable, "control plane unavailable") + }} + _, err := NewActorWorkflow(store, unavailable, make(chan struct{}, 1), time.Hour).Revert(t.Context(), crashed) + require.Error(t, err, "the first attempt leaves an admitted recovery for Start to retry") + + const idleSessions = 4 // one per quiescence worker + for range idleSessions { + idle, _, err := store.CreateSession(t.Context(), &apiv1.Session{Id: uuid.NewString(), Creator: crashed.Creator, Agent: crashed.Agent}, uuid.NewString()) + require.NoError(t, err) + idle, err = NewActorWorkflow(store, base, make(chan struct{}, 1), time.Hour).Create(t.Context(), idle) + require.NoError(t, err) + request := stageQuiescenceTask(t, store.Client, idle, a2a.TaskStateCompleted) + require.NoError(t, store.SettleSessionTask(t.Context(), idle.Id, request.TaskId, request.Version)) + } + + actors := &blockedSuspendActors{lifecycleTestActors: base} + startQuiescenceWorker(t, NewActorWorkflow(store, actors, make(chan struct{}, 1), 20*time.Millisecond)) + require.Eventually(t, func() bool { return actors.entered.Load() == idleSessions }, 5*time.Second, 10*time.Millisecond) + require.Eventually(t, func() bool { + current, err := store.GetSessionByID(t.Context(), crashed.Id) + return err == nil && current.State == apiv1.RuntimeState_RUNTIME_STATE_READY + }, 5*time.Second, 10*time.Millisecond) + actor, err := base.GetActor(t.Context(), "team-a", substrate.ActorName(crashed.Id)) + require.NoError(t, err) + require.Equal(t, ateapipb.ActorState_ACTOR_STATE_SUSPENDED, actor.GetStatus().GetState()) + require.EqualValues(t, idleSessions, actors.entered.Load(), "idle work stayed blocked while recovery completed") +} From e87db8d3f42d2fc0444ec24bf6deace995397f17 Mon Sep 17 00:00:00 2001 From: Ada Lundhe Date: Tue, 3 Nov 2026 09:59:12 -0800 Subject: [PATCH 11/12] fix(session): revert crashed runtimes through execute, reading before the claim Crash recovery had its own copy of execute: claim, GetActor, verify, revert, finish. It claimed before its reads, and a released claim keeps executor_id. One failed GetActor therefore made Delete refuse the Session with "begin Session operation: resource conflict" until recovery succeeded. A deleted Actor made that permanent. Execute also sent every operation on a crashed Session into the recovery copy, so even a Delete that superseded the recovery would have returned a conflict. Recovery now runs through execute like suspend, resume and delete. substrate.PrepareActorRevert does the read-only checks before the claim, and ApplyActorTransition issues RevertActor, which moves onto LifecycleClient. The checks are unchanged: the Actor keeps its UID, ends SUSPENDED, and restores the durable snapshot observed before the revert. A SUSPENDED Actor still completes a recovery whose revert response was lost. A failed preparation keeps the RESUME intent instead of clearing it, because admission already failed the interrupted turn. Nothing was issued, so Delete can supersede it. Rejected: keeping the separate path and only moving its reads above the claim. That fixes the stuck Delete, but fixes to execute would still have to land twice. Verified: TestDeleteSupersedesUnissuedCrashRecovery covers an unreadable and a missing Actor; both fail on the previous code with "resource conflict" and pass here. TestActorRevertRestoresTheObservedSnapshotInPlace covers the revert transition's states and its UID and snapshot checks. go test for ./internal/service/..., ./internal/substrate/... and ./internal/database/... passes against Postgres; golangci-lint reports 0 issues. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Ada Lundhe --- go/core/internal/service/session/recovery.go | 63 --------------- .../internal/service/session/recovery_test.go | 48 ++++++++++++ go/core/internal/service/session/workflow.go | 31 ++++++-- go/core/internal/substrate/lifecycle.go | 51 ++++++++++++ go/core/internal/substrate/lifecycle_test.go | 77 +++++++++++++++++++ 5 files changed, 199 insertions(+), 71 deletions(-) diff --git a/go/core/internal/service/session/recovery.go b/go/core/internal/service/session/recovery.go index 27922bc3ed..cecde11d45 100644 --- a/go/core/internal/service/session/recovery.go +++ b/go/core/internal/service/session/recovery.go @@ -4,11 +4,8 @@ import ( "context" "errors" "fmt" - "time" - "github.com/google/uuid" ateapipb "go.ate.dev/substrate/api/ateapi/v1" - "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" @@ -92,63 +89,3 @@ func (w *ActorWorkflow) Revert(ctx context.Context, session *apiv1.Session) (*ap } return w.execute(ctx, operation) } - -// recoverCrashedSession reconciles lost revert responses without replacing the Actor. -func (w *ActorWorkflow) recoverCrashedSession(ctx context.Context, operation *database.SessionOperation) (_ *apiv1.Session, err error) { - session := operation.Instance - if session.Operation != apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME { - return nil, database.ErrConflict - } - executorID := uuid.New() - claimed, err := w.store.ClaimSessionOperation(ctx, session.Id, operation.ID, executorID) - if err != nil { - return nil, err - } - if !claimed { - current, err := w.store.GetSessionOperation(ctx, session.Id, operation.ID) - if err != nil { - return nil, err - } - return operationOutcome(current) - } - defer func() { - if err != nil { - releaseCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) - defer cancel() - err = errors.Join(err, w.store.ReleaseRuntimeOperation(releaseCtx, session.Id, operation.ID, executorID)) - } - }() - revision, err := w.store.GetRuntimeRevision(ctx, session.PreparedRevision) - if err != nil { - return nil, err - } - atespace, name := revision.ActorTemplateAtespace, substrate.ActorName(session.Id) - actor, err := w.actors.GetActor(ctx, atespace, name) - if err != nil { - return nil, err - } - if err := w.verifyActor(ctx, session, revision, actor); err != nil { - return nil, err - } - before := actor - beforeSnapshot := proto.CloneOf(substrate.LatestDurableActorSnapshot(actor)) - if beforeSnapshot != nil { - beforeSnapshot.Locality = "" - } - switch actor.GetStatus().GetState() { - case ateapipb.ActorState_ACTOR_STATE_CRASHED, ateapipb.ActorState_ACTOR_STATE_REVERTING: - actor, err = w.actors.RevertActor(ctx, atespace, name) - if err != nil { - return nil, fmt.Errorf("revert crashed Actor %s/%s: %w", atespace, name, err) - } - case ateapipb.ActorState_ACTOR_STATE_SUSPENDED: - default: - return nil, fmt.Errorf("crashed Actor %s/%s has unexpected state %s", atespace, name, actor.GetStatus().GetState()) - } - if !validActorIdentity(actor, revision, name) || actor.GetMetadata().GetUid() != before.GetMetadata().GetUid() || actor.GetStatus().GetState() != ateapipb.ActorState_ACTOR_STATE_SUSPENDED || !proto.Equal(beforeSnapshot, substrate.LatestActorSnapshot(actor)) { - return nil, fmt.Errorf("revert Actor %s/%s returned unexpected identity or state", atespace, name) - } - finishCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) - defer cancel() - return w.store.FinishSessionOperation(finishCtx, session.Id, operation.ID, executorID, "", actor.GetMetadata().GetUid(), "") -} diff --git a/go/core/internal/service/session/recovery_test.go b/go/core/internal/service/session/recovery_test.go index 0a7ef5a51c..3769f82e22 100644 --- a/go/core/internal/service/session/recovery_test.go +++ b/go/core/internal/service/session/recovery_test.go @@ -587,3 +587,51 @@ func TestStartRetriesAdmittedRecoveryWhileIdleWorkBlocks(t *testing.T) { require.Equal(t, ateapipb.ActorState_ACTOR_STATE_SUSPENDED, actor.GetStatus().GetState()) require.EqualValues(t, idleSessions, actors.entered.Load(), "idle work stayed blocked while recovery completed") } + +// A recovery whose Actor cannot be read has issued nothing, so it must not take +// a claim that would make Delete refuse the Session. +func TestDeleteSupersedesUnissuedCrashRecovery(t *testing.T) { + for _, test := range []struct { + name string + gone bool + readErr error + }{ + {name: "actor read fails", readErr: status.Error(codes.Unavailable, "control plane unavailable")}, + {name: "actor is gone", gone: true}, + } { + t.Run(test.name, func(t *testing.T) { + store, base, session, _, _ := crashedTurnFixture(t) + key := actorKey("team-a", substrate.ActorName(session.Id)) + _, err := store.BeginCrashedSessionRecovery(t.Context(), session.Id, revertGeneration(t, store, session.Id), base.actors[key].GetMetadata().GetUid(), database.SessionTaskSnapshot{}) + require.NoError(t, err) + if test.gone { + delete(base.actors, key) + } + actors := &recoveryTestActors{lifecycleTestActors: base, readErr: test.readErr} + actors.revert = func(context.Context, string, string) (*ateapipb.Actor, error) { + t.Fatal("an unread Actor was reverted") + return nil, nil + } + workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute) + admitted, err := store.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + _, err = workflow.Revert(t.Context(), admitted) + require.Error(t, err) + _, err = workflow.recoverSessionPage(t.Context(), "") + require.NoError(t, err) + + pending, err := store.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_FAILED, pending.State) + require.Equal(t, apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME, pending.Operation, "a failed read keeps the recovery for the next scan") + deleted, err := NewActorWorkflow(store, base, make(chan struct{}, 1), time.Minute).Delete(t.Context(), pending) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_DELETED, deleted.State) + _, err = base.GetActor(t.Context(), "team-a", substrate.ActorName(session.Id)) + require.Equal(t, codes.NotFound, status.Code(err)) + page, err := store.ListSessionRecoveryCandidates(t.Context(), "", 100) + require.NoError(t, err) + require.Empty(t, page) + }) + } +} diff --git a/go/core/internal/service/session/workflow.go b/go/core/internal/service/session/workflow.go index ea99fa4eb9..8698bf9adb 100644 --- a/go/core/internal/service/session/workflow.go +++ b/go/core/internal/service/session/workflow.go @@ -33,7 +33,6 @@ type workflowStore interface { type actorClient interface { substrate.LifecycleClient - RevertActor(context.Context, string, string) (*ateapipb.Actor, error) } // ActorWorkflow runs the imperative Substrate operations behind Session @@ -133,9 +132,6 @@ func (w *ActorWorkflow) execute(ctx context.Context, operation *database.Session if operation.Instance.Operation == apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE { return operationOutcome(operation) } - if operation.Instance.State == apiv1.RuntimeState_RUNTIME_STATE_FAILED && operation.Instance.Failure.GetReason() == database.SessionRuntimeCrashed { - return w.recoverCrashedSession(ctx, operation) - } kind := operation.Instance.Operation session := operation.Instance revision, err := w.store.GetRuntimeRevision(ctx, session.GetPreparedRevision()) @@ -164,11 +160,15 @@ func (w *ActorWorkflow) execute(ctx context.Context, operation *database.Session creation.Snapshot = &substrate.ActorSnapshot{Tag: &ateapipb.ObjectRef{Atespace: snapshot.Atespace, Name: tagName}, URI: snapshot.URI, Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES} } } - prepare := substrate.PrepareActorTransition - if operation.ExecutorID != uuid.Nil { - prepare = substrate.PrepareActorRetry + var transition *substrate.ActorTransition + switch { + case crashRecovery(operation): + transition, err = substrate.PrepareActorRevert(ctx, w.actors, binding) + case operation.ExecutorID != uuid.Nil: + transition, err = substrate.PrepareActorRetry(ctx, w.actors, binding, kind, creation) + default: + transition, err = substrate.PrepareActorTransition(ctx, w.actors, binding, kind, creation) } - transition, err := prepare(ctx, w.actors, binding, kind, creation) if err != nil { return w.failPreparation(ctx, operation, err) } @@ -221,6 +221,12 @@ func (w *ActorWorkflow) failPreparation(ctx context.Context, admitted *database. if admitted.ExecutorID != uuid.Nil { return nil, cause // An earlier attempt may have issued work; retain its intent. } + if crashRecovery(admitted) { + // Admission already failed the interrupted turn, so clearing the operation + // would leave a FAILED Session that nothing retries. Keep the intent for the + // next scan. No claim was taken, so Delete can still supersede it. + return nil, cause + } finishCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) defer cancel() _, err := w.store.FinishSessionOperation(finishCtx, admitted.Instance.Id, admitted.ID, uuid.Nil, "", "", "lifecycle preparation failed") @@ -234,6 +240,15 @@ func (w *ActorWorkflow) failPreparation(ctx context.Context, admitted *database. return nil, errors.Join(cause, err) } +// crashRecovery reports a RESUME admitted after Substrate confirmed a crash. +// It reverts the recorded Actor in place instead of resuming it. +func crashRecovery(operation *database.SessionOperation) bool { + session := operation.Instance + return session.State == apiv1.RuntimeState_RUNTIME_STATE_FAILED && + session.Operation == apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME && + session.Failure.GetReason() == database.SessionRuntimeCrashed +} + func operationOutcome(operation *database.SessionOperation) (*apiv1.Session, error) { if operation.Instance.Operation == apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE { return operation.Instance, nil diff --git a/go/core/internal/substrate/lifecycle.go b/go/core/internal/substrate/lifecycle.go index 49450c2830..773de26c6b 100644 --- a/go/core/internal/substrate/lifecycle.go +++ b/go/core/internal/substrate/lifecycle.go @@ -19,6 +19,7 @@ type LifecycleClient interface { CheckActorEgressPolicy(context.Context, string, string) error EnsureActorEgressPolicy(context.Context, string, string, string) error ResumeActor(context.Context, string, string) (*ateapipb.Actor, error) + RevertActor(context.Context, string, string) (*ateapipb.Actor, error) SuspendActor(context.Context, string, string) (*ateapipb.Actor, error) DeleteActor(context.Context, string, string) error } @@ -51,6 +52,9 @@ type ActorTransition struct { operation apiv1.RuntimeOperation actor *ateapipb.Actor creation *ActorCreation + // revert restores a crashed Actor to its last durable snapshot. The + // operation stays RESUME because that is what the Session records. + revert bool } // ActorUID returns the observed runtime identity, including after creation. @@ -118,6 +122,26 @@ func prepareActorTransition(ctx context.Context, actors LifecycleClient, binding return &ActorTransition{binding: binding, operation: operation, actor: actor, creation: creation}, nil } +// PrepareActorRevert performs the read-only checks for crash recovery before +// the store's durable claim. Reading first keeps a failed or missing Actor from +// leaving a claim behind, so Delete can still supersede the unissued recovery. +// A SUSPENDED Actor means an earlier revert completed and its response was lost. +func PrepareActorRevert(ctx context.Context, actors LifecycleClient, binding ActorBinding) (*ActorTransition, error) { + actor, err := actors.GetActor(ctx, binding.Atespace, binding.Name) + if err != nil { + return nil, fmt.Errorf("get Actor %s/%s: %w", binding.Atespace, binding.Name, err) + } + if !binding.Matches(actor) { + return nil, fmt.Errorf("actor %s/%s identity or template changed", binding.Atespace, binding.Name) + } + switch actor.GetStatus().GetState() { + case ateapipb.ActorState_ACTOR_STATE_CRASHED, ateapipb.ActorState_ACTOR_STATE_REVERTING, ateapipb.ActorState_ACTOR_STATE_SUSPENDED: + default: + return nil, fmt.Errorf("crashed Actor %s/%s has unexpected state %s", binding.Atespace, binding.Name, actor.GetStatus().GetState()) + } + return &ActorTransition{binding: binding, operation: apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME, actor: actor, revert: true}, nil +} + // ApplyActorTransition requires durable execution authority. Errors retain the // claim's durable intent and pins. Another bounded attempt can retry that intent. func ApplyActorTransition(ctx context.Context, actors LifecycleClient, transition *ActorTransition) error { @@ -125,6 +149,9 @@ func ApplyActorTransition(ctx context.Context, actors LifecycleClient, transitio return err } binding, actor := transition.binding, transition.actor + if transition.revert { + return transition.applyRevert(ctx, actors) + } var err error switch transition.operation { case apiv1.RuntimeOperation_RUNTIME_OPERATION_CREATE: @@ -202,6 +229,30 @@ func ApplyActorTransition(ctx context.Context, actors LifecycleClient, transitio return nil } +// applyRevert must not replace the Actor or roll it past the snapshot observed +// before the claim. A changed UID or snapshot means something other than this +// recovery acted on the runtime, so the Session must not adopt the result. +func (t *ActorTransition) applyRevert(ctx context.Context, actors LifecycleClient) error { + binding, actor := t.binding, t.actor + before := proto.CloneOf(LatestDurableActorSnapshot(actor)) + if before != nil { + // Revert restores onto any node, so the locality of the source copy is + // expected to change. + before.Locality = "" + } + if state := actor.GetStatus().GetState(); state == ateapipb.ActorState_ACTOR_STATE_CRASHED || state == ateapipb.ActorState_ACTOR_STATE_REVERTING { + reverted, err := actors.RevertActor(ctx, binding.Atespace, binding.Name) + if err != nil { + return fmt.Errorf("revert crashed Actor %s/%s: %w", binding.Atespace, binding.Name, err) + } + actor = reverted + } + if t.checkResult(actor, ateapipb.ActorState_ACTOR_STATE_SUSPENDED) != nil || !proto.Equal(before, LatestActorSnapshot(actor)) { + return fmt.Errorf("revert Actor %s/%s returned unexpected identity or state", binding.Atespace, binding.Name) + } + return nil +} + // Matches checks that the actor belongs to the expected runtime and template. func (b ActorBinding) Matches(actor *ateapipb.Actor) bool { return actor.GetMetadata().GetName() == b.Name && actor.GetMetadata().GetAtespace() == b.Atespace && actor.GetMetadata().GetUid() != "" && diff --git a/go/core/internal/substrate/lifecycle_test.go b/go/core/internal/substrate/lifecycle_test.go index d712914d9f..70efdc88cc 100644 --- a/go/core/internal/substrate/lifecycle_test.go +++ b/go/core/internal/substrate/lifecycle_test.go @@ -23,6 +23,8 @@ type transitionActors struct { checkErr error resumeUID string suspendSnapshot *ateapipb.Snapshot + revertUID string + revertSnapshot *ateapipb.Snapshot } func (a *transitionActors) GetActor(context.Context, string, string) (*ateapipb.Actor, error) { @@ -70,6 +72,26 @@ func (a *transitionActors) SuspendActor(context.Context, string, string) (*ateap return proto.CloneOf(a.actor), nil } +// RevertActor restores the latest durable snapshot unless the test overrides +// the identity or snapshot it returns. +func (a *transitionActors) RevertActor(context.Context, string, string) (*ateapipb.Actor, error) { + a.calls = append(a.calls, "revert") + a.actor.Status.State = ateapipb.ActorState_ACTOR_STATE_SUSPENDED + snapshot := proto.CloneOf(LatestDurableActorSnapshot(a.actor)) + if a.revertSnapshot != nil { + snapshot = proto.CloneOf(a.revertSnapshot) + } + a.actor.Status.Snapshots = nil + if snapshot != nil { + snapshot.Locality = "" + a.actor.Status.Snapshots = []*ateapipb.Snapshot{snapshot} + } + if a.revertUID != "" { + a.actor.Metadata.Uid = a.revertUID + } + return proto.CloneOf(a.actor), nil +} + func (a *transitionActors) DeleteActor(context.Context, string, string) error { a.calls = append(a.calls, "delete") a.actor = nil @@ -230,3 +252,58 @@ func TestActorSuspendRequiresDurableVolumesSnapshot(t *testing.T) { }) } } + +func TestActorRevertRestoresTheObservedSnapshotInPlace(t *testing.T) { + durable := func(generation int32, uri string) *ateapipb.Snapshot { + return &ateapipb.Snapshot{Generation: generation, Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES, Locality: "lost-node", + DurableSnapshot: &ateapipb.SnapshotStorage{Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_COMPLETED, Object: &ateapipb.ObjectSnapshot{SnapshotUri: uri}}} + } + for _, test := range []struct { + name string + missing bool + state ateapipb.ActorState + revertUID string + revertSnapshot *ateapipb.Snapshot + wantPrepareErr bool + wantApplyErr bool + wantCalls []string + }{ + {name: "crashed", state: ateapipb.ActorState_ACTOR_STATE_CRASHED, wantCalls: []string{"revert"}}, + {name: "revert in progress", state: ateapipb.ActorState_ACTOR_STATE_REVERTING, wantCalls: []string{"revert"}}, + {name: "earlier revert response lost", state: ateapipb.ActorState_ACTOR_STATE_SUSPENDED}, + {name: "running actor", state: ateapipb.ActorState_ACTOR_STATE_RUNNING, wantPrepareErr: true}, + {name: "actor gone", missing: true, wantPrepareErr: true}, + {name: "replaced actor", state: ateapipb.ActorState_ACTOR_STATE_CRASHED, revertUID: "another-actor", wantApplyErr: true, wantCalls: []string{"revert"}}, + {name: "different snapshot", state: ateapipb.ActorState_ACTOR_STATE_CRASHED, revertSnapshot: durable(1, "s3://snapshots/older"), wantApplyErr: true, wantCalls: []string{"revert"}}, + } { + t.Run(test.name, func(t *testing.T) { + actors := &transitionActors{revertUID: test.revertUID, revertSnapshot: test.revertSnapshot} + if !test.missing { + actors.actor = &ateapipb.Actor{ + Metadata: &ateapipb.ResourceMetadata{Atespace: "team", Name: "instance", Uid: "actor-uid"}, + ActorTemplate: &ateapipb.ObjectRef{Atespace: "team", Name: "revision"}, + Status: &ateapipb.ActorStatus{State: test.state, Snapshots: []*ateapipb.Snapshot{durable(2, "s3://snapshots/latest")}}, + } + if test.state == ateapipb.ActorState_ACTOR_STATE_SUSPENDED { + actors.actor.Status.Snapshots[0].Locality = "" + } + } + binding := ActorBinding{Revision: "revision", Atespace: "team", Name: "instance", TemplateAtespace: "team", TemplateName: "revision"} + transition, err := PrepareActorRevert(t.Context(), actors, binding) + if test.wantPrepareErr { + require.Error(t, err) + require.Empty(t, actors.calls, "a failed read issues nothing") + return + } + require.NoError(t, err) + require.Equal(t, "actor-uid", transition.ActorUID()) + err = ApplyActorTransition(t.Context(), actors, transition) + if test.wantApplyErr { + require.Error(t, err) + } else { + require.NoError(t, err) + } + require.Equal(t, test.wantCalls, actors.calls) + }) + } +} From 72d68a3966cc80073ed7d88972676b856e75cfc2 Mon Sep 17 00:00:00 2001 From: Ada Lundhe Date: Wed, 4 Nov 2026 08:39:22 -0800 Subject: [PATCH 12/12] test(session): use SnapshotStorage state in recovery tests Main renamed SnapshotStorage.status to state (SnapshotStorageState). Update the recovery and substrate snapshot tests to the new field. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Ada Lundhe --- go/core/internal/service/session/recovery_test.go | 6 +++--- go/core/internal/substrate/lifecycle_test.go | 2 +- go/core/internal/substrate/snapshot_test.go | 6 +++--- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/go/core/internal/service/session/recovery_test.go b/go/core/internal/service/session/recovery_test.go index 3769f82e22..842c4a3b98 100644 --- a/go/core/internal/service/session/recovery_test.go +++ b/go/core/internal/service/session/recovery_test.go @@ -434,11 +434,11 @@ func TestExplicitRevertPreservesOnlyCapturedWaits(t *testing.T) { require.NoError(t, store.FinishSessionQuiescence(t.Context(), work, &database.SessionTaskSnapshot{ Atespace: "team-a", URI: "s3://snapshots/waiting", Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES.String(), })) - actors.actors[actorKey("team-a", substrate.ActorName(session.Id))].Status.Snapshots = []*ateapipb.Snapshot{{Generation: 1, Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES, DurableSnapshot: &ateapipb.SnapshotStorage{Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_COMPLETED, Object: &ateapipb.ObjectSnapshot{SnapshotUri: "s3://snapshots/waiting"}}}} + actors.actors[actorKey("team-a", substrate.ActorName(session.Id))].Status.Snapshots = []*ateapipb.Snapshot{{Generation: 1, Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES, DurableSnapshot: &ateapipb.SnapshotStorage{State: ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_COMPLETED, Object: &ateapipb.ObjectSnapshot{SnapshotUri: "s3://snapshots/waiting"}}}} actor := actors.actors[actorKey("team-a", substrate.ActorName(session.Id))] actor.Status.Snapshots[0].Locality = "lost-node" actor.Status.Snapshots = append(actor.Status.Snapshots, &ateapipb.Snapshot{ - Generation: 2, DurableSnapshot: &ateapipb.SnapshotStorage{Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_PENDING}, + Generation: 2, DurableSnapshot: &ateapipb.SnapshotStorage{State: ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_PENDING}, }) } _, err = NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute).Revert(t.Context(), session) @@ -475,7 +475,7 @@ func TestExplicitRevertRetainsWaitCapturedByLaterTurn(t *testing.T) { require.NoError(t, store.FinishSessionQuiescence(t.Context(), work, &database.SessionTaskSnapshot{ Atespace: "team-a", URI: "s3://snapshots/later-turn", Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES.String(), })) - actors.actors[actorKey("team-a", substrate.ActorName(session.Id))].Status.Snapshots = []*ateapipb.Snapshot{{Generation: 1, Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES, DurableSnapshot: &ateapipb.SnapshotStorage{Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_COMPLETED, Object: &ateapipb.ObjectSnapshot{SnapshotUri: "s3://snapshots/later-turn"}}}} + actors.actors[actorKey("team-a", substrate.ActorName(session.Id))].Status.Snapshots = []*ateapipb.Snapshot{{Generation: 1, Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES, DurableSnapshot: &ateapipb.SnapshotStorage{State: ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_COMPLETED, Object: &ateapipb.ObjectSnapshot{SnapshotUri: "s3://snapshots/later-turn"}}}} _, err = NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute).Revert(t.Context(), session) require.NoError(t, err) retained, err := store.GetSettledSessionTask(t.Context(), session.Id, string(waiting.ID), nil) diff --git a/go/core/internal/substrate/lifecycle_test.go b/go/core/internal/substrate/lifecycle_test.go index 70efdc88cc..21e9eaa3dd 100644 --- a/go/core/internal/substrate/lifecycle_test.go +++ b/go/core/internal/substrate/lifecycle_test.go @@ -256,7 +256,7 @@ func TestActorSuspendRequiresDurableVolumesSnapshot(t *testing.T) { func TestActorRevertRestoresTheObservedSnapshotInPlace(t *testing.T) { durable := func(generation int32, uri string) *ateapipb.Snapshot { return &ateapipb.Snapshot{Generation: generation, Fidelity: ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES, Locality: "lost-node", - DurableSnapshot: &ateapipb.SnapshotStorage{Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_COMPLETED, Object: &ateapipb.ObjectSnapshot{SnapshotUri: uri}}} + DurableSnapshot: &ateapipb.SnapshotStorage{State: ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_COMPLETED, Object: &ateapipb.ObjectSnapshot{SnapshotUri: uri}}} } for _, test := range []struct { name string diff --git a/go/core/internal/substrate/snapshot_test.go b/go/core/internal/substrate/snapshot_test.go index dfda60e0cf..4e923131a8 100644 --- a/go/core/internal/substrate/snapshot_test.go +++ b/go/core/internal/substrate/snapshot_test.go @@ -43,13 +43,13 @@ func TestActorSnapshotGenerationAndDurability(t *testing.T) { func TestRevertSnapshotSkipsIncompleteGenerations(t *testing.T) { durable := &ateapipb.Snapshot{Generation: 2, DurableSnapshot: &ateapipb.SnapshotStorage{ - Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_COMPLETED, + State: ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_COMPLETED, Object: &ateapipb.ObjectSnapshot{SnapshotUri: "s3://snapshots/durable"}, }} for _, latest := range []*ateapipb.Snapshot{ {Generation: 3, Locality: "node-1"}, - {Generation: 3, DurableSnapshot: &ateapipb.SnapshotStorage{Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_PENDING}}, - {Generation: 3, DurableSnapshot: &ateapipb.SnapshotStorage{Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_COMPLETED}}, + {Generation: 3, DurableSnapshot: &ateapipb.SnapshotStorage{State: ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_PENDING}}, + {Generation: 3, DurableSnapshot: &ateapipb.SnapshotStorage{State: ateapipb.SnapshotStorageState_SNAPSHOT_STORAGE_STATE_COMPLETED}}, } { actor := &ateapipb.Actor{Status: &ateapipb.ActorStatus{Snapshots: []*ateapipb.Snapshot{latest, durable}}} require.Same(t, durable, LatestDurableActorSnapshot(actor))