codex_memory/mcp_server/

handlers.rs

//! Simple MCP request handlers
use crate::chunking::FileChunker;
use crate::error::Result;
use crate::models::{SearchParams, SearchStrategy};
use crate::storage::Storage;
use serde_json::{json, Value};
use std::path::Path;
use std::sync::Arc;
use uuid::Uuid;

/// Minimal MCP request handlers
pub struct MCPHandlers {
    storage: Arc<Storage>,
}

impl MCPHandlers {
    /// Create new handlers with storage backend
    pub fn new(storage: Arc<Storage>) -> Self {
        Self { storage }
    }

    /// Handle tool calls
    pub async fn handle_tool_call(&self, tool_name: &str, params: Value) -> Result<Value> {
        match tool_name {
            "store_memory" => self.handle_store_memory(params).await,
            "get_memory" => self.handle_get_memory(params).await,
            "delete_memory" => self.handle_delete_memory(params).await,
            "get_statistics" => self.handle_get_statistics().await,
            "store_file" => self.handle_store_file(params).await,
            "search_memory" => self.handle_search_memory(params).await,
            _ => Err(crate::error::Error::Other(format!(
                "Unknown tool: {}",
                tool_name
            ))),
        }
    }

    async fn handle_store_memory(&self, params: Value) -> Result<Value> {
        let content = params["content"]
            .as_str()
            .ok_or_else(|| crate::error::Error::Other("Missing content parameter".to_string()))?;

        // Context is required
        let context = params["context"]
            .as_str()
            .ok_or_else(|| {
                crate::error::Error::Other("Missing required context parameter".to_string())
            })?
            .to_string();

        // Summary is required
        let summary = params["summary"]
            .as_str()
            .ok_or_else(|| {
                crate::error::Error::Other("Missing required summary parameter".to_string())
            })?
            .to_string();

        // Tags are required
        let tags = params["tags"]
            .as_array()
            .ok_or_else(|| {
                crate::error::Error::Other("Missing required tags parameter".to_string())
            })?
            .iter()
            .filter_map(|v| v.as_str().map(String::from))
            .collect::<Vec<_>>();

        let id = self
            .storage
            .store(content, context, summary, Some(tags))
            .await?;

        Ok(json!({
            "id": id.to_string(),
            "message": "Memory stored successfully"
        }))
    }

    async fn handle_get_memory(&self, params: Value) -> Result<Value> {
        let id_str = params["id"]
            .as_str()
            .ok_or_else(|| crate::error::Error::Other("Missing id parameter".to_string()))?;

        let id = Uuid::parse_str(id_str)
            .map_err(|e| crate::error::Error::Other(format!("Invalid UUID: {}", e)))?;

        match self.storage.get(id).await? {
            Some(memory) => Ok(serde_json::to_value(memory)?),
            None => Err(crate::error::Error::Other(format!(
                "Memory not found: {}",
                id
            ))),
        }
    }

    async fn handle_delete_memory(&self, params: Value) -> Result<Value> {
        let id_str = params["id"]
            .as_str()
            .ok_or_else(|| crate::error::Error::Other("Missing id parameter".to_string()))?;

        let id = Uuid::parse_str(id_str)
            .map_err(|e| crate::error::Error::Other(format!("Invalid UUID: {}", e)))?;

        let deleted = self.storage.delete(id).await?;

        Ok(json!({
            "deleted": deleted,
            "message": if deleted { "Memory deleted successfully" } else { "Memory not found" }
        }))
    }

    async fn handle_get_statistics(&self) -> Result<Value> {
        let stats = self.storage.stats().await?;
        Ok(serde_json::to_value(stats)?)
    }

    async fn handle_store_file(&self, params: Value) -> Result<Value> {
        let file_path = params["file_path"]
            .as_str()
            .ok_or_else(|| crate::error::Error::Other("Missing file_path parameter".to_string()))?;

        let chunk_size = params
            .get("chunk_size")
            .and_then(|v| v.as_u64())
            .unwrap_or(8000) as usize;

        let overlap = params
            .get("overlap")
            .and_then(|v| v.as_u64())
            .unwrap_or(200) as usize;

        // Parse chunking strategy
        // Note: Semantic chunking strategies have been moved to codex-dreams
        // This now uses simple byte-based chunking with overlap

        let tags = params.get("tags").and_then(|v| v.as_array()).map(|arr| {
            arr.iter()
                .filter_map(|v| v.as_str().map(String::from))
                .collect::<Vec<_>>()
        });

        // Read the file
        let content = tokio::fs::read_to_string(file_path)
            .await
            .map_err(|e| crate::error::Error::Other(format!("Failed to read file: {}", e)))?;

        // Extract filename for context
        let filename = Path::new(file_path)
            .file_name()
            .and_then(|n| n.to_str())
            .unwrap_or("unknown");

        // Use semantic chunking to preserve meaning boundaries
        let content_len = content.len();
        let mut stored_ids = Vec::new();

        // Create simple file chunker
        let chunker = FileChunker::new(chunk_size, overlap);
        let chunks = chunker.chunk_content(&content)?;

        if chunks.len() == 1 {
            // File fits in a single chunk
            let context = format!("Content from file: {}", filename);
            let summary = format!(
                "Complete content of {} ({} characters)",
                filename, content_len
            );

            let id = self
                .storage
                .store(&content, context, summary, tags.clone())
                .await?;

            stored_ids.push(id.to_string());
        } else {
            // Multiple semantic chunks needed
            let parent_id = Uuid::new_v4();
            let total_chunks = chunks.len();

            for (index, chunk) in chunks.into_iter().enumerate() {
                let chunk_num = index + 1;

                let context = format!(
                    "Chunk {} of {} from file: {}",
                    chunk_num, total_chunks, filename
                );

                let summary = format!(
                    "Part {} of {} from {} (bytes {}-{} of {})",
                    chunk_num,
                    total_chunks,
                    filename,
                    chunk.start_byte,
                    chunk.end_byte,
                    content_len
                );

                let mut chunk_tags = tags.clone().unwrap_or_default();
                chunk_tags.push(format!("chunk_{}", chunk_num));
                chunk_tags.push(format!("file_{}", filename));
                chunk_tags.push("byte_chunked".to_string());

                let id = self
                    .storage
                    .store_chunk(
                        &chunk.content,
                        context,
                        summary,
                        Some(chunk_tags),
                        chunk_num as i32,
                        total_chunks as i32,
                        parent_id,
                    )
                    .await?;

                stored_ids.push(id.to_string());
            }
        }

        Ok(json!({
            "file_path": file_path,
            "file_size": content_len,
            "chunks_created": stored_ids.len(),
            "chunk_ids": stored_ids,
            "message": format!("Successfully ingested {} as {} chunk(s)", filename, stored_ids.len())
        }))
    }

