From 812ee4bfec2a2b6fdf0a3442d3ea21d19ac8be70 Mon Sep 17 00:00:00 2001
From: mawkone <mark@getacquired.com>
Date: Mon, 15 Jun 2026 12:23:45 -0700
Subject: [PATCH] Polish thinking streaming UI and ordering

---
 vibn-frontend/app/api/chat/route.ts           |  35 ++++--
 .../components/vibn-chat/chat-panel.tsx       | 102 +++++++++++++++---
 2 files changed, 112 insertions(+), 25 deletions(-)

diff --git a/vibn-frontend/app/api/chat/route.ts b/vibn-frontend/app/api/chat/route.ts
index 4b180d89..9af6133e 100644
--- a/vibn-frontend/app/api/chat/route.ts
+++ b/vibn-frontend/app/api/chat/route.ts
@@ -1182,13 +1182,39 @@ export async function POST(request: Request) {
             error: undefined as string | undefined,
           };
 
+          let currentTimelineKind: "thought" | "text" | null = null;
+          let currentTimelineText = "";
+
+          const flushTimeline = () => {
+            if (!currentTimelineKind) return;
+            if (currentTimelineKind === "thought") {
+              assistantTimeline.push({ kind: "thought", text: currentTimelineText });
+            } else if (currentTimelineKind === "text") {
+              assistantText += (assistantText ? "\n\n" : "") + currentTimelineText;
+              assistantTextSegments.push(currentTimelineText);
+              assistantTimeline.push({ kind: "text", text: currentTimelineText });
+            }
+            currentTimelineKind = null;
+            currentTimelineText = "";
+          };
+
           for await (const chunk of stream) {
             if (aborted) break;
             
             if (chunk.type === "thinking_delta" && chunk.text) {
+              if (currentTimelineKind !== "thought") {
+                flushTimeline();
+                currentTimelineKind = "thought";
+              }
+              currentTimelineText += chunk.text;
               resp.thoughts += chunk.text;
               emit({ type: "thinking_delta", text: chunk.text });
             } else if (chunk.type === "text_delta" && chunk.text) {
+              if (currentTimelineKind !== "text") {
+                flushTimeline();
+                currentTimelineKind = "text";
+              }
+              currentTimelineText += chunk.text;
               resp.text += chunk.text;
               emit({ type: "text_delta", text: chunk.text });
             } else if (chunk.type === "tool_calls" && chunk.toolCalls) {
@@ -1197,16 +1223,9 @@ export async function POST(request: Request) {
               resp.error = chunk.error;
             }
           }
+          flushTimeline();
 
-          // If the model produced any thoughts or text, record them in the timeline once stream is complete.
-          // (The UI handles the delta-rendering live, but we save the complete chunk to Postgres).
-          if (resp.thoughts) {
-            assistantTimeline.push({ kind: "thought", text: resp.thoughts });
-          }
           if (resp.text) {
-            assistantText += (assistantText ? "\n\n" : "") + resp.text;
-            assistantTextSegments.push(resp.text);
-            assistantTimeline.push({ kind: "text", text: resp.text });
             roundsSinceText = 0;
             toolCallsSinceText = 0;
           } else if (resp.toolCalls.length) {
diff --git a/vibn-frontend/components/vibn-chat/chat-panel.tsx b/vibn-frontend/components/vibn-chat/chat-panel.tsx
index 5477e15d..6c9d4840 100644
--- a/vibn-frontend/components/vibn-chat/chat-panel.tsx
+++ b/vibn-frontend/components/vibn-chat/chat-panel.tsx
@@ -635,23 +635,7 @@ function Timeline({ entries }: { entries: TimelineEntry[] }) {
     <div style={{ marginBottom: 6 }}>
       {items.map((item, i) => {
         if (item.kind === "thought") {
-          return (
-            <div
-              key={i}
-              style={{
-                fontSize: "0.85rem",
-                color: "#6b7280",
-                fontStyle: "italic",
-                padding: "8px 12px",
-                background: "#f9fafb",
-                borderLeft: "2px solid #d1d5db",
-                marginBottom: 8,
-                whiteSpace: "pre-wrap",
-              }}
-            >
-              Thinking: {item.text}
-            </div>
-          );
+          return <TimelineThought key={i} text={item.text} />;
         }
         if (item.kind === "text") {
           return <TimelineText key={i} text={item.text} />;
@@ -2903,3 +2887,87 @@ export function ChatPanel({
     </div>
   );
 }
+
+function TimelineThought({ text }: { text: string }) {
+  // Auto-expand if the thought is actively streaming (i.e., less than a full turn old)
+  // but let the user collapse it. To keep it simple, we default to false (collapsed),
+  // but if the component mounts and text is very short (just starting to stream), we could
+  // expand it. A better UX is to collapse it by default but allow expanding. 
+  // However, you asked that the final message display before the final thinking hides the response.
+  // Wait, the prompt: "most tools collapse the thinking text once the next action starts to make the chat less noisy. But also gives the user to the option to expand it."
+  
+  // We can track if we are the "latest" thought, but an easier way is to just use a ref
+  // to see if we're actively receiving new props (streaming).
+  
+  const [expanded, setExpanded] = React.useState(true);
+  const textLenRef = React.useRef(text.length);
+  
+  React.useEffect(() => {
+    // If text stops growing for a bit, auto-collapse
+    let t = setTimeout(() => {
+      if (text.length === textLenRef.current) {
+        setExpanded(false);
+      }
+    }, 1500);
+    return () => clearTimeout(t);
+  }, [text]);
+
+  React.useEffect(() => {
+    textLenRef.current = text.length;
+  }, [text]);
+
+  return (
+    <div
+      style={{
+        margin: "8px 0",
+        fontFamily: "var(--font-inter),ui-sans-serif,sans-serif",
+      }}
+    >
+      <button
+        onClick={() => setExpanded((v) => !v)}
+        style={{
+          display: "flex",
+          alignItems: "center",
+          gap: 6,
+          background: "transparent",
+          border: "none",
+          padding: 0,
+          cursor: "pointer",
+          fontSize: "0.85rem",
+          color: "#9ca3af",
+          fontWeight: 500,
+        }}
+      >
+        <span
+          style={{
+            transform: expanded ? "rotate(90deg)" : "rotate(0deg)",
+            transition: "transform 0.15s ease",
+            display: "inline-block",
+            fontSize: "0.6rem",
+          }}
+        >
+          ▶
+        </span>
+        Thinking...
+      </button>
+      {expanded && (
+        <div
+          style={{
+            marginTop: 6,
+            marginLeft: 4,
+            padding: "8px 12px",
+            borderLeft: "2px solid #e5e7eb",
+            fontSize: "0.85rem",
+            color: "#6b7280",
+            fontStyle: "italic",
+            background: "#f9fafb",
+            whiteSpace: "pre-wrap",
+            borderRadius: "0 4px 4px 0",
+          }}
+        >
+          {text.trim()}
+        </div>
+      )}
+    </div>
+  );
+}