@assistant-ui/ai-sdk 0.0.5 → 0.0.7

package/src/runtime/useChatRuntime.ts

@@ -9,6 +9,7 @@ import {
 } from "@assistant-ui/core/react";
 import { useAui, useAuiState } from "@assistant-ui/store";
 import { useChatThread, type ChatThreadOptions } from "./useChatThread";
+import { AI_SDK_SDK } from "./sdkIdentity";
 
 export type UseChatRuntimeOptions<UI_MESSAGE extends UIMessage = UIMessage> =
   ChatThreadOptions<UI_MESSAGE> & {
@@ -38,7 +39,7 @@ export const useChatRuntime = <UI_MESSAGE extends UIMessage = UIMessage>({
   onThreadIdChange,
   ...options
 }: UseChatRuntimeOptions<UI_MESSAGE> = {}): AssistantRuntime => {
-  const cloudAdapter = useCloudThreadListAdapter({ cloud });
+  const cloudAdapter = useCloudThreadListAdapter({ cloud, sdk: AI_SDK_SDK });
   return useRemoteThreadListRuntime({
     runtimeHook: function RuntimeHook() {
       return useChatThreadRuntime(options);

package/src/runtime/useChatThread.test.ts

@@ -0,0 +1,74 @@
+// @vitest-environment jsdom
+
+import { describe, expect, it, vi } from "vitest";
+import { resource, useResource, flushTapSync } from "@assistant-ui/tap";
+import { useState } from "react";
+import {
+  RuntimeAdapter,
+  runtimeAdapterTransformScopes,
+} from "@assistant-ui/core/store";
+import {
+  attachTransformScopes,
+  AuiConfig,
+  createAssistantClient,
+} from "@assistant-ui/store/client";
+import { useChatThread, type ChatThreadEnvironment } from "./useChatThread";
+import {
+  createCancellableTransport,
+  nextTask,
+} from "./__tests__/controlled-transport";
+
+const createHost = (
+  env: Pick<ChatThreadEnvironment, "stopOnClientDestroy">,
+) => {
+  const useHost = (options: Parameters<typeof useChatThread>[0]) => {
+    const [threadListItem] = useState(() => ({
+      initialize: async () => ({ remoteId: "main", externalId: undefined }),
+    }));
+    const runtime = useChatThread(options, {
+      id: "main",
+      isMainThread: true,
+      getThreadListItem: () => threadListItem,
+      ...env,
+    });
+    return useResource(RuntimeAdapter(runtime));
+  };
+  attachTransformScopes(useHost, runtimeAdapterTransformScopes);
+  return resource(useHost);
+};
+
+const streamThenDestroy = async (
+  env: Pick<ChatThreadEnvironment, "stopOnClientDestroy">,
+) => {
+  const { transport, getCancelCount, close } = createCancellableTransport();
+  const Host = createHost(env);
+  const handle = createAssistantClient(
+    AuiConfig({ threads: Host({ transport }) }),
+  );
+  handle.subscribe(() => {});
+  const aui = handle.getClient();
+
+  try {
+    flushTapSync(() => aui.composer.setText("stop me"));
+    flushTapSync(() => aui.composer.send());
+    await vi.waitFor(() => {
+      expect(aui.thread.getState().isRunning).toBe(true);
+    });
+  } finally {
+    handle.destroy();
+  }
+  await nextTask();
+  const cancelCount = getCancelCount();
+  if (cancelCount === 0) close();
+  return cancelCount;
+};
+
+describe("useChatThread", () => {
+  it("stops an in-flight chat on client destroy when stopOnClientDestroy is omitted", async () => {
+    expect(await streamThenDestroy({})).toBe(1);
+  });
+
+  it("leaves an in-flight chat running on client destroy when stopOnClientDestroy is false", async () => {
+    expect(await streamThenDestroy({ stopOnClientDestroy: false })).toBe(0);
+  });
+});

package/src/runtime/useChatThread.ts

@@ -1,6 +1,6 @@
 "use client";
 
-import { useChat, type Chat, type UIMessage } from "@ai-sdk/react";
+import { Chat, useChat, type UIMessage } from "@ai-sdk/react";
 import type { MessageRepository } from "@assistant-ui/core/internal";
 import {
   pickExternalStoreSharedOptions,
@@ -26,6 +26,7 @@ import {
   useEffect,
   useMemo,
   useRef,
+  useState,
   useSyncExternalStore,
 } from "react";
 import { useResourceCleanup } from "./useResourceCleanup";
@@ -38,6 +39,7 @@ export type ChatThreadOptions<UI_MESSAGE extends UIMessage = UIMessage> =
       toCreateMessage?: CustomToCreateMessageFunction;
       onResume?: AISDKRuntimeAdapter["onResume"];
       onResumeToolCall?: AISDKRuntimeAdapter["onResumeToolCall"];
+      onRespondToToolApproval?: AISDKRuntimeAdapter["onRespondToToolApproval"];
       /**
        * Called when an automatic resumable stream reconnect fails. Use this to
        * surface a toast, report telemetry, or mark the thread as needing a
@@ -140,6 +142,7 @@ export const splitChatThreadOptions = <UI_MESSAGE extends UIMessage>(
     suggestions: _suggestions,
     onResume,
     onResumeToolCall,
+    onRespondToToolApproval,
     onResumeError,
     joinStrategy,
     messageRepository,
@@ -158,6 +161,7 @@ export const splitChatThreadOptions = <UI_MESSAGE extends UIMessage>(
     toCreateMessage,
     onResume,
     onResumeToolCall,
+    onRespondToToolApproval,
     onResumeError,
     joinStrategy,
     messageRepository,
@@ -166,6 +170,30 @@ export const splitChatThreadOptions = <UI_MESSAGE extends UIMessage>(
   };
 };
 
+type ChatCallbacks<UI_MESSAGE extends UIMessage> = Pick<
+  ChatInit<UI_MESSAGE>,
+  "onToolCall" | "onData" | "onFinish" | "onError" | "sendAutomaticallyWhen"
+>;
+
+/**
+ * Constructs a `Chat` whose callbacks read the latest options through
+ * `callbacksRef`, the forwarding `useChat` applies only to a chat it
+ * constructs itself.
+ */
+export const createChat = <UI_MESSAGE extends UIMessage>(
+  init: ChatInit<UI_MESSAGE>,
+  callbacksRef: { readonly current: ChatCallbacks<UI_MESSAGE> | undefined },
+): Chat<UI_MESSAGE> =>
+  new Chat<UI_MESSAGE>({
+    ...init,
+    onToolCall: (arg) => callbacksRef.current?.onToolCall?.(arg),
+    onData: (arg) => callbacksRef.current?.onData?.(arg),
+    onFinish: (arg) => callbacksRef.current?.onFinish?.(arg),
+    onError: (arg) => callbacksRef.current?.onError?.(arg),
+    sendAutomaticallyWhen: (arg) =>
+      callbacksRef.current?.sendAutomaticallyWhen?.(arg) ?? false,
+  });
+
 export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
   options: ChatThreadOptions<UI_MESSAGE> | undefined,
   env: ChatThreadEnvironment<UI_MESSAGE>,
@@ -177,6 +205,7 @@ export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
     toCreateMessage,
     onResume,
     onResumeToolCall,
+    onRespondToToolApproval,
     onResumeError,
     joinStrategy,
     messageRepository,
@@ -188,7 +217,7 @@ export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
     id,
     isMainThread,
     getThreadListItem,
-    stopOnClientDestroy = false,
+    stopOnClientDestroy = true,
     chat: externalChat,
     messageRepositoryInstance,
   } = env;
@@ -197,12 +226,21 @@ export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
   const sourceTransport = transportOptions ?? defaultTransport;
   const transport = useDynamicChatTransport(sourceTransport);
 
+  const latestChatOptionsRef = useRef(chatOptions);
+  useEffect(() => {
+    latestChatOptionsRef.current = chatOptions;
+  });
+  // `useChat` stops a chat it constructs whenever it unmounts, and a
+  // resource's soft unmount runs that cleanup, so the thread owns its chat.
+  const [ownedChat] = useState(
+    () =>
+      externalChat ??
+      createChat({ ...chatOptions, id, transport }, latestChatOptionsRef),
+  );
+
   const chat = useChat({
-    ...chatOptions,
-    id,
-    transport,
+    chat: externalChat ?? ownedChat,
     ...(throttle !== undefined && { throttle }),
-    ...(externalChat !== undefined && { chat: externalChat }),
   });
 
   useResourceCleanup(stopOnClientDestroy, () => {
@@ -215,6 +253,7 @@ export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
     ...(toCreateMessage && { toCreateMessage }),
     ...(onResume && { onResume }),
     ...(onResumeToolCall && { onResumeToolCall }),
+    ...(onRespondToToolApproval && { onRespondToToolApproval }),
     ...(joinStrategy && { joinStrategy }),
     ...(messageRepository && { messageRepository }),
     ...(messageRepositoryInstance && {

package/src/runtime/useExternalHistory.test.ts

@@ -498,6 +498,52 @@ describe("useExternalHistory persistence", () => {
     };
   };
 
+  it("persists a message that lands while the thread is idle", async () => {
+    const { append, step } = createPersistenceHarness(false);
+    const message = createAssistantMessage(
+      { type: "complete", reason: "stop" },
+      [{ id: "spoken-1", parts: ["spoken"] }],
+    );
+
+    await step({ messages: [message] });
+
+    await waitFor(() => expect(append).toHaveBeenCalledTimes(1));
+    expect(append).toHaveBeenCalledWith({
+      parentId: null,
+      message: { id: "spoken-1", parts: ["spoken"] },
+    });
+  });
+
+  it("schedules one idle persistence attempt per landed message when append fails", async () => {
+    const { append, step, flush } = createPersistenceHarness(false);
+    append.mockRejectedValueOnce(new Error("offline"));
+    const first = createAssistantMessage(
+      { type: "complete", reason: "stop" },
+      [{ id: "spoken-1", parts: ["first"] }],
+      "assistant-a",
+    );
+
+    await step({ messages: [first] });
+    await waitFor(() => expect(append).toHaveBeenCalledTimes(1));
+    await step({ messages: [first] });
+    await step({ messages: [first] });
+    await flush();
+    expect(append).toHaveBeenCalledTimes(1);
+
+    const second = createAssistantMessage(
+      { type: "complete", reason: "stop" },
+      [{ id: "spoken-2", parts: ["second"] }],
+      "assistant-b",
+    );
+    await step({ messages: [first, second] });
+    await waitFor(() => expect(append).toHaveBeenCalledTimes(3));
+    expect(append.mock.calls.map(([item]) => item)).toEqual([
+      { parentId: null, message: { id: "spoken-1", parts: ["first"] } },
+      { parentId: null, message: { id: "spoken-1", parts: ["first"] } },
+      { parentId: "spoken-1", message: { id: "spoken-2", parts: ["second"] } },
+    ]);
+  });
+
   it("persists a settled turn when the history adapter becomes active after it", async () => {
     let listener: (() => void) | undefined;
     let isRunning = false;
@@ -780,6 +826,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

@@ -50,6 +50,11 @@ const isAwaitingToolApproval = (message: ThreadMessage) =>
   message.status?.type === "requires-action" &&
   message.status.reason === "tool-calls";
 
+const isTerminalMessage = (message: ThreadMessage) =>
+  message.status === undefined ||
+  message.status.type === "complete" ||
+  message.status.type === "incomplete";
+
 const encodeContent = <TMessage>(
   storageFormatAdapter: MessageFormatAdapter<TMessage, any>,
   item: MessageFormatItem<TMessage>,
@@ -185,6 +190,14 @@ export const useExternalHistory = <TMessage>(
     if (!formatAdapter) return;
     const adapter = formatAdapter;
 
+    const idleScheduledInnerIds = new Set<string>();
+    const unpersistedInnerIds = (message: ThreadMessage) =>
+      isTerminalMessage(message)
+        ? getExternalStoreMessages<TMessage>(message)
+            .map((innerMessage) => storageFormatAdapter.getId(innerMessage))
+            .filter((innerId) => !persistedInnerMessages.current.has(innerId))
+        : [];
+
     const unsubscribe = runtimeRef.current.thread.subscribe(() => {
       const threadState = runtimeRef.current.thread.getState();
       const { isRunning } = threadState;
@@ -220,12 +233,16 @@ export const useExternalHistory = <TMessage>(
         return;
       }
 
-      // Only act on the true→false transition
-      if (!wasRunning) return;
-
-      // Record step boundary offset (synchronous for accuracy)
-      if (runStartRef.current != null) {
-        stepBoundariesRef.current.push(Date.now() - runStartRef.current);
+      if (wasRunning) {
+        // Record step boundary offset (synchronous for accuracy)
+        if (runStartRef.current != null) {
+          stepBoundariesRef.current.push(Date.now() - runStartRef.current);
+        }
+      } else {
+        const pending = threadState.messages.flatMap(unpersistedInnerIds);
+        if (pending.every((innerId) => idleScheduledInnerIds.has(innerId)))
+          return;
+        for (const innerId of pending) idleScheduledInnerIds.add(innerId);
       }
 
       // Debounce: wait one macrotask so agentic step flickers are absorbed
@@ -313,10 +330,7 @@ export const useExternalHistory = <TMessage>(
           for (const message of messages) {
             const innerMessages = getExternalStoreMessages<TMessage>(message);
 
-            const isTerminal =
-              message.status === undefined ||
-              message.status.type === "complete" ||
-              message.status.type === "incomplete";
+            const isTerminal = isTerminalMessage(message);
             const isAwaitingToolCalls = isAwaitingToolApproval(message);
             // A paused message's later content can only reach storage via update, so it is persisted early only when the adapter supports update.
             const isReady =
@@ -371,7 +385,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

@@ -623,8 +623,14 @@ describe("AISDKToolkit", () => {
       toolCallId: "call-docs-search",
       messages: [],
     };
+    const execute = toolSet.docs_search?.execute as
+      | ((
+          args: { query: string },
+          options: typeof executeOptions,
+        ) => Promise<string>)
+      | undefined;
     await expect(
-      toolSet.docs_search?.execute?.({ query: "assistant-ui" }, executeOptions),
+      execute?.({ query: "assistant-ui" }, executeOptions),
     ).resolves.toBe("docs result");
     expect(docsExecute).toHaveBeenCalledWith(
       { query: "assistant-ui" },

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/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"}