@assistant-ui/react-ag-ui 0.1.5 → 0.1.6

package/src/runtime/AgUiThreadController.ts

@@ -60,6 +60,7 @@ import {
   isPlainObject,
   RunAggregator,
   tryParseJSON,
+  type RetainedMessageState,
 } from "./adapter/run-aggregator";
 import {
   fromAgUiSnapshotMessages,
@@ -250,6 +251,17 @@ type CoreOptions = {
 
 const FALLBACK_USER_STATUS = { type: "complete", reason: "unknown" } as const;
 
+const snapshotRecordToolCallIds = (record: unknown): Set<string> => {
+  const toolCalls = (record as { toolCalls?: unknown } | undefined)?.toolCalls;
+  return new Set(
+    Array.isArray(toolCalls)
+      ? toolCalls
+          .map((call: { id?: unknown } | undefined) => call?.id)
+          .filter((id): id is string => typeof id === "string")
+      : [],
+  );
+};
+
 const textOf = (content: ThreadAssistantMessage["content"]) =>
   content.flatMap((part) => (part.type === "text" ? [part.text] : [])).join("");
 
@@ -1768,6 +1780,12 @@ export class AgUiThreadController {
     let settlingRun = false;
     let latestSnapshot: ThreadAssistantMessage | undefined;
     let snapshotToolCallIds = new Set<string>();
+    const snapshotResultIds = new Set<string>();
+    let announcingMove = false;
+    const retainedSnapshots = new Map<
+      string,
+      { message: ThreadAssistantMessage; toolCallIds: ReadonlySet<string> }
+    >();
     const ensureAssistant = (allowRecreate = false): string => {
       const cached = assistantMessageId;
       if (cached !== undefined && this.session.tryGetMessage(cached))
@@ -1834,6 +1852,8 @@ export class AgUiThreadController {
       }
       for (const [messageId, openToolCallIds] of heldMessages) {
         const existing = this.session.tryGetMessage(messageId);
+        if (!canWaitForTools && this.isTerminalStatus(existing?.message.status))
+          continue;
         if (existing)
           this.markPendingAssistantHistory(messageId, existing.parentId);
         let requiresToolResult = false;
@@ -2166,6 +2186,85 @@ export class AgUiThreadController {
       }
     };
 
+    // A message a snapshot moved the head past keeps streaming until nothing
+    // on it is open; held, it settles with the run.
+    const applyRetainedUpdate = (
+      update: ChatModelRunResult,
+      { messageId, streaming, openToolCallIds }: RetainedMessageState,
+    ) => {
+      const item = this.session.tryGetMessage(messageId);
+      if (item?.message.role !== "assistant") return;
+      this.markPendingAssistantHistory(messageId, item.parentId);
+      let snapshot = retainedSnapshots.get(messageId);
+      if (!snapshot) {
+        snapshot = {
+          message: item.message as ThreadAssistantMessage,
+          toolCallIds: new Set(
+            Array.from(
+              iterateToolCallParts(item.message.content),
+              (part) => part.toolCallId,
+            ),
+          ),
+        };
+        retainedSnapshots.set(messageId, snapshot);
+      }
+      const content =
+        update.content !== undefined
+          ? this.mergeSnapshotAssistantContent(
+              snapshot.message.content,
+              update.content as ThreadAssistantMessage["content"],
+              update.content.some((part) => part.type === "reasoning"),
+              update.content.some(isActivityPart),
+              snapshot.toolCallIds,
+            )
+          : update.content;
+      // A call only the snapshot carried stays open until a result arrives.
+      const unresolved = [
+        ...new Set([
+          ...openToolCallIds,
+          ...Array.from(iterateToolCallParts(content ?? item.message.content))
+            .filter(
+              (part) =>
+                snapshot.toolCallIds.has(part.toolCallId) &&
+                !isResolvedToolCall(part),
+            )
+            .map((part) => part.toolCallId),
+        ]),
+      ].filter((id) => !snapshotResultIds.has(id));
+      // A terminal event settles a held message with the run instead.
+      const held =
+        streaming ||
+        unresolved.length > 0 ||
+        (settlingRun && heldMessages.has(messageId));
+      if (held) heldMessages.set(messageId, new Set(unresolved));
+      else heldMessages.delete(messageId);
+      this.updateAssistantMessage(messageId, {
+        ...update,
+        ...(content !== undefined ? { content } : {}),
+        status: held
+          ? { type: "running" }
+          : { type: "complete", reason: "unknown" },
+      });
+      const current = this.session.tryGetMessage(messageId)?.message;
+      if (current?.role === "assistant")
+        retainedSnapshots.set(messageId, { ...snapshot, message: current });
+      // The run is still streaming, which the message it streams into shows
+      // as soon as any event arrives after the move.
+      const head =
+        assistantMessageId === undefined
+          ? undefined
+          : this.session.tryGetMessage(assistantMessageId)?.message;
+      if (
+        !settlingRun &&
+        !announcingMove &&
+        head?.role === "assistant" &&
+        head.status?.type !== "running" &&
+        !runStartMessageIds.has(head.id) &&
+        !aggregator.isRetained(head.id)
+      )
+        this.updateAssistantMessage(head.id, { status: { type: "running" } });
+    };
+
     const abortController = new AbortController();
     const abortSignal = abortController.signal;
     const aggregator = new RunAggregator({
@@ -2268,10 +2367,24 @@ export class AgUiThreadController {
             ? current
             : undefined;
       },
+      onRetainedMessage: applyRetainedUpdate,
+      canRetainMessage: (messageId) =>
+        messageId !== assistantMessageId &&
+        messageId !== this.session.headId &&
+        !runStartMessageIds.has(messageId) &&
+        this.session.tryGetMessage(messageId)?.message.role === "assistant",
       onServerMessageId: (serverId) => {
         adoptServerMessageId(serverId, false);
       },
       onTextMessageStart: (serverId, openToolCallIds) => {
+        const outgoingId =
+          assistantMessageId !== undefined &&
+          !assistantCollided &&
+          !runStartMessageIds.has(assistantMessageId) &&
+          !replayContexts.has(assistantMessageId) &&
+          this.session.hasMessage(assistantMessageId)
+            ? assistantMessageId
+            : undefined;
         const owners = this.getTrailingAssistantTurn().filter((message) =>
           replayContexts.has(message.id),
         );
@@ -2293,6 +2406,7 @@ export class AgUiThreadController {
         }
         frameOwnerIds.clear();
         adoptServerMessageId(serverId, true, frameOpenIds);
+        return outgoingId;
       },
     });
     const settleRunSubagents = (status: MessageStatus) => {
@@ -2345,6 +2459,17 @@ export class AgUiThreadController {
           openToolCallIds.delete(event.toolCallId);
         }
       }
+      // A result a snapshot delivers answers its call the way a streamed
+      // TOOL_CALL_RESULT does.
+      if (event.type === "MESSAGES_SNAPSHOT") {
+        for (const record of event.messages) {
+          if (record?.role !== "tool" || typeof record.toolCallId !== "string")
+            continue;
+          snapshotResultIds.add(record.toolCallId);
+          for (const openToolCallIds of heldMessages.values())
+            openToolCallIds.delete(record.toolCallId);
+        }
+      }
       settlingRun =
         event.type === "RUN_FINISHED" ||
         event.type === "RUN_ERROR" ||
@@ -2360,6 +2485,7 @@ export class AgUiThreadController {
           joinStrategy,
           runStartMessageIds,
           hasRunUserBoundary,
+          heldMessages,
           candidateContexts.size ? reconcileReplaySnapshot : undefined,
         );
         if (event.type === "MESSAGES_SNAPSHOT") {
@@ -2406,18 +2532,31 @@ export class AgUiThreadController {
           );
           latestSnapshot =
             imported?.role === "assistant" && source ? imported : undefined;
-          snapshotToolCallIds = new Set(
-            Array.isArray(source?.toolCalls)
-              ? source.toolCalls
-                  .map((call: { id?: unknown }) => call.id)
-                  .filter((id: unknown): id is string => typeof id === "string")
-              : [],
-          );
+          snapshotToolCallIds = snapshotRecordToolCallIds(source);
           if (!replayEnded && assistantMessageId === undefined) {
             const tail = this.session.getMessages().at(-1);
             snapshotAssistantId =
               tail?.role === "assistant" ? tail.id : undefined;
           }
+          for (const record of event.messages) {
+            if (
+              record?.role !== "assistant" ||
+              !aggregator.isRetained(record.id)
+            )
+              continue;
+            const message = this.session.tryGetMessage(record.id)?.message;
+            if (message?.role !== "assistant") continue;
+            retainedSnapshots.set(record.id, {
+              message,
+              toolCallIds: snapshotRecordToolCallIds(record),
+            });
+          }
+          announcingMove = true;
+          try {
+            aggregator.announceRetained();
+          } finally {
+            announcingMove = false;
+          }
         }
         if (event.type === "RUN_ERROR") {
           settleRunSubagents({
@@ -2516,10 +2655,11 @@ export class AgUiThreadController {
       const current = assistantMessageId
         ? this.session.tryGetMessage(assistantMessageId)?.message
         : undefined;
+      const currentRunning =
+        current?.role === "assistant" && current.status?.type === "running";
       if (
         this.abortController === abortController &&
-        current?.role === "assistant" &&
-        current.status?.type === "running"
+        (currentRunning || heldMessages.size > 0)
       ) {
         persistRunAssistants(true);
         // Finalization is transport cleanup, not proof of a terminal event.
@@ -2529,10 +2669,11 @@ export class AgUiThreadController {
           error: "The connection ended before the answer finished.",
         };
         settleHeldMessages(status);
-        applyUpdate({
-          status,
-          content: settleNestedContent(current.content, status),
-        });
+        if (currentRunning)
+          applyUpdate({
+            status,
+            content: settleNestedContent(current.content, status),
+          });
       }
       if (this.abortController === abortController && !runFinished) {
         // The run ended without RUN_FINISHED, so no sub-agent call it left
@@ -3371,6 +3512,7 @@ export class AgUiThreadController {
     joinStrategy: UseAgUiRuntimeOptions["joinStrategy"],
     runStartMessageIds: ReadonlySet<string>,
     hasRunUserBoundary: boolean,
+    heldMessages: ReadonlyMap<string, unknown>,
     reconcileSnapshot?: SnapshotMessageReconciler,
   ): string | undefined {
     if (
@@ -3426,6 +3568,7 @@ export class AgUiThreadController {
           joinStrategy,
           runStartMessageIds,
           hasRunUserBoundary,
+          heldMessages,
           reconcileSnapshot,
         );
         const headId = this.session.headId;
@@ -3444,17 +3587,22 @@ export class AgUiThreadController {
         ) {
           const nextActiveAssistantId =
             head?.role === "assistant" ? (headId ?? undefined) : undefined;
-          if (
-            activeAssistantId !== nextActiveAssistantId &&
-            activeAssistantId
-          ) {
+          const previousActiveId =
+            activeAssistantId !== nextActiveAssistantId
+              ? activeAssistantId
+              : undefined;
+          const retained =
+            nextActiveAssistantId !== undefined &&
+            joinStrategy !== "concat-content" &&
+            aggregator.rebaseMessage(nextActiveAssistantId, previousActiveId);
+          if (previousActiveId && !retained) {
             const previous =
-              this.session.tryGetMessage(activeAssistantId)?.message;
+              this.session.tryGetMessage(previousActiveId)?.message;
             if (
               previous?.role === "assistant" &&
               previous.status?.type === "running"
             ) {
-              this.updateAssistantMessage(activeAssistantId, {
+              this.updateAssistantMessage(previousActiveId, {
                 status: FALLBACK_USER_STATUS,
                 content: settleNestedContent(
                   previous.content,
@@ -3464,14 +3612,15 @@ export class AgUiThreadController {
             }
           }
           activeAssistantId = nextActiveAssistantId;
-          if (activeAssistantId && joinStrategy !== "concat-content") {
-            aggregator.rebaseMessage(activeAssistantId);
-          }
-          if (hasNewPendingToolCallAnswer && head) {
-            this.updateAssistantMessage(head.id, {
-              status: { type: "running" },
-            });
-          }
+        }
+        if (
+          hasNewPendingToolCallAnswer &&
+          head &&
+          head.id === activeAssistantId
+        ) {
+          this.updateAssistantMessage(head.id, {
+            status: { type: "running" },
+          });
         }
         this.updateActiveAssistantState(activeAssistantId);
         break;
@@ -3698,6 +3847,7 @@ export class AgUiThreadController {
     joinStrategy: UseAgUiRuntimeOptions["joinStrategy"],
     runStartMessageIds: ReadonlySet<string>,
     hasRunUserBoundary: boolean,
+    heldMessages: ReadonlyMap<string, unknown>,
     reconcileSnapshot?: SnapshotMessageReconciler,
   ) {
     try {
@@ -3754,7 +3904,8 @@ export class AgUiThreadController {
           const snapshot =
             reconcileSnapshot?.(convertedMessage, previous) ?? convertedMessage;
           const preservesRunningStatus =
-            convertedMessage.id === activeAssistantId &&
+            (convertedMessage.id === activeAssistantId ||
+              heldMessages.has(convertedMessage.id)) &&
             previous?.status?.type === "running" &&
             snapshot.status?.type === "requires-action" &&
             snapshot.status.reason === "tool-calls";

package/src/runtime/adapter/run-aggregator.routing.test.ts

@@ -0,0 +1,359 @@
+import { describe, expect, it, vi } from "vitest";
+import type { ChatModelRunResult } from "@assistant-ui/core";
+import type { JoinStrategy } from "@assistant-ui/core/react";
+import { RunAggregator, type RetainedMessageState } from "./run-aggregator";
+import type { AgUiEvent } from "../types";
+
+const setup = (joinStrategy?: JoinStrategy) => {
+  const current: ChatModelRunResult[] = [];
+  const retained: {
+    update: ChatModelRunResult;
+    state: RetainedMessageState;
+  }[] = [];
+  const aggregator = new RunAggregator({
+    showThinking: true,
+    joinStrategy,
+    logger: { debug: vi.fn(), error: vi.fn() },
+    emit: (update) => current.push(update),
+    onRetainedMessage: (update, state) => retained.push({ update, state }),
+    onTextMessageStart: () => {},
+  });
+  const send = (...events: AgUiEvent[]) => {
+    for (const event of events) aggregator.handle(event);
+  };
+  send(
+    { type: "RUN_STARTED", runId: "run" },
+    { type: "TEXT_MESSAGE_START", messageId: "answer" },
+    { type: "TEXT_MESSAGE_CONTENT", messageId: "answer", delta: "a1" },
+    {
+      type: "TOOL_CALL_START",
+      toolCallId: "task",
+      toolCallName: "task",
+      parentMessageId: "answer",
+    },
+    {
+      type: "SUBAGENT_STARTED",
+      subagentRunId: "sub",
+      name: "worker",
+      parentToolCallId: "task",
+    },
+  );
+  aggregator.rebaseMessage("other", "answer");
+  aggregator.announceRetained();
+  const lastRetained = () => retained.at(-1);
+  const lastCurrent = () => current.at(-1)?.content ?? [];
+  const retainedTask = () =>
+    lastRetained()?.update.content?.find(
+      (part) => part.type === "tool-call" && part.toolCallId === "task",
+    );
+  return {
+    aggregator,
+    send,
+    current,
+    retained,
+    lastRetained,
+    lastCurrent,
+    retainedTask,
+  };
+};
+
+describe("RunAggregator routing after the current message moves", () => {
+  it("announces the moved-off message once with what is still open", () => {
+    const { aggregator, retained } = setup();
+    aggregator.announceRetained();
+
+    expect(retained).toHaveLength(1);
+    expect(retained[0]?.state).toEqual({
+      messageId: "answer",
+      streaming: true,
+      openToolCallIds: ["task"],
+    });
+    expect(aggregator.isRetained("answer")).toBe(true);
+    expect(aggregator.hasToolCall("task")).toBe(true);
+  });
+
+  it("routes root text naming the moved-off message to it", () => {
+    const { send, lastRetained, current } = setup();
+    const emitted = current.length;
+    send({ type: "TEXT_MESSAGE_CONTENT", messageId: "answer", delta: "a2" });
+
+    expect(lastRetained()?.update.content?.[0]).toEqual({
+      type: "text",
+      text: "a1a2",
+    });
+    expect(current).toHaveLength(emitted);
+  });
+
+  it("closes the text frame on TEXT_MESSAGE_END", () => {
+    const { send, lastRetained } = setup();
+    send({ type: "TEXT_MESSAGE_END", messageId: "answer" });
+
+    expect(lastRetained()?.state.streaming).toBe(true);
+    send(
+      { type: "SUBAGENT_FINISHED", subagentRunId: "sub" },
+      {
+        type: "TOOL_CALL_RESULT",
+        messageId: "tool",
+        toolCallId: "task",
+        content: "done",
+      },
+    );
+    expect(lastRetained()?.state).toEqual({
+      messageId: "answer",
+      streaming: false,
+      openToolCallIds: [],
+    });
+  });
+
+  it("routes TOOL_CALL_START by parentMessageId, and to the current message without one", () => {
+    const { send, lastRetained, lastCurrent } = setup();
+    send(
+      {
+        type: "TOOL_CALL_START",
+        toolCallId: "late",
+        toolCallName: "late",
+        parentMessageId: "answer",
+      },
+      { type: "TOOL_CALL_START", toolCallId: "next", toolCallName: "next" },
+    );
+
+    expect(lastRetained()?.state.openToolCallIds).toEqual(["task", "late"]);
+    expect(
+      lastCurrent().filter((part) => part.type === "tool-call"),
+    ).toMatchObject([{ toolCallId: "next" }]);
+  });
+
+  it("routes TOOL_CALL_ARGS, END, and RESULT by the call they name", () => {
+    const { send, retainedTask, current } = setup();
+    const emitted = current.length;
+    send(
+      { type: "TOOL_CALL_ARGS", toolCallId: "task", delta: '{"q":1}' },
+      { type: "TOOL_CALL_END", toolCallId: "task" },
+      {
+        type: "TOOL_CALL_RESULT",
+        messageId: "tool",
+        toolCallId: "task",
+        content: "done",
+      },
+    );
+
+    expect(retainedTask()).toMatchObject({ args: { q: 1 }, result: "done" });
+    expect(current).toHaveLength(emitted);
+  });
+
+  it("routes sub-agent events by run id and nests a new run under its call", () => {
+    const { send, retainedTask, current } = setup();
+    const emitted = current.length;
+    send(
+      {
+        type: "TEXT_MESSAGE_CONTENT",
+        messageId: "s1",
+        delta: "working",
+        subagentRunId: "sub",
+      },
+      {
+        type: "SUBAGENT_STARTED",
+        subagentRunId: "sub2",
+        name: "helper",
+        parentToolCallId: "task",
+      },
+      { type: "SUBAGENT_ERROR", subagentRunId: "sub2", message: "boom" },
+      { type: "SUBAGENT_FINISHED", subagentRunId: "sub" },
+    );
+
+    const task = retainedTask();
+    expect(task?.type === "tool-call" && task.messages).toMatchObject([
+      {
+        id: "sub",
+        status: { type: "complete" },
+        content: [{ type: "text", text: "working" }],
+      },
+      { id: "sub2", status: { type: "incomplete", reason: "error" } },
+    ]);
+    expect(current).toHaveLength(emitted);
+  });
+
+  it("routes an MCP app snapshot by its tool call", () => {
+    const { send, retainedTask } = setup();
+    send({
+      type: "ACTIVITY_SNAPSHOT",
+      messageId: "act",
+      activityType: "mcp-apps",
+      content: { toolCallId: "task", resourceUri: "ui://app/view" },
+    } as AgUiEvent);
+
+    expect(retainedTask()).toMatchObject({
+      mcp: { app: { resourceUri: "ui://app/view" } },
+    });
+  });
+
+  it("routes reasoning and activity deltas to the message that holds them", () => {
+    const { aggregator, send, lastRetained, current } = setup();
+    aggregator.rebaseMessage("answer", "other");
+    send(
+      { type: "REASONING_MESSAGE_START", messageId: "r" },
+      { type: "REASONING_MESSAGE_CONTENT", messageId: "r", delta: "think" },
+      {
+        type: "ACTIVITY_SNAPSHOT",
+        messageId: "act",
+        activityType: "progress",
+        content: { step: 1 },
+      } as AgUiEvent,
+    );
+    aggregator.rebaseMessage("other", "answer");
+    const emitted = current.length;
+    send({ type: "REASONING_MESSAGE_CONTENT", messageId: "r", delta: "ing" }, {
+      type: "ACTIVITY_DELTA",
+      messageId: "act",
+      activityType: "progress",
+      patch: [{ op: "replace", path: "/step", value: 2 }],
+    } as AgUiEvent);
+
+    expect(current).toHaveLength(emitted);
+    expect(lastRetained()?.update.content).toEqual(
+      expect.arrayContaining([
+        expect.objectContaining({ type: "reasoning", text: "thinking" }),
+        expect.objectContaining({ type: "data", data: { step: 2 } }),
+      ]),
+    );
+  });
+
+  it("replaces an activity snapshot on the message that holds it", () => {
+    const { aggregator, send, lastRetained, current } = setup();
+    aggregator.rebaseMessage("answer", "other");
+    send({
+      type: "ACTIVITY_SNAPSHOT",
+      messageId: "act",
+      activityType: "progress",
+      content: { step: 1 },
+    } as AgUiEvent);
+    aggregator.rebaseMessage("other", "answer");
+    const emitted = current.length;
+    send({
+      type: "ACTIVITY_SNAPSHOT",
+      messageId: "act",
+      activityType: "progress",
+      content: { step: 2 },
+    } as AgUiEvent);
+
+    expect(
+      lastRetained()?.update.content?.filter((part) => part.type === "data"),
+    ).toEqual([expect.objectContaining({ data: { step: 2 } })]);
+    expect(current).toHaveLength(emitted);
+  });
+
+  it("keeps parsing args a snapshot started", () => {
+    const { send, retainedTask } = setup();
+    send(
+      {
+        type: "MESSAGES_SNAPSHOT",
+        messages: [
+          {
+            id: "answer",
+            role: "assistant",
+            content: "a1",
+            toolCalls: [
+              {
+                id: "task",
+                type: "function",
+                function: { name: "task", arguments: '{"q":' },
+              },
+            ],
+          },
+        ],
+      },
+      { type: "TOOL_CALL_ARGS", toolCallId: "task", delta: '"x"}' },
+    );
+
+    expect(retainedTask()).toMatchObject({
+      argsText: '{"q":"x"}',
+      args: { q: "x" },
+    });
+  });
+
+  it("keeps root reasoning and custom events on the current message", () => {
+    const { send, lastCurrent, retained } = setup();
+    const announced = retained.length;
+    send(
+      { type: "REASONING_MESSAGE_START", messageId: "r" },
+      { type: "REASONING_MESSAGE_CONTENT", messageId: "r", delta: "hmm" },
+      { type: "CUSTOM", name: "note", value: 1 },
+    );
+
+    expect(lastCurrent().map((part) => part.type)).toEqual([
+      "reasoning",
+      "data",
+    ]);
+    expect(retained).toHaveLength(announced);
+  });
+
+  it("applies MESSAGES_SNAPSHOT text and args to every message", () => {
+    const { send, retainedTask, lastRetained } = setup();
+    send(
+      {
+        type: "MESSAGES_SNAPSHOT",
+        messages: [
+          {
+            id: "answer",
+            role: "assistant",
+            content: "a1 fixed",
+            toolCalls: [
+              {
+                id: "task",
+                type: "function",
+                function: { name: "task", arguments: '{"q":2}' },
+              },
+            ],
+          },
+        ],
+      },
+      { type: "TEXT_MESSAGE_CONTENT", messageId: "answer", delta: "!" },
+    );
+
+    expect(lastRetained()?.update.content?.[0]).toEqual({
+      type: "text",
+      text: "a1 fixed!",
+    });
+    expect(retainedTask()).toMatchObject({ argsText: '{"q":2}' });
+  });
+
+  it("settles running sub-agents on moved-off messages before the current emit", () => {
+    const { send, retained, current, retainedTask } = setup();
+    send({ type: "RUN_ERROR", message: "boom" });
+
+    const task = retainedTask();
+    expect(task?.type === "tool-call" && task.messages?.[0]?.status).toEqual({
+      type: "incomplete",
+      reason: "error",
+      error: "boom",
+    });
+    expect(retained.at(-1)?.update.status).toBeUndefined();
+    expect(current.at(-1)?.status).toMatchObject({ type: "incomplete" });
+  });
+
+  it("resumes a moved-off message that becomes current again", () => {
+    const { aggregator, send, lastCurrent } = setup();
+    aggregator.rebaseMessage("answer", "other");
+    send({ type: "TEXT_MESSAGE_CONTENT", messageId: "answer", delta: "a2" });
+
+    expect(lastCurrent()[0]).toEqual({ type: "text", text: "a1a2" });
+    expect(aggregator.isRetained("other")).toBe(true);
+  });
+
+  it("forgets moved-off messages when a new run starts", () => {
+    const { aggregator, send } = setup();
+    send({ type: "RUN_STARTED", runId: "next" });
+
+    expect(aggregator.isRetained("answer")).toBe(false);
+    expect(aggregator.hasToolCall("task")).toBe(false);
+  });
+
+  it("control: concat-content routes everything to the current message", () => {
+    const { send, lastCurrent, retained } = setup("concat-content");
+    const announced = retained.length;
+    send({ type: "TEXT_MESSAGE_CONTENT", messageId: "answer", delta: "a2" });
+
+    expect(retained).toHaveLength(announced);
+    expect(lastCurrent()).toContainEqual({ type: "text", text: "a2" });
+  });
+});