@assistant-ui/react 0.10.49 → 0.11.0

package/src/primitives/actionBar/useActionBarFloatStatus.tsx

@@ -1,11 +1,6 @@
 "use client";
 
-import {
-  useMessageRuntime,
-  useMessageUtilsStore,
-} from "../../context/react/MessageContext";
-import { useThreadRuntime } from "../../context/react/ThreadContext";
-import { useCombinedStore } from "../../utils/combined/useCombinedStore";
+import { useAssistantState } from "../../context";
 
 export enum HideAndFloatStatus {
   Hidden = "hidden",
@@ -24,32 +19,25 @@ export const useActionBarFloatStatus = ({
   autohide,
   autohideFloat,
 }: UseActionBarFloatStatusProps) => {
-  const threadRuntime = useThreadRuntime();
-  const messageRuntime = useMessageRuntime();
-  const messageUtilsStore = useMessageUtilsStore();
-
-  return useCombinedStore(
-    [threadRuntime, messageRuntime, messageUtilsStore],
-    (t, m, mu) => {
-      if (hideWhenRunning && t.isRunning) return HideAndFloatStatus.Hidden;
-
-      const autohideEnabled =
-        autohide === "always" || (autohide === "not-last" && !m.isLast);
-
-      // normal status
-      if (!autohideEnabled) return HideAndFloatStatus.Normal;
-
-      // hidden status
-      if (!mu.isHovering) return HideAndFloatStatus.Hidden;
-
-      // floating status
-      if (
-        autohideFloat === "always" ||
-        (autohideFloat === "single-branch" && m.branchCount <= 1)
-      )
-        return HideAndFloatStatus.Floating;
-
-      return HideAndFloatStatus.Normal;
-    },
-  );
+  return useAssistantState(({ thread, message }) => {
+    if (hideWhenRunning && thread.isRunning) return HideAndFloatStatus.Hidden;
+
+    const autohideEnabled =
+      autohide === "always" || (autohide === "not-last" && !message.isLast);
+
+    // normal status
+    if (!autohideEnabled) return HideAndFloatStatus.Normal;
+
+    // hidden status
+    if (!message.isHovering) return HideAndFloatStatus.Hidden;
+
+    // floating status
+    if (
+      autohideFloat === "always" ||
+      (autohideFloat === "single-branch" && message.branchCount <= 1)
+    )
+      return HideAndFloatStatus.Floating;
+
+    return HideAndFloatStatus.Normal;
+  });
 };

package/src/primitives/assistantModal/AssistantModalRoot.tsx

@@ -3,7 +3,7 @@
 import { FC, useEffect, useState } from "react";
 import * as PopoverPrimitive from "@radix-ui/react-popover";
 import { ScopedProps, usePopoverScope } from "./scope";
-import { useThreadRuntime } from "../../context";
+import { useAssistantApi } from "../../context";
 
 export namespace AssistantModalPrimitiveRoot {
   export type Props = PopoverPrimitive.PopoverProps & {
@@ -21,14 +21,14 @@ const useAssistantModalOpenState = ({
   const state = useState(defaultOpen);
 
   const [, setOpen] = state;
-  const threadRuntime = useThreadRuntime();
+  const api = useAssistantApi();
   useEffect(() => {
     if (!unstable_openOnRunStart) return undefined;
 
-    return threadRuntime.unstable_on("run-start", () => {
+    return api.on("thread.run-start", () => {
       setOpen(true);
     });
-  }, [unstable_openOnRunStart, setOpen, threadRuntime]);
+  }, [unstable_openOnRunStart, setOpen, api]);
 
   return state;
 };

package/src/primitives/attachment/AttachmentName.tsx

@@ -1,7 +1,7 @@
 "use client";
 
 import type { FC } from "react";
-import { useAttachment } from "../../context/react/AttachmentContext";
+import { useAssistantState } from "../../context";
 
 export namespace AttachmentPrimitiveName {
   export type Props = Record<string, never>;
@@ -10,7 +10,7 @@ export namespace AttachmentPrimitiveName {
 export const AttachmentPrimitiveName: FC<
   AttachmentPrimitiveName.Props
 > = () => {
-  const name = useAttachment((a) => a.name);
+  const name = useAssistantState(({ attachment }) => attachment.name);
   return <>{name}</>;
 };
 

package/src/primitives/attachment/AttachmentRemove.tsx

@@ -6,14 +6,14 @@ import {
   createActionButton,
 } from "../../utils/createActionButton";
 import { useCallback } from "react";
-import { useAttachmentRuntime } from "../../context/react/AttachmentContext";
+import { useAssistantApi } from "../../context";
 
 const useAttachmentRemove = () => {
-  const attachmentRuntime = useAttachmentRuntime();
+  const api = useAssistantApi();
 
   const handleRemoveAttachment = useCallback(() => {
-    attachmentRuntime.remove();
-  }, [attachmentRuntime]);
+    api.attachment().remove();
+  }, [api]);
 
   return handleRemoveAttachment;
 };

package/src/primitives/attachment/AttachmentThumb.tsx

@@ -1,7 +1,7 @@
 "use client";
 
 import { ComponentPropsWithoutRef, forwardRef, type ComponentRef } from "react";
-import { useAttachment } from "../../context/react/AttachmentContext";
+import { useAssistantState } from "../../context";
 import { Primitive } from "@radix-ui/react-primitive";
 
 type PrimitiveDivProps = ComponentPropsWithoutRef<typeof Primitive.div>;
@@ -15,7 +15,10 @@ export const AttachmentPrimitiveThumb = forwardRef<
   AttachmentPrimitiveThumb.Element,
   AttachmentPrimitiveThumb.Props
 >((props, ref) => {
-  const ext = useAttachment((a) => a.name.split(".").pop());
+  const ext = useAssistantState(({ attachment }) => {
+    const parts = attachment.name.split(".");
+    return parts.length > 1 ? parts.pop()! : "";
+  });
   return (
     <Primitive.div {...props} ref={ref}>
       .{ext}

package/src/primitives/branchPicker/BranchPickerCount.tsx

@@ -1,10 +1,10 @@
 "use client";
 
 import type { FC } from "react";
-import { useMessage } from "../../context/react/MessageContext";
+import { useAssistantState } from "../../context";
 
 const useBranchPickerCount = () => {
-  const branchCount = useMessage((s) => s.branchCount);
+  const branchCount = useAssistantState(({ message }) => message.branchCount);
   return branchCount;
 };
 

package/src/primitives/branchPicker/BranchPickerNext.tsx

@@ -6,18 +6,17 @@ import {
   createActionButton,
 } from "../../utils/createActionButton";
 import { useCallback } from "react";
-import {
-  useMessage,
-  useMessageRuntime,
-} from "../../context/react/MessageContext";
+import { useAssistantState, useAssistantApi } from "../../context";
 
 const useBranchPickerNext = () => {
-  const messageRuntime = useMessageRuntime();
-  const disabled = useMessage((m) => m.branchNumber >= m.branchCount);
+  const api = useAssistantApi();
+  const disabled = useAssistantState(
+    ({ message }) => message.branchNumber >= message.branchCount,
+  );
 
   const callback = useCallback(() => {
-    messageRuntime.switchToBranch({ position: "next" });
-  }, [messageRuntime]);
+    api.message().switchToBranch({ position: "next" });
+  }, [api]);
 
   if (disabled) return null;
   return callback;

package/src/primitives/branchPicker/BranchPickerNumber.tsx

@@ -1,10 +1,10 @@
 "use client";
 
 import type { FC } from "react";
-import { useMessage } from "../../context/react/MessageContext";
+import { useAssistantState } from "../../context";
 
 const useBranchPickerNumber = () => {
-  const branchNumber = useMessage((s) => s.branchNumber);
+  const branchNumber = useAssistantState(({ message }) => message.branchNumber);
   return branchNumber;
 };
 

package/src/primitives/branchPicker/BranchPickerPrevious.tsx

@@ -6,10 +6,7 @@ import {
   createActionButton,
 } from "../../utils/createActionButton";
 import { useCallback } from "react";
-import {
-  useMessage,
-  useMessageRuntime,
-} from "../../context/react/MessageContext";
+import { useAssistantState, useAssistantApi } from "../../context";
 
 /**
  * Hook that provides navigation to the previous branch functionality.
@@ -33,12 +30,14 @@ import {
  * ```
  */
 const useBranchPickerPrevious = () => {
-  const messageRuntime = useMessageRuntime();
-  const disabled = useMessage((m) => m.branchNumber <= 1);
+  const api = useAssistantApi();
+  const disabled = useAssistantState(
+    ({ message }) => message.branchNumber <= 1,
+  );
 
   const callback = useCallback(() => {
-    messageRuntime.switchToBranch({ position: "previous" });
-  }, [messageRuntime]);
+    api.message().switchToBranch({ position: "previous" });
+  }, [api]);
 
   if (disabled) return null;
   return callback;

package/src/primitives/composer/ComposerAddAttachment.tsx

@@ -6,7 +6,7 @@ import {
   createActionButton,
 } from "../../utils/createActionButton";
 import { useCallback } from "react";
-import { useComposer, useComposerRuntime } from "../../context";
+import { useAssistantState, useAssistantApi } from "../../context";
 
 const useComposerAddAttachment = ({
   multiple = true,
@@ -14,16 +14,16 @@ const useComposerAddAttachment = ({
   /** allow selecting multiple files */
   multiple?: boolean | undefined;
 } = {}) => {
-  const disabled = useComposer((c) => !c.isEditing);
+  const disabled = useAssistantState(({ composer }) => !composer.isEditing);
+  const api = useAssistantApi();
 
-  const composerRuntime = useComposerRuntime();
   const callback = useCallback(() => {
     const input = document.createElement("input");
     input.type = "file";
     input.multiple = multiple;
     input.hidden = true;
 
-    const attachmentAccept = composerRuntime.getAttachmentAccept();
+    const attachmentAccept = api.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) {
-        composerRuntime.addAttachment(file);
+        api.composer().addAttachment(file);
       }
 
       document.body.removeChild(input);
@@ -47,7 +47,7 @@ const useComposerAddAttachment = ({
     };
 
     input.click();
-  }, [composerRuntime, multiple]);
+  }, [api, multiple]);
 
   if (disabled) return null;
   return callback;

package/src/primitives/composer/ComposerAttachmentDropzone.tsx

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

package/src/primitives/composer/ComposerAttachments.tsx

@@ -2,9 +2,10 @@
 
 import { ComponentType, type FC, memo, useMemo } from "react";
 import { Attachment } from "../../types";
-import { useComposer, useComposerRuntime } from "../../context";
-import { useThreadComposerAttachment } from "../../context/react/AttachmentContext";
-import { AttachmentRuntimeProvider } from "../../context/providers/AttachmentRuntimeProvider";
+import {
+  useAssistantState,
+  ComposerAttachmentByIndexProvider,
+} from "../../context";
 
 export namespace ComposerPrimitiveAttachments {
   export type Props = {
@@ -40,10 +41,10 @@ const getComponent = (
 const AttachmentComponent: FC<{
   components: ComposerPrimitiveAttachments.Props["components"];
 }> = ({ components }) => {
-  const Component = useThreadComposerAttachment((a) =>
-    getComponent(components, a),
-  );
+  const attachment = useAssistantState(({ attachment }) => attachment);
+  if (!attachment) return null;
 
+  const Component = getComponent(components, attachment);
   if (!Component) return null;
   return <Component />;
 };
@@ -75,16 +76,10 @@ export namespace ComposerPrimitiveAttachmentByIndex {
 export const ComposerPrimitiveAttachmentByIndex: FC<ComposerPrimitiveAttachmentByIndex.Props> =
   memo(
     ({ index, components }) => {
-      const composerRuntime = useComposerRuntime();
-      const runtime = useMemo(
-        () => composerRuntime.getAttachmentByIndex(index),
-        [composerRuntime, index],
-      );
-
       return (
-        <AttachmentRuntimeProvider runtime={runtime}>
+        <ComposerAttachmentByIndexProvider index={index}>
           <AttachmentComponent components={components} />
-        </AttachmentRuntimeProvider>
+        </ComposerAttachmentByIndexProvider>
       );
     },
     (prev, next) =>
@@ -101,7 +96,9 @@ ComposerPrimitiveAttachmentByIndex.displayName =
 export const ComposerPrimitiveAttachments: FC<
   ComposerPrimitiveAttachments.Props
 > = ({ components }) => {
-  const attachmentsCount = useComposer((s) => s.attachments.length);
+  const attachmentsCount = useAssistantState(
+    (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 { useComposer, useComposerRuntime } from "../../context";
+import { useAssistantState, useAssistantApi } from "../../context";
 
 const useComposerCancel = () => {
-  const composerRuntime = useComposerRuntime();
-  const disabled = useComposer((c) => !c.canCancel);
+  const api = useAssistantApi();
+  const disabled = useAssistantState(({ composer }) => !composer.canCancel);
 
   const callback = useCallback(() => {
-    composerRuntime.cancel();
-  }, [composerRuntime]);
+    api.composer().cancel();
+  }, [api]);
 
   if (disabled) return null;
   return callback;

package/src/primitives/composer/ComposerIf.tsx

@@ -1,7 +1,7 @@
 "use client";
 
 import type { FC, PropsWithChildren } from "react";
-import { useComposer } from "../../context/react/ComposerContext";
+import { useAssistantState } from "../../context";
 import type { RequireAtLeastOne } from "../../utils/RequireAtLeastOne";
 
 type ComposerIfFilters = {
@@ -11,7 +11,7 @@ type ComposerIfFilters = {
 export type UseComposerIfProps = RequireAtLeastOne<ComposerIfFilters>;
 
 const useComposerIf = (props: UseComposerIfProps) => {
-  return useComposer((composer) => {
+  return useAssistantState(({ composer }) => {
     if (props.editing === true && !composer.isEditing) return false;
     if (props.editing === false && composer.isEditing) return false;
 

package/src/primitives/composer/ComposerInput.tsx

@@ -14,14 +14,9 @@ import {
 import TextareaAutosize, {
   type TextareaAutosizeProps,
 } from "react-textarea-autosize";
-import {
-  useComposer,
-  useComposerRuntime,
-} from "../../context/react/ComposerContext";
-import { useThread, useThreadRuntime } from "../../context/react/ThreadContext";
 import { useEscapeKeydown } from "@radix-ui/react-use-escape-keydown";
 import { useOnScrollToBottom } from "../../utils/hooks/useOnScrollToBottom";
-import { useThreadListItemRuntime } from "../../context/react/ThreadListItemContext";
+import { useAssistantState, useAssistantApi } from "../../context";
 
 export namespace ComposerPrimitiveInput {
   export type Element = HTMLTextAreaElement;
@@ -102,26 +97,26 @@ export const ComposerPrimitiveInput = forwardRef<
     },
     forwardedRef,
   ) => {
-    const threadListItemRuntime = useThreadListItemRuntime();
-    const threadRuntime = useThreadRuntime();
-    const composerRuntime = useComposerRuntime();
+    const api = useAssistantApi();
 
-    const value = useComposer((c) => {
-      if (!c.isEditing) return "";
-      return c.text;
+    const value = useAssistantState(({ composer }) => {
+      if (!composer.isEditing) return "";
+      return composer.text;
     });
 
     const Component = asChild ? Slot : TextareaAutosize;
 
-    const isDisabled = Boolean(useThread((t) => t.isDisabled) || disabledProp);
+    const isDisabled =
+      useAssistantState(({ thread }) => thread.isDisabled) || disabledProp;
     const textareaRef = useRef<HTMLTextAreaElement>(null);
     const ref = useComposedRefs(forwardedRef, textareaRef);
 
     useEscapeKeydown((e) => {
       if (!cancelOnEscape) return;
 
-      if (composerRuntime.getState().canCancel) {
-        composerRuntime.cancel();
+      const composer = api.composer();
+      if (composer.getState().canCancel) {
+        composer.cancel();
         e.preventDefault();
       }
     });
@@ -133,7 +128,7 @@ export const ComposerPrimitiveInput = forwardRef<
       if (e.nativeEvent.isComposing) return;
 
       if (e.key === "Enter" && e.shiftKey === false) {
-        const { isRunning } = threadRuntime.getState();
+        const isRunning = api.thread().getState().isRunning;
 
         if (!isRunning) {
           e.preventDefault();
@@ -145,14 +140,14 @@ export const ComposerPrimitiveInput = forwardRef<
 
     const handlePaste = async (e: ClipboardEvent<HTMLTextAreaElement>) => {
       if (!addAttachmentOnPaste) return;
-      const threadCapabilities = threadRuntime.getState().capabilities;
+      const threadCapabilities = api.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) => composerRuntime.addAttachment(file)),
+            files.map((file) => api.composer().addAttachment(file)),
           );
         } catch (error) {
           console.error("Error adding attachment:", error);
@@ -172,29 +167,33 @@ export const ComposerPrimitiveInput = forwardRef<
     useEffect(() => focus(), [focus]);
 
     useOnScrollToBottom(() => {
-      if (composerRuntime.type === "thread" && unstable_focusOnScrollToBottom) {
+      if (
+        api.composer().getState().type === "thread" &&
+        unstable_focusOnScrollToBottom
+      ) {
         focus();
       }
     });
 
     useEffect(() => {
-      if (composerRuntime.type !== "thread" || !unstable_focusOnRunStart)
+      if (
+        api.composer().getState().type !== "thread" ||
+        !unstable_focusOnRunStart
+      )
         return undefined;
 
-      return threadRuntime.unstable_on("run-start", focus);
-    }, [unstable_focusOnRunStart, focus, composerRuntime, threadRuntime]);
+      return api.on("thread.run-start", focus);
+    }, [unstable_focusOnRunStart, focus, api]);
 
     useEffect(() => {
-      if (composerRuntime.type !== "thread" || !unstable_focusOnThreadSwitched)
+      if (
+        api.composer().getState().type !== "thread" ||
+        !unstable_focusOnThreadSwitched
+      )
         return undefined;
 
-      return threadListItemRuntime.unstable_on("switched-to", focus);
-    }, [
-      unstable_focusOnThreadSwitched,
-      focus,
-      composerRuntime,
-      threadListItemRuntime,
-    ]);
+      return api.on("thread-list-item.switched-to", focus);
+    }, [unstable_focusOnThreadSwitched, focus, api]);
 
     return (
       <Component
@@ -204,8 +203,9 @@ export const ComposerPrimitiveInput = forwardRef<
         ref={ref as React.ForwardedRef<HTMLTextAreaElement>}
         disabled={isDisabled}
         onChange={composeEventHandlers(onChange, (e) => {
-          if (!composerRuntime.getState().isEditing) return;
-          return composerRuntime.setText(e.target.value);
+          if (!api.composer().getState().isEditing) return;
+          api.composer().setText(e.target.value);
+          api.flushSync();
         })}
         onKeyDown={composeEventHandlers(onKeyDown, handleKeyPress)}
         onPaste={composeEventHandlers(onPaste, handlePaste)}

package/src/primitives/composer/ComposerSend.tsx

@@ -6,22 +6,18 @@ import {
   createActionButton,
 } from "../../utils/createActionButton";
 import { useCallback } from "react";
-import { useCombinedStore } from "../../utils/combined/useCombinedStore";
-import { useThreadRuntime } from "../../context/react/ThreadContext";
-import { useComposerRuntime } from "../../context";
+import { useAssistantState, useAssistantApi } from "../../context";
 
 export const useComposerSend = () => {
-  const composerRuntime = useComposerRuntime();
-  const threadRuntime = useThreadRuntime();
+  const api = useAssistantApi();
 
-  const disabled = useCombinedStore(
-    [threadRuntime, composerRuntime],
-    (t, c) => t.isRunning || !c.isEditing || c.isEmpty,
+  const disabled = useAssistantState(
+    (s) => s.thread.isRunning || !s.composer.isEditing || s.composer.isEmpty,
   );
 
   const callback = useCallback(() => {
-    composerRuntime.send();
-  }, [composerRuntime]);
+    api.composer().send();
+  }, [api]);
 
   if (disabled) 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 { useMessage } from "../../context/react/MessageContext";
+import { useAssistantState } from "../../context";
 
 export namespace ErrorPrimitiveMessage {
   export type Element = ComponentRef<typeof Primitive.span>;
@@ -13,9 +13,10 @@ export const ErrorPrimitiveMessage = forwardRef<
   ErrorPrimitiveMessage.Element,
   ErrorPrimitiveMessage.Props
 >(({ children, ...props }, forwardRef) => {
-  const error = useMessage((m) => {
-    return m.status?.type === "incomplete" && m.status.reason === "error"
-      ? m.status.error
+  const error = useAssistantState(({ message }) => {
+    return message.status?.type === "incomplete" &&
+      message.status.reason === "error"
+      ? message.status.error
       : undefined;
   });
 

package/src/primitives/message/MessageAttachments.tsx

@@ -1,9 +1,10 @@
 "use client";
 
 import { ComponentType, type FC, memo, useMemo } from "react";
-import { useMessage, useMessageRuntime } from "../../context";
-import { useMessageAttachment } from "../../context/react/AttachmentContext";
-import { AttachmentRuntimeProvider } from "../../context/providers/AttachmentRuntimeProvider";
+import {
+  useAssistantState,
+  MessageAttachmentByIndexProvider,
+} from "../../context";
 import { CompleteAttachment } from "../../types";
 
 export namespace MessagePrimitiveAttachments {
@@ -40,8 +41,10 @@ const getComponent = (
 const AttachmentComponent: FC<{
   components: MessagePrimitiveAttachments.Props["components"];
 }> = ({ components }) => {
-  const Component = useMessageAttachment((a) => getComponent(components, a));
+  const attachment = useAssistantState(({ attachment }) => attachment);
+  if (!attachment) return null;
 
+  const Component = getComponent(components, attachment as CompleteAttachment);
   if (!Component) return null;
   return <Component />;
 };
@@ -73,16 +76,10 @@ export namespace MessagePrimitiveAttachmentByIndex {
 export const MessagePrimitiveAttachmentByIndex: FC<MessagePrimitiveAttachmentByIndex.Props> =
   memo(
     ({ index, components }) => {
-      const messageRuntime = useMessageRuntime();
-      const runtime = useMemo(
-        () => messageRuntime.getAttachmentByIndex(index),
-        [messageRuntime, index],
-      );
-
       return (
-        <AttachmentRuntimeProvider runtime={runtime}>
+        <MessageAttachmentByIndexProvider index={index}>
           <AttachmentComponent components={components} />
-        </AttachmentRuntimeProvider>
+        </MessageAttachmentByIndexProvider>
       );
     },
     (prev, next) =>
@@ -99,7 +96,7 @@ MessagePrimitiveAttachmentByIndex.displayName =
 export const MessagePrimitiveAttachments: FC<
   MessagePrimitiveAttachments.Props
 > = ({ components }) => {
-  const attachmentsCount = useMessage((message) => {
+  const attachmentsCount = useAssistantState(({ message }) => {
     if (message.role !== "user") return 0;
     return message.attachments.length;
   });