diff --git a/.specs/cloud-agent-session.md b/.specs/cloud-agent-session.md index 09c69d1754..d552557383 100644 --- a/.specs/cloud-agent-session.md +++ b/.specs/cloud-agent-session.md @@ -69,15 +69,19 @@ repository. commands. 3. Preparation that only acquired and booted an environment -- warm reuse, no real provisioning -- MUST NOT leave a completed preparation row. -4. Running and failed preparation MUST always be visible. Failed preparation - MUST show an error and a way to open details. +4. Running and failed preparation MUST always be visible. Failed preparation, + its triggering failed message, and its safe error MUST remain visible, with + a way to open details. 5. Setup commands MUST run on the first prepare and again on every rebuild. A failing or timed-out setup command MUST fail preparation and the turn. 6. Follow-up turns MUST NOT show preparation unless the environment was rebuilt. 7. Preparation output MUST NOT reveal tokens or secret values. -8. The composer MUST stay disabled until the environment is ready, and MUST say - which state it is waiting on. +8. The composer MUST stay disabled while preparation or finalization runs, and + MUST say which state it is waiting on. After preparation failure settles the + turn, the composer MUST be restored when the session is writable and its + transport permits sending. A later submission MUST use fresh message and + preparation-attempt identities. ### Turn diff --git a/packages/cloud-agent-sdk/src/base-connection.ts b/packages/cloud-agent-sdk/src/base-connection.ts index 7c41a7371c..49a1d3271e 100644 --- a/packages/cloud-agent-sdk/src/base-connection.ts +++ b/packages/cloud-agent-sdk/src/base-connection.ts @@ -57,6 +57,7 @@ export type Connection = { disconnect: () => void; reconnectWithRefreshedAuth?: () => void; retryReconnect: () => void; + recoverAfterSuccessfulMutation: () => void; destroy: () => void; }; @@ -80,13 +81,16 @@ export function createBaseConnection(config: BaseConnectionConfig): Connec let authRefreshAttempted = false; let connected = false; let reconnectAttempt = 0; - let exhausted = false; + let exhaustionReason: 'retry-limit' | 'auth-failure' | null = null; let generation = 0; let hasConnectedOnce = false; let stalenessTimeoutId: ReturnType | null = null; let lastMessageTime = 0; let hiddenAt = 0; let preconnectAuthRefreshAttempted = false; + // Coalesce successful mutations into one retry budget at a time. A later mutation may + // request another bounded cycle after that budget exhausts without any inbound event. + let mutationRecovery: 'idle' | 'pending' | 'active' = 'idle'; const stalenessTimeoutMs = config.stalenessTimeoutMs ?? DEFAULT_STALENESS_TIMEOUT_MS; const maxReconnectAttempts = config.maxReconnectAttempts ?? MAX_RECONNECT_ATTEMPTS; @@ -109,8 +113,8 @@ export function createBaseConnection(config: BaseConnectionConfig): Connec } function clearExhausted(): void { - if (exhausted) { - exhausted = false; + if (exhaustionReason !== null) { + exhaustionReason = null; config.onReconnectExhaustionChange?.(false); } } @@ -178,8 +182,14 @@ export function createBaseConnection(config: BaseConnectionConfig): Connec if (destroyed || intentionalDisconnect || expectedGeneration !== generation) return; if (attempt >= maxReconnectAttempts) { - if (!exhausted) { - exhausted = true; + if (mutationRecovery === 'pending') { + mutationRecovery = 'active'; + retryReconnect(); + return; + } + mutationRecovery = 'idle'; + if (exhaustionReason === null) { + exhaustionReason = 'retry-limit'; config.onReconnectExhaustionChange?.(true); } return; @@ -262,6 +272,7 @@ export function createBaseConnection(config: BaseConnectionConfig): Connec // Reset auth refresh flag on successful message authRefreshAttempted = false; reconnectAttempt = 0; + mutationRecovery = 'idle'; clearExhausted(); if (!connected) { @@ -313,9 +324,10 @@ export function createBaseConnection(config: BaseConnectionConfig): Connec // Already tried refreshing auth and still failing - stop retrying. // The current physical route is gone even though no new socket follows. if (isAuthFailure && authRefreshAttempted) { + mutationRecovery = 'idle'; notifyReplacingConnection(expectedGeneration); - if (!exhausted) { - exhausted = true; + if (exhaustionReason === null) { + exhaustionReason = 'auth-failure'; config.onReconnectExhaustionChange?.(true); } return; @@ -457,6 +469,7 @@ export function createBaseConnection(config: BaseConnectionConfig): Connec destroyed = false; authRefreshAttempted = false; preconnectAuthRefreshAttempted = false; + mutationRecovery = 'idle'; connected = false; reconnectAttempt = 0; clearExhausted(); @@ -473,6 +486,7 @@ export function createBaseConnection(config: BaseConnectionConfig): Connec intentionalDisconnect = true; generation += 1; preconnectAuthRefreshAttempted = false; + mutationRecovery = 'idle'; clearReconnectTimer(); clearStalenessTimeout(); @@ -519,10 +533,28 @@ export function createBaseConnection(config: BaseConnectionConfig): Connec void refreshAndConnect(generation); } + function recoverAfterSuccessfulMutation() { + if ( + destroyed || + intentionalDisconnect || + connected || + exhaustionReason === 'auth-failure' || + mutationRecovery !== 'idle' + ) + return; + if (exhaustionReason === 'retry-limit') { + mutationRecovery = 'active'; + retryReconnect(); + return; + } + mutationRecovery = 'pending'; + } + function destroy() { destroyed = true; generation += 1; preconnectAuthRefreshAttempted = false; + mutationRecovery = 'idle'; clearReconnectTimer(); clearStalenessTimeout(); @@ -537,7 +569,14 @@ export function createBaseConnection(config: BaseConnectionConfig): Connec connected = false; } - return { connect, disconnect, reconnectWithRefreshedAuth, retryReconnect, destroy }; + return { + connect, + disconnect, + reconnectWithRefreshedAuth, + retryReconnect, + recoverAfterSuccessfulMutation, + destroy, + }; } export function createBrowserLifecycleHooks(): ConnectionLifecycleHooks { diff --git a/packages/cloud-agent-sdk/src/cloud-agent-transport-recovery.test.ts b/packages/cloud-agent-sdk/src/cloud-agent-transport-recovery.test.ts new file mode 100644 index 0000000000..e26462d163 --- /dev/null +++ b/packages/cloud-agent-sdk/src/cloud-agent-transport-recovery.test.ts @@ -0,0 +1,397 @@ +import { createCloudAgentTransport } from './cloud-agent-transport'; +import { createServiceState } from './service-state'; +import { createEventHelpers } from './__fixtures__/helpers'; +import type { ChatEvent, ServiceEvent } from './normalizer'; +import type { CloudAgentApi, Transport, TransportSendInput } from './transport'; +import { cloudAgentId, kiloId, makeSnapshot } from './test-helpers'; + +type TestSocket = { + url: string; + readyState: number; + onopen: (() => void) | null; + onmessage: ((event: MessageEvent) => void) | null; + onclose: ((event: CloseEvent) => void) | null; + close: jest.Mock; +}; + +const sockets: TestSocket[] = []; +const originalWebSocket = globalThis.WebSocket; +const { createEvent, kilocode, resetCounter } = createEventHelpers(); +const input = { + messageId: 'fresh-message', + payload: { + type: 'prompt', + prompt: 'fresh demand', + mode: 'code', + model: { providerID: 'kilo', modelID: 'fake-deterministic' }, + }, +} satisfies TransportSendInput; + +beforeEach(() => { + jest.useFakeTimers(); + jest.spyOn(Math, 'random').mockReturnValue(0); + resetCounter(); + sockets.length = 0; + const constructor = Object.assign( + jest.fn((url: string) => { + const socket: TestSocket = { + url, + readyState: 0, + onopen: null, + onmessage: null, + onclose: null, + close: jest.fn(), + }; + sockets.push(socket); + return socket; + }), + { OPEN: 1, CLOSED: 3 } + ); + Object.defineProperty(globalThis, 'WebSocket', { configurable: true, value: constructor }); +}); + +afterEach(() => { + Object.defineProperty(globalThis, 'WebSocket', { + configurable: true, + value: originalWebSocket, + }); + jest.useRealTimers(); + jest.restoreAllMocks(); +}); + +function latestSocket(): TestSocket { + const socket = sockets.at(-1); + if (!socket) throw new Error('Expected a stream socket'); + return socket; +} + +function closeSocket(code = 1006): void { + const socket = latestSocket(); + socket.readyState = 3; + socket.onclose?.({ code, reason: '', wasClean: false } as CloseEvent); +} + +function receive(event: ReturnType, socket = latestSocket()): void { + socket.readyState = 1; + socket.onmessage?.({ data: JSON.stringify(event) } as MessageEvent); +} + +function createHarness() { + const chatEvents: ChatEvent[] = []; + const serviceEvents: ServiceEvent[] = []; + const state = createServiceState({ rootSessionId: 'ses-1' }); + const send = jest.fn, Parameters>( + async () => ({ accepted: true }) + ); + const api = { + send, + interrupt: jest.fn(async () => ({ success: true })), + cancelQueuedMessage: jest.fn(async () => ({ dropped: true })), + answer: jest.fn(async () => ({ success: true })), + reject: jest.fn(async () => ({ success: true })), + respondToPermission: jest.fn(async () => ({ success: true })), + } satisfies CloudAgentApi; + const getTicket = jest.fn(async () => ({ + ticket: 'local-stream-ticket', + expiresAt: Math.floor(Date.now() / 1000) + 60, + })); + const fetchSnapshot = jest.fn(async () => makeSnapshot({ id: 'ses-1' })); + const transport = createCloudAgentTransport({ + sessionId: cloudAgentId('ses-1'), + kiloSessionId: kiloId('ses-1'), + websocketBaseUrl: 'ws://localhost:9999', + getTicket, + fetchSnapshot, + api, + })({ + onChatEvent: event => chatEvents.push(event), + onServiceEvent: event => { + serviceEvents.push(event); + state.process(event); + }, + }); + return { + transport, + state, + api, + send, + getTicket, + fetchSnapshot, + chatEvents, + serviceEvents, + async submit() { + if (!transport.send) throw new Error('Expected send support'); + return transport.send(input); + }, + }; +} + +async function connect(harness: ReturnType): Promise { + harness.transport.connect(); + await jest.advanceTimersByTimeAsync(0); + receive({ ...createEvent('connected', {}), eventId: 7 }); +} + +type MutationName = 'interrupt' | 'dropQueuedMessage' | 'answer' | 'reject' | 'permission'; + +async function performMutation(transport: Transport, mutation: MutationName): Promise { + switch (mutation) { + case 'interrupt': + return transport.interrupt?.(); + case 'dropQueuedMessage': + return transport.dropQueuedMessage?.('queued-message'); + case 'answer': + return transport.answer?.({ requestId: 'question-1', answers: [['yes']] }); + case 'reject': + return transport.reject?.({ requestId: 'question-1' }); + case 'permission': + return transport.respondToPermission?.({ requestId: 'permission-1', response: 'once' }); + } +} + +async function exhaustRetries(): Promise { + const startingCount = sockets.length; + closeSocket(); + for (let attempt = 0; attempt < 8; attempt += 1) { + await jest.advanceTimersByTimeAsync(Math.min(30_000, 1000 * 2 ** attempt) / 2); + closeSocket(); + } + expect(sockets).toHaveLength(startingCount + 8); + await jest.advanceTimersByTimeAsync(600_000); + expect(sockets).toHaveLength(startingCount + 8); +} + +describe('Cloud Agent stream recovery after a backend outage', () => { + it('reopens an exhausted stream after an accepted send and replays from its cursor', async () => { + const harness = createHarness(); + await connect(harness); + const oldSocket = latestSocket(); + await exhaustRetries(); + expect(harness.state.getStatus()).toEqual({ type: 'disconnected' }); + expect(harness.serviceEvents).toContainEqual({ + type: 'stopped', + reason: 'transport-disconnected', + }); + const ticketCount = harness.getTicket.mock.calls.length; + + await expect(harness.submit()).resolves.toEqual({ accepted: true }); + await jest.advanceTimersByTimeAsync(0); + + expect(sockets).toHaveLength(10); + expect(harness.getTicket).toHaveBeenCalledTimes(ticketCount + 1); + expect(new URL(latestSocket().url).searchParams.get('fromId')).toBe('7'); + expect(harness.fetchSnapshot).toHaveBeenCalledTimes(1); + expect(harness.send).toHaveBeenCalledTimes(1); + expect(harness.send).toHaveBeenCalledWith({ + sessionId: 'ses-1', + messageId: input.messageId, + payload: { + type: 'prompt', + prompt: 'fresh demand', + mode: 'code', + model: 'fake-deterministic', + }, + }); + + const reply = { + ...kilocode('message.part.updated', { + part: { + id: 'reply-part', + messageID: 'reply', + sessionID: 'ses-1', + type: 'text', + text: 'canonical reply', + }, + }), + eventId: 8, + }; + receive(reply, oldSocket); + receive({ ...reply, sessionId: 'other-session' }); + expect(harness.chatEvents).toEqual([]); + receive({ + ...createEvent('connected', { + sessionStatus: { type: 'idle' }, + cloudStatus: { type: 'ready' }, + }), + eventId: 0, + }); + expect(harness.state.getStatus()).toEqual({ type: 'idle' }); + receive(reply); + expect(harness.chatEvents).toEqual([ + expect.objectContaining({ + type: 'message.part.updated', + part: expect.objectContaining({ text: 'canonical reply', sessionID: 'ses-1' }), + }), + ]); + harness.transport.destroy(); + }); + + it.each(['interrupt', 'dropQueuedMessage', 'answer', 'reject', 'permission'] as const)( + 'reopens an exhausted stream after a successful %s mutation', + async mutation => { + const harness = createHarness(); + await connect(harness); + await exhaustRetries(); + const ticketCount = harness.getTicket.mock.calls.length; + + await performMutation(harness.transport, mutation); + await jest.advanceTimersByTimeAsync(0); + + expect(sockets).toHaveLength(10); + expect(harness.getTicket).toHaveBeenCalledTimes(ticketCount + 1); + expect(new URL(latestSocket().url).searchParams.get('fromId')).toBe('7'); + harness.transport.destroy(); + } + ); + + it('reserves one recovery budget when a mutation succeeds before retries exhaust', async () => { + const harness = createHarness(); + await connect(harness); + closeSocket(); + for (let attempt = 0; attempt < 7; attempt += 1) { + await jest.advanceTimersByTimeAsync(Math.min(30_000, 1000 * 2 ** attempt) / 2); + closeSocket(); + } + expect(sockets).toHaveLength(8); + + await performMutation(harness.transport, 'interrupt'); + await jest.advanceTimersByTimeAsync(15_000); + closeSocket(); + await jest.advanceTimersByTimeAsync(0); + + expect(sockets).toHaveLength(10); + expect(new URL(latestSocket().url).searchParams.get('fromId')).toBe('7'); + receive(createEvent('connected', {})); + await jest.advanceTimersByTimeAsync(600_000); + expect(sockets).toHaveLength(10); + harness.transport.destroy(); + }); + + it('recovers established closure and failed handshakes within the existing retry budget', async () => { + const harness = createHarness(); + await connect(harness); + closeSocket(); + await jest.advanceTimersByTimeAsync(500); + closeSocket(); + await jest.advanceTimersByTimeAsync(1000); + receive(createEvent('connected', {})); + receive(kilocode('session.status', { sessionID: 'ses-1', status: { type: 'idle' } })); + expect(harness.serviceEvents.at(-1)).toEqual({ + type: 'session.status', + sessionId: 'ses-1', + status: { type: 'idle' }, + }); + expect(sockets).toHaveLength(3); + expect(harness.send).not.toHaveBeenCalled(); + harness.transport.destroy(); + }); + + it.each([403, 404, 503])( + 'does not restart observation when send fails with HTTP %s', + async status => { + const harness = createHarness(); + await connect(harness); + await exhaustRetries(); + const ticketCount = harness.getTicket.mock.calls.length; + harness.send.mockRejectedValueOnce(new Error(`HTTP ${status}`)); + await expect(harness.submit()).rejects.toThrow(`HTTP ${status}`); + await jest.advanceTimersByTimeAsync(600_000); + expect(sockets).toHaveLength(9); + expect(harness.getTicket).toHaveBeenCalledTimes(ticketCount); + harness.transport.destroy(); + } + ); + + it.each(['connected', 'retrying'] as const)( + 'does not replace a %s stream after send', + async state => { + const harness = createHarness(); + await connect(harness); + if (state === 'retrying') closeSocket(); + await harness.submit(); + await jest.advanceTimersByTimeAsync(0); + expect(sockets).toHaveLength(1); + expect(harness.getTicket).toHaveBeenCalledTimes(1); + harness.transport.destroy(); + } + ); + + it.each(['disconnect', 'destroy', 'connect'] as const)( + 'fences an accepted send that settles after %s', + async action => { + const harness = createHarness(); + await connect(harness); + await exhaustRetries(); + let resolveSend: (value: unknown) => void = () => {}; + harness.send.mockImplementationOnce(() => new Promise(resolve => (resolveSend = resolve))); + const pending = harness.submit(); + harness.transport[action](); + await jest.advanceTimersByTimeAsync(0); + if (action === 'connect') await exhaustRetries(); + const socketCount = sockets.length; + const ticketCount = harness.getTicket.mock.calls.length; + resolveSend({ accepted: true }); + await pending; + await jest.advanceTimersByTimeAsync(600_000); + expect(sockets).toHaveLength(socketCount); + expect(harness.getTicket).toHaveBeenCalledTimes(ticketCount); + harness.transport.destroy(); + } + ); + + it('does not revive a stream stopped by terminal authentication failure', async () => { + const harness = createHarness(); + await connect(harness); + closeSocket(4001); + await jest.advanceTimersByTimeAsync(0); + closeSocket(4001); + const ticketCount = harness.getTicket.mock.calls.length; + await harness.submit(); + await jest.advanceTimersByTimeAsync(600_000); + expect(sockets).toHaveLength(2); + expect(harness.getTicket).toHaveBeenCalledTimes(ticketCount); + harness.transport.destroy(); + }); + + it('coalesces concurrent sends and requires a new mutation after recovery exhausts', async () => { + const harness = createHarness(); + await connect(harness); + await exhaustRetries(); + let resolveTicket: (value: { ticket: string; expiresAt: number }) => void = () => {}; + harness.getTicket.mockImplementationOnce( + () => new Promise(resolve => (resolveTicket = resolve)) + ); + const ticketCount = harness.getTicket.mock.calls.length; + await Promise.all([harness.submit(), harness.submit()]); + expect(harness.getTicket).toHaveBeenCalledTimes(ticketCount + 1); + resolveTicket({ ticket: 'renewed-ticket', expiresAt: Math.floor(Date.now() / 1000) + 60 }); + await jest.advanceTimersByTimeAsync(0); + expect(sockets).toHaveLength(10); + expect(new URL(latestSocket().url).searchParams.get('ticket')).toBe('renewed-ticket'); + await exhaustRetries(); + expect(harness.send).toHaveBeenCalledTimes(2); + + const exhaustedSocketCount = sockets.length; + await harness.submit(); + await jest.advanceTimersByTimeAsync(0); + expect(sockets).toHaveLength(exhaustedSocketCount + 1); + expect(harness.send).toHaveBeenCalledTimes(3); + harness.transport.destroy(); + }); + + it.each(['disconnect', 'destroy'] as const)('fences ticket renewal after %s', async action => { + const harness = createHarness(); + await connect(harness); + await exhaustRetries(); + let resolveTicket: (value: { ticket: string; expiresAt: number }) => void = () => {}; + harness.getTicket.mockImplementationOnce( + () => new Promise(resolve => (resolveTicket = resolve)) + ); + await harness.submit(); + harness.transport[action](); + resolveTicket({ ticket: 'late-ticket', expiresAt: Math.floor(Date.now() / 1000) + 60 }); + await jest.advanceTimersByTimeAsync(600_000); + expect(sockets).toHaveLength(9); + harness.transport.destroy(); + }); +}); diff --git a/packages/cloud-agent-sdk/src/cloud-agent-transport.ts b/packages/cloud-agent-sdk/src/cloud-agent-transport.ts index 3c9707f201..8b95f1afb2 100644 --- a/packages/cloud-agent-sdk/src/cloud-agent-transport.ts +++ b/packages/cloud-agent-sdk/src/cloud-agent-transport.ts @@ -298,6 +298,15 @@ function createCloudAgentTransport(config: CloudAgentTransportConfig): Transport config.onError?.(message); } + async function runMutation(mutation: () => Promise): Promise { + const expectedGeneration = lifecycleGeneration; + const result = await mutation(); + if (expectedGeneration === lifecycleGeneration) { + connection?.recoverAfterSuccessfulMutation(); + } + return result; + } + return { connect() { closeConnection('destroy'); @@ -325,25 +334,32 @@ function createCloudAgentTransport(config: CloudAgentTransportConfig): Transport closeConnection('destroy'); }, - send: async input => - config.api.send({ - sessionId: config.sessionId, - payload: normalizeCloudAgentPayload(input.payload), - ...(input.messageId ? { messageId: input.messageId } : {}), - ...(input.attachments ? { attachments: input.attachments } : {}), - ...(input.images ? { images: input.images } : {}), - }), - interrupt: () => config.api.interrupt({ sessionId: config.sessionId }), + send: input => + runMutation(() => + config.api.send({ + sessionId: config.sessionId, + payload: normalizeCloudAgentPayload(input.payload), + ...(input.messageId ? { messageId: input.messageId } : {}), + ...(input.attachments ? { attachments: input.attachments } : {}), + ...(input.images ? { images: input.images } : {}), + }) + ), + interrupt: () => runMutation(() => config.api.interrupt({ sessionId: config.sessionId })), dropQueuedMessage: async messageId => { - if (!config.api.cancelQueuedMessage) { + const cancelQueuedMessage = config.api.cancelQueuedMessage; + if (!cancelQueuedMessage) { throw new Error('Cloud Agent cancel queued message is not configured'); } - return config.api.cancelQueuedMessage({ sessionId: config.sessionId, messageId }); + return runMutation(() => cancelQueuedMessage({ sessionId: config.sessionId, messageId })); }, - answer: payload => config.api.answer({ sessionId: config.sessionId, ...payload }), - reject: payload => config.api.reject({ sessionId: config.sessionId, ...payload }), + answer: payload => + runMutation(() => config.api.answer({ sessionId: config.sessionId, ...payload })), + reject: payload => + runMutation(() => config.api.reject({ sessionId: config.sessionId, ...payload })), respondToPermission: payload => - config.api.respondToPermission({ sessionId: config.sessionId, ...payload }), + runMutation(() => + config.api.respondToPermission({ sessionId: config.sessionId, ...payload }) + ), }; }; } diff --git a/packages/cloud-agent-sdk/src/service-state.test.ts b/packages/cloud-agent-sdk/src/service-state.test.ts index 1b12f92465..f3134bb3de 100644 --- a/packages/cloud-agent-sdk/src/service-state.test.ts +++ b/packages/cloud-agent-sdk/src/service-state.test.ts @@ -675,6 +675,153 @@ describe('createServiceState', () => { }); expect(state.getCloudStatus()).toEqual({ type: 'error', message: 'Clone failed' }); + expect(state.getPreparationAttempts()).toEqual([ + expect.objectContaining({ + id: 'attempt-1', + triggerMessageId: 'message-1', + status: 'failed', + safeError: 'Clone failed', + }), + ]); + + state.process({ + type: 'preparing', + version: 2, + attemptId: 'attempt-2', + triggerMessageId: 'message-2', + revision: 1, + timestamp: 3_000, + step: 'workspace_setup', + message: 'Preparing environment', + action: 'attempt_started', + }); + + expect(state.getPreparationAttempts()).toEqual([ + expect.objectContaining({ + id: 'attempt-1', + triggerMessageId: 'message-1', + status: 'failed', + safeError: 'Clone failed', + }), + expect.objectContaining({ + id: 'attempt-2', + triggerMessageId: 'message-2', + status: 'running', + }), + ]); + }); + + it.each([ + { + terminalAction: 'attempt_completed' as const, + terminalStep: 'ready', + expectedStatus: 'completed' as const, + expectedCloudStatus: { type: 'ready' as const }, + }, + { + terminalAction: 'attempt_failed' as const, + terminalStep: 'failed', + expectedStatus: 'failed' as const, + expectedCloudStatus: { type: 'error' as const, message: 'Clone failed' }, + safeError: 'Clone failed', + }, + ])('does not restart a $expectedStatus attempt', terminal => { + const state = createServiceState(makeConfig()); + const base = { + type: 'preparing' as const, + version: 2 as const, + attemptId: 'attempt-1', + triggerMessageId: 'message-1', + }; + + state.process({ + ...base, + revision: 1, + timestamp: 1_000, + step: 'workspace_setup', + message: 'Preparing environment', + action: 'attempt_started', + }); + state.process({ + ...base, + revision: 2, + timestamp: 2_000, + step: terminal.terminalStep, + message: 'Clone failed', + action: terminal.terminalAction, + ...(terminal.safeError === undefined ? {} : { safeError: terminal.safeError }), + }); + state.process({ + ...base, + revision: 3, + timestamp: 3_000, + step: 'workspace_setup', + message: 'Preparing environment', + action: 'attempt_started', + }); + + expect(state.getPreparationAttempts()[0]).toMatchObject({ + status: terminal.expectedStatus, + revision: 2, + ...(terminal.safeError === undefined ? {} : { safeError: terminal.safeError }), + }); + expect(state.getCloudStatus()).toEqual(terminal.expectedCloudStatus); + }); + + it('retains late terminal step snapshots without returning to preparing', () => { + const state = createServiceState(makeConfig()); + const base = { + type: 'preparing' as const, + version: 2 as const, + attemptId: 'attempt-1', + triggerMessageId: 'message-1', + }; + + state.process({ + ...base, + revision: 1, + timestamp: 1_000, + step: 'workspace_setup', + message: 'Preparing environment', + action: 'attempt_started', + }); + state.process({ + ...base, + revision: 2, + timestamp: 2_000, + step: 'ready', + message: 'Preparation complete', + action: 'attempt_completed', + }); + state.process({ + ...base, + revision: 3, + timestamp: 3_000, + step: 'setup_commands', + message: 'Preparation snapshot', + action: 'step_snapshot', + stepId: 'command:install', + stepSnapshot: { + id: 'command:install', + key: 'setup_commands', + kind: 'setup_command', + label: 'Install dependencies', + status: 'completed', + startedAt: 2_500, + completedAt: 3_000, + revision: 3, + outputTail: 'Installed dependencies', + }, + }); + + expect(state.getPreparationAttempts()[0]).toMatchObject({ + status: 'completed', + revision: 3, + steps: [ + expect.objectContaining({ id: 'command:install', outputTail: 'Installed dependencies' }), + ], + }); + expect(state.getCloudStatus()).toEqual({ type: 'ready' }); }); it('replayed snapshots of a completed attempt leave cloudStatus ready', () => { @@ -1756,6 +1903,11 @@ describe('createServiceState', () => { type: 'error', message: 'Environment preparation failed', }); + expect(state.getPendingMessages().get('m1')).toEqual({ + status: 'failed', + error: 'Environment preparation failed', + reason: 'exhausted', + }); }); it('an interrupt during preparation clears the preparing status', () => { diff --git a/packages/cloud-agent-sdk/src/service-state.ts b/packages/cloud-agent-sdk/src/service-state.ts index 89fa27cdc7..e7615025b4 100644 --- a/packages/cloud-agent-sdk/src/service-state.ts +++ b/packages/cloud-agent-sdk/src/service-state.ts @@ -413,7 +413,14 @@ function createServiceState(config: ServiceStateConfig): ServiceState { : step ); if (event.action === 'attempt_started') { - if (existing && existing.revision >= event.revision) return null; + if ( + existing && + (existing.revision >= event.revision || + existing.status === 'completed' || + existing.status === 'failed') + ) { + return null; + } const handedOff = existing?.status === 'running' ? existing : undefined; const attempt: PreparationAttempt = { id: event.attemptId, diff --git a/packages/cloud-agent-sdk/src/session-manager.test.ts b/packages/cloud-agent-sdk/src/session-manager.test.ts index a558d8a222..a9669fa7d3 100644 --- a/packages/cloud-agent-sdk/src/session-manager.test.ts +++ b/packages/cloud-agent-sdk/src/session-manager.test.ts @@ -968,6 +968,67 @@ describe('createSessionManager', () => { ).toBeNull(); }); + it('restores sending after a settled preparation failure without clearing its error', async () => { + let subscriptionCallback = (): void => { + throw new Error('Expected service state subscription callback'); + }; + let cloudStatus: CloudStatus | null = null; + mockSession.state.getCloudStatus.mockImplementation(() => cloudStatus); + mockSession.state.subscribe.mockImplementation(callback => { + subscriptionCallback = callback; + callback(); + return () => {}; + }); + mockSession.send.mockResolvedValue(undefined); + + const config = createMockConfig(); + const mgr = createSessionManager(config); + + await mgr.switchSession(kiloId('ses-1')); + await mgr.send({ + payload: { + type: 'prompt', + prompt: 'Failed preparation', + mode: 'code', + model: 'test-model', + }, + }); + const failedMessageId = mockSession.send.mock.calls[0]?.[0].messageId; + + cloudStatus = { type: 'preparing', message: 'Setting up environment...' }; + subscriptionCallback(); + expect(atomValue(config.store, mgr.atoms.canSend)).toBe(false); + + cloudStatus = { type: 'finalizing', message: 'Wrapping up...' }; + subscriptionCallback(); + expect(atomValue(config.store, mgr.atoms.canSend)).toBe(false); + + cloudStatus = { type: 'error', message: 'Clone failed' }; + subscriptionCallback(); + expect(atomValue(config.store, mgr.atoms.canSend)).toBe(true); + expect(atomValue(config.store, mgr.atoms.cloudStatus)).toEqual( + cloudStatus + ); + expect( + atomValue<{ type: string; message: string } | null>(config.store, mgr.atoms.statusIndicator) + ).toEqual(expect.objectContaining({ type: 'error', message: 'Clone failed' })); + + const accepted = await mgr.send({ + payload: { type: 'prompt', prompt: 'Retry preparation', mode: 'code', model: 'test-model' }, + }); + const retryMessageId = mockSession.send.mock.calls[1]?.[0].messageId; + expect(accepted).toBe(true); + expect(retryMessageId).toEqual(expect.stringMatching(/^msg_/)); + expect(retryMessageId).not.toBe(failedMessageId); + + mockSession.canSend = false; + subscriptionCallback(); + expect(atomValue(config.store, mgr.atoms.canSend)).toBe(false); + + mockSessionCallbacks.onResolved?.({ type: 'read-only', kiloSessionId: kiloId('ses-1') }); + expect(atomValue(config.store, mgr.atoms.canSend)).toBe(false); + }); + it('exposes active session type and remote model state from the live transport', async () => { const config = createMockConfig(); const mgr = createSessionManager(config); diff --git a/packages/cloud-agent-sdk/src/session-manager.ts b/packages/cloud-agent-sdk/src/session-manager.ts index 59d5544945..6a59144183 100644 --- a/packages/cloud-agent-sdk/src/session-manager.ts +++ b/packages/cloud-agent-sdk/src/session-manager.ts @@ -1200,7 +1200,8 @@ function createSessionManager(config: SessionManagerConfig): SessionManager { function updateCapabilityAtoms(session: CloudAgentSession): void { const cloudStatus = store.get(cloudStatusAtom); - const cloudReady = cloudStatus === null || cloudStatus.type === 'ready'; + const cloudReady = + cloudStatus === null || cloudStatus.type === 'ready' || cloudStatus.type === 'error'; const liveCanSend = session.canSend && cloudReady; if (postInterruptUnlock) { if (liveCanSend) { @@ -2156,7 +2157,7 @@ function createSessionManager(config: SessionManagerConfig): SessionManager { */ function restoreAfterInterrupt(session: CloudAgentSession): void { const cs = store.get(cloudStatusAtom); - const cloudReady = cs === null || cs.type === 'ready'; + const cloudReady = cs === null || cs.type === 'ready' || cs.type === 'error'; const readOnly = activeSessionType === 'read-only'; postInterruptUnlock = !readOnly; store.set(isStreamingAtom, false); diff --git a/services/cloud-agent-next/DEBUG.md b/services/cloud-agent-next/DEBUG.md index c10ccaecfa..b50ff2690a 100644 --- a/services/cloud-agent-next/DEBUG.md +++ b/services/cloud-agent-next/DEBUG.md @@ -167,17 +167,18 @@ The internal `getWrapperLogs` path also discovers these sandbox-side files direc ## Control-plane Diagnostics -Control-plane (`workspace_*`) sessions use verbose structured wrapper diagnostics, not the legacy raw wrapper/Kilo tarball. Records include heartbeat attempts, feed freshness, control socket and request outcomes, event-send metadata, task phases, and retirement causes. They exclude prompts, assistant/tool content, raw errors, credentials, and URLs. +Control-plane (`workspace_*`) sessions use verbose structured wrapper diagnostics plus a bounded wrapper/Kilo file archive. JSON records include heartbeat attempts, feed freshness, control socket and request outcomes, event-send metadata, task phases, and retirement causes. They exclude prompts, assistant/tool content, raw errors, credentials, and URLs. The file archive can contain prompts and code, same as legacy session tarballs. -The wrapper uploads JSON batches to the existing R2 bucket every five seconds and when a batch fills. Shutdown attempts a final flush within the existing shutdown deadline. R2 keys are: +The wrapper uploads JSON batches to the existing R2 bucket every five seconds and when a batch fills. It also uploads `/tmp/kilocode-control-wrapper.log` plus worktree Kilo log dirs (`/tmp/kilo-worktrees//.local/share/kilo/log`) as one gzip archive every 30 seconds, and again during shutdown after the JSON flush. R2 keys are: ```text logs/control////.json +logs/control////files.tar.gz ``` `sandboxId` is the logical SandboxControl ID, not the `workspace_*` session ID or physical provider allocation name. Use Worker logs to correlate these IDs. Each allocation/wrapper has separate immutable batches; sort them by the batch `sequence` and record `timestamp`, not the random batch ID. Check `droppedRecords` and `droppedTerminalRecords` for buffer overflow or rejected diagnostic records. -List and download these JSON batches directly from R2 using local tooling and the key prefix above. Uploaded batches remain available after the container disappears, subject to the bucket's retention policy. The legacy `getWrapperLogs` live-file reader and tarball retrieval do not read these JSON archives. +List and download these JSON batches and the overwriteable `files.tar.gz` object from R2 using local tooling and the key prefix above. Each wrapper incarnation keeps one file archive; later uploads replace it. Uploaded objects remain available after the container disappears, subject to the bucket's retention policy. The legacy `getWrapperLogs` live-file reader and tarball retrieval do not read these control-plane archives. Worker/DO diagnostics remain in Cloudflare logs/Axiom, not these wrapper archives. Successful `message.part.delta` forwarding is summarized in heartbeat counters and peak queue/RPC/total forwarding times instead of per-frame Worker logs. These counters and peaks reset when the DO is reconstructed. Failure and lifecycle records remain verbose, and wrapper archive logging is unchanged. Upload result markers on wrapper stderr distinguish HTTP rejection, network failure, timeout, and acceptance. An upload-only grant expires four hours after allocation launch and is not renewed; runtime credential revocation does not revoke it. Grant expiry does not delete archives. R2 retention remains governed by external bucket policy, not the session/report cleanup jobs. diff --git a/services/cloud-agent-next/src/persistence/SandboxControl.ts b/services/cloud-agent-next/src/persistence/SandboxControl.ts index 08ac2c3eee..6617d9963d 100644 --- a/services/cloud-agent-next/src/persistence/SandboxControl.ts +++ b/services/cloud-agent-next/src/persistence/SandboxControl.ts @@ -400,17 +400,10 @@ export class SandboxControl extends DurableObject { onHandshakeComplete: (identity, runtime) => this.onHandshakeComplete(identity, runtime), onReady: identity => this.onWrapperReady(identity), onHeartbeat: (payload, identity) => this.onHeartbeat(payload, identity), - onSessionEvent: (sessionIdentity, payload, identity, receiptId, receiptHash, sequence) => - this.onSessionEvent(sessionIdentity, payload, identity, receiptId, receiptHash, sequence), - onSessionPreparing: (sessionIdentity, payload, identity, receiptId, receiptHash, sequence) => - this.onSessionPreparing( - sessionIdentity, - payload, - identity, - receiptId, - receiptHash, - sequence - ), + onSessionEvent: (sessionIdentity, payload, identity, receiptId, sequence) => + this.onSessionEvent(sessionIdentity, payload, identity, receiptId, sequence), + onSessionPreparing: (sessionIdentity, payload, identity, receiptId, sequence) => + this.onSessionPreparing(sessionIdentity, payload, identity, receiptId, sequence), onOperationResult: (session, delivery, identity) => this.onOperationResult(session, delivery, identity), onNativeRuntimeRetired: (payload, identity) => this.onNativeRuntimeRetired(payload, identity), @@ -1559,6 +1552,7 @@ export class SandboxControl extends DurableObject { if (acquisition) assertAcquisitionDeadline(acquisition); this.logDiagnostic('allocation_launch', { allocationId: intent.intentId, + physicalSandboxId: intent.allocationName, phase, result: 'started', }); @@ -1576,6 +1570,7 @@ export class SandboxControl extends DurableObject { ); this.logDiagnostic('allocation_launch', { allocationId: intent.intentId, + physicalSandboxId: intent.allocationName, phase, result: 'providerRef' in created ? 'completed' : 'unresolved', durationMs: Date.now() - startedAt, @@ -1602,6 +1597,7 @@ export class SandboxControl extends DurableObject { startedAt = Date.now(); this.logDiagnostic('allocation_launch', { allocationId: intent.intentId, + physicalSandboxId: intent.allocationName, phase, result: 'started', }); @@ -1615,6 +1611,7 @@ export class SandboxControl extends DurableObject { ); this.logDiagnostic('allocation_launch', { allocationId: intent.intentId, + physicalSandboxId: intent.allocationName, phase, result: 'completed', durationMs: Date.now() - startedAt, @@ -1628,6 +1625,7 @@ export class SandboxControl extends DurableObject { phase, durationMs: Date.now() - startedAt, allocationId: physical.createIntent?.intentId, + physicalSandboxId: physical.createIntent?.allocationName, }, 'warn' ); @@ -3425,7 +3423,6 @@ export class SandboxControl extends DurableObject { payload: SessionEventPayload, connection: SandboxControlConnectionIdentity, receiptId?: string, - receiptHash?: string, sequence?: number ): Promise { const diagnostic = { @@ -3444,7 +3441,7 @@ export class SandboxControl extends DurableObject { identity, payload.type, connection, - { identity, payload, ...(receiptId ? { receiptId, receiptHash, sequence } : {}) }, + { identity, payload, ...(receiptId ? { receiptId, sequence } : {}) }, (route, fields, physical) => this.forwardSessionFrame( route, @@ -3457,7 +3454,7 @@ export class SandboxControl extends DurableObject { identity, payload, wrapperInstanceId: connection.wrapperInstanceId, - ...(receiptId ? { receiptId, receiptHash, sequence } : {}), + ...(receiptId ? { receiptId, sequence } : {}), }), receiptId !== undefined ) @@ -3474,7 +3471,6 @@ export class SandboxControl extends DurableObject { payload: SessionPreparingPayload, connection: SandboxControlConnectionIdentity, receiptId?: string, - receiptHash?: string, sequence?: number ): Promise { const diagnostic = { @@ -3493,7 +3489,7 @@ export class SandboxControl extends DurableObject { identity, 'session.preparing', connection, - { identity, payload, ...(receiptId ? { receiptId, receiptHash, sequence } : {}) }, + { identity, payload, ...(receiptId ? { receiptId, sequence } : {}) }, (route, fields, physical) => this.forwardSessionFrame( route, @@ -3506,7 +3502,7 @@ export class SandboxControl extends DurableObject { identity, payload, wrapperInstanceId: connection.wrapperInstanceId, - ...(receiptId ? { receiptId, receiptHash, sequence } : {}), + ...(receiptId ? { receiptId, sequence } : {}), }), receiptId !== undefined ) @@ -4097,6 +4093,7 @@ export class SandboxControl extends DurableObject { const stale = !sameAllocation(current, physical) || current.state === 'stopped'; this.logDiagnostic('provider_observation', { allocationId: physical.createIntent?.intentId, + physicalSandboxId: physical.createIntent?.allocationName, physicalState: physical.state, observation: result.status, result: timedOut ? 'timed_out' : failed ? 'failed' : 'completed', @@ -4379,6 +4376,7 @@ export class SandboxControl extends DurableObject { to.stopTombstone !== null || (to.state !== 'creating' && to.state !== 'running'); this.logDiagnostic('physical_committed', { allocationId: to.createIntent?.intentId ?? from.createIntent?.intentId, + physicalSandboxId: to.createIntent?.allocationName ?? from.createIntent?.allocationName, wrapperInstanceId, fromState: from.state, toState: to.state, diff --git a/services/cloud-agent-next/src/router/handlers/session-management.ts b/services/cloud-agent-next/src/router/handlers/session-management.ts index 11e29333bd..04f575c352 100644 --- a/services/cloud-agent-next/src/router/handlers/session-management.ts +++ b/services/cloud-agent-next/src/router/handlers/session-management.ts @@ -286,7 +286,10 @@ export function createSessionManagementHandlers() { }); try { - const getStub = () => resolveSessionStub(env, userId, sessionId); + const getStub = () => + sessionPlaneFromId(sessionId) === 'control' + ? getSandboxSessionStub(env, userId, sessionId) + : resolveSessionStub(env, userId, sessionId); return await withDORetry( getStub, stub => stub.cancelQueuedMessage(input.messageId), diff --git a/services/cloud-agent-next/src/router/handlers/session-queue.test.ts b/services/cloud-agent-next/src/router/handlers/session-queue.test.ts new file mode 100644 index 0000000000..3ee71a60d0 --- /dev/null +++ b/services/cloud-agent-next/src/router/handlers/session-queue.test.ts @@ -0,0 +1,107 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { TRPCError } from '@trpc/server'; +import { fetchRequestHandler } from '@trpc/server/adapters/fetch'; +import { router } from '../auth.js'; +import { createSessionManagementHandlers } from './session-management.js'; +import { createSessionSendHandlers } from './session-send.js'; +import { requireCurrentSessionAccess } from '../../session-access.js'; +import type { Env } from '../../types.js'; + +vi.mock('@cloudflare/sandbox', () => ({ getSandbox: vi.fn() })); +vi.mock('../../session-access.js', () => ({ requireCurrentSessionAccess: vi.fn() })); + +const api = router({ ...createSessionSendHandlers(), ...createSessionManagementHandlers() }); +const messageId = 'msg_000000000001AbCdEfGhIjKlMn'; + +function fixture() { + const control = { + cancelQueuedMessage: vi.fn().mockResolvedValue({ dropped: true }), + admitSubmittedMessage: vi.fn(), + }; + const legacy = { cancelQueuedMessage: vi.fn().mockResolvedValue({ dropped: false }) }; + const controlNamespace = { idFromName: vi.fn(name => name), get: vi.fn(() => control) }; + const legacyNamespace = { idFromName: vi.fn(name => name), get: vi.fn(() => legacy) }; + const env = { + SANDBOX_SESSION: controlNamespace, + CLOUD_AGENT_SESSION: legacyNamespace, + } as unknown as Env; + const call = (path: string, input: unknown, authenticated = true) => { + const request = new Request(`https://queue.test/trpc/${path}`, { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify(input), + }); + return fetchRequestHandler({ + endpoint: '/trpc', + req: request, + router: api, + createContext: () => ({ + env, + request, + userId: authenticated ? 'owner' : '', + authToken: authenticated ? 'fixture-token' : '', + }), + }); + }; + return { control, legacy, controlNamespace, legacyNamespace, call }; +} + +beforeEach(() => { + vi.mocked(requireCurrentSessionAccess).mockReset(); +}); + +describe('public queue RPC routing', () => { + it.each(['workspace', 'agent'])( + 'routes %s cancellation through its own plane without changing the result', + async plane => { + const { call, control, legacy, controlNamespace, legacyNamespace } = fixture(); + const sessionId = `${plane}_11111111-1111-4111-8111-111111111111`; + const response = await call('cancelQueuedMessage', { sessionId, messageId }); + expect(response.status).toBe(200); + expect(await response.json()).toMatchObject({ + result: { data: { dropped: plane === 'workspace' } }, + }); + const selected = plane === 'workspace' ? control : legacy; + const unused = plane === 'workspace' ? legacy : control; + expect(selected.cancelQueuedMessage).toHaveBeenCalledExactlyOnceWith(messageId); + expect(unused.cancelQueuedMessage).not.toHaveBeenCalled(); + expect( + (plane === 'workspace' ? controlNamespace : legacyNamespace).idFromName + ).toHaveBeenCalledWith(`owner:${sessionId}`); + } + ); + + it('projects the actual control admission failure contract to HTTP 429', async () => { + const { control, call } = fixture(); + control.admitSubmittedMessage.mockResolvedValue({ + success: false, + code: 'PENDING_QUEUE_FULL', + error: 'Pending message queue is full (10)', + }); + const response = await call('send', { + cloudAgentSessionId: 'workspace_11111111-1111-4111-8111-111111111111', + message: { id: messageId, prompt: 'overflow' }, + agent: { mode: 'code', model: 'test/model' }, + }); + expect(response.status).toBe(429); + expect(await response.json()).toMatchObject({ + error: { data: { code: 'TOO_MANY_REQUESTS', clientError: { retryable: true } } }, + }); + expect(control.admitSubmittedMessage).toHaveBeenCalledExactlyOnceWith( + expect.objectContaining({ + userId: 'owner', + turn: { type: 'prompt', id: messageId, prompt: 'overflow' }, + }) + ); + }); + + it('never resolves a mutation receiver for unauthenticated or forbidden access', async () => { + const { call, controlNamespace, legacyNamespace } = fixture(); + const input = { sessionId: 'workspace_11111111-1111-4111-8111-111111111111', messageId }; + expect((await call('cancelQueuedMessage', input, false)).status).toBe(401); + vi.mocked(requireCurrentSessionAccess).mockRejectedValue(new TRPCError({ code: 'FORBIDDEN' })); + expect((await call('cancelQueuedMessage', input)).status).toBe(403); + expect(controlNamespace.get).not.toHaveBeenCalled(); + expect(legacyNamespace.get).not.toHaveBeenCalled(); + }); +}); diff --git a/services/cloud-agent-next/src/sandbox-control/diagnostics.test.ts b/services/cloud-agent-next/src/sandbox-control/diagnostics.test.ts new file mode 100644 index 0000000000..d225eb5b9d --- /dev/null +++ b/services/cloud-agent-next/src/sandbox-control/diagnostics.test.ts @@ -0,0 +1,70 @@ +import { afterEach, describe, expect, it, vi } from 'vitest'; +import { CONTROL_DIAGNOSTIC_COALESCE_LIMIT, logControlDiagnostic } from './diagnostics.js'; +import { logger } from '../logger.js'; + +describe('logControlDiagnostic', () => { + const withFields = vi.spyOn(logger, 'withFields').mockReturnValue(logger); + const info = vi.spyOn(logger, 'info').mockImplementation(() => undefined); + + afterEach(() => { + withFields.mockClear(); + info.mockClear(); + }); + + it('keeps safe preparation fields and coalesces identical rejection results', () => { + const identity = `test:${crypto.randomUUID()}`; + const fields = { + sessionId: 'workspace_11111111-1111-4111-8111-111111111111', + receiptId: '22222222-2222-4222-8222-222222222222', + attemptId: 'attempt_1', + action: 'attempt_started', + revision: 4, + disposition: 'runtime_mismatch', + applied: false, + durationMs: 1, + }; + + logControlDiagnostic('session_preparing_result', fields, 'info', { + coalesceIdentity: identity, + }); + logControlDiagnostic('session_preparing_result', { ...fields, durationMs: 99 }, 'info', { + coalesceIdentity: identity, + }); + logControlDiagnostic( + 'session_preparing_result', + { ...fields, disposition: 'native_runtime_mismatch' }, + 'info', + { coalesceIdentity: identity } + ); + + expect(withFields).toHaveBeenCalledTimes(3); + expect(withFields.mock.calls[0]?.[0]).toMatchObject({ + attemptId: 'attempt_1', + action: 'attempt_started', + revision: 4, + disposition: 'runtime_mismatch', + }); + expect(withFields.mock.calls[1]?.[0]).toMatchObject({ + disposition: 'runtime_mismatch', + occurrences: 2, + }); + expect(withFields.mock.calls[2]?.[0]).toMatchObject({ + disposition: 'native_runtime_mismatch', + }); + }); + + it('evicts the oldest coalescing identity at the fixed bound', () => { + const prefix = `eviction:${crypto.randomUUID()}:`; + const fields = { applied: false, disposition: 'receipt_conflict' }; + for (let index = 0; index <= CONTROL_DIAGNOSTIC_COALESCE_LIMIT; index += 1) { + logControlDiagnostic('session_event_result', fields, 'info', { + coalesceIdentity: `${prefix}${index}`, + }); + } + withFields.mockClear(); + logControlDiagnostic('session_event_result', fields, 'info', { + coalesceIdentity: `${prefix}0`, + }); + expect(withFields).toHaveBeenCalledTimes(1); + }); +}); diff --git a/services/cloud-agent-next/src/sandbox-control/diagnostics.ts b/services/cloud-agent-next/src/sandbox-control/diagnostics.ts index 11d40d6abe..0c0f6d2668 100644 --- a/services/cloud-agent-next/src/sandbox-control/diagnostics.ts +++ b/services/cloud-agent-next/src/sandbox-control/diagnostics.ts @@ -2,6 +2,13 @@ import { withDORetry, type DORetryConfig } from '@kilocode/worker-utils'; import { logger } from '../logger.js'; export type ControlDiagnosticFields = Record; +type ControlDiagnosticOptions = { coalesceIdentity?: string }; + +export const CONTROL_DIAGNOSTIC_COALESCE_LIMIT = 128; +const coalescedDiagnostics = new Map< + string, + { fields: ControlDiagnosticFields; stableFields: string; occurrences: number } +>(); const EVENT_TYPES = new Set([ 'sandbox.ready', @@ -75,7 +82,8 @@ const DELTA_PROGRESS_EVENTS = new Set([ export function logControlDiagnostic( event: string, fields: ControlDiagnosticFields, - level: 'info' | 'warn' = 'info' + level: 'info' | 'warn' = 'info', + options?: ControlDiagnosticOptions ): void { try { if ( @@ -100,12 +108,39 @@ export function logControlDiagnostic( bounded[key] = value; } } - const scoped = logger.withFields({ - ...bounded, - logTag: 'sandbox_control', - diagnosticEvent: /^[a-z_]{1,64}$/.test(event) ? event : 'unknown', + const emit = (diagnosticFields: ControlDiagnosticFields) => { + const scoped = logger.withFields({ + ...diagnosticFields, + logTag: 'sandbox_control', + diagnosticEvent: /^[a-z_]{1,64}$/.test(event) ? event : 'unknown', + }); + scoped[level]('Sandbox control diagnostic'); + }; + if (!options?.coalesceIdentity) { + emit(bounded); + return; + } + const stableFields = JSON.stringify( + Object.entries(bounded) + .filter(([key]) => key !== 'durationMs' && key !== 'occurrences') + .sort(([left], [right]) => left.localeCompare(right)) + ); + const previous = coalescedDiagnostics.get(options.coalesceIdentity); + if (previous?.stableFields === stableFields) { + previous.occurrences = Math.min(Number.MAX_SAFE_INTEGER, previous.occurrences + 1); + return; + } + if (previous) emit({ ...previous.fields, occurrences: previous.occurrences }); + else if (coalescedDiagnostics.size >= CONTROL_DIAGNOSTIC_COALESCE_LIMIT) { + const oldest = coalescedDiagnostics.keys().next().value; + if (oldest !== undefined) coalescedDiagnostics.delete(oldest); + } + coalescedDiagnostics.set(options.coalesceIdentity, { + fields: bounded, + stableFields, + occurrences: 1, }); - scoped[level]('Sandbox control diagnostic'); + emit(bounded); } catch { return; } diff --git a/services/cloud-agent-next/src/sandbox-control/lifecycle.test.ts b/services/cloud-agent-next/src/sandbox-control/lifecycle.test.ts index 18bdaea787..ef0e75b28b 100644 --- a/services/cloud-agent-next/src/sandbox-control/lifecycle.test.ts +++ b/services/cloud-agent-next/src/sandbox-control/lifecycle.test.ts @@ -1,7 +1,5 @@ -import { createHash } from 'node:crypto'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { SandboxSession } from '../sandbox-session/SandboxSession.js'; -import { canonicalControlEventJson } from '../shared/control-event-canonical.js'; import { createMemoryEventQueries } from '../session/preparation-test-helpers.js'; import type { BillingContext } from '@kilocode/container-usage'; import { SandboxControl, type SandboxAcquisition } from '../persistence/SandboxControl.js'; @@ -440,6 +438,38 @@ afterEach(() => { }); describe('SandboxControl lifecycle boundaries', () => { + it('logs logical and derived physical allocation identities', async () => { + const h = await harness(); + const fields = vi.spyOn(logger, 'withFields').mockReturnValue(logger); + try { + await h.create(); + const physical = await h.control.getPhysicalRecord(); + expect(fields).toHaveBeenCalledWith( + expect.objectContaining({ + diagnosticEvent: 'allocation_launch', + sandboxId: SANDBOX_ID, + physicalSandboxId: physical.createIntent?.allocationName, + }) + ); + expect(fields).toHaveBeenCalledWith( + expect.objectContaining({ + diagnosticEvent: 'physical_committed', + sandboxId: SANDBOX_ID, + physicalSandboxId: physical.createIntent?.allocationName, + }) + ); + await h.control.beginStop('idle'); + await h.control.confirmStopped(); + await h.create(); + const replacement = await h.control.getPhysicalRecord(); + expect(replacement.createIntent?.allocationName).not.toBe( + physical.createIntent?.allocationName + ); + } finally { + fields.mockRestore(); + } + }); + it('recovers one eligible wrapper runtime through the production socket hooks', async () => { const h = await harness(); h.session.getControlState.mockResolvedValue({ @@ -3630,14 +3660,11 @@ describe('SandboxControl lifecycle boundaries', () => { sequence: 1, }; const receiptId = '33333333-3333-4333-8333-333333333333'; - const receiptHash = createHash('sha256') - .update(canonicalControlEventJson(publication)) - .digest('hex'); const frame = { type: 'request', requestId: 'pending-native-attach', operation: 'sandbox.event.publish', - payload: { ...publication, receiptId, receiptHash }, + payload: { ...publication, receiptId }, }; let attachment: unknown = { ...connection, @@ -3675,7 +3702,6 @@ describe('SandboxControl lifecycle boundaries', () => { payload: publication.payload, wrapperInstanceId: connection.wrapperInstanceId, receiptId, - receiptHash, sequence: 1, }); expect(send).toHaveBeenLastCalledWith( @@ -3845,9 +3871,6 @@ describe('SandboxControl lifecycle boundaries', () => { payload: { ...publication, receiptId, - receiptHash: createHash('sha256') - .update(canonicalControlEventJson(publication)) - .digest('hex'), }, }; const beforeControl = structuredClone([...h.records]); @@ -3897,9 +3920,6 @@ describe('SandboxControl lifecycle boundaries', () => { payload: { ...stalePublication, receiptId: '44444444-4444-4444-8444-444444444444', - receiptHash: createHash('sha256') - .update(canonicalControlEventJson(stalePublication)) - .digest('hex'), }, }) ); @@ -3936,9 +3956,6 @@ describe('SandboxControl lifecycle boundaries', () => { payload: { ...replacementPublication, receiptId: replacementReceiptId, - receiptHash: createHash('sha256') - .update(canonicalControlEventJson(replacementPublication)) - .digest('hex'), }, }) ); @@ -3988,7 +4005,6 @@ describe('SandboxControl lifecycle boundaries', () => { properties: { messageId: 'message_A', status: 'completed' }, }; const receiptId = '33333333-3333-4333-8333-333333333333'; - const receiptHash = 'a'.repeat(64); h.session.receiveSandboxControlEvent.mockResolvedValueOnce({ applied }); const before = structuredClone([...h.records]); h.sendRequest.mockClear(); @@ -3996,7 +4012,7 @@ describe('SandboxControl lifecycle boundaries', () => { const closeHandshakenSockets = vi.spyOn(h.socket, 'closeHandshakenSockets').mockClear(); await expect( - h.hooks.onSessionEvent?.(identity, payload, connection, receiptId, receiptHash, 1) + h.hooks.onSessionEvent?.(identity, payload, connection, receiptId, 1) ).resolves.toEqual(applied ? { applied: true } : { applied: false, retryable: false }); await h.flush(); @@ -4005,7 +4021,6 @@ describe('SandboxControl lifecycle boundaries', () => { payload, wrapperInstanceId: connection.wrapperInstanceId, receiptId, - receiptHash, sequence: 1, }); expect([...h.records]).toEqual(before); @@ -4046,7 +4061,6 @@ describe('SandboxControl lifecycle boundaries', () => { }, connection, '33333333-3333-4333-8333-333333333333', - 'a'.repeat(64), 1 ) ).resolves.toEqual({ applied: false, retryable: true }); @@ -4068,21 +4082,13 @@ describe('SandboxControl lifecycle boundaries', () => { const identity = { directory: ROUTE.directory, kiloSessionId: ROUTE.kiloSessionId }; const payload = { type: 'message.updated', properties: { id: 'message_A' } }; const receiptId = '33333333-3333-4333-8333-333333333333'; - const receiptHash = 'a'.repeat(64); h.session.receiveSandboxControlEvent.mockRejectedValue( Object.assign(new Error('Session temporarily unavailable'), { retryable: true }) ); const before = structuredClone([...h.records]); h.sendRequest.mockClear(); - const first = h.hooks.onSessionEvent?.( - identity, - payload, - connection, - receiptId, - receiptHash, - 1 - ); + const first = h.hooks.onSessionEvent?.(identity, payload, connection, receiptId, 1); await vi.advanceTimersByTimeAsync(1_000); await expect(first).resolves.toEqual({ applied: false, retryable: true }); await h.flush(); @@ -4097,7 +4103,7 @@ describe('SandboxControl lifecycle boundaries', () => { h.session.receiveSandboxControlEvent.mockResolvedValueOnce({ applied: true }); await expect( - h.hooks.onSessionEvent?.(identity, payload, connection, receiptId, receiptHash, 1) + h.hooks.onSessionEvent?.(identity, payload, connection, receiptId, 1) ).resolves.toEqual({ applied: true }); await h.flush(); expect(h.session.receiveSandboxControlEvent).toHaveBeenCalledTimes(4); @@ -4131,7 +4137,6 @@ describe('SandboxControl lifecycle boundaries', () => { payload, connection, result === 'legacy_transport_failed' ? undefined : '33333333-3333-4333-8333-333333333333', - 'a'.repeat(64), 1 ); await vi.advanceTimersByTimeAsync(0); @@ -4141,7 +4146,6 @@ describe('SandboxControl lifecycle boundaries', () => { payload, connection, '44444444-4444-4444-8444-444444444444', - 'b'.repeat(64), 2 ); await vi.advanceTimersByTimeAsync(0); diff --git a/services/cloud-agent-next/src/sandbox-control/log-routes.test.ts b/services/cloud-agent-next/src/sandbox-control/log-routes.test.ts index f58aa647f8..40c284344c 100644 --- a/services/cloud-agent-next/src/sandbox-control/log-routes.test.ts +++ b/services/cloud-agent-next/src/sandbox-control/log-routes.test.ts @@ -1,11 +1,30 @@ -import { describe, expect, it, vi } from 'vitest'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; import { Hono } from 'hono'; import type { HonoContext } from '../hono-context.js'; import type { Env } from '../types.js'; -import { CONTROL_LOG_MAX_BATCH_BYTES } from '../shared/control-diagnostics.js'; +import { + CONTROL_LOG_ARCHIVE_NAME, + CONTROL_LOG_MAX_ARCHIVE_BYTES, + CONTROL_LOG_MAX_BATCH_BYTES, + OWNED_PROCESS_CLEANUP_UNREAPED, +} from '../shared/control-diagnostics.js'; import { mintControlLogUploadGrant } from './log-upload-grant.js'; import { registerControlLogRoutes } from './log-routes.js'; +const logging = vi.hoisted(() => { + const logger = { + debug: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + withFields: vi.fn(), + }; + logger.withFields.mockReturnValue(logger); + return { logger }; +}); + +vi.mock('../logger.js', () => ({ logger: logging.logger })); + const secret = 'test-log-signing-secret'; const identity = { sandboxId: 'sandbox_test', @@ -22,13 +41,24 @@ const batch = { }; function fixture() { - const objects = new Map(); - const put = vi.fn(async (key: string, body: string, options: R2PutOptions) => { - expect(options.onlyIf).toEqual({ etagDoesNotMatch: '*' }); - if (objects.has(key)) return null; - objects.set(key, { body }); - return { key }; - }); + const objects = new Map(); + const put = vi.fn( + async (key: string, body: string | ArrayBuffer | Uint8Array, options?: R2PutOptions) => { + const stored = + typeof body === 'string' ? body : body instanceof Uint8Array ? body : new Uint8Array(body); + const condition = options?.onlyIf; + if ( + condition && + 'etagDoesNotMatch' in condition && + condition.etagDoesNotMatch === '*' && + objects.has(key) + ) { + return null; + } + objects.set(key, { body: stored, options }); + return { key }; + } + ); const env = { NEXTAUTH_SECRET: secret, R2_BUCKET: { put } } as unknown as Env; const app = new Hono(); registerControlLogRoutes(app); @@ -44,17 +74,74 @@ function fixture() { headers: { Authorization: `Bearer ${token}`, 'Content-Type': 'application/json' }, body: JSON.stringify(body), }); - return { request, upload, objects, put }; + const archivePath = `${identity.sandboxId}/${identity.allocationId}/${identity.wrapperInstanceId}/${CONTROL_LOG_ARCHIVE_NAME}`; + const uploadArchive = ( + body: Uint8Array, + path = archivePath, + token = mintControlLogUploadGrant(identity, secret), + contentType = 'application/gzip' + ) => + request(`/sandbox-logs/${path}`, { + method: 'PUT', + headers: { + Authorization: `Bearer ${token}`, + 'Content-Type': contentType, + 'Content-Length': String(body.byteLength), + }, + body, + }); + return { request, upload, uploadArchive, archivePath, objects, put }; } describe('control log routes', () => { + beforeEach(() => { + logging.logger.error.mockClear(); + logging.logger.withFields.mockClear(); + logging.logger.withFields.mockReturnValue(logging.logger); + }); + it('stores validated immutable batches without any provider or DO binding', async () => { const f = fixture(); expect((await f.upload()).status).toBe(204); const key = `logs/control/${suffix}.json`; - expect(JSON.parse(f.objects.get(key)!.body)).toEqual(batch); + const stored = f.objects.get(key)!; + expect(stored.options?.onlyIf).toEqual({ etagDoesNotMatch: '*' }); + expect(JSON.parse(stored.body as string)).toEqual(batch); expect((await f.upload({ ...batch, sequence: 42 })).status).toBe(204); - expect(JSON.parse(f.objects.get(key)!.body).sequence).toBe(0); + expect(JSON.parse(f.objects.get(key)!.body as string).sequence).toBe(0); + expect(logging.logger.error).not.toHaveBeenCalled(); + }); + + it('logs unreaped owned-process cleanup once when a new batch is stored', async () => { + const f = fixture(); + const unreaped = { + ...batch, + records: [ + { + timestamp: 100, + event: 'session.task', + fields: { + phase: 'failed', + stage: 'process_cleanup', + ok: false, + sessionId: 'workspace_test', + kiloSessionId: 'ses_test', + messageId: 'msg_test', + kind: 'execution', + detail: 'owned_process_unreaped populated=1 /workspace/test', + }, + }, + ], + }; + expect((await f.upload(unreaped)).status).toBe(204); + expect((await f.upload({ ...unreaped, sequence: 42 })).status).toBe(204); + expect(logging.logger.withFields).toHaveBeenCalledWith({ + logTag: 'owned_process_unreaped', + sessionId: 'workspace_test', + sandboxId: identity.sandboxId, + }); + expect(logging.logger.error).toHaveBeenCalledTimes(1); + expect(logging.logger.error).toHaveBeenCalledWith(OWNED_PROCESS_CLEANUP_UNREAPED); }); it('rejects cross-allocation, cross-wrapper and cross-sandbox writes', async () => { @@ -185,4 +272,58 @@ describe('control log routes', () => { expect(response.status).toBe(503); expect(await response.text()).not.toContain('private-secret'); }); + + it('stores one overwriteable gzip archive per wrapper incarnation', async () => { + const f = fixture(); + const first = new Uint8Array([0x1f, 0x8b, 1, 2, 3]); + const second = new Uint8Array([0x1f, 0x8b, 4, 5, 6]); + expect((await f.uploadArchive(first)).status).toBe(204); + expect((await f.uploadArchive(second)).status).toBe(204); + const key = `logs/control/${f.archivePath}`; + expect(f.objects.size).toBe(1); + expect(f.put).toHaveBeenCalledTimes(2); + expect(f.objects.get(key)?.options?.onlyIf).toBeUndefined(); + expect(f.objects.get(key)?.body).toEqual(second); + }); + + it('rejects json on the gzip archive route and gzip on the json route', async () => { + const f = fixture(); + expect( + ( + await f.uploadArchive( + new Uint8Array([1, 2, 3]), + f.archivePath, + mintControlLogUploadGrant(identity, secret), + 'application/json' + ) + ).status + ).toBe(415); + expect( + ( + await f.request(`/sandbox-logs/${suffix}`, { + method: 'PUT', + headers: { + Authorization: `Bearer ${mintControlLogUploadGrant(identity, secret)}`, + 'Content-Type': 'application/gzip', + }, + body: new Uint8Array([1, 2, 3]), + }) + ).status + ).toBe(415); + expect(f.put).not.toHaveBeenCalled(); + }); + + it('rejects an oversized gzip archive before writing it', async () => { + const f = fixture(); + const response = await f.request(`/sandbox-logs/${f.archivePath}`, { + method: 'PUT', + headers: { + Authorization: `Bearer ${mintControlLogUploadGrant(identity, secret)}`, + 'Content-Type': 'application/gzip', + 'Content-Length': String(CONTROL_LOG_MAX_ARCHIVE_BYTES + 1), + }, + }); + expect(response.status).toBe(413); + expect(f.put).not.toHaveBeenCalled(); + }); }); diff --git a/services/cloud-agent-next/src/sandbox-control/log-routes.ts b/services/cloud-agent-next/src/sandbox-control/log-routes.ts index ed00b0b223..44f90ed3aa 100644 --- a/services/cloud-agent-next/src/sandbox-control/log-routes.ts +++ b/services/cloud-agent-next/src/sandbox-control/log-routes.ts @@ -1,24 +1,53 @@ import type { Hono, Context } from 'hono'; import type { HonoContext } from '../hono-context.js'; import { resolveSecret } from '../auth.js'; +import { logger } from '../logger.js'; import { + CONTROL_LOG_ARCHIVE_NAME, + CONTROL_LOG_MAX_ARCHIVE_BYTES, CONTROL_LOG_MAX_BATCH_BYTES, controlLogBatchSchema, controlLogIdentitySchema, controlLogWrapperIdSchema, + isUnreapedOwnedProcessDiagnostic, + OWNED_PROCESS_CLEANUP_UNREAPED, + type ControlLogBatch, type ControlLogIdentity, } from '../shared/control-diagnostics.js'; import { validateControlLogUploadGrant } from './log-upload-grant.js'; +function reportUnreapedOwnedProcessCleanup( + identity: ControlLogIdentity, + batch: ControlLogBatch +): void { + for (const record of batch.records) { + if (!isUnreapedOwnedProcessDiagnostic(record)) continue; + logger + .withFields({ + logTag: 'owned_process_unreaped', + sessionId: record.fields.sessionId, + sandboxId: identity.sandboxId, + }) + .error(OWNED_PROCESS_CLEANUP_UNREAPED); + } +} + function archivePrefix(identity: ControlLogIdentity): string { return `logs/control/${[identity.sandboxId, identity.allocationId, identity.wrapperInstanceId] .map(encodeURIComponent) .join('/')}/`; } -async function readBoundedBody(request: Request): Promise { +function mediaType(header: string | undefined): string | undefined { + return header?.split(';')[0].trim(); +} + +async function readBoundedBytes( + request: Request, + maxBytes: number +): Promise { const stream: ReadableStream | null = request.body; - if (!stream) return ''; + if (!stream) return new Uint8Array(); const reader = stream.getReader(); const chunks: Uint8Array[] = []; let length = 0; @@ -27,7 +56,7 @@ async function readBoundedBody(request: Request): Promise { const { value, done } = await reader.read(); if (done) break; length += value.byteLength; - if (length > CONTROL_LOG_MAX_BATCH_BYTES) return undefined; + if (length > maxBytes) return undefined; chunks.push(value); } } finally { @@ -39,7 +68,13 @@ async function readBoundedBody(request: Request): Promise { bytes.set(chunk, offset); offset += chunk.byteLength; } - return new TextDecoder('utf-8', { fatal: true, ignoreBOM: false }).decode(bytes); + return bytes; +} + +function declaredLength(header: string | undefined): number | undefined | 'invalid' { + if (header === undefined) return undefined; + if (!/^\d+$/.test(header)) return 'invalid'; + return Number(header); } function routeIdentity(c: Context) { @@ -50,45 +85,90 @@ function routeIdentity(c: Context) { }); } +async function authorizeIdentity( + c: Context +): Promise<{ identity: ControlLogIdentity } | Response> { + const identity = routeIdentity(c); + if (!identity.success) return c.text('Invalid log identity', 400); + const grant = validateControlLogUploadGrant( + c.req.header('Authorization') ?? null, + await resolveSecret(c.env.NEXTAUTH_SECRET) + ); + if (!grant) return c.text('Unauthorized', 401); + if ( + grant.sandboxId !== identity.data.sandboxId || + grant.allocationId !== identity.data.allocationId || + grant.wrapperInstanceId !== identity.data.wrapperInstanceId + ) + return c.text('Log scope mismatch', 403); + return { identity: identity.data }; +} + export function registerControlLogRoutes(app: Hono): void { + app.put( + `/sandbox-logs/:sandboxId/:allocationId/:wrapperInstanceId/${CONTROL_LOG_ARCHIVE_NAME}`, + async c => { + const authorized = await authorizeIdentity(c); + if (authorized instanceof Response) return authorized; + if (mediaType(c.req.header('Content-Type')) !== 'application/gzip') { + return c.text('Expected application/gzip', 415); + } + const encoding = c.req.header('Content-Encoding'); + if (encoding && encoding !== 'identity') return c.text('Unsupported encoding', 415); + const length = declaredLength(c.req.header('Content-Length')); + if (length === 'invalid') return c.text('Invalid length', 400); + if (length !== undefined && length > CONTROL_LOG_MAX_ARCHIVE_BYTES) { + return c.text('Body too large', 413); + } + const body = await readBoundedBytes(c.req.raw, CONTROL_LOG_MAX_ARCHIVE_BYTES); + if (body === undefined) return c.text('Body too large', 413); + if (body.byteLength === 0) return c.text('Missing request body', 400); + try { + await c.env.R2_BUCKET.put( + `${archivePrefix(authorized.identity)}${CONTROL_LOG_ARCHIVE_NAME}`, + body, + { httpMetadata: { contentType: 'application/gzip' } } + ); + } catch { + return c.text('Log storage unavailable', 503); + } + return c.body(null, 204); + } + ); + app.put('/sandbox-logs/:sandboxId/:allocationId/:wrapperInstanceId/:batchId', async c => { - const identity = routeIdentity(c); + const authorized = await authorizeIdentity(c); + if (authorized instanceof Response) return authorized; const batchId = controlLogWrapperIdSchema.safeParse(c.req.param('batchId')); - if (!identity.success || !batchId.success) return c.text('Invalid log identity', 400); - const grant = validateControlLogUploadGrant( - c.req.header('Authorization') ?? null, - await resolveSecret(c.env.NEXTAUTH_SECRET) - ); - if (!grant) return c.text('Unauthorized', 401); - if ( - grant.sandboxId !== identity.data.sandboxId || - grant.allocationId !== identity.data.allocationId || - grant.wrapperInstanceId !== identity.data.wrapperInstanceId - ) - return c.text('Log scope mismatch', 403); + if (!batchId.success) return c.text('Invalid log identity', 400); - if (c.req.header('Content-Type')?.split(';')[0].trim() !== 'application/json') { + if (mediaType(c.req.header('Content-Type')) !== 'application/json') { return c.text('Expected application/json', 415); } const encoding = c.req.header('Content-Encoding'); if (encoding && encoding !== 'identity') return c.text('Unsupported encoding', 415); - const declaredLength = c.req.header('Content-Length'); - if (declaredLength && !/^\d+$/.test(declaredLength)) return c.text('Invalid length', 400); - if (Number(declaredLength) > CONTROL_LOG_MAX_BATCH_BYTES) return c.text('Body too large', 413); + const length = declaredLength(c.req.header('Content-Length')); + if (length === 'invalid') return c.text('Invalid length', 400); + if (length !== undefined && length > CONTROL_LOG_MAX_BATCH_BYTES) { + return c.text('Body too large', 413); + } - let body: unknown; + let parsed: unknown; try { - const text = await readBoundedBody(c.req.raw); - if (text === undefined) return c.text('Body too large', 413); - body = JSON.parse(text); + const bytes = await readBoundedBytes(c.req.raw, CONTROL_LOG_MAX_BATCH_BYTES); + if (bytes === undefined) return c.text('Body too large', 413); + parsed = JSON.parse( + new TextDecoder('utf-8', { fatal: true, ignoreBOM: false }).decode(bytes) + ); } catch { return c.text('Invalid log batch', 400); } - const batch = controlLogBatchSchema.safeParse(body); + const batch = controlLogBatchSchema.safeParse(parsed); if (!batch.success) return c.text('Invalid log batch', 400); + let stored: unknown; try { - await c.env.R2_BUCKET.put( - `${archivePrefix(identity.data)}${batchId.data}.json`, + stored = await c.env.R2_BUCKET.put( + `${archivePrefix(authorized.identity)}${batchId.data}.json`, JSON.stringify(batch.data), { onlyIf: { etagDoesNotMatch: '*' }, @@ -99,6 +179,7 @@ export function registerControlLogRoutes(app: Hono): void { } catch { return c.text('Log storage unavailable', 503); } + if (stored) reportUnreapedOwnedProcessCleanup(authorized.identity, batch.data); return c.body(null, 204); }); } diff --git a/services/cloud-agent-next/src/sandbox-control/socket.test.ts b/services/cloud-agent-next/src/sandbox-control/socket.test.ts index be5e720b2f..960937f4f1 100644 --- a/services/cloud-agent-next/src/sandbox-control/socket.test.ts +++ b/services/cloud-agent-next/src/sandbox-control/socket.test.ts @@ -1016,7 +1016,6 @@ describe('sandbox control socket handler', () => { wrapperInstanceId: WRAPPER_INSTANCE_ID, }); const receiptId = '123e4567-e89b-42d3-a456-426614174099'; - const receiptHash = 'a'.repeat(64); const onSessionEvent = vi.fn().mockResolvedValue({ applied: true }); const handler = createSandboxControlSocketHandler( createFakeState([ws]), @@ -1034,7 +1033,6 @@ describe('sandbox control socket handler', () => { payload: { event: 'session.event', receiptId, - receiptHash, sequence: 1, session: { directory: '/workspace/a', kiloSessionId: 'kilo_1' }, payload: { type: 'message.updated', properties: { id: 'msg_1' } }, @@ -1047,7 +1045,6 @@ describe('sandbox control socket handler', () => { { type: 'message.updated', properties: { id: 'msg_1' } }, handler.getConnectionIdentity(), receiptId, - receiptHash, 1 ); expect(JSON.parse(ws.send.mock.calls.at(-1)?.[0] as string)).toEqual({ @@ -1085,7 +1082,6 @@ describe('sandbox control socket handler', () => { payload: { event: 'session.event', receiptId: '123e4567-e89b-42d3-a456-426614174099', - receiptHash: 'a'.repeat(64), sequence: 1, session: { directory: '/workspace/a', diff --git a/services/cloud-agent-next/src/sandbox-control/socket.ts b/services/cloud-agent-next/src/sandbox-control/socket.ts index d0d984e1eb..619e473094 100644 --- a/services/cloud-agent-next/src/sandbox-control/socket.ts +++ b/services/cloud-agent-next/src/sandbox-control/socket.ts @@ -90,7 +90,6 @@ export type SandboxControlSocketHooks = { payload: SessionEventPayload, identity: SandboxControlConnectionIdentity, receiptId?: string, - receiptHash?: string, sequence?: number ): void | SandboxControlEventResult | Promise; onSessionPreparing?( @@ -98,7 +97,6 @@ export type SandboxControlSocketHooks = { payload: SessionPreparingPayload, identity: SandboxControlConnectionIdentity, receiptId?: string, - receiptHash?: string, sequence?: number ): void | SandboxControlEventResult | Promise; onOperationResult?( @@ -714,7 +712,6 @@ export function createSandboxControlSocketHandler( publication.data.payload, identity, publication.data.receiptId, - publication.data.receiptHash, publication.data.sequence ) : await hooks.onSessionPreparing?.( @@ -722,7 +719,6 @@ export function createSandboxControlSocketHandler( publication.data.payload, identity, publication.data.receiptId, - publication.data.receiptHash, publication.data.sequence ); if (!isCurrentConnection(state, ws, identity)) return; diff --git a/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts b/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts index 3c7e36c52f..86d2c1afa7 100644 --- a/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts +++ b/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts @@ -164,6 +164,7 @@ import { acceptQueuedMessage, applyMessageOutcome, assignPreparationAttemptId, + cancelPendingMessage, createSessionMessageRecord, failQueuedMessage, failWaitingMessages as applyFailWaitingMessages, @@ -180,6 +181,7 @@ import { type ControlSessionMessageInput, type SessionMessageRecord, } from './session-message-queue.js'; +import { PENDING_SESSION_MESSAGE_LIMIT } from '../session/pending-messages.js'; import { commitSessionOperationResult, dispatchSessionOperation, @@ -187,7 +189,6 @@ import { } from './session-operation.js'; import { controlEventReceiptDisposition, - hasValidControlEventReceipt, recordControlEventReceipt, bindControlEventReceiptIdentity, retireControlEventReceiptIdentity, @@ -229,6 +230,31 @@ const pendingRuntimeCleanupSchema = z.object({ type MessageRecord = SessionMessageRecord; type DispatchPhase = 'preparing' | 'attach' | 'prompt'; +type ControlEventDisposition = + | 'apply' + | 'duplicate' + | 'receipt_conflict' + | 'epoch_changed' + | 'runtime_mismatch' + | 'native_runtime_pending' + | 'native_runtime_mismatch'; +type ControlEventInput = { + identity: SessionEventIdentity; + wrapperInstanceId?: string; + receiptId?: string; + receiptHash?: string; + sequence?: number; +}; +type ControlEventEvaluationRequest = + | { + contract: 'publication_admission'; + input: ControlEventInput; + epoch: number; + publication: + | { kind: 'event' } + | { kind: 'preparing'; loadTrigger: () => SessionMessageRecord | undefined }; + } + | { contract: 'currency_recheck'; input: ControlEventInput; epoch: number }; type SandboxSessionRegistrationInput = { identity: SessionMetadata['identity']; @@ -427,6 +453,7 @@ export class SandboxSession extends DurableObject { input: SandboxControlEventInput & { wrapperInstanceId?: string } ): Promise<{ applied: boolean; retryable?: boolean }> { const startedAt = Date.now(); + const diagnosticSnapshot = this.controlEventDiagnosticSnapshot(input); const metadata = await this.getMetadata(); const epoch = this.terminalLifecycle.captureEpoch(); const result = ( @@ -434,21 +461,30 @@ export class SandboxSession extends DurableObject { disposition: string, fields: ControlDiagnosticFields = {} ) => { - logControlDiagnostic('session_event_result', { - sessionId: this.sessionId, - sandboxId: metadata?.workspace?.sandboxId, - wrapperInstanceId: input.wrapperInstanceId, - eventType: diagnosticEventType(input.payload.type), - applied, - disposition, - durationMs: Date.now() - startedAt, - ...fields, - }); + logControlDiagnostic( + 'session_event_result', + { + sessionId: this.sessionId, + sandboxId: metadata?.workspace?.sandboxId, + wrapperInstanceId: input.wrapperInstanceId, + receiptId: input.receiptId, + eventType: diagnosticEventType(input.payload.type), + applied, + disposition, + durationMs: Date.now() - startedAt, + ...diagnosticSnapshot, + ...fields, + }, + 'info', + { + ...(applied || input.receiptId === undefined + ? {} + : { coalesceIdentity: `session.event:${this.sessionId}:${input.receiptId}` }), + } + ); return { applied }; }; if (!metadata || epoch === null) return result(false, 'session_unavailable'); - if (!(await this.hasValidControlEventReceipt('session.event', input))) - return result(false, 'receipt_conflict'); const root = metadata.auth.kiloSessionId; if (input.identity.directory !== this.directory(metadata)) return result(false, 'directory_mismatch'); @@ -492,25 +528,19 @@ export class SandboxSession extends DurableObject { (root === undefined || input.identity.rootKiloSessionId !== root) ) return result(false, 'root_mismatch'); - if (!this.terminalLifecycle.isCurrent(epoch)) return result(false, 'epoch_changed'); - if (input.receiptId !== undefined) { - if (!this.isCurrentEventRuntime(input.wrapperInstanceId)) - return result(false, 'runtime_mismatch'); - const receipt = this.controlEventReceipt(input); - if (receipt === 'duplicate') return result(true, 'duplicate'); - if (receipt !== 'apply') return result(false, 'receipt_conflict'); - } - if (input.receiptId !== undefined) { - const attachment = this.nativeAttachmentEventDisposition( - input.identity, - input.wrapperInstanceId - ); - if (attachment === 'pending') - return { ...result(false, 'native_runtime_pending'), retryable: true }; - if (attachment === 'rejected') return result(false, 'native_runtime_mismatch'); - } - if (!this.isCurrentNativeEventRuntime(input.identity, input.wrapperInstanceId)) - return result(false, 'native_runtime_mismatch'); + const hasReceiptIdentity = + input.receiptId !== undefined || + input.receiptHash !== undefined || + input.sequence !== undefined; + const admission = this.evaluateControlEvent({ + contract: 'publication_admission', + input, + epoch, + publication: { kind: 'event' }, + }); + if (admission === 'native_runtime_pending') + return { ...result(false, admission), retryable: true }; + if (admission !== 'apply') return result(admission === 'duplicate', admission); const sessionId = this.requireSessionId(); if (input.payload.type === 'session.message.outcome') { const outcome = sessionMessageOutcomeSchema.safeParse(input.payload.properties); @@ -537,14 +567,21 @@ export class SandboxSession extends DurableObject { existing?.wrapperInstanceId === input.wrapperInstanceId && existing.state === outcome.data.status ) { - const receipt = this.commitControlEventReceipt( - input, - () => this.terminalLifecycle.isCurrent(epoch) && this.isCurrentReceiptRuntime(input) + const receipt = hasReceiptIdentity + ? this.commitControlEventReceipt(input, epoch) + : this.terminalLifecycle.isCurrent(epoch) + ? ('apply' as const) + : ('epoch_changed' as const); + return result( + receipt === 'apply' || receipt === 'duplicate', + receipt === 'apply' || receipt === 'duplicate' ? 'duplicate' : receipt, + diagnostic ); - return result(receipt === 'apply' || receipt === 'duplicate', 'duplicate', diagnostic); } - if (!this.isCurrentReceiptRuntime(input)) - return result(false, 'runtime_mismatch', diagnostic); + if (hasReceiptIdentity) { + const currency = this.evaluateControlEvent({ contract: 'currency_recheck', input, epoch }); + if (currency !== 'apply') return result(false, currency, diagnostic); + } const settled = applyMessageOutcome( messages, outcome.data, @@ -564,28 +601,37 @@ export class SandboxSession extends DurableObject { diagnostic ); } - let receipt: ControlEventReceiptDisposition | undefined; - const saved = this.saveMessages( - settled, - epoch, - 'wrapper_outcome', - undefined, - () => this.recordControlEventReceipt(input), - () => { - if (!this.terminalLifecycle.isCurrent(epoch) || !this.isCurrentReceiptRuntime(input)) - return false; - receipt = this.controlEventReceipt(input); - return receipt === 'apply'; - } - ); - if (!saved) { - if (receipt === 'duplicate') return result(true, 'duplicate', diagnostic); + const receipt = hasReceiptIdentity + ? this.saveMessages( + settled, + epoch, + 'wrapper_outcome', + undefined, + () => this.recordControlEventReceipt(input), + () => { + const currency = this.evaluateControlEvent({ + contract: 'currency_recheck', + input, + epoch, + }); + if (currency !== 'apply') return currency; + const current = this.controlEventReceipt(input); + return current === 'apply' + ? ('apply' as const) + : current === 'duplicate' + ? ('duplicate' as const) + : ('receipt_conflict' as const); + } + ) + : this.saveMessages(settled, epoch, 'wrapper_outcome') + ? ('apply' as const) + : ('epoch_changed' as const); + if (receipt !== 'apply') return result( - false, - receipt === undefined || receipt === 'apply' ? 'epoch_changed' : 'receipt_conflict', + receipt === 'duplicate', + receipt === 'duplicate' ? 'duplicate' : receipt, diagnostic ); - } if (this.isCurrentEventRuntime(input.wrapperInstanceId)) { this.worktreeChanges.onEvent( this.worktreeContext(metadata), @@ -601,17 +647,18 @@ export class SandboxSession extends DurableObject { } return result(true, 'outcome_applied', diagnostic); } - if (!this.isCurrentEventRuntime(input.wrapperInstanceId)) - return result(false, 'runtime_mismatch'); + const currency = this.evaluateControlEvent({ contract: 'currency_recheck', input, epoch }); + if (currency !== 'apply') return result(false, currency); if (input.payload.type === WORKTREE_CHANGED_EVENT) { if (!root || eventKiloSessionId !== root) return result(false, 'root_mismatch'); if (!input.wrapperInstanceId) return result(false, 'missing_wrapper_identity'); - const receipt = this.commitControlEventReceipt( - input, - () => this.terminalLifecycle.isCurrent(epoch) && this.isCurrentReceiptRuntime(input) - ); + const receipt = hasReceiptIdentity + ? this.commitControlEventReceipt(input, epoch) + : this.terminalLifecycle.isCurrent(epoch) + ? ('apply' as const) + : ('epoch_changed' as const); if (receipt === 'duplicate') return result(true, 'duplicate'); - if (receipt !== 'apply') return result(false, 'receipt_conflict'); + if (receipt !== 'apply') return result(false, receipt); this.worktreeChanges.onEvent( this.worktreeContext(metadata), eventKiloSessionId, @@ -626,22 +673,20 @@ export class SandboxSession extends DurableObject { message => message.state === 'accepted' || message.state === 'queued' ) ) { - if (input.receiptId === undefined) return result(false, 'no_pending_work'); - const receipt = this.commitControlEventReceipt( - input, - () => this.terminalLifecycle.isCurrent(epoch) && this.isCurrentReceiptRuntime(input) + if (!hasReceiptIdentity) return result(false, 'no_pending_work'); + const receipt = this.commitControlEventReceipt(input, epoch); + return result( + receipt === 'apply' || receipt === 'duplicate', + receipt === 'apply' || receipt === 'duplicate' ? 'no_pending_work' : receipt ); - return result(receipt === 'apply' || receipt === 'duplicate', 'no_pending_work'); } const notifications: StoredEvent[] = []; - const receipt = this.ctx.storage.transactionSync(() => { - if ( - !this.terminalLifecycle.isCurrent(epoch) || - !this.isCurrentEventRuntime(input.wrapperInstanceId) || - !this.isCurrentNativeEventRuntime(input.identity, input.wrapperInstanceId) - ) - return 'stale'; - const current = this.controlEventReceipt(input); + const receipt = this.ctx.storage.transactionSync((): ControlEventDisposition => { + const current = this.evaluateControlEvent( + hasReceiptIdentity + ? { contract: 'publication_admission', input, epoch, publication: { kind: 'event' } } + : { contract: 'currency_recheck', input, epoch } + ); if (current !== 'apply') return current; this.recordPendingInteraction(input.payload); persistSandboxControlSessionEvent({ @@ -657,7 +702,7 @@ export class SandboxSession extends DurableObject { return 'apply' as const; }); if (receipt === 'duplicate') return result(true, 'duplicate'); - if (receipt !== 'apply') return result(false, 'receipt_conflict'); + if (receipt !== 'apply') return result(false, receipt); for (const notification of notifications) this.broadcastStoredEvent(notification); this.worktreeChanges.onEvent( this.worktreeContext(metadata), @@ -723,21 +768,34 @@ export class SandboxSession extends DurableObject { receiptHash?: string; sequence?: number; }): Promise<{ applied: boolean }> { + const diagnosticSnapshot = this.controlEventDiagnosticSnapshot(input); const metadata = await this.getMetadata(); const epoch = this.terminalLifecycle.captureEpoch(); const result = (applied: boolean, disposition: string) => { - logControlDiagnostic('session_preparing_result', { - sessionId: this.sessionId, - sandboxId: metadata?.workspace?.sandboxId, - wrapperInstanceId: input.wrapperInstanceId, - applied, - disposition, - }); + logControlDiagnostic( + 'session_preparing_result', + { + sessionId: this.sessionId, + sandboxId: metadata?.workspace?.sandboxId, + wrapperInstanceId: input.wrapperInstanceId, + receiptId: input.receiptId, + attemptId: input.payload.attemptId, + action: input.payload.action, + revision: input.payload.revision, + applied, + disposition, + ...diagnosticSnapshot, + }, + 'info', + { + ...(applied || input.receiptId === undefined + ? {} + : { coalesceIdentity: `session.preparing:${this.sessionId}:${input.receiptId}` }), + } + ); return { applied }; }; if (!metadata || epoch === null) return result(false, 'session_unavailable'); - if (!(await this.hasValidControlEventReceipt('session.preparing', input))) - return result(false, 'receipt_conflict'); const root = metadata.auth.kiloSessionId; if (input.identity.directory !== this.directory(metadata)) return result(false, 'directory_mismatch'); @@ -748,45 +806,63 @@ export class SandboxSession extends DurableObject { ) { return result(false, 'root_mismatch'); } + const hasReceiptIdentity = + input.receiptId !== undefined || + input.receiptHash !== undefined || + input.sequence !== undefined; const message = this.loadMessages().find( item => item.messageId === input.payload.triggerMessageId ); - if ( - !this.terminalLifecycle.isCurrent(epoch) || - !message || - message.preparationAttemptId !== input.payload.attemptId || - (input.wrapperInstanceId !== undefined && - message.wrapperInstanceId !== input.wrapperInstanceId) - ) { - return result( - false, - !this.terminalLifecycle.isCurrent(epoch) - ? 'epoch_changed' - : !message - ? 'message_missing' - : message.preparationAttemptId !== input.payload.attemptId - ? 'attempt_mismatch' - : 'runtime_mismatch' - ); - } + if (!this.terminalLifecycle.isCurrent(epoch)) return result(false, 'epoch_changed'); + if (!message) return result(false, 'message_missing'); + if (message.preparationAttemptId !== input.payload.attemptId) + return result(false, 'attempt_mismatch'); if (message.state !== 'queued') { - const receipt = this.commitControlEventReceipt( - input, - () => this.terminalLifecycle.isCurrent(epoch) && this.isCurrentReceiptRuntime(input) + if ( + this.evaluateControlEvent({ + contract: 'publication_admission', + input, + epoch, + publication: { + kind: 'preparing', + loadTrigger: () => undefined, + }, + }) === 'duplicate' + ) + return result(true, 'duplicate'); + const receipt = hasReceiptIdentity + ? this.commitControlEventReceipt(input, epoch) + : this.terminalLifecycle.isCurrent(epoch) + ? ('apply' as const) + : ('epoch_changed' as const); + return result( + receipt === 'apply' || receipt === 'duplicate', + receipt === 'apply' || receipt === 'duplicate' ? 'already_settled' : receipt ); - return result(receipt === 'apply' || receipt === 'duplicate', 'already_settled'); } const sessionId = this.requireSessionId(); const notifications: StoredEvent[] = []; - const receipt = this.ctx.storage.transactionSync(() => { - if ( - !this.terminalLifecycle.isCurrent(epoch) || - !this.isCurrentEventRuntime(input.wrapperInstanceId) || - !this.isCurrentNativeEventRuntime(input.identity, input.wrapperInstanceId) - ) - return 'stale'; - const current = this.controlEventReceipt(input); - if (current !== 'apply') return current; + const receipt = this.ctx.storage.transactionSync((): ControlEventDisposition => { + const admission = this.evaluateControlEvent({ + contract: 'publication_admission', + input, + epoch, + publication: { + kind: 'preparing', + loadTrigger: () => { + const trigger = this.loadMessages().find( + item => item.messageId === input.payload.triggerMessageId + ); + return trigger?.state === 'queued' && + trigger.cancellation === undefined && + trigger.preparationAttemptId === input.payload.attemptId && + trigger.wrapperInstanceId === input.wrapperInstanceId + ? trigger + : undefined; + }, + }, + }); + if (admission !== 'apply' && admission !== 'native_runtime_pending') return admission; applyControlPlanePreparingEvent({ sessionId, data: input.payload, @@ -797,7 +873,7 @@ export class SandboxSession extends DurableObject { return 'apply' as const; }); if (receipt === 'duplicate') return result(true, 'duplicate'); - if (receipt !== 'apply') return result(false, 'receipt_conflict'); + if (receipt !== 'apply') return result(false, receipt); for (const notification of notifications) this.broadcastStoredEvent(notification); return result(true, 'processed'); } @@ -868,6 +944,23 @@ export class SandboxSession extends DurableObject { notifications, }); if (!ack) return undefined; + if (authorization.operation === 'session.attach') { + const attached = delivery.result.ok + ? sessionAttachResultSchema.safeParse(delivery.result.result) + : undefined; + logControlDiagnostic('session_attach_completion', { + sessionId: this.sessionId, + messageId: authorization.messageId, + attemptId: authorization.operationId, + operationId: authorization.operationId, + wrapperInstanceId: input.wrapperInstanceId, + resultState: ack.disposition, + ok: delivery.result.ok, + nativeRuntimeId: attached?.success ? attached.data.nativeRuntimeId : undefined, + errorCode: delivery.result.ok ? undefined : delivery.result.error.code, + retryable: delivery.result.ok ? undefined : delivery.result.error.retryable, + }); + } if ( authorization.operation === 'session.attach' && delivery.result.ok && @@ -1110,6 +1203,16 @@ export class SandboxSession extends DurableObject { return { state: 'reconciled' }; } + async cancelQueuedMessage(messageId: string): Promise<{ dropped: boolean }> { + const epoch = this.terminalLifecycle.captureEpoch(); + if (epoch === null || this.deletedWorktreeId) return { dropped: false }; + const result = cancelPendingMessage(this.loadMessages(), messageId); + if (!result.dropped) return { dropped: false }; + if (result.messages && !this.saveMessages(result.messages, epoch)) return { dropped: false }; + if (nextQueuedMessageId(this.loadMessages())) await this.armQueueRetry(); + return { dropped: true }; + } + async interruptExecution(): Promise<{ success: boolean; message?: string }>; async interruptExecution(input: unknown): Promise; async interruptExecution( @@ -1399,6 +1502,21 @@ export class SandboxSession extends DurableObject { ) return; this.terminalLifecycle.invalidateRuntime(input); + logControlDiagnostic('native_fence_transition', { + sessionId: this.sessionId, + sandboxId: input.sandboxId, + transition: 'retire', + oldWrapperInstanceId: current.success ? current.data.wrapperInstanceId : null, + oldNativeRuntimeId: current.success ? current.data.nativeRuntimeId : null, + newWrapperInstanceId: null, + newNativeRuntimeId: null, + confirmed: input.confirmed, + storedFenceRetained: true, + authority: + input.nativeRuntimeId === undefined + ? 'sandbox_control_invalidation' + : 'native_retirement_confirmation', + }); } async recordNativeRuntime(input: { @@ -1462,6 +1580,18 @@ export class SandboxSession extends DurableObject { authorization: authorization.data, }) ); + logControlDiagnostic('native_fence_transition', { + sessionId: this.sessionId, + sandboxId: input.sandboxId, + transition: current.success ? 'rebind' : 'bind', + oldWrapperInstanceId: current.success ? current.data.wrapperInstanceId : null, + oldNativeRuntimeId: current.success ? current.data.nativeRuntimeId : null, + newWrapperInstanceId: input.wrapperInstanceId, + newNativeRuntimeId: input.nativeRuntimeId, + attachmentEpoch: proof.attachmentEpoch, + operationId: authorization.data.operationId, + authority: 'authorized_attach_result', + }); } async isSandboxCleanupScheduled(): Promise { @@ -2136,6 +2266,16 @@ export class SandboxSession extends DurableObject { } if (validationFailure) return validationFailure; if (!intent) return { success: false, code: 'NOT_FOUND', error: 'Message not found' }; + if ( + latestMessages.filter(message => message.state === 'queued').length >= + PENDING_SESSION_MESSAGE_LIMIT + ) { + return { + success: false, + code: 'PENDING_QUEUE_FULL', + error: `Pending message queue is full (${PENDING_SESSION_MESSAGE_LIMIT})`, + }; + } const nextMessages = freezeLegacyQueuedMessages( latestMessages, latestMetadata.agent, @@ -2256,9 +2396,9 @@ export class SandboxSession extends DurableObject { 'coordinator', undefined, () => bindControlEventReceiptIdentity(this.ctx.storage.kv, runtime.data), - isCurrent + () => (isCurrent() ? 'apply' : 'epoch_changed') ); - if (saved) wrapperInstanceId = runtime.data; + if (saved === 'apply') wrapperInstanceId = runtime.data; }; const dispatch = async ( kind: 'attach' | 'prompt', @@ -2305,58 +2445,65 @@ export class SandboxSession extends DurableObject { wrapperInstanceId, dispatchDeadlineAt: deadlineAt, }; - const dispatched = await dispatchSessionOperation( - { authorization, payload }, - { - read: () => this.loadMessages(), - commit: messages => this.saveMessages(messages, epoch, 'wrapper_outcome'), - }, - { - request: (input, scope) => sandboxControlRpc(this.env, sandboxId, scope).request(input), - persistResult: delivery => - this.applySandboxOperationResult({ - session: authorization.session, - wrapperInstanceId: authorization.wrapperInstanceId, - delivery, - }), - assertAdmission: () => { - if (!this.terminalLifecycle.isCurrent(epoch)) - throw new Error('Session operation scope changed'); - const current = this.loadMessages().find(message => message.messageId === messageId); - if (!current || current.wrapperInstanceId !== wrapperInstanceId) - throw new Error('Session operation scope changed'); + let dispatched: Awaited>; + try { + dispatched = await dispatchSessionOperation( + { authorization, payload }, + { + read: () => this.loadMessages(), + commit: messages => this.saveMessages(messages, epoch, 'wrapper_outcome'), }, - assertScope: () => { - if (!this.terminalLifecycle.isCurrent(epoch)) - throw new Error('Session operation scope changed'); - const current = this.loadMessages().find(message => message.messageId === messageId); - if (!current || current.wrapperInstanceId !== wrapperInstanceId) + { + request: (input, scope) => sandboxControlRpc(this.env, sandboxId, scope).request(input), + persistResult: delivery => + this.applySandboxOperationResult({ + session: authorization.session, + wrapperInstanceId: authorization.wrapperInstanceId, + delivery, + }), + assertAdmission: () => { + if (!this.terminalLifecycle.isCurrent(epoch)) + throw new Error('Session operation scope changed'); + const current = this.loadMessages().find(message => message.messageId === messageId); + if (!current || current.wrapperInstanceId !== wrapperInstanceId) + throw new Error('Session operation scope changed'); + }, + assertScope: () => { + if (!this.terminalLifecycle.isCurrent(epoch)) + throw new Error('Session operation scope changed'); + const current = this.loadMessages().find(message => message.messageId === messageId); + if (!current || current.wrapperInstanceId !== wrapperInstanceId) + throw new Error('Session operation scope changed'); + if ( + current.operations?.[operation === 'session.attach' ? 'attach' : 'prompt'] + ?.dispatched === true + ) + return; throw new Error('Session operation scope changed'); - if ( - current.operations?.[operation === 'session.attach' ? 'attach' : 'prompt'] - ?.dispatched === true - ) - return; - throw new Error('Session operation scope changed'); - }, - defer: pending => this.ctx.waitUntil(pending), - isCurrent, - } - ); - if ( - dispatched.state === 'response' && - authorization.operation === 'session.attach' && - sandboxId - ) { - const attached = sessionAttachResultSchema.safeParse(dispatched.result); - if (attached.success && attached.data.nativeRuntimeId !== undefined) { - await this.recordNativeRuntime({ - sandboxId, + }, + defer: pending => this.ctx.waitUntil(pending), + isCurrent, + } + ); + } catch (error) { + if (operation === 'session.attach') { + logControlDiagnostic('session_attach_completion', { + sessionId: this.sessionId, + messageId, + attemptId: authorization.operationId, + operationId: authorization.operationId, wrapperInstanceId, - nativeRuntimeId: attached.data.nativeRuntimeId, - authorization, + resultState: 'failed', + ok: false, + errorCode: + error instanceof ControlRequestError && + controlErrorCodes.includes(error.code as (typeof controlErrorCodes)[number]) + ? error.code + : 'other', + retryable: error instanceof ControlRequestError ? error.retryable : false, }); } + throw error; } if (dispatched.state !== 'response' && dispatched.state !== 'completed') { if (dispatched.state === 'running' && operation === 'session.prompt') return dispatched; @@ -2364,6 +2511,16 @@ export class SandboxSession extends DurableObject { } if (operation === 'session.attach') { const attached = sessionAttachResultSchema.parse(dispatched.result); + logControlDiagnostic('session_attach_completion', { + sessionId: this.sessionId, + messageId, + attemptId: authorization.operationId, + operationId: authorization.operationId, + wrapperInstanceId, + resultState: dispatched.state, + ok: true, + nativeRuntimeId: attached.nativeRuntimeId, + }); if (attached.nativeRuntimeId !== undefined) await this.recordNativeRuntime({ sandboxId, @@ -2940,12 +3097,13 @@ export class SandboxSession extends DurableObject { private nativeAttachmentEventDisposition( identity: SessionEventIdentity, - wrapperInstanceId?: string + wrapperInstanceId?: string, + triggeringMessage?: SessionMessageRecord ): 'pending' | 'rejected' | undefined { if (identity.nativeRuntimeId === undefined) return undefined; - const messages = this.loadMessages(); - const headId = nextQueuedMessageId(messages); - const message = messages.find(item => item.messageId === headId); + const messages = triggeringMessage ? undefined : this.loadMessages(); + const message = + triggeringMessage ?? messages?.find(item => item.messageId === nextQueuedMessageId(messages)); const proof = message?.operations?.attach; if ( !proof?.dispatched || @@ -2987,49 +3145,76 @@ export class SandboxSession extends DurableObject { ); } - private isCurrentReceiptRuntime(input: { + private controlEventDiagnosticSnapshot(input: { identity: SessionEventIdentity; - receiptId?: string; - receiptHash?: string; - sequence?: number; wrapperInstanceId?: string; - }): boolean { - return ( - (input.receiptId === undefined && - input.receiptHash === undefined && - input.sequence === undefined) || - (this.isCurrentEventRuntime(input.wrapperInstanceId) && - this.isCurrentNativeEventRuntime(input.identity, input.wrapperInstanceId)) + }): ControlDiagnosticFields { + const messages = this.loadMessages(); + const current = + messages.find(message => message.state === 'accepted') ?? + messages.find(message => message.state === 'queued'); + const expectedWrapperInstanceId = + current?.wrapperInstanceId ?? + messages.findLast(message => message.wrapperInstanceId)?.wrapperInstanceId; + const fence = nativeRuntimeFenceSchema.safeParse( + this.ctx.storage.kv.get(NATIVE_RUNTIME_FENCE_KEY) ); + return { + expectedWrapperInstanceId, + fenceWrapperInstanceId: fence.success ? fence.data.wrapperInstanceId : undefined, + nativeRuntimeId: input.identity.nativeRuntimeId, + fenceNativeRuntimeId: fence.success ? fence.data.nativeRuntimeId : undefined, + }; } - private async hasValidControlEventReceipt( - event: 'session.event' | 'session.preparing', - input: { - identity: SessionEventIdentity; - payload: unknown; - receiptId?: string; - receiptHash?: string; - sequence?: number; - wrapperInstanceId?: string; + private evaluateControlEvent(request: ControlEventEvaluationRequest): ControlEventDisposition { + const { input, epoch } = request; + if (!this.terminalLifecycle.isCurrent(epoch)) return 'epoch_changed'; + if (request.contract === 'currency_recheck') { + if (!this.isCurrentEventRuntime(input.wrapperInstanceId)) return 'runtime_mismatch'; + return this.isCurrentNativeEventRuntime(input.identity, input.wrapperInstanceId) + ? 'apply' + : 'native_runtime_mismatch'; } - ): Promise { - return hasValidControlEventReceipt(event, input); + if ( + request.publication.kind === 'event' && + input.receiptId === undefined && + input.receiptHash === undefined && + input.sequence === undefined + ) { + return this.isCurrentNativeEventRuntime(input.identity, input.wrapperInstanceId) + ? 'apply' + : 'native_runtime_mismatch'; + } + if (!this.isCurrentEventRuntime(input.wrapperInstanceId)) return 'runtime_mismatch'; + const receipt = this.controlEventReceipt(input); + if (receipt === 'duplicate') return 'duplicate'; + if (receipt !== 'apply') return 'receipt_conflict'; + const trigger = + request.publication.kind === 'preparing' ? request.publication.loadTrigger() : undefined; + if (request.publication.kind === 'preparing' && !trigger) return 'native_runtime_mismatch'; + const attachment = this.nativeAttachmentEventDisposition( + input.identity, + input.wrapperInstanceId, + trigger + ); + if (attachment === 'rejected') return 'native_runtime_mismatch'; + if (attachment === 'pending') return 'native_runtime_pending'; + return this.isCurrentNativeEventRuntime(input.identity, input.wrapperInstanceId) + ? 'apply' + : 'native_runtime_mismatch'; } private commitControlEventReceipt( - input: { - receiptId?: string; - receiptHash?: string; - sequence?: number; - wrapperInstanceId?: string; - }, - isCurrent: () => boolean = () => true - ): ControlEventReceiptDisposition { + input: ControlEventInput, + epoch: number + ): ControlEventDisposition { return this.ctx.storage.transactionSync(() => { - if (!isCurrent()) return 'stale'; + const currency = this.evaluateControlEvent({ contract: 'currency_recheck', input, epoch }); + if (currency !== 'apply') return currency; const receipt = this.controlEventReceipt(input); - if (receipt !== 'apply') return receipt; + if (receipt === 'duplicate') return 'duplicate'; + if (receipt !== 'apply') return 'receipt_conflict'; this.recordControlEventReceipt(input); return 'apply'; }); @@ -3488,14 +3673,29 @@ export class SandboxSession extends DurableObject { return this.ctx.storage.kv.get(MESSAGES_KEY) ?? []; } + private saveMessages( + messages: MessageRecord[], + epoch?: number, + source?: 'coordinator' | 'wrapper_outcome' | 'operation_result', + deferredNotifications?: StoredEvent[], + onPersist?: () => void + ): boolean; + private saveMessages( + messages: MessageRecord[], + epoch: number, + source: 'coordinator' | 'wrapper_outcome' | 'operation_result', + deferredNotifications: StoredEvent[] | undefined, + onPersist: (() => void) | undefined, + beforePersist: () => ControlEventDisposition + ): ControlEventDisposition; private saveMessages( messages: MessageRecord[], epoch?: number, source: 'coordinator' | 'wrapper_outcome' | 'operation_result' = 'coordinator', deferredNotifications?: StoredEvent[], onPersist?: () => void, - beforePersist?: () => boolean - ): boolean { + beforePersist?: () => ControlEventDisposition + ): boolean | ControlEventDisposition { return this.commitSavedMessages( messages, epoch, @@ -3513,14 +3713,16 @@ export class SandboxSession extends DurableObject { source: 'coordinator' | 'wrapper_outcome' | 'operation_result', deferredNotifications?: StoredEvent[] ): boolean { - return this.commitSavedMessages( - messages, - epoch, - source, - deferredNotifications, - undefined, - undefined, - write => write() + return ( + this.commitSavedMessages( + messages, + epoch, + source, + deferredNotifications, + undefined, + undefined, + write => write() + ) === true ); } @@ -3530,22 +3732,27 @@ export class SandboxSession extends DurableObject { source: 'coordinator' | 'wrapper_outcome' | 'operation_result', deferredNotifications: StoredEvent[] | undefined, onPersist: (() => void) | undefined, - beforePersist: (() => boolean) | undefined, + beforePersist: (() => ControlEventDisposition) | undefined, enclose: (write: () => void) => void - ): boolean { + ): boolean | ControlEventDisposition { + const returnsControlEventDisposition = beforePersist !== undefined; const currentEpoch = epoch ?? this.terminalLifecycle.captureEpoch(); if ( this.deletedWorktreeId || currentEpoch === null || !this.terminalLifecycle.isCurrent(currentEpoch) ) - return false; + return returnsControlEventDisposition ? 'epoch_changed' : false; const events: StoredEvent[] = []; const committed: ControlDiagnosticFields[] = []; let persisted = false; + let disposition: ControlEventDisposition = 'epoch_changed'; const write = () => { if (!this.terminalLifecycle.isCurrent(currentEpoch)) return; - if (beforePersist?.() === false) return; + if (beforePersist) { + disposition = beforePersist(); + if (disposition !== 'apply') return; + } const before = this.loadMessages(); const previousById = new Map(before.map(message => [message.messageId, message])); const queuedHeadId = nextQueuedMessageId(before); @@ -3618,7 +3825,7 @@ export class SandboxSession extends DurableObject { persisted = true; }; enclose(write); - if (!persisted) return false; + if (!persisted) return returnsControlEventDisposition ? disposition : false; for (const fields of committed) { logControlDiagnostic('session_message_committed', { sessionId: this.sessionId, @@ -3628,7 +3835,7 @@ export class SandboxSession extends DurableObject { } if (deferredNotifications) deferredNotifications.push(...events); else for (const event of events) this.broadcastStoredEvent(event); - return true; + return returnsControlEventDisposition ? 'apply' : true; } private persistMessageLifecycleEvent(message: MessageRecord): StoredEvent | undefined { diff --git a/services/cloud-agent-next/src/sandbox-session/control-event-receipts.test.ts b/services/cloud-agent-next/src/sandbox-session/control-event-receipts.test.ts index 453bcb86b7..cce853a08d 100644 --- a/services/cloud-agent-next/src/sandbox-session/control-event-receipts.test.ts +++ b/services/cloud-agent-next/src/sandbox-session/control-event-receipts.test.ts @@ -1,6 +1,7 @@ import { describe, expect, it } from 'vitest'; import { bindControlEventReceiptIdentity, + CONTROL_EVENT_RECEIPTS_KEY, controlEventReceiptDisposition, readControlEventReceipts, recordControlEventReceipt, @@ -13,7 +14,6 @@ const runtimeB = '22222222-2222-4222-8222-222222222222'; function receipt(wrapperInstanceId: string, sequence: number) { return { receiptId: `${String(sequence).padStart(8, '0')}-1111-4111-8111-111111111111`, - receiptHash: 'a'.repeat(64), wrapperInstanceId, sequence, }; @@ -54,6 +54,43 @@ describe('control event receipts', () => { expect(readControlEventReceipts(state).activeWrapperInstanceId).toBeUndefined(); }); + it('reads legacy hashes but does not write or compare them', () => { + const state = storage(); + const first = receipt(runtimeA, 1); + state.put(CONTROL_EVENT_RECEIPTS_KEY, { + highWater: { [runtimeA]: 1 }, + retiredWrapperInstanceIds: [], + receipts: [{ ...first, receiptHash: 'a'.repeat(64) }], + }); + expect(readControlEventReceipts(state).receipts).toEqual([ + { ...first, receiptHash: 'a'.repeat(64) }, + ]); + expect(controlEventReceiptDisposition(state, { ...first, receiptHash: 'b'.repeat(64) })).toBe( + 'duplicate' + ); + expect(controlEventReceiptDisposition(state, { ...first, sequence: 2 })).toBe('conflict'); + expect(controlEventReceiptDisposition(state, { ...first, wrapperInstanceId: runtimeB })).toBe( + 'apply' + ); + + recordControlEventReceipt(state, { ...first, receiptId: crypto.randomUUID(), sequence: 2 }); + expect(readControlEventReceipts(state).receipts.at(-1)).toEqual({ + receiptId: expect.any(String), + wrapperInstanceId: runtimeA, + sequence: 2, + }); + }); + + it('rejects partial receipt identities', () => { + const state = storage(); + expect(controlEventReceiptDisposition(state, {})).toBe('apply'); + expect(controlEventReceiptDisposition(state, { receiptHash: 'a'.repeat(64) })).toBe('conflict'); + expect(controlEventReceiptDisposition(state, { sequence: 1 })).toBe('conflict'); + expect(controlEventReceiptDisposition(state, { receiptId: crypto.randomUUID() })).toBe( + 'conflict' + ); + }); + it('rejects a retired wrapper identity after pruning its receipt state', () => { const state = storage(); const first = receipt(runtimeA, 1); diff --git a/services/cloud-agent-next/src/sandbox-session/control-event-receipts.ts b/services/cloud-agent-next/src/sandbox-session/control-event-receipts.ts index 0d613b3cfc..a13409ece9 100644 --- a/services/cloud-agent-next/src/sandbox-session/control-event-receipts.ts +++ b/services/cloud-agent-next/src/sandbox-session/control-event-receipts.ts @@ -1,5 +1,4 @@ import { z } from 'zod'; -import { canonicalControlEventJson } from '../shared/control-event-canonical.js'; import { wrapperInstanceIdSchema } from '../shared/sandbox-control-protocol.js'; export const CONTROL_EVENT_RECEIPTS_KEY = 'control_event_receipts'; @@ -8,7 +7,10 @@ const CONTROL_EVENT_RECEIPT_LIMIT = 64; const controlEventReceiptSchema = z .object({ receiptId: z.string().uuid(), - receiptHash: z.string().regex(/^[a-f0-9]{64}$/), + receiptHash: z + .string() + .regex(/^[a-f0-9]{64}$/) + .optional(), wrapperInstanceId: wrapperInstanceIdSchema, sequence: z.number().int().positive().max(Number.MAX_SAFE_INTEGER), }) @@ -87,11 +89,7 @@ export function controlEventReceiptDisposition( item.receiptId === receipt.data.receiptId && item.wrapperInstanceId === receipt.data.wrapperInstanceId ); - if (stored) - return stored.receiptHash === receipt.data.receiptHash && - stored.sequence === receipt.data.sequence - ? 'duplicate' - : 'conflict'; + if (stored) return stored.sequence === receipt.data.sequence ? 'duplicate' : 'conflict'; return receipt.data.sequence <= (state.highWater[receipt.data.wrapperInstanceId] ?? 0) ? 'stale' : 'apply'; @@ -110,7 +108,14 @@ export function recordControlEventReceipt( ...current.highWater, [receipt.data.wrapperInstanceId]: receipt.data.sequence, }, - receipts: [...current.receipts, receipt.data].slice(-CONTROL_EVENT_RECEIPT_LIMIT), + receipts: [ + ...current.receipts, + { + receiptId: receipt.data.receiptId, + wrapperInstanceId: receipt.data.wrapperInstanceId, + sequence: receipt.data.sequence, + }, + ].slice(-CONTROL_EVENT_RECEIPT_LIMIT), }); } @@ -149,33 +154,3 @@ export function retireControlEventReceiptIdentity( receipts: current.receipts.filter(item => item.wrapperInstanceId !== wrapperInstanceId), }); } - -export async function hasValidControlEventReceipt( - event: 'session.event' | 'session.preparing', - input: { - identity: unknown; - payload: unknown; - } & ControlEventReceiptInput -): Promise { - if ( - input.receiptId === undefined && - input.receiptHash === undefined && - input.sequence === undefined - ) - return true; - const receipt = parseControlEventReceipt(input); - if (!receipt.success) return false; - const digest = await crypto.subtle.digest( - 'SHA-256', - new TextEncoder().encode( - canonicalControlEventJson({ - event, - session: input.identity, - payload: input.payload, - sequence: receipt.data.sequence, - }) - ) - ); - const hash = [...new Uint8Array(digest)].map(byte => byte.toString(16).padStart(2, '0')).join(''); - return hash === receipt.data.receiptHash; -} diff --git a/services/cloud-agent-next/src/sandbox-session/control-plane-preparing.test.ts b/services/cloud-agent-next/src/sandbox-session/control-plane-preparing.test.ts index a491ef217f..3b657b01ff 100644 --- a/services/cloud-agent-next/src/sandbox-session/control-plane-preparing.test.ts +++ b/services/cloud-agent-next/src/sandbox-session/control-plane-preparing.test.ts @@ -1,41 +1,54 @@ import { describe, expect, it, vi } from 'vitest'; +import { createMemoryEventQueries } from '../session/preparation-test-helpers.js'; import { applyControlPlanePreparingEvent } from './control-plane-preparing.js'; describe('applyControlPlanePreparingEvent', () => { - it('broadcasts preparing and cloud.status when the event materializes', () => { - const upsert = vi.fn(); + const event = { + version: 2, + attemptId: 'att_1', + triggerMessageId: 'msg_1', + revision: 1, + timestamp: 10, + step: 'cloning', + message: 'Cloning repository…', + action: 'attempt_started', + }; + + it('broadcasts only an accepted v2 preparing event', () => { const broadcast = vi.fn(); const applied = applyControlPlanePreparingEvent({ sessionId: 'workspace_1', - data: { - version: 2, - attemptId: 'att_1', - triggerMessageId: 'msg_1', - revision: 1, - timestamp: 10, - step: 'cloning', - message: 'Cloning repository…', - action: 'attempt_started', - }, - eventQueries: { - upsert, - insert: vi.fn(), - findByEntityId: vi.fn().mockReturnValue(undefined), - findByEntityPrefix: vi.fn().mockReturnValue([]), - } as never, + data: event, + eventQueries: createMemoryEventQueries(), broadcast, }); + expect(applied).toBe(true); + expect(broadcast).toHaveBeenCalledTimes(1); expect(broadcast).toHaveBeenCalledWith( - expect.objectContaining({ - session_id: 'workspace_1', - stream_event_type: 'preparing', - }) - ); - expect(broadcast).toHaveBeenCalledWith( - expect.objectContaining({ - stream_event_type: 'cloud.status', - }) + expect.objectContaining({ session_id: 'workspace_1', stream_event_type: 'preparing' }) ); }); + + it('does not broadcast a rejected duplicate event', () => { + const broadcast = vi.fn(); + const eventQueries = createMemoryEventQueries(); + applyControlPlanePreparingEvent({ + sessionId: 'workspace_1', + data: event, + eventQueries, + broadcast, + }); + broadcast.mockClear(); + + const applied = applyControlPlanePreparingEvent({ + sessionId: 'workspace_1', + data: event, + eventQueries, + broadcast, + }); + + expect(applied).toBe(false); + expect(broadcast).not.toHaveBeenCalled(); + }); }); diff --git a/services/cloud-agent-next/src/sandbox-session/control-plane-preparing.ts b/services/cloud-agent-next/src/sandbox-session/control-plane-preparing.ts index f1362045f6..56618927ef 100644 --- a/services/cloud-agent-next/src/sandbox-session/control-plane-preparing.ts +++ b/services/cloud-agent-next/src/sandbox-session/control-plane-preparing.ts @@ -1,9 +1,5 @@ import type { EventQueries } from '../session/queries/index.js'; -import { - cloudStatusForPreparingEvent, - materializePreparationEvent, -} from '../session/preparation-history.js'; -import type { CloudStatusData } from '../shared/protocol.js'; +import { materializePreparationEvent } from '../session/preparation-history.js'; import type { EventId } from '../types/ids.js'; import type { StoredEvent } from '../websocket/types.js'; @@ -26,20 +22,9 @@ export function applyControlPlanePreparingEvent(params: { payload: JSON.stringify(params.data), timestamp, }; + if (!materializePreparationEvent(params.eventQueries, stored, params.data)) return false; params.broadcast(stored); - const applied = materializePreparationEvent(params.eventQueries, stored, params.data); - const cloudStatus = cloudStatusForPreparingEvent(params.data, applied); - if (cloudStatus) { - params.broadcast({ - id: 0 as EventId, - execution_id: '', - session_id: params.sessionId, - stream_event_type: 'cloud.status', - payload: JSON.stringify({ cloudStatus } satisfies CloudStatusData), - timestamp, - }); - } - return applied; + return true; } function isRecord(value: unknown): value is Record { diff --git a/services/cloud-agent-next/src/sandbox-session/session-message-queue.test.ts b/services/cloud-agent-next/src/sandbox-session/session-message-queue.test.ts index c749dfe9a6..440bb27663 100644 --- a/services/cloud-agent-next/src/sandbox-session/session-message-queue.test.ts +++ b/services/cloud-agent-next/src/sandbox-session/session-message-queue.test.ts @@ -1,6 +1,4 @@ -import { createHash } from 'node:crypto'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; -import { canonicalControlEventJson } from '../shared/control-event-canonical.js'; import { normalizeCliEvent } from '../../../../packages/cloud-agent-sdk/src/normalizer'; import { createServiceState } from '../../../../packages/cloud-agent-sdk/src/service-state'; import { SandboxSession } from './SandboxSession.js'; @@ -31,6 +29,7 @@ import { } from '../shared/sandbox-control-protocol.js'; import { DEADLINE_MS } from '../sandbox-control/deadlines.js'; import { createControlPlaneCredential } from '../sandbox-control/managed-credential.js'; +import { logger } from '../logger.js'; import { SESSION_DELIVERY_TIMEOUT_MS } from './control-dispatch.js'; import { createControlStopRequest } from '../shared/control-plane-session.js'; import type { @@ -961,11 +960,36 @@ function receiptedEvent( payload, sequence, receiptId: crypto.randomUUID(), - receiptHash: createHash('sha256') - .update( - canonicalControlEventJson({ event: 'session.event', session: identity, payload, sequence }) - ) - .digest('hex'), + }; +} + +function receiptedPreparing( + sequence: number, + payload: Parameters[0]['payload'], + wrapperInstanceId = RUNTIME_ID, + nativeRuntimeId?: string +) { + const identity = { + directory: DIRECTORY, + kiloSessionId: 'kilo_root', + rootKiloSessionId: 'kilo_root', + ...(nativeRuntimeId ? { nativeRuntimeId } : {}), + }; + return { + identity, + wrapperInstanceId, + payload, + sequence, + receiptId: crypto.randomUUID(), + }; +} + +function unreceiptedPreparing(...input: Parameters) { + const preparing = receiptedPreparing(...input); + return { + identity: preparing.identity, + wrapperInstanceId: preparing.wrapperInstanceId, + payload: preparing.payload, }; } @@ -1192,6 +1216,389 @@ describe('SandboxSession orchestration', () => { vi.useRealTimers(); }); + it('reports the named runtime gate with incoming, expected, and fence identities', async () => { + const fixture = sessionFixture(); + const expectedWrapperInstanceId = RUNTIME_ID; + const wrapperInstanceId = '44444444-4444-4444-8444-444444444444'; + const fenceNativeRuntimeId = NEXT_RUNTIME_ID; + fixture.storage.kv.put('session_messages', [ + { messageId: 'queued', state: 'queued', wrapperInstanceId: expectedWrapperInstanceId }, + ]); + fixture.storage.kv.put('native_runtime_fence', { + sandboxId: SANDBOX_ID, + wrapperInstanceId: expectedWrapperInstanceId, + nativeRuntimeId: fenceNativeRuntimeId, + attachmentEpoch: 1, + authorization: { + operation: 'session.attach', + operationId: 'attempt_1', + messageId: 'queued', + session: { sessionId: SESSION_ID, kiloSessionId: 'kilo_root', directory: DIRECTORY }, + wrapperInstanceId: expectedWrapperInstanceId, + dispatchDeadlineAt: Date.now() + SESSION_DELIVERY_TIMEOUT_MS, + }, + }); + const fields = vi.spyOn(logger, 'withFields').mockReturnValue(logger); + try { + const event = receiptedEvent( + 1, + { type: 'session.updated', properties: { info: { id: 'kilo_root' } } }, + wrapperInstanceId, + '55555555-5555-4555-8555-555555555555' + ); + const before = structuredClone([...fixture.values]); + await expect(fixture.session.receiveSandboxControlEvent(event)).resolves.toEqual({ + applied: false, + }); + expect([...fixture.values]).toEqual(before); + expect(fields).toHaveBeenCalledWith( + expect.objectContaining({ + diagnosticEvent: 'session_event_result', + disposition: 'runtime_mismatch', + wrapperInstanceId, + expectedWrapperInstanceId, + nativeRuntimeId: '55555555-5555-4555-8555-555555555555', + fenceWrapperInstanceId: expectedWrapperInstanceId, + fenceNativeRuntimeId, + receiptId: event.receiptId, + }) + ); + } finally { + fields.mockRestore(); + } + }); + + it('keeps an unreceipted outcome on the early publication path', async () => { + const fixture = sessionFixture(); + fixture.storage.kv.put('session_messages', [ + { messageId: 'queued', state: 'queued', wrapperInstanceId: RUNTIME_ID }, + ]); + const fields = vi.spyOn(logger, 'withFields').mockReturnValue(logger); + try { + const input = { + identity: { + directory: DIRECTORY, + kiloSessionId: 'kilo_root', + rootKiloSessionId: 'kilo_root', + }, + wrapperInstanceId: NEXT_RUNTIME_ID, + payload: { + type: 'session.message.outcome', + properties: { messageId: 'missing', status: 'completed' }, + }, + } as const; + const before = structuredClone([...fixture.values]); + await expect(fixture.session.receiveSandboxControlEvent(input)).resolves.toEqual({ + applied: false, + }); + expect([...fixture.values]).toEqual(before); + expect(fields).toHaveBeenCalledWith( + expect.objectContaining({ + diagnosticEvent: 'session_event_result', + disposition: 'message_missing', + }) + ); + + await expect( + fixture.session.receiveSandboxControlEvent({ + ...input, + identity: { ...input.identity, nativeRuntimeId: NEXT_RUNTIME_ID }, + }) + ).resolves.toEqual({ applied: false }); + expect(fields).toHaveBeenCalledWith( + expect.objectContaining({ + diagnosticEvent: 'session_event_result', + disposition: 'native_runtime_mismatch', + }) + ); + } finally { + fields.mockRestore(); + } + }); + + it('rejects an unreceipted remaining event from a stale wrapper', async () => { + const fixture = sessionFixture(); + fixture.storage.kv.put('session_messages', [ + { messageId: 'queued', state: 'queued', wrapperInstanceId: RUNTIME_ID }, + ]); + const fields = vi.spyOn(logger, 'withFields').mockReturnValue(logger); + try { + const before = structuredClone([...fixture.values]); + const events = fixture.eventQueries.findByEntityPrefix(''); + await expect( + fixture.session.receiveSandboxControlEvent({ + identity: { + directory: DIRECTORY, + kiloSessionId: 'kilo_root', + rootKiloSessionId: 'kilo_root', + }, + wrapperInstanceId: NEXT_RUNTIME_ID, + payload: { type: 'session.updated', properties: { info: { id: 'kilo_root' } } }, + }) + ).resolves.toEqual({ applied: false }); + expect([...fixture.values]).toEqual(before); + expect(fixture.eventQueries.findByEntityPrefix('')).toEqual(events); + expect(fields).toHaveBeenCalledWith( + expect.objectContaining({ + diagnosticEvent: 'session_event_result', + disposition: 'runtime_mismatch', + }) + ); + } finally { + fields.mockRestore(); + } + }); + + it('rejects partial receipt identities without recording events or receipts', async () => { + const fixture = sessionFixture(); + fixture.storage.kv.put('session_messages', [ + { + messageId: 'queued', + state: 'queued', + wrapperInstanceId: RUNTIME_ID, + preparationAttemptId: 'attempt_1', + }, + ]); + const event = receiptedEvent(1, { + type: 'session.updated', + properties: { info: { id: 'kilo_root' } }, + }); + const preparing = receiptedPreparing(1, { + version: 2, + attemptId: 'attempt_1', + triggerMessageId: 'queued', + revision: 1, + timestamp: Date.now(), + step: 'workspace_setup', + action: 'attempt_started', + message: 'Preparing environment', + }); + const before = structuredClone([...fixture.values]); + const events = fixture.eventQueries.findByEntityPrefix(''); + const fields = vi.spyOn(logger, 'withFields').mockReturnValue(logger); + + try { + await expect( + fixture.session.receiveSandboxControlEvent({ + ...event, + receiptId: undefined, + sequence: undefined, + receiptHash: 'a'.repeat(64), + }) + ).resolves.toEqual({ applied: false }); + expect(fields).toHaveBeenCalledWith( + expect.objectContaining({ + diagnosticEvent: 'session_event_result', + disposition: 'receipt_conflict', + }) + ); + + await expect( + fixture.session.receiveSandboxControlEvent({ + ...event, + receiptId: undefined, + }) + ).resolves.toEqual({ applied: false }); + expect(fields).toHaveBeenCalledWith( + expect.objectContaining({ + diagnosticEvent: 'session_event_result', + disposition: 'receipt_conflict', + }) + ); + + await expect( + fixture.session.receiveSandboxControlPreparing({ + ...preparing, + receiptId: undefined, + sequence: undefined, + receiptHash: 'a'.repeat(64), + }) + ).resolves.toEqual({ applied: false }); + expect(fields).toHaveBeenCalledWith( + expect.objectContaining({ + diagnosticEvent: 'session_preparing_result', + disposition: 'receipt_conflict', + }) + ); + + expect([...fixture.values]).toEqual(before); + expect(fixture.eventQueries.findByEntityPrefix('')).toEqual(events); + } finally { + fields.mockRestore(); + } + }); + + it.each([ + { family: 'session.event', disposition: 'epoch_changed' }, + { family: 'session.event', disposition: 'runtime_mismatch' }, + { family: 'session.event', disposition: 'native_runtime_mismatch' }, + { family: 'session.event', disposition: 'receipt_conflict' }, + { family: 'session.preparing', disposition: 'epoch_changed' }, + { family: 'session.preparing', disposition: 'runtime_mismatch' }, + { family: 'session.preparing', disposition: 'native_runtime_mismatch' }, + { family: 'session.preparing', disposition: 'receipt_conflict' }, + ] as const)( + 'reports $family transaction recheck as $disposition', + async ({ family, disposition }) => { + const fixture = sessionFixture(); + const nativeRuntimeId = '55555555-5555-4555-8555-555555555555'; + const replacementNativeRuntimeId = '66666666-6666-4666-8666-666666666666'; + const replacementWrapperInstanceId = '77777777-7777-4777-8777-777777777777'; + const authorization: SessionOperationAuthorization = { + operation: 'session.attach', + operationId: 'attempt_1', + messageId: 'queued', + session: { sessionId: SESSION_ID, kiloSessionId: 'kilo_root', directory: DIRECTORY }, + wrapperInstanceId: RUNTIME_ID, + dispatchDeadlineAt: Date.now() + SESSION_DELIVERY_TIMEOUT_MS, + }; + const fence = (runtimeId: string) => ({ + sandboxId: SANDBOX_ID, + wrapperInstanceId: RUNTIME_ID, + nativeRuntimeId: runtimeId, + attachmentEpoch: 1, + authorization, + }); + fixture.storage.kv.put('native_runtime_fence', fence(nativeRuntimeId)); + fixture.storage.kv.put('session_messages', [ + { + messageId: 'queued', + state: 'queued', + wrapperInstanceId: RUNTIME_ID, + ...(family === 'session.preparing' + ? { + preparationAttemptId: 'attempt_1', + operations: { attach: { authorization, dispatched: true } }, + } + : {}), + }, + ]); + const eventInput = receiptedEvent( + 1, + { type: 'question.asked', properties: { id: 'question_1', sessionID: 'kilo_root' } }, + RUNTIME_ID, + nativeRuntimeId + ); + const preparingInput = receiptedPreparing( + 1, + { + version: 2, + attemptId: 'attempt_1', + triggerMessageId: 'queued', + revision: 1, + timestamp: Date.now(), + step: 'workspace_setup', + action: 'attempt_started', + message: 'Preparing environment', + }, + RUNTIME_ID, + nativeRuntimeId + ); + const input = family === 'session.event' ? eventInput : preparingInput; + const storage = fixture.storage as unknown as { + transactionSync: (callback: () => T) => T; + }; + const transactionSync = storage.transactionSync.bind(storage); + let storageAtRecheck: Array<[string, unknown]> | undefined; + let restoreEpochCheck: () => void = () => {}; + storage.transactionSync = callback => { + if (!storageAtRecheck) { + if (disposition === 'epoch_changed') { + const lifecycle = ( + fixture.session as unknown as { + terminalLifecycle: { isCurrent: (epoch: number) => boolean }; + } + ).terminalLifecycle; + const epochCheck = vi.spyOn(lifecycle, 'isCurrent').mockReturnValue(false); + restoreEpochCheck = () => epochCheck.mockRestore(); + } else if (disposition === 'runtime_mismatch') { + fixture.storage.kv.put('session_messages', [ + { + messageId: 'replacement', + state: 'queued', + wrapperInstanceId: replacementWrapperInstanceId, + }, + ]); + } else if (disposition === 'native_runtime_mismatch' && family === 'session.preparing') { + fixture.storage.kv.put('session_messages', [ + { + messageId: 'queued', + state: 'queued', + wrapperInstanceId: RUNTIME_ID, + preparationAttemptId: 'attempt_1', + operations: { + attach: { + authorization: { ...authorization, operationId: 'replacement-attach' }, + dispatched: true, + result: { + ok: false, + error: { code: 'not_ready', message: 'Attach failed', retryable: true }, + }, + }, + }, + } satisfies SessionMessageRecord, + ]); + } else if (disposition === 'native_runtime_mismatch') { + fixture.storage.kv.put('native_runtime_fence', fence(replacementNativeRuntimeId)); + } else { + fixture.storage.kv.put('control_event_receipts', { + highWater: {}, + retiredWrapperInstanceIds: [], + receipts: [ + { + receiptId: input.receiptId, + wrapperInstanceId: RUNTIME_ID, + sequence: (input.sequence ?? 0) + 1, + }, + ], + }); + } + storageAtRecheck = structuredClone([...fixture.values]); + } + return transactionSync(callback); + }; + const events = structuredClone(fixture.eventQueries.findByEntityPrefix('')); + const fields = vi.spyOn(logger, 'withFields').mockReturnValue(logger); + try { + const result = + family === 'session.event' + ? await fixture.session.receiveSandboxControlEvent(eventInput) + : await fixture.session.receiveSandboxControlPreparing(preparingInput); + expect(result).toEqual({ applied: false }); + expect([...fixture.values]).toEqual(storageAtRecheck); + expect(fixture.eventQueries.findByEntityPrefix('')).toEqual(events); + expect(fields).toHaveBeenCalledWith( + expect.objectContaining({ + diagnosticEvent: + family === 'session.event' ? 'session_event_result' : 'session_preparing_result', + disposition, + wrapperInstanceId: RUNTIME_ID, + expectedWrapperInstanceId: RUNTIME_ID, + nativeRuntimeId, + fenceWrapperInstanceId: RUNTIME_ID, + fenceNativeRuntimeId: nativeRuntimeId, + }) + ); + if (family === 'session.preparing') { + const diagnostic = fields.mock.calls + .map(([value]) => value) + .find(value => value.diagnosticEvent === 'session_preparing_result'); + expect(diagnostic).toMatchObject({ + attemptId: 'attempt_1', + action: 'attempt_started', + revision: 1, + }); + expect([diagnostic?.attemptId, diagnostic?.action, diagnostic?.revision]).not.toContain( + 'redacted' + ); + } + } finally { + restoreEpochCheck(); + fields.mockRestore(); + } + } + ); + it('persists the alarm and head budget before the first RPC and wakes the head on a fresh ID after reset', async () => { const fixture = sessionFixture(); const firstReady = deferred(); @@ -1239,7 +1646,6 @@ describe('SandboxSession orchestration', () => { }, wrapperInstanceId: RUNTIME_ID, receiptId: '11111111-1111-4111-8111-111111111111', - receiptHash: 'fe56f74e2e533bffb077a4b73ded75450cdbb4752e421e6bb77d44093fd762f8', sequence: 1, payload: { type: 'session.message.outcome', @@ -1266,7 +1672,7 @@ describe('SandboxSession orchestration', () => { await expect( fixture.session.receiveSandboxControlEvent({ ...event, - receiptHash: 'f'.repeat(64), + sequence: 2, }) ).resolves.toEqual({ applied: false }); expect(fixture.terminalEvents()).toHaveLength(1); @@ -1322,16 +1728,6 @@ describe('SandboxSession orchestration', () => { expect(fixture.control.quarantineRuntime).not.toHaveBeenCalled(); const invalid = { ...event, identity: { ...event.identity, directory: '/foreign' } }; - invalid.receiptHash = createHash('sha256') - .update( - canonicalControlEventJson({ - event: 'session.event', - session: invalid.identity, - payload, - sequence: invalid.sequence, - }) - ) - .digest('hex'); await expect(fixture.session.receiveSandboxControlEvent(invalid)).resolves.toEqual({ applied: false, }); @@ -1491,6 +1887,576 @@ describe('SandboxSession orchestration', () => { ); describe('native startup attach authority', () => { + it.each([false, true])( + 'applies preparing from its pending attach proof without changing prior fence=%s', + async priorFence => { + const fixture = sessionFixture(); + const attach = deferred(); + const nativeRuntimeId = NEXT_RUNTIME_ID; + if (priorFence) { + const authorization: SessionOperationAuthorization = { + operation: 'session.attach', + operationId: 'previous-attach', + messageId: 'previous', + session: { sessionId: SESSION_ID, kiloSessionId: 'kilo_root', directory: DIRECTORY }, + wrapperInstanceId: '11111111-1111-4111-8111-111111111111', + dispatchDeadlineAt: Date.now() + SESSION_DELIVERY_TIMEOUT_MS, + }; + fixture.storage.kv.put('session_messages', [ + { + messageId: 'previous', + state: 'completed', + wrapperInstanceId: authorization.wrapperInstanceId, + operations: { + attach: { + authorization, + dispatched: true, + completedAt: Date.now(), + attachmentEpoch: 1, + }, + }, + } satisfies SessionMessageRecord, + ]); + await fixture.session.recordNativeRuntime({ + sandboxId: SANDBOX_ID, + wrapperInstanceId: authorization.wrapperInstanceId, + nativeRuntimeId: '44444444-4444-4444-8444-444444444444', + authorization, + }); + } + const fence = structuredClone(fixture.values.get('native_runtime_fence')); + fixture.setStatus({ + physical: 'running', + connection: 'ready', + wrapperInstanceId: RUNTIME_ID, + operationResults: true, + }); + delegateRequest(fixture, 'session.attach', () => attach.promise); + await fixture.admit('a'); + await fixture.flush(); + const message = fixture.record('a'); + const authorization = message?.operations?.attach?.authorization; + const attemptId = message?.preparationAttemptId; + if (!authorization || !attemptId) throw new Error('Missing pending attach authority'); + const preparing = receiptedPreparing( + 1, + { + version: 2, + attemptId, + triggerMessageId: 'a', + revision: 1_000, + timestamp: Date.now(), + step: 'workspace_setup', + action: 'step_progress', + message: 'Preparing environment', + stepId: 'phase:workspace_setup', + detail: 'Preparing environment', + }, + RUNTIME_ID, + nativeRuntimeId + ); + await expect(fixture.session.receiveSandboxControlPreparing(preparing)).resolves.toEqual({ + applied: true, + }); + expect(fixture.values.get('native_runtime_fence')).toEqual(fence); + expect(fixture.values.get('control_event_receipts')).toBeDefined(); + const receipts = structuredClone(fixture.values.get('control_event_receipts')); + expect( + fixture.eventQueries.findByEntityPrefix(`preparation/attempt/${attemptId}`).length + ).toBeGreaterThan(0); + const prepared = fixture.eventQueries.findByEntityPrefix(''); + expect( + JSON.parse( + fixture.eventQueries.findByEntityId(`preparation/attempt/${attemptId}`)?.payload ?? '{}' + ) + ).toMatchObject({ revision: 1_000, triggerMessageId: 'a' }); + await expect(fixture.session.receiveSandboxControlPreparing(preparing)).resolves.toEqual({ + applied: true, + }); + expect(fixture.values.get('control_event_receipts')).toEqual(receipts); + expect(fixture.eventQueries.findByEntityPrefix('')).toEqual(prepared); + + const nativeEvents = [ + receiptedEvent( + 2, + { + type: 'session.status', + properties: { sessionID: 'kilo_root', status: { type: 'busy' } }, + }, + RUNTIME_ID, + nativeRuntimeId + ), + receiptedEvent( + 3, + { + type: 'session.updated', + properties: { info: { id: 'kilo_root', title: 'Native startup' } }, + }, + RUNTIME_ID, + nativeRuntimeId + ), + ]; + for (const event of nativeEvents) { + const values = structuredClone([...fixture.values]); + const events = fixture.eventQueries.findByEntityPrefix(''); + await expect(fixture.session.receiveSandboxControlEvent(event)).resolves.toEqual({ + applied: false, + retryable: true, + }); + expect([...fixture.values]).toEqual(values); + expect(fixture.eventQueries.findByEntityPrefix('')).toEqual(events); + } + expect(fixture.values.get('native_runtime_fence')).toEqual(fence); + + attach.resolve(controlResponse({ attached: true, nativeRuntimeId })); + await fixture.flush(); + expect(fixture.values.get('native_runtime_fence')).toMatchObject({ + wrapperInstanceId: RUNTIME_ID, + nativeRuntimeId, + authorization, + }); + for (const event of nativeEvents) + await expect(fixture.session.receiveSandboxControlEvent(event)).resolves.toEqual({ + applied: true, + }); + } + ); + + it('rejects unreceipted preparing from a stale wrapper before persistence', async () => { + const fixture = sessionFixture(); + const attach = deferred(); + fixture.setStatus({ + physical: 'running', + connection: 'ready', + wrapperInstanceId: RUNTIME_ID, + operationResults: true, + }); + delegateRequest(fixture, 'session.attach', () => attach.promise); + await fixture.admit('a'); + await fixture.flush(); + const attemptId = fixture.record('a')?.preparationAttemptId; + if (!attemptId) throw new Error('Missing preparation attempt'); + const preparing = unreceiptedPreparing( + 1, + { + version: 2, + attemptId, + triggerMessageId: 'a', + revision: 1, + timestamp: Date.now(), + step: 'workspace_setup', + action: 'attempt_started', + message: 'Preparing environment', + }, + NEXT_RUNTIME_ID + ); + const before = structuredClone([...fixture.values]); + const events = fixture.eventQueries.findByEntityPrefix(''); + const fields = vi.spyOn(logger, 'withFields').mockReturnValue(logger); + try { + await expect(fixture.session.receiveSandboxControlPreparing(preparing)).resolves.toEqual({ + applied: false, + }); + expect([...fixture.values]).toEqual(before); + expect(fixture.eventQueries.findByEntityPrefix('')).toEqual(events); + expect(fields).toHaveBeenCalledWith( + expect.objectContaining({ + diagnosticEvent: 'session_preparing_result', + disposition: 'runtime_mismatch', + }) + ); + } finally { + fields.mockRestore(); + } + }); + + it('rechecks an unreceipted preparing trigger inside its transaction', async () => { + const fixture = sessionFixture(); + const attach = deferred(); + fixture.setStatus({ + physical: 'running', + connection: 'ready', + wrapperInstanceId: RUNTIME_ID, + operationResults: true, + }); + delegateRequest(fixture, 'session.attach', () => attach.promise); + await fixture.admit('a'); + await fixture.flush(); + const message = fixture.record('a'); + const attemptId = message?.preparationAttemptId; + if (!message || !attemptId) throw new Error('Missing preparation attempt'); + const preparing = unreceiptedPreparing(1, { + version: 2, + attemptId, + triggerMessageId: 'a', + revision: 1, + timestamp: Date.now(), + step: 'workspace_setup', + action: 'attempt_started', + message: 'Preparing environment', + }); + const storage = fixture.storage as unknown as { + transactionSync: (callback: () => T) => T; + }; + const transactionSync = storage.transactionSync.bind(storage); + storage.transactionSync = callback => { + fixture.storage.kv.put('session_messages', [ + { ...message, cancellation: { operationId: 'cancel', deadlineAt: Date.now() } }, + ]); + return transactionSync(callback); + }; + const events = fixture.eventQueries.findByEntityPrefix(''); + const fields = vi.spyOn(logger, 'withFields').mockReturnValue(logger); + try { + await expect(fixture.session.receiveSandboxControlPreparing(preparing)).resolves.toEqual({ + applied: false, + }); + expect(fixture.eventQueries.findByEntityPrefix('')).toEqual(events); + expect(fixture.record('a')?.cancellation).toBeDefined(); + expect(fields).toHaveBeenCalledWith( + expect.objectContaining({ + diagnosticEvent: 'session_preparing_result', + disposition: 'native_runtime_mismatch', + }) + ); + } finally { + fields.mockRestore(); + } + }); + + it('persists unreceipted preparing after its pending attach proof passes', async () => { + const fixture = sessionFixture(); + const attach = deferred(); + fixture.setStatus({ + physical: 'running', + connection: 'ready', + wrapperInstanceId: RUNTIME_ID, + operationResults: true, + }); + delegateRequest(fixture, 'session.attach', () => attach.promise); + await fixture.admit('a'); + await fixture.flush(); + const attemptId = fixture.record('a')?.preparationAttemptId; + if (!attemptId) throw new Error('Missing preparation attempt'); + const preparing = unreceiptedPreparing( + 1, + { + version: 2, + attemptId, + triggerMessageId: 'a', + revision: 1, + timestamp: Date.now(), + step: 'workspace_setup', + action: 'attempt_started', + message: 'Preparing environment', + }, + RUNTIME_ID, + NEXT_RUNTIME_ID + ); + await expect(fixture.session.receiveSandboxControlPreparing(preparing)).resolves.toEqual({ + applied: true, + }); + expect( + fixture.eventQueries.findByEntityPrefix(`preparation/attempt/${attemptId}`).length + ).toBeGreaterThan(0); + expect(fixture.values.get('control_event_receipts')).toMatchObject({ receipts: [] }); + }); + + it('uses only the trigger message attach proof for pending preparation authority', async () => { + const fixture = sessionFixture(); + const attach = deferred(); + fixture.setStatus({ + physical: 'running', + connection: 'ready', + wrapperInstanceId: RUNTIME_ID, + operationResults: true, + }); + delegateRequest(fixture, 'session.attach', () => attach.promise); + await fixture.admit('trigger'); + await fixture.flush(); + const trigger = fixture.record('trigger'); + const proof = trigger?.operations?.attach; + const attemptId = trigger?.preparationAttemptId; + if (!trigger || !proof || !attemptId) throw new Error('Missing trigger attach authority'); + fixture.storage.kv.put('session_messages', [ + { + messageId: 'head', + state: 'queued', + wrapperInstanceId: RUNTIME_ID, + preparationAttemptId: 'head-attempt', + operations: { + attach: { + ...proof, + authorization: { + ...proof.authorization, + operationId: 'head-attach', + messageId: 'head', + }, + }, + }, + } satisfies SessionMessageRecord, + { ...trigger, operations: undefined }, + ]); + const preparing = receiptedPreparing( + 1, + { + version: 2, + attemptId, + triggerMessageId: 'trigger', + revision: 1, + timestamp: Date.now(), + step: 'workspace_setup', + action: 'attempt_started', + message: 'Preparing environment', + }, + RUNTIME_ID, + NEXT_RUNTIME_ID + ); + const before = structuredClone([...fixture.values]); + const events = fixture.eventQueries.findByEntityPrefix(''); + await expect(fixture.session.receiveSandboxControlPreparing(preparing)).resolves.toEqual({ + applied: false, + }); + expect([...fixture.values]).toEqual(before); + expect(fixture.eventQueries.findByEntityPrefix('')).toEqual(events); + }); + + it('applies preparation from a non-head trigger with its own pending attach proof', async () => { + const fixture = sessionFixture(); + const attach = deferred(); + fixture.setStatus({ + physical: 'running', + connection: 'ready', + wrapperInstanceId: RUNTIME_ID, + operationResults: true, + }); + delegateRequest(fixture, 'session.attach', () => attach.promise); + await fixture.admit('trigger'); + await fixture.flush(); + const trigger = fixture.record('trigger'); + const attemptId = trigger?.preparationAttemptId; + if (!trigger || !attemptId) throw new Error('Missing trigger attach authority'); + fixture.storage.kv.put('session_messages', [ + { messageId: 'head', state: 'queued', wrapperInstanceId: RUNTIME_ID }, + trigger, + ]); + await expect( + fixture.session.receiveSandboxControlPreparing( + receiptedPreparing( + 1, + { + version: 2, + attemptId, + triggerMessageId: 'trigger', + revision: 1, + timestamp: Date.now(), + step: 'workspace_setup', + action: 'attempt_started', + message: 'Preparing environment', + }, + RUNTIME_ID, + NEXT_RUNTIME_ID + ) + ) + ).resolves.toEqual({ applied: true }); + }); + + it.each([ + 'wrong_message', + 'wrong_attempt', + 'undispatched', + 'cancelled', + 'expired', + 'wrapper_mismatch', + 'scope_mismatch', + 'failed_result', + 'mismatched_result', + ] as const)('rejects pending preparation with %s attach authority', async invalid => { + const fixture = sessionFixture(); + const attach = deferred(); + fixture.setStatus({ + physical: 'running', + connection: 'ready', + wrapperInstanceId: RUNTIME_ID, + operationResults: true, + }); + delegateRequest(fixture, 'session.attach', () => attach.promise); + await fixture.admit('a'); + await fixture.flush(); + const message = fixture.record('a'); + const proof = message?.operations?.attach; + const attemptId = message?.preparationAttemptId; + if (!message || !proof || !attemptId) throw new Error('Missing pending attach authority'); + const triggerMessageId = invalid === 'wrong_message' ? 'missing' : 'a'; + const eventAttemptId = invalid === 'wrong_attempt' ? 'wrong-attempt' : attemptId; + const altered: SessionMessageRecord = { + ...message, + ...(invalid === 'cancelled' + ? { cancellation: { operationId: 'cancel', deadlineAt: Date.now() } } + : {}), + ...(invalid === 'wrapper_mismatch' ? { wrapperInstanceId: NEXT_RUNTIME_ID } : {}), + operations: + invalid === 'undispatched' + ? { attach: { ...proof, dispatched: false } } + : invalid === 'expired' + ? { + attach: { + ...proof, + authorization: { + ...proof.authorization, + dispatchDeadlineAt: Date.now() - SESSION_DELIVERY_TIMEOUT_MS, + }, + }, + } + : invalid === 'scope_mismatch' + ? { + attach: { + ...proof, + authorization: { + ...proof.authorization, + session: { ...proof.authorization.session, directory: '/wrong-directory' }, + }, + }, + } + : invalid === 'failed_result' + ? { + attach: { + ...proof, + result: { + ok: false, + error: { code: 'not_ready', message: 'Attach failed', retryable: true }, + }, + }, + } + : invalid === 'mismatched_result' + ? { + attach: { + ...proof, + result: { + ok: true, + result: { + attached: true, + nativeRuntimeId: '11111111-1111-4111-8111-111111111111', + }, + }, + }, + } + : { attach: proof }, + }; + fixture.storage.kv.put('session_messages', [altered]); + const preparing = receiptedPreparing( + 1, + { + version: 2, + attemptId: eventAttemptId, + triggerMessageId, + revision: 1, + timestamp: Date.now(), + step: 'workspace_setup', + action: 'attempt_started', + message: 'Preparing environment', + }, + RUNTIME_ID, + NEXT_RUNTIME_ID + ); + const before = structuredClone([...fixture.values]); + const events = fixture.eventQueries.findByEntityPrefix(''); + await expect(fixture.session.receiveSandboxControlPreparing(preparing)).resolves.toEqual({ + applied: false, + }); + expect([...fixture.values]).toEqual(before); + expect(fixture.eventQueries.findByEntityPrefix('')).toEqual(events); + }); + + it('acknowledges queued duplicates before attachment and settled duplicates after a fence rebound', async () => { + const fixture = sessionFixture(); + const attach = deferred(); + fixture.setStatus({ + physical: 'running', + connection: 'ready', + wrapperInstanceId: RUNTIME_ID, + operationResults: true, + }); + delegateRequest(fixture, 'session.attach', () => attach.promise); + await fixture.admit('a'); + await fixture.flush(); + const message = fixture.record('a'); + const attemptId = message?.preparationAttemptId; + const authorization = message?.operations?.attach?.authorization; + if (!message || !attemptId || !authorization) + throw new Error('Missing pending attach authority'); + const preparing = receiptedPreparing( + 1, + { + version: 2, + attemptId, + triggerMessageId: 'a', + revision: 1, + timestamp: Date.now(), + step: 'workspace_setup', + action: 'attempt_started', + message: 'Preparing environment', + }, + RUNTIME_ID, + NEXT_RUNTIME_ID + ); + await expect(fixture.session.receiveSandboxControlPreparing(preparing)).resolves.toEqual({ + applied: true, + }); + const events = fixture.eventQueries.findByEntityPrefix(''); + fixture.storage.kv.put('native_runtime_fence', { + sandboxId: SANDBOX_ID, + wrapperInstanceId: RUNTIME_ID, + nativeRuntimeId: NEXT_RUNTIME_ID, + }); + fixture.values.delete('native_runtime_fence'); + await expect(fixture.session.receiveSandboxControlPreparing(preparing)).resolves.toEqual({ + applied: true, + }); + expect(fixture.eventQueries.findByEntityPrefix('')).toEqual(events); + + fixture.storage.kv.put('session_messages', [{ ...message, state: 'accepted' }]); + fixture.storage.kv.put('native_runtime_fence', { + sandboxId: SANDBOX_ID, + wrapperInstanceId: RUNTIME_ID, + nativeRuntimeId: '11111111-1111-4111-8111-111111111111', + attachmentEpoch: 1, + authorization, + }); + const settledMessages = structuredClone(fixture.values.get('session_messages')); + const settledReceipts = structuredClone(fixture.values.get('control_event_receipts')); + const fields = vi.spyOn(logger, 'withFields').mockReturnValue(logger); + try { + await expect(fixture.session.receiveSandboxControlPreparing(preparing)).resolves.toEqual({ + applied: true, + }); + expect(fields).toHaveBeenCalledWith( + expect.objectContaining({ + diagnosticEvent: 'session_preparing_result', + disposition: 'duplicate', + }) + ); + expect(fixture.eventQueries.findByEntityPrefix('')).toEqual(events); + expect(fixture.values.get('session_messages')).toEqual(settledMessages); + expect(fixture.values.get('control_event_receipts')).toEqual(settledReceipts); + const newReceipt = receiptedPreparing(2, preparing.payload, RUNTIME_ID, NEXT_RUNTIME_ID); + const before = structuredClone([...fixture.values]); + await expect(fixture.session.receiveSandboxControlPreparing(newReceipt)).resolves.toEqual({ + applied: false, + }); + expect([...fixture.values]).toEqual(before); + expect(fixture.eventQueries.findByEntityPrefix('')).toEqual(events); + expect(fields).toHaveBeenCalledWith( + expect.objectContaining({ + diagnosticEvent: 'session_preparing_result', + disposition: 'native_runtime_mismatch', + }) + ); + } finally { + fields.mockRestore(); + } + }); + it('retries a startup event until the authorized attach response registers its native fence', async () => { const fixture = sessionFixture(); const nativeRuntimeId = NEXT_RUNTIME_ID; @@ -2084,7 +3050,7 @@ describe('SandboxSession orchestration', () => { ...originalEvent, identity: { ...originalEvent.identity, nativeRuntimeId: NEXT_RUNTIME_ID }, }) - ).resolves.toEqual({ applied: false }); + ).resolves.toEqual({ applied: true }); await fixture.session.failWaitingMessages('late-native-failure', RUNTIME_ID, nativeRuntimeId); await fixture.session.invalidateTerminalRuntime({ sandboxId: SANDBOX_ID, diff --git a/services/cloud-agent-next/src/sandbox-session/session-message-queue.ts b/services/cloud-agent-next/src/sandbox-session/session-message-queue.ts index 5af40bb96d..417c6edcac 100644 --- a/services/cloud-agent-next/src/sandbox-session/session-message-queue.ts +++ b/services/cloud-agent-next/src/sandbox-session/session-message-queue.ts @@ -352,6 +352,37 @@ export function failQueuedMessage( ); } +export function cancelPendingMessage( + messages: readonly SessionMessageRecord[], + messageId: string +): { dropped: boolean; messages?: SessionMessageRecord[] } { + const target = messages.find(message => message.messageId === messageId); + if (!target) return { dropped: false }; + if (target.state === 'cancelled' && target.failedReason === 'queued_message_cancelled') { + return { dropped: true }; + } + if ( + target.state !== 'queued' || + target.acceptedAt !== undefined || + target.unresolvedDispatch || + target.preparationAttemptId !== undefined || + target.deliveryDeadlineAt !== undefined || + target.wrapperInstanceId !== undefined || + target.operations !== undefined || + target.cancellation !== undefined + ) { + return { dropped: false }; + } + return { + dropped: true, + messages: messages.map(message => + message.messageId === messageId + ? { ...message, state: 'cancelled', failedReason: 'queued_message_cancelled' } + : message + ), + }; +} + export function acceptQueuedMessage( messages: readonly SessionMessageRecord[], messageId: string, diff --git a/services/cloud-agent-next/src/session-service.test.ts b/services/cloud-agent-next/src/session-service.test.ts index dd988ba3da..004483e548 100644 --- a/services/cloud-agent-next/src/session-service.test.ts +++ b/services/cloud-agent-next/src/session-service.test.ts @@ -2825,6 +2825,12 @@ describe('SessionService.buildWrapperSessionReadyAndPromptRequests', () => { expect(result.readyRequest.preparation?.attemptId).toBe('attempt-from-do'); }); + it('omits wrapper preparation metadata when the delivery plan does not need preparation', async () => { + const result = await buildPromptWrapperRequests(createMetadata()); + + expect(result.readyRequest).not.toHaveProperty('preparation'); + }); + it('uses direct GitLab authentication for a resumed DIND session', async () => { const result = await buildPromptWrapperRequests({ ...createMetadata({ preparedAt: 1 }), @@ -2915,6 +2921,7 @@ describe('SessionService.buildWrapperSessionReadyAndPromptRequests', () => { const result = await service.buildWrapperSessionReadyAndPromptRequests({ env, plan: { + preparation: { attemptId: 'prepare-payload' }, scope: { sessionId: 'agent_test', userId: 'user_test', @@ -2981,6 +2988,7 @@ describe('SessionService.buildWrapperSessionReadyAndPromptRequests', () => { setupCommands: ['pnpm install'], }, preparation: { + attemptId: 'prepare-payload', triggerMessageId: 'msg_018f1e2d3c4bPayloadTestAAAA', }, }); diff --git a/services/cloud-agent-next/src/session-service.ts b/services/cloud-agent-next/src/session-service.ts index 3cb68c4038..07d1e1e3f1 100644 --- a/services/cloud-agent-next/src/session-service.ts +++ b/services/cloud-agent-next/src/session-service.ts @@ -2214,13 +2214,17 @@ export class SessionService { ...(profile.runtimeSkills?.length ? { runtimeSkills: profile.runtimeSkills } : {}), }, session, - preparation: { - // Reuse the attempt the DO allocated (and may already have started - // with early sandbox-provisioning steps) so the wrapper's bootstrap - // events continue the same attempt instead of opening a second one. - attemptId: plan.preparation?.attemptId ?? crypto.randomUUID(), - triggerMessageId: turn.messageId, - }, + ...(plan.preparation + ? { + preparation: { + // Reuse the attempt the DO allocated (and may already have started + // with early sandbox-provisioning steps) so the wrapper's bootstrap + // events continue the same attempt instead of opening a second one. + attemptId: plan.preparation.attemptId, + triggerMessageId: turn.messageId, + }, + } + : {}), }; if (turn.type === 'command') { diff --git a/services/cloud-agent-next/src/session/agent-runtime.test.ts b/services/cloud-agent-next/src/session/agent-runtime.test.ts index 8cfaf18697..11e20c5548 100644 --- a/services/cloud-agent-next/src/session/agent-runtime.test.ts +++ b/services/cloud-agent-next/src/session/agent-runtime.test.ts @@ -198,17 +198,20 @@ describe('AgentRuntime', () => { }, }); - const result = await runtime.send(createPlan(), { - onProgress: (step, message) => { - progress.push({ step, message }); - }, - onWorkspaceReady: async workspace => { - workspaces.push(workspace); - }, - onAccepted: async delivery => { - accepted.push(delivery); - }, - }); + const result = await runtime.send( + { ...createPlan(), preparation: { attemptId: 'prepare-cold' } }, + { + onProgress: (step, message) => { + progress.push({ step, message }); + }, + onWorkspaceReady: async workspace => { + workspaces.push(workspace); + }, + onAccepted: async delivery => { + accepted.push(delivery); + }, + } + ); const wrapperState = await getWrapperRuntimeState(storage); const physicalLease = await getWrapperLease(storage); const [deliveredPlan] = deliveredPlans; @@ -231,6 +234,7 @@ describe('AgentRuntime', () => { wrapperGeneration: wrapperState.wrapperGeneration, wrapperConnectionId: wrapperState.wrapperConnectionId, }); + expect(deliveredPlan?.preparation).toEqual({ attemptId: 'prepare-cold' }); expect(progress).toEqual([{ step: 'kilo_server', message: 'Starting Kilo...' }]); expect(workspaces).toEqual([ready]); expect(accepted).toEqual([ @@ -248,6 +252,68 @@ describe('AgentRuntime', () => { expect(wrapperState.nextPingAt).toBe(acceptedAt + 60_000); }); + it('omits preparation progress when reusing an existing physical wrapper', async () => { + const storage = createMemoryStorage([ + [ + 'wrapper_runtime_state', + { wrapperGeneration: 3, wrapperConnectionId: 'conn_hot', wrapperRunId: 'wr_hot' }, + ], + [ + 'wrapper_lease', + { + state: 'owns_wrapper', + nextInstanceGeneration: 2, + instance: { instanceId: 'instance_hot', instanceGeneration: 1 }, + }, + ], + ]); + const execute = vi.fn( + async ( + plan: FencedWrapperDispatchRequest, + options?: { + onProgress?: (step: string, message: string) => void; + onWorkspaceReady?: (workspace: WorkspaceReady) => Promise; + } + ) => { + expect(plan.preparation).toBeUndefined(); + expect(options?.onProgress).toBeUndefined(); + await options?.onWorkspaceReady?.(createWorkspaceReady()); + return { kiloSessionId: 'kilo_runtime' }; + } + ); + const sandbox = { + discoverSessionWrappers: vi.fn().mockResolvedValue({ + status: 'present', + observed: [ + { + representation: 'process', + id: 'wrapper-hot', + port: 5_000, + instanceId: 'instance_hot', + instanceGeneration: 1, + }, + ], + }), + } as unknown as AgentSandbox; + const onProgress = vi.fn(); + const runtime = createAgentRuntime({ + storage, + env: {} as Env, + getMetadata: async () => createMetadata(), + getOrchestratorOverride: () => ({ execute }), + getSessionIdForLogs: () => 'agent_runtime', + sendToWrapper: () => false, + createAgentSandbox: () => sandbox, + }); + + await expect( + runtime.send({ ...createPlan(), preparation: { attemptId: 'prepare-hot' } }, { onProgress }) + ).resolves.toMatchObject({ success: true, outcome: 'accepted' }); + + expect(execute).toHaveBeenCalledOnce(); + expect(onProgress).not.toHaveBeenCalled(); + }); + it('fences the dispatching message until acceptance bookkeeping completes', async () => { const storage = createMemoryStorage(); const messageId = createPlan().turn.messageId; diff --git a/services/cloud-agent-next/src/session/agent-runtime.ts b/services/cloud-agent-next/src/session/agent-runtime.ts index f1d552c251..a3ca72d9b4 100644 --- a/services/cloud-agent-next/src/session/agent-runtime.ts +++ b/services/cloud-agent-next/src/session/agent-runtime.ts @@ -366,8 +366,10 @@ export function createAgentRuntime(dependencies: AgentRuntimeDependencies): Agen }) .info('AgentRuntime delivering pending message to wrapper'); + const deliveryPlan = { ...plan }; + if (!allocatedPhysicalInstance) delete deliveryPlan.preparation; const fencedPlan: FencedWrapperDispatchRequest = { - ...plan, + ...deliveryPlan, wrapper: { ...plan.wrapper, fence: { @@ -382,7 +384,7 @@ export function createAgentRuntime(dependencies: AgentRuntimeDependencies): Agen try { await getOrchestrator().execute(fencedPlan, { ...(leasedInstance ? { leasedInstance } : {}), - onProgress: hooks.onProgress, + ...(allocatedPhysicalInstance && hooks.onProgress ? { onProgress: hooks.onProgress } : {}), onWorkspaceReady: async ready => { const readyAt = Date.now(); const readyDeadlineAt = readyAt + READY_ONLY_IDLE_MS; diff --git a/services/cloud-agent-next/src/session/preparation-history.test.ts b/services/cloud-agent-next/src/session/preparation-history.test.ts index 03e9896b39..10bc92759a 100644 --- a/services/cloud-agent-next/src/session/preparation-history.test.ts +++ b/services/cloud-agent-next/src/session/preparation-history.test.ts @@ -124,21 +124,23 @@ describe('materializePreparationEvent', () => { expect(step.safeError).toBe('Setup command failed'); }); - it('restarts a terminal attempt with a fresh startedAt', () => { + it.each([ + { outcome: { status: 'completed' as const }, expected: { status: 'completed' } }, + { + outcome: { status: 'failed' as const, safeError: 'boom' }, + expected: { status: 'failed', safeError: 'boom' }, + }, + ])('does not restart a $expected.status attempt', ({ outcome, expected }) => { const eventQueries = createMemoryEventQueries(); const { attemptId, lastEventAt } = seedRunningAttempt(eventQueries, { startedAt: 1000 }); - finalizePreparationAttempt(eventQueries, attemptId, { - status: 'failed', - safeError: 'boom', - timestamp: lastEventAt, - }); - const failedRevision = readAttempt(eventQueries, attemptId).revision; + finalizePreparationAttempt(eventQueries, attemptId, { ...outcome, timestamp: lastEventAt }); + const terminalAttempt = readAttempt(eventQueries, attemptId); - materializePreparationEvent(eventQueries, storedEvent(90_000), { + const applied = materializePreparationEvent(eventQueries, storedEvent(90_000), { version: 2, attemptId, triggerMessageId: 'msg-1', - revision: failedRevision + 1, + revision: terminalAttempt.revision + 1, timestamp: 90_000, step: 'workspace_setup', message: 'Preparing environment', @@ -146,8 +148,42 @@ describe('materializePreparationEvent', () => { }); const attempt = readAttempt(eventQueries, attemptId); - expect(attempt.status).toBe('running'); - expect(attempt.startedAt).toBe(90_000); + expect(applied).toBe(false); + expect(attempt).toEqual({ ...terminalAttempt, ...expected }); + }); + + it('retains late step history without reopening a terminal attempt', () => { + const eventQueries = createMemoryEventQueries(); + const { attemptId, lastEventAt } = seedRunningAttempt(eventQueries); + finalizePreparationAttempt(eventQueries, attemptId, { + status: 'completed', + timestamp: lastEventAt, + }); + const terminalAttempt = readAttempt(eventQueries, attemptId); + + const applied = materializePreparationEvent(eventQueries, storedEvent(90_000), { + version: 2, + attemptId, + triggerMessageId: 'msg-1', + revision: terminalAttempt.revision + 1, + timestamp: 90_000, + step: 'setup_commands', + message: 'Installing dependencies', + action: 'step_started', + stepId: 'command:install', + kind: 'setup_command', + label: 'Install dependencies', + }); + + expect(applied).toBe(true); + expect(readAttempt(eventQueries, attemptId)).toMatchObject({ + status: 'completed', + revision: terminalAttempt.revision + 1, + }); + expect(readStep(eventQueries, attemptId, 'command:install')).toMatchObject({ + status: 'running', + label: 'Install dependencies', + }); }); }); @@ -342,28 +378,25 @@ describe('cloudStatusForPreparingEvent', () => { timestamp: 1000, }; - it('maps applied v2 events by action', () => { + it('never projects v2 events, including applied events with step and message data', () => { expect( cloudStatusForPreparingEvent( { ...v2, step: 'cloning', message: 'Cloning…', action: 'step_progress' }, true ) - ).toEqual({ type: 'preparing', step: 'cloning', message: 'Cloning…' }); + ).toBeNull(); expect( cloudStatusForPreparingEvent( { ...v2, step: 'ready', message: 'Preparation complete', action: 'attempt_completed' }, true ) - ).toEqual({ type: 'ready' }); + ).toBeNull(); expect( cloudStatusForPreparingEvent( { ...v2, step: 'failed', message: 'nope', action: 'attempt_failed', safeError: 'boom' }, true ) - ).toEqual({ type: 'error', message: 'boom' }); - }); - - it('suppresses the broadcast for stale v2 events', () => { + ).toBeNull(); expect( cloudStatusForPreparingEvent( { ...v2, step: 'cloning', message: 'Cloning…', action: 'step_progress' }, diff --git a/services/cloud-agent-next/src/session/preparation-history.ts b/services/cloud-agent-next/src/session/preparation-history.ts index f85060d802..33b1926ce2 100644 --- a/services/cloud-agent-next/src/session/preparation-history.ts +++ b/services/cloud-agent-next/src/session/preparation-history.ts @@ -190,7 +190,7 @@ export function materializePreparationEvent( : undefined; if (data.action === 'attempt_started') { - if (attempt && attempt.revision >= data.revision) return false; + if (attempt && (isTerminal(attempt.status) || attempt.revision >= data.revision)) return false; const handoff = attempt?.status === 'running'; attempt = { id: attemptId, @@ -213,9 +213,10 @@ export function materializePreparationEvent( return true; } - if (!attempt || data.revision <= attempt.revision || isTerminal(attempt.status)) return false; + if (!attempt || data.revision <= attempt.revision) return false; if (data.action === 'attempt_completed' || data.action === 'attempt_failed') { + if (isTerminal(attempt.status)) return false; attempt = { ...attempt, status: data.action === 'attempt_completed' ? 'completed' : 'failed', @@ -464,28 +465,17 @@ export function reconcileStalePreparationAttempts( /** * Map a 'preparing' stream event to the `cloud.status` broadcast that should - * accompany it, or null when none should be sent. Stale v2 events (ones the - * materializer rejected) must not regress a ready session back to - * 'preparing' — that strands the chat input in its disabled state. + * accompany legacy preparation, or null when none should be sent. V2 status + * comes from the materialized attempt in the SDK. */ export function cloudStatusForPreparingEvent( data: unknown, - applied: boolean + _applied: boolean ): CloudStatusData['cloudStatus'] | null { if (!isRecord(data)) return null; + if (data.version === 2) return null; const step = typeof data.step === 'string' ? { step: data.step } : {}; const message = typeof data.message === 'string' ? { message: data.message } : {}; - if (data.version === 2) { - if (!applied) return null; - if (data.action === 'attempt_completed') return { type: 'ready' }; - if (data.action === 'attempt_failed') { - return { - type: 'error', - ...(typeof data.safeError === 'string' ? { message: data.safeError } : message), - }; - } - return { type: 'preparing', ...step, ...message }; - } if (data.step === 'ready') return { type: 'ready' }; if (data.step === 'failed') return { type: 'error', ...message }; return { type: 'preparing', ...step, ...message }; diff --git a/services/cloud-agent-next/src/session/preparation-progress.ts b/services/cloud-agent-next/src/session/preparation-progress.ts index 0876967417..c3bfa0ca13 100644 --- a/services/cloud-agent-next/src/session/preparation-progress.ts +++ b/services/cloud-agent-next/src/session/preparation-progress.ts @@ -73,8 +73,10 @@ export function createPreparationProgressRecorder(options: { } function onProgress(step: string, message: string): boolean { + const existing = readPreparationAttempt(eventQueries, attemptId); + if (existing?.status === 'completed' || existing?.status === 'failed') return false; const key = step as PreparingStep; - if (!readPreparationAttempt(eventQueries, attemptId)) { + if (!existing) { emit('workspace_setup', 'Preparing environment', { action: 'attempt_started' }); } const stepId = `phase:${key}`; diff --git a/services/cloud-agent-next/src/shared/control-diagnostics.test.ts b/services/cloud-agent-next/src/shared/control-diagnostics.test.ts new file mode 100644 index 0000000000..0d73bdd7c5 --- /dev/null +++ b/services/cloud-agent-next/src/shared/control-diagnostics.test.ts @@ -0,0 +1,66 @@ +import { describe, expect, it } from 'vitest'; +import { heartbeatReasonFrom } from './sandbox-control-protocol.js'; +import { + classifyRetirementCause, + createControlDiagnosticRecord, + diagnosticDetail, +} from './control-diagnostics.js'; + +describe('classifyRetirementCause', () => { + it('maps feed machine reasons that previously became unknown', () => { + expect(classifyRetirementCause('feed_failed')).toBe('event_feed_unhealthy'); + expect(classifyRetirementCause('feed_stale')).toBe('event_feed_unhealthy'); + expect(classifyRetirementCause('feed_ended')).toBe('event_feed_unhealthy'); + }); + + it('keeps process exit distinct from unknown', () => { + expect(classifyRetirementCause('process_exited')).toBe('process_exited'); + }); + + it('classifies session event delivery failures', () => { + expect(classifyRetirementCause('Session event delivery failed')).toBe( + 'outcome_delivery_failed' + ); + expect(classifyRetirementCause('Session event delivery unconfirmed')).toBe( + 'outcome_delivery_failed' + ); + }); + + it('falls back through later reasons', () => { + expect(classifyRetirementCause('mystery', 'control_disconnected')).toBe('control_disconnected'); + expect(classifyRetirementCause('mystery')).toBe('unknown'); + }); +}); + +describe('heartbeatReasonFrom', () => { + it('passes feed and process codes through to the worker heartbeat', () => { + expect(heartbeatReasonFrom('feed_failed')).toBe('feed_failed'); + expect(heartbeatReasonFrom('process_exited')).toBe('process_exited'); + }); + + it('does not invent a machine code for human shutdown strings', () => { + expect(heartbeatReasonFrom('Wrapper received SIGTERM')).toBe('shutdown'); + }); +}); + +describe('diagnosticDetail', () => { + it('keeps a bounded reason on lifecycle records', () => { + expect(diagnosticDetail('feed_failed')).toBe('feed_failed'); + expect(diagnosticDetail(` ${'x'.repeat(200)} `)?.length).toBe(128); + const record = createControlDiagnosticRecord( + 'wrapper.lifecycle', + { + phase: 'stopping', + exitCode: 1, + retirementCause: 'event_feed_unhealthy', + detail: 'feed_failed', + }, + 1 + ); + expect(record?.fields).toMatchObject({ + phase: 'stopping', + retirementCause: 'event_feed_unhealthy', + detail: 'feed_failed', + }); + }); +}); diff --git a/services/cloud-agent-next/src/shared/control-diagnostics.ts b/services/cloud-agent-next/src/shared/control-diagnostics.ts index c1488e2ca4..8096d2800f 100644 --- a/services/cloud-agent-next/src/shared/control-diagnostics.ts +++ b/services/cloud-agent-next/src/shared/control-diagnostics.ts @@ -5,10 +5,13 @@ import { worktreeDeletePayloadSchema, } from './sandbox-control-protocol.js'; +export const OWNED_PROCESS_CLEANUP_UNREAPED = 'Owned process cleanup unreaped'; export const CONTROL_LOG_MAX_BATCH_BYTES = 256 * 1024; export const CONTROL_LOG_MAX_BATCH_RECORDS = 128; export const CONTROL_LOG_MAX_BUFFER_RECORDS = 512; export const CONTROL_LOG_MAX_RECORD_BYTES = 4096; +export const CONTROL_LOG_MAX_ARCHIVE_BYTES = 8 * 1024 * 1024; +export const CONTROL_LOG_ARCHIVE_NAME = 'files.tar.gz'; export const CONTROL_LOG_GRANT_SECONDS = 4 * 60 * 60; export const controlLogUploadResults = [ 'accepted', @@ -161,6 +164,7 @@ export const controlDiagnosticFieldsSchema = z .optional(), errorCode: z.enum([...controlErrorCodes, 'other']).optional(), retryable: z.boolean().optional(), + detail: z.string().min(1).max(128).optional(), scopeId: identifier.optional(), worktreeId: worktreeDeletePayloadSchema.shape.worktreeId.optional(), sessionId: identifier.optional(), @@ -218,6 +222,47 @@ export const controlDiagnosticRecordSchema = z .strict(); export type ControlDiagnosticRecord = z.infer; +export type RetirementCause = NonNullable; + +const retirementCauseByReason = new Map([ + ['Kilo event feed is no longer healthy', 'event_feed_unhealthy'], + ['feed_stale', 'event_feed_unhealthy'], + ['feed_reconnected', 'event_feed_unhealthy'], + ['feed_ended', 'event_feed_unhealthy'], + ['feed_failed', 'event_feed_unhealthy'], + ['process_exited', 'process_exited'], + ['credential_refresh_failed', 'credential_refresh_failed'], + ['Sandbox control connection lost', 'control_disconnected'], + ['control_disconnected', 'control_disconnected'], + ['Preparation event delivery failed', 'preparation_delivery_failed'], + ['Sandbox shutting down', 'requested_shutdown'], + ['Wrapper received SIGTERM', 'sigterm'], + ['Wrapper received SIGINT', 'sigint'], + ['Wrapper uncaught exception', 'uncaught_exception'], + ['Wrapper unhandled rejection', 'unhandled_rejection'], + ['Kilo cancellation failed', 'cancellation_failed'], + ['Kilo cancellation was not confirmed', 'cancellation_failed'], + ['Native cancellation did not settle', 'cancellation_failed'], + ['Session outcome delivery failed', 'outcome_delivery_failed'], + ['Session event delivery failed', 'outcome_delivery_failed'], + ['Execution exceeded the 60 minute limit', 'execution_deadline'], + ['Session preparation timed out', 'preparation_deadline'], +]); + +export function diagnosticDetail(value: string): string | undefined { + const detail = value.trim().slice(0, 128); + return detail === '' ? undefined : detail; +} + +export function classifyRetirementCause(...reasons: string[]): RetirementCause { + for (const reason of reasons) { + const mapped = retirementCauseByReason.get(reason); + if (mapped) return mapped; + if (reason.startsWith('feed_')) return 'event_feed_unhealthy'; + if (reason.startsWith('Session event delivery')) return 'outcome_delivery_failed'; + } + return 'unknown'; +} export const controlLogIdentitySchema = z .object({ @@ -246,6 +291,15 @@ export function diagnosticSyncStatus(value: unknown): z.infer; + +export function heartbeatReasonFrom(reason: string): SandboxKiloHeartbeatReason { + const parsed = sandboxKiloHeartbeatReasonSchema.safeParse(reason); + return parsed.success ? parsed.data : 'shutdown'; +} + export const sandboxHeartbeatPayloadSchema = z .object({ state: z.enum(['idle', 'active', 'finalizing']), @@ -238,18 +256,7 @@ export const sandboxHeartbeatPayloadSchema = z .object({ ready: z.boolean(), version: SandboxRuntimeVersionSchema.nullable().optional().catch(undefined), - reason: z - .enum([ - 'feed_stale', - 'feed_reconnected', - 'feed_ended', - 'feed_failed', - 'process_exited', - 'credential_refresh_failed', - 'control_disconnected', - 'shutdown', - ]) - .optional(), + reason: sandboxKiloHeartbeatReasonSchema.optional(), }) .strict(), sessions: z.array( @@ -659,7 +666,6 @@ export const sandboxEventPublicationPayloadSchema = z.discriminatedUnion('event' .object({ event: z.literal('session.event'), receiptId: z.string().uuid(), - receiptHash: z.string().regex(/^[a-f0-9]{64}$/), sequence: z.number().int().positive().max(Number.MAX_SAFE_INTEGER), session: sessionEventIdentitySchema, payload: sessionEventPayloadSchema, @@ -669,7 +675,6 @@ export const sandboxEventPublicationPayloadSchema = z.discriminatedUnion('event' .object({ event: z.literal('session.preparing'), receiptId: z.string().uuid(), - receiptHash: z.string().regex(/^[a-f0-9]{64}$/), sequence: z.number().int().positive().max(Number.MAX_SAFE_INTEGER), session: sessionEventIdentitySchema, payload: sessionPreparingPayloadSchema, diff --git a/services/cloud-agent-next/test/e2e/client.ts b/services/cloud-agent-next/test/e2e/client.ts index 53fc37599c..f85a831304 100644 --- a/services/cloud-agent-next/test/e2e/client.ts +++ b/services/cloud-agent-next/test/e2e/client.ts @@ -646,11 +646,14 @@ export function messageIdFromEvent(event: StreamEvent): string | undefined { export function isMessageCompleted( event: StreamEvent | null, messageId: string -): event is StreamEvent & { streamEventType: 'cloud.message.completed' } { +): event is StreamEvent { return ( event !== null && - event.streamEventType === 'cloud.message.completed' && - messageIdFromEvent(event) === messageId + ((event.streamEventType === 'cloud.message.completed' && + messageIdFromEvent(event) === messageId) || + (event.streamEventType === 'complete' && + Array.isArray(event.data.messageIds) && + event.data.messageIds.includes(messageId))) ); } @@ -756,9 +759,9 @@ export function openStream( event => messageId === undefined ? TERMINAL_STREAM_TYPES.has(event.streamEventType) - : (event.streamEventType === 'cloud.message.completed' || - event.streamEventType === 'cloud.message.failed') && - messageIdFromEvent(event) === messageId, + : isMessageCompleted(event, messageId) || + (event.streamEventType === 'cloud.message.failed' && + messageIdFromEvent(event) === messageId), timeoutMs ), get receivedCount() { diff --git a/services/cloud-agent-next/test/e2e/lifecycle.ts b/services/cloud-agent-next/test/e2e/lifecycle.ts index b0bbe1ac56..5143df32e5 100644 --- a/services/cloud-agent-next/test/e2e/lifecycle.ts +++ b/services/cloud-agent-next/test/e2e/lifecycle.ts @@ -1226,17 +1226,6 @@ export async function lifecycleCold(args: LifecycleArgs): Promise new Map()); +vi.mock('@kilocode/db/client', () => ({ getWorkerDb: () => ({}) })); +vi.mock('@kilocode/worker-utils/cloud-agent-session-access', () => ({ + queryAccessibleCloudAgentSession: async ( + _db: unknown, + input: { kiloUserId: string; cloudAgentSessionId: string } + ) => + access.get(input.cloudAgentSessionId) === input.kiloUserId + ? { kiloSessionId: 'ses_abcdefghijklmnopqrstuvwxyz', organizationId: null } + : null, +})); + +const api = router({ ...createSessionSendHandlers(), ...createSessionManagementHandlers() }); +const ownerId = 'queue-owner'; +const agent = { mode: 'code', model: 'anthropic/claude-sonnet-4' }; +const id = (index: number) => `msg_${index.toString(16).padStart(12, '0')}AbCdEfGhIjKlMn`; +type Session = ReturnType; + +async function request(path: 'send' | 'cancelQueuedMessage', input: unknown, userId = ownerId) { + const req = new Request(`https://queue.test/trpc/${path}`, { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify(input), + }); + return fetchRequestHandler({ + endpoint: '/trpc', + req, + router: api, + createContext: () => ({ + env: { ...env, HYPERDRIVE: { ...env.HYPERDRIVE, connectionString: 'postgres://unused' } }, + userId, + authToken: userId ? 'test-token' : '', + request: req, + }), + }); +} + +function send(sessionId: string, index: number, model = agent.model) { + return request('send', { + cloudAgentSessionId: sessionId, + message: { id: id(index), prompt: `queued ${index}` }, + agent: { ...agent, model }, + }); +} + +function cancel(sessionId: string, index: number, userId = ownerId) { + return request('cancelQueuedMessage', { sessionId, messageId: id(index) }, userId); +} + +function snapshot(session: Session) { + return runInDurableObject(session, async (_instance, state) => ({ + messages: state.storage.kv.get('session_messages') ?? [], + metadata: state.storage.kv.get('session_metadata'), + events: createEventQueries(drizzle(state.storage), state.storage.sql).findByFilters({}), + alarm: await state.storage.getAlarm(), + })); +} + +async function fixture() { + const sessionId = `workspace_${crypto.randomUUID()}`; + access.set(sessionId, ownerId); + const session = env.SANDBOX_SESSION.getByName(`${ownerId}:${sessionId}`); + await session.registerSession({ + identity: { sessionId, userId: ownerId }, + auth: { kiloSessionId: 'ses_abcdefghijklmnopqrstuvwxyz', kilocodeToken: 'test-token' }, + agent, + }); + await runInDurableObject(session, (_instance, state) => { + state.storage.kv.put('session_messages', [ + { + ...createSessionMessageRecord({ + turn: { type: 'prompt', messageId: id(0), prompt: 'accepted head' }, + agent, + }), + state: 'accepted', + acceptedAt: Date.now(), + lastActivityAt: Date.now(), + }, + ] satisfies SessionMessageRecord[]); + }); + const broadcast = await runInDurableObject(session, instance => { + const observed = vi.fn(instance['broadcastQueuedMessage'].bind(instance)); + instance['broadcastQueuedMessage'] = observed; + return observed; + }); + return { session, sessionId, broadcast }; +} + +beforeEach(() => { + vi.spyOn(globalThis, 'fetch').mockImplementation(async () => Response.json({ valid: true })); +}); +afterEach(async () => { + await reset(); + vi.restoreAllMocks(); + access.clear(); +}); + +describe('public control queue capacity and cancellation', () => { + it('retains the legacy public cancellation method and once-only cancellation event', async () => { + const sessionId = `agent_${crypto.randomUUID()}`; + access.set(sessionId, ownerId); + const session = env.CLOUD_AGENT_SESSION.getByName(`${ownerId}:${sessionId}`); + await session.registerSession({ + identity: { sessionId, userId: ownerId }, + auth: { kilocodeToken: 'test-token' }, + agent, + }); + const intent = { + turn: { type: 'prompt' as const, messageId: id(1), prompt: 'legacy queued' }, + agent, + }; + await runInDurableObject(session, async (_instance, state) => { + await storePendingSessionMessage( + state.storage, + createPendingSessionMessageFromIntent(intent) + ); + await putSessionMessageState(state.storage, createQueuedSessionMessageState(intent)); + }); + for (let attempt = 0; attempt < 2; attempt++) { + expect(await (await cancel(sessionId, 1)).json()).toMatchObject({ + result: { data: { dropped: true } }, + }); + } + await runInDurableObject(session, async (_instance, state) => { + expect(await listPendingSessionMessages(state.storage)).toEqual([]); + expect(await getSessionMessageState(state.storage, id(1))).toMatchObject({ + status: 'interrupted', + completionSource: 'canceled', + }); + expect( + createEventQueries(drizzle(state.storage), state.storage.sql).findByFilters({ + eventTypes: ['cloud.message.canceled'], + }) + ).toHaveLength(1); + }); + }); + + it('rejects overflow with HTTP 429 without durable or metadata side effects and allows replay at capacity', async () => { + const { session, sessionId, broadcast } = await fixture(); + for (let index = 1; index <= PENDING_SESSION_MESSAGE_LIMIT; index++) + expect((await send(sessionId, index)).status).toBe(200); + const before = await snapshot(session); + const overflow = await send(sessionId, 11, 'openai/gpt-4.1'); + expect(overflow.status).toBe(429); + expect(await overflow.json()).toMatchObject({ + error: { data: { code: 'TOO_MANY_REQUESTS', clientError: { retryable: true } } }, + }); + expect(await snapshot(session)).toEqual(before); + expect((await send(sessionId, 1)).status).toBe(200); + expect( + await session.admitSubmittedMessage({ + userId: ownerId, + turn: { type: 'prompt', id: id(0), prompt: 'accepted head' }, + agent, + }) + ).toMatchObject({ success: true, compatibilityDelivery: 'sent' }); + expect((await snapshot(session)).events).toEqual(before.events); + expect(broadcast).toHaveBeenCalledTimes(PENDING_SESSION_MESSAGE_LIMIT); + }); + + it('rechecks capacity when model validation yields while another request takes the last slot', async () => { + const { session, sessionId } = await fixture(); + for (let index = 1; index < PENDING_SESSION_MESSAGE_LIMIT; index++) + await send(sessionId, index); + let entered = false; + let released = false; + vi.mocked(globalThis.fetch).mockImplementationOnce(async () => { + entered = true; + while (!released) await new Promise(resolve => setTimeout(resolve, 1)); + return Response.json({ valid: true }); + }); + const slow = send(sessionId, 11, 'openai/gpt-4.1'); + let before: Awaited>; + try { + await vi.waitFor(() => expect(entered).toBe(true)); + expect((await send(sessionId, 10)).status).toBe(200); + before = await snapshot(session); + } finally { + released = true; + } + expect((await slow).status).toBe(429); + expect(await snapshot(session)).toEqual(before); + }); + + it('serializes concurrent admissions after asynchronous model validation', async () => { + const { session, sessionId } = await fixture(); + const responses = await Promise.all( + Array.from({ length: 16 }, (_, index) => send(sessionId, index + 1)) + ); + expect(responses.filter(response => response.status === 200)).toHaveLength( + PENDING_SESSION_MESSAGE_LIMIT + ); + expect(responses.filter(response => response.status === 429)).toHaveLength(6); + const stored = await snapshot(session); + expect(stored.messages.filter(message => message.state === 'queued')).toHaveLength( + PENDING_SESSION_MESSAGE_LIMIT + ); + for (const message of stored.messages.filter(message => message.state === 'queued')) { + expect(await session.getMessageResult(message.messageId)).toMatchObject({ + type: 'found', + result: { status: 'queued' }, + }); + } + }); + + it('cancels only the target, frees one slot, appends replacement at the tail and persists retry tombstones', async () => { + const { session, sessionId } = await fixture(); + for (let index = 1; index <= PENDING_SESSION_MESSAGE_LIMIT; index++) + await send(sessionId, index); + const before = await snapshot(session); + const responses = await Promise.all([cancel(sessionId, 4), cancel(sessionId, 4)]); + for (const response of responses) { + expect(response.status).toBe(200); + expect(await response.json()).toMatchObject({ result: { data: { dropped: true } } }); + } + const after = await snapshot(session); + expect(after.messages[0]).toEqual(before.messages[0]); + expect(after.messages.find(message => message.messageId === id(4))).toMatchObject({ + state: 'cancelled', + intent: before.messages[4].intent, + terminalAt: expect.any(Number), + }); + expect( + after.events.filter(event => event.stream_event_type === 'cloud.message.failed') + ).toHaveLength(1); + expect(JSON.parse(after.events.at(-1)?.payload ?? 'null')).toMatchObject({ + messageId: id(4), + status: 'interrupted', + accepted: false, + delivery: 'queued', + }); + expect(await session.getMessageResult(id(4))).toMatchObject({ + type: 'found', + result: { status: 'interrupted' }, + }); + expect((await send(sessionId, 11)).status).toBe(200); + expect((await send(sessionId, 12)).status).toBe(429); + expect( + (await snapshot(session)).messages + .filter(message => message.state === 'queued') + .map(message => message.messageId) + ).toEqual([1, 2, 3, 5, 6, 7, 8, 9, 10, 11].map(id)); + const persisted = await snapshot(session); + await abortAllDurableObjects(); + expect(await (await cancel(sessionId, 4)).json()).toMatchObject({ + result: { data: { dropped: true } }, + }); + expect((await send(sessionId, 4)).status).toBe(400); + expect(await snapshot(env.SANDBOX_SESSION.getByName(`${ownerId}:${sessionId}`))).toEqual( + persisted + ); + }); + + it('rolls back both the terminal event and state if cancellation persistence fails', async () => { + const { session, sessionId } = await fixture(); + await send(sessionId, 1); + const before = await snapshot(session); + await runInDurableObject(session, instance => { + const persist = instance['persistMessageLifecycleEvent'].bind(instance); + instance['persistMessageLifecycleEvent'] = vi.fn(persist).mockImplementationOnce(message => { + persist(message); + throw new Error('Injected cancellation persistence failure'); + }); + }); + expect((await cancel(sessionId, 1)).status).toBe(500); + expect(await snapshot(session)).toEqual(before); + expect(await (await cancel(sessionId, 1)).json()).toMatchObject({ + result: { data: { dropped: true } }, + }); + expect( + (await snapshot(session)).events.filter( + event => event.stream_event_type === 'cloud.message.failed' + ) + ).toHaveLength(1); + }); + + it('schedules the next queued head without dispatching cancelled work or renewing its budget', async () => { + const { session, sessionId } = await fixture(); + await send(sessionId, 1); + await send(sessionId, 2); + await runInDurableObject(session, async (_instance, state) => { + const messages = state.storage.kv.get('session_messages') ?? []; + state.storage.kv.put( + 'session_messages', + messages.filter(message => message.messageId !== id(0)) + ); + await state.storage.deleteAlarm(); + }); + const before = await snapshot(session); + expect(await (await cancel(sessionId, 1)).json()).toMatchObject({ + result: { data: { dropped: true } }, + }); + const after = await snapshot(session); + expect(after.messages.find(message => message.state === 'queued')).toEqual(before.messages[1]); + expect(after.alarm).not.toBeNull(); + expect(after.messages[0]).toMatchObject({ state: 'cancelled' }); + expect(after.messages[0].preparationAttemptId).toBeUndefined(); + }); + + it('denies unauthenticated and other-owner access and cannot cancel another session message', async () => { + const { session, sessionId } = await fixture(); + await send(sessionId, 1); + const before = await snapshot(session); + expect((await cancel(sessionId, 1, '')).status).toBe(401); + expect((await cancel(sessionId, 1, 'other-owner')).status).toBe(403); + const other = await fixture(); + expect(await (await cancel(other.sessionId, 1)).json()).toMatchObject({ + result: { data: { dropped: false } }, + }); + expect(await snapshot(session)).toEqual(before); + }); + + it.each([ + { state: 'accepted' as const, acceptedAt: 1 }, + { unresolvedDispatch: true as const }, + { preparationAttemptId: 'attempt-original', deliveryDeadlineAt: Date.now() + 60_000 }, + { wrapperInstanceId: 'wrapper-original' }, + { state: 'failed' as const, terminalAt: 1 }, + { state: 'cancelled' as const, terminalAt: 1 }, + ])('refuses accepted, ambiguous, preparing and stale targets: %j', async patch => { + const { session, sessionId } = await fixture(); + await send(sessionId, 1); + await runInDurableObject(session, (_instance, state) => { + const messages = state.storage.kv.get('session_messages') ?? []; + state.storage.kv.put( + 'session_messages', + messages.map(message => (message.messageId === id(1) ? { ...message, ...patch } : message)) + ); + }); + const before = await snapshot(session); + expect(await (await cancel(sessionId, 1)).json()).toMatchObject({ + result: { data: { dropped: false } }, + }); + expect(await (await cancel(sessionId, 99)).json()).toMatchObject({ + result: { data: { dropped: false } }, + }); + expect(await snapshot(session)).toEqual(before); + }); + + it('counts the preparing head and releases capacity on terminal settlement', async () => { + const { session, sessionId } = await fixture(); + for (let index = 1; index <= PENDING_SESSION_MESSAGE_LIMIT; index++) + await send(sessionId, index); + const acceptedHead = (await snapshot(session)).messages[0]; + await runInDurableObject(session, (_instance, state) => { + const messages = state.storage.kv.get('session_messages') ?? []; + state.storage.kv.put( + 'session_messages', + messages + .filter(message => message.messageId !== id(0)) + .map(message => + message.messageId === id(1) + ? { + ...message, + preparationAttemptId: 'original', + deliveryDeadlineAt: Date.now() + 60_000, + } + : message + ) + ); + }); + const preparing = await snapshot(session); + expect((await send(sessionId, 11)).status).toBe(429); + expect(await (await cancel(sessionId, 1)).json()).toMatchObject({ + result: { data: { dropped: false } }, + }); + expect(await snapshot(session)).toEqual(preparing); + await runInDurableObject(session, (instance, state) => { + const messages = state.storage.kv.get('session_messages') ?? []; + instance['saveMessages']([ + acceptedHead, + ...messages.map(message => + message.messageId === id(2) ? { ...message, state: 'failed' as const } : message + ), + ]); + }); + expect((await send(sessionId, 12)).status).toBe(200); + const state = await snapshot(session); + expect(state.messages.find(message => message.messageId === id(1))).toMatchObject({ + preparationAttemptId: 'original', + state: 'queued', + }); + expect((await send(sessionId, 2)).status).toBe(400); + }); +}); diff --git a/services/cloud-agent-next/test/unit/e2e/client.test.ts b/services/cloud-agent-next/test/unit/e2e/client.test.ts new file mode 100644 index 0000000000..71364a0251 --- /dev/null +++ b/services/cloud-agent-next/test/unit/e2e/client.test.ts @@ -0,0 +1,83 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +const sockets = vi.hoisted(() => ({ + instances: [] as Array<{ message: (data: string) => void }>, +})); + +vi.mock('ws', () => ({ + default: class { + private readonly handlers = new Map void>(); + + constructor() { + sockets.instances.push({ + message: data => this.handlers.get('message')?.(Buffer.from(data)), + }); + } + + on(event: string, handler: (value: unknown) => void): this { + this.handlers.set(event, handler); + return this; + } + + close(): void { + this.handlers.get('close')?.(undefined); + } + }, +})); + +import { isMessageCompleted, openStream, type StreamEvent } from '../../e2e/client.js'; + +function event(streamEventType: string, data: Record): StreamEvent { + return { + eventId: 1, + executionId: null, + sessionId: 'workspace_11111111-1111-4111-8111-111111111111', + streamEventType, + timestamp: new Date(0).toISOString(), + data, + }; +} + +describe('isMessageCompleted', () => { + beforeEach(() => { + sockets.instances = []; + }); + + it('accepts matching legacy and cloud completion events only', () => { + expect(isMessageCompleted(event('complete', { messageIds: ['message_1'] }), 'message_1')).toBe( + true + ); + expect( + isMessageCompleted(event('cloud.message.completed', { messageId: 'message_1' }), 'message_1') + ).toBe(true); + expect(isMessageCompleted(event('complete', { messageIds: ['other'] }), 'message_1')).toBe( + false + ); + expect( + isMessageCompleted(event('cloud.message.completed', { messageId: 'other' }), 'message_1') + ).toBe(false); + }); + + it.each([ + ['complete', { messageIds: ['message_1'] }], + ['cloud.message.completed', { messageId: 'message_1' }], + ])('waitForTerminal resolves matching %s events', async (streamEventType, data) => { + const stream = openStream( + { + workerUrl: 'http://worker.test', + user: { id: 'user_1', email: 'user@example.test', api_token_pepper: 'pepper' }, + nextAuthSecret: 'test-secret', + gitUrl: 'https://example.test/repo.git', + model: 'kilo/fake-deterministic', + fakeLlmUrl: 'http://fake.test', + }, + 'workspace_11111111-1111-4111-8111-111111111111' + ); + const socket = sockets.instances[0]; + if (!socket) throw new Error('Missing stream socket'); + const terminal = stream.waitForTerminal(100, 'message_1'); + socket.message(JSON.stringify(event(streamEventType, data))); + await expect(terminal).resolves.toMatchObject({ streamEventType, data }); + stream.close(); + }); +}); diff --git a/services/cloud-agent-next/wrapper/src/control/apply-attach.test.ts b/services/cloud-agent-next/wrapper/src/control/apply-attach.test.ts index 9de64b67bf..0b48ddd9c6 100644 --- a/services/cloud-agent-next/wrapper/src/control/apply-attach.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/apply-attach.test.ts @@ -248,11 +248,13 @@ describe('applySessionAttach', () => { it('fails preparation rather than falling back to main for a missing requested branch', async () => { let setupRan = false; + const diagnostics: Array> = []; const result = await applySessionAttach( { ...session, directory }, { ...payload, branch: 'missing', setupCommands: ['prepare'] }, { ...deps, + onDiagnostic: (_event, fields) => diagnostics.push(fields), runSetup: async () => { setupRan = true; return { stdout: '', stderr: '', exitCode: 0 }; @@ -264,6 +266,15 @@ describe('applySessionAttach', () => { error: { code: 'not_ready', message: 'git checkout failed', retryable: true }, }); expect(setupRan).toBe(false); + expect(diagnostics).toContainEqual( + expect.objectContaining({ + phase: 'failed', + stage: 'git_setup', + errorCode: 'not_ready', + retryable: true, + detail: 'git checkout failed', + }) + ); }); it('preserves generated-branch commits when retrying failed setup', async () => { diff --git a/services/cloud-agent-next/wrapper/src/control/apply-attach.ts b/services/cloud-agent-next/wrapper/src/control/apply-attach.ts index 33a891e366..78374d6869 100644 --- a/services/cloud-agent-next/wrapper/src/control/apply-attach.ts +++ b/services/cloud-agent-next/wrapper/src/control/apply-attach.ts @@ -2,13 +2,16 @@ import fs from 'node:fs/promises'; import path from 'node:path'; import { SANDBOX_CONTROL_ATTACH_TIMEOUT_MS, + controlErrorCodes, sessionAttachPayloadSchema, + type ControlErrorCode, type SessionAttachPayload, type SessionRequestIdentity, } from '../../../src/shared/sandbox-control-protocol.js'; import type { PreparingEventDataV2, PreparingStep } from '../../../src/shared/protocol.js'; import { CONTROL_RUNTIME_RESERVED_ENV_VARS } from '../../../src/shared/runtime-environment.js'; import { + diagnosticDetail, emitControlDiagnostic, type ControlDiagnosticRecord, type ControlDiagnosticReporter, @@ -91,10 +94,34 @@ function ok(): ControlHandlerResult { return { ok: true, result: { attached: true } }; } -function fail(code: string, message: string, retryable: boolean): ControlHandlerResult { +function fail( + code: ControlErrorCode, + message: string, + retryable: boolean +): Extract { return { ok: false, error: { code, message, retryable } }; } +function diagnosticErrorCode( + code: string +): NonNullable { + for (const value of controlErrorCodes) { + if (value === code) return value; + } + return 'other'; +} + +function attachFailureFields( + result: Extract +): Pick { + const detail = diagnosticDetail(result.error.message); + return { + errorCode: diagnosticErrorCode(result.error.code), + retryable: result.error.retryable, + ...(detail ? { detail } : {}), + }; +} + async function defaultHasGit(directory: string): Promise { try { await fs.access(path.join(directory, '.git', 'HEAD')); @@ -292,7 +319,11 @@ async function executeSessionAttach( let stage: ControlDiagnosticRecord['fields']['stage'] = 'attach_validation'; let workspaceAction: ControlDiagnosticRecord['fields']['workspaceAction']; let sessionResolution: ControlDiagnosticRecord['fields']['sessionResolution']; - const diagnostic = (phase: 'completed' | 'failed'): void => + let attachment: WorktreeKiloAttachment | undefined; + const diagnostic = ( + phase: 'completed' | 'failed', + extra: Partial = {} + ): void => emitControlDiagnostic(deps.onDiagnostic, 'control.request', { operation: 'session.attach', phase, @@ -304,19 +335,21 @@ async function executeSessionAttach( sessionResolution, elapsedMs: Math.max(0, Date.now() - startedAt), ok: phase === 'completed', + aborted: Boolean(deps.signal?.aborted || attachment?.signal.aborted), + ...extra, }); const existingDirectory = directoryForSession(session.kiloSessionId); if (existingDirectory && existingDirectory !== directory) { - diagnostic('failed'); - return fail('unauthorized', 'Session directory mismatch', false); + const result = fail('unauthorized', 'Session directory mismatch', false); + diagnostic('failed', attachFailureFields(result)); + return result; } stage = 'runtime_attach'; if (!deps.kiloRuntimes) { - diagnostic('failed'); - return fail('not_ready', 'Kilo is not ready', true); + const result = fail('not_ready', 'Kilo is not ready', true); + diagnostic('failed', attachFailureFields(result)); + return result; } - - let attachment: WorktreeKiloAttachment | undefined; const taskSignal = deps.signal ?? AbortSignal.timeout(SANDBOX_CONTROL_ATTACH_TIMEOUT_MS); try { taskSignal.throwIfAborted(); @@ -494,8 +527,8 @@ async function executeSessionAttach( ); } }); - if (workspaceFailure) { - diagnostic('failed'); + if (workspaceFailure && !workspaceFailure.ok) { + diagnostic('failed', attachFailureFields(workspaceFailure)); return workspaceFailure; } @@ -525,8 +558,9 @@ async function executeSessionAttach( if (!restored.ok) { if (restored.code !== 404 && !restored.emptySnapshot) { progress.fail('kilo_session', 'phase:kilo_session', restored.error); - diagnostic('failed'); - return fail('not_ready', 'kilo session is not ready', true); + const result = fail('not_ready', 'kilo session is not ready', true); + diagnostic('failed', attachFailureFields(result)); + return result; } stage = 'session_create'; progress.progress('kilo_session', 'phase:kilo_session', 'Starting session…'); @@ -544,8 +578,9 @@ async function executeSessionAttach( } catch { const message = signal.aborted ? 'Session attachment cancelled' : 'kilo session is not ready'; progress.fail('kilo_session', 'phase:kilo_session', message); - diagnostic('failed'); - return fail('not_ready', message, true); + const result = fail('not_ready', message, true); + diagnostic('failed', attachFailureFields(result)); + return result; } stage = 'attachment_commit'; signal.throwIfAborted(); @@ -563,15 +598,16 @@ async function executeSessionAttach( return ok(); } catch (error) { deps.onError?.(error); - diagnostic('failed'); - if (error instanceof WorktreeKiloRuntimeError || error instanceof ControlTerminalRuntimeError) { - return fail(error.code, error.message, error.retryable); - } - return fail( - 'not_ready', - taskSignal.aborted ? 'Session attachment cancelled' : 'Session attachment failed', - true - ); + const result = + error instanceof WorktreeKiloRuntimeError || error instanceof ControlTerminalRuntimeError + ? fail(error.code, error.message, error.retryable) + : fail( + 'not_ready', + taskSignal.aborted ? 'Session attachment cancelled' : 'Session attachment failed', + true + ); + diagnostic('failed', attachFailureFields(result)); + return result; } finally { attachment?.release(); } diff --git a/services/cloud-agent-next/wrapper/src/control/control-event-ordering.test.ts b/services/cloud-agent-next/wrapper/src/control/control-event-ordering.test.ts index 53655d847b..80eb327558 100644 --- a/services/cloud-agent-next/wrapper/src/control/control-event-ordering.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/control-event-ordering.test.ts @@ -115,7 +115,6 @@ describe('control event publication ordering', () => { expect(older).toEqual(original); expect(published[8]).toMatchObject({ receiptId: original.receiptId, - receiptHash: original.receiptHash, sequence: original.sequence, session: original.session, payload: original.payload, diff --git a/services/cloud-agent-next/wrapper/src/control/control-event-outbox.test.ts b/services/cloud-agent-next/wrapper/src/control/control-event-outbox.test.ts index 1847482a84..0d1375e798 100644 --- a/services/cloud-agent-next/wrapper/src/control/control-event-outbox.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/control-event-outbox.test.ts @@ -1,6 +1,4 @@ import { describe, expect, it, mock, spyOn } from 'bun:test'; -import { createHash } from 'node:crypto'; -import { canonicalControlEventJson } from '../../../src/shared/control-event-canonical'; import { ControlDeliveryError } from './sandbox-control-client'; import { createControlEventOutbox, type ControlEventPublication } from './control-event-outbox'; import { @@ -15,7 +13,7 @@ const session = { }; describe('control event outbox', () => { - it('snapshots native lifetime before replacement and binds it into the receipt hash', async () => { + it('snapshots native lifetime before replacement', async () => { const published: ControlEventPublication[] = []; const outbox = createControlEventOutbox({ publish: async publication => { @@ -42,17 +40,8 @@ describe('control event outbox', () => { expect(deadlineAt).toBeGreaterThan(Date.now()); const parsed = sandboxEventPublicationPayloadSchema.parse(wire); expect(wire).toEqual(parsed); - const hash = (value: unknown) => - createHash('sha256').update(canonicalControlEventJson(value)).digest('hex'); - const content = { - event: parsed.event, - session: parsed.session, - payload: parsed.payload, - sequence: parsed.sequence, - }; - expect(hash(content)).toBe(first.receiptHash); - expect(hash({ ...content, session: identity })).not.toBe(first.receiptHash); - expect(hash({ ...content, session })).not.toBe(first.receiptHash); + expect(wire).toEqual(expect.objectContaining({ receiptId: first.receiptId, sequence: 1 })); + expect(wire).not.toHaveProperty('receiptHash'); expect( sandboxEventPublicationPayloadSchema.safeParse({ ...wire, @@ -80,18 +69,7 @@ describe('control event outbox', () => { expect(deadlineAt).toBeGreaterThan(Date.now()); expect(wire).toEqual(sandboxEventPublicationPayloadSchema.parse(wire)); expect(wire.session).toEqual(session); - expect(wire.receiptHash).toBe( - createHash('sha256') - .update( - canonicalControlEventJson({ - event: wire.event, - session, - payload: wire.payload, - sequence: wire.sequence, - }) - ) - .digest('hex') - ); + expect(wire).not.toHaveProperty('receiptHash'); }); it('autonomously retries one stable receipt without future events or resume calls', async () => { diff --git a/services/cloud-agent-next/wrapper/src/control/control-event-outbox.ts b/services/cloud-agent-next/wrapper/src/control/control-event-outbox.ts index f047b811e1..cb8b6a35ce 100644 --- a/services/cloud-agent-next/wrapper/src/control/control-event-outbox.ts +++ b/services/cloud-agent-next/wrapper/src/control/control-event-outbox.ts @@ -1,4 +1,3 @@ -import { createHash } from 'node:crypto'; import { canonicalControlEventJson } from '../../../src/shared/control-event-canonical.js'; import { MAX_SANDBOX_CONTROL_FRAME_BYTES, @@ -13,7 +12,6 @@ const RETRY_DELAY_MS = 250; export type ControlEventPublication = { event: 'session.event' | 'session.preparing'; receiptId: string; - receiptHash: string; sequence: number; session: SessionEventIdentity; payload: unknown; @@ -31,7 +29,7 @@ export type ControlEventOutboxFailure = { export type ControlEventOutbox = { prepare( - input: Omit + input: Omit ): PreparedControlEventPublication; enqueue(publication: PreparedControlEventPublication): boolean; waitForSpace(publication: PreparedControlEventPublication): Promise; @@ -104,21 +102,18 @@ export function createControlEventOutbox(options: { }; const prepare = ( - input: Omit + input: Omit ): PreparedControlEventPublication => { const snapshot = JSON.parse( canonicalControlEventJson({ ...input, session: sessionEventIdentitySchema.parse(input.session), }) - ) as Omit; + ) as Omit; const sequence = nextSequence + 1; const receiptId = crypto.randomUUID(); - const receiptHash = createHash('sha256') - .update(canonicalControlEventJson({ ...snapshot, sequence })) - .digest('hex'); nextSequence = sequence; - const publication = { ...snapshot, sequence, receiptId, receiptHash }; + const publication = { ...snapshot, sequence, receiptId }; const bytes = Buffer.byteLength( JSON.stringify({ type: 'request', @@ -184,7 +179,6 @@ export function createControlEventOutbox(options: { { event: entry.event, receiptId: entry.receiptId, - receiptHash: entry.receiptHash, sequence: entry.sequence, session: entry.session, payload: entry.payload, diff --git a/services/cloud-agent-next/wrapper/src/control/control-event-transport.test.ts b/services/cloud-agent-next/wrapper/src/control/control-event-transport.test.ts index 3447d9dfa5..7954d5d763 100644 --- a/services/cloud-agent-next/wrapper/src/control/control-event-transport.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/control-event-transport.test.ts @@ -58,7 +58,12 @@ describe('native-scoped control event failures', () => { expect(failures).toHaveLength(1); expect(failures[0]).toMatchObject({ reason, - publication: { sequence: 1, session: { ...session, nativeRuntimeId: originalNativeId } }, + publication: { + event: 'session.event', + receiptId: expect.any(String), + sequence: 1, + session: { ...session, nativeRuntimeId: originalNativeId }, + }, }); expect(retired).not.toHaveBeenCalled(); expect(published.at(-1)?.session.nativeRuntimeId).toBe(replacement.runtimeId); diff --git a/services/cloud-agent-next/wrapper/src/control/control-event-transport.ts b/services/cloud-agent-next/wrapper/src/control/control-event-transport.ts index db59855e1c..007e4a97cb 100644 --- a/services/cloud-agent-next/wrapper/src/control/control-event-transport.ts +++ b/services/cloud-agent-next/wrapper/src/control/control-event-transport.ts @@ -30,7 +30,7 @@ export function createControlEventTransport(options: { event: EventKind; session: SessionEventIdentity; payload: unknown; - }) => Omit; + }) => Omit; sendLegacy: (payload: unknown, session: SessionEventIdentity) => boolean; onFailure: (failure: ControlEventOutboxFailure) => void; }) { diff --git a/services/cloud-agent-next/wrapper/src/control/file-log-uploader.test.ts b/services/cloud-agent-next/wrapper/src/control/file-log-uploader.test.ts new file mode 100644 index 0000000000..1aea82de9d --- /dev/null +++ b/services/cloud-agent-next/wrapper/src/control/file-log-uploader.test.ts @@ -0,0 +1,116 @@ +import { afterEach, describe, expect, it } from 'bun:test'; +import fsp from 'node:fs/promises'; +import os from 'node:os'; +import path from 'node:path'; +import { gunzipSync } from 'node:zlib'; +import { + createControlFileLogUploader, + selectControlFileLogPaths, + type ControlFileLogUploader, +} from './file-log-uploader'; + +const originalFetch = globalThis.fetch; +const temporaryDirectories: string[] = []; +const uploaders: ControlFileLogUploader[] = []; + +async function createFixture() { + const directory = await fsp.mkdtemp(path.join(os.tmpdir(), 'control-file-logs-')); + temporaryDirectories.push(directory); + const homeRoot = path.join(directory, 'kilo-worktrees'); + const kiloLogDir = path.join(homeRoot, 'home-a', '.local', 'share', 'kilo', 'log'); + const wrapperLogPath = path.join(directory, 'kilocode-control-wrapper.log'); + await fsp.mkdir(kiloLogDir, { recursive: true }); + await fsp.writeFile(wrapperLogPath, 'wrapper log\n'); + await fsp.writeFile(path.join(kiloLogDir, 'kilo.log'), 'kilo log\n'); + return { directory, homeRoot, kiloLogDir, wrapperLogPath }; +} + +function createUploader( + files: { homeRoot: string; wrapperLogPath: string }, + handler: (url: URL, init: RequestInit) => Promise +): ControlFileLogUploader { + const uploader = createControlFileLogUploader({ + uploadUrl: 'https://worker.example.com/sandbox-logs/sandbox/allocation/wrapper', + uploadGrant: 'test-upload-only-grant', + wrapperLogPath: files.wrapperLogPath, + homeRoot: files.homeRoot, + fetch: (url, init) => handler(new URL(url), init), + }); + uploaders.push(uploader); + return uploader; +} + +async function readArchive(init: RequestInit): Promise { + const bytes = await new Response(init.body).arrayBuffer(); + return gunzipSync(bytes).toString(); +} + +afterEach(async () => { + for (const uploader of uploaders.splice(0)) uploader.stop(); + globalThis.fetch = originalFetch; + await Promise.all( + temporaryDirectories + .splice(0) + .map(directory => fsp.rm(directory, { recursive: true, force: true })) + ); +}); + +describe('control file log uploader', () => { + it('packs the wrapper log and kilo log dir and puts gzip', async () => { + const files = await createFixture(); + let capturedUrl: URL | undefined; + let capturedInit: RequestInit | undefined; + let capturedArchive: string | undefined; + const uploader = createUploader(files, async (url, init) => { + capturedUrl = url; + capturedInit = init; + capturedArchive = await readArchive(init); + return new Response(null, { status: 204 }); + }); + + await uploader.uploadNow(); + + expect(capturedUrl?.pathname).toBe('/sandbox-logs/sandbox/allocation/wrapper/files.tar.gz'); + expect(new Headers(capturedInit?.headers).get('Content-Type')).toBe('application/gzip'); + expect(new Headers(capturedInit?.headers).get('Authorization')).toBe( + 'Bearer test-upload-only-grant' + ); + expect(capturedArchive).toContain('wrapper log'); + expect(capturedArchive).toContain('kilo log'); + }); + + it('keeps shutdown from waiting on a failed upload', async () => { + const files = await createFixture(); + const uploader = createUploader(files, async () => { + throw new Error('upload failed'); + }); + const settled = await Promise.race([ + uploader.finalize(50).then(() => true), + Bun.sleep(200).then(() => false), + ]); + expect(settled).toBe(true); + }); + + it('includes the wrapper log first, then newest kilo files until the cap', async () => { + const files = await createFixture(); + await fsp.rm(path.join(files.kiloLogDir, 'kilo.log')); + const older = path.join(files.kiloLogDir, 'older.log'); + const newest = path.join(files.kiloLogDir, 'newest.log'); + await fsp.writeFile(older, 'old'); + await fsp.writeFile(newest, 'new-content'); + const olderTime = new Date('2026-01-01T00:00:00Z'); + const newerTime = new Date('2026-01-02T00:00:00Z'); + await fsp.utimes(older, olderTime, olderTime); + await fsp.utimes(newest, newerTime, newerTime); + const wrapperSize = (await fsp.stat(files.wrapperLogPath)).size; + const newestSize = (await fsp.stat(newest)).size; + const selected = selectControlFileLogPaths({ + wrapperLogPath: files.wrapperLogPath, + kiloLogDirs: [files.kiloLogDir], + maxBytes: wrapperSize + newestSize, + }); + expect(selected[0]).toBe(files.wrapperLogPath); + expect(selected).toContain(newest); + expect(selected).not.toContain(older); + }); +}); diff --git a/services/cloud-agent-next/wrapper/src/control/file-log-uploader.ts b/services/cloud-agent-next/wrapper/src/control/file-log-uploader.ts new file mode 100644 index 0000000000..29e913d52b --- /dev/null +++ b/services/cloud-agent-next/wrapper/src/control/file-log-uploader.ts @@ -0,0 +1,282 @@ +import { existsSync, readdirSync, statSync } from 'node:fs'; +import os from 'node:os'; +import path from 'node:path'; +import { + CONTROL_LOG_ARCHIVE_NAME, + CONTROL_LOG_MAX_ARCHIVE_BYTES, + type ControlDiagnosticReporter, + type ControlLogUploadResult, +} from '../../../src/shared/control-diagnostics.js'; +import { createTarStream, type TarArchiveEntry } from '../log-uploader.js'; +import { logToFile, withTimeoutAndAbort } from '../utils.js'; + +export type ControlFileLogUploader = { + start: (intervalMs?: number) => void; + uploadNow: () => Promise; + finalize: (timeoutMs?: number) => Promise; + stop: () => void; +}; + +type Options = { + uploadUrl?: string; + uploadGrant?: string; + wrapperLogPath?: string; + homeRoot?: string; + fetch?: (url: string, init: RequestInit) => Promise; + onDiagnostic?: ControlDiagnosticReporter; + intervalMs?: number; + uploadTimeoutMs?: number; +}; + +const KILO_LOG_RELATIVE = path.join('.local', 'share', 'kilo', 'log'); +const DEFAULT_WRAPPER_LOG_PATH = '/tmp/kilocode-control-wrapper.log'; +const UPLOAD_TIMEOUT_MS = 15_000; +const FINAL_UPLOAD_TIMEOUT_MS = 5_000; +const INTERVAL_MS = 30_000; + +export function defaultControlWorktreeHomeRoot(): string { + return path.join(os.tmpdir(), 'kilo-worktrees'); +} + +export function listControlKiloLogDirs(homeRoot = defaultControlWorktreeHomeRoot()): string[] { + if (!existsSync(homeRoot)) return []; + const dirs: string[] = []; + for (const entry of readdirSync(homeRoot, { withFileTypes: true })) { + if (!entry.isDirectory() || entry.isSymbolicLink()) continue; + const logDir = path.join(homeRoot, entry.name, KILO_LOG_RELATIVE); + if (existsSync(logDir)) dirs.push(logDir); + } + return dirs; +} + +export function selectControlFileLogPaths(input: { + wrapperLogPath: string; + kiloLogDirs: string[]; + maxBytes?: number; +}): string[] { + const maxBytes = input.maxBytes ?? CONTROL_LOG_MAX_ARCHIVE_BYTES; + const selected: string[] = []; + let used = 0; + const add = (filePath: string, size: number): void => { + if (size > maxBytes - used) return; + selected.push(filePath); + used += size; + }; + if (existsSync(input.wrapperLogPath)) { + const st = statSync(input.wrapperLogPath); + if (st.isFile()) add(input.wrapperLogPath, st.size); + } + const kiloFiles: Array<{ path: string; size: number; mtime: number }> = []; + for (const dir of input.kiloLogDirs) { + if (!existsSync(dir)) continue; + for (const entry of readdirSync(dir, { withFileTypes: true })) { + if (!entry.isFile() || entry.isSymbolicLink()) continue; + const filePath = path.join(dir, entry.name); + const st = statSync(filePath); + if (!st.isFile()) continue; + kiloFiles.push({ path: filePath, size: st.size, mtime: st.mtimeMs }); + } + } + kiloFiles.sort((left, right) => right.mtime - left.mtime || left.path.localeCompare(right.path)); + for (const file of kiloFiles) add(file.path, file.size); + return selected; +} + +function archiveEntries( + files: string[], + wrapperLogPath: string, + homeRoot: string +): TarArchiveEntry[] { + const entries: TarArchiveEntry[] = []; + for (const filePath of files) { + if (filePath === wrapperLogPath) { + entries.push({ directory: path.dirname(filePath), name: path.basename(filePath) }); + continue; + } + const relative = path.relative(homeRoot, filePath); + if (relative.startsWith('..') || path.isAbsolute(relative)) continue; + entries.push({ directory: homeRoot, name: relative }); + } + return entries; +} + +async function collectArchive( + entries: TarArchiveEntry[], + maxBytes: number +): Promise { + const tar = createTarStream(entries); + if (!tar) return undefined; + try { + const reader = tar.stream.getReader(); + const chunks: Uint8Array[] = []; + let length = 0; + while (true) { + const { value, done } = await reader.read(); + if (done) break; + length += value.byteLength; + if (length > maxBytes) { + tar.kill(); + return undefined; + } + chunks.push(value); + } + const bytes = new Uint8Array(length); + let offset = 0; + for (const chunk of chunks) { + bytes.set(chunk, offset); + offset += chunk.byteLength; + } + return bytes; + } finally { + tar.kill(); + } +} + +export function createControlFileLogUploader(options: Options): ControlFileLogUploader { + type Upload = { promise: Promise; abort: AbortController }; + const upload = options.fetch ?? fetch; + const wrapperLogPath = options.wrapperLogPath ?? DEFAULT_WRAPPER_LOG_PATH; + const homeRoot = options.homeRoot ?? defaultControlWorktreeHomeRoot(); + let intervalId: ReturnType | undefined; + let activeUpload: Upload | undefined; + let queuedUpload: Upload | undefined; + let finalUpload: Promise | undefined; + let stopped = false; + + function reportFailure(category: ControlLogUploadResult, statusCode?: number): void { + try { + options.onDiagnostic?.('control.upload', { + phase: 'failed', + category, + statusCode: + statusCode !== undefined && statusCode >= 100 && statusCode <= 599 + ? statusCode + : undefined, + }); + } catch { + return; + } + } + + function uploadNow(): Promise { + if (finalUpload) return finalUpload; + if (stopped) return Promise.resolve(); + if (queuedUpload) return queuedUpload.promise; + const uploadUrl = options.uploadUrl; + const uploadGrant = options.uploadGrant; + if (!uploadUrl || !uploadGrant) return Promise.resolve(); + const archiveUrl = `${uploadUrl.replace(/\/$/, '')}/${CONTROL_LOG_ARCHIVE_NAME}`; + const authorization = `Bearer ${uploadGrant}`; + + const previousUpload = activeUpload; + const abort = new AbortController(); + const next: Upload = { promise: performUpload(), abort }; + if (previousUpload) queuedUpload = next; + else activeUpload = next; + return next.promise; + + async function performUpload(): Promise { + try { + await withTimeoutAndAbort( + (async () => { + await previousUpload?.promise; + abort.signal.throwIfAborted(); + if (queuedUpload === next) queuedUpload = undefined; + activeUpload = next; + + const files = selectControlFileLogPaths({ + wrapperLogPath, + kiloLogDirs: listControlKiloLogDirs(homeRoot), + }); + const entries = archiveEntries(files, wrapperLogPath, homeRoot); + if (entries.length === 0) return; + const body = await collectArchive(entries, CONTROL_LOG_MAX_ARCHIVE_BYTES); + abort.signal.throwIfAborted(); + if (!body) { + logToFile('Control file log archive exceeded the upload cap'); + return; + } + + const payload = new ArrayBuffer(body.byteLength); + new Uint8Array(payload).set(body); + const response = await upload(archiveUrl, { + method: 'PUT', + headers: { + Authorization: authorization, + 'Content-Type': 'application/gzip', + 'Content-Length': String(payload.byteLength), + }, + body: payload, + redirect: 'error', + signal: abort.signal, + }); + if (!abort.signal.aborted && response.status !== 204) { + logToFile(`Control file log upload failed: ${response.status}`); + reportFailure('http_rejection', response.status); + } + void response.body?.cancel().catch(() => undefined); + })(), + { + timeoutMs: options.uploadTimeoutMs ?? UPLOAD_TIMEOUT_MS, + timeoutMessage: 'Control file log upload timed out', + signal: abort.signal, + abortMessage: 'Control file log upload aborted', + } + ); + } catch (error) { + if (stopped || abort.signal.aborted) return; + const message = error instanceof Error ? error.message : ''; + reportFailure(message.includes('timed out') ? 'timeout' : 'network_failure'); + logToFile('Control file log upload did not complete'); + } finally { + abort.abort(); + if (activeUpload === next) activeUpload = undefined; + if (queuedUpload === next) queuedUpload = undefined; + } + } + } + + function clearUploadInterval(): void { + if (intervalId !== undefined) { + clearInterval(intervalId); + intervalId = undefined; + } + } + + function start(intervalMs = options.intervalMs ?? INTERVAL_MS): void { + stop(); + stopped = false; + finalUpload = undefined; + if (!options.uploadUrl || !options.uploadGrant) return; + void uploadNow(); + intervalId = setInterval(() => { + if (!activeUpload && !queuedUpload) void uploadNow(); + }, intervalMs); + intervalId.unref?.(); + } + + function finalize(timeoutMs = FINAL_UPLOAD_TIMEOUT_MS): Promise { + if (finalUpload) return finalUpload; + clearUploadInterval(); + finalUpload = withTimeoutAndAbort(uploadNow(), { + timeoutMs, + timeoutMessage: 'Final control file log upload timed out', + abortMessage: 'Final control file log upload aborted', + }) + .catch(() => { + reportFailure('timeout'); + logToFile('Final control file log upload timed out'); + }) + .finally(stop); + return finalUpload; + } + + function stop(): void { + stopped = true; + clearUploadInterval(); + activeUpload?.abort.abort(); + queuedUpload?.abort.abort(); + } + + return { start, uploadNow, finalize, stop }; +} diff --git a/services/cloud-agent-next/wrapper/src/control/main.ts b/services/cloud-agent-next/wrapper/src/control/main.ts index e6c4636893..7be75bdd85 100644 --- a/services/cloud-agent-next/wrapper/src/control/main.ts +++ b/services/cloud-agent-next/wrapper/src/control/main.ts @@ -1,4 +1,8 @@ -import type { SandboxHeartbeatPayload } from '../../../src/shared/sandbox-control-protocol.js'; +import { + heartbeatReasonFrom, + sessionAttachResultSchema, + type SandboxHeartbeatPayload, +} from '../../../src/shared/sandbox-control-protocol.js'; import { WRAPPER_VERSION } from '../../../src/shared/wrapper-version.js'; import { logToFile } from '../utils.js'; import { @@ -17,30 +21,21 @@ import { eventKiloSessionId, sessionEventIdentity, updateSessionSnapshots } from import { createControlTerminalRuntime } from './terminal-runtime'; import { createWorktreeKiloRuntimes } from './worktree-runtime'; import { createControlDiagnostics, type ControlDiagnostics } from './diagnostics'; -import { controlLogWrapperIdSchema } from '../../../src/shared/control-diagnostics.js'; +import { createControlFileLogUploader, type ControlFileLogUploader } from './file-log-uploader'; +import { + classifyRetirementCause, + controlLogWrapperIdSchema, + diagnosticDetail, +} from '../../../src/shared/control-diagnostics.js'; import { createWorktreeMutationNotifications } from './worktree-mutation-notifications'; import { createControlEventFailureHandler } from './control-event-transport'; +import type { ControlEventOutboxFailure } from './control-event-outbox'; -const retirementCauses = new Map([ - ['Kilo event feed is no longer healthy', 'event_feed_unhealthy'], - ['process_exited', 'process_exited'], - ['credential_refresh_failed', 'credential_refresh_failed'], - ['Sandbox control connection lost', 'control_disconnected'], - ['Preparation event delivery failed', 'preparation_delivery_failed'], - ['Sandbox shutting down', 'requested_shutdown'], - ['Wrapper received SIGTERM', 'sigterm'], - ['Wrapper received SIGINT', 'sigint'], - ['Wrapper uncaught exception', 'uncaught_exception'], - ['Wrapper unhandled rejection', 'unhandled_rejection'], - ['Kilo cancellation failed', 'cancellation_failed'], - ['Kilo cancellation was not confirmed', 'cancellation_failed'], - ['Native cancellation did not settle', 'cancellation_failed'], - ['Session outcome delivery failed', 'outcome_delivery_failed'], - ['Execution exceeded the 60 minute limit', 'execution_deadline'], - ['Session preparation timed out', 'preparation_deadline'], -]); - -function main(diagnostics: ControlDiagnostics, wrapperInstanceId: string): void { +function main( + diagnostics: ControlDiagnostics, + fileLogs: ControlFileLogUploader, + wrapperInstanceId: string +): void { const controlConfig = { SANDBOX_CONTROL_URL: process.env.SANDBOX_CONTROL_URL, SANDBOX_CONTROL_CREDENTIAL: process.env.SANDBOX_CONTROL_CREDENTIAL, @@ -98,7 +93,7 @@ function main(diagnostics: ControlDiagnostics, wrapperInstanceId: string): void return current === undefined || current.runtimeId === failure.runtimeId; }; if (failure.cleanup === 'unconfirmed' || !control?.reportNativeRuntimeRetirement) { - if (stillCurrent()) shutdown(1, failure.reason); + if (stillCurrent()) shutdown(1, failure.reason, heartbeatReasonFrom(failure.reason)); return; } void control @@ -111,10 +106,11 @@ function main(diagnostics: ControlDiagnostics, wrapperInstanceId: string): void }) .then( retired => { - if (!retired && stillCurrent()) shutdown(1, failure.reason); + if (!retired && stillCurrent()) + shutdown(1, failure.reason, heartbeatReasonFrom(failure.reason)); }, () => { - if (stillCurrent()) shutdown(1, failure.reason); + if (stillCurrent()) shutdown(1, failure.reason, heartbeatReasonFrom(failure.reason)); } ); }, @@ -154,7 +150,7 @@ function main(diagnostics: ControlDiagnostics, wrapperInstanceId: string): void }, options?.retained ? { preserveConnectionOnFailure: true } : undefined ) === true, - retireRuntime: reason => shutdown(1, reason), + retireRuntime: reason => shutdown(1, reason, heartbeatReasonFrom(reason)), onShutdown: () => shutdown(0, 'Sandbox shutting down'), }); @@ -170,6 +166,105 @@ function main(diagnostics: ControlDiagnostics, wrapperInstanceId: string): void return payload; } + function reportOutboxRetirement( + failure: ControlEventOutboxFailure, + nativeRuntimeId: string, + phase: 'started' | 'retired' | 'failed', + ok?: boolean + ): void { + const fields = { + phase, + category: + failure.publication.event === 'session.preparing' + ? ('preparing' as const) + : ('session_event' as const), + sessionId: failure.publication.session.kiloSessionId, + receiptId: failure.publication.receiptId, + sequence: failure.publication.sequence, + wrapperInstanceId, + nativeRuntimeId, + failureReason: failure.reason, + outboxExpiryRetirement: failure.reason === 'expired', + ...(ok === undefined ? {} : { ok }), + }; + diagnostics.onDiagnostic('control.event', fields); + logToFile( + `control diagnostic ${JSON.stringify({ + event: 'outbox_retirement', + publicationEvent: failure.publication.event, + ...fields, + })}` + ); + } + + type SessionAttachResult = + | { kind: 'response'; response: Awaited> } + | { kind: 'failed' }; + + function reportSessionAttachResult( + session: Parameters[1], + authorization: Parameters[4], + outcome: SessionAttachResult + ): void { + if (outcome.kind === 'failed') { + diagnostics.onDiagnostic('control.request', { + phase: 'response_failed', + operation: 'session.attach', + sessionId: session?.sessionId, + requestId: authorization?.operationId, + scopeId: session?.kiloSessionId, + wrapperInstanceId, + ok: false, + errorCode: 'other', + retryable: false, + }); + logToFile( + `control diagnostic ${JSON.stringify({ + event: 'session_attach_result', + operationId: authorization?.operationId, + attemptId: authorization?.operationId, + sessionId: session?.sessionId, + kiloSessionId: session?.kiloSessionId, + wrapperInstanceId, + ok: false, + result: 'failed', + errorCode: 'other', + retryable: false, + })}` + ); + return; + } + const { response } = outcome; + const attached = response.ok ? sessionAttachResultSchema.safeParse(response.result) : undefined; + diagnostics.onDiagnostic('control.request', { + phase: response.ok ? 'response_sent' : 'response_failed', + operation: 'session.attach', + sessionId: session?.sessionId, + requestId: authorization?.operationId, + scopeId: session?.kiloSessionId, + wrapperInstanceId, + nativeRuntimeId: attached?.success ? attached.data.nativeRuntimeId : undefined, + ok: response.ok, + errorCode: response.ok ? undefined : response.error.code, + retryable: response.ok ? undefined : response.error.retryable, + }); + logToFile( + `control diagnostic ${JSON.stringify({ + event: 'session_attach_result', + operationId: authorization?.operationId, + attemptId: authorization?.operationId, + sessionId: session?.sessionId, + kiloSessionId: session?.kiloSessionId, + wrapperInstanceId, + nativeRuntimeId: attached?.success ? attached.data.nativeRuntimeId : undefined, + ok: response.ok, + result: response.ok ? 'accepted' : 'rejected', + errorCode: response.ok ? undefined : response.error.code, + retryable: response.ok ? undefined : response.error.retryable, + })}` + ); + } + function shutdown( exitCode: number, reason: string, @@ -191,16 +286,15 @@ function main(diagnostics: ControlDiagnostics, wrapperInstanceId: string): void } }; const deadline = setTimeout(finish, KILO_CONTROL_REQUEST_TIMEOUT_MS); + const detail = diagnosticDetail(reason); diagnostics.onDiagnostic('wrapper.lifecycle', { phase: 'stopping', exitCode, - retirementCause: - retirementCauses.get(reason) ?? - retirementCauses.get(diagnosticReason) ?? - (diagnosticReason.startsWith('feed_') ? 'event_feed_unhealthy' : 'unknown'), + retirementCause: classifyRetirementCause(reason, diagnosticReason), + ...(detail ? { detail } : {}), }); void diagnostics.flush(); - logToFile(`control-plane wrapper retiring exitCode=${exitCode}`); + logToFile(`control-plane wrapper retiring exitCode=${exitCode} reason=${reason}`); const stopped = (async () => { try { control?.sendEvent?.('sandbox.heartbeat', withHeartbeatReason(buildHeartbeatPayload(deps))); @@ -222,6 +316,8 @@ function main(diagnostics: ControlDiagnostics, wrapperInstanceId: string): void .then(async () => { const remaining = KILO_CONTROL_REQUEST_TIMEOUT_MS - (Date.now() - shutdownAt) - 100; await diagnostics.finalize(Math.max(1, Math.min(4000, remaining))); + const fileRemaining = KILO_CONTROL_REQUEST_TIMEOUT_MS - (Date.now() - shutdownAt) - 100; + await fileLogs.finalize(Math.max(1, Math.min(5000, fileRemaining))); }) .finally(() => { clearTimeout(deadline); @@ -241,6 +337,7 @@ function main(diagnostics: ControlDiagnostics, wrapperInstanceId: string): void onEventReceiptFailure: createControlEventFailureHandler({ getRuntime: directory => kiloRuntimes.get(directory), onFailure: (failure, runtime) => { + reportOutboxRetirement(failure, runtime.runtimeId, 'started'); void deps.operations .retireDirectory( failure.publication.session.directory, @@ -248,7 +345,11 @@ function main(diagnostics: ControlDiagnostics, wrapperInstanceId: string): void Date.now() + KILO_CONTROL_REQUEST_TIMEOUT_MS, { runtimeId: runtime.runtimeId, client: runtime.kiloClient } ) + .then(() => { + reportOutboxRetirement(failure, runtime.runtimeId, 'retired', true); + }) .catch(() => { + reportOutboxRetirement(failure, runtime.runtimeId, 'failed', false); diagnostics.onDiagnostic('wrapper.lifecycle', { phase: 'failed' }); }); }, @@ -258,33 +359,46 @@ function main(diagnostics: ControlDiagnostics, wrapperInstanceId: string): void if (Date.now() >= deadlineAt) throw new Error('Control recovery deadline expired'); await deps.operations.drainDelivery(deadlineAt); }, - onRequest: (operation, session, payload, authorization) => { - return handleControlRequest( - operation, - session, - payload, - { - ...deps, - emitPreparing: (event, options) => { - if (!session) return; - if ( - !control?.sendEvent?.( - 'session.preparing', - event, - { - directory: session.directory, - kiloSessionId: session.kiloSessionId, - rootKiloSessionId: session.kiloSessionId, - ...(options?.nativeRuntimeId ? { nativeRuntimeId: options.nativeRuntimeId } : {}), - }, - options?.retained ? { preserveConnectionOnFailure: true } : undefined + onRequest: async (operation, session, payload, authorization) => { + try { + const response = await handleControlRequest( + operation, + session, + payload, + { + ...deps, + emitPreparing: (event, options) => { + if (!session) return; + if ( + !control?.sendEvent?.( + 'session.preparing', + event, + { + directory: session.directory, + kiloSessionId: session.kiloSessionId, + rootKiloSessionId: session.kiloSessionId, + ...(options?.nativeRuntimeId + ? { nativeRuntimeId: options.nativeRuntimeId } + : {}), + }, + options?.retained ? { preserveConnectionOnFailure: true } : undefined + ) ) - ) - throw new Error('Preparation event delivery failed'); + throw new Error('Preparation event delivery failed'); + }, }, - }, - authorization - ); + authorization + ); + if (operation === 'session.attach') { + reportSessionAttachResult(session, authorization, { kind: 'response', response }); + } + return response; + } catch (error) { + if (operation === 'session.attach') { + reportSessionAttachResult(session, authorization, { kind: 'failed' }); + } + throw error; + } }, getHeartbeatPayload: () => withHeartbeatReason(buildHeartbeatPayload(deps)), sampleHeartbeat: signal => refreshHeartbeatPayload(deps, signal).then(() => undefined), @@ -300,20 +414,32 @@ const configuredWrapperId = controlLogWrapperIdSchema.safeParse( const wrapperInstanceId = configuredWrapperId.success ? configuredWrapperId.data : crypto.randomUUID(); +const uploadUrl = process.env.CONTROL_LOG_UPLOAD_URL; +const uploadGrant = process.env.CONTROL_LOG_UPLOAD_GRANT; const diagnostics = createControlDiagnostics({ - uploadUrl: process.env.CONTROL_LOG_UPLOAD_URL, - uploadGrant: process.env.CONTROL_LOG_UPLOAD_GRANT, + uploadUrl, + uploadGrant, +}); +const fileLogs = createControlFileLogUploader({ + uploadUrl, + uploadGrant, + wrapperLogPath: process.env.WRAPPER_LOG_PATH, + onDiagnostic: diagnostics.onDiagnostic, }); delete process.env.CONTROL_LOG_UPLOAD_URL; delete process.env.CONTROL_LOG_UPLOAD_GRANT; delete process.env.CONTROL_WRAPPER_INSTANCE_ID; diagnostics.onDiagnostic('wrapper.lifecycle', { phase: 'starting' }); diagnostics.start(); +fileLogs.start(); try { - main(diagnostics, wrapperInstanceId); + main(diagnostics, fileLogs, wrapperInstanceId); } catch { diagnostics.onDiagnostic('wrapper.lifecycle', { phase: 'start_failed' }); logToFile('control-plane wrapper failed'); - void diagnostics.finalize().finally(() => process.exit(1)); + void diagnostics + .finalize() + .then(() => fileLogs.finalize()) + .finally(() => process.exit(1)); } diff --git a/services/cloud-agent-next/wrapper/src/control/operation-registry.test.ts b/services/cloud-agent-next/wrapper/src/control/operation-registry.test.ts index 179b016345..549ae8a924 100644 --- a/services/cloud-agent-next/wrapper/src/control/operation-registry.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/operation-registry.test.ts @@ -2,6 +2,10 @@ import { afterEach, beforeEach, describe, expect, it, setSystemTime, spyOn } fro import fs from 'node:fs'; import os from 'node:os'; import path from 'node:path'; +import { + OWNED_PROCESS_CLEANUP_UNREAPED, + type ControlDiagnosticFields, +} from '../../../src/shared/control-diagnostics'; import { SANDBOX_CONTROL_OPERATION_LIMIT, SANDBOX_CONTROL_OUTCOME_TIMEOUT_MS, @@ -228,3 +232,116 @@ describe('operation admission and lookup', () => { expect(handlerDeps.operations.abortTarget(session, 'msg_1')).toBe(prompt); }); }); + +describe('completed receipt prune', () => { + it('does not retire the wrapper when a completed receipt still has leftover occupancy', async () => { + const retired: string[] = []; + const diagnostics: Array<{ event: string; fields: ControlDiagnosticFields }> = []; + const handlerDeps = deps({ + sendOperationResult: (_session, delivery) => acknowledgeOperation(delivery), + retireRuntime: reason => { + retired.push(reason); + }, + onDiagnostic: (event, fields) => diagnostics.push({ event, fields }), + }); + const nativeRuntime = handlerDeps.kiloRuntimes; + if (!nativeRuntime) throw new Error('Missing native runtime'); + const nativeRetire = spyOn(nativeRuntime, 'retireRuntime').mockResolvedValue('unconfirmed'); + rememberAttachedRoot(session.kiloSessionId, session.directory); + const authorization = operationAuthorization(); + await handleControlRequest( + 'session.prompt', + session, + promptPayload, + handlerDeps, + authorization + ); + const record = onlyOperation(handlerDeps); + await record.done; + await record.waitForDelivery(); + + const release = spyOn(record, 'releaseProcessOwnership').mockReturnValue(false); + const cleaned = Promise.withResolvers(); + const cleanup = spyOn(record, 'cleanupOwnedWork').mockImplementation(async () => { + cleaned.resolve(); + return false; + }); + const requestRetirement = spyOn(record, 'requestRetirement'); + const logged = spyOn(console, 'error').mockImplementation(() => {}); + const runtime = nativeRuntime.get(session.directory); + if (!runtime) throw new Error('Missing native runtime'); + const attachStarted = Promise.withResolvers(); + const attachRelease = Promise.withResolvers(); + let attachAborted = false; + const attaching = handleControlRequest( + 'session.attach', + session, + { kilo }, + { + ...handlerDeps, + applyAttach: async (_identity, _payload, hooks) => { + if (!hooks.onRuntime) throw new Error('Missing attach runtime hook'); + hooks.onRuntime(runtime); + const signal = hooks.signal; + if (signal) { + signal.addEventListener( + 'abort', + () => { + attachAborted = true; + }, + { once: true } + ); + } + attachStarted.resolve(); + await attachRelease.promise; + if (signal?.aborted) { + return { + ok: false, + error: { + code: 'not_ready', + message: 'Session attachment cancelled', + retryable: true, + }, + }; + } + return { ok: true, result: { attached: true } }; + }, + } + ); + try { + await attachStarted.promise; + setSystemTime(authorization.dispatchDeadlineAt + SANDBOX_CONTROL_OUTCOME_TIMEOUT_MS + 1); + pruneControlOperations(handlerDeps); + await cleaned.promise; + await Promise.resolve(); + await Promise.resolve(); + + expect(requestRetirement).not.toHaveBeenCalled(); + expect(retired).toEqual([]); + expect(attachAborted).toBe(false); + expect(handlerDeps.operations.retained()).not.toContain(record); + expect( + diagnostics.some( + diagnostic => + diagnostic.event === 'session.task' && + diagnostic.fields.stage === 'process_cleanup' && + diagnostic.fields.phase === 'failed' && + diagnostic.fields.ok === false && + diagnostic.fields.messageId === authorization.messageId && + String(diagnostic.fields.detail ?? '').startsWith('owned_process_unreaped ') + ) + ).toBe(true); + expect(logged.mock.calls.some(args => args[0] === OWNED_PROCESS_CLEANUP_UNREAPED)).toBe(true); + + attachRelease.resolve(); + expect(await attaching).toMatchObject({ ok: true }); + } finally { + attachRelease.resolve(); + release.mockRestore(); + cleanup.mockRestore(); + requestRetirement.mockRestore(); + nativeRetire.mockRestore(); + logged.mockRestore(); + } + }); +}); diff --git a/services/cloud-agent-next/wrapper/src/control/operation-registry.ts b/services/cloud-agent-next/wrapper/src/control/operation-registry.ts index 2b3273533e..f0df4fe7ad 100644 --- a/services/cloud-agent-next/wrapper/src/control/operation-registry.ts +++ b/services/cloud-agent-next/wrapper/src/control/operation-registry.ts @@ -77,9 +77,15 @@ export function createOperationRegistry(deps: OperationRegistryDependencies) { if (operation.releaseProcessOwnership()) retained.delete(id); else { const deadlineAt = operation.captureCleanupDeadline(); - void operation.cleanupOwnedWork(deadlineAt).then(confirmed => { - if (!confirmed) operation.requestRetirement('Owned process cleanup failed', deadlineAt); - }); + void operation + .cleanupOwnedWork(deadlineAt) + .catch(() => false) + .then(confirmed => { + if (retained.get(id) !== operation) return; + const released = operation.releaseProcessOwnership(); + if (!confirmed || !released) operation.reportUnreapedProcessCleanup(!released); + retained.delete(id); + }); } } } diff --git a/services/cloud-agent-next/wrapper/src/control/owned-processes.test.ts b/services/cloud-agent-next/wrapper/src/control/owned-processes.test.ts index f2cb8544e9..55c4eb77c7 100644 --- a/services/cloud-agent-next/wrapper/src/control/owned-processes.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/owned-processes.test.ts @@ -1,5 +1,6 @@ import { afterEach, describe, expect, it } from 'bun:test'; import { once } from 'node:events'; +import { runProcess } from '../utils.js'; import { createOwnedProcessScope } from './owned-processes.js'; const spawned: ReturnType[] = []; @@ -16,6 +17,11 @@ afterEach(async () => { await Promise.all(spawned.splice(0).map(scope => scope.stop(Date.now() + 1_000))); }); +const descendantThatExitsMs = (ms: number) => + `const { spawn } = require('node:child_process'); const child = spawn(process.execPath, ['-e', 'setTimeout(() => {}, ${ms})'], { stdio: 'ignore' }); process.stdout.write(String(child.pid)); setTimeout(() => process.exit(0), 20);`; + +const immortalDescendant = `const { spawn } = require('node:child_process'); const child = spawn(process.execPath, ['-e', 'setInterval(() => {}, 1000)'], { stdio: 'ignore' }); process.stdout.write(String(child.pid)); setTimeout(() => process.exit(0), 20);`; + describe('owned process scopes', () => { it('coalesces cleanup and treats unavailable containment as unconfirmed', async () => { const scope = createOwnedProcessScope(); @@ -61,4 +67,91 @@ describe('owned process scopes', () => { } expect(() => process.kill(descendant, 0)).toThrow(); }); + + it('keeps occupancy after a successful parent exit until descendants are gone on Linux', async () => { + if (process.platform !== 'linux') return; + const scope = createOwnedProcessScope(); + spawned.push(scope); + let descendant = 0; + try { + const result = await scope.run(() => + runProcess(process.execPath, ['-e', immortalDescendant], { timeoutMs: 400 }) + ); + descendant = Number(result.stdout); + expect(result.exitCode).toBe(0); + expect(Number.isSafeInteger(descendant) && descendant > 0).toBe(true); + if (!scope.observesOccupancy()) return; + expect(await scope.verify(false)).toBe(false); + expect(scope.dispose()).toBe(false); + killPid(descendant); + const deadlineAt = Date.now() + 1_000; + while (Date.now() < deadlineAt && !(await scope.verify(false))) { + await Bun.sleep(25); + } + expect(await scope.verify(false)).toBe(true); + } finally { + if (descendant > 0) killPid(descendant); + } + }); + + it('waits for a short-lived descendant before treating runProcess as complete on Linux', async () => { + if (process.platform !== 'linux') return; + const scope = createOwnedProcessScope(); + spawned.push(scope); + let descendant = 0; + try { + const startedAt = Date.now(); + const result = await scope.run(() => + runProcess(process.execPath, ['-e', descendantThatExitsMs(250)], { timeoutMs: 2_000 }) + ); + descendant = Number(result.stdout); + expect(result.exitCode).toBe(0); + if (!scope.observesOccupancy()) return; + expect(Date.now() - startedAt).toBeGreaterThanOrEqual(200); + expect(await scope.verify(false)).toBe(true); + } finally { + if (descendant > 0) killPid(descendant); + } + }); + + it('does not wait for Darwin occupancy after the tracked parent exits', async () => { + if (process.platform === 'linux') return; + const scope = createOwnedProcessScope(); + spawned.push(scope); + let descendant = 0; + try { + const startedAt = Date.now(); + const result = await scope.run(() => + runProcess(process.execPath, ['-e', immortalDescendant], { timeoutMs: 2_000 }) + ); + descendant = Number(result.stdout); + expect(result.exitCode).toBe(0); + expect(Number.isSafeInteger(descendant) && descendant > 0).toBe(true); + expect(Date.now() - startedAt).toBeLessThan(1_000); + expect(await scope.verify(false)).toBe(false); + } finally { + if (descendant > 0) killPid(descendant); + } + }); + + it('does not wait out timeoutMs after close when occupancy is unobservable', async () => { + const scope = createOwnedProcessScope(); + spawned.push(scope); + const result = await scope.run(async () => { + const ungated = scope.spawn('/bin/true', [], { + cwd: process.cwd(), + env: process.env, + shell: true, + }); + await once(ungated, 'close'); + expect(scope.observesOccupancy()).toBe(false); + const startedAt = Date.now(); + const completed = await runProcess(process.execPath, ['-e', 'process.exit(0)'], { + timeoutMs: 2_000, + }); + expect(Date.now() - startedAt).toBeLessThan(1_000); + return completed; + }); + expect(result.exitCode).toBe(0); + }); }); diff --git a/services/cloud-agent-next/wrapper/src/control/owned-processes.ts b/services/cloud-agent-next/wrapper/src/control/owned-processes.ts index 31fad70664..8cdd08d8d0 100644 --- a/services/cloud-agent-next/wrapper/src/control/owned-processes.ts +++ b/services/cloud-agent-next/wrapper/src/control/owned-processes.ts @@ -1,6 +1,7 @@ import { AsyncLocalStorage } from 'node:async_hooks'; import { spawn, + spawnSync, type ChildProcessWithoutNullStreams, type SpawnOptionsWithoutStdio, } from 'node:child_process'; @@ -31,12 +32,19 @@ export type OwnedProcessScope = { run(operation: () => T): T; seal(): void; dispose(): boolean; + observesOccupancy(): boolean; captureBaseline(allowed: (argv: string[]) => boolean, deadlineAt?: number): Promise; verify(baseline?: boolean, deadlineAt?: number): Promise; stop(deadlineAt: number): Promise; }; -type ProcessIdentity = { pid: number; parent: number; group: number; identity: string }; +type ProcessIdentity = { + pid: number; + parent: number; + group: number; + identity: string; + state: string; +}; type OwnedChild = { process: ChildProcessWithoutNullStreams; identity?: string; @@ -138,10 +146,13 @@ function processIdentity(pid: number, value: string): ProcessIdentity { .slice(value.lastIndexOf(')') + 2) .trim() .split(/\s+/); + const state = fields[0]; const startedAt = fields[19]; const parent = Number(fields[1]); const group = Number(fields[2]); if ( + !state || + !/^[A-Za-z]$/.test(state) || !startedAt || !/^\d+$/.test(startedAt) || !Number.isSafeInteger(parent) || @@ -151,7 +162,11 @@ function processIdentity(pid: number, value: string): ProcessIdentity { ) { throw new Error('Process identity unavailable'); } - return { pid, parent, group, identity: `${pid}:${startedAt}` }; + return { pid, parent, group, identity: `${pid}:${startedAt}`, state }; +} + +function isLiveProcessState(state: string): boolean { + return state !== 'Z' && state !== 'X' && state !== 'x'; } function closeDescriptors(descriptors: number[]): void { @@ -164,6 +179,14 @@ function closeDescriptors(descriptors: number[]): void { } } +function isErofs(error: unknown): boolean { + return typeof error === 'object' && error !== null && 'code' in error && error.code === 'EROFS'; +} + +function remountCgroupWritable(target: string): void { + spawnSync('mount', ['-o', 'remount,rw', target], { stdio: 'ignore' }); +} + function createCgroup(): Cgroup | undefined { const descriptors: number[] = []; let created: { directory: string; dev: number; ino: number } | undefined; @@ -195,14 +218,22 @@ function createCgroup(): Cgroup | undefined { ) { throw new Error('Process containment root unavailable'); } - descriptors.push( - openSync( - path.join(parentReference, 'cgroup.procs'), - constants.O_WRONLY | constants.O_NOFOLLOW - ) - ); + const parentProcs = path.join(parentReference, 'cgroup.procs'); + try { + descriptors.push(openSync(parentProcs, constants.O_WRONLY | constants.O_NOFOLLOW)); + } catch (error) { + if (!isErofs(error)) throw error; + remountCgroupWritable(parent); + descriptors.push(openSync(parentProcs, constants.O_WRONLY | constants.O_NOFOLLOW)); + } const name = `kilo-control-${crypto.randomUUID()}`; - mkdirSync(path.join(parentReference, name)); + try { + mkdirSync(path.join(parentReference, name)); + } catch (error) { + if (!isErofs(error)) throw error; + remountCgroupWritable(parent); + mkdirSync(path.join(parentReference, name)); + } const directory = path.join(parent, name); const descriptor = openSync( path.join(parentReference, name), @@ -232,7 +263,9 @@ function createCgroup(): Cgroup | undefined { } closeDescriptors(descriptors.splice(0, 2)); return { directory, reference, dev, ino, descriptors, procs, kill }; - } catch { + } catch (error) { + const message = error instanceof Error ? error.message : 'unknown'; + console.warn(`Owned process containment unavailable: ${message}`); if (created) { try { const fresh = lstatSync(created.directory); @@ -365,26 +398,37 @@ export function createOwnedProcessScope(): OwnedProcessScope { const baseline = new Set(); const observations = new Set(); const live = (child: OwnedChild): boolean => !child.exited && child.process.pid !== undefined; + const occupancyGroup = (): Cgroup | undefined => + group !== undefined && contained ? group : undefined; + const occupancyObservable = (): boolean => occupancyGroup() !== undefined; const verify = async (allowBaseline: boolean, deadline: Deadline): Promise => { try { deadline.check(); if (!used || removed || stopped) return true; - if (!group || !contained) return false; - await assertDirectory(group, deadline); + const observed = occupancyGroup(); + if (!observed) return false; + await assertDirectory(observed, deadline); const populated = population( - await readText(path.join(group.reference, 'cgroup.events'), deadline) + await readText(path.join(observed.reference, 'cgroup.events'), deadline) ); - await assertDirectory(group, deadline); + await assertDirectory(observed, deadline); if (populated === 0) return ![...children].some(live); - if (!allowBaseline || baseline.size === 0) return false; + if (!allowBaseline || baseline.size === 0) { + if (allowBaseline || [...children].some(live)) return false; + for (const pid of (await snapshotCgroup(observed, deadline)).pids) { + const { state } = processIdentity(pid, await readText(`/proc/${pid}/stat`, deadline)); + if (isLiveProcessState(state)) return false; + } + return true; + } const identities = new Set(); - for (const pid of (await snapshotCgroup(group, deadline)).pids) { + for (const pid of (await snapshotCgroup(observed, deadline)).pids) { const { identity } = processIdentity(pid, await readText(`/proc/${pid}/stat`, deadline)); if (!baseline.has(identity)) return false; identities.add(identity); } - await assertDirectory(group, deadline); + await assertDirectory(observed, deadline); return ( identities.size > 0 && [...children] @@ -509,13 +553,14 @@ export function createOwnedProcessScope(): OwnedProcessScope { return child; }, run: operation => current.run(scope, operation), + observesOccupancy: occupancyObservable, seal() { sealed = true; }, dispose() { sealed = true; if (removed) return true; - if (used && (!group || !contained || [...children].some(live))) return false; + if (used && (!occupancyObservable() || [...children].some(live))) return false; if ( group && !removeCgroup(group, stopDeadline?.deadlineAt ?? Date.now() + OBSERVATION_TIMEOUT_MS) @@ -530,9 +575,10 @@ export function createOwnedProcessScope(): OwnedProcessScope { const deadline = createDeadline(Math.min(deadlineAt, stopDeadline?.deadlineAt ?? Infinity)); observations.add(deadline); try { - if (!group || !contained || sealed) return; + const observed = occupancyGroup(); + if (!observed || sealed) return; const entries: (ProcessIdentity & { allowed: boolean })[] = []; - for (const pid of (await snapshotCgroup(group, deadline)).pids) { + for (const pid of (await snapshotCgroup(observed, deadline)).pids) { const before = processIdentity(pid, await readText(`/proc/${pid}/stat`, deadline)); const argv = (await readText(`/proc/${pid}/cmdline`, deadline)) .split('\0') @@ -619,7 +665,7 @@ export function createOwnedProcessScope(): OwnedProcessScope { } } await signalChildren('SIGKILL', deadline); - if (!group || !contained) return false; + if (!occupancyObservable()) return false; while (true) { if (await verify(false, deadline)) return true; await deadline.wait(signal => delay(25, undefined, { signal })); diff --git a/services/cloud-agent-next/wrapper/src/control/sandbox-control-client.ts b/services/cloud-agent-next/wrapper/src/control/sandbox-control-client.ts index 9c54f9ddd1..1decee326b 100644 --- a/services/cloud-agent-next/wrapper/src/control/sandbox-control-client.ts +++ b/services/cloud-agent-next/wrapper/src/control/sandbox-control-client.ts @@ -737,7 +737,7 @@ export function createSandboxControlClient( publication: { event: 'session.event' | 'session.preparing'; receiptId: string; - receiptHash: string; + sequence: number; session: SessionEventIdentity; payload: unknown; }, diff --git a/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.test.ts b/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.test.ts index 234810d6f7..29adfe1df6 100644 --- a/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.test.ts @@ -84,6 +84,7 @@ function completion(error?: Completion['info']['error']): Completion { function fakeKilo(overrides: Partial = {}): WrapperKiloClient { return { getSession: async id => ({ id }), + getSessionDetails: async (id, directory) => ({ id, directory }), ensureSession: async () => undefined, sendPrompt: async () => completion(), sendPromptAsync: async () => {}, @@ -1870,7 +1871,8 @@ describe('production worktree deletion routes', () => { item.event.properties.status === 'cancelled' ) ).toBe(true); - expect(lookups).toEqual([directory]); + expect(lookups.length).toBeGreaterThan(0); + expect(lookups.every(value => value === directory)).toBe(true); expect(http.requests.every(request => request.directory === directory)).toBe(true); expect(http.requests).toContainEqual({ method: 'POST', @@ -4339,6 +4341,7 @@ describe('control wrapper heartbeat source policy', () => { expect(source).toContain( "onDisconnected: () => shutdown(1, 'Sandbox control connection lost', 'control_disconnected')" ); + expect(source).toContain('shutdown(1, failure.reason, heartbeatReasonFrom(failure.reason))'); expect(source).toContain( 'if (!payload.kilo.ready && heartbeatReason) payload.kilo.reason = heartbeatReason;' ); diff --git a/services/cloud-agent-next/wrapper/src/control/sandbox-control-runtime.test.ts b/services/cloud-agent-next/wrapper/src/control/sandbox-control-runtime.test.ts index 0bbe974ae0..c9b0cf6da3 100644 --- a/services/cloud-agent-next/wrapper/src/control/sandbox-control-runtime.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/sandbox-control-runtime.test.ts @@ -384,7 +384,7 @@ describe('maybeStartSandboxControlClient', () => { }); it.each(['false', 'throw'] as const)( - 'retires the connection when heartbeat delivery returns %s', + 'leaves final disconnect to reconnect exhaustion when heartbeat delivery returns %s', async failure => { const timers = spyOn(globalThis, 'setInterval'); let disconnected = 0; @@ -401,12 +401,15 @@ describe('maybeStartSandboxControlClient', () => { onDisconnected: () => { disconnected++; }, - createClient: () => ({ + createClient: clientOptions => ({ connect: async () => {}, close: () => {}, sendEvent: event => { if (event !== 'sandbox.heartbeat') return true; if (failure === 'throw') throw new Error('private transport error'); + const onConnectionLost = clientOptions.onConnectionLost; + if (!onConnectionLost) throw new Error('Missing connection-loss callback'); + onConnectionLost(); return false; }, }), @@ -414,7 +417,7 @@ describe('maybeStartSandboxControlClient', () => { ); try { await flushAsyncWork(); - expect(disconnected).toBe(1); + expect(disconnected).toBe(0); expect( timers.mock.calls.filter(([, ms]) => ms === SANDBOX_CONTROL_REPORT_INTERVAL_MS) ).toHaveLength(0); @@ -425,6 +428,40 @@ describe('maybeStartSandboxControlClient', () => { } ); + it('leaves final disconnect to reconnect exhaustion when the ready event fails', async () => { + let disconnected = 0; + const started = maybeStartSandboxControlClient( + { + SANDBOX_CONTROL_URL: 'wss://example.test/sandbox-control/sbx_1', + SANDBOX_CONTROL_CREDENTIAL: 'secret', + PROVIDER_INSTANCE_ID: 'inst_1', + }, + () => {}, + { + wrapperVersion: '2.4.0', + onDisconnected: () => { + disconnected++; + }, + createClient: clientOptions => ({ + connect: async () => {}, + close: () => {}, + sendEvent: () => { + const onConnectionLost = clientOptions.onConnectionLost; + if (!onConnectionLost) throw new Error('Missing connection-loss callback'); + onConnectionLost(); + return false; + }, + }), + } + ); + try { + await flushAsyncWork(); + expect(disconnected).toBe(0); + } finally { + started?.close(); + } + }); + it('clears reconnect after close and omits credentials from connect failure logs', async () => { const credential = 'super-secret-token'; const logs: string[] = []; diff --git a/services/cloud-agent-next/wrapper/src/control/sandbox-control-runtime.ts b/services/cloud-agent-next/wrapper/src/control/sandbox-control-runtime.ts index 9e2fce77f1..4985d0c126 100644 --- a/services/cloud-agent-next/wrapper/src/control/sandbox-control-runtime.ts +++ b/services/cloud-agent-next/wrapper/src/control/sandbox-control-runtime.ts @@ -417,14 +417,14 @@ export function maybeStartSandboxControlClient( try { if (!active.sendEvent?.('sandbox.heartbeat', payload)) { diagnostic('send_failed'); - handleDisconnected(); + handleConnectionLost(); } else { lastSentAt = Date.now(); diagnostic('sent'); } } catch { diagnostic('send_threw'); - handleDisconnected(); + handleConnectionLost(); } } @@ -432,7 +432,7 @@ export function maybeStartSandboxControlClient( if (closed || options.isReady?.() === false) return; if (sampleAbort.signal.aborted) sampleAbort = new AbortController(); if (!active.sendEvent?.('sandbox.ready', { kiloReady: true, globalFeedAttached: true })) { - handleDisconnected(); + handleConnectionLost(); return; } sendHeartbeat(active); diff --git a/services/cloud-agent-next/wrapper/src/control/session-operation-cleanup-proof.test.ts b/services/cloud-agent-next/wrapper/src/control/session-operation-cleanup-proof.test.ts new file mode 100644 index 0000000000..bab7a71348 --- /dev/null +++ b/services/cloud-agent-next/wrapper/src/control/session-operation-cleanup-proof.test.ts @@ -0,0 +1,277 @@ +import { afterEach, beforeEach, describe, expect, it, jest } from 'bun:test'; +import { fakeKilo, session } from './control-test-fixtures'; +import { createOwnedProcessScope } from './owned-processes'; +import { SessionOperationCleanup, type NativeOperationTarget } from './session-operation-cleanup'; +import { + forgetAttachedRoot, + rememberAttachedRoot, + rememberChildSession, + resetSessionDirectoryState, +} from './session-directories'; +import type { WrapperKiloClient } from '../kilo-api'; + +type Statuses = Awaited>; + +beforeEach(() => { + resetSessionDirectoryState(); + rememberAttachedRoot(session.kiloSessionId, session.directory); +}); + +afterEach(() => { + jest.useRealTimers(); + resetSessionDirectoryState(); +}); + +function fixture(overrides: Partial = {}) { + const client = fakeKilo({ getSessionStatuses: async () => ({}), ...overrides }); + const target = { runtimeId: 'native_exact', client }; + const processes = createOwnedProcessScope(); + const stop = jest.spyOn(processes, 'stop').mockResolvedValue(true); + const verify = jest.fn( + async (observed: NativeOperationTarget, deadlineAt: number) => + observed === target && Date.now() < deadlineAt + ); + const confirmed = jest.fn(); + let current = true; + const cleanup = new SessionOperationCleanup( + session, + processes, + verify, + confirmed, + observed => current && observed === target + ); + return { + cleanup, + target, + stop, + verify, + confirmed, + replaceRuntime: () => { + current = false; + }, + run: (deadlineAt = Date.now() + 150) => + cleanup.cleanup({ + deadlineAt, + target, + completionEvidence: 'unconfirmed', + cancel: () => {}, + }), + }; +} + +describe('native cancellation proof', () => { + it('accepts native empty-map idle only with exact process proof after acknowledged abort', async () => { + const abort = jest.fn(async () => true); + const f = fixture({ abortSession: abort }); + const deadline = Date.now() + 1_000; + expect(await f.run(deadline)).toBe(true); + expect(abort).toHaveBeenCalledTimes(1); + expect(f.stop).toHaveBeenCalledWith(deadline); + expect(f.verify).toHaveBeenCalledWith(f.target, deadline); + expect(f.cleanup.cleanupState).toBe('confirmed'); + expect(f.confirmed).toHaveBeenCalledTimes(1); + expect(await f.run(deadline + 10_000)).toBe(true); + expect(abort).toHaveBeenCalledTimes(1); + }); + + it('does not require an unrelated known busy root to become idle', async () => { + rememberAttachedRoot('root_b', session.directory); + const f = fixture({ getSessionStatuses: async () => ({ root_b: { type: 'busy' } }) }); + expect(await f.run()).toBe(true); + }); + + it.each(['busy', 'retry', 'offline'])( + 'waits for relevant %s status to become native idle', + async type => { + let reads = 0; + const f = fixture({ + getSessionStatuses: async () => { + reads++; + if (reads === 1) { + expect(f.verify).not.toHaveBeenCalled(); + return { [session.kiloSessionId]: { type } }; + } + return {}; + }, + }); + expect(await f.run(Date.now() + 1_000)).toBe(true); + expect(reads).toBe(2); + } + ); + + it.each(['busy', 'retry', 'offline'])( + 'does not confirm persistent relevant %s status', + async type => { + const f = fixture({ + getSessionStatuses: async () => ({ [session.kiloSessionId]: { type } }), + }); + expect(await f.run()).toBe(false); + expect(f.verify).not.toHaveBeenCalled(); + } + ); + + it.each(['child', 'unknown'])('fails closed for an active %s with no idle proof', async id => { + if (id === 'child') rememberChildSession({ childId: id, parentId: session.kiloSessionId }); + const f = fixture({ getSessionStatuses: async () => ({ [id]: { type: 'busy' } }) }); + expect(await f.run()).toBe(false); + expect(f.verify).not.toHaveBeenCalled(); + }); + + it('fails closed for a descendant in an unobserved directory', async () => { + rememberChildSession({ + childId: 'child', + parentId: session.kiloSessionId, + directory: '/other', + }); + expect(await fixture().run()).toBe(false); + }); + + it.each(['operation', 'native'])( + 'requires %s owned-process proof despite native idle', + async scope => { + const f = fixture(); + if (scope === 'operation') f.stop.mockResolvedValue(false); + else f.verify.mockResolvedValue(false); + expect(await f.run()).toBe(false); + expect(f.confirmed).not.toHaveBeenCalled(); + } + ); + + it.each(['failed', 'undefined'])( + 'does not turn a %s native status read into idle', + async failure => { + const f = fixture({ + getSessionStatuses: async () => { + if (failure === 'failed') throw new Error('status unavailable'); + return undefined as unknown as Statuses; + }, + }); + expect(await f.run()).toBe(false); + expect(f.verify).not.toHaveBeenCalled(); + } + ); + + it.each(['missing', 'wrong_id', 'wrong_directory'])( + 'rejects %s native session existence proof', + async failure => { + const f = fixture({ + getSessionDetails: async () => { + if (failure === 'missing') throw new Error('Session not found'); + return { + id: failure === 'wrong_id' ? 'other' : session.kiloSessionId, + directory: failure === 'wrong_directory' ? '/other' : session.directory, + }; + }, + }); + expect(await f.run()).toBe(false); + expect(f.verify).not.toHaveBeenCalled(); + } + ); + + it('waits for an owned descendant to disappear from the native active map', async () => { + rememberChildSession({ childId: 'child', parentId: session.kiloSessionId }); + let reads = 0; + const f = fixture({ + getSessionStatuses: async (): Promise => + ++reads === 1 ? { child: { type: 'busy' } } : {}, + }); + expect(await f.run()).toBe(true); + expect(reads).toBe(2); + }); + + it('rejects runtime replacement while process proof is pending', async () => { + const f = fixture(); + f.verify.mockImplementation(async () => { + f.replaceRuntime(); + return true; + }); + expect(await f.run()).toBe(false); + }); + + it('aborts a hung status read at the original deadline', async () => { + let readSignal: AbortSignal | undefined; + const f = fixture({ + getSessionStatuses: (_directory, signal) => { + readSignal = signal; + return new Promise(() => {}); + }, + }); + const deadline = Date.now() + 100; + expect(await f.run(deadline)).toBe(false); + expect(readSignal?.aborted).toBe(true); + expect(f.cleanup.cleanupDeadline).toBe(deadline); + expect(f.verify).not.toHaveBeenCalled(); + }); + + it('does not give status observation a new budget after a delayed abort', async () => { + const abort = Promise.withResolvers(); + let reads = 0; + const f = fixture({ + abortSession: () => abort.promise, + getSessionStatuses: async () => { + reads++; + return { [session.kiloSessionId]: { type: 'busy' } }; + }, + }); + const deadline = Date.now() + 150; + const pending = f.run(deadline); + await new Promise(resolve => setTimeout(resolve, 100)); + abort.resolve(true); + expect(await pending).toBe(false); + expect(reads).toBeGreaterThan(0); + expect(Date.now()).toBeLessThan(deadline + 75); + expect(f.cleanup.cleanupDeadline).toBe(deadline); + }); + + it('rejects stale runtime identity before abort', async () => { + const abort = jest.fn(async () => true); + const f = fixture({ abortSession: abort }); + f.replaceRuntime(); + expect(await f.run()).toBe(false); + expect(abort).not.toHaveBeenCalled(); + }); + + it('rejects runtime replacement during status observation', async () => { + const f = fixture({ + getSessionStatuses: async () => { + f.replaceRuntime(); + return {}; + }, + }); + expect(await f.run()).toBe(false); + expect(f.verify).not.toHaveBeenCalled(); + }); + + it('rejects attachment replacement during status observation', async () => { + const f = fixture({ + getSessionStatuses: async () => { + forgetAttachedRoot(session.kiloSessionId); + rememberAttachedRoot(session.kiloSessionId, session.directory); + return {}; + }, + }); + expect(await f.run()).toBe(false); + }); + + it('does not infer existence from an empty map for an unknown session', async () => { + forgetAttachedRoot(session.kiloSessionId); + expect(await fixture().run()).toBe(false); + }); + + it('retains the original cleanup deadline and abort acknowledgement while polling', async () => { + const abort = jest.fn(async () => true); + const f = fixture({ + abortSession: abort, + getSessionStatuses: async () => ({ [session.kiloSessionId]: { type: 'busy' } }), + }); + const deadline = Date.now() + 150; + const first = f.run(deadline); + expect(await f.run(deadline + 10_000)).toBe(false); + expect(await first).toBe(false); + expect(f.cleanup.cleanupDeadline).toBe(deadline); + expect(abort).toHaveBeenCalledTimes(1); + expect(f.cleanup.cleanupState).toBe('unconfirmed'); + expect(await f.run(deadline + 20_000)).toBe(false); + expect(abort).toHaveBeenCalledTimes(1); + }); +}); diff --git a/services/cloud-agent-next/wrapper/src/control/session-operation-cleanup.test.ts b/services/cloud-agent-next/wrapper/src/control/session-operation-cleanup.test.ts index c42cbeb35e..b3b9578f83 100644 --- a/services/cloud-agent-next/wrapper/src/control/session-operation-cleanup.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/session-operation-cleanup.test.ts @@ -127,7 +127,7 @@ describe('SessionOperation cleanup', () => { await operation.done; }); - it('does not treat an abort acknowledgement or missing status as quiescence', async () => { + it('does not treat acknowledged abort and native empty-map idle as owned-process proof', async () => { let aborts = 0; const client = fakeKilo({ sendPrompt: async () => completion(), @@ -302,11 +302,10 @@ describe('SessionOperation cleanup', () => { result: { status: 'aborted', quiescent: true, - runtimeRetired: true, - nativeRuntimeId: 'native_1', }, }); - expect(handlerDeps.kiloRuntimes?.get(session.directory)).toBeUndefined(); + expect(handlerDeps.kiloRuntimes?.get(session.directory)?.runtimeId).toBe('native_1'); + expect(operationB?.snapshot().local?.result).toEqual({ ok: true, result: {} }); }); it('keeps a failed local attachment unconfirmed when its captured cleanup cannot prove retirement', async () => { diff --git a/services/cloud-agent-next/wrapper/src/control/session-operation-cleanup.ts b/services/cloud-agent-next/wrapper/src/control/session-operation-cleanup.ts index 0d23888df3..03e3bb201f 100644 --- a/services/cloud-agent-next/wrapper/src/control/session-operation-cleanup.ts +++ b/services/cloud-agent-next/wrapper/src/control/session-operation-cleanup.ts @@ -1,3 +1,4 @@ +import { setTimeout as delay } from 'node:timers/promises'; import { SANDBOX_CONTROL_CLEANUP_TIMEOUT_MS, type SessionRequestIdentity, @@ -6,6 +7,7 @@ import type { WrapperKiloClient } from '../kilo-api.js'; import { withTimeoutAndAbort } from '../utils.js'; import { withKiloRequestDeadline } from './sandbox-control-runtime.js'; import type { OwnedProcessScope } from './owned-processes.js'; +import { directoriesForRoot, rootAttachmentId, rootForSession } from './session-directories.js'; export type NativeOperationTarget = Readonly<{ runtimeId: string; @@ -31,7 +33,8 @@ export class SessionOperationCleanup { target: NativeOperationTarget, deadlineAt: number ) => Promise, - private readonly onConfirmed: () => void + private readonly onConfirmed: () => void, + private readonly isCurrent: (target: NativeOperationTarget) => boolean ) {} get cleanupDeadline(): number | undefined { @@ -83,20 +86,9 @@ export class SessionOperationCleanup { const retired = await input.preClientCleanup?.(deadlineAt); return retired === 'retired' || retired === 'stale'; } + if (!this.isCurrent(target) || Date.now() >= deadlineAt) return false; if (!(await this.abortNative(target, deadlineAt))) return false; - const statuses = await withTimeoutAndAbort( - withKiloRequestDeadline(signal => - client.getSessionStatuses(this.session.directory, signal) - ), - { - timeoutMs: Math.max(1, deadlineAt - Date.now()), - timeoutMessage: 'Kilo cleanup status probe timed out', - abortMessage: 'Kilo cleanup status probe cancelled', - } - ); - const observed = Object.values(statuses); - if (observed.length === 0 || !observed.every(value => value.type === 'idle')) return false; - return this.verifyQuiescence(target, deadlineAt); + return this.observeQuiescence(target, client, deadlineAt); })() .catch(() => false) .then(confirmed => this.confirm(confirmed, deadlineAt)); @@ -112,6 +104,53 @@ export class SessionOperationCleanup { return quiescent; } + private async observeQuiescence( + target: NativeOperationTarget, + client: WrapperKiloClient, + deadlineAt: number + ): Promise { + const { kiloSessionId, directory } = this.session; + const attachment = rootAttachmentId(kiloSessionId); + const current = () => + attachment !== undefined && + rootAttachmentId(kiloSessionId) === attachment && + rootForSession(kiloSessionId, directory) === kiloSessionId && + directoriesForRoot(kiloSessionId, directory).every(value => value === directory) && + this.isCurrent(target) && + Date.now() < Math.min(deadlineAt, this.deadlineAt ?? Infinity); + if (!current()) return false; + const controller = new AbortController(); + try { + return await withTimeoutAndAbort( + withKiloRequestDeadline(async signal => { + const session = await client.getSessionDetails(kiloSessionId, directory, signal); + if (session.id !== kiloSessionId || session.directory !== directory) return false; + while (current()) { + const statuses = await client.getSessionStatuses(directory, signal); + if (!current()) return false; + let idle = true; + for (const [id, status] of Object.entries(statuses)) { + if (status.type === 'idle') continue; + const root = rootForSession(id, directory); + if (!root) return false; + if (root === kiloSessionId) idle = false; + } + if (idle) return (await this.verifyQuiescence(target, deadlineAt)) && current(); + await delay(Math.min(25, Math.max(1, deadlineAt - Date.now())), undefined, { signal }); + } + return false; + }, controller.signal), + { + timeoutMs: Math.max(1, deadlineAt - Date.now()), + timeoutMessage: 'Kilo cleanup status probe timed out', + abortMessage: 'Kilo cleanup status probe cancelled', + } + ); + } finally { + controller.abort(); + } + } + private abortNative(target: NativeOperationTarget, deadlineAt: number): Promise { if (this.nativeAbort) return this.nativeAbort; const client = target.client; diff --git a/services/cloud-agent-next/wrapper/src/control/session-operation.ts b/services/cloud-agent-next/wrapper/src/control/session-operation.ts index fd0a00f5a3..a06e50786c 100644 --- a/services/cloud-agent-next/wrapper/src/control/session-operation.ts +++ b/services/cloud-agent-next/wrapper/src/control/session-operation.ts @@ -1,6 +1,8 @@ import { isDeepStrictEqual } from 'node:util'; import { + diagnosticDetail, emitControlDiagnostic, + OWNED_PROCESS_CLEANUP_UNREAPED, type ControlDiagnosticReporter, } from '../../../src/shared/control-diagnostics.js'; import { @@ -196,7 +198,16 @@ export class SessionOperation { this.session, this.processes, deps.verifyQuiescence, - () => deps.onCleanupConfirmed() + () => deps.onCleanupConfirmed(), + target => { + const runtime = deps.getRuntime(); + return ( + runtime?.runtimeId === target.runtimeId && + runtime.kiloClient === target.client && + runtime.directory === this.session.directory && + !runtime.signal.aborted + ); + } ); this.diagnostic('started'); this.timeout = setTimeout( @@ -325,6 +336,23 @@ export class SessionOperation { this.deps.retireRuntime(reason, this.captureCleanupDeadline(deadlineAt), this.nativeTarget()); } + reportUnreapedProcessCleanup(populated: boolean): void { + const detail = diagnosticDetail( + `owned_process_unreaped populated=${populated ? '1' : '0'} ${this.session.directory}` + ); + emitControlDiagnostic(this.deps.onDiagnostic, 'session.task', { + sessionId: this.session.sessionId, + kiloSessionId: this.session.kiloSessionId, + messageId: this.messageId, + kind: this.work.operation === 'session.attach' ? 'preparation' : 'execution', + stage: 'process_cleanup', + phase: 'failed', + ok: false, + ...(detail ? { detail } : {}), + }); + console.error(OWNED_PROCESS_CLEANUP_UNREAPED); + } + confirmCleanup(confirmed: boolean, deadlineAt: number): boolean { return this.cleanupOwner.confirm(confirmed, deadlineAt); } diff --git a/services/cloud-agent-next/wrapper/src/control/worktree-feed.test.ts b/services/cloud-agent-next/wrapper/src/control/worktree-feed.test.ts index c089b51df1..4eb9df8655 100644 --- a/services/cloud-agent-next/wrapper/src/control/worktree-feed.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/worktree-feed.test.ts @@ -66,6 +66,7 @@ function fixture(rejectReconnections = false) { const attempts: ReturnType[] = []; const failures: string[] = []; const events: KiloFeedEvent[] = []; + const diagnostics: Array<{ phase?: string; detail?: string }> = []; const start = spyOn(controlRuntime, 'startSandboxControlEventFeed').mockImplementation( async options => { if (rejectReconnections && attempts.length > 0) throw new Error('Feed unavailable'); @@ -80,8 +81,13 @@ function fixture(rejectReconnections = false) { runtimeId === source.runtimeId && kiloClient === source.kiloClient, onEvent: event => events.push(event), onFailure: reason => failures.push(reason), + onDiagnostic: (_event, fields) => + diagnostics.push({ + phase: typeof fields.phase === 'string' ? fields.phase : undefined, + detail: typeof fields.detail === 'string' ? fields.detail : undefined, + }), }); - return { attempts, events, failures, feed, source, start }; + return { attempts, diagnostics, events, failures, feed, source, start }; } const cleanups: Array<() => void> = []; @@ -312,6 +318,12 @@ describe('createWorktreeFeed', () => { await waitFor(() => h.failures.length === 1); expect(h.start).toHaveBeenCalledTimes(1 + SANDBOX_CONTROL_RECOVERY_MAX_ATTEMPTS); expect(h.failures).toEqual(['feed_ended']); + expect(h.diagnostics).toEqual( + expect.arrayContaining([ + { phase: 'retry_scheduled', detail: 'feed_ended' }, + { phase: 'failed', detail: 'feed_ended' }, + ]) + ); expect(h.source.signal.aborted).toBe(false); expect(h.feed.prepareForNewWork()).toBe(false); }); diff --git a/services/cloud-agent-next/wrapper/src/control/worktree-feed.ts b/services/cloud-agent-next/wrapper/src/control/worktree-feed.ts index 3ebc4f6712..2e9529fe76 100644 --- a/services/cloud-agent-next/wrapper/src/control/worktree-feed.ts +++ b/services/cloud-agent-next/wrapper/src/control/worktree-feed.ts @@ -1,6 +1,7 @@ import { setTimeout as delay } from 'node:timers/promises'; import { createKiloClient as createKiloEventClient } from '@kilocode/sdk/v2/client'; import { + diagnosticDetail, emitControlDiagnostic, type ControlDiagnosticReporter, } from '../../../src/shared/control-diagnostics.js'; @@ -173,11 +174,24 @@ export function createWorktreeFeed(options: { } } + function feedDiagnostic( + phase: 'retry_scheduled' | 'failed', + reason: KiloEventFeedError['reason'] + ): void { + const detail = diagnosticDetail(reason); + emitControlDiagnostic(options.onDiagnostic, 'control.feed', { + phase, + scopeId, + ...(detail ? { detail } : {}), + }); + } + function unavailable(current: Recovery): void { if (!isCurrent() || recovery !== current) return; recovery = undefined; state = 'unavailable'; options.onStateChange?.(); + feedDiagnostic('failed', current.reason); options.onFailure(current.reason); } @@ -220,6 +234,7 @@ export function createWorktreeFeed(options: { state = 'recovering'; closeActive(); options.onStateChange?.(); + feedDiagnostic('retry_scheduled', reason); void retry(current); } diff --git a/services/cloud-agent-next/wrapper/src/control/worktree-runtime.test.ts b/services/cloud-agent-next/wrapper/src/control/worktree-runtime.test.ts index 4326cbc09b..a75feda600 100644 --- a/services/cloud-agent-next/wrapper/src/control/worktree-runtime.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/worktree-runtime.test.ts @@ -185,7 +185,10 @@ function createKiloStub(health: unknown = { healthy: true, version: '7.4.20' }) return Response.json(completion); } if (request.method === 'GET' && url.pathname.startsWith('/session/')) { - return Response.json({ id: decodeURIComponent(url.pathname.slice('/session/'.length)) }); + return Response.json({ + id: decodeURIComponent(url.pathname.slice('/session/'.length)), + directory: url.searchParams.get('directory'), + }); } if (request.method === 'POST' && url.pathname === '/pty') { return Response.json({ diff --git a/services/cloud-agent-next/wrapper/src/kilo-api.test.ts b/services/cloud-agent-next/wrapper/src/kilo-api.test.ts index fc72af638a..319c76cb38 100644 --- a/services/cloud-agent-next/wrapper/src/kilo-api.test.ts +++ b/services/cloud-agent-next/wrapper/src/kilo-api.test.ts @@ -569,6 +569,39 @@ describe('createWrapperKiloClient generated SDK HTTP boundary', () => { } ); + it.each( + [null, [], '', { ses_1: null }, { ses_1: {} }, { ses_1: { type: 1 } }].map(body => ({ body })) + )('rejects malformed successful session status maps: %j', async ({ body }) => { + const result: unknown = await createClient(startStub(200, body).url) + .getSessionStatuses() + .catch(error => error); + expect(result).toBeInstanceOf(Error); + }); + + it.each([ + { status: 204, body: undefined }, + { status: 500, body: {} }, + ])('rejects unavailable session status: %j', async ({ status, body }) => { + const result: unknown = await createClient(startStub(status, body).url) + .getSessionStatuses() + .catch(error => error); + expect(result).toBeInstanceOf(Error); + }); + + it('validates exact session existence for cleanup without importing', async () => { + const directory = '/workspace/exact'; + const client = createClient(startStub(200, { id: 'ses_1', directory }).url); + expect(await client.getSessionDetails('ses_1', directory)).toEqual({ id: 'ses_1', directory }); + for (const lookup of [ + () => client.getSessionDetails('ses_other', directory), + () => client.getSessionDetails('ses_1', '/other'), + () => createClient(startStub(404, {}).url).getSessionDetails('ses_1', directory), + ]) { + const result: unknown = await lookup().catch(error => error); + expect(result).toBeInstanceOf(Error); + } + }); + it('accepts genuinely empty successful session status and pending-input results', async () => { expect(await createClient(startStub(200, {}).url).getSessionStatuses()).toEqual({}); expect(await createClient(startStub(200, []).url).getQuestions()).toEqual([]); diff --git a/services/cloud-agent-next/wrapper/src/kilo-api.ts b/services/cloud-agent-next/wrapper/src/kilo-api.ts index 9badda8ced..d65149461d 100644 --- a/services/cloud-agent-next/wrapper/src/kilo-api.ts +++ b/services/cloud-agent-next/wrapper/src/kilo-api.ts @@ -14,9 +14,15 @@ import { type SessionCommandResponse, type SessionPromptResponse, } from '@kilocode/sdk/v2'; +import { z } from 'zod'; import { logToFile } from './utils.js'; import { toSlashCommandInfo, type SlashCommandInfo } from '../../src/shared/slash-commands.js'; +const sessionStatusesSchema = z.record( + z.string().min(1), + z.object({ type: z.string().min(1) }).passthrough() +); + function isRecord(value: unknown): value is Record { return typeof value === 'object' && value !== null; } @@ -237,6 +243,11 @@ type PromptOptions = { export type WrapperKiloClient = { createSession: (opts?: { title?: string }) => Promise<{ id: string }>; getSession: (sessionId: string) => Promise<{ id: string }>; + getSessionDetails: ( + sessionId: string, + directory: string, + signal?: AbortSignal + ) => Promise<{ id: string; directory: string }>; ensureSession: (sessionId: string, directory: string, signal?: AbortSignal) => Promise; sendPrompt: (opts: PromptOptions) => Promise; sendPromptAsync: (opts: PromptOptions) => Promise; @@ -387,6 +398,15 @@ export function createWrapperKiloClient( return { id: data.id }; }, + getSessionDetails: async (sessionId, directory, signal) => { + const result = await v2Client.session.get({ sessionID: sessionId, directory }, { signal }); + const data = requireSdkData(result, 'Session cleanup lookup'); + if (data.id !== sessionId || data.directory !== directory) { + throw new Error('Session cleanup lookup returned an invalid session'); + } + return { id: data.id, directory: data.directory }; + }, + ensureSession: async (sessionId, directory, signal) => { const lookupTimeout = AbortSignal.timeout(5_000); const lookupSignal = signal ? AbortSignal.any([signal, lookupTimeout]) : lookupTimeout; @@ -546,7 +566,9 @@ export function createWrapperKiloClient( getSessionStatuses: async (directory = workspacePath, signal) => { const result = await v2Client.session.status({ directory }, { signal }); - return requireSdkData(result, 'Session status'); + const parsed = sessionStatusesSchema.safeParse(requireSdkData(result, 'Session status')); + if (!parsed.success) throw new Error('Session status returned an invalid map'); + return parsed.data; }, getQuestions: async (directory = workspacePath, signal) => { diff --git a/services/cloud-agent-next/wrapper/src/log-uploader.ts b/services/cloud-agent-next/wrapper/src/log-uploader.ts index e1238da9dc..dc975cc160 100644 --- a/services/cloud-agent-next/wrapper/src/log-uploader.ts +++ b/services/cloud-agent-next/wrapper/src/log-uploader.ts @@ -1,6 +1,6 @@ import { randomUUID } from 'node:crypto'; import { existsSync } from 'node:fs'; -import { basename, dirname } from 'node:path'; +import { basename, dirname, join } from 'node:path'; import { spawn } from 'node:child_process'; import { logToFile, withTimeoutAndAbort } from './utils.js'; @@ -36,20 +36,24 @@ export function createLogArchiveId(wrapperRunId: string): string { return `${wrapperRunId}--${randomUUID()}`; } +export type TarArchiveEntry = { directory: string; name: string }; + type TarStream = { stream: ReadableStream; kill: () => void; }; -function createTarStream(paths: Array): TarStream | undefined { - const existing = paths.filter(f => existsSync(f)); +export function createTarStream(paths: Array): TarStream | undefined { + const existing = paths + .map(path => + typeof path === 'string' ? { directory: dirname(path), name: basename(path) } : path + ) + .filter(entry => existsSync(join(entry.directory, entry.name))); if (existing.length === 0) return undefined; - // Use -C parent basename for each path so the archive contains relative names, not full paths. - // Works for both files and directories. const tarArgs = ['czf', '-']; - for (const f of existing) { - tarArgs.push('-C', dirname(f), basename(f)); + for (const entry of existing) { + tarArgs.push('-C', entry.directory, entry.name); } const proc = spawn('tar', tarArgs, { stdio: ['ignore', 'pipe', 'pipe'] }); const { stdout, stderr: stderrStream } = proc; diff --git a/services/cloud-agent-next/wrapper/src/utils.ts b/services/cloud-agent-next/wrapper/src/utils.ts index 41c9d6e24e..d627eebbb4 100644 --- a/services/cloud-agent-next/wrapper/src/utils.ts +++ b/services/cloud-agent-next/wrapper/src/utils.ts @@ -1,6 +1,7 @@ import { spawn } from 'child_process'; import { appendFileSync } from 'fs'; -import { currentOwnedProcessScope } from './control/owned-processes.js'; +import { setTimeout as delay } from 'node:timers/promises'; +import { currentOwnedProcessScope, type OwnedProcessScope } from './control/owned-processes.js'; export type ExecResult = { stdout: string; @@ -45,6 +46,18 @@ const EXEC_HARD_TIMEOUT_MESSAGE = 'exec hard timeout reached'; const EXEC_ABORTED_MESSAGE = 'exec aborted'; const DEFAULT_MAX_OUTPUT_BYTES = 64 * 1_024; const TRUNCATION_MARKER = 'output truncated'; +const OWNED_TREE_OBSERVATION_MS = 1_000; + +async function waitForOwnedTree(scope: OwnedProcessScope, deadlineAt: number): Promise { + if (process.platform !== 'linux' || !scope.observesOccupancy() || Date.now() >= deadlineAt) + return; + while (Date.now() < deadlineAt) { + if (await scope.verify(false, deadlineAt)) return; + const remaining = deadlineAt - Date.now(); + if (remaining <= 0) return; + await delay(Math.min(25, remaining)); + } +} export type TerminationReason = 'timeout' | 'inactivity_timeout' | 'hard_timeout' | 'abort'; @@ -128,8 +141,9 @@ export function runProcess( ? { env: { ...process.env, ...opts.env } } : {}), }; + const owned = currentOwnedProcessScope(); const proc = - currentOwnedProcessScope()?.spawn(command, args, options) ?? + owned?.spawn(command, args, options) ?? spawn(command, args, { ...options, detached: true, @@ -276,31 +290,45 @@ export function runProcess( opts.signal.addEventListener('abort', abortHandler, { once: true }); } } + let finishing = false; proc.on('close', (code, signal) => { - if (settled) return; + if (settled || finishing) return; if (terminationReason !== null) { waitForTerminatedGroup(); return; } - settled = true; + finishing = true; clearTimers(); removeAbortHandler(); - resolve({ - stdout, - stderr, - exitCode: code ?? (signal === null ? 0 : 1), - elapsedMs: Date.now() - startedAt, - ...(stdoutTruncated ? { stdoutTruncated: true } : {}), - ...(stderrTruncated ? { stderrTruncated: true } : {}), - }); - }); - proc.on('error', err => { - if (!settled) { + const exitCode = code ?? (signal === null ? 0 : 1); + const complete = (): void => { + if (settled) return; settled = true; - clearTimers(); - removeAbortHandler(); - reject(err); + resolve({ + stdout, + stderr, + exitCode, + elapsedMs: Date.now() - startedAt, + ...(stdoutTruncated ? { stdoutTruncated: true } : {}), + ...(stderrTruncated ? { stderrTruncated: true } : {}), + }); + }; + if (!owned || process.platform !== 'linux') { + complete(); + return; } + const remainingMs = + opts?.timeoutMs !== undefined + ? Math.max(0, opts.timeoutMs - (Date.now() - startedAt)) + : OWNED_TREE_OBSERVATION_MS; + void waitForOwnedTree(owned, Date.now() + remainingMs).then(complete, complete); + }); + proc.on('error', err => { + if (settled || finishing) return; + settled = true; + clearTimers(); + removeAbortHandler(); + reject(err); }); }); }