@assistant-ui/react 0.15.23 → 0.15.25

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

@@ -1,7 +1,13 @@
 // @vitest-environment jsdom
 
 import { render, screen, waitFor } from "@testing-library/react";
-import { useEffect, type FC, type ReactNode } from "react";
+import {
+  useEffect,
+  useState,
+  type FC,
+  type ReactNode,
+  type PropsWithChildren,
+} from "react";
 import { describe, expect, it } from "vitest";
 import type { ThreadMessageLike } from "@assistant-ui/core";
 import {
@@ -63,10 +69,12 @@ const Example = ({
   content,
   Message: MessageComponent = Message,
   extra,
+  isRunning = false,
 }: {
   content: ThreadMessageLike["content"];
   Message?: FC;
   extra?: ReactNode;
+  isRunning?: boolean;
 }) => {
   const messages: ThreadMessageLike[] = [
     { id: "message", role: "assistant", content },
@@ -74,6 +82,7 @@ const Example = ({
   const runtime = useExternalStoreRuntime({
     messages,
     convertMessage: (message) => message,
+    isRunning,
     onNew: async () => {},
   });
   return (
@@ -88,6 +97,267 @@ const Example = ({
 };
 
 describe("MessagePrimitive.Unstable_PartsGroupedByParentId", () => {
+  it.each(["tool-call", "reasoning"] as const)(
+    "keeps wrapper and mounted state when a tool call is appended after %s",
+    (firstType) => {
+      let mounts = 0;
+      const Group = ({ children }: PropsWithChildren) => {
+        const [mount] = useState(() => ++mounts);
+        return <section data-mount={mount}>{children}</section>;
+      };
+      const Text = ({ text }: { text: string }) => {
+        const [seed] = useState(text);
+        return <span>{`${seed}:${text}`}</span>;
+      };
+      const Tool = ({ argsText }: { argsText: string }) => (
+        <Text text={argsText} />
+      );
+      const Message = partsMessage({
+        Reasoning: Text,
+        Group,
+        tools: { Fallback: Tool },
+      });
+      const first =
+        firstType === "tool-call"
+          ? {
+              type: "tool-call" as const,
+              toolCallId: "t2",
+              toolName: "task",
+              args: {},
+              argsText: "draft",
+              parentId: "parent",
+            }
+          : {
+              type: "reasoning" as const,
+              id: "r1",
+              text: "draft",
+              parentId: "parent",
+            };
+      const view = render(
+        <Example Message={Message} content={[first]} isRunning />,
+      );
+      const wrapper = view.container.querySelector("section");
+      const streamed =
+        first.type === "tool-call"
+          ? { ...first, argsText: "streamed" }
+          : { ...first, text: "streamed" };
+      const appended = {
+        type: "tool-call" as const,
+        toolCallId: "t1",
+        toolName: "task",
+        args: {},
+        argsText: "appended",
+        parentId: "parent",
+      };
+      view.rerender(
+        <Example Message={Message} content={[streamed, appended]} isRunning />,
+      );
+      expect(view.container.querySelector("section")).toBe(wrapper);
+      expect(wrapper?.dataset.mount).toBe("1");
+      expect(
+        Array.from(
+          view.container.querySelectorAll("span"),
+          (el) => el.textContent,
+        ),
+      ).toEqual(["draft:streamed", "appended:appended"]);
+      view.rerender(
+        <Example
+          Message={Message}
+          content={[
+            streamed,
+            { ...appended, argsText: "updated" },
+            { ...appended, toolCallId: "t0", argsText: "last" },
+          ]}
+          isRunning
+        />,
+      );
+      expect(view.container.querySelector("section")).toBe(wrapper);
+      expect(
+        Array.from(
+          view.container.querySelectorAll("span"),
+          (el) => el.textContent,
+        ),
+      ).toEqual(["draft:streamed", "appended:updated", "last:last"]);
+    },
+  );
+
+  it("keeps parent names and part ids with key delimiters separate", () => {
+    let mounts = 0;
+    const Group = ({ children }: PropsWithChildren) => {
+      const [mount] = useState(() => ++mounts);
+      return <section data-mount={mount}>{children}</section>;
+    };
+    const Text = ({ text }: { text: string }) => {
+      const [seed] = useState(text);
+      return <span>{`${seed}:${text}`}</span>;
+    };
+    const Message = partsMessage({ Text, Group });
+    const first = {
+      type: "text" as const,
+      id: "x-id:text:y",
+      text: "first",
+      parentId: "a",
+    };
+    const second = {
+      type: "text" as const,
+      id: "y",
+      text: "second",
+      parentId: "a-id:text:x",
+    };
+    const view = render(
+      <Example Message={Message} content={[first, second]} />,
+    );
+    view.rerender(
+      <Example
+        Message={Message}
+        content={[
+          { ...second, text: "second updated" },
+          { ...first, text: "first updated" },
+        ]}
+      />,
+    );
+    expect(
+      Array.from(
+        view.container.querySelectorAll("section"),
+        (el) => el.dataset.mount,
+      ),
+    ).toEqual(["2", "1"]);
+    expect(
+      Array.from(
+        view.container.querySelectorAll("span"),
+        (el) => el.textContent,
+      ),
+    ).toEqual(["second:second updated", "first:first updated"]);
+  });
+
+  it.each([undefined, "parent"])(
+    "keeps wrapper and leaf state when later identified parts swap with parent=%s",
+    (parentId) => {
+      let mounts = 0;
+      const Group = ({ children }: PropsWithChildren) => {
+        const [mount] = useState(() => ++mounts);
+        return <section data-mount={mount}>{children}</section>;
+      };
+      const Text = ({ text }: { text: string }) => {
+        const [seed] = useState(text);
+        return <span>{`${seed}:${text}`}</span>;
+      };
+      const Message = partsMessage({ Text, Group });
+      const first = {
+        type: "text" as const,
+        id: "p1",
+        text: "first",
+        ...(parentId !== undefined && { parentId }),
+      };
+      const second = {
+        type: "text" as const,
+        id: "p2",
+        text: "second",
+        ...(parentId !== undefined && { parentId }),
+      };
+      const third = {
+        ...second,
+        id: "p3",
+        text: "third",
+      };
+      const view = render(
+        <Example Message={Message} content={[first, second, third]} />,
+      );
+      view.rerender(
+        <Example
+          Message={Message}
+          content={[
+            { ...first, text: "first updated" },
+            { ...third, text: "third updated" },
+            { ...second, text: "second updated" },
+          ]}
+        />,
+      );
+      expect(
+        Array.from(
+          view.container.querySelectorAll("span"),
+          (el) => el.textContent,
+        ),
+      ).toEqual([
+        "first:first updated",
+        "third:third updated",
+        "second:second updated",
+      ]);
+      expect(
+        Array.from(
+          view.container.querySelectorAll("section"),
+          (el) => el.dataset.mount,
+        ),
+      ).toEqual(parentId === undefined ? ["1", "3", "2"] : ["1"]);
+    },
+  );
+
+  it("keys unidentified groups by index and keeps parent names separate", () => {
+    let mounts = 0;
+    const Group = ({
+      groupKey,
+      children,
+    }: PropsWithChildren<{ groupKey: string | undefined }>) => {
+      const [mount] = useState(() => ++mounts);
+      return (
+        <section data-mount={mount} data-parent={groupKey}>
+          {children}
+        </section>
+      );
+    };
+    const Message = partsMessage({ Group });
+    const parts = [
+      { type: "text" as const, text: "first" },
+      { type: "text" as const, text: "second" },
+    ];
+    const view = render(<Example Message={Message} content={parts} />);
+    view.rerender(
+      <Example
+        Message={Message}
+        content={[
+          { type: "text", text: "parented", parentId: "ungrouped" },
+          ...parts,
+        ]}
+      />,
+    );
+    expect(
+      Array.from(
+        view.container.querySelectorAll("section"),
+        (el) => el.dataset.mount,
+      ),
+    ).toEqual(["3", "2", "4"]);
+  });
+
+  it("keeps a text seed while streaming and resets it for a replacement id", () => {
+    const SeededText = ({ text }: { text: string }) => {
+      const [seed] = useState(text);
+      return <span>{`${seed}:${text}`}</span>;
+    };
+    const Message = partsMessage({ Text: SeededText });
+    const view = render(
+      <Example
+        Message={Message}
+        content={[{ type: "text", id: "p1", text: "old" }]}
+      />,
+    );
+
+    view.rerender(
+      <Example
+        Message={Message}
+        content={[{ type: "text", id: "p1", text: "old streamed" }]}
+      />,
+    );
+    expect(view.container.textContent).toBe("old:old streamed");
+
+    view.rerender(
+      <Example
+        Message={Message}
+        content={[{ type: "text", id: "p2", text: "new" }]}
+      />,
+    );
+    expect(view.container.textContent).toBe("new:new");
+  });
+
   it("keeps parent IDs separate from ungrouped parts across content updates", () => {
     const { rerender } = render(
       <Example

package/src/primitives/message/MessagePartsGrouped.tsx

@@ -27,6 +27,7 @@ import type {
 } from "@assistant-ui/core/react";
 import { MessagePartPrimitiveInProgress } from "../messagePart/MessagePartInProgress";
 import type { MessagePartStatus } from "@assistant-ui/core";
+import { getMessagePartKeys } from "@assistant-ui/core/internal";
 
 type MessagePartGroup = {
   groupKey: string | undefined;
@@ -73,15 +74,16 @@ const groupMessagePartsByParentId: GroupingFunction = (
 
 const useMessagePartsGrouped = (
   groupingFunction: GroupingFunction,
-): MessagePartGroup[] => {
+): { groups: MessagePartGroup[]; partKeys: string[] } => {
   const parts = useAuiState((s) => s.message.parts);
 
-  return useMemo(() => {
+  const groups = useMemo(() => {
     if (parts.length === 0) {
       return [];
     }
     return groupingFunction(parts);
   }, [parts, groupingFunction]);
+  return { groups, partKeys: getMessagePartKeys(parts) };
 };
 
 export namespace MessagePrimitiveUnstable_PartsGrouped {
@@ -463,7 +465,8 @@ export const MessagePrimitiveUnstable_PartsGrouped: FC<
   MessagePrimitiveUnstable_PartsGrouped.Props
 > = ({ groupingFunction, components }) => {
   const contentLength = useAuiState((s) => s.message.parts.length);
-  const messageGroups = useMessagePartsGrouped(groupingFunction);
+  const { groups: messageGroups, partKeys } =
+    useMessagePartsGrouped(groupingFunction);
 
   const partsElements = useMemo(() => {
     if (contentLength === 0) {
@@ -472,16 +475,21 @@ export const MessagePrimitiveUnstable_PartsGrouped: FC<
 
     return messageGroups.map((group, groupIndex) => {
       const GroupComponent = components?.Group ?? defaultComponents.Group;
+      const identity = partKeys[group.indices[0]!];
 
       return (
         <GroupComponent
-          key={`group-${groupIndex}-${group.groupKey ?? "ungrouped"}`}
+          key={JSON.stringify([
+            "group",
+            group.groupKey ?? null,
+            identity?.includes(":") ? `id:${identity}` : groupIndex,
+          ])}
           groupKey={group.groupKey}
           indices={group.indices}
         >
           {group.indices.map((partIndex) => (
             <MessagePart
-              key={partIndex}
+              key={partKeys[partIndex]}
               partIndex={partIndex}
               components={components}
             />
@@ -489,7 +497,7 @@ export const MessagePrimitiveUnstable_PartsGrouped: FC<
         </GroupComponent>
       );
     });
-  }, [messageGroups, components, contentLength]);
+  }, [messageGroups, components, contentLength, partKeys]);
 
   return <>{partsElements}</>;
 };

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

@@ -2,14 +2,26 @@
 
 import { act, render } from "@testing-library/react";
 import { createContext, useContext } from "react";
+import type { ReactNode } from "react";
+import { hydrateRoot } from "react-dom/client";
+import { renderToString } from "react-dom/server";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
-import type { ThreadMessageLike } from "@assistant-ui/core";
+import type {
+  ChatModelAdapter,
+  ThreadMessageLike,
+  ThreadMessage,
+} from "@assistant-ui/core";
 import {
   AssistantRuntimeProvider,
+  ThreadListItemRuntimeProvider,
   useExternalStoreRuntime,
+  useLocalRuntime,
 } from "@assistant-ui/core/react";
 import { useAuiState } from "@assistant-ui/store";
-import { ThreadPrimitiveMessageByIndex } from "../thread/ThreadMessages";
+import {
+  ThreadPrimitiveMessageByIndex,
+  ThreadPrimitiveMessages,
+} from "../thread/ThreadMessages";
 import { ThreadPrimitiveRoot } from "../thread/ThreadRoot";
 import { ThreadPrimitiveViewport } from "../thread/ThreadViewport";
 import { MessagePrimitiveRoot } from "./MessageRoot";
@@ -22,6 +34,14 @@ const messages: ThreadMessageLike[] = [
   },
 ];
 
+const initialMessages: readonly ThreadMessageLike[] = [
+  { role: "user", content: "Hydrated message" },
+];
+
+const chatModel: ChatModelAdapter = {
+  run: async () => ({ content: [] }),
+};
+
 class TestResizeObserver {
   observe() {}
   unobserve() {}
@@ -62,6 +82,64 @@ const Example = ({ visible = true }: { visible?: boolean }) => {
   );
 };
 
+const HydrationMessage = () => (
+  <MessagePrimitiveRoot data-testid="hydration-message" />
+);
+
+const LocalThread = ({
+  onRuntime,
+  turnAnchor,
+}: {
+  onRuntime: (runtime: ReturnType<typeof useLocalRuntime>) => void;
+  turnAnchor: "bottom" | "top";
+}) => {
+  const runtime = useLocalRuntime(chatModel, { initialMessages });
+  onRuntime(runtime);
+
+  return (
+    <AssistantRuntimeProvider runtime={runtime}>
+      <ThreadPrimitiveRoot>
+        <ThreadPrimitiveViewport
+          scrollToBottomOnInitialize={false}
+          turnAnchor={turnAnchor}
+        >
+          <ThreadPrimitiveMessages components={{ Message: HydrationMessage }} />
+        </ThreadPrimitiveViewport>
+      </ThreadPrimitiveRoot>
+    </AssistantRuntimeProvider>
+  );
+};
+
+const NestedLocalRuntime = ({
+  onRuntime,
+  turnAnchor,
+}: {
+  onRuntime: (runtime: ReturnType<typeof useLocalRuntime>) => void;
+  turnAnchor: "bottom" | "top";
+}) => {
+  const host = useExternalStoreRuntime<ThreadMessage>({
+    messages: [],
+    onNew: async () => {},
+  });
+
+  return (
+    <AssistantRuntimeProvider runtime={host}>
+      <ThreadListItemRuntimeProvider runtime={host.threads.mainItem}>
+        <LocalThread onRuntime={onRuntime} turnAnchor={turnAnchor} />
+      </ThreadListItemRuntimeProvider>
+    </AssistantRuntimeProvider>
+  );
+};
+
+const renderOnServer = (node: ReactNode) => {
+  vi.stubGlobal("document", undefined);
+  try {
+    return renderToString(node);
+  } finally {
+    vi.unstubAllGlobals();
+  }
+};
+
 beforeEach(() => {
   vi.stubGlobal("ResizeObserver", TestResizeObserver);
 });
@@ -82,6 +160,51 @@ const mockHoveredElement = () => {
 };
 
 describe("MessagePrimitiveRoot", () => {
+  it.each(["bottom", "top"] as const)(
+    "hydrates nested local message ids with %s anchoring",
+    async (turnAnchor) => {
+      let serverRuntime: ReturnType<typeof useLocalRuntime> | undefined;
+      const html = renderOnServer(
+        <NestedLocalRuntime
+          onRuntime={(runtime) => (serverRuntime = runtime)}
+          turnAnchor={turnAnchor}
+        />,
+      );
+      const serverMessageId = serverRuntime!.thread.getState().messages[0]!.id;
+      const container = document.createElement("div");
+      container.innerHTML = html;
+      vi.stubGlobal("ResizeObserver", TestResizeObserver);
+
+      let clientRuntime: ReturnType<typeof useLocalRuntime> | undefined;
+      const errors = vi.spyOn(console, "error").mockImplementation(() => {});
+      let root: ReturnType<typeof hydrateRoot> | undefined;
+      await act(async () => {
+        root = hydrateRoot(
+          container,
+          <NestedLocalRuntime
+            onRuntime={(runtime) => (clientRuntime = runtime)}
+            turnAnchor={turnAnchor}
+          />,
+        );
+      });
+
+      const clientMessageId = clientRuntime!.thread.getState().messages[0]!.id;
+      const domMessageId = container
+        .querySelector("[data-testid='hydration-message']")!
+        .getAttribute("data-message-id");
+
+      expect(clientMessageId).toBe(serverMessageId);
+      expect(domMessageId).toBe(clientMessageId);
+      expect(
+        errors.mock.calls.some((args) =>
+          args.some((value) => String(value).includes("data-message-id")),
+        ),
+      ).toBe(false);
+
+      await act(async () => root?.unmount());
+    },
+  );
+
   it("synchronizes hover state while mounted", async () => {
     mockHoveredElement();
 

package/src/primitives/thread/ThreadLoadEarlier.tsx

@@ -0,0 +1,35 @@
+"use client";
+
+import { useCallback } from "react";
+import { useAui, useAuiState } from "@assistant-ui/store";
+import {
+  type ActionButtonElement,
+  type ActionButtonProps,
+  createActionButton,
+} from "../../utils/createActionButton";
+
+const useThreadLoadEarlier = () => {
+  const aui = useAui();
+  const hasEarlier = useAuiState((s) => s.thread.hasEarlier);
+  const loadEarlier = useCallback(() => {
+    void aui.thread.loadEarlier();
+  }, [aui]);
+  if (!hasEarlier) return null;
+  return loadEarlier;
+};
+
+export namespace ThreadPrimitiveLoadEarlier {
+  export type Element = ActionButtonElement;
+  export type Props = ActionButtonProps<typeof useThreadLoadEarlier>;
+}
+
+/**
+ * A button that loads the page before the first loaded message. Disabled
+ * while the runtime reports no earlier messages. A press while a page is
+ * loading joins that load, so the button stays enabled and keeps keyboard
+ * focus until the page lands; read `thread.isLoadingEarlier` to show progress.
+ */
+export const ThreadPrimitiveLoadEarlier = createActionButton(
+  "ThreadPrimitive.LoadEarlier",
+  useThreadLoadEarlier,
+);

package/src/primitives/thread/ThreadMessages.tsx

@@ -5,4 +5,5 @@ export {
   ThreadPrimitiveMessagesImpl,
   ThreadPrimitiveMessageByIndex,
   ThreadPrimitiveUnstable_MessageById,
+  ThreadPrimitiveRow,
 } from "@assistant-ui/core/react";

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

@@ -357,6 +357,39 @@ describe("ThreadPrimitiveRoot", () => {
     expect(event.defaultPrevented).toBe(true);
   });
 
+  it.each([
+    ["during the composition", { isComposing: true }],
+    ["after compositionend (Safari)", { keyCode: 229 }],
+  ])(
+    "keeps speech playing when Escape ends an IME composition %s",
+    async (_, init) => {
+      const speech = createSpeechAdapter();
+      const runtimeRef: RuntimeRef = { current: null };
+      render(
+        <RuntimeProvider runtimeRef={runtimeRef} speech={speech.adapter}>
+          <ThreadPrimitiveRoot>
+            <ComposerPrimitiveInput data-testid="composer" />
+          </ThreadPrimitiveRoot>
+        </RuntimeProvider>,
+      );
+      startSpeaking(runtimeRef);
+      await waitFor(() => {
+        expect(runtimeRef.current!.thread.getState().speech).toBeDefined();
+      });
+
+      const event = new KeyboardEvent("keydown", {
+        key: "Escape",
+        bubbles: true,
+        cancelable: true,
+        ...init,
+      });
+      act(() => screen.getByTestId("composer").dispatchEvent(event));
+
+      expect(speech.cancel).not.toHaveBeenCalled();
+      expect(event.defaultPrevented).toBe(false);
+    },
+  );
+
   it("lets an idle composer fall through so Escape still stops speech", async () => {
     const speech = createSpeechAdapter();
     const runtimeRef: RuntimeRef = { current: null };

package/src/primitives/thread/ThreadRoot.tsx

@@ -1,6 +1,7 @@
 "use client";
 
 import { Primitive } from "../../utils/Primitive";
+import { isCompositionKey } from "../../utils/isCompositionKey";
 import {
   type ComponentRef,
   forwardRef,
@@ -64,6 +65,8 @@ export const ThreadPrimitiveRoot = forwardRef<
   useEffect(() => {
     const handleKeyDown = (event: KeyboardEvent) => {
       if (event.key !== "Escape") return;
+      // An Escape that ends an IME composition belongs to the input method.
+      if (isCompositionKey(event)) return;
       if (event.defaultPrevented || aui.thread.source === null) return;
       const eventThreadRoot =
         escapeEventThreadRoots.get(event) ??