Revert "fix(agent-runtime): fix duplicate text in subagent streaming view"

This reverts commit 8c09e0d.

Author: jahooma

packages/agent-runtime/src/__tests__/subagent-streaming.test.ts

@@ -21,7 +21,6 @@ import { handleSpawnAgents } from '../tools/handlers/tool/spawn-agents'
 import type { AgentTemplate } from '../templates/types'
 import type { SendSubagentChunk } from '../tools/handlers/tool/spawn-agents'
 import type { CodebuffToolCall } from '@codebuff/common/tools/list'
-import type { PrintModeEvent } from '@codebuff/common/types/print-mode'
 import type { ParamsExcluding } from '@codebuff/common/types/function-params'
 import type { Mock } from 'bun:test'
 
@@ -119,7 +118,6 @@ describe('Subagent Streaming', () => {
   beforeEach(() => {
     mockSendSubagentChunk.mockClear()
     mockLoopAgentSteps.mockClear()
-    mockWriteToClient.mockClear()
   })
 
   afterAll(() => {
@@ -173,6 +171,7 @@ describe('Subagent Streaming', () => {
       2,
       expect.objectContaining({ type: 'subagent_finish' }),
     )
+    return
   })
 
   it('should include correct agentId and agentType in streaming messages', async () => {
@@ -229,75 +228,4 @@ describe('Subagent Streaming', () => {
     expect(firstCall.userInputId).toBe('test-input')
     expect(secondCall.userInputId).toBe('test-input')
   })
-
-  it('streams subagent text chunks without duplication', async () => {
-    mockLoopAgentSteps.mockImplementationOnce(async (options) => {
-      options.onResponseChunk?.('first chunk')
-      options.onResponseChunk?.({ type: 'text', text: 'second chunk' })
-      options.onResponseChunk?.({ type: 'text', text: '' })
-
-      return {
-        agentState: {
-          ...options.agentState,
-          messageHistory: [assistantMessage('Final subagent response')],
-        },
-        output: {
-          type: 'lastMessage',
-          value: [assistantMessage('Final subagent response')],
-        },
-      }
-    })
-
-    const sessionState = getInitialSessionState(mockFileContext)
-    const agentState = sessionState.mainAgentState
-
-    const parentTemplate = {
-      id: 'base',
-      spawnableAgents: ['thinker'],
-    } as unknown as AgentTemplate
-
-    const toolCall: CodebuffToolCall<'spawn_agents'> = {
-      toolName: 'spawn_agents' as const,
-      toolCallId: 'test-tool-call-id-3',
-      input: {
-        agents: [
-          {
-            agent_type: 'thinker',
-            prompt: 'Check for duplicates',
-          },
-        ],
-      },
-    }
-
-    await handleSpawnAgents({
-      ...handleSpawnAgentsBaseParams,
-      agentState,
-      agentTemplate: parentTemplate,
-      localAgentTemplates: {
-        [mockAgentTemplate.id]: mockAgentTemplate,
-      },
-      toolCall,
-    })
-
-    const streamedChunks = (
-      mockSendSubagentChunk.mock.calls as Array<
-        [Parameters<SendSubagentChunk>[0]]
-      >
-    ).map(([call]) => call.chunk)
-
-    expect(streamedChunks).toEqual(['first chunk', 'second chunk'])
-
-    const subagentLifecycleEvents = (
-      mockWriteToClient.mock.calls as Array<[string | PrintModeEvent]>
-    )
-      .map(([event]) => event)
-      .filter(
-        (event) =>
-          typeof event === 'object' &&
-          (event?.type === 'subagent_start' ||
-            event?.type === 'subagent_finish'),
-      )
-
-    expect(subagentLifecycleEvents).toHaveLength(2)
-  })
 })

packages/agent-runtime/src/__tests__/tool-stream-parser.test.ts

@@ -51,6 +51,11 @@ describe('processStreamWithTags', () => {
     }
 
     const result: string[] = []
+    const responseChunks: any[] = []
+
+    function onResponseChunk(chunk: any) {
+      responseChunks.push(chunk)
+    }
 
     function defaultProcessor(toolName: string) {
       return {
@@ -65,6 +70,7 @@ describe('processStreamWithTags', () => {
       processors,
       defaultProcessor,
       onError,
+      onResponseChunk,
     })) {
       if (chunk.type === 'text') {
         result.push(chunk.text)
@@ -111,6 +117,11 @@ describe('processStreamWithTags', () => {
     }
 
     const result: string[] = []
+    const responseChunks: any[] = []
+
+    function onResponseChunk(chunk: any) {
+      responseChunks.push(chunk)
+    }
 
     function defaultProcessor(toolName: string) {
       return {
@@ -125,6 +136,7 @@ describe('processStreamWithTags', () => {
       processors,
       defaultProcessor,
       onError,
+      onResponseChunk,
     })) {
       if (chunk.type === 'text') {
         result.push(chunk.text)
@@ -181,6 +193,11 @@ describe('processStreamWithTags', () => {
     }
 
     const result: string[] = []
+    const responseChunks: any[] = []
+
+    function onResponseChunk(chunk: any) {
+      responseChunks.push(chunk)
+    }
 
     function defaultProcessor(toolName: string) {
       return {
@@ -195,6 +212,7 @@ describe('processStreamWithTags', () => {
       processors,
       defaultProcessor,
       onError,
+      onResponseChunk,
     })) {
       if (chunk.type === 'text') {
         result.push(chunk.text)
@@ -249,6 +267,12 @@ describe('processStreamWithTags', () => {
       events.push({ name, error, type: 'error' })
     }
 
+    const responseChunks: any[] = []
+
+    function onResponseChunk(chunk: any) {
+      responseChunks.push(chunk)
+    }
+
     function defaultProcessor(toolName: string) {
       // For unknown tools, still return a processor but track the error
       events.push({
@@ -268,6 +292,7 @@ describe('processStreamWithTags', () => {
       processors,
       defaultProcessor,
       onError,
+      onResponseChunk,
     })) {
       // consume stream
     }
@@ -316,6 +341,11 @@ describe('processStreamWithTags', () => {
     }
 
     const result: string[] = []
+    const responseChunks: any[] = []
+
+    function onResponseChunk(chunk: any) {
+      responseChunks.push(chunk)
+    }
 
     function defaultProcessor(toolName: string) {
       return {
@@ -330,6 +360,7 @@ describe('processStreamWithTags', () => {
       processors,
       defaultProcessor,
       onError,
+      onResponseChunk,
     })) {
       if (chunk.type === 'text') {
         result.push(chunk.text)
@@ -382,6 +413,11 @@ describe('processStreamWithTags', () => {
     }
 
     const result: string[] = []
+    const responseChunks: any[] = []
+
+    function onResponseChunk(chunk: any) {
+      responseChunks.push(chunk)
+    }
 
     function defaultProcessor(toolName: string) {
       return {
@@ -396,6 +432,7 @@ describe('processStreamWithTags', () => {
       processors,
       defaultProcessor,
       onError,
+      onResponseChunk,
     })) {
       if (chunk.type === 'text') {
         result.push(chunk.text)
@@ -429,6 +466,11 @@ describe('processStreamWithTags', () => {
     }
 
     const result: string[] = []
+    const responseChunks: any[] = []
+
+    function onResponseChunk(chunk: any) {
+      responseChunks.push(chunk)
+    }
 
     function defaultProcessor(toolName: string) {
       return {
@@ -443,6 +485,7 @@ describe('processStreamWithTags', () => {
       processors,
       defaultProcessor,
       onError,
+      onResponseChunk,
     })) {
       if (chunk.type === 'text') {
         result.push(chunk.text)
@@ -469,6 +512,11 @@ describe('processStreamWithTags', () => {
     }
 
     const result: string[] = []
+    const responseChunks: any[] = []
+
+    function onResponseChunk(chunk: any) {
+      responseChunks.push(chunk)
+    }
 
     function defaultProcessor(toolName: string) {
       return {
@@ -483,6 +531,7 @@ describe('processStreamWithTags', () => {
       processors,
       defaultProcessor,
       onError,
+      onResponseChunk,
     })) {
       if (chunk.type === 'text') {
         result.push(chunk.text)

packages/agent-runtime/src/tool-stream-parser.ts

@@ -4,6 +4,11 @@ import type { Model } from '@codebuff/common/old-constants'
 import type { TrackEventFn } from '@codebuff/common/types/contracts/analytics'
 import type { StreamChunk } from '@codebuff/common/types/contracts/llm'
 import type { Logger } from '@codebuff/common/types/contracts/logger'
+import type {
+  PrintModeError,
+  PrintModeText,
+} from '@codebuff/common/types/print-mode'
+
 export async function* processStreamWithTools(params: {
   stream: AsyncGenerator<StreamChunk, string | null>
   processors: Record<
@@ -18,6 +23,7 @@ export async function* processStreamWithTools(params: {
     onTagEnd: (tagName: string, params: Record<string, any>) => void
   }
   onError: (tagName: string, errorMessage: string) => void
+  onResponseChunk: (chunk: PrintModeText | PrintModeError) => void
   logger: Logger
   loggerOptions?: {
     userId?: string
@@ -31,11 +37,13 @@ export async function* processStreamWithTools(params: {
     processors,
     defaultProcessor,
     onError,
+    onResponseChunk,
     logger,
     loggerOptions,
     trackEvent,
   } = params
   let streamCompleted = false
+  let buffer = ''
   let autocompleted = false
 
   function processToolCallObject(params: {
@@ -65,14 +73,31 @@ export async function* processStreamWithTools(params: {
     processor.onTagEnd(toolName, input)
   }
 
+  function flush() {
+    if (buffer) {
+      onResponseChunk({
+        type: 'text',
+        text: buffer,
+      })
+    }
+    buffer = ''
+  }
+
   function* processChunk(
     chunk: StreamChunk | undefined,
   ): Generator<StreamChunk> {
     if (chunk === undefined) {
+      flush()
       streamCompleted = true
       return
     }
 
+    if (chunk.type === 'text') {
+      buffer += chunk.text
+    } else {
+      flush()
+    }
+
     if (chunk.type === 'tool-call') {
       processToolCallObject(chunk)
     }

packages/agent-runtime/src/tools/stream-parser.ts

@@ -234,6 +234,21 @@ export async function processStream(
       model: agentTemplate.model,
       agentName: agentTemplate.id,
     },
+    onResponseChunk: (chunk) => {
+      if (chunk.type === 'text') {
+        if (chunk.text) {
+          assistantMessages.push(assistantMessage(chunk.text))
+        }
+      } else if (chunk.type === 'error') {
+        // do nothing
+      } else {
+        chunk satisfies never
+        throw new Error(
+          `Internal error: unhandled chunk type: ${(chunk as any).type}`,
+        )
+      }
+      return onResponseChunk(chunk)
+    },
   })
 
   let messageId: string | null = null
@@ -258,7 +273,6 @@ export async function processStream(
     } else if (chunk.type === 'text') {
       onResponseChunk(chunk.text)
       fullResponseChunks.push(chunk.text)
-      assistantMessages.push(assistantMessage(chunk.text))
     } else if (chunk.type === 'error') {
       onResponseChunk(chunk)