@assistant-ui/react 0.0.5 → 0.0.7

package/dist/index.d.mts

@@ -1,8 +1,7 @@
 import * as react from 'react';
-import { FC, PropsWithChildren } from 'react';
+import { FC, PropsWithChildren, ReactNode, ComponentType } from 'react';
 import { TextareaAutosizeProps } from 'react-textarea-autosize';
-import { Message } from 'ai';
-import { UseChatHelpers } from 'ai/react';
+import { UseChatHelpers, UseAssistantHelpers } from 'ai/react';
 import { UseBoundStore, StoreApi } from 'zustand';
 
 declare const ThreadRoot: react.ForwardRefExoticComponent<Pick<Omit<react.DetailedHTMLProps<react.HTMLAttributes<HTMLDivElement>, HTMLDivElement>, "ref"> & {
@@ -84,9 +83,60 @@ 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 ThreadMessageTextPart = {
+    type: "text";
+    text: string;
+};
+type ThreadMessageImagePart = {
+    type: "image";
+    image: string;
+};
+type ThreadMessageUIPart = {
+    type: "ui";
+    display: ReactNode;
+};
+type ThreadMessageToolCallPart = {
+    type: "tool-call";
+    name: string;
+    args: object;
+    result?: object;
+};
+type ThreadUserMessageContent = ThreadMessageTextPart | ThreadMessageImagePart | ThreadMessageUIPart;
+type ThreadAssistantMessageContent = ThreadMessageTextPart | ThreadMessageImagePart | ThreadMessageUIPart | ThreadMessageToolCallPart;
+type ThreadUserMessage = {
+    id: string;
+    role: "user";
+    content: ThreadUserMessageContent[];
+};
+type ThreadAssistantMessage = {
+    id: string;
+    role: "assistant";
+    content: ThreadAssistantMessageContent[];
+};
+type ThreadMessage = ThreadUserMessage | ThreadAssistantMessage;
+type CreateThreadMessage = Omit<ThreadUserMessage, "id">;
+
 type MessageProviderProps = {
     children?: React.ReactNode;
-    message: Message;
+    message: ThreadMessage;
 };
 declare const MessageProvider: FC<MessageProviderProps>;
 
@@ -108,7 +158,28 @@ type MessageIfProps = RequireAtLeastOne<MessageIfFilters> & {
 };
 declare const MessageIf: FC<MessageIfProps>;
 
-declare const MessageContent: FC;
+type MessageContentProps = {
+    components?: {
+        Text?: ComponentType<{
+            part: ThreadMessageTextPart;
+        }>;
+        Image?: ComponentType<{
+            part: ThreadMessageImagePart;
+        }>;
+        UI?: ComponentType<{
+            part: ThreadMessageUIPart;
+        }>;
+        tools?: {
+            by_name?: Record<string, ComponentType<{
+                part: ThreadMessageToolCallPart;
+            }>>;
+            Fallback?: ComponentType<{
+                part: ThreadMessageToolCallPart;
+            }>;
+        };
+    };
+};
+declare const MessageContent: FC<MessageContentProps>;
 
 declare namespace index$2 {
   export { MessageContent as Content, MessageIf as If, MessageProvider as Provider, MessageRoot as Root };
@@ -177,32 +248,26 @@ declare namespace index {
   export { ActionBarCopy as Copy, ActionBarEdit as Edit, ActionBarReload as Reload, ActionBarRoot as Root };
 }
 
-type VercelAIAssistantProviderProps = {
+type VercelAIAssistantProviderProps$1 = PropsWithChildren<{
     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;
+} | {
+    assistant: UseAssistantHelpers;
+}>;
+declare const VercelAIAssistantProvider: FC<VercelAIAssistantProviderProps$1>;
+
+type RSCMessage = {
+    id: string;
+    role: "user" | "assistant";
+    display: ReactNode;
 };
+type VercelAIAssistantProviderProps = PropsWithChildren<{
+    messages: RSCMessage[];
+    append: (message: CreateThreadMessage) => Promise<void>;
+}>;
+declare const VercelRSCAssistantProvider: FC<VercelAIAssistantProviderProps>;
 
 type MessageState = {
-    message: Message;
+    message: ThreadMessage;
     branchState: BranchState;
     isLast: boolean;
     isCopied: boolean;
@@ -227,4 +292,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, VercelAIAssistantProvider, VercelRSCAssistantProvider as unstable_VercelRSCAssistantProvider, useMessageContext as unstable_useMessageContext, useBeginMessageEdit, useCopyMessage, useGoToNextBranch, useGoToPreviousBranch, useReloadMessage };

package/dist/index.d.ts

@@ -1,8 +1,7 @@
 import * as react from 'react';
-import { FC, PropsWithChildren } from 'react';
+import { FC, PropsWithChildren, ReactNode, ComponentType } from 'react';
 import { TextareaAutosizeProps } from 'react-textarea-autosize';
-import { Message } from 'ai';
-import { UseChatHelpers } from 'ai/react';
+import { UseChatHelpers, UseAssistantHelpers } from 'ai/react';
 import { UseBoundStore, StoreApi } from 'zustand';
 
 declare const ThreadRoot: react.ForwardRefExoticComponent<Pick<Omit<react.DetailedHTMLProps<react.HTMLAttributes<HTMLDivElement>, HTMLDivElement>, "ref"> & {
@@ -84,9 +83,60 @@ 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 ThreadMessageTextPart = {
+    type: "text";
+    text: string;
+};
+type ThreadMessageImagePart = {
+    type: "image";
+    image: string;
+};
+type ThreadMessageUIPart = {
+    type: "ui";
+    display: ReactNode;
+};
+type ThreadMessageToolCallPart = {
+    type: "tool-call";
+    name: string;
+    args: object;
+    result?: object;
+};
+type ThreadUserMessageContent = ThreadMessageTextPart | ThreadMessageImagePart | ThreadMessageUIPart;
+type ThreadAssistantMessageContent = ThreadMessageTextPart | ThreadMessageImagePart | ThreadMessageUIPart | ThreadMessageToolCallPart;
+type ThreadUserMessage = {
+    id: string;
+    role: "user";
+    content: ThreadUserMessageContent[];
+};
+type ThreadAssistantMessage = {
+    id: string;
+    role: "assistant";
+    content: ThreadAssistantMessageContent[];
+};
+type ThreadMessage = ThreadUserMessage | ThreadAssistantMessage;
+type CreateThreadMessage = Omit<ThreadUserMessage, "id">;
+
 type MessageProviderProps = {
     children?: React.ReactNode;
-    message: Message;
+    message: ThreadMessage;
 };
 declare const MessageProvider: FC<MessageProviderProps>;
 
@@ -108,7 +158,28 @@ type MessageIfProps = RequireAtLeastOne<MessageIfFilters> & {
 };
 declare const MessageIf: FC<MessageIfProps>;
 
-declare const MessageContent: FC;
+type MessageContentProps = {
+    components?: {
+        Text?: ComponentType<{
+            part: ThreadMessageTextPart;
+        }>;
+        Image?: ComponentType<{
+            part: ThreadMessageImagePart;
+        }>;
+        UI?: ComponentType<{
+            part: ThreadMessageUIPart;
+        }>;
+        tools?: {
+            by_name?: Record<string, ComponentType<{
+                part: ThreadMessageToolCallPart;
+            }>>;
+            Fallback?: ComponentType<{
+                part: ThreadMessageToolCallPart;
+            }>;
+        };
+    };
+};
+declare const MessageContent: FC<MessageContentProps>;
 
 declare namespace index$2 {
   export { MessageContent as Content, MessageIf as If, MessageProvider as Provider, MessageRoot as Root };
@@ -177,32 +248,26 @@ declare namespace index {
   export { ActionBarCopy as Copy, ActionBarEdit as Edit, ActionBarReload as Reload, ActionBarRoot as Root };
 }
 
-type VercelAIAssistantProviderProps = {
+type VercelAIAssistantProviderProps$1 = PropsWithChildren<{
     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;
+} | {
+    assistant: UseAssistantHelpers;
+}>;
+declare const VercelAIAssistantProvider: FC<VercelAIAssistantProviderProps$1>;
+
+type RSCMessage = {
+    id: string;
+    role: "user" | "assistant";
+    display: ReactNode;
 };
+type VercelAIAssistantProviderProps = PropsWithChildren<{
+    messages: RSCMessage[];
+    append: (message: CreateThreadMessage) => Promise<void>;
+}>;
+declare const VercelRSCAssistantProvider: FC<VercelAIAssistantProviderProps>;
 
 type MessageState = {
-    message: Message;
+    message: ThreadMessage;
     branchState: BranchState;
     isLast: boolean;
     isCopied: boolean;
@@ -227,4 +292,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, VercelAIAssistantProvider, VercelRSCAssistantProvider as unstable_VercelRSCAssistantProvider, useMessageContext as unstable_useMessageContext, useBeginMessageEdit, useCopyMessage, useGoToNextBranch, useGoToPreviousBranch, useReloadMessage };

package/dist/index.js

@@ -35,7 +35,8 @@ __export(src_exports, {
   ComposerPrimitive: () => composer_exports,
   MessagePrimitive: () => message_exports,
   ThreadPrimitive: () => thread_exports,
-  VercelAIThreadProvider: () => VercelAIAssistantProvider,
+  VercelAIAssistantProvider: () => VercelAIAssistantProvider,
+  unstable_VercelRSCAssistantProvider: () => VercelRSCAssistantProvider,
   unstable_useMessageContext: () => useMessageContext,
   useBeginMessageEdit: () => useBeginMessageEdit,
   useCopyMessage: () => useCopyMessage,
@@ -102,10 +103,10 @@ var ThreadEmpty = ({ children }) => {
 };
 
 // src/primitives/thread/ThreadViewport.tsx
-var import_react4 = require("react");
-var import_react_primitive2 = require("@radix-ui/react-primitive");
-var import_react_compose_refs = require("@radix-ui/react-compose-refs");
 var import_primitive = require("@radix-ui/primitive");
+var import_react_compose_refs = require("@radix-ui/react-compose-refs");
+var import_react_primitive2 = require("@radix-ui/react-primitive");
+var import_react4 = require("react");
 
 // src/utils/hooks/useOnResizeContent.tsx
 var import_react3 = require("react");
@@ -138,10 +139,8 @@ var useOnResizeContent = (ref, callback) => {
     mutationObserver.observe(el, { childList: true });
     for (const child of el.children) {
       resizeObserver.observe(child);
-      console.log("observing child", child);
     }
     return () => {
-      console.log("disconnecting");
       resizeObserver.disconnect();
       mutationObserver.disconnect();
     };
@@ -165,7 +164,6 @@ var ThreadViewport = (0, import_react4.forwardRef)(({ onScroll, children, ...res
       return;
     setIsAtBottom(div.scrollHeight - div.scrollTop <= div.clientHeight + 50);
   };
-  console.log(isAtBottom);
   return /* @__PURE__ */ React.createElement(
     import_react_primitive2.Primitive.div,
     {
@@ -177,7 +175,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 +243,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(),
@@ -270,19 +268,27 @@ var useChatWithBranches = (chat) => {
     },
     [data, chat.messages, chat.setMessages]
   );
+  const reloadMaybe = "reload" in chat ? chat.reload : void 0;
   const reloadAt = (0, import_react5.useCallback)(
     async (message) => {
+      if (!reloadMaybe)
+        throw new Error("Reload not supported by Vercel AI SDK's useAssistant");
       const newMessages = sliceMessagesUntil(chat.messages, message);
       chat.setMessages(newMessages);
-      await chat.reload();
+      await reloadMaybe();
     },
-    [chat.messages, chat.setMessages, chat.reload]
+    [chat.messages, chat.setMessages, reloadMaybe]
   );
   const editAt = (0, import_react5.useCallback)(
     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 +379,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 });
       },
@@ -483,10 +496,41 @@ var MessageIf = ({ children, ...query }) => {
 };
 
 // src/primitives/message/MessageContent.tsx
-var MessageContent = () => {
+var defaultComponents = {
+  Text: ({ part }) => /* @__PURE__ */ React.createElement(React.Fragment, null, part.text),
+  Image: () => null,
+  UI: ({ part }) => part.display,
+  tools: {
+    Fallback: () => null
+  }
+};
+var MessageContent = ({
+  components: {
+    Text = defaultComponents.Text,
+    Image = defaultComponents.Image,
+    UI = defaultComponents.UI,
+    tools: { by_name = {}, Fallback = defaultComponents.tools.Fallback } = {}
+  } = {}
+}) => {
   const { useMessage } = useMessageContext();
   const content = useMessage((s) => s.message.content);
-  return /* @__PURE__ */ React.createElement(React.Fragment, null, content);
+  return /* @__PURE__ */ React.createElement(React.Fragment, null, content.map((part, i) => {
+    const key = i;
+    switch (part.type) {
+      case "text":
+        return /* @__PURE__ */ React.createElement(Text, { key, part });
+      case "image":
+        return /* @__PURE__ */ React.createElement(Image, { key, part });
+      case "ui":
+        return /* @__PURE__ */ React.createElement(UI, { key, part });
+      case "tool-call": {
+        const Tool = by_name[part.name] || Fallback;
+        return /* @__PURE__ */ React.createElement(Tool, { key, part });
+      }
+      default:
+        return null;
+    }
+  }));
 };
 
 // src/primitives/thread/ThreadMessages.tsx
@@ -515,7 +559,7 @@ var ThreadMessages = ({ components }) => {
       message: {
         id: UPCOMING_MESSAGE_ID,
         role: "assistant",
-        content: "..."
+        content: [{ type: "text", text: "..." }]
       }
     },
     /* @__PURE__ */ React.createElement(AssistantMessage, null)
@@ -810,11 +854,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 +875,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);
   };
 };
 
@@ -856,36 +902,40 @@ var useBeginMessageEdit = () => {
 var ActionBarEdit = createActionButton(useBeginMessageEdit);
 
 // src/vercel/VercelAIAssistantProvider.tsx
+var import_react18 = require("react");
+
+// src/vercel/useDummyAIAssistantContext.tsx
 var import_react17 = require("react");
 var import_zustand2 = require("zustand");
-var useAIAssistantContext = () => {
+var useDummyAIAssistantContext = () => {
   const [context] = (0, import_react17.useState)(() => {
     const useThread = (0, import_zustand2.create)()(() => ({
       messages: [],
-      setMessages: () => {
-      },
       isLoading: false,
       reload: async () => {
+        throw new Error("Not implemented");
       },
       append: async () => {
+        throw new Error("Not implemented");
       },
       stop: () => {
+        throw new Error("Not implemented");
       }
     }));
     const useComposer = (0, import_zustand2.create)()(() => ({
       isEditing: true,
       canCancel: false,
       value: "",
-      setValue: () => {
+      setValue: (value) => {
+        useComposer.setState({ value });
       },
       edit: () => {
         throw new Error("Not implemented");
       },
       send: () => {
         useThread.getState().append({
-          content: useComposer.getState().value,
           role: "user",
-          createdAt: /* @__PURE__ */ new Date()
+          content: [{ type: "text", text: useComposer.getState().value }]
         });
         useComposer.getState().setValue("");
       },
@@ -896,70 +946,156 @@ var useAIAssistantContext = () => {
     const useBranchObserver = (0, import_zustand2.create)()(() => ({
       getBranchState: () => ({
         branchId: 0,
-        branchCount: 0
+        branchCount: 1
       }),
       switchToBranch: () => {
+        throw new Error("Not implemented");
       },
       editAt: async () => {
+        throw new Error("Not implemented");
       },
       reloadAt: async () => {
+        throw new Error("Not implemented");
       }
     }));
     return { useThread, useComposer, useBranchObserver };
   });
   return context;
 };
+
+// src/vercel/VercelAIAssistantProvider.tsx
+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 VercelAIAssistantProvider = ({
-  chat,
-  children
+  children,
+  ...rest
 }) => {
-  const context = useAIAssistantContext();
-  (0, import_react17.useMemo)(() => {
+  const context = useDummyAIAssistantContext();
+  const vercel = "chat" in rest ? rest.chat : rest.assistant;
+  const messages = (0, import_react18.useMemo)(() => {
+    return vercelToCachedThreadMessages(vercel.messages);
+  }, [vercel.messages]);
+  const maybeReload = "reload" in vercel ? vercel.reload : null;
+  const reload = (0, import_react18.useCallback)(async () => {
+    if (!maybeReload)
+      throw new Error("Reload not supported");
+    await maybeReload();
+  }, [maybeReload]);
+  const append = (0, import_react18.useCallback)(
+    async (message) => {
+      if (message.content[0]?.type !== "text") {
+        throw new Error("Only text content is currently supported");
+      }
+      await vercel.append({
+        role: message.role,
+        content: message.content[0].text
+      });
+    },
+    [vercel.append]
+  );
+  const stop = (0, import_react18.useCallback)(() => {
+    const lastMessage = vercel.messages.at(-1);
+    vercel.stop();
+    if (lastMessage?.role === "user") {
+      vercel.setInput(lastMessage.content);
+    }
+  }, [vercel.messages, vercel.stop, vercel.setInput]);
+  const isLoading = "isLoading" in vercel ? vercel.isLoading : vercel.status === "in_progress";
+  (0, import_react18.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);
-          }
-        }
+        messages,
+        isLoading,
+        reload,
+        append,
+        stop
       },
       true
     );
-  }, [context, chat]);
-  (0, import_react17.useMemo)(() => {
+  }, [context, messages, reload, append, stop, isLoading]);
+  (0, import_react18.useMemo)(() => {
     context.useComposer.setState({
-      canCancel: chat.isLoading,
-      value: chat.input,
-      setValue: chat.setInput
+      canCancel: isLoading,
+      value: vercel.input,
+      setValue: vercel.setInput
     });
-  }, [context, chat.isLoading, chat.input, chat.setInput]);
-  const branches = useChatWithBranches(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, isLoading, vercel.input, vercel.setInput]);
+  const branches = useVercelAIBranches(vercel);
+  (0, import_react18.useMemo)(() => {
+    context.useBranchObserver.setState(branches, true);
   }, [context, branches]);
   return /* @__PURE__ */ React.createElement(AssistantContext.Provider, { value: context }, children);
 };
+
+// src/vercel/VercelRSCAssistantProvider.tsx
+var import_react19 = require("react");
+var ThreadMessageCache2 = /* @__PURE__ */ new WeakMap();
+var vercelToThreadMessage2 = (message) => {
+  if (message.role !== "user" && message.role !== "assistant")
+    throw new Error("Unsupported role");
+  return {
+    id: message.id,
+    role: message.role,
+    content: [{ type: "ui", display: message.display }]
+  };
+};
+var vercelToCachedThreadMessages2 = (messages) => {
+  return messages.map((m) => {
+    const cached = ThreadMessageCache2.get(m);
+    if (cached)
+      return cached;
+    const newMessage = vercelToThreadMessage2(m);
+    ThreadMessageCache2.set(m, newMessage);
+    return newMessage;
+  });
+};
+var VercelRSCAssistantProvider = ({
+  children,
+  messages: vercelMessages,
+  append: vercelAppend
+}) => {
+  const context = useDummyAIAssistantContext();
+  const messages = (0, import_react19.useMemo)(() => {
+    return vercelToCachedThreadMessages2(vercelMessages);
+  }, [vercelMessages]);
+  const append = (0, import_react19.useCallback)(
+    async (message) => {
+      if (message.content[0]?.type !== "text") {
+        throw new Error("Only text content is currently supported");
+      }
+      await vercelAppend({
+        role: message.role,
+        content: [{ type: "text", text: message.content[0].text }]
+      });
+    },
+    [vercelAppend]
+  );
+  (0, import_react19.useMemo)(() => {
+    context.useThread.setState({
+      messages,
+      append
+    });
+  }, [context, messages, append]);
+  return /* @__PURE__ */ React.createElement(AssistantContext.Provider, { value: context }, children);
+};
 // Annotate the CommonJS export names for ESM import in node:
 0 && (module.exports = {
   ActionBarPrimitive,
@@ -967,7 +1103,8 @@ var VercelAIAssistantProvider = ({
   ComposerPrimitive,
   MessagePrimitive,
   ThreadPrimitive,
-  VercelAIThreadProvider,
+  VercelAIAssistantProvider,
+  unstable_VercelRSCAssistantProvider,
   unstable_useMessageContext,
   useBeginMessageEdit,
   useCopyMessage,

package/dist/index.mjs

@@ -63,12 +63,12 @@ var ThreadEmpty = ({ children }) => {
 };
 
 // src/primitives/thread/ThreadViewport.tsx
-import { forwardRef as forwardRef2, useRef as useRef2, useState } from "react";
+import { composeEventHandlers } from "@radix-ui/primitive";
+import { useComposedRefs } from "@radix-ui/react-compose-refs";
 import {
   Primitive as Primitive2
 } from "@radix-ui/react-primitive";
-import { useComposedRefs } from "@radix-ui/react-compose-refs";
-import { composeEventHandlers } from "@radix-ui/primitive";
+import { forwardRef as forwardRef2, useRef as useRef2, useState } from "react";
 
 // src/utils/hooks/useOnResizeContent.tsx
 import { useLayoutEffect, useRef } from "react";
@@ -101,10 +101,8 @@ var useOnResizeContent = (ref, callback) => {
     mutationObserver.observe(el, { childList: true });
     for (const child of el.children) {
       resizeObserver.observe(child);
-      console.log("observing child", child);
     }
     return () => {
-      console.log("disconnecting");
       resizeObserver.disconnect();
       mutationObserver.disconnect();
     };
@@ -128,7 +126,6 @@ var ThreadViewport = forwardRef2(({ onScroll, children, ...rest }, forwardedRef)
       return;
     setIsAtBottom(div.scrollHeight - div.scrollTop <= div.clientHeight + 50);
   };
-  console.log(isAtBottom);
   return /* @__PURE__ */ React.createElement(
     Primitive2.div,
     {
@@ -140,7 +137,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 +205,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(),
@@ -233,19 +230,27 @@ var useChatWithBranches = (chat) => {
     },
     [data, chat.messages, chat.setMessages]
   );
+  const reloadMaybe = "reload" in chat ? chat.reload : void 0;
   const reloadAt = useCallback(
     async (message) => {
+      if (!reloadMaybe)
+        throw new Error("Reload not supported by Vercel AI SDK's useAssistant");
       const newMessages = sliceMessagesUntil(chat.messages, message);
       chat.setMessages(newMessages);
-      await chat.reload();
+      await reloadMaybe();
     },
-    [chat.messages, chat.setMessages, chat.reload]
+    [chat.messages, chat.setMessages, reloadMaybe]
   );
   const editAt = useCallback(
     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 });
       },
@@ -448,10 +460,41 @@ var MessageIf = ({ children, ...query }) => {
 };
 
 // src/primitives/message/MessageContent.tsx
-var MessageContent = () => {
+var defaultComponents = {
+  Text: ({ part }) => /* @__PURE__ */ React.createElement(React.Fragment, null, part.text),
+  Image: () => null,
+  UI: ({ part }) => part.display,
+  tools: {
+    Fallback: () => null
+  }
+};
+var MessageContent = ({
+  components: {
+    Text = defaultComponents.Text,
+    Image = defaultComponents.Image,
+    UI = defaultComponents.UI,
+    tools: { by_name = {}, Fallback = defaultComponents.tools.Fallback } = {}
+  } = {}
+}) => {
   const { useMessage } = useMessageContext();
   const content = useMessage((s) => s.message.content);
-  return /* @__PURE__ */ React.createElement(React.Fragment, null, content);
+  return /* @__PURE__ */ React.createElement(React.Fragment, null, content.map((part, i) => {
+    const key = i;
+    switch (part.type) {
+      case "text":
+        return /* @__PURE__ */ React.createElement(Text, { key, part });
+      case "image":
+        return /* @__PURE__ */ React.createElement(Image, { key, part });
+      case "ui":
+        return /* @__PURE__ */ React.createElement(UI, { key, part });
+      case "tool-call": {
+        const Tool = by_name[part.name] || Fallback;
+        return /* @__PURE__ */ React.createElement(Tool, { key, part });
+      }
+      default:
+        return null;
+    }
+  }));
 };
 
 // src/primitives/thread/ThreadMessages.tsx
@@ -480,7 +523,7 @@ var ThreadMessages = ({ components }) => {
       message: {
         id: UPCOMING_MESSAGE_ID,
         role: "assistant",
-        content: "..."
+        content: [{ type: "text", text: "..." }]
       }
     },
     /* @__PURE__ */ React.createElement(AssistantMessage, null)
@@ -793,11 +836,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 +857,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,36 +884,40 @@ 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 } from "react";
+
+// src/vercel/useDummyAIAssistantContext.tsx
+import { useState as useState3 } from "react";
 import { create as create2 } from "zustand";
-var useAIAssistantContext = () => {
+var useDummyAIAssistantContext = () => {
   const [context] = useState3(() => {
     const useThread = create2()(() => ({
       messages: [],
-      setMessages: () => {
-      },
       isLoading: false,
       reload: async () => {
+        throw new Error("Not implemented");
       },
       append: async () => {
+        throw new Error("Not implemented");
       },
       stop: () => {
+        throw new Error("Not implemented");
       }
     }));
     const useComposer = create2()(() => ({
       isEditing: true,
       canCancel: false,
       value: "",
-      setValue: () => {
+      setValue: (value) => {
+        useComposer.setState({ value });
       },
       edit: () => {
         throw new Error("Not implemented");
       },
       send: () => {
         useThread.getState().append({
-          content: useComposer.getState().value,
           role: "user",
-          createdAt: /* @__PURE__ */ new Date()
+          content: [{ type: "text", text: useComposer.getState().value }]
         });
         useComposer.getState().setValue("");
       },
@@ -879,77 +928,167 @@ var useAIAssistantContext = () => {
     const useBranchObserver = create2()(() => ({
       getBranchState: () => ({
         branchId: 0,
-        branchCount: 0
+        branchCount: 1
       }),
       switchToBranch: () => {
+        throw new Error("Not implemented");
       },
       editAt: async () => {
+        throw new Error("Not implemented");
       },
       reloadAt: async () => {
+        throw new Error("Not implemented");
       }
     }));
     return { useThread, useComposer, useBranchObserver };
   });
   return context;
 };
+
+// src/vercel/VercelAIAssistantProvider.tsx
+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 VercelAIAssistantProvider = ({
-  chat,
-  children
+  children,
+  ...rest
 }) => {
-  const context = useAIAssistantContext();
+  const context = useDummyAIAssistantContext();
+  const vercel = "chat" in rest ? rest.chat : rest.assistant;
+  const messages = useMemo4(() => {
+    return vercelToCachedThreadMessages(vercel.messages);
+  }, [vercel.messages]);
+  const maybeReload = "reload" in vercel ? vercel.reload : null;
+  const reload = useCallback2(async () => {
+    if (!maybeReload)
+      throw new Error("Reload not supported");
+    await maybeReload();
+  }, [maybeReload]);
+  const append = useCallback2(
+    async (message) => {
+      if (message.content[0]?.type !== "text") {
+        throw new Error("Only text content is currently supported");
+      }
+      await vercel.append({
+        role: message.role,
+        content: message.content[0].text
+      });
+    },
+    [vercel.append]
+  );
+  const stop = useCallback2(() => {
+    const lastMessage = vercel.messages.at(-1);
+    vercel.stop();
+    if (lastMessage?.role === "user") {
+      vercel.setInput(lastMessage.content);
+    }
+  }, [vercel.messages, vercel.stop, vercel.setInput]);
+  const isLoading = "isLoading" in vercel ? vercel.isLoading : vercel.status === "in_progress";
   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);
-          }
-        }
+        messages,
+        isLoading,
+        reload,
+        append,
+        stop
       },
       true
     );
-  }, [context, chat]);
+  }, [context, messages, reload, append, stop, isLoading]);
   useMemo4(() => {
     context.useComposer.setState({
-      canCancel: chat.isLoading,
-      value: chat.input,
-      setValue: chat.setInput
+      canCancel: isLoading,
+      value: vercel.input,
+      setValue: vercel.setInput
     });
-  }, [context, chat.isLoading, chat.input, chat.setInput]);
-  const branches = useChatWithBranches(chat);
+  }, [context, isLoading, vercel.input, vercel.setInput]);
+  const branches = useVercelAIBranches(vercel);
   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);
 };
+
+// src/vercel/VercelRSCAssistantProvider.tsx
+import {
+  useCallback as useCallback3,
+  useMemo as useMemo5
+} from "react";
+var ThreadMessageCache2 = /* @__PURE__ */ new WeakMap();
+var vercelToThreadMessage2 = (message) => {
+  if (message.role !== "user" && message.role !== "assistant")
+    throw new Error("Unsupported role");
+  return {
+    id: message.id,
+    role: message.role,
+    content: [{ type: "ui", display: message.display }]
+  };
+};
+var vercelToCachedThreadMessages2 = (messages) => {
+  return messages.map((m) => {
+    const cached = ThreadMessageCache2.get(m);
+    if (cached)
+      return cached;
+    const newMessage = vercelToThreadMessage2(m);
+    ThreadMessageCache2.set(m, newMessage);
+    return newMessage;
+  });
+};
+var VercelRSCAssistantProvider = ({
+  children,
+  messages: vercelMessages,
+  append: vercelAppend
+}) => {
+  const context = useDummyAIAssistantContext();
+  const messages = useMemo5(() => {
+    return vercelToCachedThreadMessages2(vercelMessages);
+  }, [vercelMessages]);
+  const append = useCallback3(
+    async (message) => {
+      if (message.content[0]?.type !== "text") {
+        throw new Error("Only text content is currently supported");
+      }
+      await vercelAppend({
+        role: message.role,
+        content: [{ type: "text", text: message.content[0].text }]
+      });
+    },
+    [vercelAppend]
+  );
+  useMemo5(() => {
+    context.useThread.setState({
+      messages,
+      append
+    });
+  }, [context, messages, append]);
+  return /* @__PURE__ */ React.createElement(AssistantContext.Provider, { value: context }, children);
+};
 export {
   actionBar_exports as ActionBarPrimitive,
   branchPicker_exports as BranchPickerPrimitive,
   composer_exports as ComposerPrimitive,
   message_exports as MessagePrimitive,
   thread_exports as ThreadPrimitive,
-  VercelAIAssistantProvider as VercelAIThreadProvider,
+  VercelAIAssistantProvider,
+  VercelRSCAssistantProvider as unstable_VercelRSCAssistantProvider,
   useMessageContext as unstable_useMessageContext,
   useBeginMessageEdit,
   useCopyMessage,

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@assistant-ui/react",
-  "version": "0.0.5",
+  "version": "0.0.7",
   "license": "MIT",
   "exports": {
     ".": {
@@ -28,7 +28,7 @@
     "@radix-ui/react-compose-refs": "^1.0.1",
     "@radix-ui/react-primitive": "^1.0.3",
     "@radix-ui/react-slot": "^1.0.2",
-    "ai": "^3.1.10",
+    "ai": "^3.1.12",
     "react-textarea-autosize": "^8.5.3",
     "zustand": "^4.5.2"
   },