@assistant-ui/ai-sdk 0.0.4 → 0.0.5

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@assistant-ui/ai-sdk",
-  "version": "0.0.4",
+  "version": "0.0.5",
   "description": "Vercel AI SDK adapter for assistant-ui",
   "keywords": [
     "ai-sdk",
@@ -46,14 +46,14 @@
   ],
   "sideEffects": false,
   "dependencies": {
-    "@ai-sdk/mcp": "^2.0.41",
-    "@ai-sdk/react": "^4.0.88",
-    "@assistant-ui/core": "^0.3.17",
-    "@assistant-ui/store": "^0.3.12",
-    "@assistant-ui/tap": "^0.9.16",
-    "ai": "^7.0.85",
+    "@ai-sdk/mcp": "^2.0.45",
+    "@ai-sdk/react": "^4.0.96",
+    "@assistant-ui/core": "^0.3.18",
+    "@assistant-ui/store": "^0.3.13",
+    "@assistant-ui/tap": "^0.9.17",
+    "ai": "^7.0.93",
     "assistant-cloud": "*",
-    "assistant-stream": "^0.3.41"
+    "assistant-stream": "^0.3.42"
   },
   "peerDependencies": {
     "@types/react": "*",
@@ -68,11 +68,11 @@
     }
   },
   "devDependencies": {
-    "@assistant-ui/x-buildutils": "0.0.25",
+    "@assistant-ui/x-buildutils": "0.0.26",
     "@testing-library/react": "^16.3.3",
     "@types/json-schema": "^7.0.15",
     "@types/react": "^19.2.18",
-    "@types/react-dom": "^19.2.5",
+    "@types/react-dom": "^19.2.7",
     "jsdom": "^30.0.1",
     "react": "^19.2.8",
     "vitest": "^4.1.11"

package/src/converters/convertMessage.test.ts

@@ -394,6 +394,108 @@ describe("AISDKMessageConverter", () => {
     });
   });
 
