@assistant-ui/ai-sdk 0.0.2 → 0.0.4

package/src/runtime/useAISDKRuntime.ts

@@ -1,6 +1,13 @@
 "use client";
 
-import { useEffect, useMemo, useRef, useState } from "react";
+import {
+  useCallback,
+  useEffect,
+  useInsertionEffect,
+  useMemo,
+  useRef,
+  useState,
+} from "react";
 import type {
   UIMessage,
   useChat,
@@ -35,7 +42,10 @@ import {
   getExternalStoreMessages,
   pickExternalStoreSharedOptions,
 } from "@assistant-ui/core";
-import { consumeSuggestionResult } from "@assistant-ui/core/internal";
+import {
+  consumeSuggestionResult,
+  MessageRepository,
+} from "@assistant-ui/core/internal";
 import type { ReadonlyJSONObject } from "assistant-stream/utils";
 import { sliceMessagesUntil } from "../utils/sliceMessagesUntil";
 import { toCreateMessage } from "../converters/toCreateMessage";
@@ -70,48 +80,63 @@ const toUIMessage = <UI_MESSAGE extends UIMessage>(
     role: createMessage.role ?? fallbackRole,
   }) as UI_MESSAGE;
 
-export type AISDKRuntimeAdapter = ExternalStoreSharedOptions & {
-  adapters?:
-    | (NonNullable<ExternalStoreAdapter["adapters"]> & {
-        history?: ThreadHistoryAdapter | undefined;
-        suggestion?: SuggestionAdapter | undefined;
-      })
-    | undefined;
-  toCreateMessage?: CustomToCreateMessageFunction;
-  /**
-   * Whether to automatically cancel pending interactive tool calls when the user sends a new message.
-   *
-   * When enabled (default), the pending tool calls will be marked as failed with an error message
-   * indicating the user cancelled the tool call by sending a new message.
-   *
-   * @default true
-   */
-  cancelPendingToolCallsOnSend?: boolean | undefined;
-  /**
-   * Called when `runtime.thread.resumeRun(config)` is invoked.
-   *
-   * When omitted, `resumeRun` throws `"Runtime does not support resuming runs."`.
-   * Provide this to bridge resume invocations into a custom replay channel
-   * (for example, an SSE reconnect endpoint keyed by turn id).
-   */
-  onResume?: ExternalStoreAdapter["onResume"];
-  /**
-   * Called when `runtime.thread.resumeToolCall(options)` is invoked for a tool call the in-process tracker does not own.
-   *
-   * When omitted, `resumeToolCall` throws `"Tool call ${toolCallId} is not waiting for resume."`.
-   * Provide this to bridge resume-tool-call invocations into a custom handler.
-   */
-  onResumeToolCall?: ExternalStoreAdapter["onResumeToolCall"];
-  /**
-   * How consecutive assistant messages are rendered.
-   *
-   * `"concat-content"` (the default) merges them into a single thread message.
-   * `"none"` keeps each assistant message as its own thread message, which is
-   * useful when a backend persists proactive or consecutive assistant messages
-   * as separate entries.
-   */
-  joinStrategy?: JoinStrategy | undefined;
-};
+export type AISDKRuntimeAdapter<UI_MESSAGE extends UIMessage = UIMessage> =
+  ExternalStoreSharedOptions & {
+    adapters?:
+      | (NonNullable<ExternalStoreAdapter["adapters"]> & {
+          history?: ThreadHistoryAdapter | undefined;
+          suggestion?: SuggestionAdapter | undefined;
+        })
+      | undefined;
+    toCreateMessage?: CustomToCreateMessageFunction;
+    unstable_messageRepositoryInstance?: MessageRepository | undefined;
+    /**
+     * Whether to automatically cancel pending interactive tool calls when the user sends a new message.
+     *
+     * When enabled (default), the pending tool calls will be marked as failed with an error message
+     * indicating the user cancelled the tool call by sending a new message.
+     *
+     * @default true
+     */
+    cancelPendingToolCallsOnSend?: boolean | undefined;
+    /**
+     * Called when `runtime.thread.resumeRun(config)` is invoked.
+     *
+     * When omitted, `resumeRun` throws `"Runtime does not support resuming runs."`.
+     * Provide this to bridge resume invocations into a custom replay channel
+     * (for example, an SSE reconnect endpoint keyed by turn id).
+     */
+    onResume?: ExternalStoreAdapter["onResume"];
+    /**
+     * Called when `runtime.thread.resumeToolCall(options)` is invoked for a tool call the in-process tracker does not own.
+     *
+     * When omitted, `resumeToolCall` throws `"Tool call ${toolCallId} is not waiting for resume."`.
+     * Provide this to bridge resume-tool-call invocations into a custom handler.
+     */
+    onResumeToolCall?: ExternalStoreAdapter["onResumeToolCall"];
+    /**
+     * How consecutive assistant messages are rendered.
+     *
+     * `"concat-content"` (the default) merges them into a single thread message.
+     * `"none"` keeps each assistant message as its own thread message, which is
+     * useful when a backend persists proactive or consecutive assistant messages
+     * as separate entries.
+     */
+    joinStrategy?: JoinStrategy | undefined;
+    /**
+     * A branch-aware AI SDK message tree seeded once when `useChat` is empty.
+     * After that seed, live updates come only from `useChat`. A later empty
+     * chat or a new object identity does not reload the tree.
+     */
+    messageRepository?: MessageFormatRepository<UI_MESSAGE>;
+    /**
+     * Called after an explicit `switchToBranch` (for example a BranchPicker
+     * click). Complements `setMessages` and does not enable switching by itself.
+     *
+     * @deprecated This API is still under active development and might change without notice.
+     */
+    unstable_onBranchChange?: ExternalStoreAdapter["unstable_onBranchChange"];
+  };
 
 const EMPTY_SUGGESTIONS: readonly ThreadSuggestion[] = [];
 
