diff --git a/packages/cloud-agent-sdk/src/normalizer.test.ts b/packages/cloud-agent-sdk/src/normalizer.test.ts index c7f7f0a125..24cec70b33 100644 --- a/packages/cloud-agent-sdk/src/normalizer.test.ts +++ b/packages/cloud-agent-sdk/src/normalizer.test.ts @@ -1812,9 +1812,32 @@ describe('normalize', () => { }); }); - it('defaults error when missing', () => { + it('uses attach_exhausted as the error when error is missing', () => { const result = normalize( - createRaw('cloud.message.failed', { messageId: 'msg', delivery: 'sent' }) + createRaw('cloud.message.failed', { + messageId: 'msg', + delivery: 'sent', + reason: 'attach_exhausted', + }) + ); + expect(result).toEqual({ + type: 'cloud.message.failed', + messageId: 'msg', + executionId: undefined, + delivery: 'sent', + error: 'attach_exhausted', + reason: 'execution', + attempts: undefined, + }); + }); + + it('defaults error for internal disconnect reasons when error is missing', () => { + const result = normalize( + createRaw('cloud.message.failed', { + messageId: 'msg', + delivery: 'sent', + reason: 'control_disconnected', + }) ); expect(result).toEqual({ type: 'cloud.message.failed', diff --git a/packages/cloud-agent-sdk/src/normalizer.ts b/packages/cloud-agent-sdk/src/normalizer.ts index 2ead98c57a..bbe9d81d6a 100644 --- a/packages/cloud-agent-sdk/src/normalizer.ts +++ b/packages/cloud-agent-sdk/src/normalizer.ts @@ -583,7 +583,11 @@ function normalizeInnerEvent(eventType: string, data: unknown): NormalizedEvent const reason: 'interrupted' | 'exhausted' | 'execution' = rawReason === 'interrupted' ? 'interrupted' : attempts != null ? 'exhausted' : 'execution'; const error = - r.data.error !== undefined ? extractErrorMessage(r.data.error) : 'Message delivery failed'; + r.data.error !== undefined + ? extractErrorMessage(r.data.error) + : rawReason === 'attach_exhausted' + ? rawReason + : 'Message delivery failed'; return { type: 'cloud.message.failed', messageId, diff --git a/services/cloud-agent-next/src/persistence/SandboxControl.ts b/services/cloud-agent-next/src/persistence/SandboxControl.ts index 6617d9963d..46c55a0655 100644 --- a/services/cloud-agent-next/src/persistence/SandboxControl.ts +++ b/services/cloud-agent-next/src/persistence/SandboxControl.ts @@ -1024,6 +1024,19 @@ export class SandboxControl extends DurableObject { !matchesRoute(route) || !socket ) { + const guard = + physical.state !== 'running' + ? 'physical_not_running' + : physical.stopTombstone + ? 'physical_stopping' + : !runtime + ? 'runtime_not_ready' + : physical.providerRef !== runtime.providerInstanceId + ? 'provider_mismatch' + : !matchesRoute(route) + ? 'route_mismatch' + : 'socket_not_ready'; + logControlDiagnostic('worktree_changes_not_ready', { guard }, 'warn'); return errorResponse(crypto.randomUUID(), 'not_ready', 'Worktree is not attached and ready'); } if ( diff --git a/services/cloud-agent-next/src/sandbox-control/diagnostics.test.ts b/services/cloud-agent-next/src/sandbox-control/diagnostics.test.ts index d225eb5b9d..ae466aac49 100644 --- a/services/cloud-agent-next/src/sandbox-control/diagnostics.test.ts +++ b/services/cloud-agent-next/src/sandbox-control/diagnostics.test.ts @@ -1,5 +1,9 @@ import { afterEach, describe, expect, it, vi } from 'vitest'; -import { CONTROL_DIAGNOSTIC_COALESCE_LIMIT, logControlDiagnostic } from './diagnostics.js'; +import { + CONTROL_DIAGNOSTIC_COALESCE_LIMIT, + diagnosticCause, + logControlDiagnostic, +} from './diagnostics.js'; import { logger } from '../logger.js'; describe('logControlDiagnostic', () => { @@ -68,3 +72,10 @@ describe('logControlDiagnostic', () => { expect(withFields).toHaveBeenCalledTimes(1); }); }); + +describe('diagnosticCause', () => { + it('sanitizes and bounds unknown causes', () => { + expect(diagnosticCause('untrusted cause/value')).toBe('untrusted_cause_value'); + expect(diagnosticCause('x'.repeat(129))).toHaveLength(128); + }); +}); diff --git a/services/cloud-agent-next/src/sandbox-control/diagnostics.ts b/services/cloud-agent-next/src/sandbox-control/diagnostics.ts index 0c0f6d2668..89d2b316c6 100644 --- a/services/cloud-agent-next/src/sandbox-control/diagnostics.ts +++ b/services/cloud-agent-next/src/sandbox-control/diagnostics.ts @@ -69,7 +69,9 @@ export function diagnosticEventType(value: string): string { } export function diagnosticCause(value: string): string { - return CAUSES.has(value) ? value.replaceAll(' ', '_') : 'other'; + return CAUSES.has(value) + ? value.replaceAll(' ', '_') + : value.replace(/[^a-zA-Z0-9_.:-]/g, '_').slice(0, 128); } const DELTA_PROGRESS_EVENTS = new Set([ diff --git a/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts b/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts index 9857dc844c..2e96c99c4a 100644 --- a/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts +++ b/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts @@ -175,6 +175,7 @@ import { matchesSessionMessageReplay, nextQueuedMessageId, recordAcceptedMessageActivity, + releaseCompletedRetryableAttach, releaseUnadmittedWaitingMessages, resolveSessionMessageIntent, streamCloudStatus, @@ -2362,6 +2363,19 @@ export class SandboxSession extends DurableObject { const queued = assigned.messages.find(message => message.messageId === messageId); const deadlineAt = queued?.deliveryDeadlineAt; if (!queued || deadlineAt === undefined) return; + if (Date.now() >= deadlineAt && !queued.operations?.prompt?.dispatched) { + await this.failDelivery( + messageId, + 'preparation_timeout', + queued.wrapperInstanceId, + queued.deliveryRetryScope + ); + return; + } + if (queued.retryNotBefore !== undefined && queued.retryNotBefore > Date.now()) { + await this.armQueueRetry(Math.min(deadlineAt, queued.retryNotBefore)); + return; + } const provider = getSandboxProvider(metadata); const acquisition = provider === 'cloudflare' ? { id: assigned.attemptId, deadlineAt } : undefined; @@ -2480,6 +2494,12 @@ export class SandboxSession extends DurableObject { ?.dispatched === true ) return; + if ( + operation === 'session.attach' && + current.operations?.retiredAttach && + sameSessionOperation(current.operations.retiredAttach.authorization, authorization) + ) + return; throw new Error('Session operation scope changed'); }, defer: pending => this.ctx.waitUntil(pending), @@ -2685,7 +2705,7 @@ export class SandboxSession extends DurableObject { ? { preparation: { attemptId: recorder.attemptId, triggerMessageId: messageId } } : {}), }; - phase = 'attach'; + phase = needsPreparation ? 'preparing' : 'attach'; await wait(() => control.attachSession({ ...(metadata.workspace?.worktreeId @@ -2862,21 +2882,41 @@ export class SandboxSession extends DurableObject { const message = this.queuedMessage(messageId, epoch, wrapperInstanceId); if (!message) return; const rejection = error instanceof ControlRequestError && error.code !== 'runtime_unhealthy'; - const scope = rejection && !message.unresolvedDispatch ? 'message' : 'runtime'; + const completedAttachFailure = + message.operations?.attach?.dispatched === true && + message.operations.attach.result?.ok === false; + const retryableCompletedAttach = + phase !== 'prompt' && rejection && isRetryableDeliveryError(error) && completedAttachFailure; + const scope = + phase === 'attach' + ? retryableCompletedAttach + ? 'message' + : 'runtime' + : rejection && !message.unresolvedDispatch + ? 'message' + : 'runtime'; if (Date.now() >= deadlineAt) { await this.failDelivery(messageId, 'preparation_timeout', wrapperInstanceId, scope); return; } const busy = rejection && error.code === 'session_busy'; + const retryNotBefore = Math.min(deadlineAt, Date.now() + QUEUE_RETRY_MS); + const released = retryableCompletedAttach + ? releaseCompletedRetryableAttach(this.loadMessages(), messageId, retryNotBefore) + : this.loadMessages(); const updated = phase === 'preparing' || busy ? undefined - : incrementDeliveryFailure(this.loadMessages(), messageId, phase); - const messages = (updated?.messages ?? this.loadMessages()).map( + : incrementDeliveryFailure(released, messageId, phase); + const messages = (updated?.messages ?? released).map( (message): MessageRecord => message.messageId === messageId ? { ...message, deliveryRetryScope: scope } : message ); if (!this.saveMessages(messages, epoch)) return; + if (retryableCompletedAttach && !updated?.exhausted) { + await this.armQueueRetry(retryNotBefore); + return; + } if (isRetryableDeliveryError(error) && !updated?.exhausted) { await this.armQueueRetry(Math.min(deadlineAt, Date.now() + QUEUE_RETRY_MS)); return; diff --git a/services/cloud-agent-next/src/sandbox-session/control-dispatch.test.ts b/services/cloud-agent-next/src/sandbox-session/control-dispatch.test.ts index a0a5eaf837..793eb1e958 100644 --- a/services/cloud-agent-next/src/sandbox-session/control-dispatch.test.ts +++ b/services/cloud-agent-next/src/sandbox-session/control-dispatch.test.ts @@ -74,10 +74,10 @@ describe('controlRequestResult', () => { describe('deliveryErrorLogFields', () => { it.each(['session_busy', 'not_ready', 'runtime_unhealthy'])( - 'logs only the allowlisted %s code and retry classification', + 'logs the public message with the allowlisted %s code and retry classification', code => { const error = Object.assign( - new ControlRequestError({ code, message: 'sensitive-message', retryable: true }), + new ControlRequestError({ code, message: 'Public control error', retryable: true }), { cause: 'sensitive-cause', stack: 'sensitive-stack', @@ -85,7 +85,11 @@ describe('deliveryErrorLogFields', () => { env: 'sensitive-env', } ); - expect(deliveryErrorLogFields(error)).toEqual({ errorCode: code, retryable: true }); + expect(deliveryErrorLogFields(error)).toEqual({ + errorCode: code, + errorMessage: 'Public control error', + retryable: true, + }); } ); @@ -94,11 +98,15 @@ describe('deliveryErrorLogFields', () => { deliveryErrorLogFields( new ControlRequestError({ code: 'sensitive-untrusted-code', - message: 'sensitive-message', + message: 'Public control error', retryable: false, }) ) - ).toEqual({ errorCode: 'unknown_control_error', retryable: false }); + ).toEqual({ + errorCode: 'unknown_control_error', + errorMessage: 'Public control error', + retryable: false, + }); }); it.each([false, true])( diff --git a/services/cloud-agent-next/src/sandbox-session/control-dispatch.ts b/services/cloud-agent-next/src/sandbox-session/control-dispatch.ts index e6796eea48..88cb1cda28 100644 --- a/services/cloud-agent-next/src/sandbox-session/control-dispatch.ts +++ b/services/cloud-agent-next/src/sandbox-session/control-dispatch.ts @@ -72,6 +72,7 @@ export function deliveryErrorLogFields(error: unknown) { error instanceof ControlRequestError ? (controlErrorCodes.find(code => code === error.code) ?? 'unknown_control_error') : 'transport_or_internal_error', + ...(error instanceof ControlRequestError ? { errorMessage: error.message } : {}), retryable: isRetryableDeliveryError(error), }; } diff --git a/services/cloud-agent-next/src/sandbox-session/recovery/failed-snapshot-uses-terminal-failure.test.ts b/services/cloud-agent-next/src/sandbox-session/recovery/failed-snapshot-uses-terminal-failure.test.ts index 35b87866c1..9d96c1ff36 100644 --- a/services/cloud-agent-next/src/sandbox-session/recovery/failed-snapshot-uses-terminal-failure.test.ts +++ b/services/cloud-agent-next/src/sandbox-session/recovery/failed-snapshot-uses-terminal-failure.test.ts @@ -22,6 +22,7 @@ describe('failed snapshot', () => { delivery: 'queued', accepted: false, reason: 'environment_failed', + error: 'environment_failed', timestamp: 99, }, }, @@ -35,6 +36,7 @@ describe('failed snapshot', () => { delivery: 'sent', accepted: true, reason: 'environment_failed', + error: 'environment_failed', timestamp: 20, }, }, 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 440bb27663..775157f538 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 @@ -40,6 +40,7 @@ import type { } from '../execution/types.js'; import { acceptQueuedMessage, + applySessionOperationResult, applyMessageOutcome, assignPreparationAttemptId, createSessionMessageRecord, @@ -50,6 +51,7 @@ import { matchesSessionMessageReplay, nextQueuedMessageId, recordAcceptedMessageActivity, + releaseCompletedRetryableAttach, resolveSessionMessageIntent, streamCloudStatus, streamQueuedSnapshots, @@ -570,6 +572,77 @@ describe('assignPreparationAttemptId', () => { }); }); +describe('releaseCompletedRetryableAttach', () => { + const authorization: SessionOperationAuthorization = { + operation: 'session.attach', + operationId: 'attach-a', + messageId: 'a', + session: { sessionId: SESSION_ID, kiloSessionId: 'kilo_root', directory: DIRECTORY }, + wrapperInstanceId: RUNTIME_ID, + dispatchDeadlineAt: 100, + }; + const failedResult = { + ok: false as const, + error: { code: 'not_ready', message: 'Wrapper is not ready', retryable: true }, + }; + + it('retires the failed attach proof before retrying', () => { + const attach = { authorization, dispatched: true, result: failedResult }; + const messages: SessionMessageRecord[] = [ + { + ...createSessionMessageRecord({ turn: promptTurn, agent: defaultAgent }), + unresolvedDispatch: true, + preparationAttemptId: authorization.operationId, + operations: { attach }, + }, + ]; + + const released = releaseCompletedRetryableAttach(messages, 'a', 200); + + expect(released[0]).toMatchObject({ + unresolvedDispatch: undefined, + preparationAttemptId: undefined, + retryNotBefore: 200, + operations: { retiredAttach: attach }, + }); + expect(released[0]?.operations?.attach).toBeUndefined(); + }); + + it('applies a duplicate result to a retired attach proof without restoring it', () => { + const completedAt = 200; + const delivery: SessionOperationDelivery = { + version: 2, + authorization, + completedAt, + result: failedResult, + events: [], + preparing: [], + }; + const messages: SessionMessageRecord[] = [ + { + ...createSessionMessageRecord({ turn: promptTurn, agent: defaultAgent }), + wrapperInstanceId: RUNTIME_ID, + operations: { + retiredAttach: { + authorization, + dispatched: true, + result: failedResult, + resultHash: 'result-hash', + completedAt, + decision: { state: 'queued', at: completedAt }, + }, + }, + }, + ]; + + const applied = applySessionOperationResult(messages, delivery, 'result-hash', completedAt + 1); + + expect(applied).toMatchObject({ disposition: 'identical' }); + expect(applied?.messages[0]?.operations?.attach).toBeUndefined(); + expect(applied?.messages[0]?.operations?.retiredAttach).toMatchObject({ authorization }); + }); +}); + describe('applyMessageOutcome', () => { it('settles only the message identified by the matching runtime', () => { const before = [{ ...msg('a', 'accepted'), wrapperInstanceId: 'runtime' }, msg('b', 'queued')]; @@ -750,6 +823,7 @@ describe('streamQueuedSnapshots', () => { delivery: 'sent', accepted: true, reason: 'prompt_exhausted', + error: 'prompt_exhausted', timestamp: 20, }, }, @@ -884,6 +958,7 @@ describe('streamQueuedSnapshots', () => { delivery: 'queued', accepted: false, reason: 'preparation_failed', + error: 'preparation_failed', timestamp: 99, }, }, @@ -897,6 +972,7 @@ describe('streamQueuedSnapshots', () => { delivery: 'sent', accepted: true, reason: 'wrapper_failed', + error: 'wrapper_failed', timestamp: 20, }, }, @@ -1813,6 +1889,202 @@ describe('SandboxSession orchestration', () => { }); }); + it('retries a completed retryable attach with a new authorization before delivering the prompt', async () => { + const fixture = sessionFixture(); + fixture.setStatus({ + physical: 'running', + connection: 'ready', + wrapperInstanceId: RUNTIME_ID, + operationResults: true, + }); + const original = fixture.control.request.getMockImplementation(); + if (!original) throw new Error('Missing control fixture'); + let attachAttempts = 0; + delegateRequest(fixture, 'session.attach', async input => + ++attachAttempts === 1 ? controlFailure(true, 'not_ready') : original(input) + ); + + await fixture.admit('a'); + await fixture.flush(); + const firstAuthorization = fixture.record('a')?.operations?.attach?.authorization; + if (!firstAuthorization) throw new Error('Missing first attach authorization'); + + const delivery: SessionOperationDelivery = { + version: 2, + authorization: firstAuthorization, + completedAt: Date.now(), + result: { + ok: false, + error: { code: 'not_ready', message: 'Wrapper is not ready', retryable: true }, + }, + events: [], + preparing: [], + }; + await fixture.session.receiveSandboxOperationResult({ + session: firstAuthorization.session, + wrapperInstanceId: RUNTIME_ID, + delivery, + }); + await fixture.flush(); + expect(fixture.record('a')).toMatchObject({ + state: 'queued', + preparationAttemptId: firstAuthorization.operationId, + operations: { attach: { dispatched: true, result: delivery.result } }, + }); + delegateRequest(fixture, 'session.operation.get', async () => + controlResponse({ state: 'completed', delivery }) + ); + + const now = Date.now(); + await fixture.fireAlarm(); + await fixture.flush(); + expect( + fixture.control.request.mock.calls.filter(([input]) => input.operation === 'session.attach') + ).toHaveLength(1); + expect( + fixture.control.request.mock.calls.filter( + ([input]) => input.operation === 'session.operation.get' + ) + ).toHaveLength(1); + const retryNotBefore = fixture.record('a')?.retryNotBefore; + if (retryNotBefore === undefined) throw new Error('Missing attach retry time'); + expect(retryNotBefore - now).toBe(5_000); + + vi.setSystemTime(retryNotBefore); + await fixture.fireAlarm(); + await fixture.flush(); + const attachRequests = fixture.control.request.mock.calls.filter( + ([input]) => input.operation === 'session.attach' + ); + expect(attachRequests).toHaveLength(2); + const secondAuthorization = attachRequests[1]?.[0].authorization; + expect(secondAuthorization?.operationId).not.toBe(firstAuthorization.operationId); + expect(fixture.record('a')?.preparationAttemptId).toBe(secondAuthorization?.operationId); + expect( + fixture.control.request.mock.calls.filter(([input]) => input.operation === 'session.prompt') + ).toHaveLength(1); + }); + + it('reconstructs a released attach retry after its retry time has passed', async () => { + const fixture = sessionFixture(); + fixture.setStatus({ + physical: 'running', + connection: 'ready', + wrapperInstanceId: RUNTIME_ID, + operationResults: true, + }); + fixture.storage.kv.put('session_messages', [ + { + ...createSessionMessageRecord({ + turn: { type: 'prompt', messageId: 'recovered', prompt: 'continue delivery' }, + agent: defaultAgent, + }), + deliveryDeadlineAt: Date.now() + SESSION_DELIVERY_TIMEOUT_MS, + retryNotBefore: Date.now() - 1, + wrapperInstanceId: RUNTIME_ID, + }, + ]); + + await fixture.fireAlarm(); + await fixture.flush(); + + const attach = fixture.control.request.mock.calls.find( + ([input]) => input.operation === 'session.attach' + )?.[0]; + expect(attach?.authorization?.operationId).toBe( + fixture.record('recovered')?.preparationAttemptId + ); + expect( + fixture.control.request.mock.calls + .map(([input]) => input.operation) + .filter(operation => operation === 'session.attach' || operation === 'session.prompt') + ).toEqual(['session.attach', 'session.prompt']); + }); + + it.each([ + { + error: { code: 'not_ready', message: 'Wrapper is not ready', retryable: true }, + quarantinesRuntime: false, + }, + { + error: { code: 'runtime_unhealthy', message: 'Wrapper is unhealthy', retryable: false }, + quarantinesRuntime: true, + }, + ] as const)( + 'counts exhausted warm attach failures and uses the required retry scope: $error.code', + async ({ error, quarantinesRuntime }) => { + const fixture = sessionFixture(); + fixture.setStatus({ + physical: 'running', + connection: 'ready', + wrapperInstanceId: RUNTIME_ID, + operationResults: true, + }); + await fixture.admit('cold'); + await fixture.flush(); + await fixture.outcome('cold', 'completed'); + fixture.reload(); + fixture.control.ensureReady.mockResolvedValue({ + physical: 'running', + connection: 'ready', + wrapperInstanceId: RUNTIME_ID, + operationResults: true, + attachment: { + ...ATTACHMENT, + kilo: { ...ATTACHMENT.kilo, containmentEnabled: false }, + }, + }); + const authorization: SessionOperationAuthorization = { + operation: 'session.attach', + operationId: 'warm-attach', + messageId: 'warm', + session: { sessionId: SESSION_ID, kiloSessionId: 'kilo_root', directory: DIRECTORY }, + wrapperInstanceId: RUNTIME_ID, + dispatchDeadlineAt: Date.now() + SESSION_DELIVERY_TIMEOUT_MS, + }; + const delivery: SessionOperationDelivery = { + version: 2, + authorization, + completedAt: Date.now(), + result: { ok: false, error }, + events: [], + preparing: [], + }; + fixture.storage.kv.put('session_messages', [ + { + ...createSessionMessageRecord({ + turn: { type: 'prompt', messageId: 'warm', prompt: 'warm retry' }, + agent: defaultAgent, + }), + wrapperInstanceId: RUNTIME_ID, + deliveryDeadlineAt: authorization.dispatchDeadlineAt, + preparationAttemptId: authorization.operationId, + attachFailures: 1, + operations: { attach: { authorization, dispatched: true, result: delivery.result } }, + }, + ]); + delegateRequest(fixture, 'session.operation.get', async () => + controlResponse({ state: 'completed', delivery }) + ); + + await fixture.fireAlarm(); + await fixture.flush(); + + expect(fixture.record('warm')).toMatchObject({ + state: 'failed', + attachFailures: 2, + failedReason: 'attach_exhausted', + }); + if (quarantinesRuntime) { + expect(fixture.control.quarantineRuntime).toHaveBeenCalledWith( + expect.objectContaining({ wrapperInstanceId: RUNTIME_ID, reason: 'attach_exhausted' }) + ); + } else { + expect(fixture.control.quarantineRuntime).not.toHaveBeenCalled(); + } + } + ); + it.each(['running', 'completed'] as const)( 'reconstructs a late accepted prompt from its original %s operation result without redispatch', async state => { @@ -3512,6 +3784,12 @@ describe('SandboxSession orchestration', () => { if (source === 'exception') first.reject(transient); else first.resolve(controlFailure(true)); await fixture.flush(); + if (operation === 'session.attach') { + expect(fixture.record('a')?.state).toBe('queued'); + expect(fixture.record('a')?.attachFailures).toBeUndefined(); + expect(fixture.terminalEvents()).toHaveLength(0); + return; + } for (let attempt = 2; attempt <= limit; attempt++) { expect(fixture.record('a')).toMatchObject({ state: 'queued', @@ -3585,7 +3863,8 @@ describe('SandboxSession orchestration', () => { ); } await fixture.flush(); - expect(fixture.record('a')).toMatchObject({ state: 'failed', failedReason: reason }); + const expectedReason = operation === 'session.attach' ? 'environment_failed' : reason; + expect(fixture.record('a')).toMatchObject({ state: 'failed', failedReason: expectedReason }); if (failure === 'permanent response') { expect(fixture.record('b')?.state).toBe('queued'); expect(fixture.terminalEvents()).toHaveLength(1); @@ -3595,7 +3874,7 @@ describe('SandboxSession orchestration', () => { expect(fixture.record('b')?.state).toBe('queued'); expect(fixture.terminalEvents()).toHaveLength(1); expect(fixture.control.quarantineRuntime).toHaveBeenCalledWith( - expect.objectContaining({ wrapperInstanceId: RUNTIME_ID, reason }) + expect.objectContaining({ wrapperInstanceId: RUNTIME_ID, reason: expectedReason }) ); await fixture.fireAlarm(); expect( @@ -3637,7 +3916,10 @@ describe('SandboxSession orchestration', () => { expect(fixture.terminalEvents()).toHaveLength(0); return; } - expect(fixture.record('a')).toMatchObject({ state: 'failed', failedReason: reason }); + expect(fixture.record('a')).toMatchObject({ + state: 'failed', + failedReason: operation === 'session.attach' ? 'environment_failed' : reason, + }); expect(fixture.record('b')?.state).toBe('queued'); expect(fixture.terminalEvents()).toHaveLength(1); expect(fixture.control.quarantineRuntime).toHaveBeenCalledOnce(); @@ -3812,11 +4094,18 @@ describe('SandboxSession orchestration', () => { sibling.reload(); await sibling.fireAlarm(); } - expect(sibling.record('rejected')).toMatchObject({ - state: 'failed', - failedReason: operation === 'session.attach' ? 'attach_exhausted' : 'prompt_exhausted', - }); - expect(sibling.terminalEvents()).toHaveLength(1); + if (operation === 'session.attach' && retryable) { + expect(sibling.record('rejected')).toMatchObject({ state: 'queued' }); + } else { + expect(sibling.record('rejected')).toMatchObject({ + state: 'failed', + failedReason: + operation === 'session.attach' ? 'environment_failed' : 'prompt_exhausted', + }); + } + expect(sibling.terminalEvents()).toHaveLength( + operation === 'session.attach' && retryable ? 0 : 1 + ); expect(await sibling.session.isSandboxCleanupScheduled()).toBe(false); expect(writer.control.quarantineRuntime).not.toHaveBeenCalled(); await writer.outcome('writer', 'completed'); @@ -5495,7 +5784,9 @@ describe('SandboxSession orchestration', () => { expect(fixture.eventQueries.findByEntityPrefix('preparation/attempt/')).toEqual( coldPreparation ); - expect(fixture.control.quarantineRuntime).not.toHaveBeenCalled(); + expect(fixture.control.quarantineRuntime).toHaveBeenCalledWith( + expect.objectContaining({ wrapperInstanceId: RUNTIME_ID, reason: 'attach_exhausted' }) + ); expect(await fixture.snapshot()).toMatchObject({ preparationSnapshots: coldPreparation }); }); 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 5b7ebea952..2baa6e3e68 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 @@ -63,10 +63,12 @@ type SessionMessageLifecycle = { attachFailures?: number; promptFailures?: number; preparationAttemptId?: string; + retryNotBefore?: number; executionDeadlineAt?: number; cancellation?: { operationId: string; deadlineAt: number }; operations?: { attach?: SessionOperationProof; + retiredAttach?: SessionOperationProof; prompt?: SessionOperationProof; }; }; @@ -306,6 +308,7 @@ export function releaseUnadmittedWaitingMessages( ...message, wrapperInstanceId: undefined, preparationAttemptId: undefined, + retryNotBefore: undefined, deliveryDeadlineAt: undefined, operations: undefined, }; @@ -314,6 +317,27 @@ export function releaseUnadmittedWaitingMessages( }; } +export function releaseCompletedRetryableAttach( + messages: readonly SessionMessageRecord[], + messageId: string, + retryNotBefore: number +): SessionMessageRecord[] { + return messages.map(message => { + const attach = message.messageId === messageId ? message.operations?.attach : undefined; + if (!attach?.dispatched || attach.result?.ok !== false) return message; + const operations = { ...message.operations }; + operations.retiredAttach = attach; + delete operations.attach; + return { + ...message, + unresolvedDispatch: undefined, + preparationAttemptId: undefined, + retryNotBefore, + ...(Object.keys(operations).length > 0 ? { operations } : { operations: undefined }), + }; + }); +} + export function incrementDeliveryFailure( messages: readonly SessionMessageRecord[], messageId: string, @@ -467,7 +491,11 @@ export function failedMessageSnapshot( delivery: accepted ? 'sent' : 'queued', accepted, reason: cancelled ? 'interrupted' : message.failedReason, - ...(cancelled ? { error: 'The message was interrupted' } : {}), + ...(cancelled + ? { error: 'The message was interrupted' } + : message.failedReason + ? { error: message.failedReason } + : {}), timestamp: message.acceptedAt ?? now, }; } @@ -518,8 +546,19 @@ export function applySessionOperationResult( const authorization = delivery.authorization; const message = messages.find(item => item.messageId === authorization.messageId); const kind = authorization.operation === 'session.attach' ? 'attach' : 'prompt'; - const proof = message?.operations?.[kind]; - const storedAuthorization = sessionOperationAuthorizationSchema.safeParse(proof?.authorization); + let proof = message?.operations?.[kind]; + let proofSlot: 'attach' | 'retiredAttach' | 'prompt' = kind; + let storedAuthorization = sessionOperationAuthorizationSchema.safeParse(proof?.authorization); + if ( + kind === 'attach' && + (!proof?.dispatched || + !storedAuthorization.success || + !sameSessionOperation(storedAuthorization.data, authorization)) + ) { + proof = message?.operations?.retiredAttach; + proofSlot = 'retiredAttach'; + storedAuthorization = sessionOperationAuthorizationSchema.safeParse(proof?.authorization); + } if ( !message || !proof?.dispatched || @@ -574,7 +613,7 @@ export function applySessionOperationResult( ...item, operations: { ...item.operations, - [kind]: { + [proofSlot]: { ...proof, result: delivery.result, resultHash, diff --git a/services/cloud-agent-next/test/integration/sandbox-control.test.ts b/services/cloud-agent-next/test/integration/sandbox-control.test.ts index 3b5340dc6b..d1c4cf81c1 100644 --- a/services/cloud-agent-next/test/integration/sandbox-control.test.ts +++ b/services/cloud-agent-next/test/integration/sandbox-control.test.ts @@ -7717,7 +7717,7 @@ describe('SandboxSession worktree changes persistence', () => { fixture.session.getMessageResult('msg_failed_reattach') ).resolves.toMatchObject({ type: 'found', - result: { status: retry === 'cancelled' ? 'interrupted' : 'failed' }, + result: { status: retry === 'cancelled' ? 'interrupted' : 'queued' }, }); await expect(fixture.session.getWorktreeChanges()).resolves.toEqual(beforeCleanup); await expect(fixture.control.getStatus()).resolves.toMatchObject({ 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 7954d5d763..a93d75e23d 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 @@ -21,6 +21,42 @@ const session = { const payload = { type: 'session.idle', properties: {} }; describe('native-scoped control event failures', () => { + it.each([ + ['session.preparing', false], + ['session.event', true], + ] as const)( + 'retires the runtime after an expired %s only when required', + async (event, retires) => { + const clock = spyOn(Date, 'now').mockReturnValue(1_000); + const runtime = { runtimeId: crypto.randomUUID() }; + const retired = mock(); + const handleFailure = createControlEventFailureHandler({ + getRuntime: () => runtime, + onFailure: retired, + }); + const reported = mock((failure: ControlEventOutboxFailure) => handleFailure(failure)); + const transport = createControlEventTransport({ + supportsReceipts: () => true, + prepare: input => input, + publish: async () => {}, + sendLegacy: () => false, + onFailure: reported, + }); + try { + expect( + transport.enqueue(event, payload, { ...session, nativeRuntimeId: runtime.runtimeId }) + ).toBe(true); + clock.mockReturnValue(31_000); + expect(await transport.resume()).toBe(true); + expect(reported).toHaveBeenCalledTimes(1); + expect(retired).toHaveBeenCalledTimes(retires ? 1 : 0); + } finally { + transport.close(); + clock.mockRestore(); + } + } + ); + it.each(['expired', 'rejected'] as const)( 'reports an immutable N1 %s publication without retiring or blocking N2', async reason => { 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 007e4a97cb..7ebbe3f199 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 @@ -14,6 +14,7 @@ export function createControlEventFailureHandler(); return (failure?: ControlEventOutboxFailure): void => { if (!failure) return; + if (failure.publication.event === 'session.preparing') return; const { directory, nativeRuntimeId } = failure.publication.session; if (!nativeRuntimeId) return; const runtime = options.getRuntime(directory); 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 29adfe1df6..6c98d37062 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 @@ -450,6 +450,16 @@ describe('handleControlRequest', () => { expect(result).toEqual({ ok: true, result: { attached: true } }); }); + it('attaches while Kilo is not ready', async () => { + const result = await handleControlRequest( + 'session.attach', + session, + { kilo }, + deps({ kiloReady: false }) + ); + expect(result).toEqual({ ok: true, result: { attached: true } }); + }); + it('registers terminal eligibility only after successful session attachment', async () => { const attached: unknown[] = []; const terminalRuntime = fakeTerminalRuntime({ @@ -666,7 +676,7 @@ describe('handleControlRequest', () => { expect(answered).toEqual([{ permissionId: 'perm_1', response: 'once' }]); }); - it('fences new work during feed recovery while preserving pending input replies and Stop', async () => { + it('fences prompts during feed recovery while allowing attachments and pending input replies', async () => { const kiloClient = fakeKilo({ getPermissions: async () => [ { @@ -695,6 +705,10 @@ describe('handleControlRequest', () => { admission: 'not-admitted', }, }); + expect(await handleControlRequest('session.attach', session, { kilo }, handlerDeps)).toEqual({ + ok: true, + result: { attached: true }, + }); expect( await handleControlRequest( 'session.permission.resolve', @@ -3380,11 +3394,6 @@ describe('control cancellation and attachments', () => { expect(signals).toHaveLength(2); expect(signals.every(signal => signal.aborted)).toBe(true); expect(terminalStopped).toBe(true); - expect( - await handleControlRequest('session.attach', session, { kilo }, handlerDeps) - ).toMatchObject({ - ok: false, - }); running.resolve(completion()); }); diff --git a/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.ts b/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.ts index 2d2e956271..eaea32d624 100644 --- a/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.ts +++ b/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.ts @@ -496,15 +496,14 @@ export async function handleControlRequest( const admission = deps.operations.admission(operation, session, payload, authorization); if (admission.kind === 'reply') return admission.result; if ( - (operation === 'session.attach' || - operation === 'session.prompt' || - operation === 'session.terminal.create') && + (operation === 'session.prompt' || operation === 'session.terminal.create') && deps.kiloRuntimes?.prepareForNewWork?.(session.directory) === false ) { return rejectBeforeAdmission('not_ready', 'Native feed recovery is in progress', true); } if ( - (deps.signal?.aborted || (!deps.kiloReady && operation !== 'session.git.summary')) && + (deps.signal?.aborted || + (!deps.kiloReady && operation !== 'session.attach' && operation !== 'session.git.summary')) && operation !== 'session.abort' && operation !== 'session.detach' ) { 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 a06e50786c..a520493903 100644 --- a/services/cloud-agent-next/wrapper/src/control/session-operation.ts +++ b/services/cloud-agent-next/wrapper/src/control/session-operation.ts @@ -431,8 +431,6 @@ export class SessionOperation { work: Extract ): Promise { this.assertCurrent(); - if (this.deps.prepareForNewWork?.() === false) - return fail('Native feed recovery is in progress', true); const result = await work.apply(this.session, work.payload, { signal: this.signal, assertCurrent: () => this.assertCurrent(),