@assistant-ui/ai-sdk 0.0.4 → 0.0.6

package/src/runtime/useChatThread.ts

@@ -188,7 +188,7 @@ export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
     id,
     isMainThread,
     getThreadListItem,
-    stopOnClientDestroy = false,
+    stopOnClientDestroy = true,
     chat: externalChat,
     messageRepositoryInstance,
   } = env;

package/src/runtime/useExternalHistory.test.ts

@@ -780,6 +780,35 @@ describe("useExternalHistory persistence", () => {
       ],
       expect.any(Object),
     );
+    expect(reportTelemetry.mock.calls[0]![1]).toMatchObject({
+      message: expect.objectContaining({
+        id: "assistant-a",
+        status: { type: "complete", reason: "stop" },
+      }),
+    });
+  });
+
+  it("reports a run that failed before any assistant message", async () => {
+    const { append, reportTelemetry, runCycle, flush } =
+      createPersistenceHarness(true);
+    const failed = createAssistantMessage(
+      {
+        type: "incomplete",
+        reason: "error",
+        error: { code: "AI_APICallError", message: "upstream failed" },
+      },
+      [],
+    );
+
+    await runCycle([failed]);
+    await flush();
+
+    expect(append).not.toHaveBeenCalled();
+    expect(reportTelemetry).toHaveBeenCalledTimes(1);
+    expect(reportTelemetry).toHaveBeenCalledWith(
+      [],
+      expect.objectContaining({ message: failed }),
+    );
   });
 
   it("restores deferred telemetry for reloaded paused messages", async () => {

package/src/runtime/useExternalHistory.ts

@@ -143,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) {
@@ -368,7 +371,10 @@ export const useExternalHistory = <TMessage>(
 
             if (deferredTelemetryIds.current.has(message.id) && isTerminal) {
               deferredTelemetryIds.current.delete(message.id);
-              adapter.reportTelemetry?.(batchItems, telemetryOptions);
+              adapter.reportTelemetry?.(batchItems, {
+                ...telemetryOptions,
+                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/transport/AssistantChatTransport.test.ts

@@ -1,5 +1,5 @@
 import type { UIMessage } from "ai";
-import { beforeEach, describe, expect, it, vi } from "vitest";
+import { describe, expect, it, vi } from "vitest";
 
 import { RESUMABLE_STREAM_ID_HEADER } from "./resumable";
 import { AssistantChatTransport } from "./AssistantChatTransport";
@@ -26,10 +26,6 @@ const sendMessagesOptions = {
 };
 
 describe("AssistantChatTransport.prepareSendMessagesRequest", () => {
-  beforeEach(() => {
-    vi.clearAllMocks();
-  });
-
   it("passes the initialized remote thread id to prepareSendMessagesRequest as options.id", async () => {
     const threadListItem = createThreadListItem("remote-thread-id");
     const captured: { id: unknown } = { id: undefined };
@@ -209,10 +205,6 @@ const wrappedFetchOf = (
   ).fetch;
 
 describe("AssistantChatTransport resumable fetch wrapper", () => {
-  beforeEach(() => {
-    vi.clearAllMocks();
-  });
-
   it("passes a 204 with a non-null empty body through untouched (WebKit)", async () => {
     const response = nullBodyStatusWithBody(204);
     const fetchMock = vi.fn(async () => response);

package/src/usage.ts

@@ -163,6 +163,14 @@ function findLatestMessageWithUsage(
   return undefined;
 }
 
+/**
+ * Reads token usage from the newest assistant message that reports any.
+ *
+ * A route attaches usage through the AI SDK's `messageMetadata` option. Because
+ * a thread message carries a fixed metadata shape, the converter moves every
+ * other key the route returns into `metadata.custom`, which is where this hook
+ * looks.
+ */
 export function useThreadTokenUsage(): ThreadTokenUsage | undefined {
   const msg = useAuiState((s) => findLatestMessageWithUsage(s.thread.messages));
   return useMemo(() => getThreadMessageTokenUsage(msg), [msg]);

package/dist/adapters/aiSDKFormatAdapter.d.ts.map

@@ -1 +0,0 @@
-{"version":3,"file":"aiSDKFormatAdapter.d.ts","names":[],"sources":["../../src/adapters/aiSDKFormatAdapter.ts"],"mappings":";;;KAQY,qBAAqB,KAAK;cAEzB,sBAAsB,qBACjC,WACA"}

package/dist/adapters/aiSDKFormatAdapter.js.map

@@ -1 +0,0 @@
-{"version":3,"file":"aiSDKFormatAdapter.js","names":["UIMessage","MessageFormatAdapter","MessageFormatItem","MessageStorageEntry","AISDKStorageFormat","Omit","aiSDKV6FormatAdapter","format","encode","message","id","parts","decode","stored","parentId","parent_id","content","getId"],"sources":["../../src/adapters/aiSDKFormatAdapter.ts"],"sourcesContent":["import type { UIMessage } from \"ai\";\nimport type {\n  MessageFormatAdapter,\n  MessageFormatItem,\n  MessageStorageEntry,\n} from \"@assistant-ui/core\";\n\n// Storage format for AI SDK messages - just the UIMessage\nexport type AISDKStorageFormat = Omit<UIMessage, \"id\">;\n\nexport const aiSDKV6FormatAdapter: MessageFormatAdapter<\n  UIMessage,\n  AISDKStorageFormat\n> = {\n  format: \"ai-sdk/v6\",\n\n  encode({\n    message: { id, parts, ...message },\n  }: MessageFormatItem<UIMessage>): AISDKStorageFormat {\n    return {\n      ...message,\n      parts,\n    };\n  },\n\n  decode(\n    stored: MessageStorageEntry<AISDKStorageFormat>,\n  ): MessageFormatItem<UIMessage> {\n    return {\n      parentId: stored.parent_id,\n      message: {\n        id: stored.id,\n        ...stored.content,\n      },\n    };\n  },\n\n  getId(message: UIMessage): string {\n    return message.id;\n  },\n};\n"],"mappings":";AAUA,MAAaM,uBAGT;CACFC,QAAQ;CAERC,OAAO,EACLC,SAAS,EAAEC,IAAIC,OAAO,GAAGF,aAC0B;EACnD,OAAO;GACL,GAAGA;GACHE;EACF;CACF;CAEAC,OACEC,QAC8B;EAC9B,OAAO;GACLC,UAAUD,OAAOE;GACjBN,SAAS;IACPC,IAAIG,OAAOH;IACX,GAAGG,OAAOG;GACZ;EACF;CACF;CAEAC,MAAMR,SAA4B;EAChC,OAAOA,QAAQC;CACjB;AACF"}