mcp-memory/server/src/handlers/vision.rs

use crate::router::McpTool;
use crate::state::MemoryState;
use crate::tools::{ClipboardAction, ClipboardTool};
use arboard::{Clipboard, ImageData};
use async_trait::async_trait;
use image::{ImageBuffer, imageops::FilterType};
use serde_json::{Value, json};
use std::borrow::Cow;
use std::sync::Arc;


static CLIPBOARD_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());

pub fn get_native_clipboard_text() -> Option<String> {
    let _guard = CLIPBOARD_LOCK.lock().unwrap_or_else(|e| e.into_inner());
    for _ in 0..3 {
        if let Ok(mut clipboard) = arboard::Clipboard::new()
            && let Ok(text) = clipboard.get_text()
            && !text.trim().is_empty()
        {
            return Some(text);
        }
        std::thread::sleep(std::time::Duration::from_millis(20));
    }
    None
}

#[cfg(target_os = "windows")]
fn get_windows_png_clipboard_image() -> Option<image::DynamicImage> {
    use windows_sys::Win32::System::DataExchange::{
        CloseClipboard, GetClipboardData, IsClipboardFormatAvailable, OpenClipboard,
        RegisterClipboardFormatW,
    };
    use windows_sys::Win32::System::Memory::{GlobalLock, GlobalSize, GlobalUnlock};
    let format_name: Vec<u16> = "PNG\0".encode_utf16().collect();
    unsafe {
        let format_id = RegisterClipboardFormatW(format_name.as_ptr());
        if format_id == 0 {
            return None;
        }
        for _ in 0..3 {
            if OpenClipboard(std::ptr::null_mut()) != 0 {
                let has_format = IsClipboardFormatAvailable(format_id) != 0;
                let mut img_opt = None;
                if has_format {
                    let handle = GetClipboardData(format_id);
                    if !handle.is_null() {
                        let size = GlobalSize(handle as _);
                        if size > 0 {
                            let ptr = GlobalLock(handle as _);
                            if !ptr.is_null() {
                                let slice = std::slice::from_raw_parts(ptr as *const u8, size);
                                img_opt = image::load_from_memory(slice).ok();
                                GlobalUnlock(handle as _);
                            }
                        }
                    }
                }
                CloseClipboard();
                if img_opt.is_some() {
                    return img_opt;
                }
                break;
            }
            std::thread::sleep(std::time::Duration::from_millis(20));
        }
    }
    None
}
pub fn get_native_clipboard_image() -> Option<image::DynamicImage> {
    let _guard = CLIPBOARD_LOCK.lock().unwrap_or_else(|e| e.into_inner());
    #[cfg(target_os = "windows")]
    if let Some(img) = get_windows_png_clipboard_image() {
        return Some(img);
    }
    for _ in 0..3 {
        if let Ok(mut clipboard) = arboard::Clipboard::new()
            && let Ok(image_data) = clipboard.get_image()
            && let Some(img) = ImageBuffer::<image::Rgba<u8>, _>::from_raw(
                image_data.width as u32,
                image_data.height as u32,
                image_data.bytes.into_owned(),
            )
        {
            return Some(image::DynamicImage::ImageRgba8(img));
        }
        std::thread::sleep(std::time::Duration::from_millis(20));
    }
    None
}

pub struct ClipboardHandler;

