codei_mcp/
manager.rs

1use std::sync::Arc;
2
3use codei_config::{load_mcp_config, McpConfig, McpServer};
4use serde_json::Value;
5use tokio::sync::Mutex;
6use tracing::{info, warn};
7
8use crate::client::{McpClient, McpToolInfo};
9use crate::error::McpError;
10
11/// A connected MCP server with its discovered tools.
12pub struct McpConnection {
13    pub server_name: String,
14    client: Arc<Mutex<McpClient>>,
15    pub tools: Vec<McpToolInfo>,
16}
17
18impl McpConnection {
19    pub async fn call_tool(
20        &self,
21        tool_name: &str,
22        arguments: Value,
23    ) -> Result<crate::client::McpToolCallResult, McpError> {
24        let mut client = self.client.lock().await;
25        client.call_tool(tool_name, arguments).await
26    }
27}
28
29/// Manages all configured MCP server connections.
30pub struct McpManager {
31    connections: Vec<Arc<McpConnection>>,
32}
33
34impl McpManager {
35    pub async fn connect_all(config: &McpConfig) -> Result<Self, McpError> {
36        let mut connections = Vec::new();
37        for server in &config.servers {
38            match Self::connect_server(server).await {
39                Ok(conn) => connections.push(conn),
40                Err(err) => {
41                    warn!(server = %server.name, %err, "failed to connect MCP server");
42                }
43            }
44        }
45        Ok(Self { connections })
46    }
47
48    pub async fn connect_from_config() -> Result<Self, McpError> {
49        let config = load_mcp_config().map_err(|err| McpError::Protocol {
50            server: "config".into(),
51            message: err.to_string(),
52        })?;
53        Self::connect_all(&config).await
54    }
55
56    async fn connect_server(server: &McpServer) -> Result<Arc<McpConnection>, McpError> {
57        let mut client = McpClient::connect(server).await?;
58        let tools = client.list_tools().await?;
59        info!(
60            server = %server.name,
61            tools = tools.len(),
62            "connected MCP server"
63        );
64        Ok(Arc::new(McpConnection {
65            server_name: server.name.clone(),
66            client: Arc::new(Mutex::new(client)),
67            tools,
68        }))
69    }
70
71    pub fn connections(&self) -> &[Arc<McpConnection>] {
72        &self.connections
73    }
74
75    pub fn is_empty(&self) -> bool {
76        self.connections.is_empty()
77    }
78
79    pub fn tool_count(&self) -> usize {
80        self.connections.iter().map(|c| c.tools.len()).sum()
81    }
82
83    /// Connect to configured servers; returns `None` when none are available.
84    pub async fn connect_optional() -> Option<Arc<Self>> {
85        match Self::connect_from_config().await {
86            Ok(manager) if !manager.is_empty() => Some(Arc::new(manager)),
87            Ok(_) => None,
88            Err(err) => {
89                warn!(%err, "MCP initialization failed");
90                None
91            }
92        }
93    }
94
95    /// Resolve `mcp_{server}_{tool}` back to connection + original tool name.
96    pub fn resolve_tool(&self, registered_name: &str) -> Option<(Arc<McpConnection>, String)> {
97        for conn in &self.connections {
98            let prefix = format!("mcp_{}_", sanitize_name(&conn.server_name));
99            if let Some(tool_name) = registered_name.strip_prefix(&prefix) {
100                return Some((Arc::clone(conn), tool_name.to_string()));
101            }
102        }
103        None
104    }
105}
106
107/// Build a stable tool name for the LLM registry.
108pub fn registered_tool_name(server_name: &str, tool_name: &str) -> String {
109    format!("mcp_{}_{}", sanitize_name(server_name), tool_name)
110}
111
112fn sanitize_name(name: &str) -> String {
113    name.chars()
114        .map(|c| {
115            if c.is_ascii_alphanumeric() || c == '_' {
116                c
117            } else {
118                '_'
119            }
120        })
121        .collect()
122}
codei_mcp/manager.rs

codei_mcp/
manager.rs
