@assistant-ui/react 0.0.12 → 0.0.14

package/dist/index.js

@@ -54,6 +54,7 @@ __export(thread_exports, {
   Messages: () => ThreadMessages,
   Root: () => ThreadRoot,
   ScrollToBottom: () => ThreadScrollToBottom,
+  Suggestion: () => ThreadSuggestion,
   Viewport: () => ThreadViewport
 });
 
@@ -74,7 +75,7 @@ var useAssistantContext = () => {
   const context = (0, import_react2.useContext)(AssistantContext);
   if (!context)
     throw new Error(
-      "useAssistantContext must be used within a AssistantProvider"
+      "This component must be used within a AssistantProvider"
     );
   return context;
 };
@@ -224,7 +225,7 @@ 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");
+    throw new Error("This component must be used within a MessageProvider");
   return context;
 };
 
@@ -329,6 +330,7 @@ var useMessageContext2 = () => {
   const [context] = (0, import_react8.useState)(() => {
     const useMessage = (0, import_zustand2.create)(() => ({
       message: null,
+      parentId: null,
       branches: [],
       isLast: false,
       isCopied: false,
@@ -343,15 +345,19 @@ var useMessageContext2 = () => {
         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 message.content[0].text;
+        const text = message.content.filter((part) => part.type === "text").map((part) => part.text).join("\n");
+        return text;
       },
       onSend: (text) => {
-        const message = useMessage.getState().message;
+        const { message, parentId } = useMessage.getState();
+        if (message.role !== "user")
+          throw new Error("Editing is only supported for user messages");
+        const nonTextParts = message.content.filter(
+          (part) => part.type !== "text" && part.type !== "ui"
+        );
         useThread.getState().append({
-          parentId: message.parentId,
-          content: [{ type: "text", text }]
+          parentId,
+          content: [{ type: "text", text }, ...nonTextParts]
         });
       }
     });
@@ -361,6 +367,7 @@ var useMessageContext2 = () => {
 };
 var MessageProvider = ({
   message,
+  parentId,
   children
 }) => {
   const { useThread } = useAssistantContext();
@@ -373,6 +380,7 @@ var MessageProvider = ({
     context.useMessage.setState(
       {
         message,
+        parentId,
         branches,
         isLast,
         isCopied,
@@ -382,7 +390,7 @@ var MessageProvider = ({
       },
       true
     );
-  }, [context, message, branches, isLast, isCopied, isHovering]);
+  }, [context, message, parentId, branches, isLast, isCopied, isHovering]);
   return /* @__PURE__ */ (0, import_jsx_runtime4.jsx)(MessageContext.Provider, { value: context, children });
 };
 
@@ -493,15 +501,21 @@ var ThreadMessages = ({ components }) => {
   if (messages.length === 0)
     return null;
   return /* @__PURE__ */ (0, import_jsx_runtime7.jsx)(import_jsx_runtime7.Fragment, { children: messages.map((message, idx) => {
-    return (
-      // biome-ignore lint/suspicious/noArrayIndexKey: fixes a11y issues with branch navigation
-      /* @__PURE__ */ (0, import_jsx_runtime7.jsxs)(MessageProvider, { message, children: [
-        /* @__PURE__ */ (0, import_jsx_runtime7.jsxs)(MessageIf, { user: true, children: [
-          /* @__PURE__ */ (0, import_jsx_runtime7.jsx)(ComposerIf, { editing: false, children: /* @__PURE__ */ (0, import_jsx_runtime7.jsx)(UserMessage, {}) }),
-          /* @__PURE__ */ (0, import_jsx_runtime7.jsx)(ComposerIf, { editing: true, children: /* @__PURE__ */ (0, import_jsx_runtime7.jsx)(EditComposer, {}) })
-        ] }),
-        /* @__PURE__ */ (0, import_jsx_runtime7.jsx)(MessageIf, { assistant: true, children: /* @__PURE__ */ (0, import_jsx_runtime7.jsx)(AssistantMessage, {}) })
-      ] }, idx)
+    const parentId = messages[idx - 1]?.id ?? null;
+    return /* @__PURE__ */ (0, import_jsx_runtime7.jsxs)(
+      MessageProvider,
+      {
+        message,
+        parentId,
+        children: [
+          /* @__PURE__ */ (0, import_jsx_runtime7.jsxs)(MessageIf, { user: true, children: [
+            /* @__PURE__ */ (0, import_jsx_runtime7.jsx)(ComposerIf, { editing: false, children: /* @__PURE__ */ (0, import_jsx_runtime7.jsx)(UserMessage, {}) }),
+            /* @__PURE__ */ (0, import_jsx_runtime7.jsx)(ComposerIf, { editing: true, children: /* @__PURE__ */ (0, import_jsx_runtime7.jsx)(EditComposer, {}) })
+          ] }),
+          /* @__PURE__ */ (0, import_jsx_runtime7.jsx)(MessageIf, { assistant: true, children: /* @__PURE__ */ (0, import_jsx_runtime7.jsx)(AssistantMessage, {}) })
+        ]
+      },
+      parentId ?? "__ROOT__"
     );
   }) });
 };
@@ -528,6 +542,33 @@ var ThreadScrollToBottom = (0, import_react10.forwardRef)(({ onClick, ...rest },
   );
 });
 
+// src/primitives/thread/ThreadSuggestion.tsx
+var import_primitive4 = require("@radix-ui/primitive");
+var import_react_primitive5 = require("@radix-ui/react-primitive");
+var import_react11 = require("react");
+var import_jsx_runtime9 = require("react/jsx-runtime");
+var ThreadSuggestion = (0, import_react11.forwardRef)(({ onClick, prompt, method, autoSend: send, ...rest }, ref) => {
+  const { useThread, useComposer } = useAssistantContext();
+  const isDisabled = useThread((t) => t.isRunning);
+  const handleApplySuggestion = () => {
+    const thread = useThread.getState();
+    const composer = useComposer.getState();
+    composer.setValue(prompt);
+    if (send && !thread.isRunning) {
+      composer.send();
+    }
+  };
+  return /* @__PURE__ */ (0, import_jsx_runtime9.jsx)(
+    import_react_primitive5.Primitive.button,
+    {
+      ...rest,
+      disabled: isDisabled,
+      ref,
+      onClick: (0, import_primitive4.composeEventHandlers)(onClick, handleApplySuggestion)
+    }
+  );
+});
+
 // src/primitives/composer/index.ts
 var composer_exports = {};
 __export(composer_exports, {
@@ -539,16 +580,16 @@ __export(composer_exports, {
 });
 
 // src/primitives/composer/ComposerRoot.tsx
-var import_primitive4 = require("@radix-ui/primitive");
+var import_primitive5 = require("@radix-ui/primitive");
 var import_react_compose_refs2 = require("@radix-ui/react-compose-refs");
-var import_react_primitive5 = require("@radix-ui/react-primitive");
-var import_react11 = require("react");
-var import_jsx_runtime9 = require("react/jsx-runtime");
-var ComposerRoot = (0, import_react11.forwardRef)(
+var import_react_primitive6 = require("@radix-ui/react-primitive");
+var import_react12 = require("react");
+var import_jsx_runtime10 = require("react/jsx-runtime");
+var ComposerRoot = (0, import_react12.forwardRef)(
   ({ onSubmit, ...rest }, forwardedRef) => {
     const { useViewport } = useAssistantContext();
     const { useComposer } = useComposerContext();
-    const formRef = (0, import_react11.useRef)(null);
+    const formRef = (0, import_react12.useRef)(null);
     const ref = (0, import_react_compose_refs2.useComposedRefs)(forwardedRef, formRef);
     const handleSubmit = (e) => {
       const composerState = useComposer.getState();
@@ -558,26 +599,26 @@ var ComposerRoot = (0, import_react11.forwardRef)(
       composerState.send();
       useViewport.getState().scrollToBottom();
     };
-    return /* @__PURE__ */ (0, import_jsx_runtime9.jsx)(
-      import_react_primitive5.Primitive.form,
+    return /* @__PURE__ */ (0, import_jsx_runtime10.jsx)(
+      import_react_primitive6.Primitive.form,
       {
         ...rest,
         ref,
-        onSubmit: (0, import_primitive4.composeEventHandlers)(onSubmit, handleSubmit)
+        onSubmit: (0, import_primitive5.composeEventHandlers)(onSubmit, handleSubmit)
       }
     );
   }
 );
 
 // src/primitives/composer/ComposerInput.tsx
-var import_primitive5 = require("@radix-ui/primitive");
+var import_primitive6 = require("@radix-ui/primitive");
 var import_react_compose_refs3 = require("@radix-ui/react-compose-refs");
 var import_react_slot = require("@radix-ui/react-slot");
-var import_react12 = require("react");
+var import_react13 = require("react");
 var import_react_textarea_autosize = __toESM(require("react-textarea-autosize"));
-var import_jsx_runtime10 = require("react/jsx-runtime");
-var ComposerInput = (0, import_react12.forwardRef)(
-  ({ autoFocus, asChild, disabled, onChange, onKeyDown, ...rest }, forwardedRef) => {
+var import_jsx_runtime11 = require("react/jsx-runtime");
+var ComposerInput = (0, import_react13.forwardRef)(
+  ({ autoFocus = false, asChild, disabled, onChange, onKeyDown, ...rest }, forwardedRef) => {
     const { useThread, useViewport } = useAssistantContext();
     const { useComposer, type } = useComposerContext();
     const value = useComposer((c) => {
@@ -603,10 +644,10 @@ var ComposerInput = (0, import_react12.forwardRef)(
         }
       }
     };
-    const textareaRef = (0, import_react12.useRef)(null);
+    const textareaRef = (0, import_react13.useRef)(null);
     const ref = (0, import_react_compose_refs3.useComposedRefs)(forwardedRef, textareaRef);
-    const autoFocusEnabled = autoFocus !== false && !disabled;
-    const focus = (0, import_react12.useCallback)(() => {
+    const autoFocusEnabled = autoFocus && !disabled;
+    const focus = (0, import_react13.useCallback)(() => {
       const textarea = textareaRef.current;
       if (!textarea || !autoFocusEnabled)
         return;
@@ -616,41 +657,41 @@ var ComposerInput = (0, import_react12.forwardRef)(
         textareaRef.current.value.length
       );
     }, [autoFocusEnabled]);
-    (0, import_react12.useEffect)(() => focus(), [focus]);
+    (0, import_react13.useEffect)(() => focus(), [focus]);
     useOnScrollToBottom(() => {
       if (type === "assistant") {
         focus();
       }
     });
-    return /* @__PURE__ */ (0, import_jsx_runtime10.jsx)(
+    return /* @__PURE__ */ (0, import_jsx_runtime11.jsx)(
       Component,
       {
         value,
         ...rest,
         ref,
         disabled,
-        onChange: (0, import_primitive5.composeEventHandlers)(onChange, (e) => {
+        onChange: (0, import_primitive6.composeEventHandlers)(onChange, (e) => {
           const composerState = useComposer.getState();
           if (!composerState.isEditing)
             return;
           return composerState.setValue(e.target.value);
         }),
-        onKeyDown: (0, import_primitive5.composeEventHandlers)(onKeyDown, handleKeyPress)
+        onKeyDown: (0, import_primitive6.composeEventHandlers)(onKeyDown, handleKeyPress)
       }
     );
   }
 );
 
 // src/primitives/composer/ComposerSend.tsx
-var import_react_primitive6 = require("@radix-ui/react-primitive");
-var import_react13 = require("react");
-var import_jsx_runtime11 = require("react/jsx-runtime");
-var ComposerSend = (0, import_react13.forwardRef)(
+var import_react_primitive7 = require("@radix-ui/react-primitive");
+var import_react14 = require("react");
+var import_jsx_runtime12 = require("react/jsx-runtime");
+var ComposerSend = (0, import_react14.forwardRef)(
   ({ disabled, ...rest }, ref) => {
     const { useComposer } = useComposerContext();
     const hasValue = useComposer((c) => c.isEditing && c.value.length > 0);
-    return /* @__PURE__ */ (0, import_jsx_runtime11.jsx)(
-      import_react_primitive6.Primitive.button,
+    return /* @__PURE__ */ (0, import_jsx_runtime12.jsx)(
+      import_react_primitive7.Primitive.button,
       {
         type: "submit",
         ...rest,
@@ -662,22 +703,22 @@ var ComposerSend = (0, import_react13.forwardRef)(
 );
 
 // src/primitives/composer/ComposerCancel.tsx
-var import_primitive6 = require("@radix-ui/primitive");
-var import_react_primitive7 = require("@radix-ui/react-primitive");
-var import_react14 = require("react");
-var import_jsx_runtime12 = require("react/jsx-runtime");
-var ComposerCancel = (0, import_react14.forwardRef)(({ onClick, ...rest }, ref) => {
+var import_primitive7 = require("@radix-ui/primitive");
+var import_react_primitive8 = require("@radix-ui/react-primitive");
+var import_react15 = require("react");
+var import_jsx_runtime13 = require("react/jsx-runtime");
+var ComposerCancel = (0, import_react15.forwardRef)(({ onClick, ...rest }, ref) => {
   const { useComposer } = useComposerContext();
   const handleCancel = () => {
     useComposer.getState().cancel();
   };
-  return /* @__PURE__ */ (0, import_jsx_runtime12.jsx)(
-    import_react_primitive7.Primitive.button,
+  return /* @__PURE__ */ (0, import_jsx_runtime13.jsx)(
+    import_react_primitive8.Primitive.button,
     {
       type: "button",
       ...rest,
       ref,
-      onClick: (0, import_primitive6.composeEventHandlers)(onClick, handleCancel)
+      onClick: (0, import_primitive7.composeEventHandlers)(onClick, handleCancel)
     }
   );
 });
@@ -693,10 +734,10 @@ __export(branchPicker_exports, {
 });
 
 // src/utils/context/combined/useCombinedStore.ts
-var import_react16 = require("react");
+var import_react17 = require("react");
 
 // src/utils/context/combined/createCombinedStore.ts
-var import_react15 = require("react");
+var import_react16 = require("react");
 var createCombinedStore = (stores) => {
   const subscribe = (callback) => {
     const unsubscribes = stores.map((store) => store.subscribe(callback));
@@ -708,13 +749,13 @@ var createCombinedStore = (stores) => {
   };
   return (selector) => {
     const getSnapshot = () => selector(...stores.map((store) => store.getState()));
-    return (0, import_react15.useSyncExternalStore)(subscribe, getSnapshot, getSnapshot);
+    return (0, import_react16.useSyncExternalStore)(subscribe, getSnapshot, getSnapshot);
   };
 };
 
 // src/utils/context/combined/useCombinedStore.ts
 var useCombinedStore = (stores, selector) => {
-  const useCombined = (0, import_react16.useMemo)(() => createCombinedStore(stores), stores);
+  const useCombined = (0, import_react17.useMemo)(() => createCombinedStore(stores), stores);
   return useCombined(selector);
 };
 
@@ -735,22 +776,22 @@ var useGoToNextBranch = () => {
 };
 
 // src/utils/createActionButton.tsx
-var import_primitive7 = require("@radix-ui/primitive");
-var import_react_primitive8 = require("@radix-ui/react-primitive");
-var import_react17 = require("react");
-var import_jsx_runtime13 = require("react/jsx-runtime");
+var import_primitive8 = require("@radix-ui/primitive");
+var import_react_primitive9 = require("@radix-ui/react-primitive");
+var import_react18 = require("react");
+var import_jsx_runtime14 = require("react/jsx-runtime");
 var createActionButton = (useActionButton) => {
-  return (0, import_react17.forwardRef)(
+  return (0, import_react18.forwardRef)(
     (props, forwardedRef) => {
       const onClick = useActionButton(props);
-      return /* @__PURE__ */ (0, import_jsx_runtime13.jsx)(
-        import_react_primitive8.Primitive.button,
+      return /* @__PURE__ */ (0, import_jsx_runtime14.jsx)(
+        import_react_primitive9.Primitive.button,
         {
           type: "button",
           disabled: !onClick,
           ...props,
           ref: forwardedRef,
-          onClick: (0, import_primitive7.composeEventHandlers)(props.onClick, onClick ?? void 0)
+          onClick: (0, import_primitive8.composeEventHandlers)(props.onClick, onClick ?? void 0)
         }
       );
     }
@@ -783,27 +824,27 @@ var useGoToPreviousBranch = () => {
 var BranchPickerPrevious = createActionButton(useGoToPreviousBranch);
 
 // src/primitives/branchPicker/BranchPickerCount.tsx
-var import_jsx_runtime14 = require("react/jsx-runtime");
+var import_jsx_runtime15 = require("react/jsx-runtime");
 var BranchPickerCount = () => {
   const { useMessage } = useMessageContext();
   const branchCount = useMessage((s) => s.branches.length);
-  return /* @__PURE__ */ (0, import_jsx_runtime14.jsx)(import_jsx_runtime14.Fragment, { children: branchCount });
+  return /* @__PURE__ */ (0, import_jsx_runtime15.jsx)(import_jsx_runtime15.Fragment, { children: branchCount });
 };
 
 // src/primitives/branchPicker/BranchPickerNumber.tsx
-var import_jsx_runtime15 = require("react/jsx-runtime");
+var import_jsx_runtime16 = require("react/jsx-runtime");
 var BranchPickerNumber = () => {
   const { useMessage } = useMessageContext();
   const branchIdx = useMessage((s) => s.branches.indexOf(s.message.id));
-  return /* @__PURE__ */ (0, import_jsx_runtime15.jsx)(import_jsx_runtime15.Fragment, { children: branchIdx + 1 });
+  return /* @__PURE__ */ (0, import_jsx_runtime16.jsx)(import_jsx_runtime16.Fragment, { children: branchIdx + 1 });
 };
 
 // src/primitives/branchPicker/BranchPickerRoot.tsx
-var import_react_primitive9 = require("@radix-ui/react-primitive");
-var import_react18 = require("react");
-var import_jsx_runtime16 = require("react/jsx-runtime");
-var BranchPickerRoot = (0, import_react18.forwardRef)(({ hideWhenSingleBranch, ...rest }, ref) => {
-  return /* @__PURE__ */ (0, import_jsx_runtime16.jsx)(MessageIf, { hasBranches: hideWhenSingleBranch ? true : void 0, children: /* @__PURE__ */ (0, import_jsx_runtime16.jsx)(import_react_primitive9.Primitive.div, { ...rest, ref }) });
+var import_react_primitive10 = require("@radix-ui/react-primitive");
+var import_react19 = require("react");
+var import_jsx_runtime17 = require("react/jsx-runtime");
+var BranchPickerRoot = (0, import_react19.forwardRef)(({ hideWhenSingleBranch, ...rest }, ref) => {
+  return /* @__PURE__ */ (0, import_jsx_runtime17.jsx)(MessageIf, { hasBranches: hideWhenSingleBranch ? true : void 0, children: /* @__PURE__ */ (0, import_jsx_runtime17.jsx)(import_react_primitive10.Primitive.div, { ...rest, ref }) });
 });
 
 // src/primitives/actionBar/index.ts
@@ -816,10 +857,10 @@ __export(actionBar_exports, {
 });
 
 // src/primitives/actionBar/ActionBarRoot.tsx
-var import_react_primitive10 = require("@radix-ui/react-primitive");
-var import_react19 = require("react");
-var import_jsx_runtime17 = require("react/jsx-runtime");
-var ActionBarRoot = (0, import_react19.forwardRef)(({ hideWhenRunning, autohide, autohideFloat, ...rest }, ref) => {
+var import_react_primitive11 = require("@radix-ui/react-primitive");
+var import_react20 = require("react");
+var import_jsx_runtime18 = require("react/jsx-runtime");
+var ActionBarRoot = (0, import_react20.forwardRef)(({ hideWhenRunning, autohide, autohideFloat, ...rest }, ref) => {
   const { useThread } = useAssistantContext();
   const { useMessage } = useMessageContext();
   const hideAndfloatStatus = useCombinedStore(
@@ -839,8 +880,8 @@ var ActionBarRoot = (0, import_react19.forwardRef)(({ hideWhenRunning, autohide,
   );
   if (hideAndfloatStatus === "hidden" /* Hidden */)
     return null;
-  return /* @__PURE__ */ (0, import_jsx_runtime17.jsx)(
-    import_react_primitive10.Primitive.div,
+  return /* @__PURE__ */ (0, import_jsx_runtime18.jsx)(
+    import_react_primitive11.Primitive.div,
     {
       "data-floating": hideAndfloatStatus === "floating" /* Floating */,
       ...rest,
@@ -879,10 +920,10 @@ var useReloadMessage = () => {
   if (disabled)
     return null;
   return () => {
-    const message = useMessage.getState().message;
+    const { message, parentId } = useMessage.getState();
     if (message.role !== "assistant")
       throw new Error("Reloading is only supported on assistant messages");
-    useThread.getState().startRun(message.parentId);
+    useThread.getState().startRun(parentId);
     useViewport.getState().scrollToBottom();
   };
 };
@@ -909,10 +950,10 @@ var useBeginMessageEdit = () => {
 var ActionBarEdit = createActionButton(useBeginMessageEdit);
 
 // src/adapters/vercel/VercelAIAssistantProvider.tsx
-var import_react22 = require("react");
+var import_react23 = require("react");
 
 // src/adapters/vercel/useDummyAIAssistantContext.tsx
-var import_react20 = require("react");
+var import_react21 = require("react");
 var import_zustand4 = require("zustand");
 
 // src/utils/context/stores/ViewportStore.tsx
@@ -958,7 +999,7 @@ var makeDummyThreadStore = () => {
   }));
 };
 var useDummyAIAssistantContext = () => {
-  const [context] = (0, import_react20.useState)(() => {
+  const [context] = (0, import_react21.useState)(() => {
     const useThread = makeDummyThreadStore();
     const useViewport = makeViewportStore();
     const useComposer = makeThreadComposerStore(useThread);
@@ -967,8 +1008,8 @@ var useDummyAIAssistantContext = () => {
   return context;
 };
 
-// src/adapters/vercel/useVercelAIBranches.tsx
-var import_react21 = require("react");
+// src/adapters/vercel/useVercelAIThreadState.tsx
+var import_react22 = require("react");
 
 // src/adapters/MessageRepository.tsx
 var import_non_secure = require("nanoid/non-secure");
@@ -996,17 +1037,26 @@ var MessageRepository = class {
     }
     return messages;
   }
-  addOrUpdateMessage(message) {
+  // TODO previousId is confusing
+  // TODO previousId does not fix children
+  // TODO cut / link operations
+  addOrUpdateMessage(parentId, message, previousId = message.id) {
     const item = this.messages.get(message.id);
     if (item) {
-      if (item.current.parentId !== message.parentId) {
-        this.deleteMessage(message.id);
-      } else {
-        item.current = message;
-        return;
+      if (item.prev?.current.id !== parentId) {
+        if ((item.prev?.current.id ?? null) !== parentId) {
+          this.deleteMessage(message.id);
+        } else {
+          item.current = message;
+          if (previousId !== message.id) {
+            this.messages.delete(previousId);
+            this.messages.set(message.id, item);
+          }
+          return;
+        }
       }
     }
-    const prev = message.parentId ? this.messages.get(message.parentId) : null;
+    const prev = parentId ? this.messages.get(parentId) : null;
     if (prev === void 0)
       throw new Error("Unexpected: Parent message not found");
     const newItem = {
@@ -1055,16 +1105,27 @@ var MessageRepository = class {
       this.head = message.prev ? findHead(message.prev) : null;
     }
   }
-  getOptimisticId = () => {
+  getOptimisticId() {
     let optimisticId;
     do {
       optimisticId = generateOptimisticId();
     } while (this.messages.has(optimisticId));
     return optimisticId;
-  };
+  }
+  commitOptimisticAppend(message) {
+    const optimisticIdUser = this.getOptimisticId();
+    this.addOrUpdateMessage(message.parentId, {
+      id: optimisticIdUser,
+      role: "user",
+      content: message.content,
+      createdAt: /* @__PURE__ */ new Date()
+    });
+    const optimisticIdAssistant = this.commitOptimisticRun(optimisticIdUser);
+    return [optimisticIdUser, optimisticIdAssistant];
+  }
   commitOptimisticRun(parentId) {
     const optimisticId = this.getOptimisticId();
-    this.addOrUpdateMessage({
+    this.addOrUpdateMessage(parentId, {
       id: optimisticId,
       role: "assistant",
       content: [
@@ -1073,7 +1134,6 @@ var MessageRepository = class {
           text: ""
         }
       ],
-      parentId,
       createdAt: /* @__PURE__ */ new Date()
     });
     return optimisticId;
@@ -1113,7 +1173,37 @@ var MessageRepository = class {
   }
 };
 
-// src/adapters/vercel/useVercelAIBranches.tsx
+// src/adapters/ThreadMessageConverter.tsx
+var ThreadMessageConverter = class {
+  constructor(converter2) {
+    this.converter = converter2;
+  }
+  cache = /* @__PURE__ */ new WeakMap();
+  convertMessages(messages) {
+    return messages.map((m) => {
+      const cached = this.cache.get(m);
+      if (cached)
+        return cached;
+      const newMessage = this.converter(m);
+      this.cache.set(m, newMessage);
+      return newMessage;
+    });
+  }
+};
+
+// src/adapters/vercel/useVercelAIThreadState.tsx
+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 }],
+    createdAt: message.createdAt ?? /* @__PURE__ */ new Date(),
+    innerMessage: message
+  };
+};
+var converter = new ThreadMessageConverter(vercelToThreadMessage);
 var sliceMessagesUntil = (messages, messageId) => {
   if (messageId == null)
     return [];
@@ -1127,195 +1217,188 @@ var sliceMessagesUntil = (messages, messageId) => {
 var hasUpcomingMessage = (isRunning, messages) => {
   return isRunning && messages[messages.length - 1]?.role !== "assistant";
 };
-var useVercelAIBranches = (chat, messages) => {
-  const [data] = (0, import_react21.useState)(() => new MessageRepository());
-  const isRunning = "isLoading" in chat ? chat.isLoading : chat.status === "in_progress";
-  const assistantOptimisticIdRef = (0, import_react21.useRef)(null);
-  const messagesEx = (0, import_react21.useMemo)(() => {
-    for (const message of messages) {
-      data.addOrUpdateMessage(message);
+var getIsRunning = (vercel) => {
+  if ("isLoading" in vercel)
+    return vercel.isLoading;
+  return vercel.status === "in_progress";
+};
+var useVercelAIThreadState = (vercel) => {
+  const [data] = (0, import_react22.useState)(() => new MessageRepository());
+  const vercelRef = (0, import_react22.useRef)(vercel);
+  vercelRef.current = vercel;
+  const isRunning = getIsRunning(vercelRef.current);
+  const assistantOptimisticIdRef = (0, import_react22.useRef)(null);
+  const messages = (0, import_react22.useMemo)(() => {
+    const vm = converter.convertMessages(vercel.messages);
+    for (let i = 0; i < vm.length; i++) {
+      const message = vm[i];
+      const parent = vm[i - 1];
+      data.addOrUpdateMessage(parent?.id ?? null, message);
     }
     if (assistantOptimisticIdRef.current) {
       data.deleteMessage(assistantOptimisticIdRef.current);
       assistantOptimisticIdRef.current = null;
     }
-    if (hasUpcomingMessage(isRunning, messages)) {
+    if (hasUpcomingMessage(isRunning, vm)) {
       assistantOptimisticIdRef.current = data.commitOptimisticRun(
-        messages.at(-1)?.id ?? null
+        vm.at(-1)?.id ?? null
       );
     }
-    data.resetHead(
-      assistantOptimisticIdRef.current ?? messages.at(-1)?.id ?? null
-    );
+    data.resetHead(assistantOptimisticIdRef.current ?? vm.at(-1)?.id ?? null);
     return data.getMessages();
-  }, [data, isRunning, messages]);
-  const getBranches = (0, import_react21.useCallback)(
+  }, [data, isRunning, vercel.messages]);
+  const getBranches = (0, import_react22.useCallback)(
     (messageId) => {
       return data.getBranches(messageId);
     },
     [data]
   );
-  const switchToBranch = (0, import_react21.useCallback)(
+  const switchToBranch = (0, import_react22.useCallback)(
     (messageId) => {
       data.switchToBranch(messageId);
-      chat.setMessages(
+      vercelRef.current.setMessages(
         data.getMessages().filter((m) => !isOptimisticId(m.id)).map((m) => m.innerMessage)
       );
     },
-    [data, chat.setMessages]
-  );
-  const reloadMaybe = "reload" in chat ? chat.reload : void 0;
-  const startRun = (0, import_react21.useCallback)(
-    async (parentId) => {
-      if (!reloadMaybe)
-        throw new Error("Reload not supported by Vercel AI SDK's useAssistant");
-      const newMessages = sliceMessagesUntil(chat.messages, parentId);
-      chat.setMessages(newMessages);
-      await reloadMaybe();
-    },
-    [chat.messages, chat.setMessages, reloadMaybe]
-  );
-  const append = (0, import_react21.useCallback)(
-    async (message) => {
-      if (message.content.length !== 1 || message.content[0]?.type !== "text")
-        throw new Error("Only text content is supported by Vercel AI SDK");
-      const newMessages = sliceMessagesUntil(chat.messages, message.parentId);
-      chat.setMessages(newMessages);
-      await chat.append({
-        role: "user",
-        content: message.content[0].text
-      });
-    },
-    [chat.messages, chat.setMessages, chat.append]
+    [data]
   );
-  return (0, import_react21.useMemo)(
+  const startRun = (0, import_react22.useCallback)(async (parentId) => {
+    const reloadMaybe = "reload" in vercelRef.current ? vercelRef.current.reload : void 0;
+    if (!reloadMaybe)
+      throw new Error("Reload not supported by Vercel AI SDK's useAssistant");
+    const newMessages = sliceMessagesUntil(
+      vercelRef.current.messages,
+      parentId
+    );
+    vercelRef.current.setMessages(newMessages);
+    await reloadMaybe();
+  }, []);
+  const append = (0, import_react22.useCallback)(async (message) => {
+    if (message.content.length !== 1 || message.content[0]?.type !== "text")
+      throw new Error("Only text content is supported by Vercel AI SDK");
+    const newMessages = sliceMessagesUntil(
+      vercelRef.current.messages,
+      message.parentId
+    );
+    vercelRef.current.setMessages(newMessages);
+    await vercelRef.current.append({
+      role: "user",
+      content: message.content[0].text
+    });
+  }, []);
+  const cancelRun = (0, import_react22.useCallback)(() => {
+    const lastMessage = vercelRef.current.messages.at(-1);
+    vercelRef.current.stop();
+    if (lastMessage?.role === "user") {
+      vercelRef.current.setInput(lastMessage.content);
+    }
+  }, []);
+  return (0, import_react22.useMemo)(
     () => ({
-      messages: messagesEx,
+      isRunning,
+      messages,
       getBranches,
       switchToBranch,
       append,
-      startRun
+      startRun,
+      cancelRun
     }),
-    [messagesEx, getBranches, switchToBranch, append, startRun]
+    [
+      isRunning,
+      messages,
+      getBranches,
+      switchToBranch,
+      append,
+      startRun,
+      cancelRun
+    ]
   );
 };
 
 // src/adapters/vercel/VercelAIAssistantProvider.tsx
-var import_jsx_runtime18 = require("react/jsx-runtime");
-var ThreadMessageCache = /* @__PURE__ */ new WeakMap();
-var vercelToThreadMessage = (message, parentId) => {
-  if (message.role !== "user" && message.role !== "assistant")
-    throw new Error("Unsupported role");
-  return {
-    parentId,
-    id: message.id,
-    role: message.role,
-    content: [{ type: "text", text: message.content }],
-    createdAt: message.createdAt ?? /* @__PURE__ */ new Date(),
-    innerMessage: message
-  };
-};
-var vercelToCachedThreadMessages = (messages) => {
-  return messages.map((m, idx) => {
-    const cached = ThreadMessageCache.get(m);
-    const parentId = messages[idx - 1]?.id ?? null;
-    if (cached && cached.parentId === parentId)
-      return cached;
-    const newMessage = vercelToThreadMessage(m, parentId);
-    ThreadMessageCache.set(m, newMessage);
-    return newMessage;
-  });
-};
+var import_jsx_runtime19 = require("react/jsx-runtime");
 var VercelAIAssistantProvider = ({
   children,
   ...rest
 }) => {
   const context = useDummyAIAssistantContext();
   const vercel = "chat" in rest ? rest.chat : rest.assistant;
-  const messages = (0, import_react22.useMemo)(() => {
-    return vercelToCachedThreadMessages(vercel.messages);
-  }, [vercel.messages]);
-  const branches = useVercelAIBranches(vercel, messages);
-  const cancelRun = (0, import_react22.useCallback)(() => {
-    const lastMessage = vercel.messages.at(-1);
-    vercel.stop();
-    if (lastMessage?.role === "user") {
-      vercel.setInput(lastMessage.content);
-    }
-  }, [vercel.messages, vercel.stop, vercel.setInput]);
-  const isRunning = "isLoading" in vercel ? vercel.isLoading : vercel.status === "in_progress";
-  (0, import_react22.useMemo)(() => {
-    context.useThread.setState(
-      {
-        messages: branches.messages,
-        isRunning,
-        getBranches: branches.getBranches,
-        switchToBranch: branches.switchToBranch,
-        append: branches.append,
-        startRun: branches.startRun,
-        cancelRun
-      },
-      true
-    );
-  }, [context, isRunning, cancelRun, branches]);
-  (0, import_react22.useMemo)(() => {
+  const threadState = useVercelAIThreadState(vercel);
+  (0, import_react23.useMemo)(() => {
+    context.useThread.setState(threadState, true);
+  }, [context, threadState]);
+  (0, import_react23.useMemo)(() => {
     context.useComposer.setState({
       value: vercel.input,
       setValue: vercel.setInput
     });
   }, [context, vercel.input, vercel.setInput]);
-  return /* @__PURE__ */ (0, import_jsx_runtime18.jsx)(AssistantContext.Provider, { value: context, children });
+  return /* @__PURE__ */ (0, import_jsx_runtime19.jsx)(AssistantContext.Provider, { value: context, children });
 };
 
 // src/adapters/vercel/VercelRSCAssistantProvider.tsx
-var import_react23 = require("react");
-var import_jsx_runtime19 = require("react/jsx-runtime");
-var ThreadMessageCache2 = /* @__PURE__ */ new WeakMap();
-var vercelToThreadMessage2 = (parentId, message) => {
+var import_react24 = require("react");
+var import_jsx_runtime20 = require("react/jsx-runtime");
+var vercelToThreadMessage2 = (message) => {
   if (message.role !== "user" && message.role !== "assistant")
     throw new Error("Unsupported role");
   return {
-    parentId,
     id: message.id,
     role: message.role,
     content: [{ type: "ui", display: message.display }],
     createdAt: message.createdAt ?? /* @__PURE__ */ new Date()
   };
 };
-var vercelToCachedThreadMessages2 = (messages) => {
-  return messages.map((m, idx) => {
-    const cached = ThreadMessageCache2.get(m);
-    const parentId = messages[idx - 1]?.id ?? null;
-    if (cached && cached.parentId === parentId)
-      return cached;
-    const newMessage = vercelToThreadMessage2(parentId, m);
-    ThreadMessageCache2.set(m, newMessage);
-    return newMessage;
-  });
-};
-var VercelRSCAssistantProvider = ({ children, messages: vercelMessages, append: vercelAppend }) => {
+var VercelRSCAssistantProvider = ({
+  children,
+  convertMessage,
+  messages: vercelMessages,
+  append: appendCallback,
+  edit,
+  reload
+}) => {
   const context = useDummyAIAssistantContext();
-  const messages = (0, import_react23.useMemo)(() => {
-    return vercelToCachedThreadMessages2(vercelMessages);
-  }, [vercelMessages]);
-  const append = (0, import_react23.useCallback)(
+  const converter2 = (0, import_react24.useMemo)(() => {
+    const rscConverter = convertMessage ?? ((m) => m);
+    return new ThreadMessageConverter((m) => {
+      return vercelToThreadMessage2(rscConverter(m));
+    });
+  }, [convertMessage]);
+  const messages = (0, import_react24.useMemo)(() => {
+    return converter2.convertMessages(vercelMessages);
+  }, [converter2, vercelMessages]);
+  const append = (0, import_react24.useCallback)(
     async (message) => {
-      if (message.parentId !== (context.useThread.getState().messages.at(-1)?.id ?? null))
-        throw new Error("Unexpected: Message editing is not supported");
-      if (message.content[0]?.type !== "text") {
-        throw new Error("Only text content is currently supported");
+      if (message.parentId !== (context.useThread.getState().messages.at(-1)?.id ?? null)) {
+        if (!edit)
+          throw new Error(
+            "Unexpected: Message editing is not supported, no edit callback was provided to VercelRSCAssistantProvider."
+          );
+        await edit(message);
+      } else {
+        await appendCallback(message);
       }
-      await vercelAppend(message);
     },
-    [context, vercelAppend]
+    [context, appendCallback, edit]
   );
-  (0, import_react23.useMemo)(() => {
+  const startRun = (0, import_react24.useCallback)(
+    async (parentId) => {
+      if (!reload)
+        throw new Error(
+          "Unexpected: Message reloading is not supported, no reload callback was provided to VercelRSCAssistantProvider."
+        );
+      await reload(parentId);
+    },
+    [reload]
+  );
+  (0, import_react24.useMemo)(() => {
     context.useThread.setState({
       messages,
-      append
+      append,
+      startRun
     });
-  }, [context, messages, append]);
-  return /* @__PURE__ */ (0, import_jsx_runtime19.jsx)(AssistantContext.Provider, { value: context, children });
+  }, [context, messages, append, startRun]);
+  return /* @__PURE__ */ (0, import_jsx_runtime20.jsx)(AssistantContext.Provider, { value: context, children });
 };
 // Annotate the CommonJS export names for ESM import in node:
 0 && (module.exports = {