@assistant-ui/react 0.0.4 → 0.0.5

package/dist/index.js

@@ -33,18 +33,15 @@ __export(src_exports, {
   ActionBarPrimitive: () => actionBar_exports,
   BranchPickerPrimitive: () => branchPicker_exports,
   ComposerPrimitive: () => composer_exports,
-  EditBarPrimitive: () => editBar_exports,
   MessagePrimitive: () => message_exports,
   ThreadPrimitive: () => thread_exports,
-  VercelAIThreadProvider: () => VercelAIThreadProvider,
+  VercelAIThreadProvider: () => VercelAIAssistantProvider,
   unstable_useMessageContext: () => useMessageContext,
   useBeginMessageEdit: () => useBeginMessageEdit,
-  useCancelMessageEdit: () => useCancelMessageEdit,
   useCopyMessage: () => useCopyMessage,
   useGoToNextBranch: () => useGoToNextBranch,
   useGoToPreviousBranch: () => useGoToPreviousBranch,
-  useReloadMessage: () => useReloadMessage,
-  useSaveMessageEdit: () => useSaveMessageEdit
+  useReloadMessage: () => useReloadMessage
 });
 module.exports = __toCommonJS(src_exports);
 
@@ -67,87 +64,22 @@ var ThreadRoot = (0, import_react.forwardRef)(
   }
 );
 
-// src/utils/context/createStoreContext.tsx
+// src/utils/context/AssistantContext.ts
 var import_react2 = require("react");
