From 4a291e7901d28ed1412cbcb71ad783c157aec54a Mon Sep 17 00:00:00 2001
From: Mark <markrachapoom@gmail.com>
Date: Fri, 13 Feb 2026 18:40:46 -0800
Subject: [PATCH] Thinking and fix the loader after chat completed, and
 minimize after streaming message.

---
 apps/web/app/components/chain-of-thought.tsx  | 14 +++++++-----
 apps/web/app/components/chat-message.tsx      | 16 +++++++++++---
 apps/web/app/components/chat-panel.tsx        |  3 ++-
 ...pi-embedded-subscribe.handlers.messages.ts | 22 +++++++++++++++++++
 4 files changed, 46 insertions(+), 9 deletions(-)

diff --git a/apps/web/app/components/chain-of-thought.tsx b/apps/web/app/components/chain-of-thought.tsx
index affb4842681..fb53e199fed 100644
--- a/apps/web/app/components/chain-of-thought.tsx
+++ b/apps/web/app/components/chain-of-thought.tsx
@@ -516,9 +516,8 @@ function groupToolSteps(tools: ToolPart[]): VisualItem[] {
 
 /* ─── Main component ─── */
 
-export function ChainOfThought({ parts }: { parts: ChainPart[] }) {
+export function ChainOfThought({ parts, isStreaming }: { parts: ChainPart[]; isStreaming?: boolean }) {
 	const [isOpen, setIsOpen] = useState(true);
-	const prevActiveRef = useRef(true);
 
 	const isActive = parts.some(
 		(p) =>
@@ -556,12 +555,17 @@ export function ChainOfThought({ parts }: { parts: ChainPart[] }) {
 		return rem > 0 ? `${m}m ${rem}s` : `${m}m`;
 	}, []);
 
+	// Collapse only when the parent stream truly ends — not on intermediate
+	// isActive flickers (e.g. gap between reasoning end and tool start).
+	const wasStreamingRef = useRef(false);
 	useEffect(() => {
-		if (prevActiveRef.current && !isActive && parts.length > 0) {
+		if (isStreaming) {
+			wasStreamingRef.current = true;
+		} else if (wasStreamingRef.current && parts.length > 0) {
+			wasStreamingRef.current = false;
 			setIsOpen(false);
 		}
-		prevActiveRef.current = isActive;
-	}, [isActive, parts.length]);
+	}, [isStreaming, parts.length]);
 
 	const statusParts = parts.filter(
 		(p): p is Extract<ChainPart, { kind: "status" }> =>
diff --git a/apps/web/app/components/chat-message.tsx b/apps/web/app/components/chat-message.tsx
index 114482ff216..48203a0554a 100644
--- a/apps/web/app/components/chat-message.tsx
+++ b/apps/web/app/components/chat-message.tsx
@@ -56,8 +56,17 @@ function groupParts(parts: UIMessage["parts"]): MessageSegment[] {
 	const segments: MessageSegment[] = [];
 	let chain: ChainPart[] = [];
 
-	const flush = () => {
+	const flush = (textFollows?: boolean) => {
 		if (chain.length > 0) {
+			// If text content follows this chain, all tools must have
+			// completed — force any stuck "running" tools to "done".
+			if (textFollows) {
+				for (const cp of chain) {
+					if (cp.kind === "tool" && cp.status === "running") {
+						cp.status = "done";
+					}
+				}
+			}
 			segments.push({ type: "chain", parts: [...chain] });
 			chain = [];
 		}
@@ -65,7 +74,7 @@ function groupParts(parts: UIMessage["parts"]): MessageSegment[] {
 
 	for (const part of parts) {
 		if (part.type === "text") {
-			flush();
+			flush(true);
 			const text = (part as { type: "text"; text: string }).text;
 			if (hasReportBlocks(text)) {
 				segments.push(
@@ -504,7 +513,7 @@ const mdComponents: Components = {
 
 /* ─── Chat message ─── */
 
-export function ChatMessage({ message }: { message: UIMessage }) {
+export function ChatMessage({ message, isStreaming }: { message: UIMessage; isStreaming?: boolean }) {
 	const isUser = message.role === "user";
 	const segments = groupParts(message.parts);
 
@@ -649,6 +658,7 @@ export function ChatMessage({ message }: { message: UIMessage }) {
 					<ChainOfThought
 						key={index}
 						parts={segment.parts}
+						isStreaming={isStreaming}
 					/>
 				);
 			})}
diff --git a/apps/web/app/components/chat-panel.tsx b/apps/web/app/components/chat-panel.tsx
index 9ed7263c6c9..584757f770e 100644
--- a/apps/web/app/components/chat-panel.tsx
+++ b/apps/web/app/components/chat-panel.tsx
@@ -1315,10 +1315,11 @@ export const ChatPanel = forwardRef<ChatPanelHandle, ChatPanelProps>(
 						<div
 							className={`${compact ? "" : "max-w-3xl mx-auto"} py-3`}
 						>
-							{messages.map((message) => (
+							{messages.map((message, i) => (
 								<ChatMessage
 									key={message.id}
 									message={message}
+									isStreaming={isStreaming && i === messages.length - 1}
 								/>
 							))}
 							<div ref={messagesEndRef} />
diff --git a/src/agents/pi-embedded-subscribe.handlers.messages.ts b/src/agents/pi-embedded-subscribe.handlers.messages.ts
index f7d140af843..95267f84103 100644
--- a/src/agents/pi-embedded-subscribe.handlers.messages.ts
+++ b/src/agents/pi-embedded-subscribe.handlers.messages.ts
@@ -64,6 +64,28 @@ export function handleMessageUpdate(
       : undefined;
   const evtType = typeof assistantRecord?.type === "string" ? assistantRecord.type : "";
 
+  // Handle native extended thinking events (Anthropic API thinking blocks).
+  // These arrive as thinking_delta / thinking_start / thinking_end from the
+  // provider adapter and must be forwarded to the agent event bus so that
+  // active-runs.ts (web UI) and other consumers receive them.
+  if (evtType === "thinking_delta" || evtType === "thinking_start" || evtType === "thinking_end") {
+    if (evtType === "thinking_delta") {
+      const thinkingDelta = typeof assistantRecord?.delta === "string" ? assistantRecord.delta : "";
+      if (thinkingDelta) {
+        emitAgentEvent({
+          runId: ctx.params.runId,
+          stream: "thinking",
+          data: { delta: thinkingDelta },
+        });
+        void ctx.params.onAgentEvent?.({
+          stream: "thinking",
+          data: { delta: thinkingDelta },
+        });
+      }
+    }
+    return;
+  }
+
   if (evtType !== "text_delta" && evtType !== "text_start" && evtType !== "text_end") {
     return;
   }