@assistant-ui/ai-sdk 0.0.6 → 0.0.8

package/src/runtime/useExternalHistory.ts

@@ -8,9 +8,14 @@ import type {
   MessageFormatItem,
   MessageFormatRepository,
   ExportedMessageRepository,
+  RespondToToolApprovalOptions,
+  Unstable_ToolInteractionLog,
 } from "@assistant-ui/core";
 import { getExternalStoreMessages } from "@assistant-ui/core";
-import { MessageRepository } from "@assistant-ui/core/internal";
+import {
+  MessageRepository,
+  readToolInteractionLog,
+} from "@assistant-ui/core/internal";
 import { useAui } from "@assistant-ui/store";
 import {
   useRef,
@@ -50,6 +55,211 @@ 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 TOOL_ARTIFACTS_METADATA_KEY = "__aui_toolArtifacts";
+const TOOL_INTERACTIONS_METADATA_KEY = "__aui_toolInteractions";
+const TOOL_APPROVAL_RESPONSES_METADATA_KEY = "__aui_toolApprovalResponses";
+
+type StoredToolApprovalResponse = Omit<
+  RespondToToolApprovalOptions,
+  "approvalId"
+>;
+
+const isRecord = (value: unknown): value is Record<string, unknown> =>
+  value !== null && typeof value === "object" && !Array.isArray(value);
+
+const collectToolArtifacts = (
+  message: ThreadMessage,
+  toolArtifacts: ReadonlyMap<string, unknown> | undefined,
+) => {
+  if (!toolArtifacts) return undefined;
+  const entries = message.content.flatMap((part) => {
+    if (part.type !== "tool-call") return [];
+    const artifact = toolArtifacts.get(part.toolCallId);
+    return artifact === undefined ? [] : [[part.toolCallId, artifact] as const];
+  });
+  return entries.length > 0 ? Object.fromEntries(entries) : undefined;
+};
+
+const collectToolInteractions = (
+  message: ThreadMessage,
+  toolInteractions:
+    | ReadonlyMap<string, Unstable_ToolInteractionLog>
+    | undefined,
+) => {
+  if (!toolInteractions) return undefined;
+  const entries = message.content.flatMap((part) => {
+    if (part.type !== "tool-call") return [];
+    const interactions = toolInteractions.get(part.toolCallId);
+    return interactions === undefined
+      ? []
+      : [[part.toolCallId, interactions] as const];
+  });
+  return entries.length > 0 ? Object.fromEntries(entries) : undefined;
+};
+
+const collectToolApprovalResponses = (
+  message: ThreadMessage,
+  toolApprovalResponses:
+    | ReadonlyMap<string, RespondToToolApprovalOptions>
+    | undefined,
+) => {
+  if (!toolApprovalResponses) return undefined;
+  const entries = message.content.flatMap((part) => {
+    if (part.type !== "tool-call" || !part.approval) return [];
+    const response = toolApprovalResponses.get(part.approval.id);
+    if (!response) return [];
+    return [
+      [
+        part.approval.id,
+        {
+          approved: response.approved,
+          ...(response.optionId != null && { optionId: response.optionId }),
+          ...(response.text != null && { text: response.text }),
+          ...(response.reason != null && { reason: response.reason }),
+        },
+      ] as const,
+    ];
+  });
+  return entries.length > 0 ? Object.fromEntries(entries) : undefined;
+};
+
+const addToolData = <TMessage>(
+  message: TMessage,
+  toolArtifacts: Record<string, unknown> | undefined,
+  toolInteractions: Record<string, Unstable_ToolInteractionLog> | undefined,
+  toolApprovalResponses: Record<string, StoredToolApprovalResponse> | undefined,
+): TMessage => {
+  if (
+    (!toolArtifacts && !toolInteractions && !toolApprovalResponses) ||
+    !isRecord(message) ||
+    !Array.isArray(message.parts)
+  )
+    return message;
+  const toolCallIds = message.parts.flatMap((part) => {
+    if (!isRecord(part) || typeof part.toolCallId !== "string") return [];
+    return [part.toolCallId];
+  });
+  const artifacts = toolArtifacts
+    ? Object.fromEntries(
+        toolCallIds.flatMap((toolCallId) =>
+          Object.hasOwn(toolArtifacts, toolCallId)
+            ? [[toolCallId, toolArtifacts[toolCallId]] as const]
+            : [],
+        ),
+      )
+    : undefined;
+  const interactions = toolInteractions
+    ? Object.fromEntries(
+        toolCallIds.flatMap((toolCallId) =>
+          Object.hasOwn(toolInteractions, toolCallId)
+            ? [[toolCallId, toolInteractions[toolCallId]] as const]
+            : [],
+        ),
+      )
+    : undefined;
+  const approvalIds = message.parts.flatMap((part) => {
+    if (!isRecord(part) || !isRecord(part.approval)) return [];
+    const approvalId = part.approval.id;
+    return typeof approvalId === "string" ? [approvalId] : [];
+  });
+  const approvalResponses = toolApprovalResponses
+    ? Object.fromEntries(
+        approvalIds.flatMap((approvalId) =>
+          Object.hasOwn(toolApprovalResponses, approvalId)
+            ? [[approvalId, toolApprovalResponses[approvalId]] as const]
+            : [],
+        ),
+      )
+    : undefined;
+  const hasArtifacts = !!artifacts && Object.keys(artifacts).length > 0;
+  const hasInteractions =
+    !!interactions && Object.keys(interactions).length > 0;
+  const hasApprovalResponses =
+    !!approvalResponses && Object.keys(approvalResponses).length > 0;
+  if (!hasArtifacts && !hasInteractions && !hasApprovalResponses)
+    return message;
+  const metadata = isRecord(message.metadata) ? message.metadata : {};
+  return {
+    ...message,
+    metadata: {
+      ...metadata,
+      ...(hasArtifacts && { [TOOL_ARTIFACTS_METADATA_KEY]: artifacts }),
+      ...(hasInteractions && {
+        [TOOL_INTERACTIONS_METADATA_KEY]: interactions,
+      }),
+      ...(hasApprovalResponses && {
+        [TOOL_APPROVAL_RESPONSES_METADATA_KEY]: approvalResponses,
+      }),
+    },
+  } as TMessage;
+};
+
+const restoreToolData = <TMessage>(
+  message: TMessage,
+  toolArtifacts: Map<string, unknown> | undefined,
+  toolInteractions: Map<string, Unstable_ToolInteractionLog> | undefined,
+  toolApprovalResponses: Map<string, RespondToToolApprovalOptions> | undefined,
+): TMessage => {
+  if (!isRecord(message) || !isRecord(message.metadata)) return message;
+  const metadata = message.metadata;
+  const hasArtifacts = Object.hasOwn(metadata, TOOL_ARTIFACTS_METADATA_KEY);
+  const hasInteractions = Object.hasOwn(
+    metadata,
+    TOOL_INTERACTIONS_METADATA_KEY,
+  );
+  const hasApprovalResponses = Object.hasOwn(
+    metadata,
+    TOOL_APPROVAL_RESPONSES_METADATA_KEY,
+  );
+  if (!hasArtifacts && !hasInteractions && !hasApprovalResponses)
+    return message;
+  const artifacts = metadata[TOOL_ARTIFACTS_METADATA_KEY];
+  if (toolArtifacts && isRecord(artifacts)) {
+    for (const [toolCallId, artifact] of Object.entries(artifacts)) {
+      toolArtifacts.set(toolCallId, artifact);
+    }
+  }
+  const interactions = metadata[TOOL_INTERACTIONS_METADATA_KEY];
+  if (toolInteractions && isRecord(interactions)) {
+    for (const [toolCallId, value] of Object.entries(interactions)) {
+      const log = readToolInteractionLog(value);
+      if (log) toolInteractions.set(toolCallId, log);
+    }
+  }
+  const approvalResponses = metadata[TOOL_APPROVAL_RESPONSES_METADATA_KEY];
+  if (toolApprovalResponses && isRecord(approvalResponses)) {
+    for (const [approvalId, value] of Object.entries(approvalResponses)) {
+      if (!isRecord(value) || typeof value.approved !== "boolean") continue;
+      toolApprovalResponses.set(approvalId, {
+        approvalId,
+        approved: value.approved,
+        ...(typeof value.optionId === "string" && {
+          optionId: value.optionId,
+        }),
+        ...(typeof value.text === "string" && { text: value.text }),
+        ...(typeof value.reason === "string" && { reason: value.reason }),
+      });
+    }
+  }
+  const {
+    [TOOL_ARTIFACTS_METADATA_KEY]: _,
+    [TOOL_INTERACTIONS_METADATA_KEY]: __,
+    [TOOL_APPROVAL_RESPONSES_METADATA_KEY]: ___,
+    ...restMetadata
+  } = metadata;
+  const { metadata: _metadata, ...restMessage } = message;
+  return (
+    Object.keys(restMetadata).length === 0
+      ? restMessage
+      : { ...restMessage, metadata: restMetadata }
+  ) as TMessage;
+};
+
 const encodeContent = <TMessage>(
   storageFormatAdapter: MessageFormatAdapter<TMessage, any>,
   item: MessageFormatItem<TMessage>,
@@ -61,6 +271,12 @@ export const useExternalHistory = <TMessage>(
   toThreadMessages: (messages: TMessage[]) => ThreadMessage[],
   storageFormatAdapter: MessageFormatAdapter<TMessage, any>,
   onSetMessages: (messages: TMessage[]) => void,
+  toolArtifacts?: Map<string, unknown>,
+  onToolArtifactsRestored?: () => void,
+  toolInteractions?: Map<string, Unstable_ToolInteractionLog>,
+  onToolInteractionsRestored?: () => void,
+  toolApprovalResponses?: Map<string, RespondToToolApprovalOptions>,
+  onToolApprovalResponsesRestored?: () => void,
 ) => {
   const loadedRef = useRef(false);
   const [itemEpoch, setItemEpoch] = useState(0);
@@ -103,17 +319,37 @@ export const useExternalHistory = <TMessage>(
     const loadHistory = async () => {
       try {
         const repo = await formatAdapter.load();
-        if (repo && repo.messages.length > 0) {
-          for (const m of repo.messages) {
+        toolArtifacts?.clear();
+        toolInteractions?.clear();
+        toolApprovalResponses?.clear();
+        const restoredMessages =
+          repo?.messages.map((item) => ({
+            ...item,
+            message: restoreToolData(
+              item.message,
+              toolArtifacts,
+              toolInteractions,
+              toolApprovalResponses,
+            ),
+          })) ?? [];
+        onToolArtifactsRestored?.();
+        onToolInteractionsRestored?.();
+        onToolApprovalResponsesRestored?.();
+        if (repo && restoredMessages.length > 0) {
+          const restoredRepo = { ...repo, messages: restoredMessages };
+          for (const [index, m] of repo.messages.entries()) {
             persistedInnerMessages.current.set(
               storageFormatAdapter.getId(m.message),
               {
-                source: m.message,
+                source: restoredMessages[index]!.message,
                 content: encodeContent(storageFormatAdapter, m),
               },
             );
           }
-          const converted = toExportedMessageRepository(toThreadMessages, repo);
+          const converted = toExportedMessageRepository(
+            toThreadMessages,
+            restoredRepo,
+          );
           runtimeRef.current.thread.import(converted);
 
           const tempRepo = new MessageRepository();
@@ -172,6 +408,12 @@ export const useExternalHistory = <TMessage>(
     aui,
     itemEpoch,
     storageFormatAdapter,
+    toolArtifacts,
+    onToolArtifactsRestored,
+    toolInteractions,
+    onToolInteractionsRestored,
+    toolApprovalResponses,
+    onToolApprovalResponsesRestored,
   ]);
 
   const runStartRef = useRef<number | null>(null);
@@ -181,10 +423,96 @@ export const useExternalHistory = <TMessage>(
   const wasRunningRef = useRef(false);
   const toolCallCountRef = useRef(0);
 
+  const persistToolInteractions = useCallback(
+    (messageId: string) => {
+      const persistence = persistInFlightRef.current.then(async () => {
+        if (!formatAdapter?.update) return;
+        const messages = runtimeRef.current.thread.getState().messages;
+        const message = messages.find(
+          (item) =>
+            item.id === messageId ||
+            getExternalStoreMessages<TMessage>(item).some(
+              (innerMessage) =>
+                storageFormatAdapter.getId(innerMessage) === messageId,
+            ),
+        );
+        if (!message) return;
+
+        const previousMessages = messages
+          .slice(0, messages.indexOf(message))
+          .flatMap(getExternalStoreMessages<TMessage>);
+        let parentId = previousMessages.at(-1)
+          ? storageFormatAdapter.getId(previousMessages.at(-1)!)
+          : null;
+        const storedToolArtifacts = collectToolArtifacts(
+          message,
+          toolArtifacts,
+        );
+        const storedToolInteractions = collectToolInteractions(
+          message,
+          toolInteractions,
+        );
+        const storedToolApprovalResponses = collectToolApprovalResponses(
+          message,
+          toolApprovalResponses,
+        );
+
+        for (const innerMessage of getExternalStoreMessages<TMessage>(
+          message,
+        )) {
+          const item = {
+            parentId,
+            message: addToolData(
+              innerMessage,
+              storedToolArtifacts,
+              storedToolInteractions,
+              storedToolApprovalResponses,
+            ),
+          };
+          const innerId = storageFormatAdapter.getId(item.message);
+          const persisted = persistedInnerMessages.current.get(innerId);
+          if (persisted) {
+            const content = encodeContent(storageFormatAdapter, item);
+            if (content === persisted.content) {
+              persisted.source = item.message;
+            } else {
+              await formatAdapter.update(item, innerId);
+              persistedInnerMessages.current.set(innerId, {
+                source: item.message,
+                content,
+              });
+            }
+          }
+          parentId = innerId;
+        }
+      });
+      persistInFlightRef.current = persistence.catch(() => {});
+      return persistence.catch((error) => {
+        console.error("Failed to persist tool data:", error);
+      });
+    },
+    [
+      formatAdapter,
+      runtimeRef,
+      storageFormatAdapter,
+      toolArtifacts,
+      toolInteractions,
+      toolApprovalResponses,
+    ],
+  );
+
   useEffect(() => {
     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 +548,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
@@ -301,22 +633,33 @@ export const useExternalHistory = <TMessage>(
           const getLastInnerId = (msgs: TMessage[]): string | null =>
             msgs.length > 0 ? storageFormatAdapter.getId(msgs.at(-1)!) : null;
 
-          const toBatchItems = (msgs: TMessage[]) =>
+          const toBatchItems = (
+            msgs: TMessage[],
+            toolArtifacts: Record<string, unknown> | undefined,
+            toolInteractions:
+              | Record<string, Unstable_ToolInteractionLog>
+              | undefined,
+            toolApprovalResponses:
+              | Record<string, StoredToolApprovalResponse>
+              | undefined,
+          ) =>
             msgs.map((msg, idx) => ({
               parentId:
                 idx === 0
                   ? lastInnerMessageId
                   : storageFormatAdapter.getId(msgs[idx - 1]!),
-              message: msg,
+              message: addToolData(
+                msg,
+                toolArtifacts,
+                toolInteractions,
+                toolApprovalResponses,
+              ),
             }));
 
           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 =
@@ -334,7 +677,12 @@ export const useExternalHistory = <TMessage>(
               deferredTelemetryIds.current.add(message.id);
             }
 
-            const batchItems = toBatchItems(innerMessages);
+            const batchItems = toBatchItems(
+              innerMessages,
+              collectToolArtifacts(message, toolArtifacts),
+              collectToolInteractions(message, toolInteractions),
+              collectToolApprovalResponses(message, toolApprovalResponses),
+            );
             for (const item of batchItems) {
               const innerId = storageFormatAdapter.getId(item.message);
               const persisted = persistedInnerMessages.current.get(innerId);
@@ -371,10 +719,16 @@ export const useExternalHistory = <TMessage>(
 
             if (deferredTelemetryIds.current.has(message.id) && isTerminal) {
               deferredTelemetryIds.current.delete(message.id);
-              adapter.reportTelemetry?.(batchItems, {
-                ...telemetryOptions,
-                message,
-              });
+              adapter.reportTelemetry?.(
+                batchItems.map((item, index) => ({
+                  ...item,
+                  message: innerMessages[index]!,
+                })),
+                {
+                  ...telemetryOptions,
+                  message,
+                },
+              );
             }
           }
         })
@@ -391,7 +745,14 @@ export const useExternalHistory = <TMessage>(
         persistSettled(false);
       }
     };
-  }, [formatAdapter, storageFormatAdapter, runtimeRef]);
+  }, [
+    formatAdapter,
+    storageFormatAdapter,
+    runtimeRef,
+    toolArtifacts,
+    toolInteractions,
+    toolApprovalResponses,
+  ]);
 
   const deleteMessage = useCallback(
     async (messageId: string) => {
@@ -434,5 +795,10 @@ export const useExternalHistory = <TMessage>(
     [formatAdapter, runtimeRef, storageFormatAdapter],
   );
 
-  return { isLoading, deleteMessage };
+  return {
+    isLoading,
+    deleteMessage,
+    persistToolInteractions,
+    persistToolApprovalResponses: persistToolInteractions,
+  };
 };

package/src/runtime/useResourceCleanup.ts

@@ -1,11 +1,19 @@
 import { useEffect, useRef } from "react";
 import { useAssistantClientDestroySignal } from "@assistant-ui/store/internal";
 
-export const useResourceCleanup = (enabled: boolean, cleanup: () => void) => {
+export const useResourceCleanup = (
+  enabled: boolean,
+  cleanup: () => void,
+  hostDestroySignal?: AbortSignal,
+) => {
   const destroySignal = useAssistantClientDestroySignal();
   const cleanupRef = useRef(cleanup);
   const enabledRef = useRef(enabled);
-  const registeredSignalRef = useRef<AbortSignal | undefined>(undefined);
+  const registrationRef = useRef<{
+    destroySignal: AbortSignal | undefined;
+    hostDestroySignal: AbortSignal | undefined;
+    controller: AbortController;
+  } | null>(null);
 
   useEffect(() => {
     cleanupRef.current = cleanup;
@@ -13,20 +21,38 @@ export const useResourceCleanup = (enabled: boolean, cleanup: () => void) => {
   });
 
   useEffect(() => {
-    if (!enabled || !destroySignal) return undefined;
-    if (registeredSignalRef.current === destroySignal) return undefined;
+    if (!enabled || (!destroySignal && !hostDestroySignal)) return undefined;
+    const current = registrationRef.current;
+    if (
+      current !== null &&
+      current.destroySignal === destroySignal &&
+      current.hostDestroySignal === hostDestroySignal
+    ) {
+      return undefined;
+    }
+    current?.controller.abort();
 
-    registeredSignalRef.current = destroySignal;
-    destroySignal.addEventListener(
-      "abort",
-      () => {
-        if (enabledRef.current) cleanupRef.current();
-      },
-      { once: true },
-    );
+    const registration = new AbortController();
+    registrationRef.current = {
+      destroySignal,
+      hostDestroySignal,
+      controller: registration,
+    };
+    const run = () => {
+      if (registration.signal.aborted) return;
+      registration.abort();
+      if (enabledRef.current) cleanupRef.current();
+    };
+    if (destroySignal?.aborted || hostDestroySignal?.aborted) {
+      run();
+      return undefined;
+    }
+    const options = { once: true, signal: registration.signal };
+    destroySignal?.addEventListener("abort", run, options);
+    hostDestroySignal?.addEventListener("abort", run, options);
 
-    // The listener must survive standalone soft unmounts so a later permanent
-    // client destroy still cleans up the retained resource state.
+    // The listeners must survive standalone soft unmounts so a later
+    // permanent destroy still cleans up the retained resource state.
     return undefined;
-  }, [destroySignal, enabled]);
+  }, [destroySignal, enabled, hostDestroySignal]);
 };

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/usage.ts

@@ -1,4 +1,4 @@
-/// <reference types="@assistant-ui/core/react" />
+/// <reference types="@assistant-ui/core/react" preserve="true" />
 import { useMemo } from "react";
 import { useAuiState } from "@assistant-ui/store";