codex_memory/mcp_server/

mod.rs

//! Minimal MCP server implementation
pub mod handlers;
pub mod tools;
pub mod transport;

// Re-export for tests
pub use handlers::MCPHandlers;

use crate::config::Config;
use crate::error::Result;
use crate::storage::Storage;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tracing::{error, info, warn};

/// Simple MCP server
pub struct MCPServer {
    _config: Config,
    handlers: Arc<MCPHandlers>,
    start_time: Instant,
    last_request: Arc<std::sync::Mutex<Instant>>,
}

impl MCPServer {
    /// Create a new MCP server
    pub fn new(config: Config, storage: Arc<Storage>) -> Self {
        let handlers = Arc::new(MCPHandlers::new(storage));
        let now = Instant::now();
        Self {
            _config: config,
            handlers,
            start_time: now,
            last_request: Arc::new(std::sync::Mutex::new(now)),
        }
    }

    /// Check if server should self-terminate due to inactivity
    fn should_terminate(&self) -> bool {
        let last_request = *self.last_request.lock().unwrap();
        let inactive_duration = last_request.elapsed();
        
        // Terminate if inactive for more than 10 minutes
        if inactive_duration > Duration::from_secs(600) {
            warn!("Server inactive for {:?}, initiating shutdown", inactive_duration);
            return true;
        }
        
        false
    }

    /// Update last request time
    fn update_last_request(&self) {
        *self.last_request.lock().unwrap() = Instant::now();
    }

    /// Log health status periodically
    async fn health_monitor(&self) {
        let mut interval = tokio::time::interval(Duration::from_secs(60)); // Every minute
        
        loop {
            interval.tick().await;
            
            if self.should_terminate() {
                error!("Health monitor detected inactivity timeout, terminating process");
                std::process::exit(1);
            }
            
            let uptime = self.start_time.elapsed();
            let last_request_ago = self.last_request.lock().unwrap().elapsed();
            
            info!("Health check: uptime={:?}, last_request={:?} ago", uptime, last_request_ago);
        }
    }

    /// Run in stdio mode for Claude Desktop
    pub async fn run_stdio(&self) -> Result<()> {
        info!("MCP server running in stdio mode");

        // Spawn health monitor task
        let health_monitor = {
            let server_clone = Self {
                _config: self._config.clone(),
                handlers: Arc::clone(&self.handlers),
                start_time: self.start_time,
                last_request: Arc::clone(&self.last_request),
            };
            tokio::spawn(async move {
                server_clone.health_monitor().await;
            })
        };

        let stdin = tokio::io::stdin();
        let stdout = tokio::io::stdout();
        let mut stdin = stdin;
        let mut stdout = stdout;

        let mut buffer = String::new();
        let mut temp_buf = [0u8; 8192]; // 8KB buffer for reading chunks

        loop {
            // Read chunk from stdin
            match stdin.read(&mut temp_buf).await {
                Ok(0) => {
                    info!("Received EOF, shutting down MCP server");
                    break; // EOF
                }
                Ok(n) => {
                    // Convert bytes to string and append to buffer
                    let chunk = String::from_utf8_lossy(&temp_buf[..n]);
                    info!(
                        "Read {} bytes: {:?}",
                        n,
                        &chunk[..std::cmp::min(100, chunk.len())]
                    );
                    buffer.push_str(&chunk);

                    // Process complete JSON objects from buffer
                    while let Some(json_end) = self.find_complete_json(&buffer) {
                        let json_str = buffer[..json_end].trim().to_string();
                        buffer.drain(..json_end);

                        if !json_str.is_empty() {
                            info!("Processing JSON request ({} chars)", json_str.len());
                            self.update_last_request();
                            let response = self.handle_request(&json_str).await;
                            stdout.write_all(response.as_bytes()).await?;
                            stdout.write_all(b"\n").await?;
                            stdout.flush().await?;
                        }
                    }
                }
                Err(e) => {
                    error!("Error reading input: {}", e);
                    break;
                }
            }
        }

        info!("MCP server shutting down gracefully");
        
        // Cancel health monitor task
        health_monitor.abort();
        
        // The pool will be dropped automatically when Storage is dropped
        Ok(())
    }

    /// Find the end of a complete JSON object in the buffer
    fn find_complete_json(&self, buffer: &str) -> Option<usize> {
        let mut brace_count = 0;
        let mut in_string = false;
        let mut escape_next = false;
        let mut start_found = false;

        for (i, ch) in buffer.char_indices() {
            if escape_next {
                escape_next = false;
                continue;
            }

            match ch {
                '\\' if in_string => escape_next = true,
                '"' => in_string = !in_string,
                '{' if !in_string => {
                    brace_count += 1;
                    start_found = true;
                }
                '}' if !in_string => {
                    brace_count -= 1;
                    if start_found && brace_count == 0 {
                        return Some(i + 1);
                    }
                }
                _ => {}
            }
        }

        None
    }

    async fn handle_request(&self, request: &str) -> String {
        let request: serde_json::Value = match serde_json::from_str(request) {
            Ok(v) => v,
            Err(e) => {
                return format!(
                    r#"{{"jsonrpc":"2.0","id":null,"error":{{"code":-32700,"message":"Parse error: {}"}}}}"#,
                    e
                );
            }
        };

        let method = request["method"].as_str().unwrap_or("");
        let params = request.get("params").cloned().unwrap_or_default();
        let id = request.get("id").cloned();

        let result = match method {
            "initialize" => Ok(serde_json::json!({
                "protocolVersion": "2024-11-05",
                "capabilities": {
                    "tools": {}
                },
                "serverInfo": {
                    "name": "codex-memory",
                    "version": env!("CARGO_PKG_VERSION")
                }
            })),
            "tools/list" => Ok(serde_json::json!({
                "tools": tools::MCPTools::get_tools_list()
            })),
            "tools/call" => {
                let tool_name = params["name"].as_str().unwrap_or("");
                let tool_params = params.get("arguments").cloned().unwrap_or_default();
                self.handlers.handle_tool_call(tool_name, tool_params).await
            }
            _ => Err(crate::error::Error::Other(format!(
                "Unknown method: {}",
                method
            ))),
        };

        match result {
            Ok(value) => {
                if let Some(id) = id {
                    format!(r#"{{"jsonrpc":"2.0","id":{},"result":{}}}"#, id, value)
                } else {
                    format!(r#"{{"jsonrpc":"2.0","result":{}}}"#, value)
                }
            }
            Err(e) => {
                if let Some(id) = id {
                    format!(
                        r#"{{"jsonrpc":"2.0","id":{},"error":{{"code":-32000,"message":"{}"}}}}"#,
                        id, e
                    )
                } else {
                    format!(
                        r#"{{"jsonrpc":"2.0","id":null,"error":{{"code":-32000,"message":"{}"}}}}"#,
                        e
                    )
                }
            }
        }
    }
}