diff --git a/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts index 86bd429fc1..2f10846b47 100644 --- a/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts +++ b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts @@ -28,6 +28,8 @@ import type { CloudAgentSession } from '../persistence/CloudAgentSession.js'; import type { Env } from '../types.js'; import { withDORetry } from '../utils/do-retry.js'; import { resolveSessionStub } from '../sandbox-session/session-stub.js'; +import { sessionPlaneFromId } from '../session-plane.js'; +import { interruptControlSession } from '../router/control-plane-session.js'; import { preflightAndAdmitPromptMessage } from '../session/queue-message.js'; import { parseBasicKiloPrompt } from './basic-prompt.js'; import { @@ -881,6 +883,19 @@ async function defaultInterruptPrompt(params: { userId: string; cloudAgentSessionId: string; }): Promise>> { + if (sessionPlaneFromId(params.cloudAgentSessionId) === 'control') { + const receipt = await interruptControlSession({ + env: params.env, + ownerId: params.userId, + sessionId: params.cloudAgentSessionId, + }); + return receipt?.state === 'rejected' + ? { success: false, message: receipt.message } + : { + success: receipt !== undefined, + ...(receipt ? {} : { message: 'No session work to interrupt' }), + }; + } return withDORetry< DurableObjectStub, Awaited> diff --git a/services/cloud-agent-next/src/persistence/SandboxControl.ts b/services/cloud-agent-next/src/persistence/SandboxControl.ts index 85ee66180b..3172a2477b 100644 --- a/services/cloud-agent-next/src/persistence/SandboxControl.ts +++ b/services/cloud-agent-next/src/persistence/SandboxControl.ts @@ -35,6 +35,11 @@ import { type SandboxControlSocketHandler, } from '../sandbox-control/socket.js'; import { SandboxControlConnectionError } from '../sandbox-control/waiters.js'; +import { + hasScopedStopMaintenanceFields, + parseScopedStopMaintenance, + stopAbortWirePayload, +} from '../sandbox-control/scoped-stop-maintenance.js'; import { createSessionForwarding, SessionForwardingError, @@ -52,6 +57,7 @@ import { sessionOperationAckSchema, sessionOperationAuthorizationSchema, sessionOperationExpiresAt, + sessionAbortPayloadSchema, sessionRequestIdentitySchema, wrapperInstanceIdSchema, type ResponseFrame, @@ -529,9 +535,17 @@ export class SandboxControl extends DurableObject { async request(input: SandboxControlOutboundRequest): Promise { await this.ensureOperationalInitialized(); if (input.operation === 'session.git.summary') return this.requestWorktreeChanges(input); + const scopedStop = + input.operation === 'session.abort' ? parseScopedStopMaintenance(input.payload) : undefined; + if ( + input.operation === 'session.abort' && + hasScopedStopMaintenanceFields(input.payload) && + scopedStop === undefined + ) + throw new Error('Invalid scoped Stop maintenance request'); const maintenance = input.operation === 'session.operation.get' || input.operation === 'session.operation.ack'; - if (!maintenance) await this.assertRequestWorktreeAdmission(input); + if (!maintenance && !scopedStop) await this.assertRequestWorktreeAdmission(input); if (input.operation === 'worktree.delete' || input.operation === 'worktree.prepareDeletion') { throw new Error('Worktree cleanup requires the deletion coordinator'); } @@ -578,7 +592,8 @@ export class SandboxControl extends DurableObject { Date.now() >= sessionOperationExpiresAt(maintenanceAuthorization.data)) ) throw new Error('Invalid session operation maintenance authorization'); - const runtime = maintenance + const usesMaintenanceChannel = maintenance || scopedStop !== undefined; + const runtime = usesMaintenanceChannel ? this.socketHandler.getConnectionIdentity() : this.readyWrapperRuntime(); if (!runtime) throw new Error('Sandbox runtime is not ready'); @@ -594,20 +609,34 @@ export class SandboxControl extends DurableObject { ) throw new Error('Sandbox wrapper runtime changed'); const isCurrent = () => { - const current = maintenance + const current = usesMaintenanceChannel ? this.socketHandler.getConnectionIdentity() : this.readyWrapperRuntime(); return current !== null && this.sameConnection(current, runtime); }; const physical = await loadPhysicalRecord(this.ctx.storage); if ( - (!maintenance && physical.state !== 'running') || - (!maintenance && physical.stopTombstone) || + ((!maintenance || scopedStop !== undefined) && physical.state !== 'running') || + ((!maintenance || scopedStop !== undefined) && physical.stopTombstone) || physical.providerRef !== runtime.providerInstanceId || !isCurrent() ) { throw new Error('Sandbox runtime is not ready'); } + if (scopedStop) { + const session = sessionRequestIdentitySchema.safeParse(input.session); + if (!session.success || expectedWrapperInstanceId === undefined) + throw new Error('Scoped Stop identity is required'); + const route = (await loadRouteTable(this.ctx.storage)).get(session.data.sessionId); + if ( + !route || + route.kiloSessionId !== session.data.kiloSessionId || + route.directory !== session.data.directory || + runtime.wrapperInstanceId !== expectedWrapperInstanceId + ) + throw new Error('Scoped Stop target is stale'); + this.assertWorktreeAdmission(route.worktreeId); + } if (input.operation === 'session.attach' || input.operation === 'session.prompt') { const payload = parseOperationPayload(input.operation, input.payload); if (!payload.ok) throw new Error(payload.error.message); @@ -664,8 +693,16 @@ export class SandboxControl extends DurableObject { if (!isCurrent()) throw new Error('Sandbox wrapper runtime changed'); }); } - if (!maintenance) await this.assertRequestWorktreeAdmission(input); + if (!usesMaintenanceChannel) await this.assertRequestWorktreeAdmission(input); if (!isCurrent()) throw new Error('Sandbox wrapper runtime changed'); + if (scopedStop) { + const payload = sessionAbortPayloadSchema.parse(input.payload); + return this.socketHandler.sendRequest({ + ...input, + payload: stopAbortWirePayload(payload, this.socketHandler.supportsScopedStopAbort()), + deadlineAt: scopedStop.cleanupDeadlineAt, + }); + } return this.socketHandler.sendRequest(input); } diff --git a/services/cloud-agent-next/src/router.test.ts b/services/cloud-agent-next/src/router.test.ts index 5a0bbec55e..624012bd6b 100644 --- a/services/cloud-agent-next/src/router.test.ts +++ b/services/cloud-agent-next/src/router.test.ts @@ -907,9 +907,20 @@ describe('router sessionId validation', () => { it('routes workspace_ interrupts to SANDBOX_SESSION', async () => { const sessionId: SessionId = 'workspace_12345678-1234-1234-1234-123456789abc'; + const controlStub = { + ...mockSessionStub, + getControlState: vi.fn().mockResolvedValue({ + version: 1, + scope: { sandboxId: 'sandbox_1' }, + targets: [{ messageId: 'message_1' }], + }), + interruptExecution: vi + .fn() + .mockImplementation(request => Promise.resolve({ ...request, state: 'confirmed' })), + }; const sandboxSession = { idFromName: vi.fn((id: string) => ({ id })), - get: vi.fn(() => mockSessionStub), + get: vi.fn(() => controlStub), }; mockContext.env.SANDBOX_SESSION = sandboxSession as unknown as TRPCContext['env']['SANDBOX_SESSION']; @@ -927,7 +938,9 @@ describe('router sessionId validation', () => { expect(result.success).toBe(true); expect(sandboxSession.idFromName).toHaveBeenCalledWith(`test-user-123:${sessionId}`); expect(cloudAgentSession.idFromName).not.toHaveBeenCalled(); - expect(mockSessionStub.interruptExecution).toHaveBeenCalled(); + expect(controlStub.interruptExecution).toHaveBeenCalledWith( + expect.objectContaining({ targets: [{ messageId: 'message_1' }] }) + ); }); }); diff --git a/services/cloud-agent-next/src/router/control-plane-session.test.ts b/services/cloud-agent-next/src/router/control-plane-session.test.ts new file mode 100644 index 0000000000..fd533f23e9 --- /dev/null +++ b/services/cloud-agent-next/src/router/control-plane-session.test.ts @@ -0,0 +1,68 @@ +import { describe, expect, it } from 'vitest'; +import { interruptControlSession } from './control-plane-session.js'; + +const OPERATION_ID = '33333333-3333-4333-8333-333333333333'; +const RUNTIME_A = '11111111-1111-4111-8111-111111111111'; +const RUNTIME_B = '22222222-2222-4222-8222-222222222222'; + +describe('interruptControlSession', () => { + it('reuses one captured Stop request after the first Durable Object reply is lost', async () => { + let stateCalls = 0; + const deliveries: unknown[] = []; + let current = { + version: 1 as const, + scope: { sandboxId: 'sandbox-a', wrapperInstanceId: RUNTIME_A }, + targets: [{ messageId: 'a', wrapperInstanceId: RUNTIME_A, executionDeadlineAt: 3_601_000 }], + }; + const getStub = () => ({ + getControlState: async () => { + stateCalls++; + return current; + }, + interruptExecution: async (request: unknown) => { + deliveries.push(structuredClone(request)); + current = { + version: 1, + scope: { sandboxId: 'sandbox-b', wrapperInstanceId: RUNTIME_B }, + targets: [ + { messageId: 'b', wrapperInstanceId: RUNTIME_B, executionDeadlineAt: 3_602_000 }, + ], + }; + return { ...(request as object), state: 'confirmed' }; + }, + }); + + const receipt = await interruptControlSession( + { env: {} as never, ownerId: 'user-a', sessionId: 'workspace-a' }, + { + getStub, + now: 1_000, + operationId: OPERATION_ID, + retry: async (operation, operationName) => { + if (operationName === 'getControlState') return operation(getStub()); + await operation(getStub()); + return operation(getStub()); + }, + } + ); + + expect(stateCalls).toBe(1); + expect(deliveries).toEqual([ + { + version: 1, + operationId: OPERATION_ID, + scope: { sandboxId: 'sandbox-a', wrapperInstanceId: RUNTIME_A }, + targets: [{ messageId: 'a', wrapperInstanceId: RUNTIME_A, executionDeadlineAt: 3_601_000 }], + cleanupDeadlineAt: 11_000, + }, + { + version: 1, + operationId: OPERATION_ID, + scope: { sandboxId: 'sandbox-a', wrapperInstanceId: RUNTIME_A }, + targets: [{ messageId: 'a', wrapperInstanceId: RUNTIME_A, executionDeadlineAt: 3_601_000 }], + cleanupDeadlineAt: 11_000, + }, + ]); + expect(receipt).toMatchObject({ operationId: OPERATION_ID, state: 'confirmed' }); + }); +}); diff --git a/services/cloud-agent-next/src/router/control-plane-session.ts b/services/cloud-agent-next/src/router/control-plane-session.ts new file mode 100644 index 0000000000..ad728f3a25 --- /dev/null +++ b/services/cloud-agent-next/src/router/control-plane-session.ts @@ -0,0 +1,53 @@ +import type { Env } from '../types.js'; +import { + controlSessionStateSchema, + controlStopReceiptSchema, + createControlStopRequest, + type ControlStopReceipt, +} from '../shared/control-plane-session.js'; +import { getSandboxSessionStub } from '../sandbox-session/session-stub.js'; +import { withDORetry } from '../utils/do-retry.js'; + +type ControlStopSession = { + getControlState: () => Promise; + interruptExecution: (request: unknown) => Promise; +}; + +type ControlSessionStopDependencies = { + getStub?: () => ControlStopSession; + retry?: ( + operation: (session: ControlStopSession) => Promise, + operationName: string + ) => Promise; + now?: number; + operationId?: string; +}; + +export async function interruptControlSession( + input: { + env: Pick; + ownerId: string; + sessionId: string; + }, + dependencies: ControlSessionStopDependencies = {} +): Promise { + const stub = + dependencies.getStub ?? + (() => getSandboxSessionStub(input.env, input.ownerId, input.sessionId)); + const retry = + dependencies.retry ?? + ((operation: (session: ControlStopSession) => Promise, operationName: string) => + withDORetry(stub, operation, operationName)); + const state = await retry(session => session.getControlState(), 'getControlState'); + if (!state) return undefined; + const request = createControlStopRequest( + controlSessionStateSchema.parse(state), + dependencies.now, + dependencies.operationId + ); + return retry( + session => + session.interruptExecution(request).then(receipt => controlStopReceiptSchema.parse(receipt)), + 'interruptControlSession' + ); +} diff --git a/services/cloud-agent-next/src/router/handlers/session-management.ts b/services/cloud-agent-next/src/router/handlers/session-management.ts index 6b1ad150d1..11e29333bd 100644 --- a/services/cloud-agent-next/src/router/handlers/session-management.ts +++ b/services/cloud-agent-next/src/router/handlers/session-management.ts @@ -14,6 +14,7 @@ import { import { withDORetry } from '../../utils/do-retry.js'; import { getSandboxSessionStub, resolveSessionStub } from '../../sandbox-session/session-stub.js'; import { sessionPlaneFromId } from '../../session-plane.js'; +import { interruptControlSession } from '../control-plane-session.js'; import { protectedProcedure, publicProcedure, internalApiProtectedProcedure } from '../auth.js'; import { sessionIdSchema, @@ -206,34 +207,45 @@ export function createSessionManagementHandlers() { }; } - // Mark session as interrupted in DO before killing processes (with retry) - // This signals the streaming generator to stop const getStub = () => resolveSessionStub(env, userId, sessionId); await withDORetry(getStub, stub => stub.markAsInterrupted(), 'markAsInterrupted'); - const interruptResult = await withDORetry( - getStub, - stub => stub.interruptExecution(), - 'interruptExecution' - ); - - if (!interruptResult.success) { + const interruptResult = + sessionPlaneFromId(sessionId) === 'control' + ? await interruptControlSession({ env, ownerId: userId, sessionId }) + : await withDORetry( + getStub, + stub => stub.interruptExecution(), + 'interruptExecution' + ); + + const success = + interruptResult !== undefined && + ('success' in interruptResult + ? interruptResult.success + : interruptResult.state !== 'rejected'); + const message = + interruptResult === undefined + ? 'No session work to interrupt' + : 'success' in interruptResult + ? interruptResult.message + : interruptResult.message; + + if (!success) { logger .withFields({ - message: - interruptResult.message ?? - 'No accepted current messages or pending queued messages', + message: message ?? 'No accepted current messages or pending queued messages', }) .info('No accepted current messages or pending queued messages to interrupt'); } logger.info('Session interruption completed'); return { - success: interruptResult.success, - message: interruptResult.success + success, + message: success ? 'Session interruption accepted' - : (interruptResult.message ?? 'No session work to interrupt'), + : (message ?? 'No session work to interrupt'), processesFound: false, }; } catch (error) { diff --git a/services/cloud-agent-next/src/sandbox-control/frames.test.ts b/services/cloud-agent-next/src/sandbox-control/frames.test.ts index 0fe43cfc28..f0445542c9 100644 --- a/services/cloud-agent-next/src/sandbox-control/frames.test.ts +++ b/services/cloud-agent-next/src/sandbox-control/frames.test.ts @@ -43,7 +43,11 @@ describe('sandbox control frames', () => { expect(sandboxHelloResultSchema.parse(helloResult())).toEqual({ protocolVersion: 1, handshakeComplete: true, - capabilities: { kiloVersionHeartbeat: true, sessionOperationResults: true }, + capabilities: { + kiloVersionHeartbeat: true, + sessionOperationResults: true, + scopedStopAbort: true, + }, }); const previous = { protocolVersion: 1, handshakeComplete: true }; expect(sandboxHelloResultSchema.parse(previous)).toEqual(previous); diff --git a/services/cloud-agent-next/src/sandbox-control/frames.ts b/services/cloud-agent-next/src/sandbox-control/frames.ts index c742960086..2a5424004a 100644 --- a/services/cloud-agent-next/src/sandbox-control/frames.ts +++ b/services/cloud-agent-next/src/sandbox-control/frames.ts @@ -181,6 +181,10 @@ export function helloResult(): SandboxHelloResult { return { protocolVersion: SANDBOX_CONTROL_PROTOCOL_VERSION, handshakeComplete: true, - capabilities: { kiloVersionHeartbeat: true, sessionOperationResults: true }, + capabilities: { + kiloVersionHeartbeat: true, + sessionOperationResults: true, + scopedStopAbort: true, + }, }; } 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 f1476d7cb4..2c4d4d1a53 100644 --- a/services/cloud-agent-next/src/sandbox-control/lifecycle.test.ts +++ b/services/cloud-agent-next/src/sandbox-control/lifecycle.test.ts @@ -2266,6 +2266,30 @@ describe('SandboxControl lifecycle boundaries', () => { expect((await h.control.getPhysicalRecord()).providerRef).toBeNull(); }); + it.each([ + { messageId: 'message_1', operationId: 'not-a-uuid', cleanupDeadlineAt: Date.now() + 1_000 }, + { + messageId: 'message_1', + operationId: '33333333-3333-4333-8333-333333333333', + cleanupDeadlineAt: Date.now(), + }, + ])('rejects invalid scoped Stop requests without forwarding them', async payload => { + const h = await harness(); + await h.create(); + const identity = await h.ready(); + + await expect( + h.control.request({ + operation: 'session.abort', + session: ROUTE, + payload, + expectedWrapperInstanceId: identity.wrapperInstanceId, + }) + ).rejects.toThrow('Invalid scoped Stop maintenance request'); + + expect(h.sendRequest).not.toHaveBeenCalled(); + }); + it('rejects malformed cleanup transfers rather than acknowledging a stale runtime', async () => { const h = await harness(); await h.create(); diff --git a/services/cloud-agent-next/src/sandbox-control/scoped-stop-maintenance.test.ts b/services/cloud-agent-next/src/sandbox-control/scoped-stop-maintenance.test.ts new file mode 100644 index 0000000000..3d2010a3b5 --- /dev/null +++ b/services/cloud-agent-next/src/sandbox-control/scoped-stop-maintenance.test.ts @@ -0,0 +1,46 @@ +import { describe, expect, it } from 'vitest'; +import { + hasScopedStopMaintenanceFields, + parseScopedStopMaintenance, + stopAbortWirePayload, +} from './scoped-stop-maintenance.js'; + +const OPERATION_ID = '33333333-3333-4333-8333-333333333333'; + +describe('scoped Stop maintenance', () => { + it('requires a current immutable cleanup bound', () => { + expect( + parseScopedStopMaintenance( + { messageId: 'a', operationId: OPERATION_ID, cleanupDeadlineAt: 11_000 }, + 1_000 + ) + ).toEqual({ messageId: 'a', operationId: OPERATION_ID, cleanupDeadlineAt: 11_000 }); + expect( + parseScopedStopMaintenance( + { messageId: 'a', operationId: OPERATION_ID, cleanupDeadlineAt: 1_000 }, + 1_000 + ) + ).toBeUndefined(); + expect( + parseScopedStopMaintenance( + { messageId: 'a', operationId: OPERATION_ID, cleanupDeadlineAt: 11_001 }, + 1_000 + ) + ).toBeUndefined(); + }); + + it('sends strict Stop fields only to a negotiated peer', () => { + const payload = { messageId: 'a', operationId: OPERATION_ID, cleanupDeadlineAt: 11_000 }; + + expect(stopAbortWirePayload(payload, true)).toEqual(payload); + expect(stopAbortWirePayload(payload, false)).toEqual({ messageId: 'a' }); + }); + + it('identifies incomplete strict Stop payloads so callers can fail closed', () => { + expect(hasScopedStopMaintenanceFields({ messageId: 'a' })).toBe(false); + expect(hasScopedStopMaintenanceFields({ messageId: 'a', operationId: OPERATION_ID })).toBe( + true + ); + expect(hasScopedStopMaintenanceFields({ cleanupDeadlineAt: 11_000 })).toBe(true); + }); +}); diff --git a/services/cloud-agent-next/src/sandbox-control/scoped-stop-maintenance.ts b/services/cloud-agent-next/src/sandbox-control/scoped-stop-maintenance.ts new file mode 100644 index 0000000000..8ef31948a8 --- /dev/null +++ b/services/cloud-agent-next/src/sandbox-control/scoped-stop-maintenance.ts @@ -0,0 +1,33 @@ +import { + SANDBOX_CONTROL_CLEANUP_TIMEOUT_MS, + sessionScopedStopAbortPayloadSchema, + type SessionAbortPayload, +} from '../shared/sandbox-control-protocol.js'; + +export function parseScopedStopMaintenance(payload: unknown, now = Date.now()) { + const parsed = sessionScopedStopAbortPayloadSchema.safeParse(payload); + if ( + !parsed.success || + parsed.data.cleanupDeadlineAt <= now || + parsed.data.cleanupDeadlineAt > now + SANDBOX_CONTROL_CLEANUP_TIMEOUT_MS + ) + return undefined; + return parsed.data; +} + +export function hasScopedStopMaintenanceFields(payload: unknown): boolean { + return ( + typeof payload === 'object' && + payload !== null && + !Array.isArray(payload) && + (Object.hasOwn(payload, 'operationId') || Object.hasOwn(payload, 'cleanupDeadlineAt')) + ); +} + +export function stopAbortWirePayload( + payload: SessionAbortPayload, + supportsScopedStopAbort: boolean +): SessionAbortPayload { + if (supportsScopedStopAbort) return payload; + return payload.messageId ? { messageId: payload.messageId } : {}; +} diff --git a/services/cloud-agent-next/src/sandbox-control/socket.test.ts b/services/cloud-agent-next/src/sandbox-control/socket.test.ts index f74972a428..41df119f84 100644 --- a/services/cloud-agent-next/src/sandbox-control/socket.test.ts +++ b/services/cloud-agent-next/src/sandbox-control/socket.test.ts @@ -322,7 +322,11 @@ describe('sandbox control socket handler', () => { result: { protocolVersion: 1, handshakeComplete: true, - capabilities: { kiloVersionHeartbeat: true, sessionOperationResults: true }, + capabilities: { + kiloVersionHeartbeat: true, + sessionOperationResults: true, + scopedStopAbort: true, + }, }, }) ); diff --git a/services/cloud-agent-next/src/sandbox-control/socket.ts b/services/cloud-agent-next/src/sandbox-control/socket.ts index 69eaaa1733..69c423af3f 100644 --- a/services/cloud-agent-next/src/sandbox-control/socket.ts +++ b/services/cloud-agent-next/src/sandbox-control/socket.ts @@ -109,6 +109,7 @@ export type SandboxControlSocketHandler = { sendRequest(input: SandboxControlOutboundRequest): Promise; hasHandshakenSocket(): boolean; supportsOperationResults(): boolean; + supportsScopedStopAbort(): boolean; getConnectionIdentity(): SandboxControlConnectionIdentity | null; getReadySocket(): WebSocket | null; closeProvisionalSockets(): void; @@ -337,6 +338,13 @@ export function createSandboxControlSocketHandler( ); }, + supportsScopedStopAbort(): boolean { + const current = currentHandshakenSocket(state); + return ( + current !== null && readAttachment(current.socket)?.capabilities?.scopedStopAbort === true + ); + }, + getConnectionIdentity(): SandboxControlConnectionIdentity | null { return currentHandshakenSocket(state)?.identity ?? null; }, @@ -815,10 +823,14 @@ export function createSandboxControlSocketHandler( const authorizationTimeout = authorization?.success ? authorization.data.dispatchDeadlineAt - Date.now() : undefined; - const timeoutMs = - authorizationTimeout === undefined - ? input.timeoutMs - : Math.max(1, Math.min(input.timeoutMs ?? authorizationTimeout, authorizationTimeout)); + const deadlineTimeout = + input.deadlineAt === undefined ? undefined : Math.max(1, input.deadlineAt - Date.now()); + const timeoutMs = [input.timeoutMs, authorizationTimeout, deadlineTimeout] + .filter((timeout): timeout is number => timeout !== undefined) + .reduce( + (shortest, timeout) => (shortest === undefined ? timeout : Math.min(shortest, timeout)), + undefined + ); const pending = waiters.wait(requestId, timeoutMs); log('socket_request_sent', { ...diagnosticConnection(current.identity), diff --git a/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts b/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts index f549ab7fa6..8225d99f07 100644 --- a/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts +++ b/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts @@ -7,6 +7,12 @@ import { TRPCError } from '@trpc/server'; import { withTimeout } from '@kilocode/worker-utils'; import { z } from 'zod'; import { diagnosticSyncStatus } from '../shared/control-diagnostics.js'; +import { + controlSessionStateSchema, + controlStopRequestSchema, + type ControlSessionState, + type ControlStopReceipt, +} from '../shared/control-plane-session.js'; import { cloudAgentWorktreeIdSchema, cloudAgentWorktreeLocationSchema, @@ -142,6 +148,8 @@ import { withDeliveryDeadline, } from './control-dispatch.js'; import { acceptedAlarmDecision } from './accepted-overdue.js'; +import { createSessionStopLifecycle } from './session-stop-lifecycle.js'; +import { progressSessionStop } from './session-stop-progress.js'; import { bootPreparingStep, provisionPreparingStep } from './preparing-steps.js'; import { createSandboxTerminalBridge, type SandboxTerminalRecord } from './terminal-bridge.js'; import { @@ -219,6 +227,8 @@ export class SandboxSession extends DurableObject { private ingestPublicationChain: Promise = Promise.resolve(); private deletedWorktreeId: CloudAgentWorktreeId | undefined; private readonly activeOperations = new Set>(); + private readonly stopProgresses = new Map>(); + private readonly stopLifecycle: ReturnType; private deletionCompletion: Promise | undefined; private readonly worktreeChanges: ReturnType; private readonly interactionRefresh: InteractionRefresh; @@ -306,6 +316,18 @@ export class SandboxSession extends DurableObject { ), waitUntil: promise => this.ctx.waitUntil(promise), }); + this.stopLifecycle = createSessionStopLifecycle({ + list: () => ctx.storage.kv.list({ prefix: 'session_stop/' }), + readLegacy: () => ctx.storage.kv.get('session_stops'), + put: (key, value) => ctx.storage.kv.put(key, value), + delete: key => { + ctx.storage.kv.delete(key); + }, + deleteLegacy: () => { + ctx.storage.kv.delete('session_stops'); + }, + transaction: callback => ctx.storage.transactionSync(callback), + }); this.interactionRefresh = createInteractionRefresh({ captureScope: () => this.captureInteractionScope(), sync: (scope, trigger) => this.syncAcceptedMessage(scope, trigger), @@ -856,7 +878,56 @@ export class SandboxSession extends DurableObject { return; } - async interruptExecution(): Promise<{ success: boolean; message?: string }> { + async getControlState(): Promise { + return this.controlSessionState(); + } + + private controlSessionState(): ControlSessionState | null { + const metadata = this.terminalLifecycle.getStoredMetadata(); + const sandboxId = metadata?.workspace?.sandboxId; + if (!metadata || !sandboxId || this.terminalLifecycle.captureEpoch() === null) return null; + const targets = this.loadMessages() + .filter( + message => + (message.state === 'queued' || message.state === 'accepted') && + message.cancellation === undefined + ) + .map(message => ({ + messageId: message.messageId, + ...(message.wrapperInstanceId ? { wrapperInstanceId: message.wrapperInstanceId } : {}), + ...(message.executionDeadlineAt + ? { executionDeadlineAt: message.executionDeadlineAt } + : {}), + })); + if (targets.length === 0) return null; + return controlSessionStateSchema.parse({ + version: 1, + scope: { + sandboxId, + ...(this.terminalLifecycle.getAttachedWrapperInstanceId() + ? { wrapperInstanceId: this.terminalLifecycle.getAttachedWrapperInstanceId() } + : {}), + }, + targets, + }); + } + + async interruptExecution(): Promise<{ success: boolean; message?: string }>; + async interruptExecution(input: unknown): Promise; + async interruptExecution( + input?: unknown + ): Promise<{ success: boolean; message?: string } | ControlStopReceipt> { + if (input === undefined) return this.interruptLegacyExecution(); + const request = controlStopRequestSchema.parse(input); + const receipt = await this.admitControlStop(request); + if (receipt.state === 'accepted') { + await this.armQueueRetry(Math.min(receipt.cleanupDeadlineAt, Date.now() + QUEUE_RETRY_MS)); + void this.scheduleStopProgress(receipt.operationId); + } + return receipt; + } + + private async interruptLegacyExecution(): Promise<{ success: boolean; message?: string }> { const epoch = this.terminalLifecycle.captureEpoch(); if (epoch === null) return { success: false, message: 'Session not found' }; const before = this.loadMessages(); @@ -937,6 +1008,116 @@ export class SandboxSession extends DurableObject { } } + async getInterruptResult(operationId: string): Promise { + return this.stopLifecycle.receipt(operationId); + } + + private async admitControlStop( + request: ReturnType + ): Promise { + const epoch = this.terminalLifecycle.captureEpoch(); + const metadata = this.terminalLifecycle.getStoredMetadata(); + const receipt = this.stopLifecycle.admit({ + messages: this.loadMessages(), + request, + currentSandboxId: metadata?.workspace?.sandboxId, + currentWrapperInstanceId: this.terminalLifecycle.getAttachedWrapperInstanceId(), + now: Date.now(), + commit: (messages, stop) => + epoch !== null && + metadata !== undefined && + this.saveMessages(messages, epoch, 'coordinator', undefined, () => { + this.stopLifecycle.persist(stop); + }), + }); + const interrupted = this.loadMessages().some( + message => + message.state === 'accepted' && message.cancellation?.operationId === request.operationId + ); + if (interrupted && metadata) + this.worktreeChanges.markInterrupted(this.worktreeContext(metadata)); + return receipt; + } + + private scheduleStopProgress(operationId: string): Promise { + const current = this.stopProgresses.get(operationId); + if (current) return current; + const pending = this.trackOperation(this.runStopProgress(operationId)).finally(() => { + this.stopProgresses.delete(operationId); + }); + this.stopProgresses.set(operationId, pending); + this.ctx.waitUntil(pending); + return pending; + } + + private async runStopProgress(operationId: string): Promise { + const stop = this.stopLifecycle.get(operationId); + const epoch = this.terminalLifecycle.captureEpoch(); + const metadata = this.terminalLifecycle.getStoredMetadata(); + const sandboxId = metadata?.workspace?.sandboxId; + const kiloSessionId = metadata?.auth.kiloSessionId; + if (!stop || stop.state !== 'accepted') return; + const updated = await progressSessionStop({ + stop, + now: Date.now, + readMessages: () => this.loadMessages(), + saveMessages: messages => epoch !== null && this.saveMessages(messages, epoch), + abort: async target => { + if ( + epoch === null || + !metadata || + !sandboxId || + sandboxId !== stop.request.scope.sandboxId || + !kiloSessionId + ) + throw new Error('Stop runtime is unavailable'); + const expected = stop.request.targets.find(item => item.messageId === target.messageId); + if (!expected?.wrapperInstanceId) { + return { status: 'unconfirmed', quiescent: false }; + } + const response = await withDeliveryDeadline( + () => + sandboxControlRpc(this.env, sandboxId).request({ + operation: 'session.abort', + session: { + sessionId: metadata.identity.sessionId, + kiloSessionId, + directory: this.directory(metadata), + }, + payload: { + messageId: target.messageId, + operationId: stop.request.operationId, + cleanupDeadlineAt: stop.request.cleanupDeadlineAt, + }, + expectedWrapperInstanceId: expected.wrapperInstanceId, + }), + stop.request.cleanupDeadlineAt + ); + if (!response.ok) throw new Error('Session abort failed'); + return sessionAbortResultSchema.parse(response.result); + }, + applyDelivery: async delivery => { + if (!delivery) return undefined; + return this.applySandboxOperationResult({ + session: delivery.authorization.session, + wrapperInstanceId: delivery.authorization.wrapperInstanceId, + delivery, + }); + }, + }); + if (!this.stopLifecycle.replace(stop, updated)) return; + if (updated.state === 'accepted') + await this.armQueueRetry( + Math.min(updated.request.cleanupDeadlineAt, Date.now() + QUEUE_RETRY_MS) + ); + else if ( + epoch !== null && + this.terminalLifecycle.isCurrent(epoch) && + nextQueuedMessageId(this.loadMessages()) + ) + await this.armQueueRetry(); + } + async answerPermission(input: { permissionId: string; response: 'once' | 'always' | 'reject'; @@ -1412,12 +1593,16 @@ export class SandboxSession extends DurableObject { async alarm(): Promise { if (this.pendingRuntimeCleanup()) await this.transferRuntimeCleanup(); + for (const stop of this.stopLifecycle.pending()) + void this.scheduleStopProgress(stop.request.operationId); const epoch = this.terminalLifecycle.captureEpoch(); if (epoch === null || this.deletedWorktreeId) return; const now = Date.now(); const messages = this.loadMessages(); if (!this.terminalLifecycle.isCurrent(epoch)) return; - const accepted = messages.find(message => message.state === 'accepted'); + const accepted = messages.find( + message => message.state === 'accepted' && message.cancellation === undefined + ); if (accepted) { const decision = acceptedAlarmDecision( accepted.acceptedAt ?? 0, @@ -2269,11 +2454,13 @@ export class SandboxSession extends DurableObject { private async armQueueRetry(when = Date.now() + QUEUE_RETRY_MS): Promise { const epoch = this.terminalLifecycle.captureEpoch(); - if (epoch === null && !this.pendingRuntimeCleanup()) return; + const hasPendingStop = this.stopLifecycle.pending().length > 0; + if (epoch === null && !this.pendingRuntimeCleanup() && !hasPendingStop) return; const existing = await this.ctx.storage.getAlarm(); if ( (epoch === null || !this.terminalLifecycle.isCurrent(epoch)) && - !this.pendingRuntimeCleanup() + !this.pendingRuntimeCleanup() && + !hasPendingStop ) return; if (existing === null || existing > when) await this.ctx.storage.setAlarm(when); @@ -2734,10 +2921,16 @@ export class SandboxSession extends DurableObject { messages: MessageRecord[], epoch?: number, source: 'coordinator' | 'wrapper_outcome' | 'operation_result' = 'coordinator', - deferredNotifications?: StoredEvent[] + deferredNotifications?: StoredEvent[], + onPersist?: () => void ): boolean { - return this.commitSavedMessages(messages, epoch, source, deferredNotifications, write => - this.ctx.storage.transactionSync(write) + return this.commitSavedMessages( + messages, + epoch, + source, + deferredNotifications, + onPersist, + write => this.ctx.storage.transactionSync(write) ); } @@ -2747,8 +2940,13 @@ export class SandboxSession extends DurableObject { source: 'coordinator' | 'wrapper_outcome' | 'operation_result', deferredNotifications?: StoredEvent[] ): boolean { - return this.commitSavedMessages(messages, epoch, source, deferredNotifications, write => - write() + return this.commitSavedMessages( + messages, + epoch, + source, + deferredNotifications, + undefined, + write => write() ); } @@ -2757,6 +2955,7 @@ export class SandboxSession extends DurableObject { epoch: number | undefined, source: 'coordinator' | 'wrapper_outcome' | 'operation_result', deferredNotifications: StoredEvent[] | undefined, + onPersist: (() => void) | undefined, enclose: (write: () => void) => void ): boolean { const currentEpoch = epoch ?? this.terminalLifecycle.captureEpoch(); @@ -2837,6 +3036,7 @@ export class SandboxSession extends DurableObject { return terminal; }); this.ctx.storage.kv.put(MESSAGES_KEY, next); + onPersist?.(); }; enclose(write); for (const fields of committed) { 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 c350c4cd3e..48ac11378d 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 @@ -26,6 +26,7 @@ import { import { DEADLINE_MS } from '../sandbox-control/deadlines.js'; import { createControlPlaneCredential } from '../sandbox-control/managed-credential.js'; import { SESSION_DELIVERY_TIMEOUT_MS } from './control-dispatch.js'; +import { createControlStopRequest } from '../shared/control-plane-session.js'; import type { AcceptedCommandTurn, AcceptedPromptTurn, @@ -3773,6 +3774,93 @@ describe('SandboxSession orchestration', () => { }); }); +describe('SandboxSession durable Stop wiring', () => { + beforeEach(() => { + vi.useFakeTimers(); + vi.setSystemTime(1_000_000); + orchestrationMocks.broadcast.mockClear(); + }); + afterEach(() => { + vi.useRealTimers(); + }); + + async function stopRequest(fixture: ReturnType, operationId: string) { + const state = await fixture.session.getControlState(); + if (!state) throw new Error('Missing control state'); + return createControlStopRequest(state, Date.now(), operationId); + } + + it('keeps the Stop receipt absent when the paired message transaction fails', async () => { + const fixture = sessionFixture(); + await fixture.admit('a'); + await fixture.flush(); + const request = await stopRequest(fixture, '33333333-3333-4333-8333-333333333333'); + const before = fixture.record('a'); + vi.spyOn(fixture.storage, 'transactionSync').mockImplementationOnce(() => { + throw new Error('transaction failed'); + }); + + await expect(fixture.session.interruptExecution(request)).rejects.toThrow('transaction failed'); + + expect(fixture.record('a')).toEqual(before); + expect(fixture.values.get(`session_stop/${request.operationId}`)).toBeUndefined(); + }); + + it('repairs an admitted receipt with its original cleanup bound after alarm scheduling fails and reloads', async () => { + const fixture = sessionFixture(); + await fixture.admit('a'); + await fixture.flush(); + const request = await stopRequest(fixture, '33333333-3333-4333-8333-333333333334'); + await fixture.storage.deleteAlarm(); + vi.spyOn(fixture.storage, 'setAlarm').mockRejectedValueOnce(new Error('alarm unavailable')); + + await expect(fixture.session.interruptExecution(request)).rejects.toThrow('alarm unavailable'); + expect(await fixture.session.getInterruptResult(request.operationId)).toMatchObject({ + state: 'accepted', + cleanupDeadlineAt: request.cleanupDeadlineAt, + }); + + fixture.reload(); + await expect(fixture.session.interruptExecution(request)).resolves.toMatchObject({ + state: 'accepted', + cleanupDeadlineAt: request.cleanupDeadlineAt, + }); + await fixture.flush(); + + const abort = fixture.control.request.mock.calls + .map(([input]) => input) + .find(input => input.operation === 'session.abort'); + expect(abort).toMatchObject({ + payload: { + messageId: 'a', + operationId: request.operationId, + cleanupDeadlineAt: request.cleanupDeadlineAt, + }, + }); + }); + + it('marks an unrecoverable Stop unconfirmed after its bound when the terminal epoch disappears', async () => { + const fixture = sessionFixture(); + await fixture.admit('a'); + await fixture.flush(); + const request = await stopRequest(fixture, '33333333-3333-4333-8333-333333333335'); + + await expect(fixture.session.interruptExecution(request)).resolves.toMatchObject({ + state: 'accepted', + }); + await fixture.flush(); + fixture.values.delete('session_metadata'); + vi.setSystemTime(request.cleanupDeadlineAt); + await fixture.fireAlarm(); + await fixture.flush(); + + expect(fixture.values.get(`session_stop/${request.operationId}`)).toMatchObject({ + state: 'unconfirmed', + request: { cleanupDeadlineAt: request.cleanupDeadlineAt }, + }); + }); +}); + describe('streamCloudStatus', () => { it('is ready for accepted work even with queued followers', () => { expect(streamCloudStatus([msg('a', 'queued')])).toEqual({ type: 'preparing' }); 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 2efbe466c9..03052648c3 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 @@ -11,6 +11,7 @@ import { import { dispatchedKilocodeModelId } from '../persistence/model-utils.js'; import type { CloudMessageFailedPayload } from '../session/message-settlement-outbox.js'; import { + SANDBOX_CONTROL_EXECUTION_TIMEOUT_MS, sessionOperationAuthorizationSchema, sameSessionOperation, type SessionMessageOutcome, @@ -38,6 +39,7 @@ export type ControlSessionMessageInput = Pick(); + let legacy: unknown; + return { + values, + setLegacy(value: unknown) { + legacy = value; + }, + lifecycle: createSessionStopLifecycle({ + list: () => [...values.entries()], + readLegacy: () => legacy, + put: (key, value) => values.set(key, structuredClone(value)), + delete: key => { + values.delete(key); + }, + deleteLegacy: () => { + legacy = undefined; + }, + transaction: callback => callback(), + }), + }; +} + +describe('Session Stop lifecycle persistence', () => { + it('keeps an admitted receipt absent when the paired message transaction fails', () => { + const store = lifecycleStore(); + const messages = [message(1)]; + + expect(() => + store.lifecycle.admit({ + messages, + request: request(1), + currentSandboxId: 'sandbox-a', + currentWrapperInstanceId: RUNTIME_A, + now: 1_000, + commit: () => false, + }) + ).toThrow('Stop intent persistence failed'); + + expect(messages).toEqual([message(1)]); + expect(store.lifecycle.receipt(request(1).operationId)).toBeNull(); + }); + + it('retains all unexpired receipts and rejects admission at capacity', () => { + const store = lifecycleStore(); + for (let index = 0; index < SANDBOX_CONTROL_OPERATION_LIMIT; index++) { + const receipt = store.lifecycle.admit({ + messages: [message(index)], + request: request(index), + currentSandboxId: 'sandbox-a', + currentWrapperInstanceId: RUNTIME_A, + now: 1_000, + commit: (_messages, stop) => { + store.lifecycle.persist(stop); + return true; + }, + }); + expect(receipt.state).toBe('accepted'); + } + + const blocked = store.lifecycle.admit({ + messages: [message(SANDBOX_CONTROL_OPERATION_LIMIT)], + request: request(SANDBOX_CONTROL_OPERATION_LIMIT), + currentSandboxId: 'sandbox-a', + currentWrapperInstanceId: RUNTIME_A, + now: 1_000, + commit: () => true, + }); + + expect(blocked).toMatchObject({ + state: 'rejected', + message: 'Stop receipt capacity is unavailable', + }); + expect(store.lifecycle.pending()).toHaveLength(SANDBOX_CONTROL_OPERATION_LIMIT); + }); + + it('migrates a retained legacy receipt and prunes it only after its cleanup bound expires', () => { + vi.useFakeTimers(); + vi.setSystemTime(1_000); + try { + const store = lifecycleStore(); + const pending = store.lifecycle.admit({ + messages: [message(1)], + request: request(1, 2_000), + currentSandboxId: 'sandbox-a', + currentWrapperInstanceId: RUNTIME_A, + now: 1_000, + commit: (_messages, stop) => { + store.lifecycle.persist(stop); + return true; + }, + }); + const persisted = store.values.values().next().value as PersistedSessionStop; + store.values.clear(); + store.setLegacy({ [pending.operationId]: { ...persisted, state: 'unconfirmed' } }); + + expect(store.lifecycle.receipt(pending.operationId)).toMatchObject({ state: 'unconfirmed' }); + vi.setSystemTime(2_000); + expect(store.lifecycle.receipt(pending.operationId)).toMatchObject({ state: 'unconfirmed' }); + vi.setSystemTime(2_001); + expect(store.lifecycle.get(pending.operationId)).toBeUndefined(); + expect(store.values).toEqual(new Map()); + } finally { + vi.useRealTimers(); + } + }); +}); diff --git a/services/cloud-agent-next/src/sandbox-session/session-stop-lifecycle.ts b/services/cloud-agent-next/src/sandbox-session/session-stop-lifecycle.ts new file mode 100644 index 0000000000..76a360ec06 --- /dev/null +++ b/services/cloud-agent-next/src/sandbox-session/session-stop-lifecycle.ts @@ -0,0 +1,124 @@ +import type { ControlStopReceipt, ControlStopRequest } from '../shared/control-plane-session.js'; +import { SANDBOX_CONTROL_OPERATION_LIMIT } from '../shared/sandbox-control-protocol.js'; +import type { SessionMessageRecord } from './session-message-queue.js'; +import { + admitSessionStop, + persistedSessionStopSchema, + persistedSessionStopsSchema, + sessionStopReceipt, + type PersistedSessionStop, +} from './session-stop.js'; + +const STOP_PREFIX = 'session_stop/'; + +export type SessionStopLifecyclePersistence = { + list: () => Iterable<[string, unknown]>; + readLegacy: () => unknown; + put: (key: string, value: PersistedSessionStop) => void; + delete: (key: string) => void; + deleteLegacy: () => void; + transaction: (callback: () => T) => T; +}; + +function key(operationId: string): string { + return `${STOP_PREFIX}${operationId}`; +} + +function isExpired(stop: PersistedSessionStop, now: number): boolean { + return stop.state !== 'accepted' && now > stop.request.cleanupDeadlineAt; +} + +export function createSessionStopLifecycle(persistence: SessionStopLifecyclePersistence) { + const entries = (): Array<[string, PersistedSessionStop]> => { + const stored = [...persistence.list()].map(([storedKey, value]) => { + const parsed = persistedSessionStopSchema.safeParse(value); + if (!parsed.success) throw new Error('Persisted Stop state is invalid'); + return [storedKey, parsed.data] as [string, PersistedSessionStop]; + }); + if (stored.length > 0) return stored; + const legacyValue = persistence.readLegacy(); + if (legacyValue === undefined) return []; + const legacy = persistedSessionStopsSchema.safeParse(legacyValue); + if (!legacy.success) throw new Error('Persisted Stop state is invalid'); + return Object.entries(legacy.data).map(([id, stop]) => [key(id), stop]); + }; + + const stops = (): PersistedSessionStop[] => entries().map(([, stop]) => stop); + + const compact = (now = Date.now()): PersistedSessionStop[] => { + const current = entries(); + const retained = current.filter(([, stop]) => !isExpired(stop, now)); + if (retained.length === current.length) return retained.map(([, stop]) => stop); + persistence.transaction(() => { + for (const [storedKey] of current) { + if (!retained.some(([retainedKey]) => retainedKey === storedKey)) + persistence.delete(storedKey); + } + for (const [storedKey, stop] of retained) persistence.put(storedKey, stop); + persistence.deleteLegacy(); + }); + return retained.map(([, stop]) => stop); + }; + + const persist = (stop: PersistedSessionStop) => { + for (const [storedKey, current] of entries()) persistence.put(storedKey, current); + persistence.put(key(stop.request.operationId), stop); + persistence.deleteLegacy(); + }; + + return { + get(operationId: string): PersistedSessionStop | undefined { + return compact().find(stop => stop.request.operationId === operationId); + }, + receipt(operationId: string): ControlStopReceipt | null { + const stop = compact().find(stop => stop.request.operationId === operationId); + return stop ? sessionStopReceipt(stop) : null; + }, + pending(): PersistedSessionStop[] { + return compact().filter(stop => stop.state === 'accepted'); + }, + admit(input: { + request: ControlStopRequest; + messages: readonly SessionMessageRecord[]; + currentSandboxId?: string; + currentWrapperInstanceId?: string; + now: number; + commit: (messages: SessionMessageRecord[], stop: PersistedSessionStop) => boolean; + }): ControlStopReceipt { + const current = compact(input.now); + const existing = current.find(stop => stop.request.operationId === input.request.operationId); + const result = admitSessionStop({ + messages: input.messages, + existing, + request: input.request, + currentSandboxId: input.currentSandboxId, + currentWrapperInstanceId: input.currentWrapperInstanceId, + now: input.now, + }); + if (!result.stop || existing) return result.receipt; + if (current.length >= SANDBOX_CONTROL_OPERATION_LIMIT) + return { + ...result.receipt, + state: 'rejected', + message: 'Stop receipt capacity is unavailable', + }; + if (!input.commit(result.messages, result.stop)) + throw new Error('Stop intent persistence failed'); + return result.receipt; + }, + replace(previous: PersistedSessionStop, next: PersistedSessionStop): boolean { + const current = stops().find( + stop => stop.request.operationId === previous.request.operationId + ); + if ( + !current || + current.request.operationId !== previous.request.operationId || + current.request.cleanupDeadlineAt !== previous.request.cleanupDeadlineAt + ) + return false; + persistence.transaction(() => persist(next)); + return true; + }, + persist, + }; +} diff --git a/services/cloud-agent-next/src/sandbox-session/session-stop-progress.test.ts b/services/cloud-agent-next/src/sandbox-session/session-stop-progress.test.ts new file mode 100644 index 0000000000..7a9c58daeb --- /dev/null +++ b/services/cloud-agent-next/src/sandbox-session/session-stop-progress.test.ts @@ -0,0 +1,317 @@ +import { describe, expect, it, vi } from 'vitest'; +import type { ControlStopRequest } from '../shared/control-plane-session.js'; +import type { + SessionOperationAck, + SessionOperationAuthorization, + SessionOperationDelivery, +} from '../shared/sandbox-control-protocol.js'; +import type { SessionMessageRecord } from './session-message-queue.js'; +import { progressSessionStop } from './session-stop-progress.js'; +import { admitSessionStop } from './session-stop.js'; + +const RUNTIME_A = '11111111-1111-4111-8111-111111111111'; +const RUNTIME_B = '22222222-2222-4222-8222-222222222222'; +const STOP_ID = '33333333-3333-4333-8333-333333333333'; + +function request(): ControlStopRequest { + return { + version: 1, + operationId: STOP_ID, + scope: { sandboxId: 'sandbox-a', wrapperInstanceId: RUNTIME_A }, + targets: [ + { + messageId: 'a', + wrapperInstanceId: RUNTIME_A, + executionDeadlineAt: 3_601_000, + }, + ], + cleanupDeadlineAt: 11_000, + }; +} + +function authorization( + messageId = 'a', + wrapperInstanceId = RUNTIME_A +): SessionOperationAuthorization { + return { + operation: 'session.prompt', + operationId: messageId, + messageId, + session: { + sessionId: 'workspace_11111111-1111-4111-8111-111111111111', + kiloSessionId: 'ses_11111111111111111111111111', + directory: '/workspace/a', + }, + wrapperInstanceId, + dispatchDeadlineAt: 100_000, + }; +} + +function deliveredCompletion( + authorization: SessionOperationAuthorization +): SessionOperationDelivery { + return { + version: 2, + authorization, + completedAt: 2_000, + result: { ok: true, result: {} }, + outcome: { messageId: authorization.messageId, status: 'completed' }, + events: [], + preparing: [], + }; +} + +function deliveredFailure(authorization: SessionOperationAuthorization): SessionOperationDelivery { + return { + version: 2, + authorization, + completedAt: 2_000, + result: { + ok: false, + error: { code: 'runtime_unhealthy', message: 'Native execution failed', retryable: false }, + }, + outcome: { + messageId: authorization.messageId, + status: 'failed', + reason: 'Native execution failed', + }, + events: [], + preparing: [], + }; +} + +function acknowledgement( + delivery: SessionOperationDelivery, + state: SessionOperationAck['decision']['state'] +): SessionOperationAck { + return { + version: 2, + authorization: delivery.authorization, + resultHash: 'a'.repeat(64), + disposition: 'applied', + decision: { state, at: 2_000 }, + }; +} + +function dispatchedMessage(authorization: SessionOperationAuthorization): SessionMessageRecord { + return { + messageId: authorization.messageId, + state: 'accepted', + wrapperInstanceId: authorization.wrapperInstanceId, + executionDeadlineAt: 3_601_000, + operations: { prompt: { authorization, dispatched: true } }, + }; +} + +describe('Session Stop progress', () => { + it('retries the captured A target without changing B authority after a lost abort response', async () => { + let messages: SessionMessageRecord[] = [ + { + messageId: 'a', + state: 'accepted', + wrapperInstanceId: RUNTIME_A, + executionDeadlineAt: 3_601_000, + }, + ]; + const admitted = admitSessionStop({ + messages, + request: request(), + currentSandboxId: 'sandbox-a', + currentWrapperInstanceId: RUNTIME_A, + now: 1_000, + }); + if (!admitted.stop) throw new Error('Stop was not admitted'); + messages = admitted.messages; + const abort = vi + .fn() + .mockRejectedValueOnce(new Error('Response was lost')) + .mockResolvedValueOnce({ status: 'aborted', quiescent: true }); + const progress = () => + progressSessionStop({ + stop: admitted.stop!, + now: () => 2_000, + readMessages: () => messages, + saveMessages: next => { + messages = next; + return true; + }, + abort, + applyDelivery: async () => undefined, + }); + + await progress(); + messages.push({ + messageId: 'b', + state: 'accepted', + wrapperInstanceId: RUNTIME_B, + executionDeadlineAt: 3_602_000, + }); + const completed = await progress(); + + expect(abort.mock.calls.map(([target]) => target.messageId)).toEqual(['a', 'a']); + expect(completed.state).toBe('confirmed'); + expect(messages).toMatchObject([ + { messageId: 'a', state: 'cancelled' }, + { messageId: 'b', state: 'accepted', wrapperInstanceId: RUNTIME_B }, + ]); + }); + + it('does not terminalize an abort response without confirmed cleanup', async () => { + let messages: SessionMessageRecord[] = [ + { + messageId: 'a', + state: 'accepted', + wrapperInstanceId: RUNTIME_A, + executionDeadlineAt: 3_601_000, + }, + ]; + const admitted = admitSessionStop({ + messages, + request: request(), + currentSandboxId: 'sandbox-a', + currentWrapperInstanceId: RUNTIME_A, + now: 1_000, + }); + if (!admitted.stop) throw new Error('Stop was not admitted'); + messages = admitted.messages; + + const progressed = await progressSessionStop({ + stop: admitted.stop, + now: () => 2_000, + readMessages: () => messages, + saveMessages: next => { + messages = next; + return true; + }, + abort: async () => ({ status: 'aborted', quiescent: false }), + applyDelivery: async () => undefined, + }); + + expect(progressed.state).toBe('accepted'); + expect(messages).toMatchObject([{ messageId: 'a', state: 'accepted' }]); + }); + + it('keeps Stop unconfirmed for a failed retained result without cleanup proof', async () => { + const original = authorization(); + const delivery = deliveredFailure(original); + let messages: SessionMessageRecord[] = [dispatchedMessage(original)]; + const admitted = admitSessionStop({ + messages, + request: request(), + currentSandboxId: 'sandbox-a', + currentWrapperInstanceId: RUNTIME_A, + now: 1_000, + }); + if (!admitted.stop) throw new Error('Stop was not admitted'); + messages = admitted.messages; + + const progressed = await progressSessionStop({ + stop: admitted.stop, + now: () => 2_000, + readMessages: () => messages, + saveMessages: next => { + messages = next; + return true; + }, + abort: async () => ({ status: 'unconfirmed', quiescent: false, delivery }), + applyDelivery: async () => { + messages = messages.map(message => + message.messageId === 'a' + ? { ...message, state: 'failed', terminalSource: 'operation_result', terminalAt: 2_000 } + : message + ); + return acknowledgement(delivery, 'failed'); + }, + }); + + expect(progressed.state).toBe('accepted'); + expect(messages).toMatchObject([ + { messageId: 'a', state: 'failed', terminalSource: 'operation_result' }, + ]); + }); + + it('does not treat a mismatched delivery as cleanup for the Stop target', async () => { + const original = authorization(); + const mismatched = deliveredCompletion(authorization('b', RUNTIME_B)); + let messages: SessionMessageRecord[] = [dispatchedMessage(original)]; + const admitted = admitSessionStop({ + messages, + request: request(), + currentSandboxId: 'sandbox-a', + currentWrapperInstanceId: RUNTIME_A, + now: 1_000, + }); + if (!admitted.stop) throw new Error('Stop was not admitted'); + messages = admitted.messages; + + const progressed = await progressSessionStop({ + stop: admitted.stop, + now: () => 2_000, + readMessages: () => messages, + saveMessages: next => { + messages = next; + return true; + }, + abort: async () => ({ status: 'unconfirmed', quiescent: false, delivery: mismatched }), + applyDelivery: async () => { + messages = messages.map(message => + message.messageId === 'a' + ? { + ...message, + state: 'completed', + terminalSource: 'wrapper_outcome', + terminalAt: 2_000, + } + : message + ); + return acknowledgement(mismatched, 'completed'); + }, + }); + + expect(progressed.state).toBe('accepted'); + }); + + it('confirms Stop when the original completed delivery is acknowledged during Stop', async () => { + const original = authorization(); + const delivery = deliveredCompletion(original); + let messages: SessionMessageRecord[] = [dispatchedMessage(original)]; + const admitted = admitSessionStop({ + messages, + request: request(), + currentSandboxId: 'sandbox-a', + currentWrapperInstanceId: RUNTIME_A, + now: 1_000, + }); + if (!admitted.stop) throw new Error('Stop was not admitted'); + messages = admitted.messages; + + const progressed = await progressSessionStop({ + stop: admitted.stop, + now: () => 2_000, + readMessages: () => messages, + saveMessages: next => { + messages = next; + return true; + }, + abort: async () => ({ status: 'unconfirmed', quiescent: false, delivery }), + applyDelivery: async () => { + messages = messages.map(message => + message.messageId === 'a' + ? { + ...message, + state: 'completed', + terminalSource: 'operation_result', + terminalAt: 2_000, + } + : message + ); + return acknowledgement(delivery, 'completed'); + }, + }); + + expect(progressed.state).toBe('confirmed'); + expect(messages).toMatchObject([ + { messageId: 'a', state: 'completed', terminalSource: 'operation_result' }, + ]); + }); +}); diff --git a/services/cloud-agent-next/src/sandbox-session/session-stop-progress.ts b/services/cloud-agent-next/src/sandbox-session/session-stop-progress.ts new file mode 100644 index 0000000000..a7d9eececb --- /dev/null +++ b/services/cloud-agent-next/src/sandbox-session/session-stop-progress.ts @@ -0,0 +1,132 @@ +import { + sameSessionOperation, + sessionOperationAuthorizationSchema, + type SessionAbortResult, + type SessionOperationAck, + type SessionOperationDelivery, +} from '../shared/sandbox-control-protocol.js'; +import type { SessionMessageRecord } from './session-message-queue.js'; +import { expireSessionStop, settleSessionStop, type PersistedSessionStop } from './session-stop.js'; + +function isTerminal(message: SessionMessageRecord | undefined): boolean { + return ( + message?.state === 'completed' || message?.state === 'failed' || message?.state === 'cancelled' + ); +} + +function isCompletedDeliveredTarget(input: { + stop: PersistedSessionStop; + target: PersistedSessionStop['targets'][number]; + message: SessionMessageRecord | undefined; + delivery: SessionOperationDelivery; + acknowledgement: SessionOperationAck | undefined; +}): boolean { + const { stop, target, message, delivery, acknowledgement } = input; + const requestedTarget = stop.request.targets.find(item => item.messageId === target.messageId); + const storedAuthorization = sessionOperationAuthorizationSchema.safeParse( + message?.operations?.prompt?.authorization + ); + return ( + requestedTarget !== undefined && + delivery.authorization.operation === 'session.prompt' && + delivery.authorization.messageId === target.messageId && + delivery.authorization.wrapperInstanceId === requestedTarget.wrapperInstanceId && + storedAuthorization.success && + sameSessionOperation(storedAuthorization.data, delivery.authorization) && + delivery.result.ok && + delivery.outcome?.status === 'completed' && + acknowledgement !== undefined && + sameSessionOperation(acknowledgement.authorization, delivery.authorization) && + (acknowledgement.disposition === 'applied' || acknowledgement.disposition === 'identical') && + acknowledgement.decision.state === 'completed' + ); +} + +export async function progressSessionStop(input: { + stop: PersistedSessionStop; + now: () => number; + readMessages: () => SessionMessageRecord[]; + saveMessages: (messages: SessionMessageRecord[]) => boolean; + abort: (target: PersistedSessionStop['targets'][number]) => Promise; + applyDelivery: ( + delivery: SessionAbortResult['delivery'] + ) => Promise; +}): Promise { + let stop = input.stop; + for (const target of stop.targets) { + const message = input.readMessages().find(item => item.messageId === target.messageId); + const awaitingCancellation = + target.state === 'pending' && + message?.cancellation?.operationId === stop.request.operationId && + message.wrapperInstanceId !== undefined; + if ((isTerminal(message) && !awaitingCancellation) || !message?.cancellation) { + stop = settleSessionStop(stop, target.messageId); + continue; + } + if (target.state === 'confirmed' || input.now() >= stop.request.cleanupDeadlineAt) continue; + try { + const result = await input.abort(target); + const acknowledgement = result.delivery + ? await input.applyDelivery(result.delivery) + : undefined; + const current = input.readMessages().find(item => item.messageId === target.messageId); + if ( + result.delivery && + isCompletedDeliveredTarget({ + stop, + target, + message: current, + delivery: result.delivery, + acknowledgement, + }) + ) { + stop = settleSessionStop(stop, target.messageId); + continue; + } + if ( + result.quiescent === true && + current?.state === 'accepted' && + current.cancellation?.operationId === stop.request.operationId + ) { + input.saveMessages( + input.readMessages().map(message => + message.messageId === target.messageId && message.state === 'accepted' + ? { + ...message, + state: 'cancelled', + terminalAt: input.now(), + terminalSource: 'coordinator', + } + : message + ) + ); + } + if ( + result.quiescent === true && + isTerminal(input.readMessages().find(item => item.messageId === target.messageId)) + ) { + stop = settleSessionStop(stop, target.messageId); + } + } catch { + continue; + } + } + const expired = expireSessionStop(stop, input.now()); + if (expired.state === 'unconfirmed' && stop.state === 'accepted') { + input.saveMessages( + input.readMessages().map(message => + message.state === 'accepted' && + message.cancellation?.operationId === stop.request.operationId + ? { + ...message, + state: 'failed', + failedReason: 'interruption_unconfirmed', + terminalAt: input.now(), + terminalSource: 'coordinator', + } + : message + ) + ); + } + return expired; +} diff --git a/services/cloud-agent-next/src/sandbox-session/session-stop.test.ts b/services/cloud-agent-next/src/sandbox-session/session-stop.test.ts new file mode 100644 index 0000000000..86a549e882 --- /dev/null +++ b/services/cloud-agent-next/src/sandbox-session/session-stop.test.ts @@ -0,0 +1,185 @@ +import { describe, expect, it } from 'vitest'; +import type { ControlStopRequest } from '../shared/control-plane-session.js'; +import type { SessionOperationAuthorization } from '../shared/sandbox-control-protocol.js'; +import { + recordSessionOperationDispatch, + type SessionMessageRecord, +} from './session-message-queue.js'; +import { admitSessionStop, type PersistedSessionStop } from './session-stop.js'; + +const RUNTIME_A = '11111111-1111-4111-8111-111111111111'; +const RUNTIME_B = '22222222-2222-4222-8222-222222222222'; +const STOP_ID = '33333333-3333-4333-8333-333333333333'; + +function message( + messageId: string, + state: SessionMessageRecord['state'], + wrapperInstanceId?: string +): SessionMessageRecord { + return { + messageId, + state, + ...(wrapperInstanceId ? { wrapperInstanceId } : {}), + ...(state === 'accepted' ? { executionDeadlineAt: 3_601_000 } : {}), + }; +} + +function stopRequest(): ControlStopRequest { + return { + version: 1, + operationId: STOP_ID, + scope: { sandboxId: 'sandbox-a', wrapperInstanceId: RUNTIME_A }, + targets: [ + { + messageId: 'a', + wrapperInstanceId: RUNTIME_A, + executionDeadlineAt: 3_601_000, + }, + ], + cleanupDeadlineAt: 11_000, + }; +} + +function admittedStop(messages: SessionMessageRecord[]): PersistedSessionStop { + const admitted = admitSessionStop({ + messages, + request: stopRequest(), + currentSandboxId: 'sandbox-a', + currentWrapperInstanceId: RUNTIME_A, + now: 1_000, + }); + if (!admitted.stop) throw new Error('Stop was not admitted'); + return admitted.stop; +} + +describe('Session Stop admission', () => { + it('retains the first target and cleanup bound when a lost response is retried after B starts', () => { + const stop = admittedStop([message('a', 'accepted', RUNTIME_A)]); + const retry = admitSessionStop({ + messages: [message('a', 'completed', RUNTIME_A), message('b', 'accepted', RUNTIME_B)], + existing: stop, + request: stopRequest(), + currentSandboxId: 'sandbox-a', + currentWrapperInstanceId: RUNTIME_B, + now: 2_000, + }); + + expect(retry.receipt).toMatchObject({ + operationId: STOP_ID, + cleanupDeadlineAt: 11_000, + targets: [{ messageId: 'a', wrapperInstanceId: RUNTIME_A }], + }); + expect(retry.messages).toEqual([ + message('a', 'completed', RUNTIME_A), + message('b', 'accepted', RUNTIME_B), + ]); + }); + + it('fails closed when a Stop targets an older wrapper incarnation', () => { + const admission = admitSessionStop({ + messages: [message('a', 'accepted', RUNTIME_A)], + request: stopRequest(), + currentSandboxId: 'sandbox-a', + currentWrapperInstanceId: RUNTIME_B, + now: 1_000, + }); + + expect(admission.receipt).toMatchObject({ + state: 'rejected', + message: 'Stop runtime scope is stale', + }); + }); + + it('preserves the first cancellation authority when a later Stop overlaps the same live target', () => { + const first = admitSessionStop({ + messages: [message('a', 'accepted', RUNTIME_A)], + request: stopRequest(), + currentSandboxId: 'sandbox-a', + currentWrapperInstanceId: RUNTIME_A, + now: 1_000, + }); + if (!first.stop) throw new Error('First Stop was not admitted'); + const second = admitSessionStop({ + messages: first.messages, + request: { + ...stopRequest(), + operationId: '44444444-4444-4444-8444-444444444444', + cleanupDeadlineAt: 12_000, + }, + currentSandboxId: 'sandbox-a', + currentWrapperInstanceId: RUNTIME_A, + now: 2_000, + }); + + expect(second.receipt).toMatchObject({ + operationId: '44444444-4444-4444-8444-444444444444', + state: 'rejected', + message: 'Stop target already has a cancellation intent', + }); + expect(first.messages).toMatchObject([ + { messageId: 'a', cancellation: { operationId: STOP_ID, deadlineAt: 11_000 } }, + ]); + }); +}); + +describe('execution deadline persistence', () => { + it('does not replace the execution bound with a later dispatch attempt after reconstruction', () => { + const authorization: SessionOperationAuthorization = { + operation: 'session.prompt', + operationId: 'message-a', + messageId: 'message-a', + session: { sessionId: 'workspace-a', kiloSessionId: 'kilo-a', directory: '/workspace/a' }, + wrapperInstanceId: RUNTIME_A, + dispatchDeadlineAt: 31_000, + }; + const first = recordSessionOperationDispatch( + [message('message-a', 'queued', RUNTIME_A)], + authorization + ); + if (!first) throw new Error('Initial dispatch proof was not recorded'); + + const reconstructed = structuredClone(first); + const replayed = recordSessionOperationDispatch(reconstructed, authorization); + + expect(replayed?.[0]).toMatchObject({ executionDeadlineAt: 3_631_000 }); + expect(replayed?.[0]?.operations?.prompt).toMatchObject({ + executionDeadlineAt: 3_631_000, + }); + }); + + it('keeps an equal-valued cleanup bound separate from execution authority', () => { + const authorization: SessionOperationAuthorization = { + operation: 'session.prompt', + operationId: 'message-a', + messageId: 'message-a', + session: { sessionId: 'workspace-a', kiloSessionId: 'kilo-a', directory: '/workspace/a' }, + wrapperInstanceId: RUNTIME_A, + dispatchDeadlineAt: 11_000, + }; + const dispatched = recordSessionOperationDispatch( + [message('message-a', 'queued', RUNTIME_A)], + authorization + ); + if (!dispatched) throw new Error('Initial dispatch proof was not recorded'); + const admission = admitSessionStop({ + messages: dispatched, + request: { + ...stopRequest(), + targets: [ + { + messageId: 'message-a', + wrapperInstanceId: RUNTIME_A, + executionDeadlineAt: 3_611_000, + }, + ], + cleanupDeadlineAt: 11_000, + }, + currentSandboxId: 'sandbox-a', + currentWrapperInstanceId: RUNTIME_A, + now: 1_000, + }); + + expect(admission.receipt).toMatchObject({ cleanupDeadlineAt: 11_000 }); + expect(admission.messages).toMatchObject([{ executionDeadlineAt: 3_611_000 }]); + }); +}); diff --git a/services/cloud-agent-next/src/sandbox-session/session-stop.ts b/services/cloud-agent-next/src/sandbox-session/session-stop.ts new file mode 100644 index 0000000000..45b75b91c6 --- /dev/null +++ b/services/cloud-agent-next/src/sandbox-session/session-stop.ts @@ -0,0 +1,176 @@ +import type { ControlStopReceipt, ControlStopRequest } from '../shared/control-plane-session.js'; +import { controlStopRequestSchema } from '../shared/control-plane-session.js'; +import { SANDBOX_CONTROL_CLEANUP_TIMEOUT_MS } from '../shared/sandbox-control-protocol.js'; +import type { SessionMessageRecord } from './session-message-queue.js'; +import { z } from 'zod'; + +export const persistedSessionStopSchema = z + .object({ + version: z.literal(1), + request: z.lazy(() => controlStopRequestSchema), + state: z.enum(['accepted', 'confirmed', 'unconfirmed']), + targets: z.array( + z + .object({ + messageId: z.string().min(1), + state: z.enum(['pending', 'confirmed']), + }) + .strict() + ), + }) + .strict(); + +export const persistedSessionStopsSchema = z.record(z.string(), persistedSessionStopSchema); + +export type PersistedSessionStop = z.infer; + +type StopAdmission = { + messages: SessionMessageRecord[]; + stop?: PersistedSessionStop; + receipt: ControlStopReceipt; +}; + +function receipt(stop: PersistedSessionStop): ControlStopReceipt { + return { + version: 1, + operationId: stop.request.operationId, + scope: structuredClone(stop.request.scope), + targets: structuredClone(stop.request.targets), + cleanupDeadlineAt: stop.request.cleanupDeadlineAt, + state: stop.state, + }; +} + +function rejected(request: ControlStopRequest, message: string): StopAdmission { + return { + messages: [], + receipt: { + version: 1, + operationId: request.operationId, + scope: structuredClone(request.scope), + targets: structuredClone(request.targets), + cleanupDeadlineAt: request.cleanupDeadlineAt, + state: 'rejected', + message, + }, + }; +} + +function sameRequest(existing: ControlStopRequest, next: ControlStopRequest): boolean { + return ( + existing.operationId === next.operationId && + existing.cleanupDeadlineAt === next.cleanupDeadlineAt && + existing.scope.sandboxId === next.scope.sandboxId && + existing.scope.wrapperInstanceId === next.scope.wrapperInstanceId && + existing.targets.length === next.targets.length && + existing.targets.every( + (target, index) => + target.messageId === next.targets[index]?.messageId && + target.wrapperInstanceId === next.targets[index]?.wrapperInstanceId && + target.executionDeadlineAt === next.targets[index]?.executionDeadlineAt + ) + ); +} + +export function admitSessionStop(input: { + messages: readonly SessionMessageRecord[]; + existing?: PersistedSessionStop; + request: ControlStopRequest; + currentSandboxId?: string; + currentWrapperInstanceId?: string; + now: number; +}): StopAdmission { + const { messages, existing, request, currentSandboxId, currentWrapperInstanceId, now } = input; + if (existing) { + if (!sameRequest(existing.request, request)) + return rejected(request, 'Stop operation conflicts'); + return { messages: [...messages], stop: existing, receipt: receipt(existing) }; + } + if ( + request.cleanupDeadlineAt <= now || + request.cleanupDeadlineAt > now + SANDBOX_CONTROL_CLEANUP_TIMEOUT_MS + ) { + return rejected(request, 'Stop cleanup deadline is invalid'); + } + if (request.scope.sandboxId !== currentSandboxId) { + return rejected(request, 'Stop sandbox scope is stale'); + } + if ( + request.scope.wrapperInstanceId !== undefined && + request.scope.wrapperInstanceId !== currentWrapperInstanceId + ) { + return rejected(request, 'Stop runtime scope is stale'); + } + const requested = new Set(); + for (const target of request.targets) { + if (requested.has(target.messageId)) return rejected(request, 'Stop targets are duplicated'); + requested.add(target.messageId); + const message = messages.find(item => item.messageId === target.messageId); + if ( + !message || + (message.state !== 'queued' && message.state !== 'accepted') || + message.cancellation !== undefined || + message.wrapperInstanceId !== target.wrapperInstanceId || + message.executionDeadlineAt !== target.executionDeadlineAt + ) { + return rejected( + request, + message?.cancellation + ? 'Stop target already has a cancellation intent' + : 'Stop target is stale' + ); + } + } + const targets: PersistedSessionStop['targets'] = request.targets.map(target => { + const message = messages.find(item => item.messageId === target.messageId); + return { + messageId: target.messageId, + state: message?.wrapperInstanceId ? 'pending' : 'confirmed', + }; + }); + const state = targets.every(target => target.state === 'confirmed') ? 'confirmed' : 'accepted'; + const stop: PersistedSessionStop = { + version: 1, + request: structuredClone(request), + state, + targets, + }; + return { + messages: messages.map(message => { + if (!requested.has(message.messageId)) return message; + const cancellation = { + operationId: request.operationId, + deadlineAt: request.cleanupDeadlineAt, + }; + return message.state === 'queued' + ? { ...message, state: 'cancelled', cancellation } + : { ...message, cancellation }; + }), + stop, + receipt: receipt(stop), + }; +} + +export function settleSessionStop( + stop: PersistedSessionStop, + messageId: string +): PersistedSessionStop { + if (stop.state !== 'accepted') return stop; + const targets = stop.targets.map(target => + target.messageId === messageId ? { ...target, state: 'confirmed' as const } : target + ); + return { + ...stop, + targets, + state: targets.every(target => target.state === 'confirmed') ? 'confirmed' : 'accepted', + }; +} + +export function expireSessionStop(stop: PersistedSessionStop, now: number): PersistedSessionStop { + if (stop.state !== 'accepted' || now < stop.request.cleanupDeadlineAt) return stop; + return { ...stop, state: 'unconfirmed' }; +} + +export function sessionStopReceipt(stop: PersistedSessionStop): ControlStopReceipt { + return receipt(stop); +} diff --git a/services/cloud-agent-next/src/shared/control-plane-session.ts b/services/cloud-agent-next/src/shared/control-plane-session.ts new file mode 100644 index 0000000000..4bdae4e0e1 --- /dev/null +++ b/services/cloud-agent-next/src/shared/control-plane-session.ts @@ -0,0 +1,70 @@ +import { z } from 'zod'; +import { + SANDBOX_CONTROL_CLEANUP_TIMEOUT_MS, + wrapperInstanceIdSchema, +} from './sandbox-control-protocol.js'; + +const timestampSchema = z.number().int().positive(); + +export const controlStopTargetSchema = z + .object({ + messageId: z.string().min(1), + wrapperInstanceId: wrapperInstanceIdSchema.optional(), + executionDeadlineAt: timestampSchema.optional(), + }) + .strict(); + +export const controlStopScopeSchema = z + .object({ + sandboxId: z.string().min(1), + wrapperInstanceId: wrapperInstanceIdSchema.optional(), + }) + .strict(); + +export const controlStopRequestSchema = z + .object({ + version: z.literal(1), + operationId: z.string().uuid(), + scope: controlStopScopeSchema, + targets: z.array(controlStopTargetSchema).min(1), + cleanupDeadlineAt: timestampSchema, + }) + .strict(); + +export const controlSessionStateSchema = z + .object({ + version: z.literal(1), + scope: controlStopScopeSchema, + targets: z.array(controlStopTargetSchema).min(1), + }) + .strict(); + +export const controlStopReceiptSchema = z + .object({ + version: z.literal(1), + operationId: z.string().uuid(), + scope: controlStopScopeSchema, + targets: z.array(controlStopTargetSchema).min(1), + cleanupDeadlineAt: timestampSchema, + state: z.enum(['accepted', 'confirmed', 'unconfirmed', 'rejected']), + message: z.string().optional(), + }) + .strict(); + +export type ControlStopRequest = z.infer; +export type ControlSessionState = z.infer; +export type ControlStopReceipt = z.infer; + +export function createControlStopRequest( + state: ControlSessionState, + now = Date.now(), + operationId = crypto.randomUUID() +): ControlStopRequest { + return { + version: 1, + operationId, + scope: structuredClone(state.scope), + targets: structuredClone(state.targets), + cleanupDeadlineAt: now + SANDBOX_CONTROL_CLEANUP_TIMEOUT_MS, + }; +} 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 95c15f6c31..5615542116 100644 --- a/services/cloud-agent-next/src/shared/sandbox-control-protocol.ts +++ b/services/cloud-agent-next/src/shared/sandbox-control-protocol.ts @@ -159,7 +159,12 @@ export const sandboxHelloPayloadSchema = z.object({ providerInstanceId: z.string().min(1).max(256), wrapperInstanceId: wrapperInstanceIdSchema.optional(), wrapperVersion: z.string().min(1).max(128).optional(), - capabilities: z.object({ sessionOperationResults: z.boolean().optional() }).optional(), + capabilities: z + .object({ + sessionOperationResults: z.boolean().optional(), + scopedStopAbort: z.boolean().optional(), + }) + .optional(), }); export const sandboxHelloResultSchema = z.object({ @@ -169,6 +174,7 @@ export const sandboxHelloResultSchema = z.object({ .object({ kiloVersionHeartbeat: z.boolean().optional(), sessionOperationResults: z.boolean().optional(), + scopedStopAbort: z.boolean().optional(), }) .optional(), }); @@ -445,12 +451,24 @@ export const sessionAbortPayloadSchema = z .object({ messageId: z.string().min(1).max(128).optional(), reason: z.string().min(1).max(256).optional(), + operationId: z.string().uuid().optional(), + cleanupDeadlineAt: z.number().int().positive().optional(), + }) + .strict(); + +export const sessionScopedStopAbortPayloadSchema = z + .object({ + messageId: z.string().min(1).max(128), + operationId: z.string().uuid(), + cleanupDeadlineAt: z.number().int().positive(), }) .strict(); export const sessionAbortResultSchema = z .object({ - status: z.enum(['aborted', 'already_idle']), + status: z.enum(['aborted', 'already_idle', 'unconfirmed']), + quiescent: z.boolean().optional(), + delivery: z.lazy(() => sessionOperationDeliverySchema).optional(), }) .strict(); @@ -773,7 +791,12 @@ export const sandboxControlSocketAttachmentSchema = z.object({ acceptedAt: z.number().int().nonnegative(), connectionId: z.string().uuid().optional(), protocolVersion: z.literal(SANDBOX_CONTROL_PROTOCOL_VERSION).optional(), - capabilities: z.object({ sessionOperationResults: z.boolean().optional() }).optional(), + capabilities: z + .object({ + sessionOperationResults: z.boolean().optional(), + scopedStopAbort: z.boolean().optional(), + }) + .optional(), providerInstanceId: z.string().min(1).max(256).optional(), wrapperInstanceId: wrapperInstanceIdSchema.optional(), observation: sandboxControlObservationSchema.optional(), 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 eaeddd5179..fb52c089e7 100644 --- a/services/cloud-agent-next/test/integration/sandbox-control.test.ts +++ b/services/cloud-agent-next/test/integration/sandbox-control.test.ts @@ -391,7 +391,11 @@ async function completeHello( result: { protocolVersion: 1, handshakeComplete: true, - capabilities: { kiloVersionHeartbeat: true, sessionOperationResults: true }, + capabilities: { + kiloVersionHeartbeat: true, + sessionOperationResults: true, + scopedStopAbort: true, + }, }, }) ); 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 6b87ae7402..9b3abdcc58 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 @@ -17,6 +17,7 @@ import { completion, createHandlerFixture, fakeKilo, + kilo, operationAuthorization, promptPayload, session, @@ -123,6 +124,25 @@ describe('operation admission and lookup', () => { expect(record.snapshot().native.completion).toEqual(completion().info); expect(record.snapshot().local?.result.ok).toBe(true); expect(record.snapshot().delivery?.state).toBe('acknowledged'); + expect( + await handleControlRequest( + 'session.abort', + session, + { + messageId: authorization.messageId, + operationId: '11111111-1111-4111-8111-111111111111', + cleanupDeadlineAt: Date.now() + 1_000, + }, + handlerDeps + ) + ).toMatchObject({ + ok: true, + result: { + status: 'unconfirmed', + quiescent: false, + delivery: { authorization }, + }, + }); expect(handlerDeps.operations.counts().active).toBe(0); expect( await handleControlRequest('session.operation.get', session, authorization, handlerDeps) @@ -182,4 +202,25 @@ describe('operation admission and lookup', () => { ).toMatchObject({ ok: false, error: { code: 'not_ready', retryable: false } }); expect(handlerDeps.operations.counts().active).toBe(0); }); + + it('aborts a retained prompt instead of a same-message attach', async () => { + const handlerDeps = deps({ + sendOperationResult: (_session, delivery) => acknowledgeOperation(delivery), + }); + const attachAuth = operationAuthorization('session.attach'); + await handleControlRequest('session.attach', session, { kilo }, handlerDeps, attachAuth); + const attached = handlerDeps.operations.retained()[0]; + if (!attached) throw new Error('Missing attach'); + await attached.done; + await attached.waitForDelivery(); + const promptAuth = operationAuthorization(); + await handleControlRequest('session.prompt', session, promptPayload, handlerDeps, promptAuth); + const prompt = handlerDeps.operations + .retained() + .find(operation => operation.kind !== 'preparation'); + if (!prompt) throw new Error('Missing prompt'); + await prompt.done; + await prompt.waitForDelivery(); + expect(handlerDeps.operations.abortTarget(session, 'msg_1')).toBe(prompt); + }); }); diff --git a/services/cloud-agent-next/wrapper/src/control/operation-registry.ts b/services/cloud-agent-next/wrapper/src/control/operation-registry.ts index 18dd62a614..508dd5ba40 100644 --- a/services/cloud-agent-next/wrapper/src/control/operation-registry.ts +++ b/services/cloud-agent-next/wrapper/src/control/operation-registry.ts @@ -180,6 +180,22 @@ export function createOperationRegistry(deps: OperationRegistryDependencies) { start, prune, active: (rootKiloSessionId: string) => active.get(rootKiloSessionId), + abortTarget(session: SessionRequestIdentity, messageId?: string) { + const current = active.get(session.kiloSessionId); + if ( + current && + isDeepStrictEqual(current.session, session) && + (messageId === undefined || current.messageId === messageId) + ) { + return current; + } + if (messageId === undefined) return undefined; + const matches = [...retained.values()].filter( + operation => + operation.messageId === messageId && isDeepStrictEqual(operation.session, session) + ); + return matches.find(operation => operation.kind !== 'preparation') ?? matches[0]; + }, hasActive: (rootKiloSessionId: string) => active.has(rootKiloSessionId), activeOperations: () => [...active.values()], retained: () => [...retained.values()], diff --git a/services/cloud-agent-next/wrapper/src/control/sandbox-control-client.test.ts b/services/cloud-agent-next/wrapper/src/control/sandbox-control-client.test.ts index c499f42552..76556e9491 100644 --- a/services/cloud-agent-next/wrapper/src/control/sandbox-control-client.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/sandbox-control-client.test.ts @@ -304,7 +304,7 @@ describe('createSandboxControlClient', () => { protocolVersion: number; wrapperVersion: string; providerInstanceId: string; - capabilities?: { sessionOperationResults?: boolean }; + capabilities?: { sessionOperationResults?: boolean; scopedStopAbort?: boolean }; }; }; expect(hello.operation).toBe('sandbox.hello'); @@ -312,7 +312,7 @@ describe('createSandboxControlClient', () => { protocolVersion: 1, wrapperVersion: '2.4.0', providerInstanceId: 'inst_1', - capabilities: { sessionOperationResults: true }, + capabilities: { sessionOperationResults: true, scopedStopAbort: true }, }); fake.respond( JSON.stringify({ diff --git a/services/cloud-agent-next/wrapper/src/control/sandbox-control-client.ts b/services/cloud-agent-next/wrapper/src/control/sandbox-control-client.ts index 7c00e9488e..ec4a27c32c 100644 --- a/services/cloud-agent-next/wrapper/src/control/sandbox-control-client.ts +++ b/services/cloud-agent-next/wrapper/src/control/sandbox-control-client.ts @@ -372,7 +372,7 @@ export function createSandboxControlClient( payload: { protocolVersion: SANDBOX_CONTROL_PROTOCOL_VERSION, providerInstanceId: options.providerInstanceId, - capabilities: { sessionOperationResults: true }, + capabilities: { sessionOperationResults: true, scopedStopAbort: true }, ...(wrapperInstanceId ? { wrapperInstanceId } : {}), ...(options.wrapperVersion ? { wrapperVersion: options.wrapperVersion } : {}), }, diff --git a/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.test.ts b/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.test.ts index 58140da0c4..3f2ee9fe55 100644 --- a/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.test.ts @@ -2231,6 +2231,18 @@ describe('owned control execution', () => { expect( await handleControlRequest('session.abort', session, { messageId: 'older' }, handlerDeps) ).toEqual({ ok: true, result: { status: 'already_idle' } }); + expect( + await handleControlRequest( + 'session.abort', + session, + { + messageId: 'older', + operationId: '11111111-1111-4111-8111-111111111111', + cleanupDeadlineAt: Date.now() + 1_000, + }, + handlerDeps + ) + ).toEqual({ ok: true, result: { status: 'unconfirmed', quiescent: false } }); expect(taskSignal?.aborted).toBe(false); expect(aborts).toBe(0); const aborting = handleControlRequest( @@ -2308,6 +2320,73 @@ describe('owned control execution', () => { } }); + it('keeps a scoped Stop unconfirmed when native completion misses a successful abort acknowledgement', async () => { + const running = Promise.withResolvers(); + const started = Promise.withResolvers(); + const { handlerDeps, retired } = runtimeDeps( + fakeKilo({ + sendPrompt: () => { + started.resolve(); + return running.promise; + }, + abortSession: async () => true, + }) + ); + try { + await handleControlRequest('session.prompt', session, promptPayload, handlerDeps); + await started.promise; + + const aborted = await handleControlRequest( + 'session.abort', + session, + { + messageId: promptPayload.messageId, + operationId: '11111111-1111-4111-8111-111111111111', + cleanupDeadlineAt: Date.now() + 20, + }, + handlerDeps + ); + expect(aborted).toEqual({ ok: true, result: { status: 'unconfirmed', quiescent: false } }); + + expect(retired).toEqual(['Native cancellation did not settle']); + } finally { + running.resolve(completion({ name: 'MessageAbortedError', data: { message: 'cancelled' } })); + await waitForTasks(handlerDeps); + } + }); + + it('keeps a scoped Stop unconfirmed when preparation fails after cancellation', async () => { + const started = Promise.withResolvers(); + const handlerDeps = deps({ + applyAttach: async (_identity, _payload, hooks) => { + started.resolve(); + const signal = hooks.signal; + if (!signal) throw new Error('Missing preparation cancellation signal'); + await new Promise(resolve => + signal.addEventListener('abort', () => resolve(), { once: true }) + ); + return { + ok: false, + error: { code: 'not_ready', message: 'Preparation failed', retryable: true }, + }; + }, + }); + const attaching = handleControlRequest('session.attach', session, {}, handlerDeps); + await started.promise; + + const aborted = await handleControlRequest( + 'session.abort', + session, + { + operationId: '11111111-1111-4111-8111-111111111111', + cleanupDeadlineAt: Date.now() + 1_000, + }, + handlerDeps + ); + expect(aborted).toEqual({ ok: true, result: { status: 'unconfirmed', quiescent: false } }); + expect(await attaching).toMatchObject({ ok: false }); + }); + it.each(['false', 'malformed', 'HTTP failure'] as const)( 'retires and rejects replacement work after an abort returns %s', async response => { diff --git a/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.ts b/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.ts index f6dec67e6a..02c1c07c37 100644 --- a/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.ts +++ b/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.ts @@ -887,17 +887,36 @@ async function handleAbort( ): Promise { const parsed = sessionAbortPayloadSchema.safeParse(payload ?? {}); if (!parsed.success) return fail('protocol_error', 'Invalid payload', false); - const task = deps.operations.active(session.kiloSessionId); - if (parsed.data.messageId && task?.messageId !== parsed.data.messageId) { - return ok({ status: 'already_idle' }); - } + const task = deps.operations.abortTarget(session, parsed.data.messageId); if (task) { - task.cancel('Session aborted', 'cancelled'); + if ( + parsed.data.cleanupDeadlineAt !== undefined && + parsed.data.cleanupDeadlineAt <= Date.now() + ) { + return ok( + parsed.data.operationId + ? { status: 'unconfirmed', quiescent: false } + : { status: 'already_idle' } + ); + } + task.cancel('Session aborted', 'cancelled', parsed.data.cleanupDeadlineAt); const result = await task.done; + if (parsed.data.operationId) { + const delivery = task.deliveryResult(); + return ok({ + status: 'unconfirmed', + quiescent: false, + ...(delivery ? { delivery } : {}), + }); + } if (!result.ok && task.kind !== 'preparation') return result; return ok({ status: 'aborted' }); } - return ok({ status: 'already_idle' }); + return ok( + parsed.data.operationId + ? { status: 'unconfirmed', quiescent: false } + : { status: 'already_idle' } + ); } async function readRootRequests( diff --git a/services/cloud-agent-next/wrapper/src/control/session-operation.ts b/services/cloud-agent-next/wrapper/src/control/session-operation.ts index 8adbd530a7..1be5d8bc00 100644 --- a/services/cloud-agent-next/wrapper/src/control/session-operation.ts +++ b/services/cloud-agent-next/wrapper/src/control/session-operation.ts @@ -131,6 +131,7 @@ export class SessionOperation { private outcome?: SessionMessageOutcome; private local?: { result: ControlHandlerResult; completedAt: number }; private delivery?: OperationResultDelivery; + private cleanupDeadlineAt?: number; constructor( session: SessionRequestIdentity, @@ -239,7 +240,10 @@ export class SessionOperation { return this.delivery?.acknowledge(ack, isCurrent) ?? Promise.resolve(false); } - cancel(reason: string, status: 'failed' | 'cancelled'): void { + cancel(reason: string, status: 'failed' | 'cancelled', cleanupDeadlineAt?: number): void { + if (cleanupDeadlineAt !== undefined) { + this.cleanupDeadlineAt = Math.min(this.cleanupDeadlineAt ?? Infinity, cleanupDeadlineAt); + } if (!this.local) this.controller.abort(new ControlTaskCancellation(status, reason)); } @@ -603,7 +607,10 @@ export class SessionOperation { }; try { diagnostic('abort_started'); - const cleanupDeadlineAt = Date.now() + KILO_CONTROL_REQUEST_TIMEOUT_MS; + const cleanupDeadlineAt = Math.min( + this.cleanupDeadlineAt ?? Infinity, + Date.now() + KILO_CONTROL_REQUEST_TIMEOUT_MS + ); const abortController = new AbortController(); const abortTimer = setTimeout( () => abortController.abort(new Error('Kilo cancellation timed out')),