@assistant-ui/react 0.15.24 → 0.15.26

package/src/primitives/composer/trigger/TriggerPopoverExports.test.ts

@@ -0,0 +1,61 @@
+import { describe, expect, expectTypeOf, it } from "vitest";
+import * as ComposerPrimitive from "../../composer";
+import {
+  defaultDirectiveFormatter,
+  unstable_defaultDirectiveFormatter,
+  type DirectiveFormatter,
+  type TriggerAdapter,
+  type TriggerCategory,
+  type TriggerItem,
+  type Unstable_DirectiveFormatter,
+  type Unstable_TriggerBehavior,
+  type Unstable_TriggerItem,
+} from "../../../index";
+import type { TriggerBehavior } from "./triggerSelectionResource";
+
+describe("stable trigger popover exports", () => {
+  it("keeps every old composer member as an alias of its stable member", () => {
+    expect(ComposerPrimitive.Unstable_TriggerPopover).toBe(
+      ComposerPrimitive.TriggerPopover,
+    );
+    expect(ComposerPrimitive.Unstable_TriggerPopover.Directive).toBe(
+      ComposerPrimitive.TriggerPopover.Directive,
+    );
+    expect(ComposerPrimitive.Unstable_TriggerPopover.Action).toBe(
+      ComposerPrimitive.TriggerPopover.Action,
+    );
+    expect(ComposerPrimitive.Unstable_TriggerPopoverRoot).toBe(
+      ComposerPrimitive.TriggerPopoverRoot,
+    );
+    expect(ComposerPrimitive.Unstable_TriggerPopoverCategories).toBe(
+      ComposerPrimitive.TriggerPopoverCategories,
+    );
+    expect(ComposerPrimitive.Unstable_TriggerPopoverCategoryItem).toBe(
+      ComposerPrimitive.TriggerPopoverCategoryItem,
+    );
+    expect(ComposerPrimitive.Unstable_TriggerPopoverItems).toBe(
+      ComposerPrimitive.TriggerPopoverItems,
+    );
+    expect(ComposerPrimitive.Unstable_TriggerPopoverItem).toBe(
+      ComposerPrimitive.TriggerPopoverItem,
+    );
+    expect(ComposerPrimitive.Unstable_TriggerPopoverBack).toBe(
+      ComposerPrimitive.TriggerPopoverBack,
+    );
+  });
+
+  it("keeps the old members' namespaced prop types", () => {
+    expectTypeOf<ComposerPrimitive.Unstable_TriggerPopoverItem.Props>().toEqualTypeOf<ComposerPrimitive.TriggerPopoverItem.Props>();
+    expectTypeOf<ComposerPrimitive.Unstable_TriggerPopoverBack.Element>().toEqualTypeOf<ComposerPrimitive.TriggerPopoverBack.Element>();
+  });
+
+  it("exports the stable adapter, item, category and formatter names beside the old ones", () => {
+    expect(unstable_defaultDirectiveFormatter).toBe(defaultDirectiveFormatter);
+    expectTypeOf<Unstable_TriggerItem>().toEqualTypeOf<TriggerItem>();
+    expectTypeOf<Unstable_DirectiveFormatter>().toEqualTypeOf<DirectiveFormatter>();
+    expectTypeOf<TriggerAdapter["categories"]>().returns.toEqualTypeOf<
+      readonly TriggerCategory[]
+    >();
+    expectTypeOf<Unstable_TriggerBehavior>().toEqualTypeOf<TriggerBehavior>();
+  });
+});

package/src/primitives/composer/trigger/TriggerPopoverItems.tsx

@@ -10,7 +10,7 @@ import {
 } from "react";
 import { composeEventHandlers } from "radix-ui/internal";
 import { useTriggerPopoverScopeContext } from "./TriggerPopover";
