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. For the common case — grouping by part.type — 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/thread.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>
  );
};

Reach for an inline groupBy={(part) => …} only when the helper isn't expressive enough — e.g. you need to branch on part.toolName, part.parentId, or part metadata.

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; leaf cases should render their own UI or return null.

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={(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

Combine multiple text parts into cohesive blocks:

<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 = {
  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
  },
} satisfies Toolkit;

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 nothing is threaded 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>;

The "mcp-app" key is deprecated in favor of "standalone-tool-call", which is a superset (it also matches MCP-app tool calls). Existing "mcp-app" usage keeps working.

Group by Custom Metadata

Use any custom metadata in your parts for grouping:

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) => {
    const priority = part.metadata?.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 uppercase 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

When using assistant-stream libraries, you can add custom metadata to 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 with a different 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");

  await researchController.addToolCallPart({
    toolName: "search",
    args: { query: "climate data" },
  });

  // Add parts with custom metadata
  controller.appendPart({
    type: "text",
    text: "High priority finding",
    priority: "high",
    category: "findings",
  });
});

API Reference

MessagePrimitive.GroupedParts

MessagePrimitiveGroupedPartsProps

groupBy : (part: PartState) => readonly `group-${string}`[]: Maps a part to a group-key path. Return `[]` to leave the part ungrouped. Group keys must start with `group-`. Prefer `groupPartByType` for the common `part.type → path` case.
children : ({ part, children }) => ReactNode: Render function called for group nodes and leaf parts. Group nodes expose children; leaf parts should render their own UI or return null.

groupPartByType helper

groupPartByType builds a groupBy from a part.type → group-key path map. 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 indices: readonly number[];
};

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
          </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[]>();

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

    return Array.from(groups.entries()).map(([key, indices]) => ({
      groupKey: key.startsWith("__ungrouped_") ? undefined : key,
      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.

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 (experimental)
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