fix(app): preserve plugin-rendered replies when collapsing turns (#8811)

* fix(app): preserve plugin-rendered replies when collapsing turns

* test(app): exercise plugin reply folding in the browser

packages/app/e2e/browser/turn-collapse.spec.ts

@@ -9,6 +9,7 @@ import { getActiveTabTestId } from "../support/helpers/launcher";
 import type { Page, Locator } from "@playwright/test";
 import { test, expect } from "../support/fixtures";
 import { expectComposerVisible } from "../support/helpers/composer";
+import { withDirectoryPlugin } from "../support/helpers/plugin-timeline";
 import { openAgentRoute, seedMockAgentWorkspace } from "../support/helpers/mock-agent";
 
 test.use({
@@ -50,15 +51,16 @@ async function seedAnsweredAgent(title: string) {
 }
 
 /** Chat rows between the latest turn's summary line and its reply, in document order. */
-async function rowsBetweenLatestSummaryAndReply(page: Page): Promise<number> {
-  return page.evaluate(() => {
+async function rowsBetweenLatestSummaryAndReply(
+  page: Page,
+  replyTestId = "assistant-message",
+): Promise<number> {
+  return page.evaluate((testId) => {
     const rows = [...document.querySelectorAll("[data-message-id]")];
     const summaryRow = rows.findLastIndex((row) => row.querySelector("[data-testid=turn-summary]"));
-    const replyRow = rows.findLastIndex((row) =>
-      row.querySelector("[data-testid=assistant-message]"),
-    );
+    const replyRow = rows.findLastIndex((row) => row.querySelector(`[data-testid="${testId}"]`));
     return replyRow - summaryRow - 1;
-  });
+  }, replyTestId);
 }
 
 async function groupedToolsInsideLatestTurn(page: Page): Promise<string[]> {
@@ -126,6 +128,67 @@ async function expandAndCollapseLongPrompt(page: Page) {
 }
 
 test.describe("Turn collapse", () => {
+  test("folds work before a plugin-rendered reply and reopens it", async ({ page }, info) => {
+    test.setTimeout(120_000);
+    await useChatDisplay(page, "collapseTurns");
+    const agent = await seedMockAgentWorkspace({
+      repoPrefix: "plugin-reply-fold-",
+      title: "Plugin reply folding",
+      model: "ten-second-stream",
+      initialPrompt: "Make the fix and run the burst script.",
+    });
+    try {
+      await agent.client.waitForFinish(agent.agentId, 30_000);
+      await withDirectoryPlugin({
+        pluginId: "reply-fold-e2e",
+        source: `import React from "react";
+import { Text } from "react-native";
+import { z } from "zod";
+
+export default function contribute(client) {
+  client.addTimelineTransformer({
+    id: "reply", query: { itemType: "assistant_message" },
+    transform({ item, phase }) {
+      if (phase === "streaming") return undefined;
+      return { items: [{ type: "plugin", kind: "reply", version: 1, data: { text: item.text } }] };
+    },
+  });
+  client.addTimelineRenderer({
+    kind: "reply", version: 1, schema: z.object({ text: z.string() }),
+    Component: ({ item }) => <Text testID="plugin-reply">{item.data.text}</Text>,
+  });
+  return () => {};
+}`,
+        run: async () => {
+          await page.setViewportSize({ width: 1280, height: 1000 });
+          await openAgentRoute(page, agent);
+          await expectComposerVisible(page);
+          for (const width of [1280, 390]) {
+            await page.setViewportSize({ width, height: 1000 });
+            const summary = page.getByTestId("turn-summary").last().getByRole("button");
+            const reply = page.getByTestId("plugin-reply").last();
+            await expect(summary).toBeVisible();
+            await expect(reply).toContainText("(end of synthetic stream)");
+            await expect.poll(() => rowsBetweenLatestSummaryAndReply(page, "plugin-reply")).toBe(0);
+            await page.screenshot({ path: info.outputPath(`plugin-reply-folded-${width}.png`) });
+            await summary.click();
+            await expect
+              .poll(() => rowsBetweenLatestSummaryAndReply(page, "plugin-reply"))
+              .toBeGreaterThan(0);
+            await expect(reply).toContainText("(end of synthetic stream)");
+            await summary.click();
+            await expect.poll(() => rowsBetweenLatestSummaryAndReply(page, "plugin-reply")).toBe(0);
+          }
+          await page.reload();
+          await expect(page.getByTestId("turn-summary").last()).toBeVisible();
+          await expect.poll(() => rowsBetweenLatestSummaryAndReply(page, "plugin-reply")).toBe(0);
+        },
+      });
+    } finally {
+      await agent.cleanup();
+    }
+  });
+
   test("folds each answered turn behind a summary line that opens it", async ({ page }) => {
     test.setTimeout(120_000);
     await page.setViewportSize({ width: 1280, height: 1000 });

packages/app/e2e/support/helpers/plugin-timeline.ts

@@ -134,7 +134,7 @@ export default function contribute(client) {
 }`;
 }
 
-async function withDirectoryPlugin(options: {
+export async function withDirectoryPlugin(options: {
   pluginId: string;
   source: string;
   run: () => Promise<void>;

packages/app/src/agent-stream/presentation.test.ts

@@ -593,6 +593,92 @@ describe("timeline presentation", () => {
   });
 });
 
+it("folds work before a plugin-rendered reply and leaves the reply visible", () => {
+  const timestamp = new Date("2026-01-01T00:00:00.000Z");
+  const tail: StreamItem[] = [
+    { kind: "user_message", id: "prompt", text: "Calculate", timestamp },
+    toolCall("one", "Bash"),
+    { kind: "assistant_message", id: "reply", text: "$$x^2$$", timestamp },
+  ];
+  const result = createStreamPresentation()({
+    ...presentationOptions,
+    transform: installedTransform(installProbe("assistant_message", "phase === 'complete'")),
+    tail,
+    head: [],
+    display: "collapseTurns",
+    isTurnActive: false,
+  });
+  expect(result.tail.map((item) => item.id)).toEqual(["prompt", "probe-plugin/reply/0"]);
+  expect([...result.turns.turnsByHostId.keys()]).toEqual(["prompt"]);
+  expect(pluginData(result.tail)).toEqual([{ text: "$$x^2$$", phase: "complete" }]);
+});
+
+it("keeps streaming replies native, then folds work without losing manual expansion", () => {
+  const timestamp = new Date("2026-01-01T00:00:00.000Z");
+  const work: StreamItem[] = [
+    { kind: "user_message", id: "prompt", text: "Calculate", timestamp },
+    toolCall("one", "Bash"),
+  ];
+  const reply: StreamItem = {
+    kind: "assistant_message",
+    id: "reply",
+    text: "$$x^2$$",
+    timestamp,
+  };
+  const input = {
+    ...presentationOptions,
+    transform: installedTransform(installProbe("assistant_message", "phase === 'complete'")),
+    display: "collapseTurns" as const,
+  };
+  const present = createStreamPresentation();
+  const live = present({ ...input, tail: work, head: [reply], isTurnActive: true });
+  expect(rows(live).map((item) => item.id)).toEqual(["prompt", "tool-one", "reply"]);
+  const completed = { ...input, tail: [...work, reply], head: [], isTurnActive: false };
+  expect(present(completed).tail.map((item) => item.id)).toEqual([
+    "prompt",
+    "probe-plugin/reply/0",
+  ]);
+  const expandedTurnHostIds = new Set(["prompt"]);
+  const opened = present({ ...completed, expandedTurnHostIds });
+  expect(opened.tail.map((item) => item.id)).toEqual([
+    "prompt",
+    "tool-one",
+    "probe-plugin/reply/0",
+  ]);
+  expect(opened.turns.nestedItemIds.has("probe-plugin/reply/0")).toBe(false);
+  expect(present({ ...completed, expandedTurnHostIds }).tail).toEqual(opened.tail);
+  expect(present(completed).tail.map((item) => item.id)).toEqual([
+    "prompt",
+    "probe-plugin/reply/0",
+  ]);
+});
+
+it.each(["full", "collapseToolCalls"] as const)(
+  "leaves plugin replies and individual work visible in %s",
+  (display) => {
+    const timestamp = new Date("2026-01-01T00:00:00.000Z");
+    const tail: StreamItem[] = [
+      { kind: "user_message", id: "prompt", text: "Calculate", timestamp },
+      toolCall("one", "Bash"),
+      { kind: "assistant_message", id: "reply", text: "$$x^2$$", timestamp },
+    ];
+    const result = createStreamPresentation()({
+      ...presentationOptions,
+      transform: installedTransform(installProbe("assistant_message")),
+      tail,
+      head: [],
+      display,
+      isTurnActive: false,
+    });
+    expect(result.tail.map((item) => item.id)).toEqual([
+      "prompt",
+      "tool-one",
+      "probe-plugin/reply/0",
+    ]);
+    expect(result.turns.turnsByHostId.size).toBe(0);
+  },
+);
+
 it("opening a turn reveals individual calls rather than another tool group", () => {
   const present = createStreamPresentation();
   const timestamp = new Date();

packages/app/src/agent-stream/turn-collapse/model.test.ts

@@ -12,6 +12,21 @@ function reply(id: string): StreamItem {
   return { kind: "assistant_message", id, text: id, timestamp: at };
 }
 
+function pluginReply(sourceId: string, part = "0"): StreamItem {
+  return {
+    kind: "plugin",
+    id: `math/${sourceId}/${part}`,
+    sourceId,
+    sourceKind: "assistant_message",
+    timestamp: at,
+    pluginId: "math",
+    pluginItemId: `${sourceId}/${part}`,
+    itemKind: "formula",
+    version: 1,
+    data: { text: "$$x^2$$" },
+  };
+}
+
 function thought(id: string): StreamItem {
   return { kind: "thought", id, text: id, timestamp: at, status: "ready" };
 }
@@ -49,6 +64,44 @@ function collapse(tail: StreamItem[], options: Partial<CollapseTurnsInput> = {})
 const ids = (items: StreamItem[]) => items.map((item) => item.id);
 
 describe("collapseTurns", () => {
+  it("keeps every card of the final reply outside the folded turn", () => {
+    const items = [
+      user("prompt"),
+      reply("acknowledgment"),
+      shell("call"),
+      pluginReply("final", "left"),
+      pluginReply("final", "right"),
+    ];
+    const result = collapse(items);
+    expect(ids(result.tail)).toEqual(["prompt", "math/final/left", "math/final/right"]);
+    expect(result.hostIdByFoldedMessageId.get("acknowledgment")).toBe("prompt");
+    expect(result.hostIdByFoldedMessageId.has("final")).toBe(false);
+    const expanded = collapse(items, { expandedHostIds: new Set(["prompt"]) });
+    expect(expanded.tail).toEqual(items);
+    expect([...expanded.nestedItemIds]).toEqual(["acknowledgment", "call"]);
+  });
+
+  it("folds intermediate plugin replies even without tool calls", () => {
+    const result = collapse([user("prompt"), pluginReply("intermediate"), pluginReply("final")]);
+    expect(ids(result.tail)).toEqual(["prompt", "math/final/0"]);
+    expect(result.hostIdByFoldedMessageId.get("intermediate")).toBe("prompt");
+  });
+
+  it.each([undefined, "thought", "tool_call"] as const)(
+    "does not treat a plugin row with source kind %s as a reply",
+    (sourceKind) => {
+      const card = { ...pluginReply("card"), sourceKind };
+      const items = [user("prompt"), shell("call"), card];
+      expect(collapse(items).tail).toEqual(items);
+      expect(collapse(items).turnsByHostId.size).toBe(0);
+    },
+  );
+
+  it("keeps a plugin-rendered reply and its work unfolded while the turn is open", () => {
+    const items = [user("prompt"), shell("call"), pluginReply("final")];
+    expect(collapse(items, { isLatestTurnOpen: true }).tail).toEqual(items);
+  });
+
   // A history row renders again only when its item changes, so a prompt that is already retained
   // when its turn closes needs its row revised to show the summary line.
   it("reports the hosts whose summary line appeared, changed or went away", () => {

packages/app/src/agent-stream/turn-collapse/model.ts

@@ -100,9 +100,25 @@ function staysVisibleWhenCollapsed(item: StreamItem): boolean {
   return isTurnTrigger(item) || item.kind === "compaction";
 }
 
-function findLastAssistantIndex(tail: StreamItem[], start: number, end: number): number {
+function isAssistantReply(item: StreamItem): boolean {
+  return (
+    item.kind === "assistant_message" ||
+    (item.kind === "plugin" && item.sourceKind === "assistant_message")
+  );
+}
+
+function findLastReplyIndex(tail: StreamItem[], start: number, end: number): number {
   for (let index = end - 1; index >= start; index -= 1) {
-    if (tail[index]?.kind === "assistant_message") return index;
+    const item = tail[index]!;
+    if (!isAssistantReply(item)) continue;
+    const messageId = getStreamItemMessageId(item);
+    // A transformer can split one reply into several cards; keep the whole reply visible.
+    while (index > start) {
+      const previous = tail[index - 1]!;
+      if (!isAssistantReply(previous) || getStreamItemMessageId(previous) !== messageId) break;
+      index -= 1;
+    }
+    return index;
   }
   return -1;
 }
@@ -125,7 +141,7 @@ export function collapseTurns(input: CollapseTurnsInput): TurnCollapse {
   const flushRun = (runEnd: number) => {
     const host = tail[runStart - 1];
     const isClosed = runEnd < tail.length || !input.isLatestTurnOpen;
-    const replyIndex = findLastAssistantIndex(tail, runStart, runEnd);
+    const replyIndex = findLastReplyIndex(tail, runStart, runEnd);
     if (
       !host ||
       (host.kind !== "user_message" && host.kind !== "tool_call") ||
@@ -137,7 +153,7 @@ export function collapseTurns(input: CollapseTurnsInput): TurnCollapse {
     }
     const folded = tail.slice(runStart, replyIndex);
     const summary = summarizeTurn(folded, input.getGroupedCalls);
-    if (!hasWork(summary) && !folded.some((item) => item.kind === "assistant_message")) {
+    if (!hasWork(summary) && !folded.some(isAssistantReply)) {
       output.push(...tail.slice(runStart, runEnd));
       return;
     }

packages/app/src/plugins/timeline/projection.test.ts

@@ -98,6 +98,8 @@ describe("plugin timeline projection", () => {
       {
         kind: "plugin",
         id: "inline-thinking/thought-1/0",
+        sourceId: "thought-1",
+        sourceKind: "thought",
         data: { text: "First", phase: "streaming" },
       },
     ]);

packages/app/src/plugins/timeline/projection.ts

@@ -106,6 +106,7 @@ function transformSourceItem(
       kind: "plugin",
       id: `${pluginItem.pluginId}/${pluginItem.id}`,
       sourceId: item.id,
+      sourceKind: item.kind,
       timestamp: item.timestamp,
       pluginId: pluginItem.pluginId,
       pluginItemId: pluginItem.id,

packages/app/src/types/stream.ts

@@ -812,6 +812,8 @@ export interface PluginTimelineStreamItem {
   id: string;
   /** The id of the stream item a timeline transformer replaced; set only by projection. */
   sourceId?: string;
+  /** The original row kind, retained by projection for turn folding. */
+  sourceKind?: StreamItem["kind"];
   timelineCursor?: TimelinePosition;
   turnId?: string;
   timestamp: Date;