@assistant-ui/react 0.12.9 → 0.12.10

package/src/primitives/attachment/AttachmentName.tsx

@@ -10,7 +10,7 @@ export namespace AttachmentPrimitiveName {
 export const AttachmentPrimitiveName: FC<
   AttachmentPrimitiveName.Props
 > = () => {
-  const name = useAuiState(({ attachment }) => attachment.name);
+  const name = useAuiState((s) => s.attachment.name);
   return <>{name}</>;
 };
 

package/src/primitives/attachment/AttachmentThumb.tsx

@@ -15,8 +15,8 @@ export const AttachmentPrimitiveThumb = forwardRef<
   AttachmentPrimitiveThumb.Element,
   AttachmentPrimitiveThumb.Props
 >((props, ref) => {
-  const ext = useAuiState(({ attachment }) => {
-    const parts = attachment.name.split(".");
+  const ext = useAuiState((s) => {
+    const parts = s.attachment.name.split(".");
     return parts.length > 1 ? parts.pop()! : "";
   });
   return (

package/src/primitives/branchPicker/BranchPickerCount.tsx

@@ -4,7 +4,7 @@ import type { FC } from "react";
 import { useAuiState } from "@assistant-ui/store";
 
 const useBranchPickerCount = () => {
-  const branchCount = useAuiState(({ message }) => message.branchCount);
+  const branchCount = useAuiState((s) => s.message.branchCount);
   return branchCount;
 };
 

package/src/primitives/branchPicker/BranchPickerNext.tsx

@@ -10,12 +10,12 @@ import { useAuiState, useAui } from "@assistant-ui/store";
 
 const useBranchPickerNext = () => {
   const aui = useAui();
-  const disabled = useAuiState(({ thread, message }) => {
+  const disabled = useAuiState((s) => {
     // Disabled if no next branch
-    if (message.branchNumber >= message.branchCount) return true;
+    if (s.message.branchNumber >= s.message.branchCount) return true;
 
     // Disabled if running and capability not supported
-    if (thread.isRunning && !thread.capabilities.switchBranchDuringRun) {
+    if (s.thread.isRunning && !s.thread.capabilities.switchBranchDuringRun) {
       return true;
     }
 

package/src/primitives/branchPicker/BranchPickerNumber.tsx

@@ -4,7 +4,7 @@ import type { FC } from "react";
 import { useAuiState } from "@assistant-ui/store";
 
 const useBranchPickerNumber = () => {
-  const branchNumber = useAuiState(({ message }) => message.branchNumber);
+  const branchNumber = useAuiState((s) => s.message.branchNumber);
   return branchNumber;
 };
 

package/src/primitives/branchPicker/BranchPickerPrevious.tsx

@@ -31,12 +31,12 @@ import { useAuiState, useAui } from "@assistant-ui/store";
  */
 const useBranchPickerPrevious = () => {
   const aui = useAui();
-  const disabled = useAuiState(({ thread, message }) => {
+  const disabled = useAuiState((s) => {
     // Disabled if no previous branch
-    if (message.branchNumber <= 1) return true;
+    if (s.message.branchNumber <= 1) return true;
 
     // Disabled if running and capability not supported
-    if (thread.isRunning && !thread.capabilities.switchBranchDuringRun) {
+    if (s.thread.isRunning && !s.thread.capabilities.switchBranchDuringRun) {
       return true;
     }
 

package/src/primitives/chainOfThought/ChainOfThoughtAccordionTrigger.tsx

@@ -10,9 +10,7 @@ import { useAuiState, useAui } from "@assistant-ui/store";
 
 const useChainOfThoughtAccordionTrigger = () => {
   const aui = useAui();
-  const collapsed = useAuiState(
-    ({ chainOfThought }) => chainOfThought.collapsed,
-  );
+  const collapsed = useAuiState((s) => s.chainOfThought.collapsed);
 
   const callback = useCallback(() => {
     aui.chainOfThought().setCollapsed(!collapsed);

package/src/primitives/chainOfThought/ChainOfThoughtParts.tsx

@@ -53,9 +53,7 @@ export namespace ChainOfThoughtPrimitiveParts {
 export const ChainOfThoughtPrimitiveParts: FC<
   ChainOfThoughtPrimitiveParts.Props
 > = ({ components }) => {
-  const partsLength = useAuiState(
-    ({ chainOfThought }) => chainOfThought.parts.length,
-  );
+  const partsLength = useAuiState((s) => s.chainOfThought.parts.length);
 
   const messageComponents = useMemo(
     () => ({

package/src/primitives/composer/ComposerAddAttachment.tsx

@@ -14,7 +14,7 @@ const useComposerAddAttachment = ({
   /** allow selecting multiple files */
   multiple?: boolean | undefined;
 } = {}) => {
-  const disabled = useAuiState(({ composer }) => !composer.isEditing);
+  const disabled = useAuiState((s) => !s.composer.isEditing);
   const aui = useAui();
 
   const callback = useCallback(() => {

package/src/primitives/composer/ComposerAttachments.tsx

@@ -39,7 +39,7 @@ const getComponent = (
 const AttachmentComponent: FC<{
   components: ComposerPrimitiveAttachments.Props["components"];
 }> = ({ components }) => {
-  const attachment = useAuiState(({ attachment }) => attachment);
+  const attachment = useAuiState((s) => s.attachment);
   if (!attachment) return null;
 
   const Component = getComponent(components, attachment);

package/src/primitives/composer/ComposerCancel.tsx

@@ -10,7 +10,7 @@ import { useAuiState, useAui } from "@assistant-ui/store";
 
 const useComposerCancel = () => {
   const aui = useAui();
-  const disabled = useAuiState(({ composer }) => !composer.canCancel);
+  const disabled = useAuiState((s) => !s.composer.canCancel);
 
   const callback = useCallback(() => {
     aui.composer().cancel();

package/src/primitives/composer/ComposerDictate.tsx

@@ -11,10 +11,10 @@ import { createActionButton } from "../../utils/createActionButton";
 const useComposerDictate = () => {
   const aui = useAui();
   const disabled = useAuiState(
-    ({ thread, composer }) =>
-      composer.dictation != null ||
-      !thread.capabilities.dictation ||
-      !composer.isEditing,
+    (s) =>
+      s.composer.dictation != null ||
+      !s.thread.capabilities.dictation ||
+      !s.composer.isEditing,
   );
 
   const callback = useCallback(() => {

package/src/primitives/composer/ComposerDictationTranscript.tsx

@@ -28,9 +28,7 @@ export const ComposerPrimitiveDictationTranscript = forwardRef<
   ComposerPrimitiveDictationTranscript.Element,
   ComposerPrimitiveDictationTranscript.Props
 >(({ children, ...props }, forwardRef) => {
-  const transcript = useAuiState(
-    ({ composer }) => composer.dictation?.transcript,
-  );
+  const transcript = useAuiState((s) => s.composer.dictation?.transcript);
 
   if (!transcript) return null;
 

package/src/primitives/composer/ComposerIf.tsx

@@ -14,11 +14,11 @@ type ComposerIfFilters = {
 export type UseComposerIfProps = RequireAtLeastOne<ComposerIfFilters>;
 
 const useComposerIf = (props: UseComposerIfProps) => {
-  return useAuiState(({ composer }) => {
-    if (props.editing === true && !composer.isEditing) return false;
-    if (props.editing === false && composer.isEditing) return false;
+  return useAuiState((s) => {
+    if (props.editing === true && !s.composer.isEditing) return false;
+    if (props.editing === false && s.composer.isEditing) return false;
 
-    const isDictating = composer.dictation != null;
+    const isDictating = s.composer.dictation != null;
     if (props.dictation === true && !isDictating) return false;
     if (props.dictation === false && isDictating) return false;
 

package/src/primitives/composer/ComposerInput.tsx

@@ -21,17 +21,13 @@ import { flushResourcesSync } from "@assistant-ui/tap";
 
 export namespace ComposerPrimitiveInput {
   export type Element = HTMLTextAreaElement;
-  export type Props = TextareaAutosizeProps & {
+
+  type BaseProps = {
     /**
      * Whether to render as a child component using Slot.
      * When true, the component will merge its props with its child.
      */
     asChild?: boolean | undefined;
-    /**
-     * Whether to submit the message when Enter is pressed (without Shift).
-     * @default true
-     */
-    submitOnEnter?: boolean | undefined;
     /**
      * Whether to cancel message composition when Escape is pressed.
      * @default true
@@ -58,6 +54,34 @@ export namespace ComposerPrimitiveInput {
      */
     addAttachmentOnPaste?: boolean | undefined;
   };
+
+  type SubmitModeProps =
+    | {
+        /**
+         * Controls how the Enter key submits messages.
+         * - "enter": Plain Enter submits (Shift+Enter for newline)
+         * - "ctrlEnter": Ctrl/Cmd+Enter submits (plain Enter for newline)
+         * - "none": Keyboard submission disabled
+         * @default "enter"
+         */
+        submitMode?: "enter" | "ctrlEnter" | "none" | undefined;
+        /**
+         * @deprecated Use `submitMode` instead
+         * @ignore
+         */
+        submitOnEnter?: never;
+      }
+    | {
+        submitMode?: never;
+        /**
+         * Whether to submit the message when Enter is pressed (without Shift).
+         * @default true
+         * @deprecated Use `submitMode` instead. Will be removed in a future version.
+         */
+        submitOnEnter?: boolean | undefined;
+      };
+
+  export type Props = TextareaAutosizeProps & BaseProps & SubmitModeProps;
 }
 
 /**
@@ -69,10 +93,16 @@ export namespace ComposerPrimitiveInput {
  *
  * @example
  * ```tsx
+ * // Ctrl/Cmd+Enter to submit (plain Enter inserts newline)
+ * <ComposerPrimitive.Input
+ *   placeholder="Type your message..."
+ *   submitMode="ctrlEnter"
+ * />
+ *
+ * // Old API (deprecated, still supported)
  * <ComposerPrimitive.Input
  *   placeholder="Type your message..."
  *   submitOnEnter={true}
- *   addAttachmentOnPaste={true}
  * />
  * ```
  */
@@ -88,7 +118,8 @@ export const ComposerPrimitiveInput = forwardRef<
       onChange,
       onKeyDown,
       onPaste,
-      submitOnEnter = true,
+      submitOnEnter,
+      submitMode,
       cancelOnEscape = true,
       unstable_focusOnRunStart = true,
       unstable_focusOnScrollToBottom = true,
@@ -100,17 +131,19 @@ export const ComposerPrimitiveInput = forwardRef<
   ) => {
     const aui = useAui();
 
-    const value = useAuiState(({ composer }) => {
-      if (!composer.isEditing) return "";
-      return composer.text;
+    const effectiveSubmitMode =
+      submitMode ?? (submitOnEnter === false ? "none" : "enter");
+
+    const value = useAuiState((s) => {
+      if (!s.composer.isEditing) return "";
+      return s.composer.text;
     });
 
     const Component = asChild ? Slot : TextareaAutosize;
 
     const isDisabled =
       useAuiState(
-        ({ thread, composer }) =>
-          thread.isDisabled || composer.dictation?.inputDisabled,
+        (s) => s.thread.isDisabled || s.composer.dictation?.inputDisabled,
       ) || disabledProp;
     const textareaRef = useRef<HTMLTextAreaElement>(null);
     const ref = useComposedRefs(forwardedRef, textareaRef);
@@ -129,17 +162,24 @@ export const ComposerPrimitiveInput = forwardRef<
     });
 
     const handleKeyPress = (e: KeyboardEvent) => {
-      if (isDisabled || !submitOnEnter) return;
+      if (isDisabled) return;
 
       // ignore IME composition events
       if (e.nativeEvent.isComposing) return;
 
-      if (e.key === "Enter" && e.shiftKey === false) {
+      if (e.key === "Enter" && !e.shiftKey) {
         const isRunning = aui.thread().getState().isRunning;
+        if (isRunning) return;
 
-        if (!isRunning) {
-          e.preventDefault();
+        let shouldSubmit = false;
+        if (effectiveSubmitMode === "ctrlEnter") {
+          shouldSubmit = e.ctrlKey || e.metaKey;
+        } else if (effectiveSubmitMode === "enter") {
+          shouldSubmit = true;
+        }
 
+        if (shouldSubmit) {
+          e.preventDefault();
           textareaRef.current?.closest("form")?.requestSubmit();
         }
       }

package/src/primitives/composer/ComposerQuote.tsx

@@ -0,0 +1,109 @@
+"use client";
+
+import { Primitive } from "@radix-ui/react-primitive";
+import {
+  type ComponentRef,
+  type ComponentPropsWithoutRef,
+  forwardRef,
+  useCallback,
+} from "react";
+import { useAui, useAuiState } from "@assistant-ui/store";
+import { composeEventHandlers } from "@radix-ui/primitive";
+
+// ---- Root ----
+
+export namespace ComposerPrimitiveQuote {
+  export type Element = ComponentRef<typeof Primitive.div>;
+  export type Props = ComponentPropsWithoutRef<typeof Primitive.div>;
+}
+
+/**
+ * Renders a container for the quoted text preview in the composer.
+ * Only renders when a quote is set.
+ *
+ * @example
+ * ```tsx
+ * <ComposerPrimitive.Quote>
+ *   <ComposerPrimitive.QuoteText />
+ *   <ComposerPrimitive.QuoteDismiss>×</ComposerPrimitive.QuoteDismiss>
+ * </ComposerPrimitive.Quote>
+ * ```
+ */
+export const ComposerPrimitiveQuote = forwardRef<
+  ComposerPrimitiveQuote.Element,
+  ComposerPrimitiveQuote.Props
+>((props, forwardedRef) => {
+  const quote = useAuiState((s) => s.composer.quote);
+  if (!quote) return null;
+
+  return <Primitive.div {...props} ref={forwardedRef} />;
+});
+
+ComposerPrimitiveQuote.displayName = "ComposerPrimitive.Quote";
+
+// ---- QuoteText ----
+
+export namespace ComposerPrimitiveQuoteText {
+  export type Element = ComponentRef<typeof Primitive.span>;
+  export type Props = ComponentPropsWithoutRef<typeof Primitive.span>;
+}
+
+/**
+ * Renders the quoted text content.
+ *
+ * @example
+ * ```tsx
+ * <ComposerPrimitive.QuoteText />
+ * ```
+ */
+export const ComposerPrimitiveQuoteText = forwardRef<
+  ComposerPrimitiveQuoteText.Element,
+  ComposerPrimitiveQuoteText.Props
+>(({ children, ...props }, forwardedRef) => {
+  const text = useAuiState((s) => s.composer.quote?.text);
+  if (!text) return null;
+
+  return (
+    <Primitive.span {...props} ref={forwardedRef}>
+      {children ?? text}
+    </Primitive.span>
+  );
+});
+
+ComposerPrimitiveQuoteText.displayName = "ComposerPrimitive.QuoteText";
+
+// ---- QuoteDismiss ----
+
+export namespace ComposerPrimitiveQuoteDismiss {
+  export type Element = ComponentRef<typeof Primitive.button>;
+  export type Props = ComponentPropsWithoutRef<typeof Primitive.button>;
+}
+
+/**
+ * A button that clears the current quote from the composer.
+ *
+ * @example
+ * ```tsx
+ * <ComposerPrimitive.QuoteDismiss>×</ComposerPrimitive.QuoteDismiss>
+ * ```
+ */
+export const ComposerPrimitiveQuoteDismiss = forwardRef<
+  ComposerPrimitiveQuoteDismiss.Element,
+  ComposerPrimitiveQuoteDismiss.Props
+>(({ onClick, ...props }, forwardedRef) => {
+  const aui = useAui();
+  const handleDismiss = useCallback(() => {
+    aui.composer().setQuote(undefined);
+  }, [aui]);
+
+  return (
+    <Primitive.button
+      type="button"
+      {...props}
+      ref={forwardedRef}
+      onClick={composeEventHandlers(onClick, handleDismiss)}
+    />
+  );
+});
+
+ComposerPrimitiveQuoteDismiss.displayName = "ComposerPrimitive.QuoteDismiss";

package/src/primitives/composer/ComposerStopDictation.tsx

@@ -11,7 +11,7 @@ import { useAuiState } from "@assistant-ui/store";
 
 const useComposerStopDictation = () => {
   const aui = useAui();
-  const isDictating = useAuiState(({ composer }) => composer.dictation != null);
+  const isDictating = useAuiState((s) => s.composer.dictation != null);
 
   const callback = useCallback(() => {
     aui.composer().stopDictation();

package/src/primitives/composer/index.ts

@@ -10,3 +10,6 @@ export { ComposerPrimitiveDictate as Dictate } from "./ComposerDictate";
 export { ComposerPrimitiveStopDictation as StopDictation } from "./ComposerStopDictation";
 export { ComposerPrimitiveDictationTranscript as DictationTranscript } from "./ComposerDictationTranscript";
 export { ComposerPrimitiveIf as If } from "./ComposerIf";
+export { ComposerPrimitiveQuote as Quote } from "./ComposerQuote";
+export { ComposerPrimitiveQuoteText as QuoteText } from "./ComposerQuote";
+export { ComposerPrimitiveQuoteDismiss as QuoteDismiss } from "./ComposerQuote";

package/src/primitives/error/ErrorMessage.tsx

@@ -13,10 +13,10 @@ export const ErrorPrimitiveMessage = forwardRef<
   ErrorPrimitiveMessage.Element,
   ErrorPrimitiveMessage.Props
 >(({ children, ...props }, forwardRef) => {
-  const error = useAuiState(({ message }) => {
-    return message.status?.type === "incomplete" &&
-      message.status.reason === "error"
-      ? message.status.error
+  const error = useAuiState((s) => {
+    return s.message.status?.type === "incomplete" &&
+      s.message.status.reason === "error"
+      ? s.message.status.error
       : undefined;
   });
 

package/src/primitives/index.ts

@@ -13,6 +13,7 @@ export * as SuggestionPrimitive from "./suggestion";
 export * as ThreadListPrimitive from "./threadList";
 export * as ThreadListItemPrimitive from "./threadListItem";
 export * as ThreadListItemMorePrimitive from "./threadListItemMore";
+export * as SelectionToolbarPrimitive from "./selectionToolbar";
 
 export { useMessagePartText } from "./messagePart/useMessagePartText";
 export { useMessagePartReasoning } from "./messagePart/useMessagePartReasoning";
@@ -22,3 +23,4 @@ export { useMessagePartImage } from "./messagePart/useMessagePartImage";
 export { useMessagePartData } from "./messagePart/useMessagePartData";
 export { useThreadViewportAutoScroll } from "./thread/useThreadViewportAutoScroll";
 export { useScrollLock } from "./reasoning";
+export { useMessageQuote } from "../hooks/useMessageQuote";

package/src/primitives/message/MessageAttachments.tsx

@@ -39,7 +39,7 @@ const getComponent = (
 const AttachmentComponent: FC<{
   components: MessagePrimitiveAttachments.Props["components"];
 }> = ({ components }) => {
-  const attachment = useAuiState(({ attachment }) => attachment);
+  const attachment = useAuiState((s) => s.attachment);
   if (!attachment) return null;
 
   const Component = getComponent(components, attachment as CompleteAttachment);
@@ -94,9 +94,9 @@ MessagePrimitiveAttachmentByIndex.displayName =
 export const MessagePrimitiveAttachments: FC<
   MessagePrimitiveAttachments.Props
 > = ({ components }) => {
-  const attachmentsCount = useAuiState(({ message }) => {
-    if (message.role !== "user") return 0;
-    return message.attachments.length;
+  const attachmentsCount = useAuiState((s) => {
+    if (s.message.role !== "user") return 0;
+    return s.message.attachments.length;
   });
 
   const attachmentElements = useMemo(() => {

package/src/primitives/message/MessageError.tsx

@@ -5,9 +5,9 @@ import { useAuiState } from "@assistant-ui/store";
 
 export const MessagePrimitiveError: FC<PropsWithChildren> = ({ children }) => {
   const hasError = useAuiState(
-    ({ message }) =>
-      message.status?.type === "incomplete" &&
-      message.status.reason === "error",
+    (s) =>
+      s.message.status?.type === "incomplete" &&
+      s.message.status.reason === "error",
   );
   return hasError ? children : null;
 };

package/src/primitives/message/MessageIf.tsx

@@ -20,7 +20,7 @@ type MessageIfFilters = {
 type UseMessageIfProps = RequireAtLeastOne<MessageIfFilters>;
 
 const useMessageIf = (props: UseMessageIfProps) => {
-  return useAuiState(({ message }) => {
+  return useAuiState((s) => {
     const {
       role,
       attachments,
@@ -30,7 +30,7 @@ const useMessageIf = (props: UseMessageIfProps) => {
       speech,
       isCopied,
       isHovering,
-    } = message;
+    } = s.message;
 
     if (props.hasBranches === true && branchCount < 2) return false;
 
@@ -64,7 +64,7 @@ const useMessageIf = (props: UseMessageIfProps) => {
 
     if (
       props.submittedFeedback !== undefined &&
-      (message.metadata.submittedFeedback?.type ?? null) !==
+      (s.message.metadata.submittedFeedback?.type ?? null) !==
         props.submittedFeedback
     )
       return false;

package/src/primitives/message/MessageParts.tsx

@@ -321,8 +321,8 @@ const ToolUIDisplay = ({
 }: {
   Fallback: ToolCallMessagePartComponent | undefined;
 } & ToolCallMessagePartProps) => {
-  const Render = useAuiState(({ tools }) => {
-    const Render = tools.tools[props.toolName] ?? Fallback;
+  const Render = useAuiState((s) => {
+    const Render = s.tools.tools[props.toolName] ?? Fallback;
     if (Array.isArray(Render)) return Render[0] ?? Fallback;
     return Render;
   });
@@ -364,7 +364,7 @@ export const MessagePartComponent: FC<MessagePartComponentProps> = ({
   } = {},
 }) => {
   const aui = useAui();
-  const part = useAuiState(({ part }) => part);
+  const part = useAuiState((s) => s.part);
 
   const type = part.type;
   if (type === "tool-call") {
@@ -503,11 +503,11 @@ const ConditionalEmptyImpl: FC<{
   components: MessagePrimitiveParts.Props["components"];
   enabled: boolean;
 }> = ({ components, enabled }) => {
-  const shouldShowEmpty = useAuiState(({ message }) => {
+  const shouldShowEmpty = useAuiState((s) => {
     if (!enabled) return false;
-    if (message.parts.length === 0) return false;
+    if (s.message.parts.length === 0) return false;
 
-    const lastPart = message.parts[message.parts.length - 1];
+    const lastPart = s.message.parts[s.message.parts.length - 1];
     return lastPart?.type !== "text" && lastPart?.type !== "reasoning";
   });
 
@@ -551,7 +551,7 @@ export const MessagePrimitiveParts: FC<MessagePrimitiveParts.Props> = ({
   components,
   unstable_showEmptyOnNonTextEnd = true,
 }) => {
-  const contentLength = useAuiState(({ message }) => message.parts.length);
+  const contentLength = useAuiState((s) => s.message.parts.length);
   const useChainOfThought = !!components?.ChainOfThought;
   const messageRanges = useMessagePartsGroups(useChainOfThought);
 

package/src/primitives/message/MessagePartsGrouped.tsx

@@ -74,7 +74,7 @@ const groupMessagePartsByParentId: GroupingFunction = (
 const useMessagePartsGrouped = (
   groupingFunction: GroupingFunction,
 ): MessagePartGroup[] => {
-  const parts = useAuiState(({ message }) => message.parts);
+  const parts = useAuiState((s) => s.message.parts);
 
   return useMemo(() => {
     if (parts.length === 0) {
@@ -219,8 +219,8 @@ const ToolUIDisplay = ({
 }: {
   Fallback: ToolCallMessagePartComponent | undefined;
 } & ToolCallMessagePartProps) => {
-  const Render = useAuiState(({ tools }) => {
-    const Render = tools.tools[props.toolName] ?? Fallback;
+  const Render = useAuiState((s) => {
+    const Render = s.tools.tools[props.toolName] ?? Fallback;
     if (Array.isArray(Render)) return Render[0] ?? Fallback;
     return Render;
   });
@@ -261,7 +261,7 @@ const MessagePartComponent: FC<MessagePartComponentProps> = ({
   } = {},
 }) => {
   const aui = useAui();
-  const part = useAuiState(({ part }) => part);
+  const part = useAuiState((s) => s.part);
 
   const type = part.type;
   if (type === "tool-call") {
@@ -426,7 +426,7 @@ const EmptyParts = memo(
 export const MessagePrimitiveUnstable_PartsGrouped: FC<
   MessagePrimitiveUnstable_PartsGrouped.Props
 > = ({ groupingFunction, components }) => {
-  const contentLength = useAuiState(({ message }) => message.parts.length);
+  const contentLength = useAuiState((s) => s.message.parts.length);
   const messageGroups = useMessagePartsGrouped(groupingFunction);
 
   const partsElements = useMemo(() => {

package/src/primitives/message/MessageRoot.tsx

@@ -60,11 +60,11 @@ const useMessageViewportRef = () => {
   // inset rules:
   // - the previous user message before the last assistant message registers its full height
   const shouldRegisterAsInset = useAuiState(
-    ({ thread, message }) =>
+    (s) =>
       turnAnchor === "top" &&
-      message.role === "user" &&
-      message.index === thread.messages.length - 2 &&
-      thread.messages.at(-1)?.role === "assistant",
+      s.message.role === "user" &&
+      s.message.index === s.thread.messages.length - 2 &&
+      s.thread.messages.at(-1)?.role === "assistant",
   );
 
   const getHeight = useCallback((el: HTMLElement) => el.offsetHeight, []);
@@ -116,10 +116,11 @@ export const MessagePrimitiveRoot = forwardRef<
     isHoveringRef,
     anchorUserMessageRef,
   );
+  const messageId = useAuiState((s) => s.message.id);
 
   return (
     <ThreadPrimitiveViewportSlack>
-      <Primitive.div {...props} ref={ref} />
+      <Primitive.div {...props} ref={ref} data-message-id={messageId} />
     </ThreadPrimitiveViewportSlack>
   );
 });

package/src/primitives/messagePart/MessagePartInProgress.tsx

@@ -11,9 +11,7 @@ export namespace MessagePartPrimitiveInProgress {
 export const MessagePartPrimitiveInProgress: FC<
   MessagePartPrimitiveInProgress.Props
 > = ({ children }) => {
-  const isInProgress = useAuiState(
-    ({ part }) => part.status.type === "running",
-  );
+  const isInProgress = useAuiState((s) => s.part.status.type === "running");
 
   return isInProgress ? children : null;
 };