use std::collections::HashMap; use std::sync::Arc; use anyhow::Context; use forge_app::domain::{ McpConfig, McpServerConfig, McpServers, ServerName, ToolCallFull, ToolDefinition, ToolName, ToolOutput, }; use forge_app::{ EnvironmentInfra, KVStore, McpClientInfra, McpConfigManager, McpServerInfra, McpService, }; use tokio::sync::{Mutex, RwLock}; use crate::mcp::tool::McpExecutor; fn generate_mcp_tool_name(server_name: &ServerName, tool_name: &ToolName) -> ToolName { let sanitized_server_name = ToolName::sanitized(server_name.to_string().as_str()); let sanitized_tool_name = tool_name.clone().into_sanitized(); ToolName::new(format!( "mcp_{sanitized_server_name}_tool_{sanitized_tool_name}" )) } #[derive(Clone)] pub struct ForgeMcpService { tools: Arc>>>>, failed_servers: Arc>>, previous_config_hash: Arc>, init_lock: Arc>, manager: Arc, infra: Arc, } #[derive(Clone)] struct ToolHolder { definition: ToolDefinition, executable: T, server_name: String, } impl ForgeMcpService where M: McpConfigManager, I: McpServerInfra + KVStore + EnvironmentInfra, C: McpClientInfra + Clone, C: From<::Client>, { pub fn new(manager: Arc, infra: Arc) -> Self { Self { tools: Default::default(), failed_servers: Default::default(), previous_config_hash: Arc::new(Mutex::new(Default::default())), init_lock: Arc::new(Mutex::new(())), manager, infra, } } async fn is_config_modified(&self, config: &McpConfig) -> bool { *self.previous_config_hash.lock().await != config.cache_key() } async fn insert_clients(&self, server_name: &ServerName, client: Arc) -> anyhow::Result<()> { let tools = client.list().await?; let mut tool_map = self.tools.write().await; for mut tool in tools.into_iter() { let actual_name = tool.name.clone(); let server = McpExecutor::new(actual_name, client.clone())?; let generated_name = generate_mcp_tool_name(server_name, &tool.name); tool.name = generated_name.clone(); tool_map.insert( generated_name, ToolHolder { definition: tool, executable: server, server_name: server_name.to_string(), }, ); } Ok(()) } async fn connect( &self, server_name: &ServerName, config: McpServerConfig, ) -> anyhow::Result<()> { let env_vars = self.infra.get_env_vars(); let environment = self.infra.get_environment(); let client = self.infra.connect(config, &env_vars, &environment).await?; let client = Arc::new(C::from(client)); self.insert_clients(server_name, client).await?; Ok(()) } async fn ensure_mcp_initialized(&self) -> anyhow::Result<()> { let raw_mcp = self.manager.read_mcp_config(None).await?; // Fast path: if config is unchanged, skip reinitialization without acquiring // the lock if !self.is_config_modified(&raw_mcp).await { return Ok(()); } // Serialise concurrent initialisations so only one caller runs update_mcp at a // time let _guard = self.init_lock.lock().await; // Double-check under the lock: a concurrent caller may have already updated if !self.is_config_modified(&raw_mcp).await { return Ok(()); } // Apply the trust gate. The prompt was already shown at startup via // init_mcp, so on first tool use the user's decision is re-applied by // calling filter_trusted again. let trusted_mcp = self.manager.filter_trusted(raw_mcp).await?; self.update_mcp(trusted_mcp).await } async fn update_mcp(&self, mcp: McpConfig) -> Result<(), anyhow::Error> { // Use the raw config hash (pre-trust-gate) so that is_config_modified always // compares against the original file hash, preventing infinite re-prompt loops // when some servers are rejected. let new_hash = mcp.cache_key(); self.clear_tools().await; // Clear failed servers map before attempting new connections self.failed_servers.write().await.clear(); let connections: Vec<_> = mcp .mcp_servers .into_iter() .filter(|v| !v.1.is_disabled()) .map(|(name, server)| async move { let conn = self .connect(&name, server) .await .context(format!("Failed to initiate MCP server: {name}")); (name, conn) }) .collect(); let results = futures::future::join_all(connections).await; for (server_name, result) in results { match result { Ok(_) => {} Err(error) => { // Format error with full chain for detailed diagnostics // Using Debug formatting with alternate flag shows the full error chain let error_string = format!("{error:?}"); self.failed_servers .write() .await .insert(server_name.clone(), error_string.clone()); } } } // Write the hash only after join_all finishes so that any waiter on // init_lock re-checks is_config_modified only once self.tools is fully // populated, preventing "Tool not found" races. *self.previous_config_hash.lock().await = new_hash; Ok(()) } async fn list(&self) -> anyhow::Result { self.ensure_mcp_initialized().await?; let tools = self.tools.read().await; let mut grouped_tools = std::collections::HashMap::new(); for tool in tools.values() { grouped_tools .entry(ServerName::from(tool.server_name.clone())) .or_insert_with(Vec::new) .push(tool.definition.clone()); } let failures = self.failed_servers.read().await.clone(); Ok(McpServers::new(grouped_tools, failures)) } async fn clear_tools(&self) { self.tools.write().await.clear() } async fn call(&self, call: ToolCallFull) -> anyhow::Result { // Ensure MCP connections are initialized before calling tools self.ensure_mcp_initialized().await?; let tools = self.tools.read().await; // Try exact match first, then fall back to legacy-format lookup for // tool calls arriving in the Claude Code `mcp__{server}__{tool}` format. let tool = tools .get(&call.name) .or_else(|| call.name.to_legacy_mcp_name().and_then(|n| tools.get(&n))) .context("Tool not found")?; tool.executable.call_tool(call.arguments.parse()?).await } /// Refresh the MCP cache by clearing cached data. /// Does NOT eagerly connect to servers - connections happen lazily /// when list() or call() is invoked, avoiding interactive OAuth during /// reload. async fn refresh_cache(&self) -> anyhow::Result<()> { // Hold init_lock so we don't race with an in-flight update_mcp: without // this, clear_tools could run while connections are still being // established, leaving waiters released into an empty tool map. let _guard = self.init_lock.lock().await; // Clear the infra cache and reset config hash to force re-init on next access self.infra.cache_clear().await?; *self.previous_config_hash.lock().await = Default::default(); self.clear_tools().await; self.failed_servers.write().await.clear(); Ok(()) } } #[async_trait::async_trait] impl McpService for ForgeMcpService where C: McpClientInfra + Clone, C: From<::Client>, { async fn get_mcp_servers(&self) -> anyhow::Result { // Apply the trust gate before computing the cache key so that rejected // servers are excluded. Using the raw config hash would allow a stale KV // cache entry (populated before a rejection) to be returned, bypassing // filter_trusted entirely and leaking rejected tools into requests. let raw_mcp = self.manager.read_mcp_config(None).await?; let trusted_mcp = self.manager.filter_trusted(raw_mcp).await?; let config_hash = trusted_mcp.cache_key(); // Check if cache is valid (exists and not expired) if let Some(cache) = self.infra.cache_get::<_, McpServers>(&config_hash).await? { return Ok(cache.clone()); } let servers = self.list().await?; self.infra.cache_set(&config_hash, &servers).await?; Ok(servers) } async fn execute_mcp(&self, call: ToolCallFull) -> anyhow::Result { self.call(call).await } async fn reload_mcp(&self) -> anyhow::Result<()> { self.refresh_cache().await } async fn init_mcp(&self) -> anyhow::Result<()> { // Run the trust gate prompt at startup so the user's decision is captured // before any tool use. The result is intentionally discarded — servers are // NOT connected here. Connections remain lazy and happen on first tool use // via ensure_mcp_initialized. let raw_mcp = self.manager.read_mcp_config(None).await?; let _ = self.manager.filter_trusted(raw_mcp).await?; Ok(()) } } #[cfg(test)] mod tests { use std::collections::BTreeMap; use std::sync::Arc; use fake::{Fake, Faker}; use forge_app::domain::{ ConfigOperation, Environment, McpConfig, McpServerConfig, Scope, ServerName, ToolCallFull, ToolDefinition, ToolName, ToolOutput, }; use forge_app::{ EnvironmentInfra, KVStore, McpClientInfra, McpConfigManager, McpServerInfra, McpService, }; use forge_config::ForgeConfig; use pretty_assertions::assert_eq; use serde::de::DeserializeOwned; use super::{ForgeMcpService, generate_mcp_tool_name}; // ── Mock MCP client ────────────────────────────────────────────────────── #[derive(Clone)] struct MockMcpClient; #[async_trait::async_trait] impl McpClientInfra for MockMcpClient { async fn list(&self) -> anyhow::Result> { Ok(vec![ToolDefinition::new("test_tool")]) } async fn call( &self, _tool_name: &ToolName, _input: serde_json::Value, ) -> anyhow::Result { Ok(ToolOutput::text("mock result")) } } // ── Mock config manager ────────────────────────────────────────────────── struct MockMcpManager; #[async_trait::async_trait] impl McpConfigManager for MockMcpManager { async fn read_mcp_config(&self, _scope: Option<&Scope>) -> anyhow::Result { let mut servers = BTreeMap::new(); servers.insert( ServerName::from("test-server".to_string()), McpServerConfig::new_stdio("echo", vec![], None), ); Ok(McpConfig { mcp_servers: servers }) } async fn write_mcp_config( &self, _config: &McpConfig, _scope: &Scope, ) -> anyhow::Result<()> { Ok(()) } async fn filter_trusted(&self, raw: McpConfig) -> anyhow::Result { // In tests all configs are implicitly trusted. Ok(raw) } } // ── Mock infrastructure ────────────────────────────────────────────────── #[derive(Clone)] struct MockInfra; #[async_trait::async_trait] impl McpServerInfra for MockInfra { type Client = MockMcpClient; async fn connect( &self, _config: McpServerConfig, _env_vars: &BTreeMap, _environment: &Environment, ) -> anyhow::Result { Ok(MockMcpClient) } } #[async_trait::async_trait] impl KVStore for MockInfra { async fn cache_get(&self, _key: &K) -> anyhow::Result> where K: std::hash::Hash + Sync, V: serde::Serialize + DeserializeOwned + Send, { Ok(None) } async fn cache_set(&self, _key: &K, _value: &V) -> anyhow::Result<()> where K: std::hash::Hash + Sync, V: serde::Serialize + Sync, { Ok(()) } async fn cache_clear(&self) -> anyhow::Result<()> { Ok(()) } } impl EnvironmentInfra for MockInfra { type Config = ForgeConfig; fn get_env_var(&self, _key: &str) -> Option { None } fn get_env_vars(&self) -> BTreeMap { BTreeMap::new() } fn get_environment(&self) -> Environment { Faker.fake() } fn get_config(&self) -> anyhow::Result { Ok(ForgeConfig::default()) } async fn update_environment(&self, _ops: Vec) -> anyhow::Result<()> { Ok(()) } } // ── Fixture ────────────────────────────────────────────────────────────── fn fixture() -> ForgeMcpService { ForgeMcpService::new(Arc::new(MockMcpManager), Arc::new(MockInfra)) } #[test] fn test_generate_mcp_tool_name_uses_legacy_format() { let fixture = ServerName::from("hugging-face".to_string()); let actual = generate_mcp_tool_name(&fixture, &ToolName::new("read-channel")); let expected = ToolName::new("mcp_hugging_face_tool_read_channel"); assert_eq!(actual, expected); } #[test] fn test_generate_mcp_tool_name_sanitizes_server_and_tool_names() { let fixture = ServerName::from("claude.ai Slack".to_string()); let actual = generate_mcp_tool_name(&fixture, &ToolName::new("Add comment")); let expected = ToolName::new("mcp_claude_ai_slack_tool_add_comment"); assert_eq!(actual, expected); } #[test] fn test_to_legacy_mcp_name_converts_claude_code_format() { let actual = ToolName::new("mcp__github__create_issue").to_legacy_mcp_name(); let expected = Some(ToolName::new("mcp_github_tool_create_issue")); assert_eq!(actual, expected); } #[test] fn test_to_legacy_mcp_name_converts_multipart_server_name() { let actual = ToolName::new("mcp__hugging_face__read_channel").to_legacy_mcp_name(); let expected = Some(ToolName::new("mcp_hugging_face_tool_read_channel")); assert_eq!(actual, expected); } #[test] fn test_to_legacy_mcp_name_returns_none_for_non_mcp_tools() { let actual = ToolName::new("read").to_legacy_mcp_name(); assert_eq!(actual, None); } #[test] fn test_to_legacy_mcp_name_returns_none_for_legacy_format() { // Already in legacy format — should not double-convert let actual = ToolName::new("mcp_github_tool_create_issue").to_legacy_mcp_name(); assert_eq!(actual, None); } // ── Concurrent initialisation test ────────────────────────────────────── /// Verify that two concurrent callers of `get_mcp_servers` do not race: /// after both futures settle, every registered tool must be callable /// without a "Tool not found" error. #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn test_concurrent_init_does_not_race() { let service = Arc::new(fixture()); let s1 = service.clone(); let s2 = service.clone(); let (r1, r2) = tokio::join!(s1.get_mcp_servers(), s2.get_mcp_servers()); r1.unwrap(); r2.unwrap(); let servers = service.get_mcp_servers().await.unwrap(); let tool_name = servers .get_servers() .values() .flat_map(|tools| tools.iter()) .next() .expect("at least one tool must be registered") .name .clone(); let call = ToolCallFull::new(tool_name); let actual = service.execute_mcp(call).await.unwrap(); let expected = ToolOutput::text("mock result"); assert_eq!(actual, expected); } }