@assistant-ui/react 0.15.18 → 0.15.19

package/src/primitives/assistantModal/AssistantModalRoot.tsx

@@ -1,9 +1,9 @@
 "use client";
 
-import { type FC, useEffect, useState } from "react";
+import { type FC, useState } from "react";
 import { Popover as PopoverPrimitive } from "radix-ui";
 import { type ScopedProps, usePopoverScope } from "./scope";
-import { useAui } from "@assistant-ui/store";
+import { useAuiEvent } from "@assistant-ui/store";
 
 export namespace AssistantModalPrimitiveRoot {
   export type Props = PopoverPrimitive.PopoverProps & {
@@ -11,54 +11,34 @@ export namespace AssistantModalPrimitiveRoot {
   };
 }
 
-const useAssistantModalOpenState = ({
-  defaultOpen = false,
-  unstable_openOnRunStart = true,
-}: {
-  defaultOpen?: boolean | undefined;
-  unstable_openOnRunStart?: boolean | undefined;
-}) => {
-  const state = useState(defaultOpen);
-
-  const [, setOpen] = state;
-  const aui = useAui();
-  useEffect(() => {
-    if (!unstable_openOnRunStart) return undefined;
-
-    return aui.on("thread.runStart", () => {
-      setOpen(true);
-    });
-  }, [unstable_openOnRunStart, aui, setOpen]);
-
-  return state;
-};
-
 export const AssistantModalPrimitiveRoot: FC<
   AssistantModalPrimitiveRoot.Props
 > = ({
   __scopeAssistantModal,
-  defaultOpen,
-  unstable_openOnRunStart,
+  defaultOpen = false,
+  unstable_openOnRunStart = true,
   open,
   onOpenChange,
   ...rest
 }: ScopedProps<AssistantModalPrimitiveRoot.Props>) => {
   const scope = usePopoverScope(__scopeAssistantModal);
 
-  const [modalOpen, setOpen] = useAssistantModalOpenState({
-    defaultOpen,
-    unstable_openOnRunStart,
-  });
+  const [modalOpen, setOpen] = useState(defaultOpen);
+  const isOpen = open ?? modalOpen;
 
   const openChangeHandler = (open: boolean) => {
     onOpenChange?.(open);
     setOpen(open);
   };
 
+  useAuiEvent("thread.runStart", () => {
+    if (unstable_openOnRunStart && !isOpen) openChangeHandler(true);
+  });
+
   return (
     <PopoverPrimitive.Root
       {...scope}
-      open={open === undefined ? modalOpen : open}
+      open={isOpen}
       onOpenChange={openChangeHandler}
       {...rest}
     />

package/src/primitives/assistantModal/AssistantModalTrigger.test.tsx

@@ -1,7 +1,6 @@
 import { renderToStaticMarkup } from "react-dom/server";
 import { describe, expect, it, vi } from "vitest";
 import { AssistantModalPrimitiveAnchor } from "./AssistantModalAnchor";
-import { AssistantModalPrimitiveContent } from "./AssistantModalContent";
 import { AssistantModalPrimitiveRoot } from "./AssistantModalRoot";
 import { AssistantModalPrimitiveTrigger } from "./AssistantModalTrigger";
 
@@ -9,9 +8,7 @@ vi.mock("@assistant-ui/store", async (importOriginal) => {
   const actual = await importOriginal<typeof import("@assistant-ui/store")>();
   return {
     ...actual,
-    useAui: () => ({
-      on: () => () => {},
-    }),
+    useAuiEvent: () => {},
   };
 });
 
@@ -65,16 +62,4 @@ describe("AssistantModalPrimitive render props", () => {
     expect(renderHtml).toContain('class="anchor"');
     expect(renderHtml).toContain("Anchor");
   });
-
-  it("accepts render on Trigger, Anchor, and Content", () => {
-    const tree = (
-      <AssistantModalPrimitiveRoot open>
-        <AssistantModalPrimitiveTrigger render={<button type="button" />} />
-        <AssistantModalPrimitiveAnchor render={<div />} />
-        <AssistantModalPrimitiveContent render={<div />} />
-      </AssistantModalPrimitiveRoot>
-    );
-
-    expect(tree).toBeTruthy();
-  });
 });

package/src/primitives/composer/ComposerAttachmentDropzone.test.tsx

@@ -74,6 +74,49 @@ describe("ComposerPrimitiveAttachmentDropzone", () => {
     vi.restoreAllMocks();
   });
 
+  it("composes a render element and keeps the drag props attached", async () => {
+    await act(async () => {
+      root.render(
+        <ComposerPrimitiveAttachmentDropzone
+          data-testid="dropzone"
+          render={<section className="child" />}
+          className="parent"
+        >
+          <div>outer</div>
+        </ComposerPrimitiveAttachmentDropzone>,
+      );
+    });
+
+    const dropzone = container.querySelector(
+      "section[data-testid='dropzone']",
+    ) as HTMLElement;
+    expect(dropzone).not.toBeNull();
+    expect(dropzone.textContent).toBe("outer");
+    expect(dropzone.className).toContain("parent");
+    expect(dropzone.className).toContain("child");
+
+    await act(async () => {
+      dropzone.dispatchEvent(createDragEvent("dragenter", ["Files"]));
+    });
+
+    expect(dropzone.getAttribute("data-dragging")).toBe("true");
+  });
+
+  it("falls back to the render element's own children", async () => {
+    await act(async () => {
+      root.render(
+        <ComposerPrimitiveAttachmentDropzone
+          data-testid="dropzone"
+          render={<section>fallback</section>}
+        />,
+      );
+    });
+
+    expect(
+      container.querySelector("section[data-testid='dropzone']")?.textContent,
+    ).toBe("fallback");
+  });
+
   it("starts all dropped attachments before awaiting completion", async () => {
     const resolvers: Array<() => void> = [];
     addAttachment.mockImplementation(

package/src/primitives/composer/ComposerAttachmentDropzone.tsx

@@ -5,7 +5,6 @@ import {
   useCallback,
   useState,
   type ReactElement,
-  cloneElement,
   isValidElement,
 } from "react";
 
@@ -13,6 +12,7 @@ import { composeEventHandlers } from "radix-ui/internal";
 import { Slot } from "radix-ui";
 import type React from "react";
 import { useAui } from "@assistant-ui/store";
+import { renderSlot } from "../../utils/Primitive";
 
 export namespace ComposerPrimitiveAttachmentDropzone {
   export type Element = HTMLDivElement;
@@ -116,15 +116,7 @@ export const ComposerPrimitiveAttachmentDropzone = forwardRef<
   };
 
   if (render && isValidElement(render)) {
-    const renderChildren =
-      children !== undefined
-        ? children
-        : (render.props as Record<string, unknown>).children;
-    return (
-      <Slot.Root {...mergedProps}>
-        {cloneElement(render, undefined, renderChildren as React.ReactNode)}
-      </Slot.Root>
-    );
+    return renderSlot(render, children, mergedProps);
   }
 
   const Comp = asChild ? Slot.Root : "div";

package/src/primitives/composer/ComposerInput.test.tsx

@@ -240,6 +240,28 @@ describe("ComposerPrimitiveInput", () => {
     return textarea;
   };
 
+  it("composes a render element with the computed input props", async () => {
+    await act(async () => {
+      root.render(
+        <form>
+          <ComposerPrimitiveInput
+            render={<textarea data-testid="custom" className="child" />}
+            className="parent"
+          />
+        </form>,
+      );
+    });
+
+    const textarea = container.querySelector(
+      "textarea[data-testid='custom']",
+    ) as HTMLTextAreaElement;
+    expect(textarea).not.toBeNull();
+    expect(textarea.className).toContain("parent");
+    expect(textarea.className).toContain("child");
+    expect(textarea.name).toBe("input");
+    expect(textarea.hasAttribute("render")).toBe(false);
+  });
+
   it("syncs setText during active composition so React 19 cannot reset the textarea", async () => {
     const textarea = await mount();
 

package/src/primitives/composer/ComposerInput.tsx

@@ -12,7 +12,6 @@ import {
   useCallback,
   useEffect,
   useRef,
-  cloneElement,
   isValidElement,
 } from "react";
 import TextareaAutosize, {
@@ -22,6 +21,7 @@ import TextareaAutosize, {
 import { useEscapeKeydown } from "radix-ui/internal";
 import { useOnScrollToBottom } from "../../utils/hooks/useOnScrollToBottom";
 import { useMediaQuery } from "../../utils/hooks/useMediaQuery";
+import { renderSlot } from "../../utils/Primitive";
 import { useAui } from "@assistant-ui/store";
 import { flushTapSync } from "@assistant-ui/tap";
 import { useComposerInputPluginRegistryOptional } from "./ComposerInputPluginContext";
@@ -434,14 +434,10 @@ export const ComposerPrimitiveInput = forwardRef<
     };
 
     if (render && isValidElement(render)) {
-      const renderChildren =
-        (rest as any).children !== undefined
-          ? ((rest as any).children as ReactNode)
-          : ((render.props as Record<string, unknown>).children as ReactNode);
-      return (
-        <Slot.Root {...inputProps}>
-          {cloneElement(render, undefined, renderChildren)}
-        </Slot.Root>
+      return renderSlot(
+        render,
+        (rest as { children?: ReactNode }).children,
+        inputProps,
       );
     }
 

package/src/primitives/composer/trigger/triggerKeyboardResource.ts

@@ -1,4 +1,4 @@
-import { useEffect, useEffectEvent, useState } from "react";
+import { useEffectEvent, useState } from "react";
 import { resource } from "@assistant-ui/tap";
 import type {
   Unstable_TriggerCategory,
@@ -58,14 +58,20 @@ const useTriggerKeyboardResource = ({
   close: () => void;
 }): TriggerKeyboardResourceOutput => {
   const [highlightedIndex, setHighlightedIndex] = useState(0);
+  const [highlightScope, setHighlightScope] = useState({
+    navigableList,
+    isSearchMode,
+    activeCategoryId,
+  });
 
-  useEffect(() => {
-    setHighlightedIndex(0);
-  }, [navigableList]);
-
-  useEffect(() => {
+  if (
+    highlightScope.navigableList !== navigableList ||
+    highlightScope.isSearchMode !== isSearchMode ||
+    highlightScope.activeCategoryId !== activeCategoryId
+  ) {
+    setHighlightScope({ navigableList, isSearchMode, activeCategoryId });
     setHighlightedIndex(0);
-  }, [isSearchMode, activeCategoryId]);
+  }
 
   const highlightIndex = useEffectEvent((index: number) => {
     if (index < 0 || index >= navigableList.length) return;

package/src/primitives/composer/trigger/triggerNavigationResource.ts

@@ -1,4 +1,4 @@
-import { useEffect, useEffectEvent, useMemo, useState } from "react";
+import { useEffectEvent, useMemo, useState } from "react";
 import { resource } from "@assistant-ui/tap";
 import type {
   Unstable_TriggerAdapter,
@@ -41,10 +41,12 @@ const useTriggerNavigationResource = ({
   open: boolean;
 }): TriggerNavigationResourceOutput => {
   const [activeCategoryId, setActiveCategoryId] = useState<string | null>(null);
+  const [wasOpen, setWasOpen] = useState(open);
 
-  useEffect(() => {
+  if (wasOpen !== open) {
+    setWasOpen(open);
     if (!open) setActiveCategoryId(null);
-  }, [open]);
+  }
 
   const categories = useMemo<readonly Unstable_TriggerCategory[]>(() => {
     if (!open || !adapter) return [];

package/src/primitives/message/MessageParts.tsx

@@ -43,29 +43,9 @@ export const MessagePrimitiveParts: FC<MessagePrimitiveParts.Props> = (
   const { components, ...rest } = props;
   const merged = components
     ? {
+        ...components,
         Text: components.Text ?? webDefaultComponents.Text,
         Image: components.Image ?? webDefaultComponents.Image,
-        Reasoning:
-          components.Reasoning ?? messagePartsDefaultComponents.Reasoning,
-        Source: components.Source ?? messagePartsDefaultComponents.Source,
-        File: components.File ?? messagePartsDefaultComponents.File,
-        Unstable_Audio:
-          components.Unstable_Audio ??
-          messagePartsDefaultComponents.Unstable_Audio,
-        ...("ChainOfThought" in components
-          ? { ChainOfThought: components.ChainOfThought }
-          : {
-              tools: components.tools,
-              data: components.data,
-              ToolGroup:
-                components.ToolGroup ?? messagePartsDefaultComponents.ToolGroup,
-              ReasoningGroup:
-                components.ReasoningGroup ??
-                messagePartsDefaultComponents.ReasoningGroup,
-            }),
-        Empty: components.Empty,
-        Quote: components.Quote,
-        generativeUI: components.generativeUI,
       }
     : webDefaultComponents;
 

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

@@ -0,0 +1,93 @@
+// @vitest-environment jsdom
+
+import { render, screen } from "@testing-library/react";
+import { describe, expect, it } from "vitest";
+import type { ThreadMessageLike } from "@assistant-ui/core";
+import {
+  AssistantRuntimeProvider,
+  useExternalStoreRuntime,
+} from "@assistant-ui/core/react";
+import { ThreadPrimitiveMessageByIndex } from "../thread/ThreadMessages";
+import { MessagePrimitiveUnstable_PartsGroupedByParentId } from "./MessagePartsGrouped";
+
+const Message = () => (
+  <MessagePrimitiveUnstable_PartsGroupedByParentId
+    components={{
+      Text: ({ text }) => <span>{text}</span>,
+      Group: ({ groupKey, indices, children }) => (
+        <section
+          data-testid="group"
+          data-parent={groupKey}
+          data-indices={indices.join(",")}
+        >
+          {children}
+        </section>
+      ),
+    }}
+  />
+);
+
+const Example = ({ content }: { content: ThreadMessageLike["content"] }) => {
+  const messages: ThreadMessageLike[] = [
+    { id: "message", role: "assistant", content },
+  ];
+  const runtime = useExternalStoreRuntime({
+    messages,
+    convertMessage: (message) => message,
+    onNew: async () => {},
+  });
+  return (
+    <AssistantRuntimeProvider runtime={runtime}>
+      <ThreadPrimitiveMessageByIndex index={0} components={{ Message }} />
+    </AssistantRuntimeProvider>
+  );
+};
+
+describe("MessagePrimitive.Unstable_PartsGroupedByParentId", () => {
+  it("keeps parent IDs separate from ungrouped parts across content updates", () => {
+    const { rerender } = render(
+      <Example
+        content={[
+          { type: "text", text: "standalone" },
+          { type: "text", text: "child", parentId: "__ungrouped_0" },
+        ]}
+      />,
+    );
+    expect(
+      screen.getAllByTestId("group").map((group) => ({
+        parent: group.getAttribute("data-parent"),
+        indices: group.getAttribute("data-indices"),
+        text: group.textContent,
+      })),
+    ).toEqual([
+      { parent: null, indices: "0", text: "standalone" },
+      { parent: "__ungrouped_0", indices: "1", text: "child" },
+    ]);
+
+    rerender(
+      <Example
+        content={[
+          { type: "text", text: "first", parentId: "__ungrouped_parent" },
+          { type: "text", text: "standalone" },
+          { type: "text", text: "last", parentId: "__ungrouped_parent" },
+          { type: "text", text: "numeric", parentId: "1" },
+          { type: "text", text: "empty", parentId: "" },
+          { type: "text", text: "trailing" },
+        ]}
+      />,
+    );
+    expect(
+      screen.getAllByTestId("group").map((group) => ({
+        parent: group.getAttribute("data-parent"),
+        indices: group.getAttribute("data-indices"),
+        text: group.textContent,
+      })),
+    ).toEqual([
+      { parent: "__ungrouped_parent", indices: "0,2", text: "firstlast" },
+      { parent: null, indices: "1", text: "standalone" },
+      { parent: "1", indices: "3", text: "numeric" },
+      { parent: "", indices: "4", text: "empty" },
+      { parent: null, indices: "5", text: "trailing" },
+    ]);
+  });
+});

package/src/primitives/message/MessagePartsGrouped.tsx

@@ -44,7 +44,7 @@ const groupMessagePartsByParentId: GroupingFunction = (
   parts: readonly any[],
 ): MessagePartGroup[] => {
   // Map maintains insertion order, so groups appear in order of first occurrence
-  const groupMap = new Map<string, number[]>();
+  const groupMap = new Map<string | number, number[]>();
 
   // Process each part in order
   for (let i = 0; i < parts.length; i++) {
@@ -52,7 +52,7 @@ const groupMessagePartsByParentId: GroupingFunction = (
     const parentId = part?.parentId as string | undefined;
 
     // For parts without parentId, assign a unique group ID to maintain their position
-    const groupId = parentId ?? `__ungrouped_${i}`;
+    const groupId = parentId ?? i;
 
     // Get or create the indices array for this group
     const indices = groupMap.get(groupId) ?? [];
@@ -64,7 +64,7 @@ const groupMessagePartsByParentId: GroupingFunction = (
   const groups: MessagePartGroup[] = [];
   for (const [groupId, indices] of groupMap) {
     // Extract parentId (undefined for ungrouped parts)
-    const groupKey = groupId.startsWith("__ungrouped_") ? undefined : groupId;
+    const groupKey = typeof groupId === "string" ? groupId : undefined;
     groups.push({ groupKey, indices });
   }
 
@@ -94,15 +94,15 @@ export namespace MessagePrimitiveUnstable_PartsGrouped {
      * ```tsx
      * // Group by parent ID (default behavior)
      * groupingFunction={(parts) => {
-     *   const groups = new Map<string, number[]>();
+     *   const groups = new Map<string | number, number[]>();
      *   parts.forEach((part, i) => {
-     *     const key = part.parentId ?? `__ungrouped_${i}`;
+     *     const key = part.parentId ?? i;
      *     const indices = groups.get(key) ?? [];
      *     indices.push(i);
      *     groups.set(key, indices);
      *   });
      *   return Array.from(groups.entries()).map(([key, indices]) => ({
-     *     key: key.startsWith("__ungrouped_") ? undefined : key,
+     *     groupKey: typeof key === "string" ? key : undefined,
      *     indices
      *   }));
      * }}

package/src/primitives/messagePart/MessagePartText.test.tsx

@@ -0,0 +1,69 @@
+// @vitest-environment jsdom
+
+import { createRef } from "react";
+import { render, screen } from "@testing-library/react";
+import { describe, expect, it } from "vitest";
+import { TextMessagePartProvider } from "../../context/providers/TextMessagePartProvider";
+import { MessagePartPrimitiveText } from "./MessagePartText";
+
+describe("MessagePartPrimitive.Text", () => {
+  it("renders text in the supplied element and composes its props and refs", () => {
+    const ref = createRef<HTMLSpanElement>();
+    const targetRef = createRef<HTMLElement>();
+    const { rerender } = render(
+      <TextMessagePartProvider text="Hello">
+        <MessagePartPrimitiveText
+          smooth={false}
+          render={<mark ref={targetRef} className="target" />}
+          className="text"
+          ref={ref}
+        />
+      </TextMessagePartProvider>,
+    );
+
+    const text = screen.getByText("Hello");
+    expect(text.tagName).toBe("MARK");
+    expect(text.classList.contains("target")).toBe(true);
+    expect(text.classList.contains("text")).toBe(true);
+    expect(text.getAttribute("data-status")).toBe("complete");
+    expect(text.hasAttribute("render")).toBe(false);
+    expect(ref.current).toBe(text);
+    expect(targetRef.current).toBe(text);
+
+    rerender(
+      <TextMessagePartProvider text="Hello again" isRunning>
+        <MessagePartPrimitiveText smooth={false} render={<mark />} />
+      </TextMessagePartProvider>,
+    );
+
+    const updated = screen.getByText("Hello again");
+    expect(updated.tagName).toBe("MARK");
+    expect(updated.getAttribute("data-status")).toBe("running");
+  });
+
+  it("keeps the default span and explicit component behavior", () => {
+    const spanRef = createRef<HTMLSpanElement>();
+    const componentRef = createRef<HTMLElement>();
+    render(
+      <TextMessagePartProvider text="Hello">
+        <MessagePartPrimitiveText ref={spanRef} />
+        <MessagePartPrimitiveText component="p" ref={componentRef} />
+        <MessagePartPrimitiveText component="p" render={<mark />} />
+        <MessagePartPrimitiveText component="p" render={undefined} />
+      </TextMessagePartProvider>,
+    );
+
+    expect(spanRef.current?.tagName).toBe("SPAN");
+    expect(spanRef.current?.getAttribute("data-status")).toBe("complete");
+    expect(componentRef.current?.tagName).toBe("P");
+    expect(screen.getAllByText("Hello").map((text) => text.tagName)).toEqual([
+      "SPAN",
+      "P",
+      "MARK",
+      "P",
+    ]);
+    expect(
+      screen.getByText("Hello", { selector: "mark" }).hasAttribute("render"),
+    ).toBe(false);
+  });
+});

package/src/primitives/messagePart/MessagePartText.tsx

@@ -1,11 +1,12 @@
 "use client";
 
-import type { Primitive } from "../../utils/Primitive";
+import { Primitive, renderSlot } from "../../utils/Primitive";
 import {
   type ComponentRef,
   forwardRef,
   type ComponentPropsWithoutRef,
   type ElementType,
+  isValidElement,
 } from "react";
 import { useMessagePartText } from "./useMessagePartText";
 import { useSmooth, type SmoothOptions } from "../../utils/smooth/useSmooth";
@@ -26,6 +27,7 @@ export namespace MessagePartPrimitiveText {
     smooth?: boolean | SmoothOptions;
     /**
      * The HTML element or React component to render as.
+     * Ignored when a valid `render` element is supplied.
      * @default "span"
      */
     component?: ElementType;
@@ -51,14 +53,25 @@ export namespace MessagePartPrimitiveText {
 export const MessagePartPrimitiveText = forwardRef<
   MessagePartPrimitiveText.Element,
   MessagePartPrimitiveText.Props
->(({ smooth = true, component: Component = "span", ...rest }, forwardedRef) => {
-  const { text, status } = useSmooth(useMessagePartText(), smooth);
+>(
+  (
+    { smooth = true, component: Component = Primitive.span, render, ...rest },
+    forwardedRef,
+  ) => {
+    const { text, status } = useSmooth(useMessagePartText(), smooth);
 
-  return (
-    <Component data-status={status.type} {...rest} ref={forwardedRef}>
-      {text}
-    </Component>
-  );
-});
+    const mergedProps = {
+      "data-status": status.type,
+      ...rest,
+      ref: forwardedRef,
+    };
+
+    if (render && isValidElement(render)) {
+      return renderSlot(render, text, mergedProps);
+    }
+
+    return <Component {...mergedProps}>{text}</Component>;
+  },
+);
 
 MessagePartPrimitiveText.displayName = "MessagePartPrimitive.Text";