hanzo-mcp 1.1.23

use crate::{Config, ToolRegistry};
use anyhow::Result;
use jsonrpc_core::{IoHandler, Params};
use jsonrpc_http_server::ServerBuilder;
use log::{debug, info, error};
use serde_json::json;
use std::path::PathBuf;
use std::sync::Arc;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::sync::RwLock;

pub struct MCPServer {
    config: Config,
    port: u16,
    roots: Vec<PathBuf>,
    tools: Arc<RwLock<ToolRegistry>>,
    handler: IoHandler,
}

impl MCPServer {
    pub fn new(config: Config, port: u16) -> Result<Self> {
        let tools = Arc::new(RwLock::new(ToolRegistry::with_defaults()));
        let mut handler = IoHandler::new();
        
        // Clone for move into closures
        let tools_clone = tools.clone();
        
        // Initialize method
        handler.add_method("initialize", move |params: Params| {
            let tools = tools_clone.clone();
            Box::pin(async move {
                debug!("Received initialize request: {:?}", params);
                
                let _tools = tools.read().await;
                
                Ok(json!({
                    "protocolVersion": "2024-11-05",
                    "serverInfo": {
                        "name": "hanzo-mcp",
                        "version": env!("CARGO_PKG_VERSION")
                    },
                    "capabilities": {
                        "tools": {},
                        "resources": {},
                        "prompts": {}
                    }
                }))
            })
        });
        
        // List tools method
        let tools_clone = tools.clone();
        handler.add_method("tools/list", move |_params: Params| {
            let tools = tools_clone.clone();
            Box::pin(async move {
                let tools = tools.read().await;
                let tool_list = tools.get_definitions();
                
                Ok(json!({
                    "tools": tool_list
                }))
            })
        });
        
        // Call tool method
        let tools_clone = tools.clone();
        handler.add_method("tools/call", move |params: Params| {
            let tools = tools_clone.clone();
            Box::pin(async move {
                let params = params.parse::<serde_json::Value>()
                    .map_err(|e| jsonrpc_core::Error::invalid_params(e.to_string()))?;
                
                let tool_name = params["name"].as_str()
                    .ok_or_else(|| jsonrpc_core::Error::invalid_params("Missing tool name"))?;
                
                let tool_params = params.get("arguments").cloned().unwrap_or(json!({}));
                
                let tools = tools.read().await;
                match tools.execute(tool_name, tool_params).await {
                    Ok(result) => {
                        let text = if result.success {
                            serde_json::to_string(&result.content).unwrap_or_default()
                        } else {
                            result.error.unwrap_or_else(|| "Unknown tool error".to_string())
                        };
                        Ok(json!({
                            "content": [{
                                "type": "text",
                                "text": text
                            }],
                            "isError": !result.success
                        }))
                    },
                    Err(e) => {
                        error!("Tool execution failed: {}", e);
                        Ok(json!({
                            "content": [{
                                "type": "text",
                                "text": format!("Error: {}", e)
                            }],
                            "isError": true
                        }))
                    }
                }
            })
        });
        
        // List resources method
        handler.add_method("resources/list", |_params: Params| {
            Box::pin(async move {
                Ok(json!({
                    "resources": []
                }))
            })
        });
        
        // List prompts method
        handler.add_method("prompts/list", |_params: Params| {
            Box::pin(async move {
                Ok(json!({
                    "prompts": []
                }))
            })
        });
        
        // Ping method for health checks
        handler.add_method("ping", |_params: Params| {
            Box::pin(async move {
                Ok(json!("pong"))
            })
        });
        
        Ok(Self {
            config,
            port,
            roots: Vec::new(),
            tools,
            handler,
        })
    }

    /// Record an allowed filesystem root the server may operate under.
    pub fn allow_root(&mut self, root: PathBuf) {
        self.roots.push(root);
    }

    /// Serve JSON-RPC over STDIO: one request per line on stdin, one response
    /// per line on stdout. Requests are dispatched through the same IoHandler
    /// the HTTP transport uses. Notifications (no id) yield no output line.
    pub async fn run_stdio(self) -> Result<()> {
        for root in &self.roots {
            info!("[MCP] allowed root: {}", root.display());
        }
        info!("[MCP] JSON-RPC over STDIO");

        let mut lines = BufReader::new(tokio::io::stdin()).lines();
        let mut out = tokio::io::stdout();

        while let Some(line) = lines.next_line().await? {
            if line.trim().is_empty() {
                continue;
            }
            if let Some(response) = self.handler.handle_request(&line).await {
                out.write_all(response.as_bytes()).await?;
                out.write_all(b"\n").await?;
                out.flush().await?;
            }
        }

        Ok(())
    }

    pub async fn run(self) -> Result<()> {
        let server = ServerBuilder::new(self.handler)
            .start_http(&format!("127.0.0.1:{}", self.port).parse()?)
            .map_err(|e| anyhow::anyhow!("Failed to start server: {}", e))?;
        
        info!("MCP Server running on http://127.0.0.1:{}", self.port);
        
        // Keep server running
        server.wait();
        
        Ok(())
    }
    
    pub async fn add_tool(&self, tool: Box<dyn crate::MCPTool>) {
        let mut tools = self.tools.write().await;
        tools.register(tool);
    }
}