Add screenshot paste support (#2679)

ui/desktop/src/components/ChatInput.tsx

@@ -2,12 +2,24 @@ import React, { useRef, useState, useEffect, useCallback } from 'react';
 import { Button } from './ui/button';
 import type { View } from '../App';
 import Stop from './ui/Stop';
-import { Attach, Send } from './icons';
+import { Attach, Send, Close } from './icons';
 import { debounce } from 'lodash';
 import BottomMenu from './bottom_menu/BottomMenu';
 import { LocalMessageStorage } from '../utils/localMessageStorage';
 import { Message } from '../types/message';
 
+interface PastedImage {
+  id: string;
+  dataUrl: string; // For immediate preview
+  filePath?: string; // Path on filesystem after saving
+  isLoading: boolean;
+  error?: string;
+}
+
+// Constants for image handling
+const MAX_IMAGES_PER_MESSAGE = 5;
+const MAX_IMAGE_SIZE_MB = 5;
+
 interface ChatInputProps {
   handleSubmit: (e: React.FormEvent) => void;
   isLoading?: boolean;
@@ -37,15 +49,28 @@ export default function ChatInput({
   const [_value, setValue] = useState(initialValue);
   const [displayValue, setDisplayValue] = useState(initialValue); // For immediate visual feedback
   const [isFocused, setIsFocused] = useState(false);
+  const [pastedImages, setPastedImages] = useState<PastedImage[]>([]);
 
   // Update internal value when initialValue changes
   useEffect(() => {
     setValue(initialValue);
     setDisplayValue(initialValue);
+
+    // Use a functional update to get the current pastedImages
+    // and perform cleanup. This avoids needing pastedImages in the deps.
+    setPastedImages((currentPastedImages) => {
+      currentPastedImages.forEach((img) => {
+        if (img.filePath) {
+          window.electron.deleteTempFile(img.filePath);
+        }
+      });
+      return []; // Return a new empty array
+    });
+
     // Reset history index when input is cleared
     setHistoryIndex(-1);
     setIsInGlobalHistory(false);
-  }, [initialValue]);
+  }, [initialValue]); // Keep only initialValue as a dependency
 
   // State to track if the IME is composing (i.e., in the middle of Japanese IME input)
   const [isComposing, setIsComposing] = useState(false);
@@ -55,6 +80,48 @@ export default function ChatInput({
   const textAreaRef = useRef<HTMLTextAreaElement>(null);
   const [processedFilePaths, setProcessedFilePaths] = useState<string[]>([]);
 
+  const handleRemovePastedImage = (idToRemove: string) => {
+    const imageToRemove = pastedImages.find((img) => img.id === idToRemove);
+    if (imageToRemove?.filePath) {
+      window.electron.deleteTempFile(imageToRemove.filePath);
+    }
+    setPastedImages((currentImages) => currentImages.filter((img) => img.id !== idToRemove));
+  };
+
+  const handleRetryImageSave = async (imageId: string) => {
+    const imageToRetry = pastedImages.find((img) => img.id === imageId);
+    if (!imageToRetry || !imageToRetry.dataUrl) return;
+
+    // Set the image to loading state
+    setPastedImages((prev) =>
+      prev.map((img) =>
+        img.id === imageId
+          ? { ...img, isLoading: true, error: undefined }
+          : img
+      )
+    );
+
+    try {
+      const result = await window.electron.saveDataUrlToTemp(imageToRetry.dataUrl, imageId);
+      setPastedImages((prev) =>
+        prev.map((img) =>
+          img.id === result.id
+            ? { ...img, filePath: result.filePath, error: result.error, isLoading: false }
+            : img
+        )
+      );
+    } catch (err) {
+      console.error('Error retrying image save:', err);
+      setPastedImages((prev) =>
+        prev.map((img) =>
+          img.id === imageId
+            ? { ...img, error: 'Failed to save image via Electron.', isLoading: false }
+            : img
+        )
+      );
+    }
+  };
+
   useEffect(() => {
     if (textAreaRef.current) {
       textAreaRef.current.focus();
@@ -111,6 +178,89 @@ export default function ChatInput({
     debouncedSetValue(val); // Debounce the actual state update
   };
 
+  const handlePaste = async (evt: React.ClipboardEvent<HTMLTextAreaElement>) => {
+    const files = Array.from(evt.clipboardData.files || []);
+    const imageFiles = files.filter(file => file.type.startsWith('image/'));
+    
+    if (imageFiles.length === 0) return;
+    
+    // Check if adding these images would exceed the limit
+    if (pastedImages.length + imageFiles.length > MAX_IMAGES_PER_MESSAGE) {
+      // Show error message to user
+      setPastedImages((prev) => [
+        ...prev,
+        {
+          id: `error-${Date.now()}`,
+          dataUrl: '',
+          isLoading: false,
+          error: `Cannot paste ${imageFiles.length} image(s). Maximum ${MAX_IMAGES_PER_MESSAGE} images per message allowed.`
+        }
+      ]);
+      
+      // Remove the error message after 3 seconds
+      setTimeout(() => {
+        setPastedImages((prev) => prev.filter(img => !img.id.startsWith('error-')));
+      }, 3000);
+      
+      return;
+    }
+    
+    evt.preventDefault();
+    
+    for (const file of imageFiles) {
+      // Check individual file size before processing
+      if (file.size > MAX_IMAGE_SIZE_MB * 1024 * 1024) {
+        const errorId = `error-${Date.now()}-${Math.random().toString(36).substring(2, 9)}`;
+        setPastedImages((prev) => [
+          ...prev,
+          {
+            id: errorId,
+            dataUrl: '',
+            isLoading: false,
+            error: `Image too large (${Math.round(file.size / (1024 * 1024))}MB). Maximum ${MAX_IMAGE_SIZE_MB}MB allowed.`
+          }
+        ]);
+        
+        // Remove the error message after 3 seconds
+        setTimeout(() => {
+          setPastedImages((prev) => prev.filter(img => img.id !== errorId));
+        }, 3000);
+        
+        continue;
+      }
+      
+      const reader = new FileReader();
+      reader.onload = async (e) => {
+        const dataUrl = e.target?.result as string;
+        if (dataUrl) {
+          const imageId = `img-${Date.now()}-${Math.random().toString(36).substring(2, 9)}`;
+          setPastedImages((prev) => [...prev, { id: imageId, dataUrl, isLoading: true }]);
+
+          try {
+            const result = await window.electron.saveDataUrlToTemp(dataUrl, imageId);
+            setPastedImages((prev) =>
+              prev.map((img) =>
+                img.id === result.id
+                  ? { ...img, filePath: result.filePath, error: result.error, isLoading: false }
+                  : img
+              )
+            );
+          } catch (err) {
+            console.error('Error saving pasted image:', err);
+            setPastedImages((prev) =>
+              prev.map((img) =>
+                img.id === imageId
+                  ? { ...img, error: 'Failed to save image via Electron.', isLoading: false }
+                  : img
+              )
+            );
+          }
+        }
+      };
+      reader.readAsDataURL(file);
+    }
+  };
+
   // Cleanup debounced functions on unmount
   useEffect(() => {
     return () => {
@@ -197,6 +347,36 @@ export default function ChatInput({
     }
   };
 
+  const performSubmit = () => {
+    const validPastedImageFilesPaths = pastedImages
+      .filter((img) => img.filePath && !img.error && !img.isLoading)
+      .map((img) => img.filePath as string);
+
+    let textToSend = displayValue.trim();
+
+    if (validPastedImageFilesPaths.length > 0) {
+      const pathsString = validPastedImageFilesPaths.join(' ');
+      textToSend = textToSend ? `${textToSend} ${pathsString}` : pathsString;
+    }
+
+    if (textToSend) {
+      if (displayValue.trim()) {
+        LocalMessageStorage.addMessage(displayValue);
+      } else if (validPastedImageFilesPaths.length > 0) {
+        LocalMessageStorage.addMessage(validPastedImageFilesPaths.join(' '));
+      }
+
+      handleSubmit(new CustomEvent('submit', { detail: { value: textToSend } }));
+
+      setDisplayValue('');
+      setValue('');
+      setPastedImages([]);
+      setHistoryIndex(-1);
+      setSavedInput('');
+      setIsInGlobalHistory(false);
+    }
+  };
+
   const handleKeyDown = (evt: React.KeyboardEvent<HTMLTextAreaElement>) => {
     // Handle history navigation first
     handleHistoryNavigation(evt);
@@ -207,6 +387,7 @@ export default function ChatInput({
         // Allow line break for Shift+Enter, or during IME composition
         return;
       }
+
       if (evt.altKey) {
         const newValue = displayValue + '\n';
         setDisplayValue(newValue);
@@ -214,44 +395,31 @@ export default function ChatInput({
         return;
       }
 
-      // Prevent default Enter behavior when loading or when not loading but has content
-      // So it won't trigger a new line
       evt.preventDefault();
-
-      // Only submit if not loading and has content
-      if (!isLoading && displayValue.trim()) {
-        // Always add to global chat storage before submitting
-        LocalMessageStorage.addMessage(displayValue);
-
-        handleSubmit(new CustomEvent('submit', { detail: { value: displayValue } }));
-        setDisplayValue('');
-        setValue('');
-        setHistoryIndex(-1);
-        setSavedInput('');
-        setIsInGlobalHistory(false);
+      const canSubmit =
+        !isLoading &&
+        (displayValue.trim() ||
+          pastedImages.some((img) => img.filePath && !img.error && !img.isLoading));
+      if (canSubmit) {
+        performSubmit();
       }
     }
   };
 
   const onFormSubmit = (e: React.FormEvent) => {
     e.preventDefault();
-    if (displayValue.trim() && !isLoading) {
-      // Always add to global chat storage before submitting
-      LocalMessageStorage.addMessage(displayValue);
-
-      handleSubmit(new CustomEvent('submit', { detail: { value: displayValue } }));
-      setDisplayValue('');
-      setValue('');
-      setHistoryIndex(-1);
-      setSavedInput('');
-      setIsInGlobalHistory(false);
+    const canSubmit =
+      !isLoading &&
+      (displayValue.trim() ||
+        pastedImages.some((img) => img.filePath && !img.error && !img.isLoading));
+    if (canSubmit) {
+      performSubmit();
     }
   };
 
   const handleFileSelect = async () => {
     const path = await window.electron.selectFileOrDirectory();
     if (path) {
-      // Append the path to existing text, with a space if there's existing text
       const newValue = displayValue.trim() ? `${displayValue.trim()} ${path}` : path;
       setDisplayValue(newValue);
       setValue(newValue);
@@ -259,6 +427,10 @@ export default function ChatInput({
     }
   };
 
+  const hasSubmittableContent =
+    displayValue.trim() || pastedImages.some((img) => img.filePath && !img.error && !img.isLoading);
+  const isAnyImageLoading = pastedImages.some((img) => img.isLoading);
+
   return (
     <div
       className={`flex flex-col relative h-auto border rounded-lg transition-colors ${
@@ -278,6 +450,7 @@ export default function ChatInput({
           onCompositionStart={handleCompositionStart}
           onCompositionEnd={handleCompositionEnd}
           onKeyDown={handleKeyDown}
+          onPaste={handlePaste}
           onFocus={() => setIsFocused(true)}
           onBlur={() => setIsFocused(false)}
           ref={textAreaRef}
@@ -290,6 +463,54 @@ export default function ChatInput({
           className="w-full pl-4 pr-[68px] outline-none border-none focus:ring-0 bg-transparent pt-3 pb-1.5 text-sm resize-none text-textStandard placeholder:text-textPlaceholder"
         />
 
+        {pastedImages.length > 0 && (
+          <div className="flex flex-wrap gap-2 p-2 border-t border-borderSubtle">
+            {pastedImages.map((img) => (
+              <div key={img.id} className="relative group w-20 h-20">
+                {img.dataUrl && (
+                  <img
+                    src={img.dataUrl} // Use dataUrl for instant preview
+                    alt={`Pasted image ${img.id}`}
+                    className={`w-full h-full object-cover rounded border ${img.error ? 'border-red-500' : 'border-borderStandard'}`}
+                  />
+                )}
+                {img.isLoading && (
+                  <div className="absolute inset-0 flex items-center justify-center bg-black bg-opacity-50 rounded">
+                    <div className="animate-spin rounded-full h-6 w-6 border-t-2 border-b-2 border-white"></div>
+                  </div>
+                )}
+                {img.error && !img.isLoading && (
+                  <div className="absolute inset-0 flex flex-col items-center justify-center bg-black bg-opacity-75 rounded p-1 text-center">
+                    <p className="text-red-400 text-[10px] leading-tight break-all mb-1">
+                      {img.error.substring(0, 50)}
+                    </p>
+                    {img.dataUrl && (
+                      <button
+                        type="button"
+                        onClick={() => handleRetryImageSave(img.id)}
+                        className="bg-blue-600 hover:bg-blue-700 text-white rounded px-1 py-0.5 text-[8px] leading-none"
+                        title="Retry saving image"
+                      >
+                        Retry
+                      </button>
+                    )}
+                  </div>
+                )}
+                {!img.isLoading && (
+                  <button
+                    type="button"
+                    onClick={() => handleRemovePastedImage(img.id)}
+                    className="absolute -top-1 -right-1 bg-gray-700 hover:bg-red-600 text-white rounded-full w-5 h-5 flex items-center justify-center text-xs leading-none opacity-0 group-hover:opacity-100 focus:opacity-100 transition-opacity z-10"
+                    aria-label="Remove image"
+                  >
+                    <Close size={14} />
+                  </button>
+                )}
+              </div>
+            ))}
+          </div>
+        )}
+
         {isLoading ? (
           <Button
             type="button"
@@ -309,12 +530,13 @@ export default function ChatInput({
             type="submit"
             size="icon"
             variant="ghost"
-            disabled={!displayValue.trim()}
+            disabled={!hasSubmittableContent || isAnyImageLoading} // Disable if no content or if images are still loading/saving
             className={`absolute right-3 top-2 transition-colors rounded-full w-7 h-7 [&_svg]:size-4 ${
-              !displayValue.trim()
+              !hasSubmittableContent || isAnyImageLoading
                 ? 'text-textSubtle cursor-not-allowed'
                 : 'bg-bgAppInverse text-textProminentInverse hover:cursor-pointer'
             }`}
+            title={isAnyImageLoading ? 'Waiting for images to save...' : 'Send'}
           >
             <Send />
           </Button>