diff --git a/packages/worker-utils/src/do-retry-scope.test.ts b/packages/worker-utils/src/do-retry-scope.test.ts new file mode 100644 index 0000000000..0d41bc6e65 --- /dev/null +++ b/packages/worker-utils/src/do-retry-scope.test.ts @@ -0,0 +1,29 @@ +import { afterEach, describe, expect, it, vi } from 'vitest'; +import { DEFAULT_DO_RETRY_CONFIG, withDORetry } from './do-retry.js'; + +afterEach(() => vi.useRealTimers()); + +describe('withDORetry scopes', () => { + it('does not start another attempt after its deadline expires', async () => { + vi.useFakeTimers(); + const operation = vi.fn(() => new Promise(() => undefined)); + const pending = withDORetry( + () => ({}), + operation, + 'scoped_operation', + { + ...DEFAULT_DO_RETRY_CONFIG, + scope: { deadlineAt: Date.now() + 100 }, + }, + { warn: () => undefined, error: () => undefined } + ); + const outcome = pending.then( + () => undefined, + error => error + ); + + await vi.advanceTimersByTimeAsync(100); + await expect(outcome).resolves.toMatchObject({ name: 'TimeoutError' }); + expect(operation).toHaveBeenCalledOnce(); + }); +}); diff --git a/packages/worker-utils/src/do-retry.ts b/packages/worker-utils/src/do-retry.ts index 7353e202c8..afde8e897f 100644 --- a/packages/worker-utils/src/do-retry.ts +++ b/packages/worker-utils/src/do-retry.ts @@ -1,11 +1,20 @@ // Cloudflare Workers provides scheduler.wait() for cooperative delays. // Not in standard webworker lib types. -declare const scheduler: undefined | { wait(ms: number): Promise }; +declare const scheduler: + | undefined + | { wait(ms: number, options?: { signal?: AbortSignal }): Promise }; + +export type DORetryScope = { + deadlineAt: number; + signal?: AbortSignal; + assertCurrent?: () => void; +}; export type DORetryConfig = { maxAttempts: number; baseBackoffMs: number; maxBackoffMs: number; + scope?: DORetryScope; }; export const DEFAULT_DO_RETRY_CONFIG: DORetryConfig = { @@ -44,11 +53,76 @@ function calculateBackoff(attempt: number, config: DORetryConfig): number { return Math.min(config.maxBackoffMs, jitteredBackoff); } -function waitMs(ms: number): Promise { +function waitMs(ms: number, signal?: AbortSignal): Promise { + signal?.throwIfAborted(); if (typeof scheduler !== 'undefined' && 'wait' in scheduler) { - return scheduler.wait(ms); + return signal ? scheduler.wait(ms, { signal }) : scheduler.wait(ms); + } + if (!signal) return new Promise(resolve => setTimeout(resolve, ms)); + return new Promise((resolve, reject) => { + const onAbort = () => { + clearTimeout(timeoutId); + signal.removeEventListener('abort', onAbort); + reject(signal.reason); + }; + const timeoutId = setTimeout(() => { + signal.removeEventListener('abort', onAbort); + resolve(); + }, ms); + signal.addEventListener('abort', onAbort, { once: true }); + if (signal.aborted) onAbort(); + }); +} + +function createRetryScope({ deadlineAt, signal, assertCurrent }: DORetryScope) { + if (!Number.isFinite(deadlineAt)) + throw new RangeError('Durable Object retry deadlineAt must be finite'); + + const controller = new AbortController(); + const deadlineError = new DOMException('Durable Object retry deadline exceeded', 'TimeoutError'); + const onAbort = () => controller.abort(signal?.reason); + signal?.addEventListener('abort', onAbort, { once: true }); + if (signal?.aborted) onAbort(); + const timeoutId = setTimeout( + () => controller.abort(deadlineError), + Math.max(0, deadlineAt - Date.now()) + ); + + return { + deadlineAt, + signal: controller.signal, + assertActive() { + if (Date.now() >= deadlineAt) controller.abort(deadlineError); + controller.signal.throwIfAborted(); + try { + assertCurrent?.(); + } catch (error) { + controller.abort(error); + throw error; + } + if (Date.now() >= deadlineAt) controller.abort(deadlineError); + controller.signal.throwIfAborted(); + }, + dispose() { + clearTimeout(timeoutId); + signal?.removeEventListener('abort', onAbort); + }, + }; +} + +async function waitWithSignal(pending: Promise, signal: AbortSignal): Promise { + let onAbort: (() => void) | undefined; + const cancelled = new Promise((_, reject) => { + onAbort = () => reject(signal.reason); + signal.addEventListener('abort', onAbort, { once: true }); + if (signal.aborted) onAbort(); + }); + + try { + return await Promise.race([pending, cancelled]); + } finally { + if (onAbort) signal.removeEventListener('abort', onAbort); } - return new Promise(resolve => setTimeout(resolve, ms)); } type DORetryLogger = { @@ -87,48 +161,68 @@ export async function withDORetry( logger: DORetryLogger = console ): Promise { let lastError: Error | undefined; + const scope = config.scope ? createRetryScope(config.scope) : undefined; - for (let attempt = 0; attempt < config.maxAttempts; attempt++) { - try { - // Create fresh stub for each attempt - const stub = getStub(); - return await operation(stub); - } catch (error) { - lastError = error instanceof Error ? error : new Error(String(error)); - - // Check if we should retry - if (!isRetryableError(error)) { - logger.warn('[do-retry] Non-retryable error', { + try { + for (let attempt = 0; attempt < config.maxAttempts; attempt++) { + scope?.assertActive(); + try { + // Create fresh stub for each attempt + const stub = getStub(); + scope?.assertActive(); + const result = scope + ? await waitWithSignal(operation(stub), scope.signal) + : await operation(stub); + scope?.assertActive(); + return result; + } catch (error) { + scope?.assertActive(); + lastError = error instanceof Error ? error : new Error(String(error)); + + // Check if we should retry + if (!isRetryableError(error)) { + logger.warn('[do-retry] Non-retryable error', { + operation: operationName, + attempt: attempt + 1, + error: lastError.message, + retryable: false, + }); + throw lastError; + } + + // Check if we have retries left + if (attempt + 1 >= config.maxAttempts) { + logger.error('[do-retry] All retry attempts exhausted', { + operation: operationName, + attempts: attempt + 1, + error: lastError.message, + }); + throw lastError; + } + + // Calculate backoff and wait + const requestedBackoffMs = calculateBackoff(attempt, config); + const backoffMs = scope + ? Math.min(requestedBackoffMs, Math.max(0, scope.deadlineAt - Date.now())) + : requestedBackoffMs; + logger.warn('[do-retry] Retrying', { operation: operationName, attempt: attempt + 1, + backoffMs: Math.round(backoffMs), error: lastError.message, - retryable: false, }); - throw lastError; - } - // Check if we have retries left - if (attempt + 1 >= config.maxAttempts) { - logger.error('[do-retry] All retry attempts exhausted', { - operation: operationName, - attempts: attempt + 1, - error: lastError.message, - }); - throw lastError; + scope?.assertActive(); + try { + await waitMs(backoffMs, scope?.signal); + } finally { + scope?.assertActive(); + } } - - // Calculate backoff and wait - const backoffMs = calculateBackoff(attempt, config); - logger.warn('[do-retry] Retrying', { - operation: operationName, - attempt: attempt + 1, - backoffMs: Math.round(backoffMs), - error: lastError.message, - }); - - await waitMs(backoffMs); } - } - throw lastError ?? new Error('Unexpected retry loop exit'); + throw lastError ?? new Error('Unexpected retry loop exit'); + } finally { + scope?.dispose(); + } } diff --git a/packages/worker-utils/src/index.ts b/packages/worker-utils/src/index.ts index d7e263447f..7c02caa312 100644 --- a/packages/worker-utils/src/index.ts +++ b/packages/worker-utils/src/index.ts @@ -1,7 +1,7 @@ export { getCachedSecret, clearSecretCacheForTest } from './cached-secret.js'; export { withDORetry, DEFAULT_DO_RETRY_CONFIG } from './do-retry.js'; -export type { DORetryConfig } from './do-retry.js'; +export type { DORetryConfig, DORetryScope } from './do-retry.js'; export { backendAuthMiddleware } from './backend-auth-middleware.js'; diff --git a/services/cloud-agent-next/src/persistence/SandboxControl.ts b/services/cloud-agent-next/src/persistence/SandboxControl.ts index 704fb9bbc4..2b72c9f1d6 100644 --- a/services/cloud-agent-next/src/persistence/SandboxControl.ts +++ b/services/cloud-agent-next/src/persistence/SandboxControl.ts @@ -17,7 +17,7 @@ import { type SandboxWorktreeCleanupInput, } from '../sandbox-control/worktree-deletion.js'; import { getSandbox } from '@cloudflare/sandbox'; -import { withTimeout } from '@kilocode/worker-utils'; +import { DEFAULT_DO_RETRY_CONFIG, withTimeout } from '@kilocode/worker-utils'; import { z } from 'zod'; import type { Env } from '../types.js'; import { resolveSecret } from '../auth.js'; @@ -297,6 +297,10 @@ export type SandboxControlStatus = { operationResults?: true; }; +export type RuntimeQuarantineResult = + | { quarantined: true; disposition: 'native_retired' | 'native_pending' | 'physical_stopping' } + | { quarantined: false; disposition: 'physical_stopped' | 'wrapper_replaced' | 'unconfirmed' }; + export class SandboxControl extends DurableObject { readonly sandboxId: string; private socketHandler: SandboxControlSocketHandler; @@ -941,8 +945,18 @@ export class SandboxControl extends DurableObject { sessionId: string; wrapperInstanceId: string; reason: string; - }): Promise<{ quarantined: boolean }> { + nativeRuntimeId?: string; + authorization?: SessionOperationAuthorization; + }): Promise { await this.ensureOperationalInitialized(); + const nativeRuntimeId = + input.nativeRuntimeId === undefined + ? undefined + : z.string().uuid().safeParse(input.nativeRuntimeId); + const authorization = + input.authorization === undefined + ? undefined + : sessionOperationAuthorizationSchema.safeParse(input.authorization); if ( typeof input.ownerId !== 'string' || !input.ownerId || @@ -952,7 +966,14 @@ export class SandboxControl extends DurableObject { !input.wrapperInstanceId || typeof input.reason !== 'string' || !input.reason || - input.reason.length > 256 + input.reason.length > 256 || + (nativeRuntimeId !== undefined && + (!nativeRuntimeId.success || + !authorization?.success || + authorization.data.operation !== 'session.attach' || + authorization.data.session.sessionId !== input.sessionId || + authorization.data.wrapperInstanceId !== input.wrapperInstanceId)) || + (nativeRuntimeId === undefined && authorization !== undefined) ) { throw new Error('Invalid sandbox quarantine request'); } @@ -961,28 +982,65 @@ export class SandboxControl extends DurableObject { loadPhysicalRecord(this.ctx.storage), ]); if (ownerId !== input.ownerId) throw new Error('Sandbox owner mismatch'); - if (physical.state === 'stopped') return { quarantined: false }; - const connection = this.activeConnection; + if (nativeRuntimeId?.success && authorization?.success) { + const retired = (await loadNativeRuntimeRetirements(this.ctx.storage)).some( + receipt => + receipt.state === 'completed' && + receipt.disposition === 'retired' && + receipt.nativeRuntimeId === nativeRuntimeId.data && + receipt.connection.wrapperInstanceId === input.wrapperInstanceId && + receipt.recipients.some( + recipient => + recipient.ownerId === input.ownerId && + recipient.sessionId === authorization.data.session.sessionId && + recipient.kiloSessionId === authorization.data.session.kiloSessionId && + recipient.directory === authorization.data.session.directory && + recipient.nativeRuntimeId === nativeRuntimeId.data + ) + ); + if (retired) return { quarantined: true, disposition: 'native_retired' }; + } + if (physical.state === 'stopped') + return { quarantined: false, disposition: 'physical_stopped' }; + const connection = this.activeConnection ?? this.socketHandler.getConnectionIdentity(); + if (connection?.wrapperInstanceId && connection.wrapperInstanceId !== input.wrapperInstanceId) { + return { quarantined: false, disposition: 'wrapper_replaced' }; + } const retirementConnection = this.nativeRetirementConnection(connection); const wrapperInstanceId = physical.stopTombstone?.wrapperInstanceId ?? connection?.wrapperInstanceId; - if (wrapperInstanceId !== input.wrapperInstanceId) return { quarantined: false }; + if (wrapperInstanceId !== input.wrapperInstanceId) + return { quarantined: false, disposition: 'unconfirmed' }; const route = (await loadRouteTable(this.ctx.storage)).get(input.sessionId); + const targetRoute = + nativeRuntimeId?.success && authorization?.success + ? route && + route.ownerId === input.ownerId && + route.kiloSessionId === authorization.data.session.kiloSessionId && + route.directory === authorization.data.session.directory && + route.nativeRuntimeId === nativeRuntimeId.data + ? route + : undefined + : route; + if (nativeRuntimeId?.success && !targetRoute) + return { quarantined: false, disposition: 'unconfirmed' }; if ( - route && - route.ownerId === input.ownerId && + targetRoute && + targetRoute.ownerId === input.ownerId && retirementConnection && - route.nativeRuntimeId !== undefined && + targetRoute.nativeRuntimeId !== undefined && !physical.stopTombstone ) { const retirement = await this.nativeRuntimeRetirement.retire({ connection: retirementConnection, physical, - route, + route: targetRoute, reason: input.reason, }); - if (retirement !== 'unavailable') return { quarantined: true }; + if (retirement === 'retired') return { quarantined: true, disposition: 'native_retired' }; + if (retirement === 'pending') return { quarantined: true, disposition: 'native_pending' }; } + if (nativeRuntimeId?.success) return { quarantined: false, disposition: 'unconfirmed' }; if (!physical.stopTombstone) { const next = beginStop(physical, input.reason, Date.now(), wrapperInstanceId); await this.persistPhysical(physical, next, input.reason); @@ -990,7 +1048,7 @@ export class SandboxControl extends DurableObject { } else { await this.repairLifecycleScheduling(physical); } - return { quarantined: true }; + return { quarantined: true, disposition: 'physical_stopping' }; } async prepareSessionCredentials(input: { @@ -3295,26 +3353,35 @@ export class SandboxControl extends DurableObject { throw new SandboxControlConnectionError('Operation result fence changed', false); }; let attempts = 0; - const ack = await withTimeout( - withDORetry( - () => getSandboxSessionStub(this.env, route.ownerId, route.sessionId), - async stub => { - await assertCurrent(); - attempts++; - const result = await stub.receiveSandboxOperationResult({ - session, - wrapperInstanceId, - delivery, - }); - await assertCurrent(); - if (!result) - throw new SandboxControlConnectionError('Operation result was not acknowledged'); - return result; + const ack = await withDORetry( + () => getSandboxSessionStub(this.env, route.ownerId, route.sessionId), + async stub => { + await assertCurrent(); + attempts++; + const result = await stub.receiveSandboxOperationResult({ + session, + wrapperInstanceId, + delivery, + }); + await assertCurrent(); + if (!result) + throw new SandboxControlConnectionError('Operation result was not acknowledged'); + return result; + }, + 'receiveSandboxOperationResult', + { + ...DEFAULT_DO_RETRY_CONFIG, + scope: { + deadlineAt: forwardDeadlineAt, + assertCurrent: () => { + if (!current()) + throw new SandboxControlConnectionError( + 'Operation result forwarding expired', + false + ); + }, }, - 'receiveSandboxOperationResult' - ), - Math.max(1, forwardDeadlineAt - Date.now()), - 'Operation result forwarding timed out' + } ); if (!current()) throw new SandboxControlConnectionError('Operation result expired', false); diff --git a/services/cloud-agent-next/src/sandbox-control/diagnostics.ts b/services/cloud-agent-next/src/sandbox-control/diagnostics.ts index fd94807f46..11d40d6abe 100644 --- a/services/cloud-agent-next/src/sandbox-control/diagnostics.ts +++ b/services/cloud-agent-next/src/sandbox-control/diagnostics.ts @@ -1,4 +1,4 @@ -import { withDORetry } from '@kilocode/worker-utils'; +import { withDORetry, type DORetryConfig } from '@kilocode/worker-utils'; import { logger } from '../logger.js'; export type ControlDiagnosticFields = Record; @@ -126,7 +126,8 @@ export function diagnosticConnection( export function withControlDORetry( getStub: () => TStub, operation: (stub: TStub) => Promise, - operationName: string + operationName: string, + config?: DORetryConfig ): Promise { const logRetry = (_message: unknown, fields: unknown) => { try { @@ -169,7 +170,7 @@ export function withControlDORetry( return; } }; - return withDORetry(getStub, operation, operationName, undefined, { + return withDORetry(getStub, operation, operationName, config, { warn: logRetry, error: logRetry, }); 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 e82356c68e..1ef3657881 100644 --- a/services/cloud-agent-next/src/sandbox-control/lifecycle.test.ts +++ b/services/cloud-agent-next/src/sandbox-control/lifecycle.test.ts @@ -1388,7 +1388,10 @@ describe('SandboxControl lifecycle boundaries', () => { expect(h.records.has('wrapper_credential_hash')).toBe(false); expect(h.records.has('active_wrapper_runtime')).toBe(false); vi.setSystemTime(Date.now() + 1_000); - await expect(h.control.quarantineRuntime(quarantine)).resolves.toEqual({ quarantined: true }); + await expect(h.control.quarantineRuntime(quarantine)).resolves.toEqual({ + quarantined: true, + disposition: 'physical_stopping', + }); expect(await loadDeadlines(h.storage)).toEqual(repaired); expect((await h.control.getPhysicalRecord()).stopTombstone?.attempts).toBe(0); await h.fireAlarm(); @@ -2208,9 +2211,15 @@ describe('SandboxControl lifecycle boundaries', () => { ); await expect( h.control.quarantineRuntime({ ...request, wrapperInstanceId: 'stale' }) - ).resolves.toEqual({ quarantined: false }); - await expect(h.control.quarantineRuntime(request)).resolves.toEqual({ quarantined: true }); - await expect(h.control.quarantineRuntime(request)).resolves.toEqual({ quarantined: true }); + ).resolves.toEqual({ quarantined: false, disposition: 'wrapper_replaced' }); + await expect(h.control.quarantineRuntime(request)).resolves.toEqual({ + quarantined: true, + disposition: 'physical_stopping', + }); + await expect(h.control.quarantineRuntime(request)).resolves.toEqual({ + quarantined: true, + disposition: 'physical_stopping', + }); await expect(observed.promise).resolves.toEqual({ attempts: 1, alarmAt: expect.any(Number) }); expect(await h.control.getPhysicalRecord()).toMatchObject({ state: 'stopping', @@ -2255,8 +2264,11 @@ describe('SandboxControl lifecycle boundaries', () => { await h.control.detachSession(ROUTE.sessionId); await expect( h.control.quarantineRuntime({ ...input, wrapperInstanceId: 'stale' }) - ).resolves.toEqual({ quarantined: false }); - await expect(transfer(input)).resolves.toEqual({ quarantined: true }); + ).resolves.toEqual({ quarantined: false, disposition: 'wrapper_replaced' }); + await expect(transfer(input)).resolves.toEqual({ + quarantined: true, + disposition: 'physical_stopping', + }); await h.flush(); expect(await h.control.getPhysicalRecord()).toMatchObject({ state: 'stopping', @@ -2268,7 +2280,10 @@ describe('SandboxControl lifecycle boundaries', () => { await h.fireAlarm(); expect(h.runtime(identity.providerInstanceId)?.state.running).toBe(false); expect(h.runtime(identity.providerInstanceId)?.destroy).toHaveBeenCalledTimes(2); - await expect(h.control.quarantineRuntime(input)).resolves.toEqual({ quarantined: false }); + await expect(h.control.quarantineRuntime(input)).resolves.toEqual({ + quarantined: false, + disposition: 'physical_stopped', + }); expect((await h.control.getPhysicalRecord()).providerRef).toBeNull(); }); @@ -2768,13 +2783,82 @@ describe('SandboxControl lifecycle boundaries', () => { vi.setSystemTime(Date.now() + DEADLINE_MS.stopAttempt); reply.reject(new Error('late lost reply')); - await expect(quarantine).resolves.toEqual({ quarantined: true }); + await expect(quarantine).resolves.toEqual({ quarantined: true, disposition: 'native_pending' }); await expect(h.control.getPhysicalRecord()).resolves.toMatchObject({ state: 'running', stopTombstone: null, }); }); + it('replays a completed native quarantine after its response is lost without selecting N2', async () => { + const h = await harness(); + await h.create(); + const identity = await h.ready(); + const nativeRuntimeId = '11111111-1111-4111-8111-111111111111'; + const replacementRuntimeId = '22222222-2222-4222-8222-222222222222'; + const [route] = await h.control.listRoutes(); + if (!route) throw new Error('Missing route'); + const authorization = { + operation: 'session.attach' as const, + operationId: '33333333-3333-4333-8333-333333333333', + messageId: 'message_1', + session: { + sessionId: route.sessionId, + kiloSessionId: route.kiloSessionId, + directory: route.directory, + }, + wrapperInstanceId: identity.wrapperInstanceId ?? '', + dispatchDeadlineAt: Date.now() + 1_000, + }; + h.records.set('session_routes', [{ ...route, nativeRuntimeId }]); + h.sendRequest.mockResolvedValueOnce({ + type: 'response', + requestId: 'retire_n1', + ok: true, + result: { status: 'aborted', quiescent: true, runtimeRetired: true, nativeRuntimeId }, + }); + + await expect( + h.control.quarantineRuntime({ + ownerId: OWNER, + sessionId: route.sessionId, + wrapperInstanceId: identity.wrapperInstanceId ?? '', + reason: 'runtime_unhealthy', + nativeRuntimeId, + authorization, + }) + ).resolves.toEqual({ quarantined: true, disposition: 'native_retired' }); + + h.records.set('session_routes', [{ ...route, nativeRuntimeId: replacementRuntimeId }]); + const beforeRetry = { + deadlines: await loadDeadlines(h.storage), + alarmAt: await h.storage.getAlarm(), + }; + h.sendRequest.mockClear(); + + await expect( + h.control.quarantineRuntime({ + ownerId: OWNER, + sessionId: route.sessionId, + wrapperInstanceId: identity.wrapperInstanceId ?? '', + reason: 'runtime_unhealthy', + nativeRuntimeId, + authorization, + }) + ).resolves.toEqual({ quarantined: true, disposition: 'native_retired' }); + + expect(h.sendRequest).not.toHaveBeenCalled(); + await expect(h.control.getPhysicalRecord()).resolves.toMatchObject({ + state: 'running', + stopTombstone: null, + }); + await expect(h.control.listRoutes()).resolves.toEqual([ + expect.objectContaining({ nativeRuntimeId: replacementRuntimeId }), + ]); + expect(await loadDeadlines(h.storage)).toEqual(beforeRetry.deadlines); + expect(await h.storage.getAlarm()).toBe(beforeRetry.alarmAt); + }); + it('coalesces concurrent native retirement commands for one runtime', async () => { const h = await harness(); await h.create(); @@ -2809,7 +2893,7 @@ describe('SandboxControl lifecycle boundaries', () => { wrapperInstanceId: identity.wrapperInstanceId ?? '', reason: 'native_cleanup_unavailable', }) - ).resolves.toEqual({ quarantined: true }); + ).resolves.toEqual({ quarantined: true, disposition: 'native_pending' }); expect(h.sendRequest).toHaveBeenCalledOnce(); reply.resolve({ type: 'response', @@ -2817,7 +2901,7 @@ describe('SandboxControl lifecycle boundaries', () => { ok: true, result: { status: 'aborted', quiescent: true, runtimeRetired: true, nativeRuntimeId }, }); - await expect(first).resolves.toEqual({ quarantined: true }); + await expect(first).resolves.toEqual({ quarantined: true, disposition: 'native_retired' }); }); it('retains an exhausted notification fence without escalating a completed retirement', async () => { @@ -2965,7 +3049,7 @@ describe('SandboxControl lifecycle boundaries', () => { wrapperInstanceId: identity.wrapperInstanceId ?? '', reason: 'native_cleanup_unavailable', }) - ).resolves.toEqual({ quarantined: true }); + ).resolves.toEqual({ quarantined: true, disposition: 'physical_stopping' }); await expect(h.control.getPhysicalRecord()).resolves.toMatchObject({ state: 'stopping', @@ -3158,7 +3242,7 @@ describe('SandboxControl lifecycle boundaries', () => { wrapperInstanceId: identity.wrapperInstanceId ?? '', reason: 'late_failure', }) - ).resolves.toEqual({ quarantined: false }); + ).resolves.toEqual({ quarantined: false, disposition: 'wrapper_replaced' }); }); it('retains a known allocation through a hanging launch and cleans it without a background replacement', async () => { diff --git a/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts b/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts index 9244519dd7..d3f1fadb3c 100644 --- a/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts +++ b/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts @@ -178,7 +178,11 @@ import { type ControlSessionMessageInput, type SessionMessageRecord, } from './session-message-queue.js'; -import { commitSessionOperationResult, dispatchSessionOperation } from './session-operation.js'; +import { + commitSessionOperationResult, + dispatchSessionOperation, + operationDispatchError, +} from './session-operation.js'; const METADATA_KEY = SANDBOX_SESSION_METADATA_KEY; const MESSAGES_KEY = 'session_messages'; @@ -206,6 +210,8 @@ const pendingRuntimeCleanupSchema = z.object({ sandboxId: z.string().min(1), wrapperInstanceId: wrapperInstanceIdSchema, reason: z.string(), + nativeRuntimeId: z.string().uuid().optional(), + authorization: sessionOperationAuthorizationSchema.optional(), }); type MessageRecord = SessionMessageRecord; @@ -1645,25 +1651,39 @@ export class SandboxSession extends DurableObject { commit: messages => this.saveMessages(messages, epoch, 'operation_result'), }, { - request: input => sandboxControlRpc(this.env, sandboxId).request(input), + request: (input, scope) => sandboxControlRpc(this.env, sandboxId, scope).request(input), persistResult: delivery => this.applySandboxOperationResult({ session: authorization.data.session, wrapperInstanceId: authorization.data.wrapperInstanceId, delivery, }), - isDispatchCurrent: () => false, - isMaintenanceCurrent: () => { + assertAdmission: () => { + const current = this.loadMessages().find(item => item.messageId === message.messageId); + if ( + this.terminalLifecycle.isCurrent(epoch) && + current?.wrapperInstanceId === authorization.data.wrapperInstanceId + ) + return; + throw new Error('Original operation scope is unavailable'); + }, + assertScope: () => { const current = this.loadMessages().find(item => item.messageId === message.messageId); - return ( + if ( this.terminalLifecycle.isCurrent(epoch) && current?.wrapperInstanceId === authorization.data.wrapperInstanceId && current.operations?.prompt?.dispatched === true - ); + ) + return; + throw new Error('Original operation scope is unavailable'); }, + defer: pending => this.ctx.waitUntil(pending), + isCurrent: () => false, } ); - return dispatched.state; + return dispatched.state === 'running' || dispatched.state === 'completed' + ? dispatched.state + : undefined; } async alarm(): Promise { @@ -2043,23 +2063,35 @@ export class SandboxSession extends DurableObject { commit: messages => this.saveMessages(messages, epoch, 'wrapper_outcome'), }, { - request: input => control.request(input), + request: (input, scope) => sandboxControlRpc(this.env, sandboxId, scope).request(input), persistResult: delivery => this.applySandboxOperationResult({ session: authorization.session, wrapperInstanceId: authorization.wrapperInstanceId, delivery, }), - isDispatchCurrent: isCurrent, - isMaintenanceCurrent: () => { - if (!this.terminalLifecycle.isCurrent(epoch)) return false; + 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) return false; - return ( + 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'); }, + defer: pending => this.ctx.waitUntil(pending), + isCurrent, } ); if ( @@ -2077,6 +2109,10 @@ export class SandboxSession extends DurableObject { }); } } + if (dispatched.state !== 'response' && dispatched.state !== 'completed') { + if (dispatched.state === 'running' && operation === 'session.prompt') return dispatched; + throw operationDispatchError(dispatched); + } return dispatched; }; await this.armQueueRetry(Math.min(deadlineAt, Date.now() + QUEUE_RETRY_MS)); @@ -2090,7 +2126,7 @@ export class SandboxSession extends DurableObject { ); return; } - if (this.pendingRuntimeCleanup()) return; + if (this.pendingRuntimeCleanup() && !(await this.transferRuntimeCleanup())) return; const intent = queued.intent; const model = dispatchedKilocodeModelId(intent?.agent.model); const control = sandboxControlRpc(this.env, sandboxId); @@ -2448,13 +2484,27 @@ export class SandboxSession extends DurableObject { const sandboxId = metadata.workspace?.sandboxId; if (!sandboxId) return; const pending = this.pendingRuntimeCleanup(); - if (pending && pending.wrapperInstanceId !== wrapperInstanceId) return; + // The first retained fence names the runtime that caused this cleanup. Never retarget it. + if (pending) return; + const nativeRuntime = nativeRuntimeFenceSchema.safeParse( + this.ctx.storage.kv.get(NATIVE_RUNTIME_FENCE_KEY) + ); + const target = + nativeRuntime.success && + nativeRuntime.data.sandboxId === sandboxId && + nativeRuntime.data.wrapperInstanceId === wrapperInstanceId + ? { + nativeRuntimeId: nativeRuntime.data.nativeRuntimeId, + authorization: nativeRuntime.data.authorization, + } + : {}; this.ctx.storage.kv.put(PENDING_RUNTIME_CLEANUP_KEY, { ownerId: metadata.identity.userId, sessionId: metadata.identity.sessionId, sandboxId, wrapperInstanceId, reason, + ...target, }); this.terminalLifecycle.invalidateRuntime({ sandboxId, wrapperInstanceId, confirmed: false }); } @@ -2478,6 +2528,12 @@ export class SandboxSession extends DurableObject { sessionId: pending.sessionId, wrapperInstanceId: pending.wrapperInstanceId, reason: pending.reason, + ...(pending.nativeRuntimeId + ? { + nativeRuntimeId: pending.nativeRuntimeId, + authorization: pending.authorization, + } + : {}), }), 'quarantineRuntime' ), @@ -2486,10 +2542,24 @@ export class SandboxSession extends DurableObject { ); if (typeof response?.quarantined !== 'boolean') throw new Error('Invalid quarantine response'); - if (this.pendingRuntimeCleanup()?.wrapperInstanceId === pending.wrapperInstanceId) { - this.ctx.storage.kv.delete(PENDING_RUNTIME_CLEANUP_KEY); + if ( + response.disposition === 'native_retired' || + response.disposition === 'physical_stopped' || + response.disposition === 'wrapper_replaced' + ) { + if (this.pendingRuntimeCleanup()?.wrapperInstanceId === pending.wrapperInstanceId) + this.ctx.storage.kv.delete(PENDING_RUNTIME_CLEANUP_KEY); + return this.pendingRuntimeCleanup() === undefined; + } + if ( + response.disposition === 'native_pending' || + response.disposition === 'physical_stopping' || + response.disposition === 'unconfirmed' + ) { + await this.armQueueRetry(Date.now() + SANDBOX_CONTROL_REQUEST_TIMEOUT_MS); + return false; } - return this.pendingRuntimeCleanup() === undefined; + throw new Error('Invalid quarantine disposition'); } catch { logger.withFields({ sessionId: this.sessionId }).warn('Runtime quarantine transfer failed'); await this.armQueueRetry(Date.now() + SANDBOX_CONTROL_REQUEST_TIMEOUT_MS); @@ -2545,7 +2615,12 @@ export class SandboxSession extends DurableObject { if (epoch === null) return; const before = this.loadMessages(); if (!this.terminalLifecycle.isCurrent(epoch)) return; - const { messages, failedIds } = applyFailWaitingMessages(before, reason, wrapperInstanceId); + const { messages, failedIds } = applyFailWaitingMessages( + before, + reason, + wrapperInstanceId, + false + ); if (failedIds.length === 0 || !this.saveMessages(messages, epoch)) return; if (this.pendingRuntimeCleanup() || nextQueuedMessageId(this.loadMessages())) await this.armQueueRetry(); 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 0778a15df9..e6796eea48 100644 --- a/services/cloud-agent-next/src/sandbox-session/control-dispatch.ts +++ b/services/cloud-agent-next/src/sandbox-session/control-dispatch.ts @@ -16,12 +16,14 @@ export const SESSION_DELIVERY_TIMEOUT_MS = export class ControlRequestError extends Error { readonly code: string; readonly retryable: boolean; + readonly admission: ControlError['admission']; constructor(error: ControlError) { super(error.message); this.name = 'ControlRequestError'; this.code = error.code; this.retryable = error.retryable; + this.admission = error.admission; } } diff --git a/services/cloud-agent-next/src/sandbox-session/control-rpc.ts b/services/cloud-agent-next/src/sandbox-session/control-rpc.ts index 159347c206..52fc7c653b 100644 --- a/services/cloud-agent-next/src/sandbox-session/control-rpc.ts +++ b/services/cloud-agent-next/src/sandbox-session/control-rpc.ts @@ -1,6 +1,12 @@ import type { VercelSandboxNetworkPolicy } from '../agent-sandbox/vercel/vercel-sandbox-rest-client.js'; import type { CredentialContainmentRequirements } from '../sandbox-control/physical-lifecycle.js'; -import type { ResponseFrame, SessionAttachPayload } from '../shared/sandbox-control-protocol.js'; +import { + SANDBOX_CONTROL_REQUEST_TIMEOUT_MS, + sessionAbortPayloadSchema, + type ResponseFrame, + type SessionAttachPayload, +} from '../shared/sandbox-control-protocol.js'; +import { DEFAULT_DO_RETRY_CONFIG, type DORetryScope } from '@kilocode/worker-utils'; import type { SandboxControlOutboundRequest } from '../sandbox-control/socket.js'; import type { AttachRouteInput } from '../sandbox-control/session-routes.js'; import type { ConnectionState, PhysicalState } from '../sandbox-control/status-projection.js'; @@ -8,8 +14,9 @@ import type { SandboxTerminalAccessInput, SandboxTerminalAccessResult, } from '../sandbox-control/terminal-billing.js'; +import type { SessionOperationAuthorization } from '../shared/sandbox-control-protocol.js'; import type { Env } from '../types.js'; -import type { SandboxAcquisition } from '../persistence/SandboxControl.js'; +import type { RuntimeQuarantineResult, SandboxAcquisition } from '../persistence/SandboxControl.js'; import type { SandboxBillingInput } from '../container-usage-context.js'; import { getSandboxControlStub } from '../sandbox-control/stub.js'; import { withDORetry } from '../utils/do-retry.js'; @@ -45,7 +52,9 @@ type SandboxControlRpc = { sessionId: string; wrapperInstanceId: string; reason: string; - }): Promise<{ quarantined: boolean }>; + nativeRuntimeId?: string; + authorization?: SessionOperationAuthorization; + }): Promise; attachSession(input: AttachRouteInput): Promise; detachSession(sessionId: string): Promise<{ existed: boolean }>; validateTerminalAccess(input: SandboxTerminalAccessInput): Promise; @@ -58,8 +67,22 @@ type SandboxControlRpc = { request(input: SandboxControlOutboundRequest): Promise; }; -export function sandboxControlRpc(env: Env, sandboxId: string): SandboxControlRpc { +export function sandboxControlRpc( + env: Env, + sandboxId: string, + scope?: DORetryScope +): SandboxControlRpc { const stub = () => getSandboxControlStub(env, sandboxId); + const config = ( + deadlineAt = Date.now() + SANDBOX_CONTROL_REQUEST_TIMEOUT_MS, + retrySafe = true + ) => ({ + ...DEFAULT_DO_RETRY_CONFIG, + maxAttempts: retrySafe ? DEFAULT_DO_RETRY_CONFIG.maxAttempts : 1, + ...(retrySafe + ? { scope: { ...scope, deadlineAt: Math.min(scope?.deadlineAt ?? Infinity, deadlineAt) } } + : {}), + }); return { prepareSessionCredentials: input => withDORetry( @@ -68,7 +91,7 @@ export function sandboxControlRpc(env: Env, sandboxId: string): SandboxControlRp 'prepareSessionCredentials' ), ensureReady: input => stub().ensureReady(input), - getStatus: () => stub().getStatus(), + getStatus: () => withDORetry(stub, control => control.getStatus(), 'getStatus', config()), quarantineRuntime: input => stub().quarantineRuntime(input), attachSession: input => stub().attachSession(input), detachSession: sessionId => @@ -79,6 +102,32 @@ export function sandboxControlRpc(env: Env, sandboxId: string): SandboxControlRp withDORetry(stub, control => control.recordTerminalActivity(input), 'recordTerminalActivity'), updateNetworkPolicy: input => withDORetry(stub, control => control.updateNetworkPolicy(input), 'updateNetworkPolicy'), - request: input => stub().request(input), + request: input => { + const deadlineAt = Math.min( + input.deadlineAt ?? Infinity, + scope?.deadlineAt ?? Infinity, + Date.now() + (input.timeoutMs ?? SANDBOX_CONTROL_REQUEST_TIMEOUT_MS) + ); + const abort = + input.operation === 'session.abort' + ? sessionAbortPayloadSchema.safeParse(input.payload) + : undefined; + const retrySafe = + [ + 'sandbox.status', + 'session.sync', + 'session.operation.get', + 'session.operation.ack', + ].includes(input.operation) || + (abort?.success === true && + abort.data.operationId !== undefined && + abort.data.messageId !== undefined); + return withDORetry( + stub, + control => control.request(input), + 'controlRequest', + config(deadlineAt, retrySafe) + ); + }, }; } diff --git a/services/cloud-agent-next/src/sandbox-session/session-delivery.test.ts b/services/cloud-agent-next/src/sandbox-session/session-delivery.test.ts new file mode 100644 index 0000000000..12e06decff --- /dev/null +++ b/services/cloud-agent-next/src/sandbox-session/session-delivery.test.ts @@ -0,0 +1,75 @@ +import { describe, expect, it, vi } from 'vitest'; +import type { + ResponseFrame, + SessionOperationDelivery, +} from '../shared/sandbox-control-protocol.js'; +import { persistSessionOperationDelivery } from './session-delivery.js'; + +const delivery: SessionOperationDelivery = { + version: 2, + authorization: { + operation: 'session.prompt', + operationId: 'delivery_operation', + messageId: 'delivery_message', + session: { + sessionId: 'workspace_delivery', + kiloSessionId: 'kilo_delivery', + directory: '/workspace/delivery', + }, + wrapperInstanceId: '11111111-1111-4111-8111-111111111111', + dispatchDeadlineAt: Date.now() + 60_000, + }, + completedAt: Date.now(), + result: { ok: true, result: { messageId: 'delivery_message', status: 'accepted' } }, + outcome: { messageId: 'delivery_message', status: 'completed' }, + events: [], + preparing: [], +}; + +function response(): ResponseFrame { + return { type: 'response', requestId: 'ack', ok: true, result: { acknowledged: true } }; +} + +describe('persistSessionOperationDelivery', () => { + it('persists before starting the acknowledgement handoff', async () => { + const request = vi.fn(async () => response()); + const persistResult = vi.fn(async () => ({ + version: 2 as const, + authorization: delivery.authorization, + resultHash: 'result', + disposition: 'applied' as const, + decision: { state: 'completed' as const, at: delivery.completedAt }, + })); + + await expect( + persistSessionOperationDelivery(delivery, Date.now() + 60_000, { + request, + persistResult, + assertScope: () => undefined, + defer: pending => void pending, + }) + ).resolves.toBe('persisted'); + await Promise.resolve(); + expect(persistResult.mock.invocationCallOrder[0]).toBeLessThan( + request.mock.invocationCallOrder[0] + ); + expect(request).toHaveBeenCalledWith( + expect.objectContaining({ operation: 'session.operation.ack' }), + expect.any(Object) + ); + }); + + it('does not acknowledge an unpersisted result', async () => { + const request = vi.fn(async () => response()); + + await expect( + persistSessionOperationDelivery(delivery, Date.now() + 60_000, { + request, + persistResult: async () => undefined, + assertScope: () => undefined, + defer: pending => void pending, + }) + ).resolves.toBe('unverified'); + expect(request).not.toHaveBeenCalled(); + }); +}); diff --git a/services/cloud-agent-next/src/sandbox-session/session-delivery.ts b/services/cloud-agent-next/src/sandbox-session/session-delivery.ts new file mode 100644 index 0000000000..d63c8caa07 --- /dev/null +++ b/services/cloud-agent-next/src/sandbox-session/session-delivery.ts @@ -0,0 +1,48 @@ +import type { DORetryScope } from '@kilocode/worker-utils'; +import type { SandboxControlOutboundRequest } from '../sandbox-control/socket.js'; +import { + SANDBOX_CONTROL_OUTCOME_TIMEOUT_MS, + type ResponseFrame, + type SessionOperationAck, + type SessionOperationDelivery, +} from '../shared/sandbox-control-protocol.js'; +import { controlRequestResult, withDeliveryDeadline } from './control-dispatch.js'; + +export type SessionOperationDeliveryEffects = { + request: (input: SandboxControlOutboundRequest, scope: DORetryScope) => Promise; + persistResult: (delivery: SessionOperationDelivery) => Promise; + assertScope: () => void; + defer: (pending: Promise) => void; +}; + +export async function persistSessionOperationDelivery( + delivery: SessionOperationDelivery, + operationDeadlineAt: number, + effects: SessionOperationDeliveryEffects +): Promise<'persisted' | 'unverified'> { + const acknowledgement = await effects.persistResult(delivery); + if (!acknowledgement) return 'unverified'; + const acknowledgementDeadlineAt = Math.min( + operationDeadlineAt, + delivery.completedAt + SANDBOX_CONTROL_OUTCOME_TIMEOUT_MS + ); + if (Date.now() >= acknowledgementDeadlineAt) return 'persisted'; + effects.defer( + withDeliveryDeadline(async () => { + effects.assertScope(); + controlRequestResult( + await effects.request( + { + operation: 'session.operation.ack', + session: delivery.authorization.session, + payload: acknowledgement, + expectedWrapperInstanceId: delivery.authorization.wrapperInstanceId, + deadlineAt: acknowledgementDeadlineAt, + }, + { deadlineAt: acknowledgementDeadlineAt, assertCurrent: effects.assertScope } + ) + ); + }, acknowledgementDeadlineAt).catch(() => undefined) + ); + return 'persisted'; +} 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 1f8850cc0e..f3fcdb8375 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 @@ -15,13 +15,17 @@ import type { sandboxControlRpc } from './control-rpc.js'; import type { SandboxControlOutboundRequest } from '../sandbox-control/socket.js'; import { SANDBOX_CONTROL_ATTACH_TIMEOUT_MS, + SANDBOX_CONTROL_OUTCOME_TIMEOUT_MS, SANDBOX_CONTROL_REQUEST_TIMEOUT_MS, + sessionOperationExpiresAt, + sessionOperationResultHash, sessionPromptPayloadSchema, sessionGitSummaryPayloadSchema, type ResponseFrame, type SessionAttachPayload, type SessionMessageOutcome, type SessionOperationAuthorization, + type SessionOperationDelivery, } from '../shared/sandbox-control-protocol.js'; import { DEADLINE_MS } from '../sandbox-control/deadlines.js'; import { createControlPlaneCredential } from '../sandbox-control/managed-credential.js'; @@ -921,6 +925,7 @@ const ATTACHMENT = { type Control = ReturnType; type ControlStatus = Awaited>; +type RuntimeQuarantineResult = Awaited>; function deferred() { let resolve: (value: T) => void = () => undefined; @@ -1032,9 +1037,14 @@ function sessionFixture(overrides: Partial = {}, sharedControl? ), attachSession: vi.fn(async () => ({})), detachSession: vi.fn(async () => ({ existed: true })), - quarantineRuntime: vi.fn(async (_input: Parameters[0]) => ({ - quarantined: true, - })), + quarantineRuntime: vi.fn( + async ( + _input: Parameters[0] + ): Promise => ({ + quarantined: true, + disposition: 'physical_stopping', + }) + ), validateTerminalAccess: vi.fn(async () => ({ allowed: true })), recordTerminalActivity: vi.fn(async () => ({ allowed: true })), prepareSessionCredentials: vi.fn(async () => ({})), @@ -1189,6 +1199,186 @@ describe('SandboxSession orchestration', () => { expect(fixture.record('b')?.deliveryDeadlineAt).toBe(Date.now() + SESSION_DELIVERY_TIMEOUT_MS); }); + it('dispatches the first normal attach and prompt once with operation receipts enabled', async () => { + const fixture = sessionFixture(); + fixture.setStatus({ + physical: 'running', + connection: 'ready', + wrapperInstanceId: RUNTIME_ID, + operationResults: true, + }); + + await fixture.admit('a'); + await fixture.flush(); + + const operations = fixture.control.request.mock.calls + .map(([input]) => input.operation) + .filter(operation => operation.startsWith('session.')); + expect(operations).toEqual(['session.attach', 'session.prompt']); + expect(fixture.record('a')).toMatchObject({ + state: 'accepted', + operations: { + attach: { dispatched: true }, + prompt: { dispatched: true }, + }, + }); + }); + + it.each(['running', 'completed'] as const)( + 'reconstructs a late accepted prompt from its original %s operation result without redispatch', + async state => { + const fixture = sessionFixture(); + fixture.setStatus({ + physical: 'running', + connection: 'ready', + wrapperInstanceId: RUNTIME_ID, + operationResults: true, + }); + await fixture.admit('a'); + await fixture.flush(); + const authorization = fixture.record('a')?.operations?.prompt?.authorization; + if (!authorization) throw new Error('Missing prompt operation authorization'); + const completedAt = authorization.dispatchDeadlineAt + 1; + const delivery: SessionOperationDelivery = { + version: 2, + authorization, + completedAt, + result: { ok: true, result: { messageId: 'a', status: 'accepted' } }, + outcome: { messageId: 'a', status: 'completed' }, + events: [], + preparing: [], + }; + const resultHash = await sessionOperationResultHash(delivery); + delegateRequest(fixture, 'session.operation.get', async input => { + expect(input).toMatchObject({ + expectedWrapperInstanceId: RUNTIME_ID, + payload: authorization, + }); + expect(input.deadlineAt).toBe(sessionOperationExpiresAt(authorization)); + return controlResponse( + state === 'running' ? { state, authorization } : { state, delivery } + ); + }); + delegateRequest(fixture, 'session.operation.ack', async input => { + expect(state).toBe('completed'); + expect(input).toMatchObject({ + expectedWrapperInstanceId: RUNTIME_ID, + payload: { + version: 2, + authorization, + resultHash, + disposition: 'applied', + decision: { state: 'completed', at: completedAt }, + }, + }); + expect(input.deadlineAt).toBe(completedAt + SANDBOX_CONTROL_OUTCOME_TIMEOUT_MS); + return controlResponse({ acknowledged: true }); + }); + + vi.setSystemTime(authorization.dispatchDeadlineAt + 1); + fixture.reload(); + await fixture.fireAlarm(); + await fixture.flush(); + + expect( + fixture.control.request.mock.calls.filter(([input]) => input.operation === 'session.prompt') + ).toHaveLength(1); + expect( + fixture.control.request.mock.calls.filter( + ([input]) => input.operation === 'session.operation.get' + ) + ).toHaveLength(1); + expect( + fixture.control.request.mock.calls.filter( + ([input]) => input.operation === 'session.operation.ack' + ) + ).toHaveLength(state === 'completed' ? 1 : 0); + expect(fixture.record('a')?.state).toBe(state === 'completed' ? 'completed' : 'accepted'); + } + ); + + it('allows B to reattach on the same wrapper only after native retirement is confirmed', async () => { + const fixture = sessionFixture(); + const nativeRuntimeId = '11111111-1111-4111-8111-111111111111'; + fixture.setStatus({ + physical: 'running', + connection: 'ready', + wrapperInstanceId: RUNTIME_ID, + operationResults: true, + }); + await fixture.admit('a'); + await fixture.flush(); + const authorization = fixture.record('a')?.operations?.attach?.authorization; + if (!authorization) throw new Error('Missing attach operation authorization'); + await fixture.session.recordNativeRuntime({ + sandboxId: fixture.metadata.workspace?.sandboxId ?? '', + wrapperInstanceId: RUNTIME_ID, + nativeRuntimeId, + authorization, + }); + fixture.control.quarantineRuntime.mockImplementation(async input => { + expect(input).toMatchObject({ + sessionId: SESSION_ID, + wrapperInstanceId: RUNTIME_ID, + reason: 'runtime_unhealthy', + nativeRuntimeId, + authorization, + }); + await fixture.session.invalidateTerminalRuntime({ + sandboxId: fixture.metadata.workspace?.sandboxId ?? '', + wrapperInstanceId: RUNTIME_ID, + nativeRuntimeId, + confirmed: true, + }); + return { quarantined: true, disposition: 'native_retired' }; + }); + delegateRequest(fixture, 'session.sync', async () => + controlResponse({ status: { type: 'idle' }, questions: [], permissions: [] }) + ); + vi.advanceTimersByTime(DEADLINE_MS.acceptedOverdue); + + await fixture.fireAlarm(); + await fixture.flush(); + await fixture.admit('b'); + await fixture.flush(); + + expect(fixture.record('a')).toMatchObject({ + state: 'failed', + failedReason: 'runtime_unhealthy', + }); + expect(fixture.record('b')).toMatchObject({ state: 'accepted', wrapperInstanceId: RUNTIME_ID }); + expect( + fixture.control.request.mock.calls.filter(([input]) => input.operation === 'session.attach') + ).toHaveLength(2); + expect( + fixture.control.request.mock.calls.filter(([input]) => input.operation === 'session.prompt') + ).toHaveLength(2); + expect(await fixture.session.isSandboxCleanupScheduled()).toBe(false); + }); + + it('keeps cleanup pending when a target-scoped quarantine is unconfirmed without a successor', async () => { + const fixture = sessionFixture(); + const attach = deferred(); + delegateRequest(fixture, 'session.attach', () => attach.promise); + fixture.control.quarantineRuntime.mockResolvedValue({ + quarantined: false, + disposition: 'unconfirmed', + }); + await fixture.admit('a'); + await fixture.flush(); + + await expect(fixture.session.interruptExecution()).resolves.toEqual({ success: true }); + expect(fixture.record('a')?.state).toBe('cancelled'); + expect(fixture.record('b')).toBeUndefined(); + expect(await fixture.session.isSandboxCleanupScheduled()).toBe(true); + + attach.resolve(controlResponse({ attached: true })); + await fixture.fireAlarm(); + await fixture.flush(); + expect(await fixture.session.isSandboxCleanupScheduled()).toBe(true); + expect(fixture.control.quarantineRuntime).toHaveBeenCalledTimes(2); + }); + it.each(['cloudflare', 'vercel'] as const)( 'sends the expected runtime fence for cold and warm handoffs on %s', async provider => { @@ -1378,7 +1568,7 @@ describe('SandboxSession orchestration', () => { operation === 'ensureReady' ? DEADLINE_MS.startup : SANDBOX_CONTROL_REQUEST_TIMEOUT_MS ); expect(fixture.record('a')?.state).toBe('failed'); - expect(fixture.record('b')?.state).toBe('failed'); + expect(fixture.record('b')?.state).toBe('queued'); expect(fixture.control.quarantineRuntime).toHaveBeenCalledWith( expect.objectContaining({ wrapperInstanceId: RUNTIME_ID }) ); @@ -1532,8 +1722,8 @@ describe('SandboxSession orchestration', () => { expect(fixture.record('b')?.state).toBe('accepted'); } else { expect(fixture.record('a')?.failedReason).toBe(reason); - expect(fixture.record('b')?.state).toBe('failed'); - expect(fixture.terminalEvents()).toHaveLength(2); + expect(fixture.record('b')?.state).toBe('queued'); + expect(fixture.terminalEvents()).toHaveLength(1); expect(fixture.control.quarantineRuntime).toHaveBeenCalledWith( expect.objectContaining({ wrapperInstanceId: RUNTIME_ID, reason }) ); @@ -1575,8 +1765,8 @@ describe('SandboxSession orchestration', () => { expect(fixture.control.quarantineRuntime).not.toHaveBeenCalled(); return; } - expect(fixture.record('b')?.state).toBe('failed'); - expect(fixture.terminalEvents()).toHaveLength(2); + expect(fixture.record('b')?.state).toBe('queued'); + expect(fixture.terminalEvents()).toHaveLength(1); expect(fixture.control.quarantineRuntime).toHaveBeenCalledWith( expect.objectContaining({ wrapperInstanceId: RUNTIME_ID, reason }) ); @@ -1621,8 +1811,8 @@ describe('SandboxSession orchestration', () => { return; } expect(fixture.record('a')).toMatchObject({ state: 'failed', failedReason: reason }); - expect(fixture.record('b')?.state).toBe('failed'); - expect(fixture.terminalEvents()).toHaveLength(2); + expect(fixture.record('b')?.state).toBe('queued'); + expect(fixture.terminalEvents()).toHaveLength(1); expect(fixture.control.quarantineRuntime).toHaveBeenCalledOnce(); await fixture.fireAlarm(); expect( @@ -1650,7 +1840,7 @@ describe('SandboxSession orchestration', () => { writer.control.quarantineRuntime.mockImplementation(async input => { writer.setStatus({ physical: 'stopped', connection: 'disconnected' }); await writer.session.failWaitingMessages(input.reason, input.wrapperInstanceId); - return { quarantined: true }; + return { quarantined: false, disposition: 'physical_stopped' as const }; }); return { writer, sibling }; } @@ -1964,7 +2154,7 @@ describe('SandboxSession orchestration', () => { }); fixture.control.quarantineRuntime.mockImplementation(async () => { remoteWorkRunning = false; - return { quarantined: true }; + return { quarantined: true, disposition: 'physical_stopping' }; }); await fixture.admit('waiting'); await fixture.flush(); @@ -2277,7 +2467,7 @@ describe('SandboxSession orchestration', () => { await fixture.admit('b'); await fixture.flush(); expect(fixture.record('a')?.state).toBe('failed'); - expect(fixture.record('b')?.state).toBe('failed'); + expect(fixture.record('b')?.state).toBe('queued'); expect(fixture.control.quarantineRuntime).toHaveBeenCalledWith({ ownerId: 'user_1', sessionId: SESSION_ID, @@ -2361,7 +2551,7 @@ describe('SandboxSession orchestration', () => { const fixture = sessionFixture(); const attach = deferred(); const abort = deferred(); - const quarantine = deferred<{ quarantined: boolean }>(); + const quarantine = deferred(); if (phase === 'preparing') delegateRequest(fixture, 'session.attach', () => attach.promise); delegateRequest(fixture, 'session.abort', () => abort.promise); fixture.control.quarantineRuntime.mockImplementation(() => quarantine.promise); @@ -2392,13 +2582,14 @@ describe('SandboxSession orchestration', () => { reason: 'preparation_interrupted', }); attach.resolve(controlResponse({ attached: true })); - quarantine.resolve({ quarantined: true }); + quarantine.resolve({ quarantined: true, disposition: 'physical_stopping' }); } await expect(interruption).resolves.toEqual({ success: true }); await fixture.flush(); expect(fixture.record('a')?.state).toBe('cancelled'); expect(fixture.terminalEvents()).toHaveLength(1); if (phase === 'preparing') { + expect(await fixture.session.isSandboxCleanupScheduled()).toBe(true); expect( fixture.control.request.mock.calls.some(([input]) => input.operation === 'session.prompt') ).toBe(false); @@ -2410,7 +2601,7 @@ describe('SandboxSession orchestration', () => { const fixture = sessionFixture(); const attach = deferred(); delegateRequest(fixture, 'session.attach', () => attach.promise); - const cleanup = deferred<{ quarantined: boolean }>(); + const cleanup = deferred(); fixture.control.quarantineRuntime.mockImplementation(() => cleanup.promise); await fixture.admit('a'); await fixture.flush(); @@ -2422,7 +2613,7 @@ describe('SandboxSession orchestration', () => { expect(fixture.record('b')?.state).toBe('queued'); expect(await fixture.session.isSandboxCleanupScheduled()).toBe(true); attach.resolve(controlResponse({ attached: true })); - cleanup.resolve({ quarantined: true }); + cleanup.resolve({ quarantined: true, disposition: 'physical_stopping' }); await interrupt; await fixture.flush(); expect( @@ -2574,17 +2765,18 @@ describe('SandboxSession orchestration', () => { expect(acquisition.id).not.toBe(oldAcquisition.id); expect(fixture.control.ensureReady).toHaveBeenCalledOnce(); fixture.reload(); - fixture.control.quarantineRuntime.mockImplementation(async () => { - fixture.setStatus({ physical: 'stopping', connection: 'disconnected' }); - return { quarantined: true }; + fixture.control.quarantineRuntime.mockResolvedValueOnce({ + quarantined: true, + disposition: 'physical_stopping', }); await fixture.fireAlarm(); - expect(await fixture.session.isSandboxCleanupScheduled()).toBe(false); + expect(await fixture.session.isSandboxCleanupScheduled()).toBe(true); expect(fixture.record('b')?.state).toBe('queued'); - expect(fixture.control.ensureReady).toHaveBeenLastCalledWith( - expect.objectContaining({ acquisition }) - ); fixture.setStatus({ physical: 'stopped', connection: 'disconnected' }); + fixture.control.quarantineRuntime.mockResolvedValueOnce({ + quarantined: false, + disposition: 'physical_stopped', + }); fixture.control.ensureReady.mockImplementationOnce(async input => { expect(input.acquisition).toEqual(acquisition); const ready = { @@ -2596,6 +2788,9 @@ describe('SandboxSession orchestration', () => { return { ...ready, attachment: ATTACHMENT }; }); await fixture.fireAlarm(); + expect(fixture.control.ensureReady).toHaveBeenLastCalledWith( + expect.objectContaining({ acquisition }) + ); expect(fixture.record('b')).toMatchObject({ state: 'accepted', intent, @@ -2611,7 +2806,7 @@ describe('SandboxSession orchestration', () => { expect(fixture.record('a')?.state).toBe('cancelled'); expect(fixture.record('b')?.state).toBe('accepted'); expect(fixture.record('c')?.state).toBe('queued'); - expect(fixture.control.ensureReady).toHaveBeenCalledTimes(3); + expect(fixture.control.ensureReady).toHaveBeenCalledTimes(2); expect( fixture.control.request.mock.calls .filter(([input]) => input.operation === 'session.prompt') @@ -2899,24 +3094,27 @@ describe('SandboxSession orchestration', () => { if (health === 'hang') return new Promise(() => undefined); return controlResponse({ status: { type: 'idle' }, questions: [], permissions: [] }); }); - const cleanup = deferred<{ quarantined: boolean }>(); + const cleanup = deferred(); fixture.control.quarantineRuntime.mockImplementation(() => cleanup.promise); vi.setSystemTime(Date.now() + DEADLINE_MS.acceptedOverdue); const alarm = fixture.fireAlarm(); await vi.advanceTimersByTimeAsync(health === 'hang' ? SANDBOX_CONTROL_REQUEST_TIMEOUT_MS : 0); expect(fixture.record('a')?.state).toBe('failed'); - expect(fixture.record('b')?.state).toBe('failed'); - expect(fixture.terminalEvents()).toHaveLength(2); + expect(fixture.record('b')?.state).toBe('queued'); + expect(fixture.terminalEvents()).toHaveLength(1); expect(await fixture.session.isSandboxCleanupScheduled()).toBe(true); await fixture.admit('c'); await fixture.flush(); expect(fixture.record('c')?.state).toBe('queued'); - const acquisition = fixture.acquisition('c'); + const acquisition = fixture.acquisition('b'); const ensureCount = fixture.control.ensureReady.mock.calls.length; await vi.advanceTimersByTimeAsync(SANDBOX_CONTROL_REQUEST_TIMEOUT_MS); await alarm; fixture.reload(); - fixture.control.quarantineRuntime.mockResolvedValue({ quarantined: false }); + fixture.control.quarantineRuntime.mockResolvedValue({ + quarantined: false, + disposition: 'physical_stopped', + }); fixture.setStatus({ physical: 'stopped', connection: 'disconnected' }); fixture.control.ensureReady.mockImplementationOnce(async () => { const ready = { @@ -2929,7 +3127,7 @@ describe('SandboxSession orchestration', () => { }); await fixture.fireAlarm(); expect(await fixture.session.isSandboxCleanupScheduled()).toBe(false); - expect(fixture.record('c')).toMatchObject({ + expect(fixture.record('b')).toMatchObject({ state: 'accepted', wrapperInstanceId: NEXT_RUNTIME_ID, }); @@ -2937,6 +3135,9 @@ describe('SandboxSession orchestration', () => { expect(fixture.control.ensureReady).toHaveBeenLastCalledWith( expect.objectContaining({ acquisition }) ); + await fixture.outcome('b', 'completed', NEXT_RUNTIME_ID); + await fixture.flush(); + expect(fixture.record('c')?.state).toBe('accepted'); await fixture.outcome('c', 'completed', NEXT_RUNTIME_ID); await fixture.flush(); await fixture.admit('d'); @@ -2945,7 +3146,7 @@ describe('SandboxSession orchestration', () => { await fixture.session.failWaitingMessages('delayed_old_failure', RUNTIME_ID); expect(fixture.record('d')?.state).toBe('accepted'); expect(fixture.record('e')?.state).toBe('queued'); - cleanup.resolve({ quarantined: true }); + cleanup.resolve({ quarantined: true, disposition: 'physical_stopping' }); await fixture.flush(); expect(fixture.record('d')?.state).toBe('accepted'); } @@ -2981,7 +3182,10 @@ describe('SandboxSession orchestration', () => { it('does not hand off a prompt when the runtime changes during attachment', async () => { const fixture = sessionFixture(); - fixture.control.quarantineRuntime.mockResolvedValue({ quarantined: false }); + fixture.control.quarantineRuntime.mockResolvedValue({ + quarantined: false, + disposition: 'wrapper_replaced', + }); delegateRequest(fixture, 'session.attach', async () => { fixture.setStatus({ physical: 'running', @@ -3858,11 +4062,11 @@ describe('SandboxSession orchestration', () => { }); it.each(['rejected', 'malformed', 'error', 'timeout'] as const)( - 'quarantines an accepted runtime on %s abort and fails a follow-up admitted during cancellation', + 'quarantines an accepted runtime on %s abort and keeps a follow-up queued during cancellation', async failure => { const fixture = sessionFixture(); const abort = deferred(); - const cleanup = deferred<{ quarantined: boolean }>(); + const cleanup = deferred(); delegateRequest(fixture, 'session.abort', () => abort.promise); fixture.control.quarantineRuntime.mockImplementation(() => cleanup.promise); await fixture.admit('a'); @@ -3888,15 +4092,14 @@ describe('SandboxSession orchestration', () => { failedReason: 'runtime_unhealthy', }); expect(fixture.record('b')).toMatchObject({ - state: 'failed', - failedReason: 'runtime_unhealthy', + state: 'queued', }); expect(await fixture.session.isSandboxCleanupScheduled()).toBe(true); expect(fixture.control.quarantineRuntime).toHaveBeenCalledWith( expect.objectContaining({ wrapperInstanceId: RUNTIME_ID, reason: 'runtime_unhealthy' }) ); - expect(fixture.terminalEvents()).toHaveLength(2); - cleanup.resolve({ quarantined: true }); + expect(fixture.terminalEvents()).toHaveLength(1); + cleanup.resolve({ quarantined: true, disposition: 'physical_stopping' }); await expect(interruption).resolves.toEqual({ success: false, message: 'The session runtime could not be interrupted', @@ -3906,15 +4109,20 @@ describe('SandboxSession orchestration', () => { connection: 'ready', wrapperInstanceId: NEXT_RUNTIME_ID, }); + fixture.control.quarantineRuntime.mockResolvedValue({ + quarantined: false, + disposition: 'wrapper_replaced', + }); await fixture.admit('c'); await fixture.flush(); abort.resolve(controlResponse({ status: 'aborted' })); await fixture.flush(); expect(fixture.record('a')?.state).toBe('failed'); - expect(fixture.record('c')).toMatchObject({ + expect(fixture.record('b')).toMatchObject({ state: 'accepted', wrapperInstanceId: NEXT_RUNTIME_ID, }); + expect(fixture.record('c')?.state).toBe('queued'); } ); 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 bed50c5c2d..337641d1ce 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 @@ -251,13 +251,15 @@ export function assignPreparationAttemptId( export function failWaitingMessages( messages: readonly SessionMessageRecord[], reason: string, - wrapperInstanceId?: string + wrapperInstanceId?: string, + includeUnassigned = true ): { messages: SessionMessageRecord[]; failedIds: string[] } { const head = messages.find(message => message.state === 'accepted') ?? messages.find(message => message.state === 'queued'); const failUnassigned = - wrapperInstanceId === undefined || head?.wrapperInstanceId === wrapperInstanceId; + wrapperInstanceId === undefined || + (includeUnassigned && head?.wrapperInstanceId === wrapperInstanceId); const failedIds: string[] = []; return { messages: messages.map(message => { @@ -524,7 +526,8 @@ export function applySessionOperationResult( export function recordSessionOperationDispatch( messages: readonly SessionMessageRecord[], - authorization: SessionOperationAuthorization + authorization: SessionOperationAuthorization, + dispatched = true ): SessionMessageRecord[] | undefined { const message = messages.find(item => item.messageId === authorization.messageId); const kind = authorization.operation === 'session.attach' ? 'attach' : 'prompt'; @@ -544,14 +547,14 @@ export function recordSessionOperationDispatch( item.messageId === message.messageId ? { ...item, - unresolvedDispatch: true, + unresolvedDispatch: dispatched ? true : undefined, deliveryRetryScope: undefined, operations: { ...item.operations, [kind]: { authorization: structuredClone(authorization), - dispatched: true, - ...(kind === 'prompt' + dispatched, + ...(kind === 'prompt' && dispatched ? { executionDeadlineAt: item.executionDeadlineAt ?? @@ -560,7 +563,7 @@ export function recordSessionOperationDispatch( : {}), }, }, - ...(kind === 'prompt' + ...(kind === 'prompt' && dispatched ? { executionDeadlineAt: item.executionDeadlineAt ?? diff --git a/services/cloud-agent-next/src/sandbox-session/session-operation.test.ts b/services/cloud-agent-next/src/sandbox-session/session-operation.test.ts index 7846bce933..b196f37046 100644 --- a/services/cloud-agent-next/src/sandbox-session/session-operation.test.ts +++ b/services/cloud-agent-next/src/sandbox-session/session-operation.test.ts @@ -1,5 +1,6 @@ import { describe, expect, it, vi } from 'vitest'; import { + SANDBOX_CONTROL_OUTCOME_TIMEOUT_MS, sessionOperationResultHash, type ResponseFrame, type SessionOperationAuthorization, @@ -78,8 +79,10 @@ describe('dispatchSessionOperation', () => { { request, persistResult: async () => undefined, - isDispatchCurrent: () => true, - isMaintenanceCurrent: () => true, + assertAdmission: () => undefined, + assertScope: () => undefined, + defer: pending => void pending, + isCurrent: () => true, } ) ).resolves.toEqual({ @@ -90,35 +93,39 @@ describe('dispatchSessionOperation', () => { }); it('uses the retained result and exact acknowledgement after a lost prompt response', async () => { - const dispatched = recordSessionOperationDispatch(messages(), authorization); + const lateAuthorization = { ...authorization, dispatchDeadlineAt: Date.now() - 1_000 }; + const dispatched = recordSessionOperationDispatch(messages(), lateAuthorization); if (!dispatched) throw new Error('Failed to create dispatch proof'); let stored = dispatched; const delivery: SessionOperationDelivery = { version: 2, - authorization, + authorization: lateAuthorization, completedAt: Date.now(), - result: { ok: true, result: { messageId: authorization.messageId, status: 'accepted' } }, - outcome: { messageId: authorization.messageId, status: 'completed' }, + result: { ok: true, result: { messageId: lateAuthorization.messageId, status: 'accepted' } }, + outcome: { messageId: lateAuthorization.messageId, status: 'completed' }, events: [], preparing: [], }; const ack = { version: 2 as const, - authorization, + authorization: lateAuthorization, resultHash: await sessionOperationResultHash(delivery), disposition: 'applied' as const, decision: { state: 'completed' as const, at: delivery.completedAt }, }; const request = vi.fn(async (input: SandboxControlOutboundRequest) => { - if (input.operation === 'session.operation.get') + if (input.operation === 'session.operation.get') { + expect(input.deadlineAt).toBeGreaterThan(Date.now()); return response({ state: 'completed', delivery }); + } expect(input).toMatchObject({ operation: 'session.operation.ack', payload: ack }); + expect(input.deadlineAt).toBe(delivery.completedAt + SANDBOX_CONTROL_OUTCOME_TIMEOUT_MS); return response({ acknowledged: true }); }); await expect( dispatchSessionOperation( - { authorization, payload }, + { authorization: lateAuthorization, payload }, { read: () => stored, commit: next => { @@ -129,11 +136,13 @@ describe('dispatchSessionOperation', () => { { request, persistResult: async () => ack, - isDispatchCurrent: () => false, - isMaintenanceCurrent: () => true, + assertAdmission: () => undefined, + assertScope: () => undefined, + defer: pending => void pending, + isCurrent: () => false, } ) - ).resolves.toEqual({ state: 'completed' }); + ).resolves.toMatchObject({ state: 'completed' }); expect(request.mock.calls.map(([input]) => input.operation)).toEqual([ 'session.operation.get', 'session.operation.ack', @@ -157,11 +166,13 @@ describe('dispatchSessionOperation', () => { { request, persistResult: async () => undefined, - isDispatchCurrent: () => false, - isMaintenanceCurrent: () => true, + assertAdmission: () => undefined, + assertScope: () => undefined, + defer: pending => void pending, + isCurrent: () => false, } ) - ).resolves.toEqual({ state: 'running' }); + ).resolves.toMatchObject({ state: 'running' }); expect(request.mock.calls.map(([input]) => input.operation)).toEqual(['session.operation.get']); }); @@ -179,14 +190,177 @@ describe('dispatchSessionOperation', () => { { request, persistResult: async () => undefined, - isDispatchCurrent: () => true, - isMaintenanceCurrent: () => true, + assertAdmission: () => undefined, + assertScope: () => undefined, + defer: pending => void pending, + isCurrent: () => true, } ) - ).rejects.toThrow('Original session operation is missing'); + ).resolves.toEqual({ state: 'uncertain', reason: 'missing' }); expect(request.mock.calls.map(([input]) => input.operation)).toEqual(['session.operation.get']); }); + it('does not replay a prompt after its admission response is lost before application', async () => { + let stored = messages(); + const request = vi.fn(async (input: SandboxControlOutboundRequest) => { + if (input.operation === 'session.prompt') + throw Object.assign(new Error('Prompt admission response was lost'), { retryable: true }); + if (input.operation === 'session.operation.get') return response({ state: 'missing' }); + throw new Error(`Unexpected operation ${input.operation}`); + }); + const effects = { + request, + persistResult: async () => undefined, + assertAdmission: () => undefined, + assertScope: () => undefined, + defer: (pending: Promise) => void pending, + isCurrent: () => true, + }; + + await expect( + dispatchSessionOperation( + { authorization, payload }, + { read: () => stored, commit: next => ((stored = next), true) }, + effects + ) + ).resolves.toEqual({ state: 'uncertain', reason: 'transport', error: expect.any(Error) }); + await expect( + dispatchSessionOperation( + { authorization, payload }, + { read: () => stored, commit: next => ((stored = next), true) }, + effects + ) + ).resolves.toEqual({ state: 'uncertain', reason: 'missing' }); + expect(request.mock.calls.map(([input]) => input.operation)).toEqual([ + 'session.prompt', + 'session.operation.get', + ]); + }); + + it('keeps dispatch proof after an unmarked busy rejection', async () => { + let stored = messages(); + const request = vi.fn( + async (): Promise => ({ + type: 'response', + requestId: crypto.randomUUID(), + ok: false, + error: { code: 'session_busy', message: 'busy after admission', retryable: true }, + }) + ); + + await expect( + dispatchSessionOperation( + { authorization, payload }, + { read: () => stored, commit: next => ((stored = next), true) }, + { + request, + persistResult: async () => undefined, + assertAdmission: () => undefined, + assertScope: () => undefined, + defer: pending => void pending, + isCurrent: () => true, + } + ) + ).resolves.toMatchObject({ state: 'rejected', error: { code: 'session_busy' } }); + expect(stored[0]?.operations?.prompt?.dispatched).toBe(true); + }); + + it('clears dispatch proof only after an explicit before-admission rejection', async () => { + let stored = messages(); + const request = vi.fn( + async (): Promise => ({ + type: 'response', + requestId: crypto.randomUUID(), + ok: false, + error: { + code: 'session_busy', + message: 'receipt capacity is unavailable', + retryable: true, + admission: 'not-admitted', + }, + }) + ); + + await expect( + dispatchSessionOperation( + { authorization, payload }, + { read: () => stored, commit: next => ((stored = next), true) }, + { + request, + persistResult: async () => undefined, + assertAdmission: () => undefined, + assertScope: () => undefined, + defer: pending => void pending, + isCurrent: () => true, + } + ) + ).resolves.toMatchObject({ state: 'rejected', error: { code: 'session_busy' } }); + expect(stored[0]).toMatchObject({ + unresolvedDispatch: undefined, + operations: { prompt: { authorization, dispatched: false } }, + }); + }); + + it('persists a recovered result before its acknowledgement can be retried', async () => { + const dispatched = recordSessionOperationDispatch(messages(), authorization); + if (!dispatched) throw new Error('Failed to create dispatch proof'); + let stored = dispatched; + const delivery: SessionOperationDelivery = { + version: 2, + authorization, + completedAt: Date.now(), + result: { ok: true, result: { messageId: authorization.messageId, status: 'accepted' } }, + outcome: { messageId: authorization.messageId, status: 'completed' }, + events: [], + preparing: [], + }; + const ack = { + version: 2 as const, + authorization, + resultHash: await sessionOperationResultHash(delivery), + disposition: 'applied' as const, + decision: { state: 'completed' as const, at: delivery.completedAt }, + }; + const request = vi.fn(async (input: SandboxControlOutboundRequest) => { + if (input.operation === 'session.operation.get') + return response({ state: 'completed', delivery }); + if (input.operation === 'session.operation.ack') + throw Object.assign(new Error('Acknowledgement response was lost'), { retryable: true }); + throw new Error(`Unexpected operation ${input.operation}`); + }); + + await expect( + dispatchSessionOperation( + { authorization, payload }, + { read: () => stored, commit: next => ((stored = next), true) }, + { + request, + persistResult: async receipt => { + const applied = applySessionOperationResult( + stored, + receipt, + await sessionOperationResultHash(receipt), + Date.now() + ); + if (!applied) return undefined; + stored = applied.messages; + return ack; + }, + assertAdmission: () => undefined, + assertScope: () => undefined, + defer: pending => void pending, + isCurrent: () => true, + } + ) + ).resolves.toMatchObject({ state: 'completed' }); + await Promise.resolve(); + expect(stored).toMatchObject([{ state: 'completed', terminalSource: 'operation_result' }]); + expect(request.mock.calls.map(([input]) => input.operation)).toEqual([ + 'session.operation.get', + 'session.operation.ack', + ]); + }); + it('keeps the first canonical result through duplicates and conflicts', async () => { const dispatched = recordSessionOperationDispatch(messages(), authorization); if (!dispatched) throw new Error('Failed to create dispatch proof'); diff --git a/services/cloud-agent-next/src/sandbox-session/session-operation.ts b/services/cloud-agent-next/src/sandbox-session/session-operation.ts index eb92488123..27e0241c79 100644 --- a/services/cloud-agent-next/src/sandbox-session/session-operation.ts +++ b/services/cloud-agent-next/src/sandbox-session/session-operation.ts @@ -1,37 +1,39 @@ +import type { DORetryScope } from '@kilocode/worker-utils'; +import type { SandboxControlOutboundRequest } from '../sandbox-control/socket.js'; +import type { EventQueries } from '../session/queries/index.js'; +import type { StoredEvent } from '../websocket/types.js'; import { SANDBOX_CONTROL_ATTACH_TIMEOUT_MS, - SANDBOX_CONTROL_OUTCOME_TIMEOUT_MS, SANDBOX_CONTROL_REQUEST_TIMEOUT_MS, + sameSessionOperation, sessionAttachPayloadSchema, sessionAttachResultSchema, - sessionOperationAuthorizationSchema, sessionOperationAckSchema, + sessionOperationAuthorizationSchema, sessionOperationExpiresAt, sessionOperationLookupResultSchema, sessionPromptPayloadSchema, sessionPromptResultSchema, - sameSessionOperation, + type ControlError, type ResponseFrame, type SessionOperationAck, type SessionOperationAuthorization, type SessionOperationDelivery, } from '../shared/sandbox-control-protocol.js'; -import type { SandboxControlOutboundRequest } from '../sandbox-control/socket.js'; -import type { EventQueries } from '../session/queries/index.js'; -import type { StoredEvent } from '../websocket/types.js'; import { applySessionOperationResult, completeSessionOperationAttachment, recordSessionOperationDispatch, type SessionMessageRecord, } from './session-message-queue.js'; -import { persistSandboxControlSessionEvent } from './sandbox-control-event.js'; import { applyControlPlanePreparingEvent } from './control-plane-preparing.js'; +import { persistSandboxControlSessionEvent } from './sandbox-control-event.js'; import { ControlRequestError, controlRequestResult, withDeliveryDeadline, } from './control-dispatch.js'; +import { persistSessionOperationDelivery } from './session-delivery.js'; type OperationMessages = { read: () => SessionMessageRecord[]; @@ -39,152 +41,213 @@ type OperationMessages = { }; export type SessionOperationEffects = { - request: (input: SandboxControlOutboundRequest) => Promise; + request: (input: SandboxControlOutboundRequest, scope: DORetryScope) => Promise; persistResult: (delivery: SessionOperationDelivery) => Promise; - isDispatchCurrent: () => boolean; - isMaintenanceCurrent: () => boolean; + assertAdmission: () => void; + assertScope: () => void; + defer: (pending: Promise) => void; }; +type RunningOperation = Extract< + ReturnType, + { state: 'running' } +>; +type CompletedOperation = Extract< + ReturnType, + { state: 'completed' } +>; +type UncertainOperation = { + state: 'uncertain'; + reason: 'missing' | 'unverified' | 'transport'; + error?: unknown; +}; +type RejectedOperation = { state: 'rejected'; error: ControlError }; +export type SessionOperationObservation = + | RunningOperation + | CompletedOperation + | UncertainOperation + | RejectedOperation; export type SessionOperationDispatch = | { state: 'response'; result: unknown } - | { state: 'running' } - | { state: 'completed' }; + | { state: 'completed'; result: unknown } + | RunningOperation + | UncertainOperation + | RejectedOperation; -function uncertainOperation(reason: string): ControlRequestError { - return new ControlRequestError({ code: 'runtime_unhealthy', message: reason, retryable: false }); +function rejectedOrUncertain(error: unknown): RejectedOperation | UncertainOperation { + return error instanceof ControlRequestError + ? { + state: 'rejected', + error: { code: error.code, message: error.message, retryable: error.retryable }, + } + : { state: 'uncertain', reason: 'transport', error }; +} + +function rejectedBeforeAdmission(error: unknown): error is ControlRequestError { + return error instanceof ControlRequestError && error.admission === 'not-admitted'; +} + +export async function reconcileSessionOperation( + original: SessionOperationAuthorization, + deadlineAt: number, + effects: SessionOperationEffects +): Promise { + const authorization = sessionOperationAuthorizationSchema.parse(original); + const operationDeadlineAt = Math.min(deadlineAt, sessionOperationExpiresAt(authorization)); + const scope = { deadlineAt: operationDeadlineAt, assertCurrent: effects.assertScope }; + try { + effects.assertScope(); + const response = await withDeliveryDeadline( + () => + effects.request( + { + operation: 'session.operation.get', + session: authorization.session, + payload: authorization, + expectedWrapperInstanceId: authorization.wrapperInstanceId, + timeoutMs: Math.max( + 1, + Math.min(SANDBOX_CONTROL_REQUEST_TIMEOUT_MS, operationDeadlineAt - Date.now()) + ), + deadlineAt: operationDeadlineAt, + }, + scope + ), + operationDeadlineAt + ); + effects.assertScope(); + if (Date.now() >= operationDeadlineAt) throw new Error('Operation observation expired'); + const lookup = sessionOperationLookupResultSchema.parse(controlRequestResult(response)); + if (lookup.state === 'missing') return { state: 'uncertain', reason: 'missing' }; + const observed = + lookup.state === 'completed' ? lookup.delivery.authorization : lookup.authorization; + if (!sameSessionOperation(observed, authorization)) + throw new Error('Operation observation identity changed'); + if (lookup.state === 'completed') { + if ( + (await persistSessionOperationDelivery(lookup.delivery, operationDeadlineAt, effects)) === + 'unverified' + ) + return { state: 'uncertain', reason: 'unverified' }; + } + return lookup; + } catch (error) { + return rejectedOrUncertain(error); + } } export async function dispatchSessionOperation( input: { authorization: SessionOperationAuthorization; payload: unknown }, messages: OperationMessages, - effects: SessionOperationEffects + effects: SessionOperationEffects & { isCurrent: () => boolean } ): Promise { const authorization = sessionOperationAuthorizationSchema.parse(input.authorization); const kind = authorization.operation === 'session.attach' ? 'attach' : 'prompt'; - const timeoutMs = - kind === 'attach' ? SANDBOX_CONTROL_ATTACH_TIMEOUT_MS : SANDBOX_CONTROL_REQUEST_TIMEOUT_MS; - const assertDispatchCurrent = () => { - if (!effects.isDispatchCurrent() || Date.now() >= authorization.dispatchDeadlineAt) - throw uncertainOperation('Session operation dispatch authority expired'); + const deadlineAt = authorization.dispatchDeadlineAt; + const current = () => effects.isCurrent() && Date.now() < deadlineAt; + const assertAdmissionCurrent = () => { + effects.assertAdmission(); + if (!current()) throw new Error('Session delivery is no longer authorized'); }; - const assertMaintenanceCurrent = () => { - if (!effects.isMaintenanceCurrent() || Date.now() >= sessionOperationExpiresAt(authorization)) - throw uncertainOperation('Session operation maintenance authority expired'); + const assertDispatchedCurrent = () => { + effects.assertScope(); + if (!current()) throw new Error('Session delivery is no longer authorized'); }; - const existing = messages.read().find(message => message.messageId === authorization.messageId); - const proof = existing?.operations?.[kind]; - if ( - proof && - !sameSessionOperation( - sessionOperationAuthorizationSchema.parse(proof.authorization), - authorization - ) - ) - throw uncertainOperation('Session operation authorization changed'); - - if (proof?.dispatched) { - assertMaintenanceCurrent(); - const lookup = sessionOperationLookupResultSchema.parse( - controlRequestResult( - await withDeliveryDeadline( - () => - effects.request({ - operation: 'session.operation.get', - session: authorization.session, - payload: authorization, - expectedWrapperInstanceId: authorization.wrapperInstanceId, - deadlineAt: sessionOperationExpiresAt(authorization), - timeoutMs: SANDBOX_CONTROL_REQUEST_TIMEOUT_MS, - }), - sessionOperationExpiresAt(authorization), - SANDBOX_CONTROL_REQUEST_TIMEOUT_MS - ) - ) - ); - assertMaintenanceCurrent(); - if (lookup.state === 'missing') - throw uncertainOperation('Original session operation is missing'); - if (lookup.state === 'running') { - if (!sameSessionOperation(lookup.authorization, authorization)) - throw uncertainOperation('Original session operation identity changed'); - return { state: 'running' }; + const record = (dispatched: boolean) => { + if (!current()) return false; + const next = recordSessionOperationDispatch(messages.read(), authorization, dispatched); + return next !== undefined && messages.commit(next); + }; + try { + const message = messages.read().find(item => item.messageId === authorization.messageId); + const proof = message?.operations?.[kind]; + if (proof && !sameSessionOperation(proof.authorization, authorization)) + throw new Error('Original operation authorization changed'); + if (proof?.dispatched) { + effects.assertScope(); + const lookup = await reconcileSessionOperation( + authorization, + sessionOperationExpiresAt(authorization), + effects + ); + if (lookup.state !== 'completed') return lookup; + if (!lookup.delivery.result.ok) + return { state: 'rejected', error: lookup.delivery.result.error }; + if (kind === 'attach') { + const completed = completeSessionOperationAttachment(messages.read(), authorization); + if (!completed || !messages.commit(completed)) + return { state: 'uncertain', reason: 'unverified' }; + } + return { state: 'completed', result: lookup.delivery.result.result }; } - if (!sameSessionOperation(lookup.delivery.authorization, authorization)) - throw uncertainOperation('Original session operation identity changed'); - const ack = await effects.persistResult(lookup.delivery); - if (!ack) throw uncertainOperation('Original session operation result was not verified'); - assertMaintenanceCurrent(); - controlRequestResult( - await withDeliveryDeadline( - () => - effects.request({ - operation: 'session.operation.ack', - session: authorization.session, - payload: ack, - expectedWrapperInstanceId: authorization.wrapperInstanceId, - deadlineAt: Math.min( - sessionOperationExpiresAt(authorization), - lookup.delivery.completedAt + SANDBOX_CONTROL_OUTCOME_TIMEOUT_MS - ), - timeoutMs: SANDBOX_CONTROL_REQUEST_TIMEOUT_MS, - }), - Math.min( - sessionOperationExpiresAt(authorization), - lookup.delivery.completedAt + SANDBOX_CONTROL_OUTCOME_TIMEOUT_MS - ), - SANDBOX_CONTROL_REQUEST_TIMEOUT_MS - ) + assertAdmissionCurrent(); + const payload = structuredClone( + kind === 'attach' + ? sessionAttachPayloadSchema.parse(input.payload) + : sessionPromptPayloadSchema.parse(input.payload) ); - if (kind === 'attach') { - const attached = sessionAttachResultSchema.safeParse( - lookup.delivery.result.ok ? lookup.delivery.result.result : undefined + if (!record(true)) throw new Error('Session operation proof could not be persisted'); + assertDispatchedCurrent(); + try { + const timeoutMs = + kind === 'attach' ? SANDBOX_CONTROL_ATTACH_TIMEOUT_MS : SANDBOX_CONTROL_REQUEST_TIMEOUT_MS; + const response = await withDeliveryDeadline( + () => + effects.request( + { + operation: authorization.operation, + authorization, + session: authorization.session, + expectedWrapperInstanceId: authorization.wrapperInstanceId, + payload, + timeoutMs, + deadlineAt, + }, + { deadlineAt, assertCurrent: assertDispatchedCurrent } + ), + deadlineAt, + timeoutMs ); - if (!attached.success) throw uncertainOperation('Original attachment result is invalid'); - const completed = completeSessionOperationAttachment(messages.read(), authorization); - if (!completed || !messages.commit(completed)) - throw uncertainOperation('Original attachment result was not persisted'); - return { state: 'response', result: attached.data }; + const result = controlRequestResult(response); + assertDispatchedCurrent(); + if (kind === 'attach') { + const attached = sessionAttachResultSchema.parse(result); + const completed = completeSessionOperationAttachment(messages.read(), authorization); + if (!completed || !messages.commit(completed)) + return { state: 'uncertain', reason: 'unverified' }; + return { state: 'response', result: attached }; + } + const prompt = sessionPromptResultSchema.parse(result); + if (prompt.messageId !== authorization.messageId) + throw new Error('Prompt response message identity mismatch'); + return { state: 'response', result: prompt }; + } catch (error) { + if (rejectedBeforeAdmission(error)) record(false); + return rejectedOrUncertain(error); } - return { state: 'completed' }; + } catch (error) { + return rejectedOrUncertain(error); } +} - const payload = - kind === 'attach' - ? sessionAttachPayloadSchema.parse(input.payload) - : sessionPromptPayloadSchema.parse(input.payload); - assertDispatchCurrent(); - const recorded = recordSessionOperationDispatch(messages.read(), authorization); - if (!recorded || !messages.commit(recorded)) - throw uncertainOperation('Session operation dispatch proof was not persisted'); - assertDispatchCurrent(); - const result = controlRequestResult( - await withDeliveryDeadline( - () => - effects.request({ - operation: authorization.operation, - authorization, - session: authorization.session, - expectedWrapperInstanceId: authorization.wrapperInstanceId, - payload, - deadlineAt: authorization.dispatchDeadlineAt, - timeoutMs, - }), - authorization.dispatchDeadlineAt, - timeoutMs - ) - ); - assertDispatchCurrent(); - if (kind === 'attach') { - const attached = sessionAttachResultSchema.parse(result); - const completed = completeSessionOperationAttachment(messages.read(), authorization); - if (!completed || !messages.commit(completed)) - throw uncertainOperation('Session attachment response was not persisted'); - return { state: 'response', result: attached }; - } - const prompt = sessionPromptResultSchema.parse(result); - if (prompt.messageId !== authorization.messageId) - throw uncertainOperation('Prompt response message identity mismatch'); - return { state: 'response', result: prompt }; +export function operationDispatchError( + result: Exclude +): ControlRequestError { + if (result.state === 'rejected') return new ControlRequestError(result.error); + if (result.state === 'running') + return new ControlRequestError({ + code: 'session_busy', + message: 'Original preparation is running', + retryable: true, + }); + return new ControlRequestError({ + code: 'runtime_unhealthy', + message: + result.reason === 'transport' + ? 'Operation admission acknowledgement is unconfirmed' + : 'Original operation outcome is unconfirmed', + retryable: result.reason === 'transport', + }); } export function commitSessionOperationResult(input: { diff --git a/services/cloud-agent-next/src/shared/sandbox-control-protocol.ts b/services/cloud-agent-next/src/shared/sandbox-control-protocol.ts index 2e55fb7123..530617d798 100644 --- a/services/cloud-agent-next/src/shared/sandbox-control-protocol.ts +++ b/services/cloud-agent-next/src/shared/sandbox-control-protocol.ts @@ -114,6 +114,7 @@ export const controlErrorSchema = z.object({ code: z.string().min(1), message: z.string(), retryable: z.boolean(), + admission: z.literal('not-admitted').optional(), }); export const requestFrameSchema = z.object({ 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 c6523db0ec..3b5340dc6b 100644 --- a/services/cloud-agent-next/test/integration/sandbox-control.test.ts +++ b/services/cloud-agent-next/test/integration/sandbox-control.test.ts @@ -353,6 +353,7 @@ function sendHello( providerInstanceId?: string; wrapperInstanceId?: string; sessionOperationResults?: boolean; + nativeRuntimeRetirement?: boolean; } = {} ): void { ws.send( @@ -367,6 +368,14 @@ function sendHello( ...(identity.sessionOperationResults ? { capabilities: { sessionOperationResults: true } } : {}), + ...(identity.nativeRuntimeRetirement + ? { + capabilities: { + sessionOperationResults: true, + nativeRuntimeRetirement: true, + }, + } + : {}), ...(identity.wrapperInstanceId ? { wrapperInstanceId: identity.wrapperInstanceId } : {}), }, }) @@ -380,6 +389,7 @@ async function completeHello( providerInstanceId?: string; wrapperInstanceId?: string; sessionOperationResults?: boolean; + nativeRuntimeRetirement?: boolean; } = {} ): Promise { sendHello(ws, requestId, identity); @@ -423,7 +433,10 @@ type TerminalRuntimeFixture = { wrapperInstanceId?: string; }; -async function initializeTerminalRuntime(fixture: TerminalRuntimeFixture) { +async function initializeTerminalRuntime( + fixture: TerminalRuntimeFixture, + capabilities: { sessionOperationResults?: boolean; nativeRuntimeRetirement?: boolean } = {} +) { const credential = generateSandboxCredential(); await seedCredential(credential, fixture.sandboxId); const control = env.SANDBOX_CONTROL.getByName(fixture.sandboxId); @@ -451,6 +464,7 @@ async function initializeTerminalRuntime(fixture: TerminalRuntimeFixture) { await completeHello(socket, `hello_${fixture.sandboxId}`, { providerInstanceId: providerRef, ...(fixture.wrapperInstanceId ? { wrapperInstanceId: fixture.wrapperInstanceId } : {}), + ...capabilities, }); return { control, credential, socket, providerRef, ...provider }; } @@ -4912,7 +4926,7 @@ describe('SandboxControl recovery watchdogs', () => { wrapperInstanceId, reason: 'preparation_interrupted', }) - ).resolves.toEqual({ quarantined: true }); + ).resolves.toEqual({ quarantined: true, disposition: 'physical_stopping' }); await runInDurableObject(control, async (_instance, state) => { expect(await state.storage.getAlarm()).toBe(repairedAt); }); @@ -8372,9 +8386,12 @@ describe('SandboxSession control-plane regressions', () => { providerRef: cloudflareRef(fixture.sandboxId), }); await runInDurableObject(session, (_instance, state) => { - expect(state.storage.kv.get('pending_runtime_cleanup')).toBeUndefined(); + expect(state.storage.kv.get('pending_runtime_cleanup')).toEqual(cleanup); }); - expect(acquisitions.at(-1)?.acquisition).toEqual(acquisitionB); + expect((await admissionState(session)).messages[1]).toEqual(b); + expect(acquisitions.map(input => input.acquisition?.id)).toEqual([ + preparing.preparationAttemptId, + ]); expect(provider.create).not.toHaveBeenCalled(); await runInDurableObject(control, async (_instance, state) => { expect(await state.storage.get('acquisition_receipts')).toEqual([ @@ -8389,6 +8406,10 @@ describe('SandboxSession control-plane regressions', () => { providerRef: null, }); await expect(runDurableObjectAlarm(session)).resolves.toBe(true); + await runInDurableObject(session, (_instance, state) => { + expect(state.storage.kv.get('pending_runtime_cleanup')).toBeUndefined(); + }); + expect(acquisitions.at(-1)?.acquisition).toEqual(acquisitionB); expect(provider.launch).toHaveBeenCalledTimes(1); const launch = provider.launch.mock.calls[0]; if (!launch) throw new Error('Expected one replacement launch for B'); @@ -8414,7 +8435,7 @@ describe('SandboxSession control-plane regressions', () => { wrapperInstanceId: fixture.wrapperInstanceId, reason: 'late_cancelled_runtime_cleanup', }) - ).resolves.toEqual({ quarantined: false }); + ).resolves.toEqual({ quarantined: false, disposition: 'wrapper_replaced' }); expect((await admissionState(session)).messages).toMatchObject([ { messageId: 'msg_ffffffffffff00000000000001', state: 'cancelled' }, { @@ -8424,12 +8445,8 @@ describe('SandboxSession control-plane regressions', () => { }, ]); const continuations = acquisitions.filter(input => input.acquisition?.id === acquisitionB.id); - expect(continuations).toHaveLength(3); - expect(continuations.map(input => input.acquisition)).toEqual([ - acquisitionB, - acquisitionB, - acquisitionB, - ]); + expect(continuations).toHaveLength(2); + expect(continuations.map(input => input.acquisition)).toEqual([acquisitionB, acquisitionB]); expect(continuations.every(input => input.allowCreate === undefined)).toBe(true); expect( replacementRequests @@ -8450,6 +8467,193 @@ describe('SandboxSession control-plane regressions', () => { } }); + it('retries a lost completed native cleanup after Session reload without selecting N2', async () => { + const { fixture, session: originalSession } = messageFixture(); + let session = originalSession; + const nativeRuntimeId = '11111111-1111-4111-8111-111111111111'; + const replacementRuntimeId = '22222222-2222-4222-8222-222222222222'; + const { control, socket, provider } = await initializeTerminalRuntime(fixture, { + nativeRuntimeRetirement: true, + }); + let dropCleanupResponse = true; + const requests: RequestFrame[] = []; + socket.addEventListener('message', event => { + const request = requestFrameSchema.parse(JSON.parse(String(event.data))); + requests.push(request); + let result: unknown; + switch (request.operation) { + case 'session.attach': + result = { + attached: true, + nativeRuntimeId: + requests.filter(candidate => candidate.operation === 'session.attach').length === 1 + ? nativeRuntimeId + : replacementRuntimeId, + }; + break; + case 'session.prompt': + result = { + messageId: sessionPromptPayloadSchema.parse(request.payload).messageId, + status: 'accepted', + }; + break; + case 'session.operation.get': + result = { state: 'missing' }; + break; + case 'session.sync': + result = { status: { type: 'idle' }, questions: [], permissions: [] }; + break; + case 'session.abort': + result = { + status: 'aborted', + quiescent: true, + runtimeRetired: true, + nativeRuntimeId, + }; + break; + default: + throw new Error(`Unexpected control request: ${request.operation}`); + } + socket.send( + JSON.stringify({ type: 'response', requestId: request.requestId, ok: true, result }) + ); + }); + await runInDurableObject(control, instance => { + const prototype = Object.getPrototypeOf(instance) as typeof instance; + const quarantine = instance.quarantineRuntime.bind(instance); + vi.spyOn(prototype, 'quarantineRuntime').mockImplementation(async input => { + const result = await quarantine(input); + if (dropCleanupResponse) throw new Error('lost cleanup response'); + return result; + }); + }); + try { + signalWrapperReady(socket); + await waitForWrapperReady(fixture); + await session.createSessionWithInitialAdmission({ + identity: { sessionId: fixture.sessionId, userId: fixture.ownerId }, + auth: { kiloSessionId: ROOT_ID, kilocodeToken: KILO_TOKEN }, + agent: agentA, + workspace: { sandboxId: fixture.sandboxId, workspacePath: '/workspace/terminal' }, + message: { + initialTurn: { type: 'prompt', messageId: INITIAL_MESSAGE_ID, prompt: 'A' }, + }, + }); + await waitForAccepted(session, INITIAL_MESSAGE_ID); + const attached = sessionOperationAuthorizationSchema.parse({ + operation: 'session.attach', + operationId: '33333333-3333-4333-8333-333333333333', + messageId: INITIAL_MESSAGE_ID, + session: { + sessionId: fixture.sessionId, + kiloSessionId: ROOT_ID, + directory: '/workspace/terminal', + }, + wrapperInstanceId: fixture.wrapperInstanceId, + dispatchDeadlineAt: Date.now() + SESSION_DELIVERY_TIMEOUT_MS, + }); + await runInDurableObject(session, (_instance, state) => { + state.storage.kv.put('native_runtime_fence', { + sandboxId: fixture.sandboxId, + wrapperInstanceId: fixture.wrapperInstanceId, + nativeRuntimeId, + attachmentEpoch: 1, + authorization: attached, + }); + }); + await expect(control.listRoutes()).resolves.toEqual([ + expect.objectContaining({ nativeRuntimeId }), + ]); + await runInDurableObject(session, (_instance, state) => { + const messages = state.storage.kv.get('session_messages') ?? []; + state.storage.kv.put( + 'session_messages', + messages.map(message => + message.messageId === INITIAL_MESSAGE_ID + ? { + ...message, + acceptedAt: Date.now() - DEADLINE_MS.acceptedOverdue, + lastActivityAt: Date.now() - DEADLINE_MS.acceptedOverdue, + } + : message + ) + ); + }); + await expect(runDurableObjectAlarm(session)).resolves.toBe(true); + await vi.waitFor(async () => { + await expect(control.listRoutes()).resolves.toEqual([ + expect.not.objectContaining({ nativeRuntimeId }), + ]); + }); + await runInDurableObject(session, (_instance, state) => { + expect(state.storage.kv.get('pending_runtime_cleanup')).toMatchObject({ + nativeRuntimeId, + authorization: attached, + }); + }); + await expect(control.getPhysicalRecord()).resolves.toMatchObject({ + state: 'running', + stopTombstone: null, + }); + expect(requests.filter(request => request.operation === 'session.abort')).toHaveLength(1); + + const replacementAuthorization = { + operation: 'session.attach' as const, + operationId: '33333333-3333-4333-8333-333333333333', + messageId: 'msg_018f1e2d3c4bAbCdEfGhIjKlMo', + session: { + sessionId: fixture.sessionId, + kiloSessionId: ROOT_ID, + directory: '/workspace/terminal', + }, + wrapperInstanceId: fixture.wrapperInstanceId ?? '', + dispatchDeadlineAt: Date.now() + SESSION_DELIVERY_TIMEOUT_MS, + }; + await expect( + control.request({ + operation: 'session.attach', + authorization: replacementAuthorization, + session: replacementAuthorization.session, + expectedWrapperInstanceId: fixture.wrapperInstanceId, + payload: {}, + }) + ).resolves.toMatchObject({ ok: true, result: { nativeRuntimeId: replacementRuntimeId } }); + await expect(control.listRoutes()).resolves.toEqual([ + expect.objectContaining({ nativeRuntimeId: replacementRuntimeId }), + ]); + const beforeRetry = await runInDurableObject(control, async (_instance, state) => ({ + deadlines: await loadDeadlines(state.storage), + alarmAt: await state.storage.getAlarm(), + })); + + await expect( + runInDurableObject(session, (_instance, state) => state.abort('reload after lost cleanup')) + ).rejects.toThrow('reload after lost cleanup'); + session = env.SANDBOX_SESSION.getByName(`${fixture.ownerId}:${fixture.sessionId}`); + dropCleanupResponse = false; + await expect(runDurableObjectAlarm(session)).resolves.toBe(true); + + await runInDurableObject(session, (_instance, state) => { + expect(state.storage.kv.get('pending_runtime_cleanup')).toBeUndefined(); + }); + expect(requests.filter(request => request.operation === 'session.abort')).toHaveLength(1); + expect(provider.stop).not.toHaveBeenCalled(); + await expect(control.getPhysicalRecord()).resolves.toMatchObject({ + state: 'running', + stopTombstone: null, + }); + await expect(control.listRoutes()).resolves.toEqual([ + expect.objectContaining({ nativeRuntimeId: replacementRuntimeId }), + ]); + await runInDurableObject(control, async (_instance, state) => { + expect(await loadDeadlines(state.storage)).toEqual(beforeRetry.deadlines); + expect(await state.storage.getAlarm()).toBe(beforeRetry.alarmAt); + }); + } finally { + socket.close(); + } + }); + it.each([ { sandboxProvider: 'cloudflare', operation: 'session.attach' }, { sandboxProvider: 'cloudflare', operation: 'session.prompt' }, diff --git a/services/cloud-agent-next/test/integration/session-observation.test.ts b/services/cloud-agent-next/test/integration/session-observation.test.ts index 06195711d1..aa2d01621e 100644 --- a/services/cloud-agent-next/test/integration/session-observation.test.ts +++ b/services/cloud-agent-next/test/integration/session-observation.test.ts @@ -44,7 +44,7 @@ async function fixture(instance: SandboxSession, state: DurableObjectState) { const control = { getStatus: vi.fn(async () => ({ connection: 'ready', physical: 'running', wrapperInstanceId })), request: vi.fn(() => pending.promise), - quarantineRuntime: vi.fn(async () => ({ quarantined: true })), + quarantineRuntime: vi.fn(async () => ({ quarantined: true, disposition: 'physical_stopping' })), }; const originalEnv = instance['env']; Object.assign(instance, { diff --git a/services/cloud-agent-next/wrapper/src/control/control-handler-result.ts b/services/cloud-agent-next/wrapper/src/control/control-handler-result.ts index dd112e3e5e..20a60042aa 100644 --- a/services/cloud-agent-next/wrapper/src/control/control-handler-result.ts +++ b/services/cloud-agent-next/wrapper/src/control/control-handler-result.ts @@ -2,12 +2,12 @@ import type { ControlError } from '../../../src/shared/sandbox-control-protocol. export type ControlHandlerResult = | { ok: true; result: unknown; admission?: never } - | { ok: false; error: ControlError; admission?: 'not-admitted' }; + | { ok: false; error: ControlError; admission?: never }; export function rejectBeforeAdmission( code: ControlError['code'], message: string, retryable: boolean ): ControlHandlerResult { - return { ok: false, error: { code, message, retryable }, admission: 'not-admitted' }; + return { ok: false, error: { code, message, retryable, admission: 'not-admitted' } }; } 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 a66bdf4bde..de5d1de2e8 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 @@ -187,7 +187,10 @@ describe('operation admission and lookup', () => { ).toMatchObject({ ok: true, result: { state: 'completed' } }); expect( await handleControlRequest('session.prompt', session, promptPayload, handlerDeps, first) - ).toMatchObject({ ok: false, error: { code: 'session_busy' } }); + ).toMatchObject({ + ok: false, + error: { code: 'session_busy', admission: 'not-admitted' }, + }); setSystemTime(first.dispatchDeadlineAt + SANDBOX_CONTROL_OUTCOME_TIMEOUT_MS + 1); pruneControlOperations(handlerDeps); expect(handlerDeps.operations.counts().retained).toBe(0);