Polish streaming thinking traces

2026-08-30 15:26:31 -07:00
parent 922172fc60
commit 26d917e96a
10 changed files with 1128 additions and 122 deletions
@@ -81,6 +81,16 @@ import {
  type SidebarSelection,
 } from "@/lib/sidebar-selection";
 import { buildQuickQuestionRequest } from "@/lib/quick-question";
+import {
+  appendStreamingDelta,
+  clearStreamingAssistantTrace,
+  createStreamingAssistantMetadata,
+  createStreamingAttemptState,
+  finalizeStreamingAssistant,
+  promoteStreamingAssistantTrace,
+  setStreamingAssistantSegment,
+  transitionStreamingToolCall,
+} from "@/lib/chat-stream-presentation";
 import { cn } from "@/lib/utils";

 type DraftSelectionKind = "chat" | "search";
@@ -2486,7 +2496,7 @@ export default function App() {
      role: "assistant",
      content: "",
      name: null,
-      metadata: null,
+      metadata: createStreamingAssistantMetadata(),
    };

    let chatId = draftKind === "chat" ? null : selectedItem?.kind === "chat" ? selectedItem.id : null;
@@ -2591,7 +2601,7 @@ export default function App() {

    while (true) {
      let streamErrorMessage: string | null = null;
-      let replayedAssistantText = "";
+      let streamingAttempt = createStreamingAttemptState();
      const abortController = new AbortController();
      chatStreamAbortRefs.current.set(chatId, abortController);

@@ -2609,45 +2619,39 @@ export default function App() {
              if (payload.chatId !== chatId) return;
            },
            onToolCall: (payload) => {
+              const transition = transitionStreamingToolCall(streamingAttempt, payload);
+              streamingAttempt = transition.attempt;
              setPendingChatStates((current) => {
                const pendingState = current[chatId];
                if (!pendingState) return current;
+                const messages = transition.startedNewToolCall
+                  ? promoteStreamingAssistantTrace(pendingState.messages, "temp-assistant-", transition.traceText ?? "")
+                  : pendingState.messages;
                return {
                  ...current,
                  [chatId]: {
-                    messages: upsertOptimisticToolMessage(pendingState.messages, payload, "temp-assistant-"),
+                    messages: upsertOptimisticToolMessage(messages, payload, "temp-assistant-"),
                  },
                };
              });
            },
            onDelta: (payload) => {
              if (!payload.text) return;
-              replayedAssistantText += payload.text;
+              streamingAttempt = appendStreamingDelta(streamingAttempt, payload.text);
+              const segmentText = streamingAttempt.segmentText;
              setPendingChatStates((current) => {
                const pendingState = current[chatId];
                if (!pendingState) return current;
-                let updated = false;
-                const nextMessages = pendingState.messages.map((message, index, all) => {
-                  const isTarget = index === all.length - 1 && message.id.startsWith("temp-assistant-");
-                  if (!isTarget) return message;
-                  updated = true;
-                  return { ...message, content: replayedAssistantText };
-                });
-                return updated ? { ...current, [chatId]: { messages: nextMessages } } : current;
+                const nextMessages = setStreamingAssistantSegment(pendingState.messages, "temp-assistant-", segmentText);
+                return nextMessages === pendingState.messages ? current : { ...current, [chatId]: { messages: nextMessages } };
              });
            },
            onDone: (payload) => {
              setPendingChatStates((current) => {
                const pendingState = current[chatId];
                if (!pendingState) return current;
-                let updated = false;
-                const nextMessages = pendingState.messages.map((message, index, all) => {
-                  const isTarget = index === all.length - 1 && message.id.startsWith("temp-assistant-");
-                  if (!isTarget) return message;
-                  updated = true;
-                  return { ...message, content: payload.text };
-                });
-                return updated ? { ...current, [chatId]: { messages: nextMessages } } : current;
+                const nextMessages = finalizeStreamingAssistant(pendingState.messages, "temp-assistant-", payload.text);
+                return nextMessages === pendingState.messages ? current : { ...current, [chatId]: { messages: nextMessages } };
              });
            },
            onError: (payload) => {
@@ -2896,6 +2900,7 @@ export default function App() {
        const abortController = new AbortController();
        chatStreamAbortRefs.current.set(chatId, abortController);
        let streamErrorMessage: string | null = null;
+        let streamingAttempt = createStreamingAttemptState();

        try {
          const baseChat = await getChat(chatId);
@@ -2912,7 +2917,7 @@ export default function App() {
                  role: "assistant",
                  content: "",
                  name: null,
-                  metadata: null,
+                  metadata: createStreamingAssistantMetadata(),
                }),
              },
            };
@@ -2922,44 +2927,39 @@ export default function App() {
            chatId,
            {
              onToolCall: (payload) => {
+                const transition = transitionStreamingToolCall(streamingAttempt, payload);
+                streamingAttempt = transition.attempt;
                setPendingChatStates((current) => {
                  const pendingState = current[chatId];
                  if (!pendingState) return current;
+                  const messages = transition.startedNewToolCall
+                    ? promoteStreamingAssistantTrace(pendingState.messages, "temp-assistant-", transition.traceText ?? "")
+                    : pendingState.messages;
                  return {
                    ...current,
                    [chatId]: {
-                      messages: upsertOptimisticToolMessage(pendingState.messages, payload, "temp-assistant-"),
+                      messages: upsertOptimisticToolMessage(messages, payload, "temp-assistant-"),
                    },
                  };
                });
              },
              onDelta: (payload) => {
                if (!payload.text) return;
+                streamingAttempt = appendStreamingDelta(streamingAttempt, payload.text);
+                const segmentText = streamingAttempt.segmentText;
                setPendingChatStates((current) => {
                  const pendingState = current[chatId];
                  if (!pendingState) return current;
-                  let updated = false;
-                  const nextMessages = pendingState.messages.map((message, index, all) => {
-                    const isTarget = index === all.length - 1 && message.id.startsWith("temp-assistant-");
-                    if (!isTarget) return message;
-                    updated = true;
-                    return { ...message, content: message.content + payload.text };
-                  });
-                  return updated ? { ...current, [chatId]: { messages: nextMessages } } : current;
+                  const nextMessages = setStreamingAssistantSegment(pendingState.messages, "temp-assistant-", segmentText);
+                  return nextMessages === pendingState.messages ? current : { ...current, [chatId]: { messages: nextMessages } };
                });
              },
              onDone: (payload) => {
                setPendingChatStates((current) => {
                  const pendingState = current[chatId];
                  if (!pendingState) return current;
-                  let updated = false;
-                  const nextMessages = pendingState.messages.map((message, index, all) => {
-                    const isTarget = index === all.length - 1 && message.id.startsWith("temp-assistant-");
-                    if (!isTarget) return message;
-                    updated = true;
-                    return { ...message, content: payload.text };
-                  });
-                  return updated ? { ...current, [chatId]: { messages: nextMessages } } : current;
+                  const nextMessages = finalizeStreamingAssistant(pendingState.messages, "temp-assistant-", payload.text);
+                  return nextMessages === pendingState.messages ? current : { ...current, [chatId]: { messages: nextMessages } };
                });
              },
              onError: (payload) => {
@@ -3211,7 +3211,7 @@ export default function App() {
      role: "assistant",
      content: "",
      name: null,
-      metadata: null,
+      metadata: createStreamingAssistantMetadata(),
    };

    quickQuestionAbortRef.current?.abort();
@@ -3225,6 +3225,7 @@ export default function App() {
    setIsQuickQuestionSending(true);

    let streamErrorMessage: string | null = null;
+    let streamingAttempt = createStreamingAttemptState();

    try {
      await runQuickQuestionStream(
@@ -3235,33 +3236,26 @@ export default function App() {
        }),
        {
          onToolCall: (payload) => {
+            const transition = transitionStreamingToolCall(streamingAttempt, payload);
+            streamingAttempt = transition.attempt;
            setQuickQuestionMessages((current) => {
-              return upsertOptimisticToolMessage(current, payload, "temp-assistant-quick-");
+              const messages = transition.startedNewToolCall
+                ? promoteStreamingAssistantTrace(current, "temp-assistant-quick-", transition.traceText ?? "")
+                : current;
+              return upsertOptimisticToolMessage(messages, payload, "temp-assistant-quick-");
            });
          },
          onDelta: (payload) => {
            if (!payload.text) return;
+            streamingAttempt = appendStreamingDelta(streamingAttempt, payload.text);
+            const segmentText = streamingAttempt.segmentText;
            setQuickQuestionMessages((current) => {
-              let updated = false;
-              const nextMessages = current.map((message, index, all) => {
-                const isTarget = index === all.length - 1 && message.id.startsWith("temp-assistant-quick-");
-                if (!isTarget) return message;
-                updated = true;
-                return { ...message, content: message.content + payload.text };
-              });
-              return updated ? nextMessages : current;
+              return setStreamingAssistantSegment(current, "temp-assistant-quick-", segmentText);
            });
          },
          onDone: (payload) => {
            setQuickQuestionMessages((current) => {
-              let updated = false;
-              const nextMessages = current.map((message, index, all) => {
-                const isTarget = index === all.length - 1 && message.id.startsWith("temp-assistant-quick-");
-                if (!isTarget) return message;
-                updated = true;
-                return { ...message, content: payload.text };
-              });
-              return updated ? nextMessages : current;
+              return finalizeStreamingAssistant(current, "temp-assistant-quick-", payload.text);
            });
          },
          onError: (payload) => {
@@ -3276,6 +3270,7 @@ export default function App() {
      }
    } catch (err) {
      if (abortController.signal.aborted) return;
+      setQuickQuestionMessages((current) => clearStreamingAssistantTrace(current, "temp-assistant-quick-"));
      const message = err instanceof Error ? err.message : String(err);
      if (message.includes("bearer token")) {
        handleAuthFailure(message);