fix(tui): buffer streaming messages by runId to prevent render ordering issues

Fixes #1172

- Add per-runId message buffering in ChatLog
- Separate thinking stream from content stream handling
- Ensure proper sequencing (thinking always before content)
- Model-agnostic: works with or without thinking tokens

src/tui/tui-event-handlers.ts

@@ -1,6 +1,12 @@
 import type { TUI } from "@mariozechner/pi-tui";
 import type { ChatLog } from "./components/chat-log.js";
-import { asString, extractTextFromMessage, resolveFinalAssistantText } from "./tui-formatters.js";
+import {
+  asString,
+  extractTextFromMessage,
+  extractThinkingFromMessage,
+  extractContentFromMessage,
+  resolveFinalAssistantText,
+} from "./tui-formatters.js";
 import type { AgentEvent, ChatEvent, TuiStateAccess } from "./tui-types.js";
 
 type EventHandlerContext = {
@@ -11,12 +17,25 @@ type EventHandlerContext = {
   refreshSessionInfo?: () => Promise<void>;
 };
 
+/**
+ * Per-run stream buffer for tracking thinking/content separately.
+ * Enables proper sequencing regardless of network arrival order.
+ */
+interface RunStreamBuffer {
+  thinkingText: string;
+  contentText: string;
+  lastUpdateMs: number;
+}
+
 export function createEventHandlers(context: EventHandlerContext) {
   const { chatLog, tui, state, setActivityStatus, refreshSessionInfo } = context;
   const finalizedRuns = new Map<string, number>();
+  // FIXED: Per-run stream buffers for proper isolation
+  const runBuffers = new Map<string, RunStreamBuffer>();
 
   const noteFinalizedRun = (runId: string) => {
     finalizedRuns.set(runId, Date.now());
+    runBuffers.delete(runId); // Clean up buffer
     if (finalizedRuns.size <= 200) return;
     const keepUntil = Date.now() - 10 * 60 * 1000;
     for (const [key, ts] of finalizedRuns) {
@@ -31,6 +50,22 @@ export function createEventHandlers(context: EventHandlerContext) {
     }
   };
 
+  /**
+   * Get or create a stream buffer for a specific runId.
+   */
+  const getOrCreateBuffer = (runId: string): RunStreamBuffer => {
+    let buffer = runBuffers.get(runId);
+    if (!buffer) {
+      buffer = {
+        thinkingText: "",
+        contentText: "",
+        lastUpdateMs: Date.now(),
+      };
+      runBuffers.set(runId, buffer);
+    }
+    return buffer;
+  };
+
   const handleChatEvent = (payload: unknown) => {
     if (!payload || typeof payload !== "object") return;
     const evt = payload as ChatEvent;
@@ -40,11 +75,33 @@ export function createEventHandlers(context: EventHandlerContext) {
       if (evt.state === "final") return;
     }
     if (evt.state === "delta") {
-      const text = extractTextFromMessage(evt.message, {
-        includeThinking: state.showThinking,
+      const buffer = getOrCreateBuffer(evt.runId);
+
+      // FIXED: Extract thinking and content SEPARATELY for proper sequencing
+      // This is model-agnostic: models without thinking blocks just return empty string
+      const thinkingText = extractThinkingFromMessage(evt.message);
+      const contentText = extractContentFromMessage(evt.message);
+
+      // Update buffer with new content
+      // In streaming, we typically receive the full accumulated text each time
+      if (thinkingText) {
+        buffer.thinkingText = thinkingText;
+      }
+      if (contentText) {
+        buffer.contentText = contentText;
+      }
+      buffer.lastUpdateMs = Date.now();
+
+      // Skip render if both are empty
+      if (!buffer.thinkingText && !buffer.contentText) return;
+
+      // FIXED: Pass separated streams to ChatLog for proper sequencing
+      chatLog.updateAssistant("", evt.runId, {
+        thinkingText: buffer.thinkingText,
+        contentText: buffer.contentText,
+        showThinking: state.showThinking,
       });
-      if (!text) return;
-      chatLog.updateAssistant(text, evt.runId);
+
       setActivityStatus("streaming");
     }
     if (evt.state === "final") {
@@ -54,11 +111,23 @@ export function createEventHandlers(context: EventHandlerContext) {
             ? ((evt.message as Record<string, unknown>).stopReason as string)
             : ""
           : "";
-      const text = extractTextFromMessage(evt.message, {
-        includeThinking: state.showThinking,
-      });
+
+      // FIXED: Extract final content with proper thinking handling
+      const thinkingText = extractThinkingFromMessage(evt.message);
+      const contentText = extractContentFromMessage(evt.message);
+
+      // Compose final text with proper ordering (thinking before content)
+      const parts: string[] = [];
+      if (state.showThinking && thinkingText.trim()) {
+        parts.push(`[thinking]\n${thinkingText}`);
+      }
+      if (contentText.trim()) {
+        parts.push(contentText);
+      }
+      const finalComposed = parts.join("\n\n").trim();
+
       const finalText = resolveFinalAssistantText({
-        finalText: text,
+        finalText: finalComposed,
         streamedText: chatLog.getStreamingText(evt.runId),
       });
       chatLog.finalizeAssistant(finalText, evt.runId);
@@ -70,12 +139,14 @@ export function createEventHandlers(context: EventHandlerContext) {
     }
     if (evt.state === "aborted") {
       chatLog.addSystem("run aborted");
+      runBuffers.delete(evt.runId);
       state.activeChatRunId = null;
       setActivityStatus("aborted");
       void refreshSessionInfo?.();
     }
     if (evt.state === "error") {
       chatLog.addSystem(`run error: ${evt.errorMessage ?? "unknown"}`);
+      runBuffers.delete(evt.runId);
       state.activeChatRunId = null;
       setActivityStatus("error");
       void refreshSessionInfo?.();