From 8a90f6d1ba01c97ed5c4e01f835a783958e581b7 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sat, 12 Sep 2026 16:17:36 -0400 Subject: [PATCH 1/3] feat: Add Trusted VM Command Policy (#187) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 🛰️ feat: Add trusted VM command policy * docs: clarify trusted VM socket boundary --- docs/adr/001-stateful-code-environments.md | 13 +++- packages/code/README.md | 35 +++++++++ packages/code/package.json | 4 + packages/code/src/cli.test.ts | 40 ++++++++++ packages/code/src/cli.ts | 24 +++++- packages/code/src/index.ts | 1 + packages/code/src/native-policy.test.ts | 63 ++++++++++++++++ packages/code/src/native-policy.ts | 86 ++++++++++++++++++++++ packages/code/src/native-process.test.ts | 30 ++++++++ packages/code/src/native-process.ts | 2 + packages/code/src/native-sandbox.test.ts | 50 ++++++++++++- packages/code/src/native-sandbox.ts | 29 ++++++-- 12 files changed, 364 insertions(+), 13 deletions(-) create mode 100644 packages/code/src/native-policy.test.ts create mode 100644 packages/code/src/native-policy.ts 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; } From 31def177aa3ca2b4651bc413f96cf63d6522d707 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sat, 12 Sep 2026 20:31:55 -0400 Subject: [PATCH 2/3] fix: Retry Clean Cancelled BYOM Settlements Through Stop Grace (#188) A clean atomic workspace mutation rejection was settled with retries cut off at the original execution deadline. When Stop arrived near that deadline, process-tree termination finished after it, so the first settlement attempt was aborted immediately while Code API was still draining the cancellation. The client received ASSIGNMENT_EXPIRED and the durable mutation guard stayed armed. Route clean workspace mutation rejections through the known-clean rejection recovery path: a transient-retrying heartbeat and settlement retries through the rejection acknowledgement grace, floored at the bridge cancellation settlement grace. Worker shutdown still fails closed. Share the grace constant from the protocol module so the bridge and worker stay aligned. Closes #173 --- packages/code/src/protocol.ts | 2 + packages/code/src/worker.ts | 26 ++- packages/code/src/workspace-worker.test.ts | 202 +++++++++++++++++++++ service/src/bridge/store.ts | 4 +- 4 files changed, 227 insertions(+), 7 deletions(-) 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/store.ts b/service/src/bridge/store.ts index 32205571..5e4249a9 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; @@ -1924,7 +1924,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( From 118eb9b3ad704e752acc6ffa4b50fb9f1b3bb0ee Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sat, 12 Sep 2026 21:15:39 -0400 Subject: [PATCH 3/3] fix: Release Unassigned Workspace Slots When Dispatch Cleanup Fails (#189) Closes #170 --- service/src/bridge/concurrent-store.test.ts | 40 +++++++++++++++++++++ service/src/bridge/store.ts | 40 +++++++++++++++------ 2 files changed, 69 insertions(+), 11 deletions(-) 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 5e4249a9..b09d3792 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -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', - ); - } } } @@ -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,