@assistant-ui/react 0.12.9 → 0.12.11

package/src/model-context/{useAssistantTool.tsx → useAssistantTool.ts}

@@ -3,13 +3,12 @@
 import { useEffect } from "react";
 import { useAui } from "@assistant-ui/store";
 import type { ToolCallMessagePartComponent } from "../types/MessagePartComponentTypes";
-import type { Tool } from "assistant-stream";
+import type { AssistantToolProps as CoreAssistantToolProps } from "@assistant-ui/core";
 
 export type AssistantToolProps<
   TArgs extends Record<string, unknown>,
   TResult,
-> = Tool<TArgs, TResult> & {
-  toolName: string;
+> = CoreAssistantToolProps<TArgs, TResult> & {
   render?: ToolCallMessagePartComponent<TArgs, TResult> | undefined;
 };
 

package/src/primitives/actionBar/ActionBarCopy.tsx

@@ -37,15 +37,16 @@ const useActionBarPrimitiveCopy = ({
   copiedDuration?: number | undefined;
 } = {}) => {
   const aui = useAui();
-  const hasCopyableContent = useAuiState(({ message }) => {
+  const hasCopyableContent = useAuiState((s) => {
     return (
-      (message.role !== "assistant" || message.status?.type !== "running") &&
-      message.parts.some((c) => c.type === "text" && c.text.length > 0)
+      (s.message.role !== "assistant" ||
+        s.message.status?.type !== "running") &&
+      s.message.parts.some((c) => c.type === "text" && c.text.length > 0)
     );
   });
 
-  const isEditing = useAuiState(({ composer }) => composer.isEditing);
-  const composerValue = useAuiState(({ composer }) => composer.text);
+  const isEditing = useAuiState((s) => s.composer.isEditing);
+  const composerValue = useAuiState((s) => s.composer.text);
 
   const callback = useCallback(() => {
     const valueToCopy = isEditing ? composerValue : aui.message().getCopyText();
@@ -89,7 +90,7 @@ export const ActionBarPrimitiveCopy = forwardRef<
   ActionBarPrimitiveCopy.Element,
   ActionBarPrimitiveCopy.Props
 >(({ copiedDuration, onClick, disabled, ...props }, forwardedRef) => {
-  const isCopied = useAuiState(({ message }) => message.isCopied);
+  const isCopied = useAuiState((s) => s.message.isCopied);
   const callback = useActionBarPrimitiveCopy({ copiedDuration });
   return (
     <Primitive.button

package/src/primitives/actionBar/{ActionBarEdit.tsx → ActionBarEdit.ts}

@@ -31,7 +31,7 @@ import { useAuiState, useAui } from "@assistant-ui/store";
  */
 const useActionBarEdit = () => {
   const aui = useAui();
-  const disabled = useAuiState(({ composer }) => composer.isEditing);
+  const disabled = useAuiState((s) => s.composer.isEditing);
 
   const callback = useCallback(() => {
     aui.composer().beginEdit();

package/src/primitives/actionBar/ActionBarExportMarkdown.tsx

@@ -14,10 +14,11 @@ const useActionBarExportMarkdown = ({
   onExport?: ((content: string) => void | Promise<void>) | undefined;
 } = {}) => {
   const aui = useAui();
-  const hasExportableContent = useAuiState(({ message }) => {
+  const hasExportableContent = useAuiState((s) => {
     return (
-      (message.role !== "assistant" || message.status?.type !== "running") &&
-      message.parts.some((c) => c.type === "text" && c.text.length > 0)
+      (s.message.role !== "assistant" ||
+        s.message.status?.type !== "running") &&
+      s.message.parts.some((c) => c.type === "text" && c.text.length > 0)
     );
   });
 

package/src/primitives/actionBar/{ActionBarSpeak.tsx → ActionBarSpeak.ts}

@@ -14,10 +14,11 @@ const useActionBarSpeak = () => {
     aui.message().speak();
   }, [aui]);
 
-  const hasSpeakableContent = useAuiState(({ message }) => {
+  const hasSpeakableContent = useAuiState((s) => {
     return (
-      (message.role !== "assistant" || message.status?.type !== "running") &&
-      message.parts.some((c) => c.type === "text" && c.text.length > 0)
+      (s.message.role !== "assistant" ||
+        s.message.status?.type !== "running") &&
+      s.message.parts.some((c) => c.type === "text" && c.text.length > 0)
     );
   });
 

package/src/primitives/actionBar/ActionBarStopSpeaking.tsx

@@ -10,7 +10,7 @@ import { useAuiState, useAui } from "@assistant-ui/store";
 
 const useActionBarStopSpeaking = () => {
   const aui = useAui();
-  const isSpeaking = useAuiState(({ message }) => message.speech != null);
+  const isSpeaking = useAuiState((s) => s.message.speech != null);
 
   const callback = useCallback(() => {
     aui.message().stopSpeaking();

package/src/primitives/actionBar/{useActionBarFloatStatus.tsx → useActionBarFloatStatus.ts}

@@ -19,22 +19,22 @@ export const useActionBarFloatStatus = ({
   autohide,
   autohideFloat,
 }: UseActionBarFloatStatusProps) => {
-  return useAuiState(({ thread, message }) => {
-    if (hideWhenRunning && thread.isRunning) return HideAndFloatStatus.Hidden;
+  return useAuiState((s) => {
+    if (hideWhenRunning && s.thread.isRunning) return HideAndFloatStatus.Hidden;
 
     const autohideEnabled =
-      autohide === "always" || (autohide === "not-last" && !message.isLast);
+      autohide === "always" || (autohide === "not-last" && !s.message.isLast);
 
     // normal status
     if (!autohideEnabled) return HideAndFloatStatus.Normal;
 
     // hidden status
-    if (!message.isHovering) return HideAndFloatStatus.Hidden;
+    if (!s.message.isHovering) return HideAndFloatStatus.Hidden;
 
     // floating status
     if (
       autohideFloat === "always" ||
-      (autohideFloat === "single-branch" && message.branchCount <= 1)
+      (autohideFloat === "single-branch" && s.message.branchCount <= 1)
     )
       return HideAndFloatStatus.Floating;
 

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 → BranchPickerNext.ts}

@@ -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 → BranchPickerPrevious.ts}

@@ -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 → ChainOfThoughtAccordionTrigger.ts}

@@ -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 → ComposerAddAttachment.ts}

@@ -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

@@ -1,7 +1,7 @@
 "use client";
 
+import type { Attachment } from "@assistant-ui/core";
 import { ComponentType, type FC, memo, useMemo } from "react";
-import { Attachment } from "../../types";
 import { useAuiState } from "@assistant-ui/store";
 import { ComposerAttachmentByIndexProvider } from "../../context/providers";
 
@@ -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 → ComposerCancel.ts}

@@ -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 → ComposerDictate.ts}

@@ -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 → ComposerIf.ts}

@@ -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 → ComposerStopDictation.ts}

@@ -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,5 @@ export { useMessagePartImage } from "./messagePart/useMessagePartImage";
 export { useMessagePartData } from "./messagePart/useMessagePartData";
 export { useThreadViewportAutoScroll } from "./thread/useThreadViewportAutoScroll";
 export { useScrollLock } from "./reasoning";
+export { useMessageQuote } from "../hooks/useMessageQuote";
+export { useMessageTiming } from "../hooks/useMessageTiming";

package/src/primitives/message/MessageAttachments.tsx

@@ -1,9 +1,9 @@
 "use client";
 
+import type { CompleteAttachment } from "@assistant-ui/core";
 import { ComponentType, type FC, memo, useMemo } from "react";
 import { useAuiState } from "@assistant-ui/store";
 import { MessageAttachmentByIndexProvider } from "../../context/providers";
-import { CompleteAttachment } from "../../types";
 
 export namespace MessagePrimitiveAttachments {
   export type Props = {
@@ -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 → MessageError.ts}

@@ -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;
 };