Branch

Manages multiple versions of AI messages, allowing users to navigate between different response branches.

The Branch component manages multiple versions of AI messages, allowing users to navigate between different response branches. It provides a clean, modern interface with customizable themes and keyboard-accessible navigation buttons.

What are the key strategies for optimizing React performance?

Here's the first response to your question. This approach focuses on performance optimization.

"use client";

import {
  Branch,
  BranchMessages,
  BranchNext,
  BranchPage,
  BranchPrevious,
  BranchSelector,
} from "@/components/ai-elements/elements/branch";
import { Message, MessageAvatar, MessageContent } from "@/components/ai-elements/elements/message";
import { nanoid } from "nanoid";

const userMessages = [
  {
    id: nanoid(),
    content: "What are the key strategies for optimizing React performance?",
  },
  {
    id: nanoid(),
    content: "How can I improve the performance of my React application?",
  },
  {
    id: nanoid(),
    content: "What performance optimization techniques should I use in React?",
  },
];

const assistantMessages = [
  {
    id: nanoid(),
    content:
      "Here's the first response to your question. This approach focuses on performance optimization.",
  },
  {
    id: nanoid(),
    content:
      "Here's an alternative response. This approach emphasizes code readability and maintainability over pure performance.",
  },
  {
    id: nanoid(),
    content:
      "And here's a third option. This balanced approach considers both performance and maintainability, making it suitable for most use cases.",
  },
];

const Example = () => {
  const handleBranchChange = (branchIndex: number) => {
    console.log("Branch changed to:", branchIndex);
  };

  return (
    <div style={{ height: "300px" }}>
      <Branch defaultBranch={0} onBranchChange={handleBranchChange}>
        <BranchMessages>
          {userMessages.map((message) => (
            <Message from="user" key={message.id}>
              <MessageContent>{message.content}</MessageContent>
              <MessageAvatar
                name="Hayden Bleasel"
                src="https://github.com/haydenbleasel.png"
              />
            </Message>
          ))}
        </BranchMessages>
        <BranchSelector from="user">
          <BranchPrevious />
          <BranchPage />
          <BranchNext />
        </BranchSelector>
      </Branch>

      <Branch defaultBranch={0} onBranchChange={handleBranchChange}>
        <BranchMessages>
          {assistantMessages.map((message) => (
            <Message from="assistant" key={message.id}>
              <MessageContent>{message.content}</MessageContent>
              <MessageAvatar name="AI" src="https://github.com/openai.png" />
            </Message>
          ))}
        </BranchMessages>
        <BranchSelector from="assistant">
          <BranchPrevious />
          <BranchPage />
          <BranchNext />
        </BranchSelector>
      </Branch>
    </div>
  );
};

export default Example;

"use client";import {  Branch,  BranchMessages,  BranchNext,  BranchPage,  BranchPrevious,  BranchSelector,} from "@/components/ai-elements/elements/branch";import { Message, MessageAvatar, MessageContent } from "@/components/ai-elements/elements/message";import { nanoid } from "nanoid";const userMessages = [  {    id: nanoid(),    content: "What are the key strategies for optimizing React performance?",  },  {    id: nanoid(),    content: "How can I improve the performance of my React application?",  },  {    id: nanoid(),    content: "What performance optimization techniques should I use in React?",  },];const assistantMessages = [  {    id: nanoid(),    content:      "Here's the first response to your question. This approach focuses on performance optimization.",  },  {    id: nanoid(),    content:      "Here's an alternative response. This approach emphasizes code readability and maintainability over pure performance.",  },  {    id: nanoid(),    content:      "And here's a third option. This balanced approach considers both performance and maintainability, making it suitable for most use cases.",  },];const Example = () => {  const handleBranchChange = (branchIndex: number) => {    console.log("Branch changed to:", branchIndex);  };  return (    <div style={{ height: "300px" }}>      <Branch defaultBranch={0} onBranchChange={handleBranchChange}>        <BranchMessages>          {userMessages.map((message) => (            <Message from="user" key={message.id}>              <MessageContent>{message.content}</MessageContent>              <MessageAvatar                name="Hayden Bleasel"                src="https://github.com/haydenbleasel.png"              />            </Message>          ))}        </BranchMessages>        <BranchSelector from="user">          <BranchPrevious />          <BranchPage />          <BranchNext />        </BranchSelector>      </Branch>      <Branch defaultBranch={0} onBranchChange={handleBranchChange}>        <BranchMessages>          {assistantMessages.map((message) => (            <Message from="assistant" key={message.id}>              <MessageContent>{message.content}</MessageContent>              <MessageAvatar name="AI" src="https://github.com/openai.png" />            </Message>          ))}        </BranchMessages>        <BranchSelector from="assistant">          <BranchPrevious />          <BranchPage />          <BranchNext />        </BranchSelector>      </Branch>    </div>  );};export default Example;

