diff --git a/docs/remote-bridge/README.md b/docs/remote-bridge/README.md index 8d10f02..5fbb5b3 100644 --- a/docs/remote-bridge/README.md +++ b/docs/remote-bridge/README.md @@ -86,6 +86,52 @@ locally, proves possession on every request, and rotates its short-lived credential before expiry. `CODEAPI_BRIDGE_AUTH_MODE=static` remains available for non-hardened development compatibility only. +To expose an existing checkout as a worker-local workspace, start the CLI with +an explicit directory and logical ID: + +```bash +librechat-code run \ + --worker-dir /srv/checkouts/librechat \ + --workspace-id primary \ + --workspace-name LibreChat +``` + +The worker advertises only the workspace ID, optional display name, and +supported operations. Its host path is never registered with Code API. An +authenticated caller can execute the initial read-only operations through: + +```bash +curl -fsS https://code.example.com/v1/workspace-tools/execute \ + -H "Authorization: Bearer $LIBRECHAT_JWT" \ + -H 'Content-Type: application/json' \ + --data '{ + "protocolVersion":1, + "operation":"read_file", + "workspaceId":"primary", + "path":"README.md", + "startLine":1, + "maxLines":200 + }' +``` + +The endpoint uses the same authenticated principal-bound worker selection, +tenant fence, lease deadline, cancellation, and settlement lifecycle as remote +sandbox execution. Requests must name a workspace and operation advertised by +that worker. Results are validated against the originating request before they +leave Code API, and are bounded to 1 MiB/500 lines for reads or 200 matches for +searches. Absolute paths, traversal, backslashes, symlink escapes, unexpected +fields, and host roots are rejected. + +The workspace root can be an existing project, a Git repository, or an empty +directory; Git is not required. This boundary keeps that directory local to the +operator's machine, but the selected file contents, search matches, and later +tool results necessarily cross the outbound bridge to Code API and the model. +Treat them as explicit tool outputs, apply the same retention and audit policy +as chat content, and do not register a directory containing secrets. The +default operations are read-only; future mutation and shell operations must be +gated by LibreChat's tool-approval hooks in addition to worker capability +checks. + Stateful deployments must also set `LIBRECHAT_CODE_STATEFUL_WORKSPACE=true` and route the CLI's `{runtimeSessionId}` endpoint template to an isolated, persistent local runner per session. A single sandbox endpoint is stateless and diff --git a/packages/code/README.md b/packages/code/README.md index 499f85c..076f831 100644 --- a/packages/code/README.md +++ b/packages/code/README.md @@ -191,7 +191,9 @@ API did not commit. ## 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. +coding-agent access to workspace directories on the worker machine. A workspace +may be an existing project, a Git repository, or a newly created empty +directory; Git is optional. `LocalWorkspaceTools` registers opaque workspace IDs with optional display names and exposes bounded `read_file` and literal `search_text` operations. Only IDs, names, protocol version, and supported operations appear in worker @@ -207,11 +209,11 @@ 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. -Register one repository already present on the worker machine with the -Cursor-style worker-directory option: +Register one directory already present on the worker machine with the +worker-directory option: ```bash -librechat-code run --worker-dir /path/to/repository +librechat-code run --worker-dir /path/to/workspace ``` The default public workspace ID is `primary` and the default display name is @@ -223,7 +225,7 @@ 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, +workspace 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 diff --git a/packages/code/src/protocol.ts b/packages/code/src/protocol.ts index 94d5d0e..286b812 100644 --- a/packages/code/src/protocol.ts +++ b/packages/code/src/protocol.ts @@ -1,5 +1,3 @@ -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; @@ -7,6 +5,12 @@ export const BRIDGE_RUNTIME_MAX_COUNT = 32; export const BRIDGE_RUNTIME_MAX_LENGTH = 64; export const BRIDGE_WORKSPACE_MAX_COUNT = 32; export const BRIDGE_WORKSPACE_NAME_MAX_LENGTH = 128; +export const BRIDGE_WORKSPACE_PATH_MAX_LENGTH = 4096; +export const BRIDGE_WORKSPACE_QUERY_MAX_LENGTH = 4096; +export const BRIDGE_WORKSPACE_READ_MAX_BYTES = 1024 * 1024; +export const BRIDGE_WORKSPACE_READ_MAX_LINES = 500; +export const BRIDGE_WORKSPACE_SEARCH_MAX_RESULTS = 200; +export const BRIDGE_WORKSPACE_SEARCH_TEXT_MAX_LENGTH = 2000; export type BridgeProtocolVersion = typeof BRIDGE_PROTOCOL_VERSION; @@ -23,6 +27,99 @@ export interface BridgeWorkspaceToolCapabilities { workspaces: BridgeWorkspaceDescriptor[]; } +export interface WorkspaceReadFileRequest { + protocolVersion: BridgeProtocolVersion; + operation: 'read_file'; + workspaceId: string; + path: string; + startLine?: number; + maxLines?: number; +} + +export interface WorkspaceReadFileResult { + protocolVersion: BridgeProtocolVersion; + operation: 'read_file'; + workspaceId: string; + path: string; + content: string; + startLine: number; + endLine: number; + truncated: boolean; + nextStartLine?: number; +} + +export interface WorkspaceSearchTextRequest { + protocolVersion: BridgeProtocolVersion; + operation: 'search_text'; + workspaceId: string; + query: string; + path?: string; + maxResults?: number; +} + +export interface WorkspaceSearchMatch { + path: string; + line: number; + column: number; + text: string; +} + +export interface WorkspaceSearchTextResult { + protocolVersion: BridgeProtocolVersion; + operation: 'search_text'; + workspaceId: string; + matches: WorkspaceSearchMatch[]; + truncated: boolean; +} + +export type WorkspaceToolRequest = + | WorkspaceReadFileRequest + | WorkspaceSearchTextRequest; +export type WorkspaceToolResult = + | WorkspaceReadFileResult + | WorkspaceSearchTextResult; + +const WORKSPACE_READ_REQUEST_KEYS = new Set([ + 'protocolVersion', + 'operation', + 'workspaceId', + 'path', + 'startLine', + 'maxLines', +]); +const WORKSPACE_SEARCH_REQUEST_KEYS = new Set([ + 'protocolVersion', + 'operation', + 'workspaceId', + 'query', + 'path', + 'maxResults', +]); +const WORKSPACE_READ_RESULT_KEYS = new Set([ + 'protocolVersion', + 'operation', + 'workspaceId', + 'path', + 'content', + 'startLine', + 'endLine', + 'truncated', + 'nextStartLine', +]); +const WORKSPACE_SEARCH_RESULT_KEYS = new Set([ + 'protocolVersion', + 'operation', + 'workspaceId', + 'matches', + 'truncated', +]); +const WORKSPACE_SEARCH_MATCH_KEYS = new Set([ + 'path', + 'line', + 'column', + 'text', +]); + export interface BridgeWorkerCapabilities { statefulWorkspace: boolean; sandboxProfile: string; @@ -106,6 +203,35 @@ export interface BridgeRejectedSettlement { incarnationId: string; status: 'rejected'; error: string; + errorCode?: WorkspaceToolErrorCode; +} + +export type WorkspaceToolErrorCode = + | 'INVALID_PATH' + | 'INVALID_REQUEST' + | 'READ_LIMIT_EXCEEDED' + | 'REGISTRATION_INVALID' + | 'EXECUTION_ABORTED' + | 'SEARCH_TIMEOUT' + | 'SEARCH_UNAVAILABLE'; + +const WORKSPACE_TOOL_ERROR_CODES = new Set([ + 'INVALID_PATH', + 'INVALID_REQUEST', + 'READ_LIMIT_EXCEEDED', + 'REGISTRATION_INVALID', + 'EXECUTION_ABORTED', + 'SEARCH_TIMEOUT', + 'SEARCH_UNAVAILABLE', +]); + +export function isWorkspaceToolErrorCode( + value: unknown, +): value is WorkspaceToolErrorCode { + return ( + typeof value === 'string' && + WORKSPACE_TOOL_ERROR_CODES.has(value as WorkspaceToolErrorCode) + ); } export type BridgeSettlement = @@ -140,6 +266,169 @@ export function isValidBridgeWorkerId(workerId: string): boolean { return BRIDGE_WORKER_ID_PATTERN.test(workerId); } +function isSafePortableRelativePath(value: unknown): value is string { + if ( + typeof value !== 'string' || + value.length === 0 || + value.length > BRIDGE_WORKSPACE_PATH_MAX_LENGTH || + Buffer.from(value).toString('utf8') !== value || + value.includes('\0') || + value.includes('\\') || + value.startsWith('/') || + /^[A-Za-z]:/.test(value) + ) { + return false; + } + return value.split('/').every((segment) => segment !== '..'); +} + +function normalizePortableRelativePath(value: string): string { + return ( + value + .split('/') + .filter((segment) => segment.length > 0 && segment !== '.') + .join('/') || '.' + ); +} + +function isWithinRequestedPath(candidate: string, requested?: string): boolean { + if (requested == null) return true; + const normalizedCandidate = normalizePortableRelativePath(candidate); + const normalizedRequested = normalizePortableRelativePath(requested); + return ( + normalizedRequested === '.' || + normalizedCandidate === normalizedRequested || + normalizedCandidate.startsWith(`${normalizedRequested}/`) + ); +} + +function hasOnlyKeys( + value: Record, + allowed: ReadonlySet, +): boolean { + return Object.keys(value).every((key) => allowed.has(key)); +} + +export function isWorkspaceToolRequest( + value: unknown, +): value is WorkspaceToolRequest { + if (typeof value !== 'object' || value === null) return false; + const request = value as Record; + if ( + request.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + typeof request.workspaceId !== 'string' || + !isValidBridgeWorkerId(request.workspaceId) + ) { + return false; + } + if (request.operation === 'read_file') { + return ( + hasOnlyKeys(request, WORKSPACE_READ_REQUEST_KEYS) && + isSafePortableRelativePath(request.path) && + (request.startLine === undefined || + (Number.isSafeInteger(request.startLine) && + Number(request.startLine) >= 1)) && + (request.maxLines === undefined || + (Number.isSafeInteger(request.maxLines) && + Number(request.maxLines) >= 1 && + Number(request.maxLines) <= BRIDGE_WORKSPACE_READ_MAX_LINES)) + ); + } + if (request.operation === 'search_text') { + return ( + hasOnlyKeys(request, WORKSPACE_SEARCH_REQUEST_KEYS) && + typeof request.query === 'string' && + request.query.length > 0 && + request.query.length <= BRIDGE_WORKSPACE_QUERY_MAX_LENGTH && + Buffer.from(request.query).toString('utf8') === request.query && + new TextEncoder().encode(request.query).byteLength <= + BRIDGE_WORKSPACE_SEARCH_TEXT_MAX_LENGTH && + !request.query.includes('\0') && + !request.query.includes('\n') && + !request.query.includes('\r') && + (request.path === undefined || + isSafePortableRelativePath(request.path)) && + (request.maxResults === undefined || + (Number.isSafeInteger(request.maxResults) && + Number(request.maxResults) >= 1 && + Number(request.maxResults) <= BRIDGE_WORKSPACE_SEARCH_MAX_RESULTS)) + ); + } + return false; +} + +export function isWorkspaceToolResult( + request: WorkspaceToolRequest, + value: unknown, +): value is WorkspaceToolResult { + if (typeof value !== 'object' || value === null) return false; + const result = value as Record; + if ( + result.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + result.operation !== request.operation || + result.workspaceId !== request.workspaceId || + typeof result.truncated !== 'boolean' + ) { + return false; + } + + 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 reportedLineCount = + 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; + return ( + hasOnlyKeys(result, WORKSPACE_READ_RESULT_KEYS) && + result.path === request.path && + isSafePortableRelativePath(result.path) && + content !== null && + new TextEncoder().encode(content).byteLength <= + BRIDGE_WORKSPACE_READ_MAX_BYTES && + result.startLine === startLine && + Number.isSafeInteger(result.endLine) && + Number(result.endLine) >= startLine - 1 && + Number(result.endLine) < startLine + maxLines && + reportedLineCount >= 0 && + reportedLineCount <= maxLines && + (content.length !== 0 || reportedLineCount <= 1) && + actualLineCount === reportedLineCount && + (result.truncated === true + ? Number.isSafeInteger(result.nextStartLine) && + Number(result.nextStartLine) === Number(result.endLine) + 1 && + Number(result.nextStartLine) > startLine + : result.nextStartLine === undefined) + ); + } + + if (!Array.isArray(result.matches)) return false; + const maxResults = request.maxResults ?? 50; + return ( + hasOnlyKeys(result, WORKSPACE_SEARCH_RESULT_KEYS) && + result.matches.length <= maxResults && + result.matches.every((match) => { + if (typeof match !== 'object' || match === null) return false; + const candidate = match as Record; + return ( + hasOnlyKeys(candidate, WORKSPACE_SEARCH_MATCH_KEYS) && + isSafePortableRelativePath(candidate.path) && + isWithinRequestedPath(candidate.path, request.path) && + Number.isSafeInteger(candidate.line) && + Number(candidate.line) >= 1 && + Number.isSafeInteger(candidate.column) && + Number(candidate.column) >= 1 && + typeof candidate.text === 'string' && + candidate.text.length <= BRIDGE_WORKSPACE_SEARCH_TEXT_MAX_LENGTH && + candidate.text.includes(request.query) + ); + }) + ); +} + export function isValidBridgeWorkspaceToolCapabilities( value: unknown, ): value is BridgeWorkspaceToolCapabilities { diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 302b290..2752674 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -7,7 +7,7 @@ import { } from './protocol.js'; import { EndpointRuntimeSupervisor } from './runtime.js'; import { signBridgeRequest } from './identity.js'; -import { isWorkspaceToolRequest } from './workspace.js'; +import { isWorkspaceToolRequest, WorkspaceToolError } from './workspace.js'; import type { BridgeAssignment, @@ -790,6 +790,10 @@ export class BridgeWorker { leaseToken: assignment.leaseToken, incarnationId: this.incarnationId, status: 'rejected', + ...(assignment.executionKind === 'workspace_tool' && + error instanceof WorkspaceToolError + ? { errorCode: error.code } + : {}), error: (error instanceof Error ? error.message diff --git a/packages/code/src/workspace-worker.test.ts b/packages/code/src/workspace-worker.test.ts index e206311..352cbe3 100644 --- a/packages/code/src/workspace-worker.test.ts +++ b/packages/code/src/workspace-worker.test.ts @@ -2,9 +2,65 @@ import assert from 'node:assert/strict'; import test from 'node:test'; import { BridgeWorker } from './worker.js'; +import { WorkspaceToolError } from './workspace.js'; const incarnationId = 'incarnation-00000001'; +test('worker preserves bounded workspace rejection codes', async () => { + let settlement: Record | undefined; + const workspaceCapabilities = { + protocolVersion: 1 as const, + operations: ['search_text' 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() { + throw new WorkspaceToolError( + 'Workspace search timed out', + 'SEARCH_TIMEOUT', + ); + }, + }, + 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-timeout', + workerId: 'vm-1', + incarnationId, + generation: 1, + 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(settlement?.status, 'rejected'); + assert.equal(settlement?.errorCode, 'SEARCH_TIMEOUT'); +}); + test('worker executes a workspace tool assignment locally without acquiring a sandbox', async () => { const requests: Array<{ url: string; init?: RequestInit }> = []; const workspaceRequests: object[] = []; diff --git a/packages/code/src/workspace.test.ts b/packages/code/src/workspace.test.ts index 0fd9ad2..e8ee53b 100644 --- a/packages/code/src/workspace.test.ts +++ b/packages/code/src/workspace.test.ts @@ -6,7 +6,11 @@ import { join } from 'node:path'; import test from 'node:test'; import { promisify } from 'node:util'; -import { LocalWorkspaceTools, WorkspaceToolError } from './workspace.js'; +import { + isWorkspaceToolResult, + LocalWorkspaceTools, + WorkspaceToolError, +} from './workspace.js'; const execFileAsync = promisify(execFile); @@ -67,7 +71,7 @@ test('rejects traversal outside a registered workspace without leaking its host }), (error: unknown) => { assert.ok(error instanceof Error); - assert.match(error.message, /invalid workspace path/i); + assert.match(error.message, /invalid workspace/i); assert.equal(error.message.includes(parent), false); return true; }, @@ -210,6 +214,35 @@ test('search does not read an explicitly targeted escaping symlink', async (t) = ); }); +test('search preserves an in-workspace symlink namespace in returned paths', async (t) => { + const root = await mkdtemp(join(tmpdir(), 'librechat-code-workspace-')); + t.after(() => rm(root, { recursive: true, force: true })); + await mkdir(join(root, 'src')); + await writeFile(join(root, 'src', 'app.ts'), 'const needle = true;'); + await symlink(join(root, 'src'), join(root, 'alias')); + const tools = await LocalWorkspaceTools.create({ + workspaces: [{ id: 'primary', root }], + }); + + const result = await tools.execute({ + protocolVersion: 1, + operation: 'search_text', + workspaceId: 'primary', + query: 'needle', + path: 'alias', + }); + + if (result.operation !== 'search_text') assert.fail('expected search result'); + assert.deepEqual(result.matches, [ + { + path: 'alias/app.ts', + line: 1, + column: 7, + text: 'const needle = true;', + }, + ]); +}); + test('search returns a bounded match for a very long line', async (t) => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-workspace-')); t.after(() => rm(root, { recursive: true, force: true })); @@ -245,7 +278,8 @@ test('search rejects multiline literal queries', async (t) => { workspaceId: 'primary', query: 'first\nsecond', }), - /invalid workspace search/i, + (error: unknown) => + error instanceof WorkspaceToolError && error.code === 'INVALID_REQUEST', ); }); @@ -263,7 +297,8 @@ test('search rejects queries larger than its bounded preview', async (t) => { workspaceId: 'primary', query: 'a'.repeat(2001), }), - /invalid workspace search/i, + (error: unknown) => + error instanceof WorkspaceToolError && error.code === 'INVALID_REQUEST', ); }); @@ -513,6 +548,31 @@ test('bounds bytes read from a workspace file', async (t) => { ); }); +test('bounds workspace reads after UTF-16 decoding', async (t) => { + const root = await mkdtemp(join(tmpdir(), 'librechat-code-workspace-')); + t.after(() => rm(root, { recursive: true, force: true })); + const utf16 = Buffer.from('\u4e00'.repeat(400_000), 'utf16le'); + await writeFile( + join(root, 'large-utf16.txt'), + Buffer.concat([Buffer.from([0xff, 0xfe]), utf16]), + ); + const tools = await LocalWorkspaceTools.create({ + workspaces: [{ id: 'primary', root }], + }); + + await assert.rejects( + tools.execute({ + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'primary', + path: 'large-utf16.txt', + }), + (error: unknown) => + error instanceof WorkspaceToolError && + error.code === 'READ_LIMIT_EXCEEDED', + ); +}); + test('rejects ambiguous workspace registrations', async (t) => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-workspace-')); t.after(() => rm(root, { recursive: true, force: true })); @@ -575,3 +635,114 @@ test('does not start workspace I/O after its execution is aborted', async (t) => /workspace tool execution aborted/i, ); }); + +test('validates workspace results against the originating request', () => { + const request = { + protocolVersion: 1 as const, + operation: 'read_file' as const, + workspaceId: 'primary', + path: 'README.md', + }; + const result = { + protocolVersion: 1 as const, + operation: 'read_file' as const, + workspaceId: 'primary', + path: 'README.md', + content: '# LibreChat', + startLine: 1, + endLine: 1, + truncated: false, + }; + + assert.equal(isWorkspaceToolResult(request, result), true); + assert.equal( + isWorkspaceToolResult(request, { ...result, path: '/Users/operator/key' }), + false, + ); + assert.equal( + isWorkspaceToolResult(request, { + ...result, + root: '/Users/operator/private', + }), + false, + ); + assert.equal( + isWorkspaceToolResult(request, { + ...result, + truncated: true, + }), + false, + ); + assert.equal( + isWorkspaceToolResult(request, { + ...result, + content: '', + endLine: 0, + truncated: true, + nextStartLine: 1, + }), + false, + ); + assert.equal( + isWorkspaceToolResult(request, { + ...result, + nextStartLine: 2, + }), + false, + ); + assert.equal( + isWorkspaceToolResult( + { ...request, maxLines: 1 }, + { ...result, content: 'first\nsecond' }, + ), + false, + ); + assert.equal( + isWorkspaceToolResult(request, { + ...result, + matches: [ + { path: 'src/index.ts', line: 1, column: 1, text: 'unrelated' }, + ], + }), + false, + ); +}); + +test('validates search result paths against the requested scope', () => { + const request = { + protocolVersion: 1 as const, + operation: 'search_text' as const, + workspaceId: 'primary', + query: 'needle', + path: './src', + }; + const result = { + protocolVersion: 1 as const, + operation: 'search_text' as const, + workspaceId: 'primary', + matches: [ + { path: 'src/index.ts', line: 1, column: 1, text: 'needle' }, + ], + truncated: false, + }; + + assert.equal(isWorkspaceToolResult(request, result), true); + assert.equal( + isWorkspaceToolResult(request, { + ...result, + matches: [ + { path: 'src-old/index.ts', line: 1, column: 1, text: 'needle' }, + ], + }), + false, + ); + assert.equal( + isWorkspaceToolResult(request, { + ...result, + matches: [ + { path: 'secrets.env', line: 1, column: 1, text: 'needle' }, + ], + }), + false, + ); +}); diff --git a/packages/code/src/workspace.ts b/packages/code/src/workspace.ts index bde330a..3e6c27c 100644 --- a/packages/code/src/workspace.ts +++ b/packages/code/src/workspace.ts @@ -7,16 +7,39 @@ import type { FileHandle } from 'node:fs/promises'; import { BRIDGE_PROTOCOL_VERSION, - isValidBridgeWorkerId, + BRIDGE_WORKSPACE_READ_MAX_BYTES, + BRIDGE_WORKSPACE_READ_MAX_LINES, + BRIDGE_WORKSPACE_SEARCH_MAX_RESULTS, + BRIDGE_WORKSPACE_SEARCH_TEXT_MAX_LENGTH, isValidBridgeWorkspaceToolCapabilities, + isWorkspaceToolRequest, + isWorkspaceToolResult, } from './protocol.js'; import type { - BridgeProtocolVersion, BridgeWorkspaceDescriptor, BridgeWorkspaceToolCapabilities, + WorkspaceReadFileRequest, + WorkspaceReadFileResult, + WorkspaceSearchMatch, + WorkspaceSearchTextRequest, + WorkspaceSearchTextResult, + WorkspaceToolRequest, + WorkspaceToolErrorCode, + WorkspaceToolResult, } from './protocol.js'; +export { isWorkspaceToolRequest, isWorkspaceToolResult }; +export type { + WorkspaceReadFileRequest, + WorkspaceReadFileResult, + WorkspaceSearchMatch, + WorkspaceSearchTextRequest, + WorkspaceSearchTextResult, + WorkspaceToolRequest, + WorkspaceToolResult, +}; + export interface LocalWorkspaceConfig { id: string; name?: string; @@ -35,59 +58,6 @@ export interface WorkspaceToolExecutor { ): Promise; } -export interface WorkspaceReadFileRequest { - protocolVersion: BridgeProtocolVersion; - operation: 'read_file'; - workspaceId: string; - path: string; - startLine?: number; - maxLines?: number; -} - -export interface WorkspaceReadFileResult { - protocolVersion: BridgeProtocolVersion; - operation: 'read_file'; - workspaceId: string; - path: string; - content: string; - startLine: number; - endLine: number; - truncated: boolean; - nextStartLine?: number; -} - -export interface WorkspaceSearchTextRequest { - protocolVersion: BridgeProtocolVersion; - operation: 'search_text'; - workspaceId: string; - query: string; - path?: string; - maxResults?: number; -} - -export interface WorkspaceSearchMatch { - path: string; - line: number; - column: number; - text: string; -} - -export interface WorkspaceSearchTextResult { - protocolVersion: BridgeProtocolVersion; - operation: 'search_text'; - workspaceId: string; - matches: WorkspaceSearchMatch[]; - truncated: boolean; -} - -export type WorkspaceToolRequest = - WorkspaceReadFileRequest | WorkspaceSearchTextRequest; -export type WorkspaceToolResult = - WorkspaceReadFileResult | WorkspaceSearchTextResult; - -const MAX_READ_LINES = 500; -const MAX_READ_BYTES = 1024 * 1024; -const MAX_SEARCH_PREVIEW_LENGTH = 2000; const MAX_SEARCH_CANDIDATE_BYTES = 1024 * 1024; const MAX_SEARCH_CANDIDATES = 20_000; const SEARCH_TIMEOUT_MS = 10_000; @@ -140,60 +110,13 @@ function sliceWithoutSplittingSurrogates( export class WorkspaceToolError extends Error { constructor( message: string, - public readonly code: - | 'INVALID_PATH' - | 'INVALID_REQUEST' - | 'READ_LIMIT_EXCEEDED' - | 'REGISTRATION_INVALID' - | 'EXECUTION_ABORTED' - | 'SEARCH_TIMEOUT' - | 'SEARCH_UNAVAILABLE', + public readonly code: WorkspaceToolErrorCode, ) { super(message); this.name = 'WorkspaceToolError'; } } -export function isWorkspaceToolRequest( - value: unknown, -): value is WorkspaceToolRequest { - if (typeof value !== 'object' || value === null) return false; - const request = value as Record; - if ( - request.protocolVersion !== BRIDGE_PROTOCOL_VERSION || - typeof request.workspaceId !== 'string' || - !isValidBridgeWorkerId(request.workspaceId) - ) { - return false; - } - if (request.operation === 'read_file') { - return ( - typeof request.path === 'string' && - request.path.length > 0 && - request.path.length <= 4096 && - isUtf8ScalarString(request.path) && - (request.startLine === undefined || - Number.isSafeInteger(request.startLine)) && - (request.maxLines === undefined || Number.isSafeInteger(request.maxLines)) - ); - } - if (request.operation === 'search_text') { - return ( - typeof request.query === 'string' && - request.query.length > 0 && - request.query.length <= 4096 && - (request.path === undefined || - (typeof request.path === 'string' && - request.path.length > 0 && - request.path.length <= 4096 && - isUtf8ScalarString(request.path))) && - (request.maxResults === undefined || - Number.isSafeInteger(request.maxResults)) - ); - } - return false; -} - function isWithinRoot(root: string, candidate: string): boolean { const relativePath = relative(root, candidate); return !( @@ -243,15 +166,15 @@ async function readConfinedFileBuffer( ) { throw new Error('Invalid workspace path'); } - if (openedFile.size > MAX_READ_BYTES) { + if (openedFile.size > BRIDGE_WORKSPACE_READ_MAX_BYTES) { throw new WorkspaceToolError( 'Workspace file exceeds read limit', 'READ_LIMIT_EXCEEDED', ); } - const buffer = Buffer.allocUnsafe(MAX_READ_BYTES + 1); + const buffer = Buffer.allocUnsafe(BRIDGE_WORKSPACE_READ_MAX_BYTES + 1); let bytesRead = 0; - while (bytesRead <= MAX_READ_BYTES) { + while (bytesRead <= BRIDGE_WORKSPACE_READ_MAX_BYTES) { const result = await handle.read( buffer, bytesRead, @@ -261,7 +184,7 @@ async function readConfinedFileBuffer( if (result.bytesRead === 0) break; bytesRead += result.bytesRead; } - if (bytesRead > MAX_READ_BYTES) { + if (bytesRead > BRIDGE_WORKSPACE_READ_MAX_BYTES) { throw new WorkspaceToolError( 'Workspace file exceeds read limit', 'READ_LIMIT_EXCEEDED', @@ -280,7 +203,16 @@ async function readConfinedFile( root: string, requestedPath: string, ): Promise { - return decodeWorkspaceText(await readConfinedFileBuffer(root, requestedPath)); + const decoded = decodeWorkspaceText( + await readConfinedFileBuffer(root, requestedPath), + ); + if (Buffer.byteLength(decoded, 'utf8') > BRIDGE_WORKSPACE_READ_MAX_BYTES) { + throw new WorkspaceToolError( + 'Workspace file exceeds read limit', + 'READ_LIMIT_EXCEEDED', + ); + } + return decoded; } interface SearchCandidates { @@ -304,6 +236,8 @@ async function listSearchCandidates( '--files', '--no-config', '--no-follow', + '--path-separator', + '/', '--null', '--max-filesize', '1M', @@ -414,7 +348,7 @@ async function searchWorkspace( if ( !request.query || request.query.length > 4096 || - encodedQuery.length > MAX_SEARCH_PREVIEW_LENGTH || + encodedQuery.length > BRIDGE_WORKSPACE_SEARCH_TEXT_MAX_LENGTH || encodedQuery.toString('utf8') !== request.query || request.query.includes('\0') || request.query.includes('\n') || @@ -423,7 +357,11 @@ async function searchWorkspace( throw new WorkspaceToolError('Invalid workspace search', 'INVALID_REQUEST'); } const maxResults = request.maxResults ?? 50; - if (!Number.isSafeInteger(maxResults) || maxResults < 1 || maxResults > 200) { + if ( + !Number.isSafeInteger(maxResults) || + maxResults < 1 || + maxResults > BRIDGE_WORKSPACE_SEARCH_MAX_RESULTS + ) { throw new WorkspaceToolError('Invalid workspace search', 'INVALID_REQUEST'); } @@ -438,6 +376,12 @@ async function searchWorkspace( if (!isWithinRoot(root, canonicalTarget)) throw new WorkspaceToolError('Invalid workspace path', 'INVALID_PATH'); const canonicalSearchPath = relative(root, canonicalTarget) || '.'; + const portableCanonicalSearchPath = canonicalSearchPath.split(sep).join('/'); + const normalizedRequestedResultPath = request.path + ?.split('/') + .filter((segment) => segment.length > 0 && segment !== '.') + .join('/'); + const requestedResultPath = normalizedRequestedResultPath || undefined; const deadline = Date.now() + SEARCH_TIMEOUT_MS; const candidates = await listSearchCandidates( @@ -461,9 +405,16 @@ async function searchWorkspace( 'SEARCH_TIMEOUT', ); } - const path = candidate.startsWith(`.${sep}`) - ? candidate.slice(2) - : candidate; + const path = candidate.startsWith('./') ? candidate.slice(2) : candidate; + const resultPath = + requestedResultPath == null + ? path + : portableCanonicalSearchPath === '.' + ? `${requestedResultPath}/${path}` + : path === portableCanonicalSearchPath || + path.startsWith(`${portableCanonicalSearchPath}/`) + ? `${requestedResultPath}${path.slice(portableCanonicalSearchPath.length)}` + : path; let content: Buffer; try { content = await readConfinedFileBuffer(root, path); @@ -511,19 +462,24 @@ async function searchWorkspace( 0, column - Math.floor( - (MAX_SEARCH_PREVIEW_LENGTH - request.query.length) / 2, + (BRIDGE_WORKSPACE_SEARCH_TEXT_MAX_LENGTH - + request.query.length) / + 2, ), ), - Math.max(0, line.length - MAX_SEARCH_PREVIEW_LENGTH), + Math.max( + 0, + line.length - BRIDGE_WORKSPACE_SEARCH_TEXT_MAX_LENGTH, + ), ); matches.push({ - path, + path: resultPath, line: lineNumber, column: column + 1, text: sliceWithoutSplittingSurrogates( line, previewStart, - MAX_SEARCH_PREVIEW_LENGTH, + BRIDGE_WORKSPACE_SEARCH_TEXT_MAX_LENGTH, ), }); } @@ -626,7 +582,7 @@ export class LocalWorkspaceTools implements WorkspaceToolExecutor { startLine < 1 || !Number.isSafeInteger(maxLines) || maxLines < 1 || - maxLines > MAX_READ_LINES + maxLines > BRIDGE_WORKSPACE_READ_MAX_LINES ) { throw new WorkspaceToolError('Invalid workspace read', 'INVALID_REQUEST'); } diff --git a/service/src/api-server.ts b/service/src/api-server.ts index 1e4634a..4a09857 100644 --- a/service/src/api-server.ts +++ b/service/src/api-server.ts @@ -19,6 +19,7 @@ import { localAuth } from './auth/local'; import serviceRouter from './service/router'; import programmaticRouter from './service/programmatic-router'; import bridgeRouter from './bridge'; +import workspaceToolsRouter from './workspace-tools'; import { connection } from './queue'; import { metricsHandler } from './metrics'; import { httpMetricsMiddleware } from './middleware/httpMetrics'; @@ -55,6 +56,7 @@ app.get('/v1/health', async (_, res) => { v1.use('/bridge', bridgeRouter); v1.use(isLocalMode ? localAuth : apiKeyAuth); +v1.use(workspaceToolsRouter); v1.use(serviceRouter); v1.use(programmaticRouter); diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index 996dd16..0905993 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -10,6 +10,7 @@ import { BRIDGE_PROTOCOL_VERSION, isValidBridgeWorkerCapabilities, isValidBridgeWorkerId, + isWorkspaceToolErrorCode, } from '../../../packages/code/src/protocol'; import { BridgePairingError, RedisBridgePairingStore } from './pairing'; import { BridgeStoreError, RedisBridgeStore } from './store'; @@ -113,7 +114,12 @@ function isSettlement(value: unknown): value is CodeBridgeSettlement { return false; } if (value.status === 'rejected') { - return typeof value.error === 'string' && value.error.length <= 4096; + return ( + typeof value.error === 'string' && + value.error.length <= 4096 && + (value.errorCode === undefined || + isWorkspaceToolErrorCode(value.errorCode)) + ); } return ( value.status === 'fulfilled' && diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index afc8a72..649224c 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -6,9 +6,15 @@ import type { BridgeAssignment, BridgeSettlement, BridgeWorkerRegistration, + WorkspaceToolRequest, + WorkspaceToolResult, } from '../../../packages/code/src/protocol'; -import { BRIDGE_PROTOCOL_VERSION } from '../../../packages/code/src/protocol'; +import { + BRIDGE_PROTOCOL_VERSION, + isWorkspaceToolRequest, + isWorkspaceToolResult, +} from '../../../packages/code/src/protocol'; import type { BridgeWorkerBinding } from './pairing'; const PREFIX = 'codeapi:bridge:v1'; @@ -24,6 +30,10 @@ export type CodeBridgeSettlement = BridgeSettlement< run?: t.ExecuteResponse['run']; } >; +export type CodeBridgeWorkspaceSettlement = BridgeSettlement; +type AnyCodeBridgeSettlement = + | CodeBridgeSettlement + | CodeBridgeWorkspaceSettlement; export class BridgeStoreError extends Error { constructor( @@ -37,7 +47,9 @@ export class BridgeStoreError extends Error { | 'WORKER_FENCED' | 'WORKER_QUARANTINED' | 'WORKSPACE_QUARANTINED' - | 'WORKER_MISMATCH', + | 'WORKER_MISMATCH' + | 'ASSIGNMENT_INVALID' + | 'RESULT_INVALID', message: string, ) { super(message); @@ -56,6 +68,20 @@ export interface RegisteredBridgeWorker extends BridgeWorkerRegistration { identityId?: string; } +function supportsWorkspaceTool( + registration: RegisteredBridgeWorker, + request: WorkspaceToolRequest, +): boolean { + const capabilities = registration.capabilities.workspaceTools; + return ( + capabilities != null && + capabilities.operations.includes(request.operation) && + capabilities.workspaces.some( + (workspace) => workspace.id === request.workspaceId, + ) + ); +} + function workerKey(workerId: string): string { return `${PREFIX}:worker:${encodeURIComponent(workerId)}`; } @@ -442,12 +468,45 @@ export class RedisBridgeStore { } } + async dispatchWorkspaceTool(args: { + workerId: string; + tenantId?: string; + requireTenantBinding?: boolean; + request: WorkspaceToolRequest; + deadlineAtMs: number; + signal: AbortSignal; + }): Promise { + if (!isWorkspaceToolRequest(args.request)) { + throw new BridgeStoreError( + 'ASSIGNMENT_INVALID', + 'Invalid workspace tool request', + ); + } + const settlement = (await this.dispatch({ + ...args, + body: {} as t.PayloadBody, + headers: {}, + workspaceRequest: args.request, + })) as unknown as CodeBridgeWorkspaceSettlement; + if ( + settlement.status === 'fulfilled' && + !isWorkspaceToolResult(args.request, settlement.result) + ) { + throw new BridgeStoreError( + 'RESULT_INVALID', + 'Bridge worker returned an invalid workspace tool result', + ); + } + return settlement; + } + async dispatch(args: { workerId: string; tenantId?: string; requireTenantBinding?: boolean; body: t.PayloadBody; headers: Record; + workspaceRequest?: WorkspaceToolRequest; runtimeSessionId?: string; deadlineAtMs: number; signal: AbortSignal; @@ -489,6 +548,15 @@ export class RedisBridgeStore { `Bridge worker ${args.workerId} does not provide a stateful workspace`, ); } + if ( + args.workspaceRequest != null && + !supportsWorkspaceTool(registration, args.workspaceRequest) + ) { + throw new BridgeStoreError( + 'WORKER_MISMATCH', + `Bridge worker ${args.workerId} does not advertise the requested workspace tool`, + ); + } if ( args.runtimeSessionId !== undefined && (await this.dispatchCommand( @@ -549,10 +617,17 @@ export class RedisBridgeStore { : {}), expiresAt: new Date(args.deadlineAtMs).toISOString(), runtimeSessionId: args.runtimeSessionId, - request: { - body: args.body, - headers: args.headers, - }, + ...(args.workspaceRequest != null + ? { + executionKind: 'workspace_tool' as const, + request: args.workspaceRequest, + } + : { + request: { + body: args.body, + headers: args.headers, + }, + }), }; let queued = false; for (let attempt = 0; attempt < 8 && !queued; attempt += 1) { @@ -589,6 +664,15 @@ export class RedisBridgeStore { `Bridge worker ${args.workerId} does not provide a stateful workspace`, ); } + if ( + args.workspaceRequest != null && + !supportsWorkspaceTool(replacement.registration, args.workspaceRequest) + ) { + throw new BridgeStoreError( + 'WORKER_MISMATCH', + `Bridge worker ${args.workerId} no longer advertises the requested workspace tool`, + ); + } registration = replacement.registration; readyToken = replacement.readyToken; } @@ -948,7 +1032,7 @@ export class RedisBridgeStore { async settle( workerId: string, assignmentId: string, - settlement: CodeBridgeSettlement, + settlement: AnyCodeBridgeSettlement, signal?: AbortSignal, identityId?: string, ): Promise { @@ -1430,7 +1514,7 @@ export class RedisBridgeStore { private async commitPendingWorkspace( assignment: StoredAssignment, - settlement: CodeBridgeSettlement, + settlement: AnyCodeBridgeSettlement, deadlineAtMs: number, signal: AbortSignal, ): Promise { diff --git a/service/src/bridge/workspace-store.test.ts b/service/src/bridge/workspace-store.test.ts new file mode 100644 index 0000000..182a2da --- /dev/null +++ b/service/src/bridge/workspace-store.test.ts @@ -0,0 +1,159 @@ +import { afterEach, expect, test } from 'bun:test'; +import RedisMock from 'ioredis-mock'; + +import type Redis from 'ioredis'; + +import { BRIDGE_PROTOCOL_VERSION } from '../../../packages/code/src/protocol'; +import { RedisBridgeStore } from './store'; + +const redis = new RedisMock() as unknown as Redis; +const store = new RedisBridgeStore(redis); +const incarnationId = 'incarnation-00000001'; + +afterEach(async () => { + await redis.flushall(); +}); + +test('dispatches a workspace tool only to a worker advertising its workspace and operation', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'workspace-worker', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + workspaceTools: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operations: ['read_file'], + workspaces: [{ id: 'primary', name: 'LibreChat' }], + }, + }, + }); + const request = { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operation: 'read_file' as const, + workspaceId: 'primary', + path: 'README.md', + }; + const completion = store.dispatchWorkspaceTool({ + workerId: 'workspace-worker', + tenantId: 'tenant-1', + request, + deadlineAtMs: Date.now() + 5_000, + signal: new AbortController().signal, + }); + + const assignment = await store.lease('workspace-worker', incarnationId, 1_000); + expect(assignment).toMatchObject({ + executionKind: 'workspace_tool', + request, + }); + await store.settle('workspace-worker', assignment?.assignmentId ?? '', { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: assignment?.generation ?? 0, + leaseToken: assignment?.leaseToken ?? '', + incarnationId, + status: 'fulfilled', + result: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operation: 'read_file', + workspaceId: 'primary', + path: 'README.md', + content: '# LibreChat', + startLine: 1, + endLine: 1, + truncated: false, + }, + }); + + await expect(completion).resolves.toMatchObject({ + status: 'fulfilled', + result: { content: '# LibreChat' }, + }); +}); + +test('rejects a workspace tool that the selected worker did not advertise', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'workspace-worker', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + workspaceTools: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operations: ['read_file'], + workspaces: [{ id: 'primary' }], + }, + }, + }); + + await expect( + store.dispatchWorkspaceTool({ + workerId: 'workspace-worker', + request: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operation: 'search_text', + workspaceId: 'primary', + query: 'needle', + }, + 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('rejects a fulfilled workspace settlement that violates the result contract', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'workspace-worker', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + workspaceTools: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operations: ['read_file'], + workspaces: [{ id: 'primary' }], + }, + }, + }); + const completion = store.dispatchWorkspaceTool({ + workerId: 'workspace-worker', + request: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operation: 'read_file', + workspaceId: 'primary', + path: 'README.md', + }, + deadlineAtMs: Date.now() + 5_000, + signal: new AbortController().signal, + }); + const assignment = await store.lease('workspace-worker', incarnationId, 1_000); + await store.settle('workspace-worker', assignment?.assignmentId ?? '', { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: assignment?.generation ?? 0, + leaseToken: assignment?.leaseToken ?? '', + incarnationId, + status: 'fulfilled', + result: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operation: 'read_file', + workspaceId: 'primary', + path: 'README.md', + content: 'safe', + startLine: 1, + endLine: 1, + truncated: false, + root: '/Users/operator/private', + } as never, + }); + + await expect(completion).rejects.toMatchObject({ + code: 'RESULT_INVALID', + }); +}); diff --git a/service/src/local-api.ts b/service/src/local-api.ts index df35cb5..871bc35 100644 --- a/service/src/local-api.ts +++ b/service/src/local-api.ts @@ -11,6 +11,7 @@ import express, { json, Router } from 'express'; import serviceRouter from './service/router'; import programmaticRouter from './service/programmatic-router'; import bridgeRouter from './bridge'; +import workspaceToolsRouter from './workspace-tools'; import { requestErrorLogger, requestNotFoundLogger } from './middleware/request-error-logger'; import { executionProfileMiddleware } from './middleware/execution-profile'; import { localAuth } from './auth/local'; @@ -52,6 +53,7 @@ app.get('/v1/health', async (_, res) => { v1.use('/bridge', bridgeRouter); v1.use(localAuth); +v1.use(workspaceToolsRouter); v1.use(serviceRouter); v1.use(programmaticRouter); app.use('/v1', v1); diff --git a/service/src/service-api.ts b/service/src/service-api.ts index 79db08d..b5b7d52 100644 --- a/service/src/service-api.ts +++ b/service/src/service-api.ts @@ -6,6 +6,7 @@ import { executionProfileMiddleware } from './middleware/execution-profile'; import serviceRouter from './service/router'; import programmaticRouter from './service/programmatic-router'; import bridgeRouter from './bridge'; +import workspaceToolsRouter from './workspace-tools'; import { connection } from './queue'; import { env } from './config'; import logger from './logger'; @@ -32,6 +33,7 @@ app.get('/v1/health', async (_, res) => { v1.use('/bridge', bridgeRouter); v1.use(apiKeyAuth); +v1.use(workspaceToolsRouter); v1.use(serviceRouter); v1.use(programmaticRouter); diff --git a/service/src/workspace-tools/index.ts b/service/src/workspace-tools/index.ts new file mode 100644 index 0000000..6f105f1 --- /dev/null +++ b/service/src/workspace-tools/index.ts @@ -0,0 +1,20 @@ +import { Router } from 'express'; + +import { bridgeStore } from '../bridge'; +import { env } from '../config'; +import { executionLimiter } from '../middleware/limits'; +import { createWorkspaceToolsRouter } from './router'; + +const router = Router(); +router.use('/workspace-tools/execute', executionLimiter); +router.use( + createWorkspaceToolsRouter({ + store: bridgeStore, + backend: env.SANDBOX_BACKEND, + configuredWorkerId: env.BRIDGE_WORKER_ID, + dynamicWorkers: env.BRIDGE_DYNAMIC_WORKERS, + timeoutMs: env.JOB_TIMEOUT, + }), +); + +export default router; diff --git a/service/src/workspace-tools/router.test.ts b/service/src/workspace-tools/router.test.ts new file mode 100644 index 0000000..aa9c210 --- /dev/null +++ b/service/src/workspace-tools/router.test.ts @@ -0,0 +1,203 @@ +import { createServer } from 'node:http'; +import type { Server } from 'node:http'; + +import { afterEach, expect, test } from 'bun:test'; +import express, { json } from 'express'; + +import { applyPrincipal } from '../auth/principal'; +import { BridgeStoreError } from '../bridge/store'; +import { bridgeStoreStatus, createWorkspaceToolsRouter } from './router'; + +let server: Server | undefined; + +afterEach(() => { + server?.close(); + server = undefined; +}); + +test('maps invalid worker results to an upstream failure', () => { + expect(bridgeStoreStatus(new BridgeStoreError('RESULT_INVALID', 'invalid worker result'))).toBe(502); +}); + +test('rejects new workspace dispatches while the service is shutting down', async () => { + let dispatched = false; + const app = express(); + app.use(json()); + app.use((req, _res, next) => { + applyPrincipal(req, { + userId: 'user-1', + tenantId: 'tenant-1', + principalSource: 'librechat_jwt', + codeWorkerId: 'user-worker', + }); + next(); + }); + app.use( + createWorkspaceToolsRouter({ + backend: 'remote-bridge', + configuredWorkerId: 'shared-worker', + dynamicWorkers: true, + isShuttingDown: () => true, + store: { + async dispatchWorkspaceTool() { + dispatched = true; + throw new Error('must not dispatch'); + }, + }, + }), + ); + server = createServer(app); + await new Promise((resolve) => server?.listen(0, '127.0.0.1', resolve)); + const address = server.address(); + if (address == null || typeof address === 'string') { + throw new Error('Expected TCP listener'); + } + + const response = await fetch(`http://127.0.0.1:${address.port}/workspace-tools/execute`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'primary', + path: 'README.md', + }), + }); + + expect(response.status).toBe(503); + expect(dispatched).toBe(false); +}); + +test.each([ + ['SEARCH_TIMEOUT', 504], + ['SEARCH_UNAVAILABLE', 503], +] as const)('maps worker %s rejections to HTTP %i', async (errorCode, expectedStatus) => { + const app = express(); + app.use(json()); + app.use((req, _res, next) => { + applyPrincipal(req, { + userId: 'user-1', + tenantId: 'tenant-1', + principalSource: 'librechat_jwt', + codeWorkerId: 'user-worker', + }); + next(); + }); + app.use( + createWorkspaceToolsRouter({ + backend: 'remote-bridge', + configuredWorkerId: 'shared-worker', + dynamicWorkers: true, + store: { + async dispatchWorkspaceTool() { + return { + protocolVersion: 1, + generation: 1, + leaseToken: 'lease-token', + incarnationId: 'incarnation-1', + status: 'rejected', + error: 'search failed', + errorCode, + }; + }, + }, + }), + ); + server = createServer(app); + await new Promise((resolve) => server?.listen(0, '127.0.0.1', resolve)); + const address = server.address(); + if (address == null || typeof address === 'string') { + throw new Error('Expected TCP listener'); + } + + const response = await fetch(`http://127.0.0.1:${address.port}/workspace-tools/execute`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + protocolVersion: 1, + operation: 'search_text', + workspaceId: 'primary', + query: 'needle', + }), + }); + + expect(response.status).toBe(expectedStatus); + await expect(response.json()).resolves.toMatchObject({ code: errorCode }); +}); + +test('dispatches an authenticated workspace tool request to the principal-bound worker', async () => { + let dispatchArgs: Record | undefined; + const app = express(); + app.use(json()); + app.use((req, _res, next) => { + applyPrincipal(req, { + userId: 'user-1', + tenantId: 'tenant-1', + principalSource: 'librechat_jwt', + codeWorkerId: 'user-worker', + }); + next(); + }); + app.use( + createWorkspaceToolsRouter({ + backend: 'remote-bridge', + configuredWorkerId: 'shared-worker', + dynamicWorkers: true, + store: { + async dispatchWorkspaceTool(args) { + dispatchArgs = args as unknown as Record; + return { + protocolVersion: 1, + generation: 1, + leaseToken: 'lease-token', + incarnationId: 'incarnation-1', + status: 'fulfilled', + result: { + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'primary', + path: 'README.md', + content: '# LibreChat', + startLine: 1, + endLine: 1, + truncated: false, + }, + }; + }, + }, + }), + ); + server = createServer(app); + await new Promise((resolve) => server?.listen(0, '127.0.0.1', resolve)); + const address = server.address(); + if (address == null || typeof address === 'string') { + throw new Error('Expected TCP listener'); + } + + const request = { + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'primary', + path: 'README.md', + }; + const response = await fetch(`http://127.0.0.1:${address.port}/workspace-tools/execute`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-LibreChat-Code-Worker-ID': 'user-worker', + }, + body: JSON.stringify(request), + }); + + expect(response.status).toBe(200); + await expect(response.json()).resolves.toMatchObject({ + operation: 'read_file', + content: '# LibreChat', + }); + expect(dispatchArgs).toMatchObject({ + workerId: 'user-worker', + tenantId: 'tenant-1', + requireTenantBinding: true, + request, + }); +}); diff --git a/service/src/workspace-tools/router.ts b/service/src/workspace-tools/router.ts new file mode 100644 index 0000000..13a4bd9 --- /dev/null +++ b/service/src/workspace-tools/router.ts @@ -0,0 +1,130 @@ +import { Router } from 'express'; + +import type { RequestHandler, Response } from 'express'; +import type { AuthenticatedRequest } from '../types'; +import type { RedisBridgeStore } from '../bridge/store'; + +import { getPrincipalOrReject } from '../auth/principal'; +import { BridgeStoreError } from '../bridge/store'; +import { checkServiceShutDown } from '../lifecycle'; +import { isWorkspaceToolRequest } from '../../../packages/code/src/protocol'; +import { + CODEAPI_BRIDGE_WORKER_HEADER, + BridgeWorkerSelectionError, + resolveBridgeWorkerSelection, +} from '../bridge/selection'; + +interface WorkspaceToolsRouterOptions { + store: Pick; + backend: 'http' | 'lambda-microvm' | 'remote-bridge'; + configuredWorkerId: string; + dynamicWorkers: boolean; + timeoutMs?: number; + isShuttingDown?: () => boolean; +} + +function asyncRoute(handler: (req: AuthenticatedRequest, res: Response) => Promise): RequestHandler { + return (req, res, next) => { + void handler(req as AuthenticatedRequest, res).catch(next); + }; +} + +export function bridgeStoreStatus(error: BridgeStoreError): number { + if (error.code === 'WORKER_UNAUTHORIZED') return 403; + if (error.code === 'ASSIGNMENT_INVALID') return 400; + if (error.code === 'RESULT_INVALID') return 502; + if (error.code === 'ASSIGNMENT_EXPIRED') return 504; + if (error.code === 'WORKER_OFFLINE' || error.code === 'WORKER_BUSY') { + return 503; + } + return 409; +} + +export function createWorkspaceToolsRouter(options: WorkspaceToolsRouterOptions): Router { + const router = Router(); + + router.post( + '/workspace-tools/execute', + asyncRoute(async (req, res) => { + const principal = getPrincipalOrReject(req, res); + if (!principal) return; + if ((options.isShuttingDown ?? checkServiceShutDown)()) { + res.status(503).json({ error: 'Service is shutting down' }); + return; + } + if (!isWorkspaceToolRequest(req.body)) { + res.status(400).json({ + error: 'Invalid workspace tool request', + }); + return; + } + + let selection: { workerId: string; explicit: boolean } | undefined; + try { + selection = resolveBridgeWorkerSelection({ + backend: options.backend, + configuredWorkerId: options.configuredWorkerId, + dynamicWorkers: options.dynamicWorkers, + requestedWorkerId: req.header(CODEAPI_BRIDGE_WORKER_HEADER), + trustedWorkerId: principal.codeWorkerId, + }); + } catch (error) { + if (error instanceof BridgeWorkerSelectionError) { + res.status(error.status).json({ error: error.message }); + return; + } + throw error; + } + if (selection == null) { + res.status(503).json({ + error: 'Workspace tools require the remote-bridge backend', + }); + return; + } + + const controller = new AbortController(); + const abort = () => controller.abort(); + req.once('aborted', abort); + const abortClosedResponse = () => { + if (!res.writableEnded) abort(); + }; + res.once('close', abortClosedResponse); + try { + const settlement = await options.store.dispatchWorkspaceTool({ + workerId: selection.workerId, + tenantId: principal.tenantId, + requireTenantBinding: + selection.explicit && (options.dynamicWorkers || selection.workerId !== options.configuredWorkerId), + request: req.body, + deadlineAtMs: Date.now() + Math.max(1, options.timeoutMs ?? 30_000), + signal: controller.signal, + }); + if (settlement.status === 'rejected') { + let status = 422; + if (settlement.errorCode === 'SEARCH_TIMEOUT') status = 504; + if (settlement.errorCode === 'SEARCH_UNAVAILABLE') status = 503; + res.status(status).json({ + error: settlement.error, + code: settlement.errorCode ?? 'WORKSPACE_TOOL_REJECTED', + }); + return; + } + res.status(200).json(settlement.result); + } catch (error) { + if (error instanceof BridgeStoreError) { + res.status(bridgeStoreStatus(error)).json({ + error: error.message, + code: error.code, + }); + return; + } + throw error; + } finally { + req.removeListener('aborted', abort); + res.removeListener('close', abortClosedResponse); + } + }), + ); + + return router; +}