@assistant-ui/ai-sdk 0.0.8 → 0.0.9

package/src/runtime/useAISDKRuntime.ts

@@ -20,6 +20,7 @@ import {
   useRuntimeAdapters,
   type JoinStrategy,
 } from "@assistant-ui/core/react";
+import { useReplaySafeEffect } from "@assistant-ui/store/internal";
 import type {
   SuggestionAdapter,
   ThreadSuggestion,
@@ -114,6 +115,13 @@ export type AISDKRuntimeAdapter<UI_MESSAGE extends UIMessage = UIMessage> =
       | undefined;
     toCreateMessage?: CustomToCreateMessageFunction;
     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.
+     */
+    unstable_hostApprovalOwner?: object | undefined;
     /**
      * Whether to automatically cancel pending interactive tool calls when the user sends a new message.
      *
@@ -141,7 +149,7 @@ export type AISDKRuntimeAdapter<UI_MESSAGE extends UIMessage = UIMessage> =
     /**
      * 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.
+     * 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. With `unstable_hostApprovalOwner`, in-memory answers survive runtime remounts over the same chat; without an owner or history adapter, they last only as long as the runtime. Bring the resumed run back into the chat, for example with `resumeStream`, and have the endpoint refuse a second resume.
      *
      * 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.
      */
@@ -250,7 +258,7 @@ const useGeneratedSuggestions = (
     })();
   }, [hasAdapter, isRunning]);
 
