@assistant-ui/react 0.0.3 → 0.0.5

package/dist/index.mjs

@@ -25,92 +25,22 @@ var ThreadRoot = forwardRef(
   }
 );
 
-// src/utils/context/createStoreContext.tsx
-import {
-  createContext,
-  useContext,
-  useMemo,
-  useState
-} from "react";
-import { useSyncExternalStoreWithSelector } from "use-sync-external-store/with-selector";
-
-// src/utils/context/shallow.ts
-function shallow(objA, objB) {
-  if (Object.is(objA, objB)) {
-    return true;
-  }
-  if (typeof objA !== "object" || objA === null || typeof objB !== "object" || objB === null) {
-    return false;
-  }
-  const keysA = Object.keys(objA);
-  if (keysA.length !== Object.keys(objB).length) {
-    return false;
-  }
-  for (const keyA of keysA) {
-    if (!Object.prototype.hasOwnProperty.call(objB, keyA) || !Object.is(objA[keyA], objB[keyA])) {
-      return false;
-    }
-  }
-  return true;
-}
-
-// src/utils/context/createStoreContext.tsx
-var createStoreContext = (providerName) => {
-  const context = createContext(null);
-  const StoreProvider = ({ children, ...rest }) => {
-    const unstableContext = rest;
-    const [store] = useState(() => {
-      let state = unstableContext;
-      const listeners = /* @__PURE__ */ new Set();
-      return {
-        subscribe: (cb) => {
-          listeners.add(cb);
-          return () => listeners.delete(cb);
-        },
-        emit: () => {
-          for (const listener of listeners) {
-            listener();
-          }
-        },
-        snapshot: () => {
-          return state;
-        },
-        setState: (value) => {
-          state = value;
-          store.emit();
-        }
-      };
-    });
-    useMemo(
-      () => store.setState(unstableContext),
-      Object.values(unstableContext)
-    );
-    return /* @__PURE__ */ React.createElement(context.Provider, { value: store }, children);
-  };
-  const useStoreContext = (consumerName, selector) => {
-    const store = useContext(context);
-    if (!store)
-      throw new Error(
-        `${consumerName} can only be used inside ${providerName}.`
-      );
-    return useSyncExternalStoreWithSelector(
-      store.subscribe,
-      store.snapshot,
-      store.snapshot,
-      selector,
-      shallow
+// src/utils/context/AssistantContext.ts
+import { createContext, useContext } from "react";
+var AssistantContext = createContext(null);
+var useAssistantContext = () => {
+  const context = useContext(AssistantContext);
+  if (!context)
+    throw new Error(
+      "useAssistantContext must be used within a AssistantProvider"
     );
-  };
-  return [StoreProvider, useStoreContext];
+  return context;
 };
 
-// src/utils/context/ThreadContext.ts
-var [ThreadContextProvider, useThreadContext] = createStoreContext("Thread.Provider");
-
 // src/primitives/thread/ThreadIf.tsx
 var useThreadIf = (props) => {
-  return useThreadContext("Thread.If", (s) => {
-    const thread = s.chat;
+  const { useThread } = useAssistantContext();
+  return useThread((thread) => {
     if (props.empty === true && thread.messages.length !== 0)
       return false;
     if (props.empty === false && thread.messages.length === 0)
@@ -133,7 +63,7 @@ var ThreadEmpty = ({ children }) => {
 };
 
 // src/primitives/thread/ThreadViewport.tsx
-import { forwardRef as forwardRef2, useRef as useRef2, useState as useState2 } from "react";
+import { forwardRef as forwardRef2, useRef as useRef2, useState } from "react";
 import {
   Primitive as Primitive2
 } from "@radix-ui/react-primitive";
@@ -185,7 +115,7 @@ var useOnResizeContent = (ref, callback) => {
 var ThreadViewport = forwardRef2(({ onScroll, children, ...rest }, forwardedRef) => {
   const divRef = useRef2(null);
   const ref = useComposedRefs(forwardedRef, divRef);
-  const [isAtBottom, setIsAtBottom] = useState2(true);
+  const [isAtBottom, setIsAtBottom] = useState(true);
   useOnResizeContent(divRef, () => {
     const div = divRef.current;
     if (!div || !isAtBottom)
@@ -211,7 +141,7 @@ var ThreadViewport = forwardRef2(({ onScroll, children, ...rest }, forwardedRef)
 });
 
 // src/utils/hooks/useBranches.tsx
-import { useCallback, useMemo as useMemo2, useRef as useRef3 } from "react";
+import { useCallback, useMemo, useRef as useRef3 } from "react";
 var ROOT_ID = "__ROOT_ID__";
 var UPCOMING_MESSAGE_ID = "__UPCOMING_MESSAGE_ID__";
 var updateBranchData = (data, messages) => {
@@ -319,80 +249,154 @@ var useChatWithBranches = (chat) => {
     },
     [chat.messages, chat.setMessages, chat.append]
   );
-  return useMemo2(
+  return useMemo(
     () => ({
-      ...chat,
       getBranchState,
       switchToBranch,
       editAt,
       reloadAt
     }),
-    [chat, getBranchState, switchToBranch, editAt, reloadAt]
+    [getBranchState, switchToBranch, editAt, reloadAt]
   );
 };
+var hasUpcomingMessage = (thread) => {
+  return thread.isLoading && thread.messages[thread.messages.length - 1]?.role !== "assistant";
+};
+
+// src/utils/context/ComposerState.ts
+import { useContext as useContext3 } from "react";
+
+// src/utils/context/MessageContext.ts
+import { createContext as createContext2, useContext as useContext2 } from "react";
+var MessageContext = createContext2(null);
+var useMessageContext = () => {
+  const context = useContext2(MessageContext);
+  if (!context)
+    throw new Error("useMessageContext must be used within a MessageProvider");
+  return context;
+};
+
+// src/utils/context/ComposerState.ts
+var useComposerContext = () => {
+  const { useComposer: useAssisstantComposer } = useAssistantContext();
+  const { useComposer: useMessageComposer } = useContext3(MessageContext) ?? {};
+  return { useComposer: useMessageComposer ?? useAssisstantComposer };
+};
+
+// src/primitives/composer/ComposerIf.tsx
+var useComposerIf = (props) => {
+  const { useComposer } = useComposerContext();
+  return useComposer((composer) => {
+    if (props.editing === true && !composer.isEditing)
+      return false;
+    if (props.editing === false && composer.isEditing)
+      return false;
+    return true;
+  });
+};
+var ComposerIf = ({ children, ...query }) => {
+  const result = useComposerIf(query);
+  return result ? children : null;
+};
 
 // src/primitives/message/index.ts
 var message_exports = {};
 __export(message_exports, {
   Content: () => MessageContent,
-  EditableContent: () => MessageEditableContent,
   If: () => MessageIf,
   Provider: () => MessageProvider,
   Root: () => MessageRoot
 });
 
 // src/primitives/message/MessageProvider.tsx
-import { useMemo as useMemo3, useState as useState3 } from "react";
-
-// src/utils/context/MessageContext.ts
-var [MessageContextProvider, useMessageContext] = createStoreContext("Thread.Provider");
-
-// src/primitives/message/MessageProvider.tsx
+import { useMemo as useMemo2, useState as useState2 } from "react";
+import { create } from "zustand";
+import { useShallow } from "zustand/react/shallow";
+var getIsLast = (thread, message) => {
+  const hasUpcoming = hasUpcomingMessage(thread);
+  return hasUpcoming ? message.id === UPCOMING_MESSAGE_ID : thread.messages[thread.messages.length - 1]?.id === message.id;
+};
+var useMessageContext2 = () => {
+  const { useBranchObserver } = useAssistantContext();
+  const [context] = useState2(() => {
+    const useMessage = create(() => ({
+      message: null,
+      isLast: false,
+      isCopied: false,
+      isHovering: false,
+      setIsCopied: () => {
+      },
+      setIsHovering: () => {
+      },
+      branchState: {
+        branchId: 0,
+        branchCount: 0
+      }
+    }));
+    const useComposer = create((set, get) => ({
+      isEditing: false,
+      canCancel: true,
+      edit: () => set({
+        isEditing: true,
+        value: useMessage.getState().message.content
+      }),
+      cancel: () => set({ isEditing: false }),
+      send: () => {
+        const message = useMessage.getState().message;
+        useBranchObserver.getState().editAt(message, {
+          ...message,
+          id: void 0,
+          // remove id to create a new message
+          content: get().value
+        });
+        set({ isEditing: false });
+      },
+      value: "",
+      setValue: (value) => set({ value })
+    }));
+    return { useMessage, useComposer };
+  });
+  return context;
+};
 var MessageProvider = ({
   message,
   children
 }) => {
-  const getBranchState = useThreadContext(
-    "Message.Provider",
-    (s) => s.chat.getBranchState
-  );
-  const [editState, setEditState] = useState3({
-    isEditing: false
-  });
-  const [isCopied, setIsCopied] = useState3(false);
-  const [isHovering, setIsHovering] = useState3(false);
-  const branchState = useMemo3(
-    () => getBranchState(message),
-    [getBranchState, message]
-  );
-  return /* @__PURE__ */ React.createElement(
-    MessageContextProvider,
-    {
-      message,
-      editState,
-      setEditState,
-      branchState,
-      isCopied,
-      setIsCopied,
-      isHovering,
-      setIsHovering
-    },
-    children
+  const { useThread, useBranchObserver } = useAssistantContext();
+  const context = useMessageContext2();
+  const branchState = useBranchObserver(
+    useShallow((b) => b.getBranchState(message))
   );
+  const isLast = useThread((thread) => getIsLast(thread, message));
+  const [isCopied, setIsCopied] = useState2(false);
+  const [isHovering, setIsHovering] = useState2(false);
+  useMemo2(() => {
+    context.useMessage.setState(
+      {
+        message,
+        isLast,
+        isCopied,
+        isHovering,
+        setIsCopied,
+        setIsHovering,
+        branchState
+      },
+      true
+    );
+  }, [context, message, isLast, isCopied, isHovering, branchState]);
+  return /* @__PURE__ */ React.createElement(MessageContext.Provider, { value: context }, children);
 };
 
 // src/primitives/message/MessageRoot.tsx
-import { forwardRef as forwardRef3 } from "react";
+import { composeEventHandlers as composeEventHandlers2 } from "@radix-ui/primitive";
 import {
   Primitive as Primitive3
 } from "@radix-ui/react-primitive";
-import { composeEventHandlers as composeEventHandlers2 } from "@radix-ui/primitive";
+import { forwardRef as forwardRef3 } from "react";
 var MessageRoot = forwardRef3(
   ({ onMouseEnter, onMouseLeave, ...rest }, ref) => {
-    const setIsHovering = useMessageContext(
-      "Message.Root",
-      (s) => s.setIsHovering
-    );
+    const { useMessage } = useMessageContext();
+    const setIsHovering = useMessage((s) => s.setIsHovering);
     const handleMouseEnter = () => {
       setIsHovering(true);
     };
@@ -413,22 +417,22 @@ var MessageRoot = forwardRef3(
 
 // src/primitives/message/MessageIf.tsx
 var useMessageIf = (props) => {
-  const thread = useThreadContext("Message.If", (s) => s.chat);
-  return useMessageContext(
-    "Message.If",
-    ({ message, editState: { isEditing }, isCopied, isHovering }) => {
-      const { branchCount } = thread.getBranchState(message);
+  const { useMessage } = useMessageContext();
+  return useMessage(
+    ({
+      message,
+      isLast,
+      isCopied,
+      isHovering,
+      branchState: { branchCount }
+    }) => {
       if (props.hasBranches === true && branchCount < 2)
         return false;
       if (props.user && message.role !== "user")
         return false;
       if (props.assistant && message.role !== "assistant")
         return false;
-      if (props.editing === true && !isEditing)
-        return false;
-      if (props.editing === false && isEditing)
-        return false;
-      if (props.unstable_hoveringOrLast === true && !isHovering && thread.messages[thread.messages.length - 1]?.id !== message.id)
+      if (props.lastOrHover === true && !isHovering && !isLast)
         return false;
       if (props.copied === true && !isCopied)
         return false;
@@ -445,60 +449,32 @@ var MessageIf = ({ children, ...query }) => {
 
 // src/primitives/message/MessageContent.tsx
 var MessageContent = () => {
-  const content = useMessageContext(
-    "Message.Content",
-    (s) => s.message.content
-  );
+  const { useMessage } = useMessageContext();
+  const content = useMessage((s) => s.message.content);
   return /* @__PURE__ */ React.createElement(React.Fragment, null, content);
 };
 
-// src/primitives/message/MessageEditableContent.tsx
-import { forwardRef as forwardRef4 } from "react";
-import { composeEventHandlers as composeEventHandlers3 } from "@radix-ui/primitive";
-import TextareaAutosize from "react-textarea-autosize";
-var MessageEditableContent = forwardRef4(({ onChange, value, ...rest }, forwardedRef) => {
-  const [editState, setEditState] = useMessageContext(
-    "Message.EditableContent",
-    (s) => [s.editState, s.setEditState]
-  );
-  const handleChange = (e) => {
-    setEditState({ isEditing: true, value: e.target.value });
-  };
-  if (!editState.isEditing)
-    throw new Error(
-      "Message.EditableContent may only be rendered when edit mode is enabled. Consider wrapping the component in <Message.If editing>."
-    );
-  return /* @__PURE__ */ React.createElement(
-    TextareaAutosize,
-    {
-      ...rest,
-      ref: forwardedRef,
-      onChange: composeEventHandlers3(onChange, handleChange),
-      value: editState.value || value
-    }
-  );
-});
-
 // src/primitives/thread/ThreadMessages.tsx
 var getComponents = (components) => {
   return {
-    EditingUserMessage: components.EditingUserMessage ?? components.UserMessage ?? components.Message,
+    EditComposer: components.EditComposer ?? components.UserMessage ?? components.Message,
     UserMessage: components.UserMessage ?? components.Message,
     AssistantMessage: components.AssistantMessage ?? components.Message
   };
 };
 var ThreadMessages = ({ components }) => {
-  const chat = useThreadContext("Thread.Messages", (s) => s.chat);
-  const messages = chat.messages;
-  const { UserMessage, EditingUserMessage, AssistantMessage } = getComponents(components);
+  const { useThread } = useAssistantContext();
+  const thread = useThread();
+  const messages = thread.messages;
+  const { UserMessage, EditComposer, AssistantMessage } = getComponents(components);
   if (messages.length === 0)
     return null;
   return /* @__PURE__ */ React.createElement(React.Fragment, null, messages.map((message, idx) => {
     return (
       // biome-ignore lint/suspicious/noArrayIndexKey: fixes a11y issues with branch navigation
-      /* @__PURE__ */ React.createElement(MessageProvider, { key: idx, message }, /* @__PURE__ */ React.createElement(MessageIf, { user: true, editing: false }, /* @__PURE__ */ React.createElement(UserMessage, null)), /* @__PURE__ */ React.createElement(MessageIf, { user: true, editing: true }, /* @__PURE__ */ React.createElement(EditingUserMessage, null)), /* @__PURE__ */ React.createElement(MessageIf, { assistant: true }, /* @__PURE__ */ React.createElement(AssistantMessage, null)))
+      /* @__PURE__ */ React.createElement(MessageProvider, { key: idx, message }, /* @__PURE__ */ React.createElement(MessageIf, { user: true }, /* @__PURE__ */ React.createElement(ComposerIf, { editing: false }, /* @__PURE__ */ React.createElement(UserMessage, null)), /* @__PURE__ */ React.createElement(ComposerIf, { editing: true }, /* @__PURE__ */ React.createElement(EditComposer, null))), /* @__PURE__ */ React.createElement(MessageIf, { assistant: true }, /* @__PURE__ */ React.createElement(AssistantMessage, null)))
     );
-  }), chat.isLoading && chat.messages[chat.messages.length - 1]?.role !== "assistant" && /* @__PURE__ */ React.createElement(
+  }), hasUpcomingMessage(thread) && /* @__PURE__ */ React.createElement(
     MessageProvider,
     {
       message: {
@@ -514,22 +490,29 @@ var ThreadMessages = ({ components }) => {
 // src/primitives/composer/index.ts
 var composer_exports = {};
 __export(composer_exports, {
+  Cancel: () => ComposerCancel,
+  If: () => ComposerIf,
   Input: () => ComposerInput,
   Root: () => ComposerRoot,
-  Send: () => ComposerSend,
-  Stop: () => ComposerStop
+  Send: () => ComposerSend
 });
 
 // src/primitives/composer/ComposerRoot.tsx
-import { createContext as createContext2, forwardRef as forwardRef5, useContext as useContext2, useMemo as useMemo4, useRef as useRef4 } from "react";
+import { composeEventHandlers as composeEventHandlers3 } from "@radix-ui/primitive";
+import { useComposedRefs as useComposedRefs2 } from "@radix-ui/react-compose-refs";
 import {
   Primitive as Primitive4
 } from "@radix-ui/react-primitive";
-import { composeEventHandlers as composeEventHandlers4 } from "@radix-ui/primitive";
-import { useComposedRefs as useComposedRefs2 } from "@radix-ui/react-compose-refs";
-var ComposerContext = createContext2(null);
-var useComposerContext = () => {
-  const context = useContext2(ComposerContext);
+import {
+  createContext as createContext3,
+  forwardRef as forwardRef4,
+  useContext as useContext4,
+  useMemo as useMemo3,
+  useRef as useRef4
+} from "react";
+var ComposerFormContext = createContext3(null);
+var useComposerFormContext = () => {
+  const context = useContext4(ComposerFormContext);
   if (!context) {
     throw new Error(
       "Composer compound components cannot be rendered outside the Composer component"
@@ -537,15 +520,12 @@ var useComposerContext = () => {
   }
   return context;
 };
-var ComposerRoot = forwardRef5(
+var ComposerRoot = forwardRef4(
   ({ onSubmit, ...rest }, forwardedRef) => {
-    const handleSubmit = useThreadContext(
-      "Composer.Root",
-      (s) => s.chat.handleSubmit
-    );
+    const { useComposer } = useComposerContext();
     const formRef = useRef4(null);
     const ref = useComposedRefs2(forwardedRef, formRef);
-    const composerContextValue = useMemo4(
+    const composerContextValue = useMemo3(
       () => ({
         submit: () => formRef.current?.dispatchEvent(
           new Event("submit", { cancelable: true, bubbles: true })
@@ -553,108 +533,115 @@ var ComposerRoot = forwardRef5(
       }),
       []
     );
-    return /* @__PURE__ */ React.createElement(ComposerContext.Provider, { value: composerContextValue }, /* @__PURE__ */ React.createElement(
+    const handleSubmit = (e) => {
+      const composerState = useComposer.getState();
+      if (!composerState.isEditing)
+        return;
+      e.preventDefault();
+      composerState.send();
+    };
+    return /* @__PURE__ */ React.createElement(ComposerFormContext.Provider, { value: composerContextValue }, /* @__PURE__ */ React.createElement(
       Primitive4.form,
       {
         ...rest,
         ref,
-        onSubmit: composeEventHandlers4(onSubmit, handleSubmit)
+        onSubmit: composeEventHandlers3(onSubmit, handleSubmit)
       }
     ));
   }
 );
 
 // src/primitives/composer/ComposerInput.tsx
-import { forwardRef as forwardRef6 } from "react";
+import { composeEventHandlers as composeEventHandlers4 } from "@radix-ui/primitive";
 import { Slot } from "@radix-ui/react-slot";
-import { composeEventHandlers as composeEventHandlers5 } from "@radix-ui/primitive";
-import TextareaAutosize2 from "react-textarea-autosize";
-var ComposerInput = forwardRef6(({ asChild, onChange, onKeyDown, ...rest }, forwardedRef) => {
-  const chat = useThreadContext(
-    "Composer.Input",
-    ({ chat: { input, handleInputChange, isLoading } }) => ({
-      input,
-      handleInputChange,
-      isLoading
-    })
-  );
-  const Component = asChild ? Slot : TextareaAutosize2;
-  const composer = useComposerContext();
+import { forwardRef as forwardRef5 } from "react";
+import TextareaAutosize from "react-textarea-autosize";
+var ComposerInput = forwardRef5(({ asChild, disabled, onChange, onKeyDown, ...rest }, forwardedRef) => {
+  const { useThread } = useAssistantContext();
+  const isLoading = useThread((t) => t.isLoading);
+  const { useComposer } = useComposerContext();
+  const value = useComposer((c) => {
+    if (!c.isEditing)
+      return "";
+    return c.value;
+  });
+  const Component = asChild ? Slot : TextareaAutosize;
+  const composerForm = useComposerFormContext();
   const handleKeyPress = (e) => {
-    if (chat.isLoading || rest.disabled)
+    if (disabled)
+      return;
+    if (e.key === "Escape") {
+      useComposer.getState().cancel();
+    }
+    if (isLoading)
       return;
     if (e.key === "Enter" && e.shiftKey === false) {
       e.preventDefault();
-      composer.submit();
+      composerForm.submit();
     }
   };
   return /* @__PURE__ */ React.createElement(
     Component,
     {
-      value: chat.input,
+      value,
       ...rest,
       ref: forwardedRef,
-      onChange: composeEventHandlers5(onChange, chat.handleInputChange),
-      onKeyDown: composeEventHandlers5(onKeyDown, handleKeyPress)
+      disabled,
+      onChange: composeEventHandlers4(onChange, (e) => {
+        const composerState = useComposer.getState();
+        if (!composerState.isEditing)
+          return;
+        return composerState.setValue(e.target.value);
+      }),
+      onKeyDown: composeEventHandlers4(onKeyDown, handleKeyPress)
     }
   );
 });
 
 // src/primitives/composer/ComposerSend.tsx
-import { forwardRef as forwardRef7 } from "react";
 import {
   Primitive as Primitive5
 } from "@radix-ui/react-primitive";
-var ComposerSend = forwardRef7(
+import { forwardRef as forwardRef6 } from "react";
+var ComposerSend = forwardRef6(
   ({ disabled, ...rest }, ref) => {
-    const input = useThreadContext("Composer.Send", (s) => s.chat.input);
+    const { useComposer } = useComposerContext();
+    const hasValue = useComposer((c) => c.isEditing && c.value.length > 0);
     return /* @__PURE__ */ React.createElement(
       Primitive5.button,
       {
         type: "submit",
         ...rest,
         ref,
-        disabled: disabled || input.length === 0
+        disabled: disabled || !hasValue
       }
     );
   }
 );
 
-// src/utils/createActionButton.tsx
-import { forwardRef as forwardRef8 } from "react";
+// src/primitives/composer/ComposerCancel.tsx
+import { composeEventHandlers as composeEventHandlers5 } from "@radix-ui/primitive";
 import {
   Primitive as Primitive6
 } from "@radix-ui/react-primitive";
-import { composeEventHandlers as composeEventHandlers6 } from "@radix-ui/primitive";
-var createActionButton = (useActionButton) => {
-  return forwardRef8(
-    (props, forwardedRef) => {
-      const onClick = useActionButton(props);
-      return /* @__PURE__ */ React.createElement(
-        Primitive6.button,
-        {
-          type: "button",
-          disabled: !onClick,
-          ...props,
-          ref: forwardedRef,
-          onClick: composeEventHandlers6(props.onClick, onClick ?? void 0)
-        }
-      );
+import { forwardRef as forwardRef7 } from "react";
+var ComposerCancel = forwardRef7(({ disabled, onClick, ...rest }, ref) => {
+  const { useComposer } = useComposerContext();
+  const hasValue = useComposer((c) => c.canCancel);
+  const handleClose = () => {
+    useComposer.getState().cancel();
+  };
+  return /* @__PURE__ */ React.createElement(
+    Primitive6.button,
+    {
+      type: "button",
+      ...rest,
+      ref,
+      onClick: composeEventHandlers5(onClick, handleClose),
+      disabled: disabled || !hasValue
     }
   );
-};
-
-// src/primitives/composer/ComposerStop.tsx
-var useComposerStop = () => {
-  const [isLoading, stop] = useThreadContext("Composer.Stop", (s) => [
-    s.chat.isLoading,
-    s.chat.stop
-  ]);
-  if (!isLoading)
-    return null;
-  return stop;
-};
-var ComposerStop = createActionButton(useComposerStop);
+});
 
 // src/primitives/branchPicker/index.ts
 var branchPicker_exports = {};
@@ -668,52 +655,66 @@ __export(branchPicker_exports, {
 
 // src/actions/useGoToNextBranch.tsx
 var useGoToNextBranch = () => {
-  const switchToBranch = useThreadContext(
-    "BranchPicker.Next",
-    (s) => s.chat.switchToBranch
+  const { useThread, useBranchObserver } = useAssistantContext();
+  const { useComposer, useMessage } = useMessageContext();
+  const isLoading = useThread((s) => s.isLoading);
+  const isEditing = useComposer((s) => s.isEditing);
+  const hasNext = useMessage(
+    ({ branchState: { branchId, branchCount } }) => branchId + 1 < branchCount
   );
-  const context = useMessageContext("BranchPicker.Next", (s) => {
-    const {
-      message: message2,
-      editState: { isEditing },
-      branchState: { branchId: branchId2, branchCount }
-    } = s;
-    if (isEditing || branchCount <= 1 || branchId2 + 1 >= branchCount)
-      return null;
-    return { message: message2, branchId: branchId2 };
-  });
-  if (!context)
+  if (isLoading || isEditing || !hasNext)
     return null;
-  const { message, branchId } = context;
   return () => {
-    switchToBranch(message, branchId + 1);
+    const {
+      message,
+      branchState: { branchId }
+    } = useMessage.getState();
+    useBranchObserver.getState().switchToBranch(message, branchId + 1);
   };
 };
 
+// src/utils/createActionButton.tsx
+import { forwardRef as forwardRef8 } from "react";
+import {
+  Primitive as Primitive7
+} from "@radix-ui/react-primitive";
+import { composeEventHandlers as composeEventHandlers6 } from "@radix-ui/primitive";
+var createActionButton = (useActionButton) => {
+  return forwardRef8(
+    (props, forwardedRef) => {
+      const onClick = useActionButton(props);
+      return /* @__PURE__ */ React.createElement(
+        Primitive7.button,
+        {
+          type: "button",
+          disabled: !onClick,
+          ...props,
+          ref: forwardedRef,
+          onClick: composeEventHandlers6(props.onClick, onClick ?? void 0)
+        }
+      );
+    }
+  );
+};
+
 // src/primitives/branchPicker/BranchPickerNext.tsx
 var BranchPickerNext = createActionButton(useGoToNextBranch);
 
 // src/actions/useGoToPreviousBranch.tsx
 var useGoToPreviousBranch = () => {
-  const switchToBranch = useThreadContext(
-    "BranchPicker.Previous",
-    (s) => s.chat.switchToBranch
-  );
-  const context = useMessageContext("BranchPicker.Previous", (s) => {
-    const {
-      message: message2,
-      editState: { isEditing },
-      branchState: { branchId: branchId2, branchCount }
-    } = s;
-    if (isEditing || branchCount <= 1 || branchId2 <= 0)
-      return null;
-    return { message: message2, branchId: branchId2 };
-  });
-  if (!context)
+  const { useThread, useBranchObserver } = useAssistantContext();
+  const { useComposer, useMessage } = useMessageContext();
+  const isLoading = useThread((s) => s.isLoading);
+  const isEditing = useComposer((s) => s.isEditing);
+  const hasNext = useMessage(({ branchState: { branchId } }) => branchId > 0);
+  if (isLoading || isEditing || !hasNext)
     return null;
-  const { message, branchId } = context;
   return () => {
-    switchToBranch(message, branchId - 1);
+    const {
+      message,
+      branchState: { branchId }
+    } = useMessage.getState();
+    useBranchObserver.getState().switchToBranch(message, branchId - 1);
   };
 };
 
@@ -722,29 +723,25 @@ var BranchPickerPrevious = createActionButton(useGoToPreviousBranch);
 
 // src/primitives/branchPicker/BranchPickerCount.tsx
 var BranchPickerCount = () => {
-  const branchCount = useMessageContext(
-    "BranchPicker.Count",
-    (s) => s.branchState.branchCount
-  );
+  const { useMessage } = useMessageContext();
+  const branchCount = useMessage((s) => s.branchState.branchCount);
   return /* @__PURE__ */ React.createElement(React.Fragment, null, branchCount);
 };
 
 // src/primitives/branchPicker/BranchPickerNumber.tsx
 var BranchPickerNumber = () => {
-  const branchId = useMessageContext(
-    "BranchPicker.Number",
-    (s) => s.branchState.branchId
-  );
+  const { useMessage } = useMessageContext();
+  const branchId = useMessage((s) => s.branchState.branchId);
   return /* @__PURE__ */ React.createElement(React.Fragment, null, branchId + 1);
 };
 
 // src/primitives/branchPicker/BranchPickerRoot.tsx
 import {
-  Primitive as Primitive7
+  Primitive as Primitive8
 } from "@radix-ui/react-primitive";
 import { forwardRef as forwardRef9 } from "react";
 var BranchPickerRoot = forwardRef9(({ hideWhenSingleBranch, ...rest }, ref) => {
-  return /* @__PURE__ */ React.createElement(MessageIf, { hasBranches: hideWhenSingleBranch ? true : void 0 }, /* @__PURE__ */ React.createElement(Primitive7.div, { ...rest, ref }));
+  return /* @__PURE__ */ React.createElement(MessageIf, { hasBranches: hideWhenSingleBranch ? true : void 0 }, /* @__PURE__ */ React.createElement(Primitive8.div, { ...rest, ref }));
 });
 
 // src/primitives/actionBar/index.ts
@@ -758,29 +755,48 @@ __export(actionBar_exports, {
 
 // src/primitives/actionBar/ActionBarRoot.tsx
 import {
-  Primitive as Primitive8
+  Primitive as Primitive9
 } from "@radix-ui/react-primitive";
 import { forwardRef as forwardRef10 } from "react";
-var ActionBarRoot = forwardRef10(({ ...rest }, ref) => {
-  return /* @__PURE__ */ React.createElement(Primitive8.div, { ...rest, ref });
+var ActionBarRoot = forwardRef10(({ hideWhenBusy, autohide, autohideFloat, ...rest }, ref) => {
+  const { useThread } = useAssistantContext();
+  const { useMessage } = useMessageContext();
+  const hideAndfloatStatus = useMessage((m) => {
+    const autohideEnabled = autohide === "always" || autohide === "not-last" && !m.isLast;
+    if (!autohideEnabled)
+      return "normal" /* Normal */;
+    if (!m.isHovering)
+      return "hidden" /* Hidden */;
+    if (autohideFloat === "always" || autohideFloat === "single-branch" && m.branchState.branchCount <= 1)
+      return "floating" /* Floating */;
+    return "normal" /* Normal */;
+  });
+  const busy = useThread((t) => t.isLoading);
+  if (hideWhenBusy && busy)
+    return null;
+  if (hideAndfloatStatus === "hidden" /* Hidden */)
+    return null;
+  return /* @__PURE__ */ React.createElement(
+    Primitive9.div,
+    {
+      "data-floating": hideAndfloatStatus === "floating" /* Floating */,
+      ...rest,
+      ref
+    }
+  );
 });
 
 // src/actions/useCopyMessage.tsx
 var useCopyMessage = ({ copiedDuration = 3e3 }) => {
-  const context = useMessageContext("ActionBar.Copy", (s) => {
-    const {
-      editState: { isEditing },
-      message: { content: content2 },
-      setIsCopied: setIsCopied2
-    } = s;
-    if (isEditing)
-      return null;
-    return { content: content2, setIsCopied: setIsCopied2 };
-  });
-  if (!context)
+  const { useMessage, useComposer } = useMessageContext();
+  const isEditing = useComposer((s) => s.isEditing);
+  if (isEditing)
     return null;
-  const { content, setIsCopied } = context;
   return () => {
+    const {
+      message: { content },
+      setIsCopied
+    } = useMessage.getState();
     navigator.clipboard.writeText(content);
     setIsCopied(true);
     setTimeout(() => setIsCopied(false), copiedDuration);
@@ -792,19 +808,15 @@ var ActionBarCopy = createActionButton(useCopyMessage);
 
 // src/actions/useReloadMessage.tsx
 var useReloadMessage = () => {
-  const [isLoading, reloadAt] = useThreadContext("ActionBar.Reload", (s) => [
-    s.chat.isLoading,
-    s.chat.reloadAt
-  ]);
-  const message = useMessageContext("ActionBar.Reload", (s) => {
-    const message2 = s.message;
-    if (message2.role !== "assistant" || isLoading)
-      return null;
-    return message2;
-  });
-  if (!message)
+  const { useThread, useBranchObserver } = useAssistantContext();
+  const { useMessage } = useMessageContext();
+  const isLoading = useThread((s) => s.isLoading);
+  const isAssistant = useMessage((s) => s.message.role === "assistant");
+  if (isLoading || !isAssistant)
     return null;
-  return () => reloadAt(message);
+  return () => {
+    useBranchObserver.getState().reloadAt(useMessage.getState().message);
+  };
 };
 
 // src/primitives/actionBar/ActionBarReload.tsx
@@ -812,116 +824,136 @@ var ActionBarReload = createActionButton(useReloadMessage);
 
 // src/actions/useBeginMessageEdit.tsx
 var useBeginMessageEdit = () => {
-  const context = useMessageContext("ActionBar.Edit", (s) => {
-    const {
-      message: { content: content2 },
-      editState: { isEditing },
-      setEditState: setEditState2
-    } = s;
-    if (isEditing)
-      return null;
-    return { content: content2, setEditState: setEditState2 };
-  });
-  if (!context)
+  const { useMessage, useComposer } = useMessageContext();
+  const isUser = useMessage((s) => s.message.role === "user");
+  const isEditing = useComposer((s) => s.isEditing);
+  if (!isUser || isEditing)
     return null;
-  const { content, setEditState } = context;
   return () => {
-    setEditState({ isEditing: true, value: content });
+    const { edit } = useComposer.getState();
+    edit();
   };
 };
 
 // src/primitives/actionBar/ActionBarEdit.tsx
 var ActionBarEdit = createActionButton(useBeginMessageEdit);
 
-// src/primitives/editBar/index.ts
-var editBar_exports = {};
-__export(editBar_exports, {
-  Cancel: () => EditBarCancel,
-  Root: () => EditBarRoot,
-  Save: () => EditBarSave
-});
-
-// src/primitives/editBar/EditBarRoot.tsx
-import {
-  Primitive as Primitive9
-} from "@radix-ui/react-primitive";
-import { forwardRef as forwardRef11 } from "react";
-var EditBarRoot = forwardRef11(
-  ({ ...rest }, ref) => {
-    return /* @__PURE__ */ React.createElement(Primitive9.div, { ...rest, ref });
-  }
-);
-
-// src/actions/useSaveMessageEdit.tsx
-var useSaveMessageEdit = () => {
-  const chat = useThreadContext("EditBar.Save", (s) => s.chat);
-  const context = useMessageContext("EditBar.Save", (s) => {
-    const { message: message2, editState, setEditState: setEditState2 } = s;
-    if (!editState.isEditing)
-      return null;
-    return { message: message2, content: editState.value, setEditState: setEditState2 };
-  });
-  if (!context)
-    return null;
-  const { message, content, setEditState } = context;
-  return () => {
-    chat.editAt(message, {
-      ...message,
-      id: void 0,
-      // remove id to create a new message
-      content
-    });
-    setEditState({ isEditing: false });
-  };
-};
-
-// src/primitives/editBar/EditBarSave.tsx
-var EditBarSave = createActionButton(useSaveMessageEdit);
-
-// src/actions/useCancelMessageEdit.tsx
-var useCancelMessageEdit = () => {
-  const context = useMessageContext("EditBar.Cancel", (s) => {
-    const {
-      editState: { isEditing },
-      setEditState: setEditState2
-    } = s;
-    if (!isEditing)
-      return null;
-    return { setEditState: setEditState2 };
+// src/vercel/VercelAIAssistantProvider.tsx
+import { useMemo as useMemo4, useState as useState3 } from "react";
+import { create as create2 } from "zustand";
+var useAIAssistantContext = () => {
+  const [context] = useState3(() => {
+    const useThread = create2()(() => ({
+      messages: [],
+      setMessages: () => {
+      },
+      isLoading: false,
+      reload: async () => {
+      },
+      append: async () => {
+      },
+      stop: () => {
+      }
+    }));
+    const useComposer = create2()(() => ({
+      isEditing: true,
+      canCancel: false,
+      value: "",
+      setValue: () => {
+      },
+      edit: () => {
+        throw new Error("Not implemented");
+      },
+      send: () => {
+        useThread.getState().append({
+          content: useComposer.getState().value,
+          role: "user",
+          createdAt: /* @__PURE__ */ new Date()
+        });
+        useComposer.getState().setValue("");
+      },
+      cancel: () => {
+        useThread.getState().stop();
+      }
+    }));
+    const useBranchObserver = create2()(() => ({
+      getBranchState: () => ({
+        branchId: 0,
+        branchCount: 0
+      }),
+      switchToBranch: () => {
+      },
+      editAt: async () => {
+      },
+      reloadAt: async () => {
+      }
+    }));
+    return { useThread, useComposer, useBranchObserver };
   });
-  if (!context)
-    return null;
-  const { setEditState } = context;
-  return () => {
-    setEditState({ isEditing: false });
-  };
+  return context;
 };
-
-// src/primitives/editBar/EditBarCancel.tsx
-var EditBarCancel = createActionButton(useCancelMessageEdit);
-
-// src/vercel/VercelAIThreadProvider.tsx
-var VercelAIThreadProvider = ({
+var VercelAIAssistantProvider = ({
   chat,
   children
 }) => {
+  const context = useAIAssistantContext();
+  useMemo4(() => {
+    context.useThread.setState(
+      {
+        messages: chat.messages,
+        setMessages: (value) => {
+          chat.setMessages(value);
+        },
+        isLoading: chat.isLoading,
+        reload: async () => {
+          await chat.reload();
+        },
+        append: async (message) => {
+          await chat.append(message);
+        },
+        stop: () => {
+          const lastMessage = chat.messages.at(-1);
+          chat.stop();
+          if (lastMessage?.role === "user") {
+            chat.setInput(lastMessage.content);
+          }
+        }
+      },
+      true
+    );
+  }, [context, chat]);
+  useMemo4(() => {
+    context.useComposer.setState({
+      canCancel: chat.isLoading,
+      value: chat.input,
+      setValue: chat.setInput
+    });
+  }, [context, chat.isLoading, chat.input, chat.setInput]);
   const branches = useChatWithBranches(chat);
-  return /* @__PURE__ */ React.createElement(ThreadContextProvider, { chat: branches }, children);
+  useMemo4(() => {
+    context.useBranchObserver.setState(
+      {
+        getBranchState: (message) => branches.getBranchState(message),
+        switchToBranch: (message, branchId) => branches.switchToBranch(message, branchId),
+        editAt: async (message, newMessage) => branches.editAt(message, newMessage),
+        reloadAt: async (message) => branches.reloadAt(message)
+      },
+      true
+    );
+  }, [context, branches]);
+  return /* @__PURE__ */ React.createElement(AssistantContext.Provider, { value: context }, children);
 };
 export {
   actionBar_exports as ActionBarPrimitive,
   branchPicker_exports as BranchPickerPrimitive,
   composer_exports as ComposerPrimitive,
-  editBar_exports as EditBarPrimitive,
   message_exports as MessagePrimitive,
   thread_exports as ThreadPrimitive,
-  VercelAIThreadProvider,
+  VercelAIAssistantProvider as VercelAIThreadProvider,
   useMessageContext as unstable_useMessageContext,
   useBeginMessageEdit,
-  useCancelMessageEdit,
   useCopyMessage,
   useGoToNextBranch,
   useGoToPreviousBranch,
-  useReloadMessage,
-  useSaveMessageEdit
+  useReloadMessage
 };