@assistant-ui/core 0.1.5 → 0.1.8

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

@@ -1,23 +1,31 @@
-import type { CompleteAttachment } from "../../../types";
-import { ComponentType, type FC, memo, useMemo } from "react";
-import { useAuiState } from "@assistant-ui/store";
-import { MessageAttachmentByIndexProvider } from "../../providers";
+import type { CompleteAttachment } from "../../../types/attachment";
+import { ComponentType, type FC, type ReactNode, memo, useMemo } from "react";
+import { RenderChildrenWithAccessor, useAuiState } from "@assistant-ui/store";
+import { MessageAttachmentByIndexProvider } from "../../providers/AttachmentByIndexProvider";
+
+type MessageAttachmentsComponentConfig = {
+  Image?: ComponentType | undefined;
+  Document?: ComponentType | undefined;
+  File?: ComponentType | undefined;
+  Attachment?: ComponentType | undefined;
+};
 
 export namespace MessagePrimitiveAttachments {
-  export type Props = {
-    components:
-      | {
-          Image?: ComponentType | undefined;
-          Document?: ComponentType | undefined;
-          File?: ComponentType | undefined;
-          Attachment?: ComponentType | undefined;
-        }
-      | undefined;
-  };
+  export type Props =
+    | {
+        /** @deprecated Use the children render function instead. */
+        components: MessageAttachmentsComponentConfig;
+        children?: never;
+      }
+    | {
+        /** Render function called for each attachment. Receives the attachment. */
+        children: (value: { attachment: CompleteAttachment }) => ReactNode;
+        components?: never;
+      };
 }
 
 const getComponent = (
-  components: MessagePrimitiveAttachments.Props["components"],
+  components: MessageAttachmentsComponentConfig | undefined,
   attachment: CompleteAttachment,
 ) => {
   const type = attachment.type;
@@ -34,7 +42,7 @@ const getComponent = (
 };
 
 const AttachmentComponent: FC<{
-  components: MessagePrimitiveAttachments.Props["components"];
+  components: MessageAttachmentsComponentConfig | undefined;
 }> = ({ components }) => {
   const attachment = useAuiState((s) => s.attachment);
   if (!attachment) return null;
@@ -47,7 +55,7 @@ const AttachmentComponent: FC<{
 export namespace MessagePrimitiveAttachmentByIndex {
   export type Props = {
     index: number;
-    components?: MessagePrimitiveAttachments.Props["components"];
+    components?: MessageAttachmentsComponentConfig;
   };
 }
 
@@ -74,25 +82,56 @@ export const MessagePrimitiveAttachmentByIndex: FC<MessagePrimitiveAttachmentByI
 MessagePrimitiveAttachmentByIndex.displayName =
   "MessagePrimitive.AttachmentByIndex";
 
-export const MessagePrimitiveAttachments: FC<
-  MessagePrimitiveAttachments.Props
-> = ({ components }) => {
+const MessagePrimitiveAttachmentsInner: FC<{
+  children: (value: { attachment: CompleteAttachment }) => ReactNode;
+}> = ({ children }) => {
   const attachmentsCount = useAuiState((s) => {
     if (s.message.role !== "user") return 0;
     return s.message.attachments.length;
   });
 
-  const attachmentElements = useMemo(() => {
-    return Array.from({ length: attachmentsCount }, (_, index) => (
-      <MessagePrimitiveAttachmentByIndex
-        key={index}
-        index={index}
-        components={components}
-      />
-    ));
-  }, [attachmentsCount, components]);
+  return useMemo(
+    () =>
+      Array.from({ length: attachmentsCount }, (_, index) => (
+        <MessageAttachmentByIndexProvider key={index} index={index}>
+          <RenderChildrenWithAccessor
+            getItemState={(aui) =>
+              aui.message().attachment({ index }).getState()
+            }
+          >
+            {(getItem) =>
+              children({
+                get attachment() {
+                  return getItem() as CompleteAttachment;
+                },
+              })
+            }
+          </RenderChildrenWithAccessor>
+        </MessageAttachmentByIndexProvider>
+      )),
+    [attachmentsCount, children],
+  );
+};
 
-  return attachmentElements;
+export const MessagePrimitiveAttachments: FC<
+  MessagePrimitiveAttachments.Props
+> = ({ components, children }) => {
+  if (components) {
+    return (
+      <MessagePrimitiveAttachmentsInner>
+        {({ attachment }) => {
+          const Component = getComponent(components, attachment);
+          if (!Component) return null;
+          return <Component />;
+        }}
+      </MessagePrimitiveAttachmentsInner>
+    );
+  }
+  return (
+    <MessagePrimitiveAttachmentsInner>
+      {children}
+    </MessagePrimitiveAttachmentsInner>
+  );
 };
 
 MessagePrimitiveAttachments.displayName = "MessagePrimitive.Attachments";

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);