    async fn handle_search_memory(&self, params: Value) -> Result<Value> {
        let query = params["query"]
            .as_str()
            .ok_or_else(|| crate::error::Error::Other("Missing query parameter".to_string()))?
            .to_string();

        // Parse optional parameters with defaults
        let tag_filter = params
            .get("tag_filter")
            .and_then(|v| v.as_array())
            .map(|arr| {
                arr.iter()
                    .filter_map(|v| v.as_str().map(String::from))
                    .collect::<Vec<_>>()
            });

        let use_tag_embedding = params
            .get("use_tag_embedding")
            .and_then(|v| v.as_bool())
            .unwrap_or(true);

        let use_content_embedding = params
            .get("use_content_embedding")
            .and_then(|v| v.as_bool())
            .unwrap_or(true);

        let similarity_threshold = params
            .get("similarity_threshold")
            .and_then(|v| v.as_f64())
            .unwrap_or(0.7)
            .clamp(0.0, 1.0);

        let max_results = params
            .get("max_results")
            .and_then(|v| v.as_u64())
            .unwrap_or(10)
            .clamp(1, 100) as usize;

        let search_strategy = params
            .get("search_strategy")
            .and_then(|v| v.as_str())
            .map(|s| match s {
                "tags_first" => SearchStrategy::TagsFirst,
                "content_first" => SearchStrategy::ContentFirst,
                _ => SearchStrategy::Hybrid,
            })
            .unwrap_or(SearchStrategy::Hybrid);

        let boost_recent = params
            .get("boost_recent")
            .and_then(|v| v.as_bool())
            .unwrap_or(false);

        let tag_weight = params
            .get("tag_weight")
            .and_then(|v| v.as_f64())
            .unwrap_or(0.4)
            .clamp(0.0, 1.0);

        let content_weight = params
            .get("content_weight")
            .and_then(|v| v.as_f64())
            .unwrap_or(0.6)
            .clamp(0.0, 1.0);

        // Create search parameters
        let search_params = SearchParams {
            query,
            tag_filter,
            use_tag_embedding,
            use_content_embedding,
            similarity_threshold,
            max_results,
            search_strategy,
            boost_recent,
            tag_weight,
            content_weight,
        };

        // Perform the progressive search
        let search_start = std::time::Instant::now();
        let search_result_with_metadata = self
            .storage
            .search_memories_progressive_with_metadata(search_params.clone())
            .await?;
        let search_duration = search_start.elapsed();

        // Format results for JSON response
        let formatted_results: Vec<Value> = search_result_with_metadata
            .results
            .iter()
            .map(|result| {
                json!({
                    "id": result.memory.id,
                    "content": result.memory.content,
                    "context": result.memory.context,
                    "summary": result.memory.summary,
                    "tags": result.memory.tags,
                    "chunk_index": result.memory.chunk_index,
                    "total_chunks": result.memory.total_chunks,
                    "parent_id": result.memory.parent_id,
                    "created_at": result.memory.created_at,
                    "updated_at": result.memory.updated_at,
                    "tag_similarity": result.tag_similarity,
                    "content_similarity": result.content_similarity,
                    "combined_score": result.combined_score,
                    "semantic_cluster": result.semantic_cluster
                })
            })
            .collect();

        Ok(json!({
            "results": formatted_results,
            "search_metadata": {
                "query": search_params.query,
                "total_results": search_result_with_metadata.results.len(),
                "search_strategy": match search_params.search_strategy {
                    SearchStrategy::TagsFirst => "tags_first",
                    SearchStrategy::ContentFirst => "content_first",
                    SearchStrategy::Hybrid => "hybrid",
                },
                "similarity_threshold": search_params.similarity_threshold,
                "max_results": search_params.max_results,
                "tag_filter": search_params.tag_filter,
                "use_tag_embedding": search_params.use_tag_embedding,
                "use_content_embedding": search_params.use_content_embedding,
                "boost_recent": search_params.boost_recent,
                "tag_weight": search_params.tag_weight,
                "content_weight": search_params.content_weight,
                "search_time_ms": search_duration.as_millis(),
                "progressive_search": {
                    "stage_used": search_result_with_metadata.metadata.stage_used,
                    "stage_description": search_result_with_metadata.metadata.stage_description,
                    "threshold_used": search_result_with_metadata.metadata.threshold_used
                }
            }
        }))
    }
}