diff --git a/api/server/services/Endpoints/agents/skillDeps.js b/api/server/services/Endpoints/agents/skillDeps.js index 29b5e70ed6..8fe49aedc3 100644 --- a/api/server/services/Endpoints/agents/skillDeps.js +++ b/api/server/services/Endpoints/agents/skillDeps.js @@ -7,6 +7,9 @@ const { readWorkspaceFile, searchWorkspace, listWorkspaceFiles, + writeWorkspaceFile, + previewWorkspaceEdit, + editWorkspaceFile, readSandboxFile, readSandboxImage, writeSandboxFile, @@ -386,6 +389,9 @@ const skillToolDeps = { readWorkspaceFile, searchWorkspace, listWorkspaceFiles, + writeWorkspaceFile, + previewWorkspaceEdit, + editWorkspaceFile, /** * `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 eac12f7033..bf848b7045 100644 --- a/api/server/services/Endpoints/agents/skillDeps.spec.js +++ b/api/server/services/Endpoints/agents/skillDeps.spec.js @@ -8,6 +8,9 @@ const mockCreateDeploymentSkillMethods = jest.fn((methods) => methods); const mockReadWorkspaceFile = jest.fn(); const mockSearchWorkspace = jest.fn(); const mockListWorkspaceFiles = jest.fn(); +const mockWriteWorkspaceFile = jest.fn(); +const mockPreviewWorkspaceEdit = jest.fn(); +const mockEditWorkspaceFile = jest.fn(); jest.mock('~/server/services/Files/strategies', () => ({ getStrategyFunctions: (...args) => mockGetStrategyFunctions(...args), @@ -24,6 +27,9 @@ jest.mock('~/server/services/Files/Code/process', () => ({ readWorkspaceFile: (...args) => mockReadWorkspaceFile(...args), searchWorkspace: (...args) => mockSearchWorkspace(...args), listWorkspaceFiles: (...args) => mockListWorkspaceFiles(...args), + writeWorkspaceFile: (...args) => mockWriteWorkspaceFile(...args), + previewWorkspaceEdit: (...args) => mockPreviewWorkspaceEdit(...args), + editWorkspaceFile: (...args) => mockEditWorkspaceFile(...args), writeSandboxFile: jest.fn(), })); @@ -103,6 +109,15 @@ describe('skillDeps saveSkillFileContent', () => { expect(mockListWorkspaceFiles).toHaveBeenCalledWith({ path: 'src' }); }); + it('exposes attached-workspace mutations to agent handlers', () => { + getSkillToolDeps().writeWorkspaceFile({ path: 'src/new.ts' }); + getSkillToolDeps().previewWorkspaceEdit({ path: 'src/app.ts' }); + getSkillToolDeps().editWorkspaceFile({ path: 'src/app.ts' }); + expect(mockWriteWorkspaceFile).toHaveBeenCalledWith({ path: 'src/new.ts' }); + expect(mockPreviewWorkspaceEdit).toHaveBeenCalledWith({ path: 'src/app.ts' }); + expect(mockEditWorkspaceFile).toHaveBeenCalledWith({ 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 bae25bbd41..2b5968c117 100644 --- a/api/server/services/Files/Code/process.js +++ b/api/server/services/Files/Code/process.js @@ -1798,6 +1798,97 @@ async function listWorkspaceFiles({ }); } +/** Writes a UTF-8 file in the workspace registered by an attached worker. */ +async function writeWorkspaceFile({ + file_path, + content, + overwrite, + workspace_id, + codeApiBaseUrl, + executionProfile, + bridgeWorkerId, + req, + signal, +}) { + const authHeaders = await getCodeApiAuthHeaders(req, bridgeWorkerId); + return executeWorkspaceTool({ + baseURL: codeApiBaseUrl, + authHeaders: { + ...authHeaders, + ...codeExecutionHeaders({ executionProfile, bridgeWorkerId }), + }, + request: { + protocolVersion: 1, + operation: 'write_file', + workspaceId: workspace_id, + path: file_path, + content, + overwrite, + }, + ...(signal ? { signal } : {}), + }); +} + +/** Applies an ordered exact-edit batch in one attached-worker mutation. */ +async function editWorkspaceFile({ + file_path, + edits, + expected_base_sha256, + workspace_id, + codeApiBaseUrl, + executionProfile, + bridgeWorkerId, + req, + signal, +}) { + const authHeaders = await getCodeApiAuthHeaders(req, bridgeWorkerId); + return executeWorkspaceTool({ + baseURL: codeApiBaseUrl, + authHeaders: { + ...authHeaders, + ...codeExecutionHeaders({ executionProfile, bridgeWorkerId }), + }, + request: { + protocolVersion: 1, + operation: 'edit_file', + workspaceId: workspace_id, + path: file_path, + edits, + ...(expected_base_sha256 ? { expectedBaseSha256: expected_base_sha256 } : {}), + }, + ...(signal ? { signal } : {}), + }); +} + +/** Previews an ordered exact-edit batch without mutating the attached workspace. */ +async function previewWorkspaceEdit({ + file_path, + edits, + workspace_id, + codeApiBaseUrl, + executionProfile, + bridgeWorkerId, + req, + signal, +}) { + const authHeaders = await getCodeApiAuthHeaders(req, bridgeWorkerId); + return executeWorkspaceTool({ + baseURL: codeApiBaseUrl, + authHeaders: { + ...authHeaders, + ...codeExecutionHeaders({ executionProfile, bridgeWorkerId }), + }, + request: { + protocolVersion: 1, + operation: 'preview_edit', + workspaceId: workspace_id, + path: file_path, + edits, + }, + ...(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 @@ -2067,6 +2158,9 @@ module.exports = { readWorkspaceFile, searchWorkspace, listWorkspaceFiles, + writeWorkspaceFile, + previewWorkspaceEdit, + editWorkspaceFile, readSandboxFile, readSandboxImage, writeSandboxFile, diff --git a/api/server/services/Files/Code/process.spec.js b/api/server/services/Files/Code/process.spec.js index 1f7deefa19..aef4233c37 100644 --- a/api/server/services/Files/Code/process.spec.js +++ b/api/server/services/Files/Code/process.spec.js @@ -251,6 +251,9 @@ const { readWorkspaceFile, searchWorkspace, listWorkspaceFiles, + writeWorkspaceFile, + previewWorkspaceEdit, + editWorkspaceFile, readSandboxImage, writeSandboxFile, primeFiles, @@ -1988,6 +1991,137 @@ describe('Code Process', () => { }); }); + describe('workspace mutations', () => { + it('forwards create-only writes to the selected attached worker', async () => { + const controller = new AbortController(); + const result = { + protocolVersion: 1, + operation: 'write_file', + workspaceId: 'primary', + path: 'src/new.ts', + created: true, + bytesWritten: 5, + }; + getCodeApiAuthHeaders.mockResolvedValueOnce({ Authorization: 'Bearer workspace-token' }); + mockExecuteWorkspaceTool.mockResolvedValueOnce(result); + + await expect( + writeWorkspaceFile({ + file_path: 'src/new.ts', + content: 'ready', + overwrite: false, + workspace_id: 'primary', + codeApiBaseUrl: 'https://attached-code.example.com/v1', + executionProfile: 'stateful', + bridgeWorkerId: 'worker-user-1', + req: mockReq, + signal: controller.signal, + }), + ).resolves.toBe(result); + + 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: 'write_file', + workspaceId: 'primary', + path: 'src/new.ts', + content: 'ready', + overwrite: false, + }, + signal: controller.signal, + }); + }); + + it('forwards an edit batch as one attached-worker request', async () => { + const result = { + protocolVersion: 1, + operation: 'edit_file', + workspaceId: 'primary', + path: 'src/app.ts', + replacements: 2, + bytesWritten: 20, + }; + getCodeApiAuthHeaders.mockResolvedValueOnce({ Authorization: 'Bearer workspace-token' }); + mockExecuteWorkspaceTool.mockResolvedValueOnce(result); + + await expect( + editWorkspaceFile({ + file_path: 'src/app.ts', + edits: [ + { oldText: 'draft', newText: 'ready' }, + { oldText: 'false', newText: 'true' }, + ], + workspace_id: 'primary', + codeApiBaseUrl: 'https://attached-code.example.com/v1', + executionProfile: 'stateful', + bridgeWorkerId: 'worker-user-1', + req: mockReq, + expected_base_sha256: 'a'.repeat(64), + }), + ).resolves.toBe(result); + + expect(mockExecuteWorkspaceTool).toHaveBeenCalledWith( + expect.objectContaining({ + request: { + protocolVersion: 1, + operation: 'edit_file', + workspaceId: 'primary', + path: 'src/app.ts', + edits: [ + { oldText: 'draft', newText: 'ready' }, + { oldText: 'false', newText: 'true' }, + ], + expectedBaseSha256: 'a'.repeat(64), + }, + }), + ); + }); + + it('forwards a non-mutating attached-workspace edit preview', async () => { + const result = { + protocolVersion: 1, + operation: 'preview_edit', + workspaceId: 'primary', + path: 'src/app.ts', + content: 'prefix SECRET', + hasUtf8Bom: false, + baseSha256: 'b'.repeat(64), + replacements: 1, + bytesWritten: 13, + }; + mockExecuteWorkspaceTool.mockResolvedValueOnce(result); + + await expect( + previewWorkspaceEdit({ + file_path: 'src/app.ts', + edits: [{ oldText: ' SEC', newText: ' SECRET' }], + workspace_id: 'primary', + codeApiBaseUrl: 'https://attached-code.example.com/v1', + executionProfile: 'stateful', + req: mockReq, + }), + ).resolves.toBe(result); + + expect(mockExecuteWorkspaceTool).toHaveBeenCalledWith( + expect.objectContaining({ + request: { + protocolVersion: 1, + operation: 'preview_edit', + workspaceId: 'primary', + path: 'src/app.ts', + edits: [{ oldText: ' SEC', newText: ' SECRET' }], + }, + }), + ); + }); + }); + 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 fc28cee788..447b0ac880 100644 --- a/packages/api/src/agents/handlers.spec.ts +++ b/packages/api/src/agents/handlers.spec.ts @@ -15,6 +15,7 @@ import type { CodeExecutionContext } from './execution'; import { createToolExecuteHandler, ToolExecuteOptions } from './handlers'; import { markSandboxReady } from './prewarm'; import { ContentFilterError } from '../middleware/contentFilter'; +import { WorkspaceToolHttpError } from '../code/workspace'; function createMockTool( name: string, @@ -4064,6 +4065,318 @@ describe('createToolExecuteHandler', () => { ]); }); + it('creates a file atomically in an attached workspace', async () => { + const writeWorkspaceFile = jest.fn(async () => ({ + protocolVersion: 1 as const, + operation: 'write_file' as const, + workspaceId: 'primary', + path: 'src/new.ts', + created: true, + bytesWritten: 20, + })); + const handler = makeSandboxAuthoringHandler( + { writeWorkspaceFile }, + { + codeExecutionContext: { + baseUrl: 'https://code.example.com', + codeSessionKey: 'attached-session', + executionProfile: 'stateful', + statefulSessions: true, + environmentType: 'attached', + bridgeWorkerId: 'user-worker', + }, + }, + ); + + const [result] = await invokeHandler(handler, [ + { + id: 'call_create_workspace', + name: 'create_file', + args: { + path: 'workspace/src/new.ts', + content: 'export const ok = 1;', + }, + }, + ]); + + expect(result).toMatchObject({ + status: 'success', + artifact: { + path: 'workspace/src/new.ts', + created: true, + bytes_written: 20, + }, + }); + expect(writeWorkspaceFile).toHaveBeenCalledWith({ + file_path: 'src/new.ts', + content: 'export const ok = 1;', + overwrite: false, + workspace_id: 'primary', + codeApiBaseUrl: 'https://code.example.com', + executionProfile: 'stateful', + bridgeWorkerId: 'user-worker', + req, + }); + }); + + it('surfaces an attached create-only conflict without retrying as an overwrite', async () => { + const writeWorkspaceFile = jest.fn(async () => { + throw new WorkspaceToolHttpError('rejected', 409); + }); + const handler = makeSandboxAuthoringHandler( + { writeWorkspaceFile }, + { + codeExecutionContext: { + baseUrl: 'https://code.example.com', + codeSessionKey: 'attached-session', + executionProfile: 'stateful', + statefulSessions: true, + environmentType: 'attached', + }, + }, + ); + + const [result] = await invokeHandler(handler, [ + { + id: 'call_create_workspace_conflict', + name: 'create_file', + args: { path: 'workspace/existing.txt', content: 'replacement' }, + }, + ]); + + expect(result.status).toBe('error'); + expect(result.errorMessage).toContain('overwrite: true'); + expect(writeWorkspaceFile).toHaveBeenCalledTimes(1); + }); + + it('rejects attached workspace writes above the worker protocol limit', async () => { + const writeWorkspaceFile = jest.fn(); + const handler = makeSandboxAuthoringHandler( + { writeWorkspaceFile }, + { + codeExecutionContext: { + baseUrl: 'https://code.example.com', + codeSessionKey: 'attached-session', + executionProfile: 'stateful', + statefulSessions: true, + environmentType: 'attached', + }, + }, + ); + + const [result] = await invokeHandler(handler, [ + { + id: 'call_create_workspace_oversized', + name: 'create_file', + args: { + path: 'workspace/large.txt', + content: 'x'.repeat(1024 * 1024 + 1), + }, + }, + ]); + + expect(result.status).toBe('error'); + expect(result.errorMessage).toContain('1 MiB'); + expect(writeWorkspaceFile).not.toHaveBeenCalled(); + }); + + it('sends attached edit batches as one atomic workspace mutation', async () => { + const editWorkspaceFile = jest.fn(async () => ({ + protocolVersion: 1 as const, + operation: 'edit_file' as const, + workspaceId: 'primary', + path: 'src/app.ts', + replacements: 2, + bytesWritten: 24, + })); + const handler = makeSandboxAuthoringHandler( + { editWorkspaceFile }, + { + codeExecutionContext: { + baseUrl: 'https://code.example.com', + codeSessionKey: 'attached-session', + executionProfile: 'stateful', + statefulSessions: true, + environmentType: 'attached', + bridgeWorkerId: 'user-worker', + }, + }, + ); + + const [result] = await invokeHandler(handler, [ + { + id: 'call_edit_workspace', + name: 'edit_file', + args: { + path: 'workspace/src/app.ts', + edits: [ + { old_text: 'draft', new_text: 'ready' }, + { old_text: 'false', new_text: 'true' }, + ], + }, + }, + ]); + + expect(result).toMatchObject({ + status: 'success', + artifact: { + path: 'workspace/src/app.ts', + edits: 2, + strategies: ['exact', 'exact'], + }, + }); + expect(editWorkspaceFile).toHaveBeenCalledWith({ + file_path: 'src/app.ts', + edits: [ + { oldText: 'draft', newText: 'ready' }, + { oldText: 'false', newText: 'true' }, + ], + workspace_id: 'primary', + codeApiBaseUrl: 'https://code.example.com', + executionProfile: 'stateful', + bridgeWorkerId: 'user-worker', + req, + }); + }); + + it('blocks protected attached edit content before worker dispatch', async () => { + const previewWorkspaceEdit = jest.fn(async () => ({ + protocolVersion: 1 as const, + operation: 'preview_edit' as const, + workspaceId: 'primary', + path: 'src/app.ts', + content: 'prefix ATTACHED-SECRET suffix', + hasUtf8Bom: false, + baseSha256: 'a'.repeat(64), + replacements: 1, + bytesWritten: 29, + })); + const editWorkspaceFile = jest.fn(); + const filteredReq = { + user: { id: 'user-1' }, + config: { + filters: { + files: { + pii: { + fields: ['content'], + starterPatterns: [], + customPatterns: [ + { + id: 'attached-secret', + label: 'attached secret', + regex: 'ATTACHED-SECRET', + }, + ], + }, + }, + }, + }, + } as never; + const handler = makeSandboxAuthoringHandler( + { previewWorkspaceEdit, editWorkspaceFile }, + { + req: filteredReq, + codeExecutionContext: { + baseUrl: 'https://code.example.com', + codeSessionKey: 'attached-session', + executionProfile: 'stateful', + statefulSessions: true, + environmentType: 'attached', + }, + }, + ); + + const [result] = await invokeHandler(handler, [ + { + id: 'call_edit_workspace_filtered', + name: 'edit_file', + args: { + path: 'workspace/src/app.ts', + old_text: ' suffix', + new_text: '-SECRET suffix', + }, + }, + ]); + + expect(result.status).toBe('error'); + expect(result.errorMessage).toContain('content_filter_block'); + expect(previewWorkspaceEdit).toHaveBeenCalledTimes(1); + expect(editWorkspaceFile).not.toHaveBeenCalled(); + }); + + it('commits inspected attached edits with the preview revision fence', async () => { + const previewWorkspaceEdit = jest.fn(async () => ({ + protocolVersion: 1 as const, + operation: 'preview_edit' as const, + workspaceId: 'primary', + path: 'src/app.ts', + content: 'const state = "ready";', + hasUtf8Bom: false, + baseSha256: 'b'.repeat(64), + replacements: 1, + bytesWritten: 22, + })); + const editWorkspaceFile = jest.fn(async () => ({ + protocolVersion: 1 as const, + operation: 'edit_file' as const, + workspaceId: 'primary', + path: 'src/app.ts', + replacements: 1, + bytesWritten: 22, + })); + const protectedReq = { + user: { id: 'user-1' }, + config: { + filters: { + files: { + pii: { + fields: ['content'], + starterPatterns: [], + customPatterns: [ + { + id: 'blocked-placeholder', + label: 'blocked placeholder', + regex: 'NEVER-MATCH-THIS', + }, + ], + }, + }, + }, + }, + } as never; + const handler = makeSandboxAuthoringHandler( + { previewWorkspaceEdit, editWorkspaceFile }, + { + req: protectedReq, + codeExecutionContext: { + baseUrl: 'https://code.example.com', + codeSessionKey: 'attached-session', + executionProfile: 'stateful', + statefulSessions: true, + environmentType: 'attached', + bridgeWorkerId: 'user-worker', + }, + }, + ); + + const [result] = await invokeHandler(handler, [ + { + id: 'call_edit_workspace_inspected', + name: 'edit_file', + args: { + path: 'workspace/src/app.ts', + old_text: 'draft', + new_text: 'ready', + }, + }, + ]); + + expect(result.status).toBe('success'); + expect(editWorkspaceFile).toHaveBeenCalledWith( + expect.objectContaining({ expected_base_sha256: 'b'.repeat(64) }), + ); + }); + it('contains a file-artifact policy rejection to its call without rejecting the batch', async () => { const detectorLabel = 'generated-file bearer token'; const detectorRule = 'generated-file-bearer'; diff --git a/packages/api/src/agents/handlers.ts b/packages/api/src/agents/handlers.ts index 11935afe39..8e6cb93d16 100644 --- a/packages/api/src/agents/handlers.ts +++ b/packages/api/src/agents/handlers.ts @@ -18,16 +18,19 @@ import type { import type { StructuredToolInterface } from '@librechat/agents/langchain/tools'; import type { CodeEnvRef, PtcToolCallEvent } from 'librechat-data-provider'; import type { ValidationIssue } from '@librechat/data-schemas'; +import type { + WorkspaceEditResult, + WorkspacePreviewEditResult, + WorkspaceListResult, + WorkspaceReadResult, + WorkspaceSearchResult, + WorkspaceWriteResult, +} from '~/code/workspace'; import type { BackgroundToolDeadClaimRecovery, BackgroundToolWakeupAdmission, BackgroundToolWakeupRegistration, } from './backgroundCompletion'; -import type { - WorkspaceListResult, - WorkspaceReadResult, - WorkspaceSearchResult, -} from '~/code/workspace'; import type { SkillFileRecord, PrimeSkillFilesResult } from './skillFiles'; import type { BackgroundToolResultState } from './harvest'; import type { CodeExecutionContext } from './execution'; @@ -87,6 +90,11 @@ import { runOutsideTracing, getSafeErrorMetadata, } from '~/utils'; +import { + WorkspaceToolHttpError, + WORKSPACE_EDIT_MAX_COUNT, + WORKSPACE_WRITE_MAX_BYTES, +} from '~/code/workspace'; import { hasIntentArg, stripIntentArg, @@ -485,6 +493,41 @@ export interface ToolExecuteOptions { req?: ServerRequest; signal?: AbortSignal; }) => Promise; + /** Writes a UTF-8 file within an attached worker's logical workspace. */ + writeWorkspaceFile?: (params: { + file_path: string; + content: string; + overwrite: boolean; + workspace_id: string; + codeApiBaseUrl: string; + executionProfile: CodeExecutionContext['executionProfile']; + bridgeWorkerId?: string; + req?: ServerRequest; + signal?: AbortSignal; + }) => Promise; + /** Previews exact replacements without mutating an attached worker workspace. */ + previewWorkspaceEdit?: (params: { + file_path: string; + edits: Array<{ oldText: string; newText: string }>; + workspace_id: string; + codeApiBaseUrl: string; + executionProfile: CodeExecutionContext['executionProfile']; + bridgeWorkerId?: string; + req?: ServerRequest; + signal?: AbortSignal; + }) => Promise; + /** Applies exact replacements atomically within an attached worker workspace. */ + editWorkspaceFile?: (params: { + file_path: string; + edits: Array<{ oldText: string; newText: string }>; + expected_base_sha256?: string; + workspace_id: string; + 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 @@ -3560,6 +3603,202 @@ async function writeBundledSkillFile({ }); } +function attachedWorkspaceAuthoringPath( + tc: ToolCallRequest, + filePath: string, +): { filePath: string } | ToolExecuteResult { + if (!filePath.startsWith('workspace/')) { + return errorResult(tc, 'Attached environment file paths must use "workspace/{relativePath}".'); + } + const relativePath = filePath.slice('workspace/'.length); + const pathError = invalidSandboxAuthoringPath(relativePath); + return pathError ? errorResult(tc, pathError) : { filePath: relativePath }; +} + +function attachedWorkspaceMutationParams( + codeExecutionContext: CodeExecutionContext, + req: ServerRequest | undefined, + signal: AbortSignal | undefined, +): { + workspace_id: string; + codeApiBaseUrl: string; + executionProfile: CodeExecutionContext['executionProfile']; + bridgeWorkerId?: string; + req?: ServerRequest; + signal?: AbortSignal; +} { + return { + workspace_id: 'primary', + codeApiBaseUrl: codeExecutionContext.baseUrl, + executionProfile: codeExecutionContext.executionProfile, + ...(codeExecutionContext.bridgeWorkerId + ? { bridgeWorkerId: codeExecutionContext.bridgeWorkerId } + : {}), + ...(req ? { req } : {}), + ...(signal ? { signal } : {}), + }; +} + +async function handleAttachedWorkspaceCreateFileCall({ + tc, + options, + req, + filePath, + content, + overwrite, + codeExecutionContext, + signal, +}: { + tc: ToolCallRequest; + options: ToolExecuteOptions; + req?: ServerRequest; + filePath: string; + content: string; + overwrite: boolean; + codeExecutionContext: CodeExecutionContext; + signal?: AbortSignal; +}): AuthoringResult { + if (!options.writeWorkspaceFile) { + return errorResult(tc, 'Attached workspace file writing is not configured.'); + } + const path = attachedWorkspaceAuthoringPath(tc, filePath); + if ('status' in path) return path; + if (new TextEncoder().encode(content).byteLength > WORKSPACE_WRITE_MAX_BYTES) { + return errorResult(tc, 'Attached workspace files are limited to 1 MiB per write.'); + } + const filtered = filteredFileResult(tc, req, path.filePath, content); + if (filtered != null) return filtered; + + try { + const result = await options.writeWorkspaceFile({ + file_path: path.filePath, + content, + overwrite, + ...attachedWorkspaceMutationParams(codeExecutionContext, req, signal), + }); + const action = result.created ? 'Created' : 'Updated'; + return successResult(tc, `${action} workspace/${path.filePath} (${content.length} chars).`, { + path: `workspace/${path.filePath}`, + [HOST_FILE_AUTHORING_ARTIFACT_KEY]: true, + bytes_written: result.bytesWritten, + created: result.created, + }); + } catch (error) { + if (error instanceof WorkspaceToolHttpError && error.upstreamStatus === 409 && !overwrite) { + return errorResult(tc, 'File already exists. Pass overwrite: true to replace.'); + } + if (signal?.aborted === true && isAbortError(error)) throw error; + logger.warn('[file_authoring] Attached workspace write failed', getSafeErrorMetadata(error)); + return errorResult(tc, `Failed to write "workspace/${path.filePath}".`); + } +} + +async function handleAttachedWorkspaceEditFileCall({ + tc, + options, + req, + filePath, + edits, + codeExecutionContext, + signal, +}: { + tc: ToolCallRequest; + options: ToolExecuteOptions; + req?: ServerRequest; + filePath: string; + edits: TextEdit[]; + codeExecutionContext: CodeExecutionContext; + signal?: AbortSignal; +}): AuthoringResult { + if (!options.editWorkspaceFile) { + return errorResult(tc, 'Attached workspace file editing is not configured.'); + } + const path = attachedWorkspaceAuthoringPath(tc, filePath); + if ('status' in path) return path; + if (edits.length > WORKSPACE_EDIT_MAX_COUNT) { + return errorResult( + tc, + `Attached workspace edits are limited to ${WORKSPACE_EDIT_MAX_COUNT} replacements per call.`, + ); + } + const editBytes = edits.reduce( + (bytes, edit) => + bytes + + new TextEncoder().encode(edit.old_text).byteLength + + new TextEncoder().encode(edit.new_text).byteLength, + 0, + ); + if (editBytes > WORKSPACE_WRITE_MAX_BYTES) { + return errorResult(tc, 'Attached workspace edit text is limited to 1 MiB per call.'); + } + const filteredName = filteredFileNameResult(tc, req, path.filePath); + if (filteredName != null) return filteredName; + + try { + const workspaceEdits = edits.map((edit) => ({ + oldText: edit.old_text, + newText: edit.new_text, + })); + let expectedBaseSha256: string | undefined; + if (hasActiveFileFieldPolicy(req?.config?.filters, ['content', 'extracted_text'])) { + if (!options.previewWorkspaceEdit) { + return errorResult( + tc, + 'Attached workspace editing requires an updated BYOM worker while file-content protections are enabled.', + ); + } + let preview: WorkspacePreviewEditResult; + try { + preview = await options.previewWorkspaceEdit({ + file_path: path.filePath, + edits: workspaceEdits, + ...attachedWorkspaceMutationParams(codeExecutionContext, req, signal), + }); + } catch (error) { + if (signal?.aborted === true && isAbortError(error)) throw error; + if (error instanceof WorkspaceToolHttpError && error.upstreamStatus === 400) { + return errorResult( + tc, + 'This attached environment must update its LibreChat Code worker before protected files can be edited.', + ); + } + throw error; + } + const filteredContent = filteredFileResult(tc, req, path.filePath, preview.content); + if (filteredContent != null) return filteredContent; + expectedBaseSha256 = preview.baseSha256; + } + const result = await options.editWorkspaceFile({ + file_path: path.filePath, + edits: workspaceEdits, + ...(expectedBaseSha256 ? { expected_base_sha256: expectedBaseSha256 } : {}), + ...attachedWorkspaceMutationParams(codeExecutionContext, req, signal), + }); + return successResult( + tc, + `Updated workspace/${path.filePath} with ${result.replacements} exact replacement${result.replacements === 1 ? '' : 's'}.`, + { + path: `workspace/${path.filePath}`, + [HOST_FILE_AUTHORING_ARTIFACT_KEY]: true, + bytes_written: result.bytesWritten, + created: false, + edits: result.replacements, + strategies: Array.from({ length: result.replacements }, () => 'exact'), + }, + ); + } catch (error) { + if (error instanceof WorkspaceToolHttpError && error.upstreamStatus === 409) { + return errorResult( + tc, + `The requested text did not match exactly once in "workspace/${path.filePath}". Re-read the file and retry.`, + ); + } + if (signal?.aborted === true && isAbortError(error)) throw error; + logger.warn('[file_authoring] Attached workspace edit failed', getSafeErrorMetadata(error)); + return errorResult(tc, `Failed to edit "workspace/${path.filePath}".`); + } +} + async function handleSandboxCreateFileCall({ tc, options, @@ -3569,6 +3808,7 @@ async function handleSandboxCreateFileCall({ overwrite, sandboxContext, codeExecutionContext, + signal, }: { tc: ToolCallRequest; options: ToolExecuteOptions; @@ -3578,7 +3818,20 @@ async function handleSandboxCreateFileCall({ overwrite: boolean; sandboxContext?: SandboxSessionContext; codeExecutionContext?: CodeExecutionContext; + signal?: AbortSignal; }): AuthoringResult { + if (codeExecutionContext?.environmentType === 'attached') { + return await handleAttachedWorkspaceCreateFileCall({ + tc, + options, + req, + filePath, + content, + overwrite, + codeExecutionContext, + signal, + }); + } const pathError = invalidSandboxAuthoringPath(filePath); if (pathError) { return errorResult(tc, pathError); @@ -3620,6 +3873,7 @@ async function handleSandboxEditFileCall({ edits, sandboxContext, codeExecutionContext, + signal, }: { tc: ToolCallRequest; options: ToolExecuteOptions; @@ -3628,7 +3882,19 @@ async function handleSandboxEditFileCall({ edits: TextEdit[]; sandboxContext?: SandboxSessionContext; codeExecutionContext?: CodeExecutionContext; + signal?: AbortSignal; }): AuthoringResult { + if (codeExecutionContext?.environmentType === 'attached') { + return await handleAttachedWorkspaceEditFileCall({ + tc, + options, + req, + filePath, + edits, + codeExecutionContext, + signal, + }); + } const pathError = invalidSandboxAuthoringPath(filePath); if (pathError) { return errorResult(tc, pathError); @@ -3688,6 +3954,7 @@ async function handleCreateFileCall( req?: ServerRequest, sourceConfigurable?: Record, sandboxContext?: SandboxSessionContext, + signal?: AbortSignal, ): AuthoringResult { const args = tc.args as { path?: unknown; content?: unknown; overwrite?: unknown }; if (typeof args.path !== 'string' || args.path.length === 0) { @@ -3722,6 +3989,7 @@ async function handleCreateFileCall( overwrite, sandboxContext, codeExecutionContext: getCodeExecutionContext(mergedConfigurable), + signal, }); } @@ -3801,6 +4069,7 @@ async function handleEditFileCall( options: ToolExecuteOptions, req?: ServerRequest, sandboxContext?: SandboxSessionContext, + signal?: AbortSignal, ): AuthoringResult { const args = tc.args as { path?: unknown; @@ -3832,6 +4101,7 @@ async function handleEditFileCall( edits, sandboxContext, codeExecutionContext: getCodeExecutionContext(mergedConfigurable), + signal, }); } @@ -6036,6 +6306,7 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand req, sourceConfigurable, sandboxContext, + runSignal, ); } else if (tc.name === EDIT_FILE_TOOL_NAME && isFileAuthoringCall) { handlerResult = await handleEditFileCall( @@ -6044,6 +6315,7 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand options, req, sandboxContext, + runSignal, ); } else { handlerResult = errorResult(tc, `Tool ${tc.name} not found`); diff --git a/packages/api/src/agents/initialize.ts b/packages/api/src/agents/initialize.ts index dc43ff1ccc..8d58ab35e9 100644 --- a/packages/api/src/agents/initialize.ts +++ b/packages/api/src/agents/initialize.ts @@ -1537,6 +1537,7 @@ export async function initializeAgent( toolRegistry, toolDefinitions, includeSkillFileInstructions: skillAuthoringAvailable, + workspaceTools: attachedWorkspaceTools, }); toolDefinitions = fileAuthoringResult.toolDefinitions; /** File authoring is owned by whichever capability switched it on — diff --git a/packages/api/src/agents/tools.spec.ts b/packages/api/src/agents/tools.spec.ts index 7b24a7a1cf..0ebdabfc17 100644 --- a/packages/api/src/agents/tools.spec.ts +++ b/packages/api/src/agents/tools.spec.ts @@ -870,6 +870,25 @@ describe('registerFileAuthoringTools', () => { expect(filePathDescription(editFile)).not.toContain('rename skills'); }); + it('registers attached-workspace paths and atomic edit semantics', () => { + const result = registerFileAuthoringTools({ + toolRegistry: makeRegistry(), + toolDefinitions: [], + includeSkillFileInstructions: false, + workspaceTools: true, + }); + const createFile = result.toolDefinitions.find((d) => d.name === 'create_file'); + const editFile = result.toolDefinitions.find((d) => d.name === 'edit_file'); + + expect(createFile?.description).toContain('workspace/{relativePath}'); + expect(createFile?.description).not.toContain('/mnt/data/'); + expect(editFile?.description).toContain('entire batch commits atomically'); + expect(filePathDescription(createFile)).toContain('workspace/{relativePath}'); + expect(filePathDescription(editFile)).toContain('workspace/{relativePath}'); + expect(isFileAuthoringToolDefinition(createFile)).toBe(true); + expect(isFileAuthoringToolDefinition(editFile)).toBe(true); + }); + it('is idempotent across repeated registration calls', () => { const toolRegistry = makeRegistry(); const first = registerFileAuthoringTools({ @@ -897,7 +916,6 @@ describe('registerFileAuthoringTools', () => { toolDefinitions: codeOnly.toolDefinitions, includeSkillFileInstructions: true, }); - expect(upgraded.registered).toEqual([]); expect(upgraded.toolDefinitions.find((d) => d.name === 'create_file')?.description).toContain( 'skills/', @@ -922,12 +940,19 @@ describe('registerFileAuthoringTools', () => { toolDefinitions: codeOnly.toolDefinitions, includeSkillFileInstructions: true, }); + const attached = registerFileAuthoringTools({ + toolRegistry: makeRegistry(), + toolDefinitions: [], + includeSkillFileInstructions: true, + workspaceTools: true, + }); expect( maxToolDescriptionLength([ ...skillAware.toolDefinitions, ...codeOnly.toolDefinitions, ...upgraded.toolDefinitions, + ...attached.toolDefinitions, ]), ).toBeLessThanOrEqual(TOOL_DESCRIPTION_ADVISORY_MAX_LENGTH); }); diff --git a/packages/api/src/agents/tools.ts b/packages/api/src/agents/tools.ts index 3a7e02bd9e..75b8bc0047 100644 --- a/packages/api/src/agents/tools.ts +++ b/packages/api/src/agents/tools.ts @@ -400,6 +400,8 @@ export interface RegisterFileAuthoringToolsParams { * descriptions stay focused on code-execution sandbox files. */ includeSkillFileInstructions?: boolean; + /** When true, non-skill paths use the attached worker's workspace/ namespace. */ + workspaceTools?: boolean; } /** @@ -709,6 +711,70 @@ const CODE_EDIT_FILE_DEF: LCTool = Object.freeze({ responseFormat: 'content_and_artifact' as LCTool['responseFormat'], }) as LCTool; +const ATTACHED_CODE_CREATE_FILE_DESCRIPTION = `Create a new file in the selected attached environment, or overwrite one with explicit intent. + +Use a path in the form "workspace/{relativePath}". Requires overwrite: true to replace an existing file and refuses otherwise. The workspace may be an existing project, a Git repository, or an empty directory. + +Very long content can exceed the streamed tool-argument limit (64 KB by default). The attached workspace also limits each write to 1 MiB. Keep each call bounded.`; + +const ATTACHED_CODE_EDIT_FILE_DESCRIPTION = `Apply one or more ordered exact text replacements to an existing file in the selected attached environment. + +Use a path in the form "workspace/{relativePath}". Every old_text must match exactly one location at its step in the batch. Up to 100 replacements and 1 MiB of edit text are allowed; the entire batch commits atomically or makes no change.`; + +const ATTACHED_SKILL_CREATE_FILE_DESCRIPTION = `${SKILL_CREATE_FILE_DESCRIPTION.replace( + 'Non-skills paths target the code-execution sandbox when enabled. Prefer /mnt/data/{file}.', + 'For the selected attached environment, non-skill paths must use "workspace/{relativePath}".', +)}`; + +const ATTACHED_SKILL_EDIT_FILE_DESCRIPTION = `Apply targeted text replacements to an existing file. + +For skills/{skillName}/... paths, exact matching falls back to whitespace-tolerant matching when needed and the result includes a unified diff. Keep SKILL.md YAML frontmatter name equal to {skillName}; create a new skills/{newName}/SKILL.md to rename a skill. + +For workspace/{relativePath} paths in the selected attached environment, every old_text must match exactly one location at its step. There is no whitespace-tolerant fallback. Up to 100 replacements and 1 MiB of edit text commit atomically, and the result is a write summary rather than a unified diff.`; + +function attachedFileAuthoringParameters( + parameters: LCTool['parameters'], + includeSkillFileInstructions: boolean, +): LCTool['parameters'] { + const properties = parameters?.properties ?? {}; + return Object.freeze({ + ...parameters, + properties: { + ...properties, + path: { + type: 'string', + description: includeSkillFileInstructions + ? 'Use "skills/{skillName}/..." for a LibreChat skill file, or "workspace/{relativePath}" for a file in the selected attached environment.' + : 'Use "workspace/{relativePath}" for a file in the selected attached environment.', + }, + }, + }) as LCTool['parameters']; +} + +const ATTACHED_CODE_CREATE_FILE_DEF: LCTool = Object.freeze({ + ...CODE_CREATE_FILE_DEF, + description: ATTACHED_CODE_CREATE_FILE_DESCRIPTION, + parameters: attachedFileAuthoringParameters(CODE_CREATE_FILE_PARAMETERS, false), +}) as LCTool; + +const ATTACHED_CODE_EDIT_FILE_DEF: LCTool = Object.freeze({ + ...CODE_EDIT_FILE_DEF, + description: ATTACHED_CODE_EDIT_FILE_DESCRIPTION, + parameters: attachedFileAuthoringParameters(CODE_EDIT_FILE_PARAMETERS, false), +}) as LCTool; + +const ATTACHED_SKILL_CREATE_FILE_DEF: LCTool = Object.freeze({ + ...SKILL_CREATE_FILE_DEF, + description: ATTACHED_SKILL_CREATE_FILE_DESCRIPTION, + parameters: attachedFileAuthoringParameters(SKILL_CREATE_FILE_PARAMETERS, true), +}) as LCTool; + +const ATTACHED_SKILL_EDIT_FILE_DEF: LCTool = Object.freeze({ + ...SKILL_EDIT_FILE_DEF, + description: ATTACHED_SKILL_EDIT_FILE_DESCRIPTION, + parameters: attachedFileAuthoringParameters(SKILL_EDIT_FILE_PARAMETERS, true), +}) as LCTool; + function buildReadFileDef(includeSkillFileInstructions: boolean, workspaceTools: boolean): LCTool { if (workspaceTools) { return includeSkillFileInstructions @@ -718,7 +784,15 @@ function buildReadFileDef(includeSkillFileInstructions: boolean, workspaceTools: return includeSkillFileInstructions ? READ_FILE_DEF : CODE_READ_FILE_DEF; } -function buildFileAuthoringDefs(includeSkillFileInstructions: boolean): LCTool[] { +function buildFileAuthoringDefs( + includeSkillFileInstructions: boolean, + workspaceTools: boolean, +): LCTool[] { + if (workspaceTools) { + return includeSkillFileInstructions + ? [ATTACHED_SKILL_CREATE_FILE_DEF, ATTACHED_SKILL_EDIT_FILE_DEF] + : [ATTACHED_CODE_CREATE_FILE_DEF, ATTACHED_CODE_EDIT_FILE_DEF]; + } return includeSkillFileInstructions ? [SKILL_CREATE_FILE_DEF, SKILL_EDIT_FILE_DEF] : [CODE_CREATE_FILE_DEF, CODE_EDIT_FILE_DEF]; @@ -734,10 +808,16 @@ function isCodeOnlyReadFileDef(def: LCTool | undefined): boolean { function isCodeOnlyFileAuthoringDef(def: LCTool | undefined): boolean { if (def?.name === CREATE_FILE_TOOL_NAME) { - return def.description === CODE_CREATE_FILE_DESCRIPTION; + return ( + def.description === CODE_CREATE_FILE_DESCRIPTION || + def.description === ATTACHED_CODE_CREATE_FILE_DESCRIPTION + ); } if (def?.name === EDIT_FILE_TOOL_NAME) { - return def.description === CODE_EDIT_FILE_DESCRIPTION; + return ( + def.description === CODE_EDIT_FILE_DESCRIPTION || + def.description === ATTACHED_CODE_EDIT_FILE_DESCRIPTION + ); } return false; } @@ -746,13 +826,17 @@ export function isFileAuthoringToolDefinition(def: LCTool | undefined): boolean if (def?.name === CREATE_FILE_TOOL_NAME) { return ( def.description === CODE_CREATE_FILE_DESCRIPTION || - def.description === SKILL_CREATE_FILE_DESCRIPTION + def.description === SKILL_CREATE_FILE_DESCRIPTION || + def.description === ATTACHED_CODE_CREATE_FILE_DESCRIPTION || + def.description === ATTACHED_SKILL_CREATE_FILE_DESCRIPTION ); } if (def?.name === EDIT_FILE_TOOL_NAME) { return ( def.description === CODE_EDIT_FILE_DESCRIPTION || - def.description === SKILL_EDIT_FILE_DESCRIPTION + def.description === SKILL_EDIT_FILE_DESCRIPTION || + def.description === ATTACHED_CODE_EDIT_FILE_DESCRIPTION || + def.description === ATTACHED_SKILL_EDIT_FILE_DESCRIPTION ); } return false; @@ -903,9 +987,14 @@ export function registerCodeExecutionTools( export function registerFileAuthoringTools( params: RegisterFileAuthoringToolsParams, ): RegisterFileAuthoringToolsResult { - const { toolRegistry, toolDefinitions, includeSkillFileInstructions = true } = params; + const { + toolRegistry, + toolDefinitions, + includeSkillFileInstructions = true, + workspaceTools = false, + } = params; - const candidates = buildFileAuthoringDefs(includeSkillFileInstructions); + const candidates = buildFileAuthoringDefs(includeSkillFileInstructions, workspaceTools); const toolNames = candidates.map((def) => def.name); const inputDefinitions = toolDefinitions ?? []; let workingDefinitions = inputDefinitions; diff --git a/packages/api/src/code/workspace.spec.ts b/packages/api/src/code/workspace.spec.ts index 422ed54633..63f1b6ab98 100644 --- a/packages/api/src/code/workspace.spec.ts +++ b/packages/api/src/code/workspace.spec.ts @@ -562,7 +562,7 @@ describe('executeWorkspaceTool', () => { }); test('accepts a bounded command result from the selected attached worker', async () => { - const fetchImpl: CodeBridgeFetch = jest.fn(async () => + const fetchImpl = jest.fn(async () => Response.json({ protocolVersion: 1, operation: 'execute_command', @@ -634,4 +634,187 @@ describe('executeWorkspaceTool', () => { }), ).rejects.toMatchObject({ reason: 'invalid' }); }); + + test('validates atomic create-only workspace writes', async () => { + const fetchImpl = jest.fn(async () => + Response.json({ + protocolVersion: 1, + operation: 'write_file', + workspaceId: 'primary', + path: 'src/new.ts', + created: true, + bytesWritten: 5, + }), + ); + + await expect( + executeWorkspaceTool({ + baseURL: 'https://code.example.com/v1', + authHeaders: { Authorization: 'Bearer jwt' }, + request: { + protocolVersion: 1, + operation: 'write_file', + workspaceId: 'primary', + path: 'src/new.ts', + content: 'ready', + overwrite: false, + }, + fetchImpl, + }), + ).resolves.toMatchObject({ created: true, bytesWritten: 5 }); + + fetchImpl.mockResolvedValueOnce( + Response.json({ + protocolVersion: 1, + operation: 'write_file', + workspaceId: 'primary', + path: 'src/new.ts', + created: false, + bytesWritten: 5, + }), + ); + await expect( + executeWorkspaceTool({ + baseURL: 'https://code.example.com/v1', + authHeaders: {}, + request: { + protocolVersion: 1, + operation: 'write_file', + workspaceId: 'primary', + path: 'src/new.ts', + content: 'ready', + overwrite: false, + }, + fetchImpl, + }), + ).rejects.toMatchObject({ reason: 'invalid' }); + }); + + test('validates bounded atomic workspace edit batches', async () => { + const fetchImpl: CodeBridgeFetch = jest.fn(async () => + Response.json({ + protocolVersion: 1, + operation: 'edit_file', + workspaceId: 'primary', + path: 'src/app.ts', + replacements: 2, + bytesWritten: 18, + }), + ); + const request = { + protocolVersion: 1 as const, + operation: 'edit_file' as const, + workspaceId: 'primary', + path: 'src/app.ts', + edits: [ + { oldText: 'false', newText: 'true' }, + { oldText: 'draft', newText: 'ready' }, + ], + }; + + await expect( + executeWorkspaceTool({ + baseURL: 'https://code.example.com/v1', + authHeaders: {}, + request, + fetchImpl, + }), + ).resolves.toMatchObject({ replacements: 2 }); + + await expect( + executeWorkspaceTool({ + baseURL: 'https://code.example.com/v1', + authHeaders: {}, + request: { ...request, edits: [] }, + fetchImpl, + }), + ).rejects.toMatchObject({ reason: 'invalid' }); + await expect( + executeWorkspaceTool({ + baseURL: 'https://code.example.com/v1', + authHeaders: {}, + request: { ...request, edits: undefined } as unknown as typeof request, + fetchImpl, + }), + ).rejects.toMatchObject({ reason: 'invalid' }); + expect(fetchImpl).toHaveBeenCalledTimes(1); + }); + + test('validates exact edit previews and revision-fenced commits', async () => { + const edits = [{ oldText: ' suffix', newText: 'RET suffix' }]; + const baseSha256 = 'a'.repeat(64); + const fetchImpl: CodeBridgeFetch = jest + .fn() + .mockResolvedValueOnce( + Response.json({ + protocolVersion: 1, + operation: 'preview_edit', + workspaceId: 'primary', + path: 'src/app.ts', + content: 'prefix SECRET suffix', + hasUtf8Bom: false, + baseSha256, + replacements: 1, + bytesWritten: 20, + }), + ) + .mockResolvedValueOnce( + Response.json({ + protocolVersion: 1, + operation: 'edit_file', + workspaceId: 'primary', + path: 'src/app.ts', + replacements: 1, + bytesWritten: 20, + }), + ); + + await expect( + executeWorkspaceTool({ + baseURL: 'https://code.example.com/v1', + authHeaders: {}, + request: { + protocolVersion: 1, + operation: 'preview_edit', + workspaceId: 'primary', + path: 'src/app.ts', + edits, + }, + fetchImpl, + }), + ).resolves.toMatchObject({ content: 'prefix SECRET suffix', baseSha256 }); + + await expect( + executeWorkspaceTool({ + baseURL: 'https://code.example.com/v1', + authHeaders: {}, + request: { + protocolVersion: 1, + operation: 'edit_file', + workspaceId: 'primary', + path: 'src/app.ts', + edits, + expectedBaseSha256: baseSha256, + }, + fetchImpl, + }), + ).resolves.toMatchObject({ replacements: 1 }); + + await expect( + executeWorkspaceTool({ + baseURL: 'https://code.example.com/v1', + authHeaders: {}, + request: { + protocolVersion: 1, + operation: 'edit_file', + workspaceId: 'primary', + path: 'src/app.ts', + edits, + expectedBaseSha256: 'invalid', + }, + fetchImpl, + }), + ).rejects.toMatchObject({ reason: 'invalid' }); + expect(fetchImpl).toHaveBeenCalledTimes(2); + }); }); diff --git a/packages/api/src/code/workspace.ts b/packages/api/src/code/workspace.ts index e16aa3b2d6..84d275bb7e 100644 --- a/packages/api/src/code/workspace.ts +++ b/packages/api/src/code/workspace.ts @@ -8,6 +8,8 @@ const MAX_READ_LINES = 500; const MAX_SEARCH_RESULTS = 200; const MAX_SEARCH_TEXT_LENGTH = 2000; const MAX_LIST_RESULTS = 500; +export const WORKSPACE_WRITE_MAX_BYTES: number = 1024 * 1024; +export const WORKSPACE_EDIT_MAX_COUNT: number = 100; const MAX_COMMAND_BYTES = 32 * 1024; const DEFAULT_COMMAND_TIMEOUT_MS = 30_000; const MAX_COMMAND_TIMEOUT_MS = 5 * 60_000; @@ -54,6 +56,34 @@ const COMMAND_RESULT_KEYS = new Set([ 'truncated', 'timedOut', ]); +const WRITE_RESULT_KEYS = new Set([ + 'protocolVersion', + 'operation', + 'workspaceId', + 'path', + 'created', + 'bytesWritten', +]); +const EDIT_RESULT_KEYS = new Set([ + 'protocolVersion', + 'operation', + 'workspaceId', + 'path', + 'replacements', + 'bytesWritten', +]); +const PREVIEW_EDIT_RESULT_KEYS = new Set([ + 'protocolVersion', + 'operation', + 'workspaceId', + 'path', + 'content', + 'hasUtf8Bom', + 'baseSha256', + 'replacements', + 'bytesWritten', +]); +const TEXT_EDIT_KEYS = new Set(['oldText', 'newText']); export interface WorkspaceReadRequest { protocolVersion: 1; @@ -92,10 +122,44 @@ export interface WorkspaceExecuteCommandRequest { maxOutputBytes?: number; } +export interface WorkspaceWriteRequest { + protocolVersion: 1; + operation: 'write_file'; + workspaceId: string; + path: string; + content: string; + overwrite?: boolean; +} + +export interface WorkspaceTextEdit { + oldText: string; + newText: string; +} + +export interface WorkspaceEditRequest { + protocolVersion: 1; + operation: 'edit_file'; + workspaceId: string; + path: string; + edits: WorkspaceTextEdit[]; + expectedBaseSha256?: string; +} + +export interface WorkspacePreviewEditRequest { + protocolVersion: 1; + operation: 'preview_edit'; + workspaceId: string; + path: string; + edits: WorkspaceTextEdit[]; +} + export type WorkspaceToolRequest = | WorkspaceReadRequest | WorkspaceSearchRequest | WorkspaceListRequest + | WorkspaceWriteRequest + | WorkspacePreviewEditRequest + | WorkspaceEditRequest | WorkspaceExecuteCommandRequest; export interface WorkspaceReadResult { @@ -139,10 +203,43 @@ export interface WorkspaceExecuteCommandResult { timedOut: boolean; } +export interface WorkspaceWriteResult { + protocolVersion: 1; + operation: 'write_file'; + workspaceId: string; + path: string; + created: boolean; + bytesWritten: number; +} + +export interface WorkspaceEditResult { + protocolVersion: 1; + operation: 'edit_file'; + workspaceId: string; + path: string; + replacements: number; + bytesWritten: number; +} + +export interface WorkspacePreviewEditResult { + protocolVersion: 1; + operation: 'preview_edit'; + workspaceId: string; + path: string; + content: string; + hasUtf8Bom: boolean; + baseSha256: string; + replacements: number; + bytesWritten: number; +} + export type WorkspaceToolResult = | WorkspaceReadResult | WorkspaceSearchResult | WorkspaceListResult + | WorkspaceWriteResult + | WorkspacePreviewEditResult + | WorkspaceEditResult | WorkspaceExecuteCommandResult; export class WorkspaceToolHttpError extends Error { @@ -186,6 +283,28 @@ function isUtf8StringWithinBytes(value: unknown, maximum: number): value is stri ); } +function areValidWorkspaceEdits(edits: unknown): edits is WorkspaceTextEdit[] { + if (!Array.isArray(edits)) return false; + if (edits.length < 1 || edits.length > WORKSPACE_EDIT_MAX_COUNT) return false; + let bytes = 0; + for (const edit of edits) { + if ( + !isRecord(edit) || + !hasOnlyKeys(edit, TEXT_EDIT_KEYS) || + !isUtf8StringWithinBytes(edit.oldText, WORKSPACE_WRITE_MAX_BYTES) || + edit.oldText.length === 0 || + !isUtf8StringWithinBytes(edit.newText, WORKSPACE_WRITE_MAX_BYTES) + ) { + return false; + } + bytes += + new TextEncoder().encode(edit.oldText).byteLength + + new TextEncoder().encode(edit.newText).byteLength; + if (bytes > WORKSPACE_WRITE_MAX_BYTES) return false; + } + return true; +} + function hasOnlyKeys(value: Record, allowed: ReadonlySet): boolean { return Object.keys(value).every((key) => allowed.has(key)); } @@ -296,6 +415,23 @@ function isValidRequest(request: WorkspaceToolRequest): boolean { isPositiveInteger(request.maxOutputBytes, MAX_COMMAND_OUTPUT_BYTES)) ); } + if (request.operation === 'write_file') { + return ( + isSafePath(request.path) && + isUtf8StringWithinBytes(request.content, WORKSPACE_WRITE_MAX_BYTES) && + (request.overwrite === undefined || typeof request.overwrite === 'boolean') + ); + } + if (request.operation === 'preview_edit') { + return isSafePath(request.path) && areValidWorkspaceEdits(request.edits); + } + if (request.operation === 'edit_file') { + return ( + isSafePath(request.path) && + areValidWorkspaceEdits(request.edits) && + (request.expectedBaseSha256 == null || /^[a-f0-9]{64}$/.test(request.expectedBaseSha256)) + ); + } if (request.operation !== 'search_text') { return false; } @@ -317,8 +453,7 @@ function isValidResult( !isRecord(value) || value.protocolVersion !== 1 || value.operation !== request.operation || - value.workspaceId !== request.workspaceId || - typeof value.truncated !== 'boolean' + value.workspaceId !== request.workspaceId ) { return false; } @@ -339,6 +474,7 @@ function isValidResult( value.path === request.path && isSafePath(value.path) && content != null && + typeof value.truncated === 'boolean' && new TextEncoder().encode(content).byteLength <= MAX_READ_BYTES && value.startLine === startLine && Number.isSafeInteger(value.endLine) && @@ -358,6 +494,7 @@ function isValidResult( const maxResults = request.maxResults ?? 100; if ( !hasOnlyKeys(value, LIST_RESULT_KEYS) || + typeof value.truncated !== 'boolean' || !Array.isArray(value.paths) || value.paths.length > maxResults ) { @@ -384,6 +521,7 @@ function isValidResult( const outputLimit = request.maxOutputBytes ?? DEFAULT_COMMAND_OUTPUT_BYTES; return ( hasOnlyKeys(value, COMMAND_RESULT_KEYS) && + typeof value.truncated === 'boolean' && stdout != null && stderr != null && Buffer.from(stdout).toString('utf8') === stdout && @@ -404,9 +542,46 @@ function isValidResult( : value.timedOut === false && value.signal == null) ); } + if (request.operation === 'write_file') { + return ( + hasOnlyKeys(value, WRITE_RESULT_KEYS) && + value.path === request.path && + typeof value.created === 'boolean' && + (request.overwrite !== false || value.created === true) && + Number.isSafeInteger(value.bytesWritten) && + Number(value.bytesWritten) === new TextEncoder().encode(request.content).byteLength + ); + } + if (request.operation === 'edit_file') { + return ( + hasOnlyKeys(value, EDIT_RESULT_KEYS) && + value.path === request.path && + value.replacements === request.edits.length && + Number.isSafeInteger(value.bytesWritten) && + Number(value.bytesWritten) >= 0 && + Number(value.bytesWritten) <= WORKSPACE_WRITE_MAX_BYTES + ); + } + if (request.operation === 'preview_edit') { + const content = typeof value.content === 'string' ? value.content : null; + return ( + hasOnlyKeys(value, PREVIEW_EDIT_RESULT_KEYS) && + value.path === request.path && + content != null && + Buffer.from(content).toString('utf8') === content && + typeof value.hasUtf8Bom === 'boolean' && + /^[a-f0-9]{64}$/.test(typeof value.baseSha256 === 'string' ? value.baseSha256 : '') && + value.replacements === request.edits.length && + Number.isSafeInteger(value.bytesWritten) && + Number(value.bytesWritten) === + new TextEncoder().encode(content).byteLength + (value.hasUtf8Bom ? 3 : 0) && + Number(value.bytesWritten) <= WORKSPACE_WRITE_MAX_BYTES + ); + } const maxResults = request.maxResults ?? 50; return ( hasOnlyKeys(value, SEARCH_RESULT_KEYS) && + typeof value.truncated === 'boolean' && Array.isArray(value.matches) && value.matches.length <= maxResults && value.matches.every(