@assistant-ui/react 0.15.24 → 0.15.26

package/src/primitives/thread/useThreadViewportAutoScroll.ts

@@ -1,8 +1,8 @@
 "use client";
 
 import { useComposedRefs } from "radix-ui/internal";
-import { useCallback, useLayoutEffect, useRef, type RefCallback } from "react";
-import { useAuiEvent, useAuiState } from "@assistant-ui/store";
+import { useCallback, useRef, type RefCallback } from "react";
+import { useAui, useAuiEvent, useAuiState } from "@assistant-ui/store";
 import {
   isUserScrollUp,
   isViewportAtBottom,
@@ -13,12 +13,21 @@ import { useOnScrollToBottom } from "../../utils/hooks/useOnScrollToBottom";
 import { useManagedRef } from "../../utils/hooks/useManagedRef";
 import { writableStore } from "../../context/ReadonlyStore";
 import { useThreadViewportStore } from "../../context/react/ThreadViewportContext";
-
-// Enter and Space activate a focused control, which is how a collapsible tool
-// call expands without a pointer event ever firing. No other key changes thread
-// content: the keys that scroll the viewport reach handleScroll, which already
-// clears the intent when the user scrolls up.
-const ACTIVATION_KEYS = new Set(["Enter", " "]);
+import { useIsomorphicLayoutEffect } from "../../utils/useIsomorphicLayoutEffect";
+
+const SCROLL_KEYS = new Set([
+  " ",
+  "ArrowUp",
+  "ArrowDown",
+  "PageUp",
+  "PageDown",
+  "Home",
+  "End",
+]);
+
+// Enter and Space activate focused controls, while navigation keys can
+// interrupt scrolling without producing a scroll event.
+const INTENT_CANCEL_KEYS = new Set(["Enter", ...SCROLL_KEYS]);
 
 // A control that consumes the activation key itself instead of acting on thread
 // content. `contenteditable="false"` marks a non-editable island inside an
@@ -33,6 +42,85 @@ const TEXT_ENTRY_SELECTOR = [
     ":not([type='range']):not([type='file']):not([type='color'])",
 ].join(", ");
 
+const MESSAGE_SELECTOR = "[data-message-id]";
+
+const DISCLOSURE_SELECTOR = "[aria-expanded], summary";
+
+// A popup trigger (menu, dialog, listbox) also carries aria-expanded, but what
+// it opens is portaled away from the thread, so only disclosures count.
+const isOpeningDisclosure = (target: Element): boolean => {
+  const control = target.closest(DISCLOSURE_SELECTOR);
+  if (!control) return false;
+  if (control.tagName === "SUMMARY")
+    return control.parentElement?.hasAttribute("open") === false;
+  return (
+    control.getAttribute("aria-expanded") === "false" &&
+    (control.getAttribute("aria-haspopup") ?? "false") === "false" &&
+    control.getAttribute("role") !== "combobox"
+  );
+};
+
+type MessageOffset = { readonly id: string; readonly offset: number };
+
+const offsetInViewport = (element: Element, viewport: Element) =>
+  element.getBoundingClientRect().top - viewport.getBoundingClientRect().top;
+
+const measureFirstMessage = (viewport: HTMLElement): MessageOffset | null => {
+  const element = viewport.querySelector<HTMLElement>(MESSAGE_SELECTOR);
+  const id = element?.dataset["messageId"];
+  return element && id
+    ? { id, offset: offsetInViewport(element, viewport) }
+    : null;
+};
+
+/** The first message that starts in view, else the one spanning its top edge. */
+const measureReaderMessage = (viewport: HTMLElement): MessageOffset | null => {
+  const { top, bottom } = viewport.getBoundingClientRect();
+  let spanning: MessageOffset | null = null;
+  for (const element of viewport.querySelectorAll<HTMLElement>(
+    MESSAGE_SELECTOR,
+  )) {
+    const rect = element.getBoundingClientRect();
+    if (rect.bottom <= top) continue;
+    if (rect.top >= bottom) break;
+    const id = element.dataset["messageId"];
+    if (!id) continue;
+    if (rect.top >= top) return { id, offset: rect.top - top };
+    spanning ??= { id, offset: rect.top - top };
+  }
+  return spanning;
+};
+
+const findMessage = (viewport: HTMLElement, id: string) => {
+  for (const element of viewport.querySelectorAll<HTMLElement>(
+    MESSAGE_SELECTOR,
+  )) {
+    if (element.dataset["messageId"] === id) return element;
+  }
+  return null;
+};
+
+/**
+ * Scrolls so the message sits at `anchor.offset` again, recording the
+ * position it scrolls to in `heldScrollTop` before the scroll event can
+ * report it; false when the message is gone.
+ */
+const keepMessageAt = (
+  viewport: HTMLElement,
+  anchor: MessageOffset,
+  heldScrollTop: { current: number },
+) => {
+  const element = findMessage(viewport, anchor.id);
+  if (!element) return false;
+  const delta = offsetInViewport(element, viewport) - anchor.offset;
+  heldScrollTop.current = viewport.scrollTop + delta;
+  if (Math.abs(delta) >= 1) {
+    viewport.scrollTo({ top: heldScrollTop.current, behavior: "instant" });
+  }
+  heldScrollTop.current = viewport.scrollTop;
+  return true;
+};
+
 export namespace useThreadViewportAutoScroll {
   export type Options = {
     /**
@@ -72,8 +160,11 @@ export const useThreadViewportAutoScroll = <TElement extends HTMLElement>({
   scrollToBottomOnInitialize = true,
   scrollToBottomOnThreadSwitch = true,
 }: useThreadViewportAutoScroll.Options): RefCallback<TElement> => {
+  const aui = useAui();
   const divRef = useRef<TElement>(null);
   const hasMessages = useAuiState((s) => s.thread.messages.length > 0);
+  const firstMessageId = useAuiState((s) => s.thread.messages[0]?.id);
+  const threadId = useAuiState((s) => s.threadListItem.id);
   const isRunning = useAuiState((s) => s.thread.isRunning);
   const initializeScrollRequestedRef = useRef(false);
   const scheduledFrameRef = useRef<number | null>(null);
@@ -87,15 +178,28 @@ export const useThreadViewportAutoScroll = <TElement extends HTMLElement>({
   const lastScrollHeight = useRef<number>(0);
   const lastObservedScrollHeight = useRef<number>(0);
   const lastObservedClientHeight = useRef<number>(0);
+  const firstMessageRef = useRef<MessageOffset | null>(null);
+  // The row the reader sees when earlier messages land, held until a gesture
+  // or a scroll this hook did not make, so content above it that settles
+  // afterwards (lazy layout, late media) does not move it on engines without
+  // native scroll anchoring.
+  const prependAnchorRef = useRef<MessageOffset | null>(null);
+  const heldScrollTopRef = useRef(0);
 
   // Pending bottom-scroll intent. Planted by initialize/run-start/switch/button
   // triggers, cleared when handleScroll confirms we reached bottom, or when the
   // user actively scrolls up while content size is stable.
   const scrollingToBottomBehaviorRef = useRef<ScrollBehavior | null>(null);
   const followBottomRef = useRef(autoScroll);
+  // Set by expanding a disclosure; suspends follow without clearing its intent
+  // until a scroll gesture reaches an overflowing bottom (content that fits
+  // reads as at the bottom), a run starts, the thread changes, or the reader
+  // asks for the bottom again.
+  const followPausedRef = useRef(false);
+  const scrolledSincePauseRef = useRef(false);
   const previousAutoScrollRef = useRef(autoScroll);
 
-  useLayoutEffect(() => {
+  useIsomorphicLayoutEffect(() => {
     const previousAutoScroll = previousAutoScrollRef.current;
     previousAutoScrollRef.current = autoScroll;
     if (previousAutoScroll || !autoScroll) return;
@@ -109,8 +213,12 @@ export const useThreadViewportAutoScroll = <TElement extends HTMLElement>({
     if (!div) return;
 
     followBottomRef.current = true;
+    followPausedRef.current = false;
     scrollingToBottomBehaviorRef.current = behavior;
+    prependAnchorRef.current = null;
     div.scrollTo({ top: div.scrollHeight, behavior });
+    lastScrollTop.current = div.scrollTop;
+    lastScrollHeight.current = div.scrollHeight;
   }, []);
 
   const cancelScheduledFrame = useCallback(() => {
@@ -131,7 +239,10 @@ export const useThreadViewportAutoScroll = <TElement extends HTMLElement>({
     [cancelScheduledFrame, scrollToBottom],
   );
 
-  useLayoutEffect(() => () => cancelScheduledFrame(), [cancelScheduledFrame]);
+  useIsomorphicLayoutEffect(
+    () => () => cancelScheduledFrame(),
+    [cancelScheduledFrame],
+  );
 
   const hasActiveTopAnchor = useCallback(() => {
     const state = threadViewportStore.getState();
@@ -142,15 +253,24 @@ export const useThreadViewportAutoScroll = <TElement extends HTMLElement>({
     );
   }, [threadViewportStore]);
 
-  const handleScroll = () => {
+  const onScroll = () => {
     const div = divRef.current;
     if (!div) return;
 
+    if (
+      prependAnchorRef.current &&
+      Math.abs(div.scrollTop - heldScrollTopRef.current) >= 1
+    ) {
+      prependAnchorRef.current = null;
+    }
+
     const isAtBottom = threadViewportStore.getState().isAtBottom;
     const newIsAtBottom = isViewportAtBottom(div);
 
     const isInFlightDownwardScroll =
-      !newIsAtBottom && lastScrollTop.current < div.scrollTop;
+      scrollingToBottomBehaviorRef.current !== null &&
+      !newIsAtBottom &&
+      lastScrollTop.current < div.scrollTop;
     if (isInFlightDownwardScroll) {
       // no-op: a smooth scroll-to-bottom fires many midpoint scroll events
       // before landing, don't flicker isAtBottom or clear intent mid-animation
@@ -168,6 +288,7 @@ export const useThreadViewportAutoScroll = <TElement extends HTMLElement>({
         // keep intent alive until content can actually scroll
         if (viewportOverflows(div)) {
           scrollingToBottomBehaviorRef.current = null;
+          if (scrolledSincePauseRef.current) followPausedRef.current = false;
         }
         if (autoScroll) followBottomRef.current = true;
       } else if (userScrolledUp) {
@@ -188,9 +309,15 @@ export const useThreadViewportAutoScroll = <TElement extends HTMLElement>({
 
     lastScrollTop.current = div.scrollTop;
     lastScrollHeight.current = div.scrollHeight;
+    firstMessageRef.current = measureFirstMessage(div);
   };
+  const onScrollRef = useRef(onScroll);
+  useIsomorphicLayoutEffect(() => {
+    onScrollRef.current = onScroll;
+  });
+  const handleScroll = useCallback(() => onScrollRef.current(), []);
 
-  const resizeRef = useOnResizeContent(() => {
+  const onResize = () => {
     const div = divRef.current;
     if (!div) return;
 
@@ -204,55 +331,199 @@ export const useThreadViewportAutoScroll = <TElement extends HTMLElement>({
     lastObservedScrollHeight.current = scrollHeight;
     lastObservedClientHeight.current = clientHeight;
 
-    const scrollBehavior = scrollingToBottomBehaviorRef.current;
-    if (scrollBehavior && hasActiveTopAnchor()) {
-      // Let the top-anchor reserve own scrolling while a run starts to avoid a bottom-scroll race.
-      scrollingToBottomBehaviorRef.current = null;
-    } else if (scrollBehavior) {
-      scrollToBottom(scrollBehavior);
-    } else if (
-      autoScroll &&
-      !(isRunning && hasActiveTopAnchor()) &&
-      followBottomRef.current
-    ) {
-      scrollToBottom("instant");
+    // A scheduled frame owns the pending behavior; resize callbacks must not apply it early.
+    if (scheduledFrameRef.current === null) {
+      const scrollBehavior = scrollingToBottomBehaviorRef.current;
+      if (scrollBehavior && hasActiveTopAnchor()) {
+        // Let the top-anchor reserve own scrolling while a run starts to avoid a bottom-scroll race.
+        scrollingToBottomBehaviorRef.current = null;
+      } else if (scrollBehavior) {
+        scrollToBottom(scrollBehavior);
+      } else if (
+        autoScroll &&
+        !(isRunning && hasActiveTopAnchor()) &&
+        followBottomRef.current &&
+        !followPausedRef.current
+      ) {
+        scrollToBottom("instant");
+      } else if (
+        prependAnchorRef.current &&
+        !keepMessageAt(div, prependAnchorRef.current, heldScrollTopRef)
+      ) {
+        prependAnchorRef.current = null;
+      }
     }
 
     handleScroll();
+  };
+  const onResizeRef = useRef(onResize);
+  useIsomorphicLayoutEffect(() => {
+    onResizeRef.current = onResize;
   });
+  const resizeRef = useOnResizeContent(() => onResizeRef.current());
+
+  const scrollRef = useManagedRef<HTMLElement>(
+    useCallback(
+      (el) => {
+        // A user gesture invalidates pending bottom-scroll intent; otherwise an
+        // intent kept alive by a non-overflowing thread (see handleScroll) hijacks
+        // the next content growth, e.g. expanding a collapsible tool call. Keyboard
+        // activation reaches that same content without ever emitting a pointer
+        // event, so it has to cancel the intent too.
+        const cancelPendingScrollToBottom = () => {
+          // With nothing pending, the last scroll event already published isAtBottom.
+          if (
+            scrollingToBottomBehaviorRef.current === null &&
+            scheduledFrameRef.current === null
+          )
+            return;
+          // A scheduled frame re-plants the intent when it runs, so clearing the
+          // ref alone leaves the gesture undone.
+          cancelScheduledFrame();
+          scrollingToBottomBehaviorRef.current = null;
+          handleScroll();
+        };
+        // The composer renders inside the viewport, so its keystrokes bubble here;
+        // cancellation keys only represent a gesture on thread content when they
+        // originate outside text entry.
+        const cancelOnKeyDown = (event: KeyboardEvent) => {
+          if (!INTENT_CANCEL_KEYS.has(event.key)) return;
+          const target = event.target as Element | null;
+          if (target?.closest?.(TEXT_ENTRY_SELECTOR)) return;
+          cancelPendingScrollToBottom();
+        };
+        // Expanding a disclosure inside a message, such as a reasoning block or a
+        // tool call card, means the reader is inspecting it, so bottom follow
+        // pauses until the reader scrolls or asks for the bottom again. Every
+        // other interaction, including copying or selecting text, keeps
+        // following. Activation is read from click, which keyboard activation also
+        // dispatches, in the capture phase so the state is read before any handler
+        // on the disclosure flips it or stops the event.
+        const pauseFollowOnExpand = (event: MouseEvent) => {
+          const target = event.target as Element | null;
+          if (
+            !target?.closest?.(MESSAGE_SELECTOR) ||
+            target.closest(TEXT_ENTRY_SELECTOR) ||
+            !isOpeningDisclosure(target)
+          )
+            return;
+          followPausedRef.current = true;
+          scrolledSincePauseRef.current = false;
+          cancelPendingScrollToBottom();
+        };
+        const noteScrollGesture = (event: Event) => {
+          if (event instanceof KeyboardEvent) {
+            if (!SCROLL_KEYS.has(event.key)) return;
+            if (
+              (event.target as Element | null)?.closest?.(TEXT_ENTRY_SELECTOR)
+            )
+              return;
+          } else if (event.type === "pointerdown" && event.target !== el) {
+            return;
+          }
+          scrolledSincePauseRef.current = true;
+        };
+        const releasePrependAnchor = () => {
+          prependAnchorRef.current = null;
+        };
+        const gestures = [
+          "pointerdown",
+          "wheel",
+          "touchstart",
+          "keydown",
+        ] as const;
+        const resumeGestures = [
+          "pointerdown",
+          "wheel",
+          "touchmove",
+          "keydown",
+        ] as const;
+        el.addEventListener("scroll", handleScroll);
+        el.addEventListener("pointerdown", cancelPendingScrollToBottom);
+        el.addEventListener("wheel", cancelPendingScrollToBottom, {
+          passive: true,
+        });
+        el.addEventListener("touchstart", cancelPendingScrollToBottom, {
+          passive: true,
+        });
+        el.addEventListener("keydown", cancelOnKeyDown);
+        el.addEventListener("click", pauseFollowOnExpand, { capture: true });
+        for (const gesture of resumeGestures) {
+          el.addEventListener(gesture, noteScrollGesture, { passive: true });
+        }
+        for (const gesture of gestures) {
+          el.addEventListener(gesture, releasePrependAnchor, { passive: true });
+        }
+        return () => {
+          el.removeEventListener("scroll", handleScroll);
+          el.removeEventListener("pointerdown", cancelPendingScrollToBottom);
+          el.removeEventListener("wheel", cancelPendingScrollToBottom);
+          el.removeEventListener("touchstart", cancelPendingScrollToBottom);
+          el.removeEventListener("keydown", cancelOnKeyDown);
+          el.removeEventListener("click", pauseFollowOnExpand, {
+            capture: true,
+          });
+          for (const gesture of resumeGestures) {
+            el.removeEventListener(gesture, noteScrollGesture);
+          }
+          for (const gesture of gestures) {
+            el.removeEventListener(gesture, releasePrependAnchor);
+          }
+        };
+      },
+      [cancelScheduledFrame, handleScroll],
+    ),
+  );
 
-  const scrollRef = useManagedRef<HTMLElement>((el) => {
-    // A user gesture invalidates pending bottom-scroll intent; otherwise an
-    // intent kept alive by a non-overflowing thread (see handleScroll) hijacks
-    // the next content growth, e.g. expanding a collapsible tool call. Keyboard
-    // activation reaches that same content without ever emitting a pointer
-    // event, so it has to cancel the intent too.
-    const cancelPendingScrollToBottom = () => {
-      // A scheduled frame re-plants the intent when it runs, so clearing the
-      // ref alone leaves the gesture undone.
-      cancelScheduledFrame();
-      scrollingToBottomBehaviorRef.current = null;
-    };
-    // The composer renders inside the viewport, so its keystrokes bubble here;
-    // only an activation key aimed at something other than a text field is a
-    // gesture on thread content.
-    const cancelOnKeyDown = (event: KeyboardEvent) => {
-      if (!ACTIVATION_KEYS.has(event.key)) return;
-      const target = event.target as Element | null;
-      if (target?.closest?.(TEXT_ENTRY_SELECTOR)) return;
-      cancelPendingScrollToBottom();
-    };
-    el.addEventListener("scroll", handleScroll);
-    el.addEventListener("pointerdown", cancelPendingScrollToBottom);
-    el.addEventListener("keydown", cancelOnKeyDown);
-    return () => {
-      el.removeEventListener("scroll", handleScroll);
-      el.removeEventListener("pointerdown", cancelPendingScrollToBottom);
-      el.removeEventListener("keydown", cancelOnKeyDown);
-    };
-  });
+  // Earlier messages prepended above the reader put the old first message back
+  // where the reader saw it, so growth below it in the same commit (a
+  // streaming reply) is not mistaken for the page; where the browser already
+  // anchored the scroll, this is a no-op.
+  const previousRef = useRef({ threadId, firstMessageId });
+  useIsomorphicLayoutEffect(() => {
+    const previous = previousRef.current;
+    previousRef.current = { threadId, firstMessageId };
+    if (previous.threadId !== threadId) {
+      prependAnchorRef.current = null;
+      return;
+    }
+    const div = divRef.current;
+    const previousFirstMessageId = previous.firstMessageId;
+    if (
+      !div ||
+      previousFirstMessageId === undefined ||
+      previousFirstMessageId === firstMessageId ||
+      (autoScroll && followBottomRef.current && !followPausedRef.current) ||
+      scrollingToBottomBehaviorRef.current !== null
+    )
+      return;
+    const prepended = aui.thread
+      .getState()
+      .messages.some(
+        (message, index) => index > 0 && message.id === previousFirstMessageId,
+      );
+    if (!prepended) return;
+
+    const anchor = firstMessageRef.current;
+    if (
+      anchor?.id === previousFirstMessageId &&
+      keepMessageAt(div, anchor, heldScrollTopRef)
+    ) {
+      prependAnchorRef.current = measureReaderMessage(div) ?? anchor;
+    } else {
+      // Message roots that render no `data-message-id` leave only the total
+      // growth to go on, which counts any growth below the reader too.
+      const top =
+        div.scrollHeight - (lastScrollHeight.current - lastScrollTop.current);
+      if (Math.abs(div.scrollTop - top) >= 1) {
+        div.scrollTo({ top, behavior: "instant" });
+      }
+    }
+    lastScrollTop.current = div.scrollTop;
+    lastScrollHeight.current = div.scrollHeight;
+  }, [aui, autoScroll, firstMessageId, threadId]);
 
-  useLayoutEffect(() => {
+  useIsomorphicLayoutEffect(() => {
     if (!scrollToBottomOnInitialize) return;
     if (!hasMessages) {
       initializeScrollRequestedRef.current = false;
@@ -272,12 +543,15 @@ export const useThreadViewportAutoScroll = <TElement extends HTMLElement>({
   });
 
   useAuiEvent("thread.runStart", () => {
+    prependAnchorRef.current = null;
+    followPausedRef.current = false;
     if (!scrollToBottomOnRunStart) return;
     if (threadViewportStore.getState().turnAnchor === "top") return;
     scheduleScrollToBottom("auto");
   });
 
   useAuiEvent("threads.selectionChanged", () => {
+    followPausedRef.current = false;
     if (!scrollToBottomOnThreadSwitch) return;
     scheduleScrollToBottom("instant");
   });

package/src/primitives/thread.ts

@@ -7,7 +7,9 @@ export { ThreadPrimitiveViewportFooter as ViewportFooter } from "./thread/Thread
 export { ThreadPrimitiveMessages as Messages } from "./thread/ThreadMessages";
 export { ThreadPrimitiveMessageByIndex as MessageByIndex } from "./thread/ThreadMessages";
 export { ThreadPrimitiveUnstable_MessageById as Unstable_MessageById } from "./thread/ThreadMessages";
+export { ThreadPrimitiveRow as Row } from "./thread/ThreadMessages";
 export { ThreadPrimitiveScrollToBottom as ScrollToBottom } from "./thread/ThreadScrollToBottom";
+export { ThreadPrimitiveLoadEarlier as LoadEarlier } from "./thread/ThreadLoadEarlier";
 export { ThreadPrimitiveSuggestion as Suggestion } from "./thread/ThreadSuggestion";
 export {
   ThreadPrimitiveSuggestions as Suggestions,

package/src/sandbox-host/SandboxHost.test.tsx

@@ -416,6 +416,40 @@ describe("SandboxHost", () => {
     expect(onError).not.toHaveBeenCalled();
   });
 
+  it("cancels pending renders on replacement and unmount", async () => {
+    renderHtmlMock.mockImplementation(() => new Promise(() => {}));
+
+    await act(async () => {
+      root.render(
+        <SandboxHost
+          content={{ html: "first" }}
+          contentKey="first"
+          createBridge={() => ({ onMessage: vi.fn(), dispose: vi.fn() })}
+        />,
+      );
+    });
+    const firstSignal = renderHtmlMock.mock.calls[0]![2].signal as AbortSignal;
+
+    await act(async () => {
+      root.render(
+        <SandboxHost
+          content={{ html: "second" }}
+          contentKey="second"
+          createBridge={() => ({ onMessage: vi.fn(), dispose: vi.fn() })}
+        />,
+      );
+    });
+    const secondSignal = renderHtmlMock.mock.calls[1]![2].signal as AbortSignal;
+
+    expect(firstSignal.aborted).toBe(true);
+    expect(secondSignal.aborted).toBe(false);
+
+    await act(async () => {
+      root.unmount();
+    });
+    expect(secondSignal.aborted).toBe(true);
+  });
+
   it("contains failures thrown by onError", async () => {
     const renderError = new Error("render failed");
     const callbackError = new Error("error callback failed");

package/src/sandbox-host/SandboxHost.tsx

@@ -1,12 +1,7 @@
 "use client";
 
-import {
-  type CSSProperties,
-  useEffect,
-  useLayoutEffect,
-  useRef,
-  useState,
-} from "react";
+import { type CSSProperties, useEffect, useRef, useState } from "react";
+import { useIsomorphicLayoutEffect } from "../utils/useIsomorphicLayoutEffect";
 import {
   isShimLoadError,
   type RenderedFrame,
@@ -95,7 +90,7 @@ export function SandboxHost({
     createBridge,
     onError,
   });
-  useLayoutEffect(() => {
+  useIsomorphicLayoutEffect(() => {
     liveRef.current = { content, sandbox, createBridge, onError };
   }, [content, sandbox, createBridge, onError]);
 
@@ -107,6 +102,7 @@ export function SandboxHost({
     let frame: RenderedFrame | null = null;
     let bridge: SandboxBridge | null = null;
     let onMessage: ((event: MessageEvent) => void) | null = null;
+    const renderController = new AbortController();
 
     const { content: liveContent, sandbox: sb } = liveRef.current;
 
@@ -129,10 +125,12 @@ export function SandboxHost({
       ...(sb?.salt !== undefined && { salt: sb.salt }),
     });
 
-    const renderOpts =
-      sb?.unsafeDocumentWrite !== undefined
-        ? { unsafeDocumentWrite: sb.unsafeDocumentWrite }
-        : undefined;
+    const renderOpts = {
+      signal: renderController.signal,
+      ...(sb?.unsafeDocumentWrite !== undefined && {
+        unsafeDocumentWrite: sb.unsafeDocumentWrite,
+      }),
+    };
 
     scf
       .renderHtml(liveContent.html, container, renderOpts)
@@ -216,6 +214,7 @@ export function SandboxHost({
       const bridgeToDispose = bridge;
       bridge = null;
       if (bridgeToDispose) runCleanup(() => bridgeToDispose.dispose());
+      runCleanup(() => renderController.abort());
       const frameToDispose = frame;
       frame = null;
       if (frameToDispose) runCleanup(() => frameToDispose.dispose());

package/src/tests/toolCallTiming.test.tsx

@@ -47,8 +47,15 @@ const ElapsedProbe: FC = () => {
 };
 
 const StallProbe: FC = () => {
-  const { stalled } = unstable_useMessageStallDetection({ thresholdMs: 2000 });
-  return <span data-testid="stalled">{String(stalled)}</span>;
+  const { stalled, stalledForMs } = unstable_useMessageStallDetection({
+    thresholdMs: 2000,
+  });
+  return (
+    <>
+      <span data-testid="stalled">{String(stalled)}</span>
+      <span data-testid="stalled-for">{stalledForMs}</span>
+    </>
+  );
 };
 
 const RuntimeProvider: FC<
@@ -224,6 +231,16 @@ describe("unstable_useMessageStallDetection", () => {
     });
 
     expect(screen.getByTestId("stalled").textContent).toBe("true");
+    const stalledFor = () =>
+      Number(screen.getByTestId("stalled-for").textContent);
+    const atStall = stalledFor();
+    expect(atStall).toBeGreaterThanOrEqual(2000);
+
+    await act(async () => {
+      await vi.advanceTimersByTimeAsync(1000);
+    });
+
+    expect(stalledFor()).toBeGreaterThan(atStall);
   });
 
   it("never stalls on settled messages", async () => {