@assistant-ui/react 0.0.5 → 0.0.6

package/dist/index.d.mts

@@ -1,7 +1,6 @@
 import * as react from 'react';
 import { FC, PropsWithChildren } from 'react';
 import { TextareaAutosizeProps } from 'react-textarea-autosize';
-import { Message } from 'ai';
 import { UseChatHelpers } from 'ai/react';
 import { UseBoundStore, StoreApi } from 'zustand';
 
@@ -84,9 +83,49 @@ declare namespace index$3 {
   export { ComposerCancel as Cancel, ComposerIf as If, ComposerInput as Input, ComposerRoot as Root, ComposerSend as Send };
 }
 
+type BranchState = {
+    branchId: number;
+    branchCount: number;
+};
+
+type ComposerState = {
+    isEditing: boolean;
+    canCancel: boolean;
+    edit: () => void;
+    send: () => void;
+    cancel: () => void;
+    value: string;
+    setValue: (value: string) => void;
+};
+type ComposerStore = {
+    useComposer: UseBoundStore<StoreApi<ComposerState>>;
+};
+
+type TextContent = {
+    type: "text";
+    text: string;
+};
+type ImageContent = {
+    type: "image";
+    image: string;
+};
+type ThreadUserMessageContent = TextContent | ImageContent;
+type ThreadAssistantMessageContent = TextContent | ImageContent;
+type ThreadUserMessage = {
+    id: string;
+    role: "user";
+    content: ThreadUserMessageContent[];
+};
+type ThreadAssistantMessage = {
+    id: string;
+    role: "assistant";
+    content: ThreadAssistantMessageContent[];
+};
+type ThreadMessage = ThreadUserMessage | ThreadAssistantMessage;
+
 type MessageProviderProps = {
     children?: React.ReactNode;
-    message: Message;
+    message: ThreadMessage;
 };
 declare const MessageProvider: FC<MessageProviderProps>;
 
@@ -177,32 +216,14 @@ declare namespace index {
   export { ActionBarCopy as Copy, ActionBarEdit as Edit, ActionBarReload as Reload, ActionBarRoot as Root };
 }
 
