switch cli to use reasoning stream

Author: charleslien

cli/src/hooks/use-send-message.ts

@@ -557,8 +557,12 @@ export const useSendMessage = ({
         )
       }
 
-      const appendRootTextChunk = (delta: string) => {
-        if (!delta) {
+      const appendRootChunk = (
+        delta:
+          | { type: 'text'; text: string }
+          | { type: 'reasoning'; text: string },
+      ) => {
+        if (!delta.text) {
           return
         }
 
@@ -571,10 +575,14 @@ export const useSendMessage = ({
             const blocks: ContentBlock[] = msg.blocks ? [...msg.blocks] : []
             const lastBlock = blocks[blocks.length - 1]
 
-            if (lastBlock && lastBlock.type === 'text') {
+            if (
+              lastBlock &&
+              lastBlock.type === 'text' &&
+              delta.type === lastBlock.textType
+            ) {
               const updatedBlock: ContentBlock = {
                 ...lastBlock,
-                content: lastBlock.content + delta,
+                content: lastBlock.content + delta.text,
               }
               return {
                 ...msg,
@@ -584,7 +592,15 @@ export const useSendMessage = ({
 
             return {
               ...msg,
-              blocks: [...blocks, { type: 'text', content: delta }],
+              blocks: [
+                ...blocks,
+                {
+                  type: 'text',
+                  content: delta.text,
+                  textType: delta.type,
+                  ...(delta.type === 'reasoning' && { color: 'grey' }),
+                },
+              ],
             }
           }),
         )
@@ -627,7 +643,30 @@ export const useSendMessage = ({
           maxAgentSteps: 40,
 
           handleStreamChunk: (event) => {
-            if (typeof event !== 'string') {
+            logger.info({ event }, 'asdf event')
+            if (typeof event === 'string' || event.type === 'reasoning_chunk') {
+              const eventObj:
+                | { type: 'text'; text: string }
+                | { type: 'reasoning'; text: string } =
+                typeof event === 'string'
+                  ? { type: 'text', text: event }
+                  : { type: 'reasoning', text: event.chunk }
+              if (!hasReceivedContent) {
+                hasReceivedContent = true
+                setIsWaitingForResponse(false)
+              }
+
+              if (!eventObj.text) {
+                return
+              }
+
+              if (eventObj.type === 'text') {
+                rootStreamBufferRef.current =
+                  (rootStreamBufferRef.current ?? '') + eventObj.text
+              }
+              rootStreamSeenRef.current = true
+              appendRootChunk(eventObj)
+            } else if (event.type === 'subagent_chunk') {
               const { agentId, chunk } = event
 
               const previous =
@@ -642,21 +681,9 @@ export const useSendMessage = ({
                 content: chunk,
               })
               return
+            } else {
+              event satisfies never
             }
-
-            if (!hasReceivedContent) {
-              hasReceivedContent = true
-              setIsWaitingForResponse(false)
-            }
-
-            const previous = rootStreamBufferRef.current ?? ''
-            if (!event) {
-              return
-            }
-
-            rootStreamBufferRef.current = previous + event
-            rootStreamSeenRef.current = true
-            appendRootTextChunk(event)
           },
 
           handleEvent: (event: any) => {
@@ -727,7 +754,7 @@ export const useSendMessage = ({
                 )
                 rootStreamBufferRef.current = previous + text
 
-                appendRootTextChunk(text)
+                appendRootChunk({ type: 'text', text })
               }
               return
             }

cli/src/types/chat.ts

@@ -12,6 +12,7 @@ export type ContentBlock =
       marginTop?: number
       marginBottom?: number
       status?: 'running' | 'complete'
+      textType?: 'reasoning' | 'text'
     }
   | {
       type: 'html'

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

@@ -1,9 +1,4 @@
-import {
-  endToolTag,
-  startToolTag,
-  toolNameParam,
-  toolNames,
-} from '@codebuff/common/tools/constants'
+import { toolNames } from '@codebuff/common/tools/constants'
 import { buildArray } from '@codebuff/common/util/array'
 import { generateCompactId } from '@codebuff/common/util/string'
 import { cloneDeep } from 'lodash'
@@ -181,7 +176,6 @@ export async function processStreamWithTools(
     },
   })
 
-  let reasoning = false
   let messageId: string | null = null
   while (true) {
     const { value: chunk, done } = await streamWithTags.next()
@@ -191,18 +185,8 @@ export async function processStreamWithTools(
     }
 
     if (chunk.type === 'reasoning') {
-      if (!reasoning) {
-        reasoning = true
-        onResponseChunk(`\n\n${startToolTag}{
-  ${JSON.stringify(toolNameParam)}: "think_deeply",
-  "thought": "`)
-      }
-      onResponseChunk(JSON.stringify(chunk.text).slice(1, -1))
+      onResponseChunk(chunk)
     } else if (chunk.type === 'text') {
-      if (reasoning) {
-        reasoning = false
-        onResponseChunk(`"\n}${endToolTag}\n\n`)
-      }
       onResponseChunk(chunk.text)
       fullResponseChunks.push(chunk.text)
     } else if (chunk.type === 'error') {

sdk/src/run.ts

@@ -64,6 +64,10 @@ export type CodebuffClientOptions = {
           agentId: string
           agentType: string
           chunk: string
+        }
+      | {
+          type: 'reasoning_chunk'
+          chunk: string
         },
   ) => void | Promise<void>
 
@@ -151,13 +155,31 @@ export async function run({
 
   const buffers: Record<string | 0, string> = { 0: '' }
 
+  let reasoning = ''
+  async function flushReasoning() {
+    if (reasoning) {
+      await handleEvent?.({
+        type: 'reasoning',
+        text: reasoning,
+      })
+    }
+    reasoning = ''
+  }
   const onResponseChunk = async (
     action: ServerAction<'response-chunk'>,
   ): Promise<void> => {
     checkAborted(signal)
     const { chunk } = action
     if (typeof chunk !== 'string') {
-      await handleEvent?.(chunk)
+      if (chunk.type === 'reasoning') {
+        handleStreamChunk?.({
+          type: 'reasoning_chunk',
+          chunk: chunk.text,
+        })
+      } else {
+        await flushReasoning()
+        await handleEvent?.(chunk)
+      }
       return
     }
 
@@ -173,7 +195,10 @@ export async function run({
           break
         }
 
-        await handleStreamChunk(value.chunk)
+        if (value.chunk) {
+          await flushReasoning()
+          await handleStreamChunk(value.chunk)
+        }
       }
     }
   }