@assistant-ui/ai-sdk 0.0.2 → 0.0.4

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@assistant-ui/ai-sdk",
-  "version": "0.0.2",
+  "version": "0.0.4",
   "description": "Vercel AI SDK adapter for assistant-ui",
   "keywords": [
     "ai-sdk",
@@ -46,14 +46,14 @@
   ],
   "sideEffects": false,
   "dependencies": {
-    "@ai-sdk/mcp": "^2.0.32",
-    "@ai-sdk/react": "^4.0.69",
-    "@assistant-ui/core": "^0.3.15",
-    "@assistant-ui/store": "^0.3.10",
-    "@assistant-ui/tap": "^0.9.14",
-    "ai": "^7.0.66",
+    "@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",
     "assistant-cloud": "*",
-    "assistant-stream": "^0.3.39"
+    "assistant-stream": "^0.3.41"
   },
   "peerDependencies": {
     "@types/react": "*",
@@ -68,14 +68,14 @@
     }
   },
   "devDependencies": {
-    "@testing-library/react": "^16.3.2",
+    "@assistant-ui/x-buildutils": "0.0.25",
+    "@testing-library/react": "^16.3.3",
     "@types/json-schema": "^7.0.15",
     "@types/react": "^19.2.18",
-    "@types/react-dom": "^19.2.4",
+    "@types/react-dom": "^19.2.5",
     "jsdom": "^30.0.1",
     "react": "^19.2.8",
-    "vitest": "^4.1.10",
-    "@assistant-ui/x-buildutils": "0.0.23"
+    "vitest": "^4.1.11"
   },
   "publishConfig": {
     "access": "public",

package/src/converters/convertMessage.test.ts

@@ -26,6 +26,28 @@ describe("AISDKMessageConverter", () => {
     expect(converted[1]?.metadata.isOptimistic).toBe(true);
   });
 
+  it("keeps metadata outside the thread shape reachable under custom", () => {
+    const converted = AISDKMessageConverter.toThreadMessages([
+      {
+        id: "a1",
+        role: "assistant",
+        parts: [{ type: "text", text: "yo" }],
+        metadata: {
+          usage: { inputTokens: 40, outputTokens: 2 },
+          modelId: "gpt-5.6-luna",
+          custom: { source: "route" },
+        },
+      },
+    ] as any);
+
+    expect(converted[0]?.metadata.custom).toEqual({
+      usage: { inputTokens: 40, outputTokens: 2 },
+      modelId: "gpt-5.6-luna",
+      source: "route",
+    });
+    expect(converted[0]?.metadata).not.toHaveProperty("usage");
+  });
+
   it("does not flag messages when no optimistic id is provided", () => {
     const converted = AISDKMessageConverter.toThreadMessages([
       { id: "a1", role: "assistant", parts: [{ type: "text", text: "yo" }] },

package/src/converters/convertMessage.ts

@@ -31,6 +31,30 @@ import {
 import { unwrapModelContentEnvelope } from "./modelContentEnvelope";
 
 type MessageMetadata = ThreadMessageLike["metadata"];
+
+const THREAD_METADATA_KEYS = new Set([
+  "unstable_state",
+  "unstable_annotations",
+  "unstable_data",
+  "steps",
+  "timing",
+  "submittedFeedback",
+  "isOptimistic",
+  "custom",
+]);
+
+const toThreadMetadata = (metadata: unknown): MessageMetadata => {
+  if (!metadata || typeof metadata !== "object") return undefined;
+  const result: Record<string, unknown> = {};
+  const extra: Record<string, unknown> = {};
+  for (const [key, value] of Object.entries(metadata)) {
+    (THREAD_METADATA_KEYS.has(key) ? result : extra)[key] = value;
+  }
+  if (Object.keys(extra).length > 0) {
+    result.custom = { ...extra, ...(result.custom as object | undefined) };
+  }
+  return result as MessageMetadata;
+};
 export type AISDKMessageConverterMetadata =
   useExternalMessageConverter.Metadata & {
     toolArgsKeyOrderCache?: Map<string, Map<string, string[]>>;
@@ -427,7 +451,7 @@ export const AISDKMessageConverter = unstable_createMessageConverter(
                 status: { type: "complete" as const },
               };
             }),
-          metadata: message.metadata as MessageMetadata,
+          metadata: toThreadMetadata(message.metadata),
         };
 
       case "system":
@@ -442,7 +466,7 @@ export const AISDKMessageConverter = unstable_createMessageConverter(
           createdAt,
           content,
           metadata: {
-            ...(message.metadata as MessageMetadata),
+            ...toThreadMetadata(message.metadata),
             ...(timing && { timing }),
             ...(isOptimistic && { isOptimistic: true }),
           },

package/src/converters/toCreateMessage.ts

@@ -1,8 +1,7 @@
 import type { AppendMessage } from "@assistant-ui/core";
 import {
-  httpUrlPattern,
-  parseDataUrl,
   resolveFileMediaType,
+  resolveFilePartSource,
   resolveImageMediaType,
   toMediaWireUrl,
 } from "@assistant-ui/core/internal";
@@ -72,11 +71,13 @@ export const toCreateMessage = <UI_MESSAGE extends UIMessage = UIMessage>(
         // envelope is rebuilt from the typed format rather than forwarded.
         const mediaType = `audio/${part.audio.format}`;
         const data = part.audio.data;
+        const source = resolveFilePartSource({ data, mimeType: mediaType });
         return {
           type: "file",
-          url: httpUrlPattern.test(data)
-            ? data
-            : `data:${mediaType};base64,${parseDataUrl(data)?.data ?? data}`,
+          url:
+            source.kind === "url"
+              ? source.url
+              : `data:${mediaType};base64,${source.data}`,
           mediaType,
           ...(part.filename && { filename: part.filename }),
         };

package/src/model-context/injectInteractableContext.ts

@@ -23,7 +23,7 @@ import {
  * @example
  * ```ts
  * import { convertToModelMessages, streamText } from "ai";
- * import { unstable_injectInteractableContext } from "@assistant-ui/react-ai-sdk";
+ * import { unstable_injectInteractableContext } from "@assistant-ui/ai-sdk";
  *
  * export async function POST(req: Request) {
  *   const { messages } = await req.json();

package/src/model-context/injectQuoteContext.ts

@@ -22,7 +22,7 @@ const getQuoteText = (metadata: unknown): string | undefined => {
  * @example
  * ```ts
  * import { convertToModelMessages, streamText } from "ai";
- * import { injectQuoteContext } from "@assistant-ui/react-ai-sdk";
+ * import { injectQuoteContext } from "@assistant-ui/ai-sdk";
  *
  * export async function POST(req: Request) {
  *   const { messages } = await req.json();

package/src/runtime/AISDKThreads.cloud.test.ts

@@ -95,7 +95,7 @@ describe("AISDKThreads cloud", () => {
     handle.destroy();
   });
 
-  it("stops an in-flight cloud chat when switching away", async () => {
+  it("keeps an in-flight cloud chat running across a switch and stops it on delete", async () => {
     const chat = createCancellableTransport();
     const handle = createAssistantClient(
       AuiConfig({
@@ -126,6 +126,16 @@ describe("AISDKThreads cloud", () => {
       await vi.waitFor(() => {
         expect(handle.getClient().threads.getState().mainThreadId).toBe("t2");
       });
+      expect(chat.getCancelCount()).toBe(0);
+      await vi.waitFor(() => {
+        expect(
+          handle.getClient().threads.item({ id: "t1" }).getState().isRunning,
+        ).toBe(true);
+      });
+
+      flushTapSync(() =>
+        handle.getClient().threads.item({ id: "t1" }).delete(),
+      );
       await vi.waitFor(() => {
         expect(chat.getCancelCount()).toBe(1);
       });

package/src/runtime/AISDKThreads.test.ts

@@ -546,3 +546,105 @@ describe("AISDKThreads", () => {
     }
   });
 });
+
+describe("AISDKThreads branch isolation", () => {
+  const completeRound = async (
+    handle: ReturnType<typeof createAssistantClient>,
+    emit: (...chunks: never[]) => void,
+    close: () => void,
+    question: string,
+    answer: string,
+  ) => {
+    const aui = handle.getClient();
+    flushTapSync(() => aui.composer.setText(question));
+    flushTapSync(() => aui.composer.send());
+    await vi.waitFor(() => {
+      expect(aui.thread.getState().messages.length).toBeGreaterThan(0);
+    });
+    emit(...(textReply(answer) as never[]));
+    close();
+    await vi.waitFor(() => {
+      expect(handle.getClient().thread.getState().isRunning).toBe(false);
+    });
+  };
+
+  it("keeps switching between populated threads free of cross-thread branches", async () => {
+    const { transport, emit, close } = createControlledTransport();
+    const handle = createAssistantClient(
+      AuiConfig({ threads: AISDKThreads({ transport }) }),
+    );
+    handle.subscribe(() => {});
+    const aui = handle.getClient();
+
+    await completeRound(
+      handle,
+      emit as never,
+      close,
+      "thread a question",
+      "thread a answer",
+    );
+
+    flushTapSync(() => aui.threads.switchToNewThread());
+    await completeRound(
+      handle,
+      emit as never,
+      close,
+      "thread b question",
+      "thread b answer",
+    );
+
+    flushTapSync(() => aui.threads.switchToThread("main"));
+    await vi.waitFor(() => {
+      expect(threadText(handle as never)).toEqual([
+        "thread a question",
+        "thread a answer",
+      ]);
+    });
+    expect(aui.thread.message({ index: 0 }).getState().branchCount).toBe(1);
+    expect(aui.thread.message({ index: 1 }).getState().branchCount).toBe(1);
+
+    handle.destroy();
+  });
+
+  it("preserves intra-thread branches across a switch through an empty thread", async () => {
+    const { transport, emit, close } = createControlledTransport();
+    const handle = createAssistantClient(
+      AuiConfig({ threads: AISDKThreads({ transport }) }),
+    );
+    handle.subscribe(() => {});
+    const aui = handle.getClient();
+
+    await completeRound(
+      handle,
+      emit as never,
+      close,
+      "branchy question",
+      "first answer",
+    );
+
+    flushTapSync(() => aui.thread.message({ index: 1 }).reload());
+    await vi.waitFor(() => {
+      expect(handle.getClient().thread.getState().isRunning).toBe(true);
+    });
+    emit(...(textReply("second answer") as never[]));
+    close();
+    await vi.waitFor(() => {
+      expect(handle.getClient().thread.getState().isRunning).toBe(false);
+      expect(aui.thread.message({ index: 1 }).getState().branchCount).toBe(2);
+    });
+
+    flushTapSync(() => aui.threads.switchToNewThread());
+    expect(handle.getClient().thread.getState().messages).toHaveLength(0);
+
+    flushTapSync(() => handle.getClient().threads.switchToThread("main"));
+    await vi.waitFor(() => {
+      expect(threadText(handle as never)).toEqual([
+        "branchy question",
+        "second answer",
+      ]);
+    });
+    expect(aui.thread.message({ index: 1 }).getState().branchCount).toBe(2);
+
+    handle.destroy();
+  });
+});

package/src/runtime/AISDKThreads.ts

@@ -24,6 +24,7 @@ import {
   useChatThread,
   type ChatThreadOptions,
 } from "./useChatThread";
+import { MessageRepository } from "@assistant-ui/core/internal";
 import { useResourceCleanup } from "./useResourceCleanup";
 
 export type AISDKThreadsOptions<UI_MESSAGE extends UIMessage = UIMessage> =
@@ -41,9 +42,9 @@ export type AISDKThreadsOptions<UI_MESSAGE extends UIMessage = UIMessage> =
       | undefined;
     /**
      * When set, the thread list is a `RemoteThreadList` backed by this
-     * assistant-cloud. Omit it to keep the in-memory list. The thread
-     * factory is keyed so cloud history reloads on a switch, and an
-     * in-flight run does not continue in the background.
+     * assistant-cloud. Omit it to keep the in-memory list. Every visited
+     * cloud thread stays mounted, so an in-flight run continues after a
+     * switch and stops on delete; per-thread history loads once per thread.
      */
     cloud?: AssistantCloud | undefined;
     /**
@@ -64,6 +65,7 @@ type AISDKThreadChatOptions<UI_MESSAGE extends UIMessage = UIMessage> = Omit<
 type ChatEntry<UI_MESSAGE extends UIMessage> = {
   chat: Chat<UI_MESSAGE>;
   transport: ChatTransport<UI_MESSAGE>;
+  repository: MessageRepository;
 };
 
 const createChatEntry = <UI_MESSAGE extends UIMessage>(
@@ -84,6 +86,7 @@ const createChatEntry = <UI_MESSAGE extends UIMessage>(
   return {
     chat: new Chat<UI_MESSAGE>({ ...chatInit, id: threadId, transport }),
     transport,
+    repository: new MessageRepository(),
   };
 };
 
@@ -113,7 +116,7 @@ const useAISDKChatThread = <UI_MESSAGE extends UIMessage = UIMessage>({
   const [owned] = useState(() =>
     cloud ? createChatEntry(threadId, options) : undefined,
   );
-  const { chat, transport } =
+  const { chat, transport, repository } =
     owned ?? getOrCreateChatEntry(threadId, options, chats);
 
   useEffect(() => {
@@ -145,6 +148,7 @@ const useAISDKChatThread = <UI_MESSAGE extends UIMessage = UIMessage>({
             : undefined
           : fallbackItem,
       chat,
+      messageRepositoryInstance: repository,
       stopOnClientDestroy: cloud,
     },
   );
@@ -190,6 +194,7 @@ const useAISDKThreads = <UI_MESSAGE extends UIMessage = UIMessage>(
     bindCloud
       ? RemoteThreadList({
           adapter: cloudAdapter,
+          backgroundThreads: true,
           thread,
           threadId,
           onThreadIdChange,
@@ -212,11 +217,13 @@ const useAISDKThreads = <UI_MESSAGE extends UIMessage = UIMessage>(
  * per-thread orchestration as {@link AISDKChat} inside the client's own
  * resource tree, so it works with any `AssistantClient` host, React or not.
  * Without `cloud`, threads live in memory for the client's lifetime and keep
- * their history across switches; each thread's chat id is its thread id.
- * With `cloud`, the list is a `RemoteThreadList` and the factory is keyed so
- * cloud history reloads on a switch. The store entry mounts only the visible
- * thread, so a switch cancels an in-flight run. Model context is
- * registered on the visible thread only.
+ * their history across switches; only the visible thread is mounted, and a
+ * switched-away chat keeps streaming into its stored state until it settles
+ * or the thread is deleted. With `cloud`, the list is a
+ * `RemoteThreadList` with `backgroundThreads`: every visited thread stays
+ * mounted with its own history, a run continues after a switch and stops on
+ * delete, and a freshly created thread titles itself. Model context is
+ * registered on every mounted thread.
  */
 export const AISDKThreads = resource(useAISDKThreads);
 

package/src/runtime/useAISDKRuntime.denied-tool.test.tsx

@@ -0,0 +1,122 @@
+// @vitest-environment jsdom
+
+import { renderHook } from "@testing-library/react";
+import type { ExternalStoreAdapter } from "@assistant-ui/core";
+import { describe, expect, it, vi } from "vitest";
+
+const mocks = vi.hoisted(() => ({
+  adapter: undefined as ExternalStoreAdapter | undefined,
+}));
+
+vi.mock("@assistant-ui/core/react", async (importOriginal) => {
+  const original =
+    await importOriginal<typeof import("@assistant-ui/core/react")>();
+  return {
+    ...original,
+    useExternalStoreRuntime: vi.fn((adapter: ExternalStoreAdapter) => {
+      mocks.adapter = adapter;
+      return {};
+    }),
+    useRuntimeAdapters: vi.fn(() => ({})),
+  };
+});
+
+vi.mock("./useExternalHistory", async (importOriginal) => {
+  const original =
+    await importOriginal<typeof import("./useExternalHistory")>();
+  return {
+    ...original,
+    useExternalHistory: vi.fn(() => ({
+      isLoading: false,
+      deleteMessage: vi.fn().mockResolvedValue(undefined),
+    })),
+  };
+});
+
+import { useAISDKRuntime } from "./useAISDKRuntime";
+
+describe("useAISDKRuntime pending tool completion", () => {
+  const setup = (toolPart: Record<string, unknown>) => {
+    const messages = [
+      {
+        id: "m1",
+        role: "assistant",
+        parts: [toolPart],
+      },
+    ];
+    let updated: any[] | undefined;
+    const setMessages = vi.fn(
+      (updater: (m: typeof messages) => typeof messages) => {
+        updated = typeof updater === "function" ? updater(messages) : updater;
+      },
+    );
+    const chat = {
+      id: "chat-1",
+      status: "ready",
+      error: undefined,
+      messages,
+      setMessages,
+      sendMessage: vi.fn().mockResolvedValue(undefined),
+      regenerate: vi.fn(),
+      addToolOutput: vi.fn(),
+      addToolApprovalResponse: vi.fn(),
+      stop: vi.fn(),
+    };
+
+    renderHook(() => useAISDKRuntime(chat as never));
+
+    return {
+      send: () =>
+        mocks.adapter!.onNew!({
+          role: "user",
+          parentId: "m1",
+          sourceId: null,
+          runConfig: {},
+          content: [{ type: "text", text: "next" }],
+          attachments: [],
+          createdAt: new Date(0),
+          metadata: { custom: {} },
+        } as never),
+      getUpdated: () => updated,
+    };
+  };
+
+  it("preserves a denied tool call and its approval record on the next send", async () => {
+    const { send, getUpdated } = setup({
+      type: "tool-search",
+      toolCallId: "call_1",
+      state: "output-denied",
+      input: { query: "x" },
+      approval: { id: "appr_1", approved: false, reason: "too dangerous" },
+    });
+
+    await send();
+
+    const updated = getUpdated();
+    const part = updated?.at(-1)?.parts?.[0];
+    expect(part?.state).toBe("output-denied");
+    expect(part?.approval).toEqual({
+      id: "appr_1",
+      approved: false,
+      reason: "too dangerous",
+    });
+    expect(part?.errorText).toBeUndefined();
+  });
+
+  it("still cancels a genuinely pending tool call on the next send", async () => {
+    const { send, getUpdated } = setup({
+      type: "tool-search",
+      toolCallId: "call_1",
+      state: "input-available",
+      input: { query: "x" },
+    });
+
+    await send();
+
+    const part = getUpdated()?.at(-1)?.parts?.[0];
+    expect(part?.state).toBe("output-error");
+    expect(part?.errorText).toBe(
+      "User cancelled tool call by sending a new message.",
+    );
+  });
+});