@assistant-ui/react-ag-ui 0.0.58 → 0.0.60

package/src/runtime/adapter/conversions.ts

@@ -5,6 +5,7 @@ import { generateId } from "@assistant-ui/core";
 import type {
   ThreadMessageLike as CoreThreadMessageLike,
   PartProviderMetadata,
+  ReasoningMessagePart,
   ThreadMessage,
   ToolCallMessagePartMcpMetadata,
   ToolModelContentPart,
@@ -44,14 +45,14 @@ type AttachmentLike = {
 };
 
 type ThreadMessageLike = {
-  id?: string;
+  id?: string | undefined;
   role: string;
   content: unknown;
   metadata?: unknown;
-  name?: string;
-  toolCallId?: string;
-  error?: string;
-  attachments?: readonly AttachmentLike[];
+  name?: string | undefined;
+  toolCallId?: string | undefined;
+  error?: string | undefined;
+  attachments?: readonly AttachmentLike[] | undefined;
 };
 
 type NormalizedThreadMessageLike = ThreadMessageLike & { id: string };
@@ -733,6 +734,112 @@ function attachA2uiSurfaces(
   return { ...message, content: [...preserved, ...a2uiParts] };
 }
 
+// A tool call that arrives without a `parentMessageId` has no record to join, so
+// the client opens one keyed by the call id. An assistant record whose own id is
+// a call it carries is that container: an id the client invented for the call
+// rather than a boundary an agent drew. The prose the run answers with lands in
+// a record of its own, and nothing on the wire binds the two to the run that
+// produced them.
+function isSyntheticToolCallContainer(message: CoreThreadMessageLike): boolean {
+  if (message.role !== "assistant" || !Array.isArray(message.content))
+    return false;
+  let carriesOwnId = false;
+  for (const part of message.content) {
+    if (!isObject(part)) continue;
+    // The container is opened empty, so text on the record means an agent
+    // addressed it and the id is one it chose.
+    if (part.type === "text") return false;
+    if (
+      part.type === "tool-call" &&
+      getString(part, "toolCallId") === message.id
+    ) {
+      carriesOwnId = true;
+    }
+  }
+  return carriesOwnId;
+}
+
+function carriesText(message: CoreThreadMessageLike): boolean {
+  return (
+    Array.isArray(message.content) &&
+    message.content.some((part) => isObject(part) && part.type === "text")
+  );
+}
+
+// What the run put after the container folds onto it: another container for a
+// parallel call, or the parts of the answer. A record carrying a tool call an
+// agent addressed is a boundary of its own and keeps its message.
+function foldsOntoToolCallContainer(message: CoreThreadMessageLike): boolean {
+  if (message.role !== "assistant" || !Array.isArray(message.content))
+    return false;
+  if (message.content.length === 0) return false;
+  return (
+    isSyntheticToolCallContainer(message) ||
+    !message.content.some((part) => isObject(part) && part.type === "tool-call")
+  );
+}
+
+function readCustomMetadata(metadata: unknown): Record<string, unknown> {
+  if (!isObject(metadata) || !isObject(metadata.custom)) return {};
+  return metadata.custom;
+}
+
+function readNamespacedMetadata(metadata: unknown): Record<string, unknown> {
+  const namespaced = readCustomMetadata(metadata)[AG_UI_METADATA_NAMESPACE];
+  return isObject(namespaced) ? namespaced : {};
+}
+
+function foldOntoToolCallContainer(
+  container: CoreThreadMessageLike,
+  message: CoreThreadMessageLike,
+): CoreThreadMessageLike {
+  const interrupts = [
+    ...(readPersistedInterrupts(container.metadata) ?? []),
+    ...(readPersistedInterrupts(message.metadata) ?? []),
+  ];
+  // An entry that rode the folded record sat ahead of content that now follows
+  // the container's calls, so it is replayed after the merged record and its
+  // tool results rather than ahead of a message that no longer starts there.
+  const opaqueReasoning = [
+    ...readOpaqueReasoning(container.metadata),
+    ...readOpaqueReasoning(message.metadata).map((entry) => ({
+      ...entry,
+      after: true,
+    })),
+  ];
+  const namespaced = {
+    ...readNamespacedMetadata(container.metadata),
+    ...readNamespacedMetadata(message.metadata),
+    ...(interrupts.length > 0 ? { interrupts } : {}),
+    ...(opaqueReasoning.length > 0 ? { opaqueReasoning } : {}),
+  } satisfies AgUiCustomMetadata & Record<string, unknown>;
+  const custom = {
+    ...readCustomMetadata(container.metadata),
+    ...readCustomMetadata(message.metadata),
+    ...(Object.keys(namespaced).length > 0
+      ? { [AG_UI_METADATA_NAMESPACE]: namespaced }
+      : {}),
+  };
+  const metadata = {
+    ...(isObject(container.metadata) ? container.metadata : {}),
+    ...(isObject(message.metadata) ? message.metadata : {}),
+    ...(Object.keys(custom).length > 0 ? { custom } : {}),
+  };
+  return {
+    ...container,
+    ...message,
+    // The container's id names a call rather than the turn, so the record the
+    // agent addressed its prose to wins: that is the id the next run has to
+    // carry for the agent to recognize the message it wrote.
+    id: carriesText(message) ? message.id : container.id,
+    content: [
+      ...(Array.isArray(container.content) ? container.content : []),
+      ...(Array.isArray(message.content) ? message.content : []),
+    ],
+    ...(Object.keys(metadata).length > 0 ? { metadata } : {}),
+  } as CoreThreadMessageLike;
+}
+
 export type FromAgUiMessagesOptions = {
   /**
    * Whether to convert `reasoning` messages into visible reasoning parts.
@@ -754,6 +861,40 @@ export function fromAgUiMessages(
   // neighbouring message instead of being dropped. The anchor is the index the
   // next pushed message will occupy, which is where it sat on the wire.
   const opaqueReasoning: (AgUiOpaqueReasoning & { anchor: number })[] = [];
+  // A reasoning record belongs to the assistant record that follows it, which is
+  // the binding the export writes and the one the reference integrations read
+  // back. Holding it until that assistant arrives rebuilds the pair as the
+  // single message the live run produced.
+  const pendingReasoning: {
+    id: string | undefined;
+    part: ReasoningMessagePart;
+  }[] = [];
+
+  const flushPendingReasoning = () => {
+    for (const { id, part } of pendingReasoning) {
+      converted.push({
+        id: id ?? generateId(),
+        role: "assistant",
+        content: [part],
+      });
+    }
+    pendingReasoning.length = 0;
+  };
+
+  const withPendingReasoning = (
+    message: CoreThreadMessageLike,
+  ): CoreThreadMessageLike => {
+    if (pendingReasoning.length === 0) return message;
+    const parts = pendingReasoning.map(({ part }) => part);
+    pendingReasoning.length = 0;
+    return {
+      ...message,
+      content: [
+        ...parts,
+        ...(Array.isArray(message.content) ? message.content : []),
+      ],
+    };
+  };
 
   for (const rawMessage of messages) {
     if (!isObject(rawMessage)) continue;
@@ -761,6 +902,7 @@ export function fromAgUiMessages(
     if (!role) continue;
 
     if (role === "tool") {
+      flushPendingReasoning();
       const toolCallId = getToolCallId(rawMessage) ?? `tool-${generateId()}`;
       const toolMessageId = getString(rawMessage, "id");
       const modelContent = extractText(rawMessage.content);
@@ -872,6 +1014,11 @@ export function fromAgUiMessages(
       const operations = activityContent?.["a2ui_operations"];
       if (!Array.isArray(operations)) continue;
 
+      // A surface belongs to the turn that painted it, and held reasoning is a
+      // nearer antecedent than the previous turn's assistant record, so it is
+      // released here rather than folded past this record.
+      flushPendingReasoning();
+
       let ownerIndex = -1;
       for (let i = converted.length - 1; i >= 0; i--) {
         const candidate = converted[i];
@@ -903,7 +1050,9 @@ export function fromAgUiMessages(
     }
 
     if (role === "assistant") {
-      converted.push(toAssistantSnapshotMessage(rawMessage));
+      converted.push(
+        withPendingReasoning(toAssistantSnapshotMessage(rawMessage)),
+      );
       continue;
     }
 
@@ -916,6 +1065,9 @@ export function fromAgUiMessages(
         // accept the next run.
         const opaqueId = getString(rawMessage, "id");
         if (opaqueId?.trim() && encryptedValue?.trim()) {
+          // The anchor counts materialized messages, so anything still held
+          // takes its own slot rather than folding past this record.
+          flushPendingReasoning();
           opaqueReasoning.push({
             id: opaqueId,
             encryptedValue,
@@ -939,31 +1091,37 @@ export function fromAgUiMessages(
         }
         continue;
       }
-      converted.push({
-        id: getString(rawMessage, "id") ?? generateId(),
-        role: "assistant",
-        content: [
-          {
-            type: "reasoning",
-            text,
-            ...(encryptedValue !== undefined
-              ? {
-                  providerMetadata: {
-                    [AG_UI_METADATA_NAMESPACE]: { encryptedValue },
-                  },
-                }
-              : {}),
-          },
-        ],
+      const rawReasoningId = getString(rawMessage, "id");
+      // A blank id still wins over a synthesized one on export, which would put
+      // an unaddressable record on the wire.
+      const reasoningId = rawReasoningId?.trim() ? rawReasoningId : undefined;
+      // The fold costs the record its own message id, so it rides the part
+      // instead: the export re-emits each block under the id it arrived with.
+      const meta = {
+        ...(reasoningId !== undefined ? { reasoningId } : {}),
+        ...(encryptedValue !== undefined ? { encryptedValue } : {}),
+      };
+      pendingReasoning.push({
+        id: reasoningId,
+        part: {
+          type: "reasoning",
+          text,
+          ...(Object.keys(meta).length > 0
+            ? { providerMetadata: { [AG_UI_METADATA_NAMESPACE]: meta } }
+            : {}),
+        },
       });
       continue;
     }
 
     if (role === "user" || role === "system" || role === "developer") {
+      flushPendingReasoning();
       converted.push(toUserOrSystemSnapshotMessage(role, rawMessage));
     }
   }
 
+  flushPendingReasoning();
+
   for (const { anchor, ...entry } of opaqueReasoning) {
     // Nothing followed it on the wire, so it trails the last message instead.
     const trailing = anchor >= converted.length;
@@ -976,8 +1134,24 @@ export function fromAgUiMessages(
     ]);
   }
 
-  for (let i = 0; i < converted.length; i++) {
-    const message = converted[i]!;
+  // A turn that called a tool before answering sits on the wire as a container
+  // and the records that follow it, which the live run renders as one message.
+  const folded: CoreThreadMessageLike[] = [];
+  for (const message of converted) {
+    const previous = folded[folded.length - 1];
+    if (
+      previous !== undefined &&
+      isSyntheticToolCallContainer(previous) &&
+      foldsOntoToolCallContainer(message)
+    ) {
+      folded[folded.length - 1] = foldOntoToolCallContainer(previous, message);
+      continue;
+    }
+    folded.push(message);
+  }
+
+  for (let i = 0; i < folded.length; i++) {
+    const message = folded[i]!;
     if (message.role !== "assistant") continue;
 
     const hasInterrupt =
@@ -992,7 +1166,7 @@ export function fromAgUiMessages(
       );
 
     if (hasInterrupt || hasPendingToolCall) {
-      converted[i] = {
+      folded[i] = {
         ...message,
         status: getAutoStatus(
           false,
@@ -1005,7 +1179,7 @@ export function fromAgUiMessages(
     }
   }
 
-  return converted;
+  return folded;
 }
 
 function convertAssistantMessage(

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

@@ -87,6 +87,7 @@ type ToolCallState = {
   toolCallId: string;
   toolCallName: string;
   argsText: string;
+  argsTextScanner?: JSONContainerScanner | undefined;
   parsedArgs: Record<string, unknown> | undefined;
   result: unknown;
   isError: boolean | undefined;
@@ -99,6 +100,80 @@ type ToolCallState = {
   subagentRunId?: string;
 };
 
+type JSONContainerScanner = {
+  state: "pending" | "open" | "complete" | "invalid";
+  stack: ("{" | "[")[];
+  inString: boolean;
+  escaped: boolean;
+};
+
+const createJSONContainerScanner = (): JSONContainerScanner => ({
+  state: "pending",
+  stack: [],
+  inString: false,
+  escaped: false,
+});
+
+const isJSONWhitespace = (char: string): boolean =>
+  char === " " || char === "\t" || char === "\n" || char === "\r";
+
+const scanJSONContainerDelta = (
+  scanner: JSONContainerScanner,
+  delta: string,
+): boolean => {
+  let completed = false;
+
+  for (const char of delta) {
+    if (scanner.state === "invalid") continue;
+    if (scanner.state === "complete") {
+      if (isJSONWhitespace(char)) continue;
+      scanner.state = "invalid";
+      continue;
+    }
+    if (scanner.state === "pending") {
+      if (isJSONWhitespace(char)) continue;
+      if (char !== "{" && char !== "[") {
+        scanner.state = "invalid";
+        continue;
+      }
+      scanner.state = "open";
+      scanner.stack.push(char);
+      continue;
+    }
+    if (scanner.inString) {
+      if (scanner.escaped) {
+        scanner.escaped = false;
+      } else if (char === "\\") {
+        scanner.escaped = true;
+      } else if (char === '"') {
+        scanner.inString = false;
+      }
+      continue;
+    }
+    if (char === '"') {
+      scanner.inString = true;
+      continue;
+    }
+    if (char === "{" || char === "[") {
+      scanner.stack.push(char);
+      continue;
+    }
+    if (char === "}" || char === "]") {
+      const expected = char === "}" ? "{" : "[";
+      if (scanner.stack.pop() !== expected) {
+        scanner.state = "invalid";
+        continue;
+      }
+      if (scanner.stack.length === 0) {
+        scanner.state = "complete";
+        completed = true;
+      }
+    }
+  }
+
+  return completed && scanner.state === "complete";
+};
+
 export const MCP_APPS_ACTIVITY_TYPE = "mcp-apps";
 
 export const A2UI_SURFACE_ACTIVITY_TYPE = "a2ui-surface";
@@ -816,7 +891,17 @@ export class RunAggregator {
   private appendToolArgs(id: string | undefined, delta: string) {
     const entry = id ? this.toolCalls.get(id) : undefined;
     if (!entry) return;
+    if (!entry.argsTextScanner) {
+      entry.argsTextScanner = createJSONContainerScanner();
+    }
     entry.argsText += delta;
+    const shouldParse = scanJSONContainerDelta(entry.argsTextScanner, delta);
+    if (!shouldParse) {
+      if (entry.argsTextScanner.state === "invalid") {
+        entry.parsedArgs = undefined;
+      }
+      return;
+    }
     try {
       const parsed = JSON.parse(entry.argsText);
       if (parsed && typeof parsed === "object") {

package/src/runtime/types.ts

@@ -56,10 +56,36 @@ export type UseAgUiRuntimeAdapters = {
   threadList?: UseAgUiThreadListAdapter;
 };
 
+export type AgUiResumeTranscript = "full" | "appended";
+
 export type UseAgUiRuntimeOptions = ExternalStoreSharedOptions & {
   agent: AbstractAgent;
   logger?: Partial<Logger>;
   showThinking?: boolean;
+  /**
+   * What `messages` carries on a resume run, meaning a `RunAgentInput` that
+   * also carries `resume`. The AG-UI interrupt spec leaves this undefined: its
+   * contract rules constrain the thread id, interrupt coverage, idempotency,
+   * and expiry, and none of them mentions `messages`. Hosts therefore disagree,
+   * and both readings are conformant.
+   *
+   * `"full"` sends the whole thread, which is what `@ag-ui/client` itself does.
+   * Hosts that resume from a checkpoint ignore the transcript, and hosts that
+   * rebuild the interrupted run from it need it.
+   *
+   * `"appended"` sends only what was appended locally after the interrupted
+   * assistant message, which is nothing for an approval or a denial and the new
+   * user turn for `steerAway`. Choose it for a host that owns the thread and
+   * seeds a resume request from its own stored snapshot, because such a host
+   * appends the request body to that snapshot and a re-sent transcript
+   * duplicates the interrupted turn in its stored history. A host that rebuilds
+   * the run from `messages` has nothing to resume from under `"appended"`, and
+   * a host that derives its outgoing `MESSAGES_SNAPSHOT` from the request body
+   * emits a truncated one.
+   *
+   * Defaults to `"full"`.
+   */
+  resumeTranscript?: AgUiResumeTranscript | undefined;
   /**
    * When the user sends, edits, or reloads a message while client-side tool
    * calls are still pending, automatically cancel the unresolved tool calls

package/src/useAgUiRuntime.queue.test.tsx

@@ -1,6 +1,7 @@
 // @vitest-environment jsdom
 
 import { describe, it, expect, vi, afterEach } from "vitest";
+import { startTransition, useEffect } from "react";
 import {
   act,
   cleanup,
@@ -409,6 +410,59 @@ describe("useAgUiRuntime unstable_enableMessageQueue", () => {
     );
   });
 
+  it("does not clear the committed queue from a suspended render", async () => {
+    const { agent, runAgent } = gatedAgent();
+    let runtime: AssistantRuntime | undefined;
+    const never = new Promise<void>(() => {});
+
+    const Harness = ({
+      enabled,
+      suspend,
+    }: {
+      enabled: boolean;
+      suspend: boolean;
+    }) => {
+      const nextRuntime = useAgUiRuntime({
+        agent,
+        unstable_enableMessageQueue: enabled,
+      });
+      useEffect(() => {
+        runtime = nextRuntime;
+      }, [nextRuntime]);
+      if (suspend) throw never;
+      return null;
+    };
+
+    const view = render(<Harness enabled suspend={false} />);
+    await waitFor(() => expect(runtime).toBeDefined());
+    mount(runtime!);
+
+    await act(async () => {
+      await runtime!.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "first" }],
+      });
+    });
+    await waitFor(() => expect(runAgent).toHaveBeenCalledTimes(1));
+
+    await act(async () => {
+      await runtime!.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "second" }],
+        parentId: runtime!.thread.getState().messages.at(-1)?.id ?? null,
+      });
+    });
+    await waitFor(() =>
+      expect(screen.getByTestId("queued").textContent).toBe("second"),
+    );
+
+    act(() => {
+      startTransition(() => view.rerender(<Harness enabled={false} suspend />));
+    });
+
+    expect(screen.getByTestId("queued").textContent).toBe("second");
+  });
+
   it("leaves the queue capability off when the flag is not set", () => {
     const { agent } = gatedAgent();