diff --git a/src/mgr/postgres-store.ts b/src/mgr/postgres-store.ts index 9c6cb70..41dd02a 100644 --- a/src/mgr/postgres-store.ts +++ b/src/mgr/postgres-store.ts @@ -6,7 +6,7 @@ import { redactJson } from "../common/redaction.js"; import type { BackendProfile, BackendTurnResult, CancelRequestRecord, CancelStage, CancelTargetKind, CommandRecord, CommandState, CreateCommandInput, CreateQueueTaskInput, CreateRunInput, EventType, FailureKind, JsonRecord, JsonValue, KafkaEventOutboxRecord, ListGcExpiredSessionsInput, QueueAttemptListResult, QueueAttemptRecord, QueueAttemptRef, QueueCommanderSnapshot, QueueReadCursorRecord, QueueRetryActivation, QueueRetryReservation, QueueStats, QueueTaskListResult, QueueTaskRecord, QueueTaskState, RunEvent, RunnerDispatchCompletion, RunnerDispatchIntentRecord, RunnerJobRecord, RunnerRecord, RunRecord, RunStatus, SessionEventPage, SessionListResult, SessionReadCursorRecord, SessionRecord, SessionRef, SessionStoragePatch, SessionSummary, TerminalStatus, UpsertSessionInput } from "../common/types.js"; import { newId, nowIso, stableHash } from "../common/validation.js"; import type { ActivateQueueRetryAttemptInput, AgentRunStore, CreateRunIdentity, DurableQueueClaim, EventOutboxConfig, ListQueueAttemptsInput, ListQueueTasksInput, ListSessionsInput, RecordQueueRetryAttemptFailureInput, ReserveQueueRetryAttemptInput, RunnerRetentionFenceFacts, RunnerRetentionFenceInput, SaveRunnerJobInput, SessionEventPageInput, SessionSteerAdmission, SessionSteerAdmissionInput, SessionTurnAdmission, SessionTurnAdmissionInput, StoreHealth, UpdateQueueTaskAttemptInput } from "./store.js"; -import { activeSessionAdmissionConflict, assertQueueAttemptActivationReplay, assertQueueRetryable, assertQueueTaskPayloadHash, assertRunCreateReplay, assertSessionBoundary, assertSessionCommandReplay, assertSessionSteerAdmissionContract, assertSessionTurnAdmissionContract, buildQueueStats, buildQueueTaskSummary, buildSessionSummary, cancelStagePayload, clampQueueLimit, clampSessionLimit, commandStateFromTerminal, fenceLateEventForCancelledRun, invalidSessionAdmissionProjection, isLeaseExpired, isSessionOutputEvent, isTerminalCommandState, isTerminalQueueAttemptState, isTerminalQueueTaskState, isTerminalRunStatus, latestQueueAttempt, lateWriteRejectedPayload, nextQueueRetryIndex, parseQueueCursor, parseSessionCursor, queueAttemptRef, queueRetryReservation, queueRetryStateError, queueTaskMatchesCommander, queueTaskPayloadHash, queueTaskSort, receivableActiveSessionTurn, runCreatedEventPayload, runnerDispatchFailureKind, runnerDispatchFailureReason, sessionListFilters, sessionMatchesListState, sessionRefFromRecord, sessionSort, sessionTitleFromCommand, statusFromTerminal, summarizeSessionRef, titleFromMetadata } from "./store.js"; +import { activeSessionAdmissionConflict, assertQueueAttemptActivationReplay, assertQueueRetryable, assertQueueTaskPayloadHash, assertRunCreateReplay, assertSessionBoundary, assertSessionCommandReplay, assertSessionSteerAdmissionContract, assertSessionTurnAdmissionContract, buildQueueStats, buildQueueTaskSummary, buildSessionSummary, cancelStagePayload, clampQueueLimit, clampSessionLimit, commandStateFromTerminal, fenceLateEventForCancelledRun, invalidSessionAdmissionProjection, isLeaseExpired, isSessionOutputEvent, isTerminalCommandState, isTerminalQueueAttemptState, isTerminalQueueTaskState, isTerminalRunStatus, latestQueueAttempt, lateWriteRejectedPayload, nextQueueRetryIndex, parseQueueCursor, parseSessionCursor, queueAttemptRef, queueRetryReservation, queueRetryStateError, queueTaskMatchesCommander, queueTaskPayloadHash, queueTaskSort, receivableActiveSessionTurn, runCreatedEventPayload, runnerDispatchFailureKind, runnerDispatchFailureReason, sessionListFilters, sessionMatchesListState, sessionRefFromRecord, sessionSort, sessionTitleFromCommand, staleActiveSessionTurn, statusFromTerminal, summarizeSessionRef, titleFromMetadata } from "./store.js"; import { backendCapabilitiesSqlValues, mergeBackendCapability } from "../common/backend-profiles.js"; import { normalizeRunEventPayload, requireEventType, requireExternallyAppendableEventType, userMessagePayloadForCommand } from "../common/events.js"; import { buildKafkaEventOutboxRecord, canonicalAgentRunEventPartitionKey } from "./event-outbox.js"; @@ -679,8 +679,12 @@ CREATE TABLE IF NOT EXISTS agentrun_schema_migrations ( try { return await this.replayOrRecoverSessionTurnWithClient(client, run, command, input, "replayed"); } catch (error) { - if (error instanceof AgentRunError && error.failureKind === "schema-invalid") throw activeSessionAdmissionConflict(input.sessionId, run, command); - throw error; + if (error instanceof AgentRunError && error.failureKind === "schema-invalid") { + if (!staleActiveSessionTurn(session, run, command)) throw activeSessionAdmissionConflict(input.sessionId, run, command); + await this.terminalizeStaleSessionAdmissionWithClient(client, run, command); + } else { + throw error; + } } } } @@ -1968,6 +1972,59 @@ CREATE TABLE IF NOT EXISTS agentrun_schema_migrations ( return { run, command: attachRunnerDispatchIntent(command, recovered), disposition: "recovered", recoveredPriorPartialWrite: true, partialWrite: false }; } + private async terminalizeStaleSessionAdmissionWithClient(client: PoolClient, run: RunRecord, command: CommandRecord): Promise { + const at = nowIso(); + const reason = "session-turn-admission-stale-lease"; + const message = "session turn admission replaced a stale runner lease"; + const updatedCommands = await client.query( + "UPDATE agentrun_commands SET state = 'failed', updated_at = $2 WHERE run_id = $1 AND state NOT IN ('completed', 'failed', 'cancelled') RETURNING *", + [run.id, at], + ); + await fenceActiveRunnerDispatchIntents(client, { runId: run.id }, message, at); + await client.query( + `UPDATE agentrun_runs + SET status = 'failed', terminal_status = 'failed', failure_kind = 'infra-failed', failure_message = $2, + claimed_by = NULL, lease_expires_at = NULL, updated_at = $3 + WHERE id = $1`, + [run.id, message, at], + ); + await client.query("DELETE FROM agentrun_leases WHERE run_id = $1", [run.id]); + await this.appendEventWithLockedRun(client, run.id, "backend_status", { + phase: "session-turn-admission-stale-terminalized", + reason, + commandId: command.id, + claimedBy: run.claimedBy, + leaseExpiresAt: run.leaseExpiresAt, + failureKind: "infra-failed", + valuesPrinted: false, + }); + for (const row of updatedCommands.rows) { + const item = commandFromRow(row); + await this.appendEventWithLockedRun(client, run.id, "backend_status", { + phase: "command-terminal", + reason, + commandId: item.id, + state: "failed", + terminalStatus: "failed", + failureKind: "infra-failed", + message, + threadId: null, + turnId: null, + }); + } + await this.appendEventWithLockedRun(client, run.id, "terminal_status", { terminalStatus: "failed", reason, failureKind: "infra-failed", message }); + await this.touchSessionForRun(client, run, { + executionState: "terminal", + activeRunId: null, + activeCommandId: null, + lastRunId: run.id, + lastCommandId: command.id, + terminalStatus: "failed", + failureKind: "infra-failed", + lastActivityAt: at, + }, { bumpVersion: true, at }); + } + private async runnerDispatchIntentForCommandWithClient(client: PoolClient, commandId: string): Promise { const result = await client.query("SELECT * FROM agentrun_runner_dispatch_intents WHERE command_id = $1", [commandId]); return result.rows[0] ? runnerDispatchIntentFromRow(result.rows[0]) : null; diff --git a/src/mgr/store.ts b/src/mgr/store.ts index af4c827..596a2b8 100644 --- a/src/mgr/store.ts +++ b/src/mgr/store.ts @@ -374,6 +374,7 @@ export class MemoryAgentRunStore implements AgentRunStore { : null; if (replay?.run) return this.replayOrRecoverSessionTurn(replay.run, replay.command, input, "replayed"); + let staleAdmission: { run: RunRecord; command: CommandRecord } | null = null; if (session?.activeRunId || session?.activeCommandId) { if (!session.activeRunId || !session.activeCommandId) throw invalidSessionAdmissionProjection(input.sessionId); const run = this.runs.get(session.activeRunId); @@ -383,30 +384,36 @@ export class MemoryAgentRunStore implements AgentRunStore { try { return this.replayOrRecoverSessionTurn(run, command, input, "replayed"); } catch (error) { - if (error instanceof AgentRunError && error.failureKind === "schema-invalid") throw activeSessionAdmissionConflict(input.sessionId, run, command); - throw error; + if (error instanceof AgentRunError && error.failureKind === "schema-invalid") { + if (!staleActiveSessionTurn(session, run, command)) throw activeSessionAdmissionConflict(input.sessionId, run, command); + staleAdmission = { run, command }; + } else { + throw error; + } } } } - const sessionBefore = session ? { ...session } : null; + const runsBefore = new Map(this.runs); + const commandsBefore = new Map(this.commands); + const eventsBefore = new Map(Array.from(this.eventsByRun, ([runId, events]) => [runId, [...events]])); + const sessionsBefore = new Map(this.sessions); + const dispatchIntentsBefore = new Map(this.runnerDispatchIntents); + const outboxBefore = new Map(this.kafkaEventOutbox); const sessionVersionBefore = this.sessionVersion; const kafkaOutboxSeqBefore = this.kafkaOutboxSeq; - let run: RunRecord | null = null; try { - run = this.createRun(input.run); + if (staleAdmission) this.terminalizeStaleSessionAdmission(staleAdmission.run, staleAdmission.command); + const run = this.createRun(input.run); const command = this.createCommand(run.id, input.command); return { run, command, disposition: "created", recoveredPriorPartialWrite: false, partialWrite: false }; } catch (error) { - if (run) { - this.runs.delete(run.id); - this.eventsByRun.delete(run.id); - for (const [commandId, command] of this.commands) if (command.runId === run.id) this.commands.delete(commandId); - for (const [intentId, intent] of this.runnerDispatchIntents) if (intent.runId === run.id) this.runnerDispatchIntents.delete(intentId); - for (const [outboxId, item] of this.kafkaEventOutbox) if (item.runId === run.id) this.kafkaEventOutbox.delete(outboxId); - } - if (sessionBefore) this.sessions.set(input.sessionId, sessionBefore); - else this.sessions.delete(input.sessionId); + restoreMap(this.runs, runsBefore); + restoreMap(this.commands, commandsBefore); + restoreMap(this.eventsByRun, eventsBefore); + restoreMap(this.sessions, sessionsBefore); + restoreMap(this.runnerDispatchIntents, dispatchIntentsBefore); + restoreMap(this.kafkaEventOutbox, outboxBefore); this.sessionVersion = sessionVersionBefore; this.kafkaOutboxSeq = kafkaOutboxSeqBefore; throw error; @@ -440,6 +447,59 @@ export class MemoryAgentRunStore implements AgentRunStore { return { run, command: attachRunnerDispatchIntent(command, recovered), disposition: "recovered", recoveredPriorPartialWrite: true, partialWrite: false }; } + private terminalizeStaleSessionAdmission(run: RunRecord, command: CommandRecord): void { + const at = nowIso(); + const reason = "session-turn-admission-stale-lease"; + const message = "session turn admission replaced a stale runner lease"; + const updatedCommands = Array.from(this.commands.values()).filter((item) => item.runId === run.id && !isTerminalCommandState(item.state)); + for (const item of updatedCommands) this.commands.set(item.id, { ...item, state: "failed", updatedAt: at }); + this.cancelRunnerDispatchIntentsForRun(run.id, message, at); + const terminalRun: RunRecord = { + ...run, + status: "failed", + terminalStatus: "failed", + failureKind: "infra-failed", + failureMessage: message, + claimedBy: null, + leaseExpiresAt: null, + updatedAt: at, + }; + this.runs.set(run.id, terminalRun); + this.appendEvent(run.id, "backend_status", { + phase: "session-turn-admission-stale-terminalized", + reason, + commandId: command.id, + claimedBy: run.claimedBy, + leaseExpiresAt: run.leaseExpiresAt, + failureKind: "infra-failed", + valuesPrinted: false, + }); + for (const item of updatedCommands) { + this.appendEvent(run.id, "backend_status", { + phase: "command-terminal", + reason, + commandId: item.id, + state: "failed", + terminalStatus: "failed", + failureKind: "infra-failed", + message, + threadId: null, + turnId: null, + }); + } + this.appendEvent(run.id, "terminal_status", { terminalStatus: "failed", reason, failureKind: "infra-failed", message }); + this.touchSessionForRun(terminalRun, { + executionState: "terminal", + activeRunId: null, + activeCommandId: null, + lastRunId: run.id, + lastCommandId: command.id, + terminalStatus: "failed", + failureKind: "infra-failed", + lastActivityAt: at, + }, { bumpVersion: true, at }); + } + getRunnerDispatchIntent(commandId: string): RunnerDispatchIntentRecord | null { return Array.from(this.runnerDispatchIntents.values()).find((intent) => intent.commandId === commandId) ?? null; } @@ -1702,6 +1762,17 @@ export function receivableActiveSessionTurn(session: SessionRecord, run: RunReco return { reason: "active-turn-running", leaseExpired: false }; } +export function staleActiveSessionTurn(session: SessionRecord, run: RunRecord, command: CommandRecord, now = Date.now()): { reason: "session-turn-admission-stale-lease"; leaseExpired: true } | null { + if (session.activeRunId !== run.id || session.activeCommandId !== command.id) return null; + if (run.sessionRef?.sessionId !== session.sessionId || command.runId !== run.id) return null; + if (isTerminalRunStatus(run.status) || isTerminalCommandState(command.state)) return null; + if (command.type !== "turn" || !run.claimedBy) return null; + if (run.status !== "claimed" && run.status !== "running") return null; + const leaseExpiresAt = Date.parse(run.leaseExpiresAt ?? ""); + if (Number.isFinite(leaseExpiresAt) && leaseExpiresAt > now) return null; + return { reason: "session-turn-admission-stale-lease", leaseExpired: true }; +} + export function assertSessionCommandReplay(existing: CommandRecord, input: CreateCommandInput): void { if (existing.type !== input.type || existing.payloadHash !== stableHash(input.payload) || (existing.idempotencyKey ?? null) !== (input.idempotencyKey ?? null)) { throw new AgentRunError("schema-invalid", "session turn admission replay does not match the existing command", { @@ -1927,6 +1998,11 @@ function isDueKafkaOutboxItem(item: KafkaEventOutboxRecord | undefined, now: num return Boolean(item && Date.parse(item.nextAttemptAt) <= now && (!item.leaseExpiresAt || Date.parse(item.leaseExpiresAt) <= now)); } +function restoreMap(target: Map, snapshot: Map): void { + target.clear(); + for (const [key, value] of snapshot) target.set(key, value); +} + export function isSessionOutputEvent(event: RunEvent): boolean { return event.type === "assistant_progress" || event.type === "assistant_message" || event.type === "command_output" || event.type === "diff" || event.type === "error" || event.type === "terminal_status"; } diff --git a/src/selftest/cases/66-session-turn-admission.ts b/src/selftest/cases/66-session-turn-admission.ts index 90daa50..8ca4b6c 100644 --- a/src/selftest/cases/66-session-turn-admission.ts +++ b/src/selftest/cases/66-session-turn-admission.ts @@ -21,6 +21,7 @@ const selfTest: SelfTestCase = async (context) => { await assertHttpAdmissionDoesNotCallKubectl(fakeKubectl); await assertDisabledDispatcherFailsBeforeFormalFacts(fakeKubectl); assertMemoryAdmissionRollback(); + await assertMemoryStaleAdmissionRollback(); assertLegacyHalfCommitRecovery(); await assertRestartVersionForwardRecovery(fakeKubectl); assertMemoryReplayAndConflict(); @@ -29,6 +30,7 @@ const selfTest: SelfTestCase = async (context) => { await assertSuccessfulSameSessionContinuity(fakeKubectl); await assertSessionFollowUpInheritsAipodResources(fakeKubectl); await assertActiveSessionSendSteersAtomically(fakeKubectl); + await assertStaleSessionAdmissionCreatesFreshTurn(fakeKubectl); await assertSessionFollowUpTerminalRaceCreatesNewRunner(fakeKubectl); } finally { restoreEnv("AGENTRUN_SELFTEST_KUBECTL_MODE", previousMode); @@ -39,6 +41,7 @@ const selfTest: SelfTestCase = async (context) => { "session-send-durable-admission-no-sync-kubectl", "session-send-dispatcher-disabled-prewrite-failure", "memory-session-turn-admission-rollback", + "memory-stale-session-turn-admission-rollback", "legacy-half-commit-recovered-in-place", "restart-version-forward-recovery-keeps-frozen-run-contract", "session-turn-idempotent-replay-and-active-conflict", @@ -47,6 +50,8 @@ const selfTest: SelfTestCase = async (context) => { "successful-same-session-continuity", "session-follow-up-inherits-aipod-workspace-bundles-and-tool-credentials", "active-session-send-steers-under-session-admission-lock", + "stale-claimed-session-admission-creates-fresh-turn", + "runner-lost-session-send-creates-fresh-runner-admission", "session-follow-up-terminal-race-creates-new-runner-without-empty-old-command", ], }; @@ -141,6 +146,29 @@ function assertMemoryAdmissionRollback(): void { assert.throws(() => store.getRun(store.createdRunId ?? ""), /was not found/u); } +async function assertMemoryStaleAdmissionRollback(): Promise { + const store = new FailingStaleReplacementStore({ eventOutbox: { enabled: true, topic: "agentrun.event.v1", source: "session-admission-selftest" } }); + const sessionId = "ses_session_admission_stale_rollback"; + const staleRun = store.createRun(runInput(sessionId)); + const staleCommand = store.createCommand(staleRun.id, validateCreateCommand({ type: "turn", payload: { prompt: "stale admission before memory rollback" } })); + store.claimRun(staleRun.id, "runner_memory_stale_rollback", 1); + store.ackCommand(staleCommand.id); + await delay(5); + const beforeEvents = store.listEvents(staleRun.id, 0, 100); + const beforeOutbox = store.kafkaEventOutboxStatus(); + assert.throws( + () => store.admitSessionTurn(admissionInput(sessionId, "reject stale replacement after write", "session-admission-stale-rollback")), + /simulated stale replacement failure/u, + ); + assert.equal(store.getRun(staleRun.id).status, "claimed"); + assert.equal(store.getRun(staleRun.id).claimedBy, "runner_memory_stale_rollback"); + assert.equal(store.getCommand(staleCommand.id).state, "acknowledged"); + assert.equal(store.getSession(sessionId)?.activeRunId, staleRun.id); + assert.equal(store.getSession(sessionId)?.activeCommandId, staleCommand.id); + assert.deepEqual(store.listEvents(staleRun.id, 0, 100), beforeEvents); + assert.deepEqual(store.kafkaEventOutboxStatus(), beforeOutbox); +} + function assertLegacyHalfCommitRecovery(): void { const store = new MemoryAgentRunStore(); const sessionId = "ses_session_admission_legacy"; @@ -400,6 +428,63 @@ async function assertActiveSessionSendSteersAtomically(kubectlCommand: string): } } +async function assertStaleSessionAdmissionCreatesFreshTurn(kubectlCommand: string): Promise { + const pendingStore = new MemoryAgentRunStore(); + const pendingSessionId = "sess_artificer_stale_claimed"; + const stalePending = pendingStore.createRun(artificerRunInput(pendingSessionId)); + const stalePendingCommand = pendingStore.createCommand(stalePending.id, validateCreateCommand({ type: "turn", payload: { prompt: "runner disappeared before ack" } })); + pendingStore.claimRun(stalePending.id, "runner_artificer_stale_claimed", 1); + await delay(5); + const pendingReplacement = pendingStore.admitSessionTurn(admissionInput(pendingSessionId, "replace stale claimed admission", "session-stale-claimed-replacement")); + assert.notEqual(pendingReplacement.run.id, stalePending.id); + assert.equal(pendingStore.getRun(stalePending.id).status, "failed"); + assert.equal(pendingStore.getRun(stalePending.id).claimedBy, null); + assert.equal(pendingStore.getCommand(stalePendingCommand.id).state, "failed"); + assert.equal(pendingStore.getSession(pendingSessionId)?.activeRunId, pendingReplacement.run.id); + + const store = new MemoryAgentRunStore(); + const sessionId = "sess_artificer_runner_lost"; + const staleRun = store.createRun(artificerRunInput(sessionId)); + const staleCommand = store.createCommand(staleRun.id, validateCreateCommand({ type: "turn", payload: { prompt: "runner disappeared after ack" } })); + store.claimRun(staleRun.id, "runner_artificer_lost", 1); + store.ackCommand(staleCommand.id); + await delay(5); + const server = await startManagerServer({ + store, + sourceCommit: "selftest", + runnerJobDefaults: defaults(kubectlCommand), + runnerDispatcherOptions: { ...dispatcherOptions(3), intervalMs: 60_000 }, + kafkaOutboxRelayOptions: { enabled: false }, + }); + try { + const client = new ManagerClient(server.baseUrl); + const response = await client.post(`/api/v1/sessions/${sessionId}/send`, sessionSendBody(sessionId, "continue after runner loss", "session-runner-lost-replacement")) as JsonRecord; + const nextRun = response.run as RunRecord; + const nextCommand = response.command as CommandRecord; + const admission = response.runnerAdmission as JsonRecord; + assert.equal(response.decision, "turn"); + assert.equal(response.mutation, true); + assert.equal(response.partialWrite, false); + assert.equal(response.activeBefore, null); + assert.notEqual(nextRun.id, staleRun.id); + assert.notEqual(nextCommand.id, staleCommand.id); + assert.equal(admission.disposition, "created"); + assert.equal(admission.state, "pending"); + assert.equal(admission.mutation, true); + assert.equal(admission.partialWrite, false); + assert.equal(store.getRun(staleRun.id).status, "failed"); + assert.equal(store.getRun(staleRun.id).failureKind, "infra-failed"); + assert.equal(store.getRun(staleRun.id).claimedBy, null); + assert.equal(store.getRun(staleRun.id).leaseExpiresAt, null); + assert.equal(store.getCommand(staleCommand.id).state, "failed"); + assert.equal(store.listEvents(staleRun.id, 0, 100).some((event) => event.payload.reason === "session-turn-admission-stale-lease"), true); + assert.equal(store.getSession(sessionId)?.activeRunId, nextRun.id); + assert.equal(store.getSession(sessionId)?.activeCommandId, nextCommand.id); + } finally { + await closeServer(server.server); + } +} + function assertTerminalizedAdmission(store: MemoryAgentRunStore, runId: string, commandId: string, sessionId: string, reason: string): void { const intent = store.getRunnerDispatchIntent(commandId); assert.equal(intent?.state, "failed"); @@ -551,6 +636,14 @@ class FailingSessionCommandStore extends MemoryAgentRunStore { } } +class FailingStaleReplacementStore extends MemoryAgentRunStore { + override createCommand(runId: string, input: CreateCommandInput): CommandRecord { + const command = super.createCommand(runId, input); + if (input.payload.prompt === "reject stale replacement after write") throw new Error("simulated stale replacement failure"); + return command; + } +} + class TerminalizingSteerStore extends MemoryAgentRunStore { override admitSessionSteer(input: SessionSteerAdmissionInput): SessionSteerAdmission | null { const session = this.getSession(input.sessionId); diff --git a/src/selftest/integration/postgres-session-turn-admission.ts b/src/selftest/integration/postgres-session-turn-admission.ts index c72ee2e..9922f33 100644 --- a/src/selftest/integration/postgres-session-turn-admission.ts +++ b/src/selftest/integration/postgres-session-turn-admission.ts @@ -12,8 +12,12 @@ const suffix = `${Date.now()}_${process.pid}`; const concurrentSessionId = `ses_pg_admission_${suffix}`; const legacySessionId = `ses_pg_admission_legacy_${suffix}`; const rollbackSessionId = `ses_pg_admission_rollback_${suffix}`; +const staleRollbackSessionId = `ses_pg_admission_stale_rollback_${suffix}`; const resourceSessionId = `ses_pg_admission_resource_${suffix}`; const terminalRaceSessionId = `ses_pg_admission_terminal_race_${suffix}`; +const staleClaimedSessionId = `ses_pg_admission_stale_claimed_${suffix}`; +const runnerLostSessionId = `ses_pg_admission_runner_lost_${suffix}`; +const freshActiveSessionId = `ses_pg_admission_fresh_active_${suffix}`; const eventOutbox = { enabled: true, topic: "agentrun.event.v1", source: "session-admission-selftest" } as const; const primary = await createPostgresAgentRunStore({ connectionString, poolMax: 4, eventOutbox }); const concurrent = await createPostgresAgentRunStore({ connectionString, poolMax: 4, eventOutbox }); @@ -22,6 +26,9 @@ const reopenedStores: Array undefined); await cleanupSession(legacySessionId).catch(() => undefined); await cleanupSession(rollbackSessionId).catch(() => undefined); + await cleanupSession(staleRollbackSessionId).catch(() => undefined); await cleanupSession(resourceSessionId).catch(() => undefined); await cleanupSession(terminalRaceSessionId).catch(() => undefined); + await cleanupSession(staleClaimedSessionId).catch(() => undefined); + await cleanupSession(runnerLostSessionId).catch(() => undefined); + await cleanupSession(freshActiveSessionId).catch(() => undefined); await Promise.all(reopenedStores.map(async (store) => await store.close())); await Promise.all([primary.close(), concurrent.close(), control.end()]); } @@ -107,6 +121,73 @@ async function assertConcurrentAdmissionAndRestartRecovery(): Promise { assert.equal(session?.lastCommandId, next.command.id); } +async function assertStaleClaimedAdmissionCreatesFreshTurn(): Promise { + const stale = await primary.admitSessionTurn(admissionInput(staleClaimedSessionId, "postgres stale claimed before ack", `pg-stale-claimed-${suffix}`)); + await primary.claimRun(stale.run.id, `runner_pg_stale_claimed_${suffix}`, 1); + await delay(10); + const replacement = await primary.admitSessionTurn(admissionInput(staleClaimedSessionId, "postgres replace stale claimed", `pg-stale-claimed-replacement-${suffix}`)); + assert.equal(replacement.disposition, "created"); + assert.notEqual(replacement.run.id, stale.run.id); + const staleRun = await primary.getRun(stale.run.id); + const staleCommand = await primary.getCommand(stale.command.id); + assert.equal(staleRun.status, "failed"); + assert.equal(staleRun.failureKind, "infra-failed"); + assert.equal(staleRun.claimedBy, null); + assert.equal(staleRun.leaseExpiresAt, null); + assert.equal(staleCommand.state, "failed"); + assert.equal((await primary.getRunnerDispatchIntent(stale.command.id))?.state, "cancelled"); + assert.equal((await primary.listEvents(stale.run.id, 0, 100)).some((event) => event.payload.reason === "session-turn-admission-stale-lease"), true); + const session = await primary.getSession(staleClaimedSessionId); + assert.equal(session?.activeRunId, replacement.run.id); + assert.equal(session?.activeCommandId, replacement.command.id); +} + +async function assertRunnerLostConcurrentAdmissionCreatesOneFreshTurn(): Promise { + const stale = await primary.admitSessionTurn(admissionInput(runnerLostSessionId, "postgres runner lost after ack", `pg-runner-lost-${suffix}`)); + await primary.claimRun(stale.run.id, `runner_pg_lost_${suffix}`, 1); + await primary.ackCommand(stale.command.id); + await delay(10); + const replacementInput = admissionInput(runnerLostSessionId, "postgres continue after runner loss", `pg-runner-lost-replacement-${suffix}`); + const [first, second] = await Promise.all([primary.admitSessionTurn(replacementInput), concurrent.admitSessionTurn(replacementInput)]); + assert.equal(first.run.id, second.run.id); + assert.equal(first.command.id, second.command.id); + assert.deepEqual(new Set([first.disposition, second.disposition]), new Set(["created", "replayed"])); + assert.notEqual(first.run.id, stale.run.id); + assert.equal((await primary.getRun(stale.run.id)).status, "failed"); + assert.equal((await primary.getCommand(stale.command.id)).state, "failed"); + const counts = await control.query( + `SELECT + (SELECT count(*)::int FROM agentrun_runs WHERE session_ref->>'sessionId' = $1) AS run_count, + (SELECT count(*)::int FROM agentrun_commands c JOIN agentrun_runs r ON r.id = c.run_id WHERE r.session_ref->>'sessionId' = $1) AS command_count, + (SELECT count(*)::int FROM agentrun_runner_dispatch_intents i JOIN agentrun_runs r ON r.id = i.run_id WHERE r.session_ref->>'sessionId' = $1) AS intent_count`, + [runnerLostSessionId], + ); + assert.deepEqual(counts.rows[0], { run_count: 2, command_count: 2, intent_count: 2 }); + const session = await primary.getSession(runnerLostSessionId); + assert.equal(session?.activeRunId, first.run.id); + assert.equal(session?.activeCommandId, first.command.id); +} + +async function assertFreshActiveAdmissionRemainsProtected(): Promise { + const active = await primary.admitSessionTurn(admissionInput(freshActiveSessionId, "postgres fresh active turn", `pg-fresh-active-${suffix}`)); + await primary.claimRun(active.run.id, `runner_pg_fresh_${suffix}`, 60_000); + await primary.ackCommand(active.command.id); + const steer = await primary.admitSessionSteer({ + sessionId: freshActiveSessionId, + command: validateCreateCommand({ type: "steer", payload: { prompt: "postgres steer fresh active turn" }, idempotencyKey: `pg-fresh-active-steer-${suffix}` }), + }); + assert.equal(steer?.run.id, active.run.id); + assert.equal(steer?.targetCommand.id, active.command.id); + await assert.rejects( + () => concurrent.admitSessionTurn(admissionInput(freshActiveSessionId, "postgres conflicting fresh turn", `pg-fresh-active-conflict-${suffix}`)), + (error) => error instanceof AgentRunError + && error.failureKind === "runner-lease-conflict" + && error.details?.reason === "session-turn-admission-active", + ); + assert.equal((await primary.getRun(active.run.id)).status, "claimed"); + assert.equal((await primary.getCommand(active.command.id)).state, "acknowledged"); +} + async function assertLegacyHalfCommitRecoveryAfterRestart(): Promise { const input = admissionInput(legacySessionId, "postgres recover legacy half commit", `pg-admission-legacy-${suffix}`); const run = await primary.createRun(input.run); @@ -133,9 +214,15 @@ async function assertLegacyHalfCommitRecoveryAfterRestart(): Promise { } async function assertAdmissionTransactionRollback(): Promise { + const stale = await primary.admitSessionTurn(admissionInput(staleRollbackSessionId, "postgres stale admission before rollback", `pg-admission-stale-before-rollback-${suffix}`)); + await primary.claimRun(stale.run.id, `runner_pg_stale_rollback_${suffix}`, 1); + await primary.ackCommand(stale.command.id); + await delay(10); await installRollbackTrigger(); const input = admissionInput(rollbackSessionId, `postgres admission rollback ${suffix}`, `pg-admission-rollback-${suffix}`); await assert.rejects(() => primary.admitSessionTurn(input), /agentrun selftest reject session admission/u); + const staleReplacementInput = admissionInput(staleRollbackSessionId, `postgres admission rollback ${suffix}`, `pg-admission-stale-rollback-${suffix}`); + await assert.rejects(() => primary.admitSessionTurn(staleReplacementInput), /agentrun selftest reject session admission/u); await dropRollbackTrigger(); const counts = await control.query( `SELECT @@ -148,6 +235,23 @@ async function assertAdmissionTransactionRollback(): Promise { [rollbackSessionId], ); assert.deepEqual(counts.rows[0], { session_count: 0, run_count: 0, command_count: 0, event_count: 0, intent_count: 0, outbox_count: 0 }); + const staleRun = await primary.getRun(stale.run.id); + const staleCommand = await primary.getCommand(stale.command.id); + const staleSession = await primary.getSession(staleRollbackSessionId); + assert.equal(staleRun.status, "claimed"); + assert.equal(staleRun.claimedBy, `runner_pg_stale_rollback_${suffix}`); + assert.equal(staleCommand.state, "acknowledged"); + assert.equal(staleSession?.activeRunId, stale.run.id); + assert.equal(staleSession?.activeCommandId, stale.command.id); + assert.equal((await primary.listEvents(stale.run.id, 0, 100)).some((event) => event.payload.reason === "session-turn-admission-stale-lease"), false); + const staleCounts = await control.query( + `SELECT + (SELECT count(*)::int FROM agentrun_runs WHERE session_ref->>'sessionId' = $1) AS run_count, + (SELECT count(*)::int FROM agentrun_commands c JOIN agentrun_runs r ON r.id = c.run_id WHERE r.session_ref->>'sessionId' = $1) AS command_count, + (SELECT count(*)::int FROM agentrun_runner_dispatch_intents i JOIN agentrun_runs r ON r.id = i.run_id WHERE r.session_ref->>'sessionId' = $1) AS intent_count`, + [staleRollbackSessionId], + ); + assert.deepEqual(staleCounts.rows[0], { run_count: 1, command_count: 1, intent_count: 1 }); } async function assertLatestSessionResourceRunSkipsJsonNull(): Promise {