@@ -125,9 +150,13 @@ const useGeneratedSuggestions = (
   const controllerRef = useRef<AbortController | null>(null);
   const wasRunningRef = useRef(false);
   const messagesRef = useRef(messages);
-  messagesRef.current = messages;
+  useInsertionEffect(() => {
+    messagesRef.current = messages;
+  }, [messages]);
   const adapterRef = useRef(suggestionAdapter);
-  adapterRef.current = suggestionAdapter;
+  useInsertionEffect(() => {
+    adapterRef.current = suggestionAdapter;
+  }, [suggestionAdapter]);
   const hasAdapter = suggestionAdapter != null;
 
   useEffect(() => {
@@ -188,9 +217,11 @@ const useGeneratedSuggestions = (
   return suggestions;
 };
 
+const NO_CANCELLED_MESSAGE_IDS: ReadonlySet<string> = new Set();
+
 export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
   chatHelpers: ReturnType<typeof useChat<UI_MESSAGE>>,
-  adapter: AISDKRuntimeAdapter = {},
+  adapter: AISDKRuntimeAdapter<UI_MESSAGE> = {},
 ) => {
   const {
     adapters,
@@ -199,12 +230,18 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     onResume,
     onResumeToolCall,
     joinStrategy,
+    messageRepository,
+    unstable_onBranchChange,
   } = adapter;
   const suggestionAdapter = adapters?.suggestion;
   const contextAdapters = useRuntimeAdapters();
   const [toolStatuses, setToolStatuses] = useState<
     Record<string, ToolExecutionStatus>
   >({});
+  const [cancelledMessages, setCancelledMessages] = useState<{
+    chatId: string;
+    ids: ReadonlySet<string>;
+  } | null>(null);
   const toolArgsKeyOrderCacheRef = useRef<Map<string, Map<string, string[]>>>(
     new Map(),
   );
@@ -217,10 +254,10 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
   const hasExecutingTools = Object.values(toolStatuses).some(
     (s) => s?.type === "executing",
   );
-  const isRunning =
-    chatHelpers.status === "submitted" ||
-    chatHelpers.status === "streaming" ||
-    hasExecutingTools;
+  const providerIsRunning =
+    chatHelpers.status === "submitted" || chatHelpers.status === "streaming";
+  const isRunning = providerIsRunning || hasExecutingTools;
+  const wasProviderRunningRef = useRef(providerIsRunning);
 
   const messageTiming = useStreamingTiming(chatHelpers.messages, isRunning);
 
@@ -230,6 +267,42 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
   const optimisticMessageId =
     isRunning && lastMessage?.role === "assistant" ? lastMessage.id : undefined;
 
+  const cancelledMessageIds =
+    cancelledMessages?.chatId === chatHelpers.id
+      ? cancelledMessages.ids
+      : NO_CANCELLED_MESSAGE_IDS;
+
+  const retractCancellation = useCallback(
+    (chatId: string, messageId: string) => {
+      setCancelledMessages((prev) => {
+        if (prev?.chatId !== chatId || !prev.ids.has(messageId)) return prev;
+        const ids = new Set(prev.ids);
+        ids.delete(messageId);
+        return { chatId, ids };
+      });
+    },
+    [],
+  );
+
+  // A provider run that resumes the stopped response retracts its cancellation;
+  // a run that starts a new response leaves the stopped one marked.
+  const resumedMessageId =
+    providerIsRunning && lastMessage?.role === "assistant"
+      ? lastMessage.id
+      : undefined;
+
+  useEffect(() => {
+    const wasProviderRunning = wasProviderRunningRef.current;
+    wasProviderRunningRef.current = providerIsRunning;
+    if (wasProviderRunning || !resumedMessageId) return;
+    retractCancellation(chatHelpers.id, resumedMessageId);
+  }, [
+    providerIsRunning,
+    resumedMessageId,
+    chatHelpers.id,
+    retractCancellation,
+  ]);
+
   const messages = AISDKMessageConverter.useThreadMessages({
     isRunning,
     messages: chatHelpers.messages,
@@ -243,11 +316,29 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
         mcpAppMetadataCache: mcpAppMetadataCacheRef.current,
         ...(optimisticMessageId && { optimisticMessageId }),
         ...(chatHelpers.error && { error: chatHelpers.error.message }),
+        ...(cancelledMessageIds.size > 0 && { cancelledMessageIds }),
       }),
-      [toolStatuses, messageTiming, optimisticMessageId, chatHelpers.error],
+      [
+        toolStatuses,
+        messageTiming,
+        optimisticMessageId,
+        chatHelpers.error,
+        cancelledMessageIds,
+      ],
     ),
   });
 
