@ulearning/ai-sdk-adapter 0.0.0-stage → 1.0.0-next.0

package/src/runtime/useExternalHistory.ts

@@ -0,0 +1,902 @@
+"use client";
+
+import type {
+  AssistantRuntime,
+  ThreadAssistantMessage,
+  ThreadHistoryAdapter,
+  ThreadMessage,
+  MessageFormatAdapter,
+  MessageFormatItem,
+  MessageFormatRepository,
+  ExportedMessageRepository,
+  RespondToToolApprovalOptions,
+  Unstable_ToolInteractionLog,
+} from "@ulearning/ai-core";
+import { getExternalStoreMessages } from "@ulearning/ai-core";
+import {
+  MessageRepository,
+  withParentsFirst,
+} from "@ulearning/ai-core/internal";
+import { useAui } from "@ulearning/ai-store";
+import {
+  useRef,
+  useEffect,
+  useState,
+  useCallback,
+  useMemo,
+  useSyncExternalStore,
+} from "@ulearning/ai-tap";
+import {
+  addToolData,
+  collectToolApprovalResponses,
+  collectToolArtifacts,
+  collectToolInteractions,
+  restoreToolData,
+  type StoredToolApprovalResponse,
+} from "./toolHistoryCodec";
+
+export const toExportedMessageRepository = <TMessage>(
+  toThreadMessages: (messages: TMessage[]) => ThreadMessage[],
+  messages: MessageFormatRepository<TMessage>,
+): ExportedMessageRepository => {
+  const converted = messages.messages.flatMap((m) => {
+    const message = toThreadMessages([m.message])[0];
+    if (!message) {
+      console.warn("Skipping a stored message that could not be loaded.");
+      return [];
+    }
+    return [{ ...m, message }];
+  });
+  const survivingIds = new Set<string>();
+  const survivors = withParentsFirst(converted, () => false).filter(
+    ({ parentId, message }) => {
+      if (parentId && !survivingIds.has(parentId)) return false;
+      survivingIds.add(message.id);
+      return true;
+    },
+  );
+
+  return {
+    headId:
+      messages.headId && survivingIds.has(messages.headId)
+        ? messages.headId
+        : null,
+    messages: survivors,
+  };
+};
+
+const hasUnansweredApproval = (message: ThreadMessage) =>
+  message.content.some(
+    (part) =>
+      part.type === "tool-call" &&
+      part.approval != null &&
+      part.approval.approved === undefined &&
+      part.approval.resolution === undefined,
+  );
+
+// Core reports a run paused on an unanswered approval as an interrupt, not as tool calls.
+const isAwaitingToolApproval = (message: ThreadMessage) =>
+  message.status?.type === "requires-action" &&
+  (message.status.reason === "tool-calls" ||
+    (message.status.reason === "interrupt" && hasUnansweredApproval(message)));
+
+// A run that settles without an assistant message after its prompt leaves nothing to report on its own.
+const isUnansweredRun = (
+  messages: readonly ThreadMessage[],
+  promptId: string,
+) => {
+  const promptIndex = messages.findIndex((message) => message.id === promptId);
+  return (
+    promptIndex === -1 ||
+    !messages
+      .slice(promptIndex + 1)
+      .some((message) => message.role === "assistant")
+  );
+};
+
+const createUnansweredRunMessage = (
+  reason: "cancelled" | "other",
+): ThreadAssistantMessage => ({
+  id: "unanswered-run",
+  role: "assistant",
+  content: [],
+  status: { type: "incomplete", reason },
+  createdAt: new Date(),
+  metadata: {
+    unstable_state: null,
+    unstable_annotations: [],
+    unstable_data: [],
+    steps: [],
+    custom: {},
+  },
+});
+
+const isTerminalMessage = (message: ThreadMessage) =>
+  message.status === undefined ||
+  message.status.type === "complete" ||
+  message.status.type === "incomplete";
+
+const encodeContent = <TMessage>(
+  storageFormatAdapter: MessageFormatAdapter<TMessage, any>,
+  item: MessageFormatItem<TMessage>,
+) => JSON.stringify(storageFormatAdapter.encode(item));
+
+export const useExternalHistory = <TMessage>(
+  runtimeRef: { current: AssistantRuntime },
+  historyAdapter: ThreadHistoryAdapter | undefined,
+  toThreadMessages: (messages: TMessage[]) => ThreadMessage[],
+  storageFormatAdapter: MessageFormatAdapter<TMessage, any>,
+  onSetMessages: (messages: TMessage[]) => void,
+  getToolArtifacts?: () => ReadonlyMap<string, unknown>,
+  onToolArtifactsRestored?: (next: Map<string, unknown>) => void,
+  getToolInteractions?: () => ReadonlyMap<string, Unstable_ToolInteractionLog>,
+  onToolInteractionsRestored?: (
+    next: Map<string, Unstable_ToolInteractionLog>,
+  ) => void,
+  toolApprovalResponses?: Map<string, RespondToToolApprovalOptions>,
+  onToolApprovalResponsesRestored?: () => void,
+) => {
+  const loadedRef = useRef(false);
+
+  const aui = useAui();
+  const optionalThreadListItem = useCallback(
+    () => (aui.threadListItem.source ? aui.threadListItem : null),
+    [aui],
+  );
+  const getRemoteId = useCallback(
+    () => optionalThreadListItem()?.getState().remoteId,
+    [optionalThreadListItem],
+  );
+  const remoteId = useSyncExternalStore(
+    aui.subscribe,
+    getRemoteId,
+    getRemoteId,
+  );
+
+  const [hasLoaded, setHasLoaded] = useState(false);
+  const historyReadyRef = useRef(false);
+  const deferredPersistRef = useRef(false);
+  const persistWhenReadyRef = useRef<() => void>(() => {});
+  const loadInFlightRef = useRef<Promise<void> | null>(null);
+
+  useEffect(() => {
+    historyReadyRef.current = hasLoaded;
+    if (hasLoaded) persistWhenReadyRef.current();
+  }, [hasLoaded]);
+
+  const historyIds = useRef(new Set<string>());
+  const deferredTelemetryIds = useRef(new Set<string>());
+  // A deleted message stays in the chat until its storage delete resolves, so passes queued before it leaves skip it, and its release is queued behind them.
+  const deletedMessageIds = useRef(new Map<string, boolean>());
+  // `content` is a snapshot taken at write time rather than a re-encode of `source`, because a retained message object can be mutated in place by the runtime that produced it.
+  const persistedInnerMessages = useRef(
+    new Map<string, { source: TMessage; content: string }>(),
+  );
+
+  const onSetMessagesRef = useRef(onSetMessages);
+  useEffect(() => {
+    onSetMessagesRef.current = onSetMessages;
+  });
+
+  const formatAdapter = useMemo(() => {
+    if (!historyAdapter) return undefined;
+    if (!historyAdapter.withFormat) {
+      throw new Error(
+        "useAISDKRuntime: ThreadHistoryAdapter is missing the required `withFormat` method.",
+      );
+    }
+    return historyAdapter.withFormat<TMessage, any>(storageFormatAdapter);
+  }, [historyAdapter, storageFormatAdapter]);
+
+  const isLoading = formatAdapter != null && !hasLoaded;
+
+  useEffect(() => {
+    if (!formatAdapter || loadedRef.current) return undefined;
+
+    const loadHistory = async () => {
+      try {
+        const repo = await formatAdapter.load();
+        if (optionalThreadListItem()?.getState().remoteId !== remoteId) return;
+        // A resumed stream or a send may have appended live messages while
+        // the load was in flight. The history imports beneath that live tail
+        // so neither side of the race is dropped.
+        const liveMessages = runtimeRef.current.thread.getState().messages;
+        const restoredArtifacts = new Map(
+          liveMessages.length > 0 ? getToolArtifacts?.() : undefined,
+        );
+        const restoredInteractions = new Map(
+          liveMessages.length > 0 ? getToolInteractions?.() : undefined,
+        );
+        if (liveMessages.length === 0) {
+          toolApprovalResponses?.clear();
+        }
+        const restoredMessages =
+          repo?.messages.map((item) => ({
+            ...item,
+            message: restoreToolData(
+              item.message,
+              restoredArtifacts,
+              restoredInteractions,
+              toolApprovalResponses,
+            ),
+          })) ?? [];
+        onToolArtifactsRestored?.(restoredArtifacts);
+        onToolInteractionsRestored?.(restoredInteractions);
+        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: restoredMessages[index]!.message,
+                content: encodeContent(storageFormatAdapter, m),
+              },
+            );
+          }
+          const converted = toExportedMessageRepository(
+            toThreadMessages,
+            restoredRepo,
+          );
+          const liveById = new Map(
+            liveMessages.map((message) => [message.id, message]),
+          );
+          const loadedMessages = converted.messages.map((item) => ({
+            ...item,
+            message: liveById.get(item.message.id) ?? item.message,
+          }));
+          const loadedIds = new Set(loadedMessages.map((m) => m.message.id));
+          const liveTail = liveMessages.filter((m) => !loadedIds.has(m.id));
+          const anchorId =
+            converted.headId ?? converted.messages.at(-1)?.message.id ?? null;
+          const merged =
+            liveTail.length === 0
+              ? { ...converted, messages: loadedMessages }
+              : {
+                  headId: liveTail.at(-1)!.id,
+                  messages: [
+                    ...loadedMessages,
+                    ...liveTail.map((message, index) => ({
+                      parentId:
+                        index === 0 ? anchorId : liveTail[index - 1]!.id,
+                      message,
+                    })),
+                  ],
+                };
+          runtimeRef.current.thread.import(merged);
+
+          const tempRepo = new MessageRepository();
+          tempRepo.import(merged);
+          const messages = tempRepo.getMessages();
+
+          onSetMessagesRef.current(
+            messages.flatMap(getExternalStoreMessages<TMessage>),
+          );
+
+          historyIds.current = new Set();
+          for (const m of loadedMessages) {
+            historyIds.current.add(m.message.id);
+            if (isAwaitingToolApproval(m.message)) {
+              deferredTelemetryIds.current.add(m.message.id);
+            }
+          }
+        }
+      } catch (error) {
+        console.error("Failed to load message history:", error);
+      } finally {
+        setHasLoaded(true);
+      }
+    };
+
+    formatAdapter.pin?.();
+
+    if (!remoteId) {
+      // eslint-disable-next-line react-hooks/set-state-in-effect
+      setHasLoaded(true);
+      return undefined;
+    }
+
+    const threadState = runtimeRef.current.thread.getState();
+    if (threadState.isRunning || threadState.messages.length > 0) {
+      loadedRef.current = true;
+      setHasLoaded(true);
+      return undefined;
+    }
+
+    loadedRef.current = true;
+    historyReadyRef.current = false;
+    setHasLoaded(false);
+    loadInFlightRef.current = loadHistory();
+    return undefined;
+  }, [
+    formatAdapter,
+    toThreadMessages,
+    runtimeRef,
+    optionalThreadListItem,
+    remoteId,
+    storageFormatAdapter,
+    getToolArtifacts,
+    onToolArtifactsRestored,
+    getToolInteractions,
+    onToolInteractionsRestored,
+    toolApprovalResponses,
+    onToolApprovalResponsesRestored,
+  ]);
+
+  const runStartRef = useRef<number | null>(null);
+  const persistTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
+  const persistInFlightRef = useRef<Promise<void>>(Promise.resolve());
+  const stepBoundariesRef = useRef<number[]>([]);
+  const wasRunningRef = useRef(false);
+  const toolCallCountRef = useRef(0);
+  const queuedPromptIdRef = useRef<string | null>(null);
+  const runPromptIdRef = useRef<string | null>(null);
+  const runPromptMessagesRef = useRef<TMessage[]>([]);
+  const settledPromptIdRef = useRef<string | null>(null);
+  const settledPromptMessagesRef = useRef<TMessage[]>([]);
+  const runCancelledRef = useRef(false);
+  const settledRunCancelledRef = useRef(false);
+
+  const markRunCancelled = useCallback(() => {
+    if (wasRunningRef.current) runCancelledRef.current = true;
+  }, []);
+
+  const releaseRemovedMessages = useCallback(
+    (messages: readonly ThreadMessage[]) => {
+      for (const [messageId, releaseQueued] of deletedMessageIds.current) {
+        if (
+          releaseQueued ||
+          messages.some((message) => message.id === messageId)
+        )
+          continue;
+        deletedMessageIds.current.set(messageId, true);
+        persistInFlightRef.current = persistInFlightRef.current.then(() => {
+          deletedMessageIds.current.delete(messageId);
+        });
+      }
+    },
+    [],
+  );
+
+  const clearDeletedMessage = useCallback(
+    (messageId: string, items: MessageFormatItem<TMessage>[]) => {
+      historyIds.current.delete(messageId);
+      deferredTelemetryIds.current.delete(messageId);
+      deletedMessageIds.current.set(messageId, false);
+      releaseRemovedMessages(runtimeRef.current.thread.getState().messages);
+      for (const item of items) {
+        persistedInnerMessages.current.delete(
+          storageFormatAdapter.getId(item.message),
+        );
+      }
+    },
+    [releaseRemovedMessages, runtimeRef, storageFormatAdapter],
+  );
+
+  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))
+          .filter((item) => !deletedMessageIds.current.has(item.id))
+          .flatMap(getExternalStoreMessages<TMessage>);
+        let parentId = previousMessages.at(-1)
+          ? storageFormatAdapter.getId(previousMessages.at(-1)!)
+          : null;
+        const storedToolArtifacts = collectToolArtifacts(
+          message,
+          getToolArtifacts?.(),
+        );
+        const storedToolInteractions = collectToolInteractions(
+          message,
+          getToolInteractions?.(),
+        );
+        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,
+      getToolArtifacts,
+      getToolInteractions,
+      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();
+      releaseRemovedMessages(threadState.messages);
+      const { isRunning } = threadState;
+      const wasRunning = wasRunningRef.current;
+      wasRunningRef.current = isRunning;
+
+      // Track step boundaries by content changes (more reliable than isRunning)
+      if (runStartRef.current != null) {
+        const lastMsg = threadState.messages.at(-1);
+        if (lastMsg?.role === "assistant") {
+          const currentToolCallCount = lastMsg.content.filter(
+            (p) => p.type === "tool-call",
+          ).length;
+          while (toolCallCountRef.current < currentToolCallCount) {
+            stepBoundariesRef.current.push(Date.now() - runStartRef.current);
+            toolCallCountRef.current++;
+          }
+        }
+      }
+
+      if (isRunning) {
+        if (runStartRef.current == null) {
+          runStartRef.current = Date.now();
+          stepBoundariesRef.current = [];
+          toolCallCountRef.current = 0;
+          adapter.pin?.();
+        }
+        const prompt = threadState.messages.findLast(
+          (message) => message.role === "user",
+        );
+        if (prompt) {
+          runPromptIdRef.current = prompt.id;
+          runPromptMessagesRef.current =
+            getExternalStoreMessages<TMessage>(prompt);
+        }
+        if (historyReadyRef.current) persistRunPrompt(threadState.messages);
+        // Cancel any pending persist — isRunning went back to true
+        if (persistTimerRef.current) {
+          clearTimeout(persistTimerRef.current);
+          persistTimerRef.current = null;
+        }
+        return;
+      }
+
+      if (wasRunning) {
+        settledPromptIdRef.current = runPromptIdRef.current;
+        settledPromptMessagesRef.current = runPromptMessagesRef.current;
+        settledRunCancelledRef.current = runCancelledRef.current;
+        runPromptIdRef.current = null;
+        runPromptMessagesRef.current = [];
+        runCancelledRef.current = false;
+        queuedPromptIdRef.current = null;
+        // 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
+      if (persistTimerRef.current) clearTimeout(persistTimerRef.current);
+      persistTimerRef.current = setTimeout(() => {
+        persistSettled(false);
+      }, 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,
+      ignoreHistoryReady = false,
+    ) {
+      persistTimerRef.current = null;
+      if (!historyReadyRef.current && !ignoreHistoryReady) {
+        deferredPersistRef.current = true;
+        return;
+      }
+      const latest = runtimeRef.current.thread.getState();
+      if (!ignoreRunning && latest.isRunning) return;
+      deferredPersistRef.current = false;
+
+      const boundaries = stepBoundariesRef.current;
+      const durationMs = boundaries.length > 0 ? boundaries.at(-1) : undefined;
+
+      // Fallback: if only 1 boundary but message has multiple steps, distribute evenly
+      if (boundaries.length === 1 && durationMs != null) {
+        const lastAssistant = latest.messages.findLast(
+          (m) => m.role === "assistant",
+        );
+        if (lastAssistant) {
+          const tcCount = lastAssistant.content.filter(
+            (p) => p.type === "tool-call",
+          ).length;
+          if (tcCount > 0) {
+            const totalSteps = tcCount + 1;
+            const stepDur = durationMs / totalSteps;
+            boundaries.length = 0;
+            for (let i = 0; i < totalSteps; i++) {
+              boundaries.push(Math.round((i + 1) * stepDur));
+            }
+          }
+        }
+      }
+
+      // Build per-step timestamps when there are multiple steps
+      const stepTimestamps =
+        boundaries.length > 1
+          ? boundaries.map((endMs, i) => ({
+              start_ms: i === 0 ? 0 : boundaries[i - 1]!,
+              end_ms: endMs,
+            }))
+          : undefined;
+
+      runStartRef.current = null;
+      stepBoundariesRef.current = [];
+      const promptId = settledPromptIdRef.current;
+      const promptMessages = settledPromptMessagesRef.current;
+      const runCancelled = settledRunCancelledRef.current;
+      settledPromptIdRef.current = null;
+      settledPromptMessagesRef.current = [];
+      settledRunCancelledRef.current = false;
+
+      const telemetryOptions = {
+        ...(durationMs != null ? { durationMs } : undefined),
+        ...(stepTimestamps != null ? { stepTimestamps } : undefined),
+      };
+
+      persistInFlightRef.current = persistInFlightRef.current
+        .then(async () => {
+          const { messages } = latest;
+          let lastInnerMessageId: string | null = null;
+          let lastUserMessageId: string | undefined;
+          let writeFailure: { error: unknown } | undefined;
+
+          const getLastInnerId = (msgs: TMessage[]): string | null =>
+            msgs.length > 0 ? storageFormatAdapter.getId(msgs.at(-1)!) : null;
+
+          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: addToolData(
+                msg,
+                toolArtifacts,
+                toolInteractions,
+                toolApprovalResponses,
+              ),
+            }));
+
+          for (const message of messages) {
+            if (deletedMessageIds.current.has(message.id)) continue;
+            if (message.role === "user") lastUserMessageId = message.id;
+            const innerMessages = getExternalStoreMessages<TMessage>(message);
+
+            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 =
+              isTerminal ||
+              (isAwaitingToolCalls && adapter.update !== undefined);
+
+            if (!isReady) {
+              lastInnerMessageId =
+                getLastInnerId(innerMessages) ?? lastInnerMessageId;
+              continue;
+            }
+
+            if (!historyIds.current.has(message.id)) {
+              historyIds.current.add(message.id);
+              deferredTelemetryIds.current.add(message.id);
+            }
+
+            const batchItems = toBatchItems(
+              innerMessages,
+              collectToolArtifacts(message, getToolArtifacts?.()),
+              collectToolInteractions(message, getToolInteractions?.()),
+              collectToolApprovalResponses(message, toolApprovalResponses),
+            );
+            let messageWriteFailure = writeFailure;
+            for (const item of batchItems) {
+              // Rows are stored parent first, so nothing after a failed append can be stored in this pass.
+              if (writeFailure) break;
+              const innerId = storageFormatAdapter.getId(item.message);
+              const persisted = persistedInnerMessages.current.get(innerId);
+              if (!persisted) {
+                try {
+                  await appendIfUnpersisted(item);
+                } catch (error) {
+                  console.error("Failed to persist message history:", error);
+                  writeFailure = { error };
+                  messageWriteFailure = writeFailure;
+                }
+              } else if (
+                persisted.source !== item.message &&
+                durationMs !== undefined &&
+                adapter.update
+              ) {
+                const content = encodeContent(storageFormatAdapter, item);
+                if (content === persisted.content) {
+                  persisted.source = item.message;
+                } else {
+                  try {
+                    await adapter.update(item, innerId);
+                    persistedInnerMessages.current.set(innerId, {
+                      source: item.message,
+                      content,
+                    });
+                  } catch (error) {
+                    // A failed update leaves the stale baseline behind so the next run stop retries it.
+                    messageWriteFailure ??= { error };
+                  }
+                }
+              }
+            }
+
+            lastInnerMessageId =
+              getLastInnerId(innerMessages) ?? lastInnerMessageId;
+
+            if (deferredTelemetryIds.current.has(message.id) && isTerminal) {
+              deferredTelemetryIds.current.delete(message.id);
+              adapter.reportTelemetry?.(
+                batchItems.map((item, index) => ({
+                  ...item,
+                  message: innerMessages[index]!,
+                })),
+                {
+                  ...telemetryOptions,
+                  ...(lastUserMessageId !== undefined
+                    ? { userMessageId: lastUserMessageId }
+                    : undefined),
+                  message,
+                  ...(messageWriteFailure && {
+                    persistenceError: messageWriteFailure.error,
+                  }),
+                },
+              );
+            }
+          }
+
+          if (promptId !== null && isUnansweredRun(messages, promptId)) {
+            const shouldDeletePrompt =
+              runCancelled &&
+              !messages.some((message) => message.id === promptId);
+            adapter.reportTelemetry?.([], {
+              ...telemetryOptions,
+              ...(!(shouldDeletePrompt && adapter.delete)
+                ? { userMessageId: promptId }
+                : undefined),
+              message: createUnansweredRunMessage(
+                runCancelled ? "cancelled" : "other",
+              ),
+            });
+            if (shouldDeletePrompt && adapter.delete) {
+              const items = promptMessages.map((message) => ({
+                parentId: null,
+                message,
+              }));
+              try {
+                await adapter.delete(items);
+                clearDeletedMessage(promptId, items);
+              } catch (error) {
+                console.error("Failed to delete message history:", error);
+              }
+            }
+          }
+        })
+        .catch((error) => {
+          console.error("Failed to persist message history:", error);
+        });
+    }
+
+    persistWhenReadyRef.current = () => {
+      const state = runtimeRef.current.thread.getState();
+      if (state.messages.length === 0 && settledPromptIdRef.current === null)
+        return;
+      if (state.isRunning) {
+        persistRunPrompt(state.messages);
+        return;
+      }
+      if (persistTimerRef.current) clearTimeout(persistTimerRef.current);
+      persistSettled(true);
+    };
+
+    function persistRunPrompt(messages: readonly ThreadMessage[]) {
+      const promptIndex = messages.findLastIndex(
+        (message) => message.role === "user",
+      );
+      const prompt = messages[promptIndex];
+      if (!prompt || prompt.id === queuedPromptIdRef.current) return;
+      queuedPromptIdRef.current = prompt.id;
+      const innerMessages = getExternalStoreMessages<TMessage>(prompt);
+      if (
+        innerMessages.every((innerMessage) =>
+          persistedInnerMessages.current.has(
+            storageFormatAdapter.getId(innerMessage),
+          ),
+        )
+      )
+        return;
+      const previousInnerMessage = messages
+        .slice(0, promptIndex)
+        .flatMap(getExternalStoreMessages<TMessage>)
+        .at(-1);
+
+      const persistence = persistInFlightRef.current.then(async () => {
+        if (deletedMessageIds.current.has(prompt.id)) return;
+        let parentId = previousInnerMessage
+          ? storageFormatAdapter.getId(previousInnerMessage)
+          : null;
+        for (const innerMessage of innerMessages) {
+          // Rows are appended parent first, so a prompt whose parent is not stored yet waits for the settled pass.
+          if (
+            parentId !== null &&
+            !persistedInnerMessages.current.has(parentId)
+          )
+            return;
+          await appendIfUnpersisted({ parentId, message: innerMessage });
+          parentId = storageFormatAdapter.getId(innerMessage);
+        }
+      });
+      persistInFlightRef.current = persistence.catch((error) => {
+        console.error("Failed to persist message history:", error);
+      });
+    }
+
+    async function appendIfUnpersisted(item: MessageFormatItem<TMessage>) {
+      const innerId = storageFormatAdapter.getId(item.message);
+      if (persistedInnerMessages.current.has(innerId)) return;
+      await adapter.append(item);
+      persistedInnerMessages.current.set(innerId, {
+        source: item.message,
+        content: encodeContent(storageFormatAdapter, item),
+      });
+    }
+
+    return () => {
+      persistWhenReadyRef.current = () => {};
+      unsubscribe();
+      const hasPendingTimer = persistTimerRef.current !== null;
+      if (persistTimerRef.current) {
+        clearTimeout(persistTimerRef.current);
+        persistTimerRef.current = null;
+      }
+      if (
+        !historyReadyRef.current &&
+        (deferredPersistRef.current || hasPendingTimer)
+      ) {
+        if (loadInFlightRef.current) {
+          void loadInFlightRef.current.then(() => persistSettled(true, true));
+        } else {
+          persistSettled(true, true);
+        }
+      } else if (hasPendingTimer) {
+        persistSettled(false);
+      }
+    };
+  }, [
+    formatAdapter,
+    storageFormatAdapter,
+    runtimeRef,
+    getToolArtifacts,
+    getToolInteractions,
+    toolApprovalResponses,
+    releaseRemovedMessages,
+    clearDeletedMessage,
+  ]);
+
+  const deleteMessage = useCallback(
+    async (messageId: string) => {
+      const deleteMessages = formatAdapter?.delete?.bind(formatAdapter);
+      if (!deleteMessages) return;
+
+      const messages = runtimeRef.current.thread.getState().messages;
+      const messageIndex = messages.findIndex((m) => m.id === messageId);
+      if (messageIndex === -1) return;
+
+      const previousInnerMessages = messages
+        .slice(0, messageIndex)
+        .flatMap(getExternalStoreMessages<TMessage>);
+      let parentId = previousInnerMessages.at(-1)
+        ? storageFormatAdapter.getId(previousInnerMessages.at(-1)!)
+        : null;
+      const itemsToDelete = getExternalStoreMessages<TMessage>(
+        messages[messageIndex]!,
+      ).map((message) => {
+        const item = { parentId, message };
+        parentId = storageFormatAdapter.getId(message);
+        return item;
+      });
+
+      const deletion = persistInFlightRef.current.then(async () => {
+        await deleteMessages(itemsToDelete);
+        clearDeletedMessage(messageId, itemsToDelete);
+      });
+
+      persistInFlightRef.current = deletion.catch(() => {});
+      await deletion;
+    },
+    [formatAdapter, runtimeRef, storageFormatAdapter, clearDeletedMessage],
+  );
+
+  return {
+    isLoading,
+    deleteMessage,
+    markRunCancelled,
+    persistToolInteractions,
+    persistToolApprovalResponses: persistToolInteractions,
+  };
+};