From 31e379bc09dd7973d3359ca6a3e50fe0f8ed707a Mon Sep 17 00:00:00 2001
From: wwayne <wayne.wang0821@gmail.com>
Date: Thu, 30 May 2024 12:58:55 +0800
Subject: [PATCH] fix(vscode): copy function for vscode and remove codeblock
 download

---
 clients/vscode/src/chat/ChatViewProvider.ts   | 16 +++++++-
 ee/tabby-ui/components/ui/codeblock.tsx       | 39 +------------------
 .../lib/hooks/use-copy-to-clipboard.tsx       | 29 +++++++-------
 3 files changed, 28 insertions(+), 56 deletions(-)

diff --git a/clients/vscode/src/chat/ChatViewProvider.ts b/clients/vscode/src/chat/ChatViewProvider.ts
index e1510b0f4466..fdf4949c9a13 100644
--- a/clients/vscode/src/chat/ChatViewProvider.ts
+++ b/clients/vscode/src/chat/ChatViewProvider.ts
@@ -62,8 +62,15 @@ export class ChatViewProvider implements WebviewViewProvider {
     });
 
     webviewView.webview.onDidReceiveMessage(async (message) => {
-      if (message.action === "rendered") {
-        await this.initChatPage();
+      switch (message.action) {
+        case "rendered": {
+          await this.initChatPage();
+          return;
+        }
+        case "copy": {
+          env.clipboard.writeText(message.data);
+          return;
+        }
       }
     });
 
@@ -165,6 +172,10 @@ export class ChatViewProvider implements WebviewViewProvider {
                     syncTheme();
                     return;
                   }
+                  if (event.data.action === 'copy') {
+                    vscode.postMessage(event.data);
+                    return;
+                  }
 
                   if (event.data.data) {
                     chatIframe.contentWindow.postMessage(event.data.data[0], "${endpoint}");
@@ -180,6 +191,7 @@ export class ChatViewProvider implements WebviewViewProvider {
           <iframe
             id="chat"
             src="${endpoint}/chat?from=vscode"
+            allow="clipboard-read; clipboard-write"
             onload="iframeLoaded(this)" />
         </body>
       </html>
diff --git a/ee/tabby-ui/components/ui/codeblock.tsx b/ee/tabby-ui/components/ui/codeblock.tsx
index 0875dda6abd7..7e384a38e496 100644
--- a/ee/tabby-ui/components/ui/codeblock.tsx
+++ b/ee/tabby-ui/components/ui/codeblock.tsx
@@ -9,7 +9,7 @@ import { coldarkDark } from 'react-syntax-highlighter/dist/cjs/styles/prism'
 
 import { useCopyToClipboard } from '@/lib/hooks/use-copy-to-clipboard'
 import { Button } from '@/components/ui/button'
-import { IconCheck, IconCopy, IconDownload } from '@/components/ui/icons'
+import { IconCheck, IconCopy } from '@/components/ui/icons'
 
 interface Props {
   language: string
@@ -59,34 +59,6 @@ export const generateRandomString = (length: number, lowercase = false) => {
 const CodeBlock: FC<Props> = memo(({ language, value }) => {
   const { isCopied, copyToClipboard } = useCopyToClipboard({ timeout: 2000 })
 
-  const downloadAsFile = () => {
-    if (typeof window === 'undefined') {
-      return
-    }
-    const fileExtension = programmingLanguages[language] || '.file'
-    const suggestedFileName = `file-${generateRandomString(
-      3,
-      true
-    )}${fileExtension}`
-    const fileName = window.prompt('Enter file name' || '', suggestedFileName)
-
-    if (!fileName) {
-      // User pressed cancel on prompt.
-      return
-    }
-
-    const blob = new Blob([value], { type: 'text/plain' })
-    const url = URL.createObjectURL(blob)
-    const link = document.createElement('a')
-    link.download = fileName
-    link.href = url
-    link.style.display = 'none'
-    document.body.appendChild(link)
-    link.click()
-    document.body.removeChild(link)
-    URL.revokeObjectURL(url)
-  }
-
   const onCopy = () => {
     if (isCopied) return
     copyToClipboard(value)
@@ -97,15 +69,6 @@ const CodeBlock: FC<Props> = memo(({ language, value }) => {
       <div className="flex w-full items-center justify-between bg-zinc-800 px-6 py-2 pr-4 text-zinc-100">
         <span className="text-xs lowercase">{language}</span>
         <div className="flex items-center space-x-1">
-          <Button
-            variant="ghost"
-            className="hover:bg-[#3C382F] hover:text-[#F4F4F5] focus-visible:ring-1 focus-visible:ring-slate-700 focus-visible:ring-offset-0"
-            onClick={downloadAsFile}
-            size="icon"
-          >
-            <IconDownload />
-            <span className="sr-only">Download</span>
-          </Button>
           <Button
             variant="ghost"
             size="icon"
diff --git a/ee/tabby-ui/lib/hooks/use-copy-to-clipboard.tsx b/ee/tabby-ui/lib/hooks/use-copy-to-clipboard.tsx
index 7b5d48779b02..da734de8fbab 100644
--- a/ee/tabby-ui/lib/hooks/use-copy-to-clipboard.tsx
+++ b/ee/tabby-ui/lib/hooks/use-copy-to-clipboard.tsx
@@ -2,16 +2,13 @@
 
 import * as React from 'react'
 import copy from 'copy-to-clipboard'
-import { toast } from 'sonner'
 
 export interface useCopyToClipboardProps {
   timeout?: number
-  onError?: (e?: any) => void
 }
 
 export function useCopyToClipboard({
-  timeout = 2000,
-  onError
+  timeout = 2000
 }: useCopyToClipboardProps) {
   const [isCopied, setIsCopied] = React.useState<Boolean>(false)
 
@@ -22,15 +19,6 @@ export function useCopyToClipboard({
     }, timeout)
   }
 
-  const onCopyError = (error?: any) => {
-    if (typeof onError === 'function') {
-      onError?.(error)
-      return
-    }
-
-    toast.error('Failed to copy.')
-  }
-
   const copyToClipboard = (value: string) => {
     if (typeof window === 'undefined') return
     if (!value) return
@@ -39,15 +27,24 @@ export function useCopyToClipboard({
       navigator.clipboard
         .writeText(value)
         .then(onCopySuccess)
-        .catch(onCopyError)
+        .catch(() => {})
     } else {
       const copyResult = copy(value)
       if (copyResult) {
         onCopySuccess()
-      } else {
-        onCopyError()
       }
     }
+
+    // When component inside an iframe sandbox(VSCode)
+    // We need to notify parent environment to handle the copy event
+    parent.postMessage(
+      {
+        action: 'copy',
+        data: value
+      },
+      '*'
+    )
+    onCopySuccess()
   }
 
   return { isCopied, copyToClipboard }