-var import_with_selector = require("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 = (0, import_react2.createContext)(null);
-  const StoreProvider = ({ children, ...rest }) => {
-    const unstableContext = rest;
-    const [store] = (0, import_react2.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();
-        }
-      };
-    });
-    (0, import_react2.useMemo)(
-      () => store.setState(unstableContext),
-      Object.values(unstableContext)
-    );
-    return /* @__PURE__ */ React.createElement(context.Provider, { value: store }, children);
-  };
-  const useStoreContext = (consumerName, selector) => {
-    const store = (0, import_react2.useContext)(context);
-    if (!store)
-      throw new Error(
-        `${consumerName} can only be used inside ${providerName}.`
-      );
-    return (0, import_with_selector.useSyncExternalStoreWithSelector)(
-      store.subscribe,
-      store.snapshot,
-      store.snapshot,
-      selector,
-      shallow
+var AssistantContext = (0, import_react2.createContext)(null);
+var useAssistantContext = () => {
+  const context = (0, import_react2.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)
@@ -356,79 +288,150 @@ var useChatWithBranches = (chat) => {
   );
   return (0, import_react5.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
+var import_react7 = require("react");
+
+// src/utils/context/MessageContext.ts
+var import_react6 = require("react");
+var MessageContext = (0, import_react6.createContext)(null);
+var useMessageContext = () => {
+  const context = (0, import_react6.useContext)(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 } = (0, import_react7.useContext)(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
-var import_react6 = require("react");
-
-// src/utils/context/MessageContext.ts
-var [MessageContextProvider, useMessageContext] = createStoreContext("Thread.Provider");
-
-// src/primitives/message/MessageProvider.tsx
+var import_react8 = require("react");
+var import_zustand = require("zustand");
+var import_shallow = require("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] = (0, import_react8.useState)(() => {
+    const useMessage = (0, import_zustand.create)(() => ({
+      message: null,
+      isLast: false,
+      isCopied: false,
+      isHovering: false,
+      setIsCopied: () => {
+      },
+      setIsHovering: () => {
+      },
+      branchState: {
+        branchId: 0,
+        branchCount: 0
+      }
+    }));
+    const useComposer = (0, import_zustand.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] = (0, import_react6.useState)({
-    isEditing: false
-  });
-  const [isCopied, setIsCopied] = (0, import_react6.useState)(false);
-  const [isHovering, setIsHovering] = (0, import_react6.useState)(false);
-  const branchState = (0, import_react6.useMemo)(
-    () => 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(
+    (0, import_shallow.useShallow)((b) => b.getBranchState(message))
   );
+  const isLast = useThread((thread) => getIsLast(thread, message));
+  const [isCopied, setIsCopied] = (0, import_react8.useState)(false);
+  const [isHovering, setIsHovering] = (0, import_react8.useState)(false);
+  (0, import_react8.useMemo)(() => {
+    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
-var import_react7 = require("react");
-var import_react_primitive3 = require("@radix-ui/react-primitive");
 var import_primitive2 = require("@radix-ui/primitive");
-var MessageRoot = (0, import_react7.forwardRef)(
+var import_react_primitive3 = require("@radix-ui/react-primitive");
+var import_react9 = require("react");
+var MessageRoot = (0, import_react9.forwardRef)(
   ({ 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);
     };
@@ -448,27 +451,23 @@ var MessageRoot = (0, import_react7.forwardRef)(
 );
 
 // src/primitives/message/MessageIf.tsx
-var isLast = (thread, message) => {
-  const hasUpcoming = hasUpcomingMessage(thread);
-  return hasUpcoming ? message.id === UPCOMING_MESSAGE_ID : thread.messages[thread.messages.length - 1]?.id === message.id;
-};
 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.lastOrHover === true && !isHovering && !isLast(thread, message))
+      if (props.lastOrHover === true && !isHovering && !isLast)
         return false;
       if (props.copied === true && !isCopied)
         return false;
@@ -485,60 +484,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
-var import_react8 = require("react");
-var import_primitive3 = require("@radix-ui/primitive");
-var import_react_textarea_autosize = __toESM(require("react-textarea-autosize"));
-var MessageEditableContent = (0, import_react8.forwardRef)(({ 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(
-    import_react_textarea_autosize.default,
-    {
-      ...rest,
-      ref: forwardedRef,
-      onChange: (0, import_primitive3.composeEventHandlers)(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)))
     );
-  }), hasUpcomingMessage(chat) && /* @__PURE__ */ React.createElement(
+  }), hasUpcomingMessage(thread) && /* @__PURE__ */ React.createElement(
     MessageProvider,
     {
       message: {
@@ -554,20 +525,21 @@ 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
-var import_react9 = require("react");
-var import_react_primitive4 = require("@radix-ui/react-primitive");
-var import_primitive4 = require("@radix-ui/primitive");
+var import_primitive3 = require("@radix-ui/primitive");
 var import_react_compose_refs2 = require("@radix-ui/react-compose-refs");
-var ComposerContext = (0, import_react9.createContext)(null);
-var useComposerContext = () => {
-  const context = (0, import_react9.useContext)(ComposerContext);
+var import_react_primitive4 = require("@radix-ui/react-primitive");
+var import_react10 = require("react");
+var ComposerFormContext = (0, import_react10.createContext)(null);
+var useComposerFormContext = () => {
+  const context = (0, import_react10.useContext)(ComposerFormContext);
   if (!context) {
     throw new Error(
       "Composer compound components cannot be rendered outside the Composer component"
@@ -575,15 +547,12 @@ var useComposerContext = () => {
   }
   return context;
 };
-var ComposerRoot = (0, import_react9.forwardRef)(
+var ComposerRoot = (0, import_react10.forwardRef)(
   ({ onSubmit, ...rest }, forwardedRef) => {
-    const handleSubmit = useThreadContext(
-      "Composer.Root",
-      (s) => s.chat.handleSubmit
-    );
-    const formRef = (0, import_react9.useRef)(null);
+    const { useComposer } = useComposerContext();
+    const formRef = (0, import_react10.useRef)(null);
     const ref = (0, import_react_compose_refs2.useComposedRefs)(forwardedRef, formRef);
-    const composerContextValue = (0, import_react9.useMemo)(
+    const composerContextValue = (0, import_react10.useMemo)(
       () => ({
         submit: () => formRef.current?.dispatchEvent(
           new Event("submit", { cancelable: true, bubbles: true })
@@ -591,104 +560,111 @@ var ComposerRoot = (0, import_react9.forwardRef)(
       }),
       []
     );
-    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(
       import_react_primitive4.Primitive.form,
       {
         ...rest,
         ref,
-        onSubmit: (0, import_primitive4.composeEventHandlers)(onSubmit, handleSubmit)
+        onSubmit: (0, import_primitive3.composeEventHandlers)(onSubmit, handleSubmit)
       }
     ));
   }
 );
 
 // src/primitives/composer/ComposerInput.tsx
-var import_react10 = require("react");
+var import_primitive4 = require("@radix-ui/primitive");
 var import_react_slot = require("@radix-ui/react-slot");
-var import_primitive5 = require("@radix-ui/primitive");
-var import_react_textarea_autosize2 = __toESM(require("react-textarea-autosize"));
-var ComposerInput = (0, import_react10.forwardRef)(({ asChild, onChange, onKeyDown, ...rest }, forwardedRef) => {
-  const chat = useThreadContext(
-    "Composer.Input",
-    ({ chat: { input, handleInputChange, isLoading } }) => ({
-      input,
-      handleInputChange,
-      isLoading
-    })
-  );
-  const Component = asChild ? import_react_slot.Slot : import_react_textarea_autosize2.default;
-  const composer = useComposerContext();
+var import_react11 = require("react");
+var import_react_textarea_autosize = __toESM(require("react-textarea-autosize"));
+var ComposerInput = (0, import_react11.forwardRef)(({ 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 ? import_react_slot.Slot : import_react_textarea_autosize.default;
+  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: (0, import_primitive5.composeEventHandlers)(onChange, chat.handleInputChange),
-      onKeyDown: (0, import_primitive5.composeEventHandlers)(onKeyDown, handleKeyPress)
+      disabled,
+      onChange: (0, import_primitive4.composeEventHandlers)(onChange, (e) => {
+        const composerState = useComposer.getState();
+        if (!composerState.isEditing)
+          return;
+        return composerState.setValue(e.target.value);
+      }),
+      onKeyDown: (0, import_primitive4.composeEventHandlers)(onKeyDown, handleKeyPress)
     }
   );
 });
 
 // src/primitives/composer/ComposerSend.tsx
-var import_react11 = require("react");
 var import_react_primitive5 = require("@radix-ui/react-primitive");
-var ComposerSend = (0, import_react11.forwardRef)(
+var import_react12 = require("react");
+var ComposerSend = (0, import_react12.forwardRef)(
   ({ 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(
       import_react_primitive5.Primitive.button,
       {
         type: "submit",
         ...rest,
         ref,
-        disabled: disabled || input.length === 0
+        disabled: disabled || !hasValue
       }
     );
   }
 );
 
-// src/utils/createActionButton.tsx
-var import_react12 = require("react");
+// src/primitives/composer/ComposerCancel.tsx
+var import_primitive5 = require("@radix-ui/primitive");
 var import_react_primitive6 = require("@radix-ui/react-primitive");
-var import_primitive6 = require("@radix-ui/primitive");
-var createActionButton = (useActionButton) => {
-  return (0, import_react12.forwardRef)(
-    (props, forwardedRef) => {
-      const onClick = useActionButton(props);
-      return /* @__PURE__ */ React.createElement(
-        import_react_primitive6.Primitive.button,
-        {
-          type: "button",
-          disabled: !onClick,
-          ...props,
-          ref: forwardedRef,
-          onClick: (0, import_primitive6.composeEventHandlers)(props.onClick, onClick ?? void 0)
-        }
-      );
+var import_react13 = require("react");
+var ComposerCancel = (0, import_react13.forwardRef)(({ disabled, onClick, ...rest }, ref) => {
+  const { useComposer } = useComposerContext();
+  const hasValue = useComposer((c) => c.canCancel);
+  const handleClose = () => {
+    useComposer.getState().cancel();
+  };
+  return /* @__PURE__ */ React.createElement(
+    import_react_primitive6.Primitive.button,
+    {
+      type: "button",
+      ...rest,
+      ref,
+      onClick: (0, import_primitive5.composeEventHandlers)(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 = {};
@@ -702,52 +678,64 @@ __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
+var import_react14 = require("react");
+var import_react_primitive7 = require("@radix-ui/react-primitive");
+var import_primitive6 = require("@radix-ui/primitive");
+var createActionButton = (useActionButton) => {
+  return (0, import_react14.forwardRef)(
+    (props, forwardedRef) => {
+      const onClick = useActionButton(props);
+      return /* @__PURE__ */ React.createElement(
+        import_react_primitive7.Primitive.button,
+        {
+          type: "button",
+          disabled: !onClick,
+          ...props,
+          ref: forwardedRef,
+          onClick: (0, import_primitive6.composeEventHandlers)(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);
   };
 };
 
@@ -756,27 +744,23 @@ 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
-var import_react_primitive7 = require("@radix-ui/react-primitive");
-var import_react13 = require("react");
-var BranchPickerRoot = (0, import_react13.forwardRef)(({ hideWhenSingleBranch, ...rest }, ref) => {
-  return /* @__PURE__ */ React.createElement(MessageIf, { hasBranches: hideWhenSingleBranch ? true : void 0 }, /* @__PURE__ */ React.createElement(import_react_primitive7.Primitive.div, { ...rest, ref }));
+var import_react_primitive8 = require("@radix-ui/react-primitive");
+var import_react15 = require("react");
+var BranchPickerRoot = (0, import_react15.forwardRef)(({ hideWhenSingleBranch, ...rest }, ref) => {
+  return /* @__PURE__ */ React.createElement(MessageIf, { hasBranches: hideWhenSingleBranch ? true : void 0 }, /* @__PURE__ */ React.createElement(import_react_primitive8.Primitive.div, { ...rest, ref }));
 });
 
 // src/primitives/actionBar/index.ts
@@ -789,28 +773,47 @@ __export(actionBar_exports, {
 });
 
 // src/primitives/actionBar/ActionBarRoot.tsx
-var import_react_primitive8 = require("@radix-ui/react-primitive");
-var import_react14 = require("react");
-var ActionBarRoot = (0, import_react14.forwardRef)(({ hideWhenBusy, hideWhenNotLastOrHover, ...rest }, ref) => {
-  return /* @__PURE__ */ React.createElement(ThreadIf, { busy: hideWhenBusy ? false : void 0 }, /* @__PURE__ */ React.createElement(MessageIf, { lastOrHover: hideWhenNotLastOrHover ? true : void 0 }, /* @__PURE__ */ React.createElement(import_react_primitive8.Primitive.div, { ...rest, ref })));
+var import_react_primitive9 = require("@radix-ui/react-primitive");
+var import_react16 = require("react");
+var ActionBarRoot = (0, import_react16.forwardRef)(({ 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(
+    import_react_primitive9.Primitive.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);
@@ -822,19 +825,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
@@ -842,115 +841,137 @@ 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
-var import_react_primitive9 = require("@radix-ui/react-primitive");
-var import_react15 = require("react");
-var EditBarRoot = (0, import_react15.forwardRef)(
-  ({ ...rest }, ref) => {
-    return /* @__PURE__ */ React.createElement(import_react_primitive9.Primitive.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
+var import_react17 = require("react");
+var import_zustand2 = require("zustand");
+var useAIAssistantContext = () => {
+  const [context] = (0, import_react17.useState)(() => {
+    const useThread = (0, import_zustand2.create)()(() => ({
+      messages: [],
+      setMessages: () => {
+      },
+      isLoading: false,
+      reload: async () => {
+      },
+      append: async () => {
+      },
+      stop: () => {
+      }
+    }));
+    const useComposer = (0, import_zustand2.create)()(() => ({
+      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 = (0, import_zustand2.create)()(() => ({
+      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();
+  (0, import_react17.useMemo)(() => {
+    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]);
+  (0, import_react17.useMemo)(() => {
+    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);
+  (0, import_react17.useMemo)(() => {
+    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);
 };
 // Annotate the CommonJS export names for ESM import in node:
 0 && (module.exports = {
   ActionBarPrimitive,
   BranchPickerPrimitive,
   ComposerPrimitive,
-  EditBarPrimitive,
   MessagePrimitive,
   ThreadPrimitive,
   VercelAIThreadProvider,
   unstable_useMessageContext,
   useBeginMessageEdit,
-  useCancelMessageEdit,
   useCopyMessage,
   useGoToNextBranch,
   useGoToPreviousBranch,
-  useReloadMessage,
-  useSaveMessageEdit
+  useReloadMessage
 });