From 2b34722b26ae5ed0262332bc4f1695f620097331 Mon Sep 17 00:00:00 2001 From: Shashank Ram Date: Tue, 10 Nov 2026 14:14:33 -0800 Subject: [PATCH 1/2] runtime/session: suspend actors after abandoned delivery \# Description Canceled delivery can leave a restored actor running without a turn, holding worker capacity indefinitely. - Retain expired dispatch attempts for the existing lifecycle worker to process. - Claim suspension only when no active turn or newer delivery owns the Session. - Preserve snapshots and logical readiness so subsequent messages resume normally. - Cover cancellation, competing lifecycle work, retries, and subsequent delivery. \# Change Type ``` /kind fix ``` \# Changelog ```release-note Actors left running after canceled message delivery now suspend and release worker capacity automatically. ``` \# Additional Notes Signed-off-by: Shashank Ram --- docs/architecture/runtime-and-lifecycle.md | 27 +- go/core/internal/a2agateway/routing_test.go | 2 +- .../scheduledrun/controller_test.go | 2 +- go/core/internal/database/session_deletion.go | 3 +- .../database/session_dispatch_cleanup.go | 202 ++++++++ .../internal/database/session_operations.go | 9 + .../internal/database/session_quiescence.go | 35 +- go/core/internal/database/session_recovery.go | 5 +- go/core/internal/database/session_tasks.go | 28 +- .../internal/database/task_finalization.go | 5 +- go/core/internal/database/task_store_test.go | 438 ++++++++++++++++++ .../scheduledrun_controller_test.go | 2 +- .../internal/grpcserver/task_store_test.go | 4 +- .../service/session/expiration_test.go | 2 +- .../internal/service/session/interactions.go | 19 +- .../service/session/interactions_test.go | 14 +- .../internal/service/session/quiescence.go | 105 ++++- .../service/session/quiescence_test.go | 92 ++++ .../service/session/quiescence_wake_test.go | 283 +++++++++++ go/core/internal/service/session/workflow.go | 26 +- go/core/pkg/app/app.go | 2 +- .../core/000004_session_dispatch_cleanup.sql | 22 + go/core/pkg/migrations/runner_test.go | 24 +- go/core/test/byoa2a/main.go | 60 +++ go/core/test/e2e/README.md | 10 + .../e2e/session_resuming_deletion_test.go | 85 ++++ 26 files changed, 1451 insertions(+), 55 deletions(-) create mode 100644 go/core/internal/database/session_dispatch_cleanup.go create mode 100644 go/core/pkg/migrations/core/000004_session_dispatch_cleanup.sql diff --git a/docs/architecture/runtime-and-lifecycle.md b/docs/architecture/runtime-and-lifecycle.md index a76cc94697..8312db1816 100644 --- a/docs/architecture/runtime-and-lifecycle.md +++ b/docs/architecture/runtime-and-lifecycle.md @@ -208,6 +208,31 @@ leave an actor running until a surviving replica's next recovery scan; increasin the interval increases that delay and the time until its snapshot is available. Ordinary committed settlements do not wait for the recovery interval. +The existing Session dispatch ID and expiry also track abandoned delivery. +Revocation expires the reservation immediately, retains its ID, and wakes the +same workers. An unrevoked reservation becomes eligible at its existing expiry. +This covers cancellation after actor restore but before a turn starts, when no +task boundary exists to trigger suspension. Accepted execution clears the +reservation; saving continuation input without leaving a waiting state retains +it. An expired reservation cannot authorize a late first task write. + +For abandoned dispatches, the worker first reads and verifies the actor without +blocking new messages. Only a RUNNING actor can be claimed for suspension. The +claim atomically rechecks the dispatch generation, absence of active tasks, and +other lifecycle work. A newer message supersedes unclaimed cleanup; a claimed +suspension blocks new execution until its snapshot is recorded and the dispatch +reservation and claim are cleared. Cleanup preserves +the existing history and logical READY state. When history exists, its captured +head receives the snapshot without adding a task event. + +RESUMING or SUSPENDED actors and failed preparation reads remain unclaimed. Each +worker visits candidates in Session ID order once per scan, then revisits them on +the existing worker wakeups and recovery scans. A SUSPENDED observation alone +cannot prove that canceled ingress will not finish restoring later. Such cleanup +can remain pending until an actor runs, a new message supersedes it, or an explicit +lifecycle operation retires it. This is cancellation cleanup, not general idle +actor detection. + An empty claim is followed by a check for unclaimed work on ready sessions. If dispatch, checkpoint creation, a lifecycle operation, or a row lock blocks that work, workers retry within one second until it becomes eligible or is superseded. @@ -215,7 +240,7 @@ Database errors also use this short retry. Sessions that are not READY and claim work do not keep the short retry active. Idle workers make one claim attempt and one pending-work check per recovery interval, rather than querying every second. Other background workers and connection-pool idle settings still affect whether -PostgreSQL can become idle. +PostgreSQL can become idle. Each scan also checks for abandoned dispatch cleanup. Unclaimed idle work survives API restarts. A claim for possibly issued runtime work never expires: losing the worker does not prove that the suspend stopped. Uncertain diff --git a/go/core/internal/a2agateway/routing_test.go b/go/core/internal/a2agateway/routing_test.go index 7a062f307d..7daa71808e 100644 --- a/go/core/internal/a2agateway/routing_test.go +++ b/go/core/internal/a2agateway/routing_test.go @@ -54,7 +54,7 @@ func newTestSessions(store *gatewayTestStore, authorizer auth.Authorizer) *sessi } func newTestGateway(store *gatewayTestStore, authorizer auth.Authorizer, dialer *gatewayTestDialer, url string) a2asrv.RequestHandler { - interactions := sessionsvc.NewInteractionService(store, gatewayTestAgents{store, authorizer}, newTestSessions(store, authorizer), false) + interactions := sessionsvc.NewInteractionService(store, gatewayTestAgents{store, authorizer}, newTestSessions(store, authorizer), false, nil) return New(interactions, dialer, url) } diff --git a/go/core/internal/controller/scheduledrun/controller_test.go b/go/core/internal/controller/scheduledrun/controller_test.go index 77ef1b23e7..b1a79d9bc0 100644 --- a/go/core/internal/controller/scheduledrun/controller_test.go +++ b/go/core/internal/controller/scheduledrun/controller_test.go @@ -338,7 +338,7 @@ func TestControllerDispatchUsesBoundedInProcessAuthority(t *testing.T) { store := &scopedControllerStore{session: &apiv1.Session{Id: id, ContextId: id, Creator: "owner", Agent: &apiv1.AgentProvenance{Ref: &apiv1.ResourceReference{Namespace: "team", Name: "assistant"}}, State: apiv1.RuntimeState_RUNTIME_STATE_READY, Operation: apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE}} policy := &controllerSessionPolicy{} sessions := sessionsvc.NewService(store, policy, nil) - interactions := sessionsvc.NewInteractionService(store, nil, sessions, false) + interactions := sessionsvc.NewInteractionService(store, nil, sessions, false, nil) ctx, cancel := context.WithTimeout(t.Context(), 2*time.Second) defer cancel() gateway := scopedControllerGateway{interactions: interactions, cancel: cancel, agent: types.NamespacedName{Namespace: "team", Name: "assistant"}} diff --git a/go/core/internal/database/session_deletion.go b/go/core/internal/database/session_deletion.go index 7b7b22dd46..49f6c3fc3b 100644 --- a/go/core/internal/database/session_deletion.go +++ b/go/core/internal/database/session_deletion.go @@ -48,7 +48,7 @@ func beginSessionDeletion(ctx context.Context, tx pgx.Tx, row sessionRow, reason if err != nil { return nil, err } - if err := execSQL(ctx, tx, `UPDATE session SET dispatch_id = NULL, dispatch_expires_at = NULL WHERE id = $1`, row.ID); err != nil { + if err := retireSessionDispatchCleanup(ctx, tx, row.ID); err != nil { return nil, err } if err := execSQL(ctx, tx, `UPDATE session SET deletion_reason = COALESCE(deletion_reason, $2) WHERE id = $1`, row.ID, string(reason)); err != nil { @@ -68,6 +68,7 @@ func admitForcedSessionDeletion(ctx context.Context, tx pgx.Tx, row sessionRow) SELECT EXISTS (SELECT 1 FROM session_checkpoint WHERE source_session_id = $1 AND state = 'CREATING') OR EXISTS (SELECT 1 FROM session_task_event WHERE history_id = $2 AND quiescence_pending AND quiescence_executor_id IS NOT NULL) + OR EXISTS (SELECT 1 FROM session WHERE id = $1 AND dispatch_cleanup_executor_id IS NOT NULL) `, pgx.RowTo[bool], row.ID, row.HistoryID) if err != nil { return err diff --git a/go/core/internal/database/session_dispatch_cleanup.go b/go/core/internal/database/session_dispatch_cleanup.go new file mode 100644 index 0000000000..a43577bd28 --- /dev/null +++ b/go/core/internal/database/session_dispatch_cleanup.go @@ -0,0 +1,202 @@ +package database + +import ( + "context" + "errors" + "fmt" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" +) + +// SessionDispatchCleanup identifies an abandoned delivery's runtime cleanup. +// An unclaimed candidate grants no execution authority. Claim returns a new +// value whose executor fences runtime work until its known outcome is stored. +type SessionDispatchCleanup struct { + Session *apiv1.Session + DispatchID uuid.UUID + ActorUID string + ExecutorID uuid.UUID + boundary *dispatchCleanupBoundary +} + +type dispatchCleanupBoundary struct { + TaskID string + Version int64 + State string +} + +// GetPendingSessionDispatchCleanup observes the next revoked or expired delivery +// after the Session cursor without holding locks during runtime preparation. +// Missing or blocked work returns ErrNotFound. Claim revalidates the generation; +// advancing the cursor lets a worker skip an actor that is not ready to suspend. +func (c *Client) GetPendingSessionDispatchCleanup(ctx context.Context, afterSessionID string) (*SessionDispatchCleanup, error) { + type candidate struct { + SessionID string + DispatchID uuid.UUID + ActorUID string + } + row, err := queryOne(ctx, c.db, ` + SELECT s.id::text AS session_id, s.dispatch_id, s.actor_uid + FROM session s JOIN runtime_instance r USING (id) + WHERE s.dispatch_id IS NOT NULL + AND s.dispatch_cleanup_executor_id IS NULL AND s.actor_uid IS NOT NULL + AND ($1::text = '' OR s.id > NULLIF($1::text, '')::uuid) + AND r.state = 'RUNTIME_STATE_READY' AND r.operation = 'RUNTIME_OPERATION_NONE' + AND s.dispatch_expires_at <= clock_timestamp() + AND NOT EXISTS (SELECT 1 FROM session_task t WHERE t.history_id = s.history_id + AND t.state NOT IN ('TASK_STATE_COMPLETED', 'TASK_STATE_CANCELED', 'TASK_STATE_FAILED', + 'TASK_STATE_REJECTED', 'TASK_STATE_INPUT_REQUIRED', 'TASK_STATE_AUTH_REQUIRED')) + AND NOT EXISTS (SELECT 1 FROM session_task_event e WHERE e.history_id = s.history_id + AND (NOT e.published OR (e.quiescence_pending AND e.quiescence_executor_id IS NOT NULL))) + AND NOT EXISTS (SELECT 1 FROM session_checkpoint WHERE source_session_id = s.id AND state = 'CREATING') + ORDER BY s.id LIMIT 1 + `, pgx.RowToStructByName[candidate], afterSessionID) + if err != nil { + return nil, notFoundOr(err) + } + session, err := readSession(ctx, c.db, row.SessionID) + if err != nil { + return nil, notFoundOr(err) + } + value, err := toSession(session) + if err != nil { + return nil, err + } + return &SessionDispatchCleanup{Session: value, DispatchID: row.DispatchID, ActorUID: row.ActorUID}, nil +} + +// ClaimSessionDispatchCleanup grants exclusive runtime authority for this exact +// delivery and actor. Admission and history capture commit together. A stale or +// blocked candidate returns ErrNotFound; issued work is never reclaimed by time. +func (c *Client) ClaimSessionDispatchCleanup(ctx context.Context, candidate *SessionDispatchCleanup) (*SessionDispatchCleanup, error) { + if candidate == nil || candidate.Session == nil || candidate.DispatchID == uuid.Nil || candidate.ActorUID == "" { + return nil, ErrFailedPrecondition + } + var result *SessionDispatchCleanup + err := c.withTx(ctx, func(tx pgx.Tx) error { + session, err := lockSession(ctx, tx, candidate.Session.Id) + if err != nil { + return notFoundOr(err) + } + executor := uuid.New() + tag, err := tx.Exec(ctx, ` + UPDATE session s SET dispatch_cleanup_executor_id = $4 + FROM runtime_instance r + WHERE s.id = $1 AND r.id = s.id AND s.dispatch_id = $2 AND s.actor_uid = $3 + AND s.dispatch_cleanup_executor_id IS NULL + AND r.state = 'RUNTIME_STATE_READY' AND r.operation = 'RUNTIME_OPERATION_NONE' + AND s.dispatch_expires_at <= clock_timestamp() + AND NOT EXISTS (SELECT 1 FROM session_task t WHERE t.history_id = s.history_id + AND t.state NOT IN ('TASK_STATE_COMPLETED', 'TASK_STATE_CANCELED', 'TASK_STATE_FAILED', + 'TASK_STATE_REJECTED', 'TASK_STATE_INPUT_REQUIRED', 'TASK_STATE_AUTH_REQUIRED')) + AND NOT EXISTS (SELECT 1 FROM session_task_event e WHERE e.history_id = s.history_id + AND (NOT e.published OR (e.quiescence_pending AND e.quiescence_executor_id IS NOT NULL))) + AND NOT EXISTS (SELECT 1 FROM session_checkpoint WHERE source_session_id = s.id AND state = 'CREATING') + `, session.ID, candidate.DispatchID, candidate.ActorUID, executor) + if err != nil { + return err + } + if tag.RowsAffected() != 1 { + return ErrNotFound + } + value, err := toSession(session) + if err != nil { + return err + } + result = &SessionDispatchCleanup{Session: value, DispatchID: candidate.DispatchID, ActorUID: candidate.ActorUID, ExecutorID: executor} + boundary, err := queryOne(ctx, tx, ` + SELECT e.task_id, e.sequence AS version, t.state FROM session_task_event e + JOIN session_task t ON t.history_id = e.history_id AND t.id = e.task_id + WHERE e.history_id = $1 ORDER BY e.sequence DESC LIMIT 1 + `, pgx.RowToStructByName[dispatchCleanupBoundary], session.HistoryID) + if err == nil { + result.boundary = &boundary + return nil + } + if !errors.Is(err, pgx.ErrNoRows) { + return err + } + hasTasks, err := queryOne(ctx, tx, `SELECT EXISTS (SELECT 1 FROM session_task WHERE history_id = $1)`, pgx.RowTo[bool], session.HistoryID) + if err != nil { + return err + } + if hasTasks { + return fmt.Errorf("session task history is missing: %w", ErrFailedPrecondition) + } + return nil + }) + return result, err +} + +// FinishSessionDispatchCleanup records the claimed snapshot without inventing +// history or changing logical Session state. It releases admission atomically +// with snapshot metadata and obsolete task-quiescence work. Completed retries +// do nothing; a newer cleanup generation or owner returns ErrConflict. +func (c *Client) FinishSessionDispatchCleanup(ctx context.Context, work *SessionDispatchCleanup, snapshot *SessionTaskSnapshot) error { + if work == nil || work.Session == nil || work.DispatchID == uuid.Nil || work.ExecutorID == uuid.Nil { + return ErrFailedPrecondition + } + if snapshot == nil || snapshot.Atespace == "" || snapshot.URI == "" || snapshot.Fidelity == "" { + return fmt.Errorf("dispatch cleanup requires a runtime snapshot: %w", ErrFailedPrecondition) + } + return c.withTx(ctx, func(tx pgx.Tx) error { + session, err := lockSession(ctx, tx, work.Session.Id) + if err != nil { + return notFoundOr(err) + } + type claim struct { + ID *uuid.UUID + Executor *uuid.UUID + ActorUID *string + } + current, err := queryOne(ctx, tx, ` + SELECT dispatch_id AS id, dispatch_cleanup_executor_id AS executor, actor_uid + FROM session WHERE id = $1 + `, pgx.RowToStructByName[claim], session.ID) + if err != nil { + return err + } + if current.ID == nil { + return nil + } + if *current.ID != work.DispatchID || current.Executor == nil || *current.Executor != work.ExecutorID || current.ActorUID == nil || *current.ActorUID != work.ActorUID { + return ErrConflict + } + if target := work.boundary; target != nil { + latest, err := queryOne(ctx, tx, `SELECT MAX(sequence) FROM session_task_event WHERE history_id = $1`, pgx.RowTo[int64], session.HistoryID) + if err != nil { + return err + } + stored, err := readSessionTask(ctx, tx, session.HistoryID, target.TaskID) + if err != nil { + return err + } + if latest != target.Version || stored.State != target.State { + return ErrConflict + } + if err := saveSessionTaskSnapshot(ctx, tx, session.HistoryID, target.TaskID, target.Version, snapshot); err != nil { + return err + } + } + if err := execSQL(ctx, tx, ` + UPDATE session_task_event SET quiescence_pending = FALSE + WHERE history_id = $1 AND quiescence_pending AND quiescence_executor_id IS NULL + `, session.HistoryID); err != nil { + return err + } + return retireSessionDispatchCleanup(ctx, tx, session.ID) + }) +} + +// retireSessionDispatchCleanup removes compensation after execution or an +// authoritative lifecycle outcome takes responsibility. Callers hold the Session +// lock and must have checked admission or established runtime termination. +func retireSessionDispatchCleanup(ctx context.Context, tx pgx.Tx, sessionID uuid.UUID) error { + return execSQL(ctx, tx, ` + UPDATE session SET dispatch_id = NULL, dispatch_expires_at = NULL, dispatch_cleanup_executor_id = NULL + WHERE id = $1 + `, sessionID) +} diff --git a/go/core/internal/database/session_operations.go b/go/core/internal/database/session_operations.go index b629034aee..12b17628b9 100644 --- a/go/core/internal/database/session_operations.go +++ b/go/core/internal/database/session_operations.go @@ -150,6 +150,15 @@ func (c *Client) FinishSessionOperation(ctx context.Context, sessionID string, i return err } if failure == "" { + // Successful lifecycle work supersedes abandoned delivery from the + // prior runtime state. Retire its dispatch and suspension claim in + // this transaction so stale cleanup cannot suspend a resumed actor. + // Creation has no prior dispatch; failures leave cleanup intact. + if result.Operation != apiv1.RuntimeOperation_RUNTIME_OPERATION_CREATE { + if err := retireSessionDispatchCleanup(ctx, tx, row.ID); err != nil { + return err + } + } switch result.Operation { case apiv1.RuntimeOperation_RUNTIME_OPERATION_CREATE: if authority == "" || actorUID == "" { diff --git a/go/core/internal/database/session_quiescence.go b/go/core/internal/database/session_quiescence.go index a1905cdb13..b3a052a8e9 100644 --- a/go/core/internal/database/session_quiescence.go +++ b/go/core/internal/database/session_quiescence.go @@ -56,6 +56,7 @@ func (c *Client) ClaimSessionQuiescence(ctx context.Context) (*SessionQuiescence AND i.state = 'RUNTIME_STATE_READY' AND i.operation = 'RUNTIME_OPERATION_NONE' AND (i.dispatch_expires_at IS NULL OR i.dispatch_expires_at <= clock_timestamp()) + AND NOT EXISTS (SELECT 1 FROM session WHERE id = i.id AND dispatch_cleanup_executor_id IS NOT NULL) AND NOT EXISTS (SELECT 1 FROM session_checkpoint WHERE source_session_id = i.id AND state = 'CREATING') ORDER BY e.sequence LIMIT 1 FOR UPDATE OF i SKIP LOCKED `, pgx.RowToStructByName[candidate]) @@ -89,6 +90,7 @@ func (c *Client) ClaimSessionQuiescence(ctx context.Context) (*SessionQuiescence WHERE sequence = $1 AND published AND quiescence_pending AND quiescence_executor_id IS NULL AND NOT EXISTS (SELECT 1 FROM session WHERE id = $3 AND dispatch_expires_at > clock_timestamp()) + AND NOT EXISTS (SELECT 1 FROM session WHERE id = $3 AND dispatch_cleanup_executor_id IS NOT NULL) AND NOT EXISTS (SELECT 1 FROM session_checkpoint WHERE source_session_id = $3 AND state = 'CREATING') `, result.Version, result.ExecutorID, session.ID) if err != nil { @@ -138,23 +140,28 @@ func (c *Client) FinishSessionQuiescence(ctx context.Context, work *SessionQuies if stored.State != string(work.State) { return ErrConflict } - if snapshot != nil { - if err := execSQL(ctx, tx, ` - UPDATE session_task SET snapshot_atespace = $3, snapshot_uri = $4, - snapshot_fidelity = $5, history_sequence = $6 - WHERE history_id = $1 AND id = $2 - `, session.HistoryID, work.TaskID, snapshot.Atespace, snapshot.URI, snapshot.Fidelity, work.Version); err != nil { - return err - } - if err := execSQL(ctx, tx, ` - UPDATE session_task_event SET snapshot_atespace = $2, snapshot_uri = $3, snapshot_fidelity = $4 - WHERE sequence = $1 - `, work.Version, snapshot.Atespace, snapshot.URI, snapshot.Fidelity); err != nil { - return err - } + if err := saveSessionTaskSnapshot(ctx, tx, session.HistoryID, work.TaskID, work.Version, snapshot); err != nil { + return err } return execSQL(ctx, tx, ` UPDATE session_task_event SET quiescence_pending = FALSE WHERE sequence = $1 `, work.Version) }) } + +// saveSessionTaskSnapshot associates a confirmed runtime snapshot with its +// captured history version and head task. The caller holds the Session lock and +// has validated the version under exclusive lifecycle authority. +func saveSessionTaskSnapshot(ctx context.Context, tx pgx.Tx, historyID uuid.UUID, taskID string, version int64, snapshot *SessionTaskSnapshot) error { + if err := execSQL(ctx, tx, ` + UPDATE session_task SET snapshot_atespace = $3, snapshot_uri = $4, + snapshot_fidelity = $5, history_sequence = $6 + WHERE history_id = $1 AND id = $2 + `, historyID, taskID, snapshot.Atespace, snapshot.URI, snapshot.Fidelity, version); err != nil { + return err + } + return execSQL(ctx, tx, ` + UPDATE session_task_event SET snapshot_atespace = $2, snapshot_uri = $3, snapshot_fidelity = $4 + WHERE sequence = $1 + `, version, snapshot.Atespace, snapshot.URI, snapshot.Fidelity) +} diff --git a/go/core/internal/database/session_recovery.go b/go/core/internal/database/session_recovery.go index d591f33b5b..a0bb345ee1 100644 --- a/go/core/internal/database/session_recovery.go +++ b/go/core/internal/database/session_recovery.go @@ -116,7 +116,10 @@ func (c *Client) BeginCrashedSessionRecovery(ctx context.Context, generation *Se 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 { + if err := execSQL(ctx, tx, ` + UPDATE session SET data = $2, dispatch_id = NULL, dispatch_expires_at = NULL, dispatch_cleanup_executor_id = NULL + WHERE id = $1 + `, row.ID, row.Data); err != nil { return err } if err := settleCrashedSessionTasks(ctx, tx, row, snapshot); err != nil { diff --git a/go/core/internal/database/session_tasks.go b/go/core/internal/database/session_tasks.go index a51317f2c4..11e8c1960a 100644 --- a/go/core/internal/database/session_tasks.go +++ b/go/core/internal/database/session_tasks.go @@ -26,6 +26,14 @@ var ErrDispatchBusy = errors.New("session temporarily unavailable for dispatch") // ErrMessageAccepted means an earlier dispatch already persisted this initial input. var ErrMessageAccepted = errors.New("message already accepted") +// sessionDispatchLifetime is how long a send attempt has to start a turn. +// During this window, automatic suspension waits for delivery to finish. If no +// turn starts before expiry, the dispatch token can no longer start one, and the +// lifecycle worker may suspend the idle actor. Starting a turn consumes the +// reservation; this timeout does not stop a turn that is already running. +// Canceling delivery expires the reservation immediately. +const sessionDispatchLifetime = 2 * time.Minute + // ReserveSessionDispatch prevents pause/suspend between routing a request // and its first persisted task. The attempt expires after two minutes; a late // first save must fail before native execution. This is not a native work lease. @@ -68,15 +76,17 @@ func (c *Client) ReserveSessionDispatch(ctx context.Context, sessionID string, d return err } return execSQL(ctx, tx, ` - UPDATE session SET dispatch_id = $2, dispatch_expires_at = clock_timestamp() + INTERVAL '2 minutes' + UPDATE session SET dispatch_id = $2, dispatch_expires_at = clock_timestamp() + $3 * INTERVAL '1 millisecond' WHERE id = $1 - `, session.ID, dispatchID) + `, session.ID, dispatchID, sessionDispatchLifetime.Milliseconds()) }) } // RevokeSessionDispatch fences an unused attempt. True proves its first // save did not commit and cannot commit later, so retrying the send is safe. // False does not prove acceptance: callers must preserve ambiguous outcomes. +// Repeated revocations retain that proof while the expired reservation remains. +// Keeping the expired reservation also schedules abandoned runtime cleanup. func (c *Client) RevokeSessionDispatch(ctx context.Context, sessionID string, dispatchID uuid.UUID, messageID string) (bool, error) { var notAccepted bool err := c.withTx(ctx, func(tx pgx.Tx) error { @@ -85,7 +95,7 @@ func (c *Client) RevokeSessionDispatch(ctx context.Context, sessionID string, di return notFoundOr(err) } tag, err := tx.Exec(ctx, ` - UPDATE session SET dispatch_id = NULL, dispatch_expires_at = NULL + UPDATE session SET dispatch_expires_at = LEAST(dispatch_expires_at, clock_timestamp()) WHERE id = $1 AND dispatch_id = $2 `, session.ID, dispatchID) if err != nil || tag.RowsAffected() == 0 { @@ -186,6 +196,7 @@ func (c *Client) writeRuntimeTask(ctx context.Context, sessionID string, expecte } } wasWaiting := stored != nil && (stored.Status.State == a2a.TaskStateInputRequired || stored.Status.State == a2a.TaskStateAuthRequired) + remainsWaiting := wasWaiting && stored.Status.State == task.Status.State if session.State != "RUNTIME_STATE_READY" || session.Operation != "RUNTIME_OPERATION_NONE" { return fmt.Errorf("session cannot accept runtime updates during a lifecycle operation: %w", ErrConflict) } @@ -194,7 +205,7 @@ func (c *Client) writeRuntimeTask(ctx context.Context, sessionID string, expecte UPDATE session SET dispatch_id = CASE WHEN $3 THEN dispatch_id END, dispatch_expires_at = CASE WHEN $3 THEN dispatch_expires_at END WHERE id = $1 AND dispatch_id = $2 AND dispatch_expires_at > clock_timestamp() - `, session.ID, dispatchID, stored != nil && stored.Status.State == task.Status.State) + `, session.ID, dispatchID, remainsWaiting) if err != nil { return err } @@ -208,6 +219,14 @@ func (c *Client) writeRuntimeTask(ctx context.Context, sessionID string, expecte // New execution supersedes idle work that no worker has claimed yet. // Keep the marker for idempotent acknowledgements of the older boundary. if stored == nil || wasWaiting { + // A reply can be saved before the task resumes. Keep the dispatch while + // the task is still waiting, so a canceled delivery can still trigger + // actor suspension. + if !remainsWaiting { + if err := retireSessionDispatchCleanup(ctx, tx, session.ID); err != nil { + return err + } + } if err := execSQL(ctx, tx, ` UPDATE session_task_event SET quiescence_pending = FALSE WHERE history_id = $1 AND quiescence_pending @@ -826,6 +845,7 @@ func requireSettledRuntime(ctx context.Context, db dbExecutor, historyID uuid.UU AND (NOT published OR (quiescence_pending AND quiescence_executor_id IS NOT NULL))) OR EXISTS (SELECT 1 FROM session WHERE history_id = $1 AND dispatch_expires_at > clock_timestamp() AND ($2::text = '' OR dispatch_id::text <> $2)) + OR EXISTS (SELECT 1 FROM session WHERE history_id = $1 AND dispatch_cleanup_executor_id IS NOT NULL) `, pgx.RowTo[bool], historyID, dispatchID) if err != nil { return err diff --git a/go/core/internal/database/task_finalization.go b/go/core/internal/database/task_finalization.go index 2f80fb7c79..8ffbafba8f 100644 --- a/go/core/internal/database/task_finalization.go +++ b/go/core/internal/database/task_finalization.go @@ -54,7 +54,10 @@ func publishSavedSessionResults(ctx context.Context, tx pgx.Tx, session sessionR // retireSessionQuiescence drops idle-work markers for a runtime that no longer // needs a later suspension. func retireSessionQuiescence(ctx context.Context, tx pgx.Tx, session sessionRow) error { - return execSQL(ctx, tx, `UPDATE session_task_event SET quiescence_pending = FALSE WHERE history_id = $1 AND quiescence_pending`, session.HistoryID) + 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 + } + return retireSessionDispatchCleanup(ctx, tx, session.ID) } // settleSessionTask publishes one saved boundary under the session lock. Runtime diff --git a/go/core/internal/database/task_store_test.go b/go/core/internal/database/task_store_test.go index 1ba79bae67..e634b2597a 100644 --- a/go/core/internal/database/task_store_test.go +++ b/go/core/internal/database/task_store_test.go @@ -540,3 +540,441 @@ func TestCheckpointRejectsOlderPendingTask(t *testing.T) { require.ErrorIs(t, err, ErrFailedPrecondition) require.ErrorContains(t, err, "all tasks to be terminal") } + +func TestDispatchCleanupFencesTasklessSession(t *testing.T) { + t.Run("Given a restored actor without a turn, when its dispatch is revoked, assert that cleanup fences late execution and preserves empty history", func(t *testing.T) { + client, session := dispatchCleanupFixture(t) + ctx := t.Context() + id := uuid.New() + require.NoError(t, client.ReserveSessionDispatch(ctx, session.Id, id, "input")) + _, err := client.GetPendingSessionDispatchCleanup(ctx, "") + require.ErrorIs(t, err, ErrNotFound, "live delivery must not become cleanup") + revoked, err := client.RevokeSessionDispatch(ctx, session.Id, id, "input") + require.NoError(t, err) + require.True(t, revoked) + revoked, err = client.RevokeSessionDispatch(ctx, session.Id, id, "input") + require.NoError(t, err) + require.True(t, revoked, "repeated revocation must preserve proof that input was never accepted") + task := newSessionTask(uuid.NewString(), "input") + task.ContextID = session.ContextId + _, err = client.CreateRuntimeTask(ctx, session.Id, taskMutationHash("late"), task, id.String()) + require.ErrorIs(t, err, ErrFailedPrecondition, "revocation fences late input before cleanup claims") + candidate, err := client.GetPendingSessionDispatchCleanup(ctx, "") + require.NoError(t, err) + require.Equal(t, id, candidate.DispatchID) + replaced := *candidate + replaced.ActorUID = "another-actor" + _, err = client.ClaimSessionDispatchCleanup(ctx, &replaced) + require.ErrorIs(t, err, ErrNotFound, "the claim must remain bound to the observed actor") + work, err := client.ClaimSessionDispatchCleanup(ctx, candidate) + require.NoError(t, err) + require.NotEqual(t, uuid.Nil, work.ExecutorID) + revoked, err = client.RevokeSessionDispatch(ctx, session.Id, id, "input") + require.NoError(t, err) + require.True(t, revoked, "repeated revocation remains safe after cleanup is claimed") + _, err = client.ClaimSessionDispatchCleanup(ctx, candidate) + require.ErrorIs(t, err, ErrNotFound) + require.ErrorIs(t, client.ReserveSessionDispatch(ctx, session.Id, uuid.New(), "next"), ErrDispatchBusy) + for _, token := range []string{id.String(), ""} { + _, err = client.CreateRuntimeTask(ctx, session.Id, taskMutationHash("late"), task, token) + require.ErrorIs(t, err, ErrFailedPrecondition) + } + _, err = client.BeginSessionOperation(ctx, session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_SUSPEND, false) + require.ErrorIs(t, err, ErrFailedPrecondition) + _, err = client.BeginSessionOperation(ctx, session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE, true) + require.ErrorIs(t, err, ErrConflict) + require.ErrorIs(t, client.FinishSessionDispatchCleanup(ctx, work, nil), ErrFailedPrecondition) + require.NoError(t, client.FinishSessionDispatchCleanup(ctx, work, dispatchCleanupSnapshot())) + require.NoError(t, client.FinishSessionDispatchCleanup(ctx, work, dispatchCleanupSnapshot())) + revoked, err = client.RevokeSessionDispatch(ctx, session.Id, id, "input") + require.NoError(t, err) + require.False(t, revoked, "completed cleanup no longer retains the dispatch proof") + _, err = client.CreateRuntimeTask(ctx, session.Id, taskMutationHash("late"), task, id.String()) + require.ErrorIs(t, err, ErrFailedPrecondition, "completed cleanup must not revive the old dispatch") + _, err = client.GetPendingSessionDispatchCleanup(ctx, "") + require.ErrorIs(t, err, ErrNotFound) + tasks, total, err := client.ListSessionTasks(ctx, session.Id, "", a2a.TaskStateUnspecified, nil, 10, nil) + require.NoError(t, err) + require.Empty(t, tasks) + require.Zero(t, total) + current, err := client.GetSessionByID(ctx, session.Id) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, current.State) + require.NoError(t, client.ReserveSessionDispatch(ctx, session.Id, uuid.New(), "next")) + }) +} + +func TestDispatchCleanupRecoversExpiredDelivery(t *testing.T) { + t.Run("Given a gateway that never revoked its delivery, when its reservation expires, assert that cleanup fences late acceptance", func(t *testing.T) { + client, session := dispatchCleanupFixture(t) + ctx := t.Context() + id := uuid.New() + require.NoError(t, client.ReserveSessionDispatch(ctx, session.Id, id, "input")) + // Advance durable deadlines without waiting for the production lease. + _, err := client.db.Exec(ctx, `UPDATE session SET dispatch_expires_at = clock_timestamp() - INTERVAL '1 second' WHERE id = $1`, session.Id) + require.NoError(t, err) + candidate, err := client.GetPendingSessionDispatchCleanup(ctx, "") + require.NoError(t, err) + revoked, err := client.RevokeSessionDispatch(ctx, session.Id, id, "input") + require.NoError(t, err) + require.True(t, revoked, "expiry must preserve proof that the input was never accepted") + work, err := client.ClaimSessionDispatchCleanup(ctx, candidate) + require.NoError(t, err) + task := newSessionTask(uuid.NewString(), "input") + task.ContextID = session.ContextId + _, err = client.CreateRuntimeTask(ctx, session.Id, taskMutationHash("late"), task, id.String()) + require.ErrorIs(t, err, ErrFailedPrecondition) + require.NoError(t, client.FinishSessionDispatchCleanup(ctx, work, dispatchCleanupSnapshot())) + }) +} + +func TestDispatchCleanupRejectsStalePreparation(t *testing.T) { + t.Run("Given a replacement delivery, when old cleanup preparation finishes, assert that it cannot claim or retire the newer obligation", func(t *testing.T) { + client, session := dispatchCleanupFixture(t) + ctx := t.Context() + first := abandonedDispatchFixture(t, client, session) + second := abandonedDispatchFixture(t, client, session) + require.NotEqual(t, first.DispatchID, second.DispatchID) + _, err := client.ClaimSessionDispatchCleanup(ctx, first) + require.ErrorIs(t, err, ErrNotFound) + revoked, err := client.RevokeSessionDispatch(ctx, session.Id, first.DispatchID, "input") + require.NoError(t, err) + require.False(t, revoked) + work, err := client.ClaimSessionDispatchCleanup(ctx, second) + require.NoError(t, err) + require.NoError(t, client.FinishSessionDispatchCleanup(ctx, work, dispatchCleanupSnapshot())) + third := abandonedDispatchFixture(t, client, session) + require.ErrorIs(t, client.FinishSessionDispatchCleanup(ctx, work, dispatchCleanupSnapshot()), ErrConflict) + pending, err := client.GetPendingSessionDispatchCleanup(ctx, "") + require.NoError(t, err) + require.Equal(t, third.DispatchID, pending.DispatchID) + require.NoError(t, client.ReserveSessionDispatch(ctx, session.Id, uuid.New(), "next"), "unclaimed cleanup must not block sends") + }) +} + +func TestDispatchCleanupRacesNewDelivery(t *testing.T) { + t.Run("Given pending cleanup, when a new delivery and cleanup claim race, assert that exactly one acquires execution authority", func(t *testing.T) { + client, session := dispatchCleanupFixture(t) + ctx := t.Context() + for range 12 { + candidate := abandonedDispatchFixture(t, client, session) + start := make(chan struct{}) + claimed := make(chan *SessionDispatchCleanup, 1) + claimErrors, sendErrors := make(chan error, 1), make(chan error, 1) + next := uuid.New() + go func() { + <-start + work, err := client.ClaimSessionDispatchCleanup(ctx, candidate) + claimed <- work + claimErrors <- err + }() + go func() { + <-start + sendErrors <- client.ReserveSessionDispatch(ctx, session.Id, next, "next") + }() + close(start) + work, claimErr, sendErr := <-claimed, <-claimErrors, <-sendErrors + if claimErr == nil { + require.ErrorIs(t, sendErr, ErrDispatchBusy) + require.NoError(t, client.FinishSessionDispatchCleanup(ctx, work, dispatchCleanupSnapshot())) + } else { + require.ErrorIs(t, claimErr, ErrNotFound) + require.NoError(t, sendErr) + _, err := client.RevokeSessionDispatch(ctx, session.Id, next, "next") + require.NoError(t, err) + } + } + }) +} + +func TestDispatchCleanupRetainsAcceptedWaitingInput(t *testing.T) { + t.Run("Given input saved on a waiting task, when delivery is revoked before native execution, assert that cleanup preserves receipt and snapshot coverage", func(t *testing.T) { + client := NewClient(setupTestDB(t), "public") + ctx := t.Context() + sessionFixture(t, client, ctx, "team-a", "revision", "assistant", "kagent") + session, waiting := waitingTaskFixture(t, client) + id := uuid.New() + require.NoError(t, client.ReserveSessionDispatch(ctx, session.Id, id, "")) + _, version, err := client.GetVersionedSessionTask(ctx, session.Id, string(waiting.ID)) + require.NoError(t, err) + input := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("continue")) + input.TaskID, input.ContextID = waiting.ID, waiting.ContextID + waiting.History = append(waiting.History, input) + version, err = client.UpdateSessionTask(ctx, session.Id, version, taskMutationHash("input"), waiting, input, id.String()) + require.NoError(t, err) + revoked, err := client.RevokeSessionDispatch(ctx, session.Id, id, input.ID) + require.NoError(t, err) + require.False(t, revoked, "a persisted receipt must not become retryable") + revoked, err = client.RevokeSessionDispatch(ctx, session.Id, id, input.ID) + require.NoError(t, err) + require.False(t, revoked, "repeating revocation must not make saved input retryable") + candidate, err := client.GetPendingSessionDispatchCleanup(ctx, "") + require.NoError(t, err) + work, err := client.ClaimSessionDispatchCleanup(ctx, candidate) + require.NoError(t, err) + revoked, err = client.RevokeSessionDispatch(ctx, session.Id, id, input.ID) + require.NoError(t, err) + require.False(t, revoked, "claiming cleanup must preserve the input receipt") + active := *waiting + active.Status.State = a2a.TaskStateWorking + _, err = client.UpdateSessionTask(ctx, session.Id, version, taskMutationHash("late active"), &active, &active, id.String()) + require.ErrorIs(t, err, ErrFailedPrecondition) + _, _, err = client.ReserveSessionCheckpoint(ctx, &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(waiting.ID)}, session.Creator, uuid.NewString()) + require.ErrorIs(t, err, ErrFailedPrecondition) + snapshot := dispatchCleanupSnapshot() + require.NoError(t, client.FinishSessionDispatchCleanup(ctx, work, snapshot)) + stored, afterVersion, err := client.GetVersionedSessionTask(ctx, session.Id, string(waiting.ID)) + require.NoError(t, err) + require.Equal(t, version, afterVersion, "cleanup must not append a task event") + require.Equal(t, a2a.TaskStateInputRequired, stored.Status.State) + require.Equal(t, input.ID, stored.History[len(stored.History)-1].ID) + _, err = client.BeginCrashedSessionRecovery(ctx, observedGeneration(t, client, session.Id, "actor-"+session.Id), *snapshot) + require.NoError(t, err) + stored, err = client.GetSettledSessionTask(ctx, session.Id, string(waiting.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateInputRequired, stored.Status.State, "the latest cleanup snapshot must cover the parked wait") + }) +} + +func TestDispatchCleanupTransfersToActiveTask(t *testing.T) { + t.Run("Given a reserved delivery, when its initial active task commits, assert that dispatch cleanup is retired", func(t *testing.T) { + client, session := dispatchCleanupFixture(t) + ctx := t.Context() + candidate := abandonedDispatchFixture(t, client, session) + id := uuid.New() + require.NoError(t, client.ReserveSessionDispatch(ctx, session.Id, id, "accepted")) + task := newSessionTask(uuid.NewString(), "accepted") + task.ContextID = session.ContextId + _, err := client.CreateRuntimeTask(ctx, session.Id, taskMutationHash("accepted"), task, id.String()) + require.NoError(t, err) + revoked, err := client.RevokeSessionDispatch(ctx, session.Id, id, "accepted") + require.NoError(t, err) + require.False(t, revoked) + _, err = client.ClaimSessionDispatchCleanup(ctx, candidate) + require.ErrorIs(t, err, ErrNotFound) + _, err = client.GetPendingSessionDispatchCleanup(ctx, "") + require.ErrorIs(t, err, ErrNotFound) + }) +} + +func TestDispatchCleanupCoordinatesTaskQuiescence(t *testing.T) { + for _, taskFirst := range []bool{false, true} { + name := "Given settled task work, when dispatch cleanup claims first, assert that one snapshot satisfies both obligations" + if taskFirst { + name = "Given abandoned dispatch work, when task quiescence claims first, assert that cleanup waits and later retires safely" + } + t.Run(name, func(t *testing.T) { + client, session := dispatchCleanupFixture(t) + ctx := t.Context() + task := newSessionTask(uuid.NewString(), "completed") + task.ContextID, task.Status.State = session.ContextId, a2a.TaskStateCompleted + version, err := client.CreateRuntimeTask(ctx, session.Id, taskMutationHash("completed"), task, "") + require.NoError(t, err) + require.NoError(t, client.SettleSessionTask(ctx, session.Id, string(task.ID), version)) + candidate := abandonedDispatchFixture(t, client, session) + if taskFirst { + normal, err := client.ClaimSessionQuiescence(ctx) + require.NoError(t, err) + _, err = client.ClaimSessionDispatchCleanup(ctx, candidate) + require.ErrorIs(t, err, ErrNotFound) + require.NoError(t, client.FinishSessionQuiescence(ctx, normal, dispatchCleanupSnapshot())) + } + work, err := client.ClaimSessionDispatchCleanup(ctx, candidate) + require.NoError(t, err) + _, err = client.ClaimSessionQuiescence(ctx) + require.ErrorIs(t, err, ErrNotFound) + require.NoError(t, client.FinishSessionDispatchCleanup(ctx, work, dispatchCleanupSnapshot())) + _, err = client.ClaimSessionQuiescence(ctx) + require.ErrorIs(t, err, ErrNotFound) + _, err = client.GetPendingSessionDispatchCleanup(ctx, "") + require.ErrorIs(t, err, ErrNotFound) + _, afterVersion, err := client.GetVersionedSessionTask(ctx, session.Id, string(task.ID)) + require.NoError(t, err) + require.Equal(t, version, afterVersion) + }) + } +} + +func TestDispatchCleanupCoordinatesCheckpoints(t *testing.T) { + t.Run("Given a completed head with a snapshot and claimed cleanup, when a checkpoint is requested, assert that it waits for the new snapshot", func(t *testing.T) { + client, session := dispatchCleanupFixture(t) + ctx := t.Context() + task := newSessionTask(uuid.NewString(), "completed") + task.ContextID, task.Status.State = session.ContextId, a2a.TaskStateCompleted + require.NoError(t, saveRuntimeTask(t, client, session.Id, task, task, dispatchCleanupSnapshot())) + candidate := abandonedDispatchFixture(t, client, session) + work, err := client.ClaimSessionDispatchCleanup(ctx, candidate) + require.NoError(t, err) + + checkpoint := &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID)} + requestID := uuid.NewString() + _, _, err = client.ReserveSessionCheckpoint(ctx, checkpoint, session.Creator, requestID) + require.ErrorIs(t, err, ErrSnapshotPending, "an older valid snapshot must not bypass issued cleanup") + + const snapshotURI = "s3://snapshots/cleanup-checkpoint" + snapshot := dispatchCleanupSnapshot() + snapshot.URI = snapshotURI + require.NoError(t, client.FinishSessionDispatchCleanup(ctx, work, snapshot)) + reserved, captured, err := client.ReserveSessionCheckpoint(ctx, checkpoint, session.Creator, requestID) + require.NoError(t, err) + require.Equal(t, apiv1.CheckpointState_CHECKPOINT_STATE_CREATING, reserved.State) + require.Equal(t, snapshot, captured, "retry must capture the completed cleanup snapshot") + }) + t.Run("Given a checkpoint being created, when cleanup attempts to claim the actor, assert that it waits for snapshot retention", func(t *testing.T) { + client, session := dispatchCleanupFixture(t) + ctx := t.Context() + task := newSessionTask(uuid.NewString(), "completed") + task.ContextID, task.Status.State = session.ContextId, a2a.TaskStateCompleted + require.NoError(t, saveRuntimeTask(t, client, session.Id, task, task, dispatchCleanupSnapshot())) + candidate := abandonedDispatchFixture(t, client, session) + checkpoint, _, err := client.ReserveSessionCheckpoint(ctx, &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID)}, session.Creator, uuid.NewString()) + require.NoError(t, err) + + _, err = client.GetPendingSessionDispatchCleanup(ctx, "") + require.ErrorIs(t, err, ErrNotFound) + _, err = client.ClaimSessionDispatchCleanup(ctx, candidate) + require.ErrorIs(t, err, ErrNotFound, "preparation before checkpoint creation must not grant runtime authority") + + const ( + checkpointTag = "checkpoint-tag" + retainedSnapshotURI = "s3://snapshots/checkpoint-retained" + ) + _, err = client.FinalizeSessionCheckpoint(ctx, checkpoint.Id, checkpointTag, retainedSnapshotURI, "") + require.NoError(t, err) + pending, err := client.GetPendingSessionDispatchCleanup(ctx, "") + require.NoError(t, err) + require.Equal(t, candidate.DispatchID, pending.DispatchID) + work, err := client.ClaimSessionDispatchCleanup(ctx, pending) + require.NoError(t, err) + require.NoError(t, client.FinishSessionDispatchCleanup(ctx, work, dispatchCleanupSnapshot())) + }) +} + +func TestDispatchCleanupCapturesForkedHistory(t *testing.T) { + t.Run("Given forked history without a task quiescence marker, when an abandoned delivery is cleaned up, assert that the existing history receives its snapshot", func(t *testing.T) { + client, source := dispatchCleanupFixture(t) + ctx := t.Context() + task := newSessionTask(uuid.NewString(), "completed") + task.ContextID, task.Status.State = source.ContextId, a2a.TaskStateCompleted + require.NoError(t, saveRuntimeTask(t, client, source.Id, task, task, dispatchCleanupSnapshot())) + checkpoint, _, err := client.ReserveSessionCheckpoint(ctx, &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: source.Id, HeadTaskId: string(task.ID)}, source.Creator, uuid.NewString()) + require.NoError(t, err) + _, err = client.FinalizeSessionCheckpoint(ctx, checkpoint.Id, "fork-tag", "s3://snapshots/fork-retained", "") + require.NoError(t, err) + fork, _, err := client.ForkSession(ctx, checkpoint.Id, source.Creator, uuid.NewString(), uuid.NewString()) + require.NoError(t, err) + fork, err = markSessionReady(ctx, client, fork.Id, "fork.example") + require.NoError(t, err) + inherited, _, err := client.ListSessionTasks(ctx, fork.Id, "", a2a.TaskStateUnspecified, nil, 10, nil) + require.NoError(t, err) + require.Len(t, inherited, 1) + _, version, err := client.GetVersionedSessionTask(ctx, fork.Id, string(inherited[0].ID)) + require.NoError(t, err) + candidate := abandonedDispatchFixture(t, client, fork) + work, err := client.ClaimSessionDispatchCleanup(ctx, candidate) + require.NoError(t, err) + snapshot := dispatchCleanupSnapshot() + require.NoError(t, client.FinishSessionDispatchCleanup(ctx, work, snapshot)) + _, captured, err := client.ReserveSessionCheckpoint(ctx, &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: fork.Id, HeadTaskId: string(inherited[0].ID)}, fork.Creator, uuid.NewString()) + require.NoError(t, err) + require.Equal(t, snapshot, captured) + _, afterVersion, err := client.GetVersionedSessionTask(ctx, fork.Id, string(inherited[0].ID)) + require.NoError(t, err) + require.Equal(t, version, afterVersion) + }) +} + +func TestDispatchCleanupRetiresAfterRuntimeRecovery(t *testing.T) { + t.Run("Given issued dispatch cleanup, when the same actor is reverted after a confirmed crash, assert that stale cleanup cannot claim the recovered runtime", func(t *testing.T) { + client, session := dispatchCleanupFixture(t) + ctx := t.Context() + candidate := abandonedDispatchFixture(t, client, session) + work, err := client.ClaimSessionDispatchCleanup(ctx, candidate) + require.NoError(t, err) + completeCrashedRecovery(t, client, session.Id, candidate.ActorUID) + _, err = client.GetPendingSessionDispatchCleanup(ctx, "") + require.ErrorIs(t, err, ErrNotFound) + next := abandonedDispatchFixture(t, client, session) + _, err = client.ClaimSessionDispatchCleanup(ctx, candidate) + require.ErrorIs(t, err, ErrNotFound) + require.ErrorIs(t, client.FinishSessionDispatchCleanup(ctx, work, dispatchCleanupSnapshot()), ErrConflict) + claimed, err := client.ClaimSessionDispatchCleanup(ctx, next) + require.NoError(t, err) + require.NoError(t, client.FinishSessionDispatchCleanup(ctx, claimed, dispatchCleanupSnapshot())) + }) + t.Run("Given pending cleanup, when explicit suspension and resume complete, assert that the old obligation is retired", func(t *testing.T) { + client, session := dispatchCleanupFixture(t) + ctx := t.Context() + candidate := abandonedDispatchFixture(t, client, session) + _, err := finishSessionOperation(ctx, client, session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_SUSPEND, "") + require.NoError(t, err) + _, err = finishSessionOperation(ctx, client, session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME, "") + require.NoError(t, err) + _, err = client.GetPendingSessionDispatchCleanup(ctx, "") + require.ErrorIs(t, err, ErrNotFound) + _, err = client.ClaimSessionDispatchCleanup(ctx, candidate) + require.ErrorIs(t, err, ErrNotFound) + }) +} + +func TestDispatchCleanupCursorSkipsUnreadyActor(t *testing.T) { + t.Run("Given two pending actors, when a worker leaves the first unclaimed, assert that its cursor reaches the next and a new pass revisits the first", func(t *testing.T) { + client, first := dispatchCleanupFixture(t) + ctx := t.Context() + second, _, err := client.CreateSession(ctx, newSessionRequest(uuid.NewString(), "assistant", "kagent", ""), uuid.NewString()) + require.NoError(t, err) + second, err = markSessionReady(ctx, client, second.Id, "second.example") + require.NoError(t, err) + if second.Id < first.Id { + first, second = second, first + } + for _, session := range []*apiv1.Session{first, second} { + id := uuid.New() + require.NoError(t, client.ReserveSessionDispatch(ctx, session.Id, id, "pending")) + _, err := client.RevokeSessionDispatch(ctx, session.Id, id, "pending") + require.NoError(t, err) + } + candidate, err := client.GetPendingSessionDispatchCleanup(ctx, "") + require.NoError(t, err) + require.Equal(t, first.Id, candidate.Session.Id) + next, err := client.GetPendingSessionDispatchCleanup(ctx, candidate.Session.Id) + require.NoError(t, err) + require.Equal(t, second.Id, next.Session.Id) + _, err = client.GetPendingSessionDispatchCleanup(ctx, next.Session.Id) + require.ErrorIs(t, err, ErrNotFound) + again, err := client.GetPendingSessionDispatchCleanup(ctx, "") + require.NoError(t, err) + require.Equal(t, candidate.DispatchID, again.DispatchID) + require.NoError(t, client.ReserveSessionDispatch(ctx, first.Id, uuid.New(), "new delivery")) + next, err = client.GetPendingSessionDispatchCleanup(ctx, "") + require.NoError(t, err) + require.Equal(t, second.Id, next.Session.Id, "a live dispatch is ineligible without an extra retry timestamp") + }) +} + +func dispatchCleanupFixture(t *testing.T) (*Client, *apiv1.Session) { + t.Helper() + 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) + return client, session +} + +func abandonedDispatchFixture(t *testing.T, client *Client, session *apiv1.Session) *SessionDispatchCleanup { + t.Helper() + id := uuid.New() + require.NoError(t, client.ReserveSessionDispatch(t.Context(), session.Id, id, "abandoned")) + _, err := client.RevokeSessionDispatch(t.Context(), session.Id, id, "abandoned") + require.NoError(t, err) + work, err := client.GetPendingSessionDispatchCleanup(t.Context(), "") + require.NoError(t, err) + require.Equal(t, id, work.DispatchID) + return work +} + +func dispatchCleanupSnapshot() *SessionTaskSnapshot { + return &SessionTaskSnapshot{Atespace: "team-a", URI: "s3://snapshots/dispatch-cleanup", Fidelity: "SNAPSHOT_FIDELITY_VOLUMES"} +} diff --git a/go/core/internal/grpcserver/scheduledrun_controller_test.go b/go/core/internal/grpcserver/scheduledrun_controller_test.go index 244ad99763..12419b372c 100644 --- a/go/core/internal/grpcserver/scheduledrun_controller_test.go +++ b/go/core/internal/grpcserver/scheduledrun_controller_test.go @@ -356,7 +356,7 @@ func TestScheduledRunControllerThroughGRPC(t *testing.T) { if tc.want == apiv1.ScheduledRunExecutionState_SCHEDULED_RUN_EXECUTION_STATE_SUCCEEDED && !tc.stream { controllerStore.loseTaskLink = true } - interactions := sessionsvc.NewInteractionService(store, nil, sessionsvc.NewService(store, scheduledControllerAuthorizer{}, nil), false) + interactions := sessionsvc.NewInteractionService(store, nil, sessionsvc.NewService(store, scheduledControllerAuthorizer{}, nil), false, nil) controller := scheduledrun.NewController(controllerStore, workflow, a2agateway.New(interactions, dialer, "http://gateway.test"), time.Second) go func() { done <- controller.Start(ctx) }() t.Cleanup(func() { cancel(); require.NoError(t, <-done) }) diff --git a/go/core/internal/grpcserver/task_store_test.go b/go/core/internal/grpcserver/task_store_test.go index 24082a062d..bdb9cdb24f 100644 --- a/go/core/internal/grpcserver/task_store_test.go +++ b/go/core/internal/grpcserver/task_store_test.go @@ -287,7 +287,7 @@ func TestRuntimeTaskStoreThroughGRPC(t *testing.T) { var httpTransport a2aclient.Transport for i := range gateways { publicListener := bufconn.Listen(DefaultMaxMessageSize) - interactions := sessionsvc.NewInteractionService(store, nil, sessionsvc.NewService(store, &auth.NoopAuthorizer{}, nil), true) + interactions := sessionsvc.NewInteractionService(store, nil, sessionsvc.NewService(store, &auth.NoopAuthorizer{}, nil), true, nil) gateway := a2agateway.New(interactions, taskStoreRuntimeDialer{runtimeListener}, "http://gateway.test") public, err := New(Config{ Listener: publicListener, SystemService: testSystemService(), @@ -443,7 +443,7 @@ func TestRuntimeTaskStoreThroughGRPC(t *testing.T) { storedConfigs, err := store.ListTaskPushConfigs(t.Context(), id, parkedID, "", 10) require.NoError(t, err) require.Len(t, storedConfigs, 2, "stored configs: %+v", storedConfigs) - managedService := sessionsvc.NewInteractionService(store, nil, sessionsvc.NewService(store, &auth.NoopAuthorizer{}, nil), true) + managedService := sessionsvc.NewInteractionService(store, nil, sessionsvc.NewService(store, &auth.NoopAuthorizer{}, nil), true, nil) serviceCtx := auth.AuthSessionTo(publicCtx, &authimpl.SimpleSession{P: auth.Principal{ID: "alice"}}) servicePage, err := managedService.ListTaskPushConfigs(serviceCtx, types.NamespacedName{Namespace: "team-a", Name: "assistant"}, &a2a.ListTaskPushConfigRequest{TaskID: a2a.TaskID(parkedID), PageSize: 1}) diff --git a/go/core/internal/service/session/expiration_test.go b/go/core/internal/service/session/expiration_test.go index 4acd1e36dc..3f6f9277f6 100644 --- a/go/core/internal/service/session/expiration_test.go +++ b/go/core/internal/service/session/expiration_test.go @@ -107,7 +107,7 @@ func TestSessionExpirationRetainsCheckpointHistory(t *testing.T) { store, source := lifecycleFixture(t) actors := &lifecycleTestActors{actors: map[string]*ateapipb.Actor{}} _, checkpointID := lifecycleForkFixture(t, store, actors, source) - interactions := NewInteractionService(store, nil, NewService(store, serviceTestAuthorizer{}, nil), false) + interactions := NewInteractionService(store, nil, NewService(store, serviceTestAuthorizer{}, nil), false, nil) agent := types.NamespacedName{Namespace: source.GetAgent().Ref.Namespace, Name: source.GetAgent().Ref.Name} // Real ownership filtering must hide an existing conversation in exactly // the same way as a missing or expired conversation, without reading history. diff --git a/go/core/internal/service/session/interactions.go b/go/core/internal/service/session/interactions.go index 464db60622..86667acfc1 100644 --- a/go/core/internal/service/session/interactions.go +++ b/go/core/internal/service/session/interactions.go @@ -48,14 +48,15 @@ type agentService interface { // state. It never connects to actors: the gateway owns dispatch and observation. // These operations form the persistence boundary for a separate API service. type InteractionService struct { - store interactionStore - agents agentService - sessions *Service - allowPushHTTP bool + store interactionStore + agents agentService + sessions *Service + allowPushHTTP bool + quiescenceWake chan<- struct{} } -func NewInteractionService(store interactionStore, agents agentService, sessions *Service, allowPushHTTP bool) *InteractionService { - return &InteractionService{store: store, agents: agents, sessions: sessions, allowPushHTTP: allowPushHTTP} +func NewInteractionService(store interactionStore, agents agentService, sessions *Service, allowPushHTTP bool, quiescenceWake chan<- struct{}) *InteractionService { + return &InteractionService{store: store, agents: agents, sessions: sessions, allowPushHTTP: allowPushHTTP, quiescenceWake: quiescenceWake} } // PreparedSend either grants a new dispatch or returns an already accepted task. @@ -257,6 +258,12 @@ func (s *InteractionService) RevokeSend(ctx context.Context, agent types.Namespa if err != nil { return false, interactionStoreError(ctx, err) } + // Revocation can leave cleanup even when a continuation saved its input + // without starting execution. The durable work, not the hint, owns recovery. + select { + case s.quiescenceWake <- struct{}{}: + default: + } return revoked, nil } diff --git a/go/core/internal/service/session/interactions_test.go b/go/core/internal/service/session/interactions_test.go index 21421bc391..caef038f3e 100644 --- a/go/core/internal/service/session/interactions_test.go +++ b/go/core/internal/service/session/interactions_test.go @@ -59,7 +59,7 @@ func TestPrepareSendStampsInboundTimeline(t *testing.T) { if test.accepted { tasks.dispatchErr = database.ErrMessageAccepted } - service := NewInteractionService(tasks, nil, NewService(&serviceTestStore{getResult: stored}, serviceTestAuthorizer{}, nil), false) + service := NewInteractionService(tasks, nil, NewService(&serviceTestStore{getResult: stored}, serviceTestAuthorizer{}, nil), false, nil) message := a2atype.NewMessage(a2atype.MessageRoleUser, a2atype.NewTextPart("Approved")) message.ID, message.ContextID, message.TaskID = original.ID, id, test.taskID apia2a.SetTimelinePosition(message, time.Unix(1, 0)) @@ -99,7 +99,7 @@ func TestPrepareSendReservesTaskOpeningMessages(t *testing.T) { Operation: apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE, Agent: &apiv1.AgentProvenance{Ref: &apiv1.ResourceReference{Namespace: agent.Namespace, Name: agent.Name}}} tasks := &interactionTestStore{task: &a2atype.Task{ID: "task", ContextID: id, Status: a2atype.TaskStatus{State: a2atype.TaskStateInputRequired}}} - service := NewInteractionService(tasks, nil, NewService(&serviceTestStore{getResult: stored}, serviceTestAuthorizer{}, nil), false) + service := NewInteractionService(tasks, nil, NewService(&serviceTestStore{getResult: stored}, serviceTestAuthorizer{}, nil), false, nil) message := a2atype.NewMessage(a2atype.MessageRoleUser, a2atype.NewTextPart("hello")) message.ID, message.ContextID, message.TaskID = "opening", id, test.taskID _, err := service.PrepareSend(serviceTestContext("alice"), agent, &a2atype.SendMessageRequest{Message: message}) @@ -244,7 +244,7 @@ func TestInteractionsEnforcePermissionsWithoutGateway(t *testing.T) { sessionStore.getResult = nil } tasks := &interactionTestStore{task: &a2atype.Task{ID: "task", ContextID: id, Status: a2atype.TaskStatus{State: a2atype.TaskStateCompleted}}} - service := NewInteractionService(tasks, nil, NewService(sessionStore, authorizer, nil), false) + service := NewInteractionService(tasks, nil, NewService(sessionStore, authorizer, nil), false, nil) want := test.want _, shared := auth.ShareContextFrom(test.ctx) if shared && operation.verb != auth.VerbGet { @@ -273,7 +273,7 @@ func TestInteractionsEnforcePermissionsWithoutGateway(t *testing.T) { } func TestInteractionsRequireAgentForUnfilteredTaskList(t *testing.T) { - service := NewInteractionService(nil, nil, nil, false) + service := NewInteractionService(nil, nil, nil, false, nil) for _, agent := range []types.NamespacedName{ {}, {Namespace: "team-a"}, @@ -298,7 +298,7 @@ func TestExistingConversationDoesNotRequireCurrentAgentGrant(t *testing.T) { ctx := serviceTestContext("alice") message := a2atype.NewMessage(a2atype.MessageRoleUser, a2atype.NewTextPart("continue")) message.ContextID = id - _, err := NewInteractionService(nil, nil, sessions, false).messageSession(ctx, agent, message) + _, err := NewInteractionService(nil, nil, sessions, false, nil).messageSession(ctx, agent, message) if denySession { if !errors.Is(err, a2atype.ErrUnauthorized) { t.Fatalf("error = %v, want Session denial", err) @@ -334,7 +334,7 @@ func TestScheduledMessageIDPrefixIsReserved(t *testing.T) { sessions := NewService(&serviceTestStore{getResult: stored}, &recordingAuthorizer{}, serviceTestWorkflow{}) message := a2atype.NewMessage(a2atype.MessageRoleUser, a2atype.NewTextPart("continue")) message.ID, message.ContextID = test.id, id - _, err := NewInteractionService(nil, nil, sessions, false).messageSession(test.ctx, agent, message) + _, err := NewInteractionService(nil, nil, sessions, false, nil).messageSession(test.ctx, agent, message) if test.wantOK { require.NoError(t, err) } else { @@ -350,7 +350,7 @@ func TestScheduledMessageIDPrefixCannotStartAConversation(t *testing.T) { message := a2atype.NewMessage(a2atype.MessageRoleUser, a2atype.NewTextPart("hello")) message.ID = scheduling.MessageIDPrefix + uuid.NewString() agent := types.NamespacedName{Namespace: "team-a", Name: "assistant"} - _, err := NewInteractionService(nil, nil, sessions, false).PrepareSend(serviceTestContext("alice"), agent, &a2atype.SendMessageRequest{Message: message}) + _, err := NewInteractionService(nil, nil, sessions, false, nil).PrepareSend(serviceTestContext("alice"), agent, &a2atype.SendMessageRequest{Message: message}) require.ErrorIs(t, err, a2atype.ErrInvalidParams) require.Nil(t, store.createInput) } diff --git a/go/core/internal/service/session/quiescence.go b/go/core/internal/service/session/quiescence.go index 283e70c7cc..624b7bbe86 100644 --- a/go/core/internal/service/session/quiescence.go +++ b/go/core/internal/service/session/quiescence.go @@ -3,9 +3,13 @@ package session import ( "context" "errors" + "fmt" "sync" "time" + ateapipb "go.ate.dev/substrate/api/ateapi/v1" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" "sigs.k8s.io/controller-runtime/pkg/manager" "github.com/kagent-dev/kagent/go/core/internal/database" @@ -37,13 +41,11 @@ func (w *ActorWorkflow) Start(ctx context.Context) error { workers.Go(func() { timer := time.NewTimer(w.recoveryInterval) defer timer.Stop() + var afterSessionID string for ctx.Err() == nil { - work, err := w.store.ClaimSessionQuiescence(ctx) + nextSessionID, err := w.quiesceNextSession(ctx, afterSessionID) + afterSessionID = nextSessionID if err == nil { - // Hand off a hint before runtime I/O so a single settlement - // burst can enlist all workers, even if this actor is slow. - w.wakeQuiescence() - w.quiesceIdleSession(ctx, work) continue } delay := min(time.Second, w.recoveryInterval) @@ -65,6 +67,7 @@ func (w *ActorWorkflow) Start(ctx context.Context) error { case <-timer.C: case <-w.quiescenceWake: } + afterSessionID = "" } }) } @@ -72,6 +75,84 @@ func (w *ActorWorkflow) Start(ctx context.Context) error { return nil } +// quiesceNextSession processes one abandoned delivery or task awaiting suspension. +// It returns the last Session ID scanned; ErrNotFound means neither has work. +func (w *ActorWorkflow) quiesceNextSession(ctx context.Context, afterSessionID string) (string, error) { + cleanup, err := w.store.GetPendingSessionDispatchCleanup(ctx, afterSessionID) + if err == nil { + // Visit each candidate once per scan. An unfinished restore must not + // busy-loop or hide later Sessions while it waits for the next scan. + return cleanup.Session.Id, w.quiesceAbandonedDispatch(ctx, cleanup) + } + if !errors.Is(err, database.ErrNotFound) { + return afterSessionID, err + } + work, err := w.store.ClaimSessionQuiescence(ctx) + if err != nil { + return afterSessionID, err + } + // Hand off a hint before runtime I/O so a burst can enlist all workers. + w.wakeQuiescence() + w.quiesceIdleSession(ctx, work) + return afterSessionID, nil +} + +// quiesceAbandonedDispatch suspends a running actor after canceled or expired +// delivery. New sends remain allowed until the worker claims suspension, which +// checks that no turn or newer delivery has started. +func (w *ActorWorkflow) quiesceAbandonedDispatch(ctx context.Context, candidate *database.SessionDispatchCleanup) error { + runtimeCtx, cancel := context.WithTimeout(ctx, database.RuntimeOperationTimeout) + defer cancel() + revision, actor, err := w.prepareQuiescence(runtimeCtx, candidate.Session) + if err == nil && actor.GetMetadata().GetUid() != candidate.ActorUID { + err = fmt.Errorf("runtime Actor changed before abandoned dispatch cleanup") + } + if err != nil { + if ctx.Err() == nil { + logging.FromContext(ctx).ErrorContext(ctx, "prepare abandoned dispatch cleanup", "session_id", candidate.Session.Id, "error", err) + } + return nil + } + if actor.GetStatus().GetState() != ateapipb.ActorState_ACTOR_STATE_RUNNING { + // SUSPENDED is not proof that canceled ingress cannot resume later. + // RESUMING cannot be suspended yet. Neither warrants an issued claim. + return nil + } + work, err := w.store.ClaimSessionDispatchCleanup(runtimeCtx, candidate) + if errors.Is(err, database.ErrNotFound) || errors.Is(err, database.ErrConflict) { + return nil // A new send or another lifecycle worker won admission. + } + if err != nil { + return err + } + w.wakeQuiescence() + var snapshot *database.SessionTaskSnapshot + for { + snapshot, err = w.suspendQuiescentActor(runtimeCtx, work.Session, revision, actor) + if status.Code(err) != codes.Aborted { + break + } + // Substrate permits retrying an aborted workflow. Retain the same + // claim regardless of earlier effects so no new send races the retry. + select { + case <-runtimeCtx.Done(): + err = runtimeCtx.Err() + case <-time.After(100 * time.Millisecond): + continue + } + break + } + if err != nil { + logging.FromContext(ctx).ErrorContext(ctx, "abandoned dispatch suspension outcome unknown", "session_id", work.Session.Id, "dispatch_id", work.DispatchID, "error", err) + return nil // A possibly issued call cannot safely release admission. + } + w.recordQuiescence(ctx, work.Session.Id, func(finishCtx context.Context) error { + return w.store.FinishSessionDispatchCleanup(finishCtx, work, snapshot) + }) + return nil +} + +// quiesceIdleSession suspends an actor for a claimed task and records its snapshot. 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) @@ -82,16 +163,22 @@ func (w *ActorWorkflow) quiesceIdleSession(ctx context.Context, work *database.S logging.FromContext(ctx).ErrorContext(ctx, "runtime boundary outcome unknown", "session_id", work.Session.Id, "version", work.Version, "error", err) return } - // Keep a known snapshot until its reference is stored. Retry database failures - // without repeating runtime work; give shutdown one bounded completion attempt. + w.recordQuiescence(ctx, work.Session.Id, func(finishCtx context.Context) error { + return w.store.FinishSessionQuiescence(finishCtx, work, snapshot) + }) +} + +// recordQuiescence retries saving a completed suspension without suspending again. +// Each write is bounded and may finish after shutdown; further retries stop. +func (w *ActorWorkflow) recordQuiescence(ctx context.Context, sessionID string, finish func(context.Context) error) { for delay := 100 * time.Millisecond; ; delay = min(2*delay, 5*time.Second) { finishCtx, finishCancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) - err = w.store.FinishSessionQuiescence(finishCtx, work, snapshot) + err := finish(finishCtx) finishCancel() if err == nil { return } - logging.FromContext(ctx).ErrorContext(ctx, "record idle runtime outcome", "session_id", work.Session.Id, "version", work.Version, "error", err) + logging.FromContext(ctx).ErrorContext(ctx, "record idle runtime outcome", "session_id", sessionID, "error", err) if errors.Is(err, database.ErrNotFound) || errors.Is(err, database.ErrConflict) { return } diff --git a/go/core/internal/service/session/quiescence_test.go b/go/core/internal/service/session/quiescence_test.go index 99e8852d72..562b97e1dc 100644 --- a/go/core/internal/service/session/quiescence_test.go +++ b/go/core/internal/service/session/quiescence_test.go @@ -18,6 +18,7 @@ import ( apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" "github.com/kagent-dev/kagent/go/core/internal/database" + "github.com/kagent-dev/kagent/go/core/internal/substrate" ) type idleQuiescenceStore struct { @@ -32,6 +33,10 @@ type idleQuiescenceStore struct { var _ workflowStore = (*idleQuiescenceStore)(nil) +func (*idleQuiescenceStore) GetPendingSessionDispatchCleanup(context.Context, string) (*database.SessionDispatchCleanup, error) { + return nil, database.ErrNotFound +} + func (s *idleQuiescenceStore) ClaimSessionQuiescence(context.Context) (*database.SessionQuiescence, error) { s.claims.Add(1) if s.claimErr != nil { @@ -69,6 +74,93 @@ func TestConfiguredQuiescenceInterval(t *testing.T) { }) } +func TestDispatchCleanupScanSkipsActorsUntilRecovery(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + // Given two pending cleanups whose actors are not ready to suspend. + store := &dispatchScanStore{ + idleQuiescenceStore: &idleQuiescenceStore{}, + revision: &database.RuntimeRevision{Revision: "revision-1", ActorTemplateAtespace: "team-a", ActorTemplateName: "assistant"}, + } + base := &lifecycleTestActors{actors: map[string]*ateapipb.Actor{}} + for i, state := range []ateapipb.ActorState{ateapipb.ActorState_ACTOR_STATE_RESUMING, ateapipb.ActorState_ACTOR_STATE_SUSPENDED} { + id := uuid.UUID{15: byte(i + 1)}.String() + session := &apiv1.Session{Id: id, PreparedRevision: store.revision.Revision} + actor, err := base.CreateActor(t.Context(), store.revision.ActorTemplateAtespace, substrate.ActorName(id), store.revision.ActorTemplateAtespace, store.revision.ActorTemplateName) + require.NoError(t, err) + base.actors[actorKey(actor.Metadata.Atespace, actor.Metadata.Name)].Status.State = state + store.pending = append(store.pending, &database.SessionDispatchCleanup{Session: session, DispatchID: uuid.New(), ActorUID: actor.Metadata.Uid}) + } + release := make(chan struct{}) + close(release) + actors := &gatedQuiescenceActors{lifecycleTestActors: base, release: release, calls: map[string]int{}} + actorReads := func(name string) int { + actors.mu.Lock() + defer actors.mu.Unlock() + return actors.calls[name] + } + ctx, cancel := context.WithCancel(t.Context()) + done := make(chan error, 1) + go func() { done <- NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute).Start(ctx) }() + defer func() { cancel(); require.NoError(t, <-done) }() + + // When the workers scan, assert that the first skipped actor does not + // hide the later actor or immediately wake another scan. + synctest.Wait() + require.EqualValues(t, 12, store.lookups.Load(), "each worker should visit both candidates and reach the end of its scan") + for _, work := range store.pending { + require.Equal(t, 4, actorReads(substrate.ActorName(work.Session.Id))) + } + time.Sleep(59 * time.Second) + synctest.Wait() + require.EqualValues(t, 12, store.lookups.Load(), "skipped actors must not cause a busy loop between recovery scans") + + // The durable markers are unchanged. The next recovery scan must + // nevertheless revisit both actors, including the first skipped one. + time.Sleep(time.Second) + synctest.Wait() + require.EqualValues(t, 24, store.lookups.Load()) + for _, work := range store.pending { + require.Equal(t, 8, actorReads(substrate.ActorName(work.Session.Id))) + } + }) +} + +type dispatchScanStore struct { + *idleQuiescenceStore + revision *database.RuntimeRevision + pending []*database.SessionDispatchCleanup + lookups atomic.Int32 +} + +var _ workflowStore = (*dispatchScanStore)(nil) + +func (s *dispatchScanStore) GetPendingSessionDispatchCleanup(_ context.Context, afterSessionID string) (*database.SessionDispatchCleanup, error) { + if s.lookups.Add(1) > 100 { + // Bound a broken scan so the regression fails its assertions instead + // of hanging synctest.Wait in a goroutine that never becomes idle. + return nil, status.Error(codes.Unavailable, "cleanup scan did not advance") + } + for _, work := range s.pending { + if work.Session.Id > afterSessionID { + return work, nil + } + } + return nil, database.ErrNotFound +} + +func (s *dispatchScanStore) GetRuntimeRevision(context.Context, string) (*database.RuntimeRevision, error) { + return s.revision, nil +} + +func (s *dispatchScanStore) GetSessionForRuntime(_ context.Context, sessionID, actorUID string) (*apiv1.Session, error) { + for _, work := range s.pending { + if work.Session.Id == sessionID && work.ActorUID == actorUID { + return work.Session, nil + } + } + return nil, database.ErrNotFound +} + func TestQuiescenceWakeRacesIdleScan(t *testing.T) { synctest.Test(t, func(t *testing.T) { ctx, cancel := context.WithCancel(t.Context()) diff --git a/go/core/internal/service/session/quiescence_wake_test.go b/go/core/internal/service/session/quiescence_wake_test.go index 691235c5d1..57de067555 100644 --- a/go/core/internal/service/session/quiescence_wake_test.go +++ b/go/core/internal/service/session/quiescence_wake_test.go @@ -13,6 +13,9 @@ import ( "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" + "k8s.io/apimachinery/pkg/types" apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" runtimev1 "github.com/kagent-dev/kagent/go/api/gen/kagent/runtime/v1" @@ -278,3 +281,283 @@ func TestQuiescenceCancellationRetainsClaim(t *testing.T) { require.NoError(t, err) require.False(t, pending, "a retained claim must not be treated as retryable work") } + +func TestAbandonedDispatchWakesLifecycleWorker(t *testing.T) { + // Given a running actor whose first delivery has not accepted a task. + store, actors, session := dispatchCleanupFixture(t) + wake := make(chan struct{}, 1) + writes := &observedDispatchCleanupStore{observedQuiescenceStore: &observedQuiescenceStore{lifecycleTestStore: store}} + startQuiescenceWorker(t, NewActorWorkflow(writes, actors, wake, time.Hour)) + require.Eventually(t, func() bool { return writes.idle.Load() >= 4 }, 5*time.Second, time.Millisecond) + message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("hello")) + message.ContextID = session.ContextId + dispatchID := uuid.New() + require.NoError(t, store.ReserveSessionDispatch(t.Context(), session.Id, dispatchID, message.ID)) + service := NewInteractionService(store, nil, NewService(store, serviceTestAuthorizer{}, nil), false, wake) + + // When gateway cleanup revokes the delivery, its commit wakes the worker. + agent := types.NamespacedName{Namespace: session.Agent.Ref.Namespace, Name: session.Agent.Ref.Name} + notAccepted, err := service.RevokeSend(serviceTestContext(session.Creator), agent, message, dispatchID) + require.NoError(t, err) + require.True(t, notAccepted) + + // Assert that suspension does not wait for the recovery interval, create a + // synthetic task, or change the Session's logical availability. + require.Eventually(t, func() bool { return writes.finishedCleanup.Load() == 1 }, 5*time.Second, time.Millisecond) + actor, err := actors.GetActor(t.Context(), agent.Namespace, substrate.ActorName(session.Id)) + require.NoError(t, err) + require.Equal(t, ateapipb.ActorState_ACTOR_STATE_SUSPENDED, actor.Status.State) + stored, err := store.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, stored.State) + tasks, total, err := store.ListSessionTasks(t.Context(), session.Id, "", "", nil, 10, nil) + require.NoError(t, err) + require.Empty(t, tasks) + require.Zero(t, total) + require.NoError(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "next-message")) +} + +func TestAbandonedDispatchWaitsForRestore(t *testing.T) { + for _, state := range []ateapipb.ActorState{ateapipb.ActorState_ACTOR_STATE_SUSPENDED, ateapipb.ActorState_ACTOR_STATE_RESUMING} { + t.Run("Given an actor in "+state.String()+", when restore finishes, assert that cleanup suspends it", func(t *testing.T) { + store, base, session := dispatchCleanupFixture(t) + base.mu.Lock() + base.actors[actorKey(session.Agent.Ref.Namespace, substrate.ActorName(session.Id))].Status.State = state + base.mu.Unlock() + abandonDispatch(t, store, session) + actors := &dispatchCleanupTestActors{lifecycleTestActors: base} + writes := &observedDispatchCleanupStore{observedQuiescenceStore: &observedQuiescenceStore{lifecycleTestStore: store}} + startQuiescenceWorker(t, NewActorWorkflow(writes, actors, make(chan struct{}, 1), 20*time.Millisecond)) + require.Eventually(t, func() bool { return actors.readCalls.Load() > 0 }, 5*time.Second, time.Millisecond) + require.Zero(t, actors.suspendCalls.Load()) + // Observing an unfinished restore must not leave an issued claim that + // prevents a replacement delivery from acquiring the Session. + dispatchID := uuid.New() + require.NoError(t, store.ReserveSessionDispatch(t.Context(), session.Id, dispatchID, "replacement")) + _, err := store.RevokeSessionDispatch(t.Context(), session.Id, dispatchID, "replacement") + require.NoError(t, err) + _, err = base.ResumeActor(t.Context(), session.Agent.Ref.Namespace, substrate.ActorName(session.Id)) + require.NoError(t, err) + require.Eventually(t, func() bool { return writes.finishedCleanup.Load() == 1 }, 5*time.Second, time.Millisecond) + require.EqualValues(t, 1, actors.suspendCalls.Load()) + }) + } +} + +func TestSavedWaitingInputStillWakesDispatchCleanup(t *testing.T) { + // Given a continuation whose input was saved without starting execution. + store, actors, session := dispatchCleanupFixture(t) + boundary := stageQuiescenceTask(t, store.Client, session, a2a.TaskStateInputRequired) + require.NoError(t, store.SettleSessionTask(t.Context(), session.Id, boundary.TaskId, boundary.Version)) + waiting, version, err := store.GetVersionedSessionTask(t.Context(), session.Id, boundary.TaskId) + require.NoError(t, err) + message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("continue")) + message.ContextID, message.TaskID = session.ContextId, waiting.ID + dispatchID := uuid.New() + require.NoError(t, store.ReserveSessionDispatch(t.Context(), session.Id, dispatchID, "")) + waiting.History = append(waiting.History, message) + hash := sha256.Sum256([]byte(message.ID)) + _, err = store.UpdateSessionTask(t.Context(), session.Id, version, hash[:], waiting, message, dispatchID.String()) + require.NoError(t, err) + wake := make(chan struct{}, 1) + writes := &observedDispatchCleanupStore{observedQuiescenceStore: &observedQuiescenceStore{lifecycleTestStore: store}} + startQuiescenceWorker(t, NewActorWorkflow(writes, actors, wake, time.Hour)) + require.Eventually(t, func() bool { return writes.idle.Load() >= 4 }, 5*time.Second, time.Millisecond) + service := NewInteractionService(store, nil, NewService(store, serviceTestAuthorizer{}, nil), false, wake) + + // When revocation finds that accepted input, it must still signal cleanup. + agent := types.NamespacedName{Namespace: session.Agent.Ref.Namespace, Name: session.Agent.Ref.Name} + notAccepted, err := service.RevokeSend(serviceTestContext(session.Creator), agent, message, dispatchID) + require.NoError(t, err) + require.False(t, notAccepted, "persisted input must not be reported safe to resend") + + // Assert that the actor suspends while the accepted waiting task survives. + require.Eventually(t, func() bool { return writes.finishedCleanup.Load() == 1 }, 5*time.Second, time.Millisecond) + stored, err := store.GetSessionTask(t.Context(), session.Id, string(waiting.ID), nil) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateInputRequired, stored.Status.State) + require.Equal(t, message.ID, stored.History[len(stored.History)-1].ID) + actor, err := actors.GetActor(t.Context(), agent.Namespace, substrate.ActorName(session.Id)) + require.NoError(t, err) + require.Equal(t, ateapipb.ActorState_ACTOR_STATE_SUSPENDED, actor.Status.State) + require.NoError(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "next-message")) +} + +func TestAbandonedDispatchReadFailureKeepsAdmissionOpen(t *testing.T) { + // Given a due cleanup whose actor cannot currently be read. + store, base, session := dispatchCleanupFixture(t) + work := abandonDispatch(t, store, session) + actors := &dispatchCleanupTestActors{lifecycleTestActors: base, readErr: status.Error(codes.Unavailable, "actor lookup unavailable")} + workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute) + + // When preparation fails, assert that no suspension was issued and a new + // delivery can supersede the deferred cleanup. + require.NoError(t, workflow.quiesceAbandonedDispatch(t.Context(), work)) + require.Zero(t, actors.suspendCalls.Load()) + require.NoError(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "next-message")) +} + +func TestAbandonedDispatchCancellationBeforeClaimKeepsAdmissionOpen(t *testing.T) { + store, base, session := dispatchCleanupFixture(t) + work := abandonDispatch(t, store, session) + actors := &gatedQuiescenceActors{lifecycleTestActors: base, release: make(chan struct{}), calls: map[string]int{}} + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + finished := make(chan error, 1) + go func() { + finished <- NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute).quiesceAbandonedDispatch(ctx, work) + }() + require.Eventually(t, func() bool { return actors.entered.Load() == 1 }, 5*time.Second, time.Millisecond) + + // Cancel the worker while it is still reading the actor, before it obtains + // permission to mutate it. The canceled attempt must not close admission. + cancel() + select { + case err := <-finished: + require.NoError(t, err) + case <-time.After(5 * time.Second): + t.Fatal("canceled cleanup preparation did not stop") + } + require.NoError(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "next-message")) +} + +func TestAbandonedDispatchCannotSuspendNewDelivery(t *testing.T) { + store, base, session := dispatchCleanupFixture(t) + work := abandonDispatch(t, store, session) + dispatchID := uuid.New() + var once sync.Once + actors := &dispatchCleanupTestActors{lifecycleTestActors: base, afterRead: func() { + // A new delivery wins after the worker reads RUNNING but before it + // claims permission to suspend that actor. + once.Do(func() { + require.NoError(t, store.ReserveSessionDispatch(t.Context(), session.Id, dispatchID, "next-message")) + }) + }} + workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute) + require.NoError(t, workflow.quiesceAbandonedDispatch(t.Context(), work)) + + require.Zero(t, actors.suspendCalls.Load()) + message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("next")) + message.ID, message.ContextID = "next-message", session.ContextId + task := a2a.NewSubmittedTask(message, message) + hash := sha256.Sum256([]byte(task.ID)) + _, err := store.CreateRuntimeTask(t.Context(), session.Id, hash[:], task, dispatchID.String()) + require.NoError(t, err, "the winning delivery must still be able to accept its task") +} + +func TestAbandonedDispatchRetriesKnownFailures(t *testing.T) { + for _, scenario := range []struct { + name string + abortedCalls int32 + finishFailures int32 + }{ + {name: "Substrate lease contention", abortedCalls: 2}, + {name: "database completion failure", finishFailures: 2}, + } { + t.Run("Given "+scenario.name+", when cleanup runs, assert that one suspension completes", func(t *testing.T) { + store, base, session := dispatchCleanupFixture(t) + work := abandonDispatch(t, store, session) + actors := &dispatchCleanupTestActors{lifecycleTestActors: base, abortedCalls: scenario.abortedCalls} + writes := &observedDispatchCleanupStore{observedQuiescenceStore: &observedQuiescenceStore{lifecycleTestStore: store}, finishFailures: scenario.finishFailures} + require.NoError(t, NewActorWorkflow(writes, actors, make(chan struct{}, 1), time.Minute).quiesceAbandonedDispatch(t.Context(), work)) + + require.EqualValues(t, 1, writes.finishedCleanup.Load()) + require.Equal(t, scenario.abortedCalls+1, actors.suspendCalls.Load()) + require.Equal(t, scenario.finishFailures+1, writes.finishAttempts.Load()) + require.NoError(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "next-message")) + }) + } +} + +func TestAbandonedDispatchRetainsUncertainSuspension(t *testing.T) { + // Given a suspension whose result is lost after it changes the actor. + store, base, session := dispatchCleanupFixture(t) + work := abandonDispatch(t, store, session) + actors := &dispatchCleanupTestActors{lifecycleTestActors: base, mutationErr: status.Error(codes.Unavailable, "suspend response lost")} + require.NoError(t, NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute).quiesceAbandonedDispatch(t.Context(), work)) + + // Assert that neither a new delivery nor another cleanup worker can race + // the potentially issued call by taking over its permission. + require.EqualValues(t, 1, actors.suspendCalls.Load()) + require.ErrorIs(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "next-message"), database.ErrDispatchBusy) + _, err := store.GetPendingSessionDispatchCleanup(t.Context(), "") + require.ErrorIs(t, err, database.ErrNotFound) +} + +func dispatchCleanupFixture(t *testing.T) (*lifecycleTestStore, *lifecycleTestActors, *apiv1.Session) { + t.Helper() + store, session := lifecycleFixture(t) + actors := &lifecycleTestActors{actors: map[string]*ateapipb.Actor{}} + session, err := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute).Create(t.Context(), session) + require.NoError(t, err) + _, err = actors.ResumeActor(t.Context(), session.Agent.Ref.Namespace, substrate.ActorName(session.Id)) + require.NoError(t, err) + return store, actors, session +} + +func abandonDispatch(t *testing.T, store *lifecycleTestStore, session *apiv1.Session) *database.SessionDispatchCleanup { + t.Helper() + dispatchID := uuid.New() + require.NoError(t, store.ReserveSessionDispatch(t.Context(), session.Id, dispatchID, "abandoned-message")) + notAccepted, err := store.RevokeSessionDispatch(t.Context(), session.Id, dispatchID, "abandoned-message") + require.NoError(t, err) + require.True(t, notAccepted) + work, err := store.GetPendingSessionDispatchCleanup(t.Context(), "") + require.NoError(t, err) + return work +} + +type observedDispatchCleanupStore struct { + *observedQuiescenceStore + finishedCleanup atomic.Int32 + finishAttempts atomic.Int32 + finishFailures int32 +} + +var _ workflowStore = (*observedDispatchCleanupStore)(nil) + +func (s *observedDispatchCleanupStore) FinishSessionDispatchCleanup(ctx context.Context, work *database.SessionDispatchCleanup, snapshot *database.SessionTaskSnapshot) error { + if s.finishAttempts.Add(1) <= s.finishFailures { + return status.Error(codes.Unavailable, "database unavailable") + } + err := s.Client.FinishSessionDispatchCleanup(ctx, work, snapshot) + if err == nil { + s.finishedCleanup.Add(1) + } + return err +} + +type dispatchCleanupTestActors struct { + *lifecycleTestActors + afterRead func() + readErr error + mutationErr error + abortedCalls int32 + readCalls atomic.Int32 + suspendCalls atomic.Int32 +} + +var _ actorClient = (*dispatchCleanupTestActors)(nil) + +func (a *dispatchCleanupTestActors) GetActor(ctx context.Context, space, name string) (*ateapipb.Actor, error) { + a.readCalls.Add(1) + if a.readErr != nil { + return nil, a.readErr + } + actor, err := a.lifecycleTestActors.GetActor(ctx, space, name) + if a.afterRead != nil { + a.afterRead() + } + return actor, err +} + +func (a *dispatchCleanupTestActors) SuspendActor(ctx context.Context, space, name string) (*ateapipb.Actor, error) { + if a.suspendCalls.Add(1) <= a.abortedCalls { + return nil, status.Error(codes.Aborted, "actor lease is held") + } + actor, err := a.lifecycleTestActors.SuspendActor(ctx, space, name) + if a.mutationErr != nil { + return nil, a.mutationErr + } + return actor, err +} diff --git a/go/core/internal/service/session/workflow.go b/go/core/internal/service/session/workflow.go index 14e3c61783..ccd7b18d46 100644 --- a/go/core/internal/service/session/workflow.go +++ b/go/core/internal/service/session/workflow.go @@ -20,6 +20,9 @@ type workflowStore interface { ClaimSessionQuiescence(context.Context) (*database.SessionQuiescence, error) HasPendingSessionQuiescence(context.Context) (bool, error) FinishSessionQuiescence(context.Context, *database.SessionQuiescence, *database.SessionTaskSnapshot) error + GetPendingSessionDispatchCleanup(context.Context, string) (*database.SessionDispatchCleanup, error) + ClaimSessionDispatchCleanup(context.Context, *database.SessionDispatchCleanup) (*database.SessionDispatchCleanup, error) + FinishSessionDispatchCleanup(context.Context, *database.SessionDispatchCleanup, *database.SessionTaskSnapshot) error GetSessionForRuntime(context.Context, string, string) (*apiv1.Session, error) GetSessionCheckpointSnapshot(context.Context, string, string) (*database.SessionTaskSnapshot, string, error) GetRuntimeRevision(context.Context, string) (*database.RuntimeRevision, error) @@ -54,18 +57,35 @@ func NewActorWorkflow(store workflowStore, actors actorClient, quiescenceWake ch // logical READY state and records its external snapshot URI. Only a checkpoint // retains a copy after the Actor advances. func (w *ActorWorkflow) Quiesce(ctx context.Context, session *apiv1.Session) (*database.SessionTaskSnapshot, error) { + revision, current, err := w.prepareQuiescence(ctx, session) + if err != nil { + return nil, err + } + return w.suspendQuiescentActor(ctx, session, revision, current) +} + +// prepareQuiescence loads the Session's runtime configuration and verifies its +// actor's identity without changing the actor or blocking new sends. +func (w *ActorWorkflow) prepareQuiescence(ctx context.Context, session *apiv1.Session) (*database.RuntimeRevision, *ateapipb.Actor, error) { revision, err := w.store.GetRuntimeRevision(ctx, session.GetPreparedRevision()) if err != nil { - return nil, fmt.Errorf("load prepared revision: %w", err) + return nil, nil, fmt.Errorf("load prepared revision: %w", err) } atespace, name := revision.ActorTemplateAtespace, substrate.ActorName(session.GetId()) current, err := w.actors.GetActor(ctx, atespace, name) if err != nil { - return nil, err + return nil, nil, err } if err := w.verifyActor(ctx, session, revision, current); err != nil { - return nil, err + return nil, nil, err } + return revision, current, nil +} + +// suspendQuiescentActor suspends the verified actor and returns its durable volume +// snapshot. The caller must prevent new turns until the snapshot is recorded. +func (w *ActorWorkflow) suspendQuiescentActor(ctx context.Context, session *apiv1.Session, revision *database.RuntimeRevision, current *ateapipb.Actor) (*database.SessionTaskSnapshot, error) { + atespace, name := revision.ActorTemplateAtespace, substrate.ActorName(session.GetId()) actor, err := w.actors.SuspendActor(ctx, atespace, name) if err != nil { return nil, fmt.Errorf("suspend Actor %s/%s: %w", atespace, name, err) diff --git a/go/core/pkg/app/app.go b/go/core/pkg/app/app.go index d992c9a2a9..a98d0f9104 100644 --- a/go/core/pkg/app/app.go +++ b/go/core/pkg/app/app.go @@ -384,7 +384,7 @@ func Run(ctx context.Context, input *configv1.ControllerConfig, opts Options) er return err } agents := kubecrud.NewService(manager.GetClient(), authorizer, &kagentv1alpha3.Agent{}, &kagentv1alpha3.AgentList{}, "Agent") - interactions := sessionsvc.NewInteractionService(store, agents, sessions, cfg.GetPush().GetAllowHttp()) + interactions := sessionsvc.NewInteractionService(store, agents, sessions, cfg.GetPush().GetAllowHttp(), quiescenceWake) pushSender := sessionsvc.NewHTTPPushSender(cfg.GetPush().GetSendTimeout().AsDuration(), cfg.GetPush().GetAllowHttp(), cfg.GetPush().GetAllowPrivateNetworks()) pushIssuer := cfg.GetPush().GetIssuer() diff --git a/go/core/pkg/migrations/core/000004_session_dispatch_cleanup.sql b/go/core/pkg/migrations/core/000004_session_dispatch_cleanup.sql new file mode 100644 index 0000000000..3ee79fb9d9 --- /dev/null +++ b/go/core/pkg/migrations/core/000004_session_dispatch_cleanup.sql @@ -0,0 +1,22 @@ +-- +goose Up +ALTER TABLE session + ADD COLUMN dispatch_cleanup_executor_id UUID, + ADD CONSTRAINT session_dispatch_cleanup_claim CHECK ( + dispatch_cleanup_executor_id IS NULL OR dispatch_id IS NOT NULL); + +CREATE INDEX session_dispatch_cleanup_pending_idx ON session (id) + WHERE dispatch_id IS NOT NULL AND dispatch_cleanup_executor_id IS NULL; + +-- +goose Down +-- +goose StatementBegin +DO $$ BEGIN + IF EXISTS (SELECT 1 FROM session WHERE dispatch_cleanup_executor_id IS NOT NULL) THEN + RAISE EXCEPTION 'cannot remove dispatch cleanup while issued work is unresolved'; + END IF; +END $$; +-- +goose StatementEnd + +DROP INDEX session_dispatch_cleanup_pending_idx; +ALTER TABLE session + DROP CONSTRAINT session_dispatch_cleanup_claim, + DROP COLUMN dispatch_cleanup_executor_id; diff --git a/go/core/pkg/migrations/runner_test.go b/go/core/pkg/migrations/runner_test.go index 10fae8fbd3..5834150260 100644 --- a/go/core/pkg/migrations/runner_test.go +++ b/go/core/pkg/migrations/runner_test.go @@ -320,7 +320,7 @@ func TestBuiltinMigrationsRoundTrip(t *testing.T) { for _, source := range sources { expected := []int64{0, 1} if source.Name == "core" { - expected = append(expected, 2, 3) + expected = append(expected, 2, 3, 4) } if versions := testVersions(t, dsn, source.TrackingTable); !slices.Equal(versions, expected) { t.Fatalf("%s versions = %v, want %v", source.Name, versions, expected) @@ -337,6 +337,9 @@ func TestBuiltinMigrationsRoundTrip(t *testing.T) { if !testColumnExists(t, dsn, "session", "deletion_reason") { t.Fatal("session deletion reason is missing") } + if !testColumnExists(t, dsn, "session", "dispatch_cleanup_executor_id") { + t.Fatal("session dispatch cleanup executor is missing") + } if !testTableExists(t, dsn, "agent_definition") || testTableExists(t, dsn, "agent_template_harness_pair") { t.Fatal("initial migration must create explicit Agent definitions without legacy pairs") } @@ -360,6 +363,25 @@ func TestBuiltinMigrationsRoundTrip(t *testing.T) { VALUES ($1, 'user', 'schedule', $2, $3)`, scheduleID, make([]byte, 32), []byte{}) execSQL(t, dsn, `INSERT INTO scheduled_run_execution (id, scheduled_run_id, manual_request_id, session_id, data, deadline) VALUES ($1, $2, 'manual', $3, $4, statement_timestamp() + interval '1 minute')`, executionID, scheduleID, sessionID, []byte{}) + t.Run("Given issued dispatch cleanup, when rollback is requested, assert that the admission fence remains installed", func(t *testing.T) { + execSQL(t, dsn, `UPDATE session SET dispatch_id = $1, dispatch_expires_at = statement_timestamp(), dispatch_cleanup_executor_id = $1 WHERE id = $1`, sessionID) + defer execSQL(t, dsn, `UPDATE session SET dispatch_id = NULL, dispatch_expires_at = NULL, dispatch_cleanup_executor_id = NULL WHERE id = $1`, sessionID) + for _, source := range sources { + if source.Name != "core" { + continue + } + err := WithProvider(t.Context(), config, source, func(provider *goose.Provider) error { + _, err := provider.DownTo(t.Context(), 3) + return err + }) + if err == nil || !strings.Contains(err.Error(), "issued work is unresolved") { + t.Fatalf("rollback with issued cleanup = %v", err) + } + if !testColumnExists(t, dsn, "session", "dispatch_cleanup_executor_id") { + t.Fatal("failed rollback removed the cleanup fence") + } + } + }) for _, source := range slices.Backward(sources) { if err := WithProvider(context.Background(), config, source, func(provider *goose.Provider) error { _, err := provider.DownTo(context.Background(), 0) diff --git a/go/core/test/byoa2a/main.go b/go/core/test/byoa2a/main.go index 238eabe13b..5f2cd97480 100644 --- a/go/core/test/byoa2a/main.go +++ b/go/core/test/byoa2a/main.go @@ -3,6 +3,8 @@ package main import ( "cmp" "context" + "errors" + "flag" "fmt" "iter" "os" @@ -17,6 +19,53 @@ import ( type executor struct{} +const ( + blockFirstRestoredSendFlag = "block-first-restored-send" + restoredSendMarker = "/data/e2e-restored-send" + restoredSendArmed = "armed" + restoredSendReleased = "released" +) + +type canceledSendInterceptor struct { + a2asrv.PassthroughCallInterceptor +} + +var _ a2asrv.CallInterceptor = canceledSendInterceptor{} + +// Before blocks message delivery until cancellation, before any task is saved. +func (canceledSendInterceptor) Before(ctx context.Context, _ *a2asrv.CallContext, request *a2asrv.Request) (context.Context, any, error) { + if _, sending := request.Payload.(*a2atype.SendMessageRequest); sending { + <-ctx.Done() + return ctx, nil, ctx.Err() + } + return ctx, nil, nil +} + +// restoredSendGate arms the golden snapshot, blocks sends on its first restore, +// and lets subsequent restores execute normally. The marker advances before +// readiness, so RUNNING is a deterministic barrier for canceling an unaccepted +// send. Its persisted value also proves the canceled actor was checkpointed. +func restoredSendGate() ([]a2asrv.RequestHandlerOption, error) { + marker, err := os.ReadFile(restoredSendMarker) + if errors.Is(err, os.ErrNotExist) { + return nil, os.WriteFile(restoredSendMarker, []byte(restoredSendArmed), 0o600) + } + if err != nil { + return nil, err + } + switch string(marker) { + case restoredSendArmed: + if err := os.WriteFile(restoredSendMarker, []byte(restoredSendReleased), 0o600); err != nil { + return nil, err + } + return []a2asrv.RequestHandlerOption{a2asrv.WithCallInterceptors(canceledSendInterceptor{})}, nil + case restoredSendReleased: + return nil, nil + default: + return nil, fmt.Errorf("unexpected restored-send marker %q", marker) + } +} + func (executor) Execute(_ context.Context, request *a2asrv.ExecutorContext) iter.Seq2[a2atype.Event, error] { return func(yield func(a2atype.Event, error) bool) { if !yield(a2atype.NewSubmittedTask(request, request.Message), nil) { @@ -33,6 +82,8 @@ func (executor) Cancel(context.Context, *a2asrv.ExecutorContext) iter.Seq2[a2aty } func main() { + blockFirstRestoredSend := flag.Bool(blockFirstRestoredSendFlag, false, "block sends before persistence on the first restored activation") + flag.Parse() if err := config.MaterializeFromBootstrap(context.Background(), "/config"); err != nil { fmt.Fprintln(os.Stderr, err) os.Exit(1) @@ -48,9 +99,18 @@ func main() { fmt.Fprintln(os.Stderr, err) os.Exit(1) } + var handlerOpts []a2asrv.RequestHandlerOption + if *blockFirstRestoredSend { + handlerOpts, err = restoredSendGate() + if err != nil { + logger.Error("failed to prepare restored-send gate", "error", err) + return + } + } application, err := app.New(app.AppConfig{ControlPlaneURL: runtime.ControlPlaneUrl, MaxContentLength: runtime.MaxContentLength, AgentCard: a2atype.AgentCard{Name: "opaque-byo", Version: "v1", Capabilities: a2atype.AgentCapabilities{Streaming: true}}, Port: fmt.Sprint(cmp.Or(runtime.GetPort(), 80)), AppName: cmp.Or(runtime.Name, "opaque-byo"), Logger: logger, + HandlerOpts: handlerOpts, }, executor{}) if err != nil { logger.Error("failed to create A2A application", "error", err) diff --git a/go/core/test/e2e/README.md b/go/core/test/e2e/README.md index 34e3c4d8d2..72de1942e1 100644 --- a/go/core/test/e2e/README.md +++ b/go/core/test/e2e/README.md @@ -60,6 +60,16 @@ The test requires the actor to remain `RESUMING` before ordinary Session deletio then verifies Session and Actor absence through kagent APIs. Backend cleanup retries are allowed only after the actor reaches `DELETING` or is absent. +`TestSessionSuspendsAfterCanceledDispatch` clones the `byo-e2e` Harness with the +test runtime's `--block-first-restored-send` flag. Its first restored activation +blocks before accepting a task. The test cancels after the Actor reaches +`RUNNING`, then requires suspension, released worker capacity, and empty task +history. A later invocation must complete and suspend normally. A marker in the +runtime's durable directory releases the gate on that later restore, so the +successful invocation also verifies that cleanup saved the runtime state. +Build the BYO fixture with `make build-byo-a2a` and install its digest in +`byo-e2e` before running this test; the gate is disabled for other tests. + The suite exercises the public API against a clean Kind installation. It does not reconcile Kubernetes resources itself: installation creates the Harness fixtures, and each test owns the templates and API resources it creates. diff --git a/go/core/test/e2e/session_resuming_deletion_test.go b/go/core/test/e2e/session_resuming_deletion_test.go index ae90d5e034..939157cc63 100644 --- a/go/core/test/e2e/session_resuming_deletion_test.go +++ b/go/core/test/e2e/session_resuming_deletion_test.go @@ -6,6 +6,8 @@ import ( "testing" "time" + a2atype "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2apb/v1/pbconv" "github.com/stretchr/testify/require" ateapipb "go.ate.dev/substrate/api/ateapi/v1" "google.golang.org/grpc/codes" @@ -20,6 +22,89 @@ import ( "github.com/kagent-dev/kagent/go/core/internal/substrate" ) +func TestSessionSuspendsAfterCanceledDispatch(t *testing.T) { + t.Parallel() + // Given a runtime whose first restored activation blocks before accepting + // input, while later restores use the ordinary executor and TaskStore. + target := interactionTarget(t) + kube := interactionKubeClient(t) + base := &v1alpha3.Harness{} + require.NoError(t, kube.Get(t.Context(), ctrlclient.ObjectKey{Namespace: "kagent", Name: "byo-e2e"}, base)) + harness := &v1alpha3.Harness{ + ObjectMeta: metav1.ObjectMeta{GenerateName: "canceled-dispatch-", Namespace: "kagent"}, + Spec: *base.Spec.DeepCopy(), + } + harness.Spec.Workload.Args = append(harness.Spec.Workload.Args, "--block-first-restored-send") + require.NoError(t, kube.Create(t.Context(), harness)) + t.Cleanup(func() { + if err := kube.Delete(context.Background(), harness); err != nil && !apierrors.IsNotFound(err) { + t.Errorf("delete canceled-dispatch Harness: %v", err) + } + }) + template := &v1alpha3.AgentTemplate{ + ObjectMeta: metav1.ObjectMeta{GenerateName: "canceled-dispatch-", Namespace: "kagent"}, + Spec: v1alpha3.AgentTemplateSpec{Description: "Canceled dispatch lifecycle fixture"}, + } + createAndWaitInteractionTemplateForHarness(t, kube, template, harness.Name) + fixture := newInteractionFixtureForHarnessTemplate(t, target, harness.Name, template.Name) + actorName := substrate.ActorName(fixture.sessionID) + initial, err := findSubstrateActor(fixture.ctx, fixture.system, "kagent", actorName) + require.NoError(t, err) + require.NotNil(t, initial) + require.Equal(t, ateapipb.ActorState_ACTOR_STATE_SUSPENDED, initial.GetStatus().GetState()) + + // When the caller cancels after restore, no task has reached persistence. + _, request := newMessageRequest(t, "hello") + request.Tenant, request.Message.ContextId = fixture.tenant, fixture.sessionID + sendCtx, cancel := context.WithCancel(fixture.ctx) + defer cancel() + sent := make(chan error, 1) + go func() { + _, err := fixture.client.SendMessage(sendCtx, request) + sent <- err + }() + require.NoError(t, wait.PollUntilContextTimeout(fixture.ctx, 100*time.Millisecond, time.Minute, true, func(ctx context.Context) (bool, error) { + select { + case err := <-sent: + return false, fmt.Errorf("send ended before cancellation: %v", err) + default: + } + actor, err := findSubstrateActor(ctx, fixture.system, "kagent", actorName) + if err != nil { + return false, err + } + return actor != nil && actor.GetStatus().GetState() == ateapipb.ActorState_ACTOR_STATE_RUNNING, nil + })) + cancel() + select { + case err := <-sent: + require.Equal(t, codes.Canceled, status.Code(err)) + case <-time.After(5 * time.Second): + t.Fatal("send did not return after cancellation") + } + + // Assert that turnless cleanup releases the allocation without inventing + // task history, and preserves the snapshot needed for the next invocation. + assertActorSuspended(t, fixture) + suspended, err := findSubstrateActor(fixture.ctx, fixture.system, "kagent", actorName) + require.NoError(t, err) + require.NotNil(t, suspended) + require.Nil(t, suspended.GetStatus().GetWorkerAssignment()) + listRequest, err := pbconv.ToProtoListTasksRequest(&a2atype.ListTasksRequest{Tenant: fixture.tenant, ContextID: fixture.contextID}) + require.NoError(t, err) + listed, err := fixture.client.ListTasks(fixture.ctx, listRequest) + require.NoError(t, err) + require.Empty(t, listed.GetTasks()) + current, err := fixture.sessions.GetSession(fixture.ctx, &apiv1.GetSessionRequest{SessionId: fixture.sessionID}) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, current.GetSession().GetState()) + + _, _, task := fixture.send(t, "hello") + require.Equal(t, a2atype.TaskStateCompleted, task.Status.State) + require.Contains(t, taskText(task), "BYO agent response") + assertActorSuspended(t, fixture) +} + // The Go ADK startup delay reproduces the router timeout independently of the // model protocol. The installed Harness is left intact. func TestSessionDeletionAfterInterruptedResume(t *testing.T) { From d30cb0bd468b0a91f2d9ee60a1e0c9d3563a9b67 Mon Sep 17 00:00:00 2001 From: Shashank Ram Date: Wed, 11 Nov 2026 11:09:12 -0800 Subject: [PATCH 2/2] runtime/session: recover interrupted actor resumes \# Description Canceled message delivery can leave an actor stuck resuming without a turn, holding worker capacity until another caller resumes it. - Resume interrupted actors before suspending under the existing cleanup claim. - Queue bounded retries with backoff capped at 15 seconds to free workers. - Preserve ownership across retries and allow explicit force deletion of cleanup. - Cover cancellation, lost replies, competing lifecycle work, and later delivery. \# Change Type ``` /kind fix ``` \# Changelog ```release-note Actors left resuming after canceled message delivery now recover and suspend automatically. ``` \# Additional Notes Signed-off-by: Shashank Ram --- docs/architecture/runtime-and-lifecycle.md | 37 +- go/core/internal/database/session_deletion.go | 9 +- .../database/session_deletion_test.go | 58 ++ .../database/session_dispatch_cleanup.go | 32 +- go/core/internal/database/session_recovery.go | 11 +- go/core/internal/database/task_store_test.go | 104 ++- .../internal/service/session/quiescence.go | 297 +++++-- .../service/session/quiescence_test.go | 33 +- .../service/session/quiescence_wake_test.go | 762 ++++++++++++++++-- go/core/internal/service/session/workflow.go | 22 +- go/core/test/e2e/README.md | 11 +- .../e2e/session_resuming_deletion_test.go | 122 +++ 12 files changed, 1349 insertions(+), 149 deletions(-) diff --git a/docs/architecture/runtime-and-lifecycle.md b/docs/architecture/runtime-and-lifecycle.md index 8312db1816..48296f8cec 100644 --- a/docs/architecture/runtime-and-lifecycle.md +++ b/docs/architecture/runtime-and-lifecycle.md @@ -217,15 +217,25 @@ reservation; saving continuation input without leaving a waiting state retains it. An expired reservation cannot authorize a late first task write. For abandoned dispatches, the worker first reads and verifies the actor without -blocking new messages. Only a RUNNING actor can be claimed for suspension. The +blocking new messages. RUNNING and RESUMING actors can be claimed for cleanup. The claim atomically rechecks the dispatch generation, absence of active tasks, and other lifecycle work. A newer message supersedes unclaimed cleanup; a claimed -suspension blocks new execution until its snapshot is recorded and the dispatch +cleanup blocks new execution until its snapshot is recorded and the dispatch reservation and claim are cleared. Cleanup preserves the existing history and logical READY state. When history exists, its captured head receives the snapshot without adding a task event. -RESUMING or SUSPENDED actors and failed preparation reads remain unclaimed. Each +After claiming, the worker reads the actor again. If it is RESUMING, the worker +retries resume before using the existing suspension path. Once resume succeeds +or suspension starts, that owner only attempts suspension and snapshot recording. +If another suspension completed before this owner began, it releases its claim +but retains the expired dispatch so a later ingress resume is still discovered. +Transient failures use fresh bounded calls and backoff capped at 15 seconds. Claimed retries +and recovery scans share a queue with four workers, so repeatedly failing actors +yield to other Sessions. The same owner and completed phases survive requeueing; +database retries never repeat a completed runtime operation. + +SUSPENDED actors and failed preparation reads remain unclaimed. Each worker visits candidates in Session ID order once per scan, then revisits them on the existing worker wakeups and recovery scans. A SUSPENDED observation alone cannot prove that canceled ingress will not finish restoring later. Such cleanup @@ -248,6 +258,13 @@ 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. +Explicit force deletion may supersede dispatch cleanup, including an uncertain +claim. It atomically enters DELETING and revokes the dispatch, permanently closing +execution admission. Queued cleanup stops when it observes deletion; an already +issued resume or suspend cannot recreate the deleted Actor. Ordinary deletion +waits for cleanup. Checkpoint and claimed task-boundary work still block force +deletion. + `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 @@ -256,8 +273,9 @@ replacements cannot authorize a revert. A runtime that is not crashed, including a Session at rest outside READY, has nothing to recover, so RevertSession returns the Session unchanged. That is also what a retry sees after a revert whose response was lost. A missing or replaced Actor returns -`FAILED_PRECONDITION`, because no retry restores it. A transport failure returns `UNAVAILABLE`. Only an active -lifecycle operation returns `ABORTED`. +`FAILED_PRECONDITION`, because no retry restores it. A transport failure returns +`UNAVAILABLE`; conflicting lifecycle work or claimed dispatch cleanup returns +`ABORTED`. Admission records FAILED/REVERT with a RuntimeCrashed failure under the Session lock. REVERT is its own `RuntimeOperation`: only a confirmed crash starts one, and @@ -283,10 +301,17 @@ Actor UID and returns the Session to READY, so only the generation exposes a stale observation. The recorded UID changes only through a lifecycle operation, which also replaces the generation. -An issued idle suspension does not exclude recovery. It cannot finish against the +An issued task-boundary 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. +Claimed dispatch cleanup does exclude recovery because its owner may still retry +resume and suspension against the same Actor UID. Its owner can release the claim +after a fresh CRASHED observation only if all mutation outcomes were conclusive +and it stops issuing calls. A later successful resume does not erase uncertainty +from an earlier timeout or lost response. Such claims remain fenced; explicit +force deletion is available when the Session cannot safely be recovered. + Execution 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 diff --git a/go/core/internal/database/session_deletion.go b/go/core/internal/database/session_deletion.go index 49f6c3fc3b..497c9f6cac 100644 --- a/go/core/internal/database/session_deletion.go +++ b/go/core/internal/database/session_deletion.go @@ -43,7 +43,9 @@ func beginSessionDeletion(ctx context.Context, tx pgx.Tx, row sessionRow, reason return nil, err } } - // Record deletion and revoke dispatch together so late input cannot start a turn. + // DELETING closes admission permanently, so force can retire dispatch cleanup + // without exposing a newer turn to its in-flight calls. Actor deletion fences + // those calls; the Session tombstone prevents this runtime from being reused. operation, err := startSessionOperation(ctx, tx, row, kind) if err != nil { return nil, err @@ -58,8 +60,8 @@ func beginSessionDeletion(ctx context.Context, tx pgx.Tx, row sessionRow, reason } // admitForcedSessionDeletion checks work that force must not interrupt: -// checkpoint snapshot copying and idle suspension already claimed by a worker. -// It skips turn, dispatch, and runtime-cleanup checks. These are kagent work +// checkpoint snapshot copying and claimed task-boundary suspension. +// It skips turn, dispatch, and dispatch-cleanup checks. These are kagent work // records, not the Substrate actor state; RESUMING alone does not block deletion. // The caller holds the session lock. runtimeOperationNeeded checks other // lifecycle conflicts before this helper runs. @@ -68,7 +70,6 @@ func admitForcedSessionDeletion(ctx context.Context, tx pgx.Tx, row sessionRow) SELECT EXISTS (SELECT 1 FROM session_checkpoint WHERE source_session_id = $1 AND state = 'CREATING') OR EXISTS (SELECT 1 FROM session_task_event WHERE history_id = $2 AND quiescence_pending AND quiescence_executor_id IS NOT NULL) - OR EXISTS (SELECT 1 FROM session WHERE id = $1 AND dispatch_cleanup_executor_id IS NOT NULL) `, pgx.RowTo[bool], row.ID, row.HistoryID) if err != nil { return err diff --git a/go/core/internal/database/session_deletion_test.go b/go/core/internal/database/session_deletion_test.go index 0cb12b1588..bc7dd780cf 100644 --- a/go/core/internal/database/session_deletion_test.go +++ b/go/core/internal/database/session_deletion_test.go @@ -178,6 +178,64 @@ func TestForceDeletionRetainsIdleLifecycleOwnership(t *testing.T) { require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, current.State) } +func TestForceDeletionSupersedesDispatchCleanup(t *testing.T) { + // Given a waiting task whose abandoned dispatch has an issued cleanup claim. + client, session, task, version := deletionTaskFixture(t, a2a.TaskStateInputRequired) + ctx := t.Context() + require.NoError(t, client.SettleSessionTask(ctx, session.Id, string(task.ID), version)) + candidate := abandonedDispatchFixture(t, client, session) + work, err := client.ClaimSessionDispatchCleanup(ctx, candidate) + require.NoError(t, err) + _, err = client.BeginSessionOperation(ctx, session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE, false) + require.ErrorIs(t, err, ErrFailedPrecondition, "ordinary deletion must still wait for issued cleanup") + + // When force deletion is admitted, it retires cleanup while permanently + // closing admission. A stalled cleanup owner cannot keep deletion blocked. + operation, err := client.BeginSessionOperation(ctx, session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE, true) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_DELETING, operation.Instance.State) + _, err = client.GetPendingSessionDispatchCleanup(ctx, "") + require.ErrorIs(t, err, ErrNotFound) + _, err = client.ClaimSessionDispatchCleanup(ctx, candidate) + require.ErrorIs(t, err, ErrNotFound) + require.ErrorIs(t, client.ReleaseSessionDispatchCleanup(ctx, work), ErrConflict) + require.NoError(t, client.FinishSessionDispatchCleanup(ctx, work, dispatchCleanupSnapshot()), "retired cleanup completion must do nothing") + require.ErrorIs(t, client.ReserveSessionDispatch(ctx, session.Id, uuid.New(), "late-input"), ErrConflict) + for _, token := range []string{work.DispatchID.String(), ""} { + _, err = client.UpdateSessionTask(ctx, session.Id, version, taskMutationHash("late update"), task, task, token) + require.ErrorIs(t, err, ErrConflict) + } + row, err := readSession(ctx, client.db, session.Id) + require.NoError(t, err) + stored, err := readSessionTask(ctx, client.db, row.HistoryID, string(task.ID)) + require.NoError(t, err) + require.Nil(t, stored.SnapshotURI, "late cleanup must not record a snapshot after deletion took ownership") + current, err := client.GetSessionOperation(ctx, session.Id, operation.ID) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE, current.Instance.Operation) + + // A failed preparation keeps DELETING and a new deletion generation. Late + // cleanup callbacks cannot release its fence or replace that operation. + _, err = client.FinishSessionOperation(ctx, session.Id, operation.ID, uuid.Nil, "", "", "preparation failed") + require.NoError(t, err) + retry, err := client.BeginSessionOperation(ctx, session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE, false) + require.NoError(t, err) + require.NotEqual(t, operation.ID, retry.ID) + require.ErrorIs(t, client.ReleaseSessionDispatchCleanup(ctx, work), ErrConflict) + require.NoError(t, client.FinishSessionDispatchCleanup(ctx, work, dispatchCleanupSnapshot())) + current, err = client.GetSessionOperation(ctx, session.Id, retry.ID) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_DELETING, current.Instance.State) + require.ErrorIs(t, client.ReserveSessionDispatch(ctx, session.Id, uuid.New(), "late-input"), ErrConflict) + finishAdmittedDeletion(t, client, retry) + require.ErrorIs(t, client.ReleaseSessionDispatchCleanup(ctx, work), ErrConflict) + require.NoError(t, client.FinishSessionDispatchCleanup(ctx, work, dispatchCleanupSnapshot())) + deleted, err := client.BeginSessionOperation(ctx, session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE, false) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_DELETED, deleted.Instance.State) + require.ErrorIs(t, client.ReserveSessionDispatch(ctx, session.Id, uuid.New(), "late-input"), ErrConflict) +} + func TestDeletionSerializesWithMessageAdmission(t *testing.T) { client, session := deletionSessionFixture(t) start := make(chan struct{}) diff --git a/go/core/internal/database/session_dispatch_cleanup.go b/go/core/internal/database/session_dispatch_cleanup.go index a43577bd28..72e036b5a1 100644 --- a/go/core/internal/database/session_dispatch_cleanup.go +++ b/go/core/internal/database/session_dispatch_cleanup.go @@ -13,7 +13,8 @@ import ( // SessionDispatchCleanup identifies an abandoned delivery's runtime cleanup. // An unclaimed candidate grants no execution authority. Claim returns a new -// value whose executor fences runtime work until its known outcome is stored. +// value whose executor fences runtime work until the owner records completion +// or conclusively stops and releases it. type SessionDispatchCleanup struct { Session *apiv1.Session DispatchID uuid.UUID @@ -131,6 +132,35 @@ func (c *Client) ClaimSessionDispatchCleanup(ctx context.Context, candidate *Ses return result, err } +// ReleaseSessionDispatchCleanup releases only the exact claiming owner and +// preserves the expired dispatch for retry or crash recovery. The caller must +// permanently stop runtime mutations with conclusive outcomes before releasing; +// a timeout or transport error does not establish that remote work has ended. +// Missing Sessions return ErrNotFound. Stale, mismatched, or already released +// claims return ErrConflict without changing the current owner's authority. +func (c *Client) ReleaseSessionDispatchCleanup(ctx context.Context, work *SessionDispatchCleanup) error { + if work == nil || work.Session == nil || work.DispatchID == uuid.Nil || work.ExecutorID == uuid.Nil || work.ActorUID == "" { + return ErrFailedPrecondition + } + return c.withTx(ctx, func(tx pgx.Tx) error { + session, err := lockSession(ctx, tx, work.Session.Id) + if err != nil { + return notFoundOr(err) + } + tag, err := tx.Exec(ctx, ` + UPDATE session SET dispatch_cleanup_executor_id = NULL + WHERE id = $1 AND dispatch_id = $2 AND dispatch_cleanup_executor_id = $3 AND actor_uid = $4 + `, session.ID, work.DispatchID, work.ExecutorID, work.ActorUID) + if err != nil { + return err + } + if tag.RowsAffected() != 1 { + return ErrConflict + } + return nil + }) +} + // FinishSessionDispatchCleanup records the claimed snapshot without inventing // history or changing logical Session state. It releases admission atomically // with snapshot metadata and obsolete task-quiescence work. Completed retries diff --git a/go/core/internal/database/session_recovery.go b/go/core/internal/database/session_recovery.go index a0bb345ee1..7ee82b3208 100644 --- a/go/core/internal/database/session_recovery.go +++ b/go/core/internal/database/session_recovery.go @@ -65,7 +65,8 @@ func (c *Client) GetSessionGeneration(ctx context.Context, id string) (*SessionG // BeginCrashedSessionRecovery admits recovery of a confirmed crash and preserves // durable results. generation must be read before the crash was observed: a // concurrent revert keeps the Actor UID and returns to READY, so only a newer -// lifecycle operation shows that the observation is stale. +// lifecycle operation shows that the observation is stale. An issued dispatch +// cleanup blocks recovery until its owner completes or conclusively stops. func (c *Client) BeginCrashedSessionRecovery(ctx context.Context, generation *SessionGeneration, snapshot SessionTaskSnapshot) (*SessionOperation, error) { var operation *SessionOperation err := c.withTx(ctx, func(tx pgx.Tx) error { @@ -91,12 +92,12 @@ func (c *Client) BeginCrashedSessionRecovery(ctx context.Context, generation *Se if current != generation.OperationID { return ErrConflict } - // An issued idle suspension does not block recovery. It cannot finish against - // the crashed Actor and its worker keeps the claim when suspension fails, so - // waiting would block forever. Settlement below retires the claim; a late - // finish changes nothing. + // Task-boundary suspension keeps its existing recovery semantics. Dispatch + // cleanup may still resume and then suspend the Actor, so its owner must + // finish or conclusively stop before recovery can retire the reservation. busy, err := queryOne(ctx, tx, ` SELECT EXISTS (SELECT 1 FROM session_checkpoint WHERE source_session_id = $1 AND state = 'CREATING') + OR EXISTS (SELECT 1 FROM session WHERE id = $1 AND dispatch_cleanup_executor_id IS NOT NULL) `, pgx.RowTo[bool], row.ID) if err != nil { return err diff --git a/go/core/internal/database/task_store_test.go b/go/core/internal/database/task_store_test.go index e634b2597a..793014fc79 100644 --- a/go/core/internal/database/task_store_test.go +++ b/go/core/internal/database/task_store_test.go @@ -581,8 +581,8 @@ func TestDispatchCleanupFencesTasklessSession(t *testing.T) { } _, err = client.BeginSessionOperation(ctx, session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_SUSPEND, false) require.ErrorIs(t, err, ErrFailedPrecondition) - _, err = client.BeginSessionOperation(ctx, session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE, true) - require.ErrorIs(t, err, ErrConflict) + _, err = client.BeginSessionOperation(ctx, session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE, false) + require.ErrorIs(t, err, ErrFailedPrecondition) require.ErrorIs(t, client.FinishSessionDispatchCleanup(ctx, work, nil), ErrFailedPrecondition) require.NoError(t, client.FinishSessionDispatchCleanup(ctx, work, dispatchCleanupSnapshot())) require.NoError(t, client.FinishSessionDispatchCleanup(ctx, work, dispatchCleanupSnapshot())) @@ -885,12 +885,23 @@ func TestDispatchCleanupCapturesForkedHistory(t *testing.T) { } func TestDispatchCleanupRetiresAfterRuntimeRecovery(t *testing.T) { - t.Run("Given issued dispatch cleanup, when the same actor is reverted after a confirmed crash, assert that stale cleanup cannot claim the recovered runtime", func(t *testing.T) { + t.Run("Given issued dispatch cleanup, when crash recovery is requested, assert that only owner release permits recovery", func(t *testing.T) { client, session := dispatchCleanupFixture(t) ctx := t.Context() candidate := abandonedDispatchFixture(t, client, session) work, err := client.ClaimSessionDispatchCleanup(ctx, candidate) require.NoError(t, err) + _, err = client.BeginCrashedSessionRecovery(ctx, observedGeneration(t, client, session.Id, candidate.ActorUID), SessionTaskSnapshot{}) + require.ErrorIs(t, err, ErrConflict) + current, err := client.GetSessionByID(ctx, 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.ErrorIs(t, client.ReserveSessionDispatch(ctx, session.Id, uuid.New(), "next-message"), ErrDispatchBusy) + require.NoError(t, client.ReleaseSessionDispatchCleanup(ctx, work)) + pending, err := client.GetPendingSessionDispatchCleanup(ctx, "") + require.NoError(t, err) + require.Equal(t, work.DispatchID, pending.DispatchID, "release must retain the expired cleanup obligation") completeCrashedRecovery(t, client, session.Id, candidate.ActorUID) _, err = client.GetPendingSessionDispatchCleanup(ctx, "") require.ErrorIs(t, err, ErrNotFound) @@ -902,6 +913,17 @@ func TestDispatchCleanupRetiresAfterRuntimeRecovery(t *testing.T) { require.NoError(t, err) require.NoError(t, client.FinishSessionDispatchCleanup(ctx, claimed, dispatchCleanupSnapshot())) }) + t.Run("Given unclaimed dispatch cleanup, when crash recovery completes, assert that the obligation is retired", func(t *testing.T) { + client, session := dispatchCleanupFixture(t) + ctx := t.Context() + candidate := abandonedDispatchFixture(t, client, session) + completeCrashedRecovery(t, client, session.Id, candidate.ActorUID) + _, err := client.GetPendingSessionDispatchCleanup(ctx, "") + require.ErrorIs(t, err, ErrNotFound) + _, err = client.ClaimSessionDispatchCleanup(ctx, candidate) + require.ErrorIs(t, err, ErrNotFound) + require.NoError(t, client.ReserveSessionDispatch(ctx, session.Id, uuid.New(), "next-message")) + }) t.Run("Given pending cleanup, when explicit suspension and resume complete, assert that the old obligation is retired", func(t *testing.T) { client, session := dispatchCleanupFixture(t) ctx := t.Context() @@ -917,6 +939,82 @@ func TestDispatchCleanupRetiresAfterRuntimeRecovery(t *testing.T) { }) } +func TestDispatchCleanupReleaseRequiresOwner(t *testing.T) { + for _, scenario := range []struct { + name string + change func(*SessionDispatchCleanup) + }{ + {name: "another dispatch", change: func(work *SessionDispatchCleanup) { work.DispatchID = uuid.New() }}, + {name: "another executor", change: func(work *SessionDispatchCleanup) { work.ExecutorID = uuid.New() }}, + {name: "another actor", change: func(work *SessionDispatchCleanup) { work.ActorUID = uuid.NewString() }}, + } { + t.Run("Given claimed cleanup, when "+scenario.name+" releases it, assert that the original owner retains admission", func(t *testing.T) { + client, session := dispatchCleanupFixture(t) + ctx := t.Context() + candidate := abandonedDispatchFixture(t, client, session) + work, err := client.ClaimSessionDispatchCleanup(ctx, candidate) + require.NoError(t, err) + mismatch := *work + scenario.change(&mismatch) + require.ErrorIs(t, client.ReleaseSessionDispatchCleanup(ctx, &mismatch), ErrConflict) + require.ErrorIs(t, client.ReleaseSessionDispatchCleanup(ctx, candidate), ErrFailedPrecondition, "an unclaimed candidate cannot release ownership") + require.ErrorIs(t, client.ReserveSessionDispatch(ctx, session.Id, uuid.New(), "next-message"), ErrDispatchBusy) + require.NoError(t, client.ReleaseSessionDispatchCleanup(ctx, work)) + pending, err := client.GetPendingSessionDispatchCleanup(ctx, "") + require.NoError(t, err) + require.Equal(t, work.DispatchID, pending.DispatchID) + task := newSessionTask(uuid.NewString(), "late-input") + task.ContextID = session.ContextId + _, err = client.CreateRuntimeTask(ctx, session.Id, taskMutationHash("late-input"), task, work.DispatchID.String()) + require.ErrorIs(t, err, ErrFailedPrecondition, "release must not restore the expired dispatch's execution permission") + }) + } + t.Run("Given claimed cleanup, when release names another Session, assert that the original owner retains admission", func(t *testing.T) { + client, session := dispatchCleanupFixture(t) + ctx := t.Context() + candidate := abandonedDispatchFixture(t, client, session) + work, err := client.ClaimSessionDispatchCleanup(ctx, candidate) + require.NoError(t, err) + other, _, err := client.CreateSession(ctx, newSessionRequest(uuid.NewString(), "assistant", "kagent", ""), uuid.NewString()) + require.NoError(t, err) + other, err = markSessionReady(ctx, client, other.Id, "other-runtime.example") + require.NoError(t, err) + mismatch := *work + mismatch.Session = other + require.ErrorIs(t, client.ReleaseSessionDispatchCleanup(ctx, &mismatch), ErrConflict) + require.ErrorIs(t, client.ReserveSessionDispatch(ctx, session.Id, uuid.New(), "next-message"), ErrDispatchBusy) + require.NoError(t, client.ReleaseSessionDispatchCleanup(ctx, work)) + pending, err := client.GetPendingSessionDispatchCleanup(ctx, "") + require.NoError(t, err) + require.Equal(t, work.DispatchID, pending.DispatchID) + }) +} + +func TestDispatchCleanupReleaseRejectsStaleOwner(t *testing.T) { + t.Run("Given cleanup reclaimed after release, when the old owner releases again, assert that the replacement owner stays fenced", func(t *testing.T) { + client, session := dispatchCleanupFixture(t) + ctx := t.Context() + candidate := abandonedDispatchFixture(t, client, session) + first, err := client.ClaimSessionDispatchCleanup(ctx, candidate) + require.NoError(t, err) + require.NoError(t, client.ReleaseSessionDispatchCleanup(ctx, first)) + require.ErrorIs(t, client.ReleaseSessionDispatchCleanup(ctx, first), ErrConflict, "released owners no longer have authority") + second, err := client.ClaimSessionDispatchCleanup(ctx, candidate) + require.NoError(t, err) + require.NotEqual(t, first.ExecutorID, second.ExecutorID) + require.ErrorIs(t, client.ReleaseSessionDispatchCleanup(ctx, first), ErrConflict) + require.ErrorIs(t, client.ReserveSessionDispatch(ctx, session.Id, uuid.New(), "next-message"), ErrDispatchBusy) + require.ErrorIs(t, client.FinishSessionDispatchCleanup(ctx, first, dispatchCleanupSnapshot()), ErrConflict) + require.NoError(t, client.FinishSessionDispatchCleanup(ctx, second, dispatchCleanupSnapshot())) + latest := abandonedDispatchFixture(t, client, session) + third, err := client.ClaimSessionDispatchCleanup(ctx, latest) + require.NoError(t, err) + require.ErrorIs(t, client.ReleaseSessionDispatchCleanup(ctx, second), ErrConflict, "an old generation cannot release newer cleanup") + require.ErrorIs(t, client.ReserveSessionDispatch(ctx, session.Id, uuid.New(), "next-message"), ErrDispatchBusy) + require.NoError(t, client.FinishSessionDispatchCleanup(ctx, third, dispatchCleanupSnapshot())) + }) +} + func TestDispatchCleanupCursorSkipsUnreadyActor(t *testing.T) { t.Run("Given two pending actors, when a worker leaves the first unclaimed, assert that its cursor reaches the next and a new pass revisits the first", func(t *testing.T) { client, first := dispatchCleanupFixture(t) diff --git a/go/core/internal/service/session/quiescence.go b/go/core/internal/service/session/quiescence.go index 624b7bbe86..f4d824c802 100644 --- a/go/core/internal/service/session/quiescence.go +++ b/go/core/internal/service/session/quiescence.go @@ -10,15 +10,37 @@ import ( ateapipb "go.ate.dev/substrate/api/ateapi/v1" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" + "k8s.io/client-go/util/workqueue" "sigs.k8s.io/controller-runtime/pkg/manager" "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" ) var _ manager.Runnable = (*ActorWorkflow)(nil) var _ manager.LeaderElectionRunnable = (*ActorWorkflow)(nil) +const dispatchCleanupMaxRetryDelay = 15 * time.Second + +// quiescenceWork is either a recovery scan cursor or one exclusively owned +// dispatch. Its stable queue identity prevents concurrent attempts by workers. +type quiescenceWork struct { + afterSessionID string + cleanup *dispatchCleanup +} + +// dispatchCleanup retains the owner and completed phases between bounded +// attempts. A finish callback records a known outcome without more actor calls. +type dispatchCleanup struct { + work *database.SessionDispatchCleanup + revision *database.RuntimeRevision + suspending bool + ambiguousMutation bool + finish func(context.Context) error + retryDelay time.Duration +} + // Every API replica can process idle work; PostgreSQL grants each claim once. func (*ActorWorkflow) NeedLeaderElection() bool { return false } @@ -36,71 +58,104 @@ func (w *ActorWorkflow) wakeQuiescence() { // Startup and periodic recovery scans discover work missed by local signals. // Workers drain claims with bounded concurrency; task reads never wait for them. func (w *ActorWorkflow) Start(ctx context.Context) error { + queue := workqueue.NewTypedDelayingQueue[*quiescenceWork]() + scans := make([]*quiescenceWork, 4) var workers sync.WaitGroup - for range 4 { + for i := range scans { + scans[i] = &quiescenceWork{} + queue.Add(scans[i]) workers.Go(func() { - timer := time.NewTimer(w.recoveryInterval) - defer timer.Stop() - var afterSessionID string - for ctx.Err() == nil { - nextSessionID, err := w.quiesceNextSession(ctx, afterSessionID) - afterSessionID = nextSessionID - if err == nil { - continue - } - delay := min(time.Second, w.recoveryInterval) - if errors.Is(err, database.ErrNotFound) { - pending, pendingErr := w.store.HasPendingSessionQuiescence(ctx) - if pendingErr != nil { - err = pendingErr - } else if !pending { - delay = w.recoveryInterval - } - } - if !errors.Is(err, database.ErrNotFound) && ctx.Err() == nil { - logging.FromContext(ctx).ErrorContext(ctx, "claim runtime boundary", "error", err) - } - timer.Reset(delay) - select { - case <-ctx.Done(): + for { + work, shutdown := queue.Get() + if shutdown { return - case <-timer.C: - case <-w.quiescenceWake: } - afterSessionID = "" + if ctx.Err() == nil { + w.processQuiescence(ctx, queue, work) + } + queue.Done(work) } }) } - workers.Wait() - return nil + for { + select { + case <-ctx.Done(): + queue.ShutDown() + workers.Wait() + return nil + case <-w.quiescenceWake: + for _, scan := range scans { + queue.Add(scan) + } + } + } +} + +// processQuiescence yields between claimed attempts so failing actors cannot +// occupy every worker indefinitely. Scan tokens share the same fair queue. +func (w *ActorWorkflow) processQuiescence(ctx context.Context, queue workqueue.TypedDelayingInterface[*quiescenceWork], work *quiescenceWork) { + if work.cleanup != nil { + done, delay, err := w.retryDispatchCleanup(ctx, work.cleanup) + if err != nil { + logging.FromContext(ctx).ErrorContext(ctx, "abandoned dispatch suspension outcome unknown", "session_id", work.cleanup.work.Session.Id, "dispatch_id", work.cleanup.work.DispatchID, "error", err) + } + if !done { + queue.AddAfter(work, delay) + } + return + } + nextSessionID, cleanup, err := w.quiesceNextSession(ctx, work.afterSessionID) + work.afterSessionID = nextSessionID + if cleanup != nil { + queue.Add(&quiescenceWork{cleanup: cleanup}) + } + if err == nil { + queue.Add(work) + return + } + delay := min(time.Second, w.recoveryInterval) + if errors.Is(err, database.ErrNotFound) { + pending, pendingErr := w.store.HasPendingSessionQuiescence(ctx) + if pendingErr != nil { + err = pendingErr + } else if !pending { + delay = w.recoveryInterval + } + } + if !errors.Is(err, database.ErrNotFound) && ctx.Err() == nil { + logging.FromContext(ctx).ErrorContext(ctx, "claim runtime boundary", "error", err) + } + work.afterSessionID = "" + queue.AddAfter(work, delay) } // quiesceNextSession processes one abandoned delivery or task awaiting suspension. -// It returns the last Session ID scanned; ErrNotFound means neither has work. -func (w *ActorWorkflow) quiesceNextSession(ctx context.Context, afterSessionID string) (string, error) { +// It returns the scan cursor and any claimed dispatch for queued execution; +// ErrNotFound means neither has work. +func (w *ActorWorkflow) quiesceNextSession(ctx context.Context, afterSessionID string) (string, *dispatchCleanup, error) { cleanup, err := w.store.GetPendingSessionDispatchCleanup(ctx, afterSessionID) if err == nil { // Visit each candidate once per scan. An unfinished restore must not // busy-loop or hide later Sessions while it waits for the next scan. - return cleanup.Session.Id, w.quiesceAbandonedDispatch(ctx, cleanup) + claimed, err := w.claimAbandonedDispatch(ctx, cleanup) + return cleanup.Session.Id, claimed, err } if !errors.Is(err, database.ErrNotFound) { - return afterSessionID, err + return afterSessionID, nil, err } work, err := w.store.ClaimSessionQuiescence(ctx) if err != nil { - return afterSessionID, err + return afterSessionID, nil, err } // Hand off a hint before runtime I/O so a burst can enlist all workers. w.wakeQuiescence() w.quiesceIdleSession(ctx, work) - return afterSessionID, nil + return afterSessionID, nil, nil } -// quiesceAbandonedDispatch suspends a running actor after canceled or expired -// delivery. New sends remain allowed until the worker claims suspension, which -// checks that no turn or newer delivery has started. -func (w *ActorWorkflow) quiesceAbandonedDispatch(ctx context.Context, candidate *database.SessionDispatchCleanup) error { +// claimAbandonedDispatch verifies the actor and claims abandoned delivery before +// queueing runtime work. The claim prevents new execution through all attempts. +func (w *ActorWorkflow) claimAbandonedDispatch(ctx context.Context, candidate *database.SessionDispatchCleanup) (*dispatchCleanup, error) { runtimeCtx, cancel := context.WithTimeout(ctx, database.RuntimeOperationTimeout) defer cancel() revision, actor, err := w.prepareQuiescence(runtimeCtx, candidate.Session) @@ -111,45 +166,155 @@ func (w *ActorWorkflow) quiesceAbandonedDispatch(ctx context.Context, candidate if ctx.Err() == nil { logging.FromContext(ctx).ErrorContext(ctx, "prepare abandoned dispatch cleanup", "session_id", candidate.Session.Id, "error", err) } - return nil + return nil, nil } - if actor.GetStatus().GetState() != ateapipb.ActorState_ACTOR_STATE_RUNNING { + state := actor.GetStatus().GetState() + if state != ateapipb.ActorState_ACTOR_STATE_RUNNING && state != ateapipb.ActorState_ACTOR_STATE_RESUMING { // SUSPENDED is not proof that canceled ingress cannot resume later. - // RESUMING cannot be suspended yet. Neither warrants an issued claim. - return nil + return nil, nil } work, err := w.store.ClaimSessionDispatchCleanup(runtimeCtx, candidate) + cancel() if errors.Is(err, database.ErrNotFound) || errors.Is(err, database.ErrConflict) { - return nil // A new send or another lifecycle worker won admission. + return nil, nil // A new send or another lifecycle worker won admission. } if err != nil { - return err + return nil, err } w.wakeQuiescence() - var snapshot *database.SessionTaskSnapshot - for { - snapshot, err = w.suspendQuiescentActor(runtimeCtx, work.Session, revision, actor) - if status.Code(err) != codes.Aborted { - break + return &dispatchCleanup{work: work, revision: revision}, nil +} + +// retryDispatchCleanup performs one bounded runtime or persistence attempt. +// A retry retains the same claim and completed phases but yields its worker. +// Terminal errors leave the claim fenced; only force deletion may supersede it. +func (w *ActorWorkflow) retryDispatchCleanup(ctx context.Context, cleanup *dispatchCleanup) (done bool, delay time.Duration, err error) { + if err := ctx.Err(); err != nil && cleanup.finish == nil { + return true, 0, err + } + retryable := true + if cleanup.finish == nil { + attemptCtx, cancel := context.WithTimeout(ctx, database.RuntimeOperationTimeout) + current, readErr := w.store.GetSessionForRuntime(attemptCtx, cleanup.work.Session.Id, cleanup.work.ActorUID) + err = readErr + if errors.Is(err, database.ErrNotFound) || (err == nil && !lifecycleIdle(current)) { + cancel() + // Force deletion may supersede this claim. It never permits a new + // turn or recreates the actor, even if an old call is still in flight. + return true, 0, nil } - // Substrate permits retrying an aborted workflow. Retain the same - // claim regardless of earlier effects so no new send races the retry. - select { - case <-runtimeCtx.Done(): - err = runtimeCtx.Err() - case <-time.After(100 * time.Millisecond): - continue + if err == nil { + err = w.advanceDispatchCleanup(attemptCtx, cleanup) + retryable = retryableCleanupCall(err) } - break + cancel() + } + if err == nil && cleanup.finish != nil { + // Attempt to persist a known outcome before yielding, including on + // shutdown. Failed writes requeue without repeating runtime work. + finishCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + err = cleanup.finish(finishCtx) + cancel() + if err == nil { + return true, 0, nil + } + retryable = !errors.Is(err, database.ErrNotFound) && !errors.Is(err, database.ErrConflict) + } + if err != nil && !retryable { + return true, 0, err + } + if ctx.Err() != nil { + return true, 0, ctx.Err() + } + if err == nil { + cleanup.retryDelay = 0 + return false, 0, nil + } + logging.FromContext(ctx).DebugContext(ctx, "retry abandoned dispatch cleanup", "session_id", cleanup.work.Session.Id, "dispatch_id", cleanup.work.DispatchID, "error", err) + if cleanup.retryDelay == 0 { + cleanup.retryDelay = 100 * time.Millisecond + } else { + cleanup.retryDelay = min(2*cleanup.retryDelay, dispatchCleanupMaxRetryDelay) } + return false, cleanup.retryDelay, nil +} + +// advanceDispatchCleanup chooses the first action from a fresh owned observation. +// Once resume succeeds or suspension starts, it never resumes again. A known +// terminal outcome switches permanently to database completion or owner release. +func (w *ActorWorkflow) advanceDispatchCleanup(ctx context.Context, cleanup *dispatchCleanup) error { + work, revision := cleanup.work, cleanup.revision + atespace, name := revision.ActorTemplateAtespace, substrate.ActorName(work.Session.Id) + actor, err := w.actors.GetActor(ctx, atespace, name) if err != nil { - logging.FromContext(ctx).ErrorContext(ctx, "abandoned dispatch suspension outcome unknown", "session_id", work.Session.Id, "dispatch_id", work.DispatchID, "error", err) - return nil // A possibly issued call cannot safely release admission. + return err + } + if !validActorIdentity(actor, revision, name) || actor.GetMetadata().GetUid() != work.ActorUID { + return fmt.Errorf("runtime Actor changed during abandoned dispatch cleanup") + } + state := actor.GetStatus().GetState() + previousSuspension := !cleanup.suspending && (state == ateapipb.ActorState_ACTOR_STATE_SUSPENDING || state == ateapipb.ActorState_ACTOR_STATE_SUSPENDED) + if state == ateapipb.ActorState_ACTOR_STATE_CRASHED || previousSuspension { + if cleanup.ambiguousMutation { + return fmt.Errorf("runtime Actor reached %s after an uncertain cleanup call", state) + } + // Another suspension may finish between preparation and claiming. Keep + // the expired dispatch so a later canceled-ingress resume is still found. + // A confirmed crash instead permits recovery once this owner stops. + cleanup.finish = func(finishCtx context.Context) error { + return w.store.ReleaseSessionDispatchCleanup(finishCtx, work) + } + return nil + } + if !cleanup.suspending && state == ateapipb.ActorState_ACTOR_STATE_RESUMING { + err = w.resumeQuiescentActor(ctx, work.Session, revision, actor) + if err == nil { + cleanup.suspending = true + } + } else { + if state != ateapipb.ActorState_ACTOR_STATE_RUNNING && state != ateapipb.ActorState_ACTOR_STATE_SUSPENDING && state != ateapipb.ActorState_ACTOR_STATE_SUSPENDED { + return fmt.Errorf("runtime Actor cannot continue abandoned dispatch cleanup in state %s", state) + } + cleanup.suspending = true + var snapshot *database.SessionTaskSnapshot + snapshot, err = w.suspendQuiescentActor(ctx, work.Session, revision, actor) + if err == nil { + cleanup.finish = func(finishCtx context.Context) error { + return w.store.FinishSessionDispatchCleanup(finishCtx, work, snapshot) + } + } + } + // A later successful Resume may use Substrate's RUNNING fast path; + // it cannot prove an earlier timed-out mutation finished remotely. + if ctx.Err() != nil || !conclusiveCleanupCall(err) { + cleanup.ambiguousMutation = true + } + return err +} + +// conclusiveCleanupCall identifies completed Substrate replies, not absence of +// effects. It must never erase uncertainty from an earlier mutation attempt. +func conclusiveCleanupCall(err error) bool { + switch status.Code(err) { + case codes.OK, codes.Aborted, codes.FailedPrecondition: + return true + default: + return false + } +} + +// retryableCleanupCall includes Internal because Substrate can wrap interrupted +// restores with that code. Local identity and snapshot validation errors stop. +func retryableCleanupCall(err error) bool { + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return true + } + switch status.Code(err) { + case codes.Aborted, codes.Unavailable, codes.DeadlineExceeded, codes.Canceled, codes.Internal, codes.ResourceExhausted, codes.FailedPrecondition: + return true + default: + return false } - w.recordQuiescence(ctx, work.Session.Id, func(finishCtx context.Context) error { - return w.store.FinishSessionDispatchCleanup(finishCtx, work, snapshot) - }) - return nil } // quiesceIdleSession suspends an actor for a claimed task and records its snapshot. @@ -168,7 +333,7 @@ func (w *ActorWorkflow) quiesceIdleSession(ctx context.Context, work *database.S }) } -// recordQuiescence retries saving a completed suspension without suspending again. +// recordQuiescence retries recording a known outcome without further actor calls. // Each write is bounded and may finish after shutdown; further retries stop. func (w *ActorWorkflow) recordQuiescence(ctx context.Context, sessionID string, finish func(context.Context) error) { for delay := 100 * time.Millisecond; ; delay = min(2*delay, 5*time.Second) { diff --git a/go/core/internal/service/session/quiescence_test.go b/go/core/internal/service/session/quiescence_test.go index 562b97e1dc..ada7fa8f79 100644 --- a/go/core/internal/service/session/quiescence_test.go +++ b/go/core/internal/service/session/quiescence_test.go @@ -19,6 +19,7 @@ import ( apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" "github.com/kagent-dev/kagent/go/core/internal/database" "github.com/kagent-dev/kagent/go/core/internal/substrate" + "github.com/kagent-dev/kagent/go/pkg/logging" ) type idleQuiescenceStore struct { @@ -31,6 +32,31 @@ type idleQuiescenceStore struct { beforeIdle func() } +// quiesceAbandonedDispatch drives the same queued steps synchronously for tests +// focused on a single cleanup. Worker scheduling is tested through Start. +func (w *ActorWorkflow) quiesceAbandonedDispatch(ctx context.Context, candidate *database.SessionDispatchCleanup) error { + cleanup, err := w.claimAbandonedDispatch(ctx, candidate) + if err != nil || cleanup == nil { + return err + } + for { + done, delay, err := w.retryDispatchCleanup(ctx, cleanup) + if err != nil { + logging.FromContext(ctx).ErrorContext(ctx, "test cleanup stopped", "error", err) + } + if done { + return nil + } + if delay > 0 { + select { + case <-ctx.Done(): + return nil + case <-time.After(delay): + } + } + } +} + var _ workflowStore = (*idleQuiescenceStore)(nil) func (*idleQuiescenceStore) GetPendingSessionDispatchCleanup(context.Context, string) (*database.SessionDispatchCleanup, error) { @@ -82,7 +108,7 @@ func TestDispatchCleanupScanSkipsActorsUntilRecovery(t *testing.T) { revision: &database.RuntimeRevision{Revision: "revision-1", ActorTemplateAtespace: "team-a", ActorTemplateName: "assistant"}, } base := &lifecycleTestActors{actors: map[string]*ateapipb.Actor{}} - for i, state := range []ateapipb.ActorState{ateapipb.ActorState_ACTOR_STATE_RESUMING, ateapipb.ActorState_ACTOR_STATE_SUSPENDED} { + for i, state := range []ateapipb.ActorState{ateapipb.ActorState_ACTOR_STATE_CRASHED, ateapipb.ActorState_ACTOR_STATE_SUSPENDED} { id := uuid.UUID{15: byte(i + 1)}.String() session := &apiv1.Session{Id: id, PreparedRevision: store.revision.Revision} actor, err := base.CreateActor(t.Context(), store.revision.ActorTemplateAtespace, substrate.ActorName(id), store.revision.ActorTemplateAtespace, store.revision.ActorTemplateName) @@ -172,14 +198,15 @@ func TestQuiescenceWakeRacesIdleScan(t *testing.T) { done := make(chan error, 1) go func() { done <- workflow.Start(ctx) }() synctest.Wait() - require.EqualValues(t, 5, store.claims.Load(), "a signal during an empty scan must trigger another claim") + claimsAfterWake := store.claims.Load() + require.Greater(t, claimsAfterWake, int32(4), "a signal during an empty scan must trigger another claim") var senders sync.WaitGroup for range 100 { senders.Go(workflow.wakeQuiescence) } senders.Wait() synctest.Wait() - require.Greater(t, store.claims.Load(), int32(5)) + require.Greater(t, store.claims.Load(), claimsAfterWake) cancel() require.NoError(t, <-done) for range 100 { diff --git a/go/core/internal/service/session/quiescence_wake_test.go b/go/core/internal/service/session/quiescence_wake_test.go index 57de067555..ae3b7082f5 100644 --- a/go/core/internal/service/session/quiescence_wake_test.go +++ b/go/core/internal/service/session/quiescence_wake_test.go @@ -318,32 +318,579 @@ func TestAbandonedDispatchWakesLifecycleWorker(t *testing.T) { } func TestAbandonedDispatchWaitsForRestore(t *testing.T) { - for _, state := range []ateapipb.ActorState{ateapipb.ActorState_ACTOR_STATE_SUSPENDED, ateapipb.ActorState_ACTOR_STATE_RESUMING} { - t.Run("Given an actor in "+state.String()+", when restore finishes, assert that cleanup suspends it", func(t *testing.T) { + t.Run("Given a suspended actor, when restore finishes, assert that cleanup suspends it", func(t *testing.T) { + store, base, session := dispatchCleanupFixture(t) + setDispatchCleanupActorState(base, session, ateapipb.ActorState_ACTOR_STATE_SUSPENDED) + abandonDispatch(t, store, session) + actors := &dispatchCleanupTestActors{lifecycleTestActors: base} + writes := &observedDispatchCleanupStore{observedQuiescenceStore: &observedQuiescenceStore{lifecycleTestStore: store}} + startQuiescenceWorker(t, NewActorWorkflow(writes, actors, make(chan struct{}, 1), 20*time.Millisecond)) + require.Eventually(t, func() bool { return actors.readCalls.Load() > 0 }, 5*time.Second, time.Millisecond) + require.Zero(t, actors.suspendCalls.Load()) + // Observing an unfinished restore must not leave an issued claim that + // prevents a replacement delivery from acquiring the Session. + dispatchID := uuid.New() + require.NoError(t, store.ReserveSessionDispatch(t.Context(), session.Id, dispatchID, "replacement")) + _, err := store.RevokeSessionDispatch(t.Context(), session.Id, dispatchID, "replacement") + require.NoError(t, err) + _, err = base.ResumeActor(t.Context(), session.Agent.Ref.Namespace, substrate.ActorName(session.Id)) + require.NoError(t, err) + require.Eventually(t, func() bool { return writes.finishedCleanup.Load() == 1 }, 5*time.Second, time.Millisecond) + require.EqualValues(t, 1, actors.suspendCalls.Load()) + }) +} + +func TestAbandonedDispatchCompletesInterruptedResume(t *testing.T) { + // Given a canceled delivery with an actor stuck RESUMING and no task. + store, base, session := dispatchCleanupFixture(t) + setDispatchCleanupActorState(base, session, ateapipb.ActorState_ACTOR_STATE_RESUMING) + work := abandonDispatch(t, store, session) + actors := &dispatchCleanupTestActors{lifecycleTestActors: base, resumeAttempt: func(context.Context, int32) error { + require.ErrorIs(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "racing-message"), database.ErrDispatchBusy, + "cleanup must claim admission before retrying resume") + return nil + }} + writes := &observedDispatchCleanupStore{observedQuiescenceStore: &observedQuiescenceStore{lifecycleTestStore: store}} + + // When cleanup retries the interrupted resume, it must then suspend under + // the same claim instead of waiting for another client to drive restoration. + require.NoError(t, NewActorWorkflow(writes, actors, make(chan struct{}, 1), time.Minute).quiesceAbandonedDispatch(t.Context(), work)) + + // Assert that no turn was invented and the next delivery can be accepted. + require.EqualValues(t, 1, actors.resumeCalls.Load()) + require.EqualValues(t, 1, actors.suspendCalls.Load()) + require.EqualValues(t, 1, writes.finishedCleanup.Load()) + actor, err := base.GetActor(t.Context(), session.Agent.Ref.Namespace, substrate.ActorName(session.Id)) + require.NoError(t, err) + require.Equal(t, ateapipb.ActorState_ACTOR_STATE_SUSPENDED, actor.Status.State) + stored, err := store.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, stored.State) + tasks, total, err := store.ListSessionTasks(t.Context(), session.Id, "", "", nil, 10, nil) + require.NoError(t, err) + require.Empty(t, tasks) + require.Zero(t, total) + dispatchID := uuid.New() + message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("next")) + message.ContextID = session.ContextId + require.NoError(t, store.ReserveSessionDispatch(t.Context(), session.Id, dispatchID, message.ID)) + task := a2a.NewSubmittedTask(message, message) + hash := sha256.Sum256([]byte(task.ID)) + _, err = store.CreateRuntimeTask(t.Context(), session.Id, hash[:], task, dispatchID.String()) + require.NoError(t, err) +} + +func TestAbandonedDispatchUsesOwnedActorState(t *testing.T) { + // Given an actor observed RUNNING before cleanup owns admission, whose + // state changes to RESUMING before the owner's first read. + store, base, session := dispatchCleanupFixture(t) + work := abandonDispatch(t, store, session) + var changeState sync.Once + actors := &dispatchCleanupTestActors{lifecycleTestActors: base, afterRead: func() { + changeState.Do(func() { + setDispatchCleanupActorState(base, session, ateapipb.ActorState_ACTOR_STATE_RESUMING) + }) + }} + writes := &observedDispatchCleanupStore{observedQuiescenceStore: &observedQuiescenceStore{lifecycleTestStore: store}} + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + + // When cleanup owns the actor, assert that the fresh state selects resume + // before suspension, instead of trusting the pre-claim observation. + require.NoError(t, NewActorWorkflow(writes, actors, make(chan struct{}, 1), time.Minute).quiesceAbandonedDispatch(ctx, work)) + require.EqualValues(t, 1, actors.resumeCalls.Load()) + require.EqualValues(t, 1, actors.suspendCalls.Load()) + require.EqualValues(t, 1, writes.cleanupClaims.Load()) + require.EqualValues(t, 1, writes.finishedCleanup.Load()) + require.NoError(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "next-message")) +} + +func TestAbandonedDispatchPreservesCleanupAfterConcurrentSuspension(t *testing.T) { + for _, state := range []ateapipb.ActorState{ateapipb.ActorState_ACTOR_STATE_SUSPENDED, ateapipb.ActorState_ACTOR_STATE_SUSPENDING} { + t.Run(state.String(), func(t *testing.T) { + // Given an actor observed RUNNING before cleanup owns admission, whose + // state changes to a concurrent suspension before the owner's first read. store, base, session := dispatchCleanupFixture(t) - base.mu.Lock() - base.actors[actorKey(session.Agent.Ref.Namespace, substrate.ActorName(session.Id))].Status.State = state - base.mu.Unlock() - abandonDispatch(t, store, session) - actors := &dispatchCleanupTestActors{lifecycleTestActors: base} + work := abandonDispatch(t, store, session) + var changeState sync.Once + actors := &dispatchCleanupTestActors{lifecycleTestActors: base, afterRead: func() { + changeState.Do(func() { + setDispatchCleanupActorState(base, session, state) + }) + }} writes := &observedDispatchCleanupStore{observedQuiescenceStore: &observedQuiescenceStore{lifecycleTestStore: store}} - startQuiescenceWorker(t, NewActorWorkflow(writes, actors, make(chan struct{}, 1), 20*time.Millisecond)) - require.Eventually(t, func() bool { return actors.readCalls.Load() > 0 }, 5*time.Second, time.Millisecond) + workflow := NewActorWorkflow(writes, actors, make(chan struct{}, 1), time.Minute) + + // When cleanup sees the concurrent suspension, assert that it releases its + // unissued claim while retaining the same abandoned delivery for another scan. + require.NoError(t, workflow.quiesceAbandonedDispatch(t.Context(), work)) + require.Zero(t, actors.resumeCalls.Load()) require.Zero(t, actors.suspendCalls.Load()) - // Observing an unfinished restore must not leave an issued claim that - // prevents a replacement delivery from acquiring the Session. - dispatchID := uuid.New() - require.NoError(t, store.ReserveSessionDispatch(t.Context(), session.Id, dispatchID, "replacement")) - _, err := store.RevokeSessionDispatch(t.Context(), session.Id, dispatchID, "replacement") + require.Zero(t, writes.finishedCleanup.Load()) + require.EqualValues(t, 1, writes.cleanupClaims.Load()) + pending, err := store.GetPendingSessionDispatchCleanup(t.Context(), "") require.NoError(t, err) + require.Equal(t, work.DispatchID, pending.DispatchID) + require.Equal(t, work.Session.Id, pending.Session.Id) + + // When concurrent suspension finishes and the original ingress resumes, + // assert that the preserved delivery is rediscovered and suspended. + setDispatchCleanupActorState(base, session, ateapipb.ActorState_ACTOR_STATE_SUSPENDED) _, err = base.ResumeActor(t.Context(), session.Agent.Ref.Namespace, substrate.ActorName(session.Id)) require.NoError(t, err) - require.Eventually(t, func() bool { return writes.finishedCleanup.Load() == 1 }, 5*time.Second, time.Millisecond) + require.NoError(t, workflow.quiesceAbandonedDispatch(t.Context(), pending)) + require.Zero(t, actors.resumeCalls.Load()) + require.EqualValues(t, 1, actors.suspendCalls.Load()) + require.EqualValues(t, 2, writes.cleanupClaims.Load()) + require.EqualValues(t, 1, writes.finishedCleanup.Load()) + _, err = store.GetPendingSessionDispatchCleanup(t.Context(), "") + require.ErrorIs(t, err, database.ErrNotFound) + require.NoError(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "next-message")) + }) + } +} + +func TestAbandonedDispatchRejectsChangedResumeIdentity(t *testing.T) { + // Given a cleanup whose resume response belongs to a different actor UID. + store, base, session := dispatchCleanupFixture(t) + setDispatchCleanupActorState(base, session, ateapipb.ActorState_ACTOR_STATE_RESUMING) + work := abandonDispatch(t, store, session) + const replacementActorUID = "replacement-actor-uid" + actors := &dispatchCleanupTestActors{lifecycleTestActors: base, resumedActorUID: replacementActorUID} + writes := &observedDispatchCleanupStore{observedQuiescenceStore: &observedQuiescenceStore{lifecycleTestStore: store}} + + // When resume returns that identity, assert that cleanup neither suspends + // the actor nor releases the claim after the uncertain runtime outcome. + require.NoError(t, NewActorWorkflow(writes, actors, make(chan struct{}, 1), time.Minute).quiesceAbandonedDispatch(t.Context(), work)) + require.EqualValues(t, 1, actors.resumeCalls.Load()) + require.Zero(t, actors.suspendCalls.Load()) + require.Zero(t, writes.finishedCleanup.Load()) + require.ErrorIs(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "next-message"), database.ErrDispatchBusy) +} + +func TestAbandonedDispatchRetriesInterruptedResume(t *testing.T) { + for _, code := range []codes.Code{codes.Aborted, codes.Unavailable, codes.DeadlineExceeded, codes.Canceled, codes.Internal, codes.ResourceExhausted, codes.FailedPrecondition} { + t.Run("Given a resume returning "+code.String()+", when cleanup retries, assert that the actor suspends under one claim", func(t *testing.T) { + store, base, session := dispatchCleanupFixture(t) + setDispatchCleanupActorState(base, session, ateapipb.ActorState_ACTOR_STATE_RESUMING) + work := abandonDispatch(t, store, session) + var firstDeadline time.Time + actors := &dispatchCleanupTestActors{lifecycleTestActors: base, resumeAttempt: func(ctx context.Context, call int32) error { + require.NoError(t, ctx.Err()) + deadline, bounded := ctx.Deadline() + require.True(t, bounded, "each runtime attempt must be bounded") + require.ErrorIs(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "racing-message"), database.ErrDispatchBusy) + if call == 1 { + firstDeadline = deadline + return status.Error(code, "temporary resume failure") + } + require.True(t, deadline.After(firstDeadline), "a retry needs a fresh attempt deadline") + return nil + }} + writes := &observedDispatchCleanupStore{observedQuiescenceStore: &observedQuiescenceStore{lifecycleTestStore: store}} + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + stop := time.AfterFunc(5*time.Second, cancel) + defer stop.Stop() + require.NoError(t, NewActorWorkflow(writes, actors, make(chan struct{}, 1), time.Minute).quiesceAbandonedDispatch(ctx, work)) + + require.EqualValues(t, 2, actors.resumeCalls.Load()) require.EqualValues(t, 1, actors.suspendCalls.Load()) + require.EqualValues(t, 1, writes.cleanupClaims.Load()) + require.EqualValues(t, 1, writes.finishedCleanup.Load()) + require.NoError(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "next-message")) + }) + } +} + +func TestAbandonedDispatchSuspendRetryDoesNotResumeAgain(t *testing.T) { + for _, code := range []codes.Code{codes.Aborted, codes.Unavailable, codes.DeadlineExceeded, codes.Canceled, codes.Internal, codes.ResourceExhausted, codes.FailedPrecondition} { + t.Run("Given a successful resume and suspend returning "+code.String()+", when cleanup retries, assert that it only retries suspension", func(t *testing.T) { + store, base, session := dispatchCleanupFixture(t) + setDispatchCleanupActorState(base, session, ateapipb.ActorState_ACTOR_STATE_RESUMING) + work := abandonDispatch(t, store, session) + actors := &dispatchCleanupTestActors{lifecycleTestActors: base, suspendAttempt: func(_ context.Context, call int32) error { + require.ErrorIs(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "racing-message"), database.ErrDispatchBusy) + if call == 1 { + setDispatchCleanupActorState(base, session, ateapipb.ActorState_ACTOR_STATE_SUSPENDING) + return status.Error(code, "temporary suspension failure") + } + return nil + }} + writes := &observedDispatchCleanupStore{observedQuiescenceStore: &observedQuiescenceStore{lifecycleTestStore: store}} + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + require.NoError(t, NewActorWorkflow(writes, actors, make(chan struct{}, 1), time.Minute).quiesceAbandonedDispatch(ctx, work)) + + require.EqualValues(t, 1, actors.resumeCalls.Load()) + require.EqualValues(t, 2, actors.suspendCalls.Load()) + require.EqualValues(t, 1, writes.cleanupClaims.Load()) + require.EqualValues(t, 1, writes.finishedCleanup.Load()) + require.NoError(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "next-message")) }) } } +func TestAbandonedDispatchRetriesLostSuspensionResponse(t *testing.T) { + // Given an interrupted resume whose cleanup suspends the actor but loses + // the response, leaving a valid snapshot behind an uncertain call outcome. + store, base, session := dispatchCleanupFixture(t) + setDispatchCleanupActorState(base, session, ateapipb.ActorState_ACTOR_STATE_RESUMING) + work := abandonDispatch(t, store, session) + actors := &dispatchCleanupTestActors{lifecycleTestActors: base, suspendAttempt: func(ctx context.Context, call int32) error { + require.ErrorIs(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "racing-message"), database.ErrDispatchBusy) + if call == 1 { + _, err := base.SuspendActor(ctx, session.Agent.Ref.Namespace, substrate.ActorName(session.Id)) + require.NoError(t, err) + return status.Error(codes.Unavailable, "suspension response lost") + } + return nil + }} + writes := &observedDispatchCleanupStore{observedQuiescenceStore: &observedQuiescenceStore{lifecycleTestStore: store}} + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + + // When cleanup retries against SUSPENDED, assert that the same owner + // records completion and admits the next message without resuming again. + require.NoError(t, NewActorWorkflow(writes, actors, make(chan struct{}, 1), time.Minute).quiesceAbandonedDispatch(ctx, work)) + require.EqualValues(t, 1, actors.resumeCalls.Load()) + require.EqualValues(t, 2, actors.suspendCalls.Load()) + require.EqualValues(t, 1, writes.cleanupClaims.Load()) + require.EqualValues(t, 1, writes.finishedCleanup.Load()) + require.NoError(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "next-message")) +} + +func TestAbandonedDispatchRetriesDoNotStarveSettledTasks(t *testing.T) { + // Given four interrupted restores that fail indefinitely, filling every + // lifecycle worker, and a later completed turn on another Session. + const failedRestores = 4 + const abandonedMessage = "abandoned-message" + store, first := lifecycleFixture(t) + base := &lifecycleTestActors{actors: map[string]*ateapipb.Actor{}} + creator := NewActorWorkflow(store, base, make(chan struct{}, 1), time.Hour) + var sessions []*apiv1.Session + for i := range failedRestores { + session := first + if i != 0 { + var err error + session, _, err = store.CreateSession(t.Context(), &apiv1.Session{Id: uuid.NewString(), Creator: "alice", Agent: first.GetAgent()}, uuid.NewString()) + require.NoError(t, err) + } + session, err := creator.Create(t.Context(), session) + require.NoError(t, err) + setDispatchCleanupActorState(base, session, ateapipb.ActorState_ACTOR_STATE_RESUMING) + dispatchID := uuid.New() + require.NoError(t, store.ReserveSessionDispatch(t.Context(), session.Id, dispatchID, abandonedMessage)) + revoked, err := store.RevokeSessionDispatch(t.Context(), session.Id, dispatchID, abandonedMessage) + require.NoError(t, err) + require.True(t, revoked) + sessions = append(sessions, session) + } + var failResume atomic.Bool + failResume.Store(true) + actors := &dispatchCleanupTestActors{lifecycleTestActors: base, resumeAttempt: func(context.Context, int32) error { + if failResume.Load() { + return status.Error(codes.Unavailable, "restore is temporarily unavailable") + } + return nil + }} + writes := &observedDispatchCleanupStore{observedQuiescenceStore: &observedQuiescenceStore{lifecycleTestStore: store}} + workflow := NewActorWorkflow(writes, actors, make(chan struct{}, 1), time.Hour) + startQuiescenceWorker(t, workflow) + require.Eventually(t, func() bool { return writes.cleanupClaims.Load() == failedRestores }, 5*time.Second, time.Millisecond) + + healthy, _, err := store.CreateSession(t.Context(), &apiv1.Session{Id: uuid.NewString(), Creator: "alice", Agent: first.GetAgent()}, uuid.NewString()) + require.NoError(t, err) + healthy, err = creator.Create(t.Context(), healthy) + require.NoError(t, err) + _, err = base.ResumeActor(t.Context(), healthy.Agent.Ref.Namespace, substrate.ActorName(healthy.Id)) + require.NoError(t, err) + request := stageQuiescenceTask(t, store.Client, healthy, a2a.TaskStateCompleted) + require.NoError(t, store.SettleSessionTask(t.Context(), healthy.Id, request.TaskId, request.Version)) + + // When the completed turn wakes the workers, assert that its suspension + // finishes while each failing restore retains its original admission claim. + workflow.wakeQuiescence() + require.Eventually(t, func() bool { return writes.finished.Load() == 1 }, 3*time.Second, time.Millisecond, + "retrying restores must yield workers to another Session's completed turn") + require.Zero(t, writes.finishedCleanup.Load()) + require.EqualValues(t, failedRestores, writes.cleanupClaims.Load()) + for _, session := range sessions { + require.ErrorIs(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "racing-message"), database.ErrDispatchBusy) + } + actor, err := base.GetActor(t.Context(), healthy.Agent.Ref.Namespace, substrate.ActorName(healthy.Id)) + require.NoError(t, err) + require.Equal(t, ateapipb.ActorState_ACTOR_STATE_SUSPENDED, actor.GetStatus().GetState()) + + // When restoration becomes available, assert that the retained owners + // finish cleanup without releasing and reacquiring the claims between tries. + failResume.Store(false) + workflow.wakeQuiescence() + require.Eventually(t, func() bool { return writes.finishedCleanup.Load() == failedRestores }, 10*time.Second, time.Millisecond) + require.EqualValues(t, failedRestores, writes.cleanupClaims.Load()) + for _, session := range sessions { + actor, err := base.GetActor(t.Context(), session.Agent.Ref.Namespace, substrate.ActorName(session.Id)) + require.NoError(t, err) + require.Equal(t, ateapipb.ActorState_ACTOR_STATE_SUSPENDED, actor.GetStatus().GetState()) + require.NoError(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "next-message")) + } +} + +func TestAbandonedDispatchRejectsReplacementBeforeRetry(t *testing.T) { + // Given a failed resume followed by an externally replaced actor. + store, base, session := dispatchCleanupFixture(t) + setDispatchCleanupActorState(base, session, ateapipb.ActorState_ACTOR_STATE_RESUMING) + work := abandonDispatch(t, store, session) + const replacementActorUID = "replacement-actor-uid" + actors := &dispatchCleanupTestActors{lifecycleTestActors: base, resumeAttempt: func(context.Context, int32) error { + base.mu.Lock() + base.actors[actorKey(session.Agent.Ref.Namespace, substrate.ActorName(session.Id))].Metadata.Uid = replacementActorUID + base.mu.Unlock() + return status.Error(codes.Unavailable, "resume response lost") + }} + writes := &observedDispatchCleanupStore{observedQuiescenceStore: &observedQuiescenceStore{lifecycleTestStore: store}} + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + + // When cleanup rereads before retrying, assert that it rejects the changed + // identity without mutating the replacement or releasing the uncertain claim. + require.NoError(t, NewActorWorkflow(writes, actors, make(chan struct{}, 1), time.Minute).quiesceAbandonedDispatch(ctx, work)) + require.EqualValues(t, 1, actors.resumeCalls.Load()) + require.Zero(t, actors.suspendCalls.Load()) + require.Zero(t, writes.finishedCleanup.Load()) + require.ErrorIs(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "next-message"), database.ErrDispatchBusy) +} + +func TestAbandonedDispatchCrashRecoveryWaitsForConclusiveCleanup(t *testing.T) { + // Given a resume whose server confirms a crash, but has not returned yet. + store, base, session := dispatchCleanupFixture(t) + _, err := base.SuspendActor(t.Context(), session.Agent.Ref.Namespace, substrate.ActorName(session.Id)) + require.NoError(t, err) + setDispatchCleanupActorState(base, session, ateapipb.ActorState_ACTOR_STATE_RESUMING) + work := abandonDispatch(t, store, session) + entered, release := make(chan struct{}), make(chan struct{}) + actors := &dispatchCleanupTestActors{lifecycleTestActors: base, resumeAttempt: func(ctx context.Context, _ int32) error { + setDispatchCleanupActorState(base, session, ateapipb.ActorState_ACTOR_STATE_CRASHED) + close(entered) + select { + case <-release: + return status.Error(codes.FailedPrecondition, "actor crashed during resume") + case <-ctx.Done(): + return ctx.Err() + } + }} + writes := &observedDispatchCleanupStore{observedQuiescenceStore: &observedQuiescenceStore{lifecycleTestStore: store}} + workflow := NewActorWorkflow(writes, actors, make(chan struct{}, 1), time.Minute) + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + done := make(chan error, 1) + go func() { done <- workflow.quiesceAbandonedDispatch(ctx, work) }() + select { + case <-entered: + case <-ctx.Done(): + t.Fatal("cleanup did not start resume") + } + + // When recovery races an issued call, assert that it cannot retire the + // cleanup claim. Only the owner may release it after a conclusive result. + _, err = workflow.Revert(t.Context(), session) + require.ErrorIs(t, err, database.ErrConflict) + close(release) + select { + case err := <-done: + require.NoError(t, err) + case <-ctx.Done(): + t.Fatal("cleanup did not finish after the conclusive crash") + } + require.EqualValues(t, 1, actors.resumeCalls.Load()) + require.Zero(t, actors.suspendCalls.Load()) + require.Zero(t, writes.finishedCleanup.Load(), "a crash is not a successful cleanup snapshot") + recovered, err := workflow.Revert(t.Context(), session) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, recovered.State) + require.NoError(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "next-message")) +} + +func TestAbandonedDispatchUncertainCrashAllowsForceDeletion(t *testing.T) { + // Given an interrupted restore whose Internal reply leaves a crashed + // actor and an uncertain cleanup owner, ordinary recovery must stay fenced. + store, base, session := dispatchCleanupFixture(t) + _, err := base.SuspendActor(t.Context(), session.Agent.Ref.Namespace, substrate.ActorName(session.Id)) + require.NoError(t, err) + setDispatchCleanupActorState(base, session, ateapipb.ActorState_ACTOR_STATE_RESUMING) + candidate := abandonDispatch(t, store, session) + deleteEntered, releaseDelete := make(chan struct{}), make(chan struct{}) + actors := &gatedDispatchDeletionActors{ + dispatchCleanupTestActors: &dispatchCleanupTestActors{lifecycleTestActors: base, resumeAttempt: func(context.Context, int32) error { + setDispatchCleanupActorState(base, session, ateapipb.ActorState_ACTOR_STATE_CRASHED) + return status.Error(codes.Internal, "restore failed after the actor crashed") + }}, + deleteEntered: deleteEntered, + releaseDelete: releaseDelete, + } + writes := &observedDispatchCleanupStore{observedQuiescenceStore: &observedQuiescenceStore{lifecycleTestStore: store}} + workflow := NewActorWorkflow(writes, actors, make(chan struct{}, 1), time.Minute) + cleanup, err := workflow.claimAbandonedDispatch(t.Context(), candidate) + require.NoError(t, err) + require.NotNil(t, cleanup) + done, delay, err := workflow.retryDispatchCleanup(t.Context(), cleanup) + require.NoError(t, err) + require.False(t, done, "an Internal reply must retain ownership for retry") + require.Positive(t, delay, "failed restore must wait before retrying") + require.ErrorIs(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "next-message"), database.ErrDispatchBusy) + _, err = workflow.Revert(t.Context(), session) + require.ErrorIs(t, err, database.ErrConflict) + _, err = workflow.Delete(t.Context(), session, false) + require.ErrorIs(t, err, database.ErrFailedPrecondition) + + // When force deletion takes ownership, gate its backend call so a queued + // cleanup retry observes DELETING before the actor has been removed. + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + deleted := make(chan *apiv1.Session, 1) + deleteErrors := make(chan error, 1) + deleteFinished := make(chan struct{}) + go func() { + defer close(deleteFinished) + result, err := workflow.Delete(ctx, session, true) + deleted <- result + deleteErrors <- err + }() + t.Cleanup(func() { + cancel() + <-deleteFinished + }) + select { + case <-deleteEntered: + case err := <-deleteErrors: + t.Fatalf("force deletion did not reach the actor: %v", err) + case <-ctx.Done(): + t.Fatal("force deletion did not start") + } + current, err := store.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_DELETING, current.State) + reads := actors.readCalls.Load() + done, _, err = workflow.retryDispatchCleanup(t.Context(), cleanup) + require.NoError(t, err) + require.True(t, done, "queued cleanup must stop after deletion takes ownership") + require.Equal(t, reads, actors.readCalls.Load(), "cleanup must not contact the runtime after observing deletion") + require.EqualValues(t, 1, actors.resumeCalls.Load()) + require.Zero(t, actors.suspendCalls.Load()) + + // Assert that force deletion finishes while the obsolete cleanup owner + // neither releases admission nor records a successful suspension. + close(releaseDelete) + select { + case err := <-deleteErrors: + require.NoError(t, err) + case <-ctx.Done(): + t.Fatal("force deletion did not finish") + } + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_DELETED, (<-deleted).State) + _, err = base.GetActor(t.Context(), session.Agent.Ref.Namespace, substrate.ActorName(session.Id)) + require.Equal(t, codes.NotFound, status.Code(err)) + _, err = store.GetSessionByID(t.Context(), session.Id) + require.ErrorIs(t, err, database.ErrNotFound) + require.Zero(t, writes.finishedCleanup.Load()) +} + +func TestAbandonedDispatchStopsAfterExternalDeletionStarts(t *testing.T) { + // Given an interrupted resume whose next observation is an externally + // deleting actor, cleanup no longer has a supported transition to retry. + store, base, session := dispatchCleanupFixture(t) + setDispatchCleanupActorState(base, session, ateapipb.ActorState_ACTOR_STATE_RESUMING) + work := abandonDispatch(t, store, session) + actors := &dispatchCleanupTestActors{lifecycleTestActors: base, resumeAttempt: func(context.Context, int32) error { + setDispatchCleanupActorState(base, session, ateapipb.ActorState_ACTOR_STATE_DELETING) + return status.Error(codes.Unavailable, "resume response lost") + }} + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + + // When cleanup rereads that state, assert that it stops without issuing + // another mutation or releasing the uncertain claim. + require.NoError(t, NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute).quiesceAbandonedDispatch(ctx, work)) + require.NoError(t, ctx.Err(), "unsupported actor state must stop retries before the context expires") + require.EqualValues(t, 1, actors.resumeCalls.Load()) + require.Zero(t, actors.suspendCalls.Load()) + require.ErrorIs(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "next-message"), database.ErrDispatchBusy) +} + +func TestAbandonedDispatchCrashKeepsPriorAmbiguity(t *testing.T) { + for _, phase := range []string{"resume", "suspend after a successful resume"} { + t.Run("Given a lost resume response, when the actor crashes during "+phase+", assert that its uncertain claim remains", func(t *testing.T) { + store, base, session := dispatchCleanupFixture(t) + setDispatchCleanupActorState(base, session, ateapipb.ActorState_ACTOR_STATE_RESUMING) + work := abandonDispatch(t, store, session) + actors := &dispatchCleanupTestActors{lifecycleTestActors: base, resumeAttempt: func(_ context.Context, call int32) error { + if call == 1 { + if phase == "resume" { + setDispatchCleanupActorState(base, session, ateapipb.ActorState_ACTOR_STATE_CRASHED) + } + return status.Error(codes.Unavailable, "resume response lost") + } + return nil + }, suspendAttempt: func(context.Context, int32) error { + setDispatchCleanupActorState(base, session, ateapipb.ActorState_ACTOR_STATE_CRASHED) + return status.Error(codes.FailedPrecondition, "actor crashed before suspension") + }} + writes := &observedDispatchCleanupStore{observedQuiescenceStore: &observedQuiescenceStore{lifecycleTestStore: store}} + workflow := NewActorWorkflow(writes, actors, make(chan struct{}, 1), time.Minute) + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + require.NoError(t, workflow.quiesceAbandonedDispatch(ctx, work)) + + require.Zero(t, writes.finishedCleanup.Load()) + require.EqualValues(t, 1, writes.cleanupClaims.Load()) + require.ErrorIs(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "next-message"), database.ErrDispatchBusy) + _, err := workflow.Revert(t.Context(), session) + require.ErrorIs(t, err, database.ErrConflict, "a later successful call must not erase the earlier uncertain outcome") + _, err = store.GetPendingSessionDispatchCleanup(t.Context(), "") + require.ErrorIs(t, err, database.ErrNotFound) + }) + } +} + +func TestAbandonedDispatchRecordsSuspensionDuringShutdown(t *testing.T) { + // Given a cleanup whose successful suspension races worker shutdown. + store, base, session := dispatchCleanupFixture(t) + work := abandonDispatch(t, store, session) + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + actors := &dispatchCleanupTestActors{lifecycleTestActors: base, afterSuspend: func(*ateapipb.Actor) { + cancel() + }} + writes := &observedDispatchCleanupStore{observedQuiescenceStore: &observedQuiescenceStore{lifecycleTestStore: store}} + + // When shutdown occurs as the successful runtime response arrives, assert + // that the known snapshot is recorded and admission opens for the next turn. + require.NoError(t, NewActorWorkflow(writes, actors, make(chan struct{}, 1), time.Minute).quiesceAbandonedDispatch(ctx, work)) + require.ErrorIs(t, ctx.Err(), context.Canceled) + require.EqualValues(t, 1, actors.suspendCalls.Load()) + require.EqualValues(t, 1, writes.finishedCleanup.Load()) + require.NoError(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "next-message")) +} + +func TestAbandonedDispatchRejectsMissingSnapshot(t *testing.T) { + // Given a suspension response without the durable snapshot it promised. + store, base, session := dispatchCleanupFixture(t) + work := abandonDispatch(t, store, session) + actors := &dispatchCleanupTestActors{lifecycleTestActors: base, afterSuspend: func(actor *ateapipb.Actor) { + actor.Status.Snapshots = nil + }} + writes := &observedDispatchCleanupStore{observedQuiescenceStore: &observedQuiescenceStore{lifecycleTestStore: store}} + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + + // When cleanup validates the result, assert that it retains admission and + // stops retrying an invalid runtime response. + require.NoError(t, NewActorWorkflow(writes, actors, make(chan struct{}, 1), time.Minute).quiesceAbandonedDispatch(ctx, work)) + require.EqualValues(t, 1, actors.suspendCalls.Load()) + require.Zero(t, writes.finishedCleanup.Load()) + require.ErrorIs(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "next-message"), database.ErrDispatchBusy) +} + func TestSavedWaitingInputStillWakesDispatchCleanup(t *testing.T) { // Given a continuation whose input was saved without starting execution. store, actors, session := dispatchCleanupFixture(t) @@ -422,27 +969,33 @@ func TestAbandonedDispatchCancellationBeforeClaimKeepsAdmissionOpen(t *testing.T } func TestAbandonedDispatchCannotSuspendNewDelivery(t *testing.T) { - store, base, session := dispatchCleanupFixture(t) - work := abandonDispatch(t, store, session) - dispatchID := uuid.New() - var once sync.Once - actors := &dispatchCleanupTestActors{lifecycleTestActors: base, afterRead: func() { - // A new delivery wins after the worker reads RUNNING but before it - // claims permission to suspend that actor. - once.Do(func() { - require.NoError(t, store.ReserveSessionDispatch(t.Context(), session.Id, dispatchID, "next-message")) - }) - }} - workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute) - require.NoError(t, workflow.quiesceAbandonedDispatch(t.Context(), work)) + for _, state := range []ateapipb.ActorState{ateapipb.ActorState_ACTOR_STATE_RUNNING, ateapipb.ActorState_ACTOR_STATE_RESUMING} { + t.Run("Given an actor in "+state.String()+", when a newer send wins before cleanup claims, assert that cleanup issues no runtime mutations", func(t *testing.T) { + store, base, session := dispatchCleanupFixture(t) + setDispatchCleanupActorState(base, session, state) + work := abandonDispatch(t, store, session) + dispatchID := uuid.New() + var once sync.Once + actors := &dispatchCleanupTestActors{lifecycleTestActors: base, afterRead: func() { + // A new delivery wins after the worker reads the actor but before it + // claims permission to resume or suspend that actor. + once.Do(func() { + require.NoError(t, store.ReserveSessionDispatch(t.Context(), session.Id, dispatchID, "next-message")) + }) + }} + workflow := NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute) + require.NoError(t, workflow.quiesceAbandonedDispatch(t.Context(), work)) - require.Zero(t, actors.suspendCalls.Load()) - message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("next")) - message.ID, message.ContextID = "next-message", session.ContextId - task := a2a.NewSubmittedTask(message, message) - hash := sha256.Sum256([]byte(task.ID)) - _, err := store.CreateRuntimeTask(t.Context(), session.Id, hash[:], task, dispatchID.String()) - require.NoError(t, err, "the winning delivery must still be able to accept its task") + require.Zero(t, actors.resumeCalls.Load()) + require.Zero(t, actors.suspendCalls.Load()) + message := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("next")) + message.ID, message.ContextID = "next-message", session.ContextId + task := a2a.NewSubmittedTask(message, message) + hash := sha256.Sum256([]byte(task.ID)) + _, err := store.CreateRuntimeTask(t.Context(), session.Id, hash[:], task, dispatchID.String()) + require.NoError(t, err, "the winning delivery must still be able to accept its task") + }) + } } func TestAbandonedDispatchRetriesKnownFailures(t *testing.T) { @@ -457,7 +1010,12 @@ func TestAbandonedDispatchRetriesKnownFailures(t *testing.T) { t.Run("Given "+scenario.name+", when cleanup runs, assert that one suspension completes", func(t *testing.T) { store, base, session := dispatchCleanupFixture(t) work := abandonDispatch(t, store, session) - actors := &dispatchCleanupTestActors{lifecycleTestActors: base, abortedCalls: scenario.abortedCalls} + actors := &dispatchCleanupTestActors{lifecycleTestActors: base, suspendAttempt: func(_ context.Context, call int32) error { + if call <= scenario.abortedCalls { + return status.Error(codes.Aborted, "actor lease is held") + } + return nil + }} writes := &observedDispatchCleanupStore{observedQuiescenceStore: &observedQuiescenceStore{lifecycleTestStore: store}, finishFailures: scenario.finishFailures} require.NoError(t, NewActorWorkflow(writes, actors, make(chan struct{}, 1), time.Minute).quiesceAbandonedDispatch(t.Context(), work)) @@ -470,18 +1028,57 @@ func TestAbandonedDispatchRetriesKnownFailures(t *testing.T) { } func TestAbandonedDispatchRetainsUncertainSuspension(t *testing.T) { - // Given a suspension whose result is lost after it changes the actor. - store, base, session := dispatchCleanupFixture(t) - work := abandonDispatch(t, store, session) - actors := &dispatchCleanupTestActors{lifecycleTestActors: base, mutationErr: status.Error(codes.Unavailable, "suspend response lost")} - require.NoError(t, NewActorWorkflow(store, actors, make(chan struct{}, 1), time.Minute).quiesceAbandonedDispatch(t.Context(), work)) + for _, phase := range []string{"resume", "suspend"} { + t.Run("Given an uncertain "+phase+" call, when the worker shuts down, assert that its claim stays fenced", func(t *testing.T) { + store, base, session := dispatchCleanupFixture(t) + setDispatchCleanupActorState(base, session, ateapipb.ActorState_ACTOR_STATE_RESUMING) + work := abandonDispatch(t, store, session) + actors := &dispatchCleanupTestActors{lifecycleTestActors: base} + entered := make(chan struct{}) + mutation := base.SuspendActor + if phase == "resume" { + mutation = base.ResumeActor + } + interrupted := func(ctx context.Context, _ int32) error { + if _, err := mutation(ctx, session.Agent.Ref.Namespace, substrate.ActorName(session.Id)); err != nil { + return err + } + close(entered) + <-ctx.Done() + return ctx.Err() + } + if phase == "resume" { + actors.resumeAttempt = interrupted + } else { + actors.suspendAttempt = interrupted + } + writes := &observedDispatchCleanupStore{observedQuiescenceStore: &observedQuiescenceStore{lifecycleTestStore: store}} + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + done := make(chan error, 1) + go func() { + done <- NewActorWorkflow(writes, actors, make(chan struct{}, 1), time.Minute).quiesceAbandonedDispatch(ctx, work) + }() + select { + case <-entered: + case <-time.After(5 * time.Second): + t.Fatal("cleanup did not reach the runtime mutation") + } + cancel() + select { + case err := <-done: + require.NoError(t, err) + case <-time.After(5 * time.Second): + t.Fatal("shutdown did not stop cleanup") + } - // Assert that neither a new delivery nor another cleanup worker can race - // the potentially issued call by taking over its permission. - require.EqualValues(t, 1, actors.suspendCalls.Load()) - require.ErrorIs(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "next-message"), database.ErrDispatchBusy) - _, err := store.GetPendingSessionDispatchCleanup(t.Context(), "") - require.ErrorIs(t, err, database.ErrNotFound) + require.EqualValues(t, 1, writes.cleanupClaims.Load()) + require.Zero(t, writes.finishedCleanup.Load()) + require.ErrorIs(t, store.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "next-message"), database.ErrDispatchBusy) + _, err := store.GetPendingSessionDispatchCleanup(t.Context(), "") + require.ErrorIs(t, err, database.ErrNotFound) + }) + } } func dispatchCleanupFixture(t *testing.T) (*lifecycleTestStore, *lifecycleTestActors, *apiv1.Session) { @@ -495,6 +1092,12 @@ func dispatchCleanupFixture(t *testing.T) (*lifecycleTestStore, *lifecycleTestAc return store, actors, session } +func setDispatchCleanupActorState(actors *lifecycleTestActors, session *apiv1.Session, state ateapipb.ActorState) { + actors.mu.Lock() + defer actors.mu.Unlock() + actors.actors[actorKey(session.Agent.Ref.Namespace, substrate.ActorName(session.Id))].Status.State = state +} + func abandonDispatch(t *testing.T, store *lifecycleTestStore, session *apiv1.Session) *database.SessionDispatchCleanup { t.Helper() dispatchID := uuid.New() @@ -509,6 +1112,7 @@ func abandonDispatch(t *testing.T, store *lifecycleTestStore, session *apiv1.Ses type observedDispatchCleanupStore struct { *observedQuiescenceStore + cleanupClaims atomic.Int32 finishedCleanup atomic.Int32 finishAttempts atomic.Int32 finishFailures int32 @@ -516,6 +1120,14 @@ type observedDispatchCleanupStore struct { var _ workflowStore = (*observedDispatchCleanupStore)(nil) +func (s *observedDispatchCleanupStore) ClaimSessionDispatchCleanup(ctx context.Context, candidate *database.SessionDispatchCleanup) (*database.SessionDispatchCleanup, error) { + work, err := s.Client.ClaimSessionDispatchCleanup(ctx, candidate) + if err == nil { + s.cleanupClaims.Add(1) + } + return work, err +} + func (s *observedDispatchCleanupStore) FinishSessionDispatchCleanup(ctx context.Context, work *database.SessionDispatchCleanup, snapshot *database.SessionTaskSnapshot) error { if s.finishAttempts.Add(1) <= s.finishFailures { return status.Error(codes.Unavailable, "database unavailable") @@ -527,14 +1139,33 @@ func (s *observedDispatchCleanupStore) FinishSessionDispatchCleanup(ctx context. return err } +type gatedDispatchDeletionActors struct { + *dispatchCleanupTestActors + deleteEntered chan struct{} + releaseDelete <-chan struct{} +} + +func (a *gatedDispatchDeletionActors) DeleteActor(ctx context.Context, space, name, uid string) error { + close(a.deleteEntered) + select { + case <-a.releaseDelete: + return a.lifecycleTestActors.DeleteActor(ctx, space, name, uid) + case <-ctx.Done(): + return ctx.Err() + } +} + type dispatchCleanupTestActors struct { *lifecycleTestActors - afterRead func() - readErr error - mutationErr error - abortedCalls int32 - readCalls atomic.Int32 - suspendCalls atomic.Int32 + afterRead func() + resumeAttempt func(context.Context, int32) error + suspendAttempt func(context.Context, int32) error + afterSuspend func(*ateapipb.Actor) + resumedActorUID string + readErr error + readCalls atomic.Int32 + resumeCalls atomic.Int32 + suspendCalls atomic.Int32 } var _ actorClient = (*dispatchCleanupTestActors)(nil) @@ -552,12 +1183,29 @@ func (a *dispatchCleanupTestActors) GetActor(ctx context.Context, space, name st } func (a *dispatchCleanupTestActors) SuspendActor(ctx context.Context, space, name string) (*ateapipb.Actor, error) { - if a.suspendCalls.Add(1) <= a.abortedCalls { - return nil, status.Error(codes.Aborted, "actor lease is held") + call := a.suspendCalls.Add(1) + if a.suspendAttempt != nil { + if err := a.suspendAttempt(ctx, call); err != nil { + return nil, err + } } actor, err := a.lifecycleTestActors.SuspendActor(ctx, space, name) - if a.mutationErr != nil { - return nil, a.mutationErr + if err == nil && a.afterSuspend != nil { + a.afterSuspend(actor) + } + return actor, err +} + +func (a *dispatchCleanupTestActors) ResumeActor(ctx context.Context, space, name string) (*ateapipb.Actor, error) { + call := a.resumeCalls.Add(1) + if a.resumeAttempt != nil { + if err := a.resumeAttempt(ctx, call); err != nil { + return nil, err + } + } + actor, err := a.lifecycleTestActors.ResumeActor(ctx, space, name) + if err == nil && a.resumedActorUID != "" { + actor.Metadata.Uid = a.resumedActorUID } return actor, err } diff --git a/go/core/internal/service/session/workflow.go b/go/core/internal/service/session/workflow.go index ccd7b18d46..ce0150a7c7 100644 --- a/go/core/internal/service/session/workflow.go +++ b/go/core/internal/service/session/workflow.go @@ -22,6 +22,7 @@ type workflowStore interface { FinishSessionQuiescence(context.Context, *database.SessionQuiescence, *database.SessionTaskSnapshot) error GetPendingSessionDispatchCleanup(context.Context, string) (*database.SessionDispatchCleanup, error) ClaimSessionDispatchCleanup(context.Context, *database.SessionDispatchCleanup) (*database.SessionDispatchCleanup, error) + ReleaseSessionDispatchCleanup(context.Context, *database.SessionDispatchCleanup) error FinishSessionDispatchCleanup(context.Context, *database.SessionDispatchCleanup, *database.SessionTaskSnapshot) error GetSessionForRuntime(context.Context, string, string) (*apiv1.Session, error) GetSessionCheckpointSnapshot(context.Context, string, string) (*database.SessionTaskSnapshot, string, error) @@ -82,6 +83,23 @@ func (w *ActorWorkflow) prepareQuiescence(ctx context.Context, session *apiv1.Se return revision, current, nil } +// resumeQuiescentActor completes an interrupted restore under the cleanup claim. +// It verifies the returned actor before the caller proceeds to suspension. +func (w *ActorWorkflow) resumeQuiescentActor(ctx context.Context, session *apiv1.Session, revision *database.RuntimeRevision, current *ateapipb.Actor) error { + atespace, name := revision.ActorTemplateAtespace, substrate.ActorName(session.GetId()) + actor, err := w.actors.ResumeActor(ctx, atespace, name) + if err != nil { + return fmt.Errorf("resume Actor %s/%s: %w", atespace, name, err) + } + if actor.GetStatus().GetState() != ateapipb.ActorState_ACTOR_STATE_RUNNING { + return fmt.Errorf("resume Actor %s/%s returned status %s", atespace, name, actor.GetStatus().GetState()) + } + if !validActorIdentity(actor, revision, name) || actor.GetMetadata().GetUid() != current.GetMetadata().GetUid() { + return fmt.Errorf("resume actor %s/%s returned invalid identity", atespace, name) + } + return nil +} + // suspendQuiescentActor suspends the verified actor and returns its durable volume // snapshot. The caller must prevent new turns until the snapshot is recorded. func (w *ActorWorkflow) suspendQuiescentActor(ctx context.Context, session *apiv1.Session, revision *database.RuntimeRevision, current *ateapipb.Actor) (*database.SessionTaskSnapshot, error) { @@ -126,8 +144,8 @@ func (w *ActorWorkflow) Suspend(ctx context.Context, session *apiv1.Session) (*a } // Delete closes admission before stopping and deleting compute. It can supersede -// unissued work and any crash recovery, but never deletes a session while another -// issued call is uncertain. +// unissued work and crash recovery. Force also supersedes dispatch cleanup: +// deletion prevents new execution even if an old resume or suspend is uncertain. func (w *ActorWorkflow) Delete(ctx context.Context, session *apiv1.Session, force bool) (*apiv1.Session, error) { return w.run(ctx, session.GetId(), apiv1.RuntimeOperation_RUNTIME_OPERATION_DELETE, force) } diff --git a/go/core/test/e2e/README.md b/go/core/test/e2e/README.md index 72de1942e1..744cd1f135 100644 --- a/go/core/test/e2e/README.md +++ b/go/core/test/e2e/README.md @@ -57,8 +57,9 @@ parking budget to 5 seconds and request timeout to 10 seconds. For local runs, include this file in `install.substrate.valuesFiles` when installing the test cluster. These overrides keep the test independent of the chart's defaults. The test requires the actor to remain `RESUMING` before ordinary Session deletion, -then verifies Session and Actor absence through kagent APIs. Backend cleanup -retries are allowed only after the actor reaches `DELETING` or is absent. +then verifies Session and Actor absence through kagent APIs. If automatic +dispatch cleanup claims the actor first, deletion waits for that cleanup to +finish. Backend deletion retries require `DELETING` or an absent actor. `TestSessionSuspendsAfterCanceledDispatch` clones the `byo-e2e` Harness with the test runtime's `--block-first-restored-send` flag. Its first restored activation @@ -70,6 +71,12 @@ successful invocation also verifies that cleanup saved the runtime state. Build the BYO fixture with `make build-byo-a2a` and install its digest in `byo-e2e` before running this test; the gate is disabled for other tests. +`TestSessionSuspendsAfterCanceledResume` clones `byo-e2e` with a startup wrapper +that delays only the first restored activation. It cancels while the actor is +`RESUMING` with a worker assignment, then requires automatic resume and suspension +with no tasks and no remaining assignment. A subsequent message must complete +normally. The installed Harness and ordinary runtime behavior are unchanged. + The suite exercises the public API against a clean Kind installation. It does not reconcile Kubernetes resources itself: installation creates the Harness fixtures, and each test owns the templates and API resources it creates. diff --git a/go/core/test/e2e/session_resuming_deletion_test.go b/go/core/test/e2e/session_resuming_deletion_test.go index 939157cc63..6ea7da4dae 100644 --- a/go/core/test/e2e/session_resuming_deletion_test.go +++ b/go/core/test/e2e/session_resuming_deletion_test.go @@ -105,6 +105,123 @@ func TestSessionSuspendsAfterCanceledDispatch(t *testing.T) { assertActorSuspended(t, fixture) } +func TestSessionSuspendsAfterCanceledResume(t *testing.T) { + t.Parallel() + // Given a golden runtime that starts normally, but delays readiness for ten + // seconds on its first restore. Cleanup's snapshot preserves the second marker + // so the next invocation starts without the delay. + const delayedRestore = `resume_marker=/data/e2e-canceled-resume +if [ -f "$resume_marker" ] && [ ! -f "$resume_marker.released" ]; then + touch "$resume_marker.released" + sleep 10 +fi +touch "$resume_marker" +exec /app` + target := interactionTarget(t) + kube := interactionKubeClient(t) + base := &v1alpha3.Harness{} + require.NoError(t, kube.Get(t.Context(), ctrlclient.ObjectKey{Namespace: "kagent", Name: "byo-e2e"}, base)) + harness := &v1alpha3.Harness{ + ObjectMeta: metav1.ObjectMeta{GenerateName: "canceled-resume-", Namespace: "kagent"}, + Spec: *base.Spec.DeepCopy(), + } + harness.Spec.Workload.Command = []string{"/bin/sh", "-c"} + harness.Spec.Workload.Args = []string{delayedRestore} + require.NoError(t, kube.Create(t.Context(), harness)) + t.Cleanup(func() { + if err := kube.Delete(context.Background(), harness); err != nil && !apierrors.IsNotFound(err) { + t.Errorf("delete canceled-resume Harness: %v", err) + } + }) + template := &v1alpha3.AgentTemplate{ + ObjectMeta: metav1.ObjectMeta{GenerateName: "canceled-resume-", Namespace: "kagent"}, + Spec: v1alpha3.AgentTemplateSpec{Description: "Canceled resume lifecycle fixture"}, + } + createAndWaitInteractionTemplateForHarness(t, kube, template, harness.Name) + fixture := newInteractionFixtureForHarnessTemplate(t, target, harness.Name, template.Name) + actorName := substrate.ActorName(fixture.sessionID) + initial, err := findSubstrateActor(fixture.ctx, fixture.system, "kagent", actorName) + require.NoError(t, err) + require.NotNil(t, initial) + require.Equal(t, ateapipb.ActorState_ACTOR_STATE_SUSPENDED, initial.GetStatus().GetState()) + + // When the caller cancels during restore, no task can have reached the + // runtime. A read of actor inventory must not itself finish that restore. + _, request := newMessageRequest(t, "hello") + request.Tenant, request.Message.ContextId = fixture.tenant, fixture.sessionID + sendCtx, cancel := context.WithCancel(fixture.ctx) + defer cancel() + sent := make(chan error, 1) + go func() { + _, err := fixture.client.SendMessage(sendCtx, request) + sent <- err + }() + require.NoError(t, wait.PollUntilContextTimeout(fixture.ctx, 100*time.Millisecond, 30*time.Second, true, func(ctx context.Context) (bool, error) { + select { + case err := <-sent: + return false, fmt.Errorf("send ended before cancellation: %v", err) + default: + } + actor, err := findSubstrateActor(ctx, fixture.system, "kagent", actorName) + if err != nil { + return false, err + } + return actor != nil && actor.GetStatus().GetState() == ateapipb.ActorState_ACTOR_STATE_RESUMING && actor.GetStatus().GetWorkerAssignment() != nil, nil + })) + cancel() + select { + case err := <-sent: + require.Equal(t, codes.Canceled, status.Code(err)) + case <-time.After(5 * time.Second): + t.Fatal("send did not return after cancellation") + } + interrupted, err := findSubstrateActor(fixture.ctx, fixture.system, "kagent", actorName) + require.NoError(t, err) + require.NotNil(t, interrupted) + listRequest, err := pbconv.ToProtoListTasksRequest(&a2atype.ListTasksRequest{Tenant: fixture.tenant, ContextID: fixture.contextID}) + require.NoError(t, err) + listed, err := fixture.client.ListTasks(fixture.ctx, listRequest) + require.NoError(t, err) + t.Logf("canceled Session %s actor kagent/%s: state=%s worker_assigned=%t tasks=%d", fixture.sessionID, actorName, + interrupted.GetStatus().GetState(), interrupted.GetStatus().GetWorkerAssignment() != nil, len(listed.GetTasks())) + require.Empty(t, listed.GetTasks()) + + // Assert that automatic cleanup finishes the interrupted restore and suspends + // it without accepting input or changing the Session's logical availability. + lastActor := interrupted + err = wait.PollUntilContextTimeout(fixture.ctx, 100*time.Millisecond, 30*time.Second, true, func(ctx context.Context) (bool, error) { + actor, err := findSubstrateActor(ctx, fixture.system, "kagent", actorName) + if err != nil { + return false, err + } + if actor == nil { + return false, fmt.Errorf("actor disappeared during canceled-resume cleanup") + } + if actor.GetStatus().GetState() != lastActor.GetStatus().GetState() { + t.Logf("Session %s actor kagent/%s: state=%s worker_assigned=%t", fixture.sessionID, actorName, + actor.GetStatus().GetState(), actor.GetStatus().GetWorkerAssignment() != nil) + } + lastActor = actor + snapshot := substrate.LatestActorSnapshot(actor) + return actor.GetStatus().GetState() == ateapipb.ActorState_ACTOR_STATE_SUSPENDED && actor.GetStatus().GetWorkerAssignment() == nil && + substrate.SnapshotURI(snapshot) != "" && snapshot.GetFidelity() == ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_VOLUMES, nil + }) + t.Logf("cleanup result for Session %s actor kagent/%s: state=%s worker_assigned=%t error=%v", fixture.sessionID, actorName, + lastActor.GetStatus().GetState(), lastActor.GetStatus().GetWorkerAssignment() != nil, err) + listed, listErr := fixture.client.ListTasks(fixture.ctx, listRequest) + require.NoError(t, listErr) + t.Logf("Session %s tasks after canceled-resume cleanup: %d", fixture.sessionID, len(listed.GetTasks())) + require.NoError(t, err, "canceled restore must automatically reach Suspended with a durable snapshot") + require.Empty(t, listed.GetTasks()) + current, err := fixture.sessions.GetSession(fixture.ctx, &apiv1.GetSessionRequest{SessionId: fixture.sessionID}) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, current.GetSession().GetState()) + _, _, task := fixture.send(t, "hello") + require.Equal(t, a2atype.TaskStateCompleted, task.Status.State) + require.Contains(t, taskText(task), "BYO agent response") + assertActorSuspended(t, fixture) +} + // The Go ADK startup delay reproduces the router timeout independently of the // model protocol. The installed Harness is left intact. func TestSessionDeletionAfterInterruptedResume(t *testing.T) { @@ -163,6 +280,11 @@ func TestSessionDeletionAfterInterruptedResume(t *testing.T) { require.NoError(t, wait.PollUntilContextTimeout(fixture.ctx, 100*time.Millisecond, 30*time.Second, true, func(ctx context.Context) (bool, error) { var err error deleted, err = fixture.sessions.DeleteSession(ctx, &apiv1.DeleteSessionRequest{SessionId: fixture.sessionID}) + if status.Code(err) == codes.FailedPrecondition { + // Automatic canceled-dispatch cleanup may own admission until its + // interrupted restore and suspension finish. Deletion must then win. + return false, nil + } if status.Code(err) != codes.Unavailable { return err == nil, err }