@assistant-ui/ai-sdk 0.0.6 → 0.0.8

package/src/runtime/useAISDKRuntime.ts

@@ -14,7 +14,7 @@ import type {
   CreateUIMessage,
   UseChatHelpers,
 } from "@ai-sdk/react";
-import { isToolUIPart, generateId } from "ai";
+import { isToolUIPart, generateId, getToolName } from "ai";
 import {
   useExternalStoreRuntime,
   useRuntimeAdapters,
@@ -37,12 +37,15 @@ import type {
   AppendMessage,
   RunConfig,
   McpAppMetadata,
+  RespondToToolApprovalOptions,
+  Unstable_ToolInteractionLog,
 } from "@assistant-ui/core";
 import {
   getExternalStoreMessages,
   pickExternalStoreSharedOptions,
 } from "@assistant-ui/core";
 import {
+  appendToolInteraction,
   consumeSuggestionResult,
   MessageRepository,
 } from "@assistant-ui/core/internal";
@@ -52,7 +55,10 @@ import { sliceMessagesUntil } from "../utils/sliceMessagesUntil";
 import { toCreateMessage } from "../converters/toCreateMessage";
 import { vercelAttachmentAdapter } from "../adapters/vercelAttachmentAdapter";
 import { getVercelAIMessages } from "../utils/getVercelAIMessages";
-import { AISDKMessageConverter } from "../converters/convertMessage";
+import {
+  AISDKMessageConverter,
+  type AISDKMessageConverterMetadata,
+} from "../converters/convertMessage";
 import { wrapModelContentEnvelope } from "../converters/modelContentEnvelope";
 import {
   type AISDKStorageFormat,
@@ -81,6 +87,23 @@ const toUIMessage = <UI_MESSAGE extends UIMessage>(
     role: createMessage.role ?? fallbackRole,
   }) as UI_MESSAGE;
 
+const toVoiceTranscriptUIMessage = <UI_MESSAGE extends UIMessage>(
+  message: ThreadMessage,
+): UI_MESSAGE =>
+  ({
+    id: message.id,
+    role: message.role,
+    parts: message.content
+      .filter((part) => part.type === "text")
+      .map((part) => ({ type: "text", text: part.text })),
+    metadata: {
+      ...(message.metadata.modality && { modality: message.metadata.modality }),
+      ...(Object.keys(message.metadata.custom).length > 0 && {
+        custom: message.metadata.custom,
+      }),
+    },
+  }) as UI_MESSAGE;
+
 export type AISDKRuntimeAdapter<UI_MESSAGE extends UIMessage = UIMessage> =
   ExternalStoreSharedOptions & {
     adapters?:
@@ -115,6 +138,24 @@ export type AISDKRuntimeAdapter<UI_MESSAGE extends UIMessage = UIMessage> =
      * Provide this to bridge resume-tool-call invocations into a custom handler.
      */
     onResumeToolCall?: ExternalStoreAdapter["onResumeToolCall"];
+    /**
+     * Answers tool approval requests through a host-owned channel instead of the AI SDK's `addToolApprovalResponse`.
+     *
+     * Called for every approval request in the thread with the complete response, including option and free-form answers. Hand requests the host does not own to `respondViaAISDK`, which is what runs when this option is omitted. The answer applies to the approval when the handler starts and is removed if it throws. It is never written into the `useChat` messages, so `sendAutomaticallyWhen` cannot forward it. With a history adapter, the answer is stored with its message once the handler resolves and returns on reload, and a second response to the same request rejects; without one it lasts as long as this runtime, and a runtime mounted again over the same chat shows the request open until the resumed run records its resolution in the chat.
+     *
+     * While a handler is set, an approval's `display`, `allowFreeform`, `dismissible` and `options` reach the renderer, because the handler can receive answers the AI SDK cannot carry. A stream declares them through the `approvalDescriptor` of its `tool-approval-request` chunk, the one approval field the AI SDK keeps opaque; the converter reads the request and answer fields from that descriptor when the approval itself lacks them.
+     */
+    onRespondToToolApproval?:
+      | ((
+          response: RespondToToolApprovalOptions,
+          context: {
+            toolCallId: string;
+            toolName: string;
+            /** Sends this response through the AI SDK's `addToolApprovalResponse`, which carries only `approved` and `reason`. */
+            respondViaAISDK: () => Promise<void>;
+          },
+        ) => Promise<void> | void)
+      | undefined;
     /**
      * How consecutive assistant messages are rendered.
      *
@@ -220,6 +261,11 @@ const useGeneratedSuggestions = (
 
 const NO_CANCELLED_MESSAGE_IDS: ReadonlySet<string> = new Set();
 
+const NO_TOOL_APPROVAL_RESPONSES: ReadonlyMap<
+  string,
+  RespondToToolApprovalOptions
+> = new Map();
+
 const toChatError = (error: Error): AssistantError => {
   const code = (error as { code?: unknown }).code;
   return {
@@ -243,6 +289,7 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     cancelPendingToolCallsOnSend = true,
     onResume,
     onResumeToolCall,
+    onRespondToToolApproval: customOnRespondToToolApproval,
     joinStrategy,
     messageRepository,
     unstable_onBranchChange,
@@ -256,14 +303,36 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     chatId: string;
     ids: ReadonlySet<string>;
   } | null>(null);
+  const [toolApprovalResponses, setToolApprovalResponses] = useState<
+    ReadonlyMap<string, RespondToToolApprovalOptions>
+  >(NO_TOOL_APPROVAL_RESPONSES);
+  const [toolArtifactEpoch, setToolArtifactEpoch] = useState(0);
+  const [toolInteractionEpoch, setToolInteractionEpoch] = useState(0);
+  const hostApprovalIdsRef = useRef(new Set<string>());
+  const toolApprovalResponsesRef = useRef<
+    Map<string, RespondToToolApprovalOptions>
+  >(new Map());
   const toolArgsKeyOrderCacheRef = useRef<Map<string, Map<string, string[]>>>(
     new Map(),
   );
   const toolLastInputCacheRef = useRef<Map<string, ReadonlyJSONObject>>(
     new Map(),
   );
+  const toolArgsTextCacheRef = useRef<
+    WeakMap<ReadonlyJSONObject, Map<string, string>>
+  >(new WeakMap());
   const mcpAppMetadataCacheRef = useRef<Map<string, McpAppMetadata>>(new Map());
+  const toolArtifactsRef = useRef<Map<string, unknown>>(new Map());
+  const toolInteractionsRef = useRef<Map<string, Unstable_ToolInteractionLog>>(
+    new Map(),
+  );
   const lastRunConfigRef = useRef<RunConfig | undefined>(undefined);
+  const markToolArtifactsChanged = useCallback(() => {
+    setToolArtifactEpoch((epoch) => epoch + 1);
+  }, []);
+  const markToolInteractionsChanged = useCallback(() => {
+    setToolInteractionEpoch((epoch) => epoch + 1);
+  }, []);
 
   const hasExecutingTools = Object.values(toolStatuses).some(
     (s) => s?.type === "executing",
@@ -285,6 +354,25 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     cancelledMessages?.chatId === chatHelpers.id
       ? cancelledMessages.ids
       : NO_CANCELLED_MESSAGE_IDS;
+  const supportsRichToolApprovalResponses =
+    customOnRespondToToolApproval != null;
+
+  const toThreadMessages = useCallback(
+    (sourceMessages: UI_MESSAGE[]) => {
+      const metadata: AISDKMessageConverterMetadata = {
+        supportsRichToolApprovalResponses,
+        toolArtifacts: toolArtifactsRef.current,
+        toolInteractions: toolInteractionsRef.current,
+        toolApprovalResponses: toolApprovalResponsesRef.current,
+      };
+      return AISDKMessageConverter.toThreadMessages(
+        sourceMessages,
+        false,
+        metadata,
+      );
+    },
+    [supportsRichToolApprovalResponses],
+  );
 
   const retractCancellation = useCallback(
     (chatId: string, messageId: string) => {
@@ -321,18 +409,23 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     isRunning,
     messages: chatHelpers.messages,
     joinStrategy,
-    metadata: useMemo(
+    metadata: useMemo<AISDKMessageConverterMetadata>(
       () => ({
         toolStatuses,
         messageTiming,
         toolArgsKeyOrderCache: toolArgsKeyOrderCacheRef.current,
+        toolArgsTextCache: toolArgsTextCacheRef.current,
         toolLastInputCache: toolLastInputCacheRef.current,
         mcpAppMetadataCache: mcpAppMetadataCacheRef.current,
+        toolArtifacts: toolArtifactsRef.current,
+        toolInteractions: toolInteractionsRef.current,
+        supportsRichToolApprovalResponses,
         ...(optimisticMessageId && { optimisticMessageId }),
         ...(chatHelpers.error && {
           error: toChatError(chatHelpers.error),
         }),
         ...(cancelledMessageIds.size > 0 && { cancelledMessageIds }),
+        ...(toolApprovalResponses.size > 0 && { toolApprovalResponses }),
       }),
       [
         toolStatuses,
@@ -340,6 +433,10 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
         optimisticMessageId,
         chatHelpers.error,
         cancelledMessageIds,
+        toolApprovalResponses,
+        supportsRichToolApprovalResponses,
+        toolArtifactEpoch,
+        toolInteractionEpoch,
       ],
     ),
   });
@@ -347,13 +444,11 @@ 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[],
+      toThreadMessages,
       messageRepository,
     );
     return converted.messages.length > 0 ? converted : undefined;
-  }, [messageRepository]);
+  }, [messageRepository, toThreadMessages]);
 
   const generatedSuggestions = useGeneratedSuggestions(
     suggestionAdapter,
@@ -367,12 +462,15 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     },
   }));
 
-  const { isLoading, deleteMessage: deleteHistoryMessage } = useExternalHistory(
+  const {
+    isLoading,
+    deleteMessage: deleteHistoryMessage,
+    persistToolInteractions,
+    persistToolApprovalResponses,
+  } = useExternalHistory(
     runtimeRef,
     adapters?.history ?? contextAdapters?.history,
-    AISDKMessageConverter.toThreadMessages as (
-      messages: UI_MESSAGE[],
-    ) => ThreadMessage[],
+    toThreadMessages,
     aiSDKV6FormatAdapter as MessageFormatAdapter<
       UI_MESSAGE,
       AISDKStorageFormat
@@ -380,6 +478,17 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     (messages) => {
       chatHelpers.setMessages(messages);
     },
+    toolArtifactsRef.current,
+    markToolArtifactsChanged,
+    toolInteractionsRef.current,
+    markToolInteractionsChanged,
+    toolApprovalResponsesRef.current,
+    () => {
+      hostApprovalIdsRef.current = new Set(
+        toolApprovalResponsesRef.current.keys(),
+      );
+      setToolApprovalResponses(new Map(toolApprovalResponsesRef.current));
+    },
   );
 
   const {
@@ -402,38 +511,112 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     if (!cancelPendingToolCallsOnSend) return;
 
     // The runtime auto-aborts in-flight tool invocations when a new run
-    // is dispatched (append() / startRun()). All we need to do here is
-    // mark any tool without a result as cancelled in the UI message list.
-
-    // Mark any tool without a result as cancelled (uses setMessages to avoid triggering sendAutomaticallyWhen)
+    // is dispatched (append() / startRun()), so this only has to mark the
+    // abandoned tools cancelled in the UI message list. Every non-terminal
+    // tool call qualifies wherever it sits: the run that produced it is over,
+    // and a staged `startRun: false` message can sit between it and the tail.
+    // Uses setMessages to avoid triggering sendAutomaticallyWhen.
     chatHelpers.setMessages((messages) => {
-      const lastMessage = messages.at(-1);
-      if (lastMessage?.role !== "assistant") return messages;
-
       let hasChanges = false;
-      const parts = lastMessage.parts?.map((part) => {
-        if (!isToolUIPart(part)) return part;
-        if (
-          part.state === "output-available" ||
-          part.state === "output-error" ||
-          part.state === "output-denied"
-        )
-          return part;
 
+      const next = messages.map((message) => {
+        if (message.role !== "assistant") return message;
+
+        let messageChanged = false;
+        const parts = message.parts?.map((part) => {
+          if (!isToolUIPart(part)) return part;
+          if (
+            part.state === "output-available" ||
+            part.state === "output-error" ||
+            part.state === "output-denied"
+          )
+            return part;
+
+          messageChanged = true;
+          const { approval: _approval, ...rest } = part;
+          return {
+            ...rest,
+            state: "output-error" as const,
+            errorText: "User cancelled tool call by sending a new message.",
+          };
+        });
+
+        if (!messageChanged) return message;
         hasChanges = true;
-        const { approval: _approval, ...rest } = part;
-        return {
-          ...rest,
-          state: "output-error" as const,
-          errorText: "User cancelled tool call by sending a new message.",
-        };
+        return { ...message, parts };
       });
 
       if (!hasChanges) return messages;
-      return [...messages.slice(0, -1), { ...lastMessage, parts }];
+      return next;
     });
   };
 
+  const respondViaAISDK = ({
+    approvalId,
+    approved,
+    reason,
+  }: RespondToToolApprovalOptions) =>
+    Promise.resolve(
+      chatHelpers.addToolApprovalResponse({
+        id: approvalId,
+        approved,
+        ...(reason != null && { reason }),
+        options: { metadata: lastRunConfigRef.current },
+      }),
+    );
+
+  const respondViaHost = async (
+    onRespond: NonNullable<AISDKRuntimeAdapter["onRespondToToolApproval"]>,
+    response: RespondToToolApprovalOptions,
+  ) => {
+    const { approvalId } = response;
+    const requested = chatHelpers.messages
+      .flatMap((message) =>
+        message.parts.flatMap((part) =>
+          isToolUIPart(part) ? [{ messageId: message.id, part }] : [],
+        ),
+      )
+      .find(
+        ({ part }) =>
+          part.state === "approval-requested" &&
+          part.approval.id === approvalId,
+      );
+    if (!requested || hostApprovalIdsRef.current.has(approvalId))
+      throw new Error(
+        `Tool approval ${approvalId} is not waiting for a response.`,
+      );
+
+    // A host answer stays out of the useChat messages, where sendAutomaticallyWhen would forward it to the chat route.
+    const applyResponse = (applied: boolean) => {
+      if (applied) hostApprovalIdsRef.current.add(approvalId);
+      else hostApprovalIdsRef.current.delete(approvalId);
+      if (applied) toolApprovalResponsesRef.current.set(approvalId, response);
+      else toolApprovalResponsesRef.current.delete(approvalId);
+      setToolApprovalResponses(new Map(toolApprovalResponsesRef.current));
+    };
+
+    applyResponse(true);
+    try {
+      await onRespond(response, {
+        toolCallId: requested.part.toolCallId,
+        toolName: getToolName(requested.part),
+        respondViaAISDK: async () => {
+          try {
+            await respondViaAISDK(response);
+          } finally {
+            applyResponse(false);
+          }
+        },
+      });
+    } catch (error) {
+      if (hostApprovalIdsRef.current.has(approvalId)) applyResponse(false);
+      throw error;
+    }
+    if (hostApprovalIdsRef.current.has(approvalId)) {
+      await persistToolApprovalResponses(requested.messageId);
+    }
+  };
+
   const hasSeededRepositoryRef = useRef(false);
   const shouldFeedRepository =
     exportedMessageRepository != null &&
@@ -441,6 +624,7 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     messages.length === 0;
 
   const runtime = useExternalStoreRuntime({
+    unstable_persistsHistory: true,
     isRunning: providerIsRunning,
     ...(shouldFeedRepository
       ? { messageRepository: exportedMessageRepository }
@@ -461,6 +645,11 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
           .filter(Boolean)
           .flat(),
       ),
+    onVoiceTranscript: (message: ThreadMessage) =>
+      chatHelpers.setMessages((current) => [
+        ...current,
+        toVoiceTranscriptUIMessage<UI_MESSAGE>(message),
+      ]),
     onExportExternalState: (): MessageFormatRepository<UI_MESSAGE> => {
       const exported = runtimeRef.current.thread.export();
 
@@ -501,10 +690,7 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     },
     onLoadExternalState: (repo: MessageFormatRepository<UI_MESSAGE>) => {
       // Convert MessageFormatRepository to ExportedMessageRepository
-      const exportedRepo = toExportedMessageRepository(
-        AISDKMessageConverter.toThreadMessages,
-        repo,
-      );
+      const exportedRepo = toExportedMessageRepository(toThreadMessages, repo);
 
       // Import into the thread's MessageRepository
       runtimeRef.current.thread.import(exportedRepo);
@@ -584,6 +770,34 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
 
       await deleteHistoryMessage(messageId);
 
+      let removedToolArtifact = false;
+      let removedToolInteractions = false;
+      let removedToolApprovalResponse = false;
+      let removedHostApprovalId = false;
+      for (const part of threadMessages[messageIndex]!.content) {
+        if (part.type === "tool-call") {
+          removedToolArtifact =
+            toolArtifactsRef.current.delete(part.toolCallId) ||
+            removedToolArtifact;
+          removedToolInteractions =
+            toolInteractionsRef.current.delete(part.toolCallId) ||
+            removedToolInteractions;
+          if (part.approval) {
+            removedToolApprovalResponse =
+              toolApprovalResponsesRef.current.delete(part.approval.id) ||
+              removedToolApprovalResponse;
+            removedHostApprovalId =
+              hostApprovalIdsRef.current.delete(part.approval.id) ||
+              removedHostApprovalId;
+          }
+        }
+      }
+      if (removedToolArtifact) markToolArtifactsChanged();
+      if (removedToolInteractions) markToolInteractionsChanged();
+      if (removedToolApprovalResponse || removedHostApprovalId) {
+        setToolApprovalResponses(new Map(toolApprovalResponsesRef.current));
+      }
+
       const deleteIds = new Set(
         getExternalStoreMessages<UI_MESSAGE>(threadMessages[messageIndex]!).map(
           (message) => message.id,
@@ -605,8 +819,13 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
       toolName,
       result,
       isError,
+      artifact,
       modelContent,
     }) => {
+      if (artifact !== undefined) {
+        toolArtifactsRef.current.set(toolCallId, artifact);
+        markToolArtifactsChanged();
+      }
       const options = { metadata: lastRunConfigRef.current };
       if (isError) {
         return Promise.resolve(
@@ -634,15 +853,24 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
         );
       }
     },
-    onRespondToToolApproval: ({ approvalId, approved, reason }) =>
-      Promise.resolve(
-        chatHelpers.addToolApprovalResponse({
-          id: approvalId,
-          approved,
-          ...(reason != null && { reason }),
-          options: { metadata: lastRunConfigRef.current },
-        }),
-      ),
+    onRespondToToolApproval: customOnRespondToToolApproval
+      ? (response) => respondViaHost(customOnRespondToToolApproval, response)
+      : respondViaAISDK,
+    unstable_onRecordToolInteraction: ({
+      messageId,
+      toolCallId,
+      interaction,
+    }) => {
+      toolInteractionsRef.current.set(
+        toolCallId,
+        appendToolInteraction(
+          toolInteractionsRef.current.get(toolCallId),
+          interaction,
+        ),
+      );
+      markToolInteractionsChanged();
+      return persistToolInteractions(messageId);
+    },
     ...pickExternalStoreSharedOptions(adapter),
     ...(adapter.unstable_messageRepositoryInstance && {
       unstable_messageRepositoryInstance: