diff --git a/docker-compose.yaml b/docker-compose.yaml index db517b4c..a4a43cdd 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -15,6 +15,7 @@ services: - CODEAPI_BRIDGE_DYNAMIC_WORKERS=${CODEAPI_BRIDGE_DYNAMIC_WORKERS:-true} - CODEAPI_BRIDGE_MAX_WORKSPACE_LEASE_SLOTS=${CODEAPI_BRIDGE_MAX_WORKSPACE_LEASE_SLOTS:-1} - CODEAPI_BRIDGE_LANE_GIT=${CODEAPI_BRIDGE_LANE_GIT:-false} + - CODEAPI_BRIDGE_IMPORT_FILE_MAX_BYTES=${CODEAPI_BRIDGE_IMPORT_FILE_MAX_BYTES:-104857600} - CODEAPI_BRIDGE_WORKER_ID=${CODEAPI_BRIDGE_WORKER_ID:-} - CODEAPI_BRIDGE_RECOVERY_SERVER_ID=${CODEAPI_BRIDGE_RECOVERY_SERVER_ID:-} - CODEAPI_BRIDGE_ENROLLMENT_TTL_SECONDS=${CODEAPI_BRIDGE_ENROLLMENT_TTL_SECONDS:-0} diff --git a/docs/adr/0014-workspace-file-import.md b/docs/adr/0014-workspace-file-import.md new file mode 100644 index 00000000..e1aac4a7 --- /dev/null +++ b/docs/adr/0014-workspace-file-import.md @@ -0,0 +1,65 @@ +# 0014: Import chat uploads into attached workspaces + +Status: Accepted + +## Context + +A LibreChat chat can run on a personal machine through the remote bridge. Files +the user uploads in that chat stay in Code API storage: only programmatic Bash +stages them, per run, in a private scratch directory. `write_file` accepts +1 MiB of UTF-8 text, so binaries such as archives, images, and PDFs cannot reach +the workspace at all. + +## Decision + +Add one additive workspace operation, `import_file`, under protocol version 1. + +- **Request.** `{ path, source: { storage_session_id, id }, bytes, sha256, + overwrite? }` plus the usual workspace selectors. `sha256` is lowercase hex. +- **Ownership.** Code API, not the caller, decides whether `source` may be read. + It applies the `GET /download/:session_id/:fileId` rule for a user-kind + object: the live `session:` key must name the caller's storage key, and an + upload record must exist for that key, session, and object. +- **Grant.** Code API asks the egress gateway for a grant whose only input is + that object, with no upload capacity, two reads, and an expiry of the + request's admission plus execution budget plus 30 seconds. The relayed + request carries the grant and opaque handles in place of the raw ids. A + caller-supplied `egressGrant` is rejected. Synchronous requests revoke the + grant when they finish; durable requests let it expire. +- **Worker.** The worker streams the object from + `LIBRECHAT_CODE_FILE_RELAY_UPSTREAM`, the route programmatic inputs already + use, and holds no storage credential. It applies `write_file`'s path rules, + refuses `.git` segments and their filesystem aliases, refuses symlinked + parents, creates missing parents as 0755, enforces `bytes` while streaming + into a same-directory temporary file, verifies size and SHA-256, syncs, and + installs the file atomically with mode 0644. The size and hash are checked on + the synced temporary file itself. A create links exclusively; an overwrite + rechecks the inspected inode immediately before its rename, as `write_file` + does. +- **Idempotency.** Identical existing bytes return `unchanged: true` and + `bytesWritten: 0` without a transfer. Different bytes return `EDIT_CONFLICT` + unless `overwrite: true`. +- **Limits.** Code API (`CODEAPI_BRIDGE_IMPORT_FILE_MAX_BYTES`) and the worker + (`LIBRECHAT_CODE_IMPORT_FILE_MAX_BYTES`) each enforce a per-file ceiling, + 100 MiB by default, as `WRITE_LIMIT_EXCEEDED`. The worker also requires the + file plus 64 MiB to fit in the target filesystem's free space. +- **Policy.** `import_file` is a file write. It shares the mutation guard with + `write_file`, is never advertised for a read-only root, and a read-only root + refuses it with `WRITE_DISABLED`. +- **Negotiation.** Code API lists `import_file` in + `supportedWorkspaceToolOperations` only when `EGRESS_GATEWAY_URL` is set. A + worker advertises it by default whenever writes are enabled, a relay + upstream is configured, and the Code API listed it; operators opt out with + `--no-workspace-imports` or `LIBRECHAT_CODE_ALLOW_WORKSPACE_IMPORTS=false`. + LibreChat sends it only when advertised. LibreChat releases that predate + LibreChat-AI/LibreChat#17376 reject a worker status listing an unknown + operation, so LibreChat is updated first and the opt-out is the escape + hatch for workers serving an older one. + +## Consequences + +Uploads land in the workspace without a model round trip, and older Code APIs, +workers, and LibreChat releases keep their current behavior. Deployments +without an egress gateway do not offer the operation. Parent directories are +created only beneath a descriptor-held root, as for `write_file`; a legacy +pathname-only root reports a missing parent instead. diff --git a/docs/remote-bridge/README.md b/docs/remote-bridge/README.md index dc7cb98b..40465e00 100644 --- a/docs/remote-bridge/README.md +++ b/docs/remote-bridge/README.md @@ -172,6 +172,11 @@ whitespace-tolerant matching and multi-location replacement (see text, require an existing in-workspace parent directory, reject symlinks, and commit atomically. The worker capability is an enforcement boundary; LibreChat should still route every mutation through its configurable tool-approval hooks. +Writable workers with a file relay upstream also offer `import_file` by +default (opt out with `LIBRECHAT_CODE_ALLOW_WORKSPACE_IMPORTS=false`), which +copies a chat upload the caller owns into the workspace through a one-object +egress grant (ADR 0014). It is a file write: a worker without writes enabled +never advertises it and refuses it with `WRITE_DISABLED`. 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 diff --git a/packages/code/README.md b/packages/code/README.md index 0437ecfb..8f911996 100644 --- a/packages/code/README.md +++ b/packages/code/README.md @@ -630,6 +630,27 @@ with `--allow-workspace-writes` or `overwrite: false` to require an atomic create that returns `EDIT_CONFLICT` if the target already exists. Code API dispatches that mode only after the worker and server negotiate `create` in `writeFileModes`. + +`import_file` copies one stored chat upload into a writable workspace as a +binary-safe file. It is on by default whenever the worker can perform it: +writes are enabled and `LIBRECHAT_CODE_FILE_RELAY_UPSTREAM` names the +deployment's egress gateway. It is advertised only to a Code API that lists it +in `supportedWorkspaceToolOperations` (one that can mint egress grants). Opt out +with `--no-workspace-imports` or `LIBRECHAT_CODE_ALLOW_WORKSPACE_IMPORTS=false`. +Update LibreChat before workers: LibreChat releases that predate +LibreChat-AI/LibreChat#17376 reject a worker status that lists an operation +they do not know, so a worker serving one of those must opt out. Code API checks that the caller owns the object +and relays a short-lived grant scoped to that object alone; the worker holds no +storage credential. The worker applies `write_file`'s path rules, refuses any +`.git` segment (including case, Unicode-ignorable, trailing-dot and `GIT~1` +aliases) and symlinked parents, creates missing parents as 0755 under a held +root, streams into a same-directory temporary file with `bytes` as a hard cap, +verifies size and SHA-256, syncs, and installs the file atomically with mode +0644. Identical existing bytes return `unchanged: true` without a transfer; +different bytes return `EDIT_CONFLICT` unless `overwrite: true`. A failure +removes the temporary file and any parents it created. The per-file ceiling is +`LIBRECHAT_CODE_IMPORT_FILE_MAX_BYTES` (`--import-file-max-bytes`, default +100 MiB), and the filesystem must keep 64 MiB free after the file lands. `edit_file` accepts either the legacy `oldText`/`newText` pair or an ordered `edits` array; every exact replacement is validated before the updated file is installed as one atomic mutation. Code API dispatches the batch form only after diff --git a/packages/code/src/cli.ts b/packages/code/src/cli.ts index bee3697a..ae2fa3bc 100644 --- a/packages/code/src/cli.ts +++ b/packages/code/src/cli.ts @@ -39,6 +39,7 @@ import { } from './storage.js'; import { BridgeWorker } from './worker.js'; import { LocalWorkspaceTools, SandboxWorkspaceTools } from './workspace.js'; +import { createWorkspaceFileImporter } from './workspace-import.js'; import { DockerRuntimeSupervisor, EndpointRuntimeSupervisor, @@ -78,6 +79,7 @@ import type { RuntimeSupervisor } from './runtime.js'; import type { GitHubCredentialProvider } from './github.js'; import type { WorkspaceToolExecutor } from './workspace.js'; import { + BRIDGE_WORKSPACE_IMPORT_FILE_DEFAULT_MAX_BYTES, BRIDGE_WORKSPACE_NAME_MAX_LENGTH, BridgeProtocolError, isValidBridgeWorkerCapabilities, @@ -914,10 +916,25 @@ async function run( ]), ) : undefined; + // A chat upload reaches the workspace only through the same relay upstream + // programmatic inputs use; without one the worker never advertises import_file. + const workspaceImporter = createWorkspaceFileImporter({ + optOut: args.includes('--no-workspace-imports'), + enabled: process.env.LIBRECHAT_CODE_ALLOW_WORKSPACE_IMPORTS, + writable: allowWorkspaceWrites, + upstreamUrl: fileRelayUpstream, + maxBytes: positiveInteger( + 'LIBRECHAT_CODE_IMPORT_FILE_MAX_BYTES', + option(args, '--import-file-max-bytes') ?? + process.env.LIBRECHAT_CODE_IMPORT_FILE_MAX_BYTES, + BRIDGE_WORKSPACE_IMPORT_FILE_DEFAULT_MAX_BYTES, + ), + }); const localWorkspaceTools = workerDirectory ? await LocalWorkspaceTools.create({ workspaces: roots, repositoryInstructions: args.includes('--repository-instructions'), + importer: workspaceImporter, }) : undefined; let workspaceTools: WorkspaceToolExecutor | undefined = localWorkspaceTools; @@ -1298,6 +1315,7 @@ async function run( }, repositoryInstructions: args.includes('--repository-instructions'), writable: root.writable ?? false, + importer: workspaceImporter, }, ]), ), @@ -1333,6 +1351,7 @@ async function run( command: nativeOptionsForWorkspace(root.id), repositoryInstructions: args.includes('--repository-instructions'), writable: root.writable ?? false, + importer: workspaceImporter, }, ]), ), diff --git a/packages/code/src/linked-worktrees.ts b/packages/code/src/linked-worktrees.ts index 225abebd..609bfbd8 100644 --- a/packages/code/src/linked-worktrees.ts +++ b/packages/code/src/linked-worktrees.ts @@ -18,7 +18,7 @@ import type { WorkspaceToolRequest, WorkspaceToolResult, } from './protocol.js'; -import type { WorkspaceToolExecutor } from './workspace.js'; +import type { WorkspaceFileImporter, WorkspaceToolExecutor } from './workspace.js'; /** Linked worktrees are only admitted from this directory beneath a checkout. */ export const LINKED_WORKTREE_DIRECTORY = '.worktrees'; @@ -60,6 +60,7 @@ export interface LinkedWorktreeSource { command?: NativeProcessSandboxOptions; repositoryInstructions: boolean; writable: boolean; + importer?: WorkspaceFileImporter; } export interface VerifiedLinkedWorktree { @@ -324,6 +325,7 @@ export class LinkedWorktreeWorkspaceTools implements WorkspaceToolExecutor, Link fingerprint, value: LocalWorkspaceTools.create({ repositoryInstructions: source.repositoryInstructions, + importer: source.importer, workspaces: [ { id: internalId, identity: lane.identity, root: lane.root, writable: source.writable }, ], diff --git a/packages/code/src/native-programmatic.ts b/packages/code/src/native-programmatic.ts index 9b6eea9e..7a4852af 100644 --- a/packages/code/src/native-programmatic.ts +++ b/packages/code/src/native-programmatic.ts @@ -15,7 +15,11 @@ import { isSafePortableRelativePath, isSupportedBridgeArtifactName, } from './protocol.js'; -import { validateFileRelayUpstream } from './relay.js'; +import { + FILE_RELAY_GRANT_HEADER, + fileRelayObjectUrl, + validateFileRelayUpstream, +} from './relay.js'; import { WorkspaceToolError } from './workspace.js'; import type { @@ -25,7 +29,6 @@ import type { } from './protocol.js'; import type { NativeSrtWorkspaceCommandSandbox } from './native-sandbox.js'; -const EGRESS_GRANT_HEADER = 'X-CodeAPI-Egress-Grant'; const EXECUTION_MAIN_FILE = 'main.sh'; const EXECUTION_HISTORY_FILE = '_ptc_history.json'; const EXECUTION_CONTROL_FILE = '_ptc_pending_result.json'; @@ -327,12 +330,9 @@ export class NativeWorkspaceProgrammaticExecutor { const timer = setTimeout(() => controller.abort(), transferTimeoutMs); try { const response = await this.fetchImpl( - new URL( - `sessions/${encodeURIComponent(file.storage_session_id)}/objects/${encodeURIComponent(file.id)}`, - `${this.upstream.toString().replace(/\/+$/, '')}/`, - ), + fileRelayObjectUrl(this.upstream, file.storage_session_id, file.id), { - headers: { [EGRESS_GRANT_HEADER]: grant }, + headers: { [FILE_RELAY_GRANT_HEADER]: grant }, redirect: 'error', signal: controller.signal, }, @@ -704,14 +704,11 @@ export class NativeWorkspaceProgrammaticExecutor { let response: Response; try { response = await this.fetchImpl( - new URL( - `sessions/${encodeURIComponent(outputSessionId!)}/objects/${id}`, - `${this.upstream.toString().replace(/\/+$/, '')}/`, - ), + fileRelayObjectUrl(this.upstream, outputSessionId!, id), { method: 'PUT', headers: { - [EGRESS_GRANT_HEADER]: grant!, + [FILE_RELAY_GRANT_HEADER]: grant!, 'Content-Type': bridgeArtifactMediaType(name), 'Content-Length': String(bytes.byteLength), 'X-Original-Filename': encodeURIComponent(name), diff --git a/packages/code/src/protocol.ts b/packages/code/src/protocol.ts index 1451986d..0fcb6090 100644 --- a/packages/code/src/protocol.ts +++ b/packages/code/src/protocol.ts @@ -35,6 +35,12 @@ export function programmaticTransferReserveMs(jobTimeoutMs: number): number { export const BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILE_BYTES = 10 * 1024 * 1024; export const BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_HISTORY_BYTES = 40_000_000; export const BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_TOTAL_BYTES = 100 * 1024 * 1024; +/** Default per-file `import_file` ceiling; Code API and each worker may configure their own. */ +export const BRIDGE_WORKSPACE_IMPORT_FILE_DEFAULT_MAX_BYTES = 100 * 1024 * 1024; +/** Longest execution budget Code API grants one `import_file` transfer. */ +export const BRIDGE_WORKSPACE_IMPORT_FILE_TIMEOUT_MS = 5 * 60_000; +const BRIDGE_WORKSPACE_IMPORT_FILE_HANDLE_MAX_LENGTH = 32_768; +const BRIDGE_WORKSPACE_IMPORT_FILE_GRANT_MAX_LENGTH = 256 * 1024; /** How long Code API drains a clean rejection after Stop cancels a workspace mutation. */ export const BRIDGE_CANCELLED_WORKSPACE_SETTLEMENT_GRACE_MS = 5_000; @@ -313,7 +319,20 @@ export type BridgeWorkspaceToolOperation = | 'write_file' | 'preview_edit' | 'edit_file' - | 'execute_command'; + | 'execute_command' + | 'import_file'; + +/** Operations that may change workspace contents and therefore arm the durable mutation guard. */ +export function isWorkspaceMutationOperation( + operation: BridgeWorkspaceToolOperation, +): boolean { + return ( + operation === 'write_file' || + operation === 'edit_file' || + operation === 'execute_command' || + operation === 'import_file' + ); +} export type WorkspaceWriteFileMode = 'replace' | 'create'; export type WorkspaceEditFileMode = 'single' | 'batch'; @@ -674,6 +693,54 @@ export interface WorkspaceExecuteCommandResult { laneGit?: WorkspaceLaneGit; } +/** A stored object the caller owns, named exactly as Code API file storage names it. */ +export interface WorkspaceImportFileSource { + storage_session_id: string; + id: string; +} + +/** + * Copies one stored object into the workspace as a binary-safe file. The + * worker streams it through its file relay upstream, refuses `.git` paths and + * symlinked parents, verifies `bytes` and `sha256`, and installs it atomically + * with mode 0644. + */ +export interface WorkspaceImportFileRequest { + protocolVersion: BridgeProtocolVersion; + operation: 'import_file'; + workspaceId: string; + workspaceInstanceId?: string; + /** Linked worktree lane at `.worktrees/`; paths and cwd are relative to it. */ + worktree?: string; + path: string; + source: WorkspaceImportFileSource; + /** Exact object size, enforced as a hard cap while streaming. */ + bytes: number; + /** Lowercase hex SHA-256 of the object's bytes. */ + sha256: string; + /** False or omitted refuses to replace a file whose content differs. */ + overwrite?: boolean; + /** + * Set only by Code API when it relays the request to a worker: a short-lived + * read grant scoped to `source`, whose ids are then opaque grant handles. + * Callers never send it. + */ + egressGrant?: string; +} + +export interface WorkspaceImportFileResult { + protocolVersion: BridgeProtocolVersion; + operation: 'import_file'; + workspaceId: string; + path: string; + created: boolean; + /** True when the target already held these exact bytes; nothing was written. */ + unchanged?: boolean; + /** Bytes installed, or 0 when `unchanged`. */ + bytesWritten: number; + sha256: string; +} + export type WorkspaceToolRequest = | WorkspaceReadFileRequest | WorkspaceSearchTextRequest @@ -681,7 +748,8 @@ export type WorkspaceToolRequest = | WorkspaceWriteFileRequest | WorkspacePreviewEditRequest | WorkspaceEditFileRequest - | WorkspaceExecuteCommandRequest; + | WorkspaceExecuteCommandRequest + | WorkspaceImportFileRequest; export type WorkspaceToolResult = | WorkspaceReadFileResult | WorkspaceSearchTextResult @@ -689,7 +757,8 @@ export type WorkspaceToolResult = | WorkspaceWriteFileResult | WorkspacePreviewEditResult | WorkspaceEditFileResult - | WorkspaceExecuteCommandResult; + | WorkspaceExecuteCommandResult + | WorkspaceImportFileResult; const WORKSPACE_READ_REQUEST_KEYS = new Set([ 'instructionSha256', @@ -771,6 +840,30 @@ const WORKSPACE_COMMAND_REQUEST_KEYS = new Set([ 'timeoutMs', 'maxOutputBytes', ]); +const WORKSPACE_IMPORT_REQUEST_KEYS = new Set([ + 'protocolVersion', + 'operation', + 'workspaceId', + 'workspaceInstanceId', + 'worktree', + 'path', + 'source', + 'bytes', + 'sha256', + 'overwrite', + 'egressGrant', +]); +const WORKSPACE_IMPORT_SOURCE_KEYS = new Set(['storage_session_id', 'id']); +const WORKSPACE_IMPORT_RESULT_KEYS = new Set([ + 'protocolVersion', + 'operation', + 'workspaceId', + 'path', + 'created', + 'unchanged', + 'bytesWritten', + 'sha256', +]); const WORKSPACE_READ_RESULT_KEYS = new Set([ 'protocolVersion', 'operation', @@ -1275,6 +1368,94 @@ export function isSafePortableRelativePath(value: unknown): value is string { return value.split('/').every(segment => segment !== '..'); } +/** + * Whether a path segment names Git metadata on any host filesystem: case + * folding (macOS, Windows), ignorable code points (HFS+), trailing dots and + * spaces (Windows), NTFS stream syntax such as `.git::$INDEX_ALLOCATION`, and + * the NTFS short name all alias `.git`. + */ +function isDotGitPathSegment(segment: string): boolean { + const folded = segment + .split(':', 1)[0]! + .toLowerCase() + .replace(/\p{Default_Ignorable_Code_Point}/gu, '') + .replace(/[. ]+$/, ''); + return folded === '.git' || folded === 'git~1'; +} + +/** `write_file` path rules, plus no `.git` segment and a file-like final segment. */ +export function isWorkspaceImportFilePath(value: unknown): value is string { + if (!isSafePortableRelativePath(value)) return false; + const segments = value.split('/'); + const last = segments[segments.length - 1]; + return ( + last !== undefined && + last !== '' && + last !== '.' && + !segments.some(isDotGitPathSegment) + ); +} + +function isWorkspaceImportFileHandle(value: unknown): value is string { + return ( + typeof value === 'string' && + value.length > 0 && + value.length <= BRIDGE_WORKSPACE_IMPORT_FILE_HANDLE_MAX_LENGTH && + !/[\0\r\n]/.test(value) + ); +} + +function isWorkspaceImportFileRequestShape( + request: Record, +): boolean { + const source = request.source; + return ( + hasOnlyKeys(request, WORKSPACE_IMPORT_REQUEST_KEYS) && + isWorkspaceImportFilePath(request.path) && + typeof source === 'object' && + source !== null && + !Array.isArray(source) && + hasOnlyKeys(source as Record, WORKSPACE_IMPORT_SOURCE_KEYS) && + isWorkspaceImportFileHandle((source as Record).storage_session_id) && + isWorkspaceImportFileHandle((source as Record).id) && + Number.isSafeInteger(request.bytes) && + Number(request.bytes) >= 0 && + typeof request.sha256 === 'string' && + /^[a-f0-9]{64}$/.test(request.sha256) && + (request.overwrite === undefined || typeof request.overwrite === 'boolean') && + (request.egressGrant === undefined || + (typeof request.egressGrant === 'string' && + request.egressGrant.length > 0 && + request.egressGrant.length <= BRIDGE_WORKSPACE_IMPORT_FILE_GRANT_MAX_LENGTH && + !/[\0\r\n]/.test(request.egressGrant))) + ); +} + +/** + * The caller-facing `import_file` request LibreChat sends to Code API. A + * caller may never supply `egressGrant`; Code API mints it when relaying. + */ +export function isBridgeWorkspaceImportFileRequest( + value: unknown, +): value is WorkspaceImportFileRequest & { egressGrant?: never } { + return ( + isWorkspaceToolRequest(value) && + value.operation === 'import_file' && + value.egressGrant === undefined + ); +} + +/** The relayed `import_file` request a worker executes: it must carry Code API's grant. */ +export function isRelayedWorkspaceImportFileRequest( + value: unknown, +): value is WorkspaceImportFileRequest & { egressGrant: string } { + return ( + isWorkspaceToolRequest(value) && + value.operation === 'import_file' && + value.egressGrant !== undefined + ); +} + function normalizePortableRelativePath(value: string): string { return ( value @@ -1573,6 +1754,9 @@ export function isWorkspaceToolRequest( BRIDGE_WORKSPACE_COMMAND_MAX_OUTPUT_BYTES)) ); } + if (request.operation === 'import_file') { + return isWorkspaceImportFileRequestShape(request); + } return false; } @@ -1706,6 +1890,26 @@ export function isWorkspaceToolResult( ); } + if (request.operation === 'import_file') { + if ( + !hasOnlyKeys(result, WORKSPACE_IMPORT_RESULT_KEYS) || + result.path !== request.path || + result.sha256 !== request.sha256 || + typeof result.created !== 'boolean' || + (result.unchanged !== undefined && typeof result.unchanged !== 'boolean') || + !Number.isSafeInteger(result.bytesWritten) + ) { + return false; + } + if (result.unchanged === true) { + return result.created === false && result.bytesWritten === 0; + } + return ( + result.bytesWritten === request.bytes && + (request.overwrite === true || result.created === true) + ); + } + if (request.operation === 'edit_file') { const replacements = request.edits?.length ?? 1; return ( @@ -1808,7 +2012,7 @@ export function isValidBridgeWorkspaceToolCapabilities( capabilities.protocolVersion !== BRIDGE_PROTOCOL_VERSION || !Array.isArray(capabilities.operations) || capabilities.operations.length < 1 || - capabilities.operations.length > 7 || + capabilities.operations.length > 8 || !capabilities.operations.every( operation => operation === 'read_file' || @@ -1817,7 +2021,8 @@ export function isValidBridgeWorkspaceToolCapabilities( operation === 'write_file' || operation === 'preview_edit' || operation === 'edit_file' || - operation === 'execute_command', + operation === 'execute_command' || + operation === 'import_file', ) || new Set(capabilities.operations).size !== capabilities.operations.length || diff --git a/packages/code/src/relay.ts b/packages/code/src/relay.ts index 07aa8714..1cab6e65 100644 --- a/packages/code/src/relay.ts +++ b/packages/code/src/relay.ts @@ -34,6 +34,21 @@ const LOCAL_HTTP_HOSTS = new Set([ class RelayPayloadTooLargeError extends Error {} class UpstreamPayloadTooLargeError extends Error {} +/** Header carrying Code API's egress grant on every file relay object request. */ +export const FILE_RELAY_GRANT_HEADER = 'X-CodeAPI-Egress-Grant'; + +/** `/sessions//objects/`, the egress gateway's object route. */ +export function fileRelayObjectUrl( + upstream: URL, + storageSessionId: string, + objectId: string, +): URL { + return new URL( + `sessions/${encodeURIComponent(storageSessionId)}/objects/${encodeURIComponent(objectId)}`, + `${upstream.toString().replace(/\/+$/, '')}/`, + ); +} + export function validateFileRelayUpstream(value: string): URL { const upstream = new URL(value); if ( diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 8978d4cd..5bfee37d 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -7,6 +7,7 @@ import { bridgeWorkerPath, isBridgeWorkspaceProgrammaticRequest, isValidLinkedWorktreeName, + isWorkspaceMutationOperation, isWorkspaceToolResult, workspaceIsolationKey, workspaceIsolationKeysConflict, @@ -564,10 +565,7 @@ export class BridgeWorker { } if ( options.capabilities.workspaceTools?.operations.some( - (operation) => - operation === 'write_file' || - operation === 'edit_file' || - operation === 'execute_command', + isWorkspaceMutationOperation, ) === true && options.workspaceMutationQuarantine == null && options.workspaceQuarantines == null && @@ -1885,10 +1883,9 @@ export class BridgeWorker { 'Workspace listing feature is not advertised', ); } - const isMutation = - workspaceRequest.operation === 'write_file' || - workspaceRequest.operation === 'edit_file' || - workspaceRequest.operation === 'execute_command'; + const isMutation = isWorkspaceMutationOperation( + workspaceRequest.operation, + ); if (isMutation) { this.mutationGuardArmed = true; try { @@ -2566,9 +2563,7 @@ export class BridgeWorker { (assignment.executionKind === 'workspace_programmatic' || (assignment.executionKind === 'workspace_tool' && isWorkspaceToolRequest(assignment.request) && - (assignment.request.operation === 'write_file' || - assignment.request.operation === 'edit_file' || - assignment.request.operation === 'execute_command'))); + isWorkspaceMutationOperation(assignment.request.operation))); if (signal?.aborted === true) { if (assignment.runtimeSessionId != null || fulfilledWorkspaceMutation) { throw await this.quarantineWorkspace( diff --git a/packages/code/src/workspace-import.test.ts b/packages/code/src/workspace-import.test.ts new file mode 100644 index 00000000..18d37b86 --- /dev/null +++ b/packages/code/src/workspace-import.test.ts @@ -0,0 +1,565 @@ +import assert from 'node:assert/strict'; +import { createHash, randomBytes } from 'node:crypto'; +import { createServer } from 'node:http'; +import { chmod, lstat, mkdir, mkdtemp, readdir, readFile, realpath, rename, rm, stat, symlink, writeFile } from 'node:fs/promises'; +import { setTimeout as delay } from 'node:timers/promises'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import test from 'node:test'; + +import type { IncomingMessage, ServerResponse } from 'node:http'; +import type { AddressInfo } from 'node:net'; +import type { TestContext } from 'node:test'; + +import { + isBridgeWorkspaceImportFileRequest, + isRelayedWorkspaceImportFileRequest, + isValidBridgeWorkspaceToolCapabilities, + isWorkspaceImportFilePath, + isWorkspaceToolRequest, + isWorkspaceToolResult, +} from './protocol.js'; +import { captureWorkspaceRootIdentity } from './root-identity.js'; +import { BridgeWorker } from './worker.js'; +import { createWorkspaceFileImporter, HttpWorkspaceFileImporter } from './workspace-import.js'; +import { LocalWorkspaceTools, WorkspaceToolError } from './workspace.js'; + +import type { WorkspaceImportFileRequest, WorkspaceToolRequest } from './protocol.js'; +import type { HttpWorkspaceFileImporterOptions } from './workspace-import.js'; + +const GRANT = 'ceg1.grant-for-one-object'; +const SESSION = 'session-handle'; +const OBJECT = 'object-handle'; +const incarnationId = 'incarnation-00000001'; + +const sha256 = (bytes: Uint8Array): string => createHash('sha256').update(bytes).digest('hex'); + +function importRequest( + bytes: Buffer, + overrides: Partial = {}, +): WorkspaceImportFileRequest { + return { + protocolVersion: 1, + operation: 'import_file', + workspaceId: 'primary', + path: 'uploads/archive.zip', + source: { storage_session_id: SESSION, id: OBJECT }, + bytes: bytes.byteLength, + sha256: sha256(bytes), + egressGrant: GRANT, + ...overrides, + }; +} + +interface EgressStub { + url: string; + hits: number; + body: Buffer; + /** Replace the default response (exact body, declared length, valid grant). */ + respond?: (req: IncomingMessage, res: ServerResponse) => void; +} + +/** A local egress gateway: serves one object to one grant, like the real object route. */ +async function egressStub(t: TestContext, body: Buffer): Promise { + const stub: EgressStub = { url: '', hits: 0, body }; + const server = createServer((req, res) => { + stub.hits += 1; + if (stub.respond) return stub.respond(req, res); + if ( + req.method !== 'GET' || + req.url !== `/v1/egress/sessions/${SESSION}/objects/${OBJECT}` || + req.headers['x-codeapi-egress-grant'] !== GRANT + ) { + res.writeHead(403).end(); + return; + } + res.writeHead(200, { 'content-length': String(stub.body.byteLength) }).end(stub.body); + }); + await new Promise((resolve) => server.listen(0, '127.0.0.1', resolve)); + t.after(() => new Promise((resolve) => { + server.closeAllConnections(); + server.close(() => resolve()); + })); + stub.url = `http://127.0.0.1:${(server.address() as AddressInfo).port}/v1/egress`; + return stub; +} + +async function workspace( + t: TestContext, + stub: EgressStub, + options: { writable?: boolean; held?: boolean; importer?: Partial } = {}, +) { + const root = await realpath(await mkdtemp(join(tmpdir(), 'librechat-code-import-'))); + t.after(() => rm(root, { recursive: true, force: true })); + const tools = await LocalWorkspaceTools.create({ + importer: new HttpWorkspaceFileImporter({ upstreamUrl: stub.url, ...options.importer }), + workspaces: [{ + id: 'primary', + root, + writable: options.writable ?? true, + ...(options.held === false ? {} : { identity: await captureWorkspaceRootIdentity(root) }), + }], + }); + return { root, tools }; +} + +async function rejection(promise: Promise): Promise { + try { + await promise; + } catch (error) { + assert.ok(error instanceof WorkspaceToolError, String(error)); + return error; + } + assert.fail('expected a WorkspaceToolError'); +} + +async function temporaryFiles(directory: string): Promise { + return (await readdir(directory).catch(() => [] as string[])).filter((name) => name.endsWith('.tmp')); +} + +test('import requests follow write_file path rules and refuse Git metadata', () => { + const bytes = Buffer.from('payload'); + assert.equal(isWorkspaceToolRequest(importRequest(bytes)), true); + for (const path of [ + '../escape.bin', '/etc/passwd', 'C:/x.bin', 'a\\b.bin', 'dir/', 'dir/.', '.git', '.git/config', + 'src/.git/hooks/pre-commit', '.GIT/config', '.Git./x', 'a/.git /x', 'GIT~1/config', '.g\u200cit/config', + '.git::$INDEX_ALLOCATION/config', '.GIT:stream', '.git. ::$INDEX_ALLOCATION/hooks/x', 'git~1::$INDEX_ALLOCATION/x', + ]) { + assert.equal(isWorkspaceImportFilePath(path), false, path); + assert.equal(isWorkspaceToolRequest(importRequest(bytes, { path })), false, path); + } + for (const path of ['archive.zip', 'uploads/photo.png', '.gitignore', '.github/workflow.yml', 'docs/.gitkeep', 'notes/10:30 standup.md']) { + assert.equal(isWorkspaceImportFilePath(path), true, path); + } + for (const invalid of [ + { sha256: 'A'.repeat(64) }, + { sha256: 'a'.repeat(63) }, + { bytes: -1 }, + { bytes: 1.5 }, + { source: { storage_session_id: SESSION } }, + { source: { storage_session_id: SESSION, id: OBJECT, name: 'x' } }, + { source: { storage_session_id: 'a\nb', id: OBJECT } }, + { overwrite: 'yes' }, + { egressGrant: '' }, + ] as Array>) { + assert.equal(isWorkspaceToolRequest({ ...importRequest(bytes), ...invalid }), false, JSON.stringify(invalid)); + } + assert.equal(isWorkspaceToolRequest({ ...importRequest(bytes), content: 'x' }), false); + + const { egressGrant: _grant, ...callerRequest } = importRequest(bytes); + assert.equal(isBridgeWorkspaceImportFileRequest(callerRequest), true); + assert.equal(isBridgeWorkspaceImportFileRequest(importRequest(bytes)), false, 'callers never supply the grant'); + assert.equal(isRelayedWorkspaceImportFileRequest(importRequest(bytes)), true); + assert.equal(isRelayedWorkspaceImportFileRequest(callerRequest), false); +}); + +test('import results must echo the request and account for every byte', () => { + const bytes = Buffer.from('payload'); + const request = importRequest(bytes); + const result = { + protocolVersion: 1, + operation: 'import_file', + workspaceId: 'primary', + path: request.path, + created: true, + bytesWritten: bytes.byteLength, + sha256: request.sha256, + }; + assert.equal(isWorkspaceToolResult(request, result), true); + assert.equal(isWorkspaceToolResult(request, { ...result, created: false, unchanged: true, bytesWritten: 0 }), true); + assert.equal(isWorkspaceToolResult(request, { ...result, created: false, unchanged: true }), false); + assert.equal(isWorkspaceToolResult(request, { ...result, unchanged: true, bytesWritten: 0 }), false); + assert.equal(isWorkspaceToolResult(request, { ...result, created: false }), false, 'a create cannot replace'); + assert.equal(isWorkspaceToolResult({ ...request, overwrite: true }, { ...result, created: false }), true); + assert.equal(isWorkspaceToolResult(request, { ...result, bytesWritten: 3 }), false); + assert.equal(isWorkspaceToolResult(request, { ...result, sha256: 'b'.repeat(64) }), false); + assert.equal(isWorkspaceToolResult(request, { ...result, path: 'other.zip' }), false); + assert.equal(isWorkspaceToolResult(request, { ...result, content: 'x' }), false); +}); + +test('a worker imports by default whenever it can, and an explicit opt-out removes it', async (t) => { + const upstreamUrl = 'https://code.example.test/egress'; + const enabled = createWorkspaceFileImporter({ optOut: false, writable: true, upstreamUrl }); + assert.ok(enabled instanceof HttpWorkspaceFileImporter, 'on by default'); + for (const value of ['true', ' TRUE ', '']) { + assert.ok(createWorkspaceFileImporter({ optOut: false, enabled: value, writable: true, upstreamUrl })); + } + assert.equal(createWorkspaceFileImporter({ optOut: true, writable: true, upstreamUrl }), undefined, '--no-workspace-imports'); + assert.equal(createWorkspaceFileImporter({ optOut: false, enabled: 'false', writable: true, upstreamUrl }), undefined); + assert.equal(createWorkspaceFileImporter({ optOut: false, enabled: 'False', writable: true, upstreamUrl }), undefined); + assert.equal(createWorkspaceFileImporter({ optOut: false, writable: true }), undefined, 'no relay upstream'); + assert.equal(createWorkspaceFileImporter({ optOut: false, writable: true, upstreamUrl: ' ' }), undefined); + assert.equal(createWorkspaceFileImporter({ optOut: false, writable: false, upstreamUrl }), undefined, 'writes disabled'); + assert.throws(() => createWorkspaceFileImporter({ optOut: false, enabled: 'yes', writable: true, upstreamUrl }), /must be true or false/); + assert.equal(createWorkspaceFileImporter({ optOut: false, writable: true, upstreamUrl, maxBytes: 4096 })?.maxBytes, 4096); + + const root = await realpath(await mkdtemp(join(tmpdir(), 'librechat-code-import-'))); + t.after(() => rm(root, { recursive: true, force: true })); + const advertised = await LocalWorkspaceTools.create({ importer: enabled, workspaces: [{ id: 'primary', root, writable: true }] }); + assert.ok(advertised.capabilities.operations.includes('import_file')); + const optedOut = await LocalWorkspaceTools.create({ + importer: createWorkspaceFileImporter({ optOut: true, writable: true, upstreamUrl }), + workspaces: [{ id: 'primary', root, writable: true }], + }); + assert.equal(optedOut.capabilities.operations.includes('import_file'), false); +}); + +test('import_file is advertised only beside writes, on writable roots, with an importer', async (t) => { + const stub = await egressStub(t, Buffer.from('x')); + const writable = await realpath(await mkdtemp(join(tmpdir(), 'librechat-code-import-'))); + const readOnly = await realpath(await mkdtemp(join(tmpdir(), 'librechat-code-import-'))); + t.after(() => Promise.all([rm(writable, { recursive: true }), rm(readOnly, { recursive: true })])); + const importer = new HttpWorkspaceFileImporter({ upstreamUrl: stub.url }); + + const mixed = await LocalWorkspaceTools.create({ + importer, + workspaces: [{ id: 'primary', root: writable, writable: true }, { id: 'docs', root: readOnly }], + }); + assert.ok(mixed.capabilities.operations.includes('import_file')); + assert.ok(mixed.capabilities.workspaces[0]?.operations?.includes('import_file')); + assert.equal(mixed.capabilities.workspaces[1]?.operations?.includes('import_file'), false); + assert.equal(isValidBridgeWorkspaceToolCapabilities(mixed.capabilities), true); + const denied = await rejection(mixed.execute({ ...importRequest(Buffer.from('x')), workspaceId: 'docs' })); + assert.equal(denied.code, 'WRITE_DISABLED'); + assert.equal(stub.hits, 0); + + const allReadOnly = await LocalWorkspaceTools.create({ importer, workspaces: [{ id: 'docs', root: readOnly }] }); + assert.equal(allReadOnly.capabilities.operations.includes('import_file'), false); + const withoutRelay = await LocalWorkspaceTools.create({ workspaces: [{ id: 'primary', root: writable, writable: true }] }); + assert.equal(withoutRelay.capabilities.operations.includes('import_file'), false); + assert.equal((await rejection(withoutRelay.execute(importRequest(Buffer.from('x'))))).code, 'WRITE_UNAVAILABLE'); +}); + +test('imports binary bytes atomically into new 0755 parents as a 0644 file, whatever the umask', async (t) => { + const bytes = Buffer.concat([Buffer.from([0, 0xff, 0x50, 0x4b, 3, 4, 0]), randomBytes(256 * 1024)]); + const stub = await egressStub(t, bytes); + const { root, tools } = await workspace(t, stub); + const request = importRequest(bytes, { path: 'uploads/2026/archive.zip' }); + const previousUmask = process.umask(0o077); + t.after(() => process.umask(previousUmask)); + + const result = await tools.execute(request); + assert.deepEqual(result, { + protocolVersion: 1, + operation: 'import_file', + workspaceId: 'primary', + path: request.path, + created: true, + bytesWritten: bytes.byteLength, + sha256: request.sha256, + }); + assert.equal(isWorkspaceToolResult(request, result), true); + const target = join(root, 'uploads', '2026', 'archive.zip'); + assert.deepEqual(await readFile(target), bytes); + assert.equal((await stat(target)).mode & 0o777, 0o644); + assert.equal((await stat(join(root, 'uploads'))).mode & 0o777, 0o755); + assert.equal((await stat(join(root, 'uploads', '2026'))).mode & 0o777, 0o755); + assert.deepEqual(await temporaryFiles(join(root, 'uploads', '2026')), []); + assert.equal(stub.hits, 1); +}); + +test('re-importing identical bytes is idempotent and never downloads again', async (t) => { + const bytes = randomBytes(4096); + const stub = await egressStub(t, bytes); + const { root, tools } = await workspace(t, stub); + const request = importRequest(bytes, { path: 'photo.png' }); + await tools.execute(request); + const installed = await stat(join(root, 'photo.png')); + + const again = await tools.execute(request); + assert.deepEqual(again, { + protocolVersion: 1, + operation: 'import_file', + workspaceId: 'primary', + path: 'photo.png', + created: false, + unchanged: true, + bytesWritten: 0, + sha256: request.sha256, + }); + assert.equal(isWorkspaceToolResult(request, again), true); + assert.equal(stub.hits, 1); + assert.equal((await stat(join(root, 'photo.png'))).ino, installed.ino); +}); + +test('different existing content conflicts unless overwrite is explicit', async (t) => { + const bytes = Buffer.from('new report'); + const stub = await egressStub(t, bytes); + const { root, tools } = await workspace(t, stub); + const target = join(root, 'report.pdf'); + await writeFile(target, 'old report'); + await chmod(target, 0o755); + + const conflict = await rejection(tools.execute(importRequest(bytes, { path: 'report.pdf' }))); + assert.equal(conflict.code, 'EDIT_CONFLICT'); + assert.equal(conflict.mutationMayHaveCommitted, false); + assert.equal(await readFile(target, 'utf8'), 'old report'); + assert.equal(stub.hits, 0, 'a conflict is decided before any transfer'); + + const replaced = await tools.execute(importRequest(bytes, { path: 'report.pdf', overwrite: true })); + assert.equal((replaced as { created: boolean }).created, false); + assert.equal(await readFile(target, 'utf8'), 'new report'); + assert.equal((await stat(target)).mode & 0o777, 0o644, 'an import is never executable'); +}); + +test('a hash or size mismatch installs nothing and removes the parents it created', async (t) => { + const bytes = Buffer.from('expected bytes'); + const stub = await egressStub(t, Buffer.from('tampered bytes')); + const { root, tools } = await workspace(t, stub); + + const mismatch = await rejection(tools.execute(importRequest(bytes, { path: 'fresh/dir/file.bin' }))); + assert.equal(mismatch.code, 'INVALID_REQUEST'); + assert.equal(mismatch.mutationMayHaveCommitted, false); + assert.match(mismatch.message, /SHA-256/); + assert.deepEqual(await readdir(root), [], 'created parents are rolled back'); + + stub.respond = (_req, res) => { + res.writeHead(200); + res.write(bytes); + res.end(Buffer.from('and then some')); + }; + const oversize = await rejection(tools.execute(importRequest(bytes, { path: 'file.bin' }))); + assert.equal(oversize.code, 'INVALID_REQUEST'); + assert.match(oversize.message, /larger than its declared size/); + + stub.respond = (_req, res) => res.writeHead(200, { 'content-length': '3' }).end('abc'); + const declared = await rejection(tools.execute(importRequest(bytes, { path: 'file.bin' }))); + assert.equal(declared.code, 'INVALID_REQUEST'); + + stub.respond = (_req, res) => res.writeHead(200).end(bytes.subarray(0, 4)); + const short = await rejection(tools.execute(importRequest(bytes, { path: 'file.bin' }))); + assert.match(short.message, /smaller than its declared size/); + assert.deepEqual(await readdir(root), []); +}); + +test('the staged file is verified after it is synced, not only the bytes that arrived', async (t) => { + const bytes = randomBytes(64 * 1024); + const stub = await egressStub(t, bytes); + const { root } = await workspace(t, stub); + const relay = new HttpWorkspaceFileImporter({ upstreamUrl: stub.url }); + const tools = await LocalWorkspaceTools.create({ + importer: { + maxBytes: relay.maxBytes, + availableBytes: (directory) => relay.availableBytes(directory), + async open(request, signal) { + const chunks = await relay.open(request, signal); + return (async function* tamper() { + let first = true; + for await (const chunk of chunks) { + yield chunk; + if (!first) continue; + first = false; + // Another local writer rewrites already-staged bytes before the transfer finishes. + const [staged] = await temporaryFiles(root); + assert.ok(staged); + await writeFile(join(root, staged), Buffer.alloc(chunk.byteLength, 0x41), { flag: 'r+' }); + } + })(); + }, + }, + workspaces: [{ id: 'primary', root, writable: true, identity: await captureWorkspaceRootIdentity(root) }], + }); + const tampered = await rejection(tools.execute(importRequest(bytes, { path: 'payload.bin' }))); + assert.equal(tampered.code, 'INVALID_REQUEST'); + assert.match(tampered.message, /SHA-256/); + assert.deepEqual(await readdir(root), []); +}); + +test('rollback never removes a directory that replaced one the import created', async (t) => { + const bytes = randomBytes(32 * 1024); + const stub = await egressStub(t, Buffer.concat([bytes.subarray(0, 1024), Buffer.alloc(bytes.byteLength - 1024)])); + const { root } = await workspace(t, stub); + const relay = new HttpWorkspaceFileImporter({ upstreamUrl: stub.url }); + const tools = await LocalWorkspaceTools.create({ + importer: { + maxBytes: relay.maxBytes, + availableBytes: (directory) => relay.availableBytes(directory), + async open(request, signal) { + const chunks = await relay.open(request, signal); + return (async function* swap() { + for await (const chunk of chunks) yield chunk; + // Another local process moves the new parent aside and makes its own at that path. + await rename(join(root, 'incoming'), join(root, 'moved')); + await mkdir(join(root, 'incoming')); + })(); + }, + }, + workspaces: [{ id: 'primary', root, writable: true, identity: await captureWorkspaceRootIdentity(root) }], + }); + const failed = await rejection(tools.execute(importRequest(bytes, { path: 'incoming/file.bin' }))); + assert.equal(failed.code, 'INVALID_REQUEST'); + assert.equal(failed.mutationMayHaveCommitted, true, 'an unremovable created parent is reported, not hidden'); + assert.equal((await stat(join(root, 'incoming'))).isDirectory(), true); +}); + +test('the stall timeout covers waiting on the source, not slow consumer writes', async (t) => { + const bytes = randomBytes(8 * 1024); + const stub = await egressStub(t, bytes); + stub.respond = (_req, res) => { + res.writeHead(200, { 'content-length': String(bytes.byteLength) }); + res.write(bytes.subarray(0, 4096)); + setTimeout(() => res.end(bytes.subarray(4096)), 10); + }; + const importer = new HttpWorkspaceFileImporter({ upstreamUrl: stub.url, stallTimeoutMs: 50 }); + const received: Buffer[] = []; + for await (const chunk of await importer.open(importRequest(bytes) as WorkspaceImportFileRequest & { egressGrant: string })) { + received.push(Buffer.from(chunk)); + await delay(150); + } + assert.deepEqual(Buffer.concat(received), bytes); +}); + +test('file and free-space limits fail before anything is written', async (t) => { + const bytes = randomBytes(2048); + const stub = await egressStub(t, bytes); + const small = await workspace(t, stub, { importer: { maxBytes: 1024 } }); + const tooLarge = await rejection(small.tools.execute(importRequest(bytes, { path: 'big.bin' }))); + assert.equal(tooLarge.code, 'WRITE_LIMIT_EXCEEDED'); + assert.equal(stub.hits, 0); + + const full = await workspace(t, stub, { importer: { availableBytes: async () => 4096 } }); + const noSpace = await rejection(full.tools.execute(importRequest(bytes, { path: 'big.bin' }))); + assert.equal(noSpace.code, 'WRITE_LIMIT_EXCEEDED'); + assert.match(noSpace.message, /free space/); + assert.equal(stub.hits, 0); + assert.deepEqual(await readdir(full.root), []); +}); + +test('imports refuse symlinked parents, Git metadata, and requests without a grant', async (t) => { + const bytes = Buffer.from('payload'); + const stub = await egressStub(t, bytes); + const { root, tools } = await workspace(t, stub); + const outside = await realpath(await mkdtemp(join(tmpdir(), 'librechat-code-import-outside-'))); + t.after(() => rm(outside, { recursive: true, force: true })); + await mkdir(join(root, 'real')); + await symlink(join(root, 'real'), join(root, 'alias')); + await symlink(outside, join(root, 'escape')); + await mkdir(join(root, '.git')); + + for (const path of ['alias/file.bin', 'escape/file.bin']) { + const refused = await rejection(tools.execute(importRequest(bytes, { path }))); + assert.equal(refused.code, 'INVALID_PATH', path); + } + assert.deepEqual(await readdir(join(root, 'real')), []); + assert.deepEqual(await readdir(outside), []); + + const git = await rejection(tools.execute(importRequest(bytes, { path: '.git/hooks/post-checkout' }))); + assert.equal(git.code, 'INVALID_REQUEST'); + assert.deepEqual(await readdir(join(root, '.git')), []); + + await symlink(join(root, 'real', 'target.bin'), join(root, 'link.bin')); + const linkTarget = await rejection(tools.execute(importRequest(bytes, { path: 'link.bin', overwrite: true }))); + assert.equal(linkTarget.code, 'INVALID_PATH'); + assert.equal((await lstat(join(root, 'link.bin'))).isSymbolicLink(), true); + + const { egressGrant: _grant, ...ungranted } = importRequest(bytes, { path: 'file.bin' }); + assert.equal((await rejection(tools.execute(ungranted))).code, 'INVALID_REQUEST'); + assert.equal(stub.hits, 0); +}); + +test('an unavailable or refusing relay fails cleanly as WRITE_UNAVAILABLE', async (t) => { + const bytes = Buffer.from('payload'); + const stub = await egressStub(t, bytes); + const { root, tools } = await workspace(t, stub); + const refused = await rejection(tools.execute(importRequest(bytes, { egressGrant: 'ceg1.another-object' }))); + assert.equal(refused.code, 'WRITE_UNAVAILABLE'); + assert.equal(refused.mutationMayHaveCommitted, false); + assert.deepEqual(await readdir(root), []); + + stub.respond = (_req, res) => { + res.writeHead(200, { 'content-length': String(bytes.byteLength) }); + res.write(bytes.subarray(0, 2)); + }; + const stalled = await workspace(t, stub, { importer: { stallTimeoutMs: 50 } }); + const timedOut = await rejection(stalled.tools.execute(importRequest(bytes))); + assert.equal(timedOut.code, 'WRITE_UNAVAILABLE'); + assert.match(timedOut.message, /stalled/); + assert.deepEqual(await readdir(stalled.root), []); +}); + +test('a worker negotiates import_file and executes it as a guarded mutation', async (t) => { + const bytes = randomBytes(8192); + const stub = await egressStub(t, bytes); + const { root, tools } = await workspace(t, stub); + const lifecycle: string[] = []; + const registrations: string[][] = []; + const settlements: Array> = []; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { statefulWorkspace: false, sandboxProfile: 'anthropic-srt', runtimes: [], workspaceTools: tools.capabilities }, + workspaceTools: tools, + workspaceMutationQuarantine: { + async assertAvailable() {}, + async arm() { lifecycle.push('arm'); }, + async clear() { lifecycle.push('clear'); }, + async quarantine() { lifecycle.push('quarantine'); }, + }, + fetchImpl: async (input, init) => { + const url = String(input); + const body = (init?.body == null ? {} : JSON.parse(String(init.body))) as Record; + if (url.endsWith('/register')) { + registrations.push((body.capabilities as { workspaceTools: { operations: string[] } }).workspaceTools.operations); + return Response.json({ + protocolVersion: 1, workerId: 'vm-1', incarnationId, registeredAt: new Date().toISOString(), leaseTtlMs: 60_000, + supportedWorkspaceToolOperations: ['read_file', 'search_text', 'list_files', 'write_file', 'preview_edit', 'edit_file', 'import_file'], + supportedWorkspaceWriteFileModes: ['replace', 'create'], + }); + } + if (!url.endsWith('/settle')) return Response.json({ protocolVersion: 1, cancelled: false }); + lifecycle.push('settle'); + settlements.push(body); + return Response.json({ protocolVersion: 1, accepted: true }); + }, + }); + await worker.register(); + assert.ok(registrations.at(-1)?.includes('import_file')); + + const request: WorkspaceToolRequest = importRequest(bytes, { path: 'chat/upload.bin' }); + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-import', + workerId: 'vm-1', + incarnationId, + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 10_000).toISOString(), + remainingMs: 10_000, + executionKind: 'workspace_tool', + request, + }); + assert.deepEqual(lifecycle, ['arm', 'settle', 'clear']); + assert.equal(settlements[0]?.status, 'fulfilled'); + assert.deepEqual(await readFile(join(root, 'chat', 'upload.bin')), bytes); +}); + +test('a worker never advertises import_file to a Code API that cannot relay it', async (t) => { + const stub = await egressStub(t, Buffer.from('x')); + const { tools } = await workspace(t, stub); + const registrations: string[][] = []; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { statefulWorkspace: false, sandboxProfile: 'anthropic-srt', runtimes: [], workspaceTools: tools.capabilities }, + workspaceTools: tools, + workspaceMutationQuarantine: { async assertAvailable() {}, async arm() {}, async clear() {}, async quarantine() {} }, + fetchImpl: async (_input, init) => { + const body = JSON.parse(String(init?.body)) as { capabilities: { workspaceTools: { operations: string[] } } }; + registrations.push(body.capabilities.workspaceTools.operations); + return Response.json({ + protocolVersion: 1, workerId: 'vm-1', incarnationId, registeredAt: new Date().toISOString(), leaseTtlMs: 60_000, + supportedWorkspaceToolOperations: ['read_file', 'search_text', 'list_files', 'write_file', 'preview_edit', 'edit_file'], + supportedWorkspaceWriteFileModes: ['replace', 'create'], + }); + }, + }); + await worker.register(); + assert.ok(registrations.length > 0); + for (const operations of registrations) assert.equal(operations.includes('import_file'), false); +}); diff --git a/packages/code/src/workspace-import.ts b/packages/code/src/workspace-import.ts new file mode 100644 index 00000000..e4c927d3 --- /dev/null +++ b/packages/code/src/workspace-import.ts @@ -0,0 +1,188 @@ +import { statfs } from 'node:fs/promises'; + +import { + BRIDGE_WORKSPACE_IMPORT_FILE_DEFAULT_MAX_BYTES, + BRIDGE_WORKSPACE_PROGRAMMATIC_TRANSFER_TIMEOUT_MS, +} from './protocol.js'; +import { + FILE_RELAY_GRANT_HEADER, + fileRelayObjectUrl, + validateFileRelayUpstream, +} from './relay.js'; +import { WorkspaceToolError } from './workspace.js'; + +import type { WorkspaceImportFileRequest } from './protocol.js'; +import type { WorkspaceFileImporter } from './workspace.js'; + +export interface HttpWorkspaceFileImporterOptions { + /** The egress gateway this worker already uses for programmatic inputs. */ + upstreamUrl: string; + maxBytes?: number; + /** Longest wait for response headers or for the next body chunk. */ + stallTimeoutMs?: number; + fetchImpl?: typeof fetch; + availableBytes?: (directory: string) => Promise; +} + +/** + * The importer a worker runs with, if any. `import_file` is on by default + * whenever the worker can perform it: writes enabled and a relay upstream to + * reach stored objects. `--no-workspace-imports` or + * `LIBRECHAT_CODE_ALLOW_WORKSPACE_IMPORTS=false` opts out; any other value of + * that variable fails startup rather than silently choosing either way. + */ +export function createWorkspaceFileImporter(input: { + optOut: boolean; + enabled?: string; + writable: boolean; + upstreamUrl?: string; + maxBytes?: number; +}): HttpWorkspaceFileImporter | undefined { + const enabled = input.enabled?.trim().toLowerCase() ?? ''; + if (enabled !== '' && enabled !== 'true' && enabled !== 'false') { + throw new Error('LIBRECHAT_CODE_ALLOW_WORKSPACE_IMPORTS must be true or false'); + } + const upstreamUrl = input.upstreamUrl?.trim() ?? ''; + if (input.optOut || enabled === 'false' || !input.writable || upstreamUrl === '') { + return undefined; + } + return new HttpWorkspaceFileImporter({ upstreamUrl, maxBytes: input.maxBytes }); +} + +async function filesystemAvailableBytes(directory: string): Promise { + const stats = await statfs(directory); + return Number(stats.bavail) * Number(stats.bsize); +} + +function importAborted(): WorkspaceToolError { + return new WorkspaceToolError('Workspace import aborted', 'EXECUTION_ABORTED'); +} + +function importUnavailable(message: string): WorkspaceToolError { + return new WorkspaceToolError(message, 'WRITE_UNAVAILABLE'); +} + +/** + * Streams one granted object from the file relay upstream. The worker holds no + * storage credential: Code API's grant authorizes exactly this object, briefly. + */ +export class HttpWorkspaceFileImporter implements WorkspaceFileImporter { + readonly maxBytes: number; + private readonly upstream: URL; + private readonly stallTimeoutMs: number; + private readonly fetchImpl: typeof fetch; + private readonly freeBytes: (directory: string) => Promise; + + constructor(options: HttpWorkspaceFileImporterOptions) { + this.upstream = validateFileRelayUpstream(options.upstreamUrl); + this.maxBytes = options.maxBytes ?? BRIDGE_WORKSPACE_IMPORT_FILE_DEFAULT_MAX_BYTES; + this.stallTimeoutMs = options.stallTimeoutMs ?? BRIDGE_WORKSPACE_PROGRAMMATIC_TRANSFER_TIMEOUT_MS; + if (!Number.isSafeInteger(this.maxBytes) || this.maxBytes < 1) { + throw new RangeError('Workspace import file limit must be a positive safe integer'); + } + if (!Number.isSafeInteger(this.stallTimeoutMs) || this.stallTimeoutMs < 1) { + throw new RangeError('Workspace import stall timeout must be a positive safe integer'); + } + this.fetchImpl = options.fetchImpl ?? fetch; + this.freeBytes = options.availableBytes ?? filesystemAvailableBytes; + } + + availableBytes(directory: string): Promise { + return this.freeBytes(directory); + } + + async open( + request: WorkspaceImportFileRequest & { egressGrant: string }, + signal?: AbortSignal, + ): Promise> { + if (signal?.aborted) throw importAborted(); + const controller = new AbortController(); + const abort = (): void => controller.abort(signal?.reason); + signal?.addEventListener('abort', abort, { once: true }); + let stalled = false; + let timer: ReturnType | undefined; + const arm = (): void => { + clearTimeout(timer); + timer = setTimeout(() => { + stalled = true; + controller.abort(); + }, this.stallTimeoutMs); + }; + const release = (): void => { + clearTimeout(timer); + signal?.removeEventListener('abort', abort); + }; + const failure = (): WorkspaceToolError => + signal?.aborted + ? importAborted() + : importUnavailable(stalled ? 'Import source transfer stalled' : 'Import source download failed'); + + arm(); + let response: Response; + try { + response = await this.fetchImpl( + fileRelayObjectUrl(this.upstream, request.source.storage_session_id, request.source.id), + { + headers: { + [FILE_RELAY_GRANT_HEADER]: request.egressGrant, + 'Accept-Encoding': 'identity', + }, + redirect: 'error', + signal: controller.signal, + }, + ); + } catch { + release(); + throw failure(); + } + // A compressed transfer declares its encoded length; only the decoded bytes are checked then. + const encoding = response.headers.get('content-encoding')?.trim().toLowerCase(); + const declaredLength = encoding == null || encoding === 'identity' + ? response.headers.get('content-length') + : null; + let rejection: WorkspaceToolError | undefined; + if (!response.ok) { + rejection = importUnavailable(`Import source download failed with HTTP ${response.status}`); + } else if (declaredLength != null && declaredLength !== String(request.bytes)) { + rejection = new WorkspaceToolError( + 'Import source size does not match the declared size', + 'INVALID_REQUEST', + ); + } else if (response.body == null && request.bytes > 0) { + rejection = importUnavailable('Import source returned no content'); + } + if (rejection != null) { + release(); + await response.body?.cancel().catch(() => undefined); + throw rejection; + } + const body = response.body; + async function* chunks(): AsyncGenerator { + if (body == null) { + release(); + return; + } + const reader = body.getReader(); + try { + for (;;) { + let next: Awaited>; + // The stall timer covers only waiting on the source, never the consumer's writes. + arm(); + try { + next = await reader.read(); + } catch { + throw failure(); + } finally { + clearTimeout(timer); + } + if (next.done) return; + yield next.value; + } + } finally { + release(); + await reader.cancel().catch(() => undefined); + } + } + return chunks(); + } +} diff --git a/packages/code/src/workspace-instances.ts b/packages/code/src/workspace-instances.ts index 68e76a10..7927152b 100644 --- a/packages/code/src/workspace-instances.ts +++ b/packages/code/src/workspace-instances.ts @@ -12,12 +12,13 @@ import type { WorkspaceToolRequest, WorkspaceToolResult, } from './protocol.js'; -import type { WorkspaceToolExecutor } from './workspace.js'; +import type { WorkspaceFileImporter, WorkspaceToolExecutor } from './workspace.js'; interface WorkspaceInstanceSource { command?: NativeProcessSandboxOptions; repositoryInstructions: boolean; writable: boolean; + importer?: WorkspaceFileImporter; } export interface GitWorktreeWorkspaceToolsOptions { @@ -119,6 +120,7 @@ export class GitWorktreeWorkspaceTools implements WorkspaceToolExecutor { identity: instance.identity, value: LocalWorkspaceTools.create({ repositoryInstructions: source.repositoryInstructions, + importer: source.importer, workspaces: [ { id: internalId, diff --git a/packages/code/src/workspace.ts b/packages/code/src/workspace.ts index d740d48a..0f6cbe34 100644 --- a/packages/code/src/workspace.ts +++ b/packages/code/src/workspace.ts @@ -1,6 +1,6 @@ import { createHash, randomBytes } from 'node:crypto'; import { constants } from 'node:fs'; -import { deferWorkspaceRootCleanup, holdsWorkspaceRoot, link, lstat, mkdir, open, realpath, rename, stat, unlink, spawn, withWorkspaceRoot, WorkspaceRootAccessError } from './root-access.js'; +import { deferWorkspaceRootCleanup, holdsWorkspaceRoot, link, lstat, mkdir, open, realpath, rename, rmdir, stat, unlink, spawn, withWorkspaceRoot, WorkspaceRootAccessError } from './root-access.js'; import { basename, dirname, isAbsolute, relative, resolve, sep } from 'node:path'; import type { FileHandle } from 'node:fs/promises'; @@ -17,6 +17,7 @@ import { BRIDGE_WORKSPACE_SEARCH_TEXT_MAX_LENGTH, comparePortableRelativePaths, isSafePortableRelativePath, + isWorkspaceImportFilePath, isValidBridgeWorkspaceToolCapabilities, isWorkspaceToolRequest, isWorkspaceToolResult, @@ -42,6 +43,8 @@ import type { WorkspacePreviewEditResult, WorkspaceExecuteCommandRequest, WorkspaceExecuteCommandResult, + WorkspaceImportFileRequest, + WorkspaceImportFileResult, WorkspaceListFilesRequest, WorkspaceListFilesResult, WorkspaceSearchMatch, @@ -65,6 +68,8 @@ export type { WorkspacePreviewEditResult, WorkspaceExecuteCommandRequest, WorkspaceExecuteCommandResult, + WorkspaceImportFileRequest, + WorkspaceImportFileResult, WorkspaceListFilesRequest, WorkspaceListFilesResult, WorkspaceSearchMatch, @@ -85,9 +90,27 @@ export interface LocalWorkspaceConfig { writable?: boolean; } +/** + * Fetches a stored object for `import_file`. Only a worker with a configured + * file relay upstream has one, and `import_file` is advertised only then. + */ +export interface WorkspaceFileImporter { + /** Largest file this worker imports; larger requests fail before any transfer. */ + readonly maxBytes: number; + /** Streams the granted object. Failures are WorkspaceToolErrors. */ + open( + request: WorkspaceImportFileRequest & { egressGrant: string }, + signal?: AbortSignal, + ): Promise>; + /** Bytes a new file in `directory` may still use on its filesystem. */ + availableBytes(directory: string): Promise; +} + export interface LocalWorkspaceToolsOptions { workspaces: LocalWorkspaceConfig[]; repositoryInstructions?: boolean; + /** Enables `import_file` on writable workspaces. */ + importer?: WorkspaceFileImporter; } export interface WorkspaceToolExecutor { @@ -144,6 +167,12 @@ const READ_OPERATIONS = [ 'list_files', ] as const; const WRITE_OPERATIONS = ['write_file', 'preview_edit', 'edit_file'] as const; +const IMPORT_OPERATIONS = ['import_file'] as const; +const IMPORT_DIRECTORY_MODE = 0o755; +const IMPORT_FILE_MODE = 0o644; +const IMPORT_HASH_CHUNK_BYTES = 1024 * 1024; +/** Headroom left on the target filesystem after an import lands. */ +export const WORKSPACE_IMPORT_FREE_SPACE_RESERVE_BYTES = 64 * 1024 * 1024; function isUtf8ScalarString(value: string): boolean { return Buffer.from(value).toString('utf8') === value; @@ -682,6 +711,7 @@ async function classifyMissingWorkspacePath( async function createMissingParentDirectories( root: string, candidate: string, + mode = 0o777, ): Promise { const created: string[] = []; if (!holdsWorkspaceRoot()) return created; @@ -699,7 +729,7 @@ async function createMissingParentDirectories( let createdHere = false; if (entry == null) { try { - await mkdir(current, 0o777); + await mkdir(current, mode); created.push(current); createdHere = true; } catch (error) { @@ -1070,6 +1100,373 @@ async function writeWorkspaceFile( }; } +interface ImportTargetIdentity { + dev: bigint | number; + ino: bigint | number; +} + +type ImportTarget = + | { state: 'missing' } + | { state: 'unchanged' } + | ({ state: 'different' } & ImportTargetIdentity); + +function importIntegrityError(message: string): WorkspaceToolError { + return new WorkspaceToolError(message, 'INVALID_REQUEST'); +} + +/** Size and SHA-256 of everything an open file holds now, read through its descriptor. */ +async function hashOpenFile( + handle: FileHandle, + signal?: AbortSignal, +): Promise<{ bytes: number; sha256: string }> { + const hash = createHash('sha256'); + const chunk = Buffer.allocUnsafe(IMPORT_HASH_CHUNK_BYTES); + let position = 0; + for (;;) { + throwIfAborted(signal); + const { bytesRead } = await handle.read(chunk, 0, chunk.byteLength, position); + if (bytesRead === 0) break; + hash.update(chunk.subarray(0, bytesRead)); + position += bytesRead; + } + return { bytes: position, sha256: hash.digest('hex') }; +} + +/** Hash an existing target only when its size already matches the import. */ +async function inspectImportTarget( + root: string, + request: WorkspaceImportFileRequest, + signal?: AbortSignal, +): Promise { + try { + return await withConfinedFile(root, request.path, async (handle, size) => { + const { dev, ino } = await handle.stat(); + if (size !== request.bytes) return { state: 'different', dev, ino }; + const existing = await hashOpenFile(handle, signal); + return existing.bytes === request.bytes && existing.sha256 === request.sha256 + ? { state: 'unchanged' } + : { state: 'different', dev, ino }; + }); + } catch (error) { + if (error instanceof WorkspaceToolError && error.code === 'NOT_FOUND') { + return { state: 'missing' }; + } + throw error; + } +} + +/** + * Undo parent directories this import created; true when every one is gone + * again. A directory is removed only while it is still the inode this import + * created, so a replacement made by another local process is never deleted. + */ +async function removeCreatedDirectories( + created: readonly string[], + identities: readonly ImportTargetIdentity[], +): Promise { + if (identities.length !== created.length) return false; + for (let index = created.length - 1; index >= 0; index -= 1) { + const directory = created[index]!; + try { + const current = await lstat(directory); + if ( + current.isSymbolicLink() || + !current.isDirectory() || + current.dev !== identities[index]!.dev || + current.ino !== identities[index]!.ino + ) { + return false; + } + await rmdir(directory); + await syncWorkspaceDirectory(dirname(directory)); + } catch { + return false; + } + } + return true; +} + +/** + * Record each directory this import created, then give it the documented mode + * regardless of the worker's umask. `identities` fills as each is opened. + */ +async function setImportDirectoryModes( + created: readonly string[], + identities: ImportTargetIdentity[], +): Promise { + for (const directory of created) { + const handle = await open( + directory, + constants.O_RDONLY | constants.O_DIRECTORY | constants.O_NOFOLLOW, + ); + try { + const { dev, ino } = await handle.stat(); + identities.push({ dev, ino }); + await handle.chmod(IMPORT_DIRECTORY_MODE); + } finally { + await handle.close(); + } + } +} + +async function writeImportChunk(handle: FileHandle, chunk: Uint8Array, position: number): Promise { + let offset = 0; + while (offset < chunk.byteLength) { + const { bytesWritten } = await handle.write(chunk, offset, chunk.byteLength - offset, position + offset); + if (bytesWritten < 1) throw new Error('short workspace write'); + offset += bytesWritten; + } +} + +/** Confirm the installed entry is the staged inode, inside the root, after syncing its directory. */ +async function confirmInstalledImport( + root: string, + installTarget: string, + staged: ImportTargetIdentity, +): Promise { + let handle: FileHandle | undefined; + try { + await syncWorkspaceDirectory(dirname(installTarget)); + handle = await open( + installTarget, + constants.O_RDONLY | constants.O_NOFOLLOW | constants.O_NONBLOCK, + ); + const [installed, canonicalPath] = await Promise.all([handle.stat(), realpath(installTarget)]); + if ( + !installed.isFile() || + !isWithinRoot(root, canonicalPath) || + installed.dev !== staged.dev || + installed.ino !== staged.ino + ) { + throw new Error('installed workspace import changed'); + } + } catch { + throw new WorkspaceToolError( + 'Workspace import durability could not be confirmed', + 'WRITE_UNAVAILABLE', + true, + ); + } finally { + await handle?.close().catch(() => undefined); + } +} + +/** + * Stream one object into a same-directory temporary file, verify the synced + * file's size and SHA-256 through its own descriptor, then install it + * atomically. A create links exclusively so a racing file is never replaced. + * An overwrite rechecks the inspected inode immediately before its rename, the + * same check-then-rename `write_file` uses; bridge mutations of one root are + * serialized, so only an out-of-band local writer can race that last step. + */ +async function installImportedFile( + root: string, + candidate: string, + request: WorkspaceImportFileRequest & { egressGrant: string }, + importer: WorkspaceFileImporter, + target: Exclude, + signal?: AbortSignal, +): Promise<{ created: boolean }> { + const parent = dirname(candidate); + let canonicalParent: string; + let parentIdentity: Awaited>; + try { + canonicalParent = await realpath(parent); + parentIdentity = await verifyDirectoryPathHasNoSymlinks(root, parent); + if (!isWithinRoot(root, canonicalParent) || !parentIdentity.isDirectory()) { + throw new WorkspaceToolError('Invalid workspace path', 'INVALID_PATH'); + } + } catch (error) { + if (isMissingEntry(error)) throw await classifyMissingWorkspacePath(root, parent); + throw classifyWritePathValidationError(error); + } + let available: number; + try { + available = await importer.availableBytes(canonicalParent); + } catch { + throw new WorkspaceToolError('Workspace free space could not be determined', 'WRITE_UNAVAILABLE'); + } + if (available < request.bytes + WORKSPACE_IMPORT_FREE_SPACE_RESERVE_BYTES) { + throw new WorkspaceToolError( + 'Workspace filesystem does not have enough free space for this file', + 'WRITE_LIMIT_EXCEEDED', + ); + } + + const temporary = resolve(canonicalParent, `.librechat-code-${randomBytes(18).toString('hex')}.tmp`); + const installTarget = resolve(canonicalParent, basename(candidate)); + let handle: FileHandle | undefined; + let temporaryNeedsCleanup = false; + try { + handle = await open( + temporary, + constants.O_RDWR | constants.O_CREAT | constants.O_EXCL | constants.O_NOFOLLOW, + 0o600, + ); + temporaryNeedsCleanup = true; + let written = 0; + for await (const chunk of await importer.open(request, signal)) { + throwIfAborted(signal); + if (written + chunk.byteLength > request.bytes) { + throw importIntegrityError('Imported file is larger than its declared size'); + } + await writeImportChunk(handle, chunk, written); + written += chunk.byteLength; + } + throwIfAborted(signal); + if (written !== request.bytes) { + throw importIntegrityError('Imported file is smaller than its declared size'); + } + await handle.chmod(IMPORT_FILE_MODE); + await handle.sync(); + // Verify what will be installed, not what arrived: the staged file is a workspace entry. + const staged = await hashOpenFile(handle, signal); + if (staged.bytes !== request.bytes) { + throw importIntegrityError('Imported file does not match its declared size'); + } + if (staged.sha256 !== request.sha256) { + throw importIntegrityError('Imported file does not match its declared SHA-256'); + } + const stagedStat = await handle.stat(); + const stagedIdentity = { dev: stagedStat.dev, ino: stagedStat.ino }; + if (process.platform === 'win32') { + await handle.close(); + handle = undefined; + } + + const [currentParent, currentParentIdentity] = await Promise.all([ + realpath(parent), + verifyDirectoryPathHasNoSymlinks(root, parent), + ]); + if ( + currentParent !== canonicalParent || + currentParentIdentity.dev !== parentIdentity.dev || + currentParentIdentity.ino !== parentIdentity.ino + ) { + throw new WorkspaceToolError('Invalid workspace path', 'INVALID_PATH'); + } + let current: Awaited> | undefined; + try { + current = await lstat(candidate); + } catch (error) { + if (!isMissingEntry(error)) throw error; + } + if ( + target.state === 'missing' + ? current != null + : current == null || current.isSymbolicLink() || current.dev !== target.dev || current.ino !== target.ino + ) { + throw new WorkspaceToolError( + 'Workspace file changed before import could be committed', + 'EDIT_CONFLICT', + ); + } + throwIfAborted(signal); + if (target.state === 'different') { + await rename(temporary, installTarget); + temporaryNeedsCleanup = false; + } else { + try { + await link(temporary, installTarget); + } catch (error) { + if ((error as NodeJS.ErrnoException).code === 'EEXIST') { + throw new WorkspaceToolError('Workspace file already exists', 'EDIT_CONFLICT'); + } + throw error; + } + try { + await unlink(temporary); + temporaryNeedsCleanup = false; + } catch { + throw new WorkspaceToolError( + 'Workspace import cleanup could not be confirmed', + 'WRITE_UNAVAILABLE', + true, + ); + } + } + await confirmInstalledImport(root, installTarget, stagedIdentity); + return { created: target.state === 'missing' }; + } catch (error) { + if (error instanceof WorkspaceToolError) throw error; + throw new WorkspaceToolError('Workspace storage is unavailable', 'WRITE_UNAVAILABLE'); + } finally { + await handle?.close().catch(() => undefined); + if (temporaryNeedsCleanup) await removeImportStaging(temporary); + } +} + +/** A staging file that cannot be removed is upload bytes left in the workspace: never a clean failure. */ +async function removeImportStaging(temporary: string): Promise { + try { + await unlink(temporary); + } catch (error) { + if (isMissingEntry(error)) return; + throw new WorkspaceToolError( + 'Workspace import staging file could not be removed', + 'WRITE_UNAVAILABLE', + true, + ); + } +} + +async function importWorkspaceFile( + root: string, + request: WorkspaceImportFileRequest, + importer: WorkspaceFileImporter, + signal?: AbortSignal, +): Promise { + throwIfAborted(signal); + if (!isWorkspaceImportFilePath(request.path)) { + throw new WorkspaceToolError('Invalid workspace path', 'INVALID_PATH'); + } + const { egressGrant } = request; + if (egressGrant === undefined) { + throw new WorkspaceToolError('Import source grant is unavailable', 'INVALID_REQUEST'); + } + if (request.bytes > importer.maxBytes) { + throw new WorkspaceToolError('Imported file exceeds the worker file limit', 'WRITE_LIMIT_EXCEEDED'); + } + const result = (created: boolean, bytesWritten: number, unchanged = false): WorkspaceImportFileResult => ({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operation: 'import_file', + workspaceId: request.workspaceId, + path: request.path, + created, + ...(unchanged ? { unchanged: true } : {}), + bytesWritten, + sha256: request.sha256, + }); + const candidate = resolveWorkspacePath(root, request.path); + const target = await inspectImportTarget(root, request, signal); + if (target.state === 'unchanged') return result(false, 0, true); + if (target.state === 'different' && request.overwrite !== true) { + throw new WorkspaceToolError( + 'Workspace file already exists with different content', + 'EDIT_CONFLICT', + ); + } + const directories = await createMissingParentDirectories(root, candidate, IMPORT_DIRECTORY_MODE); + const identities: ImportTargetIdentity[] = []; + try { + await setImportDirectoryModes(directories, identities); + const { created } = await installImportedFile( + root, + candidate, + { ...request, egressGrant }, + importer, + target, + signal, + ); + return result(created, request.bytes); + } catch (error) { + const failure = classifyWritePathValidationError(error); + if (directories.length === 0 || failure.mutationMayHaveCommitted) throw failure; + if (await removeCreatedDirectories(directories, identities)) throw failure; + throw afterCreatedDirectories(directories, failure); + } +} + async function editWorkspaceFile( root: string, request: WorkspaceEditFileRequest, @@ -1878,6 +2275,7 @@ export class LocalWorkspaceTools implements WorkspaceToolExecutor { editFileModes?: BridgeWorkspaceToolCapabilities['editFileModes'], editFileFeatures?: BridgeWorkspaceToolCapabilities['editFileFeatures'], listFileFeatures?: BridgeWorkspaceToolCapabilities['listFileFeatures'], + private readonly importer?: WorkspaceFileImporter, ) { this.capabilities = { protocolVersion: BRIDGE_PROTOCOL_VERSION, @@ -1897,9 +2295,13 @@ export class LocalWorkspaceTools implements WorkspaceToolExecutor { const anyWritable = options.workspaces.some( (workspace) => workspace.writable === true, ); + const mutations = [ + ...WRITE_OPERATIONS, + ...(options.importer != null ? IMPORT_OPERATIONS : []), + ]; const operations: BridgeWorkspaceToolCapabilities['operations'] = [ ...READ_OPERATIONS, - ...(anyWritable ? WRITE_OPERATIONS : []), + ...(anyWritable ? mutations : []), ]; const workspaces: BridgeWorkspaceDescriptor[] = options.workspaces.map( (workspace) => ({ @@ -1909,7 +2311,7 @@ export class LocalWorkspaceTools implements WorkspaceToolExecutor { ? { operations: [ ...READ_OPERATIONS, - ...(workspace.writable === true ? WRITE_OPERATIONS : []), + ...(workspace.writable === true ? mutations : []), ], } : {}), @@ -1959,6 +2361,7 @@ export class LocalWorkspaceTools implements WorkspaceToolExecutor { capabilities.editFileModes, capabilities.editFileFeatures, capabilities.listFileFeatures, + anyWritable ? options.importer : undefined, ); tools.repositoryInstructions = options.repositoryInstructions === true; return tools; @@ -2000,6 +2403,22 @@ export class LocalWorkspaceTools implements WorkspaceToolExecutor { } const { root } = workspace; + if (request.operation === 'import_file') { + if (!workspace.writable) { + throw new WorkspaceToolError( + 'Workspace mutations are disabled by the worker', + 'WRITE_DISABLED', + ); + } + if (this.importer == null) { + throw new WorkspaceToolError( + 'Workspace file import is not configured on this worker', + 'WRITE_UNAVAILABLE', + ); + } + return importWorkspaceFile(root, request, this.importer, signal); + } + if ( request.operation === 'write_file' || request.operation === 'preview_edit' || diff --git a/service/openapi.yml b/service/openapi.yml index 07d2d7e3..929fab71 100644 --- a/service/openapi.yml +++ b/service/openapi.yml @@ -72,7 +72,12 @@ components: protocolVersion: { type: integer, enum: [1] } operation: type: string - enum: [list_files, search_text, read_file, write_file, preview_edit, edit_file, execute_command] + enum: [list_files, search_text, read_file, write_file, preview_edit, edit_file, execute_command, import_file] + description: >- + `import_file` copies a stored object the caller owns into the + workspace: `path`, `source` ({ storage_session_id, id }), `bytes`, + lowercase hex `sha256`, optional `overwrite`. Send it only when the + worker advertises it. Callers never send `egressGrant`. workspaceId: { type: string } additionalProperties: true diff --git a/service/src/bridge/enabled.ts b/service/src/bridge/enabled.ts index 540a844a..e0cba569 100644 --- a/service/src/bridge/enabled.ts +++ b/service/src/bridge/enabled.ts @@ -7,3 +7,8 @@ export function isBridgeEnabled(): boolean { || env.BRIDGE_DYNAMIC_WORKERS || env.BRIDGE_TOKEN.length > 0; } + +/** `import_file` relays need the egress gateway to mint each one-object read grant. */ +export function isWorkspaceImportEnabled(): boolean { + return env.EGRESS_GATEWAY_URL.trim() !== ''; +} diff --git a/service/src/bridge/index.ts b/service/src/bridge/index.ts index f7ba03b4..93370dfa 100644 --- a/service/src/bridge/index.ts +++ b/service/src/bridge/index.ts @@ -3,7 +3,7 @@ import { env } from '../config'; import { RedisBridgePairingStore } from './pairing'; import { createBridgeRouter } from './router'; import { RedisBridgeStore } from './store'; -import { isBridgeEnabled } from './enabled'; +import { isBridgeEnabled, isWorkspaceImportEnabled } from './enabled'; export const bridgeStore = new RedisBridgeStore( connection, @@ -39,4 +39,5 @@ export default createBridgeRouter({ configuredWorkerId: env.BRIDGE_WORKER_ID, allowDynamicWorkers: env.BRIDGE_DYNAMIC_WORKERS, maxCommandTimeoutMs: env.JOB_TIMEOUT, + workspaceImportFile: isWorkspaceImportEnabled(), }); diff --git a/service/src/bridge/router.test.ts b/service/src/bridge/router.test.ts index 5e76876e..a478f623 100644 --- a/service/src/bridge/router.test.ts +++ b/service/src/bridge/router.test.ts @@ -216,6 +216,53 @@ describe('paired bridge HTTP API', () => { }); } + for (const workspaceImportFile of [false, true]) { + test(`registration ${workspaceImportFile ? 'offers' : 'omits'} import_file when grants ${workspaceImportFile ? 'can' : 'cannot'} be minted`, async () => { + const app = express(); + app.use(json()); + app.use( + '/v1/bridge', + createBridgeRouter({ + store: new RedisBridgeStore(redis), + pairings: new RedisBridgePairingStore(redis), + authMode: 'static', + adminToken: 'strong-administrator-bootstrap-token', + configuredWorkerId: 'vm-1', + workspaceImportFile, + }), + ); + 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}/v1/bridge/workers/register`, { + method: 'POST', + headers: { + Authorization: 'Bearer strong-administrator-bootstrap-token', + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'anthropic-srt', + runtimes: [], + workspaceTools: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operations: ['read_file', 'write_file', 'import_file'], + workspaces: [{ id: 'primary' }], + }, + }, + }), + }); + expect(response.status).toBe(200); + const body = (await response.json()) as { supportedWorkspaceToolOperations: string[] }; + expect(body.supportedWorkspaceToolOperations.includes('import_file')).toBe(workspaceImportFile); + }); + } + test('rejects a malformed optional binding for a configured worker', async () => { const app = express(); app.use(json()); diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index d035b032..c3b5aead 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -45,6 +45,8 @@ export interface BridgeRouterOptions { configuredWorkerId?: string; allowDynamicWorkers?: boolean; maxCommandTimeoutMs?: number; + /** Offer `import_file`; true only when this Code API can mint its egress grant. */ + workspaceImportFile?: boolean; } function sameToken(left: string, right: string): boolean { @@ -606,6 +608,7 @@ router.post( 'preview_edit', 'edit_file', 'execute_command', + ...(options.workspaceImportFile === true ? ['import_file' as const] : []), ], supportedWorkspaceWriteFileModes: ['replace', 'create'], supportedWorkspaceEditFileModes: ['single', 'batch'], diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index 0ca74c45..f9919f49 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -15,6 +15,7 @@ import { BRIDGE_PROTOCOL_VERSION, isValidBridgeWorkerCapabilities, isValidBridgeWorkerId, + isWorkspaceMutationOperation, isWorkspaceToolRequest, isWorkspaceToolResult, workspaceIsolationKey, @@ -169,7 +170,7 @@ export interface BridgeWorkerStatus { } function isWorkspaceMutation(request: WorkspaceToolRequest): boolean { - return request.operation === 'write_file' || request.operation === 'edit_file' || request.operation === 'execute_command'; + return isWorkspaceMutationOperation(request.operation); } function supportsWorkspaceTool( diff --git a/service/src/config.ts b/service/src/config.ts index b3e397b8..f8fad667 100644 --- a/service/src/config.ts +++ b/service/src/config.ts @@ -7,6 +7,7 @@ import { resolveExecutionProfile, resolveExecutionProfileSource, } from './execution-profile'; +import { BRIDGE_WORKSPACE_IMPORT_FILE_DEFAULT_MAX_BYTES } from '../../packages/code/src/protocol'; export const languageConfig: Record = { [Languages.bash]: { language: 'bash', version: '5.2.0', fileName: 'script.sh' }, @@ -425,6 +426,14 @@ export const env = { * Off by default: enable only after the LibreChat talking to this service accepts the field. */ BRIDGE_LANE_GIT: process.env.CODEAPI_BRIDGE_LANE_GIT === 'true', + /** + * Largest stored object relayed into a worker by `import_file`. The operation + * is offered only when EGRESS_GATEWAY_URL can mint its one-object grant. + */ + BRIDGE_IMPORT_FILE_MAX_BYTES: resolvePositiveIntEnv( + process.env.CODEAPI_BRIDGE_IMPORT_FILE_MAX_BYTES, + BRIDGE_WORKSPACE_IMPORT_FILE_DEFAULT_MAX_BYTES, + ), /** Outbound worker selected by the remote-bridge backend. */ BRIDGE_WORKER_ID: process.env.CODEAPI_BRIDGE_WORKER_ID ?? '', /** Static compatibility auth or short-lived proof-of-possession credentials. */ diff --git a/service/src/session-ownership.ts b/service/src/session-ownership.ts index bf4ba6dd..1841ed06 100644 --- a/service/src/session-ownership.ts +++ b/service/src/session-ownership.ts @@ -94,7 +94,7 @@ export type SessionOwnershipResult = * simply gone. */ export async function authorizeSessionOwnership( - store: SessionOwnershipStore, + store: Pick, args: { session_id: string; expectedSessionKey: string; diff --git a/service/src/workspace-tools/import.test.ts b/service/src/workspace-tools/import.test.ts new file mode 100644 index 00000000..7aa22aef --- /dev/null +++ b/service/src/workspace-tools/import.test.ts @@ -0,0 +1,295 @@ +import { createServer } from 'node:http'; +import type { Server } from 'node:http'; + +import { afterEach, expect, test } from 'bun:test'; +import express, { json } from 'express'; +import RedisMock from 'ioredis-mock'; +import type Redis from 'ioredis'; + +import type * as t from '../types'; +import type { CodeApiPrincipal } from '../auth/principal'; +import type { WorkspaceImportFileRequest, WorkspaceToolRequest } from '../../../packages/code/src/protocol'; + +import { applyPrincipal } from '../auth/principal'; +import { RedisBridgeStore } from '../bridge/store'; +import { resolveSessionKey } from '../session-key'; +import { recordSessionOwnership } from '../session-ownership'; +import { openEgressGrant, openEgressHandle, prepareSandboxEgress } from '../egress-grant'; +import { createWorkspaceToolsRouter } from './router'; +import { RedisWorkspaceRequests } from './requests'; +import { + WorkspaceImportError, + authorizeWorkspaceImportSource, + createWorkspaceImportService, +} from './import'; + +const SECRET = 'workspace-import-test-secret-of-sufficient-length'; +const SESSION = 'sessAAAAAAAAAAAAAAAAA'; +const OBJECT = 'fileAAAAAAAAAAAAAAAAA'; +const OTHER_SESSION = 'sessBBBBBBBBBBBBBBBBB'; +const OTHER_OBJECT = 'fileBBBBBBBBBBBBBBBBB'; +const owner: CodeApiPrincipal = { userId: 'user-1', tenantId: 'tenant-1', principalSource: 'librechat_jwt', codeWorkerId: 'worker' }; +const stranger: CodeApiPrincipal = { userId: 'user-2', tenantId: 'tenant-1', principalSource: 'librechat_jwt', codeWorkerId: 'worker' }; + +const importBody = (overrides: Record = {}): WorkspaceImportFileRequest => ({ + protocolVersion: 1, + operation: 'import_file', + workspaceId: 'primary', + path: 'uploads/report.pdf', + source: { storage_session_id: SESSION, id: OBJECT }, + bytes: 4096, + sha256: 'a'.repeat(64), + ...overrides, +}); + +function principalRequest(principal: CodeApiPrincipal): t.AuthenticatedRequest { + const req = { headers: {} } as t.AuthenticatedRequest; + applyPrincipal(req, principal); + return req; +} + +/** Uploads as LibreChat's code-environment upload does: user-kind session key plus upload record. */ +async function upload(redis: Redis, principal: CodeApiPrincipal, sessionId: string, objectId: string): Promise { + const sessionKey = resolveSessionKey(principalRequest(principal), { kind: 'user', id: principal.userId }); + await recordSessionOwnership(redis, sessionId, sessionKey); + await redis.set(`upload:${sessionKey}${sessionId}${objectId}`, '1'); +} + +function importService(redis: Redis, minted: string[] = [], revoked: string[] = [], maxBytes = 1024 * 1024) { + return createWorkspaceImportService({ + maxBytes, + store: redis, + async createGrant({ payload, claims }) { + const grantId = `grant-${minted.length + 1}`; + minted.push(grantId); + return { grant_id: grantId, ...prepareSandboxEgress({ payload, claims, grantId, secret: SECRET }) }; + }, + async revokeGrant(grantId) { + revoked.push(grantId); + }, + }); +} + +// ioredis-mock instances share one keyspace across test files; clear it like the bridge suites do. +const redis = new RedisMock() as unknown as Redis; +let server: Server | undefined; +afterEach(async () => { + server?.close(); + server = undefined; + await redis.flushall(); +}); + +async function listen(app: express.Express): Promise { + 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('Missing listener'); + return `http://127.0.0.1:${address.port}`; +} + +function appFor(router: express.Router): express.Express { + const app = express(); + app.use(json()); + app.use((req, _res, next) => { + applyPrincipal(req, req.header('X-Test-User') === 'user-2' ? stranger : owner); + next(); + }); + app.use(router); + return app; +} + +function post(url: string, body: unknown, headers: Record = {}, signal?: AbortSignal): Promise { + return fetch(url, { method: 'POST', headers: { 'Content-Type': 'application/json', ...headers }, body: JSON.stringify(body), signal }); +} + +test('ownership follows the download rule and rejects another principal\'s object', async () => { + await upload(redis, owner, SESSION, OBJECT); + await upload(redis, stranger, OTHER_SESSION, OTHER_OBJECT); + const req = principalRequest(owner); + + expect(await authorizeWorkspaceImportSource({ req, store: redis, source: { storage_session_id: SESSION, id: OBJECT } })) + .toBe(resolveSessionKey(req, { kind: 'user', id: 'user-1' })); + const denials = [ + { storage_session_id: OTHER_SESSION, id: OTHER_OBJECT }, + { storage_session_id: SESSION, id: OTHER_OBJECT }, + { storage_session_id: OTHER_SESSION, id: OBJECT }, + ]; + for (const source of denials) { + const error = await authorizeWorkspaceImportSource({ req, store: redis, source }).catch(caught => caught); + expect(error).toBeInstanceOf(WorkspaceImportError); + expect(error).toMatchObject({ status: 403, code: 'IMPORT_SOURCE_UNAUTHORIZED' }); + } + const malformed = await authorizeWorkspaceImportSource({ + req, store: redis, source: { storage_session_id: '../escape', id: OBJECT }, + }).catch(caught => caught); + expect(malformed).toMatchObject({ status: 400, code: 'INVALID_WORKSPACE_TOOL_REQUEST' }); + + await redis.del(`session:${SESSION}`); + const expired = await authorizeWorkspaceImportSource({ req, store: redis, source: { storage_session_id: SESSION, id: OBJECT } }) + .catch(caught => caught); + expect(expired).toMatchObject({ status: 403, code: 'IMPORT_SOURCE_UNAUTHORIZED' }); +}); + +test('the grant reads exactly one object, briefly, and can upload nothing', async () => { + await upload(redis, owner, SESSION, OBJECT); + const imports = importService(redis); + const req = principalRequest(owner); + const sessionKey = await imports.authorize(req, importBody()); + const before = Math.floor(Date.now() / 1000); + const prepared = await imports.prepare({ req, request: importBody(), sessionKey, ttlSeconds: 60 }); + + const grant = openEgressGrant(prepared.request.egressGrant, SECRET); + expect(grant.input_files).toEqual([{ id: OBJECT, session_id: SESSION, name: 'report.pdf' }]); + expect(grant.read_sessions).toEqual([SESSION]); + expect(grant).toMatchObject({ max_upload_bytes: 0, max_output_files: 0, max_requests: 2, session_key: sessionKey }); + expect(grant.exp - before).toBeGreaterThanOrEqual(90); + expect(grant.exp - before).toBeLessThanOrEqual(91); + expect(prepared.request.source.id).not.toBe(OBJECT); + expect(prepared.request.source.storage_session_id).not.toBe(SESSION); + expect(openEgressHandle(prepared.request.source.id, SECRET)).toMatchObject({ typ: 'object', object_id: OBJECT, session_id: SESSION }); + expect(openEgressHandle(prepared.request.source.storage_session_id, SECRET)).toMatchObject({ typ: 'session', dir: 'read', session_id: SESSION }); + const { egressGrant: _grant, source: _source, ...unchanged } = prepared.request; + const { source: _callerSource, ...callerFields } = importBody(); + expect(unchanged).toEqual(callerFields); +}); + +test('a synchronous import relays the granted request and revokes its grant', async () => { + await upload(redis, owner, SESSION, OBJECT); + await upload(redis, stranger, OTHER_SESSION, OTHER_OBJECT); + const minted: string[] = []; + const revoked: string[] = []; + const dispatched: Array<{ request: WorkspaceToolRequest; executionTimeoutMs?: number }> = []; + const url = await listen(appFor(createWorkspaceToolsRouter({ + backend: 'remote-bridge', configuredWorkerId: 'worker', dynamicWorkers: false, timeoutMs: 120_000, + imports: importService(redis, minted, revoked, 8192), + store: { async dispatchWorkspaceTool(args) { + dispatched.push({ request: args.request, executionTimeoutMs: args.executionTimeoutMs }); + const request = args.request as WorkspaceImportFileRequest; + return { protocolVersion: 1, generation: 1, leaseToken: 'lease', incarnationId: 'incarnation', status: 'fulfilled', result: { + protocolVersion: 1, operation: 'import_file', workspaceId: 'primary', path: request.path, + created: true, bytesWritten: request.bytes, sha256: request.sha256, + } }; + } }, + }))); + const endpoint = `${url}/workspace-tools/execute`; + + const response = await post(endpoint, importBody()); + expect(response.status).toBe(200); + expect(await response.json()).toMatchObject({ operation: 'import_file', created: true, bytesWritten: 4096 }); + expect(dispatched).toHaveLength(1); + const relayed = dispatched[0]!.request as WorkspaceImportFileRequest; + expect(typeof relayed.egressGrant).toBe('string'); + expect(relayed.source.id).not.toBe(OBJECT); + expect(dispatched[0]!.executionTimeoutMs).toBe(120_000); + await Bun.sleep(0); + expect(revoked).toEqual(minted); + + const callerGrant = await post(endpoint, { ...importBody(), egressGrant: 'ceg1.forged' }); + expect(callerGrant.status).toBe(400); + const foreign = await post(endpoint, importBody({ source: { storage_session_id: OTHER_SESSION, id: OTHER_OBJECT } })); + expect(foreign.status).toBe(403); + expect(await foreign.json()).toMatchObject({ code: 'IMPORT_SOURCE_UNAUTHORIZED' }); + const stolen = await post(endpoint, importBody(), { 'X-Test-User': 'user-2' }); + expect(stolen.status).toBe(403); + const oversized = await post(endpoint, importBody({ bytes: 8193 })); + expect(oversized.status).toBe(413); + expect(await oversized.json()).toMatchObject({ code: 'WRITE_LIMIT_EXCEEDED' }); + expect(dispatched).toHaveLength(1); + expect(minted).toHaveLength(1); +}); + +test('a caller that disconnects while its grant is minted never dispatches, and the grant is revoked', async () => { + await upload(redis, owner, SESSION, OBJECT); + const minted: string[] = []; + const revoked: string[] = []; + const base = importService(redis, minted, revoked); + const grantRequested = Promise.withResolvers(); + const releaseGrant = Promise.withResolvers(); + let dispatched = 0; + const url = await listen(appFor(createWorkspaceToolsRouter({ + backend: 'remote-bridge', configuredWorkerId: 'worker', dynamicWorkers: false, + imports: { ...base, async prepare(args) { + grantRequested.resolve(); + await releaseGrant.promise; + return base.prepare(args); + } }, + store: { async dispatchWorkspaceTool(): Promise { dispatched++; throw new Error('must not dispatch'); } }, + }))); + const caller = new AbortController(); + const response = post(`${url}/workspace-tools/execute`, importBody(), {}, caller.signal).catch(() => undefined); + await grantRequested.promise; + caller.abort(); + await response; + await Bun.sleep(50); + releaseGrant.resolve(); + for (let attempt = 0; attempt < 100 && revoked.length === 0; attempt++) await Bun.sleep(10); + expect(dispatched).toBe(0); + expect(revoked).toEqual(minted); + expect(minted).toHaveLength(1); +}); + +test('a Code API that cannot mint grants refuses import_file without dispatching', async () => { + let dispatched = 0; + const url = await listen(appFor(createWorkspaceToolsRouter({ + backend: 'remote-bridge', configuredWorkerId: 'worker', dynamicWorkers: false, + store: { async dispatchWorkspaceTool(): Promise { dispatched++; throw new Error('must not dispatch'); } }, + }))); + const response = await post(`${url}/workspace-tools/execute`, importBody()); + expect(response.status).toBe(503); + expect(await response.json()).toMatchObject({ code: 'IMPORT_UNAVAILABLE' }); + expect(dispatched).toBe(0); +}); + +test('a durable import fingerprints the caller request and mints its grant once', async () => { + await upload(redis, owner, SESSION, OBJECT); + const bridge = new RedisBridgeStore(redis); + await bridge.register({ protocolVersion: 1, workerId: 'worker', incarnationId: 'incarnation-00000001', capabilities: { + sandboxProfile: 'anthropic-srt', statefulWorkspace: false, runtimes: [], + workspaceTools: { protocolVersion: 1, operations: ['read_file', 'import_file'], workspaces: [{ id: 'primary' }] }, + } }); + const minted: string[] = []; + const revoked: string[] = []; + const url = await listen(appFor(createWorkspaceToolsRouter({ + backend: 'remote-bridge', configuredWorkerId: 'worker', dynamicWorkers: false, + imports: importService(redis, minted, revoked), + requests: new RedisWorkspaceRequests(redis, bridge), + store: { async dispatchWorkspaceTool(): Promise { throw new Error('durable requests never dispatch synchronously'); } }, + }))); + const headers = { 'X-LibreChat-Workspace-Request-Id': 'import-request-000001' }; + for (let attempt = 0; attempt < 2; attempt++) { + const response = await post(`${url}/workspace-tools/requests`, importBody(), headers); + expect(response.status).toBe(202); + expect(await response.json()).toMatchObject({ state: 'queued', workspaceId: 'primary' }); + } + expect(minted).toHaveLength(1); + + // A retry of the same request answers from its record even once the upload is gone. + await redis.del(`session:${SESSION}`); + const retried = await post(`${url}/workspace-tools/requests`, importBody(), headers); + expect(retried.status).toBe(202); + expect(minted).toHaveLength(1); + + const keys = (await redis.keys('codeapi:workspace-requests:v1:*')).filter(key => !key.endsWith(':active')); + const records = await Promise.all(keys.map(key => redis.hget(key, 'record'))); + const stored = records.flatMap(record => record == null ? [] : [JSON.parse(record).request as WorkspaceImportFileRequest]); + expect(stored.length).toBeGreaterThan(0); + for (const request of stored) { + expect(openEgressGrant(request.egressGrant!, SECRET).grant_id).toBe(minted[0]); + expect(request.source.id).not.toBe(OBJECT); + } + + // Racing submissions of one new request ID: only the adopted grant survives. + const raceHeaders = { 'X-LibreChat-Workspace-Request-Id': 'import-request-000002' }; + await upload(redis, owner, SESSION, OBJECT); + const before = minted.length; + const raced = await Promise.all([0, 1].map(() => post(`${url}/workspace-tools/requests`, importBody(), raceHeaders))); + expect(raced.map(response => response.status)).toEqual([202, 202]); + const racedGrants = minted.slice(before); + expect(racedGrants.length).toBeGreaterThan(0); + expect(revoked).toHaveLength(racedGrants.length - 1); + + const mintedBeforeConflict = minted.length; + const conflict = await post(`${url}/workspace-tools/requests`, importBody({ bytes: 1 }), headers); + expect(conflict.status).toBe(409); + expect(minted).toHaveLength(mintedBeforeConflict); +}); diff --git a/service/src/workspace-tools/import.ts b/service/src/workspace-tools/import.ts new file mode 100644 index 00000000..6f22b55b --- /dev/null +++ b/service/src/workspace-tools/import.ts @@ -0,0 +1,182 @@ +import { posix } from 'path'; +import { nanoid } from 'nanoid'; + +import type * as t from '../types'; +import type { GatewayPreparedEgress } from '../egress-gateway-client'; +import type { ExecutionManifestClaims } from '../execution-manifest'; +import type { SessionOwnershipStore } from '../session-ownership'; +import type { WorkspaceImportFileRequest } from '../../../packages/code/src/protocol'; + +import { isValidId } from '../utils'; +import { authorizeSessionOwnership } from '../session-ownership'; +import { resolveSessionKey, SessionKeyResolutionError } from '../session-key'; +import { buildExecutionManifestClaims } from '../execution-manifest-claims'; + +/** One download plus one retry of the same object; nothing else is readable. */ +const IMPORT_GRANT_MAX_REQUESTS = 2; +/** Clock-skew slack beyond the request's own admission and execution budget. */ +const IMPORT_GRANT_TTL_SLACK_SECONDS = 30; + +export type WorkspaceImportErrorCode = + | 'INVALID_WORKSPACE_TOOL_REQUEST' + | 'IMPORT_SOURCE_UNAUTHORIZED' + | 'IMPORT_UNAVAILABLE'; + +export class WorkspaceImportError extends Error { + constructor( + readonly status: 400 | 403 | 503, + readonly code: WorkspaceImportErrorCode, + message: string, + ) { + super(message); + this.name = 'WorkspaceImportError'; + } +} + +/** Redis reads used to prove ownership: the session cache key and the upload record. */ +export interface WorkspaceImportSourceStore extends Pick { + exists(key: string): Promise; +} + +export type WorkspaceImportGrantFactory = (args: { + payload: t.PayloadBody; + claims: ExecutionManifestClaims; +}) => Promise; + +export interface PreparedWorkspaceImport { + grantId: string; + request: WorkspaceImportFileRequest & { egressGrant: string }; +} + +export interface WorkspaceImportService { + /** Largest file Code API relays, checked before ownership or dispatch. */ + readonly maxBytes: number; + /** Proves the caller owns `source` and returns the storage session key that owns it. */ + authorize(req: t.AuthenticatedRequest, request: WorkspaceImportFileRequest): Promise; + /** Mints the one-object grant and rewrites `source` into its opaque handles. */ + prepare(args: { + req: t.AuthenticatedRequest; + request: WorkspaceImportFileRequest; + sessionKey: string; + ttlSeconds: number; + }): Promise; + revoke(grantId: string): Promise; +} + +const unauthorized = (): WorkspaceImportError => + new WorkspaceImportError(403, 'IMPORT_SOURCE_UNAUTHORIZED', 'Unauthorized import source'); + +/** + * Applies the `GET /download/:session_id/:fileId` rule for a user-kind object: + * the live session cache key must name the caller's own storage key, and the + * object must have been uploaded under that key and session. + */ +export async function authorizeWorkspaceImportSource(args: { + req: t.AuthenticatedRequest; + source: WorkspaceImportFileRequest['source']; + store: WorkspaceImportSourceStore; +}): Promise { + const { req, source, store } = args; + if (!isValidId(source.storage_session_id) || !isValidId(source.id)) { + throw new WorkspaceImportError(400, 'INVALID_WORKSPACE_TOOL_REQUEST', 'Invalid import source'); + } + const userId = req.codeApiAuthContext?.userId; + if (!userId) throw unauthorized(); + let sessionKey: string; + try { + sessionKey = resolveSessionKey(req, { kind: 'user', id: userId }); + } catch (error) { + if (error instanceof SessionKeyResolutionError && error.status === 400) throw unauthorized(); + throw error; + } + const ownership = await authorizeSessionOwnership(store, { + session_id: source.storage_session_id, + expectedSessionKey: sessionKey, + allowExpiredCache: false, + }); + if (!ownership.authorized) throw unauthorized(); + const uploaded = await store.exists(`upload:${sessionKey}${source.storage_session_id}${source.id}`); + if (uploaded !== 1) throw unauthorized(); + return sessionKey; +} + +/** Claims for a read-only grant over exactly one object, with no upload capacity. */ +export function workspaceImportGrantClaims(args: { + req: t.AuthenticatedRequest; + request: WorkspaceImportFileRequest; + sessionKey: string; + ttlSeconds: number; + nowSeconds?: number; +}): { payload: t.PayloadBody; claims: ExecutionManifestClaims } { + const now = args.nowSeconds ?? Math.floor(Date.now() / 1000); + const payload: t.PayloadBody = { + language: 'workspace-import', + version: '1', + files: [{ + id: args.request.source.id, + storage_session_id: args.request.source.storage_session_id, + name: posix.basename(args.request.path), + }], + }; + const claims = buildExecutionManifestClaims({ + req: args.req, + executionId: `import_${nanoid()}`, + userId: args.req.codeApiAuthContext?.userId ?? '', + sessionKey: args.sessionKey, + outputSessionId: `import_${nanoid()}`, + payload, + nowSeconds: now, + }); + return { + payload, + claims: { + ...claims, + max_upload_bytes: 0, + max_output_files: 0, + max_requests: IMPORT_GRANT_MAX_REQUESTS, + tool_call_socket: false, + exp: now + Math.max(1, Math.ceil(args.ttlSeconds)) + IMPORT_GRANT_TTL_SLACK_SECONDS, + }, + }; +} + +function isPayloadFileRef(file: t.PayloadBody['files'][number] | undefined): file is t.PayloadFileRef { + return file != null && 'id' in file && typeof file.id === 'string' && + typeof file.storage_session_id === 'string'; +} + +export function createWorkspaceImportService(options: { + maxBytes: number; + store: WorkspaceImportSourceStore; + createGrant: WorkspaceImportGrantFactory; + revokeGrant: (grantId: string) => Promise; +}): WorkspaceImportService { + if (!Number.isSafeInteger(options.maxBytes) || options.maxBytes < 1) { + throw new RangeError('Workspace import file limit must be a positive safe integer'); + } + return { + maxBytes: options.maxBytes, + authorize: (req, request) => authorizeWorkspaceImportSource({ req, source: request.source, store: options.store }), + async prepare({ req, request, sessionKey, ttlSeconds }) { + let prepared: GatewayPreparedEgress; + try { + prepared = await options.createGrant(workspaceImportGrantClaims({ req, request, sessionKey, ttlSeconds })); + } catch { + throw new WorkspaceImportError(503, 'IMPORT_UNAVAILABLE', 'Import source grant could not be issued'); + } + const [file] = prepared.payload.files; + if (!isPayloadFileRef(file) || prepared.payload.files.length !== 1 || !prepared.egressGrantToken) { + throw new WorkspaceImportError(503, 'IMPORT_UNAVAILABLE', 'Import source grant is invalid'); + } + return { + grantId: prepared.grant_id, + request: { + ...request, + source: { storage_session_id: file.storage_session_id, id: file.id }, + egressGrant: prepared.egressGrantToken, + }, + }; + }, + revoke: options.revokeGrant, + }; +} diff --git a/service/src/workspace-tools/index.ts b/service/src/workspace-tools/index.ts index 3c7caee2..58c043d9 100644 --- a/service/src/workspace-tools/index.ts +++ b/service/src/workspace-tools/index.ts @@ -9,7 +9,9 @@ import { connection } from '../queue'; import logger from '../logger'; import { checkServiceStartUp, checkServiceShutDown } from '../lifecycle'; import { withSpan } from '../telemetry'; -import { isBridgeEnabled } from '../bridge/enabled'; +import { isBridgeEnabled, isWorkspaceImportEnabled } from '../bridge/enabled'; +import { createWorkspaceImportService } from './import'; +import { createGatewayEgressGrant, revokeGatewayEgressGrant } from '../egress-gateway-client'; import { workspaceRequestCoordinationScope } from './coordination'; const router = Router(); @@ -53,6 +55,14 @@ router.use(['/workspace-tools/execute', '/workspace-tools/requests'], (req, res, if (req.method === 'POST') executionLimiter(req, res, next); else next(); }); +const imports = isWorkspaceImportEnabled() + ? createWorkspaceImportService({ + maxBytes: env.BRIDGE_IMPORT_FILE_MAX_BYTES, + store: connection, + createGrant: args => createGatewayEgressGrant(args), + revokeGrant: grantId => revokeGatewayEgressGrant({ grantId, reason: 'completed' }), + }) + : undefined; router.use( createWorkspaceToolsRouter({ store: bridgeStore, @@ -61,6 +71,7 @@ router.use( configuredWorkerId: env.BRIDGE_WORKER_ID, dynamicWorkers: env.BRIDGE_DYNAMIC_WORKERS, timeoutMs: env.JOB_TIMEOUT, + imports, }), ); diff --git a/service/src/workspace-tools/requests.ts b/service/src/workspace-tools/requests.ts index 2988c7cc..902651fc 100644 --- a/service/src/workspace-tools/requests.ts +++ b/service/src/workspace-tools/requests.ts @@ -99,6 +99,13 @@ export class RedisWorkspaceRequests { request: WorkspaceToolRequest; queueWaitMs: number; executionTimeoutMs: number; + /** + * Derives the request a worker receives, such as an import carrying its + * egress grant. Runs only for a new request ID; the fingerprint, and so + * idempotent retries, cover the caller's request. `release` undoes the + * preparation when the queue does not adopt this submission. + */ + prepare?: () => Promise<{ request: WorkspaceToolRequest; release?: () => Promise }>; }): Promise { if (!WORKSPACE_REQUEST_ID_PATTERN.test(args.requestId)) { throw new BridgeStoreError('ASSIGNMENT_INVALID', 'Invalid durable workspace request ID'); @@ -117,21 +124,31 @@ export class RedisWorkspaceRequests { const registration = await this.bridge.prepareDurableWorkspaceTool({ ...args, tenantId: args.owner.tenantId, }); + const prepared = args.prepare == null ? undefined : await args.prepare(); + const release = async (): Promise => { await prepared?.release?.().catch(() => undefined); }; + const request = prepared?.request ?? args.request; const now = Date.now(); const record: StoredWorkspaceRequest = { id: `durable-${randomBytes(18).toString('base64url')}`, key, fingerprint, workerId: args.workerId, tenantId: args.owner.tenantId, requireTenantBinding: args.requireTenantBinding, - request: args.request, registration, createdAtMs: now, + request, registration, createdAtMs: now, queueDeadlineAtMs: now + args.queueWaitMs, executionTimeoutMs: args.executionTimeoutMs, state: 'queued', }; const queue = new BridgeAdmissionQueue(this.redis); - const accepted = await queue.submit({ - workerId: record.workerId, id: record.id, deadlineAtMs: record.queueDeadlineAtMs, - workspaceId: workspaceAdmissionId(record.request.workspaceId, record.request.workspaceInstanceId, record.request.worktree), - key, activeKey: this.activeKey, fingerprint, record: JSON.stringify(record), retentionMs: RETENTION_MS, - }); + let accepted: Awaited>; + try { + accepted = await queue.submit({ + workerId: record.workerId, id: record.id, deadlineAtMs: record.queueDeadlineAtMs, + workspaceId: workspaceAdmissionId(record.request.workspaceId, record.request.workspaceInstanceId, record.request.worktree), + key, activeKey: this.activeKey, fingerprint, record: JSON.stringify(record), retentionMs: RETENTION_MS, + }); + } catch (error) { + await release(); + throw error; + } + if (accepted !== 'accepted') await release(); if (accepted === 'conflict') throw new WorkspaceRequestConflict('Request ID was already used for different work'); if (accepted === 'full') throw new BridgeStoreError('WORKER_QUEUE_FULL', 'Bridge worker pending request limit reached'); const submitted = (await this.read(key))!; diff --git a/service/src/workspace-tools/router.ts b/service/src/workspace-tools/router.ts index d4d65d14..bf31e272 100644 --- a/service/src/workspace-tools/router.ts +++ b/service/src/workspace-tools/router.ts @@ -4,14 +4,19 @@ import type { RequestHandler, Response } from 'express'; import type { AuthenticatedRequest } from '../types'; import type { RedisBridgeStore } from '../bridge/store'; import type { WorkspaceToolRequest } from '../../../packages/code/src/protocol'; +import type { PreparedWorkspaceImport, WorkspaceImportService } from './import'; +import logger from '../logger'; import { getWorkspaceToolOutcome } from './outcome'; +import { WorkspaceImportError } from './import'; import { getPrincipalOrReject } from '../auth/principal'; import { BridgeStoreError } from '../bridge/store'; import { checkServiceShutDown } from '../lifecycle'; import { isWorkspaceToolRequest, + isBridgeWorkspaceImportFileRequest, BRIDGE_WORKSPACE_COMMAND_DEFAULT_TIMEOUT_MS, + BRIDGE_WORKSPACE_IMPORT_FILE_TIMEOUT_MS, } from '../../../packages/code/src/protocol'; import { CODEAPI_BRIDGE_WORKER_HEADER, @@ -35,6 +40,8 @@ interface WorkspaceToolsRouterOptions { timeoutMs?: number; queueTimeoutMs?: number; isShuttingDown?: () => boolean; + /** Present only when this Code API can mint egress grants for `import_file`. */ + imports?: WorkspaceImportService; } function asyncRoute(handler: (req: AuthenticatedRequest, res: Response) => Promise): RequestHandler { @@ -82,13 +89,28 @@ export function createWorkspaceToolsRouter(options: WorkspaceToolsRouterOptions) res.status(503).json({ error: 'Service is shutting down' }); return; } - if (!isWorkspaceToolRequest(req.body)) { + if ( + !isWorkspaceToolRequest(req.body) || + (req.body.operation === 'import_file' && !isBridgeWorkspaceImportFileRequest(req.body)) + ) { outcome.errorCode = 'INVALID_WORKSPACE_TOOL_REQUEST'; res.status(400).json({ error: 'Invalid workspace tool request', }); return; } + if (req.body.operation === 'import_file') { + if (options.imports == null) { + outcome.errorCode = 'IMPORT_UNAVAILABLE'; + res.status(503).json({ error: 'Workspace file import is not available on this Code API', code: 'IMPORT_UNAVAILABLE' }); + return; + } + if (req.body.bytes > options.imports.maxBytes) { + outcome.errorCode = 'WRITE_LIMIT_EXCEEDED'; + res.status(413).json({ error: 'Imported file exceeds the Code API file limit', code: 'WRITE_LIMIT_EXCEEDED' }); + return; + } + } const advertisedQueueWait = req.header(WORKSPACE_QUEUE_WAIT_HEADER); if (advertisedQueueWait !== undefined && !/^[1-9]\d*$/.test(advertisedQueueWait)) { outcome.errorCode = 'INVALID_WORKSPACE_QUEUE_WAIT'; @@ -121,9 +143,14 @@ export function createWorkspaceToolsRouter(options: WorkspaceToolsRouterOptions) options.timeoutMs ?? Number.MAX_SAFE_INTEGER, ) } : principalRequest; - const executionBudgetMs = request.operation === 'execute_command' - ? request.timeoutMs! + 5_000 - : Math.min(options.timeoutMs ?? 30_000, 30_000); + let executionBudgetMs = Math.min(options.timeoutMs ?? 30_000, 30_000); + if (request.operation === 'execute_command') executionBudgetMs = request.timeoutMs! + 5_000; + if (request.operation === 'import_file') { + executionBudgetMs = Math.min( + options.timeoutMs ?? BRIDGE_WORKSPACE_IMPORT_FILE_TIMEOUT_MS, + BRIDGE_WORKSPACE_IMPORT_FILE_TIMEOUT_MS, + ); + } outcome.deadlineBudgetMs = queueBudgetMs + executionBudgetMs; let selection: { workerId: string; explicit: boolean } | undefined; @@ -152,20 +179,47 @@ export function createWorkspaceToolsRouter(options: WorkspaceToolsRouterOptions) } outcome.workerId = selection.workerId; + const sendImportError = (error: unknown): boolean => { + if (!(error instanceof WorkspaceImportError)) return false; + outcome.errorCode = error.code; + res.status(error.status).json({ error: error.message, code: error.code }); + return true; + }; + const imports = options.imports; + const importRequest = request.operation === 'import_file' ? request : undefined; + /** Ownership and the grant are settled only for work about to be queued or dispatched, + * so a durable retry of an existing request never depends on the upload still existing. */ + const prepareImport = imports == null || importRequest == null + ? undefined + : async (): Promise => { + const sessionKey = await imports.authorize(req, importRequest); + // The grant must outlive admission and execution, never more. + const ttlSeconds = Math.ceil((queueBudgetMs + executionBudgetMs) / 1000); + return imports.prepare({ req, request: importRequest, sessionKey, ttlSeconds }); + }; + if (durable) { if (options.requests == null) { res.status(404).json({ error: 'Durable workspace requests are unavailable' }); return; } try { + const prepare = prepareImport; const status = await options.requests.submit({ owner: principal, requestId: req.header('X-LibreChat-Workspace-Request-Id') ?? '', workerId: selection.workerId, requireTenantBinding: selection.explicit && (options.dynamicWorkers || selection.workerId !== options.configuredWorkerId), request, queueWaitMs: queueBudgetMs, executionTimeoutMs: executionBudgetMs, + ...(prepare == null || imports == null ? {} : { + prepare: async () => { + const prepared = await prepare(); + return { request: prepared.request, release: () => imports.revoke(prepared.grantId) }; + }, + }), }); res.status(202).json(status); } catch (error) { + if (sendImportError(error)) return; if (error instanceof WorkspaceRequestConflict) { res.status(409).json({ error: error.message, code: 'REQUEST_CONFLICT' }); } else if (error instanceof BridgeStoreError) { @@ -182,7 +236,24 @@ export function createWorkspaceToolsRouter(options: WorkspaceToolsRouterOptions) if (!res.writableEnded) abort(); }; res.once('close', abortClosedResponse); + // Events that fired before these listeners existed are visible only on the socket. + if (req.socket?.destroyed === true) abort(); + let preparedImport: PreparedWorkspaceImport | undefined; try { + if (prepareImport != null) { + try { + preparedImport = await prepareImport(); + } catch (error) { + if (sendImportError(error)) return; + throw error; + } + // A caller that left while ownership or the grant was settled must not start a mutation. + if (controller.signal.aborted) { + outcome.errorCode = 'CLIENT_DISCONNECTED'; + return; + } + } + const dispatchRequest: WorkspaceToolRequest = preparedImport?.request ?? request; outcome.dispatchPending = true; const dispatchStartedAt = performance.now(); const settlement = await options.store.dispatchWorkspaceTool({ @@ -190,7 +261,7 @@ export function createWorkspaceToolsRouter(options: WorkspaceToolsRouterOptions) tenantId: principal.tenantId, requireTenantBinding: selection.explicit && (options.dynamicWorkers || selection.workerId !== options.configuredWorkerId), - request, + request: dispatchRequest, deadlineAtMs: Date.now() + queueBudgetMs, executionTimeoutMs: executionBudgetMs, signal: controller.signal, @@ -243,6 +314,12 @@ export function createWorkspaceToolsRouter(options: WorkspaceToolsRouterOptions) outcome.flush(); req.removeListener('aborted', abort); res.removeListener('close', abortClosedResponse); + if (preparedImport != null && options.imports != null) { + const { grantId } = preparedImport; + void options.imports.revoke(grantId).catch(error => { + logger.warn('Failed to revoke workspace import grant', { grantId, error }); + }); + } } }); router.post('/workspace-tools/execute', execute(false)); diff --git a/tests/compose-bridge-config.cjs b/tests/compose-bridge-config.cjs index bafea2fe..79268b52 100644 --- a/tests/compose-bridge-config.cjs +++ b/tests/compose-bridge-config.cjs @@ -16,6 +16,7 @@ function render(overrides) { CODEAPI_BRIDGE_TOKEN: '', CODEAPI_BRIDGE_MAX_WORKSPACE_LEASE_SLOTS: '', CODEAPI_BRIDGE_LANE_GIT: '', + CODEAPI_BRIDGE_IMPORT_FILE_MAX_BYTES: '', CODEAPI_BRIDGE_RECOVERY_SERVER_ID: '', CODEAPI_BRIDGE_ENROLLMENT_TTL_SECONDS: '', CODEAPI_BRIDGE_RECOVERY_CHALLENGE_TTL_SECONDS: '', @@ -36,6 +37,7 @@ for (const overrides of [ CODEAPI_BRIDGE_WORKER_ID: 'test-worker', CODEAPI_BRIDGE_MAX_WORKSPACE_LEASE_SLOTS: '4', CODEAPI_BRIDGE_LANE_GIT: 'true', + CODEAPI_BRIDGE_IMPORT_FILE_MAX_BYTES: '1048576', CODEAPI_BRIDGE_RECOVERY_SERVER_ID: 'https://code.example.test', CODEAPI_BRIDGE_ENROLLMENT_TTL_SECONDS: '86400', CODEAPI_BRIDGE_RECOVERY_CHALLENGE_TTL_SECONDS: '90', @@ -61,6 +63,10 @@ for (const overrides of [ assert.equal(env.CODEAPI_BRIDGE_RECOVERY_MAX_ATTEMPTS_PER_MINUTE, overrides.CODEAPI_BRIDGE_RECOVERY_MAX_ATTEMPTS_PER_MINUTE ?? '30'); assert.equal(env.CODEAPI_BRIDGE_RECOVERY_MAX_UNTRUSTED_PER_MINUTE, overrides.CODEAPI_BRIDGE_RECOVERY_MAX_UNTRUSTED_PER_MINUTE ?? '240'); } + assert.equal( + config.services.api.environment.CODEAPI_BRIDGE_IMPORT_FILE_MAX_BYTES, + overrides.CODEAPI_BRIDGE_IMPORT_FILE_MAX_BYTES ?? '104857600', + ); for (const name of ['egress_gateway', 'sandbox-runner']) { assert.equal(config.services[name].environment.CODEAPI_BRIDGE_TOKEN, undefined); assert.equal(config.services[name].environment.CODEAPI_BRIDGE_RECOVERY_SERVER_ID, undefined);