diff --git a/src/components/ai-elements/code-block.tsx b/src/components/ai-elements/code-block.tsx
index 09cbda4..e001b3c 100644
--- a/src/components/ai-elements/code-block.tsx
+++ b/src/components/ai-elements/code-block.tsx
@@ -1,6 +1,6 @@
 "use client";
 
-import type { ComponentProps, CSSProperties, HTMLAttributes } from "react";
+import type { ComponentProps, CSSProperties, HTMLAttributes, ReactNode } from "react";
 import type {
   BundledLanguage,
   BundledTheme,
@@ -20,8 +20,20 @@ import {
 import { cn } from "@/lib/utils";
 import { useThemeFamily } from "@/lib/theme/context";
 import { resolveShikiTheme, resolveShikiThemes, SHIKI_DEFAULT_LIGHT, SHIKI_DEFAULT_DARK } from "@/lib/theme/code-themes";
-import { Check, Copy } from "@phosphor-icons/react";
+import type { Icon } from "@phosphor-icons/react";
 import {
+  Check,
+  Copy,
+  CaretDown,
+  CaretUp,
+  FileCode,
+  Terminal,
+  Code,
+  File,
+  Hash,
+} from "@phosphor-icons/react";
+import {
+  createElement,
   createContext,
   memo,
   useCallback,
@@ -33,6 +45,24 @@ import {
 } from "react";
 import { createHighlighter } from "shiki";
 
+// ── Collapse/expand constants ──────────────────────────────────────────
+const COLLAPSE_THRESHOLD = 20;
+const VISIBLE_LINES = 10;
+
+// ── Terminal language detection ────────────────────────────────────────
+const TERMINAL_LANGUAGES = new Set(["bash", "sh", "shell", "terminal", "zsh", "console"]);
+
+// ── Language icon mapping ──────────────────────────────────────────────
+function getLanguageIcon(language: string): Icon {
+  const lower = language.toLowerCase();
+  if (TERMINAL_LANGUAGES.has(lower)) return Terminal;
+  if (["typescript", "tsx", "javascript", "jsx"].includes(lower)) return Code;
+  if (["json", "yaml", "yml", "toml", "xml"].includes(lower)) return Code;
+  if (["python", "ruby", "go", "rust", "java", "c", "cpp"].includes(lower)) return Hash;
+  if (["css", "scss", "html"].includes(lower)) return File;
+  return FileCode;
+}
+
 // Shiki uses bitflags for font styles: 1=italic, 2=bold, 4=underline
 // biome-ignore lint/suspicious/noBitwiseOperators: shiki bitflag check
  
@@ -104,8 +134,9 @@ const LineSpan = ({
 // Types
 type CodeBlockProps = HTMLAttributes<HTMLDivElement> & {
   code: string;
-  language: BundledLanguage;
+  language: BundledLanguage | (string & {});
   showLineNumbers?: boolean;
+  filename?: string;
 };
 
 interface TokenizedCode {
@@ -116,11 +147,13 @@ interface TokenizedCode {
 
 interface CodeBlockContextType {
   code: string;
+  language: string;
 }
 
 // Context
 const CodeBlockContext = createContext<CodeBlockContextType>({
   code: "",
+  language: "text",
 });
 
 // Highlighter cache keyed by "lang:lightTheme:darkTheme"
@@ -411,48 +444,140 @@ export const CodeBlockContent = ({
   code,
   language,
   showLineNumbers = false,
+  collapsible = false,
+  maxCollapsedLines = VISIBLE_LINES,
 }: {
   code: string;
   language: BundledLanguage;
   showLineNumbers?: boolean;
+  collapsible?: boolean;
+  maxCollapsedLines?: number;
 }) => {
   const { light: lightTheme, dark: darkTheme } = useShikiThemes();
+  const [expanded, setExpanded] = useState(false);
+  const codeContainerRef = useRef<HTMLDivElement>(null);
+  const [animatingHeight, setAnimatingHeight] = useState<string | undefined>(undefined);
+
+  const lines = useMemo(() => code.split("\n"), [code]);
+  const totalLines = lines.length;
+  const isCollapsible = collapsible && totalLines > COLLAPSE_THRESHOLD;
+
+  const displayCode = useMemo(() => {
+    if (!isCollapsible || expanded) return code;
+    return lines.slice(0, maxCollapsedLines).join("\n");
+  }, [code, lines, isCollapsible, expanded, maxCollapsedLines]);
 
   // Memoized raw tokens for immediate display
-  const rawTokens = useMemo(() => createRawTokens(code), [code]);
+  const rawTokens = useMemo(() => createRawTokens(displayCode), [displayCode]);
 
   // Try to get cached result synchronously, otherwise use raw tokens
   const syncTokenized = useMemo(
-    () => highlightCode(code, language, undefined, lightTheme, darkTheme) ?? rawTokens,
-    [code, language, rawTokens, lightTheme, darkTheme]
+    () => highlightCode(displayCode, language, undefined, lightTheme, darkTheme) ?? rawTokens,
+    [displayCode, language, rawTokens, lightTheme, darkTheme]
   );
 
   // Track async highlighting results keyed by code+language+themes to avoid stale state
   const [asyncResult, setAsyncResult] = useState<{ key: string; tokens: TokenizedCode } | null>(null);
-  const resultKey = `${code}:${language}:${lightTheme}:${darkTheme}`;
+  const resultKey = `${displayCode}:${language}:${lightTheme}:${darkTheme}`;
 
   useEffect(() => {
     let cancelled = false;
 
     // Subscribe to async highlighting result
-    highlightCode(code, language, (result) => {
+    highlightCode(displayCode, language, (result) => {
       if (!cancelled) {
-        setAsyncResult({ key: `${code}:${language}:${lightTheme}:${darkTheme}`, tokens: result });
+        setAsyncResult({ key: `${displayCode}:${language}:${lightTheme}:${darkTheme}`, tokens: result });
       }
     }, lightTheme, darkTheme);
 
     return () => {
       cancelled = true;
     };
-  }, [code, language, lightTheme, darkTheme]);
+  }, [displayCode, language, lightTheme, darkTheme]);
 
   // Only use async result if it matches the current code+language+themes
   const tokenized = (asyncResult && asyncResult.key === resultKey) ? asyncResult.tokens : syncTokenized;
 
+  const isTerminal = TERMINAL_LANGUAGES.has(language.toLowerCase());
+
+  const handleToggleExpand = () => {
+    const container = codeContainerRef.current;
+    if (!container) {
+      setExpanded(!expanded);
+      return;
+    }
+    const currentHeight = container.scrollHeight;
+    if (!expanded) {
+      setAnimatingHeight(`${currentHeight}px`);
+      setExpanded(true);
+      requestAnimationFrame(() => {
+        const fullHeight = container.scrollHeight;
+        setAnimatingHeight(`${fullHeight}px`);
+        setTimeout(() => setAnimatingHeight(undefined), 300);
+      });
+    } else {
+      setAnimatingHeight(`${currentHeight}px`);
+      requestAnimationFrame(() => {
+        const collapsedH = maxCollapsedLines * 1.5 + 1.5;
+        setAnimatingHeight(`${collapsedH}rem`);
+        setTimeout(() => {
+          setExpanded(false);
+          setAnimatingHeight(undefined);
+        }, 300);
+      });
+    }
+  };
+
   return (
-    <div className="relative overflow-auto">
-      <CodeBlockBody showLineNumbers={showLineNumbers} tokenized={tokenized} />
-    </div>
+    <>
+      <div
+        ref={codeContainerRef}
+        className="relative transition-[max-height] duration-300 ease-in-out overflow-hidden"
+        style={{
+          maxHeight: animatingHeight ?? (!isCollapsible || expanded ? undefined : `${maxCollapsedLines * 1.5 + 1.5}rem`),
+        }}
+      >
+        <div className="relative overflow-auto">
+          <CodeBlockBody showLineNumbers={showLineNumbers} tokenized={tokenized} />
+        </div>
+
+        {/* Gradient overlay for collapsed state */}
+        {isCollapsible && !expanded && (
+          <div className={cn(
+            "absolute bottom-0 left-0 right-0 h-16 pointer-events-none",
+            isTerminal
+              ? "bg-gradient-to-t from-[#0a0a0a] to-transparent"
+              : "bg-gradient-to-t from-muted to-transparent"
+          )} />
+        )}
+      </div>
+
+      {/* Expand/Collapse button */}
+      {isCollapsible && (
+        <button
+          onClick={handleToggleExpand}
+          type="button"
+          className={cn(
+            "flex w-full items-center justify-center gap-1.5 py-1.5 text-xs transition-colors",
+            isTerminal
+              ? "bg-zinc-950 text-zinc-400 hover:text-zinc-200"
+              : "bg-muted text-muted-foreground hover:text-foreground"
+          )}
+        >
+          {expanded ? (
+            <>
+              <CaretUp size={12} />
+              <span>Collapse</span>
+            </>
+          ) : (
+            <>
+              <CaretDown size={12} />
+              <span>Expand all {totalLines} lines</span>
+            </>
+          )}
+        </button>
+      )}
+    </>
   );
 };
 
@@ -460,26 +585,164 @@ export const CodeBlock = ({
   code,
   language,
   showLineNumbers = false,
+  filename,
   className,
   children,
   ...props
 }: CodeBlockProps) => {
-  const contextValue = useMemo(() => ({ code }), [code]);
+  const contextValue = useMemo(() => ({ code, language }), [code, language]);
+  const isTerminal = TERMINAL_LANGUAGES.has(language.toLowerCase());
+
+  // When children are provided, use the composable API (caller controls header).
+  // Otherwise, render a default header with language icon, copy, copy-as-markdown.
+  const hasCustomChildren = children != null;
 
   return (
     <CodeBlockContext.Provider value={contextValue}>
-      <CodeBlockContainer className={className} language={language} {...props}>
-        {children}
+      <CodeBlockContainer
+        className={cn(
+          hasCustomChildren ? undefined : "not-prose my-3",
+          isTerminal && !hasCustomChildren && "border-zinc-700/50",
+          className,
+        )}
+        language={language}
+        {...props}
+      >
+        {hasCustomChildren ? (
+          children
+        ) : (
+          <CodeBlockDefaultHeader
+            language={language}
+            filename={filename}
+            isTerminal={isTerminal}
+          />
+        )}
         <CodeBlockContent
           code={code}
-          language={language}
+          language={language as BundledLanguage}
           showLineNumbers={showLineNumbers}
+          collapsible={!hasCustomChildren}
         />
       </CodeBlockContainer>
     </CodeBlockContext.Provider>
   );
 };
 
+/** Default header rendered when CodeBlock has no children (non-composable usage). */
+const CodeBlockDefaultHeader = ({
+  language,
+  filename,
+  isTerminal,
+}: {
+  language: string;
+  filename?: string;
+  isTerminal: boolean;
+}) => {
+  const { code: contextCode, language: contextLanguage } = useContext(CodeBlockContext);
+  const [copied, setCopied] = useState(false);
+  const [copiedMarkdown, setCopiedMarkdown] = useState(false);
+
+  const handleCopy = async () => {
+    try {
+      await navigator.clipboard.writeText(contextCode);
+      setCopied(true);
+      setTimeout(() => setCopied(false), 2000);
+    } catch {
+      // clipboard not available
+    }
+  };
+
+  const handleCopyMarkdown = async () => {
+    try {
+      const markdown = `\`\`\`${contextLanguage}\n${contextCode}\n\`\`\``;
+      await navigator.clipboard.writeText(markdown);
+      setCopiedMarkdown(true);
+      setTimeout(() => setCopiedMarkdown(false), 2000);
+    } catch {
+      // clipboard not available
+    }
+  };
+
+  const langIcon = getLanguageIcon(language);
+
+  return (
+    <div className={cn(
+      "flex items-center justify-between px-4 py-1.5 text-xs border-b",
+      isTerminal
+        ? "bg-zinc-950 text-zinc-400"
+        : "bg-muted text-muted-foreground"
+    )}>
+      <div className="flex items-center gap-2 min-w-0">
+        {createElement(langIcon, { size: 14, className: cn(
+          "shrink-0",
+          isTerminal ? "text-green-400" : "text-muted-foreground",
+        ) })}
+        {filename && (
+          <span className={cn(
+            "truncate font-medium",
+            isTerminal ? "text-zinc-300" : "text-foreground"
+          )}>{filename}</span>
+        )}
+        {filename && <span className="text-muted-foreground/50">|</span>}
+        <span className={cn(
+          "rounded px-1.5 py-0.5",
+          isTerminal
+            ? "bg-zinc-700/50 text-green-400"
+            : "bg-accent text-accent-foreground"
+        )}>{language.toUpperCase()}</span>
+      </div>
+      <div className="flex items-center gap-1 ml-2 shrink-0">
+        <button
+          onClick={handleCopy}
+          type="button"
+          className={cn(
+            "flex items-center gap-1 rounded px-1.5 py-0.5 transition-colors",
+            isTerminal
+              ? "text-zinc-400 hover:text-zinc-200 hover:bg-zinc-700/50"
+              : "text-muted-foreground hover:text-foreground hover:bg-accent"
+          )}
+          title="Copy code"
+        >
+          {copied ? (
+            <>
+              <Check size={12} />
+              <span>Copied</span>
+            </>
+          ) : (
+            <>
+              <Copy size={12} />
+              <span>Copy</span>
+            </>
+          )}
+        </button>
+        <button
+          onClick={handleCopyMarkdown}
+          type="button"
+          className={cn(
+            "flex items-center gap-1 rounded px-1.5 py-0.5 transition-colors",
+            isTerminal
+              ? "text-zinc-400 hover:text-zinc-200 hover:bg-zinc-700/50"
+              : "text-muted-foreground hover:text-foreground hover:bg-accent"
+          )}
+          title="Copy as Markdown"
+        >
+          {copiedMarkdown ? (
+            <>
+              <Check size={12} />
+              <span>Copied</span>
+            </>
+          ) : (
+            <>
+              <FileCode size={12} />
+              <span>Markdown</span>
+            </>
+          )}
+        </button>
+      </div>
+    </div>
+  );
+};
+
 export type CodeBlockCopyButtonProps = ComponentProps<typeof Button> & {
   onCopy?: () => void;
   onError?: (error: Error) => void;
@@ -591,3 +854,12 @@ export type CodeBlockLanguageSelectorItemProps = ComponentProps<
 export const CodeBlockLanguageSelectorItem = (
   props: CodeBlockLanguageSelectorItemProps
 ) => <SelectItem {...props} />;
+
+// ── InlineCode ─────────────────────────────────────────────────────────
+export function InlineCode({ children }: { children: ReactNode }) {
+  return (
+    <code className="rounded bg-muted px-1.5 py-0.5 text-sm font-mono">
+      {children}
+    </code>
+  );
+}
diff --git a/src/components/chat/CodeBlock.tsx b/src/components/chat/CodeBlock.tsx
deleted file mode 100644
index 9ecfa45..0000000
--- a/src/components/chat/CodeBlock.tsx
+++ /dev/null
@@ -1,293 +0,0 @@
-'use client';
-
-import { useState, useMemo, useRef, createElement } from 'react';
-import { Prism as SyntaxHighlighter } from 'react-syntax-highlighter';
-import { vscDarkPlus } from 'react-syntax-highlighter/dist/esm/styles/prism';
-import { useThemeFamily } from '@/lib/theme/context';
-import { resolveCodeTheme, resolvePrismStyle } from '@/lib/theme/code-themes';
-import { useTheme } from 'next-themes';
-import type { Icon } from "@phosphor-icons/react";
-import {
-  Copy,
-  Check,
-  CaretDown,
-  CaretUp,
-  FileCode,
-  Terminal,
-  Code,
-  File,
-  Hash,
-} from "@phosphor-icons/react";
-import { cn } from '@/lib/utils';
-
-const COLLAPSE_THRESHOLD = 20;
-const VISIBLE_LINES = 10;
-
-const TERMINAL_LANGUAGES = new Set(['bash', 'sh', 'shell', 'terminal', 'zsh', 'console']);
-
-function getLanguageIcon(language: string): Icon {
-  const lower = language.toLowerCase();
-  if (TERMINAL_LANGUAGES.has(lower)) return Terminal;
-  if (['typescript', 'tsx', 'javascript', 'jsx'].includes(lower)) return Code;
-  if (['json', 'yaml', 'yml', 'toml', 'xml'].includes(lower)) return Code;
-  if (['python', 'ruby', 'go', 'rust', 'java', 'c', 'cpp'].includes(lower)) return Hash;
-  if (['css', 'scss', 'html'].includes(lower)) return File;
-  return FileCode;
-}
-
-interface CodeBlockProps {
-  code: string;
-  language?: string;
-  filename?: string;
-  showLineNumbers?: boolean;
-  maxCollapsedLines?: number;
-}
-
-/** Resolve the Prism syntax highlighting style for CodeBlock. */
-function useCodeBlockTheme(isTerminal: boolean, isDark: boolean) {
-  const { family, families } = useThemeFamily();
-  if (isTerminal) return vscDarkPlus;
-  const codeTheme = resolveCodeTheme(families, family);
-  return resolvePrismStyle(codeTheme, isDark);
-}
-
-export function CodeBlock({
-  code,
-  language = 'text',
-  filename,
-  showLineNumbers = true,
-  maxCollapsedLines = VISIBLE_LINES,
-}: CodeBlockProps) {
-  const [copied, setCopied] = useState(false);
-  const [copiedMarkdown, setCopiedMarkdown] = useState(false);
-  const [expanded, setExpanded] = useState(false);
-  const codeContainerRef = useRef<HTMLDivElement>(null);
-  const [animatingHeight, setAnimatingHeight] = useState<string | undefined>(undefined);
-  const { resolvedTheme } = useTheme();
-  const isDark = resolvedTheme === 'dark';
-
-  const lines = useMemo(() => code.split('\n'), [code]);
-  const totalLines = lines.length;
-  const isCollapsible = totalLines > COLLAPSE_THRESHOLD;
-  const isTerminal = TERMINAL_LANGUAGES.has(language.toLowerCase());
-
-  const displayCode = useMemo(() => {
-    if (!isCollapsible || expanded) return code;
-    return lines.slice(0, maxCollapsedLines).join('\n');
-  }, [code, lines, isCollapsible, expanded, maxCollapsedLines]);
-
-  const handleCopy = async () => {
-    try {
-      await navigator.clipboard.writeText(code);
-      setCopied(true);
-      setTimeout(() => setCopied(false), 2000);
-    } catch {
-      // clipboard not available
-    }
-  };
-
-  const handleCopyMarkdown = async () => {
-    try {
-      const markdown = `\`\`\`${language}\n${code}\n\`\`\``;
-      await navigator.clipboard.writeText(markdown);
-      setCopiedMarkdown(true);
-      setTimeout(() => setCopiedMarkdown(false), 2000);
-    } catch {
-      // clipboard not available
-    }
-  };
-
-  const handleToggleExpand = () => {
-    const container = codeContainerRef.current;
-    if (!container) {
-      setExpanded(!expanded);
-      return;
-    }
-    const currentHeight = container.scrollHeight;
-    if (!expanded) {
-      setAnimatingHeight(`${currentHeight}px`);
-      setExpanded(true);
-      requestAnimationFrame(() => {
-        const fullHeight = container.scrollHeight;
-        setAnimatingHeight(`${fullHeight}px`);
-        setTimeout(() => setAnimatingHeight(undefined), 300);
-      });
-    } else {
-      setAnimatingHeight(`${currentHeight}px`);
-      requestAnimationFrame(() => {
-        const collapsedH = maxCollapsedLines * 1.5 + 1.5;
-        setAnimatingHeight(`${collapsedH}rem`);
-        setTimeout(() => {
-          setExpanded(false);
-          setAnimatingHeight(undefined);
-        }, 300);
-      });
-    }
-  };
-
-  const languageIcon = getLanguageIcon(language);
-
-  const theme = useCodeBlockTheme(isTerminal, isDark);
-
-  return (
-    <div className={cn(
-      "relative group not-prose my-3 rounded-lg overflow-hidden",
-      isTerminal
-        ? "border border-zinc-700/50"
-        : "border border-border"
-    )}>
-      {/* Header bar */}
-      <div className={cn(
-        "flex items-center justify-between px-4 py-1.5 text-xs",
-        isTerminal
-          ? "bg-zinc-950 text-zinc-400"
-          : "bg-muted text-muted-foreground"
-      )}>
-        <div className="flex items-center gap-2 min-w-0">
-          {createElement(languageIcon, { size: 14, className: cn(
-            "shrink-0",
-            isTerminal ? "text-green-400" : "text-muted-foreground",
-          ) })}
-          {filename && (
-            <span className={cn(
-              "truncate font-medium",
-              isTerminal ? "text-zinc-300" : "text-foreground"
-            )}>{filename}</span>
-          )}
-          {filename && <span className="text-muted-foreground/50">|</span>}
-          <span className={cn(
-            "rounded px-1.5 py-0.5",
-            isTerminal
-              ? "bg-zinc-700/50 text-green-400"
-              : "bg-accent text-accent-foreground"
-          )}>{language.toUpperCase()}</span>
-        </div>
-        <div className="flex items-center gap-1 ml-2 shrink-0">
-          <button
-            onClick={handleCopy}
-            className={cn(
-              "flex items-center gap-1 rounded px-1.5 py-0.5 transition-colors",
-              isTerminal
-                ? "text-zinc-400 hover:text-zinc-200 hover:bg-zinc-700/50"
-                : "text-muted-foreground hover:text-foreground hover:bg-accent"
-            )}
-            title="Copy code"
-          >
-            {copied ? (
-              <>
-                <Check size={12} />
-                <span>Copied</span>
-              </>
-            ) : (
-              <>
-                <Copy size={12} />
-                <span>Copy</span>
-              </>
-            )}
-          </button>
-          <button
-            onClick={handleCopyMarkdown}
-            className={cn(
-              "flex items-center gap-1 rounded px-1.5 py-0.5 transition-colors",
-              isTerminal
-                ? "text-zinc-400 hover:text-zinc-200 hover:bg-zinc-700/50"
-                : "text-muted-foreground hover:text-foreground hover:bg-accent"
-            )}
-            title="Copy as Markdown"
-          >
-            {copiedMarkdown ? (
-              <>
-                <Check size={12} />
-                <span>Copied</span>
-              </>
-            ) : (
-              <>
-                <FileCode size={12} />
-                <span>Markdown</span>
-              </>
-            )}
-          </button>
-        </div>
-      </div>
-
-      {/* Code area */}
-      <div
-        ref={codeContainerRef}
-        className="relative transition-[max-height] duration-300 ease-in-out overflow-hidden"
-        style={{
-          maxHeight: animatingHeight ?? (!isCollapsible || expanded ? undefined : `${maxCollapsedLines * 1.5 + 1.5}rem`),
-        }}
-      >
-        <SyntaxHighlighter
-          style={theme}
-          language={language}
-          PreTag="div"
-          showLineNumbers={showLineNumbers && !isTerminal}
-          lineNumberStyle={{
-            minWidth: '2.5em',
-            paddingRight: '1em',
-            color: 'var(--muted-foreground)',
-            opacity: 0.4,
-            userSelect: 'none',
-          }}
-          customStyle={{
-            margin: 0,
-            borderRadius: 0,
-            fontSize: '0.8125rem',
-            lineHeight: '1.5',
-            padding: isTerminal ? '0.75rem 1rem' : '0.75rem 0.5rem',
-            background: isTerminal ? '#0a0a0a' : undefined,
-            overflow: 'auto',
-          }}
-          wrapLines
-        >
-          {expanded ? code : displayCode}
-        </SyntaxHighlighter>
-
-        {/* Gradient overlay for collapsed state */}
-        {isCollapsible && !expanded && (
-          <div className={cn(
-            "absolute bottom-0 left-0 right-0 h-16 pointer-events-none",
-            isTerminal
-              ? "bg-gradient-to-t from-[#0a0a0a] to-transparent"
-              : "bg-gradient-to-t from-muted to-transparent"
-          )} />
-        )}
-      </div>
-
-      {/* Expand/Collapse button */}
-      {isCollapsible && (
-        <button
-          onClick={handleToggleExpand}
-          className={cn(
-            "flex w-full items-center justify-center gap-1.5 py-1.5 text-xs transition-colors",
-            isTerminal
-              ? "bg-zinc-950 text-zinc-400 hover:text-zinc-200"
-              : "bg-muted text-muted-foreground hover:text-foreground"
-          )}
-        >
-          {expanded ? (
-            <>
-              <CaretUp size={12} />
-              <span>Collapse</span>
-            </>
-          ) : (
-            <>
-              <CaretDown size={12} />
-              <span>Expand all {totalLines} lines</span>
-            </>
-          )}
-        </button>
-      )}
-    </div>
-  );
-}
-
-// Inline code component for reuse
-export function InlineCode({ children }: { children: React.ReactNode }) {
-  return (
-    <code className="rounded bg-muted px-1.5 py-0.5 text-sm font-mono">
-      {children}
-    </code>
-  );
-}
diff --git a/src/components/chat/ToolCallBlock.tsx b/src/components/chat/ToolCallBlock.tsx
index dba355a..d6a474c 100644
--- a/src/components/chat/ToolCallBlock.tsx
+++ b/src/components/chat/ToolCallBlock.tsx
@@ -15,7 +15,7 @@ import {
   XCircle,
 } from "@phosphor-icons/react";
 import { cn } from '@/lib/utils';
-import { CodeBlock } from './CodeBlock';
+import { CodeBlock } from '@/components/ai-elements/code-block';
 
 type ToolStatus = 'running' | 'success' | 'error';