riz/mcp-memory
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(server): refactor handlers, router, state management, and memory tools

Browse Source
This commit is contained in:
Riz Ashraf committed 2026-10-02 07:27:37 +01:00
1 parent 87ddb01063
commit a083719cf1
36 files changed
+1899 -597
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

server/src/handlers/vision.rs
+118 -73
View File
@@ -1,12 +1,12 @@
use crate::router::McpTool;
use crate::state::MemoryState;
use crate::tools::{ReadClipboardTool, WriteClipboardTool};
use async_trait::async_trait;
use image::{imageops::FilterType, ImageBuffer};
use serde_json::{json, Value};
use std::sync::Arc;
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;
pub struct WriteClipboardHandler;
@@ -24,28 +24,19 @@ impl McpTool for WriteClipboardHandler {
}
async fn execute(&self, args: Value, _state: Arc<MemoryState>) -> crate::error::Result<String> {
let tool_args: WriteClipboardTool =
serde_json::from_value(args).map_err(|e| crate::error::AppError::Internal(format!("Invalid args: {}", e)))?;
let tool_args: WriteClipboardTool = serde_json::from_value(args)
.map_err(|e| crate::error::AppError::Internal(format!("Invalid args: {}", e)))?;
tokio::task::spawn_blocking(move || {
let mut msgs = Vec::new();
if let Ok(mut clipboard) = Clipboard::new() {
if let Some(text) = &tool_args.text {
if clipboard.set_text(text).is_ok() {
msgs.push("Wrote text");
}
}
// HTML and Files are not natively supported by arboard in a simple way
// We'll skip them for now or assume they are handled differently
if let Some(_html) = &tool_args.html {
// Not supported via arboard
}
if let Some(_files) = &tool_args.files {
// Not supported via arboard
}
if let Ok(mut clipboard) = Clipboard::new()
&& let Some(text) = &tool_args.text
&& clipboard.set_text(text).is_ok()
{
msgs.push("Wrote text");
}
// Handle arboard for image
if let Some(image_path) = &tool_args.image_path {
match image::open(image_path) {
@@ -61,18 +52,28 @@ impl McpTool for WriteClipboardHandler {
if clipboard.set_image(img_data).is_ok() {
msgs.push("Wrote image");
} else {
return Err(crate::error::AppError::Internal("Failed to write image to clipboard".to_string()));
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))),
}
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(", ")))
Ok(format!(
"Successfully populated clipboard: {}",
msgs.join(", ")
))
}
})
.await
@@ -95,54 +96,87 @@ impl McpTool for ReadClipboardHandler {
)
}
async fn execute(&self, _args: Value, _state: Arc<MemoryState>) -> crate::error::Result<String> {
let result = tokio::task::spawn_blocking(move || -> crate::error::Result<serde_json::Value> {
let mut out = serde_json::Map::new();
async fn execute(&self, _args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let result =
tokio::task::spawn_blocking(move || -> crate::error::Result<serde_json::Value> {
let mut out = serde_json::Map::new();
if let Ok(mut clipboard) = arboard::Clipboard::new() {
if let Ok(text) = clipboard.get_text() {
if !text.trim().is_empty() {
out.insert("text".into(), json!(text));
}
if let Ok(mut clipboard) = arboard::Clipboard::new()
&& let Ok(text) = clipboard.get_text()
&& !text.trim().is_empty()
{
out.insert("text".into(), json!(text));
}
}
if let Ok(mut clipboard) = arboard::Clipboard::new()
&& let Ok(image_data) = clipboard.get_image()
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(),
) {
let mut dynamic_img = image::DynamicImage::ImageRgba8(img);
let max_dim = 1024;
if dynamic_img.width() > max_dim || dynamic_img.height() > max_dim {
dynamic_img = dynamic_img.resize(max_dim, max_dim, FilterType::Lanczos3);
}
let rgb_img = dynamic_img.into_rgb8();
)
{
let mut dynamic_img = image::DynamicImage::ImageRgba8(img);
let max_dim = 1024;
if dynamic_img.width() > max_dim || dynamic_img.height() > max_dim {
dynamic_img = dynamic_img.resize(max_dim, max_dim, FilterType::Lanczos3);
}
let rgb_img = dynamic_img.into_rgb8();
let cache_dir = dirs::home_dir()
.unwrap_or_default()
.join(".gemini/mcp_memory/clipboard");
let _ = std::fs::create_dir_all(&cache_dir);
let cache_dir = dirs::home_dir()
.unwrap_or_default()
.join(".gemini/mcp_memory/clipboard");
let _ = std::fs::create_dir_all(&cache_dir);
let timestamp = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_secs();
let file_path = cache_dir.join(format!("clipboard_{}.jpg", timestamp));
let timestamp = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_secs();
let file_path = cache_dir.join(format!("clipboard_{}.jpg", timestamp));
if rgb_img.save_with_format(&file_path, image::ImageFormat::Jpeg).is_ok() {
out.insert("image_path".into(), json!(file_path.to_string_lossy().to_string()));
if rgb_img
.save_with_format(&file_path, image::ImageFormat::Jpeg)
.is_ok()
{
let path_str = file_path.to_string_lossy().to_string();
out.insert("image_path".into(), json!(path_str));
// Read image bytes for base64 encoding if needed by vision
if let Ok(bytes) = std::fs::read(&file_path) {
use base64::Engine;
let b64 = base64::engine::general_purpose::STANDARD.encode(&bytes);
out.insert("image_base64".into(), json!(b64));
}
}
}
Ok(Value::Object(out))
})
.await
.map_err(|e| crate::error::AppError::Internal(format!("Task panic: {}", e)))??;
Ok(Value::Object(out))
})
.await
.map_err(|e| crate::error::AppError::Internal(format!("Task panic: {}", e)))??;
Ok::<String, crate::error::AppError>(serde_json::to_string_pretty(&result)?)
let mut final_obj = result;
if let Some(b64) = final_obj.get("image_base64").and_then(|v| v.as_str()) {
let b64_str = b64.to_string();
// Remove huge base64 string from final user output
if let Some(obj) = final_obj.as_object_mut() {
obj.remove("image_base64");
}
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
&& let Some(obj) = final_obj.as_object_mut()
{
obj.insert("image_analysis".to_string(), json!(analysis.trim()));
}
}
Ok::<String, crate::error::AppError>(serde_json::to_string_pretty(&final_obj)?)
}
}
use crate::tools::ToggleClipboardWatchModeTool;
@@ -163,13 +197,17 @@ impl McpTool for ToggleClipboardWatchModeHandler {
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let tool_args: ToggleClipboardWatchModeTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let tool_args: ToggleClipboardWatchModeTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
let mut watch_mode = state.clipboard_watch_mode.write().await;
*watch_mode = tool_args.enable;
if tool_args.enable {
Ok("Clipboard watch mode enabled. Changes will be ingested as StickyNotes.".to_string())
Ok(
"Clipboard watch mode enabled. Changes will be ingested as StickyNotes."
.to_string(),
)
} else {
Ok("Clipboard watch mode disabled.".to_string())
}
@@ -179,9 +217,9 @@ impl McpTool for ToggleClipboardWatchModeHandler {
#[cfg(test)]
mod tests {
use super::*;
use tempfile::tempdir;
use std::sync::Arc;
use serde_json::json;
use std::sync::Arc;
use tempfile::tempdir;
#[tokio::test]
async fn test_toggle_clipboard_watch_mode() {
@@ -193,7 +231,8 @@ mod tests {
"enable": true
});
let result = handler.execute(args, state.clone())
let result = handler
.execute(args, state.clone())
.await
.map_err(|e| format!("Failed to toggle clipboard: {}", e))
.unwrap();
@@ -211,13 +250,18 @@ mod tests {
"text": "test_text"
});
let result = handler.execute(args, state)
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"));
assert!(
result.contains("Successfully populated")
|| result.contains("No valid clipboard data")
|| result.contains("Failed to write image")
);
}
#[tokio::test]
@@ -226,11 +270,12 @@ mod tests {
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
let handler = ReadClipboardHandler;
let result = handler.execute(json!({}), state)
let result = handler
.execute(json!({}), 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());