#[async_trait]
impl McpTool for ClipboardHandler {
    fn name(&self) -> &'static str {
        "clipboard"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<ClipboardTool>(
            "clipboard",
            "Consolidated OS clipboard management (read, write)",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
        let req: ClipboardTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        match req.action {
            ClipboardAction::Read => {
                let (mut out, b64_opt) =
                    tokio::task::spawn_blocking(move || -> crate::error::Result<(serde_json::Map<String, Value>, Option<String>)> {
                        let mut out = serde_json::Map::new();
                        let mut b64_opt = None;

                        if let Some(text) = get_native_clipboard_text() {
                            out.insert("text".into(), json!(text));
                        }

                        if let Some(dynamic_img) = get_native_clipboard_image() {
                            let mut img = dynamic_img;
                            let max_dim = 1024;
                            if img.width() > max_dim || img.height() > max_dim {
                                img = img.resize(max_dim, max_dim, FilterType::Lanczos3);
                            }
                            let rgb_img = img.into_rgb8();

                            let mut jpeg_bytes = std::io::Cursor::new(Vec::new());
                            if rgb_img.write_to(&mut jpeg_bytes, image::ImageFormat::Jpeg).is_ok() {
                                let bytes = jpeg_bytes.into_inner();
                                use base64::Engine;
                                let b64 = base64::engine::general_purpose::STANDARD.encode(&bytes);
                                b64_opt = Some(b64);

                                let cache_dir = dirs::home_dir()
                                    .unwrap_or_default()
                                    .join(".gemini/mcp_memory/clipboard");
                                let _ = std::fs::create_dir_all(&cache_dir);
                                let file_path = cache_dir.join("clipboard_latest.jpg");
                                if std::fs::write(&file_path, &bytes).is_ok() {
                                    let path_str = file_path.to_string_lossy().to_string();
                                    out.insert("image_path".into(), json!(path_str));
                                    #[cfg(target_os = "windows")]
                                    if let Some(rest) = path_str.strip_prefix("C:\\") {
                                        let wsl_path = format!("/mnt/c/{}", rest.replace('\\', "/"));
                                        out.insert("image_path_wsl".into(), json!(wsl_path));
                                    } else if let Some(rest) = path_str.strip_prefix("c:\\") {
                                        let wsl_path = format!("/mnt/c/{}", rest.replace('\\', "/"));
                                        out.insert("image_path_wsl".into(), json!(wsl_path));
                                    }
                                }
                            }
                        }
                        Ok((out, b64_opt))
                    })
                    .await
                    .map_err(|e| crate::error::AppError::Internal(format!("Task panic: {}", e)))??;

                if let Some(b64_str) = b64_opt {
                    if state.ollama.is_available().await
                        && let Ok(analysis) = state
                            .ollama
                            .generate_vision(
                                "Describe the content, text, or UI layout of this image concisely:",
                                &b64_str,
                            )
                            .await
                    {
                        out.insert("image_analysis".to_string(), json!(analysis.trim()));
                    }
                }

                state.record_activity("clipboard", "Read contents from OS clipboard", None);
                Ok::<String, crate::error::AppError>(serde_json::to_string_pretty(&Value::Object(out))?)
            }
            ClipboardAction::Write => {
                let text_opt = req.text;
                let image_path_opt = req.image_path;
                let res = tokio::task::spawn_blocking(move || {
                    let _guard = CLIPBOARD_LOCK.lock().unwrap_or_else(|e| e.into_inner());
                    let mut msgs = Vec::new();

                    if let Some(text) = &text_opt {
                        let mut written = false;
                        if let Ok(mut clipboard) = Clipboard::new() {
                            for _ in 0..3 {
                                if clipboard.set_text(text).is_ok() {
                                    written = true;
                                    break;
                                }
                                std::thread::sleep(std::time::Duration::from_millis(20));
                            }
                        }
                        if written {
                            msgs.push("Wrote text");
                        }
                    }

                    if let Some(image_path) = &image_path_opt {
                        match image::open(image_path) {
                            Ok(img) => {
                                let img = if img.width() > 2048 || img.height() > 2048 {
                                    img.resize(2048, 2048, FilterType::Triangle)
                                } else {
                                    img
                                };
                                let rgba = img.into_rgba8();
                                let (w, h) = rgba.dimensions();
                                let raw_bytes = rgba.into_raw();
                                let mut written = false;
                                if let Ok(mut clipboard) = arboard::Clipboard::new() {
                                    for _ in 0..3 {
                                        let img_data = ImageData {
                                            width: w as usize,
                                            height: h as usize,
                                            bytes: Cow::Borrowed(&raw_bytes),
                                        };
                                        if clipboard.set_image(img_data).is_ok() {
                                            written = true;
                                            break;
                                        }
                                        std::thread::sleep(std::time::Duration::from_millis(20));
                                    }
                                }
                                if written {
                                    msgs.push("Wrote image");
                                } else {
                                    return Err(crate::error::AppError::Internal(
                                        "Failed to write image to clipboard".to_string(),
                                    ));
                                }
                            }
                            Err(e) => {
                                return Err(crate::error::AppError::Internal(format!(
                                    "Failed to open image {}: {}",
                                    image_path, e
                                )));
                            }
                        }
                    }

                    if msgs.is_empty() {
                        Ok("No valid clipboard data provided or failed to write.".to_string())
                    } else {
                        Ok(format!(
                            "Successfully populated clipboard: {}",
                            msgs.join(", ")
                        ))
                    }
                })
                .await
                .map_err(|e| crate::error::AppError::Internal(format!("Task panic: {}", e)))??;

                state.record_activity("clipboard", &res, None);
                Ok(res)
            }
        }
    }
}


#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;
    use std::sync::Arc;
    use tempfile::tempdir;


    #[tokio::test]
    async fn test_write_clipboard() {
        let dir = tempdir().unwrap();
        let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
        let handler = ClipboardHandler;

        let args = json!({
            "action": "write",
            "text": "test_text"
        });

        let result = handler
            .execute(args, state)
            .await
            .map_err(|e| format!("Failed to write clipboard: {}", e))
            .unwrap();

        // Either successfully wrote, or failed to open clipboard (expected in CI)
        assert!(
            result.contains("Successfully populated")
                || result.contains("No valid clipboard data")
                || result.contains("Failed to write image")
        );
    }

    #[tokio::test]
    async fn test_read_clipboard() {
        let dir = tempdir().unwrap();
        let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
        let handler = ClipboardHandler;

        let result = handler
            .execute(json!({"action": "read"}), state)
            .await
            .map_err(|e| format!("Failed to read clipboard: {}", e))
            .unwrap();

        // Returns a JSON string, possibly {}
        let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
        assert!(parsed.is_object());
    }

    #[tokio::test]
    async fn test_read_clipboard_empty() {
        let dir = tempfile::tempdir().unwrap();
        let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
        let handler = ClipboardHandler;

        let result = handler
            .execute(serde_json::json!({"action": "read"}), state)
            .await
            .map_err(|e| format!("Failed to read clipboard: {}", e))
            .unwrap();

        let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
        assert!(parsed.is_object());
    }

    #[test]
    fn test_no_subprocess_clipboard_regression() {
        let vision_src = include_str!("vision.rs");
        let code_only = vision_src.split("mod tests").next().unwrap_or(vision_src);
        let forbidden_cmd = format!("Command::{}{}", "n", "ew");
        let forbidden_ps = format!("power{}", "shell");
        let forbidden_wl = format!("wl-{}", "paste");
        let forbidden_xc = format!("x{}", "clip");

        assert!(
            !code_only.contains(&forbidden_cmd),
            "Regression detected: vision.rs must not spawn subprocesses!"
        );
        assert!(
            !code_only.contains(&forbidden_ps),
            "Regression detected: vision.rs must not invoke powershell!"
        );
        assert!(
            !code_only.contains(&forbidden_wl),
            "Regression detected: vision.rs must not invoke wl-paste!"
        );
        assert!(
            !code_only.contains(&forbidden_xc),
            "Regression detected: vision.rs must not invoke xclip!"
        );
    }
}