@assistant-ui/react 0.10.50 → 0.11.0

package/src/primitives/message/MessageError.tsx

@@ -1,11 +1,13 @@
 "use client";
 
 import { FC, PropsWithChildren } from "react";
-import { useMessage } from "../../context/react/MessageContext";
+import { useAssistantState } from "../../context";
 
 export const MessagePrimitiveError: FC<PropsWithChildren> = ({ children }) => {
-  const hasError = useMessage(
-    (m) => m.status?.type === "incomplete" && m.status.reason === "error",
+  const hasError = useAssistantState(
+    ({ message }) =>
+      message.status?.type === "incomplete" &&
+      message.status.reason === "error",
   );
   return hasError ? <>{children}</> : null;
 };

package/src/primitives/message/MessageIf.tsx

@@ -1,12 +1,8 @@
 "use client";
 
 import type { FC, PropsWithChildren } from "react";
-import {
-  useMessageRuntime,
-  useMessageUtilsStore,
-} from "../../context/react/MessageContext";
+import { useAssistantState } from "../../context";
 import type { RequireAtLeastOne } from "../../utils/RequireAtLeastOne";
-import { useCombinedStore } from "../../utils/combined/useCombinedStore";
 
 type MessageIfFilters = {
   user: boolean | undefined;
@@ -24,62 +20,57 @@ type MessageIfFilters = {
 type UseMessageIfProps = RequireAtLeastOne<MessageIfFilters>;
 
 const useMessageIf = (props: UseMessageIfProps) => {
-  const messageRuntime = useMessageRuntime();
-  const messageUtilsStore = useMessageUtilsStore();
+  return useAssistantState(({ message }) => {
+    const {
+      role,
+      attachments,
+      parts,
+      branchCount,
+      isLast,
+      speech,
+      submittedFeedback,
+      isCopied,
+      isHovering,
+    } = message;
 
-  return useCombinedStore(
-    [messageRuntime, messageUtilsStore],
-    (
-      {
-        role,
-        attachments,
-        content,
-        branchCount,
-        isLast,
-        speech,
-        submittedFeedback,
-      },
-      { isCopied, isHovering },
-    ) => {
-      if (props.hasBranches === true && branchCount < 2) return false;
+    if (props.hasBranches === true && branchCount < 2) return false;
 
-      if (props.user && role !== "user") return false;
-      if (props.assistant && role !== "assistant") return false;
-      if (props.system && role !== "system") return false;
+    if (props.user && role !== "user") return false;
+    if (props.assistant && role !== "assistant") return false;
+    if (props.system && role !== "system") return false;
 
-      if (props.lastOrHover === true && !isHovering && !isLast) return false;
-      if (props.last !== undefined && props.last !== isLast) return false;
+    if (props.lastOrHover === true && !isHovering && !isLast) return false;
+    if (props.last !== undefined && props.last !== isLast) return false;
 
-      if (props.copied === true && !isCopied) return false;
-      if (props.copied === false && isCopied) return false;
+    if (props.copied === true && !isCopied) return false;
+    if (props.copied === false && isCopied) return false;
 
-      if (props.speaking === true && speech == null) return false;
-      if (props.speaking === false && speech != null) return false;
+    if (props.speaking === true && speech == null) return false;
+    if (props.speaking === false && speech != null) return false;
 
-      if (
-        props.hasAttachments === true &&
-        (role !== "user" || !attachments.length)
-      )
-        return false;
-      if (
-        props.hasAttachments === false &&
-        role === "user" &&
-        !!attachments.length
-      )
-        return false;
+    if (
+      props.hasAttachments === true &&
+      (role !== "user" || !attachments?.length)
+    )
+      return false;
+    if (
+      props.hasAttachments === false &&
+      role === "user" &&
+      !!attachments?.length
+    )
+      return false;
 
-      if (props.hasContent === true && content.length === 0) return false;
-      if (props.hasContent === false && content.length > 0) return false;
+    if (props.hasContent === true && parts.length === 0) return false;
+    if (props.hasContent === false && parts.length > 0) return false;
 
-      if (
-        props.submittedFeedback !== undefined &&
-        (submittedFeedback?.type ?? null) !== props.submittedFeedback
-      )
-        return false;
+    if (
+      props.submittedFeedback !== undefined &&
+      (submittedFeedback?.type ?? null) !== props.submittedFeedback
+    )
+      return false;
 
-      return true;
-    },
-  );
+    return true;
+  });
 };
 
 export namespace MessagePrimitiveIf {

package/src/primitives/message/MessageParts.tsx

@@ -8,16 +8,11 @@ import {
   useMemo,
 } from "react";
 import {
+  useAssistantState,
+  useAssistantApi,
+  PartByIndexProvider,
   TextMessagePartProvider,
-  useMessagePart,
-  useMessagePartRuntime,
-  useToolUIs,
 } from "../../context";
-import {
-  useMessage,
-  useMessageRuntime,
-} from "../../context/react/MessageContext";
-import { MessagePartRuntimeProvider } from "../../context/providers/MessagePartRuntimeProvider";
 import { MessagePartPrimitiveText } from "../messagePart/MessagePartText";
 import { MessagePartPrimitiveImage } from "../messagePart/MessagePartImage";
 import type {
@@ -86,8 +81,8 @@ const groupMessageParts = (
 };
 
 const useMessagePartsGroups = (): MessagePartRange[] => {
-  const messageTypes = useMessage(
-    useShallow((m) => m.content.map((c) => c.type)),
+  const messageTypes = useAssistantState(
+    useShallow((s) => s.message.parts.map((c: any) => c.type)),
   );
 
   return useMemo(() => {
@@ -200,7 +195,11 @@ const ToolUIDisplay = ({
 }: {
   Fallback: ToolCallMessagePartComponent | undefined;
 } & ToolCallMessagePartProps) => {
-  const Render = useToolUIs((s) => s.getToolUI(props.toolName)) ?? Fallback;
+  const Render = useAssistantState(({ toolUIs }) => {
+    const Render = toolUIs[props.toolName] ?? Fallback;
+    if (Array.isArray(Render)) return Render[0];
+    return Render;
+  });
   if (!Render) return null;
   return <Render {...props} />;
 };
@@ -237,20 +236,19 @@ const MessagePartComponent: FC<MessagePartComponentProps> = ({
     tools = {},
   } = {},
 }) => {
-  const MessagePartRuntime = useMessagePartRuntime();
-
-  const part = useMessagePart();
+  const api = useAssistantApi();
+  const part = useAssistantState(({ part }) => part);
 
   const type = part.type;
   if (type === "tool-call") {
-    const addResult = (result: any) => MessagePartRuntime.addToolResult(result);
+    const addResult = (result: any) => api.part().addToolResult(result);
     if ("Override" in tools)
       return <tools.Override {...part} addResult={addResult} />;
     const Tool = tools.by_name?.[part.toolName] ?? tools.Fallback;
     return <ToolUIDisplay {...part} Fallback={Tool} addResult={addResult} />;
   }
 
-  if (part.status.type === "requires-action")
+  if (part.status?.type === "requires-action")
     throw new Error("Encountered unexpected requires-action status");
 
   switch (type) {
@@ -306,16 +304,10 @@ export namespace MessagePrimitivePartByIndex {
 export const MessagePrimitivePartByIndex: FC<MessagePrimitivePartByIndex.Props> =
   memo(
     ({ index, components }) => {
-      const messageRuntime = useMessageRuntime();
-      const runtime = useMemo(
-        () => messageRuntime.getMessagePartByIndex(index),
-        [messageRuntime, index],
-      );
-
       return (
-        <MessagePartRuntimeProvider runtime={runtime}>
+        <PartByIndexProvider index={index}>
           <MessagePartComponent components={components} />
-        </MessagePartRuntimeProvider>
+        </PartByIndexProvider>
       );
     },
     (prev, next) =>
@@ -332,10 +324,6 @@ export const MessagePrimitivePartByIndex: FC<MessagePrimitivePartByIndex.Props>
 
 MessagePrimitivePartByIndex.displayName = "MessagePrimitive.PartByIndex";
 
-const COMPLETE_STATUS: MessagePartStatus = Object.freeze({
-  type: "complete",
-});
-
 const EmptyPartFallback: FC<{
   status: MessagePartStatus;
   component: TextMessagePartComponent;
@@ -347,9 +335,14 @@ const EmptyPartFallback: FC<{
   );
 };
 
+const COMPLETE_STATUS: MessagePartStatus = Object.freeze({
+  type: "complete",
+});
+
 const EmptyPartsImpl: FC<MessagePartComponentProps> = ({ components }) => {
-  const status =
-    useMessage((s) => s.status as MessagePartStatus) ?? COMPLETE_STATUS;
+  const status = useAssistantState(
+    (s) => (s.message.status ?? COMPLETE_STATUS) as MessagePartStatus,
+  );
 
   if (components?.Empty) return <components.Empty status={status} />;
 
@@ -395,7 +388,9 @@ const EmptyParts = memo(
 export const MessagePrimitiveParts: FC<MessagePrimitiveParts.Props> = ({
   components,
 }) => {
-  const contentLength = useMessage((s) => s.content.length);
+  const contentLength = useAssistantState(
+    ({ message }) => message.parts.length,
+  );
   const messageRanges = useMessagePartsGroups();
 
   const partsElements = useMemo(() => {

package/src/primitives/message/MessagePartsGrouped.tsx

@@ -8,16 +8,11 @@ import {
   useMemo,
 } from "react";
 import {
+  useAssistantState,
+  PartByIndexProvider,
+  useAssistantApi,
   TextMessagePartProvider,
-  useMessagePart,
-  useMessagePartRuntime,
-  useToolUIs,
 } from "../../context";
-import {
-  useMessage,
-  useMessageRuntime,
-} from "../../context/react/MessageContext";
-import { MessagePartRuntimeProvider } from "../../context/providers/MessagePartRuntimeProvider";
 import { MessagePartPrimitiveText } from "../messagePart/MessagePartText";
 import { MessagePartPrimitiveImage } from "../messagePart/MessagePartImage";
 import type {
@@ -80,7 +75,7 @@ const groupMessagePartsByParentId: GroupingFunction = (
 const useMessagePartsGrouped = (
   groupingFunction: GroupingFunction,
 ): MessagePartGroup[] => {
-  const parts = useMessage((m) => m.content);
+  const parts = useAssistantState(({ message }) => message.parts);
 
   return useMemo(() => {
     if (parts.length === 0) {
@@ -225,7 +220,11 @@ const ToolUIDisplay = ({
 }: {
   Fallback: ToolCallMessagePartComponent | undefined;
 } & ToolCallMessagePartProps) => {
-  const Render = useToolUIs((s) => s.getToolUI(props.toolName)) ?? Fallback;
+  const Render = useAssistantState(({ toolUIs }) => {
+    const Render = toolUIs[props.toolName] ?? Fallback;
+    if (Array.isArray(Render)) return Render[0];
+    return Render;
+  });
   if (!Render) return null;
   return <Render {...props} />;
 };
@@ -262,20 +261,19 @@ const MessagePartComponent: FC<MessagePartComponentProps> = ({
     tools = {},
   } = {},
 }) => {
-  const MessagePartRuntime = useMessagePartRuntime();
-
-  const part = useMessagePart();
+  const api = useAssistantApi();
+  const part = useAssistantState(({ part }) => part);
 
   const type = part.type;
   if (type === "tool-call") {
-    const addResult = (result: any) => MessagePartRuntime.addToolResult(result);
+    const addResult = (result: any) => api.part().addToolResult(result);
     if ("Override" in tools)
       return <tools.Override {...part} addResult={addResult} />;
     const Tool = tools.by_name?.[part.toolName] ?? tools.Fallback;
     return <ToolUIDisplay {...part} Fallback={Tool} addResult={addResult} />;
   }
 
-  if (part.status.type === "requires-action")
+  if (part.status?.type === "requires-action")
     throw new Error("Encountered unexpected requires-action status");
 
   switch (type) {
@@ -310,16 +308,10 @@ type MessagePartProps = {
 };
 
 const MessagePartImpl: FC<MessagePartProps> = ({ partIndex, components }) => {
-  const messageRuntime = useMessageRuntime();
-  const runtime = useMemo(
-    () => messageRuntime.getMessagePartByIndex(partIndex),
-    [messageRuntime, partIndex],
-  );
-
   return (
-    <MessagePartRuntimeProvider runtime={runtime}>
+    <PartByIndexProvider index={partIndex}>
       <MessagePartComponent components={components} />
-    </MessagePartRuntimeProvider>
+    </PartByIndexProvider>
   );
 };
 
@@ -337,10 +329,6 @@ const MessagePart = memo(
     prev.components?.Group === next.components?.Group,
 );
 
-const COMPLETE_STATUS: MessagePartStatus = Object.freeze({
-  type: "complete",
-});
-
 const EmptyPartFallback: FC<{
   status: MessagePartStatus;
   component: TextMessagePartComponent;
@@ -352,9 +340,14 @@ const EmptyPartFallback: FC<{
   );
 };
 
+const COMPLETE_STATUS: MessagePartStatus = Object.freeze({
+  type: "complete",
+});
+
 const EmptyPartsImpl: FC<MessagePartComponentProps> = ({ components }) => {
-  const status =
-    useMessage((s) => s.status as MessagePartStatus) ?? COMPLETE_STATUS;
+  const status = useAssistantState(
+    (s) => (s.message.status ?? COMPLETE_STATUS) as MessagePartStatus,
+  );
 
   if (components?.Empty) return <components.Empty status={status} />;
 
@@ -424,7 +417,9 @@ const EmptyParts = memo(
 export const MessagePrimitiveUnstable_PartsGrouped: FC<
   MessagePrimitiveUnstable_PartsGrouped.Props
 > = ({ groupingFunction, components }) => {
-  const contentLength = useMessage((s) => s.content.length);
+  const contentLength = useAssistantState(
+    ({ message }) => message.parts.length,
+  );
   const messageGroups = useMessagePartsGrouped(groupingFunction);
 
   const partsElements = useMemo(() => {

package/src/primitives/message/MessageRoot.tsx

@@ -7,33 +7,35 @@ import {
   ComponentPropsWithoutRef,
   useCallback,
 } from "react";
-import { useMessageUtilsStore } from "../../context/react/MessageContext";
+import { useAssistantApi, useAssistantState } from "../../context";
 import { useManagedRef } from "../../utils/hooks/useManagedRef";
 import { useComposedRefs } from "@radix-ui/react-compose-refs";
 
 const useIsHoveringRef = () => {
-  const messageUtilsStore = useMessageUtilsStore();
+  const api = useAssistantApi();
+  const message = useAssistantState(() => api.message());
+
   const callbackRef = useCallback(
     (el: HTMLElement) => {
-      const setIsHovering = messageUtilsStore.getState().setIsHovering;
-
       const handleMouseEnter = () => {
-        setIsHovering(true);
+        message.setIsHovering(true);
       };
       const handleMouseLeave = () => {
-        setIsHovering(false);
+        message.setIsHovering(false);
       };
 
       el.addEventListener("mouseenter", handleMouseEnter);
       el.addEventListener("mouseleave", handleMouseLeave);
 
+      if (el.matches(":hover")) message.setIsHovering(true);
+
       return () => {
         el.removeEventListener("mouseenter", handleMouseEnter);
         el.removeEventListener("mouseleave", handleMouseLeave);
-        setIsHovering(false);
+        message.setIsHovering(false);
       };
     },
-    [messageUtilsStore],
+    [message],
   );
 
   return useManagedRef(callbackRef);

package/src/primitives/messagePart/MessagePartInProgress.tsx

@@ -1,7 +1,7 @@
 "use client";
 
 import { FC, PropsWithChildren } from "react";
-import { useMessagePart } from "../../context";
+import { useAssistantState } from "../../context";
 
 export namespace MessagePartPrimitiveInProgress {
   export type Props = PropsWithChildren;
@@ -11,7 +11,9 @@ export namespace MessagePartPrimitiveInProgress {
 export const MessagePartPrimitiveInProgress: FC<
   MessagePartPrimitiveInProgress.Props
 > = ({ children }) => {
-  const isInProgress = useMessagePart((c) => c.status.type === "running");
+  const isInProgress = useAssistantState(
+    ({ part }) => part.status.type === "running",
+  );
 
   return isInProgress ? children : null;
 };

package/src/primitives/messagePart/useMessagePartFile.tsx

@@ -1,17 +1,17 @@
 "use client";
 
 import { MessagePartState } from "../../api/MessagePartRuntime";
-import { useMessagePart } from "../../context/react/MessagePartContext";
+import { useAssistantState } from "../../context";
 import { FileMessagePart } from "../../types";
 
 export const useMessagePartFile = () => {
-  const file = useMessagePart((c) => {
-    if (c.type !== "file")
+  const file = useAssistantState(({ part }) => {
+    if (part.type !== "file")
       throw new Error(
         "MessagePartFile can only be used inside file message parts.",
       );
 
-    return c as MessagePartState & FileMessagePart;
+    return part as MessagePartState & FileMessagePart;
   });
 
   return file;

package/src/primitives/messagePart/useMessagePartImage.tsx

@@ -1,17 +1,17 @@
 "use client";
 
 import { MessagePartState } from "../../api/MessagePartRuntime";
-import { useMessagePart } from "../../context/react/MessagePartContext";
+import { useAssistantState } from "../../context";
 import { ImageMessagePart } from "../../types";
 
 export const useMessagePartImage = () => {
-  const image = useMessagePart((c) => {
-    if (c.type !== "image")
+  const image = useAssistantState(({ part }) => {
+    if (part.type !== "image")
       throw new Error(
         "MessagePartImage can only be used inside image message parts.",
       );
 
-    return c as MessagePartState & ImageMessagePart;
+    return part as MessagePartState & ImageMessagePart;
   });
 
   return image;

package/src/primitives/messagePart/useMessagePartReasoning.tsx

@@ -1,17 +1,17 @@
 "use client";
 
 import { MessagePartState } from "../../api/MessagePartRuntime";
-import { useMessagePart } from "../../context/react/MessagePartContext";
+import { useAssistantState } from "../../context";
 import { ReasoningMessagePart } from "../../types";
 
 export const useMessagePartReasoning = () => {
-  const text = useMessagePart((c) => {
-    if (c.type !== "reasoning")
+  const text = useAssistantState(({ part }) => {
+    if (part.type !== "reasoning")
       throw new Error(
         "MessagePartReasoning can only be used inside reasoning message parts.",
       );
 
-    return c as MessagePartState & ReasoningMessagePart;
+    return part as MessagePartState & ReasoningMessagePart;
   });
 
   return text;

package/src/primitives/messagePart/useMessagePartSource.tsx

@@ -1,17 +1,17 @@
 "use client";
 
 import { MessagePartState } from "../../api/MessagePartRuntime";
-import { useMessagePart } from "../../context/react/MessagePartContext";
+import { useAssistantState } from "../../context";
 import { SourceMessagePart } from "../../types";
 
 export const useMessagePartSource = () => {
-  const source = useMessagePart((c) => {
-    if (c.type !== "source")
+  const source = useAssistantState(({ part }) => {
+    if (part.type !== "source")
       throw new Error(
         "MessagePartSource can only be used inside source message parts.",
       );
 
-    return c as MessagePartState & SourceMessagePart;
+    return part as MessagePartState & SourceMessagePart;
   });
 
   return source;

package/src/primitives/messagePart/useMessagePartText.tsx

@@ -1,17 +1,17 @@
 "use client";
 
 import { MessagePartState } from "../../api/MessagePartRuntime";
-import { useMessagePart } from "../../context/react/MessagePartContext";
+import { useAssistantState } from "../../context";
 import { TextMessagePart, ReasoningMessagePart } from "../../types";
 
 export const useMessagePartText = () => {
-  const text = useMessagePart((c) => {
-    if (c.type !== "text" && c.type !== "reasoning")
+  const text = useAssistantState(({ part }) => {
+    if (part.type !== "text" && part.type !== "reasoning")
       throw new Error(
         "MessagePartText can only be used inside text or reasoning message parts.",
       );
 
-    return c as MessagePartState & (TextMessagePart | ReasoningMessagePart);
+    return part as MessagePartState & (TextMessagePart | ReasoningMessagePart);
   });
 
   return text;

package/src/primitives/thread/ThreadEmpty.tsx

@@ -1,7 +1,7 @@
 "use client";
 
 import type { FC, PropsWithChildren } from "react";
-import { useThread } from "../../context";
+import { useAssistantState } from "../../context";
 
 export namespace ThreadPrimitiveEmpty {
   export type Props = PropsWithChildren;
@@ -10,7 +10,7 @@ export namespace ThreadPrimitiveEmpty {
 export const ThreadPrimitiveEmpty: FC<ThreadPrimitiveEmpty.Props> = ({
   children,
 }) => {
-  const empty = useThread((u) => u.messages.length === 0);
+  const empty = useAssistantState(({ thread }) => thread.messages.length === 0);
   return empty ? children : null;
 };
 

package/src/primitives/thread/ThreadIf.tsx

@@ -1,7 +1,7 @@
 "use client";
 
 import type { FC, PropsWithChildren } from "react";
-import { useThread } from "../../context";
+import { useAssistantState } from "../../context";
 import type { RequireAtLeastOne } from "../../utils/RequireAtLeastOne";
 
 type ThreadIfFilters = {
@@ -13,7 +13,7 @@ type ThreadIfFilters = {
 type UseThreadIfProps = RequireAtLeastOne<ThreadIfFilters>;
 
 const useThreadIf = (props: UseThreadIfProps) => {
-  return useThread((thread) => {
+  return useAssistantState(({ thread }) => {
     if (props.empty === true && thread.messages.length !== 0) return false;
     if (props.empty === false && thread.messages.length === 0) return false;
     if (props.running === true && !thread.isRunning) return false;

package/src/primitives/thread/ThreadMessages.tsx

@@ -1,9 +1,7 @@
 "use client";
 
 import { type ComponentType, type FC, memo, useMemo } from "react";
-import { useThread, useThreadRuntime } from "../../context/react/ThreadContext";
-import { MessageRuntimeProvider } from "../../context/providers/MessageRuntimeProvider";
-import { useEditComposer, useMessage } from "../../context";
+import { useAssistantState, MessageByIndexProvider } from "../../context";
 import { ThreadMessage as ThreadMessageType } from "../../types";
 
 export namespace ThreadPrimitiveMessages {
@@ -130,8 +128,10 @@ type ThreadMessageComponentProps = {
 const ThreadMessageComponent: FC<ThreadMessageComponentProps> = ({
   components,
 }) => {
-  const role = useMessage((m) => m.role);
-  const isEditing = useEditComposer((c) => c.isEditing);
+  const role = useAssistantState(({ message }) => message.role);
+  const isEditing = useAssistantState(
+    ({ message }) => message.composer.isEditing,
+  );
   const Component = getComponent(components, role, isEditing);
 
   return <Component />;
@@ -163,16 +163,10 @@ export namespace ThreadPrimitiveMessageByIndex {
 export const ThreadPrimitiveMessageByIndex: FC<ThreadPrimitiveMessageByIndex.Props> =
   memo(
     ({ index, components }) => {
-      const threadRuntime = useThreadRuntime();
-      const runtime = useMemo(
-        () => threadRuntime.getMesssageByIndex(index),
-        [threadRuntime, index],
-      );
-
       return (
-        <MessageRuntimeProvider runtime={runtime}>
+        <MessageByIndexProvider index={index}>
           <ThreadMessageComponent components={components} />
-        </MessageRuntimeProvider>
+        </MessageByIndexProvider>
       );
     },
     (prev, next) =>
@@ -203,7 +197,9 @@ ThreadPrimitiveMessageByIndex.displayName = "ThreadPrimitive.MessageByIndex";
 export const ThreadPrimitiveMessagesImpl: FC<ThreadPrimitiveMessages.Props> = ({
   components,
 }) => {
-  const messagesLength = useThread((t) => t.messages.length);
+  const messagesLength = useAssistantState(
+    ({ thread }) => thread.messages.length,
+  );
 
   const messageElements = useMemo(() => {
     if (messagesLength === 0) return null;