+  it("preserves producer-defined approval fields and gives prompt precedence", () => {
+    const descriptor = { scope: "account:deploy" };
+    const converted = AISDKMessageConverter.toThreadMessages([
+      {
+        id: "a1",
+        role: "assistant",
+        parts: [
+          {
+            type: "tool-deploy",
+            toolCallId: "tc-1",
+            state: "approval-responded",
+            input: { environment: "production" },
+            approval: {
+              id: "approval-1",
+              approved: true,
+              reason: "approved by operator",
+              prompt: "Deploy to production?",
+              descriptor,
+              requestReason: "Production access requires approval",
+              signature: "signed-approval",
+              futureField: "preserved",
+            },
+          },
+        ],
+      } as any,
+    ]);
+
+    const toolCall = converted[0]?.content.find(
+      (part): part is any => part.type === "tool-call",
+    );
+    expect(toolCall?.approval).toEqual({
+      id: "approval-1",
+      approved: true,
+      reason: "approved by operator",
+      prompt: "Deploy to production?",
+      descriptor,
+      requestReason: "Production access requires approval",
+      signature: "signed-approval",
+      futureField: "preserved",
+    });
+  });
+
+  it("drops fields the AI SDK cannot answer and uses requestReason as the prompt", () => {
+    const converted = AISDKMessageConverter.toThreadMessages([
+      {
+        id: "a1",
+        role: "assistant",
+        parts: [
+          {
+            type: "tool-deploy",
+            toolCallId: "tc-1",
+            state: "approval-requested",
+            input: {},
+            approval: {
+              id: "approval-1",
+              display: "select",
+              allowFreeform: true,
+              options: [{ id: "once", kind: "allow-once" }],
+              optionId: "once",
+              text: "an answer",
+              resolution: "cancelled",
+              requestReason: "kept",
+            },
+          },
+        ],
+      } as any,
+    ]);
+
+    const toolCall = converted[0]?.content.find(
+      (part): part is any => part.type === "tool-call",
+    );
+    expect(toolCall?.approval).toEqual({
+      id: "approval-1",
+      prompt: "kept",
+      resolution: "cancelled",
+      requestReason: "kept",
+    });
+  });
+
+  it("drops a resolution the core contract does not declare", () => {
+    const converted = AISDKMessageConverter.toThreadMessages([
+      {
+        id: "a1",
+        role: "assistant",
+        parts: [
+          {
+            type: "tool-deploy",
+            toolCallId: "tc-1",
+            state: "approval-requested",
+            input: {},
+            approval: { id: "approval-1", resolution: "whatever" },
+          },
+        ],
+      } as any,
+    ]);
+
+    const toolCall = converted[0]?.content.find(
+      (part): part is any => part.type === "tool-call",
+    );
+    expect(toolCall?.approval).toEqual({ id: "approval-1" });
+  });
+
   it("strips closing delimiters from streaming tool argsText", () => {
     const converted = AISDKMessageConverter.toThreadMessages([
       {

package/src/converters/convertMessage.ts

@@ -154,30 +154,50 @@ function extractMcpAppMetadata(
 
 function getToolApprovalAndInterrupt(
   part: {
-    approval?:
-      | {
-          id: string;
-          approved?: boolean;
-          reason?: string;
-          isAutomatic?: boolean;
-        }
-      | undefined;
+    approval?: Record<string, unknown> | undefined;
   },
   toolStatus: { type: string; payload?: unknown } | undefined,
 ): {
   approval?: NonNullable<ToolCallMessagePart["approval"]>;
   interrupt?: NonNullable<ToolCallMessagePart["interrupt"]>;
 } {
-  if (part.approval && typeof part.approval.id === "string") {
-    const { id, approved, reason, isAutomatic } = part.approval;
-    return {
-      approval: {
-        id,
-        ...(typeof approved === "boolean" && { approved }),
-        ...(typeof reason === "string" && { reason }),
-        ...(isAutomatic === true && { isAutomatic: true }),
-      },
-    };
+  if (part.approval) {
+    // The AI SDK sends only id, approved and reason back to the server, so a
+    // request shape promising any other answer would render controls whose
+    // response cannot travel.
+    const {
+      id,
+      prompt,
+      approved,
+      reason,
+      isAutomatic,
+      resolution,
+      display,
+      allowFreeform,
+      options,
+      optionId,
+      text,
+      ...additionalApprovalFields
+    } = part.approval;
+    const requestReason = additionalApprovalFields.requestReason;
+    if (typeof id === "string")
+      return {
+        approval: {
+          ...additionalApprovalFields,
+          id,
+          ...(typeof prompt === "string"
+            ? { prompt }
+            : typeof requestReason === "string"
+              ? { prompt: requestReason }
+              : {}),
+          ...(typeof approved === "boolean" && { approved }),
+          ...(typeof reason === "string" && { reason }),
+          ...(isAutomatic === true && { isAutomatic: true }),
+          ...((resolution === "cancelled" || resolution === "expired") && {
+            resolution,
+          }),
+        } as NonNullable<ToolCallMessagePart["approval"]>,
+      };
   }
 
   if (toolStatus?.type === "interrupt") {

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

@@ -1,47 +1,68 @@
 // @vitest-environment jsdom
 
-import { useEffect } from "react";
-import { render, waitFor } from "@testing-library/react";
-import { StrictMode } from "react";
+import { StrictMode, type ReactNode } from "react";
+import { act, render, waitFor } from "@testing-library/react";
 import { AuiConfig, AuiProvider, useAui } from "@assistant-ui/store";
-import { flushTapSync } from "@assistant-ui/tap";
+import type { ChatTransport, UIMessage } from "ai";
 import { describe, expect, it } from "vitest";
 import { AISDKChat } from "./AISDKChat";
-import { createCancellableTransport } from "./__tests__/controlled-transport";
+import {
+  createCancellableTransport,
+  createStreamHarness,
+} from "./__tests__/controlled-transport";
 
 describe("AISDKChat React integration", () => {
-  it("does not treat React provider unmount as client destruction", async () => {
-    const { transport, getCancelCount, close } = createCancellableTransport();
-    let started = false;
-    let isRunning = () => false;
-
-    const SendOnMount = () => {
-      const aui = useAui();
-      isRunning = () => aui.thread.getState().isRunning;
-      useEffect(() => {
-        if (started) return;
-        started = true;
-        flushTapSync(() => aui.composer.setText("keep streaming"));
-        flushTapSync(() => aui.composer.send());
-      }, [aui]);
-      return null;
-    };
+  it("aborts the in-flight transport after a real unmount", async () => {
+    const { transport, getCancelCount } = createCancellableTransport();
+    const { Probe, send, isRunning } = createStreamHarness();
 
     const view = render(
       <StrictMode>
         <AuiProvider config={AuiConfig({ threads: AISDKChat({ transport }) })}>
-          <SendOnMount />
+          <Probe />
         </AuiProvider>
       </StrictMode>,
     );
 
-    await waitFor(() => {
-      expect(isRunning()).toBe(true);
-    });
+    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 new Promise((resolve) => setTimeout(resolve, 0));
+    await waitFor(() => expect(getCancelCount()).toBe(1));
+  });
+});
+
+describe("AISDKChat legacy useAui host integration", () => {
+  const LegacyProvider = ({
+    transport,
+    children,
+  }: {
+    transport: ChatTransport<UIMessage>;
+    children: ReactNode;
+  }) => {
+    const aui = useAui(AuiConfig({ threads: AISDKChat({ transport }) }));
+    return <AuiProvider value={aui}>{children}</AuiProvider>;
+  };
+
+  it("aborts the in-flight transport after a real unmount", async () => {
+    const { transport, getCancelCount } = createCancellableTransport();
+    const { Probe, send, isRunning } = createStreamHarness();
+
+    const view = render(
+      <StrictMode>
+        <LegacyProvider transport={transport}>
+          <Probe />
+        </LegacyProvider>
+      </StrictMode>,
+    );
+
+    await act(async () => send());
+    await waitFor(() => expect(isRunning()).toBe(true));
     expect(getCancelCount()).toBe(0);
-    close();
+
+    view.unmount();
+    await waitFor(() => expect(getCancelCount()).toBe(1));
   });
 });

package/src/runtime/__tests__/controlled-transport.ts

@@ -1,4 +1,6 @@
 import type { ChatTransport, UIMessage, UIMessageChunk } from "ai";
+import { useAui, type AssistantClient } from "@assistant-ui/store";
+import { flushTapSync } from "@assistant-ui/tap";
 
 export const createControlledTransport = () => {
   let controller!: ReadableStreamDefaultController<UIMessageChunk>;
@@ -41,3 +43,22 @@ export const createCancellableTransport = () => {
     close: () => controller.close(),
   };
 };
+
+export const nextTask = () => new Promise((resolve) => setTimeout(resolve, 0));
+
+export const createStreamHarness = () => {
+  let aui: AssistantClient | undefined;
+  const Probe = () => {
+    aui = useAui();
+    return null;
+  };
+  return {
+    Probe,
+    send: () => {
+      flushTapSync(() => aui!.composer.setText("keep streaming"));
+      flushTapSync(() => aui!.composer.send());
+    },
+    isRunning: () => aui?.thread.getState().isRunning === true,
+    client: () => aui!,
+  };
+};

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

@@ -1,12 +1,19 @@
 // @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";
 
@@ -65,8 +72,134 @@ 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 (

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/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) {

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.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]);