@assistant-ui/react 0.11.57 → 0.12.0

package/src/legacy-runtime/hooks/ThreadListItemContext.ts

@@ -2,10 +2,10 @@
 
 import { ThreadListItemRuntime } from "../runtime/ThreadListItemRuntime";
 import { createStateHookForRuntime } from "../../context/react/utils/createStateHookForRuntime";
-import { useAssistantApi, useAssistantState } from "../../context/react";
+import { useAui, useAuiState } from "@assistant-ui/store";
 
 /**
- * @deprecated Use `useAssistantApi()` with `api.threadListItem()` instead. See migration guide: https://assistant-ui.com/docs/migrations/v0-12
+ * @deprecated Use `useAui()` with `aui.threadListItem()` instead. See migration guide: https://assistant-ui.com/docs/migrations/v0-12
  */
 export function useThreadListItemRuntime(options?: {
   optional?: false | undefined;
@@ -16,10 +16,10 @@ export function useThreadListItemRuntime(options?: {
 export function useThreadListItemRuntime(options?: {
   optional?: boolean | undefined;
 }) {
-  const api = useAssistantApi();
-  const runtime = useAssistantState(() =>
-    api.threadListItem.source
-      ? (api.threadListItem().__internal_getRuntime?.() ?? null)
+  const aui = useAui();
+  const runtime = useAuiState(() =>
+    aui.threadListItem.source
+      ? (aui.threadListItem().__internal_getRuntime?.() ?? null)
       : null,
   );
   if (!runtime && !options?.optional) {
@@ -29,7 +29,7 @@ export function useThreadListItemRuntime(options?: {
 }
 
 /**
- * @deprecated Use `useAssistantState(({ threadListItem }) => threadListItem)` instead. See migration guide: https://assistant-ui.com/docs/migrations/v0-12
+ * @deprecated Use `useAuiState(({ threadListItem }) => threadListItem)` instead. See migration guide: https://assistant-ui.com/docs/migrations/v0-12
  */
 export const useThreadListItem = createStateHookForRuntime(
   useThreadListItemRuntime,

package/src/legacy-runtime/runtime/ThreadListItemRuntime.ts

@@ -3,7 +3,7 @@ import { ThreadListItemRuntimePath } from "./RuntimePathTypes";
 import { SubscribableWithState } from "./subscribable/Subscribable";
 import { ThreadListRuntimeCoreBinding } from "./ThreadListRuntime";
 
-export type ThreadListItemEventType = "switched-to" | "switched-away";
+export type ThreadListItemEventType = "switchedTo" | "switchedAway";
 
 import type {
   ThreadListItemState,
@@ -120,7 +120,7 @@ export class ThreadListItemRuntimeImpl implements ThreadListItemRuntime {
 
   public unstable_on(event: ThreadListItemEventType, callback: () => void) {
     // if the runtime is bound to a specific thread, trigger if isMain is toggled
-    // if the runtime is bound to the main thread, trigger switched-to if threadId changes
+    // if the runtime is bound to the main thread, trigger switchedTo if threadId changes
 
     let prevIsMain = this._core.getState().isMain;
     let prevThreadId = this._core.getState().id;
@@ -132,8 +132,8 @@ export class ThreadListItemRuntimeImpl implements ThreadListItemRuntime {
       prevIsMain = newIsMain;
       prevThreadId = newThreadId;
 
-      if (event === "switched-to" && !newIsMain) return;
-      if (event === "switched-away" && newIsMain) return;
+      if (event === "switchedTo" && !newIsMain) return;
+      if (event === "switchedAway" && newIsMain) return;
       callback();
     });
   }

package/src/legacy-runtime/runtime-cores/assistant-transport/useAssistantTransportRuntime.tsx

@@ -31,7 +31,7 @@ import { ToolExecutionStatus, useToolInvocations } from "./useToolInvocations";
 import { toAISDKTools, getEnabledTools, createRequestHeaders } from "./utils";
 import { useRemoteThreadListRuntime } from "../remote-thread-list/useRemoteThreadListRuntime";
 import { InMemoryThreadListAdapter } from "../remote-thread-list/adapter/in-memory";
-import { useAssistantApi, useAssistantState } from "../../../context/react";
+import { useAui, useAuiState } from "@assistant-ui/store";
 import { UserExternalState } from "../../../augmentations";
 
 const symbolAssistantTransportExtras = Symbol("assistant-transport-extras");
@@ -57,10 +57,10 @@ const asAssistantTransportExtras = (
 };
 
 export const useAssistantTransportSendCommand = () => {
-  const api = useAssistantApi();
+  const aui = useAui();
 
   return (command: AssistantTransportCommand) => {
-    const extras = api.thread().getState().extras;
+    const extras = aui.thread().getState().extras;
     const transportExtras = asAssistantTransportExtras(extras);
     transportExtras.sendCommand(command);
   };
@@ -73,7 +73,7 @@ export function useAssistantTransportState<T>(
 export function useAssistantTransportState<T>(
   selector: (state: UserExternalState) => T = (t) => t as T,
 ): T | UserExternalState {
-  return useAssistantState(({ thread }) =>
+  return useAuiState(({ thread }) =>
     selector(asAssistantTransportExtras(thread.extras).state),
   );
 }
@@ -89,9 +89,7 @@ const useAssistantTransportThreadRuntime = <T,>(
     onQueue: () => runManager.schedule(),
   });
 
-  const threadId = useAssistantState(
-    ({ threadListItem }) => threadListItem.remoteId,
-  );
+  const threadId = useAuiState(({ threadListItem }) => threadListItem.remoteId);
 
   const runManager = useRunManager({
     onRun: async (signal: AbortSignal) => {

package/src/legacy-runtime/runtime-cores/assistant-transport/useToolInvocations.ts

@@ -192,6 +192,22 @@ export function useToolInvocations({
                   }
                 } else {
                   if (!content.argsText.startsWith(lastState.argsText)) {
+                    // Check if this is key reordering (both are complete JSON)
+                    // This happens when transitioning from streaming to complete state
+                    // and the provider returns keys in a different order
+                    if (
+                      isArgsTextComplete(lastState.argsText) &&
+                      isArgsTextComplete(content.argsText)
+                    ) {
+                      lastState.controller.argsText.close();
+                      lastToolStates.current[content.toolCallId] = {
+                        argsText: content.argsText,
+                        hasResult: lastState.hasResult,
+                        argsComplete: true,
+                        controller: lastState.controller,
+                      };
+                      return; // Continue to next content part
+                    }
                     throw new Error(
                       `Tool call argsText can only be appended, not updated: ${content.argsText} does not start with ${lastState.argsText}`,
                     );
@@ -249,7 +265,7 @@ export function useToolInvocations({
     if (isInitialState.current) {
       isInitialState.current = false;
     }
-  }, [state, controller, onResult]);
+  }, [state, controller]);
 
   const abort = (): Promise<void> => {
     humanInputRef.current.forEach(({ reject }) => {

package/src/legacy-runtime/runtime-cores/composer/BaseComposerRuntimeCore.tsx

@@ -204,7 +204,7 @@ export abstract class BaseComposerRuntimeCore
       upsertAttachment(await promiseOrGenerator);
     }
 
-    this._notifyEventSubscribers("attachment-add");
+    this._notifyEventSubscribers("attachmentAdd");
     this._notifySubscribers();
   }
 

package/src/legacy-runtime/runtime-cores/core/BaseThreadRuntimeCore.tsx

@@ -216,7 +216,7 @@ export abstract class BaseThreadRuntimeCore implements ThreadRuntimeCore {
   >();
 
   public unstable_on(event: ThreadRuntimeEventType, callback: () => void) {
-    if (event === "model-context-update") {
+    if (event === "modelContextUpdate") {
       return this._contextProvider.subscribe?.(callback) ?? (() => {});
     }
 

package/src/legacy-runtime/runtime-cores/core/ComposerRuntimeCore.tsx

@@ -6,7 +6,7 @@ import type {
 import type { MessageRole, RunConfig } from "../../../types/AssistantTypes";
 import type { DictationAdapter } from "../adapters/speech/SpeechAdapterTypes";
 
-export type ComposerRuntimeEventType = "send" | "attachment-add";
+export type ComposerRuntimeEventType = "send" | "attachmentAdd";
 
 /**
  * State representing an active dictation session.

package/src/legacy-runtime/runtime-cores/core/ThreadRuntimeCore.tsx

@@ -58,10 +58,10 @@ export type SubmittedFeedback = {
 };
 
 export type ThreadRuntimeEventType =
-  | "run-start"
-  | "run-end"
+  | "runStart"
+  | "runEnd"
   | "initialize"
-  | "model-context-update";
+  | "modelContextUpdate";
 
 export type StartRunConfig = {
   parentId: string | null;

package/src/legacy-runtime/runtime-cores/external-store/ExternalStoreThreadRuntimeCore.tsx

@@ -209,9 +209,9 @@ export class ExternalStoreThreadRuntimeCore
 
     if ((oldStore?.isRunning ?? false) !== (store.isRunning ?? false)) {
       if (store.isRunning) {
-        this._notifyEventSubscribers("run-start");
+        this._notifyEventSubscribers("runStart");
       } else {
-        this._notifyEventSubscribers("run-end");
+        this._notifyEventSubscribers("runEnd");
       }
     }
 

package/src/legacy-runtime/runtime-cores/external-store/ThreadMessageLike.tsx

@@ -85,12 +85,10 @@ export const fromThreadMessageLike = (
     image,
     ...rest
   }: ImageMessagePart): ImageMessagePart | null => {
-    // Accept base64 data URIs
-    const isBase64 = /^data:image\/(png|jpeg|jpg|gif|webp);base64,/.test(image);
-    // Accept HTTP/HTTPS URLs (e.g., from cloud storage)
-    const isUrl = /^https?:\/\//.test(image);
-
-    if (isBase64 || isUrl) {
+    const match = image.match(
+      /^data:image\/(png|jpeg|jpg|gif|webp);base64,(.*)$/,
+    );
+    if (match) {
       return { ...rest, image };
     }
     console.warn(`Invalid image data format detected`);

package/src/legacy-runtime/runtime-cores/external-store/createMessageConverter.tsx

@@ -1,6 +1,6 @@
 "use client";
 import { ThreadState } from "../../runtime";
-import { useAssistantApi, useAssistantState } from "../../../context";
+import { useAui, useAuiState } from "@assistant-ui/store";
 import { ThreadMessage } from "../../../types";
 import {
   useExternalMessageConverter,
@@ -57,13 +57,13 @@ export const createMessageConverter = <T extends object>(
       return first;
     },
     useOriginalMessages: () => {
-      const api = useAssistantApi();
-      const partMessages = useAssistantState((s) => {
-        if (api.part.source) return getExternalStoreMessages(s.part);
+      const aui = useAui();
+      const partMessages = useAuiState((s) => {
+        if (aui.part.source) return getExternalStoreMessages(s.part);
         return undefined;
       });
 
-      const messageMessages = useAssistantState<T[]>(({ message }) =>
+      const messageMessages = useAuiState<T[]>(({ message }) =>
         getExternalStoreMessages(message),
       );
 

package/src/legacy-runtime/runtime-cores/external-store/external-message-converter.tsx

@@ -185,7 +185,7 @@ const joinExternalMessages = (
                 ] as typeof part;
                 assistantMessage.content[existingIdx] = {
                   ...existing,
-                  text: existing.text + "\n\n" + part.text,
+                  text: `${existing.text}\n\n${part.text}`,
                   ...{
                     [symbolInnerMessage]: [
                       ...((existing as any)[symbolInnerMessage] ?? []),

package/src/legacy-runtime/runtime-cores/local/LocalThreadRuntimeCore.tsx

@@ -230,7 +230,7 @@ export class LocalThreadRuntimeCore
       createdAt: new Date(),
     };
 
-    this._notifyEventSubscribers("run-start");
+    this._notifyEventSubscribers("runStart");
 
     try {
       this._suggestions = [];
@@ -248,7 +248,7 @@ export class LocalThreadRuntimeCore
         runCallback = undefined;
       } while (shouldContinue(message, this._options.unstable_humanToolNames));
     } finally {
-      this._notifyEventSubscribers("run-end");
+      this._notifyEventSubscribers("runEnd");
     }
 
     this._suggestionsController = new AbortController();
@@ -412,6 +412,10 @@ export class LocalThreadRuntimeCore
         updateMessage({
           status: { type: "incomplete", reason: "cancelled" },
         });
+      } else if (e instanceof Error && e.name === "AbortError") {
+        updateMessage({
+          status: { type: "incomplete", reason: "cancelled" },
+        });
       } else {
         updateMessage({
           status: {

package/src/legacy-runtime/runtime-cores/local/useLocalRuntime.tsx

@@ -8,29 +8,26 @@ import { useRuntimeAdapters } from "../adapters/RuntimeAdapterProvider";
 import { useRemoteThreadListRuntime } from "../remote-thread-list/useRemoteThreadListRuntime";
 import { useCloudThreadListAdapter } from "../remote-thread-list/adapter/cloud";
 import { AssistantRuntimeImpl } from "../../../internal";
-import { useAssistantState } from "../../../context/react";
+import { useAuiState } from "@assistant-ui/store";
 
 const useLocalThreadRuntime = (
   adapter: ChatModelAdapter,
   { initialMessages, ...options }: LocalRuntimeOptions,
 ) => {
   const { modelContext, ...threadListAdapters } = useRuntimeAdapters() ?? {};
-  const opt = useMemo(
-    () => ({
-      ...options,
-      adapters: {
-        ...threadListAdapters,
-        ...options.adapters,
-        chatModel: adapter,
-      },
-    }),
-    [adapter, options, threadListAdapters],
-  );
+  const opt = {
+    ...options,
+    adapters: {
+      ...threadListAdapters,
+      ...options.adapters,
+      chatModel: adapter,
+    },
+  };
 
   const [runtime] = useState(() => new LocalRuntimeCore(opt, initialMessages));
 
   const threadIdRef = useRef<string | undefined>(undefined);
-  threadIdRef.current = useAssistantState(
+  threadIdRef.current = useAuiState(
     ({ threadListItem }) => threadListItem.remoteId,
   );
 
@@ -49,7 +46,7 @@ const useLocalThreadRuntime = (
   useEffect(() => {
     runtime.threads.getMainThreadRuntimeCore().__internal_setOptions(opt);
     runtime.threads.getMainThreadRuntimeCore().__internal_load();
-  }, [runtime, opt]);
+  });
 
   useEffect(() => {
     if (!modelContext) return undefined;

package/src/legacy-runtime/runtime-cores/remote-thread-list/RemoteThreadListHookInstanceManager.tsx

@@ -10,7 +10,8 @@ import {
   ComponentType,
 } from "react";
 import { UseBoundStore, StoreApi, create } from "zustand";
-import { useAssistantApi, ThreadListItemByIdProvider } from "../../../context";
+import { useAui } from "@assistant-ui/store";
+import { ThreadListItemByIdProvider } from "../../../context/providers";
 import { ThreadRuntimeCore, ThreadRuntimeImpl } from "../../../internal";
 import { BaseSubscribable } from "./BaseSubscribable";
 import { AssistantRuntime } from "../../runtime";
@@ -103,7 +104,7 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
       return threadBinding.outerSubscribe(updateRuntime);
     }, [threadBinding, updateRuntime]);
 
-    const api = useAssistantApi();
+    const aui = useAui();
     const initPromiseRef = useRef<Promise<unknown> | undefined>(undefined);
 
     useEffect(() => {
@@ -112,24 +113,24 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
         "__internal_setGetInitializePromise"
       ];
       if (typeof setGetInitializePromise === "function") {
-        setGetInitializePromise(() => initPromiseRef.current);
+        setGetInitializePromise.call(runtimeCore, () => initPromiseRef.current);
       }
     }, [threadBinding]);
 
     useEffect(() => {
       return runtime.threads.main.unstable_on("initialize", () => {
-        const state = api.threadListItem().getState();
+        const state = aui.threadListItem().getState();
         if (state.status === "new") {
-          initPromiseRef.current = api.threadListItem().initialize();
+          initPromiseRef.current = aui.threadListItem().initialize();
 
-          const dispose = runtime.thread.unstable_on("run-end", () => {
+          const dispose = runtime.thread.unstable_on("runEnd", () => {
             dispose();
 
-            api.threadListItem().generateTitle();
+            aui.threadListItem().generateTitle();
           });
         }
       });
-    }, [runtime, api]);
+    }, [runtime, aui]);
 
     return null;
   };

package/src/legacy-runtime/runtime-cores/remote-thread-list/useRemoteThreadListRuntime.tsx

@@ -7,7 +7,7 @@ import { RemoteThreadListOptions } from "./types";
 import { AssistantRuntimeImpl } from "../../../internal";
 import { AssistantRuntimeCore } from "../core/AssistantRuntimeCore";
 import { AssistantRuntime } from "../../runtime/AssistantRuntime";
-import { useAssistantApiImpl } from "../../../context/react/AssistantApiContext";
+import { useAui } from "@assistant-ui/store";
 
 class RemoteThreadListRuntimeCore
   extends BaseAssistantRuntimeCore
@@ -42,8 +42,8 @@ const useRemoteThreadListRuntimeImpl = (
 export const useRemoteThreadListRuntime = (
   options: RemoteThreadListOptions,
 ): AssistantRuntime => {
-  const api = useAssistantApiImpl();
-  const isNested = api.threadListItem.source !== null;
+  const aui = useAui();
+  const isNested = aui.threadListItem.source !== null;
 
   if (isNested) {
     if (!options.allowNesting) {

package/src/model-context/makeAssistantVisible.tsx

@@ -11,7 +11,7 @@ import {
   createContext,
   useContext,
 } from "react";
-import { useAssistantApi } from "../context/react/AssistantApiContext";
+import { useAui } from "@assistant-ui/store";
 import { useComposedRefs } from "@radix-ui/react-compose-refs";
 import { tool } from "./tool";
 
@@ -83,11 +83,11 @@ export const makeAssistantVisible = <T extends ComponentType<any>>(
       const clickId = useId();
       const componentRef = useRef<HTMLElement>(null);
 
-      const api = useAssistantApi();
+      const aui = useAui();
 
       const { clickable, editable } = config ?? {};
       useEffect(() => {
-        return api.modelContext().register({
+        return aui.modelContext().register({
           getModelContext: () => {
             return {
               tools: {
@@ -100,7 +100,7 @@ export const makeAssistantVisible = <T extends ComponentType<any>>(
             };
           },
         });
-      }, [isNestedReadable, api, clickable, editable]);
+      }, [isNestedReadable, aui, clickable, editable]);
 
       const ref = useComposedRefs(componentRef, outerRef);
 

package/src/model-context/useAssistantInstructions.tsx

@@ -1,7 +1,7 @@
 "use client";
 
 import { useEffect } from "react";
-import { useAssistantApi } from "../context/react/AssistantApiContext";
+import { useAui } from "@assistant-ui/store";
 
 export type AssistantInstructionsConfig = {
   disabled?: boolean | undefined;
@@ -19,7 +19,7 @@ export const useAssistantInstructions = (
   config: string | AssistantInstructionsConfig,
 ) => {
   const { instruction, disabled = false } = getInstructions(config);
-  const api = useAssistantApi();
+  const aui = useAui();
 
   useEffect(() => {
     if (disabled) return;
@@ -27,8 +27,8 @@ export const useAssistantInstructions = (
     const config = {
       system: instruction,
     };
-    return api.modelContext().register({
+    return aui.modelContext().register({
       getModelContext: () => config,
     });
-  }, [api, instruction, disabled]);
+  }, [aui, instruction, disabled]);
 };

package/src/model-context/useAssistantTool.tsx

@@ -1,7 +1,7 @@
 "use client";
 
 import { useEffect } from "react";
-import { useAssistantApi } from "../context/react/AssistantApiContext";
+import { useAui } from "@assistant-ui/store";
 import type { ToolCallMessagePartComponent } from "../types/MessagePartComponentTypes";
 import type { Tool } from "assistant-stream";
 
@@ -19,12 +19,12 @@ export const useAssistantTool = <
 >(
   tool: AssistantToolProps<TArgs, TResult>,
 ) => {
-  const api = useAssistantApi();
+  const aui = useAui();
 
   useEffect(() => {
     if (!tool.render) return undefined;
-    return api.tools().setToolUI(tool.toolName, tool.render);
-  }, [api, tool.toolName, tool.render]);
+    return aui.tools().setToolUI(tool.toolName, tool.render);
+  }, [aui, tool.toolName, tool.render]);
 
   useEffect(() => {
     const { toolName, render, ...rest } = tool;
@@ -33,8 +33,8 @@ export const useAssistantTool = <
         [toolName]: rest,
       },
     };
-    return api.modelContext().register({
+    return aui.modelContext().register({
       getModelContext: () => context,
     });
-  }, [api, tool]);
+  }, [aui, tool]);
 };

package/src/model-context/useAssistantToolUI.tsx

@@ -1,7 +1,7 @@
 "use client";
 
 import { useEffect } from "react";
-import { useAssistantApi } from "../context/react/AssistantApiContext";
+import { useAui } from "@assistant-ui/store";
 import type { ToolCallMessagePartComponent } from "../types/MessagePartComponentTypes";
 
 export type AssistantToolUIProps<TArgs, TResult> = {
@@ -12,9 +12,9 @@ export type AssistantToolUIProps<TArgs, TResult> = {
 export const useAssistantToolUI = (
   tool: AssistantToolUIProps<any, any> | null,
 ) => {
-  const api = useAssistantApi();
+  const aui = useAui();
   useEffect(() => {
     if (!tool?.toolName || !tool?.render) return undefined;
-    return api.tools().setToolUI(tool.toolName, tool.render);
-  }, [api, tool?.toolName, tool?.render]);
+    return aui.tools().setToolUI(tool.toolName, tool.render);
+  }, [aui, tool?.toolName, tool?.render]);
 };

package/src/primitives/actionBar/ActionBarCopy.tsx

@@ -5,7 +5,7 @@ import { ActionButtonProps } from "../../utils/createActionButton";
 import { composeEventHandlers } from "@radix-ui/primitive";
 import { Primitive } from "@radix-ui/react-primitive";
 import { useCallback } from "react";
-import { useAssistantState, useAssistantApi } from "../../context";
+import { useAuiState, useAui } from "@assistant-ui/store";
 
 /**
  * Hook that provides copy functionality for action bar buttons.
@@ -36,27 +36,27 @@ const useActionBarPrimitiveCopy = ({
 }: {
   copiedDuration?: number | undefined;
 } = {}) => {
-  const api = useAssistantApi();
-  const hasCopyableContent = useAssistantState(({ message }) => {
+  const aui = useAui();
+  const hasCopyableContent = useAuiState(({ message }) => {
     return (
       (message.role !== "assistant" || message.status?.type !== "running") &&
       message.parts.some((c) => c.type === "text" && c.text.length > 0)
     );
   });
 
-  const isEditing = useAssistantState(({ composer }) => composer.isEditing);
-  const composerValue = useAssistantState(({ composer }) => composer.text);
+  const isEditing = useAuiState(({ composer }) => composer.isEditing);
+  const composerValue = useAuiState(({ composer }) => composer.text);
 
   const callback = useCallback(() => {
-    const valueToCopy = isEditing ? composerValue : api.message().getCopyText();
+    const valueToCopy = isEditing ? composerValue : aui.message().getCopyText();
 
     if (!valueToCopy) return;
 
     navigator.clipboard.writeText(valueToCopy).then(() => {
-      api.message().setIsCopied(true);
-      setTimeout(() => api.message().setIsCopied(false), copiedDuration);
+      aui.message().setIsCopied(true);
+      setTimeout(() => aui.message().setIsCopied(false), copiedDuration);
     });
-  }, [api, isEditing, composerValue, copiedDuration]);
+  }, [aui, isEditing, composerValue, copiedDuration]);
 
   if (!hasCopyableContent) return null;
   return callback;
@@ -89,7 +89,7 @@ export const ActionBarPrimitiveCopy = forwardRef<
   ActionBarPrimitiveCopy.Element,
   ActionBarPrimitiveCopy.Props
 >(({ copiedDuration, onClick, disabled, ...props }, forwardedRef) => {
-  const isCopied = useAssistantState(({ message }) => message.isCopied);
+  const isCopied = useAuiState(({ message }) => message.isCopied);
   const callback = useActionBarPrimitiveCopy({ copiedDuration });
   return (
     <Primitive.button

package/src/primitives/actionBar/ActionBarEdit.tsx

@@ -6,7 +6,7 @@ import {
   createActionButton,
 } from "../../utils/createActionButton";
 import { useCallback } from "react";
-import { useAssistantState, useAssistantApi } from "../../context";
+import { useAuiState, useAui } from "@assistant-ui/store";
 
 /**
  * Hook that provides edit functionality for action bar buttons.
@@ -30,12 +30,12 @@ import { useAssistantState, useAssistantApi } from "../../context";
  * ```
  */
 const useActionBarEdit = () => {
-  const api = useAssistantApi();
-  const disabled = useAssistantState(({ composer }) => composer.isEditing);
+  const aui = useAui();
+  const disabled = useAuiState(({ composer }) => composer.isEditing);
 
   const callback = useCallback(() => {
-    api.composer().beginEdit();
-  }, [api]);
+    aui.composer().beginEdit();
+  }, [aui]);
 
   if (disabled) return null;
   return callback;

package/src/primitives/actionBar/ActionBarExportMarkdown.tsx

@@ -4,7 +4,7 @@ import { forwardRef, useCallback } from "react";
 import { ActionButtonProps } from "../../utils/createActionButton";
 import { composeEventHandlers } from "@radix-ui/primitive";
 import { Primitive } from "@radix-ui/react-primitive";
-import { useAssistantState, useAssistantApi } from "../../context";
+import { useAuiState, useAui } from "@assistant-ui/store";
 
 const useActionBarExportMarkdown = ({
   filename,
@@ -13,8 +13,8 @@ const useActionBarExportMarkdown = ({
   filename?: string | undefined;
   onExport?: ((content: string) => void | Promise<void>) | undefined;
 } = {}) => {
-  const api = useAssistantApi();
-  const hasExportableContent = useAssistantState(({ message }) => {
+  const aui = useAui();
+  const hasExportableContent = useAuiState(({ message }) => {
     return (
       (message.role !== "assistant" || message.status?.type !== "running") &&
       message.parts.some((c) => c.type === "text" && c.text.length > 0)
@@ -22,7 +22,7 @@ const useActionBarExportMarkdown = ({
   });
 
   const callback = useCallback(async () => {
-    const content = api.message().getCopyText();
+    const content = aui.message().getCopyText();
     if (!content) return;
 
     if (onExport) {
@@ -37,7 +37,7 @@ const useActionBarExportMarkdown = ({
     a.download = filename ?? `message-${Date.now()}.md`;
     a.click();
     URL.revokeObjectURL(url);
-  }, [api, filename, onExport]);
+  }, [aui, filename, onExport]);
 
   if (!hasExportableContent) return null;
   return callback;