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
@@ -0,0 +1,148 @@
+import assert from "node:assert/strict";
+import test from "node:test";
+import {
+  appendStreamingDelta,
+  asStreamingAssistantMetadata,
+  clearStreamingAssistantTrace,
+  createStreamingAssistantMetadata,
+  createStreamingAttemptState,
+  finalizeStreamingAssistant,
+  promoteStreamingAssistantTrace,
+  setStreamingAssistantSegment,
+  transitionStreamingToolCall,
+} from "../src/lib/chat-stream-presentation.ts";
+
+function assistantMessage(content = "") {
+  return {
+    id: "temp-assistant-test",
+    createdAt: "2026-08-30T00:00:00.000Z",
+    role: "assistant",
+    content,
+    name: null,
+    metadata: createStreamingAssistantMetadata(),
+  };
+}
+
+function metadata(messages) {
+  return asStreamingAssistantMetadata(messages.at(-1)?.metadata);
+}
+
+test("a tool boundary promotes the current segment and the next delta starts the answer in its place", () => {
+  let messages = [assistantMessage()];
+  let attempt = appendStreamingDelta(createStreamingAttemptState(), "I'll search for that.");
+  messages = setStreamingAssistantSegment(messages, "temp-assistant-", attempt.segmentText);
+
+  const toolTransition = transitionStreamingToolCall(attempt, {
+    toolCallId: "search-1",
+    status: "initiated",
+  });
+  attempt = toolTransition.attempt;
+  messages = promoteStreamingAssistantTrace(messages, "temp-assistant-", toolTransition.traceText ?? "");
+
+  assert.equal(messages.at(-1).content, "");
+  assert.deepEqual(metadata(messages), {
+    kind: "streaming_assistant",
+    ephemeralTrace: "I'll search for that.",
+    traceRevision: 1,
+  });
+
+  attempt = appendStreamingDelta(attempt, "Here is the answer");
+  messages = setStreamingAssistantSegment(messages, "temp-assistant-", attempt.segmentText);
+
+  assert.equal(messages.at(-1).content, "Here is the answer");
+  assert.equal(metadata(messages)?.ephemeralTrace, null);
+
+  messages = finalizeStreamingAssistant(messages, "temp-assistant-", "Here is the answer.");
+  assert.equal(messages.at(-1).content, "Here is the answer.");
+  assert.equal(metadata(messages)?.ephemeralTrace, null);
+});
+
+test("a later tool round replaces the previous trace and parallel initiated calls do not erase it", () => {
+  let messages = [assistantMessage("First trace")];
+  let attempt = appendStreamingDelta(createStreamingAttemptState(), "First trace");
+  let transition = transitionStreamingToolCall(attempt, { toolCallId: "tool-1", status: "initiated" });
+  attempt = transition.attempt;
+  messages = promoteStreamingAssistantTrace(messages, "temp-assistant-", transition.traceText ?? "");
+
+  attempt = appendStreamingDelta(attempt, "I need one more lookup.");
+  messages = setStreamingAssistantSegment(messages, "temp-assistant-", attempt.segmentText);
+  transition = transitionStreamingToolCall(attempt, { toolCallId: "tool-2", status: "initiated" });
+  attempt = transition.attempt;
+  messages = promoteStreamingAssistantTrace(messages, "temp-assistant-", transition.traceText ?? "");
+
+  assert.equal(metadata(messages)?.ephemeralTrace, "I need one more lookup.");
+  assert.equal(metadata(messages)?.traceRevision, 2);
+
+  transition = transitionStreamingToolCall(attempt, { toolCallId: "tool-3", status: "initiated" });
+  assert.equal(transition.startedNewToolCall, true);
+  assert.equal(transition.traceText, null);
+  messages = promoteStreamingAssistantTrace(messages, "temp-assistant-", transition.traceText ?? "");
+
+  assert.equal(metadata(messages)?.ephemeralTrace, "I need one more lookup.");
+  assert.equal(metadata(messages)?.traceRevision, 2);
+});
+
+test("terminal and duplicate tool events never promote or reset the in-progress segment", () => {
+  const initial = appendStreamingDelta(createStreamingAttemptState(), "candidate");
+  const completed = transitionStreamingToolCall(initial, { toolCallId: "tool-1", status: "completed" });
+
+  assert.equal(completed.startedNewToolCall, false);
+  assert.strictEqual(completed.attempt, initial);
+
+  const initiated = transitionStreamingToolCall(initial, { toolCallId: "tool-1", status: "initiated" });
+  const withNextSegment = appendStreamingDelta(initiated.attempt, "next segment");
+  const duplicate = transitionStreamingToolCall(withNextSegment, { toolCallId: "tool-1", status: "initiated" });
+
+  assert.equal(duplicate.startedNewToolCall, false);
+  assert.strictEqual(duplicate.attempt, withNextSegment);
+  assert.equal(duplicate.attempt.segmentText, "next segment");
+});
+
+test("a fresh per-attempt accumulator can replay and reclassify an already-seen tool round", () => {
+  let messages = [assistantMessage("stale partial answer")];
+  let replayAttempt = appendStreamingDelta(createStreamingAttemptState(), "Replayed trace");
+  messages = setStreamingAssistantSegment(messages, "temp-assistant-", replayAttempt.segmentText);
+
+  const replayedTool = transitionStreamingToolCall(replayAttempt, { toolCallId: "same-tool-id", status: "initiated" });
+  messages = promoteStreamingAssistantTrace(messages, "temp-assistant-", replayedTool.traceText ?? "");
+
+  assert.equal(replayedTool.startedNewToolCall, true);
+  assert.equal(messages.at(-1).content, "");
+  assert.equal(metadata(messages)?.ephemeralTrace, "Replayed trace");
+});
+
+test("message helpers preserve identity when no matching temporary assistant exists", () => {
+  const messages = [
+    {
+      id: "persisted-assistant",
+      createdAt: "2026-08-30T00:00:00.000Z",
+      role: "assistant",
+      content: "Persisted",
+      name: null,
+      metadata: null,
+    },
+  ];
+
+  assert.strictEqual(setStreamingAssistantSegment(messages, "temp-assistant-", "new"), messages);
+  assert.strictEqual(promoteStreamingAssistantTrace(messages, "temp-assistant-", "trace"), messages);
+  assert.strictEqual(finalizeStreamingAssistant(messages, "temp-assistant-", "done"), messages);
+  assert.strictEqual(clearStreamingAssistantTrace(messages, "temp-assistant-"), messages);
+});
+
+test("a terminal error clears the ephemeral trace without discarding a partial answer", () => {
+  let messages = [
+    {
+      ...assistantMessage("Partial answer"),
+      metadata: {
+        ...createStreamingAssistantMetadata(),
+        ephemeralTrace: "Still checking one source.",
+        traceRevision: 1,
+      },
+    },
+  ];
+
+  messages = clearStreamingAssistantTrace(messages, "temp-assistant-");
+
+  assert.equal(messages.at(-1).content, "Partial answer");
+  assert.equal(metadata(messages)?.ephemeralTrace, null);
+});