@assistant-ui/core 0.1.3 → 0.1.7

package/src/react/primitives/message/MessageParts.tsx

@@ -1,13 +1,21 @@
 import {
   type ComponentType,
   type FC,
+  type ReactNode,
   memo,
   PropsWithChildren,
   useMemo,
 } from "react";
-import { useAuiState, useAui } from "@assistant-ui/store";
-import { PartByIndexProvider, TextMessagePartProvider } from "../../providers";
+import {
+  RenderChildrenWithAccessor,
+  useAuiState,
+  useAui,
+} from "@assistant-ui/store";
+import type { PartState } from "../../../store/scopes/part";
+import { PartByIndexProvider } from "../../providers/PartByIndexProvider";
+import { TextMessagePartProvider } from "../../providers/TextMessagePartProvider";
 import { ChainOfThoughtByIndicesProvider } from "../../providers/ChainOfThoughtByIndicesProvider";
+import { getMessageQuote } from "../../utils/getMessageQuote";
 import type {
   Unstable_AudioMessagePartComponent,
   DataMessagePartComponent,
@@ -21,8 +29,9 @@ import type {
   FileMessagePartComponent,
   ReasoningMessagePartComponent,
   ReasoningGroupComponent,
-} from "../../types";
-import type { MessagePartStatus } from "../../../types";
+  QuoteMessagePartComponent,
+} from "../../types/MessagePartComponentTypes";
+import type { MessagePartStatus } from "../../../types/message";
 import { useShallow } from "zustand/shallow";
 
 type MessagePartRange =
@@ -161,6 +170,8 @@ export namespace MessagePrimitiveParts {
     Unstable_Audio?: Unstable_AudioMessagePartComponent | undefined;
     /** Configuration for data part rendering */
     data?: DataConfig | undefined;
+    /** Component for rendering a quoted message reference (from metadata, not parts) */
+    Quote?: QuoteMessagePartComponent | undefined;
   };
 
   type ToolsConfig =
@@ -230,24 +241,32 @@ export namespace MessagePrimitiveParts {
     ReasoningGroup?: never;
   };
 
