Revert "refactor(backend): decouple tool execution from WebSocket layer"

This reverts commit c710a2d.

Author: brandonkachen

backend/src/__tests__/run-programmatic-step.test.ts

@@ -15,10 +15,11 @@ import {
   clearAgentGeneratorCache,
   runProgrammaticStep,
 } from '../run-programmatic-step'
-import { mockFileContext } from './test-utils'
+import { mockFileContext, MockWebSocket } 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'
@@ -28,7 +29,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 { SendSubagentChunkFn } from '@codebuff/common/types/contracts/client'
+import type { WebSocket } from 'ws'
 
 const logger: Logger = {
   debug: () => {},
@@ -44,8 +45,7 @@ describe('runProgrammaticStep', () => {
   let executeToolCallSpy: any
   let getRequestContextSpy: any
   let addAgentStepSpy: any
-  let sendSubagentChunk: SendSubagentChunkFn
-  let sentSubagentChunks: Parameters<SendSubagentChunkFn>[0][]
+  let sendActionSpy: any
 
   beforeEach(() => {
     // Mock analytics
@@ -72,10 +72,10 @@ describe('runProgrammaticStep', () => {
       async () => 'test-step-id',
     )
 
-    sentSubagentChunks = []
-    sendSubagentChunk = (data) => {
-      sentSubagentChunks.push(data)
-    }
+    // Mock sendAction
+    sendActionSpy = spyOn(websocketAction, 'sendAction').mockImplementation(
+      () => {},
+    )
 
     // 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,6 +226,14 @@ 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
@@ -245,16 +253,15 @@ describe('runProgrammaticStep', () => {
       )
 
       // Check that no tool call chunk was sent for add_message
-      const addMessageToolCallChunk = sentSubagentChunks.find(
-        ({ chunk }) =>
+      const addMessageToolCallChunk = sentChunks.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 = sentSubagentChunks.find(
-        ({ chunk }) =>
-          chunk.includes('read_files') && chunk.includes('test.txt'),
+      const readFilesToolCallChunk = sentChunks.find(
+        (chunk) => chunk.includes('read_files') && chunk.includes('test.txt'),
       )
       expect(readFilesToolCallChunk).toBeDefined()
 

backend/src/__tests__/sandbox-generator.test.ts

@@ -8,13 +8,14 @@ import {
   clearAgentGeneratorCache,
   runProgrammaticStep,
 } from '../run-programmatic-step'
-import { mockFileContext } from './test-utils'
+import { mockFileContext, MockWebSocket } 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 { SendSubagentChunkFn } from '@codebuff/common/types/contracts/client'
+import type { WebSocket } from 'ws'
 
 const logger: Logger = {
   debug: () => {},
@@ -27,7 +28,6 @@ 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,
     }
   })

backend/src/__tests__/spawn-agents-permissions.test.ts

@@ -739,12 +739,12 @@ describe('Spawn Agents Permissions', () => {
           writeToClient: () => {},
           getLatestState: () => ({ messages: [] }),
           state: {
-            // Missing required fields like fingerprintId, messages, etc.
+            // Missing required fields like ws, fingerprintId, etc.
             agentTemplate: parentAgent,
             localAgentTemplates: {},
           },
         })
-      }).toThrow('Missing fingerprintId in state')
+      }).toThrow('Missing WebSocket in state')
       expect(mockLoopAgentSteps).not.toHaveBeenCalled()
     })
   })

backend/src/run-programmatic-step.ts

