Print mode JSON streaming (#223)

Author: charleslien

backend/src/__tests__/main-prompt.integration.test.ts

@@ -14,6 +14,7 @@ import { mainPrompt } from '../main-prompt'
 
 // Mock imports needed for setup within the test
 import { getToolCallString } from '@codebuff/common/constants/tools'
+import { PrintModeObject } from '@codebuff/common/types/print-mode'
 import { ProjectFileContext } from '@codebuff/common/util/file'
 import * as checkTerminalCommandModule from '../check-terminal-command'
 import * as requestFilesPrompt from '../find-files/request-files-prompt'
@@ -363,7 +364,10 @@ export function getMessagesSubset(messages: Message[], otherTokens: number) {
     } = await mainPrompt(new MockWebSocket() as unknown as WebSocket, action, {
       userId: TEST_USER_ID,
       clientSessionId: 'test-session-delete-function-integration',
-      onResponseChunk: (chunk: string) => {
+      onResponseChunk: (chunk: string | PrintModeObject) => {
+        if (typeof chunk !== 'string') {
+          return
+        }
         process.stdout.write(chunk)
       },
     })
@@ -444,7 +448,10 @@ export function getMessagesSubset(messages: Message[], otherTokens: number) {
       await mainPrompt(new MockWebSocket() as unknown as WebSocket, action, {
         userId: TEST_USER_ID,
         clientSessionId: 'test-session-delete-function-integration',
-        onResponseChunk: (chunk: string) => {
+        onResponseChunk: (chunk: string | PrintModeObject) => {
+          if (typeof chunk !== 'string') {
+            return
+          }
           process.stdout.write(chunk)
         },
       })

backend/src/main-prompt.ts

@@ -9,6 +9,7 @@ import {
 import { WebSocket } from 'ws'
 
 import { renderToolResults } from '@codebuff/common/constants/tools'
+import { PrintModeObject } from '@codebuff/common/types/print-mode'
 import { checkTerminalCommand } from './check-terminal-command'
 import { loopAgentSteps } from './run-agent-step'
 import { getAllAgentTemplates } from './templates/agent-registry'
@@ -20,7 +21,7 @@ import { requestToolCall } from './websockets/websocket-action'
 export interface MainPromptOptions {
   userId: string | undefined
   clientSessionId: string
-  onResponseChunk: (chunk: string) => void
+  onResponseChunk: (chunk: string | PrintModeObject) => void
 }
 
 export const mainPrompt = async (

backend/src/run-agent-step.ts

@@ -10,6 +10,7 @@ import {
   renderToolResults,
 } from '@codebuff/common/constants/tools'
 import { CodebuffMessage } from '@codebuff/common/types/message'
+import { PrintModeObject } from '@codebuff/common/types/print-mode'
 import {
   AgentState,
   ToolResult,
@@ -50,7 +51,7 @@ export interface AgentOptions {
   userInputId: string
   clientSessionId: string
   fingerprintId: string
-  onResponseChunk: (chunk: string) => void
+  onResponseChunk: (chunk: string | PrintModeObject) => void
 
   agentType: AgentTemplateType
   fileContext: ProjectFileContext
@@ -480,7 +481,20 @@ export const runAgentStep = async (
 
 export const loopAgentSteps = async (
   ws: WebSocket,
-  options: {
+  {
+    userInputId,
+    agentType,
+    agentState,
+    prompt,
+    params,
+    fingerprintId,
+    fileContext,
+    toolResults,
+    agentRegistry,
+    userId,
+    clientSessionId,
+    onResponseChunk,
+  }: {
     userInputId: string
     agentType: AgentTemplateType
     agentState: AgentState
@@ -493,22 +507,9 @@ export const loopAgentSteps = async (
 
     userId: string | undefined
     clientSessionId: string
-    onResponseChunk: (chunk: string) => void
+    onResponseChunk: (chunk: string | PrintModeObject) => void
   }
 ) => {
-  const {
-    agentState,
-    prompt,
-    params,
-    userId,
-    clientSessionId,
-    agentRegistry,
-    onResponseChunk,
-    userInputId,
-    fingerprintId,
-    fileContext,
-    agentType,
-  } = options
   const agentTemplate = agentRegistry[agentType]
   if (!agentTemplate) {
     throw new Error(`Agent template not found for type: ${agentType}`)

backend/src/run-programmatic-step.ts

@@ -1,3 +1,4 @@
+import { PrintModeObject } from '@codebuff/common/types/print-mode'
 import {
   AgentState,
   AgentTemplateType,
@@ -37,40 +38,44 @@ export function clearAgentGeneratorCache() {
 // Function to handle programmatic agents
 export async function runProgrammaticStep(
   agentState: AgentState,
-  params: {
+  {
+    template,
+    prompt,
+    params,
+    userId,
+    userInputId,
+    clientSessionId,
+    fingerprintId,
+    onResponseChunk,
+    agentType,
+    fileContext,
+    assistantMessage,
+    assistantPrefix,
+    ws,
+  }: {
     template: AgentTemplate
     prompt: string | undefined
     params: Record<string, any> | undefined
     userId: string | undefined
     userInputId: string
     clientSessionId: string
     fingerprintId: string
-    onResponseChunk: (chunk: string) => void
+    onResponseChunk: (chunk: string | PrintModeObject) => void
     agentType: AgentTemplateType
     fileContext: ProjectFileContext
     ws: WebSocket
   }
 ): Promise<{ agentState: AgentState; endTurn: boolean }> {
-  const {
-    template,
-    onResponseChunk,
-    ws,
-    userId,
-    userInputId,
-    clientSessionId,
-    fingerprintId,
-    fileContext,
-  } = params
   if (!template.handleSteps) {
     throw new Error('No step handler found for agent template ' + template.id)
   }
 
   logger.info(
     {
       template: template.id,
-      agentType: params.agentType,
-      prompt: params.prompt,
-      params: params.params,
+      agentType,
+      prompt,
+      params,
     },
     'Running programmatic step'
   )
@@ -88,16 +93,16 @@ export async function runProgrammaticStep(
         template.handleSteps,
         {
           agentState,
-          prompt: params.prompt,
-          params: params.params,
+          prompt,
+          params,
         }
       )
     } else {
       // Initialize native generator
       generator = template.handleSteps({
         agentState,
-        prompt: params.prompt,
-        params: params.params,
+        prompt,
+        params,
       })
       agentIdToGenerator[agentState.agentId] = generator
     }

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

@@ -12,6 +12,7 @@ import type { CodebuffToolHandlerFunction } from '../handler-function-type'
 import type { SendSubagentChunk } from './spawn-agents'
 
 import { ASYNC_AGENTS_ENABLED } from '@codebuff/common/constants'
+import { PrintModeObject } from '@codebuff/common/types/print-mode'
 import { generateCompactId } from '@codebuff/common/util/string'
 import { asyncAgentManager } from '../../../async-agent-manager'
 import { getAllAgentTemplates } from '../../../templates/agent-registry'
@@ -197,14 +198,18 @@ export const handleSpawnAgentsAsync = ((params: {
               toolResults: [],
               userId,
               clientSessionId,
-              onResponseChunk: (chunk: string) =>
+              onResponseChunk: (chunk: string | PrintModeObject) => {
+                if (typeof chunk !== 'string') {
+                  return
+                }
                 sendSubagentChunk({
                   userInputId,
                   agentId,
                   agentType,
                   chunk,
                   prompt,
-                }),
+                })
+              },
             })
 
             // Send completion message to parent if agent has appropriate output mode

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

@@ -10,6 +10,7 @@ import type { AgentTemplate } from '../../../templates/types'
 import type { CodebuffToolCall } from '../../constants'
 import type { CodebuffToolHandlerFunction } from '../handler-function-type'
 
+import { PrintModeObject } from '@codebuff/common/types/print-mode'
 import { generateCompactId } from '@codebuff/common/util/string'
 import { getAllAgentTemplates } from '../../../templates/agent-registry'
 import { logger } from '../../../util/logger'
@@ -193,7 +194,10 @@ export const handleSpawnAgents = ((params: {
           toolResults: [],
           userId,
           clientSessionId,
-          onResponseChunk: (chunk: string) => {
+          onResponseChunk: (chunk: string | PrintModeObject) => {
+            if (typeof chunk !== 'string') {
+              return
+            }
             // Send subagent streaming chunks to client
             sendSubagentChunk({
               userInputId,

backend/src/tools/stream-parser.ts

@@ -1,20 +1,23 @@
-import { ToolName, toolNames } from '@codebuff/common/constants/tools'
-import { CodebuffMessage } from '@codebuff/common/types/message'
-import {
+import type { ToolName } from '@codebuff/common/constants/tools'
+import type { CodebuffMessage } from '@codebuff/common/types/message'
+import type { PrintModeObject } from '@codebuff/common/types/print-mode'
+import type {
   AgentState,
   Subgoal,
   ToolResult,
 } from '@codebuff/common/types/session-state'
+import type { ProjectFileContext } from '@codebuff/common/util/file'
+import type { ToolCallPart } from 'ai'
+import type { WebSocket } from 'ws'
+import type { AgentTemplate } from '../templates/types'
+import type { CodebuffToolCall } from './constants'
+
+import { toolNames } from '@codebuff/common/constants/tools'
 import { buildArray } from '@codebuff/common/util/array'
-import { ProjectFileContext } from '@codebuff/common/util/file'
 import { generateCompactId } from '@codebuff/common/util/string'
-import { ToolCallPart } from 'ai'
-import { WebSocket } from 'ws'
-import { AgentTemplate } from '../templates/types'
 import { expireMessages } from '../util/messages'
 import { sendAction } from '../websockets/websocket-action'
 import { processStreamWithTags } from '../xml-stream-parser'
-import { CodebuffToolCall } from './constants'
 import { executeToolCall } from './tool-executor'
 
 export type ToolCallError = {
@@ -37,7 +40,7 @@ export async function processStreamWithTools<T extends string>(options: {
   messages: CodebuffMessage[]
   agentState: AgentState
   agentContext: Record<string, Subgoal>
-  onResponseChunk: (chunk: string) => void
+  onResponseChunk: (chunk: string | PrintModeObject) => void
   fullResponse: string
 }) {
   const {
@@ -126,6 +129,7 @@ export async function processStreamWithTools<T extends string>(options: {
         result: error,
       })
     },
+    onResponseChunk,
     {
       userId,
       model: agentTemplate.model,

backend/src/tools/tool-executor.ts

@@ -1,4 +1,5 @@
 import type { ToolName } from '@codebuff/common/constants/tools'
+import type { PrintModeObject } from '@codebuff/common/types/print-mode'
 import type { ToolResult } from '@codebuff/common/types/session-state'
 import type { ProjectFileContext } from '@codebuff/common/util/file'
 import type { WebSocket } from 'ws'
@@ -105,34 +106,30 @@ export interface ExecuteToolCallParams<T extends ToolName = ToolName> {
   clientSessionId: string
   userInputId: string
   fullResponse: string
-  onResponseChunk: (chunk: string) => void
+  onResponseChunk: (chunk: string | PrintModeObject) => void
   state: Record<string, any>
   userId: string | undefined
   autoInsertEndStepParam?: boolean
 }
 
-export function executeToolCall<T extends ToolName>(
-  options: ExecuteToolCallParams<T>
-): Promise<void> {
-  const {
-    toolName,
-    args,
-    toolCalls,
-    toolResults,
-    previousToolCallFinished,
-    ws,
-    agentTemplate,
-    fileContext,
-    agentStepId,
-    clientSessionId,
-    userInputId,
-    fullResponse,
-    onResponseChunk,
-    state,
-    userId,
-    autoInsertEndStepParam = false,
-  } = options
-
+export function executeToolCall<T extends ToolName>({
+  toolName,
+  args,
+  toolCalls,
+  toolResults,
+  previousToolCallFinished,
+  ws,
+  agentTemplate,
+  fileContext,
+  agentStepId,
+  clientSessionId,
+  userInputId,
+  fullResponse,
+  onResponseChunk,
+  state,
+  userId,
+  autoInsertEndStepParam = false,
+}: ExecuteToolCallParams<T>): Promise<void> {
   const toolCall: CodebuffToolCall<T> | ToolCallError = parseRawToolCall<T>(
     {
       toolName,
@@ -154,6 +151,13 @@ export function executeToolCall<T extends ToolName>(
     return previousToolCallFinished
   }
 
+  onResponseChunk({
+    type: 'tool_call',
+    toolCallId: toolCall.toolCallId,
+    toolName,
+    args: toolCall.args,
+  })
+
   logger.debug(
     { toolCall },
     `${toolName} (${toolCall.toolCallId}) tool call detected in stream`
@@ -226,6 +230,12 @@ export function executeToolCall<T extends ToolName>(
       return
     }
 
+    onResponseChunk({
+      type: 'tool_result',
+      toolCallId: toolResult.toolCallId,
+      result: toolResult.result,
+    })
+
     toolResults.push(toolResult)
 
     state.messages.push({

backend/src/websockets/websocket-action.ts

@@ -181,6 +181,11 @@ const onPrompt = async (
           }
         )
 
+        sendAction(ws, {
+          type: 'response-chunk',
+          userInputId: promptId,
+          chunk: { type: 'error', message: response },
+        })
         sendAction(ws, {
           type: 'response-chunk',
           userInputId: promptId,