inject requestMcpToolData

Author: charleslien

backend/src/client-wrapper.ts

@@ -3,6 +3,8 @@ 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 { 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'
 import type { ServerMessage } from '@codebuff/common/websockets/websocket-schema'
@@ -94,3 +96,48 @@ export async function requestToolCallWs(params: {
     })
   })
 }
+
+/**
+ * Requests a tool call execution from the client with timeout support
+ * @param ws - The WebSocket connection
+ * @param mcpConfig - The configuration for the MCP server
+ * @param input - Arguments for the tool (can include timeout)
+ * @returns Promise resolving to the tool execution result
+ */
+export async function requestMcpToolDataWs(
+  params: ParamsOf<RequestMcpToolDataFn> & {
+    ws: WebSocket
+  },
+): ReturnType<RequestMcpToolDataFn> {
+  const { ws, mcpConfig, toolNames } = params
+
+  return new Promise((resolve) => {
+    const requestId = generateCompactId()
+
+    // Set up timeout
+    const timeoutHandle = setTimeout(
+      () => {
+        unsubscribe()
+        resolve([])
+      },
+      45_000 + 5000, // Convert to ms and add a small buffer
+    )
+
+    // Subscribe to response
+    const unsubscribe = subscribeToAction('mcp-tool-data', (action) => {
+      if (action.requestId === requestId) {
+        clearTimeout(timeoutHandle)
+        unsubscribe()
+        resolve(action.tools)
+      }
+    })
+
+    // Send the request
+    sendAction(ws, {
+      type: 'request-mcp-tool-data',
+      mcpConfig,
+      requestId,
+      ...(toolNames && { toolNames }),
+    })
+  })
+}

backend/src/mcp/util.ts

@@ -1,20 +1,22 @@
-import { requestMcpToolData } from '../websockets/websocket-action'
-
 import type { AgentTemplate } from '../templates/types'
+import type { RequestMcpToolDataFn } from '@codebuff/common/types/contracts/client'
+import type { OptionalFields } from '@codebuff/common/types/function-params'
 import type { ProjectFileContext } from '@codebuff/common/util/file'
-import type { WebSocket } from 'ws'
 