@@ -6,17 +6,15 @@ 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,
-  SendSubagentChunkFn,
-} from '@codebuff/common/types/contracts/client'
+import type { HandleStepsLogChunkFn } from '@codebuff/common/types/contracts/client'
 import type { Logger } from '@codebuff/common/types/contracts/logger'
 import type {
   ParamsExcluding,
@@ -28,6 +26,7 @@ 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()
@@ -60,10 +59,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<
@@ -92,9 +91,9 @@ export async function runProgrammaticStep(
     userInputId,
     fingerprintId,
     onResponseChunk,
+    ws,
     localAgentTemplates,
     stepsComplete,
-    sendSubagentChunk,
     handleStepsLogChunk,
     logger,
   } = params
@@ -180,13 +179,25 @@ export async function runProgrammaticStep(
   const toolCalls: CodebuffToolCall[] = []
   const toolResults: ToolResultPart[] = []
   const state = {
+    ws,
     fingerprintId,
     userId,
     repoId,
     agentTemplate: template,
     localAgentTemplates,
     system,
-    sendSubagentChunk,
+    sendSubagentChunk: (data: {
+      userInputId: string
+      agentId: string
+      agentType: string
+      chunk: string
+      prompt?: string
+    }) => {
+      sendAction(ws, {
+        type: 'subagent-response-chunk',
+        ...data,
+      })
+    },
     agentState: cloneDeep({
       ...agentState,
       runId: agentState.runId!, // We've already verified runId exists above

backend/src/tools/handlers/tool/read-files.ts

@@ -9,6 +9,8 @@ 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: {
@@ -19,6 +21,7 @@ export const handleReadFiles = ((
     fileContext: ProjectFileContext
 
     state: {
+      ws?: WebSocket
       userId?: string
       fingerprintId?: string
       repoId?: string
@@ -36,8 +39,11 @@ export const handleReadFiles = ((
     fileContext,
     state,
   } = params
-  const { fingerprintId, userId, repoId, messages } = state
+  const { ws, 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')
   }

backend/src/tools/handlers/tool/spawn-agent-inline.ts

@@ -19,6 +19,7 @@ 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 = ((
@@ -32,6 +33,7 @@ export const handleSpawnAgentInline = ((
 
     getLatestState: () => { messages: Message[] }
     state: {
+      ws?: WebSocket
       fingerprintId?: string
       userId?: string
       agentTemplate?: AgentTemplate
@@ -54,6 +56,7 @@ export const handleSpawnAgentInline = ((
     | 'parentSystemPrompt'
     | 'onResponseChunk'
     | 'clearUserPromptMessagesAfterResponse'
+    | 'ws'
     | 'fingerprintId'
   >,
 ): { result: Promise<CodebuffToolOutput<ToolName>>; state: {} } => {
@@ -71,6 +74,7 @@ export const handleSpawnAgentInline = ((
     params: spawnParams,
   } = toolCall.input
   const {
+    ws,
     fingerprintId,
     userId,
     agentTemplate: parentAgentTemplate,
@@ -111,6 +115,7 @@ export const handleSpawnAgentInline = ((
 
     const result = await executeSubagent({
       ...params,
+      ws,
       userInputId: `${userInputId}-inline-${agentType}${childAgentState.agentId}`,
       prompt: prompt || '',
       spawnParams,

backend/src/tools/handlers/tool/spawn-agent-utils.ts

@@ -19,13 +19,16 @@ 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
@@ -50,6 +53,7 @@ export function validateSpawnState(
   toolName: string,
 ): Omit<Required<BaseSpawnState>, 'userId'> & { userId: string | undefined } {
   const {
+    ws,
     fingerprintId,
     agentTemplate: parentAgentTemplate,
     localAgentTemplates,
@@ -59,6 +63,11 @@ 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`,
@@ -87,6 +96,7 @@ export function validateSpawnState(
   }
 
   return {
+    ws,
     fingerprintId,
     userId,
     agentTemplate: parentAgentTemplate,

backend/src/tools/handlers/tool/spawn-agents.ts

@@ -18,6 +18,7 @@ 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
@@ -39,6 +40,7 @@ export const handleSpawnAgents = ((
 
     getLatestState: () => { messages: Message[] }
     state: {
+      ws?: WebSocket
       fingerprintId?: string
       userId?: string
       agentTemplate?: AgentTemplate
@@ -55,6 +57,7 @@ export const handleSpawnAgents = ((
   > &
     ParamsExcluding<
       typeof executeSubagent,
+      | 'ws'
       | 'userInputId'
       | 'prompt'
       | 'spawnParams'
@@ -90,6 +93,7 @@ export const handleSpawnAgents = ((
   }
 
   const {
+    ws,
     fingerprintId,
     userId,
     agentTemplate: parentAgentTemplate,
@@ -131,6 +135,7 @@ export const handleSpawnAgents = ((
 
           const result = await executeSubagent({
             ...params,
+            ws,
             userInputId: `${userInputId}-${agentType}${subAgentState.agentId}`,
             prompt: prompt || '',
             spawnParams,