diff --git a/crates/mcp/src/lib.rs b/crates/mcp/src/lib.rs index c3fd10f5..df2325cd 100644 --- a/crates/mcp/src/lib.rs +++ b/crates/mcp/src/lib.rs @@ -318,8 +318,31 @@ impl McpManager { qualified_tool_name: &str, arguments: Value, ) -> Result { - let (server_name, tool_name) = parse_qualified_tool_name(qualified_tool_name) - .with_context(|| format!("invalid qualified MCP tool name: {qualified_tool_name}"))?; + let parsed = parse_qualified_tool_name(qualified_tool_name) + .with_context(|| format!("invalid qualified MCP tool name: {qualified_tool_name}")); + + if let Ok((server_name, tool_name)) = &parsed + && self.clients.contains_key(server_name) + && let Ok(result) = self.call_tool(server_name, tool_name, arguments.clone()) + { + return Ok(result); + } + + for (server_name, (_, filter)) in &self.configs { + let Some(client) = self.clients.get(server_name) else { + continue; + }; + for tool in client.list_tools()? { + if !allowed_by_filter(&tool.tool_name, filter) { + continue; + } + if qualify_tool_name(server_name, &tool.tool_name) == qualified_tool_name { + return client.call_tool(&tool.tool_name, arguments); + } + } + } + + let (server_name, tool_name) = parsed?; self.call_tool(&server_name, &tool_name, arguments) } @@ -392,18 +415,29 @@ fn sanitize_component(value: &str) -> String { } fn qualify_tool_name(server: &str, tool: &str) -> String { - let mut name = format!( - "mcp__{}__{}", - sanitize_component(server), - sanitize_component(tool) - ); + let server = sanitize_component(server); + let tool = sanitize_component(tool); + let mut name = format!("mcp__{server}__{tool}"); if name.len() > 64 { let mut hasher = DefaultHasher::new(); name.hash(&mut hasher); let hash = format!("{:x}", hasher.finish()); - name.truncate(48); - name.push('_'); - name.push_str(&hash[..12]); + let suffix = format!("_{}", &hash[..12]); + let component_budget = 64 - "mcp__".len() - "__".len() - suffix.len(); + let mut server_len = server.len().min(component_budget / 2); + let mut tool_len = tool.len().min(component_budget - server_len); + let remaining = component_budget - server_len - tool_len; + if remaining > 0 { + let server_extra = (server.len() - server_len).min(remaining); + server_len += server_extra; + tool_len += (tool.len() - tool_len).min(remaining - server_extra); + } + name = format!( + "mcp__{}__{}{}", + &server[..server_len], + &tool[..tool_len], + suffix + ); } name } @@ -957,3 +991,416 @@ impl JsonRpcError { } } } + +#[cfg(test)] +mod tests { + use super::*; + + // ── InMemoryMcpClient ────────────────────────────────────────────── + + #[test] + fn in_memory_client_list_tools_returns_registered() { + let client = InMemoryMcpClient::default() + .with_tool("echo", json!({"output": "hi"})) + .with_tool("greet", json!({"msg": "hello"})); + let tools = client.list_tools().unwrap(); + assert_eq!(tools.len(), 2); + let names: Vec<&str> = tools.iter().map(|t| t.tool_name.as_str()).collect(); + assert!(names.contains(&"echo")); + assert!(names.contains(&"greet")); + } + + #[test] + fn in_memory_client_call_tool_returns_value() { + let client = InMemoryMcpClient::default().with_tool("echo", json!({"output": "hi"})); + let result = client.call_tool("echo", json!({})).unwrap(); + assert_eq!(result["output"], "hi"); + } + + #[test] + fn in_memory_client_call_tool_errors_on_missing() { + let client = InMemoryMcpClient::default(); + let err = client.call_tool("nope", json!({})).unwrap_err(); + assert!(err.to_string().contains("not found")); + } + + #[test] + fn in_memory_client_list_resources_returns_registered() { + let client = InMemoryMcpClient::default() + .with_resource("mcp://s/health", json!({"ok": true})) + .with_resource("mcp://s/caps", json!({"tools": []})); + let resources = client.list_resources().unwrap(); + assert_eq!(resources.len(), 2); + } + + #[test] + fn in_memory_client_read_resource_returns_value() { + let client = + InMemoryMcpClient::default().with_resource("mcp://s/health", json!({"ok": true})); + let result = client.read_resource("mcp://s/health").unwrap(); + assert_eq!(result["ok"], true); + } + + #[test] + fn in_memory_client_read_resource_errors_on_missing() { + let client = InMemoryMcpClient::default(); + let err = client.read_resource("mcp://s/nope").unwrap_err(); + assert!(err.to_string().contains("not found")); + } + + // ── McpManager ───────────────────────────────────────────────────── + + fn make_server_config(name: &str) -> McpServerConfig { + McpServerConfig { + name: name.to_string(), + command: "test".to_string(), + args: vec![], + env: HashMap::new(), + enabled: true, + } + } + + #[test] + fn manager_start_all_marks_ready_for_registered_clients() { + let mut manager = McpManager::default(); + manager.register_server( + make_server_config("s1"), + ToolFilter::default(), + Box::new(InMemoryMcpClient::default().with_tool("t", json!(null))), + ); + let mut events = Vec::new(); + let summary = manager.start_all(|e| events.push(e)); + assert_eq!(summary.ready, vec!["s1"]); + assert!(summary.failed.is_empty()); + assert!(events.iter().any(|event| { + event.server_name == "s1" && event.status == McpStartupStatus::Starting + })); + assert!( + events.iter().any(|event| { + event.server_name == "s1" && event.status == McpStartupStatus::Ready + }) + ); + } + + #[test] + fn manager_start_all_marks_failed_when_client_missing() { + let mut manager = McpManager::default(); + manager.register_server( + make_server_config("s1"), + ToolFilter::default(), + Box::new(InMemoryMcpClient::default()), + ); + manager.stop_server("s1").unwrap(); + let summary = manager.start_all(|_| {}); + assert!(summary.ready.is_empty()); + assert_eq!(summary.failed.len(), 1); + assert_eq!(summary.failed[0].server_name, "s1"); + } + + #[test] + fn manager_start_all_cancels_disabled_servers() { + let mut manager = McpManager::default(); + let mut cfg = make_server_config("s1"); + cfg.enabled = false; + manager.register_server( + cfg, + ToolFilter::default(), + Box::new(InMemoryMcpClient::default()), + ); + let summary = manager.start_all(|_| {}); + assert!(summary.ready.is_empty()); + assert_eq!(summary.cancelled, vec!["s1"]); + } + + #[test] + fn manager_list_tools_applies_filter() { + let mut manager = McpManager::default(); + let client = InMemoryMcpClient::default() + .with_tool("allowed", json!(null)) + .with_tool("denied", json!(null)); + manager.register_server( + make_server_config("s1"), + ToolFilter { + allow: vec!["allowed".to_string()], + deny: vec![], + }, + Box::new(client), + ); + let tools = manager.list_tools().unwrap(); + assert_eq!(tools.len(), 1); + assert_eq!(tools[0].tool_name, "allowed"); + } + + #[test] + fn manager_list_tools_deny_overrides_allow() { + let mut manager = McpManager::default(); + let client = InMemoryMcpClient::default() + .with_tool("a", json!(null)) + .with_tool("b", json!(null)); + manager.register_server( + make_server_config("s1"), + ToolFilter { + allow: vec!["a".to_string(), "b".to_string()], + deny: vec!["b".to_string()], + }, + Box::new(client), + ); + let tools = manager.list_tools().unwrap(); + assert_eq!(tools.len(), 1); + assert_eq!(tools[0].tool_name, "a"); + } + + #[test] + fn manager_call_tool_delegates_to_client() { + let mut manager = McpManager::default(); + manager.register_server( + make_server_config("s1"), + ToolFilter::default(), + Box::new(InMemoryMcpClient::default().with_tool("t", json!({"v": 42}))), + ); + let result = manager.call_tool("s1", "t", json!({})).unwrap(); + assert_eq!(result["v"], 42); + } + + #[test] + fn manager_call_tool_errors_on_missing_server() { + let manager = McpManager::default(); + let err = manager.call_tool("nope", "t", json!({})).unwrap_err(); + assert!(err.to_string().contains("not available")); + } + + #[test] + fn manager_call_qualified_tool_parses_name() { + let mut manager = McpManager::default(); + manager.register_server( + make_server_config("my_server"), + ToolFilter::default(), + Box::new(InMemoryMcpClient::default().with_tool("my_tool", json!({"ok": true}))), + ); + let result = manager + .call_qualified_tool("mcp__my_server__my_tool", json!({})) + .unwrap(); + assert_eq!(result["ok"], true); + } + + #[test] + fn manager_call_qualified_tool_handles_truncated_names() { + let long_server = "server".repeat(20); + let long_tool = "tool".repeat(20); + let mut manager = McpManager::default(); + manager.register_server( + make_server_config(&long_server), + ToolFilter::default(), + Box::new(InMemoryMcpClient::default().with_tool(&long_tool, json!({"ok": true}))), + ); + let tools = manager.list_tools().unwrap(); + let qualified = &tools[0].qualified_name; + assert!(qualified.len() <= 64); + assert!(parse_qualified_tool_name(qualified).is_ok()); + + let result = manager.call_qualified_tool(qualified, json!({})).unwrap(); + assert_eq!(result["ok"], true); + } + + #[test] + fn manager_unregister_removes_server() { + let mut manager = McpManager::default(); + manager.register_server( + make_server_config("s1"), + ToolFilter::default(), + Box::new(InMemoryMcpClient::default()), + ); + manager.unregister_server("s1").unwrap(); + assert!(manager.configs.is_empty()); + } + + #[test] + fn manager_unregister_errors_on_unknown() { + let mut manager = McpManager::default(); + let err = manager.unregister_server("nope").unwrap_err(); + assert!(err.to_string().contains("not registered")); + } + + #[test] + fn manager_stop_server_errors_on_unknown() { + let mut manager = McpManager::default(); + let err = manager.stop_server("nope").unwrap_err(); + assert!(err.to_string().contains("not running")); + } + + #[test] + fn manager_list_resources_returns_from_clients() { + let mut manager = McpManager::default(); + manager.register_server( + make_server_config("s1"), + ToolFilter::default(), + Box::new( + InMemoryMcpClient::default().with_resource("mcp://s1/health", json!({"ok": true})), + ), + ); + let resources = manager.list_resources().unwrap(); + assert_eq!(resources.len(), 1); + assert_eq!(resources[0].server_name, "s1"); + } + + #[test] + fn manager_read_resource_delegates() { + let mut manager = McpManager::default(); + manager.register_server( + make_server_config("s1"), + ToolFilter::default(), + Box::new( + InMemoryMcpClient::default().with_resource("mcp://s1/health", json!({"ok": true})), + ), + ); + let result = manager.read_resource("s1", "mcp://s1/health").unwrap(); + assert_eq!(result["ok"], true); + } + + #[test] + fn manager_update_sandbox_state_returns_notices() { + let mut manager = McpManager::default(); + manager.register_server( + make_server_config("s1"), + ToolFilter::default(), + Box::new(InMemoryMcpClient::default()), + ); + let notices = manager.update_sandbox_state("strict", "/tmp").unwrap(); + assert_eq!(notices.len(), 1); + assert_eq!(notices[0]["server_name"], "s1"); + } + + // ── Tool filter ──────────────────────────────────────────────────── + + #[test] + fn allowed_by_filter_empty_allow_permits_all() { + let filter = ToolFilter { + allow: vec![], + deny: vec![], + }; + assert!(allowed_by_filter("anything", &filter)); + } + + #[test] + fn allowed_by_filter_deny_blocks() { + let filter = ToolFilter { + allow: vec![], + deny: vec!["danger".to_string()], + }; + assert!(!allowed_by_filter("danger", &filter)); + assert!(allowed_by_filter("safe", &filter)); + } + + #[test] + fn allowed_by_filter_allow_only_permits_listed() { + let filter = ToolFilter { + allow: vec!["a".to_string()], + deny: vec![], + }; + assert!(allowed_by_filter("a", &filter)); + assert!(!allowed_by_filter("b", &filter)); + } + + // ── Helper functions ─────────────────────────────────────────────── + + #[test] + fn sanitize_component_lowercases_and_replaces_specials() { + assert_eq!(sanitize_component("My-Server.Name"), "my_server_name"); + assert_eq!(sanitize_component("ABC123"), "abc123"); + } + + #[test] + fn qualify_tool_name_produces_mcp_prefix() { + let name = qualify_tool_name("server", "tool"); + assert!(name.starts_with("mcp__server__tool")); + } + + #[test] + fn qualify_tool_name_truncates_long_names() { + let long_server = "a".repeat(100); + let name = qualify_tool_name(&long_server, "tool"); + assert!(name.len() <= 64); + assert!(parse_qualified_tool_name(&name).is_ok()); + } + + #[test] + fn parse_qualified_tool_name_round_trip() { + let qualified = qualify_tool_name("my_server", "my_tool"); + let (server, tool) = parse_qualified_tool_name(&qualified).unwrap(); + assert_eq!(server, "my_server"); + assert_eq!(tool, "my_tool"); + } + + #[test] + fn parse_qualified_tool_name_rejects_missing_prefix() { + let err = parse_qualified_tool_name("not_mcp__server__tool").unwrap_err(); + assert!(err.to_string().contains("missing mcp__ prefix")); + } + + #[test] + fn parse_qualified_tool_name_rejects_empty_segments() { + let err = parse_qualified_tool_name("mcp____tool").unwrap_err(); + assert!(err.to_string().contains("missing server segment")); + } + + #[test] + fn parse_server_from_uri_extracts_server() { + assert_eq!( + parse_server_from_uri("mcp://my-server/capabilities"), + Some("my-server".to_string()) + ); + } + + #[test] + fn parse_server_from_uri_returns_none_for_invalid() { + assert!(parse_server_from_uri("http://not-mcp").is_none()); + assert!(parse_server_from_uri("mcp:///path").is_none()); + } + + // ── JsonRpcError ─────────────────────────────────────────────────── + + #[test] + fn jsonrpc_error_codes_are_correct() { + assert_eq!(JsonRpcError::parse_error("").code, -32700); + assert_eq!(JsonRpcError::invalid_request("").code, -32600); + assert_eq!(JsonRpcError::method_not_found("x").code, -32601); + assert_eq!(JsonRpcError::invalid_params("").code, -32602); + assert_eq!(JsonRpcError::internal("").code, -32603); + } + + #[test] + fn jsonrpc_result_produces_valid_envelope() { + let result = jsonrpc_result(Some(json!(1)), json!({"ok": true})); + assert_eq!(result["jsonrpc"], "2.0"); + assert_eq!(result["id"], 1); + assert_eq!(result["result"]["ok"], true); + } + + #[test] + fn jsonrpc_error_produces_valid_envelope() { + let err = jsonrpc_error(Some(json!(2)), JsonRpcError::invalid_params("bad")); + assert_eq!(err["jsonrpc"], "2.0"); + assert_eq!(err["id"], 2); + assert_eq!(err["error"]["code"], -32602); + } + + // ── McpServerConfig serialization ────────────────────────────────── + + #[test] + fn mcp_server_config_defaults_enabled_to_true() { + let json = json!({"name": "s", "command": "cmd"}); + let config: McpServerConfig = serde_json::from_value(json).unwrap(); + assert!(config.enabled); + assert!(config.args.is_empty()); + assert!(config.env.is_empty()); + } + + #[test] + fn mcp_startup_status_serializes_with_snake_case() { + let status = McpStartupStatus::Failed { + error: "oops".to_string(), + }; + let json = serde_json::to_value(&status).unwrap(); + assert_eq!(json["failed"]["error"], "oops"); + } +}