refactor: Store pasted images in memory as base64 data URLs

123vivekr · 123vivekr · commit 4cd104b1ddcc · 2025-07-06T16:45:12.000+05:30
- Remove save_clipboard_image and cleanup_temp_images backend commands
- Update FloatingPromptInput to store pasted images as data URLs in the prompt
- Update ImagePreview component to handle both file paths and data URLs
- Update extractImagePaths to properly handle quoted data URLs
- Update handleRemoveImage to handle data URL removal

This eliminates file system operations for pasted images and stores them directly
in the prompt as base64 data URLs (e.g., @"data:image/png;base64,...").
Images are now fully self-contained within the session without creating temp files.
diff --git a/src-tauri/src/commands/claude.rs b/src-tauri/src/commands/claude.rs
@@ -922,104 +922,7 @@ pub async fn load_session_history(
     Ok(messages)
 }
 
-/// Saves a clipboard image to a temporary location for a session
-/// 
-/// This command handles pasted images by saving them to a temp directory
-/// within the project structure, making them accessible for Claude.
-#[tauri::command]
-pub async fn save_clipboard_image(
-    project_path: String,
-    session_id: String,
-    image_data: String, // Base64 encoded image data
-    mime_type: String,
-) -> Result<String, String> {
-    log::info!(
-        "Saving clipboard image for session: {} in project: {}",
-        session_id,
-        project_path
-    );
 
-    // Parse the base64 data (remove data URL prefix if present)
-    let base64_data = if image_data.starts_with("data:") {
-        image_data
-            .split(',')
-            .nth(1)
-            .ok_or("Invalid data URL format")?
-    } else {
-        &image_data
-    };
-
-    // Decode base64 to bytes
-    use base64::Engine;
-    let image_bytes = base64::engine::general_purpose::STANDARD
-        .decode(base64_data)
-        .map_err(|e| format!("Failed to decode base64 image: {}", e))?;
-
-    // Determine file extension from MIME type
-    let extension = match mime_type.as_str() {
-        "image/png" => "png",
-        "image/jpeg" | "image/jpg" => "jpg",
-        "image/gif" => "gif",
-        "image/webp" => "webp",
-        "image/svg+xml" => "svg",
-        _ => "png", // Default to PNG
-    };
-
-    // Create temp directory for the project if it doesn't exist
-    let project_path_buf = PathBuf::from(&project_path);
-    let temp_dir = project_path_buf.join(".claude_temp").join(&session_id);
-    fs::create_dir_all(&temp_dir)
-        .map_err(|e| format!("Failed to create temp directory: {}", e))?;
-
-    // Generate unique filename with timestamp
-    let timestamp = SystemTime::now()
-        .duration_since(SystemTime::UNIX_EPOCH)
-        .unwrap_or_default()
-        .as_millis();
-    let filename = format!("pasted_image_{}.{}", timestamp, extension);
-    let file_path = temp_dir.join(&filename);
-
-    // Write image to file
-    fs::write(&file_path, image_bytes)
-        .map_err(|e| format!("Failed to write image file: {}", e))?;
-
-    // Return the absolute path
-    let absolute_path = file_path
-        .canonicalize()
-        .map_err(|e| format!("Failed to get absolute path: {}", e))?
-        .to_string_lossy()
-        .to_string();
-
-    log::info!("Saved clipboard image to: {}", absolute_path);
-    Ok(absolute_path)
-}
-
-/// Cleans up temporary images for a session
-/// 
-/// This command removes the temporary directory created for pasted images
-/// when a session ends to avoid accumulating temporary files.
-#[tauri::command]
-pub async fn cleanup_temp_images(
-    project_path: String,
-    session_id: String,
-) -> Result<(), String> {
-    log::info!(
-        "Cleaning up temp images for session: {} in project: {}",
-        session_id,
-        project_path
-    );
-
-    let project_path_buf = PathBuf::from(&project_path);
-    let temp_dir = project_path_buf.join(".claude_temp").join(&session_id);
-
-    if temp_dir.exists() {
-        fs::remove_dir_all(&temp_dir)
-            .map_err(|e| format!("Failed to remove temp directory: {}", e))?;
-        log::info!("Cleaned up temp directory: {}", temp_dir.display());
-    }
-
-    Ok(())
-}
 
 /// Execute a new interactive Claude Code session with streaming output
 #[tauri::command]
