diff --git a/services/cloud-agent-next/src/persistence/SandboxControl.ts b/services/cloud-agent-next/src/persistence/SandboxControl.ts index dd9a884948..85ee66180b 100644 --- a/services/cloud-agent-next/src/persistence/SandboxControl.ts +++ b/services/cloud-agent-next/src/persistence/SandboxControl.ts @@ -34,6 +34,11 @@ import { type SandboxControlOutboundRequest, type SandboxControlSocketHandler, } from '../sandbox-control/socket.js'; +import { SandboxControlConnectionError } from '../sandbox-control/waiters.js'; +import { + createSessionForwarding, + SessionForwardingError, +} from '../sandbox-control/session-forwarding.js'; import { errorResponse, parseOperationPayload } from '../sandbox-control/frames.js'; import { generateSandboxCredential, @@ -44,14 +49,20 @@ import { import { SANDBOX_CONTROL_AUTO_PING, SANDBOX_CONTROL_AUTO_PONG, + sessionOperationAckSchema, + sessionOperationAuthorizationSchema, + sessionOperationExpiresAt, sessionRequestIdentitySchema, wrapperInstanceIdSchema, type ResponseFrame, type SessionAttachPayload, + type SessionOperationAck, + type SessionOperationDelivery, type SandboxHeartbeatPayload, type SessionEventIdentity, type SessionEventPayload, type SessionPreparingPayload, + type SessionRequestIdentity, } from '../shared/sandbox-control-protocol.js'; import { armDeadline, @@ -84,6 +95,7 @@ import { applyReportedSessionState, attachRoute, detachRoute, + getRouteBySessionId, hasActiveWork, resolveSessionEventRoute, type AttachRouteInput, @@ -247,6 +259,10 @@ type TerminalRuntimeRejection = { reason: string; }; +function sessionForwardFrameBytes(frame: unknown): number { + return new TextEncoder().encode(JSON.stringify(frame)).byteLength; +} + export type AttachSessionInput = AttachRouteInput; export type SandboxControlStatus = { @@ -255,6 +271,7 @@ export type SandboxControlStatus = { connection: ConnectionState; work: WorkState; wrapperInstanceId?: string; + operationResults?: true; }; export class SandboxControl extends DurableObject { @@ -264,7 +281,7 @@ export class SandboxControl extends DurableObject { private activeConnection: SandboxControlConnectionIdentity | null = null; private readyConnectionId: string | null = null; private providerKind: AgentSandboxProvider = 'cloudflare'; - private readonly sessionForwardChains = new Map>(); + private readonly sessionForwarding = createSessionForwarding(); private readonly forwarding = { enqueued: 0, settled: 0, @@ -274,6 +291,7 @@ export class SandboxControl extends DurableObject { dropped: 0, notApplied: 0, failed: 0, + bufferedBytes: 0, maxQueueWaitMs: 0, maxRpcWaitMs: 0, maxTotalForwardMs: 0, @@ -314,6 +332,8 @@ export class SandboxControl extends DurableObject { this.onSessionEvent(sessionIdentity, payload, identity), onSessionPreparing: (sessionIdentity, payload, identity) => this.onSessionPreparing(sessionIdentity, payload, identity), + onOperationResult: (session, delivery, identity) => + this.onOperationResult(session, delivery, identity), onSocketClosed: (handshakeComplete, identity) => this.onSocketClosed(handshakeComplete, identity), }); @@ -509,7 +529,9 @@ export class SandboxControl extends DurableObject { async request(input: SandboxControlOutboundRequest): Promise { await this.ensureOperationalInitialized(); if (input.operation === 'session.git.summary') return this.requestWorktreeChanges(input); - await this.assertRequestWorktreeAdmission(input); + const maintenance = + input.operation === 'session.operation.get' || input.operation === 'session.operation.ack'; + if (!maintenance) await this.assertRequestWorktreeAdmission(input); if (input.operation === 'worktree.delete' || input.operation === 'worktree.prepareDeletion') { throw new Error('Worktree cleanup requires the deletion coordinator'); } @@ -517,7 +539,48 @@ export class SandboxControl extends DurableObject { input.expectedWrapperInstanceId === undefined ? undefined : wrapperInstanceIdSchema.parse(input.expectedWrapperInstanceId); - const runtime = this.readyWrapperRuntime(); + const authorization = input.authorization + ? sessionOperationAuthorizationSchema.safeParse(input.authorization) + : undefined; + const maintenanceAck = + input.operation === 'session.operation.ack' + ? sessionOperationAckSchema.safeParse(input.payload) + : undefined; + const maintenanceAuthorization = + input.operation === 'session.operation.get' + ? sessionOperationAuthorizationSchema.safeParse(input.payload) + : maintenanceAck?.success + ? sessionOperationAuthorizationSchema.safeParse(maintenanceAck.data.authorization) + : undefined; + if ( + authorization && + (!authorization.success || + !this.socketHandler.supportsOperationResults() || + (input.operation !== 'session.attach' && input.operation !== 'session.prompt') || + authorization.data.operation !== input.operation || + input.session === undefined || + authorization.data.session.sessionId !== input.session.sessionId || + authorization.data.session.kiloSessionId !== input.session.kiloSessionId || + authorization.data.session.directory !== input.session.directory || + (expectedWrapperInstanceId !== undefined && + authorization.data.wrapperInstanceId !== expectedWrapperInstanceId) || + Date.now() >= authorization.data.dispatchDeadlineAt) + ) + throw new Error('Invalid session operation authorization'); + if ( + maintenance && + (!maintenanceAuthorization || + !maintenanceAuthorization.success || + input.session === undefined || + maintenanceAuthorization.data.session.sessionId !== input.session.sessionId || + maintenanceAuthorization.data.session.kiloSessionId !== input.session.kiloSessionId || + maintenanceAuthorization.data.session.directory !== input.session.directory || + Date.now() >= sessionOperationExpiresAt(maintenanceAuthorization.data)) + ) + throw new Error('Invalid session operation maintenance authorization'); + const runtime = maintenance + ? this.socketHandler.getConnectionIdentity() + : this.readyWrapperRuntime(); if (!runtime) throw new Error('Sandbox runtime is not ready'); if ( expectedWrapperInstanceId !== undefined && @@ -525,14 +588,21 @@ export class SandboxControl extends DurableObject { ) { throw new Error('Sandbox wrapper runtime changed'); } + if ( + authorization?.success && + authorization.data.wrapperInstanceId !== runtime.wrapperInstanceId + ) + throw new Error('Sandbox wrapper runtime changed'); const isCurrent = () => { - const current = this.readyWrapperRuntime(); + const current = maintenance + ? this.socketHandler.getConnectionIdentity() + : this.readyWrapperRuntime(); return current !== null && this.sameConnection(current, runtime); }; const physical = await loadPhysicalRecord(this.ctx.storage); if ( - physical.state !== 'running' || - physical.stopTombstone || + (!maintenance && physical.state !== 'running') || + (!maintenance && physical.stopTombstone) || physical.providerRef !== runtime.providerInstanceId || !isCurrent() ) { @@ -594,7 +664,7 @@ export class SandboxControl extends DurableObject { if (!isCurrent()) throw new Error('Sandbox wrapper runtime changed'); }); } - await this.assertRequestWorktreeAdmission(input); + if (!maintenance) await this.assertRequestWorktreeAdmission(input); if (!isCurrent()) throw new Error('Sandbox wrapper runtime changed'); return this.socketHandler.sendRequest(input); } @@ -1399,7 +1469,7 @@ export class SandboxControl extends DurableObject { } existed = result.existed; if (existed) { - this.sessionForwardChains.delete(sessionId); + this.sessionForwarding.delete(sessionId); await this.appendLog(routeTransition(Date.now(), 'detach', sessionId)); } if (!hasActiveWork(result.table)) { @@ -1500,8 +1570,8 @@ export class SandboxControl extends DurableObject { } return { value: sessionIds, changed: sessionIds.length > 0 }; }); - await Promise.allSettled(detached.flatMap(id => this.sessionForwardChains.get(id) ?? [])); - for (const id of detached) this.sessionForwardChains.delete(id); + await Promise.allSettled(detached.flatMap(id => this.sessionForwarding.get(id) ?? [])); + for (const id of detached) this.sessionForwarding.delete(id); if ( !journal.destroyed && (await this.fenceAndCheckWorktreeExclusivity( @@ -1835,6 +1905,10 @@ export class SandboxControl extends DurableObject { ...(physical.state === 'running' && runtime?.wrapperInstanceId ? { wrapperInstanceId: runtime.wrapperInstanceId } : {}), + ...(typeof this.socketHandler.supportsOperationResults === 'function' && + this.socketHandler.supportsOperationResults() + ? { operationResults: true as const } + : {}), }; } @@ -2588,8 +2662,7 @@ export class SandboxControl extends DurableObject { const physical = await loadPhysicalRecord(this.ctx.storage); if ( !this.isCurrentConnection(identity) || - physical.state !== 'running' || - physical.stopTombstone || + physical.state === 'stopped' || physical.providerRef !== identity.providerInstanceId ) return; @@ -2817,14 +2890,25 @@ export class SandboxControl extends DurableObject { this.recordForwardDrop('missing_identity', diagnostic); return; } - await this.forwardRoutedSessionFrame(identity, payload.type, connection, (route, fields) => - this.forwardSessionFrame(route, connection, fields, 'receiveSandboxControlEvent', stub => - stub.receiveSandboxControlEvent({ - identity, - payload, - wrapperInstanceId: connection.wrapperInstanceId, - }) - ) + await this.forwardRoutedSessionFrame( + identity, + payload.type, + connection, + { identity, payload }, + (route, fields, physical) => + this.forwardSessionFrame( + route, + physical, + connection, + fields, + 'receiveSandboxControlEvent', + stub => + stub.receiveSandboxControlEvent({ + identity, + payload, + wrapperInstanceId: connection.wrapperInstanceId, + }) + ) ); } @@ -2849,9 +2933,11 @@ export class SandboxControl extends DurableObject { identity, 'session.preparing', connection, - (route, fields) => + { identity, payload }, + (route, fields, physical) => this.forwardSessionFrame( route, + physical, connection, fields, 'receiveSandboxControlPreparing', @@ -2865,11 +2951,198 @@ export class SandboxControl extends DurableObject { ); } + private async onOperationResult( + session: SessionRequestIdentity, + delivery: SessionOperationDelivery, + identity: SandboxControlConnectionIdentity + ): Promise { + const connection = this.socketHandler.getConnectionIdentity(); + if (!connection || connection.connectionId !== identity.connectionId) return undefined; + const authorization = delivery.authorization; + if ( + !identity.wrapperInstanceId || + authorization.wrapperInstanceId !== identity.wrapperInstanceId || + session.sessionId !== authorization.session.sessionId || + session.kiloSessionId !== authorization.session.kiloSessionId || + session.directory !== authorization.session.directory + ) + return undefined; + const deadlineAt = sessionOperationExpiresAt(authorization); + const current = () => this.isCurrentConnection(connection) && Date.now() < deadlineAt; + if (!current()) return undefined; + const [table, physical, ownerId] = await Promise.all([ + loadRouteTable(this.ctx.storage), + loadPhysicalRecord(this.ctx.storage), + this.readOwner(), + ]); + const route = getRouteBySessionId(table, session.sessionId); + if ( + !route || + route.ownerId !== ownerId || + route.kiloSessionId !== session.kiloSessionId || + route.directory !== session.directory || + this.runtimeDeleted || + this.exclusiveDeletionWorktreeId !== undefined || + (route.worktreeId !== undefined && this.deletingWorktrees.has(route.worktreeId)) || + !current() + ) + throw new SandboxControlConnectionError('Operation result route is not current', false); + if ( + physical.state === 'stopped' || + physical.providerRef !== connection.providerInstanceId || + !this.matchesWorktreeContainment(physical) || + !current() + ) + throw new SandboxControlConnectionError('Operation result runtime is not current', false); + let frameBytes: number; + try { + frameBytes = sessionForwardFrameBytes({ session, delivery }); + } catch { + throw new SandboxControlConnectionError('Operation result cannot be serialized', false); + } + const diagnostic = { + ...diagnosticConnection(connection), + eventType: diagnosticEventType('session.operation.result'), + }; + const queuedAt = Date.now(); + this.forwarding.enqueued++; + const fields = { + ...diagnostic, + sessionId: route.sessionId, + forwardSequence: this.forwarding.enqueued, + queuedAt, + }; + this.logDiagnostic('forward_enqueued', { ...fields, ...this.forwarding }); + const forwardDeadlineAt = Math.min(deadlineAt, Date.now() + DEADLINE_MS.stopAttempt); + const next = this.sessionForwarding.enqueueFenced({ + sessionId: route.sessionId, + bytes: frameBytes, + deadlineAt: forwardDeadlineAt, + fence: async () => current(), + forward: async () => { + const queueWaitMs = Date.now() - queuedAt; + this.forwarding.maxQueueWaitMs = Math.max(this.forwarding.maxQueueWaitMs, queueWaitMs); + this.logDiagnostic('forward_started', { ...fields, queueWaitMs, ...this.forwarding }); + try { + if (!current()) + throw new SandboxControlConnectionError('Operation result expired', false); + const wrapperInstanceId = identity.wrapperInstanceId; + if (!wrapperInstanceId) + throw new SandboxControlConnectionError( + 'Operation result wrapper is not current', + false + ); + const assertCurrent = async () => { + if (!current() || Date.now() >= forwardDeadlineAt) + throw new SandboxControlConnectionError('Operation result forwarding expired', false); + const [routes, nextPhysical] = await Promise.all([ + loadRouteTable(this.ctx.storage), + loadPhysicalRecord(this.ctx.storage), + ]); + const nextRoute = routes.get(route.sessionId); + if ( + !current() || + !nextRoute || + nextRoute.ownerId !== route.ownerId || + nextRoute.kiloSessionId !== session.kiloSessionId || + nextRoute.directory !== session.directory || + nextRoute.worktreeId !== route.worktreeId || + Date.now() >= forwardDeadlineAt || + !sameAllocation(nextPhysical, physical) || + nextPhysical.state === 'stopped' || + nextPhysical.providerRef !== connection.providerInstanceId || + !this.matchesWorktreeContainment(nextPhysical) || + this.runtimeDeleted || + this.exclusiveDeletionWorktreeId !== undefined || + (nextRoute.worktreeId !== undefined && + this.deletingWorktrees.has(nextRoute.worktreeId)) || + !current() + ) + 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; + }, + 'receiveSandboxOperationResult' + ), + Math.max(1, forwardDeadlineAt - Date.now()), + 'Operation result forwarding timed out' + ); + if (!current()) + throw new SandboxControlConnectionError('Operation result expired', false); + const parsed = sessionOperationAckSchema.safeParse(ack); + if (!parsed.success) + throw new SandboxControlConnectionError( + 'Operation result acknowledgement is invalid', + false + ); + this.logDiagnostic('forward_result', { + ...fields, + operation: 'receiveSandboxOperationResult', + attempts, + result: 'delivered', + }); + return parsed.data; + } catch (error) { + this.forwarding.failed++; + this.logDiagnostic( + 'forward_result', + { + ...fields, + operation: 'receiveSandboxOperationResult', + result: 'failed', + }, + 'warn' + ); + throw error instanceof SandboxControlConnectionError + ? error + : error instanceof SessionForwardingError + ? new SandboxControlConnectionError(error.message, error.retryable) + : new SandboxControlConnectionError('Operation result forwarding failed'); + } finally { + this.forwarding.settled++; + const totalForwardMs = Date.now() - queuedAt; + this.forwarding.maxTotalForwardMs = Math.max( + this.forwarding.maxTotalForwardMs, + totalForwardMs + ); + this.logDiagnostic('forward_settled', { ...fields, totalForwardMs, ...this.forwarding }); + } + }, + }); + const delivered = next.catch(error => { + throw error instanceof SessionForwardingError + ? new SandboxControlConnectionError(error.message, error.retryable) + : error; + }); + this.ctx.waitUntil(delivered.catch(() => undefined)); + return delivered; + } + private async forwardRoutedSessionFrame( identity: SessionEventIdentity, eventType: string, connection: SandboxControlConnectionIdentity, - forward: (route: SessionRoute, diagnostic: ControlDiagnosticFields) => Promise + frame: unknown, + forward: ( + route: SessionRoute, + diagnostic: ControlDiagnosticFields, + physical: PhysicalRecord + ) => Promise ): Promise { const diagnostic = { ...diagnosticConnection(connection), @@ -2885,14 +3158,26 @@ export class SandboxControl extends DurableObject { this.recordForwardDrop('unroutable', { ...diagnostic, routeCount: table.size }); return; } - + const physical = await loadPhysicalRecord(this.ctx.storage); + if ( + physical.state !== 'running' || + physical.stopTombstone || + physical.providerRef !== connection.providerInstanceId || + !this.matchesWorktreeContainment(physical) || + !this.isCurrentConnection(connection) + ) { + this.recordForwardDrop('runtime_not_current', diagnostic); + return; + } + let frameBytes: number; + try { + frameBytes = sessionForwardFrameBytes(frame); + } catch { + this.recordForwardDrop('forwarding_frame_invalid', diagnostic); + return; + } const queuedAt = Date.now(); this.forwarding.enqueued++; - this.forwarding.waiting++; - this.forwarding.highWater = Math.max( - this.forwarding.highWater, - this.forwarding.waiting + this.forwarding.inFlight - ); const fields = { ...diagnostic, sessionId: route.sessionId, @@ -2900,12 +3185,12 @@ export class SandboxControl extends DurableObject { queuedAt, }; this.logDiagnostic('forward_enqueued', { ...fields, ...this.forwarding }); - const previous = this.sessionForwardChains.get(route.sessionId) ?? Promise.resolve(); - const next = previous - .catch(() => undefined) - .then(async () => { - this.forwarding.waiting--; - this.forwarding.inFlight++; + const next = this.sessionForwarding.enqueueFenced({ + sessionId: route.sessionId, + bytes: frameBytes, + deadlineAt: Date.now() + DEADLINE_MS.stopAttempt, + fence: async () => this.isCurrentConnection(connection), + forward: async () => { const queueWaitMs = Date.now() - queuedAt; this.forwarding.maxQueueWaitMs = Math.max(this.forwarding.maxQueueWaitMs, queueWaitMs); this.logDiagnostic('forward_started', { @@ -2914,10 +3199,9 @@ export class SandboxControl extends DurableObject { ...this.forwarding, }); try { - if (this.isCurrentConnection(connection)) await forward(route, fields); + if (this.isCurrentConnection(connection)) await forward(route, fields, physical); else this.recordForwardDrop('stale_before_send', fields); } finally { - this.forwarding.inFlight--; this.forwarding.settled++; const totalForwardMs = Date.now() - queuedAt; this.forwarding.maxTotalForwardMs = Math.max( @@ -2930,9 +3214,18 @@ export class SandboxControl extends DurableObject { ...this.forwarding, }); } - }); - this.sessionForwardChains.set(route.sessionId, next); - this.ctx.waitUntil(next); + }, + }); + this.ctx.waitUntil( + next.catch(error => { + this.recordForwardDrop( + error instanceof SessionForwardingError && !error.retryable + ? 'forwarding_frame_rejected' + : 'forwarding_capacity_exhausted', + fields + ); + }) + ); } private recordForwardDrop(reason: string, fields: ControlDiagnosticFields): void { @@ -2940,14 +3233,43 @@ export class SandboxControl extends DurableObject { this.logDiagnostic('forward_dropped', { ...fields, reason, ...this.forwarding }); } + private async isCurrentSessionForward( + route: SessionRoute, + connection: SandboxControlConnectionIdentity, + expectedPhysical: PhysicalRecord + ): Promise { + if (!this.isCurrentConnection(connection)) return false; + const [routes, physical] = await Promise.all([ + loadRouteTable(this.ctx.storage), + loadPhysicalRecord(this.ctx.storage), + ]); + const current = routes.get(route.sessionId); + return ( + current?.ownerId === route.ownerId && + current.kiloSessionId === route.kiloSessionId && + current.directory === route.directory && + current.worktreeId === route.worktreeId && + sameAllocation(physical, expectedPhysical) && + physical.state === 'running' && + !physical.stopTombstone && + physical.providerRef === connection.providerInstanceId && + this.matchesWorktreeContainment(physical) && + !this.runtimeDeleted && + !this.exclusiveDeletionWorktreeId && + !(route.worktreeId && this.deletingWorktrees.has(route.worktreeId)) && + this.isCurrentConnection(connection) + ); + } + private async forwardSessionFrame( route: SessionRoute, + physical: PhysicalRecord, connection: SandboxControlConnectionIdentity, diagnostic: ControlDiagnosticFields, operation: 'receiveSandboxControlEvent' | 'receiveSandboxControlPreparing', send: (stub: ReturnType) => Promise<{ applied: boolean }> ): Promise { - if (!this.isCurrentConnection(connection)) { + if (!(await this.isCurrentSessionForward(route, connection, physical))) { this.recordForwardDrop('stale_before_send', diagnostic); return; } @@ -2958,14 +3280,19 @@ export class SandboxControl extends DurableObject { const delivered = await withTimeout( withDORetry( () => getSandboxSessionStub(this.env, route.ownerId, route.sessionId), - stub => { - skipped = !this.isCurrentConnection(connection); + async stub => { + skipped = !(await this.isCurrentSessionForward(route, connection, physical)); if (skipped) { this.recordForwardDrop('stale_retry', diagnostic); return Promise.resolve({ applied: true }); } attempts++; - return send(stub); + const result = await send(stub); + if (!(await this.isCurrentSessionForward(route, connection, physical))) { + this.recordForwardDrop('stale_after_send', diagnostic); + return { applied: false }; + } + return result; }, operation ), 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 d36bf30aaa..0fe43cfc28 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,7 @@ describe('sandbox control frames', () => { expect(sandboxHelloResultSchema.parse(helloResult())).toEqual({ protocolVersion: 1, handshakeComplete: true, - capabilities: { kiloVersionHeartbeat: true }, + capabilities: { kiloVersionHeartbeat: true, sessionOperationResults: 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 f4f783719a..c742960086 100644 --- a/services/cloud-agent-next/src/sandbox-control/frames.ts +++ b/services/cloud-agent-next/src/sandbox-control/frames.ts @@ -24,6 +24,8 @@ import { sessionTerminalConnectPayloadSchema, sessionTerminalCreatePayloadSchema, sessionTerminalResizePayloadSchema, + sessionOperationAuthorizationSchema, + sessionOperationAckSchema, worktreeDeletePayloadSchema, type ControlError, type ControlErrorCode, @@ -57,6 +59,8 @@ const REQUEST_PAYLOAD_SCHEMAS: Record = { 'session.terminal.resize': sessionTerminalResizePayloadSchema, 'session.terminal.close': sessionTerminalClosePayloadSchema, 'session.terminal.connect': sessionTerminalConnectPayloadSchema, + 'session.operation.get': sessionOperationAuthorizationSchema, + 'session.operation.ack': sessionOperationAckSchema, }; const EVENT_PAYLOAD_SCHEMAS: Record = { @@ -177,6 +181,6 @@ export function helloResult(): SandboxHelloResult { return { protocolVersion: SANDBOX_CONTROL_PROTOCOL_VERSION, handshakeComplete: true, - capabilities: { kiloVersionHeartbeat: true }, + capabilities: { kiloVersionHeartbeat: true, sessionOperationResults: true }, }; } diff --git a/services/cloud-agent-next/src/sandbox-control/session-forwarding.test.ts b/services/cloud-agent-next/src/sandbox-control/session-forwarding.test.ts new file mode 100644 index 0000000000..84add1a29d --- /dev/null +++ b/services/cloud-agent-next/src/sandbox-control/session-forwarding.test.ts @@ -0,0 +1,120 @@ +import { afterEach, describe, expect, it, vi } from 'vitest'; +import { createSessionForwarding } from './session-forwarding.js'; + +describe('createSessionForwarding', () => { + afterEach(() => { + vi.useRealTimers(); + }); + it('keeps results behind earlier events for the same session', async () => { + const forwarding = createSessionForwarding(); + const releaseEvent = Promise.withResolvers(); + const event = vi.fn(async () => releaseEvent.promise); + const result = vi.fn(async () => undefined); + + const first = forwarding.enqueue('workspace_1', event); + const second = forwarding.enqueue('workspace_1', result); + await Promise.resolve(); + expect(result).not.toHaveBeenCalled(); + + releaseEvent.resolve(); + await expect(first).resolves.toBeUndefined(); + await expect(second).resolves.toBeUndefined(); + expect(event).toHaveBeenCalledTimes(1); + expect(result).toHaveBeenCalledTimes(1); + }); + + it('continues a session chain after a failed forwarding attempt', async () => { + const forwarding = createSessionForwarding(); + const failure = forwarding.enqueue('workspace_1', async () => { + throw new Error('forwarding failed'); + }); + const recovery = vi.fn(async () => 'acknowledged'); + const next = forwarding.enqueue('workspace_1', recovery); + + await expect(failure).rejects.toThrow('forwarding failed'); + await expect(next).resolves.toBe('acknowledged'); + expect(recovery).toHaveBeenCalledTimes(1); + }); + + it('does not call a fenced delivery after its deadline', async () => { + const forwarding = createSessionForwarding(); + const forward = vi.fn(async () => 'acknowledged'); + + await expect( + forwarding.enqueueFenced({ + sessionId: 'workspace_1', + bytes: 1, + deadlineAt: Date.now() - 1, + fence: async () => true, + forward, + }) + ).rejects.toMatchObject({ retryable: false }); + expect(forward).not.toHaveBeenCalled(); + }); + + it('rechecks the deadline after a delayed fence', async () => { + vi.useFakeTimers(); + vi.setSystemTime(0); + const forwarding = createSessionForwarding(); + const fence = Promise.withResolvers(); + const forward = vi.fn(async () => 'acknowledged'); + const pending = forwarding.enqueueFenced({ + sessionId: 'workspace_1', + bytes: 1, + deadlineAt: 1, + fence: async () => fence.promise, + forward, + }); + + await Promise.resolve(); + vi.setSystemTime(2); + fence.resolve(true); + await expect(pending).rejects.toMatchObject({ retryable: false }); + expect(forward).not.toHaveBeenCalled(); + }); + + it('rechecks the deadline after forwarding before returning a result', async () => { + vi.useFakeTimers(); + vi.setSystemTime(0); + const forwarding = createSessionForwarding(); + const finalFence = Promise.withResolvers(); + let fenceCalls = 0; + const pending = forwarding.enqueueFenced({ + sessionId: 'workspace_1', + bytes: 1, + deadlineAt: 1, + fence: () => (++fenceCalls === 1 ? Promise.resolve(true) : finalFence.promise), + forward: async () => 'acknowledged', + }); + + await Promise.resolve(); + vi.setSystemTime(2); + finalFence.resolve(true); + await expect(pending).rejects.toMatchObject({ retryable: false }); + }); + + it('releases capacity after a pre-start fence rejection', async () => { + const forwarding = createSessionForwarding(); + await expect( + forwarding.enqueueFenced({ + sessionId: 'workspace_1', + bytes: 1, + deadlineAt: Date.now() + 1_000, + fence: async () => false, + forward: async () => 'unreachable', + }) + ).rejects.toMatchObject({ retryable: false }); + expect(forwarding.stats()).toMatchObject({ waiting: 0, inFlight: 0, bufferedBytes: 0 }); + + await expect( + forwarding.enqueueFenced({ + sessionId: 'workspace_1', + bytes: 1, + deadlineAt: Date.now() + 1_000, + fence: async () => true, + forward: async () => 'acknowledged', + }) + ).resolves.toBe('acknowledged'); + expect(forwarding.stats()).toMatchObject({ waiting: 0, inFlight: 0, bufferedBytes: 0 }); + }); +}); diff --git a/services/cloud-agent-next/src/sandbox-control/session-forwarding.ts b/services/cloud-agent-next/src/sandbox-control/session-forwarding.ts new file mode 100644 index 0000000000..b242a7581a --- /dev/null +++ b/services/cloud-agent-next/src/sandbox-control/session-forwarding.ts @@ -0,0 +1,108 @@ +import { + MAX_SANDBOX_CONTROL_FRAME_BYTES, + SANDBOX_CONTROL_OPERATION_LIMIT, +} from '../shared/sandbox-control-protocol.js'; + +const MAX_SESSION_FORWARD_BYTES = 4 * MAX_SANDBOX_CONTROL_FRAME_BYTES; + +export class SessionForwardingError extends Error { + constructor( + message: string, + readonly retryable: boolean + ) { + super(message); + this.name = 'SessionForwardingError'; + } +} + +export type SessionForwardingStats = { + waiting: number; + inFlight: number; + bufferedBytes: number; + highWater: number; +}; + +export type FencedSessionForward = { + sessionId: string; + bytes: number; + deadlineAt: number; + fence: () => Promise; + forward: () => Promise; +}; + +export type SessionForwarding = { + enqueue: (sessionId: string, forward: () => Promise) => Promise; + enqueueFenced: (input: FencedSessionForward) => Promise; + stats: () => SessionForwardingStats; + get: (sessionId: string) => Promise | undefined; + values: () => IterableIterator>; + delete: (sessionId: string) => void; +}; + +export function createSessionForwarding(): SessionForwarding { + const chains = new Map>(); + const stats: SessionForwardingStats = { + waiting: 0, + inFlight: 0, + bufferedBytes: 0, + highWater: 0, + }; + + const enqueue = (sessionId: string, forward: () => Promise): Promise => { + const previous = chains.get(sessionId) ?? Promise.resolve(); + const next = previous.catch(() => undefined).then(forward); + chains.set( + sessionId, + next.then( + () => undefined, + () => undefined + ) + ); + return next; + }; + + return { + enqueue, + enqueueFenced(input: FencedSessionForward): Promise { + if (input.bytes > MAX_SANDBOX_CONTROL_FRAME_BYTES) + return Promise.reject(new SessionForwardingError('Forwarded frame is too large', false)); + if ( + stats.waiting + stats.inFlight >= SANDBOX_CONTROL_OPERATION_LIMIT || + stats.bufferedBytes + input.bytes > MAX_SESSION_FORWARD_BYTES + ) + return Promise.reject( + new SessionForwardingError('Forwarding capacity is unavailable', true) + ); + stats.waiting++; + stats.bufferedBytes += input.bytes; + stats.highWater = Math.max(stats.highWater, stats.waiting + stats.inFlight); + return enqueue(input.sessionId, async () => { + stats.waiting--; + stats.inFlight++; + try { + if ( + Date.now() >= input.deadlineAt || + !(await input.fence()) || + Date.now() >= input.deadlineAt + ) + throw new SessionForwardingError('Forwarding fence changed', false); + const result = await input.forward(); + if ( + Date.now() >= input.deadlineAt || + !(await input.fence()) || + Date.now() >= input.deadlineAt + ) + throw new SessionForwardingError('Forwarding fence changed', false); + return result; + } finally { + stats.inFlight--; + stats.bufferedBytes -= input.bytes; + } + }); + }, + stats: () => ({ ...stats }), + get: sessionId => chains.get(sessionId), + values: () => chains.values(), + delete: sessionId => chains.delete(sessionId), + }; +} 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 1b2fd782ae..f74972a428 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,7 @@ describe('sandbox control socket handler', () => { result: { protocolVersion: 1, handshakeComplete: true, - capabilities: { kiloVersionHeartbeat: true }, + capabilities: { kiloVersionHeartbeat: true, sessionOperationResults: true }, }, }) ); diff --git a/services/cloud-agent-next/src/sandbox-control/socket.ts b/services/cloud-agent-next/src/sandbox-control/socket.ts index 65697a2e8d..69eaaa1733 100644 --- a/services/cloud-agent-next/src/sandbox-control/socket.ts +++ b/services/cloud-agent-next/src/sandbox-control/socket.ts @@ -14,6 +14,8 @@ import { SANDBOX_HELLO_DEADLINE_MS, sandboxControlSocketAttachmentSchema, sandboxControlObservationSchema, + sessionOperationAuthorizationSchema, + sessionOperationDeliverySchema, type SandboxControlObservation, sessionRequestIdentitySchema, type ControlOperation, @@ -23,6 +25,9 @@ import { type SandboxHeartbeatPayload, type SessionEventIdentity, type SessionEventPayload, + type SessionOperationAck, + type SessionOperationAuthorization, + type SessionOperationDelivery, type SessionPreparingPayload, type SessionRequestIdentity, } from '../shared/sandbox-control-protocol.js'; @@ -51,8 +56,10 @@ export type SandboxControlOutboundRequest = { operation: Exclude; session?: SessionRequestIdentity; payload: unknown; + authorization?: SessionOperationAuthorization; timeoutMs?: number; expectedWrapperInstanceId?: string; + deadlineAt?: number; }; export type SandboxControlConnectionIdentity = { @@ -82,6 +89,11 @@ export type SandboxControlSocketHooks = { payload: SessionPreparingPayload, identity: SandboxControlConnectionIdentity ): void | Promise; + onOperationResult?( + session: SessionRequestIdentity, + delivery: SessionOperationDelivery, + identity: SandboxControlConnectionIdentity + ): Promise | SessionOperationAck | undefined; onSocketClosed?( handshakeComplete: boolean, identity?: SandboxControlConnectionIdentity @@ -96,6 +108,7 @@ export type SandboxControlSocketHandler = { closeHandshakenSockets(code: number, reason: string): void; sendRequest(input: SandboxControlOutboundRequest): Promise; hasHandshakenSocket(): boolean; + supportsOperationResults(): boolean; getConnectionIdentity(): SandboxControlConnectionIdentity | null; getReadySocket(): WebSocket | null; closeProvisionalSockets(): void; @@ -316,6 +329,14 @@ export function createSandboxControlSocketHandler( return currentHandshakenSocket(state) !== null; }, + supportsOperationResults(): boolean { + const current = currentHandshakenSocket(state); + return ( + current !== null && + readAttachment(current.socket)?.capabilities?.sessionOperationResults === true + ); + }, + getConnectionIdentity(): SandboxControlConnectionIdentity | null { return currentHandshakenSocket(state)?.identity ?? null; }, @@ -478,6 +499,7 @@ export function createSandboxControlSocketHandler( connectionId: identity.connectionId, protocolVersion: SANDBOX_CONTROL_PROTOCOL_VERSION, providerInstanceId: identity.providerInstanceId, + ...(payload.capabilities ? { capabilities: payload.capabilities } : {}), ...(identity.wrapperInstanceId ? { wrapperInstanceId: identity.wrapperInstanceId } : {}), }; const superseded: WebSocket[] = []; @@ -616,6 +638,66 @@ export function createSandboxControlSocketHandler( return; } + if (frame.operation === 'session.operation.result') { + let parsedSession: SessionRequestIdentity; + let parsedDelivery: SessionOperationDelivery; + try { + parsedSession = sessionRequestIdentitySchema.parse(frame.session); + parsedDelivery = sessionOperationDeliverySchema.parse(frame.payload); + } catch { + sendJson( + ws, + errorResponse( + frame.requestId, + 'protocol_error', + 'Invalid operation result payload', + false + ) + ); + return; + } + if (parsedDelivery.authorization.wrapperInstanceId !== identity.wrapperInstanceId) { + sendJson( + ws, + errorResponse(frame.requestId, 'unauthorized', 'Operation source mismatch', false) + ); + return; + } + try { + const ack = await hooks.onOperationResult?.(parsedSession, parsedDelivery, identity); + if (!isCurrentConnection(state, ws, identity)) return; + sendJson( + ws, + ack + ? okResponse(frame.requestId, ack) + : errorResponse( + frame.requestId, + 'not_ready', + 'Operation result was not acknowledged', + true + ) + ); + } catch (error) { + if (isCurrentConnection(state, ws, identity)) { + const permanent = error instanceof SandboxControlConnectionError && !error.retryable; + try { + sendJson( + ws, + errorResponse( + frame.requestId, + permanent ? 'unauthorized' : 'not_ready', + 'Operation result delivery failed', + !permanent + ) + ); + } catch { + return; + } + } + } + return; + } + if (!isControlOperation(frame.operation)) { sendJson(ws, errorResponse(frame.requestId, 'unknown_operation', 'Unknown operation')); return; @@ -683,6 +765,24 @@ export function createSandboxControlSocketHandler( if (!payload.ok) { throw new Error(payload.error.message); } + const authorization = input.authorization + ? sessionOperationAuthorizationSchema.safeParse(input.authorization) + : undefined; + if ( + authorization && + (!authorization.success || + (input.operation !== 'session.attach' && input.operation !== 'session.prompt') || + authorization.data.operation !== input.operation || + !input.session || + authorization.data.session.sessionId !== input.session.sessionId || + authorization.data.session.kiloSessionId !== input.session.kiloSessionId || + authorization.data.session.directory !== input.session.directory || + (input.expectedWrapperInstanceId !== undefined && + authorization.data.wrapperInstanceId !== input.expectedWrapperInstanceId) || + Date.now() >= authorization.data.dispatchDeadlineAt) + ) { + throw new SandboxControlConnectionError('Invalid session operation authorization', false); + } if ( isSessionOperation(input.operation) && !sessionRequestIdentitySchema.safeParse(input.session).success @@ -710,8 +810,16 @@ export function createSandboxControlSocketHandler( operation: input.operation, payload: payload.payload, ...(input.session ? { session: input.session } : {}), + ...(input.authorization ? { authorization: input.authorization } : {}), }; - const pending = waiters.wait(requestId, input.timeoutMs); + 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 pending = waiters.wait(requestId, timeoutMs); log('socket_request_sent', { ...diagnosticConnection(current.identity), requestId, diff --git a/services/cloud-agent-next/src/sandbox-control/waiters.ts b/services/cloud-agent-next/src/sandbox-control/waiters.ts index e616f650ce..71f43a4d3c 100644 --- a/services/cloud-agent-next/src/sandbox-control/waiters.ts +++ b/services/cloud-agent-next/src/sandbox-control/waiters.ts @@ -5,9 +5,11 @@ import { export class SandboxControlConnectionError extends Error { readonly code = 'not_ready'; - readonly retryable = true; - constructor(message: string) { + constructor( + message: string, + readonly retryable = true + ) { super(message); this.name = 'SandboxControlConnectionError'; } diff --git a/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts b/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts index 12a7a39d03..f549ab7fa6 100644 --- a/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts +++ b/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts @@ -102,16 +102,24 @@ import { } from '../session/preparation-history.js'; import { SANDBOX_CONTROL_ATTACH_TIMEOUT_MS, + SANDBOX_CONTROL_OUTCOME_TIMEOUT_MS, SANDBOX_CONTROL_REQUEST_TIMEOUT_MS, controlErrorCodes, sessionAttachResultSchema, sessionMessageOutcomeSchema, + sessionOperationAuthorizationSchema, + sessionOperationDeliverySchema, + sessionOperationExpiresAt, + sessionOperationResultHash, sessionPromptResultSchema, sessionSyncResultSchema, sessionPermissionResolveResultSchema, sessionQuestionResolveResultSchema, sessionAbortResultSchema, wrapperInstanceIdSchema, + type SessionOperationAck, + type SessionOperationAuthorization, + type SessionRequestIdentity, type SessionSyncResult, type SessionEventIdentity, type SessionPreparingPayload, @@ -161,6 +169,7 @@ import { type ControlSessionMessageInput, type SessionMessageRecord, } from './session-message-queue.js'; +import { commitSessionOperationResult, dispatchSessionOperation } from './session-operation.js'; const METADATA_KEY = SANDBOX_SESSION_METADATA_KEY; const MESSAGES_KEY = 'session_messages'; @@ -440,6 +449,13 @@ export class SandboxSession extends DurableObject { fromState: existing?.state, outcome: outcome.data.status, }; + const proof = existing?.operations?.prompt ?? existing?.operations?.attach; + if (proof?.dispatched) + return result( + false, + proof.resultHash === undefined ? 'operation_result_pending' : 'operation_result_required', + diagnostic + ); if ( existing?.wrapperInstanceId === input.wrapperInstanceId && existing.state === outcome.data.status @@ -632,6 +648,89 @@ export class SandboxSession extends DurableObject { return result(true, 'processed'); } + receiveSandboxOperationResult(input: { + session: SessionRequestIdentity; + wrapperInstanceId: string; + delivery: unknown; + }): Promise { + return this.trackOperation(this.applySandboxOperationResult(input)); + } + + private async applySandboxOperationResult(input: { + session: SessionRequestIdentity; + wrapperInstanceId: string; + delivery: unknown; + }): Promise { + const parsed = sessionOperationDeliverySchema.safeParse(input.delivery); + if (!parsed.success) return undefined; + const delivery = parsed.data; + const authorization = delivery.authorization; + const deadlineAt = Math.min( + sessionOperationExpiresAt(authorization), + delivery.completedAt + SANDBOX_CONTROL_OUTCOME_TIMEOUT_MS + ); + if ( + delivery.completedAt > Date.now() + SANDBOX_CONTROL_REQUEST_TIMEOUT_MS || + authorization.wrapperInstanceId !== input.wrapperInstanceId || + authorization.session.sessionId !== input.session.sessionId || + authorization.session.kiloSessionId !== input.session.kiloSessionId || + authorization.session.directory !== input.session.directory || + Date.now() >= sessionOperationExpiresAt(authorization) + ) + return undefined; + const metadata = await this.getMetadata(); + const epoch = this.terminalLifecycle.captureEpoch(); + if ( + !metadata || + epoch === null || + input.session.sessionId !== this.sessionId || + input.session.kiloSessionId !== metadata.auth.kiloSessionId || + input.session.directory !== this.directory(metadata) + ) + return undefined; + const hash = await sessionOperationResultHash(delivery); + const notifications: StoredEvent[] = []; + const ack = commitSessionOperationResult({ + storage: this.ctx.storage, + delivery, + hash, + deadlineAt, + isCurrent: () => { + const current = this.terminalLifecycle.getStoredMetadata(); + return ( + this.terminalLifecycle.isCurrent(epoch) && + current !== null && + input.session.sessionId === this.sessionId && + input.session.kiloSessionId === current.auth.kiloSessionId && + input.session.directory === this.directory(current) + ); + }, + messages: { + read: () => this.loadMessages(), + commit: messages => + this.saveMessagesInCurrentTransaction(messages, epoch, 'operation_result', notifications), + }, + eventQueries: this.eventQueries, + notifications, + }); + if (!ack) return undefined; + for (const notification of notifications) this.broadcastStoredEvent(notification); + if (ack.disposition === 'applied' && delivery.outcome) { + if (this.isCurrentEventRuntime(input.wrapperInstanceId)) + this.worktreeChanges.onEvent( + this.worktreeContext(metadata), + metadata.auth.kiloSessionId, + 'session.message.outcome', + delivery.outcome + ); + await this.armQueueRetry(); + const nextId = nextQueuedMessageId(this.loadMessages()); + if (nextId && this.terminalLifecycle.isCurrent(epoch)) + this.ctx.waitUntil(this.dispatchQueued(nextId)); + } + return ack; + } + async closeOrgStreams(organizationId: string): Promise { const metadata = this.terminalLifecycle.getStoredMetadata(); if (!metadata?.identity.orgId || metadata.identity.orgId !== organizationId) return 0; @@ -1272,6 +1371,45 @@ export class SandboxSession extends DurableObject { return { success: false, code: 'BAD_REQUEST', error: 'Prepared admission is legacy-only' }; } + private async observeAcceptedOperation( + message: MessageRecord, + epoch: number + ): Promise<'running' | 'completed' | 'response' | undefined> { + const authorization = sessionOperationAuthorizationSchema.safeParse( + message.operations?.prompt?.authorization + ); + if (!authorization.success || !message.operations?.prompt?.dispatched) return undefined; + const metadata = await this.getMetadata(); + const sandboxId = metadata?.workspace?.sandboxId; + if (!metadata || !sandboxId) throw new Error('Accepted runtime is unavailable'); + const dispatched = await dispatchSessionOperation( + { authorization: authorization.data, payload: undefined }, + { + read: () => this.loadMessages(), + commit: messages => this.saveMessages(messages, epoch, 'operation_result'), + }, + { + request: input => sandboxControlRpc(this.env, sandboxId).request(input), + persistResult: delivery => + this.applySandboxOperationResult({ + session: authorization.data.session, + wrapperInstanceId: authorization.data.wrapperInstanceId, + delivery, + }), + isDispatchCurrent: () => false, + isMaintenanceCurrent: () => { + const current = this.loadMessages().find(item => item.messageId === message.messageId); + return ( + this.terminalLifecycle.isCurrent(epoch) && + current?.wrapperInstanceId === authorization.data.wrapperInstanceId && + current.operations?.prompt?.dispatched === true + ); + }, + } + ); + return dispatched.state; + } + async alarm(): Promise { if (this.pendingRuntimeCleanup()) await this.transferRuntimeCleanup(); const epoch = this.terminalLifecycle.captureEpoch(); @@ -1309,6 +1447,14 @@ export class SandboxSession extends DurableObject { ); logControlDiagnostic('accepted_reconciliation', { ...diagnostic, phase: 'started' }); try { + if (accepted.operations?.prompt?.dispatched) { + diagnostic.stage = 'operation_receipt'; + const observed = await this.observeAcceptedOperation(accepted, epoch); + if (observed === 'running' || observed === 'completed') { + report('healthy'); + return; + } + } const snapshot = await this.interactionRefresh.refresh(scope, 'accepted_alarm'); if ( !snapshot || @@ -1561,6 +1707,13 @@ export class SandboxSession extends DurableObject { const recordRuntime = (identity: string | undefined) => { const runtime = wrapperInstanceIdSchema.safeParse(identity); if (!runtime.success) return; + const current = this.loadMessages().find(message => message.messageId === messageId); + if ( + current?.wrapperInstanceId !== undefined && + current.wrapperInstanceId !== runtime.data && + (current.operations?.attach?.dispatched || current.operations?.prompt?.dispatched) + ) + return; wrapperInstanceId = runtime.data; this.saveMessages( this.loadMessages().map(message => @@ -1610,9 +1763,50 @@ export class SandboxSession extends DurableObject { throw error; } }; + const dispatchAuthorized = async ( + operation: SessionOperationAuthorization['operation'], + payload: unknown + ) => { + if (!wrapperInstanceId) throw new Error('Wrapper identity is missing'); + const authorization: SessionOperationAuthorization = { + operation, + operationId: operation === 'session.attach' ? assigned.attemptId : messageId, + messageId, + session: { sessionId, kiloSessionId, directory: this.directory(metadata) }, + wrapperInstanceId, + dispatchDeadlineAt: deadlineAt, + }; + const dispatched = await dispatchSessionOperation( + { authorization, payload }, + { + read: () => this.loadMessages(), + commit: messages => this.saveMessages(messages, epoch, 'wrapper_outcome'), + }, + { + request: input => control.request(input), + persistResult: delivery => + this.applySandboxOperationResult({ + session: authorization.session, + wrapperInstanceId: authorization.wrapperInstanceId, + delivery, + }), + isDispatchCurrent: isCurrent, + isMaintenanceCurrent: () => { + if (!this.terminalLifecycle.isCurrent(epoch)) return false; + const current = this.loadMessages().find(message => message.messageId === messageId); + if (!current || current.wrapperInstanceId !== wrapperInstanceId) return false; + return ( + current.operations?.[operation === 'session.attach' ? 'attach' : 'prompt'] + ?.dispatched === true + ); + }, + } + ); + return dispatched; + }; await this.armQueueRetry(Math.min(deadlineAt, Date.now() + QUEUE_RETRY_MS)); if (!isCurrent()) return; - if (Date.now() >= deadlineAt) { + if (Date.now() >= deadlineAt && !queued.operations?.prompt?.dispatched) { await this.failDelivery( messageId, 'preparation_timeout', @@ -1741,6 +1935,16 @@ export class SandboxSession extends DurableObject { return; } if (!wrapperInstanceId) throw new Error('Wrapper identity is missing'); + const operationResults = status.operationResults === true; + const proofBacked = + queued.operations?.attach?.dispatched === true || + queued.operations?.prompt?.dispatched === true; + if (proofBacked && !operationResults) + throw new ControlRequestError({ + code: 'runtime_unhealthy', + message: 'Wrapper operation receipt capability is unavailable', + retryable: false, + }); const needsPreparation = this.terminalLifecycle.getAttachedWrapperInstanceId() !== wrapperInstanceId; if (needsPreparation || status.attachment?.kilo?.containmentEnabled === false) { @@ -1770,23 +1974,30 @@ export class SandboxSession extends DurableObject { await this.compensateSessionAttachment(metadata); return; } - await dispatch('attach', () => - wait( - async () => - sessionAttachResultSchema.parse( - controlRequestResult( - await control.request({ - operation: 'session.attach', - session, - expectedWrapperInstanceId: wrapperInstanceId, - payload: attachPayload, - timeoutMs: SANDBOX_CONTROL_ATTACH_TIMEOUT_MS, - }) - ) - ), - SANDBOX_CONTROL_ATTACH_TIMEOUT_MS - ) - ); + if (operationResults) { + if ((await dispatchAuthorized('session.attach', attachPayload)).state === 'running') { + await this.armQueueRetry(Math.min(deadlineAt, Date.now() + QUEUE_RETRY_MS)); + return; + } + } else { + await dispatch('attach', () => + wait( + async () => + sessionAttachResultSchema.parse( + controlRequestResult( + await control.request({ + operation: 'session.attach', + session, + expectedWrapperInstanceId: wrapperInstanceId, + payload: attachPayload, + timeoutMs: SANDBOX_CONTROL_ATTACH_TIMEOUT_MS, + }) + ) + ), + SANDBOX_CONTROL_ATTACH_TIMEOUT_MS + ) + ); + } if (!isCurrent()) { if (!this.terminalLifecycle.isCurrent(epoch)) await this.compensateSessionAttachment(metadata); @@ -1805,45 +2016,58 @@ export class SandboxSession extends DurableObject { recorder.finalize({ status: 'completed' }); this.worktreeChanges.attached(preparationGeneration, this.worktreeContext(metadata)); phase = 'prompt'; - await dispatch('prompt', async () => { - const prompt = await wait(async () => - controlRequestResult( - await control.request({ - operation: 'session.prompt', - session, - expectedWrapperInstanceId: wrapperInstanceId, - payload: { - messageId, - turn: - intent.turn.type === 'command' - ? { - type: 'command', - command: intent.turn.command, - arguments: intent.turn.arguments, - } - : { type: 'prompt', prompt: intent.turn.prompt }, - agent: { - mode: intent.agent.mode, - ...(model !== undefined ? { model } : {}), - ...(intent.agent.variant !== undefined ? { variant: intent.agent.variant } : {}), - }, - ...(intent.finalization ? { finalization: intent.finalization } : {}), - ...(attachments.length ? { attachments } : {}), - }, - }) - ) - ); - const result = sessionPromptResultSchema.parse(prompt); - if (result.messageId !== messageId) - throw new Error('Prompt response message identity mismatch'); - }); + const promptPayload = { + messageId, + turn: + intent.turn.type === 'command' + ? { + type: 'command' as const, + command: intent.turn.command, + arguments: intent.turn.arguments, + } + : { type: 'prompt' as const, prompt: intent.turn.prompt }, + agent: { + mode: intent.agent.mode, + ...(model !== undefined ? { model } : {}), + ...(intent.agent.variant !== undefined ? { variant: intent.agent.variant } : {}), + }, + ...(intent.finalization ? { finalization: intent.finalization } : {}), + ...(attachments.length ? { attachments } : {}), + }; + if (operationResults) { + const dispatched = await dispatchAuthorized('session.prompt', promptPayload); + if (dispatched.state === 'completed') return; + if (dispatched.state === 'running') { + const accepted = acceptQueuedMessage(this.loadMessages(), messageId, Date.now()); + if (!accepted || !this.saveMessages(accepted, epoch)) return; + await this.armQueueRetry(Date.now() + DEADLINE_MS.acceptedAlarmCap); + return; + } + } else { + await dispatch('prompt', async () => { + const prompt = await wait(async () => + controlRequestResult( + await control.request({ + operation: 'session.prompt', + session, + expectedWrapperInstanceId: wrapperInstanceId, + payload: promptPayload, + }) + ) + ); + const result = sessionPromptResultSchema.parse(prompt); + if (result.messageId !== messageId) + throw new Error('Prompt response message identity mismatch'); + }); + } if (!isCurrent()) { if (!this.terminalLifecycle.isCurrent(epoch)) await this.compensateSessionAttachment(metadata); return; } const accepted = acceptQueuedMessage(this.loadMessages(), messageId, Date.now()); - if (!accepted || !this.saveMessages(accepted, epoch)) return; + if (!accepted) return; + if (!this.saveMessages(accepted, epoch)) return; await this.armQueueRetry(Date.now() + DEADLINE_MS.acceptedAlarmCap); } catch (error) { if (!isCurrent()) { @@ -2509,7 +2733,31 @@ export class SandboxSession extends DurableObject { private saveMessages( messages: MessageRecord[], epoch?: number, - source: 'coordinator' | 'wrapper_outcome' = 'coordinator' + source: 'coordinator' | 'wrapper_outcome' | 'operation_result' = 'coordinator', + deferredNotifications?: StoredEvent[] + ): boolean { + return this.commitSavedMessages(messages, epoch, source, deferredNotifications, write => + this.ctx.storage.transactionSync(write) + ); + } + + private saveMessagesInCurrentTransaction( + messages: MessageRecord[], + epoch: number | undefined, + source: 'coordinator' | 'wrapper_outcome' | 'operation_result', + deferredNotifications?: StoredEvent[] + ): boolean { + return this.commitSavedMessages(messages, epoch, source, deferredNotifications, write => + write() + ); + } + + private commitSavedMessages( + messages: MessageRecord[], + epoch: number | undefined, + source: 'coordinator' | 'wrapper_outcome' | 'operation_result', + deferredNotifications: StoredEvent[] | undefined, + enclose: (write: () => void) => void ): boolean { const currentEpoch = epoch ?? this.terminalLifecycle.captureEpoch(); if ( @@ -2520,7 +2768,7 @@ export class SandboxSession extends DurableObject { return false; const events: StoredEvent[] = []; const committed: ControlDiagnosticFields[] = []; - this.ctx.storage.transactionSync(() => { + const write = () => { const before = this.loadMessages(); const previousById = new Map(before.map(message => [message.messageId, message])); const queuedHeadId = nextQueuedMessageId(before); @@ -2544,7 +2792,11 @@ export class SandboxSession extends DurableObject { } return message; } - const terminal = { ...message, terminalAt: message.terminalAt ?? now }; + const terminal = { + ...message, + terminalAt: message.terminalAt ?? now, + terminalSource: message.terminalSource ?? source, + }; if (previous?.state === 'accepted' || queuedHeadId === message.messageId) { const interactions = this.readPendingInteractions(); this.ctx.storage.kv.put(PENDING_INTERACTIONS_KEY, { @@ -2585,7 +2837,8 @@ export class SandboxSession extends DurableObject { return terminal; }); this.ctx.storage.kv.put(MESSAGES_KEY, next); - }); + }; + enclose(write); for (const fields of committed) { logControlDiagnostic('session_message_committed', { sessionId: this.sessionId, @@ -2593,7 +2846,8 @@ export class SandboxSession extends DurableObject { ...fields, }); } - for (const event of events) this.broadcastStoredEvent(event); + if (deferredNotifications) deferredNotifications.push(...events); + else for (const event of events) this.broadcastStoredEvent(event); return true; } 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 6dcdfca266..159347c206 100644 --- a/services/cloud-agent-next/src/sandbox-session/control-rpc.ts +++ b/services/cloud-agent-next/src/sandbox-session/control-rpc.ts @@ -31,12 +31,14 @@ type SandboxControlRpc = { connection: ConnectionState; physical: PhysicalState; wrapperInstanceId?: string; + operationResults?: true; attachment?: SessionAttachPayload; }>; getStatus(): Promise<{ connection: ConnectionState; physical: PhysicalState; wrapperInstanceId?: string; + operationResults?: true; }>; quarantineRuntime(input: { ownerId: string; 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 4cc3daddbb..c350c4cd3e 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 @@ -21,6 +21,7 @@ import { type ResponseFrame, type SessionAttachPayload, type SessionMessageOutcome, + type SessionOperationAuthorization, } from '../shared/sandbox-control-protocol.js'; import { DEADLINE_MS } from '../sandbox-control/deadlines.js'; import { createControlPlaneCredential } from '../sandbox-control/managed-credential.js'; @@ -241,6 +242,7 @@ describe('createSessionMessageRecord', () => { lastActivityAt: 20, wrapperInstanceId: 'runtime', terminalAt: 30, + terminalSource: 'wrapper_outcome', }, ]); }); @@ -2491,6 +2493,61 @@ describe('SandboxSession orchestration', () => { ).toHaveLength(0); }); + it('persists a retained operation result without nested transactionSync', async () => { + const fixture = sessionFixture(); + const kiloSessionId = fixture.metadata.auth.kiloSessionId; + if (!kiloSessionId) throw new Error('Missing Kilo session ID'); + const authorization = { + operation: 'session.prompt', + operationId: 'a', + messageId: 'a', + session: { sessionId: SESSION_ID, kiloSessionId, directory: DIRECTORY }, + wrapperInstanceId: RUNTIME_ID, + dispatchDeadlineAt: Date.now() + 60_000, + } satisfies SessionOperationAuthorization; + fixture.values.set('session_messages', [ + { + messageId: 'a', + state: 'accepted', + wrapperInstanceId: RUNTIME_ID, + operations: { prompt: { authorization, dispatched: true } }, + } as SessionMessageRecord, + ]); + let depth = 0; + const transactionSync = fixture.storage.transactionSync.bind(fixture.storage); + fixture.storage.transactionSync = (callback: () => T): T => { + if (depth > 0) throw new Error('nested transactionSync'); + depth += 1; + try { + return transactionSync(callback); + } finally { + depth -= 1; + } + }; + await expect( + fixture.session.receiveSandboxOperationResult({ + session: authorization.session, + wrapperInstanceId: authorization.wrapperInstanceId, + delivery: { + version: 2, + authorization, + completedAt: Date.now(), + result: { ok: true, result: { messageId: 'a', status: 'accepted' } }, + outcome: { messageId: 'a', status: 'completed' }, + events: [ + { + type: 'autocommit_completed', + properties: { success: true, messageId: 'a' }, + timestamp: new Date().toISOString(), + }, + ], + preparing: [], + }, + }) + ).resolves.toMatchObject({ disposition: 'applied' }); + expect(fixture.record('a')?.state).toBe('completed'); + }); + it('delivers a follow-up after awaited cancel, failed quarantine transfer, reset, and old-runtime cleanup', async () => { const fixture = sessionFixture(); const attach = deferred(); 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 6ec592f827..2efbe466c9 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 @@ -10,9 +10,17 @@ import { } from '../execution/types.js'; import { dispatchedKilocodeModelId } from '../persistence/model-utils.js'; import type { CloudMessageFailedPayload } from '../session/message-settlement-outbox.js'; -import type { SessionMessageOutcome } from '../shared/sandbox-control-protocol.js'; +import { + sessionOperationAuthorizationSchema, + sameSessionOperation, + type SessionMessageOutcome, + type SessionOperationAck, + type SessionOperationAuthorization, + type SessionOperationDelivery, +} from '../shared/sandbox-control-protocol.js'; export type SessionMessageState = 'queued' | 'accepted' | 'completed' | 'failed' | 'cancelled'; +export type SessionMessageTerminalSource = 'coordinator' | 'wrapper_outcome' | 'operation_result'; export type ControlCommandAgentSelection = Omit & { model?: string }; @@ -27,6 +35,15 @@ export type ControlSessionMessageInput = Pick item.messageId === outcome.messageId); if ( @@ -295,6 +318,7 @@ export function applyMessageOutcome( unresolvedDispatch: undefined, acceptedAt: item.acceptedAt ?? now, terminalAt: now, + terminalSource, ...(outcome.reason ? { failedReason: outcome.reason } : {}), } : item @@ -404,3 +428,150 @@ export function streamCloudStatus( if (messages.some(message => message.state === 'queued')) return { type: 'preparing' }; return messages.length > 0 ? { type: 'ready' } : null; } + +export function applySessionOperationResult( + messages: readonly SessionMessageRecord[], + delivery: SessionOperationDelivery, + resultHash: string, + now: number +): + | { + messages: SessionMessageRecord[]; + disposition: SessionOperationAck['disposition']; + decision: SessionOperationAck['decision']; + } + | undefined { + const authorization = delivery.authorization; + const message = messages.find(item => item.messageId === authorization.messageId); + const kind = authorization.operation === 'session.attach' ? 'attach' : 'prompt'; + const proof = message?.operations?.[kind]; + const storedAuthorization = sessionOperationAuthorizationSchema.safeParse(proof?.authorization); + if ( + !message || + !proof?.dispatched || + !storedAuthorization.success || + message.wrapperInstanceId !== authorization.wrapperInstanceId || + !sameSessionOperation(storedAuthorization.data, authorization) + ) + return undefined; + if (message.state !== 'queued' && message.state !== 'accepted') { + if (message.terminalAt === undefined) return undefined; + return { + messages: [...messages], + disposition: + proof.resultHash === resultHash + ? 'identical' + : message.terminalSource === 'coordinator' + ? 'superseded' + : 'already_final', + decision: { state: message.state, at: message.terminalAt }, + }; + } + if (proof.resultHash !== undefined) { + return proof.resultHash === resultHash && proof.decision + ? { messages: [...messages], disposition: 'identical', decision: proof.decision } + : undefined; + } + const applied = delivery.outcome + ? applyMessageOutcome( + messages, + delivery.outcome, + authorization.wrapperInstanceId, + now, + 'operation_result' + ) + : [...messages]; + if (!applied) return undefined; + const resultMessage = applied.find(item => item.messageId === message.messageId); + if (!resultMessage) return undefined; + const decision = { + state: resultMessage.state, + at: resultMessage.terminalAt ?? delivery.completedAt, + }; + return { + messages: applied.map(item => + item.messageId === message.messageId + ? { + ...item, + operations: { + ...item.operations, + [kind]: { + ...proof, + result: delivery.result, + resultHash, + completedAt: delivery.completedAt, + decision, + }, + }, + } + : item + ), + disposition: 'applied', + decision, + }; +} + +export function recordSessionOperationDispatch( + messages: readonly SessionMessageRecord[], + authorization: SessionOperationAuthorization +): SessionMessageRecord[] | undefined { + const message = messages.find(item => item.messageId === authorization.messageId); + const kind = authorization.operation === 'session.attach' ? 'attach' : 'prompt'; + const proof = message?.operations?.[kind]; + const storedAuthorization = sessionOperationAuthorizationSchema.safeParse(proof?.authorization); + if ( + !message || + nextQueuedMessageId(messages) !== message.messageId || + message.wrapperInstanceId !== authorization.wrapperInstanceId || + (proof && + (!storedAuthorization.success || + !sameSessionOperation(storedAuthorization.data, authorization))) + ) + return undefined; + return messages.map(item => + item.messageId === message.messageId + ? { + ...item, + unresolvedDispatch: true, + deliveryRetryScope: undefined, + operations: { + ...item.operations, + [kind]: { + authorization: structuredClone(authorization), + dispatched: true, + }, + }, + } + : item + ); +} + +export function completeSessionOperationAttachment( + messages: readonly SessionMessageRecord[], + authorization: SessionOperationAuthorization +): SessionMessageRecord[] | undefined { + const message = messages.find(item => item.messageId === authorization.messageId); + const proof = message?.operations?.attach; + const storedAuthorization = sessionOperationAuthorizationSchema.safeParse(proof?.authorization); + if ( + authorization.operation !== 'session.attach' || + !message || + !proof?.dispatched || + !storedAuthorization.success || + !sameSessionOperation(storedAuthorization.data, authorization) || + nextQueuedMessageId(messages) !== message.messageId + ) + return undefined; + return messages.map(item => + item.messageId === message.messageId + ? { + ...item, + unresolvedDispatch: undefined, + operations: { + ...item.operations, + attach: { ...proof, completedAt: proof.completedAt ?? Date.now() }, + }, + } + : item + ); +} 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 new file mode 100644 index 0000000000..7846bce933 --- /dev/null +++ b/services/cloud-agent-next/src/sandbox-session/session-operation.test.ts @@ -0,0 +1,357 @@ +import { describe, expect, it, vi } from 'vitest'; +import { + sessionOperationResultHash, + type ResponseFrame, + 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, + createSessionMessageRecord, + recordSessionOperationDispatch, + type SessionMessageRecord, +} from './session-message-queue.js'; +import { commitSessionOperationResult, dispatchSessionOperation } from './session-operation.js'; + +const authorization: SessionOperationAuthorization = { + operation: 'session.prompt', + operationId: 'msg_operation_1', + messageId: 'msg_operation_1', + session: { + sessionId: 'workspace_operation_1', + kiloSessionId: 'kilo_operation_1', + directory: '/workspace/operation', + }, + wrapperInstanceId: '11111111-1111-4111-8111-111111111111', + dispatchDeadlineAt: Date.now() + 60_000, +}; + +const payload = { + messageId: authorization.messageId, + turn: { type: 'prompt' as const, prompt: 'durably deliver this prompt' }, + agent: { mode: 'code' as const, model: 'kilo/openai/gpt-4.1' }, +}; + +function response(result: unknown): ResponseFrame { + return { type: 'response', requestId: crypto.randomUUID(), ok: true, result }; +} + +function messages(): SessionMessageRecord[] { + return [ + { + ...createSessionMessageRecord({ + turn: { type: 'prompt', messageId: authorization.messageId, prompt: payload.turn.prompt }, + agent: payload.agent, + }), + wrapperInstanceId: authorization.wrapperInstanceId, + }, + ]; +} + +describe('dispatchSessionOperation', () => { + it('writes the immutable prompt proof before the first wrapper request', async () => { + let stored = messages(); + const request = vi.fn(async (input: SandboxControlOutboundRequest) => { + expect(input).toMatchObject({ operation: 'session.prompt', authorization, payload }); + expect(stored).toMatchObject([ + { + unresolvedDispatch: true, + operations: { prompt: { dispatched: true, authorization } }, + }, + ]); + return response({ messageId: authorization.messageId, status: 'accepted' }); + }); + + await expect( + dispatchSessionOperation( + { authorization, payload }, + { + read: () => stored, + commit: next => { + stored = next; + return true; + }, + }, + { + request, + persistResult: async () => undefined, + isDispatchCurrent: () => true, + isMaintenanceCurrent: () => true, + } + ) + ).resolves.toEqual({ + state: 'response', + result: { messageId: authorization.messageId, status: 'accepted' }, + }); + expect(request).toHaveBeenCalledTimes(1); + }); + + it('uses the retained result and exact acknowledgement after a lost prompt response', 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 }); + expect(input).toMatchObject({ operation: 'session.operation.ack', payload: ack }); + return response({ acknowledged: true }); + }); + + await expect( + dispatchSessionOperation( + { authorization, payload }, + { + read: () => stored, + commit: next => { + stored = next; + return true; + }, + }, + { + request, + persistResult: async () => ack, + isDispatchCurrent: () => false, + isMaintenanceCurrent: () => true, + } + ) + ).resolves.toEqual({ state: 'completed' }); + expect(request.mock.calls.map(([input]) => input.operation)).toEqual([ + 'session.operation.get', + 'session.operation.ack', + ]); + }); + + it('keeps a positively running operation without replaying its prompt', async () => { + const dispatched = recordSessionOperationDispatch(messages(), authorization); + if (!dispatched) throw new Error('Failed to create dispatch proof'); + const request = vi.fn(async (_input: SandboxControlOutboundRequest) => + response({ + state: 'running', + authorization, + }) + ); + + await expect( + dispatchSessionOperation( + { authorization, payload }, + { read: () => dispatched, commit: () => true }, + { + request, + persistResult: async () => undefined, + isDispatchCurrent: () => false, + isMaintenanceCurrent: () => true, + } + ) + ).resolves.toEqual({ state: 'running' }); + expect(request.mock.calls.map(([input]) => input.operation)).toEqual(['session.operation.get']); + }); + + it('does not replay a mutation when the retained operation is missing', async () => { + const dispatched = recordSessionOperationDispatch(messages(), authorization); + if (!dispatched) throw new Error('Failed to create dispatch proof'); + const request = vi.fn(async (_input: SandboxControlOutboundRequest) => + response({ state: 'missing' }) + ); + + await expect( + dispatchSessionOperation( + { authorization, payload }, + { read: () => dispatched, commit: () => true }, + { + request, + persistResult: async () => undefined, + isDispatchCurrent: () => true, + isMaintenanceCurrent: () => true, + } + ) + ).rejects.toThrow('Original session operation is missing'); + expect(request.mock.calls.map(([input]) => input.operation)).toEqual(['session.operation.get']); + }); + + 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'); + 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 resultHash = await sessionOperationResultHash(delivery); + const applied = applySessionOperationResult(dispatched, delivery, resultHash, Date.now()); + if (!applied) throw new Error('Failed to apply operation result'); + expect(applied).toMatchObject({ + disposition: 'applied', + messages: [ + { + state: 'completed', + unresolvedDispatch: undefined, + terminalSource: 'operation_result', + operations: { prompt: { resultHash } }, + }, + ], + }); + + expect( + applySessionOperationResult(applied.messages, delivery, resultHash, Date.now()) + ).toMatchObject({ + disposition: 'identical', + }); + expect( + applySessionOperationResult( + applied.messages, + { + ...delivery, + result: { + ok: false, + error: { code: 'git_failed', message: 'conflict', retryable: false }, + }, + }, + 'f'.repeat(64), + Date.now() + ) + ).toMatchObject({ disposition: 'already_final' }); + expect(applied.messages[0]?.operations?.prompt?.resultHash).toBe(resultHash); + }); + + it.each(['event', 'message'] as const)( + 'does not acknowledge or publish a failed %s transaction', + async failure => { + 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: [ + { + type: 'autocommit_completed', + properties: { success: true, messageId: authorization.messageId }, + timestamp: new Date().toISOString(), + }, + ], + preparing: [], + }; + const hash = await sessionOperationResultHash(delivery); + const notifications: StoredEvent[] = []; + const commit = vi.fn((next: SessionMessageRecord[]) => { + if (failure === 'message') return false; + stored = next; + return true; + }); + const eventQueries = { + upsert: vi.fn(() => { + if (failure === 'event') throw new Error('event write failed'); + return 1; + }), + insert: vi.fn(() => 1), + } as unknown as EventQueries; + + expect(() => + commitSessionOperationResult({ + storage: { transactionSync: callback => callback() }, + delivery, + hash, + deadlineAt: Date.now() + 1_000, + isCurrent: () => true, + messages: { read: () => stored, commit }, + eventQueries, + notifications, + }) + ).toThrow(failure === 'event' ? 'event write failed' : 'Operation result was not persisted'); + expect(notifications).toEqual([]); + expect(stored).toEqual(dispatched); + + const acknowledgement = commitSessionOperationResult({ + storage: { transactionSync: callback => callback() }, + delivery, + hash, + deadlineAt: Date.now() + 1_000, + isCurrent: () => true, + messages: { + read: () => stored, + commit: next => { + stored = next; + return true; + }, + }, + eventQueries: { upsert: () => 1, insert: () => 1 } as unknown as EventQueries, + notifications, + }); + expect(acknowledgement).toMatchObject({ disposition: 'applied', resultHash: hash }); + expect(stored[0]?.operations?.prompt?.resultHash).toBe(hash); + expect(notifications).toHaveLength(1); + } + ); + + it('invokes message commit inside the only result transaction', async () => { + const dispatched = recordSessionOperationDispatch(messages(), authorization); + if (!dispatched) throw new Error('Failed to create dispatch proof'); + let stored = dispatched; + let depth = 0; + let commitDepth: number | undefined; + const storage = { + transactionSync: (callback: () => T): T => { + depth += 1; + try { + return callback(); + } finally { + depth -= 1; + } + }, + }; + 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 hash = await sessionOperationResultHash(delivery); + const acknowledgement = commitSessionOperationResult({ + storage, + delivery, + hash, + deadlineAt: Date.now() + 1_000, + isCurrent: () => true, + messages: { + read: () => stored, + commit: next => { + commitDepth = depth; + stored = next; + return true; + }, + }, + notifications: [], + }); + expect(acknowledgement).toMatchObject({ disposition: 'applied', resultHash: hash }); + expect(commitDepth).toBe(1); + }); +}); diff --git a/services/cloud-agent-next/src/sandbox-session/session-operation.ts b/services/cloud-agent-next/src/sandbox-session/session-operation.ts new file mode 100644 index 0000000000..eb92488123 --- /dev/null +++ b/services/cloud-agent-next/src/sandbox-session/session-operation.ts @@ -0,0 +1,243 @@ +import { + SANDBOX_CONTROL_ATTACH_TIMEOUT_MS, + SANDBOX_CONTROL_OUTCOME_TIMEOUT_MS, + SANDBOX_CONTROL_REQUEST_TIMEOUT_MS, + sessionAttachPayloadSchema, + sessionAttachResultSchema, + sessionOperationAuthorizationSchema, + sessionOperationAckSchema, + sessionOperationExpiresAt, + sessionOperationLookupResultSchema, + sessionPromptPayloadSchema, + sessionPromptResultSchema, + sameSessionOperation, + 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 { + ControlRequestError, + controlRequestResult, + withDeliveryDeadline, +} from './control-dispatch.js'; + +type OperationMessages = { + read: () => SessionMessageRecord[]; + commit: (messages: SessionMessageRecord[]) => boolean; +}; + +export type SessionOperationEffects = { + request: (input: SandboxControlOutboundRequest) => Promise; + persistResult: (delivery: SessionOperationDelivery) => Promise; + isDispatchCurrent: () => boolean; + isMaintenanceCurrent: () => boolean; +}; + +export type SessionOperationDispatch = + | { state: 'response'; result: unknown } + | { state: 'running' } + | { state: 'completed' }; + +function uncertainOperation(reason: string): ControlRequestError { + return new ControlRequestError({ code: 'runtime_unhealthy', message: reason, retryable: false }); +} + +export async function dispatchSessionOperation( + input: { authorization: SessionOperationAuthorization; payload: unknown }, + messages: OperationMessages, + effects: SessionOperationEffects +): 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 assertMaintenanceCurrent = () => { + if (!effects.isMaintenanceCurrent() || Date.now() >= sessionOperationExpiresAt(authorization)) + throw uncertainOperation('Session operation maintenance authority expired'); + }; + 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' }; + } + 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 + ) + ); + if (kind === 'attach') { + const attached = sessionAttachResultSchema.safeParse( + lookup.delivery.result.ok ? lookup.delivery.result.result : undefined + ); + 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 }; + } + return { state: 'completed' }; + } + + 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 commitSessionOperationResult(input: { + storage: Pick; + delivery: SessionOperationDelivery; + hash: string; + deadlineAt: number; + isCurrent: () => boolean; + messages: OperationMessages; + eventQueries?: EventQueries; + notifications: StoredEvent[]; +}): SessionOperationAck | undefined { + const { delivery, messages, notifications } = input; + const authorization = delivery.authorization; + const committedNotifications: StoredEvent[] = []; + const acknowledgement = input.storage.transactionSync(() => { + if (!input.isCurrent() || Date.now() >= sessionOperationExpiresAt(authorization)) return; + const current = messages.read(); + const message = current.find(item => item.messageId === authorization.messageId); + if ( + Date.now() >= input.deadlineAt && + (message?.state === 'queued' || message?.state === 'accepted') + ) + return; + const applied = applySessionOperationResult(current, delivery, input.hash, Date.now()); + if (!applied) return; + if (applied.disposition === 'applied') { + if (input.eventQueries) { + for (const payload of delivery.events) + persistSandboxControlSessionEvent({ + sessionId: authorization.session.sessionId, + payload, + eventQueries: input.eventQueries, + broadcast: event => committedNotifications.push(event), + }); + for (const data of delivery.preparing) + applyControlPlanePreparingEvent({ + sessionId: authorization.session.sessionId, + data, + eventQueries: input.eventQueries, + broadcast: event => committedNotifications.push(event), + }); + } + if (!messages.commit(applied.messages)) throw new Error('Operation result was not persisted'); + } + return sessionOperationAckSchema.parse({ + version: 2, + authorization, + resultHash: input.hash, + disposition: applied.disposition, + decision: applied.decision, + }); + }); + if (acknowledgement) notifications.push(...committedNotifications); + return acknowledgement; +} 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 abd19c3c63..95c15f6c31 100644 --- a/services/cloud-agent-next/src/shared/sandbox-control-protocol.ts +++ b/services/cloud-agent-next/src/shared/sandbox-control-protocol.ts @@ -29,6 +29,11 @@ export const SANDBOX_CONTROL_REQUEST_TIMEOUT_MS = 30_000; export const SANDBOX_CONTROL_ATTACH_TIMEOUT_MS = 8 * 60_000; export const SANDBOX_CONTROL_EXECUTION_TIMEOUT_MS = 60 * 60_000; +export const SANDBOX_CONTROL_CLEANUP_TIMEOUT_MS = 10_000; +export const SANDBOX_CONTROL_OPERATION_LIMIT = 32; +export const SANDBOX_CONTROL_OUTCOME_TIMEOUT_MS = 90_000; +export const SANDBOX_CONTROL_OUTCOME_RETRY_MS = 1_000; +export const SANDBOX_CONTROL_RECOVERY_MAX_ATTEMPTS = 3; export const SANDBOX_OPERATIONS = [ 'sandbox.hello', @@ -51,6 +56,8 @@ export const SESSION_OPERATIONS = [ 'session.terminal.resize', 'session.terminal.close', 'session.terminal.connect', + 'session.operation.get', + 'session.operation.ack', ] as const; export const SANDBOX_EVENTS = ['sandbox.ready', 'sandbox.heartbeat'] as const; @@ -115,6 +122,7 @@ export const requestFrameSchema = z.object({ operation: z.string().min(1), session: sessionRequestIdentitySchema.optional(), payload: z.unknown(), + authorization: z.lazy(() => sessionOperationAuthorizationSchema).optional(), }); export const responseFrameSchema = z.object({ @@ -151,12 +159,18 @@ 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(), }); export const sandboxHelloResultSchema = z.object({ protocolVersion: z.literal(SANDBOX_CONTROL_PROTOCOL_VERSION), handshakeComplete: z.literal(true), - capabilities: z.object({ kiloVersionHeartbeat: z.boolean().optional() }).optional(), + capabilities: z + .object({ + kiloVersionHeartbeat: z.boolean().optional(), + sessionOperationResults: z.boolean().optional(), + }) + .optional(), }); export type SandboxHelloPayload = z.infer; @@ -591,6 +605,150 @@ export type SessionTerminalConnectResult = z.infer; export type SessionPreparingPayload = z.infer; +export const sessionOperationAuthorizationSchema = z + .object({ + operation: z.enum(['session.attach', 'session.prompt']), + operationId: requestIdSchema, + messageId: requestIdSchema, + session: sessionRequestIdentitySchema, + wrapperInstanceId: wrapperInstanceIdSchema, + dispatchDeadlineAt: z.number().int().positive().max(Number.MAX_SAFE_INTEGER), + }) + .strict() + .refine(value => value.operation !== 'session.prompt' || value.operationId === value.messageId); + +export type SessionOperationAuthorization = z.infer; + +export function sameSessionOperation( + left: SessionOperationAuthorization, + right: SessionOperationAuthorization +): boolean { + return ( + left.operation === right.operation && + left.operationId === right.operationId && + left.messageId === right.messageId && + left.session.sessionId === right.session.sessionId && + left.session.kiloSessionId === right.session.kiloSessionId && + left.session.directory === right.session.directory && + left.wrapperInstanceId === right.wrapperInstanceId && + left.dispatchDeadlineAt === right.dispatchDeadlineAt + ); +} + +export function sessionOperationExpiresAt(authorization: SessionOperationAuthorization): number { + return ( + authorization.dispatchDeadlineAt + + (authorization.operation === 'session.attach' + ? SANDBOX_CONTROL_ATTACH_TIMEOUT_MS + : SANDBOX_CONTROL_EXECUTION_TIMEOUT_MS) + + SANDBOX_CONTROL_OUTCOME_TIMEOUT_MS + ); +} + +export const sessionOperationResultSchema = z.discriminatedUnion('ok', [ + z.object({ ok: z.literal(true), result: z.unknown() }).strict(), + z.object({ ok: z.literal(false), error: controlErrorSchema }).strict(), +]); + +export const sessionOperationDeliverySchema = z + .object({ + version: z.literal(2), + authorization: sessionOperationAuthorizationSchema, + completedAt: z.number().int().nonnegative().max(Number.MAX_SAFE_INTEGER), + result: sessionOperationResultSchema, + outcome: sessionMessageOutcomeSchema.optional(), + assistantMessageId: requestIdSchema.optional(), + events: z.array(sessionEventPayloadSchema).max(8), + preparing: z.array(sessionPreparingPayloadSchema).max(64), + }) + .strict() + .refine(value => + value.authorization.operation === 'session.prompt' + ? value.outcome?.messageId === value.authorization.messageId && value.preparing.length === 0 + : value.outcome === undefined && value.events.length === 0 + ) + .refine(value => value.completedAt < sessionOperationExpiresAt(value.authorization)) + .refine(value => + value.events.every( + event => + (event.type === 'autocommit_completed' || event.type === 'status') && + typeof event.properties.messageId === 'string' && + (event.properties.messageId === value.authorization.messageId || + event.properties.messageId === value.assistantMessageId) + ) + ) + .refine(value => + value.preparing.every( + event => + event.attemptId === value.authorization.operationId && + event.triggerMessageId === value.authorization.messageId + ) + ) + .refine(value => { + try { + return ( + new TextEncoder().encode(JSON.stringify(value)).byteLength <= + MAX_SANDBOX_CONTROL_FRAME_BYTES - 4096 + ); + } catch { + return false; + } + }); + +export type SessionOperationDelivery = z.infer; + +export const sessionOperationLookupResultSchema = z.discriminatedUnion('state', [ + z.object({ state: z.literal('missing') }).strict(), + z + .object({ + state: z.literal('running'), + authorization: sessionOperationAuthorizationSchema, + }) + .strict(), + z + .object({ + state: z.literal('completed'), + delivery: sessionOperationDeliverySchema, + }) + .strict(), +]); + +export type SessionOperationLookupResult = z.infer; + +export const sessionOperationAckSchema = z + .object({ + version: z.literal(2), + authorization: sessionOperationAuthorizationSchema, + resultHash: z.string().regex(/^[a-f0-9]{64}$/), + disposition: z.enum(['applied', 'identical', 'already_final', 'superseded']), + decision: z + .object({ + state: z.enum(['queued', 'accepted', 'completed', 'failed', 'cancelled']), + at: z.number().int().nonnegative().max(Number.MAX_SAFE_INTEGER), + }) + .strict(), + }) + .strict() + .refine( + value => + (value.authorization.operation !== 'session.prompt' && + value.disposition !== 'already_final' && + value.disposition !== 'superseded') || + ['completed', 'failed', 'cancelled'].includes(value.decision.state) + ); + +export type SessionOperationAck = z.infer; + +export async function sessionOperationResultHash( + delivery: SessionOperationDelivery +): Promise { + const bytes = new TextEncoder().encode( + JSON.stringify(sessionOperationDeliverySchema.parse(delivery)) + ); + const hash = await crypto.subtle.digest('SHA-256', bytes); + return Array.from(new Uint8Array(hash), byte => byte.toString(16).padStart(2, '0')).join(''); +} + const observationTimestampSchema = z.number().int().nonnegative().max(8_640_000_000_000_000); export const sandboxControlObservationSchema = z @@ -615,6 +773,7 @@ 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(), 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 3b0236e13a..eaeddd5179 100644 --- a/services/cloud-agent-next/test/integration/sandbox-control.test.ts +++ b/services/cloud-agent-next/test/integration/sandbox-control.test.ts @@ -106,6 +106,9 @@ import { createEventQueries } from '../../src/session/queries/index.js'; import { throwAdmissionError } from '../../src/session/queue-message.js'; import { requestFrameSchema, + responseFrameSchema, + sessionOperationAckSchema, + sessionOperationAuthorizationSchema, sessionPromptPayloadSchema, SANDBOX_CONTROL_AUTO_PING, SANDBOX_CONTROL_AUTO_PONG, @@ -113,6 +116,7 @@ import { type RequestFrame, type ResponseFrame, type SessionAttachPayload, + type SessionOperationDelivery, sandboxControlSocketAttachmentSchema, type SandboxHeartbeatPayload, } from '../../src/shared/sandbox-control-protocol.js'; @@ -345,7 +349,11 @@ function persistedSessionEvents(state: DurableObjectState, eventTypes: string[]) function sendHello( ws: WebSocket, requestId: string, - identity: { providerInstanceId?: string; wrapperInstanceId?: string } = {} + identity: { + providerInstanceId?: string; + wrapperInstanceId?: string; + sessionOperationResults?: boolean; + } = {} ): void { ws.send( JSON.stringify({ @@ -356,6 +364,9 @@ function sendHello( protocolVersion: 1, providerInstanceId: identity.providerInstanceId ?? cloudflareRef(socketSandboxIds.get(ws) ?? sandboxId), + ...(identity.sessionOperationResults + ? { capabilities: { sessionOperationResults: true } } + : {}), ...(identity.wrapperInstanceId ? { wrapperInstanceId: identity.wrapperInstanceId } : {}), }, }) @@ -365,7 +376,11 @@ function sendHello( async function completeHello( ws: WebSocket, requestId: string, - identity: { providerInstanceId?: string; wrapperInstanceId?: string } = {} + identity: { + providerInstanceId?: string; + wrapperInstanceId?: string; + sessionOperationResults?: boolean; + } = {} ): Promise { sendHello(ws, requestId, identity); await expect(nextMessage(ws)).resolves.toBe( @@ -376,7 +391,7 @@ async function completeHello( result: { protocolVersion: 1, handshakeComplete: true, - capabilities: { kiloVersionHeartbeat: true }, + capabilities: { kiloVersionHeartbeat: true, sessionOperationResults: true }, }, }) ); @@ -6427,7 +6442,7 @@ describe('SandboxControl passive status', () => { }, }) ); - await Promise.all(fresh['sessionForwardChains'].values()); + await Promise.all(fresh['sessionForwarding'].values()); expect(await fresh.getSandboxStatus(statusInput)).toMatchObject({ status: 'active', estimatedSleepAt: null, @@ -6692,7 +6707,7 @@ const savedWorktreeSnapshot: WorktreeChangesSnapshot = { capturedAt: '2026-08-20T10:00:00.000Z', }; -async function worktreeFixture() { +async function worktreeFixture(options: { sessionOperationResults?: boolean } = {}) { const suffix = crypto.randomUUID(); const userId = `user_worktree_${suffix}`; const sessionId = `workspace_${suffix}` as const; @@ -6779,18 +6794,32 @@ async function worktreeFixture() { await control.prepareSessionCredentials({ ownerId: userId, sessionId }); await control.attachSession({ sessionId, kiloSessionId, directory, worktreeId, ownerId: userId }); let ws = await connect(credential, sandboxId); - await completeHello(ws, `hello_${suffix}`, { wrapperInstanceId }); + await completeHello(ws, `hello_${suffix}`, { + wrapperInstanceId, + sessionOperationResults: options.sessionOperationResults, + }); const captures: RequestFrame[] = []; const inbox: RequestFrame[] = []; const captureWaiters: ((request: RequestFrame) => void)[] = []; const prompts: RequestFrame[] = []; const promptSeen = Promise.withResolvers(); const aborts: RequestFrame[] = []; + const resultWaiters = new Map void>(); let nextAttach: ((request: RequestFrame) => void) | undefined; function receive(client: WebSocket): void { client.addEventListener('message', event => { - const parsed = requestFrameSchema.safeParse(JSON.parse(String(event.data))); + const frame = JSON.parse(String(event.data)); + const response = responseFrameSchema.safeParse(frame); + if (response.success) { + const resolve = resultWaiters.get(response.data.requestId); + if (resolve) { + resultWaiters.delete(response.data.requestId); + resolve(response.data); + } + return; + } + const parsed = requestFrameSchema.safeParse(frame); if (!parsed.success) return; const request = parsed.data; if (request.operation === 'session.git.summary') { @@ -6882,6 +6911,20 @@ async function worktreeFixture() { reply(request: RequestFrame, result: unknown): void { ws.send(JSON.stringify({ type: 'response', requestId: request.requestId, ok: true, result })); }, + async sendOperationResult(delivery: SessionOperationDelivery): Promise { + const requestId = crypto.randomUUID(); + const response = new Promise(resolve => resultWaiters.set(requestId, resolve)); + ws.send( + JSON.stringify({ + type: 'request', + requestId, + operation: 'session.operation.result', + session: delivery.authorization.session, + payload: delivery, + }) + ); + return response; + }, fail( request: RequestFrame, retryable = false, @@ -6958,6 +7001,141 @@ function captureRevision(request: RequestFrame): number { return (request.payload as { revision: number }).revision; } +describe('SandboxSession operation authorization admission', () => { + it('persists dispatch proof before the capability-gated attach and prompt reach the wrapper socket', async () => { + const fetchMock = vi + .spyOn(globalThis, 'fetch') + .mockImplementation(async () => Response.json({ valid: true })); + const fixture = await worktreeFixture({ sessionOperationResults: true }); + const messageId = 'msg_operation_authorization'; + try { + const attach = fixture.holdNextAttach(); + await expect( + fixture.session.admitSubmittedMessage({ + userId: fixture.userId, + turn: { type: 'prompt', id: messageId, prompt: 'persist before egress' }, + }) + ).resolves.toMatchObject({ success: true, messageId }); + + const attachRequest = await attach; + expect(attachRequest).toMatchObject({ + operation: 'session.attach', + session: { + sessionId: fixture.sessionId, + kiloSessionId: fixture.kiloSessionId, + directory: fixture.directory, + }, + authorization: { + operation: 'session.attach', + messageId, + wrapperInstanceId: fixture.wrapperInstanceId, + }, + }); + await runInDurableObject(fixture.session, async (_instance, state) => { + const messages = state.storage.kv.get('session_messages') ?? []; + expect(messages).toMatchObject([ + { + messageId, + state: 'queued', + unresolvedDispatch: true, + operations: { + attach: { + dispatched: true, + authorization: { + operation: 'session.attach', + messageId, + wrapperInstanceId: fixture.wrapperInstanceId, + }, + }, + }, + }, + ]); + }); + + fixture.reply(attachRequest, { attached: true }); + await fixture.promptSeen; + expect(fixture.prompts).toHaveLength(1); + expect(fixture.prompts[0]).toMatchObject({ + operation: 'session.prompt', + authorization: { + operation: 'session.prompt', + operationId: messageId, + messageId, + wrapperInstanceId: fixture.wrapperInstanceId, + }, + }); + } finally { + fixture.close(); + fetchMock.mockRestore(); + } + }); + + it('returns the exact durable acknowledgement when the wrapper repeats a completed prompt result', async () => { + const fetchMock = vi + .spyOn(globalThis, 'fetch') + .mockImplementation(async () => Response.json({ valid: true })); + const fixture = await worktreeFixture({ sessionOperationResults: true }); + const messageId = 'msg_operation_result_ack'; + try { + await expect( + fixture.session.admitSubmittedMessage({ + userId: fixture.userId, + turn: { type: 'prompt', id: messageId, prompt: 'retain this completed result' }, + }) + ).resolves.toMatchObject({ success: true, messageId }); + await fixture.promptSeen; + const prompt = fixture.prompts[0]; + if (!prompt) throw new Error('Missing prompt operation request'); + const authorization = sessionOperationAuthorizationSchema.parse(prompt.authorization); + const delivery: SessionOperationDelivery = { + version: 2, + authorization, + completedAt: Date.now(), + result: { ok: true, result: { messageId, status: 'accepted' } }, + outcome: { messageId, status: 'completed' }, + events: [], + preparing: [], + }; + + const first = await fixture.sendOperationResult(delivery); + const second = await fixture.sendOperationResult(delivery); + const firstAck = sessionOperationAckSchema.parse(first.ok ? first.result : undefined); + expect(first).toMatchObject({ + ok: true, + result: { + authorization, + disposition: 'applied', + decision: { state: 'completed' }, + resultHash: expect.stringMatching(/^[a-f0-9]{64}$/), + }, + }); + expect(second).toMatchObject({ + ok: true, + result: { + authorization, + disposition: 'identical', + resultHash: firstAck.resultHash, + }, + }); + await runInDurableObject(fixture.session, async (_instance, state) => { + const messages = state.storage.kv.get('session_messages') ?? []; + expect(messages).toMatchObject([ + { + messageId, + state: 'completed', + terminalSource: 'operation_result', + operations: { prompt: { resultHash: firstAck.resultHash } }, + }, + ]); + }); + expect(fixture.prompts).toHaveLength(1); + } finally { + fixture.close(); + fetchMock.mockRestore(); + } + }); +}); + describe('SandboxSession worktree changes persistence', () => { beforeEach(() => { vi.spyOn(globalThis, 'fetch').mockImplementation(async () => Response.json({ valid: true })); @@ -9296,6 +9474,7 @@ describe('SandboxSession control-plane regressions', () => { preparationAttemptId: expect.any(String), deliveryDeadlineAt: expect.any(Number), terminalAt: expect.any(Number), + terminalSource: 'coordinator', }); expect(delivered.messages.slice(1)).toMatchObject([ { messageId: 'msg_model_less', state: 'accepted' }, diff --git a/services/cloud-agent-next/test/unit/wrapper/worktree-credential-refresh.test.ts b/services/cloud-agent-next/test/unit/wrapper/worktree-credential-refresh.test.ts index 8a74198d39..3f16d4c4c1 100644 --- a/services/cloud-agent-next/test/unit/wrapper/worktree-credential-refresh.test.ts +++ b/services/cloud-agent-next/test/unit/wrapper/worktree-credential-refresh.test.ts @@ -17,17 +17,22 @@ import { createSessionActivityRegistry, handleControlRequest, refreshHeartbeatPayload, + type ControlHandlerResult, type HandlerDeps, } from '../../../wrapper/src/control/sandbox-control-handlers.js'; import { startSandboxControlEventFeed } from '../../../wrapper/src/control/sandbox-control-runtime.js'; import type * as ControlRuntimeModule from '../../../wrapper/src/control/sandbox-control-runtime.js'; +import type * as KiloApiModule from '../../../wrapper/src/kilo-api.js'; import type * as UtilsModule from '../../../wrapper/src/utils.js'; vi.mock('node:fs/promises', () => ({ default: { mkdir: vi.fn(async () => undefined), writeFile: vi.fn(async () => undefined) }, })); vi.mock('@kilocode/sdk', () => ({ createKiloClient: vi.fn(() => ({})) })); -vi.mock('../../../wrapper/src/kilo-api.js', () => ({ createWrapperKiloClient: vi.fn() })); +vi.mock('../../../wrapper/src/kilo-api.js', async importOriginal => ({ + ...(await importOriginal()), + createWrapperKiloClient: vi.fn(), +})); vi.mock('../../../wrapper/src/control/sandbox-control-runtime.js', async importOriginal => ({ ...(await importOriginal()), startSandboxControlEventFeed: vi.fn(async () => ({ isFresh: () => true })), @@ -268,14 +273,14 @@ describe('direct worktree credential refresh', () => { const activity = createSessionActivityRegistry(); activity.attach(identity.kiloSessionId); activity.attach(sibling.kiloSessionId); - const deps: HandlerDeps = { + const emitSessionEvent = vi.fn(); + const deps: HandlerDeps = createControlHandlerDeps({ kiloRuntimes: f.registry, version: 'test', kiloReady: true, sessions: [], - tasks: new Map(), activity, - emitSessionEvent: vi.fn(), + emitSessionEvent, retireRuntime: vi.fn(), applyAttach: (session, payload, options) => applySessionAttach(session, payload, { @@ -283,7 +288,7 @@ describe('direct worktree credential refresh', () => { hasBootstrapMarker: async () => true, sessionExists: async () => true, }), - }; + }); const refreshing = handleControlRequest( 'session.attach', identity, @@ -306,8 +311,8 @@ describe('direct worktree credential refresh', () => { deps ) ).toMatchObject({ ok: false, error: { code: 'session_busy', retryable: true } }); - expect(deps.tasks.has(sibling.kiloSessionId)).toBe(false); - expect(deps.emitSessionEvent).not.toHaveBeenCalled(); + expect(deps.operations.hasActive(sibling.kiloSessionId)).toBe(false); + expect(emitSessionEvent).not.toHaveBeenCalled(); expect(runtime.signal.aborted).toBe(false); } finally { release.resolve({}); @@ -331,29 +336,66 @@ describe('direct worktree credential refresh', () => { if (condition === 'stale') fresh = false; else vi.spyOn(f.registry, 'get').mockReturnValue(undefined); if (condition === 'unattached') f.registry.detach(sibling); - const controller = new AbortController(); - if (condition === 'aborted') controller.abort(); - const deps: HandlerDeps = { + const emitSessionEvent = vi.fn(); + const retireRuntime = vi.fn(); + const deps: HandlerDeps = createControlHandlerDeps({ kiloRuntimes: f.registry, version: 'test', kiloReady: true, sessions: [], - tasks: new Map(), - emitSessionEvent: vi.fn(), - retireRuntime: vi.fn(), - }; + emitSessionEvent, + retireRuntime, + }); + const held = Promise.withResolvers(); + let operation: + | { cancel: (r: string, s: 'cancelled') => void; done: Promise } + | undefined; if (condition !== 'no-task') { - deps.tasks.set(identity.kiloSessionId, { - kind: condition === 'execution' ? 'execution' : 'preparation', - messageId: 'other-message', - session: { - ...identity, - directory: condition === 'other-directory' ? '/workspace/other' : identity.directory, - }, - controller, - signal: controller.signal, - done: Promise.resolve({ ok: true, result: {} }), - }); + const taskSession = { + ...identity, + directory: condition === 'other-directory' ? '/workspace/other' : identity.directory, + }; + if (condition === 'execution') { + const runtimeLifetime = new AbortController(); + const fakeRuntime = { + directory: taskSession.directory, + scopeId: 'fake', + env: {}, + kiloClient: { + sendPrompt: () => held.promise, + abortSession: async () => true, + getSessionDetails: async (id: string) => ({ id }), + }, + signal: runtimeLifetime.signal, + }; + operation = deps.operations.start( + taskSession, + undefined, + { + operation: 'session.prompt', + payload: { + messageId: 'other-message', + turn: { type: 'prompt', prompt: 'block' }, + agent: { mode: 'code', model: 'kilo/test' }, + }, + runtime: fakeRuntime as any, + }, + { emitSessionEvent: () => {} } + ); + } else { + operation = deps.operations.start( + taskSession, + undefined, + { + operation: 'session.attach', + payload: {} as any, + apply: () => held.promise, + onAttached: () => {}, + }, + { emitSessionEvent: () => {} } + ); + } + if (condition === 'aborted') operation.cancel('test-abort', 'cancelled'); } expect( await handleControlRequest( @@ -367,9 +409,12 @@ describe('direct worktree credential refresh', () => { deps ) ).toMatchObject({ ok: false, error: { code: 'not_ready', retryable: true } }); - expect(deps.tasks.has(sibling.kiloSessionId)).toBe(false); - expect(deps.emitSessionEvent).not.toHaveBeenCalled(); - expect(deps.retireRuntime).not.toHaveBeenCalled(); + expect(deps.operations.hasActive(sibling.kiloSessionId)).toBe(false); + expect(emitSessionEvent).not.toHaveBeenCalled(); + expect(retireRuntime).not.toHaveBeenCalled(); + // Clean up held operations + held.resolve({ ok: true, result: {} }); + if (operation) await operation.done.catch(() => {}); } ); @@ -578,7 +623,6 @@ describe('direct worktree credential refresh', () => { version: 'test', kiloReady: true, sessions: [], - tasks: new Map(), activity, emitSessionEvent: vi.fn(), retireRuntime: vi.fn(), @@ -734,12 +778,11 @@ describe('direct worktree credential refresh', () => { activity.attach(identity.kiloSessionId); activity.attach(sibling.kiloSessionId); activity.markActive(sibling.kiloSessionId); - const deps: HandlerDeps = { + const deps: HandlerDeps = createControlHandlerDeps({ kiloRuntimes: f.registry, version: 'test', kiloReady: true, sessions: [], - tasks: new Map(), activity, emitSessionEvent: vi.fn(), retireRuntime: vi.fn(), @@ -749,25 +792,49 @@ describe('direct worktree credential refresh', () => { hasBootstrapMarker: async () => true, sessionExists: async () => true, }), - }; - const controller = new AbortController(); - deps.tasks.set(sibling.kiloSessionId, { - kind: 'execution', - messageId: 'active-sibling-message', - session: sibling, - controller, - signal: controller.signal, - done: Promise.resolve({ ok: true, result: {} }), }); + const held = Promise.withResolvers(); + const runtimeLifetime = new AbortController(); + const fakeRuntime = { + directory: sibling.directory, + scopeId: 'fake', + env: {}, + kiloClient: { + sendPrompt: () => held.promise, + abortSession: async () => true, + getSessionDetails: async (id: string) => ({ id }), + }, + signal: runtimeLifetime.signal, + }; + const siblingOp = deps.operations.start( + sibling, + undefined, + { + operation: 'session.prompt', + payload: { + messageId: 'active-sibling-message', + turn: { type: 'prompt', prompt: 'block' }, + agent: { mode: 'code', model: 'kilo/test' }, + }, + runtime: fakeRuntime as any, + }, + { emitSessionEvent: () => {} } + ); const payload = { kilo: auth, env: { ...originalEnv, GH_TOKEN: 'github-renewed' } }; expect(await handleControlRequest('session.attach', identity, payload, deps)).toMatchObject({ ok: false, error: { code: 'session_busy', retryable: true }, }); - expect(controller.signal.aborted).toBe(false); + expect(siblingOp.signal.aborted).toBe(false); expect(runtime.signal.aborted).toBe(false); expect(f.close).not.toHaveBeenCalled(); - deps.tasks.delete(sibling.kiloSessionId); + // Remove the sibling operation (equivalent to old deps.tasks.delete) + siblingOp.cancel('test-cleanup', 'cancelled'); + held.resolve({ ok: true, result: {} }); + await siblingOp.done.catch(() => {}); + // Operation completion reconciles activity to idle; restore active state + // to match the original test which only removed the task without touching activity + activity.markActive(sibling.kiloSessionId); expect(await handleControlRequest('session.attach', identity, payload, deps)).toMatchObject({ ok: false, error: { code: 'session_busy', retryable: true }, @@ -780,7 +847,7 @@ describe('direct worktree credential refresh', () => { }); expect(f.close).toHaveBeenCalledTimes(1); expect(runtime.env.GH_TOKEN).toBe('github-renewed'); - expect(deps.retireRuntime).not.toHaveBeenCalled(); + runtimeLifetime.abort(); }); it('fails attach instead of accepting stale Git auth when origin refresh fails', async () => { diff --git a/services/cloud-agent-next/wrapper/src/control/apply-attach.ts b/services/cloud-agent-next/wrapper/src/control/apply-attach.ts index faf5180b5d..3e35bd97cf 100644 --- a/services/cloud-agent-next/wrapper/src/control/apply-attach.ts +++ b/services/cloud-agent-next/wrapper/src/control/apply-attach.ts @@ -48,7 +48,10 @@ const SETUP_COMMAND_INACTIVITY_TIMEOUT_MS = 4 * 60_000; const SETUP_COMMAND_HARD_TIMEOUT_MS = 300_000; const workspacePreparations = new Map>(); -export type AttachPreparingEmitter = (event: PreparingEventDataV2) => void; +export type AttachPreparingEmitter = ( + event: PreparingEventDataV2, + options?: { retained?: true } +) => void; export type ApplyAttachDeps = { onDiagnostic?: ControlDiagnosticReporter; 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 new file mode 100644 index 0000000000..dd112e3e5e --- /dev/null +++ b/services/cloud-agent-next/wrapper/src/control/control-handler-result.ts @@ -0,0 +1,13 @@ +import type { ControlError } from '../../../src/shared/sandbox-control-protocol.js'; + +export type ControlHandlerResult = + | { ok: true; result: unknown; admission?: never } + | { ok: false; error: ControlError; admission?: 'not-admitted' }; + +export function rejectBeforeAdmission( + code: ControlError['code'], + message: string, + retryable: boolean +): ControlHandlerResult { + return { ok: false, error: { code, message, retryable }, admission: 'not-admitted' }; +} diff --git a/services/cloud-agent-next/wrapper/src/control/control-test-fixtures.ts b/services/cloud-agent-next/wrapper/src/control/control-test-fixtures.ts new file mode 100644 index 0000000000..323201ea2a --- /dev/null +++ b/services/cloud-agent-next/wrapper/src/control/control-test-fixtures.ts @@ -0,0 +1,163 @@ +import fs from 'node:fs'; +import path from 'node:path'; +import { + sessionOperationResultHash, + type SessionOperationAuthorization, + type SessionOperationDelivery, + type SessionOperationAck, +} from '../../../src/shared/sandbox-control-protocol'; +import type { WrapperKiloClient } from '../kilo-api'; +import { applySessionAttach } from './apply-attach'; +import { createControlHandlerDeps, type HandlerDeps } from './sandbox-control-handlers'; +import { + buildWorktreeKiloEnvironment, + type WorktreeKiloAuth, + type WorktreeKiloRuntime, +} from './worktree-runtime'; + +export const session = { + sessionId: 'ses_1', + kiloSessionId: 'kilo_1', + directory: '/workspace', +}; + +export type Completion = Awaited>; + +export function completion(error?: Completion['info']['error']): Completion { + return { + info: { + id: 'assistant_1', + sessionID: session.kiloSessionId, + parentID: 'msg_1', + role: 'assistant', + time: { created: 1, completed: 2 }, + modelID: 'kilo/example', + providerID: 'kilo', + mode: 'code', + agent: 'code', + path: { cwd: session.directory, root: session.directory }, + cost: 0, + tokens: { input: 1, output: 1, reasoning: 0, cache: { read: 0, write: 0 } }, + ...(error ? { error } : {}), + }, + parts: [], + }; +} + +export function fakeKilo(overrides: Partial = {}): WrapperKiloClient { + return { + getSession: async (id: string) => ({ id }), + ensureSession: async () => undefined, + sendPrompt: async () => completion(), + sendPromptAsync: async () => {}, + sendCommand: async () => completion(), + summarizeSession: async () => true, + generateCommitMessage: async () => ({ message: 'Apply normal control turn' }), + abortSession: async () => true, + answerPermission: async () => true, + answerQuestion: async () => true, + rejectQuestion: async () => true, + getSessionDetails: async (id: string, directory = session.directory) => ({ id, directory }), + getSessionStatuses: async () => ({}), + getQuestions: async () => [], + getPermissions: async () => [], + ...overrides, + } as WrapperKiloClient; +} + +export const kilo: WorktreeKiloAuth = { + scopeId: 'worktree_1', + token: 'guest-kilo-token', + targets: { + backendBaseUrl: 'https://backend.example.test', + providerBaseUrl: 'https://provider.example.test', + sessionIngestBaseUrl: 'https://ingest.example.test', + }, +}; + +export const promptPayload = { + messageId: 'msg_1', + turn: { type: 'prompt', prompt: 'hello' }, + agent: { mode: 'architect', model: 'kilo/example', variant: 'high' }, +} as const; + +export function createHandlerFixture( + homeRoot: string, + overrides: Partial> & { kiloClient?: WrapperKiloClient } = {}, + identity = session +): HandlerDeps { + const { kiloClient, ...rest } = overrides; + const client = Object.hasOwn(overrides, 'kiloClient') ? kiloClient : fakeKilo(); + const nativeLifetime = new AbortController(); + const runtime: WorktreeKiloRuntime | undefined = client + ? { + scopeId: kilo.scopeId, + directory: identity.directory, + env: buildWorktreeKiloEnvironment( + identity.directory, + fs.mkdtempSync(path.join(homeRoot, 'worktree-')), + kilo, + {}, + {} + ), + kiloClient: client, + signal: nativeLifetime.signal, + } + : undefined; + return createControlHandlerDeps({ + kiloRuntimes: runtime + ? { + attach: () => { + nativeLifetime.signal.throwIfAborted(); + return { + ready: Promise.resolve(runtime), + signal: runtime.signal, + commit: () => {}, + release: () => {}, + }; + }, + detach: () => true, + deleteDirectory: async () => {}, + get: directory => + directory === runtime.directory && !nativeLifetime.signal.aborted ? runtime : undefined, + isHealthy: () => true, + shutdown: () => {}, + } + : undefined, + version: '2.4.0', + kiloReady: true, + sessions: [], + emitSessionEvent: () => {}, + retireRuntime: () => {}, + applyAttach: (session, payload, deps) => + applySessionAttach(session, payload, { ...deps, sessionExists: async () => true }), + ...rest, + }); +} + +export function operationAuthorization( + operation: SessionOperationAuthorization['operation'] = 'session.prompt', + messageId = 'msg_1', + identity = session +): SessionOperationAuthorization { + return { + operation, + operationId: operation === 'session.prompt' ? messageId : `prepare_${messageId}`, + messageId, + session: { ...identity }, + wrapperInstanceId: crypto.randomUUID(), + dispatchDeadlineAt: Date.now() + 60_000, + }; +} + +export async function acknowledgeOperation( + delivery: SessionOperationDelivery +): Promise { + return { + version: 2, + authorization: delivery.authorization, + resultHash: await sessionOperationResultHash(delivery), + disposition: 'applied', + decision: { state: delivery.outcome?.status ?? 'queued', at: delivery.completedAt }, + }; +} diff --git a/services/cloud-agent-next/wrapper/src/control/main.ts b/services/cloud-agent-next/wrapper/src/control/main.ts index c7911836bf..470f3a754c 100644 --- a/services/cloud-agent-next/wrapper/src/control/main.ts +++ b/services/cloud-agent-next/wrapper/src/control/main.ts @@ -101,20 +101,28 @@ function main(diagnostics: ControlDiagnostics, wrapperInstanceId: string): void return !shuttingDown && kiloRuntimes.isHealthy(); }, sessions: [], - tasks: new Map(), activity: createSessionActivityRegistry(), signal: abort.signal, ...(terminalRuntime ? { terminalRuntime } : {}), - emitSessionEvent: (session, payload) => { + sendOperationResult: (session, delivery, signal, deadlineAt) => { + if (!control?.sendOperationResult) + throw new Error('Sandbox control operation result delivery unavailable'); + return control.sendOperationResult(session, delivery, signal, deadlineAt); + }, + emitSessionEvent: (session, payload, options) => { if ( - !control?.sendEvent?.('session.event', payload, { - directory: session.directory, - kiloSessionId: session.kiloSessionId, - rootKiloSessionId: session.kiloSessionId, - }) - ) { + !control?.sendEvent?.( + 'session.event', + payload, + { + directory: session.directory, + kiloSessionId: session.kiloSessionId, + rootKiloSessionId: session.kiloSessionId, + }, + options?.retained ? { preserveConnectionOnFailure: true } : undefined + ) + ) throw new Error('Sandbox control event delivery failed'); - } }, retireRuntime: reason => shutdown(1, reason), onShutdown: () => shutdown(0, 'Sandbox shutting down'), @@ -176,6 +184,7 @@ function main(diagnostics: ControlDiagnostics, wrapperInstanceId: string): void terminalRuntime?.shutdown(); } finally { await tasks; + await deps.operations.drainDelivery(shutdownAt + KILO_CONTROL_REQUEST_TIMEOUT_MS); } })(); void stopped @@ -200,22 +209,32 @@ function main(diagnostics: ControlDiagnostics, wrapperInstanceId: string): void isReady: () => deps.kiloReady, onConnected: () => diagnostics.onDiagnostic('wrapper.lifecycle', { phase: 'ready', ok: true }), onDisconnected: () => shutdown(1, 'Sandbox control connection lost', 'control_disconnected'), - onRequest: (operation, session, payload) => - handleControlRequest(operation, session, payload, { - ...deps, - emitPreparing: event => { - if (!session) return; - if ( - !control?.sendEvent?.('session.preparing', event, { - directory: session.directory, - kiloSessionId: session.kiloSessionId, - rootKiloSessionId: session.kiloSessionId, - }) - ) { - shutdown(1, 'Preparation event delivery failed', 'control_disconnected'); - } + onRequest: (operation, session, payload, authorization) => + handleControlRequest( + operation, + session, + payload, + { + ...deps, + emitPreparing: (event, options) => { + if (!session) return; + if ( + !control?.sendEvent?.( + 'session.preparing', + event, + { + directory: session.directory, + kiloSessionId: session.kiloSessionId, + rootKiloSessionId: session.kiloSessionId, + }, + options?.retained ? { preserveConnectionOnFailure: true } : undefined + ) + ) + throw new Error('Preparation event delivery failed'); + }, }, - }), + authorization + ), getHeartbeatPayload: () => withHeartbeatReason(buildHeartbeatPayload(deps)), sampleHeartbeat: signal => refreshHeartbeatPayload(deps, signal).then(() => undefined), }); diff --git a/services/cloud-agent-next/wrapper/src/control/operation-intent.ts b/services/cloud-agent-next/wrapper/src/control/operation-intent.ts new file mode 100644 index 0000000000..afe42111fd --- /dev/null +++ b/services/cloud-agent-next/wrapper/src/control/operation-intent.ts @@ -0,0 +1,47 @@ +import { + sessionAttachPayloadSchema, + sessionPromptPayloadSchema, +} from '../../../src/shared/sandbox-control-protocol.js'; + +export function operationIntent(operation: 'session.attach' | 'session.prompt', payload: unknown) { + if (operation === 'session.prompt') { + const { attachments, ...intent } = sessionPromptPayloadSchema.parse(payload); + return { + ...intent, + attachments: attachments?.map(({ filename, mime, localPath }) => ({ + filename, + mime, + localPath, + })), + }; + } + const { kilo, git, env, snapshotIdentity, directory, branch, setupCommands, preparation } = + sessionAttachPayloadSchema.parse(payload); + const credentials = new Set([ + 'KILOCODE_TOKEN', + 'KILOCODE_ORGANIZATION_ID', + 'GH_TOKEN', + 'GITHUB_TOKEN', + 'GITLAB_TOKEN', + 'GITLAB_OAUTH_TOKEN', + 'BITBUCKET_TOKEN', + 'BITBUCKET_APP_PASSWORD', + ]); + return { + snapshotIdentity, + directory, + branch, + setupCommands, + preparation, + kilo: kilo + ? { + scopeId: kilo.scopeId, + organizationId: kilo.organizationId, + containmentEnabled: kilo.containmentEnabled !== false, + targets: kilo.targets, + } + : undefined, + git: git ? { url: git.url, platform: git.platform } : undefined, + env: Object.fromEntries(Object.entries(env ?? {}).filter(([key]) => !credentials.has(key))), + }; +} 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 new file mode 100644 index 0000000000..6b87ae7402 --- /dev/null +++ b/services/cloud-agent-next/wrapper/src/control/operation-registry.test.ts @@ -0,0 +1,185 @@ +import { afterEach, beforeEach, describe, expect, it, setSystemTime, spyOn } from 'bun:test'; +import fs from 'node:fs'; +import os from 'node:os'; +import path from 'node:path'; +import { + SANDBOX_CONTROL_OPERATION_LIMIT, + SANDBOX_CONTROL_OUTCOME_TIMEOUT_MS, + sessionOperationLookupResultSchema, +} from '../../../src/shared/sandbox-control-protocol'; +import { + handleControlRequest, + pruneControlOperations, + type HandlerDeps, +} from './sandbox-control-handlers'; +import { + acknowledgeOperation, + completion, + createHandlerFixture, + fakeKilo, + operationAuthorization, + promptPayload, + session, + type Completion, +} from './control-test-fixtures'; +import { rememberAttachedRoot, resetSessionDirectoryState } from './session-directories'; +import { resetDirectoryOperationState } from './worktree-operations'; + +let homeRoot: string; + +beforeEach(() => { + resetSessionDirectoryState(); + resetDirectoryOperationState(); + rememberAttachedRoot(session.kiloSessionId, session.directory); + homeRoot = fs.mkdtempSync(path.join(os.tmpdir(), 'operation-registry-test-')); +}); + +afterEach(() => { + setSystemTime(); + fs.rmSync(homeRoot, { recursive: true, force: true }); +}); + +function deps(overrides: Parameters[1] = {}): HandlerDeps { + return createHandlerFixture(homeRoot, overrides); +} + +function onlyOperation(handlerDeps: HandlerDeps) { + const records = handlerDeps.operations.retained(); + expect(records).toHaveLength(1); + const record = records[0]; + if (!record) throw new Error('Missing operation record'); + return record; +} + +describe('operation admission and lookup', () => { + it('looks up the same operation during and after completion and rejects changed intent', async () => { + const running = Promise.withResolvers(); + const started = Promise.withResolvers(); + let submissions = 0; + const handlerDeps = deps({ + kiloClient: fakeKilo({ + sendPrompt: async () => { + submissions++; + started.resolve(); + return running.promise; + }, + }), + sendOperationResult: (_session, delivery) => acknowledgeOperation(delivery), + }); + rememberAttachedRoot(session.kiloSessionId, session.directory); + const authorization = operationAuthorization(); + await handleControlRequest( + 'session.prompt', + session, + promptPayload, + handlerDeps, + authorization + ); + await started.promise; + const record = onlyOperation(handlerDeps); + expect(handlerDeps.operations.active(session.kiloSessionId)).toBe(record); + const lookup = await handleControlRequest( + 'session.operation.get', + session, + authorization, + handlerDeps + ); + expect(lookup).toMatchObject({ ok: true, result: { state: 'running', authorization } }); + if (!lookup.ok) throw new Error('Missing running operation receipt'); + expect(sessionOperationLookupResultSchema.parse(lookup.result)).toEqual({ + state: 'running', + authorization, + }); + expect( + await handleControlRequest( + 'session.prompt', + session, + promptPayload, + handlerDeps, + authorization + ) + ).toMatchObject({ ok: true, result: { status: 'existing' } }); + expect( + await handleControlRequest( + 'session.prompt', + session, + { ...promptPayload, agent: { ...promptPayload.agent, variant: 'low' } }, + handlerDeps, + authorization + ) + ).toMatchObject({ ok: false, error: { code: 'idempotency_conflict' } }); + const prepareHistory = spyOn(handlerDeps.kiloRuntimes!, 'isHealthy').mockReturnValue(true); + try { + expect( + await handleControlRequest('session.operation.get', session, authorization, handlerDeps) + ).toMatchObject({ ok: true, result: { state: 'running', authorization } }); + expect(submissions).toBe(1); + } finally { + prepareHistory.mockRestore(); + running.resolve(completion()); + } + await record.done; + await record.waitForDelivery(); + expect(record.snapshot().native.completion).toEqual(completion().info); + expect(record.snapshot().local?.result.ok).toBe(true); + expect(record.snapshot().delivery?.state).toBe('acknowledged'); + expect(handlerDeps.operations.counts().active).toBe(0); + expect( + await handleControlRequest('session.operation.get', session, authorization, handlerDeps) + ).toMatchObject({ + ok: true, + result: { state: 'completed', delivery: { outcome: { status: 'completed' } } }, + }); + await handleControlRequest( + 'session.prompt', + session, + promptPayload, + handlerDeps, + authorization + ); + expect(submissions).toBe(1); + }); + + it('preserves lookup at capacity and rejects expired replay after safe pruning', async () => { + const handlerDeps = deps({ + sendOperationResult: (_session, delivery) => acknowledgeOperation(delivery), + }); + rememberAttachedRoot(session.kiloSessionId, session.directory); + const first = operationAuthorization(); + for (let index = 0; index < SANDBOX_CONTROL_OPERATION_LIMIT; index++) { + const messageId = `message_${index}`; + const authorization = { ...first, messageId, operationId: messageId }; + await handleControlRequest( + 'session.prompt', + session, + { ...promptPayload, messageId }, + handlerDeps, + authorization + ); + const record = [...handlerDeps.operations.retained()].at(-1); + if (!record) throw new Error('Missing admitted record'); + await record.done; + await record.waitForDelivery(); + } + const original = { ...first, messageId: 'message_0', operationId: 'message_0' }; + expect( + await handleControlRequest('session.operation.get', session, original, handlerDeps) + ).toMatchObject({ ok: true, result: { state: 'completed' } }); + expect( + await handleControlRequest('session.prompt', session, promptPayload, handlerDeps, first) + ).toMatchObject({ ok: false, error: { code: 'session_busy' } }); + setSystemTime(first.dispatchDeadlineAt + SANDBOX_CONTROL_OUTCOME_TIMEOUT_MS + 1); + pruneControlOperations(handlerDeps); + expect(handlerDeps.operations.counts().retained).toBe(0); + expect( + await handleControlRequest( + 'session.prompt', + session, + { ...promptPayload, messageId: original.messageId }, + handlerDeps, + original + ) + ).toMatchObject({ ok: false, error: { code: 'not_ready', retryable: false } }); + expect(handlerDeps.operations.counts().active).toBe(0); + }); +}); diff --git a/services/cloud-agent-next/wrapper/src/control/operation-registry.ts b/services/cloud-agent-next/wrapper/src/control/operation-registry.ts new file mode 100644 index 0000000000..18dd62a614 --- /dev/null +++ b/services/cloud-agent-next/wrapper/src/control/operation-registry.ts @@ -0,0 +1,204 @@ +import { isDeepStrictEqual } from 'node:util'; +import { + SANDBOX_CONTROL_OPERATION_LIMIT, + sessionOperationAuthorizationSchema, + sessionOperationAckSchema, + sessionOperationExpiresAt, + type SessionOperationAuthorization, + type SessionRequestIdentity, +} from '../../../src/shared/sandbox-control-protocol.js'; +import { rejectBeforeAdmission } from './control-handler-result.js'; +import type { WorktreeKiloRuntimes } from './worktree-runtime.js'; +import { + SessionOperation, + type ControlHandlerResult, + type SessionOperationDependencies, + type SessionOperationWork, +} from './session-operation.js'; + +type OperationRegistryDependencies = { + native: Pick; + onStarted: (session: SessionRequestIdentity, preparation: boolean) => void; + onCompleted: (session: SessionRequestIdentity) => void; + retireRuntime: (reason: string) => void; +}; + +type OperationEffects = Pick< + SessionOperationDependencies, + 'signal' | 'emitSessionEvent' | 'sendOperationResult' | 'onDiagnostic' +>; + +type Admission = + | { kind: 'continue' } + | { kind: 'reply'; result: ControlHandlerResult | Promise }; + +function key(authorization: SessionOperationAuthorization): string { + return JSON.stringify([ + authorization.session.sessionId, + authorization.operation, + authorization.operationId, + ]); +} + +function ok(result: unknown): ControlHandlerResult { + return { ok: true, result }; +} + +function fail(code: string, message: string, retryable: boolean): ControlHandlerResult { + return { ok: false, error: { code, message, retryable } }; +} + +export function createOperationRegistry(deps: OperationRegistryDependencies) { + const active = new Map(); + const retained = new Map(); + + function prune(now = Date.now()): void { + for (const [id, operation] of retained) { + if (!operation.canPrune(now)) continue; + retained.delete(id); + } + } + + function admission( + operation: string, + session: SessionRequestIdentity, + payload: unknown, + authorization?: SessionOperationAuthorization + ): Admission { + if (operation !== 'session.operation.get' && !authorization) return { kind: 'continue' }; + const reply = (result: ControlHandlerResult | Promise): Admission => ({ + kind: 'reply', + result, + }); + const parsed = sessionOperationAuthorizationSchema.safeParse( + operation === 'session.operation.get' ? payload : authorization + ); + if (!parsed.success) + return reply( + rejectBeforeAdmission('protocol_error', 'Invalid operation authorization', false) + ); + const target = parsed.data; + if (!isDeepStrictEqual(target.session, session)) + return reply(rejectBeforeAdmission('unauthorized', 'Operation target mismatch', false)); + if (operation !== 'session.operation.get' && target.operation !== operation) + return reply( + rejectBeforeAdmission('unauthorized', 'Operation authorization mismatch', false) + ); + if (Date.now() >= sessionOperationExpiresAt(target)) + return reply(rejectBeforeAdmission('not_ready', 'Operation authorization expired', false)); + prune(); + const existing = retained.get(key(target)); + if (existing) { + if (!existing.matchesAuthorization(target) || !isDeepStrictEqual(existing.session, session)) + return reply( + rejectBeforeAdmission('idempotency_conflict', 'Operation identity mismatch', false) + ); + if (operation === 'session.operation.get') { + const delivery = existing.deliveryResult(); + return reply( + ok( + delivery + ? { state: 'completed', delivery } + : { + state: 'running', + authorization: target, + } + ) + ); + } + try { + if (!existing.matchesIntent(payload)) + return reply( + rejectBeforeAdmission('idempotency_conflict', 'Operation intent mismatch', false) + ); + } catch { + return reply(rejectBeforeAdmission('protocol_error', 'Invalid payload', false)); + } + return reply( + operation === 'session.attach' + ? existing.done + : ok({ messageId: target.messageId, status: 'existing' }) + ); + } + if (operation === 'session.operation.get') return reply(ok({ state: 'missing' })); + if (Date.now() >= target.dispatchDeadlineAt) + return reply( + rejectBeforeAdmission('not_ready', 'Operation dispatch authorization expired', false) + ); + if (retained.size >= SANDBOX_CONTROL_OPERATION_LIMIT) + return reply( + rejectBeforeAdmission('session_busy', 'Operation receipt capacity is unavailable', true) + ); + return { kind: 'continue' }; + } + + async function acknowledge( + session: SessionRequestIdentity, + payload: unknown + ): Promise { + const parsed = sessionOperationAckSchema.safeParse(payload); + if (!parsed.success || !isDeepStrictEqual(parsed.data.authorization.session, session)) + return fail('unauthorized', 'Invalid operation acknowledgement', false); + const id = key(parsed.data.authorization); + const operation = retained.get(id); + if (!operation) return fail('unauthorized', 'Operation acknowledgement is not current', false); + return (await operation.acknowledge(parsed.data, () => retained.get(id) === operation)) + ? ok({ acknowledged: true }) + : fail('unauthorized', 'Operation acknowledgement does not match the result', false); + } + + function start( + session: SessionRequestIdentity, + authorization: SessionOperationAuthorization | undefined, + work: SessionOperationWork, + effects: OperationEffects + ): SessionOperation { + const identity = Object.freeze({ ...session }); + const operation = new SessionOperation(identity, authorization, work, { + ...effects, + isCurrent: () => active.get(identity.kiloSessionId) === operation, + getRuntime: () => deps.native.get(identity.directory), + retireRuntime: reason => deps.retireRuntime(reason), + onLocalCompletion: retain => { + if (active.get(identity.kiloSessionId) === operation) { + active.delete(identity.kiloSessionId); + deps.onCompleted(identity); + } + if (authorization && !retain && retained.get(key(authorization)) === operation) + retained.delete(key(authorization)); + }, + }); + if (authorization) retained.set(key(authorization), operation); + active.set(identity.kiloSessionId, operation); + deps.onStarted(identity, work.operation === 'session.attach'); + return operation; + } + + return { + admission, + acknowledge, + start, + prune, + active: (rootKiloSessionId: string) => active.get(rootKiloSessionId), + hasActive: (rootKiloSessionId: string) => active.has(rootKiloSessionId), + activeOperations: () => [...active.values()], + retained: () => [...retained.values()], + counts: () => ({ active: active.size, retained: retained.size }), + async drainDelivery(deadlineAt: number): Promise { + const timeout = Math.max(0, deadlineAt - Date.now()); + let timer: ReturnType | undefined; + try { + await Promise.race([ + Promise.allSettled([...retained.values()].map(operation => operation.waitForDelivery())), + new Promise(resolve => { + timer = setTimeout(resolve, timeout); + }), + ]); + } finally { + if (timer) clearTimeout(timer); + } + }, + }; +} + +export type OperationRegistry = ReturnType; diff --git a/services/cloud-agent-next/wrapper/src/control/operation-result-delivery.test.ts b/services/cloud-agent-next/wrapper/src/control/operation-result-delivery.test.ts new file mode 100644 index 0000000000..ef3f09f1a9 --- /dev/null +++ b/services/cloud-agent-next/wrapper/src/control/operation-result-delivery.test.ts @@ -0,0 +1,156 @@ +import { describe, expect, it, mock, spyOn } from 'bun:test'; +import { + sessionOperationExpiresAt, + type SessionOperationAck, + type SessionOperationDelivery, +} from '../../../src/shared/sandbox-control-protocol'; +import { acknowledgeOperation, operationAuthorization } from './control-test-fixtures'; +import { createOperationResultDelivery } from './operation-result-delivery'; +import { ControlDeliveryError } from './sandbox-control-client'; + +function result(): SessionOperationDelivery { + return { + version: 2, + authorization: operationAuthorization(), + completedAt: Date.now(), + result: { ok: true, result: {} }, + outcome: { messageId: 'msg_1', status: 'completed' }, + events: [], + preparing: [], + }; +} + +describe('sealed operation result delivery', () => { + it('seals producer data and starts delivery once without exposing mutable result state', async () => { + const original = result(); + const expected = structuredClone(original); + const send = mock(async (payload: SessionOperationDelivery) => acknowledgeOperation(payload)); + const delivery = createOperationResultDelivery(original, Date.now() + 1_000, send); + original.outcome = { messageId: 'msg_1', status: 'failed', reason: 'caller mutation' }; + original.events.push({ + type: 'status', + properties: { messageId: 'msg_1', message: 'caller mutation' }, + }); + const first = delivery.start(); + expect(delivery.start()).toBe(first); + await delivery.drain(); + expect(send).toHaveBeenCalledTimes(1); + expect(send.mock.calls[0]?.[0]).toEqual(expected); + const snapshot = delivery.snapshot(); + snapshot.payload.outcome = { messageId: 'msg_1', status: 'cancelled' }; + expect(delivery.result()).toEqual(expected); + expect(delivery.status().state).toBe('acknowledged'); + }); + + it('retains an exact acknowledgement when cancelled transport waiting later rejects', async () => { + const held = Promise.withResolvers(); + const sending = Promise.withResolvers(); + const payload = result(); + const delivery = createOperationResultDelivery( + payload, + Date.now() + 1_000, + (_payload, signal) => { + sending.resolve(signal); + return held.promise; + } + ); + void delivery.start(); + const signal = await sending.promise; + const ack = await acknowledgeOperation(payload); + try { + expect(await delivery.acknowledge({ ...ack, resultHash: '0'.repeat(64) }, () => true)).toBe( + false + ); + expect(await delivery.acknowledge(ack, () => true)).toBe(true); + expect(signal.aborted).toBe(true); + held.reject(new Error('Transport completed late')); + await delivery.drain(); + expect(delivery.snapshot().acknowledgement).toEqual(ack); + expect(delivery.status().state).toBe('acknowledged'); + expect(delivery.result()).toEqual(payload); + } finally { + held.resolve(ack); + await delivery.drain(); + } + }); + + it('rechecks current identity and authorization expiry after hashing an explicit acknowledgement', async () => { + const payload = result(); + const delivery = createOperationResultDelivery(payload, Date.now() + 1_000); + const ack = await acknowledgeOperation(payload); + let current = true; + const superseded = delivery.acknowledge(ack, () => current); + current = false; + expect(await superseded).toBe(false); + const clock = spyOn(Date, 'now'); + try { + const expired = delivery.acknowledge(ack, () => true); + clock.mockReturnValue(sessionOperationExpiresAt(payload.authorization)); + expect(await expired).toBe(false); + expect(delivery.status().state).toBe('pending'); + expect(delivery.result()).toEqual(payload); + } finally { + clock.mockRestore(); + } + }); + + it('stops retrying after a permanent rejection but retains the unacknowledged result', async () => { + let attempts = 0; + const payload = result(); + const deadlineAt = Date.now() + 500; + const delivery = createOperationResultDelivery(payload, deadlineAt, async () => { + attempts++; + throw new ControlDeliveryError('Control delivery was not acknowledged', false); + }); + void delivery.start(); + await delivery.drain(); + expect(attempts).toBe(1); + expect(delivery.status().state).toBe('exhausted'); + expect(delivery.result()).toEqual(payload); + }); + + it('retries transient failures but stops at the attempt limit', async () => { + let attempts = 0; + const payload = result(); + const deadlineAt = Date.now() + 4_500; + const delivery = createOperationResultDelivery(payload, deadlineAt, async () => { + attempts++; + throw new ControlDeliveryError('Control transport unavailable', true); + }); + void delivery.start(); + await delivery.drain(); + expect(attempts).toBe(3); + expect(delivery.status().state).toBe('exhausted'); + expect(delivery.result()).toEqual(payload); + }); + + it('exhausts the original delivery deadline without replaying after a late valid acknowledgement', async () => { + const payload = result(); + const now = Date.now(); + const deadlineAt = now + 1_000; + const clock = spyOn(Date, 'now').mockReturnValue(now); + const held = Promise.withResolvers(); + const sending = Promise.withResolvers(); + const send = mock(() => { + sending.resolve(); + return held.promise; + }); + const delivery = createOperationResultDelivery(payload, deadlineAt, send); + try { + void delivery.start(); + await sending.promise; + const ack = await acknowledgeOperation(payload); + clock.mockReturnValue(deadlineAt); + held.resolve(ack); + await delivery.drain(); + await delivery.start(); + expect(delivery.status()).toEqual({ state: 'exhausted', deadlineAt }); + expect(send).toHaveBeenCalledTimes(1); + expect(delivery.result()).toEqual(payload); + } finally { + held.resolve(await acknowledgeOperation(payload)); + await delivery.drain(); + clock.mockRestore(); + } + }); +}); diff --git a/services/cloud-agent-next/wrapper/src/control/operation-result-delivery.ts b/services/cloud-agent-next/wrapper/src/control/operation-result-delivery.ts new file mode 100644 index 0000000000..4721ed366b --- /dev/null +++ b/services/cloud-agent-next/wrapper/src/control/operation-result-delivery.ts @@ -0,0 +1,133 @@ +import { setTimeout as delay } from 'node:timers/promises'; +import { + SANDBOX_CONTROL_OUTCOME_RETRY_MS, + SANDBOX_CONTROL_RECOVERY_MAX_ATTEMPTS, + sameSessionOperation, + sessionOperationAckSchema, + sessionOperationExpiresAt, + sessionOperationResultHash, + type SessionOperationAck, + type SessionOperationDelivery, +} from '../../../src/shared/sandbox-control-protocol.js'; +import { withTimeoutAndAbort } from '../utils.js'; + +export type OperationResultSender = ( + delivery: SessionOperationDelivery, + signal: AbortSignal, + deadlineAt: number +) => Promise; + +export function createOperationResultDelivery( + result: SessionOperationDelivery, + deadlineAt: number, + send?: OperationResultSender +) { + const payload = structuredClone(result); + let state: 'pending' | 'acknowledged' | 'exhausted' = 'pending'; + let acknowledgement: SessionOperationAck | undefined; + let controller: AbortController | undefined; + let pending: Promise | undefined; + + async function acknowledge(ack: SessionOperationAck, isCurrent: () => boolean): Promise { + if ( + !sameSessionOperation(payload.authorization, ack.authorization) || + Date.now() >= sessionOperationExpiresAt(payload.authorization) + ) + return false; + const hash = await sessionOperationResultHash(payload); + if ( + !isCurrent() || + Date.now() >= sessionOperationExpiresAt(payload.authorization) || + hash !== ack.resultHash + ) + return false; + acknowledgement = structuredClone(ack); + state = 'acknowledged'; + controller?.abort(); + return true; + } + + async function deliver(): Promise { + if (state !== 'pending' || !send) return; + const deliveryController = new AbortController(); + controller = deliveryController; + const signal = deliveryController.signal; + const timeout = setTimeout( + () => deliveryController.abort(), + Math.max(0, deadlineAt - Date.now()) + ); + timeout.unref(); + try { + const hash = await sessionOperationResultHash(payload); + for (let attempt = 0; attempt < SANDBOX_CONTROL_RECOVERY_MAX_ATTEMPTS; attempt++) { + if (signal.aborted || Date.now() >= deadlineAt) break; + try { + const ack = sessionOperationAckSchema.parse( + await withTimeoutAndAbort(send(structuredClone(payload), signal, deadlineAt), { + signal, + timeoutMs: Math.max(1, deadlineAt - Date.now()), + timeoutMessage: 'Operation result delivery expired', + abortMessage: 'Operation result delivery cancelled', + }) + ); + if ( + !signal.aborted && + Date.now() < deadlineAt && + sameSessionOperation(ack.authorization, payload.authorization) && + ack.resultHash === hash + ) { + acknowledgement = ack; + state = 'acknowledged'; + return; + } + } catch (error: unknown) { + if (signal.aborted) break; + if ( + error instanceof Error && + 'retryable' in error && + (error as { retryable: boolean }).retryable === false + ) + break; + } + if (attempt + 1 < SANDBOX_CONTROL_RECOVERY_MAX_ATTEMPTS) { + await delay( + Math.min( + SANDBOX_CONTROL_OUTCOME_RETRY_MS * 2 ** attempt, + Math.max(0, deadlineAt - Date.now()) + ), + undefined, + { signal } + ); + } + } + if (acknowledgement === undefined && Date.now() < deadlineAt) + await delay(deadlineAt - Date.now(), undefined, { signal }); + } catch { + if (acknowledgement === undefined) state = 'exhausted'; + } finally { + clearTimeout(timeout); + deliveryController.abort(); + if (controller === deliveryController) controller = undefined; + if (acknowledgement === undefined) state = 'exhausted'; + } + } + + return { + start(): Promise { + pending ??= Promise.resolve().then(deliver); + return pending; + }, + acknowledge, + status: () => ({ state, deadlineAt }), + result: () => structuredClone(payload), + snapshot: () => ({ + payload: structuredClone(payload), + deadlineAt, + state, + acknowledgement: acknowledgement ? structuredClone(acknowledgement) : undefined, + }), + drain: () => pending ?? Promise.resolve(), + }; +} + +export type OperationResultDelivery = ReturnType; diff --git a/services/cloud-agent-next/wrapper/src/control/retained-operation-notifications.test.ts b/services/cloud-agent-next/wrapper/src/control/retained-operation-notifications.test.ts new file mode 100644 index 0000000000..e5ece7e87b --- /dev/null +++ b/services/cloud-agent-next/wrapper/src/control/retained-operation-notifications.test.ts @@ -0,0 +1,165 @@ +import { describe, expect, it } from 'bun:test'; +import { + MAX_SANDBOX_CONTROL_FRAME_BYTES, + sessionOperationDeliverySchema, + type SessionOperationAuthorization, + type SessionPreparingPayload, +} from '../../../src/shared/sandbox-control-protocol.js'; +import type { PreparingEventDataV2 } from '../../../src/shared/protocol.js'; +import { createRetainedOperationNotifications } from './retained-operation-notifications.js'; + +function authorization(): SessionOperationAuthorization { + return { + operation: 'session.attach', + operationId: 'prepare_msg_1', + messageId: 'msg_1', + session: { + sessionId: 'ses_1', + kiloSessionId: 'kilo_1', + directory: '/workspace', + }, + wrapperInstanceId: '00000000-0000-4000-8000-000000000000', + dispatchDeadlineAt: Date.now() + 60_000, + }; +} + +function delivery(preparing: SessionPreparingPayload[]) { + return { + version: 2 as const, + authorization: authorization(), + completedAt: Date.now(), + result: { ok: true as const, result: { attached: true } }, + events: [], + preparing, + }; +} + +function retainedBytes( + snapshot: ReturnType['snapshot']> +) { + return [...snapshot.events, ...snapshot.preparing].reduce( + (total, payload) => total + Buffer.byteLength(JSON.stringify(payload)), + 0 + ); +} + +function optionalStep(revision: number, message: string, metadata: string): PreparingEventDataV2 { + return { + version: 2, + attemptId: 'prepare_msg_1', + triggerMessageId: 'msg_1', + revision, + timestamp: revision, + step: 'workspace_setup', + message, + action: 'step_started', + stepId: `step_${revision}`, + kind: 'phase', + label: metadata, + command: metadata, + }; +} + +function completedStep(revision: number): PreparingEventDataV2 { + return { + version: 2, + attemptId: 'prepare_msg_1', + triggerMessageId: 'msg_1', + revision, + timestamp: revision, + step: 'workspace_setup', + message: `Completed ${revision}`, + action: 'step_completed', + stepId: `step_${revision}`, + }; +} + +describe('retained operation notifications', () => { + it('evicts multiple optional entries for a larger valid attempt terminal', () => { + const recorder = createRetainedOperationNotifications(); + const musicalSymbolGClef = String.fromCodePoint(0x1d11e); + const message = musicalSymbolGClef.repeat(900); + const metadata = 'm'.repeat(3_000); + let retained = 0; + for (let revision = 0; revision < 64; revision++) { + if (!recorder.retainPreparing(optionalStep(revision, message, metadata))) break; + retained++; + } + const before = recorder.snapshot(); + const terminalText = musicalSymbolGClef.repeat(2_048); + const terminal = recorder.retainPreparing({ + version: 2, + attemptId: 'prepare_msg_1', + triggerMessageId: 'msg_1', + revision: retained, + timestamp: retained, + step: 'workspace_setup', + message: terminalText, + action: 'attempt_failed', + safeError: terminalText, + }); + const after = recorder.snapshot(); + + expect(retained).toBeGreaterThan(1); + expect(retainedBytes(before)).toBeGreaterThan( + Math.floor(MAX_SANDBOX_CONTROL_FRAME_BYTES / 2) * 0.99 + ); + expect(terminal).toBeDefined(); + expect(before.preparing.length - after.preparing.length).toBe(1); + expect(after.events).toHaveLength(0); + expect(after.preparing.length).toBeLessThanOrEqual(64); + expect(retainedBytes(after)).toBeLessThanOrEqual( + Math.floor(MAX_SANDBOX_CONTROL_FRAME_BYTES / 2) + ); + expect(after.preparing).toContainEqual( + expect.objectContaining({ action: 'attempt_failed', safeError: terminalText }) + ); + const wire = delivery(after.preparing); + expect(sessionOperationDeliverySchema.parse(wire)).toEqual(wire); + }); + + it.each(['attempt_completed', 'attempt_failed'] as const)( + 'reserves a final slot after 64 completed steps for %s', + action => { + const recorder = createRetainedOperationNotifications(); + for (let revision = 0; revision < 64; revision++) + recorder.retainPreparing(completedStep(revision)); + const terminal = + action === 'attempt_failed' + ? recorder.retainPreparing({ + version: 2, + attemptId: 'prepare_msg_1', + triggerMessageId: 'msg_1', + revision: 64, + timestamp: 64, + step: 'workspace_setup', + message: 'Preparation failed', + action, + safeError: 'Command failed', + }) + : recorder.retainPreparing({ + version: 2, + attemptId: 'prepare_msg_1', + triggerMessageId: 'msg_1', + revision: 64, + timestamp: 64, + step: 'workspace_setup', + message: 'Preparation completed', + action, + }); + const snapshot = recorder.snapshot(); + + expect(terminal).toBeDefined(); + expect(snapshot.preparing).toHaveLength(64); + expect(snapshot.preparing.filter(event => event.action === 'step_completed')).toHaveLength( + 63 + ); + expect(snapshot.preparing).toContainEqual(expect.objectContaining({ action })); + expect(retainedBytes(snapshot)).toBeLessThanOrEqual( + Math.floor(MAX_SANDBOX_CONTROL_FRAME_BYTES / 2) + ); + const wire = delivery(snapshot.preparing); + expect(sessionOperationDeliverySchema.parse(wire)).toEqual(wire); + } + ); +}); diff --git a/services/cloud-agent-next/wrapper/src/control/retained-operation-notifications.ts b/services/cloud-agent-next/wrapper/src/control/retained-operation-notifications.ts new file mode 100644 index 0000000000..6c9e02d156 --- /dev/null +++ b/services/cloud-agent-next/wrapper/src/control/retained-operation-notifications.ts @@ -0,0 +1,296 @@ +import { + MAX_SANDBOX_CONTROL_FRAME_BYTES, + sessionPreparingPayloadSchema, + type SessionEventPayload, + type SessionPreparingPayload, +} from '../../../src/shared/sandbox-control-protocol.js'; +import type { PreparingEventDataV2 } from '../../../src/shared/protocol.js'; + +const MAX_RETAINED_DELIVERY_EVENTS = 8; +const MAX_RETAINED_PREPARING_EVENTS = 64; +const MAX_RETAINED_DELIVERY_BYTES = Math.floor(MAX_SANDBOX_CONTROL_FRAME_BYTES / 2); +const RESERVED_TERMINAL_EVENTS = 1; +const RESERVED_TERMINAL_PREPARING_EVENTS = 8; +const RESERVED_ATTEMPT_TERMINALS = 1; +const MAX_EVENT_MESSAGE_LENGTH = 4_096; +const MAX_COMMIT_HASH_LENGTH = 128; +const MAX_TIMESTAMP_LENGTH = 128; + +type Entry = + | { kind: 'event'; payload: SessionEventPayload; terminal: boolean; bytes: number } + | { + kind: 'preparing'; + payload: SessionPreparingPayload; + terminal: boolean; + attemptTerminal: boolean; + bytes: number; + }; + +const preparingActions = new Set([ + 'attempt_started', + 'attempt_completed', + 'attempt_failed', + 'step_started', + 'step_completed', + 'step_failed', +]); + +const terminalPreparingActions = new Set([ + 'attempt_completed', + 'attempt_failed', + 'step_completed', + 'step_failed', +]); + +function boundedString(value: string, limit: number): string { + return value.length <= limit ? value : value.slice(0, limit); +} + +function boundedTimestamp(timestamp: string | undefined): string | undefined { + return timestamp === undefined ? undefined : boundedString(timestamp, MAX_TIMESTAMP_LENGTH); +} + +function projectFinalizationEvent(payload: SessionEventPayload): + | { + payload: SessionEventPayload; + terminal: boolean; + } + | undefined { + if (payload.type === 'autocommit_completed') { + const { properties } = payload; + if (typeof properties.success !== 'boolean' || typeof properties.messageId !== 'string') + return undefined; + const projected: Record = { + success: properties.success, + messageId: properties.messageId, + }; + if (typeof properties.skipped === 'boolean') projected.skipped = properties.skipped; + if (typeof properties.commitHash === 'string') + projected.commitHash = boundedString(properties.commitHash, MAX_COMMIT_HASH_LENGTH); + if (typeof properties.message === 'string') + projected.message = boundedString(properties.message, MAX_EVENT_MESSAGE_LENGTH); + if (typeof properties.commitMessage === 'string') + projected.commitMessage = boundedString(properties.commitMessage, MAX_EVENT_MESSAGE_LENGTH); + return { + payload: { + type: payload.type, + properties: projected, + ...(boundedTimestamp(payload.timestamp) + ? { timestamp: boundedTimestamp(payload.timestamp) } + : {}), + }, + terminal: true, + }; + } + if (payload.type === 'status') { + const { properties } = payload; + if (typeof properties.message !== 'string' || typeof properties.messageId !== 'string') + return undefined; + return { + payload: { + type: payload.type, + properties: { + message: boundedString(properties.message, MAX_EVENT_MESSAGE_LENGTH), + messageId: properties.messageId, + }, + ...(boundedTimestamp(payload.timestamp) + ? { timestamp: boundedTimestamp(payload.timestamp) } + : {}), + }, + terminal: false, + }; + } + return undefined; +} + +function projectPreparing(payload: PreparingEventDataV2): + | { + payload: PreparingEventDataV2; + terminal: boolean; + attemptTerminal: boolean; + } + | undefined { + if (!preparingActions.has(payload.action)) return undefined; + const common = { + version: payload.version, + attemptId: payload.attemptId, + triggerMessageId: payload.triggerMessageId, + revision: payload.revision, + timestamp: payload.timestamp, + step: payload.step, + message: boundedString(payload.message, MAX_EVENT_MESSAGE_LENGTH), + }; + let projected: PreparingEventDataV2; + switch (payload.action) { + case 'attempt_started': + case 'attempt_completed': + projected = { ...common, action: payload.action }; + break; + case 'attempt_failed': + projected = { + ...common, + action: payload.action, + safeError: boundedString(payload.safeError, MAX_EVENT_MESSAGE_LENGTH), + }; + break; + case 'step_started': + projected = { + ...common, + action: payload.action, + stepId: payload.stepId, + kind: payload.kind, + label: boundedString(payload.label, MAX_EVENT_MESSAGE_LENGTH), + ...(payload.command === undefined + ? {} + : { command: boundedString(payload.command, MAX_EVENT_MESSAGE_LENGTH) }), + ...(payload.commandIndex === undefined ? {} : { commandIndex: payload.commandIndex }), + ...(payload.commandCount === undefined ? {} : { commandCount: payload.commandCount }), + }; + break; + case 'step_completed': + projected = { + ...common, + action: payload.action, + stepId: payload.stepId, + ...(payload.exitCode === undefined ? {} : { exitCode: payload.exitCode }), + }; + break; + case 'step_failed': + projected = { + ...common, + action: payload.action, + stepId: payload.stepId, + safeError: boundedString(payload.safeError, MAX_EVENT_MESSAGE_LENGTH), + ...(payload.exitCode === undefined ? {} : { exitCode: payload.exitCode }), + }; + break; + default: + return undefined; + } + if (!sessionPreparingPayloadSchema.safeParse(projected).success) return undefined; + return { + payload: projected, + terminal: terminalPreparingActions.has(projected.action), + attemptTerminal: + projected.action === 'attempt_completed' || projected.action === 'attempt_failed', + }; +} + +function serializedBytes( + payload: SessionEventPayload | SessionPreparingPayload +): number | undefined { + try { + return Buffer.byteLength(JSON.stringify(payload)); + } catch { + return undefined; + } +} + +export function isRetainedOperationPreparing(payload: PreparingEventDataV2): boolean { + return preparingActions.has(payload.action); +} + +export function createRetainedOperationNotifications() { + const entries: Entry[] = []; + let bytes = 0; + + function count(kind: Entry['kind']): number { + return entries.filter(entry => entry.kind === kind).length; + } + + function terminalCount(kind: Entry['kind']): number { + return entries.filter(entry => entry.kind === kind && entry.terminal).length; + } + + function removeOldestOptional(kind?: Entry['kind']): boolean { + const index = entries.findIndex( + entry => !entry.terminal && (kind === undefined || entry.kind === kind) + ); + if (index === -1) return false; + const [removed] = entries.splice(index, 1); + if (removed) bytes -= removed.bytes; + return true; + } + + function removeOldestStepTerminal(): boolean { + const index = entries.findIndex( + entry => entry.kind === 'preparing' && entry.terminal && !entry.attemptTerminal + ); + if (index === -1) return false; + const [removed] = entries.splice(index, 1); + if (removed) bytes -= removed.bytes; + return true; + } + + function attemptTerminalCount(): number { + return entries.filter(entry => entry.kind === 'preparing' && entry.attemptTerminal).length; + } + + function retain( + kind: EntryKind, + payload: Extract['payload'], + terminal: boolean, + attemptTerminal = false + ): boolean { + const size = serializedBytes(payload); + if (size === undefined || size > MAX_RETAINED_DELIVERY_BYTES) return false; + const limit = kind === 'event' ? MAX_RETAINED_DELIVERY_EVENTS : MAX_RETAINED_PREPARING_EVENTS; + const reserved = + kind === 'event' ? RESERVED_TERMINAL_EVENTS : RESERVED_TERMINAL_PREPARING_EVENTS; + if (terminal) { + const terminalLimit = + kind === 'preparing' && !attemptTerminal + ? limit - Math.max(0, RESERVED_ATTEMPT_TERMINALS - attemptTerminalCount()) + : limit; + while (count(kind) >= terminalLimit) { + if (!removeOldestOptional(kind) && !(attemptTerminal && removeOldestStepTerminal())) + return false; + } + while (bytes + size > MAX_RETAINED_DELIVERY_BYTES) { + if (!removeOldestOptional() && !(attemptTerminal && removeOldestStepTerminal())) + return false; + } + } else { + const optionalLimit = limit - Math.max(0, reserved - terminalCount(kind)); + if (count(kind) >= optionalLimit || bytes + size > MAX_RETAINED_DELIVERY_BYTES) return false; + } + bytes += size; + entries.push({ + kind, + payload: structuredClone(payload), + terminal, + ...(kind === 'preparing' ? { attemptTerminal } : {}), + bytes: size, + } as Entry); + return true; + } + + return { + retainFinalization(payload: SessionEventPayload): SessionEventPayload | undefined { + const projected = projectFinalizationEvent(payload); + if (!projected || !retain('event', projected.payload, projected.terminal)) return undefined; + return structuredClone(projected.payload); + }, + retainPreparing(payload: PreparingEventDataV2): PreparingEventDataV2 | undefined { + const projected = projectPreparing(payload); + if ( + !projected || + !retain('preparing', projected.payload, projected.terminal, projected.attemptTerminal) + ) + return undefined; + return structuredClone(projected.payload); + }, + snapshot(): { events: SessionEventPayload[]; preparing: SessionPreparingPayload[] } { + return { + events: entries + .filter((entry): entry is Extract => entry.kind === 'event') + .map(entry => structuredClone(entry.payload)), + preparing: entries + .filter( + (entry): entry is Extract => entry.kind === 'preparing' + ) + .map(entry => structuredClone(entry.payload)), + }; + }, + }; +} 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 b6aa3b048e..c499f42552 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 @@ -1,7 +1,7 @@ import { describe, expect, it, mock, spyOn } from 'bun:test'; import { z } from 'zod'; import { helloResult } from '../../../src/sandbox-control/frames'; -import { buildHeartbeatPayload } from './sandbox-control-handlers'; +import { buildHeartbeatPayload, createControlHandlerDeps } from './sandbox-control-handlers'; import { MAX_SANDBOX_CONTROL_FRAME_BYTES, SANDBOX_CONTROL_REQUEST_TIMEOUT_MS, @@ -119,25 +119,26 @@ const previousHeartbeatSchema = z .strict(); function versionHeartbeat(version: string | null) { - return buildHeartbeatPayload({ - kiloRuntimes: { - kiloCliVersion: version, - attach() { - throw new Error('Unexpected attach'); + return buildHeartbeatPayload( + createControlHandlerDeps({ + kiloRuntimes: { + kiloCliVersion: version, + attach() { + throw new Error('Unexpected attach'); + }, + detach: () => false, + deleteDirectory: async () => {}, + get: () => undefined, + isHealthy: () => true, + shutdown() {}, }, - detach: () => false, - deleteDirectory: async () => {}, - get: () => undefined, - isHealthy: () => true, - shutdown() {}, - }, - version: '2.4.0', - kiloReady: true, - sessions: [], - tasks: new Map(), - emitSessionEvent() {}, - retireRuntime() {}, - }); + version: '2.4.0', + kiloReady: true, + sessions: [], + emitSessionEvent() {}, + retireRuntime() {}, + }) + ); } describe('heartbeat version rollout compatibility', () => { @@ -299,13 +300,19 @@ describe('createSandboxControlClient', () => { const hello = JSON.parse(fake.sent[0] ?? '{}') as { requestId: string; operation: string; - payload: { protocolVersion: number; wrapperVersion: string; providerInstanceId: string }; + payload: { + protocolVersion: number; + wrapperVersion: string; + providerInstanceId: string; + capabilities?: { sessionOperationResults?: boolean }; + }; }; expect(hello.operation).toBe('sandbox.hello'); expect(hello.payload).toEqual({ protocolVersion: 1, wrapperVersion: '2.4.0', providerInstanceId: 'inst_1', + capabilities: { sessionOperationResults: 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 630df7d63f..7c00e9488e 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 @@ -17,11 +17,17 @@ import { sandboxHelloResultSchema, sandboxHeartbeatPayloadSchema, sessionEventPayloadSchema, + sessionOperationDeliverySchema, + sessionOperationAckSchema, type ControlError, type EventFrame, + type ResponseFrame, type SessionEventPayload, + type SessionOperationAck, + type SessionOperationDelivery, type RequestFrame, type SessionEventIdentity, + type SessionOperationAuthorization, type SessionRequestIdentity, } from '../../../src/shared/sandbox-control-protocol.js'; @@ -33,7 +39,8 @@ type WebSocketCtor = new ( export type SandboxControlRequestHandler = ( operation: string, session: SessionRequestIdentity | undefined, - payload: unknown + payload: unknown, + authorization?: SessionOperationAuthorization ) => Promise<{ ok: boolean; result?: unknown; error?: ControlError }>; export type SandboxControlClientOptions = { @@ -50,14 +57,37 @@ export type SandboxControlClientOptions = { reconnectDelayMs?: (attempt: number) => number; }; +export class ControlDeliveryError extends Error { + constructor( + message: string, + readonly retryable: boolean + ) { + super(message); + } +} + +const PERMANENT_CONTROL_ERRORS = new Set([ + 'unauthorized', + 'protocol_error', + 'unknown_operation', + 'idempotency_conflict', +]); + export type SandboxControlClient = { connect(): Promise; close(): void; sendEvent?( event: string, payload: unknown, - session?: { directory: string; kiloSessionId?: string; rootKiloSessionId?: string } + session?: { directory: string; kiloSessionId?: string; rootKiloSessionId?: string }, + options?: { preserveConnectionOnFailure?: boolean } ): boolean; + sendOperationResult?( + session: SessionRequestIdentity, + delivery: SessionOperationDelivery, + signal: AbortSignal, + deadlineAt: number + ): Promise; }; type ClientState = @@ -146,6 +176,24 @@ export function createSandboxControlClient( const wrapperInstanceId = options.wrapperInstanceId; let state: ClientState = { kind: 'idle' }; let eventSequence = 0; + const pendingRequests = new Map< + string, + { resolve: (frame: ResponseFrame) => void; reject: (reason: unknown) => void } + >(); + + function settleResponse(frame: ResponseFrame): void { + const waiter = pendingRequests.get(frame.requestId); + if (!waiter) return; + pendingRequests.delete(frame.requestId); + waiter.resolve(frame); + } + + function rejectPendingRequests(reason: string): void { + for (const [id, waiter] of pendingRequests) { + pendingRequests.delete(id); + waiter.reject(new ControlDeliveryError(reason, true)); + } + } const diagnostic = (phase: string, ws?: WebSocket): void => emitControlDiagnostic(options.onDiagnostic, 'control.socket', { phase, @@ -157,6 +205,7 @@ export function createSandboxControlClient( if (state.kind !== 'ready' || state.socket !== ws) return; const current = state; diagnostic('retired', ws); + rejectPendingRequests('Sandbox control connection closed'); state = { kind: 'closed' }; current.dispose(); options.onDisconnected?.(); @@ -182,7 +231,12 @@ export function createSandboxControlClient( requestDiagnostic('received'); let outcome: Awaited>; try { - outcome = await options.onRequest(request.operation, request.session, request.payload); + outcome = await options.onRequest( + request.operation, + request.session, + request.payload, + request.authorization + ); } catch { outcome = { ok: false, @@ -318,6 +372,7 @@ export function createSandboxControlClient( payload: { protocolVersion: SANDBOX_CONTROL_PROTOCOL_VERSION, providerInstanceId: options.providerInstanceId, + capabilities: { sessionOperationResults: true }, ...(wrapperInstanceId ? { wrapperInstanceId } : {}), ...(options.wrapperVersion ? { wrapperVersion: options.wrapperVersion } : {}), }, @@ -344,7 +399,8 @@ export function createSandboxControlClient( if (!parsed.success) return; const frame = parsed.data; if (state.kind === 'ready') { - if (frame.type === 'request') void dispatchRequest(ws, frame); + if (frame.type === 'response') settleResponse(frame); + else if (frame.type === 'request') void dispatchRequest(ws, frame); return; } if (Date.now() >= deadlineAt) { @@ -487,13 +543,116 @@ export function createSandboxControlClient( close(): void { const current = state; diagnostic('closed', current.kind === 'ready' ? current.socket : undefined); + rejectPendingRequests('Sandbox control client closed'); state = { kind: 'closed' }; if (current.kind === 'starting') current.abort.abort(new Error('sandbox control client closed')); else if (current.kind === 'ready') current.dispose(); }, - sendEvent(event: string, payload: unknown, session?: SessionEventIdentity): boolean { + async sendOperationResult( + session: SessionRequestIdentity, + delivery: SessionOperationDelivery, + signal: AbortSignal, + deadlineAt: number + ): Promise { + signal.throwIfAborted(); + if (Date.now() >= deadlineAt) + throw new ControlDeliveryError('Control delivery expired', false); + if (state.kind !== 'ready' || state.socket.readyState !== 1) + throw new ControlDeliveryError('Control transport unavailable', true); + const { socket } = state; + const requestId = crypto.randomUUID(); + let payload: SessionOperationDelivery; + try { + payload = sessionOperationDeliverySchema.parse(delivery); + } catch { + throw new ControlDeliveryError('Operation result payload is invalid', false); + } + const frame: RequestFrame = { + type: 'request', + requestId, + operation: 'session.operation.result', + session, + payload, + }; + let serialized: string; + try { + serialized = JSON.stringify(frame); + } catch { + throw new ControlDeliveryError('Operation result cannot be serialized', false); + } + if (Buffer.byteLength(serialized) > MAX_SANDBOX_CONTROL_FRAME_BYTES) + throw new ControlDeliveryError('Control delivery exceeds the frame budget', false); + const pending = new Promise((resolve, reject) => { + const timeout = setTimeout( + () => { + pendingRequests.delete(requestId); + reject(new ControlDeliveryError('Operation result delivery timed out', true)); + }, + Math.max(1, Math.min(SANDBOX_CONTROL_REQUEST_TIMEOUT_MS, deadlineAt - Date.now())) + ); + timeout.unref(); + pendingRequests.set(requestId, { + resolve: frame => { + clearTimeout(timeout); + resolve(frame); + }, + reject: reason => { + clearTimeout(timeout); + reject(reason); + }, + }); + }); + const onAbort = () => { + const waiter = pendingRequests.get(requestId); + if (waiter) { + pendingRequests.delete(requestId); + waiter.reject(new ControlDeliveryError('Control delivery cancelled', false)); + } + }; + signal.addEventListener('abort', onAbort, { once: true }); + try { + signal.throwIfAborted(); + socket.send(serialized); + } catch { + const waiter = pendingRequests.get(requestId); + if (waiter) { + pendingRequests.delete(requestId); + waiter.reject(new ControlDeliveryError('Control delivery publication failed', true)); + } + } + try { + const response = await pending; + signal.throwIfAborted(); + if ( + Date.now() >= deadlineAt || + state.kind !== 'ready' || + state.socket !== socket || + socket.readyState !== 1 + ) + throw new ControlDeliveryError('Control delivery acknowledgement is stale', true); + if (!response.ok) + throw new ControlDeliveryError( + 'Control delivery was not acknowledged', + response.error?.retryable === true && !PERMANENT_CONTROL_ERRORS.has(response.error.code) + ); + try { + return sessionOperationAckSchema.parse(response.result); + } catch { + throw new ControlDeliveryError('Control delivery acknowledgement is invalid', false); + } + } finally { + signal.removeEventListener('abort', onAbort); + } + }, + + sendEvent( + event: string, + payload: unknown, + session?: SessionEventIdentity, + deliveryOptions?: { preserveConnectionOnFailure?: boolean } + ): boolean { eventSequence += 1; const category = event === 'session.event' @@ -526,7 +685,7 @@ export function createSandboxControlClient( const { socket } = state; if (socket.readyState !== 1) { eventDiagnostic('send_failed'); - retireConnection(socket); + if (!deliveryOptions?.preserveConnectionOnFailure) retireConnection(socket); return false; } try { @@ -539,7 +698,7 @@ export function createSandboxControlClient( return true; } catch { eventDiagnostic('send_failed'); - retireConnection(socket); + if (!deliveryOptions?.preserveConnectionOnFailure) retireConnection(socket); return false; } }, 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 376bd97692..58140da0c4 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 @@ -153,7 +153,6 @@ function deps( version: '2.4.0', kiloReady: true, sessions: [], - tasks: new Map(), emitSessionEvent: () => {}, retireRuntime: () => {}, applyAttach: (session, payload, deps) => @@ -167,8 +166,26 @@ function runtimeDeps(kiloClient: WrapperKiloClient) { const events: SessionEventPayload[] = []; const retired: string[] = []; let shuttingDown = false; - const handlerDeps: HandlerDeps = { - ...deps({ kiloClient }), + const handlerDeps: HandlerDeps = createControlHandlerDeps({ + ...(() => { + const base = deps({ kiloClient }); + return { + kiloRuntimes: base.kiloRuntimes, + worktreeCleanupClient: base.worktreeCleanupClient, + version: base.version, + sessions: base.sessions, + activity: base.activity, + nativeObservations: base.nativeObservations, + terminalRuntime: base.terminalRuntime, + applyAttach: base.applyAttach, + materializeAttachments: base.materializeAttachments, + runAutoCommit: base.runAutoCommit, + collectWorktreeChanges: base.collectWorktreeChanges, + onDiagnostic: base.onDiagnostic, + onShutdown: base.onShutdown, + emitPreparing: base.emitPreparing, + }; + })(), get kiloReady() { return !shuttingDown; }, @@ -181,12 +198,12 @@ function runtimeDeps(kiloClient: WrapperKiloClient) { void cancelControlTasks(handlerDeps, reason, 'failed'); abort.abort(); }, - }; + }); return { handlerDeps, events, retired }; } async function waitForTasks(handlerDeps: HandlerDeps): Promise { - await Promise.all([...handlerDeps.tasks.values()].map(task => task.done)); + await Promise.all(handlerDeps.operations.activeOperations().map(task => task.done)); } const promptPayload = { @@ -390,7 +407,7 @@ describe('handleControlRequest', () => { deps({ kiloClient }) ); - expect(result).toEqual({ + expect(result).toMatchObject({ ok: false, error: { code: 'protocol_error', message: 'Invalid payload', retryable: false }, }); @@ -523,7 +540,7 @@ describe('handleControlRequest', () => { ).toMatchObject({ ok: false }); expect(activity.snapshots()).toEqual([]); expect(rootForSession(session.kiloSessionId)).toBeUndefined(); - expect(handlerDeps.tasks.size).toBe(0); + expect(handlerDeps.operations.counts().active).toBe(0); }); it('marks only accepted root work active and leaves rejected sibling prompts idle', async () => { @@ -549,7 +566,7 @@ describe('handleControlRequest', () => { { ...promptPayload, messageId: 'msg_2', agent: { mode: 'code' } }, handlerDeps ) - ).toEqual({ + ).toMatchObject({ ok: false, error: { code: 'protocol_error', message: 'Invalid payload', retryable: false }, }); @@ -593,7 +610,7 @@ describe('handleControlRequest', () => { { env: { [name]: 'sensitive-value' } }, deps({ kiloClient }) ); - expect(result).toEqual({ + expect(result).toMatchObject({ ok: false, error: { code: 'protocol_error', @@ -918,14 +935,16 @@ describe('handleControlRequest', () => { ) ).toEqual({ ok: true, result: { messageId: cancellingId, status: 'accepted' } }); await pending.started.promise; - expect( - await handleControlRequest( - 'session.abort', - identity, - { messageId: cancellingId }, - handlerDeps - ) - ).toEqual({ + const aborting = handleControlRequest( + 'session.abort', + identity, + { messageId: cancellingId }, + handlerDeps + ); + pending.completion.resolve( + completion({ name: 'MessageAbortedError', data: { message: 'cancelled' } }) + ); + expect(await aborting).toEqual({ ok: true, result: { status: 'aborted' }, }); @@ -938,7 +957,7 @@ describe('handleControlRequest', () => { signal: expect.any(AbortSignal), }, }); - expect(handlerDeps.tasks.has(identity.kiloSessionId)).toBe(false); + expect(handlerDeps.operations.hasActive(identity.kiloSessionId)).toBe(false); } finally { pending.completion.resolve(completion()); await waitForTasks(handlerDeps); @@ -988,7 +1007,7 @@ describe('handleControlRequest', () => { let confirmed = false; const handlerDeps = deps({ kiloClient: fakeKilo({ abortSession: async () => confirmed }) }); const runtime = handlerDeps.kiloRuntimes?.get(session.directory); - expect(handlerDeps.tasks.size).toBe(0); + expect(handlerDeps.operations.counts().active).toBe(0); expect(await handleControlRequest('session.detach', session, {}, handlerDeps)).toMatchObject({ ok: false, error: { code: 'not_ready' }, @@ -1045,7 +1064,7 @@ describe('handleControlRequest', () => { try { await handleControlRequest('session.prompt', session, promptPayload, handlerDeps); await started.promise; - const task = handlerDeps.tasks.get(session.kiloSessionId); + const task = handlerDeps.operations.active(session.kiloSessionId); const snapshots = [...handlerDeps.sessions]; expect( await handleControlRequest( @@ -1062,7 +1081,7 @@ describe('handleControlRequest', () => { expect(rootForSession('child_1')).toBe(session.kiloSessionId); expect(rootForSession('child_2')).toBe(sibling.kiloSessionId); expect(directoryForSession(session.kiloSessionId)).toBe(session.directory); - expect(handlerDeps.tasks.get(session.kiloSessionId)).toBe(task); + expect(handlerDeps.operations.active(session.kiloSessionId)).toBe(task); expect(task?.signal.aborted).toBe(false); expect(handlerDeps.sessions).toEqual(snapshots); expect(activity.snapshots()).toEqual([ @@ -1116,10 +1135,12 @@ describe('handleControlRequest', () => { }); expect(shutdowns).toBe(0); expect(runtimes.get(session.directory)).toBe(runtime); - expect(handlerDeps.tasks.get(sibling.kiloSessionId)?.signal.aborted).toBe(false); + expect(handlerDeps.operations.active(sibling.kiloSessionId)?.signal.aborted).toBe(false); expect(events).toEqual([]); expect(aborted).toEqual([session.kiloSessionId]); - expect(await handleControlRequest('session.abort', sibling, {}, handlerDeps)).toEqual({ + const aborting = handleControlRequest('session.abort', sibling, {}, handlerDeps); + running.resolve(completion({ name: 'MessageAbortedError', data: { message: 'cancelled' } })); + expect(await aborting).toEqual({ ok: true, result: { status: 'aborted' }, }); @@ -1130,7 +1151,7 @@ describe('handleControlRequest', () => { properties: { messageId: promptPayload.messageId, status: 'cancelled', - reason: 'Session aborted', + reason: 'Kilo execution ended with MessageAbortedError', }, }, ]); @@ -1704,7 +1725,7 @@ describe('production worktree deletion routes', () => { lookups.length = 0; expect(rootForSession(sessionId(2))).toBeUndefined(); expect(rootForSession(sessionId(3))).toBeUndefined(); - const siblingTask = handlerDeps.tasks.get(sibling.kiloSessionId); + const siblingTask = handlerDeps.operations.active(sibling.kiloSessionId); preparation = handleControlRequest( 'worktree.prepareDeletion', undefined, @@ -1715,8 +1736,8 @@ describe('production worktree deletion routes', () => { return result; }); await Promise.all([...abortStarted.values()].map(item => item.promise)); - expect(handlerDeps.tasks.get(first.kiloSessionId)?.signal.aborted).toBe(true); - expect(handlerDeps.tasks.get(second.kiloSessionId)?.signal.aborted).toBe(true); + expect(handlerDeps.operations.active(first.kiloSessionId)?.signal.aborted).toBe(true); + expect(handlerDeps.operations.active(second.kiloSessionId)?.signal.aborted).toBe(true); expect(siblingTask?.signal.aborted).toBe(false); expect(prepared).toBe(false); expect(http.requests).toEqual([]); @@ -1728,10 +1749,16 @@ describe('production worktree deletion routes', () => { }); expect(attach).not.toHaveBeenCalled(); abortResponses.get(first.kiloSessionId)?.resolve(true); - await handlerDeps.tasks.get(first.kiloSessionId)?.done; + completionResponses + .get(first.kiloSessionId) + ?.resolve(completion({ name: 'MessageAbortedError', data: { message: 'cancelled' } })); + await handlerDeps.operations.active(first.kiloSessionId)?.done; expect(prepared).toBe(false); expect(http.requests).toEqual([]); abortResponses.get(second.kiloSessionId)?.resolve(true); + completionResponses + .get(second.kiloSessionId) + ?.resolve(completion({ name: 'MessageAbortedError', data: { message: 'cancelled' } })); expect( await withTimeoutAndAbort(preparation, { timeoutMs: 1_000, @@ -1745,8 +1772,8 @@ describe('production worktree deletion routes', () => { sessionIds: [sessionId(1), sessionId(6), sessionId(2), sessionId(3)], }, }); - expect(handlerDeps.tasks.size).toBe(1); - expect(handlerDeps.tasks.get(sibling.kiloSessionId)).toBe(siblingTask); + expect(handlerDeps.operations.counts().active).toBe(1); + expect(handlerDeps.operations.active(sibling.kiloSessionId)).toBe(siblingTask); expect(siblingTask?.signal.aborted).toBe(false); expect(aborted.toSorted()).toEqual([first.kiloSessionId, second.kiloSessionId]); expect(outcomes.map(item => item.id).toSorted()).toEqual([ @@ -2152,7 +2179,7 @@ describe('owned control execution', () => { updateSessionSnapshots(event, handlerDeps.sessions); } expect(events).toEqual([]); - expect(handlerDeps.tasks.size).toBe(1); + expect(handlerDeps.operations.counts().active).toBe(1); finished.resolve(completion()); await waitForTasks(handlerDeps); expect(events[0]?.properties).toEqual({ messageId: 'msg_1', status: 'completed' }); @@ -2218,22 +2245,27 @@ describe('owned control execution', () => { await handleControlRequest('session.prompt', session, payload, handlerDeps) ).toMatchObject({ ok: false }); expect(events).toEqual([]); - expect(handlerDeps.tasks.get(session.kiloSessionId)?.messageId).toBe('newer'); + expect(handlerDeps.operations.active(session.kiloSessionId)?.messageId).toBe('newer'); remoteStopped.resolve(true); + running.resolve(completion({ name: 'MessageAbortedError', data: { message: 'cancelled' } })); expect(await aborting).toEqual({ ok: true, result: { status: 'aborted' } }); expect(events).toEqual([ { type: 'session.message.outcome', - properties: { messageId: 'newer', status: 'cancelled', reason: 'Session aborted' }, + properties: { + messageId: 'newer', + status: 'cancelled', + reason: 'Kilo execution ended with MessageAbortedError', + }, }, ]); - expect(handlerDeps.tasks.size).toBe(0); + expect(handlerDeps.operations.counts().active).toBe(0); expect(await handleControlRequest('session.prompt', session, payload, handlerDeps)).toEqual({ ok: true, result: { messageId: 'msg_1', status: 'accepted' }, }); await replacementStarted.promise; - const replacementTask = handlerDeps.tasks.get(session.kiloSessionId); + const replacementTask = handlerDeps.operations.active(session.kiloSessionId); expect(replacementTask?.messageId).toBe('msg_1'); const oldCompletion = completion(); running.resolve({ @@ -2241,7 +2273,7 @@ describe('owned control execution', () => { info: { ...oldCompletion.info, parentID: 'newer' }, }); await new Promise(resolve => setImmediate(resolve)); - expect(handlerDeps.tasks.get(session.kiloSessionId)).toBe(replacementTask); + expect(handlerDeps.operations.active(session.kiloSessionId)).toBe(replacementTask); expect(replacementTask?.signal.aborted).toBe(false); expect(finalized).toEqual([]); expect(events).toHaveLength(1); @@ -2305,7 +2337,7 @@ describe('owned control execution', () => { expect(retired).toEqual(['Kilo cancellation failed']); expect(handlerDeps.signal?.aborted).toBe(true); expect(buildHeartbeatPayload(handlerDeps).kilo.ready).toBe(false); - expect(handlerDeps.tasks.size).toBe(0); + expect(handlerDeps.operations.counts().active).toBe(0); expect( await handleControlRequest( 'session.prompt', @@ -2361,7 +2393,7 @@ describe('owned control execution', () => { ).toMatchObject({ ok: false, error: { code: 'not_ready' } }); expect(retired).toEqual(['Kilo cancellation failed']); expect(handlerDeps.signal?.aborted).toBe(true); - expect(handlerDeps.tasks.size).toBe(0); + expect(handlerDeps.operations.counts().active).toBe(0); expect( await handleControlRequest( 'session.prompt', @@ -2408,7 +2440,7 @@ describe('owned control execution', () => { ); const abortSignal = await abortStarted.promise; expect(abortSignal.aborted).toBe(false); - expect(handlerDeps.tasks.get(session.kiloSessionId)?.messageId).toBe('msg_1'); + expect(handlerDeps.operations.active(session.kiloSessionId)?.messageId).toBe('msg_1'); expect(events).toEqual([]); expect( await handleControlRequest( @@ -2418,11 +2450,13 @@ describe('owned control execution', () => { handlerDeps ) ).toMatchObject({ ok: false }); - const deadline = timers.mock.calls.find( - ([, ms]) => ms === KILO_CONTROL_REQUEST_TIMEOUT_MS - )?.[0]; - if (typeof deadline !== 'function') throw new Error('Missing abort request deadline'); - deadline(); + const deadlines = timers.mock.calls + .filter(([, ms]) => ms === KILO_CONTROL_REQUEST_TIMEOUT_MS) + .slice(-2) + .map(([callback]) => callback); + if (deadlines.length !== 2 || deadlines.some(callback => typeof callback !== 'function')) + throw new Error('Missing abort request deadlines'); + for (const deadline of deadlines) deadline(); expect(await aborting).toMatchObject({ ok: false, error: { code: 'not_ready' } }); expect(abortSignal.aborted).toBe(true); expect(retired).toEqual(['Kilo cancellation failed']); @@ -2430,7 +2464,7 @@ describe('owned control execution', () => { remoteStopped.resolve(true); running.resolve(completion()); await new Promise(resolve => setImmediate(resolve)); - expect(handlerDeps.tasks.size).toBe(0); + expect(handlerDeps.operations.counts().active).toBe(0); expect(events).toHaveLength(1); expect(buildHeartbeatPayload(handlerDeps).kilo.ready).toBe(false); expect( @@ -2461,6 +2495,12 @@ describe('owned control execution', () => { started.resolve(); return running.promise; }, + abortSession: async () => { + running.resolve( + completion({ name: 'MessageAbortedError', data: { message: 'cancelled' } }) + ); + return true; + }, }), emitSessionEvent: (_session, event) => events.push(event), retireRuntime: reason => { @@ -2483,14 +2523,14 @@ describe('owned control execution', () => { const deadline = timers.mock.calls.find( ([, ms]) => ms === SANDBOX_CONTROL_EXECUTION_TIMEOUT_MS )?.[0]; - if (typeof deadline !== 'function') throw new Error('missing owned execution deadline'); + if (typeof deadline !== 'function') throw new Error('Missing owned execution deadline'); deadline(); await waitForTasks(handlerDeps); expect(retired).toEqual(['Execution exceeded the 60 minute limit']); expect(events[0]?.properties).toEqual({ messageId: 'msg_1', - status: 'failed', - reason: 'Execution exceeded the 60 minute limit', + status: 'cancelled', + reason: 'Kilo execution ended with MessageAbortedError', }); expect(buildHeartbeatPayload(handlerDeps).kilo.ready).toBe(false); expect( @@ -2793,7 +2833,7 @@ describe('control finalization and compact', () => { if (typeof deadline !== 'function') throw new Error('Missing owned execution deadline'); deadline(); expect(signal.aborted).toBe(true); - expect(handlerDeps.tasks.size).toBe(1); + expect(handlerDeps.operations.counts().active).toBe(1); stopped.resolve({ success: false }); await waitForTasks(handlerDeps); expect(retired).toEqual(['Execution exceeded the 60 minute limit']); @@ -2891,7 +2931,7 @@ describe('control finalization and compact', () => { handlerDeps ) ).toMatchObject({ ok: false, error: { code: 'protocol_error' } }); - expect(handlerDeps.tasks.size).toBe(0); + expect(handlerDeps.operations.counts().active).toBe(0); expect( await handleControlRequest('session.prompt', session, payload, handlerDeps) ).toMatchObject({ ok: true }); @@ -2979,7 +3019,7 @@ describe('control cancellation and attachments', () => { ); expect(cancelled).toMatchObject({ ok: true }); expect(await siblingAttach).toMatchObject({ ok: false }); - expect(handlerDeps.tasks.has(sibling.kiloSessionId)).toBe(false); + expect(handlerDeps.operations.hasActive(sibling.kiloSessionId)).toBe(false); expect(signal.aborted).toBe(false); expect(buildHeartbeatPayload(handlerDeps).activeKiloSessions).toBe(1); @@ -2999,7 +3039,7 @@ describe('control cancellation and attachments', () => { await new Promise(resolve => setImmediate(resolve)); expect(signal.aborted).toBe(true); expect(cancellationSettled).toBe(false); - expect(handlerDeps.tasks.has(session.kiloSessionId)).toBe(true); + expect(handlerDeps.operations.hasActive(session.kiloSessionId)).toBe(true); expect(setupCommands).toEqual(['first']); stopped.resolve({ exitCode: 0, stdout: '', stderr: '' }); @@ -3030,6 +3070,12 @@ describe('control cancellation and attachments', () => { started.resolve(); return running.promise; }, + abortSession: async () => { + running.resolve( + completion({ name: 'MessageAbortedError', data: { message: 'cancelled' } }) + ); + return true; + }, }), applyAttach: (identity, payload, dependencies) => applySessionAttach(identity, payload, { @@ -3044,7 +3090,7 @@ describe('control cancellation and attachments', () => { terminalRuntime: fakeTerminalRuntime({ detachSession: async () => { expect(workSignal?.aborted).toBe(true); - expect(handlerDeps.tasks.size).toBe(0); + expect(handlerDeps.operations.counts().active).toBe(0); detached = true; }, }), @@ -3096,7 +3142,7 @@ describe('control cancellation and attachments', () => { }, terminalRuntime: fakeTerminalRuntime({ shutdown: () => { - expect(handlerDeps.tasks.size).toBe(0); + expect(handlerDeps.operations.counts().active).toBe(0); terminalStopped = true; }, }), @@ -3111,6 +3157,12 @@ describe('control cancellation and attachments', () => { if (opts.signal) signals.push(opts.signal); return running.promise; }, + abortSession: async () => { + running.resolve( + completion({ name: 'MessageAbortedError', data: { message: 'cancelled' } }) + ); + return true; + }, }), }, secondSession @@ -3139,7 +3191,7 @@ describe('control cancellation and attachments', () => { handlerDeps ) ).toEqual({ ok: true, result: { messageId: promptPayload.messageId, status: 'accepted' } }); - expect(handlerDeps.tasks.size).toBe(2); + expect(handlerDeps.operations.counts().active).toBe(2); expect(await handleControlRequest('sandbox.shutdown', undefined, {}, handlerDeps)).toEqual({ ok: true, result: { shuttingDown: true }, @@ -3203,11 +3255,11 @@ describe('control cancellation and attachments', () => { ).toMatchObject({ ok: false, }); - expect(handlerDeps.tasks.size).toBe(1); + expect(handlerDeps.operations.counts().active).toBe(1); stopped.resolve({ exitCode: 0, stdout: '', stderr: '' }); expect(await attaching).toMatchObject({ ok: false }); expect(markerWritten).toBe(false); - expect(handlerDeps.tasks.size).toBe(0); + expect(handlerDeps.operations.counts().active).toBe(0); expect(retired).toEqual(['Session preparation timed out']); expect(buildHeartbeatPayload(handlerDeps).kilo.ready).toBe(false); expect( @@ -4194,7 +4246,7 @@ describe('buildHeartbeatPayload', () => { running.resolve(completion()); await finalizing.promise; } - const task = handlerDeps.tasks.get(session.kiloSessionId); + const task = handlerDeps.operations.active(session.kiloSessionId); expect(task?.kind).toBe(phase); expect(await refreshHeartbeatPayload(handlerDeps)).toMatchObject({ state: phase === 'finalizing' ? 'finalizing' : 'active', @@ -4210,7 +4262,7 @@ describe('buildHeartbeatPayload', () => { ], }); expect(activity.snapshots().every(snapshot => snapshot.state === 'idle')).toBe(true); - expect(handlerDeps.tasks.get(session.kiloSessionId)).toBe(task); + expect(handlerDeps.operations.active(session.kiloSessionId)).toBe(task); expect(events).toEqual([]); running.resolve(completion()); finalized.resolve({ success: true }); @@ -4390,7 +4442,7 @@ describe('refreshHeartbeatPayload', () => { kilo: { ready: true }, sessions: [{ kiloSessionId: session.kiloSessionId, state: 'idle', idleForMs: 0 }], }); - expect(handlerDeps.tasks.size).toBe(0); + expect(handlerDeps.operations.counts().active).toBe(0); expect(events).toEqual([ { type: 'session.message.outcome', @@ -4459,7 +4511,7 @@ describe('refreshHeartbeatPayload', () => { { kiloSessionId: session.kiloSessionId, state: 'idle', idleForMs: 0 }, ], }); - expect(handlerDeps.tasks.size).toBe(0); + expect(handlerDeps.operations.counts().active).toBe(0); } finally { statuses.resolve({}); await refresh; @@ -4975,7 +5027,7 @@ describe('session.git.summary', () => { expect(directories).toEqual([session.directory, session.directory]); expect(activity.snapshots()).toEqual(snapshots); expect(handlerDeps.sessions).toEqual([]); - expect(handlerDeps.tasks.size).toBe(0); + expect(handlerDeps.operations.counts().active).toBe(0); }); it('drains in-flight capture before deletion and rejects late results without fencing another worktree', async () => { @@ -5026,7 +5078,7 @@ describe('session.git.summary', () => { }); await deletion; expect(fenced).toBe(true); - expect(handlerDeps.tasks.size).toBe(0); + expect(handlerDeps.operations.counts().active).toBe(0); expect(handlerDeps.sessions).toEqual([]); } finally { capture.resolve(captured); @@ -5084,7 +5136,7 @@ describe('session.git.summary', () => { expect(rootForSession(sibling.kiloSessionId, sibling.directory)).toBe( sibling.kiloSessionId ); - expect(handlerDeps.tasks.size).toBe(0); + expect(handlerDeps.operations.counts().active).toBe(0); } finally { capture.resolve(captured); await request; 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 fad71def27..f6dec67e6a 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 @@ -6,15 +6,11 @@ import { type ControlDiagnosticReporter, } from '../../../src/shared/control-diagnostics.js'; import { - SANDBOX_CONTROL_ATTACH_TIMEOUT_MS, - SANDBOX_CONTROL_EXECUTION_TIMEOUT_MS, SESSION_OPERATIONS, sandboxShutdownPayloadSchema, sessionAbortPayloadSchema, sessionAttachPayloadSchema, sessionDetachPayloadSchema, - sessionMessageOutcomeSchema, - sessionEventPayloadSchema, sessionGitSummaryPayloadSchema, sessionPermissionResolvePayloadSchema, sessionPromptPayloadSchema, @@ -31,16 +27,18 @@ import { worktreeDeletePayloadSchema, type SandboxHeartbeatPayload, type SessionEventPayload, - type SessionMessageOutcome, + type SessionOperationAuthorization, + type SessionOperationDelivery, + type SessionOperationAck, type SessionPromptPayload, type SessionRequestIdentity, } from '../../../src/shared/sandbox-control-protocol.js'; import { CONTROL_RUNTIME_RESERVED_ENV_VARS } from '../../../src/shared/runtime-environment.js'; import { isKiloServerUnreachableError, type WrapperKiloClient } from '../kilo-api.js'; -import { materializeMessageAttachments } from '../session-bootstrap.js'; -import { runAutoCommit } from '../auto-commit.js'; -import type { IngestEvent } from '../../../src/shared/protocol.js'; -import { withTimeoutAndAbort } from '../utils.js'; +import type { materializeMessageAttachments } from '../session-bootstrap.js'; +import type { runAutoCommit } from '../auto-commit.js'; +import { rejectBeforeAdmission, type ControlHandlerResult } from './control-handler-result.js'; +import { createOperationRegistry, type OperationRegistry } from './operation-registry.js'; import { applySessionAttach, type AttachPreparingEmitter } from './apply-attach'; import { directoriesForRoot, @@ -48,7 +46,10 @@ import { forgetAttachedRoot, rootForSession, } from './session-directories'; -import { withKiloRequestDeadline } from './sandbox-control-runtime'; +import { + KILO_CONTROL_REQUEST_TIMEOUT_MS, + withKiloRequestDeadline, +} from './sandbox-control-runtime'; import { ControlTerminalRuntimeError, type ControlTerminalRuntime } from './terminal-runtime.js'; import { WorktreeKiloRuntimeError, @@ -70,23 +71,15 @@ import { import { collectWorktreeChanges } from './worktree-changes'; import { createNativeObservations, type NativeObservations } from './native-observations.js'; +export type { ControlHandlerResult } from './control-handler-result.js'; +export type { SessionOperation as OwnedSessionTask } from './session-operation.js'; + export type HandlerSessionSnapshot = { kiloSessionId: string; lastActivityAt: number; pendingInputs?: Set; }; -type TaskIdentity = - | { kind: 'preparation'; messageId?: string } - | { kind: 'execution' | 'finalizing'; messageId: string }; - -export type OwnedSessionTask = TaskIdentity & { - session: SessionRequestIdentity; - controller: AbortController; - signal: AbortSignal; - done: Promise; -}; - type SessionActivity = { revision: symbol; state: 'idle' | 'active' | 'finalizing'; @@ -215,14 +208,24 @@ export function createSessionActivityRegistry( export type HandlerDeps = { kiloRuntimes?: WorktreeKiloRuntimes; worktreeCleanupClient?: WorktreeKiloCleanupClient; + operations: OperationRegistry; version: string; kiloReady: boolean; sessions: HandlerSessionSnapshot[]; - tasks: Map; + sendOperationResult?: ( + session: SessionRequestIdentity, + delivery: SessionOperationDelivery, + signal: AbortSignal, + deadlineAt: number + ) => Promise; signal?: AbortSignal; activity?: SessionActivityRegistry; nativeObservations?: NativeObservations; - emitSessionEvent: (session: SessionRequestIdentity, payload: SessionEventPayload) => void; + emitSessionEvent: ( + session: SessionRequestIdentity, + payload: SessionEventPayload, + options?: { retained?: true } + ) => void; retireRuntime: (reason: string) => void; onShutdown?: () => void; onDiagnostic?: ControlDiagnosticReporter; @@ -234,19 +237,6 @@ export type HandlerDeps = { collectWorktreeChanges?: typeof collectWorktreeChanges; }; -export type ControlHandlerResult = - | { ok: true; result: unknown } - | { ok: false; error: { code: string; message: string; retryable: boolean } }; - -class ControlTaskCancellation extends Error { - constructor( - readonly status: 'failed' | 'cancelled', - message: string - ) { - super(message); - } -} - const SESSION_OPERATION_SET = new Set(SESSION_OPERATIONS); function ok(result: unknown): ControlHandlerResult { @@ -261,13 +251,31 @@ function kiloFailure(error: unknown): ControlHandlerResult { return fail('not_ready', 'Kilo request failed', isKiloServerUnreachableError(error)); } +function operationEffects(session: SessionRequestIdentity, deps: HandlerDeps) { + const send = deps.sendOperationResult; + return { + signal: deps.signal, + onDiagnostic: deps.onDiagnostic, + emitSessionEvent: (event: SessionEventPayload, options?: { retained?: true }) => + deps.emitSessionEvent(session, event, options), + sendOperationResult: send + ? (delivery: SessionOperationDelivery, signal: AbortSignal, deadlineAt: number) => + send(session, delivery, signal, deadlineAt) + : undefined, + }; +} + +export function pruneControlOperations(deps: HandlerDeps, now = Date.now()): void { + deps.operations.prune(now); +} + export function buildHeartbeatPayload(deps: HandlerDeps): SandboxHeartbeatPayload { const now = Date.now(); const snapshots = new Map( deps.activity?.snapshots().map(snapshot => [snapshot.kiloSessionId, snapshot]) ); for (const snapshot of deps.sessions) { - const task = deps.tasks.get(snapshot.kiloSessionId); + const task = deps.operations.active(snapshot.kiloSessionId); if (!task && snapshots.has(snapshot.kiloSessionId)) continue; const waitingOn = task?.kind === 'preparation' @@ -294,7 +302,7 @@ export function buildHeartbeatPayload(deps: HandlerDeps): SandboxHeartbeatPayloa ? 'finalizing' : 'active', activeKiloSessions: active.length, - pendingMessages: deps.tasks.size, + pendingMessages: deps.operations.counts().active, kilo: { ready: deps.kiloReady && !deps.signal?.aborted, ...(deps.kiloRuntimes?.kiloCliVersion !== undefined @@ -313,7 +321,28 @@ export async function refreshHeartbeatPayload( return buildHeartbeatPayload(deps); } -export function createControlHandlerDeps(deps: HandlerDeps): HandlerDeps { +export function createControlHandlerDeps(input: Omit): HandlerDeps { + const deps: HandlerDeps = Object.assign(input, { + operations: createOperationRegistry({ + native: { get: directory => input.kiloRuntimes?.get(directory) }, + onStarted: (session, preparation) => { + if (!preparation) deps.activity?.markActive(session.kiloSessionId); + const snapshot = deps.sessions.find(item => item.kiloSessionId === session.kiloSessionId); + if (snapshot) snapshot.lastActivityAt = Date.now(); + else + deps.sessions.push({ kiloSessionId: session.kiloSessionId, lastActivityAt: Date.now() }); + }, + onCompleted: session => { + deps.activity?.reconcile({}, [session.kiloSessionId]); + const snapshot = deps.sessions.find(item => item.kiloSessionId === session.kiloSessionId); + if (snapshot) { + snapshot.lastActivityAt = Date.now(); + delete snapshot.pendingInputs; + } + }, + retireRuntime: reason => deps.retireRuntime(reason), + }), + }); if (!deps.nativeObservations && deps.activity && deps.kiloRuntimes) { deps.nativeObservations = createNativeObservations({ get signal() { @@ -332,82 +361,13 @@ export function createControlHandlerDeps(deps: HandlerDeps): HandlerDeps { return deps; } -function startSessionTask( - session: SessionRequestIdentity, - identity: TaskIdentity, - deps: HandlerDeps, - run: (task: OwnedSessionTask) => Promise -): OwnedSessionTask { - const startedAt = Date.now(); - const diagnostic = (phase: string): void => - emitControlDiagnostic(deps.onDiagnostic, 'session.task', { - sessionId: session.sessionId, - kiloSessionId: session.kiloSessionId, - messageId: identity.messageId, - kind: identity.kind, - phase, - elapsedMs: Date.now() - startedAt, - }); - diagnostic('started'); - const completion = Promise.withResolvers(); - const controller = new AbortController(); - const task: OwnedSessionTask = { - ...identity, - session, - controller, - signal: deps.signal ? AbortSignal.any([controller.signal, deps.signal]) : controller.signal, - done: completion.promise, - }; - deps.tasks.set(session.kiloSessionId, task); - if (identity.kind !== 'preparation') deps.activity?.markActive(session.kiloSessionId); - const snapshot = deps.sessions.find(item => item.kiloSessionId === session.kiloSessionId); - if (snapshot) { - snapshot.lastActivityAt = Date.now(); - } else { - deps.sessions.push({ kiloSessionId: session.kiloSessionId, lastActivityAt: Date.now() }); - } - const timeout = setTimeout( - () => { - const reason = - identity.kind !== 'preparation' - ? 'Execution exceeded the 60 minute limit' - : 'Session preparation timed out'; - diagnostic('deadline_expired'); - controller.abort(new ControlTaskCancellation('failed', reason)); - deps.retireRuntime(reason); - }, - identity.kind !== 'preparation' - ? SANDBOX_CONTROL_EXECUTION_TIMEOUT_MS - : SANDBOX_CONTROL_ATTACH_TIMEOUT_MS - ); - timeout.unref(); - void Promise.resolve() - .then(() => run(task)) - .catch(kiloFailure) - .then(result => { - clearTimeout(timeout); - if (deps.tasks.get(session.kiloSessionId) === task) { - deps.tasks.delete(session.kiloSessionId); - deps.activity?.reconcile({}, [session.kiloSessionId]); - const snapshot = deps.sessions.find(item => item.kiloSessionId === session.kiloSessionId); - if (snapshot) { - snapshot.lastActivityAt = Date.now(); - delete snapshot.pendingInputs; - } - } - diagnostic(result.ok ? 'finished' : 'failed'); - completion.resolve(result); - }); - return task; -} - export async function cancelControlTasks( deps: HandlerDeps, reason: string, status: 'failed' | 'cancelled' = 'cancelled' ): Promise { - const tasks = [...deps.tasks.values()]; - for (const task of tasks) task.controller.abort(new ControlTaskCancellation(status, reason)); + const tasks = deps.operations.activeOperations(); + for (const task of tasks) task.cancel(reason, status); await Promise.all(tasks.map(task => task.done)); } @@ -415,7 +375,8 @@ export async function handleControlRequest( operation: string, session: SessionRequestIdentity | undefined, payload: unknown, - deps: HandlerDeps + deps: HandlerDeps, + authorization?: SessionOperationAuthorization ): Promise { if (operation === 'sandbox.status') { const heartbeat = buildHeartbeatPayload(deps); @@ -427,11 +388,13 @@ export async function handleControlRequest( }); } if (operation === 'sandbox.shutdown') { + const deadlineAt = Date.now() + KILO_CONTROL_REQUEST_TIMEOUT_MS; if (!sandboxShutdownPayloadSchema.safeParse(payload).success) { return fail('protocol_error', 'Invalid payload', false); } deps.onShutdown?.(); await cancelControlTasks(deps, 'Sandbox shutting down'); + await deps.operations.drainDelivery(deadlineAt); deps.terminalRuntime?.shutdown(); deps.kiloRuntimes?.shutdown(); return ok({ shuttingDown: true }); @@ -465,12 +428,10 @@ export async function handleControlRequest( const fenced = fenceDirectoryOperations(input.directory); diagnostic('started', 'deletion_fence'); failureStage = 'task_cancellation'; - const tasks = [...deps.tasks.values()].filter( - task => task.session.directory === input.directory - ); - for (const task of tasks) { - task.controller.abort(new ControlTaskCancellation('cancelled', 'Worktree deleted')); - } + const tasks = deps.operations + .activeOperations() + .filter(task => task.session.directory === input.directory); + for (const task of tasks) task.cancel('Worktree deleted', 'cancelled'); const results = await Promise.all(tasks.map(task => task.done)); failureStage = 'deletion_fence'; await fenced; @@ -518,6 +479,9 @@ export async function handleControlRequest( if (!session) { return fail('protocol_error', 'session identity is required', false); } + if (operation === 'session.operation.ack') return deps.operations.acknowledge(session, payload); + const admission = deps.operations.admission(operation, session, payload, authorization); + if (admission.kind === 'reply') return admission.result; if ( (deps.signal?.aborted || (!deps.kiloReady && operation !== 'session.git.summary')) && operation !== 'session.abort' && @@ -535,7 +499,7 @@ export async function handleControlRequest( errorCode: 'not_ready', retryable: true, aborted: deps.signal?.aborted ?? false, - ownedTask: deps.tasks.has(session.kiloSessionId), + ownedTask: deps.operations.hasActive(session.kiloSessionId), statusQueryPending: false, questionQueryPending: false, permissionQueryPending: false, @@ -552,7 +516,7 @@ export async function handleControlRequest( ) { return fail('unauthorized', 'Session directory mismatch', false); } - const current = deps.tasks.get(session.kiloSessionId); + const current = deps.operations.active(session.kiloSessionId); if ( current && (current.session.directory !== session.directory || @@ -563,7 +527,7 @@ export async function handleControlRequest( try { assertDirectoryActive(session.directory); - return await handleSessionControlRequest(operation, session, payload, deps); + return await handleSessionControlRequest(operation, session, payload, deps, authorization); } catch { return fail('not_ready', 'Worktree is being deleted', false); } @@ -573,15 +537,16 @@ async function handleSessionControlRequest( operation: string, session: SessionRequestIdentity, payload: unknown, - deps: HandlerDeps + deps: HandlerDeps, + authorization?: SessionOperationAuthorization ): Promise { switch (operation) { case 'session.attach': - return handleAttach(session, payload, deps); + return handleAttach(session, payload, deps, authorization); case 'session.detach': return handleDetach(session, payload, deps); case 'session.prompt': - return handlePrompt(session, payload, deps); + return handlePrompt(session, payload, deps, authorization); case 'session.abort': return handleAbort(session, payload, deps); case 'session.permission.resolve': @@ -659,49 +624,56 @@ function terminalFailure(error: unknown): ControlHandlerResult { async function handleAttach( session: SessionRequestIdentity, payload: unknown, - deps: HandlerDeps + deps: HandlerDeps, + authorization?: SessionOperationAuthorization ): Promise { const parsed = sessionAttachPayloadSchema.safeParse(payload ?? {}); - if (!parsed.success) return fail('protocol_error', 'Invalid payload', false); + if (!parsed.success) return rejectBeforeAdmission('protocol_error', 'Invalid payload', false); if ( parsed.data.env && CONTROL_RUNTIME_RESERVED_ENV_VARS.some(name => Object.hasOwn(parsed.data.env ?? {}, name)) ) { - return fail('protocol_error', 'Reserved control runtime environment variable', false); + return rejectBeforeAdmission( + 'protocol_error', + 'Reserved control runtime environment variable', + false + ); } - if (deps.tasks.has(session.kiloSessionId)) { - return fail('session_busy', 'Session has work in progress', true); + if (deps.operations.hasActive(session.kiloSessionId)) { + return rejectBeforeAdmission('session_busy', 'Session has work in progress', true); } - const task = startSessionTask( + const task = deps.operations.start( session, - { kind: 'preparation', messageId: parsed.data.preparation?.triggerMessageId }, - deps, - async owned => { - const result = await (deps.applyAttach ?? applySessionAttach)(session, parsed.data, { - onDiagnostic: deps.onDiagnostic, - kiloRuntimes: deps.kiloRuntimes, - signal: owned.signal, - canRefreshCredentials: () => - !owned.signal.aborted && - ![...deps.tasks.values()].some( - task => task !== owned && task.session.directory === session.directory - ) && - !(deps.activity?.snapshots() ?? []).some( - snapshot => - snapshot.state !== 'idle' && - directoryForSession(snapshot.kiloSessionId) === session.directory - ), - ...(deps.terminalRuntime ? { terminalRuntime: deps.terminalRuntime } : {}), - ...(deps.emitPreparing ? { emitPreparing: deps.emitPreparing } : {}), - }); - if (owned.signal.aborted) { - return fail('not_ready', 'Session attachment cancelled', true); - } - if (result.ok) deps.activity?.attach(session.kiloSessionId); - return result; - } + authorization, + { + operation: 'session.attach', + payload: parsed.data, + apply: (identity, payload, hooks) => + (deps.applyAttach ?? applySessionAttach)(identity, payload, { + ...hooks, + onDiagnostic: deps.onDiagnostic, + kiloRuntimes: deps.kiloRuntimes, + canRefreshCredentials: () => + !deps.operations + .activeOperations() + .some( + task => + task.session.directory === session.directory && + task.session.kiloSessionId !== session.kiloSessionId + ) && + !(deps.activity?.snapshots() ?? []).some( + snapshot => + snapshot.state !== 'idle' && + directoryForSession(snapshot.kiloSessionId) === session.directory + ), + ...(deps.terminalRuntime ? { terminalRuntime: deps.terminalRuntime } : {}), + }), + onAttached: () => deps.activity?.attach(session.kiloSessionId), + emitPreparing: deps.emitPreparing, + }, + operationEffects(session, deps) ); return task.done; } @@ -714,9 +686,9 @@ async function handleDetach( if (!sessionDetachPayloadSchema.safeParse(payload).success) { return fail('protocol_error', 'Invalid payload', false); } - const task = deps.tasks.get(session.kiloSessionId); + const task = deps.operations.active(session.kiloSessionId); if (task) { - task.controller.abort(new ControlTaskCancellation('cancelled', 'Session detached')); + task.cancel('Session detached', 'cancelled'); const result = await task.done; if (!result.ok && task.kind !== 'preparation') return result; } @@ -824,7 +796,8 @@ function validAttachmentPaths( function handlePrompt( session: SessionRequestIdentity, payload: unknown, - deps: HandlerDeps + deps: HandlerDeps, + authorization?: SessionOperationAuthorization ): ControlHandlerResult { const runtime = sessionKiloRuntime(session, deps); if (!runtime) { @@ -832,38 +805,42 @@ function handlePrompt( deps.kiloRuntimes?.isHealthy() && directoryForSession(session.kiloSessionId) === session.directory && rootForSession(session.kiloSessionId) === session.kiloSessionId && - [...deps.tasks.values()].some( - task => - task.kind === 'preparation' && - task.session.directory === session.directory && - !task.signal.aborted - ) + deps.operations + .activeOperations() + .some( + task => + task.kind === 'preparation' && + task.session.directory === session.directory && + !task.signal.aborted + ) ) { - return fail('session_busy', 'Worktree has preparation in progress', true); + return rejectBeforeAdmission('session_busy', 'Worktree has preparation in progress', true); } - return missingKilo(); + return rejectBeforeAdmission('not_ready', 'Kilo is not ready', true); } const parsed = sessionPromptPayloadSchema.safeParse(payload); - if (!parsed.success) return fail('protocol_error', 'Invalid payload', false); + if (!parsed.success) return rejectBeforeAdmission('protocol_error', 'Invalid payload', false); const request = parsed.data; + if (authorization && authorization.messageId !== request.messageId) + return rejectBeforeAdmission('idempotency_conflict', 'Message identity mismatch', false); if ( request.turn.type === 'command' && request.turn.command === 'compact' && !request.agent.model ) { - return fail('protocol_error', 'Model is required for compact', false); + return rejectBeforeAdmission('protocol_error', 'Model is required for compact', false); } if (!validAttachmentPaths(session, request)) { - return fail('protocol_error', 'Invalid attachment path', false); + return rejectBeforeAdmission('protocol_error', 'Invalid attachment path', false); } if (request.turn.type === 'command' && request.attachments?.length) { - return fail( + return rejectBeforeAdmission( 'protocol_error', 'Command attachments are not supported by the control runtime', false ); } - const existing = deps.tasks.get(session.kiloSessionId); + const existing = deps.operations.active(session.kiloSessionId); if (existing) { if ( existing.kind !== 'preparation' && @@ -872,16 +849,19 @@ function handlePrompt( ) { return ok({ messageId: request.messageId, status: 'existing' }); } - return fail('session_busy', 'Session has work in progress', true); + return rejectBeforeAdmission('session_busy', 'Session has work in progress', true); } - startSessionTask( + deps.operations.start( session, - { kind: 'execution', messageId: request.messageId }, + authorization, { - ...deps, - signal: deps.signal ? AbortSignal.any([deps.signal, runtime.signal]) : runtime.signal, + operation: 'session.prompt', + payload: request, + runtime, + materializeAttachments: deps.materializeAttachments, + runAutoCommit: deps.runAutoCommit, }, - task => executePrompt(task, request, runtime, deps) + operationEffects(session, deps) ); return ok({ messageId: request.messageId, status: 'accepted' }); } @@ -900,214 +880,6 @@ async function abortKiloSession( if (aborted !== true) throw new Error('Kilo cancellation was not confirmed'); } -function emitFinalizationEvent( - session: SessionRequestIdentity, - event: IngestEvent, - deps: HandlerDeps -): void { - deps.emitSessionEvent( - session, - sessionEventPayloadSchema.parse({ - type: event.streamEventType, - properties: event.data, - timestamp: event.timestamp, - }) - ); -} - -async function summarizeOwnedSession( - task: OwnedSessionTask, - kiloClient: WrapperKiloClient, - model: { providerID?: string; modelID: string }, - auto?: boolean -): Promise { - const success = await withTimeoutAndAbort( - kiloClient.summarizeSession({ - sessionId: task.session.kiloSessionId, - directory: task.session.directory, - signal: task.signal, - model, - ...(auto === undefined ? {} : { auto }), - }), - { - signal: task.signal, - timeoutMs: SANDBOX_CONTROL_EXECUTION_TIMEOUT_MS, - timeoutMessage: 'Execution exceeded the 60 minute limit', - abortMessage: 'Execution cancelled', - } - ); - if (!success) throw new Error('Session summarization failed'); -} - -async function executePrompt( - task: OwnedSessionTask, - request: SessionPromptPayload, - runtime: WorktreeKiloRuntime, - deps: HandlerDeps -): Promise { - const { session, signal } = task; - const { kiloClient, env } = runtime; - const { messageId, turn, agent } = request; - const startedAt = Date.now(); - const diagnostic = (phase: string, status?: SessionMessageOutcome['status']): void => - emitControlDiagnostic(deps.onDiagnostic, 'session.execution', { - sessionId: session.sessionId, - kiloSessionId: session.kiloSessionId, - messageId, - phase, - status, - elapsedMs: Date.now() - startedAt, - aborted: signal.aborted, - }); - let outcome: SessionMessageOutcome; - let result = ok({}); - let failureReason = 'Kilo execution failed'; - const emitStatus = (message: string): void => - emitFinalizationEvent( - session, - { - streamEventType: 'status', - data: { message, messageId }, - timestamp: new Date().toISOString(), - }, - deps - ); - try { - signal.throwIfAborted(); - let completion: Awaited> | undefined; - const options = { - sessionId: session.kiloSessionId, - directory: session.directory, - signal, - messageId, - agent: agent.mode, - ...(agent.variant ? { variant: agent.variant } : {}), - }; - const deadline = { - signal, - timeoutMs: SANDBOX_CONTROL_EXECUTION_TIMEOUT_MS, - timeoutMessage: 'Execution exceeded the 60 minute limit', - abortMessage: 'Execution cancelled', - }; - if (turn.type === 'prompt') { - if (agent.model === undefined) throw new Error('Prompt model is required'); - const message = await (deps.materializeAttachments ?? materializeMessageAttachments)( - { id: messageId, prompt: turn.prompt, parts: turn.parts, attachments: request.attachments }, - { signal } - ); - signal.throwIfAborted(); - diagnostic('prompt_started'); - completion = await withTimeoutAndAbort( - kiloClient.sendPrompt({ - ...options, - prompt: message.prompt, - ...(message.parts ? { parts: message.parts } : {}), - model: { providerID: 'kilo', modelID: agent.model }, - }), - deadline - ); - diagnostic('prompt_completed'); - } else if (turn.command === 'compact') { - if (!agent.model) throw new Error('Model is required for compact'); - failureReason = 'Context condensation failed'; - emitStatus('Condensing context...'); - diagnostic('compact_started'); - await summarizeOwnedSession(task, kiloClient, { providerID: 'kilo', modelID: agent.model }); - diagnostic('compact_completed'); - signal.throwIfAborted(); - emitStatus('Context condensed successfully'); - } else { - diagnostic('command_started'); - completion = await withTimeoutAndAbort( - kiloClient.sendCommand({ - ...options, - command: turn.command, - args: turn.arguments, - ...(agent.model !== undefined - ? { model: { providerID: 'kilo', modelID: agent.model } } - : {}), - }), - deadline - ); - diagnostic('command_completed'); - } - signal.throwIfAborted(); - const error = completion?.info.error; - if (!error && (request.finalization?.autoCommit || request.finalization?.condenseOnComplete)) { - task.kind = 'finalizing'; - diagnostic('finalization_started'); - if (request.finalization.autoCommit) { - failureReason = 'Auto-commit failed'; - diagnostic('autocommit_started'); - const committed = await (deps.runAutoCommit ?? runAutoCommit)({ - workspacePath: session.directory, - kiloClient, - env, - messageId: completion?.info.id ?? messageId, - signal, - onEvent: event => emitFinalizationEvent(session, event, deps), - }); - signal.throwIfAborted(); - if (!committed.success) throw new Error('Auto-commit failed'); - diagnostic('autocommit_completed'); - } - if (request.finalization.condenseOnComplete) { - failureReason = 'Context condensation failed'; - const model = agent.model - ? { providerID: 'kilo', modelID: agent.model } - : completion - ? { providerID: completion.info.providerID, modelID: completion.info.modelID } - : undefined; - if (!model) throw new Error('Model is required for condensation'); - emitStatus('Condensing context...'); - diagnostic('condense_started'); - await summarizeOwnedSession(task, kiloClient, model, true); - diagnostic('condense_completed'); - signal.throwIfAborted(); - emitStatus('Context condensed successfully'); - } - } - outcome = error - ? { - messageId, - status: error.name === 'MessageAbortedError' ? 'cancelled' : 'failed', - reason: `Kilo execution ended with ${error.name}`, - } - : { messageId, status: 'completed' }; - } catch { - diagnostic('execution_failed'); - const cancellation: unknown = signal.reason; - outcome = { - messageId, - status: cancellation instanceof ControlTaskCancellation ? cancellation.status : 'failed', - reason: - cancellation instanceof ControlTaskCancellation ? cancellation.message : failureReason, - }; - try { - diagnostic('abort_started'); - await abortKiloSession(session, kiloClient); - diagnostic('abort_completed'); - } catch (error) { - diagnostic('abort_failed'); - deps.retireRuntime('Kilo cancellation failed'); - result = kiloFailure(error); - } - } - try { - diagnostic('outcome_sending', outcome.status); - deps.emitSessionEvent(session, { - type: 'session.message.outcome', - properties: sessionMessageOutcomeSchema.parse(outcome), - }); - diagnostic('outcome_sent', outcome.status); - } catch { - diagnostic('outcome_failed', outcome.status); - deps.retireRuntime('Session outcome delivery failed'); - return fail('not_ready', 'Session outcome delivery failed', false); - } - return result; -} - async function handleAbort( session: SessionRequestIdentity, payload: unknown, @@ -1115,12 +887,12 @@ async function handleAbort( ): Promise { const parsed = sessionAbortPayloadSchema.safeParse(payload ?? {}); if (!parsed.success) return fail('protocol_error', 'Invalid payload', false); - const task = deps.tasks.get(session.kiloSessionId); + const task = deps.operations.active(session.kiloSessionId); if (parsed.data.messageId && task?.messageId !== parsed.data.messageId) { return ok({ status: 'already_idle' }); } if (task) { - task.controller.abort(new ControlTaskCancellation('cancelled', 'Session aborted')); + task.cancel('Session aborted', 'cancelled'); const result = await task.done; if (!result.ok && task.kind !== 'preparation') return result; return ok({ status: 'aborted' }); @@ -1258,7 +1030,7 @@ async function handleSync( elapsedMs: Math.max(0, Date.now() - startedAt), ok: phase === 'completed', aborted: deps.signal?.aborted ?? false, - ownedTask: deps.tasks.has(session.kiloSessionId), + ownedTask: deps.operations.hasActive(session.kiloSessionId), statusQueryPending: queries.sync_status, questionQueryPending: queries.sync_questions, permissionQueryPending: queries.sync_permissions, @@ -1350,7 +1122,7 @@ async function handleSync( if (!questions.complete || !permissions.complete) { throw new Error('Native requests contain unresolved ancestry'); } - const ownedTask = deps.tasks.has(session.kiloSessionId); + const ownedTask = deps.operations.hasActive(session.kiloSessionId); const status = ownedTask ? { type: 'busy' } : (statuses[session.kiloSessionId] ?? { type: 'idle' }); diff --git a/services/cloud-agent-next/wrapper/src/control/session-operation-cleanup.test.ts b/services/cloud-agent-next/wrapper/src/control/session-operation-cleanup.test.ts new file mode 100644 index 0000000000..2033cfced7 --- /dev/null +++ b/services/cloud-agent-next/wrapper/src/control/session-operation-cleanup.test.ts @@ -0,0 +1,101 @@ +import { afterEach, beforeEach, describe, expect, it, jest } from 'bun:test'; +import fs from 'node:fs'; +import os from 'node:os'; +import path from 'node:path'; +import { handleControlRequest, type HandlerDeps } from './sandbox-control-handlers'; +import { KILO_CONTROL_REQUEST_TIMEOUT_MS } from './sandbox-control-runtime'; +import { + acknowledgeOperation, + completion, + createHandlerFixture, + fakeKilo, + operationAuthorization, + promptPayload, + session, +} from './control-test-fixtures'; +import { rememberAttachedRoot, resetSessionDirectoryState } from './session-directories'; +import { resetDirectoryOperationState } from './worktree-operations'; + +let homeRoot: string; + +beforeEach(() => { + resetSessionDirectoryState(); + resetDirectoryOperationState(); + rememberAttachedRoot(session.kiloSessionId, session.directory); + homeRoot = fs.mkdtempSync(path.join(os.tmpdir(), 'session-operation-cleanup-test-')); +}); + +afterEach(() => { + jest.useRealTimers(); + fs.rmSync(homeRoot, { recursive: true, force: true }); +}); + +function deps(overrides: Parameters[1] = {}): HandlerDeps { + return createHandlerFixture(homeRoot, overrides); +} + +describe('SessionOperation cleanup', () => { + it('keeps an unconfirmed native receipt after its original cleanup allowance expires', async () => { + jest.useFakeTimers(); + const started = Promise.withResolvers(); + const original = Promise.withResolvers>(); + const abortStarted = Promise.withResolvers(); + let submissions = 0; + const handlerDeps = deps({ + kiloClient: fakeKilo({ + sendPrompt: () => { + submissions++; + started.resolve(); + return original.promise; + }, + abortSession: async () => { + abortStarted.resolve(); + return true; + }, + }), + sendOperationResult: (_session, delivery) => acknowledgeOperation(delivery), + }); + const authorization = operationAuthorization(); + await handleControlRequest( + 'session.prompt', + session, + promptPayload, + handlerDeps, + authorization + ); + await started.promise; + const record = handlerDeps.operations.retained()[0]; + if (!record) throw new Error('Missing operation record'); + record.cancel('Session aborted', 'cancelled'); + await abortStarted.promise; + await Promise.resolve(); + await Promise.resolve(); + jest.advanceTimersByTime(KILO_CONTROL_REQUEST_TIMEOUT_MS); + await Promise.resolve(); + jest.advanceTimersByTime(0); + await record.done; + await record.waitForDelivery(); + const sealed = record.snapshot(); + + expect(sealed.native.state).toBe('unknown'); + expect(sealed.delivery?.state).toBe('acknowledged'); + handlerDeps.operations.prune(Number.MAX_SAFE_INTEGER); + expect(handlerDeps.operations.retained()).toEqual([record]); + expect( + await handleControlRequest( + 'session.prompt', + session, + promptPayload, + handlerDeps, + authorization + ) + ).toEqual({ ok: true, result: { messageId: 'msg_1', status: 'existing' } }); + expect(submissions).toBe(1); + + original.resolve(completion()); + await Promise.resolve(); + expect(record.snapshot().native).toEqual(sealed.native); + expect(record.snapshot().local).toEqual(sealed.local); + expect(record.deliveryResult()).toEqual(sealed.delivery?.payload); + }); +}); diff --git a/services/cloud-agent-next/wrapper/src/control/session-operation.test.ts b/services/cloud-agent-next/wrapper/src/control/session-operation.test.ts new file mode 100644 index 0000000000..41cc8653e9 --- /dev/null +++ b/services/cloud-agent-next/wrapper/src/control/session-operation.test.ts @@ -0,0 +1,714 @@ +import { afterEach, beforeEach, describe, expect, it, setSystemTime, spyOn } from 'bun:test'; +import fs from 'node:fs'; +import os from 'node:os'; +import path from 'node:path'; +import { + SANDBOX_CONTROL_EXECUTION_TIMEOUT_MS, + sessionOperationDeliverySchema, + type SessionEventPayload, + type SessionOperationAck, + type SessionOperationDelivery, +} from '../../../src/shared/sandbox-control-protocol'; +import type { AutoCommitResult } from '../auto-commit'; +import { + buildHeartbeatPayload, + handleControlRequest, + type HandlerDeps, +} from './sandbox-control-handlers'; +import { + acknowledgeOperation, + completion, + createHandlerFixture, + fakeKilo, + operationAuthorization, + promptPayload, + session, + type Completion, +} from './control-test-fixtures'; +import { rememberAttachedRoot, resetSessionDirectoryState } from './session-directories'; +import { resetDirectoryOperationState } from './worktree-operations'; + +let homeRoot: string; + +beforeEach(() => { + resetSessionDirectoryState(); + resetDirectoryOperationState(); + rememberAttachedRoot(session.kiloSessionId, session.directory); + homeRoot = fs.mkdtempSync(path.join(os.tmpdir(), 'session-operation-test-')); +}); + +afterEach(() => { + setSystemTime(); + fs.rmSync(homeRoot, { recursive: true, force: true }); +}); + +function deps(overrides: Parameters[1] = {}): HandlerDeps { + return createHandlerFixture(homeRoot, overrides); +} + +function onlyOperation(handlerDeps: HandlerDeps) { + const records = handlerDeps.operations.retained(); + expect(records).toHaveLength(1); + const record = records[0]; + if (!record) throw new Error('Missing operation record'); + return record; +} + +describe('operation results and delivery', () => { + it.each([{ data: undefined }, { data: null }, { data: [] }, { data: 'invalid' }, { data: 1 }])( + 'rejects malformed finalization notification data without changing the producer result: %j', + async ({ data }) => { + const notifications: SessionEventPayload[] = []; + const handlerDeps = deps({ + runAutoCommit: async options => { + options.onEvent({ + streamEventType: 'autocommit_completed', + data, + timestamp: new Date().toISOString(), + }); + return { success: true }; + }, + emitSessionEvent: (_session, event) => notifications.push(event), + sendOperationResult: (_session, delivery) => acknowledgeOperation(delivery), + }); + await handleControlRequest( + 'session.prompt', + session, + { ...promptPayload, finalization: { autoCommit: true } }, + handlerDeps, + operationAuthorization() + ); + const record = onlyOperation(handlerDeps); + await record.done; + await record.waitForDelivery(); + expect(record.snapshot().finalization.autoCommit).toEqual({ + state: 'completed', + result: { success: true }, + }); + expect(record.snapshot().outcome?.status).toBe('completed'); + expect(record.snapshot().events).toEqual([]); + expect(record.snapshot().delivery?.payload.events).toEqual([]); + expect(record.snapshot().delivery?.state).toBe('acknowledged'); + expect(notifications.filter(event => event.type === 'autocommit_completed')).toEqual([]); + } + ); + + it('does not retire the wrapper when aborted native work rejects after confirmed cancellation', async () => { + const prompt = Promise.withResolvers(); + let retired: string | undefined; + const handlerDeps = deps({ + kiloClient: fakeKilo({ + sendPrompt: () => prompt.promise, + abortSession: async () => true, + }), + retireRuntime: reason => { + retired = reason; + }, + sendOperationResult: (_session, delivery) => acknowledgeOperation(delivery), + }); + rememberAttachedRoot(session.kiloSessionId, session.directory); + await handleControlRequest( + 'session.prompt', + session, + promptPayload, + handlerDeps, + operationAuthorization() + ); + const record = onlyOperation(handlerDeps); + record.cancel('User stop', 'cancelled'); + const aborted = Object.assign(new Error('aborted'), { name: 'MessageAbortedError' }); + prompt.reject(aborted); + await record.done; + await record.waitForDelivery(); + expect(retired).toBeUndefined(); + expect(record.snapshot().outcome?.status).toBe('cancelled'); + }); + + it('aborts an already-awaiting finalizer and retains its late original result without false quiescence', async () => { + const entered = Promise.withResolvers(); + const finalized = Promise.withResolvers(); + let finalizerSignal: AbortSignal | undefined; + const handlerDeps = deps({ + runAutoCommit: async options => { + finalizerSignal = options.signal; + entered.resolve(); + const result = await finalized.promise; + options.onEvent({ + streamEventType: 'autocommit_completed', + data: { + success: result.success, + commitHash: 'original-commit', + messageId: options.messageId, + }, + timestamp: new Date().toISOString(), + }); + return result; + }, + sendOperationResult: (_session, delivery) => acknowledgeOperation(delivery), + }); + rememberAttachedRoot(session.kiloSessionId, session.directory); + const authorization = operationAuthorization(); + await handleControlRequest( + 'session.prompt', + session, + { ...promptPayload, finalization: { autoCommit: true } }, + handlerDeps, + authorization + ); + await entered.promise; + const record = onlyOperation(handlerDeps); + record.cancel('Late work cancellation', 'cancelled'); + expect(finalizerSignal?.aborted).toBe(true); + expect(record.snapshot().local).toBeUndefined(); + expect(handlerDeps.operations.active(session.kiloSessionId)).toBe(record); + finalized.resolve({ success: true }); + await record.done; + await record.waitForDelivery(); + expect(record.snapshot().finalization.autoCommit).toEqual({ + state: 'completed', + result: { success: true }, + }); + expect(record.snapshot().events).toEqual([ + { + type: 'autocommit_completed', + properties: { success: true, commitHash: 'original-commit', messageId: 'assistant_1' }, + timestamp: expect.any(String), + }, + ]); + expect(record.snapshot().native.completion).toEqual(completion().info); + expect(record.snapshot().outcome?.status).toBe('completed'); + expect(handlerDeps.operations.counts().active).toBe(0); + }); + + it('releases execution before ACK and does not let Stop or an old execution timer rewrite completion', async () => { + const acknowledgement = Promise.withResolvers(); + const sending = Promise.withResolvers(); + const handlerDeps = deps({ + sendOperationResult: (_session, delivery) => { + sending.resolve(delivery); + return acknowledgement.promise; + }, + }); + rememberAttachedRoot(session.kiloSessionId, session.directory); + const authorization = operationAuthorization(); + const timers = spyOn(globalThis, 'setTimeout'); + const cleared = spyOn(globalThis, 'clearTimeout'); + try { + await handleControlRequest( + 'session.prompt', + session, + promptPayload, + handlerDeps, + authorization + ); + const record = onlyOperation(handlerDeps); + await record.done; + const delivery = await sending.promise; + const original = structuredClone(record.snapshot().local); + expect(cleared).toHaveBeenCalledWith(timers.mock.results[0]?.value); + expect(buildHeartbeatPayload(handlerDeps)).toMatchObject({ + state: 'idle', + pendingMessages: 0, + }); + setSystemTime(Date.now() + SANDBOX_CONTROL_EXECUTION_TIMEOUT_MS + 1); + expect(record.snapshot().delivery?.payload).toEqual(delivery); + acknowledgement.resolve(await acknowledgeOperation(delivery)); + await record.waitForDelivery(); + expect(record.snapshot().local).toEqual(original); + expect(record.snapshot().outcome).toEqual({ messageId: 'msg_1', status: 'completed' }); + expect(handlerDeps.operations.counts().active).toBe(0); + } finally { + for (const record of handlerDeps.operations.retained()) { + const delivery = record.deliveryResult(); + if (delivery) { + acknowledgement.resolve(await acknowledgeOperation(delivery)); + await record.waitForDelivery(); + } + } + timers.mockRestore(); + cleared.mockRestore(); + } + }); + + it('drains a retained result before completing sandbox shutdown', async () => { + const acknowledgement = Promise.withResolvers(); + const sending = Promise.withResolvers(); + const handlerDeps = deps({ + sendOperationResult: (_session, delivery) => { + sending.resolve(delivery); + return acknowledgement.promise; + }, + }); + const authorization = operationAuthorization(); + await handleControlRequest( + 'session.prompt', + session, + promptPayload, + handlerDeps, + authorization + ); + const record = onlyOperation(handlerDeps); + await record.done; + const delivery = await sending.promise; + let settled = false; + const shutdown = handleControlRequest('sandbox.shutdown', undefined, {}, handlerDeps).then( + result => { + settled = true; + return result; + } + ); + + await Promise.resolve(); + expect(settled).toBe(false); + acknowledgement.resolve(await acknowledgeOperation(delivery)); + expect(await shutdown).toEqual({ ok: true, result: { shuttingDown: true } }); + await record.waitForDelivery(); + }); + + it('retires held delivery only for an exact lookup-driven durable acknowledgement', async () => { + const held = Promise.withResolvers(); + const sending = Promise.withResolvers(); + const handlerDeps = deps({ + sendOperationResult: (_session, delivery) => { + sending.resolve(delivery); + return held.promise; + }, + }); + rememberAttachedRoot(session.kiloSessionId, session.directory); + const authorization = operationAuthorization(); + await handleControlRequest( + 'session.prompt', + session, + promptPayload, + handlerDeps, + authorization + ); + const record = onlyOperation(handlerDeps); + await record.done; + const delivery = await sending.promise; + const ack = await acknowledgeOperation(delivery); + try { + expect( + await handleControlRequest( + 'session.operation.ack', + session, + { ...ack, resultHash: '0'.repeat(64) }, + handlerDeps + ) + ).toMatchObject({ ok: false }); + expect( + await handleControlRequest( + 'session.operation.ack', + session, + { ...ack, authorization: { ...authorization, wrapperInstanceId: crypto.randomUUID() } }, + handlerDeps + ) + ).toMatchObject({ ok: false }); + expect(record.snapshot().delivery?.state).toBe('pending'); + record.cancel('Late work cancellation', 'cancelled'); + expect( + await handleControlRequest('session.operation.ack', session, ack, handlerDeps) + ).toEqual({ ok: true, result: { acknowledged: true } }); + await record.waitForDelivery(); + expect(record.snapshot().delivery?.acknowledgement).toEqual(ack); + expect(record.snapshot().delivery?.state).toBe('acknowledged'); + expect(record.snapshot().outcome?.status).toBe('completed'); + expect(handlerDeps.operations.counts().active).toBe(0); + } finally { + held.resolve(ack); + } + }); + + it.each(['prompt', 'command', 'compact'] as const)( + 'does not replay an unknown native %s outcome', + async kind => { + let submissions = 0; + let commits = 0; + const unavailable = async () => { + submissions++; + throw new Error('Mutation response unavailable'); + }; + const handlerDeps = deps({ + kiloClient: fakeKilo({ + sendPrompt: unavailable, + sendCommand: unavailable, + summarizeSession: unavailable, + }), + runAutoCommit: async () => { + commits++; + return { success: true }; + }, + sendOperationResult: (_session, delivery) => acknowledgeOperation(delivery), + }); + rememberAttachedRoot(session.kiloSessionId, session.directory); + const authorization = operationAuthorization(); + const payload = { + ...promptPayload, + turn: + kind === 'prompt' + ? promptPayload.turn + : { + type: 'command', + command: kind === 'compact' ? 'compact' : 'review', + arguments: '', + }, + finalization: { autoCommit: true }, + }; + await handleControlRequest('session.prompt', session, payload, handlerDeps, authorization); + const record = onlyOperation(handlerDeps); + await record.done; + await record.waitForDelivery(); + expect(record.snapshot().native.state).toBe('unknown'); + expect(record.snapshot().outcome).toMatchObject({ + status: 'failed', + reason: 'Kilo execution outcome is unconfirmed', + }); + await handleControlRequest('session.prompt', session, payload, handlerDeps, authorization); + expect(submissions).toBe(1); + expect(commits).toBe(0); + } + ); + + it('stops work after a cancelled wait but gives retained delivery a separate lifetime', async () => { + const materialized = Promise.withResolvers(); + let submissions = 0; + let deliveryWasCancelled: boolean | undefined; + const handlerDeps = deps({ + materializeAttachments: async message => { + await materialized.promise; + return message; + }, + kiloClient: fakeKilo({ + sendPrompt: async () => { + submissions++; + return completion(); + }, + }), + sendOperationResult: (_session, delivery, signal) => { + deliveryWasCancelled = signal.aborted; + return acknowledgeOperation(delivery); + }, + }); + rememberAttachedRoot(session.kiloSessionId, session.directory); + const authorization = operationAuthorization(); + await handleControlRequest( + 'session.prompt', + session, + promptPayload, + handlerDeps, + authorization + ); + const record = onlyOperation(handlerDeps); + record.cancel('Session aborted', 'cancelled'); + materialized.resolve(); + await record.done; + await record.waitForDelivery(); + expect(record.signal.aborted).toBe(true); + expect(deliveryWasCancelled).toBe(false); + expect(record.snapshot().outcome?.status).toBe('cancelled'); + expect(submissions).toBe(0); + }); + + it('normalizes a late native MessageAbortedError to cancelled before sealing', async () => { + const entered = Promise.withResolvers(); + const aborted = Promise.withResolvers(); + const handlerDeps = deps({ + kiloClient: fakeKilo({ + sendPrompt: async () => { + entered.resolve(); + await aborted.promise; + return completion({ name: 'MessageAbortedError', data: { message: 'User aborted' } }); + }, + }), + sendOperationResult: (_session, delivery) => acknowledgeOperation(delivery), + }); + rememberAttachedRoot(session.kiloSessionId, session.directory); + const authorization = operationAuthorization(); + await handleControlRequest( + 'session.prompt', + session, + promptPayload, + handlerDeps, + authorization + ); + await entered.promise; + const record = onlyOperation(handlerDeps); + record.cancel('Session aborted', 'cancelled'); + aborted.resolve(); + await record.done; + await record.waitForDelivery(); + expect(record.snapshot().outcome).toMatchObject({ + status: 'cancelled', + reason: 'Kilo execution ended with MessageAbortedError', + }); + expect(record.snapshot().native.completion?.error?.name).toBe('MessageAbortedError'); + expect(record.snapshot().delivery?.state).toBe('acknowledged'); + }); + + it.each([ + [undefined, 'completed'], + [{ name: 'MessageAbortedError', data: { message: 'cancelled' } }, 'cancelled'], + ] as const)( + 'waits for a native %s result that arrives after the abort acknowledgement', + async (nativeError, status) => { + const started = Promise.withResolvers(); + const original = Promise.withResolvers>(); + const abortAcknowledged = Promise.withResolvers(); + const handlerDeps = deps({ + kiloClient: fakeKilo({ + sendPrompt: () => { + started.resolve(); + return original.promise; + }, + abortSession: async () => { + abortAcknowledged.resolve(); + setTimeout(() => original.resolve(completion(nativeError)), 125); + return true; + }, + }), + sendOperationResult: (_session, delivery) => acknowledgeOperation(delivery), + }); + const authorization = operationAuthorization(); + await handleControlRequest( + 'session.prompt', + session, + promptPayload, + handlerDeps, + authorization + ); + await started.promise; + const record = onlyOperation(handlerDeps); + const aborting = handleControlRequest( + 'session.abort', + session, + { messageId: 'msg_1' }, + handlerDeps + ); + await abortAcknowledged.promise; + expect(record.locallyComplete).toBe(false); + expect(await aborting).toEqual({ ok: true, result: { status: 'aborted' } }); + await record.waitForDelivery(); + + expect(record.snapshot().outcome?.status).toBe(status); + expect(record.snapshot().delivery?.state).toBe('acknowledged'); + } + ); + + it.each([ + ['false', () => false], + [ + 'throws', + () => { + throw new Error('live event transport unavailable'); + }, + ], + ])('retains a bounded completion when live publication %s', async (_name, emitSessionEvent) => { + let finalizations = 0; + const handlerDeps = deps({ + runAutoCommit: async options => { + finalizations++; + for (let index = 0; index < 8; index++) { + options.onEvent({ + streamEventType: 'status', + data: { message: `Optional status ${index}`, messageId: options.messageId }, + timestamp: new Date().toISOString(), + }); + } + options.onEvent({ + streamEventType: 'autocommit_completed', + data: { + success: true, + messageId: options.messageId, + commitHash: 'commit_123', + message: 'push failed '.repeat(100_000), + commitMessage: 'subject '.repeat(100_000), + ignoredMetadata: { tooLarge: 'metadata '.repeat(100_000) }, + }, + timestamp: new Date().toISOString(), + }); + return { success: true }; + }, + emitSessionEvent: () => emitSessionEvent(), + sendOperationResult: (_session, delivery) => acknowledgeOperation(delivery), + }); + const authorization = operationAuthorization(); + await handleControlRequest( + 'session.prompt', + session, + { ...promptPayload, finalization: { autoCommit: true } }, + handlerDeps, + authorization + ); + const record = onlyOperation(handlerDeps); + await record.done; + await record.waitForDelivery(); + const delivery = record.deliveryResult(); + if (!delivery) throw new Error('Missing retained completion'); + const completionEvent = record + .snapshot() + .events.find(event => event.type === 'autocommit_completed'); + + expect(record.snapshot().outcome?.status).toBe('completed'); + expect(finalizations).toBe(1); + expect(completionEvent).toMatchObject({ + properties: { + success: true, + messageId: 'assistant_1', + commitHash: 'commit_123', + }, + }); + expect(String(completionEvent?.properties.message).length).toBeLessThanOrEqual(4_096); + expect(String(completionEvent?.properties.commitMessage).length).toBeLessThanOrEqual(4_096); + expect(completionEvent?.properties).not.toHaveProperty('ignoredMetadata'); + expect(sessionOperationDeliverySchema.parse(delivery)).toEqual(delivery); + + expect( + await handleControlRequest( + 'session.prompt', + session, + { ...promptPayload, finalization: { autoCommit: true } }, + handlerDeps, + authorization + ) + ).toEqual({ ok: true, result: { messageId: 'msg_1', status: 'existing' } }); + expect(finalizations).toBe(1); + }); + + it('publishes auto-commit progress while retaining its completion', async () => { + const events: SessionEventPayload[] = []; + let finalizations = 0; + const handlerDeps = deps({ + runAutoCommit: async options => { + finalizations++; + options.onEvent({ + streamEventType: 'autocommit_started', + data: { message: 'Committing changes', messageId: options.messageId }, + timestamp: new Date().toISOString(), + }); + options.onEvent({ + streamEventType: 'autocommit_completed', + data: { + success: true, + message: 'Changes committed', + messageId: options.messageId, + commitHash: 'commit_123', + }, + timestamp: new Date().toISOString(), + }); + return { success: true }; + }, + emitSessionEvent: (_session, event) => events.push(event), + sendOperationResult: (_session, delivery) => acknowledgeOperation(delivery), + }); + await handleControlRequest( + 'session.prompt', + session, + { ...promptPayload, finalization: { autoCommit: true } }, + handlerDeps, + operationAuthorization() + ); + const record = onlyOperation(handlerDeps); + await record.done; + await record.waitForDelivery(); + const delivery = record.deliveryResult(); + if (!delivery) throw new Error('Missing retained completion'); + + expect(events.map(event => event.type)).toEqual(['autocommit_started', 'autocommit_completed']); + expect(record.snapshot().events).toEqual([ + expect.objectContaining({ + type: 'autocommit_completed', + properties: expect.objectContaining({ commitHash: 'commit_123' }), + }), + ]); + expect(record.snapshot().outcome?.status).toBe('completed'); + expect(sessionOperationDeliverySchema.parse(delivery)).toEqual(delivery); + expect(finalizations).toBe(1); + }); + + it('retains a terminal preparation action after progress fills the optional slots', async () => { + const handlerDeps = deps({ + applyAttach: async (_session, _payload, hooks) => { + for (let revision = 0; revision < 64; revision++) { + hooks.emitPreparing?.({ + version: 2, + attemptId: 'prepare_msg_1', + triggerMessageId: 'msg_1', + revision, + timestamp: revision, + step: 'workspace_setup', + message: `Preparing ${revision}`, + action: 'step_started', + stepId: `step_${revision}`, + kind: 'phase', + label: 'Setup', + }); + } + hooks.emitPreparing?.({ + version: 2, + attemptId: 'prepare_msg_1', + triggerMessageId: 'msg_1', + revision: 64, + timestamp: 64, + step: 'workspace_setup', + message: 'Preparation completed', + action: 'attempt_completed', + }); + return { ok: true, result: { attached: true } }; + }, + sendOperationResult: (_session, delivery) => acknowledgeOperation(delivery), + }); + const authorization = operationAuthorization('session.attach'); + await handleControlRequest('session.attach', session, {}, handlerDeps, authorization); + const record = onlyOperation(handlerDeps); + await record.done; + await record.waitForDelivery(); + const delivery = record.deliveryResult(); + if (!delivery) throw new Error('Missing retained preparation'); + + expect(record.snapshot().preparing).toHaveLength(57); + expect(record.snapshot().preparing).toContainEqual( + expect.objectContaining({ action: 'attempt_completed', message: 'Preparation completed' }) + ); + expect(sessionOperationDeliverySchema.parse(delivery)).toEqual(delivery); + }); + + it('retains a bounded failed completion without changing its failed outcome', async () => { + const handlerDeps = deps({ + runAutoCommit: async options => { + options.onEvent({ + streamEventType: 'autocommit_completed', + data: { + success: false, + messageId: options.messageId, + commitHash: 'commit_123', + message: 'git push failed '.repeat(100_000), + }, + timestamp: new Date().toISOString(), + }); + return { success: false, error: 'git push failed' }; + }, + sendOperationResult: (_session, delivery) => acknowledgeOperation(delivery), + }); + await handleControlRequest( + 'session.prompt', + session, + { ...promptPayload, finalization: { autoCommit: true } }, + handlerDeps, + operationAuthorization() + ); + const record = onlyOperation(handlerDeps); + await record.done; + await record.waitForDelivery(); + + expect(record.snapshot().outcome).toMatchObject({ + status: 'failed', + reason: 'Auto-commit failed', + }); + expect(record.snapshot().events).toContainEqual( + expect.objectContaining({ + type: 'autocommit_completed', + properties: expect.objectContaining({ success: false, commitHash: 'commit_123' }), + }) + ); + }); +}); diff --git a/services/cloud-agent-next/wrapper/src/control/session-operation.ts b/services/cloud-agent-next/wrapper/src/control/session-operation.ts new file mode 100644 index 0000000000..8adbd530a7 --- /dev/null +++ b/services/cloud-agent-next/wrapper/src/control/session-operation.ts @@ -0,0 +1,717 @@ +import { isDeepStrictEqual } from 'node:util'; +import { + emitControlDiagnostic, + type ControlDiagnosticReporter, +} from '../../../src/shared/control-diagnostics.js'; +import { + SANDBOX_CONTROL_ATTACH_TIMEOUT_MS, + SANDBOX_CONTROL_EXECUTION_TIMEOUT_MS, + SANDBOX_CONTROL_OUTCOME_TIMEOUT_MS, + sessionOperationExpiresAt, + sameSessionOperation, + sessionMessageOutcomeSchema, + sessionEventPayloadSchema, + sessionOperationDeliverySchema, + type SessionOperationAuthorization, + type SessionOperationDelivery, + type SessionOperationAck, + type SessionAttachPayload, + type SessionEventPayload, + type SessionMessageOutcome, + type SessionPromptPayload, + type SessionRequestIdentity, +} from '../../../src/shared/sandbox-control-protocol.js'; +import type { IngestEvent } from '../../../src/shared/protocol.js'; +import { isKiloServerUnreachableError, type WrapperKiloClient } from '../kilo-api.js'; +import { materializeMessageAttachments } from '../session-bootstrap.js'; +import { runAutoCommit, type AutoCommitResult } from '../auto-commit.js'; +import { withTimeoutAndAbort } from '../utils.js'; +import type { AttachPreparingEmitter } from './apply-attach.js'; +import { KILO_CONTROL_REQUEST_TIMEOUT_MS } from './sandbox-control-runtime.js'; +import type { WorktreeKiloRuntime } from './worktree-runtime.js'; +import { operationIntent } from './operation-intent.js'; +import { + createOperationResultDelivery, + type OperationResultDelivery, + type OperationResultSender, +} from './operation-result-delivery.js'; +import { + createRetainedOperationNotifications, + isRetainedOperationPreparing, +} from './retained-operation-notifications.js'; + +import type { ControlHandlerResult } from './control-handler-result.js'; +export type { ControlHandlerResult } from './control-handler-result.js'; + +type NativeCompletion = Awaited>; +type NativeResult = { + state: 'not_started' | 'pending' | 'completed' | 'unknown'; + completion?: NativeCompletion['info']; + result?: boolean; + error?: unknown; +}; +type Finalization = { + autoCommit?: + | { state: 'running' } + | { state: 'completed'; result: AutoCommitResult } + | { state: 'unknown'; error: unknown }; + condensation?: { + state: 'running' | 'completed' | 'unknown'; + result?: boolean; + error?: unknown; + }; +}; + +export type SessionOperationWork = + | { + operation: 'session.attach'; + payload: SessionAttachPayload; + apply: ( + session: SessionRequestIdentity, + payload: SessionAttachPayload, + hooks: { signal: AbortSignal; emitPreparing?: AttachPreparingEmitter } + ) => Promise; + onAttached: () => void; + emitPreparing?: AttachPreparingEmitter; + } + | { + operation: 'session.prompt'; + payload: SessionPromptPayload; + runtime: WorktreeKiloRuntime; + materializeAttachments?: typeof materializeMessageAttachments; + runAutoCommit?: typeof runAutoCommit; + }; + +export type SessionOperationDependencies = { + signal?: AbortSignal; + isCurrent: () => boolean; + getRuntime: () => WorktreeKiloRuntime | undefined; + retireRuntime: (reason: string) => void; + emitSessionEvent: (payload: SessionEventPayload, options?: { retained?: true }) => void; + sendOperationResult?: OperationResultSender; + onLocalCompletion: (retain: boolean) => void; + onDiagnostic?: ControlDiagnosticReporter; +}; + +class ControlTaskCancellation extends Error { + constructor( + readonly status: 'failed' | 'cancelled', + message: string + ) { + super(message); + } +} + +function fail(message: string, retryable: boolean): ControlHandlerResult { + return { ok: false, error: { code: 'not_ready', message, retryable } }; +} + +function kiloFailure(error: unknown): ControlHandlerResult { + return fail('Kilo request failed', isKiloServerUnreachableError(error)); +} + +export class SessionOperation { + readonly session: Readonly; + readonly authorization?: Readonly; + readonly messageId?: string; + readonly executionDeadlineAt: number; + readonly signal: AbortSignal; + readonly done: Promise; + private readonly controller = new AbortController(); + private readonly completion = Promise.withResolvers(); + private readonly intent: ReturnType; + private readonly startedAt = Date.now(); + private readonly timeout: ReturnType; + private phase: 'preparation' | 'execution' | 'finalizing'; + private client?: WrapperKiloClient; + private native: NativeResult = { state: 'not_started' }; + private nativePending?: Promise; + private finalization: Finalization = {}; + private readonly retainedNotifications = createRetainedOperationNotifications(); + private outcome?: SessionMessageOutcome; + private local?: { result: ControlHandlerResult; completedAt: number }; + private delivery?: OperationResultDelivery; + + constructor( + session: SessionRequestIdentity, + authorization: SessionOperationAuthorization | undefined, + private readonly work: SessionOperationWork, + private readonly deps: SessionOperationDependencies + ) { + this.session = Object.freeze({ ...session }); + this.authorization = authorization + ? Object.freeze({ ...authorization, session: Object.freeze({ ...authorization.session }) }) + : undefined; + this.work = + work.operation === 'session.attach' + ? { ...work, payload: structuredClone(work.payload) } + : { ...work, payload: structuredClone(work.payload) }; + this.phase = work.operation === 'session.attach' ? 'preparation' : 'execution'; + this.messageId = + work.operation === 'session.attach' + ? (authorization?.messageId ?? work.payload.preparation?.triggerMessageId) + : work.payload.messageId; + this.intent = structuredClone(operationIntent(work.operation, work.payload)); + this.executionDeadlineAt = + this.phase === 'preparation' + ? Math.min( + this.startedAt + SANDBOX_CONTROL_ATTACH_TIMEOUT_MS, + authorization?.dispatchDeadlineAt ?? Infinity + ) + : this.startedAt + SANDBOX_CONTROL_EXECUTION_TIMEOUT_MS; + const signals = [this.controller.signal]; + if (deps.signal) signals.push(deps.signal); + if (work.operation === 'session.prompt') signals.push(work.runtime.signal); + this.signal = AbortSignal.any(signals); + this.done = this.completion.promise; + this.diagnostic('started'); + this.timeout = setTimeout( + () => this.expire(), + Math.max(0, this.executionDeadlineAt - this.startedAt) + ); + this.timeout.unref(); + void Promise.resolve() + .then(() => + this.work.operation === 'session.attach' ? this.attach(this.work) : this.execute(this.work) + ) + .catch((error: unknown) => { + this.recordUncertainty(error); + const cancellation: unknown = this.signal.reason; + return cancellation instanceof ControlTaskCancellation && cancellation.status === 'failed' + ? fail(cancellation.message, true) + : kiloFailure(error); + }) + .then(result => this.complete(result)); + } + + get kind() { + return this.phase; + } + get locallyComplete() { + return this.local !== undefined; + } + + snapshot() { + const retained = this.retainedNotifications.snapshot(); + return { + kind: this.phase, + native: structuredClone(this.native), + finalization: structuredClone(this.finalization), + events: retained.events, + preparing: retained.preparing, + outcome: this.outcome ? structuredClone(this.outcome) : undefined, + local: this.local ? structuredClone(this.local) : undefined, + delivery: this.delivery?.snapshot(), + }; + } + + deliveryResult(): SessionOperationDelivery | undefined { + return this.delivery?.result(); + } + + waitForDelivery(): Promise { + return this.delivery?.drain() ?? Promise.resolve(); + } + + matchesAuthorization(authorization: SessionOperationAuthorization): boolean { + return ( + this.authorization !== undefined && sameSessionOperation(this.authorization, authorization) + ); + } + + matchesIntent(payload: unknown): boolean { + return isDeepStrictEqual(this.intent, operationIntent(this.work.operation, payload)); + } + + canPrune(now: number): boolean { + const delivery = this.delivery?.status(); + return ( + this.authorization !== undefined && + this.local !== undefined && + this.native.state !== 'pending' && + this.native.state !== 'unknown' && + delivery?.state === 'acknowledged' && + now >= Math.max(this.authorization.dispatchDeadlineAt, delivery.deadlineAt) + ); + } + + acknowledge(ack: SessionOperationAck, isCurrent: () => boolean): Promise { + return this.delivery?.acknowledge(ack, isCurrent) ?? Promise.resolve(false); + } + + cancel(reason: string, status: 'failed' | 'cancelled'): void { + if (!this.local) this.controller.abort(new ControlTaskCancellation(status, reason)); + } + + private diagnostic(phase: string): void { + emitControlDiagnostic(this.deps.onDiagnostic, 'session.task', { + sessionId: this.session.sessionId, + kiloSessionId: this.session.kiloSessionId, + messageId: this.messageId, + kind: this.work.operation === 'session.attach' ? 'preparation' : 'execution', + phase, + elapsedMs: Date.now() - this.startedAt, + }); + } + + private expire(): void { + if (this.local) return; + const reason = + this.work.operation === 'session.attach' + ? 'Session preparation timed out' + : 'Execution exceeded the 60 minute limit'; + this.diagnostic('deadline_expired'); + this.cancel(reason, 'failed'); + this.deps.retireRuntime(reason); + } + + private assertCurrent(): void { + if (Date.now() >= this.executionDeadlineAt) this.expire(); + this.signal.throwIfAborted(); + if (!this.deps.isCurrent()) throw new Error('Operation execution authority expired'); + } + + private recordUncertainty(error: unknown): void { + if (this.native.state === 'pending') this.native = { state: 'unknown', error }; + if (this.finalization.autoCommit?.state === 'running') + this.finalization.autoCommit = { state: 'unknown', error }; + if (this.finalization.condensation?.state === 'running') + this.finalization.condensation = { state: 'unknown', error }; + } + + private captureClient(client: WrapperKiloClient): void { + if (this.client) return; + this.client = client; + } + + private async attach( + work: Extract + ): Promise { + this.assertCurrent(); + const result = await work.apply(this.session, work.payload, { + signal: this.signal, + emitPreparing: event => { + const retained = + this.authorization && isRetainedOperationPreparing(event) + ? this.retainedNotifications.retainPreparing(event) + : undefined; + if (this.authorization && isRetainedOperationPreparing(event) && !retained) return; + try { + work.emitPreparing?.( + retained ?? event, + this.authorization ? { retained: true } : undefined + ); + } catch { + if (!this.authorization) throw new Error('Preparation event delivery failed'); + } + }, + }); + if (this.signal.aborted) { + const cancellation: unknown = this.signal.reason; + return fail( + cancellation instanceof ControlTaskCancellation && cancellation.status === 'failed' + ? cancellation.message + : 'Session attachment cancelled', + true + ); + } + if (result.ok) work.onAttached(); + return result; + } + + private emitFinalizationEvent(event: IngestEvent): void { + const payload = sessionEventPayloadSchema.safeParse({ + type: event.streamEventType, + properties: event.data, + timestamp: event.timestamp, + }); + if (!payload.success) return; + const requiresRetention = + payload.data.type === 'autocommit_completed' || payload.data.type === 'status'; + const retained = + this.authorization && requiresRetention + ? this.retainedNotifications.retainFinalization(payload.data) + : undefined; + if (this.authorization && requiresRetention && !retained) return; + try { + this.deps.emitSessionEvent( + retained ?? payload.data, + this.authorization ? { retained: true } : undefined + ); + } catch { + if (!this.authorization) throw new Error('Finalization event delivery failed'); + } + } + + private observeNative(pending: Promise): Promise { + this.native.state = 'pending'; + const observed = pending.then( + completion => { + if (!this.local) + this.native = { state: 'completed', completion: structuredClone(completion.info) }; + return completion; + }, + (error: unknown) => { + if (!this.local) this.native = { state: 'unknown', error }; + throw error; + } + ); + this.nativePending = observed; + return observed; + } + + private async waitForNativeAfterAbort( + pending: Promise, + deadlineAt: number + ): Promise { + try { + await withTimeoutAndAbort( + pending.then( + () => undefined, + () => undefined + ), + { + timeoutMs: Math.max(0, deadlineAt - Date.now()), + timeoutMessage: 'Native cancellation did not settle', + abortMessage: 'Native cancellation interrupted', + } + ); + } catch (error) { + this.recordUncertainty(error); + this.deps.retireRuntime('Native cancellation did not settle'); + } + } + + private async summarize( + client: WrapperKiloClient, + model: { providerID?: string; modelID: string }, + auto?: boolean + ): Promise { + this.signal.throwIfAborted(); + const remaining = this.executionDeadlineAt - Date.now(); + if (remaining <= 0) throw new Error('Execution exceeded the 60 minute limit'); + this.finalization.condensation = { state: 'running' }; + if (auto === undefined) this.native.state = 'pending'; + const pending = client + .summarizeSession({ + sessionId: this.session.kiloSessionId, + directory: this.session.directory, + signal: this.signal, + model, + ...(auto === undefined ? {} : { auto }), + }) + .then( + result => { + if (!this.local) { + this.finalization.condensation = { state: 'completed', result }; + if (auto === undefined) this.native = { state: 'completed', result }; + } + return result; + }, + (error: unknown) => { + if (!this.local) { + this.finalization.condensation = { state: 'unknown', error }; + if (auto === undefined) this.native = { state: 'unknown', error }; + } + throw error; + } + ); + this.nativePending = pending; + const success = await withTimeoutAndAbort(pending, { + signal: this.signal, + timeoutMs: remaining, + timeoutMessage: 'Execution exceeded the 60 minute limit', + abortMessage: 'Execution cancelled', + }); + if (!success) throw new Error('Session summarization failed'); + } + + private async execute( + work: Extract + ): Promise { + const { session, signal } = this; + const request = work.payload; + const { runtime } = work; + const { kiloClient, env } = runtime; + this.captureClient(kiloClient); + const assertCurrent = (submitting = false) => { + signal.throwIfAborted(); + if ( + !this.deps.isCurrent() || + runtime.kiloClient !== kiloClient || + Date.now() >= this.executionDeadlineAt || + (submitting && this.deps.getRuntime() !== runtime) || + (submitting && this.authorization && Date.now() >= this.authorization.dispatchDeadlineAt) + ) + throw new Error('Operation execution authority expired'); + }; + const { messageId, turn, agent } = request; + const startedAt = Date.now(); + const diagnostic = (phase: string, status?: SessionMessageOutcome['status']): void => + emitControlDiagnostic(this.deps.onDiagnostic, 'session.execution', { + sessionId: session.sessionId, + kiloSessionId: session.kiloSessionId, + messageId, + phase, + status, + elapsedMs: Date.now() - startedAt, + aborted: signal.aborted, + }); + let outcome: SessionMessageOutcome; + let result: ControlHandlerResult = { ok: true, result: {} }; + let failureReason = 'Kilo execution failed'; + const emitStatus = (message: string): void => + this.emitFinalizationEvent({ + streamEventType: 'status', + data: { message, messageId }, + timestamp: new Date().toISOString(), + }); + try { + assertCurrent(true); + let completion: NativeCompletion | undefined; + const options = { + sessionId: session.kiloSessionId, + directory: session.directory, + signal, + messageId, + agent: agent.mode, + ...(agent.variant ? { variant: agent.variant } : {}), + }; + const deadline = { + signal, + timeoutMs: Math.max(1, this.executionDeadlineAt - Date.now()), + timeoutMessage: 'Execution exceeded the 60 minute limit', + abortMessage: 'Execution cancelled', + }; + if (turn.type === 'prompt') { + if (agent.model === undefined) throw new Error('Prompt model is required'); + const message = await (work.materializeAttachments ?? materializeMessageAttachments)( + { + id: messageId, + prompt: turn.prompt, + parts: turn.parts, + attachments: request.attachments, + }, + { signal } + ); + assertCurrent(true); + diagnostic('prompt_started'); + this.native.state = 'pending'; + completion = await withTimeoutAndAbort( + this.observeNative( + kiloClient.sendPrompt({ + ...options, + prompt: message.prompt, + ...(message.parts ? { parts: message.parts } : {}), + model: { providerID: 'kilo', modelID: agent.model }, + }) + ), + { ...deadline, timeoutMs: Math.max(1, this.executionDeadlineAt - Date.now()) } + ); + diagnostic('prompt_completed'); + } else if (turn.command === 'compact') { + if (!agent.model) throw new Error('Model is required for compact'); + failureReason = 'Context condensation failed'; + emitStatus('Condensing context...'); + diagnostic('compact_started'); + await this.summarize(kiloClient, { providerID: 'kilo', modelID: agent.model }); + diagnostic('compact_completed'); + signal.throwIfAborted(); + emitStatus('Context condensed successfully'); + } else { + assertCurrent(true); + diagnostic('command_started'); + this.native.state = 'pending'; + completion = await withTimeoutAndAbort( + this.observeNative( + kiloClient.sendCommand({ + ...options, + command: turn.command, + args: turn.arguments, + ...(agent.model !== undefined + ? { model: { providerID: 'kilo', modelID: agent.model } } + : {}), + }) + ), + deadline + ); + diagnostic('command_completed'); + } + assertCurrent(); + const error = completion?.info.error; + if ( + !error && + (request.finalization?.autoCommit || request.finalization?.condenseOnComplete) + ) { + this.phase = 'finalizing'; + diagnostic('finalization_started'); + if (request.finalization.autoCommit) { + failureReason = 'Auto-commit failed'; + assertCurrent(); + diagnostic('autocommit_started'); + this.finalization.autoCommit = { state: 'running' }; + const committed = await (work.runAutoCommit ?? runAutoCommit)({ + workspacePath: session.directory, + kiloClient, + env, + messageId: completion?.info.id ?? messageId, + signal, + onEvent: event => this.emitFinalizationEvent(event), + }); + this.finalization.autoCommit = { state: 'completed', result: structuredClone(committed) }; + assertCurrent(); + if (!committed.success) throw new Error('Auto-commit failed'); + diagnostic('autocommit_completed'); + } + if (request.finalization.condenseOnComplete) { + failureReason = 'Context condensation failed'; + const model = agent.model + ? { providerID: 'kilo', modelID: agent.model } + : completion + ? { providerID: completion.info.providerID, modelID: completion.info.modelID } + : undefined; + if (!model) throw new Error('Model is required for condensation'); + emitStatus('Condensing context...'); + assertCurrent(); + diagnostic('condense_started'); + await this.summarize(kiloClient, model, true); + diagnostic('condense_completed'); + signal.throwIfAborted(); + emitStatus('Context condensed successfully'); + } + } + outcome = error + ? { + messageId, + status: error.name === 'MessageAbortedError' ? 'cancelled' : 'failed', + reason: `Kilo execution ended with ${error.name}`, + } + : { messageId, status: 'completed' }; + } catch (error) { + diagnostic('execution_failed'); + this.recordUncertainty(error); + const cancellation: unknown = signal.reason; + outcome = { + messageId, + status: cancellation instanceof ControlTaskCancellation ? cancellation.status : 'failed', + reason: + cancellation instanceof ControlTaskCancellation + ? cancellation.message + : this.authorization && this.native.state === 'unknown' + ? 'Kilo execution outcome is unconfirmed' + : failureReason, + }; + try { + diagnostic('abort_started'); + const cleanupDeadlineAt = Date.now() + KILO_CONTROL_REQUEST_TIMEOUT_MS; + const abortController = new AbortController(); + const abortTimer = setTimeout( + () => abortController.abort(new Error('Kilo cancellation timed out')), + Math.max(0, cleanupDeadlineAt - Date.now()) + ); + let aborted: boolean; + try { + aborted = await withTimeoutAndAbort( + kiloClient.abortSession({ + sessionId: session.kiloSessionId, + directory: session.directory, + signal: abortController.signal, + }), + { + timeoutMs: Math.max(0, cleanupDeadlineAt - Date.now()), + timeoutMessage: 'Kilo cancellation timed out', + abortMessage: 'Kilo cancellation interrupted', + } + ); + } finally { + clearTimeout(abortTimer); + } + if (aborted !== true) throw new Error('Kilo cancellation was not confirmed'); + const pendingNative = this.nativePending; + if (pendingNative) await this.waitForNativeAfterAbort(pendingNative, cleanupDeadlineAt); + diagnostic('abort_completed'); + } catch (error) { + diagnostic('abort_failed'); + this.deps.retireRuntime('Kilo cancellation failed'); + result = kiloFailure(error); + } + const original = this.native.completion; + if (original?.error) + outcome = { + messageId, + status: original.error.name === 'MessageAbortedError' ? 'cancelled' : 'failed', + reason: `Kilo execution ended with ${original.error.name}`, + }; + else if ( + this.native.state === 'unknown' && + this.native.error instanceof Error && + this.native.error.name === 'MessageAbortedError' + ) + outcome = { + messageId, + status: 'cancelled', + reason: 'Kilo execution ended with MessageAbortedError', + }; + else if ( + this.native.state === 'completed' && + this.native.result !== false && + (!request.finalization?.autoCommit || + (this.finalization.autoCommit?.state === 'completed' && + this.finalization.autoCommit.result.success)) && + (!request.finalization?.condenseOnComplete || + (this.finalization.condensation?.state === 'completed' && + this.finalization.condensation.result === true)) + ) + outcome = { messageId, status: 'completed' }; + } + this.outcome = sessionMessageOutcomeSchema.parse(outcome); + if (!this.authorization) { + try { + diagnostic('outcome_sending', outcome.status); + this.deps.emitSessionEvent({ type: 'session.message.outcome', properties: this.outcome }); + diagnostic('outcome_sent', outcome.status); + } catch { + diagnostic('outcome_failed', outcome.status); + this.deps.retireRuntime('Session outcome delivery failed'); + return fail('Session outcome delivery failed', false); + } + } + return result; + } + + private complete(result: ControlHandlerResult): void { + result = result.ok ? { ok: true, result: result.result } : { ok: false, error: result.error }; + this.local = { result: structuredClone(result), completedAt: Date.now() }; + clearTimeout(this.timeout); + const retain = + result.ok || + result.error.code !== 'session_busy' || + this.native.state !== 'not_started' || + this.signal.aborted; + this.deps.onLocalCompletion(retain); + this.diagnostic(result.ok ? 'finished' : 'failed'); + this.completion.resolve(result); + if (this.authorization && retain) { + const retained = this.retainedNotifications.snapshot(); + const delivery = sessionOperationDeliverySchema.parse({ + version: 2, + authorization: this.authorization, + completedAt: this.local.completedAt, + result, + ...(this.outcome ? { outcome: this.outcome } : {}), + ...(this.native.completion ? { assistantMessageId: this.native.completion.id } : {}), + events: retained.events, + preparing: retained.preparing, + }); + this.delivery = createOperationResultDelivery( + delivery, + Math.min( + this.local.completedAt + SANDBOX_CONTROL_OUTCOME_TIMEOUT_MS, + sessionOperationExpiresAt(this.authorization) + ), + this.deps.sendOperationResult + ); + void this.delivery.start(); + } + } +} diff --git a/services/cloud-agent-next/wrapper/src/control/worktree-mutation-notifications.test.ts b/services/cloud-agent-next/wrapper/src/control/worktree-mutation-notifications.test.ts index 3228b7c3ae..6541128ad5 100644 --- a/services/cloud-agent-next/wrapper/src/control/worktree-mutation-notifications.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/worktree-mutation-notifications.test.ts @@ -14,8 +14,8 @@ import type { WrapperKiloClient } from '../kilo-api'; import { eventKiloSessionId, sessionEventIdentity, unfilteredKiloEvents } from './feed'; import { buildHeartbeatPayload, + createControlHandlerDeps, createSessionActivityRegistry, - type HandlerDeps, type HandlerSessionSnapshot, } from './sandbox-control-handlers'; import { @@ -596,12 +596,14 @@ describe('worktree mutation notifications', () => { const activity = createSessionActivityRegistry(); activity.attach('root'); activity.attach('sibling'); - const deps = { + const deps = createControlHandlerDeps({ sessions: h.sessions, - tasks: new Map(), activity, + version: 'test', kiloReady: true, - } as HandlerDeps; + emitSessionEvent: () => {}, + retireRuntime: () => {}, + }); const snapshots = structuredClone(h.sessions); const heartbeat = buildHeartbeatPayload(deps); const routed = []; diff --git a/services/cloud-agent-next/wrapper/src/control/worktree-runtime.test.ts b/services/cloud-agent-next/wrapper/src/control/worktree-runtime.test.ts index c7a0eea475..585f343a59 100644 --- a/services/cloud-agent-next/wrapper/src/control/worktree-runtime.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/worktree-runtime.test.ts @@ -16,6 +16,7 @@ import { } from '../../../src/shared/sandbox-control-protocol'; import { buildHeartbeatPayload, + createControlHandlerDeps, handleControlRequest, type HandlerDeps, } from './sandbox-control-handlers'; @@ -272,16 +273,15 @@ function createHandlerDeps(registry: WorktreeKiloRuntimes): HandlerDeps { getKiloRuntime: directory => registry.get(directory), }); terminalRuntimes.push(terminalRuntime); - return { + return createControlHandlerDeps({ kiloRuntimes: registry, terminalRuntime, version: 'test', kiloReady: true, sessions: [], - tasks: new Map(), emitSessionEvent: () => {}, retireRuntime: () => {}, - }; + }); } beforeEach(() => { @@ -768,13 +768,12 @@ describe('worktree Kilo runtime registry', () => { wrapperInstanceId: crypto.randomUUID(), getKiloRuntime: directory => harness.registry.get(directory), }); - const deps: HandlerDeps = { + const deps: HandlerDeps = createControlHandlerDeps({ kiloRuntimes: harness.registry, terminalRuntime: terminals, version: 'test', kiloReady: true, sessions: [], - tasks: new Map(), emitSessionEvent: (identity, event) => { if (event.type === 'session.message.outcome') { const { messageId, status } = sessionMessageOutcomeSchema.parse(event.properties); @@ -782,8 +781,9 @@ describe('worktree Kilo runtime registry', () => { } }, retireRuntime: () => {}, - }; - const waitForTasks = () => Promise.all([...deps.tasks.values()].map(task => task.done)); + }); + const waitForTasks = () => + Promise.all(deps.operations.activeOperations().map(task => task.done)); try { const attached = await Promise.all( identities.map((identity, index) =>