From 9e96f51e092516eaeed1aeed6e69896f14d8753f Mon Sep 17 00:00:00 2001
From: Shine1i <wasimysdev@gmail.com>
Date: Mon, 9 Mar 2026 13:07:54 +0100
Subject: [PATCH] feat(studio): add support for code block actions including
 copy and download options in markdown blocks

---
 .../components/assistant-ui/markdown-text.tsx | 131 +++++++++++++++++-
 1 file changed, 127 insertions(+), 4 deletions(-)

diff --git a/studio/frontend/src/components/assistant-ui/markdown-text.tsx b/studio/frontend/src/components/assistant-ui/markdown-text.tsx
index 252faea6c3..c967b33c27 100644
--- a/studio/frontend/src/components/assistant-ui/markdown-text.tsx
+++ b/studio/frontend/src/components/assistant-ui/markdown-text.tsx
@@ -7,18 +7,69 @@ import { HugeiconsIcon } from "@hugeicons/react";
 import { code } from "@streamdown/code";
 import { math } from "@streamdown/math";
 import { mermaid } from "@streamdown/mermaid";
+import { DownloadIcon } from "lucide-react";
 import { Block, type BlockProps, Streamdown } from "streamdown";
 import { useEffect, useRef, useState } from "react";
 import "katex/dist/katex.min.css";
 import { AudioPlayer } from "./audio-player";
 
-const { withSmoothContextProvider, useSmoothStatus } = INTERNAL;
+const { withSmoothContextProvider } = INTERNAL;
 
 function getMermaidSource(blockContent: string): string | null {
   const source = blockContent.match(/```mermaid\s*([\s\S]*?)```/i)?.[1]?.trim();
   return source && source.length > 0 ? source : null;
 }
 
+function getCodeFence(
+  blockContent: string,
+): { language: string | null; source: string } | null {
+  const match = blockContent
+    .trimEnd()
+    .match(/^```([^\n`]*)\n([\s\S]*?)\n?```$/);
+  if (!match) return null;
+
+  return {
+    language: match[1]?.trim() || null,
+    source: match[2],
+  };
+}
+
+function getCodeFilename(language: string | null) {
+  const extByLanguage: Record<string, string> = {
+    bash: "sh",
+    javascript: "js",
+    js: "js",
+    json: "json",
+    jsx: "jsx",
+    markdown: "md",
+    md: "md",
+    python: "py",
+    py: "py",
+    shell: "sh",
+    sh: "sh",
+    sql: "sql",
+    ts: "ts",
+    tsx: "tsx",
+    typescript: "ts",
+    yaml: "yml",
+    yml: "yml",
+  };
+
+  const normalized = language?.toLowerCase();
+  const ext = normalized ? extByLanguage[normalized] || normalized : "txt";
+  return `snippet.${ext}`;
+}
+
+function downloadTextFile(filename: string, text: string) {
+  const blob = new Blob([text], { type: "text/plain;charset=utf-8" });
+  const url = URL.createObjectURL(blob);
+  const anchor = document.createElement("a");
+  anchor.href = url;
+  anchor.download = filename;
+  anchor.click();
+  URL.revokeObjectURL(url);
+}
+
 const COPY_RESET_MS = 2000;
 
 function MermaidCopyButton({ source }: { source: string }) {
@@ -55,9 +106,68 @@ function MermaidCopyButton({ source }: { source: string }) {
   );
 }
 
+function CodeBlockActions({
+  disabled,
+  language,
+  source,
+}: {
+  disabled: boolean;
+  language: string | null;
+  source: string;
+}) {
+  const [copied, setCopied] = useState(false);
+  const resetTimeoutRef = useRef<ReturnType<typeof setTimeout> | null>(null);
+
+  useEffect(() => {
+    return () => {
+      if (resetTimeoutRef.current) {
+        clearTimeout(resetTimeoutRef.current);
+      }
+    };
+  }, []);
+
+  return (
+    <div className="pointer-events-none absolute top-3.5 right-3 z-20 flex items-center justify-end">
+      <div className="pointer-events-auto flex shrink-0 items-center gap-2 rounded-md border border-sidebar bg-sidebar/80 px-1.5 py-1 supports-[backdrop-filter]:bg-sidebar/70 supports-[backdrop-filter]:backdrop-blur">
+        <button
+          type="button"
+          className="cursor-pointer p-1 text-muted-foreground transition-all hover:text-foreground disabled:cursor-not-allowed disabled:opacity-50"
+          title="Copy code"
+          disabled={disabled}
+          onClick={() => {
+            if (!copyToClipboard(source)) return;
+            setCopied(true);
+            if (resetTimeoutRef.current) {
+              clearTimeout(resetTimeoutRef.current);
+            }
+            resetTimeoutRef.current = setTimeout(() => {
+              setCopied(false);
+              resetTimeoutRef.current = null;
+            }, COPY_RESET_MS);
+          }}
+        >
+          <HugeiconsIcon icon={copied ? Tick02Icon : Copy02Icon} className="size-3.5" />
+        </button>
+        <button
+          type="button"
+          className="cursor-pointer p-1 text-muted-foreground transition-all hover:text-foreground disabled:cursor-not-allowed disabled:opacity-50"
+          title="Download file"
+          disabled={disabled}
+          onClick={() => {
+            downloadTextFile(getCodeFilename(language), source);
+          }}
+        >
+          <DownloadIcon className="size-3.5" />
+        </button>
+      </div>
+    </div>
+  );
+}
+
 function StreamdownBlock(props: BlockProps) {
   const hasMermaidFence = props.content.includes("```mermaid");
   const mermaidSource = getMermaidSource(props.content);
+  const codeFence = getCodeFence(props.content);
 
   if (props.isIncomplete && hasMermaidFence) {
     return (
@@ -69,20 +179,32 @@ function StreamdownBlock(props: BlockProps) {
 
   if (mermaidSource) {
     return (
-      <div className="relative">
+      <div className="relative isolate">
         <Block {...props} />
         <MermaidCopyButton source={mermaidSource} />
       </div>
     );
   }
 
+  if (codeFence) {
+    return (
+      <div className="relative isolate">
+        <Block {...props} />
+        <CodeBlockActions
+          disabled={props.isIncomplete}
+          language={codeFence.language}
+          source={codeFence.source}
+        />
+      </div>
+    );
+  }
+
   return <Block {...props} />;
 }
 const AUDIO_PLAYER_RE = /<audio-player\s+src="([^"]+)"\s*\/>/;
 
 const MarkdownTextImpl = () => {
-  const { text } = useMessagePartText();
-  const status = useSmoothStatus();
+  const { text, status } = useMessagePartText();
 
   const audioMatch = text.match(AUDIO_PLAYER_RE);
   if (audioMatch) {
@@ -96,6 +218,7 @@ const MarkdownTextImpl = () => {
         isAnimating={status.type === "running"}
         plugins={{ code, math, mermaid }}
         controls={{
+          code: false,
           mermaid: {
             fullscreen: true,
             download: true,