@assistant-ui/react 0.12.17 → 0.12.20

package/src/index.ts

@@ -16,18 +16,367 @@ export {
   type AssistantEventCallback,
 } from "@assistant-ui/store";
 
-export * from "./legacy-runtime/runtime";
-export * from "./legacy-runtime/cloud";
-export * from "./legacy-runtime/runtime-cores";
+// ============================================================================
+// legacy-runtime/runtime
+// ============================================================================
+// Re-export public runtime types from @assistant-ui/core
+export type {
+  AssistantRuntime,
+  ThreadRuntime,
+  ThreadState,
+  CreateAppendMessage,
+  CreateStartRunConfig,
+  CreateResumeRunConfig,
+  MessageRuntime,
+  MessageState,
+  MessagePartRuntime,
+  MessagePartState,
+  ComposerRuntime,
+  ThreadComposerRuntime,
+  EditComposerRuntime,
+  EditComposerState,
+  ThreadComposerState,
+  ComposerState,
+  AttachmentRuntime,
+  AttachmentState,
+  ThreadListRuntime,
+  ThreadListState,
+  ThreadListItemRuntime,
+  ThreadListItemState,
+} from "@assistant-ui/core";
+
+// ============================================================================
+// legacy-runtime/cloud
+// ============================================================================
+export { useCloudThreadListRuntime } from "./legacy-runtime/cloud/useCloudThreadListRuntime";
+export { AssistantCloud } from "assistant-cloud";
+
+// ============================================================================
+// legacy-runtime/runtime-cores
+// ============================================================================
+
+// --- adapters/attachment ---
+export type { AttachmentAdapter } from "@assistant-ui/core";
+export {
+  SimpleImageAttachmentAdapter,
+  SimpleTextAttachmentAdapter,
+  CompositeAttachmentAdapter,
+} from "@assistant-ui/core";
+export { CloudFileAttachmentAdapter } from "./legacy-runtime/runtime-cores/adapters/attachment/CloudFileAttachmentAdapter";
+
+// --- adapters/feedback ---
+export type { FeedbackAdapter } from "@assistant-ui/core";
+
+// --- adapters/speech ---
+export type {
+  SpeechSynthesisAdapter,
+  DictationAdapter,
+} from "@assistant-ui/core";
+export {
+  WebSpeechSynthesisAdapter,
+  WebSpeechDictationAdapter,
+} from "@assistant-ui/core";
+
+// --- adapters/suggestion ---
+export type { SuggestionAdapter } from "@assistant-ui/core";
+
+// --- adapters/RuntimeAdapterProvider ---
+export {
+  RuntimeAdapterProvider,
+  useRuntimeAdapters,
+  type RuntimeAdapters,
+} from "@assistant-ui/core/react";
 
