fix(stories): replay queued and stage-less tool calls faithfully

The sample replay gave a queued call (pending, no approval record) a
synthetic running frame, so it ran and then went back to pending. It
is now shown queued from the start, like a waiting or running call.

A tool batch also emitted one frame per call, so a call with no
earlier form produced a frame identical to the one before it. The
batch now finishes only the calls that change, one at a time, and
leaves the rest in their final form, so every frame differs from the
previous one.

Author: pratikbuilds

stories/chat/samples/play-turn.test.ts

@@ -7,7 +7,12 @@ import type {
   PartTool,
   PartToolBatch,
 } from "../../../src/lib/conversation.js";
-import { assistant, generatedImages, toolStorm } from "./conversations.js";
+import {
+  assistant,
+  generatedImages,
+  samples,
+  toolStorm,
+} from "./conversations.js";
 import { playTurn, turnFrames } from "./play-turn.js";
 
 function turn(parts: readonly ConversationPart[]): ConversationMessage {
@@ -37,6 +42,20 @@ function lastPart(frame: ConversationMessage | undefined) {
   return frame?.parts.at(-1);
 }
 
+function statusesOf(part: ConversationPart | undefined) {
+  return part?.kind === "tool-batch"
+    ? part.tools.map((call) => call.status)
+    : null;
+}
+
+/** How many frames show exactly the parts of the frame before them. */
+function repeatedFrames(frames: readonly ConversationMessage[]): number {
+  return frames.filter(
+    (frame, index) =>
+      index > 0 && Bun.deepEquals(frame.parts, frames[index - 1]?.parts),
+  ).length;
+}
+
 describe("turnFrames", () => {
   test("starts empty and streaming, and ends with the message itself", () => {
     const message = turn([{ kind: "text", text: "Done." }]);
@@ -132,6 +151,57 @@ describe("turnFrames", () => {
     expect(new Set(statuses)).toEqual(new Set(["approval-requested"]));
   });
 
+  test("shows a queued call in its final form at once, never as running", () => {
+    const queued = tool({ status: "pending", output: undefined });
+    expect(turnFrames(turn([queued])).map(lastPart)).toEqual([
+      undefined,
+      queued,
+    ]);
+
+    const batch: PartToolBatch = {
+      kind: "tool-batch",
+      id: "batch",
+      tools: [tool({ toolCallId: "a" }), { ...queued, toolCallId: "b" }],
+    };
+    const queuedStatuses = turnFrames(turn([batch])).flatMap((frame) => {
+      const part = lastPart(frame);
+      return part?.kind === "tool-batch" ? [part.tools[1]?.status] : [];
+    });
+    expect(new Set(queuedStatuses)).toEqual(new Set(["pending"]));
+  });
+
+  test("never repeats a frame after a batch call that has no stage", () => {
+    const batch: PartToolBatch = {
+      kind: "tool-batch",
+      id: "batch",
+      tools: [
+        tool({ toolCallId: "a" }),
+        tool({ toolCallId: "b", status: "running", output: undefined }),
+        tool({ toolCallId: "c" }),
+        tool({ toolCallId: "d", status: "running", output: undefined }),
+      ],
+    };
+    const frames = turnFrames(turn([batch, { kind: "text", text: "ok" }]));
+    expect(repeatedFrames(frames)).toBe(0);
+    expect(frames.map((frame) => statusesOf(frame.parts[0]))).toEqual([
+      null,
+      ["running", "running", "running", "running"],
+      ["output-available", "running", "running", "running"],
+      ["output-available", "running", "output-available", "running"],
+      ["output-available", "running", "output-available", "running"],
+    ]);
+  });
+
+  test("never repeats a frame in any sample message", () => {
+    const messages = samples.flatMap((sample) => [
+      ...sample.messages,
+      ...(sample.thread?.replies ?? []),
+    ]);
+    for (const message of messages) {
+      expect(repeatedFrames(turnFrames(message))).toBe(0);
+    }
+  });
+
   test("starts every call in a batch together and finishes them one by one", () => {
     const batch: PartToolBatch = {
       kind: "tool-batch",

stories/chat/samples/play-turn.ts

@@ -30,7 +30,8 @@ function prefixes(value: string): readonly string[] {
 /**
  * The forms a tool call takes before its final one. A call with an approval
  * record first waits on a person and never shows as running; a call without
- * one runs, then settles.
+ * one runs, then settles. A call still queued, waiting or running is already
+ * in its earliest form, so it has none.
  */
 function toolStages(part: PartTool): readonly PartTool[] {
   if (part.approval !== undefined) {
@@ -46,7 +47,11 @@ function toolStages(part: PartTool): readonly PartTool[] {
       },
     ];
   }
-  if (part.status === "approval-requested" || part.status === "running") {
+  if (
+    part.status === "pending" ||
+    part.status === "approval-requested" ||
+    part.status === "running"
+  ) {
     return [];
   }
   return [
@@ -96,14 +101,21 @@ function partStages(part: ConversationPart): readonly ConversationPart[] {
     case "tool-trace":
       return toolStages(part);
     case "tool-batch": {
-      // Every call starts together, then they finish one at a time.
-      const started = part.tools.map((tool) => toolStages(tool)[0] ?? tool);
-      return part.tools.map((_, finished) => ({
-        ...part,
-        tools: part.tools.map((tool, index) =>
-          index < finished ? tool : (started[index] ?? tool),
-        ),
-      }));
+      // Every call starts together, then the ones that change finish one at
+      // a time; a call with no earlier form shows its final one throughout.
+      const starts = part.tools.flatMap((tool, index) => {
+        const [start] = toolStages(tool);
+        return start === undefined ? [] : [{ index, start }];
+      });
+      return starts.map((_, finished) => {
+        const unfinished = new Map(
+          starts.slice(finished).map(({ index, start }) => [index, start]),
+        );
+        return {
+          ...part,
+          tools: part.tools.map((tool, index) => unfinished.get(index) ?? tool),
+        };
+      });
     }
     case "image":
       return imageStages(part);