feat: add copy button for tool call blocks (#463)

* feat: add copy button for tool call blocks

* refactor: simplify copy state updates with helper function

---------

Co-authored-by: binge_c-admin <totchinaa@gmail.com>

components/chat-message-display.tsx

@@ -230,6 +230,12 @@ export function ChatMessageDisplay({
     const [expandedTools, setExpandedTools] = useState<Record<string, boolean>>(
         {},
     )
+    const [copiedToolCallId, setCopiedToolCallId] = useState<string | null>(
+        null,
+    )
+    const [copyFailedToolCallId, setCopyFailedToolCallId] = useState<
+        string | null
+    >(null)
     const [copiedMessageId, setCopiedMessageId] = useState<string | null>(null)
     const [copyFailedMessageId, setCopyFailedMessageId] = useState<
         string | null
@@ -245,12 +251,38 @@ export function ChatMessageDisplay({
         Record<string, boolean>
     >({})
 
-    const copyMessageToClipboard = async (messageId: string, text: string) => {
+    const setCopyState = (
-        try {
+        messageId: string,
-            await navigator.clipboard.writeText(text)
+        isToolCall: boolean,
-
+        isSuccess: boolean,
+    ) => {
+        if (isSuccess) {
+            if (isToolCall) {
+                setCopiedToolCallId(messageId)
+                setTimeout(() => setCopiedToolCallId(null), 2000)
+            } else {
                 setCopiedMessageId(messageId)
                 setTimeout(() => setCopiedMessageId(null), 2000)
+            }
+        } else {
+            if (isToolCall) {
+                setCopyFailedToolCallId(messageId)
+                setTimeout(() => setCopyFailedToolCallId(null), 2000)
+            } else {
+                setCopyFailedMessageId(messageId)
+                setTimeout(() => setCopyFailedMessageId(null), 2000)
+            }
+        }
+    }
+
+    const copyMessageToClipboard = async (
+        messageId: string,
+        text: string,
+        isToolCall = false,
+    ) => {
+        try {
+            await navigator.clipboard.writeText(text)
+            setCopyState(messageId, isToolCall, true)
         } catch (err) {
             // Fallback for non-secure contexts (HTTP) or permission denied
             const textarea = document.createElement("textarea")
@@ -266,13 +298,11 @@ export function ChatMessageDisplay({
                 if (!success) {
                     throw new Error("Copy command failed")
                 }
-                setCopiedMessageId(messageId)
+                setCopyState(messageId, isToolCall, true)
-                setTimeout(() => setCopiedMessageId(null), 2000)
             } catch (fallbackErr) {
                 console.error("Failed to copy message:", fallbackErr)
                 toast.error(dict.chat.failedToCopyDetail)
-                setCopyFailedMessageId(messageId)
+                setCopyState(messageId, isToolCall, false)
-                setTimeout(() => setCopyFailedMessageId(null), 2000)
             } finally {
                 document.body.removeChild(textarea)
             }
@@ -641,6 +671,7 @@ export function ChatMessageDisplay({
         const { state, input, output } = part
         const isExpanded = expandedTools[callId] ?? true
         const toolName = part.type?.replace("tool-", "")
+        const isCopied = copiedToolCallId === callId
 
         const toggleExpanded = () => {
             setExpandedTools((prev) => ({
@@ -662,6 +693,35 @@ export function ChatMessageDisplay({
             }
         }
 
+        const handleCopy = () => {
+            let textToCopy = ""
+
+            if (input && typeof input === "object") {
+                if (input.xml) {
+                    textToCopy = input.xml
+                } else if (
+                    input.operations &&
+                    Array.isArray(input.operations)
+                ) {
+                    textToCopy = JSON.stringify(input.operations, null, 2)
+                } else if (Object.keys(input).length > 0) {
+                    textToCopy = JSON.stringify(input, null, 2)
+                }
+            }
+
+            if (
+                output &&
+                toolName === "get_shape_library" &&
+                typeof output === "string"
+            ) {
+                textToCopy = output
+            }
+
+            if (textToCopy) {
+                copyMessageToClipboard(callId, textToCopy, true)
+            }
+        }
+
         return (
             <div
                 key={callId}
@@ -681,9 +741,32 @@ export function ChatMessageDisplay({
                             <div className="h-4 w-4 border-2 border-primary border-t-transparent rounded-full animate-spin" />
                         )}
                         {state === "output-available" && (
+                            <>
                                 <span className="text-xs font-medium text-green-600 bg-green-50 px-2 py-0.5 rounded-full">
-                                Complete
+                                    {dict.tools.complete}
                                 </span>
+                                {isExpanded && (
+                                    <button
+                                        type="button"
+                                        onClick={handleCopy}
+                                        className="p-1 rounded hover:bg-muted transition-colors"
+                                        title={
+                                            copiedToolCallId === callId
+                                                ? dict.chat.copied
+                                                : copyFailedToolCallId ===
+                                                    callId
+                                                  ? dict.chat.failedToCopy
+                                                  : dict.chat.copyResponse
+                                        }
+                                    >
+                                        {isCopied ? (
+                                            <Check className="w-4 h-4 text-green-600" />
+                                        ) : (
+                                            <Copy className="w-4 h-4 text-muted-foreground" />
+                                        )}
+                                    </button>
+                                )}
+                            </>
                         )}
                         {state === "output-error" &&
                             (() => {