From dffc1aa5fba9880b99ea2f324ea6788af3bde526 Mon Sep 17 00:00:00 2001
From: Glen Maddern <glen@glenmaddern.com>
Date: Wed, 25 Jun 2025 08:28:43 +1000
Subject: [PATCH] Make thinking block collapse immediately when tool call
 starts
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- End reasoning phase as soon as first tool-call event is detected
- Set reasoningEndTime and stop reasoning streaming immediately
- Creates linear flow: thinking expands → tool call (thinking collapses) → tool result → final response
- More aggressive/responsive UX that doesn't wait for final response to collapse thinking
---
 .../chat-ui/src/hooks/useStreamResponse.ts    | 22 +++++++++++++++++++
 examples/example-server/vengabus/src/index.ts | 12 +++++++---
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/examples/chat-ui/src/hooks/useStreamResponse.ts b/examples/chat-ui/src/hooks/useStreamResponse.ts
index 080aba3..c185b28 100644
--- a/examples/chat-ui/src/hooks/useStreamResponse.ts
+++ b/examples/chat-ui/src/hooks/useStreamResponse.ts
@@ -273,6 +273,28 @@ export const useStreamResponse = ({
                         processedEvents.add(eventKey)
 
                         debugLog(`[useStreamResponse] Tool call event:`, event)
+                        
+                        // End reasoning phase immediately when we see a tool call for linear flow
+                        if (reasoningStartTime && !reasoningEndTime) {
+                            reasoningEndTime = Date.now()
+                            debugLog(`[useStreamResponse] Reasoning ended due to tool call at:`, reasoningEndTime, 'Duration:', reasoningEndTime - reasoningStartTime, 'ms')
+                            
+                            // Update reasoning end time and stop streaming immediately
+                            setConversations((prev) => {
+                                const updated = [...prev]
+                                const conv = updated.find((c) => c.id === conversationId)
+                                if (conv && assistantMessageIndex >= 0) {
+                                    const assistantMessage = conv.messages[assistantMessageIndex]
+                                    if (assistantMessage && hasContent(assistantMessage) && assistantMessage.role === 'assistant') {
+                                        assistantMessage.reasoningEndTime = reasoningEndTime
+                                        assistantMessage.isReasoningStreaming = false
+                                        debugLog(`[useStreamResponse] Stopped reasoning streaming due to tool call`)
+                                    }
+                                }
+                                return updated
+                            })
+                        }
+                        
                         if (event.toolName) {
                             setConversations((prev) => {
                                 const updated = [...prev]
diff --git a/examples/example-server/vengabus/src/index.ts b/examples/example-server/vengabus/src/index.ts
index af12712..0dcffde 100644
--- a/examples/example-server/vengabus/src/index.ts
+++ b/examples/example-server/vengabus/src/index.ts
@@ -35,9 +35,15 @@ mcpServer.registerTool("get_vengabus_times",
     description: "This checks to see when the next Vengabus is, or whether there is one for the user's current location. Please display all the information to the user in a tabulated form.",
     inputSchema: { }
   },
-  async () => ({
-    content: [{ type: "text", text: `Next Vengabus: imminent. Current user's route: New York to San Francisco. Route name: "Intercity Disco".` }]
-  })
+  async () => {
+    await scheduler.wait(100)
+    return ({
+      content: [{
+        type: 'text',
+        text: `Next Vengabus: imminent. Current user's route: New York to San Francisco. Route name: "Intercity Disco".`
+      }]
+    })
+  }
 );
 
 app.post('/mcp', async (c) => {