-import type { Unstable_TriggerItem } from "@assistant-ui/core";
+import type { TriggerItem } from "@assistant-ui/core";
 
 export namespace ComposerPrimitiveTriggerPopoverItems {
   export type Element = ComponentRef<typeof Primitive.div>;
@@ -18,7 +18,7 @@ export namespace ComposerPrimitiveTriggerPopoverItems {
     ComponentPropsWithoutRef<typeof Primitive.div>,
     "children"
   > & {
-    children: (items: readonly Unstable_TriggerItem[]) => ReactNode;
+    children: (items: readonly TriggerItem[]) => ReactNode;
   };
 }
 
@@ -53,7 +53,7 @@ ComposerPrimitiveTriggerPopoverItems.displayName =
 export namespace ComposerPrimitiveTriggerPopoverItem {
   export type Element = ComponentRef<typeof Primitive.button>;
   export type Props = ComponentPropsWithoutRef<typeof Primitive.button> & {
-    item: Unstable_TriggerItem;
+    item: TriggerItem;
     index?: number | undefined;
   };
 }

package/src/primitives/composer/trigger/TriggerPopoverRootContext.tsx

@@ -294,21 +294,21 @@ const TriggerPopoverRootInner: FC<
  *
  * @example
  * ```tsx
- * <ComposerPrimitive.Unstable_TriggerPopoverRoot>
- *   <ComposerPrimitive.Unstable_TriggerPopover char="@" adapter={mentionAdapter}>
- *     <ComposerPrimitive.Unstable_TriggerPopover.Directive formatter={formatter} />
+ * <ComposerPrimitive.TriggerPopoverRoot>
+ *   <ComposerPrimitive.TriggerPopover char="@" adapter={mentionAdapter}>
+ *     <ComposerPrimitive.TriggerPopover.Directive formatter={formatter} />
  *     ...
- *   </ComposerPrimitive.Unstable_TriggerPopover>
+ *   </ComposerPrimitive.TriggerPopover>
  *
- *   <ComposerPrimitive.Unstable_TriggerPopover char="/" adapter={slashAdapter}>
- *     <ComposerPrimitive.Unstable_TriggerPopover.Action onExecute={handler} />
+ *   <ComposerPrimitive.TriggerPopover char="/" adapter={slashAdapter}>
+ *     <ComposerPrimitive.TriggerPopover.Action onExecute={handler} />
  *     ...
- *   </ComposerPrimitive.Unstable_TriggerPopover>
+ *   </ComposerPrimitive.TriggerPopover>
  *
  *   <ComposerPrimitive.Root>
  *     <ComposerPrimitive.Input />
  *   </ComposerPrimitive.Root>
- * </ComposerPrimitive.Unstable_TriggerPopoverRoot>
+ * </ComposerPrimitive.TriggerPopoverRoot>
  * ```
  */
 export const ComposerPrimitiveTriggerPopoverRoot: FC<

package/src/primitives/composer/useComposerCancelWithFocus.ts

@@ -0,0 +1,57 @@
+"use client";
+
+import { useCallback } from "react";
+import { useAui } from "@assistant-ui/store";
+import { useThreadRootElementRef } from "../thread/ThreadRootElementContext";
+import { COMPOSER_INPUT_SELECTOR } from "./composerInputSelector";
+
+const COMPOSER_SELECTOR = "[data-aui-composer-type]";
+const THREAD_SELECTOR = "[data-aui-thread-root]";
+
+export const useComposerCancelWithFocus = (cancel: () => void) => {
+  const aui = useAui();
+  const threadRootRef = useThreadRootElementRef();
+
+  const cancelWithFocus = useCallback(
+    (source: HTMLElement | null) => {
+      const root = threadRootRef?.current;
+      const doc = source?.ownerDocument;
+      const active = doc?.activeElement;
+      const composer = source?.closest(COMPOSER_SELECTOR);
+      const shouldRestore =
+        aui.composer.getState().type === "edit" &&
+        root &&
+        active &&
+        composer?.getAttribute("data-aui-composer-type") === "edit" &&
+        composer.contains(active) &&
+        composer.closest(THREAD_SELECTOR) === root;
+      const candidates = shouldRestore
+        ? Array.from(
+            root.querySelectorAll<HTMLElement>(COMPOSER_INPUT_SELECTOR),
+          ).filter(
+            (input) =>
+              input.closest(THREAD_SELECTOR) === root &&
+              input
+                .closest(COMPOSER_SELECTOR)
+                ?.getAttribute("data-aui-composer-type") === "thread",
+          )
+        : [];
+
+      cancel();
+
+      if (!doc || !active || !shouldRestore) return;
+      if (doc.activeElement !== active && doc.activeElement !== doc.body)
+        return;
+
+      for (const input of candidates) {
+        if (!input.isConnected || input.closest("[hidden], [inert]")) continue;
+        input.focus({ preventScroll: true });
+        if (doc.activeElement !== active && doc.activeElement !== doc.body)
+          break;
+      }
+    },
+    [aui, cancel, threadRootRef],
+  );
+
+  return cancelWithFocus;
+};

package/src/primitives/composer.ts

@@ -1,3 +1,5 @@
+import type { TriggerMatch, TriggerMatcher } from "./composer/trigger";
+
 export { ComposerPrimitiveRoot as Root } from "./composer/ComposerRoot";
 export { ComposerPrimitiveInput as Input } from "./composer/ComposerInput";
 export { ComposerPrimitiveSend as Send } from "./composer/ComposerSend";
@@ -15,18 +17,71 @@ export { ComposerPrimitiveQuoteText as QuoteText } from "./composer/ComposerQuot
 export { ComposerPrimitiveQuoteDismiss as QuoteDismiss } from "./composer/ComposerQuote";
 export { ComposerPrimitiveQueue as Queue } from "./composer/ComposerQueue";
 
-export { ComposerPrimitiveTriggerPopover as Unstable_TriggerPopover } from "./composer/trigger";
-export { ComposerPrimitiveTriggerPopoverRoot as Unstable_TriggerPopoverRoot } from "./composer/trigger";
-export { ComposerPrimitiveTriggerPopoverCategories as Unstable_TriggerPopoverCategories } from "./composer/trigger";
-export { ComposerPrimitiveTriggerPopoverCategoryItem as Unstable_TriggerPopoverCategoryItem } from "./composer/trigger";
-export { ComposerPrimitiveTriggerPopoverItems as Unstable_TriggerPopoverItems } from "./composer/trigger";
-export { ComposerPrimitiveTriggerPopoverItem as Unstable_TriggerPopoverItem } from "./composer/trigger";
-export { ComposerPrimitiveTriggerPopoverBack as Unstable_TriggerPopoverBack } from "./composer/trigger";
-export { useTriggerPopoverRootContext as unstable_useTriggerPopoverRootContext } from "./composer/trigger";
-export { useTriggerPopoverRootContextOptional as unstable_useTriggerPopoverRootContextOptional } from "./composer/trigger";
-export { useTriggerPopoverScopeContext as unstable_useTriggerPopoverScopeContext } from "./composer/trigger";
-export { useTriggerPopoverScopeContextOptional as unstable_useTriggerPopoverScopeContextOptional } from "./composer/trigger";
-export { useTriggerPopoverTriggers as unstable_useTriggerPopoverTriggers } from "./composer/trigger";
-export { useTriggerPopoverTriggersOptional as unstable_useTriggerPopoverTriggersOptional } from "./composer/trigger";
-export type { RegisteredTrigger as Unstable_RegisteredTrigger } from "./composer/trigger";
-export type { TriggerMatcher as Unstable_TriggerMatcher } from "./composer/trigger";
+export { ComposerPrimitiveTriggerPopover as TriggerPopover } from "./composer/trigger";
+export {
+  /** @deprecated Use `ComposerPrimitive.TriggerPopover` instead. */
+  ComposerPrimitiveTriggerPopover as Unstable_TriggerPopover,
+} from "./composer/trigger";
+export { ComposerPrimitiveTriggerPopoverRoot as TriggerPopoverRoot } from "./composer/trigger";
+export {
+  /** @deprecated Use `ComposerPrimitive.TriggerPopoverRoot` instead. */
+  ComposerPrimitiveTriggerPopoverRoot as Unstable_TriggerPopoverRoot,
+} from "./composer/trigger";
+export { ComposerPrimitiveTriggerPopoverCategories as TriggerPopoverCategories } from "./composer/trigger";
+export {
+  /** @deprecated Use `ComposerPrimitive.TriggerPopoverCategories` instead. */
+  ComposerPrimitiveTriggerPopoverCategories as Unstable_TriggerPopoverCategories,
+} from "./composer/trigger";
+export { ComposerPrimitiveTriggerPopoverCategoryItem as TriggerPopoverCategoryItem } from "./composer/trigger";
+export {
+  /** @deprecated Use `ComposerPrimitive.TriggerPopoverCategoryItem` instead. */
+  ComposerPrimitiveTriggerPopoverCategoryItem as Unstable_TriggerPopoverCategoryItem,
+} from "./composer/trigger";
+export { ComposerPrimitiveTriggerPopoverItems as TriggerPopoverItems } from "./composer/trigger";
+export {
+  /** @deprecated Use `ComposerPrimitive.TriggerPopoverItems` instead. */
+  ComposerPrimitiveTriggerPopoverItems as Unstable_TriggerPopoverItems,
+} from "./composer/trigger";
+export { ComposerPrimitiveTriggerPopoverItem as TriggerPopoverItem } from "./composer/trigger";
+export {
+  /** @deprecated Use `ComposerPrimitive.TriggerPopoverItem` instead. */
+  ComposerPrimitiveTriggerPopoverItem as Unstable_TriggerPopoverItem,
+} from "./composer/trigger";
+export { ComposerPrimitiveTriggerPopoverBack as TriggerPopoverBack } from "./composer/trigger";
+export {
+  /** @deprecated Use `ComposerPrimitive.TriggerPopoverBack` instead. */
+  ComposerPrimitiveTriggerPopoverBack as Unstable_TriggerPopoverBack,
+} from "./composer/trigger";
+export {
+  /** @deprecated Experimental since 2026-04-15. Not scheduled for removal; the API may change in any release. */
+  useTriggerPopoverRootContext as unstable_useTriggerPopoverRootContext,
+} from "./composer/trigger";
+export {
+  /** @deprecated Experimental since 2026-04-15. Not scheduled for removal; the API may change in any release. */
+  useTriggerPopoverRootContextOptional as unstable_useTriggerPopoverRootContextOptional,
+} from "./composer/trigger";
+export {
+  /** @deprecated Experimental since 2026-04-15. Not scheduled for removal; the API may change in any release. */
+  useTriggerPopoverScopeContext as unstable_useTriggerPopoverScopeContext,
+} from "./composer/trigger";
+export {
+  /** @deprecated Experimental since 2026-04-15. Not scheduled for removal; the API may change in any release. */
+  useTriggerPopoverScopeContextOptional as unstable_useTriggerPopoverScopeContextOptional,
+} from "./composer/trigger";
+export {
+  /** @deprecated Experimental since 2026-04-15. Not scheduled for removal; the API may change in any release. */
+  useTriggerPopoverTriggers as unstable_useTriggerPopoverTriggers,
+} from "./composer/trigger";
+export {
+  /** @deprecated Experimental since 2026-04-15. Not scheduled for removal; the API may change in any release. */
+  useTriggerPopoverTriggersOptional as unstable_useTriggerPopoverTriggersOptional,
+} from "./composer/trigger";
+export type {
+  /** @deprecated Experimental since 2026-04-15. Not scheduled for removal; the API may change in any release. */
+  RegisteredTrigger as Unstable_RegisteredTrigger,
+} from "./composer/trigger";
+export type { TriggerMatch, TriggerMatcher } from "./composer/trigger";
+/** @deprecated Use `TriggerMatcher` instead. */
+export type Unstable_TriggerMatcher = TriggerMatcher;
+/** @deprecated Use `TriggerMatch` instead. */
+export type Unstable_TriggerMatch = TriggerMatch;

package/src/primitives/message/MessagePartsGrouped.test.tsx

@@ -12,10 +12,12 @@ import { describe, expect, it } from "vitest";
 import type { ThreadMessageLike } from "@assistant-ui/core";
 import {
   AssistantRuntimeProvider,
+  Tools,
   useAssistantDataUI,
   useExternalStoreRuntime,
 } from "@assistant-ui/core/react";
-import { useAui } from "@assistant-ui/store";
+import { AuiConfig, useAui } from "@assistant-ui/store";
+import { resource } from "@assistant-ui/tap";
 import { ThreadPrimitiveMessageByIndex } from "../thread/ThreadMessages";
 import {
   type MessagePrimitiveUnstable_PartsGrouped,
@@ -48,6 +50,11 @@ const partsMessage =
 const Named = () => <b>named</b>;
 const Fallback = () => <i>fallback</i>;
 const GlobalFallback = () => <b>global-fallback</b>;
+const Mcp = () => <b>mcp</b>;
+const McpApp = resource(function McpApp() {
+  return { render: Mcp };
+});
+const mcpConfig = AuiConfig({ tools: Tools({ mcpApp: McpApp() }) });
 
 const RegisterFallbackDataUI: FC<{ render: typeof GlobalFallback }> = ({
   render,
@@ -69,11 +76,13 @@ const Example = ({
   content,
   Message: MessageComponent = Message,
   extra,
+  config,
   isRunning = false,
 }: {
   content: ThreadMessageLike["content"];
   Message?: FC;
   extra?: ReactNode;
+  config?: AuiConfig;
   isRunning?: boolean;
 }) => {
   const messages: ThreadMessageLike[] = [
@@ -86,7 +95,10 @@ const Example = ({
     onNew: async () => {},
   });
   return (
-    <AssistantRuntimeProvider runtime={runtime}>
+    <AssistantRuntimeProvider
+      runtime={runtime}
+      {...(config === undefined ? {} : { config })}
+    >
       {extra}
       <ThreadPrimitiveMessageByIndex
         index={0}
@@ -439,6 +451,17 @@ describe("MessagePrimitive.Unstable_PartsGroupedByParentId", () => {
     },
   );
 
+  it("falls back to the default text component for an undefined slot", () => {
+    render(
+      <Example
+        content={[{ type: "text", text: "hello" }]}
+        Message={partsMessage({ Text: undefined })}
+      />,
+    );
+
+    expect(screen.getByText("hello")).toBeTruthy();
+  });
+
   it("renders tool and data UIs registered under an inherited name", () => {
     const { container } = render(
       <Example
@@ -461,6 +484,87 @@ describe("MessagePrimitive.Unstable_PartsGroupedByParentId", () => {
     expect(container.innerHTML).toBe("<b>named</b><b>named</b>");
   });
 
+  it("renders tools.mcpApp for a tool call with a ui:// resource", () => {
+    const { container } = render(
+      <Example
+        content={[
+          {
+            type: "tool-call",
+            toolCallId: "call",
+            toolName: "show_chart",
+            args: {},
+            mcp: { app: { resourceUri: "ui://chart" } },
+          },
+        ]}
+        Message={partsMessage({ tools: { Fallback } })}
+        config={mcpConfig}
+      />,
+    );
+
+    expect(container.innerHTML).toBe("<b>mcp</b>");
+  });
+
+  it("prefers an inline by_name component over tools.mcpApp", () => {
+    const { container } = render(
+      <Example
+        content={[
+          {
+            type: "tool-call",
+            toolCallId: "call",
+            toolName: "show_chart",
+            args: {},
+            mcp: { app: { resourceUri: "ui://chart" } },
+          },
+        ]}
+        Message={partsMessage({
+          tools: { by_name: { show_chart: Named }, Fallback },
+        })}
+        config={mcpConfig}
+      />,
+    );
+
+    expect(container.innerHTML).toBe("<b>named</b>");
+  });
+
+  it("uses inline Fallback when the tool call has no ui:// resource", () => {
+    const { container } = render(
+      <Example
+        content={[
+          {
+            type: "tool-call",
+            toolCallId: "call",
+            toolName: "show_chart",
+            args: {},
+          },
+        ]}
+        Message={partsMessage({ tools: { Fallback } })}
+        config={mcpConfig}
+      />,
+    );
+
+    expect(container.innerHTML).toBe("<i>fallback</i>");
+  });
+
+  it("uses inline Fallback when the resource URI is not an MCP App URI", () => {
+    const { container } = render(
+      <Example
+        content={[
+          {
+            type: "tool-call",
+            toolCallId: "call",
+            toolName: "show_chart",
+            args: {},
+            mcp: { app: { resourceUri: "https://example.com/chart" } },
+          },
+        ]}
+        Message={partsMessage({ tools: { Fallback } })}
+        config={mcpConfig}
+      />,
+    );
+
+    expect(container.innerHTML).toBe("<i>fallback</i>");
+  });
+
   it("uses dataRenderers.fallbacks[0] before inline data.Fallback", async () => {
     const { container } = render(
       <Example

package/src/primitives/message/MessagePartsGrouped.tsx

@@ -7,23 +7,23 @@ import {
   type PropsWithChildren,
   useMemo,
 } from "react";
-import { useAuiState, useAui } from "@assistant-ui/store";
+import { useAuiState } from "@assistant-ui/store";
 import { PartByIndexProvider } from "../../context/providers/PartByIndexProvider";
 import { TextMessagePartProvider } from "../../context/providers/TextMessagePartProvider";
 import { MessagePartPrimitiveText } from "../messagePart/MessagePartText";
 import { MessagePartPrimitiveImage } from "../messagePart/MessagePartImage";
-import type {
-  Unstable_AudioMessagePartComponent,
-  DataMessagePartComponent,
-  DataMessagePartProps,
-  EmptyMessagePartComponent,
-  TextMessagePartComponent,
-  ImageMessagePartComponent,
-  SourceMessagePartComponent,
-  ToolCallMessagePartComponent,
-  ToolCallMessagePartProps,
-  FileMessagePartComponent,
-  ReasoningMessagePartComponent,
+import {
+  MessagePartComponent,
+  type Unstable_AudioMessagePartComponent,
+  type DataMessagePartComponent,
+  type EmptyMessagePartComponent,
+  type TextMessagePartComponent,
+  type ImageMessagePartComponent,
+  type SourceMessagePartComponent,
+  type ToolCallMessagePartComponent,
+  type ToolCallMessagePartProps,
+  type FileMessagePartComponent,
+  type ReasoningMessagePartComponent,
 } from "@assistant-ui/core/react";
 import { MessagePartPrimitiveInProgress } from "../messagePart/MessagePartInProgress";
 import type { MessagePartStatus } from "@assistant-ui/core";
@@ -86,6 +86,13 @@ const useMessagePartsGrouped = (
   return { groups, partKeys: getMessagePartKeys(parts) };
 };
 
+/**
+ * @deprecated Prefer `<MessagePrimitive.GroupedParts>` for adjacent
+ * grouping — it dispatches all rendering through one `switch (part.type)`
+ * and supports nested group paths. Keep this primitive only for
+ * non-adjacent clustering (e.g., gathering parts with the same parent-id
+ * across the message).
+ */
 export namespace MessagePrimitiveUnstable_PartsGrouped {
   export type Props = {
     /**
@@ -231,34 +238,6 @@ export namespace MessagePrimitiveUnstable_PartsGrouped {
   };
 }
 
-const ToolUIDisplay = ({
-  Fallback,
-  ...props
-}: {
-  Fallback: ToolCallMessagePartComponent | undefined;
-} & ToolCallMessagePartProps) => {
-  const Render = useAuiState(
-    (s) => s.tools.toolUIs[props.toolName]?.[0]?.render ?? Fallback,
-  );
-  if (!Render) return null;
-  return <Render {...props} />;
-};
-
-const DataUIDisplay = ({
-  Fallback,
-  ...props
-}: {
-  Fallback: DataMessagePartComponent | undefined;
-} & DataMessagePartProps) => {
-  const Render = useAuiState((s) => {
-    const named = s.dataRenderers.renderers[props.name]?.[0];
-    if (named) return named;
-    return s.dataRenderers.fallbacks[0] ?? Fallback;
-  });
-  if (!Render) return null;
-  return <Render {...props} />;
-};
-
 const defaultComponents = {
   Text: () => (
     <p style={{ whiteSpace: "pre-line" }}>
@@ -280,104 +259,17 @@ type MessagePartComponentProps = {
   components: MessagePrimitiveUnstable_PartsGrouped.Props["components"];
 };
 
-const MessagePartComponent: FC<MessagePartComponentProps> = ({
-  components: {
-    Text = defaultComponents.Text,
-    Reasoning = defaultComponents.Reasoning,
-    Image = defaultComponents.Image,
-    Source = defaultComponents.Source,
-    File = defaultComponents.File,
-    Unstable_Audio: Audio = defaultComponents.Unstable_Audio,
-    tools = {},
-    data,
-  } = {},
-}) => {
-  const aui = useAui();
-  const part = useAuiState((s) => s.part);
-
-  const type = part.type;
-  if (type === "tool-call") {
-    const addResult = aui.part.addToolResult;
-    const resume = aui.part.resumeToolCall;
-    const respondToApproval = aui.part.respondToToolApproval;
-    const unstable_recordInteraction = aui.part.unstable_recordInteraction;
-    if ("Override" in tools)
-      return (
-        <tools.Override
-          {...part}
-          addResult={addResult}
-          resume={resume}
-          respondToApproval={respondToApproval}
-          {...(unstable_recordInteraction && { unstable_recordInteraction })}
-        />
-      );
-    const Tool =
-      (tools.by_name && Object.hasOwn(tools.by_name, part.toolName)
-        ? tools.by_name[part.toolName]
-        : undefined) ?? tools.Fallback;
-    return (
-      <ToolUIDisplay
-        {...part}
-        Fallback={Tool}
-        addResult={addResult}
-        resume={resume}
-        respondToApproval={respondToApproval}
-        {...(unstable_recordInteraction && { unstable_recordInteraction })}
-      />
-    );
-  }
-
-  if (part.status?.type === "requires-action")
-    throw new Error("Encountered unexpected requires-action status");
-
-  switch (type) {
-    case "text":
-      return <Text {...part} />;
-
-    case "reasoning":
-      return <Reasoning {...part} />;
-
-    case "source":
-      return <Source {...part} />;
-
-    case "image":
-      return <Image {...part} />;
-
-    case "file":
-      return <File {...part} />;
-
-    case "audio":
-      return <Audio {...part} />;
-
-    case "data": {
-      const Data =
-        (data?.by_name && Object.hasOwn(data.by_name, part.name)
-          ? data.by_name[part.name]
-          : undefined) ?? data?.Fallback;
-      return <DataUIDisplay {...part} Fallback={Data} />;
-    }
-
-    default:
-      console.warn(`Unknown message part type: ${type}`);
-      return null;
-  }
-};
-
 type MessagePartProps = {
   partIndex: number;
   components: MessagePrimitiveUnstable_PartsGrouped.Props["components"];
 };
 
-const MessagePartImpl: FC<MessagePartProps> = ({ partIndex, components }) => {
-  return (
+const MessagePart = memo(
+  ({ partIndex, components }: MessagePartProps) => (
     <PartByIndexProvider index={partIndex}>
       <MessagePartComponent components={components} />
     </PartByIndexProvider>
-  );
-};
-
-const MessagePart = memo(
-  MessagePartImpl,
+  ),
   (prev, next) =>
     prev.partIndex === next.partIndex &&
     prev.components?.Text === next.components?.Text &&
@@ -387,8 +279,7 @@ const MessagePart = memo(
     prev.components?.File === next.components?.File &&
     prev.components?.Unstable_Audio === next.components?.Unstable_Audio &&
     prev.components?.tools === next.components?.tools &&
-    prev.components?.data === next.components?.data &&
-    prev.components?.Group === next.components?.Group,
+    prev.components?.data === next.components?.data,
 );
 
 const EmptyPartFallback: FC<{
@@ -473,6 +364,17 @@ export const MessagePrimitiveUnstable_PartsGrouped: FC<
       return <EmptyParts components={components} />;
     }
 
+    const leafComponents = {
+      ...components,
+      Text: components?.Text ?? defaultComponents.Text,
+      Reasoning: components?.Reasoning ?? defaultComponents.Reasoning,
+      Image: components?.Image ?? defaultComponents.Image,
+      Source: components?.Source ?? defaultComponents.Source,
+      File: components?.File ?? defaultComponents.File,
+      Unstable_Audio:
+        components?.Unstable_Audio ?? defaultComponents.Unstable_Audio,
+    };
+
     return messageGroups.map((group, groupIndex) => {
       const GroupComponent = components?.Group ?? defaultComponents.Group;
       const identity = partKeys[group.indices[0]!];
@@ -491,7 +393,7 @@ export const MessagePrimitiveUnstable_PartsGrouped: FC<
             <MessagePart
               key={partKeys[partIndex]}
               partIndex={partIndex}
-              components={components}
+              components={leafComponents}
             />
           ))}
         </GroupComponent>