@assistant-ui/react-ag-ui 0.0.63 → 0.0.65

package/src/runtime/adapter/conversions.ts

@@ -6,14 +6,12 @@ import type {
   ThreadMessageLike as CoreThreadMessageLike,
   PartProviderMetadata,
   ReasoningMessagePart,
-  ThreadMessage,
-  ToolCallMessagePartMcpMetadata,
-  ToolModelContentPart,
 } from "@assistant-ui/core";
 import {
   getAutoStatus,
   parseDataUrl,
   resolveFilePartSource,
+  walkToolCallTree,
 } from "@assistant-ui/core/internal";
 import { type Tool, toToolsJSONSchema } from "assistant-stream";
 import type { ReadonlyJSONObject } from "assistant-stream/utils";
@@ -99,29 +97,17 @@ export type AgUiMessage =
       error?: string;
     };
 
-type ToolCallPart = {
-  type: "tool-call";
-  toolCallId?: string;
-  toolName: string;
-  argsText?: string;
-  args?: ReadonlyJSONObject;
+type CoreToolCallPart = Extract<
+  Exclude<CoreThreadMessageLike["content"], string>[number],
+  { type: "tool-call" }
+>;
+
+type ToolCallPart = Omit<CoreToolCallPart, "result" | "isError"> & {
   result?: unknown;
-  artifact?: unknown;
-  isError?: boolean;
-  modelContent?: readonly ToolModelContentPart[];
+  isError?: boolean | undefined;
   unstable_toolMessageId?: string;
-  mcp?: ToolCallMessagePartMcpMetadata;
-  messages?: readonly ThreadMessage[];
-  approval?: CoreToolCallPartApproval;
 };
 
-type CoreToolCallPartApproval = NonNullable<
-  Extract<
-    Exclude<CoreThreadMessageLike["content"], string>[number],
-    { type: "tool-call" }
-  >["approval"]
->;
-
 const isObject = (value: unknown): value is Record<string, unknown> =>
   typeof value === "object" && value !== null;
 
@@ -215,6 +201,9 @@ function normalizeToolCall(part: ToolCallPart): {
   };
 }
 