-  useEffect(() => {
+  useReplaySafeEffect(() => {
     return () => {
       controllerRef.current?.abort();
     };
@@ -261,11 +269,188 @@ const useGeneratedSuggestions = (
 
 const NO_CANCELLED_MESSAGE_IDS: ReadonlySet<string> = new Set();
 
+const NO_SUPERSEDED_APPROVAL_PROJECTION = Object.freeze({
+  approvalIds: Object.freeze(new Set<string>()),
+  statusMessageIds: Object.freeze(new Set<string>()),
+});
+
 const NO_TOOL_APPROVAL_RESPONSES: ReadonlyMap<
   string,
   RespondToToolApprovalOptions
 > = new Map();
 
+/**
+ * A host answer is deliberately kept out of the `useChat` messages, so nothing
+ * in the chat records it. Held in runtime state it would die with the runtime,
+ * and a runtime mounted again over the same chat would show the request open
+ * and take a second answer. Keyed on the chat instead, the answer lives as
+ * long as the chat it belongs to, and is collected with it.
+ */
+type OwnedApproval = {
+  response: RespondToToolApprovalOptions;
+};
+
+const hostToolApprovalsByChat = new WeakMap<
+  object,
+  Map<string, OwnedApproval>
+>();
+
+const toApprovalResponses = (
+  owned: ReadonlyMap<string, OwnedApproval> | undefined,
+): ReadonlyMap<string, RespondToToolApprovalOptions> =>
+  owned && owned.size > 0
+    ? new Map([...owned].map(([id, entry]) => [id, entry.response]))
+    : NO_TOOL_APPROVAL_RESPONSES;
+
+/**
+ * The answers live on the owner, but each mounted runtime renders them from
+ * its own state, so a write has to be announced: the runtime that performed it
+ * 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 subscribeToHostApprovals = (owner: object, listener: () => void) => {
+  const listeners = hostApprovalListenersByChat.get(owner) ?? new Set();
+  hostApprovalListenersByChat.set(owner, listeners);
+  listeners.add(listener);
+  return () => {
+    listeners.delete(listener);
+  };
+};
+
+const notifyHostApprovals = (owner: object) => {
+  for (const listener of [...(hostApprovalListenersByChat.get(owner) ?? [])]) {
+    listener();
+  }
+};
+
+const getSupersededApprovalProjection = <UI_MESSAGE extends UIMessage>(
+  messages: readonly UI_MESSAGE[],
+  hostApprovalIds: ReadonlySet<string>,
+  joinStrategy: JoinStrategy | undefined,
+  isRunning: boolean,
+) => {
+  const approvalIds = new Set<string>();
+  const statusMessageIds = new Set<string>();
+  const lastIndex = messages.length - 1;
+  let lastAssistant: UI_MESSAGE | undefined;
+  let hasSupersededApproval = false;
+  let previousWasVoice = false;
+
+  const flush = () => {
+    const hasOpenToolPart =
+      lastAssistant === messages[lastIndex] &&
+      lastAssistant?.parts?.some((part) => {
+        if (
+          !isToolUIPart(part) ||
+          part.state === "output-available" ||
+          part.state === "output-error" ||
+          part.state === "output-denied"
+        )
+          return false;
+
+        const approval = (
+          part as {
+            approval?: {
+              resolution?: unknown;
+              descriptor?: unknown;
+            };
+          }
+        ).approval;
+        const resolution =
+          approval?.resolution ??
+          (approval?.descriptor as { resolution?: unknown } | undefined)
+            ?.resolution;
+        return resolution !== "cancelled" && resolution !== "expired";
+      });
+    if (
+      lastAssistant &&
+      hasSupersededApproval &&
+      !hasOpenToolPart &&
+      !(isRunning && lastAssistant === messages[lastIndex])
+    ) {
+      statusMessageIds.add(lastAssistant.id);
+    }
+    lastAssistant = undefined;
+    hasSupersededApproval = false;
+    previousWasVoice = false;
+  };
+
+  for (let index = 0; index < messages.length; index++) {
+    const message = messages[index]!;
+    if (message.role !== "assistant") {
+      flush();
+    } else {
+      const isVoice =
+        (message.metadata as { modality?: unknown } | undefined)?.modality ===
+        "voice";
+      if (isVoice || previousWasVoice || joinStrategy === "none") flush();
+      lastAssistant = message;
+      previousWasVoice = isVoice;
+    }
+
+    for (const part of message.parts ?? []) {
+      if (!isToolUIPart(part) || part.state !== "approval-requested") continue;
+
+      const approval = part.approval;
+      if (!approval) continue;
+
+      const approvalId = approval.id;
+      if (hostApprovalIds.has(approvalId)) continue;
+
+      const resolution =
+        (approval as { resolution?: unknown }).resolution ??
+        (approval.descriptor as { resolution?: unknown } | null | undefined)
+          ?.resolution;
+      if (resolution === "cancelled" || resolution === "expired") continue;
+
+      if (index !== lastIndex) {
+        approvalIds.add(approvalId);
+        if (message.role === "assistant") hasSupersededApproval = true;
+      }
+    }
+  }
+  flush();
+
+  if (approvalIds.size === 0 && statusMessageIds.size === 0)
+    return NO_SUPERSEDED_APPROVAL_PROJECTION;
+
+  return { approvalIds, statusMessageIds };
+};
+
+const findRawToolMessageIndex = <UI_MESSAGE extends UIMessage>(
+  messages: readonly UI_MESSAGE[],
+  messageId: string,
+  toolCallId: string,
+  joinStrategy: JoinStrategy | undefined,
+) => {
+  const containsToolCall = (message: UI_MESSAGE) =>
+    message.parts?.some(
+      (part) => isToolUIPart(part) && part.toolCallId === toolCallId,
+    ) === true;
+
+  const messageIndex = messages.findIndex(
+    (message) => message.id === messageId,
+  );
+  if (messageIndex === -1) return -1;
+  if (containsToolCall(messages[messageIndex]!)) return messageIndex;
+  if (joinStrategy === "none" || messages[messageIndex]?.role !== "assistant")
+    return -1;
+
+  let start = messageIndex;
+  while (start > 0 && messages[start - 1]?.role === "assistant") start--;
+
+  let end = messageIndex;
+  while (end + 1 < messages.length && messages[end + 1]?.role === "assistant")
+    end++;
+
+  for (let index = start; index <= end; index++) {
+    if (containsToolCall(messages[index]!)) return index;
+  }
+  return -1;
+};
+
 const toChatError = (error: Error): AssistantError => {
   const code = (error as { code?: unknown }).code;
   return {
@@ -303,15 +488,70 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     chatId: string;
     ids: ReadonlySet<string>;
   } | null>(null);
+  // 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 [toolApprovalResponses, setToolApprovalResponses] = useState<
     ReadonlyMap<string, RespondToToolApprovalOptions>
-  >(NO_TOOL_APPROVAL_RESPONSES);
+  >(() => toApprovalResponses(ownedApprovals));
   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 toolApprovalResponsesRef = useRef(
+    new Map(toApprovalResponses(ownedApprovals)),
+  );
+  const ownedApprovalIdsRef = useRef(new Set(ownedApprovals?.keys()));
+  const hostApprovalIdsRef = useRef(new Set<string>(ownedApprovals?.keys()));
+
+  // The owner's record is shared, so this runtime re-reads it whenever it is
+  // written rather than only at mount: the write may come from a runtime that
+  // has since unmounted, or from another runtime mounted over the same owner.
+  useEffect(() => {
+    if (!approvalOwner || !ownedApprovals) return undefined;
+    const sync = () => {
+      for (const id of ownedApprovalIdsRef.current)
+        toolApprovalResponsesRef.current.delete(id);
+      for (const [id, entry] of ownedApprovals)
+        toolApprovalResponsesRef.current.set(id, entry.response);
+      ownedApprovalIdsRef.current = new Set(ownedApprovals.keys());
+      hostApprovalIdsRef.current = new Set(
+        toolApprovalResponsesRef.current.keys(),
+      );
+      setToolApprovalResponses((prev) => {
+        const next = new Map(toolApprovalResponsesRef.current);
+        const unchanged =
+          prev.size === next.size &&
+          [...next].every(([id, response]) => prev.get(id) === response);
+        return unchanged ? prev : next;
+      });
+    };
+    const unsubscribe = subscribeToHostApprovals(approvalOwner, sync);
+    // The state was seeded during render, so a write landing between then and
+    // this subscription would otherwise never be seen.
+    sync();
+    return unsubscribe;
+  }, [approvalOwner, ownedApprovals]);
+
+  // A runtime kept mounted across a change of owner must not carry the
+  // previous chat's answers: a reused approval id would render as already
+  // answered and reject a genuine response.
+  const lastApprovalOwnerRef = useRef(approvalOwner);
+  if (lastApprovalOwnerRef.current !== approvalOwner) {
+    lastApprovalOwnerRef.current = approvalOwner;
+    hostApprovalIdsRef.current = new Set<string>(ownedApprovals?.keys());
+    toolApprovalResponsesRef.current = new Map(
+      toApprovalResponses(ownedApprovals),
+    );
+    ownedApprovalIdsRef.current = new Set(ownedApprovals?.keys());
+    setToolApprovalResponses(new Map(toolApprovalResponsesRef.current));
+  }
   const toolArgsKeyOrderCacheRef = useRef<Map<string, Map<string, string[]>>>(
     new Map(),
   );
@@ -356,6 +596,36 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
       : NO_CANCELLED_MESSAGE_IDS;
   const supportsRichToolApprovalResponses =
     customOnRespondToToolApproval != null;
+  const supersededApprovalProjectionRef = useRef(
+    NO_SUPERSEDED_APPROVAL_PROJECTION,
+  );
+  const supersededApprovalProjection = useMemo(
+    () => {
+      const projection = getSupersededApprovalProjection(
+        chatHelpers.messages,
+        hostApprovalIdsRef.current,
+        joinStrategy,
+        isRunning,
+      );
+      const previous = supersededApprovalProjectionRef.current;
+      if (
+        projection.approvalIds.size === previous.approvalIds.size &&
+        projection.statusMessageIds.size === previous.statusMessageIds.size &&
+        [...projection.approvalIds].every((id) =>
+          previous.approvalIds.has(id),
+        ) &&
+        [...projection.statusMessageIds].every((id) =>
+          previous.statusMessageIds.has(id),
+        )
+      )
+        return previous;
+
+      supersededApprovalProjectionRef.current = projection;
+      return projection;
+    },
+    // oxlint-disable-next-line react/exhaustive-deps -- hostApprovalIdsRef changes alongside toolApprovalResponses, which invalidates the projection
+    [chatHelpers.messages, joinStrategy, isRunning, toolApprovalResponses],
+  );
 
   const toThreadMessages = useCallback(
     (sourceMessages: UI_MESSAGE[]) => {
@@ -420,6 +690,9 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
         toolArtifacts: toolArtifactsRef.current,
         toolInteractions: toolInteractionsRef.current,
         supportsRichToolApprovalResponses,
+        cancelledToolApprovalIds: supersededApprovalProjection.approvalIds,
+        cancelledStatusMessageIds:
+          supersededApprovalProjection.statusMessageIds,
         ...(optimisticMessageId && { optimisticMessageId }),
         ...(chatHelpers.error && {
           error: toChatError(chatHelpers.error),
@@ -435,6 +708,7 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
         cancelledMessageIds,
         toolApprovalResponses,
         supportsRichToolApprovalResponses,
+        supersededApprovalProjection,
         toolArtifactEpoch,
         toolInteractionEpoch,
       ],
@@ -484,6 +758,9 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     markToolInteractionsChanged,
     toolApprovalResponsesRef.current,
     () => {
+      for (const [id, entry] of ownedApprovals ?? []) {
+        toolApprovalResponsesRef.current.set(id, entry.response);
+      }
       hostApprovalIdsRef.current = new Set(
         toolApprovalResponsesRef.current.keys(),
       );
@@ -579,7 +856,7 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
       .find(
         ({ part }) =>
           part.state === "approval-requested" &&
-          part.approval.id === approvalId,
+          part.approval?.id === approvalId,
       );
     if (!requested || hostApprovalIdsRef.current.has(approvalId))
       throw new Error(
@@ -587,7 +864,35 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
       );
 
     // A host answer stays out of the useChat messages, where sendAutomaticallyWhen would forward it to the chat route.
+    // The owner can change while a response is in flight, and the id set is a
+    // ref that is reseeded when it does. Both writes are therefore scoped to
+    // the record this response started under, so a rollback never reaches a
+    // different chat's state.
+    const startedWith = ownedApprovals;
+    const startedOwner = approvalOwner;
+    const startedResponses = toolApprovalResponsesRef.current;
+    // Whether this response is currently applied, tracked here rather than
+    // read back from the id ref: that ref follows the chat on screen and is
+    // reseeded when the owner changes, so it cannot answer for this response.
+    let isApplied = false;
     const applyResponse = (applied: boolean) => {
+      isApplied = applied;
+      // The captured record is always corrected, so a rollback reaches the
+      // chat the response belongs to even after the owner moved on.
+      if (applied)
+        startedWith?.set(approvalId, {
+          response,
+        });
+      else startedWith?.delete(approvalId);
+
+      if (startedOwner) {
+        // Every runtime mounted over that owner re-reads the record, including
+        // one mounted after this response started.
+        notifyHostApprovals(startedOwner);
+        return;
+      }
+
+      if (lastApprovalOwnerRef.current !== startedOwner) return;
       if (applied) hostApprovalIdsRef.current.add(approvalId);
       else hostApprovalIdsRef.current.delete(approvalId);
       if (applied) toolApprovalResponsesRef.current.set(approvalId, response);
@@ -609,10 +914,12 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
         },
       });
     } catch (error) {
-      if (hostApprovalIdsRef.current.has(approvalId)) applyResponse(false);
+      if (isApplied) applyResponse(false);
       throw error;
     }
-    if (hostApprovalIdsRef.current.has(approvalId)) {
+    const remainsApplied =
+      startedWith?.has(approvalId) ?? startedResponses.has(approvalId);
+    if (isApplied && remainsApplied) {
       await persistToolApprovalResponses(requested.messageId);
     }
   };
@@ -783,6 +1090,7 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
             toolInteractionsRef.current.delete(part.toolCallId) ||
             removedToolInteractions;
           if (part.approval) {
+            ownedApprovals?.delete(part.approval.id);
             removedToolApprovalResponse =
               toolApprovalResponsesRef.current.delete(part.approval.id) ||
               removedToolApprovalResponse;
@@ -795,6 +1103,7 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
       if (removedToolArtifact) markToolArtifactsChanged();
       if (removedToolInteractions) markToolInteractionsChanged();
       if (removedToolApprovalResponse || removedHostApprovalId) {
+        if (approvalOwner) notifyHostApprovals(approvalOwner);
         setToolApprovalResponses(new Map(toolApprovalResponsesRef.current));
       }
 
@@ -815,6 +1124,7 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
       await chatHelpers.regenerate({ metadata: config.runConfig });
     },
     onAddToolResult: ({
+      messageId,
       toolCallId,
       toolName,
       result,
@@ -826,6 +1136,62 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
         toolArtifactsRef.current.set(toolCallId, artifact);
         markToolArtifactsChanged();
       }
+
+      const targetIndex = findRawToolMessageIndex(
+        chatHelpers.messages,
+        messageId,
+        toolCallId,
+        joinStrategy,
+      );
+
+      const errorText =
+        typeof result === "string" ? result : JSON.stringify(result);
+      const output =
+        !isError && modelContent !== undefined
+          ? wrapModelContentEnvelope(result, modelContent)
+          : result;
+
+      if (targetIndex >= 0 && targetIndex !== chatHelpers.messages.length - 1) {
+        const target = chatHelpers.messages[targetIndex]!;
+        const targetPart = target.parts.find(
+          (part) => 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 (
+          targetPart?.state === "output-error" ||
+          targetPart?.state === "output-denied" ||
+          (targetPart?.state === "output-available" && !targetPart.preliminary)
+        )
+          return Promise.resolve();
+
+        const targetMessageId = target.id;
+        chatHelpers.setMessages((current) =>
+          current.map((message) => {
+            if (message.id !== targetMessageId) return message;
+
+            return {
+              ...message,
+              parts: message.parts.map((part) => {
+                if (!isToolUIPart(part) || part.toolCallId !== toolCallId)
+                  return part;
+
+                const { preliminary: _preliminary, ...finalPart } =
+                  part as typeof part & { preliminary?: boolean };
+                return {
+                  ...finalPart,
+                  state: isError
+                    ? ("output-error" as const)
+                    : ("output-available" as const),
+                  output: isError ? undefined : output,
+                  errorText: isError ? errorText : undefined,
+                } as typeof part;
+              }),
+            };
+          }),
+        );
+        return Promise.resolve();
+      }
+
       const options = { metadata: lastRunConfigRef.current };
       if (isError) {
         return Promise.resolve(
@@ -833,29 +1199,32 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
             state: "output-error",
             tool: toolName ?? toolCallId,
             toolCallId,
-            errorText:
-              typeof result === "string" ? result : JSON.stringify(result),
-            options,
-          }),
-        );
-      } else {
-        const output =
-          modelContent !== undefined
-            ? wrapModelContentEnvelope(result, modelContent)
-            : result;
-        return Promise.resolve(
-          chatHelpers.addToolOutput({
-            tool: toolName,
-            toolCallId,
-            output,
+            errorText,
             options,
           }),
         );
       }
+
+      return Promise.resolve(
+        chatHelpers.addToolOutput({
+          tool: toolName,
+          toolCallId,
+          output,
+          options,
+        }),
+      );
+    },
+    onRespondToToolApproval: (response) => {
+      if (supersededApprovalProjection.approvalIds.has(response.approvalId))
+        return Promise.reject(
+          new Error(
+            `Tool approval ${response.approvalId} is not waiting for a response.`,
+          ),
+        );
+      return customOnRespondToToolApproval
+        ? respondViaHost(customOnRespondToToolApproval, response)
+        : respondViaAISDK(response);
     },
-    onRespondToToolApproval: customOnRespondToToolApproval
-      ? (response) => respondViaHost(customOnRespondToToolApproval, response)
-      : respondViaAISDK,
     unstable_onRecordToolInteraction: ({
       messageId,
       toolCallId,

package/src/runtime/useChatRuntime.fast-refresh.test.tsx

@@ -0,0 +1,94 @@
+// @vitest-environment jsdom
+
+import { act, type ReactNode } from "react";
+import { afterAll, afterEach, expect, it, vi } from "vitest";
+
+type Family = { current: unknown };
+type RendererInternals = {
+  setRefreshHandler: (resolve: (type: unknown) => Family | undefined) => void;
+  scheduleRefresh: (
+    root: unknown,
+    update: { staleFamilies: Set<Family>; updatedFamilies: Set<Family> },
+  ) => void;
+};
+
+// Fast Refresh drives React through the renderer internals handed to the DevTools hook, so the hook has to exist before react-dom loads.
+let renderer: RendererInternals | undefined;
+const fiberRoots = new Set<unknown>();
+vi.stubGlobal("IS_REACT_ACT_ENVIRONMENT", true);
+vi.stubGlobal("__REACT_DEVTOOLS_GLOBAL_HOOK__", {
+  supportsFiber: true,
+  inject: (internals: RendererInternals) => {
+    renderer = internals;
+    return 1;
+  },
+  onScheduleFiberRoot: () => {},
+  onCommitFiberRoot: (_id: number, root: unknown) => fiberRoots.add(root),
+  onCommitFiberUnmount: () => {},
+});
+const { cleanup, render, waitFor } = await import("@testing-library/react");
+const { AssistantRuntimeProvider } = await import("@assistant-ui/core/react");
+const { useChatRuntime } = await import("./useChatRuntime");
+const { createCancellableTransport, createStreamHarness } =
+  await import("./__tests__/controlled-transport");
+
+afterEach(cleanup);
+afterAll(() => vi.unstubAllGlobals());
+
+it("keeps a running thread across a Fast Refresh of its host and aborts it on unmount", async () => {
+  const { transport, getCancelCount } = createCancellableTransport();
+  const { Probe, send, isRunning, client } = createStreamHarness();
+  const runtimes: unknown[] = [];
+  let rendered: string | undefined;
+  const createHost =
+    (name: string) =>
+    ({ children }: { children: ReactNode }) => {
+      rendered = name;
+      const runtime = useChatRuntime({ transport });
+      runtimes.push(runtime);
+      return (
+        <AssistantRuntimeProvider runtime={runtime}>
+          {children}
+        </AssistantRuntimeProvider>
+      );
+    };
+  const Before = createHost("before");
+  const After = createHost("after");
+
+  const view = render(
+    <Before>
+      <Probe />
+    </Before>,
+  );
+  await act(async () => send());
+  await waitFor(() => expect(isRunning()).toBe(true));
+  const runtime = runtimes.at(-1);
+  const threadId = client().threads.getState().mainThreadId;
+  const messages = client().thread.getState().messages;
+  const error = vi.spyOn(console, "error");
+
+  const family: Family = { current: After };
+  renderer!.setRefreshHandler((type) =>
+    type === Before || type === After ? family : undefined,
+  );
+  await act(async () => {
+    for (const fiberRoot of fiberRoots) {
+      renderer!.scheduleRefresh(fiberRoot, {
+        staleFamilies: new Set(),
+        updatedFamilies: new Set([family]),
+      });
+    }
+  });
+  await act(async () => {});
+
+  expect(rendered).toBe("after");
+  expect(runtimes.at(-1)).toBe(runtime);
+  expect(isRunning()).toBe(true);
+  expect(getCancelCount()).toBe(0);
+  expect(client().threads.getState().mainThreadId).toBe(threadId);
+  expect(client().thread.getState().messages).toEqual(messages);
+  expect(error).not.toHaveBeenCalled();
+
+  view.unmount();
+  await waitFor(() => expect(getCancelCount()).toBe(1));
+});