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) {