Add QOL features to tool call code

ui/src/components/Artifact/CodeArtifact.tsx

@@ -1,16 +1,14 @@
 /**
  * CodeArtifact - Syntax-Highlighted Code Output
  *
- * Renders code output from tools like code_interpreter with syntax highlighting.
- * Uses a simple pre/code block approach that works with the prose styling.
+ * Renders code output from tools like code_interpreter with syntax highlighting
+ * via Shiki. Uses the shared HighlightedCode component.
  */
 
-import { memo, useState } from "react";
-import { Copy, Check } from "lucide-react";
+import { memo } from "react";
 
 import type { Artifact, CodeArtifactData } from "@/components/chat-types";
-import { Button } from "@/components/Button/Button";
-import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/Tooltip/Tooltip";
+import { HighlightedCode } from "@/components/HighlightedCode/HighlightedCode";
 import { cn } from "@/utils/cn";
 
 export interface CodeArtifactProps {
@@ -28,64 +26,21 @@ function isCodeArtifactData(data: unknown): data is CodeArtifactData {
 }
 
 function CodeArtifactComponent({ artifact, className }: CodeArtifactProps) {
-  const [copied, setCopied] = useState(false);
-
-  // Validate and extract data
   if (!isCodeArtifactData(artifact.data)) {
     return <div className="p-4 text-sm text-muted-foreground">Invalid code artifact data</div>;
   }
 
   const { code, language } = artifact.data;
 
-  const handleCopy = async () => {
-    await navigator.clipboard.writeText(code);
-    setCopied(true);
-    setTimeout(() => setCopied(false), 2000);
-  };
-
   return (
-    <div className={cn("relative group", className)}>
-      {/* Copy button */}
-      <div className="absolute right-2 top-2 opacity-0 group-hover:opacity-100 transition-opacity z-10">
-        <Tooltip>
-          <TooltipTrigger asChild>
-            <Button
-              variant="secondary"
-              size="sm"
-              className="h-7 w-7 p-0"
-              onClick={handleCopy}
-              aria-label={copied ? "Copied" : "Copy code"}
-            >
-              {copied ? <Check className="h-3.5 w-3.5" /> : <Copy className="h-3.5 w-3.5" />}
-            </Button>
-          </TooltipTrigger>
-          <TooltipContent>{copied ? "Copied!" : "Copy code"}</TooltipContent>
-        </Tooltip>
-      </div>
-
-      {/* Language badge */}
-      {language && (
-        <div className="absolute left-2 top-2 z-10">
-          <span className="text-[10px] font-mono text-muted-foreground bg-muted/80 px-1.5 py-0.5 rounded">
-            {language}
-          </span>
-        </div>
-      )}
-
-      {/* Code block */}
-      {/* eslint-disable jsx-a11y/no-noninteractive-tabindex -- scrollable region needs keyboard access (axe: scrollable-region-focusable) */}
-      <pre
-        tabIndex={0}
-        className={cn(
-          "p-4 pt-8 overflow-x-auto text-sm font-mono",
-          "bg-muted/50 text-foreground",
-          "max-h-[400px] overflow-y-auto"
-        )}
-      >
-        <code className="whitespace-pre">{code}</code>
-      </pre>
-      {/* eslint-enable jsx-a11y/no-noninteractive-tabindex */}
-    </div>
+    <HighlightedCode
+      code={code}
+      language={language}
+      showCopy
+      showLanguage
+      maxHeight="400px"
+      className={cn("rounded-md overflow-hidden", className)}
+    />
   );
 }
 

ui/src/components/Artifact/HtmlArtifact.tsx

@@ -6,11 +6,12 @@
  */
 
 import { memo, useState, useRef, useEffect } from "react";
-import { Code2, Eye, Copy, Check, ExternalLink, Maximize2, X } from "lucide-react";
+import { Code2, Eye, ExternalLink, Maximize2, X } from "lucide-react";
 
 import type { Artifact } from "@/components/chat-types";
 import { Button } from "@/components/Button/Button";
 import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/Tooltip/Tooltip";
+import { HighlightedCode } from "@/components/HighlightedCode/HighlightedCode";
 import { cn } from "@/utils/cn";
 
 export interface HtmlArtifactProps {
@@ -69,7 +70,6 @@ function wrapHtml(content: string): string {
 
 function HtmlArtifactComponent({ artifact, className }: HtmlArtifactProps) {
   const [viewMode, setViewMode] = useState<"preview" | "source">("preview");
-  const [copied, setCopied] = useState(false);
   const [isFullscreen, setIsFullscreen] = useState(false);
   const iframeRef = useRef<HTMLIFrameElement>(null);
 
@@ -87,12 +87,6 @@ function HtmlArtifactComponent({ artifact, className }: HtmlArtifactProps) {
     return <div className="p-4 text-sm text-muted-foreground">Invalid HTML artifact data</div>;
   }
 
-  const handleCopy = async () => {
-    await navigator.clipboard.writeText(html);
-    setCopied(true);
-    setTimeout(() => setCopied(false), 2000);
-  };
-
   const handleOpenInNewTab = () => {
     const blob = new Blob([wrapHtml(html)], { type: "text/html" });
     const url = URL.createObjectURL(blob);
@@ -129,21 +123,6 @@ function HtmlArtifactComponent({ artifact, className }: HtmlArtifactProps) {
 
           <div className="flex-1" />
 
-          <Tooltip>
-            <TooltipTrigger asChild>
-              <Button
-                variant="ghost"
-                size="sm"
-                className="h-6 w-6 p-0"
-                onClick={handleCopy}
-                aria-label={copied ? "Copied" : "Copy HTML"}
-              >
-                {copied ? <Check className="h-3.5 w-3.5" /> : <Copy className="h-3.5 w-3.5" />}
-              </Button>
-            </TooltipTrigger>
-            <TooltipContent>{copied ? "Copied!" : "Copy HTML"}</TooltipContent>
-          </Tooltip>
-
           <Tooltip>
             <TooltipTrigger asChild>
               <Button
@@ -184,9 +163,7 @@ function HtmlArtifactComponent({ artifact, className }: HtmlArtifactProps) {
             className="w-full h-[300px] border-0 bg-white"
           />
         ) : (
-          <pre className="p-4 overflow-x-auto text-xs font-mono text-foreground max-h-[300px] overflow-y-auto bg-muted/30">
-            <code>{html}</code>
-          </pre>
+          <HighlightedCode code={html} language="html" showCopy maxHeight="300px" />
         )}
       </div>
 

ui/src/components/HighlightedCode/HighlightedCode.stories.tsx

@@ -0,0 +1,100 @@
+import type { Meta, StoryObj } from "@storybook/react";
+import { HighlightedCode } from "./HighlightedCode";
+
+const meta = {
+  title: "Chat/HighlightedCode",
+  component: HighlightedCode,
+  parameters: {
+    layout: "padded",
+  },
+} satisfies Meta<typeof HighlightedCode>;
+
+export default meta;
+type Story = StoryObj<typeof meta>;
+
+export const Python: Story = {
+  args: {
+    code: `import pandas as pd
+import numpy as np
+
+def analyze(data: list[float]) -> dict:
+    """Compute summary statistics."""
+    arr = np.array(data)
+    return {
+        "mean": float(arr.mean()),
+        "std": float(arr.std()),
+        "median": float(np.median(arr)),
+    }
+
+result = analyze([1, 2, 3, 4, 5])
+print(result)`,
+    language: "python",
+    showLanguage: true,
+  },
+};
+
+export const JavaScript: Story = {
+  args: {
+    code: `const fetchUsers = async () => {
+  const response = await fetch('/api/users');
+  const data = await response.json();
+  return data.users.filter(u => u.active);
+};
+
+fetchUsers().then(console.log);`,
+    language: "javascript",
+    showLanguage: true,
+  },
+};
+
+export const SQL: Story = {
+  args: {
+    code: `SELECT u.name, COUNT(o.id) AS order_count
+FROM users u
+LEFT JOIN orders o ON o.user_id = u.id
+WHERE u.created_at > '2024-01-01'
+GROUP BY u.name
+ORDER BY order_count DESC
+LIMIT 10;`,
+    language: "sql",
+    showLanguage: true,
+  },
+};
+
+export const Compact: Story = {
+  args: {
+    code: `data = [1, 2, 3, 4, 5]
+result = sum(data) / len(data)
+print(f"Average: {result}")`,
+    language: "python",
+    compact: true,
+    showCopy: false,
+  },
+};
+
+export const WithCopyButton: Story = {
+  args: {
+    code: `console.log("Hello, world!");`,
+    language: "javascript",
+    showCopy: true,
+  },
+};
+
+export const WithMaxHeight: Story = {
+  args: {
+    code: Array(50)
+      .fill(null)
+      .map((_, i) => `console.log("Line ${i + 1}");`)
+      .join("\n"),
+    language: "javascript",
+    showLanguage: true,
+    maxHeight: "200px",
+  },
+};
+
+export const UnknownLanguage: Story = {
+  args: {
+    code: "Some plain text content\nwith multiple lines\nand no highlighting",
+    language: "custom-lang",
+  },
+};