@assistant-ui/ai-sdk 0.0.3 → 0.0.5

package/src/runtime/useChatRuntime.integration.test.tsx

@@ -1,13 +1,21 @@
 // @vitest-environment jsdom
 
-import { render, screen, waitFor } from "@testing-library/react";
+import { act, render, screen, waitFor } from "@testing-library/react";
 import { AssistantRuntimeProvider } from "@assistant-ui/core/react";
-import { useAuiState } from "@assistant-ui/store";
-import type { UIMessage } from "ai";
-import { StrictMode, useState } from "react";
+import { AuiConfig, AuiProvider, useAuiState } from "@assistant-ui/store";
+import type { AssistantRuntime } from "@assistant-ui/core";
+import { AISDKChat } from "./AISDKChat";
+import type { ChatTransport, UIMessage } from "ai";
+import { Activity, StrictMode, useState, type ReactNode } from "react";
 import { describe, expect, it } from "vitest";
 import { AssistantChatTransport } from "../transport/AssistantChatTransport";
+import {
+  createCancellableTransport,
+  createStreamHarness,
+  nextTask,
+} from "./__tests__/controlled-transport";
 import { useChatRuntime } from "./useChatRuntime";
+import { useThreadTokenUsage } from "../usage";
 
 const messages: UIMessage[] = [
   {
@@ -64,4 +72,175 @@ describe("useChatRuntime integration", () => {
       );
     });
   });
+
+  it("aborts a deleted thread's stream while the host stays mounted", async () => {
+    const { transport, getCancelCount } = createCancellableTransport();
+    const { Probe, send, isRunning, client } = createStreamHarness();
+
+    const view = render(
+      <StrictMode>
+        <StreamingApp transport={transport} probe={<Probe />} />
+      </StrictMode>,
+    );
+
+    await act(async () => send());
+    await waitFor(() => expect(isRunning()).toBe(true));
+
+    await act(async () => client().threadListItem.delete());
+
+    await waitFor(() => expect(getCancelCount()).toBe(1));
+    view.unmount();
+  });
+
+  it("aborts the in-flight transport after a real unmount", async () => {
+    const { transport, getCancelCount } = createCancellableTransport();
+    const { Probe, send, isRunning } = createStreamHarness();
+
+    const view = render(
+      <StrictMode>
+        <StreamingApp transport={transport} probe={<Probe />} />
+      </StrictMode>,
+    );
+
+    await act(async () => send());
+    await waitFor(() => expect(isRunning()).toBe(true));
+    // the Strict Mode double mount already ran a host cleanup by now
+    expect(getCancelCount()).toBe(0);
+
+    view.unmount();
+    await waitFor(() => expect(getCancelCount()).toBe(1));
+  });
+
+  it("keeps streaming while hidden and aborts when the hidden host unmounts", async () => {
+    const { transport, getCancelCount } = createCancellableTransport();
+    const { Probe, send, isRunning } = createStreamHarness();
+
+    let setMode: ((mode: "visible" | "hidden") => void) | undefined;
+    const Shell = () => {
+      const [mode, set] = useState<"visible" | "hidden">("visible");
+      setMode = set;
+      return (
+        <Activity mode={mode}>
+          <StreamingApp transport={transport} probe={<Probe />} />
+        </Activity>
+      );
+    };
+
+    const view = render(
+      <StrictMode>
+        <Shell />
+      </StrictMode>,
+    );
+
+    await act(async () => send());
+    await waitFor(() => expect(isRunning()).toBe(true));
+
+    await act(async () => setMode?.("hidden"));
+    await act(nextTask);
+    expect(getCancelCount()).toBe(0);
+    expect(isRunning()).toBe(true);
+
+    await act(async () => setMode?.("visible"));
+    await act(nextTask);
+    expect(getCancelCount()).toBe(0);
+    expect(isRunning()).toBe(true);
+
+    await act(async () => setMode?.("hidden"));
+    view.unmount();
+    await waitFor(() => expect(getCancelCount()).toBe(1));
+  });
+
+  it("aborts a nested runtime's stream when the provider above it unmounts", async () => {
+    const outer = createCancellableTransport();
+    const { transport, getCancelCount } = createCancellableTransport();
+    let nested: AssistantRuntime | undefined;
+
+    // allowNesting: the inner useChatRuntime runs its thread hook directly, as
+    // a plain React hook under the provider rather than inside a tap resource.
+    const NestedChat = () => {
+      nested = useChatRuntime({ transport });
+      return null;
+    };
+
+    const view = render(
+      <StrictMode>
+        <AuiProvider
+          config={AuiConfig({
+            threads: AISDKChat({ transport: outer.transport }),
+          })}
+        >
+          <NestedChat />
+        </AuiProvider>
+      </StrictMode>,
+    );
+
+    await waitFor(() => expect(nested).toBeDefined());
+    await act(async () => {
+      await nested!.thread.append("keep streaming");
+    });
+    await waitFor(() => expect(nested!.thread.getState().isRunning).toBe(true));
+
+    view.unmount();
+    await waitFor(() => expect(getCancelCount()).toBe(1));
+  });
+});
+
+const StreamingApp = ({
+  transport,
+  probe,
+}: {
+  transport: ChatTransport<UIMessage>;
+  probe: ReactNode;
+}) => {
+  const runtime = useChatRuntime({ transport });
+  return (
+    <AssistantRuntimeProvider runtime={runtime}>
+      {probe}
+    </AssistantRuntimeProvider>
+  );
+};
+
+const UsageProbe = () => {
+  const usage = useThreadTokenUsage();
+  return (
+    <output data-testid="total-tokens">{usage?.totalTokens ?? "none"}</output>
+  );
+};
+
+const UsageApp = () => {
+  const [transport] = useState(
+    () => new AssistantChatTransport({ api: "/api/chat" }),
+  );
+  const runtime = useChatRuntime({
+    messages: [
+      ...messages,
+      {
+        id: "assistant-with-usage",
+        role: "assistant",
+        parts: [{ type: "text", text: "Hi" }],
+        metadata: { usage: { inputTokens: 40, outputTokens: 2 } },
+      },
+    ],
+    transport,
+  });
+
+  return (
+    <AssistantRuntimeProvider runtime={runtime}>
+      <UsageProbe />
+    </AssistantRuntimeProvider>
+  );
+};
+
+describe("useThreadTokenUsage through useChatRuntime", () => {
+  it("reads usage from the message metadata a server attached", async () => {
+    render(
+      <StrictMode>
+        <UsageApp />
+      </StrictMode>,
+    );
+
+    await waitFor(() => {
+      expect(screen.getByTestId("total-tokens").textContent).toBe("42");
+    });
+  });
 });

