tool: enable custom icons

apps/docs/content/components/(chatbot)/tool.mdx

@@ -176,6 +176,7 @@ export async function POST(req: Request) {
 ## Features
 
 - Collapsible interface for showing/hiding tool details
+- Customizable tool icons, with support for mapping icons by tool name
 - Visual status indicators with icons and badges
 - Support for multiple tool execution states (pending, running, completed, error)
 - Formatted parameter display with JSON syntax highlighting
@@ -211,6 +212,19 @@ Shows a tool that encountered an error during execution. Opens by default to dis
 
 <Preview path="tool-output-error" />
 
+### Custom Icons
+
+Pass a function to the `icon` prop to render a custom icon. The function receives an object with:
+
+- `type`: the tool part type
+- `state`: the current execution state
+- `toolName`: the derived tool name
+- `className`: default sizing and color styles
+
+Use `toolName` to map different icons by tool type. When the function returns `null`, the default wrench icon is used.
+
+<Preview path="tool-custom-icons" />
+
 ## Props
 
 ### `<Tool />`
@@ -229,6 +243,10 @@ Shows a tool that encountered an error during execution. Opens by default to dis
 
 <TypeTable
   type={{
+    icon: {
+      description: "Custom icon for the tool header. Pass a ReactNode for a static icon, or a function that receives a ToolIconProps object and returns a ReactNode.",
+      type: "ReactNode | ((props: ToolIconProps) => ReactNode)",
+    },
     title: {
       description: "Custom title to display instead of the derived tool name.",
       type: "string",

packages/elements/src/tool.tsx

@@ -1,13 +1,15 @@
 "use client";
 
+import type { DynamicToolUIPart, ToolUIPart } from "ai";
+import type { ComponentProps, ReactNode } from "react";
+
 import { Badge } from "@repo/shadcn-ui/components/ui/badge";
 import {
   Collapsible,
   CollapsibleContent,
   CollapsibleTrigger,
 } from "@repo/shadcn-ui/components/ui/collapsible";
 import { cn } from "@repo/shadcn-ui/lib/utils";
-import type { DynamicToolUIPart, ToolUIPart } from "ai";
 import {
   CheckCircleIcon,
   ChevronDownIcon,
@@ -16,7 +18,6 @@ import {
   WrenchIcon,
   XCircleIcon,
 } from "lucide-react";
-import type { ComponentProps, ReactNode } from "react";
 import { isValidElement } from "react";
 
 import { CodeBlock } from "./code-block";
@@ -32,7 +33,15 @@ export const Tool = ({ className, ...props }: ToolProps) => (
 
 export type ToolPart = ToolUIPart | DynamicToolUIPart;
 
+export type ToolIconProps = {
+  type: ToolPart["type"];
+  state: ToolPart["state"];
+  toolName: string;
+  className: string;
+};
+
 export type ToolHeaderProps = {
+  icon?: ReactNode | ((props: ToolIconProps) => ReactNode);
   title?: string;
   className?: string;
 } & (
@@ -73,6 +82,7 @@ export const getStatusBadge = (status: ToolPart["state"]) => (
 
 export const ToolHeader = ({
   className,
+  icon,
   title,
   type,
   state,
@@ -82,6 +92,12 @@ export const ToolHeader = ({
   const derivedName =
     type === "dynamic-tool" ? toolName : type.split("-").slice(1).join("-");
 
+  const iconClassName = "size-4 shrink-0 text-muted-foreground";
+  const resolvedIcon =
+    typeof icon === "function"
+      ? icon({ type, state, toolName: derivedName, className: iconClassName })
+      : icon;
+
   return (
     <CollapsibleTrigger
       className={cn(
@@ -91,7 +107,7 @@ export const ToolHeader = ({
       {...props}
     >
       <div className="flex items-center gap-2">
-        <WrenchIcon className="size-4 text-muted-foreground" />
+        {resolvedIcon ?? <WrenchIcon className={iconClassName} />}
         <span className="font-medium text-sm">{title ?? derivedName}</span>
         {getStatusBadge(state)}
       </div>

packages/examples/src/tool-custom-icons.tsx

@@ -0,0 +1,77 @@
+"use client";
+
+import type { ReactNode } from "react";
+import { Tool, ToolContent, ToolHeader, ToolInput } from "@repo/elements/tool";
+import { cn } from "@repo/shadcn-ui/lib/utils";
+import { DatabaseIcon, SearchIcon, WrenchIcon } from "lucide-react";
+import type { LucideProps } from "lucide-react";
+
+const toolIcons: Record<string, (props: LucideProps) => ReactNode> = {
+  database_query: DatabaseIcon,
+  web_search: SearchIcon,
+};
+
+const renderIcon = ({ toolName, className }: { toolName: string; className: string }) => {
+  const Icon = toolIcons[toolName];
+  return Icon ? <Icon className={className} /> : null;
+};
+
+const Example = () => (
+  <div className="space-y-4">
+    <Tool>
+      <ToolHeader
+        icon={renderIcon}
+        state="output-available"
+        title="database_query"
+        type="tool-database_query"
+      />
+      <ToolContent>
+        <ToolInput
+          input={{
+            query: "SELECT COUNT(*) FROM users WHERE created_at >= ?",
+            params: ["2024-01-01"],
+          }}
+        />
+      </ToolContent>
+    </Tool>
+    <Tool>
+      <ToolHeader
+        icon={renderIcon}
+        state="output-available"
+        title="web_search"
+        type="tool-web_search"
+      />
+      <ToolContent>
+        <ToolInput input={{ query: "latest AI news" }} />
+      </ToolContent>
+    </Tool>
+    <Tool>
+      <ToolHeader
+        icon={<WrenchIcon className="size-4 shrink-0 text-foreground" />}
+        state="input-available"
+        title="custom_style"
+        type="tool-custom_style"
+      />
+      <ToolContent>
+        <ToolInput
+          input={{
+            query: "SELECT COUNT(*) FROM large_table",
+          }}
+        />
+      </ToolContent>
+    </Tool>
+    <Tool>
+      <ToolHeader
+        icon={renderIcon}
+        state="output-available"
+        title="unknown_tool"
+        type="tool-unknown_tool"
+      />
+      <ToolContent>
+        <ToolInput input={{ foo: "bar" }} />
+      </ToolContent>
+    </Tool>
+  </div>
+);
+
+export default Example;