Message Part Grouping

Organize message parts into custom groups with flexible grouping functions.

Message parts grouped by type:

Search: climate data

Search: renewable energy stats

Analyzing trends...

Generating summary...

Basic Usage

For adjacent grouping, use MessagePrimitive.GroupedParts. To group by part.type, the common case, pass groupPartByType to map each part type to a group-key path. Adjacent parts that share a group key are coalesced and rendered through the group case.

/components/assistant-ui/elements/thread.aui.tsx

import { MessagePrimitive, groupPartByType } from "@assistant-ui/react";

const AssistantMessage: FC = () => {
  return (
    <MessagePrimitive.Root className="...">
      <div className="...">
        <MessagePrimitive.GroupedParts
          groupBy={groupPartByType({
            "tool-call": ["group-tool"],
          })}
        >
          {({ part, children }) => {
            switch (part.type) {
              case "group-tool":
                return <div className="group">{children}</div>;
              case "tool-call":
                return part.toolUI ?? null;
              default:
                return null;
            }
          }}
        </MessagePrimitive.GroupedParts>
      </div>
      <AssistantActionBar />
      <BranchPicker className="..." />
    </MessagePrimitive.Root>
  );
};

Use an inline groupBy={(part) => …} only when the helper cannot express the grouping, for example when you branch on part.parentId or on a data part's payload.

How Adjacent Grouping Works

groupBy receives each part and returns either null for an ungrouped leaf or a group-key path such as ["group-chainOfThought", "group-tool"]. Group keys must start with "group-" so your render function can distinguish synthetic groups from real part types.

The render function receives both group nodes and leaf parts through { part, children }. Only group cases should render children. For a leaf, returning null renders its registered tool or data UI when one exists. Return an empty fragment (<></>) when you want to suppress a registered UI, or return your own component for a custom leaf renderer. Text and other unregistered leaves render nothing when null is returned.

A group wrapper uses its group name and its first part's unique identity. If the first part has no unique identity, the wrapper uses its structural sibling path. Appending a tool call to a group that starts with anonymous reasoning keeps the wrapper mounted. Reordering later identified parts within the group preserves their state; changing the first member's identity remounts the group. Unstable_PartsGrouped follows the same first-member rule, with the group's index as its structural fallback.

Warning

MessagePrimitive.Unstable_PartsGrouped remains available for rare non-adjacent clustering where one group needs to collect parts from different positions in the message. For normal consecutive reasoning/tool grouping, use GroupedParts.

Use Cases & Examples

Group by Parent ID

Group adjacent content that shares the same parent relationship:

import { useState, type PropsWithChildren } from "react";
import { ChevronDownIcon, ChevronUpIcon } from "lucide-react";
import { MessagePrimitive } from "@assistant-ui/react";

function ParentGroup({
  id,
  count,
  children,
}: PropsWithChildren<{ id: string; count: number }>) {
  const [collapsed, setCollapsed] = useState(false);

  return (
    <div className="my-2 overflow-hidden rounded-lg border">
      <button
        type="button"
        onClick={() => setCollapsed((value) => !value)}
        className="hover:bg-muted/50 flex w-full items-center justify-between p-3"
      >
        <span>
          Group {id} ({count} items)
        </span>
        {collapsed ? <ChevronDownIcon /> : <ChevronUpIcon />}
      </button>
      {!collapsed && <div className="border-t p-3">{children}</div>}
    </div>
  );
}

<MessagePrimitive.GroupedParts
  groupBy={(part) => {
    if (!part.parentId) return [];
    return [`group-parent-${part.parentId}`];
  }}
>
  {({ part, children }) => {
    if (part.type.startsWith("group-parent-")) {
      const id = part.type.replace("group-parent-", "");
      return (
        <ParentGroup id={id} count={part.indices.length}>
          {children}
        </ParentGroup>
      );
    }

    if (part.type === "text") return <MarkdownText />;
    if (part.type === "tool-call") return part.toolUI ?? <ToolFallback {...part} />;
    return null;
  }}
</MessagePrimitive.GroupedParts>

Group by Tool Name