+const isExportableNestedToolCall = (part: { readonly toolCallId?: unknown }) =>
+  typeof part.toolCallId === "string" && !part.toolCallId.startsWith("a2ui:");
+
 function extractText(content: unknown): string {
   if (typeof content === "string") return content;
   if (!Array.isArray(content)) return "";
@@ -1316,7 +1305,16 @@ function convertAssistantMessage(
     part: ToolCallPart;
   }[] = [];
   for (const { part } of toolCalls) {
-    collectNestedToolCalls(part, nestedToolCalls);
+    for (const { part: nestedToolCall } of walkToolCallTree(
+      part.messages ?? [],
+      { shouldDescend: isExportableNestedToolCall },
+    )) {
+      if (!isExportableNestedToolCall(nestedToolCall)) continue;
+      nestedToolCalls.push({
+        ...normalizeToolCall(nestedToolCall),
+        part: nestedToolCall,
+      });
+    }
   }
 
   converted.push({
@@ -1348,28 +1346,6 @@ function convertAssistantMessage(
   }
 }
 
-function collectNestedToolCalls(
-  part: ToolCallPart,
-  out: { id: string; call: AgUiToolCall; part: ToolCallPart }[],
-): void {
-  for (const nested of part.messages ?? []) {
-    if (!isObject(nested) || nested.role !== "assistant") continue;
-    const nestedContent = Array.isArray(nested.content) ? nested.content : [];
-    for (const nestedPart of nestedContent) {
-      if (!isObject(nestedPart) || nestedPart.type !== "tool-call") continue;
-      const nestedToolCall = nestedPart as ToolCallPart;
-      if (
-        typeof nestedToolCall.toolCallId !== "string" ||
-        nestedToolCall.toolCallId.startsWith("a2ui:")
-      ) {
-        continue;
-      }
-      out.push({ ...normalizeToolCall(nestedToolCall), part: nestedToolCall });
-      collectNestedToolCalls(nestedToolCall, out);
-    }
-  }
-}
-
 function emitToolResult(
   toolCallId: string,
   part: ToolCallPart,

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

@@ -3,6 +3,7 @@
 import {
   isMcpAppUri,
   type ChatModelRunResult,
+  type MessagePartTiming,
   type MessageStatus,
   type MessageTiming,
   type ThreadAssistantMessagePart,
@@ -217,7 +218,10 @@ export class RunAggregator {
   private interrupts: AgUiInterrupt[] | undefined;
   private readonly textParts = new Map<string, { buffer: string }>();
   private readonly activeTextMessageIdByScope = new Map<string, string>();
-  private readonly reasoningParts = new Map<string, string>(); // key → buffer
+  private readonly reasoningParts = new Map<
+    string,
+    { buffer: string; timing?: MessagePartTiming }
+  >();
   private readonly reasoningSignatures = new Map<string, string>();
   private readonly reasoningSignatureIds = new Map<string, string>();
   // Signatures captured while thinking is hidden have no block to live on;
@@ -396,6 +400,7 @@ export class RunAggregator {
           this.scopeOf("subagentRunId" in event ? event : {}),
           "messageId" in event ? event.messageId : undefined,
           event.type === "REASONING_MESSAGE_START",
+          event.timestamp,
         );
         break;
       case "REASONING_ENCRYPTED_VALUE":
@@ -454,11 +459,11 @@ export class RunAggregator {
         break;
       case "THINKING_TEXT_MESSAGE_END":
       case "THINKING_END":
-        this.handleReasoningEnd(ROOT_SCOPE);
+        this.handleReasoningEnd(ROOT_SCOPE, event.timestamp);
         break;
       case "REASONING_MESSAGE_END":
       case "REASONING_END":
-        this.handleReasoningEnd(this.scopeOf(event));
+        this.handleReasoningEnd(this.scopeOf(event), event.timestamp);
         break;
 
       case "TOOL_CALL_START": {
@@ -1222,7 +1227,8 @@ export class RunAggregator {
 
       if (part.kind === "reasoning") {
         if (this.showThinking) {
-          const buffer = this.reasoningParts.get(part.key) ?? "";
+          const entry = this.reasoningParts.get(part.key);
+          const buffer = entry?.buffer ?? "";
           const isActive =
             this.activeReasoningKeyByScope.get(rawScope) === part.key;
           if (buffer.length > 0 || isActive) {
@@ -1235,6 +1241,7 @@ export class RunAggregator {
             snapshot.push({
               type: "reasoning",
               text: buffer,
+              ...(entry?.timing ? { timing: entry.timing } : {}),
               ...(Object.keys(meta).length > 0
                 ? { providerMetadata: { [AG_UI_METADATA_NAMESPACE]: meta } }
                 : {}),
@@ -1441,6 +1448,7 @@ export class RunAggregator {
     scope: string,
     messageId?: string,
     isMessageId = false,
+    timestamp?: number,
   ): void {
     if (!this.showThinking) {
       // Hidden-reasoning bookkeeping stays root-only for this task — a
@@ -1473,13 +1481,23 @@ export class RunAggregator {
       this.reasoningMessageIds.set(key, messageId);
     }
     if (!this.reasoningParts.has(key)) {
-      this.reasoningParts.set(key, "");
+      this.reasoningParts.set(key, { buffer: "" });
       this.partOrder.push(
         scope === ROOT_SCOPE
           ? { kind: "reasoning", key }
           : { kind: "reasoning", key, subagentRunId: scope },
       );
     }
+    const entry = this.reasoningParts.get(key);
+    if (entry?.timing) {
+      entry.timing = { startedAt: entry.timing.startedAt };
+    } else if (
+      entry &&
+      typeof timestamp === "number" &&
+      Number.isFinite(timestamp)
+    ) {
+      entry.timing = { startedAt: timestamp };
+    }
     this.activeReasoningKeyByScope.set(scope, key);
     this.emit();
   }
@@ -1506,15 +1524,25 @@ export class RunAggregator {
     }
     const key = this.activeReasoningKeyByScope.get(scope);
     if (!key) return;
-    this.reasoningParts.set(key, (this.reasoningParts.get(key) ?? "") + delta);
+    const entry = this.reasoningParts.get(key);
+    if (entry) entry.buffer += delta;
     this.emit();
   }
 
-  private handleReasoningEnd(scope: string): void {
+  private handleReasoningEnd(scope: string, timestamp?: number): void {
     if (!this.showThinking) {
       if (scope === ROOT_SCOPE) this.hiddenActiveReasoning = "none";
       return;
     }
+    const key = this.activeReasoningKeyByScope.get(scope);
+    const entry = key ? this.reasoningParts.get(key) : undefined;
+    if (
+      entry?.timing &&
+      typeof timestamp === "number" &&
+      Number.isFinite(timestamp)
+    ) {
+      entry.timing = { ...entry.timing, completedAt: timestamp };
+    }
     this.activeReasoningKeyByScope.delete(scope);
     this.emit();
   }

package/src/runtime/adapter/subscriber.ts

@@ -188,6 +188,7 @@ export const createAgUiSubscriber = (
       dispatch({ type: "RUN_FINISHED", runId });
     },
     onRunFailed: ({ error }) => {
+      if (runFinishedDispatched) return;
       runFinishedDispatched = true;
       onRunFailed?.(error);
       if (isAbortError(error)) {

package/src/runtime/event-parser.ts

@@ -146,6 +146,10 @@ export const parseAgUiEvent = (
 
   const getString = (key: string) =>
     isString(payload[key]) ? (payload[key] as string) : undefined;
+  const timestamp =
+    typeof payload.timestamp === "number" && Number.isFinite(payload.timestamp)
+      ? payload.timestamp
+      : undefined;
 
   switch (typeValue) {
     case "RUN_STARTED": {
@@ -219,24 +223,31 @@ export const parseAgUiEvent = (
     case "THINKING_START":
       return withOptional(
         { type: "THINKING_START" as const },
-        { title: getString("title") },
+        { title: getString("title"), timestamp },
       );
     case "THINKING_TEXT_MESSAGE_START":
-      return { type: "THINKING_TEXT_MESSAGE_START" };
+      return withOptional(
+        { type: "THINKING_TEXT_MESSAGE_START" as const },
+        { timestamp },
+      );
     case "THINKING_TEXT_MESSAGE_CONTENT": {
       const delta = getString("delta") ?? "";
       return { type: "THINKING_TEXT_MESSAGE_CONTENT", delta };
     }
     case "THINKING_TEXT_MESSAGE_END":
-      return { type: "THINKING_TEXT_MESSAGE_END" };
+      return withOptional(
+        { type: "THINKING_TEXT_MESSAGE_END" as const },
+        { timestamp },
+      );
     case "THINKING_END":
-      return { type: "THINKING_END" };
+      return withOptional({ type: "THINKING_END" as const }, { timestamp });
     case "REASONING_START":
       return withOptional(
         { type: "REASONING_START" as const },
         {
           messageId: getString("messageId"),
           subagentRunId: getString("subagentRunId"),
+          timestamp,
         },
       );
     case "REASONING_MESSAGE_START":
@@ -245,6 +256,7 @@ export const parseAgUiEvent = (
         {
           messageId: getString("messageId"),
           subagentRunId: getString("subagentRunId"),
+          timestamp,
         },
       );
     case "REASONING_MESSAGE_CONTENT": {
@@ -263,6 +275,7 @@ export const parseAgUiEvent = (
         {
           messageId: getString("messageId"),
           subagentRunId: getString("subagentRunId"),
+          timestamp,
         },
       );
     case "REASONING_ENCRYPTED_VALUE": {
@@ -292,6 +305,7 @@ export const parseAgUiEvent = (
         {
           messageId: getString("messageId"),
           subagentRunId: getString("subagentRunId"),
+          timestamp,
         },
       );
     case "TOOL_CALL_START": {

package/src/runtime/types.ts

@@ -180,16 +180,22 @@ export type AgUiEvent =
       delta: string;
       subagentRunId?: string;
     }
-  | { type: "THINKING_START"; title?: string }
-  | { type: "THINKING_TEXT_MESSAGE_START" }
+  | { type: "THINKING_START"; title?: string; timestamp?: number }
+  | { type: "THINKING_TEXT_MESSAGE_START"; timestamp?: number }
   | { type: "THINKING_TEXT_MESSAGE_CONTENT"; delta: string }
-  | { type: "THINKING_TEXT_MESSAGE_END" }
-  | { type: "THINKING_END" }
-  | { type: "REASONING_START"; messageId?: string; subagentRunId?: string }
+  | { type: "THINKING_TEXT_MESSAGE_END"; timestamp?: number }
+  | { type: "THINKING_END"; timestamp?: number }
+  | {
+      type: "REASONING_START";
+      messageId?: string;
+      subagentRunId?: string;
+      timestamp?: number;
+    }
   | {
       type: "REASONING_MESSAGE_START";
       messageId?: string;
       subagentRunId?: string;
+      timestamp?: number;
     }
   | {
       type: "REASONING_MESSAGE_CONTENT";
@@ -201,6 +207,7 @@ export type AgUiEvent =
       type: "REASONING_MESSAGE_END";
       messageId?: string;
       subagentRunId?: string;
+      timestamp?: number;
     }
   | {
       type: "REASONING_ENCRYPTED_VALUE";
@@ -209,7 +216,12 @@ export type AgUiEvent =
       encryptedValue: string;
       subagentRunId?: string;
     }
-  | { type: "REASONING_END"; messageId?: string; subagentRunId?: string }
+  | {
+      type: "REASONING_END";
+      messageId?: string;
+      subagentRunId?: string;
+      timestamp?: number;
+    }
   | {
       type: "TOOL_CALL_START";
       toolCallId: string;

package/src/useAgUiRuntime.fast-refresh.test.tsx

@@ -3,7 +3,10 @@
 import { act, Activity, version, type ComponentType } from "react";
 import { afterAll, afterEach, expect, it, vi } from "vitest";
 import type { HttpAgent } from "@ag-ui/client";
-import type { AssistantRuntime } from "@assistant-ui/core";
+import type {
+  AssistantRuntime,
+  MessageQueueController,
+} from "@assistant-ui/core";
 import { AgUiThreadRuntimeCore } from "./runtime/AgUiThreadRuntimeCore";
 import type { UseAgUiRuntimeOptions } from "./runtime/types";
 import { useAgUiRuntime, type AgUiAssistantRuntime } from "./useAgUiRuntime";
@@ -11,6 +14,34 @@ import { useAgUiRuntime, type AgUiAssistantRuntime } from "./useAgUiRuntime";
 const onReact18 = version.startsWith("18.");
 
 const base = vi.hoisted(() => ({ version: 0 }));
+const queue = vi.hoisted(() => ({
+  controller: null as MessageQueueController | null,
+  busy: vi.fn(),
+  idle: vi.fn(),
+}));
+vi.mock("@assistant-ui/core", async (importOriginal) => {
+  const actual = await importOriginal<typeof import("@assistant-ui/core")>();
+  return {
+    ...actual,
+    createMessageQueue: (
+      ...args: Parameters<typeof actual.createMessageQueue>
+    ) => {
+      const controller = actual.createMessageQueue(...args);
+      queue.controller = {
+        ...controller,
+        notifyBusy: () => {
+          queue.busy();
+          controller.notifyBusy();
+        },
+        notifyIdle: () => {
+          queue.idle();
+          controller.notifyIdle();
+        },
+      };
+      return queue.controller;
+    },
+  };
+});
 vi.mock("@assistant-ui/core/react", async (importOriginal) => {
   const actual =
     await importOriginal<typeof import("@assistant-ui/core/react")>();
@@ -61,6 +92,7 @@ afterEach(() => {
   fiberRoots.clear();
   vi.restoreAllMocks();
   base.version = 0;
+  queue.controller = null;
 });
 afterAll(() => vi.unstubAllGlobals());
 
@@ -112,6 +144,80 @@ it("keeps the wrapper across refresh and replaces it when the base runtime chang
   expect(runtime).not.toBe(wrapper);
 });
 
+it("emits queue edges once across refresh and dispatches the queued message", async () => {
+  let finishFirst!: () => void;
+  const firstRun = new Promise<void>((resolve) => {
+    finishFirst = resolve;
+  });
+  const runAgent = vi.fn(
+    async (_input: unknown, subscriber: { onRunFinalized?: () => void }) => {
+      if (runAgent.mock.calls.length === 1) await firstRun;
+      subscriber.onRunFinalized?.();
+    },
+  );
+  options = {
+    agent: { runAgent, abortRun: vi.fn() } as unknown as HttpAgent,
+    unstable_enableMessageQueue: true,
+  };
+  const { Before, After } = createProbe();
+  const view = render(<Before />);
+  expect(queue.idle).toHaveBeenCalledTimes(1);
+
+  await act(async () => {
+    await runtime.thread.append({
+      role: "user",
+      content: [{ type: "text", text: "first" }],
+    });
+  });
+  await waitFor(() => expect(runAgent).toHaveBeenCalledTimes(1));
+  expect(queue.busy).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,
+    });
+  });
+  expect(runAgent).toHaveBeenCalledTimes(1);
+  expect(
+    [
+      ...queue.controller!.adapter.items,
+      ...queue.controller!.adapter.steerItems,
+    ].map((item) => item.prompt),
+  ).toEqual(["second"]);
+
+  await refresh(Before, After);
+  expect(queue.busy).toHaveBeenCalledTimes(1);
+  expect(queue.idle).toHaveBeenCalledTimes(1);
+
+  await act(async () => {
+    finishFirst();
+    await firstRun;
+  });
+  await waitFor(() => expect(runAgent).toHaveBeenCalledTimes(2));
+  await waitFor(() =>
+    expect([
+      ...queue.controller!.adapter.items,
+      ...queue.controller!.adapter.steerItems,
+    ]).toHaveLength(0),
+  );
+  await waitFor(() => expect(queue.busy).toHaveBeenCalledTimes(2));
+  await waitFor(() => expect(queue.idle).toHaveBeenCalledTimes(3));
+
+  await refresh(Before, After);
+  expect(queue.busy).toHaveBeenCalledTimes(2);
+  expect(queue.idle).toHaveBeenCalledTimes(3);
+
+  const firstController = queue.controller;
+  options = { ...options, unstable_enableMessageQueue: false };
+  view.rerender(<After />);
+  options = { ...options, unstable_enableMessageQueue: true };
+  view.rerender(<After />);
+  expect(queue.controller).not.toBe(firstController);
+  expect(queue.idle).toHaveBeenCalledTimes(4);
+});
+
 const expectInFlightRunAcrossRefresh = async (
   exit: "hidden" | "unmount",
   inActivity: boolean,

package/src/useAgUiRuntime.queue.test.tsx

@@ -96,11 +96,67 @@ describe("useAgUiRuntime unstable_enableMessageQueue", () => {
     );
   });
 
+  it("holds queued sends while sending is disabled", async () => {
+    const { agent, runAgent, release } = gatedAgent();
+
+    const { result, rerender } = renderHook(
+      ({ isSendDisabled }) =>
+        useAgUiRuntime({
+          agent,
+          unstable_enableMessageQueue: true,
+          isSendDisabled,
+        }),
+      { initialProps: { isSendDisabled: false } },
+    );
+    mount(result.current);
+
+    await act(async () => {
+      await result.current.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "first" }],
+      });
+    });
+    await waitFor(() => expect(runAgent).toHaveBeenCalledTimes(1));
+
+    await act(async () => {
+      await result.current.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "second" }],
+        parentId: result.current.thread.getState().messages.at(-1)?.id ?? null,
+      });
+    });
+    await waitFor(() =>
+      expect(screen.getByTestId("queued").textContent).toBe("second"),
+    );
+    rerender({ isSendDisabled: true });
+
+    await act(async () => {
+      release();
+    });
+    await waitFor(() =>
+      expect(result.current.thread.getState().isRunning).toBe(false),
+    );
+    expect(runAgent).toHaveBeenCalledTimes(1);
+    expect(screen.getByTestId("queued").textContent).toBe("second");
+
+    rerender({ isSendDisabled: false });
+    await waitFor(() => expect(runAgent).toHaveBeenCalledTimes(2));
+    await waitFor(() =>
+      expect(screen.getByTestId("queued").textContent).toBe(""),
+    );
+  });
+
   it("preserves queued sends when the active run is cancelled", async () => {
     const { agent, runAgent, release } = gatedAgent();
 
-    const { result } = renderHook(() =>
-      useAgUiRuntime({ agent, unstable_enableMessageQueue: true }),
+    const { result, rerender } = renderHook(
+      ({ isSendDisabled }) =>
+        useAgUiRuntime({
+          agent,
+          unstable_enableMessageQueue: true,
+          isSendDisabled,
+        }),
+      { initialProps: { isSendDisabled: false } },
     );
     mount(result.current);
 
@@ -135,6 +191,11 @@ describe("useAgUiRuntime unstable_enableMessageQueue", () => {
       expect(screen.getByTestId("queued").textContent).toBe("second"),
     );
 
+    rerender({ isSendDisabled: true });
+    rerender({ isSendDisabled: false });
+    expect(runAgent).toHaveBeenCalledTimes(1);
+    expect(screen.getByTestId("queued").textContent).toBe("second");
+
     await act(async () => {
       await result.current.thread.append({
         role: "user",

package/src/useAgUiRuntime.ts

@@ -114,6 +114,10 @@ export function useAgUiRuntime(
   // A dispatch whose count never moves was never observed as busy, so no
   // falling edge is coming to release the queue.
   const busyEdgesRef = useRef(0);
+  const lastQueueEdgeRef = useRef<{
+    controller: MessageQueueController;
+    busy: boolean;
+  } | null>(null);
   if (options.unstable_enableMessageQueue && !queueRef.current) {
     queueRef.current = createMessageQueue({
       run: (message) => {
@@ -172,11 +176,26 @@ export function useAgUiRuntime(
   // refusal would lose the message instead of keeping it visible.
   const queueBusy = isRunning || core.getPendingInterrupts() !== null;
   useEffect(() => {
+    if (options.isSendDisabled === true) {
+      queueController?.hold();
+    } else {
+      queueController?.release();
+    }
+  }, [options.isSendDisabled, queueController]);
+
+  useEffect(() => {
+    if (!queueController) return;
+    if (
+      lastQueueEdgeRef.current?.controller === queueController &&
+      lastQueueEdgeRef.current.busy === queueBusy
+    )
+      return;
+    lastQueueEdgeRef.current = { controller: queueController, busy: queueBusy };
     if (queueBusy) {
       busyEdgesRef.current++;
-      queueController?.notifyBusy();
+      queueController.notifyBusy();
     } else {
-      queueController?.notifyIdle();
+      queueController.notifyIdle();
     }
   }, [queueBusy, queueController]);