Installation

npx ai-elements@latest add branch

npx ai-elements@latest add branch

npx shadcn@latest add @ai-elements/branch

npx shadcn@latest add @ai-elements/branch

"use client";

import { Button } from "@repo/shadcn-ui/components/ui/button";
import { cn } from "@repo/shadcn-ui/lib/utils";
import type { UIMessage } from "ai";
import { ChevronLeftIcon, ChevronRightIcon } from "lucide-react";
import type { ComponentProps, HTMLAttributes, ReactElement } from "react";
import { createContext, useContext, useEffect, useState } from "react";

type BranchContextType = {
  currentBranch: number;
  totalBranches: number;
  goToPrevious: () => void;
  goToNext: () => void;
  branches: ReactElement[];
  setBranches: (branches: ReactElement[]) => void;
};

const BranchContext = createContext<BranchContextType | null>(null);

const useBranch = () => {
  const context = useContext(BranchContext);

  if (!context) {
    throw new Error("Branch components must be used within Branch");
  }

  return context;
};

export type BranchProps = HTMLAttributes<HTMLDivElement> & {
  defaultBranch?: number;
  onBranchChange?: (branchIndex: number) => void;
};

export const Branch = ({
  defaultBranch = 0,
  onBranchChange,
  className,
  ...props
}: BranchProps) => {
  const [currentBranch, setCurrentBranch] = useState(defaultBranch);
  const [branches, setBranches] = useState<ReactElement[]>([]);

  const handleBranchChange = (newBranch: number) => {
    setCurrentBranch(newBranch);
    onBranchChange?.(newBranch);
  };

  const goToPrevious = () => {
    const newBranch =
      currentBranch > 0 ? currentBranch - 1 : branches.length - 1;
    handleBranchChange(newBranch);
  };

  const goToNext = () => {
    const newBranch =
      currentBranch < branches.length - 1 ? currentBranch + 1 : 0;
    handleBranchChange(newBranch);
  };

  const contextValue: BranchContextType = {
    currentBranch,
    totalBranches: branches.length,
    goToPrevious,
    goToNext,
    branches,
    setBranches,
  };

  return (
    <BranchContext.Provider value={contextValue}>
      <div
        className={cn("grid w-full gap-2 [&>div]:pb-0", className)}
        {...props}
      />
    </BranchContext.Provider>
  );
};

export type BranchMessagesProps = HTMLAttributes<HTMLDivElement>;

export const BranchMessages = ({ children, ...props }: BranchMessagesProps) => {
  const { currentBranch, setBranches, branches } = useBranch();
  const childrenArray = Array.isArray(children) ? children : [children];

  // Use useEffect to update branches when they change
  useEffect(() => {
    if (branches.length !== childrenArray.length) {
      setBranches(childrenArray);
    }
  }, [childrenArray, branches, setBranches]);

  return childrenArray.map((branch, index) => (
    <div
      className={cn(
        "grid gap-2 overflow-hidden [&>div]:pb-0",
        index === currentBranch ? "block" : "hidden"
      )}
      key={branch.key}
      {...props}
    >
      {branch}
    </div>
  ));
};

export type BranchSelectorProps = HTMLAttributes<HTMLDivElement> & {
  from: UIMessage["role"];
};

export const BranchSelector = ({
  className,
  from,
  ...props
}: BranchSelectorProps) => {
  const { totalBranches } = useBranch();

  // Don't render if there's only one branch
  if (totalBranches <= 1) {
    return null;
  }

  return (
    <div
      className={cn(
        "flex items-center gap-2 self-end px-10",
        from === "assistant" ? "justify-start" : "justify-end",
        className
      )}
      {...props}
    />
  );
};