-export async function getMCPToolData({
-  ws,
-  toolNames,
-  mcpServers,
-  writeTo,
-}: {
-  ws: WebSocket
-  toolNames: AgentTemplate['toolNames']
-  mcpServers: AgentTemplate['mcpServers']
-  writeTo: ProjectFileContext['customToolDefinitions']
-}): Promise<ProjectFileContext['customToolDefinitions']> {
+export async function getMCPToolData(
+  params: OptionalFields<
+    {
+      toolNames: AgentTemplate['toolNames']
+      mcpServers: AgentTemplate['mcpServers']
+      writeTo: ProjectFileContext['customToolDefinitions']
+      requestMcpToolData: RequestMcpToolDataFn
+    },
+    'writeTo'
+  >,
+): Promise<ProjectFileContext['customToolDefinitions']> {
+  const withDefaults = { writeTo: {}, ...params }
+  const { toolNames, mcpServers, writeTo, requestMcpToolData } = withDefaults
+
   const requestedToolsByMcp: Record<string, string[] | undefined> = {}
   for (const t of toolNames) {
     if (!t.includes('/')) {
@@ -28,13 +30,11 @@ export async function getMCPToolData({
     requestedToolsByMcp[mcpName].push(toolName)
   }
 
-  writeTo ??= {}
   const promises: Promise<any>[] = []
   for (const [mcpName, mcpConfig] of Object.entries(mcpServers)) {
     promises.push(
       (async () => {
         const mcpData = await requestMcpToolData({
-          ws,
           mcpConfig,
           toolNames: requestedToolsByMcp[mcpName] ?? null,
         })

backend/src/run-agent-step.ts

@@ -35,10 +35,7 @@ import type {
 } from '@codebuff/common/types/contracts/database'
 import type { Logger } from '@codebuff/common/types/contracts/logger'
 import type { ParamsExcluding } from '@codebuff/common/types/function-params'
-import type {
-  AssistantMessage,
-  Message,
-} from '@codebuff/common/types/messages/codebuff-message'
+import type { Message } from '@codebuff/common/types/messages/codebuff-message'
 import type {
   ToolResultPart,
   TextPart,
@@ -53,8 +50,6 @@ import type {
 import type { ProjectFileContext } from '@codebuff/common/util/file'
 import type { WebSocket } from 'ws'
 
-
-
 export const runAgentStep = async (
   params: {
     ws: WebSocket
@@ -95,6 +90,10 @@ export const runAgentStep = async (
       | 'agentState'
       | 'agentTemplates'
       | 'additionalToolDefinitions'
+    > &
+    ParamsExcluding<
+      typeof getMCPToolData,
+      'toolNames' | 'mcpServers' | 'writeTo'
     >,
 ): Promise<{
   agentState: AgentState
@@ -204,7 +203,7 @@ export const runAgentStep = async (
         ),
       )
       return getMCPToolData({
-        ws,
+        ...params,
         toolNames: agentTemplate.toolNames,
         mcpServers: agentTemplate.mcpServers,
         writeTo: additionalToolDefinitions,
@@ -435,6 +434,10 @@ export const loopAgentSteps = async (
       | 'promptType'
       | 'agentTemplates'
       | 'additionalToolDefinitions'
+    > &
+    ParamsExcluding<
+      typeof getMCPToolData,
+      'toolNames' | 'mcpServers' | 'writeTo'
     >,
 ): Promise<{
   agentState: AgentState
@@ -495,7 +498,7 @@ export const loopAgentSteps = async (
         ),
       )
       return getMCPToolData({
-        ws,
+        ...params,
         toolNames: agentTemplate.toolNames,
         mcpServers: agentTemplate.mcpServers,
         writeTo: additionalToolDefinitions,
@@ -522,7 +525,7 @@ export const loopAgentSteps = async (
               ),
             )
             return getMCPToolData({
-              ws,
+              ...params,
               toolNames: agentTemplate.toolNames,
               mcpServers: agentTemplate.mcpServers,
               writeTo: additionalToolDefinitions,
@@ -767,5 +770,3 @@ export const loopAgentSteps = async (
     }
   }
 }
-
-

backend/src/tools/tool-executor.ts

@@ -33,7 +33,6 @@ import type {
   customToolDefinitionsSchema,
   ProjectFileContext,
 } from '@codebuff/common/util/file'
-import type { WebSocket } from 'ws'
 
 export type CustomToolCall = {
   toolName: string
@@ -121,7 +120,6 @@ export type ExecuteToolCallParams<T extends string = ToolName> = {
   toolResults: ToolResultPart[]
   toolResultsToAddAfterStream: ToolResultPart[]
   previousToolCallFinished: Promise<void>
-  ws: WebSocket
   agentTemplate: AgentTemplate
   fileContext: ProjectFileContext
   agentStepId: string
@@ -146,7 +144,6 @@ export function executeToolCall<T extends ToolName>(
     toolResults,
     toolResultsToAddAfterStream,
     previousToolCallFinished,
-    ws,
     agentTemplate,
     fileContext,
     agentStepId,
@@ -159,6 +156,7 @@ export function executeToolCall<T extends ToolName>(
     autoInsertEndStepParam = false,
     excludeToolFromMessageHistory = false,
     requestToolCall,
+    requestMcpToolData,
     logger,
   } = params
   const toolCall: CodebuffToolCall<T> | ToolCallError = parseRawToolCall<T>({
@@ -378,7 +376,6 @@ export async function executeCustomToolCall(
     toolResults,
     toolResultsToAddAfterStream,
     previousToolCallFinished,
-    ws,
     agentTemplate,
     fileContext,
     userInputId,
@@ -391,7 +388,7 @@ export async function executeCustomToolCall(
   } = params
   const toolCall: CustomToolCall | ToolCallError = parseRawCustomToolCall({
     customToolDefs: await getMCPToolData({
-      ws,
+      ...params,
       toolNames: agentTemplate.toolNames,
       mcpServers: agentTemplate.mcpServers,
       writeTo: cloneDeep(fileContext.customToolDefinitions),

backend/src/websockets/middleware.ts

@@ -15,7 +15,7 @@ import { eq } from 'drizzle-orm'
 import { getUserInfoFromApiKey } from './auth'
 import { updateRequestContext } from './request-context'
 import { sendAction } from './websocket-action'
-import { requestToolCallWs } from '../client-wrapper'
+import { requestMcpToolDataWs, requestToolCallWs } from '../client-wrapper'
 import { withAppContext } from '../context/app-context'
 import { BACKEND_AGENT_RUNTIME_IMPL } from '../impl/agent-runtime'
 import { checkAuth } from '../util/check-auth'
@@ -149,6 +149,7 @@ export class WebSocketMiddleware {
 
       const scopedDeps: AgentRuntimeScopedDeps = {
         requestToolCall: (params) => requestToolCallWs({ ...params, ws }),
+        requestMcpToolData: (params) => requestMcpToolDataWs({ ...params, ws }),
       }
 
       // Use the new combined context - much cleaner!

backend/src/websockets/websocket-action.ts

@@ -28,7 +28,6 @@ import type {
 import type { GetUserInfoFromApiKeyFn } from '@codebuff/common/types/contracts/database'
 import type { Logger } from '@codebuff/common/types/contracts/logger'
 import type { ParamsExcluding } from '@codebuff/common/types/function-params'
-import type { MCPConfig } from '@codebuff/common/types/mcp'
 import type { ClientMessage } from '@codebuff/common/websockets/websocket-schema'
 import type { WebSocket } from 'ws'
 
@@ -456,56 +455,3 @@ export async function requestOptionalFile(params: {
   const file = await requestFile({ ws, filePath })
   return toOptionalFile(file)
 }
-
-/**
- * Requests a tool call execution from the client with timeout support
- * @param ws - The WebSocket connection
- * @param mcpConfig - The configuration for the MCP server
- * @param input - Arguments for the tool (can include timeout)
- * @returns Promise resolving to the tool execution result
- */
-export async function requestMcpToolData({
-  ws,
-  mcpConfig,
-  toolNames,
-}: {
-  ws: WebSocket
-  mcpConfig: MCPConfig
-  toolNames: string[] | null
-}): Promise<
-  {
-    name: string
-    description?: string
-    inputSchema: unknown
-  }[]
-> {
-  return new Promise((resolve) => {
-    const requestId = generateCompactId()
-
-    // Set up timeout
-    const timeoutHandle = setTimeout(
-      () => {
-        unsubscribe()
-        resolve([])
-      },
-      45_000 + 5000, // Convert to ms and add a small buffer
-    )
-
-    // Subscribe to response
-    const unsubscribe = subscribeToAction('mcp-tool-data', (action) => {
-      if (action.requestId === requestId) {
-        clearTimeout(timeoutHandle)
-        unsubscribe()
-        resolve(action.tools)
-      }
-    })
-
-    // Send the request
-    sendAction(ws, {
-      type: 'request-mcp-tool-data',
-      mcpConfig,
-      requestId,
-      ...(toolNames && { toolNames }),
-    })
-  })
-}

common/src/testing/impl/agent-runtime.ts

@@ -44,4 +44,7 @@ export const TEST_AGENT_RUNTIME_SCOPED_IMPL: AgentRuntimeScopedDeps =
     requestToolCall: () => {
       throw new Error('requestToolCall not implemented in test runtime')
     },
+    requestMcpToolData: () => {
+      throw new Error('requestMcpToolData not implemented in test runtime')
+    },
   })

common/src/types/contracts/agent-runtime.ts

@@ -1,4 +1,4 @@
-import type { RequestToolCallFn } from './client'
+import type { RequestMcpToolDataFn, RequestToolCallFn } from './client'
 import type {
   AddAgentStepFn,
   FetchAgentFromDatabaseFn,
@@ -33,4 +33,5 @@ export type AgentRuntimeDeps = {
 export type AgentRuntimeScopedDeps = {
   // Client (WebSocket)
   requestToolCall: RequestToolCallFn
+  requestMcpToolData: RequestMcpToolDataFn
 }

common/src/types/contracts/client.ts

@@ -9,3 +9,14 @@ export type RequestToolCallFn = (params: {
 }) => Promise<{
   output: ToolResultOutput[]
 }>
+
+export type RequestMcpToolDataFn = (params: {
+  mcpConfig: MCPConfig
+  toolNames: string[] | null
+}) => Promise<
+  {
+    name: string
+    description?: string
+    inputSchema: unknown
+  }[]
+>

evals/scaffolding.ts

@@ -3,7 +3,10 @@ import { EventEmitter } from 'events'
 import fs from 'fs'
 import path from 'path'
 
-import { requestToolCallWs } from '@codebuff/backend/client-wrapper'
+import {
+  requestMcpToolDataWs,
+  requestToolCallWs,
+} from '@codebuff/backend/client-wrapper'
 import { runAgentStep } from '@codebuff/backend/run-agent-step'
 import { assembleLocalAgentTemplates } from '@codebuff/backend/templates/agent-registry'
 import { getFileTokenScores } from '@codebuff/code-map/parse'
@@ -183,6 +186,8 @@ export async function runAgentStepScaffolding(
 
   const agentRuntimeScopedImpl: AgentRuntimeScopedDeps = {
     requestToolCall: (params) => requestToolCallWs({ ...params, ws: mockWs }),
+    requestMcpToolData: (params) =>
+      requestMcpToolDataWs({ ...params, ws: mockWs }),
   }
   const result = await runAgentStep({
     ...EVALS_AGENT_RUNTIME_IMPL,