+// --- adapters/thread-history ---
+export type {
+  ThreadHistoryAdapter,
+  GenericThreadHistoryAdapter,
+  MessageFormatAdapter,
+  MessageFormatItem,
+  MessageFormatRepository,
+  MessageStorageEntry,
+} from "@assistant-ui/core";
+
+// --- assistant-transport ---
+export {
+  useAssistantTransportRuntime,
+  useAssistantTransportSendCommand,
+  useAssistantTransportState,
+} from "./legacy-runtime/runtime-cores/assistant-transport/useAssistantTransportRuntime";
+export type {
+  AssistantTransportConnectionMetadata,
+  AssistantTransportCommand,
+  AssistantTransportProtocol,
+  SendCommandsRequestBody,
+} from "./legacy-runtime/runtime-cores/assistant-transport/types";
+
+// --- core ---
+export type {
+  AddToolResultOptions,
+  SubmitFeedbackOptions,
+  ThreadSuggestion,
+  DictationState,
+} from "@assistant-ui/core";
+
+// --- external-store ---
+export type { ThreadMessageLike } from "@assistant-ui/core";
+export {
+  getExternalStoreMessage,
+  getExternalStoreMessages,
+  bindExternalStoreMessage,
+} from "@assistant-ui/core";
+export type {
+  ExternalStoreAdapter,
+  ExternalStoreMessageConverter,
+  ExternalStoreThreadListAdapter,
+  ExternalStoreThreadData,
+} from "@assistant-ui/core";
+export { useExternalStoreRuntime } from "./legacy-runtime/runtime-cores/external-store/useExternalStoreRuntime";
+export {
+  useExternalMessageConverter,
+  convertExternalMessages as unstable_convertExternalMessages,
+} from "./legacy-runtime/runtime-cores/external-store/external-message-converter";
+export { createMessageConverter as unstable_createMessageConverter } from "./legacy-runtime/runtime-cores/external-store/createMessageConverter";
+
+// --- local ---
+export type {
+  ChatModelAdapter,
+  ChatModelRunOptions,
+  ChatModelRunResult,
+  ChatModelRunUpdate,
+  LocalRuntimeOptionsBase,
+} from "@assistant-ui/core";
+export { useLocalRuntime } from "./legacy-runtime/runtime-cores/local/useLocalRuntime";
+/**
+ * @deprecated Use `useLocalRuntime` instead.
+ */
+export { useLocalRuntime as useLocalThreadRuntime } from "./legacy-runtime/runtime-cores/local/useLocalRuntime";
+export type { LocalRuntimeOptions } from "./legacy-runtime/runtime-cores/local/LocalRuntimeOptions";
+
+// --- remote-thread-list ---
+export { useRemoteThreadListRuntime } from "./legacy-runtime/runtime-cores/remote-thread-list/useRemoteThreadListRuntime";
+/** @deprecated Use `useRemoteThreadListRuntime` instead. */
+export { useRemoteThreadListRuntime as unstable_useRemoteThreadListRuntime } from "./legacy-runtime/runtime-cores/remote-thread-list/useRemoteThreadListRuntime";
+export { useCloudThreadListAdapter } from "./legacy-runtime/runtime-cores/remote-thread-list/adapter/cloud";
+/** @deprecated Use `useCloudThreadListAdapter` instead. */
+export { useCloudThreadListAdapter as unstable_useCloudThreadListAdapter } from "./legacy-runtime/runtime-cores/remote-thread-list/adapter/cloud";
+export type { RemoteThreadListAdapter } from "@assistant-ui/core";
+/** @deprecated Use `RemoteThreadListAdapter` instead. */
+export type { RemoteThreadListAdapter as unstable_RemoteThreadListAdapter } from "@assistant-ui/core";
+export { InMemoryThreadListAdapter } from "@assistant-ui/core";
+/** @deprecated Use `InMemoryThreadListAdapter` instead. */
+export { InMemoryThreadListAdapter as unstable_InMemoryThreadListAdapter } from "@assistant-ui/core";
+
+// Re-export from @assistant-ui/core (runtime-cores root)
+export type { ExportedMessageRepositoryItem } from "@assistant-ui/core";
+export { ExportedMessageRepository } from "@assistant-ui/core";
+
+// ============================================================================
+// context
+// ============================================================================
 export * from "./context";
-export * from "./model-context";
-export * from "./primitives";
-export * from "./types";
-export * from "./devtools";
-export * from "./client";
 