export type BranchPreviousProps = ComponentProps<typeof Button>;

export const BranchPrevious = ({
  className,
  children,
  ...props
}: BranchPreviousProps) => {
  const { goToPrevious, totalBranches } = useBranch();

  return (
    <Button
      aria-label="Previous branch"
      className={cn(
        "size-7 shrink-0 rounded-full text-muted-foreground transition-colors",
        "hover:bg-accent hover:text-foreground",
        "disabled:pointer-events-none disabled:opacity-50",
        className
      )}
      disabled={totalBranches <= 1}
      onClick={goToPrevious}
      size="icon"
      type="button"
      variant="ghost"
      {...props}
    >
      {children ?? <ChevronLeftIcon size={14} />}
    </Button>
  );
};

export type BranchNextProps = ComponentProps<typeof Button>;

export const BranchNext = ({
  className,
  children,
  ...props
}: BranchNextProps) => {
  const { goToNext, totalBranches } = useBranch();

  return (
    <Button
      aria-label="Next branch"
      className={cn(
        "size-7 shrink-0 rounded-full text-muted-foreground transition-colors",
        "hover:bg-accent hover:text-foreground",
        "disabled:pointer-events-none disabled:opacity-50",
        className
      )}
      disabled={totalBranches <= 1}
      onClick={goToNext}
      size="icon"
      type="button"
      variant="ghost"
      {...props}
    >
      {children ?? <ChevronRightIcon size={14} />}
    </Button>
  );
};

export type BranchPageProps = HTMLAttributes<HTMLSpanElement>;

export const BranchPage = ({ className, ...props }: BranchPageProps) => {
  const { currentBranch, totalBranches } = useBranch();

  return (
    <span
      className={cn(
        "font-medium text-muted-foreground text-xs tabular-nums",
        className
      )}
      {...props}
    >
      {currentBranch + 1} of {totalBranches}
    </span>
  );
};

"use client";import { Button } from "@repo/shadcn-ui/components/ui/button";import { cn } from "@repo/shadcn-ui/lib/utils";import type { UIMessage } from "ai";import { ChevronLeftIcon, ChevronRightIcon } from "lucide-react";import type { ComponentProps, HTMLAttributes, ReactElement } from "react";import { createContext, useContext, useEffect, useState } from "react";type BranchContextType = {  currentBranch: number;  totalBranches: number;  goToPrevious: () => void;  goToNext: () => void;  branches: ReactElement[];  setBranches: (branches: ReactElement[]) => void;};const BranchContext = createContext<BranchContextType | null>(null);const useBranch = () => {  const context = useContext(BranchContext);  if (!context) {    throw new Error("Branch components must be used within Branch");  }  return context;};export type BranchProps = HTMLAttributes<HTMLDivElement> & {  defaultBranch?: number;  onBranchChange?: (branchIndex: number) => void;};export const Branch = ({  defaultBranch = 0,  onBranchChange,  className,  ...props}: BranchProps) => {  const [currentBranch, setCurrentBranch] = useState(defaultBranch);  const [branches, setBranches] = useState<ReactElement[]>([]);  const handleBranchChange = (newBranch: number) => {    setCurrentBranch(newBranch);    onBranchChange?.(newBranch);  };  const goToPrevious = () => {    const newBranch =      currentBranch > 0 ? currentBranch - 1 : branches.length - 1;    handleBranchChange(newBranch);  };  const goToNext = () => {    const newBranch =      currentBranch < branches.length - 1 ? currentBranch + 1 : 0;    handleBranchChange(newBranch);  };  const contextValue: BranchContextType = {    currentBranch,    totalBranches: branches.length,    goToPrevious,    goToNext,    branches,    setBranches,  };  return (    <BranchContext.Provider value={contextValue}>      <div        className={cn("grid w-full gap-2 [&>div]:pb-0", className)}        {...props}      />    </BranchContext.Provider>  );};export type BranchMessagesProps = HTMLAttributes<HTMLDivElement>;export const BranchMessages = ({ children, ...props }: BranchMessagesProps) => {  const { currentBranch, setBranches, branches } = useBranch();  const childrenArray = Array.isArray(children) ? children : [children];  // Use useEffect to update branches when they change  useEffect(() => {    if (branches.length !== childrenArray.length) {      setBranches(childrenArray);    }  }, [childrenArray, branches, setBranches]);  return childrenArray.map((branch, index) => (    <div      className={cn(        "grid gap-2 overflow-hidden [&>div]:pb-0",        index === currentBranch ? "block" : "hidden"      )}      key={branch.key}      {...props}    >      {branch}    </div>  ));};export type BranchSelectorProps = HTMLAttributes<HTMLDivElement> & {  from: UIMessage["role"];};export const BranchSelector = ({  className,  from,  ...props}: BranchSelectorProps) => {  const { totalBranches } = useBranch();  // Don't render if there's only one branch  if (totalBranches <= 1) {    return null;  }  return (    <div      className={cn(        "flex items-center gap-2 self-end px-10",        from === "assistant" ? "justify-start" : "justify-end",        className      )}      {...props}    />  );};export type BranchPreviousProps = ComponentProps<typeof Button>;export const BranchPrevious = ({  className,  children,  ...props}: BranchPreviousProps) => {  const { goToPrevious, totalBranches } = useBranch();  return (    <Button      aria-label="Previous branch"      className={cn(        "size-7 shrink-0 rounded-full text-muted-foreground transition-colors",        "hover:bg-accent hover:text-foreground",        "disabled:pointer-events-none disabled:opacity-50",        className      )}      disabled={totalBranches <= 1}      onClick={goToPrevious}      size="icon"      type="button"      variant="ghost"      {...props}    >      {children ?? <ChevronLeftIcon size={14} />}    </Button>  );};export type BranchNextProps = ComponentProps<typeof Button>;export const BranchNext = ({  className,  children,  ...props}: BranchNextProps) => {  const { goToNext, totalBranches } = useBranch();  return (    <Button      aria-label="Next branch"      className={cn(        "size-7 shrink-0 rounded-full text-muted-foreground transition-colors",        "hover:bg-accent hover:text-foreground",        "disabled:pointer-events-none disabled:opacity-50",        className      )}      disabled={totalBranches <= 1}      onClick={goToNext}      size="icon"      type="button"      variant="ghost"      {...props}    >      {children ?? <ChevronRightIcon size={14} />}    </Button>  );};export type BranchPageProps = HTMLAttributes<HTMLSpanElement>;export const BranchPage = ({ className, ...props }: BranchPageProps) => {  const { currentBranch, totalBranches } = useBranch();  return (    <span      className={cn(        "font-medium text-muted-foreground text-xs tabular-nums",        className      )}      {...props}    >      {currentBranch + 1} of {totalBranches}    </span>  );};

Usage

import {
  Branch,
  BranchMessages,
  BranchNext,
  BranchPage,
  BranchPrevious,
  BranchSelector,
} from '@/components/ai-elements/branch';

<Branch defaultBranch={0}>
  <BranchMessages>
    <Message from="user">
      <MessageContent>Hello</MessageContent>
    </Message>
    <Message from="user">
      <MessageContent>Hi!</MessageContent>
    </Message>
  </BranchMessages>
  <BranchSelector from="user">
    <BranchPrevious />
    <BranchPage />
    <BranchNext />
  </BranchSelector>
</Branch>

Usage with AI SDK

Branching is an advanced use case that you can implement yourself to suit your application's needs. While the AI SDK does not provide built-in support for branching, you have full flexibility to design and manage multiple response paths as required.

Features

Context-based state management for multiple message branches
Navigation controls for moving between branches (previous/next)
Uses CSS to prevent re-rendering of branches when switching
Branch counter showing current position (e.g., "1 of 3")
Automatic branch tracking and synchronization
Callbacks for branch change and navigation using onBranchChange
Support for custom branch change callbacks
Responsive design with mobile-friendly controls
Clean, modern styling with customizable themes
Keyboard-accessible navigation buttons

Installation

Usage

Usage with AI SDK

Features

Props

`<Branch />`

`<BranchMessages />`

`<BranchSelector />`

`<BranchPrevious />`

`<BranchNext />`

`<BranchPage />`

On this page