Organize adjacent tool calls by tool name:

<MessagePrimitive.GroupedParts
  groupBy={groupPartByType({
    "tool-call:task": ["group-subagents"],
    "tool-call": ["group-tool"],
  })}
>
  {({ part, children }) => {
    switch (part.type) {
      case "group-subagents":
        return (
          <div className="tool-group my-2 rounded-lg border">
            <div className="bg-muted/50 px-4 py-2 text-sm font-medium">
              Task calls ({part.indices.length} calls)
            </div>
            <div className="p-4">{children}</div>
          </div>
        );
      case "group-tool":
        return <div className="tool-group my-2 rounded-lg border p-4">{children}</div>;
      case "text":
        return <MarkdownText />;
      case "tool-call":
        return part.toolUI ?? <ToolFallback {...part} />;
      default:
        return null;
    }
  }}
</MessagePrimitive.GroupedParts>

Use an inline groupBy when the group key must be derived dynamically from the tool name or other part data:

<MessagePrimitive.GroupedParts
  groupBy={(part) => {
    if (part.type !== "tool-call") return [];
    return [`group-tool-${part.toolName}`];
  }}
>
  {({ part, children }) => {
    if (part.type.startsWith("group-tool-")) {
      const toolName = part.type.replace("group-tool-", "");
      return (
        <div className="tool-group my-2 rounded-lg border">
          <div className="bg-muted/50 px-4 py-2 text-sm font-medium">
            Tool: {toolName} ({part.indices.length} calls)
          </div>
          <div className="p-4">{children}</div>
        </div>
      );
    }

    if (part.type === "text") return <MarkdownText />;
    if (part.type === "tool-call") return part.toolUI ?? <ToolFallback {...part} />;
    return null;
  }}
</MessagePrimitive.GroupedParts>

Group Consecutive Text Parts

Wrap adjacent text parts in one block:

<MessagePrimitive.GroupedParts
  groupBy={groupPartByType({
    text: ["group-text-block"],
  })}
>
  {({ part, children }) => {
    switch (part.type) {
      case "group-text-block":
        return (
          <div className="prose prose-sm my-2 rounded-lg bg-gray-50 p-4">
            {children}
          </div>
        );
      case "text":
        return <MarkdownText />;
      case "tool-call":
        return part.toolUI ?? <ToolFallback {...part} />;
      default:
        return null;
    }
  }}
</MessagePrimitive.GroupedParts>

Group by Content Type

Separate different types of content for distinct visual treatment:

<MessagePrimitive.GroupedParts
  groupBy={groupPartByType({
    text: ["group-content-text"],
    "tool-call": ["group-content-tools"],
    reasoning: ["group-content-reasoning"],
  })}
>
  {({ part, children }) => {
    switch (part.type) {
      case "group-content-text":
        return <div className="space-y-2">{children}</div>;
      case "group-content-tools":
        return <div className="my-2 rounded-lg border p-3">{children}</div>;
      case "group-content-reasoning":
        return <div className="my-2 text-muted-foreground">{children}</div>;
      case "text":
        return <MarkdownText />;
      case "reasoning":
        return <Reasoning {...part} />;
      case "tool-call":
        return part.toolUI ?? <ToolFallback {...part} />;
      default:
        return null;
    }
  }}
</MessagePrimitive.GroupedParts>

Keep Standalone Tool UIs Out of the Trace

Tool UIs fall into three buckets: prompting the user (human-in-the-loop), informing the user (generative UI), and traces of what the model is doing (routine tool calls). The first two should be surfaced on their own, while the last belongs folded into a collapsible chain-of-thought group.

Mark a tool with display: "standalone" to keep its UI out of the grouped trace. human tools and MCP apps are standalone automatically; every other tool defaults to "inline" and opts in explicitly:

const toolkit = defineToolkit({
  ask_user: { type: "human", render: AskUI }, // standalone (forced)
  search_web: { type: "frontend", render: SearchUI }, // inline trace (default)
  checkout: {
    type: "frontend",
    render: CheckoutUI,
    display: "standalone", // opt in
  },
});