+// ============================================================================
+// model-context
+// ============================================================================
+// Re-export shared from core/react
+export {
+  makeAssistantTool,
+  type AssistantTool,
+  makeAssistantToolUI,
+  type AssistantToolUI,
+  makeAssistantDataUI,
+  type AssistantDataUI,
+  useAssistantTool,
+  type AssistantToolProps,
+  useAssistantToolUI,
+  type AssistantToolUIProps,
+  useAssistantDataUI,
+  type AssistantDataUIProps,
+  useAssistantInstructions,
+  useInlineRender,
+  type Toolkit,
+  type ToolDefinition,
+  Tools,
+  DataRenderers,
+} from "@assistant-ui/core/react";
+
+// Core pass-through (unchanged)
+export type {
+  ModelContext,
+  ModelContextProvider,
+  LanguageModelConfig,
+  LanguageModelV1CallSettings,
+} from "@assistant-ui/core";
+
+export { mergeModelContexts } from "@assistant-ui/core";
+
+export type { Tool } from "assistant-stream";
+
+export { tool } from "@assistant-ui/core";
+
+export { Suggestions, type SuggestionConfig } from "@assistant-ui/core/store";
+export type {
+  QueueItemState,
+  QueueItemMethods,
+} from "@assistant-ui/core/store";
+export type { ComposerSendOptions } from "@assistant-ui/core/store";
+
+// React-only (stays)
+export { makeAssistantVisible } from "./model-context/makeAssistantVisible";
+
+// --- model-context/registry ---
+export { ModelContextRegistry } from "@assistant-ui/core";
+export type {
+  ModelContextRegistryToolHandle,
+  ModelContextRegistryInstructionHandle,
+  ModelContextRegistryProviderHandle,
+} from "@assistant-ui/core";
+
+// --- model-context/frame ---
+export { AssistantFrameHost } from "@assistant-ui/core";
+export { AssistantFrameProvider } from "@assistant-ui/core";
+export type {
+  SerializedTool,
+  SerializedModelContext,
+  FrameMessageType,
+  FrameMessage,
+} from "@assistant-ui/core";
+export { FRAME_MESSAGE_CHANNEL } from "@assistant-ui/core";
+export { useAssistantFrameHost } from "./model-context/frame/useAssistantFrameHost";
+
+// ============================================================================
+// primitives
+// ============================================================================
+export * as ActionBarPrimitive from "./primitives/actionBar";
+export * as ActionBarMorePrimitive from "./primitives/actionBarMore";
+export * as AssistantModalPrimitive from "./primitives/assistantModal";
+export * as AttachmentPrimitive from "./primitives/attachment";
+export * as BranchPickerPrimitive from "./primitives/branchPicker";
+export * as ChainOfThoughtPrimitive from "./primitives/chainOfThought";
+export * as ComposerPrimitive from "./primitives/composer";
+export * as QueueItemPrimitive from "./primitives/queueItem";
+export * as MessagePartPrimitive from "./primitives/messagePart";
+export * as ErrorPrimitive from "./primitives/error";
+export * as MessagePrimitive from "./primitives/message";
+export * as ThreadPrimitive from "./primitives/thread";
+export * as SuggestionPrimitive from "./primitives/suggestion";
+export * as ThreadListPrimitive from "./primitives/threadList";
+export * as ThreadListItemPrimitive from "./primitives/threadListItem";
+export * as ThreadListItemMorePrimitive from "./primitives/threadListItemMore";
+export * as SelectionToolbarPrimitive from "./primitives/selectionToolbar";
+
+export { useMessagePartText } from "./primitives/messagePart/useMessagePartText";
+export { useMessagePartReasoning } from "./primitives/messagePart/useMessagePartReasoning";
+export { useMessagePartSource } from "./primitives/messagePart/useMessagePartSource";
+export { useMessagePartFile } from "./primitives/messagePart/useMessagePartFile";
+export { useMessagePartImage } from "./primitives/messagePart/useMessagePartImage";
+export { useMessagePartData } from "./primitives/messagePart/useMessagePartData";
+export { useThreadViewportAutoScroll } from "./primitives/thread/useThreadViewportAutoScroll";
+export { useScrollLock } from "./primitives/reasoning/useScrollLock";
+export { useMessageQuote } from "./hooks/useMessageQuote";
+export { useMessageTiming } from "./hooks/useMessageTiming";
+
+// ============================================================================
+// types
+// ============================================================================
+// Re-export core types from @assistant-ui/core
+export type {
+  Attachment,
+  PendingAttachment,
+  CompleteAttachment,
+  AttachmentStatus,
+  AppendMessage,
+  TextMessagePart,
+  ReasoningMessagePart,
+  SourceMessagePart,
+  ImageMessagePart,
+  FileMessagePart,
+  DataMessagePart,
+  Unstable_AudioMessagePart,
+  ToolCallMessagePart,
+  MessageStatus,
+  MessagePartStatus,
+  ToolCallMessagePartStatus,
+  MessageTiming,
+  ThreadUserMessagePart,
+  ThreadAssistantMessagePart,
+  ThreadSystemMessage,
+  ThreadAssistantMessage,
+  ThreadUserMessage,
+  ThreadMessage,
+  Unsubscribe,
+  QuoteInfo,
+  CreateAttachment,
+} from "@assistant-ui/core";
+
+// React component types (from core/react)
+export type {
+  EmptyMessagePartComponent,
+  EmptyMessagePartProps,
+  TextMessagePartComponent,
+  TextMessagePartProps,
+  ReasoningMessagePartComponent,
+  ReasoningMessagePartProps,
+  SourceMessagePartComponent,
+  SourceMessagePartProps,
+  ImageMessagePartComponent,
+  ImageMessagePartProps,
+  FileMessagePartComponent,
+  FileMessagePartProps,
+  Unstable_AudioMessagePartComponent,
+  Unstable_AudioMessagePartProps,
+  DataMessagePartComponent,
+  DataMessagePartProps,
+  ToolCallMessagePartComponent,
+  ToolCallMessagePartProps,
+  ReasoningGroupProps,
+  ReasoningGroupComponent,
+  QuoteMessagePartComponent,
+  QuoteMessagePartProps,
+} from "@assistant-ui/core/react";
+
+// Thread list item types
+export type { ThreadListItemStatus } from "@assistant-ui/core";
+
+// ============================================================================
+// devtools
+// ============================================================================
+export { DevToolsHooks, DevToolsProviderApi } from "./devtools/DevToolsHooks";
+
+// ============================================================================
+// client
+// ============================================================================
+// Note: Tools, DataRenderers, Suggestions, SuggestionConfig already exported above from model-context
+export { ModelContext as ModelContextClient } from "@assistant-ui/core/store";
+export { ChainOfThoughtClient } from "@assistant-ui/core/store";
+export {
+  ExternalThread,
+  type ExternalThreadProps,
+  type ExternalThreadMessage,
+  type ExternalThreadQueueAdapter,
+} from "./client/ExternalThread";
+export {
+  InMemoryThreadList,
+  type InMemoryThreadListProps,
+} from "./client/InMemoryThreadList";
+export { SingleThreadList } from "./client/SingleThreadList";
+
+// ============================================================================
+// internal & augmentations
+// ============================================================================
 export * as INTERNAL from "./internal";
