@assistant-ui/react 0.15.21 → 0.15.23

package/src/primitives/message/MessagePartsGrouped.tool-interaction.test.tsx

@@ -0,0 +1,97 @@
+// @vitest-environment jsdom
+import { cleanup, render } from "@testing-library/react";
+import { useEffect, type ReactNode } from "react";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import {
+  MessagePrimitiveUnstable_PartsGroupedByParentId,
+  type MessagePrimitiveUnstable_PartsGrouped,
+} from "./MessagePartsGrouped";
+
+const fixture = vi.hoisted(() => ({
+  recordInteraction: vi.fn(),
+  part: {
+    type: "tool-call",
+    toolCallId: "call-1",
+    toolName: "weather",
+    args: {},
+    argsText: "{}",
+    status: { type: "complete" },
+  },
+  state: {
+    message: { parts: [] as unknown[] },
+    tools: { toolUIs: {} as Record<string, unknown> },
+    dataRenderers: { renderers: {}, fallbacks: [] },
+  },
+}));
+
+vi.mock("@assistant-ui/store", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@assistant-ui/store")>()),
+  useAui: () => ({
+    part: {
+      addToolResult: vi.fn(),
+      resumeToolCall: vi.fn(),
+      respondToToolApproval: vi.fn(),
+      unstable_recordInteraction: fixture.recordInteraction,
+    },
+  }),
+  useAuiState: <T,>(selector: (state: unknown) => T) =>
+    selector({ ...fixture.state, part: fixture.part }),
+}));
+
+vi.mock(
+  "../../context/providers/PartByIndexProvider",
+  async (importOriginal) => ({
+    ...(await importOriginal<
+      typeof import("../../context/providers/PartByIndexProvider")
+    >()),
+    PartByIndexProvider: ({ children }: { children: ReactNode }) => children,
+  }),
+);
+
+const interaction = { type: "action" as const, payload: { choice: "retry" } };
+
+afterEach(cleanup);
+
+beforeEach(() => {
+  fixture.recordInteraction.mockReset();
+  fixture.state.message.parts = [fixture.part];
+  fixture.state.tools.toolUIs = {};
+});
+
+describe("MessagePrimitive.Unstable_PartsGrouped", () => {
+  it.each(["override", "registered"] as const)(
+    "forwards unstable_recordInteraction to the %s tool UI",
+    async (kind) => {
+      const Tool = ({
+        unstable_recordInteraction,
+      }: {
+        unstable_recordInteraction?:
+          | ((input: typeof interaction) => Promise<void>)
+          | undefined;
+      }) => {
+        useEffect(() => {
+          void unstable_recordInteraction?.(interaction);
+        }, [unstable_recordInteraction]);
+        return null;
+      };
+      const components = (
+        kind === "override" ? { tools: { Override: Tool } } : {}
+      ) satisfies MessagePrimitiveUnstable_PartsGrouped.Props["components"];
+      if (kind === "registered") {
+        fixture.state.tools.toolUIs = { weather: [{ render: Tool }] };
+      }
+
+      render(
+        <MessagePrimitiveUnstable_PartsGroupedByParentId
+          components={components}
+        />,
+      );
+
+      await vi.waitFor(() =>
+        expect(fixture.recordInteraction).toHaveBeenCalledExactlyOnceWith(
+          interaction,
+        ),
+      );
+    },
+  );
+});

package/src/primitives/message/MessagePartsGrouped.tsx