The synthetic "standalone-tool-call" key on groupPartByType matches all of these. MessagePrimitive.GroupedParts passes the live tool-UI registry to groupBy as a second context argument, and the helper reads it to resolve the registry-driven cases. MCP-app calls are detected from the part alone, so you pass nothing else in:

<MessagePrimitive.GroupedParts
  groupBy={groupPartByType({
    reasoning: ["group-chainOfThought", "group-reasoning"],
    "tool-call": ["group-chainOfThought", "group-tool"],
    "standalone-tool-call": [], // ungrouped — rendered on its own
  })}
>
  {({ part, children }) => {
    switch (part.type) {
      case "group-chainOfThought":
        return <div className="chain-of-thought">{children}</div>;
      case "tool-call":
        return part.toolUI ?? <ToolFallback {...part} />;
      // ...
      default:
        return null;
    }
  }}
</MessagePrimitive.GroupedParts>;

Render Tool Calls as Flat Rows

The previous section keeps "tool-call" folded into the chain-of-thought and only lifts standalone tool UIs out. To keep reasoning collapsed in a group while every tool call renders as its own row in the message flow, map both tool keys to []:

<MessagePrimitive.GroupedParts
  groupBy={groupPartByType({
    reasoning: ["group-thought", "group-reasoning"],
    "tool-call": [],
    "standalone-tool-call": [],
  })}
>
  {({ part, children }) => {
    switch (part.type) {
      case "group-thought":
      case "group-reasoning":
        return <ReasoningBlock>{children}</ReasoningBlock>;
      case "reasoning":
        return <Reasoning {...part} />;
      case "tool-call":
        return part.toolUI ?? <ToolFallback {...part} />;
      case "text":
        return <MarkdownText />;
      case "indicator":
        return <LoadingDots />;
      default:
        return null;
    }
  }}
</MessagePrimitive.GroupedParts>

Adjacent reasoning parts still coalesce into one group-thought, while each tool call is a leaf at the same level as text. The "standalone-tool-call" entry is mapped to [] too so human-in-the-loop and generative-UI tools stay flat alongside routine calls instead of falling through to "tool-call". Tool-call leaves keep their stable identity keys across re-renders, so rows don't remount as the message streams.

Warning

This renders tool calls flat in the message flow: each leaf sits where the part appears, in order. It does not pin, portal, float, or otherwise move tool UIs out of the flow into a sidebar or fixed region. For out-of-flow placement you still own the layout outside GroupedParts.

For MCP Apps, keep one host-owned container mounted for the lifetime of the app. Switching between an inline element and a portal, or changing a portal's target, remounts the frame. Reparenting its live iframe with appendChild or insertBefore also reloads the document and loses in-app state.

In React hosts, create and insert the movable container imperatively under a stable host element; do not render it as a JSX child. Moving a React-owned DOM node behind React's back can make later reconciliation fail.

Feature-detect moveBefore on Element.prototype, not Node.prototype. When available, use it to move the container between connected parents without reloading the frame. Both the container and destination parent must be connected before the move; otherwise moveBefore throws HierarchyRequestError. Insert the container normally on its first mount. In browsers without moveBefore, leave the container in place and reposition it with CSS, or treat relocation as a reload and reinitialize the app.

Group by Custom Metadata

Message parts have no free-form metadata field. Carry custom values in a data part and group on its payload:

const priorityStyles = {
  high: "border-red-500 bg-red-50",
  normal: "border-gray-300 bg-white",
  low: "border-gray-200 bg-gray-50",
};

<MessagePrimitive.GroupedParts
  groupBy={(part) => {
    if (part.type !== "data") return [];
    const priority = part.data?.priority;
    if (!priority) return [];
    return [`group-priority-${priority}`];
  }}
>
  {({ part, children }) => {
    if (part.type.startsWith("group-priority-")) {
      const priority = part.type.replace("group-priority-", "");
      return (
        <div
          className={`my-2 rounded-lg border-2 p-4 ${priorityStyles[priority] || ""}`}
        >
          <div className="mb-2 text-xs font-semibold text-gray-600">
            {priority} Priority
          </div>
          {children}
        </div>
      );
    }

    if (part.type === "text") return <MarkdownText />;
    if (part.type === "tool-call") return part.toolUI ?? <ToolFallback {...part} />;
    return null;
  }}