diff --git a/src-tauri/src/main.rs b/src-tauri/src/main.rs
@@ -18,13 +18,13 @@ use commands::agents::{
 };
 use commands::claude::{
     cancel_claude_execution, check_auto_checkpoint, check_claude_version, cleanup_old_checkpoints,
-    cleanup_temp_images, clear_checkpoint_manager, continue_claude_code, create_checkpoint, execute_claude_code,
+    clear_checkpoint_manager, continue_claude_code, create_checkpoint, execute_claude_code,
     find_claude_md_files, fork_from_checkpoint, get_checkpoint_diff, get_checkpoint_settings,
     get_checkpoint_state_stats, get_claude_session_output, get_claude_settings, get_project_sessions,
     get_recently_modified_files, get_session_timeline, get_system_prompt, list_checkpoints,
     list_directory_contents, list_projects, list_running_claude_sessions, load_session_history,
     open_new_session, read_claude_md_file, restore_checkpoint, resume_claude_code,
-    save_claude_md_file, save_claude_settings, save_clipboard_image, save_system_prompt, search_files,
+    save_claude_md_file, save_claude_settings, save_system_prompt, search_files,
     track_checkpoint_message, track_session_messages, update_checkpoint_settings,
     get_hooks_config, update_hooks_config, validate_hook_command,
     ClaudeProcessState,
@@ -101,8 +101,6 @@ fn main() {
             find_claude_md_files,
             read_claude_md_file,
             save_claude_md_file,
-            save_clipboard_image,
-            cleanup_temp_images,
             load_session_history,
             execute_claude_code,
             continue_claude_code,
diff --git a/src/components/ClaudeCodeSession.tsx b/src/components/ClaudeCodeSession.tsx
@@ -20,7 +20,6 @@ import { api, type Session } from "@/lib/api";
 import { cn } from "@/lib/utils";
 import { open } from "@tauri-apps/plugin-dialog";
 import { listen, type UnlistenFn } from "@tauri-apps/api/event";
-import { invoke } from "@tauri-apps/api/core";
 import { StreamMessage } from "./StreamMessage";
 import { FloatingPromptInput, type FloatingPromptInputRef } from "./FloatingPromptInput";
 import { ErrorBoundary } from "./ErrorBoundary";
@@ -828,16 +827,6 @@ export const ClaudeCodeSession: React.FC<ClaudeCodeSessionProps> = ({
         api.clearCheckpointManager(effectiveSession.id).catch(err => {
           console.error("Failed to clear checkpoint manager:", err);
         });
-        
-        // Clean up temporary images
-        if (projectPath) {
-          invoke('cleanup_temp_images', {
-            projectPath,
-            sessionId: effectiveSession.id
-          }).catch((err: any) => {
-            console.error("Failed to cleanup temp images:", err);
-          });
-        }
       }
     };
   }, [effectiveSession, projectPath]);
diff --git a/src/components/FloatingPromptInput.tsx b/src/components/FloatingPromptInput.tsx
@@ -19,7 +19,6 @@ import { FilePicker } from "./FilePicker";
 import { ImagePreview } from "./ImagePreview";
 import { type FileEntry } from "@/lib/api";
 import { getCurrentWebviewWindow } from "@tauri-apps/api/webviewWindow";
