diff --git a/docs/adr/001-stateful-code-environments.md b/docs/adr/001-stateful-code-environments.md index 8b9830df..cfa44bf4 100644 --- a/docs/adr/001-stateful-code-environments.md +++ b/docs/adr/001-stateful-code-environments.md @@ -56,8 +56,14 @@ worker replacement; the UI and operator documentation must not imply otherwise. revocation. - Pairing codes and credentials are stored by digest where lookup permits. - One configured worker has at most one active fenced assignment. -- Sandbox isolation and default-deny egress remain mandatory; pairing secures - the transport identity but does not make the host a sandbox. +- Sandbox isolation and default-deny egress remain the mandatory default; + pairing secures the transport identity but does not make the host a sandbox. + An operator may explicitly delegate network and local-socket restrictions to + an approved outer VM boundary through a named, digested worker policy. That + delegation retains direct workspace filesystem rules, cancellation, and + resource limits. The operator is responsible for preventing permitted host + services (for example, a privileged container socket) from bypassing those + rules and exposing worker identity or credential material. - A compromised worker can lie about advertised capabilities. Capability labels and policy digests are audit signals until enforcement is coupled to an attested sandbox or trusted host policy. @@ -65,7 +71,8 @@ worker replacement; the UI and operator documentation must not imply otherwise. ## Consequences - `@librechat/code` owns the provider-neutral protocol, identity handling, and - worker CLI; Code API owns enrollment, scheduling, and execution policy. + worker CLI, including machine-local execution-policy presets; Code API owns + enrollment, scheduling, and execution policy. - LibreChat owns environment persistence, ownership, RBAC, and user experience. - The Agents SDK keeps only its adapter until a second concrete consumer proves which coding-tool abstractions are genuinely provider neutral. diff --git a/packages/code/README.md b/packages/code/README.md index 819d9f5c..e46d6fd6 100644 --- a/packages/code/README.md +++ b/packages/code/README.md @@ -211,6 +211,41 @@ LIBRECHAT_CODE_COMMAND_SANDBOX=native-srt librechat-code run \ --worker-dir /path/to/project --allow-workspace-commands ``` +### Trusted VM command policy + +The native SRT backend can be made intentionally permissive when the selected +machine already supplies an administrator-approved outer security boundary. +The `trusted-vm` preset keeps SRT's direct filesystem rules, credential +masking, private scratch storage, cancellation, time limits, and output limits, +while allowing unmatched outbound destinations, local port binding, and Unix +sockets: + +```bash +librechat-code run \ + --worker-dir /home/ubuntu/src \ + --allow-workspace-writes \ + --allow-workspace-commands \ + --command-policy-preset trusted-vm +``` + +`LIBRECHAT_CODE_COMMAND_POLICY_PRESET=trusted-vm` is the environment equivalent. +The default is `restricted`, which preserves the default-deny network policy. +The preset configures `native-srt`; it is not an unsandboxed host-shell +backend. It is rejected unless native workspace commands are enabled. Its +normalized effective controls are included in the worker policy digest, and +the worker advertises `anthropic-srt:trusted-vm` unless an operator supplied a +custom sandbox profile label. + +Treat this preset as delegation to the machine's outer security controls. Any +outbound destination can receive workspace data, local listeners can accept +connections reachable under host policy, and Unix socket access may expose +powerful host services such as a container daemon. A socket that grants host +privilege can bypass SRT's filesystem rules and reach worker or GitHub identity +material; the outer VM boundary must prevent that path or explicitly accept +that trust. Register only the intended source root. Worker identity, +mutation-quarantine state, and configured GitHub App key files must remain +outside it. + ## Docker runtime supervisor (optional hardened adapter) `DockerRuntimeSupervisor` is the first self-contained local OCI adapter. It diff --git a/packages/code/package.json b/packages/code/package.json index a819af9d..452a8be7 100644 --- a/packages/code/package.json +++ b/packages/code/package.json @@ -35,6 +35,10 @@ "types": "./dist/native-sandbox.d.ts", "import": "./dist/native-sandbox.js" }, + "./native-policy": { + "types": "./dist/native-policy.d.ts", + "import": "./dist/native-policy.js" + }, "./github": { "types": "./dist/github.d.ts", "import": "./dist/github.js" diff --git a/packages/code/src/cli.test.ts b/packages/code/src/cli.test.ts index f978fce5..18456072 100644 --- a/packages/code/src/cli.test.ts +++ b/packages/code/src/cli.test.ts @@ -94,6 +94,46 @@ test('CLI rejects an unknown command sandbox before entering the run loop', () = ); }); +test('CLI rejects an unknown native SRT command policy preset', () => { + const result = spawnSync( + process.execPath, + [fileURLToPath(new URL('./cli.js', import.meta.url))], + { + encoding: 'utf8', + env: { + ...process.env, + LIBRECHAT_CODE_URL: 'https://code.example/v1', + LIBRECHAT_CODE_WORKER_TOKEN: 'worker-secret', + LIBRECHAT_CODE_WORKER_ID: 'engineering-vm', + LIBRECHAT_CODE_COMMAND_POLICY_PRESET: 'host-shell', + }, + }, + ); + + assert.notEqual(result.status, 0); + assert.match(result.stderr, /must be restricted or trusted-vm/); +}); + +test('CLI refuses a permissive policy when native commands are unavailable', () => { + const result = spawnSync( + process.execPath, + [fileURLToPath(new URL('./cli.js', import.meta.url))], + { + encoding: 'utf8', + env: { + ...process.env, + LIBRECHAT_CODE_URL: 'https://code.example/v1', + LIBRECHAT_CODE_WORKER_TOKEN: 'worker-secret', + LIBRECHAT_CODE_WORKER_ID: 'engineering-vm', + LIBRECHAT_CODE_COMMAND_POLICY_PRESET: 'trusted-vm', + }, + }, + ); + + assert.notEqual(result.status, 0); + assert.match(result.stderr, /requires native-srt workspace commands/); +}); + test('CLI rejects incomplete GitHub App authentication before worker registration', () => { const result = spawnSync( process.execPath, diff --git a/packages/code/src/cli.ts b/packages/code/src/cli.ts index 27a8d2d9..fdd28e9f 100644 --- a/packages/code/src/cli.ts +++ b/packages/code/src/cli.ts @@ -29,6 +29,10 @@ import { import { RuntimeWorkspaceCommandSandbox } from './workspace-runtime.js'; import { NativeProcessWorkspaceCommandSandbox } from './native-process.js'; import { NativeWorkspaceCommandPool } from './native-pool.js'; +import { + resolveNativeSrtCommandPolicy, + serializeNativeSrtCommandPolicy, +} from './native-policy.js'; import { workspaceMutationGuard } from './workspace-guards.js'; import type { NativeProcessSandboxOptions } from './native-process.js'; import type { LocalWorkspaceConfig } from './workspace.js'; @@ -404,6 +408,19 @@ async function run( 'LIBRECHAT_CODE_COMMAND_SANDBOX must be native-srt or runtime', ); } + const commandPolicy = resolveNativeSrtCommandPolicy( + option(args, '--command-policy-preset') ?? + process.env.LIBRECHAT_CODE_COMMAND_POLICY_PRESET?.trim().toLowerCase() ?? + 'restricted', + ); + if ( + commandPolicy.preset !== 'restricted' && + (!allowWorkspaceCommands || commandSandboxMode !== 'native-srt') + ) { + throw new Error( + 'A permissive command policy preset requires native-srt workspace commands', + ); + } const github = runtimeSessionId == null ? githubCredentials() @@ -756,6 +773,7 @@ async function run( }); const nativeOptions: NativeProcessSandboxOptions = { workspaceRoot: canonicalWorkerDirectory!, + commandPolicy, protectedPaths: [ identityPath, ...rootQuarantinePaths.values(), @@ -820,7 +838,9 @@ async function run( sandboxProfile: process.env.LIBRECHAT_CODE_SANDBOX_PROFILE ?? (allowWorkspaceCommands && commandSandboxMode === 'native-srt' - ? 'anthropic-srt' + ? commandPolicy.preset === 'restricted' + ? 'anthropic-srt' + : `anthropic-srt:${commandPolicy.preset}` : runtimeMode.startsWith('docker') ? 'oci-docker' : 'nsjail'), @@ -829,7 +849,7 @@ async function run( .update(policy) .update( allowWorkspaceCommands && commandSandboxMode === 'native-srt' - ? `\0native-srt\0${commandAllowedDomains.join('\0')}\0${github.policyIdentity}` + ? `\0native-srt\0${serializeNativeSrtCommandPolicy(commandPolicy)}\0${commandAllowedDomains.join('\0')}\0${github.policyIdentity}` : '', ) .digest('hex'), diff --git a/packages/code/src/index.ts b/packages/code/src/index.ts index 712e7487..5363f0c0 100644 --- a/packages/code/src/index.ts +++ b/packages/code/src/index.ts @@ -5,6 +5,7 @@ export * from './storage.js'; export * from './runtime.js'; export * from './workspace.js'; export * from './workspace-runtime.js'; +export * from './native-policy.js'; export * from './native-sandbox.js'; export * from './native-process.js'; export * from './github.js'; diff --git a/packages/code/src/native-policy.test.ts b/packages/code/src/native-policy.test.ts new file mode 100644 index 00000000..4e9f4691 --- /dev/null +++ b/packages/code/src/native-policy.test.ts @@ -0,0 +1,63 @@ +import assert from 'node:assert/strict'; +import test from 'node:test'; + +import { + normalizeNativeSrtCommandPolicy, + resolveNativeSrtCommandPolicy, + serializeNativeSrtCommandPolicy, +} from './native-policy.js'; + +test('restricted remains the default native SRT command policy', () => { + assert.deepEqual(resolveNativeSrtCommandPolicy(), { + version: 1, + preset: 'restricted', + network: { + outbound: 'allowlist', + allowLocalBinding: false, + allowAllUnixSockets: false, + }, + }); +}); + +test('trusted-vm resolves to explicit permissive network controls', () => { + assert.deepEqual(resolveNativeSrtCommandPolicy('trusted-vm'), { + version: 1, + preset: 'trusted-vm', + network: { + outbound: 'unrestricted', + allowLocalBinding: true, + allowAllUnixSockets: true, + }, + }); +}); + +test('unknown and forged native policies fail closed', () => { + assert.throws( + () => resolveNativeSrtCommandPolicy('host-shell'), + /must be restricted or trusted-vm/, + ); + assert.throws( + () => + normalizeNativeSrtCommandPolicy({ + ...resolveNativeSrtCommandPolicy('restricted'), + network: { + ...resolveNativeSrtCommandPolicy('restricted').network, + allowAllUnixSockets: true, + }, + }), + /does not match its preset/, + ); +}); + +test('serialized policy is stable and includes effective controls', () => { + const first = serializeNativeSrtCommandPolicy( + resolveNativeSrtCommandPolicy('trusted-vm'), + ); + const second = serializeNativeSrtCommandPolicy( + resolveNativeSrtCommandPolicy('trusted-vm'), + ); + assert.equal(first, second); + assert.match(first, /"outbound":"unrestricted"/); + assert.match(first, /"allowLocalBinding":true/); + assert.match(first, /"allowAllUnixSockets":true/); +}); diff --git a/packages/code/src/native-policy.ts b/packages/code/src/native-policy.ts new file mode 100644 index 00000000..5d289bfb --- /dev/null +++ b/packages/code/src/native-policy.ts @@ -0,0 +1,86 @@ +export const NATIVE_SRT_COMMAND_POLICY_PRESETS = [ + 'restricted', + 'trusted-vm', +] as const; + +export type NativeSrtCommandPolicyPreset = + typeof NATIVE_SRT_COMMAND_POLICY_PRESETS[number]; + +export interface NativeSrtCommandPolicy { + version: 1; + preset: NativeSrtCommandPolicyPreset; + network: { + outbound: 'allowlist' | 'unrestricted'; + allowLocalBinding: boolean; + allowAllUnixSockets: boolean; + }; +} + +const PRESETS: Record = { + restricted: { + version: 1, + preset: 'restricted', + network: { + outbound: 'allowlist', + allowLocalBinding: false, + allowAllUnixSockets: false, + }, + }, + 'trusted-vm': { + version: 1, + preset: 'trusted-vm', + network: { + outbound: 'unrestricted', + allowLocalBinding: true, + allowAllUnixSockets: true, + }, + }, +}; + +function isPreset(value: unknown): value is NativeSrtCommandPolicyPreset { + return ( + typeof value === 'string' && + NATIVE_SRT_COMMAND_POLICY_PRESETS.some((preset) => preset === value) + ); +} + +/** Resolve a named convenience preset into the explicit policy SRT enforces. */ +export function resolveNativeSrtCommandPolicy( + preset: unknown = 'restricted', +): NativeSrtCommandPolicy { + if (!isPreset(preset)) { + throw new Error( + 'Native SRT command policy preset must be restricted or trusted-vm', + ); + } + const policy = PRESETS[preset]; + return { ...policy, network: { ...policy.network } }; +} + +/** Validate a programmatic policy and return canonical preset-owned values. */ +export function normalizeNativeSrtCommandPolicy( + policy?: NativeSrtCommandPolicy, +): NativeSrtCommandPolicy { + const normalized = resolveNativeSrtCommandPolicy( + policy?.preset ?? 'restricted', + ); + if ( + policy !== undefined && + (policy.version !== normalized.version || + policy.network?.outbound !== normalized.network.outbound || + policy.network?.allowLocalBinding !== + normalized.network.allowLocalBinding || + policy.network?.allowAllUnixSockets !== + normalized.network.allowAllUnixSockets) + ) { + throw new Error('Native SRT command policy does not match its preset'); + } + return normalized; +} + +/** Stable policy material used in the bridge capability digest. */ +export function serializeNativeSrtCommandPolicy( + policy: NativeSrtCommandPolicy, +): string { + return JSON.stringify(normalizeNativeSrtCommandPolicy(policy)); +} diff --git a/packages/code/src/native-process.test.ts b/packages/code/src/native-process.test.ts index d77ec559..8dcfd4c1 100644 --- a/packages/code/src/native-process.test.ts +++ b/packages/code/src/native-process.test.ts @@ -108,6 +108,36 @@ test('executor bootstrap excludes bridge credentials and Node injection variable await sandbox.close(); }); +test('executor forwards the resolved command policy without worker credentials', async () => { + const fake = fixture(); + const sandbox = new NativeProcessWorkspaceCommandSandbox( + { + workspaceRoot: '/workspace', + commandPolicy: { + version: 1, + preset: 'trusted-vm', + network: { + outbound: 'unrestricted', + allowLocalBinding: true, + allowAllUnixSockets: true, + }, + }, + }, + fake.fork, + ); + await sandbox.prepare(); + assert.deepEqual(fake.messages[0].options.commandPolicy, { + version: 1, + preset: 'trusted-vm', + network: { + outbound: 'unrestricted', + allowLocalBinding: true, + allowAllUnixSockets: true, + }, + }); + await sandbox.close(); +}); + test('executor hands credentials over IPC only for the current command', async () => { const fake = fixture(); const sandbox = new NativeProcessWorkspaceCommandSandbox( diff --git a/packages/code/src/native-process.ts b/packages/code/src/native-process.ts index fe4e1aa8..96d89355 100644 --- a/packages/code/src/native-process.ts +++ b/packages/code/src/native-process.ts @@ -182,6 +182,7 @@ export class NativeProcessWorkspaceCommandSandbox child.on('disconnect', lost); const { workspaceRoot, + commandPolicy, protectedPaths, allowedDomains, homeDirectory, @@ -192,6 +193,7 @@ export class NativeProcessWorkspaceCommandSandbox { options: { workspaceRoot, + commandPolicy, protectedPaths, allowedDomains, homeDirectory, diff --git a/packages/code/src/native-sandbox.test.ts b/packages/code/src/native-sandbox.test.ts index e2dc4bda..cf790965 100644 --- a/packages/code/src/native-sandbox.test.ts +++ b/packages/code/src/native-sandbox.test.ts @@ -19,7 +19,10 @@ import { join } from 'node:path'; import { PassThrough } from 'node:stream'; import test from 'node:test'; -import type { SandboxRuntimeConfig } from '@anthropic-ai/sandbox-runtime'; +import type { + SandboxAskCallback, + SandboxRuntimeConfig, +} from '@anthropic-ai/sandbox-runtime'; import type { ChildProcessWithoutNullStreams } from 'node:child_process'; import { NativeSrtWorkspaceCommandSandbox } from './native-sandbox.js'; @@ -46,6 +49,7 @@ function fakeManager( } = {}, ) { let config: SandboxRuntimeConfig | undefined; + let askCallback: SandboxAskCallback | undefined; let reset = false; let credentialSeenDuringWrap: string | undefined; let gitLfsRequiredSeenDuringWrap: string | undefined; @@ -55,8 +59,12 @@ function fakeManager( async checkDependenciesAsync() { return { warnings: [], errors: options.dependencyErrors ?? [] }; }, - async initialize(value: SandboxRuntimeConfig) { + async initialize( + value: SandboxRuntimeConfig, + callback?: SandboxAskCallback, + ) { config = value; + askCallback = callback; if (options.initializeError) throw options.initializeError; }, async wrapWithSandboxArgv(command: string) { @@ -107,6 +115,9 @@ function fakeManager( get config() { return config; }, + get askCallback() { + return askCallback; + }, get reset() { return reset; }, @@ -276,6 +287,8 @@ test('initializes SRT with a default-deny network and scrubbed worker credential assert.deepEqual(fake.config?.network.allowedDomains, []); assert.equal(fake.config?.network.strictAllowlist, true); assert.equal(fake.config?.network.allowAllUnixSockets, false); + assert.equal(fake.config?.network.allowLocalBinding, false); + assert.equal(fake.askCallback, undefined); assert.deepEqual(fake.config?.filesystem.allowRead, [ canonicalRoot, scratchDirectory, @@ -305,6 +318,39 @@ test('initializes SRT with a default-deny network and scrubbed worker credential await assert.rejects(access(scratchDirectory!)); }); +test('trusted-vm permits unmatched egress and local development sockets', async (t) => { + const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); + t.after(() => rm(root, { recursive: true, force: true })); + const fake = fakeManager(); + const sandbox = new NativeSrtWorkspaceCommandSandbox({ + workspaceRoot: root, + commandPolicy: { + version: 1, + preset: 'trusted-vm', + network: { + outbound: 'unrestricted', + allowLocalBinding: true, + allowAllUnixSockets: true, + }, + }, + manager: fake.manager, + }); + t.after(() => sandbox.close()); + + await sandbox.prepare(); + + assert.equal(fake.config?.network.strictAllowlist, false); + assert.equal(fake.config?.network.allowLocalBinding, true); + assert.equal(fake.config?.network.allowAllUnixSockets, true); + assert.equal( + await fake.askCallback?.({ host: 'packages.example', port: 443 }), + true, + ); + assert.deepEqual(fake.config?.filesystem.allowWrite.slice(0, 1), [ + await realpath(root), + ]); +}); + test('provides an isolated scratch directory to commands and restores the host environment', async (t) => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); diff --git a/packages/code/src/native-sandbox.ts b/packages/code/src/native-sandbox.ts index 139a10d6..171ea663 100644 --- a/packages/code/src/native-sandbox.ts +++ b/packages/code/src/native-sandbox.ts @@ -41,7 +41,12 @@ import type { ChildProcessWithoutNullStreams, SpawnOptionsWithoutStdio, } from 'node:child_process'; -import type { SandboxRuntimeConfig } from '@anthropic-ai/sandbox-runtime'; +import type { + SandboxAskCallback, + SandboxRuntimeConfig, +} from '@anthropic-ai/sandbox-runtime'; +import { normalizeNativeSrtCommandPolicy } from './native-policy.js'; +import type { NativeSrtCommandPolicy } from './native-policy.js'; import type { WorkspaceExecuteCommandRequest, WorkspaceExecuteCommandResult, @@ -124,7 +129,10 @@ const HOST_TEMPORARY_ROOT = tmpdir(); interface NativeSandboxManager { isSupportedPlatform(): boolean; checkDependenciesAsync(): Promise<{ warnings: string[]; errors: string[] }>; - initialize(config: SandboxRuntimeConfig): Promise; + initialize( + config: SandboxRuntimeConfig, + sandboxAskCallback?: SandboxAskCallback, + ): Promise; wrapWithSandboxArgv( command: string, binShell?: string, @@ -153,6 +161,7 @@ type SpawnCommand = ( export interface NativeSrtWorkspaceCommandSandboxOptions { workspaceRoot: string; + commandPolicy?: NativeSrtCommandPolicy; /** Trusted worker files that must never become workspace-readable or writable. */ protectedPaths?: string[]; allowedDomains?: string[]; @@ -369,13 +378,18 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox 'REGISTRATION_INVALID', ); } + const commandPolicy = normalizeNativeSrtCommandPolicy( + this.options.commandPolicy, + ); + const unrestrictedNetwork = + commandPolicy.network.outbound === 'unrestricted'; const config: SandboxRuntimeConfig = { network: { allowedDomains: [...(this.options.allowedDomains ?? [])], deniedDomains: [], - strictAllowlist: true, - allowAllUnixSockets: false, - allowLocalBinding: false, + strictAllowlist: !unrestrictedNetwork, + allowAllUnixSockets: commandPolicy.network.allowAllUnixSockets, + allowLocalBinding: commandPolicy.network.allowLocalBinding, ...(this.options.maskedEnvironment ? { tlsTerminate: {} } : {}), }, filesystem: { @@ -439,7 +453,10 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox enableWeakerNetworkIsolation: false, git: { safeDirectories: [root] }, }; - await this.manager.initialize(config); + await this.manager.initialize( + config, + unrestrictedNetwork ? async () => true : undefined, + ); this.canonicalRoot = root; } diff --git a/packages/code/src/protocol.ts b/packages/code/src/protocol.ts index 44cfdfd0..c1dad949 100644 --- a/packages/code/src/protocol.ts +++ b/packages/code/src/protocol.ts @@ -20,6 +20,8 @@ export const BRIDGE_WORKSPACE_COMMAND_MAX_TIMEOUT_MS = 5 * 60_000; export const BRIDGE_WORKSPACE_COMMAND_DEFAULT_OUTPUT_BYTES = 256 * 1024; export const BRIDGE_WORKSPACE_COMMAND_MAX_OUTPUT_BYTES = 1024 * 1024; export const BRIDGE_WORKSPACE_COMMAND_SIGNAL_MAX_LENGTH = 32; +/** How long Code API drains a clean rejection after Stop cancels a workspace mutation. */ +export const BRIDGE_CANCELLED_WORKSPACE_SETTLEMENT_GRACE_MS = 5_000; export type BridgeProtocolVersion = typeof BRIDGE_PROTOCOL_VERSION; diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index fde84768..19f219da 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -1,6 +1,7 @@ import { randomBytes } from 'node:crypto'; import { + BRIDGE_CANCELLED_WORKSPACE_SETTLEMENT_GRACE_MS, BRIDGE_PROTOCOL_VERSION, BridgeProtocolError, bridgeWorkerPath, @@ -1619,7 +1620,13 @@ export class BridgeWorker { assignment.runtimeSessionId != null && settlement.status === 'rejected' && (!sandboxStarted || sandboxRejectedExecution); - if (knownCleanStatefulRejection) { + // An armed mutation reaches settlement as rejected only after an atomic + // failure that does not require quarantine. Code API accepts that + // rejection after expiry and drains it for its own grace after Stop, so a + // Stop near the deadline must not cut off retries at the deadline. + const knownCleanWorkspaceRejection = + workspaceMutationArmed && settlement.status === 'rejected'; + if (knownCleanStatefulRejection || knownCleanWorkspaceRejection) { heartbeatController.abort(); await heartbeat; const recoveryHeartbeatController = new AbortController(); @@ -1627,15 +1634,24 @@ export class BridgeWorker { recoveryHeartbeatController.signal, true, ).catch(() => undefined); + const rejectionAckGraceMs = Math.max( + 0, + this.options.rejectionAckGraceMs ?? REJECTION_ACK_GRACE_MS, + ); try { await this.settleWithRetry( assignment, settlement, localDeadlineAtMs + - Math.max( - 0, - this.options.rejectionAckGraceMs ?? REJECTION_ACK_GRACE_MS, - ), + (knownCleanStatefulRejection + ? rejectionAckGraceMs + : Math.max( + rejectionAckGraceMs, + BRIDGE_CANCELLED_WORKSPACE_SETTLEMENT_GRACE_MS, + )), + // Stateful rejections outlive shutdown; workspace guards still + // fail closed when the worker itself stops. + knownCleanStatefulRejection ? undefined : signal, ); } finally { recoveryHeartbeatController.abort(); diff --git a/packages/code/src/workspace-worker.test.ts b/packages/code/src/workspace-worker.test.ts index 14bf2d61..8f9fcd45 100644 --- a/packages/code/src/workspace-worker.test.ts +++ b/packages/code/src/workspace-worker.test.ts @@ -1480,6 +1480,208 @@ test('worker clears quarantine after a command cancellation confirms process ter assert.deepEqual(lifecycle, ['arm', 'execute', 'settle', 'clear']); }); +test('worker retries a clean Stop rejection near its deadline through the cancellation grace', async () => { + const lifecycle: string[] = []; + const settlements: Array> = []; + const remainingMs = 100; + const startedAt = Date.now(); + const baseCapabilities = { + protocolVersion: 1 as const, + operations: ['read_file' as const], + workspaces: [{ id: 'primary', operations: ['read_file' as const] }], + }; + const workspaceTools = new SandboxWorkspaceTools({ + workspaceTools: { + capabilities: baseCapabilities, + mutationFailuresAreAtomic: true, + async execute() { throw new Error('base executor must not run'); }, + }, + commandWorkspaces: ['primary'], + commandSandbox: { + mutationFailuresAreAtomic: true, + async execute(_request, signal) { + lifecycle.push('execute'); + await new Promise((resolve) => { + if (signal?.aborted) return resolve(); + signal?.addEventListener('abort', () => resolve(), { once: true }); + }); + lifecycle.push('stop'); + // Process-group termination is confirmed after the original deadline. + await new Promise((resolve) => setTimeout(resolve, remainingMs)); + throw new WorkspaceToolError( + 'Workspace command execution aborted', + 'EXECUTION_ABORTED', + true, + false, + ); + }, + }, + }); + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + workspaceTools: workspaceTools.capabilities, + }, + workspaceTools, + workspaceMutationQuarantine: mutationQuarantine( + () => lifecycle.push('quarantine'), + () => lifecycle.push('arm'), + () => lifecycle.push('clear'), + ), + cancellationPollIntervalMs: 5, + fetchImpl: async (input, init) => { + if (String(input).endsWith('/cancellation')) { + return Response.json({ + protocolVersion: 1, + cancelled: Date.now() >= startedAt + remainingMs / 2, + }); + } + if (!String(input).endsWith('/settle')) { + return Response.json({ protocolVersion: 1, accepted: true }); + } + lifecycle.push('settle'); + settlements.push({ + ...(JSON.parse(String(init?.body)) as Record), + attemptedAt: Date.now(), + }); + // Settlement delivery takes a real transport turn and honors its deadline. + await new Promise((resolve, reject) => { + const timer = setTimeout(resolve, 10); + init?.signal?.addEventListener( + 'abort', + () => { + clearTimeout(timer); + reject(new DOMException('aborted', 'AbortError')); + }, + { once: true }, + ); + }); + if (settlements.length === 1) { + return Response.json( + { error: 'Bridge settlement temporarily unavailable' }, + { status: 503 }, + ); + } + return Response.json({ protocolVersion: 1, accepted: true }); + }, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-command-stopped-near-deadline', + workerId: 'vm-1', + incarnationId, + generation: 4, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(startedAt + remainingMs).toISOString(), + remainingMs, + executionKind: 'workspace_tool', + request: { + protocolVersion: 1, + operation: 'execute_command', + workspaceId: 'primary', + command: 'sleep 30; touch delayed.txt', + }, + }); + + assert.deepEqual(lifecycle, [ + 'arm', + 'execute', + 'stop', + 'settle', + 'settle', + 'clear', + ]); + assert.ok(Number(settlements[0]?.attemptedAt) > startedAt + remainingMs); + assert.equal(settlements[1]?.status, 'rejected'); + assert.equal(settlements[1]?.errorCode, 'EXECUTION_ABORTED'); +}); + +test('worker keeps quarantine armed when shutdown interrupts a clean command rejection', async () => { + const lifecycle: string[] = []; + const controller = new AbortController(); + const baseCapabilities = { + protocolVersion: 1 as const, + operations: ['read_file' as const], + workspaces: [{ id: 'primary', operations: ['read_file' as const] }], + }; + const workspaceTools = new SandboxWorkspaceTools({ + workspaceTools: { + capabilities: baseCapabilities, + mutationFailuresAreAtomic: true, + async execute() { throw new Error('base executor must not run'); }, + }, + commandWorkspaces: ['primary'], + commandSandbox: { + mutationFailuresAreAtomic: true, + async execute() { + lifecycle.push('execute'); + controller.abort(new Error('shutdown')); + throw new WorkspaceToolError( + 'Workspace command execution aborted', + 'EXECUTION_ABORTED', + true, + false, + ); + }, + }, + }); + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + workspaceTools: workspaceTools.capabilities, + }, + workspaceTools, + workspaceMutationQuarantine: mutationQuarantine( + () => lifecycle.push('quarantine'), + () => lifecycle.push('arm'), + () => lifecycle.push('clear'), + ), + fetchImpl: async () => { + lifecycle.push('settle'); + return Response.json({ protocolVersion: 1, accepted: true }); + }, + }); + + await assert.rejects( + worker.executeAndSettle( + { + protocolVersion: 1, + assignmentId: 'assignment-command-shutdown-cleanly', + workerId: 'vm-1', + incarnationId, + generation: 4, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 5_000).toISOString(), + executionKind: 'workspace_tool', + request: { + protocolVersion: 1, + operation: 'execute_command', + workspaceId: 'primary', + command: 'sleep 30', + }, + }, + controller.signal, + ), + /shutdown/, + ); + assert.deepEqual(lifecycle, ['arm', 'execute']); +}); + test('worker retains quarantine when an atomic executor cannot confirm durability', async () => { const lifecycle: string[] = []; const workspaceCapabilities = { diff --git a/service/src/bridge/concurrent-store.test.ts b/service/src/bridge/concurrent-store.test.ts index fe53c44f..736b7216 100644 --- a/service/src/bridge/concurrent-store.test.ts +++ b/service/src/bridge/concurrent-store.test.ts @@ -384,6 +384,46 @@ test('queued cancellation never leases and does not block another root', async ( ).toBeUndefined(); }); +test('unassigned slot releases even when dispatch cleanup fails', async () => { + await register(); + const originalIncr = redis.incr.bind(redis); + const originalSet = redis.set.bind(redis); + const set = originalSet as (...args: unknown[]) => unknown; + redis.incr = ((key: string) => + key.endsWith(':generation') + ? Promise.reject(new Error('injected generation outage')) + : originalIncr(key)) as typeof redis.incr; + redis.set = ((key: string, ...args: unknown[]) => + key.endsWith(':cancelled') + ? Promise.reject(new Error('injected cancellation outage')) + : set(key, ...args)) as typeof redis.set; + try { + // The reservation succeeds, then dispatch fails before storing an assignment. + await expect(dispatch('a')).rejects.toThrow('injected cancellation outage'); + } finally { + redis.incr = originalIncr; + redis.set = originalSet; + } + expect( + await redis.hlen(`codeapi:bridge:v1:worker:${workerId}:workspace-slots`), + ).toBe(0); + expect( + await redis.get(`codeapi:bridge:v1:worker:${workerId}:lock`), + ).toBeNull(); + const next = dispatch('a'); + const assignment = (await store.lease( + workerId, + incarnationId, + 1000, + undefined, + undefined, + 0, + ))!; + expect(assignment.request).toMatchObject({ workspaceId: 'a' }); + await settle(assignment); + await expect(next).resolves.toMatchObject({ status: 'rejected' }); +}); + test('late quarantine releases its slot after caller cancellation and retains only its root fence', async () => { await register(); const controller = new AbortController(); diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index 32205571..b09d3792 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -11,6 +11,7 @@ import type { } from '../../../packages/code/src/protocol'; import { + BRIDGE_CANCELLED_WORKSPACE_SETTLEMENT_GRACE_MS, BRIDGE_PROTOCOL_VERSION, isValidBridgeWorkerCapabilities, isValidBridgeWorkerId, @@ -23,7 +24,6 @@ import { BridgeWorkspaceSlots } from './slots'; const PREFIX = 'codeapi:bridge:v1'; const POLL_INTERVAL_MS = 100; -const CANCELLED_WORKSPACE_SETTLEMENT_GRACE_MS = 5_000; const DEFAULT_WORKER_TTL_SECONDS = 60; const DEFAULT_REDIS_COMMAND_TIMEOUT_MS = 1_000; @@ -1060,20 +1060,15 @@ export class RedisBridgeStore { // already committed result rather than turning cleanup availability // into a client-visible failure that could prompt duplicate work. } + } else if (workspaceSlots != null && assignment == null) { + await this.cleanupUnassignedSlot( + args.workerId, + lockIncarnationId, + assignmentId, + ); } else { await this.cleanupDispatch(args.workerId, assignmentId, assignment); } - if (workspaceSlots != null && assignment == null) { - await boundedCommand( - workspaceSlots.release( - args.workerId, - lockIncarnationId, - assignmentId, - ), - this.redisCommandTimeoutMs, - 'Bridge unassigned slot cleanup', - ); - } } } @@ -1924,7 +1919,7 @@ export class RedisBridgeStore { // Give Stop its own grace so a near-timeout cancellation is not // misclassified as an ambiguous timeout. const cancellationDeadlineAtMs = - Date.now() + CANCELLED_WORKSPACE_SETTLEMENT_GRACE_MS; + Date.now() + BRIDGE_CANCELLED_WORKSPACE_SETTLEMENT_GRACE_MS; let cancellationPollMs = POLL_INTERVAL_MS; while (Date.now() < cancellationDeadlineAtMs) { const raw = await boundedCommand( @@ -2138,6 +2133,29 @@ export class RedisBridgeStore { ]); } + private async cleanupUnassignedSlot( + workerId: string, + incarnationId: string, + assignmentId: string, + ): Promise { + // No stored assignment owns this reservation, so a cancellation outage + // must not leave the slot and its root busy until TTL expiry. + const [cleanup, release] = await Promise.allSettled([ + this.cleanupDispatch(workerId, assignmentId, undefined), + boundedCommand( + new BridgeWorkspaceSlots(this.redis).release( + workerId, + incarnationId, + assignmentId, + ), + this.redisCommandTimeoutMs, + 'Bridge unassigned slot cleanup', + ), + ]); + if (cleanup.status === 'rejected') throw cleanup.reason; + if (release.status === 'rejected') throw release.reason; + } + private async commitPendingWorkspace( assignment: StoredAssignment, settlement: AnyCodeBridgeSettlement,