package/src/runtime/useChatRuntime.ts

@@ -29,6 +29,7 @@ const useChatThreadRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     isMainThread,
     getThreadListItem: () =>
       aui.threadListItem.source ? aui.threadListItem : undefined,
+    stopOnClientDestroy: true,
   });
 };
 

package/src/runtime/useChatThread.ts

@@ -1,6 +1,7 @@
 "use client";
 
 import { useChat, type Chat, type UIMessage } from "@ai-sdk/react";
+import type { MessageRepository } from "@assistant-ui/core/internal";
 import {
   pickExternalStoreSharedOptions,
   type AssistantRuntime,
@@ -60,6 +61,12 @@ export type ChatThreadEnvironment<UI_MESSAGE extends UIMessage = UIMessage> = {
    * from the instance.
    */
   chat?: Chat<UI_MESSAGE> | undefined;
+  /**
+   * An externally owned per-thread message repository. Hosts that route
+   * multiple threads through one mounting pass a distinct instance per
+   * thread so histories and branches stay isolated.
+   */
+  messageRepositoryInstance?: MessageRepository | undefined;
 };
 
 const useDynamicChatTransport = <UI_MESSAGE extends UIMessage = UIMessage>(
@@ -183,6 +190,7 @@ export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
     getThreadListItem,
     stopOnClientDestroy = false,
     chat: externalChat,
+    messageRepositoryInstance,
   } = env;
 
   const defaultTransport = useMemo(() => new AssistantChatTransport(), []);
@@ -209,6 +217,9 @@ export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
     ...(onResumeToolCall && { onResumeToolCall }),
     ...(joinStrategy && { joinStrategy }),
     ...(messageRepository && { messageRepository }),
+    ...(messageRepositoryInstance && {
+      unstable_messageRepositoryInstance: messageRepositoryInstance,
+    }),
     ...(unstable_onBranchChange && { unstable_onBranchChange }),
   });
 