@@ -249,9 +249,9 @@ const DataUIDisplay = ({
   Fallback: DataMessagePartComponent | undefined;
 } & DataMessagePartProps) => {
   const Render = useAuiState((s) => {
-    const Render = s.dataRenderers.renderers[props.name] ?? Fallback;
-    if (Array.isArray(Render)) return Render[0] ?? Fallback;
-    return Render;
+    const named = s.dataRenderers.renderers[props.name]?.[0];
+    if (named) return named;
+    return s.dataRenderers.fallbacks[0] ?? Fallback;
   });
   if (!Render) return null;
   return <Render {...props} />;
@@ -298,6 +298,7 @@ const MessagePartComponent: FC<MessagePartComponentProps> = ({
     const addResult = aui.part.addToolResult;
     const resume = aui.part.resumeToolCall;
     const respondToApproval = aui.part.respondToToolApproval;
+    const unstable_recordInteraction = aui.part.unstable_recordInteraction;
     if ("Override" in tools)
       return (
         <tools.Override
@@ -305,9 +306,13 @@ const MessagePartComponent: FC<MessagePartComponentProps> = ({
           addResult={addResult}
           resume={resume}
           respondToApproval={respondToApproval}
+          {...(unstable_recordInteraction && { unstable_recordInteraction })}
         />
       );
-    const Tool = tools.by_name?.[part.toolName] ?? tools.Fallback;
+    const Tool =
+      (tools.by_name && Object.hasOwn(tools.by_name, part.toolName)
+        ? tools.by_name[part.toolName]
+        : undefined) ?? tools.Fallback;
     return (
       <ToolUIDisplay
         {...part}
@@ -315,6 +320,7 @@ const MessagePartComponent: FC<MessagePartComponentProps> = ({
         addResult={addResult}
         resume={resume}
         respondToApproval={respondToApproval}
+        {...(unstable_recordInteraction && { unstable_recordInteraction })}
       />
     );
   }
@@ -342,7 +348,10 @@ const MessagePartComponent: FC<MessagePartComponentProps> = ({
       return <Audio {...part} />;
 
     case "data": {
-      const Data = data?.by_name?.[part.name] ?? data?.Fallback;
+      const Data =
+        (data?.by_name && Object.hasOwn(data.by_name, part.name)
+          ? data.by_name[part.name]
+          : undefined) ?? data?.Fallback;
       return <DataUIDisplay {...part} Fallback={Data} />;
     }
 

package/src/primitives/message/MessageRoot.test.tsx

@@ -0,0 +1,108 @@
+// @vitest-environment jsdom
+
+import { act, render } from "@testing-library/react";
+import { createContext, useContext } from "react";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import type { ThreadMessageLike } from "@assistant-ui/core";
+import {
+  AssistantRuntimeProvider,
+  useExternalStoreRuntime,
+} from "@assistant-ui/core/react";
+import { useAuiState } from "@assistant-ui/store";
+import { ThreadPrimitiveMessageByIndex } from "../thread/ThreadMessages";
+import { ThreadPrimitiveRoot } from "../thread/ThreadRoot";
+import { ThreadPrimitiveViewport } from "../thread/ThreadViewport";
+import { MessagePrimitiveRoot } from "./MessageRoot";
+
+const messages: ThreadMessageLike[] = [
+  {
+    id: "message-1",
+    role: "assistant",
+    content: [{ type: "text", text: "Hello" }],
+  },
+];
+
+class TestResizeObserver {
+  observe() {}
+  unobserve() {}
+  disconnect() {}
+}
+
+const VisibilityContext = createContext(true);
+
+const Message = () => {
+  const visible = useContext(VisibilityContext);
+  const isHovering = useAuiState((s) => s.message.isHovering);
+
+  return (
+    <>
+      {visible ? <MessagePrimitiveRoot /> : null}
+      <span data-testid="hover-state" data-hovering={isHovering} />
+    </>
+  );
+};
+
+const Example = ({ visible = true }: { visible?: boolean }) => {
+  const runtime = useExternalStoreRuntime({
+    messages,
+    convertMessage: (message) => message,
+    onNew: async () => {},
+  });
+
+  return (
+    <AssistantRuntimeProvider runtime={runtime}>
+      <VisibilityContext.Provider value={visible}>
+        <ThreadPrimitiveRoot>
+          <ThreadPrimitiveViewport scrollToBottomOnInitialize={false}>
+            <ThreadPrimitiveMessageByIndex index={0} components={{ Message }} />
+          </ThreadPrimitiveViewport>
+        </ThreadPrimitiveRoot>
+      </VisibilityContext.Provider>
+    </AssistantRuntimeProvider>
+  );
+};
+
+beforeEach(() => {
+  vi.stubGlobal("ResizeObserver", TestResizeObserver);
+});
+
+afterEach(() => {
+  vi.unstubAllGlobals();
+  vi.restoreAllMocks();
+});
+
+const mockHoveredElement = () => {
+  const matches = HTMLElement.prototype.matches;
+  vi.spyOn(HTMLElement.prototype, "matches").mockImplementation(function (
+    this: HTMLElement,
+    selector,
+  ) {
+    return selector === ":hover" || matches.call(this, selector);
+  });
+};
+
+describe("MessagePrimitiveRoot", () => {
+  it("synchronizes hover state while mounted", async () => {
+    mockHoveredElement();
+
+    const view = render(<Example />);
+    await act(() => Promise.resolve());
+
+    expect(view.getByTestId("hover-state").getAttribute("data-hovering")).toBe(
+      "true",
+    );
+    view.unmount();
+  });
+
+  it("does not restore hover state after unmount", async () => {
+    mockHoveredElement();
+
+    const view = render(<Example />);
+    view.rerender(<Example visible={false} />);
+    await act(() => Promise.resolve());
+
+    expect(view.getByTestId("hover-state").getAttribute("data-hovering")).toBe(
+      "false",
+    );
+  });
+});

package/src/primitives/message/MessageRoot.tsx

@@ -27,6 +27,7 @@ const useIsHoveringRef = () => {
 
   const callbackRef = useCallback(
     (el: HTMLElement) => {
+      let active = true;
       const handleMouseEnter = () => {
         message.setIsHovering(true);
       };
@@ -39,10 +40,13 @@ const useIsHoveringRef = () => {
 
       if (el.matches(":hover")) {
         // TODO this is needed for SSR to work, figure out why
-        queueMicrotask(() => message.setIsHovering(true));
+        queueMicrotask(() => {
+          if (active) message.setIsHovering(true);
+        });
       }
 
       return () => {
+        active = false;
         el.removeEventListener("mouseenter", handleMouseEnter);
         el.removeEventListener("mouseleave", handleMouseLeave);
         message.setIsHovering(false);

package/src/primitives/reasoning/useScrollLock.test.tsx

@@ -126,4 +126,51 @@ describe("useScrollLock", () => {
 
     expect(scroller.style.paddingRight).toBe("");
   });
+
+  it("locks the current scroll container after the element is reparented", () => {
+    vi.useFakeTimers();
+    const first = document.createElement("div");
+    const second = document.createElement("div");
+    first.style.overflowY = "auto";
+    second.style.overflowY = "auto";
+    document.body.append(first, second);
+
+    const animated = document.createElement("div");
+    first.appendChild(animated);
+    const ref = { current: animated };
+    const { result, unmount } = renderHook(() => useScrollLock(ref, 200));
+
+    result.current();
+    expect(first.style.scrollbarWidth).toBe("none");
+
+    second.appendChild(animated);
+    result.current();
+
+    expect(first.style.scrollbarWidth).toBe("");
+    expect(second.style.scrollbarWidth).toBe("none");
+
+    unmount();
+    expect(second.style.scrollbarWidth).toBe("");
+  });
+
+  it("forgets the old cleanup after moving outside a scroll container", () => {
+    const scroller = document.createElement("div");
+    const plainContainer = document.createElement("div");
+    scroller.style.overflowY = "auto";
+    document.body.append(scroller, plainContainer);
+
+    const animated = document.createElement("div");
+    scroller.appendChild(animated);
+    const ref = { current: animated };
+    const { result, unmount } = renderHook(() => useScrollLock(ref, 200));
+
+    result.current();
+    plainContainer.appendChild(animated);
+    result.current();
+    scroller.style.scrollbarWidth = "thin";
+
+    unmount();
+
+    expect(scroller.style.scrollbarWidth).toBe("thin");
+  });
 });

package/src/primitives/reasoning/useScrollLock.ts

@@ -2,6 +2,16 @@
 
 import { type RefObject, useCallback, useEffect, useRef } from "react";
 
+const findScrollableAncestor = (element: HTMLElement | null) => {
+  let current = element;
+  while (current) {
+    const { overflowY } = getComputedStyle(current);
+    if (overflowY === "scroll" || overflowY === "auto") return current;
+    current = current.parentElement;
+  }
+  return null;
+};
+
 /**
  * Locks scroll position during collapsible/height animations and hides scrollbar.
  *
@@ -32,7 +42,6 @@ export const useScrollLock = <T extends HTMLElement = HTMLElement>(
   animatedElementRef: RefObject<T | null>,
   animationDuration: number,
 ) => {
-  const scrollContainerRef = useRef<HTMLElement | null>(null);
   const cleanupRef = useRef<(() => void) | null>(null);
 
   useEffect(() => {
@@ -44,22 +53,11 @@ export const useScrollLock = <T extends HTMLElement = HTMLElement>(
   const lockScroll = useCallback(() => {
     cleanupRef.current?.();
 
-    (function findScrollableAncestor() {
-      if (scrollContainerRef.current || !animatedElementRef.current) return;
-
-      let el: HTMLElement | null = animatedElementRef.current;
-      while (el) {
-        const { overflowY } = getComputedStyle(el);
-        if (overflowY === "scroll" || overflowY === "auto") {
-          scrollContainerRef.current = el;
-          break;
-        }
-        el = el.parentElement;
-      }
-    })();
-
-    const scrollContainer = scrollContainerRef.current;
-    if (!scrollContainer) return;
+    const scrollContainer = findScrollableAncestor(animatedElementRef.current);
+    if (!scrollContainer) {
+      cleanupRef.current = null;
+      return;
+    }
 
     const scrollPosition = scrollContainer.scrollTop;
     const scrollbarWidth = scrollContainer.style.scrollbarWidth;

package/src/primitives/selectionToolbar/SelectionToolbarRoot.test.tsx

@@ -1,32 +1,54 @@
 /** @vitest-environment jsdom */
-import type { MouseEvent } from "react";
+import type { MouseEvent, ReactNode } from "react";
 import { fireEvent, render } from "@testing-library/react";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
-import type * as GetSelectionMessageIdModule from "../../utils/getSelectionMessageId";
+import type * as StoreModule from "@assistant-ui/store";
+import { ThreadPrimitiveRoot } from "../thread/ThreadRoot";
 import { SelectionToolbarPrimitiveRoot } from "./SelectionToolbarRoot";
 
-vi.mock("../../utils/getSelectionMessageId", async (importOriginal) => ({
-  ...(await importOriginal<typeof GetSelectionMessageIdModule>()),
-  getSelectionMessageId: () => "m1",
+const h = vi.hoisted(() => ({
+  aui: {
+    thread: {
+      source: null,
+      getState: () => ({ speech: undefined }),
+      stopSpeaking: vi.fn(),
+    },
+  },
 }));
 
-const fakeSelection = {
-  isCollapsed: false,
-  toString: () => "selected text",
-  getRangeAt: () => ({
-    getBoundingClientRect: () => ({ top: 100, left: 50, width: 20 }) as DOMRect,
-  }),
-} as unknown as Selection;
+vi.mock("@assistant-ui/store", async (importOriginal) => ({
+  ...(await importOriginal<typeof StoreModule>()),
+  useAui: () => h.aui,
+}));
+
+let selectionMessage: HTMLDivElement;
 
 beforeEach(() => {
+  selectionMessage = document.createElement("div");
+  selectionMessage.dataset.messageId = "m1";
+  selectionMessage.textContent = "selected text";
+  document.body.append(selectionMessage);
+  const selectedNode = selectionMessage.firstChild;
   vi.spyOn(window, "requestAnimationFrame").mockImplementation((cb) => {
     cb(0);
     return 0;
   });
-  vi.spyOn(window, "getSelection").mockReturnValue(fakeSelection);
+  vi.spyOn(window, "getSelection").mockReturnValue({
+    isCollapsed: false,
+    anchorNode: selectedNode,
+    focusNode: selectedNode,
+    rangeCount: 1,
+    toString: () => "selected text",
+    getRangeAt: () => ({
+      commonAncestorContainer: selectedNode,
+      getBoundingClientRect: () =>
+        ({ top: 100, left: 50, width: 20 }) as DOMRect,
+    }),
+  } as unknown as Selection);
 });
 
 afterEach(() => {
+  selectionMessage.remove();
   vi.restoreAllMocks();
 });
 
@@ -73,6 +95,124 @@ describe("SelectionToolbarPrimitiveRoot onMouseDown composition", () => {
   });
 });
 
+describe("SelectionToolbarPrimitiveRoot selection changes", () => {
+  it("opens from a selectionchange event without a mouse or key event", () => {
+    render(<SelectionToolbarPrimitiveRoot data-testid="toolbar" />);
+
+    fireEvent(document, new Event("selectionchange"));
+
+    expect(document.querySelector('[data-testid="toolbar"]')).not.toBeNull();
+  });
+
+  it("closes an open toolbar when the selection collapses", () => {
+    render(<SelectionToolbarPrimitiveRoot data-testid="toolbar" />);
+    fireEvent(document, new Event("selectionchange"));
+    expect(document.querySelector('[data-testid="toolbar"]')).not.toBeNull();
+
+    vi.mocked(window.getSelection).mockReturnValueOnce({
+      isCollapsed: true,
+    } as Selection);
+    fireEvent(document, new Event("selectionchange"));
+
+    expect(document.querySelector('[data-testid="toolbar"]')).toBeNull();
+  });
+
+  it("opens only inside the thread that owns the selection", () => {
+    const { getByTestId } = render(
+      <>
+        <ThreadPrimitiveRoot>
+          <div data-message-id="m1">
+            <span data-testid="first-message">first</span>
+          </div>
+          <SelectionToolbarPrimitiveRoot data-testid="first-toolbar" />
+        </ThreadPrimitiveRoot>
+        <ThreadPrimitiveRoot>
+          <div data-message-id="m1">
+            <span data-testid="second-message">second</span>
+          </div>
+          <SelectionToolbarPrimitiveRoot data-testid="second-toolbar" />
+        </ThreadPrimitiveRoot>
+      </>,
+    );
+    const selectedNode = getByTestId("first-message").firstChild;
+    vi.mocked(window.getSelection).mockReturnValue({
+      isCollapsed: false,
+      anchorNode: selectedNode,
+      focusNode: selectedNode,
+      rangeCount: 1,
+      toString: () => "first",
+      getRangeAt: () => ({
+        commonAncestorContainer: selectedNode,
+        getBoundingClientRect: () =>
+          ({ top: 100, left: 50, width: 20 }) as DOMRect,
+      }),
+    } as unknown as Selection);
+
+    fireEvent(document, new Event("selectionchange"));
+
+    expect(
+      document.querySelector('[data-testid="first-toolbar"]'),
+    ).not.toBeNull();
+    expect(document.querySelector('[data-testid="second-toolbar"]')).toBeNull();
+
+    const secondSelectedNode = getByTestId("second-message").firstChild;
+    vi.mocked(window.getSelection).mockReturnValue({
+      isCollapsed: false,
+      anchorNode: secondSelectedNode,
+      focusNode: secondSelectedNode,
+      rangeCount: 1,
+      toString: () => "second",
+      getRangeAt: () => ({
+        commonAncestorContainer: secondSelectedNode,
+        getBoundingClientRect: () =>
+          ({ top: 100, left: 50, width: 20 }) as DOMRect,
+      }),
+    } as unknown as Selection);
+
+    fireEvent(document, new Event("selectionchange"));
+
+    expect(document.querySelector('[data-testid="first-toolbar"]')).toBeNull();
+    expect(
+      document.querySelector('[data-testid="second-toolbar"]'),
+    ).not.toBeNull();
+  });
+
+  it("warns when a custom thread root does not forward its ref", () => {
+    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+    const RefDroppingRoot = ({ children }: { children?: ReactNode }) => (
+      <div>{children}</div>
+    );
+    const { getByTestId } = render(
+      <ThreadPrimitiveRoot render={<RefDroppingRoot />}>
+        <div data-message-id="m1">
+          <span data-testid="message">text</span>
+        </div>
+        <SelectionToolbarPrimitiveRoot data-testid="toolbar" />
+      </ThreadPrimitiveRoot>,
+    );
+    const selectedNode = getByTestId("message").firstChild;
+    vi.mocked(window.getSelection).mockReturnValue({
+      isCollapsed: false,
+      anchorNode: selectedNode,
+      focusNode: selectedNode,
+      rangeCount: 1,
+      toString: () => "text",
+      getRangeAt: () => ({
+        commonAncestorContainer: selectedNode,
+        getBoundingClientRect: () =>
+          ({ top: 100, left: 50, width: 20 }) as DOMRect,
+      }),
+    } as unknown as Selection);
+
+    fireEvent(document, new Event("selectionchange"));
+
+    expect(document.querySelector('[data-testid="toolbar"]')).not.toBeNull();
+    expect(warn).toHaveBeenCalledWith(
+      "[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.",
+    );
+  });
+});
+
 describe("SelectionToolbarPrimitiveRoot frame cleanup", () => {
   // Defer the frame instead of running it inline, so the window between the
   // selection event and the frame is observable.
@@ -106,7 +246,7 @@ describe("SelectionToolbarPrimitiveRoot frame cleanup", () => {
     const { unmount } = render(<SelectionToolbarPrimitiveRoot />);
 
     fireEvent.mouseUp(document);
-    fireEvent.keyUp(document);
+    fireEvent(document, new Event("selectionchange"));
     expect(frames).toHaveLength(2);
     expect(cancelAnimationFrame).toHaveBeenCalledWith(1);
 
@@ -115,4 +255,78 @@ describe("SelectionToolbarPrimitiveRoot frame cleanup", () => {
     expect(cancelAnimationFrame).toHaveBeenCalledTimes(2);
     expect(cancelAnimationFrame).toHaveBeenLastCalledWith(2);
   });
+
+  it("waits until mouseup to measure a drag selection", () => {
+    const { frames } = deferFrames();
+    render(<SelectionToolbarPrimitiveRoot />);
+
+    fireEvent.mouseDown(document);
+    fireEvent(document, new Event("selectionchange"));
+    fireEvent(document, new Event("selectionchange"));
+    expect(frames).toHaveLength(0);
+
+    fireEvent.mouseUp(document);
+    expect(frames).toHaveLength(1);
+  });
+
+  it("recovers when a drag ends without mouseup", () => {
+    const { frames } = deferFrames();
+    render(<SelectionToolbarPrimitiveRoot />);
+
+    fireEvent.mouseDown(document);
+    fireEvent(document, new Event("selectionchange"));
+    expect(frames).toHaveLength(0);
+
+    fireEvent(document, new Event("dragend"));
+    fireEvent(document, new Event("selectionchange"));
+    expect(frames).toHaveLength(2);
+  });
+
+  it("recovers when the window blurs during a drag", () => {
+    const { frames } = deferFrames();
+    render(<SelectionToolbarPrimitiveRoot />);
+
+    fireEvent.mouseDown(document);
+    fireEvent(document, new Event("selectionchange"));
+    expect(frames).toHaveLength(0);
+
+    fireEvent.blur(window);
+    fireEvent(document, new Event("selectionchange"));
+    expect(frames).toHaveLength(1);
+  });
+
+  it("recovers when a context menu consumes the mouseup of a press", () => {
+    const { frames } = deferFrames();
+    render(<SelectionToolbarPrimitiveRoot />);
+
+    fireEvent.mouseDown(document, { button: 2 });
+    fireEvent.contextMenu(document);
+    fireEvent(document, new Event("selectionchange"));
+    expect(frames).toHaveLength(1);
+  });
+
+  it("cancels a queued selection frame when the page scrolls", () => {
+    const { frames, cancelAnimationFrame } = deferFrames();
+    render(<SelectionToolbarPrimitiveRoot />);
+
+    fireEvent(document, new Event("selectionchange"));
+    fireEvent.scroll(document);
+
+    expect(frames).toHaveLength(1);
+    expect(cancelAnimationFrame).toHaveBeenCalledWith(1);
+  });
+
+  it("cancels a queued frame when the selection collapses", () => {
+    const { frames, cancelAnimationFrame } = deferFrames();
+    render(<SelectionToolbarPrimitiveRoot />);
+
+    fireEvent(document, new Event("selectionchange"));
+    vi.mocked(window.getSelection).mockReturnValueOnce({
+      isCollapsed: true,
+    } as Selection);
+    fireEvent(document, new Event("selectionchange"));
+
+    expect(frames).toHaveLength(1);
+    expect(cancelAnimationFrame).toHaveBeenCalledWith(1);
+  });
 });