codex_memory/mcp_server/

handlers.rs

//! Simple MCP request handlers
use crate::chunking::FileChunker;
use crate::error::Result;
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,
            _ => 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())
        }))
    }
}