@assistant-ui/ai-sdk 0.0.10 → 0.0.12

package/src/runtime/useAISDKRuntime.ts

@@ -7,6 +7,7 @@ import {
   useMemo,
   useRef,
   useState,
+  useSyncExternalStore,
 } from "react";
 import type {
   UIMessage,
@@ -60,7 +61,7 @@ import {
   AISDKMessageConverter,
   type AISDKMessageConverterMetadata,
 } from "../converters/convertMessage";
-import { wrapModelContentEnvelope } from "../converters/modelContentEnvelope";
+import { wrapModelContentEnvelope } from "assistant-stream/internal";
 import {
   type AISDKStorageFormat,
   aiSDKV6FormatAdapter,
@@ -114,12 +115,15 @@ export type AISDKRuntimeAdapter<UI_MESSAGE extends UIMessage = UIMessage> =
         })
       | undefined;
     toCreateMessage?: CustomToCreateMessageFunction;
+    /** @deprecated Experimental since 2026-08-28. Not scheduled for removal; the API may change in any release. */
     unstable_messageRepositoryInstance?: MessageRepository | undefined;
     /**
      * The object a host answer belongs to, normally the `Chat` the runtime
      * renders. A host answer never reaches the `useChat` messages, so without
      * an owner it lives only as long as this runtime. With an owner, it
      * survives runtime remounts over the same chat.
+     *
+     * @deprecated Experimental since 2026-09-29. Not scheduled for removal; the API may change in any release.
      */
     unstable_hostApprovalOwner?: object | undefined;
     /**
@@ -183,7 +187,7 @@ export type AISDKRuntimeAdapter<UI_MESSAGE extends UIMessage = UIMessage> =
      * 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.
+     * @deprecated Experimental since 2026-06-23. Not scheduled for removal; the API may change in any release.
      */
     unstable_onBranchChange?: ExternalStoreAdapter["unstable_onBranchChange"];
   };
