From b7fd4ba2a2a1e87551f12f58d08916eb412420f0 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Fri, 4 Sep 2026 09:58:46 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=91=93=20feat:=20Read=20Attached=20Worksp?= =?UTF-8?q?ace=20Files=20(#15524)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(code): read attached workspace files * style(api): sort workspace handler import * fix(code): abort attached workspace reads * test(code): type workspace fetch doubles * style(code): sort workspace bridge imports * fix(code): Preserve workspace read boundaries * style(code): Flatten workspace line validation * fix(code): Bound attached workspace responses * fix(code): Require canonical workspace paths * fix(code): fence workspace bridge operations --- .../services/Endpoints/agents/skillDeps.js | 2 + .../Endpoints/agents/skillDeps.spec.js | 8 + api/server/services/Files/Code/process.js | 48 +++ .../services/Files/Code/process.spec.js | 54 +++ packages/api/src/agents/handlers.spec.ts | 306 +++++++++++++++ packages/api/src/agents/handlers.ts | 181 ++++++++- packages/api/src/agents/initialize.ts | 5 + packages/api/src/agents/skills.ts | 4 + packages/api/src/agents/tools.spec.ts | 47 +++ packages/api/src/agents/tools.ts | 59 ++- packages/api/src/code/index.ts | 1 + packages/api/src/code/workspace.spec.ts | 358 ++++++++++++++++++ packages/api/src/code/workspace.ts | 301 +++++++++++++++ 13 files changed, 1366 insertions(+), 8 deletions(-) create mode 100644 packages/api/src/code/workspace.spec.ts create mode 100644 packages/api/src/code/workspace.ts diff --git a/api/server/services/Endpoints/agents/skillDeps.js b/api/server/services/Endpoints/agents/skillDeps.js index b58983da9a..7b7b6aa219 100644 --- a/api/server/services/Endpoints/agents/skillDeps.js +++ b/api/server/services/Endpoints/agents/skillDeps.js @@ -4,6 +4,7 @@ const { batchUploadCodeEnvFiles } = require('~/server/services/Files/Code/crud') const { getSessionInfo, checkIfActive, + readWorkspaceFile, readSandboxFile, readSandboxImage, writeSandboxFile, @@ -380,6 +381,7 @@ const skillToolDeps = { updateSkillFileCodeEnvIds: deploymentSkillMethods.updateSkillFileCodeEnvIds, getSkillFileByPath: deploymentSkillMethods.getSkillFileByPath, updateSkillFileContent: deploymentSkillMethods.updateSkillFileContent, + readWorkspaceFile, /** * `read_file` falls back to a sandbox `cat` for `/mnt/data/...` paths * and for `{firstSegment}/...` paths whose first segment isn't a known diff --git a/api/server/services/Endpoints/agents/skillDeps.spec.js b/api/server/services/Endpoints/agents/skillDeps.spec.js index 3782a4664a..084e505803 100644 --- a/api/server/services/Endpoints/agents/skillDeps.spec.js +++ b/api/server/services/Endpoints/agents/skillDeps.spec.js @@ -5,6 +5,7 @@ const mockGetFileStrategy = jest.fn(); const mockGetStorageMetadata = jest.fn(); const mockResolveRequestTenantId = jest.fn(); const mockCreateDeploymentSkillMethods = jest.fn((methods) => methods); +const mockReadWorkspaceFile = jest.fn(); jest.mock('~/server/services/Files/strategies', () => ({ getStrategyFunctions: (...args) => mockGetStrategyFunctions(...args), @@ -18,6 +19,7 @@ jest.mock('~/server/services/Files/Code/process', () => ({ getSessionInfo: jest.fn(), checkIfActive: jest.fn(), readSandboxFile: jest.fn(), + readWorkspaceFile: (...args) => mockReadWorkspaceFile(...args), writeSandboxFile: jest.fn(), })); @@ -79,6 +81,12 @@ describe('skillDeps saveSkillFileContent', () => { mockDb.getSkillFileByPath.mockResolvedValue(null); }); + it('exposes the stable attached-workspace reader to agent handlers', () => { + expect(getSkillToolDeps().readWorkspaceFile).toBeDefined(); + getSkillToolDeps().readWorkspaceFile({ file_path: 'src/app.ts' }); + expect(mockReadWorkspaceFile).toHaveBeenCalledWith({ file_path: 'src/app.ts' }); + }); + it('cleans up the uploaded object when metadata upsert returns no row', async () => { mockDb.upsertSkillFile.mockResolvedValue(null); diff --git a/api/server/services/Files/Code/process.js b/api/server/services/Files/Code/process.js index 1b90948e88..6862840b5a 100644 --- a/api/server/services/Files/Code/process.js +++ b/api/server/services/Files/Code/process.js @@ -28,6 +28,7 @@ const { getCodeExecutionBaseUrl, buildCodeEnvDownloadQuery, codeExecutionHeaders, + executeWorkspaceTool, claimCodeDestination, createCodeDestinationSet, CODE_OUTPUT_PREFLIGHT_MAX_BYTES, @@ -1663,6 +1664,52 @@ async function readSandboxFile({ return { content: String(result.stdout) }; } +/** + * Reads a bounded range from the workspace directory registered by an attached worker. + * The authenticated worker route is derived from the selected environment and + * the host path remains private to the worker. + * + * @param {Object} params + * @param {string} params.file_path + * @param {string} params.workspace_id + * @param {number} params.start_line + * @param {number} params.max_lines + * @param {string} params.codeApiBaseUrl + * @param {'default' | 'stateful'} params.executionProfile + * @param {string} [params.bridgeWorkerId] + * @param {ServerRequest} [params.req] + * @param {AbortSignal} [params.signal] + */ +async function readWorkspaceFile({ + file_path, + workspace_id, + start_line, + max_lines, + codeApiBaseUrl, + executionProfile, + bridgeWorkerId, + req, + signal, +}) { + const authHeaders = await getCodeApiAuthHeaders(req, bridgeWorkerId); + return executeWorkspaceTool({ + baseURL: codeApiBaseUrl, + authHeaders: { + ...authHeaders, + ...codeExecutionHeaders({ executionProfile, bridgeWorkerId }), + }, + request: { + protocolVersion: 1, + operation: 'read_file', + workspaceId: workspace_id, + path: file_path, + startLine: start_line, + maxLines: max_lines, + }, + ...(signal ? { signal } : {}), + }); +} + /** * Reads a small code artifact as base64 so `read_file` can surface it to * vision-capable models. Reuses bytes fetched by the current request's @@ -1929,6 +1976,7 @@ module.exports = { getSessionInfo, processCodeOutput, prepareCodeOutputForInspection, + readWorkspaceFile, readSandboxFile, readSandboxImage, writeSandboxFile, diff --git a/api/server/services/Files/Code/process.spec.js b/api/server/services/Files/Code/process.spec.js index 7865a6eac2..b6cf887bba 100644 --- a/api/server/services/Files/Code/process.spec.js +++ b/api/server/services/Files/Code/process.spec.js @@ -48,6 +48,7 @@ const mockExtractCodeArtifactInspectionText = jest.fn(async () => ({ complete: false, })); const mockExtractCodeArtifactText = jest.fn(async () => null); +const mockExecuteWorkspaceTool = jest.fn(); const mockGetExtractedTextFormat = jest.fn((_name, _mime, text) => (text == null ? null : 'text')); /* `hasOfficeHtmlPath` gates the persist-then-render split: when true, processCodeOutput * returns `{ file, finalize }` with the file persisted at `status: 'pending'` @@ -76,6 +77,7 @@ jest.mock('@librechat/api', () => { sanitizeArtifactPath: jest.fn((name) => name), flattenArtifactPath: jest.fn((name) => name.replace(/\//g, '__')), createAxiosInstance: jest.fn(() => mockAxios), + executeWorkspaceTool: (...args) => mockExecuteWorkspaceTool(...args), getCodeApiAuthHeaders: jest.fn(async () => ({})), /* Windowing, sizing and rate-limit policy are real code in * `packages/api` with their own tests (`files/code/image.spec.ts`, @@ -246,6 +248,7 @@ const { prepareCodeOutputForInspection, getSessionInfo, readSandboxFile, + readWorkspaceFile, readSandboxImage, writeSandboxFile, primeFiles, @@ -1838,6 +1841,57 @@ describe('Code Process', () => { }); }); + describe('readWorkspaceFile', () => { + it('forwards authenticated reads to the selected attached worker', async () => { + const controller = new AbortController(); + const result = { + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'primary', + path: 'src/app.ts', + content: 'const ready = true;', + startLine: 1, + endLine: 1, + truncated: false, + }; + getCodeApiAuthHeaders.mockResolvedValueOnce({ Authorization: 'Bearer workspace-token' }); + mockExecuteWorkspaceTool.mockResolvedValueOnce(result); + + await expect( + readWorkspaceFile({ + file_path: 'src/app.ts', + workspace_id: 'primary', + start_line: 1, + max_lines: 200, + codeApiBaseUrl: 'https://attached-code.example.com/v1', + executionProfile: 'stateful', + bridgeWorkerId: 'worker-user-1', + req: mockReq, + signal: controller.signal, + }), + ).resolves.toBe(result); + + expect(getCodeApiAuthHeaders).toHaveBeenCalledWith(mockReq, 'worker-user-1'); + expect(mockExecuteWorkspaceTool).toHaveBeenCalledWith({ + baseURL: 'https://attached-code.example.com/v1', + authHeaders: { + Authorization: 'Bearer workspace-token', + 'X-CodeAPI-Expected-Profile': 'stateful', + 'X-LibreChat-Code-Worker-ID': 'worker-user-1', + }, + request: { + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'primary', + path: 'src/app.ts', + startLine: 1, + maxLines: 200, + }, + signal: controller.signal, + }); + }); + }); + describe('readSandboxFile', () => { /** * `readSandboxFile` shells `cat ` through the sandbox diff --git a/packages/api/src/agents/handlers.spec.ts b/packages/api/src/agents/handlers.spec.ts index 0bdb3cfd8d..fcf9ff0d4a 100644 --- a/packages/api/src/agents/handlers.spec.ts +++ b/packages/api/src/agents/handlers.spec.ts @@ -4693,6 +4693,7 @@ describe('createToolExecuteHandler', () => { skillAuthoringAvailable?: boolean; codeExecutionContext?: CodeExecutionContext; req?: unknown; + readWorkspaceFile?: ToolExecuteOptions['readWorkspaceFile']; readSandboxFile?: ToolExecuteOptions['readSandboxFile']; readSandboxImage?: ToolExecuteOptions['readSandboxImage']; getSkillByName?: ToolExecuteOptions['getSkillByName']; @@ -4714,11 +4715,316 @@ describe('createToolExecuteHandler', () => { loadTools, getSkillByName: params.getSkillByName, getAuthorSkillByName: params.getAuthorSkillByName, + readWorkspaceFile: params.readWorkspaceFile, readSandboxFile: params.readSandboxFile, readSandboxImage: params.readSandboxImage, }); } + it('routes explicit workspace paths to the selected attached worker', async () => { + const readWorkspaceFile = jest.fn(async () => ({ + protocolVersion: 1 as const, + operation: 'read_file' as const, + workspaceId: 'primary', + path: 'src/app.ts', + content: 'const ready = true;', + startLine: 1, + endLine: 1, + truncated: false, + })); + const readSandboxFile = jest.fn(); + const handler = makeReadFileHandler({ + codeEnvAvailable: true, + accessibleSkillIds: skillsInScope(), + codeExecutionContext: { + baseUrl: 'https://code.example.com/v1', + codeSessionKey: 'execute_code:stateful:attached', + executionProfile: 'stateful', + environmentType: 'attached', + bridgeWorkerId: 'personal-worker-1', + statefulSessions: true, + }, + readWorkspaceFile, + readSandboxFile, + }); + + const [result] = await invokeHandler(handler, [ + { + id: 'call_workspace_read', + name: Constants.READ_FILE, + args: { path: 'workspace/src/app.ts' }, + }, + ]); + + expect(readWorkspaceFile).toHaveBeenCalledWith({ + file_path: 'src/app.ts', + workspace_id: 'primary', + start_line: 1, + max_lines: 200, + codeApiBaseUrl: 'https://code.example.com/v1', + executionProfile: 'stateful', + bridgeWorkerId: 'personal-worker-1', + }); + expect(readSandboxFile).not.toHaveBeenCalled(); + expect(result).toMatchObject({ + status: 'success', + content: '1 | const ready = true;', + }); + }); + + it('forwards the run abort signal to attached workspace reads', async () => { + const readWorkspaceFile = jest.fn(async () => ({ + protocolVersion: 1 as const, + operation: 'read_file' as const, + workspaceId: 'primary', + path: 'notes.txt', + content: 'ready', + startLine: 1, + endLine: 1, + truncated: false, + })); + const handler = makeReadFileHandler({ + codeEnvAvailable: true, + codeExecutionContext: { + baseUrl: 'https://code.example.com/v1', + codeSessionKey: 'execute_code:stateful:attached', + executionProfile: 'stateful', + environmentType: 'attached', + statefulSessions: true, + }, + readWorkspaceFile, + }); + const controller = new AbortController(); + + await new Promise((resolve, reject) => { + handler.handle('on_tool_execute', { + toolCalls: [ + { + id: 'call_workspace_abort', + name: Constants.READ_FILE, + args: { path: 'workspace/notes.txt' }, + }, + ], + signal: controller.signal, + resolve, + reject, + } as ToolExecuteBatchRequest); + }); + + expect(readWorkspaceFile).toHaveBeenCalledWith( + expect.objectContaining({ signal: controller.signal }), + ); + }); + + it('uses bounded pagination for attached workspace reads', async () => { + const readWorkspaceFile = jest.fn(async () => ({ + protocolVersion: 1 as const, + operation: 'read_file' as const, + workspaceId: 'primary', + path: 'notes.txt', + content: 'third\nfourth', + startLine: 3, + endLine: 4, + truncated: true, + nextStartLine: 5, + })); + const handler = makeReadFileHandler({ + codeEnvAvailable: true, + codeExecutionContext: { + baseUrl: 'https://code.example.com/v1', + codeSessionKey: 'execute_code:stateful:attached', + executionProfile: 'stateful', + environmentType: 'attached', + bridgeWorkerId: 'personal-worker-1', + statefulSessions: true, + }, + readWorkspaceFile, + }); + + const [result] = await invokeHandler(handler, [ + { + id: 'call_workspace_page', + name: Constants.READ_FILE, + args: { path: 'workspace/notes.txt', start_line: 3, max_lines: 2 }, + }, + ]); + + expect(readWorkspaceFile).toHaveBeenCalledWith( + expect.objectContaining({ start_line: 3, max_lines: 2 }), + ); + expect(result.content).toContain('3 | third\n4 | fourth'); + expect(result.content).toContain('start_line 5'); + }); + + it('continues from the first line omitted by local workspace truncation', async () => { + const firstLine = 'a'.repeat(140_000); + const secondLine = 'b'.repeat(140_000); + const readWorkspaceFile = jest.fn(async () => ({ + protocolVersion: 1 as const, + operation: 'read_file' as const, + workspaceId: 'primary', + path: 'notes.txt', + content: `${firstLine}\n${secondLine}`, + startLine: 10, + endLine: 11, + truncated: true, + nextStartLine: 12, + })); + const handler = makeReadFileHandler({ + codeEnvAvailable: true, + codeExecutionContext: { + baseUrl: 'https://code.example.com/v1', + codeSessionKey: 'execute_code:stateful:attached', + executionProfile: 'stateful', + environmentType: 'attached', + statefulSessions: true, + }, + readWorkspaceFile, + }); + + const [result] = await invokeHandler(handler, [ + { + id: 'call_workspace_local_page', + name: Constants.READ_FILE, + args: { path: 'workspace/notes.txt', start_line: 10, max_lines: 2 }, + }, + ]); + + expect(result.content).toContain('10 | '); + expect(result.content).not.toContain('11 | '); + expect(result.content).toContain('start_line 11'); + expect(result.content).not.toContain('start_line 12'); + }); + + it('applies local workspace truncation in UTF-8 bytes', async () => { + const firstLine = '界'.repeat(80_000); + const secondLine = '界'.repeat(10_000); + const readWorkspaceFile = jest.fn(async () => ({ + protocolVersion: 1 as const, + operation: 'read_file' as const, + workspaceId: 'primary', + path: 'multibyte.txt', + content: `${firstLine}\n${secondLine}`, + startLine: 1, + endLine: 2, + truncated: false, + })); + const handler = makeReadFileHandler({ + codeEnvAvailable: true, + codeExecutionContext: { + baseUrl: 'https://code.example.com/v1', + codeSessionKey: 'execute_code:stateful:attached', + executionProfile: 'stateful', + environmentType: 'attached', + statefulSessions: true, + }, + readWorkspaceFile, + }); + + const [result] = await invokeHandler(handler, [ + { + id: 'call_workspace_multibyte', + name: Constants.READ_FILE, + args: { path: 'workspace/multibyte.txt', max_lines: 2 }, + }, + ]); + + expect(result.content).toContain('1 | '); + expect(result.content).not.toContain('2 | '); + expect(result.content).toContain('start_line 2'); + expect(Buffer.byteLength(result.content as string, 'utf8')).toBeLessThan(262_300); + }); + + it('rejects workspace paths unless the selected environment is attached', async () => { + const readWorkspaceFile = jest.fn(); + const handler = makeReadFileHandler({ + codeEnvAvailable: true, + codeExecutionContext: { + baseUrl: 'https://code.example.com/v1', + codeSessionKey: 'execute_code:stateful:managed', + executionProfile: 'stateful', + environmentType: 'managed', + statefulSessions: true, + }, + readWorkspaceFile, + }); + + const [result] = await invokeHandler(handler, [ + { + id: 'call_managed_workspace_read', + name: Constants.READ_FILE, + args: { path: 'workspace/src/app.ts' }, + }, + ]); + + expect(result).toMatchObject({ + status: 'error', + errorMessage: 'workspace/ paths require an attached code environment.', + }); + expect(readWorkspaceFile).not.toHaveBeenCalled(); + }); + + it('rejects empty and unbounded attached workspace reads before dispatch', async () => { + const readWorkspaceFile = jest.fn(); + const handler = makeReadFileHandler({ + codeEnvAvailable: true, + codeExecutionContext: { + baseUrl: 'https://code.example.com/v1', + codeSessionKey: 'execute_code:stateful:attached', + executionProfile: 'stateful', + environmentType: 'attached', + bridgeWorkerId: 'personal-worker-1', + statefulSessions: true, + }, + readWorkspaceFile, + }); + + const [emptyPath] = await invokeHandler(handler, [ + { id: 'call_empty_workspace', name: Constants.READ_FILE, args: { path: 'workspace/' } }, + ]); + const [unbounded] = await invokeHandler(handler, [ + { + id: 'call_unbounded_workspace', + name: Constants.READ_FILE, + args: { path: 'workspace/src/app.ts', max_lines: 501 }, + }, + ]); + + expect(emptyPath.errorMessage).toContain('relative path'); + expect(unbounded.errorMessage).toContain('between 1 and 500'); + expect(readWorkspaceFile).not.toHaveBeenCalled(); + }); + + it('does not expose attached-worker failure details to the model', async () => { + const readWorkspaceFile = jest.fn(async () => { + throw new Error('upstream response contained /Users/operator/private'); + }); + const handler = makeReadFileHandler({ + codeEnvAvailable: true, + codeExecutionContext: { + baseUrl: 'https://code.example.com/v1', + codeSessionKey: 'execute_code:stateful:attached', + executionProfile: 'stateful', + environmentType: 'attached', + bridgeWorkerId: 'personal-worker-1', + statefulSessions: true, + }, + readWorkspaceFile, + }); + + const [result] = await invokeHandler(handler, [ + { + id: 'call_failed_workspace', + name: Constants.READ_FILE, + args: { path: 'workspace/src/app.ts' }, + }, + ]); + + expect(result.status).toBe('error'); + expect(result.errorMessage).toContain('could not be read'); + expect(result.errorMessage).not.toContain('/Users/operator/private'); + }); + it('routes /mnt/data/ paths to the sandbox fallback when codeEnv is available', async () => { const readSandboxFile = jest.fn(async () => ({ content: 'hello-world' })); const handler = makeReadFileHandler({ diff --git a/packages/api/src/agents/handlers.ts b/packages/api/src/agents/handlers.ts index 5b6b39ca8c..abfadfe056 100644 --- a/packages/api/src/agents/handlers.ts +++ b/packages/api/src/agents/handlers.ts @@ -24,6 +24,7 @@ import type { BackgroundToolWakeupRegistration, } from './backgroundCompletion'; import type { SkillFileRecord, PrimeSkillFilesResult } from './skillFiles'; +import type { WorkspaceReadResult } from '~/code/workspace'; import type { BackgroundToolResultState } from './harvest'; import type { CodeExecutionContext } from './execution'; import type { TextContentFragment } from '~/protection'; @@ -442,6 +443,18 @@ export interface ToolExecuteOptions { relativePath: string, update: { content?: string; isBinary?: boolean }, ) => Promise; + /** Reads a bounded text range from an attached worker's logical workspace. */ + readWorkspaceFile?: (params: { + file_path: string; + workspace_id: string; + start_line: number; + max_lines: number; + codeApiBaseUrl: string; + executionProfile: CodeExecutionContext['executionProfile']; + bridgeWorkerId?: string; + req?: ServerRequest; + signal?: AbortSignal; + }) => Promise; /** * Reads a code-execution sandbox file by shelling `cat` through the * sandbox `/exec` endpoint. The host implementation supplies the @@ -522,6 +535,19 @@ export interface ToolExecuteOptions { const MAX_READABLE_BYTES = 262_144; const MAX_BINARY_BYTES = 5 * 1024 * 1024; + +function truncateUtf8(value: string, maxBytes: number): string { + const bytes = Buffer.from(value, 'utf8'); + if (bytes.byteLength <= maxBytes) { + return value; + } + let end = maxBytes; + while (end > 0 && (bytes[end] & 0xc0) === 0x80) { + end -= 1; + } + return bytes.subarray(0, end).toString('utf8'); +} + /** * Inline ceiling for images pulled out of the code-execution sandbox — * deliberately tighter than {@link MAX_BINARY_BYTES}, which governs the @@ -754,10 +780,12 @@ function getValueShape(value: unknown): string { return typeof value; } -function addLineNumbers(content: string): string { +function addLineNumbers(content: string, startLine = 1): string { const lines = content.split('\n'); - const w = String(lines.length).length; - return lines.map((l, i) => `${String(i + 1).padStart(w, ' ')} | ${l}`).join('\n'); + const w = String(startLine + lines.length - 1).length; + return lines + .map((line, index) => `${String(startLine + index).padStart(w, ' ')} | ${line}`) + .join('\n'); } type AuthoringSkill = NonNullable< @@ -2143,6 +2171,120 @@ async function handleSandboxFileFallback( } } +async function handleWorkspaceFileRead( + tc: ToolCallRequest, + filePath: string, + options: ToolExecuteOptions, + req: ServerRequest | undefined, + codeExecutionContext: CodeExecutionContext, + signal?: AbortSignal, +): Promise { + const { readWorkspaceFile } = options; + if (!readWorkspaceFile) { + return { + toolCallId: tc.id, + status: 'error', + content: '', + errorMessage: 'Attached workspace reading is not configured.', + }; + } + const args = tc.args as { start_line?: number; max_lines?: number }; + const startLine = args.start_line ?? 1; + const maxLines = args.max_lines ?? 200; + if (filePath.length === 0) { + return { + toolCallId: tc.id, + status: 'error', + content: '', + errorMessage: 'A relative path after workspace/ is required.', + }; + } + if ( + !Number.isSafeInteger(startLine) || + startLine < 1 || + !Number.isSafeInteger(maxLines) || + maxLines < 1 || + maxLines > 500 + ) { + return { + toolCallId: tc.id, + status: 'error', + content: '', + errorMessage: 'start_line must be positive and max_lines must be between 1 and 500.', + }; + } + const filteredName = filteredFileNameResult(tc, req, filePath); + if (filteredName != null) { + return filteredName; + } + + try { + const result = await readWorkspaceFile({ + file_path: filePath, + workspace_id: 'primary', + start_line: startLine, + max_lines: maxLines, + codeApiBaseUrl: codeExecutionContext.baseUrl, + executionProfile: codeExecutionContext.executionProfile, + ...(codeExecutionContext.bridgeWorkerId + ? { bridgeWorkerId: codeExecutionContext.bridgeWorkerId } + : {}), + ...(req ? { req } : {}), + ...(signal ? { signal } : {}), + }); + const filtered = filteredFileResult(tc, req, filePath, result.content); + if (filtered != null) { + return filtered; + } + if (looksBinary(result.content)) { + return { + toolCallId: tc.id, + status: 'error', + content: '', + errorMessage: `"${filePath}" appears to be a binary file and cannot be read as text.`, + }; + } + let payload = result.content; + let locallyTruncated = false; + let localNextStartLine: number | undefined; + if (Buffer.byteLength(payload, 'utf8') > MAX_READABLE_BYTES) { + payload = truncateUtf8(payload, MAX_READABLE_BYTES); + locallyTruncated = true; + const lastCompleteLine = payload.lastIndexOf('\n'); + if (lastCompleteLine >= 0) { + payload = payload.slice(0, lastCompleteLine); + localNextStartLine = result.startLine + payload.split('\n').length; + } + } + let numbered = addLineNumbers(payload, result.startLine); + if (locallyTruncated) { + numbered += + localNextStartLine != null + ? `\n\n[truncated at ${MAX_READABLE_BYTES} bytes; more content is available; call read_file again with path "workspace/${filePath}" and start_line ${localNextStartLine}]` + : `\n\n[the line was truncated at ${MAX_READABLE_BYTES} bytes and cannot be paged by line]`; + } else if (result.truncated && result.nextStartLine != null) { + numbered += `\n\n[more content is available; call read_file again with path "workspace/${filePath}" and start_line ${result.nextStartLine}]`; + } + return { + toolCallId: tc.id, + status: 'success', + content: numbered, + }; + } catch (error) { + if (signal?.aborted === true && isAbortError(error)) throw error; + logger.warn( + '[handleWorkspaceFileRead] Attached workspace read failed', + getSafeErrorMetadata(error), + ); + return { + toolCallId: tc.id, + status: 'error', + content: '', + errorMessage: `"${filePath}" could not be read from the attached workspace.`, + }; + } +} + function sandboxSessionContext( tc: ToolCallRequest, override?: SandboxSessionContext, @@ -3580,6 +3722,7 @@ async function handleReadFileCall( options: ToolExecuteOptions, req?: ServerRequest, onSandboxReadSuccess?: () => void, + signal?: AbortSignal, ): Promise { const { getSkillByName, getSkillFileByPath, getStrategyFunctions, updateSkillFileContent } = options; @@ -3597,6 +3740,25 @@ async function handleReadFileCall( const codeExecutionContext = getCodeExecutionContext(mergedConfigurable); let accessibleIds = (mergedConfigurable?.accessibleSkillIds as Types.ObjectId[]) ?? []; + if (args.path.startsWith('workspace/')) { + if (!codeEnvAvailable || codeExecutionContext?.environmentType !== 'attached') { + return { + toolCallId: tc.id, + status: 'error', + content: '', + errorMessage: 'workspace/ paths require an attached code environment.', + }; + } + return handleWorkspaceFileRead( + tc, + args.path.slice('workspace/'.length), + options, + req, + codeExecutionContext, + signal, + ); + } + /** * Short-circuit absolute code-env paths: the path can never be a skill * reference (skill paths are relative `{skillName}/...`), and consulting @@ -5633,6 +5795,7 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand () => { sandboxReadSucceeded = true; }, + runSignal, ); } else if (tc.name === CREATE_FILE_TOOL_NAME && isFileAuthoringCall) { handlerResult = await handleCreateFileCall( @@ -5673,10 +5836,18 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand }); return filteredError; } - logger.error(`[ON_TOOL_EXECUTE] Tool ${tc.name} error`, { + const context = { ...logContext, toolCallArgsShape: getValueShape(tc.args), - }); + }; + if (runSignal?.aborted === true && isAbortError(toolError)) { + logger.debug( + `[ON_TOOL_EXECUTE] Tool ${tc.name} cancelled by run abort`, + context, + ); + } else { + logger.error(`[ON_TOOL_EXECUTE] Tool ${tc.name} error`, context); + } return { toolCallId: tc.id, status: 'error' as const, diff --git a/packages/api/src/agents/initialize.ts b/packages/api/src/agents/initialize.ts index 61b762eb0a..dc43ff1ccc 100644 --- a/packages/api/src/agents/initialize.ts +++ b/packages/api/src/agents/initialize.ts @@ -996,6 +996,8 @@ export async function initializeAgent( agentId: agent.id, conversationId, }); + const attachedWorkspaceTools = + effectiveCodeEnvAvailable && codeExecutionContext.environmentType === 'attached'; const requestFileIds = [ ...new Set( requestFiles @@ -1475,6 +1477,7 @@ export async function initializeAgent( includeSkillFileInstructions: false, enableToolOutputReferences: effectiveCodeEnvAvailable, statefulSessions: effectiveStatefulSessions, + workspaceTools: attachedWorkspaceTools, }); toolDefinitions = codeExecResult.toolDefinitions; recordCapabilityToolNames(AgentCapabilities.execute_code, codeExecResult.toolNames); @@ -1523,6 +1526,7 @@ export async function initializeAgent( includeBash: false, includeSkillFileInstructions: true, enableToolOutputReferences: effectiveCodeEnvAvailable, + workspaceTools: attachedWorkspaceTools, }); toolDefinitions = skillReadResult.toolDefinitions; recordCapabilityToolNames(AgentCapabilities.skills, skillReadResult.toolNames); @@ -1683,6 +1687,7 @@ export async function initializeAgent( listSkillsByAccess: db?.listSkillsByAccess, codeEnvAvailable: effectiveCodeEnvAvailable, statefulSessions: effectiveStatefulSessions, + workspaceTools: attachedWorkspaceTools, userId: user?.id, skillStates: params.skillStates, defaultActiveOnShare: params.defaultActiveOnShare, diff --git a/packages/api/src/agents/skills.ts b/packages/api/src/agents/skills.ts index 338ff36829..a495d39c9f 100644 --- a/packages/api/src/agents/skills.ts +++ b/packages/api/src/agents/skills.ts @@ -430,6 +430,8 @@ export interface InjectSkillCatalogParams { codeEnvAvailable?: boolean; /** When true, bash_tool registers with the hedged stateful-session description. */ statefulSessions?: boolean; + /** When true, read_file exposes the attached worker's workspace namespace. */ + workspaceTools?: boolean; /** Current user ID — used to determine skill ownership for active-state resolution. */ userId?: string; /** Per-user skill overrides: `{ [skillId]: boolean }`. Missing entries use the default. */ @@ -656,6 +658,7 @@ export async function injectSkillCatalog( listSkillsByAccess, codeEnvAvailable, statefulSessions, + workspaceTools, userId, skillStates, defaultActiveOnShare = false, @@ -813,6 +816,7 @@ export async function injectSkillCatalog( includeBash: codeEnvAvailable === true, enableToolOutputReferences: codeEnvAvailable === true, statefulSessions: statefulSessions === true, + workspaceTools: workspaceTools === true, }); workingDefs = codeExecResult.toolDefinitions; diff --git a/packages/api/src/agents/tools.spec.ts b/packages/api/src/agents/tools.spec.ts index 7903d302eb..6b10f044fd 100644 --- a/packages/api/src/agents/tools.spec.ts +++ b/packages/api/src/agents/tools.spec.ts @@ -474,6 +474,26 @@ describe('registerCodeExecutionTools', () => { expect(JSON.stringify(readFile?.parameters)).not.toContain('{skillName}'); }); + it('advertises explicit workspace paths and pagination for attached environments', () => { + const result = registerCodeExecutionTools({ + toolRegistry: makeRegistry(), + toolDefinitions: [], + includeBash: true, + includeSkillFileInstructions: false, + workspaceTools: true, + }); + + const readFile = result.toolDefinitions.find((definition) => definition.name === 'read_file'); + expect(readFile?.description).toContain('workspace/'); + expect(readFile?.description).toContain('attached'); + expect(readFile?.parameters).toMatchObject({ + properties: { + start_line: { type: 'integer' }, + max_lines: { type: 'integer', maximum: 500 }, + }, + }); + }); + it('upgrades a code-only read_file definition when skills are enabled later in the run', () => { const toolRegistry = makeRegistry(); const codeOnly = registerCodeExecutionTools({ @@ -497,6 +517,33 @@ describe('registerCodeExecutionTools', () => { expect(toolRegistry.get('read_file')?.description).toBe(readFile?.description); }); + it('preserves attached workspace instructions when skills upgrade read_file', () => { + const toolRegistry = makeRegistry(); + const codeOnly = registerCodeExecutionTools({ + toolRegistry, + toolDefinitions: [], + includeBash: true, + includeSkillFileInstructions: false, + workspaceTools: true, + }); + const upgraded = registerCodeExecutionTools({ + toolRegistry, + toolDefinitions: codeOnly.toolDefinitions, + includeBash: false, + includeSkillFileInstructions: true, + workspaceTools: true, + }); + + const readFile = upgraded.toolDefinitions.find( + (definition) => definition.name === 'read_file', + ); + expect(readFile?.description).toContain('skills/{skillName}/'); + expect(readFile?.description).toContain('workspace/'); + expect(readFile?.parameters).toMatchObject({ + properties: { max_lines: { maximum: 500 } }, + }); + }); + it('preserves pre-existing unrelated tool definitions', () => { const toolRegistry = makeRegistry(); const existing: LCTool[] = [ diff --git a/packages/api/src/agents/tools.ts b/packages/api/src/agents/tools.ts index afcae5c8e9..84fad0a118 100644 --- a/packages/api/src/agents/tools.ts +++ b/packages/api/src/agents/tools.ts @@ -341,6 +341,11 @@ export interface RegisterCodeExecutionToolsParams { * prompted to discover skills that are not enabled for the run. */ includeSkillFileInstructions?: boolean; + /** + * When `true`, `read_file` advertises the explicit `workspace/` namespace + * backed by the selected attached worker, including bounded line pagination. + */ + workspaceTools?: boolean; /** * When `true`, the registered `bash_tool` description includes the * LLM-facing `{{toolturn}}` reference syntax guide so the @@ -406,6 +411,8 @@ const CODE_READ_FILE_DESCRIPTION = `Read a known file from the code-execution sa Use for text, CSV, JSON, Markdown, logs, small source files, and images at paths returned by tool output, just written, or under /mnt/data/. Do not run ls/find just to rediscover known paths. Use bash_tool for other binary files, large files, transforms, metadata, or true filesystem discovery. /tmp is per-call scratch and unavailable later.`; +const ATTACHED_WORKSPACE_READ_FILE_INSTRUCTIONS = `For an attached environment, use "workspace/{relativePath}" to read a file from the workspace directory registered on the worker. It may be an existing project, a Git repository, or an empty directory; Git is not required. The worker's host path stays private. Use start_line and max_lines for bounded pagination.`; + const CODE_READ_FILE_PARAMETERS: LCTool['parameters'] = Object.freeze({ type: 'object', properties: { @@ -418,6 +425,29 @@ const CODE_READ_FILE_PARAMETERS: LCTool['parameters'] = Object.freeze({ required: ['path'], }) as LCTool['parameters']; +const ATTACHED_WORKSPACE_READ_FILE_PARAMETERS: LCTool['parameters'] = Object.freeze({ + type: 'object', + properties: { + path: { + type: 'string', + description: + 'Use "workspace/{relativePath}" for a file in the attached worker workspace directory, or a code-execution sandbox path such as "/mnt/data/result.csv".', + }, + start_line: { + type: 'integer', + minimum: 1, + description: 'Optional one-based line at which to start reading a workspace text file.', + }, + max_lines: { + type: 'integer', + minimum: 1, + maximum: 500, + description: 'Optional maximum number of workspace text-file lines to return.', + }, + }, + required: ['path'], +}) as LCTool['parameters']; + const CODE_READ_FILE_DEF: LCTool = Object.freeze({ name: ReadFileToolDefinition.name, description: CODE_READ_FILE_DESCRIPTION, @@ -425,6 +455,21 @@ const CODE_READ_FILE_DEF: LCTool = Object.freeze({ responseFormat: ReadFileToolDefinition.responseFormat, }) as LCTool; +function createAttachedWorkspaceReadFileDef(includeSkillFileInstructions: boolean): LCTool { + const baseDescription = includeSkillFileInstructions + ? SKILL_READ_FILE_DESCRIPTION + : CODE_READ_FILE_DESCRIPTION; + return Object.freeze({ + name: ReadFileToolDefinition.name, + description: `${baseDescription}\n\n${ATTACHED_WORKSPACE_READ_FILE_INSTRUCTIONS}`, + parameters: ATTACHED_WORKSPACE_READ_FILE_PARAMETERS, + responseFormat: ReadFileToolDefinition.responseFormat, + }) as LCTool; +} + +const ATTACHED_CODE_READ_FILE_DEF = createAttachedWorkspaceReadFileDef(false); +const ATTACHED_SKILL_READ_FILE_DEF = createAttachedWorkspaceReadFileDef(true); + const SKILL_CREATE_FILE_PARAMETERS: LCTool['parameters'] = Object.freeze({ type: 'object', properties: { @@ -597,7 +642,12 @@ const CODE_EDIT_FILE_DEF: LCTool = Object.freeze({ responseFormat: 'content_and_artifact' as LCTool['responseFormat'], }) as LCTool; -function buildReadFileDef(includeSkillFileInstructions: boolean): LCTool { +function buildReadFileDef(includeSkillFileInstructions: boolean, workspaceTools: boolean): LCTool { + if (workspaceTools) { + return includeSkillFileInstructions + ? ATTACHED_SKILL_READ_FILE_DEF + : ATTACHED_CODE_READ_FILE_DEF; + } return includeSkillFileInstructions ? READ_FILE_DEF : CODE_READ_FILE_DEF; } @@ -609,7 +659,9 @@ function buildFileAuthoringDefs(includeSkillFileInstructions: boolean): LCTool[] function isCodeOnlyReadFileDef(def: LCTool | undefined): boolean { return ( - def?.name === ReadFileToolDefinition.name && def?.description === CODE_READ_FILE_DESCRIPTION + def?.name === ReadFileToolDefinition.name && + (def?.description === CODE_READ_FILE_DESCRIPTION || + def?.description === ATTACHED_CODE_READ_FILE_DEF.description) ); } @@ -701,11 +753,12 @@ export function registerCodeExecutionTools( toolDefinitions, includeBash, includeSkillFileInstructions = true, + workspaceTools = false, enableToolOutputReferences = false, statefulSessions = false, } = params; - const readFileDef = buildReadFileDef(includeSkillFileInstructions); + const readFileDef = buildReadFileDef(includeSkillFileInstructions, workspaceTools); const candidates: LCTool[] = includeBash ? [readFileDef, buildBashToolDef({ enableToolOutputReferences, statefulSessions })] : [readFileDef]; diff --git a/packages/api/src/code/index.ts b/packages/api/src/code/index.ts index 787a58b1de..7b1bdf229f 100644 --- a/packages/api/src/code/index.ts +++ b/packages/api/src/code/index.ts @@ -3,3 +3,4 @@ export * from './http'; export * from './config'; export * from './bridge'; export * from './lifecycle'; +export * from './workspace'; diff --git a/packages/api/src/code/workspace.spec.ts b/packages/api/src/code/workspace.spec.ts new file mode 100644 index 0000000000..b2806747be --- /dev/null +++ b/packages/api/src/code/workspace.spec.ts @@ -0,0 +1,358 @@ +import type { CodeBridgeFetch } from './bridge'; +import { executeWorkspaceTool } from './workspace'; + +describe('executeWorkspaceTool', () => { + test('sends an authenticated bounded read to the selected attached worker', async () => { + const fetchImpl = jest.fn().mockResolvedValue( + new Response( + JSON.stringify({ + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'primary', + path: 'src/app.ts', + content: 'const ready = true;', + startLine: 1, + endLine: 1, + truncated: false, + }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ), + ); + + await expect( + executeWorkspaceTool({ + baseURL: 'https://code.example.com/v1/', + authHeaders: { + Authorization: 'Bearer jwt', + 'X-LibreChat-Code-Worker-ID': 'user-worker', + }, + request: { + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'primary', + path: 'src/app.ts', + startLine: 1, + maxLines: 200, + }, + fetchImpl, + }), + ).resolves.toMatchObject({ content: 'const ready = true;' }); + + expect(fetchImpl).toHaveBeenCalledWith( + 'https://code.example.com/v1/workspace-tools/execute', + expect.objectContaining({ + method: 'POST', + headers: expect.objectContaining({ + Authorization: 'Bearer jwt', + 'Content-Type': 'application/json', + 'X-LibreChat-Code-Worker-ID': 'user-worker', + }), + }), + ); + }); + + test('combines the caller abort signal with the workspace request timeout', async () => { + const controller = new AbortController(); + let requestSignal: AbortSignal | undefined; + const fetchImpl: CodeBridgeFetch = jest.fn(async (_url, init) => { + requestSignal = init?.signal ?? undefined; + return new Response( + JSON.stringify({ + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'primary', + path: 'notes.txt', + content: 'ready', + startLine: 1, + endLine: 1, + truncated: false, + }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + }); + + await executeWorkspaceTool({ + baseURL: 'https://code.example.com/v1', + authHeaders: { Authorization: 'Bearer jwt' }, + signal: controller.signal, + request: { + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'primary', + path: 'notes.txt', + }, + fetchImpl, + }); + + expect(requestSignal?.aborted).toBe(false); + controller.abort(); + expect(requestSignal?.aborted).toBe(true); + }); + + test('preserves caller cancellation instead of relabeling it as a transport failure', async () => { + const controller = new AbortController(); + const fetchImpl: CodeBridgeFetch = jest.fn( + async (_url, init) => + await new Promise((_resolve, reject) => { + init?.signal?.addEventListener('abort', () => reject(init.signal?.reason), { + once: true, + }); + }), + ); + + const request = executeWorkspaceTool({ + baseURL: 'https://code.example.com/v1', + authHeaders: { Authorization: 'Bearer jwt' }, + signal: controller.signal, + request: { + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'primary', + path: 'notes.txt', + }, + fetchImpl, + }); + controller.abort(); + + await expect(request).rejects.toMatchObject({ name: 'AbortError' }); + }); + + test('preserves caller cancellation while reading the response body', async () => { + const controller = new AbortController(); + const fetchImpl: CodeBridgeFetch = jest.fn(async (_url, init) => { + const body = new ReadableStream({ + start(streamController) { + init?.signal?.addEventListener( + 'abort', + () => streamController.error(init.signal?.reason), + { once: true }, + ); + }, + }); + return new Response(body, { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }); + }); + + const request = executeWorkspaceTool({ + baseURL: 'https://code.example.com/v1', + authHeaders: { Authorization: 'Bearer jwt' }, + signal: controller.signal, + request: { + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'primary', + path: 'notes.txt', + }, + fetchImpl, + }); + await new Promise((resolve) => setImmediate(resolve)); + controller.abort(); + + await expect(request).rejects.toMatchObject({ name: 'AbortError' }); + }); + + test('rejects a malformed worker result before it reaches the model', async () => { + const fetchImpl = jest.fn().mockResolvedValue( + new Response( + JSON.stringify({ + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'primary', + path: '/Users/operator/.ssh/id_ed25519', + content: 'secret', + startLine: 1, + endLine: 1, + truncated: false, + }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ), + ); + + await expect( + executeWorkspaceTool({ + baseURL: 'https://code.example.com/v1', + authHeaders: { Authorization: 'Bearer jwt' }, + request: { + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'primary', + path: 'src/app.ts', + }, + fetchImpl, + }), + ).rejects.toMatchObject({ reason: 'invalid' }); + }); + + test.each(['./src/app.ts', 'src/./app.ts', 'src//app.ts'])( + 'rejects the non-canonical request path %s before dispatch', + async (path) => { + const fetchImpl = jest.fn(); + + await expect( + executeWorkspaceTool({ + baseURL: 'https://code.example.com/v1', + authHeaders: { Authorization: 'Bearer jwt' }, + request: { + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'primary', + path, + }, + fetchImpl, + }), + ).rejects.toMatchObject({ reason: 'invalid' }); + expect(fetchImpl).not.toHaveBeenCalled(); + }, + ); + + test('rejects read content that exceeds its declared line range', async () => { + const fetchImpl = jest.fn().mockResolvedValue( + new Response( + JSON.stringify({ + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'primary', + path: 'src/app.ts', + content: 'first\nsecond', + startLine: 1, + endLine: 1, + truncated: false, + }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ), + ); + + await expect( + executeWorkspaceTool({ + baseURL: 'https://code.example.com/v1', + authHeaders: { Authorization: 'Bearer jwt' }, + request: { + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'primary', + path: 'src/app.ts', + maxLines: 1, + }, + fetchImpl, + }), + ).rejects.toMatchObject({ reason: 'invalid' }); + }); + + test('rejects unexpected result fields that could disclose worker metadata', async () => { + const fetchImpl = jest.fn().mockResolvedValue( + new Response( + JSON.stringify({ + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'primary', + path: 'src/app.ts', + content: 'safe', + startLine: 1, + endLine: 1, + truncated: false, + root: '/Users/operator/private', + }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ), + ); + + await expect( + executeWorkspaceTool({ + baseURL: 'https://code.example.com/v1', + authHeaders: { Authorization: 'Bearer jwt' }, + request: { + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'primary', + path: 'src/app.ts', + }, + fetchImpl, + }), + ).rejects.toMatchObject({ reason: 'invalid' }); + }); + + test('surfaces bounded upstream failures without returning their body', async () => { + const cancel = jest.fn(); + const fetchImpl = jest.fn().mockResolvedValue( + new Response( + new ReadableStream({ + cancel, + }), + { status: 503 }, + ), + ); + + await expect( + executeWorkspaceTool({ + baseURL: 'https://code.example.com/v1', + authHeaders: { Authorization: 'Bearer jwt' }, + request: { + protocolVersion: 1, + operation: 'search_text', + workspaceId: 'primary', + query: 'needle', + }, + fetchImpl, + }), + ).rejects.toMatchObject({ reason: 'rejected', upstreamStatus: 503 }); + expect(cancel).toHaveBeenCalledTimes(1); + }); + + test('rejects an oversized response before parsing worker-controlled JSON', async () => { + const cancel = jest.fn(); + const json = jest.fn().mockResolvedValue({ + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'primary', + path: 'src/app.ts', + content: 'safe', + startLine: 1, + endLine: 1, + truncated: false, + }); + const fetchImpl = jest.fn().mockResolvedValue({ + ok: true, + status: 200, + headers: new Headers({ 'Content-Length': String(5 * 1024 * 1024) }), + body: new ReadableStream({ cancel }), + json, + } as unknown as Response); + + await expect( + executeWorkspaceTool({ + baseURL: 'https://code.example.com/v1', + authHeaders: { Authorization: 'Bearer jwt' }, + request: { + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'primary', + path: 'src/app.ts', + }, + fetchImpl, + }), + ).rejects.toMatchObject({ reason: 'invalid' }); + expect(json).not.toHaveBeenCalled(); + expect(cancel).toHaveBeenCalledTimes(1); + }); + + test('rejects unrecognized workspace operations before dispatch', async () => { + const fetchImpl = jest.fn(); + + await expect( + executeWorkspaceTool({ + baseURL: 'https://code.example.com/v1', + authHeaders: { Authorization: 'Bearer jwt' }, + request: { + protocolVersion: 1, + operation: 'delete_file', + workspaceId: 'primary', + query: 'ignored', + } as never, + fetchImpl, + }), + ).rejects.toMatchObject({ reason: 'invalid' }); + expect(fetchImpl).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/api/src/code/workspace.ts b/packages/api/src/code/workspace.ts new file mode 100644 index 0000000000..60e1fae0b2 --- /dev/null +++ b/packages/api/src/code/workspace.ts @@ -0,0 +1,301 @@ +import type { CodeBridgeFetch } from './bridge'; + +const WORKSPACE_TOOL_TIMEOUT_MS = 30_000; +const MAX_PATH_LENGTH = 4096; +const MAX_QUERY_LENGTH = 4096; +const MAX_READ_BYTES = 1024 * 1024; +const MAX_READ_LINES = 500; +const MAX_SEARCH_RESULTS = 200; +const MAX_SEARCH_TEXT_LENGTH = 2000; +const MAX_RESPONSE_BYTES = 4 * 1024 * 1024; +const READ_RESULT_KEYS = new Set([ + 'protocolVersion', + 'operation', + 'workspaceId', + 'path', + 'content', + 'startLine', + 'endLine', + 'truncated', + 'nextStartLine', +]); +const SEARCH_RESULT_KEYS = new Set([ + 'protocolVersion', + 'operation', + 'workspaceId', + 'matches', + 'truncated', +]); +const SEARCH_MATCH_KEYS = new Set(['path', 'line', 'column', 'text']); + +export interface WorkspaceReadRequest { + protocolVersion: 1; + operation: 'read_file'; + workspaceId: string; + path: string; + startLine?: number; + maxLines?: number; +} + +export interface WorkspaceSearchRequest { + protocolVersion: 1; + operation: 'search_text'; + workspaceId: string; + query: string; + path?: string; + maxResults?: number; +} + +export type WorkspaceToolRequest = WorkspaceReadRequest | WorkspaceSearchRequest; + +export interface WorkspaceReadResult { + protocolVersion: 1; + operation: 'read_file'; + workspaceId: string; + path: string; + content: string; + startLine: number; + endLine: number; + truncated: boolean; + nextStartLine?: number; +} + +export interface WorkspaceSearchResult { + protocolVersion: 1; + operation: 'search_text'; + workspaceId: string; + matches: Array<{ path: string; line: number; column: number; text: string }>; + truncated: boolean; +} + +export type WorkspaceToolResult = WorkspaceReadResult | WorkspaceSearchResult; + +export class WorkspaceToolHttpError extends Error { + constructor( + public readonly reason: 'rejected' | 'invalid' | 'timeout' | 'failed', + public readonly upstreamStatus?: number, + ) { + super(`Workspace tool request ${reason}`); + this.name = 'WorkspaceToolHttpError'; + } +} + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value != null; +} + +function isSafePath(value: unknown): value is string { + return ( + typeof value === 'string' && + value.length > 0 && + value.length <= MAX_PATH_LENGTH && + !value.includes('\0') && + !value.includes('\\') && + !value.startsWith('/') && + !/^[A-Za-z]:/.test(value) && + value.split('/').every((segment) => segment !== '' && segment !== '.' && segment !== '..') + ); +} + +function isPositiveInteger(value: unknown, maximum: number): boolean { + return Number.isSafeInteger(value) && Number(value) >= 1 && Number(value) <= maximum; +} + +function hasOnlyKeys(value: Record, allowed: ReadonlySet): boolean { + return Object.keys(value).every((key) => allowed.has(key)); +} + +async function readBoundedJson(response: Response, signal?: AbortSignal): Promise { + const declaredLength = Number(response.headers.get('content-length')); + if (Number.isFinite(declaredLength) && declaredLength > MAX_RESPONSE_BYTES) { + await response.body?.cancel().catch(() => undefined); + throw new WorkspaceToolHttpError('invalid'); + } + + if (!response.body) { + throw new WorkspaceToolHttpError('invalid'); + } + const reader = response.body.getReader(); + const decoder = new TextDecoder(); + let byteLength = 0; + let body = ''; + try { + while (true) { + const { done, value } = await reader.read(); + if (done) break; + byteLength += value.byteLength; + if (byteLength > MAX_RESPONSE_BYTES) { + await reader.cancel(); + throw new WorkspaceToolHttpError('invalid'); + } + body += decoder.decode(value, { stream: true }); + } + body += decoder.decode(); + return JSON.parse(body) as unknown; + } catch (error) { + if (error instanceof WorkspaceToolHttpError) throw error; + if ( + signal?.aborted === true && + (error === signal.reason || + (isRecord(error) && (error.name === 'AbortError' || error.name === 'TimeoutError'))) + ) { + throw signal.reason ?? error; + } + if (isRecord(error) && (error.name === 'AbortError' || error.name === 'TimeoutError')) { + throw error; + } + throw new WorkspaceToolHttpError('invalid'); + } finally { + reader.releaseLock(); + } +} + +function isValidRequest(request: WorkspaceToolRequest): boolean { + if ( + request.protocolVersion !== 1 || + !/^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$/.test(request.workspaceId) + ) { + return false; + } + if (request.operation === 'read_file') { + return ( + isSafePath(request.path) && + (request.startLine == null || + isPositiveInteger(request.startLine, Number.MAX_SAFE_INTEGER)) && + (request.maxLines == null || isPositiveInteger(request.maxLines, MAX_READ_LINES)) + ); + } + if (request.operation !== 'search_text') { + return false; + } + return ( + typeof request.query === 'string' && + request.query.length > 0 && + request.query.length <= MAX_QUERY_LENGTH && + !request.query.includes('\0') && + (request.path == null || isSafePath(request.path)) && + (request.maxResults == null || isPositiveInteger(request.maxResults, MAX_SEARCH_RESULTS)) + ); +} + +function isValidResult( + request: WorkspaceToolRequest, + value: unknown, +): value is WorkspaceToolResult { + if ( + !isRecord(value) || + value.protocolVersion !== 1 || + value.operation !== request.operation || + value.workspaceId !== request.workspaceId || + typeof value.truncated !== 'boolean' + ) { + return false; + } + if (request.operation === 'read_file') { + const startLine = request.startLine ?? 1; + const maxLines = request.maxLines ?? 200; + const content = typeof value.content === 'string' ? value.content : null; + const reportedLineCount = + Number.isSafeInteger(value.endLine) && Number(value.endLine) >= startLine - 1 + ? Number(value.endLine) - startLine + 1 + : -1; + let actualLineCount = -1; + if (content != null) { + actualLineCount = content.length === 0 ? reportedLineCount : content.split('\n').length; + } + return ( + hasOnlyKeys(value, READ_RESULT_KEYS) && + value.path === request.path && + isSafePath(value.path) && + content != null && + new TextEncoder().encode(content).byteLength <= MAX_READ_BYTES && + value.startLine === startLine && + Number.isSafeInteger(value.endLine) && + Number(value.endLine) >= startLine - 1 && + Number(value.endLine) < startLine + maxLines && + reportedLineCount >= 0 && + reportedLineCount <= maxLines && + (content.length !== 0 || reportedLineCount <= 1) && + actualLineCount === reportedLineCount && + (value.truncated === true + ? Number.isSafeInteger(value.nextStartLine) && + Number(value.nextStartLine) === Number(value.endLine) + 1 + : value.nextStartLine == null) + ); + } + const maxResults = request.maxResults ?? 50; + return ( + hasOnlyKeys(value, SEARCH_RESULT_KEYS) && + Array.isArray(value.matches) && + value.matches.length <= maxResults && + value.matches.every( + (match) => + isRecord(match) && + hasOnlyKeys(match, SEARCH_MATCH_KEYS) && + isSafePath(match.path) && + isPositiveInteger(match.line, Number.MAX_SAFE_INTEGER) && + isPositiveInteger(match.column, Number.MAX_SAFE_INTEGER) && + typeof match.text === 'string' && + match.text.length <= MAX_SEARCH_TEXT_LENGTH, + ) + ); +} + +export async function executeWorkspaceTool({ + baseURL, + authHeaders, + request, + signal, + fetchImpl = fetch, +}: { + baseURL: string; + authHeaders: Record; + request: WorkspaceToolRequest; + signal?: AbortSignal; + fetchImpl?: CodeBridgeFetch; +}): Promise { + if (!isValidRequest(request)) { + throw new WorkspaceToolHttpError('invalid'); + } + try { + const timeoutSignal = AbortSignal.timeout(WORKSPACE_TOOL_TIMEOUT_MS); + const requestSignal = + signal != null && typeof AbortSignal.any === 'function' + ? AbortSignal.any([signal, timeoutSignal]) + : timeoutSignal; + const response = await fetchImpl( + `${baseURL.trim().replace(/\/+$/, '')}/workspace-tools/execute`, + { + method: 'POST', + headers: { + ...authHeaders, + 'Content-Type': 'application/json', + }, + body: JSON.stringify(request), + redirect: 'error', + signal: requestSignal, + }, + ); + if (!response.ok) { + await response.body?.cancel().catch(() => undefined); + throw new WorkspaceToolHttpError('rejected', response.status); + } + const result = await readBoundedJson(response, requestSignal); + if (!isValidResult(request, result)) { + throw new WorkspaceToolHttpError('invalid'); + } + return result; + } catch (error) { + if (error instanceof WorkspaceToolHttpError) throw error; + if ( + signal?.aborted === true && + (error === signal.reason || (isRecord(error) && error.name === 'AbortError')) + ) { + throw error; + } + if (error instanceof Error && error.name === 'TimeoutError') { + throw new WorkspaceToolHttpError('timeout'); + } + throw new WorkspaceToolHttpError('failed'); + } +}