feat(chat): add copy buttons to code blocks and agent prose messages (#3238)

Author: MattPua

packages/ui/src/features/sessions/components/chat-thread/ChatMarkdown.tsx

@@ -1,3 +1,4 @@
+import { Check, Copy } from "@phosphor-icons/react";
 import {
   Heading,
   Separator,
@@ -14,11 +15,41 @@ import {
   splitMarkdownBlocks,
 } from "@posthog/ui/features/editor/components/splitMarkdownBlocks";
 import { HighlightedCode } from "@posthog/ui/primitives/HighlightedCode";
-import { memo, useMemo } from "react";
+import { useCopy } from "@posthog/ui/primitives/useCopy";
+import { IconButton } from "@radix-ui/themes";
+import { memo, type ReactNode, useMemo } from "react";
 import Markdown, { type Components } from "react-markdown";
 import rehypeSanitize from "rehype-sanitize";
 import remarkGfm from "remark-gfm";
 
+function ChatCodeBlock({
+  code,
+  children,
+}: {
+  code: string;
+  children: ReactNode;
+}) {
+  const { copied, copy } = useCopy();
+
+  return (
+    <div className="group relative">
+      <pre className="overflow-x-auto rounded-lg border border-border bg-muted/50 p-3 pr-10 text-sm leading-[1.5]">
+        {children}
+      </pre>
+      <IconButton
+        size="1"
+        variant="ghost"
+        color={copied ? "green" : "gray"}
+        onClick={() => copy(code)}
+        className="absolute top-1 right-1 cursor-pointer opacity-0 transition-opacity group-hover:opacity-100"
+        aria-label="Copy code"
+      >
+        {copied ? <Check size={14} /> : <Copy size={14} />}
+      </IconButton>
+    </div>
+  );
+}
+
 /**
  * The chat thread's own markdown renderer — intentionally separate from the app-wide
  * `MarkdownRenderer` (which carries PostHog deeplink handling, Radix Text wrappers, and other
@@ -47,16 +78,24 @@ const components: Components = {
   ),
   li: ({ children }) => <li className="text-sm">{children}</li>,
   code: ({ className, children }) => {
+    const text = String(children).replace(/\n$/, "");
     const match = /language-(\w+)/.exec(className ?? "");
-    if (match) {
-      // Fenced block with a language → Shiki-highlighted (theme-aware). The `pre` renderer
-      // below provides the box; HighlightedCode renders the colored <code> inside it.
+    // Fenced blocks (carry a language, or span multiple lines) render as a boxed, copyable
+    // block; short inline spans stay inline. `pre` below is a passthrough so the box lives here,
+    // where the raw code string is in hand.
+    if (match || text.includes("\n")) {
       return (
-        <HighlightedCode
-          code={String(children).replace(/\n$/, "")}
-          language={match[1]}
-          className="rounded-sm bg-muted/50 text-xs"
-        />
+        <ChatCodeBlock code={text}>
+          {match ? (
+            <HighlightedCode
+              code={text}
+              language={match[1]}
+              className="text-xs"
+            />
+          ) : (
+            <code className="font-mono text-xs">{text}</code>
+          )}
+        </ChatCodeBlock>
       );
     }
     return (
@@ -65,11 +104,7 @@ const components: Components = {
       </code>
     );
   },
-  pre: ({ children }) => (
-    <pre className="overflow-x-auto rounded-lg border border-border bg-muted/50 p-3 text-sm leading-[1.5]">
-      {children}
-    </pre>
-  ),
+  pre: ({ children }) => <>{children}</>,
   h1: ({ children }) => (
     <Heading size="xl" className="font-bold">
       {children}
@@ -153,9 +188,9 @@ export const ChatStreamingMarkdown = memo(function ChatStreamingMarkdown({
               {openFence.before.trim() ? (
                 <ChatMarkdown content={openFence.before} />
               ) : null}
-              <pre className="overflow-x-auto rounded-lg border border-border bg-muted/50 p-3 text-sm leading-[1.5]">
+              <ChatCodeBlock code={openFence.code}>
                 <code className="font-mono text-xs">{openFence.code}</code>
-              </pre>
+              </ChatCodeBlock>
             </div>
           );
         }

packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx

@@ -1,4 +1,11 @@
-import { CaretDown, ChatCircle, FileText, Scroll } from "@phosphor-icons/react";
+import {
+  CaretDown,
+  ChatCircle,
+  Check,
+  Copy,
+  FileText,
+  Scroll,
+} from "@phosphor-icons/react";
 import { WorkerPoolContextProvider } from "@pierre/diffs/react";
 import { useService } from "@posthog/di/react";
 import {
@@ -64,10 +71,12 @@ import {
 } from "@posthog/ui/features/sessions/useSessionTaskId";
 import { useSettingsStore } from "@posthog/ui/features/settings/settingsStore";
 import { SkillButtonActionMessage } from "@posthog/ui/features/skill-buttons/components/SkillButtonActionMessage";
+import { useCopy } from "@posthog/ui/primitives/useCopy";
 import {
   DIFF_WORKER_FACTORY,
   type DiffWorkerFactory,
 } from "@posthog/ui/shell/diffWorkerHost";
+import { IconButton, Tooltip } from "@radix-ui/themes";
 import { AnimatePresence, motion, useReducedMotion } from "framer-motion";
 import {
   memo,
@@ -540,8 +549,10 @@ const AgentProse = memo(function AgentProse({
   isStreaming?: boolean;
 }) {
   const smoothed = useSmoothedText(text);
+  const { copied, copy } = useCopy();
+
   return (
-    <ChatMessage align="start">
+    <ChatMessage align="start" className="group/msg">
       <ChatMessageContent className="gap-1">
         <ChatBubble variant="ghost">
           <ChatBubbleContent>
@@ -552,6 +563,22 @@ const AgentProse = memo(function AgentProse({
             )}
           </ChatBubbleContent>
         </ChatBubble>
+        {isStreaming ? null : (
+          <ChatMessageFooter className="opacity-0 transition-opacity group-hover/msg:opacity-100">
+            <Tooltip content={copied ? "Copied!" : "Copy message"}>
+              <IconButton
+                size="1"
+                variant="ghost"
+                color={copied ? "green" : "gray"}
+                onClick={() => copy(text)}
+                className="cursor-pointer"
+                aria-label="Copy message"
+              >
+                {copied ? <Check size={14} /> : <Copy size={14} />}
+              </IconButton>
+            </Tooltip>
+          </ChatMessageFooter>
+        )}
       </ChatMessageContent>
     </ChatMessage>
   );

packages/ui/src/primitives/useCopy.ts

@@ -0,0 +1,26 @@
+import { useCallback, useState } from "react";
+
+/**
+ * Copy-to-clipboard with a transient `copied` flag for button feedback. Clipboard writes can reject
+ * (blocked permission, unfocused document, insecure context) — the rejection is swallowed and
+ * `copied` stays false, so the button never reports a success that didn't happen.
+ */
+export function useCopy(resetMs = 2000): {
+  copied: boolean;
+  copy: (text: string) => void;
+} {
+  const [copied, setCopied] = useState(false);
+  const copy = useCallback(
+    (text: string) => {
+      navigator.clipboard.writeText(text).then(
+        () => {
+          setCopied(true);
+          setTimeout(() => setCopied(false), resetMs);
+        },
+        () => {},
+      );
+    },
+    [resetMs],
+  );
+  return { copied, copy };
+}