inject handleStepsLogChunk

Author: charleslien

backend/src/client-wrapper.ts

@@ -6,6 +6,7 @@ import { subscribeToAction } from './websockets/websocket-action'
 
 import type { ServerAction } from '@codebuff/common/actions'
 import type {
+  HandleStepsLogChunkFn,
   RequestFilesFn,
   RequestMcpToolDataFn,
   RequestOptionalFileFn,
@@ -197,7 +198,6 @@ export function sendSubagentChunkWs(
 ): ReturnType<SendSubagentChunkFn> {
   const { ws, userInputId, agentId, agentType, chunk, prompt } = params
   return sendAction(ws, {
-    ...params,
     type: 'subagent-response-chunk',
     userInputId,
     agentId,
@@ -206,3 +206,19 @@ export function sendSubagentChunkWs(
     prompt,
   })
 }
+
+export function handleStepsLogChunkWs(
+  params: {
+    ws: WebSocket
+  } & ParamsOf<HandleStepsLogChunkFn>,
+): ReturnType<HandleStepsLogChunkFn> {
+  const { ws, userInputId, runId, level, data, message } = params
+  return sendAction(ws, {
+    type: 'handlesteps-log-chunk',
+    userInputId,
+    agentId: runId,
+    level,
+    data,
+    message,
+  })
+}

backend/src/run-programmatic-step.ts

@@ -14,14 +14,18 @@ import type {
   StepGenerator,
   PublicAgentState,
 } from '@codebuff/common/types/agent-template'
+import type { HandleStepsLogChunkFn } from '@codebuff/common/types/contracts/client'
+import type { Logger } from '@codebuff/common/types/contracts/logger'
+import type {
+  ParamsExcluding,
+  ParamsOf,
+} from '@codebuff/common/types/function-params'
 import type {
   ToolResultOutput,
   ToolResultPart,
 } 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 { ParamsExcluding, ParamsOf } from '@codebuff/common/types/function-params'
-import type { Logger } from '@codebuff/common/types/contracts/logger'
 import type { WebSocket } from 'ws'
 
 // Global sandbox manager for QuickJS contexts
@@ -59,6 +63,7 @@ export async function runProgrammaticStep(
     localAgentTemplates: Record<string, AgentTemplate>
     stepsComplete: boolean
     stepNumber: number
+    handleStepsLogChunk: HandleStepsLogChunkFn
     logger: Logger
   } & ParamsExcluding<
     typeof executeToolCall,
@@ -89,6 +94,7 @@ export async function runProgrammaticStep(
     ws,
     localAgentTemplates,
     stepsComplete,
+    handleStepsLogChunk,
     logger,
   } = params
   let { stepNumber } = params
@@ -111,10 +117,9 @@ export async function runProgrammaticStep(
       (level: 'debug' | 'info' | 'warn' | 'error') =>
       (data: any, msg?: string) => {
         logger[level](data, msg) // Log to backend
-        sendAction(ws, {
-          type: 'handlesteps-log-chunk',
+        handleStepsLogChunk({
           userInputId,
-          agentId: agentState.agentId,
+          runId: agentState.runId ?? 'undefined',
           level,
           data,
           message: msg,

backend/src/websockets/middleware.ts

@@ -16,6 +16,7 @@ import { getUserInfoFromApiKey } from './auth'
 import { updateRequestContext } from './request-context'
 import { sendAction } from './websocket-action'
 import {
+  handleStepsLogChunkWs,
   requestFilesWs,
   requestMcpToolDataWs,
   requestOptionalFileWs,
@@ -154,6 +155,8 @@ export class WebSocketMiddleware {
           : undefined
 
       const scopedDeps: AgentRuntimeScopedDeps = {
+        handleStepsLogChunk: (params) =>
+          handleStepsLogChunkWs({ ...params, ws }),
         requestToolCall: (params) => requestToolCallWs({ ...params, ws }),
         requestMcpToolData: (params) => requestMcpToolDataWs({ ...params, ws }),
         requestFiles: (params) => requestFilesWs({ ...params, ws }),

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

@@ -41,6 +41,9 @@ export const TEST_AGENT_RUNTIME_IMPL = Object.freeze<AgentRuntimeDeps>({
 export const TEST_AGENT_RUNTIME_SCOPED_IMPL =
   Object.freeze<AgentRuntimeScopedDeps>({
     // Database
+    handleStepsLogChunk: () => {
+      throw new Error('handleStepsLogChunk not implemented in test runtime')
+    },
     requestToolCall: () => {
       throw new Error('requestToolCall not implemented in test runtime')
     },

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

@@ -1,4 +1,5 @@
 import type {
+  HandleStepsLogChunkFn,
   RequestFilesFn,
   RequestMcpToolDataFn,
   RequestOptionalFileFn,
@@ -38,6 +39,7 @@ export type AgentRuntimeDeps = {
 
 export type AgentRuntimeScopedDeps = {
   // Client (WebSocket)
+  handleStepsLogChunk: HandleStepsLogChunkFn
   requestToolCall: RequestToolCallFn
   requestMcpToolData: RequestMcpToolDataFn
   requestFiles: RequestFilesFn

common/src/types/contracts/client.ts

@@ -36,3 +36,11 @@ export type SendSubagentChunkFn = (params: {
   chunk: string
   prompt?: string | undefined
 }) => void
+
+export type HandleStepsLogChunkFn = (params: {
+  userInputId: string
+  runId: string
+  level: 'debug' | 'info' | 'warn' | 'error'
+  data: unknown
+  message?: string
+}) => void