diff --git a/packages/code/README.md b/packages/code/README.md index cd48adb..499f85c 100644 --- a/packages/code/README.md +++ b/packages/code/README.md @@ -188,7 +188,7 @@ accepting another assignment. Reset or discard that session's local runner before restarting the worker; its workspace may contain mutations that Code API did not commit. -## Local workspace tools (library preview) +## Local workspace tools (bridge preview) `@librechat/code/workspace` provides the provider-neutral foundation for coding-agent access to repositories that already live on the worker machine. @@ -207,9 +207,28 @@ and stops after a bounded global result count. The worker process still belongs inside the trusted BYOM boundary and should receive filesystem access only to roots the operator intentionally registers. -This release exposes the library protocol only. The subsequent bridge-dispatch -layer will route signed, deadline-bound workspace tool assignments to it; until -that layer is configured, the CLI does not advertise or execute these tools. +Register one repository already present on the worker machine with the +Cursor-style worker-directory option: + +```bash +librechat-code run --worker-dir /path/to/repository +``` + +The default public workspace ID is `primary` and the default display name is +the directory basename. Operators can use `--workspace-id` and +`--workspace-name`, or `LIBRECHAT_CODE_WORKER_DIR`, +`LIBRECHAT_CODE_WORKSPACE_ID`, and `LIBRECHAT_CODE_WORKSPACE_NAME`, to set them +explicitly. `rg` must be installed on the worker for `search_text`. + +The worker advertises these capabilities only when a directory is configured +and executes matching assignments under the bridge's existing lease, +deadline, cancellation, credential-refresh, and settlement fencing. The +repository itself remains on the worker. As with Cursor's self-hosted agents, +text deliberately selected by `read_file` or `search_text` crosses the outbound +bridge so the remote agent/model can reason over it. Host paths are never part +of that payload. The Code API workspace-tool endpoint is delivered as a +dependent layer; deployments without it continue to use sandbox assignments +unchanged. After discarding or resetting that session's local runner, acknowledge recovery with `librechat-code reset-workspace `. The command uses the diff --git a/packages/code/src/cli.ts b/packages/code/src/cli.ts index 9aa6f58..a8cb1f9 100644 --- a/packages/code/src/cli.ts +++ b/packages/code/src/cli.ts @@ -1,7 +1,7 @@ #!/usr/bin/env node import { createHash, createHmac, randomBytes } from 'node:crypto'; import { readFileSync } from 'node:fs'; -import { resolve } from 'node:path'; +import { basename, resolve } from 'node:path'; import { pairBridgeWorker } from './pairing.js'; import { startFileRelay } from './relay.js'; @@ -13,7 +13,9 @@ import { } from './storage.js'; import { BridgeWorker } from './worker.js'; import { DockerRuntimeSupervisor, EndpointRuntimeSupervisor } from './runtime.js'; +import { LocalWorkspaceTools } from './workspace.js'; import { + BRIDGE_WORKSPACE_NAME_MAX_LENGTH, isValidBridgeWorkerCapabilities, isValidBridgeWorkerId, } from './protocol.js'; @@ -65,6 +67,14 @@ function option(args: string[], name: string): string | undefined { return args.find((value) => value.startsWith(`${name}=`))?.slice(name.length + 1); } +function defaultWorkspaceName(workerDirectory: string, workspaceId: string): string { + const directoryName = basename(resolve(workerDirectory)); + return directoryName.trim().length > 0 && + directoryName.length <= BRIDGE_WORKSPACE_NAME_MAX_LENGTH + ? directoryName + : workspaceId; +} + async function pair(args: string[]): Promise { const codeApiUrl = required('instance URL', args[1]); const code = required('one-time pairing code', args[2]); @@ -117,7 +127,7 @@ async function relay(): Promise { await handle.close(); } -async function run(runtimeSessionId?: string): Promise { +async function run(runtimeSessionId?: string, args: string[] = []): Promise { const configuredWorkerId = process.env.LIBRECHAT_CODE_WORKER_ID?.trim(); const configuredIdentityPath = process.env.LIBRECHAT_CODE_IDENTITY_FILE?.trim(); const configuredToken = process.env.LIBRECHAT_CODE_WORKER_TOKEN?.trim(); @@ -187,6 +197,29 @@ async function run(runtimeSessionId?: string): Promise { runtimeMode === 'docker-macos-nsjail' && runtimeSessionId == null && (fileRelayUpstream?.length ?? 0) > 0; + const workerDirectory = + runtimeSessionId == null + ? option(args, '--worker-dir') ?? + process.env.LIBRECHAT_CODE_WORKER_DIR?.trim() + : undefined; + const workspaceId = + option(args, '--workspace-id') ?? + process.env.LIBRECHAT_CODE_WORKSPACE_ID?.trim() ?? + 'primary'; + const workspaceTools = workerDirectory + ? await LocalWorkspaceTools.create({ + workspaces: [ + { + id: workspaceId, + name: + option(args, '--workspace-name') ?? + process.env.LIBRECHAT_CODE_WORKSPACE_NAME?.trim() ?? + defaultWorkspaceName(workerDirectory, workspaceId), + root: workerDirectory, + }, + ], + }) + : undefined; const capabilities = { statefulWorkspace, sandboxProfile: @@ -195,6 +228,7 @@ async function run(runtimeSessionId?: string): Promise { runtimes: list(process.env.LIBRECHAT_CODE_RUNTIMES), policyDigest: createHash('sha256').update(policy).digest('hex'), ...(fileRelayEnabled ? { requiresReadyConfirmation: true } : {}), + ...(workspaceTools ? { workspaceTools: workspaceTools.capabilities } : {}), }; if (!isValidBridgeWorkerCapabilities(capabilities)) { throw new Error( @@ -330,6 +364,7 @@ async function run(runtimeSessionId?: string): Promise { statefulWorkspace, }), capabilities, + workspaceTools, onIdentityChange: pairedIdentity && identityPath ? async (identity) => { @@ -401,7 +436,7 @@ async function main(): Promise { if (args[0] && args[0] !== 'run') { throw new Error(`Unknown command: ${args[0]}`); } - await run(); + await run(undefined, args.slice(1)); } main().catch((error: Error) => { process.stderr.write(`librechat-code: ${error.message}\n`); diff --git a/packages/code/src/protocol.ts b/packages/code/src/protocol.ts index 8a069c3..94d5d0e 100644 --- a/packages/code/src/protocol.ts +++ b/packages/code/src/protocol.ts @@ -1,3 +1,5 @@ +import type { WorkspaceToolRequest } from './workspace.js'; + export const BRIDGE_PROTOCOL_VERSION = 1 as const; export const BRIDGE_WORKER_ID_PATTERN = /^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$/; export const BRIDGE_SANDBOX_PROFILE_MAX_LENGTH = 128; @@ -77,7 +79,8 @@ export interface BridgeAssignment { /** Server-calculated execution budget at lease time; avoids VM clock skew. */ remainingMs?: number; runtimeSessionId?: string; - request: BridgeSandboxRequest; + executionKind?: 'sandbox' | 'workspace_tool'; + request: BridgeSandboxRequest | WorkspaceToolRequest; } export interface BridgeLeaseResponse { diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 81a2d91..302b290 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -7,10 +7,12 @@ import { } from './protocol.js'; import { EndpointRuntimeSupervisor } from './runtime.js'; import { signBridgeRequest } from './identity.js'; +import { isWorkspaceToolRequest } from './workspace.js'; import type { BridgeAssignment, BridgeLeaseResponse, + BridgeSandboxRequest, BridgeSettlement, BridgeSettlementResponse, BridgeWorkerCapabilities, @@ -18,6 +20,7 @@ import type { BridgeWorkerRegistrationResponse, } from './protocol.js'; import type { RuntimeLease, RuntimeSupervisor } from './runtime.js'; +import type { WorkspaceToolExecutor } from './workspace.js'; export interface BridgeWorkerOptions { codeApiUrl: string; @@ -28,6 +31,7 @@ export interface BridgeWorkerOptions { sandboxEndpoint?: string; runtimeSupervisor?: RuntimeSupervisor; capabilities: BridgeWorkerCapabilities; + workspaceTools?: WorkspaceToolExecutor; leaseWaitMs?: number; leaseTransportGraceMs?: number; registrationTransportTimeoutMs?: number; @@ -114,6 +118,25 @@ function errorCode(value: object): string | undefined { return undefined; } +function workspaceCapabilitiesMatch( + advertised: NonNullable, + executor: NonNullable, +): boolean { + return ( + advertised.protocolVersion === executor.protocolVersion && + advertised.operations.length === executor.operations.length && + advertised.operations.every( + (operation, index) => operation === executor.operations[index], + ) && + advertised.workspaces.length === executor.workspaces.length && + advertised.workspaces.every( + (workspace, index) => + workspace.id === executor.workspaces[index]?.id && + workspace.name === executor.workspaces[index]?.name, + ) + ); +} + export class BridgeWorkspaceQuarantinedError extends Error { constructor( message: string, @@ -147,6 +170,20 @@ export class BridgeWorker { if (options.runtimeSupervisor == null && !options.sandboxEndpoint?.trim()) { throw new BridgeProtocolError('Bridge worker requires a runtime supervisor'); } + if ( + (options.workspaceTools == null) !== + (options.capabilities.workspaceTools == null) || + (options.workspaceTools != null && + options.capabilities.workspaceTools != null && + !workspaceCapabilitiesMatch( + options.capabilities.workspaceTools, + options.workspaceTools.capabilities, + )) + ) { + throw new BridgeProtocolError( + 'Workspace tool capabilities require a matching executor', + ); + } this.fetchImpl = options.fetchImpl ?? fetch; this.codeApiUrl = normalizedBaseUrl(options.codeApiUrl); this.runtimeSupervisor = @@ -622,57 +659,114 @@ export class BridgeWorker { credentialMaintenanceError = error; executionController.abort(); }); - runtimeLease = await this.runtimeSupervisor.acquire( - assignment, - executionController.signal, - ); + let payload: object = {}; + if (assignment.executionKind === 'workspace_tool') { + if (this.options.workspaceTools == null) { + throw new BridgeProtocolError( + 'Worker does not provide local workspace tools', + ); + } + if (!isWorkspaceToolRequest(assignment.request)) { + throw new BridgeProtocolError('Invalid workspace tool request'); + } + const workspaceRequest = assignment.request; + const advertised = this.options.workspaceTools.capabilities; + if (!advertised.operations.includes(workspaceRequest.operation)) { + throw new BridgeProtocolError( + 'Workspace tool operation is not advertised', + ); + } + if ( + !advertised.workspaces.some( + (workspace) => workspace.id === workspaceRequest.workspaceId, + ) + ) { + throw new BridgeProtocolError('Workspace is not advertised'); + } + payload = await this.options.workspaceTools.execute( + workspaceRequest, + executionController.signal, + ); + if (executionController.signal.aborted) { + throw ( + executionController.signal.reason ?? + new DOMException('aborted', 'AbortError') + ); + } + if (Date.now() >= localDeadlineAtMs) { + throw new BridgeProtocolError( + 'Bridge assignment expired during workspace execution', + ); + } + } else { + runtimeLease = await this.runtimeSupervisor.acquire( + assignment, + executionController.signal, + ); + if (executionController.signal.aborted) { + throw ( + executionController.signal.reason ?? + new DOMException('aborted', 'AbortError') + ); + } + const sandboxRequest = assignment.request as BridgeSandboxRequest; + const headers = { + ...sandboxRequest.headers, + ...(runtimeLease.sessionId + ? { 'X-Runtime-Session-Id': runtimeLease.sessionId } + : {}), + }; + const sandboxRequestBody = JSON.stringify(sandboxRequest.body); + if (Date.now() >= localDeadlineAtMs) { + throw new BridgeProtocolError( + 'Bridge assignment expired before sandbox execution', + ); + } + sandboxStarted = true; + const response = await this.executeRuntime( + runtimeLease, + sandboxRequestBody, + { + ...headers, + 'Content-Type': 'application/json', + }, + executionController.signal, + ); + try { + payload = JSON.parse(response.body) as object; + } catch (error) { + if (response.status >= 200 && response.status < 300) throw error; + } + if (response.status < 200 || response.status >= 300) { + sandboxRejectedExecution = + response.status >= 400 && + response.status < 500 && + response.status !== 408 && + response.status !== 429 && + errorMessage(payload) !== 'session_workspace_dirty'; + throw new BridgeProtocolError( + errorMessage(payload) ?? + `Sandbox rejected execution with HTTP ${response.status}`, + response.status, + ); + } + } + cancellationController.abort(); + await cancellationWatcher; if (executionController.signal.aborted) { - throw executionController.signal.reason ?? new DOMException('aborted', 'AbortError'); + throw ( + executionController.signal.reason ?? + new DOMException('aborted', 'AbortError') + ); } - const headers = { - ...assignment.request.headers, - ...(runtimeLease.sessionId - ? { 'X-Runtime-Session-Id': runtimeLease.sessionId } - : {}), - }; - const sandboxRequestBody = JSON.stringify(assignment.request.body); if (Date.now() >= localDeadlineAtMs) { throw new BridgeProtocolError( - 'Bridge assignment expired before sandbox execution', + 'Bridge assignment expired while draining cancellation', ); } - sandboxStarted = true; - const response = await this.executeRuntime( - runtimeLease, - sandboxRequestBody, - { - ...headers, - 'Content-Type': 'application/json', - }, - executionController.signal, - ); - let payload: object = {}; - try { - payload = JSON.parse(response.body) as object; - } catch (error) { - if (response.status >= 200 && response.status < 300) throw error; - } if (credentialMaintenanceError != null) { throw credentialMaintenanceError; } - if (response.status < 200 || response.status >= 300) { - sandboxRejectedExecution = - response.status >= 400 && - response.status < 500 && - response.status !== 408 && - response.status !== 429 && - errorMessage(payload) !== 'session_workspace_dirty'; - throw new BridgeProtocolError( - errorMessage(payload) ?? - `Sandbox rejected execution with HTTP ${response.status}`, - response.status, - ); - } if (heartbeatError != null) throw heartbeatError; settlement = { protocolVersion: BRIDGE_PROTOCOL_VERSION, @@ -1005,7 +1099,9 @@ export class BridgeWorker { if (signal.aborted || executionController.signal.aborted) return; const pollController = new AbortController(); const abortPoll = (): void => pollController.abort(); - signal.addEventListener('abort', abortPoll, { once: true }); + executionController.signal.addEventListener('abort', abortPoll, { + once: true, + }); const timeout = setTimeout( abortPoll, Math.max( @@ -1028,14 +1124,14 @@ export class BridgeWorker { return; } } catch (error) { - if (signal.aborted) return; if (error instanceof BridgeProtocolError && error.status === 404) { executionController.abort(); return; } + if (signal.aborted) return; } finally { clearTimeout(timeout); - signal.removeEventListener('abort', abortPoll); + executionController.signal.removeEventListener('abort', abortPoll); } } } diff --git a/packages/code/src/workspace-cli.test.ts b/packages/code/src/workspace-cli.test.ts new file mode 100644 index 0000000..4ac486e --- /dev/null +++ b/packages/code/src/workspace-cli.test.ts @@ -0,0 +1,118 @@ +import assert from 'node:assert/strict'; +import { spawn, spawnSync } from 'node:child_process'; +import { once } from 'node:events'; +import { mkdtemp, mkdir, rm } from 'node:fs/promises'; +import { createServer } from 'node:http'; +import { join } from 'node:path'; +import { tmpdir } from 'node:os'; +import { fileURLToPath } from 'node:url'; +import test from 'node:test'; + +test('CLI validates a configured worker directory before registration', () => { + const result = spawnSync( + process.execPath, + [ + fileURLToPath(new URL('./cli.js', import.meta.url)), + 'run', + '--worker-dir', + '/definitely/missing/librechat-code-workspace', + ], + { + encoding: 'utf8', + env: { + ...process.env, + LIBRECHAT_CODE_URL: 'https://code.example/v1', + LIBRECHAT_CODE_WORKER_TOKEN: 'worker-secret', + LIBRECHAT_CODE_WORKER_ID: 'engineering-vm', + }, + }, + ); + + assert.notEqual(result.status, 0); + assert.match(result.stderr, /invalid workspace registration/i); +}); + +test('CLI falls back to the workspace ID when the directory basename is invalid', async (t) => { + const root = await mkdtemp(join(tmpdir(), 'librechat-code-cli-')); + const workspaceRoot = join(root, ' '); + await mkdir(workspaceRoot); + t.after(() => rm(root, { recursive: true, force: true })); + let resolveRegistration: ((value: Record) => void) | undefined; + const registration = new Promise>((resolve) => { + resolveRegistration = resolve; + }); + const server = createServer((request, response) => { + const chunks: Buffer[] = []; + request.on('data', (chunk: Buffer) => chunks.push(chunk)); + request.on('end', () => { + const body = JSON.parse(Buffer.concat(chunks).toString('utf8')) as Record< + string, + unknown + >; + if (request.url?.endsWith('/bridge/workers/register')) { + resolveRegistration?.(body); + response.setHeader('Content-Type', 'application/json'); + response.end( + JSON.stringify({ + protocolVersion: 1, + workerId: body.workerId, + incarnationId: body.incarnationId, + registeredAt: new Date().toISOString(), + leaseTtlMs: 60_000, + }), + ); + return; + } + response.setHeader('Content-Type', 'application/json'); + response.end(JSON.stringify({ protocolVersion: 1, serverElapsedMs: 0 })); + }); + }); + await new Promise((resolve) => server.listen(0, '127.0.0.1', resolve)); + t.after(() => server.close()); + const address = server.address(); + if (address == null || typeof address === 'string') { + assert.fail('expected TCP listener'); + } + + const child = spawn( + process.execPath, + [ + fileURLToPath(new URL('./cli.js', import.meta.url)), + 'run', + '--worker-dir', + workspaceRoot, + '--workspace-id', + 'root-workspace', + ], + { + env: { + ...process.env, + LIBRECHAT_CODE_URL: `http://127.0.0.1:${address.port}/v1`, + LIBRECHAT_CODE_WORKER_TOKEN: 'worker-secret', + LIBRECHAT_CODE_WORKER_ID: 'engineering-vm', + LIBRECHAT_CODE_SANDBOX_ENDPOINT: + 'http://127.0.0.1:2000/api/v2', + }, + stdio: 'ignore', + }, + ); + t.after(() => child.kill()); + + const body = await Promise.race([ + registration, + new Promise((_, reject) => + setTimeout(() => reject(new Error('registration timed out')), 2_000), + ), + ]); + child.kill(); + await once(child, 'exit'); + + assert.deepEqual( + (body.capabilities as Record).workspaceTools, + { + protocolVersion: 1, + operations: ['read_file', 'search_text'], + workspaces: [{ id: 'root-workspace', name: 'root-workspace' }], + }, + ); +}); diff --git a/packages/code/src/workspace-worker.test.ts b/packages/code/src/workspace-worker.test.ts new file mode 100644 index 0000000..e206311 --- /dev/null +++ b/packages/code/src/workspace-worker.test.ts @@ -0,0 +1,684 @@ +import assert from 'node:assert/strict'; +import test from 'node:test'; + +import { BridgeWorker } from './worker.js'; + +const incarnationId = 'incarnation-00000001'; + +test('worker executes a workspace tool assignment locally without acquiring a sandbox', async () => { + const requests: Array<{ url: string; init?: RequestInit }> = []; + const workspaceRequests: object[] = []; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId, + runtimeSupervisor: { + async acquire() { + throw new Error('workspace tools must not acquire a sandbox'); + }, + async reset() {}, + async quarantine() {}, + }, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + workspaceTools: { + protocolVersion: 1, + operations: ['read_file', 'search_text'], + workspaces: [{ id: 'primary', name: 'LibreChat' }], + }, + }, + workspaceTools: { + capabilities: { + protocolVersion: 1, + operations: ['read_file', 'search_text'], + workspaces: [{ id: 'primary', name: 'LibreChat' }], + }, + async execute(request) { + workspaceRequests.push(request); + return { + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'primary', + path: 'README.md', + content: '# LibreChat', + startLine: 1, + endLine: 1, + truncated: false, + }; + }, + }, + fetchImpl: async (input, init) => { + requests.push({ url: String(input), init }); + return Response.json({ protocolVersion: 1, accepted: true }); + }, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-workspace-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_tool', + request: { + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'primary', + path: 'README.md', + }, + }); + + assert.deepEqual(workspaceRequests, [ + { + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'primary', + path: 'README.md', + }, + ]); + assert.equal(requests.length, 1); + assert.deepEqual(JSON.parse(String(requests[0].init?.body)), { + protocolVersion: 1, + generation: 4, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + incarnationId, + status: 'fulfilled', + result: { + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'primary', + path: 'README.md', + content: '# LibreChat', + startLine: 1, + endLine: 1, + truncated: false, + }, + }); +}); + +test('worker refuses to advertise workspace tools without a matching executor', () => { + assert.throws( + () => + new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + workspaceTools: { + protocolVersion: 1, + operations: ['read_file'], + workspaces: [{ id: 'primary' }], + }, + }, + }), + /workspace tool capabilities require a matching executor/i, + ); +}); + +test('worker compares workspace capabilities structurally', () => { + assert.doesNotThrow( + () => + new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + workspaceTools: { + protocolVersion: 1, + operations: ['read_file'], + workspaces: [{ id: 'primary', name: 'LibreChat' }], + }, + }, + workspaceTools: { + capabilities: { + operations: ['read_file'], + workspaces: [{ name: 'LibreChat', id: 'primary' }], + protocolVersion: 1, + }, + async execute() { + throw new Error('not executed'); + }, + }, + }), + ); +}); + +test('worker rejects a workspace result returned after its deadline', async () => { + const settlements: Array> = []; + const workspaceCapabilities = { + protocolVersion: 1 as const, + operations: ['read_file' 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: 'nsjail', + runtimes: ['bash'], + workspaceTools: workspaceCapabilities, + }, + workspaceTools: { + capabilities: workspaceCapabilities, + async execute(request, signal) { + await new Promise((resolve) => { + if (signal?.aborted) return resolve(); + signal?.addEventListener('abort', () => resolve(), { once: true }); + }); + return { + protocolVersion: 1, + operation: 'read_file', + workspaceId: request.workspaceId, + path: 'README.md', + content: '# late', + startLine: 1, + endLine: 1, + truncated: false, + }; + }, + }, + fetchImpl: async (_input, init) => { + if (init?.body != null) { + settlements.push( + JSON.parse(String(init.body)) as Record, + ); + } + return Response.json({ protocolVersion: 1, accepted: true }); + }, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-workspace-deadline', + workerId: 'vm-1', + incarnationId, + generation: 4, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 20).toISOString(), + remainingMs: 20, + executionKind: 'workspace_tool', + request: { + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'primary', + path: 'README.md', + }, + }); + + assert.equal(settlements.length, 1); + assert.equal(settlements[0]?.status, 'rejected'); + assert.match(String(settlements[0]?.error), /aborted|expired/i); +}); + +test('worker drains a completed cancellation poll before fulfilling workspace work', async () => { + const settlements: Array> = []; + let finishExecution: (() => void) | undefined; + let finishCancellation: (() => void) | undefined; + let markPollStarted: (() => void) | undefined; + const pollStarted = new Promise((resolve) => { + markPollStarted = resolve; + }); + const workspaceCapabilities = { + protocolVersion: 1 as const, + operations: ['read_file' 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: 'nsjail', + runtimes: ['bash'], + workspaceTools: workspaceCapabilities, + }, + workspaceTools: { + capabilities: workspaceCapabilities, + async execute(request) { + await new Promise((resolve) => { + finishExecution = resolve; + }); + return { + protocolVersion: 1, + operation: 'read_file', + workspaceId: request.workspaceId, + path: 'README.md', + content: '# cancelled', + startLine: 1, + endLine: 1, + truncated: false, + }; + }, + }, + cancellationPollIntervalMs: 1, + fetchImpl: async (input, init) => { + if (String(input).endsWith('/cancellation')) { + markPollStarted?.(); + return await new Promise((resolve) => { + finishCancellation = () => + resolve(Response.json({ protocolVersion: 1, cancelled: true })); + }); + } + if (String(input).endsWith('/settle') && init?.body != null) { + settlements.push( + JSON.parse(String(init.body)) as Record, + ); + } + return Response.json({ protocolVersion: 1, accepted: true }); + }, + }); + + const completion = worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-workspace-cancelled', + workerId: 'vm-1', + incarnationId, + generation: 4, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 5_000).toISOString(), + remainingMs: 5_000, + executionKind: 'workspace_tool', + request: { + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'primary', + path: 'README.md', + }, + }); + + await pollStarted; + finishExecution?.(); + finishCancellation?.(); + await completion; + + assert.equal(settlements.length, 1); + assert.equal(settlements[0]?.status, 'rejected'); + assert.match(String(settlements[0]?.error), /aborted/i); +}); + +test('worker drains a cancellation response body before fulfilling workspace work', async () => { + const settlements: Array> = []; + let finishExecution: (() => void) | undefined; + let markHeadersReceived: (() => void) | undefined; + const headersReceived = new Promise((resolve) => { + markHeadersReceived = resolve; + }); + const workspaceCapabilities = { + protocolVersion: 1 as const, + operations: ['read_file' 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: 'nsjail', + runtimes: ['bash'], + workspaceTools: workspaceCapabilities, + }, + workspaceTools: { + capabilities: workspaceCapabilities, + async execute(request) { + await new Promise((resolve) => { + finishExecution = resolve; + }); + return { + protocolVersion: 1, + operation: 'read_file', + workspaceId: request.workspaceId, + path: 'README.md', + content: '# cancelled', + startLine: 1, + endLine: 1, + truncated: false, + }; + }, + }, + cancellationPollIntervalMs: 1, + fetchImpl: async (input, init) => { + if (String(input).endsWith('/cancellation')) { + const response = new Response( + new ReadableStream({ + start(controller) { + let finished = false; + init?.signal?.addEventListener( + 'abort', + () => { + if (finished) return; + finished = true; + controller.error(new DOMException('aborted', 'AbortError')); + }, + { once: true }, + ); + setTimeout(() => { + if (!init?.signal?.aborted && !finished) { + finished = true; + controller.enqueue( + new TextEncoder().encode( + JSON.stringify({ protocolVersion: 1, cancelled: true }), + ), + ); + controller.close(); + } + }, 0); + }, + }), + { headers: { 'Content-Type': 'application/json' } }, + ); + markHeadersReceived?.(); + return response; + } + if (String(input).endsWith('/settle') && init?.body != null) { + settlements.push( + JSON.parse(String(init.body)) as Record, + ); + } + return Response.json({ protocolVersion: 1, accepted: true }); + }, + }); + + const completion = worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-workspace-cancelled-body', + workerId: 'vm-1', + incarnationId, + generation: 4, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 5_000).toISOString(), + remainingMs: 5_000, + executionKind: 'workspace_tool', + request: { + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'primary', + path: 'README.md', + }, + }); + + await headersReceived; + finishExecution?.(); + await completion; + + assert.equal(settlements.length, 1); + assert.equal(settlements[0]?.status, 'rejected'); + assert.match(String(settlements[0]?.error), /aborted/i); +}); + +test('worker rechecks its deadline after draining cancellation', async () => { + const settlements: Array> = []; + let finishExecution: (() => void) | undefined; + let releaseBody: (() => void) | undefined; + let markHeadersReceived: (() => void) | undefined; + const headersReceived = new Promise((resolve) => { + markHeadersReceived = resolve; + }); + const bodyReleased = new Promise((resolve) => { + releaseBody = resolve; + }); + const workspaceCapabilities = { + protocolVersion: 1 as const, + operations: ['read_file' 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: 'nsjail', + runtimes: ['bash'], + workspaceTools: workspaceCapabilities, + }, + workspaceTools: { + capabilities: workspaceCapabilities, + async execute(request) { + await new Promise((resolve) => { + finishExecution = resolve; + }); + return { + protocolVersion: 1, + operation: 'read_file', + workspaceId: request.workspaceId, + path: 'README.md', + content: '# late', + startLine: 1, + endLine: 1, + truncated: false, + }; + }, + }, + cancellationPollIntervalMs: 1, + fetchImpl: async (input, init) => { + if (String(input).endsWith('/cancellation')) { + markHeadersReceived?.(); + return { + ok: true, + status: 200, + async json() { + await bodyReleased; + const blockedUntil = Date.now() + 60; + while (Date.now() < blockedUntil) { + // Model synchronous body parsing that crosses the deadline. + } + return { protocolVersion: 1, cancelled: false }; + }, + } as Response; + } + if (String(input).endsWith('/settle') && init?.body != null) { + settlements.push( + JSON.parse(String(init.body)) as Record, + ); + } + return Response.json({ protocolVersion: 1, accepted: true }); + }, + }); + + const completion = worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-workspace-drain-deadline', + workerId: 'vm-1', + incarnationId, + generation: 4, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 50).toISOString(), + remainingMs: 50, + executionKind: 'workspace_tool', + request: { + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'primary', + path: 'README.md', + }, + }); + + await headersReceived; + finishExecution?.(); + await new Promise((resolve) => setImmediate(resolve)); + releaseBody?.(); + await completion; + + assert.equal(settlements.length, 1); + assert.equal(settlements[0]?.status, 'rejected'); + assert.match(String(settlements[0]?.error), /expired/i); +}); + +test('worker preserves a drained 404 cancellation response', async () => { + const settlements: Array> = []; + let finishExecution: (() => void) | undefined; + let releaseBody: (() => void) | undefined; + let markHeadersReceived: (() => void) | undefined; + const headersReceived = new Promise((resolve) => { + markHeadersReceived = resolve; + }); + const bodyReleased = new Promise((resolve) => { + releaseBody = resolve; + }); + const workspaceCapabilities = { + protocolVersion: 1 as const, + operations: ['read_file' 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: 'nsjail', + runtimes: ['bash'], + workspaceTools: workspaceCapabilities, + }, + workspaceTools: { + capabilities: workspaceCapabilities, + async execute(request) { + await new Promise((resolve) => { + finishExecution = resolve; + }); + return { + protocolVersion: 1, + operation: 'read_file', + workspaceId: request.workspaceId, + path: 'README.md', + content: '# cancelled', + startLine: 1, + endLine: 1, + truncated: false, + }; + }, + }, + cancellationPollIntervalMs: 1, + fetchImpl: async (input, init) => { + if (String(input).endsWith('/cancellation')) { + markHeadersReceived?.(); + return { + ok: false, + status: 404, + async json() { + await bodyReleased; + return {}; + }, + } as Response; + } + if (String(input).endsWith('/settle') && init?.body != null) { + settlements.push( + JSON.parse(String(init.body)) as Record, + ); + } + return Response.json({ protocolVersion: 1, accepted: true }); + }, + }); + + const completion = worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-workspace-cancelled-404', + workerId: 'vm-1', + incarnationId, + generation: 4, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 5_000).toISOString(), + remainingMs: 5_000, + executionKind: 'workspace_tool', + request: { + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'primary', + path: 'README.md', + }, + }); + + await headersReceived; + finishExecution?.(); + await new Promise((resolve) => setImmediate(resolve)); + releaseBody?.(); + await completion; + + assert.equal(settlements.length, 1); + assert.equal(settlements[0]?.status, 'rejected'); + assert.match(String(settlements[0]?.error), /aborted/i); +}); + +test('worker rejects workspace operations outside its advertised capability', async () => { + let executions = 0; + let settlement: Record | undefined; + const workspaceCapabilities = { + protocolVersion: 1 as const, + operations: ['read_file' 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: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + workspaceTools: workspaceCapabilities, + }, + workspaceTools: { + capabilities: workspaceCapabilities, + async execute() { + executions += 1; + throw new Error('must not execute'); + }, + }, + fetchImpl: async (_input, init) => { + settlement = JSON.parse(String(init?.body)) as Record; + return Response.json({ protocolVersion: 1, accepted: true }); + }, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-workspace-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_tool', + request: { + protocolVersion: 1, + operation: 'search_text', + workspaceId: 'primary', + query: 'needle', + }, + }); + + assert.equal(executions, 0); + assert.equal(settlement?.status, 'rejected'); + assert.match(String(settlement?.error), /operation is not advertised/i); +}); diff --git a/packages/code/src/workspace.ts b/packages/code/src/workspace.ts index b56179c..bde330a 100644 --- a/packages/code/src/workspace.ts +++ b/packages/code/src/workspace.ts @@ -27,6 +27,14 @@ export interface LocalWorkspaceToolsOptions { workspaces: LocalWorkspaceConfig[]; } +export interface WorkspaceToolExecutor { + capabilities: BridgeWorkspaceToolCapabilities; + execute( + request: WorkspaceToolRequest, + signal?: AbortSignal, + ): Promise; +} + export interface WorkspaceReadFileRequest { protocolVersion: BridgeProtocolVersion; operation: 'read_file'; @@ -535,7 +543,7 @@ async function searchWorkspace( }; } -export class LocalWorkspaceTools { +export class LocalWorkspaceTools implements WorkspaceToolExecutor { readonly capabilities: BridgeWorkspaceToolCapabilities; private constructor(