diff --git a/docs/remote-bridge/README.md b/docs/remote-bridge/README.md index ab7e2f64..99a0dece 100644 --- a/docs/remote-bridge/README.md +++ b/docs/remote-bridge/README.md @@ -149,6 +149,17 @@ implementation and an allowlist of workspace IDs, preserves per-workspace operation restrictions, validates bounded results, and treats an unknown command failure as an uncertain mutation. +Selected attached workspaces on macOS, Linux, and WSL2 can also advertise Bash +Programmatic Tool Calling. Native Windows workers do not advertise Bash PTC. +Code API then runs each replay iteration through the same workspace-scoped +native SRT executor. Source code operates in the selected local root, while +replay metadata, injected skills and attachments, and generated artifacts are +staged in an execution-private data directory and removed after settlement. +Only authorized file references and returned artifacts cross the relay; the +repository is never uploaded to Code API. This capability is advertised only +when native SRT commands and a file-relay upstream are both configured, so +older or partially configured workers continue to fail closed. + Native SRT is the MVP and default command backend on a user's chosen laptop or VM. It uses Seatbelt on macOS, bubblewrap/seccomp on Linux, and the SRT restricted-account helper on Windows. It confines writes to the registered @@ -157,9 +168,11 @@ credentials, and denies network egress by default. Startup fails closed when the platform dependencies are unavailable; there is no unsandboxed fallback. Use `LIBRECHAT_CODE_COMMAND_ALLOWED_DOMAINS` for an explicit comma-separated egress allowlist. -Linux hosts must provide Bash at `/bin/bash`, `bubblewrap`, `socat`, and -`ripgrep`; macOS uses system facilities. Windows requires SRT's one-time -restricted-account setup. +Linux hosts must provide `bubblewrap`, `socat`, and `ripgrep`; macOS uses +system facilities. Bash Programmatic Tool Calling additionally requires Bash +5.2 or newer and `jq` on `PATH` on macOS, Linux, and WSL2. The worker resolves +the compatible shell from `PATH` rather than assuming `/bin/bash`. Windows +requires SRT's one-time restricted-account setup. The optional `docker-nsjail` adapter enables a stronger container boundary with `--allow-workspace-commands` (or diff --git a/packages/code/README.md b/packages/code/README.md index d06fbf9a..87d9dded 100644 --- a/packages/code/README.md +++ b/packages/code/README.md @@ -143,14 +143,48 @@ policy. This matches the personal-machine SRT trust model; use the Docker/NsJail backend or a dedicated VM boundary when hard teardown of adversarial process trees is required. -Linux hosts need Bash at `/bin/bash`, `bubblewrap`, `socat`, and `ripgrep`; macOS uses system -facilities. Follow SRT's one-time restricted-account setup when using Windows. +Linux hosts need `bubblewrap`, `socat`, and `ripgrep`; macOS uses system +facilities. Bash Programmatic Tool Calling additionally requires Bash 5.2 or +newer and `jq` on `PATH` on macOS, Linux, and WSL2. The worker resolves that +shell explicitly instead of assuming `/bin/bash`, which remains Bash 3.2 on +many macOS hosts. Follow SRT's one-time restricted-account setup when using Windows. An operator may allow explicit egress destinations with the comma-separated `LIBRECHAT_CODE_COMMAND_ALLOWED_DOMAINS` setting. Treat that as a security policy: an allowed destination can receive workspace data. The normalized allowlist is included in the worker policy digest. Tool approval hooks remain the user-facing allow/deny boundary for each invocation. +When Code API negotiates `bash` programmatic execution for a selected +workspace, the same native SRT executor also supports replay-mode Programmatic +Tool Calling on macOS, Linux, and WSL2 workers. Native Windows does not +advertise this Bash capability. The repository remains the command working directory. Generated +PTC scripts, replay history, skill files, chat attachments, and returned +artifacts use an owner-only per-execution directory under the worker's private +SRT scratch root, exposed to code as `LIBRECHAT_CODE_DATA_DIR`. That directory +is removed after every iteration and is never placed in the repository. + +Replay probes run against a disposable copy-on-write snapshot with network and +socket access denied, including under `trusted-vm`. External effects must not +repeat while discovering pending tools. Use registered tools for network-dependent +replay control flow; the final commit pass runs once under the configured policy. +Each probe's SRT proxy session is revoked before restoring the commit policy; +per-command network overrides alone do not restrict SRT's session-level proxies. +Probe failures do not quarantine the real workspace. Once the commit pass starts, +its fence remains until result restoration succeeds; uncertain finalization +quarantines only that workspace. + +Reference inputs and artifact outputs travel only through the configured +`LIBRECHAT_CODE_FILE_RELAY_UPSTREAM`, using Code API's execution-scoped opaque +egress grant. The worker rejects redirects and bounds each transfer to 10 MiB, +each execution to 100 files and 100 MiB total, and transfer concurrency to four. +Caller inputs are limited to 98 files, reserving two for the script and replay +history. Code API reserves one third of the job budget for all transfer batches +and negotiates each transfer's deadline before signing the request. +Its parent process keeps a 64-entry/32-MiB LRU input cache keyed by a stable, +Code-API-authorized digest; sandboxed commands cannot read that cache. Requests +against one workspace remain serialized, while negotiated lease slots allow +different registered roots to execute concurrently. + The native sandbox preserves standard `HTTP_PROXY`, `HTTPS_PROXY`, `ALL_PROXY`, and `NO_PROXY` names (including lowercase forms), plus Windows process and profile variables on Windows. SRT remains responsible for the final sandbox environment diff --git a/packages/code/src/cli.ts b/packages/code/src/cli.ts index a87f1a5a..289b0dc4 100644 --- a/packages/code/src/cli.ts +++ b/packages/code/src/cli.ts @@ -408,6 +408,11 @@ async function run( 'LIBRECHAT_CODE_COMMAND_SANDBOX must be native-srt or runtime', ); } + const nativeProgrammaticEnabled = + allowWorkspaceCommands && + commandSandboxMode === 'native-srt' && + process.platform !== 'win32' && + (fileRelayUpstream?.length ?? 0) > 0; const commandPolicy = resolveNativeSrtCommandPolicy( option(args, '--command-policy-preset') ?? process.env.LIBRECHAT_CODE_COMMAND_POLICY_PRESET?.trim().toLowerCase() ?? @@ -780,6 +785,9 @@ async function run( github.privateKeyPath, ].filter((path): path is string => path != null), allowedDomains: commandAllowedDomains, + ...(nativeProgrammaticEnabled + ? { programmaticFileUpstream: fileRelayUpstream } + : {}), ...(github.provider ? { maskedEnvironment: { @@ -819,6 +827,9 @@ async function run( workspaceTools = new SandboxWorkspaceTools({ workspaceTools, commandWorkspaces: roots.map((root) => root.id), + ...(nativeProgrammaticEnabled + ? { programmaticLanguages: ['bash'] } + : {}), commandSandbox: nativeCommandSandbox ?? new RuntimeWorkspaceCommandSandbox({ @@ -878,6 +889,9 @@ async function run( runtimeSupervisor, capabilities, workspaceTools, + ...(nativeProgrammaticEnabled && nativeCommandSandbox + ? { workspaceProgrammatic: nativeCommandSandbox } + : {}), ...(workspaceLeaseSlots > 1 || roots.length > 1 ? { workspaceQuarantines: new Map( diff --git a/packages/code/src/index.ts b/packages/code/src/index.ts index 5363f0c0..6f94b190 100644 --- a/packages/code/src/index.ts +++ b/packages/code/src/index.ts @@ -7,6 +7,7 @@ export * from './workspace.js'; export * from './workspace-runtime.js'; export * from './native-policy.js'; export * from './native-sandbox.js'; +export * from './native-programmatic.js'; export * from './native-process.js'; export * from './github.js'; export * from './worker.js'; diff --git a/packages/code/src/native-pool.test.ts b/packages/code/src/native-pool.test.ts index 3a24db34..cad30a52 100644 --- a/packages/code/src/native-pool.test.ts +++ b/packages/code/src/native-pool.test.ts @@ -13,6 +13,31 @@ const request = (workspaceId: string): WorkspaceExecuteCommandRequest => ({ workspaceId, command: 'fixture', }); + +test('native pool preflights every registered root with bounded concurrency', async () => { + const prepared: string[] = []; + let active = 0; + let peak = 0; + const pool = new NativeWorkspaceCommandPool(roots, 2, (options) => ({ + async prepare() { + active += 1; + peak = Math.max(peak, active); + await new Promise(resolve => setTimeout(resolve, 5)); + prepared.push(options.workspaceRoot); + active -= 1; + }, + async close() {}, + async execute() { + throw new Error('unreachable'); + }, + })); + + await pool.prepare(); + assert.deepEqual(prepared.sort(), ['/fixture/a', '/fixture/b', '/fixture/c']); + assert.equal(peak, 2); + await pool.close(); +}); + test('a known-clean executor failure is retired without replaying the command', async () => { let created = 0; let executed = 0; diff --git a/packages/code/src/native-pool.ts b/packages/code/src/native-pool.ts index 28f155e2..766409b1 100644 --- a/packages/code/src/native-pool.ts +++ b/packages/code/src/native-pool.ts @@ -2,6 +2,7 @@ import { NativeProcessWorkspaceCommandSandbox } from './native-process.js'; import { WorkspaceToolError } from './workspace.js'; import type { NativeProcessSandboxOptions } from './native-process.js'; import type { + BridgeWorkspaceProgrammaticRequest, WorkspaceExecuteCommandRequest, WorkspaceExecuteCommandResult, } from './protocol.js'; @@ -10,7 +11,10 @@ interface Entry { sandbox: Pick< NativeProcessWorkspaceCommandSandbox, 'prepare' | 'execute' | 'close' - >; + > & + Partial< + Pick + >; busy: boolean; } @@ -92,12 +96,25 @@ export class NativeWorkspaceCommandPool { } async prepare(): Promise { - const entry = await this.allocate(this.roots.keys().next().value!); - try { - await entry.sandbox.prepare(); - } finally { - entry.busy = false; - } + const workspaceIds = [...this.roots.keys()]; + let next = 0; + await Promise.all( + Array.from( + { length: Math.min(this.capacity, workspaceIds.length) }, + async () => { + for (;;) { + const index = next++; + if (index >= workspaceIds.length) return; + const entry = await this.allocate(workspaceIds[index]!); + try { + await entry.sandbox.prepare(); + } finally { + entry.busy = false; + } + } + }, + ), + ); } async execute( @@ -136,6 +153,51 @@ export class NativeWorkspaceCommandPool { } } + async executeProgrammatic( + workspaceId: string, + request: BridgeWorkspaceProgrammaticRequest, + signal?: AbortSignal, + ): Promise { + const entry = await this.allocate(workspaceId); + let enteredExecutor = false; + try { + if (signal?.aborted) + throw new WorkspaceToolError( + 'Programmatic execution cancelled before dispatch', + 'EXECUTION_ABORTED', + ); + enteredExecutor = true; + if (!entry.sandbox.executeProgrammatic) { + throw new WorkspaceToolError( + 'Native programmatic executor is unavailable', + 'COMMAND_UNAVAILABLE', + ); + } + return await entry.sandbox.executeProgrammatic( + workspaceId, + request, + signal, + ); + } catch (error) { + if ( + enteredExecutor && + error instanceof WorkspaceToolError && + !error.mutationMayHaveCommitted + ) { + try { + await entry.sandbox.close(); + if (this.entries.get(workspaceId) === entry) + this.entries.delete(workspaceId); + } catch { + /* Retain ownership for subsequent cleanup/shutdown. */ + } + } + throw error; + } finally { + entry.busy = false; + } + } + async close(): Promise { this.closing = true; await this.allocation; diff --git a/packages/code/src/native-process-child.ts b/packages/code/src/native-process-child.ts index 05ffedfe..2e8beb38 100644 --- a/packages/code/src/native-process-child.ts +++ b/packages/code/src/native-process-child.ts @@ -1,11 +1,16 @@ import { NativeSrtWorkspaceCommandSandbox } from './native-sandbox.js'; +import { NativeWorkspaceProgrammaticExecutor } from './native-programmatic.js'; import { WorkspaceToolError } from './workspace.js'; import type { NativeSrtWorkspaceCommandSandboxOptions } from './native-sandbox.js'; -import type { WorkspaceExecuteCommandRequest } from './protocol.js'; +import type { + BridgeWorkspaceProgrammaticRequest, + WorkspaceExecuteCommandRequest, +} from './protocol.js'; // This entrypoint is private to a forked trusted executor. No HTTP listener, // argv credentials, bridge token, or persisted pairing material is required. let sandbox: NativeSrtWorkspaceCommandSandbox | undefined; +let programmaticExecutor: NativeWorkspaceProgrammaticExecutor | undefined; let active: { id: string; controller: AbortController } | undefined; let busy = false; let credentials: Record = {}; @@ -44,11 +49,14 @@ process.on('message', async (raw: unknown) => { NativeSrtWorkspaceCommandSandboxOptions, 'maskedEnvironment' > & { + programmaticFileUpstream?: string; variables?: NonNullable< NativeSrtWorkspaceCommandSandboxOptions['maskedEnvironment'] >['variables']; }; request: WorkspaceExecuteCommandRequest; + programmaticRequest?: BridgeWorkspaceProgrammaticRequest; + workspaceId?: string; credentials?: Record; wrappedCommand?: string; }; @@ -62,7 +70,8 @@ process.on('message', async (raw: unknown) => { try { let result: unknown; if (message.type === 'prepare' && !sandbox) { - const { variables, ...options } = message.options; + const { variables, programmaticFileUpstream, ...options } = + message.options; sandbox = new NativeSrtWorkspaceCommandSandbox({ ...options, ...(variables @@ -80,11 +89,33 @@ process.on('message', async (raw: unknown) => { : {}), }); await sandbox.prepare(); + programmaticExecutor = programmaticFileUpstream + ? new NativeWorkspaceProgrammaticExecutor({ + sandbox, + upstreamUrl: programmaticFileUpstream, + }) + : undefined; + await programmaticExecutor?.prepare(); } else if (message.type === 'execute' && sandbox) { active = { id: message.id, controller: new AbortController() }; credentials = message.credentials ?? {}; wrappedCommand = message.wrappedCommand; result = await sandbox.execute(message.request, active.controller.signal); + } else if ( + message.type === 'programmatic' && + sandbox && + programmaticExecutor && + message.programmaticRequest && + typeof message.workspaceId === 'string' + ) { + active = { id: message.id, controller: new AbortController() }; + credentials = message.credentials ?? {}; + wrappedCommand = message.wrappedCommand; + result = await programmaticExecutor.execute( + message.programmaticRequest, + message.workspaceId, + active.controller.signal, + ); } else if (message.type === 'close' && sandbox) { await sandbox.close(); } else throw new Error('Invalid executor state'); diff --git a/packages/code/src/native-process.test.ts b/packages/code/src/native-process.test.ts index 800dcf11..e96f1e65 100644 --- a/packages/code/src/native-process.test.ts +++ b/packages/code/src/native-process.test.ts @@ -1,13 +1,29 @@ import assert from 'node:assert/strict'; import { EventEmitter } from 'node:events'; import test from 'node:test'; +import { mkdtemp, rm, symlink, writeFile } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; import type { ChildProcess, ForkOptions } from 'node:child_process'; import { NativeProcessWorkspaceCommandSandbox, nativeExecutorEnvironment, + trustedProgrammaticExecutable, } from './native-process.js'; import { WorkspaceToolError } from './workspace.js'; +test('preflight rejects relative and workspace-controlled executables including symlinks', async t => { + const root = await mkdtemp(join(tmpdir(), 'native-ptc-path-')); + const outside = await mkdtemp(join(tmpdir(), 'native-ptc-link-')); + t.after(async () => { await rm(root, { recursive: true, force: true }); await rm(outside, { recursive: true, force: true }); }); + const executable = join(root, 'bash'); + await writeFile(executable, '#!/bin/sh\nexit 0\n', { mode: 0o700 }); + await symlink(executable, join(outside, 'bash')); + await assert.rejects(trustedProgrammaticExecutable('./bash', root), /absolute/); + await assert.rejects(trustedProgrammaticExecutable(executable, root), /outside the workspace/); + await assert.rejects(trustedProgrammaticExecutable(join(outside, 'bash'), root), /outside the workspace/); +}); + const request = { protocolVersion: 1 as const, operation: 'execute_command' as const, @@ -43,13 +59,19 @@ function fixture( queueMicrotask(() => { if (message.type === 'prepare' && prepare) return prepare(child, message); - if (message.type === 'execute' && execute) + if ( + (message.type === 'execute' || message.type === 'programmatic') && + execute + ) return execute(child, message); if (message.type === 'cancel') return; child.emit('message', { id: message.id, ok: true, - ...(message.type === 'execute' ? { result } : {}), + ...(message.type === 'execute' || + message.type === 'programmatic' + ? { result } + : {}), }); }); return true; @@ -171,8 +193,88 @@ test('executor hands credentials over IPC only for the current command', async ( await sandbox.close(); }); +test('programmatic executor resolves and scopes credentials to its command', async () => { + const fake = fixture(); + const sandbox = new NativeProcessWorkspaceCommandSandbox( + { + workspaceRoot: tmpdir(), + programmaticFileUpstream: 'http://127.0.0.1:3190', + maskedEnvironment: { + variables: [{ name: 'TOKEN', injectHosts: ['github.com'] }], + async resolve() { + return { TOKEN: 'per-programmatic-secret' }; + }, + wrapCommand(command) { + return `wrapped ${command}`; + }, + }, + }, + fake.fork, + ); + const programmaticRequest = { + headers: {}, + body: { + language: 'bash' as const, + version: '5.2.0', + session_id: 'session', + files: [{ name: 'main.sh', content: 'git status' }], + }, + }; + await sandbox.executeProgrammatic('primary', programmaticRequest); + assert.equal( + JSON.stringify(fake.options).includes('per-programmatic-secret'), + false, + ); + const message = fake.messages.find( + candidate => candidate.type === 'programmatic', + )!; + assert.deepEqual(message.credentials, { TOKEN: 'per-programmatic-secret' }); + assert.equal( + message.wrappedCommand, + 'wrapped exec "$LIBRECHAT_CODE_BASH_PATH" "$LIBRECHAT_CODE_DATA_DIR/main.sh"', + ); + await sandbox.close(); +}); + +test('programmatic executor preserves a child-reported pre-dispatch failure', async () => { + const fake = fixture((child, message) => + child.emit('message', { + id: message.id, + ok: false, + code: 'COMMAND_UNAVAILABLE', + errorMessage: 'Programmatic input download failed', + mutation: false, + requiresQuarantine: false, + }), + ); + const sandbox = new NativeProcessWorkspaceCommandSandbox( + { + workspaceRoot: tmpdir(), + programmaticFileUpstream: 'http://127.0.0.1:3190', + }, + fake.fork, + ); + + await assert.rejects( + sandbox.executeProgrammatic('primary', { + headers: {}, + body: { + language: 'bash', + version: '5.2.0', + session_id: 'session', + files: [{ name: 'main.sh', content: 'echo ready' }], + }, + }), + (error: unknown) => + error instanceof WorkspaceToolError && + !error.mutationMayHaveCommitted && + !error.requiresQuarantine, + ); + await sandbox.close(); +}); + test('executor loss after dispatch is an uncertain mutation and is never replayed', async () => { - const fake = fixture((child) => child.emit('exit', 1)); + const fake = fixture(child => child.emit('exit', 1)); const sandbox = new NativeProcessWorkspaceCommandSandbox( { workspaceRoot: '/workspace' }, fake.fork, @@ -180,16 +282,17 @@ test('executor loss after dispatch is an uncertain mutation and is never replaye await assert.rejects( sandbox.execute(request), (error: unknown) => - error instanceof WorkspaceToolError && error.mutationMayHaveCommitted, + error instanceof WorkspaceToolError && + error.mutationMayHaveCommitted, ); await assert.rejects(sandbox.execute(request), /unavailable/); - assert.equal(fake.messages.filter((m) => m.type === 'execute').length, 1); + assert.equal(fake.messages.filter(m => m.type === 'execute').length, 1); await sandbox.close(); }); test('executor cancellation targets the active request and preserves mutation certainty', async () => { let dispatched!: () => void; - const dispatch = new Promise((resolve) => { + const dispatch = new Promise(resolve => { dispatched = resolve; }); const fake = fixture(() => dispatched()); @@ -202,7 +305,7 @@ test('executor cancellation targets the active request and preserves mutation ce await dispatch; await assert.rejects(sandbox.execute(request), /unavailable/); controller.abort(); - const command = fake.messages.find((m) => m.type === 'execute')!; + const command = fake.messages.find(m => m.type === 'execute')!; assert.deepEqual(fake.messages.at(-1), { type: 'cancel', id: command.id }); fake.child.emit('message', { id: command.id, @@ -224,7 +327,7 @@ test('executor cancellation targets the active request and preserves mutation ce test('executor ignores a cleanup exemption on non-cancellation failures', async () => { let dispatched!: () => void; - const dispatch = new Promise((resolve) => { + const dispatch = new Promise(resolve => { dispatched = resolve; }); const fake = fixture(() => dispatched()); @@ -234,7 +337,7 @@ test('executor ignores a cleanup exemption on non-cancellation failures', async ); const execution = sandbox.execute(request); await dispatch; - const command = fake.messages.find((message) => message.type === 'execute')!; + const command = fake.messages.find(message => message.type === 'execute')!; fake.child.emit('message', { id: command.id, ok: false, @@ -269,7 +372,8 @@ test('executor rejects mismatched results as uncertain and fences subsequent com await assert.rejects( sandbox.execute(request), (error: unknown) => - error instanceof WorkspaceToolError && error.mutationMayHaveCommitted, + error instanceof WorkspaceToolError && + error.mutationMayHaveCommitted, ); await assert.rejects(sandbox.execute(request), /unavailable/); await sandbox.close(); @@ -277,7 +381,7 @@ test('executor rejects mismatched results as uncertain and fences subsequent com test('executor close drains an active command before closing IPC', async () => { let dispatched!: () => void; - const dispatch = new Promise((resolve) => { + const dispatch = new Promise(resolve => { dispatched = resolve; }); const fake = fixture(() => dispatched()); @@ -288,16 +392,16 @@ test('executor close drains an active command before closing IPC', async () => { const execution = sandbox.execute(request); await dispatch; const closing = sandbox.close(); - await new Promise((resolve) => setImmediate(resolve)); + await new Promise(resolve => setImmediate(resolve)); assert.equal( - fake.messages.some((m) => m.type === 'close'), + fake.messages.some(m => m.type === 'close'), false, ); - const command = fake.messages.find((m) => m.type === 'execute')!; + const command = fake.messages.find(m => m.type === 'execute')!; fake.child.emit('message', { id: command.id, ok: true, result }); assert.deepEqual(await execution, result); await closing; - assert.equal(fake.messages.filter((m) => m.type === 'close').length, 1); + assert.equal(fake.messages.filter(m => m.type === 'close').length, 1); await assert.rejects(sandbox.execute(request), /unavailable/); }); @@ -321,7 +425,7 @@ test('executor close resolves when the child exits during the close handshake', }, }); await sandbox.close(); - assert.equal(fake.messages.filter((m) => m.type === 'close').length, 1); + assert.equal(fake.messages.filter(m => m.type === 'close').length, 1); await assert.rejects(sandbox.execute(request), /unavailable/); }); @@ -365,7 +469,7 @@ test('executor close skips the handshake once the child is already lost', async fake.child.emit('exit', 1, null); await sandbox.close(); assert.equal( - fake.messages.some((m) => m.type === 'close'), + fake.messages.some(m => m.type === 'close'), false, ); await assert.rejects(sandbox.execute(request), /unavailable/); @@ -377,17 +481,20 @@ test('executor startup loss is not reported as an applied mutation', async () => { workspaceRoot: '/workspace' }, (path, args, options) => { const child = fake.fork(path, args, options); - queueMicrotask(() => child.emit('error', new Error('startup failed'))); + queueMicrotask(() => + child.emit('error', new Error('startup failed')), + ); return child; }, ); await assert.rejects( sandbox.execute(request), (error: unknown) => - error instanceof WorkspaceToolError && !error.mutationMayHaveCommitted, + error instanceof WorkspaceToolError && + !error.mutationMayHaveCommitted, ); assert.equal( - fake.messages.some((m) => m.type === 'execute'), + fake.messages.some(m => m.type === 'execute'), false, ); await sandbox.close(); @@ -409,7 +516,7 @@ test('executor shutdown receipt fences reuse before the OS exit event', async () ); await assert.rejects(sandbox.execute(request)); await assert.rejects(sandbox.execute(request), /unavailable/); - assert.equal(fake.messages.filter((m) => m.type === 'execute').length, 1); + assert.equal(fake.messages.filter(m => m.type === 'execute').length, 1); await sandbox.close(); }); @@ -431,7 +538,8 @@ test('executor preserves bounded startup diagnostics and conventional host setti ok: false, mutation: false, code: 'COMMAND_UNAVAILABLE', - errorMessage: 'Native sandbox dependencies are unavailable: bubblewrap', + errorMessage: + 'Native sandbox dependencies are unavailable: bubblewrap', }), ); const sandbox = new NativeProcessWorkspaceCommandSandbox( @@ -465,7 +573,10 @@ test('executor matches POSIX names exactly and folds names only on Windows', () https_proxy: 'http://proxy:8080', }); assert.deepEqual( - nativeExecutorEnvironment({ Path: 'C:\\bin', Temp: 'C:\\temp' }, 'win32'), + nativeExecutorEnvironment( + { Path: 'C:\\bin', Temp: 'C:\\temp' }, + 'win32', + ), { Path: 'C:\\bin', Temp: 'C:\\temp' }, ); }); @@ -486,7 +597,8 @@ test('executor classifies every pre-dispatch setup failure as mutation-atomic', return {}; }, wrapCommand(command) { - if (failure === 'wrapper') throw new Error('wrapper failed'); + if (failure === 'wrapper') + throw new Error('wrapper failed'); return command; }, }, @@ -503,10 +615,12 @@ test('executor classifies every pre-dispatch setup failure as mutation-atomic', error instanceof WorkspaceToolError && !error.mutationMayHaveCommitted && error.code === - (failure === 'abort' ? 'EXECUTION_ABORTED' : 'COMMAND_UNAVAILABLE'), + (failure === 'abort' + ? 'EXECUTION_ABORTED' + : 'COMMAND_UNAVAILABLE'), ); assert.equal( - fake.messages.some((m) => m.type === 'execute'), + fake.messages.some(m => m.type === 'execute'), false, ); await sandbox.close(); diff --git a/packages/code/src/native-process.ts b/packages/code/src/native-process.ts index e1a02498..d2ea0d72 100644 --- a/packages/code/src/native-process.ts +++ b/packages/code/src/native-process.ts @@ -1,11 +1,21 @@ -import { fork } from 'node:child_process'; +import { execFile, fork } from 'node:child_process'; import { randomUUID } from 'node:crypto'; +import { constants as fsConstants } from 'node:fs'; +import { access, realpath } from 'node:fs/promises'; +import { isAbsolute, join, relative, sep } from 'node:path'; +import { promisify } from 'node:util'; import { WorkspaceToolError } from './workspace.js'; -import { isWorkspaceToolRequest, isWorkspaceToolResult } from './protocol.js'; +import { NATIVE_PROGRAMMATIC_COMMAND } from './native-programmatic.js'; +import { + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILES, + isWorkspaceToolRequest, + isWorkspaceToolResult, +} from './protocol.js'; import type { ChildProcess, ForkOptions } from 'node:child_process'; import type { NativeSrtWorkspaceCommandSandboxOptions } from './native-sandbox.js'; import type { WorkspaceCommandSandbox } from './workspace.js'; import type { + BridgeWorkspaceProgrammaticRequest, WorkspaceExecuteCommandRequest, WorkspaceExecuteCommandResult, } from './protocol.js'; @@ -13,7 +23,84 @@ import type { export type NativeProcessSandboxOptions = Omit< NativeSrtWorkspaceCommandSandboxOptions, 'manager' | 'spawnCommand' | 'platform' ->; +> & { + /** Hardened Code API egress gateway used for execution-scoped files. */ + programmaticFileUpstream?: string; +}; + +const execFileAsync = promisify(execFile); + +async function systemProgrammaticExecutable( + name: string, + workspaceRoot: string, +): Promise { + // Preflight runs outside SRT. Never execute a workspace-controlled PATH + // entry (including cwd, node_modules/.bin, or a symlink to another root). + for (const directory of ['/opt/homebrew/bin', '/usr/local/bin', '/usr/bin', '/bin', '/home/linuxbrew/.linuxbrew/bin']) { + const candidate = join(directory, name); + try { + const canonical = await trustedProgrammaticExecutable(candidate, workspaceRoot); + if (!['/opt/homebrew/', '/usr/local/', '/usr/bin/', '/bin/', '/home/linuxbrew/.linuxbrew/'].some(root => canonical.startsWith(root))) continue; + return canonical; + } catch { + // Continue through the bounded PATH entries. + } + } +} + +export async function trustedProgrammaticExecutable(candidate: string, workspaceRoot: string): Promise { + if (!isAbsolute(candidate)) throw new Error('Programmatic executable must be absolute'); + const [canonical, root] = await Promise.all([realpath(candidate), realpath(workspaceRoot)]); + const path = relative(root, canonical); + if (path === '' || (!isAbsolute(path) && path !== '..' && !path.startsWith(`..${sep}`))) { + throw new Error('Programmatic executable must be outside the workspace'); + } + await access(canonical, fsConstants.X_OK); + return canonical; +} + +async function resolveProgrammaticShell( + options: NativeProcessSandboxOptions, +): Promise { + const environment = options.environment ?? process.env; + const shellPath = + options.shellPath != null + ? await trustedProgrammaticExecutable(options.shellPath, options.workspaceRoot) + : await systemProgrammaticExecutable('bash', options.workspaceRoot); + const jqPath = await systemProgrammaticExecutable('jq', options.workspaceRoot); + if (!shellPath || !jqPath) { + throw new WorkspaceToolError( + 'Native programmatic execution requires trusted host installations of Bash 5.2 or newer and jq', + 'COMMAND_UNAVAILABLE', + ); + } + try { + const [{ stdout: bashVersion }] = await Promise.all([ + execFileAsync(shellPath, ['--version'], { + env: nativeExecutorEnvironment(environment), + timeout: 5_000, + }), + execFileAsync(jqPath, ['--version'], { + env: nativeExecutorEnvironment(environment), + timeout: 5_000, + }), + ]); + const match = /version\s+(\d+)\.(\d+)/i.exec(bashVersion); + if ( + !match || + Number(match[1]) < 5 || + (Number(match[1]) === 5 && Number(match[2]) < 2) + ) { + throw new Error('unsupported Bash version'); + } + } catch { + throw new WorkspaceToolError( + 'Native programmatic execution requires trusted host installations of Bash 5.2 or newer and jq', + 'COMMAND_UNAVAILABLE', + ); + } + return shellPath; +} /** Only OS discovery and conventional proxy settings cross into the executor. * In particular, never inherit NODE_OPTIONS, bridge identity, or app secrets. */ @@ -77,20 +164,22 @@ export function nativeExecutorEnvironment( * deadline, as opposed to a failure the executor reported explicitly. */ class NativeExecutorUnavailableError extends WorkspaceToolError { constructor(mutation: boolean) { - super('Native executor is unavailable', 'COMMAND_UNAVAILABLE', mutation); + super( + 'Native executor is unavailable', + 'COMMAND_UNAVAILABLE', + mutation, + ); this.name = 'NativeExecutorUnavailableError'; } } /** One persistent, process-isolated SRT manager per workspace. No automatic * restart/replay: losing IPC after execution starts is an ambiguous mutation. */ -export class NativeProcessWorkspaceCommandSandbox - implements WorkspaceCommandSandbox -{ +export class NativeProcessWorkspaceCommandSandbox implements WorkspaceCommandSandbox { readonly mutationFailuresAreAtomic = true as const; private child?: ChildProcess; private ready?: Promise; - private active?: Promise; + private active?: Promise; private closing?: Promise; private failed = false; private terminationTimer?: ReturnType; @@ -122,12 +211,17 @@ export class NativeProcessWorkspaceCommandSandbox } private async start(): Promise { + const programmaticShellPath = this.options.programmaticFileUpstream + ? await resolveProgrammaticShell(this.options) + : this.options.shellPath; const child = this.forkExecutor( new URL('./native-process-child.js', import.meta.url), [], { execArgv: [], - env: nativeExecutorEnvironment(this.options.environment ?? process.env), + env: nativeExecutorEnvironment( + this.options.environment ?? process.env, + ), stdio: ['ignore', 'ignore', 'ignore', 'ipc'], serialization: 'json', }, @@ -165,6 +259,8 @@ export class NativeProcessWorkspaceCommandSandbox const processTerminationConfirmed = code === 'EXECUTION_ABORTED' && message.requiresQuarantine === false; + const mutationMayHaveCommitted = + pending.mutation && message.mutation !== false; pending.reject( new WorkspaceToolError( typeof message.errorMessage === 'string' && @@ -172,8 +268,8 @@ export class NativeProcessWorkspaceCommandSandbox ? message.errorMessage : 'Native executor request failed', code, - pending.mutation && message.mutation !== false, - pending.mutation && !processTerminationConfirmed, + mutationMayHaveCommitted, + mutationMayHaveCommitted && !processTerminationConfirmed, ), ); } @@ -192,6 +288,7 @@ export class NativeProcessWorkspaceCommandSandbox allowedDomains, homeDirectory, shellPath, + programmaticFileUpstream, } = this.options; await this.rpc( 'prepare', @@ -202,13 +299,14 @@ export class NativeProcessWorkspaceCommandSandbox protectedPaths, allowedDomains, homeDirectory, - shellPath, + shellPath: programmaticShellPath ?? shellPath, + programmaticFileUpstream, variables: this.options.maskedEnvironment?.variables, }, }, 30_000, false, - ).catch((error) => { + ).catch(error => { this.failed = true; this.terminate(); throw error; @@ -223,7 +321,10 @@ export class NativeProcessWorkspaceCommandSandbox !isWorkspaceToolRequest(request) || request.operation !== 'execute_command' ) { - throw new WorkspaceToolError('Invalid native command', 'INVALID_REQUEST'); + throw new WorkspaceToolError( + 'Invalid native command', + 'INVALID_REQUEST', + ); } if (this.active || this.closing || this.failed) throw this.unavailable(false); @@ -236,12 +337,114 @@ export class NativeProcessWorkspaceCommandSandbox } } + async executeProgrammatic( + workspaceId: string, + request: BridgeWorkspaceProgrammaticRequest, + signal?: AbortSignal, + ): Promise { + if (this.active || this.closing || this.failed) + throw this.unavailable(false); + if (!this.options.programmaticFileUpstream) { + throw new WorkspaceToolError( + 'Native programmatic file transport is unavailable', + 'COMMAND_UNAVAILABLE', + ); + } + const active = this.executeProgrammaticOnce( + request, + workspaceId, + signal, + ); + this.active = active; + try { + return await active; + } finally { + this.active = undefined; + } + } + + private async executeProgrammaticOnce( + request: BridgeWorkspaceProgrammaticRequest, + workspaceId: string, + signal?: AbortSignal, + ): Promise { + if (signal?.aborted) + throw new WorkspaceToolError( + 'Programmatic execution aborted', + 'EXECUTION_ABORTED', + ); + let credentials: Record | undefined; + let wrappedCommand: string | undefined; + try { + await this.prepare(); + if (signal?.aborted) throw new Error('aborted'); + credentials = await this.options.maskedEnvironment?.resolve(signal); + if (signal?.aborted) throw new Error('aborted'); + wrappedCommand = this.options.maskedEnvironment?.wrapCommand?.( + NATIVE_PROGRAMMATIC_COMMAND, + process.platform, + ); + if (signal?.aborted) throw new Error('aborted'); + } catch (error) { + if (signal?.aborted) { + throw new WorkspaceToolError( + 'Programmatic execution aborted', + 'EXECUTION_ABORTED', + ); + } + throw error instanceof WorkspaceToolError + ? new WorkspaceToolError(error.message, error.code, false) + : new WorkspaceToolError( + 'Native programmatic executor setup failed before dispatch', + 'COMMAND_UNAVAILABLE', + ); + } + const result = await this.rpc( + 'programmatic', + { + programmaticRequest: request, + workspaceId, + credentials, + wrappedCommand, + }, + (request.body.run_timeout ?? 30_000) * + ((request.body.replay_tool_count ?? 0) > 0 ? 2 : 1) + + (Math.ceil( + request.body.files.filter(file => 'id' in file).length / 4, + ) + + Math.ceil( + (request.body.max_output_files ?? + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILES) / 4, + )) * + (request.body.transfer_timeout_ms ?? 30_000) + + 5_000, + true, + signal, + ); + if (signal?.aborted) { + throw new WorkspaceToolError( + 'Programmatic execution aborted', + 'EXECUTION_ABORTED', + true, + ); + } + if (typeof result !== 'object' || result === null) { + this.failed = true; + this.terminate(); + throw this.unavailable(true); + } + return result; + } + private async executeOnce( request: WorkspaceExecuteCommandRequest, signal?: AbortSignal, ): Promise { if (signal?.aborted) - throw new WorkspaceToolError('Command aborted', 'EXECUTION_ABORTED'); + throw new WorkspaceToolError( + 'Command aborted', + 'EXECUTION_ABORTED', + ); let credentials: Record | undefined; let wrappedCommand: string | undefined; try { @@ -258,7 +461,10 @@ export class NativeProcessWorkspaceCommandSandbox // No execute RPC has been sent: setup, token refresh and wrapping cannot // have mutated the workspace. Do not quarantine it for setup failures. if (signal?.aborted) - throw new WorkspaceToolError('Command aborted', 'EXECUTION_ABORTED'); + throw new WorkspaceToolError( + 'Command aborted', + 'EXECUTION_ABORTED', + ); throw error instanceof WorkspaceToolError ? new WorkspaceToolError(error.message, error.code, false) : new WorkspaceToolError( @@ -324,7 +530,7 @@ export class NativeProcessWorkspaceCommandSandbox reject(this.unavailable(mutation)); }; try { - child.send({ type, id, ...payload }, (error) => { + child.send({ type, id, ...payload }, error => { if (error) sendFailed(); }); } catch { @@ -353,9 +559,12 @@ export class NativeProcessWorkspaceCommandSandbox await this.ready?.catch(() => undefined); try { if (this.child?.connected && !this.failed) - await this.rpc('close', {}, 10_000, false).catch((error: unknown) => { - if (!(error instanceof NativeExecutorUnavailableError)) throw error; - }); + await this.rpc('close', {}, 10_000, false).catch( + (error: unknown) => { + if (!(error instanceof NativeExecutorUnavailableError)) + throw error; + }, + ); } finally { this.failed = true; this.terminate(); diff --git a/packages/code/src/native-programmatic-live.test.ts b/packages/code/src/native-programmatic-live.test.ts new file mode 100644 index 00000000..2bc22356 --- /dev/null +++ b/packages/code/src/native-programmatic-live.test.ts @@ -0,0 +1,41 @@ +import assert from 'node:assert/strict'; +import { createServer } from 'node:http'; +import { mkdtemp, readFile, rm } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import test from 'node:test'; +import type { AddressInfo } from 'node:net'; +import { NativeProcessWorkspaceCommandSandbox } from './native-process.js'; +import { resolveNativeSrtCommandPolicy } from './native-policy.js'; + +test('real SRT prevents speculative network effects under trusted-vm', { + skip: process.env.LIBRECHAT_CODE_LIVE_SRT_TESTS !== '1', + timeout: 30_000, +}, async () => { + const root = await mkdtemp(join(tmpdir(), 'native-ptc-effects-')); + let effects = 0; + const server = createServer((_req, res) => { effects += 1; res.end('ok'); }); + await new Promise(resolve => server.listen(0, '127.0.0.1', resolve)); + const port = (server.address() as AddressInfo).port; + const executor = new NativeProcessWorkspaceCommandSandbox({ + workspaceRoot: root, + commandPolicy: resolveNativeSrtCommandPolicy('trusted-vm'), + programmaticFileUpstream: `http://127.0.0.1:${port}`, + }); + try { + await executor.prepare(); + const result = await executor.executeProgrammatic('primary', { headers: {}, body: { + language: 'bash', version: '5.2.0', session_id: 'isolated-canary', replay_tool_count: 1, + run_timeout: 5000, + files: [{ name: 'main.sh', content: `curl --noproxy '*' --connect-timeout 1 --max-time 2 -s -X POST http://127.0.0.1:${port}/effect >/dev/null\nprintf once >> commit.txt\n` }], + } }) as { run: { code: number } }; + assert.equal(result.run.code, 0); + assert.equal(effects, 1, 'probe must not emit a network effect'); + assert.equal(await readFile(join(root, 'commit.txt'), 'utf8'), 'once'); + } finally { + try { await executor.close(); } finally { + await new Promise(resolve => server.close(() => resolve())); + await rm(root, { recursive: true, force: true }); + } + } +}); diff --git a/packages/code/src/native-programmatic.test.ts b/packages/code/src/native-programmatic.test.ts new file mode 100644 index 00000000..82650f95 --- /dev/null +++ b/packages/code/src/native-programmatic.test.ts @@ -0,0 +1,608 @@ +import assert from 'node:assert/strict'; +import { createHash } from 'node:crypto'; +import { createServer } from 'node:http'; +import { mkdir, mkdtemp, readFile, readdir, rm, writeFile } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import test from 'node:test'; + +import { NativeWorkspaceProgrammaticExecutor } from './native-programmatic.js'; +import { WorkspaceToolError } from './workspace.js'; + +import type { AddressInfo } from 'node:net'; +import type { BridgeWorkspaceProgrammaticRequest } from './protocol.js'; + +test('stages skill files privately and returns generated artifacts', async () => { + const scratch = await mkdtemp(join(tmpdir(), 'native-ptc-test-')); + const uploads = new Map(); + let downloadCount = 0; + const server = createServer(async (req, res) => { + assert.equal(req.headers['x-codeapi-egress-grant'], 'grant'); + if (req.method === 'GET') { + downloadCount += 1; + assert.match( + req.url ?? '', + /\/sessions\/input-session\/objects\/skill-file$/, + ); + res.end('skill-value'); + return; + } + assert.equal(req.method, 'PUT'); + assert.equal(req.headers['content-type'], 'text/plain'); + const chunks: Buffer[] = []; + for await (const chunk of req) chunks.push(Buffer.from(chunk)); + uploads.set( + decodeURIComponent(req.headers['x-original-filename'] as string), + Buffer.concat(chunks), + ); + res.statusCode = 200; + res.end(); + }); + await new Promise(resolve => server.listen(0, '127.0.0.1', resolve)); + const address = server.address() as AddressInfo; + let observedDataDirectory = ''; + const executor = new NativeWorkspaceProgrammaticExecutor({ + upstreamUrl: `http://127.0.0.1:${address.port}`, + sandbox: { + async createExecutionDirectory() { + return await mkdtemp(join(scratch, 'execution-')); + }, + async executeProgrammatic(request, dataDirectory) { + observedDataDirectory = dataDirectory; + assert.equal( + await readFile( + join(dataDirectory, 'skills/example/reference.txt'), + 'utf8', + ), + 'skill-value', + ); + await writeFile(join(dataDirectory, 'result.txt'), 'artifact'); + return { + protocolVersion: 1, + operation: 'execute_command' as const, + workspaceId: request.workspaceId, + exitCode: 0, + stdout: 'done\n', + stderr: '', + truncated: false, + timedOut: false, + }; + }, + }, + }); + const request: BridgeWorkspaceProgrammaticRequest = { + headers: {}, + body: { + language: 'bash', + version: '5.2.0', + execution_id: 'execution-one', + session_id: 'execution-session', + output_session_id: 'output-session', + egress_grant: 'grant', + files: [ + { name: 'main.sh', content: 'printf done' }, + { name: '_ptc_history.json', content: '{}' }, + { + name: 'skills/example/reference.txt', + id: 'skill-file', + storage_session_id: 'input-session', + input_cache_key: createHash('sha256') + .update('stable-authorized-input-identity') + .digest('hex'), + }, + ], + }, + }; + try { + const result = await executor.execute(request, 'primary'); + const replay = await executor.execute(request, 'primary'); + assert.equal(result.run.stdout, 'done\n'); + assert.equal(replay.run.stdout, 'done\n'); + assert.equal(result.session_id, 'output-session'); + assert.equal(downloadCount, 1); + await executor.execute( + { + ...request, + body: { ...request.body, execution_id: 'execution-two' }, + }, + 'primary', + ); + assert.equal(downloadCount, 2); + assert.equal(result.files.length, 1); + assert.equal(result.files[0]?.name, 'result.txt'); + assert.equal(uploads.get('result.txt')?.toString(), 'artifact'); + assert.deepEqual( + await readdir(observedDataDirectory).catch(() => []), + [], + ); + } finally { + await new Promise(resolve => server.close(() => resolve())); + await rm(scratch, { recursive: true, force: true }); + } +}); + +test('reports unsupported and rejected artifacts without invalidating a completed command', async () => { + const scratch = await mkdtemp(join(tmpdir(), 'native-ptc-artifact-test-')); + const uploads = new Map(); + const server = createServer(async (req, res) => { + const name = decodeURIComponent(req.headers['x-original-filename'] as string); + uploads.set(name, req.headers['content-type']); + for await (const _chunk of req) { + // Drain the bounded request body before responding. + } + res.statusCode = name === 'image.png' ? 503 : 200; + res.end(); + }); + await new Promise(resolve => server.listen(0, '127.0.0.1', resolve)); + const address = server.address() as AddressInfo; + const executor = new NativeWorkspaceProgrammaticExecutor({ + upstreamUrl: `http://127.0.0.1:${address.port}`, + sandbox: { + async createExecutionDirectory() { + return await mkdtemp(join(scratch, 'execution-')); + }, + async executeProgrammatic(_request, dataDirectory) { + await writeFile(join(dataDirectory, '_ptc_report.csv'), 'a,b\n1,2\n'); + await writeFile(join(dataDirectory, 'image.png'), 'not-a-real-png'); + await writeFile(join(dataDirectory, 'model.bin'), 'unsupported'); + return { + protocolVersion: 1, + operation: 'execute_command' as const, + workspaceId: 'primary', + exitCode: 0, + stdout: 'done\n', + stderr: '', + truncated: false, + timedOut: false, + }; + }, + }, + }); + try { + const result = await executor.execute( + { + headers: {}, + body: { + language: 'bash', + version: '5.2.0', + session_id: 'execution-session', + output_session_id: 'output-session', + egress_grant: 'grant', + files: [{ name: 'main.sh', content: 'printf done' }], + }, + }, + 'primary', + ); + assert.equal(result.run.code, 0); + assert.deepEqual(result.files.map(file => file.name), ['_ptc_report.csv']); + assert.deepEqual(result.artifact_delivery, { + code: 'artifact_delivery_failed', + status: 'partial', + attempted: 3, + delivered: 1, + failed: 2, + }); + assert.equal(uploads.get('_ptc_report.csv'), 'text/csv'); + assert.equal(uploads.get('image.png'), 'image/png'); + assert.equal(uploads.has('model.bin'), false); + } finally { + await new Promise(resolve => server.close(() => resolve())); + await rm(scratch, { recursive: true, force: true }); + } +}); + +test('reports artifact transport failure without quarantining a completed command', async () => { + const scratch = await mkdtemp(join(tmpdir(), 'native-ptc-artifact-transport-test-')); + const executor = new NativeWorkspaceProgrammaticExecutor({ + upstreamUrl: 'http://127.0.0.1:1', + fetchImpl: async () => { + throw new TypeError('transport unavailable'); + }, + sandbox: { + async createExecutionDirectory() { + return await mkdtemp(join(scratch, 'execution-')); + }, + async executeProgrammatic(_request, dataDirectory) { + await writeFile(join(dataDirectory, 'result.txt'), 'artifact'); + return { + protocolVersion: 1, + operation: 'execute_command' as const, + workspaceId: 'primary', + exitCode: 0, + stdout: 'done\n', + stderr: '', + truncated: false, + timedOut: false, + }; + }, + }, + }); + try { + const result = await executor.execute( + { + headers: {}, + body: { + language: 'bash', + version: '5.2.0', + session_id: 'execution-session', + output_session_id: 'output-session', + egress_grant: 'grant', + files: [{ name: 'main.sh', content: 'printf done' }], + }, + }, + 'primary', + ); + assert.equal(result.run.code, 0); + assert.deepEqual(result.files, []); + assert.deepEqual(result.artifact_delivery, { + code: 'artifact_delivery_failed', + status: 'failed', + attempted: 1, + delivered: 0, + failed: 1, + }); + } finally { + await rm(scratch, { recursive: true, force: true }); + } +}); + +test('preflights copy-on-write isolation and removes its private snapshot', async () => { + const scratch = await mkdtemp(join(tmpdir(), 'native-ptc-preflight-test-')); + let executionDirectory = ''; + let probes = 0; + const executor = new NativeWorkspaceProgrammaticExecutor({ + upstreamUrl: 'http://127.0.0.1:1', + sandbox: { + async createExecutionDirectory() { + executionDirectory = await mkdtemp(join(scratch, 'execution-')); + return executionDirectory; + }, + async createProgrammaticProbeWorkspace(directory) { + probes += 1; + const workspace = join(directory, 'workspace'); + await mkdir(workspace); + return workspace; + }, + async executeProgrammatic() { + throw new Error('unreachable'); + }, + }, + }); + try { + await executor.prepare(); + assert.equal(probes, 1); + assert.deepEqual(await readdir(executionDirectory).catch(() => []), []); + } finally { + await rm(scratch, { recursive: true, force: true }); + } +}); + +test('keeps replay probes read-only and commits the script exactly once', async () => { + const scratch = await mkdtemp(join(tmpdir(), 'native-ptc-probe-test-')); + const phases: boolean[] = []; + const executor = new NativeWorkspaceProgrammaticExecutor({ + upstreamUrl: 'http://127.0.0.1:1', + sandbox: { + async createExecutionDirectory() { + return await mkdtemp(join(scratch, 'execution-')); + }, + async createProgrammaticProbeWorkspace(executionDirectory) { + const workspace = join(executionDirectory, 'workspace'); + await mkdir(workspace); + return workspace; + }, + async executeProgrammatic( + _request, + dataDirectory, + _signal, + options, + ) { + phases.push(options?.probe === true); + if (options?.probe === true) { + assert.match(options.workspaceRoot ?? '', /\/workspace$/); + } + return { + protocolVersion: 1, + operation: 'execute_command' as const, + workspaceId: 'primary', + exitCode: options?.probe ? 1 : 0, + stdout: options?.probe ? 'probe\n' : 'commit\n', + stderr: options?.probe ? 'expected probe denial\n' : '', + truncated: false, + timedOut: false, + }; + }, + }, + }); + try { + const result = await executor.execute( + { + headers: {}, + body: { + language: 'bash', + version: '5.2.0', + execution_id: 'probe-then-commit', + replay_tool_count: 1, + session_id: 'execution-session', + files: [ + { name: 'main.sh', content: 'printf done' }, + { name: '_ptc_history.json', content: '{}' }, + ], + }, + }, + 'primary', + ); + assert.deepEqual(phases, [true, false]); + assert.equal(result.run.stdout, 'commit\n'); + } finally { + await rm(scratch, { recursive: true, force: true }); + } +}); + +test('returns pending calls from the private control file even when stdout truncates', async () => { + const scratch = await mkdtemp(join(tmpdir(), 'native-ptc-control-test-')); + let phases = 0; + const executor = new NativeWorkspaceProgrammaticExecutor({ + upstreamUrl: 'http://127.0.0.1:1', + sandbox: { + async createExecutionDirectory() { + return await mkdtemp(join(scratch, 'execution-')); + }, + async createProgrammaticProbeWorkspace(executionDirectory) { + const workspace = join(executionDirectory, 'workspace'); + await mkdir(workspace); + return workspace; + }, + async executeProgrammatic( + _request, + dataDirectory, + _signal, + options, + ) { + assert.equal(options?.probe, true); + phases += 1; + await writeFile( + join(dataDirectory, '_ptc_pending_result.json'), + JSON.stringify({ + pending: [ + { + call_id: 'call_001', + tool_name: 'lookup', + input: {}, + }, + ], + }), + ); + return { + protocolVersion: 1, + operation: 'execute_command' as const, + workspaceId: 'primary', + exitCode: 0, + stdout: 'x'.repeat(256 * 1024), + stderr: '', + truncated: true, + timedOut: false, + }; + }, + }, + }); + try { + const result = await executor.execute( + { + headers: {}, + body: { + language: 'bash', + version: '5.2.0', + execution_id: 'truncated-control', + replay_tool_count: 1, + session_id: 'execution-session', + files: [ + { name: 'main.sh', content: 'lookup "{}"' }, + { name: '_ptc_history.json', content: '{}' }, + ], + }, + }, + 'primary', + ); + assert.equal(phases, 1); + assert.equal(result.run.stdout, ''); + assert.equal(result.run.stderr, ''); + assert.deepEqual(JSON.parse(result.pending_tool_calls_payload ?? ''), { + pending: [{ call_id: 'call_001', tool_name: 'lookup', input: {} }], + }); + } finally { + await rm(scratch, { recursive: true, force: true }); + } +}); + +test('rejects traversal before creating execution state', async () => { + let allocated = false; + const executor = new NativeWorkspaceProgrammaticExecutor({ + upstreamUrl: 'http://127.0.0.1:1', + sandbox: { + async createExecutionDirectory() { + allocated = true; + return '/unused'; + }, + async executeProgrammatic() { + throw new Error('unreachable'); + }, + }, + }); + await assert.rejects( + executor.execute( + { + headers: {}, + body: { + language: 'bash', + version: '5.2.0', + session_id: 'execution-session', + files: [{ name: '../main.sh', content: 'echo unsafe' }], + }, + }, + 'primary', + ), + /Invalid selected-workspace programmatic request/, + ); + assert.equal(allocated, false); +}); + +test('rejects artifacts above the negotiated byte ceiling before upload', async () => { + const scratch = await mkdtemp(join(tmpdir(), 'native-ptc-output-limit-')); + let uploads = 0; + const executor = new NativeWorkspaceProgrammaticExecutor({ + upstreamUrl: 'http://127.0.0.1:1', + fetchImpl: async () => { + uploads += 1; + return new Response(); + }, + sandbox: { + async createExecutionDirectory() { + return await mkdtemp(join(scratch, 'execution-')); + }, + async executeProgrammatic(request, dataDirectory) { + await writeFile(join(dataDirectory, 'artifact.txt'), 'too large'); + return { + protocolVersion: 1, + operation: 'execute_command' as const, + workspaceId: request.workspaceId, + exitCode: 0, + stdout: '', + stderr: '', + truncated: false, + timedOut: false, + }; + }, + }, + }); + try { + await assert.rejects( + executor.execute( + { + headers: {}, + body: { + language: 'bash', + version: '5.2.0', + session_id: 'execution-session', + output_session_id: 'output-session', + egress_grant: 'grant', + max_output_file_bytes: 4, + files: [{ name: 'main.sh', content: 'printf done' }], + }, + }, + 'primary', + ), + /exceeds the file limit/, + ); + assert.equal(uploads, 0); + } finally { + await rm(scratch, { recursive: true, force: true }); + } +}); + +for (const failure of ['truncated', 'process-error']) test(`a failed speculative probe does not quarantine the real workspace (${failure})`, async t => { + const scratch = await mkdtemp(join(tmpdir(), 'native-ptc-failed-probe-')); + t.after(() => rm(scratch, { recursive: true, force: true })); + const executor = new NativeWorkspaceProgrammaticExecutor({ + upstreamUrl: 'http://127.0.0.1:1', + sandbox: { + async createExecutionDirectory() { return await mkdtemp(join(scratch, 'execution-')); }, + async createProgrammaticProbeWorkspace(directory) { const root = join(directory, 'workspace'); await mkdir(root); return root; }, + async executeProgrammatic(request, _directory, _signal, options) { + assert.equal(options?.probe, true); + if (failure === 'process-error') throw new WorkspaceToolError('probe output exceeded its limit', 'COMMAND_UNAVAILABLE', true, true); + return { protocolVersion: 1, operation: 'execute_command', workspaceId: request.workspaceId, + exitCode: 0, stdout: '', stderr: '', truncated: true, timedOut: false }; + }, + }, + }); + await assert.rejects(executor.execute({ headers: {}, body: { + language: 'bash', version: '5.2.0', session_id: 'session', replay_tool_count: 1, + files: [{ name: 'main.sh', content: 'true' }], + } }, 'primary'), (error: unknown) => { + assert.match(String(error), /probe output exceeded/); + assert.equal((error as { requiresQuarantine: boolean }).requiresQuarantine, false); + assert.equal((error as { mutationMayHaveCommitted: boolean }).mutationMayHaveCommitted, false); + return true; + }); +}); + +test('unchanged inputs do not consume the negotiated output budget', async t => { + const scratch = await mkdtemp(join(tmpdir(), 'native-ptc-unchanged-')); + t.after(() => rm(scratch, { recursive: true, force: true })); + const executor = new NativeWorkspaceProgrammaticExecutor({ + upstreamUrl: 'http://127.0.0.1:1', + sandbox: { + async createExecutionDirectory() { return await mkdtemp(join(scratch, 'execution-')); }, + async executeProgrammatic(request) { + return { protocolVersion: 1, operation: 'execute_command', workspaceId: request.workspaceId, + exitCode: 0, stdout: '', stderr: '', truncated: false, timedOut: false }; + }, + }, + }); + const result = await executor.execute({ headers: {}, body: { + language: 'bash', version: '5.2.0', session_id: 'session', max_output_file_bytes: 1, + files: [{ name: 'main.sh', content: 'true' }, { name: 'input.txt', content: 'unchanged input' }], + } }, 'primary'); + assert.deepEqual(result.files, []); +}); + +test('stops admitting downloads and drains in-flight transfers before cleanup', async () => { + const scratch = await mkdtemp(join(tmpdir(), 'native-ptc-transfer-test-')); + let executionDirectory = ''; + let requestCount = 0; + const server = createServer((req, res) => { + requestCount += 1; + if (requestCount === 1) { + res.statusCode = 503; + res.end(); + return; + } + setTimeout(() => res.end('in-flight'), 25); + }); + await new Promise(resolve => server.listen(0, '127.0.0.1', resolve)); + const address = server.address() as AddressInfo; + const executor = new NativeWorkspaceProgrammaticExecutor({ + upstreamUrl: `http://127.0.0.1:${address.port}`, + sandbox: { + async createExecutionDirectory() { + executionDirectory = await mkdtemp(join(scratch, 'execution-')); + return executionDirectory; + }, + async executeProgrammatic() { + throw new Error('unreachable'); + }, + }, + }); + const request: BridgeWorkspaceProgrammaticRequest = { + headers: {}, + body: { + language: 'bash', + version: '5.2.0', + session_id: 'execution-session', + output_session_id: 'output-session', + egress_grant: 'grant', + files: [ + ...Array.from({ length: 8 }, (_, index) => ({ + name: `inputs/${index}.txt`, + id: `input-${index}`, + storage_session_id: 'input-session', + })), + { name: 'main.sh', content: 'printf done' }, + ], + }, + }; + try { + const startedAt = performance.now(); + await assert.rejects( + executor.execute(request, 'primary'), + /Programmatic input download failed with HTTP 503/, + ); + assert.ok(performance.now() - startedAt >= 20); + assert.ok(requestCount <= 4); + assert.deepEqual(await readdir(executionDirectory).catch(() => []), []); + } finally { + await new Promise(resolve => server.close(() => resolve())); + await rm(scratch, { recursive: true, force: true }); + } +}); diff --git a/packages/code/src/native-programmatic.ts b/packages/code/src/native-programmatic.ts new file mode 100644 index 00000000..42974b79 --- /dev/null +++ b/packages/code/src/native-programmatic.ts @@ -0,0 +1,816 @@ +import { createHash, randomBytes } from 'node:crypto'; +import { constants } from 'node:fs'; +import { cp, mkdir, open, readdir, rm, writeFile } from 'node:fs/promises'; +import { dirname, join, relative, sep } from 'node:path'; + +import { + BRIDGE_PROTOCOL_VERSION, + BRIDGE_WORKSPACE_COMMAND_DEFAULT_OUTPUT_BYTES, + BRIDGE_WORKSPACE_COMMAND_MAX_TIMEOUT_MS, + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILES, + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILE_BYTES, + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_TOTAL_BYTES, + bridgeArtifactMediaType, + isBridgeWorkspaceProgrammaticRequest, + isSafePortableRelativePath, + isSupportedBridgeArtifactName, +} from './protocol.js'; +import { validateFileRelayUpstream } from './relay.js'; +import { WorkspaceToolError } from './workspace.js'; + +import type { + BridgeProgrammaticPayloadFile, + BridgeWorkspaceProgrammaticRequest, + WorkspaceExecuteCommandResult, +} from './protocol.js'; +import type { NativeSrtWorkspaceCommandSandbox } from './native-sandbox.js'; + +const EGRESS_GRANT_HEADER = 'X-CodeAPI-Egress-Grant'; +const EXECUTION_MAIN_FILE = 'main.sh'; +const EXECUTION_HISTORY_FILE = '_ptc_history.json'; +const EXECUTION_CONTROL_FILE = '_ptc_pending_result.json'; +export const NATIVE_PROGRAMMATIC_COMMAND = + 'exec "$LIBRECHAT_CODE_BASH_PATH" "$LIBRECHAT_CODE_DATA_DIR/main.sh"'; +const TRANSFER_TIMEOUT_MS = 30_000; +const TRANSFER_CONCURRENCY = 4; +const MAX_WALK_ENTRIES = 2_000; +const INPUT_CACHE_MAX_ENTRIES = 64; +const INPUT_CACHE_MAX_BYTES = 32 * 1024 * 1024; +const CONTROL_PAYLOAD_MAX_BYTES = 512 * 1024; + +type ProgrammaticFileResult = { + id: string; + name: string; + storage_session_id: string; + modified_from?: { id: string; storage_session_id: string }; +}; + +type ProgrammaticResult = { + language: 'bash'; + version: string; + session_id: string; + files: ProgrammaticFileResult[]; + artifact_delivery?: { + code: 'artifact_delivery_failed'; + status: 'partial' | 'failed'; + attempted: number; + delivered: number; + failed: number; + }; + pending_tool_calls_payload?: string; + run: { + stdout: string; + stderr: string; + code: number | null; + signal: string | null; + output: string; + memory: null; + message: string | null; + status: string | null; + cpu_time: null; + wall_time: number; + }; +}; + +type InputBaseline = { + sha256: string; + source?: { id: string; storage_session_id: string }; +}; + +function sha256(value: Uint8Array): string { + return createHash('sha256').update(value).digest('hex'); +} + +function outputFileId(): string { + return randomBytes(18).toString('base64url').slice(0, 21); +} + +function localPath(root: string, name: string): string { + if (!isSafePortableRelativePath(name)) { + throw new WorkspaceToolError( + 'Invalid programmatic file path', + 'INVALID_PATH', + ); + } + const path = join(root, ...name.split('/')); + const child = relative(root, path); + if (child === '' || child === '..' || child.startsWith(`..${sep}`)) { + throw new WorkspaceToolError( + 'Invalid programmatic file path', + 'INVALID_PATH', + ); + } + return path; +} + +async function readBoundedResponse( + response: Response, + signal: AbortSignal, +): Promise { + const declaredLength = response.headers.get('content-length'); + if ( + declaredLength != null && + (!/^\d+$/.test(declaredLength) || + Number(declaredLength) > + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILE_BYTES) + ) { + await response.body?.cancel(); + throw new WorkspaceToolError( + 'Programmatic input exceeds the file limit', + 'READ_LIMIT_EXCEEDED', + ); + } + if (!response.body) return Buffer.alloc(0); + const reader = response.body.getReader(); + const chunks: Buffer[] = []; + let bytes = 0; + try { + for (;;) { + if (signal.aborted) throw signal.reason; + const { done, value } = await reader.read(); + if (done) break; + bytes += value.byteLength; + if (bytes > BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILE_BYTES) { + throw new WorkspaceToolError( + 'Programmatic input exceeds the file limit', + 'READ_LIMIT_EXCEEDED', + ); + } + chunks.push(Buffer.from(value)); + } + } finally { + await reader.cancel().catch(() => undefined); + } + return Buffer.concat(chunks, bytes); +} + +async function mapConcurrent( + values: readonly T[], + concurrency: number, + action: (value: T) => Promise, +): Promise { + const results = new Array(values.length); + let next = 0; + let failed = false; + let failure: unknown; + await Promise.all( + Array.from( + { length: Math.min(concurrency, values.length) }, + async () => { + for (;;) { + if (failed) return; + const index = next++; + if (index >= values.length) return; + try { + results[index] = await action(values[index]!); + } catch (error) { + if (!failed) { + failed = true; + failure = error; + } + return; + } + } + }, + ), + ); + if (failed) throw failure; + return results; +} + +async function listRegularFiles(root: string): Promise { + const files: string[] = []; + const pending = ['']; + let entries = 0; + while (pending.length > 0) { + const directory = pending.pop()!; + for (const entry of await readdir(join(root, directory), { + withFileTypes: true, + })) { + if (++entries > MAX_WALK_ENTRIES) { + throw new WorkspaceToolError( + 'Programmatic output contains too many entries', + 'WRITE_LIMIT_EXCEEDED', + ); + } + const name = directory ? `${directory}/${entry.name}` : entry.name; + if (!isSafePortableRelativePath(name)) continue; + if (entry.isSymbolicLink()) continue; + if (entry.isDirectory()) pending.push(name); + else if (entry.isFile()) files.push(name); + } + } + return files.sort(); +} + +export interface NativeWorkspaceProgrammaticOptions { + sandbox: Pick< + NativeSrtWorkspaceCommandSandbox, + 'createExecutionDirectory' | 'executeProgrammatic' + > & + Partial< + Pick + >; + upstreamUrl: string; + fetchImpl?: typeof fetch; +} + +/** + * Executes one replay-mode Bash PTC iteration in an attached workspace. + * Program code and injected files live in a private SRT scratch directory; + * the selected repository remains the command cwd and is never used as a + * transport cache. + */ +export class NativeWorkspaceProgrammaticExecutor { + private readonly upstream: URL; + private readonly fetchImpl: typeof fetch; + /** Parent-process cache: sandboxed children cannot inspect this memory. */ + private readonly inputCache = new Map< + string, + { bytes: Buffer; lastUsed: number } + >(); + private inputCacheBytes = 0; + + constructor(private readonly options: NativeWorkspaceProgrammaticOptions) { + this.upstream = validateFileRelayUpstream(options.upstreamUrl); + this.fetchImpl = options.fetchImpl ?? fetch; + } + + /** + * Prove copy-on-write isolation before the worker advertises Bash PTC. + * The probe uses the exact registered root and private scratch path that a + * real replay will use, then removes the snapshot before registration. + */ + async prepare(signal?: AbortSignal): Promise { + const createProbeWorkspace = + this.options.sandbox.createProgrammaticProbeWorkspace; + if (createProbeWorkspace == null) { + throw new WorkspaceToolError( + 'Selected-workspace PTC probe isolation is unavailable', + 'COMMAND_UNAVAILABLE', + ); + } + const executionDirectory = + await this.options.sandbox.createExecutionDirectory(); + try { + await createProbeWorkspace.call( + this.options.sandbox, + executionDirectory, + signal, + ); + } finally { + await rm(executionDirectory, { recursive: true, force: true }); + } + } + + private cacheKey( + executionId: string | undefined, + file: Extract, + ): string | undefined { + return executionId && file.input_cache_key + ? `${executionId}:${file.input_cache_key}` + : undefined; + } + + private cachedInput(key: string): Buffer | undefined { + const cached = this.inputCache.get(key); + if (!cached) return undefined; + cached.lastUsed = Date.now(); + return cached.bytes; + } + + private cacheInput(key: string, bytes: Buffer): void { + if (bytes.byteLength > INPUT_CACHE_MAX_BYTES) return; + const existing = this.inputCache.get(key); + if (existing) this.inputCacheBytes -= existing.bytes.byteLength; + while ( + this.inputCache.size >= INPUT_CACHE_MAX_ENTRIES || + this.inputCacheBytes + bytes.byteLength > INPUT_CACHE_MAX_BYTES + ) { + let oldestKey: string | undefined; + let oldestAt = Number.POSITIVE_INFINITY; + for (const [candidate, value] of this.inputCache) { + if (value.lastUsed < oldestAt) { + oldestAt = value.lastUsed; + oldestKey = candidate; + } + } + if (!oldestKey) break; + this.inputCacheBytes -= + this.inputCache.get(oldestKey)!.bytes.byteLength; + this.inputCache.delete(oldestKey); + } + this.inputCache.set(key, { bytes, lastUsed: Date.now() }); + this.inputCacheBytes += bytes.byteLength; + } + + private async downloadInput( + file: Extract, + grant: string, + executionId: string | undefined, + signal?: AbortSignal, + transferTimeoutMs = TRANSFER_TIMEOUT_MS, + ): Promise { + const key = this.cacheKey(executionId, file); + const cached = key ? this.cachedInput(key) : undefined; + if (cached) return cached; + const controller = new AbortController(); + const abort = (): void => controller.abort(signal?.reason); + signal?.addEventListener('abort', abort, { once: true }); + const timer = setTimeout(() => controller.abort(), transferTimeoutMs); + try { + const response = await this.fetchImpl( + new URL( + `sessions/${encodeURIComponent(file.storage_session_id)}/objects/${encodeURIComponent(file.id)}`, + `${this.upstream.toString().replace(/\/+$/, '')}/`, + ), + { + headers: { [EGRESS_GRANT_HEADER]: grant }, + redirect: 'error', + signal: controller.signal, + }, + ); + if (!response.ok) { + await response.body?.cancel(); + throw new WorkspaceToolError( + `Programmatic input download failed with HTTP ${response.status}`, + 'COMMAND_UNAVAILABLE', + ); + } + const bytes = await readBoundedResponse( + response, + controller.signal, + ); + if (key) this.cacheInput(key, bytes); + return bytes; + } finally { + clearTimeout(timer); + signal?.removeEventListener('abort', abort); + } + } + + async execute( + request: BridgeWorkspaceProgrammaticRequest, + workspaceId: string, + signal?: AbortSignal, + ): Promise { + if (!isBridgeWorkspaceProgrammaticRequest(request)) { + throw new WorkspaceToolError( + 'Invalid selected-workspace programmatic request', + 'INVALID_REQUEST', + ); + } + if (signal?.aborted) { + throw new WorkspaceToolError( + 'Programmatic execution aborted', + 'EXECUTION_ABORTED', + ); + } + const grant = request.body.egress_grant; + const refFiles = request.body.files.filter( + ( + file, + ): file is Extract => + 'id' in file, + ); + if (refFiles.length > 0 && !grant) { + throw new WorkspaceToolError( + 'Programmatic input grant is unavailable', + 'INVALID_REQUEST', + ); + } + const executionDirectory = + await this.options.sandbox.createExecutionDirectory(); + const inputDirectory = join(executionDirectory, 'inputs'); + let dataDirectory = join(executionDirectory, 'final'); + const baselines = new Map(); + let totalInputBytes = 0; + const startedAt = performance.now(); + let commandDispatched = false; + try { + await mkdir(inputDirectory, { mode: 0o700 }); + await mapConcurrent( + request.body.files, + TRANSFER_CONCURRENCY, + async (file): Promise => { + const bytes = + 'content' in file + ? Buffer.from(file.content) + : await this.downloadInput( + file, + grant!, + request.body.execution_id, + signal, + request.body.transfer_timeout_ms, + ); + totalInputBytes += bytes.byteLength; + if ( + totalInputBytes > + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_TOTAL_BYTES + ) { + throw new WorkspaceToolError( + 'Programmatic inputs exceed the total byte limit', + 'READ_LIMIT_EXCEEDED', + ); + } + const path = localPath(inputDirectory, file.name); + await mkdir(dirname(path), { + recursive: true, + mode: 0o700, + }); + await writeFile(path, bytes, { flag: 'wx', mode: 0o600 }); + baselines.set(file.name, { + sha256: sha256(bytes), + ...('id' in file + ? { + source: { + id: file.id, + storage_session_id: + file.storage_session_id, + }, + } + : {}), + }); + }, + ); + + const run = async ( + directory: string, + probe: boolean, + workspaceRoot?: string, + ): Promise => { + await cp(inputDirectory, directory, { + recursive: true, + force: false, + errorOnExist: true, + mode: constants.COPYFILE_FICLONE, + }); + if (!probe) commandDispatched = true; + return await this.options.sandbox.executeProgrammatic( + { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operation: 'execute_command', + workspaceId, + command: NATIVE_PROGRAMMATIC_COMMAND, + timeoutMs: Math.min( + request.body.run_timeout ?? + BRIDGE_WORKSPACE_COMMAND_MAX_TIMEOUT_MS, + BRIDGE_WORKSPACE_COMMAND_MAX_TIMEOUT_MS, + ), + maxOutputBytes: + BRIDGE_WORKSPACE_COMMAND_DEFAULT_OUTPUT_BYTES, + }, + directory, + signal, + { probe, workspaceRoot }, + ); + }; + + const readPending = async ( + directory: string, + ): Promise => { + try { + const path = join(directory, EXECUTION_CONTROL_FILE); + const handle = await open( + path, + constants.O_RDONLY | constants.O_NOFOLLOW, + ); + try { + const metadata = await handle.stat(); + if ( + !metadata.isFile() || + metadata.size === 0 || + metadata.size > CONTROL_PAYLOAD_MAX_BYTES + ) { + throw new WorkspaceToolError( + 'Native programmatic control frame is invalid', + 'COMMAND_UNAVAILABLE', + ); + } + return await handle.readFile('utf8'); + } finally { + await handle.close(); + } + } catch (error) { + if ((error as NodeJS.ErrnoException).code === 'ENOENT') + return; + throw error; + } + }; + + if ((request.body.replay_tool_count ?? 0) > 0) { + const probeDirectory = join(executionDirectory, 'probe'); + const createProbeWorkspace = + this.options.sandbox.createProgrammaticProbeWorkspace; + if (createProbeWorkspace == null) { + throw new WorkspaceToolError( + 'Selected-workspace PTC probe isolation is unavailable', + 'COMMAND_UNAVAILABLE', + ); + } + const probeWorkspace = await createProbeWorkspace.call( + this.options.sandbox, + executionDirectory, + signal, + ); + const probeResult = await run( + probeDirectory, + true, + probeWorkspace, + ); + const pending = await readPending(probeDirectory); + if (pending) { + return this.result( + request, + { + ...probeResult, + /** Probe output is speculative and the script will + * run once under its real policy after tool + * resolution. Never duplicate it or expose + * expected read-only policy denials to callers. */ + stdout: '', + stderr: '', + }, + [], + performance.now() - startedAt, + pending, + ); + } + if (probeResult.truncated) { + throw new WorkspaceToolError( + 'Native programmatic probe output exceeded its limit', + 'WRITE_LIMIT_EXCEEDED', + ); + } + if ( + probeResult.timedOut || + probeResult.signal + ) { + return this.result( + request, + probeResult, + [], + performance.now() - startedAt, + ); + } + /** A read-only probe commonly exits non-zero after it reaches + * an intentional workspace write denial. With no pending call, + * run the script once under its real policy so ordinary writes + * and their resulting exit status are evaluated exactly once. */ + } + + const commandResult = await run(dataDirectory, false); + if (await readPending(dataDirectory)) { + throw new WorkspaceToolError( + 'Native programmatic commit pass requested an unexpected replay tool', + 'COMMAND_UNAVAILABLE', + true, + true, + ); + } + if (commandResult.truncated) { + throw new WorkspaceToolError( + 'Native programmatic output exceeded its limit', + 'WRITE_LIMIT_EXCEEDED', + true, + true, + ); + } + + const outputSessionId = request.body.output_session_id; + const outputNames = (await listRegularFiles(dataDirectory)).filter( + name => + name !== EXECUTION_MAIN_FILE && + name !== EXECUTION_HISTORY_FILE && + name !== EXECUTION_CONTROL_FILE && + !name.startsWith('skills/'), + ); + const changed: Array<{ + name: string; + bytes: Buffer; + source?: { id: string; storage_session_id: string }; + }> = []; + let totalOutputBytes = 0; + for (const name of outputNames) { + const path = localPath(dataDirectory, name); + const baseline = baselines.get(name); + const maxOutputFileBytes = Math.min( + request.body.max_output_file_bytes ?? BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILE_BYTES, + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILE_BYTES, + ); + let bytes: Buffer; + const handle = await open( + path, + constants.O_RDONLY | constants.O_NOFOLLOW, + ); + try { + const metadata = await handle.stat(); + if (!metadata.isFile()) continue; + // An unchanged input is not an output. It may legitimately + // exceed the negotiated output ceiling, but never the + // protocol's bounded input limit. + if (metadata.size > (baseline ? BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILE_BYTES : maxOutputFileBytes)) { + throw new WorkspaceToolError( + 'Programmatic output exceeds the file limit', + 'WRITE_LIMIT_EXCEEDED', + ); + } + bytes = await handle.readFile(); + } finally { + await handle.close(); + } + if (baseline?.sha256 === sha256(bytes)) continue; + if (bytes.byteLength > maxOutputFileBytes) { + throw new WorkspaceToolError( + 'Programmatic output exceeds the file limit', + 'WRITE_LIMIT_EXCEEDED', + ); + } + totalOutputBytes += bytes.byteLength; + if ( + totalOutputBytes > + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_TOTAL_BYTES + ) { + throw new WorkspaceToolError( + 'Programmatic outputs exceed the total byte limit', + 'WRITE_LIMIT_EXCEEDED', + ); + } + changed.push({ name, bytes, source: baseline?.source }); + } + const maxOutputFiles = Math.min( + request.body.max_output_files ?? + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILES, + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILES, + ); + if (changed.length > maxOutputFiles) { + throw new WorkspaceToolError( + 'Programmatic output contains too many files', + 'WRITE_LIMIT_EXCEEDED', + ); + } + const uploadable = changed.filter(({ name }) => + isSupportedBridgeArtifactName(name), + ); + if (uploadable.length > 0 && (!grant || !outputSessionId)) { + throw new WorkspaceToolError( + 'Programmatic output grant is unavailable', + 'COMMAND_UNAVAILABLE', + ); + } + const uploadResults = await mapConcurrent( + uploadable, + TRANSFER_CONCURRENCY, + async ({ name, bytes, source }): Promise => { + const id = outputFileId(); + const controller = new AbortController(); + const abort = (): void => controller.abort(signal?.reason); + signal?.addEventListener('abort', abort, { once: true }); + const timer = setTimeout( + () => controller.abort(), + request.body.transfer_timeout_ms ?? TRANSFER_TIMEOUT_MS, + ); + try { + let response: Response; + try { + response = await this.fetchImpl( + new URL( + `sessions/${encodeURIComponent(outputSessionId!)}/objects/${id}`, + `${this.upstream.toString().replace(/\/+$/, '')}/`, + ), + { + method: 'PUT', + headers: { + [EGRESS_GRANT_HEADER]: grant!, + 'Content-Type': bridgeArtifactMediaType(name), + 'Content-Length': String(bytes.byteLength), + 'X-Original-Filename': encodeURIComponent(name), + }, + body: new Uint8Array(bytes), + redirect: 'error', + signal: controller.signal, + }, + ); + } catch (error) { + if (signal?.aborted) throw error; + return undefined; + } + await response.body?.cancel(); + if (!response.ok) { + return undefined; + } + return { + id, + name, + storage_session_id: outputSessionId!, + ...(source ? { modified_from: source } : {}), + }; + } finally { + clearTimeout(timer); + signal?.removeEventListener('abort', abort); + } + }, + ); + const files = uploadResults.filter( + (file): file is ProgrammaticFileResult => file != null, + ); + const artifactDelivery = + files.length < changed.length + ? { + code: 'artifact_delivery_failed' as const, + status: files.length > 0 ? ('partial' as const) : ('failed' as const), + attempted: changed.length, + delivered: files.length, + failed: changed.length - files.length, + } + : undefined; + return this.result( + request, + commandResult, + files, + performance.now() - startedAt, + undefined, + artifactDelivery, + ); + } catch (error) { + if (!commandDispatched) { + // The low-level command runner classifies any launched process as a + // possible mutation. A probe can only mutate its disposable snapshot, + // so translate that classification at this ownership boundary. + throw new WorkspaceToolError( + error instanceof Error ? error.message : 'Programmatic preparation failed', + error instanceof WorkspaceToolError ? error.code : 'COMMAND_UNAVAILABLE', + false, + false, + ); + } + if (error instanceof WorkspaceToolError) { + if ( + error.mutationMayHaveCommitted || + error.requiresQuarantine + ) { + throw error; + } + throw new WorkspaceToolError( + error.message, + error.code, + true, + true, + ); + } + throw new WorkspaceToolError( + 'Native programmatic execution failed after dispatch', + 'COMMAND_UNAVAILABLE', + true, + true, + ); + } finally { + try { + await rm(executionDirectory, { recursive: true, force: true }); + } catch { + throw new WorkspaceToolError( + 'Native programmatic execution cleanup failed', + 'COMMAND_UNAVAILABLE', + commandDispatched, + commandDispatched, + ); + } + } + } + + private result( + request: BridgeWorkspaceProgrammaticRequest, + command: WorkspaceExecuteCommandResult, + files: ProgrammaticFileResult[], + elapsedMs: number, + pendingToolCallsPayload?: string, + artifactDelivery?: ProgrammaticResult['artifact_delivery'], + ): ProgrammaticResult { + return { + language: 'bash', + version: request.body.version, + // Code API masks the execution session separately from the writable + // output bucket. Sandbox results must identify the output bucket so the + // gateway can restore it to the caller-owned session after upload. + session_id: + request.body.output_session_id ?? request.body.session_id, + files, + ...(artifactDelivery ? { artifact_delivery: artifactDelivery } : {}), + ...(pendingToolCallsPayload + ? { pending_tool_calls_payload: pendingToolCallsPayload } + : {}), + run: { + stdout: command.stdout, + stderr: command.stderr, + code: command.exitCode, + signal: command.signal ?? null, + output: `${command.stdout}${command.stderr}`, + memory: null, + message: command.timedOut ? 'Execution timed out' : null, + status: command.timedOut ? 'timeout' : null, + cpu_time: null, + wall_time: elapsedMs / 1000, + }, + }; + } +} diff --git a/packages/code/src/native-sandbox.test.ts b/packages/code/src/native-sandbox.test.ts index cf790965..3252ee50 100644 --- a/packages/code/src/native-sandbox.test.ts +++ b/packages/code/src/native-sandbox.test.ts @@ -7,6 +7,7 @@ import { mkdtemp, mkdir, open, + readFile, realpath, rename, rm, @@ -25,7 +26,10 @@ import type { } from '@anthropic-ai/sandbox-runtime'; import type { ChildProcessWithoutNullStreams } from 'node:child_process'; -import { NativeSrtWorkspaceCommandSandbox } from './native-sandbox.js'; +import { + CopyOnWriteCloneUnavailableError, + NativeSrtWorkspaceCommandSandbox, +} from './native-sandbox.js'; import { restoreScratchTraversal } from './native-scratch.js'; import { WorkspaceToolError } from './workspace.js'; @@ -54,6 +58,8 @@ function fakeManager( let credentialSeenDuringWrap: string | undefined; let gitLfsRequiredSeenDuringWrap: string | undefined; let scratchSelectorSeenDuringWrap: string | undefined; + let networkSeenDuringWrap: SandboxRuntimeConfig['network'] | undefined; + let customConfigSeenDuringWrap: Partial | undefined; const manager = { isSupportedPlatform: () => true, async checkDependenciesAsync() { @@ -67,19 +73,30 @@ function fakeManager( askCallback = callback; if (options.initializeError) throw options.initializeError; }, - async wrapWithSandboxArgv(command: string) { + updateConfig(value: SandboxRuntimeConfig) { config = value; }, + async wrapWithSandboxArgv( + command: string, + _binShell?: string, + customConfig?: Partial, + ) { await options.beforeWrap?.(); - credentialSeenDuringWrap = process.env.LIBRECHAT_CODE_TEST_CREDENTIAL; + networkSeenDuringWrap = config?.network; + customConfigSeenDuringWrap = customConfig; + credentialSeenDuringWrap = + process.env.LIBRECHAT_CODE_TEST_CREDENTIAL; gitLfsRequiredSeenDuringWrap = process.env.GIT_CONFIG_VALUE_3; scratchSelectorSeenDuringWrap = process.env.CLAUDE_CODE_TMPDIR; const ambientGitEnvironment = Object.fromEntries( Object.entries(process.env).filter( - ([name, value]) => name.startsWith('GIT_CONFIG_') && value != null, + ([name, value]) => + name.startsWith('GIT_CONFIG_') && value != null, ), ); let gitEnvironment = ambientGitEnvironment; if (options.appendGitSafeDirectory) { - const index = Number(ambientGitEnvironment.GIT_CONFIG_COUNT ?? '0'); + const index = Number( + ambientGitEnvironment.GIT_CONFIG_COUNT ?? '0', + ); gitEnvironment = { ...(options.inheritedGitEnvironment ?? {}), GIT_CONFIG_COUNT: String(index + 1), @@ -130,10 +147,118 @@ function fakeManager( get scratchSelectorSeenDuringWrap() { return scratchSelectorSeenDuringWrap; }, + get networkSeenDuringWrap() { return networkSeenDuringWrap; }, + get customConfigSeenDuringWrap() { + return customConfigSeenDuringWrap; + }, }; } -test('exclusive lifecycle rejects a second workspace sharing an SRT manager', async (t) => { +for (const trustedVm of [false, true]) test(`programmatic probe denies real-workspace writes and external effects (trusted=${trustedVm})`, 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, + manager: fake.manager, + allowedDomains: ['api.example.com'], + ...(trustedVm ? { commandPolicy: { version: 1 as const, preset: 'trusted-vm' as const, + network: { outbound: 'unrestricted' as const, allowLocalBinding: true, allowAllUnixSockets: true }, + } } : {}), + }); + const dataDirectory = await sandbox.createExecutionDirectory(); + await sandbox.executeProgrammatic(request, dataDirectory, undefined, { + probe: true, + }); + assert.deepEqual(fake.customConfigSeenDuringWrap?.network, { + allowedDomains: [], + deniedDomains: [], + strictAllowlist: true, + allowUnixSockets: [], + allowAllUnixSockets: false, + allowLocalBinding: false, + }); + assert.deepEqual(fake.networkSeenDuringWrap, fake.customConfigSeenDuringWrap?.network); + assert.equal(fake.config?.network.strictAllowlist, !trustedVm); + assert.equal(fake.reset, true, 'probe proxy session must be revoked before restoring policy'); + assert.deepEqual(fake.customConfigSeenDuringWrap?.filesystem?.allowWrite, [ + await realpath(dataDirectory), + ]); + assert.equal( + fake.scratchSelectorSeenDuringWrap, + await realpath(dataDirectory), + ); + assert.ok( + fake.customConfigSeenDuringWrap?.filesystem?.denyWrite?.includes( + await realpath(root), + ), + ); + await sandbox.close(); +}); + +test('probe network cleanup failure fences executor reuse', async t => { + const root = await mkdtemp(join(tmpdir(), 'librechat-probe-cleanup-')); + t.after(() => rm(root, { recursive: true, force: true })); + const fake = fakeManager(); + const sandbox = new NativeSrtWorkspaceCommandSandbox({ workspaceRoot: root, manager: fake.manager }); + const directory = await sandbox.createExecutionDirectory(); + const reset = fake.manager.reset; + fake.manager.reset = async () => { throw new Error('proxy shutdown failed'); }; + await assert.rejects(sandbox.executeProgrammatic(request, directory, undefined, { probe: true }), /probe network cleanup failed/); + await assert.rejects(sandbox.execute(request)); + fake.manager.reset = reset; + await sandbox.close(); +}); + +test('programmatic probes use a copy-on-write workspace without mutating the project', async t => { + const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); + t.after(() => rm(root, { recursive: true, force: true })); + await writeFile(join(root, 'state.txt'), 'original'); + const fake = fakeManager(); + const sandbox = new NativeSrtWorkspaceCommandSandbox({ + workspaceRoot: root, + manager: fake.manager, + }); + t.after(() => sandbox.close()); + const executionDirectory = await sandbox.createExecutionDirectory(); + let snapshot: string; + try { + snapshot = await sandbox.createProgrammaticProbeWorkspace(executionDirectory); + } catch (error) { + if (error instanceof CopyOnWriteCloneUnavailableError) { + t.skip('host filesystem does not support copy-on-write cloning'); + return; + } + throw error; + } + await writeFile(join(snapshot, 'state.txt'), 'probe-only'); + assert.equal(await readFile(join(root, 'state.txt'), 'utf8'), 'original'); + assert.equal(await readFile(join(snapshot, 'state.txt'), 'utf8'), 'probe-only'); +}); + +test('programmatic probes do not hide clone implementation failures as unsupported filesystems', async t => { + const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); + t.after(() => rm(root, { recursive: true, force: true })); + const sandbox = new NativeSrtWorkspaceCommandSandbox({ + workspaceRoot: root, + manager: fakeManager().manager, + spawnCommand() { + throw Object.assign(new Error('spawn /bin/cp ENOENT'), { code: 'ENOENT' }); + }, + }); + t.after(() => sandbox.close()); + const executionDirectory = await sandbox.createExecutionDirectory(); + + await assert.rejects( + sandbox.createProgrammaticProbeWorkspace(executionDirectory), + (error: unknown) => + error instanceof WorkspaceToolError && + !(error instanceof CopyOnWriteCloneUnavailableError) && + error.message === 'Copy-on-write workspace clone failed unexpectedly', + ); +}); + +test('exclusive lifecycle rejects a second workspace sharing an SRT manager', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const fake = fakeManager(); @@ -164,15 +289,15 @@ test('exclusive lifecycle rejects a second workspace sharing an SRT manager', as assert.equal((await second.execute(request)).stdout, 'hello'); }); -test('exclusive lifecycle rejects overlapping commands and waits before resetting', async (t) => { +test('exclusive lifecycle rejects overlapping commands and waits before resetting', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); let entered!: () => void; - const wrapping = new Promise((resolve) => { + const wrapping = new Promise(resolve => { entered = resolve; }); let release!: () => void; - const gate = new Promise((resolve) => { + const gate = new Promise(resolve => { release = resolve; }); const fake = fakeManager({ @@ -191,7 +316,7 @@ test('exclusive lifecycle rejects overlapping commands and waits before resettin await assert.rejects(sandbox.execute(request), /active command/); const closing = sandbox.close(); const secondClose = sandbox.close(); - await new Promise((resolve) => setImmediate(resolve)); + await new Promise(resolve => setImmediate(resolve)); assert.equal(fake.reset, false); await assert.rejects(sandbox.prepare(), /closing/); release(); @@ -200,7 +325,7 @@ test('exclusive lifecycle rejects overlapping commands and waits before resettin assert.equal(fake.reset, true); }); -test('exclusive lifecycle retains ownership after a failed reset until cleanup succeeds', async (t) => { +test('exclusive lifecycle retains ownership after a failed reset until cleanup succeeds', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const fake = fakeManager(); @@ -226,16 +351,16 @@ test('exclusive lifecycle retains ownership after a failed reset until cleanup s await second.close(); }); -test('exclusive lifecycle waits for initialization before resetting the manager', async (t) => { +test('exclusive lifecycle waits for initialization before resetting the manager', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const fake = fakeManager(); let entered!: () => void; - const initializing = new Promise((resolve) => { + const initializing = new Promise(resolve => { entered = resolve; }); let release!: () => void; - const gate = new Promise((resolve) => { + const gate = new Promise(resolve => { release = resolve; }); fake.manager.initialize = async () => { @@ -249,7 +374,7 @@ test('exclusive lifecycle waits for initialization before resetting the manager' const preparing = sandbox.prepare(); await initializing; const closing = sandbox.close(); - await new Promise((resolve) => setImmediate(resolve)); + await new Promise(resolve => setImmediate(resolve)); assert.equal(fake.reset, false); release(); await preparing; @@ -257,7 +382,7 @@ test('exclusive lifecycle waits for initialization before resetting the manager' assert.equal(fake.reset, true); }); -test('initializes SRT with a default-deny network and scrubbed worker credentials', async (t) => { +test('initializes SRT with a default-deny network and scrubbed worker credentials', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); const identity = join(tmpdir(), 'librechat-code-identity.json'); t.after(() => rm(root, { recursive: true, force: true })); @@ -301,7 +426,7 @@ test('initializes SRT with a default-deny network and scrubbed worker credential assert.ok(fake.config?.filesystem.denyRead.includes(canonicalHome)); assert.ok(fake.config?.filesystem.denyWrite.includes(canonicalIdentity)); assert.ok( - fake.config?.filesystem.denyWrite.some((path) => + fake.config?.filesystem.denyWrite.some(path => path.endsWith('/tmp/claude'), ), ); @@ -318,7 +443,7 @@ 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) => { +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(); @@ -351,7 +476,7 @@ test('trusted-vm permits unmatched egress and local development sockets', async ]); }); -test('provides an isolated scratch directory to commands and restores the host environment', async (t) => { +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 })); const originalTmpdir = process.env.TMPDIR; @@ -380,7 +505,7 @@ test('provides an isolated scratch directory to commands and restores the host e await assert.rejects(access(result.stdout)); }); -test('removes scratch storage when SRT initialization fails', async (t) => { +test('removes scratch storage when SRT initialization fails', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const fake = fakeManager({ initializeError: new Error('init failed') }); @@ -396,7 +521,7 @@ test('removes scratch storage when SRT initialization fails', async (t) => { assert.equal(fake.reset, true); }); -test('rejects workspaces nested inside SRT shared scratch storage', async (t) => { +test('rejects workspaces nested inside SRT shared scratch storage', async t => { if (process.platform === 'win32') return; const sharedRoot = '/tmp/claude'; await mkdir(sharedRoot, { recursive: true }); @@ -431,17 +556,17 @@ test('rejects a workspace that contains worker scratch storage', async () => { await sandbox.close(); }); -test('keeps concurrent sandbox scratch directories independent', async (t) => { +test('keeps concurrent sandbox scratch directories independent', async t => { const firstRoot = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); const secondRoot = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(firstRoot, { recursive: true, force: true })); t.after(() => rm(secondRoot, { recursive: true, force: true })); let releaseWrap!: () => void; let wrapStarted!: () => void; - const wrapStartedPromise = new Promise((resolve) => { + const wrapStartedPromise = new Promise(resolve => { wrapStarted = resolve; }); - const holdWrap = new Promise((resolve) => { + const holdWrap = new Promise(resolve => { releaseWrap = resolve; }); const firstFake = fakeManager({ @@ -478,7 +603,7 @@ test('keeps concurrent sandbox scratch directories independent', async (t) => { await secondSandbox.close(); }); -test('removes scratch storage after a command revokes traversal permissions', async (t) => { +test('removes scratch storage after a command revokes traversal permissions', async t => { if (process.platform === 'win32') return; const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); @@ -498,7 +623,7 @@ test('removes scratch storage after a command revokes traversal permissions', as await assert.rejects(access(result.stdout)); }); -test('scratch traversal never follows a descendant replaced after inspection', async (t) => { +test('scratch traversal never follows a descendant replaced after inspection', async t => { if (process.platform === 'win32') return; const root = await mkdtemp(join(tmpdir(), 'librechat-code-scratch-race-')); const outside = await mkdtemp(join(tmpdir(), 'librechat-code-outside-')); @@ -529,7 +654,7 @@ test('scratch traversal never follows a descendant replaced after inspection', a assert.equal((await stat(outsideChild)).mode & 0o777, 0o711); }); -test('scratch traversal removes command-created Darwin ACLs', async (t) => { +test('scratch traversal removes command-created Darwin ACLs', async t => { if (process.platform !== 'darwin') return; const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); @@ -548,7 +673,7 @@ test('scratch traversal removes command-created Darwin ACLs', async (t) => { await assert.rejects(access(result.stdout)); }); -test('scratch traversal bounds descriptors and work across a deep tree', async (t) => { +test('scratch traversal bounds descriptors and work across a deep tree', async t => { if (process.platform === 'win32') return; const root = await mkdtemp(join(tmpdir(), 'librechat-code-scratch-depth-')); t.after(() => rm(root, { recursive: true, force: true })); @@ -565,12 +690,17 @@ test('scratch traversal bounds descriptors and work across a deep tree', async ( await restoreScratchTraversal(rootHandle); - assert.equal((await stat(directories[directories.length - 1])).mode & 0o777, 0o700); + assert.equal( + (await stat(directories[directories.length - 1])).mode & 0o777, + 0o700, + ); }); -test('scratch traversal rejects trees beyond its recovery depth limit', async (t) => { +test('scratch traversal rejects trees beyond its recovery depth limit', async t => { if (process.platform === 'win32') return; - const root = await mkdtemp(join(tmpdir(), 'librechat-code-scratch-depth-limit-')); + const root = await mkdtemp( + join(tmpdir(), 'librechat-code-scratch-depth-limit-'), + ); t.after(() => rm(root, { recursive: true, force: true })); let directory = root; for (let depth = 0; depth < 129; depth += 1) { @@ -586,7 +716,7 @@ test('scratch traversal rejects trees beyond its recovery depth limit', async (t ); }); -test('does not replace scratch state while cleanup remains pending', async (t) => { +test('does not replace scratch state while cleanup remains pending', async t => { if (process.platform === 'win32') return; const workspace = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); const retained = await mkdtemp(join(tmpdir(), 'librechat-code-retained-')); @@ -642,7 +772,7 @@ const windowsEnvironment = { }; for (const platform of ['darwin', 'linux', 'win32'] as const) { - test(`preserves required ${platform} environment names without allowing credentials`, async (t) => { + test(`preserves required ${platform} environment names without allowing credentials`, async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const fake = fakeManager(); @@ -658,23 +788,43 @@ for (const platform of ['darwin', 'linux', 'win32'] as const) { LD_PRELOAD: '/host/private.so', }; const sandbox = new NativeSrtWorkspaceCommandSandbox({ - workspaceRoot: root, platform, allowedDomains: ['github.com'], + workspaceRoot: root, + platform, + allowedDomains: ['github.com'], environment: { - ...proxyEnvironment, ...windowsEnvironment, ...credentials, + ...proxyEnvironment, + ...windowsEnvironment, + ...credentials, HtTp_PrOxY: 'http://mixed-case.invalid:8080', - PATH: '/usr/bin', LC_ALL: 'C.UTF-8', + PATH: '/usr/bin', + LC_ALL: 'C.UTF-8', }, manager: fake.manager, }); t.after(() => sandbox.close()); await sandbox.prepare(); - const denied = new Set(fake.config?.credentials?.envVars - ?.filter(({ mode }) => mode === 'deny').map(({ name }) => name)); - for (const name of [...Object.keys(proxyEnvironment), 'PATH', 'LC_ALL']) { - assert.equal(denied.has(name), false, `${name} must remain available`); + const denied = new Set( + fake.config?.credentials?.envVars + ?.filter(({ mode }) => mode === 'deny') + .map(({ name }) => name), + ); + for (const name of [ + ...Object.keys(proxyEnvironment), + 'PATH', + 'LC_ALL', + ]) { + assert.equal( + denied.has(name), + false, + `${name} must remain available`, + ); } for (const name of Object.keys(windowsEnvironment)) { - assert.equal(denied.has(name), platform !== 'win32', `${name} must be platform-specific`); + assert.equal( + denied.has(name), + platform !== 'win32', + `${name} must be platform-specific`, + ); } assert.equal(denied.has('HtTp_PrOxY'), platform !== 'win32'); for (const name of Object.keys(credentials)) { @@ -685,12 +835,14 @@ for (const platform of ['darwin', 'linux', 'win32'] as const) { }); } -test('uses SRT proxy values without restoring inherited proxies or credentials', async (t) => { +test('uses SRT proxy values without restoring inherited proxies or credentials', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const wrappedEnvironment = { - HTTP_PROXY: 'http://localhost:3128', HTTPS_PROXY: 'http://localhost:3128', - ALL_PROXY: 'http://localhost:3128', NO_PROXY: 'localhost', + HTTP_PROXY: 'http://localhost:3128', + HTTPS_PROXY: 'http://localhost:3128', + ALL_PROXY: 'http://localhost:3128', + NO_PROXY: 'localhost', }; const fake = fakeManager({ wrappedEnvironment }); const sandbox = new NativeSrtWorkspaceCommandSandbox({ @@ -700,14 +852,19 @@ test('uses SRT proxy values without restoring inherited proxies or credentials', }); t.after(() => sandbox.close()); const result = await sandbox.execute({ - ...request, maxOutputBytes: 256, - command: 'printf "%s|%s|%s|%s|%s" "$HTTP_PROXY" "$HTTPS_PROXY" "$ALL_PROXY" "$NO_PROXY" "${GITHUB_TOKEN-unset}"', + ...request, + maxOutputBytes: 256, + command: + 'printf "%s|%s|%s|%s|%s" "$HTTP_PROXY" "$HTTPS_PROXY" "$ALL_PROXY" "$NO_PROXY" "${GITHUB_TOKEN-unset}"', }); assert.equal(result.exitCode, 0); - assert.equal(result.stdout, `${Object.values(wrappedEnvironment).join('|')}|unset`); + assert.equal( + result.stdout, + `${Object.values(wrappedEnvironment).join('|')}|unset`, + ); }); -test('masks a host credential for only its injection host and restores the parent environment', async (t) => { +test('masks a host credential for only its injection host and restores the parent environment', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const fake = fakeManager(); @@ -731,7 +888,8 @@ test('masks a host credential for only its injection host and restores the paren ], async resolve() { return { - LIBRECHAT_CODE_TEST_CREDENTIAL: 'Authorization: Bearer real-secret', + LIBRECHAT_CODE_TEST_CREDENTIAL: + 'Authorization: Bearer real-secret', }; }, }, @@ -759,7 +917,7 @@ test('masks a host credential for only its injection host and restores the paren }); }); -test('serializes credential handoff across concurrent sandbox instances', async (t) => { +test('serializes credential handoff across concurrent sandbox instances', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const original = process.env.LIBRECHAT_CODE_TEST_CREDENTIAL; @@ -770,11 +928,11 @@ test('serializes credential handoff across concurrent sandbox instances', async else process.env.LIBRECHAT_CODE_TEST_CREDENTIAL = original; }); let firstEntered!: () => void; - const firstEnteredPromise = new Promise((resolve) => { + const firstEnteredPromise = new Promise(resolve => { firstEntered = resolve; }); let releaseFirst!: () => void; - const firstGate = new Promise((resolve) => { + const firstGate = new Promise(resolve => { releaseFirst = resolve; }); let secondEntered = false; @@ -817,7 +975,7 @@ test('serializes credential handoff across concurrent sandbox instances', async const secondExecution = sandbox(second.manager, 'second-secret').execute( request, ); - await new Promise((resolve) => setImmediate(resolve)); + await new Promise(resolve => setImmediate(resolve)); assert.equal(secondEntered, false); releaseFirst(); await firstExecution; @@ -828,7 +986,7 @@ test('serializes credential handoff across concurrent sandbox instances', async assert.equal(process.env.LIBRECHAT_CODE_TEST_CREDENTIAL, undefined); }); -test('isolates Git from host-level global and system configuration', async (t) => { +test('isolates Git from host-level global and system configuration', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const sandbox = new NativeSrtWorkspaceCommandSandbox({ @@ -844,7 +1002,7 @@ test('isolates Git from host-level global and system configuration', async (t) = assert.equal(result.stdout, '/dev/null|1'); }); -test('restores trusted Git LFS filters without reading host Git configuration', async (t) => { +test('restores trusted Git LFS filters without reading host Git configuration', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const fake = fakeManager({ @@ -884,7 +1042,7 @@ test('restores trusted Git LFS filters without reading host Git configuration', assert.ok(!denied?.includes('GIT_CONFIG_VALUE_0')); }); -test('filters environment names case-insensitively only on Windows', async (t) => { +test('filters environment names case-insensitively only on Windows', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const fake = fakeManager(); @@ -923,7 +1081,7 @@ test('filters environment names case-insensitively only on Windows', async (t) = assert.ok(!denied?.includes('git_config_count')); }); -test('fails closed when the configured POSIX shell is unavailable', async (t) => { +test('fails closed when the configured POSIX shell is unavailable', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const sandbox = new NativeSrtWorkspaceCommandSandbox({ @@ -942,7 +1100,7 @@ test('fails closed when the configured POSIX shell is unavailable', async (t) => ); }); -test('fails closed when SRT dependencies are unavailable', async (t) => { +test('fails closed when SRT dependencies are unavailable', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const fake = fakeManager({ dependencyErrors: ['bubblewrap missing'] }); @@ -960,7 +1118,7 @@ test('fails closed when SRT dependencies are unavailable', async (t) => { ); }); -test('refuses workspace roots that expose worker home or control files', async (t) => { +test('refuses workspace roots that expose worker home or control files', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); const controlDirectory = join(root, '.control'); await mkdir(controlDirectory); @@ -985,7 +1143,7 @@ test('refuses workspace roots that expose worker home or control files', async ( ); }); -test('executes in the canonical workspace and bounds aggregate output', async (t) => { +test('executes in the canonical workspace and bounds aggregate output', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); await mkdir(join(root, 'src')); t.after(() => rm(root, { recursive: true, force: true })); @@ -1015,7 +1173,7 @@ test('executes in the canonical workspace and bounds aggregate output', async (t ); }); -test('rejects an escaping or unavailable command working directory', async (t) => { +test('rejects an escaping or unavailable command working directory', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const sandbox = new NativeSrtWorkspaceCommandSandbox({ @@ -1026,11 +1184,12 @@ test('rejects an escaping or unavailable command working directory', async (t) = await assert.rejects( sandbox.execute({ ...request, cwd: '..' }), (error: unknown) => - error instanceof WorkspaceToolError && error.code === 'INVALID_REQUEST', + error instanceof WorkspaceToolError && + error.code === 'INVALID_REQUEST', ); }); -test('terminates detached command descendants before returning', async (t) => { +test('terminates detached command descendants before returning', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const sandbox = new NativeSrtWorkspaceCommandSandbox({ @@ -1043,15 +1202,15 @@ test('terminates detached command descendants before returning', async (t) => { command: '(sleep 0.2; printf late > late.txt) >/dev/null 2>&1 &', }); assert.equal(result.exitCode, 0); - await new Promise((resolve) => setTimeout(resolve, 350)); + await new Promise(resolve => setTimeout(resolve, 350)); await assert.rejects(access(join(root, 'late.txt'))); }); -test('reports cancellation after command start as a potentially committed mutation', async (t) => { +test('reports cancellation after command start as a potentially committed mutation', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); let commandStarted!: () => void; - const commandStartedPromise = new Promise((resolve) => { + const commandStartedPromise = new Promise(resolve => { commandStarted = resolve; }); const sandbox = new NativeSrtWorkspaceCommandSandbox({ @@ -1081,7 +1240,7 @@ test('reports cancellation after command start as a potentially committed mutati ); }); -test('closes stdin immediately when the command protocol provides no input', async (t) => { +test('closes stdin immediately when the command protocol provides no input', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const sandbox = new NativeSrtWorkspaceCommandSandbox({ @@ -1098,7 +1257,7 @@ test('closes stdin immediately when the command protocol provides no input', asy assert.equal(result.timedOut, false); }); -test('maps platform-native exit statuses into the bridge protocol range', async (t) => { +test('maps platform-native exit statuses into the bridge protocol range', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const spawnCommand = () => { @@ -1123,7 +1282,7 @@ test('maps platform-native exit statuses into the bridge protocol range', async assert.equal(result.exitCode, 1); }); -test('cleans allocated command state exactly once on every execution exit', async (t) => { +test('cleans allocated command state exactly once on every execution exit', async t => { for (const outcome of [ 'abort-before-spawn', 'spawn-throw', @@ -1134,8 +1293,12 @@ test('cleans allocated command state exactly once on every execution exit', asyn 'wrap-throw', ] as const) { for (const cleanupThrows of [false, true]) { - await t.test(`${outcome}, cleanup throws: ${cleanupThrows}`, async (t) => { - const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); + await t.test( + `${outcome}, cleanup throws: ${cleanupThrows}`, + async t => { + const root = await mkdtemp( + join(tmpdir(), 'librechat-code-native-'), + ); t.after(() => rm(root, { recursive: true, force: true })); const controller = new AbortController(); let cleanupCalls = 0; @@ -1143,9 +1306,11 @@ test('cleans allocated command state exactly once on every execution exit', asyn let allocated = false; const fake = fakeManager({ async beforeWrap() { - if (outcome === 'wrap-throw') throw new Error('wrap failed'); + if (outcome === 'wrap-throw') + throw new Error('wrap failed'); allocated = true; - if (outcome === 'abort-before-spawn') controller.abort(); + if (outcome === 'abort-before-spawn') + controller.abort(); }, }); fake.manager.cleanupAfterCommand = () => { @@ -1160,13 +1325,17 @@ test('cleans allocated command state exactly once on every execution exit', asyn spawnCommand() { spawnCalls += 1; assert.equal(allocated, true); - if (outcome === 'spawn-throw') throw new Error('spawn failed'); - const child = new EventEmitter() as ChildProcessWithoutNullStreams; + if (outcome === 'spawn-throw') + throw new Error('spawn failed'); + const child = + new EventEmitter() as ChildProcessWithoutNullStreams; let closeQueued = false; const close = () => { if (!closeQueued) { closeQueued = true; - queueMicrotask(() => child.emit('close', null, 'SIGKILL')); + queueMicrotask(() => + child.emit('close', null, 'SIGKILL'), + ); } return true; }; @@ -1180,7 +1349,10 @@ test('cleans allocated command state exactly once on every execution exit', asyn queueMicrotask(() => { assert.equal(cleanupCalls, 0); if (outcome === 'error') { - child.emit('error', new Error('spawn failed')); + child.emit( + 'error', + new Error('spawn failed'), + ); } else if (outcome === 'abort-after-spawn') { controller.abort(); } else if (outcome === 'close') { @@ -1196,28 +1368,43 @@ test('cleans allocated command state exactly once on every execution exit', asyn ); if (outcome === 'close' || outcome === 'timeout') { const result = await execution; - assert.equal(result.exitCode, outcome === 'close' ? 0 : null); + assert.equal( + result.exitCode, + outcome === 'close' ? 0 : null, + ); assert.equal(result.timedOut, outcome === 'timeout'); } else { - await assert.rejects(execution, (error: unknown) => + await assert.rejects( + execution, + (error: unknown) => error instanceof WorkspaceToolError && - error.code === (outcome.startsWith('abort') + error.code === + (outcome.startsWith('abort') ? 'EXECUTION_ABORTED' : 'COMMAND_UNAVAILABLE') && - error.mutationMayHaveCommitted === (outcome === 'abort-after-spawn') && + error.mutationMayHaveCommitted === + (outcome === 'abort-after-spawn') && error.requiresQuarantine === - (outcome === 'abort-after-spawn' && process.platform === 'win32'), + (outcome === 'abort-after-spawn' && + process.platform === 'win32'), ); } assert.equal( spawnCalls, - outcome === 'abort-before-spawn' || outcome === 'wrap-throw' ? 0 : 1, + outcome === 'abort-before-spawn' || + outcome === 'wrap-throw' + ? 0 + : 1, + ); + assert.equal( + cleanupCalls, + outcome === 'wrap-throw' ? 0 : 1, ); - assert.equal(cleanupCalls, outcome === 'wrap-throw' ? 0 : 1); assert.equal(allocated, false); await sandbox.close(); assert.equal(fake.reset, true); - }); + }, + ); } } }); diff --git a/packages/code/src/native-sandbox.ts b/packages/code/src/native-sandbox.ts index 171ea663..8d150d9f 100644 --- a/packages/code/src/native-sandbox.ts +++ b/packages/code/src/native-sandbox.ts @@ -11,14 +11,7 @@ import { sep, } from 'node:path'; import { constants as fsConstants } from 'node:fs'; -import { - access, - mkdtemp, - open, - realpath, - rm, - stat, -} from 'node:fs/promises'; +import { access, mkdtemp, open, realpath, rm, stat } from 'node:fs/promises'; import type { FileHandle } from 'node:fs/promises'; import { SandboxManager } from '@anthropic-ai/sandbox-runtime'; @@ -141,8 +134,12 @@ interface NativeSandboxManager { cwd?: string, options?: { commandId?: string; commandText?: string }, ): Promise<{ argv: string[]; env: NodeJS.ProcessEnv }>; - annotateStderrWithSandboxFailures(commandId: string, stderr: string): string; + annotateStderrWithSandboxFailures( + commandId: string, + stderr: string, + ): string; cleanupAfterCommand(): void; + updateConfig?(config: SandboxRuntimeConfig): void; reset(): Promise; } @@ -224,13 +221,16 @@ function deniedEnvironmentNames( platform: NodeJS.Platform, ): string[] { return Object.keys(environment) - .filter((name) => { + .filter(name => { const normalized = platform === 'win32' ? name.toUpperCase() : name; return ( normalized.startsWith('LIBRECHAT_CODE_') || (!SAFE_CHILD_ENV_NAMES.has(normalized) && !PROXY_CHILD_ENV_NAMES.has(normalized) && - !(platform === 'win32' && WINDOWS_CHILD_ENV_NAMES.has(normalized)) && + !( + platform === 'win32' && + WINDOWS_CHILD_ENV_NAMES.has(normalized) + ) && !normalized.startsWith('LC_')) ); }) @@ -244,6 +244,36 @@ function normalizedEnvironmentName( return platform === 'win32' ? name.toUpperCase() : name; } +/** Distinguishes an unsupported host filesystem from an implementation fault. */ +export class CopyOnWriteCloneUnavailableError extends WorkspaceToolError { + constructor() { + super( + 'Selected-workspace PTC requires copy-on-write filesystem cloning', + 'COMMAND_UNAVAILABLE', + ); + this.name = 'CopyOnWriteCloneUnavailableError'; + } +} + +function isCopyOnWriteUnsupported( + error: unknown, + platform: NodeJS.Platform, +): boolean { + if (platform === 'win32') return true; + if ( + typeof error === 'object' && + error !== null && + 'code' in error && + (error.code === 'ENOTSUP' || error.code === 'EOPNOTSUPP') + ) { + return true; + } + return ( + error instanceof Error && + error.message.toLowerCase().includes('operation not supported') + ); +} + export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox { readonly mutationFailuresAreAtomic = true as const; private readonly manager: NativeSandboxManager; @@ -252,6 +282,9 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox private readonly platform: NodeJS.Platform; private initialized?: Promise; private canonicalRoot?: string; + private runtimeConfig?: SandboxRuntimeConfig; + private denyReadPaths: string[] = []; + private denyWritePaths: string[] = []; private scratchDirectory?: string; private scratchHandle?: FileHandle; private execution?: Promise; @@ -288,7 +321,7 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox ); } managerOwners.set(this.manager, this); - this.initialized = this.initializeOnce().catch(async (error) => { + this.initialized = this.initializeOnce().catch(async error => { await this.manager.reset().catch(() => { this.resetFailed = true; }); @@ -314,7 +347,9 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox 'COMMAND_UNAVAILABLE', ); } - const home = await canonicalPath(this.options.homeDirectory ?? homedir()); + const home = await canonicalPath( + this.options.homeDirectory ?? homedir(), + ); if (isWithin(root, home)) { throw new WorkspaceToolError( 'Native sandbox workspace cannot contain the worker home directory', @@ -324,7 +359,7 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox const protectedPaths = await Promise.all( (this.options.protectedPaths ?? []).map(canonicalPath), ); - if (protectedPaths.some((path) => isWithin(root, path))) { + if (protectedPaths.some(path => isWithin(root, path))) { throw new WorkspaceToolError( 'Native sandbox workspace cannot contain worker control files', 'REGISTRATION_INVALID', @@ -338,13 +373,16 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox const inheritedWritablePaths = [ ...sharedScratchPaths, ...(await Promise.all( - [join(home, '.npm', '_logs'), join(home, '.claude', 'debug')].map( - canonicalPath, - ), + [ + join(home, '.npm', '_logs'), + join(home, '.claude', 'debug'), + ].map(canonicalPath), )), ]; - const deniedInheritedWritablePaths = [...new Set(inheritedWritablePaths)]; - if (deniedInheritedWritablePaths.some((path) => isWithin(path, root))) { + const deniedInheritedWritablePaths = [ + ...new Set(inheritedWritablePaths), + ]; + if (deniedInheritedWritablePaths.some(path => isWithin(path, root))) { throw new WorkspaceToolError( 'Native sandbox workspace cannot be inside an inherited writable path', 'REGISTRATION_INVALID', @@ -359,7 +397,10 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox } if (this.platform !== 'win32') { try { - await access(this.options.shellPath ?? '/bin/bash', fsConstants.X_OK); + await access( + this.options.shellPath ?? '/bin/bash', + fsConstants.X_OK, + ); } catch { throw new WorkspaceToolError( `Native sandbox shell is unavailable: ${this.options.shellPath ?? '/bin/bash'}`, @@ -383,35 +424,40 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox ); const unrestrictedNetwork = commandPolicy.network.outbound === 'unrestricted'; - const config: SandboxRuntimeConfig = { - network: { + const network: SandboxRuntimeConfig['network'] = { allowedDomains: [...(this.options.allowedDomains ?? [])], deniedDomains: [], strictAllowlist: !unrestrictedNetwork, allowAllUnixSockets: commandPolicy.network.allowAllUnixSockets, allowLocalBinding: commandPolicy.network.allowLocalBinding, ...(this.options.maskedEnvironment ? { tlsTerminate: {} } : {}), - }, + }; + const config: SandboxRuntimeConfig = { + network, filesystem: { denyRead: [ home, - ...sharedScratchPaths.filter((path) => + ...sharedScratchPaths.filter(path => deniedInheritedWritablePaths.includes(path), ), ], allowRead: [ root, - ...(canonicalScratchDirectory ? [canonicalScratchDirectory] : []), + ...(canonicalScratchDirectory + ? [canonicalScratchDirectory] + : []), ], allowWrite: [ root, - ...(canonicalScratchDirectory ? [canonicalScratchDirectory] : []), + ...(canonicalScratchDirectory + ? [canonicalScratchDirectory] + : []), ], denyWrite: [...protectedPaths, ...deniedInheritedWritablePaths], allowGitConfig: false, }, credentials: { - files: protectedPaths.map((path) => ({ + files: protectedPaths.map(path => ({ path, mode: 'deny' as const, })), @@ -424,15 +470,18 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox }, this.platform, ) - .filter((name) => { + .filter(name => { const normalized = normalizedEnvironmentName( name, this.platform, ); return ( - !Object.hasOwn(TRUSTED_GIT_ENVIRONMENT, normalized) && + !Object.hasOwn( + TRUSTED_GIT_ENVIRONMENT, + normalized, + ) && !this.options.maskedEnvironment?.variables.some( - (variable) => + variable => normalizedEnvironmentName( variable.name, this.platform, @@ -440,12 +489,16 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox ) ); }) - .map((name) => ({ name, mode: 'deny' as const })), - ...(this.options.maskedEnvironment?.variables.map((variable) => ({ + .map(name => ({ name, mode: 'deny' as const })), + ...(this.options.maskedEnvironment?.variables.map( + variable => ({ ...variable, mode: 'mask' as const, - ...(variable.extract ? { onExtractNoMatch: 'error' as const } : {}), - })) ?? []), + ...(variable.extract + ? { onExtractNoMatch: 'error' as const } + : {}), + }), + ) ?? []), ], }, allowAppleEvents: false, @@ -458,6 +511,17 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox unrestrictedNetwork ? async () => true : undefined, ); this.canonicalRoot = root; + this.runtimeConfig = config; + this.denyReadPaths = [ + home, + ...sharedScratchPaths.filter(path => + deniedInheritedWritablePaths.includes(path), + ), + ]; + this.denyWritePaths = [ + ...protectedPaths, + ...deniedInheritedWritablePaths, + ]; } async execute( @@ -479,9 +543,248 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox } } + /** + * Allocate an owner-only execution directory that is already inside this + * sandbox's allowlist. The caller must remove the returned directory after + * the execution settles. It is intentionally unavailable on native Windows + * until the restricted-account TEMP directory can be opened and verified by + * the trusted parent process. + */ + async createExecutionDirectory(): Promise { + await this.initialize(); + if (!this.scratchDirectory || this.platform === 'win32') { + throw new WorkspaceToolError( + 'Native programmatic execution storage is unavailable', + 'COMMAND_UNAVAILABLE', + ); + } + return await mkdtemp(join(this.scratchDirectory, 'execution-')); + } + + /** + * Clone the current workspace into private scratch for a side-effect- + * equivalent replay probe. Platform clone flags are intentionally strict: + * silently falling back to a byte copy would make every tool-bearing run + * consume time and disk proportional to the repository size. + */ + async createProgrammaticProbeWorkspace( + executionDirectory: string, + signal?: AbortSignal, + ): Promise { + await this.initialize(); + const scratchDirectory = this.scratchDirectory; + const root = this.canonicalRoot; + let parent: string; + try { + parent = await realpath(executionDirectory); + if ( + !scratchDirectory || + !root || + !isWithin(scratchDirectory, parent) || + !(await stat(parent)).isDirectory() + ) { + throw new Error('invalid execution directory'); + } + } catch { + throw new WorkspaceToolError( + 'Programmatic execution directory is unavailable', + 'INVALID_PATH', + ); + } + if (signal?.aborted) { + throw new WorkspaceToolError( + 'Programmatic execution aborted', + 'EXECUTION_ABORTED', + ); + } + const destination = join(parent, 'workspace'); + try { + if (this.platform === 'win32') { + throw new Error('copy-on-write cloning is unavailable on Windows'); + } + const args = + this.platform === 'darwin' + ? ['-cR', root, destination] + : ['--archive', '--reflink=always', root, destination]; + await new Promise((resolveCopy, rejectCopy) => { + const child = this.spawnCommand('/bin/cp', args, { + env: { + PATH: this.environment.PATH, + LANG: this.environment.LANG, + LC_ALL: this.environment.LC_ALL, + }, + signal, + }); + let stderr = Buffer.alloc(0); + child.stderr.on('data', (chunk: Buffer) => { + if (stderr.byteLength < 4_096) { + stderr = Buffer.concat([stderr, chunk]).subarray(0, 4_096); + } + }); + child.once('error', rejectCopy); + child.once('close', code => { + if (code === 0) resolveCopy(); + else { + rejectCopy( + new Error( + `copy-on-write clone failed (${code ?? 'signal'}): ${boundedUtf8(stderr, 4_096)}`, + ), + ); + } + }); + }); + return await realpath(destination); + } catch (error) { + await rm(destination, { recursive: true, force: true }).catch( + () => undefined, + ); + if (signal?.aborted) { + throw new WorkspaceToolError( + 'Programmatic execution aborted', + 'EXECUTION_ABORTED', + ); + } + if (isCopyOnWriteUnsupported(error, this.platform)) { + throw new CopyOnWriteCloneUnavailableError(); + } + throw new WorkspaceToolError( + 'Copy-on-write workspace clone failed unexpectedly', + 'COMMAND_UNAVAILABLE', + ); + } + } + + /** Run a generated program from a verified private execution directory. */ + async executeProgrammatic( + request: WorkspaceExecuteCommandRequest, + dataDirectory: string, + signal?: AbortSignal, + options?: { probe?: boolean; workspaceRoot?: string }, + ): Promise { + if (this.execution || this.closing) { + throw new WorkspaceToolError( + 'Native sandbox already has an active command or is closing', + 'COMMAND_UNAVAILABLE', + ); + } + await this.initialize(); + const scratchDirectory = this.scratchDirectory; + let canonicalDataDirectory: string; + let canonicalWorkspaceRoot: string | undefined; + try { + canonicalDataDirectory = await realpath(dataDirectory); + canonicalWorkspaceRoot = options?.workspaceRoot + ? await realpath(options.workspaceRoot) + : undefined; + if ( + !scratchDirectory || + !isWithin(scratchDirectory, canonicalDataDirectory) || + !(await stat(canonicalDataDirectory)).isDirectory() || + (canonicalWorkspaceRoot != null && + (!isWithin(scratchDirectory, canonicalWorkspaceRoot) || + !(await stat(canonicalWorkspaceRoot)).isDirectory())) + ) { + throw new Error('invalid execution directory'); + } + } catch { + throw new WorkspaceToolError( + 'Programmatic execution directory is unavailable', + 'INVALID_PATH', + ); + } + const execute = () => this.executeExclusive( + request, + signal, + { + LIBRECHAT_CODE_DATA_DIR: canonicalDataDirectory, + LIBRECHAT_CODE_CONTROL_PATH: join( + canonicalDataDirectory, + '_ptc_pending_result.json', + ), + LIBRECHAT_CODE_BASH_PATH: this.options.shellPath ?? '/bin/bash', + PTC_HISTORY_PATH: join( + canonicalDataDirectory, + '_ptc_history.json', + ), + TMPDIR: canonicalDataDirectory, + }, + options?.probe + ? { + filesystem: { + allowRead: [ + canonicalWorkspaceRoot ?? this.canonicalRoot!, + canonicalDataDirectory, + ], + allowWrite: [ + ...(canonicalWorkspaceRoot != null + ? [canonicalWorkspaceRoot] + : []), + canonicalDataDirectory, + ], + denyRead: this.denyReadPaths, + denyWrite: [ + this.canonicalRoot!, + ...this.denyWritePaths, + ], + }, + network: { + // A probe is speculative, even on a trusted VM. + // Copy-on-write protects files, not remote mutations. + allowedDomains: [], + deniedDomains: [], + strictAllowlist: true, + allowUnixSockets: [], + allowAllUnixSockets: false, + allowLocalBinding: false, + }, + } + : undefined, + canonicalDataDirectory, + canonicalWorkspaceRoot, + ); + const execution = options?.probe ? this.withProbeNetwork(execute) : execute(); + this.execution = execution; + try { + return await execution; + } finally { + this.execution = undefined; + } + } + + private async withProbeNetwork(execute: () => Promise): Promise { + const config = this.runtimeConfig; + if (!config || !this.manager.updateConfig) { + throw new WorkspaceToolError('Native probe network isolation is unavailable', 'COMMAND_UNAVAILABLE'); + } + // SRT's proxies and Unix/local socket rules read session configuration, + // not wrapWithSandboxArgv's per-command override. + this.manager.updateConfig({ ...config, network: { + allowedDomains: [], deniedDomains: [], strictAllowlist: true, + allowUnixSockets: [], allowAllUnixSockets: false, allowLocalBinding: false, + } }); + try { + return await execute(); + } finally { + try { + // Revoke the probe's proxy endpoints and credentials before restoring + // network access. A lingering probe must never inherit the commit's + // permissive proxy session through a live updateConfig. + await this.manager.reset(); + await this.manager.initialize(config, config.network.strictAllowlist ? undefined : async () => true); + } catch { + this.resetFailed = true; + throw new WorkspaceToolError('Native probe network cleanup failed', 'COMMAND_UNAVAILABLE'); + } + } + } + private async executeExclusive( request: WorkspaceExecuteCommandRequest, signal?: AbortSignal, + trustedEnvironment?: NodeJS.ProcessEnv, + customConfig?: Partial, + sandboxScratchDirectory?: string, + workspaceRoot?: string, ): Promise { if ( !isWorkspaceToolRequest(request) || @@ -499,7 +802,7 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox ); } await this.initialize(); - const root = this.canonicalRoot!; + const root = workspaceRoot ?? this.canonicalRoot!; let cwd: string; try { cwd = await realpath(resolve(root, request.cwd ?? '.')); @@ -528,7 +831,7 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox { ...TRUSTED_GIT_ENVIRONMENT, ...(credentialEnvironment ?? {}), - ...this.scratchSelectorEnvironment(), + ...this.scratchSelectorEnvironment(sandboxScratchDirectory), }, () => this.manager.wrapWithSandboxArgv( @@ -536,7 +839,7 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox this.platform === 'win32' ? undefined : (this.options.shellPath ?? '/bin/bash'), - undefined, + customConfig, signal, cwd, { commandId, commandText: request.command }, @@ -561,7 +864,14 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox 'EXECUTION_ABORTED', ); } - return await this.runWrapped(request, wrapped, cwd, commandId, signal); + return await this.runWrapped( + request, + wrapped, + cwd, + commandId, + signal, + trustedEnvironment, + ); } finally { // A successful wrap owns command state even when no child is spawned. try { @@ -578,7 +888,7 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox ): Promise { const previousMutation = hostEnvironmentMutationQueue; let releaseMutation!: () => void; - hostEnvironmentMutationQueue = new Promise((resolve) => { + hostEnvironmentMutationQueue = new Promise(resolve => { releaseMutation = resolve; }); await previousMutation; @@ -604,31 +914,41 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox cwd: string, commandId: string, signal?: AbortSignal, + trustedEnvironment?: NodeJS.ProcessEnv, ): Promise { const outputLimit = - request.maxOutputBytes ?? BRIDGE_WORKSPACE_COMMAND_DEFAULT_OUTPUT_BYTES; + request.maxOutputBytes ?? + BRIDGE_WORKSPACE_COMMAND_DEFAULT_OUTPUT_BYTES; const timeoutMs = request.timeoutMs ?? BRIDGE_WORKSPACE_COMMAND_DEFAULT_TIMEOUT_MS; return await new Promise( (resolvePromise, reject) => { let child: ChildProcessWithoutNullStreams; try { - child = this.spawnCommand(wrapped.argv[0], wrapped.argv.slice(1), { + child = this.spawnCommand( + wrapped.argv[0], + wrapped.argv.slice(1), + { cwd, env: { ...wrapped.env, ...this.scratchEnvironment(), + ...trustedEnvironment, ...TRUSTED_GIT_CONFIG_ENTRIES, GIT_CONFIG_COUNT: - wrapped.env.GIT_CONFIG_COUNT ?? TRUSTED_GIT_CONFIG_COUNT, + wrapped.env.GIT_CONFIG_COUNT ?? + TRUSTED_GIT_CONFIG_COUNT, GIT_CONFIG_GLOBAL: - this.platform === 'win32' ? 'NUL' : '/dev/null', + this.platform === 'win32' + ? 'NUL' + : '/dev/null', GIT_CONFIG_NOSYSTEM: '1', }, detached: this.platform !== 'win32', shell: false, windowsHide: true, - }); + }, + ); child.stdin.end(); } catch { reject( @@ -654,10 +974,15 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox const accepted = chunk.subarray(0, remaining); target.push(accepted); outputBytes += accepted.byteLength; - if (accepted.byteLength !== chunk.byteLength) truncated = true; + if (accepted.byteLength !== chunk.byteLength) + truncated = true; }; - child.stdout.on('data', (chunk: Buffer) => append(stdout, chunk)); - child.stderr.on('data', (chunk: Buffer) => append(stderr, chunk)); + child.stdout.on('data', (chunk: Buffer) => + append(stdout, chunk), + ); + child.stderr.on('data', (chunk: Buffer) => + append(stderr, chunk), + ); const abort = (): void => { if (settled) return; this.killCommandTree(child); @@ -706,7 +1031,10 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox ); return; } - const stdoutValue = boundedUtf8(Buffer.concat(stdout), outputLimit); + const stdoutValue = boundedUtf8( + Buffer.concat(stdout), + outputLimit, + ); const stderrBudget = Math.max( 0, outputLimit - Buffer.byteLength(stdoutValue), @@ -714,7 +1042,8 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox const rawStderr = Buffer.concat(stderr).toString('utf8'); let annotatedStderr = rawStderr; try { - annotatedStderr = this.manager.annotateStderrWithSandboxFailures( + annotatedStderr = + this.manager.annotateStderrWithSandboxFailures( commandId, rawStderr, ); @@ -730,12 +1059,15 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox operation: 'execute_command', workspaceId: request.workspaceId, exitCode: - timedOut || childSignal ? null : this.protocolExitCode(code), + timedOut || childSignal + ? null + : this.protocolExitCode(code), ...(childSignal ? { signal: childSignal } : {}), stdout: stdoutValue, stderr: stderrValue, truncated: - truncated || Buffer.byteLength(annotatedStderr) > stderrBudget, + truncated || + Buffer.byteLength(annotatedStderr) > stderrBudget, timedOut, }); }); @@ -775,7 +1107,7 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox ); } const canonicalTemporaryRoot = await canonicalPath(HOST_TEMPORARY_ROOT); - const sharedScratchRoot = sharedScratchPaths.find((path) => + const sharedScratchRoot = sharedScratchPaths.find(path => isWithin(path, canonicalTemporaryRoot), ); const scratchDirectory = await mkdtemp( @@ -798,7 +1130,9 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox true, ); if (((await scratchHandle.stat()).mode & 0o777) !== 0o700) { - throw new Error('Native sandbox scratch directory is not private'); + throw new Error( + 'Native sandbox scratch directory is not private', + ); } this.scratchHandle = scratchHandle; } catch (error) { @@ -829,11 +1163,13 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox : { TMPDIR: scratchDirectory }; } - private scratchSelectorEnvironment(): NodeJS.ProcessEnv { - const scratchDirectory = this.scratchDirectory; + private scratchSelectorEnvironment( + selectedDirectory = this.scratchDirectory, + ): NodeJS.ProcessEnv { + const scratchDirectory = selectedDirectory; if (!scratchDirectory) return {}; return Object.fromEntries( - SRT_SCRATCH_SELECTOR_NAMES.map((name) => [name, scratchDirectory]), + SRT_SCRATCH_SELECTOR_NAMES.map(name => [name, scratchDirectory]), ); } diff --git a/packages/code/src/protocol.test.ts b/packages/code/src/protocol.test.ts index fd426783..08cba1c7 100644 --- a/packages/code/src/protocol.test.ts +++ b/packages/code/src/protocol.test.ts @@ -1,8 +1,11 @@ import assert from 'node:assert/strict'; import test from 'node:test'; import { + bridgeArtifactMediaType, bridgeWorkerPath, comparePortableRelativePaths, + isBridgeWorkspaceProgrammaticRequest, + isSupportedBridgeArtifactName, isValidBridgeWorkerCapabilities, isValidBridgeWorkerId, isWorkspaceToolRequest, @@ -13,6 +16,23 @@ import type { WorkspacePreviewEditRequest, } from './protocol.js'; +test('accepts gateway directory markers as artifacts', () => { + assert.equal(isSupportedBridgeArtifactName('.dirkeep'), true); + assert.equal(isSupportedBridgeArtifactName('nested/.dirkeep'), true); + assert.equal(isSupportedBridgeArtifactName('nested/.dirkeep.exe'), false); +}); + +test('rejects caller-supplied programmatic control payloads', () => { + for (const name of ['_ptc_pending_result.json', '_PTC_PENDING_RESULT.JSON', 'nested/_ptc_pending_result.json']) { + assert.equal(isBridgeWorkspaceProgrammaticRequest({ + headers: {}, + body: { language: 'bash', version: '5.2.0', session_id: 'session', files: [ + { name: 'main.sh', content: 'true' }, { name, content: '{}' }, + ] }, + }), false); + } +}); + const validSingleEditRequest: WorkspaceEditFileRequest = { protocolVersion: 1, operation: 'edit_file', @@ -592,3 +612,167 @@ test('workspace capabilities allow per-workspace operation restrictions', () => false, ); }); + +test('workspace programmatic capability is closed to Bash command roots', () => { + const workspaceTools = { + protocolVersion: 1, + operations: ['execute_command'], + programmaticLanguages: ['bash'], + workspaces: [{ id: 'project-a' }], + }; + assert.equal( + isValidBridgeWorkerCapabilities({ + statefulWorkspace: false, + sandboxProfile: 'anthropic-srt', + runtimes: [], + workspaceTools, + }), + true, + ); + assert.equal( + isValidBridgeWorkerCapabilities({ + statefulWorkspace: false, + sandboxProfile: 'anthropic-srt', + runtimes: [], + workspaceTools: { ...workspaceTools, operations: ['read_file'] }, + }), + false, + ); + assert.equal( + isValidBridgeWorkerCapabilities({ + statefulWorkspace: false, + sandboxProfile: 'anthropic-srt', + runtimes: [], + workspaceTools: { ...workspaceTools, programmaticLanguages: ['python'] }, + }), + false, + ); +}); + +test('workspace programmatic requests accept only stable input cache identities', () => { + const request = { + headers: {}, + body: { + language: 'bash', + version: '5.2', + execution_id: 'execution_1', + replay_tool_count: 2, + max_output_files: 50, + max_output_file_bytes: 10_000_000, + session_id: 'session-1', + files: [ + { name: 'main.sh', content: 'echo ready' }, + { + name: 'skills/example.txt', + id: 'file-1', + storage_session_id: 'storage-1', + input_cache_key: 'a'.repeat(64), + }, + ], + }, + }; + assert.equal(isBridgeWorkspaceProgrammaticRequest(request), true); + assert.equal( + isBridgeWorkspaceProgrammaticRequest({ + ...request, + body: { + ...request.body, + files: [request.body.files[0], { ...request.body.files[1], input_cache_key: '../cache' }], + }, + }), + false, + ); + for (const body of [ + { ...request.body, execution_id: '../execution' }, + { ...request.body, replay_tool_count: -1 }, + { ...request.body, replay_tool_count: 257 }, + { ...request.body, max_output_files: -1 }, + { ...request.body, max_output_files: 101 }, + { ...request.body, max_output_file_bytes: 0 }, + { ...request.body, max_output_file_bytes: 10 * 1024 * 1024 + 1 }, + ]) { + assert.equal( + isBridgeWorkspaceProgrammaticRequest({ ...request, body }), + false, + ); + } +}); + +test('workspace programmatic history can use the bounded replay aggregate budget', () => { + const history = 'h'.repeat(10 * 1024 * 1024 + 1); + const body = { + language: 'bash', + version: '5.2', + session_id: 'session-1', + files: [ + { name: 'main.sh', content: 'echo ready' }, + { name: '_ptc_history.json', content: history }, + ], + }; + assert.equal(isBridgeWorkspaceProgrammaticRequest({ headers: {}, body }), true); + assert.equal( + isBridgeWorkspaceProgrammaticRequest({ + headers: {}, + body: { + ...body, + files: [ + { name: 'main.sh', content: history }, + { name: '_ptc_history.json', content: '{}' }, + ], + }, + }), + false, + ); +}); + +test('workspace programmatic requests reject non-canonical file paths', () => { + for (const name of ['./main.sh', 'scripts//main.sh', 'scripts/./main.sh', '.']) { + assert.equal( + isBridgeWorkspaceProgrammaticRequest({ + headers: {}, + body: { + language: 'bash', + version: '5.2', + session_id: 'session-1', + files: [ + { name: 'main.sh', content: 'echo ready' }, + { name, content: 'data' }, + ], + }, + }), + false, + name, + ); + } +}); + +test('workspace programmatic requests reject ancestor-descendant input conflicts', () => { + for (const names of [ + ['main.sh', 'main.sh/data.txt'], + ['main.sh', 'assets', 'assets/logo.png'], + ['main.sh', 'deep/path/file.txt', 'deep'], + ]) { + assert.equal( + isBridgeWorkspaceProgrammaticRequest({ + headers: {}, + body: { + language: 'bash', + version: '5.2', + session_id: 'session-1', + files: names.map(name => ({ name, content: 'data' })), + }, + }), + false, + names.join(', '), + ); + } +}); + +test('bridge artifact policy and media types match the hardened gateway contract', () => { + assert.equal(isSupportedBridgeArtifactName('reports/result.json'), true); + assert.equal(isSupportedBridgeArtifactName('preview.png'), true); + assert.equal(isSupportedBridgeArtifactName('model.bin'), false); + assert.equal(bridgeArtifactMediaType('preview.png'), 'image/png'); + assert.equal(bridgeArtifactMediaType('reports/result.json'), 'application/json'); + assert.equal(bridgeArtifactMediaType('Dockerfile'), 'application/octet-stream'); +}); diff --git a/packages/code/src/protocol.ts b/packages/code/src/protocol.ts index c1dad949..b92d21ac 100644 --- a/packages/code/src/protocol.ts +++ b/packages/code/src/protocol.ts @@ -20,9 +20,144 @@ 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; +export const BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILES = 100; +export const BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_INPUT_FILES = BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILES - 2; +export const BRIDGE_WORKSPACE_PROGRAMMATIC_TRANSFER_CONCURRENCY = 4; +export const BRIDGE_WORKSPACE_PROGRAMMATIC_TRANSFER_TIMEOUT_MS = 30_000; + +/** Reserve a bounded share for all input/output batches, not per-file grants. */ +export function programmaticTransferReserveMs(jobTimeoutMs: number): number { + return Math.max(1, Math.floor(jobTimeoutMs / 3)); +} +export const BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILE_BYTES = 10 * 1024 * 1024; +export const BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_HISTORY_BYTES = 40_000_000; +export const BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_TOTAL_BYTES = 100 * 1024 * 1024; /** How long Code API drains a clean rejection after Stop cancels a workspace mutation. */ export const BRIDGE_CANCELLED_WORKSPACE_SETTLEMENT_GRACE_MS = 5_000; +/** + * Artifact names accepted by the hardened egress gateway. Keep this policy in + * the bridge protocol package so a remote worker can reject unsupported output + * locally instead of discovering the mismatch only after mutating a workspace. + */ +const BRIDGE_ARTIFACT_EXTENSIONS = new Set([ + '.c', '.cs', '.cpp', '.go', '.java', '.js', '.kt', '.kts', '.lua', + '.php', '.pl', '.ps1', '.py', '.r', '.rb', '.rs', '.scala', '.sh', + '.sql', '.swift', '.ts', '.jsx', '.tsx', '.groovy', + '.css', '.htm', '.html', '.less', '.sass', '.scss', '.svg', '.svelte', '.vue', + '.adoc', '.asciidoc', '.md', '.rst', '.tex', '.txt', '.wiki', + '.csv', '.json', '.bson', '.json5', '.jsonl', '.parquet', '.tsv', + '.xml', '.yaml', '.yml', + '.ics', '.ical', '.ifb', '.icalendar', + '.conf', '.env', '.gitignore', '.ini', '.properties', '.toml', + '.doc', '.docx', '.pdf', '.ppt', '.pptx', '.xls', '.xlsx', + '.odt', '.ods', '.odp', '.rtf', + '.avif', '.bmp', '.gif', '.ico', '.jpeg', '.jpg', '.png', + '.tif', '.tiff', '.webp', + '.eot', '.ttf', '.woff', '.woff2', + '.7z', '.bz2', '.gz', '.gzip', '.rar', '.tar', '.zip', + '.tf', '.tfvars', '.tfstate', '.hcl', + '.dockerfile', '.Dockerfile', '.dockerignore', + '.helmignore', '.helmfile', '.jenkinsfile', '.vagrantfile', + '.eslintrc', '.prettierrc', '.editorconfig', '.nomad', + '.bat', '.cmd', '.deb', '.log', '.rpm', '.vbs', +]); + +function portableBasename(name: string): string { + return name.slice(name.lastIndexOf('/') + 1); +} + +/** Apply the gateway's extension allowlist without importing service code. */ +export function isSupportedBridgeArtifactName(name: string): boolean { + const basename = portableBasename(name); + if (basename === '.dirkeep') return true; + const dot = basename.lastIndexOf('.'); + const extension = dot > 0 ? basename.slice(dot).toLowerCase() : ''; + const dottedBasename = `.${basename}`; + return ( + (extension !== '' && BRIDGE_ARTIFACT_EXTENSIONS.has(extension)) || + BRIDGE_ARTIFACT_EXTENSIONS.has(basename) || + BRIDGE_ARTIFACT_EXTENSIONS.has(basename.toLowerCase()) || + (extension === '' && + (BRIDGE_ARTIFACT_EXTENSIONS.has(dottedBasename) || + BRIDGE_ARTIFACT_EXTENSIONS.has(dottedBasename.toLowerCase()))) + ); +} + +const BRIDGE_ARTIFACT_MEDIA_TYPES: Readonly> = { + '.avif': 'image/avif', + '.bmp': 'image/bmp', + '.bz2': 'application/x-bzip2', + '.c': 'text/x-c', + '.conf': 'text/plain', + '.cpp': 'text/x-c++src', + '.css': 'text/css', + '.csv': 'text/csv', + '.doc': 'application/msword', + '.docx': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document', + '.gif': 'image/gif', + '.gz': 'application/gzip', + '.gzip': 'application/gzip', + '.htm': 'text/html', + '.html': 'text/html', + '.ico': 'image/x-icon', + '.ics': 'text/calendar', + '.ifb': 'text/calendar', + '.ical': 'text/calendar', + '.icalendar': 'text/calendar', + '.ini': 'text/plain', + '.java': 'text/x-java-source', + '.jpeg': 'image/jpeg', + '.jpg': 'image/jpeg', + '.js': 'text/javascript', + '.json': 'application/json', + '.json5': 'application/json5', + '.jsonl': 'application/x-ndjson', + '.jsx': 'text/jsx', + '.log': 'text/plain', + '.md': 'text/markdown', + '.odt': 'application/vnd.oasis.opendocument.text', + '.ods': 'application/vnd.oasis.opendocument.spreadsheet', + '.odp': 'application/vnd.oasis.opendocument.presentation', + '.parquet': 'application/vnd.apache.parquet', + '.pdf': 'application/pdf', + '.png': 'image/png', + '.ppt': 'application/vnd.ms-powerpoint', + '.pptx': 'application/vnd.openxmlformats-officedocument.presentationml.presentation', + '.py': 'text/x-python', + '.rst': 'text/x-rst', + '.rtf': 'application/rtf', + '.sh': 'application/x-sh', + '.sql': 'application/sql', + '.svg': 'image/svg+xml', + '.tar': 'application/x-tar', + '.tex': 'application/x-tex', + '.tif': 'image/tiff', + '.tiff': 'image/tiff', + '.toml': 'application/toml', + '.ts': 'text/typescript', + '.tsx': 'text/tsx', + '.tsv': 'text/tab-separated-values', + '.txt': 'text/plain', + '.webp': 'image/webp', + '.woff': 'font/woff', + '.woff2': 'font/woff2', + '.xls': 'application/vnd.ms-excel', + '.xlsx': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet', + '.xml': 'application/xml', + '.yaml': 'application/yaml', + '.yml': 'application/yaml', + '.zip': 'application/zip', +}; + +/** Infer a safe response media type from an already-validated artifact name. */ +export function bridgeArtifactMediaType(name: string): string { + const basename = portableBasename(name).toLowerCase(); + const dot = basename.lastIndexOf('.'); + const extension = dot > 0 ? basename.slice(dot) : basename; + return BRIDGE_ARTIFACT_MEDIA_TYPES[extension] ?? 'application/octet-stream'; +} + export type BridgeProtocolVersion = typeof BRIDGE_PROTOCOL_VERSION; export type BridgeWorkspaceToolOperation = @@ -38,6 +173,7 @@ export type WorkspaceWriteFileMode = 'replace' | 'create'; export type WorkspaceEditFileMode = 'single' | 'batch'; export type WorkspaceEditFileFeature = 'expected_base_sha256'; export type WorkspaceListFileFeature = 'after_path'; +export type WorkspaceProgrammaticLanguage = 'bash'; export interface BridgeWorkspaceDescriptor { id: string; @@ -58,6 +194,8 @@ export interface BridgeWorkspaceToolCapabilities { editFileFeatures?: WorkspaceEditFileFeature[]; /** Omitted by workers that cannot continue a bounded file listing. */ listFileFeatures?: WorkspaceListFileFeature[]; + /** Languages that can execute PTC replay inside a selected workspace. */ + programmaticLanguages?: WorkspaceProgrammaticLanguage[]; } export interface WorkspaceReadFileRequest { @@ -153,8 +291,7 @@ interface WorkspaceEditFileRequestBase { expectedBaseSha256?: string; } -export interface WorkspaceSingleEditFileRequest - extends WorkspaceEditFileRequestBase { +export interface WorkspaceSingleEditFileRequest extends WorkspaceEditFileRequestBase { /** Legacy single-edit form. */ oldText: string; /** Legacy single-edit form. */ @@ -162,8 +299,7 @@ export interface WorkspaceSingleEditFileRequest edits?: never; } -export interface WorkspaceBatchEditFileRequest - extends WorkspaceEditFileRequestBase { +export interface WorkspaceBatchEditFileRequest extends WorkspaceEditFileRequestBase { /** Ordered exact replacements applied atomically as one file mutation. */ edits: WorkspaceTextEdit[]; oldText?: never; @@ -171,8 +307,7 @@ export interface WorkspaceBatchEditFileRequest } export type WorkspaceEditFileRequest = - | WorkspaceSingleEditFileRequest - | WorkspaceBatchEditFileRequest; + WorkspaceSingleEditFileRequest | WorkspaceBatchEditFileRequest; export interface WorkspaceTextEdit { oldText: string; @@ -195,23 +330,20 @@ interface WorkspacePreviewEditRequestBase { path: string; } -export interface WorkspaceSinglePreviewEditRequest - extends WorkspacePreviewEditRequestBase { +export interface WorkspaceSinglePreviewEditRequest extends WorkspacePreviewEditRequestBase { oldText: string; newText: string; edits?: never; } -export interface WorkspaceBatchPreviewEditRequest - extends WorkspacePreviewEditRequestBase { +export interface WorkspaceBatchPreviewEditRequest extends WorkspacePreviewEditRequestBase { edits: WorkspaceTextEdit[]; oldText?: never; newText?: never; } export type WorkspacePreviewEditRequest = - | WorkspaceSinglePreviewEditRequest - | WorkspaceBatchPreviewEditRequest; + WorkspaceSinglePreviewEditRequest | WorkspaceBatchPreviewEditRequest; export interface WorkspacePreviewEditResult { protocolVersion: BridgeProtocolVersion; @@ -392,12 +524,7 @@ const WORKSPACE_COMMAND_RESULT_KEYS = new Set([ 'truncated', 'timedOut', ]); -const WORKSPACE_SEARCH_MATCH_KEYS = new Set([ - 'path', - 'line', - 'column', - 'text', -]); +const WORKSPACE_SEARCH_MATCH_KEYS = new Set(['path', 'line', 'column', 'text']); export interface BridgeWorkerCapabilities { /** Opt-in protocol: maximum concurrently leased independent workspace roots. */ @@ -437,6 +564,8 @@ export interface BridgeWorkerRegistrationResponse { supportedWorkspaceEditFileFeatures?: WorkspaceEditFileFeature[]; /** Listing features this Code API can safely route to a capability-aware worker. */ supportedWorkspaceListFileFeatures?: WorkspaceListFileFeature[]; + /** PTC languages this Code API can safely route into a selected workspace. */ + supportedWorkspaceProgrammaticLanguages?: WorkspaceProgrammaticLanguage[]; } /** Administrator-visible liveness for a configured worker. Credentials, @@ -469,6 +598,37 @@ export interface BridgeSandboxRequest { headers: Record; } +export type BridgeProgrammaticPayloadFile = + | { name: string; content: string } + | { + name: string; + id: string; + storage_session_id: string; + input_cache_key?: string; + }; + +export interface BridgeWorkspaceProgrammaticBody { + language: 'bash'; + version: string; + /** Stable identity shared by every replay iteration of one execution. */ + execution_id?: string; + /** Declared replay tools; zero allows the worker to skip the probe pass. */ + replay_tool_count?: number; + run_timeout?: number; + transfer_timeout_ms?: number; + /** Manifest-bound upload ceiling negotiated by Code API. */ + max_output_files?: number; + /** Effective per-file upload ceiling negotiated by Code API. */ + max_output_file_bytes?: number; + files: BridgeProgrammaticPayloadFile[]; + session_id: string; + output_session_id?: string; + egress_grant?: string; +} + +export type BridgeWorkspaceProgrammaticRequest = + BridgeSandboxRequest; + export interface BridgeAssignment { workspaceLeaseSlot?: number; protocolVersion: BridgeProtocolVersion; @@ -481,7 +641,9 @@ export interface BridgeAssignment { /** Server-calculated execution budget at lease time; avoids VM clock skew. */ remainingMs?: number; runtimeSessionId?: string; - executionKind?: 'sandbox' | 'workspace_tool'; + executionKind?: 'sandbox' | 'workspace_tool' | 'workspace_programmatic'; + /** Selected workspace for workspace-scoped programmatic execution. */ + workspaceId?: string; request: BridgeSandboxRequest | WorkspaceToolRequest; } @@ -589,6 +751,140 @@ export function isValidBridgeWorkerId(workerId: string): boolean { return BRIDGE_WORKER_ID_PATTERN.test(workerId); } +export function isBridgeWorkspaceProgrammaticRequest( + value: unknown, +): value is BridgeWorkspaceProgrammaticRequest { + if (typeof value !== 'object' || value === null) return false; + const request = value as Record; + if ( + typeof request.headers !== 'object' || + request.headers === null || + !Object.values(request.headers).every( + entry => typeof entry === 'string', + ) || + typeof request.body !== 'object' || + request.body === null + ) { + return false; + } + const body = request.body as Record; + if ( + body.language !== 'bash' || + typeof body.version !== 'string' || + body.version.length === 0 || + body.version.length > BRIDGE_RUNTIME_MAX_LENGTH || + (body.execution_id !== undefined && + (typeof body.execution_id !== 'string' || + !/^[A-Za-z0-9_-]{1,128}$/.test(body.execution_id))) || + (body.replay_tool_count !== undefined && + (!Number.isSafeInteger(body.replay_tool_count) || + Number(body.replay_tool_count) < 0 || + Number(body.replay_tool_count) > 256)) || + (body.max_output_files !== undefined && + (!Number.isSafeInteger(body.max_output_files) || + Number(body.max_output_files) < 0 || + Number(body.max_output_files) > + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILES)) || + (body.max_output_file_bytes !== undefined && + (!Number.isSafeInteger(body.max_output_file_bytes) || + Number(body.max_output_file_bytes) < 1 || + Number(body.max_output_file_bytes) > + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILE_BYTES)) || + typeof body.session_id !== 'string' || + body.session_id.length === 0 || + body.session_id.length > 32_768 || + /[\0\r\n]/.test(body.session_id) || + (body.output_session_id !== undefined && + (typeof body.output_session_id !== 'string' || + body.output_session_id.length === 0 || + body.output_session_id.length > 32_768 || + /[\0\r\n]/.test(body.output_session_id))) || + (body.egress_grant !== undefined && + (typeof body.egress_grant !== 'string' || + body.egress_grant.length === 0 || + body.egress_grant.length > 256 * 1024)) || + (body.transfer_timeout_ms !== undefined && + (!Number.isSafeInteger(body.transfer_timeout_ms) || + Number(body.transfer_timeout_ms) < 1 || + Number(body.transfer_timeout_ms) > BRIDGE_WORKSPACE_PROGRAMMATIC_TRANSFER_TIMEOUT_MS)) || + (body.run_timeout !== undefined && + (!Number.isSafeInteger(body.run_timeout) || + Number(body.run_timeout) < 1 || + Number(body.run_timeout) > + BRIDGE_WORKSPACE_COMMAND_MAX_TIMEOUT_MS)) || + !Array.isArray(body.files) || + body.files.length < 1 || + body.files.length > BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILES + ) { + return false; + } + let inlineBytes = 0; + const names = new Set(); + for (const rawFile of body.files) { + if (typeof rawFile !== 'object' || rawFile === null) return false; + const file = rawFile as Record; + if ( + !isSafePortableRelativePath(file.name) || + file.name === '.' || + portableBasename(file.name).toLowerCase() === '_ptc_pending_result.json' || + normalizePortableRelativePath(file.name) !== file.name || + names.has(file.name) + ) { + return false; + } + names.add(file.name); + if (typeof file.content === 'string') { + inlineBytes += Buffer.byteLength(file.content); + if ( + Object.keys(file).some( + key => key !== 'name' && key !== 'content', + ) || + Buffer.byteLength(file.content) > + (file.name === '_ptc_history.json' + ? BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_HISTORY_BYTES + : BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILE_BYTES) + ) { + return false; + } + continue; + } + if ( + typeof file.id !== 'string' || + file.id.length === 0 || + file.id.length > 32_768 || + /[\0\r\n]/.test(file.id) || + typeof file.storage_session_id !== 'string' || + file.storage_session_id.length === 0 || + file.storage_session_id.length > 32_768 || + /[\0\r\n]/.test(file.storage_session_id) || + Object.keys(file).some( + key => + key !== 'name' && + key !== 'id' && + key !== 'storage_session_id' && + key !== 'input_cache_key', + ) || + (file.input_cache_key !== undefined && + (typeof file.input_cache_key !== 'string' || + !/^[a-f0-9]{64}$/.test(file.input_cache_key))) + ) { + return false; + } + } + for (const name of names) { + const segments = name.split('/'); + let ancestor = ''; + for (let index = 0; index < segments.length - 1; index += 1) { + ancestor = ancestor ? `${ancestor}/${segments[index]}` : segments[index]!; + if (names.has(ancestor)) return false; + } + } + return ( + names.has('main.sh') && + inlineBytes <= BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_TOTAL_BYTES + ); +} + export function isSafePortableRelativePath(value: unknown): value is string { if ( typeof value !== 'string' || @@ -602,20 +898,23 @@ export function isSafePortableRelativePath(value: unknown): value is string { ) { return false; } - return value.split('/').every((segment) => segment !== '..'); + return value.split('/').every(segment => segment !== '..'); } function normalizePortableRelativePath(value: string): string { return ( value .split('/') - .filter((segment) => segment.length > 0 && segment !== '.') + .filter(segment => segment.length > 0 && segment !== '.') .join('/') || '.' ); } /** Compare path segments in ripgrep's sorted, depth-first traversal order. */ -export function comparePortableRelativePaths(left: string, right: string): number { +export function comparePortableRelativePaths( + left: string, + right: string, +): number { const encoder = new TextEncoder(); const leftSegments = left.split('/'); const rightSegments = right.split('/'); @@ -645,9 +944,14 @@ function isWithinRequestedPath(candidate: string, requested?: string): boolean { ); } -function isValidWorkspaceEditRequest(request: Record): boolean { +function isValidWorkspaceEditRequest( + request: Record, +): boolean { const hasBatch = request.edits !== undefined; - if (hasBatch && (request.oldText !== undefined || request.newText !== undefined)) { + if ( + hasBatch && + (request.oldText !== undefined || request.newText !== undefined) + ) { return false; } const edits = hasBatch @@ -665,7 +969,10 @@ function isValidWorkspaceEditRequest(request: Record): boolean if ( typeof edit !== 'object' || edit === null || - !hasOnlyKeys(edit as Record, WORKSPACE_TEXT_EDIT_KEYS) + !hasOnlyKeys( + edit as Record, + WORKSPACE_TEXT_EDIT_KEYS, + ) ) { return false; } @@ -673,9 +980,11 @@ function isValidWorkspaceEditRequest(request: Record): boolean if ( typeof candidate.oldText !== 'string' || candidate.oldText.length === 0 || - Buffer.from(candidate.oldText).toString('utf8') !== candidate.oldText || + Buffer.from(candidate.oldText).toString('utf8') !== + candidate.oldText || typeof candidate.newText !== 'string' || - Buffer.from(candidate.newText).toString('utf8') !== candidate.newText + Buffer.from(candidate.newText).toString('utf8') !== + candidate.newText ) { return false; } @@ -698,7 +1007,7 @@ function hasOnlyKeys( value: Record, allowed: ReadonlySet, ): boolean { - return Object.keys(value).every((key) => allowed.has(key)); + return Object.keys(value).every(key => allowed.has(key)); } export function isWorkspaceToolRequest( @@ -723,7 +1032,8 @@ export function isWorkspaceToolRequest( (request.maxLines === undefined || (Number.isSafeInteger(request.maxLines) && Number(request.maxLines) >= 1 && - Number(request.maxLines) <= BRIDGE_WORKSPACE_READ_MAX_LINES)) + Number(request.maxLines) <= + BRIDGE_WORKSPACE_READ_MAX_LINES)) ); } if (request.operation === 'search_text') { @@ -743,7 +1053,8 @@ export function isWorkspaceToolRequest( (request.maxResults === undefined || (Number.isSafeInteger(request.maxResults) && Number(request.maxResults) >= 1 && - Number(request.maxResults) <= BRIDGE_WORKSPACE_SEARCH_MAX_RESULTS)) + Number(request.maxResults) <= + BRIDGE_WORKSPACE_SEARCH_MAX_RESULTS)) ); } if (request.operation === 'list_files') { @@ -753,12 +1064,14 @@ export function isWorkspaceToolRequest( isSafePortableRelativePath(request.path)) && (request.afterPath === undefined || (isSafePortableRelativePath(request.afterPath) && - normalizePortableRelativePath(request.afterPath) === request.afterPath && + normalizePortableRelativePath(request.afterPath) === + request.afterPath && isWithinRequestedPath(request.afterPath, request.path))) && (request.maxResults === undefined || (Number.isSafeInteger(request.maxResults) && Number(request.maxResults) >= 1 && - Number(request.maxResults) <= BRIDGE_WORKSPACE_LIST_MAX_RESULTS)) + Number(request.maxResults) <= + BRIDGE_WORKSPACE_LIST_MAX_RESULTS)) ); } if (request.operation === 'write_file') { @@ -799,7 +1112,8 @@ export function isWorkspaceToolRequest( !request.command.includes('\0') && new TextEncoder().encode(request.command).byteLength <= BRIDGE_WORKSPACE_COMMAND_MAX_BYTES && - (request.cwd === undefined || isSafePortableRelativePath(request.cwd)) && + (request.cwd === undefined || + isSafePortableRelativePath(request.cwd)) && (request.timeoutMs === undefined || (Number.isSafeInteger(request.timeoutMs) && Number(request.timeoutMs) >= 1 && @@ -838,13 +1152,19 @@ export function isWorkspaceToolResult( if (request.operation === 'read_file') { const startLine = request.startLine ?? 1; const maxLines = request.maxLines ?? 200; - const content = typeof result.content === 'string' ? result.content : null; + const content = + typeof result.content === 'string' ? result.content : null; const reportedLineCount = - Number.isSafeInteger(result.endLine) && Number(result.endLine) >= startLine - 1 + Number.isSafeInteger(result.endLine) && + Number(result.endLine) >= startLine - 1 ? Number(result.endLine) - startLine + 1 : -1; const actualLineCount = - content === null ? -1 : content.length === 0 ? reportedLineCount : content.split('\n').length; + content === null + ? -1 + : content.length === 0 + ? reportedLineCount + : content.split('\n').length; return ( hasOnlyKeys(result, WORKSPACE_READ_RESULT_KEYS) && result.path === request.path && @@ -899,7 +1219,10 @@ export function isWorkspaceToolResult( (enforcesPaginationContract && (normalizedPath !== path || (previousPath !== undefined && - comparePortableRelativePaths(normalizedPath, previousPath) <= 0))) + comparePortableRelativePaths( + normalizedPath, + previousPath, + ) <= 0))) ) { return false; } @@ -909,7 +1232,8 @@ export function isWorkspaceToolResult( if (!enforcesPaginationContract) { return result.nextAfterPath === undefined; } - if (result.truncated !== true) return result.nextAfterPath === undefined; + if (result.truncated !== true) + return result.nextAfterPath === undefined; return ( result.paths.length > 0 && result.nextAfterPath === result.paths[result.paths.length - 1] @@ -942,7 +1266,8 @@ export function isWorkspaceToolResult( if (request.operation === 'preview_edit') { const replacements = request.edits?.length ?? 1; - const content = typeof result.content === 'string' ? result.content : null; + const content = + typeof result.content === 'string' ? result.content : null; return ( hasOnlyKeys(result, WORKSPACE_PREVIEW_EDIT_RESULT_KEYS) && result.path === request.path && @@ -964,7 +1289,8 @@ export function isWorkspaceToolResult( const stdout = typeof result.stdout === 'string' ? result.stdout : null; const stderr = typeof result.stderr === 'string' ? result.stderr : null; const outputLimit = - request.maxOutputBytes ?? BRIDGE_WORKSPACE_COMMAND_DEFAULT_OUTPUT_BYTES; + request.maxOutputBytes ?? + BRIDGE_WORKSPACE_COMMAND_DEFAULT_OUTPUT_BYTES; return ( hasOnlyKeys(result, WORKSPACE_COMMAND_RESULT_KEYS) && stdout !== null && @@ -980,7 +1306,8 @@ export function isWorkspaceToolResult( Number(result.exitCode) <= 255)) && (result.signal === undefined || (typeof result.signal === 'string' && - result.signal.length <= BRIDGE_WORKSPACE_COMMAND_SIGNAL_MAX_LENGTH && + result.signal.length <= + BRIDGE_WORKSPACE_COMMAND_SIGNAL_MAX_LENGTH && /^SIG[A-Z0-9]+$/.test(result.signal))) && typeof result.truncated === 'boolean' && typeof result.timedOut === 'boolean' && @@ -995,7 +1322,7 @@ export function isWorkspaceToolResult( return ( hasOnlyKeys(result, WORKSPACE_SEARCH_RESULT_KEYS) && result.matches.length <= maxResults && - result.matches.every((match) => { + result.matches.every(match => { if (typeof match !== 'object' || match === null) return false; const candidate = match as Record; return ( @@ -1007,7 +1334,8 @@ export function isWorkspaceToolResult( Number.isSafeInteger(candidate.column) && Number(candidate.column) >= 1 && typeof candidate.text === 'string' && - candidate.text.length <= BRIDGE_WORKSPACE_SEARCH_TEXT_MAX_LENGTH && + candidate.text.length <= + BRIDGE_WORKSPACE_SEARCH_TEXT_MAX_LENGTH && candidate.text.includes(request.query) ); }) @@ -1025,7 +1353,7 @@ export function isValidBridgeWorkspaceToolCapabilities( capabilities.operations.length < 1 || capabilities.operations.length > 7 || !capabilities.operations.every( - (operation) => + operation => operation === 'read_file' || operation === 'search_text' || operation === 'list_files' || @@ -1034,7 +1362,8 @@ export function isValidBridgeWorkspaceToolCapabilities( operation === 'edit_file' || operation === 'execute_command', ) || - new Set(capabilities.operations).size !== capabilities.operations.length || + new Set(capabilities.operations).size !== + capabilities.operations.length || !Array.isArray(capabilities.workspaces) || capabilities.workspaces.length < 1 || capabilities.workspaces.length > BRIDGE_WORKSPACE_MAX_COUNT @@ -1049,7 +1378,7 @@ export function isValidBridgeWorkspaceToolCapabilities( capabilities.writeFileModes.length > 2 || !capabilities.operations.includes('write_file') || !capabilities.writeFileModes.every( - (mode) => mode === 'replace' || mode === 'create', + mode => mode === 'replace' || mode === 'create', ) || new Set(capabilities.writeFileModes).size !== capabilities.writeFileModes.length) @@ -1065,7 +1394,7 @@ export function isValidBridgeWorkspaceToolCapabilities( (!capabilities.operations.includes('edit_file') && !capabilities.operations.includes('preview_edit')) || !capabilities.editFileModes.every( - (mode) => mode === 'single' || mode === 'batch', + mode => mode === 'single' || mode === 'batch', ) || new Set(capabilities.editFileModes).size !== capabilities.editFileModes.length) @@ -1093,13 +1422,23 @@ export function isValidBridgeWorkspaceToolCapabilities( return false; } + if ( + capabilities.programmaticLanguages !== undefined && + (!Array.isArray(capabilities.programmaticLanguages) || + capabilities.programmaticLanguages.length !== 1 || + !capabilities.operations.includes('execute_command') || + capabilities.programmaticLanguages[0] !== 'bash') + ) { + return false; + } + const workspaceIds = new Set(); - return capabilities.workspaces.every((workspace) => { + return capabilities.workspaces.every(workspace => { if (typeof workspace !== 'object' || workspace === null) return false; const descriptor = workspace as Record; if ( Object.keys(descriptor).some( - (key) => key !== 'id' && key !== 'name' && key !== 'operations', + key => key !== 'id' && key !== 'name' && key !== 'operations', ) || typeof descriptor.id !== 'string' || !isValidBridgeWorkerId(descriptor.id) || @@ -1107,17 +1446,21 @@ export function isValidBridgeWorkspaceToolCapabilities( (descriptor.name !== undefined && (typeof descriptor.name !== 'string' || descriptor.name.trim().length === 0 || - descriptor.name.length > BRIDGE_WORKSPACE_NAME_MAX_LENGTH)) || + descriptor.name.length > + BRIDGE_WORKSPACE_NAME_MAX_LENGTH)) || (descriptor.operations !== undefined && (!Array.isArray(descriptor.operations) || descriptor.operations.length < 1 || descriptor.operations.length > (capabilities.operations as unknown[]).length || descriptor.operations.some( - (operation) => - !(capabilities.operations as unknown[]).includes(operation), + operation => + !(capabilities.operations as unknown[]).includes( + operation, + ), ) || - new Set(descriptor.operations).size !== descriptor.operations.length)) + new Set(descriptor.operations).size !== + descriptor.operations.length)) ) { return false; } @@ -1139,11 +1482,12 @@ export function isValidBridgeWorkerCapabilities( typeof capabilities.statefulWorkspace === 'boolean' && typeof capabilities.sandboxProfile === 'string' && capabilities.sandboxProfile.trim().length > 0 && - capabilities.sandboxProfile.length <= BRIDGE_SANDBOX_PROFILE_MAX_LENGTH && + capabilities.sandboxProfile.length <= + BRIDGE_SANDBOX_PROFILE_MAX_LENGTH && Array.isArray(capabilities.runtimes) && capabilities.runtimes.length <= BRIDGE_RUNTIME_MAX_COUNT && capabilities.runtimes.every( - (runtime) => + runtime => typeof runtime === 'string' && runtime.length > 0 && runtime.length <= BRIDGE_RUNTIME_MAX_LENGTH, diff --git a/packages/code/src/worker-slots.test.ts b/packages/code/src/worker-slots.test.ts index ba65a3cf..6e8cc8b1 100644 --- a/packages/code/src/worker-slots.test.ts +++ b/packages/code/src/worker-slots.test.ts @@ -312,3 +312,51 @@ test('a local cleanup handoff preserves the new assignment owner and remaining b assert.equal(executed, true); assert.equal(internals.activeWorkspaceAssignments.size, 0); }); + +test('programmatic work on an independent workspace bypasses another root cleanup', async () => { + const worker = new BridgeWorker({ + codeApiUrl: 'http://localhost:1', + token: 'fixture', + workerId: 'worker', + sandboxEndpoint: 'http://localhost:2', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'fixture', + runtimes: [], + }, + }); + const internals = worker as unknown as { + activeWorkspaceAssignments: Map< + string, + { id: string; done: Promise } + >; + executeOwned: (assignment: BridgeAssignment) => Promise; + }; + internals.activeWorkspaceAssignments.set('a', { + id: 'previous', + done: new Promise(() => {}), + }); + let executed = false; + internals.executeOwned = async () => { + executed = true; + assert.equal(internals.activeWorkspaceAssignments.get('b')?.id, 'next'); + }; + await worker.executeAndSettle({ + assignmentId: 'next', + executionKind: 'workspace_programmatic', + workspaceId: 'b', + remainingMs: 1_000, + request: { + headers: {}, + body: { + language: 'bash', + version: '5.2', + session_id: 'session', + files: [{ name: 'main.sh', content: 'echo ready' }], + }, + }, + } as BridgeAssignment); + assert.equal(executed, true); + assert.equal(internals.activeWorkspaceAssignments.has('b'), false); + assert.equal(internals.activeWorkspaceAssignments.get('a')?.id, 'previous'); +}); diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 19f219da..a0517bd9 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -5,6 +5,7 @@ import { BRIDGE_PROTOCOL_VERSION, BridgeProtocolError, bridgeWorkerPath, + isBridgeWorkspaceProgrammaticRequest, isWorkspaceToolResult, } from './protocol.js'; import { EndpointRuntimeSupervisor } from './runtime.js'; @@ -21,6 +22,7 @@ import type { BridgeWorkerCredentialResponse, BridgeWorkerRegistrationResponse, BridgeWorkspaceToolOperation, + BridgeWorkspaceProgrammaticRequest, } from './protocol.js'; import type { RuntimeLease, RuntimeSupervisor } from './runtime.js'; import type { WorkspaceToolExecutor } from './workspace.js'; @@ -35,6 +37,13 @@ export interface BridgeWorkerOptions { runtimeSupervisor?: RuntimeSupervisor; capabilities: BridgeWorkerCapabilities; workspaceTools?: WorkspaceToolExecutor; + workspaceProgrammatic?: { + executeProgrammatic( + workspaceId: string, + request: BridgeWorkspaceProgrammaticRequest, + signal?: AbortSignal, + ): Promise; + }; workspaceMutationQuarantine?: WorkspaceMutationQuarantine; /** Required per-root durable guards when opting into concurrent workspace leases. */ workspaceQuarantines?: ReadonlyMap; @@ -164,6 +173,11 @@ function workspaceCapabilitiesMatch( (advertised.listFileFeatures?.every( (feature, index) => feature === executor.listFileFeatures?.[index], ) ?? executor.listFileFeatures == null) && + advertised.programmaticLanguages?.length === + executor.programmaticLanguages?.length && + (advertised.programmaticLanguages?.every( + (language, index) => language === executor.programmaticLanguages?.[index], + ) ?? executor.programmaticLanguages == null) && advertised.workspaces.length === executor.workspaces.length && advertised.workspaces.every( (workspace, index) => @@ -223,6 +237,7 @@ function registrationCompatibleCapabilities( editFileModes: _editFileModes, editFileFeatures: _editFileFeatures, listFileFeatures: _listFileFeatures, + programmaticLanguages: _programmaticLanguages, ...compatibleWorkspaceTools } = workspaceTools; return { @@ -302,11 +317,16 @@ function supportedWorkspaceCapabilities( const listFileFeatures = desired.listFileFeatures?.filter((feature) => registration.supportedWorkspaceListFileFeatures?.includes(feature), ); + const programmaticLanguages = desired.programmaticLanguages?.filter( + (language) => + registration.supportedWorkspaceProgrammaticLanguages?.includes(language), + ); const { writeFileModes: _writeFileModes, editFileModes: _editFileModes, editFileFeatures: _editFileFeatures, listFileFeatures: _listFileFeatures, + programmaticLanguages: _programmaticLanguages, ...compatibleDesired } = desired; return { @@ -327,6 +347,10 @@ function supportedWorkspaceCapabilities( ...(operations.includes('list_files') && listFileFeatures?.length ? { listFileFeatures } : {}), + ...(operations.includes('execute_command') && + programmaticLanguages?.length + ? { programmaticLanguages } + : {}), }, }; } @@ -406,6 +430,16 @@ export class BridgeWorker { 'Workspace tool capabilities require a matching executor', ); } + if ( + (options.workspaceProgrammatic != null) !== + (options.capabilities.workspaceTools?.programmaticLanguages?.includes( + 'bash', + ) === true) + ) { + throw new BridgeProtocolError( + 'Workspace programmatic capability requires a matching executor', + ); + } if ( options.capabilities.workspaceTools?.operations.some( (operation) => @@ -1130,11 +1164,7 @@ export class BridgeWorker { assignment: BridgeAssignment, signal?: AbortSignal, ): Promise { - const root = - assignment.executionKind === 'workspace_tool' && - isWorkspaceToolRequest(assignment.request) - ? assignment.request.workspaceId - : undefined; + const root = this.assignmentWorkspaceId(assignment); const waitingAt = Date.now(); while (root != null && this.activeWorkspaceAssignments.has(root)) { const active = this.activeWorkspaceAssignments.get(root)!; @@ -1212,14 +1242,31 @@ export class BridgeWorker { private workspaceGuard( assignment: BridgeAssignment, ): WorkspaceMutationQuarantine | undefined { - return assignment.executionKind === 'workspace_tool' && - isWorkspaceToolRequest(assignment.request) - ? (this.options.workspaceQuarantines?.get( - assignment.request.workspaceId, - ) ?? this.options.workspaceMutationQuarantine) + const workspaceId = this.assignmentWorkspaceId(assignment); + return workspaceId != null + ? (this.options.workspaceQuarantines?.get(workspaceId) ?? + this.options.workspaceMutationQuarantine) : this.options.workspaceMutationQuarantine; } + private assignmentWorkspaceId( + assignment: BridgeAssignment, + ): string | undefined { + if ( + assignment.executionKind === 'workspace_tool' && + isWorkspaceToolRequest(assignment.request) + ) { + return assignment.request.workspaceId; + } + if ( + assignment.executionKind === 'workspace_programmatic' && + typeof assignment.workspaceId === 'string' + ) { + return assignment.workspaceId; + } + return undefined; + } + private async executeOwned( assignment: BridgeAssignment, signal?: AbortSignal, @@ -1472,6 +1519,76 @@ export class BridgeWorker { 'Bridge assignment expired during workspace execution', ); } + } else if (assignment.executionKind === 'workspace_programmatic') { + const workspaceId = assignment.workspaceId; + if ( + workspaceId == null || + this.options.workspaceProgrammatic == null || + !isBridgeWorkspaceProgrammaticRequest(assignment.request) + ) { + throw new BridgeProtocolError( + 'Worker does not provide valid selected-workspace programmatic execution', + ); + } + try { + if (this.quarantinedWorkspaces.has(workspaceId)) { + throw new Error('Workspace requires an explicit quarantine reset'); + } + if (this.options.workspaceQuarantines != null) + await guard?.assertAvailable(); + } catch (error) { + throw new BridgeWorkspaceQuarantinedError( + 'Workspace is quarantined', + error, + ); + } + const advertised = this.activeCapabilities.workspaceTools; + const workspace = advertised?.workspaces.find( + (candidate) => candidate.id === workspaceId, + ); + if ( + workspace == null || + !advertised?.operations.includes('execute_command') || + (workspace.operations != null && + !workspace.operations.includes('execute_command')) || + !advertised.programmaticLanguages?.includes('bash') + ) { + throw new BridgeProtocolError( + 'Selected-workspace programmatic execution is not advertised', + ); + } + this.mutationGuardArmed = true; + try { + this.armedWorkspaces.add(workspaceId); + await guard!.arm( + 'Workspace programmatic execution is pending settlement', + assignment.assignmentId, + ); + workspaceMutationArmed = true; + } catch (error) { + this.mutationGuardArmed = false; + throw new BridgeWorkspaceQuarantinedError( + 'Workspace mutation quarantine could not be armed before execution', + error, + ); + } + payload = await this.options.workspaceProgrammatic.executeProgrammatic( + workspaceId, + assignment.request, + executionController.signal, + ); + workspaceMutationApplied = true; + if (executionController.signal.aborted) { + throw ( + executionController.signal.reason ?? + new DOMException('aborted', 'AbortError') + ); + } + if (Date.now() >= localDeadlineAtMs) { + throw new BridgeProtocolError( + 'Bridge assignment expired during programmatic execution', + ); + } } else { runtimeLease = await this.runtimeSupervisor.acquire( assignment, @@ -1581,7 +1698,8 @@ export class BridgeWorker { leaseToken: assignment.leaseToken, incarnationId: this.incarnationId, status: 'rejected', - ...(assignment.executionKind === 'workspace_tool' && + ...((assignment.executionKind === 'workspace_tool' || + assignment.executionKind === 'workspace_programmatic') && error instanceof WorkspaceToolError ? { errorCode: error.code } : {}), @@ -1693,12 +1811,8 @@ export class BridgeWorker { clearTimeout(timer); } } - if ( - assignment.executionKind === 'workspace_tool' && - isWorkspaceToolRequest(assignment.request) - ) { - this.armedWorkspaces.delete(assignment.request.workspaceId); - } + const workspaceId = this.assignmentWorkspaceId(assignment); + if (workspaceId != null) this.armedWorkspaces.delete(workspaceId); this.mutationGuardArmed = false; } catch (error) { throw new BridgeWorkspaceQuarantinedError( @@ -1957,11 +2071,12 @@ export class BridgeWorker { const fulfilledWorkspaceMutation = workspaceMutationApplied && settlement.status === 'fulfilled' && - assignment.executionKind === 'workspace_tool' && - isWorkspaceToolRequest(assignment.request) && - (assignment.request.operation === 'write_file' || - assignment.request.operation === 'edit_file' || - assignment.request.operation === 'execute_command'); + (assignment.executionKind === 'workspace_programmatic' || + (assignment.executionKind === 'workspace_tool' && + isWorkspaceToolRequest(assignment.request) && + (assignment.request.operation === 'write_file' || + assignment.request.operation === 'edit_file' || + assignment.request.operation === 'execute_command'))); if (signal?.aborted === true) { if (assignment.runtimeSessionId != null || fulfilledWorkspaceMutation) { throw await this.quarantineWorkspace( diff --git a/packages/code/src/workspace-worker.test.ts b/packages/code/src/workspace-worker.test.ts index 8f9fcd45..fa08a84c 100644 --- a/packages/code/src/workspace-worker.test.ts +++ b/packages/code/src/workspace-worker.test.ts @@ -950,6 +950,88 @@ test('worker executes a workspace tool assignment locally without acquiring a sa }); }); +test('worker executes programmatic Bash in the selected workspace and preserves its fence', async () => { + const programmaticRequests: object[] = []; + const quarantineEvents: string[] = []; + const workspaceCapabilities = { + protocolVersion: 1 as const, + operations: ['execute_command' as const], + programmaticLanguages: ['bash' as const], + workspaces: [{ id: 'primary' }], + }; + 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: false, + sandboxProfile: 'anthropic-srt', + runtimes: [], + workspaceTools: workspaceCapabilities, + }, + workspaceTools: { + capabilities: workspaceCapabilities, + mutationFailuresAreAtomic: true, + async execute() { + throw new Error('workspace tool executor must not run'); + }, + }, + workspaceProgrammatic: { + async executeProgrammatic(workspaceId, request) { + programmaticRequests.push({ workspaceId, request }); + return { + session_id: 'session-1', + language: 'bash', + version: '5.2', + files: [], + run: { stdout: 'ready\n', stderr: '', code: 0, signal: null }, + }; + }, + }, + workspaceQuarantines: new Map([ + [ + 'primary', + mutationQuarantine( + (reason) => quarantineEvents.push(`quarantine:${reason}`), + (reason) => quarantineEvents.push(`arm:${reason}`), + () => quarantineEvents.push('clear'), + ), + ], + ]), + fetchImpl: async () => Response.json({ protocolVersion: 1, accepted: true }), + }); + const request = { + body: { + language: 'bash' as const, + version: '5.2', + session_id: 'session-1', + files: [{ name: 'main.sh', content: 'echo ready' }], + }, + headers: {}, + }; + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-programmatic-1', + 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_programmatic', + workspaceId: 'primary', + request, + }); + + assert.deepEqual(programmaticRequests, [{ workspaceId: 'primary', request }]); + assert.deepEqual(quarantineEvents, [ + 'arm:Workspace programmatic execution is pending settlement', + 'clear', + ]); +}); + test('worker stops after Code API rejects a fulfilled workspace mutation', async () => { let quarantinedReason: string | undefined; let armed = 0; diff --git a/packages/code/src/workspace.ts b/packages/code/src/workspace.ts index 5bfc07fe..dfda49fc 100644 --- a/packages/code/src/workspace.ts +++ b/packages/code/src/workspace.ts @@ -112,6 +112,8 @@ export interface SandboxWorkspaceToolsOptions { commandSandbox: WorkspaceCommandSandbox; /** Workspace IDs whose sandbox is configured and may run commands. */ commandWorkspaces: string[]; + /** Optional execution-scoped languages supplied by the same command sandbox. */ + programmaticLanguages?: BridgeWorkspaceToolCapabilities['programmaticLanguages']; } const MAX_SEARCH_CANDIDATE_BYTES = 1024 * 1024; @@ -1657,6 +1659,9 @@ export class SandboxWorkspaceTools implements WorkspaceToolExecutor { ...(base.listFileFeatures != null ? { listFileFeatures: base.listFileFeatures } : {}), + ...(options.programmaticLanguages?.length + ? { programmaticLanguages: [...options.programmaticLanguages] } + : {}), workspaces: base.workspaces.map((workspace) => ({ ...workspace, operations: [ diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index 25b0bdaf..369b306c 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -414,6 +414,7 @@ router.post( supportedWorkspaceEditFileModes: ['single', 'batch'], supportedWorkspaceEditFileFeatures: ['expected_base_sha256'], supportedWorkspaceListFileFeatures: ['after_path'], + supportedWorkspaceProgrammaticLanguages: ['bash'], }); } catch (error) { if (error instanceof BridgeStoreError) { diff --git a/service/src/bridge/selection.ts b/service/src/bridge/selection.ts index 0959279f..f3926e8a 100644 --- a/service/src/bridge/selection.ts +++ b/service/src/bridge/selection.ts @@ -1,4 +1,5 @@ export const CODEAPI_BRIDGE_WORKER_HEADER = 'X-LibreChat-Code-Worker-ID'; +export const CODEAPI_BRIDGE_WORKSPACE_HEADER = 'X-LibreChat-Code-Workspace-ID'; export const BRIDGE_WORKER_ID_PATTERN = /^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$/; export class BridgeWorkerSelectionError extends Error { diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index b09d3792..d6b91469 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -151,6 +151,26 @@ function supportsWorkspaceTool( return true; } +function supportsWorkspaceProgrammatic( + registration: RegisteredBridgeWorker, + workspaceId: string, + language: string, +): boolean { + const capabilities = registration.capabilities.workspaceTools; + const workspace = capabilities?.workspaces.find( + (candidate) => candidate.id === workspaceId, + ); + return ( + workspace != null && + capabilities?.operations.includes('execute_command') === true && + (workspace.operations == null || + workspace.operations.includes('execute_command')) && + capabilities.programmaticLanguages?.includes( + language as 'bash', + ) === true + ); +} + function workerKey(workerId: string): string { return `${PREFIX}:worker:${encodeURIComponent(workerId)}`; } @@ -700,6 +720,7 @@ export class RedisBridgeStore { body: t.PayloadBody; headers: Record; workspaceRequest?: WorkspaceToolRequest; + workspaceId?: string; runtimeSessionId?: string; deadlineAtMs: number; executionTimeoutMs?: number; @@ -709,6 +730,12 @@ export class RedisBridgeStore { registration: RegisteredBridgeWorker, ) => Promise; }): Promise { + if (args.workspaceRequest != null && args.workspaceId != null) { + throw new BridgeStoreError( + 'ASSIGNMENT_INVALID', + 'A bridge assignment cannot be both a workspace tool and programmatic execution', + ); + } if ( args.executionTimeoutMs !== undefined && (args.workspaceRequest == null || @@ -773,6 +800,19 @@ export class RedisBridgeStore { `Bridge worker ${args.workerId} does not advertise the requested workspace tool`, ); } + if ( + args.workspaceId != null && + !supportsWorkspaceProgrammatic( + registration, + args.workspaceId, + args.body.language, + ) + ) { + throw new BridgeStoreError( + 'WORKER_MISMATCH', + `Bridge worker ${args.workerId} does not advertise programmatic execution for the selected workspace`, + ); + } if ( args.runtimeSessionId !== undefined && (await this.dispatchCommand( @@ -799,14 +839,16 @@ export class RedisBridgeStore { const lockIncarnationId = registration.incarnationId; let assignment: StoredAssignment | undefined; let workspaceLeaseSlot: number | undefined; + const selectedWorkspaceId = + args.workspaceRequest?.workspaceId ?? args.workspaceId; const workspaceSlots = - args.workspaceRequest != null && + selectedWorkspaceId != null && (registration.capabilities.workspaceLeaseSlots ?? 1) > 1 ? new BridgeWorkspaceSlots(this.redis) : undefined; let resultCommitted = false; const admission = - args.workspaceRequest == null + selectedWorkspaceId == null ? undefined : new BridgeAdmissionQueue(this.redis); try { @@ -820,7 +862,7 @@ export class RedisBridgeStore { args.deadlineAtMs, workspaceSlots == null ? undefined - : args.workspaceRequest?.workspaceId, + : selectedWorkspaceId, ), args, 'Bridge admission enqueue', @@ -855,7 +897,7 @@ export class RedisBridgeStore { workerId: args.workerId, incarnationId: lockIncarnationId, assignmentId, - workspaceId: args.workspaceRequest!.workspaceId, + workspaceId: selectedWorkspaceId!, capacity: registration.capabilities.workspaceLeaseSlots!, expiresAtMs: Date.now() + ttlSeconds * 1000, }), @@ -910,7 +952,14 @@ export class RedisBridgeStore { ); } if ( - !supportsWorkspaceTool(current.registration, args.workspaceRequest!) + (args.workspaceRequest != null && + !supportsWorkspaceTool(current.registration, args.workspaceRequest)) || + (args.workspaceId != null && + !supportsWorkspaceProgrammatic( + current.registration, + args.workspaceId, + args.body.language, + )) ) { throw new BridgeStoreError( 'WORKER_MISMATCH', @@ -935,11 +984,14 @@ export class RedisBridgeStore { generation, leaseToken, leaseTokenHash: tokenHash(leaseToken), + ...(selectedWorkspaceId == null ? {} : { + workspaceFence: `native-workspace:${selectedWorkspaceId}`, + }), ...(workspaceLeaseSlot === undefined ? {} : { workspaceLeaseSlot, - workspaceFence: `native-workspace:${args.workspaceRequest!.workspaceId}`, + workspaceFence: `native-workspace:${selectedWorkspaceId!}`, }), ...(registration.identityId != null ? { workerIdentityId: registration.identityId } @@ -951,6 +1003,15 @@ export class RedisBridgeStore { executionKind: 'workspace_tool' as const, request: args.workspaceRequest, } + : args.workspaceId != null + ? { + executionKind: 'workspace_programmatic' as const, + workspaceId: args.workspaceId, + request: { + body: args.body, + headers: args.headers, + }, + } : { request: { body: args.body, @@ -1011,6 +1072,19 @@ export class RedisBridgeStore { `Bridge worker ${args.workerId} no longer advertises the requested workspace tool`, ); } + if ( + args.workspaceId != null && + !supportsWorkspaceProgrammatic( + replacement.registration, + args.workspaceId, + args.body.language, + ) + ) { + throw new BridgeStoreError( + 'WORKER_MISMATCH', + `Bridge worker ${args.workerId} no longer advertises programmatic execution for the selected workspace`, + ); + } registration = replacement.registration; readyToken = replacement.readyToken; } @@ -1034,11 +1108,24 @@ export class RedisBridgeStore { resultCommitted = true; return result; } catch (error) { - if (args.runtimeSessionId !== undefined) { + if (assignment.workspaceFence != null) { + // Native roots retain their own fence through result restoration. + // Do not quarantine unrelated roots or invalidate the worker lease. + await boundedCommand(this.redis.eval( + [ + "if redis.call('GET', KEYS[1]) ~= ARGV[1] then return 0 end", + "redis.call('SET', KEYS[1], 'quarantined:' .. ARGV[1])", + 'return 1', + ].join('\n'), + 1, + workspaceQuarantineKey(args.workerId, assignment.workspaceFence), + assignment.assignmentId, + ), this.redisCommandTimeoutMs, 'Bridge native workspace finalization quarantine'); + } else if (assignmentWorkspace(assignment) !== undefined) { await this.quarantine( args.workerId, assignment.incarnationId, - args.runtimeSessionId, + assignmentWorkspace(assignment)!, ); } throw error; @@ -1904,10 +1991,11 @@ export class RedisBridgeStore { : undefined; const cancelledMutation = signal.aborted && - workspaceRequest != null && - (workspaceRequest.operation === 'write_file' || - workspaceRequest.operation === 'edit_file' || - workspaceRequest.operation === 'execute_command'); + (assignment.executionKind === 'workspace_programmatic' || + (workspaceRequest != null && + (workspaceRequest.operation === 'write_file' || + workspaceRequest.operation === 'edit_file' || + workspaceRequest.operation === 'execute_command'))); if (cancelledMutation) { try { // Keep the acknowledged assignment available long enough for the diff --git a/service/src/bridge/workspace-store.test.ts b/service/src/bridge/workspace-store.test.ts index 5b0186d0..f248aa40 100644 --- a/service/src/bridge/workspace-store.test.ts +++ b/service/src/bridge/workspace-store.test.ts @@ -77,6 +77,123 @@ test('dispatches a workspace tool only to a worker advertising its workspace and }); }); +for (const finalizationFails of [false, true]) test(`single-slot programmatic finalization retains the workspace fence (failure=${finalizationFails})`, async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'workspace-worker', + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'anthropic-srt', + runtimes: [], + workspaceTools: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operations: ['execute_command'], + programmaticLanguages: ['bash'], + workspaces: [{ id: 'primary' }], + }, + }, + }); + const body = { + language: 'bash', + version: '5.2', + session_id: 'session-1', + files: [{ name: 'main.sh', content: 'echo ready' }], + }; + const completion = store.dispatch({ + workerId: 'workspace-worker', + body, + headers: {}, + workspaceId: 'primary', + deadlineAtMs: Date.now() + 5_000, + signal: new AbortController().signal, + finalize: async settlement => { + if (finalizationFails) throw new Error('artifact restoration failed'); + return settlement; + }, + }); + + const assignment = await store.lease('workspace-worker', incarnationId, 1_000); + expect(assignment).toMatchObject({ + executionKind: 'workspace_programmatic', + workspaceId: 'primary', + request: { body }, + }); + await store.settle('workspace-worker', assignment?.assignmentId ?? '', { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: assignment?.generation ?? 0, + leaseToken: assignment?.leaseToken ?? '', + incarnationId, + status: 'fulfilled', + result: { + session_id: 'session-1', + language: 'bash', + version: '5.2', + files: [], + run: { + stdout: 'ready\n', + stderr: '', + code: 0, + signal: null, + output: 'ready\n', + memory: null, + message: null, + status: null, + cpu_time: null, + wall_time: 0.01, + }, + }, + }); + + if (finalizationFails) { + await expect(completion).rejects.toThrow('artifact restoration failed'); + await expect(store.dispatch({ workerId: 'workspace-worker', body, headers: {}, + workspaceId: 'primary', deadlineAtMs: Date.now() + 1000, + signal: new AbortController().signal, + })).rejects.toMatchObject({ code: 'WORKSPACE_QUARANTINED' }); + return; + } + await expect(completion).resolves.toMatchObject({ + status: 'fulfilled', + result: { session_id: 'session-1' }, + }); +}); + +test('rejects programmatic execution without the workspace capability', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'workspace-worker', + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'anthropic-srt', + runtimes: [], + workspaceTools: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operations: ['execute_command'], + workspaces: [{ id: 'primary' }], + }, + }, + }); + + await expect( + store.dispatch({ + workerId: 'workspace-worker', + body: { + language: 'bash', + version: '5.2', + session_id: 'session-2', + files: [{ name: 'main.sh', content: 'echo denied' }], + }, + headers: {}, + workspaceId: 'primary', + deadlineAtMs: Date.now() + 1_000, + signal: new AbortController().signal, + }), + ).rejects.toMatchObject({ code: 'WORKER_MISMATCH' }); + expect(await redis.keys('codeapi:bridge:v1:assignment:*')).toHaveLength(0); +}); + test('drains an acknowledged workspace mutation cancellation before releasing it', async () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, diff --git a/service/src/egress-gateway.ts b/service/src/egress-gateway.ts index f86ce656..873db173 100644 --- a/service/src/egress-gateway.ts +++ b/service/src/egress-gateway.ts @@ -45,6 +45,7 @@ import { parseBoundedContentLength } from './http-limits'; import { validateEgressGatewayHardenedConfig } from './secure-startup'; import { isOpaqueObjectContentDisposition } from './file-metadata'; import { mapObjectDetails } from './file-object-resolver'; +import { isSupportedBridgeArtifactName } from '../../packages/code/src/protocol'; export const app: Express = express(); app.disable('x-powered-by'); @@ -52,29 +53,6 @@ validateEgressGatewayHardenedConfig(); app.use(traceHttpRequest('codeapi.egress_gateway.request')); app.use(httpMetricsMiddleware); -const SUPPORTED_OUTPUT_EXTENSIONS = new Set([ - '.c', '.cs', '.cpp', '.go', '.java', '.js', '.kt', '.kts', '.lua', - '.php', '.pl', '.ps1', '.py', '.r', '.rb', '.rs', '.scala', '.sh', - '.sql', '.swift', '.ts', '.jsx', '.tsx', '.groovy', - '.css', '.htm', '.html', '.less', '.sass', '.scss', '.svg', '.svelte', '.vue', - '.adoc', '.asciidoc', '.md', '.rst', '.tex', '.txt', '.wiki', - '.csv', '.json', '.bson', '.json5', '.jsonl', '.parquet', '.tsv', - '.xml', '.yaml', '.yml', - '.ics', '.ical', '.ifb', '.icalendar', - '.conf', '.env', '.gitignore', '.ini', '.properties', '.toml', - '.doc', '.docx', '.pdf', '.ppt', '.pptx', '.xls', '.xlsx', - '.odt', '.ods', '.odp', '.rtf', - '.avif', '.bmp', '.gif', '.ico', '.jpeg', '.jpg', '.png', - '.tif', '.tiff', '.webp', - '.eot', '.ttf', '.woff', '.woff2', - '.7z', '.bz2', '.gz', '.gzip', '.rar', '.tar', '.zip', - '.tf', '.tfvars', '.tfstate', '.hcl', - '.dockerfile', '.Dockerfile', '.dockerignore', - '.helmignore', '.helmfile', '.jenkinsfile', '.vagrantfile', - '.eslintrc', '.prettierrc', '.editorconfig', '.nomad', - '.bat', '.cmd', '.deb', '.log', '.rpm', '.vbs', -]); - type EgressAuditFields = { execHash?: string; requestExecHash?: string; @@ -248,18 +226,7 @@ function assertOutputFilenameAllowed(name: string): void { throw new EgressGrantError('malformed', 'Output filename must be canonical'); } if (!isDirkeepName(name)) { - const basename = path.posix.basename(name); - const ext = path.posix.extname(basename).toLowerCase(); - const dottedBasename = `.${basename}`; - const allowed = - (ext !== '' && SUPPORTED_OUTPUT_EXTENSIONS.has(ext)) || - SUPPORTED_OUTPUT_EXTENSIONS.has(basename) || - SUPPORTED_OUTPUT_EXTENSIONS.has(basename.toLowerCase()) || - (ext === '' && ( - SUPPORTED_OUTPUT_EXTENSIONS.has(dottedBasename) || - SUPPORTED_OUTPUT_EXTENSIONS.has(dottedBasename.toLowerCase()) - )); - if (!allowed) { + if (!isSupportedBridgeArtifactName(name)) { throw new EgressGrantError('scope_mismatch', 'Output filename extension is not supported'); } } diff --git a/service/src/egress-grant.test.ts b/service/src/egress-grant.test.ts index 32ef2988..cd26fcd7 100644 --- a/service/src/egress-grant.test.ts +++ b/service/src/egress-grant.test.ts @@ -5,6 +5,7 @@ import { env } from './config'; import { normalizeEgressGatewayUrl, normalizeProgrammaticTimeoutMs, + normalizeSelectedWorkspaceProgrammaticTimeoutMs, prepareSandboxJobSecurity, refreshEgressGrantClaims, timeoutMsToGrantSeconds, @@ -478,6 +479,13 @@ describe('egress encrypted grants and handles', () => { expect(() => normalizeProgrammaticTimeoutMs(0, 300000)).toThrow('timeout must be a positive number'); }); + test('budgets both selected-workspace replay passes inside the worker deadline', () => { + expect(normalizeSelectedWorkspaceProgrammaticTimeoutMs(undefined, 300_000)).toBe(82_500); + expect(normalizeSelectedWorkspaceProgrammaticTimeoutMs(120_000, 300_000)).toBe(82_500); + expect(normalizeSelectedWorkspaceProgrammaticTimeoutMs(300_000, 300_000)).toBe(82_500); + expect(normalizeSelectedWorkspaceProgrammaticTimeoutMs(10_000, 20_000)).toBe(2_167); + }); + test('normalizes the gateway callback URL for sandbox-originated PTC', () => { expect(normalizeEgressGatewayUrl(' http://egress-gateway:3190/// ')).toBe('http://egress-gateway:3190'); expect(() => normalizeEgressGatewayUrl(' ')).toThrow('EGRESS_GATEWAY_URL is required'); diff --git a/service/src/preamble-bash.test.ts b/service/src/preamble-bash.test.ts index 815040d3..e35a77f5 100644 --- a/service/src/preamble-bash.test.ts +++ b/service/src/preamble-bash.test.ts @@ -1,10 +1,19 @@ import { execFileSync } from 'child_process'; -import { mkdtempSync, rmSync, writeFileSync } from 'fs'; +import { + mkdirSync, + mkdtempSync, + readFileSync, + rmSync, + writeFileSync, +} from 'fs'; import { tmpdir } from 'os'; import { join } from 'path'; import { describe, expect, test } from 'bun:test'; import { extractPendingFromStdout, type LCTool } from './preamble'; -import { generateBashReplayPostamble, generateBashReplayPreamble } from './preamble-bash'; +import { + generateBashReplayPostamble, + generateBashReplayPreamble, +} from './preamble-bash'; interface BashRunResult { stdout: string; @@ -60,9 +69,13 @@ function assemble(userCode: string, toolSet: LCTool[] = tools): string { ].join('\n'); } -function runBash(script: string, options: number | BashRunOptions = {}): BashRunResult { - const timeoutMs = typeof options === 'number' ? options : options.timeoutMs ?? 3000; - const history = typeof options === 'number' ? {} : options.history ?? {}; +function runBash( + script: string, + options: number | BashRunOptions = {}, +): BashRunResult { + const timeoutMs = + typeof options === 'number' ? options : (options.timeoutMs ?? 3000); + const history = typeof options === 'number' ? {} : (options.history ?? {}); const dir = mkdtempSync(join(tmpdir(), 'ptc-bash-unit-')); const file = join(dir, 'main.sh'); const historyPath = join(dir, 'history.json'); @@ -99,46 +112,136 @@ function pendingNames(stdout: string): string[] { return (parsed.pending ?? []).map(call => call.tool_name).sort(); } +describe('generateBashReplayPreamble - private runtime directory', () => { + test('creates every replay tempfile beneath TMPDIR', () => { + const dir = mkdtempSync(join(tmpdir(), 'ptc-bash-private-tmp-')); + const dataDir = join(dir, 'data'); + const runtimeDir = join(dir, 'runtime'); + const file = join(dir, 'main.sh'); + const historyPath = join(dataDir, 'history.json'); + mkdirSync(dataDir, { recursive: true }); + mkdirSync(runtimeDir, { recursive: true }); + writeFileSync(historyPath, '{}'); + writeFileSync( + file, + assemble(` +printf '%s\\n' "$_PTC_PENDING_FILE" "$_PTC_ERROR_FILE" "$_PTC_COUNTER_FILE" +`), + { mode: 0o755 }, + ); + + try { + const stdout = execFileSync('bash', [file], { + env: { + ...process.env, + PTC_HISTORY_PATH: historyPath, + TMPDIR: runtimeDir, + }, + encoding: 'utf8', + }); + const paths = stdout.trim().split('\n'); + expect(paths).toHaveLength(3); + expect( + paths.every(value => value.startsWith(`${runtimeDir}/`)), + ).toBe(true); + expect( + generateBashReplayPreamble({ executionId, tools }), + ).not.toContain('mktemp -t'); + } finally { + rmSync(dir, { recursive: true, force: true }); + } + }); + + test('persists pending calls through the private native control path', () => { + const dir = mkdtempSync(join(tmpdir(), 'ptc-bash-control-')); + const file = join(dir, 'main.sh'); + const historyPath = join(dir, 'history.json'); + const controlPath = join(dir, 'control.json'); + writeFileSync(file, assemble(`get_weather '{"city":"Paris"}'`), { + mode: 0o755, + }); + writeFileSync(historyPath, '{}'); + try { + execFileSync('bash', [file], { + env: { + ...process.env, + PTC_HISTORY_PATH: historyPath, + LIBRECHAT_CODE_CONTROL_PATH: controlPath, + TMPDIR: dir, + }, + encoding: 'utf8', + }); + expect(JSON.parse(readFileSync(controlPath, 'utf8'))).toMatchObject( + { + pending: [ + { + call_id: 'call_001', + tool_name: 'get_weather', + input: { city: 'Paris' }, + }, + ], + }, + ); + } finally { + rmSync(dir, { recursive: true, force: true }); + } + }); +}); + describe('generateBashReplayPreamble - command substitution pending emission', () => { test('emits ClickHouse-style object input with SQL quotes from double-quoted JSON', () => { - const run = runBash(assemble(` + const run = runBash( + assemble( + ` SVC="45886e06-932b-4cff-bb49-3f7281d80717" result=$(run_select_query_mcp_ClickHouse "{\\"serviceId\\":\\"$SVC\\",\\"query\\":\\"SELECT name, round(avg(tempAvg)/10.0, 2) AS avg_temp_c FROM system.columns WHERE database='default' AND table='uk_prices_3' AND tempAvg != -9999\\"}") echo "AFTER: $result" -`, [clickHouseTool])); +`, + [clickHouseTool], + ), + ); const parsed = extractPendingFromStdout(run.stdout, executionId); expect(run.exitCode).toBe(0); expect(parsed.pending).toHaveLength(1); - expect(parsed.pending?.[0]?.tool_name).toBe('run_select_query_mcp_ClickHouse'); + expect(parsed.pending?.[0]?.tool_name).toBe( + 'run_select_query_mcp_ClickHouse', + ); expect(parsed.pending?.[0]?.input).toEqual({ serviceId: '45886e06-932b-4cff-bb49-3f7281d80717', - query: - "SELECT name, round(avg(tempAvg)/10.0, 2) AS avg_temp_c FROM system.columns WHERE database='default' AND table='uk_prices_3' AND tempAvg != -9999", + query: "SELECT name, round(avg(tempAvg)/10.0, 2) AS avg_temp_c FROM system.columns WHERE database='default' AND table='uk_prices_3' AND tempAvg != -9999", }); expect(parsed.stdout).not.toContain('AFTER'); }); test('emits ClickHouse-style object input with shell-escaped SQL quotes', () => { - const run = runBash(assemble(` + const run = runBash( + assemble( + ` result=$(run_select_query_mcp_ClickHouse '{"serviceId":"45886e06-932b-4cff-bb49-3f7281d80717","query":"SELECT name, type FROM system.columns WHERE database='"'"'default'"'"' AND table='"'"'uk_prices_3'"'"' ORDER BY position"}') echo "AFTER: $result" -`, [clickHouseTool])); +`, + [clickHouseTool], + ), + ); const parsed = extractPendingFromStdout(run.stdout, executionId); expect(run.exitCode).toBe(0); expect(parsed.pending).toHaveLength(1); - expect(parsed.pending?.[0]?.tool_name).toBe('run_select_query_mcp_ClickHouse'); + expect(parsed.pending?.[0]?.tool_name).toBe( + 'run_select_query_mcp_ClickHouse', + ); expect(parsed.pending?.[0]?.input).toEqual({ serviceId: '45886e06-932b-4cff-bb49-3f7281d80717', - query: - "SELECT name, type FROM system.columns WHERE database='default' AND table='uk_prices_3' ORDER BY position", + query: "SELECT name, type FROM system.columns WHERE database='default' AND table='uk_prices_3' ORDER BY position", }); expect(parsed.stdout).not.toContain('AFTER'); }); test('batches parallel ClickHouse-style command substitutions into one pending block', () => { - const run = runBash(assemble(` + const run = runBash( + assemble( + ` SVC="45886e06-932b-4cff-bb49-3f7281d80717" { @@ -158,7 +261,11 @@ SVC="45886e06-932b-4cff-bb49-3f7281d80717" wait echo "AFTER" -`, [clickHouseTool]), 3000); +`, + [clickHouseTool], + ), + 3000, + ); const parsed = extractPendingFromStdout(run.stdout, executionId); expect(run.exitCode).toBe(0); @@ -168,7 +275,11 @@ echo "AFTER" 'run_select_query_mcp_ClickHouse', 'run_select_query_mcp_ClickHouse', ]); - expect(parsed.pending?.map(call => (call.input as { query: string }).query).sort()).toEqual([ + expect( + parsed.pending + ?.map(call => (call.input as { query: string }).query) + .sort(), + ).toEqual([ "SELECT name, engine, total_rows, formatReadableSize(total_bytes) AS size, sorting_key, partition_key FROM system.tables WHERE database='default' AND name IN ('uk_prices_3','weather_noaa_mt')", "SELECT name, type, comment FROM system.columns WHERE database='default' AND table='uk_prices_3' ORDER BY position", "SELECT name, type, comment FROM system.columns WHERE database='default' AND table='weather_noaa_mt' ORDER BY position", @@ -177,12 +288,14 @@ echo "AFTER" }); test('emits a command-substitution tool call before later user code while another job is running', () => { - const run = runBash(assemble(` + const run = runBash( + assemble(` sleep 0.2 & result=$(get_weather '{"city":"Madrid"}') echo "AFTER: $result" wait -`)); +`), + ); const parsed = extractPendingFromStdout(run.stdout, executionId); expect(run.signal).not.toBe('SIGTERM'); @@ -194,12 +307,15 @@ wait }); test('batches background and command-substitution tool calls before command-substitution side effects', () => { - const run = runBash(assemble(` + const run = runBash( + assemble(` get_weather '{"city":"Oslo"}' & result=$(calculate '{"expression":"2+3"}') echo "SIDE_EFFECT: $result" wait -`), 1500); +`), + 1500, + ); const parsed = extractPendingFromStdout(run.stdout, executionId); expect(run.signal).not.toBe('SIGTERM'); @@ -210,10 +326,13 @@ wait }); test('waits for background compound commands that invoke tools later', () => { - const run = runBash(assemble(` + const run = runBash( + assemble(` (sleep 0.2; get_weather '{"city":"Paris"}') & echo "AFTER LAUNCH" -`), 1500); +`), + 1500, + ); const parsed = extractPendingFromStdout(run.stdout, executionId); expect(run.signal).not.toBe('SIGTERM'); @@ -225,10 +344,13 @@ echo "AFTER LAUNCH" }); test('does not wait for unrelated background commands with tool names as arguments', () => { - const run = runBash(assemble(` + const run = runBash( + assemble(` bash -c 'sleep 2' get_weather & echo "DONE" -`), 700); +`), + 700, + ); const parsed = extractPendingFromStdout(run.stdout, executionId); expect(run.signal).not.toBe('SIGTERM'); @@ -238,14 +360,17 @@ echo "DONE" }); test('does not treat arithmetic expansion as command substitution while batching background tools', () => { - const run = runBash(assemble(` + const run = runBash( + assemble(` get_weather '{"city":"Oslo"}' & sleep 0.1 x=$((1+1)) calculate '{"expression":"2+3"}' & wait echo "DONE $x" -`), 1500); +`), + 1500, + ); const parsed = extractPendingFromStdout(run.stdout, executionId); expect(run.signal).not.toBe('SIGTERM'); @@ -256,12 +381,15 @@ echo "DONE $x" }); test('handles backtick command substitution without waiting for unrelated background jobs', () => { - const run = runBash(assemble(` + const run = runBash( + assemble(` sleep 5 & result=\`get_weather '{"city":"Porto"}'\` echo "AFTER: $result" wait -`), 1500); +`), + 1500, + ); const parsed = extractPendingFromStdout(run.stdout, executionId); expect(run.signal).not.toBe('SIGTERM'); @@ -281,7 +409,10 @@ wait echo "DONE" `; const firstRun = runBash(assemble(userCode)); - const firstParsed = extractPendingFromStdout(firstRun.stdout, executionId); + const firstParsed = extractPendingFromStdout( + firstRun.stdout, + executionId, + ); expect(firstRun.exitCode).toBe(0); expect(firstParsed.pending).toHaveLength(2); @@ -302,7 +433,10 @@ echo "DONE" }), ); const replayRun = runBash(assemble(userCode), { history }); - const replayParsed = extractPendingFromStdout(replayRun.stdout, executionId); + const replayParsed = extractPendingFromStdout( + replayRun.stdout, + executionId, + ); expect(replayRun.exitCode).toBe(0); expect(replayParsed.pending).toBeNull(); expect(replayParsed.stdout).toContain('"slot":"first"'); @@ -326,12 +460,15 @@ echo "DONE" }, }; - const run = runBash(assemble(` + const run = runBash( + assemble(` get_weather '{"city":"Paris"}' printf '\\n' get_weather '{"city":"Paris"}' printf '\\nDONE\\n' -`), { history }); +`), + { history }, + ); const parsed = extractPendingFromStdout(run.stdout, executionId); expect(run.exitCode).toBe(0); @@ -352,12 +489,15 @@ printf '\\nDONE\\n' }, }; - const run = runBash(assemble(` + const run = runBash( + assemble(` get_weather '{"city":"Paris"}' printf '\\n' calculate '{"expression":"2+3"}' printf '\\nDONE\\n' -`), { history }); +`), + { history }, + ); const parsed = extractPendingFromStdout(run.stdout, executionId); expect(run.signal).not.toBe('SIGTERM'); @@ -385,12 +525,15 @@ printf '\\nDONE\\n' }, }; - const run = runBash(assemble(` + const run = runBash( + assemble(` get_weather '{"city":"Paris"}' printf '\\n' calculate '{"expression":"2+3"}' printf '\\nDONE\\n' -`), { history }); +`), + { history }, + ); const parsed = extractPendingFromStdout(run.stdout, executionId); expect(run.exitCode).toBe(0); @@ -418,12 +561,15 @@ printf '\\nDONE\\n' }, }; - const run = runBash(assemble(` + const run = runBash( + assemble(` get_weather '{"city":"Paris"}' printf '\\n' calculate '{"expression":"2+3"}' printf '\\nDONE\\n' -`), { history }); +`), + { history }, + ); const parsed = extractPendingFromStdout(run.stdout, executionId); expect(run.exitCode).toBe(0); diff --git a/service/src/preamble-bash.ts b/service/src/preamble-bash.ts index 4ec3fc05..c442b495 100644 --- a/service/src/preamble-bash.ts +++ b/service/src/preamble-bash.ts @@ -1,8 +1,5 @@ import type { LCTool } from './preamble'; -import { - buildScopedSentinel, - PTC_HISTORY_SANDBOX_PATH, -} from './ptc-constants'; +import { buildScopedSentinel, PTC_HISTORY_SANDBOX_PATH } from './ptc-constants'; export interface BashReplayPreambleConfig { executionId: string; @@ -42,12 +39,55 @@ export class BashToolNameCollisionError extends Error { } const BASH_RESERVED = new Set([ - 'if', 'then', 'else', 'elif', 'fi', 'case', 'esac', 'for', 'select', - 'while', 'until', 'do', 'done', 'in', 'function', 'time', 'coproc', - 'return', 'exit', 'break', 'continue', 'shift', 'export', 'readonly', - 'local', 'declare', 'typeset', 'unset', 'alias', 'unalias', 'source', - 'echo', 'printf', 'read', 'cd', 'pwd', 'kill', 'trap', 'wait', 'eval', - 'exec', 'jobs', 'bg', 'fg', 'set', 'let', 'test', 'true', 'false', + 'if', + 'then', + 'else', + 'elif', + 'fi', + 'case', + 'esac', + 'for', + 'select', + 'while', + 'until', + 'do', + 'done', + 'in', + 'function', + 'time', + 'coproc', + 'return', + 'exit', + 'break', + 'continue', + 'shift', + 'export', + 'readonly', + 'local', + 'declare', + 'typeset', + 'unset', + 'alias', + 'unalias', + 'source', + 'echo', + 'printf', + 'read', + 'cd', + 'pwd', + 'kill', + 'trap', + 'wait', + 'eval', + 'exec', + 'jobs', + 'bg', + 'fg', + 'set', + 'let', + 'test', + 'true', + 'false', ]); function normalizeBashFunctionName(name: string): string { @@ -60,10 +100,7 @@ function normalizeBashFunctionName(name: string): string { * the end-of-preamble `readonly -f` lockdown runs. Compared case- * insensitively because the `_PTC_` prefix is used for variables and * `_ptc_` for functions, and both live in the same identifier space. */ - if ( - BASH_RESERVED.has(normalized) || - /^_ptc_/i.test(normalized) - ) { + if (BASH_RESERVED.has(normalized) || /^_ptc_/i.test(normalized)) { normalized = normalized + '_tool'; } if (normalized === '') normalized = 'tool'; @@ -95,9 +132,12 @@ function escapeForBashEre(s: string): string { * Users capture results via command substitution; input is passed as a single * JSON object string argument (validated by jq). */ -export function generateBashReplayPreamble(config: BashReplayPreambleConfig): string { +export function generateBashReplayPreamble( + config: BashReplayPreambleConfig, +): string { const { executionId, tools } = config; - const { start: scopedStart, end: scopedEnd } = buildScopedSentinel(executionId); + const { start: scopedStart, end: scopedEnd } = + buildScopedSentinel(executionId); let preamble = `#!/bin/bash # ============================================================================ @@ -109,20 +149,25 @@ _PTC_EXECUTION_ID="${executionId}" _PTC_SENTINEL_START="${scopedStart}" _PTC_SENTINEL_END="${scopedEnd}" _PTC_HISTORY_PATH="\${PTC_HISTORY_PATH:-${PTC_HISTORY_SANDBOX_PATH}}" -_PTC_PENDING_FILE="$(mktemp -t _ptc_pending.XXXXXX 2>/dev/null || mktemp /tmp/_ptc_pending.XXXXXX)" -_PTC_ERROR_FILE="$(mktemp -t _ptc_error.XXXXXX 2>/dev/null || mktemp /tmp/_ptc_error.XXXXXX)" -_PTC_CONSUMED_FILE="$(mktemp -t _ptc_consumed.XXXXXX 2>/dev/null || mktemp /tmp/_ptc_consumed.XXXXXX)" -_PTC_SAW_BARE_TOOL_FILE="$(mktemp -t _ptc_saw_tool.XXXXXX 2>/dev/null || mktemp /tmp/_ptc_saw_tool.XXXXXX)" -_PTC_PRE_TOOL_JOBS_FILE="$(mktemp -t _ptc_pre_tool_jobs.XXXXXX 2>/dev/null || mktemp /tmp/_ptc_pre_tool_jobs.XXXXXX)" -_PTC_PRE_TOOL_JOBS_READY_FILE="$(mktemp -t _ptc_pre_tool_jobs_ready.XXXXXX 2>/dev/null || mktemp /tmp/_ptc_pre_tool_jobs_ready.XXXXXX)" -_PTC_TOOL_JOBS_FILE="$(mktemp -t _ptc_tool_jobs.XXXXXX 2>/dev/null || mktemp /tmp/_ptc_tool_jobs.XXXXXX)" -_PTC_WAIT_RAN_FILE="$(mktemp -t _ptc_wait_ran.XXXXXX 2>/dev/null || mktemp /tmp/_ptc_wait_ran.XXXXXX)" -_PTC_SUPPRESS_SUBSHELL_TOOL_FILE="$(mktemp -t _ptc_suppress_subshell_tool.XXXXXX 2>/dev/null || mktemp /tmp/_ptc_suppress_subshell_tool.XXXXXX)" -_PTC_SUPPRESS_SUBSHELL_TOOL_CLEAR_FILE="$(mktemp -t _ptc_suppress_subshell_tool_clear.XXXXXX 2>/dev/null || mktemp /tmp/_ptc_suppress_subshell_tool_clear.XXXXXX)" +_PTC_CONTROL_PATH="\${LIBRECHAT_CODE_CONTROL_PATH:-}" +_PTC_RUNTIME_DIR="\${TMPDIR:-/tmp}" +_ptc_mktemp() { + mktemp "\${_PTC_RUNTIME_DIR%/}/$1.XXXXXX" +} +_PTC_PENDING_FILE="$(_ptc_mktemp _ptc_pending)" +_PTC_ERROR_FILE="$(_ptc_mktemp _ptc_error)" +_PTC_CONSUMED_FILE="$(_ptc_mktemp _ptc_consumed)" +_PTC_SAW_BARE_TOOL_FILE="$(_ptc_mktemp _ptc_saw_tool)" +_PTC_PRE_TOOL_JOBS_FILE="$(_ptc_mktemp _ptc_pre_tool_jobs)" +_PTC_PRE_TOOL_JOBS_READY_FILE="$(_ptc_mktemp _ptc_pre_tool_jobs_ready)" +_PTC_TOOL_JOBS_FILE="$(_ptc_mktemp _ptc_tool_jobs)" +_PTC_WAIT_RAN_FILE="$(_ptc_mktemp _ptc_wait_ran)" +_PTC_SUPPRESS_SUBSHELL_TOOL_FILE="$(_ptc_mktemp _ptc_suppress_subshell_tool)" +_PTC_SUPPRESS_SUBSHELL_TOOL_CLEAR_FILE="$(_ptc_mktemp _ptc_suppress_subshell_tool_clear)" # Counter must persist across subshells (command substitution) so call_ids # stay deterministic across cached/uncached calls. Bash variables set in a # subshell don't propagate back, so we use a file. -_PTC_COUNTER_FILE="$(mktemp -t _ptc_counter.XXXXXX 2>/dev/null || mktemp /tmp/_ptc_counter.XXXXXX)" +_PTC_COUNTER_FILE="$(_ptc_mktemp _ptc_counter)" _PTC_LOCK_DIR="\${_PTC_PENDING_FILE}.lock" printf '0' > "$_PTC_COUNTER_FILE" : > "$_PTC_CONSUMED_FILE" @@ -243,7 +288,7 @@ _ptc_prune_finished_tool_jobs() { return 0 fi local _ptc_tmp_file - _ptc_tmp_file="$(mktemp -t _ptc_tool_jobs_live.XXXXXX 2>/dev/null || mktemp /tmp/_ptc_tool_jobs_live.XXXXXX)" + _ptc_tmp_file="$(_ptc_mktemp _ptc_tool_jobs_live)" while IFS= read -r _ptc_pid; do [ -n "$_ptc_pid" ] || continue if kill -0 "$_ptc_pid" 2>/dev/null; then @@ -330,6 +375,17 @@ _ptc_maybe_emit_pending() { trap - DEBUG EXIT exit 1 fi + # Native BYOM workers use this private execution-scoped control file so a + # large stdout stream cannot truncate away the replay frame. Other + # backends continue to consume the stdout sentinel below. + if [ -n "$_PTC_CONTROL_PATH" ]; then + printf '%s' "$_ptc_payload" > "$_PTC_CONTROL_PATH" || { + printf 'failed to persist pending PTC tool calls\n' >&2 + _ptc_cleanup_tempfiles + trap - DEBUG EXIT + exit 1 + } + fi if [ "\${BASH_SUBSHELL:-0}" -eq 1 ]; then trap - DEBUG EXIT exit 0 @@ -512,7 +568,7 @@ _ptc_call_tool() { # Large input can exceed ARG_MAX via --argjson; write once, reuse path below. local _ptc_input_tmp - _ptc_input_tmp="$(mktemp -t _ptc_input.XXXXXX 2>/dev/null || mktemp /tmp/_ptc_input.XXXXXX)" + _ptc_input_tmp="$(_ptc_mktemp _ptc_input)" printf '%s' "$_ptc_input" > "$_ptc_input_tmp" local _ptc_matches @@ -668,8 +724,12 @@ exit $_ptc_user_exit_code function generateBashToolStub(tool: LCTool): string { const fnName = normalizeBashFunctionName(tool.name); - const desc = (tool.description ?? '').split('\n').map(l => `# ${l}`).join('\n'); - const nameComment = fnName !== tool.name ? `# Original tool name: ${tool.name}\n` : ''; + const desc = (tool.description ?? '') + .split('\n') + .map(l => `# ${l}`) + .join('\n'); + const nameComment = + fnName !== tool.name ? `# Original tool name: ${tool.name}\n` : ''; const escapedToolName = escapeForBashDoubleQuote(tool.name); return `${nameComment}${desc ? desc + '\n' : ''}${fnName}() { local _default_input='{}' @@ -683,7 +743,8 @@ function generateBashToolStub(tool: LCTool): string { } function generateBashPendingDeferHelper(tools: readonly LCTool[]): string { - const toolNamesPattern = tools + const toolNamesPattern = + tools .map(tool => normalizeBashFunctionName(tool.name)) .map(escapeForBashEre) .join('|') || 'a^'; diff --git a/service/src/preamble.test.ts b/service/src/preamble.test.ts index bf50e3ec..803e1839 100644 --- a/service/src/preamble.test.ts +++ b/service/src/preamble.test.ts @@ -2,6 +2,7 @@ import { describe, expect, test } from 'bun:test'; import { buildScopedSentinel, createProgrammaticPayload, + extractPendingFromControlPayload, extractPendingFromStdout, generatePreamble, } from './preamble'; @@ -33,7 +34,9 @@ describe('generatePreamble — Unix-vs-TCP transport gate', () => { expect(preamble).toMatch(/AF_UNIX/); expect(preamble).toMatch(/\.connect\(_TOOL_CALL_SOCKET\)/); /* Regression guard against reintroducing the user-spoofable check. */ - expect(preamble).not.toMatch(/if\s+os\.path\.exists\(_TOOL_CALL_SOCKET\)/); + expect(preamble).not.toMatch( + /if\s+os\.path\.exists\(_TOOL_CALL_SOCKET\)/, + ); }); test('caches the probe verdict at module load (before user code can plant a spoof)', () => { @@ -41,10 +44,14 @@ describe('generatePreamble — Unix-vs-TCP transport gate', () => { /* The probe call must appear at top level of the preamble, NOT * inside _do_request. Otherwise a user could plant a regular file * at the path between calls and flip the gate per-request. */ - const probeCallIdx = preamble.indexOf('_USE_TOOL_CALL_SOCKET = _probe_tool_call_socket()'); + const probeCallIdx = preamble.indexOf( + '_USE_TOOL_CALL_SOCKET = _probe_tool_call_socket()', + ); expect(probeCallIdx).toBeGreaterThan(-1); /* _do_request must consult the cached verdict, not re-probe. */ - const doReqMatch = preamble.match(/def\s+_do_request[\s\S]*?(?=\ndef\s|\nclass\s|\Z)/); + const doReqMatch = preamble.match( + /def\s+_do_request[\s\S]*?(?=\ndef\s|\nclass\s|\Z)/, + ); expect(doReqMatch).not.toBeNull(); expect(doReqMatch![0]).toContain('_USE_TOOL_CALL_SOCKET'); expect(doReqMatch![0]).not.toContain('_probe_tool_call_socket('); @@ -55,7 +62,9 @@ describe('generatePreamble — Unix-vs-TCP transport gate', () => { /* The fallback must still construct the URL from _CALLBACK_URL and * delegate to _tcp_request. Without this, runners without the * proxy bind-mount would have no way to reach the orchestrator. */ - const doReqMatch = preamble.match(/def\s+_do_request[\s\S]*?(?=\ndef\s|\nclass\s|\Z)/); + const doReqMatch = preamble.match( + /def\s+_do_request[\s\S]*?(?=\ndef\s|\nclass\s|\Z)/, + ); expect(doReqMatch).not.toBeNull(); expect(doReqMatch![0]).toContain('_CALLBACK_URL + path'); expect(doReqMatch![0]).toContain('_tcp_request('); @@ -67,24 +76,47 @@ describe('generatePreamble — Unix-vs-TCP transport gate', () => { * exported. The path must remain hardcoded so the preamble does * not depend on env-var injection. */ expect(preamble).toContain('_TOOL_CALL_SOCKET = "/tmp/tcs.sock"'); - expect(preamble).not.toMatch(/os\.environ\.get\(['"]TOOL_CALL_SOCKET['"]/); + expect(preamble).not.toMatch( + /os\.environ\.get\(['"]TOOL_CALL_SOCKET['"]/, + ); expect(preamble).not.toMatch(/os\.environ\[['"]TOOL_CALL_SOCKET['"]\]/); }); }); describe('extractPendingFromStdout — input hash metadata', () => { + test('normalizes native control payload hashes instead of trusting the sandbox', () => { + const forgedHash = hashToolInput({ resource: 'B' }); + const expectedHash = hashToolInput({ resource: 'A' }); + const pending = extractPendingFromControlPayload( + JSON.stringify({ + pending: [ + { + call_id: 'call_001', + tool_name: 'authorize', + input: { resource: 'A' }, + input_hash: forgedHash, + }, + ], + }), + ); + expect(pending?.[0]?.input_hash).toBe(expectedHash); + expect(pending?.[0]?.input_hash).not.toBe(forgedHash); + }); + test('ignores sandbox-supplied input_hash and uses the parsed input hash', () => { const executionId = 'exec_hash_guard'; const { start, end } = buildScopedSentinel(executionId); const forgedHash = hashToolInput({ resource: 'B' }); const expectedHash = hashToolInput({ resource: 'A' }); const payload = { - pending: [{ + pending: [ + { call_id: 'call_001', tool_name: 'authorize', input: { resource: 'A' }, input_hash: forgedHash, - }], + }, + ], }; const parsed = extractPendingFromStdout( diff --git a/service/src/preamble.ts b/service/src/preamble.ts index 9ff75dc7..94685c6c 100644 --- a/service/src/preamble.ts +++ b/service/src/preamble.ts @@ -2,7 +2,10 @@ import fs from 'fs'; import path from 'path'; import type * as t from './types'; import { planLimits } from './config'; -import { generateBashReplayPreamble, generateBashReplayPostamble } from './preamble-bash'; +import { + generateBashReplayPreamble, + generateBashReplayPostamble, +} from './preamble-bash'; import { PTC_HISTORY_FILENAME, PTC_HISTORY_SANDBOX_PATH, @@ -11,10 +14,16 @@ import { buildScopedSentinel, isReservedPtcFilename, } from './ptc-constants'; -import { hashToolInput, pendingInputHashesFromRawPayload } from './tool-input-signature'; +import { + hashToolInput, + pendingInputHashesFromRawPayload, +} from './tool-input-signature'; // Load async matplotlib template for programmatic tool calling -const templateCodeAsync = fs.readFileSync(path.join(__dirname, 'matplotlib-async.py'), 'utf8'); +const templateCodeAsync = fs.readFileSync( + path.join(__dirname, 'matplotlib-async.py'), + 'utf8', +); // ============================================================================= // Programmatic Tool Calling Types & Preamble Generation @@ -88,11 +97,41 @@ function normalizePythonFunctionName(name: string): string { // Python keywords to avoid const pythonKeywords = new Set([ - 'False', 'None', 'True', 'and', 'as', 'assert', 'async', 'await', - 'break', 'class', 'continue', 'def', 'del', 'elif', 'else', 'except', - 'finally', 'for', 'from', 'global', 'if', 'import', 'in', 'is', - 'lambda', 'nonlocal', 'not', 'or', 'pass', 'raise', 'return', - 'try', 'while', 'with', 'yield' + 'False', + 'None', + 'True', + 'and', + 'as', + 'assert', + 'async', + 'await', + 'break', + 'class', + 'continue', + 'def', + 'del', + 'elif', + 'else', + 'except', + 'finally', + 'for', + 'from', + 'global', + 'if', + 'import', + 'in', + 'is', + 'lambda', + 'nonlocal', + 'not', + 'or', + 'pass', + 'raise', + 'return', + 'try', + 'while', + 'with', + 'yield', ]); if (pythonKeywords.has(normalized)) { @@ -137,11 +176,14 @@ function jsonSchemaToPythonType(schema: JsonSchemaProperty): string { * Sort property names so required parameters come before optional ones. * Uses a Set for O(1) lookups instead of repeated array includes() calls. */ -function getSortedPropertyNames(propertyNames: string[], required: string[]): string[] { +function getSortedPropertyNames( + propertyNames: string[], + required: string[], +): string[] { const requiredSet = new Set(required); return [ ...propertyNames.filter(name => requiredSet.has(name)), - ...propertyNames.filter(name => !requiredSet.has(name)) + ...propertyNames.filter(name => !requiredSet.has(name)), ]; } @@ -155,7 +197,10 @@ function schemaToParams(schema?: JsonSchema): string { const required = schema.required ?? []; const requiredSet = new Set(required); - const sortedNames = getSortedPropertyNames(Object.keys(schema.properties), required); + const sortedNames = getSortedPropertyNames( + Object.keys(schema.properties), + required, + ); const params: string[] = []; @@ -192,7 +237,11 @@ function inferReturnType(description?: string): string { const desc = description.toLowerCase(); - if (desc.includes('returns list') || desc.includes('returns array') || desc.includes('list of')) { + if ( + desc.includes('returns list') || + desc.includes('returns array') || + desc.includes('list of') + ) { return 'List[Dict[str, Any]]'; } if (desc.includes('returns dict') || desc.includes('returns object')) { @@ -225,7 +274,10 @@ function generateDocstring(tool: LCTool): string { doc += '\n\n Parameters:'; const required = tool.parameters.required ?? []; const requiredSet = new Set(required); - const sortedNames = getSortedPropertyNames(Object.keys(tool.parameters.properties), required); + const sortedNames = getSortedPropertyNames( + Object.keys(tool.parameters.properties), + required, + ); for (const name of sortedNames) { const propSchema = tool.parameters.properties[name]; @@ -253,7 +305,8 @@ function generateToolStub(tool: LCTool): string { const pythonFunctionName = normalizePythonFunctionName(tool.name); // If name was changed, add a comment - const nameComment = pythonFunctionName !== tool.name + const nameComment = + pythonFunctionName !== tool.name ? ` # Original tool name: ${tool.name}\n` : ''; @@ -451,7 +504,8 @@ async def _execute_tool_internal_async(tool_name: str, tool_input: Dict[str, Any */ export function generateReplayPreamble(config: ReplayPreambleConfig): string { const { executionId, tools } = config; - const { start: scopedStart, end: scopedEnd } = buildScopedSentinel(executionId); + const { start: scopedStart, end: scopedEnd } = + buildScopedSentinel(executionId); let preamble = ` # ============================================================================ @@ -564,6 +618,63 @@ export interface ExtractPendingResult { }> | null; } +export function extractPendingFromControlPayload( + rawPayload: string, +): ExtractPendingResult['pending'] { + let parsed: { pending?: unknown } | null = null; + try { + parsed = JSON.parse(rawPayload) as { pending?: unknown }; + } catch { + return null; + } + + const pendingField = parsed?.pending; + if (!Array.isArray(pendingField)) return null; + + const rawInputHashes = pendingInputHashesFromRawPayload(rawPayload); + type PendingWithIndex = { + c: { call_id: string; tool_name: string; input: unknown }; + index: number; + }; + const isPendingWithIndex = (entry: { + c: unknown; + index: number; + }): entry is PendingWithIndex => { + const { c } = entry; + return ( + c != null && + typeof c === 'object' && + typeof (c as { call_id?: unknown }).call_id === 'string' && + typeof (c as { tool_name?: unknown }).tool_name === 'string' + ); + }; + return pendingField + .map((c, index) => ({ c, index })) + .filter(isPendingWithIndex) + .map(({ c, index }) => { + const callSite = (c as { call_site?: unknown }).call_site; + const rawInputHash = rawInputHashes[index]; + const hasObjectInput = + c.input != null && typeof c.input === 'object'; + const input = (hasObjectInput ? c.input : {}) as Record< + string, + unknown + >; + return { + call_id: c.call_id, + tool_name: c.tool_name, + input, + input_hash: + hasObjectInput && typeof rawInputHash === 'string' + ? rawInputHash + : hashToolInput(input), + ...(typeof callSite === 'string' + ? { call_site: callSite } + : {}), + }; + }); +} + /** * Locate the last line whose trimmed content exactly equals `marker`. * Using full-line anchoring prevents user-provided tool payloads that happen @@ -581,7 +692,8 @@ function findSentinelLine( for (let i = lines.length - 1; i >= searchFromLine; i--) { if (lines[i].trim() === marker) { const startOffset = lineStartOffsets[i]; - const endOffset = i + 1 < lineStartOffsets.length + const endOffset = + i + 1 < lineStartOffsets.length ? lineStartOffsets[i + 1] - 1 : startOffset + lines[i].length; return { line: i, startOffset, endOffset }; @@ -619,48 +731,8 @@ export function extractPendingFromStdout( const payloadLines = lines.slice(startLine.line + 1, endLine.line); const rawPayload = payloadLines.join('\n').trim(); - let parsed: { pending?: unknown } | null = null; - try { - parsed = JSON.parse(rawPayload) as { pending?: unknown }; - } catch { - return { stdout, pending: null }; - } - - const pendingField = parsed?.pending; - if (!Array.isArray(pendingField)) return { stdout, pending: null }; - - const rawInputHashes = pendingInputHashesFromRawPayload(rawPayload); - type PendingWithIndex = { - c: { call_id: string; tool_name: string; input: unknown }; - index: number; - }; - const isPendingWithIndex = (entry: { c: unknown; index: number }): entry is PendingWithIndex => { - const { c } = entry; - return ( - c != null && - typeof c === 'object' && - typeof (c as { call_id?: unknown }).call_id === 'string' && - typeof (c as { tool_name?: unknown }).tool_name === 'string' - ); - }; - const pending = pendingField - .map((c, index) => ({ c, index })) - .filter(isPendingWithIndex) - .map(({ c, index }) => { - const callSite = (c as { call_site?: unknown }).call_site; - const rawInputHash = rawInputHashes[index]; - const hasObjectInput = c.input != null && typeof c.input === 'object'; - const input = (hasObjectInput ? c.input : {}) as Record; - return { - call_id: c.call_id, - tool_name: c.tool_name, - input, - input_hash: hasObjectInput && typeof rawInputHash === 'string' - ? rawInputHash - : hashToolInput(input), - ...(typeof callSite === 'string' ? { call_site: callSite } : {}), - }; - }); + const pending = extractPendingFromControlPayload(rawPayload); + if (pending == null) return { stdout, pending: null }; /** Strip only the sentinel block and leave every other byte of user * stdout untouched. Both the Python and bash preambles defensively @@ -674,7 +746,10 @@ export function extractPendingFromStdout( * emission, and anything else that depends on byte-accurate stdout. */ const rawHead = stdout.slice(0, startLine.startOffset); const head = rawHead.endsWith('\n') ? rawHead.slice(0, -1) : rawHead; - const tailStart = endLine.endOffset < stdout.length ? endLine.endOffset + 1 : stdout.length; + const tailStart = + endLine.endOffset < stdout.length + ? endLine.endOffset + 1 + : stdout.length; const tail = stdout.slice(tailStart); const cleaned = head + tail; @@ -688,11 +763,14 @@ export function extractPendingFromStdout( function wrapUserCodeInAsync(userCode: string): string { const lines = userCode.split('\n'); - let wrapped = '# ============================================================================\n'; + let wrapped = + '# ============================================================================\n'; wrapped += '# USER CODE BEGINS BELOW\n'; - wrapped += '# ============================================================================\n\n'; + wrapped += + '# ============================================================================\n\n'; wrapped += 'async def __user_main__():\n'; - wrapped += ' """Auto-generated wrapper for user code to support top-level await"""\n'; + wrapped += + ' """Auto-generated wrapper for user code to support top-level await"""\n'; // Indent all user code for (const line of lines) { @@ -743,10 +821,21 @@ const PROGRAMMATIC_RUN_TIMEOUT = 300000; // 5 minutes wall time * Create a payload for programmatic tool calling execution * Combines the tool preamble with user code */ -export function createProgrammaticPayload(options: CreateProgrammaticPayloadOptions): t.PayloadBody { +export function createProgrammaticPayload( + options: CreateProgrammaticPayloadOptions, +): t.PayloadBody { const { - req, session_id, execution_id, callbackUrl, callbackToken, tools, timeout, - mode = 'blocking', history, codeOverride, filesOverride, + req, + session_id, + execution_id, + callbackUrl, + callbackToken, + tools, + timeout, + mode = 'blocking', + history, + codeOverride, + filesOverride, language = 'python', } = options; const body = req.body as t.ProgrammaticRequestBody; @@ -762,7 +851,14 @@ export function createProgrammaticPayload(options: CreateProgrammaticPayloadOpti throw new Error('bash PTC is only supported in replay mode'); } return buildBashPayload({ - req, execution_id, session_id, tools, userCode, files, history, timeout, + req, + execution_id, + session_id, + tools, + userCode, + files, + history, + timeout, }); } @@ -771,7 +867,9 @@ export function createProgrammaticPayload(options: CreateProgrammaticPayloadOpti preamble = generateReplayPreamble({ executionId: execution_id, tools }); } else { if (!callbackUrl || !callbackToken) { - throw new Error('blocking PTC mode requires callbackUrl and callbackToken'); + throw new Error( + 'blocking PTC mode requires callbackUrl and callbackToken', + ); } preamble = generatePreamble({ callbackUrl, @@ -781,15 +879,21 @@ export function createProgrammaticPayload(options: CreateProgrammaticPayloadOpti }); } - const isPyPlot = userCode.includes('import matplotlib') || userCode.includes('import seaborn'); + const isPyPlot = + userCode.includes('import matplotlib') || + userCode.includes('import seaborn'); let finalCode: string; if (isPyPlot) { - const indentedUserCode = userCode.trim().split('\n').map(line => ` ${line}`).join('\n'); + const indentedUserCode = userCode + .trim() + .split('\n') + .map(line => ` ${line}`) + .join('\n'); const wrappedUserCode = templateCodeAsync.replace( /# BEGIN USER CODE\n[\s\S]*?# END USER CODE/, - `# BEGIN USER CODE\n${indentedUserCode}\n # END USER CODE` + `# BEGIN USER CODE\n${indentedUserCode}\n # END USER CODE`, ); finalCode = preamble + '\n' + wrappedUserCode; } else { @@ -797,7 +901,9 @@ export function createProgrammaticPayload(options: CreateProgrammaticPayloadOpti finalCode = preamble + wrappedUserCode; } - const run_memory_limit = planLimits[req.planId ?? '']?.run_memory_limit ?? planLimits.default.run_memory_limit; + const run_memory_limit = + planLimits[req.planId ?? '']?.run_memory_limit ?? + planLimits.default.run_memory_limit; const run_timeout = timeout ?? PROGRAMMATIC_RUN_TIMEOUT; const payload: t.PayloadBody = { @@ -809,8 +915,8 @@ export function createProgrammaticPayload(options: CreateProgrammaticPayloadOpti files: [ { name: 'main.py', - content: finalCode - } + content: finalCode, + }, ], session_id, }; @@ -851,13 +957,27 @@ function buildBashPayload(args: { history?: Record; timeout?: number; }): t.PayloadBody { - const { req, execution_id, session_id, tools, userCode, files, history, timeout } = args; - - const preamble = generateBashReplayPreamble({ executionId: execution_id, tools }); + const { + req, + execution_id, + session_id, + tools, + userCode, + files, + history, + timeout, + } = args; + + const preamble = generateBashReplayPreamble({ + executionId: execution_id, + tools, + }); const postamble = generateBashReplayPostamble(); const finalCode = preamble + userCode + '\n' + postamble; - const run_memory_limit = planLimits[req.planId ?? '']?.run_memory_limit ?? planLimits.default.run_memory_limit; + const run_memory_limit = + planLimits[req.planId ?? '']?.run_memory_limit ?? + planLimits.default.run_memory_limit; const run_timeout = timeout ?? PROGRAMMATIC_RUN_TIMEOUT; const payload: t.PayloadBody = { @@ -865,6 +985,8 @@ function buildBashPayload(args: { run_timeout, language: 'bash', version: '5.2.0', + execution_id, + replay_tool_count: tools.length, files: [ { name: 'main.sh', diff --git a/service/src/ptc-constants.test.ts b/service/src/ptc-constants.test.ts new file mode 100644 index 00000000..5028e385 --- /dev/null +++ b/service/src/ptc-constants.test.ts @@ -0,0 +1,9 @@ +import { expect, test } from 'bun:test'; +import { isReservedPtcFilename } from './ptc-constants'; + +test('reserves replay inputs and output control channels after normalization', () => { + for (const name of ['_ptc_history.json', '_ptc_pending_result.json', '_PTC_PENDING_RESULT.JSON', 'sub/../_ptc_pending_result.json', 'sub\\_ptc_pending_result.json']) { + expect(isReservedPtcFilename(name)).toBe(true); + } + expect(isReservedPtcFilename('_ptc_data.csv')).toBe(false); +}); diff --git a/service/src/ptc-constants.ts b/service/src/ptc-constants.ts index 5009c3c3..b7b3999b 100644 --- a/service/src/ptc-constants.ts +++ b/service/src/ptc-constants.ts @@ -15,8 +15,9 @@ export const PTC_HISTORY_SANDBOX_PATH = `/mnt/data/${PTC_HISTORY_FILENAME}`; * Returns `true` for any filename the submission layer must refuse. * * Two things make a name "reserved": - * 1. Its post-normalization basename is `_ptc_history.json` — the single - * runtime fixture the replay preamble injects into the submission dir. + * 1. Its post-normalization basename is `_ptc_history.json` or + * `_ptc_pending_result.json`, compared case-insensitively for macOS. + * These are the replay input and output control channels. * Any user-supplied file with that exact basename would shadow our * injected history and silently corrupt replay correctness, so we * reject it on the request path. The bash preamble's `_ptc_pending.*` @@ -61,7 +62,7 @@ export function isReservedPtcFilename(name: string): boolean { } if (escapes) return true; const basename = segments.length > 0 ? segments[segments.length - 1] : ''; - return basename === PTC_HISTORY_FILENAME; + return [PTC_HISTORY_FILENAME, '_ptc_pending_result.json'].includes(basename.toLowerCase()); } /** diff --git a/service/src/sandbox-backend/remote-bridge.test.ts b/service/src/sandbox-backend/remote-bridge.test.ts index a65bc89e..697271ee 100644 --- a/service/src/sandbox-backend/remote-bridge.test.ts +++ b/service/src/sandbox-backend/remote-bridge.test.ts @@ -62,6 +62,34 @@ describe('RemoteBridgeSandboxBackend', () => { }); }); + test('preserves an authenticated selected workspace on remote dispatch', async () => { + let dispatched: Parameters[0] | undefined; + const store = { + dispatch: async ( + args: Parameters[0], + ): ReturnType => { + dispatched = args; + return { + protocolVersion: 1 as const, + generation: 1, + leaseToken: 'a'.repeat(32), + incarnationId: 'incarnation-00000001', + status: 'fulfilled' as const, + result: { session_id: 'session-1', language: 'bash', version: '5.2', files: [] }, + }; + }, + } satisfies Pick; + const backend = new RemoteBridgeSandboxBackend(store, 'default-vm'); + + await backend.execute(request(), { ...context(), workspaceId: 'project-a' }); + + expect(dispatched).toMatchObject({ + workerId: 'user-vm', + workspaceId: 'project-a', + requireTenantBinding: true, + }); + }); + test('maps tenant authorization rejection to a bridge backend error', async () => { const store = { dispatch: async (): ReturnType => { diff --git a/service/src/sandbox-backend/remote-bridge.ts b/service/src/sandbox-backend/remote-bridge.ts index 0bee0038..6e06eda8 100644 --- a/service/src/sandbox-backend/remote-bridge.ts +++ b/service/src/sandbox-backend/remote-bridge.ts @@ -60,6 +60,7 @@ export class RemoteBridgeSandboxBackend implements SandboxBackend { (this.dynamicWorkers || ctx.bridgeWorkerId !== this.workerId), body: req.body, headers: req.headers, + ...(ctx.workspaceId != null ? { workspaceId: ctx.workspaceId } : {}), runtimeSessionId: ctx.runtimeSessionId, deadlineAtMs: ctx.deadlineAtMs ?? Date.now() + env.JOB_TIMEOUT, signal: ctx.signal, diff --git a/service/src/sandbox-backend/types.ts b/service/src/sandbox-backend/types.ts index fbaa2d20..e21982d9 100644 --- a/service/src/sandbox-backend/types.ts +++ b/service/src/sandbox-backend/types.ts @@ -39,6 +39,8 @@ export interface SandboxExecuteContext { canonicalUserId?: string; /** Trusted API-selected outbound worker. Presence requires a tenant-bound credential. */ bridgeWorkerId?: string; + /** Trusted selected workspace for native replay-mode PTC. */ + workspaceId?: string; /** Stable identifier for this queued iteration, used to derive an idempotent * stateless launch token. PTC replay reuses one executionId across every * iteration, so the executionId alone cannot separate them; the request body @@ -63,6 +65,7 @@ export type SandboxRawResponse = t.ExecuteResponse & { session_id: string; files?: t.FileRefs; run?: t.ExecuteResponse['run']; + pending_tool_calls_payload?: string; }; export interface SandboxBackend { diff --git a/service/src/sandbox-dispatch.test.ts b/service/src/sandbox-dispatch.test.ts index 435be9aa..5b1ef749 100644 --- a/service/src/sandbox-dispatch.test.ts +++ b/service/src/sandbox-dispatch.test.ts @@ -12,8 +12,10 @@ import { } from './execution-manifest'; const SECRET = 'test-secret'; -const PRIVATE_KEY = 'MC4CAQAwBQYDK2VwBCIEIBoxzSJjQ5jTVyuohHtlD+uDGqv/tZ6hQS2CmxuOg2Wn'; -const PUBLIC_KEY = 'MCowBQYDK2VwAyEAeY3PRoTS3adfU6E3gQUB5hSZdrdMSw6OrKkH4UhYh0U='; +const PRIVATE_KEY = + 'MC4CAQAwBQYDK2VwBCIEIBoxzSJjQ5jTVyuohHtlD+uDGqv/tZ6hQS2CmxuOg2Wn'; +const PUBLIC_KEY = + 'MCowBQYDK2VwAyEAeY3PRoTS3adfU6E3gQUB5hSZdrdMSw6OrKkH4UhYh0U='; function payload(overrides: Partial = {}): t.PayloadBody { return { @@ -25,14 +27,22 @@ function payload(overrides: Partial = {}): t.PayloadBody { }; } -function claims(overrides: Partial = {}): ExecutionManifestClaims { +function claims( + overrides: Partial = {}, +): ExecutionManifestClaims { return { v: EXECUTION_MANIFEST_VERSION, exec_id: 'exec_123', tenant_id: 'tenant_abc', user_id: 'user_123', session_key: 'tenant:tenant_abc:user:user_123', - input_files: [{ id: 'file_123', session_id: 'sess_input', name: 'inputs/data.csv' }], + input_files: [ + { + id: 'file_123', + session_id: 'sess_input', + name: 'inputs/data.csv', + }, + ], read_sessions: ['sess_input'], output_session_id: 'sess_output', max_upload_bytes: 1024, @@ -46,6 +56,20 @@ function claims(overrides: Partial = {}): ExecutionMani } describe('sandbox execute request dispatch', () => { + test('budgets every input and output batch before signing the request', () => { + const request = buildSandboxExecuteRequest({ + payload: payload({ files: Array.from({ length: 9 }, (_, index) => ({ name: `${index}.txt`, id: `file_${index}`, storage_session_id: 'input' })) }), + programmaticTransferReserveMs: 60_000, + executionManifestClaims: claims({ max_output_files: 10 }), + executionManifestSecret: SECRET, + executionManifestTtlSeconds: 300, + nowSeconds: 1_000, + }); + // Three download batches plus three upload batches share one reserve. + expect(request.body.transfer_timeout_ms).toBe(10_000); + const verified = verifyExecutionManifest(request.body.execution_manifest!, SECRET, { nowSeconds: 1_000 }); + expect(verified.execute_body_sha256).toBe(executionManifestBodySha256(request.body)); + }); test('keeps large egress grants out of HTTP headers', () => { const largeGrant = `ceg1.${'a'.repeat(24_000)}`; const request = buildSandboxExecuteRequest({ @@ -64,18 +88,27 @@ describe('sandbox execute request dispatch', () => { const request = buildSandboxExecuteRequest({ payload: payload(), executionManifestClaims: claims(), + maxOutputFileBytes: 1_000, executionManifestSecret: SECRET, executionManifestTtlSeconds: 300, nowSeconds: 1_000, }); expect(request.headers[EXECUTION_MANIFEST_HEADER]).toBeUndefined(); + expect(request.body.max_output_files).toBe(10); + expect(request.body.max_output_file_bytes).toBe(1_000); expect(request.body.execution_manifest).toEqual(expect.any(String)); - expect(verifyExecutionManifest(request.body.execution_manifest!, SECRET, { nowSeconds: 1_100 })).toEqual(claims({ + expect( + verifyExecutionManifest(request.body.execution_manifest!, SECRET, { + nowSeconds: 1_100, + }), + ).toEqual( + claims({ execute_body_sha256: executionManifestBodySha256(request.body), iat: 1_000, exp: 1_300, - })); + }), + ); }); test('signs execution manifests with a private key when configured', () => { @@ -88,11 +121,19 @@ describe('sandbox execute request dispatch', () => { nowSeconds: 1_000, }); - expect(verifyExecutionManifestWithPublicKey(request.body.execution_manifest!, PUBLIC_KEY, { nowSeconds: 1_100 })).toEqual(claims({ + expect( + verifyExecutionManifestWithPublicKey( + request.body.execution_manifest!, + PUBLIC_KEY, + { nowSeconds: 1_100 }, + ), + ).toEqual( + claims({ execute_body_sha256: executionManifestBodySha256(request.body), iat: 1_000, exp: 1_300, - })); + }), + ); }); test('binds body-carried egress grants into signed execution manifests', () => { @@ -106,10 +147,16 @@ describe('sandbox execute request dispatch', () => { }); expect(request.body.egress_grant).toBe('ceg1.sealed-grant'); - expect(verifyExecutionManifest(request.body.execution_manifest!, SECRET, { nowSeconds: 1_100 })).toEqual(claims({ + expect( + verifyExecutionManifest(request.body.execution_manifest!, SECRET, { + nowSeconds: 1_100, + }), + ).toEqual( + claims({ execute_body_sha256: executionManifestBodySha256(request.body), iat: 1_000, exp: 1_300, - })); + }), + ); }); }); diff --git a/service/src/sandbox-dispatch.ts b/service/src/sandbox-dispatch.ts index e3066905..340834b5 100644 --- a/service/src/sandbox-dispatch.ts +++ b/service/src/sandbox-dispatch.ts @@ -1,5 +1,14 @@ import type * as t from './types'; -import { executionManifestBodySha256, signExecutionManifestWithKey, type ExecutionManifestClaims } from './execution-manifest'; +import { + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILES, + BRIDGE_WORKSPACE_PROGRAMMATIC_TRANSFER_CONCURRENCY, + BRIDGE_WORKSPACE_PROGRAMMATIC_TRANSFER_TIMEOUT_MS, +} from '../../packages/code/src/protocol'; +import { + executionManifestBodySha256, + signExecutionManifestWithKey, + type ExecutionManifestClaims, +} from './execution-manifest'; interface BuildSandboxExecuteRequestArgs { payload: t.PayloadBody; @@ -9,6 +18,8 @@ interface BuildSandboxExecuteRequestArgs { executionManifestSecret: string; executionManifestTtlSeconds: number; nowSeconds?: number; + maxOutputFileBytes?: number; + programmaticTransferReserveMs?: number; } interface SandboxExecuteRequest { @@ -21,15 +32,31 @@ interface SandboxExecuteRequest { * ride in the JSON body instead of HTTP headers. Otherwise skill-heavy jobs can * fail with 431 before sandbox-runner reaches capability validation. */ -export function buildSandboxExecuteRequest(args: BuildSandboxExecuteRequestArgs): SandboxExecuteRequest { +export function buildSandboxExecuteRequest( + args: BuildSandboxExecuteRequestArgs, +): SandboxExecuteRequest { const body: t.PayloadBody = { ...args.payload }; - const headers: Record = { 'Content-Type': 'application/json' }; + const headers: Record = { + 'Content-Type': 'application/json', + }; if (args.egressGrantToken) { body.egress_grant = args.egressGrantToken; } + if (args.maxOutputFileBytes != null) { + body.max_output_file_bytes = args.maxOutputFileBytes; + } + if (args.programmaticTransferReserveMs != null) { + const batches = Math.max(1, + Math.ceil(body.files.filter(file => 'id' in file).length / BRIDGE_WORKSPACE_PROGRAMMATIC_TRANSFER_CONCURRENCY) + + Math.ceil((args.executionManifestClaims?.max_output_files ?? BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILES) / BRIDGE_WORKSPACE_PROGRAMMATIC_TRANSFER_CONCURRENCY), + ); + body.transfer_timeout_ms = Math.min(BRIDGE_WORKSPACE_PROGRAMMATIC_TRANSFER_TIMEOUT_MS, + Math.max(1, Math.floor(args.programmaticTransferReserveMs / batches))); + } if (args.executionManifestClaims) { + body.max_output_files = args.executionManifestClaims.max_output_files; const nowSeconds = args.nowSeconds ?? Math.floor(Date.now() / 1000); body.execution_manifest = signExecutionManifestWithKey( { diff --git a/service/src/sandbox-egress.ts b/service/src/sandbox-egress.ts index f6549c09..e47806a2 100644 --- a/service/src/sandbox-egress.ts +++ b/service/src/sandbox-egress.ts @@ -3,6 +3,7 @@ import { env } from './config'; import { createGatewayPtcCallbackToken } from './egress-gateway-client'; import type { ExecutionManifestClaims } from './execution-manifest'; import type * as t from './types'; +import { programmaticTransferReserveMs } from '../../packages/code/src/protocol'; export type SandboxJobSecurity = { payload: t.PayloadBody; @@ -65,6 +66,9 @@ export function timeoutMsToGrantSeconds(timeoutMs: number): number { } const DEFAULT_PROGRAMMATIC_TIMEOUT_MS = 300000; +const SELECTED_WORKSPACE_REPLAY_PASSES = 2; +const SELECTED_WORKSPACE_SETTLEMENT_RESERVE_MS = 5_000; +const SELECTED_WORKSPACE_MAX_QUEUE_RESERVE_MS = 30_000; export function normalizeProgrammaticTimeoutMs( rawTimeout: unknown, @@ -80,6 +84,31 @@ export function normalizeProgrammaticTimeoutMs( return Math.min(Math.ceil(rawTimeout), maxTimeout); } +/** + * Selected-workspace Bash replay may run one read-only probe and one commit + * pass in its final iteration. Bound each pass so both plus settlement reserve + * fit inside the worker-owned JOB_TIMEOUT instead of advertising a duration + * the assignment cannot complete. + */ +export function normalizeSelectedWorkspaceProgrammaticTimeoutMs( + rawTimeout: unknown, + jobTimeoutMs = env.JOB_TIMEOUT, +): number { + const totalBudget = Math.max(1, Math.floor(jobTimeoutMs)); + const queueReserve = Math.min( + SELECTED_WORKSPACE_MAX_QUEUE_RESERVE_MS, + Math.floor(totalBudget / 5), + ); + const executionBudget = Math.max( + 1, + totalBudget - queueReserve - SELECTED_WORKSPACE_SETTLEMENT_RESERVE_MS - programmaticTransferReserveMs(totalBudget), + ); + return normalizeProgrammaticTimeoutMs( + rawTimeout, + Math.max(1, Math.floor(executionBudget / SELECTED_WORKSPACE_REPLAY_PASSES)), + ); +} + export async function sealPtcCallbackTokenForGateway(args: { executionId: string; sessionId: string; diff --git a/service/src/service/programmatic-router.ts b/service/src/service/programmatic-router.ts index 9896518c..4de4a4d8 100644 --- a/service/src/service/programmatic-router.ts +++ b/service/src/service/programmatic-router.ts @@ -11,12 +11,19 @@ import { connection, getExecutionQueueBinding, } from '../queue'; -import { createProgrammaticPayload, extractPendingFromStdout } from '../preamble'; +import { + createProgrammaticPayload, + extractPendingFromControlPayload, + extractPendingFromStdout, +} from '../preamble'; import { findBashToolNameCollision } from '../preamble-bash'; import type { LCTool } from '../preamble'; import { isReservedPtcFilename } from '../ptc-constants'; import { internalServiceHeaders } from '../internal-service-auth'; -import { resolveOutputBucketSessionKey, SessionKeyResolutionError } from '../session-key'; +import { + resolveOutputBucketSessionKey, + SessionKeyResolutionError, +} from '../session-key'; import { getCredentialId, getPrincipalOrReject } from '../auth/principal'; import { getExecutionIdentity } from '../execution-identity'; import { PROGRAMMATIC_RUNTIME_SESSION_EXEMPTION } from '../runtime-session/job-policy'; @@ -34,15 +41,25 @@ import { publicExecutionFailure } from '../utils'; import { normalizeEgressGatewayUrl, normalizeProgrammaticTimeoutMs, + normalizeSelectedWorkspaceProgrammaticTimeoutMs, prepareSandboxJobSecurity, sealPtcCallbackTokenForGateway, timeoutMsToGrantSeconds, } from '../sandbox-egress'; import { findUnregisteredToolCall } from '../tool-scope'; import { summarizeRequestedFiles } from '../execution-log'; -import { pollBlockingExecution, type BlockingPendingState } from './blocking-poll'; -import { clearSessionOwnership, recordSessionOwnership } from '../session-ownership'; -import { FileRefAuthorizationError, authorizeRequestedFiles } from './file-authorization'; +import { + pollBlockingExecution, + type BlockingPendingState, +} from './blocking-poll'; +import { + clearSessionOwnership, + recordSessionOwnership, +} from '../session-ownership'; +import { + FileRefAuthorizationError, + authorizeRequestedFiles, +} from './file-authorization'; import { buildReplayExecutionState, resolveReplayStateSandboxBackend, @@ -50,8 +67,10 @@ import { import { BridgeWorkerSelectionError, CODEAPI_BRIDGE_WORKER_HEADER, + CODEAPI_BRIDGE_WORKSPACE_HEADER, resolveBridgeWorkerSelection, } from '../bridge/selection'; +import { isValidBridgeWorkerId, BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_INPUT_FILES } from '../../../packages/code/src/protocol'; import logger from '../logger'; import { type ExecutionState, @@ -145,26 +164,41 @@ async function retryToolCallServerRequest( ): Promise { let lastError: Error | undefined; - for (let attempt = 1; attempt <= TOOL_CALL_SERVER_RETRY_ATTEMPTS; attempt++) { + for ( + let attempt = 1; + attempt <= TOOL_CALL_SERVER_RETRY_ATTEMPTS; + attempt++ + ) { try { return await requestFn(); } catch (error) { lastError = error as Error; if (axios.isAxiosError(error)) { - if (error.response && error.response.status >= 400 && error.response.status < 500) { + if ( + error.response && + error.response.status >= 400 && + error.response.status < 500 + ) { throw error; } } if (attempt < TOOL_CALL_SERVER_RETRY_ATTEMPTS) { - logger.warn(`${context} failed (attempt ${attempt}/${TOOL_CALL_SERVER_RETRY_ATTEMPTS}), retrying...`, { + logger.warn( + `${context} failed (attempt ${attempt}/${TOOL_CALL_SERVER_RETRY_ATTEMPTS}), retrying...`, + { error: lastError.message, - }); - await new Promise(resolve => setTimeout(resolve, TOOL_CALL_SERVER_RETRY_DELAY * attempt)); + }, + ); + await new Promise(resolve => + setTimeout(resolve, TOOL_CALL_SERVER_RETRY_DELAY * attempt), + ); } } } - logger.error(`${context} failed after ${TOOL_CALL_SERVER_RETRY_ATTEMPTS} attempts`); + logger.error( + `${context} failed after ${TOOL_CALL_SERVER_RETRY_ATTEMPTS} attempts`, + ); throw lastError; } @@ -179,7 +213,9 @@ setInterval(() => { }, STALE_CLEANUP_INTERVAL_MS); function generateContinuationToken(execution_id: string): string { - return Buffer.from(JSON.stringify({ execution_id, ts: Date.now() })).toString('base64'); + return Buffer.from( + JSON.stringify({ execution_id, ts: Date.now() }), + ).toString('base64'); } /** Map a replay-continuation HTTP status to its operational outcome @@ -200,14 +236,21 @@ function classifyContinuationOutcome(statusCode: number): string { * timestamp is older than the execution-state TTL — without this, the * `ts` field was dead data and a client could replay an ancient token * against a freshly-reused-execution-id window. */ -function decodeContinuationToken(token: string): { execution_id: string } | null { +function decodeContinuationToken( + token: string, +): { execution_id: string } | null { try { - const parsed: unknown = JSON.parse(Buffer.from(token, 'base64').toString('utf-8')); + const parsed: unknown = JSON.parse( + Buffer.from(token, 'base64').toString('utf-8'), + ); if (parsed === null || typeof parsed !== 'object') { return null; } const candidate = parsed as { execution_id?: unknown; ts?: unknown }; - if (typeof candidate.execution_id !== 'string' || candidate.execution_id.length === 0) { + if ( + typeof candidate.execution_id !== 'string' || + candidate.execution_id.length === 0 + ) { return null; } if (typeof candidate.ts === 'number' && Number.isFinite(candidate.ts)) { @@ -226,13 +269,17 @@ function decodeContinuationToken(token: string): { execution_id: string } | null // Blocking mode (legacy path) // --------------------------------------------------------------------------- -function waitForExecutionState(execution_id: string, timeout: number): ReturnType { +function waitForExecutionState( + execution_id: string, + timeout: number, +): ReturnType { return pollBlockingExecution(execution_id, timeout, { getExecutionState, getBlockingResult, - getPending: async (id) => { + getPending: async id => { const response = await retryToolCallServerRequest( - () => axios.get( + () => + axios.get( `${env.TOOL_CALL_SERVER_URL}/sessions/${id}/pending`, { headers: internalServiceHeaders() }, ), @@ -240,7 +287,8 @@ function waitForExecutionState(execution_id: string, timeout: number): ReturnTyp ); return response.data; }, - isNotFound: (error) => axios.isAxiosError(error) && error.response?.status === 404, + isNotFound: error => + axios.isAxiosError(error) && error.response?.status === 404, sleep: () => new Promise(resolve => setTimeout(resolve, POLL_INTERVAL)), now: Date.now, }); @@ -295,7 +343,8 @@ async function runReplayIteration( }); if (DEBUG_MODE) { - const firstFile = rawPayload.files[0] as { content?: string } | undefined; + const firstFile = rawPayload.files[0] as + { content?: string } | undefined; logger.debug('Replay enqueue details', { execution_id: state.execution_id, historySize: Object.keys(history).length, @@ -320,7 +369,9 @@ async function runReplayIteration( state.executionProfile ?? env.EXECUTION_PROFILE, state.executionProfileSource ?? env.EXECUTION_PROFILE_SOURCE, ); - const job = await queue.add(Jobs.execute, { + const job = await queue.add( + Jobs.execute, + { code: state.userCode ?? '', userId, payload: sandboxSecurity.payload, @@ -332,17 +383,24 @@ async function runReplayIteration( canonicalUserId: state.canonicalUserId, executionProfile: state.executionProfile ?? env.EXECUTION_PROFILE, sandboxBackend: replayBackend, - ...(state.bridgeWorkerId != null ? { bridgeWorkerId: state.bridgeWorkerId } : {}), + ...(state.bridgeWorkerId != null + ? { bridgeWorkerId: state.bridgeWorkerId } + : {}), + ...(state.workspaceId != null + ? { workspaceId: state.workspaceId } + : {}), runtimeSessionMode: 'stateless', runtimeSessionExemption: PROGRAMMATIC_RUNTIME_SESSION_EXEMPTION, executionManifestClaims: sandboxSecurity.executionManifestClaims, egressGrantClaims: sandboxSecurity.egressGrantClaims, egressGrantToken: sandboxSecurity.egressGrantToken, - }, { + }, + { removeOnComplete: { age: 60, count: 1 }, removeOnFail: { age: 180, count: 1 }, attempts: 1, - }); + }, + ); jobsSubmitted.inc({ language }); return job.waitUntilFinished(events, JOB_COMPLETION_WAIT_TIMEOUT_MS); @@ -365,18 +423,24 @@ async function handleReplayInitial( apiKeyId: string; userId: string; bridgeWorkerId?: string; + workspaceId?: string; }, ): Promise { - const { apiKeyId, userId, bridgeWorkerId } = params; - const { - code, - tools, - user_id, - files, - } = req.body as t.ProgrammaticRequestBody; + const { apiKeyId, userId, bridgeWorkerId, workspaceId } = params; + const { code, tools, user_id, files } = + req.body as t.ProgrammaticRequestBody; let timeout: number; try { - timeout = normalizeProgrammaticTimeoutMs((req.body as t.ProgrammaticRequestBody).timeout); + if (workspaceId != null && Array.isArray(files) && files.length > BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_INPUT_FILES) { + throw new Error(`Selected-workspace execution allows at most ${BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_INPUT_FILES} input files; main and replay history occupy two reserved slots`); + } + timeout = workspaceId != null + ? normalizeSelectedWorkspaceProgrammaticTimeoutMs( + (req.body as t.ProgrammaticRequestBody).timeout, + ) + : normalizeProgrammaticTimeoutMs( + (req.body as t.ProgrammaticRequestBody).timeout, + ); } catch (error) { res.status(400).json({ error: (error as Error).message }); return; @@ -400,14 +464,23 @@ async function handleReplayInitial( }); return; } - const language: 'python' | 'bash' = requestedLanguage === 'bash' ? 'bash' : 'python'; + const language: 'python' | 'bash' = + requestedLanguage === 'bash' ? 'bash' : 'python'; + if (workspaceId != null && language !== 'bash') { + res.status(400).json({ + error: 'Selected-workspace programmatic execution supports bash only', + }); + return; + } if (!code) { res.status(400).json({ error: 'Missing required field: code' }); return; } - if (!tools || !Array.isArray(tools) || tools.length === 0) { - res.status(400).json({ error: 'Missing required field: tools (must be a non-empty array)' }); + if (!Array.isArray(tools) || (tools.length === 0 && workspaceId == null)) { + res.status(400).json({ + error: 'Missing required field: tools (must be non-empty unless a selected workspace executes bash)', + }); return; } if (tools.length > MAX_TOOLS_PER_REQUEST) { @@ -442,7 +515,8 @@ async function handleReplayInitial( files, store: connection, }); - (req.body as t.ProgrammaticRequestBody).files = authorizedFiles.length > 0 ? authorizedFiles : undefined; + (req.body as t.ProgrammaticRequestBody).files = + authorizedFiles.length > 0 ? authorizedFiles : undefined; } catch (error) { if (sendFileRefAuthorizationError(error, res, req)) return; logger.error('Error authorizing replay file refs:', error); @@ -456,7 +530,14 @@ async function handleReplayInitial( try { sessionKey = resolveOutputBucketSessionKey(req); } catch (error) { - if (sendSessionKeyResolutionError(error, res, req, 'programmatic /exec: resolveOutputBucketSessionKey')) { + if ( + sendSessionKeyResolutionError( + error, + res, + req, + 'programmatic /exec: resolveOutputBucketSessionKey', + ) + ) { return; } throw error; @@ -466,9 +547,9 @@ async function handleReplayInitial( const execution_id = nanoid(); const authContext = req.codeApiAuthContext; const identity = getExecutionIdentity(req, userId); - const isPyPlot = language === 'python' && ( - code.includes('import matplotlib') || code.includes('import seaborn') - ); + const isPyPlot = + language === 'python' && + (code.includes('import matplotlib') || code.includes('import seaborn')); await recordSessionOwnership(connection, session_id, sessionKey); @@ -487,6 +568,7 @@ async function handleReplayInitial( timeout, language, bridgeWorkerId, + workspaceId, executionProfile: env.EXECUTION_PROFILE, executionProfileSource: env.EXECUTION_PROFILE_SOURCE, sandboxBackend: resolveReplayStateSandboxBackend({ @@ -506,13 +588,16 @@ async function handleReplayInitial( await setExecutionState(state); } catch (err) { if (err instanceof ExecutionStateTooLargeError) { - logger.warn('Rejecting replay request: ExecutionState exceeds Redis cap', { + logger.warn( + 'Rejecting replay request: ExecutionState exceeds Redis cap', + { execution_id, userId, apiKeyId, bytes: err.bytes, cap: err.cap, - }); + }, + ); await clearSessionOwnership(connection, session_id).catch(() => {}); ptcReplayStateOversize.inc(); res.status(413).json({ @@ -568,9 +653,15 @@ async function handleReplayContinuation( * adds outcome plumbing through `runAndRespond`. */ const startMs = performance.now(); res.once('finish', () => { - const labels = { mode: 'replay' as const, outcome: classifyContinuationOutcome(res.statusCode) }; + const labels = { + mode: 'replay' as const, + outcome: classifyContinuationOutcome(res.statusCode), + }; ptcReplayContinuations.inc(labels); - ptcReplayContinuationDuration.observe(labels, (performance.now() - startMs) / 1000); + ptcReplayContinuationDuration.observe( + labels, + (performance.now() - startMs) / 1000, + ); }); /** Reject oversized batches before we spend any CPU on per-entry @@ -621,9 +712,14 @@ async function handleReplayContinuation( call_site: emitted.call_site, }; }); - const deltaOrError = await computeToolHistoryDelta(state.execution_id, enrichedResults); + const deltaOrError = await computeToolHistoryDelta( + state.execution_id, + enrichedResults, + ); if ('error' in deltaOrError) { - res.status(deltaOrError.status ?? 400).json({ error: deltaOrError.error }); + res.status(deltaOrError.status ?? 400).json({ + error: deltaOrError.error, + }); return; } const delta = deltaOrError; @@ -642,7 +738,9 @@ async function handleReplayContinuation( }); if (!pre.ok) { if (pre.status === 403) { - logger.warn('Unauthorized replay continuation request rejected', { + logger.warn( + 'Unauthorized replay continuation request rejected', + { execution_id: state.execution_id, requestUserId: userId, requestApiKeyId: apiKeyId, @@ -650,7 +748,8 @@ async function handleReplayContinuation( executionUserId: state.userId, executionApiKeyId: state.apiKeyId, executionTenantId: state.tenantId, - }); + }, + ); } if (pre.cleanupOnReject === true) { await cleanupExecution(state.execution_id, 'replay'); @@ -672,7 +771,10 @@ async function handleReplayContinuation( * Redis MULTI/EXEC so counters and the hash can't drift out of sync * on a partial failure. */ state.callCount = (state.callCount ?? 0) + delta.newCallIds.length; - state.historyBytes = Math.max(0, (state.historyBytes ?? 0) + delta.bytesDelta); + state.historyBytes = Math.max( + 0, + (state.historyBytes ?? 0) + delta.bytesDelta, + ); state.lastActivity = Date.now(); try { await commitToolHistoryAndState(state, delta); @@ -687,14 +789,19 @@ async function handleReplayContinuation( * forward is a fresh execution with smaller inputs. Reap the * old execution to free the lock and Redis keys, then return * an actionable 413 instead of a generic 500. */ - logger.warn('Replay continuation rejected: ExecutionState exceeds Redis cap', { + logger.warn( + 'Replay continuation rejected: ExecutionState exceeds Redis cap', + { execution_id: state.execution_id, bytes: err.bytes, cap: err.cap, callCount: state.callCount, historyBytes: state.historyBytes, - }); - await cleanupExecution(state.execution_id, 'replay').catch(() => {}); + }, + ); + await cleanupExecution(state.execution_id, 'replay').catch( + () => {}, + ); ptcReplayStateOversize.inc(); res.status(413).json({ status: 'error', @@ -715,10 +822,13 @@ async function handleReplayContinuation( * the throw bubble to the top-level catch and become an opaque * 500 — clients (and load balancers) treat 5xx classes very * differently for retry policy. */ - logger.error('Failed to commit replay continuation; returning retryable 503', { + logger.error( + 'Failed to commit replay continuation; returning retryable 503', + { execution_id: state.execution_id, err: (err as Error).message, - }); + }, + ); res.status(503).json({ status: 'error', error: 'Failed to persist replay continuation; please retry the same request', @@ -764,11 +874,15 @@ async function runAndRespond( try { result = await runReplayIteration(req, state, apiKeyId, userId); } catch (err) { - logger.error('Replay iteration failed', { execution_id: state.execution_id, err }); + logger.error('Replay iteration failed', { + execution_id: state.execution_id, + err, + }); await cleanupExecution(state.execution_id, 'replay'); if (!isDisconnected()) { const publicFailure = publicExecutionFailure(err); - const message = publicFailure?.body.message ?? (err as Error).message; + const message = + publicFailure?.body.message ?? (err as Error).message; res.status(200).json({ status: 'error', error: message !== '' ? message : 'Sandbox execution failed', @@ -786,10 +900,19 @@ async function runAndRespond( return; } - const { stdout: cleanStdout, pending } = extractPendingFromStdout( + const extracted = extractPendingFromStdout( result.stdout, state.execution_id, ); + const cleanStdout = extracted.stdout; + const controlPayload = result.pending_tool_calls_payload; + const hasControlPayload = typeof controlPayload === 'string'; + const controlPending = hasControlPayload + ? extractPendingFromControlPayload(controlPayload) + : null; + const pending = hasControlPayload + ? (controlPending ?? []) + : extracted.pending; if (pending != null) { if (pending.length === 0) { @@ -806,7 +929,10 @@ async function runAndRespond( }); return; } - const unregisteredToolCall = findUnregisteredToolCall(pending, state.tools); + const unregisteredToolCall = findUnregisteredToolCall( + pending, + state.tools, + ); if (unregisteredToolCall != null) { logger.warn('Sandbox requested unregistered replay tool call', { execution_id: state.execution_id, @@ -866,11 +992,16 @@ async function runAndRespond( await setExecutionState(state); await refreshExecutionTtl(state.execution_id); } catch (err) { - logger.error('Failed to persist execution state before continuation; aborting', { + logger.error( + 'Failed to persist execution state before continuation; aborting', + { execution_id: state.execution_id, err: (err as Error).message, - }); - await cleanupExecution(state.execution_id, 'replay').catch(() => {}); + }, + ); + await cleanupExecution(state.execution_id, 'replay').catch( + () => {}, + ); if (!isDisconnected()) { if (err instanceof ExecutionStateTooLargeError) { /** A continuation that pushes `emittedCallIds` past the @@ -887,8 +1018,7 @@ async function runAndRespond( } else { res.status(503).json({ status: 'error', - error: - 'Failed to persist replay state; please retry the request from scratch', + error: 'Failed to persist replay state; please retry the request from scratch', session_id: state.session_id, }); } @@ -912,7 +1042,8 @@ async function runAndRespond( if (!isSandboxRunSuccess(result)) { await cleanupExecution(state.execution_id, 'replay'); - const errorMessage = result.message != null && result.message !== '' + const errorMessage = + result.message != null && result.message !== '' ? result.message : `Sandbox exited with code ${result.code ?? 'unknown'}`; res.status(200).json({ @@ -941,7 +1072,10 @@ async function runAndRespond( // Request entrypoint // --------------------------------------------------------------------------- -router.post('/exec/programmatic', executionLimiter, async (req: t.AuthenticatedRequest, res) => { +router.post( + '/exec/programmatic', + executionLimiter, + async (req: t.AuthenticatedRequest, res) => { const principal = getPrincipalOrReject(req, res); if (!principal) return; const apiKeyId = getCredentialId(req); @@ -954,13 +1088,12 @@ router.post('/exec/programmatic', executionLimiter, async (req: t.AuthenticatedR return res.status(503).json({ error: 'Service is starting up' }); } - const { - continuation_token, - tool_results, - } = req.body as t.ProgrammaticRequestBody; + const { continuation_token, tool_results } = + req.body as t.ProgrammaticRequestBody; const rawBody = req.body as Record; const requestedLanguage: unknown = rawBody.language ?? rawBody.lang; let bridgeWorkerId: string | undefined; + let workspaceId: string | undefined; if (continuation_token == null || continuation_token === '') { try { const bridgeSelection = resolveBridgeWorkerSelection({ @@ -970,12 +1103,35 @@ router.post('/exec/programmatic', executionLimiter, async (req: t.AuthenticatedR requestedWorkerId: req.header(CODEAPI_BRIDGE_WORKER_HEADER), trustedWorkerId: principal.codeWorkerId, }); - bridgeWorkerId = bridgeSelection?.explicit === true - ? bridgeSelection.workerId - : undefined; + bridgeWorkerId = + bridgeSelection?.explicit === true || + (bridgeSelection != null && !env.BRIDGE_DYNAMIC_WORKERS) + ? bridgeSelection.workerId + : undefined; + const requestedWorkspaceId = req + .header(CODEAPI_BRIDGE_WORKSPACE_HEADER) + ?.trim(); + if ( + requestedWorkspaceId != null && + requestedWorkspaceId !== '' + ) { + if (bridgeWorkerId == null) { + return res.status(400).json({ + error: 'Workspace selection requires an authenticated bridge worker', + }); + } + if (!isValidBridgeWorkerId(requestedWorkspaceId)) { + return res + .status(400) + .json({ error: 'Invalid code workspace ID' }); + } + workspaceId = requestedWorkspaceId; + } } catch (error) { if (error instanceof BridgeWorkerSelectionError) { - return res.status(error.status).json({ error: error.message }); + return res + .status(error.status) + .json({ error: error.message }); } throw error; } @@ -1013,7 +1169,9 @@ router.post('/exec/programmatic', executionLimiter, async (req: t.AuthenticatedR } const decoded = decodeContinuationToken(continuation_token); if (!decoded) { - return res.status(400).json({ error: 'Invalid continuation token' }); + return res + .status(400) + .json({ error: 'Invalid continuation token' }); } const existing = await getExecutionState(decoded.execution_id); if (existing?.mode === 'replay') { @@ -1036,9 +1194,23 @@ router.post('/exec/programmatic', executionLimiter, async (req: t.AuthenticatedR }); } if (env.PTC_MODE === 'replay') { - return await handleReplayInitial(req, res, { apiKeyId, userId, bridgeWorkerId }); + return await handleReplayInitial(req, res, { + apiKeyId, + userId, + bridgeWorkerId, + workspaceId, + }); } - return await handleBlocking(req, res, { apiKeyId, userId, bridgeWorkerId }); + if (workspaceId != null) { + return res.status(400).json({ + error: 'Selected-workspace programmatic execution requires replay mode', + }); + } + return await handleBlocking(req, res, { + apiKeyId, + userId, + bridgeWorkerId, + }); } catch (err) { logger.error(`[${INSTANCE_ID}] Programmatic routing error:`, err); if (!res.headersSent) { @@ -1046,7 +1218,8 @@ router.post('/exec/programmatic', executionLimiter, async (req: t.AuthenticatedR } return; } -}); + }, +); // --------------------------------------------------------------------------- // Blocking-mode handler (extracted from the original implementation). @@ -1059,46 +1232,47 @@ async function handleBlocking( params: { apiKeyId: string; userId: string; bridgeWorkerId?: string }, ): Promise> { const { apiKeyId, userId, bridgeWorkerId } = params; - const { - code, - tools, - user_id, - files, - continuation_token, - tool_results, - } = req.body as t.ProgrammaticRequestBody; + const { code, tools, user_id, files, continuation_token, tool_results } = + req.body as t.ProgrammaticRequestBody; let timeout: number; try { - timeout = normalizeProgrammaticTimeoutMs((req.body as t.ProgrammaticRequestBody).timeout); + timeout = normalizeProgrammaticTimeoutMs( + (req.body as t.ProgrammaticRequestBody).timeout, + ); } catch (error) { return res.status(400).json({ error: (error as Error).message }); } // CASE 1: Continuation - if (continuation_token != null && continuation_token !== '' && tool_results) { + if ( + continuation_token != null && + continuation_token !== '' && + tool_results + ) { const decoded = decodeContinuationToken(continuation_token); if (!decoded) { - return res.status(400).json({ error: 'Invalid continuation token' }); + return res + .status(400) + .json({ error: 'Invalid continuation token' }); } const { execution_id } = decoded; const execution = await getExecutionState(execution_id); if (!execution) { - return res.status(404).json({ error: 'Execution not found or expired' }); + return res + .status(404) + .json({ error: 'Execution not found or expired' }); } const identity = getExecutionIdentity(req, userId); if ( execution.userId !== userId || (execution.apiKeyId != null && execution.apiKeyId !== apiKeyId) || - ( - execution.tenantId != null && - execution.tenantId !== identity.storageNamespace - ) || - ( - execution.authContextHash != null && - execution.authContextHash !== req.codeApiAuthContext?.authContextHash - ) + (execution.tenantId != null && + execution.tenantId !== identity.storageNamespace) || + (execution.authContextHash != null && + execution.authContextHash !== + req.codeApiAuthContext?.authContextHash) ) { logger.warn('Unauthorized blocking continuation request rejected', { execution_id, @@ -1122,14 +1296,19 @@ async function handleBlocking( try { await retryToolCallServerRequest( - () => axios.post(`${env.TOOL_CALL_SERVER_URL}/sessions/${execution_id}/results`, { + () => + axios.post( + `${env.TOOL_CALL_SERVER_URL}/sessions/${execution_id}/results`, + { results: tool_results.map(r => ({ call_id: r.call_id, result: r.result, is_error: r.is_error ?? false, error_message: r.error_message, })), - }, { headers: internalServiceHeaders() }), + }, + { headers: internalServiceHeaders() }, + ), 'Submit tool results', ); @@ -1174,14 +1353,21 @@ async function handleBlocking( return res.status(400).json({ error: 'Missing required field: code' }); } if (!tools || !Array.isArray(tools) || tools.length === 0) { - return res.status(400).json({ error: 'Missing required field: tools (must be a non-empty array)' }); + return res + .status(400) + .json({ + error: 'Missing required field: tools (must be a non-empty array)', + }); } if (tools.length > MAX_TOOLS_PER_REQUEST) { - logger.warn(`Too many tools provided: ${tools.length}, limit is ${MAX_TOOLS_PER_REQUEST}`, { + logger.warn( + `Too many tools provided: ${tools.length}, limit is ${MAX_TOOLS_PER_REQUEST}`, + { execution_id: 'pre-creation', userId, toolCount: tools.length, - }); + }, + ); return res.status(400).json({ error: `Too many tools provided (${tools.length}). Maximum is ${MAX_TOOLS_PER_REQUEST}.`, }); @@ -1202,7 +1388,8 @@ async function handleBlocking( files, store: connection, }); - (req.body as t.ProgrammaticRequestBody).files = authorizedFiles.length > 0 ? authorizedFiles : undefined; + (req.body as t.ProgrammaticRequestBody).files = + authorizedFiles.length > 0 ? authorizedFiles : undefined; } catch (error) { if (sendFileRefAuthorizationError(error, res, req)) return; logger.error('Error authorizing programmatic file refs:', error); @@ -1215,7 +1402,14 @@ async function handleBlocking( try { sessionKey = resolveOutputBucketSessionKey(req); } catch (error) { - if (sendSessionKeyResolutionError(error, res, req, 'programmatic /exec-blocking: resolveOutputBucketSessionKey')) { + if ( + sendSessionKeyResolutionError( + error, + res, + req, + 'programmatic /exec-blocking: resolveOutputBucketSessionKey', + ) + ) { return; } throw error; @@ -1270,24 +1464,34 @@ async function handleBlocking( try { callbackUrl = normalizeEgressGatewayUrl(env.EGRESS_GATEWAY_URL); } catch (error) { - logger.error('Blocking PTC requires egress gateway callback URL:', error); + logger.error( + 'Blocking PTC requires egress gateway callback URL:', + error, + ); await cleanupExecution(execution_id, 'blocking'); - return res.status(503).json({ error: 'Egress gateway unavailable' }); + return res + .status(503) + .json({ error: 'Egress gateway unavailable' }); } let callbackToken: string; try { const toolCallResponse = await retryToolCallServerRequest( - () => axios.post<{ + () => + axios.post<{ success: boolean; callback_token: string; - }>(`${env.TOOL_CALL_SERVER_URL}/sessions`, { + }>( + `${env.TOOL_CALL_SERVER_URL}/sessions`, + { execution_id, session_id, timeout, tools, - }, { headers: internalServiceHeaders() }), + }, + { headers: internalServiceHeaders() }, + ), 'Create Tool Call Server session', ); @@ -1299,9 +1503,14 @@ async function handleBlocking( allowedToolNames: tools.map(tool => tool.name), }); } catch (error) { - logger.error('Failed to create Tool Call Server session or callback token:', error); + logger.error( + 'Failed to create Tool Call Server session or callback token:', + error, + ); await cleanupExecution(execution_id, 'blocking'); - return res.status(503).json({ error: 'Tool Call Server unavailable' }); + return res + .status(503) + .json({ error: 'Tool Call Server unavailable' }); } let rawPayload: t.PayloadBody; @@ -1316,10 +1525,15 @@ async function handleBlocking( timeout, }); } catch (error) { - logger.error('Failed to create payload', { execution_id, error: (error as Error).message }); + logger.error('Failed to create payload', { + execution_id, + error: (error as Error).message, + }); await cleanupExecution(execution_id, 'blocking'); return res.status(400).json({ - error: (error as Error).message || 'Failed to generate code payload', + error: + (error as Error).message || + 'Failed to generate code payload', }); } const sandboxSecurity = prepareSandboxJobSecurity({ @@ -1331,7 +1545,9 @@ async function handleBlocking( payload: rawPayload, }); - const job = await pyQueue.add(Jobs.execute, { + const job = await pyQueue.add( + Jobs.execute, + { code, userId, payload: sandboxSecurity.payload, @@ -1350,18 +1566,24 @@ async function handleBlocking( ...(bridgeWorkerId != null ? { bridgeWorkerId } : {}), runtimeSessionMode: 'stateless', runtimeSessionExemption: PROGRAMMATIC_RUNTIME_SESSION_EXEMPTION, - executionManifestClaims: sandboxSecurity.executionManifestClaims, + executionManifestClaims: + sandboxSecurity.executionManifestClaims, egressGrantClaims: sandboxSecurity.egressGrantClaims, egressGrantToken: sandboxSecurity.egressGrantToken, - }, { + }, + { removeOnComplete: { age: 60, count: 1 }, removeOnFail: { age: 180, count: 1 }, attempts: 1, jobId: session_id, - }); + }, + ); jobsSubmitted.inc({ language: 'python' }); - logger.info('Job queued, polling for tool calls', { execution_id, session_id }); + logger.info('Job queued, polling for tool calls', { + execution_id, + session_id, + }); let clientDisconnected = false; req.on('close', async () => { @@ -1372,21 +1594,27 @@ async function handleBlocking( await job.remove(); await cleanupExecution(execution_id, 'blocking'); } catch (error) { - logger.error('Error cleaning up after client disconnect:', error); + logger.error( + 'Error cleaning up after client disconnect:', + error, + ); } }); job.waitUntilFinished(pyQueueEvents, JOB_COMPLETION_WAIT_TIMEOUT_MS) - .then(async (result) => { + .then(async result => { if (clientDisconnected) return; await setExecutionResult(execution_id, result); }) - .catch(async (error) => { + .catch(async error => { if (clientDisconnected) return; await setExecutionError(execution_id, error); }); - const state = await waitForExecutionState(execution_id, Math.min(timeout, MAX_POLL_TIME)); + const state = await waitForExecutionState( + execution_id, + Math.min(timeout, MAX_POLL_TIME), + ); if (state.status === 'waiting' && state.pending_calls) { return res.status(200).json({ @@ -1416,7 +1644,10 @@ async function handleBlocking( session_id, }); } catch (error) { - logger.error(`[${INSTANCE_ID}] Session ID: ${session_id} | Execution ID: ${execution_id} | Error:`, error); + logger.error( + `[${INSTANCE_ID}] Session ID: ${session_id} | Execution ID: ${execution_id} | Error:`, + error, + ); await cleanupExecution(execution_id, 'blocking'); return res.status(500).json({ error: 'Internal server error' }); } diff --git a/service/src/service/programmatic-state.test.ts b/service/src/service/programmatic-state.test.ts index fc84d8f8..81405021 100644 --- a/service/src/service/programmatic-state.test.ts +++ b/service/src/service/programmatic-state.test.ts @@ -83,6 +83,7 @@ describe('buildReplayExecutionState', () => { const state = build({ authContext, bridgeWorkerId: 'code-user_123', + workspaceId: 'project-a', sandboxBackend: 'remote-bridge', executionProfile: 'stateful', executionProfileSource: 'explicit', @@ -102,6 +103,7 @@ describe('buildReplayExecutionState', () => { authContextHash: 'hash_123', apiKeyId: 'key_legacy', bridgeWorkerId: 'code-user_123', + workspaceId: 'project-a', sandboxBackend: 'remote-bridge', executionProfile: 'stateful', executionProfileSource: 'explicit', diff --git a/service/src/service/programmatic-state.ts b/service/src/service/programmatic-state.ts index 25571fed..e6bda59a 100644 --- a/service/src/service/programmatic-state.ts +++ b/service/src/service/programmatic-state.ts @@ -38,6 +38,7 @@ export interface BuildReplayExecutionStateParams { timeout: number; language: 'python' | 'bash'; bridgeWorkerId?: string; + workspaceId?: string; sandboxBackend?: SandboxBackendName; executionProfile: ExecutionProfile; executionProfileSource: ExecutionProfileSource; @@ -66,6 +67,7 @@ export function buildReplayExecutionState( authContextHash: identity.authContextHash, apiKeyId: params.apiKeyId, bridgeWorkerId: params.bridgeWorkerId, + workspaceId: params.workspaceId, sandboxBackend: params.sandboxBackend, executionProfile: params.executionProfile, executionProfileSource: params.executionProfileSource, diff --git a/service/src/service/replay-state.ts b/service/src/service/replay-state.ts index 2254ee21..3b65cedf 100644 --- a/service/src/service/replay-state.ts +++ b/service/src/service/replay-state.ts @@ -116,6 +116,8 @@ export interface ExecutionState { apiKeyId?: string; /** Authenticated worker selection retained across every replay iteration. */ bridgeWorkerId?: string; + /** Selected workspace retained and bound across every replay iteration. */ + workspaceId?: string; /** Original queue/backend target retained across replay continuations. */ sandboxBackend?: SandboxBackendName; /** Original producer profile retained so continuations use the same queue. */ diff --git a/service/src/types/service.ts b/service/src/types/service.ts index 2a90eac7..f0a6da3b 100644 --- a/service/src/types/service.ts +++ b/service/src/types/service.ts @@ -3,7 +3,10 @@ import type { Request } from 'express'; import type { ExecutionManifestClaims } from '../execution-manifest'; import type { ExecutionIdentity } from '../execution-identity'; import type { CodeApiPrincipal } from '../auth/principal'; -import type { ExecutionProfile, SandboxBackendName } from '../execution-profile'; +import type { + ExecutionProfile, + SandboxBackendName, +} from '../execution-profile'; import { Jobs } from '@/enum/service'; /** @@ -149,7 +152,11 @@ export interface RequestBody { runtime_session_hint?: string; } -export type CreatePayload = { req: AuthenticatedRequest, session_id: string; isPyPlot?: boolean }; +export type CreatePayload = { + req: AuthenticatedRequest; + session_id: string; + isPyPlot?: boolean; +}; export interface FileObject { name: string; id: string; @@ -160,10 +167,12 @@ export interface FileObject { size?: number; lastModified?: string; etag?: string; - metadata?: { + metadata?: + | { 'content-type': string; 'original-filename': string; - } | undefined; + } + | undefined; versionId?: string | null; contentType?: string; } @@ -184,6 +193,14 @@ export type PayloadFileRef = { export interface PayloadBody { language: string; version: string; + /** Stable identity shared by all replay iterations of one execution. */ + execution_id?: string; + replay_tool_count?: number; + /** Manifest-bound upload ceiling exposed to remote workers. */ + max_output_files?: number; + /** Effective per-file ceiling after manifest and gateway policy intersect. */ + max_output_file_bytes?: number; + transfer_timeout_ms?: number; run_memory_limit?: number; run_timeout?: number; run_cpu_time?: number; @@ -238,6 +255,8 @@ export type ExecuteResult = { message?: string | null; status?: string | null; wall_time?: number | null; + /** Trusted worker control channel; avoids losing replay calls to stdout truncation. */ + pending_tool_calls_payload?: string; }; export interface LanguageConfig { @@ -265,6 +284,8 @@ export type JobData = { canonicalUserId?: string; /** Trusted dynamic outbound worker selection. */ bridgeWorkerId?: string; + /** Trusted selected workspace for native replay-mode PTC. */ + workspaceId?: string; /** Producer deployment identity. Optional only for pre-profile queued jobs. */ executionProfile?: ExecutionProfile; /** Required sandbox transport. Optional only for jobs queued before fencing. */ diff --git a/service/src/workers.ts b/service/src/workers.ts index ad20fd0f..dc2f491b 100644 --- a/service/src/workers.ts +++ b/service/src/workers.ts @@ -21,6 +21,7 @@ import { validateQueuedExecutionProfile, validateQueuedSandboxBackend, } from './execution-profile'; +import { BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILE_BYTES, programmaticTransferReserveMs } from '../../packages/code/src/protocol'; const { INSTANCE_ID } = env; const WORKER_ID = `${INSTANCE_ID}-${process.pid}`; @@ -91,9 +92,15 @@ async function processJobInner(job: t.ExecuteJob): Promise { const delivery = prepareInputDelivery(payload, sandboxPayload); const sandboxRequest = buildSandboxExecuteRequest({ + ...(job.data.workspaceId == null ? {} : { programmaticTransferReserveMs: programmaticTransferReserveMs(env.JOB_TIMEOUT) }), payload: delivery.payload, egressGrantToken, executionManifestClaims, + maxOutputFileBytes: Math.min( + executionManifestClaims?.max_upload_bytes ?? env.EGRESS_GATEWAY_MAX_FILE_BYTES, + env.EGRESS_GATEWAY_MAX_FILE_BYTES, + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILE_BYTES, + ), executionManifestPrivateKey: env.EXECUTION_MANIFEST_PRIVATE_KEY, executionManifestSecret: env.EXECUTION_MANIFEST_SECRET, executionManifestTtlSeconds: env.EXECUTION_MANIFEST_TTL_SECONDS, @@ -149,6 +156,7 @@ async function processJobInner(job: t.ExecuteJob): Promise { tenantId: job.data.tenantId, canonicalUserId: job.data.canonicalUserId, bridgeWorkerId: job.data.bridgeWorkerId, + workspaceId: job.data.workspaceId, runtimeSessionId: runtimeSession.runtimeSessionId, runtimeSessionMode: runtimeSession.runtimeSessionMode, /* Stateful backends run this as a commit barrier after user code but @@ -184,6 +192,9 @@ async function processJobInner(job: t.ExecuteJob): Promise { : {}), stdout, stderr, + ...(responseData.pending_tool_calls_payload != null + ? { pending_tool_calls_payload: responseData.pending_tool_calls_payload } + : {}), }; if (run) {