@assistant-ui/react-ag-ui 0.1.1 → 0.1.3

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

@@ -22,6 +22,7 @@ import {
 import jsonpatch, { type Operation } from "fast-json-patch";
 import { readMcpAppResourceUri, toolContentText } from "../mcp-tool-result";
 import { projectAgUiToolApprovals } from "./tool-approval";
+import { withAgUiPartMessageMetadata } from "./conversions";
 import {
   AG_UI_METADATA_NAMESPACE,
   A2UI_SURFACE_ACTIVITY_TYPE,
@@ -30,15 +31,12 @@ import {
   type AgUiEvent,
   type AgUiInterrupt,
   type AgUiOpaqueReasoning,
+  MAX_SUBAGENT_DEPTH,
 } from "../types";
 import type { Logger } from "../logger";
 
 const ROOT_SCOPE = "";
 
-// Defensive guard against a malformed or cyclic parentToolCallId /
-// parentSubagentRunId chain, not a realistic nesting depth.
-const MAX_SUBAGENT_DEPTH = 16;
-
 type SubagentRunState = {
   subagentRunId: string;
   name: string;
@@ -93,6 +91,7 @@ type ToolCallState = {
   mcpAppServerId?: string;
   artifact?: unknown;
   modelContent?: ToolModelContentPart[];
+  resultMetadata?: Record<string, unknown>;
   snapshotResultApplied: boolean;
   subagentRunId?: string;
 };
@@ -209,6 +208,7 @@ export class RunAggregator {
   private interrupts: AgUiInterrupt[] | undefined;
   private pendingToolCallIds: string[] | undefined;
   private readonly textParts = new Map<string, { buffer: string }>();
+  private readonly pendingSnapshotText = new Map<string, string>();
   private readonly activeTextMessageIdByScope = new Map<string, string>();
   private readonly reasoningParts = new Map<
     string,
@@ -267,6 +267,17 @@ export class RunAggregator {
     return this.toolCalls.has(toolCallId);
   }
 
+  rebaseMessage(messageId: string): void {
+    const snapshotText = this.pendingSnapshotText.get(messageId);
+    this.resetMessageParts();
+    this.lastTextMessageId = messageId;
+    this.reportedServerMessageId = messageId;
+    if (snapshotText !== undefined) {
+      this.pendingSnapshotText.set(messageId, snapshotText);
+    }
+    this.startTextMessage(ROOT_SCOPE, messageId);
+  }
+
   handle(event: AgUiEvent): void {
     switch (event.type) {
       case "RUN_STARTED": {
@@ -282,6 +293,20 @@ export class RunAggregator {
         this.emit();
         break;
       }
+      case "MESSAGES_SNAPSHOT": {
+        this.pendingSnapshotText.clear();
+        for (const message of event.messages) {
+          if (
+            message.role !== "assistant" ||
+            typeof message.content !== "string"
+          )
+            continue;
+          const part = this.textParts.get(message.id);
+          if (part) part.buffer = message.content;
+          else this.pendingSnapshotText.set(message.id, message.content);
+        }
+        break;
+      }
       case "RUN_FINISHED": {
         if (event.outcome?.type === "cancelled") {
           this.interrupts = undefined;
@@ -368,7 +393,7 @@ export class RunAggregator {
         const incomingId = event.messageId;
         const scope = this.scopeOf(event);
         if (scope === ROOT_SCOPE) {
-          this.beginDistinctTextMessage(incomingId);
+          this.beginDistinctTextMessage(incomingId, true);
           this.reportServerMessageId(incomingId);
           if (incomingId) this.lastTextMessageId = incomingId;
         }
@@ -494,6 +519,7 @@ export class RunAggregator {
               : undefined,
           event.messageId,
           event.mcpResult,
+          event.metadata,
         );
         this.emit();
         break;
@@ -793,6 +819,7 @@ export class RunAggregator {
 
   private resetMessageParts(): void {
     this.clearTextParts();
+    this.pendingSnapshotText.clear();
     this.reasoningParts.clear();
     this.reasoningSignatures.clear();
     this.reasoningSignatureIds.clear();
@@ -837,7 +864,10 @@ export class RunAggregator {
     return scope === ROOT_SCOPE ? id : `${scope}\u0000${id}`;
   }
 
-  private beginDistinctTextMessage(messageId: string | undefined): void {
+  private beginDistinctTextMessage(
+    messageId: string | undefined,
+    preserveSnapshotText = false,
+  ): void {
     if (
       this.joinStrategy === "concat-content" ||
       !messageId ||
@@ -854,7 +884,13 @@ export class RunAggregator {
     const openToolCallIds = Array.from(this.toolCalls.values())
       .filter((call) => call.result === undefined)
       .map((call) => call.toolCallId);
+    const snapshotText = preserveSnapshotText
+      ? this.pendingSnapshotText.get(messageId)
+      : undefined;
     this.resetMessageParts();
+    if (snapshotText !== undefined) {
+      this.pendingSnapshotText.set(messageId, snapshotText);
+    }
     this.onTextMessageStart(messageId, openToolCallIds);
     this.restoreReasoning(carried);
   }
@@ -946,15 +982,21 @@ export class RunAggregator {
   private startTextMessage(scope: string, messageId?: string): string {
     const id = messageId ?? this.generateTextKey();
     const key = this.ensureTextPart(scope, id);
+    if (scope === ROOT_SCOPE && messageId) {
+      const snapshotText = this.pendingSnapshotText.get(messageId);
+      const part = this.textParts.get(key);
+      if (snapshotText !== undefined && part) {
+        part.buffer = snapshotText;
+        this.pendingSnapshotText.delete(messageId);
+      }
+    }
     this.activeTextMessageIdByScope.set(scope, key);
     return key;
   }
 
   private resolveTextMessageId(scope: string, messageId?: string): string {
     if (messageId) {
-      const key = this.ensureTextPart(scope, messageId);
-      this.activeTextMessageIdByScope.set(scope, key);
-      return key;
+      return this.startTextMessage(scope, messageId);
     }
 
     const active = this.activeTextMessageIdByScope.get(scope);
@@ -1098,6 +1140,7 @@ export class RunAggregator {
     isError?: boolean,
     toolMessageId?: string,
     mcpResult?: Record<string, unknown>,
+    metadata?: Record<string, unknown>,
   ) {
     if (!id) return;
     const modelText = toolContentText(content);
@@ -1149,6 +1192,9 @@ export class RunAggregator {
     if (toolMessageId) {
       entry.toolMessageId = toolMessageId;
     }
+    if (metadata !== undefined) {
+      entry.resultMetadata = { ...entry.resultMetadata, ...metadata };
+    }
     this.lastResolvedToolCallIdByScope.set(scope, id);
   }
 
@@ -1299,6 +1345,14 @@ export class RunAggregator {
           ...(entry.toolMessageId
             ? { unstable_toolMessageId: entry.toolMessageId }
             : {}),
+          ...(entry.resultMetadata !== undefined
+            ? {
+                providerMetadata: withAgUiPartMessageMetadata(
+                  undefined,
+                  entry.resultMetadata,
+                ),
+              }
+            : {}),
           ...(nestedMessages.length > 0 ? { messages: nestedMessages } : {}),
         } as ToolCallMessagePart & { unstable_toolMessageId?: string };
         snapshot.push(toolPart);

package/src/runtime/adapter/subscriber.test.ts

@@ -4,7 +4,7 @@ import { createAgUiSubscriber } from "./subscriber";
 const createSubscriber = () => {
   const dispatch = vi.fn();
   const logger = { debug: vi.fn(), error: vi.fn() };
-  const subscriber = createAgUiSubscriber({ dispatch, runId: "run-1", logger });
+  const subscriber = createAgUiSubscriber({ dispatch, logger });
   return { subscriber, dispatch, logger };
 };
 

package/src/runtime/adapter/subscriber.ts

@@ -60,7 +60,6 @@ const ensureEvent = (
 
 type SubscriberOptions = {
   dispatch: Dispatch;
-  runId: string;
   logger?: Logger;
   onRunFailed?: (error: Error) => void;
 };
@@ -68,7 +67,7 @@ type SubscriberOptions = {
 export const createAgUiSubscriber = (
   options: SubscriberOptions,
 ): Subscriber => {
-  const { dispatch, runId, logger, onRunFailed } = options;
+  const { dispatch, logger, onRunFailed } = options;
   let runFinishedDispatched = false;
   const dispatchIfValid = (raw: unknown, type: AgUiEvent["type"]) => {
     const event = ensureEvent(raw, type, logger);
@@ -140,10 +139,6 @@ export const createAgUiSubscriber = (
       logger?.debug?.("[agui] step boundary ignored", event),
     onStepFinishedEvent: ({ event }) =>
       logger?.debug?.("[agui] step boundary ignored", event),
-    onRunFinalized: () => {
-      if (runFinishedDispatched) return;
-      dispatch({ type: "RUN_FINISHED", runId });
-    },
     onRunFailed: ({ error }) => {
       if (runFinishedDispatched) return;
       runFinishedDispatched = true;

package/src/runtime/event-parser.ts

@@ -344,6 +344,9 @@ export const parseAgUiEvent = (
           messageId: getString("messageId"),
           role: payload.role === "tool" ? "tool" : undefined,
           mcpResult: parseMcpToolCallResult(payload, content),
+          metadata: isPlainObject(payload.metadata)
+            ? payload.metadata
+            : undefined,
           subagentRunId: getString("subagentRunId"),
         },
       );

package/src/runtime/types.ts

@@ -26,6 +26,10 @@ export const AG_UI_METADATA_NAMESPACE = "agui";
 export const MCP_APPS_ACTIVITY_TYPE = "mcp-apps";
 export const A2UI_SURFACE_ACTIVITY_TYPE = "a2ui-surface";
 
+// Defensive guard against a malformed or cyclic parentToolCallId /
+// parentSubagentRunId chain, not a realistic nesting depth.
+export const MAX_SUBAGENT_DEPTH = 16;
+
 export type AgUiOpaqueReasoning = {
   id: string;
   encryptedValue: string;
@@ -278,6 +282,7 @@ export type AgUiEvent =
       content: string | ContentPart[];
       role?: "tool";
       mcpResult?: McpToolCallResult;
+      metadata?: Record<string, unknown>;
       subagentRunId?: string;
     }
   | {

package/src/useAgUiRuntime.cloud-run-leases.test.tsx

@@ -0,0 +1,384 @@
+// @vitest-environment jsdom
+
+import { act, cleanup, render, waitFor } from "@testing-library/react";
+import { HttpAgent, type RunAgentInput } from "@ag-ui/client";
+import {
+  AssistantRuntimeProvider,
+  useCloudThreadListAdapter,
+  useRemoteThreadListRuntime,
+} from "@assistant-ui/core/react";
+import type { AssistantRuntime } from "@assistant-ui/core";
+import { useAui, type AssistantClient } from "@assistant-ui/store";
+import { AssistantCloud } from "assistant-cloud";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { useAgUiRuntime } from "./useAgUiRuntime";
+
+// The real AssistantCloud client runs against this in-memory server.
+type Row = {
+  id: string;
+  parent_id: string | null;
+  format: string;
+  content: unknown;
+  external_id: string | null;
+  run?: { status: "running" | "ended"; last_heartbeat_ms_ago: number };
+};
+
+type Call = { method: string; path: string; body: any };
+
+const createServer = (
+  initialRows: Row[] = [],
+  heartbeatIntervalMs = 20_000,
+) => {
+  const rows = [...initialRows];
+  const calls: Call[] = [];
+  let created = 0;
+  let leases = 0;
+  const at = new Date(0).toISOString();
+  const json = (value: unknown, status = 200) =>
+    new Response(JSON.stringify(value), {
+      status,
+      headers: { "content-type": "application/json" },
+    });
+  const fetchMock = vi.fn(
+    async (input: RequestInfo | URL, init?: RequestInit) => {
+      const url = new URL(String(input));
+      const method = init?.method ?? "GET";
+      const path = url.pathname.replace(/^\/v1/, "");
+      const body = init?.body ? JSON.parse(String(init.body)) : undefined;
+      calls.push({ method, path, body });
+      if (method === "GET" && path === "/threads") {
+        return json({
+          threads: rows.length
+            ? [
+                {
+                  id: "remote-1",
+                  title: "t",
+                  is_archived: false,
+                  external_id: null,
+                  metadata: null,
+                  project_id: "p",
+                  workspace_id: "w",
+                  created_at: at,
+                  updated_at: at,
+                  last_message_at: at,
+                },
+              ]
+            : [],
+        });
+      }
+      if (method === "POST" && path === "/threads")
+        return json({ thread_id: "remote-1" });
+      if (method === "GET" && path === "/threads/remote-1/messages") {
+        return json({
+          messages: [...rows].reverse().map((row, height) => ({
+            ...row,
+            height,
+            created_at: at,
+            updated_at: at,
+            run: row.run ?? null,
+          })),
+        });
+      }
+      if (method === "POST" && path === "/threads/remote-1/messages") {
+        const id = `msg_${++created}`;
+        rows.push({
+          id,
+          parent_id: body.parent_id,
+          format: body.format,
+          content: body.content,
+          external_id: body.external_id ?? null,
+          ...(body.open_run
+            ? { run: { status: "running", last_heartbeat_ms_ago: 0 } }
+            : undefined),
+        });
+        return json({
+          message_id: id,
+          ...(body.open_run
+            ? {
+                lease_id: `lease_${++leases}`,
+                heartbeat_interval_ms: heartbeatIntervalMs,
+              }
+            : undefined),
+        });
+      }
+      const write = path.match(/^\/threads\/remote-1\/messages\/([^/]+)$/);
+      if (method === "PUT" && write) {
+        const row = rows.find((r) => r.id === write[1]);
+        if (row) row.content = body.content;
+        return json({});
+      }
+      if (/^\/threads\/remote-1\/runs\/[^/]+\/heartbeat$/.test(path))
+        return json({ heartbeat_interval_ms: heartbeatIntervalMs });
+      if (/^\/threads\/remote-1\/runs\/[^/]+\/end$/.test(path)) return json({});
+      return json({});
+    },
+  );
+  return { rows, calls, fetchMock };
+};
+
+const createsWith = (calls: Call[], openRun: boolean) =>
+  calls.filter(
+    (c) =>
+      c.method === "POST" &&
+      c.path === "/threads/remote-1/messages" &&
+      !!c.body.open_run === openRun,
+  );
+
+type Subscriber = Record<
+  string,
+  ((payload: { event: any }) => void) | undefined
+>;
+
+// An AG-UI agent whose runs are scripted: it streams one text message, under an id the backend picks, and finishes when the test says so.
+class ScriptedAgent extends HttpAgent {
+  runs: { finish: () => void; input: RunAgentInput }[] = [];
+  constructor() {
+    super({ url: "http://agent.test" });
+  }
+  override async runAgent(
+    parameters?: Partial<RunAgentInput>,
+    subscriber?: Subscriber,
+  ): Promise<any> {
+    const input = parameters as RunAgentInput;
+    const id = `agent-msg-${this.runs.length + 1}`;
+    const emit = (name: string, event: object) =>
+      subscriber?.[name]?.({ event });
+    emit("onRunStartedEvent", {
+      type: "RUN_STARTED",
+      threadId: input.threadId,
+      runId: input.runId,
+    });
+    emit("onTextMessageStartEvent", {
+      type: "TEXT_MESSAGE_START",
+      messageId: id,
+      role: "assistant",
+    });
+    emit("onTextMessageContentEvent", {
+      type: "TEXT_MESSAGE_CONTENT",
+      messageId: id,
+      delta: "partial",
+    });
+    await new Promise<void>((resolve) =>
+      this.runs.push({ finish: resolve, input }),
+    );
+    emit("onTextMessageContentEvent", {
+      type: "TEXT_MESSAGE_CONTENT",
+      messageId: id,
+      delta: " answer",
+    });
+    emit("onTextMessageEndEvent", { type: "TEXT_MESSAGE_END", messageId: id });
+    emit("onRunFinishedEvent", {
+      type: "RUN_FINISHED",
+      threadId: input.threadId,
+      runId: input.runId,
+    });
+    return { result: undefined, newMessages: [] };
+  }
+  override abortRun() {}
+}
+
+const renderApp = async (staleAfter?: number) => {
+  const cloud = new AssistantCloud({
+    baseUrl: "http://cloud.test",
+    apiKey: "key",
+    userId: "user",
+    workspaceId: "workspace",
+    telemetry: false,
+  });
+  const agent = new ScriptedAgent();
+  let aui!: AssistantClient;
+  let runtime!: AssistantRuntime;
+  const Capture = () => {
+    aui = useAui();
+    return null;
+  };
+  const App = () => {
+    const adapter = useCloudThreadListAdapter({
+      cloud,
+      ...(staleAfter !== undefined ? { staleAfter } : undefined),
+    });
+    runtime = useRemoteThreadListRuntime({
+      adapter,
+      runtimeHook: function RuntimeHook() {
+        return useAgUiRuntime({ agent });
+      },
+    });
+    return (
+      <AssistantRuntimeProvider runtime={runtime}>
+        <Capture />
+      </AssistantRuntimeProvider>
+    );
+  };
+  render(<App />);
+  await act(async () => {
+    await aui.threads.getLoadThreadsPromise();
+  });
+  return { agent, getAui: () => aui, getRuntime: () => runtime };
+};
+
+const assistantRow = (
+  last_heartbeat_ms_ago: number,
+  status: "running" | "ended" = "running",
+): Row => ({
+  id: "a1",
+  parent_id: "u1",
+  format: "aui/v0",
+  content: {
+    role: "assistant",
+    content: [{ type: "text", text: "half an answer" }],
+    status: { type: "incomplete", reason: "cancelled" },
+  },
+  external_id: "a1",
+  run: { status, last_heartbeat_ms_ago },
+});
+const userRow: Row = {
+  id: "u1",
+  parent_id: null,
+  format: "aui/v0",
+  content: { role: "user", content: [{ type: "text", text: "hello" }] },
+  external_id: "u1",
+};
+
+let server: ReturnType<typeof createServer>;
+const useServer = (rows?: Row[], heartbeatIntervalMs?: number) => {
+  server = createServer(rows, heartbeatIntervalMs);
+  vi.stubGlobal("fetch", server.fetchMock);
+};
+
+beforeEach(() => {
+  vi.spyOn(console, "warn").mockImplementation(() => {});
+});
+afterEach(() => {
+  cleanup();
+  vi.unstubAllGlobals();
+  vi.restoreAllMocks();
+});
+
+describe("useAgUiRuntime with Assistant Cloud run leases", () => {
+  it("opens a lease when the run starts, heartbeats on the server's interval, writes as its holder and ends it when the run finishes", async () => {
+    useServer([], 30);
+    const { agent, getAui } = await renderApp();
+
+    await act(async () => {
+      getAui().thread.append("hello");
+    });
+    await waitFor(() => expect(agent.runs).toHaveLength(1));
+    await waitFor(() =>
+      expect(createsWith(server.calls, true)).toHaveLength(1),
+    );
+    const open = createsWith(server.calls, true)[0]!;
+    expect(open.body).toMatchObject({
+      format: "aui/v0",
+      parent_id: "msg_1",
+      content: { role: "assistant" },
+    });
+
+    await waitFor(() =>
+      expect(
+        server.calls.filter((c) => c.path.endsWith("/runs/lease_1/heartbeat"))
+          .length,
+      ).toBeGreaterThanOrEqual(2),
+    );
+
+    await act(async () => {
+      agent.runs[0]!.finish();
+    });
+    await waitFor(() =>
+      expect(
+        server.calls.some(
+          (c) => c.method === "POST" && c.path.endsWith("/runs/lease_1/end"),
+        ),
+      ).toBe(true),
+    );
+    const end = server.calls.find((c) => c.path.endsWith("/runs/lease_1/end"))!;
+    expect(end.body).toEqual({ reason: "done" });
+    const final = server.calls.filter(
+      (c) => c.method === "PUT" && c.body.lease_id === "lease_1",
+    );
+    expect(final.length).toBeGreaterThan(0);
+    expect(JSON.stringify(final.at(-1)!.body.content)).toContain(
+      "partial answer",
+    );
+    expect(createsWith(server.calls, true)).toHaveLength(1);
+    const reply = getAui().thread.getState().messages[1]!;
+    expect(reply.id).toBe("agent-msg-1");
+  });
+
+  it("shows a reloaded run quiet for longer than staleAfter as interrupted; Resume posts a sibling with its own lease and runs the agent again", async () => {
+    useServer([userRow, assistantRow(45_000)]);
+    const { agent, getAui, getRuntime } = await renderApp();
+    await act(async () => {
+      await getRuntime().threads.switchToThread("remote-1");
+    });
+    await waitFor(() =>
+      expect(getAui().thread.getState().messages).toHaveLength(2),
+    );
+    const stale = getAui().thread.message({ index: 1 });
+    await waitFor(() => expect(stale.getState().runState).toBe("stale"));
+
+    await act(async () => {
+      stale.reload();
+    });
+    await waitFor(() => expect(agent.runs).toHaveLength(1));
+    await waitFor(() =>
+      expect(createsWith(server.calls, true)).toHaveLength(1),
+    );
+    expect(createsWith(server.calls, true)[0]!.body).toMatchObject({
+      parent_id: "u1",
+      open_run: true,
+    });
+    await act(async () => {
+      agent.runs[0]!.finish();
+    });
+    await waitFor(() =>
+      expect(
+        server.calls.some((c) => c.path.endsWith("/runs/lease_1/end")),
+      ).toBe(true),
+    );
+    expect(
+      server.calls.some((c) => c.method === "PUT" && c.path.endsWith("/a1")),
+    ).toBe(false);
+    const messages = getAui().thread.getState().messages;
+    expect(messages.map((m) => m.role)).toEqual(["user", "assistant"]);
+    expect(messages[1]!.id).not.toBe("a1");
+  });
+
+  it("shows a reloaded run with a recent beat as in progress, and staleAfter moves the line", async () => {
+    useServer([userRow, assistantRow(5_000)]);
+    const { getAui, getRuntime } = await renderApp();
+    await act(async () => {
+      await getRuntime().threads.switchToThread("remote-1");
+    });
+    await waitFor(() =>
+      expect(getAui().thread.message({ index: 1 }).getState().runState).toBe(
+        "in-progress",
+      ),
+    );
+    cleanup();
+
+    useServer([userRow, assistantRow(5_000)]);
+    const short = await renderApp(2_000);
+    await act(async () => {
+      await short.getRuntime().threads.switchToThread("remote-1");
+    });
+    await waitFor(() =>
+      expect(
+        short.getAui().thread.message({ index: 1 }).getState().runState,
+      ).toBe("stale"),
+    );
+  });
+
+  it("loads an ended run as a plain message", async () => {
+    useServer([userRow, assistantRow(90_000, "ended")]);
+    const { getAui, getRuntime } = await renderApp();
+    await act(async () => {
+      await getRuntime().threads.switchToThread("remote-1");
+    });
+    await waitFor(() =>
+      expect(getAui().thread.getState().messages).toHaveLength(2),
+    );
+    expect(
+      getAui().thread.message({ index: 1 }).getState().runState,
+    ).toBeUndefined();
+  });
+});