fix: group parallel tool results in stream retry replay (#3072)

src/__tests__/retry_replay_utils.test.ts

+import { describe, it, expect } from "vitest";
+import {
+  type RetryReplayEvent,
+  buildRetryReplayMessages,
+  maybeCaptureRetryReplayEvent,
+  maybeCaptureRetryReplayText,
+  toToolResultOutput,
+} from "@/pro/main/ipc/handlers/local_agent/retry_replay_utils";
+
+// ---------------------------------------------------------------------------
+// buildRetryReplayMessages
+// ---------------------------------------------------------------------------
+
+describe("buildRetryReplayMessages", () => {
+  it("returns empty array when no events are provided", () => {
+    expect(buildRetryReplayMessages([])).toEqual([]);
+  });
+
+  it("replays a single completed tool exchange", () => {
+    const events: RetryReplayEvent[] = [
+      { type: "assistant-text", text: "Let me read that file." },
+      {
+        type: "tool-call",
+        toolCallId: "call-1",
+        toolName: "readFile",
+        input: { path: "foo.ts" },
+      },
+      {
+        type: "tool-result",
+        toolCallId: "call-1",
+        toolName: "readFile",
+        output: "file contents",
+      },
+    ];
+
+    const messages = buildRetryReplayMessages(events);
+    expect(messages).toHaveLength(2);
+    expect(messages[0].role).toBe("assistant");
+    expect(messages[0].content).toEqual([
+      { type: "text", text: "Let me read that file." },
+      {
+        type: "tool-call",
+        toolCallId: "call-1",
+        toolName: "readFile",
+        input: { path: "foo.ts" },
+      },
+    ]);
+    expect(messages[1].role).toBe("tool");
+    expect(messages[1].content).toEqual([
+      {
+        type: "tool-result",
+        toolCallId: "call-1",
+        toolName: "readFile",
+        output: { type: "text", value: "file contents" },
+      },
+    ]);
+  });
+
+  it("groups parallel tool results into a single tool message (fixes #3070)", () => {
+    // Simulates 3 parallel tool calls: all calls emitted, then all results.
+    const events: RetryReplayEvent[] = [
+      { type: "assistant-text", text: "Reading files..." },
+      {
+        type: "tool-call",
+        toolCallId: "call-A",
+        toolName: "readFile",
+        input: { path: "a.ts" },
+      },
+      {
+        type: "tool-call",
+        toolCallId: "call-B",
+        toolName: "readFile",
+        input: { path: "b.ts" },
+      },
+      {
+        type: "tool-call",
+        toolCallId: "call-C",
+        toolName: "readFile",
+        input: { path: "c.ts" },
+      },
+      {
+        type: "tool-result",
+        toolCallId: "call-A",
+        toolName: "readFile",
+        output: "contents-a",
+      },
+      {
+        type: "tool-result",
+        toolCallId: "call-B",
+        toolName: "readFile",
+        output: "contents-b",
+      },
+      {
+        type: "tool-result",
+        toolCallId: "call-C",
+        toolName: "readFile",
+        output: "contents-c",
+      },
+    ];
+
+    const messages = buildRetryReplayMessages(events);
+
+    // Should produce exactly: assistant[text, callA, callB, callC], tool[resultA, resultB, resultC]
+    expect(messages).toHaveLength(2);
+
+    expect(messages[0].role).toBe("assistant");
+    const assistantContent = messages[0].content as Array<{
+      type: string;
+      toolCallId?: string;
+    }>;
+    expect(assistantContent).toHaveLength(4); // text + 3 calls
+    expect(assistantContent[0].type).toBe("text");
+    expect(assistantContent[1].toolCallId).toBe("call-A");
+    expect(assistantContent[2].toolCallId).toBe("call-B");
+    expect(assistantContent[3].toolCallId).toBe("call-C");
+
+    expect(messages[1].role).toBe("tool");
+    const toolContent = messages[1].content as Array<{
+      type: string;
+      toolCallId: string;
+    }>;
+    expect(toolContent).toHaveLength(3); // all 3 results grouped
+    expect(toolContent[0].toolCallId).toBe("call-A");
+    expect(toolContent[1].toolCallId).toBe("call-B");
+    expect(toolContent[2].toolCallId).toBe("call-C");
+  });
+
+  it("handles sequential tool calls followed by parallel calls", () => {
+    const events: RetryReplayEvent[] = [
+      // Sequential call
+      { type: "assistant-text", text: "Step 1" },
+      {
+        type: "tool-call",
+        toolCallId: "seq-1",
+        toolName: "readFile",
+        input: {},
+      },
+      {
+        type: "tool-result",
+        toolCallId: "seq-1",
+        toolName: "readFile",
+        output: "result-1",
+      },
+      // Parallel calls
+      { type: "assistant-text", text: "Step 2" },
+      {
+        type: "tool-call",
+        toolCallId: "par-A",
+        toolName: "readFile",
+        input: {},
+      },
+      {
+        type: "tool-call",
+        toolCallId: "par-B",
+        toolName: "readFile",
+        input: {},
+      },
+      {
+        type: "tool-result",
+        toolCallId: "par-A",
+        toolName: "readFile",
+        output: "result-A",
+      },
+      {
+        type: "tool-result",
+        toolCallId: "par-B",
+        toolName: "readFile",
+        output: "result-B",
+      },
+    ];
+
+    const messages = buildRetryReplayMessages(events);
+    expect(messages).toHaveLength(4);
+
+    // Sequential: assistant[text, seq-1] → tool[seq-1-result]
+    expect(messages[0].role).toBe("assistant");
+    expect(messages[1].role).toBe("tool");
+    expect((messages[1].content as unknown[]).length).toBe(1);
+
+    // Parallel: assistant[text, par-A, par-B] → tool[par-A-result, par-B-result]
+    expect(messages[2].role).toBe("assistant");
+    expect(messages[3].role).toBe("tool");
+    expect((messages[3].content as unknown[]).length).toBe(2);
+  });
+
+  it("excludes incomplete tool exchanges (call without result)", () => {
+    const events: RetryReplayEvent[] = [
+      { type: "assistant-text", text: "Working..." },
+      {
+        type: "tool-call",
+        toolCallId: "complete",
+        toolName: "readFile",
+        input: {},
+      },
+      {
+        type: "tool-result",
+        toolCallId: "complete",
+        toolName: "readFile",
+        output: "done",
+      },
+      { type: "assistant-text", text: "More work..." },
+      {
+        type: "tool-call",
+        toolCallId: "incomplete",
+        toolName: "writeFile",
+        input: {},
+      },
+      // No tool-result for "incomplete" — stream died
+    ];
+
+    const messages = buildRetryReplayMessages(events);
+    // Only the completed exchange should appear
+    expect(messages).toHaveLength(3); // assistant, tool, assistant (trailing text)
+
+    expect(messages[0].role).toBe("assistant");
+    const assistantContent = messages[0].content as Array<{
+      type: string;
+      toolCallId?: string;
+    }>;
+    expect(assistantContent.some((c) => c.toolCallId === "incomplete")).toBe(
+      false,
+    );
+
+    expect(messages[2].role).toBe("assistant");
+    expect(messages[2].content).toEqual([
+      { type: "text", text: "More work..." },
+    ]);
+  });
+
+  it("excludes incomplete parallel calls mixed with complete ones", () => {
+    const events: RetryReplayEvent[] = [
+      {
+        type: "tool-call",
+        toolCallId: "call-A",
+        toolName: "readFile",
+        input: {},
+      },
+      {
+        type: "tool-call",
+        toolCallId: "call-B",
+        toolName: "readFile",
+        input: {},
+      },
+      {
+        type: "tool-result",
+        toolCallId: "call-A",
+        toolName: "readFile",
+        output: "result-A",
+      },
+      // call-B has no result (stream died mid-batch)
+    ];
+
+    const messages = buildRetryReplayMessages(events);
+    expect(messages).toHaveLength(2);
+
+    // Only call-A should be in the assistant message
+    const assistantContent = messages[0].content as Array<{
+      type: string;
+      toolCallId?: string;
+    }>;
+    expect(assistantContent).toHaveLength(1);
+    expect(assistantContent[0].toolCallId).toBe("call-A");
+
+    // Only result-A in the tool message
+    const toolContent = messages[1].content as Array<{
+      type: string;
+      toolCallId: string;
+    }>;
+    expect(toolContent).toHaveLength(1);
+    expect(toolContent[0].toolCallId).toBe("call-A");
+  });
+
+  it("skips whitespace-only text events", () => {
+    const events: RetryReplayEvent[] = [
+      { type: "assistant-text", text: "   " },
+      {
+        type: "tool-call",
+        toolCallId: "c1",
+        toolName: "readFile",
+        input: {},
+      },
+      {
+        type: "tool-result",
+        toolCallId: "c1",
+        toolName: "readFile",
+        output: "ok",
+      },
+    ];
+
+    const messages = buildRetryReplayMessages(events);
+    expect(messages).toHaveLength(2);
+    // The assistant message should only have the tool-call, no whitespace text
+    const assistantContent = messages[0].content as Array<{ type: string }>;
+    expect(assistantContent).toHaveLength(1);
+    expect(assistantContent[0].type).toBe("tool-call");
+  });
+});
+
+// ---------------------------------------------------------------------------
+// maybeCaptureRetryReplayEvent
+// ---------------------------------------------------------------------------
+
+describe("maybeCaptureRetryReplayEvent", () => {
+  it("captures tool-call events", () => {
+    const events: RetryReplayEvent[] = [];
+    maybeCaptureRetryReplayEvent(events, {
+      type: "tool-call",
+      toolCallId: "tc-1",
+      toolName: "readFile",
+      input: { path: "x.ts" },
+    });
+    expect(events).toHaveLength(1);
+    expect(events[0]).toEqual({
+      type: "tool-call",
+      toolCallId: "tc-1",
+      toolName: "readFile",
+      input: { path: "x.ts" },
+    });
+  });
+
+  it("deduplicates tool-call events by toolCallId", () => {
+    const events: RetryReplayEvent[] = [];
+    const part = {
+      type: "tool-call",
+      toolCallId: "tc-1",
+      toolName: "readFile",
+      input: {},
+    };
+    maybeCaptureRetryReplayEvent(events, part);
+    maybeCaptureRetryReplayEvent(events, part);
+    expect(events).toHaveLength(1);
+  });
+
+  it("captures tool-result events", () => {
+    const events: RetryReplayEvent[] = [];
+    maybeCaptureRetryReplayEvent(events, {
+      type: "tool-result",
+      toolCallId: "tc-1",
+      toolName: "readFile",
+      output: "data",
+    });
+    expect(events).toHaveLength(1);
+    expect(events[0].type).toBe("tool-result");
+  });
+
+  it("ignores non-object or untyped parts", () => {
+    const events: RetryReplayEvent[] = [];
+    maybeCaptureRetryReplayEvent(events, null);
+    maybeCaptureRetryReplayEvent(events, "string");
+    maybeCaptureRetryReplayEvent(events, 42);
+    maybeCaptureRetryReplayEvent(events, { noType: true });
+    expect(events).toHaveLength(0);
+  });
+});
+
+// ---------------------------------------------------------------------------
+// maybeCaptureRetryReplayText
+// ---------------------------------------------------------------------------
+
+describe("maybeCaptureRetryReplayText", () => {
+  it("appends new text event", () => {
+    const events: RetryReplayEvent[] = [];
+    maybeCaptureRetryReplayText(events, "hello");
+    expect(events).toHaveLength(1);
+    expect(events[0]).toEqual({ type: "assistant-text", text: "hello" });
+  });
+
+  it("concatenates to existing trailing text event", () => {
+    const events: RetryReplayEvent[] = [
+      { type: "assistant-text", text: "hel" },
+    ];
+    maybeCaptureRetryReplayText(events, "lo");
+    expect(events).toHaveLength(1);
+    expect(events[0]).toEqual({ type: "assistant-text", text: "hello" });
+  });
+
+  it("ignores empty text", () => {
+    const events: RetryReplayEvent[] = [];
+    maybeCaptureRetryReplayText(events, "");
+    expect(events).toHaveLength(0);
+  });
+
+  it("does nothing when events is null", () => {
+    // Should not throw
+    maybeCaptureRetryReplayText(null, "hello");
+  });
+});
+
+// ---------------------------------------------------------------------------
+// toToolResultOutput
+// ---------------------------------------------------------------------------
+
+describe("toToolResultOutput", () => {
+  it("wraps string values directly", () => {
+    expect(toToolResultOutput("hello")).toEqual({
+      type: "text",
+      value: "hello",
+    });
+  });
+
+  it("JSON-stringifies objects", () => {
+    expect(toToolResultOutput({ key: "val" })).toEqual({
+      type: "text",
+      value: '{"key":"val"}',
+    });
+  });
+
+  it("handles non-serializable values", () => {
+    const circular: Record<string, unknown> = {};
+    circular.self = circular;
+    const result = toToolResultOutput(circular);
+    expect(result.type).toBe("text");
+    expect(typeof result.value).toBe("string");
+  });
+});

src/pro/main/ipc/handlers/local_agent/local_agent_handler.ts

@@ -82,6 +82,12 @@ import {
 import { getPostCompactionMessages } from "@/ipc/handlers/compaction/compaction_utils";
 import { DEFAULT_MAX_TOOL_CALL_STEPS } from "@/constants/settings_constants";
 import { DyadError, DyadErrorKind } from "@/errors/dyad_error";
+import {
+  type RetryReplayEvent,
+  maybeCaptureRetryReplayEvent,
+  maybeCaptureRetryReplayText,
+  maybeAppendRetryReplayForRetry,
+} from "./retry_replay_utils";
 
 const logger = log.scope("local_agent_handler");
 const PLANNING_QUESTIONNAIRE_TOOL_NAME = "planning_questionnaire";
@@ -122,24 +128,6 @@ interface ToolStreamingEntry {
 }
 const toolStreamingEntries = new Map<string, ToolStreamingEntry>();
 
-type RetryReplayEvent =
-  | {
-      type: "assistant-text";
-      text: string;
-    }
-  | {
-      type: "tool-call";
-      toolCallId: string;
-      toolName: string;
-      input: unknown;
-    }
-  | {
-      type: "tool-result";
-      toolCallId: string;
-      toolName: string;
-      output: unknown;
-    };
-
 function getOrCreateStreamingEntry(
   id: string,
   toolName?: string,
@@ -1401,201 +1389,10 @@ function shouldRetryTransientStreamError(params: {
   );
 }
 
-function maybeCaptureRetryReplayEvent(
-  retryReplayEvents: RetryReplayEvent[],
-  part: unknown,
-): void {
-  if (!isRecord(part) || typeof part.type !== "string") {
-    return;
-  }
-
-  if (
-    part.type === "tool-call" &&
-    typeof part.toolCallId === "string" &&
-    typeof part.toolName === "string"
-  ) {
-    // Keep one emitted tool-call event per toolCallId.
-    if (
-      retryReplayEvents.some(
-        (event) =>
-          event.type === "tool-call" && event.toolCallId === part.toolCallId,
-      )
-    ) {
-      return;
-    }
-
-    retryReplayEvents.push({
-      type: "tool-call",
-      toolCallId: part.toolCallId,
-      toolName: part.toolName,
-      input:
-        typeof part.input === "object" && part.input !== null ? part.input : {},
-    });
-    return;
-  }
-
-  if (
-    part.type === "tool-result" &&
-    typeof part.toolCallId === "string" &&
-    typeof part.toolName === "string"
-  ) {
-    // Keep one emitted tool-result event per toolCallId.
-    if (
-      retryReplayEvents.some(
-        (event) =>
-          event.type === "tool-result" && event.toolCallId === part.toolCallId,
-      )
-    ) {
-      return;
-    }
-
-    retryReplayEvents.push({
-      type: "tool-result",
-      toolCallId: part.toolCallId,
-      toolName: part.toolName,
-      output: part.output,
-    });
-  }
-}
-
-function maybeAppendRetryReplayForRetry(params: {
-  retryReplayEvents: RetryReplayEvent[];
-  currentMessageHistoryRef: ModelMessage[];
-  accumulatedAiMessagesRef: ModelMessage[];
-  onCurrentMessageHistoryUpdate: (next: ModelMessage[]) => void;
-}) {
-  const {
-    retryReplayEvents,
-    currentMessageHistoryRef,
-    accumulatedAiMessagesRef,
-    onCurrentMessageHistoryUpdate,
-  } = params;
-  const replayMessages: ModelMessage[] = [];
-  const pendingAssistantParts: Array<
-    | { type: "text"; text: string }
-    | {
-        type: "tool-call";
-        toolCallId: string;
-        toolName: string;
-        input: unknown;
-      }
-  > = [];
-  const toolCallsWithResult = new Set<string>();
-  const toolResultsWithCall = new Set<string>();
-
-  for (const event of retryReplayEvents) {
-    if (event.type === "tool-call") {
-      toolResultsWithCall.add(event.toolCallId);
-      continue;
-    }
-    if (event.type === "tool-result") {
-      toolCallsWithResult.add(event.toolCallId);
-    }
-  }
-
-  const completedToolExchangeIds = new Set(
-    [...toolCallsWithResult].filter((toolCallId) =>
-      toolResultsWithCall.has(toolCallId),
-    ),
-  );
-
-  const flushPendingAssistantMessage = () => {
-    if (pendingAssistantParts.length === 0) {
-      return;
-    }
-    replayMessages.push({
-      role: "assistant",
-      content: [...pendingAssistantParts],
-    });
-    pendingAssistantParts.length = 0;
-  };
-
-  for (const event of retryReplayEvents) {
-    if (event.type === "assistant-text") {
-      if (!event.text.trim()) {
-        continue;
-      }
-      pendingAssistantParts.push({ type: "text", text: event.text });
-      continue;
-    }
-
-    if (event.type === "tool-call") {
-      if (!completedToolExchangeIds.has(event.toolCallId)) {
-        continue;
-      }
-      pendingAssistantParts.push({
-        type: "tool-call",
-        toolCallId: event.toolCallId,
-        toolName: event.toolName,
-        input: event.input,
-      });
-      continue;
-    }
-
-    if (!completedToolExchangeIds.has(event.toolCallId)) {
-      continue;
-    }
-    flushPendingAssistantMessage();
-    replayMessages.push({
-      role: "tool",
-      content: [
-        {
-          type: "tool-result",
-          toolCallId: event.toolCallId,
-          toolName: event.toolName,
-          output: toToolResultOutput(event.output),
-        },
-      ],
-    });
-  }
-  flushPendingAssistantMessage();
-
-  if (replayMessages.length === 0) {
-    return;
-  }
-
-  onCurrentMessageHistoryUpdate([
-    ...currentMessageHistoryRef,
-    ...replayMessages,
-  ]);
-  accumulatedAiMessagesRef.push(...replayMessages);
-}
-
-function maybeCaptureRetryReplayText(
-  retryReplayEvents: RetryReplayEvent[] | null,
-  text: string,
-): void {
-  if (!retryReplayEvents || text.length === 0) {
-    return;
-  }
-
-  const lastEvent = retryReplayEvents[retryReplayEvents.length - 1];
-  if (lastEvent?.type === "assistant-text") {
-    lastEvent.text += text;
-    return;
-  }
-
-  retryReplayEvents.push({
-    type: "assistant-text",
-    text,
-  });
-}
-
 async function delay(ms: number): Promise<void> {
   await new Promise<void>((resolve) => setTimeout(resolve, ms));
 }
 
-function toToolResultOutput(value: unknown): { type: "text"; value: string } {
-  if (typeof value === "string") {
-    return { type: "text", value };
-  }
-  try {
-    return { type: "text", value: JSON.stringify(value) };
-  } catch {
-    return { type: "text", value: String(value) };
-  }
-}
-
 async function updateResponseInDb(messageId: number, content: string) {
   await db
     .update(messages)

src/pro/main/ipc/handlers/local_agent/retry_replay_utils.ts

+/**
+ * Utilities for building replay messages when retrying after a transient
+ * stream termination. Extracted for testability.
+ */
+
+import type { ModelMessage } from "ai";
+
+export type RetryReplayEvent =
+  | {
+      type: "assistant-text";
+      text: string;
+    }
+  | {
+      type: "tool-call";
+      toolCallId: string;
+      toolName: string;
+      input: unknown;
+    }
+  | {
+      type: "tool-result";
+      toolCallId: string;
+      toolName: string;
+      output: unknown;
+    };
+
+export function toToolResultOutput(value: unknown): {
+  type: "text";
+  value: string;
+} {
+  if (typeof value === "string") {
+    return { type: "text", value };
+  }
+  try {
+    return { type: "text", value: JSON.stringify(value) };
+  } catch {
+    return { type: "text", value: String(value) };
+  }
+}
+
+export function maybeCaptureRetryReplayEvent(
+  retryReplayEvents: RetryReplayEvent[],
+  part: unknown,
+): void {
+  if (
+    !part ||
+    typeof part !== "object" ||
+    !("type" in part) ||
+    typeof (part as Record<string, unknown>).type !== "string"
+  ) {
+    return;
+  }
+
+  const record = part as Record<string, unknown>;
+
+  if (
+    record.type === "tool-call" &&
+    typeof record.toolCallId === "string" &&
+    typeof record.toolName === "string"
+  ) {
+    if (
+      retryReplayEvents.some(
+        (event) =>
+          event.type === "tool-call" && event.toolCallId === record.toolCallId,
+      )
+    ) {
+      return;
+    }
+
+    retryReplayEvents.push({
+      type: "tool-call",
+      toolCallId: record.toolCallId,
+      toolName: record.toolName,
+      input:
+        typeof record.input === "object" && record.input !== null
+          ? record.input
+          : {},
+    });
+    return;
+  }
+
+  if (
+    record.type === "tool-result" &&
+    typeof record.toolCallId === "string" &&
+    typeof record.toolName === "string"
+  ) {
+    if (
+      retryReplayEvents.some(
+        (event) =>
+          event.type === "tool-result" &&
+          event.toolCallId === record.toolCallId,
+      )
+    ) {
+      return;
+    }
+
+    retryReplayEvents.push({
+      type: "tool-result",
+      toolCallId: record.toolCallId,
+      toolName: record.toolName,
+      output: record.output,
+    });
+  }
+}
+
+export function maybeCaptureRetryReplayText(
+  retryReplayEvents: RetryReplayEvent[] | null,
+  text: string,
+): void {
+  if (!retryReplayEvents || text.length === 0) {
+    return;
+  }
+
+  const lastEvent = retryReplayEvents[retryReplayEvents.length - 1];
+  if (lastEvent?.type === "assistant-text") {
+    lastEvent.text += text;
+    return;
+  }
+
+  retryReplayEvents.push({
+    type: "assistant-text",
+    text,
+  });
+}
+
+/**
+ * Builds replay messages from captured stream events for retry after a
+ * transient stream termination. Only includes completed tool exchanges
+ * (tool-call + tool-result pairs).
+ */
+export function buildRetryReplayMessages(
+  retryReplayEvents: RetryReplayEvent[],
+): ModelMessage[] {
+  const replayMessages: ModelMessage[] = [];
+  const pendingAssistantParts: Array<
+    | { type: "text"; text: string }
+    | {
+        type: "tool-call";
+        toolCallId: string;
+        toolName: string;
+        input: unknown;
+      }
+  > = [];
+  const toolCallsWithResult = new Set<string>();
+  const toolResultsWithCall = new Set<string>();
+
+  for (const event of retryReplayEvents) {
+    if (event.type === "tool-call") {
+      toolResultsWithCall.add(event.toolCallId);
+      continue;
+    }
+    if (event.type === "tool-result") {
+      toolCallsWithResult.add(event.toolCallId);
+    }
+  }
+
+  const completedToolExchangeIds = new Set(
+    [...toolCallsWithResult].filter((toolCallId) =>
+      toolResultsWithCall.has(toolCallId),
+    ),
+  );
+
+  const flushPendingAssistantMessage = () => {
+    if (pendingAssistantParts.length === 0) {
+      return;
+    }
+    replayMessages.push({
+      role: "assistant",
+      content: [...pendingAssistantParts],
+    });
+    pendingAssistantParts.length = 0;
+  };
+
+  for (const event of retryReplayEvents) {
+    if (event.type === "assistant-text") {
+      if (!event.text.trim()) {
+        continue;
+      }
+      pendingAssistantParts.push({ type: "text", text: event.text });
+      continue;
+    }
+
+    if (event.type === "tool-call") {
+      if (!completedToolExchangeIds.has(event.toolCallId)) {
+        continue;
+      }
+      pendingAssistantParts.push({
+        type: "tool-call",
+        toolCallId: event.toolCallId,
+        toolName: event.toolName,
+        input: event.input,
+      });
+      continue;
+    }
+
+    if (!completedToolExchangeIds.has(event.toolCallId)) {
+      continue;
+    }
+    flushPendingAssistantMessage();
+    // Merge consecutive tool-result messages so parallel tool results stay
+    // grouped with the preceding assistant message's tool-call blocks.
+    // The Anthropic API requires every tool_use in an assistant message to
+    // have its tool_result in the immediately following message.
+    const lastReplayMsg = replayMessages[replayMessages.length - 1];
+    if (
+      lastReplayMsg?.role === "tool" &&
+      Array.isArray(lastReplayMsg.content)
+    ) {
+      lastReplayMsg.content.push({
+        type: "tool-result",
+        toolCallId: event.toolCallId,
+        toolName: event.toolName,
+        output: toToolResultOutput(event.output),
+      });
+    } else {
+      replayMessages.push({
+        role: "tool",
+        content: [
+          {
+            type: "tool-result",
+            toolCallId: event.toolCallId,
+            toolName: event.toolName,
+            output: toToolResultOutput(event.output),
+          },
+        ],
+      });
+    }
+  }
+  flushPendingAssistantMessage();
+
+  return replayMessages;
+}
+
+export function maybeAppendRetryReplayForRetry(params: {
+  retryReplayEvents: RetryReplayEvent[];
+  currentMessageHistoryRef: ModelMessage[];
+  accumulatedAiMessagesRef: ModelMessage[];
+  onCurrentMessageHistoryUpdate: (next: ModelMessage[]) => void;
+}) {
+  const {
+    retryReplayEvents,
+    currentMessageHistoryRef,
+    accumulatedAiMessagesRef,
+    onCurrentMessageHistoryUpdate,
+  } = params;
+
+  const replayMessages = buildRetryReplayMessages(retryReplayEvents);
+
+  if (replayMessages.length === 0) {
+    return;
+  }
+
+  onCurrentMessageHistoryUpdate([
+    ...currentMessageHistoryRef,
+    ...replayMessages,
+  ]);
+  accumulatedAiMessagesRef.push(...replayMessages);
+}