inject requestFiles and requestOptionalFile

Author: charleslien

backend/src/__tests__/cost-aggregation.integration.test.ts

@@ -213,15 +213,15 @@ describe('Cost Aggregation Integration Tests', () => {
     }
 
     // Mock file reading
-    spyOn(websocketAction, 'requestFiles').mockImplementation(
-      async (params: { ws: any; filePaths: string[] }) => {
-        const results: Record<string, string | null> = {}
-        params.filePaths.forEach((path) => {
-          results[path] = path === 'hello.txt' ? 'Hello, World!' : null
-        })
-        return results
-      },
-    )
+    agentRuntimeScopedImpl.requestFiles = async (params: {
+      filePaths: string[]
+    }) => {
+      const results: Record<string, string | null> = {}
+      params.filePaths.forEach((path) => {
+        results[path] = path === 'hello.txt' ? 'Hello, World!' : null
+      })
+      return results
+    }
 
     // Mock live user input checking
     const liveUserInputs = await import('../live-user-inputs')

backend/src/__tests__/main-prompt.integration.test.ts

@@ -27,7 +27,6 @@ import * as getDocumentationForQueryModule from '../get-documentation-for-query'
 import * as liveUserInputs from '../live-user-inputs'
 import { mainPrompt } from '../main-prompt'
 import * as processFileBlockModule from '../process-file-block'
-import * as websocketAction from '../websockets/websocket-action'
 
 import type { AgentTemplate } from '@codebuff/common/types/agent-template'
 import type {
@@ -117,28 +116,24 @@ describe('mainPrompt', () => {
     mockAgentStream('Test response')
 
     // Mock websocket actions
-    spyOn(websocketAction, 'requestFiles').mockImplementation(
-      async (params: { ws: any; filePaths: string[] }) => {
-        const results: Record<string, string | null> = {}
-        params.filePaths.forEach((p) => {
-          if (p === 'test.txt') {
-            results[p] = 'mock content for test.txt'
-          } else {
-            results[p] = null
-          }
-        })
-        return results
-      },
-    )
-
-    spyOn(websocketAction, 'requestFile').mockImplementation(
-      async (params: { ws: any; filePath: string }) => {
-        if (params.filePath === 'test.txt') {
-          return 'mock content for test.txt'
+    agentRuntimeScopedImpl.requestFiles = async ({ filePaths }) => {
+      const results: Record<string, string | null> = {}
+      filePaths.forEach((p) => {
+        if (p === 'test.txt') {
+          results[p] = 'mock content for test.txt'
+        } else {
+          results[p] = null
         }
-        return null
-      },
-    )
+      })
+      return results
+    }
+
+    agentRuntimeScopedImpl.requestOptionalFile = async ({ filePath }) => {
+      if (filePath === 'test.txt') {
+        return 'mock content for test.txt'
+      }
+      return null
+    }
 
     agentRuntimeScopedImpl.requestToolCall = mock(
       async ({

backend/src/__tests__/main-prompt.test.ts

@@ -20,7 +20,6 @@ import {
 
 import { MockWebSocket, mockFileContext } from './test-utils'
 import { processStreamWithTools } from '../tools/stream-parser'
-import * as websocketAction from '../websockets/websocket-action'
 
 import type { AgentTemplate } from '../templates/types'
 import type {
@@ -72,8 +71,8 @@ describe('malformed tool call error handling', () => {
     )
 
     // Mock websocket actions
-    spyOn(websocketAction, 'requestFiles').mockImplementation(async () => ({}))
-    spyOn(websocketAction, 'requestFile').mockImplementation(async () => null)
+    agentRuntimeScopedImpl.requestFiles = async () => ({})
+    agentRuntimeScopedImpl.requestOptionalFile = async () => null
     agentRuntimeScopedImpl.requestToolCall = async () => ({
       output: [
         {

backend/src/__tests__/malformed-tool-call.test.ts

@@ -15,7 +15,6 @@ import {
 } from 'bun:test'
 
 import { loopAgentSteps } from '../run-agent-step'
-import * as websocketAction from '../websockets/websocket-action'
 
 import type { AgentTemplate } from '../templates/types'
 import type {
@@ -126,15 +125,13 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
     }
 
     // Mock file operations
-    spyOn(websocketAction, 'requestFiles').mockImplementation(
-      async (params: { ws: any; filePaths: string[] }) => {
-        const results: Record<string, string | null> = {}
-        params.filePaths.forEach((path) => {
-          results[path] = null
-        })
-        return results
-      },
-    )
+    agentRuntimeScopedImpl.requestFiles = async ({ filePaths }) => {
+      const results: Record<string, string | null> = {}
+      filePaths.forEach((path) => {
+        results[path] = null
+      })
+      return results
+    }
 
     agentRuntimeScopedImpl.requestToolCall = async () => ({
       output: [

backend/src/__tests__/prompt-caching-subagents.test.ts

@@ -90,24 +90,8 @@ describe('read_docs tool with researcher agent', () => {
     mockedFunctions.push({ name: 'bigquery.insertTrace', spy: insertTraceSpy })
 
     // Mock websocket actions
-    const requestFilesSpy = spyOn(
-      websocketAction,
-      'requestFiles',
-    ).mockImplementation(async () => ({}))
-    mockedFunctions.push({
-      name: 'websocketAction.requestFiles',
-      spy: requestFilesSpy,
-    })
-
-    const requestFileSpy = spyOn(
-      websocketAction,
-      'requestFile',
-    ).mockImplementation(async () => null)
-    mockedFunctions.push({
-      name: 'websocketAction.requestFile',
-      spy: requestFileSpy,
-    })
-
+    agentRuntimeScopedImpl.requestFiles = async () => ({})
+    agentRuntimeScopedImpl.requestOptionalFile = async () => null
     agentRuntimeScopedImpl.requestToolCall = async () => ({
       output: [
         {

backend/src/__tests__/read-docs-tool.test.ts

@@ -24,18 +24,24 @@ import * as liveUserInputs from '../live-user-inputs'
 import { runAgentStep } from '../run-agent-step'
 import { clearAgentGeneratorCache } from '../run-programmatic-step'
 import { asUserMessage } from '../util/messages'
-import * as websocketAction from '../websockets/websocket-action'
 
 import type { AgentTemplate } from '../templates/types'
-import type { AgentRuntimeDeps } from '@codebuff/common/types/contracts/agent-runtime'
+import type {
+  AgentRuntimeDeps,
+  AgentRuntimeScopedDeps,
+} from '@codebuff/common/types/contracts/agent-runtime'
 import type { ProjectFileContext } from '@codebuff/common/util/file'
 import type { WebSocket } from 'ws'
 
 describe('runAgentStep - set_output tool', () => {
   let testAgent: AgentTemplate
-  let agentRuntimeImpl: AgentRuntimeDeps = { ...TEST_AGENT_RUNTIME_IMPL }
+  let agentRuntimeImpl: AgentRuntimeDeps
+  let agentRuntimeScopedImpl: AgentRuntimeScopedDeps
 
   beforeEach(async () => {
+    agentRuntimeImpl = { ...TEST_AGENT_RUNTIME_IMPL }
+    agentRuntimeScopedImpl = { ...TEST_AGENT_RUNTIME_SCOPED_IMPL }
+
     // Create a test agent that supports set_output
     testAgent = {
       id: 'test-set-output-agent',
@@ -78,32 +84,27 @@ describe('runAgentStep - set_output tool', () => {
     spyOn(liveUserInputs, 'startUserInput').mockImplementation(() => {})
     spyOn(liveUserInputs, 'setSessionConnected').mockImplementation(() => {})
 
-    spyOn(websocketAction, 'requestFiles').mockImplementation(
-      async (params: { ws: any; filePaths: string[] }) => {
-        const results: Record<string, string | null> = {}
-        params.filePaths.forEach((p) => {
-          if (p === 'src/auth.ts') {
-            results[p] = 'export function authenticate() { return true; }'
-          } else if (p === 'src/user.ts') {
-            results[p] = 'export interface User { id: string; name: string; }'
-          } else {
-            results[p] = null
-          }
-        })
-        return results
-      },
-    )
-
-    spyOn(websocketAction, 'requestFile').mockImplementation(
-      async (params: { ws: any; filePath: string }) => {
-        if (params.filePath === 'src/auth.ts') {
-          return 'export function authenticate() { return true; }'
-        } else if (params.filePath === 'src/user.ts') {
-          return 'export interface User { id: string; name: string; }'
+    agentRuntimeScopedImpl.requestFiles = async ({ filePaths }) => {
+      const results: Record<string, string | null> = {}
+      filePaths.forEach((p) => {
+        if (p === 'src/auth.ts') {
+          results[p] = 'export function authenticate() { return true; }'
+        } else if (p === 'src/user.ts') {
+          results[p] = 'export interface User { id: string; name: string; }'
+        } else {
+          results[p] = null
         }
-        return null
-      },
-    )
+      })
+      return results
+    }
+    agentRuntimeScopedImpl.requestOptionalFile = async ({ filePath }) => {
+      if (filePath === 'src/auth.ts') {
+        return 'export function authenticate() { return true; }'
+      } else if (filePath === 'src/user.ts') {
+        return 'export interface User { id: string; name: string; }'
+      }
+      return null
+    }
 
     // Don't mock requestToolCall for integration test - let real tool execution happen
 
@@ -116,7 +117,6 @@ describe('runAgentStep - set_output tool', () => {
 
   afterEach(() => {
     mock.restore()
-    agentRuntimeImpl = { ...TEST_AGENT_RUNTIME_IMPL }
   })
 
   afterAll(() => {
@@ -363,19 +363,17 @@ describe('runAgentStep - set_output tool', () => {
     }
 
     // Mock requestFiles to return test file content
-    spyOn(websocketAction, 'requestFiles').mockImplementation(
-      async (params: { ws: any; filePaths: string[] }) => {
-        const results: Record<string, string | null> = {}
-        params.filePaths.forEach((p) => {
-          if (p === 'src/test.ts') {
-            results[p] = 'export function testFunction() { return "test"; }'
-          } else {
-            results[p] = null
-          }
-        })
-        return results
-      },
-    )
+    agentRuntimeScopedImpl.requestFiles = async ({ filePaths }) => {
+      const results: Record<string, string | null> = {}
+      filePaths.forEach((p) => {
+        if (p === 'src/test.ts') {
+          results[p] = 'export function testFunction() { return "test"; }'
+        } else {
+          results[p] = null
+        }
+      })
+      return results
+    }
 
     // Mock the LLM stream to return a response that doesn't end the turn
     agentRuntimeImpl.promptAiSdkStream = async function* ({}) {

backend/src/__tests__/run-agent-step-tools.test.ts

@@ -28,7 +28,6 @@ import { MockWebSocket, mockFileContext } from './test-utils'
 import * as linkupApi from '../llm-apis/linkup-api'
 import { runAgentStep } from '../run-agent-step'
 import { assembleLocalAgentTemplates } from '../templates/agent-registry'
-import * as websocketAction from '../websockets/websocket-action'
 
 import type {
   AgentRuntimeDeps,
@@ -64,8 +63,8 @@ describe('web_search tool with researcher agent', () => {
     )
 
     // Mock websocket actions
-    spyOn(websocketAction, 'requestFiles').mockImplementation(async () => ({}))
-    spyOn(websocketAction, 'requestFile').mockImplementation(async () => null)
+    agentRuntimeScopedImpl.requestFiles = async () => ({})
+    agentRuntimeScopedImpl.requestOptionalFile = async () => null
     agentRuntimeScopedImpl.requestToolCall = async () => ({
       output: [
         {

backend/src/__tests__/web-search-tool.test.ts

@@ -1,9 +1,15 @@
+import { toOptionalFile } from '@codebuff/common/old-constants'
+import { ensureEndsWithNewline } from '@codebuff/common/util/file'
 import { generateCompactId } from '@codebuff/common/util/string'
 
 import { subscribeToAction } from './websockets/websocket-action'
 
 import type { ServerAction } from '@codebuff/common/actions'
-import type { RequestMcpToolDataFn } from '@codebuff/common/types/contracts/client'
+import type {
+  RequestFilesFn,
+  RequestMcpToolDataFn,
+  RequestOptionalFileFn,
+} from '@codebuff/common/types/contracts/client'
 import type { ParamsOf } from '@codebuff/common/types/function-params'
 import type { MCPConfig } from '@codebuff/common/types/mcp'
 import type { ToolResultOutput } from '@codebuff/common/types/messages/content-part'
@@ -141,3 +147,44 @@ export async function requestMcpToolDataWs(
     })
   })
 }
+
+/**
+ * Requests multiple files from the client
+ * @param ws - The WebSocket connection
+ * @param filePaths - Array of file paths to request
+ * @returns Promise resolving to an object mapping file paths to their contents
+ */
+export async function requestFilesWs(
+  params: {
+    ws: WebSocket
+  } & ParamsOf<RequestFilesFn>,
+): ReturnType<RequestFilesFn> {
+  const { ws, filePaths } = params
+  return new Promise<Record<string, string | null>>((resolve) => {
+    const requestId = generateCompactId()
+    const unsubscribe = subscribeToAction('read-files-response', (action) => {
+      for (const [filename, contents] of Object.entries(action.files)) {
+        action.files[filename] = ensureEndsWithNewline(contents)
+      }
+      if (action.requestId === requestId) {
+        unsubscribe()
+        resolve(action.files)
+      }
+    })
+    sendAction(ws, {
+      type: 'read-files',
+      filePaths,
+      requestId,
+    })
+  })
+}
+
+export async function requestOptionalFileWs(
+  params: {
+    ws: WebSocket
+  } & ParamsOf<RequestOptionalFileFn>,
+): ReturnType<RequestOptionalFileFn> {
+  const { ws, filePath } = params
+  const files = await requestFilesWs({ ws, filePaths: [filePath] })
+  return toOptionalFile(files[filePath] ?? null)
+}

backend/src/client-wrapper.ts

@@ -1,19 +1,20 @@
 import { uniq } from 'lodash'
-import type { WebSocket } from 'ws'
 
-import { requestFiles } from './websockets/websocket-action'
+import type { RequestFilesFn } from '@codebuff/common/types/contracts/client'
 
-export async function getFileReadingUpdates(
-  ws: WebSocket,
-  requestedFiles: string[],
-): Promise<
+export async function getFileReadingUpdates(params: {
+  requestFiles: RequestFilesFn
+  requestedFiles: string[]
+}): Promise<
   {
     path: string
     content: string
   }[]
 > {
+  const { requestFiles, requestedFiles } = params
+
   const allFilePaths = uniq(requestedFiles)
-  const loadedFiles = await requestFiles({ ws, filePaths: allFilePaths })
+  const loadedFiles = await requestFiles({ filePaths: allFilePaths })
 
   const addedFiles = allFilePaths
     .filter(