</MessagePrimitive.GroupedParts>

Integration with Assistant Streams

With the assistant-stream libraries, you can set a parent ID on parts and send custom values in data parts:

Python (assistant-stream)

from assistant_stream import create_run

async def my_run(controller):
    # Add parts with custom parentId
    research_controller = controller.with_parent_id("research-123")

    tool = await research_controller.add_tool_call("search")
    tool.append_args_text('{"query": "climate data"}')
    tool.set_response("climate data results")
    research_controller.append_text("Key findings from the research:")

    # Add text without a parent_id
    controller.append_text("High priority finding")

TypeScript (assistant-stream)

import { createAssistantStream } from "assistant-stream";

const stream = createAssistantStream(async (controller) => {
  // Add parts with parentId
  const researchController = controller.withParentId("research-123");

  researchController.addToolCallPart({
    toolName: "search",
    args: { query: "climate data" },
    response: { result: "climate data results" },
  });
  researchController.appendText("Key findings from the research:");

  // Add a data part with custom values
  controller.appendData({
    type: "data",
    name: "finding",
    data: {
      text: "High priority finding",
      priority: "high",
      category: "findings",
    },
  });
});

API Reference

MessagePrimitive.GroupedParts

MessagePrimitiveGroupedPartsProps

groupBy(part: PartState, context: GroupByContext) => readonly `group-${string}`[] | null: Maps a part to a group-key path. Return `[]` or `null` to leave the part ungrouped. Group keys must start with `group-`. Prefer `groupPartByType` for the common `part.type → path` case.
indicator"never" | "empty" | "no-text" | "always"= "no-text": When to call the render function with a trailing `indicator` part. It is only emitted while the message is running: `empty` when the message has no parts yet, `no-text` when it has no parts or the last part is not text or reasoning, `always` whenever it is running.
children({ part, children }) => ReactNode: Render function called for group nodes and leaf parts. Group nodes expose children. Returning null for a tool or data leaf renders its registered UI; an empty fragment suppresses it. Other leaves render nothing when null is returned.

groupPartByType helper

groupPartByType builds a groupBy from a map of part.type, or tool-call:<name> for one tool, to a group-key path. Part types missing from the map are left ungrouped. The returned function carries an internal memo fingerprint so the tree survives unrelated re-renders.

import { groupPartByType } from "@assistant-ui/react";

groupPartByType({
  reasoning: ["group-thought", "group-reasoning"],
  "tool-call": ["group-thought", "group-tool"],
});

GroupPart Type

type GroupPart<TKey extends `group-${string}` = `group-${string}`> = {
  readonly type: TKey;
  readonly status: MessagePartStatus | ToolCallMessagePartStatus;
  readonly counts: GroupCounts;
  readonly indices: readonly number[];
  readonly timing?: MessagePartTiming;
};

type GroupCounts = {
  readonly running: number;
  readonly complete: number;
  readonly incomplete: number;
  readonly requiresAction: number;
};

timing spans the group's timed parts (reasoning and tool calls that carry timing): the earliest startedAt, plus the latest completedAt once every timed part has finished. It is absent when no part in the group carries timing.

Best Practices

Keep grouping local: GroupedParts groups adjacent runs. If the same key appears again later, it becomes a new group in that position.
Handle ungrouped parts: Always include leaf cases for the part types your message can render.
Only render children for groups: Leaf parts receive a sentinel children value that throws if rendered accidentally.
Use group- prefixes: Synthetic group keys must start with group- so they cannot collide with real part types.

Common Patterns

Nested Grouping

Create hierarchical adjacent groups by returning multiple keys:

<MessagePrimitive.GroupedParts
  groupBy={(part) => {
    if (part.type === "reasoning")
      return ["group-chainOfThought", "group-reasoning"];
    if (part.type === "tool-call")
      return ["group-chainOfThought", `group-tool-${part.toolName}`];
    return [];
  }}