-  export type Props = {
-    /**
-     * Component configuration for rendering different types of message content.
-     *
-     * Use either `Reasoning`/`tools`/`ToolGroup`/`ReasoningGroup` for standard rendering,
-     * or `ChainOfThought` to group all reasoning and tool-call parts into a single
-     * collapsible component. These two modes are mutually exclusive.
-     */
-    components?: StandardComponents | ChainOfThoughtComponents | undefined;
-    /**
-     * When enabled, shows the Empty component if the last part in the message
-     * is anything other than Text or Reasoning.
-     *
-     * @experimental This API is experimental and may change in future versions.
-     * @default true
-     */
-    unstable_showEmptyOnNonTextEnd?: boolean | undefined;
-  };
+  export type Props =
+    | {
+        /**
+         * Component configuration for rendering different types of message content.
+         *
+         * Use either `Reasoning`/`tools`/`ToolGroup`/`ReasoningGroup` for standard rendering,
+         * or `ChainOfThought` to group all reasoning and tool-call parts into a single
+         * collapsible component. These two modes are mutually exclusive.
+         */
+        components?: StandardComponents | ChainOfThoughtComponents | undefined;
+        /**
+         * When enabled, shows the Empty component if the last part in the message
+         * is anything other than Text or Reasoning.
+         *
+         * @experimental This API is experimental and may change in future versions.
+         * @default true
+         */
+        unstable_showEmptyOnNonTextEnd?: boolean | undefined;
+        children?: never;
+      }
+    | {
+        /** Render function called for each part. Receives the enriched part state. */
+        children: (value: { part: EnrichedPartState }) => ReactNode;
+        components?: never;
+        unstable_showEmptyOnNonTextEnd?: never;
+      };
 }
 
 const ToolUIDisplay = ({
@@ -460,6 +479,158 @@ const ConditionalEmpty = memo(
     prev.components?.Text === next.components?.Text,
 );
 
+const QuoteRendererImpl: FC<{ Quote: QuoteMessagePartComponent }> = ({
+  Quote,
+}) => {
+  const quoteInfo = useAuiState(getMessageQuote);
+  if (!quoteInfo) return null;
+  return <Quote text={quoteInfo.text} messageId={quoteInfo.messageId} />;
+};
+
+const QuoteRenderer = memo(QuoteRendererImpl);
+
+/**
+ * Stable propless component that renders the registered tool UI for the
+ * current part context. Reads tool registry and part state from context.
+ */
+const RegisteredToolUI: FC = () => {
+  const aui = useAui();
+  const part = useAuiState((s) => s.part);
+  const Render = useAuiState((s) => {
+    if (s.part.type !== "tool-call") return null;
+    const entry = s.tools.tools[s.part.toolName];
+    if (Array.isArray(entry)) return entry[0] ?? null;
+    return entry ?? null;
+  });
+
+  if (!Render || part.type !== "tool-call") return null;
+
+  return (
+    <Render
+      {...part}
+      addResult={aui.part().addToolResult}
+      resume={aui.part().resumeToolCall}
+    />
+  );
+};
+
+/**
+ * Stable propless component that renders the registered data renderer UI
+ * for the current part context.
+ */
+const RegisteredDataRendererUI: FC = () => {
+  const part = useAuiState((s) => s.part);
+  const Render = useAuiState((s) => {
+    if (s.part.type !== "data") return null;
+    const entry = s.dataRenderers.renderers[s.part.name];
+    if (Array.isArray(entry)) return entry[0] ?? null;
+    return entry ?? null;
+  });
+
+  if (!Render || part.type !== "data") return null;
+
+  return <Render {...(part as DataMessagePartProps)} />;
+};
+
+/**
+ * Fallback component rendered when the children render function returns null.
+ * Renders registered tool/data UIs via context.
+ * For all other part types, renders nothing.
+ *
+ * This allows users to write:
+ *   {({ part }) => {
+ *     if (part.type === "text") return <MyText />;
+ *     return null; // tool UIs and data UIs still render via registry
+ *   }}
+ *
+ * To explicitly render nothing (suppressing registered UIs), return <></>.
+ */
+const DefaultPartFallback: FC = () => {
+  const partType = useAuiState((s) => s.part.type);
+
+  if (partType === "tool-call") return <RegisteredToolUI />;
+  if (partType === "data") return <RegisteredDataRendererUI />;
+
+  return null;
+};
+
+export type { PartState };
+
+/**
+ * Enriched part state passed to children render functions.
+ *
+ * For tool-call parts, adds `toolUI`, `addResult`, and `resume`.
+ * For data parts, adds `dataRendererUI`.
+ */
+export type EnrichedPartState =
+  | (Extract<PartState, { type: "tool-call" }> & {
+      /** The registered tool UI element, or null if none registered. */
+      readonly toolUI: ReactNode;
+      /** Add a tool result to this tool call. */
+      addResult: ToolCallMessagePartProps["addResult"];
+      /** Resume a tool call waiting for human input. */
+      resume: ToolCallMessagePartProps["resume"];
+    })
+  | (Extract<PartState, { type: "data" }> & {
+      /** The registered data renderer UI element, or null if none registered. */
+      readonly dataRendererUI: ReactNode;
+    })
+  | Exclude<PartState, { type: "tool-call" } | { type: "data" }>;
+
+const MessagePrimitivePartsInner: FC<{
+  children: (value: { part: EnrichedPartState }) => ReactNode;
+}> = ({ children }) => {
+  const aui = useAui();
+  const contentLength = useAuiState((s) => s.message.parts.length);
+
+  // biome-ignore lint/correctness/useExhaustiveDependencies: aui accessors are stable refs
+  return useMemo(
+    () =>
+      Array.from({ length: contentLength }, (_, index) => (
+        <PartByIndexProvider key={index} index={index}>
+          <RenderChildrenWithAccessor
+            getItemState={(aui) => aui.message().part({ index }).getState()}
+          >
+            {(getItem) => {
+              const result = children({
+                get part() {
+                  const state = getItem();
+                  if (state.type === "tool-call") {
+                    const entry = aui.tools().getState().tools[state.toolName];
+                    const hasUI = Array.isArray(entry) ? !!entry[0] : !!entry;
+                    const partMethods = aui.message().part({ index });
+                    return {
+                      ...state,
+                      toolUI: hasUI ? <RegisteredToolUI /> : null,
+                      addResult: partMethods.addToolResult,
+                      resume: partMethods.resumeToolCall,
+                    };
+                  }
+                  if (state.type === "data") {
+                    const entry = aui.dataRenderers().getState().renderers[
+                      state.name
+                    ];
+                    const hasUI = Array.isArray(entry) ? !!entry[0] : !!entry;
+                    return {
+                      ...state,
+                      dataRendererUI: hasUI ? (
+                        <RegisteredDataRendererUI />
+                      ) : null,
+                    };
+                  }
+                  return state;
+                },
+              });
+              if (result !== null) return result;
+              return <DefaultPartFallback />;
+            }}
+          </RenderChildrenWithAccessor>
+        </PartByIndexProvider>
+      )),
+    [contentLength, children],
+  );
+};
+
 /**
  * Renders the parts of a message with support for multiple content types.
  *
@@ -469,7 +640,25 @@ const ConditionalEmpty = memo(
 export const MessagePrimitiveParts: FC<MessagePrimitiveParts.Props> = ({
   components,
   unstable_showEmptyOnNonTextEnd = true,
+  children,
 }) => {
+  if (children) {
+    return <MessagePrimitivePartsInner>{children}</MessagePrimitivePartsInner>;
+  }
+  return (
+    <MessagePrimitivePartsCompat
+      components={components}
+      unstable_showEmptyOnNonTextEnd={unstable_showEmptyOnNonTextEnd}
+    />
+  );
+};
+
+MessagePrimitiveParts.displayName = "MessagePrimitive.Parts";
+
+const MessagePrimitivePartsCompat: FC<{
+  components: MessagePrimitiveParts.Props["components"];
+  unstable_showEmptyOnNonTextEnd: boolean;
+}> = ({ components, unstable_showEmptyOnNonTextEnd }) => {
   const contentLength = useAuiState((s) => s.message.parts.length);
   const useChainOfThought = !!components?.ChainOfThought;
   const messageRanges = useMessagePartsGroups(useChainOfThought);
@@ -549,6 +738,7 @@ export const MessagePrimitiveParts: FC<MessagePrimitiveParts.Props> = ({
 
   return (
     <>
+      {components?.Quote && <QuoteRenderer Quote={components.Quote} />}
       {partsElements}
       <ConditionalEmpty
         components={components}
@@ -557,5 +747,3 @@ export const MessagePrimitiveParts: FC<MessagePrimitiveParts.Props> = ({
     </>
   );
 };
-
-MessagePrimitiveParts.displayName = "MessagePrimitive.Parts";

package/src/react/primitives/message/MessageQuote.tsx

@@ -0,0 +1,37 @@
+import { type FC, type ReactNode, memo } from "react";
+import { useAuiState } from "@assistant-ui/store";
+import { getMessageQuote } from "../../utils/getMessageQuote";
+import type { QuoteInfo } from "../../../types/quote";
+
+export namespace MessagePrimitiveQuote {
+  export type Props = {
+    /** Render function called when a quote is present. Receives quote info. */
+    children: (value: QuoteInfo) => ReactNode;
+  };
+}
+
+/**
+ * Renders a quote block if the message has quote metadata.
+ * Place this above `MessagePrimitive.Parts` in your message layout.
+ *
+ * @example
+ * ```tsx
+ * <MessagePrimitive.Quote>
+ *   {({ text, messageId }) => <QuoteBlock text={text} messageId={messageId} />}
+ * </MessagePrimitive.Quote>
+ * <MessagePrimitive.Parts>
+ *   {({ part }) => { ... }}
+ * </MessagePrimitive.Parts>
+ * ```
+ */
+const MessagePrimitiveQuoteImpl: FC<MessagePrimitiveQuote.Props> = ({
+  children,
+}) => {
+  const quoteInfo = useAuiState(getMessageQuote);
+  if (!quoteInfo) return null;
+  return <>{children(quoteInfo)}</>;
+};
+
+export const MessagePrimitiveQuote = memo(MessagePrimitiveQuoteImpl);
+
+MessagePrimitiveQuote.displayName = "MessagePrimitive.Quote";