+  const exportedMessageRepository = useMemo(() => {
+    if (!messageRepository) return undefined;
+    const converted = toExportedMessageRepository(
+      AISDKMessageConverter.toThreadMessages as (
+        messages: UI_MESSAGE[],
+      ) => ThreadMessage[],
+      messageRepository,
+    );
+    return converted.messages.length > 0 ? converted : undefined;
+  }, [messageRepository]);
+
   const generatedSuggestions = useGeneratedSuggestions(
     suggestionAdapter,
     messages,
@@ -306,7 +397,11 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
       let hasChanges = false;
       const parts = lastMessage.parts?.map((part) => {
         if (!isToolUIPart(part)) return part;
-        if (part.state === "output-available" || part.state === "output-error")
+        if (
+          part.state === "output-available" ||
+          part.state === "output-error" ||
+          part.state === "output-denied"
+        )
           return part;
 
         hasChanges = true;
@@ -323,9 +418,17 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     });
   };
 
+  const hasSeededRepositoryRef = useRef(false);
+  const shouldFeedRepository =
+    exportedMessageRepository != null &&
+    !hasSeededRepositoryRef.current &&
+    messages.length === 0;
+
   const runtime = useExternalStoreRuntime({
-    isRunning,
-    messages,
+    isRunning: providerIsRunning,
+    ...(shouldFeedRepository
+      ? { messageRepository: exportedMessageRepository }
+      : { messages }),
     unstable_enableToolInvocations: true,
     setToolStatuses,
     setMessages: (messages) =>
@@ -391,10 +494,27 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
       runtimeRef.current.thread.import(exportedRepo);
     },
     onCancel: async () => {
+      const message = chatHelpers.messages.at(-1);
+      const cancelledId =
+        isRunning && message?.role === "assistant" ? message.id : undefined;
+      if (cancelledId) {
+        const liveIds = new Set(chatHelpers.messages.map((m) => m.id));
+        setCancelledMessages((prev) => {
+          const kept =
+            prev?.chatId === chatHelpers.id
+              ? [...prev.ids].filter((id) => liveIds.has(id))
+              : [];
+          return {
+            chatId: chatHelpers.id,
+            ids: new Set([...kept, cancelledId]),
+          };
+        });
+      }
       try {
         await chatHelpers.stop();
       } catch (error) {
         if (!(error instanceof Error && error.name === "AbortError")) {
+          if (cancelledId) retractCancellation(chatHelpers.id, cancelledId);
           throw error;
         }
       }
@@ -508,9 +628,14 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
         }),
       ),
     ...pickExternalStoreSharedOptions(adapter),
