@assistant-ui/react 0.15.21 → 0.15.23

package/src/primitives/selectionToolbar/SelectionToolbarRoot.tsx

@@ -9,10 +9,12 @@ import {
   forwardRef,
   useContext,
   useEffect,
+  useRef,
   useState,
 } from "react";
 import { createPortal } from "react-dom";
 import { getSelectionMessageId } from "../../utils/getSelectionMessageId";
+import { useThreadRootElementRef } from "../thread/ThreadRootElementContext";
 
 type SelectionInfo = {
   text: string;
@@ -33,9 +35,9 @@ export namespace SelectionToolbarPrimitiveRoot {
 /**
  * A floating toolbar that appears when text is selected within a message.
  *
- * Listens for mouse and keyboard selection events, validates that the
- * selection is within a single message, and renders a positioned portal
- * near the selection. Prevents mousedown from clearing the selection.
+ * Listens for browser selection changes, validates that the selection is
+ * within a single message, and renders a positioned portal near the
+ * selection. Prevents mousedown from clearing the selection.
  *
  * @example
  * ```tsx
@@ -49,10 +51,13 @@ export const SelectionToolbarPrimitiveRoot = forwardRef<
   SelectionToolbarPrimitiveRoot.Props
 >(({ onMouseDown, style, ...props }, forwardedRef) => {
   const [info, setInfo] = useState<SelectionInfo | null>(null);
+  const threadRootRef = useThreadRootElementRef();
+  const warnedAboutMissingThreadRootRef = useRef(false);
 
   useEffect(() => {
     // Read the selection on the next frame so the browser has settled it.
     let pendingFrame: number | null = null;
+    let isMouseDragging = false;
 
     const checkSelection = () => {
       if (pendingFrame !== null) cancelAnimationFrame(pendingFrame);
@@ -70,7 +75,19 @@ export const SelectionToolbarPrimitiveRoot = forwardRef<
           return;
         }
 
-        const messageId = getSelectionMessageId(sel);
+        if (threadRootRef && !threadRootRef.current) {
+          if (
+            process.env.NODE_ENV !== "production" &&
+            !warnedAboutMissingThreadRootRef.current
+          ) {
+            warnedAboutMissingThreadRootRef.current = true;
+            console.warn(
+              "[SelectionToolbarPrimitive.Root] ThreadPrimitive.Root did not provide a DOM element, so the selection cannot be scoped to its thread. Ensure a custom root child forwards its ref.",
+            );
+          }
+        }
+
+        const messageId = getSelectionMessageId(sel, threadRootRef?.current);
         if (!messageId) {
           setInfo(null);
           return;
@@ -82,30 +99,61 @@ export const SelectionToolbarPrimitiveRoot = forwardRef<
       });
     };
 
-    const handleSelectionCollapse = () => {
+    const handleSelectionChange = () => {
       const sel = window.getSelection();
       if (!sel || sel.isCollapsed) {
+        if (pendingFrame !== null) {
+          cancelAnimationFrame(pendingFrame);
+          pendingFrame = null;
+        }
         setInfo(null);
+        return;
       }
+
+      if (!isMouseDragging) checkSelection();
     };
 
     const handleScroll = () => {
+      if (pendingFrame !== null) {
+        cancelAnimationFrame(pendingFrame);
+        pendingFrame = null;
+      }
       setInfo(null);
     };
 
-    document.addEventListener("mouseup", checkSelection);
-    document.addEventListener("keyup", checkSelection);
-    document.addEventListener("selectionchange", handleSelectionCollapse);
+    const handleMouseDown = () => {
+      isMouseDragging = true;
+    };
+
+    const handleMouseUp = () => {
+      isMouseDragging = false;
+      checkSelection();
+    };
+
+    const handleMouseCancel = () => {
+      isMouseDragging = false;
+    };
+
+    document.addEventListener("mousedown", handleMouseDown, true);
+    document.addEventListener("mouseup", handleMouseUp, true);
+    document.addEventListener("dragend", handleMouseUp, true);
+    // A context menu can consume the mouseup of the press that opened it.
+    document.addEventListener("contextmenu", handleMouseCancel, true);
+    window.addEventListener("blur", handleMouseCancel);
+    document.addEventListener("selectionchange", handleSelectionChange);
     document.addEventListener("scroll", handleScroll, true);
 
     return () => {
       if (pendingFrame !== null) cancelAnimationFrame(pendingFrame);
-      document.removeEventListener("mouseup", checkSelection);
-      document.removeEventListener("keyup", checkSelection);
-      document.removeEventListener("selectionchange", handleSelectionCollapse);
+      document.removeEventListener("mousedown", handleMouseDown, true);
+      document.removeEventListener("mouseup", handleMouseUp, true);
+      document.removeEventListener("dragend", handleMouseUp, true);
+      document.removeEventListener("contextmenu", handleMouseCancel, true);
+      window.removeEventListener("blur", handleMouseCancel);
+      document.removeEventListener("selectionchange", handleSelectionChange);
       document.removeEventListener("scroll", handleScroll, true);
     };
-  }, []);
+  }, [threadRootRef]);
 
   if (!info) return null;
 

package/src/primitives/thread/ThreadRoot.test.tsx

@@ -94,6 +94,7 @@ const dispatchEscape = (target: EventTarget) => {
     key: "Escape",
     bubbles: true,
     cancelable: true,
+    composed: true,
   });
   act(() => target.dispatchEvent(event));
   return event;
@@ -143,6 +144,135 @@ describe("ThreadPrimitiveRoot", () => {
     expect(event.defaultPrevented).toBe(false);
   });
 
+  it("does not stop speech in another mounted thread", async () => {
+    const firstSpeech = createSpeechAdapter();
+    const secondSpeech = createSpeechAdapter();
+    const firstRuntimeRef: RuntimeRef = { current: null };
+    const secondRuntimeRef: RuntimeRef = { current: null };
+    render(
+      <>
+        <RuntimeProvider
+          runtimeRef={firstRuntimeRef}
+          speech={firstSpeech.adapter}
+        >
+          <ThreadPrimitiveRoot>
+            <button data-testid="first-thread-control" />
+          </ThreadPrimitiveRoot>
+        </RuntimeProvider>
+        <RuntimeProvider
+          runtimeRef={secondRuntimeRef}
+          speech={secondSpeech.adapter}
+        >
+          <ThreadPrimitiveRoot>
+            <button data-testid="second-thread-control" />
+          </ThreadPrimitiveRoot>
+        </RuntimeProvider>
+      </>,
+    );
+    startSpeaking(firstRuntimeRef);
+    await waitFor(() => {
+      expect(firstRuntimeRef.current!.thread.getState().speech).toBeDefined();
+    });
+
+    const event = dispatchEscape(screen.getByTestId("second-thread-control"));
+
+    expect(firstSpeech.cancel).not.toHaveBeenCalled();
+    expect(secondSpeech.cancel).not.toHaveBeenCalled();
+    expect(event.defaultPrevented).toBe(false);
+  });
+
+  it("stops speech in the thread containing the Escape target", async () => {
+    const firstSpeech = createSpeechAdapter();
+    const secondSpeech = createSpeechAdapter();
+    const firstRuntimeRef: RuntimeRef = { current: null };
+    const secondRuntimeRef: RuntimeRef = { current: null };
+    render(
+      <>
+        <RuntimeProvider
+          runtimeRef={firstRuntimeRef}
+          speech={firstSpeech.adapter}
+        >
+          <ThreadPrimitiveRoot>
+            <button data-testid="first-thread-control" />
+          </ThreadPrimitiveRoot>
+        </RuntimeProvider>
+        <RuntimeProvider
+          runtimeRef={secondRuntimeRef}
+          speech={secondSpeech.adapter}
+        >
+          <ThreadPrimitiveRoot>
+            <button data-testid="second-thread-control" />
+          </ThreadPrimitiveRoot>
+        </RuntimeProvider>
+      </>,
+    );
+    startSpeaking(firstRuntimeRef);
+    startSpeaking(secondRuntimeRef);
+    await waitFor(() => {
+      expect(firstRuntimeRef.current!.thread.getState().speech).toBeDefined();
+      expect(secondRuntimeRef.current!.thread.getState().speech).toBeDefined();
+    });
+
+    const event = dispatchEscape(screen.getByTestId("second-thread-control"));
+
+    expect(firstSpeech.cancel).not.toHaveBeenCalled();
+    expect(secondSpeech.cancel).toHaveBeenCalledOnce();
+    expect(event.defaultPrevented).toBe(true);
+  });
+
+  it("stops speech in a thread mounted inside a closed shadow root", async () => {
+    const firstSpeech = createSpeechAdapter();
+    const secondSpeech = createSpeechAdapter();
+    const firstRuntimeRef: RuntimeRef = { current: null };
+    const secondRuntimeRef: RuntimeRef = { current: null };
+    render(
+      <RuntimeProvider
+        runtimeRef={firstRuntimeRef}
+        speech={firstSpeech.adapter}
+      >
+        <ThreadPrimitiveRoot />
+      </RuntimeProvider>,
+    );
+    const shadowHost = document.createElement("div");
+    const shadowRoot = shadowHost.attachShadow({ mode: "closed" });
+    const shadowContainer = document.createElement("div");
+    shadowRoot.append(shadowContainer);
+    document.body.append(shadowHost);
+    const shadowView = render(
+      <RuntimeProvider
+        runtimeRef={secondRuntimeRef}
+        speech={secondSpeech.adapter}
+      >
+        <ThreadPrimitiveRoot>
+          <button data-testid="shadow-thread-control" />
+        </ThreadPrimitiveRoot>
+      </RuntimeProvider>,
+      { container: shadowContainer },
+    );
+
+    try {
+      startSpeaking(firstRuntimeRef);
+      startSpeaking(secondRuntimeRef);
+      await waitFor(() => {
+        expect(firstRuntimeRef.current!.thread.getState().speech).toBeDefined();
+        expect(
+          secondRuntimeRef.current!.thread.getState().speech,
+        ).toBeDefined();
+      });
+
+      const event = dispatchEscape(
+        shadowContainer.querySelector("[data-testid=shadow-thread-control]")!,
+      );
+
+      expect(firstSpeech.cancel).not.toHaveBeenCalled();
+      expect(secondSpeech.cancel).toHaveBeenCalledOnce();
+      expect(event.defaultPrevented).toBe(true);
+    } finally {
+      shadowView.unmount();
+      shadowHost.remove();
+    }
+  });
+
   it("stops speech from outside the only mounted thread", async () => {
     const speech = createSpeechAdapter();
     const runtimeRef: RuntimeRef = { current: null };

package/src/primitives/thread/ThreadRoot.tsx

@@ -5,9 +5,15 @@ import {
   type ComponentRef,
   forwardRef,
   type ComponentPropsWithoutRef,
+  type KeyboardEventHandler,
   useEffect,
+  useRef,
 } from "react";
 import { useAui } from "@assistant-ui/store";
+import { useComposedRefs } from "radix-ui/internal";
+import { ThreadRootElementContext } from "./ThreadRootElementContext";
+
+const escapeEventThreadRoots = new WeakMap<KeyboardEvent, Element>();
 
 export namespace ThreadPrimitiveRoot {
   export type Element = ComponentRef<typeof Primitive.div>;
@@ -24,8 +30,9 @@ export namespace ThreadPrimitiveRoot {
  * This component serves as the foundational wrapper for all thread-related components.
  * It provides the basic structure and context needed for thread functionality.
  *
- * While this component is mounted, an unhandled Escape keydown stops active speech, even if
- * the action bar that started it is no longer mounted.
+ * While this component is mounted, an unhandled Escape keydown inside this thread stops its
+ * active speech, even if the action bar that started it is no longer mounted. Escape outside
+ * every thread retains the document-level fallback.
  *
  * @example
  * ```tsx
@@ -43,11 +50,31 @@ export const ThreadPrimitiveRoot = forwardRef<
   ThreadPrimitiveRoot.Props
 >((props, ref) => {
   const aui = useAui();
+  const rootRef = useRef<ThreadPrimitiveRoot.Element>(null);
+  const composedRef = useComposedRefs(ref, rootRef);
+  const handleRootKeyDown: KeyboardEventHandler<ThreadPrimitiveRoot.Element> = (
+    event,
+  ) => {
+    if (event.key === "Escape" && rootRef.current) {
+      escapeEventThreadRoots.set(event.nativeEvent, rootRef.current);
+    }
+    props.onKeyDown?.(event);
+  };
 
   useEffect(() => {
     const handleKeyDown = (event: KeyboardEvent) => {
       if (event.key !== "Escape") return;
       if (event.defaultPrevented || aui.thread.source === null) return;
+      const eventThreadRoot =
+        escapeEventThreadRoots.get(event) ??
+        event
+          .composedPath()
+          .find(
+            (target) =>
+              target instanceof Element &&
+              target.hasAttribute("data-aui-thread-root"),
+          );
+      if (eventThreadRoot && eventThreadRoot !== rootRef.current) return;
       if (aui.thread.getState().speech == null) return;
       event.preventDefault();
       try {
@@ -70,7 +97,16 @@ export const ThreadPrimitiveRoot = forwardRef<
     };
   }, [aui]);
 
-  return <Primitive.div {...props} ref={ref} />;
+  return (
+    <ThreadRootElementContext.Provider value={rootRef}>
+      <Primitive.div
+        {...props}
+        data-aui-thread-root=""
+        ref={composedRef}
+        onKeyDown={handleRootKeyDown}
+      />
+    </ThreadRootElementContext.Provider>
+  );
 });
 
 ThreadPrimitiveRoot.displayName = "ThreadPrimitive.Root";

package/src/primitives/thread/ThreadRootElementContext.ts

@@ -0,0 +1,10 @@
+"use client";
+
+import { createContext, type RefObject, useContext } from "react";
+
+export const ThreadRootElementContext = createContext<
+  RefObject<HTMLElement | null> | undefined
+>(undefined);
+
+export const useThreadRootElementRef = () =>
+  useContext(ThreadRootElementContext);

package/src/primitives/thread/topAnchor/mountTopAnchorReserve.test.ts

@@ -7,6 +7,10 @@ import {
 } from "./mountTopAnchorReserve";
 
 class ResizeObserverMock {
+  static callbacks: (() => void)[] = [];
+  constructor(callback: () => void) {
+    ResizeObserverMock.callbacks.push(callback);
+  }
   observe = vi.fn();
   disconnect = vi.fn();
 }
@@ -48,6 +52,7 @@ const activeTopAnchorTurn = { anchorId: "user-1", targetId: "assistant-1" };
 describe("mountTopAnchorReserve", () => {
   beforeEach(() => {
     vi.useFakeTimers();
+    ResizeObserverMock.callbacks = [];
     vi.stubGlobal("ResizeObserver", ResizeObserverMock);
     vi.stubGlobal("MutationObserver", MutationObserverMock);
   });
@@ -331,6 +336,197 @@ describe("mountTopAnchorReserve", () => {
     expect(reserve.style.height).toBe("0px");
   });
 
+  const mountPinnedViewport = () => {
+    const viewport = document.createElement("div");
+    const anchor = document.createElement("div");
+    const target = document.createElement("div");
+    const reserveHost = document.createElement("div");
+    reserveHost.append(target);
+    document.body.append(reserveHost);
+
+    let naturalScrollHeight = 560;
+
+    defineReadonlyNumber(viewport, "offsetTop", 0);
+    defineReadonlyNumber(viewport, "clientHeight", 400);
+    Object.defineProperty(viewport, "scrollHeight", {
+      configurable: true,
+      get: () => {
+        const reserve = document.querySelector<HTMLElement>(
+          "[data-aui-top-anchor-reserve]",
+        );
+        return (
+          naturalScrollHeight + Number.parseFloat(reserve?.style.height || "0")
+        );
+      },
+    });
+    defineReadonlyNumber(anchor, "offsetTop", 220);
+    defineReadonlyNumber(anchor, "offsetHeight", 64);
+    anchor.dataset.messageId = "msg-1";
+    viewport.scrollTo = vi.fn();
+
+    const { store, setState } = makeStore({
+      turnAnchor: "top",
+      element: { viewport, anchor, target },
+      targetConfig: numericClamp,
+      topAnchorTurn: activeTopAnchorTurn,
+    });
+
+    const unmount = mountTopAnchorReserve(store);
+    vi.runOnlyPendingTimers();
+
+    const reserve = reserveHost.querySelector<HTMLElement>(
+      "[data-aui-top-anchor-reserve]",
+    )!;
+    Object.defineProperty(reserve, "offsetHeight", {
+      configurable: true,
+      get: () => Number.parseFloat(reserve.style.height || "0"),
+    });
+    vi.runOnlyPendingTimers();
+
+    expect(viewport.scrollTo).toHaveBeenCalledTimes(1);
+    expect(viewport.scrollTo).toHaveBeenLastCalledWith({
+      top: 220,
+      behavior: "smooth",
+    });
+
+    viewport.scrollTop = 220;
+    viewport.dispatchEvent(new Event("scroll"));
+
+    return {
+      viewport,
+      anchor,
+      target,
+      store,
+      setState,
+      unmount,
+      notifyLayout: () => ResizeObserverMock.callbacks.at(-1)!(),
+      setNaturalScrollHeight: (height: number) => {
+        naturalScrollHeight = height;
+      },
+    };
+  };
+
+  it("restores the pre-clamp position when the old offset becomes unreachable", () => {
+    const { viewport, setNaturalScrollHeight } = mountPinnedViewport();
+
+    setNaturalScrollHeight(400);
+    viewport.scrollTop = 60;
+    viewport.dispatchEvent(new Event("scroll"));
+    setNaturalScrollHeight(560);
+    vi.runOnlyPendingTimers();
+
+    expect(viewport.scrollTo).toHaveBeenCalledTimes(2);
+    expect(viewport.scrollTo).toHaveBeenLastCalledWith({
+      top: 220,
+      behavior: "instant",
+    });
+  });
+
+  it("does not undo a reachable scrollIntoView movement", () => {
+    const { viewport, target, notifyLayout } = mountPinnedViewport();
+    target.scrollIntoView = vi.fn(() => {
+      viewport.scrollTop = 100;
+      viewport.dispatchEvent(new Event("scroll"));
+    });
+
+    target.scrollIntoView();
+    notifyLayout();
+    vi.runOnlyPendingTimers();
+
+    expect(viewport.scrollTo).toHaveBeenCalledTimes(1);
+  });
+
+  it("does not undo reachable focus scrolling", () => {
+    const { viewport, target, notifyLayout } = mountPinnedViewport();
+    const input = document.createElement("input");
+    target.append(input);
+
+    input.focus();
+    viewport.scrollTop = 100;
+    viewport.dispatchEvent(new Event("scroll"));
+    notifyLayout();
+    vi.runOnlyPendingTimers();
+
+    expect(viewport.scrollTo).toHaveBeenCalledTimes(1);
+  });
+
+  it("does not restore an out-of-range movement that did not land at the maximum", () => {
+    const { viewport, setNaturalScrollHeight, notifyLayout } =
+      mountPinnedViewport();
+
+    setNaturalScrollHeight(400);
+    viewport.scrollTop = 40;
+    viewport.dispatchEvent(new Event("scroll"));
+    setNaturalScrollHeight(560);
+    notifyLayout();
+    vi.runOnlyPendingTimers();
+
+    expect(viewport.scrollTo).toHaveBeenCalledTimes(1);
+  });
+
+  it("restores at most once per anchor turn", () => {
+    const { viewport, setNaturalScrollHeight, notifyLayout } =
+      mountPinnedViewport();
+
+    setNaturalScrollHeight(400);
+    viewport.scrollTop = 60;
+    viewport.dispatchEvent(new Event("scroll"));
+    setNaturalScrollHeight(560);
+    vi.runOnlyPendingTimers();
+    expect(viewport.scrollTo).toHaveBeenCalledTimes(2);
+
+    viewport.scrollTop = 220;
+    viewport.dispatchEvent(new Event("scroll"));
+    setNaturalScrollHeight(400);
+    viewport.scrollTop = 60;
+    viewport.dispatchEvent(new Event("scroll"));
+    setNaturalScrollHeight(560);
+    notifyLayout();
+    vi.runOnlyPendingTimers();
+
+    expect(viewport.scrollTo).toHaveBeenCalledTimes(2);
+  });
+
+  it("drops a pending restore across an anchor-gap thread transition", () => {
+    const { viewport, setState, setNaturalScrollHeight } =
+      mountPinnedViewport();
+
+    setNaturalScrollHeight(400);
+    viewport.scrollTop = 60;
+    viewport.dispatchEvent(new Event("scroll"));
+
+    setState({
+      turnAnchor: "top",
+      element: { viewport, anchor: null, target: null },
+      targetConfig: null,
+      topAnchorTurn: activeTopAnchorTurn,
+    });
+    vi.runOnlyPendingTimers();
+
+    const nextAnchor = document.createElement("div");
+    const nextTarget = document.createElement("div");
+    document.body.append(nextTarget);
+    defineReadonlyNumber(nextAnchor, "offsetTop", 480);
+    defineReadonlyNumber(nextAnchor, "offsetHeight", 64);
+    nextAnchor.dataset.messageId = "msg-2";
+
+    setState({
+      turnAnchor: "top",
+      element: { viewport, anchor: nextAnchor, target: nextTarget },
+      targetConfig: numericClamp,
+      topAnchorTurn: { anchorId: "user-2", targetId: "assistant-2" },
+    });
+    setNaturalScrollHeight(560);
+    vi.runOnlyPendingTimers();
+    vi.runOnlyPendingTimers();
+
+    expect(viewport.scrollTo).toHaveBeenCalledTimes(2);
+    expect(viewport.scrollTo).toHaveBeenLastCalledWith({
+      top: 480,
+      behavior: "smooth",
+    });
+  });
+
   it("does not repeat the smooth top-anchor scroll for the same message", () => {
     const viewport = document.createElement("div");
     const anchor = document.createElement("div");