@assistant-ui/ai-sdk 0.0.1 → 0.0.3

package/src/runtime/useAISDKRuntime.ts

@@ -35,7 +35,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 +73,62 @@ 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;
+    /**
+     * 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[] = [];
 
@@ -190,7 +207,7 @@ const useGeneratedSuggestions = (
 
 export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
   chatHelpers: ReturnType<typeof useChat<UI_MESSAGE>>,
-  adapter: AISDKRuntimeAdapter = {},
+  adapter: AISDKRuntimeAdapter<UI_MESSAGE> = {},
 ) => {
   const {
     adapters,
@@ -199,6 +216,8 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     onResume,
     onResumeToolCall,
     joinStrategy,
+    messageRepository,
+    unstable_onBranchChange,
   } = adapter;
   const suggestionAdapter = adapters?.suggestion;
   const contextAdapters = useRuntimeAdapters();
@@ -217,10 +236,9 @@ 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 messageTiming = useStreamingTiming(chatHelpers.messages, isRunning);
 
@@ -248,6 +266,17 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     ),
   });
 
+  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 +335,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 +356,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) =>
@@ -511,6 +552,7 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     ...(suggestionAdapter ? { suggestions: generatedSuggestions } : {}),
     ...(onResume && { onResume }),
     ...(onResumeToolCall && { onResumeToolCall }),
+    ...(unstable_onBranchChange && { unstable_onBranchChange }),
     adapters: {
       attachments: vercelAttachmentAdapter,
       ...contextAdapters,
@@ -520,5 +562,22 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     isLoading,
   });
 
+  const setMessagesRef = useRef(chatHelpers.setMessages);
+  setMessagesRef.current = 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.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

@@ -32,6 +32,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 +45,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> = {
@@ -122,6 +125,7 @@ export const splitChatThreadOptions = <UI_MESSAGE extends UIMessage>(
   const {
     adapters,
     transport,
+    throttle,
     toCreateMessage,
     isDisabled: _isDisabled,
     isSendDisabled: _isSendDisabled,
@@ -131,6 +135,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 +147,14 @@ export const splitChatThreadOptions = <UI_MESSAGE extends UIMessage>(
   return {
     adapters,
     transport,
+    throttle,
     toCreateMessage,
     onResume,
     onResumeToolCall,
     onResumeError,
     joinStrategy,
+    messageRepository,
+    unstable_onBranchChange,
     chatInit,
   };
 };
@@ -157,11 +166,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);
 
@@ -181,6 +193,7 @@ export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
     ...chatOptions,
     id,
     transport,
+    ...(throttle !== undefined && { throttle }),
     ...(externalChat !== undefined && { chat: externalChat }),
   });
 
@@ -195,6 +208,8 @@ export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
     ...(onResume && { onResume }),
     ...(onResumeToolCall && { onResumeToolCall }),
     ...(joinStrategy && { joinStrategy }),
+    ...(messageRepository && { messageRepository }),
+    ...(unstable_onBranchChange && { unstable_onBranchChange }),
   });
 
   if (sourceTransport instanceof AssistantChatTransport) {

package/src/runtime/useExternalHistory.test.ts

@@ -392,6 +392,7 @@ describe("useExternalHistory persistence", () => {
     options?: {
       loadMessages?: MessageFormatRepository<InnerMessage>;
       toThreadMessages?: (messages: InnerMessage[]) => ThreadMessage[];
+      initialIsRunning?: boolean;
     },
   ) => {
     const append = vi.fn(
@@ -426,7 +427,7 @@ describe("useExternalHistory persistence", () => {
     };
 
     let listener: (() => void) | undefined;
-    let isRunning = false;
+    let isRunning = options?.initialIsRunning ?? false;
     let messages: ThreadMessage[] = [];
     const getState = vi.fn(() => ({ isRunning, messages }));
     const thread = {
@@ -497,6 +498,83 @@ describe("useExternalHistory persistence", () => {
     };
   };
 
+  it("persists a settled turn when the history adapter becomes active after it", async () => {
+    let listener: (() => void) | undefined;
+    let isRunning = false;
+    let messages: ThreadMessage[] = [];
+    const append = vi.fn(async () => {});
+    const formattedAdapter = {
+      load: vi.fn().mockResolvedValue({ messages: [] }),
+      append,
+    };
+    const historyAdapter: ThreadHistoryAdapter = {
+      load: vi.fn(),
+      append: vi.fn(),
+      withFormat: vi.fn().mockReturnValue(formattedAdapter),
+    };
+    let activeAdapter: ThreadHistoryAdapter | undefined;
+    const thread = {
+      subscribe: (next: () => void) => {
+        listener = next;
+        return () => {};
+      },
+      getState: () => ({ isRunning, messages }),
+      import: vi.fn(),
+      export: vi.fn(() => ({ headId: null, messages: [] })),
+    } as unknown as AssistantRuntime["thread"];
+    const persistenceRuntimeRef = {
+      current: { thread } as AssistantRuntime,
+    };
+    const message = createAssistantMessage(
+      { type: "complete", reason: "stop" },
+      [{ id: "inner-1", parts: ["answer"] }],
+    );
+
+    mocks.hasThreadListItem = true;
+    mocks.remoteId = "remote-thread";
+
+    const { rerender } = renderHook(() =>
+      useExternalHistory(
+        persistenceRuntimeRef,
+        activeAdapter,
+        () => [],
+        persistenceStorageFormat,
+        () => {},
+      ),
+    );
+
+    await act(async () => {
+      messages = [message];
+      isRunning = true;
+      listener?.();
+      isRunning = false;
+      listener?.();
+    });
+
+    activeAdapter = historyAdapter;
+    await act(async () => rerender());
+
+    await waitFor(() => expect(append).toHaveBeenCalledTimes(1));
+    expect(append).toHaveBeenCalledWith({
+      parentId: null,
+      message: { id: "inner-1", parts: ["answer"] },
+    });
+  });
+
+  it("persists a turn that is already running when the subscription starts", async () => {
+    const { append, step } = createPersistenceHarness(false, {
+      initialIsRunning: true,
+    });
+    const message = createAssistantMessage(
+      { type: "complete", reason: "stop" },
+      [{ id: "inner-1", parts: ["answer"] }],
+    );
+
+    await step({ isRunning: false, messages: [message] });
+
+    await waitFor(() => expect(append).toHaveBeenCalledTimes(1));
+  });
+
   it("retries a failed append on the next persistence pass", async () => {
     const consoleError = vi
       .spyOn(console, "error")

package/src/runtime/useExternalHistory.ts

@@ -234,6 +234,19 @@ export const useExternalHistory = <TMessage>(
       }, 0);
     });
 
+    const initialThreadState = runtimeRef.current.thread.getState();
+    wasRunningRef.current = initialThreadState.isRunning;
+    if (initialThreadState.isRunning) {
+      if (runStartRef.current == null) {
+        runStartRef.current = Date.now();
+        stepBoundariesRef.current = [];
+        toolCallCountRef.current = 0;
+        adapter.pin?.();
+      }
+    } else if (initialThreadState.messages.length > 0) {
+      persistSettled(true);
+    }
+
     function persistSettled(ignoreRunning: boolean) {
       persistTimerRef.current = null;
       const latest = runtimeRef.current.thread.getState();