package/src/runtime/useExternalHistory.test.ts

@@ -1180,6 +1180,70 @@ describe("useExternalHistory persistence", () => {
     expect(reportTelemetry).not.toHaveBeenCalled();
   });
 
+  it("skips updates for a persisted message restored by a branch switch", async () => {
+    const { append, update, runCycle, flush } = createPersistenceHarness(true);
+    const completeStatus: ThreadAssistantMessage["status"] = {
+      type: "complete",
+      reason: "stop",
+    };
+    const createUserMessage = (
+      id: string,
+      inner: InnerMessage,
+    ): ThreadMessage => {
+      const message: ThreadMessage = {
+        id,
+        role: "user",
+        content: [{ type: "text", text: "hi" }],
+        attachments: [],
+        createdAt: new Date(),
+        metadata: { custom: {} },
+      };
+      bindExternalStoreMessage(message, [inner]);
+      return message;
+    };
+    const branchA = () => [
+      createUserMessage("user-a", { id: "inner-user-a", parts: ["question"] }),
+      createAssistantMessage(
+        completeStatus,
+        [{ id: "inner-assistant-a", parts: ["answer"] }],
+        "assistant-a",
+      ),
+    ];
+
+    await runCycle(branchA());
+    await waitFor(() => expect(append).toHaveBeenCalledTimes(2));
+
+    await runCycle([
+      createUserMessage("user-b", { id: "inner-user-b", parts: ["edited"] }),
+      createAssistantMessage(
+        completeStatus,
+        [{ id: "inner-assistant-b", parts: ["answer"] }],
+        "assistant-b",
+      ),
+    ]);
+    await waitFor(() => expect(append).toHaveBeenCalledTimes(4));
+
+    append.mockClear();
+    update.mockClear();
+
+    await runCycle([
+      ...branchA(),
+      createUserMessage("user-c", { id: "inner-user-c", parts: ["follow-up"] }),
+      createAssistantMessage(
+        completeStatus,
+        [{ id: "inner-assistant-c", parts: ["answer"] }],
+        "assistant-c",
+      ),
+    ]);
+    await flush();
+
+    expect(update).not.toHaveBeenCalled();
+    expect(append.mock.calls.map(([item]) => item.message.id)).toEqual([
+      "inner-user-c",
+      "inner-assistant-c",
+    ]);
+  });
+
   it("absorbs agentic flickers without losing change detection", async () => {
     const { append, update, reportTelemetry, runCycle, flush, step } =
       createPersistenceHarness(true);

package/src/runtime/useExternalHistory.ts

@@ -5,6 +5,7 @@ import type {
   ThreadHistoryAdapter,
   ThreadMessage,
   MessageFormatAdapter,
+  MessageFormatItem,
   MessageFormatRepository,
   ExportedMessageRepository,
 } from "@assistant-ui/core";
@@ -49,15 +50,10 @@ const isAwaitingToolApproval = (message: ThreadMessage) =>
   message.status?.type === "requires-action" &&
   message.status.reason === "tool-calls";
 
-const snapshotExternalMessages = <TMessage>(
-  messages: readonly ThreadMessage[],
-) =>
-  new Map<string, TMessage[]>(
-    messages.map((message) => [
-      message.id,
-      [...getExternalStoreMessages<TMessage>(message)],
-    ]),
-  );
+const encodeContent = <TMessage>(
+  storageFormatAdapter: MessageFormatAdapter<TMessage, any>,
+  item: MessageFormatItem<TMessage>,
+) => JSON.stringify(storageFormatAdapter.encode(item));
 
 export const useExternalHistory = <TMessage>(
   runtimeRef: RefObject<AssistantRuntime>,
@@ -78,9 +74,11 @@ export const useExternalHistory = <TMessage>(
   const [hasLoaded, setHasLoaded] = useState(false);
 
   const historyIds = useRef(new Set<string>());
-  const persistedInnerIds = useRef(new Set<string>());
   const deferredTelemetryIds = useRef(new Set<string>());
-  const persistedExternalMessages = useRef(new Map<string, TMessage[]>());
+  // `content` is a snapshot taken at write time rather than a re-encode of `source`, because a retained message object can be mutated in place by the runtime that produced it.
+  const persistedInnerMessages = useRef(
+    new Map<string, { source: TMessage; content: string }>(),
+  );
 
   const onSetMessagesRef = useRef(onSetMessages);
   useEffect(() => {
@@ -107,8 +105,12 @@ export const useExternalHistory = <TMessage>(
         const repo = await formatAdapter.load();
         if (repo && repo.messages.length > 0) {
           for (const m of repo.messages) {
-            persistedInnerIds.current.add(
+            persistedInnerMessages.current.set(
               storageFormatAdapter.getId(m.message),
+              {
+                source: m.message,
+                content: encodeContent(storageFormatAdapter, m),
+              },
             );
           }
           const converted = toExportedMessageRepository(toThreadMessages, repo);
@@ -129,10 +131,6 @@ export const useExternalHistory = <TMessage>(
               deferredTelemetryIds.current.add(m.message.id);
             }
           }
-          persistedExternalMessages.current =
-            snapshotExternalMessages<TMessage>(
-              converted.messages.map((m) => m.message),
-            );
         }
       } catch (error) {
         console.error("Failed to load message history:", error);
@@ -145,6 +143,9 @@ export const useExternalHistory = <TMessage>(
 
     const remoteId = optionalThreadListItem()?.getState().remoteId;
     if (!remoteId) {
+      // History loads asynchronously against the thread list item; without a
+      // remote id there is nothing to await, so the flag settles here.
+      // eslint-disable-next-line react-hooks/set-state-in-effect
       setHasLoaded(true);
       return aui.subscribe(() => {
         if (optionalThreadListItem()?.getState().remoteId) {
@@ -294,23 +295,8 @@ export const useExternalHistory = <TMessage>(
 
       persistInFlightRef.current = persistInFlightRef.current
         .then(async () => {
-          const changedRunMessageIds = new Set<string>();
-          for (const message of latest.messages) {
-            const externalMessages =
-              getExternalStoreMessages<TMessage>(message);
-            const previous = persistedExternalMessages.current.get(message.id);
-            if (
-              previous === undefined ||
-              previous.length !== externalMessages.length ||
-              externalMessages.some((item, index) => item !== previous[index])
-            ) {
-              changedRunMessageIds.add(message.id);
-            }
-          }
-
           const { messages } = latest;
           let lastInnerMessageId: string | null = null;
-          const failedUpdateIds = new Set<string>();
 
           const getLastInnerId = (msgs: TMessage[]): string | null =>
             msgs.length > 0 ? storageFormatAdapter.getId(msgs.at(-1)!) : null;
@@ -343,13 +329,7 @@ export const useExternalHistory = <TMessage>(
               continue;
             }
 
-            const isPersistedMessage = historyIds.current.has(message.id);
-            if (isPersistedMessage && !changedRunMessageIds.has(message.id)) {
-              lastInnerMessageId =
-                getLastInnerId(innerMessages) ?? lastInnerMessageId;
-              continue;
-            }
-            if (!isPersistedMessage) {
+            if (!historyIds.current.has(message.id)) {
               historyIds.current.add(message.id);
               deferredTelemetryIds.current.add(message.id);
             }
@@ -357,15 +337,31 @@ export const useExternalHistory = <TMessage>(
             const batchItems = toBatchItems(innerMessages);
             for (const item of batchItems) {
               const innerId = storageFormatAdapter.getId(item.message);
-              if (!persistedInnerIds.current.has(innerId)) {
+              const persisted = persistedInnerMessages.current.get(innerId);
+              if (!persisted) {
                 await adapter.append(item);
-                persistedInnerIds.current.add(innerId);
-              } else if (durationMs !== undefined) {
-                try {
-                  await adapter.update?.(item, innerId);
-                } catch {
-                  // A failed update drops the message from the refreshed baseline so it retries on the next run stop.
-                  failedUpdateIds.add(message.id);
+                persistedInnerMessages.current.set(innerId, {
+                  source: item.message,
+                  content: encodeContent(storageFormatAdapter, item),
+                });
+              } else if (
+                persisted.source !== item.message &&
+                durationMs !== undefined &&
+                adapter.update
+              ) {
+                const content = encodeContent(storageFormatAdapter, item);
+                if (content === persisted.content) {
+                  persisted.source = item.message;
+                } else {
+                  try {
+                    await adapter.update(item, innerId);
+                    persistedInnerMessages.current.set(innerId, {
+                      source: item.message,
+                      content,
+                    });
+                  } catch {
+                    // A failed update leaves the stale baseline behind so the next run stop retries it.
+                  }
                 }
               }
             }
@@ -378,14 +374,6 @@ export const useExternalHistory = <TMessage>(
               adapter.reportTelemetry?.(batchItems, telemetryOptions);
             }
           }
-
-          const nextSnapshot = snapshotExternalMessages<TMessage>(
-            latest.messages,
-          );
-          for (const id of failedUpdateIds) {
-            nextSnapshot.delete(id);
-          }
-          persistedExternalMessages.current = nextSnapshot;
         })
         .catch((error) => {
           console.error("Failed to persist message history:", error);
@@ -430,9 +418,8 @@ export const useExternalHistory = <TMessage>(
 
         historyIds.current.delete(messageId);
         deferredTelemetryIds.current.delete(messageId);
-        persistedExternalMessages.current.delete(messageId);
         for (const item of itemsToDelete) {
-          persistedInnerIds.current.delete(
+          persistedInnerMessages.current.delete(
             storageFormatAdapter.getId(item.message),
           );
         }

package/src/tools/generativeTools.test.ts

@@ -139,6 +139,30 @@ describe("AISDKToolkit", () => {
     mocks.createMCPClient.mockReset();
   });
 
+  it("preserves prototype-named MCP tools", async () => {
+    const prototypeTool = { inputSchema: {} };
+    mocks.tools.mockResolvedValue(
+      Object.fromEntries([["__proto__", prototypeTool]]),
+    );
+    mocks.createMCPClient.mockResolvedValue({
+      tools: mocks.tools,
+      close: mocks.close,
+    });
+
+    const toolkit = new AISDKToolkit({
+      toolkit: {
+        docs: {
+          type: "mcp",
+          server: { type: "http", url: "http://localhost:3001/mcp" },
+        },
+      },
+    });
+
+    const tools = await toolkit.tools();
+    expect(Object.hasOwn(tools, "__proto__")).toBe(true);
+    expect(tools["__proto__"]).toBe(prototypeTool);
+  });
+
   it("loads MCP tools through pooled clients", async () => {
     mocks.tools.mockResolvedValue({ echo: { inputSchema: {} } });
     mocks.createMCPClient.mockResolvedValue({
@@ -454,6 +478,61 @@ describe("AISDKToolkit", () => {
     }
   });
 
+  it("does not evict a replacement client after an older listing timeout", async () => {
+    vi.useFakeTimers();
+    const oldClient = {
+      tools: vi.fn(() => never()),
+      close: vi.fn().mockResolvedValue(undefined),
+    };
+    const replacementClient = {
+      tools: vi.fn().mockResolvedValue({ echo: { inputSchema: {} } }),
+      close: vi.fn().mockResolvedValue(undefined),
+    };
+    mocks.createMCPClient
+      .mockResolvedValueOnce(oldClient)
+      .mockResolvedValue(replacementClient);
+
+    const toolkit = new AISDKToolkit({
+      toolkit: {
+        docs: {
+          type: "mcp",
+          server: {
+            type: "http",
+            url: "http://localhost:3001/mcp",
+            connectionTimeout: 100,
+          },
+        },
+      },
+    });
+
+    try {
+      const first = toolkit.tools();
+      const firstRejection = expect(first).rejects.toThrow(
+        /timed out while listing tools/,
+      );
+      await vi.advanceTimersByTimeAsync(50);
+
+      const second = toolkit.tools();
+      const secondRejection = expect(second).rejects.toThrow(
+        /timed out while listing tools/,
+      );
+      await vi.advanceTimersByTimeAsync(50);
+      await firstRejection;
+
+      await expect(toolkit.tools()).resolves.toHaveProperty("echo");
+      expect(mocks.createMCPClient).toHaveBeenCalledTimes(2);
+
+      await vi.advanceTimersByTimeAsync(50);
+      await secondRejection;
+
+      await expect(toolkit.tools()).resolves.toHaveProperty("echo");
+      expect(mocks.createMCPClient).toHaveBeenCalledTimes(2);
+      expect(oldClient.close).toHaveBeenCalledTimes(1);
+    } finally {
+      vi.useRealTimers();
+    }
+  });
+
   it("includes the MCP toolkit entry name when listing tools fails", async () => {
     const error = new Error("list failed");
     mocks.tools.mockRejectedValue(error);

package/src/tools/generativeTools.ts

@@ -227,11 +227,8 @@ export class AISDKToolkit {
         )
         .map(async ([name, tool]) => {
           const startedAt = Date.now();
-          const client = await this.#mcpClient(
-            name,
-            tool.server,
-            startedAt,
-          ).catch((error: unknown) => {
+          const clientPromise = this.#mcpClient(name, tool.server, startedAt);
+          const client = await clientPromise.catch((error: unknown) => {
             if (error instanceof MCPConnectionTimeoutError) throw error;
             throw toMcpToolkitError(name, "connect", error);
           });
@@ -245,8 +242,10 @@ export class AISDKToolkit {
             return [name, tool, tools] as const;
           } catch (error) {
             if (error instanceof MCPConnectionTimeoutError) {
-              this.#mcpClients.delete(name);
-              void client.close().catch(() => {});
+              if (this.#mcpClients.get(name) === clientPromise) {
+                this.#mcpClients.delete(name);
+                void client.close().catch(() => {});
+              }
               throw error;
             }
             throw toMcpToolkitError(name, "list tools", error);
@@ -254,7 +253,7 @@ export class AISDKToolkit {
         }),
     );
 
-    const tools: ToolSet = {};
+    const tools = Object.create(null) as ToolSet;
     const toolSources = new Map<string, string>();
     for (const [serverName, mcpTool, toolSet] of toolSets) {
       for (const [toolName, tool] of Object.entries(toolSet)) {

package/src/usage.test.ts

@@ -1,5 +1,5 @@
 import { describe, expect, it } from "vitest";
-import { getLatestThreadTokenUsage, getThreadMessageTokenUsage } from "./usage";
+import { getThreadMessageTokenUsage } from "./usage";
 
 function msg(metadata: unknown): { role: "assistant"; metadata: unknown } {
   return {
@@ -9,6 +9,22 @@ function msg(metadata: unknown): { role: "assistant"; metadata: unknown } {
 }
 
 describe("getThreadMessageTokenUsage", () => {
+  it("reads usage from custom.usage", () => {
+    const usage = getThreadMessageTokenUsage(
+      msg({
+        custom: {
+          usage: { inputTokens: 4, outputTokens: 6 },
+        },
+      }),
+    );
+
+    expect(usage).toEqual({
+      totalTokens: 10,
+      inputTokens: 4,
+      outputTokens: 6,
+    });
+  });
+
   it("does not double-count reasoning/cached in fallback totalTokens", () => {
     const usage = getThreadMessageTokenUsage(
       msg({
@@ -157,25 +173,27 @@ describe("getThreadMessageTokenUsage", () => {
   });
 });
 
-describe("getLatestThreadTokenUsage", () => {
-  it("falls back to the latest assistant message with usage", () => {
-    const usage = getLatestThreadTokenUsage([
+describe("getThreadMessageTokenUsage", () => {
+  it("returns token usage from an earlier assistant message", () => {
+    const messages = [
       { role: "assistant", metadata: { usage: { totalTokens: 100 } } },
       { role: "user", metadata: {} },
       { role: "assistant", metadata: {} },
-    ]);
+    ];
+    const usage = getThreadMessageTokenUsage(messages[0]);
 
     expect(usage).toEqual({ totalTokens: 100 });
   });
 
-  it("prefers the newest assistant message when it has usage", () => {
-    const usage = getLatestThreadTokenUsage([
+  it("returns token usage from the newest assistant message", () => {
+    const messages = [
       { role: "assistant", metadata: { usage: { totalTokens: 100 } } },
       {
         role: "assistant",
         metadata: { usage: { inputTokens: 40, outputTokens: 2 } },
       },
-    ]);
+    ];
+    const usage = getThreadMessageTokenUsage(messages[1]);
 
     expect(usage).toEqual({
       totalTokens: 42,