+    ...(adapter.unstable_messageRepositoryInstance && {
+      unstable_messageRepositoryInstance:
+        adapter.unstable_messageRepositoryInstance,
+    }),
     ...(suggestionAdapter ? { suggestions: generatedSuggestions } : {}),
     ...(onResume && { onResume }),
     ...(onResumeToolCall && { onResumeToolCall }),
+    ...(unstable_onBranchChange && { unstable_onBranchChange }),
     adapters: {
       attachments: vercelAttachmentAdapter,
       ...contextAdapters,
@@ -520,5 +645,24 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     isLoading,
   });
 
+  const setMessagesRef = useRef(chatHelpers.setMessages);
+  useInsertionEffect(() => {
+    setMessagesRef.current = chatHelpers.setMessages;
+  }, [chatHelpers.setMessages]);
+
+  useEffect(() => {
+    if (hasSeededRepositoryRef.current) return;
+    if (!exportedMessageRepository) return;
+    if (chatHelpers.messages.length > 0) {
+      hasSeededRepositoryRef.current = true;
+      return;
+    }
+    const tempRepo = new MessageRepository();
+    tempRepo.import(exportedMessageRepository);
+    setMessagesRef.current(
+      tempRepo.getMessages().flatMap(getExternalStoreMessages<UI_MESSAGE>),
+    );
+    hasSeededRepositoryRef.current = true;
+  }, [exportedMessageRepository, chatHelpers.messages.length]);
   return runtime;
 };

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

@@ -8,6 +8,7 @@ import { StrictMode, useState } from "react";
 import { describe, expect, it } from "vitest";
 import { AssistantChatTransport } from "../transport/AssistantChatTransport";
 import { useChatRuntime } from "./useChatRuntime";
+import { useThreadTokenUsage } from "../usage";
 
 const messages: UIMessage[] = [
   {
@@ -65,3 +66,48 @@ describe("useChatRuntime integration", () => {
     });
   });
 });
+
+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.test.ts

@@ -96,6 +96,22 @@ describe("useChatRuntime", () => {
     window.sessionStorage.clear();
   });
 
