diff --git a/packages/code/README.md b/packages/code/README.md index 3e8bf72d..f53741a8 100644 --- a/packages/code/README.md +++ b/packages/code/README.md @@ -271,8 +271,19 @@ names and exposes bounded `read_file`, literal `search_text`, and deterministic can explicitly add confined `write_file` and exact-match `edit_file` operations with `--allow-workspace-writes` or `LIBRECHAT_CODE_ALLOW_WORKSPACE_WRITES=true`. -Only IDs, names, protocol version, and supported operations appear in worker -capabilities; absolute host paths remain local to the worker process. +`write_file` preserves its overwrite behavior by default; callers can set +`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`. +`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 +the worker and server negotiate `batch` in `editFileModes`. +Revision-fenced edits likewise require the negotiated +`expected_base_sha256` entry in `editFileFeatures`. +Only IDs, names, protocol version, supported operations, and negotiated write +modes appear in worker capabilities; absolute host paths remain local to the +worker process. The protocol also defines a bounded `execute_command` request and result for a sandbox-backed executor. Commands are treated as workspace mutations and cannot @@ -320,9 +331,14 @@ bounded set of ignored-aware candidates with configuration and symlink following disabled. It then opens and verifies each candidate through the same confined 1 MiB read boundary before matching locally. File listing invokes `rg` without a shell, with configuration and symlink following disabled. Both operations -stop after bounded global result counts. The worker process still belongs inside -the trusted BYOM boundary and should receive filesystem access only to roots the -operator intentionally registers. +stop after bounded global result counts. A truncated `list_files` result includes +`nextAfterPath`; pass that value back as `afterPath` with the same workspace and +path to continue deterministically beyond the 500-file protocol ceiling. +Continuation is advertised and negotiated as the `after_path` list-file feature, +so mixed Code API and worker versions keep the legacy bounded response shape +during rolling upgrades. The worker process still belongs inside the trusted +BYOM boundary and should receive filesystem access only to roots the operator +intentionally registers. Writes are limited to 1 MiB of UTF-8 text and require an existing directory inside the registered root. They reject traversal, symlink targets, and diff --git a/packages/code/src/protocol.test.ts b/packages/code/src/protocol.test.ts index 26a9928f..fd426783 100644 --- a/packages/code/src/protocol.test.ts +++ b/packages/code/src/protocol.test.ts @@ -2,11 +2,66 @@ import assert from 'node:assert/strict'; import test from 'node:test'; import { bridgeWorkerPath, + comparePortableRelativePaths, isValidBridgeWorkerCapabilities, isValidBridgeWorkerId, isWorkspaceToolRequest, isWorkspaceToolResult, } from './protocol.js'; +import type { + WorkspaceEditFileRequest, + WorkspacePreviewEditRequest, +} from './protocol.js'; + +const validSingleEditRequest: WorkspaceEditFileRequest = { + protocolVersion: 1, + operation: 'edit_file', + workspaceId: 'primary', + path: 'notes.txt', + oldText: 'before', + newText: 'after', +}; +const validBatchEditRequest: WorkspaceEditFileRequest = { + protocolVersion: 1, + operation: 'edit_file', + workspaceId: 'primary', + path: 'notes.txt', + edits: [{ oldText: 'before', newText: 'after' }], +}; +// @ts-expect-error An edit request must choose a complete single or batch form. +const invalidEmptyEditRequest: WorkspaceEditFileRequest = { + protocolVersion: 1, + operation: 'edit_file', + workspaceId: 'primary', + path: 'notes.txt', +}; +// @ts-expect-error Single and batch edit forms are mutually exclusive. +const invalidMixedEditRequest: WorkspaceEditFileRequest = { + ...validSingleEditRequest, + edits: validBatchEditRequest.edits, +}; +void invalidEmptyEditRequest; +void invalidMixedEditRequest; + +// @ts-expect-error A preview request must choose a complete single or batch form. +const invalidEmptyPreviewRequest: WorkspacePreviewEditRequest = { + protocolVersion: 1, + operation: 'preview_edit', + workspaceId: 'primary', + path: 'notes.txt', +}; +// @ts-expect-error Single and batch preview forms are mutually exclusive. +const invalidMixedPreviewRequest: WorkspacePreviewEditRequest = { + protocolVersion: 1, + operation: 'preview_edit', + workspaceId: 'primary', + path: 'notes.txt', + oldText: 'before', + newText: 'after', + edits: [{ oldText: 'before', newText: 'after' }], +}; +void invalidEmptyPreviewRequest; +void invalidMixedPreviewRequest; test('bridgeWorkerPath encodes worker-controlled path segments', () => { assert.equal( @@ -70,6 +125,91 @@ test('bridge worker capabilities accept only bounded public workspace descriptor }; assert.equal(isValidBridgeWorkerCapabilities(valid), true); + assert.equal( + isValidBridgeWorkerCapabilities({ + ...valid, + workspaceTools: { + ...valid.workspaceTools, + operations: ['read_file', 'write_file'], + writeFileModes: ['replace', 'create'], + }, + }), + true, + ); + assert.equal( + isValidBridgeWorkerCapabilities({ + ...valid, + workspaceTools: { + ...valid.workspaceTools, + operations: ['read_file', 'preview_edit'], + editFileModes: ['single', 'batch'], + }, + }), + true, + ); + assert.equal( + isValidBridgeWorkerCapabilities({ + ...valid, + workspaceTools: { + ...valid.workspaceTools, + operations: ['read_file', 'edit_file'], + editFileModes: ['single', 'batch'], + editFileFeatures: ['expected_base_sha256'], + }, + }), + true, + ); + assert.equal( + isValidBridgeWorkerCapabilities({ + ...valid, + workspaceTools: { + ...valid.workspaceTools, + editFileFeatures: ['expected_base_sha256'], + }, + }), + false, + ); + assert.equal( + isValidBridgeWorkerCapabilities({ + ...valid, + workspaceTools: { + ...valid.workspaceTools, + operations: ['read_file', 'list_files'], + listFileFeatures: ['after_path'], + }, + }), + true, + ); + assert.equal( + isValidBridgeWorkerCapabilities({ + ...valid, + workspaceTools: { + ...valid.workspaceTools, + listFileFeatures: ['after_path'], + }, + }), + false, + ); + assert.equal( + isValidBridgeWorkerCapabilities({ + ...valid, + workspaceTools: { + ...valid.workspaceTools, + editFileModes: ['batch'], + }, + }), + false, + ); + assert.equal( + isValidBridgeWorkerCapabilities({ + ...valid, + workspaceTools: { + ...valid.workspaceTools, + writeFileModes: ['create'], + }, + }), + false, + ); assert.equal( isValidBridgeWorkerCapabilities({ ...valid, @@ -99,6 +239,7 @@ test('workspace file listing accepts only bounded portable requests and results' workspaceId: 'primary', path: 'src', maxResults: 20, + afterPath: 'src/app.ts', }; assert.equal(isWorkspaceToolRequest(request), true); assert.equal( @@ -106,15 +247,65 @@ test('workspace file listing accepts only bounded portable requests and results' false, ); assert.equal(isWorkspaceToolRequest({ ...request, maxResults: 501 }), false); + assert.equal( + isWorkspaceToolRequest({ ...request, afterPath: 'outside/app.ts' }), + false, + ); const result = { protocolVersion: 1 as const, operation: 'list_files' as const, workspaceId: 'primary', - paths: ['src/app.ts', 'src/worker.ts'], - truncated: false, + paths: ['src/worker.ts', 'src/z.ts'], + truncated: true, + nextAfterPath: 'src/z.ts', }; assert.equal(isWorkspaceToolResult(request, result), true); + assert.equal( + isWorkspaceToolResult( + { ...request, afterPath: undefined }, + { + ...result, + paths: ['src//z.ts', 'src/worker.ts'], + nextAfterPath: undefined, + }, + {}, + ), + true, + ); + assert.equal( + isWorkspaceToolResult( + { ...request, afterPath: undefined }, + { + ...result, + paths: ['src//z.ts', 'src/worker.ts'], + nextAfterPath: undefined, + }, + { listFileFeatures: ['after_path'] }, + ), + false, + ); + assert.equal( + isWorkspaceToolResult(request, { + ...result, + paths: ['src/\uE000.ts', 'src/\u{10000}.ts'], + truncated: false, + nextAfterPath: undefined, + }), + true, + ); + assert.equal( + isWorkspaceToolResult(request, { ...result, nextAfterPath: 'src/worker.ts' }), + false, + ); + assert.equal( + isWorkspaceToolResult(request, { + ...result, + paths: ['src/z.ts', 'src/worker.ts'], + nextAfterPath: 'src/worker.ts', + }), + false, + ); assert.equal( isWorkspaceToolResult(request, { ...result, @@ -140,6 +331,14 @@ test('workspace file listing accepts only bounded portable requests and results' }), false, ); + assert.equal( + isWorkspaceToolResult(request, { + ...result, + paths: ['src//worker.ts'], + nextAfterPath: 'src//worker.ts', + }), + false, + ); assert.equal( isWorkspaceToolResult(request, { ...result, @@ -149,6 +348,12 @@ test('workspace file listing accepts only bounded portable requests and results' ); }); +test('workspace path ordering matches sorted depth-first traversal', () => { + assert.ok(comparePortableRelativePaths('src/app.ts', 'src.ts') < 0); + assert.ok(comparePortableRelativePaths('src/app.ts', 'src/worker.ts') < 0); + assert.ok(comparePortableRelativePaths('src.ts', 'src/app.ts') > 0); +}); + test('workspace mutations accept bounded UTF-8 requests and exact result shapes', () => { const writeRequest = { protocolVersion: 1 as const, @@ -156,8 +361,13 @@ test('workspace mutations accept bounded UTF-8 requests and exact result shapes' workspaceId: 'primary', path: 'notes.txt', content: 'hello', + overwrite: false, }; assert.equal(isWorkspaceToolRequest(writeRequest), true); + assert.equal( + isWorkspaceToolRequest({ ...writeRequest, overwrite: 'false' }), + false, + ); assert.equal( isWorkspaceToolRequest({ ...writeRequest, @@ -176,6 +386,17 @@ test('workspace mutations accept bounded UTF-8 requests and exact result shapes' }), true, ); + assert.equal( + isWorkspaceToolResult(writeRequest, { + protocolVersion: 1, + operation: 'write_file', + workspaceId: 'primary', + path: 'notes.txt', + created: false, + bytesWritten: 5, + }), + false, + ); const editRequest = { protocolVersion: 1 as const, @@ -198,6 +419,80 @@ test('workspace mutations accept bounded UTF-8 requests and exact result shapes' }), true, ); + const batchEditRequest = { + protocolVersion: 1 as const, + operation: 'edit_file' as const, + workspaceId: 'primary', + path: 'notes.txt', + edits: [ + { oldText: 'hello', newText: 'goodbye' }, + { oldText: 'world', newText: 'BYOM' }, + ], + }; + assert.equal(isWorkspaceToolRequest(batchEditRequest), true); + assert.equal( + isWorkspaceToolRequest({ ...batchEditRequest, oldText: 'mixed' }), + false, + ); + assert.equal(isWorkspaceToolRequest({ ...batchEditRequest, edits: [] }), false); + assert.equal( + isWorkspaceToolRequest({ + ...batchEditRequest, + edits: Array.from({ length: 101 }, () => ({ oldText: 'a', newText: 'b' })), + }), + false, + ); + assert.equal( + isWorkspaceToolRequest({ + ...batchEditRequest, + edits: [{ oldText: 'a'.repeat(600_000), newText: 'b'.repeat(600_000) }], + }), + false, + ); + assert.equal( + isWorkspaceToolResult(batchEditRequest, { + protocolVersion: 1, + operation: 'edit_file', + workspaceId: 'primary', + path: 'notes.txt', + replacements: 2, + bytesWritten: 12, + }), + true, + ); + const previewRequest = { + ...batchEditRequest, + operation: 'preview_edit' as const, + }; + assert.equal(isWorkspaceToolRequest(previewRequest), true); + assert.equal( + isWorkspaceToolResult(previewRequest, { + protocolVersion: 1, + operation: 'preview_edit', + workspaceId: 'primary', + path: 'notes.txt', + content: 'goodbye BYOM', + hasUtf8Bom: false, + baseSha256: 'a'.repeat(64), + replacements: 2, + bytesWritten: 12, + }), + true, + ); + assert.equal( + isWorkspaceToolRequest({ + ...editRequest, + expectedBaseSha256: 'b'.repeat(64), + }), + true, + ); + assert.equal( + isWorkspaceToolRequest({ + ...editRequest, + expectedBaseSha256: 'not-a-sha', + }), + false, + ); }); test('workspace commands require bounded sandbox inputs and outputs', () => { diff --git a/packages/code/src/protocol.ts b/packages/code/src/protocol.ts index 28e99f90..3bbd83cb 100644 --- a/packages/code/src/protocol.ts +++ b/packages/code/src/protocol.ts @@ -9,6 +9,7 @@ export const BRIDGE_WORKSPACE_PATH_MAX_LENGTH = 4096; export const BRIDGE_WORKSPACE_QUERY_MAX_LENGTH = 4096; export const BRIDGE_WORKSPACE_READ_MAX_BYTES = 1024 * 1024; export const BRIDGE_WORKSPACE_WRITE_MAX_BYTES = 1024 * 1024; +export const BRIDGE_WORKSPACE_EDIT_MAX_EDITS = 100; export const BRIDGE_WORKSPACE_READ_MAX_LINES = 500; export const BRIDGE_WORKSPACE_SEARCH_MAX_RESULTS = 200; export const BRIDGE_WORKSPACE_SEARCH_TEXT_MAX_LENGTH = 2000; @@ -27,9 +28,15 @@ export type BridgeWorkspaceToolOperation = | 'search_text' | 'list_files' | 'write_file' + | 'preview_edit' | 'edit_file' | 'execute_command'; +export type WorkspaceWriteFileMode = 'replace' | 'create'; +export type WorkspaceEditFileMode = 'single' | 'batch'; +export type WorkspaceEditFileFeature = 'expected_base_sha256'; +export type WorkspaceListFileFeature = 'after_path'; + export interface BridgeWorkspaceDescriptor { id: string; name?: string; @@ -41,6 +48,14 @@ export interface BridgeWorkspaceToolCapabilities { protocolVersion: BridgeProtocolVersion; operations: BridgeWorkspaceToolOperation[]; workspaces: BridgeWorkspaceDescriptor[]; + /** Omitted by legacy workers, which only accept replacement writes. */ + writeFileModes?: WorkspaceWriteFileMode[]; + /** Omitted by legacy workers, which only accept single exact replacements. */ + editFileModes?: WorkspaceEditFileMode[]; + /** Omitted by workers that cannot fence edits against a preview revision. */ + editFileFeatures?: WorkspaceEditFileFeature[]; + /** Omitted by workers that cannot continue a bounded file listing. */ + listFileFeatures?: WorkspaceListFileFeature[]; } export interface WorkspaceReadFileRequest { @@ -94,6 +109,8 @@ export interface WorkspaceListFilesRequest { workspaceId: string; path?: string; maxResults?: number; + /** Continue strictly after this canonical path from a previous page. */ + afterPath?: string; } export interface WorkspaceListFilesResult { @@ -102,6 +119,8 @@ export interface WorkspaceListFilesResult { workspaceId: string; paths: string[]; truncated: boolean; + /** Last returned path; pass as afterPath to fetch the next page. */ + nextAfterPath?: string; } export interface WorkspaceWriteFileRequest { @@ -110,6 +129,8 @@ export interface WorkspaceWriteFileRequest { workspaceId: string; path: string; content: string; + /** False requires an atomic create and refuses to replace an existing file. */ + overwrite?: boolean; } export interface WorkspaceWriteFileResult { @@ -121,11 +142,37 @@ export interface WorkspaceWriteFileResult { bytesWritten: number; } -export interface WorkspaceEditFileRequest { +interface WorkspaceEditFileRequestBase { protocolVersion: BridgeProtocolVersion; operation: 'edit_file'; workspaceId: string; path: string; + /** Refuses the mutation unless current file bytes match this preview revision. */ + expectedBaseSha256?: string; +} + +export interface WorkspaceSingleEditFileRequest + extends WorkspaceEditFileRequestBase { + /** Legacy single-edit form. */ + oldText: string; + /** Legacy single-edit form. */ + newText: string; + edits?: never; +} + +export interface WorkspaceBatchEditFileRequest + extends WorkspaceEditFileRequestBase { + /** Ordered exact replacements applied atomically as one file mutation. */ + edits: WorkspaceTextEdit[]; + oldText?: never; + newText?: never; +} + +export type WorkspaceEditFileRequest = + | WorkspaceSingleEditFileRequest + | WorkspaceBatchEditFileRequest; + +export interface WorkspaceTextEdit { oldText: string; newText: string; } @@ -135,7 +182,44 @@ export interface WorkspaceEditFileResult { operation: 'edit_file'; workspaceId: string; path: string; - replacements: 1; + replacements: number; + bytesWritten: number; +} + +interface WorkspacePreviewEditRequestBase { + protocolVersion: BridgeProtocolVersion; + operation: 'preview_edit'; + workspaceId: string; + path: string; +} + +export interface WorkspaceSinglePreviewEditRequest + extends WorkspacePreviewEditRequestBase { + oldText: string; + newText: string; + edits?: never; +} + +export interface WorkspaceBatchPreviewEditRequest + extends WorkspacePreviewEditRequestBase { + edits: WorkspaceTextEdit[]; + oldText?: never; + newText?: never; +} + +export type WorkspacePreviewEditRequest = + | WorkspaceSinglePreviewEditRequest + | WorkspaceBatchPreviewEditRequest; + +export interface WorkspacePreviewEditResult { + protocolVersion: BridgeProtocolVersion; + operation: 'preview_edit'; + workspaceId: string; + path: string; + content: string; + hasUtf8Bom: boolean; + baseSha256: string; + replacements: number; bytesWritten: number; } @@ -169,6 +253,7 @@ export type WorkspaceToolRequest = | WorkspaceSearchTextRequest | WorkspaceListFilesRequest | WorkspaceWriteFileRequest + | WorkspacePreviewEditRequest | WorkspaceEditFileRequest | WorkspaceExecuteCommandRequest; export type WorkspaceToolResult = @@ -176,6 +261,7 @@ export type WorkspaceToolResult = | WorkspaceSearchTextResult | WorkspaceListFilesResult | WorkspaceWriteFileResult + | WorkspacePreviewEditResult | WorkspaceEditFileResult | WorkspaceExecuteCommandResult; @@ -201,6 +287,7 @@ const WORKSPACE_LIST_REQUEST_KEYS = new Set([ 'workspaceId', 'path', 'maxResults', + 'afterPath', ]); const WORKSPACE_WRITE_REQUEST_KEYS = new Set([ 'protocolVersion', @@ -208,6 +295,7 @@ const WORKSPACE_WRITE_REQUEST_KEYS = new Set([ 'workspaceId', 'path', 'content', + 'overwrite', ]); const WORKSPACE_EDIT_REQUEST_KEYS = new Set([ 'protocolVersion', @@ -216,7 +304,19 @@ const WORKSPACE_EDIT_REQUEST_KEYS = new Set([ 'path', 'oldText', 'newText', + 'edits', + 'expectedBaseSha256', ]); +const WORKSPACE_PREVIEW_EDIT_REQUEST_KEYS = new Set([ + 'protocolVersion', + 'operation', + 'workspaceId', + 'path', + 'oldText', + 'newText', + 'edits', +]); +const WORKSPACE_TEXT_EDIT_KEYS = new Set(['oldText', 'newText']); const WORKSPACE_COMMAND_REQUEST_KEYS = new Set([ 'protocolVersion', 'operation', @@ -250,6 +350,7 @@ const WORKSPACE_LIST_RESULT_KEYS = new Set([ 'workspaceId', 'paths', 'truncated', + 'nextAfterPath', ]); const WORKSPACE_WRITE_RESULT_KEYS = new Set([ 'protocolVersion', @@ -267,6 +368,17 @@ const WORKSPACE_EDIT_RESULT_KEYS = new Set([ 'replacements', 'bytesWritten', ]); +const WORKSPACE_PREVIEW_EDIT_RESULT_KEYS = new Set([ + 'protocolVersion', + 'operation', + 'workspaceId', + 'path', + 'content', + 'hasUtf8Bom', + 'baseSha256', + 'replacements', + 'bytesWritten', +]); const WORKSPACE_COMMAND_RESULT_KEYS = new Set([ 'protocolVersion', 'operation', @@ -311,6 +423,14 @@ export interface BridgeWorkerRegistrationResponse { leaseTtlMs: number; /** Operations this Code API can dispatch after the worker advertises them. */ supportedWorkspaceToolOperations?: BridgeWorkspaceToolOperation[]; + /** Write modes this Code API can safely route to a capability-aware worker. */ + supportedWorkspaceWriteFileModes?: WorkspaceWriteFileMode[]; + /** Edit modes this Code API can safely route to a capability-aware worker. */ + supportedWorkspaceEditFileModes?: WorkspaceEditFileMode[]; + /** Edit features this Code API can safely route to a capability-aware worker. */ + supportedWorkspaceEditFileFeatures?: WorkspaceEditFileFeature[]; + /** Listing features this Code API can safely route to a capability-aware worker. */ + supportedWorkspaceListFileFeatures?: WorkspaceListFileFeature[]; } export interface BridgePairingRedemption { @@ -476,6 +596,26 @@ function normalizePortableRelativePath(value: string): string { ); } +/** Compare path segments in ripgrep's sorted, depth-first traversal order. */ +export function comparePortableRelativePaths(left: string, right: string): number { + const encoder = new TextEncoder(); + const leftSegments = left.split('/'); + const rightSegments = right.split('/'); + const segmentCount = Math.min(leftSegments.length, rightSegments.length); + for (let segmentIndex = 0; segmentIndex < segmentCount; segmentIndex += 1) { + const leftBytes = encoder.encode(leftSegments[segmentIndex]); + const rightBytes = encoder.encode(rightSegments[segmentIndex]); + const byteCount = Math.min(leftBytes.length, rightBytes.length); + for (let byteIndex = 0; byteIndex < byteCount; byteIndex += 1) { + const difference = leftBytes[byteIndex] - rightBytes[byteIndex]; + if (difference !== 0) return difference; + } + const lengthDifference = leftBytes.length - rightBytes.length; + if (lengthDifference !== 0) return lengthDifference; + } + return leftSegments.length - rightSegments.length; +} + function isWithinRequestedPath(candidate: string, requested?: string): boolean { if (requested == null) return true; const normalizedCandidate = normalizePortableRelativePath(candidate); @@ -487,6 +627,55 @@ function isWithinRequestedPath(candidate: string, requested?: string): boolean { ); } +function isValidWorkspaceEditRequest(request: Record): boolean { + const hasBatch = request.edits !== undefined; + if (hasBatch && (request.oldText !== undefined || request.newText !== undefined)) { + return false; + } + const edits = hasBatch + ? request.edits + : [{ oldText: request.oldText, newText: request.newText }]; + if ( + !Array.isArray(edits) || + edits.length < 1 || + edits.length > BRIDGE_WORKSPACE_EDIT_MAX_EDITS + ) { + return false; + } + let totalBytes = 0; + for (const edit of edits) { + if ( + typeof edit !== 'object' || + edit === null || + !hasOnlyKeys(edit as Record, WORKSPACE_TEXT_EDIT_KEYS) + ) { + return false; + } + const candidate = edit as Record; + if ( + typeof candidate.oldText !== 'string' || + candidate.oldText.length === 0 || + Buffer.from(candidate.oldText).toString('utf8') !== candidate.oldText || + typeof candidate.newText !== 'string' || + Buffer.from(candidate.newText).toString('utf8') !== candidate.newText + ) { + return false; + } + const oldBytes = new TextEncoder().encode(candidate.oldText).byteLength; + const newBytes = new TextEncoder().encode(candidate.newText).byteLength; + totalBytes += oldBytes + newBytes; + if ( + (hasBatch && totalBytes > BRIDGE_WORKSPACE_WRITE_MAX_BYTES) || + (!hasBatch && + (oldBytes > BRIDGE_WORKSPACE_WRITE_MAX_BYTES || + newBytes > BRIDGE_WORKSPACE_WRITE_MAX_BYTES)) + ) { + return false; + } + } + return true; +} + function hasOnlyKeys( value: Record, allowed: ReadonlySet, @@ -544,6 +733,10 @@ export function isWorkspaceToolRequest( hasOnlyKeys(request, WORKSPACE_LIST_REQUEST_KEYS) && (request.path === undefined || isSafePortableRelativePath(request.path)) && + (request.afterPath === undefined || + (isSafePortableRelativePath(request.afterPath) && + normalizePortableRelativePath(request.afterPath) === request.afterPath && + isWithinRequestedPath(request.afterPath, request.path))) && (request.maxResults === undefined || (Number.isSafeInteger(request.maxResults) && Number(request.maxResults) >= 1 && @@ -557,22 +750,26 @@ export function isWorkspaceToolRequest( typeof request.content === 'string' && Buffer.from(request.content).toString('utf8') === request.content && new TextEncoder().encode(request.content).byteLength <= - BRIDGE_WORKSPACE_WRITE_MAX_BYTES + BRIDGE_WORKSPACE_WRITE_MAX_BYTES && + (request.overwrite === undefined || + typeof request.overwrite === 'boolean') + ); + } + if (request.operation === 'preview_edit') { + return ( + hasOnlyKeys(request, WORKSPACE_PREVIEW_EDIT_REQUEST_KEYS) && + isSafePortableRelativePath(request.path) && + isValidWorkspaceEditRequest(request) ); } if (request.operation === 'edit_file') { return ( hasOnlyKeys(request, WORKSPACE_EDIT_REQUEST_KEYS) && isSafePortableRelativePath(request.path) && - typeof request.oldText === 'string' && - request.oldText.length > 0 && - Buffer.from(request.oldText).toString('utf8') === request.oldText && - new TextEncoder().encode(request.oldText).byteLength <= - BRIDGE_WORKSPACE_WRITE_MAX_BYTES && - typeof request.newText === 'string' && - Buffer.from(request.newText).toString('utf8') === request.newText && - new TextEncoder().encode(request.newText).byteLength <= - BRIDGE_WORKSPACE_WRITE_MAX_BYTES + (request.expectedBaseSha256 === undefined || + (typeof request.expectedBaseSha256 === 'string' && + /^[a-f0-9]{64}$/.test(request.expectedBaseSha256))) && + isValidWorkspaceEditRequest(request) ); } if (request.operation === 'execute_command') { @@ -603,6 +800,7 @@ export function isWorkspaceToolRequest( export function isWorkspaceToolResult( request: WorkspaceToolRequest, value: unknown, + capabilities?: Pick, ): value is WorkspaceToolResult { if (typeof value !== 'object' || value === null) return false; const result = value as Record; @@ -662,6 +860,14 @@ export function isWorkspaceToolResult( return false; } const normalizedPaths = new Set(); + const normalizedAfterPath = + request.afterPath === undefined + ? undefined + : normalizePortableRelativePath(request.afterPath); + const enforcesPaginationContract = + capabilities === undefined || + capabilities.listFileFeatures?.includes('after_path') === true; + let previousPath = normalizedAfterPath; for (const path of result.paths) { if ( !isSafePortableRelativePath(path) || @@ -670,10 +876,26 @@ export function isWorkspaceToolResult( return false; } const normalizedPath = normalizePortableRelativePath(path); - if (normalizedPaths.has(normalizedPath)) return false; + if ( + normalizedPaths.has(normalizedPath) || + (enforcesPaginationContract && + (normalizedPath !== path || + (previousPath !== undefined && + comparePortableRelativePaths(normalizedPath, previousPath) <= 0))) + ) { + return false; + } normalizedPaths.add(normalizedPath); + previousPath = normalizedPath; } - return true; + if (!enforcesPaginationContract) { + return result.nextAfterPath === undefined; + } + if (result.truncated !== true) return result.nextAfterPath === undefined; + return ( + result.paths.length > 0 && + result.nextAfterPath === result.paths[result.paths.length - 1] + ); } if (request.operation === 'write_file') { @@ -681,6 +903,7 @@ export function isWorkspaceToolResult( hasOnlyKeys(result, WORKSPACE_WRITE_RESULT_KEYS) && result.path === request.path && typeof result.created === 'boolean' && + (request.overwrite !== false || result.created === true) && Number.isSafeInteger(result.bytesWritten) && Number(result.bytesWritten) === new TextEncoder().encode(request.content).byteLength @@ -688,16 +911,37 @@ export function isWorkspaceToolResult( } if (request.operation === 'edit_file') { + const replacements = request.edits?.length ?? 1; return ( hasOnlyKeys(result, WORKSPACE_EDIT_RESULT_KEYS) && result.path === request.path && - result.replacements === 1 && + result.replacements === replacements && Number.isSafeInteger(result.bytesWritten) && Number(result.bytesWritten) >= 0 && Number(result.bytesWritten) <= BRIDGE_WORKSPACE_WRITE_MAX_BYTES ); } + if (request.operation === 'preview_edit') { + const replacements = request.edits?.length ?? 1; + const content = typeof result.content === 'string' ? result.content : null; + return ( + hasOnlyKeys(result, WORKSPACE_PREVIEW_EDIT_RESULT_KEYS) && + result.path === request.path && + content !== null && + Buffer.from(content).toString('utf8') === content && + typeof result.hasUtf8Bom === 'boolean' && + typeof result.baseSha256 === 'string' && + /^[a-f0-9]{64}$/.test(result.baseSha256) && + result.replacements === replacements && + Number.isSafeInteger(result.bytesWritten) && + Number(result.bytesWritten) === + new TextEncoder().encode(content).byteLength + + (result.hasUtf8Bom ? 3 : 0) && + Number(result.bytesWritten) <= BRIDGE_WORKSPACE_WRITE_MAX_BYTES + ); + } + if (request.operation === 'execute_command') { const stdout = typeof result.stdout === 'string' ? result.stdout : null; const stderr = typeof result.stderr === 'string' ? result.stderr : null; @@ -761,13 +1005,14 @@ export function isValidBridgeWorkspaceToolCapabilities( capabilities.protocolVersion !== BRIDGE_PROTOCOL_VERSION || !Array.isArray(capabilities.operations) || capabilities.operations.length < 1 || - capabilities.operations.length > 6 || + capabilities.operations.length > 7 || !capabilities.operations.every( (operation) => operation === 'read_file' || operation === 'search_text' || operation === 'list_files' || operation === 'write_file' || + operation === 'preview_edit' || operation === 'edit_file' || operation === 'execute_command', ) || @@ -779,6 +1024,57 @@ export function isValidBridgeWorkspaceToolCapabilities( return false; } + if ( + capabilities.writeFileModes !== undefined && + (!Array.isArray(capabilities.writeFileModes) || + capabilities.writeFileModes.length < 1 || + capabilities.writeFileModes.length > 2 || + !capabilities.operations.includes('write_file') || + !capabilities.writeFileModes.every( + (mode) => mode === 'replace' || mode === 'create', + ) || + new Set(capabilities.writeFileModes).size !== + capabilities.writeFileModes.length) + ) { + return false; + } + + if ( + capabilities.editFileModes !== undefined && + (!Array.isArray(capabilities.editFileModes) || + capabilities.editFileModes.length < 1 || + capabilities.editFileModes.length > 2 || + (!capabilities.operations.includes('edit_file') && + !capabilities.operations.includes('preview_edit')) || + !capabilities.editFileModes.every( + (mode) => mode === 'single' || mode === 'batch', + ) || + new Set(capabilities.editFileModes).size !== + capabilities.editFileModes.length) + ) { + return false; + } + + if ( + capabilities.editFileFeatures !== undefined && + (!Array.isArray(capabilities.editFileFeatures) || + capabilities.editFileFeatures.length !== 1 || + !capabilities.operations.includes('edit_file') || + capabilities.editFileFeatures[0] !== 'expected_base_sha256') + ) { + return false; + } + + if ( + capabilities.listFileFeatures !== undefined && + (!Array.isArray(capabilities.listFileFeatures) || + capabilities.listFileFeatures.length !== 1 || + !capabilities.operations.includes('list_files') || + capabilities.listFileFeatures[0] !== 'after_path') + ) { + return false; + } + const workspaceIds = new Set(); return capabilities.workspaces.every((workspace) => { if (typeof workspace !== 'object' || workspace === null) return false; diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 0b6f5c13..657ba726 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -19,6 +19,7 @@ import type { BridgeWorkerCapabilities, BridgeWorkerCredentialResponse, BridgeWorkerRegistrationResponse, + BridgeWorkspaceToolOperation, } from './protocol.js'; import type { RuntimeLease, RuntimeSupervisor } from './runtime.js'; import type { WorkspaceToolExecutor } from './workspace.js'; @@ -137,6 +138,24 @@ function workspaceCapabilitiesMatch( advertised.operations.every( (operation, index) => operation === executor.operations[index], ) && + advertised.writeFileModes?.length === executor.writeFileModes?.length && + (advertised.writeFileModes?.every( + (mode, index) => mode === executor.writeFileModes?.[index], + ) ?? executor.writeFileModes == null) && + advertised.editFileModes?.length === executor.editFileModes?.length && + (advertised.editFileModes?.every( + (mode, index) => mode === executor.editFileModes?.[index], + ) ?? executor.editFileModes == null) && + advertised.editFileFeatures?.length === + executor.editFileFeatures?.length && + (advertised.editFileFeatures?.every( + (feature, index) => feature === executor.editFileFeatures?.[index], + ) ?? executor.editFileFeatures == null) && + advertised.listFileFeatures?.length === + executor.listFileFeatures?.length && + (advertised.listFileFeatures?.every( + (feature, index) => feature === executor.listFileFeatures?.[index], + ) ?? executor.listFileFeatures == null) && advertised.workspaces.length === executor.workspaces.length && advertised.workspaces.every( (workspace, index) => @@ -191,10 +210,17 @@ function registrationCompatibleCapabilities( const { workspaceTools: _workspaceTools, ...compatible } = capabilities; return compatible; } + const { + writeFileModes: _writeFileModes, + editFileModes: _editFileModes, + editFileFeatures: _editFileFeatures, + listFileFeatures: _listFileFeatures, + ...compatibleWorkspaceTools + } = workspaceTools; return { ...capabilities, workspaceTools: { - ...workspaceTools, + ...compatibleWorkspaceTools, operations, workspaces, }, @@ -208,10 +234,50 @@ function supportedWorkspaceCapabilities( const desired = capabilities.workspaceTools; const supported = registration.supportedWorkspaceToolOperations; if (desired == null || !Array.isArray(supported)) return undefined; - const operations = desired.operations.filter((operation) => + let operations = desired.operations.filter((operation) => supported.includes(operation), ); if (operations.length === 0) return undefined; + let writeFileModes: typeof desired.writeFileModes; + if (operations.includes('write_file')) { + const desiredModes = desired.writeFileModes ?? ['replace']; + const serverModes = registration.supportedWorkspaceWriteFileModes ?? [ + 'replace', + ]; + const commonModes = desiredModes.filter((mode) => + serverModes.includes(mode), + ); + if (commonModes.length === 0) { + operations = operations.filter((operation) => operation !== 'write_file'); + } else if (registration.supportedWorkspaceWriteFileModes != null) { + writeFileModes = commonModes; + } + } + const editOperations = new Set([ + 'preview_edit', + 'edit_file', + ]); + let editFileModes: typeof desired.editFileModes; + if (operations.some((operation) => editOperations.has(operation))) { + const desiredModes = desired.editFileModes ?? ['single']; + const serverModes = registration.supportedWorkspaceEditFileModes ?? [ + 'single', + ]; + const commonModes = desiredModes.filter((mode) => + serverModes.includes(mode), + ); + if (commonModes.length === 0) { + operations = operations.filter( + (operation) => !editOperations.has(operation), + ); + } else if (registration.supportedWorkspaceEditFileModes != null) { + editFileModes = commonModes; + } + } + if (operations.length === 0) return undefined; + const supportsEditRequests = operations.some((operation) => + editOperations.has(operation), + ); const workspaces = desired.workspaces.flatMap((workspace) => { if (workspace.operations == null) return [workspace]; const workspaceOperations = workspace.operations.filter((operation) => @@ -222,12 +288,37 @@ function supportedWorkspaceCapabilities( : [{ ...workspace, operations: workspaceOperations }]; }); if (workspaces.length === 0) return undefined; + const editFileFeatures = desired.editFileFeatures?.filter((feature) => + registration.supportedWorkspaceEditFileFeatures?.includes(feature), + ); + const listFileFeatures = desired.listFileFeatures?.filter((feature) => + registration.supportedWorkspaceListFileFeatures?.includes(feature), + ); + const { + writeFileModes: _writeFileModes, + editFileModes: _editFileModes, + editFileFeatures: _editFileFeatures, + listFileFeatures: _listFileFeatures, + ...compatibleDesired + } = desired; return { ...capabilities, workspaceTools: { - ...desired, + ...compatibleDesired, operations, workspaces, + ...(operations.includes('write_file') && writeFileModes?.length + ? { writeFileModes } + : {}), + ...(supportsEditRequests && editFileModes?.length + ? { editFileModes } + : {}), + ...(operations.includes('edit_file') && editFileFeatures?.length + ? { editFileFeatures } + : {}), + ...(operations.includes('list_files') && listFileFeatures?.length + ? { listFileFeatures } + : {}), }, }; } @@ -875,6 +966,52 @@ export class BridgeWorker { 'Workspace tool operation is not advertised for workspace', ); } + if (workspaceRequest.operation === 'write_file') { + const mode = + workspaceRequest.overwrite === false ? 'create' : 'replace'; + const modes = advertised.writeFileModes; + if ( + (workspaceRequest.overwrite !== undefined && modes == null) || + (modes != null && !modes.includes(mode)) + ) { + throw new BridgeProtocolError( + 'Workspace write mode is not advertised', + ); + } + } + if ( + workspaceRequest.operation === 'preview_edit' || + workspaceRequest.operation === 'edit_file' + ) { + const mode = workspaceRequest.edits === undefined ? 'single' : 'batch'; + const modes = advertised.editFileModes; + if ( + (modes == null && mode !== 'single') || + (modes != null && !modes.includes(mode)) + ) { + throw new BridgeProtocolError( + 'Workspace edit mode is not advertised', + ); + } + if ( + workspaceRequest.operation === 'edit_file' && + workspaceRequest.expectedBaseSha256 !== undefined && + !advertised.editFileFeatures?.includes('expected_base_sha256') + ) { + throw new BridgeProtocolError( + 'Workspace edit feature is not advertised', + ); + } + } + if ( + workspaceRequest.operation === 'list_files' && + workspaceRequest.afterPath !== undefined && + !advertised.listFileFeatures?.includes('after_path') + ) { + throw new BridgeProtocolError( + 'Workspace listing feature is not advertised', + ); + } const isMutation = workspaceRequest.operation === 'write_file' || workspaceRequest.operation === 'edit_file' || @@ -898,6 +1035,17 @@ export class BridgeWorker { workspaceRequest, executionController.signal, ); + if ( + workspaceRequest.operation === 'list_files' && + !advertised.listFileFeatures?.includes('after_path') && + 'nextAfterPath' in payload + ) { + const { + nextAfterPath: _nextAfterPath, + ...compatiblePayload + } = payload; + payload = compatiblePayload; + } workspaceMutationApplied = isMutation; if (isMutation && !isWorkspaceToolResult(workspaceRequest, payload)) { throw new BridgeProtocolError( diff --git a/packages/code/src/workspace-cli.test.ts b/packages/code/src/workspace-cli.test.ts index 169bdf50..c1bc7fce 100644 --- a/packages/code/src/workspace-cli.test.ts +++ b/packages/code/src/workspace-cli.test.ts @@ -151,8 +151,12 @@ test('CLI advertises explicitly enabled writes without exposing the workspace ro 'search_text', 'list_files', 'write_file', + 'preview_edit', 'edit_file', ], + supportedWorkspaceWriteFileModes: ['replace', 'create'], + supportedWorkspaceEditFileModes: ['single', 'batch'], + supportedWorkspaceEditFileFeatures: ['expected_base_sha256'], }), ); return; @@ -215,8 +219,12 @@ test('CLI advertises explicitly enabled writes without exposing the workspace ro 'search_text', 'list_files', 'write_file', + 'preview_edit', 'edit_file', ], + writeFileModes: ['replace', 'create'], + editFileModes: ['single', 'batch'], + editFileFeatures: ['expected_base_sha256'], workspaces: [ { id: 'root-workspace', @@ -226,6 +234,7 @@ test('CLI advertises explicitly enabled writes without exposing the workspace ro 'search_text', 'list_files', 'write_file', + 'preview_edit', 'edit_file', ], }, diff --git a/packages/code/src/workspace-worker.test.ts b/packages/code/src/workspace-worker.test.ts index c98182b0..43c820fe 100644 --- a/packages/code/src/workspace-worker.test.ts +++ b/packages/code/src/workspace-worker.test.ts @@ -15,9 +15,13 @@ const listWorkspaceCapabilities = { 'list_files' as const, ], workspaces: [{ id: 'primary' }], + listFileFeatures: ['after_path' as const], }; -function registrationResponse(supportsList: boolean): Response { +function registrationResponse( + supportsList: boolean, + supportsPagination = false, +): Response { return Response.json({ protocolVersion: 1, workerId: 'vm-1', @@ -31,6 +35,9 @@ function registrationResponse(supportsList: boolean): Response { 'search_text', 'list_files', ], + ...(supportsPagination + ? { supportedWorkspaceListFileFeatures: ['after_path'] } + : {}), } : {}), }); @@ -131,6 +138,67 @@ test('worker re-registers list_files after the Code API advertises support', asy ]); }); +test('worker omits pagination fields until Code API negotiates them', async () => { + let settlement: Record | undefined; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + workspaceTools: listWorkspaceCapabilities, + }, + workspaceTools: { + capabilities: listWorkspaceCapabilities, + async execute() { + return { + protocolVersion: 1 as const, + operation: 'list_files' as const, + workspaceId: 'primary', + paths: ['first.txt'], + truncated: true, + nextAfterPath: 'first.txt', + }; + }, + }, + fetchImpl: async (input, init) => { + if (String(input).endsWith('/register')) return registrationResponse(true); + settlement = JSON.parse(String(init?.body)) as Record; + return Response.json({ protocolVersion: 1, accepted: true }); + }, + }); + + await worker.register(); + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-list-legacy-consumer', + workerId: 'vm-1', + incarnationId, + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 5_000).toISOString(), + executionKind: 'workspace_tool', + request: { + protocolVersion: 1, + operation: 'list_files', + workspaceId: 'primary', + maxResults: 1, + }, + }); + + assert.deepEqual(settlement?.result, { + protocolVersion: 1, + operation: 'list_files', + workspaceId: 'primary', + paths: ['first.txt'], + truncated: true, + }); +}); + test('worker omits restricted workspaces that legacy registration would widen', async () => { const registrations: Array<{ operations: string[]; @@ -220,6 +288,7 @@ test('worker promotes only operations understood by an older Code API', async () const registrations: Array<{ operations: string[]; workspaces: Array>; + writeFileModes?: string[]; }> = []; const workspaceCapabilities = { protocolVersion: 1 as const, @@ -230,6 +299,7 @@ test('worker promotes only operations understood by an older Code API', async () 'write_file' as const, 'edit_file' as const, ], + writeFileModes: ['replace' as const, 'create' as const], workspaces: [ { id: 'primary', @@ -268,6 +338,7 @@ test('worker promotes only operations understood by an older Code API', async () workspaceTools: { operations: string[]; workspaces: Array>; + writeFileModes?: string[]; }; }; }; @@ -312,6 +383,7 @@ test('worker retains per-workspace restrictions during partial mutation promotio const registrations: Array<{ operations: string[]; workspaces: Array>; + writeFileModes?: string[]; }> = []; const workspaceCapabilities = { protocolVersion: 1 as const, @@ -322,6 +394,7 @@ test('worker retains per-workspace restrictions during partial mutation promotio 'write_file' as const, 'edit_file' as const, ], + writeFileModes: ['replace' as const, 'create' as const], workspaces: [ { id: 'readonly', @@ -368,6 +441,7 @@ test('worker retains per-workspace restrictions during partial mutation promotio workspaceTools: { operations: string[]; workspaces: Array>; + writeFileModes?: string[]; }; }; }; @@ -384,6 +458,7 @@ test('worker retains per-workspace restrictions during partial mutation promotio 'list_files', 'write_file', ], + supportedWorkspaceWriteFileModes: ['replace', 'create'], }); }, }); @@ -393,6 +468,7 @@ test('worker retains per-workspace restrictions during partial mutation promotio assert.deepEqual(registrations[1], { protocolVersion: 1, operations: ['read_file', 'search_text', 'list_files', 'write_file'], + writeFileModes: ['replace', 'create'], workspaces: [ { id: 'readonly', @@ -411,6 +487,220 @@ test('worker retains per-workspace restrictions during partial mutation promotio }); }); +test('worker omits write modes not negotiated by an older Code API', async () => { + const registrations: Array> = []; + const workspaceCapabilities = { + protocolVersion: 1 as const, + operations: ['read_file' as const, 'write_file' as const], + writeFileModes: ['replace' as const, 'create' as const], + workspaces: [ + { + id: 'primary', + operations: ['read_file' as const, 'write_file' as const], + }, + ], + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + workspaceTools: workspaceCapabilities, + }, + workspaceTools: { + capabilities: workspaceCapabilities, + async execute() { + throw new Error('not executed'); + }, + }, + workspaceMutationQuarantine: mutationQuarantine(), + fetchImpl: async (_input, init) => { + const body = JSON.parse(String(init?.body)) as { + capabilities: { workspaceTools?: Record }; + }; + registrations.push(body.capabilities.workspaceTools ?? {}); + return Response.json({ + protocolVersion: 1, + workerId: 'vm-1', + incarnationId, + registeredAt: new Date().toISOString(), + leaseTtlMs: 60_000, + supportedWorkspaceToolOperations: ['read_file', 'write_file'], + }); + }, + }); + + await worker.register(); + + assert.deepEqual(registrations, [ + { + protocolVersion: 1, + operations: ['read_file'], + workspaces: [{ id: 'primary' }], + }, + { + protocolVersion: 1, + operations: ['read_file', 'write_file'], + workspaces: [ + { id: 'primary', operations: ['read_file', 'write_file'] }, + ], + }, + ]); +}); + +test('worker advertises only edit modes and features negotiated by Code API', async () => { + const registrations: Array> = []; + const workspaceCapabilities = { + protocolVersion: 1 as const, + operations: ['read_file' as const, 'edit_file' as const], + editFileModes: ['single' as const, 'batch' as const], + editFileFeatures: ['expected_base_sha256' as const], + workspaces: [ + { + id: 'primary', + operations: ['read_file' as const, 'edit_file' as const], + }, + ], + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + workspaceTools: workspaceCapabilities, + }, + workspaceTools: { + capabilities: workspaceCapabilities, + async execute() { + throw new Error('not executed'); + }, + }, + workspaceMutationQuarantine: mutationQuarantine(), + fetchImpl: async (_input, init) => { + const body = JSON.parse(String(init?.body)) as { + capabilities: { workspaceTools?: Record }; + }; + registrations.push(body.capabilities.workspaceTools ?? {}); + return Response.json({ + protocolVersion: 1, + workerId: 'vm-1', + incarnationId, + registeredAt: new Date().toISOString(), + leaseTtlMs: 60_000, + supportedWorkspaceToolOperations: ['read_file', 'edit_file'], + supportedWorkspaceEditFileModes: ['single', 'batch'], + }); + }, + }); + + await worker.register(); + + assert.deepEqual(registrations, [ + { + protocolVersion: 1, + operations: ['read_file'], + workspaces: [{ id: 'primary' }], + }, + { + protocolVersion: 1, + operations: ['read_file', 'edit_file'], + editFileModes: ['single', 'batch'], + workspaces: [ + { id: 'primary', operations: ['read_file', 'edit_file'] }, + ], + }, + ]); +}); + +test('worker drops file operations when no request mode is compatible', async () => { + const registrations: Array> = []; + const workspaceCapabilities = { + protocolVersion: 1 as const, + operations: [ + 'read_file' as const, + 'write_file' as const, + 'preview_edit' as const, + ], + writeFileModes: ['create' as const], + editFileModes: ['batch' as const], + workspaces: [ + { + id: 'primary', + operations: [ + 'read_file' as const, + 'write_file' as const, + 'preview_edit' as const, + ], + }, + ], + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + workspaceTools: workspaceCapabilities, + }, + workspaceTools: { + capabilities: workspaceCapabilities, + async execute() { + throw new Error('not executed'); + }, + }, + workspaceMutationQuarantine: mutationQuarantine(), + fetchImpl: async (_input, init) => { + const body = JSON.parse(String(init?.body)) as { + capabilities: { workspaceTools?: Record }; + }; + registrations.push(body.capabilities.workspaceTools ?? {}); + return Response.json({ + protocolVersion: 1, + workerId: 'vm-1', + incarnationId, + registeredAt: new Date().toISOString(), + leaseTtlMs: 60_000, + supportedWorkspaceToolOperations: [ + 'read_file', + 'write_file', + 'preview_edit', + ], + supportedWorkspaceWriteFileModes: ['replace'], + supportedWorkspaceEditFileModes: ['single'], + }); + }, + }); + + await worker.register(); + + assert.deepEqual(registrations, [ + { + protocolVersion: 1, + operations: ['read_file'], + workspaces: [{ id: 'primary' }], + }, + { + protocolVersion: 1, + operations: ['read_file'], + workspaces: [{ id: 'primary', operations: ['read_file'] }], + }, + ]); +}); + test('worker retains per-workspace restrictions during read-only promotion', async () => { const registrations: Array<{ operations: string[]; @@ -2126,3 +2416,61 @@ test('worker rejects workspace operations outside its advertised capability', as assert.equal(settlement?.status, 'rejected'); assert.match(String(settlement?.error), /operation is not advertised/i); }); + +test('worker rejects legacy replacement writes outside its advertised mode', async () => { + let executions = 0; + let settlement: Record | undefined; + const workspaceCapabilities = { + protocolVersion: 1 as const, + operations: ['write_file' as const], + writeFileModes: ['create' as const], + workspaces: [{ id: 'primary' }], + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + workspaceTools: workspaceCapabilities, + }, + workspaceTools: { + capabilities: workspaceCapabilities, + async execute() { + executions += 1; + throw new Error('must not execute'); + }, + }, + workspaceMutationQuarantine: mutationQuarantine(), + fetchImpl: async (_input, init) => { + settlement = JSON.parse(String(init?.body)) as Record; + return Response.json({ protocolVersion: 1, accepted: true }); + }, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-workspace-replace-mode', + workerId: 'vm-1', + incarnationId, + generation: 4, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 5_000).toISOString(), + executionKind: 'workspace_tool', + request: { + protocolVersion: 1, + operation: 'write_file', + workspaceId: 'primary', + path: 'notes.txt', + content: 'blocked', + }, + }); + + assert.equal(executions, 0); + assert.equal(settlement?.status, 'rejected'); + assert.match(String(settlement?.error), /write mode is not advertised/i); +}); diff --git a/packages/code/src/workspace.test.ts b/packages/code/src/workspace.test.ts index e0cacb48..25894773 100644 --- a/packages/code/src/workspace.test.ts +++ b/packages/code/src/workspace.test.ts @@ -201,6 +201,23 @@ test('lists workspace files deterministically with a hard result bound', async ( workspaceId: 'primary', paths: ['docs/guide.md', 'src/app.ts'], truncated: true, + nextAfterPath: 'src/app.ts', + }, + ); + assert.deepEqual( + await tools.execute({ + protocolVersion: 1, + operation: 'list_files', + workspaceId: 'primary', + maxResults: 2, + afterPath: 'src/app.ts', + }), + { + protocolVersion: 1, + operation: 'list_files', + workspaceId: 'primary', + paths: ['src/worker.ts'], + truncated: false, }, ); assert.deepEqual( @@ -221,6 +238,96 @@ test('lists workspace files deterministically with a hard result bound', async ( ); }); +test('continues a workspace listing beyond the protocol result ceiling', async (t) => { + const root = await mkdtemp(join(tmpdir(), 'librechat-code-workspace-')); + t.after(() => rm(root, { recursive: true, force: true })); + await Promise.all( + Array.from({ length: 501 }, (_, index) => + writeFile(join(root, `file-${String(index).padStart(3, '0')}.txt`), 'x'), + ), + ); + const tools = await LocalWorkspaceTools.create({ + workspaces: [{ id: 'primary', root }], + }); + + const firstPage = await tools.execute({ + protocolVersion: 1, + operation: 'list_files', + workspaceId: 'primary', + maxResults: 500, + }); + assert.equal(firstPage.operation, 'list_files'); + if (firstPage.operation !== 'list_files') assert.fail('expected list result'); + assert.equal(firstPage.paths.length, 500); + assert.equal(firstPage.truncated, true); + assert.equal(firstPage.nextAfterPath, 'file-499.txt'); + + assert.deepEqual( + await tools.execute({ + protocolVersion: 1, + operation: 'list_files', + workspaceId: 'primary', + maxResults: 500, + afterPath: firstPage.nextAfterPath, + }), + { + protocolVersion: 1, + operation: 'list_files', + workspaceId: 'primary', + paths: ['file-500.txt'], + truncated: false, + }, + ); +}); + +test( + 'continues listings across directory and file prefix siblings', + async (t) => { + const root = await mkdtemp(join(tmpdir(), 'librechat-code-workspace-')); + t.after(() => rm(root, { recursive: true, force: true })); + await mkdir(join(root, 'src')); + await writeFile(join(root, 'src', 'app.ts'), 'nested'); + await writeFile(join(root, 'src.ts'), 'sibling'); + const tools = await LocalWorkspaceTools.create({ + workspaces: [{ id: 'primary', root }], + }); + + const firstPage = await tools.execute({ + protocolVersion: 1, + operation: 'list_files', + workspaceId: 'primary', + maxResults: 1, + }); + assert.deepEqual(firstPage, { + protocolVersion: 1, + operation: 'list_files', + workspaceId: 'primary', + paths: ['src/app.ts'], + truncated: true, + nextAfterPath: 'src/app.ts', + }); + if (firstPage.operation !== 'list_files') { + assert.fail('expected list result'); + } + assert.deepEqual( + await tools.execute({ + protocolVersion: 1, + operation: 'list_files', + workspaceId: 'primary', + maxResults: 1, + afterPath: firstPage.nextAfterPath, + }), + { + protocolVersion: 1, + operation: 'list_files', + workspaceId: 'primary', + paths: ['src.ts'], + truncated: false, + }, + ); + }, +); + test('rejects listing through a directory symlink that leaves the workspace', async (t) => { const parent = await mkdtemp( join(tmpdir(), 'librechat-code-workspace-parent-'), @@ -767,6 +874,7 @@ test('advertises workspace IDs and names without exposing host roots', async (t) protocolVersion: 1, operations: ['read_file', 'search_text', 'list_files'], workspaces: [{ id: 'primary', name: 'LibreChat' }], + listFileFeatures: ['after_path'], }); assert.equal(JSON.stringify(tools.capabilities).includes(root), false); }); @@ -805,6 +913,7 @@ test('writable workspaces create, replace, and exactly edit files', async (t) => 'search_text', 'list_files', 'write_file', + 'preview_edit', 'edit_file', ], workspaces: [ @@ -816,10 +925,15 @@ test('writable workspaces create, replace, and exactly edit files', async (t) => 'search_text', 'list_files', 'write_file', + 'preview_edit', 'edit_file', ], }, ], + writeFileModes: ['replace', 'create'], + editFileModes: ['single', 'batch'], + editFileFeatures: ['expected_base_sha256'], + listFileFeatures: ['after_path'], }); await tools.execute({ protocolVersion: 1, @@ -847,6 +961,183 @@ test('writable workspaces create, replace, and exactly edit files', async (t) => assert.equal(await readFile(join(root, 'notes.txt'), 'utf8'), 'hello BYOM'); }); +test('workspace writes can require an atomic create without replacement', async (t) => { + const root = await mkdtemp(join(tmpdir(), 'librechat-code-workspace-')); + t.after(() => rm(root, { recursive: true, force: true })); + await writeFile(join(root, 'existing.txt'), 'preserve me'); + const tools = await LocalWorkspaceTools.create({ + workspaces: [{ id: 'primary', root, writable: true }], + }); + + await assert.rejects( + tools.execute({ + protocolVersion: 1, + operation: 'write_file', + workspaceId: 'primary', + path: 'existing.txt', + content: 'replace me', + overwrite: false, + }), + (error: unknown) => + error instanceof WorkspaceToolError && error.code === 'EDIT_CONFLICT', + ); + assert.equal(await readFile(join(root, 'existing.txt'), 'utf8'), 'preserve me'); + + const created = await tools.execute({ + protocolVersion: 1, + operation: 'write_file', + workspaceId: 'primary', + path: 'created.txt', + content: 'new file', + overwrite: false, + }); + assert.deepEqual(created, { + protocolVersion: 1, + operation: 'write_file', + workspaceId: 'primary', + path: 'created.txt', + created: true, + bytesWritten: 8, + }); + assert.equal(await readFile(join(root, 'created.txt'), 'utf8'), 'new file'); + + const competingWrites = await Promise.allSettled([ + tools.execute({ + protocolVersion: 1, + operation: 'write_file', + workspaceId: 'primary', + path: 'raced.txt', + content: 'first', + overwrite: false, + }), + tools.execute({ + protocolVersion: 1, + operation: 'write_file', + workspaceId: 'primary', + path: 'raced.txt', + content: 'second', + overwrite: false, + }), + ]); + assert.equal( + competingWrites.filter((result) => result.status === 'fulfilled').length, + 1, + ); + const rejected = competingWrites.find( + (result): result is PromiseRejectedResult => result.status === 'rejected', + ); + assert.ok(rejected?.reason instanceof WorkspaceToolError); + assert.equal(rejected.reason.code, 'EDIT_CONFLICT'); + assert.match(await readFile(join(root, 'raced.txt'), 'utf8'), /^(first|second)$/); +}); + +test('workspace batch edits commit all replacements atomically', async (t) => { + const root = await mkdtemp(join(tmpdir(), 'librechat-code-workspace-')); + t.after(() => rm(root, { recursive: true, force: true })); + await writeFile(join(root, 'batch.txt'), 'alpha beta gamma'); + const tools = await LocalWorkspaceTools.create({ + workspaces: [{ id: 'primary', root, writable: true }], + }); + + const result = await tools.execute({ + protocolVersion: 1, + operation: 'edit_file', + workspaceId: 'primary', + path: 'batch.txt', + edits: [ + { oldText: 'alpha', newText: 'one' }, + { oldText: 'gamma', newText: 'three' }, + ], + }); + assert.deepEqual(result, { + protocolVersion: 1, + operation: 'edit_file', + workspaceId: 'primary', + path: 'batch.txt', + replacements: 2, + bytesWritten: 14, + }); + assert.equal(await readFile(join(root, 'batch.txt'), 'utf8'), 'one beta three'); + + await assert.rejects( + tools.execute({ + protocolVersion: 1, + operation: 'edit_file', + workspaceId: 'primary', + path: 'batch.txt', + edits: [ + { oldText: 'one', newText: 'partial' }, + { oldText: 'missing', newText: 'never' }, + ], + }), + (error: unknown) => + error instanceof WorkspaceToolError && error.code === 'EDIT_CONFLICT', + ); + assert.equal(await readFile(join(root, 'batch.txt'), 'utf8'), 'one beta three'); +}); + +test('workspace edit previews are non-mutating and fence the commit revision', async (t) => { + const root = await mkdtemp(join(tmpdir(), 'librechat-code-workspace-')); + t.after(() => rm(root, { recursive: true, force: true })); + await writeFile(join(root, 'preview.txt'), 'prefix SEC suffix'); + const tools = await LocalWorkspaceTools.create({ + workspaces: [{ id: 'primary', root, writable: true }], + }); + + const preview = await tools.execute({ + protocolVersion: 1, + operation: 'preview_edit', + workspaceId: 'primary', + path: 'preview.txt', + oldText: ' suffix', + newText: 'RET suffix', + }); + assert.equal(preview.operation, 'preview_edit'); + assert.equal(preview.content, 'prefix SECRET suffix'); + assert.equal(preview.hasUtf8Bom, false); + assert.match(preview.baseSha256, /^[a-f0-9]{64}$/); + assert.equal(await readFile(join(root, 'preview.txt'), 'utf8'), 'prefix SEC suffix'); + + await writeFile(join(root, 'preview.txt'), 'changed SEC suffix'); + await assert.rejects( + tools.execute({ + protocolVersion: 1, + operation: 'edit_file', + workspaceId: 'primary', + path: 'preview.txt', + oldText: ' suffix', + newText: 'RET suffix', + expectedBaseSha256: preview.baseSha256, + }), + (error: unknown) => + error instanceof WorkspaceToolError && error.code === 'EDIT_CONFLICT', + ); + assert.equal(await readFile(join(root, 'preview.txt'), 'utf8'), 'changed SEC suffix'); +}); + +test('workspace edit previews strip a UTF-8 BOM while retaining its byte count', async (t) => { + const root = await mkdtemp(join(tmpdir(), 'librechat-code-workspace-')); + t.after(() => rm(root, { recursive: true, force: true })); + await writeFile(join(root, 'bom.txt'), Buffer.from('\ufeffbefore', 'utf8')); + const tools = await LocalWorkspaceTools.create({ + workspaces: [{ id: 'primary', root, writable: true }], + }); + + const preview = await tools.execute({ + protocolVersion: 1, + operation: 'preview_edit', + workspaceId: 'primary', + path: 'bom.txt', + oldText: 'before', + newText: 'after', + }); + assert.equal(preview.operation, 'preview_edit'); + assert.equal(preview.content, 'after'); + assert.equal(preview.hasUtf8Bom, true); + assert.equal(preview.bytesWritten, 8); + assert.equal(await readFile(join(root, 'bom.txt'), 'utf8'), '\ufeffbefore'); +}); + test('workspace mutations sync the containing directory after replacement', async (t) => { if (process.platform === 'win32') { t.skip('Directory fsync is unavailable on Windows'); @@ -889,6 +1180,42 @@ test('workspace mutations sync the containing directory after replacement', asyn assert.equal(syncCalls, 5); }); +test('atomic creates remove staging before syncing the directory', async (t) => { + if (process.platform === 'win32') { + t.skip('Directory fsync is unavailable on Windows'); + return; + } + const root = await mkdtemp(join(tmpdir(), 'librechat-code-workspace-')); + t.after(() => rm(root, { recursive: true, force: true })); + const tools = await LocalWorkspaceTools.create({ + workspaces: [{ id: 'primary', root, writable: true }], + }); + const probe = await open(root, 'r'); + const fileHandlePrototype = Object.getPrototypeOf(probe) as { + sync(): Promise; + }; + await probe.close(); + const originalSync = fileHandlePrototype.sync; + let syncCalls = 0; + t.mock.method(fileHandlePrototype, 'sync', async function (this: FileHandle) { + syncCalls += 1; + if (syncCalls === 2) { + assert.deepEqual(await readdir(root), ['created.txt']); + } + await originalSync.call(this); + }); + + await tools.execute({ + protocolVersion: 1, + operation: 'write_file', + workspaceId: 'primary', + path: 'created.txt', + content: 'durable create', + overwrite: false, + }); + assert.equal(syncCalls, 2); +}); + test('workspace mutations report uncertain commit when directory sync fails', async (t) => { if (process.platform === 'win32') { t.skip('Directory fsync is unavailable on Windows'); @@ -1588,9 +1915,16 @@ test('composes sandboxed commands without exposing them on unconfigured workspac 'search_text', 'list_files', 'write_file', + 'preview_edit', 'edit_file', 'execute_command', ]); + assert.deepEqual(tools.capabilities.writeFileModes, ['replace', 'create']); + assert.deepEqual(tools.capabilities.editFileModes, ['single', 'batch']); + assert.deepEqual(tools.capabilities.editFileFeatures, [ + 'expected_base_sha256', + ]); + assert.deepEqual(tools.capabilities.listFileFeatures, ['after_path']); assert.deepEqual( tools.capabilities.workspaces.find(({ id }) => id === 'sandboxed')?.operations, tools.capabilities.operations, diff --git a/packages/code/src/workspace.ts b/packages/code/src/workspace.ts index 3f8ed7b6..727955cd 100644 --- a/packages/code/src/workspace.ts +++ b/packages/code/src/workspace.ts @@ -1,7 +1,7 @@ import { spawn } from 'node:child_process'; -import { randomBytes } from 'node:crypto'; +import { createHash, randomBytes } from 'node:crypto'; import { constants } from 'node:fs'; -import { lstat, open, realpath, rename, stat, unlink } from 'node:fs/promises'; +import { link, lstat, open, realpath, rename, stat, unlink } from 'node:fs/promises'; import { basename, dirname, isAbsolute, relative, resolve, sep } from 'node:path'; import type { FileHandle } from 'node:fs/promises'; @@ -14,6 +14,7 @@ import { BRIDGE_WORKSPACE_LIST_MAX_RESULTS, BRIDGE_WORKSPACE_SEARCH_MAX_RESULTS, BRIDGE_WORKSPACE_SEARCH_TEXT_MAX_LENGTH, + comparePortableRelativePaths, isSafePortableRelativePath, isValidBridgeWorkspaceToolCapabilities, isWorkspaceToolRequest, @@ -27,6 +28,8 @@ import type { WorkspaceReadFileResult, WorkspaceEditFileRequest, WorkspaceEditFileResult, + WorkspacePreviewEditRequest, + WorkspacePreviewEditResult, WorkspaceExecuteCommandRequest, WorkspaceExecuteCommandResult, WorkspaceListFilesRequest, @@ -47,6 +50,8 @@ export type { WorkspaceReadFileResult, WorkspaceEditFileRequest, WorkspaceEditFileResult, + WorkspacePreviewEditRequest, + WorkspacePreviewEditResult, WorkspaceExecuteCommandRequest, WorkspaceExecuteCommandResult, WorkspaceListFilesRequest, @@ -119,7 +124,7 @@ const READ_OPERATIONS = [ 'search_text', 'list_files', ] as const; -const WRITE_OPERATIONS = ['write_file', 'edit_file'] as const; +const WRITE_OPERATIONS = ['write_file', 'preview_edit', 'edit_file'] as const; function isUtf8ScalarString(value: string): boolean { return Buffer.from(value).toString('utf8') === value; @@ -455,6 +460,7 @@ async function atomicWriteConfinedFile( content: Buffer, signal?: AbortSignal, expected?: { dev: bigint | number; ino: bigint | number; content: Buffer }, + allowOverwrite = true, ): Promise<{ created: boolean }> { throwIfAborted(signal); if (content.byteLength > BRIDGE_WORKSPACE_WRITE_MAX_BYTES) { @@ -488,6 +494,12 @@ async function atomicWriteConfinedFile( if (existing?.isSymbolicLink() || (existing != null && !existing.isFile())) { throw new WorkspaceToolError('Invalid workspace path', 'INVALID_PATH'); } + if (!allowOverwrite && existing != null) { + throw new WorkspaceToolError( + 'Workspace file already exists', + 'EDIT_CONFLICT', + ); + } if ( expected != null && (existing == null || @@ -506,6 +518,7 @@ async function atomicWriteConfinedFile( ); const installTarget = resolve(canonicalParent, basename(candidate)); let handle: FileHandle | undefined; + let temporaryNeedsCleanup = true; let staged: { dev: bigint | number; ino: bigint | number; content: Buffer }; try { handle = await open( @@ -588,10 +601,36 @@ async function atomicWriteConfinedFile( staged, signal, ); + temporaryNeedsCleanup = false; return { created: false }; } throwIfAborted(signal); - await rename(temporary, installTarget); + if (allowOverwrite) { + 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 create cleanup could not be confirmed', + 'WRITE_UNAVAILABLE', + true, + ); + } + } await confirmInstalledMutation(root, installTarget, staged); return { created: existing == null }; } catch (error) { @@ -602,7 +641,9 @@ async function atomicWriteConfinedFile( ); } finally { await handle?.close().catch(() => undefined); - await unlink(temporary).catch(() => undefined); + if (temporaryNeedsCleanup) { + await unlink(temporary).catch(() => undefined); + } } } @@ -617,6 +658,8 @@ async function writeWorkspaceFile( request.path, content, signal, + undefined, + request.overwrite !== false, ); return { protocolVersion: BRIDGE_PROTOCOL_VERSION, @@ -660,34 +703,17 @@ async function editWorkspaceFile( ); } const original = await readBoundedEditFile(opened); - const hasBom = - original[0] === 0xef && original[1] === 0xbb && original[2] === 0xbf; - const body = hasBom ? original.subarray(3) : original; - const text = body.toString('utf8'); - if (!Buffer.from(text, 'utf8').equals(body)) { - throw new WorkspaceToolError( - 'Workspace file is not UTF-8 text', - 'INVALID_REQUEST', - ); - } - const first = text.indexOf(request.oldText); if ( - first < 0 || - text.indexOf(request.oldText, first + 1) >= 0 + request.expectedBaseSha256 !== undefined && + createHash('sha256').update(original).digest('hex') !== + request.expectedBaseSha256 ) { throw new WorkspaceToolError( - 'Workspace edit must match exactly once', + 'Workspace file changed after edit preview', 'EDIT_CONFLICT', ); } - const updatedText = - text.slice(0, first) + - request.newText + - text.slice(first + request.oldText.length); - const updatedBody = Buffer.from(updatedText, 'utf8'); - const updated = hasBom - ? Buffer.concat([Buffer.from([0xef, 0xbb, 0xbf]), updatedBody]) - : updatedBody; + const { updated, replacements } = applyWorkspaceEdits(original, request); await atomicWriteConfinedFile( root, request.path, @@ -704,7 +730,7 @@ async function editWorkspaceFile( operation: 'edit_file', workspaceId: request.workspaceId, path: request.path, - replacements: 1, + replacements, bytesWritten: updated.byteLength, }; } catch (error) { @@ -715,6 +741,84 @@ async function editWorkspaceFile( } } +function applyWorkspaceEdits( + original: Buffer, + request: WorkspaceEditFileRequest | WorkspacePreviewEditRequest, +): { updated: Buffer; replacements: number } { + const hasBom = + original[0] === 0xef && original[1] === 0xbb && original[2] === 0xbf; + const body = hasBom ? original.subarray(3) : original; + const text = body.toString('utf8'); + if (!Buffer.from(text, 'utf8').equals(body)) { + throw new WorkspaceToolError( + 'Workspace file is not UTF-8 text', + 'INVALID_REQUEST', + ); + } + const edits = request.edits ?? [ + { oldText: request.oldText ?? '', newText: request.newText ?? '' }, + ]; + let updatedText = text; + for (const edit of edits) { + const first = updatedText.indexOf(edit.oldText); + if (first < 0 || updatedText.indexOf(edit.oldText, first + 1) >= 0) { + throw new WorkspaceToolError( + 'Workspace edit must match exactly once', + 'EDIT_CONFLICT', + ); + } + updatedText = + updatedText.slice(0, first) + + edit.newText + + updatedText.slice(first + edit.oldText.length); + } + const updatedBody = Buffer.from(updatedText, 'utf8'); + const updated = hasBom + ? Buffer.concat([Buffer.from([0xef, 0xbb, 0xbf]), updatedBody]) + : updatedBody; + if (updated.byteLength > BRIDGE_WORKSPACE_WRITE_MAX_BYTES) { + throw new WorkspaceToolError( + 'Workspace file exceeds write limit', + 'WRITE_LIMIT_EXCEEDED', + ); + } + return { updated, replacements: edits.length }; +} + +async function previewWorkspaceEdit( + root: string, + request: WorkspacePreviewEditRequest, + signal?: AbortSignal, +): Promise { + const original = await readConfinedFileBuffer(root, request.path); + if (signal?.aborted) { + throw new WorkspaceToolError( + 'Workspace tool execution aborted', + 'EXECUTION_ABORTED', + ); + } + const { updated, replacements } = applyWorkspaceEdits(original, request); + if (signal?.aborted) { + throw new WorkspaceToolError( + 'Workspace tool execution aborted', + 'EXECUTION_ABORTED', + ); + } + const hasUtf8Bom = + updated[0] === 0xef && updated[1] === 0xbb && updated[2] === 0xbf; + return { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operation: 'preview_edit', + workspaceId: request.workspaceId, + path: request.path, + content: decodeWorkspaceText(updated), + hasUtf8Bom, + baseSha256: createHash('sha256').update(original).digest('hex'), + replacements, + bytesWritten: updated.byteLength, + }; +} + interface SearchCandidates { paths: string[]; truncated: boolean; @@ -1049,6 +1153,7 @@ async function listWorkspaceFiles( .filter((segment) => segment.length > 0 && segment !== '.') .join('/'); const requestedResultPath = normalizedRequestedResultPath || undefined; + const afterPath = request.afterPath; const candidates: Array<{ filesystemPath: string; resultPath: string }> = []; let truncated = false; @@ -1127,6 +1232,12 @@ async function listWorkspaceFiles( ? `${requestedResultPath}${normalizedPath.slice(portableCanonicalListPath.length)}` : normalizedPath; if (!isSafePortableRelativePath(resultPath)) return; + if ( + afterPath !== undefined && + comparePortableRelativePaths(resultPath, afterPath) <= 0 + ) { + return; + } candidates.push({ filesystemPath: normalizedPath, resultPath }); }; @@ -1234,6 +1345,9 @@ async function listWorkspaceFiles( workspaceId: request.workspaceId, paths, truncated, + ...(truncated && paths.length > 0 + ? { nextAfterPath: paths[paths.length - 1] } + : {}), }; } @@ -1298,11 +1412,19 @@ export class LocalWorkspaceTools implements WorkspaceToolExecutor { private readonly roots: ReadonlyMap, operations: BridgeWorkspaceToolCapabilities['operations'], workspaces: BridgeWorkspaceDescriptor[], + writeFileModes?: BridgeWorkspaceToolCapabilities['writeFileModes'], + editFileModes?: BridgeWorkspaceToolCapabilities['editFileModes'], + editFileFeatures?: BridgeWorkspaceToolCapabilities['editFileFeatures'], + listFileFeatures?: BridgeWorkspaceToolCapabilities['listFileFeatures'], ) { this.capabilities = { protocolVersion: BRIDGE_PROTOCOL_VERSION, operations, workspaces, + ...(writeFileModes != null ? { writeFileModes } : {}), + ...(editFileModes != null ? { editFileModes } : {}), + ...(editFileFeatures != null ? { editFileFeatures } : {}), + ...(listFileFeatures != null ? { listFileFeatures } : {}), }; } @@ -1335,6 +1457,12 @@ export class LocalWorkspaceTools implements WorkspaceToolExecutor { protocolVersion: BRIDGE_PROTOCOL_VERSION, operations, workspaces, + ...(anyWritable ? { writeFileModes: ['replace', 'create'] } : {}), + ...(anyWritable ? { editFileModes: ['single', 'batch'] } : {}), + ...(anyWritable + ? { editFileFeatures: ['expected_base_sha256'] } + : {}), + listFileFeatures: ['after_path'], }; if (!isValidBridgeWorkspaceToolCapabilities(capabilities)) { throw new WorkspaceToolError( @@ -1358,7 +1486,15 @@ export class LocalWorkspaceTools implements WorkspaceToolExecutor { writable: workspace.writable === true, }); } - return new LocalWorkspaceTools(roots, operations, workspaces); + return new LocalWorkspaceTools( + roots, + operations, + workspaces, + capabilities.writeFileModes, + capabilities.editFileModes, + capabilities.editFileFeatures, + capabilities.listFileFeatures, + ); } async execute( @@ -1383,7 +1519,11 @@ export class LocalWorkspaceTools implements WorkspaceToolExecutor { } const { root } = workspace; - if (request.operation === 'write_file' || request.operation === 'edit_file') { + if ( + request.operation === 'write_file' || + request.operation === 'preview_edit' || + request.operation === 'edit_file' + ) { if (!workspace.writable) { throw new WorkspaceToolError( 'Workspace mutations are disabled by the worker', @@ -1396,8 +1536,11 @@ export class LocalWorkspaceTools implements WorkspaceToolExecutor { 'EXECUTION_ABORTED', ); } - return request.operation === 'write_file' - ? writeWorkspaceFile(root, request, signal) + if (request.operation === 'write_file') { + return writeWorkspaceFile(root, request, signal); + } + return request.operation === 'preview_edit' + ? previewWorkspaceEdit(root, request, signal) : editWorkspaceFile(root, request, signal); } @@ -1482,6 +1625,18 @@ export class SandboxWorkspaceTools implements WorkspaceToolExecutor { this.capabilities = { protocolVersion: BRIDGE_PROTOCOL_VERSION, operations: [...new Set([...base.operations, 'execute_command' as const])], + ...(base.writeFileModes != null + ? { writeFileModes: base.writeFileModes } + : {}), + ...(base.editFileModes != null + ? { editFileModes: base.editFileModes } + : {}), + ...(base.editFileFeatures != null + ? { editFileFeatures: base.editFileFeatures } + : {}), + ...(base.listFileFeatures != null + ? { listFileFeatures: base.listFileFeatures } + : {}), workspaces: base.workspaces.map((workspace) => ({ ...workspace, operations: [ diff --git a/service/src/bridge/router.test.ts b/service/src/bridge/router.test.ts index 2c0d8998..11b0eb31 100644 --- a/service/src/bridge/router.test.ts +++ b/service/src/bridge/router.test.ts @@ -284,9 +284,14 @@ describe('paired bridge HTTP API', () => { 'search_text', 'list_files', 'write_file', + 'preview_edit', 'edit_file', 'execute_command', ], + supportedWorkspaceWriteFileModes: ['replace', 'create'], + supportedWorkspaceEditFileModes: ['single', 'batch'], + supportedWorkspaceEditFileFeatures: ['expected_base_sha256'], + supportedWorkspaceListFileFeatures: ['after_path'], }); const crossDeploymentRevoke = await fetch( diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index b764bb39..7fa87b93 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -384,9 +384,14 @@ router.post( 'search_text', 'list_files', 'write_file', + 'preview_edit', 'edit_file', 'execute_command', ], + supportedWorkspaceWriteFileModes: ['replace', 'create'], + supportedWorkspaceEditFileModes: ['single', 'batch'], + supportedWorkspaceEditFileFeatures: ['expected_base_sha256'], + supportedWorkspaceListFileFeatures: ['after_path'], }); } catch (error) { if (error instanceof BridgeStoreError) { diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index b0a5c2fa..03694b1b 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -76,13 +76,41 @@ function supportsWorkspaceTool( const workspace = capabilities?.workspaces.find( (candidate) => candidate.id === request.workspaceId, ); - return ( + const supportsOperation = capabilities != null && capabilities.operations.includes(request.operation) && workspace != null && (workspace.operations == null || - workspace.operations.includes(request.operation)) - ); + workspace.operations.includes(request.operation)); + if (!supportsOperation) { + return supportsOperation; + } + if (request.operation === 'list_files' && request.afterPath !== undefined) { + return capabilities?.listFileFeatures?.includes('after_path') === true; + } + if (request.operation === 'write_file') { + const mode = request.overwrite === false ? 'create' : 'replace'; + const modes = capabilities?.writeFileModes; + return request.overwrite === undefined && modes == null + ? true + : modes?.includes(mode) === true; + } + if ( + request.operation === 'preview_edit' || + request.operation === 'edit_file' + ) { + const mode = request.edits === undefined ? 'single' : 'batch'; + const modes = capabilities?.editFileModes; + const supportsMode = modes == null ? mode === 'single' : modes.includes(mode); + if (request.operation === 'preview_edit') return supportsMode; + return ( + supportsMode && + (request.expectedBaseSha256 === undefined || + capabilities?.editFileFeatures?.includes('expected_base_sha256') === + true) + ); + } + return true; } function workerKey(workerId: string): string { @@ -485,22 +513,28 @@ export class RedisBridgeStore { 'Invalid workspace tool request', ); } - const settlement = (await this.dispatch({ + return (await this.dispatch({ ...args, body: {} as t.PayloadBody, headers: {}, workspaceRequest: args.request, + finalize: async (settlement, registration) => { + if ( + settlement.status === 'fulfilled' && + !isWorkspaceToolResult( + args.request, + settlement.result, + registration.capabilities.workspaceTools, + ) + ) { + throw new BridgeStoreError( + 'RESULT_INVALID', + 'Bridge worker returned an invalid workspace tool result', + ); + } + return settlement; + }, })) as unknown as CodeBridgeWorkspaceSettlement; - if ( - settlement.status === 'fulfilled' && - !isWorkspaceToolResult(args.request, settlement.result) - ) { - throw new BridgeStoreError( - 'RESULT_INVALID', - 'Bridge worker returned an invalid workspace tool result', - ); - } - return settlement; } async dispatch(args: { @@ -515,6 +549,7 @@ export class RedisBridgeStore { signal: AbortSignal; finalize?: ( settlement: CodeBridgeSettlement, + registration: RegisteredBridgeWorker, ) => Promise; }): Promise { this.assertDispatchActive(args.signal, args.deadlineAtMs); @@ -694,7 +729,7 @@ export class RedisBridgeStore { const result = args.finalize == null ? settlement - : await args.finalize(settlement); + : await args.finalize(settlement, registration); await this.commitPendingWorkspace( assignment, settlement, diff --git a/service/src/bridge/workspace-store.test.ts b/service/src/bridge/workspace-store.test.ts index 44ab66fe..a23274d8 100644 --- a/service/src/bridge/workspace-store.test.ts +++ b/service/src/bridge/workspace-store.test.ts @@ -44,7 +44,11 @@ test('dispatches a workspace tool only to a worker advertising its workspace and signal: new AbortController().signal, }); - const assignment = await store.lease('workspace-worker', incarnationId, 1_000); + const assignment = await store.lease( + 'workspace-worker', + incarnationId, + 1_000, + ); expect(assignment).toMatchObject({ executionKind: 'workspace_tool', request, @@ -106,6 +110,142 @@ test('rejects a workspace tool that the selected worker did not advertise', asyn expect(await redis.keys('codeapi:bridge:v1:assignment:*')).toHaveLength(0); }); +test('rejects listing continuation without the negotiated feature', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'workspace-worker', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + workspaceTools: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operations: ['list_files'], + workspaces: [{ id: 'primary' }], + }, + }, + }); + + await expect( + store.dispatchWorkspaceTool({ + workerId: 'workspace-worker', + request: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operation: 'list_files', + workspaceId: 'primary', + maxResults: 10, + afterPath: 'src/app.ts', + }, + deadlineAtMs: Date.now() + 1_000, + signal: new AbortController().signal, + }), + ).rejects.toMatchObject({ code: 'WORKER_MISMATCH' }); + expect(await redis.keys('codeapi:bridge:v1:assignment:*')).toHaveLength(0); +}); + +test('rejects pagination fields from a worker without the negotiated feature', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'workspace-worker', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + workspaceTools: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operations: ['list_files'], + workspaces: [{ id: 'primary' }], + }, + }, + }); + const completion = store.dispatchWorkspaceTool({ + workerId: 'workspace-worker', + request: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operation: 'list_files', + workspaceId: 'primary', + maxResults: 1, + }, + deadlineAtMs: Date.now() + 5_000, + signal: new AbortController().signal, + }); + + const assignment = await store.lease('workspace-worker', incarnationId, 1_000); + await store.settle('workspace-worker', assignment?.assignmentId ?? '', { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: assignment?.generation ?? 0, + leaseToken: assignment?.leaseToken ?? '', + incarnationId, + status: 'fulfilled', + result: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operation: 'list_files', + workspaceId: 'primary', + paths: ['first.txt'], + truncated: true, + nextAfterPath: 'first.txt', + }, + }); + + await expect(completion).rejects.toMatchObject({ code: 'RESULT_INVALID' }); +}); + +test('accepts a legacy truncated listing without a pagination cursor', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'workspace-worker', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + workspaceTools: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operations: ['list_files'], + workspaces: [{ id: 'primary' }], + }, + }, + }); + const completion = store.dispatchWorkspaceTool({ + workerId: 'workspace-worker', + request: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operation: 'list_files', + workspaceId: 'primary', + maxResults: 2, + }, + deadlineAtMs: Date.now() + 5_000, + signal: new AbortController().signal, + }); + + const assignment = await store.lease( + 'workspace-worker', + incarnationId, + 1_000, + ); + await store.settle('workspace-worker', assignment?.assignmentId ?? '', { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: assignment?.generation ?? 0, + leaseToken: assignment?.leaseToken ?? '', + incarnationId, + status: 'fulfilled', + result: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operation: 'list_files', + workspaceId: 'primary', + paths: ['src//z.ts', 'src/a.ts'], + truncated: true, + }, + }); + + await expect(completion).resolves.toMatchObject({ + status: 'fulfilled', + result: { paths: ['src//z.ts', 'src/a.ts'], truncated: true }, + }); +}); + test('rejects an operation omitted from the selected workspace capability', async () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, @@ -143,6 +283,181 @@ test('rejects an operation omitted from the selected workspace capability', asyn expect(await redis.keys('codeapi:bridge:v1:assignment:*')).toHaveLength(0); }); +test('rejects create-only writes from workers without the negotiated mode', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'workspace-worker', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + workspaceTools: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operations: ['write_file'], + workspaces: [{ id: 'primary' }], + }, + }, + }); + + await expect( + store.dispatchWorkspaceTool({ + workerId: 'workspace-worker', + request: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operation: 'write_file', + workspaceId: 'primary', + path: 'notes.txt', + content: 'create me', + overwrite: false, + }, + deadlineAtMs: Date.now() + 1_000, + signal: new AbortController().signal, + }), + ).rejects.toMatchObject({ code: 'WORKER_MISMATCH' }); + expect(await redis.keys('codeapi:bridge:v1:assignment:*')).toHaveLength(0); +}); + +test('rejects legacy replacement writes from create-only workers', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'workspace-worker', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + workspaceTools: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operations: ['write_file'], + writeFileModes: ['create'], + workspaces: [{ id: 'primary' }], + }, + }, + }); + + await expect( + store.dispatchWorkspaceTool({ + workerId: 'workspace-worker', + request: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operation: 'write_file', + workspaceId: 'primary', + path: 'notes.txt', + content: 'replace me', + }, + deadlineAtMs: Date.now() + 1_000, + signal: new AbortController().signal, + }), + ).rejects.toMatchObject({ code: 'WORKER_MISMATCH' }); + expect(await redis.keys('codeapi:bridge:v1:assignment:*')).toHaveLength(0); +}); + +test('rejects batch edits from workers without the negotiated mode', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'workspace-worker', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + workspaceTools: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operations: ['edit_file'], + workspaces: [{ id: 'primary' }], + }, + }, + }); + + await expect( + store.dispatchWorkspaceTool({ + workerId: 'workspace-worker', + request: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operation: 'edit_file', + workspaceId: 'primary', + path: 'notes.txt', + edits: [{ oldText: 'before', newText: 'after' }], + }, + deadlineAtMs: Date.now() + 1_000, + signal: new AbortController().signal, + }), + ).rejects.toMatchObject({ code: 'WORKER_MISMATCH' }); + expect(await redis.keys('codeapi:bridge:v1:assignment:*')).toHaveLength(0); +}); + +test('rejects fenced edits from workers without the negotiated feature', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'workspace-worker', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + workspaceTools: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operations: ['edit_file'], + editFileModes: ['single'], + workspaces: [{ id: 'primary' }], + }, + }, + }); + + await expect( + store.dispatchWorkspaceTool({ + workerId: 'workspace-worker', + request: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operation: 'edit_file', + workspaceId: 'primary', + path: 'notes.txt', + oldText: 'before', + newText: 'after', + expectedBaseSha256: 'a'.repeat(64), + }, + deadlineAtMs: Date.now() + 1_000, + signal: new AbortController().signal, + }), + ).rejects.toMatchObject({ code: 'WORKER_MISMATCH' }); + expect(await redis.keys('codeapi:bridge:v1:assignment:*')).toHaveLength(0); +}); + +test('rejects batch previews from workers without the negotiated mode', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'workspace-worker', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + workspaceTools: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operations: ['preview_edit'], + workspaces: [{ id: 'primary' }], + }, + }, + }); + + await expect( + store.dispatchWorkspaceTool({ + workerId: 'workspace-worker', + request: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operation: 'preview_edit', + workspaceId: 'primary', + path: 'notes.txt', + edits: [{ oldText: 'before', newText: 'after' }], + }, + deadlineAtMs: Date.now() + 1_000, + signal: new AbortController().signal, + }), + ).rejects.toMatchObject({ code: 'WORKER_MISMATCH' }); + expect(await redis.keys('codeapi:bridge:v1:assignment:*')).toHaveLength(0); +}); + test('rejects a fulfilled workspace settlement that violates the result contract', async () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION,