div]:pb-0", className)}
+ {...props}
+ />
+
+ );
+};
+
+export type MessageBranchContentProps = HTMLAttributes
;
+
+export const MessageBranchContent = ({
+ children,
+ ...props
+}: MessageBranchContentProps) => {
+ const { currentBranch, setBranches, branches } = useMessageBranch();
+ const childrenArray = useMemo(
+ () => (Array.isArray(children) ? 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]:pb-0",
+ index === currentBranch ? "block" : "hidden"
+ )}
+ key={branch.key}
+ {...props}
+ >
+ {branch}
+
+ ));
+};
+
+export type MessageBranchSelectorProps = ComponentProps;
+
+export const MessageBranchSelector = ({
+ className,
+ ...props
+}: MessageBranchSelectorProps) => {
+ const { totalBranches } = useMessageBranch();
+
+ // Don't render if there's only one branch
+ if (totalBranches <= 1) {
+ return null;
+ }
+
+ return (
+ *:not(:first-child)]:rounded-l-md [&>*:not(:last-child)]:rounded-r-md",
+ className
+ )}
+ orientation="horizontal"
+ {...props}
+ />
+ );
+};
+
+export type MessageBranchPreviousProps = ComponentProps;
+
+export const MessageBranchPrevious = ({
+ children,
+ ...props
+}: MessageBranchPreviousProps) => {
+ const { goToPrevious, totalBranches } = useMessageBranch();
+
+ return (
+
+ );
+};
+
+export type MessageBranchNextProps = ComponentProps;
+
+export const MessageBranchNext = ({
+ children,
+ ...props
+}: MessageBranchNextProps) => {
+ const { goToNext, totalBranches } = useMessageBranch();
+
+ return (
+
+ );
+};
+
+export type MessageBranchPageProps = HTMLAttributes;
+
+export const MessageBranchPage = ({
+ className,
+ ...props
+}: MessageBranchPageProps) => {
+ const { currentBranch, totalBranches } = useMessageBranch();
+
+ return (
+
+ {currentBranch + 1} of {totalBranches}
+
+ );
+};
+
+export type MessageResponseProps = ComponentProps;
+
+const streamdownPlugins = { cjk, code, math, mermaid };
+
+export const MessageResponse = memo(
+ ({ className, ...props }: MessageResponseProps) => (
+ *:first-child]:mt-0 [&>*:last-child]:mb-0",
+ className
+ )}
+ plugins={streamdownPlugins}
+ {...props}
+ />
+ ),
+ (prevProps, nextProps) =>
+ prevProps.children === nextProps.children &&
+ nextProps.isAnimating === prevProps.isAnimating
+);
+
+MessageResponse.displayName = "MessageResponse";
+
+export type MessageToolbarProps = ComponentProps<"div">;
+
+export const MessageToolbar = ({
+ className,
+ children,
+ ...props
+}: MessageToolbarProps) => (
+
+ {children}
+
+);
diff --git a/packages/ui/src/components/ai-elements/tool.tsx b/packages/ui/src/components/ai-elements/tool.tsx
new file mode 100644
index 0000000..ab99d50
--- /dev/null
+++ b/packages/ui/src/components/ai-elements/tool.tsx
@@ -0,0 +1,173 @@
+"use client";
+
+import { Badge } from "@code/ui/components/badge";
+import {
+ Collapsible,
+ CollapsibleContent,
+ CollapsibleTrigger,
+} from "@code/ui/components/collapsible";
+import { cn } from "@code/ui/lib/utils";
+import type { DynamicToolUIPart, ToolUIPart } from "ai";
+import {
+ CheckCircleIcon,
+ ChevronDownIcon,
+ CircleIcon,
+ ClockIcon,
+ WrenchIcon,
+ XCircleIcon,
+} from "lucide-react";
+import type { ComponentProps, ReactNode } from "react";
+import { isValidElement } from "react";
+
+import { CodeBlock } from "./code-block";
+
+export type ToolProps = ComponentProps;
+
+export const Tool = ({ className, ...props }: ToolProps) => (
+
+);
+
+export type ToolPart = ToolUIPart | DynamicToolUIPart;
+
+export type ToolHeaderProps = {
+ title?: string;
+ className?: string;
+} & (
+ | { type: ToolUIPart["type"]; state: ToolUIPart["state"]; toolName?: never }
+ | {
+ type: DynamicToolUIPart["type"];
+ state: DynamicToolUIPart["state"];
+ toolName: string;
+ }
+);
+
+const statusLabels: Record = {
+ "approval-requested": "Awaiting Approval",
+ "approval-responded": "Responded",
+ "input-available": "Running",
+ "input-streaming": "Pending",
+ "output-available": "Completed",
+ "output-denied": "Denied",
+ "output-error": "Error",
+};
+
+const statusIcons: Record = {
+ "approval-requested": ,
+ "approval-responded": ,
+ "input-available": ,
+ "input-streaming": ,
+ "output-available": ,
+ "output-denied": ,
+ "output-error": ,
+};
+
+export const getStatusBadge = (status: ToolPart["state"]) => (
+
+ {statusIcons[status]}
+ {statusLabels[status]}
+
+);
+
+export const ToolHeader = ({
+ className,
+ title,
+ type,
+ state,
+ toolName,
+ ...props
+}: ToolHeaderProps) => {
+ const derivedName =
+ type === "dynamic-tool" ? toolName : type.split("-").slice(1).join("-");
+
+ return (
+
+
+
+ {title ?? derivedName}
+ {getStatusBadge(state)}
+
+
+
+ );
+};
+
+export type ToolContentProps = ComponentProps;
+
+export const ToolContent = ({ className, ...props }: ToolContentProps) => (
+
+);
+
+export type ToolInputProps = ComponentProps<"div"> & {
+ input: ToolPart["input"];
+};
+
+export const ToolInput = ({ className, input, ...props }: ToolInputProps) => (
+
+);
+
+export type ToolOutputProps = ComponentProps<"div"> & {
+ output: ToolPart["output"];
+ errorText: ToolPart["errorText"];
+};
+
+export const ToolOutput = ({
+ className,
+ output,
+ errorText,
+ ...props
+}: ToolOutputProps) => {
+ if (!(output || errorText)) {
+ return null;
+ }
+
+ let Output = {output as ReactNode}
;
+
+ if (typeof output === "object" && !isValidElement(output)) {
+ Output = (
+
+ );
+ } else if (typeof output === "string") {
+ Output = ;
+ }
+
+ return (
+
+
+ {errorText ? "Error" : "Result"}
+
+
+ {errorText &&
{errorText}
}
+ {Output}
+
+
+
+ {label}
+ {state}
+
+);
+
+interface MobileChatToolCallProps extends ComponentProps<"div"> {
+ detail?: ReactNode;
+ icon?: ReactNode;
+ status: string;
+ tone?: "error" | "neutral" | "success";
+ toolName: string;
+}
+
+const MobileChatToolCall = ({
+ children,
+ className,
+ detail,
+ icon,
+ status,
+ tone = "neutral",
+ toolName,
+ ...props
+}: MobileChatToolCallProps) => {
+ let content: ReactNode;
+ if (children) {
+ content = {children}
;
+ } else if (detail) {
+ content = (
+
+ );
+ }
+
+ return (
+
+
+
+ {icon ?? }
+
+
+ {toolName}
+
+
+ {status}
+
+
+ {content}
+
+
+ {icon ?? }
+
+
+
+ {title}
+
+ {source ? (
+
+ {source}
+
+ ) : null}
+
+
+);
+
+interface MobileChatComposerDockProps extends ComponentProps<"div"> {
+ errorMessage?: string;
+ statusMessage?: ReactNode;
+}
+
+const MobileChatComposerDock = ({
+ children,
+ className,
+ errorMessage,
+ statusMessage,
+ ...props
+}: MobileChatComposerDockProps) => {
+ const message = errorMessage ?? statusMessage;
+
+ return (
+
+ {message ? (
+
+ {message}
+
+ ) : null}
+ {children}
+