refactor(backend): decouple tool execution from WebSocket layer

brandonkachen · brandonkachen · commit c710a2dfdb8c · 2025-10-15T14:41:50.000-07:00
Replace direct WebSocket usage in tool handlers and programmatic steps
with `sendSubagentChunk` callback function. This improves testability
and separates transport concerns from business logic.

🤖 Generated with Codebuff Co-Authored-By: Codebuff
&lt;noreply@codebuff.com&gt;
diff --git a/backend/src/__tests__/run-programmatic-step.test.ts b/backend/src/__tests__/run-programmatic-step.test.ts
@@ -15,11 +15,10 @@ import {
   clearAgentGeneratorCache,
   runProgrammaticStep,
 } from '../run-programmatic-step'
-import { mockFileContext, MockWebSocket } from './test-utils'
+import { mockFileContext } from './test-utils'
 import * as agentRun from '../agent-run'
 import * as toolExecutor from '../tools/tool-executor'
 import * as requestContext from '../websockets/request-context'
-import * as websocketAction from '../websockets/websocket-action'
 
 import type { AgentTemplate, StepGenerator } from '../templates/types'
 import type { PublicAgentState } from '@codebuff/common/types/agent-template'
@@ -29,7 +28,7 @@ import type {
 } from '@codebuff/common/types/messages/content-part'
 import type { AgentState } from '@codebuff/common/types/session-state'
 import type { Logger } from '@codebuff/common/types/contracts/logger'
