@assistant-ui/react-ag-ui 0.1.7 → 0.1.9

package/src/runtime/adapter/conversions.subagent-messages.test.ts

@@ -407,6 +407,123 @@ describe("fromAgUiSnapshotMessages subagent attribution", () => {
     ]);
   });
 
+  it("retains each nested record's pending-call admission when adjacent records are joined", () => {
+    const converted = fromAgUiSnapshotMessages(
+      [
+        { id: "a", role: "assistant", content: "", toolCalls: [call("spawn")] },
+        { id: "t", role: "tool", toolCallId: "spawn", content: '"ok"' },
+        {
+          id: "s1",
+          role: "assistant",
+          content: "",
+          toolCalls: [call("one"), call("excluded")],
+          subagentRunId: "run",
+          metadata: {
+            custom: { agui: { pendingToolCallIds: ["one", "future"] } },
+          },
+        },
+        {
+          id: "s2",
+          role: "assistant",
+          content: "",
+          toolCalls: [call("two")],
+          subagentRunId: "run",
+          metadata: { custom: { agui: { pendingToolCallIds: ["two"] } } },
+        },
+      ],
+      { joinAdjacentAssistantMessages: true },
+      new Map([["run", "spawn"]]),
+    ) as (Message & {
+      metadata?: { custom?: { agui?: { pendingToolCallIds?: string[] } } };
+    })[];
+    const nested = findCall(converted, "spawn")?.messages?.[0] as Message & {
+      metadata?: { custom?: { agui?: { pendingToolCallIds?: string[] } } };
+    };
+    expect(findCall(converted, "spawn")?.messages).toHaveLength(1);
+    expect(nested.metadata?.custom?.agui?.pendingToolCallIds).toEqual([
+      "one",
+      "future",
+      "two",
+    ]);
+    expect(converted[0]?.metadata?.custom?.agui?.pendingToolCallIds).toEqual([
+      "one",
+      "future",
+      "two",
+    ]);
+  });
+
+  it.each([false, true])(
+    "does not merge pending lists while folding a nested run when joinAdjacentAssistantMessages is %s",
+    (joinAdjacentAssistantMessages) => {
+      // An empty record does not open a turn, so the adjacent-record join
+      // leaves both records for the subagent fold.
+      const converted = fromAgUiSnapshotMessages(
+        [
+          {
+            id: "a",
+            role: "assistant",
+            content: "",
+            toolCalls: [call("spawn")],
+          },
+          { id: "t", role: "tool", toolCallId: "spawn", content: '"ok"' },
+          {
+            id: "s1",
+            role: "assistant",
+            content: "",
+            subagentRunId: "run",
+            metadata: {
+              custom: {
+                agui: {
+                  pendingToolCallIds: ["stale-a"],
+                  interrupts: [{ id: "int-a", reason: "external" }],
+                },
+              },
+            },
+          },
+          {
+            id: "s2",
+            role: "assistant",
+            content: "done",
+            subagentRunId: "run",
+            metadata: {
+              custom: { agui: { pendingToolCallIds: ["stale-b"] } },
+            },
+          },
+        ],
+        { joinAdjacentAssistantMessages },
+        new Map([["run", "spawn"]]),
+      ) as (Message & {
+        metadata?: {
+          custom?: {
+            agui?: {
+              pendingToolCallIds?: string[];
+              interrupts?: { id: string }[];
+            };
+          };
+        };
+      })[];
+      const nested = findCall(converted, "spawn")?.messages?.[0] as Message & {
+        metadata?: {
+          custom?: {
+            agui?: {
+              pendingToolCallIds?: string[];
+              interrupts?: { id: string }[];
+            };
+          };
+        };
+      };
+      expect(texts(nested)).toEqual(["done"]);
+      expect(
+        nested.metadata?.custom?.agui?.interrupts?.map(
+          (interrupt) => interrupt.id,
+        ),
+      ).toEqual(["int-a"]);
+      expect(nested.metadata?.custom?.agui?.pendingToolCallIds).toEqual([
+        "stale-b",
+      ]);
+    },
+  );
+
   it("keeps a deeply nested pending call visible through resolved spawning calls", () => {
     const messages = [
       { id: "a1", role: "assistant", content: "", toolCalls: [call("top")] },

package/src/runtime/adapter/conversions.ts

@@ -1078,6 +1078,7 @@ function readNamespacedMetadata(metadata: unknown): Record<string, unknown> {
 function foldTurnRecords(
   previous: CoreThreadMessageLike,
   message: CoreThreadMessageLike,
+  mergePendingToolCallIds: boolean,
 ): CoreThreadMessageLike {
   const interrupts = [
     ...(readPersistedInterrupts(previous.metadata) ?? []),
@@ -1095,10 +1096,21 @@ function foldTurnRecords(
       trailsCalls ? { ...entry, after: true } : entry,
     ),
   ];
+  const previousPending = readPersistedPendingToolCallIds(previous.metadata);
+  const messagePending = readPersistedPendingToolCallIds(message.metadata);
+  const pendingToolCallIds =
+    !mergePendingToolCallIds ||
+    (previousPending === undefined && messagePending === undefined)
+      ? []
+      : [
+          ...admittedUnresolvedToolCallIds(previous, false),
+          ...admittedUnresolvedToolCallIds(message, false),
+        ];
   const namespaced = {
     ...readNamespacedMetadata(previous.metadata),
     ...readNamespacedMetadata(message.metadata),
     ...(interrupts.length > 0 ? { interrupts } : {}),
+    ...(pendingToolCallIds.length > 0 ? { pendingToolCallIds } : {}),
     ...(opaqueReasoning.length > 0 ? { opaqueReasoning } : {}),
   } satisfies AgUiCustomMetadata & Record<string, unknown>;
   const custom = {
@@ -1134,6 +1146,16 @@ export type FromAgUiMessagesOptions = {
    * Matches the `showThinking` option of `useAgUiRuntime`. Defaults to `true`.
    */
   showThinking?: boolean;
+  /**
+   * Whether to join every run of adjacent assistant records into one message,
+   * for a backend whose history splits one turn across records with ids of
+   * their own. The merged message keeps the id of the record carrying the
+   * answer, with the parts in wire order. Records that belong to separate runs
+   * are joined too when no user message separates them. Matches the
+   * `joinAdjacentAssistantMessages` option of `useAgUiRuntime`. Defaults to
+   * `false`.
+   */
+  joinAdjacentAssistantMessages?: boolean;
 };
 
 const readSubagentRunId = (record: Record<string, unknown>) => {
@@ -1259,7 +1281,9 @@ export function fromAgUiSnapshotMessages(
       )
       .reduce<CoreThreadMessageLike | undefined>(
         (merged, message) =>
-          merged ? foldTurnRecords(merged, message) : message,
+          // Pending-call admission is taken from the original records below.
+          // Merging lists in this fold would union ids that admission skips.
+          merged ? foldTurnRecords(merged, message, false) : message,
         undefined,
       );
     const content = Array.isArray(folded?.content) ? folded.content : [];
@@ -1418,6 +1442,8 @@ function convertAgUiMessages(
   options?: FromAgUiMessagesOptions,
 ): CoreThreadMessageLike[] {
   const showThinking = options?.showThinking ?? true;
+  const joinAdjacentAssistantMessages =
+    options?.joinAdjacentAssistantMessages ?? false;
   const converted: CoreThreadMessageLike[] = [];
   const a2uiBuckets = new Map<string, A2uiState>();
   const a2uiBucketOwnerIndices = new Map<string, number>();
@@ -1779,7 +1805,9 @@ function convertAgUiMessages(
   // a container beside the records around it, which the live run renders as
   // one message: the container joins the assistant record ahead of it, and
   // what follows a container joins it until the turn carries text, where a
-  // further text record opens a message of its own as it does live.
+  // further text record opens a message of its own as it does live. An app
+  // that knows its backend splits a turn across records it named itself opts
+  // into joining every assistant record that follows one opening a turn.
   const folded: CoreThreadMessageLike[] = [];
   for (const message of converted) {
     const previous = folded[folded.length - 1];
@@ -1787,10 +1815,17 @@ function convertAgUiMessages(
       previous !== undefined &&
       (isSyntheticToolCallContainer(message)
         ? opensTurn(previous)
-        : isSyntheticToolCallContainer(previous) &&
-          answersToolCallContainer(message))
+        : (isSyntheticToolCallContainer(previous) &&
+            answersToolCallContainer(message)) ||
+          (joinAdjacentAssistantMessages &&
+            message.role === "assistant" &&
+            opensTurn(previous)))
     ) {
-      folded[folded.length - 1] = foldTurnRecords(previous, message);
+      folded[folded.length - 1] = foldTurnRecords(
+        previous,
+        message,
+        joinAdjacentAssistantMessages,
+      );
       continue;
     }
     folded.push(message);

package/src/runtime/types.ts

@@ -94,6 +94,12 @@ export type UseAgUiRuntimeOptions = ExternalStoreSharedOptions & {
   logger?: Partial<Logger>;
   showThinking?: boolean;
   joinStrategy?: JoinStrategy;
+  /**
+   * Join adjacent assistant records of an imported snapshot into one message.
+   * See `FromAgUiMessagesOptions.joinAdjacentAssistantMessages`. Defaults to
+   * `false`.
+   */
+  joinAdjacentAssistantMessages?: boolean;
   /**
    * Return context entries for each run. `messages` is the exact AG-UI
    * transcript sent in the run input, after transcript selection and message

package/src/useAgUiRuntime.connect.test.tsx

@@ -0,0 +1,169 @@
+// @vitest-environment jsdom
+
+import { useState } from "react";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { act, cleanup, renderHook, waitFor } from "@testing-library/react";
+import {
+  HttpAgent,
+  runHttpRequest,
+  transformHttpEventStream,
+  type RunAgentInput,
+} from "@ag-ui/client";
+import type { ThreadMessage } from "@assistant-ui/core";
+import { useAgUiRuntime } from "./useAgUiRuntime";
+
+const encoder = new TextEncoder();
+const frame = (event: object) =>
+  encoder.encode(`data: ${JSON.stringify(event)}\n\n`);
+
+// Each thread's run stays open until the request is aborted.
+const openRun = (threadId: string, init?: RequestInit) =>
+  new Response(
+    new ReadableStream<Uint8Array>({
+      start(controller) {
+        controller.enqueue(
+          frame({ type: "RUN_STARTED", threadId, runId: `run-${threadId}` }),
+        );
+        controller.enqueue(
+          frame({
+            type: "TEXT_MESSAGE_START",
+            messageId: `answer-${threadId}`,
+            role: "assistant",
+          }),
+        );
+        controller.enqueue(
+          frame({
+            type: "TEXT_MESSAGE_CONTENT",
+            messageId: `answer-${threadId}`,
+            delta: `streaming ${threadId}`,
+          }),
+        );
+        init?.signal?.addEventListener("abort", () => controller.close(), {
+          once: true,
+        });
+      },
+    }),
+    { headers: { "Content-Type": "text/event-stream" } },
+  );
+
+class ReconnectingAgent extends HttpAgent {
+  readonly connected: string[] = [];
+
+  override connectAgent(...args: Parameters<HttpAgent["connectAgent"]>) {
+    this.abortController = new AbortController();
+    return super.connectAgent(...args);
+  }
+
+  protected override connect(input: RunAgentInput) {
+    this.connected.push(input.threadId);
+    return transformHttpEventStream(
+      runHttpRequest(() =>
+        this.fetch(`${this.url}/${input.threadId}/stream`, {
+          signal: this.abortController.signal,
+        }),
+      ),
+    );
+  }
+}
+
+const question = (id: string): ThreadMessage => ({
+  id: `question-${id}`,
+  role: "user",
+  content: [{ type: "text", text: `asked in ${id}` }],
+  attachments: [],
+  createdAt: new Date(0),
+  metadata: { custom: {} },
+});
+
+afterEach(() => {
+  cleanup();
+});
+
+function renderDurable() {
+  const agent = new ReconnectingAgent({
+    url: "https://agent.test",
+    threadId: "initial",
+    fetch: async (url: string, init?: RequestInit) =>
+      openRun(new URL(url).pathname.split("/")[1]!, init),
+  });
+  const abortRun = vi.spyOn(agent, "abortRun");
+  const detachActiveRun = vi.spyOn(agent, "detachActiveRun");
+  const view = renderHook(() => {
+    const [threadId, setThreadId] = useState("initial");
+    return useAgUiRuntime({
+      agent,
+      adapters: {
+        threadList: {
+          threadId,
+          onSwitchToThread: async (id) => {
+            agent.threadId = id;
+            setThreadId(id);
+            return { messages: [question(id)], unstable_resume: true };
+          },
+          onSwitchToNewThread: () => {
+            agent.threadId = "thread-new";
+            setThreadId("thread-new");
+          },
+        },
+      },
+    });
+  });
+  return { agent, abortRun, detachActiveRun, ...view };
+}
+
+const lastText = (messages: readonly ThreadMessage[]) => {
+  const part = messages.at(-1)?.content[0];
+  return part?.type === "text" ? part.text : undefined;
+};
+
+describe("useAgUiRuntime with an agent that implements connect()", () => {
+  it("leaves a running thread by detaching, then rejoins the next thread's run", async () => {
+    const { agent, abortRun, detachActiveRun, result } = renderDurable();
+
+    await act(async () => {
+      await result.current.threads.switchToThread("thread-a");
+    });
+    await waitFor(() =>
+      expect(lastText(result.current.thread.getState().messages)).toBe(
+        "streaming thread-a",
+      ),
+    );
+
+    await act(async () => {
+      await result.current.threads.switchToThread("thread-b");
+    });
+
+    await waitFor(() =>
+      expect(lastText(result.current.thread.getState().messages)).toBe(
+        "streaming thread-b",
+      ),
+    );
+    expect(agent.connected).toEqual(["thread-a", "thread-b"]);
+    expect(
+      result.current.thread.getState().messages.map((message) => message.id),
+    ).toEqual(["question-thread-b", "answer-thread-b"]);
+    expect(detachActiveRun).toHaveBeenCalled();
+    expect(abortRun).not.toHaveBeenCalled();
+
+    await act(async () => {
+      await result.current.threads.switchToNewThread();
+    });
+    expect(abortRun).not.toHaveBeenCalled();
+  });
+
+  it("aborts the run when the user cancels", async () => {
+    const { abortRun, result } = renderDurable();
+    await act(async () => {
+      await result.current.threads.switchToThread("thread-a");
+    });
+    await waitFor(() =>
+      expect(result.current.thread.getState().isRunning).toBe(true),
+    );
+
+    act(() => {
+      result.current.thread.cancelRun();
+    });
+
+    expect(abortRun).toHaveBeenCalledTimes(1);
+  });
+});

package/src/useAgUiRuntime.external-store.test.tsx

@@ -1,7 +1,7 @@
 // @vitest-environment jsdom
 
 import { act, cleanup, renderHook } from "@testing-library/react";
-import type { HttpAgent } from "@ag-ui/client";
+import { HttpAgent } from "@ag-ui/client";
 import { afterEach, expect, it, vi } from "vitest";
 import { useAgUiRuntime } from "./useAgUiRuntime";
 
@@ -88,3 +88,48 @@ it("passes joinStrategy to the active run", async () => {
     { type: "text", text: "B" },
   ]);
 });
+
+it("batches text deltas emitted from one response before re-rendering the host", async () => {
+  const deltas = Array.from({ length: 20 }, (_, index) => String(index));
+  const events = [
+    { type: "RUN_STARTED", threadId: "main", runId: "run" },
+    { type: "TEXT_MESSAGE_START", messageId: "answer" },
+    ...deltas.map((delta) => ({
+      type: "TEXT_MESSAGE_CONTENT",
+      messageId: "answer",
+      delta,
+    })),
+    { type: "TEXT_MESSAGE_END", messageId: "answer" },
+    { type: "RUN_FINISHED", threadId: "main", runId: "run" },
+  ];
+  const agent = new HttpAgent({
+    url: "http://agent.test/agent",
+    fetch: async () =>
+      new Response(
+        events.map((event) => `data: ${JSON.stringify(event)}\n\n`).join(""),
+        { headers: { "content-type": "text/event-stream" } },
+      ),
+  });
+  let renders = 0;
+  const { result } = renderHook(() => {
+    renders++;
+    return useAgUiRuntime({ agent });
+  });
+
+  const baseline = renders;
+  const pending = result.current.thread.append({
+    role: "user",
+    content: [{ type: "text", text: "go" }],
+  });
+  await pending;
+  await new Promise((resolve) => setTimeout(resolve, 0));
+
+  const assistant = result.current.thread
+    .getState()
+    .messages.find((message) => message.role === "assistant");
+  expect(assistant).toMatchObject({
+    content: [{ type: "text", text: deltas.join("") }],
+    status: { type: "complete" },
+  });
+  expect(renders - baseline).toBeLessThanOrEqual(6);
+});

package/src/useAgUiRuntime.ts

@@ -76,6 +76,8 @@ export function useAgUiRuntime(
       logger,
       showThinking: options.showThinking ?? true,
       joinStrategy: options.joinStrategy,
+      joinAdjacentAssistantMessages:
+        options.joinAdjacentAssistantMessages ?? false,
       getRunContext: options.getRunContext,
       resumeTranscript: options.resumeTranscript,
       autoCancelPendingToolCalls: options.autoCancelPendingToolCalls,
@@ -98,6 +100,8 @@ export function useAgUiRuntime(
       logger,
       showThinking: options.showThinking ?? true,
       joinStrategy: options.joinStrategy,
+      joinAdjacentAssistantMessages:
+        options.joinAdjacentAssistantMessages ?? false,
       getRunContext: options.getRunContext,
       resumeTranscript: options.resumeTranscript,
       autoCancelPendingToolCalls: options.autoCancelPendingToolCalls,
@@ -223,12 +227,14 @@ export function useAgUiRuntime(
     const abandonPreviousThread = () => {
       queueRef.current?.notifyCancelled();
       if (core.isRunning()) {
-        void core.cancel().catch((error: unknown) => {
+        try {
+          core.leave();
+        } catch (error) {
           logger.error?.(
             "[agui] cancelling the run on thread switch failed",
             error,
           );
-        });
+        }
       }
       queueRef.current?.clear();
     };