-type VercelAIAssistantProviderProps = {
+type VercelAIChatAssistantProviderProps = {
     chat: UseChatHelpers;
     children: React.ReactNode;
 };
-declare const VercelAIAssistantProvider: FC<VercelAIAssistantProviderProps>;
-
-type ComposerState = {
-    isEditing: boolean;
-    canCancel: boolean;
-    edit: () => void;
-    send: () => void;
-    cancel: () => void;
-    value: string;
-    setValue: (value: string) => void;
-};
-type ComposerStore = {
-    useComposer: UseBoundStore<StoreApi<ComposerState>>;
-};
-
-type BranchState = {
-    branchId: number;
-    branchCount: number;
-};
+declare const VercelAIChatAssistantProvider: FC<VercelAIChatAssistantProviderProps>;
 
 type MessageState = {
-    message: Message;
+    message: ThreadMessage;
     branchState: BranchState;
     isLast: boolean;
     isCopied: boolean;
@@ -227,4 +248,4 @@ declare const useGoToNextBranch: () => (() => void) | null;
 
 declare const useGoToPreviousBranch: () => (() => void) | null;
 
-export { index as ActionBarPrimitive, index$1 as BranchPickerPrimitive, index$3 as ComposerPrimitive, index$2 as MessagePrimitive, index$4 as ThreadPrimitive, VercelAIAssistantProvider as VercelAIThreadProvider, useMessageContext as unstable_useMessageContext, useBeginMessageEdit, useCopyMessage, useGoToNextBranch, useGoToPreviousBranch, useReloadMessage };
+export { index as ActionBarPrimitive, index$1 as BranchPickerPrimitive, index$3 as ComposerPrimitive, index$2 as MessagePrimitive, index$4 as ThreadPrimitive, VercelAIChatAssistantProvider as VercelAIAssistantProvider, useMessageContext as unstable_useMessageContext, useBeginMessageEdit, useCopyMessage, useGoToNextBranch, useGoToPreviousBranch, useReloadMessage };

package/dist/index.d.ts

@@ -1,7 +1,6 @@
 import * as react from 'react';
 import { FC, PropsWithChildren } from 'react';
 import { TextareaAutosizeProps } from 'react-textarea-autosize';
-import { Message } from 'ai';
 import { UseChatHelpers } from 'ai/react';
 import { UseBoundStore, StoreApi } from 'zustand';
 
@@ -84,9 +83,49 @@ declare namespace index$3 {
   export { ComposerCancel as Cancel, ComposerIf as If, ComposerInput as Input, ComposerRoot as Root, ComposerSend as Send };
 }
 
+type BranchState = {
+    branchId: number;
+    branchCount: number;
+};
+
+type ComposerState = {
+    isEditing: boolean;
+    canCancel: boolean;
+    edit: () => void;
+    send: () => void;
+    cancel: () => void;
+    value: string;
+    setValue: (value: string) => void;
+};
+type ComposerStore = {
+    useComposer: UseBoundStore<StoreApi<ComposerState>>;
+};
+
+type TextContent = {
+    type: "text";
+    text: string;
+};
+type ImageContent = {
+    type: "image";
+    image: string;
+};
+type ThreadUserMessageContent = TextContent | ImageContent;
+type ThreadAssistantMessageContent = TextContent | ImageContent;
+type ThreadUserMessage = {
+    id: string;
+    role: "user";
+    content: ThreadUserMessageContent[];
+};
+type ThreadAssistantMessage = {
+    id: string;
+    role: "assistant";
+    content: ThreadAssistantMessageContent[];
+};
+type ThreadMessage = ThreadUserMessage | ThreadAssistantMessage;
+
 type MessageProviderProps = {
     children?: React.ReactNode;
-    message: Message;
+    message: ThreadMessage;
 };
 declare const MessageProvider: FC<MessageProviderProps>;
 
@@ -177,32 +216,14 @@ declare namespace index {
   export { ActionBarCopy as Copy, ActionBarEdit as Edit, ActionBarReload as Reload, ActionBarRoot as Root };
 }
 
-type VercelAIAssistantProviderProps = {
+type VercelAIChatAssistantProviderProps = {
     chat: UseChatHelpers;
     children: React.ReactNode;
 };
-declare const VercelAIAssistantProvider: FC<VercelAIAssistantProviderProps>;
-
-type ComposerState = {
-    isEditing: boolean;
-    canCancel: boolean;
-    edit: () => void;
-    send: () => void;
-    cancel: () => void;
-    value: string;
-    setValue: (value: string) => void;
-};
-type ComposerStore = {
-    useComposer: UseBoundStore<StoreApi<ComposerState>>;
-};
-
-type BranchState = {
-    branchId: number;
-    branchCount: number;
-};
+declare const VercelAIChatAssistantProvider: FC<VercelAIChatAssistantProviderProps>;
 
 type MessageState = {
-    message: Message;
+    message: ThreadMessage;
     branchState: BranchState;
     isLast: boolean;
     isCopied: boolean;
@@ -227,4 +248,4 @@ declare const useGoToNextBranch: () => (() => void) | null;
 
 declare const useGoToPreviousBranch: () => (() => void) | null;
 
-export { index as ActionBarPrimitive, index$1 as BranchPickerPrimitive, index$3 as ComposerPrimitive, index$2 as MessagePrimitive, index$4 as ThreadPrimitive, VercelAIAssistantProvider as VercelAIThreadProvider, useMessageContext as unstable_useMessageContext, useBeginMessageEdit, useCopyMessage, useGoToNextBranch, useGoToPreviousBranch, useReloadMessage };
+export { index as ActionBarPrimitive, index$1 as BranchPickerPrimitive, index$3 as ComposerPrimitive, index$2 as MessagePrimitive, index$4 as ThreadPrimitive, VercelAIChatAssistantProvider as VercelAIAssistantProvider, useMessageContext as unstable_useMessageContext, useBeginMessageEdit, useCopyMessage, useGoToNextBranch, useGoToPreviousBranch, useReloadMessage };

package/dist/index.js

@@ -35,7 +35,7 @@ __export(src_exports, {
   ComposerPrimitive: () => composer_exports,
   MessagePrimitive: () => message_exports,
   ThreadPrimitive: () => thread_exports,
-  VercelAIThreadProvider: () => VercelAIAssistantProvider,
+  VercelAIAssistantProvider: () => VercelAIChatAssistantProvider,
   unstable_useMessageContext: () => useMessageContext,
   useBeginMessageEdit: () => useBeginMessageEdit,
   useCopyMessage: () => useCopyMessage,
@@ -177,7 +177,7 @@ var ThreadViewport = (0, import_react4.forwardRef)(({ onScroll, children, ...res
   );
 });
 
-// src/utils/hooks/useBranches.tsx
+// src/vercel/useVercelAIBranches.tsx
 var import_react5 = require("react");
 var ROOT_ID = "__ROOT_ID__";
 var UPCOMING_MESSAGE_ID = "__UPCOMING_MESSAGE_ID__";
@@ -245,7 +245,7 @@ var sliceMessagesUntil = (messages, message) => {
     throw new Error("Unexpected: Message not found");
   return messages.slice(0, messageIdx);
 };
-var useChatWithBranches = (chat) => {
+var useVercelAIBranches = (chat) => {
   const data = (0, import_react5.useRef)({
     parentMap: /* @__PURE__ */ new Map(),
     branchMap: /* @__PURE__ */ new Map(),
@@ -282,7 +282,12 @@ var useChatWithBranches = (chat) => {
     async (message, newMessage) => {
       const newMessages = sliceMessagesUntil(chat.messages, message);
       chat.setMessages(newMessages);
-      await chat.append(newMessage);
+      if (newMessage.content[0]?.type !== "text")
+        throw new Error("Only text content is currently supported");
+      await chat.append({
+        role: "user",
+        content: newMessage.content[0].text
+      });
     },
     [chat.messages, chat.setMessages, chat.append]
   );
@@ -373,18 +378,25 @@ var useMessageContext2 = () => {
     const useComposer = (0, import_zustand.create)((set, get) => ({
       isEditing: false,
       canCancel: true,
-      edit: () => set({
-        isEditing: true,
-        value: useMessage.getState().message.content
-      }),
+      edit: () => {
+        const message = useMessage.getState().message;
+        if (message.role !== "user")
+          throw new Error("Editing is only supported for user messages");
+        if (message.content[0]?.type !== "text")
+          throw new Error("Editing is only supported for text-only messages");
+        return set({
+          isEditing: true,
+          value: message.content[0].text
+        });
+      },
       cancel: () => set({ isEditing: false }),
       send: () => {
         const message = useMessage.getState().message;
+        if (message.role !== "user")
+          throw new Error("Editing is only supported for user messages");
         useBranchObserver.getState().editAt(message, {
-          ...message,
-          id: void 0,
-          // remove id to create a new message
-          content: get().value
+          role: "user",
+          content: [{ type: "text", text: get().value }]
         });
         set({ isEditing: false });
       },
@@ -486,7 +498,9 @@ var MessageIf = ({ children, ...query }) => {
 var MessageContent = () => {
   const { useMessage } = useMessageContext();
   const content = useMessage((s) => s.message.content);
-  return /* @__PURE__ */ React.createElement(React.Fragment, null, content);
+  if (content[0]?.type !== "text")
+    throw new Error("Unsupported message content type");
+  return /* @__PURE__ */ React.createElement(React.Fragment, null, content[0].text);
 };
 
 // src/primitives/thread/ThreadMessages.tsx
@@ -515,7 +529,7 @@ var ThreadMessages = ({ components }) => {
       message: {
         id: UPCOMING_MESSAGE_ID,
         role: "assistant",
-        content: "..."
+        content: [{ type: "text", text: "..." }]
       }
     },
     /* @__PURE__ */ React.createElement(AssistantMessage, null)
@@ -810,11 +824,10 @@ var useCopyMessage = ({ copiedDuration = 3e3 }) => {
   if (isEditing)
     return null;
   return () => {
-    const {
-      message: { content },
-      setIsCopied
-    } = useMessage.getState();
-    navigator.clipboard.writeText(content);
+    const { message, setIsCopied } = useMessage.getState();
+    if (message.content[0]?.type !== "text")
+      throw new Error("Copying is only supported for text-only messages");
+    navigator.clipboard.writeText(message.content[0].text);
     setIsCopied(true);
     setTimeout(() => setIsCopied(false), copiedDuration);
   };
@@ -832,7 +845,10 @@ var useReloadMessage = () => {
   if (isLoading || !isAssistant)
     return null;
   return () => {
-    useBranchObserver.getState().reloadAt(useMessage.getState().message);
+    const message = useMessage.getState().message;
+    if (message.role !== "assistant")
+      throw new Error("Reloading is only supported on assistant messages");
+    useBranchObserver.getState().reloadAt(message);
   };
 };
 
@@ -862,8 +878,6 @@ var useAIAssistantContext = () => {
   const [context] = (0, import_react17.useState)(() => {
     const useThread = (0, import_zustand2.create)()(() => ({
       messages: [],
-      setMessages: () => {
-      },
       isLoading: false,
       reload: async () => {
       },
@@ -883,9 +897,8 @@ var useAIAssistantContext = () => {
       },
       send: () => {
         useThread.getState().append({
-          content: useComposer.getState().value,
           role: "user",
-          createdAt: /* @__PURE__ */ new Date()
+          content: [{ type: "text", text: useComposer.getState().value }]
         });
         useComposer.getState().setValue("");
       },
@@ -909,36 +922,65 @@ var useAIAssistantContext = () => {
   });
   return context;
 };
-var VercelAIAssistantProvider = ({
-  chat,
-  children
-}) => {
+var ThreadMessageCache = /* @__PURE__ */ new WeakMap();
+var vercelToThreadMessage = (message) => {
+  if (message.role !== "user" && message.role !== "assistant")
+    throw new Error("Unsupported role");
+  return {
+    id: message.id,
+    role: message.role,
+    content: [{ type: "text", text: message.content }]
+  };
+};
+var vercelToCachedThreadMessages = (messages) => {
+  return messages.map((m) => {
+    const cached = ThreadMessageCache.get(m);
+    if (cached)
+      return cached;
+    const newMessage = vercelToThreadMessage(m);
+    ThreadMessageCache.set(m, newMessage);
+    return newMessage;
+  });
+};
+var VercelAIChatAssistantProvider = ({ chat, children }) => {
   const context = useAIAssistantContext();
+  const messages = (0, import_react17.useMemo)(() => {
+    return vercelToCachedThreadMessages(chat.messages);
+  }, [chat.messages]);
+  const reload = (0, import_react17.useCallback)(async () => {
+    await chat.reload();
+  }, [chat.reload]);
+  const append = (0, import_react17.useCallback)(
+    async (message) => {
+      if (message.content[0]?.type !== "text") {
+        throw new Error("Only text content is currently supported");
+      }
+      await chat.append({
+        role: message.role,
+        content: message.content[0].text
+      });
+    },
+    [chat.append]
+  );
+  const stop = (0, import_react17.useCallback)(() => {
+    const lastMessage = chat.messages.at(-1);
+    chat.stop();
+    if (lastMessage?.role === "user") {
+      chat.setInput(lastMessage.content);
+    }
+  }, [chat.messages, chat.stop, chat.setInput]);
   (0, import_react17.useMemo)(() => {
     context.useThread.setState(
       {
-        messages: chat.messages,
-        setMessages: (value) => {
-          chat.setMessages(value);
-        },
+        messages,
         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);
-          }
-        }
+        reload,
+        append,
+        stop
       },
       true
     );
-  }, [context, chat]);
+  }, [context, messages, reload, append, stop, chat.isLoading]);
   (0, import_react17.useMemo)(() => {
     context.useComposer.setState({
       canCancel: chat.isLoading,
@@ -946,17 +988,9 @@ var VercelAIAssistantProvider = ({
       setValue: chat.setInput
     });
   }, [context, chat.isLoading, chat.input, chat.setInput]);
-  const branches = useChatWithBranches(chat);
+  const branches = useVercelAIBranches(chat);
   (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.useBranchObserver.setState(branches, true);
   }, [context, branches]);
   return /* @__PURE__ */ React.createElement(AssistantContext.Provider, { value: context }, children);
 };
@@ -967,7 +1001,7 @@ var VercelAIAssistantProvider = ({
   ComposerPrimitive,
   MessagePrimitive,
   ThreadPrimitive,
-  VercelAIThreadProvider,
+  VercelAIAssistantProvider,
   unstable_useMessageContext,
   useBeginMessageEdit,
   useCopyMessage,

package/dist/index.mjs

@@ -140,7 +140,7 @@ var ThreadViewport = forwardRef2(({ onScroll, children, ...rest }, forwardedRef)
   );
 });
 
-// src/utils/hooks/useBranches.tsx
+// src/vercel/useVercelAIBranches.tsx
 import { useCallback, useMemo, useRef as useRef3 } from "react";
 var ROOT_ID = "__ROOT_ID__";
 var UPCOMING_MESSAGE_ID = "__UPCOMING_MESSAGE_ID__";
@@ -208,7 +208,7 @@ var sliceMessagesUntil = (messages, message) => {
     throw new Error("Unexpected: Message not found");
   return messages.slice(0, messageIdx);
 };
-var useChatWithBranches = (chat) => {
+var useVercelAIBranches = (chat) => {
   const data = useRef3({
     parentMap: /* @__PURE__ */ new Map(),
     branchMap: /* @__PURE__ */ new Map(),
@@ -245,7 +245,12 @@ var useChatWithBranches = (chat) => {
     async (message, newMessage) => {
       const newMessages = sliceMessagesUntil(chat.messages, message);
       chat.setMessages(newMessages);
-      await chat.append(newMessage);
+      if (newMessage.content[0]?.type !== "text")
+        throw new Error("Only text content is currently supported");
+      await chat.append({
+        role: "user",
+        content: newMessage.content[0].text
+      });
     },
     [chat.messages, chat.setMessages, chat.append]
   );
@@ -336,18 +341,25 @@ var useMessageContext2 = () => {
     const useComposer = create((set, get) => ({
       isEditing: false,
       canCancel: true,
-      edit: () => set({
-        isEditing: true,
-        value: useMessage.getState().message.content
-      }),
+      edit: () => {
+        const message = useMessage.getState().message;
+        if (message.role !== "user")
+          throw new Error("Editing is only supported for user messages");
+        if (message.content[0]?.type !== "text")
+          throw new Error("Editing is only supported for text-only messages");
+        return set({
+          isEditing: true,
+          value: message.content[0].text
+        });
+      },
       cancel: () => set({ isEditing: false }),
       send: () => {
         const message = useMessage.getState().message;
+        if (message.role !== "user")
+          throw new Error("Editing is only supported for user messages");
         useBranchObserver.getState().editAt(message, {
-          ...message,
-          id: void 0,
-          // remove id to create a new message
-          content: get().value
+          role: "user",
+          content: [{ type: "text", text: get().value }]
         });
         set({ isEditing: false });
       },
@@ -451,7 +463,9 @@ var MessageIf = ({ children, ...query }) => {
 var MessageContent = () => {
   const { useMessage } = useMessageContext();
   const content = useMessage((s) => s.message.content);
-  return /* @__PURE__ */ React.createElement(React.Fragment, null, content);
+  if (content[0]?.type !== "text")
+    throw new Error("Unsupported message content type");
+  return /* @__PURE__ */ React.createElement(React.Fragment, null, content[0].text);
 };
 
 // src/primitives/thread/ThreadMessages.tsx
@@ -480,7 +494,7 @@ var ThreadMessages = ({ components }) => {
       message: {
         id: UPCOMING_MESSAGE_ID,
         role: "assistant",
-        content: "..."
+        content: [{ type: "text", text: "..." }]
       }
     },
     /* @__PURE__ */ React.createElement(AssistantMessage, null)
@@ -793,11 +807,10 @@ var useCopyMessage = ({ copiedDuration = 3e3 }) => {
   if (isEditing)
     return null;
   return () => {
-    const {
-      message: { content },
-      setIsCopied
-    } = useMessage.getState();
-    navigator.clipboard.writeText(content);
+    const { message, setIsCopied } = useMessage.getState();
+    if (message.content[0]?.type !== "text")
+      throw new Error("Copying is only supported for text-only messages");
+    navigator.clipboard.writeText(message.content[0].text);
     setIsCopied(true);
     setTimeout(() => setIsCopied(false), copiedDuration);
   };
@@ -815,7 +828,10 @@ var useReloadMessage = () => {
   if (isLoading || !isAssistant)
     return null;
   return () => {
-    useBranchObserver.getState().reloadAt(useMessage.getState().message);
+    const message = useMessage.getState().message;
+    if (message.role !== "assistant")
+      throw new Error("Reloading is only supported on assistant messages");
+    useBranchObserver.getState().reloadAt(message);
   };
 };
 
@@ -839,14 +855,12 @@ var useBeginMessageEdit = () => {
 var ActionBarEdit = createActionButton(useBeginMessageEdit);
 
 // src/vercel/VercelAIAssistantProvider.tsx
-import { useMemo as useMemo4, useState as useState3 } from "react";
+import { useCallback as useCallback2, 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 () => {
       },
@@ -866,9 +880,8 @@ var useAIAssistantContext = () => {
       },
       send: () => {
         useThread.getState().append({
-          content: useComposer.getState().value,
           role: "user",
-          createdAt: /* @__PURE__ */ new Date()
+          content: [{ type: "text", text: useComposer.getState().value }]
         });
         useComposer.getState().setValue("");
       },
@@ -892,36 +905,65 @@ var useAIAssistantContext = () => {
   });
   return context;
 };
-var VercelAIAssistantProvider = ({
-  chat,
-  children
-}) => {
+var ThreadMessageCache = /* @__PURE__ */ new WeakMap();
+var vercelToThreadMessage = (message) => {
+  if (message.role !== "user" && message.role !== "assistant")
+    throw new Error("Unsupported role");
+  return {
+    id: message.id,
+    role: message.role,
+    content: [{ type: "text", text: message.content }]
+  };
+};
+var vercelToCachedThreadMessages = (messages) => {
+  return messages.map((m) => {
+    const cached = ThreadMessageCache.get(m);
+    if (cached)
+      return cached;
+    const newMessage = vercelToThreadMessage(m);
+    ThreadMessageCache.set(m, newMessage);
+    return newMessage;
+  });
+};
+var VercelAIChatAssistantProvider = ({ chat, children }) => {
   const context = useAIAssistantContext();
+  const messages = useMemo4(() => {
+    return vercelToCachedThreadMessages(chat.messages);
+  }, [chat.messages]);
+  const reload = useCallback2(async () => {
+    await chat.reload();
+  }, [chat.reload]);
+  const append = useCallback2(
+    async (message) => {
+      if (message.content[0]?.type !== "text") {
+        throw new Error("Only text content is currently supported");
+      }
+      await chat.append({
+        role: message.role,
+        content: message.content[0].text
+      });
+    },
+    [chat.append]
+  );
+  const stop = useCallback2(() => {
+    const lastMessage = chat.messages.at(-1);
+    chat.stop();
+    if (lastMessage?.role === "user") {
+      chat.setInput(lastMessage.content);
+    }
+  }, [chat.messages, chat.stop, chat.setInput]);
   useMemo4(() => {
     context.useThread.setState(
       {
-        messages: chat.messages,
-        setMessages: (value) => {
-          chat.setMessages(value);
-        },
+        messages,
         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);
-          }
-        }
+        reload,
+        append,
+        stop
       },
       true
     );
-  }, [context, chat]);
+  }, [context, messages, reload, append, stop, chat.isLoading]);
   useMemo4(() => {
     context.useComposer.setState({
       canCancel: chat.isLoading,
@@ -929,17 +971,9 @@ var VercelAIAssistantProvider = ({
       setValue: chat.setInput
     });
   }, [context, chat.isLoading, chat.input, chat.setInput]);
-  const branches = useChatWithBranches(chat);
+  const branches = useVercelAIBranches(chat);
   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.useBranchObserver.setState(branches, true);
   }, [context, branches]);
   return /* @__PURE__ */ React.createElement(AssistantContext.Provider, { value: context }, children);
 };
@@ -949,7 +983,7 @@ export {
   composer_exports as ComposerPrimitive,
   message_exports as MessagePrimitive,
   thread_exports as ThreadPrimitive,
-  VercelAIAssistantProvider as VercelAIThreadProvider,
+  VercelAIChatAssistantProvider as VercelAIAssistantProvider,
   useMessageContext as unstable_useMessageContext,
   useBeginMessageEdit,
   useCopyMessage,

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@assistant-ui/react",
-  "version": "0.0.5",
+  "version": "0.0.6",
   "license": "MIT",
   "exports": {
     ".": {