diff --git a/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.test.ts b/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.test.ts new file mode 100644 index 0000000000..f96e83a343 --- /dev/null +++ b/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.test.ts @@ -0,0 +1,572 @@ +import { describe, expect, it, vi } from 'vitest'; +import type { Env, SandboxInstance } from '../../types.js'; +import type { SessionMetadata } from '../../persistence/session-metadata.js'; +import { WrapperClient } from '../../kilo/wrapper-client.js'; +import { WRAPPER_VERSION } from '../../shared/wrapper-version.js'; +import type { EnsureWrapperRequest } from '../protocol.js'; +import { CloudflareAgentSandbox } from './cloudflare-agent-sandbox.js'; +import { + SandboxCapacityInspectionError, + WorkspaceCapacityAdmissionRejectedError, +} from '../../workspace-errors.js'; + +vi.mock('@cloudflare/sandbox', () => ({ getSandbox: vi.fn() })); + +function metadata(options?: { devcontainer?: boolean }): SessionMetadata { + return { + metadataSchemaVersion: 2, + identity: { sessionId: 'agent_cloudflare', userId: 'user_cloudflare', orgId: 'org_cloudflare' }, + auth: {}, + workspace: { + sandboxId: options?.devcontainer ? 'dind-abcdef' : 'ses-abcdef', + }, + ...(options?.devcontainer + ? { + devcontainer: { + workspacePath: '/workspace/cloudflare', + innerWorkspaceFolder: '/workspaces/repo', + wrapperPort: 4173, + configPath: '.devcontainer/devcontainer.json', + }, + } + : {}), + lifecycle: { version: 1, timestamp: 1 }, + }; +} + +function ensureRequest(options?: { + devcontainer?: boolean; + leased?: boolean; +}): EnsureWrapperRequest { + const sandboxId = options?.devcontainer ? 'dind-abcdef' : 'ses-abcdef'; + const sessionMetadata = metadata(options); + return { + plan: { + scope: { sessionId: 'agent_cloudflare', userId: 'user_cloudflare', orgId: 'org_cloudflare' }, + turn: { + type: 'prompt', + messageId: 'msg_018f1e2d3c4bCloudflareAAAA', + prompt: 'Run in Cloudflare', + }, + agent: { mode: 'code', model: 'test-model' }, + workspace: { sandboxId, metadata: sessionMetadata }, + wrapper: { + kiloSessionId: 'kilo_cloudflare', + fence: { + wrapperRunId: 'wr_cloudflare', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_cloudflare', + }, + }, + }, + ...(options?.leased + ? { leasedInstance: { instanceId: 'instance_cloudflare', instanceGeneration: 3 } } + : {}), + prepared: { + ready: { + workspacePath: '/workspace/cloudflare', + sandboxId, + sessionHome: '/home/agent_cloudflare', + branchName: 'session/agent_cloudflare', + kiloSessionId: 'kilo_cloudflare', + }, + context: { workspacePath: '/workspace/cloudflare' }, + }, + }; +} + +describe('CloudflareAgentSandbox', () => { + it('starts an ordinary bootstrap wrapper through the adapter', async () => { + const bootstrapSession = {}; + const createSession = vi.fn().mockResolvedValue(bootstrapSession); + const ensureBootstrapWrapper = vi + .spyOn(WrapperClient, 'ensureBootstrapWrapper') + .mockResolvedValueOnce({ client: {} as WrapperClient }); + const sandbox = new CloudflareAgentSandbox({} as Env, metadata(), { + resolveSandbox: () => + ({ + exec: vi.fn().mockResolvedValue({ exitCode: 0, stdout: 'exists\n', stderr: '' }), + createSession, + }) as unknown as SandboxInstance, + }); + + await expect(sandbox.ensureWrapper(ensureRequest())).resolves.toMatchObject({ + status: 'wrapper-running', + }); + expect(createSession).toHaveBeenCalledWith({ + name: 'agent_cloudflare-bootstrap', + env: {}, + cwd: '/', + }); + expect(ensureBootstrapWrapper).toHaveBeenCalledWith(expect.anything(), bootstrapSession, { + agentSessionId: 'agent_cloudflare', + userId: 'user_cloudflare', + }); + ensureBootstrapWrapper.mockRestore(); + }); + + it('types ENOSPC during the cold bootstrap probe as sandbox unusable', async () => { + const createSession = vi.fn(); + const sandbox = new CloudflareAgentSandbox({} as Env, metadata(), { + resolveSandbox: () => + ({ + exec: vi.fn().mockResolvedValue({ + exitCode: 1, + stdout: '', + stderr: 'ENOSPC: no space left on device', + }), + createSession, + }) as unknown as SandboxInstance, + }); + + await expect(sandbox.ensureWrapper(ensureRequest())).rejects.toBeInstanceOf( + SandboxCapacityInspectionError + ); + expect(createSession).not.toHaveBeenCalled(); + }); + + it('rejects cold bootstrap admission before creating a wrapper session', async () => { + const createSession = vi.fn(); + const ensureBootstrapWrapper = vi.spyOn(WrapperClient, 'ensureBootstrapWrapper'); + const exec = vi + .fn() + .mockResolvedValueOnce({ exitCode: 1, stdout: '', stderr: '' }) + .mockResolvedValueOnce({ exitCode: 0, stdout: '536870912 10485760000\n', stderr: '' }) + .mockResolvedValueOnce({ exitCode: 1, stdout: '', stderr: 'no sessions' }) + .mockResolvedValueOnce({ exitCode: 0, stdout: '536870912 10485760000\n', stderr: '' }); + const sandbox = new CloudflareAgentSandbox({} as Env, metadata(), { + resolveSandbox: () => ({ exec, createSession }) as unknown as SandboxInstance, + }); + + await expect(sandbox.ensureWrapper(ensureRequest())).rejects.toBeInstanceOf( + WorkspaceCapacityAdmissionRejectedError + ); + expect(createSession).not.toHaveBeenCalled(); + expect(ensureBootstrapWrapper).not.toHaveBeenCalled(); + ensureBootstrapWrapper.mockRestore(); + }); + + it('passes a leased physical identity into bootstrap startup', async () => { + const bootstrapSession = {}; + const ensureBootstrapWrapper = vi + .spyOn(WrapperClient, 'ensureBootstrapWrapper') + .mockResolvedValueOnce({ client: {} as WrapperClient }); + const sandbox = new CloudflareAgentSandbox({} as Env, metadata(), { + resolveSandbox: () => + ({ + exec: vi.fn().mockResolvedValue({ exitCode: 0, stdout: 'exists\n', stderr: '' }), + createSession: vi.fn().mockResolvedValue(bootstrapSession), + }) as unknown as SandboxInstance, + }); + + await sandbox.ensureWrapper(ensureRequest({ leased: true })); + + expect(ensureBootstrapWrapper).toHaveBeenCalledWith(expect.anything(), bootstrapSession, { + agentSessionId: 'agent_cloudflare', + userId: 'user_cloudflare', + leasedInstance: { instanceId: 'instance_cloudflare', instanceGeneration: 3 }, + }); + ensureBootstrapWrapper.mockRestore(); + }); + + it('does not start a devcontainer wrapper after cold workspace admission is rejected', async () => { + const rejection = new WorkspaceCapacityAdmissionRejectedError({ + availableMB: 512, + thresholdMB: 2048, + cleaned: 0, + skipped: 1, + }); + const prepareWorkspace = vi.fn().mockRejectedValue(rejection); + const ensureWrapper = vi.spyOn(WrapperClient, 'ensureWrapper'); + const sandbox = new CloudflareAgentSandbox({} as Env, metadata({ devcontainer: true }), { + resolveSandbox: () => ({}) as SandboxInstance, + sessionService: { prepareWorkspace } as never, + }); + + await expect(sandbox.ensureWrapper(ensureRequest({ devcontainer: true }))).rejects.toBe( + rejection + ); + expect(ensureWrapper).not.toHaveBeenCalled(); + ensureWrapper.mockRestore(); + }); + + it('prepares and starts a devcontainer wrapper through the adapter', async () => { + const devcontainer = { + containerId: 'container-dev', + innerWorkspaceFolder: '/workspaces/repo', + workspacePath: '/workspace/cloudflare', + agentSessionId: 'agent_cloudflare', + overrideConfigPath: '/tmp/devcontainer.json', + teardown: vi.fn(), + }; + const request = ensureRequest({ devcontainer: true, leased: true }); + const prepareWorkspace = vi.fn().mockResolvedValue({ + context: { workspacePath: '/workspace/cloudflare' }, + ready: { + ...request.prepared.ready, + devcontainer: { + workspacePath: '/workspace/cloudflare', + innerWorkspaceFolder: '/workspaces/repo', + wrapperPort: 4173, + configPath: '.devcontainer/devcontainer.json', + }, + }, + runtimeEnv: {}, + session: {}, + devcontainer, + }); + const ensureWrapper = vi.spyOn(WrapperClient, 'ensureWrapper').mockResolvedValueOnce({ + client: {} as WrapperClient, + sessionId: 'kilo_cloudflare', + }); + const sandbox = new CloudflareAgentSandbox({} as Env, metadata({ devcontainer: true }), { + resolveSandbox: () => ({}) as SandboxInstance, + sessionService: { prepareWorkspace } as never, + }); + + await expect(sandbox.ensureWrapper(request)).resolves.toMatchObject({ + status: 'session-ready', + kiloSessionId: 'kilo_cloudflare', + }); + expect(ensureWrapper).toHaveBeenCalledWith( + expect.anything(), + expect.anything(), + expect.objectContaining({ + fixedPort: 4173, + devcontainer, + leasedInstance: { instanceId: 'instance_cloudflare', instanceGeneration: 3 }, + }) + ); + ensureWrapper.mockRestore(); + }); + + it('gets an existing running wrapper without provisioning compute', async () => { + const getSession = vi.fn().mockResolvedValue({}); + const createSession = vi.fn(); + const sandbox = new CloudflareAgentSandbox({} as Env, metadata(), { + resolveSandbox: () => + ({ + listProcesses: vi.fn().mockResolvedValue([ + { + id: 'wrapper-1', + command: 'kilocode-wrapper WRAPPER_PORT=5000 --agent-session agent_cloudflare', + status: 'running', + }, + ]), + getSession, + createSession, + }) as unknown as SandboxInstance, + }); + + await expect(sandbox.getRunningWrapper()).resolves.toBeInstanceOf(WrapperClient); + expect(getSession).toHaveBeenCalledWith('agent_cloudflare-bootstrap'); + expect(createSession).not.toHaveBeenCalled(); + }); + + it('returns a terminal client only for a healthy live wrapper', async () => { + const containerFetch = vi.fn().mockResolvedValue( + Response.json({ + healthy: true, + state: 'idle', + version: WRAPPER_VERSION, + sessionId: 'kilo-cloudflare', + }) + ); + const sandbox = new CloudflareAgentSandbox({} as Env, metadata(), { + resolveSandbox: () => + ({ + listProcesses: vi.fn().mockResolvedValue([ + { + id: 'wrapper-1', + command: 'kilocode-wrapper WRAPPER_PORT=5000 --agent-session agent_cloudflare', + status: 'running', + }, + ]), + containerFetch, + }) as unknown as SandboxInstance, + }); + + await expect(sandbox.getRunningTerminalClient()).resolves.toMatchObject({ status: 'ready' }); + }); + + it('distinguishes an unhealthy live wrapper from an absent terminal wrapper', async () => { + const sandbox = new CloudflareAgentSandbox({} as Env, metadata(), { + resolveSandbox: () => + ({ + listProcesses: vi.fn().mockResolvedValue([ + { + id: 'wrapper-1', + command: 'kilocode-wrapper WRAPPER_PORT=5000 --agent-session agent_cloudflare', + status: 'running', + }, + ]), + containerFetch: vi.fn().mockResolvedValue( + Response.json({ + healthy: false, + state: 'idle', + version: WRAPPER_VERSION, + sessionId: 'kilo-cloudflare', + }) + ), + }) as unknown as SandboxInstance, + }); + + await expect(sandbox.getRunningTerminalClient()).resolves.toEqual({ status: 'unhealthy' }); + }); + + it('discovers all tagged and legacy physical wrappers for its session', async () => { + const sandbox = new CloudflareAgentSandbox({} as Env, metadata(), { + resolveSandbox: () => + ({ + listProcesses: vi.fn().mockResolvedValue([ + { + id: 'wrapper-tagged', + command: + 'WRAPPER_PORT=5000 kilocode-wrapper --agent-session agent_cloudflare --wrapper-instance-id instance_1 --wrapper-instance-generation 2', + status: 'running', + }, + { + id: 'wrapper-legacy', + command: 'WRAPPER_PORT=5001 kilocode-wrapper --agent-session agent_cloudflare', + status: 'running', + }, + ]), + exec: vi.fn().mockResolvedValue({ exitCode: 0, stdout: '' }), + }) as unknown as SandboxInstance, + }); + + await expect(sandbox.discoverSessionWrappers()).resolves.toEqual({ + status: 'present', + observed: [ + { + representation: 'process', + id: 'wrapper-tagged', + port: 5000, + instanceId: 'instance_1', + instanceGeneration: 2, + }, + { representation: 'process', id: 'wrapper-legacy', port: 5001 }, + ], + }); + }); + + it('does not report lifecycle absence when physical inspection fails', async () => { + const sandbox = new CloudflareAgentSandbox({} as Env, metadata(), { + resolveSandbox: () => + ({ + listProcesses: vi.fn().mockRejectedValue(new Error('sandbox unavailable')), + }) as unknown as SandboxInstance, + }); + + await expect(sandbox.discoverSessionWrappers()).resolves.toMatchObject({ + status: 'inspection-failed', + error: expect.stringContaining('sandbox unavailable'), + }); + }); + + it('does not require Docker discovery for standard sandboxes', async () => { + const exec = vi.fn().mockRejectedValue(new Error('docker unavailable')); + const sandbox = new CloudflareAgentSandbox({} as Env, metadata(), { + resolveSandbox: () => + ({ listProcesses: vi.fn().mockResolvedValue([]), exec }) as unknown as SandboxInstance, + }); + + await expect(sandbox.discoverSessionWrappers()).resolves.toEqual({ status: 'absent' }); + expect(exec).not.toHaveBeenCalled(); + }); + + it('requires container discovery for a DIND sandbox even before resolved devcontainer metadata exists', async () => { + const unresolvedDindMetadata = { + ...metadata(), + workspace: { sandboxId: 'dind-unresolved' }, + } satisfies SessionMetadata; + const sandbox = new CloudflareAgentSandbox({} as Env, unresolvedDindMetadata, { + resolveSandbox: () => + ({ + listProcesses: vi.fn().mockResolvedValue([]), + exec: vi.fn().mockRejectedValue(new Error('docker inspection unavailable')), + }) as unknown as SandboxInstance, + }); + + await expect(sandbox.discoverSessionWrappers()).resolves.toMatchObject({ + status: 'inspection-failed', + error: expect.stringContaining('docker inspection unavailable'), + }); + }); + + it('stops remaining session wrappers before confirming an instance target is absent', async () => { + const stopObservedWrappers = vi.fn().mockResolvedValue(undefined); + const listProcesses = vi + .fn() + .mockResolvedValueOnce([ + { + id: 'wrapper-legacy', + command: 'WRAPPER_PORT=5001 kilocode-wrapper --agent-session agent_cloudflare', + status: 'running', + }, + ]) + .mockResolvedValueOnce([]); + const sandbox = new CloudflareAgentSandbox({} as Env, metadata(), { + resolveSandbox: () => ({ listProcesses }) as unknown as SandboxInstance, + stopObservedWrappers, + stopObservationDelaysMs: [0], + sleep: vi.fn().mockResolvedValue(undefined), + }); + + await expect( + sandbox.stopWrappers({ + target: { + kind: 'instance', + instance: { instanceId: 'instance_gone', instanceGeneration: 1 }, + }, + attemptId: 'attempt_residual', + reason: 'session-delete', + }) + ).resolves.toEqual({ status: 'absent' }); + expect(stopObservedWrappers).toHaveBeenCalledWith(expect.anything(), 'agent_cloudflare', [ + { representation: 'process', id: 'wrapper-legacy', port: 5001 }, + ]); + }); + + it('force stops a targeted wrapper that remains after graceful termination and confirms absence', async () => { + const listProcesses = vi + .fn() + .mockResolvedValueOnce([ + { + id: 'wrapper-target', + command: + 'WRAPPER_PORT=5000 kilocode-wrapper --agent-session agent_cloudflare --wrapper-instance-id instance_1 --wrapper-instance-generation 2', + status: 'running', + }, + ]) + .mockResolvedValueOnce([ + { + id: 'wrapper-target', + command: + 'WRAPPER_PORT=5000 kilocode-wrapper --agent-session agent_cloudflare --wrapper-instance-id instance_1 --wrapper-instance-generation 2', + status: 'running', + }, + ]) + .mockResolvedValueOnce([]); + const exec = vi.fn().mockResolvedValue({ exitCode: 0, stdout: '' }); + const sandbox = new CloudflareAgentSandbox({} as Env, metadata(), { + resolveSandbox: () => ({ listProcesses, exec }) as unknown as SandboxInstance, + sleep: vi.fn().mockResolvedValue(undefined), + stopObservationDelaysMs: [0], + }); + + await expect( + sandbox.stopWrappers({ + target: { kind: 'instance', instance: { instanceId: 'instance_1', instanceGeneration: 2 } }, + attemptId: 'attempt_1', + reason: 'readiness-failed', + }) + ).resolves.toEqual({ status: 'absent', stoppedInstanceIds: ['instance_1'] }); + expect(exec).toHaveBeenCalledWith(expect.stringContaining('pkill -f --')); + expect(exec).toHaveBeenCalledWith(expect.stringContaining('pkill -9 -f --')); + expect(exec).toHaveBeenCalledWith(expect.stringContaining('--agent-session agent_cloudflare')); + }); + + it('returns still-present when targeted forceful cleanup remains observable', async () => { + const observedProcess = { + id: 'wrapper-target', + command: + 'WRAPPER_PORT=5000 kilocode-wrapper --agent-session agent_cloudflare --wrapper-instance-id instance_1 --wrapper-instance-generation 2', + status: 'running', + }; + const listProcesses = vi.fn().mockResolvedValue([observedProcess]); + const sandbox = new CloudflareAgentSandbox({} as Env, metadata(), { + resolveSandbox: () => + ({ + listProcesses, + exec: vi.fn().mockResolvedValue({ exitCode: 0, stdout: '' }), + }) as unknown as SandboxInstance, + sleep: vi.fn().mockResolvedValue(undefined), + stopObservationDelaysMs: [0], + }); + + await expect( + sandbox.stopWrappers({ + target: { kind: 'instance', instance: { instanceId: 'instance_1', instanceGeneration: 2 } }, + attemptId: 'attempt_remaining', + reason: 'readiness-failed', + }) + ).resolves.toMatchObject({ status: 'still-present' }); + }); + + it('returns inspection-failed from stop when post-stop inspection cannot prove absence', async () => { + const listProcesses = vi + .fn() + .mockResolvedValueOnce([ + { + id: 'wrapper-target', + command: + 'WRAPPER_PORT=5000 kilocode-wrapper --agent-session agent_cloudflare --wrapper-instance-id instance_1 --wrapper-instance-generation 2', + status: 'running', + }, + ]) + .mockRejectedValueOnce(new Error('cannot re-observe')); + const sandbox = new CloudflareAgentSandbox({} as Env, metadata(), { + resolveSandbox: () => + ({ + listProcesses, + exec: vi.fn().mockResolvedValue({ exitCode: 0 }), + }) as unknown as SandboxInstance, + sleep: vi.fn().mockResolvedValue(undefined), + stopObservationDelaysMs: [0], + }); + + await expect( + sandbox.stopWrappers({ + target: { kind: 'session' }, + attemptId: 'attempt_2', + reason: 'unexpected-wrapper', + }) + ).resolves.toMatchObject({ status: 'inspection-failed' }); + }); + + it('renews and health-checks the existing runtime', async () => { + const renewActivityTimeout = vi.fn().mockResolvedValue(undefined); + const listProcesses = vi.fn().mockResolvedValue([]); + const sandbox = new CloudflareAgentSandbox({} as Env, metadata(), { + resolveSandbox: () => ({ renewActivityTimeout, listProcesses }) as unknown as SandboxInstance, + }); + + await sandbox.keepAlive(); + await sandbox.probeHealth(); + + expect(renewActivityTimeout).toHaveBeenCalledOnce(); + expect(listProcesses).toHaveBeenCalledOnce(); + }); + + it('deletes session resources without destroying a shared sandbox', async () => { + const destroy = vi.fn(); + const deleteSession = vi.fn().mockResolvedValue(undefined); + const sandbox = new CloudflareAgentSandbox({} as Env, metadata(), { + resolveSandbox: () => + ({ + getSession: vi.fn().mockResolvedValue({ + exec: vi.fn().mockResolvedValue({ exitCode: 0, stdout: '', stderr: '' }), + }), + deleteSession, + destroy, + }) as unknown as SandboxInstance, + }); + + await sandbox.delete('explicit'); + + expect(deleteSession).toHaveBeenCalledWith('agent_cloudflare'); + expect(destroy).not.toHaveBeenCalled(); + }); + + it('maps infrastructure recovery to destructive Cloudflare sandbox replacement', async () => { + const destroy = vi.fn().mockResolvedValue(undefined); + const sandbox = new CloudflareAgentSandbox({} as Env, metadata(), { + resolveSandbox: () => ({ destroy }) as unknown as SandboxInstance, + }); + + await sandbox.delete('recovery'); + + expect(destroy).toHaveBeenCalledOnce(); + }); +}); diff --git a/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.ts b/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.ts new file mode 100644 index 0000000000..023113a43f --- /dev/null +++ b/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.ts @@ -0,0 +1,402 @@ +import type { + AgentSandbox, + EnsureWrapperRequest, + StopWrappersResult, + TerminalClientResult, + WrapperLogs, + WrapperObservation, + WrapperStopTarget, +} from '../protocol.js'; +import type { + Env, + SandboxId, + SandboxInstance, + SessionId as ServiceSessionId, +} from '../../types.js'; +import type { SessionMetadata } from '../../persistence/session-metadata.js'; +import type { SandboxDeleteReason, WrapperStopReason } from '../protocol.js'; +import { getSandbox } from '@cloudflare/sandbox'; +import { SANDBOX_SLEEP_AFTER_SECONDS } from '../../core/lease.js'; +import { generateSandboxId, getSandboxNamespace } from '../../sandbox-id.js'; +import { SessionService } from '../../session-service.js'; +import { WrapperClient, WrapperContainerClient } from '../../kilo/wrapper-client.js'; +import { + discoverSessionWrappers, + findWrapperForSession, + stopObservedWrappers, +} from '../../kilo/wrapper-manager.js'; +import { + checkDiskAndCleanBeforeSetup, + cleanupWorkspace, + getSessionHomePath, + getSessionWorkspacePath, +} from '../../workspace.js'; +import { + FAST_SANDBOX_COMMAND_TIMEOUT_MS, + logSandboxOperationTimeout, + timedExec, +} from '../../sandbox-timeout-logging.js'; +import { SANDBOX_WORKSPACE_PROBE_TIMEOUT_MESSAGE } from '../../sandbox-recovery.js'; +import { withTimeout } from '@kilocode/worker-utils'; +import { WRAPPER_VERSION } from '../../shared/wrapper-version.js'; +import { ExecutionError } from '../../execution/errors.js'; +import { + isSandboxFilesystemUnusableError, + SandboxCapacityInspectionError, +} from '../../workspace-errors.js'; + +const PREPARE_WORKSPACE_TIMEOUT_MS = 10 * 60 * 1000; +const DEFAULT_STOP_OBSERVATION_DELAYS_MS = [100, 500, 1_000]; + +function withWorkspacePreparationTimeout(operation: Promise, step: string): Promise { + return withTimeout( + operation, + PREPARE_WORKSPACE_TIMEOUT_MS, + `Workspace preparation timed out during ${step} after ${PREPARE_WORKSPACE_TIMEOUT_MS / 1000}s`, + () => + logSandboxOperationTimeout({ + operation: `workspace.prepare:${step}`, + timeoutMs: PREPARE_WORKSPACE_TIMEOUT_MS, + timeoutLayer: 'outer', + }) + ); +} + +export type CloudflareAgentSandboxDependencies = { + resolveSandbox?: (sandboxId: SandboxId, options?: { sleepAfter?: number }) => SandboxInstance; + sessionService?: SessionService; + stopObservedWrappers?: typeof stopObservedWrappers; + sleep?: (ms: number) => Promise; + stopObservationDelaysMs?: number[]; +}; + +export class CloudflareAgentSandbox implements AgentSandbox { + private readonly sessionService: SessionService; + private readonly resolveSandbox: ( + sandboxId: SandboxId, + options?: { sleepAfter?: number } + ) => SandboxInstance; + private readonly stopObserved: typeof stopObservedWrappers; + private readonly sleep: (ms: number) => Promise; + private readonly stopObservationDelaysMs: number[]; + private sandboxIdPromise?: Promise; + + constructor( + private readonly env: Env, + private readonly metadata: SessionMetadata, + dependencies: CloudflareAgentSandboxDependencies = {} + ) { + this.sessionService = dependencies.sessionService ?? new SessionService(); + this.resolveSandbox = + dependencies.resolveSandbox ?? + ((sandboxId, options) => + options + ? getSandbox(getSandboxNamespace(this.env, sandboxId), sandboxId, options) + : getSandbox(getSandboxNamespace(this.env, sandboxId), sandboxId)); + this.stopObserved = dependencies.stopObservedWrappers ?? stopObservedWrappers; + this.sleep = dependencies.sleep ?? (ms => new Promise(resolve => setTimeout(resolve, ms))); + this.stopObservationDelaysMs = + dependencies.stopObservationDelaysMs ?? DEFAULT_STOP_OBSERVATION_DELAYS_MS; + } + + private resolveSandboxId(): Promise { + if (!this.sandboxIdPromise) { + this.sandboxIdPromise = this.metadata.workspace?.sandboxId + ? Promise.resolve(this.metadata.workspace.sandboxId) + : generateSandboxId( + this.env.PER_SESSION_SANDBOX_ORG_IDS, + this.metadata.identity.orgId, + this.metadata.identity.userId, + this.metadata.identity.sessionId, + this.metadata.identity.botId + ); + } + return this.sandboxIdPromise; + } + + private async getSandbox(options?: { sleepAfter?: number }): Promise { + return this.resolveSandbox(await this.resolveSandboxId(), options); + } + + private async workspaceHasGit(sandbox: SandboxInstance, workspacePath: string): Promise { + const timeoutMs = FAST_SANDBOX_COMMAND_TIMEOUT_MS; + try { + const result = await withTimeout( + timedExec( + sandbox, + `test -d '${workspacePath}/.git' && echo exists`, + 'execution.wrapperBootstrap.repoExists' + ), + timeoutMs, + `${SANDBOX_WORKSPACE_PROBE_TIMEOUT_MESSAGE} after ${timeoutMs}ms`, + () => + logSandboxOperationTimeout({ + operation: 'execution.wrapperBootstrap.repoExists', + timeoutMs, + timeoutLayer: 'outer', + }) + ); + if (result.exitCode !== 0 && isSandboxFilesystemUnusableError(result.stderr)) { + throw new SandboxCapacityInspectionError( + 'Workspace admission probe cannot run because the sandbox filesystem is unusable', + new Error(result.stderr) + ); + } + return result.stdout?.includes('exists') ?? false; + } catch (error) { + if (isSandboxFilesystemUnusableError(error)) { + throw new SandboxCapacityInspectionError( + 'Workspace admission probe cannot run because the sandbox filesystem is unusable', + error + ); + } + throw error; + } + } + + private requiresPreparedDevcontainerRuntime(request: EnsureWrapperRequest): boolean { + return ( + request.plan.workspace.metadata.workspace?.devcontainerRequested === true || + request.plan.workspace.metadata.devcontainer !== undefined + ); + } + + private usesDevcontainerRuntime(): boolean { + return ( + this.metadata.workspace?.sandboxId?.startsWith('dind-') === true || + this.metadata.workspace?.devcontainerRequested === true || + this.metadata.devcontainer !== undefined + ); + } + + private existingWrapperSessionName(): string { + const sessionId = this.metadata.identity.sessionId; + return this.usesDevcontainerRuntime() ? sessionId : `${sessionId}-bootstrap`; + } + + async ensureWrapper(request: EnsureWrapperRequest) { + const { plan, prepared } = request; + const { sessionId, userId, orgId } = plan.scope; + this.sandboxIdPromise = Promise.resolve(plan.workspace.sandboxId as SandboxId); + const sandboxId = await this.resolveSandboxId(); + const sandbox = await this.getSandbox({ sleepAfter: SANDBOX_SLEEP_AFTER_SECONDS }); + + if (this.requiresPreparedDevcontainerRuntime(request)) { + const preparedWorkspace = await withWorkspacePreparationTimeout( + this.sessionService.prepareWorkspace({ + sandbox, + sandboxId, + orgId, + userId, + sessionId: sessionId as ServiceSessionId, + kilocodeModel: plan.agent.model, + env: this.env, + metadata: plan.workspace.metadata, + onProgress: request.onProgress, + }), + 'devcontainer workspace preparation' + ); + if (!preparedWorkspace.devcontainer || !preparedWorkspace.ready.devcontainer) { + throw ExecutionError.workspaceSetupFailed( + 'Devcontainer workspace preparation did not resolve runtime metadata' + ); + } + let wrapper: Awaited>; + try { + wrapper = await WrapperClient.ensureWrapper(sandbox, preparedWorkspace.session, { + agentSessionId: sessionId, + userId, + workspacePath: preparedWorkspace.context.workspacePath, + sessionId: plan.wrapper.kiloSessionId, + runtimeEnv: preparedWorkspace.runtimeEnv, + devcontainer: preparedWorkspace.devcontainer, + fixedPort: preparedWorkspace.ready.devcontainer.wrapperPort, + ...(request.leasedInstance ? { leasedInstance: request.leasedInstance } : {}), + }); + } catch (error) { + throw ExecutionError.wrapperStartFailed( + `Failed to start devcontainer wrapper: ${error instanceof Error ? error.message : String(error)}`, + error + ); + } + return { + status: 'session-ready' as const, + client: wrapper.client, + ready: preparedWorkspace.ready, + kiloSessionId: wrapper.sessionId, + }; + } + + const workspaceWarm = await this.workspaceHasGit(sandbox, prepared.context.workspacePath); + if (!workspaceWarm) { + request.onProgress?.('disk_check', 'Checking disk space...'); + await checkDiskAndCleanBeforeSetup(sandbox, orgId, userId, sessionId); + } + request.onProgress?.('kilo_server', 'Starting Kilo...'); + const bootstrapSession = await sandbox.createSession({ + name: `${sessionId}-bootstrap`, + env: {}, + cwd: '/', + }); + const wrapper = await WrapperClient.ensureBootstrapWrapper(sandbox, bootstrapSession, { + agentSessionId: sessionId, + userId, + ...(request.leasedInstance ? { leasedInstance: request.leasedInstance } : {}), + }); + return { status: 'wrapper-running' as const, client: wrapper.client }; + } + + async discoverSessionWrappers(): Promise { + return discoverSessionWrappers(await this.getSandbox(), this.metadata.identity.sessionId, { + inspectContainers: this.usesDevcontainerRuntime(), + }); + } + + private async observeTarget(_target: WrapperStopTarget): Promise { + // The lease is session-scoped: confirming absence must account for every + // physical wrapper carrying this logical session marker, including duplicates. + return this.discoverSessionWrappers(); + } + + async stopWrappers(request: { + target: WrapperStopTarget; + attemptId: string; + reason: WrapperStopReason; + }): Promise { + const sandbox = await this.getSandbox(); + const initial = await this.observeTarget(request.target); + if (initial.status !== 'present') return initial; + + try { + await this.stopObserved(sandbox, this.metadata.identity.sessionId, initial.observed); + } catch (error) { + return { status: 'still-present', observed: initial.observed, error: String(error) }; + } + + let latest: WrapperObservation = initial; + for (const delayMs of this.stopObservationDelaysMs) { + await this.sleep(delayMs); + latest = await this.observeTarget(request.target); + if (latest.status !== 'present') return latest; + } + + try { + await this.stopObserved(sandbox, this.metadata.identity.sessionId, latest.observed, { + force: true, + }); + } catch (error) { + return { status: 'still-present', observed: latest.observed, error: String(error) }; + } + + const final = await this.observeTarget(request.target); + if (final.status === 'inspection-failed') return final; + if (final.status === 'present') return { status: 'still-present', observed: final.observed }; + const stoppedInstanceIds = initial.observed.flatMap(observed => + observed.instanceId ? [observed.instanceId] : [] + ); + return stoppedInstanceIds.length > 0 ? { status: 'absent', stoppedInstanceIds } : final; + } + + async probeHealth(): Promise { + const sandbox = await this.getSandbox(); + await sandbox.listProcesses(); + } + + async getRunningWrapper(): Promise { + const sandbox = await this.getSandbox({ sleepAfter: SANDBOX_SLEEP_AFTER_SECONDS }); + const wrapper = await findWrapperForSession(sandbox, this.metadata.identity.sessionId); + if (!wrapper) return null; + const session = await sandbox.getSession(this.existingWrapperSessionName()); + return new WrapperClient({ session, port: wrapper.port }); + } + + async getRunningTerminalClient(): Promise { + const sandbox = await this.getSandbox({ sleepAfter: SANDBOX_SLEEP_AFTER_SECONDS }); + const wrapper = await findWrapperForSession(sandbox, this.metadata.identity.sessionId); + if (!wrapper) return { status: 'not-running' }; + const client = new WrapperContainerClient({ sandbox, port: wrapper.port }); + try { + const health = await client.health(); + if (!health.healthy || health.version !== WRAPPER_VERSION) return { status: 'unhealthy' }; + } catch { + return { status: 'unhealthy' }; + } + return { status: 'ready', client }; + } + + async readWrapperLogs(): Promise { + const sandbox = await this.getSandbox({ sleepAfter: SANDBOX_SLEEP_AFTER_SECONDS }); + const session = await sandbox.getSession(this.existingWrapperSessionName()); + const logPaths: string[] = []; + const wrapperFiles = await session.listFiles('/tmp').catch(() => undefined); + if (wrapperFiles?.success) { + for (const file of wrapperFiles.files) { + if ( + file.type === 'file' && + file.name.startsWith('kilocode-wrapper-') && + file.name.endsWith('.log') + ) { + logPaths.push(file.absolutePath); + } + } + } + const sessionHome = getSessionHomePath(this.metadata.identity.sessionId); + const cliFiles = await session + .listFiles(`${sessionHome}/.local/share/kilo/log`, { recursive: true }) + .catch(() => undefined); + if (cliFiles?.success) { + for (const file of cliFiles.files) { + if (file.type === 'file') logPaths.push(file.absolutePath); + } + } + const files: Record = {}; + const contents = await Promise.allSettled( + logPaths.map(async path => ({ + path, + content: (await session.readFile(path, { encoding: 'utf-8' })).content, + })) + ); + for (const content of contents) { + if (content.status === 'fulfilled') files[content.value.path] = content.value.content; + } + let processes: WrapperLogs['processes']; + try { + processes = (await sandbox.listProcesses()).map(process => ({ + pid: Number.parseInt(process.id, 10) || 0, + command: process.command, + status: process.status, + })); + } catch { + processes = undefined; + } + return { files, processes }; + } + + async keepAlive(): Promise { + const sandbox = await this.getSandbox(); + await Promise.resolve(sandbox.renewActivityTimeout()); + } + + async delete(reason: SandboxDeleteReason): Promise { + const sandbox = await this.getSandbox(); + if (reason === 'recovery') { + await sandbox.destroy(); + return; + } + try { + const session = await sandbox.getSession(this.metadata.identity.sessionId); + await cleanupWorkspace( + session, + getSessionWorkspacePath( + this.metadata.identity.orgId, + this.metadata.identity.userId, + this.metadata.identity.sessionId + ), + getSessionHomePath(this.metadata.identity.sessionId) + ); + } catch { + // Cleanup remains best effort before session resource deletion. + } + await sandbox.deleteSession(this.metadata.identity.sessionId); + } +} diff --git a/services/cloud-agent-next/src/agent-sandbox/factory.test.ts b/services/cloud-agent-next/src/agent-sandbox/factory.test.ts new file mode 100644 index 0000000000..2dde48278b --- /dev/null +++ b/services/cloud-agent-next/src/agent-sandbox/factory.test.ts @@ -0,0 +1,23 @@ +import { describe, expect, it, vi } from 'vitest'; +import type { Env } from '../types.js'; +import type { SessionMetadata } from '../persistence/session-metadata.js'; +import { createAgentSandbox } from './factory.js'; +import { CloudflareAgentSandbox } from './cloudflare/cloudflare-agent-sandbox.js'; + +vi.mock('@cloudflare/sandbox', () => ({ getSandbox: vi.fn() })); + +function metadata(): SessionMetadata { + return { + metadataSchemaVersion: 2, + identity: { sessionId: 'agent_sandbox', userId: 'user_sandbox' }, + auth: {}, + workspace: { sandboxId: 'ses-abcdef' }, + lifecycle: { version: 1, timestamp: 1 }, + }; +} + +describe('AgentSandbox factory', () => { + it('constructs the Cloudflare runtime adapter', () => { + expect(createAgentSandbox({} as Env, metadata())).toBeInstanceOf(CloudflareAgentSandbox); + }); +}); diff --git a/services/cloud-agent-next/src/agent-sandbox/factory.ts b/services/cloud-agent-next/src/agent-sandbox/factory.ts new file mode 100644 index 0000000000..b542f10b45 --- /dev/null +++ b/services/cloud-agent-next/src/agent-sandbox/factory.ts @@ -0,0 +1,8 @@ +import type { SessionMetadata } from '../persistence/session-metadata.js'; +import type { Env } from '../types.js'; +import type { AgentSandbox } from './protocol.js'; +import { CloudflareAgentSandbox } from './cloudflare/cloudflare-agent-sandbox.js'; + +export function createAgentSandbox(env: Env, metadata: SessionMetadata): AgentSandbox { + return new CloudflareAgentSandbox(env, metadata); +} diff --git a/services/cloud-agent-next/src/agent-sandbox/protocol.ts b/services/cloud-agent-next/src/agent-sandbox/protocol.ts new file mode 100644 index 0000000000..e8f874c221 --- /dev/null +++ b/services/cloud-agent-next/src/agent-sandbox/protocol.ts @@ -0,0 +1,101 @@ +import type { WrapperClient } from '../kilo/wrapper-client.js'; +import type { TerminalWrapperClient } from '../terminal/access.js'; +import type { + FencedLegacyExecutionRequest, + FencedWrapperDispatchRequest, + WorkspaceReady, +} from '../execution/types.js'; + +export type SandboxDeleteReason = 'explicit' | 'retention-expired' | 'recovery'; + +export type WrapperInstanceLease = { + instanceId: string; + instanceGeneration: number; +}; + +export type ObservedWrapper = { + representation: 'process' | 'container'; + id: string; + port?: number; + instanceId?: string; + instanceGeneration?: number; +}; + +export type WrapperObservation = + | { status: 'absent' } + | { status: 'present'; observed: ObservedWrapper[] } + | { status: 'inspection-failed'; error: string }; + +export type WrapperStopTarget = + | { kind: 'instance'; instance: WrapperInstanceLease } + | { kind: 'session' }; + +export type WrapperStopReason = + | 'readiness-failed' + | 'startup-failed' + | 'unhealthy-wrapper' + | 'terminal-failed' + | 'terminal-interrupted' + | 'idle-timeout' + | 'keep-warm-expired' + | 'user-interrupt' + | 'session-delete' + | 'unexpected-wrapper' + | 'observation-failed'; + +export type StopWrappersResult = + | { status: 'absent'; stoppedInstanceIds?: string[] } + | { status: 'still-present'; observed: ObservedWrapper[]; error?: string } + | { status: 'inspection-failed'; error: string }; + +export type TerminalClientResult = + | { status: 'ready'; client: TerminalWrapperClient } + | { status: 'not-running' } + | { status: 'unhealthy' }; + +export type WrapperLogs = { + files: Record; + processes?: Array<{ pid: number; command: string; status: string }>; +}; + +export type EnsureWrapperRequest = { + plan: FencedWrapperDispatchRequest | FencedLegacyExecutionRequest; + leasedInstance?: WrapperInstanceLease; + prepared: { + ready: WorkspaceReady; + context: { workspacePath: string }; + }; + onProgress?: (step: string, message: string) => void; +}; + +export type EnsuredWrapper = + | { + status: 'wrapper-running'; + client: WrapperClient; + } + | { + status: 'session-ready'; + client: WrapperClient; + ready: WorkspaceReady; + kiloSessionId: string; + }; + +/** + * Product-specific runtime seam for one Cloud Agent session. + * Provider process, filesystem, and raw sandbox APIs remain private to adapters. + */ +export type AgentSandbox = { + ensureWrapper(request: EnsureWrapperRequest): Promise; + discoverSessionWrappers(): Promise; + stopWrappers(request: { + target: WrapperStopTarget; + attemptId: string; + reason: WrapperStopReason; + }): Promise; + probeHealth(): Promise; + getRunningWrapper(): Promise; + getRunningTerminalClient(): Promise; + readWrapperLogs(): Promise; + keepAlive(): Promise; + delete(reason: SandboxDeleteReason): Promise; +}; diff --git a/services/cloud-agent-next/src/execution/orchestrator.test.ts b/services/cloud-agent-next/src/execution/orchestrator.test.ts index 95d7410a7d..85c0f6894a 100644 --- a/services/cloud-agent-next/src/execution/orchestrator.test.ts +++ b/services/cloud-agent-next/src/execution/orchestrator.test.ts @@ -1,87 +1,44 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'; -import type * as WrapperClientModule from '../kilo/wrapper-client.js'; -import type { Env, SandboxInstance } from '../types.js'; +import type { AgentSandbox, WrapperInstanceLease } from '../agent-sandbox/protocol.js'; +import type { Env } from '../types.js'; +import type { ExecutionError } from './errors.js'; import type { FencedWrapperDispatchRequest } from './types.js'; +import { + SandboxCapacityInspectionError, + WorkspaceCapacityAdmissionRejectedError, +} from '../workspace-errors.js'; -const { - ensureBootstrapWrapperMock, - ensureWrapperMock, - ensureSessionReadyMock, - promptMock, - commandMock, - buildWrapperSessionReadyAndPromptRequestsMock, - prepareWorkspaceMock, -} = vi.hoisted(() => ({ - ensureBootstrapWrapperMock: vi.fn(), - ensureWrapperMock: vi.fn(), - ensureSessionReadyMock: vi.fn(), - promptMock: vi.fn(), - commandMock: vi.fn(), +const { buildWrapperSessionReadyAndPromptRequestsMock } = vi.hoisted(() => ({ buildWrapperSessionReadyAndPromptRequestsMock: vi.fn(), - prepareWorkspaceMock: vi.fn(), })); vi.mock('../session-service.js', () => ({ SessionService: class SessionService { buildWrapperSessionReadyAndPromptRequests = buildWrapperSessionReadyAndPromptRequestsMock; - prepareWorkspace = prepareWorkspaceMock; }, })); -vi.mock('../kilo/wrapper-client.js', async importActual => { - const actual = await importActual(); - return { - ...actual, - WrapperClient: { - ensureBootstrapWrapper: ensureBootstrapWrapperMock, - ensureWrapper: ensureWrapperMock, - }, - }; -}); - import { ExecutionOrchestrator } from './orchestrator.js'; const baseMetadata = { metadataSchemaVersion: 2, - identity: { - sessionId: 'agent_test', - userId: 'user_test', - }, - auth: { - kiloSessionId: 'kilo_existing', - kilocodeToken: 'kilo_token', - }, - lifecycle: { - version: 1, - timestamp: 1, - }, + identity: { sessionId: 'agent_test', userId: 'user_test' }, + auth: { kiloSessionId: 'kilo_existing', kilocodeToken: 'kilo_token' }, + lifecycle: { version: 1, timestamp: 1 }, } satisfies FencedWrapperDispatchRequest['workspace']['metadata']; const basePlan = { - scope: { - sessionId: 'agent_test', - userId: 'user_test', - }, + scope: { sessionId: 'agent_test', userId: 'user_test' }, turn: { type: 'prompt', messageId: 'msg_018f1e2d3c4bOrchestratorAAAA', prompt: 'Review this change', }, - agent: { - mode: 'code', - model: 'test-model', - }, - workspace: { - sandboxId: 'sandbox_test', - metadata: baseMetadata, - }, + agent: { mode: 'code', model: 'test-model' }, + workspace: { sandboxId: 'sandbox_test', metadata: baseMetadata }, wrapper: { kiloSessionId: 'kilo_existing', - fence: { - wrapperRunId: 'wr_test', - wrapperGeneration: 1, - wrapperConnectionId: 'conn_test', - }, + fence: { wrapperRunId: 'wr_test', wrapperGeneration: 1, wrapperConnectionId: 'conn_test' }, }, } satisfies FencedWrapperDispatchRequest; @@ -100,7 +57,6 @@ function buildPreparedRequests() { branchName: 'session/agent_test', kiloSessionId: 'kilo_existing', }; - return { type: 'prompt' as const, readyRequest: { @@ -117,214 +73,121 @@ function buildPreparedRequests() { session, }, promptRequest: { - message: { - id: basePlan.turn.messageId, - prompt: basePlan.turn.prompt, - }, - agent: { - mode: 'code', - }, + message: { id: basePlan.turn.messageId, prompt: basePlan.turn.prompt }, + agent: { mode: 'code' }, session, }, ready, - context: { - workspacePath: '/workspace/test', - }, + context: { workspacePath: '/workspace/test' }, }; } -function createOrchestrator() { - const sandbox = { - exec: vi.fn().mockResolvedValue({ exitCode: 0, stdout: 'exists\n', stderr: '' }), - createSession: vi.fn().mockResolvedValue({}), - } as unknown as SandboxInstance; +function createOrchestrator(options?: { sessionReady?: boolean }) { + const ensureSessionReady = vi.fn().mockResolvedValue({ kiloSessionId: 'kilo_ready' }); + const prompt = vi.fn().mockResolvedValue({ messageId: basePlan.turn.messageId }); + const command = vi.fn().mockResolvedValue({}); + const wrapper = { ensureSessionReady, prompt, command }; + const ensureWrapper = vi.fn().mockResolvedValue( + options?.sessionReady + ? { + status: 'session-ready', + client: wrapper, + ready: buildPreparedRequests().ready, + kiloSessionId: 'kilo_devcontainer', + } + : { status: 'wrapper-running', client: wrapper } + ); + const deleteSandbox = vi.fn().mockResolvedValue(undefined); + const agentSandbox = { + ensureWrapper, + delete: deleteSandbox, + } as unknown as AgentSandbox; const recordKiloServerActivity = vi.fn().mockResolvedValue(undefined); - const orchestrator = new ExecutionOrchestrator({ - getSandbox: vi.fn().mockResolvedValue(sandbox), + getAgentSandbox: vi.fn().mockReturnValue(agentSandbox), getSessionStub: vi.fn().mockReturnValue({ recordKiloServerActivity }), env: {} as Env, }); - - return { orchestrator, sandbox }; + return { + orchestrator, + ensureWrapper, + deleteSandbox, + ensureSessionReady, + prompt, + command, + }; } -describe('ExecutionOrchestrator split wrapper bootstrap', () => { +describe('ExecutionOrchestrator AgentSandbox delivery', () => { beforeEach(() => { vi.clearAllMocks(); - buildWrapperSessionReadyAndPromptRequestsMock.mockImplementation(async () => - buildPreparedRequests() - ); - ensureSessionReadyMock.mockResolvedValue({ kiloSessionId: 'kilo_ready' }); - promptMock.mockResolvedValue({ messageId: basePlan.turn.messageId }); - commandMock.mockResolvedValue({}); - ensureBootstrapWrapperMock.mockResolvedValue({ - client: { - ensureSessionReady: ensureSessionReadyMock, - prompt: promptMock, - command: commandMock, - }, - }); - ensureWrapperMock.mockResolvedValue({ - client: { - prompt: promptMock, - command: commandMock, - }, - sessionId: 'kilo_devcontainer', - }); + buildWrapperSessionReadyAndPromptRequestsMock.mockResolvedValue(buildPreparedRequests()); }); - it('disables interactive tools from current code-review session metadata', async () => { - const { orchestrator } = createOrchestrator(); - const plan = { - ...basePlan, - workspace: { - ...basePlan.workspace, - metadata: { - ...baseMetadata, - identity: { - ...baseMetadata.identity, - createdOnPlatform: 'code-review', - }, - }, - }, - } satisfies FencedWrapperDispatchRequest; + it('readies a wrapper before dispatching its prompt', async () => { + const { orchestrator, ensureWrapper, ensureSessionReady, prompt } = createOrchestrator(); + const onWorkspaceReady = vi.fn().mockResolvedValue(undefined); - await orchestrator.execute(plan); + await expect(orchestrator.execute(basePlan, { onWorkspaceReady })).resolves.toEqual({ + kiloSessionId: 'kilo_ready', + }); - expect(promptMock).toHaveBeenCalledWith( - expect.objectContaining({ - agent: expect.objectContaining({ - tools: { - question: false, - plan_enter: false, - plan_exit: false, - }, - }), - }) + expect(ensureWrapper).toHaveBeenCalledWith( + expect.objectContaining({ plan: basePlan, prepared: buildPreparedRequests() }) + ); + expect(ensureSessionReady).toHaveBeenCalledWith(buildPreparedRequests().readyRequest); + expect(ensureSessionReady.mock.invocationCallOrder[0]).toBeLessThan( + prompt.mock.invocationCallOrder[0] ); + expect(onWorkspaceReady).toHaveBeenCalledWith(buildPreparedRequests().ready); }); - it('does not send tool overrides for ordinary message deliveries', async () => { - const { orchestrator } = createOrchestrator(); - - await orchestrator.execute(basePlan); - - const promptRequest = promptMock.mock.calls[0]?.[0] as { - agent?: { tools?: Record }; + it('passes an authorized physical lease through to wrapper startup', async () => { + const { orchestrator, ensureWrapper } = createOrchestrator(); + const leasedInstance: WrapperInstanceLease = { + instanceId: 'instance_orchestrator', + instanceGeneration: 9, }; - expect(promptRequest.agent?.tools).toBeUndefined(); - }); - it('readies the session separately before delivering the grouped prompt', async () => { - const prepared = buildPreparedRequests(); - buildWrapperSessionReadyAndPromptRequestsMock.mockResolvedValueOnce(prepared); - const { orchestrator, sandbox } = createOrchestrator(); - const onWorkspaceReady = vi.fn().mockResolvedValue(undefined); + await orchestrator.execute(basePlan, { leasedInstance }); - await expect(orchestrator.execute(basePlan, { onWorkspaceReady })).resolves.toEqual({ - kiloSessionId: 'kilo_ready', - }); - - expect(ensureBootstrapWrapperMock).toHaveBeenCalledWith(sandbox, expect.anything(), { - agentSessionId: 'agent_test', - userId: 'user_test', - }); - expect(ensureSessionReadyMock).toHaveBeenCalledWith(prepared.readyRequest); - expect(onWorkspaceReady).toHaveBeenCalledWith(prepared.ready); - expect(promptMock).toHaveBeenCalledWith(prepared.promptRequest); + expect(ensureWrapper).toHaveBeenCalledWith(expect.objectContaining({ leasedInstance })); }); - it('forwards wrapper-resolved devcontainer readiness to persistence hooks', async () => { - const prepared = buildPreparedRequests(); - const devcontainer = { - workspacePath: '/workspace/test', - innerWorkspaceFolder: '/workspaces/repo', - wrapperPort: 4173, - configPath: '.devcontainer/devcontainer.json', - }; - buildWrapperSessionReadyAndPromptRequestsMock.mockResolvedValueOnce(prepared); - ensureSessionReadyMock.mockResolvedValueOnce({ - kiloSessionId: 'kilo_ready', - workspaceReady: { - ...prepared.ready, - devcontainer, - }, - }); - const { orchestrator } = createOrchestrator(); + it('uses already-ready devcontainer adapter output without a second ready request', async () => { + const { orchestrator, ensureSessionReady, prompt } = createOrchestrator({ sessionReady: true }); const onWorkspaceReady = vi.fn().mockResolvedValue(undefined); - await orchestrator.execute(basePlan, { onWorkspaceReady }); - - expect(onWorkspaceReady).toHaveBeenCalledWith({ - ...prepared.ready, - devcontainer, + await expect(orchestrator.execute(basePlan, { onWorkspaceReady })).resolves.toEqual({ + kiloSessionId: 'kilo_devcontainer', }); + expect(ensureSessionReady).not.toHaveBeenCalled(); + expect(onWorkspaceReady).toHaveBeenCalledWith(buildPreparedRequests().ready); + expect(prompt).toHaveBeenCalledWith(buildPreparedRequests().promptRequest); }); - it('routes devcontainer-requested deliveries through the prepared devcontainer runtime path', async () => { - const prepared = buildPreparedRequests(); - const devcontainer = { - workspacePath: '/workspace/test', - innerWorkspaceFolder: '/workspaces/repo', - wrapperPort: 4173, - configPath: '.devcontainer/devcontainer.json', - }; - const devcontainerPlan = { + it('disables interactive tools from current code-review session metadata', async () => { + const { orchestrator, prompt } = createOrchestrator(); + const plan = { ...basePlan, workspace: { - sandboxId: 'dind-abcdef', + ...basePlan.workspace, metadata: { ...baseMetadata, - workspace: { - sandboxId: 'dind-abcdef', - devcontainerRequested: true, - }, + identity: { ...baseMetadata.identity, createdOnPlatform: 'code-review' }, }, }, } satisfies FencedWrapperDispatchRequest; - buildWrapperSessionReadyAndPromptRequestsMock.mockResolvedValueOnce(prepared); - prepareWorkspaceMock.mockResolvedValueOnce({ - ready: { - ...prepared.ready, - sandboxId: 'dind-abcdef', - devcontainer, - }, - context: { - workspacePath: '/workspace/test', - }, - runtimeEnv: {}, - session: {}, - devcontainer: { - containerId: 'container-dev', - innerWorkspaceFolder: devcontainer.innerWorkspaceFolder, - workspacePath: devcontainer.workspacePath, - agentSessionId: 'agent_test', - overrideConfigPath: '/tmp/devcontainer.json', - teardown: vi.fn(), - }, - }); - const { orchestrator, sandbox } = createOrchestrator(); - const onWorkspaceReady = vi.fn().mockResolvedValue(undefined); - await expect(orchestrator.execute(devcontainerPlan, { onWorkspaceReady })).resolves.toEqual({ - kiloSessionId: 'kilo_devcontainer', - }); + await orchestrator.execute(plan); - expect(ensureBootstrapWrapperMock).not.toHaveBeenCalled(); - expect(ensureSessionReadyMock).not.toHaveBeenCalled(); - expect(ensureWrapperMock).toHaveBeenCalledWith( - sandbox, - expect.anything(), + expect(prompt).toHaveBeenCalledWith( expect.objectContaining({ - agentSessionId: 'agent_test', - fixedPort: devcontainer.wrapperPort, + agent: expect.objectContaining({ + tools: { question: false, plan_enter: false, plan_exit: false }, + }), }) ); - expect(onWorkspaceReady).toHaveBeenCalledWith( - expect.objectContaining({ sandboxId: 'dind-abcdef', devcontainer }) - ); - expect(promptMock).toHaveBeenCalledWith(prepared.promptRequest); }); it('dispatches queued command turns through wrapper command delivery', async () => { @@ -355,11 +218,58 @@ describe('ExecutionOrchestrator split wrapper bootstrap', () => { type: 'command', commandRequest, }); - const { orchestrator } = createOrchestrator(); + const { orchestrator, command, prompt } = createOrchestrator(); await orchestrator.execute(commandPlan); - expect(commandMock).toHaveBeenCalledWith(commandRequest); - expect(promptMock).not.toHaveBeenCalled(); + expect(command).toHaveBeenCalledWith(commandRequest); + expect(prompt).not.toHaveBeenCalled(); + }); + + it('keeps ordinary wrapper bootstrap failure retryable', async () => { + const { orchestrator, ensureWrapper } = createOrchestrator(); + ensureWrapper.mockRejectedValueOnce(new Error('wrapper unavailable')); + + await expect(orchestrator.execute(basePlan)).rejects.toMatchObject({ + code: 'WRAPPER_START_FAILED', + retryable: true, + } satisfies Partial); + }); + + it('does not recover the shared sandbox for plain capacity admission rejection', async () => { + const { orchestrator, ensureWrapper, deleteSandbox } = createOrchestrator(); + ensureWrapper.mockRejectedValueOnce( + new WorkspaceCapacityAdmissionRejectedError({ + availableMB: 1024, + thresholdMB: 2048, + cleaned: 0, + skipped: 1, + }) + ); + + await expect(orchestrator.execute(basePlan)).rejects.toThrow('Failed to start wrapper'); + expect(deleteSandbox).not.toHaveBeenCalled(); + }); + + it('requests provider-neutral recovery cleanup on unusable capacity inspection', async () => { + const { orchestrator, ensureWrapper, deleteSandbox } = createOrchestrator(); + ensureWrapper.mockRejectedValueOnce( + new SandboxCapacityInspectionError('Cannot inspect capacity', new Error('ENOSPC')) + ); + + await expect(orchestrator.execute(basePlan)).rejects.toThrow('Failed to start wrapper'); + expect(deleteSandbox).toHaveBeenCalledWith('recovery'); + }); + + it('requests provider-neutral recovery cleanup on infrastructure preparation failure', async () => { + const { orchestrator, ensureWrapper, deleteSandbox } = createOrchestrator(); + const sandboxError = Object.assign(new Error('HTTP Error! status: 500'), { + name: 'SandboxError', + httpStatus: 500, + }); + ensureWrapper.mockRejectedValueOnce(sandboxError); + + await expect(orchestrator.execute(basePlan)).rejects.toThrow('Failed to start wrapper'); + expect(deleteSandbox).toHaveBeenCalledWith('recovery'); }); }); diff --git a/services/cloud-agent-next/src/execution/orchestrator.ts b/services/cloud-agent-next/src/execution/orchestrator.ts index e212723c7a..0531fdd278 100644 --- a/services/cloud-agent-next/src/execution/orchestrator.ts +++ b/services/cloud-agent-next/src/execution/orchestrator.ts @@ -1,16 +1,12 @@ /** - * ExecutionOrchestrator - Handles prompt execution. + * ExecutionOrchestrator - Handles provider-neutral wrapper delivery. * - * This module handles workspace preparation and execution, called directly - * from the DO when a client sends a prompt. + * AgentSandbox obtains the usable runtime/wrapper. This module preserves the + * business sequence: prepare requests, ready the session when required, then + * hand the accepted prompt or command to the wrapper. */ -import type { - Env, - SandboxInstance, - SandboxId as ServiceSandboxId, - SessionId as ServiceSessionId, -} from '../types.js'; +import type { Env } from '../types.js'; import type { CloudAgentSession } from '../persistence/CloudAgentSession.js'; import type { ExecutionResult, @@ -21,21 +17,14 @@ import type { import { ExecutionError } from './errors.js'; import { SessionService } from '../session-service.js'; import { logger } from '../logger.js'; -import { - FAST_SANDBOX_COMMAND_TIMEOUT_MS, - logSandboxOperationTimeout, - timedExec, -} from '../sandbox-timeout-logging.js'; -import { WrapperClient, WrapperError } from '../kilo/wrapper-client.js'; +import { WrapperError } from '../kilo/wrapper-client.js'; import { withDORetry } from '../utils/do-retry.js'; import { withTimeout } from '@kilocode/worker-utils'; -import { - SANDBOX_WORKSPACE_PROBE_TIMEOUT_MESSAGE, - withPreparationInfrastructureRecovery, -} from '../sandbox-recovery.js'; -import { checkDiskAndCleanBeforeSetup } from '../workspace.js'; +import { logSandboxOperationTimeout } from '../sandbox-timeout-logging.js'; +import { withPreparationInfrastructureRecovery } from '../sandbox-recovery.js'; +import type { AgentSandbox, WrapperInstanceLease } from '../agent-sandbox/protocol.js'; -/** Maximum time allowed for workspace preparation (resume, init, fast path). */ +/** Maximum time allowed for wrapper readiness workspace preparation. */ const PREPARE_WORKSPACE_TIMEOUT_MS = 10 * 60 * 1000; const CODE_REVIEW_DISABLED_TOOLS = { @@ -58,46 +47,27 @@ function withWorkspacePreparationTimeout(operation: Promise, step: string) ); } -// --------------------------------------------------------------------------- -// Types -// --------------------------------------------------------------------------- - -/** - * Dependencies for the orchestrator (for testability via dependency injection). - */ export type OrchestratorDeps = { - /** Get a sandbox instance by ID */ - getSandbox: (sandboxId: string) => Promise; - /** Get a Durable Object stub for a session */ + getAgentSandbox: ( + plan: FencedWrapperDispatchRequest | FencedLegacyExecutionRequest + ) => AgentSandbox; getSessionStub: (userId: string, sessionId: string) => DurableObjectStub; - /** Environment bindings */ env: Env; }; -// --------------------------------------------------------------------------- -// ExecutionOrchestrator -// --------------------------------------------------------------------------- - export class ExecutionOrchestrator { - private readonly deps: OrchestratorDeps; private readonly sessionService: SessionService; - constructor(deps: OrchestratorDeps) { - this.deps = deps; + constructor(private readonly deps: OrchestratorDeps) { this.sessionService = new SessionService(); } - /** - * Execute a prompt. Handles all setup and returns immediately after prompt is sent. - * Events stream asynchronously via wrapper -> ingest WS. - * - * @throws ExecutionError with appropriate code on failure (no internal retry) - */ async execute( plan: FencedWrapperDispatchRequest | FencedLegacyExecutionRequest, options?: { onProgress?: (step: string, message: string) => void; onWorkspaceReady?: (ready: WorkspaceReady) => Promise; + leasedInstance?: WrapperInstanceLease; } ): Promise { const executionId = 'executionId' in plan ? plan.executionId : undefined; @@ -111,7 +81,6 @@ export class ExecutionOrchestrator { orgId: orgId ?? '(personal)', mode: agent.mode, }); - logger .withFields({ messageId: turn.messageId, @@ -123,151 +92,33 @@ export class ExecutionOrchestrator { }) .info('ExecutionOrchestrator starting execution'); - const sandboxId = workspace.sandboxId; - if (!sandboxId) { + if (!workspace.sandboxId) { throw ExecutionError.invalidRequest('Missing sandboxId in workspace plan'); } - let sandbox: SandboxInstance; - try { - sandbox = await this.deps.getSandbox(sandboxId); - } catch (error) { - throw ExecutionError.sandboxConnectFailed( - `Failed to connect to sandbox: ${error instanceof Error ? error.message : String(error)}`, - error - ); - } - + const sandbox = this.deps.getAgentSandbox(plan); return withPreparationInfrastructureRecovery( { - sandbox, - sandboxId, + deleteSandbox: reason => sandbox.delete(reason), + sandboxId: workspace.sandboxId, sessionId, phase: 'executionWorkspacePreparation', }, - () => this.executeWithWrapperBootstrap(sandbox, plan, options) - ); - } - - private requiresPreparedDevcontainerRuntime( - plan: FencedWrapperDispatchRequest | FencedLegacyExecutionRequest - ): boolean { - return ( - plan.workspace.metadata.workspace?.devcontainerRequested === true || - plan.workspace.metadata.devcontainer !== undefined - ); - } - - private async executeWithPreparedDevcontainerWorkspace( - sandbox: SandboxInstance, - plan: FencedWrapperDispatchRequest | FencedLegacyExecutionRequest, - prepared: Awaited>, - options?: { - onProgress?: (step: string, message: string) => void; - onWorkspaceReady?: (ready: WorkspaceReady) => Promise; - } - ): Promise { - const { sessionId, userId, orgId } = plan.scope; - const { turn, workspace, agent } = plan; - const preparedWorkspace = await withWorkspacePreparationTimeout( - this.sessionService.prepareWorkspace({ - sandbox, - sandboxId: workspace.sandboxId as ServiceSandboxId, - orgId, - userId, - sessionId: sessionId as ServiceSessionId, - kilocodeModel: agent.model, - env: this.deps.env, - metadata: workspace.metadata, - onProgress: options?.onProgress, - }), - 'devcontainer workspace preparation' - ); - - if (!preparedWorkspace.devcontainer || !preparedWorkspace.ready.devcontainer) { - throw ExecutionError.workspaceSetupFailed( - 'Devcontainer workspace preparation did not resolve runtime metadata' - ); - } - - let wrapperClient: WrapperClient; - let kiloSessionId: string; - try { - const wrapper = await WrapperClient.ensureWrapper(sandbox, preparedWorkspace.session, { - agentSessionId: sessionId, - userId, - workspacePath: preparedWorkspace.context.workspacePath, - sessionId: plan.wrapper.kiloSessionId, - runtimeEnv: preparedWorkspace.runtimeEnv, - devcontainer: preparedWorkspace.devcontainer, - fixedPort: preparedWorkspace.ready.devcontainer.wrapperPort, - }); - wrapperClient = wrapper.client; - kiloSessionId = wrapper.sessionId; - } catch (error) { - throw ExecutionError.wrapperStartFailed( - `Failed to start devcontainer wrapper: ${error instanceof Error ? error.message : String(error)}`, - error - ); - } - - await options?.onWorkspaceReady?.(preparedWorkspace.ready); - - if (prepared.type === 'command') { - await wrapperClient.command(prepared.commandRequest); - } else { - await wrapperClient.prompt(prepared.promptRequest); - } - - try { - await withDORetry( - () => this.deps.getSessionStub(userId, sessionId), - stub => stub.recordKiloServerActivity(), - 'recordKiloServerActivity' - ); - } catch { - logger - .withFields({ sessionId, messageId: turn.messageId }) - .warn('Failed to record kilo server activity'); - } - - logger - .withFields({ sessionId, messageId: turn.messageId }) - .info('ExecutionOrchestrator devcontainer execution started successfully'); - return { kiloSessionId }; - } - - private async workspaceHasGit(sandbox: SandboxInstance, workspacePath: string): Promise { - const timeoutMs = FAST_SANDBOX_COMMAND_TIMEOUT_MS; - const result = await withTimeout( - timedExec( - sandbox, - `test -d '${workspacePath}/.git' && echo exists`, - 'execution.wrapperBootstrap.repoExists' - ), - timeoutMs, - `${SANDBOX_WORKSPACE_PROBE_TIMEOUT_MESSAGE} after ${timeoutMs}ms`, - () => - logSandboxOperationTimeout({ - operation: 'execution.wrapperBootstrap.repoExists', - timeoutMs, - timeoutLayer: 'outer', - }) + () => this.executeThroughAgentSandbox(sandbox, plan, options) ); - return result.stdout?.includes('exists') ?? false; } - private async executeWithWrapperBootstrap( - sandbox: SandboxInstance, + private async executeThroughAgentSandbox( + sandbox: AgentSandbox, plan: FencedWrapperDispatchRequest | FencedLegacyExecutionRequest, options?: { onProgress?: (step: string, message: string) => void; onWorkspaceReady?: (ready: WorkspaceReady) => Promise; + leasedInstance?: WrapperInstanceLease; } ): Promise { - const { sessionId, userId, orgId } = plan.scope; + const { sessionId, userId } = plan.scope; const { turn } = plan; - const prepared = await this.sessionService.buildWrapperSessionReadyAndPromptRequests({ env: this.deps.env, plan, @@ -280,65 +131,34 @@ export class ExecutionOrchestrator { }; } - if (this.requiresPreparedDevcontainerRuntime(plan)) { - return this.executeWithPreparedDevcontainerWorkspace(sandbox, plan, prepared, options); - } - - const workspaceWarm = await this.workspaceHasGit(sandbox, prepared.context.workspacePath); - logger - .withFields({ - sessionId, - messageId: turn.messageId, - workspacePath: prepared.context.workspacePath, - workspaceWarm, - }) - .info('Workspace warmth probe completed'); - if (!workspaceWarm) { - options?.onProgress?.('disk_check', 'Checking disk space...'); - await checkDiskAndCleanBeforeSetup(sandbox, orgId, userId, sessionId); - } - - const bootstrapSession = await sandbox.createSession({ - name: `${sessionId}-bootstrap`, - env: {}, - cwd: '/', - }); - - let wrapperClient: WrapperClient; + let ensured; try { - if (!workspaceWarm) { - options?.onProgress?.('kilo_server', 'Starting Kilo...'); - } - const result = await WrapperClient.ensureBootstrapWrapper(sandbox, bootstrapSession, { - agentSessionId: sessionId, - userId, + ensured = await sandbox.ensureWrapper({ + plan, + prepared, + onProgress: options?.onProgress, + ...(options?.leasedInstance ? { leasedInstance: options.leasedInstance } : {}), }); - wrapperClient = result.client; - logger - .withFields({ sessionId, messageId: turn.messageId, workspaceWarm }) - .info('Bootstrap wrapper ready'); } catch (error) { + if (error instanceof ExecutionError) throw error; throw ExecutionError.wrapperStartFailed( `Failed to start wrapper: ${error instanceof Error ? error.message : String(error)}`, error ); } + let kiloSessionId: string; try { - const readyResult = await withWorkspacePreparationTimeout( - wrapperClient.ensureSessionReady(prepared.readyRequest), - 'wrapper readiness' - ); - logger - .withFields({ - sessionId, - messageId: turn.messageId, - kiloSessionId: readyResult.kiloSessionId, - }) - .info('Wrapper session readiness completed'); - - if (options?.onWorkspaceReady) { - await options.onWorkspaceReady( + if (ensured.status === 'session-ready') { + kiloSessionId = ensured.kiloSessionId; + await options?.onWorkspaceReady?.(ensured.ready); + } else { + const readyResult = await withWorkspacePreparationTimeout( + ensured.client.ensureSessionReady(prepared.readyRequest), + 'wrapper readiness' + ); + kiloSessionId = readyResult.kiloSessionId; + await options?.onWorkspaceReady?.( readyResult.workspaceReady ? { ...prepared.ready, ...readyResult.workspaceReady } : prepared.ready @@ -346,7 +166,7 @@ export class ExecutionOrchestrator { } if (prepared.type === 'command') { - await wrapperClient.command(prepared.commandRequest); + await ensured.client.command(prepared.commandRequest); logger .withFields({ sessionId, @@ -355,7 +175,7 @@ export class ExecutionOrchestrator { }) .info('Wrapper accepted command dispatch'); } else { - await wrapperClient.prompt(prepared.promptRequest); + await ensured.client.prompt(prepared.promptRequest); logger .withFields({ sessionId, messageId: turn.messageId }) .info('Wrapper accepted prompt dispatch'); @@ -372,9 +192,8 @@ export class ExecutionOrchestrator { .withFields({ sessionId, messageId: turn.messageId }) .warn('Failed to record kilo server activity'); } - - logger.info('ExecutionOrchestrator wrapper bootstrap execution started successfully'); - return { kiloSessionId: readyResult.kiloSessionId }; + logger.info('ExecutionOrchestrator wrapper execution started successfully'); + return { kiloSessionId }; } catch (error) { logger .withFields({ @@ -383,7 +202,7 @@ export class ExecutionOrchestrator { errorClass: error instanceof Error ? error.name : 'UnknownError', wrapperErrorCode: error instanceof WrapperError ? error.code : undefined, }) - .warn('ExecutionOrchestrator wrapper bootstrap path failed'); + .warn('ExecutionOrchestrator wrapper dispatch failed'); if (error instanceof WrapperError) { if (error.code === 'WORKSPACE_SETUP_FAILED') { throw ExecutionError.workspaceSetupFailed(error.message, error); @@ -400,16 +219,10 @@ export class ExecutionOrchestrator { } } - private getCreatedOnPlatform( - plan: FencedWrapperDispatchRequest | FencedLegacyExecutionRequest - ): string | undefined { - return plan.workspace.metadata?.identity?.createdOnPlatform; - } - private getToolOverrides( plan: FencedWrapperDispatchRequest | FencedLegacyExecutionRequest ): Record | undefined { - return this.getCreatedOnPlatform(plan) === 'code-review' + return plan.workspace.metadata.identity?.createdOnPlatform === 'code-review' ? CODE_REVIEW_DISABLED_TOOLS : undefined; } diff --git a/services/cloud-agent-next/src/kilo/wrapper-client.test.ts b/services/cloud-agent-next/src/kilo/wrapper-client.test.ts index de3ae49a81..0992271fe7 100644 --- a/services/cloud-agent-next/src/kilo/wrapper-client.test.ts +++ b/services/cloud-agent-next/src/kilo/wrapper-client.test.ts @@ -22,6 +22,7 @@ import { type WrapperTransport, } from './wrapper-client.js'; import type { ExecutionSession, SandboxInstance } from '../types.js'; +import type { WrapperInstanceLease } from '../agent-sandbox/protocol.js'; import { WRAPPER_VERSION } from '../shared/wrapper-version.js'; vi.mock('./ports.js', () => ({ @@ -983,6 +984,33 @@ describe('WrapperClient', () => { expect(startProcessCall[0]).toContain("--user-id 'test-user'"); }); + it('uses backward-compatible environment physical instance markers when startup is leased', async () => { + const session = createMockSession(createCurlError(7, 'Connection refused')); + (session.startProcess as ReturnType).mockResolvedValue({ + id: 'mock-process-id', + waitForPort: vi.fn().mockResolvedValue(undefined), + getLogs: vi.fn().mockResolvedValue({ stdout: '', stderr: '' }), + }); + const client = new WrapperClient({ session, port: defaultPort }); + const leasedInstance: WrapperInstanceLease = { + instanceId: 'instance_test', + instanceGeneration: 6, + }; + + await client.ensureRunning({ + agentSessionId, + userId, + workspacePath: '/workspace/test', + leasedInstance, + }); + + const command = (session.startProcess as ReturnType).mock.calls[0][0] as string; + expect(command).toContain("WRAPPER_INSTANCE_ID='instance_test'"); + expect(command).toContain('WRAPPER_INSTANCE_GENERATION=6'); + expect(command).not.toContain('--wrapper-instance-id'); + expect(command).not.toContain('--wrapper-instance-generation'); + }); + it('includes --session-id when provided', async () => { const session = createMockSession(createCurlError(7, 'Connection refused')); (session.startProcess as ReturnType).mockResolvedValue({ @@ -1408,6 +1436,112 @@ describe('WrapperClient', () => { expect(session.startProcess).not.toHaveBeenCalled(); }); + it('reuses a healthy wrapper only when its physical identity matches a lease', async () => { + const session = createMockSession( + createSuccessResponse({ + ...healthResponseData, + wrapperInstanceId: 'instance_current', + wrapperInstanceGeneration: 2, + }) + ); + const sandbox = createMockSandbox({ port: 5555, healthy: true }); + + await expect( + WrapperClient.ensureWrapper(sandbox, session, { + ...wrapperOptions, + leasedInstance: { instanceId: 'instance_current', instanceGeneration: 2 }, + }) + ).resolves.toMatchObject({ sessionId: 'kilo-sess-1' }); + expect(session.startProcess).not.toHaveBeenCalled(); + }); + + it('reuses an env-tagged legacy wrapper whose health does not report its lease', async () => { + const session = createMockSession(createSuccessResponse(healthResponseData)); + const sandbox = createMockSandbox({ port: 5555, healthy: true }); + (sandbox.listProcesses as ReturnType).mockResolvedValue([ + { + id: 'legacy-wrapper-id', + command: + "WRAPPER_PORT=5555 WRAPPER_INSTANCE_ID='instance_current' WRAPPER_INSTANCE_GENERATION=2 bun run /usr/local/bin/kilocode-wrapper.js --agent-session test-session", + status: 'running', + }, + ]); + + await expect( + WrapperClient.ensureWrapper(sandbox, session, { + ...wrapperOptions, + leasedInstance: { instanceId: 'instance_current', instanceGeneration: 2 }, + }) + ).resolves.toMatchObject({ sessionId: 'kilo-sess-1' }); + expect(session.startProcess).not.toHaveBeenCalled(); + }); + + it('reuses an authorized legacy devcontainer wrapper whose health omits identity', async () => { + const session = createMockSession(createSuccessResponse(healthResponseData)); + const sandbox = { + listProcesses: vi.fn().mockResolvedValue([]), + exec: vi.fn().mockImplementation((command: string) => { + if (command.startsWith('if [ -S')) { + return Promise.resolve({ exitCode: 0, stdout: '/var/run/docker.sock', stderr: '' }); + } + if (command.includes('/proc/42/environ')) { + return Promise.resolve({ + exitCode: 0, + stdout: 'WRAPPER_INSTANCE_ID=instance_current WRAPPER_INSTANCE_GENERATION=2', + stderr: '', + }); + } + if (command.includes('docker exec')) { + return Promise.resolve({ + exitCode: 0, + stdout: '42 WRAPPER_PORT=5555 kilocode-wrapper --agent-session test-session\n', + stderr: '', + }); + } + return Promise.resolve({ + exitCode: 0, + stdout: 'container-legacy\t0.0.0.0:5555->5555/tcp\tkilo.agentSession=test-session\n', + stderr: '', + }); + }), + } as unknown as SandboxInstance; + + await expect( + WrapperClient.ensureWrapper(sandbox, session, { + ...wrapperOptions, + leasedInstance: { instanceId: 'instance_current', instanceGeneration: 2 }, + devcontainer: { + containerId: 'container-legacy', + innerWorkspaceFolder: '/workspaces/test', + workspacePath: '/workspace/test', + agentSessionId: 'test-session', + overrideConfigPath: '/tmp/devcontainer.json', + teardown: vi.fn(), + }, + }) + ).resolves.toMatchObject({ sessionId: 'kilo-sess-1' }); + expect(session.startProcess).not.toHaveBeenCalled(); + }); + + it('blocks a leased replacement when a healthy wrapper has a different physical identity', async () => { + const session = createMockSession( + createSuccessResponse({ + ...healthResponseData, + wrapperInstanceId: 'instance_old', + wrapperInstanceGeneration: 1, + }) + ); + const sandbox = createMockSandbox({ port: 5555, healthy: true }); + + await expect( + WrapperClient.ensureWrapper(sandbox, session, { + ...wrapperOptions, + leasedInstance: { instanceId: 'instance_new', instanceGeneration: 2 }, + }) + ).rejects.toThrow(/does not match leased physical instance/); + expect(session.startProcess).not.toHaveBeenCalled(); + }); + it('passes Docker socket env when restarting a version-mismatched devcontainer wrapper', async () => { let healthCalls = 0; const session = createMockSession((cmd: string) => { @@ -1476,6 +1610,78 @@ describe('WrapperClient', () => { expect(session.startProcess).toHaveBeenCalled(); }); + it('accepts a tagged leased launch when a legacy wrapper omits identity from health', async () => { + let healthCalls = 0; + const session = createMockSession((cmd: string) => { + if (cmd.includes('/health')) { + healthCalls++; + if (healthCalls === 1) return createCurlError(7, 'Connection refused'); + return createSuccessResponse(healthResponseData); + } + return createCurlError(7, 'Connection refused'); + }); + (session.startProcess as ReturnType).mockResolvedValue({ + id: 'legacy-wrapper-id', + waitForPort: vi.fn().mockResolvedValue(undefined), + getLogs: vi.fn().mockResolvedValue({ stdout: '', stderr: '' }), + }); + const sandbox = createMockSandbox(null); + (sandbox.listProcesses as ReturnType) + .mockResolvedValueOnce([]) + .mockResolvedValueOnce([ + { + id: 'legacy-wrapper-id', + command: + "WRAPPER_PORT=5000 WRAPPER_INSTANCE_ID='instance_compat' WRAPPER_INSTANCE_GENERATION=3 bun run /usr/local/bin/kilocode-wrapper.js --agent-session test-session", + status: 'running', + }, + ]); + + await expect( + WrapperClient.ensureWrapper(sandbox, session, { + ...wrapperOptions, + leasedInstance: { instanceId: 'instance_compat', instanceGeneration: 3 }, + }) + ).resolves.toMatchObject({ sessionId: 'kilo-sess-1' }); + expect(sandbox.listProcesses).toHaveBeenCalledTimes(2); + }); + + it('rejects a legacy launch whose assigned marker cannot be observed', async () => { + let healthCalls = 0; + const session = createMockSession((cmd: string) => { + if (cmd.includes('/health')) { + healthCalls++; + if (healthCalls === 1) return createCurlError(7, 'Connection refused'); + return createSuccessResponse(healthResponseData); + } + return createCurlError(7, 'Connection refused'); + }); + (session.startProcess as ReturnType).mockResolvedValue({ + id: 'unverified-wrapper-id', + waitForPort: vi.fn().mockResolvedValue(undefined), + getLogs: vi.fn().mockResolvedValue({ stdout: '', stderr: '' }), + }); + + await expect( + WrapperClient.ensureWrapper(createMockSandbox(null), session, { + ...wrapperOptions, + leasedInstance: { instanceId: 'instance_compat', instanceGeneration: 3 }, + }) + ).rejects.toThrow(/did not report the leased physical instance/); + }); + + it('does not accept an untagged legacy listener on a selected leased port', async () => { + const session = createMockSession(createSuccessResponse(healthResponseData)); + + await expect( + WrapperClient.ensureWrapper(createMockSandbox(null), session, { + ...wrapperOptions, + leasedInstance: { instanceId: 'instance_compat', instanceGeneration: 3 }, + }) + ).rejects.toThrow(/did not report the leased physical instance/); + expect(session.startProcess).not.toHaveBeenCalled(); + }); + it('retries with new port on EADDRINUSE', async () => { let healthCalls = 0; const session = createMockSession((cmd: string) => { @@ -1650,6 +1856,50 @@ describe('WrapperClient', () => { expect(sandbox.exec).not.toHaveBeenCalled(); }); + it('accepts a tagged leased bootstrap launch when legacy health omits identity', async () => { + const session = createMockSession({ exitCode: 0, stdout: '{}' }); + const sandbox = createMockSandbox(null) as SandboxInstance & { + containerFetch: ReturnType; + listProcesses: ReturnType; + }; + sandbox.listProcesses.mockResolvedValueOnce([]).mockResolvedValueOnce([ + { + id: 'legacy-bootstrap-id', + command: + "WRAPPER_PORT=30000 WRAPPER_INSTANCE_ID='instance_bootstrap' WRAPPER_INSTANCE_GENERATION=4 bun run /usr/local/bin/kilocode-wrapper.js --agent-session test-session", + status: 'running', + }, + ]); + let healthCalls = 0; + sandbox.containerFetch.mockImplementation(() => { + healthCalls++; + if (healthCalls === 1) { + return Promise.resolve( + new Response(JSON.stringify({ error: 'NOT_READY', message: 'not ready' }), { + status: 503, + headers: { 'Content-Type': 'application/json' }, + }) + ); + } + return Promise.resolve( + Response.json({ + healthy: true, + state: 'idle', + version: WRAPPER_VERSION, + sessionId: 'kilo-sess-bootstrap', + }) + ); + }); + + await expect( + WrapperClient.ensureBootstrapWrapper(sandbox, session, { + ...bootstrapOptions, + leasedInstance: { instanceId: 'instance_bootstrap', instanceGeneration: 4 }, + }) + ).resolves.toBeDefined(); + expect(sandbox.listProcesses).toHaveBeenCalledTimes(2); + }); + it('replaces a pre-bootstrap wrapper reporting the previous wrapper version', async () => { const session = createMockSession({ exitCode: 0, stdout: '{}' }); const sandbox = createMockSandbox({ port: 5555, healthy: true }) as SandboxInstance & { diff --git a/services/cloud-agent-next/src/kilo/wrapper-client.ts b/services/cloud-agent-next/src/kilo/wrapper-client.ts index d97a38da73..a708c59f46 100644 --- a/services/cloud-agent-next/src/kilo/wrapper-client.ts +++ b/services/cloud-agent-next/src/kilo/wrapper-client.ts @@ -7,8 +7,10 @@ import { dirname } from 'node:path'; import type { ExecutionSession, SandboxInstance } from '../types.js'; +import type { WrapperInstanceLease } from '../agent-sandbox/protocol.js'; import { logger } from '../logger.js'; import { + discoverSessionWrappers, findWrapperForSession, findWrapperForSessionInProcesses, getWrapperSessionMarker, @@ -50,6 +52,7 @@ export type EnsureRunningOptions = { maxWaitMs?: number; workspacePath?: string; sessionId?: string; + leasedInstance?: WrapperInstanceLease; /** * Prepared session runtime environment for Kilo. This is the same env used to * create the sandbox execution session; wrapper-owned values below are layered @@ -70,6 +73,7 @@ export type EnsureWrapperOptions = { userId: string; workspacePath: string; sessionId?: string; + leasedInstance?: WrapperInstanceLease; /** See {@link EnsureRunningOptions.runtimeEnv}. */ runtimeEnv?: Record; /** See {@link EnsureRunningOptions.devcontainer}. */ @@ -88,6 +92,7 @@ export type EnsureBootstrapWrapperOptions = { userId: string; wrapperPath?: string; maxWaitMs?: number; + leasedInstance?: WrapperInstanceLease; }; export type SessionBinding = { @@ -111,6 +116,8 @@ export type WrapperHealthResponse = { state: 'idle' | 'active'; version: string; sessionId: string; + wrapperInstanceId?: string; + wrapperInstanceGeneration?: number; }; export type WrapperPty = { @@ -200,6 +207,48 @@ const ERROR_STATUS_CODES: Record = { /** Max attempts for port allocation in ensureWrapper (retry with new random port on failure) */ const MAX_PORT_ATTEMPTS = 3; +function healthMatchesLease( + health: WrapperHealthResponse, + leasedInstance: WrapperInstanceLease | undefined, + allowUnreportedIdentity = false +): boolean { + if (!leasedInstance) return true; + if ( + health.wrapperInstanceId === leasedInstance.instanceId && + health.wrapperInstanceGeneration === leasedInstance.instanceGeneration + ) { + return true; + } + return ( + allowUnreportedIdentity && + health.wrapperInstanceId === undefined && + health.wrapperInstanceGeneration === undefined + ); +} + +async function observationMatchesLease( + sandbox: SandboxInstance, + agentSessionId: string, + leasedInstance: WrapperInstanceLease, + options: { inspectContainers: boolean; expectedContainerId?: string } +): Promise { + const observation = await discoverSessionWrappers(sandbox, agentSessionId, { + inspectContainers: options.inspectContainers, + }); + if (observation.status !== 'present' || observation.observed.length !== 1) return false; + const wrapper = observation.observed[0]; + if (!wrapper) return false; + if (options.expectedContainerId !== undefined) { + if (wrapper.representation !== 'container' || wrapper.id !== options.expectedContainerId) { + return false; + } + } + return ( + wrapper.instanceId === leasedInstance.instanceId && + wrapper.instanceGeneration === leasedInstance.instanceGeneration + ); +} + function buildExportFileContent(env: Record): string { return `${validShellEnvEntries(env) .map(([key, value]) => `export ${key}=${shellQuote(value)}`) @@ -467,7 +516,7 @@ export class WrapperClient { * NOTE: This method assumes the WrapperClient was created with a port. * Port retry on EADDRINUSE is handled by the static ensureWrapper() method. */ - async ensureRunning(options: EnsureRunningOptions): Promise { + async ensureRunning(options: EnsureRunningOptions): Promise<{ started: boolean }> { const { agentSessionId, userId, @@ -475,6 +524,7 @@ export class WrapperClient { maxWaitMs = 30_000, workspacePath, sessionId, + leasedInstance, runtimeEnv, devcontainer, } = options; @@ -483,7 +533,7 @@ export class WrapperClient { try { await this.health(); logger.debug('WrapperClient: wrapper already running'); - return; // Already running + return { started: false }; } catch { // Not running, need to start logger.debug('WrapperClient: wrapper not running, starting...'); @@ -515,6 +565,12 @@ export class WrapperClient { WRAPPER_LOG_PATH: wrapperLogPath, KILO_SESSION_RETRY_LIMIT: '5', KILO_CLOUD_AGENT: '1', + ...(leasedInstance + ? { + WRAPPER_INSTANCE_ID: leasedInstance.instanceId, + WRAPPER_INSTANCE_GENERATION: String(leasedInstance.instanceGeneration), + } + : {}), }; const commandEnvParts = [ `WRAPPER_PORT=${this.port}`, @@ -522,6 +578,13 @@ export class WrapperClient { `WRAPPER_LOG_PATH=${wrapperLogPath}`, `KILO_SESSION_RETRY_LIMIT=5`, `KILO_CLOUD_AGENT=1`, + // Environment markers let pre-lease wrapper bundles launch during a rolling deploy. + ...(leasedInstance + ? [ + `WRAPPER_INSTANCE_ID=${shellQuote(leasedInstance.instanceId)}`, + `WRAPPER_INSTANCE_GENERATION=${leasedInstance.instanceGeneration}`, + ] + : []), ...dockerEnvParts, ]; const devContainerSessionHome = @@ -593,7 +656,7 @@ export class WrapperClient { clearTimeout(waitTimeoutId); logger.debug('WrapperClient: wrapper is ready', { port: this.port, processId: proc.id }); - return; + return { started: true }; } catch (error) { const startupError = error instanceof Error ? error : new Error(String(error)); @@ -714,9 +777,41 @@ export class WrapperClient { try { const healthResponse = await client.health(); if (healthResponse.version === WRAPPER_VERSION) { + let allowUnreportedIdentity = false; + if ( + options.leasedInstance && + healthResponse.wrapperInstanceId === undefined && + healthResponse.wrapperInstanceGeneration === undefined + ) { + allowUnreportedIdentity = await observationMatchesLease( + sandbox, + agentSessionId, + options.leasedInstance, + { + inspectContainers: existing.kind === 'container', + ...(existing.kind === 'container' + ? { expectedContainerId: existing.process.id } + : {}), + } + ); + } + if ( + options.leasedInstance && + !healthMatchesLease(healthResponse, options.leasedInstance, allowUnreportedIdentity) + ) { + throw new WrapperNotReadyError( + `Existing wrapper does not match leased physical instance ${options.leasedInstance.instanceId}` + ); + } return { client, sessionId: healthResponse.sessionId }; } + if (options.leasedInstance) { + throw new WrapperNotReadyError( + 'Existing leased wrapper reported an incompatible version' + ); + } + logger .withFields({ agentSessionId, @@ -761,7 +856,10 @@ export class WrapperClient { }) .warn('Failed to stop version-mismatched wrapper, starting replacement anyway'); } - } catch { + } catch (error) { + if (options.leasedInstance) { + throw error; + } logger .withFields({ agentSessionId, port }) .warn('Existing wrapper not healthy, will start new one'); @@ -784,13 +882,32 @@ export class WrapperClient { const client = new WrapperClient({ session, port }); try { - await client.ensureRunning(options); + const running = await client.ensureRunning(options); const healthResponse = await client.health(); if (healthResponse.version !== WRAPPER_VERSION) { throw new WrapperNotReadyError( `Wrapper version mismatch after startup: expected ${WRAPPER_VERSION}, got ${healthResponse.version}` ); } + let allowUnreportedIdentity = false; + if ( + running.started && + options.leasedInstance && + healthResponse.wrapperInstanceId === undefined && + healthResponse.wrapperInstanceGeneration === undefined + ) { + allowUnreportedIdentity = await observationMatchesLease( + sandbox, + agentSessionId, + options.leasedInstance, + { inspectContainers: options.devcontainer !== undefined } + ); + } + if (!healthMatchesLease(healthResponse, options.leasedInstance, allowUnreportedIdentity)) { + throw new WrapperNotReadyError( + 'Started wrapper did not report the leased physical instance' + ); + } return { client, sessionId: healthResponse.sessionId }; } catch (error) { @@ -831,10 +948,32 @@ export class WrapperClient { try { const healthResponse = await client.health(); if (healthResponse.version === WRAPPER_VERSION) { + const allowUnreportedIdentity = + options.leasedInstance !== undefined && + healthResponse.wrapperInstanceId === undefined && + healthResponse.wrapperInstanceGeneration === undefined + ? await observationMatchesLease(sandbox, agentSessionId, options.leasedInstance, { + inspectContainers: false, + }) + : false; + if ( + options.leasedInstance && + !healthMatchesLease(healthResponse, options.leasedInstance, allowUnreportedIdentity) + ) { + throw new WrapperNotReadyError( + `Existing bootstrap wrapper does not match leased physical instance ${options.leasedInstance.instanceId}` + ); + } return { client }; } + if (options.leasedInstance) { + throw new WrapperNotReadyError( + 'Existing leased bootstrap wrapper reported an incompatible version' + ); + } await sandbox.exec(`pkill -f -- '${getWrapperSessionMarker(agentSessionId)}'`); - } catch { + } catch (error) { + if (options.leasedInstance) throw error; logger .withFields({ agentSessionId, port }) .warn('Existing bootstrap wrapper not healthy, will start new one'); @@ -850,13 +989,27 @@ export class WrapperClient { transport: new ContainerFetchWrapperTransport({ sandbox, port }), }); try { - await client.ensureRunning(options); + const running = await client.ensureRunning(options); const healthResponse = await client.health(); if (healthResponse.version !== WRAPPER_VERSION) { throw new WrapperNotReadyError( `Wrapper version mismatch after startup: expected ${WRAPPER_VERSION}, got ${healthResponse.version}` ); } + const allowUnreportedIdentity = + running.started && + options.leasedInstance !== undefined && + healthResponse.wrapperInstanceId === undefined && + healthResponse.wrapperInstanceGeneration === undefined + ? await observationMatchesLease(sandbox, agentSessionId, options.leasedInstance, { + inspectContainers: false, + }) + : false; + if (!healthMatchesLease(healthResponse, options.leasedInstance, allowUnreportedIdentity)) { + throw new WrapperNotReadyError( + 'Started bootstrap wrapper did not report the leased physical instance' + ); + } return { client }; } catch (error) { lastError = error instanceof Error ? error : new Error(String(error)); diff --git a/services/cloud-agent-next/src/kilo/wrapper-manager.test.ts b/services/cloud-agent-next/src/kilo/wrapper-manager.test.ts index c186699a24..41eeb4e794 100644 --- a/services/cloud-agent-next/src/kilo/wrapper-manager.test.ts +++ b/services/cloud-agent-next/src/kilo/wrapper-manager.test.ts @@ -1,10 +1,11 @@ import { describe, expect, it, vi } from 'vitest'; import { + discoverSessionWrappers, extractPublishedWrapperPort, findWrapperContainerForSession, isWrapperLiveInProcessesOrContainers, listWrapperContainers, - stopWrapper, + stopObservedWrappers, } from './wrapper-manager.js'; const mockExec = (impl: (cmd: string) => { exitCode: number; stdout?: string }) => ({ @@ -104,12 +105,14 @@ describe('listWrapperContainers', () => { expect(await listWrapperContainers(sandbox)).toEqual([]); }); - it('skips rows with no published port', async () => { + it('retains labelled rows with no published port for lifecycle observation', async () => { const sandbox = mockExec(() => ({ exitCode: 0, stdout: 'cont-2\t\tkilo.agentSession=agent_abc\n', })); - expect(await listWrapperContainers(sandbox)).toEqual([]); + expect(await listWrapperContainers(sandbox)).toEqual([ + { containerId: 'cont-2', agentSessionId: 'agent_abc' }, + ]); }); it('parses multiple rows', async () => { @@ -136,6 +139,234 @@ describe('listWrapperContainers', () => { }); }); +describe('discoverSessionWrappers', () => { + it('returns every direct and devcontainer wrapper process tagged for the session', async () => { + const sandbox = { + listProcesses: vi.fn().mockResolvedValue([ + { + id: 'direct-one', + command: + 'WRAPPER_PORT=5010 kilocode-wrapper --agent-session agent_xyz --wrapper-instance-id instance_direct --wrapper-instance-generation 4', + status: 'running', + }, + { + id: 'direct-two', + command: 'WRAPPER_PORT=5011 kilocode-wrapper --agent-session agent_xyz', + status: 'starting', + }, + ]), + exec: vi.fn().mockImplementation((command: string) => { + if (command.includes('docker ps')) { + return Promise.resolve({ + exitCode: 0, + stdout: + 'cont-id\t0.0.0.0:5050->5050/tcp\tkilo.agentSession=agent_xyz,kilo.wrapperPort=5050\n', + }); + } + if (command.includes('/proc/42/environ')) { + return Promise.resolve({ + exitCode: 0, + stdout: 'WRAPPER_INSTANCE_ID=instance_container WRAPPER_INSTANCE_GENERATION=7', + }); + } + return Promise.resolve({ + exitCode: 0, + stdout: '42 WRAPPER_PORT=5050 kilocode-wrapper --agent-session agent_xyz\n', + }); + }), + }; + + await expect( + discoverSessionWrappers(sandbox as never, 'agent_xyz', { dockerEnv: {} }) + ).resolves.toEqual({ + status: 'present', + observed: [ + { + representation: 'process', + id: 'direct-one', + port: 5010, + instanceId: 'instance_direct', + instanceGeneration: 4, + }, + { representation: 'process', id: 'direct-two', port: 5011 }, + { + representation: 'container', + id: 'cont-id', + port: 5050, + instanceId: 'instance_container', + instanceGeneration: 7, + }, + ], + }); + }); + + it('normalizes shell-quoted physical instance markers from wrapper startup', async () => { + const sandbox = { + listProcesses: vi.fn().mockResolvedValue([ + { + id: 'quoted-instance', + command: + "WRAPPER_PORT=5010 bun run '/usr/local/bin/kilocode-wrapper.js' --agent-session agent_xyz --wrapper-instance-id 'instance_quoted' --wrapper-instance-generation 4", + status: 'running', + }, + ]), + exec: vi.fn(), + }; + + await expect( + discoverSessionWrappers(sandbox as never, 'agent_xyz', { inspectContainers: false }) + ).resolves.toEqual({ + status: 'present', + observed: [ + { + representation: 'process', + id: 'quoted-instance', + port: 5010, + instanceId: 'instance_quoted', + instanceGeneration: 4, + }, + ], + }); + }); + + it('observes backward-compatible environment physical instance markers', async () => { + const sandbox = { + listProcesses: vi.fn().mockResolvedValue([ + { + id: 'compat-instance', + command: + "WRAPPER_PORT=5010 WRAPPER_INSTANCE_ID='instance_compat' WRAPPER_INSTANCE_GENERATION=5 bun run '/usr/local/bin/kilocode-wrapper.js' --agent-session agent_xyz", + status: 'running', + }, + ]), + exec: vi.fn(), + }; + + await expect( + discoverSessionWrappers(sandbox as never, 'agent_xyz', { inspectContainers: false }) + ).resolves.toEqual({ + status: 'present', + observed: [ + { + representation: 'process', + id: 'compat-instance', + port: 5010, + instanceId: 'instance_compat', + instanceGeneration: 5, + }, + ], + }); + }); + + it('reports inspection failure instead of absence when requested Docker inspection fails', async () => { + const sandbox = { + listProcesses: vi.fn().mockResolvedValue([]), + exec: vi.fn().mockRejectedValue(new Error('docker unavailable')), + }; + + await expect( + discoverSessionWrappers(sandbox as never, 'agent_xyz', { dockerEnv: {} }) + ).resolves.toMatchObject({ + status: 'inspection-failed', + error: expect.stringContaining('docker unavailable'), + }); + }); + + it('does not require Docker inspection for a standard-sandbox lifecycle query', async () => { + const sandbox = { + listProcesses: vi.fn().mockResolvedValue([]), + exec: vi.fn().mockRejectedValue(new Error('docker unavailable')), + }; + + await expect( + discoverSessionWrappers(sandbox as never, 'agent_xyz', { inspectContainers: false }) + ).resolves.toEqual({ status: 'absent' }); + expect(sandbox.exec).not.toHaveBeenCalled(); + }); + + it('observes a direct physical wrapper even when it has no HTTP port yet', async () => { + const sandbox = { + listProcesses: vi.fn().mockResolvedValue([ + { + id: 'direct-starting', + command: + 'kilocode-wrapper --agent-session agent_xyz --wrapper-instance-id instance_starting --wrapper-instance-generation 8', + status: 'starting', + }, + ]), + exec: vi.fn(), + }; + + await expect( + discoverSessionWrappers(sandbox as never, 'agent_xyz', { inspectContainers: false }) + ).resolves.toEqual({ + status: 'present', + observed: [ + { + representation: 'process', + id: 'direct-starting', + instanceId: 'instance_starting', + instanceGeneration: 8, + }, + ], + }); + }); + + it('stops an environment-tagged direct wrapper by its logical session marker', async () => { + const sandbox = { exec: vi.fn().mockResolvedValue({ exitCode: 0 }) }; + + await stopObservedWrappers(sandbox as never, 'agent_xyz', [ + { + representation: 'process', + id: 'direct-one', + port: 5010, + instanceId: 'instance_direct', + instanceGeneration: 4, + }, + ]); + + expect(sandbox.exec).toHaveBeenCalledWith(expect.stringContaining('--agent-session agent_xyz')); + expect(sandbox.exec).not.toHaveBeenCalledWith(expect.stringContaining('WRAPPER_INSTANCE_ID')); + }); + + it('force stops a leased devcontainer wrapper process without destroying its persistent container', async () => { + const sandbox = { + exec: vi + .fn() + .mockResolvedValueOnce({ exitCode: 0, stdout: '/run/user/1000/docker.sock' }) + .mockResolvedValueOnce({ exitCode: 0, stdout: '' }), + }; + + await stopObservedWrappers( + sandbox as never, + 'agent_xyz', + [ + { + representation: 'container', + id: 'persistent-container', + port: 5050, + instanceId: 'instance_container', + instanceGeneration: 7, + }, + ], + { + force: true, + devcontainer: { + workspacePath: '/workspace/repo', + configPath: '.devcontainer/devcontainer.json', + }, + } + ); + + const command = sandbox.exec.mock.calls[1][0] as string; + expect(command).toContain('docker exec'); + expect(command).toContain('pkill -9 -f --'); + expect(command).toContain('--agent-session agent_xyz'); + expect(command).not.toContain('WRAPPER_INSTANCE_ID'); + expect(command).not.toContain('docker kill'); + }); +}); + describe('findWrapperContainerForSession', () => { it('returns null when no container matches', async () => { const sandbox = mockExec(() => ({ exitCode: 0, stdout: '' })); @@ -163,60 +394,6 @@ describe('findWrapperContainerForSession', () => { }); }); -describe('stopWrapper', () => { - it('kills only the inner wrapper process when devcontainer metadata is available', async () => { - const sandbox = { - listProcesses: vi.fn(async () => []), - exec: vi - .fn() - .mockResolvedValueOnce({ exitCode: 0, stdout: '/run/user/1000/docker.sock' }) - .mockResolvedValueOnce({ - exitCode: 0, - stdout: - 'cont-id\t0.0.0.0:5050->5050/tcp\tkilo.agentSession=agent_xyz,kilo.wrapperPort=5050\n', - }) - .mockResolvedValueOnce({ exitCode: 0, stdout: '/run/user/1000/docker.sock' }) - .mockResolvedValueOnce({ exitCode: 0, stdout: '', stderr: '' }), - }; - - await stopWrapper(sandbox as never, 'agent_xyz', { - devcontainer: { - workspacePath: '/workspace/repo', - configPath: '.devcontainer/devcontainer.json', - }, - }); - - const command = sandbox.exec.mock.calls[3][0] as string; - expect(command).toContain('devcontainer exec'); - expect(command).toContain('--workspace-folder'); - expect(command).toContain('/workspace/repo'); - // --config keeps the CLI applying our remoteUser=root override on exec. - expect(command).toContain("--config '/tmp/devcontainer-override-agent_xyz/devcontainer.json'"); - expect(command).toContain('pkill -f --'); - expect(command).not.toContain('docker kill'); - }); - - it('kills the container when devcontainer metadata is unavailable', async () => { - const sandbox = { - listProcesses: vi.fn(async () => []), - exec: vi - .fn() - .mockResolvedValueOnce({ exitCode: 0, stdout: '/run/user/1000/docker.sock' }) - .mockResolvedValueOnce({ - exitCode: 0, - stdout: - 'cont-id\t0.0.0.0:5050->5050/tcp\tkilo.agentSession=agent_xyz,kilo.wrapperPort=5050\n', - }) - .mockResolvedValueOnce({ exitCode: 0, stdout: '/run/user/1000/docker.sock' }) - .mockResolvedValueOnce({ exitCode: 0, stdout: '', stderr: '' }), - }; - - await stopWrapper(sandbox as never, 'agent_xyz'); - - expect(sandbox.exec.mock.calls[3][0]).toBe("docker kill 'cont-id'"); - }); -}); - describe('isWrapperLiveInProcessesOrContainers', () => { // The Process type from @cloudflare/sandbox has more fields than we exercise // here (kill, getLogs, etc.); cast through unknown so the unit test stays @@ -231,6 +408,16 @@ describe('isWrapperLiveInProcessesOrContainers', () => { expect(isWrapperLiveInProcessesOrContainers([baseProc], [], 'agent_xyz')).toBe(true); }); + it('protects a live direct wrapper workspace even before its port is available', () => { + const startingWithoutPort = { + id: 'p-starting', + command: 'kilocode-wrapper --agent-session agent_xyz', + status: 'starting' as const, + } as unknown as Parameters[0][number]; + + expect(isWrapperLiveInProcessesOrContainers([startingWithoutPort], [], 'agent_xyz')).toBe(true); + }); + it('returns true on a docker-label match', () => { expect( isWrapperLiveInProcessesOrContainers( diff --git a/services/cloud-agent-next/src/kilo/wrapper-manager.ts b/services/cloud-agent-next/src/kilo/wrapper-manager.ts index 8e548f577c..f78b0b2b5b 100644 --- a/services/cloud-agent-next/src/kilo/wrapper-manager.ts +++ b/services/cloud-agent-next/src/kilo/wrapper-manager.ts @@ -9,20 +9,21 @@ */ import type { SandboxInstance } from '../types.js'; +import type { ObservedWrapper, WrapperObservation } from '../agent-sandbox/protocol.js'; import { logger } from '../logger.js'; -import { - getDevContainerOverridePath, - KILO_AGENT_SESSION_LABEL, - KILO_WRAPPER_PORT_LABEL, -} from './devcontainer.js'; +import { KILO_AGENT_SESSION_LABEL, KILO_WRAPPER_PORT_LABEL } from './devcontainer.js'; import { dockerSocketEnv, resolveDockerSocketPath } from './sandbox-runtime.js'; import { shellQuote } from './utils.js'; // Re-export Process type from sandbox for consumers type Process = Awaited>[number]; -/** Command-line marker to identify which session owns a wrapper */ +/** Command markers identifying the logical session and leased physical wrapper. */ const KILO_WRAPPER_SESSION_FLAG = '--agent-session'; +const KILO_WRAPPER_INSTANCE_FLAG = '--wrapper-instance-id'; +const KILO_WRAPPER_INSTANCE_GENERATION_FLAG = '--wrapper-instance-generation'; +const KILO_WRAPPER_INSTANCE_ENV = 'WRAPPER_INSTANCE_ID='; +const KILO_WRAPPER_INSTANCE_GENERATION_ENV = 'WRAPPER_INSTANCE_GENERATION='; /** * Information about a running wrapper. @@ -64,18 +65,41 @@ export function extractWrapperPortFromCommand(command: string): number | null { * @param command - The full command string * @returns The session ID, or null if not found */ -export function extractWrapperSessionIdFromCommand(command: string): string | null { - const flagIndex = command.indexOf(KILO_WRAPPER_SESSION_FLAG); +function extractFlagValueFromCommand(command: string, flag: string): string | null { + const flagIndex = command.indexOf(flag); if (flagIndex === -1) return null; - const afterFlag = command.slice(flagIndex + KILO_WRAPPER_SESSION_FLAG.length).trimStart(); + const afterFlag = command.slice(flagIndex + flag.length).trimStart(); if (!afterFlag) return null; - const endIdx = afterFlag.indexOf(' '); - if (endIdx === -1) { - return afterFlag; + const quote = afterFlag[0]; + if (quote === "'" || quote === '"') { + const closingQuoteIndex = afterFlag.indexOf(quote, 1); + return closingQuoteIndex === -1 ? null : afterFlag.slice(1, closingQuoteIndex); } - return afterFlag.slice(0, endIdx); + + const endIdx = afterFlag.indexOf(' '); + return endIdx === -1 ? afterFlag : afterFlag.slice(0, endIdx); +} + +export function extractWrapperSessionIdFromCommand(command: string): string | null { + return extractFlagValueFromCommand(command, KILO_WRAPPER_SESSION_FLAG); +} + +export function extractWrapperInstanceIdFromCommand(command: string): string | null { + return ( + extractFlagValueFromCommand(command, KILO_WRAPPER_INSTANCE_FLAG) ?? + extractFlagValueFromCommand(command, KILO_WRAPPER_INSTANCE_ENV) + ); +} + +export function extractWrapperInstanceGenerationFromCommand(command: string): number | null { + const value = + extractFlagValueFromCommand(command, KILO_WRAPPER_INSTANCE_GENERATION_FLAG) ?? + extractFlagValueFromCommand(command, KILO_WRAPPER_INSTANCE_GENERATION_ENV); + if (!value) return null; + const generation = Number.parseInt(value, 10); + return Number.isInteger(generation) && generation >= 0 ? generation : null; } /** @@ -143,9 +167,13 @@ export async function findWrapperForSession( type LabeledWrapperRow = { containerId: string; agentSessionId: string; - port: number; + port?: number; }; +export type WrapperContainerInspection = + | { status: 'ok'; containers: LabeledWrapperRow[] } + | { status: 'inspection-failed'; error: string }; + /** Minimal exec surface — both `SandboxInstance` and `ExecutionSession` satisfy this. */ type DockerExecutor = { exec( @@ -179,28 +207,29 @@ export function extractPublishedWrapperPort(portsField: string): number | null { * spaces and arrows — survives intact. Each label key/value pair is emitted as * `Labels=k1=v1,k2=v2` so we can pull `kilo.agentSession` and the wrapper port. */ -export async function listWrapperContainers( +export async function inspectWrapperContainers( executor: DockerExecutor, options?: { dockerEnv?: Record } -): Promise { +): Promise { const cmd = `docker ps --filter label=${KILO_AGENT_SESSION_LABEL} --format '{{.ID}}\\t{{.Ports}}\\t{{.Labels}}'`; - let result: { exitCode: number; stdout?: string; stderr?: string } | undefined; + let result: { exitCode: number; stdout?: string; stderr?: string }; try { const dockerEnv = options?.dockerEnv ?? dockerSocketEnv(await resolveDockerSocketPath(executor)); result = await executor.exec(cmd, { env: dockerEnv }); } catch (error) { - logger - .withFields({ error: error instanceof Error ? error.message : String(error) }) - .debug('docker ps for wrapper containers failed'); - return []; + const message = error instanceof Error ? error.message : String(error); + logger.withFields({ error: message }).debug('docker ps for wrapper containers failed'); + return { status: 'inspection-failed', error: message }; + } + if (result.exitCode !== 0) { + return { + status: 'inspection-failed', + error: result.stderr?.trim() || `docker ps exited with code ${result.exitCode}`, + }; } - // Defensive: a missing/undefined response (or non-zero exit) means docker - // isn't reachable on this image — fall through and let process-list lookup - // (or absence of wrapper) drive the decision. - if (!result || result.exitCode !== 0) return []; - const rows: LabeledWrapperRow[] = []; + const containers: LabeledWrapperRow[] = []; for (const line of (result.stdout ?? '').split('\n')) { const trimmed = line.trim(); if (!trimmed) continue; @@ -210,10 +239,21 @@ export async function listWrapperContainers( if (!agentSessionId) continue; const port = extractPublishedWrapperPortFromLabel(labels) ?? extractPublishedWrapperPort(ports ?? ''); - if (port === null) continue; - rows.push({ containerId, agentSessionId, port }); + containers.push({ + containerId, + agentSessionId, + ...(port === null ? {} : { port }), + }); } - return rows; + return { status: 'ok', containers }; +} + +export async function listWrapperContainers( + executor: DockerExecutor, + options?: { dockerEnv?: Record } +): Promise { + const inspection = await inspectWrapperContainers(executor, options); + return inspection.status === 'ok' ? inspection.containers : []; } function extractLabelValue(labelsField: string, labelKey: string): string | null { @@ -244,13 +284,123 @@ function extractPublishedWrapperPortFromLabel(labelsField: string): number | nul * shape — `id` is the container ID, `command` carries the agent-session marker * for diagnostics. */ +export async function discoverSessionWrappers( + sandbox: SandboxInstance, + sessionId: string, + options?: { dockerEnv?: Record; inspectContainers?: boolean } +): Promise { + let processes: Process[]; + try { + processes = await sandbox.listProcesses(); + } catch (error) { + return { + status: 'inspection-failed', + error: error instanceof Error ? error.message : String(error), + }; + } + + const observed: ObservedWrapper[] = []; + const marker = `${KILO_WRAPPER_SESSION_FLAG} ${sessionId}`; + for (const proc of processes) { + if (proc.command.includes('devcontainer exec')) continue; + if (!proc.command.includes(marker) || !proc.command.includes('kilocode-wrapper')) continue; + if (proc.status !== 'running' && proc.status !== 'starting') continue; + const port = extractWrapperPortFromCommand(proc.command) ?? undefined; + const instanceId = extractWrapperInstanceIdFromCommand(proc.command) ?? undefined; + const instanceGeneration = + extractWrapperInstanceGenerationFromCommand(proc.command) ?? undefined; + observed.push({ + representation: 'process', + id: proc.id, + ...(port !== undefined ? { port } : {}), + ...(instanceId ? { instanceId } : {}), + ...(instanceGeneration !== undefined ? { instanceGeneration } : {}), + }); + } + + if (options?.inspectContainers === false) { + return observed.length === 0 ? { status: 'absent' } : { status: 'present', observed }; + } + + const containerInspection = await inspectWrapperContainers(sandbox, options); + if (containerInspection.status === 'inspection-failed') return containerInspection; + for (const container of containerInspection.containers) { + if (container.agentSessionId !== sessionId) continue; + let result: { exitCode: number; stdout?: string; stderr?: string }; + let dockerEnv: Record; + try { + dockerEnv = options?.dockerEnv ?? dockerSocketEnv(await resolveDockerSocketPath(sandbox)); + result = await sandbox.exec( + `docker exec ${shellQuote(container.containerId)} sh -c ${shellQuote('ps -eo pid=,args=')}`, + { env: dockerEnv } + ); + } catch (error) { + return { + status: 'inspection-failed', + error: error instanceof Error ? error.message : String(error), + }; + } + if (result.exitCode !== 0) { + return { + status: 'inspection-failed', + error: result.stderr?.trim() || `docker exec exited with code ${result.exitCode}`, + }; + } + for (const row of (result.stdout ?? '').split('\n')) { + const parsedRow = row.match(/^\s*(\d+)\s+(.*)$/); + const pid = parsedRow?.[1]; + const command = parsedRow?.[2] ?? row; + if (!command.includes(marker) || !command.includes('kilocode-wrapper')) continue; + + let identitySource = command; + if (pid) { + let environmentResult: { exitCode: number; stdout?: string; stderr?: string }; + try { + environmentResult = await sandbox.exec( + `docker exec ${shellQuote(container.containerId)} sh -c ${shellQuote( + `tr '\\000' ' ' < /proc/${pid}/environ` + )}`, + { env: dockerEnv } + ); + } catch (error) { + return { + status: 'inspection-failed', + error: error instanceof Error ? error.message : String(error), + }; + } + if (environmentResult.exitCode !== 0) { + return { + status: 'inspection-failed', + error: + environmentResult.stderr?.trim() || + `docker exec environment inspection exited with code ${environmentResult.exitCode}`, + }; + } + identitySource = `${environmentResult.stdout ?? ''} ${command}`; + } + const port = extractWrapperPortFromCommand(command) ?? container.port; + const instanceId = extractWrapperInstanceIdFromCommand(identitySource) ?? undefined; + const instanceGeneration = + extractWrapperInstanceGenerationFromCommand(identitySource) ?? undefined; + observed.push({ + representation: 'container', + id: container.containerId, + ...(port !== undefined ? { port } : {}), + ...(instanceId ? { instanceId } : {}), + ...(instanceGeneration !== undefined ? { instanceGeneration } : {}), + }); + } + } + return observed.length === 0 ? { status: 'absent' } : { status: 'present', observed }; +} + export async function findWrapperContainerForSession( executor: DockerExecutor, sessionId: string ): Promise { const containers = await listWrapperContainers(executor); - const match = containers.find(c => c.agentSessionId === sessionId); - if (!match) return null; + const match = containers.find(c => c.agentSessionId === sessionId && c.port !== undefined); + if (!match || match.port === undefined) return null; // Synthesise a Process-shaped record so existing call sites that read // `proc.id` / `proc.command` still work. @@ -279,7 +429,14 @@ export function isWrapperLiveInProcessesOrContainers( containers: LabeledWrapperRow[], sessionId: string ): boolean { - if (findWrapperForSessionInProcesses(processes, sessionId)) return true; + const marker = `${KILO_WRAPPER_SESSION_FLAG} ${sessionId}`; + const hasDirectWrapper = processes.some( + process => + process.command.includes(marker) && + process.command.includes('kilocode-wrapper') && + (process.status === 'running' || process.status === 'starting') + ); + if (hasDirectWrapper) return true; return containers.some(c => c.agentSessionId === sessionId); } @@ -290,73 +447,30 @@ export function getWrapperSessionMarker(sessionId: string): string { return `${KILO_WRAPPER_SESSION_FLAG} ${sessionId}`; } -/** - * Stop a running wrapper for the given session. - * Finds the wrapper process and sends SIGTERM. - * - * @param sandbox - The sandbox instance to search in - * @param sessionId - The cloud-agent session ID - */ -export async function stopWrapper( +export async function stopObservedWrappers( sandbox: SandboxInstance, sessionId: string, - options?: { devcontainer?: { workspacePath: string; configPath?: string } } + observed: ObservedWrapper[], + options?: { force?: boolean; devcontainer?: { workspacePath: string; configPath?: string } } ): Promise { - const existing = await findWrapperForSession(sandbox, sessionId); - if (!existing) { - logger.withFields({ sessionId }).debug('No wrapper found to stop'); - return; + const dockerRows = observed.filter(wrapper => wrapper.representation === 'container'); + const processRows = observed.filter(wrapper => wrapper.representation === 'process'); + const sessionMarker = getWrapperSessionMarker(sessionId); + if (processRows.length > 0) { + await sandbox.exec( + `${options?.force ? 'pkill -9' : 'pkill'} -f -- ${shellQuote(sessionMarker)}` + ); } - const { process: proc, port, kind } = existing; - logger.withFields({ sessionId, port, processId: proc.id, kind }).info('Stopping wrapper'); - try { - if (kind === 'container') { - const sessionMarker = getWrapperSessionMarker(sessionId); - const dockerEnv = dockerSocketEnv(await resolveDockerSocketPath(sandbox)); - if (options?.devcontainer) { - // The wrapper is inside a dev container — outer pkill can't see it. - // Prefer killing just the wrapper process so follow-up executions keep - // using the same devcontainer instead of falling back to the outer image. - // `--config` is required so the CLI keeps applying our remoteUser/ - // remoteEnv overrides; the path is reconstructed from sessionId - // since the override is written deterministically in - // `bringUpDevContainer`. - await sandbox.exec( - [ - 'devcontainer exec', - `--workspace-folder ${shellQuote(options.devcontainer.workspacePath)}`, - `--config ${shellQuote( - getDevContainerOverridePath( - sessionId, - options.devcontainer.workspacePath, - options.devcontainer.configPath - ) - )}`, - `--id-label ${shellQuote(`${KILO_AGENT_SESSION_LABEL}=${sessionId}`)}`, - '--', - 'sh -c', - shellQuote(`pkill -f -- ${shellQuote(sessionMarker)}`), - ].join(' '), - { env: dockerEnv } - ); - } else { - // No devcontainer metadata is available (e.g. older sessions). Kill the - // container as a last-resort cleanup rather than leaving it leaked. - await sandbox.exec(`docker kill ${shellQuote(proc.id)}`, { env: dockerEnv }); - } - } else { - const sessionMarker = getWrapperSessionMarker(sessionId); - await sandbox.exec(`pkill -f -- ${shellQuote(sessionMarker)}`); + if (dockerRows.length > 0) { + const dockerEnv = dockerSocketEnv(await resolveDockerSocketPath(sandbox)); + for (const wrapper of dockerRows) { + const pkill = options?.force ? 'pkill -9' : 'pkill'; + await sandbox.exec( + `docker exec ${shellQuote(wrapper.id)} sh -c ${shellQuote( + `${pkill} -f -- ${shellQuote(sessionMarker)}` + )}`, + { env: dockerEnv } + ); } - logger.withFields({ sessionId, port, kind }).info('Wrapper stopped'); - } catch (error) { - logger - .withFields({ - sessionId, - port, - kind, - error: error instanceof Error ? error.message : String(error), - }) - .warn('Error stopping wrapper'); } } diff --git a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts index cfb2ccead9..7d5cd06895 100644 --- a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts +++ b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts @@ -91,6 +91,7 @@ import { } from '../session/session-message-queue.js'; import { clearWrapperRuntimeIdentity, + getWrapperLease, getWrapperRuntimeState, } from '../session/wrapper-runtime-state.js'; import { @@ -111,10 +112,16 @@ import { type AgentRuntime, type AgentRuntimeAcceptedDelivery, type AgentRuntimeOrchestrator, - type AgentRuntimeStopReason, } from '../session/agent-runtime.js'; import { createWrapperSupervisor, type WrapperSupervisor } from '../session/wrapper-supervisor.js'; import { emitRunStateReport } from '../telemetry/queue-reports.js'; +import { createAgentSandbox } from '../agent-sandbox/factory.js'; +import type { + StopWrappersResult, + WrapperObservation, + WrapperStopReason, + WrapperStopTarget, +} from '../agent-sandbox/protocol.js'; // --------------------------------------------------------------------------- // Alarm Constants @@ -130,6 +137,7 @@ const EVENT_RETENTION_MS = Limits.SESSION_TTL_MS; /** Storage key for tracking last activity timestamp */ const LAST_ACTIVITY_KEY = 'last_activity'; +const EXPLICIT_DELETION_PENDING_KEY = 'explicit_deletion_pending'; /** Kilo server idle timeout: 15 minutes */ const KILO_SERVER_IDLE_TIMEOUT_MS_DEFAULT = 15 * 60 * 1000; @@ -257,6 +265,12 @@ export class CloudAgentSession extends DurableObject { private ingestHandlerSessionId?: SessionId; private sessionId?: SessionId; private orchestrator?: AgentRuntimeOrchestrator; + private physicalWrapperObserver?: () => Promise; + private physicalWrapperStopper?: (request: { + target: WrapperStopTarget; + attemptId: string; + reason: WrapperStopReason; + }) => Promise; private agentRuntime?: AgentRuntime; private messageSettlementOutbox?: MessageSettlementOutbox; private sessionMessageQueue?: SessionMessageQueue; @@ -484,6 +498,13 @@ export class CloudAgentSession extends DurableObject { sendToWrapper: (ingestTagId, command, fence) => this.sendToWrapper(ingestTagId, command, fence), getOrchestratorOverride: () => this.orchestrator, + discoverSessionWrappers: metadata => + this.physicalWrapperObserver + ? this.physicalWrapperObserver() + : this.orchestrator + ? Promise.resolve({ status: 'absent' }) + : createAgentSandbox(this.env, metadata).discoverSessionWrappers(), + requestAlarmAtOrBefore: deadline => this.scheduleAlarmAtOrBefore(deadline), }); } @@ -496,7 +517,6 @@ export class CloudAgentSession extends DurableObject { storage: this.ctx.storage, agentRuntime: { sendPing: ingestTagId => this.getAgentRuntime().sendPing(ingestTagId), - stopWrapperProcess: reason => this.stopCurrentWrapperProcess(reason), }, messageSettlementOutbox: this.getMessageSettlementOutbox(), sessionMessageQueue: this.getSessionMessageQueue(), @@ -511,6 +531,17 @@ export class CloudAgentSession extends DurableObject { hasActiveIngestConnection: async params => (await this.getIngestHandler()).hasActiveConnection(params), clearInterruptRequest: () => this.executionQueries.clearInterrupt(), + stopWrappers: async request => { + if (this.physicalWrapperStopper) return this.physicalWrapperStopper(request); + if (this.orchestrator || (!this.env.Sandbox && !this.env.SandboxSmall)) { + return { status: 'absent' }; + } + const metadata = await this.getMetadata(); + if (!metadata) + return { status: 'inspection-failed', error: 'Session metadata unavailable' }; + return createAgentSandbox(this.env, metadata).stopWrappers(request); + }, + requestAlarmAtOrBefore: deadline => this.scheduleAlarmAtOrBefore(deadline), getSessionIdForLogs: () => this.sessionId, }); } @@ -1168,12 +1199,20 @@ export class CloudAgentSession extends DurableObject { private async interruptAcceptedWrapperMessages(): Promise<{ acceptedMessageCount: number; wrapperCommandSent: boolean; + physicalWrapperStopRequested: boolean; }> { const state = await getWrapperRuntimeState(this.ctx.storage); const acceptedMessages = await listNonTerminalAcceptedMessages( this.ctx.storage, state.wrapperRunId ); + const supervisor = this.getWrapperSupervisor(); + const requiresPhysicalWrapperStop = + acceptedMessages.length > 0 || + (state.wrapperRunId !== undefined && state.wrapperConnectionId !== undefined); + if (requiresPhysicalWrapperStop) { + await supervisor.requestPhysicalWrapperStop('user-interrupt'); + } for (const msg of acceptedMessages) { const transition = await this.getMessageSettlementOutbox().persistTerminalTransition( msg.messageId, @@ -1205,30 +1244,34 @@ export class CloudAgentSession extends DurableObject { } } - const interrupt = await this.getAgentRuntime().interruptWrapper(); - if (acceptedMessages.length > 0 && state.wrapperConnectionId) { - await clearWrapperRuntimeIdentity( - this.ctx.storage, - { - wrapperGeneration: state.wrapperGeneration, - wrapperConnectionId: state.wrapperConnectionId, - }, - { incrementGeneration: true } - ); - try { - await this.stopCurrentWrapperProcess('user-interrupt'); - } catch (error) { - logger - .withFields({ - sessionId: this.sessionId, - error: error instanceof Error ? error.message : String(error), - }) - .warn('Failed to stop interrupted wrapper process after fencing'); + let wrapperCommandSent = false; + try { + wrapperCommandSent = (await this.getAgentRuntime().interruptWrapper()).commandSent; + } catch (error) { + logger + .withFields({ + sessionId: this.sessionId, + error: error instanceof Error ? error.message : String(error), + }) + .warn('Failed to signal wrapper interruption; physical cleanup will continue'); + } + if (requiresPhysicalWrapperStop) { + if (state.wrapperConnectionId) { + await clearWrapperRuntimeIdentity( + this.ctx.storage, + { + wrapperGeneration: state.wrapperGeneration, + wrapperConnectionId: state.wrapperConnectionId, + }, + { incrementGeneration: true } + ); } + await supervisor.runMaintenance(Date.now()); } return { acceptedMessageCount: acceptedMessages.length, - wrapperCommandSent: interrupt.commandSent, + wrapperCommandSent, + physicalWrapperStopRequested: requiresPhysicalWrapperStop, }; } @@ -1239,17 +1282,24 @@ export class CloudAgentSession extends DurableObject { }> { let acceptedMessageCount = 0; let wrapperCommandSent = false; + let physicalWrapperStopRequested = false; const clearedMessages = await this.getSessionMessageQueue().interruptPendingQueuedMessages( async () => { const acceptedInterruption = await this.interruptAcceptedWrapperMessages(); acceptedMessageCount = acceptedInterruption.acceptedMessageCount; wrapperCommandSent = acceptedInterruption.wrapperCommandSent; + physicalWrapperStopRequested = acceptedInterruption.physicalWrapperStopRequested; } ); await this.finalizeIdleBatchCallbackIfReady({ allowWithoutObservedIdle: true }); - if (!wrapperCommandSent && clearedMessages.length === 0 && acceptedMessageCount === 0) { + if ( + !wrapperCommandSent && + !physicalWrapperStopRequested && + clearedMessages.length === 0 && + acceptedMessageCount === 0 + ) { return { success: false, message: 'No accepted wrapper messages or pending queued messages' }; } @@ -1349,15 +1399,54 @@ export class CloudAgentSession extends DurableObject { } } + private async finalizeSessionDeletion( + reason: 'explicit' | 'retention-expired' + ): Promise { + const metadata = await this.getMetadata(); + if (!metadata) { + if ((await getWrapperLease(this.ctx.storage)).state !== 'none') { + await this.scheduleAlarmAtOrBefore(Date.now() + 1_000); + return false; + } + } else { + const supervisor = this.getWrapperSupervisor(); + await supervisor.requestPhysicalWrapperStop('session-delete', { kind: 'session' }); + await supervisor.runMaintenance(Date.now()); + if ((await getWrapperLease(this.ctx.storage)).state !== 'none') { + if (reason === 'explicit') { + await this.scheduleAlarmAtOrBefore(Date.now() + 1_000); + } + return false; + } + if (!this.orchestrator && (this.env.Sandbox || this.env.SandboxSmall)) { + await createAgentSandbox(this.env, metadata).delete(reason); + } + } + + await this.ctx.storage.deleteAlarm(); + await this.ctx.storage.deleteAll(); + return true; + } + + private async isExplicitDeletionPending(): Promise { + return (await this.ctx.storage.get(EXPLICIT_DELETION_PENDING_KEY)) === true; + } + + private async deletionPendingAdmissionFailure(): Promise { + return (await this.isExplicitDeletionPending()) + ? { success: false, code: 'NOT_FOUND', error: 'Session deletion is pending' } + : null; + } + /** - * Delete session and all associated data. + * Delete session only after physical wrapper absence has been verified. */ async deleteSession(): Promise { logger.info('Explicit DELETE requested for Durable Object'); - - // Must delete alarm before deleteAll - await this.ctx.storage.deleteAlarm(); - await this.ctx.storage.deleteAll(); + await this.ctx.storage.put(EXPLICIT_DELETION_PENDING_KEY, true); + if (!(await this.finalizeSessionDeletion('explicit'))) { + throw new Error('Session deletion pending physical wrapper cleanup'); + } } /** @@ -1366,6 +1455,9 @@ export class CloudAgentSession extends DurableObject { * delivers the first message. */ async registerSession(input: GroupedRegisterSessionInput): Promise { + if (await this.isExplicitDeletionPending()) { + return { success: false, error: 'Session deletion is pending' }; + } await this.requireSessionId(input.identity.sessionId as SessionId); const existing = await this.ctx.storage.get('metadata'); if (existing) { @@ -1476,6 +1568,8 @@ export class CloudAgentSession extends DurableObject { async createSessionWithInitialAdmission( input: CreateSessionWithInitialAdmissionInput ): Promise { + const deletionPending = await this.deletionPendingAdmissionFailure(); + if (deletionPending) return deletionPending; const initialTurn = input.message.initialTurn; const admitInitialTurn = () => this.getSessionMessageQueue().admitAcceptedMessage({ @@ -1705,6 +1799,17 @@ export class CloudAgentSession extends DurableObject { let alarmWorkFailed = false; try { + if (await this.isExplicitDeletionPending()) { + logger.withFields({ sessionId: this.sessionId }).info('Resuming explicit session deletion'); + if (await this.finalizeSessionDeletion('explicit')) { + return; + } + logger + .withFields({ sessionId: this.sessionId }) + .info('Postponing explicit session deletion until wrapper cleanup confirms absence'); + return; + } + // Check if session should be deleted due to inactivity (90 days) const lastActivity = await this.ctx.storage.get(LAST_ACTIVITY_KEY); if (lastActivity && now - lastActivity > Limits.SESSION_TTL_MS) { @@ -1712,9 +1817,12 @@ export class CloudAgentSession extends DurableObject { .withFields({ sessionId: this.sessionId, lastActivity }) .info('Deleting session due to inactivity'); - await this.ctx.storage.deleteAlarm(); - await this.ctx.storage.deleteAll(); - return; + if (await this.finalizeSessionDeletion('retention-expired')) { + return; + } + logger + .withFields({ sessionId: this.sessionId }) + .info('Postponing inactive session deletion until wrapper cleanup confirms absence'); } await this.getWrapperSupervisor().runMaintenance(now); @@ -1919,27 +2027,19 @@ export class CloudAgentSession extends DurableObject { }) .info('Stopping idle kilo server'); - const stopped = await this.getAgentRuntime().stopWrapperProcess('idle-timeout'); + await this.getWrapperSupervisor().requestPhysicalWrapperStop('idle-timeout'); const updated = { ...metadata, lifecycle: { ...metadata.lifecycle, - ...(stopped ? { kiloServerLastActivity: undefined } : {}), + kiloServerLastActivity: undefined, version: Date.now(), }, }; await this.updateMetadata(updated); - - if (stopped) { - await this.getMessageSettlementOutbox().releaseWrapperTerminalWaitForIdleBatch(); - await this.finalizeIdleBatchCallbackIfReady({ allowWithoutObservedIdle: true }); - logger - .withFields({ sessionId: this.sessionId }) - .info('Idle kilo server stopped successfully'); - return; - } - - logger.withFields({ sessionId: this.sessionId }).warn('Failed to stop idle kilo server'); + await this.getMessageSettlementOutbox().releaseWrapperTerminalWaitForIdleBatch(); + await this.finalizeIdleBatchCallbackIfReady({ allowWithoutObservedIdle: true }); + logger.withFields({ sessionId: this.sessionId }).info('Idle kilo server cleanup requested'); } /** @@ -2233,6 +2333,9 @@ export class CloudAgentSession extends DurableObject { const pendingCount = await countPendingSessionMessages(this.ctx.storage); if (pendingCount > 0) return true; + const physicalLease = await getWrapperLease(this.ctx.storage); + if (physicalLease.state !== 'none') return true; + const state = await getWrapperRuntimeState(this.ctx.storage); if (!state.wrapperConnectionId) return false; @@ -2260,10 +2363,6 @@ export class CloudAgentSession extends DurableObject { await this.getMessageSettlementOutbox().retryPendingCallbacks(now); } - private async stopCurrentWrapperProcess(reason: AgentRuntimeStopReason): Promise { - return this.getAgentRuntime().stopWrapperProcess(reason); - } - /** * Update execution heartbeat timestamp. */ @@ -2459,7 +2558,8 @@ export class CloudAgentSession extends DurableObject { async admitSubmittedMessage( request: SubmittedSessionMessageRequest ): Promise { - return this.getSessionMessageQueue().admitSubmittedMessage(request); + const deletionPending = await this.deletionPendingAdmissionFailure(); + return deletionPending ?? this.getSessionMessageQueue().admitSubmittedMessage(request); } async replayPreparedInitialMessage( @@ -2476,6 +2576,8 @@ export class CloudAgentSession extends DurableObject { async admitPreparedInitialMessage( request: LegacyRegisteredInitialAdmissionRequest ): Promise { + const deletionPending = await this.deletionPendingAdmissionFailure(); + if (deletionPending) return deletionPending; const metadata = await this.getMetadata(); if (!metadata) return { success: false, code: 'NOT_FOUND', error: 'Session not found' }; const initialMessage = metadata.initialMessage; diff --git a/services/cloud-agent-next/src/router.test.ts b/services/cloud-agent-next/src/router.test.ts index 9f2d11e6e3..f5b1e0217d 100644 --- a/services/cloud-agent-next/src/router.test.ts +++ b/services/cloud-agent-next/src/router.test.ts @@ -399,13 +399,10 @@ describe('router sessionId validation', () => { 'test-user-123', sessionId ); - expect(getSandbox).toHaveBeenCalledWith( - mockContext.env.Sandbox, - expect.stringMatching(/^org-[0-9a-f]{48}$/) - ); + expect(getSandbox).not.toHaveBeenCalled(); // eslint-disable-next-line @typescript-eslint/unbound-method const sandboxDelete = vi.mocked(mockSandbox.deleteSession); - expect(sandboxDelete).toHaveBeenCalledWith(sessionId); + expect(sandboxDelete).not.toHaveBeenCalled(); expect(cloudAgentSession.idFromName).toHaveBeenCalledWith( `${metadata.identity.userId}:${sessionId}` ); @@ -457,11 +454,7 @@ describe('router sessionId validation', () => { const result = await caller.deleteSession({ sessionId }); expect(result).toEqual({ success: true }); - // Should use usr prefix for personal accounts - expect(getSandbox).toHaveBeenCalledWith( - mockContext.env.Sandbox, - expect.stringMatching(/^usr-[0-9a-f]{48}$/) - ); + expect(getSandbox).not.toHaveBeenCalled(); }); it('should successfully delete session with botId', async () => { @@ -480,11 +473,7 @@ describe('router sessionId validation', () => { const result = await caller.deleteSession({ sessionId }); expect(result).toEqual({ success: true }); - // Should include bot suffix - expect(getSandbox).toHaveBeenCalledWith( - mockContext.env.Sandbox, - expect.stringMatching(/^bot-[0-9a-f]{48}$/) - ); + expect(getSandbox).not.toHaveBeenCalled(); }); it('should route per-session sandbox ID to SandboxSmall namespace', async () => { @@ -504,11 +493,7 @@ describe('router sessionId validation', () => { const result = await caller.deleteSession({ sessionId }); expect(result).toEqual({ success: true }); - // ses- prefixed sandbox IDs should route to SandboxSmall, not Sandbox - expect(getSandbox).toHaveBeenCalledWith( - mockContext.env.SandboxSmall, - perSessionSandboxId - ); + expect(getSandbox).not.toHaveBeenCalled(); }); }); @@ -528,8 +513,8 @@ describe('router sessionId validation', () => { }); }); - describe('sandbox deletion failure handling', () => { - it('should continue cleanup when sandbox deletion fails', async () => { + describe('provider deletion ownership', () => { + it('does not perform provider deletion outside the Durable Object', async () => { const sessionId: SessionId = 'agent_aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee'; const metadata = legacySessionMetadata({ version: 123456789, @@ -540,9 +525,7 @@ describe('router sessionId validation', () => { }); vi.mocked(fetchSessionMetadata).mockResolvedValue(metadata); - // Sandbox deletion fails - mockSandbox.deleteSession = vi.fn().mockRejectedValue(new Error('Sandbox unreachable')); - // DO cleanup succeeds + mockSandbox.deleteSession = vi.fn().mockRejectedValue(new Error('must not be called')); const deleteSessionMock = vi.mocked(cloudAgentSession.get).mockReturnValue({ deleteSession: vi.fn().mockResolvedValue(undefined), markAsInterrupted: vi.fn().mockResolvedValue(undefined), @@ -550,13 +533,12 @@ describe('router sessionId validation', () => { const result = await caller.deleteSession({ sessionId }); - // Should still succeed overall expect(result).toEqual({ success: true }); - // Should have attempted both cleanups // eslint-disable-next-line @typescript-eslint/unbound-method const sandboxDelete = vi.mocked(mockSandbox.deleteSession); - expect(sandboxDelete).toHaveBeenCalled(); - expect(deleteSessionMock().deleteSession).toHaveBeenCalledWith(); + expect(sandboxDelete).not.toHaveBeenCalled(); + expect(getSandbox).not.toHaveBeenCalled(); + expect(deleteSessionMock().deleteSession).toHaveBeenCalled(); }); }); @@ -806,6 +788,29 @@ describe('router sessionId validation', () => { caller = appRouter.createCaller(mockContext); }); + it('routes accepted interruption to the Durable Object without provider stopping', async () => { + const sessionId: SessionId = 'agent_87654321-1234-1234-1234-123456789abc'; + const metadata = legacySessionMetadata({ + version: 123456789, + sessionId, + orgId: 'org-123', + userId: 'test-user-123', + timestamp: 123456789, + }); + vi.mocked(fetchSessionMetadata).mockResolvedValue(metadata); + + const result = await caller.interruptSession({ sessionId }); + + expect(result).toEqual({ + success: true, + message: 'Session interruption accepted', + processesFound: false, + }); + expect(mockSessionStub.interruptExecution).toHaveBeenCalled(); + expect(getSandbox).not.toHaveBeenCalled(); + expect(interruptMock).not.toHaveBeenCalled(); + }); + it('short-circuits queued-only interrupts before creating a sandbox session', async () => { const sessionId: SessionId = 'agent_12345678-1234-1234-1234-123456789abc'; const metadata = legacySessionMetadata({ @@ -826,9 +831,10 @@ describe('router sessionId validation', () => { expect(result).toEqual({ success: true, - message: 'Queued session messages interrupted', + message: 'Session interruption accepted', processesFound: false, }); + expect(mockSessionStub.markAsInterrupted).toHaveBeenCalled(); expect(mockSessionStub.interruptExecution).toHaveBeenCalled(); expect(getOrCreateSessionMock).not.toHaveBeenCalled(); diff --git a/services/cloud-agent-next/src/router/handlers/session-management.ts b/services/cloud-agent-next/src/router/handlers/session-management.ts index 84fe56e5db..db20b4cb95 100644 --- a/services/cloud-agent-next/src/router/handlers/session-management.ts +++ b/services/cloud-agent-next/src/router/handlers/session-management.ts @@ -1,8 +1,8 @@ import { TRPCError } from '@trpc/server'; import * as z from 'zod'; -import { getSandbox } from '@cloudflare/sandbox'; +import { createAgentSandbox } from '../../agent-sandbox/factory.js'; import { logger, withLogTags } from '../../logger.js'; -import { generateSandboxId, getSandboxNamespace } from '../../sandbox-id.js'; +import { generateSandboxId } from '../../sandbox-id.js'; import type { SessionId, InterruptResult, TRPCContext } from '../../types.js'; import type { SandboxId } from '../../types.js'; import { @@ -10,7 +10,6 @@ import { SessionService, fetchSessionMetadata, } from '../../session-service.js'; -import { cleanupWorkspace, getSessionWorkspacePath, getSessionHomePath } from '../../workspace.js'; import { withDORetry } from '../../utils/do-retry.js'; import { protectedProcedure, publicProcedure, internalApiProtectedProcedure } from '../auth.js'; import { @@ -57,39 +56,6 @@ async function deleteSessionResources( return { success: true, message: 'Session not found or already deleted' }; } - const sandboxId: SandboxId = - metadata.workspace?.sandboxId ?? - (await generateSandboxId( - env.PER_SESSION_SANDBOX_ORG_IDS, - metadata.identity.orgId, - userId, - metadata.identity.sessionId, - metadata.identity.botId - )); - logger.setTags({ sandboxId, orgId: metadata.identity.orgId ?? '(personal)' }); - const sandbox = getSandbox(getSandboxNamespace(env, sandboxId), sandboxId); - const workspacePath = getSessionWorkspacePath(metadata.identity.orgId, userId, sessionId); - const sessionHome = getSessionHomePath(sessionId); - - try { - const session = await sandbox.getSession(sessionId); - await cleanupWorkspace(session, workspacePath, sessionHome); - logger.info('Workspace directories cleaned up'); - } catch (error) { - logger - .withFields({ error: error instanceof Error ? error.message : String(error) }) - .warn('Failed to clean up workspace directories, continuing with deletion'); - } - - await sandbox - .deleteSession(sessionId) - .then(() => logger.info('Cloudflare sandbox session deleted')) - .catch(error => { - logger - .withFields({ error: error instanceof Error ? error.message : String(error) }) - .warn('Failed to delete Cloudflare sandbox session, continuing with cleanup'); - }); - try { const doKey = `${userId}:${sessionId}`; await withDORetry( @@ -126,7 +92,6 @@ async function deleteSessionResources( * These handlers manage session lifecycle (delete, interrupt, logs) and health checks. */ export function createSessionManagementHandlers() { - const INTERRUPT_GRACE_MS = 2000; return { /** * Delete a session and clean up all associated resources. @@ -160,14 +125,9 @@ export function createSessionManagementHandlers() { }), /** - * Interrupt a running session by killing all associated kilocode processes. - * - * This endpoint allows clients to stop running executions in a session without - * deleting the session itself. Useful for canceling long-running or stuck operations. - * - * Idempotency: - * - Returns success even if no processes are found (already stopped or none running) - * - Safe to call multiple times for the same session + * Interrupt current session work through the owning Durable Object. + * The DO may signal a connected wrapper immediately and durably supervises + * physical cleanup without letting this route issue provider teardown. */ interruptSession: protectedProcedure .input( @@ -219,89 +179,14 @@ export function createSessionManagementHandlers() { .info('No accepted current messages or pending queued messages to interrupt'); } - const targetExecutionId = interruptResult.executionId; - if (!targetExecutionId) { - logger.info('Session interruption completed'); - return { - success: true, - message: 'Queued session messages interrupted', - processesFound: false, - }; - } - - const sandboxId: SandboxId = - metadata.workspace?.sandboxId ?? - (await generateSandboxId( - env.PER_SESSION_SANDBOX_ORG_IDS, - metadata.identity.orgId, - userId, - metadata.identity.sessionId, - metadata.identity.botId - )); - - logger.setTags({ sandboxId, orgId: metadata.identity.orgId ?? '(personal)' }); - - const sandbox = getSandbox(getSandboxNamespace(env, sandboxId), sandboxId); - - // Build session context for interrupt service - const sessionService = new SessionService(); - const context = sessionService.buildContext({ - sandboxId, - orgId: metadata.identity.orgId, - userId, - sessionId, - upstreamBranch: metadata.repository?.upstreamBranch, - botId: metadata.identity.botId, - }); - - await scheduler.wait(INTERRUPT_GRACE_MS); - - // Get or create the session to use for killing processes - const session = await sessionService.getOrCreateSession({ - sandbox, - context, - env, - originalToken: ctx.authToken, - originalOrgId: metadata.identity.orgId, - createdOnPlatform: metadata.identity.createdOnPlatform, - appendSystemPrompt: metadata.agent?.appendSystemPrompt, - profile: readProfileBundle(metadata), - }); - - // Kill all kilocode processes in this session - // Use pkill method as a temporary workaround for sandbox API reliability issues - const usePkill = true; - const result = await SessionService.interrupt( - sandbox, - session, - context, - usePkill, - targetExecutionId - ); - logger.info('Session interruption completed'); - - // If no processes were found but there's still a runtime execution, - // the wrapper is already dead - fail the stale execution immediately. - // Note: pkill always returns killedProcessIds: [], so we check - // processesFound instead to distinguish "killed" from "nothing to kill". - if (!result.processesFound && targetExecutionId) { - logger - .withFields({ executionId: targetExecutionId }) - .info('No processes found during interrupt - failing stale runtime execution'); - - await withDORetry( - getStub, - stub => - stub.failExecutionRpc({ - executionId: targetExecutionId, - error: 'Interrupted - no running processes found', - }), - 'failExecutionRpc' - ); - } - - return result; + return { + success: interruptResult.success, + message: interruptResult.success + ? 'Session interruption accepted' + : (interruptResult.message ?? 'No session work to interrupt'), + processesFound: false, + }; } catch (error) { const errorMsg = error instanceof Error ? error.message : String(error); logger.withFields({ error: errorMsg }).error('Session interruption failed'); @@ -489,10 +374,9 @@ export function createSessionManagementHandlers() { const activeExecutionStatus = activeMessageWork?.status; const executionHealth = activeMessageWork?.health ?? 'none'; - const sandbox = getSandbox(getSandboxNamespace(env, sandboxId), sandboxId); let sandboxStatus: 'healthy' | 'unreachable' = 'healthy'; try { - await sandbox.listProcesses(); + await createAgentSandbox(env, metadata).probeHealth(); } catch (error) { sandboxStatus = 'unreachable'; logger @@ -606,105 +490,30 @@ export function createSessionManagementHandlers() { orgId: sessionService.metadata?.identity.orgId ?? '(personal)', }); - const sandbox = getSandbox(getSandboxNamespace(env, sandboxId), sandboxId); - - // Get or create a session to read files - const context = sessionService.buildContext({ - sandboxId, - orgId: sessionService.metadata?.identity.orgId, - userId, - sessionId, - botId: sessionService.metadata?.identity.botId, - }); - - const session = await sessionService.getOrCreateSession({ - sandbox, - context, - env, - originalToken: ctx.authToken, - originalOrgId: sessionService.metadata?.identity.orgId, - createdOnPlatform: sessionService.metadata?.identity.createdOnPlatform, - appendSystemPrompt: sessionService.metadata?.agent?.appendSystemPrompt, - profile: sessionService.metadata - ? readProfileBundle(sessionService.metadata) - : undefined, - }); - - // Discover all log files from the sandbox - const logPaths: string[] = []; - - // 1. Wrapper logs: /tmp/kilocode-wrapper-*.log (one per execution) - try { - const tmpFiles = await session.listFiles('/tmp'); - if (tmpFiles.success) { - for (const f of tmpFiles.files) { - if ( - f.type === 'file' && - f.name.startsWith('kilocode-wrapper-') && - f.name.endsWith('.log') - ) { - logPaths.push(f.absolutePath); - } - } - } - } catch { - logger.debug('Could not list /tmp for wrapper logs'); - } - - // 2. CLI logs: {sessionHome}/.local/share/kilo/log/ (matches wrapper R2 uploader) - const sessionHome = getSessionHomePath(sessionId); - const cliLogsDir = `${sessionHome}/.local/share/kilo/log`; - try { - const cliFiles = await session.listFiles(cliLogsDir, { recursive: true }); - if (cliFiles.success) { - for (const f of cliFiles.files) { - if (f.type === 'file') { - logPaths.push(f.absolutePath); - } - } - } - } catch { - logger.debug('Could not list CLI logs directory', { cliLogsDir }); - } - - // Read all discovered files in parallel (best-effort per file) - const files: Record = {}; - const readResults = await Promise.allSettled( - logPaths.map(async path => { - const fileInfo = await session.readFile(path, { encoding: 'utf-8' }); - return { path, content: fileInfo.content }; - }) - ); - for (const result of readResults) { - if (result.status === 'fulfilled') { - files[result.value.path] = result.value.content; - } + const metadata = sessionService.metadata; + if (!metadata) { + throw new TRPCError({ + code: 'PRECONDITION_FAILED', + message: 'Session metadata is invalid or unavailable.', + }); } - // Fetch running processes (best-effort) - let processes: Array<{ pid: number; command: string; status: string }> | undefined; - try { - type ProcessInfo = { id: string; status: string; command: string }; - const allProcesses = (await sandbox.listProcesses()) as ProcessInfo[]; - processes = allProcesses.map((p: ProcessInfo) => ({ - pid: parseInt(p.id, 10) || 0, - command: p.command, - status: p.status, - })); - } catch (err) { - logger.debug('Could not fetch sandbox processes', { - error: err instanceof Error ? err.message : String(err), + const logs = await createAgentSandbox(env, metadata).readWrapperLogs(); + if (!logs) { + throw new TRPCError({ + code: 'PRECONDITION_FAILED', + message: 'Wrapper logs are unavailable because the session wrapper is not running', }); } logger.info('Successfully retrieved session logs', { - fileCount: Object.keys(files).length, + fileCount: Object.keys(logs.files).length, }); return { sessionId, - files, - processes, + files: logs.files, + processes: logs.processes, }; }); }), diff --git a/services/cloud-agent-next/src/router/handlers/session-questions.ts b/services/cloud-agent-next/src/router/handlers/session-questions.ts index 5ae767bcd9..649f024513 100644 --- a/services/cloud-agent-next/src/router/handlers/session-questions.ts +++ b/services/cloud-agent-next/src/router/handlers/session-questions.ts @@ -1,63 +1,29 @@ import { TRPCError } from '@trpc/server'; import * as z from 'zod'; -import { getSandbox } from '@cloudflare/sandbox'; +import { createAgentSandbox } from '../../agent-sandbox/factory.js'; import { logger, withLogTags } from '../../logger.js'; -import { generateSandboxId, getSandboxNamespace } from '../../sandbox-id.js'; -import type { SessionId, SandboxId, Env } from '../../types.js'; -import { SessionService, fetchSessionMetadata } from '../../session-service.js'; +import type { SessionId, Env } from '../../types.js'; +import { fetchSessionMetadata } from '../../session-service.js'; import { protectedProcedure } from '../auth.js'; import { sessionIdSchema } from '../schemas.js'; -import { findWrapperForSession } from '../../kilo/wrapper-manager.js'; -import { WrapperClient } from '../../kilo/wrapper-client.js'; +import type { WrapperClient } from '../../kilo/wrapper-client.js'; async function resolveWrapperClient(opts: { sessionId: SessionId; userId: string; env: Env; - authToken: string; }): Promise { - const { sessionId, userId, env, authToken } = opts; - + const { sessionId, userId, env } = opts; const metadata = await fetchSessionMetadata(env, userId, sessionId); if (!metadata) { throw new TRPCError({ code: 'NOT_FOUND', message: 'Session not found' }); } - const sandboxId: SandboxId = - metadata.workspace?.sandboxId ?? - (await generateSandboxId( - env.PER_SESSION_SANDBOX_ORG_IDS, - metadata.identity.orgId, - userId, - metadata.identity.sessionId, - metadata.identity.botId - )); - const sandbox = getSandbox(getSandboxNamespace(env, sandboxId), sandboxId); - - const wrapperInfo = await findWrapperForSession(sandbox, sessionId); - if (!wrapperInfo) { + const wrapperClient = await createAgentSandbox(env, metadata).getRunningWrapper(); + if (!wrapperClient) { throw new TRPCError({ code: 'NOT_FOUND', message: 'No wrapper found for session' }); } - - const sessionService = new SessionService(); - const context = sessionService.buildContext({ - sandboxId, - orgId: metadata.identity.orgId, - userId, - sessionId, - upstreamBranch: metadata.repository?.upstreamBranch, - botId: metadata.identity.botId, - }); - - const session = await sessionService.getOrCreateSession({ - sandbox, - context, - env, - originalToken: authToken, - originalOrgId: metadata.identity.orgId, - }); - - return new WrapperClient({ session, port: wrapperInfo.port }); + return wrapperClient; } export function createSessionQuestionHandlers() { @@ -77,12 +43,7 @@ export function createSessionQuestionHandlers() { logger.setTags({ userId, sessionId }); try { - const wrapperClient = await resolveWrapperClient({ - sessionId, - userId, - env, - authToken: ctx.authToken, - }); + const wrapperClient = await resolveWrapperClient({ sessionId, userId, env }); const result = await wrapperClient.answerQuestion(input.questionId, input.answers); logger .withFields({ questionId: input.questionId, success: result.success }) @@ -114,12 +75,7 @@ export function createSessionQuestionHandlers() { logger.setTags({ userId, sessionId }); try { - const wrapperClient = await resolveWrapperClient({ - sessionId, - userId, - env, - authToken: ctx.authToken, - }); + const wrapperClient = await resolveWrapperClient({ sessionId, userId, env }); const result = await wrapperClient.rejectQuestion(input.questionId); logger .withFields({ questionId: input.questionId, success: result.success }) @@ -153,12 +109,7 @@ export function createSessionQuestionHandlers() { logger.setTags({ userId, sessionId }); try { - const wrapperClient = await resolveWrapperClient({ - sessionId, - userId, - env, - authToken: ctx.authToken, - }); + const wrapperClient = await resolveWrapperClient({ sessionId, userId, env }); const result = await wrapperClient.answerPermission(input.permissionId, input.response); logger .withFields({ permissionId: input.permissionId, success: result.success }) diff --git a/services/cloud-agent-next/src/sandbox-recovery.test.ts b/services/cloud-agent-next/src/sandbox-recovery.test.ts index c616721b36..1fd319c366 100644 --- a/services/cloud-agent-next/src/sandbox-recovery.test.ts +++ b/services/cloud-agent-next/src/sandbox-recovery.test.ts @@ -22,7 +22,12 @@ import { withPreparationInfrastructureRecovery, } from './sandbox-recovery.js'; import { WrapperNotReadyError } from './kilo/wrapper-client.js'; -import { WorkspaceFilesystemPreparationError } from './workspace-errors.js'; +import { + SandboxCapacityInspectionError, + WorkspaceCapacityAdmissionRejectedError, + WorkspaceCapacityInspectionUnavailableError, + WorkspaceFilesystemPreparationError, +} from './workspace-errors.js'; describe('sandbox recovery', () => { it('classifies sandbox SDK internal server errors', () => { @@ -104,7 +109,7 @@ describe('sandbox recovery', () => { await expect( withPreparationInfrastructureRecovery( { - sandbox, + deleteSandbox: () => sandbox.destroy(), sandboxId: 'ses-test', sessionId: 'agent_test', phase: 'asyncPreparation', @@ -149,7 +154,7 @@ describe('sandbox recovery', () => { await expect( withPreparationInfrastructureRecovery( { - sandbox, + deleteSandbox: () => sandbox.destroy(), sandboxId: 'ses-test', sessionId: 'agent_test', phase: 'asyncPreparation', @@ -176,7 +181,7 @@ describe('sandbox recovery', () => { await expect( withPreparationInfrastructureRecovery( { - sandbox, + deleteSandbox: () => sandbox.destroy(), sandboxId: 'ses-test', sessionId: 'agent_test', phase: 'asyncPreparation', @@ -194,11 +199,82 @@ describe('sandbox recovery', () => { expect(mockInfo).toHaveBeenCalledWith('Destroyed sandbox after workspace Git probe timeout'); }); + it('destroys sandbox when capacity inspection reports filesystem unusable', async () => { + const sandbox = { destroy: vi.fn().mockResolvedValue(undefined) }; + const error = new SandboxCapacityInspectionError( + 'Disk capacity inspection cannot run because the sandbox filesystem is unusable', + new Error('ENOSPC: no space left on device') + ); + + await expect( + withPreparationInfrastructureRecovery( + { + deleteSandbox: () => sandbox.destroy(), + sandboxId: 'ses-test', + sessionId: 'agent_test', + phase: 'asyncPreparation', + }, + async () => { + throw error; + } + ) + ).rejects.toBe(error); + + expect(sandbox.destroy).toHaveBeenCalledOnce(); + expect(mockError).toHaveBeenCalledWith( + 'Sandbox capacity inspection failed; destroying unusable sandbox' + ); + }); + + it('does not destroy shared sandbox for low-capacity admission rejection', async () => { + const sandbox = { destroy: vi.fn().mockResolvedValue(undefined) }; + const error = new WorkspaceCapacityAdmissionRejectedError({ + availableMB: 900, + thresholdMB: 2048, + cleaned: 0, + skipped: 2, + }); + + const destroyed = await destroySandboxAfterPreparationInfrastructureFailure( + { + deleteSandbox: () => sandbox.destroy(), + sandboxId: 'ses-test', + sessionId: 'agent_test', + phase: 'asyncPreparation', + }, + error + ); + + expect(destroyed).toBe(false); + expect(sandbox.destroy).not.toHaveBeenCalled(); + }); + + it('does not destroy shared sandbox when capacity measurement fails without filesystem evidence', async () => { + const sandbox = { destroy: vi.fn().mockResolvedValue(undefined) }; + const error = new WorkspaceCapacityInspectionUnavailableError( + 'Workspace admission rejected because disk capacity could not be measured', + new Error('df: command not found') + ); + + const destroyed = await destroySandboxAfterPreparationInfrastructureFailure( + { + deleteSandbox: () => sandbox.destroy(), + sandboxId: 'ses-test', + sessionId: 'agent_test', + phase: 'asyncPreparation', + }, + error + ); + + expect(destroyed).toBe(false); + expect(sandbox.destroy).not.toHaveBeenCalled(); + }); + it('does not destroy sandbox for unrelated errors', async () => { const sandbox = { destroy: vi.fn().mockResolvedValue(undefined) }; const destroyed = await destroySandboxAfterInternalServerError( { - sandbox, + deleteSandbox: () => sandbox.destroy(), sandboxId: 'ses-test', sessionId: 'agent_test', phase: 'asyncPreparation', @@ -214,7 +290,7 @@ describe('sandbox recovery', () => { const sandbox = { destroy: vi.fn().mockResolvedValue(undefined) }; const destroyed = await destroySandboxAfterPreparationInfrastructureFailure( { - sandbox, + deleteSandbox: () => sandbox.destroy(), sandboxId: 'ses-test', sessionId: 'agent_test', phase: 'asyncPreparation', diff --git a/services/cloud-agent-next/src/sandbox-recovery.ts b/services/cloud-agent-next/src/sandbox-recovery.ts index 78f2cc6e5b..92a0acdd7e 100644 --- a/services/cloud-agent-next/src/sandbox-recovery.ts +++ b/services/cloud-agent-next/src/sandbox-recovery.ts @@ -1,12 +1,11 @@ import { logger } from './logger.js'; -import { WorkspaceFilesystemPreparationError } from './workspace-errors.js'; - -type DestroyableSandbox = { - destroy(): Promise; -}; +import { + SandboxCapacityInspectionError, + WorkspaceFilesystemPreparationError, +} from './workspace-errors.js'; type RecoveryContext = { - sandbox: DestroyableSandbox; + deleteSandbox(reason: 'recovery'): Promise; sandboxId: string; sessionId?: string; phase: string; @@ -27,6 +26,11 @@ type PreparationInfrastructureFailure = type: 'workspace_filesystem_preparation_error'; error: WorkspaceFilesystemPreparationError; message: string; + } + | { + type: 'sandbox_capacity_inspection_error'; + error: SandboxCapacityInspectionError; + message: string; }; function isRecord(value: unknown): value is Record { @@ -160,6 +164,23 @@ function getWorkspaceFilesystemPreparationError( return getWorkspaceFilesystemPreparationErrorWithSeen(error, new WeakSet()); } +function getSandboxCapacityInspectionErrorWithSeen( + error: unknown, + seen: WeakSet +): SandboxCapacityInspectionError | undefined { + if (!isRecord(error)) return undefined; + if (seen.has(error)) return undefined; + seen.add(error); + if (error instanceof SandboxCapacityInspectionError) return error; + return getSandboxCapacityInspectionErrorWithSeen(getNestedProperty(error, 'cause'), seen); +} + +function getSandboxCapacityInspectionError( + error: unknown +): SandboxCapacityInspectionError | undefined { + return getSandboxCapacityInspectionErrorWithSeen(error, new WeakSet()); +} + export function getPreparationInfrastructureFailure( error: unknown ): PreparationInfrastructureFailure | undefined { @@ -192,6 +213,15 @@ export function getPreparationInfrastructureFailure( }; } + const capacityInspectionError = getSandboxCapacityInspectionError(error); + if (capacityInspectionError) { + return { + type: 'sandbox_capacity_inspection_error', + error: capacityInspectionError, + message: capacityInspectionError.message, + }; + } + const workspaceError = getWorkspaceFilesystemPreparationError(error); if (workspaceError) { return { @@ -224,7 +254,7 @@ export async function destroySandboxAfterInternalServerError( .error('Sandbox returned 500 during workspace preparation; destroying sandbox'); try { - await context.sandbox.destroy(); + await context.deleteSandbox('recovery'); logger .withFields({ sandboxId: context.sandboxId, @@ -275,7 +305,7 @@ export async function destroySandboxAfterPreparationInfrastructureFailure( .error('Sandbox workspace Git probe timed out; destroying sandbox'); try { - await context.sandbox.destroy(); + await context.deleteSandbox('recovery'); logger .withFields({ sandboxId: context.sandboxId, @@ -300,6 +330,44 @@ export async function destroySandboxAfterPreparationInfrastructureFailure( } } + if (failure.type === 'sandbox_capacity_inspection_error') { + const errorMessage = getErrorMessage(failure.error); + logger + .withFields({ + sandboxId: context.sandboxId, + sessionId: context.sessionId, + phase: context.phase, + error: errorMessage, + reason: 'sandbox_filesystem_unusable', + logTag: 'sandbox_capacity_inspection_failed', + }) + .error('Sandbox capacity inspection failed; destroying unusable sandbox'); + try { + await context.deleteSandbox('recovery'); + logger + .withFields({ + sandboxId: context.sandboxId, + sessionId: context.sessionId, + phase: context.phase, + logTag: 'sandbox_capacity_inspection_destroyed', + }) + .info('Destroyed sandbox after capacity inspection failure'); + return true; + } catch (destroyError) { + logger + .withFields({ + sandboxId: context.sandboxId, + sessionId: context.sessionId, + phase: context.phase, + originalError: errorMessage, + destroyError: getErrorMessage(destroyError), + logTag: 'sandbox_capacity_inspection_destroy_failed', + }) + .error('Failed to destroy sandbox after capacity inspection failure'); + return false; + } + } + const errorMessage = getErrorMessage(failure.error); logger .withFields({ @@ -313,7 +381,7 @@ export async function destroySandboxAfterPreparationInfrastructureFailure( .error('Workspace filesystem preparation failed; destroying sandbox'); try { - await context.sandbox.destroy(); + await context.deleteSandbox('recovery'); logger .withFields({ sandboxId: context.sandboxId, diff --git a/services/cloud-agent-next/src/server.test.ts b/services/cloud-agent-next/src/server.test.ts index 28045e2dd8..16f8a96da4 100644 --- a/services/cloud-agent-next/src/server.test.ts +++ b/services/cloud-agent-next/src/server.test.ts @@ -1,16 +1,13 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'; import jwt from 'jsonwebtoken'; import type { Env } from './types.js'; -import { WRAPPER_VERSION } from './shared/wrapper-version.js'; const { - getSandboxMock, - findWrapperForSessionMock, + getRunningTerminalClientMock, consumeCloudAgentReportBatchMock, removeExpiredCloudAgentReportDataMock, } = vi.hoisted(() => ({ - getSandboxMock: vi.fn(), - findWrapperForSessionMock: vi.fn(), + getRunningTerminalClientMock: vi.fn(), consumeCloudAgentReportBatchMock: vi.fn().mockResolvedValue(undefined), removeExpiredCloudAgentReportDataMock: vi.fn().mockResolvedValue(undefined), })); @@ -33,11 +30,13 @@ vi.mock('./logger.js', () => { vi.mock('@cloudflare/sandbox', () => ({ Sandbox: class Sandbox {}, - getSandbox: getSandboxMock, + getSandbox: vi.fn(), })); -vi.mock('./kilo/wrapper-manager.js', () => ({ - findWrapperForSession: findWrapperForSessionMock, +vi.mock('./agent-sandbox/factory.js', () => ({ + createAgentSandbox: vi.fn(() => ({ + getRunningTerminalClient: getRunningTerminalClientMock, + })), })); vi.mock('./router.js', () => ({ @@ -102,8 +101,7 @@ function fetchWorker(request: Request, env: MockEnv): Promise | Respon } beforeEach(() => { - getSandboxMock.mockReset(); - findWrapperForSessionMock.mockReset(); + getRunningTerminalClientMock.mockReset(); consumeCloudAgentReportBatchMock.mockClear(); removeExpiredCloudAgentReportDataMock.mockClear(); }); @@ -216,18 +214,11 @@ describe('server /terminal', () => { }, }; const terminalResponse = new Response('proxied', { status: 200 }); - const wsConnect = vi.fn().mockResolvedValueOnce(terminalResponse); - const containerFetch = vi.fn().mockResolvedValueOnce( - Response.json({ - healthy: true, - state: 'idle', - version: WRAPPER_VERSION, - sessionId: 'session-1', - }) - ); - const sandbox = { containerFetch, wsConnect }; - getSandboxMock.mockReturnValue(sandbox); - findWrapperForSessionMock.mockResolvedValue({ port: 59954 }); + const connectTerminal = vi.fn().mockResolvedValueOnce(terminalResponse); + getRunningTerminalClientMock.mockResolvedValue({ + status: 'ready', + client: { connectTerminal }, + }); const getMetadata = vi.fn().mockResolvedValue(metadata); const fetch = vi.fn(); env.CLOUD_AGENT_SESSION.idFromName.mockReturnValue('do-id'); @@ -246,19 +237,8 @@ describe('server /terminal', () => { expect(env.CLOUD_AGENT_SESSION.idFromName).toHaveBeenCalledWith('user-1:session-1'); expect(getMetadata).toHaveBeenCalledTimes(1); expect(fetch).not.toHaveBeenCalled(); - expect(getSandboxMock).toHaveBeenCalledWith(env.Sandbox, sandboxId, expect.any(Object)); - expect(findWrapperForSessionMock).toHaveBeenCalledWith(sandbox, 'session-1'); - expect(containerFetch).toHaveBeenCalledTimes(1); - expect(containerFetch.mock.calls[0]).toEqual([ - 'http://container/health', - { method: 'GET', headers: { 'Content-Type': 'application/json' } }, - 59954, - ]); - expect(wsConnect).toHaveBeenCalledTimes(1); - const connectRequest = wsConnect.mock.calls[0]?.[0]; - expect(connectRequest).toBeInstanceOf(Request); - expect(new URL((connectRequest as Request).url).pathname).toBe('/pty/pty_123/connect'); - expect(wsConnect.mock.calls[0]?.[1]).toBe(59954); + expect(getRunningTerminalClientMock).toHaveBeenCalledOnce(); + expect(connectTerminal).toHaveBeenCalledWith('pty_123', request); }); it('rejects stream-purpose tickets', async () => { diff --git a/services/cloud-agent-next/src/session-service.test.ts b/services/cloud-agent-next/src/session-service.test.ts index e57b78c978..128c8309f8 100644 --- a/services/cloud-agent-next/src/session-service.test.ts +++ b/services/cloud-agent-next/src/session-service.test.ts @@ -72,6 +72,10 @@ import type { CloudAgentSessionState, PersistenceEnv } from './persistence/types import { parseSessionMetadata } from './persistence/session-metadata.js'; import type { ExecutionSession, SandboxInstance, SessionId } from './types.js'; import type { FencedWrapperDispatchRequest } from './execution/types.js'; +import { + SandboxCapacityInspectionError, + WorkspaceCapacityAdmissionRejectedError, +} from './workspace-errors.js'; type MockExecutionSession = ExecutionSession & { exec: ReturnType; @@ -132,13 +136,19 @@ function createSession(repoExists = false): MockExecutionSession { return { exec, gitCheckout } as unknown as MockExecutionSession; } +type TestSandbox = SandboxInstance & { + createSessionMock: ReturnType; +}; + function createSandbox( session: ExecutionSession, repoExists = false, writeFile = vi.fn().mockResolvedValue(undefined) -): SandboxInstance { +): TestSandbox { + const createSessionMock = vi.fn().mockResolvedValue(session); return { - createSession: vi.fn().mockResolvedValue(session), + createSession: createSessionMock, + createSessionMock, writeFile, mkdir: vi.fn().mockResolvedValue(undefined), exec: vi.fn(async (command: string) => { @@ -147,7 +157,7 @@ function createSandbox( } return { exitCode: 0, stdout: '', stderr: '' }; }), - } as unknown as SandboxInstance; + } as unknown as TestSandbox; } function createEnv(metadata?: CloudAgentSessionState | null): PersistenceEnv { @@ -303,6 +313,73 @@ describe('SessionService.prepareWorkspace', () => { }); }); + it('types ENOSPC during the cold devcontainer probe before provisioning', async () => { + const session = createSession(false); + const sandbox = createSandbox(session); + (sandbox.exec as ReturnType).mockResolvedValueOnce({ + exitCode: 1, + stdout: '', + stderr: 'ENOSPC: no space left on device', + }); + const metadata = { + ...createMetadata(), + workspace: { + sandboxId: 'dind-abcdef' as const, + devcontainerRequested: true, + }, + } satisfies CloudAgentSessionState; + + await expect( + new SessionService().prepareWorkspace({ + sandbox, + sandboxId: 'dind-abcdef', + userId: 'user_test', + sessionId: 'agent_test' as SessionId, + env: createEnv(), + metadata, + kilocodeModel: 'test-model', + }) + ).rejects.toBeInstanceOf(SandboxCapacityInspectionError); + + expect(workspaceMocks.setupWorkspace).not.toHaveBeenCalled(); + expect(sandbox.createSessionMock).not.toHaveBeenCalled(); + }); + + it('rejects cold devcontainer preparation before workspace or runtime provisioning when admission fails', async () => { + const session = createSession(false); + const sandbox = createSandbox(session); + const metadata = { + ...createMetadata(), + workspace: { + sandboxId: 'dind-abcdef' as const, + devcontainerRequested: true, + }, + } satisfies CloudAgentSessionState; + const rejection = new WorkspaceCapacityAdmissionRejectedError({ + availableMB: 512, + thresholdMB: 2048, + cleaned: 0, + skipped: 1, + }); + workspaceMocks.checkDiskAndCleanBeforeSetup.mockRejectedValueOnce(rejection); + + await expect( + new SessionService().prepareWorkspace({ + sandbox, + sandboxId: 'dind-abcdef', + userId: 'user_test', + sessionId: 'agent_test' as SessionId, + env: createEnv(), + metadata, + kilocodeModel: 'test-model', + }) + ).rejects.toBe(rejection); + + expect(workspaceMocks.setupWorkspace).not.toHaveBeenCalled(); + expect(sandbox.createSessionMock).not.toHaveBeenCalled(); + expect(devcontainerMocks.bringUpDevContainer).not.toHaveBeenCalled(); + }); + it('hydrates requested devcontainer metadata while preparing a cold DIND workspace', async () => { const session = createSession(false); const writeFile = vi.fn().mockResolvedValue(undefined); diff --git a/services/cloud-agent-next/src/session-service.ts b/services/cloud-agent-next/src/session-service.ts index 1bbe65248d..18408fdd74 100644 --- a/services/cloud-agent-next/src/session-service.ts +++ b/services/cloud-agent-next/src/session-service.ts @@ -69,6 +69,10 @@ import type { FencedWrapperDispatchRequest, } from './execution/types.js'; import { normalizeAgentMode } from './schema.js'; +import { + isSandboxFilesystemUnusableError, + SandboxCapacityInspectionError, +} from './workspace-errors.js'; const SETUP_COMMAND_TIMEOUT_SECONDS = 300; // 5 minutes const DEFAULT_DENIED_COMMAND_PATTERNS = ['rm -rf', 'sudo rm', 'mkfs', 'dd if=']; @@ -1838,12 +1842,28 @@ export class SessionService { executor: SandboxInstance | ExecutionSession, workspacePath: string ): Promise { - const result = await timedExec( - executor, - `test -d '${workspacePath}/.git' && echo exists`, - 'session.prepareWorkspace.repoExists' - ); - return result.stdout?.includes('exists') ?? false; + try { + const result = await timedExec( + executor, + `test -d '${workspacePath}/.git' && echo exists`, + 'session.prepareWorkspace.repoExists' + ); + if (result.exitCode !== 0 && isSandboxFilesystemUnusableError(result.stderr)) { + throw new SandboxCapacityInspectionError( + 'Workspace admission probe cannot run because the sandbox filesystem is unusable', + new Error(result.stderr) + ); + } + return result.stdout?.includes('exists') ?? false; + } catch (error) { + if (isSandboxFilesystemUnusableError(error)) { + throw new SandboxCapacityInspectionError( + 'Workspace admission probe cannot run because the sandbox filesystem is unusable', + error + ); + } + throw error; + } } private async cloneRepository( diff --git a/services/cloud-agent-next/src/session/agent-runtime.test.ts b/services/cloud-agent-next/src/session/agent-runtime.test.ts index 80a95ee31e..729d5292f0 100644 --- a/services/cloud-agent-next/src/session/agent-runtime.test.ts +++ b/services/cloud-agent-next/src/session/agent-runtime.test.ts @@ -1,12 +1,13 @@ import { describe, expect, it, vi } from 'vitest'; -import type { Env, SandboxInstance } from '../types.js'; +import type { AgentSandbox } from '../agent-sandbox/protocol.js'; +import type { Env } from '../types.js'; import type { FencedWrapperDispatchRequest, MessageDeliveryRequest, WorkspaceReady, } from '../execution/types.js'; import { createAgentRuntime } from './agent-runtime.js'; -import { getWrapperRuntimeState } from './wrapper-runtime-state.js'; +import { getWrapperLease, getWrapperRuntimeState } from './wrapper-runtime-state.js'; import type { SessionMetadata } from '../persistence/session-metadata.js'; vi.mock('@cloudflare/sandbox', () => ({ @@ -16,7 +17,8 @@ vi.mock('@cloudflare/sandbox', () => ({ type MemoryRuntimeStorage = Pick; function createMemoryStorage( - initialEntries?: Array<[string, unknown]> + initialEntries?: Array<[string, unknown]>, + onPut?: (key: string, value: unknown) => void ): MemoryRuntimeStorage & DurableObjectStorage { const store = new Map(initialEntries ?? []); return { @@ -24,6 +26,7 @@ function createMemoryStorage( return store.get(key) as T | undefined; }, async put(key: string, value: unknown) { + onPut?.(key, value); store.set(key, value); }, async delete(keys: string | string[]) { @@ -95,10 +98,15 @@ function createWorkspaceReady(): WorkspaceReady { } describe('AgentRuntime', () => { - it('fences grouped delivery, records wrapper readiness, and returns the existing queue result shape', async () => { + it('preflights cold delivery, authorizes its physical lease, and returns the existing queue result shape', async () => { const storage = createMemoryStorage(); const ready = createWorkspaceReady(); const deliveredPlans: FencedWrapperDispatchRequest[] = []; + const discoverSessionWrappers = vi.fn().mockResolvedValue({ status: 'absent' }); + const sandbox = { + discoverSessionWrappers, + } as unknown as AgentSandbox; + const alarmDeadlines: number[] = []; const orchestrator = { execute: vi.fn( async ( @@ -106,9 +114,11 @@ describe('AgentRuntime', () => { options?: { onProgress?: (step: string, message: string) => void; onWorkspaceReady?: (workspace: WorkspaceReady) => Promise; + leasedInstance?: { instanceId: string; instanceGeneration: number }; } ) => { deliveredPlans.push(plan); + expect(options?.leasedInstance).toMatchObject({ instanceGeneration: 1 }); options?.onProgress?.('kilo_server', 'Starting Kilo...'); await options?.onWorkspaceReady?.(ready); return { kiloSessionId: 'kilo_runtime' }; @@ -125,6 +135,10 @@ describe('AgentRuntime', () => { getOrchestratorOverride: () => orchestrator, getSessionIdForLogs: () => 'agent_runtime', sendToWrapper: () => false, + createAgentSandbox: () => sandbox, + requestAlarmAtOrBefore: async deadline => { + alarmDeadlines.push(deadline); + }, }); const result = await runtime.send(createPlan(), { @@ -139,8 +153,16 @@ describe('AgentRuntime', () => { }, }); const wrapperState = await getWrapperRuntimeState(storage); + const physicalLease = await getWrapperLease(storage); const [deliveredPlan] = deliveredPlans; + expect(discoverSessionWrappers).toHaveBeenCalledOnce(); + expect(alarmDeadlines).toHaveLength(2); + expect(physicalLease).toMatchObject({ + state: 'owns_wrapper', + nextInstanceGeneration: 2, + instance: { instanceGeneration: 1 }, + }); expect(result).toMatchObject({ success: true, outcome: 'accepted', @@ -166,6 +188,57 @@ describe('AgentRuntime', () => { expect(wrapperState.nextPingAt).toEqual(expect.any(Number)); }); + it('keeps an accepted new delivery supervised when physical lease acceptance persistence fails', async () => { + let rejectedAcceptedLeaseWrite = false; + const storage = createMemoryStorage(undefined, (key, value) => { + const lease = value as { state?: string; startupDeadlineAt?: number }; + if ( + key === 'wrapper_lease' && + lease.state === 'owns_wrapper' && + lease.startupDeadlineAt === undefined && + !rejectedAcceptedLeaseWrite + ) { + rejectedAcceptedLeaseWrite = true; + throw new Error('delivery_accepted lease write failed'); + } + }); + const sandbox = { + discoverSessionWrappers: vi.fn().mockResolvedValue({ status: 'absent' }), + } as unknown as AgentSandbox; + const runtime = createAgentRuntime({ + storage, + env: {} as Env, + getMetadata: async () => createMetadata(), + getOrchestratorOverride: () => ({ + execute: async (_plan, options) => { + await options?.onWorkspaceReady?.(createWorkspaceReady()); + return { kiloSessionId: 'kilo_runtime' }; + }, + }), + getSessionIdForLogs: () => 'agent_runtime', + sendToWrapper: () => false, + createAgentSandbox: () => sandbox, + }); + const onAccepted = vi.fn().mockResolvedValue(undefined); + + await expect(runtime.send(createPlan(), { onAccepted })).resolves.toMatchObject({ + success: true, + outcome: 'accepted', + }); + + expect(onAccepted).toHaveBeenCalledOnce(); + expect(rejectedAcceptedLeaseWrite).toBe(true); + await expect(getWrapperRuntimeState(storage)).resolves.toMatchObject({ + wrapperConnectionId: expect.any(String), + wrapperRunId: expect.any(String), + noOutputDeadlineAt: expect.any(Number), + }); + await expect(getWrapperLease(storage)).resolves.toMatchObject({ + state: 'owns_wrapper', + startupDeadlineAt: expect.any(Number), + }); + }); + it('preserves accepted-message liveness when a hot follow-up fails before acceptance', async () => { const storage = createMemoryStorage([ [ @@ -179,7 +252,30 @@ describe('AgentRuntime', () => { nextPingAt: 7_000, }, ], + [ + 'wrapper_lease', + { + state: 'owns_wrapper', + nextInstanceGeneration: 2, + instance: { instanceId: 'instance_hot', instanceGeneration: 1 }, + }, + ], ]); + const discoverSessionWrappers = vi.fn().mockResolvedValue({ + status: 'present', + observed: [ + { + representation: 'process', + id: 'wrapper-hot', + port: 5_000, + instanceId: 'instance_hot', + instanceGeneration: 1, + }, + ], + }); + const sandbox = { + discoverSessionWrappers, + } as unknown as AgentSandbox; const runtime = createAgentRuntime({ storage, env: {} as Env, @@ -191,10 +287,12 @@ describe('AgentRuntime', () => { }), getSessionIdForLogs: () => 'agent_runtime', sendToWrapper: () => false, + createAgentSandbox: () => sandbox, }); await expect(runtime.send(createPlan())).rejects.toThrow('hot follow-up failed'); + expect(discoverSessionWrappers).toHaveBeenCalledOnce(); await expect(getWrapperRuntimeState(storage)).resolves.toMatchObject({ wrapperGeneration: 3, wrapperConnectionId: 'conn_hot', @@ -205,8 +303,117 @@ describe('AgentRuntime', () => { }); }); - it('clears a newly allocated wrapper fence when delivery fails before readiness', async () => { + it('invalidates a warm idle physical wrapper when its reuse delivery fails', async () => { + const storage = createMemoryStorage([ + [ + 'wrapper_runtime_state', + { wrapperGeneration: 4, wrapperConnectionId: 'conn_previous', wrapperRunId: 'wr_previous' }, + ], + [ + 'wrapper_lease', + { + state: 'owns_wrapper', + nextInstanceGeneration: 2, + instance: { instanceId: 'instance_warm', instanceGeneration: 1 }, + keepWarmUntil: Date.now() + 60_000, + }, + ], + ]); + const sandbox = { + discoverSessionWrappers: vi.fn().mockResolvedValue({ + status: 'present', + observed: [ + { + representation: 'process', + id: 'wrapper-warm', + port: 5_000, + instanceId: 'instance_warm', + instanceGeneration: 1, + }, + ], + }), + } as unknown as AgentSandbox; + const runtime = createAgentRuntime({ + storage, + env: {} as Env, + getMetadata: async () => createMetadata(), + getOrchestratorOverride: () => ({ + execute: async () => { + throw new Error('warm readiness failed'); + }, + }), + getSessionIdForLogs: () => 'agent_runtime', + sendToWrapper: () => false, + createAgentSandbox: () => sandbox, + }); + + await expect(runtime.send(createPlan())).rejects.toThrow('warm readiness failed'); + + await expect(getWrapperLease(storage)).resolves.toMatchObject({ + state: 'stop_needed', + target: { kind: 'instance', instance: { instanceId: 'instance_warm' } }, + reason: 'startup-failed', + }); + await expect(getWrapperRuntimeState(storage)).resolves.not.toMatchObject({ + wrapperConnectionId: 'conn_previous', + wrapperRunId: 'wr_previous', + }); + }); + + it('releases a verified-absent owned wrapper before allocating the next generation', async () => { + const storage = createMemoryStorage([ + [ + 'wrapper_runtime_state', + { + wrapperGeneration: 7, + wrapperConnectionId: 'conn_stale', + wrapperRunId: 'wr_stale', + }, + ], + [ + 'wrapper_lease', + { + state: 'owns_wrapper', + nextInstanceGeneration: 2, + instance: { instanceId: 'instance_old', instanceGeneration: 1 }, + }, + ], + ]); + const sandbox = { + discoverSessionWrappers: vi.fn().mockResolvedValue({ status: 'absent' }), + } as unknown as AgentSandbox; + const execute = vi.fn(async (_plan, options) => { + expect(options?.leasedInstance).toMatchObject({ instanceGeneration: 2 }); + return { kiloSessionId: 'kilo_runtime' }; + }); + const runtime = createAgentRuntime({ + storage, + env: {} as Env, + getMetadata: async () => createMetadata(), + getOrchestratorOverride: () => ({ execute }), + getSessionIdForLogs: () => 'agent_runtime', + sendToWrapper: () => false, + createAgentSandbox: () => sandbox, + }); + + await runtime.send(createPlan()); + + await expect(getWrapperLease(storage)).resolves.toMatchObject({ + state: 'owns_wrapper', + nextInstanceGeneration: 3, + instance: { instanceGeneration: 2 }, + }); + await expect(getWrapperRuntimeState(storage)).resolves.not.toMatchObject({ + wrapperRunId: 'wr_stale', + wrapperConnectionId: 'conn_stale', + }); + }); + + it('stores cleanup obligation when new delivery fails before readiness', async () => { const storage = createMemoryStorage(); + const sandbox = { + discoverSessionWrappers: vi.fn().mockResolvedValue({ status: 'absent' }), + } as unknown as AgentSandbox; const runtime = createAgentRuntime({ storage, env: {} as Env, @@ -218,13 +425,113 @@ describe('AgentRuntime', () => { }), getSessionIdForLogs: () => 'agent_runtime', sendToWrapper: () => false, + createAgentSandbox: () => sandbox, }); await expect(runtime.send(createPlan())).rejects.toThrow('wrapper unavailable'); + await expect(getWrapperLease(storage)).resolves.toMatchObject({ + state: 'stop_needed', + target: { kind: 'instance', instance: { instanceGeneration: 1 } }, + reason: 'startup-failed', + }); await expect(getWrapperRuntimeState(storage)).resolves.toEqual({ wrapperGeneration: 2 }); }); + it('stores cleanup obligation when a newly leased wrapper readies but its initial dispatch fails', async () => { + const storage = createMemoryStorage(); + const sandbox = { + discoverSessionWrappers: vi.fn().mockResolvedValue({ status: 'absent' }), + } as unknown as AgentSandbox; + const runtime = createAgentRuntime({ + storage, + env: {} as Env, + getMetadata: async () => createMetadata(), + getOrchestratorOverride: () => ({ + execute: async (_plan, options) => { + await options?.onWorkspaceReady?.(createWorkspaceReady()); + throw new Error('initial prompt failed'); + }, + }), + getSessionIdForLogs: () => 'agent_runtime', + sendToWrapper: () => false, + createAgentSandbox: () => sandbox, + }); + + await expect(runtime.send(createPlan())).rejects.toThrow('initial prompt failed'); + await expect(getWrapperLease(storage)).resolves.toMatchObject({ + state: 'stop_needed', + target: { kind: 'instance', instance: { instanceGeneration: 1 } }, + reason: 'startup-failed', + }); + }); + + it.each([ + { + observation: { status: 'inspection-failed', error: 'provider unavailable' }, + reason: 'observation-failed', + }, + { observation: { status: 'present', observed: [] }, reason: 'unexpected-wrapper' }, + ])( + 'blocks cold launch after unsafe physical preflight ($reason)', + async ({ observation, reason }) => { + const storage = createMemoryStorage(); + const execute = vi.fn(); + const sandbox = { + discoverSessionWrappers: vi.fn().mockResolvedValue(observation), + } as unknown as AgentSandbox; + const runtime = createAgentRuntime({ + storage, + env: {} as Env, + getMetadata: async () => createMetadata(), + getOrchestratorOverride: () => ({ execute }), + getSessionIdForLogs: () => 'agent_runtime', + sendToWrapper: () => false, + createAgentSandbox: () => sandbox, + }); + + await expect(runtime.send(createPlan())).rejects.toThrow(/cleanup is required/i); + expect(execute).not.toHaveBeenCalled(); + await expect(getWrapperLease(storage)).resolves.toMatchObject({ + state: 'stop_needed', + target: { kind: 'session' }, + reason, + }); + } + ); + + it('blocks migrated run-fence state when no durable physical owner authorizes the visible wrapper', async () => { + const storage = createMemoryStorage([ + [ + 'wrapper_runtime_state', + { wrapperGeneration: 3, wrapperConnectionId: 'conn_legacy', wrapperRunId: 'wr_legacy' }, + ], + ]); + const execute = vi.fn(); + const sandbox = { + discoverSessionWrappers: vi.fn().mockResolvedValue({ + status: 'present', + observed: [{ representation: 'process', id: 'legacy', port: 5_000 }], + }), + } as unknown as AgentSandbox; + const runtime = createAgentRuntime({ + storage, + env: {} as Env, + getMetadata: async () => createMetadata(), + getOrchestratorOverride: () => ({ execute }), + getSessionIdForLogs: () => 'agent_runtime', + sendToWrapper: () => false, + createAgentSandbox: () => sandbox, + }); + + await expect(runtime.send(createPlan())).rejects.toThrow(/cleanup is required/i); + expect(execute).not.toHaveBeenCalled(); + await expect(getWrapperLease(storage)).resolves.toMatchObject({ + state: 'stop_needed', + target: { kind: 'session' }, + }); + }); + it('routes snapshot and interrupt commands through the current wrapper run fence', async () => { const storage = createMemoryStorage([ [ @@ -279,24 +586,20 @@ describe('AgentRuntime', () => { await expect(runtime.interruptWrapper()).resolves.toEqual({ commandSent: false }); }); - it('keeps and stops the resolved sandbox runtime through transport controls', async () => { - const renewActivityTimeout = vi.fn(); - const sandbox = { renewActivityTimeout } as unknown as SandboxInstance; - const stopRuntimeWrapper = vi.fn().mockResolvedValue(undefined); + it('keeps the runtime sandbox alive through AgentSandbox transport controls', async () => { + const keepAlive = vi.fn().mockResolvedValue(undefined); + const sandbox = { keepAlive } as unknown as AgentSandbox; const runtime = createAgentRuntime({ storage: createMemoryStorage(), env: {} as Env, getMetadata: async () => createMetadata(), getSessionIdForLogs: () => 'agent_runtime', sendToWrapper: () => false, - resolveSandbox: () => sandbox, - stopRuntimeWrapper, + createAgentSandbox: () => sandbox, }); await runtime.keepSandboxAlive(); - await expect(runtime.stopWrapperProcess('idle-timeout')).resolves.toBe(true); - expect(renewActivityTimeout).toHaveBeenCalledOnce(); - expect(stopRuntimeWrapper).toHaveBeenCalledWith(sandbox, 'agent_runtime'); + expect(keepAlive).toHaveBeenCalledOnce(); }); }); diff --git a/services/cloud-agent-next/src/session/agent-runtime.ts b/services/cloud-agent-next/src/session/agent-runtime.ts index d6483ee359..f341db23e5 100644 --- a/services/cloud-agent-next/src/session/agent-runtime.ts +++ b/services/cloud-agent-next/src/session/agent-runtime.ts @@ -1,6 +1,10 @@ -import { getSandbox } from '@cloudflare/sandbox'; -import { SANDBOX_SLEEP_AFTER_SECONDS } from '../core/lease.js'; import { ExecutionOrchestrator } from '../execution/orchestrator.js'; +import { createAgentSandbox } from '../agent-sandbox/factory.js'; +import type { + AgentSandbox, + WrapperInstanceLease, + WrapperObservation, +} from '../agent-sandbox/protocol.js'; import type { ExecutionResult, FencedWrapperDispatchRequest, @@ -9,32 +13,25 @@ import type { WorkspaceReady, } from '../execution/types.js'; import { logger } from '../logger.js'; -import { stopWrapper } from '../kilo/wrapper-manager.js'; import type { SessionMetadata } from '../persistence/session-metadata.js'; -import { generateSandboxId, getSandboxNamespace } from '../sandbox-id.js'; import type { WrapperCommand } from '../shared/protocol.js'; -import type { Env as WorkerEnv, SandboxInstance } from '../types.js'; +import type { Env as WorkerEnv } from '../types.js'; import { allocateWrapperRuntimeState, clearAllocatedWrapperRuntimeState, + clearWrapperRuntimeIdentity, + getWrapperLease, getWrapperRuntimeState, + putWrapperLease, + READY_ONLY_IDLE_MS, recordWrapperAcceptedMessage, recordWrapperReadyLease, + reduceWrapperLease, } from './wrapper-runtime-state.js'; export const WRAPPER_NO_OUTPUT_TIMEOUT_MS = 5 * 60 * 1000; export const WRAPPER_PING_INTERVAL_MS = 60_000; - -type RuntimeSandboxOptions = { - sleepAfter?: number; -}; - -type RuntimeSandboxResolver = ( - sandboxId: string, - options?: RuntimeSandboxOptions -) => SandboxInstance; - -type RuntimeWrapperStopper = (sandbox: SandboxInstance, sessionId: string) => Promise; +export const WRAPPER_STARTUP_TIMEOUT_MS = 10 * 60 * 1000; export type AgentRuntimeOrchestrator = { execute( @@ -42,6 +39,7 @@ export type AgentRuntimeOrchestrator = { options?: { onProgress?: (step: string, message: string) => void; onWorkspaceReady?: (ready: WorkspaceReady) => Promise; + leasedInstance?: WrapperInstanceLease; } ): Promise; }; @@ -57,19 +55,12 @@ export type AgentRuntimeSendHooks = { onAccepted?: (delivery: AgentRuntimeAcceptedDelivery) => Promise; }; -export type AgentRuntimeStopReason = - | 'idle-timeout' - | 'unhealthy-wrapper' - | 'keep-warm-expired' - | 'user-interrupt'; - export type AgentRuntime = { send(plan: MessageDeliveryRequest, hooks?: AgentRuntimeSendHooks): Promise; requestSnapshot(): Promise; interruptWrapper(): Promise<{ commandSent: boolean }>; sendPing(ingestTagId: string): void; keepSandboxAlive(): Promise; - stopWrapperProcess(reason: AgentRuntimeStopReason): Promise; }; export type AgentRuntimeDependencies = { @@ -83,8 +74,9 @@ export type AgentRuntimeDependencies = { fence?: { wrapperGeneration: number; wrapperConnectionId: string } ) => boolean; getOrchestratorOverride?: () => AgentRuntimeOrchestrator | undefined; - resolveSandbox?: RuntimeSandboxResolver; - stopRuntimeWrapper?: RuntimeWrapperStopper; + createAgentSandbox?: (metadata: SessionMetadata) => AgentSandbox; + discoverSessionWrappers?: (metadata: SessionMetadata) => Promise; + requestAlarmAtOrBefore?: (deadline: number) => Promise; }; function buildRuntimeAcceptanceResult( @@ -102,10 +94,9 @@ function buildRuntimeAcceptanceResult( export function createAgentRuntime(dependencies: AgentRuntimeDependencies): AgentRuntime { const { storage, env, getMetadata, getSessionIdForLogs, sendToWrapper, getOrchestratorOverride } = dependencies; - const resolveSandbox: RuntimeSandboxResolver = - dependencies.resolveSandbox ?? - ((sandboxId, options) => getSandbox(getSandboxNamespace(env, sandboxId), sandboxId, options)); - const stopRuntimeWrapper = dependencies.stopRuntimeWrapper ?? stopWrapper; + const resolveAgentSandbox = + dependencies.createAgentSandbox ?? + ((metadata: SessionMetadata) => createAgentSandbox(env, metadata)); let orchestrator: AgentRuntimeOrchestrator | undefined; function getOrchestrator(): AgentRuntimeOrchestrator { @@ -114,8 +105,7 @@ export function createAgentRuntime(dependencies: AgentRuntimeDependencies): Agen if (!orchestrator) { orchestrator = new ExecutionOrchestrator({ - getSandbox: async sandboxId => - resolveSandbox(sandboxId, { sleepAfter: SANDBOX_SLEEP_AFTER_SECONDS }), + getAgentSandbox: plan => resolveAgentSandbox(plan.workspace.metadata), getSessionStub: (userId, sessionId) => { const doKey = `${userId}:${sessionId}`; const id = env.CLOUD_AGENT_SESSION.idFromName(doKey); @@ -133,12 +123,126 @@ export function createAgentRuntime(dependencies: AgentRuntimeDependencies): Agen return runtimeState.wrapperRunId ?? null; } + async function authorizePhysicalWrapper(plan: MessageDeliveryRequest): Promise<{ + leasedInstance: WrapperInstanceLease; + allocatedPhysicalInstance: boolean; + requiresFreshRunFence: boolean; + }> { + const current = await getWrapperLease(storage); + if (current.state === 'stop_needed' || current.state === 'stopping') { + throw new Error('Wrapper cleanup is required before delivery can launch'); + } + + const observeWrappers = () => + dependencies.discoverSessionWrappers + ? dependencies.discoverSessionWrappers(plan.workspace.metadata) + : resolveAgentSandbox(plan.workspace.metadata).discoverSessionWrappers(); + let allocatable = current; + if (current.state === 'owns_wrapper') { + const observation = await observeWrappers(); + const matchingObserved = + observation.status === 'present' && + observation.observed.length === 1 && + observation.observed[0].instanceId === current.instance.instanceId && + observation.observed[0].instanceGeneration === current.instance.instanceGeneration; + if (matchingObserved) { + if (current.keepWarmUntil !== undefined) { + const startupDeadlineAt = Date.now() + WRAPPER_STARTUP_TIMEOUT_MS; + await putWrapperLease( + storage, + reduceWrapperLease(current, { + type: 'reuse', + instanceId: current.instance.instanceId, + startupDeadlineAt, + }) + ); + await dependencies.requestAlarmAtOrBefore?.(startupDeadlineAt); + } + return { + leasedInstance: current.instance, + allocatedPhysicalInstance: false, + requiresFreshRunFence: current.keepWarmUntil !== undefined, + }; + } + if (observation.status === 'absent') { + const verifiedAbsent = reduceWrapperLease(current, { + type: 'owned_absent', + instanceId: current.instance.instanceId, + }); + if (verifiedAbsent.state !== 'none') { + throw new Error('Verified wrapper absence did not release its physical lease'); + } + allocatable = verifiedAbsent; + await putWrapperLease(storage, allocatable); + } else { + const reason = + observation.status === 'inspection-failed' ? 'observation-failed' : 'unexpected-wrapper'; + const now = Date.now(); + await putWrapperLease( + storage, + reduceWrapperLease(current, { + type: 'request_stop', + target: { kind: 'session' }, + reason, + now, + }) + ); + await dependencies.requestAlarmAtOrBefore?.(now); + throw new Error('Wrapper cleanup is required before delivery can launch'); + } + } + + const observation = + allocatable === current ? await observeWrappers() : { status: 'absent' as const }; + if (observation.status !== 'absent') { + const reason = + observation.status === 'inspection-failed' ? 'observation-failed' : 'unexpected-wrapper'; + const now = Date.now(); + await putWrapperLease( + storage, + reduceWrapperLease(allocatable, { + type: 'request_stop', + target: { kind: 'session' }, + reason, + now, + }) + ); + await dependencies.requestAlarmAtOrBefore?.(now); + throw new Error('Wrapper cleanup is required before delivery can launch'); + } + + const leasedInstance = { + instanceId: `instance_${crypto.randomUUID().replace(/-/g, '')}`, + instanceGeneration: allocatable.nextInstanceGeneration, + } satisfies WrapperInstanceLease; + const startupDeadlineAt = Date.now() + WRAPPER_STARTUP_TIMEOUT_MS; + await putWrapperLease( + storage, + reduceWrapperLease(allocatable, { + type: 'allocate', + instance: leasedInstance, + startupDeadlineAt, + }) + ); + await dependencies.requestAlarmAtOrBefore?.(startupDeadlineAt); + return { leasedInstance, allocatedPhysicalInstance: true, requiresFreshRunFence: false }; + } + async function send( plan: MessageDeliveryRequest, hooks: AgentRuntimeSendHooks = {} ): Promise { const { sessionId } = plan.scope; const { turn, agent } = plan; + const { leasedInstance, allocatedPhysicalInstance, requiresFreshRunFence } = + await authorizePhysicalWrapper(plan); + const previousRuntimeState = await getWrapperRuntimeState(storage); + if ( + (allocatedPhysicalInstance || requiresFreshRunFence) && + (previousRuntimeState.wrapperConnectionId || previousRuntimeState.wrapperRunId) + ) { + await clearWrapperRuntimeIdentity(storage, {}, { incrementGeneration: true }); + } const { state: wrapperRuntimeState, allocatedNewIdentity } = await allocateWrapperRuntimeState(storage); logger @@ -169,9 +273,24 @@ export function createAgentRuntime(dependencies: AgentRuntimeDependencies): Agen let wrapperReady = false; try { await getOrchestrator().execute(fencedPlan, { + ...(leasedInstance ? { leasedInstance } : {}), onProgress: hooks.onProgress, onWorkspaceReady: async ready => { - await recordWrapperReadyLease(storage, wrapperRuntimeState); + const readyAt = Date.now(); + const readyDeadlineAt = readyAt + READY_ONLY_IDLE_MS; + await recordWrapperReadyLease(storage, wrapperRuntimeState, readyAt, readyDeadlineAt); + if (leasedInstance) { + const physicalLease = await getWrapperLease(storage); + await putWrapperLease( + storage, + reduceWrapperLease(physicalLease, { + type: 'startup_verified', + instanceId: leasedInstance.instanceId, + readyDeadlineAt, + }) + ); + await dependencies.requestAlarmAtOrBefore?.(readyDeadlineAt); + } wrapperReady = true; logger .withFields({ @@ -197,6 +316,25 @@ export function createAgentRuntime(dependencies: AgentRuntimeDependencies): Agen acceptedAt, wrapperRunId: wrapperRuntimeState.wrapperRunId, }); + try { + const acceptedLease = await getWrapperLease(storage); + await putWrapperLease( + storage, + reduceWrapperLease(acceptedLease, { + type: 'delivery_accepted', + instanceId: leasedInstance.instanceId, + }) + ); + } catch (error) { + logger + .withFields({ + sessionId, + messageId: turn.messageId, + wrapperRunId: wrapperRuntimeState.wrapperRunId, + error: error instanceof Error ? error.message : String(error), + }) + .warn('Failed to record accepted physical wrapper lease; maintenance will reconcile'); + } logger .withFields({ sessionId, @@ -219,11 +357,23 @@ export function createAgentRuntime(dependencies: AgentRuntimeDependencies): Agen errorClass: error instanceof Error ? error.name : 'UnknownError', }) .warn('AgentRuntime wrapper delivery failed'); - if (allocatedNewIdentity && !wrapperReady) { + if (allocatedNewIdentity && leasedInstance) { + const physicalLease = await getWrapperLease(storage); + const now = Date.now(); + await putWrapperLease( + storage, + reduceWrapperLease(physicalLease, { + type: 'request_stop', + target: { kind: 'instance', instance: leasedInstance }, + reason: 'startup-failed', + now, + }) + ); + await dependencies.requestAlarmAtOrBefore?.(now); await clearAllocatedWrapperRuntimeState(storage, wrapperRuntimeState); logger .withFields({ sessionId, messageId: turn.messageId }) - .debug('Cleared newly allocated runtime wrapper state after failed delivery'); + .debug('Recorded cleanup for newly allocated wrapper after failed delivery'); } else if (!allocatedNewIdentity) { const currentState = await getWrapperRuntimeState(storage); if ( @@ -280,25 +430,11 @@ export function createAgentRuntime(dependencies: AgentRuntimeDependencies): Agen sendToWrapper(ingestTagId, { type: 'ping' }); } - async function resolveRuntimeSandbox(metadata: SessionMetadata): Promise { - const sandboxId = - metadata.workspace?.sandboxId ?? - (await generateSandboxId( - env.PER_SESSION_SANDBOX_ORG_IDS, - metadata.identity.orgId, - metadata.identity.userId, - metadata.identity.sessionId, - metadata.identity.botId - )); - return resolveSandbox(sandboxId); - } - async function keepSandboxAlive(): Promise { try { const metadata = await getMetadata(); if (!metadata) return; - const sandbox = await resolveRuntimeSandbox(metadata); - await Promise.resolve(sandbox.renewActivityTimeout()); + await resolveAgentSandbox(metadata).keepAlive(); } catch (error) { logger .withFields({ @@ -309,31 +445,11 @@ export function createAgentRuntime(dependencies: AgentRuntimeDependencies): Agen } } - async function stopWrapperProcess(reason: AgentRuntimeStopReason): Promise { - try { - const metadata = await getMetadata(); - if (!metadata) return false; - const sandbox = await resolveRuntimeSandbox(metadata); - await stopRuntimeWrapper(sandbox, metadata.identity.sessionId); - return true; - } catch (error) { - logger - .withFields({ - sessionId: getSessionIdForLogs(), - reason, - error: error instanceof Error ? error.message : String(error), - }) - .warn('AgentRuntime failed to stop wrapper process'); - return false; - } - } - return { send, requestSnapshot, interruptWrapper, sendPing, keepSandboxAlive, - stopWrapperProcess, }; } diff --git a/services/cloud-agent-next/src/session/wrapper-runtime-state.test.ts b/services/cloud-agent-next/src/session/wrapper-runtime-state.test.ts new file mode 100644 index 0000000000..66dbbc3453 --- /dev/null +++ b/services/cloud-agent-next/src/session/wrapper-runtime-state.test.ts @@ -0,0 +1,209 @@ +import { describe, expect, it } from 'vitest'; +import { + emptyWrapperLease, + getWrapperLease, + nextWrapperLeaseDeadline, + putWrapperLease, + reduceWrapperLease, +} from './wrapper-runtime-state.js'; + +type MemoryStorage = Pick & DurableObjectStorage; + +function createMemoryStorage(): MemoryStorage { + const records = new Map(); + return { + async get(key: string) { + return records.get(key) as T | undefined; + }, + async put(key: string, value: unknown) { + records.set(key, value); + }, + } as MemoryStorage; +} + +const instance = { instanceId: 'instance_reducer', instanceGeneration: 1 }; + +describe('WrapperLease', () => { + it('allocates one authorized wrapper and requests targeted cleanup without losing generation', () => { + const owned = reduceWrapperLease(emptyWrapperLease(), { + type: 'allocate', + instance, + startupDeadlineAt: 2_000, + }); + expect(owned).toEqual({ + state: 'owns_wrapper', + nextInstanceGeneration: 2, + instance, + startupDeadlineAt: 2_000, + }); + + expect( + reduceWrapperLease(owned, { + type: 'request_stop', + target: { kind: 'instance', instance }, + reason: 'startup-failed', + now: 1_000, + }) + ).toEqual({ + state: 'stop_needed', + nextInstanceGeneration: 2, + target: { kind: 'instance', instance }, + reason: 'startup-failed', + requestedAt: 1_000, + nextAttemptAt: 1_000, + attempts: 0, + }); + }); + + it('retains a verified owned instance for bounded warm reuse', () => { + const owned = reduceWrapperLease(emptyWrapperLease(), { + type: 'allocate', + instance, + startupDeadlineAt: 2_000, + }); + const verified = reduceWrapperLease(owned, { + type: 'startup_verified', + instanceId: instance.instanceId, + readyDeadlineAt: 3_000, + }); + + const warm = reduceWrapperLease(verified, { + type: 'retain_warm', + instanceId: instance.instanceId, + keepWarmUntil: 20_000, + }); + + expect(warm).toEqual({ + state: 'owns_wrapper', + nextInstanceGeneration: 2, + instance, + startupDeadlineAt: undefined, + keepWarmUntil: 20_000, + }); + expect(nextWrapperLeaseDeadline(warm)).toBe(20_000); + + const reusing = reduceWrapperLease(warm, { + type: 'reuse', + instanceId: instance.instanceId, + startupDeadlineAt: 30_000, + }); + expect(reusing).toMatchObject({ + state: 'owns_wrapper', + keepWarmUntil: 20_000, + startupDeadlineAt: 30_000, + }); + expect(nextWrapperLeaseDeadline(reusing)).toBe(30_000); + expect( + reduceWrapperLease(reusing, { + type: 'startup_verified', + instanceId: instance.instanceId, + readyDeadlineAt: 31_000, + }) + ).toMatchObject({ + state: 'owns_wrapper', + keepWarmUntil: undefined, + startupDeadlineAt: 31_000, + }); + expect( + reduceWrapperLease( + reduceWrapperLease(reusing, { + type: 'startup_verified', + instanceId: instance.instanceId, + readyDeadlineAt: 31_000, + }), + { type: 'delivery_accepted', instanceId: instance.instanceId } + ) + ).toMatchObject({ state: 'owns_wrapper', startupDeadlineAt: undefined }); + }); + + it('returns an owned instance to none only after provider-verified absence', () => { + const owned = reduceWrapperLease(emptyWrapperLease(), { + type: 'allocate', + instance, + startupDeadlineAt: 2_000, + }); + + expect( + reduceWrapperLease(owned, { type: 'owned_absent', instanceId: 'stale_instance' }) + ).toEqual(owned); + expect( + reduceWrapperLease(owned, { type: 'owned_absent', instanceId: instance.instanceId }) + ).toEqual({ + state: 'none', + nextInstanceGeneration: 2, + }); + }); + + it('settles cleanup only for a matching confirmed-absent attempt and ignores stale results', () => { + const requested = reduceWrapperLease(emptyWrapperLease(), { + type: 'request_stop', + target: { kind: 'session' }, + reason: 'unexpected-wrapper', + now: 1_000, + }); + const stopping = reduceWrapperLease(requested, { + type: 'begin_stop_attempt', + attemptId: 'attempt_current', + now: 1_000, + attemptDeadlineAt: 46_000, + }); + + expect( + reduceWrapperLease(stopping, { type: 'stop_absent', attemptId: 'attempt_stale' }) + ).toEqual(stopping); + expect( + reduceWrapperLease(stopping, { + type: 'stop_attempt_expired', + attemptId: 'attempt_stale', + retryAt: 50_000, + }) + ).toEqual(stopping); + expect( + reduceWrapperLease(stopping, { type: 'stop_absent', attemptId: 'attempt_current' }) + ).toEqual({ state: 'none', nextInstanceGeneration: 1 }); + }); + + it('preserves the stop target and counter across a bounded failed attempt', () => { + const requested = reduceWrapperLease(emptyWrapperLease(), { + type: 'request_stop', + target: { kind: 'session' }, + reason: 'observation-failed', + now: 100, + }); + const stopping = reduceWrapperLease(requested, { + type: 'begin_stop_attempt', + attemptId: 'attempt_failed', + now: 100, + attemptDeadlineAt: 200, + }); + const retrying = reduceWrapperLease(stopping, { + type: 'stop_not_confirmed', + attemptId: 'attempt_failed', + retryAt: 5_200, + error: 'inspection failed', + }); + + expect(retrying).toMatchObject({ + state: 'stop_needed', + target: { kind: 'session' }, + reason: 'observation-failed', + attempts: 1, + nextAttemptAt: 5_200, + lastError: 'inspection failed', + }); + expect(nextWrapperLeaseDeadline(retrying)).toBe(5_200); + }); + + it('validates the separately persisted physical ownership record', async () => { + const storage = createMemoryStorage(); + await expect(getWrapperLease(storage)).resolves.toEqual(emptyWrapperLease()); + + const owned = reduceWrapperLease(emptyWrapperLease(), { + type: 'allocate', + instance, + startupDeadlineAt: 2_000, + }); + await putWrapperLease(storage, owned); + await expect(getWrapperLease(storage)).resolves.toEqual(owned); + }); +}); diff --git a/services/cloud-agent-next/src/session/wrapper-runtime-state.ts b/services/cloud-agent-next/src/session/wrapper-runtime-state.ts index bc2fb7b439..7e6e3ac8a7 100644 --- a/services/cloud-agent-next/src/session/wrapper-runtime-state.ts +++ b/services/cloud-agent-next/src/session/wrapper-runtime-state.ts @@ -1,6 +1,194 @@ import { z } from 'zod'; +import type { + WrapperInstanceLease, + WrapperStopReason, + WrapperStopTarget, +} from '../agent-sandbox/protocol.js'; const WRAPPER_RUNTIME_STATE_KEY = 'wrapper_runtime_state'; +const WRAPPER_LEASE_KEY = 'wrapper_lease'; + +const wrapperInstanceLeaseSchema = z.object({ + instanceId: z.string().min(1), + instanceGeneration: z.number().int().nonnegative(), +}); + +const wrapperStopTargetSchema = z.discriminatedUnion('kind', [ + z.object({ kind: z.literal('instance'), instance: wrapperInstanceLeaseSchema }), + z.object({ kind: z.literal('session') }), +]); + +const wrapperStopReasonSchema = z.enum([ + 'readiness-failed', + 'startup-failed', + 'unhealthy-wrapper', + 'terminal-failed', + 'terminal-interrupted', + 'idle-timeout', + 'keep-warm-expired', + 'user-interrupt', + 'session-delete', + 'unexpected-wrapper', + 'observation-failed', +]); + +const wrapperLeaseSchema = z.discriminatedUnion('state', [ + z.object({ state: z.literal('none'), nextInstanceGeneration: z.number().int().positive() }), + z.object({ + state: z.literal('owns_wrapper'), + nextInstanceGeneration: z.number().int().positive(), + instance: wrapperInstanceLeaseSchema, + startupDeadlineAt: z.number().int().nonnegative().optional(), + keepWarmUntil: z.number().int().nonnegative().optional(), + }), + z.object({ + state: z.literal('stop_needed'), + nextInstanceGeneration: z.number().int().positive(), + target: wrapperStopTargetSchema, + reason: wrapperStopReasonSchema, + requestedAt: z.number().int().nonnegative(), + nextAttemptAt: z.number().int().nonnegative(), + attempts: z.number().int().nonnegative(), + lastError: z.string().optional(), + }), + z.object({ + state: z.literal('stopping'), + nextInstanceGeneration: z.number().int().positive(), + target: wrapperStopTargetSchema, + reason: wrapperStopReasonSchema, + requestedAt: z.number().int().nonnegative(), + attemptId: z.string().min(1), + attemptStartedAt: z.number().int().nonnegative(), + attemptDeadlineAt: z.number().int().nonnegative(), + attempts: z.number().int().nonnegative(), + }), +]); + +export type WrapperLease = z.infer; + +export type WrapperLeaseEvent = + | { type: 'allocate'; instance: WrapperInstanceLease; startupDeadlineAt: number } + | { type: 'startup_verified'; instanceId: string; readyDeadlineAt: number } + | { type: 'delivery_accepted'; instanceId: string } + | { type: 'retain_warm'; instanceId: string; keepWarmUntil: number } + | { type: 'reuse'; instanceId: string; startupDeadlineAt: number } + | { type: 'owned_absent'; instanceId: string } + | { type: 'request_stop'; target: WrapperStopTarget; reason: WrapperStopReason; now: number } + | { type: 'begin_stop_attempt'; attemptId: string; now: number; attemptDeadlineAt: number } + | { type: 'stop_absent'; attemptId: string } + | { type: 'stop_not_confirmed'; attemptId: string; retryAt: number; error: string } + | { type: 'stop_attempt_expired'; attemptId: string; retryAt: number }; + +export const emptyWrapperLease = (): WrapperLease => ({ + state: 'none', + nextInstanceGeneration: 1, +}); + +export async function getWrapperLease(storage: DurableObjectStorage): Promise { + const parsed = wrapperLeaseSchema.safeParse(await storage.get(WRAPPER_LEASE_KEY)); + return parsed.success ? parsed.data : emptyWrapperLease(); +} + +export async function putWrapperLease( + storage: DurableObjectStorage, + lease: WrapperLease +): Promise { + await storage.put(WRAPPER_LEASE_KEY, wrapperLeaseSchema.parse(lease)); +} + +export function reduceWrapperLease(state: WrapperLease, event: WrapperLeaseEvent): WrapperLease { + switch (event.type) { + case 'allocate': + if (state.state !== 'none') return state; + return { + state: 'owns_wrapper', + nextInstanceGeneration: Math.max( + state.nextInstanceGeneration, + event.instance.instanceGeneration + 1 + ), + instance: event.instance, + startupDeadlineAt: event.startupDeadlineAt, + }; + case 'startup_verified': + if (state.state !== 'owns_wrapper' || state.instance.instanceId !== event.instanceId) + return state; + return { ...state, startupDeadlineAt: event.readyDeadlineAt, keepWarmUntil: undefined }; + case 'delivery_accepted': + if (state.state !== 'owns_wrapper' || state.instance.instanceId !== event.instanceId) + return state; + return { ...state, startupDeadlineAt: undefined, keepWarmUntil: undefined }; + case 'retain_warm': + if (state.state !== 'owns_wrapper' || state.instance.instanceId !== event.instanceId) + return state; + return { ...state, startupDeadlineAt: undefined, keepWarmUntil: event.keepWarmUntil }; + case 'reuse': + if (state.state !== 'owns_wrapper' || state.instance.instanceId !== event.instanceId) + return state; + return { ...state, startupDeadlineAt: event.startupDeadlineAt }; + case 'owned_absent': + if (state.state !== 'owns_wrapper' || state.instance.instanceId !== event.instanceId) + return state; + return { state: 'none', nextInstanceGeneration: state.nextInstanceGeneration }; + case 'request_stop': + if (state.state === 'stopping' || state.state === 'stop_needed') return state; + return { + state: 'stop_needed', + nextInstanceGeneration: state.nextInstanceGeneration, + target: event.target, + reason: event.reason, + requestedAt: event.now, + nextAttemptAt: event.now, + attempts: 0, + }; + case 'begin_stop_attempt': + if (state.state !== 'stop_needed' || event.now < state.nextAttemptAt) return state; + return { + state: 'stopping', + nextInstanceGeneration: state.nextInstanceGeneration, + target: state.target, + reason: state.reason, + requestedAt: state.requestedAt, + attemptId: event.attemptId, + attemptStartedAt: event.now, + attemptDeadlineAt: event.attemptDeadlineAt, + attempts: state.attempts + 1, + }; + case 'stop_absent': + if (state.state !== 'stopping' || state.attemptId !== event.attemptId) return state; + return { state: 'none', nextInstanceGeneration: state.nextInstanceGeneration }; + case 'stop_not_confirmed': + if (state.state !== 'stopping' || state.attemptId !== event.attemptId) return state; + return { + state: 'stop_needed', + nextInstanceGeneration: state.nextInstanceGeneration, + target: state.target, + reason: state.reason, + requestedAt: state.requestedAt, + nextAttemptAt: event.retryAt, + attempts: state.attempts, + lastError: event.error, + }; + case 'stop_attempt_expired': + if (state.state !== 'stopping' || state.attemptId !== event.attemptId) return state; + return { + state: 'stop_needed', + nextInstanceGeneration: state.nextInstanceGeneration, + target: state.target, + reason: state.reason, + requestedAt: state.requestedAt, + nextAttemptAt: event.retryAt, + attempts: state.attempts, + lastError: 'Stop attempt deadline expired', + }; + } +} + +export function nextWrapperLeaseDeadline(lease: WrapperLease): number | undefined { + if (lease.state === 'stop_needed') return lease.nextAttemptAt; + if (lease.state === 'stopping') return lease.attemptDeadlineAt; + if (lease.state !== 'owns_wrapper') return undefined; + return lease.startupDeadlineAt ?? lease.keepWarmUntil; +} export const IDLE_RECONCILIATION_GRACE_MS = 15_000; export const IDLE_KEEP_WARM_MS = 5 * 60 * 1000; diff --git a/services/cloud-agent-next/src/session/wrapper-supervisor.test.ts b/services/cloud-agent-next/src/session/wrapper-supervisor.test.ts index 5c3856439c..dcc5d99a84 100644 --- a/services/cloud-agent-next/src/session/wrapper-supervisor.test.ts +++ b/services/cloud-agent-next/src/session/wrapper-supervisor.test.ts @@ -15,7 +15,7 @@ import { type WrapperReconnectDecision, type WrapperSupervisorStorage, } from './wrapper-supervisor.js'; -import { getWrapperRuntimeState } from './wrapper-runtime-state.js'; +import { getWrapperLease, getWrapperRuntimeState } from './wrapper-runtime-state.js'; vi.mock('@cloudflare/sandbox', () => ({ getSandbox: vi.fn(), @@ -35,6 +35,14 @@ const WRAPPER_RUN_ID = 'wr_supervisor'; const WRAPPER_CONNECTION_ID = 'conn_supervisor'; const MESSAGE_ID = 'msg_018f1e2d3c4bSupvMsgAbCdEfG'; const NEWER_MESSAGE_ID = 'msg_018f1e2d3c4bNewerMsgAbCdEF'; +const OWNED_WRAPPER_LEASE: [string, unknown] = [ + 'wrapper_lease', + { + state: 'owns_wrapper', + nextInstanceGeneration: 2, + instance: { instanceId: 'instance_supervisor', instanceGeneration: 1 }, + }, +]; function createMemoryStorage( initialEntries?: Array<[string, unknown]>, @@ -102,6 +110,8 @@ function createHarness( const callbackJobs: CallbackJob[] = []; const sentPings: string[] = []; const stops: string[] = []; + const stopWrappers = vi.fn().mockResolvedValue({ status: 'absent' }); + const requestedAlarms: number[] = []; const currentMetadata = options?.metadata ?? createMetadata(); const settlementOutbox = createMessageSettlementOutbox({ storage, @@ -130,10 +140,6 @@ function createHarness( sendPing: ingestTagId => { sentPings.push(ingestTagId); }, - stopWrapperProcess: async reason => { - stops.push(reason); - return true; - }, }, messageSettlementOutbox: settlementOutbox, sessionMessageQueue: { requestPendingDrainIfNeeded }, @@ -141,6 +147,10 @@ function createHarness( getAssistantMessageForUserMessage: () => null, hasActiveIngestConnection: async () => false, clearInterruptRequest: async () => {}, + stopWrappers, + requestAlarmAtOrBefore: async deadline => { + requestedAlarms.push(deadline); + }, getSessionIdForLogs: () => currentMetadata.identity.sessionId, }); @@ -150,6 +160,8 @@ function createHarness( callbackJobs, sentPings, stops, + stopWrappers, + requestedAlarms, requestPendingDrainIfNeeded, settlementOutbox, supervisor, @@ -402,9 +414,13 @@ describe('WrapperSupervisor', () => { await expect(getWrapperRuntimeState(harness.storage)).resolves.toMatchObject({ wrapperGeneration: 5, }); + await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ + state: 'stop_needed', + reason: 'unhealthy-wrapper', + }); expect(harness.requestPendingDrainIfNeeded).not.toHaveBeenCalled(); expect(harness.events.map(event => event.streamEventType)).toEqual(['cloud.message.failed']); - expect(harness.stops).toEqual(['unhealthy-wrapper']); + expect(harness.stops).toEqual([]); }); it('rejects a stale wrapper run before reconnect grace can be cancelled', async () => { @@ -456,14 +472,18 @@ describe('WrapperSupervisor', () => { }); it.each([ - { status: 'failed' as const, expected: 'failed' as const }, - { status: 'interrupted' as const, expected: 'interrupted' as const }, + { status: 'failed' as const, expected: 'failed' as const, reason: 'terminal-failed' as const }, + { + status: 'interrupted' as const, + expected: 'interrupted' as const, + reason: 'terminal-interrupted' as const, + }, ])( - 'settles only matching-run messages on current $status terminal events', - async ({ status, expected }) => { + 'settles matching-run messages and durably requests physical stop on current $status terminal events', + async ({ status, expected, reason }) => { const otherRunId = 'wr_other_run'; const otherMessageId = NEWER_MESSAGE_ID; - const harness = createHarness([liveRuntimeState()]); + const harness = createHarness([liveRuntimeState(), OWNED_WRAPPER_LEASE]); await putSessionMessageState(harness.storage, acceptedMessage()); await putSessionMessageState(harness.storage, { ...acceptedMessage(otherMessageId), @@ -483,13 +503,56 @@ describe('WrapperSupervisor', () => { status: 'accepted', wrapperRunId: otherRunId, }); + await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ + state: 'stop_needed', + reason, + target: { + kind: 'instance', + instance: { instanceId: 'instance_supervisor', instanceGeneration: 1 }, + }, + }); + expect(harness.stops).toEqual([]); expect(harness.requestPendingDrainIfNeeded).toHaveBeenCalledOnce(); } ); - it('fails accepted current work on liveness expiry without redispatch when no wrapper-local Kilo query remains', async () => { + it('persists physical stop obligation before reading messages for a failed terminal event', async () => { + const storageRef: { current?: MemoryStorage } = {}; + let observedStopBeforeEffects = false; + const harness = createHarness([liveRuntimeState(), OWNED_WRAPPER_LEASE], { + storageHooks: { + beforeList: async prefix => { + if ( + observedStopBeforeEffects || + !prefix.startsWith('session_message:') || + !storageRef.current + ) { + return; + } + observedStopBeforeEffects = true; + await expect(getWrapperLease(storageRef.current)).resolves.toMatchObject({ + state: 'stop_needed', + reason: 'terminal-failed', + }); + }, + }, + }); + storageRef.current = harness.storage; + await putSessionMessageState(harness.storage, acceptedMessage()); + + await harness.supervisor.onTerminalEvent({ + wrapperRunId: WRAPPER_RUN_ID, + status: 'failed', + error: 'terminal event', + }); + + expect(observedStopBeforeEffects).toBe(true); + }); + + it('fails accepted current work and requests durable cleanup on liveness expiry', async () => { const harness = createHarness([ liveRuntimeState({ noOutputDeadlineAt: 9_000, nextPingAt: 30_000 }), + OWNED_WRAPPER_LEASE, ]); await putSessionMessageState(harness.storage, acceptedMessage()); @@ -506,11 +569,53 @@ describe('WrapperSupervisor', () => { failureCode: 'wrapper_no_output', }); expect(runtimeState.wrapperConnectionId).toBeUndefined(); + await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ + state: 'stop_needed', + reason: 'unhealthy-wrapper', + }); expect(harness.requestPendingDrainIfNeeded).not.toHaveBeenCalled(); - expect(harness.stops).toEqual(['unhealthy-wrapper']); + expect(harness.stops).toEqual([]); expect(harness.events.map(event => event.streamEventType)).toEqual(['cloud.message.failed']); }); + it('aggregates concurrent physical, liveness, disconnect, and idle deadlines', async () => { + const harness = createHarness([ + liveRuntimeState({ + nextPingAt: 20_000, + noOutputDeadlineAt: 50_000, + idleReconcileAfter: 30_000, + wrapperIdleDeadlineAt: 40_000, + }), + [ + 'wrapper_lease', + { + state: 'owns_wrapper', + nextInstanceGeneration: 2, + instance: { instanceId: 'instance_deadlines', instanceGeneration: 1 }, + startupDeadlineAt: 60_000, + }, + ], + [ + 'disconnect_grace', + { + wrapperRunId: WRAPPER_RUN_ID, + disconnectedAt: 5_000, + wsCloseCode: 1006, + wsCloseReason: 'lost connection', + wrapperGeneration: 4, + wrapperConnectionId: WRAPPER_CONNECTION_ID, + }, + ], + ]); + await putSessionMessageState(harness.storage, acceptedMessage()); + + const deadlines = await harness.supervisor.nextMaintenanceDeadlines(); + + expect(deadlines).toHaveLength(5); + expect(deadlines).toEqual(expect.arrayContaining([60_000, 20_000, 15_000, 30_000, 40_000])); + expect(Math.min(...deadlines)).toBe(15_000); + }); + it('reconciles accepted idle work after its root-idle deadline', async () => { const harness = createHarness([ liveRuntimeState({ @@ -574,24 +679,279 @@ describe('WrapperSupervisor', () => { }); }); - it('cleans up an idle keep-warm wrapper only after the deadline has no queued or accepted work', async () => { - const harness = createHarness([liveRuntimeState({ wrapperIdleDeadlineAt: 9_000 })]); + it('retains successful idle ownership with a bounded physical warm deadline', async () => { + const harness = createHarness([liveRuntimeState(), OWNED_WRAPPER_LEASE]); + + await harness.supervisor.onTerminalEvent({ + wrapperRunId: WRAPPER_RUN_ID, + status: 'completed', + }); + + await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ + state: 'owns_wrapper', + keepWarmUntil: expect.any(Number), + }); + await expect(getWrapperRuntimeState(harness.storage)).resolves.toEqual({ + wrapperGeneration: 5, + }); + }); + + it('requests durable cleanup when a physical keep-warm deadline expires without work', async () => { + const harness = createHarness([ + liveRuntimeState({ wrapperIdleDeadlineAt: 9_000 }), + [ + 'wrapper_lease', + { + ...(OWNED_WRAPPER_LEASE[1] as object), + keepWarmUntil: 9_000, + }, + ], + ]); await harness.supervisor.runMaintenance(10_000); const runtimeState = await getWrapperRuntimeState(harness.storage); expect(runtimeState.wrapperConnectionId).toBeUndefined(); expect(runtimeState.wrapperGeneration).toBe(5); - expect(harness.stops).toEqual(['keep-warm-expired']); + await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ + state: 'stop_needed', + reason: 'keep-warm-expired', + }); + expect(harness.stops).toEqual([]); }); - it('releases a gate-waiting callback when keep-warm cleanup abandons wrapper terminal state', async () => { - const harness = createHarness([liveRuntimeState({ wrapperIdleDeadlineAt: 9_000 })], { - metadata: { - ...createMetadata(), - finalization: { gateThreshold: 'warning' }, - }, + it('turns an expired startup allowance into verified cleanup work', async () => { + const harness = createHarness([ + [ + 'wrapper_lease', + { + state: 'owns_wrapper', + nextInstanceGeneration: 2, + instance: { instanceId: 'instance_starting', instanceGeneration: 1 }, + startupDeadlineAt: 1_000, + }, + ], + ]); + + await harness.supervisor.runMaintenance(1_001); + + expect(harness.stopWrappers).toHaveBeenCalledWith( + expect.objectContaining({ + target: { + kind: 'instance', + instance: { instanceId: 'instance_starting', instanceGeneration: 1 }, + }, + reason: 'startup-failed', + }) + ); + await expect(getWrapperLease(harness.storage)).resolves.toEqual({ + state: 'none', + nextInstanceGeneration: 2, }); + }); + + it('repairs an expired readiness deadline when accepted work already proves delivery', async () => { + const harness = createHarness([ + liveRuntimeState({ lastWrapperMessageAt: 1_000 }), + [ + 'wrapper_lease', + { + state: 'owns_wrapper', + nextInstanceGeneration: 2, + instance: { instanceId: 'instance_accepted', instanceGeneration: 1 }, + startupDeadlineAt: 1_000, + }, + ], + ]); + await putSessionMessageState(harness.storage, acceptedMessage()); + + await harness.supervisor.runMaintenance(1_001); + + await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ + state: 'owns_wrapper', + startupDeadlineAt: undefined, + }); + expect(harness.stopWrappers).not.toHaveBeenCalled(); + }); + + it('claims due cleanup before provider I/O and confirms verified absence', async () => { + const harness = createHarness([ + [ + 'wrapper_lease', + { + state: 'stop_needed', + nextInstanceGeneration: 2, + target: { + kind: 'instance', + instance: { instanceId: 'instance_stop', instanceGeneration: 1 }, + }, + reason: 'startup-failed', + requestedAt: 1_000, + nextAttemptAt: 1_000, + attempts: 0, + }, + ], + ]); + harness.stopWrappers.mockImplementationOnce(async () => { + await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ + state: 'stopping', + attemptId: expect.any(String), + }); + expect(harness.requestedAlarms).toHaveLength(1); + return { status: 'absent' }; + }); + + await harness.supervisor.runMaintenance(1_001); + + expect(harness.stopWrappers).toHaveBeenCalledOnce(); + await expect(getWrapperLease(harness.storage)).resolves.toEqual({ + state: 'none', + nextInstanceGeneration: 2, + }); + }); + + it.each([ + { result: { status: 'still-present', observed: [] }, label: 'still-present' }, + { result: { status: 'inspection-failed', error: 'unavailable' }, label: 'inspection-failed' }, + ])('retries a $label cleanup result with its target preserved', async ({ result }) => { + const harness = createHarness([ + [ + 'wrapper_lease', + { + state: 'stop_needed', + nextInstanceGeneration: 4, + target: { kind: 'session' }, + reason: 'unexpected-wrapper', + requestedAt: 1_000, + nextAttemptAt: 1_000, + attempts: 0, + }, + ], + ]); + harness.stopWrappers.mockResolvedValueOnce(result); + + await harness.supervisor.runMaintenance(1_001); + + await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ + state: 'stop_needed', + target: { kind: 'session' }, + attempts: 1, + nextAttemptAt: expect.any(Number), + }); + }); + + it('backs off unconfirmed physical cleanup retries through the saturated delay tail', async () => { + const harness = createHarness([ + [ + 'wrapper_lease', + { + state: 'stop_needed', + nextInstanceGeneration: 2, + target: { kind: 'session' }, + reason: 'unexpected-wrapper', + requestedAt: 1_000, + nextAttemptAt: 1_000, + attempts: 0, + }, + ], + ]); + harness.stopWrappers.mockResolvedValue({ status: 'still-present', observed: [] }); + const expectedDelays = [5_000, 30_000, 120_000, 300_000, 300_000]; + let now = 1_001; + + for (const [index, expectedDelay] of expectedDelays.entries()) { + await harness.supervisor.runMaintenance(now); + const lease = await getWrapperLease(harness.storage); + if (lease.state !== 'stop_needed') { + throw new Error(`Expected a retryable cleanup lease after attempt ${index + 1}`); + } + expect(lease.attempts).toBe(index + 1); + expect(lease.nextAttemptAt - now).toBe(expectedDelay); + now = lease.nextAttemptAt; + } + }); + + it('retries thrown cleanup and does not issue a parallel stop during a valid watchdog', async () => { + const harness = createHarness([ + [ + 'wrapper_lease', + { + state: 'stop_needed', + nextInstanceGeneration: 2, + target: { kind: 'session' }, + reason: 'observation-failed', + requestedAt: 1_000, + nextAttemptAt: 1_000, + attempts: 0, + }, + ], + ]); + harness.stopWrappers.mockRejectedValueOnce(new Error('stop failed')); + await harness.supervisor.runMaintenance(1_001); + await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ state: 'stop_needed' }); + + await harness.storage.put('wrapper_lease', { + state: 'stopping', + nextInstanceGeneration: 2, + target: { kind: 'session' }, + reason: 'observation-failed', + requestedAt: 1_000, + attemptId: 'inflight', + attemptStartedAt: 2_000, + attemptDeadlineAt: 30_000, + attempts: 2, + }); + await harness.supervisor.runMaintenance(20_000); + expect(harness.stopWrappers).toHaveBeenCalledOnce(); + expect(await harness.supervisor.nextMaintenanceDeadlines()).toContain(30_000); + }); + + it('expires a stale watchdog into retryable cleanup without settling a late attempt', async () => { + const harness = createHarness([ + [ + 'wrapper_lease', + { + state: 'stopping', + nextInstanceGeneration: 2, + target: { kind: 'session' }, + reason: 'observation-failed', + requestedAt: 1_000, + attemptId: 'expired', + attemptStartedAt: 1_000, + attemptDeadlineAt: 2_000, + attempts: 1, + }, + ], + ]); + + await harness.supervisor.runMaintenance(2_001); + + expect(harness.stopWrappers).not.toHaveBeenCalled(); + await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ + state: 'stop_needed', + attempts: 1, + }); + }); + + it('releases a gate-waiting callback when keep-warm cleanup abandons wrapper terminal state', async () => { + const harness = createHarness( + [ + liveRuntimeState({ wrapperIdleDeadlineAt: 9_000 }), + [ + 'wrapper_lease', + { + ...(OWNED_WRAPPER_LEASE[1] as object), + keepWarmUntil: 9_000, + }, + ], + ], + { + metadata: { + ...createMetadata(), + finalization: { gateThreshold: 'warning' }, + }, + } + ); await putSessionMessageState(harness.storage, { ...acceptedMessage(), callbackRequired: true, @@ -604,7 +964,11 @@ describe('WrapperSupervisor', () => { await harness.supervisor.runMaintenance(10_000); - expect(harness.stops).toEqual(['keep-warm-expired']); + await expect(getWrapperLease(harness.storage)).resolves.toMatchObject({ + state: 'stop_needed', + reason: 'keep-warm-expired', + }); + expect(harness.stops).toEqual([]); expect(harness.callbackJobs).toHaveLength(1); expect(harness.callbackJobs[0].payload).toMatchObject({ messageId: MESSAGE_ID, diff --git a/services/cloud-agent-next/src/session/wrapper-supervisor.ts b/services/cloud-agent-next/src/session/wrapper-supervisor.ts index 9cc2cf6213..0a7c7c6623 100644 --- a/services/cloud-agent-next/src/session/wrapper-supervisor.ts +++ b/services/cloud-agent-next/src/session/wrapper-supervisor.ts @@ -1,6 +1,11 @@ import { z } from 'zod'; import { logger } from '../logger.js'; import type { SessionMetadata } from '../persistence/session-metadata.js'; +import type { + StopWrappersResult, + WrapperStopReason, + WrapperStopTarget, +} from '../agent-sandbox/protocol.js'; import type { AgentRuntime } from './agent-runtime.js'; import { WRAPPER_NO_OUTPUT_TIMEOUT_MS, WRAPPER_PING_INTERVAL_MS } from './agent-runtime.js'; import type { MessageSettlementOutbox } from './message-settlement-outbox.js'; @@ -16,20 +21,27 @@ import { clearCurrentWrapperRuntimeLivenessState, clearWrapperIdleState, clearWrapperRuntimeIdentity, + getWrapperLease, getWrapperRuntimeState, hasCompleteWrapperIdentity, + IDLE_KEEP_WARM_MS, IDLE_RECONCILIATION_GRACE_MS, isCurrentWrapperConnection, markWrapperPingSent, + nextWrapperLeaseDeadline, + putWrapperLease, recordMeaningfulWrapperOutput, recordRootSessionIdle, recordWrapperPong, + reduceWrapperLease, type WrapperConnectionFence, type WrapperRuntimeState, } from './wrapper-runtime-state.js'; const DISCONNECT_GRACE_MS = 10_000; const WRAPPER_PING_TIMEOUT_MS = 30_000; +const WRAPPER_STOP_ATTEMPT_TIMEOUT_MS = 45_000; +const WRAPPER_STOP_RETRY_DELAYS_MS = [5_000, 30_000, 120_000, 300_000]; const DISCONNECT_GRACE_KEY = 'disconnect_grace'; const disconnectGraceStateSchema = z.object({ @@ -96,6 +108,7 @@ export type WrapperSupervisor = { ): Promise; onDisconnected(input: WrapperDisconnectedInput): Promise; onTerminalEvent(params: WrapperTerminalEvent): Promise; + requestPhysicalWrapperStop(reason: WrapperStopReason, target?: WrapperStopTarget): Promise; clearDisconnectGrace(): Promise; runMaintenance(now: number): Promise; nextMaintenanceDeadlines(): Promise; @@ -103,7 +116,7 @@ export type WrapperSupervisor = { export type WrapperSupervisorDependencies = { storage: WrapperSupervisorStorage; - agentRuntime: Pick; + agentRuntime: Pick; messageSettlementOutbox: Pick< MessageSettlementOutbox, | 'terminalizeSessionMessageOnce' @@ -127,6 +140,12 @@ export type WrapperSupervisorDependencies = { wrapperConnectionId: string; }) => Promise; clearInterruptRequest: () => Promise; + stopWrappers?: (request: { + target: WrapperStopTarget; + attemptId: string; + reason: WrapperStopReason; + }) => Promise; + requestAlarmAtOrBefore?: (deadline: number) => Promise; getSessionIdForLogs: () => string | undefined; }; @@ -189,6 +208,8 @@ export function createWrapperSupervisor( observeCorrelatedAgentActivity, hasActiveIngestConnection, clearInterruptRequest, + stopWrappers, + requestAlarmAtOrBefore, getSessionIdForLogs, } = dependencies; @@ -299,6 +320,30 @@ export function createWrapperSupervisor( ); } + async function retainPhysicalWrapperWarm(now: number): Promise { + const lease = await getWrapperLease(storage); + if (lease.state !== 'owns_wrapper') return; + const warm = reduceWrapperLease(lease, { + type: 'retain_warm', + instanceId: lease.instance.instanceId, + keepWarmUntil: now + IDLE_KEEP_WARM_MS, + }); + await putWrapperLease(storage, warm); + + const runtimeState = await getWrapperRuntimeState(storage); + if (runtimeState.wrapperConnectionId) { + await clearWrapperRuntimeIdentity( + storage, + { + wrapperGeneration: runtimeState.wrapperGeneration, + wrapperConnectionId: runtimeState.wrapperConnectionId, + }, + { incrementGeneration: true } + ); + } + await requestAlarmAtOrBefore?.(now + IDLE_KEEP_WARM_MS); + } + async function observeRootIdle( wrapperGeneration: number, wrapperConnectionId: string, @@ -361,6 +406,29 @@ export function createWrapperSupervisor( await startDisconnectGrace(input); } + async function requestPhysicalWrapperStop( + reason: WrapperStopReason, + target?: WrapperStopTarget + ): Promise { + const current = await getWrapperLease(storage); + const resolvedTarget = + target ?? + (current.state === 'owns_wrapper' + ? { kind: 'instance' as const, instance: current.instance } + : { kind: 'session' as const }); + const now = Date.now(); + const next = reduceWrapperLease(current, { + type: 'request_stop', + target: resolvedTarget, + reason, + now, + }); + if (next !== current) { + await putWrapperLease(storage, next); + await requestAlarmAtOrBefore?.(now); + } + } + async function handleUnhealthyWrapper( state: WrapperRuntimeState, error: string, @@ -375,6 +443,8 @@ export function createWrapperSupervisor( }) .warn('Handling unhealthy wrapper runtime'); + await requestPhysicalWrapperStop('unhealthy-wrapper'); + const acceptedMessages = await listNonTerminalAcceptedMessages(storage, state.wrapperRunId); for (const message of acceptedMessages) { const activityObserved = message.agentActivityObservedAt !== undefined; @@ -399,8 +469,6 @@ export function createWrapperSupervisor( state.wrapperConnectionId ); } - - await agentRuntime.stopWrapperProcess('unhealthy-wrapper'); } async function checkDisconnectGrace(now: number): Promise { @@ -408,17 +476,18 @@ export function createWrapperSupervisor( if (!graceState) return; if (now - graceState.disconnectedAt < DISCONNECT_GRACE_MS) return; - await storage.delete(DISCONNECT_GRACE_KEY); const { wrapperRunId } = graceState; const state = await getWrapperRuntimeState(storage); if ( state.wrapperRunId !== wrapperRunId || state.wrapperGeneration !== graceState.wrapperGeneration ) { + await storage.delete(DISCONNECT_GRACE_KEY); await releaseWrapperTerminalWaitForIdleBatchForWrapperRun(wrapperRunId); return; } if (state.wrapperConnectionId !== graceState.wrapperConnectionId) { + await storage.delete(DISCONNECT_GRACE_KEY); await releaseWrapperTerminalWaitForIdleBatchForWrapperRun(wrapperRunId); return; } @@ -433,6 +502,7 @@ export function createWrapperSupervisor( logger .withFields({ wrapperRunId }) .info('Wrapper reconnected during grace period — skipping failure'); + await storage.delete(DISCONNECT_GRACE_KEY); return; } @@ -441,6 +511,7 @@ export function createWrapperSupervisor( logger .withFields({ wrapperRunId }) .info('No accepted messages during grace period - skipping failure'); + await storage.delete(DISCONNECT_GRACE_KEY); await releaseWrapperTerminalWaitForIdleBatch(); return; } @@ -448,6 +519,8 @@ export function createWrapperSupervisor( logger .withFields({ wrapperRunId, messageCount: acceptedMessages.length }) .warn('Grace period expired - failing accepted messages'); + await requestPhysicalWrapperStop('unhealthy-wrapper'); + await storage.delete(DISCONNECT_GRACE_KEY); for (const message of acceptedMessages) { const activityObserved = message.agentActivityObservedAt !== undefined; await messageSettlementOutbox.terminalizeSessionMessageOnce(message.messageId, { @@ -467,7 +540,6 @@ export function createWrapperSupervisor( }, { incrementGeneration: true } ); - await agentRuntime.stopWrapperProcess('unhealthy-wrapper'); await releaseWrapperTerminalWaitForIdleBatch(); } @@ -609,8 +681,26 @@ export function createWrapperSupervisor( }) .info('Idle reconciliation processing accepted messages'); + let failedTerminalObserved = !metadata.auth.kiloSessionId; + if (!failedTerminalObserved && metadata.auth.kiloSessionId) { + failedTerminalObserved = acceptedMessages.some(message => { + const assistantMessage = getAssistantMessageForUserMessage( + metadata.identity.sessionId, + metadata.auth.kiloSessionId ?? '', + message.messageId + ); + return ( + !assistantMessage || getAssistantErrorMessage(assistantMessage.info.error) !== undefined + ); + }); + } + if (failedTerminalObserved) { + await requestPhysicalWrapperStop('terminal-failed'); + } + for (const message of acceptedMessages) { if (!metadata.auth.kiloSessionId) { + failedTerminalObserved = true; await messageSettlementOutbox.terminalizeSessionMessageOnce(message.messageId, { kind: 'failed', reason: 'missing_assistant_reply', @@ -628,6 +718,7 @@ export function createWrapperSupervisor( message.messageId ); if (!assistantMessage) { + failedTerminalObserved = true; await messageSettlementOutbox.terminalizeSessionMessageOnce(message.messageId, { kind: 'failed', reason: 'missing_assistant_reply', @@ -642,6 +733,7 @@ export function createWrapperSupervisor( await observeCorrelatedAgentActivity?.(message.messageId); const assistantError = getAssistantErrorMessage(assistantMessage.info.error); if (assistantError !== undefined) { + failedTerminalObserved = true; await messageSettlementOutbox.terminalizeSessionMessageOnce(message.messageId, { kind: 'failed', reason: 'assistant_error', @@ -660,6 +752,16 @@ export function createWrapperSupervisor( }); } + if (failedTerminalObserved) { + if (state.wrapperConnectionId) { + await clearWrapperRuntimeIdentity(storage, { + wrapperGeneration: state.wrapperGeneration, + wrapperConnectionId: state.wrapperConnectionId, + }); + } + } else { + await retainPhysicalWrapperWarm(now); + } await messageSettlementOutbox.finalizeIdleBatchCallbackIfReady(); logger .withFields({ @@ -671,9 +773,12 @@ export function createWrapperSupervisor( } async function checkKeepWarmCleanup(now: number): Promise { + const lease = await getWrapperLease(storage); + if (lease.state === 'owns_wrapper' && lease.startupDeadlineAt !== undefined) return; const wrapperState = await getWrapperRuntimeState(storage); - if (wrapperState.wrapperIdleDeadlineAt === undefined) return; - if (wrapperState.wrapperIdleDeadlineAt > now) return; + const keepWarmUntil = + lease.state === 'owns_wrapper' ? lease.keepWarmUntil : wrapperState.wrapperIdleDeadlineAt; + if (keepWarmUntil === undefined || keepWarmUntil > now) return; const pendingCount = await countPendingSessionMessages(storage); const acceptedMessages = await listNonTerminalAcceptedMessages( @@ -708,7 +813,97 @@ export function createWrapperSupervisor( ); } await releaseWrapperTerminalWaitForIdleBatch(); - await agentRuntime.stopWrapperProcess('keep-warm-expired'); + await requestPhysicalWrapperStop('keep-warm-expired'); + } + + function stopRetryAt(now: number, attempts: number): number { + const delay = + WRAPPER_STOP_RETRY_DELAYS_MS[Math.min(attempts - 1, WRAPPER_STOP_RETRY_DELAYS_MS.length - 1)]; + return now + delay; + } + + async function reconcilePhysicalCleanup(now: number): Promise { + if (!stopWrappers) return; + let lease = await getWrapperLease(storage); + if ( + lease.state === 'owns_wrapper' && + lease.startupDeadlineAt !== undefined && + now >= lease.startupDeadlineAt + ) { + const runtimeState = await getWrapperRuntimeState(storage); + if (await hasActiveWrapperWork(runtimeState)) { + lease = reduceWrapperLease(lease, { + type: 'delivery_accepted', + instanceId: lease.instance.instanceId, + }); + } else { + lease = reduceWrapperLease(lease, { + type: 'request_stop', + target: { kind: 'instance', instance: lease.instance }, + reason: 'startup-failed', + now, + }); + } + await putWrapperLease(storage, lease); + } + if (lease.state === 'stopping') { + if (now < lease.attemptDeadlineAt) return; + lease = reduceWrapperLease(lease, { + type: 'stop_attempt_expired', + attemptId: lease.attemptId, + retryAt: stopRetryAt(now, lease.attempts), + }); + await putWrapperLease(storage, lease); + return; + } + if (lease.state !== 'stop_needed' || now < lease.nextAttemptAt) return; + + const attemptId = crypto.randomUUID(); + const stopping = reduceWrapperLease(lease, { + type: 'begin_stop_attempt', + attemptId, + now, + attemptDeadlineAt: now + WRAPPER_STOP_ATTEMPT_TIMEOUT_MS, + }); + if (stopping.state !== 'stopping') return; + await putWrapperLease(storage, stopping); + await requestAlarmAtOrBefore?.(stopping.attemptDeadlineAt); + + let result: StopWrappersResult; + try { + result = await stopWrappers({ + target: stopping.target, + attemptId, + reason: stopping.reason, + }); + } catch (error) { + result = { + status: 'inspection-failed', + error: error instanceof Error ? error.message : String(error), + }; + } + + const latest = await getWrapperLease(storage); + if (result.status === 'absent') { + await putWrapperLease( + storage, + reduceWrapperLease(latest, { type: 'stop_absent', attemptId }) + ); + return; + } + const error = + result.status === 'inspection-failed' + ? result.error + : (result.error ?? 'Wrapper remains present'); + await putWrapperLease( + storage, + reduceWrapperLease(latest, { + type: 'stop_not_confirmed', + attemptId, + retryAt: stopRetryAt(now, stopping.attempts), + error, + }) + ); } async function onTerminalEvent(params: WrapperTerminalEvent): Promise { @@ -738,6 +933,9 @@ export function createWrapperSupervisor( .info('Wrapper terminal event received by supervisor'); if (status === 'failed' || status === 'interrupted') { + await requestPhysicalWrapperStop( + status === 'failed' ? 'terminal-failed' : 'terminal-interrupted' + ); const acceptedMessages = await listNonTerminalAcceptedMessages(storage, wrapperRunId); for (const message of acceptedMessages) { if (status === 'failed') { @@ -768,10 +966,7 @@ export function createWrapperSupervisor( if (status === 'completed') { const acceptedMessages = await listNonTerminalAcceptedMessages(storage, wrapperRunId); if (acceptedMessages.length === 0) { - await clearWrapperRuntimeIdentity(storage, { - wrapperGeneration: state.wrapperGeneration, - wrapperConnectionId: state.wrapperConnectionId, - }); + await retainPhysicalWrapperWarm(Date.now()); await clearInterruptRequest(); } } else { @@ -791,6 +986,7 @@ export function createWrapperSupervisor( } async function runMaintenance(now: number): Promise { + await reconcilePhysicalCleanup(now); await checkDisconnectGrace(now); await checkWrapperLiveness(now); await checkIdleReconciliation(now); @@ -799,6 +995,10 @@ export function createWrapperSupervisor( async function nextMaintenanceDeadlines(): Promise { const deadlines: number[] = []; + const physicalDeadline = nextWrapperLeaseDeadline(await getWrapperLease(storage)); + if (physicalDeadline !== undefined) { + deadlines.push(physicalDeadline); + } const livenessDeadline = await getNextWrapperLivenessDeadline(); if (livenessDeadline !== undefined) { deadlines.push(livenessDeadline); @@ -829,6 +1029,7 @@ export function createWrapperSupervisor( observeRootIdle, onDisconnected, onTerminalEvent, + requestPhysicalWrapperStop, clearDisconnectGrace, runMaintenance, nextMaintenanceDeadlines, diff --git a/services/cloud-agent-next/src/terminal/access.test.ts b/services/cloud-agent-next/src/terminal/access.test.ts index c14515ea1a..69bf31f443 100644 --- a/services/cloud-agent-next/src/terminal/access.test.ts +++ b/services/cloud-agent-next/src/terminal/access.test.ts @@ -1,12 +1,10 @@ import { describe, expect, it, vi } from 'vitest'; +import type { AgentSandbox } from '../agent-sandbox/protocol.js'; import type { CloudAgentSessionState } from '../persistence/types.js'; -import { WRAPPER_VERSION } from '../shared/wrapper-version.js'; -import type { Env, SandboxInstance } from '../types.js'; +import type { Env } from '../types.js'; import { resolveTerminalWrapperClient, validateTerminalMetadata } from './access.js'; -vi.mock('@cloudflare/sandbox', () => ({ - getSandbox: vi.fn(), -})); +vi.mock('@cloudflare/sandbox', () => ({ getSandbox: vi.fn() })); const baseMetadata = { metadataSchemaVersion: 2, @@ -97,58 +95,61 @@ describe('validateTerminalMetadata', () => { }); }); +function sandboxWithTerminalResult( + getRunningTerminalClient: AgentSandbox['getRunningTerminalClient'] +): AgentSandbox { + return { + ensureWrapper: vi.fn(), + discoverSessionWrappers: vi.fn(), + stopWrappers: vi.fn(), + probeHealth: vi.fn(), + getRunningWrapper: vi.fn(), + getRunningTerminalClient, + readWrapperLogs: vi.fn(), + keepAlive: vi.fn(), + delete: vi.fn(), + }; +} + describe('resolveTerminalWrapperClient', () => { - it('returns a healthy existing wrapper client', async () => { - const sandbox = {} as SandboxInstance; + it('returns the ready terminal client from AgentSandbox', async () => { const client = { - health: vi.fn().mockResolvedValue({ - healthy: true, - state: 'idle', - version: WRAPPER_VERSION, - sessionId: 'kilo-session-1', - }), + health: vi.fn(), createTerminal: vi.fn(), resizeTerminal: vi.fn(), closeTerminal: vi.fn(), connectTerminal: vi.fn(), }; + const getRunningTerminalClient = vi.fn().mockResolvedValue({ status: 'ready', client }); const result = await resolveTerminalWrapperClient( { - env: { PER_SESSION_SANDBOX_ORG_IDS: '' } as Env, + env: {} as Env, metadata: baseMetadata, sessionId: baseMetadata.identity.sessionId, }, { - getSandboxInstance: vi.fn().mockReturnValue(sandbox), - findWrapperForSession: vi.fn().mockResolvedValue({ port: 5050 }), - createClient: vi.fn().mockReturnValue(client), + createSandbox: vi.fn().mockReturnValue(sandboxWithTerminalResult(getRunningTerminalClient)), } ); - expect(result.success).toBe(true); - expect(result.data).toMatchObject({ - client, - sandbox, - port: 5050, - }); - expect(client.health).toHaveBeenCalledTimes(1); + expect(result).toEqual({ success: true, data: { client } }); + expect(getRunningTerminalClient).toHaveBeenCalledOnce(); }); - it('returns unavailable when no existing wrapper process is running', async () => { - const findWrapperForSession = vi.fn().mockResolvedValue(null); - const health = vi.fn(); - + it('returns unavailable when no wrapper process is running', async () => { const result = await resolveTerminalWrapperClient( { - env: { PER_SESSION_SANDBOX_ORG_IDS: '' } as Env, + env: {} as Env, metadata: baseMetadata, sessionId: baseMetadata.identity.sessionId, }, { - getSandboxInstance: vi.fn().mockReturnValue({} as SandboxInstance), - findWrapperForSession, - createClient: vi.fn().mockReturnValue({ health }), + createSandbox: vi + .fn() + .mockReturnValue( + sandboxWithTerminalResult(vi.fn().mockResolvedValue({ status: 'not-running' })) + ), } ); @@ -156,7 +157,27 @@ describe('resolveTerminalWrapperClient', () => { success: false, error: 'Terminal is unavailable because the session wrapper is not running', }); - expect(findWrapperForSession).toHaveBeenCalledTimes(1); - expect(health).not.toHaveBeenCalled(); + }); + + it('preserves the unhealthy-wrapper terminal diagnostic', async () => { + const result = await resolveTerminalWrapperClient( + { + env: {} as Env, + metadata: baseMetadata, + sessionId: baseMetadata.identity.sessionId, + }, + { + createSandbox: vi + .fn() + .mockReturnValue( + sandboxWithTerminalResult(vi.fn().mockResolvedValue({ status: 'unhealthy' })) + ), + } + ); + + expect(result).toEqual({ + success: false, + error: 'Terminal is unavailable because the session wrapper is not healthy', + }); }); }); diff --git a/services/cloud-agent-next/src/terminal/access.ts b/services/cloud-agent-next/src/terminal/access.ts index 120343c796..ff5bd35fab 100644 --- a/services/cloud-agent-next/src/terminal/access.ts +++ b/services/cloud-agent-next/src/terminal/access.ts @@ -1,15 +1,8 @@ +import { createAgentSandbox } from '../agent-sandbox/factory.js'; +import type { AgentSandbox } from '../agent-sandbox/protocol.js'; +import type { WrapperHealthResponse, WrapperPty } from '../kilo/wrapper-client.js'; import type { CloudAgentSessionState, OperationResult } from '../persistence/types.js'; -import { getSandbox } from '@cloudflare/sandbox'; -import { SANDBOX_SLEEP_AFTER_SECONDS } from '../core/lease.js'; -import { - WrapperContainerClient, - type WrapperHealthResponse, - type WrapperPty, -} from '../kilo/wrapper-client.js'; -import { findWrapperForSession } from '../kilo/wrapper-manager.js'; -import { generateSandboxId, getSandboxNamespace } from '../sandbox-id.js'; -import { WRAPPER_VERSION } from '../shared/wrapper-version.js'; -import type { Env, SandboxId, SandboxInstance } from '../types.js'; +import type { Env } from '../types.js'; const TERMINAL_SESSION_PLATFORMS = new Set(['cloud-agent', 'cloud-agent-web', 'slack']); @@ -55,21 +48,11 @@ export type TerminalWrapperClient = { }; type ResolveTerminalWrapperDeps = { - getSandboxInstance(params: { env: Env; sandboxId: SandboxId }): SandboxInstance; - findWrapperForSession( - sandbox: SandboxInstance, - sessionId: string - ): Promise<{ port: number } | null>; - createClient(params: { sandbox: SandboxInstance; port: number }): TerminalWrapperClient; + createSandbox(env: Env, metadata: CloudAgentSessionState): AgentSandbox; }; const defaultDeps: ResolveTerminalWrapperDeps = { - getSandboxInstance: ({ env, sandboxId }) => - getSandbox(getSandboxNamespace(env, sandboxId), sandboxId, { - sleepAfter: SANDBOX_SLEEP_AFTER_SECONDS, - }), - findWrapperForSession, - createClient: ({ sandbox, port }) => new WrapperContainerClient({ sandbox, port }), + createSandbox: createAgentSandbox, }; export async function resolveTerminalWrapperClient( @@ -79,62 +62,26 @@ export async function resolveTerminalWrapperClient( sessionId: string; }, deps: ResolveTerminalWrapperDeps = defaultDeps -): Promise< - OperationResult<{ - client: TerminalWrapperClient; - sandbox: SandboxInstance; - sandboxId: SandboxId; - port: number; - }> -> { +): Promise> { const metadataResult = validateTerminalMetadata(params.metadata, params.sessionId); if (!metadataResult.success || !metadataResult.data) { return { success: false, error: metadataResult.error }; } - const { metadata } = metadataResult.data; - const sandboxId = - metadata.workspace?.sandboxId ?? - (await generateSandboxId( - params.env.PER_SESSION_SANDBOX_ORG_IDS, - metadata.identity.orgId, - metadata.identity.userId, - metadata.identity.sessionId, - metadata.identity.botId - )); - const sandbox = deps.getSandboxInstance({ env: params.env, sandboxId }); - const wrapper = await deps.findWrapperForSession(sandbox, metadata.identity.sessionId); - - if (!wrapper) { + const terminal = await deps + .createSandbox(params.env, metadataResult.data.metadata) + .getRunningTerminalClient(); + if (terminal.status === 'not-running') { return { success: false, error: 'Terminal is unavailable because the session wrapper is not running', }; } - - const client = deps.createClient({ sandbox, port: wrapper.port }); - try { - const health = await client.health(); - if (!health.healthy || health.version !== WRAPPER_VERSION) { - return { - success: false, - error: 'Terminal is unavailable because the session wrapper is not healthy', - }; - } - } catch { + if (terminal.status === 'unhealthy') { return { success: false, error: 'Terminal is unavailable because the session wrapper is not healthy', }; } - - return { - success: true, - data: { - client, - sandbox, - sandboxId, - port: wrapper.port, - }, - }; + return { success: true, data: { client: terminal.client } }; } diff --git a/services/cloud-agent-next/src/workspace-errors.ts b/services/cloud-agent-next/src/workspace-errors.ts index d2842d1a4e..1fd2af2d84 100644 --- a/services/cloud-agent-next/src/workspace-errors.ts +++ b/services/cloud-agent-next/src/workspace-errors.ts @@ -9,3 +9,47 @@ export class WorkspaceFilesystemPreparationError extends Error { this.target = target; } } + +export type WorkspaceCapacityAdmissionRejectionDetails = { + availableMB: number; + thresholdMB: number; + cleaned: number; + skipped: number; +}; + +export class WorkspaceCapacityAdmissionRejectedError extends Error { + readonly availableMB: number; + readonly thresholdMB: number; + readonly cleaned: number; + readonly skipped: number; + + constructor(details: WorkspaceCapacityAdmissionRejectionDetails) { + super( + `Workspace admission rejected: ${details.availableMB} MB available below ${details.thresholdMB} MB threshold after cleanup` + ); + this.name = 'WorkspaceCapacityAdmissionRejectedError'; + this.availableMB = details.availableMB; + this.thresholdMB = details.thresholdMB; + this.cleaned = details.cleaned; + this.skipped = details.skipped; + } +} + +export class WorkspaceCapacityInspectionUnavailableError extends Error { + constructor(message: string, cause: unknown) { + super(message, { cause }); + this.name = 'WorkspaceCapacityInspectionUnavailableError'; + } +} + +export function isSandboxFilesystemUnusableError(error: unknown): boolean { + const message = error instanceof Error ? error.message : String(error); + return /\bENOSPC\b|no space left on device/i.test(message); +} + +export class SandboxCapacityInspectionError extends Error { + constructor(message: string, cause: unknown) { + super(message, { cause }); + this.name = 'SandboxCapacityInspectionError'; + } +} diff --git a/services/cloud-agent-next/src/workspace.test.ts b/services/cloud-agent-next/src/workspace.test.ts index d8030e6f36..44771d9411 100644 --- a/services/cloud-agent-next/src/workspace.test.ts +++ b/services/cloud-agent-next/src/workspace.test.ts @@ -41,7 +41,12 @@ import { LOW_DISK_THRESHOLD_MB, STALE_DIR_MIN_AGE_SECONDS, } from './workspace'; -import { WorkspaceFilesystemPreparationError } from './workspace-errors'; +import { + SandboxCapacityInspectionError, + WorkspaceCapacityAdmissionRejectedError, + WorkspaceCapacityInspectionUnavailableError, + WorkspaceFilesystemPreparationError, +} from './workspace-errors'; import type { ExecutionSession, SandboxInstance } from './types'; describe('setupWorkspace', () => { @@ -789,7 +794,7 @@ describe('disk space checking', () => { } as unknown as SandboxInstance; }); - it('cleans up sessions with no running wrapper', async () => { + it('cleans up sessions with no running wrapper and reports reclamation', async () => { const oldMtime = String(Math.floor(Date.now() / 1000) - STALE_DIR_MIN_AGE_SECONDS - 60); mockSandboxExec .mockResolvedValueOnce({ @@ -805,9 +810,10 @@ describe('disk space checking', () => { mockListProcesses.mockResolvedValue([]); - await cleanupStaleWorkspaces(fakeSandbox, '/workspace/org/user', 'agent_current-aaaa'); + await expect( + cleanupStaleWorkspaces(fakeSandbox, '/workspace/org/user', 'agent_current-aaaa') + ).resolves.toEqual({ cleaned: 1, skipped: 1 }); - // listProcesses is called exactly once (not per session) expect(mockListProcesses).toHaveBeenCalledTimes(1); const execCalls = mockSandboxExec.mock.calls.map((c: string[]) => c[0]); @@ -884,6 +890,22 @@ describe('disk space checking', () => { expect(mockListProcesses).not.toHaveBeenCalled(); }); + it('does not count failed removal as reclaimed workspace', async () => { + const oldMtime = String(Math.floor(Date.now() / 1000) - STALE_DIR_MIN_AGE_SECONDS - 60); + mockSandboxExec + .mockResolvedValueOnce({ exitCode: 0, stdout: 'agent_stale-aaaa\n', stderr: '' }) + .mockResolvedValueOnce(dockerSocketPath) + .mockResolvedValueOnce(dockerPsEmpty) + .mockResolvedValueOnce({ exitCode: 0, stdout: `${oldMtime}\n`, stderr: '' }) + .mockResolvedValueOnce({ exitCode: 1, stdout: '', stderr: 'rm failed' }) + .mockResolvedValueOnce({ exitCode: 0, stdout: '', stderr: '' }); + mockListProcesses.mockResolvedValue([]); + + await expect( + cleanupStaleWorkspaces(fakeSandbox, '/workspace/org/user', 'agent_current-aaaa') + ).resolves.toEqual({ cleaned: 0, skipped: 1 }); + }); + it('continues cleaning remaining sessions when one throws', async () => { const oldMtime = String(Math.floor(Date.now() / 1000) - STALE_DIR_MIN_AGE_SECONDS - 60); mockSandboxExec @@ -901,10 +923,9 @@ describe('disk space checking', () => { mockListProcesses.mockResolvedValue([]); - // Should not throw await expect( cleanupStaleWorkspaces(fakeSandbox, '/workspace/org/user', 'agent_current-aaaa') - ).resolves.toBeUndefined(); + ).resolves.toEqual({ cleaned: 1, skipped: 1 }); // listProcesses is called exactly once (not per session) expect(mockListProcesses).toHaveBeenCalledTimes(1); @@ -922,10 +943,9 @@ describe('disk space checking', () => { }); mockListProcesses.mockRejectedValue(new Error('sandbox unavailable')); - // Should not throw — returns early without cleaning any sessions await expect( cleanupStaleWorkspaces(fakeSandbox, '/workspace/org/user', 'agent_current-aaaa') - ).resolves.toBeUndefined(); + ).resolves.toEqual({ cleaned: 0, skipped: 1 }); // Only the ls call — no rm calls since listProcesses failed expect(mockSandboxExec).toHaveBeenCalledTimes(1); @@ -936,7 +956,7 @@ describe('disk space checking', () => { await expect( cleanupStaleWorkspaces(fakeSandbox, '/workspace/org/user', 'agent_current-aaaa') - ).resolves.toBeUndefined(); + ).resolves.toEqual({ cleaned: 0, skipped: 0 }); }); it('skips directory entries that do not match the agent_ session ID format', async () => { @@ -1057,6 +1077,19 @@ describe('disk space checking', () => { expect(mockSandboxExec).toHaveBeenCalledTimes(4); }); + it('skips all cleanup when devcontainer wrapper inspection fails', async () => { + mockSandboxExec + .mockResolvedValueOnce({ exitCode: 0, stdout: 'agent_unknown-cccc\n', stderr: '' }) + .mockResolvedValueOnce(dockerSocketPath) + .mockResolvedValueOnce({ exitCode: 1, stdout: '', stderr: 'docker unavailable' }); + mockListProcesses.mockResolvedValue([]); + + await expect( + cleanupStaleWorkspaces(fakeSandbox, '/workspace/org/user', 'agent_current-aaaa') + ).resolves.toEqual({ cleaned: 0, skipped: 1 }); + expect(mockSandboxExec.mock.calls.every(call => !call[0].includes('rm -rf'))).toBe(true); + }); + it('skips sessions with a wrapper running inside a dev container', async () => { const oldMtime = String(Math.floor(Date.now() / 1000) - STALE_DIR_MIN_AGE_SECONDS - 60); mockSandboxExec @@ -1098,91 +1131,219 @@ describe('disk space checking', () => { } as unknown as SandboxInstance; }); - it('runs cleanup when disk space is low', async () => { + it('admits setup without cleanup when capacity is adequate', async () => { + mockSandboxExec.mockResolvedValueOnce({ + exitCode: 0, + stdout: '5242880000 10485760000\n', + stderr: '', + }); + + await expect( + checkDiskAndCleanBeforeSetup(fakeSandbox, 'org-1', 'user-1', 'agent_current-aaaa') + ).resolves.toEqual({ + availableMB: 5000, + thresholdMB: LOW_DISK_THRESHOLD_MB, + cleanup: { cleaned: 0, skipped: 0 }, + }); + + expect(mockSandboxExec).toHaveBeenCalledTimes(1); + expect(mockListProcesses).not.toHaveBeenCalled(); + }); + + it('cleans low capacity and admits after a successful recheck', async () => { const oldMtime = String(Math.floor(Date.now() / 1000) - STALE_DIR_MIN_AGE_SECONDS - 60); - // df returns low disk (1024 MB avail, 10000 MB total) mockSandboxExec .mockResolvedValueOnce({ exitCode: 0, stdout: '1073741824 10485760000\n', stderr: '', - }) // df (checkDiskSpace) + }) .mockResolvedValueOnce({ exitCode: 0, stdout: 'agent_stale-1111\nagent_current-aaaa\n', stderr: '', - }) // ls sessions/ (cleanupStaleWorkspaces) + }) + .mockResolvedValueOnce({ exitCode: 0, stdout: '/run/user/1000/docker.sock', stderr: '' }) + .mockResolvedValueOnce({ exitCode: 0, stdout: '', stderr: '' }) + .mockResolvedValueOnce({ exitCode: 0, stdout: `${oldMtime}\n`, stderr: '' }) + .mockResolvedValueOnce({ exitCode: 0, stdout: '', stderr: '' }) + .mockResolvedValueOnce({ exitCode: 0, stdout: '', stderr: '' }) .mockResolvedValueOnce({ exitCode: 0, - stdout: '/run/user/1000/docker.sock', + stdout: '3145728000 10485760000\n', stderr: '', - }) // resolve Docker socket - .mockResolvedValueOnce({ exitCode: 0, stdout: '', stderr: '' }) // docker ps (listWrapperContainers) - .mockResolvedValueOnce({ exitCode: 0, stdout: `${oldMtime}\n`, stderr: '' }) // stat agent_stale-1111 - .mockResolvedValueOnce({ exitCode: 0, stdout: '', stderr: '' }) // rm workspace - .mockResolvedValueOnce({ exitCode: 0, stdout: '', stderr: '' }); // rm home - + }); mockListProcesses.mockResolvedValue([]); - await checkDiskAndCleanBeforeSetup(fakeSandbox, 'org-1', 'user-1', 'agent_current-aaaa'); + await expect( + checkDiskAndCleanBeforeSetup(fakeSandbox, 'org-1', 'user-1', 'agent_current-aaaa') + ).resolves.toEqual({ + availableMB: 3000, + thresholdMB: LOW_DISK_THRESHOLD_MB, + cleanup: { cleaned: 1, skipped: 1 }, + }); - // df was called - expect(mockSandboxExec.mock.calls[0][0]).toContain('df -B1'); - // ls was called to find sessions - expect(mockSandboxExec.mock.calls[1][0]).toContain('ls -1'); - // Docker socket was resolved before docker ps. - expect(mockSandboxExec.mock.calls[2][0]).toContain('/run/user/1000/docker.sock'); - // docker ps was called to find devcontainer-launched wrappers - expect(mockSandboxExec.mock.calls[3][0]).toContain('docker ps'); - // stat was called for the stale session - expect(mockSandboxExec.mock.calls[4][0]).toContain('stat'); - // stale session was cleaned - expect(mockSandboxExec.mock.calls[5][0]).toContain('agent_stale-1111'); + expect(mockSandboxExec.mock.calls[7][0]).toContain('df -B1'); }); - it('skips cleanup when disk space is adequate', async () => { - // df returns adequate disk (5000 MB avail, 10000 MB total) - mockSandboxExec.mockResolvedValueOnce({ - exitCode: 0, - stdout: '5242880000 10485760000\n', - stderr: '', - }); + it('types ENOSPC during stale cleanup as sandbox unusable instead of rechecking capacity', async () => { + mockSandboxExec + .mockResolvedValueOnce({ + exitCode: 0, + stdout: '1073741824 10485760000\n', + stderr: '', + }) + .mockRejectedValueOnce(new Error('ENOSPC: no space left on device')); - await checkDiskAndCleanBeforeSetup(fakeSandbox, 'org-1', 'user-1', 'agent_current-aaaa'); + await expect( + checkDiskAndCleanBeforeSetup(fakeSandbox, 'org-1', 'user-1', 'agent_current-aaaa') + ).rejects.toBeInstanceOf(SandboxCapacityInspectionError); + expect(mockSandboxExec).toHaveBeenCalledTimes(2); + }); - // Only the df call — no cleanup - expect(mockSandboxExec).toHaveBeenCalledTimes(1); - expect(mockListProcesses).not.toHaveBeenCalled(); + it('types ENOSPC returned by stale removal as sandbox unusable', async () => { + const oldMtime = String(Math.floor(Date.now() / 1000) - STALE_DIR_MIN_AGE_SECONDS - 60); + mockSandboxExec + .mockResolvedValueOnce({ exitCode: 0, stdout: '1073741824 10485760000\n', stderr: '' }) + .mockResolvedValueOnce({ exitCode: 0, stdout: 'agent_stale-aaaa\n', stderr: '' }) + .mockResolvedValueOnce({ exitCode: 0, stdout: '/run/user/1000/docker.sock', stderr: '' }) + .mockResolvedValueOnce({ exitCode: 0, stdout: '', stderr: '' }) + .mockResolvedValueOnce({ exitCode: 0, stdout: `${oldMtime}\n`, stderr: '' }) + .mockResolvedValueOnce({ exitCode: 1, stdout: '', stderr: 'No space left on device' }) + .mockResolvedValueOnce({ exitCode: 0, stdout: '', stderr: '' }); + mockListProcesses.mockResolvedValue([]); + + await expect( + checkDiskAndCleanBeforeSetup(fakeSandbox, 'org-1', 'user-1', 'agent_current-aaaa') + ).rejects.toBeInstanceOf(SandboxCapacityInspectionError); }); - it('does not throw when disk check fails', async () => { - mockSandboxExec.mockResolvedValueOnce({ - exitCode: 1, - stdout: '', - stderr: 'df: command not found', - }); + it('allows ordinary cleanup inspection failure only after a safe recheck', async () => { + mockSandboxExec + .mockResolvedValueOnce({ + exitCode: 0, + stdout: '1073741824 10485760000\n', + stderr: '', + }) + .mockRejectedValueOnce(new Error('transient list sessions failure')) + .mockResolvedValueOnce({ + exitCode: 0, + stdout: '3145728000 10485760000\n', + stderr: '', + }); - // Should not throw — logs and continues await expect( checkDiskAndCleanBeforeSetup(fakeSandbox, 'org-1', 'user-1', 'agent_current-aaaa') - ).resolves.toBeUndefined(); + ).resolves.toEqual({ + availableMB: 3000, + thresholdMB: LOW_DISK_THRESHOLD_MB, + cleanup: { cleaned: 0, skipped: 0 }, + }); }); - it('does not throw when cleanup fails', async () => { - // df returns low disk + it('rejects admission when low capacity remains after cleanup', async () => { mockSandboxExec .mockResolvedValueOnce({ exitCode: 0, stdout: '1073741824 10485760000\n', stderr: '', }) - // ls throws - .mockRejectedValueOnce(new Error('exec error')); + .mockResolvedValueOnce({ exitCode: 1, stdout: '', stderr: 'no sessions' }) + .mockResolvedValueOnce({ + exitCode: 0, + stdout: '1572864000 10485760000\n', + stderr: '', + }); + + const result = checkDiskAndCleanBeforeSetup( + fakeSandbox, + 'org-1', + 'user-1', + 'agent_current-aaaa' + ); + + await expect(result).rejects.toBeInstanceOf(WorkspaceCapacityAdmissionRejectedError); + await expect(result).rejects.toMatchObject({ + availableMB: 1500, + thresholdMB: LOW_DISK_THRESHOLD_MB, + cleaned: 0, + skipped: 0, + }); + }); + + it('continues to protect live sibling wrappers while rejecting unsafe admission', async () => { + const oldMtime = String(Math.floor(Date.now() / 1000) - STALE_DIR_MIN_AGE_SECONDS - 60); + mockSandboxExec + .mockResolvedValueOnce({ exitCode: 0, stdout: '1073741824 10485760000\n', stderr: '' }) + .mockResolvedValueOnce({ exitCode: 0, stdout: 'agent_active-bbbb\n', stderr: '' }) + .mockResolvedValueOnce({ exitCode: 0, stdout: '/run/user/1000/docker.sock', stderr: '' }) + .mockResolvedValueOnce({ exitCode: 0, stdout: '', stderr: '' }) + .mockResolvedValueOnce({ exitCode: 0, stdout: `${oldMtime}\n`, stderr: '' }) + .mockResolvedValueOnce({ exitCode: 0, stdout: '1073741824 10485760000\n', stderr: '' }); + mockListProcesses.mockResolvedValue([ + { + id: '1', + command: 'kilocode-wrapper --agent-session agent_active-bbbb WRAPPER_PORT=5001', + status: 'running', + }, + ]); + + await expect( + checkDiskAndCleanBeforeSetup(fakeSandbox, 'org-1', 'user-1', 'agent_current-aaaa') + ).rejects.toMatchObject({ skipped: 1 }); + expect(mockSandboxExec.mock.calls.every(call => !call[0].includes('rm -rf'))).toBe(true); + }); + + it('classifies ENOSPC disk inspection failures as sandbox unusable', async () => { + mockSandboxExec.mockResolvedValueOnce({ + exitCode: 1, + stdout: '', + stderr: 'cannot create temporary file: No space left on device', + }); + + await expect( + checkDiskAndCleanBeforeSetup(fakeSandbox, 'org-1', 'user-1', 'agent_current-aaaa') + ).rejects.toBeInstanceOf(SandboxCapacityInspectionError); + }); + + it('does not classify temporary-file permission failures as destructive capacity evidence', async () => { + mockSandboxExec.mockResolvedValueOnce({ + exitCode: 1, + stdout: '', + stderr: 'cannot create temporary file: Permission denied', + }); + + await expect( + checkDiskAndCleanBeforeSetup(fakeSandbox, 'org-1', 'user-1', 'agent_current-aaaa') + ).rejects.toBeInstanceOf(WorkspaceCapacityInspectionUnavailableError); + }); + + it('rejects admission without destroying a shared sandbox when disk execution is transiently unavailable', async () => { + mockSandboxExec.mockRejectedValueOnce(new Error('execution session unavailable')); + + await expect( + checkDiskAndCleanBeforeSetup(fakeSandbox, 'org-1', 'user-1', 'agent_current-aaaa') + ).rejects.toBeInstanceOf(WorkspaceCapacityInspectionUnavailableError); + }); + + it('classifies a thrown ENOSPC disk execution failure as sandbox unusable', async () => { + mockSandboxExec.mockRejectedValueOnce(new Error('ENOSPC: no space left on device')); + + await expect( + checkDiskAndCleanBeforeSetup(fakeSandbox, 'org-1', 'user-1', 'agent_current-aaaa') + ).rejects.toBeInstanceOf(SandboxCapacityInspectionError); + }); + + it('rejects admission without infrastructure recovery when a non-ENOSPC disk check returns failure', async () => { + mockSandboxExec.mockResolvedValueOnce({ + exitCode: 1, + stdout: '', + stderr: 'df: command not found', + }); - // Should not throw — cleanupStaleWorkspaces catches internally await expect( checkDiskAndCleanBeforeSetup(fakeSandbox, 'org-1', 'user-1', 'agent_current-aaaa') - ).resolves.toBeUndefined(); + ).rejects.toBeInstanceOf(WorkspaceCapacityInspectionUnavailableError); }); }); }); diff --git a/services/cloud-agent-next/src/workspace.ts b/services/cloud-agent-next/src/workspace.ts index c396b3713e..eb851c30e5 100644 --- a/services/cloud-agent-next/src/workspace.ts +++ b/services/cloud-agent-next/src/workspace.ts @@ -2,8 +2,8 @@ import type { SandboxInstance, ExecutionSession, SystemSandboxUsageEvent } from import type { ExecResult, ExecOptions } from '@cloudflare/sandbox'; import { logger } from './logger.js'; import { + inspectWrapperContainers, isWrapperLiveInProcessesOrContainers, - listWrapperContainers, } from './kilo/wrapper-manager.js'; import { DISK_CHECK_TIMEOUT_MS, @@ -16,7 +16,13 @@ import { } from './sandbox-timeout-logging.js'; import { withTimeout } from '@kilocode/worker-utils'; import { isSandboxInternalServerError } from './sandbox-recovery.js'; -import { WorkspaceFilesystemPreparationError } from './workspace-errors.js'; +import { + isSandboxFilesystemUnusableError, + SandboxCapacityInspectionError, + WorkspaceCapacityAdmissionRejectedError, + WorkspaceCapacityInspectionUnavailableError, + WorkspaceFilesystemPreparationError, +} from './workspace-errors.js'; /** * Minimal interface for running shell commands. @@ -179,37 +185,112 @@ export type SessionPaths = { sessionHome: string; }; +export type StaleWorkspaceCleanupResult = { + cleaned: number; + skipped: number; +}; + +export type WorkspaceAdmissionResult = { + availableMB: number; + thresholdMB: number; + cleanup: StaleWorkspaceCleanupResult; +}; + +function throwIfSandboxFilesystemUnusable(operation: string, error: unknown): void { + if (!isSandboxFilesystemUnusableError(error)) return; + throw new SandboxCapacityInspectionError( + `${operation} cannot run because the sandbox filesystem is unusable`, + error + ); +} + /** - * Check disk space and clean up stale workspaces if low, using the sandbox - * directly so it can run before any session or workspace directory exists. - * Errors are caught and logged so cleanup failure never blocks setup, except - * sandbox 500s which indicate a bad container that should be destroyed. + * Admit cold workspace setup only when measured capacity is safe, preserving + * conservative stale cleanup for live sibling wrappers on shared sandboxes. */ export async function checkDiskAndCleanBeforeSetup( sandbox: SandboxInstance, orgId: string | undefined, userId: string, sessionId: string -): Promise { +): Promise { try { - const diskSpace = await checkDiskSpace(sandbox); - if (diskSpace.isLow) { - logger.info('Low disk space detected before workspace setup, cleaning stale workspaces'); - await cleanupStaleWorkspaces(sandbox, getBaseWorkspacePath(orgId, userId), sessionId); + const initialCapacity = await checkDiskSpace(sandbox); + if (!initialCapacity.isLow) { + const admitted = { + availableMB: initialCapacity.availableMB, + thresholdMB: LOW_DISK_THRESHOLD_MB, + cleanup: { cleaned: 0, skipped: 0 }, + } satisfies WorkspaceAdmissionResult; + logger.withFields(admitted).info('Workspace capacity admission accepted'); + return admitted; } + + const cleanup = await cleanupStaleWorkspaces( + sandbox, + getBaseWorkspacePath(orgId, userId), + sessionId + ); + const recheckedCapacity = await checkDiskSpace(sandbox); + if (recheckedCapacity.isLow) { + const rejection = new WorkspaceCapacityAdmissionRejectedError({ + availableMB: recheckedCapacity.availableMB, + thresholdMB: LOW_DISK_THRESHOLD_MB, + cleaned: cleanup.cleaned, + skipped: cleanup.skipped, + }); + logger + .withFields({ + availableMB: rejection.availableMB, + thresholdMB: rejection.thresholdMB, + cleaned: rejection.cleaned, + skipped: rejection.skipped, + reason: 'low_capacity_after_cleanup', + }) + .warn('Workspace capacity admission rejected'); + throw rejection; + } + + const admitted = { + availableMB: recheckedCapacity.availableMB, + thresholdMB: LOW_DISK_THRESHOLD_MB, + cleanup, + } satisfies WorkspaceAdmissionResult; + logger.withFields(admitted).info('Workspace capacity admission accepted after cleanup'); + return admitted; } catch (error) { + if ( + error instanceof WorkspaceCapacityAdmissionRejectedError || + error instanceof SandboxCapacityInspectionError + ) { + throw error; + } if (isSandboxInternalServerError(error)) { logger .withFields({ error: error instanceof Error ? error.message : String(error) }) .error('Pre-setup disk check hit sandbox 500, aborting workspace setup'); throw error; } + const errorMessage = error instanceof Error ? error.message : String(error); + if (isSandboxFilesystemUnusableError(error)) { + const unusableError = new SandboxCapacityInspectionError( + 'Disk capacity inspection cannot run because the sandbox filesystem is unusable', + error + ); + logger + .withFields({ error: errorMessage, reason: 'sandbox_filesystem_unusable' }) + .error('Workspace capacity inspection cannot safely admit setup'); + throw unusableError; + } - // Log and continue — a failed disk check should not block workspace setup. - // The worst case is that mkdir fails (which it would have anyway without cleanup). + const unavailableError = new WorkspaceCapacityInspectionUnavailableError( + 'Workspace admission rejected because disk capacity could not be measured', + error + ); logger - .withFields({ error: error instanceof Error ? error.message : String(error) }) - .warn('Pre-setup disk check failed, continuing with workspace setup'); + .withFields({ error: errorMessage, reason: 'capacity_inspection_unavailable' }) + .warn('Workspace capacity admission rejected without sandbox recovery'); + throw unavailableError; } } @@ -261,11 +342,23 @@ export async function cleanupWorkspace( workspacePath: string, sessionHome: string ): Promise { + await cleanupWorkspaceBestEffort(executor, workspacePath, sessionHome); +} + +type WorkspaceCleanupResult = { + cleaned: boolean; + filesystemUnusableCause?: unknown; +}; + +async function cleanupWorkspaceBestEffort( + executor: CommandExecutor, + workspacePath: string, + sessionHome: string +): Promise { logger.setTags({ workspacePath, sessionHome }); logger.info('Cleaning up workspace directories'); try { - // Delete workspace directory const workspaceResult = await timedExec( executor, `rm -rf '${workspacePath}'`, @@ -277,7 +370,6 @@ export async function cleanupWorkspace( .warn('Failed to delete workspace directory'); } - // Delete session home directory const homeResult = await timedExec( executor, `rm -rf '${sessionHome}'`, @@ -289,25 +381,33 @@ export async function cleanupWorkspace( .warn('Failed to delete session home directory'); } - logger.info('Workspace cleanup completed'); + const cleaned = workspaceResult.exitCode === 0 && homeResult.exitCode === 0; + if (cleaned) logger.info('Workspace cleanup completed'); + const unusableResult = [workspaceResult, homeResult].find(result => + isSandboxFilesystemUnusableError(result.stderr) + ); + return unusableResult + ? { cleaned: false, filesystemUnusableCause: new Error(unusableResult.stderr) } + : { cleaned }; } catch (error) { logger .withFields({ error: error instanceof Error ? error.message : String(error) }) .warn('Workspace cleanup encountered an error'); - // Don't throw - cleanup failures shouldn't block session termination + return isSandboxFilesystemUnusableError(error) + ? { cleaned: false, filesystemUnusableCause: error } + : { cleaned: false }; } } /** * Clean up workspace directories for sessions that no longer have a running wrapper. - * Called when disk space is low to reclaim space from abandoned sessions. - * Errors are caught and logged — never rethrown — so cleanup failure never blocks setup. + * Candidate inspection fails closed: unknown or live sessions are counted as skipped. */ export async function cleanupStaleWorkspaces( sandbox: SandboxInstance, baseWorkspacePath: string, currentSessionId: string -): Promise { +): Promise { logger .withFields({ baseWorkspacePath, currentSessionId }) .info('Starting stale workspace cleanup'); @@ -321,9 +421,9 @@ export async function cleanupStaleWorkspaces( ); if (lsResult.exitCode !== 0 || !lsResult.stdout) { logger - .withFields({ stderr: lsResult.stderr }) + .withFields({ stderr: lsResult.stderr, cleaned: 0, skipped: 0 }) .info('No sessions directory or listing failed, skipping cleanup'); - return; + return { cleaned: 0, skipped: 0 }; } sessionDirs = lsResult.stdout .trim() @@ -331,10 +431,15 @@ export async function cleanupStaleWorkspaces( .map(d => d.trim()) .filter(d => /^agent_[\w-]+$/.test(d)); } catch (error) { + throwIfSandboxFilesystemUnusable('Stale workspace discovery', error); logger - .withFields({ error: error instanceof Error ? error.message : String(error) }) + .withFields({ + error: error instanceof Error ? error.message : String(error), + cleaned: 0, + skipped: 0, + }) .warn('Failed to list sessions directory, skipping cleanup'); - return; + return { cleaned: 0, skipped: 0 }; } logger.withFields({ found: sessionDirs.length }).info('Found session directories'); @@ -344,15 +449,33 @@ export async function cleanupStaleWorkspaces( try { processes = await sandbox.listProcesses(); } catch (error) { + throwIfSandboxFilesystemUnusable('Stale wrapper process inspection', error); logger - .withFields({ error: error instanceof Error ? error.message : String(error) }) + .withFields({ + error: error instanceof Error ? error.message : String(error), + cleaned: 0, + skipped: sessionDirs.length, + }) .warn('Failed to list processes, skipping cleanup'); - return; + return { cleaned: 0, skipped: sessionDirs.length }; } - // Also fetch wrapper containers (devcontainer flow). Best-effort — on the - // non-DIND outer image `docker ps` simply returns empty, which is fine. - const wrapperContainers = await listWrapperContainers(sandbox); + const wrapperContainerInspection = await inspectWrapperContainers(sandbox); + if (wrapperContainerInspection.status === 'inspection-failed') { + throwIfSandboxFilesystemUnusable( + 'Stale devcontainer wrapper inspection', + wrapperContainerInspection.error + ); + logger + .withFields({ + error: wrapperContainerInspection.error, + cleaned: 0, + skipped: sessionDirs.length, + }) + .warn('Failed to inspect devcontainer wrappers, skipping cleanup'); + return { cleaned: 0, skipped: sessionDirs.length }; + } + const wrapperContainers = wrapperContainerInspection.containers; // Get current epoch once so we can age-check directories without re-shelling per candidate const nowSeconds = Math.floor(Date.now() / 1000); @@ -413,9 +536,19 @@ export async function cleanupStaleWorkspaces( .withFields({ candidateSessionId, workspacePath, sessionHome }) .info('Removing stale session directories'); - await cleanupWorkspace(sandbox, workspacePath, sessionHome); - cleaned++; + const cleanup = await cleanupWorkspaceBestEffort(sandbox, workspacePath, sessionHome); + if (cleanup.filesystemUnusableCause !== undefined) { + throwIfSandboxFilesystemUnusable( + 'Stale workspace cleanup', + cleanup.filesystemUnusableCause + ); + } + if (cleanup.cleaned) cleaned++; + else skipped++; } catch (error) { + if (error instanceof SandboxCapacityInspectionError) throw error; + throwIfSandboxFilesystemUnusable('Stale workspace cleanup', error); + skipped++; logger .withFields({ candidateSessionId, @@ -425,7 +558,9 @@ export async function cleanupStaleWorkspaces( } } - logger.withFields({ cleaned, skipped }).info('Stale workspace cleanup complete'); + const result = { cleaned, skipped } satisfies StaleWorkspaceCleanupResult; + logger.withFields(result).info('Stale workspace cleanup complete'); + return result; } export type GitAuthorConfig = { @@ -458,18 +593,33 @@ export async function checkDiskSpace(executor: CommandExecutor): Promise candidate.id === sandbox.id) && sandboxesAfter.every(candidate => sandboxIdsBeforeFollowup.has(candidate.id)); - if (!hotResult.terminal || !noPrepare || !sameContainers) { + if (!completed || !noPrepare || !sameContainers) { return { name: 'cold-hot', conversation, diff --git a/services/cloud-agent-next/test/integration/session/callback-outbox.test.ts b/services/cloud-agent-next/test/integration/session/callback-outbox.test.ts index 584ba1e44c..5a31f3e031 100644 --- a/services/cloud-agent-next/test/integration/session/callback-outbox.test.ts +++ b/services/cloud-agent-next/test/integration/session/callback-outbox.test.ts @@ -45,7 +45,13 @@ function removeCallbackQueue(instance: CloudAgentSession): void { describe('callback outbox — missing target or queue', () => { beforeEach(async () => { const ids = await listDurableObjectIds(env.CLOUD_AGENT_SESSION); - await Promise.all(ids.map(id => env.CLOUD_AGENT_SESSION.get(id).deleteSession())); + await Promise.all( + ids.map(id => + runInDurableObject(env.CLOUD_AGENT_SESSION.get(id), instance => + instance.ctx.storage.deleteAll() + ) + ) + ); }); it('records callbackLastError when callbackTarget is missing on a callback-required message', async () => { diff --git a/services/cloud-agent-next/test/integration/session/deletion-lifecycle.test.ts b/services/cloud-agent-next/test/integration/session/deletion-lifecycle.test.ts new file mode 100644 index 0000000000..d3440b4a32 --- /dev/null +++ b/services/cloud-agent-next/test/integration/session/deletion-lifecycle.test.ts @@ -0,0 +1,291 @@ +import { env, runInDurableObject } from 'cloudflare:test'; +import { describe, expect, it } from 'vitest'; +import type { CloudAgentSession } from '../../../src/persistence/CloudAgentSession.js'; +import { listPendingSessionMessages } from '../../../src/session/pending-messages.js'; +import { getWrapperLease } from '../../../src/session/wrapper-runtime-state.js'; +import { + groupedRegisterSessionInput, + queueUserMessageInput, + registerReadySession, +} from '../../helpers/session-setup.js'; + +async function establishOwnedWrapper( + instance: CloudAgentSession, + stopStatus: 'still-present' | 'absent' = 'still-present' +): Promise { + await instance.ctx.storage.put('wrapper_lease', { + state: 'owns_wrapper', + nextInstanceGeneration: 2, + instance: { instanceId: 'instance_delete', instanceGeneration: 1 }, + }); + instance['physicalWrapperStopper'] = async () => + stopStatus === 'absent' ? { status: 'absent' } : { status: 'still-present', observed: [] }; +} + +describe('session deletion physical cleanup', () => { + it('erases Durable Object state after explicit deletion confirms wrapper absence', async () => { + const userId = 'user_delete_complete'; + const sessionId = 'agent_delete_complete'; + const stub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`) + ); + + const result = await runInDurableObject(stub, async instance => { + await registerReadySession(instance, { + sessionId, + userId, + prompt: 'delete complete', + mode: 'code', + model: 'test-model', + }); + await establishOwnedWrapper(instance, 'absent'); + await instance.ctx.storage.put('deletion_marker', true); + + await instance.deleteSession(); + return { + metadata: await instance.getMetadata(), + marker: await instance.ctx.storage.get('deletion_marker'), + alarm: await instance.ctx.storage.getAlarm(), + lease: await getWrapperLease(instance.ctx.storage), + }; + }); + + expect(result).toEqual({ + metadata: null, + marker: undefined, + alarm: null, + lease: { state: 'none', nextInstanceGeneration: 1 }, + }); + }); + + it('does not erase Durable Object state when explicit deletion cannot confirm wrapper absence', async () => { + const userId = 'user_delete_pending'; + const sessionId = 'agent_delete_pending'; + const stub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`) + ); + + const result = await runInDurableObject(stub, async instance => { + await registerReadySession(instance, { + sessionId, + userId, + prompt: 'delete pending', + mode: 'code', + model: 'test-model', + }); + await establishOwnedWrapper(instance); + + await expect(instance.deleteSession()).rejects.toThrow( + 'Session deletion pending physical wrapper cleanup' + ); + const result = { + metadata: await instance.getMetadata(), + lease: await getWrapperLease(instance.ctx.storage), + }; + await instance.ctx.storage.deleteAll(); + return result; + }); + + expect(result.metadata?.identity.sessionId).toBe(sessionId); + expect(result.lease).toMatchObject({ + state: 'stop_needed', + reason: 'session-delete', + attempts: 1, + }); + }); + + it('rejects new message admission while explicit deletion is pending', async () => { + const userId = 'user_delete_reject_admission'; + const sessionId = 'agent_delete_reject_admission'; + const stub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`) + ); + + const result = await runInDurableObject(stub, async instance => { + await registerReadySession(instance, { + sessionId, + userId, + prompt: 'delete reject admission', + mode: 'code', + model: 'test-model', + }); + await establishOwnedWrapper(instance); + + await expect(instance.deleteSession()).rejects.toThrow( + 'Session deletion pending physical wrapper cleanup' + ); + return { + submittedAdmission: await instance.admitSubmittedMessage( + queueUserMessageInput({ userId, prompt: 'must not queue after delete' }) + ), + groupedInitialAdmission: await instance.createSessionWithInitialAdmission({ + ...groupedRegisterSessionInput({ + sessionId, + userId, + prompt: 'must not replay grouped initial admission after delete', + mode: 'code', + model: 'test-model', + }), + message: { + initialTurn: { + type: 'prompt', + messageId: 'msg_018f1e2d3c4bDeleteGroupABC', + prompt: 'must not replay grouped initial admission after delete', + }, + }, + }), + preparedInitialAdmission: await instance.admitPreparedInitialMessage({ userId }), + registration: await instance.registerSession( + groupedRegisterSessionInput({ + sessionId, + userId, + prompt: 'must not register after delete', + mode: 'code', + model: 'test-model', + }) + ), + pendingMessages: await listPendingSessionMessages(instance.ctx.storage), + }; + }); + + const deletionPending = { + success: false, + code: 'NOT_FOUND', + error: 'Session deletion is pending', + }; + expect(result).toEqual({ + submittedAdmission: deletionPending, + groupedInitialAdmission: deletionPending, + preparedInitialAdmission: deletionPending, + registration: { success: false, error: 'Session deletion is pending' }, + pendingMessages: [], + }); + }); + + it('finishes pending explicit deletion from an alarm after wrapper absence is confirmed', async () => { + const userId = 'user_delete_alarm_complete'; + const sessionId = 'agent_delete_alarm_complete'; + const stub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`) + ); + + const result = await runInDurableObject(stub, async instance => { + await registerReadySession(instance, { + sessionId, + userId, + prompt: 'delete alarm complete', + mode: 'code', + model: 'test-model', + }); + await establishOwnedWrapper(instance); + await instance.ctx.storage.put('deletion_marker', true); + let stopAttempts = 0; + instance['physicalWrapperStopper'] = async () => { + stopAttempts += 1; + return stopAttempts === 1 + ? { status: 'still-present', observed: [] } + : { status: 'absent' }; + }; + + await expect(instance.deleteSession()).rejects.toThrow( + 'Session deletion pending physical wrapper cleanup' + ); + const pendingLease = await getWrapperLease(instance.ctx.storage); + if (pendingLease.state !== 'stop_needed') throw new Error('Expected pending wrapper cleanup'); + await instance.ctx.storage.put('wrapper_lease', { + ...pendingLease, + nextAttemptAt: Date.now() - 1, + }); + + await instance.alarm(); + return { + metadata: await instance.getMetadata(), + marker: await instance.ctx.storage.get('deletion_marker'), + alarm: await instance.ctx.storage.getAlarm(), + lease: await getWrapperLease(instance.ctx.storage), + }; + }); + + expect(result).toEqual({ + metadata: null, + marker: undefined, + alarm: null, + lease: { state: 'none', nextInstanceGeneration: 1 }, + }); + }); + + it('postpones retention deletion while physical wrapper cleanup remains unresolved', async () => { + const userId = 'user_ttl_pending'; + const sessionId = 'agent_ttl_pending'; + const stub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`) + ); + + const result = await runInDurableObject(stub, async instance => { + await registerReadySession(instance, { + sessionId, + userId, + prompt: 'ttl pending', + mode: 'code', + model: 'test-model', + }); + await establishOwnedWrapper(instance); + await instance.ctx.storage.put('last_activity', 1); + + await instance.alarm(); + const result = { + metadata: await instance.getMetadata(), + lease: await getWrapperLease(instance.ctx.storage), + }; + await instance.ctx.storage.deleteAll(); + return result; + }); + + expect(result.metadata?.identity.sessionId).toBe(sessionId); + expect(result.lease).toMatchObject({ + state: 'stop_needed', + reason: 'session-delete', + }); + }); + + it('retains physical cleanup backoff while retention deletion is already pending', async () => { + const userId = 'user_ttl_backoff'; + const sessionId = 'agent_ttl_backoff'; + const stub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`) + ); + + const result = await runInDurableObject(stub, async instance => { + await registerReadySession(instance, { + sessionId, + userId, + prompt: 'ttl backoff', + mode: 'code', + model: 'test-model', + }); + const nextAttemptAt = Date.now() + 30_000; + await instance.ctx.storage.put('wrapper_lease', { + state: 'stop_needed', + nextInstanceGeneration: 2, + target: { kind: 'session' }, + reason: 'session-delete', + requestedAt: Date.now(), + nextAttemptAt, + attempts: 1, + }); + await instance.ctx.storage.put('last_activity', 1); + + await instance.alarm(); + const result = { + lease: await getWrapperLease(instance.ctx.storage), + alarm: await instance.ctx.storage.getAlarm(), + nextAttemptAt, + }; + await instance.ctx.storage.deleteAll(); + return result; + }); + + expect(result.lease).toMatchObject({ state: 'stop_needed', attempts: 1 }); + expect(result.alarm).toBe(result.nextAttemptAt); + }); +}); diff --git a/services/cloud-agent-next/test/integration/session/disconnect-and-reaper.test.ts b/services/cloud-agent-next/test/integration/session/disconnect-and-reaper.test.ts index e7323af325..67591f3999 100644 --- a/services/cloud-agent-next/test/integration/session/disconnect-and-reaper.test.ts +++ b/services/cloud-agent-next/test/integration/session/disconnect-and-reaper.test.ts @@ -4,6 +4,10 @@ import { drizzle } from 'drizzle-orm/durable-sqlite'; import { createEventQueries } from '../../../src/session/queries/events.js'; import { storePendingSessionMessage } from '../../../src/session/pending-messages.js'; import { putSessionMessageState } from '../../../src/session/session-message-state.js'; +import { + getWrapperLease, + getWrapperRuntimeState, +} from '../../../src/session/wrapper-runtime-state.js'; import type { ExecutionId } from '../../../src/types/ids.js'; describe('Disconnect handling and compatibility execution RPCs', () => { @@ -81,6 +85,56 @@ describe('Disconnect handling and compatibility execution RPCs', () => { expect(result.keptWithPending).toBeDefined(); }); + it('idle cleanup preserves a completed wrapper retained for warm fenced reuse', async () => { + const userId = 'user_idle_warm_reuse'; + const sessionId = 'agent_idle_warm_reuse'; + const stub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`) + ); + + const result = await runInDurableObject(stub, async (instance, state) => { + const now = Date.now(); + const expiredActivity = now - 20 * 60 * 1000; + await instance.updateMetadata({ + version: now, + sessionId, + userId, + timestamp: now, + kiloServerLastActivity: expiredActivity, + }); + await state.storage.put('wrapper_runtime_state', { + wrapperRunId: 'wr_idle_warm_reuse', + wrapperGeneration: 1, + wrapperConnectionId: 'connection-completed', + lastWrapperMessageAt: now, + }); + await state.storage.put('wrapper_lease', { + state: 'owns_wrapper', + nextInstanceGeneration: 2, + instance: { instanceId: 'instance_idle_warm_reuse', instanceGeneration: 1 }, + }); + + await instance.handleWrapperTerminalEvent({ + wrapperRunId: 'wr_idle_warm_reuse', + status: 'completed', + }); + await instance.alarm(); + + return { + lease: await getWrapperLease(state.storage), + runtime: await getWrapperRuntimeState(state.storage), + activity: (await instance.getMetadata())?.lifecycle.kiloServerLastActivity, + }; + }); + + expect(result.lease).toMatchObject({ + state: 'owns_wrapper', + keepWarmUntil: expect.any(Number), + }); + expect(result.runtime.wrapperConnectionId).toBeUndefined(); + expect(result.activity).toBeDefined(); + }); + it('failExecutionRpc retains the public execution-record failure contract', async () => { const userId = 'user_rpc_failure'; const sessionId = 'agent_rpc_failure'; diff --git a/services/cloud-agent-next/test/integration/session/execute-directly-failure.test.ts b/services/cloud-agent-next/test/integration/session/execute-directly-failure.test.ts index 62edb12977..bd5f921572 100644 --- a/services/cloud-agent-next/test/integration/session/execute-directly-failure.test.ts +++ b/services/cloud-agent-next/test/integration/session/execute-directly-failure.test.ts @@ -13,6 +13,7 @@ import { createEventQueries } from '../../../src/session/queries/events.js'; import type { FencedWrapperDispatchRequest } from '../../../src/execution/types.js'; import { listPendingSessionMessages } from '../../../src/session/pending-messages.js'; import { + getWrapperLease, getWrapperRuntimeState, recordWrapperPong, allocateWrapperRuntimeState, @@ -28,7 +29,13 @@ import { queueUserMessageInput, registerReadySession } from '../../helpers/sessi describe('executeDirectly failure handling', () => { beforeEach(async () => { const ids = await listDurableObjectIds(env.CLOUD_AGENT_SESSION); - await Promise.all(ids.map(id => env.CLOUD_AGENT_SESSION.get(id).deleteSession())); + await Promise.all( + ids.map(id => + runInDurableObject(env.CLOUD_AGENT_SESSION.get(id), instance => + instance.ctx.storage.deleteAll() + ) + ) + ); }); it('wrapper heartbeat does not reset the no-output deadline', async () => { @@ -38,10 +45,6 @@ describe('executeDirectly failure handling', () => { const stub = env.CLOUD_AGENT_SESSION.get(doId); const result = await runInDurableObject(stub, async instance => { - instance['stopCurrentWrapperProcess'] = async () => { - throw new Error('should not stop wrapper'); - }; - await registerReadySession(instance, { sessionId, userId, @@ -200,7 +203,10 @@ describe('executeDirectly failure handling', () => { const pendingAfterAlarm = await listPendingSessionMessages(instance.ctx.storage); const executionsAfterFirstAlarm = await instance.getExecutions(); const wrapperRuntimeState = await getWrapperRuntimeState(instance.ctx.storage); + const wrapperLeaseAfterFailure = await getWrapperLease(instance.ctx.storage); + await instance.alarm(); + const wrapperLeaseAfterCleanup = await getWrapperLease(instance.ctx.storage); const retriableMessage = pendingAfterAlarm[0]; if (retriableMessage) { await instance.ctx.storage.put('pending_message:0000000000000001:retry-fix', { @@ -239,6 +245,8 @@ describe('executeDirectly failure handling', () => { executionsAfterRetry, acceptedMessages, wrapperRuntimeState, + wrapperLeaseAfterFailure, + wrapperLeaseAfterCleanup, retryEvents, }; }); @@ -258,6 +266,11 @@ describe('executeDirectly failure handling', () => { expect(result.executionsAfterFirstAlarm).toEqual([]); expect(result.wrapperRuntimeState.wrapperGeneration).toBe(2); expect(result.wrapperRuntimeState.wrapperConnectionId).toBeUndefined(); + expect(result.wrapperLeaseAfterFailure).toMatchObject({ + state: 'stop_needed', + reason: 'startup-failed', + }); + expect(result.wrapperLeaseAfterCleanup).toMatchObject({ state: 'none' }); expect(result.attemptCount).toBe(2); expect(result.pendingAfterRetry).toHaveLength(0); @@ -272,7 +285,13 @@ describe('executeDirectly failure handling', () => { describe('handleWrapperTerminalEvent — new-path identity and message preservation', () => { beforeEach(async () => { const ids = await listDurableObjectIds(env.CLOUD_AGENT_SESSION); - await Promise.all(ids.map(id => env.CLOUD_AGENT_SESSION.get(id).deleteSession())); + await Promise.all( + ids.map(id => + runInDurableObject(env.CLOUD_AGENT_SESSION.get(id), instance => + instance.ctx.storage.deleteAll() + ) + ) + ); }); it('wrapper complete does not clear wrapper runtime identity when accepted messages remain', async () => { @@ -337,7 +356,13 @@ describe('handleWrapperTerminalEvent — new-path identity and message preservat describe('new-path liveness without executionId', () => { beforeEach(async () => { const ids = await listDurableObjectIds(env.CLOUD_AGENT_SESSION); - await Promise.all(ids.map(id => env.CLOUD_AGENT_SESSION.get(id).deleteSession())); + await Promise.all( + ids.map(id => + runInDurableObject(env.CLOUD_AGENT_SESSION.get(id), instance => + instance.ctx.storage.deleteAll() + ) + ) + ); }); it('schedules liveness deadlines for accepted messages and fails them on no-output timeout', async () => { @@ -347,8 +372,6 @@ describe('new-path liveness without executionId', () => { const stub = env.CLOUD_AGENT_SESSION.get(doId); const result = await runInDurableObject(stub, async (instance, state) => { - instance['stopCurrentWrapperProcess'] = async () => {}; - await registerReadySession(instance, { sessionId, userId, @@ -431,8 +454,6 @@ describe('new-path liveness without executionId', () => { const stub = env.CLOUD_AGENT_SESSION.get(doId); const result = await runInDurableObject(stub, async (instance, state) => { - instance['stopCurrentWrapperProcess'] = async () => {}; - await registerReadySession(instance, { sessionId, userId, @@ -505,7 +526,13 @@ describe('new-path liveness without executionId', () => { describe('hot delivery failure preserves existing wrapper identity', () => { beforeEach(async () => { const ids = await listDurableObjectIds(env.CLOUD_AGENT_SESSION); - await Promise.all(ids.map(id => env.CLOUD_AGENT_SESSION.get(id).deleteSession())); + await Promise.all( + ids.map(id => + runInDurableObject(env.CLOUD_AGENT_SESSION.get(id), instance => + instance.ctx.storage.deleteAll() + ) + ) + ); }); it('failed hot delivery does not clear wrapper identity for already accepted work', async () => { @@ -538,6 +565,23 @@ describe('hot delivery failure preserves existing wrapper identity', () => { const originalRunId = wrapperState.wrapperRunId!; const originalConnectionId = wrapperState.wrapperConnectionId!; const originalGeneration = wrapperState.wrapperGeneration; + await instance.ctx.storage.put('wrapper_lease', { + state: 'owns_wrapper', + nextInstanceGeneration: 2, + instance: { instanceId: 'instance_hot_failure', instanceGeneration: 1 }, + }); + instance['physicalWrapperObserver'] = async () => ({ + status: 'present', + observed: [ + { + representation: 'process', + id: 'wrapper-hot-failure', + port: 5000, + instanceId: 'instance_hot_failure', + instanceGeneration: 1, + }, + ], + }); const acceptedMsg: SessionMessageState = { messageId: 'msg_018f1e2d3c4bHotFailAccAbCd', @@ -594,7 +638,7 @@ describe('hot delivery failure preserves existing wrapper identity', () => { expect(result.acceptedMessages[0]?.status).toBe('accepted'); }); - it('failed cold delivery clears newly allocated wrapper identity', async () => { + it('failed cold delivery fences its run and retains physical cleanup responsibility', async () => { const userId = 'user_cold_fail_identity'; const sessionId = 'agent_cold_fail_identity'; const doId = env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`); @@ -633,10 +677,12 @@ describe('hot delivery failure preserves existing wrapper identity', () => { await instance.alarm(); const wrapperRuntimeState = await getWrapperRuntimeState(instance.ctx.storage); + const wrapperLease = await getWrapperLease(instance.ctx.storage); return { preAlarmState, wrapperRuntimeState, + wrapperLease, }; }); @@ -645,5 +691,9 @@ describe('hot delivery failure preserves existing wrapper identity', () => { expect(result.wrapperRuntimeState.wrapperGeneration).toBeGreaterThan( result.preAlarmState.wrapperGeneration ); + expect(result.wrapperLease).toMatchObject({ + state: 'stop_needed', + reason: 'startup-failed', + }); }); }); diff --git a/services/cloud-agent-next/test/integration/session/execution-id-removal.test.ts b/services/cloud-agent-next/test/integration/session/execution-id-removal.test.ts index 7596e75ea6..e18b97d5af 100644 --- a/services/cloud-agent-next/test/integration/session/execution-id-removal.test.ts +++ b/services/cloud-agent-next/test/integration/session/execution-id-removal.test.ts @@ -29,7 +29,13 @@ import { queueUserMessageInput, registerReadySession } from '../../helpers/sessi describe('execution-id removal - queue and start response', () => { beforeEach(async () => { const ids = await listDurableObjectIds(env.CLOUD_AGENT_SESSION); - await Promise.all(ids.map(id => env.CLOUD_AGENT_SESSION.get(id).deleteSession())); + await Promise.all( + ids.map(id => + runInDurableObject(env.CLOUD_AGENT_SESSION.get(id), instance => + instance.ctx.storage.deleteAll() + ) + ) + ); }); it('new-path admission result omits executionId', async () => { @@ -118,7 +124,13 @@ describe('execution-id removal - queue and start response', () => { describe('execution-id removal - flush does not create execution rows', () => { beforeEach(async () => { const ids = await listDurableObjectIds(env.CLOUD_AGENT_SESSION); - await Promise.all(ids.map(id => env.CLOUD_AGENT_SESSION.get(id).deleteSession())); + await Promise.all( + ids.map(id => + runInDurableObject(env.CLOUD_AGENT_SESSION.get(id), instance => + instance.ctx.storage.deleteAll() + ) + ) + ); }); it('new-path flush does not insert an execution metadata row', async () => { @@ -385,7 +397,13 @@ describe('execution-id removal - flush does not create execution rows', () => { describe('execution-id removal - stream events do not expose fake executionIds', () => { beforeEach(async () => { const ids = await listDurableObjectIds(env.CLOUD_AGENT_SESSION); - await Promise.all(ids.map(id => env.CLOUD_AGENT_SESSION.get(id).deleteSession())); + await Promise.all( + ids.map(id => + runInDurableObject(env.CLOUD_AGENT_SESSION.get(id), instance => + instance.ctx.storage.deleteAll() + ) + ) + ); }); it('new-path message queued event payload does not contain executionId', async () => { @@ -544,7 +562,13 @@ describe('execution-id removal - stream events do not expose fake executionIds', describe('execution-id removal - ingest does not alias wrapperRunId as execution_id', () => { beforeEach(async () => { const ids = await listDurableObjectIds(env.CLOUD_AGENT_SESSION); - await Promise.all(ids.map(id => env.CLOUD_AGENT_SESSION.get(id).deleteSession())); + await Promise.all( + ids.map(id => + runInDurableObject(env.CLOUD_AGENT_SESSION.get(id), instance => + instance.ctx.storage.deleteAll() + ) + ) + ); }); it('new-path ingest events do not use wrapperRunId as execution_id in StoredEvent', async () => { diff --git a/services/cloud-agent-next/test/integration/session/hot-delivery.test.ts b/services/cloud-agent-next/test/integration/session/hot-delivery.test.ts index 6122fd65ed..0d6d1ec8e1 100644 --- a/services/cloud-agent-next/test/integration/session/hot-delivery.test.ts +++ b/services/cloud-agent-next/test/integration/session/hot-delivery.test.ts @@ -21,7 +21,10 @@ import { putSessionMessageState, type SessionMessageState, } from '../../../src/session/session-message-state.js'; -import { allocateWrapperRuntimeState } from '../../../src/session/wrapper-runtime-state.js'; +import { + allocateWrapperRuntimeState, + recordMeaningfulWrapperOutput, +} from '../../../src/session/wrapper-runtime-state.js'; import type { FencedWrapperDispatchRequest } from '../../../src/execution/types.js'; import { registerReadySession } from '../../helpers/session-setup.js'; @@ -47,6 +50,18 @@ describe('hot delivery — DO integration', () => { return { messageId: plan.turn.messageId, kiloSessionId: 'kilo_hot_test' }; }, }; + (instance as any).physicalWrapperObserver = async () => ({ + status: 'present', + observed: [ + { + representation: 'process', + id: 'wrapper-hot', + port: 4_173, + instanceId: 'instance_hot', + instanceGeneration: 1, + }, + ], + }); await registerReadySession(instance, { sessionId, @@ -61,8 +76,19 @@ describe('hot delivery — DO integration', () => { gitToken: 'git-token', }); - // Simulate a warm wrapper: allocate runtime state with connection + // Simulate a warm, physically owned wrapper with recent output. const { state: wrapperState } = await allocateWrapperRuntimeState(instance.ctx.storage); + await instance.ctx.storage.put('wrapper_lease', { + state: 'owns_wrapper', + nextInstanceGeneration: 2, + instance: { instanceId: 'instance_hot', instanceGeneration: 1 }, + }); + await recordMeaningfulWrapperOutput( + instance.ctx.storage, + wrapperState.wrapperGeneration, + wrapperState.wrapperConnectionId!, + Date.now() + ); // Add an accepted message so hasCurrentWrapper is true const acceptedMsg: SessionMessageState = { diff --git a/services/cloud-agent-next/test/integration/session/idle-reconciliation.test.ts b/services/cloud-agent-next/test/integration/session/idle-reconciliation.test.ts index 74ed211603..14111294a4 100644 --- a/services/cloud-agent-next/test/integration/session/idle-reconciliation.test.ts +++ b/services/cloud-agent-next/test/integration/session/idle-reconciliation.test.ts @@ -12,6 +12,7 @@ import { describe, it, expect, beforeEach } from 'vitest'; import { drizzle } from 'drizzle-orm/durable-sqlite'; import { createEventQueries } from '../../../src/session/queries/events.js'; import { + getWrapperLease, getWrapperRuntimeState, allocateWrapperRuntimeState, } from '../../../src/session/wrapper-runtime-state.js'; @@ -26,7 +27,13 @@ import { registerReadySession } from '../../helpers/session-setup.js'; describe('idle reconciliation scheduling', () => { beforeEach(async () => { const ids = await listDurableObjectIds(env.CLOUD_AGENT_SESSION); - await Promise.all(ids.map(id => env.CLOUD_AGENT_SESSION.get(id).deleteSession())); + await Promise.all( + ids.map(id => + runInDurableObject(env.CLOUD_AGENT_SESSION.get(id), instance => + instance.ctx.storage.deleteAll() + ) + ) + ); }); it('root session.idle records lastWrapperIdleAt and idleReconcileAfter', async () => { @@ -104,8 +111,6 @@ describe('idle reconciliation scheduling', () => { const stub = env.CLOUD_AGENT_SESSION.get(doId); const result = await runInDurableObject(stub, async instance => { - instance['stopCurrentWrapperProcess'] = async () => {}; - await registerReadySession(instance, { sessionId, userId, @@ -160,8 +165,6 @@ describe('idle reconciliation scheduling', () => { const stub = env.CLOUD_AGENT_SESSION.get(doId); const result = await runInDurableObject(stub, async (instance, state) => { - instance['stopCurrentWrapperProcess'] = async () => {}; - await registerReadySession(instance, { sessionId, userId, @@ -172,45 +175,42 @@ describe('idle reconciliation scheduling', () => { model: 'test-model', kilocodeToken: 'token-idle-after', }); - const { state: wrapperState } = await allocateWrapperRuntimeState(instance.ctx.storage); - const { wrapperRunId, wrapperConnectionId } = wrapperState; - - const acceptedMessage: SessionMessageState = { - messageId: 'msg_018f1e2d3c4b00000000000002', + const messageId = 'msg_018f1e2d3c4b00000000000002'; + await putSessionMessageState(instance.ctx.storage, { + messageId, status: 'accepted', prompt: 'hello', createdAt: Date.now(), acceptedAt: Date.now(), - wrapperRunId: wrapperRunId!, - }; - await putSessionMessageState(instance.ctx.storage, acceptedMessage); - + wrapperRunId: wrapperState.wrapperRunId!, + }); const past = Date.now() - 1; await instance.ctx.storage.put('wrapper_runtime_state', { wrapperGeneration: wrapperState.wrapperGeneration, - wrapperConnectionId, - wrapperRunId, + wrapperConnectionId: wrapperState.wrapperConnectionId, + wrapperRunId: wrapperState.wrapperRunId, lastWrapperIdleAt: past - 15_000, idleReconcileAfter: past, }); - await instance.alarm(); const nonTerminalMessages = await listNonTerminalAcceptedMessages( instance.ctx.storage, - wrapperRunId! + wrapperState.wrapperRunId! ); - const failedMessage = await getSessionMessageState( - instance.ctx.storage, - 'msg_018f1e2d3c4b00000000000002' + const failedMessage = await getSessionMessageState(instance.ctx.storage, messageId); + const events = createEventQueries( + drizzle(state.storage, { logger: false }), + state.storage.sql ); - const db = drizzle(state.storage, { logger: false }); - const eventQueries = createEventQueries(db, state.storage.sql); - const allEvents = eventQueries.findByFilters({}); - - return { nonTerminalMessages, failedMessage, allEvents }; + return { + nonTerminalMessages, + failedMessage, + failedEvents: events.findByFilters({ eventTypes: ['cloud.message.failed'] }), + lease: await getWrapperLease(instance.ctx.storage), + }; }); expect(result.nonTerminalMessages).toHaveLength(0); @@ -219,29 +219,26 @@ describe('idle reconciliation scheduling', () => { failureCode: 'missing_assistant_reply', }); - const failedEvents = result.allEvents.filter( - event => event.stream_event_type === 'cloud.message.failed' - ); - expect(failedEvents).toHaveLength(1); - const payload = JSON.parse(failedEvents[0].payload); - expect(payload).toMatchObject({ + expect(result.failedEvents).toHaveLength(1); + expect(JSON.parse(result.failedEvents[0].payload)).toMatchObject({ messageId: 'msg_018f1e2d3c4b00000000000002', status: 'failed', error: 'No assistant reply found after idle timeout', delivery: 'sent', accepted: true, + completionSource: 'idle_reconciliation', }); - expect(payload.completionSource).toBe('idle_reconciliation'); + expect(result.lease).toMatchObject({ state: 'stop_needed', reason: 'terminal-failed' }); }); it('idle reconciliation treats object-shaped assistant errors as failed replies', async () => { const userId = 'user_idle_object_error'; const sessionId = 'agent_idle_object_error'; - const doId = env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`); - const stub = env.CLOUD_AGENT_SESSION.get(doId); + const stub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`) + ); const result = await runInDurableObject(stub, async (instance, state) => { - instance['stopCurrentWrapperProcess'] = async () => {}; await registerReadySession(instance, { sessionId, userId, @@ -252,7 +249,6 @@ describe('idle reconciliation scheduling', () => { model: 'test-model', kilocodeToken: 'token-idle-object-error', }); - const { state: wrapperState } = await allocateWrapperRuntimeState(instance.ctx.storage); const messageId = 'msg_018f1e2d3c4bObjErrIdleAbCd'; await putSessionMessageState(instance.ctx.storage, { @@ -263,7 +259,6 @@ describe('idle reconciliation scheduling', () => { acceptedAt: Date.now(), wrapperRunId: wrapperState.wrapperRunId!, }); - const events = createEventQueries( drizzle(state.storage, { logger: false }), state.storage.sql @@ -287,7 +282,6 @@ describe('idle reconciliation scheduling', () => { timestamp: Date.now(), entityId: 'message/assistant_obj_error_idle', }); - const past = Date.now() - 1; await instance.ctx.storage.put('wrapper_runtime_state', { wrapperGeneration: wrapperState.wrapperGeneration, @@ -296,25 +290,23 @@ describe('idle reconciliation scheduling', () => { lastWrapperIdleAt: past - 15_000, idleReconcileAfter: past, }); - await instance.alarm(); - return events.findByFilters({ - eventTypes: ['cloud.message.failed', 'cloud.message.completed'], - }); + return { + failedEvents: events.findByFilters({ eventTypes: ['cloud.message.failed'] }), + completedEvents: events.findByFilters({ eventTypes: ['cloud.message.completed'] }), + lease: await getWrapperLease(instance.ctx.storage), + }; }); - const failedEvents = result.filter(event => event.stream_event_type === 'cloud.message.failed'); - const completedEvents = result.filter( - event => event.stream_event_type === 'cloud.message.completed' - ); - expect(failedEvents).toHaveLength(1); - expect(completedEvents).toHaveLength(0); - expect(JSON.parse(failedEvents[0].payload)).toMatchObject({ + expect(result.failedEvents).toHaveLength(1); + expect(JSON.parse(result.failedEvents[0].payload)).toMatchObject({ messageId: 'msg_018f1e2d3c4bObjErrIdleAbCd', status: 'failed', error: 'provider failed during idle', completionSource: 'idle_reconciliation', }); + expect(result.completedEvents).toHaveLength(0); + expect(result.lease).toMatchObject({ state: 'stop_needed', reason: 'terminal-failed' }); }); it('meaningful wrapper output clears idle state', async () => { @@ -458,8 +450,9 @@ describe('idle reconciliation scheduling', () => { let stopWrapperCalled = false; const result = await runInDurableObject(stub, async instance => { - instance['stopCurrentWrapperProcess'] = async () => { + instance['physicalWrapperStopper'] = async () => { stopWrapperCalled = true; + return { status: 'absent' }; }; await registerReadySession(instance, { @@ -504,8 +497,9 @@ describe('idle reconciliation scheduling', () => { let stopWrapperCalled = false; const result = await runInDurableObject(stub, async instance => { - instance['stopCurrentWrapperProcess'] = async () => { + instance['physicalWrapperStopper'] = async () => { stopWrapperCalled = true; + return { status: 'absent' }; }; await registerReadySession(instance, { @@ -533,11 +527,16 @@ describe('idle reconciliation scheduling', () => { await instance.alarm(); const runtimeState = await getWrapperRuntimeState(instance.ctx.storage); - return { runtimeState, stopWrapperCalled }; + return { + runtimeState, + lease: await getWrapperLease(instance.ctx.storage), + stopWrapperCalled, + }; }); - expect(result.stopWrapperCalled).toBe(true); + expect(result.stopWrapperCalled).toBe(false); expect(result.runtimeState.wrapperConnectionId).toBeUndefined(); + expect(result.lease).toMatchObject({ state: 'stop_needed', reason: 'keep-warm-expired' }); }); it('keep-warm cleanup clears idle state when work exists after deadline', async () => { @@ -549,8 +548,9 @@ describe('idle reconciliation scheduling', () => { let stopWrapperCalled = false; const result = await runInDurableObject(stub, async instance => { - instance['stopCurrentWrapperProcess'] = async () => { + instance['physicalWrapperStopper'] = async () => { stopWrapperCalled = true; + return { status: 'absent' }; }; await registerReadySession(instance, { diff --git a/services/cloud-agent-next/test/integration/session/message-terminalization.test.ts b/services/cloud-agent-next/test/integration/session/message-terminalization.test.ts index 819897f857..33147e3085 100644 --- a/services/cloud-agent-next/test/integration/session/message-terminalization.test.ts +++ b/services/cloud-agent-next/test/integration/session/message-terminalization.test.ts @@ -106,7 +106,13 @@ async function seedAssistantMessageWithParent( describe('message terminalization and stream events', () => { beforeEach(async () => { const ids = await listDurableObjectIds(env.CLOUD_AGENT_SESSION); - await Promise.all(ids.map(id => env.CLOUD_AGENT_SESSION.get(id).deleteSession())); + await Promise.all( + ids.map(id => + runInDurableObject(env.CLOUD_AGENT_SESSION.get(id), instance => + instance.ctx.storage.deleteAll() + ) + ) + ); }); it('alarm repairs terminal effects without duplicating a durable terminal event', async () => { diff --git a/services/cloud-agent-next/test/integration/session/pending-messages.test.ts b/services/cloud-agent-next/test/integration/session/pending-messages.test.ts index 8d1c9a3b4f..0a84d8ca9e 100644 --- a/services/cloud-agent-next/test/integration/session/pending-messages.test.ts +++ b/services/cloud-agent-next/test/integration/session/pending-messages.test.ts @@ -24,6 +24,7 @@ import { queueUserMessageInput, registerReadySession, } from '../../helpers/session-setup.js'; +import { getWrapperLease } from '../../../src/session/wrapper-runtime-state.js'; const createMessage = (overrides: Partial): PendingSessionMessage => ({ messageId: 'msg_018f1e2d3c4bAbCdEfGhIjKlMn', @@ -836,9 +837,10 @@ describe('pending session messages', () => { } originalEnsure(params); }; - instance['stopCurrentWrapperProcess'] = async () => { - throw new Error('stop failed'); - }; + instance['physicalWrapperStopper'] = async () => ({ + status: 'inspection-failed', + error: 'stop failed', + }); const interrupt = await instance.interruptExecution(); await instance.alarm(); @@ -858,7 +860,59 @@ describe('pending session messages', () => { expect(result.runtime).toEqual({ wrapperGeneration: 2 }); }); - it('interrupt with accepted work and no live socket fences and stops current wrapper runtime', async () => { + it('interrupt with accepted work preserves durable physical cleanup when absence cannot be confirmed', async () => { + const userId = 'user_pending_interrupt_cleanup'; + const sessionId = 'agent_pending_interrupt_cleanup'; + const stub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`) + ); + + const result = await runInDurableObject(stub, async instance => { + await registerReadySession(instance, { + sessionId, + userId, + kiloSessionId: '87878787-8787-4878-8878-878787878787', + prompt: 'prepared prompt', + mode: 'code', + model: 'test-model', + kilocodeToken: 'token-interrupt-cleanup', + }); + await instance.ctx.storage.put('wrapper_lease', { + state: 'owns_wrapper', + nextInstanceGeneration: 2, + instance: { instanceId: 'instance_interrupt', instanceGeneration: 1 }, + }); + await instance.ctx.storage.put('wrapper_runtime_state', { + wrapperGeneration: 1, + wrapperConnectionId: 'conn_interrupt_cleanup', + wrapperRunId: 'wr_interrupt_cleanup', + }); + await putSessionMessageState(instance.ctx.storage, { + messageId: 'msg_018f1e2d3c4bIntrCleanAbCdE', + status: 'accepted', + prompt: 'active message', + createdAt: 1, + acceptedAt: 1, + wrapperRunId: 'wr_interrupt_cleanup', + }); + instance['physicalWrapperStopper'] = async () => ({ + status: 'still-present', + observed: [], + }); + + const interrupt = await instance.interruptExecution(); + return { interrupt, lease: await getWrapperLease(instance.ctx.storage) }; + }); + + expect(result.interrupt).toEqual({ success: true, executionId: undefined }); + expect(result.lease).toMatchObject({ + state: 'stop_needed', + reason: 'user-interrupt', + attempts: 1, + }); + }); + + it('interrupt with accepted work and no live socket fences and requests current wrapper cleanup', async () => { const userId = 'user_pending_interrupt_no_socket'; const sessionId = 'agent_pending_interrupt_no_socket'; const stub = env.CLOUD_AGENT_SESSION.get( @@ -867,9 +921,9 @@ describe('pending session messages', () => { const result = await runInDurableObject(stub, async instance => { const stopped: string[] = []; - instance['stopCurrentWrapperProcess'] = async reason => { - stopped.push(reason); - return true; + instance['physicalWrapperStopper'] = async request => { + stopped.push(request.reason); + return { status: 'absent' }; }; await registerReadySession(instance, { sessionId, @@ -880,6 +934,11 @@ describe('pending session messages', () => { model: 'test-model', kilocodeToken: 'token-followup', }); + await instance.ctx.storage.put('wrapper_lease', { + state: 'owns_wrapper', + nextInstanceGeneration: 2, + instance: { instanceId: 'instance_interrupt_missing', instanceGeneration: 1 }, + }); await instance.ctx.storage.put('wrapper_runtime_state', { wrapperGeneration: 1, wrapperConnectionId: 'conn_interrupt_missing', @@ -898,12 +957,130 @@ describe('pending session messages', () => { const runtimeState = await instance.ctx.storage.get<{ wrapperGeneration: number }>( 'wrapper_runtime_state' ); - return { interrupt, runtimeState, stopped }; + return { + interrupt, + runtimeState, + stopped, + lease: await getWrapperLease(instance.ctx.storage), + }; }); expect(result.interrupt).toEqual({ success: true, executionId: undefined }); expect(result.runtimeState).toEqual({ wrapperGeneration: 2 }); expect(result.stopped).toEqual(['user-interrupt']); + expect(result.lease).toMatchObject({ state: 'none' }); + }); + + it('interrupt with a live fenced socket and no accepted work requests cleanup and fences reuse', async () => { + const userId = 'user_pending_interrupt_idle_wrapper'; + const sessionId = 'agent_pending_interrupt_idle_wrapper'; + const stub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`) + ); + + const result = await runInDurableObject(stub, async instance => { + const sentCommands: unknown[] = []; + const stopped: string[] = []; + instance.sendToWrapper = (_ingestTagId, command, _fence) => { + sentCommands.push(command); + return true; + }; + instance['physicalWrapperStopper'] = async request => { + stopped.push(request.reason); + return { status: 'absent' }; + }; + await registerReadySession(instance, { + sessionId, + userId, + kiloSessionId: '97979797-9797-4979-8979-979797979797', + prompt: 'prepared prompt', + mode: 'code', + model: 'test-model', + kilocodeToken: 'token-followup', + }); + await instance.ctx.storage.put('wrapper_lease', { + state: 'owns_wrapper', + nextInstanceGeneration: 2, + instance: { instanceId: 'instance_interrupt_idle', instanceGeneration: 1 }, + }); + await instance.ctx.storage.put('wrapper_runtime_state', { + wrapperGeneration: 1, + wrapperConnectionId: 'conn_interrupt_idle', + wrapperRunId: 'wr_interrupt_idle', + }); + + const interrupt = await instance.interruptExecution(); + const runtimeState = await instance.ctx.storage.get<{ wrapperGeneration: number }>( + 'wrapper_runtime_state' + ); + return { + interrupt, + runtimeState, + sentCommands, + stopped, + lease: await getWrapperLease(instance.ctx.storage), + }; + }); + + expect(result.interrupt).toEqual({ success: true, executionId: undefined }); + expect(result.sentCommands).toEqual([{ type: 'kill', signal: 'SIGTERM' }]); + expect(result.runtimeState).toEqual({ wrapperGeneration: 2 }); + expect(result.stopped).toEqual(['user-interrupt']); + expect(result.lease).toMatchObject({ state: 'none' }); + }); + + it('interrupt fences a live wrapper when immediate kill signaling fails', async () => { + const userId = 'user_pending_interrupt_signal_failure'; + const sessionId = 'agent_pending_interrupt_signal_failure'; + const stub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`) + ); + + const result = await runInDurableObject(stub, async instance => { + const stopped: string[] = []; + instance.sendToWrapper = () => { + throw new Error('socket send failed'); + }; + instance['physicalWrapperStopper'] = async request => { + stopped.push(request.reason); + return { status: 'absent' }; + }; + await registerReadySession(instance, { + sessionId, + userId, + kiloSessionId: '98989898-9898-4989-8989-989898989898', + prompt: 'prepared prompt', + mode: 'code', + model: 'test-model', + kilocodeToken: 'token-followup', + }); + await instance.ctx.storage.put('wrapper_lease', { + state: 'owns_wrapper', + nextInstanceGeneration: 2, + instance: { instanceId: 'instance_interrupt_signal_failure', instanceGeneration: 1 }, + }); + await instance.ctx.storage.put('wrapper_runtime_state', { + wrapperGeneration: 1, + wrapperConnectionId: 'conn_interrupt_signal_failure', + wrapperRunId: 'wr_interrupt_signal_failure', + }); + + const interrupt = await instance.interruptExecution(); + const runtimeState = await instance.ctx.storage.get<{ wrapperGeneration: number }>( + 'wrapper_runtime_state' + ); + return { + interrupt, + runtimeState, + stopped, + lease: await getWrapperLease(instance.ctx.storage), + }; + }); + + expect(result.interrupt).toEqual({ success: true, executionId: undefined }); + expect(result.runtimeState).toEqual({ wrapperGeneration: 2 }); + expect(result.stopped).toEqual(['user-interrupt']); + expect(result.lease).toMatchObject({ state: 'none' }); }); it('interrupt with a live fenced socket sends kill then fences accepted work', async () => { @@ -975,7 +1152,6 @@ describe('pending session messages', () => { ); const result = await runInDurableObject(stub, async instance => { - instance['stopCurrentWrapperProcess'] = async () => true; await registerReadySession(instance, { sessionId, userId, @@ -1073,12 +1249,30 @@ describe('pending session messages', () => { gitUrl: 'https://example.com/repo.git', gitToken: 'old-token', }); + await instance.ctx.storage.put('wrapper_lease', { + state: 'owns_wrapper', + nextInstanceGeneration: 2, + instance: { instanceId: 'instance_active_flush', instanceGeneration: 1 }, + }); + instance['physicalWrapperObserver'] = async () => ({ + status: 'present', + observed: [ + { + representation: 'process', + id: 'wrapper-active-flush', + port: 5000, + instanceId: 'instance_active_flush', + instanceGeneration: 1, + }, + ], + }); await instance.ctx.storage.put('wrapper_runtime_state', { wrapperGeneration: 1, wrapperConnectionId: 'conn_active_flush', wrapperRunId: 'wr_active_flush', lastWrapperMessageAt: Date.now(), }); + await putSessionMessageState(instance.ctx.storage, { messageId: 'msg_018f1e2d3c4bActBusyRunAbCd', status: 'accepted', diff --git a/services/cloud-agent-next/test/unit/execution/orchestrator.test.ts b/services/cloud-agent-next/test/unit/execution/orchestrator.test.ts index cd4fccf3be..7f21c4e13a 100644 --- a/services/cloud-agent-next/test/unit/execution/orchestrator.test.ts +++ b/services/cloud-agent-next/test/unit/execution/orchestrator.test.ts @@ -1,4 +1,7 @@ import { afterEach, describe, expect, it, vi } from 'vitest'; + +vi.mock('@cloudflare/sandbox', () => ({ getSandbox: vi.fn() })); + import { ExecutionError, isExecutionError, @@ -6,6 +9,7 @@ import { type RetryableErrorCode, } from '../../../src/execution/errors.js'; import { ExecutionOrchestrator } from '../../../src/execution/orchestrator.js'; +import { CloudflareAgentSandbox } from '../../../src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.js'; import { WrapperClient } from '../../../src/kilo/wrapper-client.js'; import { SessionService } from '../../../src/session-service.js'; import type { @@ -188,7 +192,10 @@ function createOrchestrator( const recordKiloServerActivity = options.recordKiloServerActivity ?? vi.fn().mockResolvedValue(undefined); return new ExecutionOrchestrator({ - getSandbox: vi.fn().mockResolvedValue(sandbox), + getAgentSandbox: plan => + new CloudflareAgentSandbox(env as Env, plan.workspace.metadata, { + resolveSandbox: () => sandbox, + }), getSessionStub: vi.fn( () => ({ @@ -416,7 +423,7 @@ describe('ExecutionOrchestrator bootstrap execution', () => { ]); }); - it('does not emit preparation progress when delivering to a warm workspace', async () => { + it('reports Kilo startup progress when delivering to a warm workspace', async () => { const { sandbox } = createMockSandbox({ workspaceWarm: true }); const { ensureSessionReady, prompt } = stubWrapperBootstrap(); const orchestrator = createOrchestrator(sandbox); @@ -424,7 +431,7 @@ describe('ExecutionOrchestrator bootstrap execution', () => { await orchestrator.execute(createExecutionPlan(), { onProgress }); - expect(onProgress).not.toHaveBeenCalled(); + expect(onProgress).toHaveBeenCalledExactlyOnceWith('kilo_server', 'Starting Kilo...'); expect(ensureSessionReady).toHaveBeenCalledOnce(); expect(prompt).toHaveBeenCalledOnce(); }); diff --git a/services/cloud-agent-next/wrapper/src/main.ts b/services/cloud-agent-next/wrapper/src/main.ts index a29067c392..394aba82d5 100644 --- a/services/cloud-agent-next/wrapper/src/main.ts +++ b/services/cloud-agent-next/wrapper/src/main.ts @@ -65,12 +65,16 @@ type StartupArgs = { agentSessionId: string; userId: string; sessionId?: string; + wrapperInstanceId?: string; + wrapperInstanceGeneration?: number; }; function parseStartupArgs(argv: string[]): StartupArgs { let agentSessionId: string | undefined; let userId: string | undefined; let sessionId: string | undefined; + let wrapperInstanceId: string | undefined; + let wrapperInstanceGeneration: number | undefined; for (let index = 0; index < argv.length; index++) { const arg = argv[index]; @@ -103,6 +107,28 @@ function parseStartupArgs(argv: string[]): StartupArgs { continue; } + if (arg === '--wrapper-instance-id') { + if (!value) { + failStartup('Missing value for --wrapper-instance-id'); + } + wrapperInstanceId = value; + index++; + continue; + } + + if (arg === '--wrapper-instance-generation') { + if (!value) { + failStartup('Missing value for --wrapper-instance-generation'); + } + const generation = Number.parseInt(value, 10); + if (!Number.isInteger(generation) || generation < 0) { + failStartup('Invalid value for --wrapper-instance-generation'); + } + wrapperInstanceGeneration = generation; + index++; + continue; + } + failStartup(`Unknown argument: ${arg}`); } @@ -114,7 +140,11 @@ function parseStartupArgs(argv: string[]): StartupArgs { failStartup('Missing required --user-id argument'); } - return { agentSessionId, userId, sessionId }; + if ((wrapperInstanceId === undefined) !== (wrapperInstanceGeneration === undefined)) { + failStartup('Wrapper instance identity requires both id and generation'); + } + + return { agentSessionId, userId, sessionId, wrapperInstanceId, wrapperInstanceGeneration }; } // --------------------------------------------------------------------------- @@ -129,11 +159,41 @@ async function main() { // is now passed in the POST /job/prompt body. const wrapperPort = getOptionalEnvInt('WRAPPER_PORT', 5000); const initialWorkspacePath = process.env.WORKSPACE_PATH; - const { - agentSessionId, - userId, - sessionId: configuredSessionId, - } = parseStartupArgs(process.argv.slice(2)); + const startupArgs = parseStartupArgs(process.argv.slice(2)); + // New bundles report env-based identity; old bundles safely ignore these rolling-deploy markers. + const envWrapperInstanceId = process.env.WRAPPER_INSTANCE_ID; + const envWrapperInstanceGenerationValue = process.env.WRAPPER_INSTANCE_GENERATION; + let envWrapperInstanceGeneration: number | undefined; + if (envWrapperInstanceGenerationValue !== undefined) { + const parsedGeneration = Number.parseInt(envWrapperInstanceGenerationValue, 10); + if (!Number.isInteger(parsedGeneration) || parsedGeneration < 0) { + failStartup('Invalid value for WRAPPER_INSTANCE_GENERATION'); + } + envWrapperInstanceGeneration = parsedGeneration; + } + if ( + startupArgs.wrapperInstanceId !== undefined && + envWrapperInstanceId !== undefined && + startupArgs.wrapperInstanceId !== envWrapperInstanceId + ) { + failStartup('Conflicting wrapper instance id configuration'); + } + if ( + startupArgs.wrapperInstanceGeneration !== undefined && + envWrapperInstanceGeneration !== undefined && + startupArgs.wrapperInstanceGeneration !== envWrapperInstanceGeneration + ) { + failStartup('Conflicting wrapper instance generation configuration'); + } + const agentSessionId = startupArgs.agentSessionId; + const userId = startupArgs.userId; + const configuredSessionId = startupArgs.sessionId; + const wrapperInstanceId = startupArgs.wrapperInstanceId ?? envWrapperInstanceId; + const wrapperInstanceGeneration = + startupArgs.wrapperInstanceGeneration ?? envWrapperInstanceGeneration; + if ((wrapperInstanceId === undefined) !== (wrapperInstanceGeneration === undefined)) { + failStartup('Wrapper instance identity requires both id and generation'); + } if (!SESSION_ID_RE.test(agentSessionId)) { failStartup(`Invalid agent session ID: ${agentSessionId}`); @@ -150,6 +210,11 @@ async function main() { if (configuredSessionId) { logToFile(`config: sessionId=${configuredSessionId}`); } + if (wrapperInstanceId !== undefined && wrapperInstanceGeneration !== undefined) { + logToFile( + `config: wrapperInstanceId=${wrapperInstanceId} wrapperInstanceGeneration=${wrapperInstanceGeneration}` + ); + } // --------------------------------------------------------------------------- // Wire up components @@ -178,6 +243,8 @@ async function main() { sessionId: kiloSessionId, agentSessionId, userId, + wrapperInstanceId, + wrapperInstanceGeneration, platform: process.env.KILO_PLATFORM, }; diff --git a/services/cloud-agent-next/wrapper/src/server.test.ts b/services/cloud-agent-next/wrapper/src/server.test.ts index ef7feba389..88235dc12f 100644 --- a/services/cloud-agent-next/wrapper/src/server.test.ts +++ b/services/cloud-agent-next/wrapper/src/server.test.ts @@ -79,6 +79,8 @@ function createTestFetch(overrides?: { sessionId: 'kilo_sess_test', agentSessionId: 'agent_00000000-0000-0000-0000-000000000000', userId: 'user_test', + wrapperInstanceId: 'instance_test', + wrapperInstanceGeneration: 8, }, { state: new WrapperState(), @@ -98,6 +100,21 @@ afterEach(async () => { await Promise.all(servers.splice(0).map(server => server.stop())); }); +describe('wrapper health', () => { + it('reports leased physical wrapper identity separately from session identity', async () => { + const { fetchHandler } = createTestFetch(); + const response = await fetchHandler(new Request('http://wrapper.test/health')); + if (!response) throw new Error('Expected health response'); + + const body = await response.json(); + expect(body).toMatchObject({ + sessionId: 'kilo_sess_test', + wrapperInstanceId: 'instance_test', + wrapperInstanceGeneration: 8, + }); + }); +}); + describe('wrapper PTY routes', () => { it('creates a workspace PTY and applies the requested size', async () => { const { fetchHandler, ptyCalls, resizeCalls } = createTestFetch(); diff --git a/services/cloud-agent-next/wrapper/src/server.ts b/services/cloud-agent-next/wrapper/src/server.ts index 19281af47b..2a86baec5c 100644 --- a/services/cloud-agent-next/wrapper/src/server.ts +++ b/services/cloud-agent-next/wrapper/src/server.ts @@ -40,6 +40,9 @@ export type ServerConfig = { agentSessionId: string; /** Stable Cloud Agent user ID, passed at wrapper startup */ userId: string; + /** Stable physical wrapper identity, present after leased startup. */ + wrapperInstanceId?: string; + wrapperInstanceGeneration?: number; /** Product surface that created the session, e.g. code-review. */ platform?: string; }; @@ -298,6 +301,10 @@ function createHealthHandler(config: ServerConfig, state: WrapperState) { state: state.isActive ? 'active' : 'idle', version: config.version, sessionId: config.sessionId, + ...(config.wrapperInstanceId ? { wrapperInstanceId: config.wrapperInstanceId } : {}), + ...(config.wrapperInstanceGeneration !== undefined + ? { wrapperInstanceGeneration: config.wrapperInstanceGeneration } + : {}), pendingMessages: state.pendingMessageIds.length, }); };