+
+// Unstable - mention
+export {
+  unstable_useToolMentionAdapter,
+  type Unstable_ToolMentionAdapterOptions,
+} from "./unstable/useToolMentionAdapter";
+export {
+  useMentionContext as unstable_useMentionContext,
+  useMentionContextOptional as unstable_useMentionContextOptional,
+  useMentionInternalContext as unstable_useMentionInternalContext,
+} from "./primitives/composer/mention";
 export type { ToolExecutionStatus } from "./internal";
 
 export type { Assistant } from "./augmentations";

package/src/internal.ts

@@ -28,4 +28,8 @@ export {
   type ToolExecutionStatus,
 } from "@assistant-ui/core/react";
 
-export * from "./utils/smooth";
+export { useSmooth } from "./utils/smooth/useSmooth";
+export {
+  useSmoothStatus,
+  withSmoothContextProvider,
+} from "./utils/smooth/SmoothContext";

package/src/legacy-runtime/AssistantRuntimeProvider.tsx

@@ -5,7 +5,7 @@ import { AssistantClient, useAui } from "@assistant-ui/store";
 import { AssistantRuntime } from "./runtime/AssistantRuntime";
 import { AssistantProviderBase } from "@assistant-ui/core/react";
 import { ThreadPrimitiveViewportProvider } from "../context/providers/ThreadViewportProvider";