>
  {({ part, children }) => {
    if (part.type === "group-chainOfThought") return <div>{children}</div>;
    if (part.type === "group-reasoning") return <ReasoningRoot>{children}</ReasoningRoot>;
    if (part.type.startsWith("group-tool-")) return <div>{children}</div>;
    if (part.type === "reasoning") return <Reasoning {...part} />;
    if (part.type === "tool-call") return part.toolUI ?? <ToolFallback {...part} />;
    if (part.type === "text") return <MarkdownText />;
    return null;
  }}
</MessagePrimitive.GroupedParts>

Dynamic Group Rendering

Adjust group appearance based on the group node's status or indices:

<MessagePrimitive.GroupedParts
  groupBy={groupPartByType({
    "tool-call": ["group-tools"],
  })}
>
  {({ part, children }) => {
    if (part.type === "group-tools") {
      const isRunning = part.status.type === "running";

      return (
        <div
          className={`my-2 rounded-lg border p-4 ${
            isRunning ? "animate-pulse border-blue-500" : ""
          }`}
        >
          <div className="mb-2 text-sm text-muted-foreground">
            {part.indices.length} tool calls ({part.counts.running} running, {part.counts.complete} complete)
          </div>
          {children}
        </div>
      );
    }

    if (part.type === "tool-call") return part.toolUI ?? <ToolFallback {...part} />;
    if (part.type === "text") return <MarkdownText />;
    return null;
  }}
</MessagePrimitive.GroupedParts>

Non-adjacent grouping (Unstable)

Use MessagePrimitive.Unstable_PartsGrouped only when you need to collect non-adjacent parts into the same rendered group, such as gathering every part with the same parent ID even when other parts appear between them. This API is unstable and should not be used for normal consecutive reasoning or tool-call grouping.

<MessagePrimitive.Unstable_PartsGrouped
  groupingFunction={(parts) => {
    const groups = new Map<string | number, number[]>();

    parts.forEach((part, index) => {
      const key = part.parentId ?? index;
      const indices = groups.get(key) ?? [];
      indices.push(index);
      groups.set(key, indices);
    });

    return Array.from(groups.entries()).map(([key, indices]) => ({
      groupKey: typeof key === "string" ? key : undefined,
      indices,
    }));
  }}
  components={{
    Text: MarkdownText,
    tools: { Fallback: ToolFallback },
    Group: ({ groupKey, children }) => {
      if (!groupKey) return <>{children}</>;
      return <div className="rounded-lg border p-3">{children}</div>;
    },
  }}
/>

MessagePrimitiveUnstable_PartsGroupedProps

groupingFunction(parts: readonly PartState[]) => MessagePartGroup[]

Function that takes all message parts and returns non-adjacent group descriptors.

components?object

Legacy component map used to render message part types and group wrappers.

Components

Empty?EmptyMessagePartComponent: Component for rendering empty messages
Text?TextMessagePartComponent: Component for rendering text content
Reasoning?ReasoningMessagePartComponent: Component for rendering reasoning content (typically hidden)
Source?SourceMessagePartComponent: Component for rendering source content
Image?ImageMessagePartComponent: Component for rendering image content
File?FileMessagePartComponent: Component for rendering file content
Unstable_Audio?Unstable_AudioMessagePartComponent: Component for rendering audio content (deprecated, render audio/* from File)
tools?object | { Override: ComponentType }: Configuration for tool call rendering. Can be an object with by_name map and Fallback component, or an Override component.
Group?ComponentType<PropsWithChildren<{ groupKey: string | undefined; indices: number[] }>>: Component for rendering grouped message parts. Receives groupKey, indices array, and children to render.

MessagePartGroup Type

type MessagePartGroup = {
  groupKey: string | undefined; // The group identifier (undefined for ungrouped parts)
  indices: number[]; // Array of part indices belonging to this group
};

Group Component Props

The Group component receives:

groupKey: The group identifier (or undefined for ungrouped parts)
indices: Array of indices for the parts in this group
children: The rendered message part components