-import { invoke } from "@tauri-apps/api/core";
 
 interface FloatingPromptInputProps {
   /**
@@ -220,33 +219,39 @@ const FloatingPromptInputInner = (
 
   // Helper function to check if a file is an image
   const isImageFile = (path: string): boolean => {
+    // Check if it's a data URL
+    if (path.startsWith('data:image/')) {
+      return true;
+    }
+    // Otherwise check file extension
     const ext = path.split('.').pop()?.toLowerCase();
     return ['png', 'jpg', 'jpeg', 'gif', 'svg', 'webp', 'ico', 'bmp'].includes(ext || '');
   };
 
   // Extract image paths from prompt text
   const extractImagePaths = (text: string): string[] => {
-    console.log('[extractImagePaths] Input text:', text);
+    console.log('[extractImagePaths] Input text length:', text.length);
     
     // Updated regex to handle both quoted and unquoted paths
-    // Pattern 1: @"path with spaces" - quoted paths
+    // Pattern 1: @"path with spaces or data URLs" - quoted paths
     // Pattern 2: @path - unquoted paths (continues until @ or end)
     const quotedRegex = /@"([^"]+)"/g;
     const unquotedRegex = /@([^@\n\s]+)/g;
     
     const pathsSet = new Set<string>(); // Use Set to ensure uniqueness
     
-    // First, extract quoted paths
+    // First, extract quoted paths (including data URLs)
     let matches = Array.from(text.matchAll(quotedRegex));
-    console.log('[extractImagePaths] Quoted matches:', matches.map(m => m[0]));
+    console.log('[extractImagePaths] Quoted matches:', matches.length);
     
     for (const match of matches) {
       const path = match[1]; // No need to trim, quotes preserve exact path
-      console.log('[extractImagePaths] Processing quoted path:', path);
+      console.log('[extractImagePaths] Processing quoted path:', path.startsWith('data:') ? 'data URL' : path);
       
-      // Convert relative path to absolute if needed
-      const fullPath = path.startsWith('/') ? path : (projectPath ? `${projectPath}/${path}` : path);
-      console.log('[extractImagePaths] Full path:', fullPath, 'Is image:', isImageFile(fullPath));
+      // For data URLs, use as-is; for file paths, convert to absolute
+      const fullPath = path.startsWith('data:') 
+        ? path 
+        : (path.startsWith('/') ? path : (projectPath ? `${projectPath}/${path}` : path));
       
       if (isImageFile(fullPath)) {
         pathsSet.add(fullPath);
@@ -256,25 +261,27 @@ const FloatingPromptInputInner = (
     // Remove quoted mentions from text to avoid double-matching
     let textWithoutQuoted = text.replace(quotedRegex, '');
     
-    // Then extract unquoted paths
+    // Then extract unquoted paths (typically file paths)
     matches = Array.from(textWithoutQuoted.matchAll(unquotedRegex));
-    console.log('[extractImagePaths] Unquoted matches:', matches.map(m => m[0]));
+    console.log('[extractImagePaths] Unquoted matches:', matches.length);
     
     for (const match of matches) {
       const path = match[1].trim();
+      // Skip if it looks like a data URL fragment (shouldn't happen with proper quoting)
+      if (path.includes('data:')) continue;
+      
       console.log('[extractImagePaths] Processing unquoted path:', path);
       
       // Convert relative path to absolute if needed
       const fullPath = path.startsWith('/') ? path : (projectPath ? `${projectPath}/${path}` : path);
-      console.log('[extractImagePaths] Full path:', fullPath, 'Is image:', isImageFile(fullPath));
       
       if (isImageFile(fullPath)) {
         pathsSet.add(fullPath);
       }
     }
 
-    const uniquePaths = Array.from(pathsSet); // Convert Set back to Array
-    console.log('[extractImagePaths] Final extracted paths (unique):', uniquePaths);
+    const uniquePaths = Array.from(pathsSet);
+    console.log('[extractImagePaths] Final extracted paths (unique):', uniquePaths.length);
     return uniquePaths;
   };
 
@@ -479,7 +486,7 @@ const FloatingPromptInputInner = (
 
   const handlePaste = async (e: React.ClipboardEvent) => {
     const items = e.clipboardData?.items;
-    if (!items || !projectPath) return;
+    if (!items) return;
 
     for (const item of items) {
       if (item.type.startsWith('image/')) {
@@ -492,24 +499,13 @@ const FloatingPromptInputInner = (
         try {
           // Convert blob to base64
           const reader = new FileReader();
-          reader.onload = async () => {
+          reader.onload = () => {
             const base64Data = reader.result as string;
             
-            // Generate a session-specific ID for the image
-            const sessionId = `paste-${Date.now()}`;
-            
-            // Save the image via Tauri command
-            const imagePath = await invoke<string>('save_clipboard_image', {
-              projectPath,
-              sessionId,
-              imageData: base64Data,
-              mimeType: item.type
-            });
-
-            // Add the image path as a mention to the prompt
+            // Add the base64 data URL directly to the prompt
             setPrompt(currentPrompt => {
-              // Wrap path in quotes if it contains spaces
-              const mention = imagePath.includes(' ') ? `@"${imagePath}"` : `@${imagePath}`;
+              // Use the data URL directly as the image reference
+              const mention = `@"${base64Data}"`;
               const newPrompt = currentPrompt + (currentPrompt.endsWith(' ') || currentPrompt === '' ? '' : ' ') + mention + ' ';
               
               // Focus the textarea and move cursor to end
@@ -548,6 +544,17 @@ const FloatingPromptInputInner = (
   const handleRemoveImage = (index: number) => {
     // Remove the corresponding @mention from the prompt
     const imagePath = embeddedImages[index];
+    
+    // For data URLs, we need to handle them specially since they're always quoted
+    if (imagePath.startsWith('data:')) {
+      // Simply remove the exact quoted data URL
+      const quotedPath = `@"${imagePath}"`;
+      const newPrompt = prompt.replace(quotedPath, '').trim();
+      setPrompt(newPrompt);
+      return;
+    }
+    
+    // For file paths, use the original logic
     const escapedPath = imagePath.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
     const escapedRelativePath = imagePath.replace(projectPath + '/', '').replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
     
diff --git a/src/components/ImagePreview.tsx b/src/components/ImagePreview.tsx
@@ -56,6 +56,16 @@ export const ImagePreview: React.FC<ImagePreviewProps> = ({
     onRemove(index);
   };
 
+  // Helper to get the image source - handles both file paths and data URLs
+  const getImageSrc = (imagePath: string): string => {
+    // If it's already a data URL, return as-is
+    if (imagePath.startsWith('data:')) {
+      return imagePath;
+    }
+    // Otherwise, convert the file path
+    return convertFileSrc(imagePath);
+  };
+
   if (displayImages.length === 0) return null;
 
   return (
@@ -83,7 +93,7 @@ export const ImagePreview: React.FC<ImagePreviewProps> = ({
                   </div>
                 ) : (
                   <img
-                    src={convertFileSrc(imagePath)}
+                    src={getImageSrc(imagePath)}
                     alt={`Preview ${index + 1}`}
                     className="w-full h-full object-cover"
                     onError={() => handleImageError(index)}
@@ -131,7 +141,7 @@ export const ImagePreview: React.FC<ImagePreviewProps> = ({
           {selectedImageIndex !== null && (
             <div className="relative w-full h-full flex items-center justify-center p-4">
               <img
-                src={convertFileSrc(displayImages[selectedImageIndex])}
+                src={getImageSrc(displayImages[selectedImageIndex])}
                 alt={`Full preview ${selectedImageIndex + 1}`}
                 className="max-w-full max-h-full object-contain"
                 onError={() => handleImageError(selectedImageIndex)}
@@ -164,4 +174,4 @@ export const ImagePreview: React.FC<ImagePreviewProps> = ({
       </Dialog>
     </>
   );
-}; 
+}; 
