@assistant-ui/react 0.11.58 → 0.12.0

package/src/primitives/composer/ComposerAddAttachment.tsx

@@ -6,7 +6,7 @@ import {
   createActionButton,
 } from "../../utils/createActionButton";
 import { useCallback } from "react";
-import { useAssistantState, useAssistantApi } from "../../context";
+import { useAuiState, useAui } from "@assistant-ui/store";
 
 const useComposerAddAttachment = ({
   multiple = true,
@@ -14,8 +14,8 @@ const useComposerAddAttachment = ({
   /** allow selecting multiple files */
   multiple?: boolean | undefined;
 } = {}) => {
-  const disabled = useAssistantState(({ composer }) => !composer.isEditing);
-  const api = useAssistantApi();
+  const disabled = useAuiState(({ composer }) => !composer.isEditing);
+  const aui = useAui();
 
   const callback = useCallback(() => {
     const input = document.createElement("input");
@@ -23,7 +23,7 @@ const useComposerAddAttachment = ({
     input.multiple = multiple;
     input.hidden = true;
 
-    const attachmentAccept = api.composer().getState().attachmentAccept;
+    const attachmentAccept = aui.composer().getState().attachmentAccept;
     if (attachmentAccept !== "*") {
       input.accept = attachmentAccept;
     }
@@ -34,7 +34,7 @@ const useComposerAddAttachment = ({
       const fileList = (e.target as HTMLInputElement).files;
       if (!fileList) return;
       for (const file of fileList) {
-        api.composer().addAttachment(file);
+        aui.composer().addAttachment(file);
       }
 
       document.body.removeChild(input);
@@ -47,7 +47,7 @@ const useComposerAddAttachment = ({
     };
 
     input.click();
-  }, [api, multiple]);
+  }, [aui, multiple]);
 
   if (disabled) return null;
   return callback;

package/src/primitives/composer/ComposerAttachmentDropzone.tsx

@@ -4,7 +4,7 @@ import { forwardRef, useCallback, useState } from "react";
 
 import { Slot } from "@radix-ui/react-slot";
 import React from "react";
-import { useAssistantApi } from "../../context";
+import { useAui } from "@assistant-ui/store";
 
 export namespace ComposerPrimitiveAttachmentDropzone {
   export type Element = HTMLDivElement;
@@ -19,7 +19,7 @@ export const ComposerPrimitiveAttachmentDropzone = forwardRef<
   ComposerPrimitiveAttachmentDropzone.Props
 >(({ disabled, asChild = false, children, ...rest }, ref) => {
   const [isDragging, setIsDragging] = useState(false);
-  const api = useAssistantApi();
+  const aui = useAui();
 
   const handleDragEnterCapture = useCallback(
     (e: React.DragEvent) => {
@@ -59,13 +59,13 @@ export const ComposerPrimitiveAttachmentDropzone = forwardRef<
       setIsDragging(false);
       for (const file of e.dataTransfer.files) {
         try {
-          await api.composer().addAttachment(file);
+          await aui.composer().addAttachment(file);
         } catch (error) {
           console.error("Failed to add attachment:", error);
         }
       }
     },
-    [disabled, api],
+    [disabled, aui],
   );
 
   const dragProps = {

package/src/primitives/composer/ComposerAttachments.tsx

@@ -2,10 +2,8 @@
 
 import { ComponentType, type FC, memo, useMemo } from "react";
 import { Attachment } from "../../types";
-import {
-  useAssistantState,
-  ComposerAttachmentByIndexProvider,
-} from "../../context";
+import { useAuiState } from "@assistant-ui/store";
+import { ComposerAttachmentByIndexProvider } from "../../context/providers";
 
 export namespace ComposerPrimitiveAttachments {
   export type Props = {
@@ -41,7 +39,7 @@ const getComponent = (
 const AttachmentComponent: FC<{
   components: ComposerPrimitiveAttachments.Props["components"];
 }> = ({ components }) => {
-  const attachment = useAssistantState(({ attachment }) => attachment);
+  const attachment = useAuiState(({ attachment }) => attachment);
   if (!attachment) return null;
 
   const Component = getComponent(components, attachment);
@@ -96,9 +94,7 @@ ComposerPrimitiveAttachmentByIndex.displayName =
 export const ComposerPrimitiveAttachments: FC<
   ComposerPrimitiveAttachments.Props
 > = ({ components }) => {
-  const attachmentsCount = useAssistantState(
-    (s) => s.composer.attachments.length,
-  );
+  const attachmentsCount = useAuiState((s) => s.composer.attachments.length);
 
   const attachmentElements = useMemo(() => {
     return Array.from({ length: attachmentsCount }, (_, index) => (

package/src/primitives/composer/ComposerCancel.tsx

@@ -6,15 +6,15 @@ import {
   createActionButton,
 } from "../../utils/createActionButton";
 import { useCallback } from "react";
-import { useAssistantState, useAssistantApi } from "../../context";
+import { useAuiState, useAui } from "@assistant-ui/store";
 
 const useComposerCancel = () => {
-  const api = useAssistantApi();
-  const disabled = useAssistantState(({ composer }) => !composer.canCancel);
+  const aui = useAui();
+  const disabled = useAuiState(({ composer }) => !composer.canCancel);
 
   const callback = useCallback(() => {
-    api.composer().cancel();
-  }, [api]);
+    aui.composer().cancel();
+  }, [aui]);
 
   if (disabled) return null;
   return callback;

package/src/primitives/composer/ComposerDictate.tsx

@@ -1,7 +1,7 @@
 "use client";
 
 import { useCallback } from "react";
-import { useAssistantState, useAssistantApi } from "../../context";
+import { useAui, useAuiState } from "@assistant-ui/store";
 import type {
   ActionButtonElement,
   ActionButtonProps,
@@ -9,8 +9,8 @@ import type {
 import { createActionButton } from "../../utils/createActionButton";
 
 const useComposerDictate = () => {
-  const api = useAssistantApi();
-  const disabled = useAssistantState(
+  const aui = useAui();
+  const disabled = useAuiState(
     ({ thread, composer }) =>
       composer.dictation != null ||
       !thread.capabilities.dictation ||
@@ -18,8 +18,8 @@ const useComposerDictate = () => {
   );
 
   const callback = useCallback(() => {
-    api.composer().startDictation();
-  }, [api]);
+    aui.composer().startDictation();
+  }, [aui]);
 
   if (disabled) return null;
   return callback;

package/src/primitives/composer/ComposerDictationTranscript.tsx

@@ -2,7 +2,7 @@
 
 import { Primitive } from "@radix-ui/react-primitive";
 import { type ComponentRef, forwardRef, ComponentPropsWithoutRef } from "react";
-import { useAssistantState } from "../../context";
+import { useAuiState } from "@assistant-ui/store";
 
 export namespace ComposerPrimitiveDictationTranscript {
   export type Element = ComponentRef<typeof Primitive.span>;
@@ -28,7 +28,7 @@ export const ComposerPrimitiveDictationTranscript = forwardRef<
   ComposerPrimitiveDictationTranscript.Element,
   ComposerPrimitiveDictationTranscript.Props
 >(({ children, ...props }, forwardRef) => {
-  const transcript = useAssistantState(
+  const transcript = useAuiState(
     ({ composer }) => composer.dictation?.transcript,
   );
 

package/src/primitives/composer/ComposerIf.tsx

@@ -1,7 +1,7 @@
 "use client";
 
 import type { FC, PropsWithChildren } from "react";
-import { useAssistantState } from "../../context";
+import { useAuiState } from "@assistant-ui/store";
 import type { RequireAtLeastOne } from "../../utils/RequireAtLeastOne";
 
 type ComposerIfFilters = {
@@ -14,7 +14,7 @@ type ComposerIfFilters = {
 export type UseComposerIfProps = RequireAtLeastOne<ComposerIfFilters>;
 
 const useComposerIf = (props: UseComposerIfProps) => {
-  return useAssistantState(({ composer }) => {
+  return useAuiState(({ composer }) => {
     if (props.editing === true && !composer.isEditing) return false;
     if (props.editing === false && composer.isEditing) return false;
 
@@ -31,7 +31,7 @@ export namespace ComposerPrimitiveIf {
 }
 
 /**
- * @deprecated Use `<AssistantIf condition={({ composer }) => ...} />` instead.
+ * @deprecated Use `<AuiIf condition={({ composer }) => ...} />` instead.
  */
 export const ComposerPrimitiveIf: FC<ComposerPrimitiveIf.Props> = ({
   children,

package/src/primitives/composer/ComposerInput.tsx

@@ -16,8 +16,8 @@ import TextareaAutosize, {
 } from "react-textarea-autosize";
 import { useEscapeKeydown } from "@radix-ui/react-use-escape-keydown";
 import { useOnScrollToBottom } from "../../utils/hooks/useOnScrollToBottom";
-import { useAssistantState, useAssistantApi } from "../../context";
-import { flushSync } from "@assistant-ui/tap";
+import { useAuiState, useAui } from "@assistant-ui/store";
+import { flushResourcesSync } from "@assistant-ui/tap";
 
 export namespace ComposerPrimitiveInput {
   export type Element = HTMLTextAreaElement;
@@ -98,9 +98,9 @@ export const ComposerPrimitiveInput = forwardRef<
     },
     forwardedRef,
   ) => {
-    const api = useAssistantApi();
+    const aui = useAui();
 
-    const value = useAssistantState(({ composer }) => {
+    const value = useAuiState(({ composer }) => {
       if (!composer.isEditing) return "";
       return composer.text;
     });
@@ -108,7 +108,7 @@ export const ComposerPrimitiveInput = forwardRef<
     const Component = asChild ? Slot : TextareaAutosize;
 
     const isDisabled =
-      useAssistantState(
+      useAuiState(
         ({ thread, composer }) =>
           thread.isDisabled || composer.dictation?.inputDisabled,
       ) || disabledProp;
@@ -121,7 +121,7 @@ export const ComposerPrimitiveInput = forwardRef<
       // Only handle ESC if it originated from within this input
       if (!textareaRef.current?.contains(e.target as Node)) return;
 
-      const composer = api.composer();
+      const composer = aui.composer();
       if (composer.getState().canCancel) {
         composer.cancel();
         e.preventDefault();
@@ -135,7 +135,7 @@ export const ComposerPrimitiveInput = forwardRef<
       if (e.nativeEvent.isComposing) return;
 
       if (e.key === "Enter" && e.shiftKey === false) {
-        const isRunning = api.thread().getState().isRunning;
+        const isRunning = aui.thread().getState().isRunning;
 
         if (!isRunning) {
           e.preventDefault();
@@ -147,14 +147,14 @@ export const ComposerPrimitiveInput = forwardRef<
 
     const handlePaste = async (e: ClipboardEvent<HTMLTextAreaElement>) => {
       if (!addAttachmentOnPaste) return;
-      const threadCapabilities = api.thread().getState().capabilities;
+      const threadCapabilities = aui.thread().getState().capabilities;
       const files = Array.from(e.clipboardData?.files || []);
 
       if (threadCapabilities.attachments && files.length > 0) {
         try {
           e.preventDefault();
           await Promise.all(
-            files.map((file) => api.composer().addAttachment(file)),
+            files.map((file) => aui.composer().addAttachment(file)),
           );
         } catch (error) {
           console.error("Error adding attachment:", error);
@@ -175,7 +175,7 @@ export const ComposerPrimitiveInput = forwardRef<
 
     useOnScrollToBottom(() => {
       if (
-        api.composer().getState().type === "thread" &&
+        aui.composer().getState().type === "thread" &&
         unstable_focusOnScrollToBottom
       ) {
         focus();
@@ -184,23 +184,23 @@ export const ComposerPrimitiveInput = forwardRef<
 
     useEffect(() => {
       if (
-        api.composer().getState().type !== "thread" ||
+        aui.composer().getState().type !== "thread" ||
         !unstable_focusOnRunStart
       )
         return undefined;
 
-      return api.on("thread.run-start", focus);
-    }, [unstable_focusOnRunStart, focus, api]);
+      return aui.on("thread.runStart", focus);
+    }, [unstable_focusOnRunStart, focus, aui]);
 
     useEffect(() => {
       if (
-        api.composer().getState().type !== "thread" ||
+        aui.composer().getState().type !== "thread" ||
         !unstable_focusOnThreadSwitched
       )
         return undefined;
 
-      return api.on("thread-list-item.switched-to", focus);
-    }, [unstable_focusOnThreadSwitched, focus, api]);
+      return aui.on("threadListItem.switchedTo", focus);
+    }, [unstable_focusOnThreadSwitched, focus, aui]);
 
     return (
       <Component
@@ -210,9 +210,9 @@ export const ComposerPrimitiveInput = forwardRef<
         ref={ref as React.ForwardedRef<HTMLTextAreaElement>}
         disabled={isDisabled}
         onChange={composeEventHandlers(onChange, (e) => {
-          if (!api.composer().getState().isEditing) return;
-          flushSync(() => {
-            api.composer().setText(e.target.value);
+          if (!aui.composer().getState().isEditing) return;
+          flushResourcesSync(() => {
+            aui.composer().setText(e.target.value);
           });
         })}
         onKeyDown={composeEventHandlers(onKeyDown, handleKeyPress)}

package/src/primitives/composer/ComposerSend.tsx

@@ -6,18 +6,18 @@ import {
   createActionButton,
 } from "../../utils/createActionButton";
 import { useCallback } from "react";
-import { useAssistantState, useAssistantApi } from "../../context";
+import { useAuiState, useAui } from "@assistant-ui/store";
 
 export const useComposerSend = () => {
-  const api = useAssistantApi();
+  const aui = useAui();
 
-  const disabled = useAssistantState(
+  const disabled = useAuiState(
     (s) => s.thread.isRunning || !s.composer.isEditing || s.composer.isEmpty,
   );
 
   const callback = useCallback(() => {
-    api.composer().send();
-  }, [api]);
+    aui.composer().send();
+  }, [aui]);
 
   if (disabled) return null;
   return callback;

package/src/primitives/composer/ComposerStopDictation.tsx

@@ -1,22 +1,21 @@
 "use client";
 
 import { useCallback } from "react";
-import { useAssistantState, useAssistantApi } from "../../context";
+import { useAui } from "@assistant-ui/store";
 import type {
   ActionButtonElement,
   ActionButtonProps,
 } from "../../utils/createActionButton";
 import { createActionButton } from "../../utils/createActionButton";
+import { useAuiState } from "@assistant-ui/store";
 
 const useComposerStopDictation = () => {
-  const api = useAssistantApi();
-  const isDictating = useAssistantState(
-    ({ composer }) => composer.dictation != null,
-  );
+  const aui = useAui();
+  const isDictating = useAuiState(({ composer }) => composer.dictation != null);
 
   const callback = useCallback(() => {
-    api.composer().stopDictation();
-  }, [api]);
+    aui.composer().stopDictation();
+  }, [aui]);
 
   if (!isDictating) return null;
   return callback;

package/src/primitives/error/ErrorMessage.tsx

@@ -2,7 +2,7 @@
 
 import { Primitive } from "@radix-ui/react-primitive";
 import { type ComponentRef, forwardRef, ComponentPropsWithoutRef } from "react";
-import { useAssistantState } from "../../context";
+import { useAuiState } from "@assistant-ui/store";
 
 export namespace ErrorPrimitiveMessage {
   export type Element = ComponentRef<typeof Primitive.span>;
@@ -13,7 +13,7 @@ export const ErrorPrimitiveMessage = forwardRef<
   ErrorPrimitiveMessage.Element,
   ErrorPrimitiveMessage.Props
 >(({ children, ...props }, forwardRef) => {
-  const error = useAssistantState(({ message }) => {
+  const error = useAuiState(({ message }) => {
     return message.status?.type === "incomplete" &&
       message.status.reason === "error"
       ? message.status.error

package/src/primitives/index.ts

@@ -12,8 +12,6 @@ export * as ThreadListPrimitive from "./threadList";
 export * as ThreadListItemPrimitive from "./threadListItem";
 export * as ThreadListItemMorePrimitive from "./threadListItemMore";
 
-export { AssistantIf } from "./assistant/AssistantIf";
-
 export { useMessagePartText } from "./messagePart/useMessagePartText";
 export { useMessagePartReasoning } from "./messagePart/useMessagePartReasoning";
 export { useMessagePartSource } from "./messagePart/useMessagePartSource";

package/src/primitives/message/MessageAttachments.tsx

@@ -1,10 +1,8 @@
 "use client";
 
 import { ComponentType, type FC, memo, useMemo } from "react";
-import {
-  useAssistantState,
-  MessageAttachmentByIndexProvider,
-} from "../../context";
+import { useAuiState } from "@assistant-ui/store";
+import { MessageAttachmentByIndexProvider } from "../../context/providers";
 import { CompleteAttachment } from "../../types";
 
 export namespace MessagePrimitiveAttachments {
@@ -41,7 +39,7 @@ const getComponent = (
 const AttachmentComponent: FC<{
   components: MessagePrimitiveAttachments.Props["components"];
 }> = ({ components }) => {
-  const attachment = useAssistantState(({ attachment }) => attachment);
+  const attachment = useAuiState(({ attachment }) => attachment);
   if (!attachment) return null;
 
   const Component = getComponent(components, attachment as CompleteAttachment);
@@ -96,7 +94,7 @@ MessagePrimitiveAttachmentByIndex.displayName =
 export const MessagePrimitiveAttachments: FC<
   MessagePrimitiveAttachments.Props
 > = ({ components }) => {
-  const attachmentsCount = useAssistantState(({ message }) => {
+  const attachmentsCount = useAuiState(({ message }) => {
     if (message.role !== "user") return 0;
     return message.attachments.length;
   });

package/src/primitives/message/MessageError.tsx

@@ -1,10 +1,10 @@
 "use client";
 
 import { FC, PropsWithChildren } from "react";
-import { useAssistantState } from "../../context";
+import { useAuiState } from "@assistant-ui/store";
 
 export const MessagePrimitiveError: FC<PropsWithChildren> = ({ children }) => {
-  const hasError = useAssistantState(
+  const hasError = useAuiState(
     ({ message }) =>
       message.status?.type === "incomplete" &&
       message.status.reason === "error",

package/src/primitives/message/MessageIf.tsx

@@ -1,7 +1,7 @@
 "use client";
 
 import type { FC, PropsWithChildren } from "react";
-import { useAssistantState } from "../../context";
+import { useAuiState } from "@assistant-ui/store";
 import type { RequireAtLeastOne } from "../../utils/RequireAtLeastOne";
 
 type MessageIfFilters = {
@@ -20,7 +20,7 @@ type MessageIfFilters = {
 type UseMessageIfProps = RequireAtLeastOne<MessageIfFilters>;
 
 const useMessageIf = (props: UseMessageIfProps) => {
-  return useAssistantState(({ message }) => {
+  return useAuiState(({ message }) => {
     const {
       role,
       attachments,
@@ -78,7 +78,7 @@ export namespace MessagePrimitiveIf {
 }
 
 /**
- * @deprecated Use `<AssistantIf condition={({ message }) => ...} />` instead.
+ * @deprecated Use `<AuiIf condition={({ message }) => ...} />` instead.
  */
 export const MessagePrimitiveIf: FC<MessagePrimitiveIf.Props> = ({
   children,

package/src/primitives/message/MessageParts.tsx

@@ -7,12 +7,11 @@ import {
   PropsWithChildren,
   useMemo,
 } from "react";
+import { useAuiState, useAui } from "@assistant-ui/store";
 import {
-  useAssistantState,
-  useAssistantApi,
   PartByIndexProvider,
   TextMessagePartProvider,
-} from "../../context";
+} from "../../context/providers";
 import { MessagePartPrimitiveText } from "../messagePart/MessagePartText";
 import { MessagePartPrimitiveImage } from "../messagePart/MessagePartImage";
 import type {
@@ -107,7 +106,7 @@ const groupMessageParts = (
 };
 
 const useMessagePartsGroups = (): MessagePartRange[] => {
-  const messageTypes = useAssistantState(
+  const messageTypes = useAuiState(
     useShallow((s) => s.message.parts.map((c: any) => c.type)),
   );
 
@@ -252,7 +251,7 @@ const ToolUIDisplay = ({
 }: {
   Fallback: ToolCallMessagePartComponent | undefined;
 } & ToolCallMessagePartProps) => {
-  const Render = useAssistantState(({ tools }) => {
+  const Render = useAuiState(({ tools }) => {
     const Render = tools.tools[props.toolName] ?? Fallback;
     if (Array.isArray(Render)) return Render[0] ?? Fallback;
     return Render;
@@ -294,13 +293,13 @@ const MessagePartComponent: FC<MessagePartComponentProps> = ({
     tools = {},
   } = {},
 }) => {
-  const api = useAssistantApi();
-  const part = useAssistantState(({ part }) => part);
+  const aui = useAui();
+  const part = useAuiState(({ part }) => part);
 
   const type = part.type;
   if (type === "tool-call") {
-    const addResult = api.part().addToolResult;
-    const resume = api.part().resumeToolCall;
+    const addResult = aui.part().addToolResult;
+    const resume = aui.part().resumeToolCall;
     if ("Override" in tools)
       return <tools.Override {...part} addResult={addResult} resume={resume} />;
     const Tool = tools.by_name?.[part.toolName] ?? tools.Fallback;
@@ -409,7 +408,7 @@ const COMPLETE_STATUS: MessagePartStatus = Object.freeze({
 });
 
 const EmptyPartsImpl: FC<MessagePartComponentProps> = ({ components }) => {
-  const status = useAssistantState(
+  const status = useAuiState(
     (s) => (s.message.status ?? COMPLETE_STATUS) as MessagePartStatus,
   );
 
@@ -457,9 +456,7 @@ const EmptyParts = memo(
 export const MessagePrimitiveParts: FC<MessagePrimitiveParts.Props> = ({
   components,
 }) => {
-  const contentLength = useAssistantState(
-    ({ message }) => message.parts.length,
-  );
+  const contentLength = useAuiState(({ message }) => message.parts.length);
   const messageRanges = useMessagePartsGroups();
 
   const partsElements = useMemo(() => {

package/src/primitives/message/MessagePartsGrouped.tsx

@@ -7,12 +7,11 @@ import {
   PropsWithChildren,
   useMemo,
 } from "react";
+import { useAuiState, useAui } from "@assistant-ui/store";
 import {
-  useAssistantState,
   PartByIndexProvider,
-  useAssistantApi,
   TextMessagePartProvider,
-} from "../../context";
+} from "../../context/providers";
 import { MessagePartPrimitiveText } from "../messagePart/MessagePartText";
 import { MessagePartPrimitiveImage } from "../messagePart/MessagePartImage";
 import type {
@@ -75,7 +74,7 @@ const groupMessagePartsByParentId: GroupingFunction = (
 const useMessagePartsGrouped = (
   groupingFunction: GroupingFunction,
 ): MessagePartGroup[] => {
-  const parts = useAssistantState(({ message }) => message.parts);
+  const parts = useAuiState(({ message }) => message.parts);
 
   return useMemo(() => {
     if (parts.length === 0) {
@@ -220,7 +219,7 @@ const ToolUIDisplay = ({
 }: {
   Fallback: ToolCallMessagePartComponent | undefined;
 } & ToolCallMessagePartProps) => {
-  const Render = useAssistantState(({ tools }) => {
+  const Render = useAuiState(({ tools }) => {
     const Render = tools.tools[props.toolName] ?? Fallback;
     if (Array.isArray(Render)) return Render[0] ?? Fallback;
     return Render;
@@ -261,13 +260,13 @@ const MessagePartComponent: FC<MessagePartComponentProps> = ({
     tools = {},
   } = {},
 }) => {
-  const api = useAssistantApi();
-  const part = useAssistantState(({ part }) => part);
+  const aui = useAui();
+  const part = useAuiState(({ part }) => part);
 
   const type = part.type;
   if (type === "tool-call") {
-    const addResult = (result: any) => api.part().addToolResult(result);
-    const resume = api.part().resumeToolCall;
+    const addResult = (result: any) => aui.part().addToolResult(result);
+    const resume = aui.part().resumeToolCall;
     if ("Override" in tools)
       return <tools.Override {...part} addResult={addResult} resume={resume} />;
     const Tool = tools.by_name?.[part.toolName] ?? tools.Fallback;
@@ -355,7 +354,7 @@ const COMPLETE_STATUS: MessagePartStatus = Object.freeze({
 });
 
 const EmptyPartsImpl: FC<MessagePartComponentProps> = ({ components }) => {
-  const status = useAssistantState(
+  const status = useAuiState(
     (s) => (s.message.status ?? COMPLETE_STATUS) as MessagePartStatus,
   );
 
@@ -427,9 +426,7 @@ const EmptyParts = memo(
 export const MessagePrimitiveUnstable_PartsGrouped: FC<
   MessagePrimitiveUnstable_PartsGrouped.Props
 > = ({ groupingFunction, components }) => {
-  const contentLength = useAssistantState(
-    ({ message }) => message.parts.length,
-  );
+  const contentLength = useAuiState(({ message }) => message.parts.length);
   const messageGroups = useMessagePartsGrouped(groupingFunction);
 
   const partsElements = useMemo(() => {

package/src/primitives/message/MessageRoot.tsx

@@ -7,7 +7,7 @@ import {
   ComponentPropsWithoutRef,
   useCallback,
 } from "react";
-import { useAssistantApi, useAssistantState } from "../../context";
+import { useAui, useAuiState } from "@assistant-ui/store";
 import { useManagedRef } from "../../utils/hooks/useManagedRef";
 import { useSizeHandle } from "../../utils/hooks/useSizeHandle";
 import { useComposedRefs } from "@radix-ui/react-compose-refs";
@@ -15,8 +15,8 @@ import { useThreadViewport } from "../../context/react/ThreadViewportContext";
 import { ThreadPrimitiveViewportSlack } from "../thread/ThreadViewportSlack";
 
 const useIsHoveringRef = () => {
-  const api = useAssistantApi();
-  const message = useAssistantState(() => api.message());
+  const aui = useAui();
+  const message = useAuiState(() => aui.message());
 
   const callbackRef = useCallback(
     (el: HTMLElement) => {
@@ -59,7 +59,7 @@ const useMessageViewportRef = () => {
 
   // inset rules:
   // - the previous user message before the last assistant message registers its full height
-  const shouldRegisterAsInset = useAssistantState(
+  const shouldRegisterAsInset = useAuiState(
     ({ thread, message }) =>
       turnAnchor === "top" &&
       message.role === "user" &&