diff --git a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts index bdd79ae423..4c96ae9e16 100644 --- a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts +++ b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts @@ -90,7 +90,11 @@ import type { import { renderExecutionTurnContent } from '../execution/types.js'; import type { Env as WorkerEnv, SandboxId } from '../types.js'; import { deriveSharedSandboxId, generateSandboxId } from '../sandbox-id.js'; -import { recordSharedSandboxFailover } from '../shared-sandbox-route.js'; +import { + recordSharedSandboxFailover, + resolveSharedSandboxAssignment, +} from '../shared-sandbox-route.js'; +import type { SHARED_SANDBOX_FAILOVER_SUFFIX } from '../shared-sandbox-route.js'; import { resolveSecret, validateStreamTicket } from '../auth.js'; import { resolveTerminalWrapperClient, type TerminalWrapperClient } from '../terminal/access.js'; @@ -438,6 +442,7 @@ export class CloudAgentSession extends DurableObject { }) => Promise; private sandboxSessionDeleter?: (reason: 'explicit' | 'retention-expired') => Promise; private ephemeralSandboxDestroyer?: () => Promise; + private postExhaustionSandboxDestroyer?: () => Promise; private sharedSandboxFailoverRecorder?: (routeKey: SandboxId) => Promise; private agentRuntime?: AgentRuntime; private sandboxLifecycle?: AgentSandboxLifecycle; @@ -768,6 +773,17 @@ export class CloudAgentSession extends DurableObject { this.sharedSandboxFailoverRecorder ? this.sharedSandboxFailoverRecorder(routeKey) : recordSharedSandboxFailover(this.env.SHARED_SANDBOX_OVERRIDES, routeKey), + resolveSharedSandboxFailover: routeKey => + resolveSharedSandboxAssignment(this.env.SHARED_SANDBOX_OVERRIDES, routeKey), + persistSharedSandboxRecoveryAssignment: input => + this.persistSharedSandboxRecoveryAssignment(input), + destroyIsolatedSandbox: async metadata => { + if (this.postExhaustionSandboxDestroyer) { + await this.postExhaustionSandboxDestroyer(); + return; + } + await createAgentSandbox(this.env, metadata).delete('recovery'); + }, requestAlarmAtOrBefore: deadline => this.scheduleAlarmAtOrBefore(deadline), getSessionIdForLogs: () => this.sessionId, }); @@ -1272,6 +1288,51 @@ export class CloudAgentSession extends DurableObject { return metadata ? parseSessionMetadata(metadata) : null; } + private async persistSharedSandboxRecoveryAssignment(input: { + routeKey: SandboxId; + sourceSandboxId: SandboxId; + replacementSandboxId: SandboxId; + suffix: typeof SHARED_SANDBOX_FAILOVER_SUFFIX; + }): Promise { + const metadata = await this.getStoredMetadata(); + if (!metadata?.workspace) throw new Error('Session workspace metadata is unavailable'); + const route = metadata.workspace.sandboxRoute; + if (route?.kind !== 'shared' || route.routeKey !== input.routeKey) { + throw new Error('Session shared sandbox route no longer matches recovery action'); + } + if ( + metadata.workspace.sandboxId === input.replacementSandboxId && + route.suffix === input.suffix + ) { + return; + } + if (metadata.workspace.sandboxId !== input.sourceSandboxId || route.suffix !== undefined) { + throw new Error('Session shared sandbox assignment no longer matches recovery action'); + } + const updated = serializeSessionMetadata({ + ...metadata, + workspace: { + ...metadata.workspace, + sandboxId: input.replacementSandboxId, + sandboxRoute: { + kind: 'shared', + routeKey: input.routeKey, + suffix: input.suffix, + }, + }, + }); + await this.ctx.storage.put('metadata', updated); + logger + .withFields({ + sessionId: metadata.identity.sessionId, + oldSandboxId: input.sourceSandboxId, + newSandboxId: input.replacementSandboxId, + routeKey: input.routeKey, + logTag: 'shared_sandbox_session_reassigned', + }) + .warn('Moved session to shared sandbox failover assignment'); + } + async getMetadata(): Promise { if (await this.hasDeletionIntent()) return null; return this.getStoredMetadata(); @@ -2412,19 +2473,6 @@ export class CloudAgentSession extends DurableObject { await this.getWrapperSupervisor().runMaintenance(now); await this.destroyEphemeralSandboxIfReady(now); - - try { - await this.getMessageSettlementOutbox().repairTerminalEffects(); - } catch (error) { - terminalEffectRepairRetryAt = Date.now() + PENDING_FLUSH_DEBOUNCE_MS; - logger - .withFields({ - sessionId: this.sessionId, - error: error instanceof Error ? error.message : String(error), - }) - .warn('Terminal effect repair failed; scheduled retry will continue recovery'); - } - await this.retryPendingCallbacks(now); await this.getSessionMessageQueue().recoverPendingInterruption(async () => { await this.interruptAcceptedWrapperMessages(); }); @@ -2437,6 +2485,34 @@ export class CloudAgentSession extends DurableObject { const flushOneResult = await this.flushOnePendingSessionMessage(); pendingFlushRetryAt = flushOneResult.retryAt; remainingPendingCount = flushOneResult.remainingPendingCount; + + const runtimeState = await getWrapperRuntimeState(this.ctx.storage); + const wrapperLease = await getWrapperLease(this.ctx.storage); + const deferTerminalEffects = + remainingPendingCount > 0 && + isWrapperDeliveryHeld(runtimeState, wrapperLease) && + !isWrapperCleanupExhausted(wrapperLease); + if (deferTerminalEffects) { + logger + .withFields({ + sessionId: this.sessionId, + remainingPendingCount, + }) + .debug('Deferring terminal effect repair while queued delivery is held'); + } else { + try { + await this.getMessageSettlementOutbox().repairTerminalEffects(); + } catch (error) { + terminalEffectRepairRetryAt = Date.now() + PENDING_FLUSH_DEBOUNCE_MS; + logger + .withFields({ + sessionId: this.sessionId, + error: error instanceof Error ? error.message : String(error), + }) + .warn('Terminal effect repair failed; scheduled retry will continue recovery'); + } + await this.retryPendingCallbacks(now); + } } catch (error) { alarmWorkFailed = true; logger diff --git a/services/cloud-agent-next/src/session/message-settlement-outbox.test.ts b/services/cloud-agent-next/src/session/message-settlement-outbox.test.ts index b8ddb78e49..303b74c109 100644 --- a/services/cloud-agent-next/src/session/message-settlement-outbox.test.ts +++ b/services/cloud-agent-next/src/session/message-settlement-outbox.test.ts @@ -720,6 +720,39 @@ describe('MessageSettlementOutbox', () => { expect(harness.callbackJobs[0].payload.messageId).toBe(firstMessageId); }); + it('repairs a released terminal wrapper-run callback while the next run remains pending', async () => { + const harness = createHarness({ hasObservedWrapperIdle: false }); + await putSessionMessageState( + harness.storage, + acceptedMessageState(firstMessageId, { url: 'https://example.com/repaired-batch' }) + ); + await storePendingSessionMessage( + harness.storage, + createPendingSessionMessage({ + messageId: secondMessageId, + role: 'user', + content: 'next prompt', + createdAt: 3_000, + }) + ); + + await harness.outbox.persistTerminalTransition( + firstMessageId, + { + kind: 'failed', + reason: 'wrapper became unhealthy', + completionSource: 'wrapper_failure', + }, + { allowIdleBatchWithoutObservedIdle: true } + ); + + await harness.outbox.repairTerminalEffects(); + await harness.outbox.repairTerminalEffects(); + + expect(harness.callbackJobs).toHaveLength(1); + expect(harness.callbackJobs[0].payload.messageId).toBe(firstMessageId); + }); + it('keeps a wrapper-run callback blocked while that run has a nonterminal accepted message', async () => { const harness = createHarness(); await putSessionMessageState( diff --git a/services/cloud-agent-next/src/session/message-settlement-outbox.ts b/services/cloud-agent-next/src/session/message-settlement-outbox.ts index d6574c7eec..ae057a5ff7 100644 --- a/services/cloud-agent-next/src/session/message-settlement-outbox.ts +++ b/services/cloud-agent-next/src/session/message-settlement-outbox.ts @@ -792,9 +792,21 @@ export function createMessageSettlementOutbox( } const batch = await getCurrentIdleBatchCallbackState(); if (batch) { - await finalizeIdleBatchCallbackIfReady({ - allowWithoutObservedIdle: batch.allowWithoutObservedIdle, - }); + const representative = batch.representativeMessageId + ? await getSessionMessageState(storage, batch.representativeMessageId) + : null; + if ( + representative?.wrapperRunId && + (batch.wrapperTerminalWaitReleasedAt !== undefined || + (batch.allowWithoutObservedIdle && + representative.completionSource === 'wrapper_failure')) + ) { + await finalizeTerminalWrapperRunCallbackIfReady(representative.wrapperRunId); + } else { + await finalizeIdleBatchCallbackIfReady({ + allowWithoutObservedIdle: batch.allowWithoutObservedIdle, + }); + } } } catch (error) { await requestAlarmAtOrBefore(Date.now() + 1_000); diff --git a/services/cloud-agent-next/src/session/wrapper-runtime-state.test.ts b/services/cloud-agent-next/src/session/wrapper-runtime-state.test.ts index 8d0f1012a4..7ba34ebdf7 100644 --- a/services/cloud-agent-next/src/session/wrapper-runtime-state.test.ts +++ b/services/cloud-agent-next/src/session/wrapper-runtime-state.test.ts @@ -252,6 +252,132 @@ describe('WrapperLease', () => { ).toEqual(exhausted); }); + it('persists one post-exhaustion action and retries it with a durable capped deadline', () => { + const expectedLease = { + target: { kind: 'session' as const }, + requestedAt: 100, + nextInstanceGeneration: 3, + }; + const prepared = reduceSandboxRecoveryState(undefined, { + type: 'prepare_post_exhaustion', + recovery: { + kind: 'isolated-destroy', + sourceSandboxId: 'ses-abcdef' as SandboxId, + expectedLease, + attempts: 0, + nextAttemptAt: 300, + }, + }); + const duplicate = reduceSandboxRecoveryState(prepared, { + type: 'prepare_post_exhaustion', + recovery: { + kind: 'isolated-destroy', + sourceSandboxId: 'ses-competing' as SandboxId, + expectedLease, + attempts: 0, + nextAttemptAt: 301, + }, + }); + const retrying = reduceSandboxRecoveryState(duplicate, { + type: 'record_post_exhaustion_retry', + expectedLease, + expectedAttempts: 0, + nextAttemptAt: 2_300, + error: 'provider unavailable', + }); + + expect(duplicate).toEqual(prepared); + expect(retrying).toMatchObject({ + postExhaustionRecovery: { + sourceSandboxId: 'ses-abcdef', + attempts: 1, + nextAttemptAt: 2_300, + lastError: 'provider unavailable', + }, + }); + expect(nextSandboxRecoveryDeadline(retrying)).toBe(2_300); + }); + + it('checkpoints a completed isolated destroy and keeps it across retry recording', () => { + const expectedLease = { + target: { kind: 'session' as const }, + requestedAt: 100, + nextInstanceGeneration: 3, + }; + const prepared = reduceSandboxRecoveryState(undefined, { + type: 'prepare_post_exhaustion', + recovery: { + kind: 'isolated-destroy', + sourceSandboxId: 'ses-abcdef' as SandboxId, + expectedLease, + attempts: 0, + nextAttemptAt: 300, + }, + }); + + const destroyed = reduceSandboxRecoveryState(prepared, { + type: 'record_post_exhaustion_destroy', + expectedLease, + destroyedAt: 400, + }); + expect(destroyed).toMatchObject({ + postExhaustionRecovery: { kind: 'isolated-destroy', destroyedAt: 400 }, + }); + expect( + reduceSandboxRecoveryState(destroyed, { + type: 'record_post_exhaustion_destroy', + expectedLease: { ...expectedLease, requestedAt: 99 }, + destroyedAt: 500, + }) + ).toEqual(destroyed); + + const retried = reduceSandboxRecoveryState(destroyed, { + type: 'record_post_exhaustion_retry', + expectedLease, + expectedAttempts: 0, + nextAttemptAt: 2_400, + error: 'settlement failed', + }); + expect(retried).toMatchObject({ + postExhaustionRecovery: { destroyedAt: 400, attempts: 1, nextAttemptAt: 2_400 }, + }); + }); + + it('clears only the exhausted lease captured by completed sandbox recovery', () => { + const requested = reduceWrapperLease( + { state: 'none', nextInstanceGeneration: 3 }, + { + type: 'request_stop', + target: { kind: 'session' }, + reason: 'unhealthy-wrapper', + now: 100, + } + ); + if (requested.state !== 'stop_needed') throw new Error('Expected cleanup request'); + const exhausted = reduceWrapperLease( + { ...requested, attempts: 5 }, + { type: 'cleanup_exhausted', now: 300, error: 'still present' } + ); + const expectedLease = { + target: { kind: 'session' as const }, + requestedAt: 100, + nextInstanceGeneration: 3, + }; + + expect( + reduceWrapperLease(exhausted, { + type: 'sandbox_recovery_completed', + expectedLease: { ...expectedLease, requestedAt: 99 }, + }) + ).toEqual(exhausted); + expect( + reduceWrapperLease(exhausted, { + type: 'sandbox_recovery_completed', + expectedLease, + }) + ).toEqual({ state: 'none', nextInstanceGeneration: 3 }); + }); + it('counts only fresh list-processes timeouts and persists publication deadlines', () => { const routeKey = `usr-${'d'.repeat(48)}` as SandboxId; expect( diff --git a/services/cloud-agent-next/src/session/wrapper-runtime-state.ts b/services/cloud-agent-next/src/session/wrapper-runtime-state.ts index 1e7eda3254..42bd9631cc 100644 --- a/services/cloud-agent-next/src/session/wrapper-runtime-state.ts +++ b/services/cloud-agent-next/src/session/wrapper-runtime-state.ts @@ -45,6 +45,34 @@ const wrapperStopReasonSchema = z.enum([ const SharedSandboxRouteKeySchema = z.custom( value => typeof value === 'string' && isGeneratedSharedSandboxId(value) ); +const SandboxIdSchema = z.custom(value => typeof value === 'string' && value.length > 0); +const exhaustedWrapperLeaseIdentitySchema = z.object({ + target: wrapperStopTargetSchema, + requestedAt: z.number().int().nonnegative(), + nextInstanceGeneration: z.number().int().positive(), +}); +const postExhaustionRecoveryCommonSchema = { + expectedLease: exhaustedWrapperLeaseIdentitySchema, + attempts: z.number().int().nonnegative(), + nextAttemptAt: z.number().int().nonnegative(), + lastError: z.string().optional(), +}; +const postExhaustionRecoverySchema = z.discriminatedUnion('kind', [ + z.object({ + kind: z.literal('isolated-destroy'), + sourceSandboxId: SandboxIdSchema, + destroyedAt: z.number().int().nonnegative().optional(), + ...postExhaustionRecoveryCommonSchema, + }), + z.object({ + kind: z.literal('shared-failover'), + routeKey: SharedSandboxRouteKeySchema, + sourceSandboxId: SandboxIdSchema, + replacementSandboxId: SandboxIdSchema.optional(), + routeSuffix: z.literal('shared-slot-v1').optional(), + ...postExhaustionRecoveryCommonSchema, + }), +]); const sharedSandboxFailoverPublicationSchema = z.discriminatedUnion('status', [ z.object({ status: z.literal('pending'), @@ -67,8 +95,9 @@ const sharedSandboxFailoverPublicationSchema = z.discriminatedUnion('status', [ ]); const sandboxRecoveryStateSchema = z .object({ - listProcessesTimeouts: z.number().int().positive(), + listProcessesTimeouts: z.number().int().nonnegative(), failoverPublication: sharedSandboxFailoverPublicationSchema.optional(), + postExhaustionRecovery: postExhaustionRecoverySchema.optional(), }) .superRefine((state, context) => { if (state.failoverPublication && state.listProcessesTimeouts < 2) { @@ -133,6 +162,8 @@ const wrapperLeaseSchema = z export type WrapperLease = z.infer; export type SandboxRecoveryState = z.infer; +export type ExhaustedWrapperLeaseIdentity = z.infer; +export type PostExhaustionRecovery = z.infer; export type SandboxRecoveryEvent = | { @@ -153,7 +184,28 @@ export type SandboxRecoveryEvent = outcome: 'recorded' | 'exhausted'; routeKey: SandboxId; expectedFailedAttempts: number; - }; + } + | { type: 'prepare_post_exhaustion'; recovery: PostExhaustionRecovery } + | { + type: 'record_post_exhaustion_assignment'; + expectedLease: ExhaustedWrapperLeaseIdentity; + routeKey: SandboxId; + replacementSandboxId: SandboxId; + routeSuffix: 'shared-slot-v1'; + } + | { + type: 'record_post_exhaustion_destroy'; + expectedLease: ExhaustedWrapperLeaseIdentity; + destroyedAt: number; + } + | { + type: 'record_post_exhaustion_retry'; + expectedLease: ExhaustedWrapperLeaseIdentity; + expectedAttempts: number; + nextAttemptAt: number; + error: string; + } + | { type: 'settle_post_exhaustion'; expectedLease: ExhaustedWrapperLeaseIdentity }; export type WrapperLeaseEvent = | { type: 'allocate'; instance: WrapperInstanceLease; startupDeadlineAt: number } @@ -167,7 +219,8 @@ export type WrapperLeaseEvent = | { type: 'stop_absent'; attemptId: string } | { type: 'stop_not_confirmed'; attemptId: string; retryAt: number; error: string } | { type: 'stop_attempt_expired'; attemptId: string; retryAt: number } - | { type: 'cleanup_exhausted'; attemptId?: string; now: number; error: string }; + | { type: 'cleanup_exhausted'; attemptId?: string; now: number; error: string } + | { type: 'sandbox_recovery_completed'; expectedLease: ExhaustedWrapperLeaseIdentity }; export const emptyWrapperLease = (): WrapperLease => ({ state: 'none', @@ -250,6 +303,7 @@ export async function recordSandboxInspectionFailure( export async function clearSettledSandboxRecovery(storage: DurableObjectStorage): Promise { const state = await getSandboxRecoveryState(storage); if ( + state?.postExhaustionRecovery || state?.failoverPublication?.status === 'pending' || (state && state.listProcessesTimeouts >= 2 && !state.failoverPublication) ) { @@ -328,15 +382,134 @@ export function reduceSandboxRecoveryState( }, }; } + case 'prepare_post_exhaustion': + if (state?.postExhaustionRecovery) return state; + return { + listProcessesTimeouts: state?.listProcessesTimeouts ?? 0, + ...state, + postExhaustionRecovery: event.recovery, + }; + case 'record_post_exhaustion_assignment': { + const recovery = state?.postExhaustionRecovery; + if ( + !state || + recovery?.kind !== 'shared-failover' || + !isSameExhaustedWrapperLeaseIdentity(recovery.expectedLease, event.expectedLease) || + recovery.routeKey !== event.routeKey + ) { + return state; + } + return { + ...state, + postExhaustionRecovery: { + ...recovery, + replacementSandboxId: event.replacementSandboxId, + routeSuffix: event.routeSuffix, + }, + }; + } + case 'record_post_exhaustion_destroy': { + const recovery = state?.postExhaustionRecovery; + if ( + !state || + recovery?.kind !== 'isolated-destroy' || + !isSameExhaustedWrapperLeaseIdentity(recovery.expectedLease, event.expectedLease) + ) { + return state; + } + return { + ...state, + postExhaustionRecovery: { ...recovery, destroyedAt: event.destroyedAt }, + }; + } + case 'record_post_exhaustion_retry': { + const recovery = state?.postExhaustionRecovery; + if ( + !state || + !recovery || + !isSameExhaustedWrapperLeaseIdentity(recovery.expectedLease, event.expectedLease) || + recovery.attempts !== event.expectedAttempts + ) { + return state; + } + return { + ...state, + postExhaustionRecovery: { + ...recovery, + attempts: recovery.attempts + 1, + nextAttemptAt: event.nextAttemptAt, + lastError: event.error, + }, + }; + } + case 'settle_post_exhaustion': { + const recovery = state?.postExhaustionRecovery; + if ( + !state || + !recovery || + !isSameExhaustedWrapperLeaseIdentity(recovery.expectedLease, event.expectedLease) + ) { + return state; + } + const { postExhaustionRecovery: _settled, ...remaining } = state; + if (remaining.listProcessesTimeouts < 2 || remaining.failoverPublication) { + return remaining; + } + return { + ...remaining, + failoverPublication: + recovery.kind === 'shared-failover' + ? { status: 'recorded', routeKey: recovery.routeKey } + : { status: 'not-applicable' }, + }; + } } } +export function exhaustedWrapperLeaseIdentity( + lease: Extract +): ExhaustedWrapperLeaseIdentity { + return { + target: lease.target, + requestedAt: lease.requestedAt, + nextInstanceGeneration: lease.nextInstanceGeneration, + }; +} + +export function isSameExhaustedWrapperLeaseIdentity( + left: ExhaustedWrapperLeaseIdentity, + right: ExhaustedWrapperLeaseIdentity +): boolean { + return ( + left.requestedAt === right.requestedAt && + left.nextInstanceGeneration === right.nextInstanceGeneration && + JSON.stringify(left.target) === JSON.stringify(right.target) + ); +} + +export function doesRecoveryMatchExhaustedLease( + recovery: PostExhaustionRecovery, + lease: WrapperLease +): boolean { + return ( + isWrapperCleanupExhausted(lease) && + isSameExhaustedWrapperLeaseIdentity( + recovery.expectedLease, + exhaustedWrapperLeaseIdentity(lease) + ) + ); +} + export function nextSandboxRecoveryDeadline( state: SandboxRecoveryState | undefined ): number | undefined { - return state?.failoverPublication?.status === 'pending' - ? state.failoverPublication.nextAttemptAt - : undefined; + const deadlines = [ + state?.failoverPublication?.status === 'pending' + ? state.failoverPublication.nextAttemptAt + : undefined, + state?.postExhaustionRecovery?.nextAttemptAt, + ].filter((deadline): deadline is number => deadline !== undefined); + return deadlines.length > 0 ? Math.min(...deadlines) : undefined; } export function reduceWrapperLease(state: WrapperLease, event: WrapperLeaseEvent): WrapperLease { @@ -442,6 +615,17 @@ export function reduceWrapperLease(state: WrapperLease, event: WrapperLeaseEvent lastError: event.error, exhaustedAt: event.now, }; + case 'sandbox_recovery_completed': + if ( + !isWrapperCleanupExhausted(state) || + !isSameExhaustedWrapperLeaseIdentity( + exhaustedWrapperLeaseIdentity(state), + event.expectedLease + ) + ) { + return state; + } + return { state: 'none', nextInstanceGeneration: state.nextInstanceGeneration }; } } diff --git a/services/cloud-agent-next/src/session/wrapper-supervisor.test.ts b/services/cloud-agent-next/src/session/wrapper-supervisor.test.ts index ec6fc666d9..6643e070f9 100644 --- a/services/cloud-agent-next/src/session/wrapper-supervisor.test.ts +++ b/services/cloud-agent-next/src/session/wrapper-supervisor.test.ts @@ -6,7 +6,7 @@ import { createMessageSettlementOutbox, type MessageSettlementOutboxStorage, } from './message-settlement-outbox.js'; -import { storePendingSessionMessage } from './pending-messages.js'; +import { listPendingSessionMessages, storePendingSessionMessage } from './pending-messages.js'; import { getSessionMessageState, putSessionMessageState, @@ -24,6 +24,8 @@ import { } from './wrapper-runtime-state.js'; import type { LatestAssistantMessage } from './types.js'; import type { SandboxId } from '../types.js'; +import { deriveSharedSandboxId } from '../sandbox-id.js'; +import { SHARED_SANDBOX_FAILOVER_SUFFIX } from '../shared-sandbox-route.js'; vi.mock('@cloudflare/sandbox', () => ({ getSandbox: vi.fn(), @@ -122,6 +124,17 @@ function createHarness( wrapperRunId: string ) => Promise; recordSharedSandboxFailover?: (routeKey: SandboxId) => Promise; + destroyIsolatedSandbox?: (metadata: SessionMetadata) => Promise; + resolveSharedSandboxFailover?: (routeKey: SandboxId) => Promise<{ + sandboxId: SandboxId; + suffix?: typeof SHARED_SANDBOX_FAILOVER_SUFFIX; + }>; + persistSharedSandboxRecoveryAssignment?: (input: { + routeKey: SandboxId; + sourceSandboxId: SandboxId; + replacementSandboxId: SandboxId; + suffix: typeof SHARED_SANDBOX_FAILOVER_SUFFIX; + }) => Promise; } ) { const getAssistantMessageForUserMessage = @@ -158,6 +171,17 @@ function createHarness( const recordSharedSandboxFailover = vi.fn( options?.recordSharedSandboxFailover ?? (async () => {}) ); + const destroyIsolatedSandbox = vi.fn(options?.destroyIsolatedSandbox ?? (async () => {})); + const resolveSharedSandboxFailover = vi.fn( + options?.resolveSharedSandboxFailover ?? + (async routeKey => ({ + sandboxId: await deriveSharedSandboxId(routeKey, SHARED_SANDBOX_FAILOVER_SUFFIX), + suffix: SHARED_SANDBOX_FAILOVER_SUFFIX, + })) + ); + const persistSharedSandboxRecoveryAssignment = vi.fn( + options?.persistSharedSandboxRecoveryAssignment ?? (async () => {}) + ); const supervisor = createWrapperSupervisor({ storage, agentRuntime: { @@ -175,6 +199,9 @@ function createHarness( options?.ensureAcceptedMessageBeforeTerminal ?? (async () => {}), stopWrappers, recordSharedSandboxFailover: routeKey => recordSharedSandboxFailover(routeKey), + resolveSharedSandboxFailover: routeKey => resolveSharedSandboxFailover(routeKey), + persistSharedSandboxRecoveryAssignment: input => persistSharedSandboxRecoveryAssignment(input), + destroyIsolatedSandbox: metadata => destroyIsolatedSandbox(metadata), requestAlarmAtOrBefore: async deadline => { requestedAlarms.push(deadline); }, @@ -191,6 +218,9 @@ function createHarness( requestedAlarms, requestPendingDrainIfNeeded, recordSharedSandboxFailover, + resolveSharedSandboxFailover, + persistSharedSandboxRecoveryAssignment, + destroyIsolatedSandbox, settlementOutbox, supervisor, }; @@ -286,13 +316,14 @@ describe('WrapperSupervisor', () => { await harness.supervisor.runMaintenance(grace.disconnectedAt + 10_001); - await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ - state: 'stop_needed', - reason: 'unhealthy-wrapper', + await expect(getWrapperLease(harness.storage)).resolves.toEqual({ + state: 'none', + nextInstanceGeneration: 2, }); await expect(getWrapperRuntimeState(harness.storage)).resolves.toEqual({ wrapperGeneration: 5, }); + expect(harness.stopWrappers).toHaveBeenCalledOnce(); }); it('starts disconnect grace while a completed gate callback still waits for wrapper terminal state', async () => { @@ -500,13 +531,13 @@ describe('WrapperSupervisor', () => { await expect(getWrapperRuntimeState(harness.storage)).resolves.toMatchObject({ wrapperGeneration: 5, }); - await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ - state: 'stop_needed', - reason: 'unhealthy-wrapper', + await expect(getWrapperLease(harness.storage)).resolves.toEqual({ + state: 'none', + nextInstanceGeneration: 1, }); - expect(harness.requestPendingDrainIfNeeded).not.toHaveBeenCalled(); + expect(harness.requestPendingDrainIfNeeded).toHaveBeenCalledOnce(); expect(harness.events.map(event => event.streamEventType)).toEqual(['cloud.message.failed']); - expect(harness.stops).toEqual([]); + expect(harness.stopWrappers).toHaveBeenCalledOnce(); }); it('rejects a stale wrapper run before reconnect grace can be cancelled', async () => { @@ -919,12 +950,16 @@ describe('WrapperSupervisor', () => { failureCode: 'wrapper_no_output', }); expect(runtimeState.wrapperConnectionId).toBeUndefined(); - await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ - state: 'stop_needed', - reason: 'unhealthy-wrapper', + await expect(getWrapperLease(harness.storage)).resolves.toEqual({ + state: 'none', + nextInstanceGeneration: 2, }); - expect(harness.requestPendingDrainIfNeeded).not.toHaveBeenCalled(); - expect(harness.stops).toEqual([]); + expect(harness.requestPendingDrainIfNeeded).toHaveBeenCalledOnce(); + expect(harness.stopWrappers).toHaveBeenCalledOnce(); + expect(harness.events).toHaveLength(0); + + await harness.settlementOutbox.repairTerminalEffects(); + expect(harness.events.map(event => event.streamEventType)).toEqual(['cloud.message.failed']); }); @@ -946,6 +981,89 @@ describe('WrapperSupervisor', () => { }); }); + it('recovers queued work after a ping timeout without waiting for terminal effects', async () => { + const pingDeadlineAt = 92_000; + const harness = createHarness([ + liveRuntimeState({ pingDeadlineAt, noOutputDeadlineAt: 332_000 }), + OWNED_WRAPPER_LEASE, + ]); + await putSessionMessageState(harness.storage, acceptedMessage()); + await storePendingSessionMessage(harness.storage, { + messageId: NEWER_MESSAGE_ID, + content: 'queued follow-up', + createdAt: 4_000, + intent: { + turn: { type: 'prompt', messageId: NEWER_MESSAGE_ID, prompt: 'queued follow-up' }, + agent: { mode: 'code', model: 'test-model' }, + }, + }); + + await harness.supervisor.runMaintenance(pingDeadlineAt); + + await expect(getSessionMessageState(harness.storage, MESSAGE_ID)).resolves.toMatchObject({ + status: 'failed', + failureReason: 'wrapper_failure', + failureStage: 'post_dispatch_no_activity', + failureCode: 'wrapper_ping_timeout', + terminalEffects: { + event: 'pending', + callback: { disposition: 'not-required' }, + push: { disposition: 'pending' }, + }, + }); + await expect(listPendingSessionMessages(harness.storage)).resolves.toEqual([ + expect.objectContaining({ messageId: NEWER_MESSAGE_ID }), + ]); + await expect(getWrapperRuntimeState(harness.storage)).resolves.toEqual({ + wrapperGeneration: 5, + }); + await expect(getWrapperLease(harness.storage)).resolves.toEqual({ + state: 'none', + nextInstanceGeneration: 2, + }); + expect(harness.stopWrappers).toHaveBeenCalledOnce(); + expect(harness.requestPendingDrainIfNeeded).toHaveBeenCalledOnce(); + await expect( + harness.supervisor.checkReconnect({ + wrapperRunId: WRAPPER_RUN_ID, + wrapperGeneration: 4, + wrapperConnectionId: WRAPPER_CONNECTION_ID, + }) + ).resolves.toEqual({ accepted: false, reason: 'stale-wrapper-run' }); + }); + + it('keeps queued work fenced when the unhealthy wrapper remains physically present', async () => { + const pingDeadlineAt = 92_000; + const harness = createHarness([ + liveRuntimeState({ pingDeadlineAt, noOutputDeadlineAt: 332_000 }), + OWNED_WRAPPER_LEASE, + ]); + harness.stopWrappers.mockResolvedValue({ status: 'still-present', observed: [] }); + await putSessionMessageState(harness.storage, acceptedMessage()); + await storePendingSessionMessage(harness.storage, { + messageId: NEWER_MESSAGE_ID, + content: 'queued follow-up', + createdAt: 4_000, + intent: { + turn: { type: 'prompt', messageId: NEWER_MESSAGE_ID, prompt: 'queued follow-up' }, + agent: { mode: 'code', model: 'test-model' }, + }, + }); + + await harness.supervisor.runMaintenance(pingDeadlineAt); + + expect(harness.stopWrappers).toHaveBeenCalledOnce(); + await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ + state: 'stop_needed', + reason: 'unhealthy-wrapper', + attempts: 1, + }); + expect(harness.requestPendingDrainIfNeeded).not.toHaveBeenCalled(); + await expect(listPendingSessionMessages(harness.storage)).resolves.toEqual([ + expect.objectContaining({ messageId: NEWER_MESSAGE_ID }), + ]); + }); + it('defers liveness failure while disconnect grace is active for the current connection', async () => { const pingDeadlineAt = 92_000; const noOutputDeadlineAt = 332_000; @@ -1093,10 +1211,11 @@ describe('WrapperSupervisor', () => { failureReason: 'wrapper_disconnected', failureCode: 'wrapper_disconnected', }); - await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ - state: 'stop_needed', - reason: 'unhealthy-wrapper', + await expect(getWrapperLease(harness.storage)).resolves.toEqual({ + state: 'none', + nextInstanceGeneration: 2, }); + expect(harness.stopWrappers).toHaveBeenCalledOnce(); }); it('includes the disconnect grace expiry deadline even when the ping deadline is earlier', async () => { @@ -1156,28 +1275,23 @@ describe('WrapperSupervisor', () => { await harness.supervisor.runMaintenance(10_000); - expect(harness.callbackJobs).toHaveLength(1); - expect(harness.callbackJobs[0].payload).toMatchObject({ - messageId: MESSAGE_ID, - status: 'completed', - }); - expect(harness.callbackJobs[0].payload.gateResult).toBeUndefined(); - const cleanupHold = await getWrapperLease(harness.storage); - expect(cleanupHold).toMatchObject({ - state: 'stop_needed', - reason: 'unhealthy-wrapper', - }); - if (cleanupHold.state !== 'stop_needed') throw new Error('Expected physical cleanup hold'); - expect(harness.requestPendingDrainIfNeeded).not.toHaveBeenCalled(); - - await harness.supervisor.runMaintenance(cleanupHold.nextAttemptAt); - + expect(harness.callbackJobs).toHaveLength(0); expect(harness.stopWrappers).toHaveBeenCalledOnce(); await expect(getWrapperLease(harness.storage)).resolves.toEqual({ state: 'none', nextInstanceGeneration: 2, }); expect(harness.requestPendingDrainIfNeeded).toHaveBeenCalledOnce(); + + await harness.settlementOutbox.repairTerminalEffects(); + await harness.settlementOutbox.repairTerminalEffects(); + + expect(harness.callbackJobs).toHaveLength(1); + expect(harness.callbackJobs[0].payload).toMatchObject({ + messageId: MESSAGE_ID, + status: 'completed', + }); + expect(harness.callbackJobs[0].payload.gateResult).toBeUndefined(); }); it('schedules the updated no-output deadline when it is the next liveness deadline', async () => { @@ -1437,6 +1551,9 @@ describe('WrapperSupervisor', () => { clearInterruptRequest: async () => {}, ensureAcceptedMessageBeforeTerminal: ensureAccepted, recordSharedSandboxFailover: async () => {}, + resolveSharedSandboxFailover: async routeKey => ({ sandboxId: routeKey }), + persistSharedSandboxRecoveryAssignment: async () => {}, + destroyIsolatedSandbox: async () => {}, requestAlarmAtOrBefore: async () => {}, getSessionIdForLogs: () => 'agent_supervisor', }); @@ -2091,7 +2208,7 @@ describe('WrapperSupervisor', () => { }); }); - it('records failover when the fifth cleanup attempt supplies the second fresh timeout', async () => { + it('completes shared recovery when the fifth cleanup attempt supplies the second timeout', async () => { const routeKey = `usr-${'c'.repeat(48)}` as SandboxId; const metadata = { ...createMetadata(), @@ -2127,15 +2244,12 @@ describe('WrapperSupervisor', () => { await harness.supervisor.runMaintenance(10_000); expect(harness.recordSharedSandboxFailover).toHaveBeenCalledWith(routeKey); - await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ - state: 'stop_needed', - attempts: 5, - exhaustedAt: expect.any(Number), - }); - await expect(getSandboxRecoveryState(harness.storage)).resolves.toMatchObject({ - listProcessesTimeouts: 2, - failoverPublication: { status: 'recorded', routeKey }, + await expect(getWrapperLease(harness.storage)).resolves.toEqual({ + state: 'none', + nextInstanceGeneration: 2, }); + await expect(getSandboxRecoveryState(harness.storage)).resolves.toBeUndefined(); + expect(harness.requestPendingDrainIfNeeded).toHaveBeenCalledOnce(); }); it('clears timeout evidence after cleanup confirms wrapper absence', async () => { @@ -2270,9 +2384,7 @@ describe('WrapperSupervisor', () => { reason: 'observation-failed', requestedAt: 1, nextAttemptAt: 3_153_600_000_001, - attempts: 5, - lastError: 'inspection failed', - exhaustedAt: 1, + attempts: 1, }, ], ['sandbox_recovery_state', { listProcessesTimeouts: 2 }], @@ -2301,7 +2413,9 @@ describe('WrapperSupervisor', () => { }); await harness.supervisor.runMaintenance(25_000); expect(publish).toHaveBeenCalledTimes(4); - await expect(harness.supervisor.nextMaintenanceDeadlines()).resolves.toEqual([]); + await expect(harness.supervisor.nextMaintenanceDeadlines()).resolves.toEqual([ + 3_153_600_000_001, + ]); }); it('quarantines cleanup after five failed attempts', async () => { @@ -2357,6 +2471,214 @@ describe('WrapperSupervisor', () => { ); }); + it('retries isolated sandbox destruction after cleanup exhaustion before releasing delivery', async () => { + const metadata = { + ...createMetadata(), + workspace: { sandboxId: 'ses-abcdef' as SandboxId }, + } satisfies SessionMetadata; + const destroy = vi + .fn<(metadata: SessionMetadata) => Promise>() + .mockRejectedValueOnce(new Error('delete unavailable')) + .mockResolvedValueOnce(undefined); + const harness = createHarness( + [ + [ + 'wrapper_lease', + { + state: 'stop_needed', + nextInstanceGeneration: 4, + target: { kind: 'session' }, + reason: 'unhealthy-wrapper', + requestedAt: 100, + nextAttemptAt: 1_000, + attempts: 4, + }, + ], + ], + { metadata, destroyIsolatedSandbox: destroy } + ); + harness.stopWrappers.mockResolvedValue({ status: 'still-present', observed: [] }); + const clock = vi.spyOn(Date, 'now').mockReturnValue(1_000); + + try { + await harness.supervisor.runMaintenance(1_000); + await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ + state: 'stop_needed', + exhaustedAt: 1_000, + }); + await expect(getSandboxRecoveryState(harness.storage)).resolves.toMatchObject({ + postExhaustionRecovery: { + kind: 'isolated-destroy', + attempts: 1, + nextAttemptAt: 3_000, + }, + }); + expect(harness.requestPendingDrainIfNeeded).not.toHaveBeenCalled(); + + clock.mockReturnValue(3_000); + await harness.supervisor.runMaintenance(3_000); + } finally { + clock.mockRestore(); + } + + expect(destroy).toHaveBeenCalledTimes(2); + await expect(getWrapperLease(harness.storage)).resolves.toEqual({ + state: 'none', + nextInstanceGeneration: 4, + }); + await expect(getSandboxRecoveryState(harness.storage)).resolves.toBeUndefined(); + expect(harness.requestPendingDrainIfNeeded).toHaveBeenCalledOnce(); + }); + + it('settles a recorded isolated destruction without destroying the sandbox again', async () => { + const metadata = { + ...createMetadata(), + workspace: { sandboxId: 'ses-abcdef' as SandboxId }, + } satisfies SessionMetadata; + const harness = createHarness( + [ + [ + 'wrapper_lease', + { + state: 'stop_needed', + nextInstanceGeneration: 4, + target: { kind: 'session' }, + reason: 'unhealthy-wrapper', + requestedAt: 100, + nextAttemptAt: 3_153_600_000_100, + attempts: 5, + exhaustedAt: 100, + }, + ], + [ + 'sandbox_recovery_state', + { + listProcessesTimeouts: 0, + postExhaustionRecovery: { + kind: 'isolated-destroy', + sourceSandboxId: 'ses-abcdef', + destroyedAt: 950, + expectedLease: { + target: { kind: 'session' }, + requestedAt: 100, + nextInstanceGeneration: 4, + }, + attempts: 1, + nextAttemptAt: 1_000, + }, + }, + ], + ], + { metadata } + ); + + await harness.supervisor.runMaintenance(1_000); + + expect(harness.destroyIsolatedSandbox).not.toHaveBeenCalled(); + await expect(getWrapperLease(harness.storage)).resolves.toEqual({ + state: 'none', + nextInstanceGeneration: 4, + }); + await expect(getSandboxRecoveryState(harness.storage)).resolves.toBeUndefined(); + expect(harness.requestPendingDrainIfNeeded).toHaveBeenCalledOnce(); + }); + + it('moves an exhausted shared session to failover without destroying shared compute', async () => { + const routeKey = `usr-${'f'.repeat(48)}` as SandboxId; + const replacementSandboxId = await deriveSharedSandboxId( + routeKey, + SHARED_SANDBOX_FAILOVER_SUFFIX + ); + const metadata = { + ...createMetadata(), + workspace: { + sandboxId: routeKey, + sandboxRoute: { kind: 'shared', routeKey }, + }, + } satisfies SessionMetadata; + const harness = createHarness( + [ + [ + 'wrapper_lease', + { + state: 'stop_needed', + nextInstanceGeneration: 5, + target: { kind: 'session' }, + reason: 'unhealthy-wrapper', + requestedAt: 100, + nextAttemptAt: 3_153_600_000_100, + attempts: 5, + exhaustedAt: 100, + }, + ], + ], + { metadata } + ); + + await harness.supervisor.runMaintenance(1_000); + + expect(harness.destroyIsolatedSandbox).not.toHaveBeenCalled(); + expect(harness.recordSharedSandboxFailover).toHaveBeenCalledWith(routeKey); + expect(harness.persistSharedSandboxRecoveryAssignment).toHaveBeenCalledWith({ + routeKey, + sourceSandboxId: routeKey, + replacementSandboxId, + suffix: SHARED_SANDBOX_FAILOVER_SUFFIX, + }); + await expect(getWrapperLease(harness.storage)).resolves.toEqual({ + state: 'none', + nextInstanceGeneration: 5, + }); + expect(harness.requestPendingDrainIfNeeded).toHaveBeenCalledOnce(); + }); + + it('settles a stale sandbox recovery action without releasing a newer wrapper lease', async () => { + const metadata = { + ...createMetadata(), + workspace: { sandboxId: 'ses-abcdef' as SandboxId }, + } satisfies SessionMetadata; + const harness = createHarness( + [ + [ + 'wrapper_lease', + { + state: 'owns_wrapper', + nextInstanceGeneration: 6, + instance: { instanceId: 'instance_newer', instanceGeneration: 5 }, + }, + ], + [ + 'sandbox_recovery_state', + { + listProcessesTimeouts: 0, + postExhaustionRecovery: { + kind: 'isolated-destroy', + sourceSandboxId: 'ses-abcdef', + expectedLease: { + target: { kind: 'session' }, + requestedAt: 100, + nextInstanceGeneration: 5, + }, + attempts: 0, + nextAttemptAt: 1_000, + }, + }, + ], + ], + { metadata } + ); + + await harness.supervisor.runMaintenance(1_000); + + expect(harness.destroyIsolatedSandbox).not.toHaveBeenCalled(); + await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ + state: 'owns_wrapper', + instance: { instanceId: 'instance_newer' }, + }); + await expect(getSandboxRecoveryState(harness.storage)).resolves.toBeUndefined(); + expect(harness.requestPendingDrainIfNeeded).not.toHaveBeenCalled(); + }); + it('quarantines the fifth cleanup attempt when its watchdog expires', async () => { const harness = createHarness([ [ diff --git a/services/cloud-agent-next/src/session/wrapper-supervisor.ts b/services/cloud-agent-next/src/session/wrapper-supervisor.ts index da9a7895ea..83c6fbe793 100644 --- a/services/cloud-agent-next/src/session/wrapper-supervisor.ts +++ b/services/cloud-agent-next/src/session/wrapper-supervisor.ts @@ -21,6 +21,7 @@ import { import type { WrapperTerminalFailureCode } from '../shared/protocol.js'; import type { LatestAssistantMessage } from './types.js'; import type { SandboxId } from '../types.js'; +import type { SharedSandboxAssignment } from '../shared-sandbox-route.js'; import { MODEL_NOT_FOUND_RUNTIME_DIAGNOSTIC_LOG_CHUNK_SIZE, MODEL_NOT_FOUND_RUNTIME_DIAGNOSTIC_MAX_SERIALIZED_BYTES, @@ -32,6 +33,8 @@ import { clearCurrentWrapperRuntimeLivenessState, clearSettledSandboxRecovery, clearWrapperRuntimeIdentity, + doesRecoveryMatchExhaustedLease, + exhaustedWrapperLeaseIdentity, getSandboxRecoveryState, getWrapperLease, getWrapperRuntimeState, @@ -55,6 +58,7 @@ import { reduceWrapperLease, resetWrapperLivenessAfterReconnect, type WrapperConnectionFence, + type PostExhaustionRecovery, type WrapperRuntimeState, WRAPPER_STOP_MAX_ATTEMPTS, } from './wrapper-runtime-state.js'; @@ -64,6 +68,8 @@ const WRAPPER_PING_TIMEOUT_MS = 30_000; const WRAPPER_STOP_ATTEMPT_TIMEOUT_MS = 45_000; const WRAPPER_STOP_RETRY_DELAYS_MS = [5_000, 10_000, 10_000, 10_000] as const; const SHARED_SANDBOX_FAILOVER_RETRY_DELAYS_MS = [2_000, 4_000, 8_000] as const; +const SANDBOX_RECOVERY_RETRY_DELAYS_MS = [2_000, 4_000, 8_000, 30_000] as const; +const SANDBOX_RECOVERY_MAX_RETRY_DELAY_MS = 30_000; const DISCONNECT_GRACE_KEY = 'disconnect_grace'; const MODEL_NOT_FOUND_SAFE_ERROR_MESSAGE = 'Assistant request failed: model not found'; @@ -147,6 +153,7 @@ export type WrapperSupervisorDependencies = { agentRuntime: Pick; messageSettlementOutbox: Pick< MessageSettlementOutbox, + | 'persistTerminalTransition' | 'terminalizeSessionMessageOnce' | 'observeWrapperTerminalForIdleBatch' | 'releaseWrapperTerminalWaitForIdleBatch' @@ -176,6 +183,14 @@ export type WrapperSupervisorDependencies = { reason: WrapperStopReason; }) => Promise; recordSharedSandboxFailover: (routeKey: SandboxId) => Promise; + resolveSharedSandboxFailover: (routeKey: SandboxId) => Promise; + persistSharedSandboxRecoveryAssignment: (input: { + routeKey: SandboxId; + sourceSandboxId: SandboxId; + replacementSandboxId: SandboxId; + suffix: 'shared-slot-v1'; + }) => Promise; + destroyIsolatedSandbox: (metadata: SessionMetadata) => Promise; requestAlarmAtOrBefore?: (deadline: number) => Promise; getSessionIdForLogs: () => string | undefined; }; @@ -360,6 +375,9 @@ export function createWrapperSupervisor( ensureAcceptedMessageBeforeTerminal, stopWrappers, recordSharedSandboxFailover, + resolveSharedSandboxFailover, + persistSharedSandboxRecoveryAssignment, + destroyIsolatedSandbox, requestAlarmAtOrBefore, getSessionIdForLogs, } = dependencies; @@ -633,23 +651,29 @@ export function createWrapperSupervisor( const acceptedMessages = await listNonTerminalAcceptedMessages(storage, state.wrapperRunId); for (const message of acceptedMessages) { const activityObserved = message.agentActivityObservedAt !== undefined; - await messageSettlementOutbox.terminalizeSessionMessageOnce(message.messageId, { - kind: 'failed', - reason: 'wrapper_failure', - error, - completionSource: 'wrapper_failure', - failureStage: activityObserved ? 'agent_activity' : 'post_dispatch_no_activity', - failureCode: activityObserved ? 'wrapper_error_after_activity' : failureCode, - }); + await messageSettlementOutbox.persistTerminalTransition( + message.messageId, + { + kind: 'failed', + reason: 'wrapper_failure', + error, + completionSource: 'wrapper_failure', + failureStage: activityObserved ? 'agent_activity' : 'post_dispatch_no_activity', + failureCode: activityObserved ? 'wrapper_error_after_activity' : failureCode, + }, + { allowIdleBatchWithoutObservedIdle: true } + ); } + logger + .withFields({ + sessionId: getSessionIdForLogs(), + wrapperRunId: state.wrapperRunId, + wrapperGeneration: state.wrapperGeneration, + acceptedMessageCount: acceptedMessages.length, + }) + .info('Unhealthy wrapper terminal states persisted'); + await messageSettlementOutbox.releaseWrapperTerminalWaitForIdleBatch(); - if (isWrapperRunFinalizing(state) && state.wrapperRunId) { - await messageSettlementOutbox.finalizeTerminalWrapperRunCallbackIfReady(state.wrapperRunId); - } else { - await messageSettlementOutbox.finalizeIdleBatchCallbackIfReady({ - allowWithoutObservedIdle: true, - }); - } if (state.wrapperConnectionId) { await clearCurrentWrapperRuntimeFailureState( @@ -1073,6 +1097,219 @@ export function createWrapperSupervisor( .error('Wrapper cleanup attempt limit exhausted'); } + function sandboxRecoveryRetryAt(now: number, attempts: number): number { + const delay = + SANDBOX_RECOVERY_RETRY_DELAYS_MS[ + Math.min(attempts, SANDBOX_RECOVERY_RETRY_DELAYS_MS.length - 1) + ]; + return now + delay; + } + + function createPostExhaustionRecovery( + metadata: SessionMetadata, + lease: Extract>, { state: 'stop_needed' }>, + now: number + ): PostExhaustionRecovery | undefined { + const sandboxId = metadata.workspace?.sandboxId; + const route = metadata.workspace?.sandboxRoute; + if (!sandboxId) return undefined; + const common = { + expectedLease: exhaustedWrapperLeaseIdentity(lease), + attempts: 0, + nextAttemptAt: now, + }; + if (route?.kind === 'shared') { + return { + kind: 'shared-failover', + routeKey: route.routeKey, + sourceSandboxId: sandboxId, + ...common, + }; + } + if (/^(ses|dind|crv)-/.test(sandboxId)) { + return { kind: 'isolated-destroy', sourceSandboxId: sandboxId, ...common }; + } + return undefined; + } + + async function settlePostExhaustionRecovery(recovery: PostExhaustionRecovery): Promise { + const latestLease = await getWrapperLease(storage); + const completedLease = reduceWrapperLease(latestLease, { + type: 'sandbox_recovery_completed', + expectedLease: recovery.expectedLease, + }); + const latestRecovery = await getSandboxRecoveryState(storage); + const settled = reduceSandboxRecoveryState(latestRecovery, { + type: 'settle_post_exhaustion', + expectedLease: recovery.expectedLease, + }); + if (completedLease === latestLease) { + if (settled) await putSandboxRecoveryState(storage, settled); + await clearSettledSandboxRecovery(storage); + return false; + } + await putWrapperLease(storage, completedLease); + if (settled) await putSandboxRecoveryState(storage, settled); + await clearSettledSandboxRecovery(storage); + await sessionMessageQueue.requestPendingDrainIfNeeded(); + return true; + } + + async function recordPostExhaustionRecoveryFailure( + recovery: PostExhaustionRecovery, + error: unknown + ): Promise { + const failedAt = Date.now(); + const errorMessage = error instanceof Error ? error.message : String(error); + const latest = await getSandboxRecoveryState(storage); + const updated = reduceSandboxRecoveryState(latest, { + type: 'record_post_exhaustion_retry', + expectedLease: recovery.expectedLease, + expectedAttempts: recovery.attempts, + nextAttemptAt: sandboxRecoveryRetryAt(failedAt, recovery.attempts), + error: errorMessage, + }); + if (updated) { + await putSandboxRecoveryState(storage, updated); + const retryAt = updated.postExhaustionRecovery?.nextAttemptAt; + if (retryAt !== undefined) await requestAlarmAtOrBefore?.(retryAt); + } + logger + .withFields({ + sessionId: getSessionIdForLogs(), + recoveryKind: recovery.kind, + sourceSandboxId: recovery.sourceSandboxId, + attempts: recovery.attempts + 1, + error: errorMessage, + logTag: 'sandbox_recovery_failed', + }) + .error('Post-exhaustion sandbox recovery failed'); + } + + async function reconcilePostExhaustionSandboxRecovery(now: number): Promise { + const lease = await getWrapperLease(storage); + let state = await getSandboxRecoveryState(storage); + let recovery = state?.postExhaustionRecovery; + + if (recovery && !doesRecoveryMatchExhaustedLease(recovery, lease)) { + const settled = reduceSandboxRecoveryState(state, { + type: 'settle_post_exhaustion', + expectedLease: recovery.expectedLease, + }); + if (settled) await putSandboxRecoveryState(storage, settled); + await clearSettledSandboxRecovery(storage); + return; + } + if (!isWrapperCleanupExhausted(lease)) return; + + if (!recovery) { + const metadata = await getMetadata(); + recovery = metadata ? createPostExhaustionRecovery(metadata, lease, now) : undefined; + if (!recovery) { + const retryAt = now + SANDBOX_RECOVERY_MAX_RETRY_DELAY_MS; + await requestAlarmAtOrBefore?.(retryAt); + logger + .withFields({ + sessionId: getSessionIdForLogs(), + sandboxId: metadata?.workspace?.sandboxId, + routeKind: metadata?.workspace?.sandboxRoute?.kind, + retryAt, + logTag: 'sandbox_recovery_metadata_unactionable', + }) + .error('Cannot classify exhausted wrapper sandbox recovery'); + return; + } + state = reduceSandboxRecoveryState(state, { + type: 'prepare_post_exhaustion', + recovery, + }); + if (!state) return; + await putSandboxRecoveryState(storage, state); + } + + if (recovery.nextAttemptAt > now) return; + const metadata = await getMetadata(); + if (!metadata || metadata.workspace?.sandboxId !== recovery.sourceSandboxId) { + if ( + recovery.kind !== 'shared-failover' || + !recovery.replacementSandboxId || + metadata?.workspace?.sandboxId !== recovery.replacementSandboxId + ) { + await recordPostExhaustionRecoveryFailure( + recovery, + new Error('Session sandbox metadata no longer matches recovery action') + ); + return; + } + } + if (!metadata) return; + + try { + if (recovery.kind === 'isolated-destroy') { + if (recovery.destroyedAt === undefined) { + await destroyIsolatedSandbox(metadata); + const latest = await getSandboxRecoveryState(storage); + const withDestroy = reduceSandboxRecoveryState(latest, { + type: 'record_post_exhaustion_destroy', + expectedLease: recovery.expectedLease, + destroyedAt: Date.now(), + }); + if (!withDestroy) return; + await putSandboxRecoveryState(storage, withDestroy); + const destroyedRecovery = withDestroy.postExhaustionRecovery; + if (destroyedRecovery?.kind !== 'isolated-destroy') return; + recovery = destroyedRecovery; + } + } else { + let replacementSandboxId = recovery.replacementSandboxId; + let routeSuffix = recovery.routeSuffix; + if (!replacementSandboxId || !routeSuffix) { + await recordSharedSandboxFailover(recovery.routeKey); + const assignment = await resolveSharedSandboxFailover(recovery.routeKey); + if (assignment.sandboxId === recovery.sourceSandboxId || !assignment.suffix) { + throw new Error('Shared sandbox failover did not resolve a replacement assignment'); + } + const latest = await getSandboxRecoveryState(storage); + const withAssignment = reduceSandboxRecoveryState(latest, { + type: 'record_post_exhaustion_assignment', + expectedLease: recovery.expectedLease, + routeKey: recovery.routeKey, + replacementSandboxId: assignment.sandboxId, + routeSuffix: assignment.suffix, + }); + if (!withAssignment) return; + await putSandboxRecoveryState(storage, withAssignment); + const assignedRecovery = withAssignment.postExhaustionRecovery; + if (assignedRecovery?.kind !== 'shared-failover') return; + recovery = assignedRecovery; + replacementSandboxId = assignment.sandboxId; + routeSuffix = assignment.suffix; + } + await persistSharedSandboxRecoveryAssignment({ + routeKey: recovery.routeKey, + sourceSandboxId: recovery.sourceSandboxId, + replacementSandboxId, + suffix: routeSuffix, + }); + } + const releasedDelivery = await settlePostExhaustionRecovery(recovery); + logger + .withFields({ + sessionId: getSessionIdForLogs(), + recoveryKind: recovery.kind, + sourceSandboxId: recovery.sourceSandboxId, + replacementSandboxId: + recovery.kind === 'shared-failover' ? recovery.replacementSandboxId : undefined, + routeKey: recovery.kind === 'shared-failover' ? recovery.routeKey : undefined, + releasedDelivery, + logTag: 'sandbox_recovery_completed', + }) + .warn('Post-exhaustion sandbox recovery completed'); + } catch (error) { + await recordPostExhaustionRecoveryFailure(recovery, error); + } + } + let sharedSandboxFailoverReconciliation: Promise | undefined; async function clearCompletedRecoveryIfWrapperAbsent(): Promise { @@ -1465,6 +1702,8 @@ export function createWrapperSupervisor( async function runMaintenance(now: number): Promise { await reconcilePhysicalCleanup(now); + await reconcilePostExhaustionSandboxRecovery(now); + const leaseAfterInitialCleanup = await getWrapperLease(storage); await reconcileSharedSandboxFailover(); await checkDisconnectGrace(now); // While the current wrapper connection is inside its disconnect grace @@ -1482,6 +1721,15 @@ export function createWrapperSupervisor( } else { await checkWrapperLiveness(now); } + const leaseAfterSupervision = await getWrapperLease(storage); + if ( + leaseAfterInitialCleanup.state !== 'stop_needed' && + leaseAfterInitialCleanup.state !== 'stopping' && + leaseAfterSupervision.state === 'stop_needed' + ) { + await reconcilePhysicalCleanup(Date.now()); + await reconcilePostExhaustionSandboxRecovery(Date.now()); + } await checkKeepWarmCleanup(now); } diff --git a/services/cloud-agent-next/test/integration/session/execute-directly-failure.test.ts b/services/cloud-agent-next/test/integration/session/execute-directly-failure.test.ts index 0c62b7405d..a5fada9845 100644 --- a/services/cloud-agent-next/test/integration/session/execute-directly-failure.test.ts +++ b/services/cloud-agent-next/test/integration/session/execute-directly-failure.test.ts @@ -13,12 +13,18 @@ import { createEventQueries } from '../../../src/session/queries/events.js'; import type { FencedWrapperDispatchRequest } from '../../../src/execution/types.js'; import { listPendingSessionMessages } from '../../../src/session/pending-messages.js'; import { + getSandboxRecoveryState, getWrapperLease, getWrapperRuntimeState, recordWrapperPong, allocateWrapperRuntimeState, recordWrapperAcceptedMessage, } from '../../../src/session/wrapper-runtime-state.js'; +import type { CallbackJob } from '../../../src/callbacks/types.js'; +import { deriveSharedSandboxId } from '../../../src/sandbox-id.js'; +import { SHARED_SANDBOX_FAILOVER_SUFFIX } from '../../../src/shared-sandbox-route.js'; +import { serializeSessionMetadata } from '../../../src/persistence/session-metadata.js'; +import type { CloudAgentSession } from '../../../src/persistence/CloudAgentSession.js'; import { getSessionMessageState, listNonTerminalAcceptedMessages, @@ -539,6 +545,512 @@ describe('new-path liveness without executionId', () => { expect(result.wrapperRuntimeState.pingDeadlineAt).toBeUndefined(); expect(result.wrapperRuntimeState.nextPingAt).toBeUndefined(); }); + + it('delivers only the queued follow-up on a fresh fence after ping-timeout recovery', async () => { + const userId = 'user_ping_recovery'; + const sessionId = 'agent_ping_recovery'; + const acceptedMessageId = 'msg_018f1e2d3c4bPingOldAbCdEfG'; + const followUpMessageId = 'msg_018f1e2d3c4bPingNewAbCdEfG'; + const doId = env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`); + const stub = env.CLOUD_AGENT_SESSION.get(doId); + + const result = await runInDurableObject(stub, async (instance, state) => { + const deliveredPlans: FencedWrapperDispatchRequest[] = []; + (instance as any).orchestrator = { + execute: async (plan: FencedWrapperDispatchRequest) => { + deliveredPlans.push(plan); + return { messageId: plan.turn.messageId, kiloSessionId: 'kilo_ping_recovery' }; + }, + }; + instance['physicalWrapperStopper'] = async () => ({ status: 'absent' }); + + await registerReadySession(instance, { + sessionId, + userId, + orgId: 'org_ping_recovery', + kiloSessionId: '12121212-1212-4121-8121-121212121212', + prompt: 'initial prompt', + mode: 'code', + model: 'test-model', + kilocodeToken: 'token-ping-recovery', + }); + + const { state: originalWrapper } = await allocateWrapperRuntimeState(instance.ctx.storage); + const originalRunId = originalWrapper.wrapperRunId!; + const originalConnectionId = originalWrapper.wrapperConnectionId!; + await instance.ctx.storage.put('wrapper_lease', { + state: 'owns_wrapper', + nextInstanceGeneration: 2, + instance: { instanceId: 'instance_ping_recovery', instanceGeneration: 1 }, + }); + await putSessionMessageState(instance.ctx.storage, { + messageId: acceptedMessageId, + status: 'accepted', + prompt: 'work accepted by the old wrapper', + createdAt: Date.now(), + acceptedAt: Date.now(), + wrapperRunId: originalRunId, + }); + + const admission = await instance.admitSubmittedMessage( + queueUserMessageInput({ + userId, + prompt: 'queued follow-up', + messageId: followUpMessageId, + }) + ); + expect(admission.success).toBe(true); + + const expiredAt = Date.now() - 1; + await instance.ctx.storage.put('wrapper_runtime_state', { + wrapperGeneration: originalWrapper.wrapperGeneration, + wrapperConnectionId: originalConnectionId, + wrapperRunId: originalRunId, + pingDeadlineAt: expiredAt, + lastHeartbeatUpdate: expiredAt - 10 * 60_000, + }); + + await instance.alarm(); + + const db = drizzle(state.storage, { logger: false }); + const allEvents = createEventQueries(db, state.storage.sql).findByFilters({}); + const supervisor = instance['getWrapperSupervisor'](); + return { + deliveredPlans, + acceptedMessage: await getSessionMessageState(instance.ctx.storage, acceptedMessageId), + followUpMessage: await getSessionMessageState(instance.ctx.storage, followUpMessageId), + pendingMessages: await listPendingSessionMessages(instance.ctx.storage), + wrapperRuntimeState: await getWrapperRuntimeState(instance.ctx.storage), + wrapperLease: await getWrapperLease(instance.ctx.storage), + failedEvents: allEvents.filter(event => event.stream_event_type === 'cloud.message.failed'), + staleReconnect: await supervisor.checkReconnect({ + wrapperRunId: originalRunId, + wrapperGeneration: originalWrapper.wrapperGeneration, + wrapperConnectionId: originalConnectionId, + }), + originalRunId, + originalGeneration: originalWrapper.wrapperGeneration, + }; + }); + + expect(result.acceptedMessage).toMatchObject({ + status: 'failed', + failureCode: 'wrapper_ping_timeout', + }); + expect(result.followUpMessage).toMatchObject({ + status: 'accepted', + }); + expect(result.pendingMessages).toHaveLength(0); + expect(result.deliveredPlans).toHaveLength(1); + expect(result.deliveredPlans[0].turn.messageId).toBe(followUpMessageId); + expect(result.deliveredPlans[0].wrapper.fence.wrapperRunId).not.toBe(result.originalRunId); + expect(result.deliveredPlans[0].wrapper.fence.wrapperGeneration).toBeGreaterThan( + result.originalGeneration + ); + expect(result.wrapperRuntimeState.wrapperRunId).toBe( + result.deliveredPlans[0].wrapper.fence.wrapperRunId + ); + expect(result.wrapperLease).toMatchObject({ state: 'owns_wrapper' }); + expect(result.failedEvents).toHaveLength(1); + expect(result.staleReconnect).toEqual({ accepted: false, reason: 'stale-wrapper-run' }); + }); + + it('defers terminal effects while physical cleanup still fences queued recovery', async () => { + const userId = 'user_ping_cleanup_hold'; + const sessionId = 'agent_ping_cleanup_hold'; + const acceptedMessageId = 'msg_018f1e2d3c4bPingHoldOldAbC'; + const followUpMessageId = 'msg_018f1e2d3c4bPingHoldNewAbC'; + const doId = env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`); + const stub = env.CLOUD_AGENT_SESSION.get(doId); + + const result = await runInDurableObject(stub, async (instance, state) => { + await registerReadySession(instance, { + sessionId, + userId, + orgId: 'org_ping_cleanup_hold', + kiloSessionId: '34343434-3434-4343-8343-343434343434', + prompt: 'initial prompt', + mode: 'code', + model: 'test-model', + kilocodeToken: 'token-ping-cleanup-hold', + }); + + const { state: originalWrapper } = await allocateWrapperRuntimeState(instance.ctx.storage); + await instance.ctx.storage.put('wrapper_lease', { + state: 'owns_wrapper', + nextInstanceGeneration: 2, + instance: { instanceId: 'instance_ping_cleanup_hold', instanceGeneration: 1 }, + }); + instance['physicalWrapperStopper'] = async () => ({ + status: 'still-present', + observed: [], + }); + await putSessionMessageState(instance.ctx.storage, { + messageId: acceptedMessageId, + status: 'accepted', + prompt: 'work accepted by the old wrapper', + createdAt: Date.now(), + acceptedAt: Date.now(), + wrapperRunId: originalWrapper.wrapperRunId!, + }); + const admission = await instance.admitSubmittedMessage( + queueUserMessageInput({ + userId, + prompt: 'queued follow-up', + messageId: followUpMessageId, + }) + ); + expect(admission.success).toBe(true); + + const expiredAt = Date.now() - 1; + await instance.ctx.storage.put('wrapper_runtime_state', { + wrapperGeneration: originalWrapper.wrapperGeneration, + wrapperConnectionId: originalWrapper.wrapperConnectionId!, + wrapperRunId: originalWrapper.wrapperRunId!, + pingDeadlineAt: expiredAt, + lastHeartbeatUpdate: expiredAt - 10 * 60_000, + }); + + await instance.alarm(); + + const db = drizzle(state.storage, { logger: false }); + const allEvents = createEventQueries(db, state.storage.sql).findByFilters({}); + return { + acceptedMessage: await getSessionMessageState(instance.ctx.storage, acceptedMessageId), + pendingMessages: await listPendingSessionMessages(instance.ctx.storage), + wrapperLease: await getWrapperLease(instance.ctx.storage), + failedEvents: allEvents.filter(event => event.stream_event_type === 'cloud.message.failed'), + }; + }); + + expect(result.acceptedMessage).toMatchObject({ + status: 'failed', + failureCode: 'wrapper_ping_timeout', + terminalEffects: { + event: 'pending', + push: { disposition: 'pending' }, + }, + }); + expect(result.pendingMessages).toEqual([ + expect.objectContaining({ messageId: followUpMessageId }), + ]); + expect(result.wrapperLease).toMatchObject({ + state: 'stop_needed', + reason: 'unhealthy-wrapper', + attempts: 1, + }); + expect(result.failedEvents).toHaveLength(0); + }); + + it('repairs terminal effects during exhausted cleanup and resumes queued work after isolated recovery', async () => { + const userId = 'user_exhausted_isolated_recovery'; + const sessionId = 'agent_exhausted_isolated_recovery'; + const acceptedMessageId = 'msg_018f1e2d3c4bExhaustOldAbCd'; + const unrelatedCallbackMessageId = 'msg_018f1e2d3c4bExhaustCbAbCdE'; + const followUpMessageId = 'msg_018f1e2d3c4bExhaustNewAbCd'; + const doId = env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`); + const stub = env.CLOUD_AGENT_SESSION.get(doId); + + const result = await runInDurableObject(stub, async (instance, state) => { + const callbackJobs: CallbackJob[] = []; + ( + instance as unknown as { + env: { CALLBACK_QUEUE: { send(job: CallbackJob): Promise } }; + } + ).env.CALLBACK_QUEUE = { + send: async job => { + callbackJobs.push(job); + }, + }; + const deliveredPlans: FencedWrapperDispatchRequest[] = []; + (instance as any).orchestrator = { + execute: async (plan: FencedWrapperDispatchRequest) => { + deliveredPlans.push(plan); + return { messageId: plan.turn.messageId, kiloSessionId: 'kilo_exhausted_recovery' }; + }, + }; + let destroyAttempts = 0; + instance['postExhaustionSandboxDestroyer'] = async () => { + destroyAttempts += 1; + if (destroyAttempts === 1) throw new Error('sandbox delete temporarily unavailable'); + }; + + const callbackTarget = { url: 'https://example.com/exhausted-recovery' }; + await registerReadySession(instance, { + sessionId, + userId, + orgId: 'org_exhausted_isolated_recovery', + kiloSessionId: '56565656-5656-4656-8656-565656565656', + prompt: 'initial prompt', + mode: 'code', + model: 'test-model', + kilocodeToken: 'token-exhausted-recovery', + sandboxId: 'ses-abcdef', + callbackTarget, + }); + + await instance.ctx.storage.put('wrapper_runtime_state', { wrapperGeneration: 4 }); + await instance.ctx.storage.put('wrapper_lease', { + state: 'stop_needed', + nextInstanceGeneration: 5, + target: { kind: 'session' }, + reason: 'unhealthy-wrapper', + requestedAt: 100, + nextAttemptAt: 3_153_600_000_100, + attempts: 5, + lastError: 'wrapper remains present', + exhaustedAt: 100, + }); + await putSessionMessageState(instance.ctx.storage, { + messageId: acceptedMessageId, + status: 'accepted', + prompt: 'accepted by the failed wrapper', + createdAt: 1, + acceptedAt: 2, + wrapperRunId: 'wr_exhausted_old', + callbackRequired: true, + callbackTarget, + }); + await instance['getMessageSettlementOutbox']().persistTerminalTransition( + acceptedMessageId, + { + kind: 'failed', + reason: 'wrapper_failure', + error: 'Wrapper liveness ping timed out', + completionSource: 'wrapper_failure', + failureStage: 'post_dispatch_no_activity', + failureCode: 'wrapper_ping_timeout', + }, + { allowIdleBatchWithoutObservedIdle: true } + ); + await putSessionMessageState(instance.ctx.storage, { + messageId: unrelatedCallbackMessageId, + status: 'failed', + prompt: 'unrelated callback retry', + createdAt: 1, + terminalAt: 2, + failureReason: 'delivery_failure', + completionSource: 'delivery_failure', + callbackRequired: true, + callbackTarget, + callbackAttempts: 1, + callbackRetryAt: 0, + terminalEffects: { + event: 'accounted', + callback: { disposition: 'accounted', allowWithoutObservedIdle: true }, + push: { disposition: 'accounted' }, + }, + }); + await instance.ctx.storage.put('idle_batch_callback:unrelated-retry', { + batchId: 'unrelated-retry', + createdAt: 1, + updatedAt: 2, + representativeMessageId: unrelatedCallbackMessageId, + finalizedAt: 2, + }); + const admission = await instance.admitSubmittedMessage( + queueUserMessageInput({ + userId, + prompt: 'queued follow-up', + messageId: followUpMessageId, + }) + ); + expect(admission.success).toBe(true); + + await instance.alarm(); + + const db = drizzle(state.storage, { logger: false }); + const eventsAfterFailure = createEventQueries(db, state.storage.sql).findByFilters({}); + const messageAfterFailure = await getSessionMessageState( + instance.ctx.storage, + acceptedMessageId + ); + const pendingAfterFailure = await listPendingSessionMessages(instance.ctx.storage); + const recoveryAfterFailure = await getSandboxRecoveryState(instance.ctx.storage); + expect(pendingAfterFailure).toEqual([ + expect.objectContaining({ messageId: followUpMessageId }), + ]); + expect(messageAfterFailure?.terminalEffects).toMatchObject({ + event: 'accounted', + callback: { disposition: 'accounted' }, + }); + expect( + eventsAfterFailure.filter(event => event.stream_event_type === 'cloud.message.failed') + ).toHaveLength(1); + expect(callbackJobs.map(job => job.payload.messageId).sort()).toEqual( + [acceptedMessageId, unrelatedCallbackMessageId].sort() + ); + expect(recoveryAfterFailure?.postExhaustionRecovery).toMatchObject({ + kind: 'isolated-destroy', + attempts: 1, + lastError: 'sandbox delete temporarily unavailable', + }); + + if (!recoveryAfterFailure?.postExhaustionRecovery) { + throw new Error('Expected persisted sandbox recovery'); + } + await instance.ctx.storage.put('sandbox_recovery_state', { + ...recoveryAfterFailure, + postExhaustionRecovery: { + ...recoveryAfterFailure.postExhaustionRecovery, + nextAttemptAt: 0, + }, + }); + await instance.alarm(); + + return { + callbackJobs, + deliveredPlans, + destroyAttempts, + acceptedMessage: await getSessionMessageState(instance.ctx.storage, acceptedMessageId), + followUpMessage: await getSessionMessageState(instance.ctx.storage, followUpMessageId), + pendingMessages: await listPendingSessionMessages(instance.ctx.storage), + wrapperLease: await getWrapperLease(instance.ctx.storage), + events: createEventQueries(db, state.storage.sql).findByFilters({}), + }; + }); + + expect(result.destroyAttempts).toBe(2); + expect(result.acceptedMessage).toMatchObject({ + status: 'failed', + wrapperRunId: 'wr_exhausted_old', + }); + expect(result.followUpMessage).toMatchObject({ status: 'accepted' }); + expect(result.pendingMessages).toHaveLength(0); + expect(result.deliveredPlans).toHaveLength(1); + expect(result.deliveredPlans[0].turn.messageId).toBe(followUpMessageId); + expect(result.deliveredPlans[0].wrapper.fence.wrapperGeneration).toBeGreaterThan(4); + expect(result.wrapperLease).toMatchObject({ state: 'owns_wrapper' }); + expect(result.callbackJobs.map(job => job.payload.messageId).sort()).toEqual( + [acceptedMessageId, unrelatedCallbackMessageId].sort() + ); + expect( + result.events.filter(event => event.stream_event_type === 'cloud.message.failed') + ).toHaveLength(1); + }); + + it('moves only the exhausted shared session to failover before queued delivery resumes', async () => { + const routeKey = `usr-${'9'.repeat(48)}` as const; + const replacementSandboxId = await deriveSharedSandboxId( + routeKey, + SHARED_SANDBOX_FAILOVER_SUFFIX + ); + const affectedUserId = 'user_exhausted_shared_recovery'; + const affectedSessionId = 'agent_exhausted_shared_recovery'; + const otherUserId = 'user_exhausted_shared_other'; + const otherSessionId = 'agent_exhausted_shared_other'; + const followUpMessageId = 'msg_018f1e2d3c4bSharedNewAbCdE'; + const affectedStub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`${affectedUserId}:${affectedSessionId}`) + ); + const otherStub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`${otherUserId}:${otherSessionId}`) + ); + + const installSharedMetadata = async (instance: CloudAgentSession) => { + const metadata = await instance.getMetadata(); + if (!metadata?.workspace) throw new Error('Expected ready workspace metadata'); + await instance.ctx.storage.put( + 'metadata', + serializeSessionMetadata({ + ...metadata, + workspace: { + ...metadata.workspace, + sandboxId: routeKey, + sandboxRoute: { kind: 'shared', routeKey }, + }, + }) + ); + }; + + await runInDurableObject(otherStub, async instance => { + await registerReadySession(instance, { + sessionId: otherSessionId, + userId: otherUserId, + kiloSessionId: '67676767-6767-4676-8676-676767676767', + prompt: 'other session', + mode: 'code', + model: 'test-model', + sandboxId: routeKey, + }); + await installSharedMetadata(instance); + }); + + const result = await runInDurableObject(affectedStub, async instance => { + const deliveredPlans: FencedWrapperDispatchRequest[] = []; + (instance as any).orchestrator = { + execute: async (plan: FencedWrapperDispatchRequest) => { + deliveredPlans.push(plan); + return { messageId: plan.turn.messageId, kiloSessionId: 'kilo_shared_recovery' }; + }, + }; + let destroyCalls = 0; + instance['postExhaustionSandboxDestroyer'] = async () => { + destroyCalls += 1; + }; + await registerReadySession(instance, { + sessionId: affectedSessionId, + userId: affectedUserId, + kiloSessionId: '78787878-7878-4787-8787-787878787878', + prompt: 'affected session', + mode: 'code', + model: 'test-model', + sandboxId: routeKey, + }); + await installSharedMetadata(instance); + await instance.ctx.storage.put('wrapper_runtime_state', { wrapperGeneration: 6 }); + await instance.ctx.storage.put('wrapper_lease', { + state: 'stop_needed', + nextInstanceGeneration: 7, + target: { kind: 'session' }, + reason: 'unhealthy-wrapper', + requestedAt: 200, + nextAttemptAt: 3_153_600_000_200, + attempts: 5, + lastError: 'wrapper remains present', + exhaustedAt: 200, + }); + const admission = await instance.admitSubmittedMessage( + queueUserMessageInput({ + userId: affectedUserId, + prompt: 'queued for shared failover', + messageId: followUpMessageId, + }) + ); + expect(admission.success).toBe(true); + + await instance.alarm(); + + return { + deliveredPlans, + destroyCalls, + metadata: await instance.getMetadata(), + pendingMessages: await listPendingSessionMessages(instance.ctx.storage), + wrapperLease: await getWrapperLease(instance.ctx.storage), + }; + }); + const otherMetadata = await runInDurableObject(otherStub, instance => instance.getMetadata()); + + expect(result.destroyCalls).toBe(0); + expect(result.metadata?.workspace).toMatchObject({ + sandboxId: replacementSandboxId, + sandboxRoute: { + kind: 'shared', + routeKey, + suffix: SHARED_SANDBOX_FAILOVER_SUFFIX, + }, + }); + expect(result.pendingMessages).toHaveLength(0); + expect(result.deliveredPlans).toHaveLength(1); + expect(result.deliveredPlans[0].turn.messageId).toBe(followUpMessageId); + expect(result.deliveredPlans[0].workspace.sandboxId).toBe(replacementSandboxId); + expect(result.deliveredPlans[0].wrapper.fence.wrapperGeneration).toBeGreaterThan(6); + expect(result.wrapperLease).toMatchObject({ state: 'owns_wrapper' }); + expect(otherMetadata?.workspace).toMatchObject({ + sandboxId: routeKey, + sandboxRoute: { kind: 'shared', routeKey }, + }); + }); }); describe('hot delivery failure preserves existing wrapper identity', () => { diff --git a/services/cloud-agent-next/test/integration/session/start-execution-v2.test.ts b/services/cloud-agent-next/test/integration/session/start-execution-v2.test.ts index 2cdd43ae56..53efb85ea2 100644 --- a/services/cloud-agent-next/test/integration/session/start-execution-v2.test.ts +++ b/services/cloud-agent-next/test/integration/session/start-execution-v2.test.ts @@ -10,6 +10,8 @@ import { storePendingSessionMessage, } from '../../../src/session/pending-messages.js'; import { listNonTerminalAcceptedMessages } from '../../../src/session/session-message-state.js'; +import { deriveSharedSandboxId } from '../../../src/sandbox-id.js'; +import { SHARED_SANDBOX_FAILOVER_SUFFIX } from '../../../src/shared-sandbox-route.js'; import { groupedRegisterSessionInput, @@ -520,6 +522,77 @@ describe('CloudAgentSession message admission', () => { expect(result.metadata?.workspace).toEqual(input.workspace); }); + it('replays registration against the durably installed shared failover assignment', async () => { + const userId = 'user_grouped_route_migrated' as const; + const sessionId = 'agent_grouped_route_migrated' as const; + const messageId = 'msg_018f1e2d3c4bRouteMigrAbCdE'; + const routeKey = `usr-${'8'.repeat(48)}` as const; + const replacementSandboxId = await deriveSharedSandboxId( + routeKey, + SHARED_SANDBOX_FAILOVER_SUFFIX + ); + const stub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`) + ); + const base = groupedRegisterSessionInput({ + sessionId, + userId, + prompt: 'replay migrated route', + mode: 'code', + model: 'test-model', + kiloSessionId: '89898989-8989-4898-8989-898989898989', + }); + const original = { + ...base, + workspace: { + sandboxId: routeKey, + sandboxProvider: 'cloudflare' as const, + sandboxRoute: { kind: 'shared' as const, routeKey }, + }, + message: { + initialTurn: { + type: 'prompt' as const, + messageId, + prompt: 'replay migrated route', + }, + }, + }; + + const result = await runInDurableObject(stub, async instance => { + const first = await instance.createSessionWithInitialAdmission(original); + await instance['persistSharedSandboxRecoveryAssignment']({ + routeKey, + sourceSandboxId: routeKey, + replacementSandboxId, + suffix: SHARED_SANDBOX_FAILOVER_SUFFIX, + }); + const replay = await instance.createSessionWithInitialAdmission({ + ...original, + workspace: { + ...original.workspace, + sandboxId: replacementSandboxId, + sandboxRoute: { + kind: 'shared', + routeKey, + suffix: SHARED_SANDBOX_FAILOVER_SUFFIX, + }, + }, + }); + return { first, replay, metadata: await instance.getMetadata() }; + }); + + expect(result.first).toMatchObject({ success: true, messageId }); + expect(result.replay).toEqual(result.first); + expect(result.metadata?.workspace).toMatchObject({ + sandboxId: replacementSandboxId, + sandboxRoute: { + kind: 'shared', + routeKey, + suffix: SHARED_SANDBOX_FAILOVER_SUFFIX, + }, + }); + }); + it('rejects readiness updates that change a shared route to another sandbox', async () => { const userId = 'user_shared_route_ready_mismatch' as const; const sessionId = 'agent_shared_route_ready_mismatch' as const;