-import type { WebSocket } from 'ws'
+import type { SendSubagentChunkFn } from '@codebuff/common/types/contracts/client'
 
 const logger: Logger = {
   debug: () => {},
@@ -45,7 +44,8 @@ describe('runProgrammaticStep', () => {
   let executeToolCallSpy: any
   let getRequestContextSpy: any
   let addAgentStepSpy: any
-  let sendActionSpy: any
+  let sendSubagentChunk: SendSubagentChunkFn
+  let sentSubagentChunks: Parameters<SendSubagentChunkFn>[0][]
 
   beforeEach(() => {
     // Mock analytics
@@ -72,10 +72,10 @@ describe('runProgrammaticStep', () => {
       async () => 'test-step-id',
     )
 
-    // Mock sendAction
-    sendActionSpy = spyOn(websocketAction, 'sendAction').mockImplementation(
-      () => {},
-    )
+    sentSubagentChunks = []
+    sendSubagentChunk = (data) => {
+      sentSubagentChunks.push(data)
+    }
 
     // Mock crypto.randomUUID
     spyOn(crypto, 'randomUUID').mockImplementation(
@@ -132,10 +132,10 @@ describe('runProgrammaticStep', () => {
       fileContext: mockFileContext,
       assistantMessage: undefined,
       assistantPrefix: undefined,
-      ws: new MockWebSocket() as unknown as WebSocket,
       localAgentTemplates: {},
       stepsComplete: false,
       stepNumber: 1,
+      sendSubagentChunk,
       logger,
     }
   })
@@ -226,14 +226,6 @@ describe('runProgrammaticStep', () => {
       mockTemplate.handleSteps = () => mockGenerator
       mockTemplate.toolNames = ['add_message', 'read_files', 'end_turn']
 
-      // Track chunks sent via sendSubagentChunk
-      const sentChunks: string[] = []
-      sendActionSpy.mockImplementation((ws: any, action: any) => {
-        if (action.type === 'subagent-response-chunk') {
-          sentChunks.push(action.chunk)
-        }
-      })
-
       const result = await runProgrammaticStep(mockParams)
 
       // Verify add_message tool was executed
@@ -253,15 +245,16 @@ describe('runProgrammaticStep', () => {
       )
 
       // Check that no tool call chunk was sent for add_message
-      const addMessageToolCallChunk = sentChunks.find(
-        (chunk) =>
+      const addMessageToolCallChunk = sentSubagentChunks.find(
+        ({ chunk }) =>
           chunk.includes('add_message') && chunk.includes('Hello world'),
       )
       expect(addMessageToolCallChunk).toBeUndefined()
 
       // Check that tool call chunk WAS sent for read_files (normal behavior)
-      const readFilesToolCallChunk = sentChunks.find(
-        (chunk) => chunk.includes('read_files') && chunk.includes('test.txt'),
+      const readFilesToolCallChunk = sentSubagentChunks.find(
+        ({ chunk }) =>
+          chunk.includes('read_files') && chunk.includes('test.txt'),
       )
       expect(readFilesToolCallChunk).toBeDefined()
 
diff --git a/backend/src/__tests__/sandbox-generator.test.ts b/backend/src/__tests__/sandbox-generator.test.ts
@@ -8,14 +8,13 @@ import {
   clearAgentGeneratorCache,
   runProgrammaticStep,
 } from '../run-programmatic-step'
-import { mockFileContext, MockWebSocket } from './test-utils'
+import { mockFileContext } from './test-utils'
 import * as agentRun from '../agent-run'
 import * as requestContext from '../websockets/request-context'
-import * as websocketAction from '../websockets/websocket-action'
 
 import type { AgentTemplate } from '../templates/types'
 import type { Logger } from '@codebuff/common/types/contracts/logger'
-import type { WebSocket } from 'ws'
+import type { SendSubagentChunkFn } from '@codebuff/common/types/contracts/client'
 
 const logger: Logger = {
   debug: () => {},
@@ -28,6 +27,7 @@ describe('QuickJS Sandbox Generator', () => {
   let mockAgentState: AgentState
   let mockParams: any
   let mockTemplate: AgentTemplate
+  let sendSubagentChunk: SendSubagentChunkFn
 
   beforeEach(() => {
     clearAgentGeneratorCache({ logger })
@@ -39,11 +39,11 @@ describe('QuickJS Sandbox Generator', () => {
     spyOn(requestContext, 'getRequestContext').mockImplementation(() => ({
       processedRepoId: 'test-repo-id',
     }))
-    spyOn(websocketAction, 'sendAction').mockImplementation(() => {})
     spyOn(crypto, 'randomUUID').mockImplementation(
       () =>
         'mock-uuid-0000-0000-0000-000000000000' as `${string}-${string}-${string}-${string}-${string}`,
     )
+    sendSubagentChunk = () => {}
 
     // Reuse common test data structure
     mockAgentState = {
@@ -91,10 +91,10 @@ describe('QuickJS Sandbox Generator', () => {
       fileContext: mockFileContext,
       assistantMessage: undefined,
       assistantPrefix: undefined,
-      ws: new MockWebSocket() as unknown as WebSocket,
       localAgentTemplates: {},
       stepsComplete: false,
       stepNumber: 1,
+      sendSubagentChunk,
       logger,
     }
   })
diff --git a/backend/src/__tests__/spawn-agents-permissions.test.ts b/backend/src/__tests__/spawn-agents-permissions.test.ts
@@ -739,12 +739,12 @@ describe('Spawn Agents Permissions', () => {
           writeToClient: () => {},
           getLatestState: () => ({ messages: [] }),
           state: {
-            // Missing required fields like ws, fingerprintId, etc.
+            // Missing required fields like fingerprintId, messages, etc.
             agentTemplate: parentAgent,
             localAgentTemplates: {},
           },
         })
-      }).toThrow('Missing WebSocket in state')
+      }).toThrow('Missing fingerprintId in state')
       expect(mockLoopAgentSteps).not.toHaveBeenCalled()
     })
   })
diff --git a/backend/src/run-programmatic-step.ts b/backend/src/run-programmatic-step.ts
@@ -6,15 +6,17 @@ import { addAgentStep } from './agent-run'
 import { executeToolCall } from './tools/tool-executor'
 import { SandboxManager } from './util/quickjs-sandbox'
 import { getRequestContext } from './websockets/request-context'
-import { sendAction } from './websockets/websocket-action'
 
 import type { CodebuffToolCall } from '@codebuff/common/tools/list'
 import type {
   AgentTemplate,
   StepGenerator,
   PublicAgentState,
 } from '@codebuff/common/types/agent-template'
-import type { HandleStepsLogChunkFn } from '@codebuff/common/types/contracts/client'
+import type {
+  HandleStepsLogChunkFn,
+  SendSubagentChunkFn,
+} from '@codebuff/common/types/contracts/client'
 import type { Logger } from '@codebuff/common/types/contracts/logger'
 import type {
   ParamsExcluding,
@@ -26,7 +28,6 @@ import type {
 } from '@codebuff/common/types/messages/content-part'
 import type { PrintModeEvent } from '@codebuff/common/types/print-mode'
 import type { AgentState } from '@codebuff/common/types/session-state'
-import type { WebSocket } from 'ws'
 
 // Global sandbox manager for QuickJS contexts
 const sandboxManager = new SandboxManager()
@@ -59,10 +60,10 @@ export async function runProgrammaticStep(
     userInputId: string
     fingerprintId: string
     onResponseChunk: (chunk: string | PrintModeEvent) => void
-    ws: WebSocket
     localAgentTemplates: Record<string, AgentTemplate>
     stepsComplete: boolean
     stepNumber: number
+    sendSubagentChunk: SendSubagentChunkFn | undefined
     handleStepsLogChunk: HandleStepsLogChunkFn
     logger: Logger
   } & ParamsExcluding<
@@ -91,9 +92,9 @@ export async function runProgrammaticStep(
     userInputId,
     fingerprintId,
     onResponseChunk,
-    ws,
     localAgentTemplates,
     stepsComplete,
+    sendSubagentChunk,
     handleStepsLogChunk,
     logger,
   } = params
@@ -179,25 +180,13 @@ export async function runProgrammaticStep(
   const toolCalls: CodebuffToolCall[] = []
   const toolResults: ToolResultPart[] = []
   const state = {
-    ws,
     fingerprintId,
     userId,
     repoId,
     agentTemplate: template,
     localAgentTemplates,
     system,
-    sendSubagentChunk: (data: {
-      userInputId: string
-      agentId: string
-      agentType: string
-      chunk: string
-      prompt?: string
-    }) => {
-      sendAction(ws, {
-        type: 'subagent-response-chunk',
-        ...data,
-      })
-    },
+    sendSubagentChunk,
     agentState: cloneDeep({
       ...agentState,
       runId: agentState.runId!, // We've already verified runId exists above
diff --git a/backend/src/tools/handlers/tool/read-files.ts b/backend/src/tools/handlers/tool/read-files.ts
@@ -9,8 +9,6 @@ import type {
 import type { ParamsExcluding } from '@codebuff/common/types/function-params'
 import type { Message } from '@codebuff/common/types/messages/codebuff-message'
 import type { ProjectFileContext } from '@codebuff/common/util/file'
-import type { WebSocket } from 'ws'
-
 type ToolName = 'read_files'
 export const handleReadFiles = ((
   params: {
@@ -21,7 +19,6 @@ export const handleReadFiles = ((
     fileContext: ProjectFileContext
 
     state: {
-      ws?: WebSocket
       userId?: string
       fingerprintId?: string
       repoId?: string
@@ -39,11 +36,8 @@ export const handleReadFiles = ((
     fileContext,
     state,
   } = params
-  const { ws, fingerprintId, userId, repoId, messages } = state
+  const { fingerprintId, userId, repoId, messages } = state
   const { paths } = toolCall.input
-  if (!ws) {
-    throw new Error('Internal error for read_files: Missing WebSocket in state')
-  }
   if (!messages) {
     throw new Error('Internal error for read_files: Missing messages in state')
   }
diff --git a/backend/src/tools/handlers/tool/spawn-agent-inline.ts b/backend/src/tools/handlers/tool/spawn-agent-inline.ts
@@ -19,7 +19,6 @@ import type { AgentState } from '@codebuff/common/types/session-state'
 import type { ProjectFileContext } from '@codebuff/common/util/file'
 import type { ParamsExcluding } from '@codebuff/common/types/function-params'
 import type { Logger } from '@codebuff/common/types/contracts/logger'
-import type { WebSocket } from 'ws'
 
 type ToolName = 'spawn_agent_inline'
 export const handleSpawnAgentInline = ((
@@ -33,7 +32,6 @@ export const handleSpawnAgentInline = ((
 
     getLatestState: () => { messages: Message[] }
     state: {
-      ws?: WebSocket
       fingerprintId?: string
       userId?: string
       agentTemplate?: AgentTemplate
@@ -56,7 +54,6 @@ export const handleSpawnAgentInline = ((
     | 'parentSystemPrompt'
     | 'onResponseChunk'
     | 'clearUserPromptMessagesAfterResponse'
-    | 'ws'
     | 'fingerprintId'
   >,
 ): { result: Promise<CodebuffToolOutput<ToolName>>; state: {} } => {
@@ -74,7 +71,6 @@ export const handleSpawnAgentInline = ((
     params: spawnParams,
   } = toolCall.input
   const {
-    ws,
     fingerprintId,
     userId,
     agentTemplate: parentAgentTemplate,
@@ -115,7 +111,6 @@ export const handleSpawnAgentInline = ((
 
     const result = await executeSubagent({
       ...params,
-      ws,
       userInputId: `${userInputId}-inline-${agentType}${childAgentState.agentId}`,
       prompt: prompt || '',
       spawnParams,
diff --git a/backend/src/tools/handlers/tool/spawn-agent-utils.ts b/backend/src/tools/handlers/tool/spawn-agent-utils.ts
@@ -19,16 +19,13 @@ import type {
   Subgoal,
 } from '@codebuff/common/types/session-state'
 import type { ProjectFileContext } from '@codebuff/common/util/file'
-import type { WebSocket } from 'ws'
-
 export interface SpawnAgentParams {
   agent_type: string
   prompt?: string
   params?: any
 }
 
 export interface BaseSpawnState {
-  ws?: WebSocket
   fingerprintId?: string
   userId?: string
   agentTemplate?: AgentTemplate
@@ -53,7 +50,6 @@ export function validateSpawnState(
   toolName: string,
 ): Omit<Required<BaseSpawnState>, 'userId'> & { userId: string | undefined } {
   const {
-    ws,
     fingerprintId,
     agentTemplate: parentAgentTemplate,
     localAgentTemplates,
@@ -63,11 +59,6 @@ export function validateSpawnState(
     system,
   } = state
 
-  if (!ws) {
-    throw new Error(
-      `Internal error for ${toolName}: Missing WebSocket in state`,
-    )
-  }
   if (!fingerprintId) {
     throw new Error(
       `Internal error for ${toolName}: Missing fingerprintId in state`,
@@ -96,7 +87,6 @@ export function validateSpawnState(
   }
 
   return {
-    ws,
     fingerprintId,
     userId,
     agentTemplate: parentAgentTemplate,
diff --git a/backend/src/tools/handlers/tool/spawn-agents.ts b/backend/src/tools/handlers/tool/spawn-agents.ts
@@ -18,7 +18,6 @@ import type { ParamsExcluding } from '@codebuff/common/types/function-params'
 import type { Message } from '@codebuff/common/types/messages/codebuff-message'
 import type { PrintModeEvent } from '@codebuff/common/types/print-mode'
 import type { AgentState } from '@codebuff/common/types/session-state'
-import type { WebSocket } from 'ws'
 
 export type SendSubagentChunk = (data: {
   userInputId: string
@@ -40,7 +39,6 @@ export const handleSpawnAgents = ((
 
     getLatestState: () => { messages: Message[] }
     state: {
-      ws?: WebSocket
       fingerprintId?: string
       userId?: string
       agentTemplate?: AgentTemplate
@@ -57,7 +55,6 @@ export const handleSpawnAgents = ((
   > &
     ParamsExcluding<
       typeof executeSubagent,
-      | 'ws'
       | 'userInputId'
       | 'prompt'
       | 'spawnParams'
@@ -93,7 +90,6 @@ export const handleSpawnAgents = ((
   }
 
   const {
-    ws,
     fingerprintId,
     userId,
     agentTemplate: parentAgentTemplate,
@@ -135,7 +131,6 @@ export const handleSpawnAgents = ((
 
           const result = await executeSubagent({
             ...params,
-            ws,
             userInputId: `${userInputId}-${agentType}${subAgentState.agentId}`,
             prompt: prompt || '',
             spawnParams,
