@assistant-ui/react 0.0.17 → 0.0.18

package/dist/index.d.mts

@@ -3,7 +3,7 @@ import { FC, ReactNode, PropsWithChildren, ComponentType } from 'react';
 import { TextareaAutosizeProps } from 'react-textarea-autosize';
 import { UseChatHelpers, UseAssistantHelpers } from 'ai/react';
 import * as react_jsx_runtime from 'react/jsx-runtime';
-import { Message } from 'ai';
+import { LanguageModel, Message } from 'ai';
 import { UseBoundStore, StoreApi } from 'zustand';
 
 declare const ThreadRoot: react.ForwardRefExoticComponent<Pick<Omit<react.DetailedHTMLProps<react.HTMLAttributes<HTMLDivElement>, HTMLDivElement>, "ref"> & {
@@ -133,7 +133,7 @@ type ToolCallContentPart = {
     result?: object;
 };
 type UserContentPart = TextContentPart | ImageContentPart | UIContentPart;
-type AssistantContentPart = TextContentPart | ImageContentPart | UIContentPart | ToolCallContentPart;
+type AssistantContentPart = TextContentPart | UIContentPart | ToolCallContentPart;
 type AppendContentPart = TextContentPart | ImageContentPart;
 type BaseMessage = {
     id: string;
@@ -312,6 +312,70 @@ type RSCMessageConverter<T> = {
 type VercelRSCAssistantProviderProps<T = VercelRSCMessage> = VercelRSCAssistantProviderBaseProps<T> & (T extends VercelRSCMessage ? object : RSCMessageConverter<T>);
 declare const VercelRSCAssistantProvider: <T extends WeakKey = VercelRSCMessage>({ children, convertMessage, messages: vercelMessages, append: appendCallback, edit, reload, }: VercelRSCAssistantProviderProps<T>) => react_jsx_runtime.JSX.Element;
 
+type MessageUpdateCallback = (parentId: string | null, message: ThreadMessage) => void;
+type StatusUpdateCallback = (isRunning: boolean) => void;
+type Unsubscribe = () => void;
+type ThreadRuntime = {
+    append(message: AppendMessage): Promise<{
+        parentId: string;
+        id: string;
+    }>;
+    startRun(parentId: string | null): Promise<{
+        id: string;
+    }>;
+    cancelRun(): void;
+    subscribeToMessageUpdates(callback: MessageUpdateCallback): Unsubscribe;
+    subscribeToStatusUpdates(callback: StatusUpdateCallback): Unsubscribe;
+};
+
+type AssistantProviderProps = {
+    runtime: ThreadRuntime;
+};
+declare const AssistantProvider: FC<PropsWithChildren<AssistantProviderProps>>;
+
+type ChatModelRunResult = {
+    content: AssistantContentPart[];
+};
+type ChatModelRunOptions = {
+    messages: ThreadMessage[];
+    abortSignal: AbortSignal;
+    onUpdate: (result: ChatModelRunResult) => void;
+};
+type ChatModelAdapter = {
+    run: (options: ChatModelRunOptions) => Promise<ChatModelRunResult>;
+};
+
+declare class LocalRuntime implements ThreadRuntime {
+    adapter: ChatModelAdapter;
+    private _messageUpdateCallbacks;
+    private _statusUpdateCallbacks;
+    private abortController;
+    private repository;
+    constructor(adapter: ChatModelAdapter);
+    append(message: AppendMessage): Promise<{
+        parentId: string;
+        id: string;
+    }>;
+    startRun(parentId: string | null): Promise<{
+        id: string;
+    }>;
+    private addOrUpdateMessage;
+    private run;
+    cancelRun(): void;
+    subscribeToMessageUpdates(callback: MessageUpdateCallback): Unsubscribe;
+    subscribeToStatusUpdates(callback: StatusUpdateCallback): Unsubscribe;
+}
+
+declare const useLocalRuntime: (adapter: ChatModelAdapter) => LocalRuntime;
+
+declare class VercelModelAdapter implements ChatModelAdapter {
+    private readonly model;
+    constructor(model: LanguageModel);
+    run({ messages, abortSignal, onUpdate }: ChatModelRunOptions): Promise<{
+        content: AssistantContentPart[];
+    }>;
+}
+
 declare const getVercelMessage: (message: ThreadMessage) => Message | undefined;
 declare const getVercelRSCMessage: <T>(message: ThreadMessage) => T | undefined;
 
@@ -346,4 +410,4 @@ declare const useGoToNextBranch: () => (() => void) | null;
 
 declare const useGoToPreviousBranch: () => (() => void) | null;
 
-export { index$1 as ActionBarPrimitive, type AppendContentPart, type AppendMessage, index$2 as BranchPickerPrimitive, index$4 as ComposerPrimitive, index as ContentPartPrimitive, type ImageContentPart, index$3 as MessagePrimitive, type VercelRSCMessage as RSCMessage, type TextContentPart, index$5 as ThreadPrimitive, VercelAIAssistantProvider, type VercelAIAssistantProviderProps, VercelRSCAssistantProvider, type VercelRSCAssistantProviderProps, getVercelMessage as unstable_getVercelMessage, getVercelRSCMessage as unstable_getVercelRSCMessage, useMessageContext as unstable_useMessageContext, useBeginMessageEdit, useCopyMessage, useGoToNextBranch, useGoToPreviousBranch, useReloadMessage };
+export { index$1 as ActionBarPrimitive, type AppendContentPart, type AppendMessage, index$2 as BranchPickerPrimitive, index$4 as ComposerPrimitive, index as ContentPartPrimitive, type ImageContentPart, index$3 as MessagePrimitive, type VercelRSCMessage as RSCMessage, type TextContentPart, index$5 as ThreadPrimitive, VercelAIAssistantProvider, type VercelAIAssistantProviderProps, VercelRSCAssistantProvider, type VercelRSCAssistantProviderProps, AssistantProvider as unstable_AssistantProvider, type ChatModelAdapter as unstable_ChatModelAdapter, type ChatModelRunOptions as unstable_ChatModelRunOptions, VercelModelAdapter as unstable_VercelModelAdapter, getVercelMessage as unstable_getVercelMessage, getVercelRSCMessage as unstable_getVercelRSCMessage, useLocalRuntime as unstable_useLocalRuntime, useMessageContext as unstable_useMessageContext, useBeginMessageEdit, useCopyMessage, useGoToNextBranch, useGoToPreviousBranch, useReloadMessage };

package/dist/index.d.ts

@@ -3,7 +3,7 @@ import { FC, ReactNode, PropsWithChildren, ComponentType } from 'react';
 import { TextareaAutosizeProps } from 'react-textarea-autosize';
 import { UseChatHelpers, UseAssistantHelpers } from 'ai/react';
 import * as react_jsx_runtime from 'react/jsx-runtime';
-import { Message } from 'ai';
+import { LanguageModel, Message } from 'ai';
 import { UseBoundStore, StoreApi } from 'zustand';
 
 declare const ThreadRoot: react.ForwardRefExoticComponent<Pick<Omit<react.DetailedHTMLProps<react.HTMLAttributes<HTMLDivElement>, HTMLDivElement>, "ref"> & {
@@ -133,7 +133,7 @@ type ToolCallContentPart = {
     result?: object;
 };
 type UserContentPart = TextContentPart | ImageContentPart | UIContentPart;
-type AssistantContentPart = TextContentPart | ImageContentPart | UIContentPart | ToolCallContentPart;
+type AssistantContentPart = TextContentPart | UIContentPart | ToolCallContentPart;
 type AppendContentPart = TextContentPart | ImageContentPart;
 type BaseMessage = {
     id: string;
@@ -312,6 +312,70 @@ type RSCMessageConverter<T> = {
 type VercelRSCAssistantProviderProps<T = VercelRSCMessage> = VercelRSCAssistantProviderBaseProps<T> & (T extends VercelRSCMessage ? object : RSCMessageConverter<T>);
 declare const VercelRSCAssistantProvider: <T extends WeakKey = VercelRSCMessage>({ children, convertMessage, messages: vercelMessages, append: appendCallback, edit, reload, }: VercelRSCAssistantProviderProps<T>) => react_jsx_runtime.JSX.Element;
 
+type MessageUpdateCallback = (parentId: string | null, message: ThreadMessage) => void;
+type StatusUpdateCallback = (isRunning: boolean) => void;
+type Unsubscribe = () => void;
+type ThreadRuntime = {
+    append(message: AppendMessage): Promise<{
+        parentId: string;
+        id: string;
+    }>;
+    startRun(parentId: string | null): Promise<{
+        id: string;
+    }>;
+    cancelRun(): void;
+    subscribeToMessageUpdates(callback: MessageUpdateCallback): Unsubscribe;
+    subscribeToStatusUpdates(callback: StatusUpdateCallback): Unsubscribe;
+};
+
+type AssistantProviderProps = {
+    runtime: ThreadRuntime;
+};
+declare const AssistantProvider: FC<PropsWithChildren<AssistantProviderProps>>;
+
+type ChatModelRunResult = {
+    content: AssistantContentPart[];
+};
+type ChatModelRunOptions = {
+    messages: ThreadMessage[];
+    abortSignal: AbortSignal;
+    onUpdate: (result: ChatModelRunResult) => void;
+};
+type ChatModelAdapter = {
+    run: (options: ChatModelRunOptions) => Promise<ChatModelRunResult>;
+};
+
+declare class LocalRuntime implements ThreadRuntime {
+    adapter: ChatModelAdapter;
+    private _messageUpdateCallbacks;
+    private _statusUpdateCallbacks;
+    private abortController;
+    private repository;
+    constructor(adapter: ChatModelAdapter);
+    append(message: AppendMessage): Promise<{
+        parentId: string;
+        id: string;
+    }>;
+    startRun(parentId: string | null): Promise<{
+        id: string;
+    }>;
+    private addOrUpdateMessage;
+    private run;
+    cancelRun(): void;
+    subscribeToMessageUpdates(callback: MessageUpdateCallback): Unsubscribe;
+    subscribeToStatusUpdates(callback: StatusUpdateCallback): Unsubscribe;
+}
+
+declare const useLocalRuntime: (adapter: ChatModelAdapter) => LocalRuntime;
+
+declare class VercelModelAdapter implements ChatModelAdapter {
+    private readonly model;
+    constructor(model: LanguageModel);
+    run({ messages, abortSignal, onUpdate }: ChatModelRunOptions): Promise<{
+        content: AssistantContentPart[];
+    }>;
+}
+
 declare const getVercelMessage: (message: ThreadMessage) => Message | undefined;
 declare const getVercelRSCMessage: <T>(message: ThreadMessage) => T | undefined;
 
@@ -346,4 +410,4 @@ declare const useGoToNextBranch: () => (() => void) | null;
 
 declare const useGoToPreviousBranch: () => (() => void) | null;
 
-export { index$1 as ActionBarPrimitive, type AppendContentPart, type AppendMessage, index$2 as BranchPickerPrimitive, index$4 as ComposerPrimitive, index as ContentPartPrimitive, type ImageContentPart, index$3 as MessagePrimitive, type VercelRSCMessage as RSCMessage, type TextContentPart, index$5 as ThreadPrimitive, VercelAIAssistantProvider, type VercelAIAssistantProviderProps, VercelRSCAssistantProvider, type VercelRSCAssistantProviderProps, getVercelMessage as unstable_getVercelMessage, getVercelRSCMessage as unstable_getVercelRSCMessage, useMessageContext as unstable_useMessageContext, useBeginMessageEdit, useCopyMessage, useGoToNextBranch, useGoToPreviousBranch, useReloadMessage };
+export { index$1 as ActionBarPrimitive, type AppendContentPart, type AppendMessage, index$2 as BranchPickerPrimitive, index$4 as ComposerPrimitive, index as ContentPartPrimitive, type ImageContentPart, index$3 as MessagePrimitive, type VercelRSCMessage as RSCMessage, type TextContentPart, index$5 as ThreadPrimitive, VercelAIAssistantProvider, type VercelAIAssistantProviderProps, VercelRSCAssistantProvider, type VercelRSCAssistantProviderProps, AssistantProvider as unstable_AssistantProvider, type ChatModelAdapter as unstable_ChatModelAdapter, type ChatModelRunOptions as unstable_ChatModelRunOptions, VercelModelAdapter as unstable_VercelModelAdapter, getVercelMessage as unstable_getVercelMessage, getVercelRSCMessage as unstable_getVercelRSCMessage, useLocalRuntime as unstable_useLocalRuntime, useMessageContext as unstable_useMessageContext, useBeginMessageEdit, useCopyMessage, useGoToNextBranch, useGoToPreviousBranch, useReloadMessage };

package/dist/index.js

@@ -38,8 +38,11 @@ __export(src_exports, {
   ThreadPrimitive: () => thread_exports,
   VercelAIAssistantProvider: () => VercelAIAssistantProvider,
   VercelRSCAssistantProvider: () => VercelRSCAssistantProvider,
+  unstable_AssistantProvider: () => AssistantProvider,
+  unstable_VercelModelAdapter: () => VercelModelAdapter,
   unstable_getVercelMessage: () => getVercelMessage,
   unstable_getVercelRSCMessage: () => getVercelRSCMessage,
+  unstable_useLocalRuntime: () => useLocalRuntime,
   unstable_useMessageContext: () => useMessageContext,
   useBeginMessageEdit: () => useBeginMessageEdit,
   useCopyMessage: () => useCopyMessage,
@@ -1090,7 +1093,7 @@ var useDummyAIAssistantContext = () => {
 var import_react_use_callback_ref3 = require("@radix-ui/react-use-callback-ref");
 var import_react25 = require("react");
 
-// src/adapters/MessageRepository.tsx
+// src/adapters/idUtils.tsx
 var import_non_secure = require("nanoid/non-secure");
 var generateId = (0, import_non_secure.customAlphabet)(
   "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz",
@@ -1098,6 +1101,8 @@ var generateId = (0, import_non_secure.customAlphabet)(
 );
 var optimisticPrefix = "__optimistic__";
 var generateOptimisticId = () => `${optimisticPrefix}${generateId()}`;
+
+// src/adapters/MessageRepository.tsx
 var findHead = (message) => {
   if (message.next) return findHead(message.next);
   return message;
@@ -1109,15 +1114,6 @@ var MessageRepository = class {
   root = {
     children: []
   };
-  getFallbackChild(p) {
-    const childId = p.children.at(-1);
-    const child = childId ? this.messages.get(childId) : null;
-    if (child === void 0)
-      throw new Error(
-        "MessageRepository(getFallbackChild): Child message not found. This is likely an internal bug in assistant-ui."
-      );
-    return child;
-  }
   performOp(newParent, child, operation) {
     const parentOrRoot = child.prev ?? this.root;
     const newParentOrRoot = newParent ?? this.root;
@@ -1127,7 +1123,14 @@ var MessageRepository = class {
         (m) => m !== child.current.id
       );
       if (child.prev?.next === child) {
-        child.prev.next = this.getFallbackChild(child.prev);
+        const fallbackId = child.prev.children.at(-1);
+        const fallback = fallbackId ? this.messages.get(fallbackId) : null;
+        if (fallback === void 0) {
+          throw new Error(
+            "MessageRepository(performOp/cut): Fallback sibling message not found. This is likely an internal bug in assistant-ui."
+          );
+        }
+        child.prev.next = fallback;
       }
     }
     if (operation !== "cut") {
@@ -1186,14 +1189,14 @@ var MessageRepository = class {
     });
     return optimisticId;
   }
-  deleteMessage(messageId, newParentId) {
+  deleteMessage(messageId, replacementId) {
     const message = this.messages.get(messageId);
-    const newParent = newParentId ? this.messages.get(newParentId) : null;
+    const replacement = replacementId ? this.messages.get(replacementId) : null;
     if (!message)
       throw new Error(
         "MessageRepository(deleteMessage): Optimistic message not found. This is likely an internal bug in assistant-ui."
       );
-    if (newParent === void 0)
+    if (replacement === void 0)
       throw new Error(
         "MessageRepository(deleteMessage): New message not found. This is likely an internal bug in assistant-ui."
       );
@@ -1203,11 +1206,11 @@ var MessageRepository = class {
         throw new Error(
           "MessageRepository(deleteMessage): Child message not found. This is likely an internal bug in assistant-ui."
         );
-      this.performOp(newParent, childMessage, "relink");
+      this.performOp(replacement, childMessage, "relink");
     }
     this.messages.delete(messageId);
     if (this.head === message) {
-      this.head = this.getFallbackChild(message.prev ?? this.root);
+      this.head = replacement;
     }
     this.performOp(null, message, "cut");
   }
@@ -1250,17 +1253,13 @@ var MessageRepository = class {
   }
 };
 
-// src/adapters/ThreadMessageConverter.tsx
+// src/adapters/ThreadMessageConverter.ts
 var ThreadMessageConverter = class {
-  constructor(converter) {
-    this.converter = converter;
-  }
   cache = /* @__PURE__ */ new WeakMap();
-  convertMessages(messages) {
+  convertMessages(converter, messages) {
     return messages.map((m) => {
       const cached = this.cache.get(m);
-      if (cached) return cached;
-      const newMessage = this.converter(m);
+      const newMessage = converter(m, cached);
       this.cache.set(m, newMessage);
       return newMessage;
     });
@@ -1333,16 +1332,17 @@ var getIsRunning = (vercel) => {
 var useVercelAIThreadState = (vercel) => {
   const [data] = (0, import_react25.useState)(() => new MessageRepository());
   const isRunning = getIsRunning(vercel);
-  const convertCallback = (0, import_react_use_callback_ref3.useCallbackRef)((message) => {
-    return vercelToThreadMessage(
-      message,
-      vercel.messages.at(-1) === message && isRunning ? "in_progress" : "done"
-    );
-  });
-  const converter = new ThreadMessageConverter(convertCallback);
+  const converter = (0, import_react25.useMemo)(() => new ThreadMessageConverter(), []);
   const assistantOptimisticIdRef = (0, import_react25.useRef)(null);
   const messages = (0, import_react25.useMemo)(() => {
-    const vm = converter.convertMessages(vercel.messages);
+    const lastMessageId = vercel.messages.at(-1)?.id;
+    const convertCallback = (message, cache) => {
+      const status = lastMessageId === message.id && isRunning ? "in_progress" : "done";
+      if (cache && (cache.role === "user" || cache.status === status))
+        return cache;
+      return vercelToThreadMessage(message, status);
+    };
+    const vm = converter.convertMessages(convertCallback, vercel.messages);
     for (let i = 0; i < vm.length; i++) {
       const message = vm[i];
       const parent = vm[i - 1];
@@ -1490,15 +1490,17 @@ var VercelRSCAssistantProvider = ({
     setIsRunning(true);
     return callback.finally(() => setIsRunning(false));
   }, []);
-  const converter = (0, import_react27.useMemo)(() => {
+  const [converter, convertCallback] = (0, import_react27.useMemo)(() => {
     const rscConverter = convertMessage ?? ((m) => m);
-    return new ThreadMessageConverter((m) => {
+    const convertCallback2 = (m, cache) => {
+      if (cache) return cache;
       return vercelToThreadMessage2(rscConverter, m);
-    });
+    };
+    return [new ThreadMessageConverter(), convertCallback2];
   }, [convertMessage]);
   const messages = (0, import_react27.useMemo)(() => {
-    return converter.convertMessages(vercelMessages);
-  }, [converter, vercelMessages]);
+    return converter.convertMessages(convertCallback, vercelMessages);
+  }, [converter, convertCallback, vercelMessages]);
   const append = (0, import_react27.useCallback)(
     async (message) => {
       if (message.parentId !== (context.useThread.getState().messages.at(-1)?.id ?? null)) {
@@ -1539,6 +1541,259 @@ var VercelRSCAssistantProvider = ({
   }, [context, messages, isRunning, append, startRun]);
   return /* @__PURE__ */ (0, import_jsx_runtime22.jsx)(AssistantContext.Provider, { value: context, children });
 };
+
+// src/adapters/core/utils/useAssistantContext.tsx
+var import_react28 = require("react");
+var import_zustand6 = require("zustand");
+
+// src/adapters/core/utils/AssistantMessageRepository.tsx
+var AssistantMessageRepository = class {
+  constructor(flushCallback) {
+    this.flushCallback = flushCallback;
+  }
+  repository = new MessageRepository();
+  getBranches(messageId) {
+    return this.repository.getBranches(messageId);
+  }
+  withModifications(callback) {
+    const res = callback(this.repository);
+    this.flushCallback(this.repository.getMessages());
+    return res;
+  }
+};
+
+// src/adapters/core/utils/useAssistantContext.tsx
+var makeThreadStore = (runtimeRef) => {
+  const repository = new AssistantMessageRepository((messages) => {
+    useThread.setState({ messages });
+  });
+  const useThread = (0, import_zustand6.create)(() => ({
+    messages: [],
+    isRunning: false,
+    getBranches: (messageId) => repository.getBranches(messageId),
+    switchToBranch: (branchId) => {
+      repository.withModifications((repository2) => {
+        repository2.switchToBranch(branchId);
+      });
+    },
+    startRun: async (parentId) => {
+      const optimisticId = repository.withModifications((repository2) => {
+        const optimisticId2 = repository2.appendOptimisticMessage(parentId, {
+          role: "assistant",
+          content: [{ type: "text", text: "" }]
+        });
+        repository2.resetHead(optimisticId2);
+        return optimisticId2;
+      });
+      const { id } = await runtimeRef.current.startRun(parentId);
+      repository.withModifications((repository2) => {
+        repository2.deleteMessage(optimisticId, id);
+      });
+    },
+    append: async (message) => {
+      const [parentOptimisticId, optimisticId] = repository.withModifications(
+        (repository2) => {
+          const parentOptimisticId2 = repository2.appendOptimisticMessage(
+            message.parentId,
+            {
+              role: "user",
+              content: message.content
+            }
+          );
+          const optimisticId2 = repository2.appendOptimisticMessage(
+            parentOptimisticId2,
+            {
+              role: "assistant",
+              content: [{ type: "text", text: "" }]
+            }
+          );
+          repository2.resetHead(optimisticId2);
+          return [parentOptimisticId2, optimisticId2];
+        }
+      );
+      const { parentId, id } = await runtimeRef.current.append(message);
+      repository.withModifications((repository2) => {
+        repository2.deleteMessage(parentOptimisticId, parentId);
+        repository2.deleteMessage(optimisticId, id);
+      });
+    },
+    cancelRun: () => runtimeRef.current.cancelRun()
+  }));
+  const onNewMessage = (parentId, message) => {
+    repository.withModifications((repository2) => {
+      repository2.addOrUpdateMessage(parentId, message);
+    });
+  };
+  const onRunningChange = (isRunning) => {
+    useThread.setState({ isRunning });
+  };
+  return {
+    useThread,
+    onNewMessage,
+    onRunningChange
+  };
+};
+var useAssistantContext2 = (runtime) => {
+  const runtimeRef = (0, import_react28.useRef)(runtime);
+  runtimeRef.current = runtime;
+  const [{ context, onNewMessage, onRunningChange }] = (0, import_react28.useState)(() => {
+    const { useThread, onNewMessage: onNewMessage2, onRunningChange: onRunningChange2 } = makeThreadStore(runtimeRef);
+    const useViewport = makeViewportStore();
+    const useComposer = makeThreadComposerStore(useThread);
+    return {
+      context: { useViewport, useThread, useComposer },
+      onNewMessage: onNewMessage2,
+      onRunningChange: onRunningChange2
+    };
+  });
+  (0, import_react28.useEffect)(() => {
+    return runtime.subscribeToMessageUpdates(onNewMessage);
+  }, [runtime, onNewMessage]);
+  (0, import_react28.useEffect)(() => {
+    return runtime.subscribeToStatusUpdates(onRunningChange);
+  }, [runtime, onRunningChange]);
+  return context;
+};
+
+// src/adapters/core/AssistantProvider.tsx
+var import_jsx_runtime23 = require("react/jsx-runtime");
+var AssistantProvider = ({ children, runtime }) => {
+  const context = useAssistantContext2(runtime);
+  return /* @__PURE__ */ (0, import_jsx_runtime23.jsx)(AssistantContext.Provider, { value: context, children });
+};
+
+// src/adapters/core/local/useLocalRuntime.tsx
+var import_react29 = require("react");
+
+// src/adapters/core/local/LocalRuntime.tsx
+var LocalRuntime = class {
+  constructor(adapter) {
+    this.adapter = adapter;
+  }
+  _messageUpdateCallbacks = /* @__PURE__ */ new Set();
+  _statusUpdateCallbacks = /* @__PURE__ */ new Set();
+  abortController = null;
+  repository = new MessageRepository();
+  async append(message) {
+    const userMessageId = generateId();
+    const userMessage = {
+      id: userMessageId,
+      role: "user",
+      content: message.content,
+      createdAt: /* @__PURE__ */ new Date()
+    };
+    this.addOrUpdateMessage(message.parentId, userMessage);
+    const { id } = await this.startRun(userMessageId);
+    return { parentId: userMessageId, id };
+  }
+  async startRun(parentId) {
+    const id = generateId();
+    this.repository.resetHead(parentId);
+    const messages = this.repository.getMessages();
+    const message = {
+      id,
+      role: "assistant",
+      status: "in_progress",
+      content: [{ type: "text", text: "" }],
+      createdAt: /* @__PURE__ */ new Date()
+    };
+    this.addOrUpdateMessage(parentId, message);
+    void this.run(parentId, messages, message);
+    return { id };
+  }
+  addOrUpdateMessage(parentId, message) {
+    const clone = { ...message };
+    this.repository.addOrUpdateMessage(parentId, clone);
+    for (const callback of this._messageUpdateCallbacks)
+      callback(parentId, clone);
+  }
+  async run(parentId, messages, message) {
+    this.cancelRun();
+    for (const callback of this._statusUpdateCallbacks) callback(true);
+    this.abortController = new AbortController();
+    try {
+      await this.adapter.run({
+        messages,
+        abortSignal: this.abortController.signal,
+        onUpdate: ({ content }) => {
+          message.content = content;
+          this.addOrUpdateMessage(parentId, message);
+        }
+      });
+      message.status = "done";
+      this.addOrUpdateMessage(parentId, message);
+    } catch (e) {
+      message.status = "error";
+      this.addOrUpdateMessage(parentId, message);
+      console.error(e);
+    } finally {
+      this.cancelRun();
+    }
+  }
+  cancelRun() {
+    if (!this.abortController) return;
+    this.abortController.abort();
+    this.abortController = null;
+    for (const callback of this._statusUpdateCallbacks) callback(false);
+  }
+  subscribeToMessageUpdates(callback) {
+    this._messageUpdateCallbacks.add(callback);
+    return () => this._messageUpdateCallbacks.delete(callback);
+  }
+  subscribeToStatusUpdates(callback) {
+    this._statusUpdateCallbacks.add(callback);
+    return () => this._statusUpdateCallbacks.delete(callback);
+  }
+};
+
+// src/adapters/core/local/useLocalRuntime.tsx
+var useLocalRuntime = (adapter) => {
+  const [runtime] = (0, import_react29.useState)(() => new LocalRuntime(adapter));
+  runtime.adapter = adapter;
+  return runtime;
+};
+
+// src/adapters/core/local/vercel/VercelModelAdapter.tsx
+var import_ai = require("ai");
+var VercelModelAdapter = class {
+  constructor(model) {
+    this.model = model;
+  }
+  async run({ messages, abortSignal, onUpdate }) {
+    const { fullStream } = await (0, import_ai.streamText)({
+      model: this.model,
+      abortSignal,
+      messages: messages.map((m) => ({
+        role: m.role,
+        content: m.content.filter((c) => c.type !== "ui")
+      }))
+    });
+    const content = [];
+    for await (const aiPart of fullStream) {
+      switch (aiPart.type) {
+        case "text-delta": {
+          let part = content.at(-1);
+          if (!part || part.type !== "text") {
+            part = { type: "text", text: "" };
+            content.push(part);
+          }
+          part.text += aiPart.textDelta;
+          break;
+        }
+        case "tool-call": {
+          content.push({
+            type: "tool-call",
+            name: aiPart.toolName,
+            args: aiPart.args
+          });
+          break;
+        }
+      }
+      onUpdate({ content });
+    }
+    return { content };
+  }
+};
 // Annotate the CommonJS export names for ESM import in node:
 0 && (module.exports = {
   ActionBarPrimitive,
@@ -1549,8 +1804,11 @@ var VercelRSCAssistantProvider = ({
   ThreadPrimitive,
   VercelAIAssistantProvider,
   VercelRSCAssistantProvider,
+  unstable_AssistantProvider,
+  unstable_VercelModelAdapter,
   unstable_getVercelMessage,
   unstable_getVercelRSCMessage,
+  unstable_useLocalRuntime,
   unstable_useMessageContext,
   useBeginMessageEdit,
   useCopyMessage,

package/dist/index.mjs

@@ -1076,7 +1076,7 @@ var useDummyAIAssistantContext = () => {
 import { useCallbackRef as useCallbackRef3 } from "@radix-ui/react-use-callback-ref";
 import { useCallback as useCallback2, useMemo as useMemo5, useRef as useRef4, useState as useState4 } from "react";
 
-// src/adapters/MessageRepository.tsx
+// src/adapters/idUtils.tsx
 import { customAlphabet } from "nanoid/non-secure";
 var generateId = customAlphabet(
   "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz",
@@ -1084,6 +1084,8 @@ var generateId = customAlphabet(
 );
 var optimisticPrefix = "__optimistic__";
 var generateOptimisticId = () => `${optimisticPrefix}${generateId()}`;
+
+// src/adapters/MessageRepository.tsx
 var findHead = (message) => {
   if (message.next) return findHead(message.next);
   return message;
@@ -1095,15 +1097,6 @@ var MessageRepository = class {
   root = {
     children: []
   };
-  getFallbackChild(p) {
-    const childId = p.children.at(-1);
-    const child = childId ? this.messages.get(childId) : null;
-    if (child === void 0)
-      throw new Error(
-        "MessageRepository(getFallbackChild): Child message not found. This is likely an internal bug in assistant-ui."
-      );
-    return child;
-  }
   performOp(newParent, child, operation) {
     const parentOrRoot = child.prev ?? this.root;
     const newParentOrRoot = newParent ?? this.root;
@@ -1113,7 +1106,14 @@ var MessageRepository = class {
         (m) => m !== child.current.id
       );
       if (child.prev?.next === child) {
-        child.prev.next = this.getFallbackChild(child.prev);
+        const fallbackId = child.prev.children.at(-1);
+        const fallback = fallbackId ? this.messages.get(fallbackId) : null;
+        if (fallback === void 0) {
+          throw new Error(
+            "MessageRepository(performOp/cut): Fallback sibling message not found. This is likely an internal bug in assistant-ui."
+          );
+        }
+        child.prev.next = fallback;
       }
     }
     if (operation !== "cut") {
@@ -1172,14 +1172,14 @@ var MessageRepository = class {
     });
     return optimisticId;
   }
-  deleteMessage(messageId, newParentId) {
+  deleteMessage(messageId, replacementId) {
     const message = this.messages.get(messageId);
-    const newParent = newParentId ? this.messages.get(newParentId) : null;
+    const replacement = replacementId ? this.messages.get(replacementId) : null;
     if (!message)
       throw new Error(
         "MessageRepository(deleteMessage): Optimistic message not found. This is likely an internal bug in assistant-ui."
       );
-    if (newParent === void 0)
+    if (replacement === void 0)
       throw new Error(
         "MessageRepository(deleteMessage): New message not found. This is likely an internal bug in assistant-ui."
       );
@@ -1189,11 +1189,11 @@ var MessageRepository = class {
         throw new Error(
           "MessageRepository(deleteMessage): Child message not found. This is likely an internal bug in assistant-ui."
         );
-      this.performOp(newParent, childMessage, "relink");
+      this.performOp(replacement, childMessage, "relink");
     }
     this.messages.delete(messageId);
     if (this.head === message) {
-      this.head = this.getFallbackChild(message.prev ?? this.root);
+      this.head = replacement;
     }
     this.performOp(null, message, "cut");
   }
@@ -1236,17 +1236,13 @@ var MessageRepository = class {
   }
 };
 
-// src/adapters/ThreadMessageConverter.tsx
+// src/adapters/ThreadMessageConverter.ts
 var ThreadMessageConverter = class {
-  constructor(converter) {
-    this.converter = converter;
-  }
   cache = /* @__PURE__ */ new WeakMap();
-  convertMessages(messages) {
+  convertMessages(converter, messages) {
     return messages.map((m) => {
       const cached = this.cache.get(m);
-      if (cached) return cached;
-      const newMessage = this.converter(m);
+      const newMessage = converter(m, cached);
       this.cache.set(m, newMessage);
       return newMessage;
     });
@@ -1319,16 +1315,17 @@ var getIsRunning = (vercel) => {
 var useVercelAIThreadState = (vercel) => {
   const [data] = useState4(() => new MessageRepository());
   const isRunning = getIsRunning(vercel);
-  const convertCallback = useCallbackRef3((message) => {
-    return vercelToThreadMessage(
-      message,
-      vercel.messages.at(-1) === message && isRunning ? "in_progress" : "done"
-    );
-  });
-  const converter = new ThreadMessageConverter(convertCallback);
+  const converter = useMemo5(() => new ThreadMessageConverter(), []);
   const assistantOptimisticIdRef = useRef4(null);
   const messages = useMemo5(() => {
-    const vm = converter.convertMessages(vercel.messages);
+    const lastMessageId = vercel.messages.at(-1)?.id;
+    const convertCallback = (message, cache) => {
+      const status = lastMessageId === message.id && isRunning ? "in_progress" : "done";
+      if (cache && (cache.role === "user" || cache.status === status))
+        return cache;
+      return vercelToThreadMessage(message, status);
+    };
+    const vm = converter.convertMessages(convertCallback, vercel.messages);
     for (let i = 0; i < vm.length; i++) {
       const message = vm[i];
       const parent = vm[i - 1];
@@ -1480,15 +1477,17 @@ var VercelRSCAssistantProvider = ({
     setIsRunning(true);
     return callback.finally(() => setIsRunning(false));
   }, []);
-  const converter = useMemo7(() => {
+  const [converter, convertCallback] = useMemo7(() => {
     const rscConverter = convertMessage ?? ((m) => m);
-    return new ThreadMessageConverter((m) => {
+    const convertCallback2 = (m, cache) => {
+      if (cache) return cache;
       return vercelToThreadMessage2(rscConverter, m);
-    });
+    };
+    return [new ThreadMessageConverter(), convertCallback2];
   }, [convertMessage]);
   const messages = useMemo7(() => {
-    return converter.convertMessages(vercelMessages);
-  }, [converter, vercelMessages]);
+    return converter.convertMessages(convertCallback, vercelMessages);
+  }, [converter, convertCallback, vercelMessages]);
   const append = useCallback3(
     async (message) => {
       if (message.parentId !== (context.useThread.getState().messages.at(-1)?.id ?? null)) {
@@ -1529,6 +1528,259 @@ var VercelRSCAssistantProvider = ({
   }, [context, messages, isRunning, append, startRun]);
   return /* @__PURE__ */ jsx22(AssistantContext.Provider, { value: context, children });
 };
+
+// src/adapters/core/utils/useAssistantContext.tsx
+import { useEffect as useEffect4, useRef as useRef5, useState as useState6 } from "react";
+import { create as create6 } from "zustand";
+
+// src/adapters/core/utils/AssistantMessageRepository.tsx
+var AssistantMessageRepository = class {
+  constructor(flushCallback) {
+    this.flushCallback = flushCallback;
+  }
+  repository = new MessageRepository();
+  getBranches(messageId) {
+    return this.repository.getBranches(messageId);
+  }
+  withModifications(callback) {
+    const res = callback(this.repository);
+    this.flushCallback(this.repository.getMessages());
+    return res;
+  }
+};
+
+// src/adapters/core/utils/useAssistantContext.tsx
+var makeThreadStore = (runtimeRef) => {
+  const repository = new AssistantMessageRepository((messages) => {
+    useThread.setState({ messages });
+  });
+  const useThread = create6(() => ({
+    messages: [],
+    isRunning: false,
+    getBranches: (messageId) => repository.getBranches(messageId),
+    switchToBranch: (branchId) => {
+      repository.withModifications((repository2) => {
+        repository2.switchToBranch(branchId);
+      });
+    },
+    startRun: async (parentId) => {
+      const optimisticId = repository.withModifications((repository2) => {
+        const optimisticId2 = repository2.appendOptimisticMessage(parentId, {
+          role: "assistant",
+          content: [{ type: "text", text: "" }]
+        });
+        repository2.resetHead(optimisticId2);
+        return optimisticId2;
+      });
+      const { id } = await runtimeRef.current.startRun(parentId);
+      repository.withModifications((repository2) => {
+        repository2.deleteMessage(optimisticId, id);
+      });
+    },
+    append: async (message) => {
+      const [parentOptimisticId, optimisticId] = repository.withModifications(
+        (repository2) => {
+          const parentOptimisticId2 = repository2.appendOptimisticMessage(
+            message.parentId,
+            {
+              role: "user",
+              content: message.content
+            }
+          );
+          const optimisticId2 = repository2.appendOptimisticMessage(
+            parentOptimisticId2,
+            {
+              role: "assistant",
+              content: [{ type: "text", text: "" }]
+            }
+          );
+          repository2.resetHead(optimisticId2);
+          return [parentOptimisticId2, optimisticId2];
+        }
+      );
+      const { parentId, id } = await runtimeRef.current.append(message);
+      repository.withModifications((repository2) => {
+        repository2.deleteMessage(parentOptimisticId, parentId);
+        repository2.deleteMessage(optimisticId, id);
+      });
+    },
+    cancelRun: () => runtimeRef.current.cancelRun()
+  }));
+  const onNewMessage = (parentId, message) => {
+    repository.withModifications((repository2) => {
+      repository2.addOrUpdateMessage(parentId, message);
+    });
+  };
+  const onRunningChange = (isRunning) => {
+    useThread.setState({ isRunning });
+  };
+  return {
+    useThread,
+    onNewMessage,
+    onRunningChange
+  };
+};
+var useAssistantContext2 = (runtime) => {
+  const runtimeRef = useRef5(runtime);
+  runtimeRef.current = runtime;
+  const [{ context, onNewMessage, onRunningChange }] = useState6(() => {
+    const { useThread, onNewMessage: onNewMessage2, onRunningChange: onRunningChange2 } = makeThreadStore(runtimeRef);
+    const useViewport = makeViewportStore();
+    const useComposer = makeThreadComposerStore(useThread);
+    return {
+      context: { useViewport, useThread, useComposer },
+      onNewMessage: onNewMessage2,
+      onRunningChange: onRunningChange2
+    };
+  });
+  useEffect4(() => {
+    return runtime.subscribeToMessageUpdates(onNewMessage);
+  }, [runtime, onNewMessage]);
+  useEffect4(() => {
+    return runtime.subscribeToStatusUpdates(onRunningChange);
+  }, [runtime, onRunningChange]);
+  return context;
+};
+
+// src/adapters/core/AssistantProvider.tsx
+import { jsx as jsx23 } from "react/jsx-runtime";
+var AssistantProvider = ({ children, runtime }) => {
+  const context = useAssistantContext2(runtime);
+  return /* @__PURE__ */ jsx23(AssistantContext.Provider, { value: context, children });
+};
+
+// src/adapters/core/local/useLocalRuntime.tsx
+import { useState as useState7 } from "react";
+
+// src/adapters/core/local/LocalRuntime.tsx
+var LocalRuntime = class {
+  constructor(adapter) {
+    this.adapter = adapter;
+  }
+  _messageUpdateCallbacks = /* @__PURE__ */ new Set();
+  _statusUpdateCallbacks = /* @__PURE__ */ new Set();
+  abortController = null;
+  repository = new MessageRepository();
+  async append(message) {
+    const userMessageId = generateId();
+    const userMessage = {
+      id: userMessageId,
+      role: "user",
+      content: message.content,
+      createdAt: /* @__PURE__ */ new Date()
+    };
+    this.addOrUpdateMessage(message.parentId, userMessage);
+    const { id } = await this.startRun(userMessageId);
+    return { parentId: userMessageId, id };
+  }
+  async startRun(parentId) {
+    const id = generateId();
+    this.repository.resetHead(parentId);
+    const messages = this.repository.getMessages();
+    const message = {
+      id,
+      role: "assistant",
+      status: "in_progress",
+      content: [{ type: "text", text: "" }],
+      createdAt: /* @__PURE__ */ new Date()
+    };
+    this.addOrUpdateMessage(parentId, message);
+    void this.run(parentId, messages, message);
+    return { id };
+  }
+  addOrUpdateMessage(parentId, message) {
+    const clone = { ...message };
+    this.repository.addOrUpdateMessage(parentId, clone);
+    for (const callback of this._messageUpdateCallbacks)
+      callback(parentId, clone);
+  }
+  async run(parentId, messages, message) {
+    this.cancelRun();
+    for (const callback of this._statusUpdateCallbacks) callback(true);
+    this.abortController = new AbortController();
+    try {
+      await this.adapter.run({
+        messages,
+        abortSignal: this.abortController.signal,
+        onUpdate: ({ content }) => {
+          message.content = content;
+          this.addOrUpdateMessage(parentId, message);
+        }
+      });
+      message.status = "done";
+      this.addOrUpdateMessage(parentId, message);
+    } catch (e) {
+      message.status = "error";
+      this.addOrUpdateMessage(parentId, message);
+      console.error(e);
+    } finally {
+      this.cancelRun();
+    }
+  }
+  cancelRun() {
+    if (!this.abortController) return;
+    this.abortController.abort();
+    this.abortController = null;
+    for (const callback of this._statusUpdateCallbacks) callback(false);
+  }
+  subscribeToMessageUpdates(callback) {
+    this._messageUpdateCallbacks.add(callback);
+    return () => this._messageUpdateCallbacks.delete(callback);
+  }
+  subscribeToStatusUpdates(callback) {
+    this._statusUpdateCallbacks.add(callback);
+    return () => this._statusUpdateCallbacks.delete(callback);
+  }
+};
+
+// src/adapters/core/local/useLocalRuntime.tsx
+var useLocalRuntime = (adapter) => {
+  const [runtime] = useState7(() => new LocalRuntime(adapter));
+  runtime.adapter = adapter;
+  return runtime;
+};
+
+// src/adapters/core/local/vercel/VercelModelAdapter.tsx
+import { streamText } from "ai";
+var VercelModelAdapter = class {
+  constructor(model) {
+    this.model = model;
+  }
+  async run({ messages, abortSignal, onUpdate }) {
+    const { fullStream } = await streamText({
+      model: this.model,
+      abortSignal,
+      messages: messages.map((m) => ({
+        role: m.role,
+        content: m.content.filter((c) => c.type !== "ui")
+      }))
+    });
+    const content = [];
+    for await (const aiPart of fullStream) {
+      switch (aiPart.type) {
+        case "text-delta": {
+          let part = content.at(-1);
+          if (!part || part.type !== "text") {
+            part = { type: "text", text: "" };
+            content.push(part);
+          }
+          part.text += aiPart.textDelta;
+          break;
+        }
+        case "tool-call": {
+          content.push({
+            type: "tool-call",
+            name: aiPart.toolName,
+            args: aiPart.args
+          });
+          break;
+        }
+      }
+      onUpdate({ content });
+    }
+    return { content };
+  }
+};
 export {
   actionBar_exports as ActionBarPrimitive,
   branchPicker_exports as BranchPickerPrimitive,
@@ -1538,8 +1790,11 @@ export {
   thread_exports as ThreadPrimitive,
   VercelAIAssistantProvider,
   VercelRSCAssistantProvider,
+  AssistantProvider as unstable_AssistantProvider,
+  VercelModelAdapter as unstable_VercelModelAdapter,
   getVercelMessage as unstable_getVercelMessage,
   getVercelRSCMessage as unstable_getVercelRSCMessage,
+  useLocalRuntime as unstable_useLocalRuntime,
   useMessageContext as unstable_useMessageContext,
   useBeginMessageEdit,
   useCopyMessage,

package/package.json

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