Skip to main content

hanzo_mcp/
server.rs

1use crate::{Config, ToolRegistry};
2use anyhow::Result;
3use jsonrpc_core::{IoHandler, Params};
4use jsonrpc_http_server::ServerBuilder;
5use log::{debug, info, error};
6use serde_json::json;
7use std::path::PathBuf;
8use std::sync::Arc;
9use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
10use tokio::sync::RwLock;
11
12pub struct MCPServer {
13    config: Config,
14    port: u16,
15    roots: Vec<PathBuf>,
16    tools: Arc<RwLock<ToolRegistry>>,
17    handler: IoHandler,
18}
19
20impl MCPServer {
21    pub fn new(config: Config, port: u16) -> Result<Self> {
22        let tools = Arc::new(RwLock::new(ToolRegistry::with_defaults()));
23        let mut handler = IoHandler::new();
24        
25        // Clone for move into closures
26        let tools_clone = tools.clone();
27        
28        // Initialize method
29        handler.add_method("initialize", move |params: Params| {
30            let tools = tools_clone.clone();
31            Box::pin(async move {
32                debug!("Received initialize request: {:?}", params);
33                
34                let _tools = tools.read().await;
35                
36                Ok(json!({
37                    "protocolVersion": "2024-11-05",
38                    "serverInfo": {
39                        "name": "hanzo-mcp",
40                        "version": env!("CARGO_PKG_VERSION")
41                    },
42                    "capabilities": {
43                        "tools": {},
44                        "resources": {},
45                        "prompts": {}
46                    }
47                }))
48            })
49        });
50        
51        // List tools method
52        let tools_clone = tools.clone();
53        handler.add_method("tools/list", move |_params: Params| {
54            let tools = tools_clone.clone();
55            Box::pin(async move {
56                let tools = tools.read().await;
57                let tool_list = tools.get_definitions();
58                
59                Ok(json!({
60                    "tools": tool_list
61                }))
62            })
63        });
64        
65        // Call tool method
66        let tools_clone = tools.clone();
67        handler.add_method("tools/call", move |params: Params| {
68            let tools = tools_clone.clone();
69            Box::pin(async move {
70                let params = params.parse::<serde_json::Value>()
71                    .map_err(|e| jsonrpc_core::Error::invalid_params(e.to_string()))?;
72                
73                let tool_name = params["name"].as_str()
74                    .ok_or_else(|| jsonrpc_core::Error::invalid_params("Missing tool name"))?;
75                
76                let tool_params = params.get("arguments").cloned().unwrap_or(json!({}));
77                
78                let tools = tools.read().await;
79                match tools.execute(tool_name, tool_params).await {
80                    Ok(result) => {
81                        let text = if result.success {
82                            serde_json::to_string(&result.content).unwrap_or_default()
83                        } else {
84                            result.error.unwrap_or_else(|| "Unknown tool error".to_string())
85                        };
86                        Ok(json!({
87                            "content": [{
88                                "type": "text",
89                                "text": text
90                            }],
91                            "isError": !result.success
92                        }))
93                    },
94                    Err(e) => {
95                        error!("Tool execution failed: {}", e);
96                        Ok(json!({
97                            "content": [{
98                                "type": "text",
99                                "text": format!("Error: {}", e)
100                            }],
101                            "isError": true
102                        }))
103                    }
104                }
105            })
106        });
107        
108        // List resources method
109        handler.add_method("resources/list", |_params: Params| {
110            Box::pin(async move {
111                Ok(json!({
112                    "resources": []
113                }))
114            })
115        });
116        
117        // List prompts method
118        handler.add_method("prompts/list", |_params: Params| {
119            Box::pin(async move {
120                Ok(json!({
121                    "prompts": []
122                }))
123            })
124        });
125        
126        // Ping method for health checks
127        handler.add_method("ping", |_params: Params| {
128            Box::pin(async move {
129                Ok(json!("pong"))
130            })
131        });
132        
133        Ok(Self {
134            config,
135            port,
136            roots: Vec::new(),
137            tools,
138            handler,
139        })
140    }
141
142    /// Record an allowed filesystem root the server may operate under.
143    pub fn allow_root(&mut self, root: PathBuf) {
144        self.roots.push(root);
145    }
146
147    /// Serve JSON-RPC over STDIO: one request per line on stdin, one response
148    /// per line on stdout. Requests are dispatched through the same IoHandler
149    /// the HTTP transport uses. Notifications (no id) yield no output line.
150    pub async fn run_stdio(self) -> Result<()> {
151        for root in &self.roots {
152            info!("[MCP] allowed root: {}", root.display());
153        }
154        info!("[MCP] JSON-RPC over STDIO");
155
156        let mut lines = BufReader::new(tokio::io::stdin()).lines();
157        let mut out = tokio::io::stdout();
158
159        while let Some(line) = lines.next_line().await? {
160            if line.trim().is_empty() {
161                continue;
162            }
163            if let Some(response) = self.handler.handle_request(&line).await {
164                out.write_all(response.as_bytes()).await?;
165                out.write_all(b"\n").await?;
166                out.flush().await?;
167            }
168        }
169
170        Ok(())
171    }
172
173    pub async fn run(self) -> Result<()> {
174        let server = ServerBuilder::new(self.handler)
175            .start_http(&format!("127.0.0.1:{}", self.port).parse()?)
176            .map_err(|e| anyhow::anyhow!("Failed to start server: {}", e))?;
177        
178        info!("MCP Server running on http://127.0.0.1:{}", self.port);
179        
180        // Keep server running
181        server.wait();
182        
183        Ok(())
184    }
185    
186    pub async fn add_tool(&self, tool: Box<dyn crate::MCPTool>) {
187        let mut tools = self.tools.write().await;
188        tools.register(tool);
189    }
190}