@assistant-ui/react 0.10.50 → 0.11.1

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

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(() => {