+  it("forwards a defined chat update throttle to useChat", () => {
+    mocks.useChat.mockReturnValue({
+      resumeStream: vi.fn(),
+      status: "ready",
+    });
+
+    renderHook(() => useChatRuntime({ throttle: 50 }));
+    renderHook(() => useChatRuntime({ throttle: undefined }));
+
+    expect(mocks.useChat).toHaveBeenNthCalledWith(
+      1,
+      expect.objectContaining({ throttle: 50 }),
+    );
+    expect(mocks.useChat.mock.calls[1]?.[0]).not.toHaveProperty("throttle");
+  });
+
   it("waits for external history to load before resuming a stream", async () => {
     mocks.state.isLoadingHistory = true;
     const resumeStream = vi.fn().mockResolvedValue(undefined);

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,
@@ -32,6 +33,7 @@ import { useResourceCleanup } from "./useResourceCleanup";
 export type ChatThreadOptions<UI_MESSAGE extends UIMessage = UIMessage> =
   ChatInit<UI_MESSAGE> &
     ExternalStoreSharedOptions & {
+      throttle?: number | undefined;
       adapters?: AISDKRuntimeAdapter["adapters"] | undefined;
       toCreateMessage?: CustomToCreateMessageFunction;
       onResume?: AISDKRuntimeAdapter["onResume"];
@@ -44,6 +46,8 @@ export type ChatThreadOptions<UI_MESSAGE extends UIMessage = UIMessage> =
        */
       onResumeError?: ((error: unknown) => void) | undefined;
       joinStrategy?: AISDKRuntimeAdapter["joinStrategy"];
+      messageRepository?: AISDKRuntimeAdapter<UI_MESSAGE>["messageRepository"];
+      unstable_onBranchChange?: AISDKRuntimeAdapter["unstable_onBranchChange"];
     };
 
 export type ChatThreadEnvironment<UI_MESSAGE extends UIMessage = UIMessage> = {
@@ -57,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>(
@@ -122,6 +132,7 @@ export const splitChatThreadOptions = <UI_MESSAGE extends UIMessage>(
   const {
     adapters,
     transport,
+    throttle,
     toCreateMessage,
     isDisabled: _isDisabled,
     isSendDisabled: _isSendDisabled,
@@ -131,6 +142,8 @@ export const splitChatThreadOptions = <UI_MESSAGE extends UIMessage>(
     onResumeToolCall,
     onResumeError,
     joinStrategy,
+    messageRepository,
+    unstable_onBranchChange,
     ...chatInit
   } = options ?? {};
   // peel guard: any shared key left in `chatInit` collapses this to `never`
@@ -141,11 +154,14 @@ export const splitChatThreadOptions = <UI_MESSAGE extends UIMessage>(
   return {
     adapters,
     transport,
+    throttle,
     toCreateMessage,
     onResume,
     onResumeToolCall,
     onResumeError,
     joinStrategy,
+    messageRepository,
+    unstable_onBranchChange,
     chatInit,
   };
 };
@@ -157,11 +173,14 @@ export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
   const {
     adapters,
     transport: transportOptions,
+    throttle,
     toCreateMessage,
     onResume,
     onResumeToolCall,
     onResumeError,
     joinStrategy,
+    messageRepository,
+    unstable_onBranchChange,
     chatInit: chatOptions,
   } = splitChatThreadOptions(options);
 
@@ -171,6 +190,7 @@ export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
     getThreadListItem,
     stopOnClientDestroy = false,
     chat: externalChat,
+    messageRepositoryInstance,
   } = env;
 
   const defaultTransport = useMemo(() => new AssistantChatTransport(), []);
@@ -181,6 +201,7 @@ export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
     ...chatOptions,
     id,
     transport,
+    ...(throttle !== undefined && { throttle }),
     ...(externalChat !== undefined && { chat: externalChat }),
   });
 
@@ -195,6 +216,11 @@ export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
     ...(onResume && { onResume }),
     ...(onResumeToolCall && { onResumeToolCall }),
     ...(joinStrategy && { joinStrategy }),
+    ...(messageRepository && { messageRepository }),
+    ...(messageRepositoryInstance && {
+      unstable_messageRepositoryInstance: messageRepositoryInstance,
+    }),
+    ...(unstable_onBranchChange && { unstable_onBranchChange }),
   });
 
   if (sourceTransport instanceof AssistantChatTransport) {