From 103accc07ae1a66ce861cb3c846c1c81291939e5 Mon Sep 17 00:00:00 2001 From: Evgeny Shurakov Date: Mon, 13 Jul 2026 15:13:34 +0200 Subject: [PATCH 1/4] CloudAgentNext - broadcast cloud.status ready after terminalization and broaden wrapper terminal error detection --- packages/db/src/schema.ts | 1 + .../worker-utils/src/cloud-agent-failure.ts | 1 + .../src/persistence/CloudAgentSession.ts | 22 +++- .../src/session/safe-failure-projection.ts | 1 + .../src/session/terminal-error-projector.ts | 1 + .../src/session/wrapper-supervisor.test.ts | 85 +++++++++++++ .../src/session/wrapper-supervisor.ts | 13 ++ .../cloud-agent-next/src/shared/protocol.ts | 6 +- .../session/disconnect-and-reaper.test.ts | 117 ++++++++++++++++++ .../test/unit/wrapper/reconnection.test.ts | 111 +++++++++++++++++ .../wrapper/src/connection.ts | 21 +++- 11 files changed, 376 insertions(+), 3 deletions(-) diff --git a/packages/db/src/schema.ts b/packages/db/src/schema.ts index c84b62b321..080f600751 100644 --- a/packages/db/src/schema.ts +++ b/packages/db/src/schema.ts @@ -5533,6 +5533,7 @@ export type CloudAgentSessionRunFailureCode = | 'wrapper_error_after_activity' | 'missing_assistant_reply' | 'payment_required' + | 'provider_error' | 'user_interrupt' | 'container_shutdown' | 'system_interrupt' diff --git a/packages/worker-utils/src/cloud-agent-failure.ts b/packages/worker-utils/src/cloud-agent-failure.ts index 23e08ffc1d..f801249bf5 100644 --- a/packages/worker-utils/src/cloud-agent-failure.ts +++ b/packages/worker-utils/src/cloud-agent-failure.ts @@ -28,6 +28,7 @@ export const CLOUD_AGENT_FAILURE_CODES = [ 'wrapper_error_after_activity', 'missing_assistant_reply', 'payment_required', + 'provider_error', 'user_interrupt', 'container_shutdown', 'system_interrupt', diff --git a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts index 001c1eefb2..267bd37af4 100644 --- a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts +++ b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts @@ -715,6 +715,15 @@ export class CloudAgentSession extends DurableObject { : recordSharedSandboxFailover(this.env.SHARED_SANDBOX_OVERRIDES, routeKey), requestAlarmAtOrBefore: deadline => this.scheduleAlarmAtOrBefore(deadline), getSessionIdForLogs: () => this.sessionId, + onSessionTerminalized: async () => { + this.broadcastVolatileEvent({ + executionId: '' as EventSourceId, + sessionId: this.sessionId ?? '', + streamEventType: 'cloud.status', + payload: JSON.stringify({ cloudStatus: { type: 'ready' } }), + timestamp: Date.now(), + }); + }, }); } @@ -3259,8 +3268,19 @@ export class CloudAgentSession extends DurableObject { } async handleWrapperTerminalEvent(params: WrapperTerminalEvent): Promise { - await this.resolveSessionId(); + const sessionId = await this.resolveSessionId(); await this.getWrapperSupervisor().onTerminalEvent(params); + + if (params.status === 'failed' || params.status === 'interrupted') { + this.broadcastVolatileEvent({ + executionId: '' as EventSourceId, + sessionId: sessionId ?? '', + streamEventType: 'cloud.status', + payload: JSON.stringify({ cloudStatus: { type: 'ready' } }), + timestamp: Date.now(), + }); + } + const metadata = await this.getMetadata(); if (!metadata) return; if (!isCodeReviewEphemeralSandboxId(metadata.workspace?.sandboxId)) return; diff --git a/services/cloud-agent-next/src/session/safe-failure-projection.ts b/services/cloud-agent-next/src/session/safe-failure-projection.ts index 52e604b295..187ddd1440 100644 --- a/services/cloud-agent-next/src/session/safe-failure-projection.ts +++ b/services/cloud-agent-next/src/session/safe-failure-projection.ts @@ -39,6 +39,7 @@ const GENERIC_FAILURE_MESSAGES = { wrapper_error_after_activity: 'Agent wrapper failed while processing the message', missing_assistant_reply: 'No assistant reply was produced', payment_required: 'Assistant request failed: insufficient credits', + provider_error: 'Assistant provider returned an invalid response', user_interrupt: 'The message was interrupted by the user', container_shutdown: 'The agent container shut down', system_interrupt: 'The message was interrupted', diff --git a/services/cloud-agent-next/src/session/terminal-error-projector.ts b/services/cloud-agent-next/src/session/terminal-error-projector.ts index 7d0b0c9773..f87c5f0a9f 100644 --- a/services/cloud-agent-next/src/session/terminal-error-projector.ts +++ b/services/cloud-agent-next/src/session/terminal-error-projector.ts @@ -9,6 +9,7 @@ const NON_RETRYABLE_FAILURE_CODES = new Set([ 'session_metadata_missing', 'model_missing', 'payment_required', + 'provider_error', 'user_interrupt', ]); 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..ebb3e2fe25 100644 --- a/services/cloud-agent-next/src/session/wrapper-supervisor.test.ts +++ b/services/cloud-agent-next/src/session/wrapper-supervisor.test.ts @@ -122,6 +122,7 @@ function createHarness( wrapperRunId: string ) => Promise; recordSharedSandboxFailover?: (routeKey: SandboxId) => Promise; + onSessionTerminalized?: () => Promise; } ) { const getAssistantMessageForUserMessage = @@ -158,6 +159,7 @@ function createHarness( const recordSharedSandboxFailover = vi.fn( options?.recordSharedSandboxFailover ?? (async () => {}) ); + const onSessionTerminalized = vi.fn(options?.onSessionTerminalized ?? (async () => {})); const supervisor = createWrapperSupervisor({ storage, agentRuntime: { @@ -179,6 +181,7 @@ function createHarness( requestedAlarms.push(deadline); }, getSessionIdForLogs: () => currentMetadata.identity.sessionId, + onSessionTerminalized, }); return { @@ -191,6 +194,7 @@ function createHarness( requestedAlarms, requestPendingDrainIfNeeded, recordSharedSandboxFailover, + onSessionTerminalized, settlementOutbox, supervisor, }; @@ -1099,6 +1103,87 @@ describe('WrapperSupervisor', () => { }); }); + it('calls onSessionTerminalized after disconnect grace terminalizes accepted work', async () => { + const harness = createHarness([ + liveRuntimeState({ pingDeadlineAt: 92_000, noOutputDeadlineAt: 332_000 }), + OWNED_WRAPPER_LEASE, + disconnectGraceForCurrentConnection(90_000), + ]); + await putSessionMessageState(harness.storage, acceptedMessage()); + + await harness.supervisor.runMaintenance(100_001); + + await expect(getSessionMessageState(harness.storage, MESSAGE_ID)).resolves.toMatchObject({ + status: 'failed', + failureCode: 'wrapper_disconnected', + }); + expect(harness.onSessionTerminalized).toHaveBeenCalledTimes(1); + }); + + it('calls onSessionTerminalized after liveness expiry terminalizes accepted work', async () => { + const acceptedAt = 2_000; + const noOutputDeadlineAt = acceptedAt + WRAPPER_NO_OUTPUT_TIMEOUT_MS; + const harness = createHarness([ + liveRuntimeState({ noOutputDeadlineAt, nextPingAt: noOutputDeadlineAt + 1 }), + OWNED_WRAPPER_LEASE, + ]); + await putSessionMessageState(harness.storage, acceptedMessage()); + + await harness.supervisor.runMaintenance(noOutputDeadlineAt); + + await expect(getSessionMessageState(harness.storage, MESSAGE_ID)).resolves.toMatchObject({ + status: 'failed', + failureCode: 'wrapper_no_output', + }); + expect(harness.onSessionTerminalized).toHaveBeenCalledTimes(1); + }); + + it('does not call onSessionTerminalized when disconnect grace expires for a stale wrapper run', async () => { + const harness = createHarness([ + liveRuntimeState({ + wrapperRunId: 'wr_newer_current', + wrapperGeneration: 5, + wrapperConnectionId: 'conn_newer', + }), + OWNED_WRAPPER_LEASE, + disconnectGraceForCurrentConnection(90_000, { + wrapperRunId: 'wr_stale_run', + wrapperGeneration: 4, + wrapperConnectionId: WRAPPER_CONNECTION_ID, + }), + ]); + await putSessionMessageState(harness.storage, { + ...acceptedMessage(), + wrapperRunId: 'wr_newer_current', + }); + + await harness.supervisor.runMaintenance(100_001); + + await expect(harness.storage.get('disconnect_grace')).resolves.toBeUndefined(); + expect(harness.onSessionTerminalized).not.toHaveBeenCalled(); + }); + + it('does not call onSessionTerminalized when liveness expiry has no accepted messages to terminalize', async () => { + const harness = createHarness([ + liveRuntimeState({ + finalizingWrapperRunId: WRAPPER_RUN_ID, + noOutputDeadlineAt: 9_000, + nextPingAt: 30_000, + }), + OWNED_WRAPPER_LEASE, + ]); + await putSessionMessageState(harness.storage, { + ...acceptedMessage(), + status: 'completed', + terminalAt: 3_000, + completionSource: 'assistant_message_event', + }); + + await harness.supervisor.runMaintenance(10_000); + + expect(harness.onSessionTerminalized).not.toHaveBeenCalled(); + }); + it('includes the disconnect grace expiry deadline even when the ping deadline is earlier', async () => { const harness = createHarness([ liveRuntimeState({ pingDeadlineAt: 92_000, noOutputDeadlineAt: 332_000 }), diff --git a/services/cloud-agent-next/src/session/wrapper-supervisor.ts b/services/cloud-agent-next/src/session/wrapper-supervisor.ts index da9a7895ea..cdfad6d8d2 100644 --- a/services/cloud-agent-next/src/session/wrapper-supervisor.ts +++ b/services/cloud-agent-next/src/session/wrapper-supervisor.ts @@ -178,6 +178,13 @@ export type WrapperSupervisorDependencies = { recordSharedSandboxFailover: (routeKey: SandboxId) => Promise; requestAlarmAtOrBefore?: (deadline: number) => Promise; getSessionIdForLogs: () => string | undefined; + /** + * Invoked after a disconnect-grace or liveness path terminalizes accepted + * work, so the DO can broadcast a `cloud.status: ready` event and re-enable + * client input. Not invoked for the normal terminal-event path, which + * broadcasts readiness from {@link handleWrapperTerminalEvent} directly. + */ + onSessionTerminalized?: () => Promise; }; function matchesDisconnectGraceFence( @@ -362,6 +369,7 @@ export function createWrapperSupervisor( recordSharedSandboxFailover, requestAlarmAtOrBefore, getSessionIdForLogs, + onSessionTerminalized, } = dependencies; async function readDisconnectGrace(): Promise { @@ -658,6 +666,10 @@ export function createWrapperSupervisor( state.wrapperConnectionId ); } + + if (acceptedMessages.length > 0) { + await onSessionTerminalized?.(); + } } async function checkDisconnectGrace(now: number): Promise { @@ -730,6 +742,7 @@ export function createWrapperSupervisor( { incrementGeneration: true } ); await releaseWrapperTerminalWaitForIdleBatchForWrapperRun(wrapperRunId); + await onSessionTerminalized?.(); } async function hasActiveWrapperWork(state: WrapperRuntimeState): Promise { diff --git a/services/cloud-agent-next/src/shared/protocol.ts b/services/cloud-agent-next/src/shared/protocol.ts index 0d28d29aae..1f2e8f6155 100644 --- a/services/cloud-agent-next/src/shared/protocol.ts +++ b/services/cloud-agent-next/src/shared/protocol.ts @@ -53,7 +53,11 @@ export type IngestEvent = { data: unknown; }; -export const WrapperTerminalFailureCodes = ['payment_required', 'model_missing'] as const; +export const WrapperTerminalFailureCodes = [ + 'payment_required', + 'model_missing', + 'provider_error', +] as const; export type WrapperTerminalFailureCode = (typeof WrapperTerminalFailureCodes)[number]; /** diff --git a/services/cloud-agent-next/test/integration/session/disconnect-and-reaper.test.ts b/services/cloud-agent-next/test/integration/session/disconnect-and-reaper.test.ts index c58fdabeaf..920b888420 100644 --- a/services/cloud-agent-next/test/integration/session/disconnect-and-reaper.test.ts +++ b/services/cloud-agent-next/test/integration/session/disconnect-and-reaper.test.ts @@ -9,6 +9,7 @@ import { getWrapperRuntimeState, } from '../../../src/session/wrapper-runtime-state.js'; import type { ExecutionId } from '../../../src/types/ids.js'; +import { registerReadySession } from '../../helpers/session-setup.js'; describe('Disconnect handling and compatibility execution RPCs', () => { it('alarm schedules the idle cadence when no current message deadlines exist', async () => { @@ -251,4 +252,120 @@ describe('Disconnect handling and compatibility execution RPCs', () => { true ); }); + + it('broadcasts cloud.status: ready after a wrapper disconnect and grace expiry', async () => { + const userId = 'user_disconnect_ready'; + const sessionId = 'agent_disconnect_ready'; + const stub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`) + ); + + const result = await runInDurableObject(stub, async (instance, state) => { + await registerReadySession(instance, { + sessionId, + userId, + kiloSessionId: 'ses_disconnect_ready', + prompt: 'disconnect ready test', + mode: 'code', + model: 'test-model', + }); + + const now = Date.now(); + await state.storage.put('wrapper_runtime_state', { + wrapperRunId: 'wr_disconnect_ready', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_disconnect_ready', + }); + await putSessionMessageState(state.storage, { + messageId: 'msg_018f1e2d3c4bDiscReadyABCDE', + status: 'accepted', + prompt: 'disconnect ready test', + createdAt: now, + acceptedAt: now, + wrapperRunId: 'wr_disconnect_ready', + }); + await state.storage.put('disconnect_grace', { + wrapperRunId: 'wr_disconnect_ready', + disconnectedAt: now - 20_000, + wsCloseCode: 1006, + wsCloseReason: 'socket closed', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_disconnect_ready', + }); + + const broadcasted: Array<{ stream_event_type: string; payload: string }> = []; + instance.broadcastEvent = event => { + broadcasted.push({ + stream_event_type: event.stream_event_type, + payload: event.payload, + }); + }; + + await instance.alarm(); + return { broadcasted }; + }); + + const cloudStatusReady = result.broadcasted.filter(event => { + if (event.stream_event_type !== 'cloud.status') return false; + const payload = JSON.parse(event.payload) as { cloudStatus?: { type?: string } }; + return payload.cloudStatus?.type === 'ready'; + }); + expect(cloudStatusReady).toHaveLength(1); + }); + + it('broadcasts cloud.status: ready after a failed wrapper terminal event', async () => { + const userId = 'user_terminal_ready'; + const sessionId = 'agent_terminal_ready'; + const stub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`) + ); + + const result = await runInDurableObject(stub, async (instance, state) => { + await registerReadySession(instance, { + sessionId, + userId, + kiloSessionId: 'ses_terminal_ready', + prompt: 'terminal ready test', + mode: 'code', + model: 'test-model', + }); + + const now = Date.now(); + await state.storage.put('wrapper_runtime_state', { + wrapperRunId: 'wr_terminal_ready', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_terminal_ready', + }); + await putSessionMessageState(state.storage, { + messageId: 'msg_018f1e2d3c4bTermReadyABCDE', + status: 'accepted', + prompt: 'terminal ready test', + createdAt: now, + acceptedAt: now, + wrapperRunId: 'wr_terminal_ready', + }); + + const broadcasted: Array<{ stream_event_type: string; payload: string }> = []; + instance.broadcastEvent = event => { + broadcasted.push({ + stream_event_type: event.stream_event_type, + payload: event.payload, + }); + }; + + await instance.handleWrapperTerminalEvent({ + wrapperRunId: 'wr_terminal_ready', + status: 'failed', + error: 'Assistant request failed', + }); + return { broadcasted }; + }); + + const cloudStatusReady = result.broadcasted.filter(event => { + if (event.stream_event_type !== 'cloud.status') return false; + const payload = JSON.parse(event.payload) as { cloudStatus?: { type?: string } }; + return payload.cloudStatus?.type === 'ready'; + }); + expect(cloudStatusReady).toHaveLength(1); + }); }); diff --git a/services/cloud-agent-next/test/unit/wrapper/reconnection.test.ts b/services/cloud-agent-next/test/unit/wrapper/reconnection.test.ts index a3e815648e..e50120b444 100644 --- a/services/cloud-agent-next/test/unit/wrapper/reconnection.test.ts +++ b/services/cloud-agent-next/test/unit/wrapper/reconnection.test.ts @@ -1406,6 +1406,117 @@ describe('ingest WS reconnection', () => { expect(callbacks.onTerminalError).not.toHaveBeenCalled(); }); + it('classifies AI_InvalidResponseDataError as a terminal provider error', async () => { + const kiloClient = createMockKiloClient({ + subscribeEvents: vi.fn().mockResolvedValue({ + stream: createEventStream([ + { + type: 'session.error', + properties: { + sessionID: 'kilo_sess_456', + error: { + name: 'AI_InvalidResponseDataError', + message: "Expected 'id' to be a string", + }, + }, + }, + ]), + }), + }); + + const manager = createManagerWithClient(kiloClient); + await openConnection(manager); + await vi.advanceTimersByTimeAsync(0); + + expect(callbacks.onTerminalError).toHaveBeenCalledWith({ + code: 'provider_error', + message: "Expected 'id' to be a string", + errorSource: 'assistant', + }); + }); + + it('classifies an AI_JSONParseError as a terminal provider error', async () => { + const kiloClient = createMockKiloClient({ + subscribeEvents: vi.fn().mockResolvedValue({ + stream: createEventStream([ + { + type: 'session.error', + properties: { + sessionID: 'kilo_sess_456', + error: { + name: 'AI_JSONParseError', + message: 'Unexpected end of JSON input', + }, + }, + }, + ]), + }), + }); + + const manager = createManagerWithClient(kiloClient); + await openConnection(manager); + await vi.advanceTimersByTimeAsync(0); + + expect(callbacks.onTerminalError).toHaveBeenCalledWith({ + code: 'provider_error', + message: 'Unexpected end of JSON input', + errorSource: 'assistant', + }); + }); + + it('classifies an expected-field pattern error as a terminal provider error', async () => { + const kiloClient = createMockKiloClient({ + subscribeEvents: vi.fn().mockResolvedValue({ + stream: createEventStream([ + { + type: 'session.error', + properties: { + sessionID: 'kilo_sess_456', + error: { + data: { message: "Expected 'choices' to be an array" }, + }, + }, + }, + ]), + }), + }); + + const manager = createManagerWithClient(kiloClient); + await openConnection(manager); + await vi.advanceTimersByTimeAsync(0); + + expect(callbacks.onTerminalError).toHaveBeenCalledWith({ + code: 'provider_error', + message: "Expected 'choices' to be an array", + errorSource: 'assistant', + }); + }); + + it('does not treat a transient provider error as terminal', async () => { + const kiloClient = createMockKiloClient({ + subscribeEvents: vi.fn().mockResolvedValue({ + stream: createEventStream([ + { + type: 'session.error', + properties: { + sessionID: 'kilo_sess_456', + error: { + name: 'ProviderError', + data: { message: 'Rate limit exceeded for provider request' }, + }, + }, + }, + ]), + }), + }); + + const manager = createManagerWithClient(kiloClient); + await openConnection(manager); + await vi.advanceTimersByTimeAsync(0); + + expect(callbacks.onTerminalError).not.toHaveBeenCalled(); + }); + it('records explicit Kilo gate results from event properties', async () => { const kiloClient = createMockKiloClient({ subscribeEvents: vi.fn().mockResolvedValue({ diff --git a/services/cloud-agent-next/wrapper/src/connection.ts b/services/cloud-agent-next/wrapper/src/connection.ts index 814b13985f..c8e369a83d 100644 --- a/services/cloud-agent-next/wrapper/src/connection.ts +++ b/services/cloud-agent-next/wrapper/src/connection.ts @@ -885,7 +885,14 @@ export function createConnectionManager( normalizedError.includes('too many requests') ); })()); - if (!code && !isExplicitAssistantFailure) return undefined; + const isProviderResponseError = + eventType === 'session.error' && + isAIProviderResponseError(terminalErrorText, properties.error); + if (!code && !isExplicitAssistantFailure && !isProviderResponseError) return undefined; + + if (isProviderResponseError && !code) { + code = 'provider_error'; + } return { ...(code ? { code } : {}), @@ -920,6 +927,18 @@ export function createConnectionManager( return /\b(model\s+(?:was\s+)?not\s+found|unknown\s+model|invalid\s+model)\b/i.test(message); } + function isAIProviderResponseError(errorText: string, error: unknown): boolean { + const normalized = errorText.toLowerCase(); + if (normalized.includes('ai_invalidresponsedataerror')) return true; + if (normalized.includes('ai_jsonparseerror')) return true; + if (isRecord(error) && typeof error.name === 'string') { + const name = error.name.toLowerCase(); + if (name === 'ai_invalidresponsedataerror' || name === 'ai_jsonparseerror') return true; + } + if (/^expected\s+['"].+['"]\s+to be\s+/i.test(errorText.trim())) return true; + return false; + } + function shouldFetchModelNotFoundDiagnostics( eventType: string, properties: Record, From 387f552b5b98975a95eff2399757f5dab739d271 Mon Sep 17 00:00:00 2001 From: Evgeny Shurakov Date: Mon, 13 Jul 2026 20:57:27 +0200 Subject: [PATCH 2/4] CloudAgentNext - address review feedback on cloud.status ready broadcast and provider error classification - onTerminalEvent returns boolean; only broadcast cloud.status: ready for accepted terminal events so stale events don't re-enable input mid-run - Tie onSessionTerminalized to the durable terminalization transition (before fallible downstream effects) and invoke unconditionally so finalizing runs with no accepted messages still publish readiness - Extend drain delay to 2s when aborted while disconnected so the buffered fatal error frame survives reconnect delivery before close() clears it - Remap provider_error to assistant_error in persistedFailureCode (matching the existing model_missing/payment_required pattern) so run-state reports validate against the queue-report schema without a DB migration --- .../src/persistence/CloudAgentSession.ts | 4 +- .../src/session/wrapper-supervisor.test.ts | 4 +- .../src/session/wrapper-supervisor.ts | 14 ++--- .../src/telemetry/queue-reports.test.ts | 1 + .../src/telemetry/queue-reports.ts | 6 +- .../session/disconnect-and-reaper.test.ts | 56 +++++++++++++++++++ .../wrapper/src/lifecycle.test.ts | 34 +++++++++++ .../cloud-agent-next/wrapper/src/lifecycle.ts | 36 ++++++++---- 8 files changed, 130 insertions(+), 25 deletions(-) diff --git a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts index 267bd37af4..f03ba21b4c 100644 --- a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts +++ b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts @@ -3269,9 +3269,9 @@ export class CloudAgentSession extends DurableObject { async handleWrapperTerminalEvent(params: WrapperTerminalEvent): Promise { const sessionId = await this.resolveSessionId(); - await this.getWrapperSupervisor().onTerminalEvent(params); + const accepted = await this.getWrapperSupervisor().onTerminalEvent(params); - if (params.status === 'failed' || params.status === 'interrupted') { + if (accepted && (params.status === 'failed' || params.status === 'interrupted')) { this.broadcastVolatileEvent({ executionId: '' as EventSourceId, sessionId: sessionId ?? '', 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 ebb3e2fe25..703c4cef97 100644 --- a/services/cloud-agent-next/src/session/wrapper-supervisor.test.ts +++ b/services/cloud-agent-next/src/session/wrapper-supervisor.test.ts @@ -1163,7 +1163,7 @@ describe('WrapperSupervisor', () => { expect(harness.onSessionTerminalized).not.toHaveBeenCalled(); }); - it('does not call onSessionTerminalized when liveness expiry has no accepted messages to terminalize', async () => { + it('calls onSessionTerminalized when liveness expiry abandons a finalizing run with no accepted messages', async () => { const harness = createHarness([ liveRuntimeState({ finalizingWrapperRunId: WRAPPER_RUN_ID, @@ -1181,7 +1181,7 @@ describe('WrapperSupervisor', () => { await harness.supervisor.runMaintenance(10_000); - expect(harness.onSessionTerminalized).not.toHaveBeenCalled(); + expect(harness.onSessionTerminalized).toHaveBeenCalledTimes(1); }); it('includes the disconnect grace expiry deadline even when the ping deadline is earlier', async () => { diff --git a/services/cloud-agent-next/src/session/wrapper-supervisor.ts b/services/cloud-agent-next/src/session/wrapper-supervisor.ts index cdfad6d8d2..4e019e2aa6 100644 --- a/services/cloud-agent-next/src/session/wrapper-supervisor.ts +++ b/services/cloud-agent-next/src/session/wrapper-supervisor.ts @@ -135,7 +135,7 @@ export type WrapperSupervisor = { ): Promise; observeFinalizing(wrapperRunId: string): Promise; onDisconnected(input: WrapperDisconnectedInput): Promise; - onTerminalEvent(params: WrapperTerminalEvent): Promise; + onTerminalEvent(params: WrapperTerminalEvent): Promise; requestPhysicalWrapperStop(reason: WrapperStopReason, target?: WrapperStopTarget): Promise; clearDisconnectGrace(): Promise; runMaintenance(now: number): Promise; @@ -650,6 +650,7 @@ export function createWrapperSupervisor( failureCode: activityObserved ? 'wrapper_error_after_activity' : failureCode, }); } + await onSessionTerminalized?.(); await messageSettlementOutbox.releaseWrapperTerminalWaitForIdleBatch(); if (isWrapperRunFinalizing(state) && state.wrapperRunId) { await messageSettlementOutbox.finalizeTerminalWrapperRunCallbackIfReady(state.wrapperRunId); @@ -666,10 +667,6 @@ export function createWrapperSupervisor( state.wrapperConnectionId ); } - - if (acceptedMessages.length > 0) { - await onSessionTerminalized?.(); - } } async function checkDisconnectGrace(now: number): Promise { @@ -733,6 +730,7 @@ export function createWrapperSupervisor( failureCode: activityObserved ? 'wrapper_error_after_activity' : 'wrapper_disconnected', }); } + await onSessionTerminalized?.(); await clearWrapperRuntimeIdentity( storage, { @@ -742,7 +740,6 @@ export function createWrapperSupervisor( { incrementGeneration: true } ); await releaseWrapperTerminalWaitForIdleBatchForWrapperRun(wrapperRunId); - await onSessionTerminalized?.(); } async function hasActiveWrapperWork(state: WrapperRuntimeState): Promise { @@ -1303,7 +1300,7 @@ export function createWrapperSupervisor( ); } - async function onTerminalEvent(params: WrapperTerminalEvent): Promise { + async function onTerminalEvent(params: WrapperTerminalEvent): Promise { const { wrapperRunId, status, @@ -1326,7 +1323,7 @@ export function createWrapperSupervisor( logger .withFields({ sessionId, wrapperRunId, status }) .warn('Ignoring non-current wrapper terminal event'); - return; + return false; } logger @@ -1474,6 +1471,7 @@ export function createWrapperSupervisor( ) { await sessionMessageQueue.requestPendingDrainIfNeeded(); } + return true; } async function runMaintenance(now: number): Promise { diff --git a/services/cloud-agent-next/src/telemetry/queue-reports.test.ts b/services/cloud-agent-next/src/telemetry/queue-reports.test.ts index 3fb2afb8f3..35193674ac 100644 --- a/services/cloud-agent-next/src/telemetry/queue-reports.test.ts +++ b/services/cloud-agent-next/src/telemetry/queue-reports.test.ts @@ -83,6 +83,7 @@ describe('Cloud Agent report emitter', () => { it.each([ ['agent_activity', 'payment_required', 'assistant_error'], ['agent_activity', 'model_missing', 'assistant_error'], + ['agent_activity', 'provider_error', 'assistant_error'], ['post_dispatch_no_activity', 'payment_required', 'wrapper_error_before_activity'], ['post_dispatch_no_activity', 'model_missing', 'wrapper_error_before_activity'], ] as const)( diff --git a/services/cloud-agent-next/src/telemetry/queue-reports.ts b/services/cloud-agent-next/src/telemetry/queue-reports.ts index e3716c6d88..02f8bd8b3c 100644 --- a/services/cloud-agent-next/src/telemetry/queue-reports.ts +++ b/services/cloud-agent-next/src/telemetry/queue-reports.ts @@ -52,7 +52,11 @@ function timestamp(value: number): string { } function persistedFailureCode(state: SessionMessageState): SessionMessageState['failureCode'] { - if (state.failureCode !== 'model_missing' && state.failureCode !== 'payment_required') { + if ( + state.failureCode !== 'model_missing' && + state.failureCode !== 'payment_required' && + state.failureCode !== 'provider_error' + ) { return state.failureCode; } if (state.failureStage === 'agent_activity') return 'assistant_error'; diff --git a/services/cloud-agent-next/test/integration/session/disconnect-and-reaper.test.ts b/services/cloud-agent-next/test/integration/session/disconnect-and-reaper.test.ts index 920b888420..930a954cfe 100644 --- a/services/cloud-agent-next/test/integration/session/disconnect-and-reaper.test.ts +++ b/services/cloud-agent-next/test/integration/session/disconnect-and-reaper.test.ts @@ -368,4 +368,60 @@ describe('Disconnect handling and compatibility execution RPCs', () => { }); expect(cloudStatusReady).toHaveLength(1); }); + + it('does not broadcast cloud.status: ready for a stale wrapper terminal event', async () => { + const userId = 'user_stale_terminal'; + const sessionId = 'agent_stale_terminal'; + const stub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`) + ); + + const result = await runInDurableObject(stub, async (instance, state) => { + await registerReadySession(instance, { + sessionId, + userId, + kiloSessionId: 'ses_stale_terminal', + prompt: 'stale terminal test', + mode: 'code', + model: 'test-model', + }); + + const now = Date.now(); + await state.storage.put('wrapper_runtime_state', { + wrapperRunId: 'wr_current', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_current', + }); + await putSessionMessageState(state.storage, { + messageId: 'msg_018f1e2d3c4bStaleTermABCDE', + status: 'accepted', + prompt: 'stale terminal test', + createdAt: now, + acceptedAt: now, + wrapperRunId: 'wr_current', + }); + + const broadcasted: Array<{ stream_event_type: string; payload: string }> = []; + instance.broadcastEvent = event => { + broadcasted.push({ + stream_event_type: event.stream_event_type, + payload: event.payload, + }); + }; + + await instance.handleWrapperTerminalEvent({ + wrapperRunId: 'wr_stale', + status: 'failed', + error: 'Assistant request failed', + }); + return { broadcasted }; + }); + + const cloudStatusReady = result.broadcasted.filter(event => { + if (event.stream_event_type !== 'cloud.status') return false; + const payload = JSON.parse(event.payload) as { cloudStatus?: { type?: string } }; + return payload.cloudStatus?.type === 'ready'; + }); + expect(cloudStatusReady).toHaveLength(0); + }); }); diff --git a/services/cloud-agent-next/wrapper/src/lifecycle.test.ts b/services/cloud-agent-next/wrapper/src/lifecycle.test.ts index 9ad2363204..0021a00923 100644 --- a/services/cloud-agent-next/wrapper/src/lifecycle.test.ts +++ b/services/cloud-agent-next/wrapper/src/lifecycle.test.ts @@ -119,4 +119,38 @@ describe('wrapper lifecycle drain races', () => { await wait(500); expect(events.filter(event => event.streamEventType === 'complete')).toHaveLength(1); }, 10_000); + + it('delays close longer when aborted while disconnected to allow reconnect delivery', async () => { + const state = new WrapperState(); + state.bindSession(sessionContext); + state.setSendToIngestFn(() => {}); + + let closed = false; + const lifecycle = createLifecycleManager( + { workspacePath: '/tmp' }, + { + state, + kiloClient: {} as WrapperKiloClient, + closeConnections: async () => { + closed = true; + }, + isConnected: () => false, + reconnectEventSubscription: () => {}, + } + ); + + state.acceptMessage('message-1', { + autoCommit: false, + condenseOnComplete: false, + }); + state.clearAllMessages(); + lifecycle.setAborted(); + lifecycle.triggerDrainAndClose(); + + await wait(500); + expect(closed).toBe(false); + + await wait(2_000); + expect(closed).toBe(true); + }, 10_000); }); diff --git a/services/cloud-agent-next/wrapper/src/lifecycle.ts b/services/cloud-agent-next/wrapper/src/lifecycle.ts index 1f1ce3eb73..5db53d8174 100644 --- a/services/cloud-agent-next/wrapper/src/lifecycle.ts +++ b/services/cloud-agent-next/wrapper/src/lifecycle.ts @@ -5,6 +5,13 @@ import { runCondenseOnComplete } from './condense-on-complete.js'; import { getCurrentBranch, logToFile } from './utils.js'; const DRAIN_DELAY_MS = 250; +/** + * When a terminal error is detected while the ingest WS is disconnected, the + * fatal error frame is buffered and only delivered after the WS reconnects + * (first attempt after ~1 s). Extend the drain window so the reconnect can + * flush the buffered terminal frame before close() clears the buffer. + */ +const TERMINAL_ERROR_DISCONNECTED_DRAIN_DELAY_MS = 2_000; const STABLE_ROOT_IDLE_MS = 3_000; const SSE_TRANSPORT_TIMEOUT_MS = 15_000; const AUTO_COMMIT_TIMEOUT_MS = 120_000; @@ -201,18 +208,23 @@ export function createLifecycleManager( }); } - drainTimeout = setTimeout(() => { - void deps - .closeConnections() - .catch(error => - logToFile(`close failed: ${error instanceof Error ? error.message : String(error)}`) - ) - .finally(() => { - isDraining = false; - drainTimeout = null; - state.clearSession(); - }); - }, DRAIN_DELAY_MS); + drainTimeout = setTimeout( + () => { + void deps + .closeConnections() + .catch(error => + logToFile(`close failed: ${error instanceof Error ? error.message : String(error)}`) + ) + .finally(() => { + isDraining = false; + drainTimeout = null; + state.clearSession(); + }); + }, + isAborted && !deps.isConnected() + ? TERMINAL_ERROR_DISCONNECTED_DRAIN_DELAY_MS + : DRAIN_DELAY_MS + ); } })(); } From 30a18913f71ebba52feb4b5722dcd9f158cf6b2b Mon Sep 17 00:00:00 2001 From: Evgeny Shurakov Date: Tue, 14 Jul 2026 11:53:18 +0200 Subject: [PATCH 3/4] CloudAgentNext - harden terminalization settlement errors and close immediately on reconnect during aborted drain --- .../src/session/wrapper-supervisor.test.ts | 24 ++++++++ .../src/session/wrapper-supervisor.ts | 55 ++++++++++++++----- .../wrapper/src/lifecycle.test.ts | 39 +++++++++++++ .../cloud-agent-next/wrapper/src/lifecycle.ts | 43 ++++++++++----- 4 files changed, 132 insertions(+), 29 deletions(-) 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 703c4cef97..6e347b49c3 100644 --- a/services/cloud-agent-next/src/session/wrapper-supervisor.test.ts +++ b/services/cloud-agent-next/src/session/wrapper-supervisor.test.ts @@ -1138,6 +1138,30 @@ describe('WrapperSupervisor', () => { expect(harness.onSessionTerminalized).toHaveBeenCalledTimes(1); }); + it('publishes readiness even when terminalization throws', async () => { + const acceptedAt = 2_000; + const noOutputDeadlineAt = acceptedAt + WRAPPER_NO_OUTPUT_TIMEOUT_MS; + const harness = createHarness([ + liveRuntimeState({ noOutputDeadlineAt, nextPingAt: noOutputDeadlineAt + 1 }), + OWNED_WRAPPER_LEASE, + ]); + await putSessionMessageState(harness.storage, acceptedMessage()); + + const originalPut = harness.storage.put.bind(harness.storage); + harness.storage.put = (async (key: string, value: unknown) => { + if (key.startsWith('session_message')) { + throw new Error('terminalize storage failure'); + } + await originalPut(key, value); + }) as typeof harness.storage.put; + + await harness.supervisor.runMaintenance(noOutputDeadlineAt); + + expect(harness.onSessionTerminalized).toHaveBeenCalledTimes(1); + const runtimeState = await getWrapperRuntimeState(harness.storage); + expect(runtimeState.wrapperConnectionId).toBeUndefined(); + }); + it('does not call onSessionTerminalized when disconnect grace expires for a stale wrapper run', async () => { const harness = createHarness([ liveRuntimeState({ diff --git a/services/cloud-agent-next/src/session/wrapper-supervisor.ts b/services/cloud-agent-next/src/session/wrapper-supervisor.ts index 4e019e2aa6..bc90a4338d 100644 --- a/services/cloud-agent-next/src/session/wrapper-supervisor.ts +++ b/services/cloud-agent-next/src/session/wrapper-supervisor.ts @@ -641,23 +641,48 @@ 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, - }); + try { + 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, + }); + } catch (terminalizeError) { + logger + .withFields({ + sessionId: getSessionIdForLogs(), + wrapperRunId: state.wrapperRunId, + messageId: message.messageId, + error: + terminalizeError instanceof Error + ? terminalizeError.message + : String(terminalizeError), + }) + .warn('Failed to terminalize accepted message during unhealthy wrapper handling'); + } } await onSessionTerminalized?.(); - await messageSettlementOutbox.releaseWrapperTerminalWaitForIdleBatch(); - if (isWrapperRunFinalizing(state) && state.wrapperRunId) { - await messageSettlementOutbox.finalizeTerminalWrapperRunCallbackIfReady(state.wrapperRunId); - } else { - await messageSettlementOutbox.finalizeIdleBatchCallbackIfReady({ - allowWithoutObservedIdle: true, - }); + try { + await messageSettlementOutbox.releaseWrapperTerminalWaitForIdleBatch(); + if (isWrapperRunFinalizing(state) && state.wrapperRunId) { + await messageSettlementOutbox.finalizeTerminalWrapperRunCallbackIfReady(state.wrapperRunId); + } else { + await messageSettlementOutbox.finalizeIdleBatchCallbackIfReady({ + allowWithoutObservedIdle: true, + }); + } + } catch (settlementError) { + logger + .withFields({ + sessionId: getSessionIdForLogs(), + wrapperRunId: state.wrapperRunId, + error: + settlementError instanceof Error ? settlementError.message : String(settlementError), + }) + .warn('Post-readiness settlement failed during unhealthy wrapper handling'); } if (state.wrapperConnectionId) { diff --git a/services/cloud-agent-next/wrapper/src/lifecycle.test.ts b/services/cloud-agent-next/wrapper/src/lifecycle.test.ts index 0021a00923..3737656887 100644 --- a/services/cloud-agent-next/wrapper/src/lifecycle.test.ts +++ b/services/cloud-agent-next/wrapper/src/lifecycle.test.ts @@ -153,4 +153,43 @@ describe('wrapper lifecycle drain races', () => { await wait(2_000); expect(closed).toBe(true); }, 10_000); + + it('closes immediately when reconnect restores during aborted drain', async () => { + const state = new WrapperState(); + state.bindSession(sessionContext); + state.setSendToIngestFn(() => {}); + + let closeCount = 0; + const lifecycle = createLifecycleManager( + { workspacePath: '/tmp' }, + { + state, + kiloClient: {} as WrapperKiloClient, + closeConnections: async () => { + closeCount++; + }, + isConnected: () => false, + reconnectEventSubscription: () => {}, + } + ); + + state.acceptMessage('message-1', { + autoCommit: false, + condenseOnComplete: false, + }); + state.clearAllMessages(); + lifecycle.setAborted(); + lifecycle.triggerDrainAndClose(); + + await wait(300); + expect(closeCount).toBe(0); + + lifecycle.onConnectionRestored(); + + await wait(100); + expect(closeCount).toBe(1); + + await wait(2_000); + expect(closeCount).toBe(1); + }, 10_000); }); diff --git a/services/cloud-agent-next/wrapper/src/lifecycle.ts b/services/cloud-agent-next/wrapper/src/lifecycle.ts index 5db53d8174..4d92c988c0 100644 --- a/services/cloud-agent-next/wrapper/src/lifecycle.ts +++ b/services/cloud-agent-next/wrapper/src/lifecycle.ts @@ -9,7 +9,9 @@ const DRAIN_DELAY_MS = 250; * When a terminal error is detected while the ingest WS is disconnected, the * fatal error frame is buffered and only delivered after the WS reconnects * (first attempt after ~1 s). Extend the drain window so the reconnect can - * flush the buffered terminal frame before close() clears the buffer. + * flush the buffered terminal frame before close() clears the buffer. If the + * reconnect completes within this window, the drain closes immediately rather + * than waiting for the full delay; this value is the fallback ceiling. */ const TERMINAL_ERROR_DISCONNECTED_DRAIN_DELAY_MS = 2_000; const STABLE_ROOT_IDLE_MS = 3_000; @@ -158,6 +160,19 @@ export function createLifecycleManager( } } + function performDrainClose(): void { + void deps + .closeConnections() + .catch(error => + logToFile(`close failed: ${error instanceof Error ? error.message : String(error)}`) + ) + .finally(() => { + isDraining = false; + drainTimeout = null; + state.clearSession(); + }); + } + function triggerDrainAndClose(): void { state.blockAdmissions(); if (isDraining) return; @@ -209,18 +224,7 @@ export function createLifecycleManager( } drainTimeout = setTimeout( - () => { - void deps - .closeConnections() - .catch(error => - logToFile(`close failed: ${error instanceof Error ? error.message : String(error)}`) - ) - .finally(() => { - isDraining = false; - drainTimeout = null; - state.clearSession(); - }); - }, + performDrainClose, isAborted && !deps.isConnected() ? TERMINAL_ERROR_DISCONNECTED_DRAIN_DELAY_MS : DRAIN_DELAY_MS @@ -286,7 +290,18 @@ export function createLifecycleManager( } armStableIdleCandidate(); }, - onConnectionRestored: armStableIdleCandidate, + onConnectionRestored: () => { + if (isDraining && isAborted && drainTimeout) { + logToFile( + 'reconnect restored during aborted drain — closing now that buffered terminal frame flushed' + ); + clearTimeout(drainTimeout); + drainTimeout = null; + performDrainClose(); + return; + } + armStableIdleCandidate(); + }, triggerDrainAndClose, signalCompletion, setAborted: () => { From 0eb422ce1889c019a9e86150054bd98b298dd3e3 Mon Sep 17 00:00:00 2001 From: Evgeny Shurakov Date: Tue, 14 Jul 2026 21:17:55 +0200 Subject: [PATCH 4/4] CloudAgentNext - preserve supervision on terminalize persistence failure and defer close during in-progress reconnect --- .../src/session/wrapper-supervisor.test.ts | 48 +++++++++++++++++-- .../src/session/wrapper-supervisor.ts | 45 +++++++++-------- .../wrapper/src/lifecycle.test.ts | 40 ++++++++++++++++ .../cloud-agent-next/wrapper/src/lifecycle.ts | 17 +++++++ services/cloud-agent-next/wrapper/src/main.ts | 1 + 5 files changed, 129 insertions(+), 22 deletions(-) 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 6e347b49c3..f32b7f0cf0 100644 --- a/services/cloud-agent-next/src/session/wrapper-supervisor.test.ts +++ b/services/cloud-agent-next/src/session/wrapper-supervisor.test.ts @@ -1138,7 +1138,7 @@ describe('WrapperSupervisor', () => { expect(harness.onSessionTerminalized).toHaveBeenCalledTimes(1); }); - it('publishes readiness even when terminalization throws', async () => { + it('preserves supervision when terminalization persistence fails', async () => { const acceptedAt = 2_000; const noOutputDeadlineAt = acceptedAt + WRAPPER_NO_OUTPUT_TIMEOUT_MS; const harness = createHarness([ @@ -1155,11 +1155,53 @@ describe('WrapperSupervisor', () => { await originalPut(key, value); }) as typeof harness.storage.put; + await expect(harness.supervisor.runMaintenance(noOutputDeadlineAt)).rejects.toThrow( + 'terminalize storage failure' + ); + + // Readiness is not published because the durable transition failed + expect(harness.onSessionTerminalized).not.toHaveBeenCalled(); + // Supervision is preserved — the wrapper runtime identity is not cleared + const runtimeState = await getWrapperRuntimeState(harness.storage); + expect(runtimeState.wrapperConnectionId).toBeDefined(); + // The message remains accepted so maintenance can retry terminalization + await expect(getSessionMessageState(harness.storage, MESSAGE_ID)).resolves.toMatchObject({ + status: 'accepted', + }); + }); + + it('publishes readiness when terminal effects fail after the durable transition', async () => { + const acceptedAt = 2_000; + const noOutputDeadlineAt = acceptedAt + WRAPPER_NO_OUTPUT_TIMEOUT_MS; + const harness = createHarness([ + liveRuntimeState({ noOutputDeadlineAt, nextPingAt: noOutputDeadlineAt + 1 }), + OWNED_WRAPPER_LEASE, + ]); + await putSessionMessageState(harness.storage, acceptedMessage()); + + // Let the durable terminal transition succeed (first session_message: + // state write), then fail on the subsequent terminal-effects write so the + // effect path throws while the message is already terminal. + let sessionMessageStatePutCount = 0; + const originalPut = harness.storage.put.bind(harness.storage); + harness.storage.put = (async (key: string, value: unknown) => { + if (key.startsWith('session_message:')) { + sessionMessageStatePutCount++; + if (sessionMessageStatePutCount > 1) { + throw new Error('terminal effects storage failure'); + } + } + await originalPut(key, value); + }) as typeof harness.storage.put; + await harness.supervisor.runMaintenance(noOutputDeadlineAt); + // Readiness is published because the durable transition succeeded expect(harness.onSessionTerminalized).toHaveBeenCalledTimes(1); - const runtimeState = await getWrapperRuntimeState(harness.storage); - expect(runtimeState.wrapperConnectionId).toBeUndefined(); + // The message is terminal despite the effect failure + await expect(getSessionMessageState(harness.storage, MESSAGE_ID)).resolves.toMatchObject({ + status: 'failed', + }); }); it('does not call onSessionTerminalized when disconnect grace expires for a stale wrapper run', async () => { diff --git a/services/cloud-agent-next/src/session/wrapper-supervisor.ts b/services/cloud-agent-next/src/session/wrapper-supervisor.ts index bc90a4338d..53d9ece431 100644 --- a/services/cloud-agent-next/src/session/wrapper-supervisor.ts +++ b/services/cloud-agent-next/src/session/wrapper-supervisor.ts @@ -13,6 +13,8 @@ import { classifyAssistantFailureMessage } from './safe-failure-projection.js'; import { countPendingSessionMessages, type SessionQueueStorage } from './pending-messages.js'; import type { SessionMessageQueue } from './session-message-queue.js'; import { + getSessionMessageState, + isTerminalMessageState, listMessagesForWrapperRun, listNonTerminalAcceptedMessages, type SessionMessageState, @@ -651,6 +653,22 @@ export function createWrapperSupervisor( failureCode: activityObserved ? 'wrapper_error_after_activity' : failureCode, }); } catch (terminalizeError) { + // terminalizeSessionMessageOnly only becomes fallible after the + // durable transition commits (event/callback effects). A throw while + // the message is still non-terminal means the durable write itself + // failed; re-throw so the wrapper runtime identity is preserved and + // maintenance can retry instead of orphaning still-accepted work. + let durableTransitionSucceeded = false; + try { + const persistedState = await getSessionMessageState(storage, message.messageId); + durableTransitionSucceeded = + persistedState !== undefined && isTerminalMessageState(persistedState); + } catch { + // Re-read failed — assume the durable transition did not succeed. + } + if (!durableTransitionSucceeded) { + throw terminalizeError; + } logger .withFields({ sessionId: getSessionIdForLogs(), @@ -661,28 +679,17 @@ export function createWrapperSupervisor( ? terminalizeError.message : String(terminalizeError), }) - .warn('Failed to terminalize accepted message during unhealthy wrapper handling'); + .warn('Failed to apply terminal effects during unhealthy wrapper handling'); } } await onSessionTerminalized?.(); - try { - await messageSettlementOutbox.releaseWrapperTerminalWaitForIdleBatch(); - if (isWrapperRunFinalizing(state) && state.wrapperRunId) { - await messageSettlementOutbox.finalizeTerminalWrapperRunCallbackIfReady(state.wrapperRunId); - } else { - await messageSettlementOutbox.finalizeIdleBatchCallbackIfReady({ - allowWithoutObservedIdle: true, - }); - } - } catch (settlementError) { - logger - .withFields({ - sessionId: getSessionIdForLogs(), - wrapperRunId: state.wrapperRunId, - error: - settlementError instanceof Error ? settlementError.message : String(settlementError), - }) - .warn('Post-readiness settlement failed during unhealthy wrapper handling'); + await messageSettlementOutbox.releaseWrapperTerminalWaitForIdleBatch(); + if (isWrapperRunFinalizing(state) && state.wrapperRunId) { + await messageSettlementOutbox.finalizeTerminalWrapperRunCallbackIfReady(state.wrapperRunId); + } else { + await messageSettlementOutbox.finalizeIdleBatchCallbackIfReady({ + allowWithoutObservedIdle: true, + }); } if (state.wrapperConnectionId) { diff --git a/services/cloud-agent-next/wrapper/src/lifecycle.test.ts b/services/cloud-agent-next/wrapper/src/lifecycle.test.ts index 3737656887..c8ec0b4f1b 100644 --- a/services/cloud-agent-next/wrapper/src/lifecycle.test.ts +++ b/services/cloud-agent-next/wrapper/src/lifecycle.test.ts @@ -192,4 +192,44 @@ describe('wrapper lifecycle drain races', () => { await wait(2_000); expect(closeCount).toBe(1); }, 10_000); + + it('defers close while reconnecting during aborted drain until reconnect gives up', async () => { + const state = new WrapperState(); + state.bindSession(sessionContext); + state.setSendToIngestFn(() => {}); + + let closed = false; + let reconnecting = true; + const lifecycle = createLifecycleManager( + { workspacePath: '/tmp' }, + { + state, + kiloClient: {} as WrapperKiloClient, + closeConnections: async () => { + closed = true; + }, + isConnected: () => false, + isReconnecting: () => reconnecting, + reconnectEventSubscription: () => {}, + } + ); + + state.acceptMessage('message-1', { + autoCommit: false, + condenseOnComplete: false, + }); + state.clearAllMessages(); + lifecycle.setAborted(); + lifecycle.triggerDrainAndClose(); + + // Past the 2-second fallback ceiling the close is still deferred because a + // reconnect is in progress with a buffered terminal frame. + await wait(2_500); + expect(closed).toBe(false); + + // Once the reconnect gives up the next poll closes. + reconnecting = false; + await wait(500); + expect(closed).toBe(true); + }, 10_000); }); diff --git a/services/cloud-agent-next/wrapper/src/lifecycle.ts b/services/cloud-agent-next/wrapper/src/lifecycle.ts index 4d92c988c0..9a4662525d 100644 --- a/services/cloud-agent-next/wrapper/src/lifecycle.ts +++ b/services/cloud-agent-next/wrapper/src/lifecycle.ts @@ -27,6 +27,14 @@ export type LifecycleDependencies = { kiloClient: WrapperKiloClient; closeConnections: () => Promise; isConnected: () => boolean; + /** + * Whether the ingest WebSocket is currently attempting to reconnect. When a + * terminal error is detected while disconnected, the fatal frame is buffered + * and only flushed once a reconnect completes; the drain close defers while a + * reconnect is in progress so that frame can be delivered instead of being + * discarded by close(). + */ + isReconnecting?: () => boolean; reconnectEventSubscription: () => void; }; @@ -161,6 +169,15 @@ export function createLifecycleManager( } function performDrainClose(): void { + // When closing on an aborted drain while disconnected, a reconnect may + // still be in progress with a buffered fatal terminal frame. Closing now + // would abort that reconnect and discard the frame before it can be + // flushed. Defer until the reconnect completes (onConnectionRestored + // closes immediately once the frame is flushed) or gives up. + if (isAborted && !deps.isConnected() && deps.isReconnecting?.()) { + drainTimeout = setTimeout(performDrainClose, DRAIN_DELAY_MS); + return; + } void deps .closeConnections() .catch(error => diff --git a/services/cloud-agent-next/wrapper/src/main.ts b/services/cloud-agent-next/wrapper/src/main.ts index 39bc1d49ae..7e3611aebd 100644 --- a/services/cloud-agent-next/wrapper/src/main.ts +++ b/services/cloud-agent-next/wrapper/src/main.ts @@ -551,6 +551,7 @@ async function main() { kiloClient: nextKiloClient, closeConnections: () => connectionManager?.close() ?? Promise.resolve(), isConnected: () => connectionManager?.isConnected() ?? false, + isReconnecting: () => connectionManager?.isReconnecting() ?? false, reconnectEventSubscription: () => connectionManager?.reconnectEventSubscription(), } );