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..f03ba21b4c 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(); - await this.getWrapperSupervisor().onTerminalEvent(params); + const sessionId = await this.resolveSessionId(); + const accepted = await this.getWrapperSupervisor().onTerminalEvent(params); + + if (accepted && (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..f32b7f0cf0 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,153 @@ 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('preserves supervision when terminalization persistence fails', 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 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); + // 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 () => { + 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('calls onSessionTerminalized when liveness expiry abandons a finalizing run with no accepted messages', 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).toHaveBeenCalledTimes(1); + }); + 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..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, @@ -135,7 +137,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; @@ -178,6 +180,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 +371,7 @@ export function createWrapperSupervisor( recordSharedSandboxFailover, requestAlarmAtOrBefore, getSessionIdForLogs, + onSessionTerminalized, } = dependencies; async function readDisconnectGrace(): Promise { @@ -633,15 +643,46 @@ 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) { + // 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(), + wrapperRunId: state.wrapperRunId, + messageId: message.messageId, + error: + terminalizeError instanceof Error + ? terminalizeError.message + : String(terminalizeError), + }) + .warn('Failed to apply terminal effects during unhealthy wrapper handling'); + } } + await onSessionTerminalized?.(); await messageSettlementOutbox.releaseWrapperTerminalWaitForIdleBatch(); if (isWrapperRunFinalizing(state) && state.wrapperRunId) { await messageSettlementOutbox.finalizeTerminalWrapperRunCallbackIfReady(state.wrapperRunId); @@ -721,6 +762,7 @@ export function createWrapperSupervisor( failureCode: activityObserved ? 'wrapper_error_after_activity' : 'wrapper_disconnected', }); } + await onSessionTerminalized?.(); await clearWrapperRuntimeIdentity( storage, { @@ -1290,7 +1332,7 @@ export function createWrapperSupervisor( ); } - async function onTerminalEvent(params: WrapperTerminalEvent): Promise { + async function onTerminalEvent(params: WrapperTerminalEvent): Promise { const { wrapperRunId, status, @@ -1313,7 +1355,7 @@ export function createWrapperSupervisor( logger .withFields({ sessionId, wrapperRunId, status }) .warn('Ignoring non-current wrapper terminal event'); - return; + return false; } logger @@ -1461,6 +1503,7 @@ export function createWrapperSupervisor( ) { await sessionMessageQueue.requestPendingDrainIfNeeded(); } + return true; } async function runMaintenance(now: number): 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/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 c58fdabeaf..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 @@ -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,176 @@ 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); + }); + + 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/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, diff --git a/services/cloud-agent-next/wrapper/src/lifecycle.test.ts b/services/cloud-agent-next/wrapper/src/lifecycle.test.ts index 9ad2363204..c8ec0b4f1b 100644 --- a/services/cloud-agent-next/wrapper/src/lifecycle.test.ts +++ b/services/cloud-agent-next/wrapper/src/lifecycle.test.ts @@ -119,4 +119,117 @@ 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); + + 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); + + 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 1f1ce3eb73..9a4662525d 100644 --- a/services/cloud-agent-next/wrapper/src/lifecycle.ts +++ b/services/cloud-agent-next/wrapper/src/lifecycle.ts @@ -5,6 +5,15 @@ 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. 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; const SSE_TRANSPORT_TIMEOUT_MS = 15_000; const AUTO_COMMIT_TIMEOUT_MS = 120_000; @@ -18,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; }; @@ -151,6 +168,28 @@ 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 => + 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; @@ -201,18 +240,12 @@ 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( + performDrainClose, + isAborted && !deps.isConnected() + ? TERMINAL_ERROR_DISCONNECTED_DRAIN_DELAY_MS + : DRAIN_DELAY_MS + ); } })(); } @@ -274,7 +307,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: () => { 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(), } );