-import { DevToolsProviderApi } from "../devtools";
+import { DevToolsProviderApi } from "../devtools/DevToolsHooks";
 
 export namespace AssistantRuntimeProvider {
   export type Props = PropsWithChildren<{

package/src/legacy-runtime/cloud/useCloudThreadListRuntime.ts

@@ -1,7 +1,7 @@
 "use client";
 
 import { AssistantCloud } from "assistant-cloud";
-import { AssistantRuntime } from "../runtime";
+import type { AssistantRuntime } from "@assistant-ui/core";
 import { useRemoteThreadListRuntime } from "../runtime-cores/remote-thread-list/useRemoteThreadListRuntime";
 import { useCloudThreadListAdapter } from "../runtime-cores/remote-thread-list/adapter/cloud";
 

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

@@ -3,7 +3,7 @@
 import { MessageRuntime } from "../runtime/MessageRuntime";
 import { useAui, useAuiState } from "@assistant-ui/store";
 import { createStateHookForRuntime } from "../../context/react/utils/createStateHookForRuntime";
-import { EditComposerRuntime } from "../runtime";
+import type { EditComposerRuntime } from "@assistant-ui/core";
 
 /**
  * @deprecated Use `useAui()` with `aui.message()` instead. See migration guide: https://assistant-ui.com/docs/migrations/v0-12

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

@@ -2,9 +2,9 @@
 
 import { useState } from "react";
 import { ThreadRuntime } from "../runtime/ThreadRuntime";
-import { ModelContext } from "../../model-context";
+import type { ModelContext } from "@assistant-ui/core";
 import { createStateHookForRuntime } from "../../context/react/utils/createStateHookForRuntime";
-import { ThreadComposerRuntime } from "../runtime";
+import type { ThreadComposerRuntime } from "@assistant-ui/core";
 import { useAui, useAuiEvent, useAuiState } from "@assistant-ui/store";
 
 /**

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

@@ -1,6 +1,9 @@
 import type { ThreadMessage } from "@assistant-ui/core";
 import { ReadonlyJSONValue } from "assistant-stream/utils";
-import { AttachmentAdapter, ThreadHistoryAdapter } from "..";
+import type {
+  AttachmentAdapter,
+  ThreadHistoryAdapter,
+} from "@assistant-ui/core";
 import { UserCommands } from "../../../augmentations";
 import type { ToolExecutionStatus } from "./useToolInvocations";
 

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

@@ -8,7 +8,7 @@ import {
 } from "assistant-stream/utils";
 import { useExternalStoreRuntime } from "../external-store/useExternalStoreRuntime";
 import { AssistantRuntime } from "../../runtime/AssistantRuntime";
-import { AddToolResultOptions } from "../core";
+import type { AddToolResultOptions } from "@assistant-ui/core";
 import { useState, useRef, useMemo } from "react";
 import {
   AssistantMessageAccumulator,

package/src/model-context/frame/SPEC_AssistantFrame.md

@@ -45,8 +45,8 @@ registry.addTool({
 const frameHost = new AssistantFrameHost(iframeWindow);
 
 // Register with assistant runtime
-const runtime = useAssistantRuntime();
-runtime.registerModelContextProvider(frameHost);
+const aui = useAui();
+aui.modelContext().register(frameHost);
 
 // The assistant now has access to tools from the iframe
 ```

package/src/primitives/actionBar/ActionBarSpeak.ts

@@ -1,29 +1,16 @@
 "use client";
 
-import { useCallback } from "react";
-import { useAuiState, useAui } from "@assistant-ui/store";
 import {
   ActionButtonElement,
   ActionButtonProps,
   createActionButton,
 } from "../../utils/createActionButton";
+import { useActionBarSpeak as useActionBarSpeakBehavior } from "@assistant-ui/core/react";
 
 const useActionBarSpeak = () => {
-  const aui = useAui();
-  const callback = useCallback(async () => {
-    aui.message().speak();
-  }, [aui]);
-
-  const hasSpeakableContent = useAuiState((s) => {
-    return (
-      (s.message.role !== "assistant" ||
-        s.message.status?.type !== "running") &&
-      s.message.parts.some((c) => c.type === "text" && c.text.length > 0)
-    );
-  });
-
-  if (!hasSpeakableContent) return null;
-  return callback;
+  const { disabled, speak } = useActionBarSpeakBehavior();
+  if (disabled) return null;
+  return speak;
 };
 
 export namespace ActionBarPrimitiveSpeak {

package/src/primitives/actionBar/ActionBarStopSpeaking.tsx

@@ -5,20 +5,12 @@ import { ActionButtonProps } from "../../utils/createActionButton";
 import { useEscapeKeydown } from "@radix-ui/react-use-escape-keydown";
 import { Primitive } from "@radix-ui/react-primitive";
 import { composeEventHandlers } from "@radix-ui/primitive";
-import { useCallback } from "react";
-import { useAuiState, useAui } from "@assistant-ui/store";
+import { useActionBarStopSpeaking as useActionBarStopSpeakingBehavior } from "@assistant-ui/core/react";
 
 const useActionBarStopSpeaking = () => {
-  const aui = useAui();
-  const isSpeaking = useAuiState((s) => s.message.speech != null);
-
-  const callback = useCallback(() => {
-    aui.message().stopSpeaking();
-  }, [aui]);
-
-  if (!isSpeaking) return null;
-
-  return callback;
+  const { disabled, stopSpeaking } = useActionBarStopSpeakingBehavior();
+  if (disabled) return null;
+  return stopSpeaking;
 };
 
 export namespace ActionBarPrimitiveStopSpeaking {

package/src/primitives/actionBar.ts

@@ -0,0 +1,9 @@
+export { ActionBarPrimitiveRoot as Root } from "./actionBar/ActionBarRoot";
+export { ActionBarPrimitiveCopy as Copy } from "./actionBar/ActionBarCopy";
+export { ActionBarPrimitiveReload as Reload } from "./actionBar/ActionBarReload";
+export { ActionBarPrimitiveEdit as Edit } from "./actionBar/ActionBarEdit";
+export { ActionBarPrimitiveSpeak as Speak } from "./actionBar/ActionBarSpeak";
+export { ActionBarPrimitiveStopSpeaking as StopSpeaking } from "./actionBar/ActionBarStopSpeaking";
+export { ActionBarPrimitiveFeedbackPositive as FeedbackPositive } from "./actionBar/ActionBarFeedbackPositive";
+export { ActionBarPrimitiveFeedbackNegative as FeedbackNegative } from "./actionBar/ActionBarFeedbackNegative";
+export { ActionBarPrimitiveExportMarkdown as ExportMarkdown } from "./actionBar/ActionBarExportMarkdown";

package/src/primitives/actionBarMore.ts

@@ -0,0 +1,5 @@
+export { ActionBarMorePrimitiveRoot as Root } from "./actionBarMore/ActionBarMoreRoot";
+export { ActionBarMorePrimitiveTrigger as Trigger } from "./actionBarMore/ActionBarMoreTrigger";
+export { ActionBarMorePrimitiveContent as Content } from "./actionBarMore/ActionBarMoreContent";
+export { ActionBarMorePrimitiveItem as Item } from "./actionBarMore/ActionBarMoreItem";
+export { ActionBarMorePrimitiveSeparator as Separator } from "./actionBarMore/ActionBarMoreSeparator";

package/src/primitives/assistantModal.ts

@@ -0,0 +1,4 @@
+export { AssistantModalPrimitiveRoot as Root } from "./assistantModal/AssistantModalRoot";
+export { AssistantModalPrimitiveTrigger as Trigger } from "./assistantModal/AssistantModalTrigger";
+export { AssistantModalPrimitiveContent as Content } from "./assistantModal/AssistantModalContent";
+export { AssistantModalPrimitiveAnchor as Anchor } from "./assistantModal/AssistantModalAnchor";

package/src/primitives/attachment/AttachmentThumb.tsx

@@ -26,4 +26,4 @@ export const AttachmentPrimitiveThumb = forwardRef<
   );
 });
 
-AttachmentPrimitiveThumb.displayName = "AttachmentPrimitive.Thumb";
+AttachmentPrimitiveThumb.displayName = "AttachmentPrimitive.unstable_Thumb";

package/src/primitives/attachment.ts

@@ -0,0 +1,4 @@
+export { AttachmentPrimitiveRoot as Root } from "./attachment/AttachmentRoot";
+export { AttachmentPrimitiveThumb as unstable_Thumb } from "./attachment/AttachmentThumb";
+export { AttachmentPrimitiveName as Name } from "./attachment/AttachmentName";
+export { AttachmentPrimitiveRemove as Remove } from "./attachment/AttachmentRemove";

package/src/primitives/branchPicker/BranchPickerNext.ts

@@ -5,29 +5,12 @@ import {
   ActionButtonProps,
   createActionButton,
 } from "../../utils/createActionButton";
-import { useCallback } from "react";
-import { useAuiState, useAui } from "@assistant-ui/store";
+import { useBranchPickerNext as useBranchPickerNextBehavior } from "@assistant-ui/core/react";
 
 const useBranchPickerNext = () => {
-  const aui = useAui();
-  const disabled = useAuiState((s) => {
-    // Disabled if no next branch
-    if (s.message.branchNumber >= s.message.branchCount) return true;
-
-    // Disabled if running and capability not supported
-    if (s.thread.isRunning && !s.thread.capabilities.switchBranchDuringRun) {
-      return true;
-    }
-
-    return false;
-  });
-
-  const callback = useCallback(() => {
-    aui.message().switchToBranch({ position: "next" });
-  }, [aui]);
-
+  const { disabled, next } = useBranchPickerNextBehavior();
   if (disabled) return null;
-  return callback;
+  return next;
 };
 
 export namespace BranchPickerPrimitiveNext {

package/src/primitives/branchPicker/BranchPickerPrevious.ts

@@ -5,50 +5,12 @@ import {
   ActionButtonProps,
   createActionButton,
 } from "../../utils/createActionButton";
-import { useCallback } from "react";
-import { useAuiState, useAui } from "@assistant-ui/store";
+import { useBranchPickerPrevious as useBranchPickerPreviousBehavior } from "@assistant-ui/core/react";
 
-/**
- * Hook that provides navigation to the previous branch functionality.
- *
- * This hook returns a callback function that switches to the previous branch
- * in the message branch tree, or null if there is no previous branch available.
- *
- * @returns A previous branch callback function, or null if navigation is disabled
- *
- * @example
- * ```tsx
- * function CustomPreviousButton() {
- *   const previous = useBranchPickerPrevious();
- *
- *   return (
- *     <button onClick={previous} disabled={!previous}>
- *       {previous ? "Previous Branch" : "No Previous Branch"}
- *     </button>
- *   );
- * }
- * ```
- */
 const useBranchPickerPrevious = () => {
-  const aui = useAui();
-  const disabled = useAuiState((s) => {
-    // Disabled if no previous branch
-    if (s.message.branchNumber <= 1) return true;
-
-    // Disabled if running and capability not supported
-    if (s.thread.isRunning && !s.thread.capabilities.switchBranchDuringRun) {
-      return true;
-    }
-
-    return false;
-  });
-
-  const callback = useCallback(() => {
-    aui.message().switchToBranch({ position: "previous" });
-  }, [aui]);
-
+  const { disabled, previous } = useBranchPickerPreviousBehavior();
   if (disabled) return null;
-  return callback;
+  return previous;
 };
 
 export namespace BranchPickerPrimitivePrevious {

package/src/primitives/branchPicker/BranchPickerRoot.tsx

@@ -2,7 +2,7 @@
 
 import { Primitive } from "@radix-ui/react-primitive";
 import { type ComponentRef, forwardRef, ComponentPropsWithoutRef } from "react";
-import { If } from "../message";
+import { MessagePrimitiveIf as If } from "../message/MessageIf";
 
 export namespace BranchPickerPrimitiveRoot {
   export type Element = ComponentRef<typeof Primitive.div>;

package/src/primitives/branchPicker.ts

@@ -0,0 +1,5 @@
+export { BranchPickerPrimitiveNext as Next } from "./branchPicker/BranchPickerNext";
+export { BranchPickerPrimitivePrevious as Previous } from "./branchPicker/BranchPickerPrevious";
+export { BranchPickerPrimitiveCount as Count } from "./branchPicker/BranchPickerCount";
+export { BranchPickerPrimitiveNumber as Number } from "./branchPicker/BranchPickerNumber";
+export { BranchPickerPrimitiveRoot as Root } from "./branchPicker/BranchPickerRoot";