| 1 | use std::collections::HashMap; |
| 2 | use std::collections::hash_map::DefaultHasher; |
| 3 | use std::hash::{Hash, Hasher}; |
| 4 | |
| 5 | use anyhow::{Context, Result, bail}; |
| 6 | use serde::de::DeserializeOwned; |
| 7 | use serde::{Deserialize, Serialize}; |
| 8 | use serde_json::{Value, json}; |
| 9 | |
| 10 | #[derive(Debug, Clone, Serialize, Deserialize)] |
| 11 | pub struct McpServerConfig { |
| 12 | pub name: String, |
| 13 | pub command: String, |
| 14 | #[serde(default)] |
| 15 | pub args: Vec<String>, |
| 16 | #[serde(default)] |
| 17 | pub env: HashMap<String, String>, |
| 18 | #[serde(default = "default_true")] |
| 19 | pub enabled: bool, |
| 20 | } |
| 21 | |
| 22 | #[derive(Debug, Clone, Serialize, Deserialize, Default)] |
| 23 | pub struct ToolFilter { |
| 24 | #[serde(default)] |
| 25 | pub allow: Vec<String>, |
| 26 | #[serde(default)] |
| 27 | pub deny: Vec<String>, |
| 28 | } |
| 29 | |
| 30 | #[derive(Debug, Clone, Serialize, Deserialize)] |
| 31 | pub struct McpServerDefinition { |
| 32 | pub config: McpServerConfig, |
| 33 | #[serde(default)] |
| 34 | pub filter: ToolFilter, |
| 35 | } |
| 36 | |
| 37 | #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] |
| 38 | #[serde(rename_all = "snake_case")] |
| 39 | pub enum McpStartupStatus { |
| 40 | Starting, |
| 41 | Ready, |
| 42 | Failed { error: String }, |
| 43 | Cancelled, |
| 44 | } |
| 45 | |
| 46 | #[derive(Debug, Clone, Serialize, Deserialize)] |
| 47 | pub struct McpStartupUpdateEvent { |
| 48 | pub server_name: String, |
| 49 | pub status: McpStartupStatus, |
| 50 | } |
| 51 | |
| 52 | #[derive(Debug, Clone, Serialize, Deserialize)] |
| 53 | pub struct McpStartupFailure { |
| 54 | pub server_name: String, |
| 55 | pub error: String, |
| 56 | } |
| 57 | |
| 58 | #[derive(Debug, Clone, Serialize, Deserialize)] |
| 59 | pub struct McpStartupCompleteEvent { |
| 60 | pub ready: Vec<String>, |
| 61 | pub failed: Vec<McpStartupFailure>, |
| 62 | pub cancelled: Vec<String>, |
| 63 | } |
| 64 | |
| 65 | #[derive(Debug, Clone, Serialize, Deserialize)] |
| 66 | pub struct McpToolDescriptor { |
| 67 | pub server_name: String, |
| 68 | pub tool_name: String, |
| 69 | pub qualified_name: String, |
| 70 | pub description: Option<String>, |
| 71 | } |
| 72 | |
| 73 | #[derive(Debug, Clone, Serialize, Deserialize)] |
| 74 | pub struct McpResourceDescriptor { |
| 75 | pub server_name: String, |
| 76 | pub uri: String, |
| 77 | pub description: Option<String>, |
| 78 | } |
| 79 | |
| 80 | pub trait McpManagedClient: Send + Sync { |
| 81 | fn list_tools(&self) -> Result<Vec<McpToolDescriptor>>; |
| 82 | fn call_tool(&self, tool_name: &str, arguments: Value) -> Result<Value>; |
| 83 | fn list_resources(&self) -> Result<Vec<McpResourceDescriptor>>; |
| 84 | fn read_resource(&self, uri: &str) -> Result<Value>; |
| 85 | } |
| 86 | |
| 87 | #[derive(Debug, Default)] |
| 88 | pub struct InMemoryMcpClient { |
| 89 | tools: HashMap<String, Value>, |
| 90 | resources: HashMap<String, Value>, |
| 91 | } |
| 92 | |
| 93 | impl InMemoryMcpClient { |
| 94 | pub fn with_tool(mut self, name: &str, sample_result: Value) -> Self { |
| 95 | self.tools.insert(name.to_string(), sample_result); |
| 96 | self |
| 97 | } |
| 98 | |
| 99 | pub fn with_resource(mut self, uri: &str, data: Value) -> Self { |
| 100 | self.resources.insert(uri.to_string(), data); |
| 101 | self |
| 102 | } |
| 103 | } |
| 104 | |
| 105 | impl McpManagedClient for InMemoryMcpClient { |
| 106 | fn list_tools(&self) -> Result<Vec<McpToolDescriptor>> { |
| 107 | Ok(self |
| 108 | .tools |
| 109 | .keys() |
| 110 | .map(|name| McpToolDescriptor { |
| 111 | server_name: "in-memory".to_string(), |
| 112 | tool_name: name.clone(), |
| 113 | qualified_name: name.clone(), |
| 114 | description: None, |
| 115 | }) |
| 116 | .collect()) |
| 117 | } |
| 118 | |
| 119 | fn call_tool(&self, tool_name: &str, _arguments: Value) -> Result<Value> { |
| 120 | self.tools |
| 121 | .get(tool_name) |
| 122 | .cloned() |
| 123 | .with_context(|| format!("tool '{tool_name}' not found")) |
| 124 | } |
| 125 | |
| 126 | fn list_resources(&self) -> Result<Vec<McpResourceDescriptor>> { |
| 127 | Ok(self |
| 128 | .resources |
| 129 | .keys() |
| 130 | .map(|uri| McpResourceDescriptor { |
| 131 | server_name: "in-memory".to_string(), |
| 132 | uri: uri.clone(), |
| 133 | description: None, |
| 134 | }) |
| 135 | .collect()) |
| 136 | } |
| 137 | |
| 138 | fn read_resource(&self, uri: &str) -> Result<Value> { |
| 139 | self.resources |
| 140 | .get(uri) |
| 141 | .cloned() |
| 142 | .with_context(|| format!("resource '{uri}' not found")) |
| 143 | } |
| 144 | } |
| 145 | |
| 146 | #[derive(Default)] |
| 147 | pub struct McpManager { |
| 148 | configs: HashMap<String, (McpServerConfig, ToolFilter)>, |
| 149 | clients: HashMap<String, Box<dyn McpManagedClient>>, |
| 150 | } |
| 151 | |
| 152 | impl McpManager { |
| 153 | pub fn register_server( |
| 154 | &mut self, |
| 155 | config: McpServerConfig, |
| 156 | filter: ToolFilter, |
| 157 | client: Box<dyn McpManagedClient>, |
| 158 | ) { |
| 159 | self.clients.insert(config.name.clone(), client); |
| 160 | self.configs.insert(config.name.clone(), (config, filter)); |
| 161 | } |
| 162 | |
| 163 | pub fn start_all<F>(&self, mut emit: F) -> McpStartupCompleteEvent |
| 164 | where |
| 165 | F: FnMut(McpStartupUpdateEvent), |
| 166 | { |
| 167 | let mut ready = Vec::new(); |
| 168 | let mut failed = Vec::new(); |
| 169 | let mut cancelled = Vec::new(); |
| 170 | for (server_name, (cfg, _)) in &self.configs { |
| 171 | if !cfg.enabled { |
| 172 | emit(McpStartupUpdateEvent { |
| 173 | server_name: server_name.clone(), |
| 174 | status: McpStartupStatus::Cancelled, |
| 175 | }); |
| 176 | cancelled.push(server_name.clone()); |
| 177 | continue; |
| 178 | } |
| 179 | emit(McpStartupUpdateEvent { |
| 180 | server_name: server_name.clone(), |
| 181 | status: McpStartupStatus::Starting, |
| 182 | }); |
| 183 | if self.clients.contains_key(server_name) { |
| 184 | emit(McpStartupUpdateEvent { |
| 185 | server_name: server_name.clone(), |
| 186 | status: McpStartupStatus::Ready, |
| 187 | }); |
| 188 | ready.push(server_name.clone()); |
| 189 | } else { |
| 190 | let error = "client not registered".to_string(); |
| 191 | emit(McpStartupUpdateEvent { |
| 192 | server_name: server_name.clone(), |
| 193 | status: McpStartupStatus::Failed { |
| 194 | error: error.clone(), |
| 195 | }, |
| 196 | }); |
| 197 | failed.push(McpStartupFailure { |
| 198 | server_name: server_name.clone(), |
| 199 | error, |
| 200 | }); |
| 201 | } |
| 202 | } |
| 203 | McpStartupCompleteEvent { |
| 204 | ready, |
| 205 | failed, |
| 206 | cancelled, |
| 207 | } |
| 208 | } |
| 209 | |
| 210 | pub fn stop_server(&mut self, server_name: &str) -> Result<()> { |
| 211 | self.clients |
| 212 | .remove(server_name) |
| 213 | .with_context(|| format!("server '{server_name}' is not running"))?; |
| 214 | Ok(()) |
| 215 | } |
| 216 | |
| 217 | pub fn unregister_server(&mut self, server_name: &str) -> Result<()> { |
| 218 | let had_config = self.configs.remove(server_name).is_some(); |
| 219 | self.clients.remove(server_name); |
| 220 | if !had_config { |
| 221 | bail!("server '{server_name}' is not registered"); |
| 222 | } |
| 223 | Ok(()) |
| 224 | } |
| 225 | |
| 226 | pub fn list_tools(&self) -> Result<Vec<McpToolDescriptor>> { |
| 227 | let mut out = Vec::new(); |
| 228 | for (server_name, (_, filter)) in &self.configs { |
| 229 | let Some(client) = self.clients.get(server_name) else { |
| 230 | continue; |
| 231 | }; |
| 232 | let tools = client.list_tools()?; |
| 233 | for tool in tools { |
| 234 | if !allowed_by_filter(&tool.tool_name, filter) { |
| 235 | continue; |
| 236 | } |
| 237 | let qualified_name = qualify_tool_name(server_name, &tool.tool_name); |
| 238 | out.push(McpToolDescriptor { |
| 239 | server_name: server_name.clone(), |
| 240 | tool_name: tool.tool_name, |
| 241 | qualified_name, |
| 242 | description: tool.description, |
| 243 | }); |
| 244 | } |
| 245 | } |
| 246 | Ok(out) |
| 247 | } |
| 248 | |
| 249 | pub fn call_tool(&self, server_name: &str, tool_name: &str, arguments: Value) -> Result<Value> { |
| 250 | let client = self |
| 251 | .clients |
| 252 | .get(server_name) |
| 253 | .with_context(|| format!("MCP server '{server_name}' not available"))?; |
| 254 | client.call_tool(tool_name, arguments) |
| 255 | } |
| 256 | |
| 257 | pub fn call_qualified_tool( |
| 258 | &self, |
| 259 | qualified_tool_name: &str, |
| 260 | arguments: Value, |
| 261 | ) -> Result<Value> { |
| 262 | let (server_name, tool_name) = parse_qualified_tool_name(qualified_tool_name) |
| 263 | .with_context(|| format!("invalid qualified MCP tool name: {qualified_tool_name}"))?; |
| 264 | self.call_tool(&server_name, &tool_name, arguments) |
| 265 | } |
| 266 | |
| 267 | pub fn list_resources(&self) -> Result<Vec<McpResourceDescriptor>> { |
| 268 | let mut out = Vec::new(); |
| 269 | for server_name in self.configs.keys() { |
| 270 | let Some(client) = self.clients.get(server_name) else { |
| 271 | continue; |
| 272 | }; |
| 273 | for mut resource in client.list_resources()? { |
| 274 | resource.server_name = server_name.clone(); |
| 275 | out.push(resource); |
| 276 | } |
| 277 | } |
| 278 | Ok(out) |
| 279 | } |
| 280 | |
| 281 | pub fn read_resource(&self, server_name: &str, uri: &str) -> Result<Value> { |
| 282 | let client = self |
| 283 | .clients |
| 284 | .get(server_name) |
| 285 | .with_context(|| format!("MCP server '{server_name}' not available"))?; |
| 286 | client.read_resource(uri) |
| 287 | } |
| 288 | |
| 289 | pub fn update_sandbox_state(&self, sandbox_mode: &str, cwd: &str) -> Result<Vec<Value>> { |
| 290 | let mut notices = Vec::new(); |
| 291 | for server_name in self.configs.keys() { |
| 292 | notices.push(json!({ |
| 293 | "server_name": server_name, |
| 294 | "method": "codex/sandbox-state/update", |
| 295 | "params": { |
| 296 | "sandbox_mode": sandbox_mode, |
| 297 | "cwd": cwd |
| 298 | } |
| 299 | })); |
| 300 | } |
| 301 | Ok(notices) |
| 302 | } |
| 303 | } |
| 304 | |
| 305 | fn default_true() -> bool { |
| 306 | true |
| 307 | } |
| 308 | |
| 309 | fn allowed_by_filter(name: &str, filter: &ToolFilter) -> bool { |
| 310 | if filter.deny.iter().any(|pattern| pattern == name) { |
| 311 | return false; |
| 312 | } |
| 313 | if filter.allow.is_empty() { |
| 314 | return true; |
| 315 | } |
| 316 | filter.allow.iter().any(|pattern| pattern == name) |
| 317 | } |
| 318 | |
| 319 | fn sanitize_component(value: &str) -> String { |
| 320 | value |
| 321 | .chars() |
| 322 | .map(|ch| { |
| 323 | if ch.is_ascii_alphanumeric() || ch == '_' { |
| 324 | ch.to_ascii_lowercase() |
| 325 | } else { |
| 326 | '_' |
| 327 | } |
| 328 | }) |
| 329 | .collect() |
| 330 | } |
| 331 | |
| 332 | fn qualify_tool_name(server: &str, tool: &str) -> String { |
| 333 | let mut name = format!( |
| 334 | "mcp__{}__{}", |
| 335 | sanitize_component(server), |
| 336 | sanitize_component(tool) |
| 337 | ); |
| 338 | if name.len() > 64 { |
| 339 | let mut hasher = DefaultHasher::new(); |
| 340 | name.hash(&mut hasher); |
| 341 | let hash = format!("{:x}", hasher.finish()); |
| 342 | name.truncate(48); |
| 343 | name.push('_'); |
| 344 | name.push_str(&hash[..12]); |
| 345 | } |
| 346 | name |
| 347 | } |
| 348 | |
| 349 | fn parse_qualified_tool_name(value: &str) -> Result<(String, String)> { |
| 350 | let Some(stripped) = value.strip_prefix("mcp__") else { |
| 351 | bail!("missing mcp__ prefix"); |
| 352 | }; |
| 353 | let mut split = stripped.splitn(2, "__"); |
| 354 | let server = split |
| 355 | .next() |
| 356 | .filter(|s| !s.is_empty()) |
| 357 | .map(ToOwned::to_owned) |
| 358 | .context("missing server segment")?; |
| 359 | let tool = split |
| 360 | .next() |
| 361 | .filter(|s| !s.is_empty()) |
| 362 | .map(ToOwned::to_owned) |
| 363 | .context("missing tool segment")?; |
| 364 | Ok((server, tool)) |
| 365 | } |
| 366 | |
| 367 | #[derive(Debug, Deserialize)] |
| 368 | struct JsonRpcRequest { |
| 369 | #[serde(default)] |
| 370 | jsonrpc: Option<String>, |
| 371 | #[serde(default)] |
| 372 | id: Option<Value>, |
| 373 | method: String, |
| 374 | #[serde(default)] |
| 375 | params: Value, |
| 376 | } |
| 377 | |
| 378 | #[derive(Debug)] |
| 379 | struct JsonRpcError { |
| 380 | code: i64, |
| 381 | message: String, |
| 382 | data: Option<Value>, |
| 383 | } |
| 384 | |
| 385 | #[derive(Debug, Deserialize)] |
| 386 | struct ToolsListParams { |
| 387 | #[serde(default)] |
| 388 | server: Option<String>, |
| 389 | } |
| 390 | |
| 391 | #[derive(Debug, Deserialize)] |
| 392 | struct ToolsCallParams { |
| 393 | #[serde(default)] |
| 394 | name: Option<String>, |
| 395 | #[serde(default)] |
| 396 | tool: Option<String>, |
| 397 | #[serde(default)] |
| 398 | server: Option<String>, |
| 399 | #[serde(default)] |
| 400 | arguments: Value, |
| 401 | } |
| 402 | |
| 403 | #[derive(Debug, Deserialize)] |
| 404 | struct ResourcesListParams { |
| 405 | #[serde(default)] |
| 406 | server: Option<String>, |
| 407 | } |
| 408 | |
| 409 | #[derive(Debug, Deserialize)] |
| 410 | struct ResourcesReadParams { |
| 411 | #[serde(default)] |
| 412 | server: Option<String>, |
| 413 | uri: String, |
| 414 | } |
| 415 | |
| 416 | #[derive(Debug, Deserialize)] |
| 417 | struct ServerRegisterParams { |
| 418 | server: McpServerConfig, |
| 419 | #[serde(default)] |
| 420 | filter: ToolFilter, |
| 421 | #[serde(default = "default_true")] |
| 422 | start: bool, |
| 423 | } |
| 424 | |
| 425 | #[derive(Debug, Deserialize)] |
| 426 | struct ServerNameParams { |
| 427 | name: String, |
| 428 | } |
| 429 | |
| 430 | struct StdioMcpState { |
| 431 | manager: McpManager, |
| 432 | definitions: HashMap<String, McpServerDefinition>, |
| 433 | running: HashMap<String, bool>, |
| 434 | lifecycle_state: String, |
| 435 | } |
| 436 | |
| 437 | pub fn run_stdio_server( |
| 438 | initial_definitions: Vec<McpServerDefinition>, |
| 439 | ) -> Result<Vec<McpServerDefinition>> { |
| 440 | use std::io::{self, BufRead, Write}; |
| 441 | |
| 442 | let stdin = io::stdin(); |
| 443 | let mut stdout = io::stdout(); |
| 444 | let mut stderr = io::stderr(); |
| 445 | let mut state = build_stdio_state(initial_definitions); |
| 446 | |
| 447 | for line in stdin.lock().lines() { |
| 448 | let line = line.context("failed to read stdio line")?; |
| 449 | if line.trim().is_empty() { |
| 450 | continue; |
| 451 | } |
| 452 | |
| 453 | let request: JsonRpcRequest = match serde_json::from_str(&line) { |
| 454 | Ok(value) => value, |
| 455 | Err(err) => { |
| 456 | let msg = jsonrpc_error( |
| 457 | None, |
| 458 | JsonRpcError::parse_error(format!("invalid json: {err}")), |
| 459 | ); |
| 460 | writeln!(stdout, "{msg}")?; |
| 461 | stdout.flush()?; |
| 462 | continue; |
| 463 | } |
| 464 | }; |
| 465 | |
| 466 | if request |
| 467 | .jsonrpc |
| 468 | .as_deref() |
| 469 | .is_some_and(|version| version != "2.0") |
| 470 | { |
| 471 | let response = jsonrpc_error( |
| 472 | request.id, |
| 473 | JsonRpcError::invalid_request("jsonrpc version must be 2.0"), |
| 474 | ); |
| 475 | writeln!(stdout, "{response}")?; |
| 476 | stdout.flush()?; |
| 477 | continue; |
| 478 | } |
| 479 | |
| 480 | let response = match dispatch_stdio_request(&mut state, &request.method, request.params) { |
| 481 | Ok((result, should_exit)) => { |
| 482 | let payload = jsonrpc_result(request.id, result); |
| 483 | writeln!(stdout, "{payload}")?; |
| 484 | stdout.flush()?; |
| 485 | if should_exit { |
| 486 | break; |
| 487 | } |
| 488 | continue; |
| 489 | } |
| 490 | Err(err) => jsonrpc_error(request.id, err), |
| 491 | }; |
| 492 | |
| 493 | writeln!(stdout, "{response}")?; |
| 494 | stdout.flush()?; |
| 495 | } |
| 496 | |
| 497 | state.lifecycle_state = "stopped".to_string(); |
| 498 | let _ = writeln!(stderr, "deepseek-mcp stdio server exited"); |
| 499 | let mut definitions: Vec<McpServerDefinition> = state.definitions.into_values().collect(); |
| 500 | definitions.sort_by(|a, b| a.config.name.cmp(&b.config.name)); |
| 501 | Ok(definitions) |
| 502 | } |
| 503 | |
| 504 | fn build_stdio_state(initial_definitions: Vec<McpServerDefinition>) -> StdioMcpState { |
| 505 | let mut manager = McpManager::default(); |
| 506 | let mut definitions = HashMap::new(); |
| 507 | let mut running = HashMap::new(); |
| 508 | |
| 509 | for definition in initial_definitions { |
| 510 | let name = definition.config.name.clone(); |
| 511 | let should_start = definition.config.enabled; |
| 512 | definitions.insert(name.clone(), definition.clone()); |
| 513 | if should_start { |
| 514 | manager.register_server( |
| 515 | definition.config.clone(), |
| 516 | definition.filter.clone(), |
| 517 | default_stdio_client(&name), |
| 518 | ); |
| 519 | running.insert(name, true); |
| 520 | } else { |
| 521 | running.insert(name, false); |
| 522 | } |
| 523 | } |
| 524 | |
| 525 | StdioMcpState { |
| 526 | manager, |
| 527 | definitions, |
| 528 | running, |
| 529 | lifecycle_state: "running".to_string(), |
| 530 | } |
| 531 | } |
| 532 | |
| 533 | fn default_stdio_client(server_name: &str) -> Box<dyn McpManagedClient> { |
| 534 | let health_uri = format!("mcp://{server_name}/health"); |
| 535 | let capabilities_uri = format!("mcp://{server_name}/capabilities"); |
| 536 | Box::new( |
| 537 | InMemoryMcpClient::default() |
| 538 | .with_tool( |
| 539 | "health", |
| 540 | json!({ |
| 541 | "status": "ok", |
| 542 | "server_name": server_name |
| 543 | }), |
| 544 | ) |
| 545 | .with_tool( |
| 546 | "capabilities", |
| 547 | json!({ |
| 548 | "tools": ["health", "capabilities"], |
| 549 | "resources": [health_uri.clone(), capabilities_uri.clone()] |
| 550 | }), |
| 551 | ) |
| 552 | .with_resource( |
| 553 | &health_uri, |
| 554 | json!({ |
| 555 | "status": "ok", |
| 556 | "server_name": server_name |
| 557 | }), |
| 558 | ) |
| 559 | .with_resource( |
| 560 | &capabilities_uri, |
| 561 | json!({ |
| 562 | "server_name": server_name, |
| 563 | "methods": [ |
| 564 | "tools/list", |
| 565 | "tools/call", |
| 566 | "resources/list", |
| 567 | "resources/read", |
| 568 | "server/list", |
| 569 | "server/register", |
| 570 | "server/start", |
| 571 | "server/stop", |
| 572 | "server/unregister" |
| 573 | ] |
| 574 | }), |
| 575 | ), |
| 576 | ) |
| 577 | } |
| 578 | |
| 579 | fn default_rpc_methods() -> Vec<&'static str> { |
| 580 | vec![ |
| 581 | "initialize", |
| 582 | "healthz", |
| 583 | "capabilities", |
| 584 | "tools/list", |
| 585 | "tools/call", |
| 586 | "resources/list", |
| 587 | "resources/read", |
| 588 | "server/list", |
| 589 | "server/register", |
| 590 | "server/start", |
| 591 | "server/stop", |
| 592 | "server/unregister", |
| 593 | "shutdown", |
| 594 | ] |
| 595 | } |
| 596 | |
| 597 | fn lifecycle_snapshot(state: &StdioMcpState) -> Value { |
| 598 | let mut servers: Vec<Value> = state |
| 599 | .definitions |
| 600 | .iter() |
| 601 | .map(|(name, definition)| { |
| 602 | let is_running = state.running.get(name).copied().unwrap_or(false); |
| 603 | json!({ |
| 604 | "name": name, |
| 605 | "enabled": definition.config.enabled, |
| 606 | "running": is_running, |
| 607 | "command": definition.config.command.clone(), |
| 608 | "args": definition.config.args.clone(), |
| 609 | }) |
| 610 | }) |
| 611 | .collect(); |
| 612 | servers.sort_by(|a, b| { |
| 613 | let a_name = a.get("name").and_then(Value::as_str).unwrap_or_default(); |
| 614 | let b_name = b.get("name").and_then(Value::as_str).unwrap_or_default(); |
| 615 | a_name.cmp(b_name) |
| 616 | }); |
| 617 | |
| 618 | let running_count = state.running.values().filter(|running| **running).count(); |
| 619 | json!({ |
| 620 | "status": state.lifecycle_state, |
| 621 | "servers": servers, |
| 622 | "counts": { |
| 623 | "defined": state.definitions.len(), |
| 624 | "running": running_count |
| 625 | } |
| 626 | }) |
| 627 | } |
| 628 | |
| 629 | fn params_or_object(params: Value) -> Value { |
| 630 | if params.is_null() { json!({}) } else { params } |
| 631 | } |
| 632 | |
| 633 | fn parse_params<T: DeserializeOwned>(params: Value) -> std::result::Result<T, JsonRpcError> { |
| 634 | serde_json::from_value(params).map_err(|err| JsonRpcError::invalid_params(err.to_string())) |
| 635 | } |
| 636 | |
| 637 | fn parse_server_from_uri(uri: &str) -> Option<String> { |
| 638 | let stripped = uri.strip_prefix("mcp://")?; |
| 639 | let server = stripped.split('/').next()?; |
| 640 | if server.is_empty() { |
| 641 | None |
| 642 | } else { |
| 643 | Some(server.to_string()) |
| 644 | } |
| 645 | } |
| 646 | |
| 647 | fn dispatch_stdio_request( |
| 648 | state: &mut StdioMcpState, |
| 649 | method: &str, |
| 650 | params: Value, |
| 651 | ) -> std::result::Result<(Value, bool), JsonRpcError> { |
| 652 | match method { |
| 653 | "initialize" | "capabilities" => Ok(( |
| 654 | json!({ |
| 655 | "server": "deepseek-mcp", |
| 656 | "transport": "stdio", |
| 657 | "methods": default_rpc_methods(), |
| 658 | "lifecycle": lifecycle_snapshot(state) |
| 659 | }), |
| 660 | false, |
| 661 | )), |
| 662 | "healthz" => Ok(( |
| 663 | json!({ |
| 664 | "status": "ok", |
| 665 | "service": "deepseek-mcp", |
| 666 | "transport": "stdio", |
| 667 | "lifecycle": lifecycle_snapshot(state) |
| 668 | }), |
| 669 | false, |
| 670 | )), |
| 671 | "tools/list" => { |
| 672 | let parsed: ToolsListParams = parse_params(params_or_object(params))?; |
| 673 | let mut tools = state |
| 674 | .manager |
| 675 | .list_tools() |
| 676 | .map_err(|err| JsonRpcError::internal(err.to_string()))?; |
| 677 | if let Some(server) = parsed.server { |
| 678 | tools.retain(|tool| tool.server_name == server); |
| 679 | } |
| 680 | Ok((json!({ "tools": tools }), false)) |
| 681 | } |
| 682 | "tools/call" => { |
| 683 | let parsed: ToolsCallParams = parse_params(params_or_object(params))?; |
| 684 | let ToolsCallParams { |
| 685 | name, |
| 686 | tool, |
| 687 | server, |
| 688 | arguments, |
| 689 | } = parsed; |
| 690 | let tool_name = name |
| 691 | .or(tool) |
| 692 | .context("missing tool name") |
| 693 | .map_err(|err| JsonRpcError::invalid_params(err.to_string()))?; |
| 694 | let arguments = if arguments.is_null() { |
| 695 | json!({}) |
| 696 | } else { |
| 697 | arguments |
| 698 | }; |
| 699 | let result = if tool_name.starts_with("mcp__") { |
| 700 | state |
| 701 | .manager |
| 702 | .call_qualified_tool(&tool_name, arguments) |
| 703 | .map_err(|err| JsonRpcError::internal(err.to_string()))? |
| 704 | } else { |
| 705 | let server = server |
| 706 | .context("missing server for unqualified tool") |
| 707 | .map_err(|err| JsonRpcError::invalid_params(err.to_string()))?; |
| 708 | state |
| 709 | .manager |
| 710 | .call_tool(&server, &tool_name, arguments) |
| 711 | .map_err(|err| JsonRpcError::internal(err.to_string()))? |
| 712 | }; |
| 713 | Ok((json!({ "result": result }), false)) |
| 714 | } |
| 715 | "resources/list" => { |
| 716 | let parsed: ResourcesListParams = parse_params(params_or_object(params))?; |
| 717 | let mut resources = state |
| 718 | .manager |
| 719 | .list_resources() |
| 720 | .map_err(|err| JsonRpcError::internal(err.to_string()))?; |
| 721 | if let Some(server) = parsed.server { |
| 722 | resources.retain(|resource| resource.server_name == server); |
| 723 | } |
| 724 | Ok((json!({ "resources": resources }), false)) |
| 725 | } |
| 726 | "resources/read" => { |
| 727 | let parsed: ResourcesReadParams = parse_params(params_or_object(params))?; |
| 728 | let ResourcesReadParams { server, uri } = parsed; |
| 729 | let server_name = server |
| 730 | .or_else(|| parse_server_from_uri(&uri)) |
| 731 | .context("missing server for resource read") |
| 732 | .map_err(|err| JsonRpcError::invalid_params(err.to_string()))?; |
| 733 | let value = state |
| 734 | .manager |
| 735 | .read_resource(&server_name, &uri) |
| 736 | .map_err(|err| JsonRpcError::internal(err.to_string()))?; |
| 737 | Ok((json!({ "resource": value }), false)) |
| 738 | } |
| 739 | "server/list" | "servers/list" => { |
| 740 | Ok((json!({ "lifecycle": lifecycle_snapshot(state) }), false)) |
| 741 | } |
| 742 | "server/register" | "servers/register" => { |
| 743 | let parsed: ServerRegisterParams = parse_params(params_or_object(params))?; |
| 744 | let name = parsed.server.name.clone(); |
| 745 | if name.trim().is_empty() { |
| 746 | return Err(JsonRpcError::invalid_params( |
| 747 | "server.name must not be empty", |
| 748 | )); |
| 749 | } |
| 750 | |
| 751 | if state.definitions.contains_key(&name) { |
| 752 | let _ = state.manager.unregister_server(&name); |
| 753 | } |
| 754 | state.definitions.insert( |
| 755 | name.clone(), |
| 756 | McpServerDefinition { |
| 757 | config: parsed.server.clone(), |
| 758 | filter: parsed.filter.clone(), |
| 759 | }, |
| 760 | ); |
| 761 | let should_run = parsed.start && parsed.server.enabled; |
| 762 | if should_run { |
| 763 | state.manager.register_server( |
| 764 | parsed.server.clone(), |
| 765 | parsed.filter.clone(), |
| 766 | default_stdio_client(&name), |
| 767 | ); |
| 768 | } |
| 769 | state.running.insert(name, should_run); |
| 770 | Ok((json!({ "lifecycle": lifecycle_snapshot(state) }), false)) |
| 771 | } |
| 772 | "server/start" | "servers/start" => { |
| 773 | let parsed: ServerNameParams = parse_params(params_or_object(params))?; |
| 774 | let definition = state |
| 775 | .definitions |
| 776 | .get(&parsed.name) |
| 777 | .cloned() |
| 778 | .with_context(|| format!("server '{}' is not defined", parsed.name)) |
| 779 | .map_err(|err| JsonRpcError::invalid_params(err.to_string()))?; |
| 780 | if !definition.config.enabled { |
| 781 | return Err(JsonRpcError::invalid_params(format!( |
| 782 | "server '{}' is disabled", |
| 783 | parsed.name |
| 784 | ))); |
| 785 | } |
| 786 | if !state.running.get(&parsed.name).copied().unwrap_or(false) { |
| 787 | state.manager.register_server( |
| 788 | definition.config.clone(), |
| 789 | definition.filter.clone(), |
| 790 | default_stdio_client(&parsed.name), |
| 791 | ); |
| 792 | state.running.insert(parsed.name, true); |
| 793 | } |
| 794 | Ok((json!({ "lifecycle": lifecycle_snapshot(state) }), false)) |
| 795 | } |
| 796 | "server/stop" | "servers/stop" => { |
| 797 | let parsed: ServerNameParams = parse_params(params_or_object(params))?; |
| 798 | if state.running.get(&parsed.name).copied().unwrap_or(false) { |
| 799 | state |
| 800 | .manager |
| 801 | .stop_server(&parsed.name) |
| 802 | .map_err(|err| JsonRpcError::internal(err.to_string()))?; |
| 803 | } |
| 804 | state.running.insert(parsed.name, false); |
| 805 | Ok((json!({ "lifecycle": lifecycle_snapshot(state) }), false)) |
| 806 | } |
| 807 | "server/unregister" | "servers/unregister" => { |
| 808 | let parsed: ServerNameParams = parse_params(params_or_object(params))?; |
| 809 | if state.definitions.remove(&parsed.name).is_none() { |
| 810 | return Err(JsonRpcError::invalid_params(format!( |
| 811 | "server '{}' is not defined", |
| 812 | parsed.name |
| 813 | ))); |
| 814 | } |
| 815 | let _ = state.manager.unregister_server(&parsed.name); |
| 816 | state.running.remove(&parsed.name); |
| 817 | Ok((json!({ "lifecycle": lifecycle_snapshot(state) }), false)) |
| 818 | } |
| 819 | "shutdown" => { |
| 820 | state.lifecycle_state = "shutting_down".to_string(); |
| 821 | Ok(( |
| 822 | json!({ |
| 823 | "ok": true, |
| 824 | "lifecycle": lifecycle_snapshot(state) |
| 825 | }), |
| 826 | true, |
| 827 | )) |
| 828 | } |
| 829 | _ => Err(JsonRpcError::method_not_found(method)), |
| 830 | } |
| 831 | } |
| 832 | |
| 833 | fn jsonrpc_result(id: Option<Value>, result: Value) -> Value { |
| 834 | json!({ |
| 835 | "jsonrpc": "2.0", |
| 836 | "id": id.unwrap_or(Value::Null), |
| 837 | "result": result |
| 838 | }) |
| 839 | } |
| 840 | |
| 841 | fn jsonrpc_error(id: Option<Value>, err: JsonRpcError) -> Value { |
| 842 | json!({ |
| 843 | "jsonrpc": "2.0", |
| 844 | "id": id.unwrap_or(Value::Null), |
| 845 | "error": { |
| 846 | "code": err.code, |
| 847 | "message": err.message, |
| 848 | "data": err.data |
| 849 | } |
| 850 | }) |
| 851 | } |
| 852 | |
| 853 | impl JsonRpcError { |
| 854 | fn parse_error(message: impl Into<String>) -> Self { |
| 855 | Self { |
| 856 | code: -32700, |
| 857 | message: message.into(), |
| 858 | data: None, |
| 859 | } |
| 860 | } |
| 861 | |
| 862 | fn invalid_request(message: impl Into<String>) -> Self { |
| 863 | Self { |
| 864 | code: -32600, |
| 865 | message: message.into(), |
| 866 | data: None, |
| 867 | } |
| 868 | } |
| 869 | |
| 870 | fn method_not_found(method: &str) -> Self { |
| 871 | Self { |
| 872 | code: -32601, |
| 873 | message: format!("unsupported method: {method}"), |
| 874 | data: None, |
| 875 | } |
| 876 | } |
| 877 | |
| 878 | fn invalid_params(message: impl Into<String>) -> Self { |
| 879 | Self { |
| 880 | code: -32602, |
| 881 | message: message.into(), |
| 882 | data: None, |
| 883 | } |
| 884 | } |
| 885 | |
| 886 | fn internal(message: impl Into<String>) -> Self { |
| 887 | Self { |
| 888 | code: -32603, |
| 889 | message: message.into(), |
| 890 | data: None, |
| 891 | } |
| 892 | } |
| 893 | } |
| 894 |