inject sendSubagentChunk

Author: charleslien

backend/src/client-wrapper.ts

@@ -9,6 +9,7 @@ import type {
   RequestFilesFn,
   RequestMcpToolDataFn,
   RequestOptionalFileFn,
+  SendSubagentChunkFn,
 } from '@codebuff/common/types/contracts/client'
 import type { ParamsOf } from '@codebuff/common/types/function-params'
 import type { MCPConfig } from '@codebuff/common/types/mcp'
@@ -188,3 +189,20 @@ export async function requestOptionalFileWs(
   const files = await requestFilesWs({ ws, filePaths: [filePath] })
   return toOptionalFile(files[filePath] ?? null)
 }
+
+export function sendSubagentChunkWs(
+  params: {
+    ws: WebSocket
+  } & ParamsOf<SendSubagentChunkFn>,
+): ReturnType<SendSubagentChunkFn> {
+  const { ws, userInputId, agentId, agentType, chunk, prompt } = params
+  return sendAction(ws, {
+    ...params,
+    type: 'subagent-response-chunk',
+    userInputId,
+    agentId,
+    agentType,
+    chunk,
+    prompt,
+  })
+}

backend/src/tools/stream-parser.ts

@@ -10,7 +10,6 @@ import { cloneDeep } from 'lodash'
 
 import { executeBatchStrReplaces } from './batch-str-replace'
 import { expireMessages } from '../util/messages'
-import { sendAction } from '../websockets/websocket-action'
 import { processStreamWithTags } from '../xml-stream-parser'
 import { executeCustomToolCall, executeToolCall } from './tool-executor'
 
@@ -19,6 +18,7 @@ import type { CustomToolCall, ExecuteToolCallParams } from './tool-executor'
 import type { AgentTemplate } from '../templates/types'
 import type { ToolName } from '@codebuff/common/tools/constants'
 import type { CodebuffToolCall } from '@codebuff/common/tools/list'
+import type { SendSubagentChunkFn } from '@codebuff/common/types/contracts/client'
 import type { StreamChunk } from '@codebuff/common/types/contracts/llm'
 import type { Logger } from '@codebuff/common/types/contracts/logger'
 import type { ParamsExcluding } from '@codebuff/common/types/function-params'
@@ -58,6 +58,7 @@ export async function processStreamWithTools(
     agentContext: Record<string, Subgoal>
     onResponseChunk: (chunk: string | PrintModeEvent) => void
     fullResponse: string
+    sendSubagentChunk: SendSubagentChunkFn
     logger: Logger
   } & Omit<
     ExecuteToolCallParams<any>,
@@ -90,6 +91,7 @@ export async function processStreamWithTools(
     system,
     agentState,
     onResponseChunk,
+    sendSubagentChunk,
     logger,
   } = params
   const fullResponseChunks: string[] = [params.fullResponse]
@@ -118,19 +120,7 @@ export async function processStreamWithTools(
     repoId,
     agentTemplate,
     localAgentTemplates,
-    sendSubagentChunk: (data: {
-      userInputId: string
-      agentId: string
-      agentType: string
-      chunk: string
-      prompt?: string
-    }) => {
-      sendAction(ws, {
-        type: 'subagent-response-chunk',
-        ...data,
-      })
-    },
-
+    sendSubagentChunk,
     agentState,
     agentContext,
     messages,

backend/src/websockets/middleware.ts

@@ -20,6 +20,7 @@ import {
   requestMcpToolDataWs,
   requestOptionalFileWs,
   requestToolCallWs,
+  sendSubagentChunkWs,
 } from '../client-wrapper'
 import { withAppContext } from '../context/app-context'
 import { BACKEND_AGENT_RUNTIME_IMPL } from '../impl/agent-runtime'
@@ -158,6 +159,7 @@ export class WebSocketMiddleware {
         requestFiles: (params) => requestFilesWs({ ...params, ws }),
         requestOptionalFile: (params) =>
           requestOptionalFileWs({ ...params, ws }),
+        sendSubagentChunk: (params) => sendSubagentChunkWs({ ...params, ws }),
       }
 
       // Use the new combined context - much cleaner!

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

@@ -3,6 +3,7 @@ import type {
   RequestMcpToolDataFn,
   RequestOptionalFileFn,
   RequestToolCallFn,
+  SendSubagentChunkFn,
 } from './client'
 import type {
   AddAgentStepFn,
@@ -41,4 +42,5 @@ export type AgentRuntimeScopedDeps = {
   requestMcpToolData: RequestMcpToolDataFn
   requestFiles: RequestFilesFn
   requestOptionalFile: RequestOptionalFileFn
+  sendSubagentChunk: SendSubagentChunkFn
 }

common/src/types/contracts/client.ts

@@ -28,3 +28,11 @@ export type RequestFilesFn = (params: {
 export type RequestOptionalFileFn = (params: {
   filePath: string
 }) => Promise<string | null>
+
+export type SendSubagentChunkFn = (params: {
+  userInputId: string
+  agentId: string
+  agentType: string
+  chunk: string
+  prompt?: string | undefined
+}) => void

evals/scaffolding.ts

@@ -4,8 +4,11 @@ import fs from 'fs'
 import path from 'path'
 
 import {
+  requestFilesWs,
   requestMcpToolDataWs,
+  requestOptionalFileWs,
   requestToolCallWs,
+  sendSubagentChunkWs,
 } from '@codebuff/backend/client-wrapper'
 import { runAgentStep } from '@codebuff/backend/run-agent-step'
 import { assembleLocalAgentTemplates } from '@codebuff/backend/templates/agent-registry'
@@ -28,7 +31,6 @@ import type {
   SDKAssistantMessage,
   SDKUserMessage,
 } from '@anthropic-ai/claude-code'
-import type { requestFilesWs as originalRequestFiles } from '@codebuff/backend/websockets/websocket-action'
 import type { ClientToolCall } from '@codebuff/common/tools/list'
 import type { AgentRuntimeScopedDeps } from '@codebuff/common/types/contracts/agent-runtime'
 import type {
@@ -80,7 +82,7 @@ export function createFileReadingMock(projectRoot: string) {
         files[filePath] = readMockFile(projectRoot, filePath)
       }
       return Promise.resolve(files)
-    }) satisfies typeof originalRequestFiles,
+    }) satisfies typeof requestFilesWs,
     requestToolCall: (async (params: {
       ws: WebSocket
       userInputId: string
@@ -188,6 +190,11 @@ export async function runAgentStepScaffolding(
     requestToolCall: (params) => requestToolCallWs({ ...params, ws: mockWs }),
     requestMcpToolData: (params) =>
       requestMcpToolDataWs({ ...params, ws: mockWs }),
+    requestFiles: (params) => requestFilesWs({ ...params, ws: mockWs }),
+    requestOptionalFile: (params) =>
+      requestOptionalFileWs({ ...params, ws: mockWs }),
+    sendSubagentChunk: (params) =>
+      sendSubagentChunkWs({ ...params, ws: mockWs }),
   }
   const result = await runAgentStep({
     ...EVALS_AGENT_RUNTIME_IMPL,