package/src/react/primitives/part/PartMessages.tsx

@@ -0,0 +1,75 @@
+import { type FC, type ReactNode, memo } from "react";
+import { useAuiState } from "@assistant-ui/store";
+import type { ThreadMessage } from "../../../types/message";
+import { ReadonlyThreadProvider } from "../../providers/ReadonlyThreadProvider";
+import {
+  ThreadPrimitiveMessages,
+  ThreadPrimitiveMessagesImpl,
+} from "../thread/ThreadMessages";
+
+export namespace PartPrimitiveMessages {
+  export type Props = {
+    components?: ThreadPrimitiveMessages.Props["components"];
+    /** Render function called for each message. Receives the message. */
+    children?: (value: { message: ThreadMessage }) => ReactNode;
+  };
+}
+
+const usePartMessages = (): readonly ThreadMessage[] | undefined => {
+  return useAuiState((s) => {
+    const part = s.part;
+    if (part.type !== "tool-call") return undefined;
+    return "messages" in part
+      ? (part.messages as readonly ThreadMessage[] | undefined)
+      : undefined;
+  });
+};
+
+/**
+ * Renders the nested messages of a tool call part (e.g. sub-agent conversation).
+ *
+ * This primitive reads `messages` from the current tool call part in the PartScope
+ * and renders them using a readonly thread context. All existing message and part
+ * primitives work inside, and parent tool UI registrations are inherited.
+ *
+ * @example
+ * ```tsx
+ * const SubAgentToolUI = makeAssistantToolUI({
+ *   toolName: "invoke_sub_agent",
+ *   render: () => (
+ *     <PartPrimitive.Messages>
+ *       {({ message }) => {
+ *         if (message.role === "user") return <MyUserMessage />;
+ *         return <MyAssistantMessage />;
+ *       }}
+ *     </PartPrimitive.Messages>
+ *   ),
+ * });
+ * ```
+ */
+export const PartPrimitiveMessagesImpl: FC<PartPrimitiveMessages.Props> = ({
+  components,
+  children,
+}) => {
+  const messages = usePartMessages();
+
+  if (!messages?.length) return null;
+
+  if (children) {
+    return (
+      <ReadonlyThreadProvider messages={messages}>
+        <ThreadPrimitiveMessagesImpl>{children}</ThreadPrimitiveMessagesImpl>
+      </ReadonlyThreadProvider>
+    );
+  }
+
+  return (
+    <ReadonlyThreadProvider messages={messages}>
+      <ThreadPrimitiveMessagesImpl components={components!} />
+    </ReadonlyThreadProvider>
+  );
+};
+
+PartPrimitiveMessagesImpl.displayName = "PartPrimitive.Messages";
+
+export const PartPrimitiveMessages = memo(PartPrimitiveMessagesImpl);