@@ -203,11 +207,25 @@ const useGeneratedSuggestions = (
   useInsertionEffect(() => {
     messagesRef.current = messages;
   }, [messages]);
+  const adapterKey = suggestionAdapter?.key;
   const adapterRef = useRef(suggestionAdapter);
+  const lastAdapterKeyRef = useRef(adapterKey);
+  const adapterReplacedRef = useRef(false);
   useInsertionEffect(() => {
+    const previousAdapter = adapterRef.current;
     adapterRef.current = suggestionAdapter;
-  }, [suggestionAdapter]);
-  const hasAdapter = suggestionAdapter != null;
+    const adapterReplaced =
+      previousAdapter !== undefined &&
+      suggestionAdapter !== undefined &&
+      lastAdapterKeyRef.current !== adapterKey;
+    adapterReplacedRef.current ||= adapterReplaced;
+    if (suggestionAdapter === undefined || adapterReplaced) {
+      controllerRef.current?.abort();
+    }
+    lastAdapterKeyRef.current = adapterKey;
+  }, [adapterKey, suggestionAdapter]);
+  const hasAdapter = suggestionAdapter !== undefined;
+  const adapterPresentRef = useRef(hasAdapter);
 
   useEffect(() => {
     const clearSuggestions = () => {
@@ -217,12 +235,18 @@ const useGeneratedSuggestions = (
     };
 
     const adapter = adapterRef.current;
+    const adapterRestored = hasAdapter && !adapterPresentRef.current;
+    adapterPresentRef.current = hasAdapter;
+    const adapterReplaced = adapterReplacedRef.current;
+    adapterReplacedRef.current = false;
     if (!adapter) {
       clearSuggestions();
       wasRunningRef.current = isRunning;
       return;
     }
 
+    if (adapterReplaced) clearSuggestions();
+
     if (isRunning) {
       if (!wasRunningRef.current) {
         clearSuggestions();
@@ -231,7 +255,7 @@ const useGeneratedSuggestions = (
       return;
     }
 
-    if (!wasRunningRef.current) return;
+    if (!wasRunningRef.current && !adapterReplaced && !adapterRestored) return;
     wasRunningRef.current = false;
 
     const currentMessages = messagesRef.current;
@@ -252,11 +276,14 @@ const useGeneratedSuggestions = (
 
         await consumeSuggestionResult(promiseOrGenerator, {
           signal,
-          onUpdate: setSuggestions,
+          onUpdate: (nextSuggestions) => {
+            if (controllerRef.current !== controller) return;
+            setSuggestions(nextSuggestions);
+          },
         });
       } catch {}
     })();
-  }, [hasAdapter, isRunning]);
+  }, [adapterKey, hasAdapter, isRunning]);
 
   useReplaySafeEffect(() => {
     return () => {
@@ -290,10 +317,25 @@ type OwnedApproval = {
   response: RespondToToolApprovalOptions;
 };
 
-const hostToolApprovalsByChat = new WeakMap<
-  object,
-  Map<string, OwnedApproval>
->();
+type ChatRuntimeState = {
+  approvals: Map<string, OwnedApproval>;
+  cancelledMessageIds: ReadonlySet<string>;
+};
+
+const chatRuntimeStates = new WeakMap<object, ChatRuntimeState>();
+
+const getChatRuntimeState = (owner: object | undefined) => {
+  if (owner === undefined) return undefined;
+  let state = chatRuntimeStates.get(owner);
+  if (state === undefined) {
+    state = {
+      approvals: new Map(),
+      cancelledMessageIds: NO_CANCELLED_MESSAGE_IDS,
+    };
+    chatRuntimeStates.set(owner, state);
+  }
+  return state;
+};
 
 const toApprovalResponses = (
   owned: ReadonlyMap<string, OwnedApproval> | undefined,
@@ -308,19 +350,19 @@ const toApprovalResponses = (
  * may already be unmounted (a rollback resolving after a remount), and another
  * runtime may be mounted over the same owner.
  */
-const hostApprovalListenersByChat = new WeakMap<object, Set<() => void>>();
+const chatRuntimeListeners = new WeakMap<object, Set<() => void>>();
 
-const subscribeToHostApprovals = (owner: object, listener: () => void) => {
-  const listeners = hostApprovalListenersByChat.get(owner) ?? new Set();
-  hostApprovalListenersByChat.set(owner, listeners);
+const subscribeToChatRuntimeState = (owner: object, listener: () => void) => {
+  const listeners = chatRuntimeListeners.get(owner) ?? new Set();
+  chatRuntimeListeners.set(owner, listeners);
   listeners.add(listener);
   return () => {
     listeners.delete(listener);
   };
 };
 
-const notifyHostApprovals = (owner: object) => {
-  for (const listener of [...(hostApprovalListenersByChat.get(owner) ?? [])]) {
+const notifyChatRuntimeState = (owner: object) => {
+  for (const listener of [...(chatRuntimeListeners.get(owner) ?? [])]) {
     listener();
   }
 };
@@ -343,6 +385,7 @@ const getSupersededApprovalProjection = <UI_MESSAGE extends UIMessage>(
       lastAssistant === messages[lastIndex] &&
       lastAssistant?.parts?.some((part) => {
         if (
+          typeof part?.type !== "string" ||
           !isToolUIPart(part) ||
           part.state === "output-available" ||
           part.state === "output-error" ||
@@ -391,7 +434,12 @@ const getSupersededApprovalProjection = <UI_MESSAGE extends UIMessage>(
     }
 
     for (const part of message.parts ?? []) {
-      if (!isToolUIPart(part) || part.state !== "approval-requested") continue;
+      if (
+        typeof part?.type !== "string" ||
+        !isToolUIPart(part) ||
+        part.state !== "approval-requested"
+      )
+        continue;
 
       const approval = part.approval;
       if (!approval) continue;
@@ -427,7 +475,10 @@ const findRawToolMessageIndex = <UI_MESSAGE extends UIMessage>(
 ) => {
   const containsToolCall = (message: UI_MESSAGE) =>
     message.parts?.some(
-      (part) => isToolUIPart(part) && part.toolCallId === toolCallId,
+      (part) =>
+        typeof part?.type === "string" &&
+        isToolUIPart(part) &&
+        part.toolCallId === toolCallId,
     ) === true;
 
   const messageIndex = messages.findIndex(
@@ -491,14 +542,22 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
   // Set by hosts that own the chat across runtime lifetimes, so a host answer
   // outlives a remount over that same chat.
   const approvalOwner = adapter.unstable_hostApprovalOwner;
-  const ownedApprovals = approvalOwner
-    ? (hostToolApprovalsByChat.get(approvalOwner) ??
-      (() => {
-        const created = new Map<string, OwnedApproval>();
-        hostToolApprovalsByChat.set(approvalOwner, created);
-        return created;
-      })())
-    : undefined;
+  const chatRuntimeState = getChatRuntimeState(approvalOwner);
+  const ownedApprovals = chatRuntimeState?.approvals;
+  const ownedCancelledMessageIds = useSyncExternalStore(
+    useCallback(
+      (listener) =>
+        approvalOwner
+          ? subscribeToChatRuntimeState(approvalOwner, listener)
+          : () => {},
+      [approvalOwner],
+    ),
+    useCallback(
+      () => chatRuntimeState?.cancelledMessageIds ?? NO_CANCELLED_MESSAGE_IDS,
+      [chatRuntimeState],
+    ),
+    () => NO_CANCELLED_MESSAGE_IDS,
+  );
   const [toolApprovalResponses, setToolApprovalResponses] = useState<
     ReadonlyMap<string, RespondToToolApprovalOptions>
   >(() => toApprovalResponses(ownedApprovals));
@@ -532,7 +591,7 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
         return unchanged ? prev : next;
       });
     };
-    const unsubscribe = subscribeToHostApprovals(approvalOwner, sync);
+    const unsubscribe = subscribeToChatRuntimeState(approvalOwner, sync);
     // The state was seeded during render, so a write landing between then and
     // this subscription would otherwise never be seen.
     sync();
@@ -559,7 +618,7 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     new Map(),
   );
   const toolArgsTextCacheRef = useRef<
-    WeakMap<ReadonlyJSONObject, Map<string, string>>
+    NonNullable<AISDKMessageConverterMetadata["toolArgsTextCache"]>
   >(new WeakMap());
   const mcpAppMetadataCacheRef = useRef<Map<string, McpAppMetadata>>(new Map());
   const toolArtifactsRef = useRef<Map<string, unknown>>(new Map());
@@ -590,8 +649,9 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
   const optimisticMessageId =
     isRunning && lastMessage?.role === "assistant" ? lastMessage.id : undefined;
 
-  const cancelledMessageIds =
-    cancelledMessages?.chatId === chatHelpers.id
+  const cancelledMessageIds = chatRuntimeState
+    ? ownedCancelledMessageIds
+    : cancelledMessages?.chatId === chatHelpers.id
       ? cancelledMessages.ids
       : NO_CANCELLED_MESSAGE_IDS;
   const supportsRichToolApprovalResponses =
@@ -646,6 +706,14 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
 
   const retractCancellation = useCallback(
     (chatId: string, messageId: string) => {
+      if (chatRuntimeState !== undefined && approvalOwner !== undefined) {
+        if (!chatRuntimeState.cancelledMessageIds.has(messageId)) return;
+        const ids = new Set(chatRuntimeState.cancelledMessageIds);
+        ids.delete(messageId);
+        chatRuntimeState.cancelledMessageIds = ids;
+        notifyChatRuntimeState(approvalOwner);
+        return;
+      }
       setCancelledMessages((prev) => {
         if (prev?.chatId !== chatId || !prev.ids.has(messageId)) return prev;
         const ids = new Set(prev.ids);
@@ -653,7 +721,7 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
         return { chatId, ids };
       });
     },
-    [],
+    [approvalOwner, chatRuntimeState],
   );
 
   // A provider run that resumes the stopped response retracts its cancellation;
@@ -801,7 +869,8 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
 
         let messageChanged = false;
         const parts = message.parts?.map((part) => {
-          if (!isToolUIPart(part)) return part;
+          if (typeof part?.type !== "string" || !isToolUIPart(part))
+            return part;
           if (
             part.state === "output-available" ||
             part.state === "output-error" ||
@@ -850,7 +919,9 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     const requested = chatHelpers.messages
       .flatMap((message) =>
         message.parts.flatMap((part) =>
-          isToolUIPart(part) ? [{ messageId: message.id, part }] : [],
+          typeof part?.type === "string" && isToolUIPart(part)
+            ? [{ messageId: message.id, part }]
+            : [],
         ),
       )
       .find(
@@ -888,7 +959,7 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
       if (startedOwner) {
         // Every runtime mounted over that owner re-reads the record, including
         // one mounted after this response started.
-        notifyHostApprovals(startedOwner);
+        notifyChatRuntimeState(startedOwner);
         return;
       }
 
@@ -924,6 +995,20 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     }
   };
 
+  // A thread message can join several AI SDK messages, so a slice through it
+  // has to end at the last of them rather than at the id the thread exposes.
+  const lastAISDKMessageId = (threadMessageId: string | null) => {
+    if (threadMessageId == null) return null;
+    const threadMessage = runtimeRef.current.thread
+      .getState()
+      .messages.find((message) => message.id === threadMessageId);
+    return (
+      (threadMessage &&
+        getVercelAIMessages<UI_MESSAGE>(threadMessage).at(-1)?.id) ??
+      threadMessageId
+    );
+  };
+
   const hasSeededRepositoryRef = useRef(false);
   const shouldFeedRepository =
     exportedMessageRepository != null &&
@@ -1008,16 +1093,26 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
         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]),
-          };
-        });
+        if (chatRuntimeState !== undefined && approvalOwner !== undefined) {
+          chatRuntimeState.cancelledMessageIds = new Set([
+            ...[...chatRuntimeState.cancelledMessageIds].filter((id) =>
+              liveIds.has(id),
+            ),
+            cancelledId,
+          ]);
+          notifyChatRuntimeState(approvalOwner);
+        } else {
+          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();
@@ -1052,9 +1147,10 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
         customToCreateMessage ?? toCreateMessage
       )<UI_MESSAGE>(message);
 
+      const parentId = lastAISDKMessageId(message.parentId);
       if (!(message.startRun ?? message.role === "user")) {
         chatHelpers.setMessages((current) => [
-          ...sliceMessagesUntil(current, message.parentId),
+          ...sliceMessagesUntil(current, parentId),
           toUIMessage<UI_MESSAGE>(createMessage, message.role),
         ]);
         return;
@@ -1062,7 +1158,7 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
 
       lastRunConfigRef.current = message.runConfig;
       chatHelpers.setMessages((current) =>
-        sliceMessagesUntil(current, message.parentId),
+        sliceMessagesUntil(current, parentId),
       );
       await chatHelpers.sendMessage(createMessage, {
         metadata: message.runConfig,
@@ -1103,7 +1199,7 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
       if (removedToolArtifact) markToolArtifactsChanged();
       if (removedToolInteractions) markToolInteractionsChanged();
       if (removedToolApprovalResponse || removedHostApprovalId) {
-        if (approvalOwner) notifyHostApprovals(approvalOwner);
+        if (approvalOwner) notifyChatRuntimeState(approvalOwner);
         setToolApprovalResponses(new Map(toolApprovalResponsesRef.current));
       }
 
@@ -1118,7 +1214,14 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     },
     onReload: async (parentId: string | null, config) => {
       lastRunConfigRef.current = config.runConfig;
-      const newMessages = sliceMessagesUntil(chatHelpers.messages, parentId);
+      const newMessages = sliceMessagesUntil(
+        chatHelpers.messages,
+        lastAISDKMessageId(parentId),
+      );
+      // regenerate drops a trailing assistant message, so the first message of
+      // the reloaded response stays for it to drop instead of the parent.
+      const reloaded = chatHelpers.messages[newMessages.length];
+      if (reloaded?.role === "assistant") newMessages.push(reloaded);
       chatHelpers.setMessages(newMessages);
 
       await chatHelpers.regenerate({ metadata: config.runConfig });
@@ -1154,7 +1257,10 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
       if (targetIndex >= 0 && targetIndex !== chatHelpers.messages.length - 1) {
         const target = chatHelpers.messages[targetIndex]!;
         const targetPart = target.parts.find(
-          (part) => isToolUIPart(part) && part.toolCallId === toolCallId,
+          (part) =>
+            typeof part?.type === "string" &&
+            isToolUIPart(part) &&
+            part.toolCallId === toolCallId,
         ) as { state?: string; preliminary?: boolean } | undefined;
         // An earlier message's settled output may already have reached the model, as the error a cancelling send writes does.
         if (
@@ -1172,7 +1278,11 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
             return {
               ...message,
               parts: message.parts.map((part) => {
-                if (!isToolUIPart(part) || part.toolCallId !== toolCallId)
+                if (
+                  typeof part?.type !== "string" ||
+                  !isToolUIPart(part) ||
+                  part.toolCallId !== toolCallId
+                )
                   return part;
 
                 const { preliminary: _preliminary, ...finalPart } =