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

package/src/runtime/adapter/conversions.ts

@@ -20,12 +20,14 @@ import type { ReadonlyJSONObject } from "assistant-stream/utils";
 import {
   AG_UI_METADATA_NAMESPACE,
   A2UI_SURFACE_ACTIVITY_TYPE,
+  MCP_APPS_ACTIVITY_TYPE,
   type AgUiCustomMetadata,
   type AgUiOpaqueReasoning,
 } from "./run-aggregator";
 import {
   applyA2uiOperations,
   convertSurfaceToUISpec,
+  surfaceToOperations,
   type A2uiState,
   type A2uiSurfaceState,
 } from "@assistant-ui/react-generative-ui/a2ui";
@@ -104,6 +106,7 @@ type ToolCallPart = {
   argsText?: string;
   args?: ReadonlyJSONObject;
   result?: unknown;
+  artifact?: unknown;
   isError?: boolean;
   modelContent?: readonly ToolModelContentPart[];
   unstable_toolMessageId?: string;
@@ -241,7 +244,9 @@ function mediaTypeForMime(mimeType: string | undefined): MediaInputType {
 // Build an AG-UI multimodal source from a data URL, raw base64 payload, or an
 // http(s) URL. A `url` source may omit the mime type; a `data` source always
 // resolves one (falling back to application/octet-stream). An explicit
-// `sourceType: "url"` on the part forces the url leg for non-http references.
+// `sourceType: "url"` on the part forces the url leg for non-http references,
+// and a data URL that is not base64 takes it too, since a `data` source
+// carries base64 bytes.
 function buildInputSource(
   value: string,
   declaredMimeType: string | undefined,
@@ -250,7 +255,8 @@ function buildInputSource(
   const source = resolveFilePartSource({
     data: value,
     mimeType: declaredMimeType ?? "application/octet-stream",
-    sourceType,
+    sourceType:
+      /^data:/i.test(value) && !parseDataUrl(value) ? "url" : sourceType,
   });
   if (source.kind === "url") {
     return declaredMimeType !== undefined
@@ -718,6 +724,7 @@ function attachA2uiSurfaces(
       args: spec as unknown as ReadonlyJSONObject,
       argsText: JSON.stringify(spec),
       result: {},
+      artifact: { a2ui: surfaceToOperations(surface) },
     });
   }
 
@@ -737,9 +744,9 @@ function attachA2uiSurfaces(
 // 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.
+// rather than a boundary an agent drew. The prose of the turn lands in records
+// of its own on either side of it, and nothing on the wire binds them to the
+// run that produced them.
 function isSyntheticToolCallContainer(message: CoreThreadMessageLike): boolean {
   if (message.role !== "assistant" || !Array.isArray(message.content))
     return false;
@@ -759,23 +766,37 @@ function isSyntheticToolCallContainer(message: CoreThreadMessageLike): boolean {
   return carriesOwnId;
 }
 
-function carriesText(message: CoreThreadMessageLike): boolean {
+function carriesPart(message: CoreThreadMessageLike, type: string): boolean {
   return (
     Array.isArray(message.content) &&
-    message.content.some((part) => isObject(part) && part.type === "text")
+    message.content.some((part) => isObject(part) && part.type === type)
+  );
+}
+
+function carriesText(message: CoreThreadMessageLike): boolean {
+  return carriesPart(message, "text");
+}
+
+// A container joins the assistant record ahead of it when that record holds
+// prose or calls of the turn. A released reasoning record holds neither: its
+// part leaves the export under the record's own id, so a message merged under
+// that id would put two records with one id on the wire.
+function opensTurn(message: CoreThreadMessageLike): boolean {
+  return (
+    message.role === "assistant" &&
+    (carriesText(message) || carriesPart(message, "tool-call"))
   );
 }
 
-// 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 {
+// The parts of the answer fold onto the container ahead of them. A record
+// carrying a tool call an agent addressed is a boundary of its own and keeps
+// its message.
+function answersToolCallContainer(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")
+  return !message.content.some(
+    (part) => isObject(part) && part.type === "tool-call",
   );
 }
 
@@ -789,51 +810,53 @@ function readNamespacedMetadata(metadata: unknown): Record<string, unknown> {
   return isObject(namespaced) ? namespaced : {};
 }
 
-function foldOntoToolCallContainer(
-  container: CoreThreadMessageLike,
+function foldTurnRecords(
+  previous: CoreThreadMessageLike,
   message: CoreThreadMessageLike,
 ): CoreThreadMessageLike {
   const interrupts = [
-    ...(readPersistedInterrupts(container.metadata) ?? []),
+    ...(readPersistedInterrupts(previous.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
+  // An entry that rode a record behind the turn's calls sat ahead of content
+  // that now follows them, so it is replayed after the merged record and its
   // tool results rather than ahead of a message that no longer starts there.
+  // One that rode the container opening the turn sat ahead of every call and
+  // stays ahead of the record.
+  const trailsCalls = carriesPart(previous, "tool-call");
   const opaqueReasoning = [
-    ...readOpaqueReasoning(container.metadata),
-    ...readOpaqueReasoning(message.metadata).map((entry) => ({
-      ...entry,
-      after: true,
-    })),
+    ...readOpaqueReasoning(previous.metadata),
+    ...readOpaqueReasoning(message.metadata).map((entry) =>
+      trailsCalls ? { ...entry, after: true } : entry,
+    ),
   ];
   const namespaced = {
-    ...readNamespacedMetadata(container.metadata),
+    ...readNamespacedMetadata(previous.metadata),
     ...readNamespacedMetadata(message.metadata),
     ...(interrupts.length > 0 ? { interrupts } : {}),
     ...(opaqueReasoning.length > 0 ? { opaqueReasoning } : {}),
   } satisfies AgUiCustomMetadata & Record<string, unknown>;
   const custom = {
-    ...readCustomMetadata(container.metadata),
+    ...readCustomMetadata(previous.metadata),
     ...readCustomMetadata(message.metadata),
     ...(Object.keys(namespaced).length > 0
       ? { [AG_UI_METADATA_NAMESPACE]: namespaced }
       : {}),
   };
   const metadata = {
-    ...(isObject(container.metadata) ? container.metadata : {}),
+    ...(isObject(previous.metadata) ? previous.metadata : {}),
     ...(isObject(message.metadata) ? message.metadata : {}),
     ...(Object.keys(custom).length > 0 ? { custom } : {}),
   };
   return {
-    ...container,
+    ...previous,
     ...message,
-    // The container's id names a call rather than the turn, so the record the
+    // A 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,
+    id: carriesText(message) ? message.id : previous.id,
     content: [
-      ...(Array.isArray(container.content) ? container.content : []),
+      ...(Array.isArray(previous.content) ? previous.content : []),
       ...(Array.isArray(message.content) ? message.content : []),
     ],
     ...(Object.keys(metadata).length > 0 ? { metadata } : {}),
@@ -856,6 +879,10 @@ export function fromAgUiMessages(
   const converted: CoreThreadMessageLike[] = [];
   const a2uiBuckets = new Map<string, A2uiState>();
   const a2uiBucketOwnerIndices = new Map<string, number>();
+  const activityParts = new Map<
+    string,
+    { ownerIndex: number; part: { type: "data"; name: string; data: unknown } }
+  >();
   // A zero-data-retention run carries its payload in encryptedValue with no
   // readable content, so the record has no part to become and rides on the
   // neighbouring message instead of being dropped. The anchor is the index the
@@ -1004,15 +1031,9 @@ export function fromAgUiMessages(
     }
 
     if (role === "activity") {
-      // Only a2ui-surface activity messages have an assistant-part equivalent
-      // to rehydrate; other activity types still have no surface to repaint.
       const activityType = getString(rawMessage, "activityType");
-      if (activityType !== A2UI_SURFACE_ACTIVITY_TYPE) continue;
-      const activityContent = isObject(rawMessage.content)
-        ? (rawMessage.content as Record<string, unknown>)
-        : null;
-      const operations = activityContent?.["a2ui_operations"];
-      if (!Array.isArray(operations)) continue;
+      if (activityType === undefined || activityType === MCP_APPS_ACTIVITY_TYPE)
+        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
@@ -1030,6 +1051,48 @@ export function fromAgUiMessages(
       if (ownerIndex === -1) continue;
 
       const owner = converted[ownerIndex]!;
+      if (activityType !== A2UI_SURFACE_ACTIVITY_TYPE) {
+        const bucketKey =
+          getString(rawMessage, "id") ?? `agui-activity:${activityType}`;
+        const part = {
+          type: "data" as const,
+          name: `agui-activity/${activityType}`,
+          data: rawMessage.content,
+        };
+        const existing = activityParts.get(bucketKey);
+        const existingOwner = existing && converted[existing.ownerIndex];
+        const existingIndex =
+          existingOwner && Array.isArray(existingOwner.content)
+            ? existingOwner.content.indexOf(existing.part)
+            : -1;
+        if (existing && existingOwner && existingIndex !== -1) {
+          const content = [...(existingOwner.content as unknown[])];
+          content[existingIndex] = part;
+          converted[existing.ownerIndex] = {
+            ...existingOwner,
+            content: content as typeof existingOwner.content,
+          };
+          activityParts.set(bucketKey, {
+            ownerIndex: existing.ownerIndex,
+            part,
+          });
+          continue;
+        }
+
+        const content = Array.isArray(owner.content) ? owner.content : [];
+        activityParts.set(bucketKey, { ownerIndex, part });
+        converted[ownerIndex] = {
+          ...owner,
+          content: [...content, part],
+        };
+        continue;
+      }
+
+      const activityContent = isObject(rawMessage.content)
+        ? (rawMessage.content as Record<string, unknown>)
+        : null;
+      const operations = activityContent?.["a2ui_operations"];
+      if (!Array.isArray(operations)) continue;
       const bucketKey = getString(rawMessage, "id") ?? "a2ui:anonymous";
       const { state } = applyA2uiOperations(new Map(), operations);
       a2uiBuckets.delete(bucketKey);
@@ -1134,17 +1197,22 @@ export function fromAgUiMessages(
     ]);
   }
 
-  // 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.
+  // A turn that called a tool without addressing a record sits on the wire as
+  // 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.
   const folded: CoreThreadMessageLike[] = [];
   for (const message of converted) {
     const previous = folded[folded.length - 1];
     if (
       previous !== undefined &&
-      isSyntheticToolCallContainer(previous) &&
-      foldsOntoToolCallContainer(message)
+      (isSyntheticToolCallContainer(message)
+        ? opensTurn(previous)
+        : isSyntheticToolCallContainer(previous) &&
+          answersToolCallContainer(message))
     ) {
-      folded[folded.length - 1] = foldOntoToolCallContainer(previous, message);
+      folded[folded.length - 1] = foldTurnRecords(previous, message);
       continue;
     }
     folded.push(message);

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

@@ -13,6 +13,7 @@ import {
 import {
   applyA2uiOperations,
   convertSurfaceToUISpec,
+  surfaceToOperations,
   type A2uiState,
   type A2uiSurfaceState,
 } from "@assistant-ui/react-generative-ui/a2ui";
@@ -49,7 +50,7 @@ type PartOrderEntry =
   | { kind: "text"; key: string; subagentRunId?: string }
   | { kind: "reasoning"; key: string; subagentRunId?: string }
   | { kind: "tool-call"; toolCallId: string }
-  | { kind: "data"; name: string; value: unknown };
+  | { kind: "data"; name: string; value: unknown; subagentRunId?: string };
 
 type BuildContext = {
   subagentsByParentToolCallId: Map<string, string[]>;
@@ -95,6 +96,7 @@ type ToolCallState = {
   toolMessageId?: string;
   mcpAppResourceUri?: string;
   mcpAppServerId?: string;
+  artifact?: unknown;
   modelContent?: ToolModelContentPart[];
   snapshotResultApplied: boolean;
   subagentRunId?: string;
@@ -246,6 +248,10 @@ export class RunAggregator {
   private readonly toolCalls = new Map<string, ToolCallState>();
   private readonly a2uiBuckets = new Map<string, A2uiState>();
   private readonly a2uiToolCallIds = new Set<string>();
+  private readonly activityParts = new Map<
+    string,
+    { kind: "data"; name: string; value: unknown; subagentRunId?: string }
+  >();
   private readonly lastResolvedToolCallIdByScope = new Map<string, string>();
   private readonly partOrder: PartOrderEntry[] = [];
   private textPartCounter = 0;
@@ -319,6 +325,9 @@ export class RunAggregator {
         break;
       }
       case "RUN_CANCELLED": {
+        // A cancel can land after RUN_FINISHED while the response is still
+        // closing, and the answer it finished is already persisted as complete.
+        if (this.status?.type === "complete") break;
         this.status = { type: "incomplete", reason: "cancelled" };
         this.closeOpenSubagentRuns(this.status);
         this.emit();
@@ -503,7 +512,10 @@ export class RunAggregator {
           this.handleA2uiActivitySnapshot(event);
           break;
         }
-        if (event.activityType !== MCP_APPS_ACTIVITY_TYPE) break;
+        if (event.activityType !== MCP_APPS_ACTIVITY_TYPE) {
+          this.handleActivitySnapshot(event);
+          break;
+        }
         const activityScope = this.scopeOf(event);
         const toolCallId = event.content.toolCallId;
         const fallbackId =
@@ -644,6 +656,42 @@ export class RunAggregator {
     this.emit();
   }
 
+  private handleActivitySnapshot(
+    event: Extract<AgUiEvent, { type: "ACTIVITY_SNAPSHOT" }>,
+  ): void {
+    const scope = this.scopeOf(event);
+    const key = this.partKey(
+      scope,
+      event.messageId !== undefined
+        ? `message:${event.messageId}`
+        : `type:${event.activityType}`,
+    );
+    const existing = this.activityParts.get(key);
+    if (existing) {
+      if (event.replace === false) return;
+      existing.name = `agui-activity/${event.activityType}`;
+      existing.value = event.content;
+    } else {
+      const part =
+        scope === ROOT_SCOPE
+          ? {
+              kind: "data" as const,
+              name: `agui-activity/${event.activityType}`,
+              value: event.content,
+            }
+          : {
+              kind: "data" as const,
+              name: `agui-activity/${event.activityType}`,
+              value: event.content,
+              subagentRunId: scope,
+            };
+      this.activityParts.set(key, part);
+      this.partOrder.push(part);
+      this.activeTextMessageIdByScope.delete(scope);
+    }
+    this.emit();
+  }
+
   private synthesizeA2uiToolCalls(): void {
     const surfaces = new Map<string, A2uiSurfaceState>();
 
@@ -671,6 +719,7 @@ export class RunAggregator {
         parsedArgs: spec,
         result: {},
         isError: undefined,
+        artifact: { a2ui: surfaceToOperations(surface) },
         snapshotResultApplied: false,
       };
       if (!this.toolCalls.has(toolCallId)) {
@@ -729,6 +778,7 @@ export class RunAggregator {
     this.toolCalls.clear();
     this.a2uiBuckets.clear();
     this.a2uiToolCallIds.clear();
+    this.activityParts.clear();
     this.lastResolvedToolCallIdByScope.clear();
     this.partOrder.length = 0;
     this.textPartCounter = 0;
@@ -1097,6 +1147,7 @@ export class RunAggregator {
           argsText: entry.argsText,
           ...(approval ? { approval } : {}),
           ...(entry.result !== undefined ? { result: entry.result } : {}),
+          ...(entry.artifact !== undefined ? { artifact: entry.artifact } : {}),
           ...(entry.modelContent !== undefined
             ? { modelContent: entry.modelContent }
             : {}),

package/src/useAgUiRuntime.history.test.tsx

@@ -1,8 +1,11 @@
 // @vitest-environment jsdom
 
 import { describe, it, expect, vi, afterEach } from "vitest";
-import { cleanup, renderHook, waitFor } from "@testing-library/react";
-import type { ThreadHistoryAdapter } from "@assistant-ui/core";
+import { act, cleanup, renderHook, waitFor } from "@testing-library/react";
+import type {
+  ThreadAssistantMessage,
+  ThreadHistoryAdapter,
+} from "@assistant-ui/core";
 import type { HttpAgent } from "@ag-ui/client";
 import { useAgUiRuntime } from "./useAgUiRuntime";
 
@@ -73,4 +76,92 @@ describe("useAgUiRuntime history", () => {
     rerender();
     expect(history.load).toHaveBeenCalledOnce();
   });
+
+  it("records a settled tool interaction through the message part runtime", async () => {
+    const assistant: ThreadAssistantMessage = {
+      id: "assistant-1",
+      role: "assistant",
+      createdAt: new Date(0),
+      status: { type: "complete", reason: "unknown" },
+      content: [
+        {
+          type: "tool-call",
+          toolCallId: "call-1",
+          toolName: "present",
+          args: {},
+          argsText: "{}",
+          result: {},
+        },
+      ],
+      metadata: {
+        unstable_state: null,
+        unstable_annotations: [],
+        unstable_data: [],
+        steps: [],
+        custom: {},
+      },
+    };
+    const update = vi.fn(async () => {});
+    const history: ThreadHistoryAdapter = {
+      load: vi.fn().mockResolvedValue({
+        headId: assistant.id,
+        messages: [{ parentId: null, message: assistant }],
+      }),
+      append: vi.fn().mockResolvedValue(undefined),
+      update,
+    };
+    const { result } = renderHook(() =>
+      useAgUiRuntime({ agent, adapters: { history } }),
+    );
+
+    await waitFor(() =>
+      expect(
+        result.current.thread.getState().messages.map((message) => message.id),
+      ).toEqual([assistant.id]),
+    );
+    const part = result.current.thread
+      .getMessageById(assistant.id)
+      .getMessagePartByToolCallId("call-1");
+    const recordInteraction = part.unstable_recordInteraction;
+    expect(recordInteraction).toBeDefined();
+
+    await act(async () => {
+      await recordInteraction!({
+        type: "action",
+        payload: { action: "confirm" },
+      });
+    });
+
+    expect(part.getState()).toMatchObject({
+      unstable_interactions: {
+        entries: [
+          {
+            type: "action",
+            payload: { action: "confirm" },
+          },
+        ],
+      },
+    });
+    expect(update).toHaveBeenCalledWith(
+      expect.objectContaining({
+        parentId: null,
+        message: expect.objectContaining({
+          id: assistant.id,
+          content: [
+            expect.objectContaining({
+              toolCallId: "call-1",
+              unstable_interactions: expect.objectContaining({
+                entries: [
+                  expect.objectContaining({
+                    type: "action",
+                    payload: { action: "confirm" },
+                  }),
+                ],
+              }),
+            }),
+          ],
+        }),
+      }),
+    );
+  });
 });

package/src/useAgUiRuntime.thread-switch.test.tsx

@@ -2,7 +2,7 @@
 
 import { useState } from "react";
 import { afterEach, describe, expect, it, vi } from "vitest";
-import { act, cleanup, renderHook } from "@testing-library/react";
+import { act, cleanup, renderHook, waitFor } from "@testing-library/react";
 import type { HttpAgent } from "@ag-ui/client";
 import type { ThreadMessage } from "@assistant-ui/core";
 import { AgUiThreadRuntimeCore } from "./runtime/AgUiThreadRuntimeCore";
@@ -15,12 +15,16 @@ type ThreadLoad = Awaited<
 
 function deferred<T>() {
   let resolve!: (value: T) => void;
-  const promise = new Promise<T>((res) => {
+  let reject!: (reason?: unknown) => void;
+  const promise = new Promise<T>((res, rej) => {
     resolve = res;
+    reject = rej;
   });
-  return { promise, resolve };
+  return { promise, resolve, reject };
 }
 
+type Subscriber = Record<string, ((payload: any) => void) | undefined>;
+
 function message(id: string): ThreadMessage {
   return {
     id,
@@ -35,15 +39,17 @@ function message(id: string): ThreadMessage {
 function renderRuntime(
   load: (id: string) => Promise<ThreadLoad>,
   create: () => Promise<void> = async () => {},
+  agentOverride?: HttpAgent,
 ) {
   const agent = {
     runAgent: vi.fn(),
     abortRun: vi.fn(),
   } as unknown as HttpAgent;
+  const runtimeAgent = agentOverride ?? agent;
   return renderHook(() => {
     const [threadId, setThreadId] = useState("initial");
     return useAgUiRuntime({
-      agent,
+      agent: runtimeAgent,
       adapters: {
         threadList: {
           threadId,
@@ -112,6 +118,80 @@ describe("useAgUiRuntime thread switching", () => {
     },
   );
 
+  it.each(["resolve", "reject"] as const)(
+    "keeps a new thread empty after an active run when creation %s",
+    async (outcome) => {
+      const load = deferred<ThreadLoad>();
+      const create = deferred<void>();
+      const run = deferred<void>();
+      let subscriber!: Subscriber;
+      const runAgent = vi.fn(async (_input: unknown, next: Subscriber) => {
+        subscriber = next;
+        await run.promise;
+      });
+      const agent = { runAgent, abortRun: vi.fn() } as unknown as HttpAgent;
+      const { result } = renderRuntime(
+        () => load.promise,
+        () => create.promise,
+        agent,
+      );
+
+      let switchA!: Promise<void>;
+      act(() => {
+        switchA = result.current.threads.switchToThread("thread-a");
+      });
+      await act(async () => {
+        load.resolve({ messages: [message("thread-a")] });
+        await switchA;
+      });
+      expect(
+        result.current.thread.export().messages.map((m) => m.message.id),
+      ).toEqual(["thread-a"]);
+
+      act(() => {
+        void result.current.thread.append("still running");
+      });
+      await waitFor(() => expect(runAgent).toHaveBeenCalledOnce());
+
+      let switchNew!: Promise<void>;
+      act(() => {
+        switchNew = result.current.threads.switchToNewThread();
+      });
+      expect(result.current.thread.export().messages).toEqual([]);
+      await waitFor(() =>
+        expect(result.current.threads.getState().mainThreadId).toBe(
+          "thread-new",
+        ),
+      );
+
+      await act(async () => {
+        if (outcome === "resolve") {
+          create.resolve();
+          await switchNew;
+        } else {
+          create.reject(new Error("create failed"));
+          await expect(switchNew).rejects.toThrow("create failed");
+        }
+      });
+
+      act(() => {
+        subscriber.onMessagesSnapshotEvent?.({
+          event: {
+            type: "MESSAGES_SNAPSHOT",
+            messages: [{ id: "late", role: "assistant", content: "too late" }],
+          },
+        });
+      });
+      expect(result.current.thread.export().messages).toEqual([]);
+
+      await act(async () => {
+        run.resolve();
+        await runAgent.mock.results[0]!.value;
+      });
+      expect(agent.abortRun).toHaveBeenCalledOnce();
+    },
+  );
+
   it("ignores a load superseded by creating a new thread", async () => {
     const resume = vi
       .spyOn(AgUiThreadRuntimeCore.prototype, "resumeInFlightRun")