@assistant-ui/core 0.3.18 → 0.3.20

package/src/react/runtimes/assistant-transport/transport-scheduling.test.ts

@@ -1,8 +1,15 @@
 // @vitest-environment jsdom
 
-import { act, render, renderHook, waitFor } from "@testing-library/react";
+import { act, renderHook, waitFor } from "@testing-library/react";
 import { describe, expect, it, vi } from "vitest";
-import { createElement, StrictMode, useLayoutEffect, useRef } from "react";
+import {
+  createElement,
+  StrictMode,
+  useEffect,
+  useLayoutEffect,
+  useRef,
+} from "react";
+import { createRoot } from "react-dom/client";
 import { useCommandQueue } from "./commandQueue";
 import { useRunManager } from "./runManager";
 import type { AssistantTransportCommand } from "./types";
@@ -85,52 +92,70 @@ const useTransportSchedulingHarness = (
 };
 
 describe("assistant transport scheduling contracts", () => {
-  it("uses current callbacks when scheduled by a descendant layout effect", async () => {
-    const onRunA = vi.fn(async () => {});
-    const onRunB = vi.fn(async () => {});
-    // act flushes passive effects before real microtasks, so preserve the
-    // browser's microtask-before-passive-effect ordering at the layout boundary.
-    const queueMicrotaskSpy = vi
-      .spyOn(globalThis, "queueMicrotask")
-      .mockImplementation((callback) => callback());
-
-    const Scheduler = ({
-      schedule,
-      enabled,
-    }: {
-      schedule: () => void;
-      enabled: boolean;
-    }) => {
-      useLayoutEffect(() => {
-        if (enabled) schedule();
-      }, [enabled, schedule]);
-      return null;
-    };
-    const Probe = ({
-      enabled,
-      onRun,
-    }: {
-      enabled: boolean;
-      onRun: (signal: AbortSignal) => Promise<void>;
-    }) => {
-      const runManager = useRunManager({ onRun });
-      return createElement(Scheduler, {
-        enabled,
-        schedule: runManager.schedule,
-      });
-    };
-
+  it("reads the committed callbacks when a run settles inside a yielded commit", async () => {
+    // A render past the scheduler's 5 ms frame budget yields before passive
+    // effects, and act would drain them first, so this renders outside act.
+    vi.stubGlobal("IS_REACT_ACT_ENVIRONMENT", false);
     try {
-      const view = render(
-        createElement(Probe, { enabled: false, onRun: onRunA }),
+      const events: string[] = [];
+      const pendingRuns: (() => void)[] = [];
+      let schedule!: () => void;
+      const settleRun = () => {
+        pendingRuns.shift()?.();
+        schedule();
+      };
+
+      const Settle = ({ onLayout }: { onLayout: (() => void) | undefined }) => {
+        useLayoutEffect(() => {
+          onLayout?.();
+        }, [onLayout]);
+        return null;
+      };
+      const Probe = ({
+        label,
+        renderMs,
+        onLayout,
+      }: {
+        label: string;
+        renderMs: number;
+        onLayout?: () => void;
+      }) => {
+        const runManager = useRunManager({
+          onRun: () => {
+            events.push(`run:${label}`);
+            return new Promise<void>((resolve) => pendingRuns.push(resolve));
+          },
+          onFinish: () => events.push(`finish:${label}`),
+        });
+        schedule = runManager.schedule;
+        useEffect(() => {
+          events.push(`passive:${label}`);
+        }, [label]);
+        const renderEnd = performance.now() + renderMs;
+        while (performance.now() < renderEnd) {}
+        return createElement(Settle, { onLayout });
+      };
+
+      const root = createRoot(document.createElement("div"));
+      root.render(createElement(Probe, { label: "A", renderMs: 0 }));
+      await vi.waitFor(() => expect(events).toEqual(["passive:A"]));
+      schedule();
+      await vi.waitFor(() => expect(events).toEqual(["passive:A", "run:A"]));
+      root.render(
+        createElement(Probe, { label: "B", renderMs: 30, onLayout: settleRun }),
       );
-      view.rerender(createElement(Probe, { enabled: true, onRun: onRunB }));
-      await act(async () => {});
-
-      expect(onRunB).toHaveBeenCalledTimes(1);
-      expect(onRunA).not.toHaveBeenCalled();
+      await vi.waitFor(() => expect(events).toContain("passive:B"));
+      root.unmount();
+
+      expect(events).toEqual([
+        "passive:A",
+        "run:A",
+        "finish:B",
+        "run:B",
+        "passive:B",
+      ]);
     } finally {
-      queueMicrotaskSpy.mockRestore();
+      vi.unstubAllGlobals();
     }
   });
 
@@ -196,7 +221,9 @@ describe("assistant transport scheduling contracts", () => {
         onRun: () => {
           throw new Error("network error");
         },
-        onError: (commands) => seen.push(commands),
+        onError: (commands) => {
+          seen.push(commands);
+        },
       }),
     );
 

package/src/react/runtimes/assistant-transport/useAssistantTransportRuntime.test.tsx

@@ -167,7 +167,13 @@ describe("useAssistantTransportRuntime", () => {
               ],
               createdAt: new Date(0),
               status: { type: "requires-action", reason: "tool-calls" },
-              metadata: { custom: {} },
+              metadata: {
+                unstable_state: {},
+                unstable_annotations: [],
+                unstable_data: [],
+                steps: [],
+                custom: {},
+              },
             },
           ],
           isRunning: meta.isSending,
@@ -203,6 +209,65 @@ describe("useAssistantTransportRuntime", () => {
     },
   );
 
+  it("forwards a tool response's modelContent to the outbound command", async () => {
+    const fetchMock = installFetch();
+    const { aui } = mountRuntime({
+      converter: (_state, meta) => ({
+        messages: [
+          {
+            id: "a1",
+            role: "assistant",
+            content: [
+              {
+                type: "tool-call",
+                toolCallId: "tc1",
+                toolName: "probe",
+                args: {},
+                argsText: "{}",
+              },
+            ],
+            createdAt: new Date(0),
+            status: { type: "requires-action", reason: "tool-calls" },
+            metadata: {
+              unstable_state: {},
+              unstable_annotations: [],
+              unstable_data: [],
+              steps: [],
+              custom: {},
+            },
+          },
+        ],
+        isRunning: meta.isSending,
+      }),
+    });
+
+    await waitFor(() =>
+      expect(
+        aui()
+          .thread.message({ id: "a1" })
+          .part({ toolCallId: "tc1" })
+          .getState().type,
+      ).toBe("tool-call"),
+    );
+
+    const modelContent = [{ type: "text" as const, text: "done" }];
+    act(() => {
+      aui()
+        .thread.message({ id: "a1" })
+        .part({ toolCallId: "tc1" })
+        .addToolResult(new ToolResponse({ result: "ok", modelContent }));
+    });
+
+    await waitFor(() => expect(fetchMock.requests).toHaveLength(1));
+    expect(fetchMock.requests[0]!.body["commands"][0]).toHaveProperty(
+      "modelContent",
+      modelContent,
+    );
+
+    act(() => fetchMock.servers[0]!.close());
+    await waitFor(() => expect(aui().thread.getState().isRunning).toBe(false));
+  });
+
   it.each(["throws", "rejects"] as const)(
     "cancels the response body when onResponse %s",
     async (failureMode) => {

package/src/react/runtimes/cloud/AssistantCloudThreadHistoryAdapter.test.ts

@@ -28,8 +28,8 @@ describe("extractAuiV0", () => {
   it("reports a terminal message once with its usage", () => {
     const result = extractAuiV0(auiV0Message({ type: "complete" }));
     expect(result?.status).toBe("completed");
-    expect(result?.inputTokens).toBe(1);
-    expect(result?.outputTokens).toBe(2);
+    expect(result?.usage?.inputTokens).toBe(1);
+    expect(result?.usage?.outputTokens).toBe(2);
   });
 
   it("sums step usage reported under the AI SDK v7 token details", () => {
@@ -57,22 +57,28 @@ describe("extractAuiV0", () => {
       },
     });
 
-    expect(result?.inputTokens).toBe(15);
-    expect(result?.outputTokens).toBe(6);
-    expect(result?.cachedInputTokens).toBe(9);
-    expect(result?.reasoningTokens).toBe(2);
+    expect(result?.usage).toEqual({
+      inputTokens: 15,
+      outputTokens: 6,
+      cachedInputTokens: 9,
+      reasoningTokens: 2,
+    });
     expect(result?.steps).toEqual([
       {
-        input_tokens: 10,
-        output_tokens: 4,
-        reasoning_tokens: 1,
-        cached_input_tokens: 6,
+        usage: {
+          inputTokens: 10,
+          outputTokens: 4,
+          inputTokenDetails: { cacheReadTokens: 6 },
+          outputTokenDetails: { reasoningTokens: 1 },
+        },
       },
       {
-        input_tokens: 5,
-        output_tokens: 2,
-        reasoning_tokens: 1,
-        cached_input_tokens: 3,
+        usage: {
+          inputTokens: 5,
+          outputTokens: 2,
+          inputTokenDetails: { cacheReadTokens: 3 },
+          outputTokenDetails: { reasoningTokens: 1 },
+        },
       },
     ]);
   });

package/src/react/runtimes/cloud/AssistantCloudThreadHistoryAdapter.test.tsx

@@ -178,6 +178,20 @@ describe("useAssistantCloudThreadHistoryAdapter", () => {
       threadId: "thread-1",
       messageId: "local-message-1",
     });
+    listeners.get("thread.toolApprovalAnswered")!({
+      threadId: "thread-1",
+      messageId: "local-message-1",
+      toolCallId: "tool-approved",
+      toolName: "send_email",
+      approved: true,
+    });
+    listeners.get("thread.toolApprovalAnswered")!({
+      threadId: "thread-1",
+      messageId: "local-message-1",
+      toolCallId: "tool-rejected",
+      toolName: "delete_account",
+      approved: false,
+    });
     listeners.get("composer.attachmentAdd")!({
       threadId: "thread-1",
       contentType: "image/png",
@@ -202,6 +216,11 @@ describe("useAssistantCloudThreadHistoryAdapter", () => {
       messageId: "local-message-1",
       reason: "error",
     });
+    listeners.get("message.error")!({
+      threadId: "thread-1",
+      messageId: "local-message-2",
+      reason: "error",
+    });
     threadState.isEmpty = true;
     threadState.suggestions = [{ prompt: "one" }, { prompt: "two" }];
     notify!();
@@ -232,6 +251,16 @@ describe("useAssistantCloudThreadHistoryAdapter", () => {
         expect.objectContaining({ kind: "message_regenerated" }),
         expect.objectContaining({ kind: "branch_switched" }),
         expect.objectContaining({ kind: "message_copied" }),
+        expect.objectContaining({
+          kind: "tool_approved",
+          message_id: "remote-message-1",
+          props: { toolCallId: "tool-approved", toolName: "send_email" },
+        }),
+        expect.objectContaining({
+          kind: "tool_rejected",
+          message_id: "remote-message-1",
+          props: { toolCallId: "tool-rejected", toolName: "delete_account" },
+        }),
         expect.objectContaining({
           kind: "attachment_added",
           props: { type: "image/png" },
@@ -250,6 +279,12 @@ describe("useAssistantCloudThreadHistoryAdapter", () => {
         expect.objectContaining({ kind: "suggestions_shown", value: 2 }),
       ]),
     );
+    expect(
+      vi
+        .mocked(cloud.events.track)
+        .mock.calls.map(([event]) => event)
+        .filter((event) => event.kind === "error_shown"),
+    ).toHaveLength(1);
   });
 
   it("refreshes formatted persistence when the Cloud client changes", async () => {
@@ -369,6 +404,31 @@ describe("useAssistantCloudThreadHistoryAdapter", () => {
     });
   });
 
+  it("forwards a feedback comment to the cloud", async () => {
+    mocks.aui = mocks.makeClient("thread-1");
+    const cloud = makeCloud();
+    const cloudRef = { current: cloud };
+    const { result } = renderHook(() =>
+      useAssistantCloudThreadHistoryAdapter(cloudRef),
+    );
+    const message = makeAssistantMessage("local-message-1");
+
+    await result.current.append({ parentId: null, message });
+    result.current.feedback.submit({
+      message,
+      type: "negative",
+      comment: "Quoted the wrong date",
+    });
+
+    await waitFor(() => {
+      expect(cloud.threads.messages.feedback).toHaveBeenCalledWith(
+        "thread-1",
+        "remote-message-1",
+        { type: "negative", comment: "Quoted the wrong date" },
+      );
+    });
+  });
+
   it("warns and skips feedback before the thread has a remote ID", async () => {
     mocks.aui = mocks.makeClient(undefined, "local-thread", "thread-1");
     const cloud = makeCloud();
@@ -501,7 +561,7 @@ describe("useAssistantCloudThreadHistoryAdapter", () => {
         },
         timing: {
           streamStartTime: 100,
-          firstTokenTime: 155.6,
+          firstTokenTime: 55.6,
           totalChunks: 1,
           toolCallCount: 0,
         },
@@ -688,6 +748,7 @@ describe("useAssistantCloudThreadHistoryAdapter", () => {
                 tool_source: "mcp",
               }),
             ],
+            finish_reason: "tool-calls",
           },
           {
             tool_calls: [
@@ -696,6 +757,7 @@ describe("useAssistantCloudThreadHistoryAdapter", () => {
                 tool_source: "frontend",
               }),
             ],
+            finish_reason: "tool-calls",
           },
         ],
         trace_id: "00112233445566778899aabbccddeeff",
@@ -705,6 +767,189 @@ describe("useAssistantCloudThreadHistoryAdapter", () => {
     );
   });
 
+  it("reports a single ai-sdk/v6 step with its tool calls", () => {
+    mocks.aui = mocks.makeClient("thread-1");
+    const cloud = makeCloud();
+    const { result } = renderHook(() =>
+      useAssistantCloudThreadHistoryAdapter({ current: cloud }),
+    );
+    const formatted = result.current.withFormat({
+      format: "ai-sdk/v6",
+      encode: ({ message }) => message,
+      decode: ({ parent_id, content }) => ({
+        parentId: parent_id,
+        message: content as { id: string },
+      }),
+      getId: (message: { id: string }) => message.id,
+    });
+
+    formatted.reportTelemetry([
+      {
+        parentId: null,
+        message: {
+          id: "message-1",
+          role: "assistant",
+          parts: [
+            { type: "step-start" },
+            {
+              type: "tool-search",
+              toolCallId: "search-1",
+              input: { query: "test" },
+              output: { result: "ok" },
+            },
+          ],
+        },
+      },
+    ]);
+
+    expect(cloud.runs.report).toHaveBeenCalledWith(
+      expect.objectContaining({
+        total_steps: 1,
+        steps: [
+          {
+            tool_calls: [
+              expect.objectContaining({
+                tool_call_id: "search-1",
+                tool_source: "frontend",
+              }),
+            ],
+            finish_reason: "tool-calls",
+          },
+        ],
+      }),
+    );
+  });
+
+  it("reads the status of an ai-sdk/v6 run from its finish reason", () => {
+    mocks.aui = mocks.makeClient("thread-1");
+    const cloud = makeCloud();
+    const { result } = renderHook(() =>
+      useAssistantCloudThreadHistoryAdapter({ current: cloud }),
+    );
+    const formatted = result.current.withFormat({
+      format: "ai-sdk/v6",
+      encode: ({ message }) => message,
+      decode: ({ parent_id, content }) => ({
+        parentId: parent_id,
+        message: content as { id: string },
+      }),
+      getId: (message: { id: string }) => message.id,
+    });
+
+    formatted.reportTelemetry([
+      {
+        parentId: null,
+        message: {
+          id: "message-1",
+          role: "assistant",
+          parts: [{ type: "text", text: "cut" }],
+          metadata: { finishReason: "length" },
+        },
+      },
+    ]);
+
+    expect(cloud.runs.report).toHaveBeenCalledWith(
+      expect.objectContaining({ status: "incomplete", outcome_type: "length" }),
+    );
+  });
+
+  it("reads the error and timing of an ai-sdk/v6 run from the thread message", () => {
+    mocks.aui = mocks.makeClient("thread-1");
+    const cloud = makeCloud();
+    const { result } = renderHook(() =>
+      useAssistantCloudThreadHistoryAdapter({ current: cloud }),
+    );
+    const formatted = result.current.withFormat({
+      format: "ai-sdk/v6",
+      encode: ({ message }) => message,
+      decode: ({ parent_id, content }) => ({
+        parentId: parent_id,
+        message: content as { id: string },
+      }),
+      getId: (message: { id: string }) => message.id,
+    });
+    const base = makeAssistantMessage("assistant-1");
+    const message: ThreadAssistantMessage = {
+      ...base,
+      status: {
+        type: "incomplete",
+        reason: "error",
+        error: Object.assign(new Error("model unavailable"), {
+          code: "model_unavailable",
+        }),
+      },
+      metadata: {
+        ...base.metadata,
+        timing: {
+          streamStartTime: 100,
+          firstTokenTime: 240,
+          totalChunks: 3,
+          toolCallCount: 0,
+        },
+      },
+    };
+
+    formatted.reportTelemetry(
+      [
+        {
+          parentId: null,
+          message: {
+            id: "message-1",
+            role: "assistant",
+            parts: [{ type: "text", text: "partial" }],
+          },
+        },
+      ],
+      { message },
+    );
+
+    expect(cloud.runs.report).toHaveBeenCalledWith(
+      expect.objectContaining({
+        status: "error",
+        error: "model unavailable",
+        error_code: "model_unavailable",
+        first_token_ms: 240,
+      }),
+    );
+  });
+
+  it("reports a run that failed before any assistant message was stored", () => {
+    mocks.aui = mocks.makeClient("thread-1");
+    const cloud = makeCloud();
+    const { result } = renderHook(() =>
+      useAssistantCloudThreadHistoryAdapter({ current: cloud }),
+    );
+    const formatted = result.current.withFormat({
+      format: "ai-sdk/v6",
+      encode: ({ message }) => message,
+      decode: ({ parent_id, content }) => ({
+        parentId: parent_id,
+        message: content as { id: string },
+      }),
+      getId: (message: { id: string }) => message.id,
+    });
+    const message: ThreadAssistantMessage = {
+      ...makeAssistantMessage("assistant-1"),
+      content: [],
+      status: {
+        type: "incomplete",
+        reason: "error",
+        error: { code: "AI_APICallError", message: "upstream failed" },
+      },
+    };
+
+    formatted.reportTelemetry([], { message, durationMs: 120 });
+
+    expect(cloud.runs.report).toHaveBeenCalledWith(
+      expect.objectContaining({
+        status: "error",
+        error: "upstream failed",
+        error_code: "AI_APICallError",
+        duration_ms: 120,
+      }),
+    );
+  });
+
   it("reports the model ID carried by aui/v0 step metadata", () => {
     mocks.aui = mocks.makeClient("thread-1");
     const cloud = makeCloud();
@@ -850,3 +1095,80 @@ describe("useAssistantCloudThreadHistoryAdapter", () => {
     );
   });
 });
+
+describe("useAssistantCloudThreadHistoryAdapter load recovery", () => {
+  const userRow = (id: string, parentId: string | null, text: string) => ({
+    id,
+    thread_id: "thread-1",
+    format: "aui/v0" as const,
+    parent_id: parentId,
+    created_at: new Date(0),
+    content: {
+      role: "user",
+      content: [{ type: "text", text }],
+      metadata: { custom: {} },
+    },
+  });
+
+  const loadHistory = async (messages: unknown[]) => {
+    mocks.aui = mocks.makeClient("thread-1");
+    const cloud = makeCloud();
+    cloud.threads.messages.list = vi.fn().mockResolvedValue({ messages });
+    const cloudRef = { current: cloud };
+    const { result } = renderHook(() =>
+      useAssistantCloudThreadHistoryAdapter(cloudRef),
+    );
+    return (await result.current.load()).messages;
+  };
+
+  it("keeps a row whose stored part is malformed, and the thread below it", async () => {
+    // The cloud lists rows newest first.
+    const messages = await loadHistory([
+      userRow("msg-3", "msg-2", "after"),
+      {
+        ...userRow("msg-2", "msg-1", "ignored"),
+        content: { role: "assistant", content: [null] },
+      },
+      userRow("msg-1", null, "hello"),
+    ]);
+
+    expect(messages.map((item) => item.message.id)).toEqual([
+      "msg-1",
+      "msg-2",
+      "msg-3",
+    ]);
+    expect(messages[1]?.message.content).toEqual([]);
+  });
+
+  it("keeps a row whose stored attachment is malformed", async () => {
+    const messages = await loadHistory([
+      {
+        ...userRow("msg-1", null, "look"),
+        content: {
+          role: "user",
+          content: [{ type: "text", text: "look" }],
+          attachments: [null],
+          metadata: { custom: {} },
+        },
+      },
+    ]);
+
+    expect(messages).toHaveLength(1);
+    expect(messages[0]?.message.content).toEqual([
+      { type: "text", text: "look" },
+    ]);
+  });
+
+  it("drops an unreadable row together with the subtree below it", async () => {
+    const messages = await loadHistory([
+      userRow("msg-3", "msg-2", "orphaned"),
+      {
+        ...userRow("msg-2", "msg-1", "ignored"),
+        content: { role: "assistant", content: null },
+      },
+      userRow("msg-1", null, "hello"),
+    ]);
+
+    expect(messages.map((item) => item.message.id)).toEqual(["msg-1"]);
+  });
+});