package/src/react/primitives/thread/ThreadMessages.tsx

@@ -1,63 +1,69 @@
-import type { ThreadMessage as ThreadMessageType } from "../../../types";
-import { type ComponentType, type FC, memo, useMemo } from "react";
-import { useAuiState } from "@assistant-ui/store";
-import { MessageByIndexProvider } from "../../providers";
+import {
+  type ComponentType,
+  type FC,
+  type ReactNode,
+  memo,
+  useMemo,
+} from "react";
+import { RenderChildrenWithAccessor, useAuiState } from "@assistant-ui/store";
+import { MessageByIndexProvider } from "../../providers/MessageByIndexProvider";
+import { MessageState } from "../../../store";
+
+type MessagesComponentConfig =
+  | {
+      /** Component used to render all message types */
+      Message: ComponentType;
+      /** Component used when editing any message type */
+      EditComposer?: ComponentType | undefined;
+      /** Component used when editing user messages specifically */
+      UserEditComposer?: ComponentType | undefined;
+      /** Component used when editing assistant messages specifically */
+      AssistantEditComposer?: ComponentType | undefined;
+      /** Component used when editing system messages specifically */
+      SystemEditComposer?: ComponentType | undefined;
+      /** Component used to render user messages specifically */
+      UserMessage?: ComponentType | undefined;
+      /** Component used to render assistant messages specifically */
+      AssistantMessage?: ComponentType | undefined;
+      /** Component used to render system messages specifically */
+      SystemMessage?: ComponentType | undefined;
+    }
+  | {
+      /** Component used to render all message types (fallback) */
+      Message?: ComponentType | undefined;
+      /** Component used when editing any message type */
+      EditComposer?: ComponentType | undefined;
+      /** Component used when editing user messages specifically */
+      UserEditComposer?: ComponentType | undefined;
+      /** Component used when editing assistant messages specifically */
+      AssistantEditComposer?: ComponentType | undefined;
+      /** Component used when editing system messages specifically */
+      SystemEditComposer?: ComponentType | undefined;
+      /** Component used to render user messages */
+      UserMessage: ComponentType;
+      /** Component used to render assistant messages */
+      AssistantMessage: ComponentType;
+      /** Component used to render system messages */
+      SystemMessage?: ComponentType | undefined;
+    };
 
 export namespace ThreadPrimitiveMessages {
-  export type Props = {
-    /**
-     * Component configuration for rendering different types of messages and composers.
-     *
-     * You can provide either:
-     * 1. A single `Message` component that handles all message types
-     * 2. Specific components for `UserMessage` and `AssistantMessage` (with optional `SystemMessage`)
-     *
-     * Optional edit composer components can be provided to customize the editing experience
-     * for different message types when users edit their messages.
-     */
-    components:
-      | {
-          /** Component used to render all message types */
-          Message: ComponentType;
-          /** Component used when editing any message type */
-          EditComposer?: ComponentType | undefined;
-          /** Component used when editing user messages specifically */
-          UserEditComposer?: ComponentType | undefined;
-          /** Component used when editing assistant messages specifically */
-          AssistantEditComposer?: ComponentType | undefined;
-          /** Component used when editing system messages specifically */
-          SystemEditComposer?: ComponentType | undefined;
-          /** Component used to render user messages specifically */
-          UserMessage?: ComponentType | undefined;
-          /** Component used to render assistant messages specifically */
-          AssistantMessage?: ComponentType | undefined;
-          /** Component used to render system messages specifically */
-          SystemMessage?: ComponentType | undefined;
-        }
-      | {
-          /** Component used to render all message types (fallback) */
-          Message?: ComponentType | undefined;
-          /** Component used when editing any message type */
-          EditComposer?: ComponentType | undefined;
-          /** Component used when editing user messages specifically */
-          UserEditComposer?: ComponentType | undefined;
-          /** Component used when editing assistant messages specifically */
-          AssistantEditComposer?: ComponentType | undefined;
-          /** Component used when editing system messages specifically */
-          SystemEditComposer?: ComponentType | undefined;
-          /** Component used to render user messages */
-          UserMessage: ComponentType;
-          /** Component used to render assistant messages */
-          AssistantMessage: ComponentType;
-          /** Component used to render system messages */
-          SystemMessage?: ComponentType | undefined;
-        };
-  };
+  export type Props =
+    | {
+        /** @deprecated Use the children render function instead. */
+        components: MessagesComponentConfig;
+        children?: never;
+      }
+    | {
+        /** Render function called for each message. Receives the message. */
+        children: (value: { message: MessageState }) => ReactNode;
+        components?: never;
+      };
 }
 
 const isComponentsSame = (
-  prev: ThreadPrimitiveMessages.Props["components"],
-  next: ThreadPrimitiveMessages.Props["components"],
+  prev: MessagesComponentConfig,
+  next: MessagesComponentConfig,
 ) => {
   return (
     prev.Message === next.Message &&
@@ -74,8 +80,8 @@ const isComponentsSame = (
 const DEFAULT_SYSTEM_MESSAGE = () => null;
 
 const getComponent = (
-  components: ThreadPrimitiveMessages.Props["components"],
-  role: ThreadMessageType["role"],
+  components: MessagesComponentConfig,
+  role: MessageState["role"],
   isEditing: boolean,
 ) => {
   switch (role) {
@@ -125,7 +131,7 @@ const getComponent = (
 };
 
 type ThreadMessageComponentProps = {
-  components: ThreadPrimitiveMessages.Props["components"];
+  components: MessagesComponentConfig;
 };
 
 const ThreadMessageComponent: FC<ThreadMessageComponentProps> = ({
@@ -140,26 +146,12 @@ const ThreadMessageComponent: FC<ThreadMessageComponentProps> = ({
 export namespace ThreadPrimitiveMessageByIndex {
   export type Props = {
     index: number;
-    components: ThreadPrimitiveMessages.Props["components"];
+    components: MessagesComponentConfig;
   };
 }
 
 /**
  * Renders a single message at the specified index in the current thread.
- *
- * This component provides message context for a specific message in the thread
- * and renders it using the provided component configuration.
- *
- * @example
- * ```tsx
- * <ThreadPrimitive.MessageByIndex
- *   index={0}
- *   components={{
- *     UserMessage: MyUserMessage,
- *     AssistantMessage: MyAssistantMessage
- *   }}
- * />
- * ```
  */
 export const ThreadPrimitiveMessageByIndex: FC<ThreadPrimitiveMessageByIndex.Props> =
   memo(
@@ -177,46 +169,68 @@ export const ThreadPrimitiveMessageByIndex: FC<ThreadPrimitiveMessageByIndex.Pro
 
 ThreadPrimitiveMessageByIndex.displayName = "ThreadPrimitive.MessageByIndex";
 
+const ThreadPrimitiveMessagesInner: FC<{
+  children: (value: { message: MessageState }) => ReactNode;
+}> = ({ children }) => {
+  const messagesLength = useAuiState((s) => s.thread.messages.length);
+
+  return useMemo(() => {
+    if (messagesLength === 0) return null;
+    return Array.from({ length: messagesLength }, (_, index) => (
+      <MessageByIndexProvider key={index} index={index}>
+        <RenderChildrenWithAccessor
+          getItemState={(aui) => aui.thread().message({ index }).getState()}
+        >
+          {(getItem) =>
+            children({
+              get message() {
+                return getItem();
+              },
+            })
+          }
+        </RenderChildrenWithAccessor>
+      </MessageByIndexProvider>
+    ));
+  }, [messagesLength, children]);
+};
+
 /**
- * Renders all messages in the current thread using the provided component configuration.
- *
- * This component automatically renders all messages in the thread, providing the appropriate
- * message context for each message. It handles different message types (user, assistant, system)
- * and supports editing mode through the provided edit composer components.
+ * Renders all messages in the current thread.
  *
  * @example
  * ```tsx
- * <ThreadPrimitive.Messages
- *   components={{
- *     UserMessage: MyUserMessage,
- *     AssistantMessage: MyAssistantMessage,
- *     EditComposer: MyEditComposer
+ * <ThreadPrimitive.Messages>
+ *   {({ message }) => {
+ *     if (message.role === "user") return <MyUserMessage />;
+ *     return <MyAssistantMessage />;
  *   }}
- * />
+ * </ThreadPrimitive.Messages>
  * ```
  */
 export const ThreadPrimitiveMessagesImpl: FC<ThreadPrimitiveMessages.Props> = ({
   components,
+  children,
 }) => {
-  const messagesLength = useAuiState((s) => s.thread.messages.length);
-
-  const messageElements = useMemo(() => {
-    if (messagesLength === 0) return null;
-    return Array.from({ length: messagesLength }, (_, index) => (
-      <ThreadPrimitiveMessageByIndex
-        key={index}
-        index={index}
-        components={components}
-      />
-    ));
-  }, [messagesLength, components]);
-
-  return messageElements;
+  if (components) {
+    return (
+      <ThreadPrimitiveMessagesInner>
+        {() => <ThreadMessageComponent components={components} />}
+      </ThreadPrimitiveMessagesInner>
+    );
+  }
+  return (
+    <ThreadPrimitiveMessagesInner>{children}</ThreadPrimitiveMessagesInner>
+  );
 };
 
 ThreadPrimitiveMessagesImpl.displayName = "ThreadPrimitive.Messages";
 
 export const ThreadPrimitiveMessages = memo(
   ThreadPrimitiveMessagesImpl,
-  (prev, next) => isComponentsSame(prev.components, next.components),
+  (prev, next) => {
+    if (prev.children || next.children) {
+      return prev.children === next.children;
+    }
+    return isComponentsSame(prev.components!, next.components!);
+  },
 );