tycode_core/mcp/
mod.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3use tokio::sync::RwLock;
4
5use crate::module::ContextComponent;
6use crate::module::Module;
7use crate::module::PromptComponent;
8use crate::module::SlashCommand;
9use crate::settings::config::{McpServerConfig, Settings};
10use crate::tools::r#trait::ToolExecutor;
11use tracing::{debug, error, info, warn};
12
13pub mod client;
14pub mod command;
15pub mod tool;
16
17#[cfg(test)]
18mod tests;
19
20use client::McpClient;
21use command::McpSlashCommand;
22use tool::McpTool;
23
24#[derive(Clone)]
25pub struct McpToolDef {
26    pub name: String,
27    pub tool: rmcp::model::Tool,
28    pub server_name: String,
29}
30
31pub(crate) struct McpModuleInner {
32    pub(crate) clients: HashMap<String, McpClient>,
33    pub(crate) tool_defs: Vec<McpToolDef>,
34}
35
36pub struct McpModule {
37    inner: Arc<RwLock<McpModuleInner>>,
38}
39
40impl McpModule {
41    pub fn empty() -> Arc<Self> {
42        Arc::new(Self {
43            inner: Arc::new(RwLock::new(McpModuleInner {
44                clients: HashMap::new(),
45                tool_defs: Vec::new(),
46            })),
47        })
48    }
49
50    pub async fn from_settings(settings: &Settings) -> anyhow::Result<Arc<Self>> {
51        let module = Self::empty();
52
53        for (name, config) in &settings.mcp_servers {
54            info!(server_name = %name, command = %config.command, "Initializing MCP server");
55            if let Err(e) = module.add_server(name.clone(), config.clone()).await {
56                error!(error = ?e, server_name = %name, "Failed to initialize MCP server");
57            }
58        }
59
60        {
61            let inner = module.inner.read().await;
62            info!(
63                servers = inner.clients.len(),
64                tools = inner.tool_defs.len(),
65                "MCP module initialized"
66            );
67        }
68
69        Ok(module)
70    }
71
72    pub async fn add_server(&self, name: String, config: McpServerConfig) -> anyhow::Result<()> {
73        debug!(server_name = %name, "Adding MCP server");
74
75        let mut client = McpClient::new(name.clone(), config).await.map_err(|e| {
76            error!(error = ?e, server_name = %name, "Failed to initialize MCP client");
77            anyhow::anyhow!("Failed to initialize MCP server '{name}': {e:?}")
78        })?;
79
80        let mcp_tools = client.list_tools().await.map_err(|e| {
81            error!(error = ?e, server_name = %name, "Failed to list MCP tools");
82            anyhow::anyhow!("Failed to list tools from MCP server '{name}': {e:?}")
83        })?;
84
85        debug!(server_name = %name, tool_count = mcp_tools.len(), "Found MCP tools");
86
87        let mut inner = self.inner.write().await;
88        for mcp_tool in mcp_tools {
89            inner.tool_defs.push(McpToolDef {
90                name: format!("mcp_{}", mcp_tool.name),
91                tool: mcp_tool,
92                server_name: name.clone(),
93            });
94        }
95        inner.clients.insert(name, client);
96        Ok(())
97    }
98
99    pub async fn remove_server(&self, name: &str) -> anyhow::Result<()> {
100        debug!(server_name = %name, "Removing MCP server");
101        let mut inner = self.inner.write().await;
102
103        if inner.clients.remove(name).is_some() {
104            inner.tool_defs.retain(|def| def.server_name != name);
105        }
106        Ok(())
107    }
108
109    pub fn get_tool_definitions(&self) -> Vec<McpToolDef> {
110        match self.inner.try_read() {
111            Ok(inner) => inner.tool_defs.clone(),
112            Err(_) => {
113                warn!("Failed to acquire read lock for MCP tool definitions");
114                Vec::new()
115            }
116        }
117    }
118}
119
120impl Module for McpModule {
121    fn prompt_components(&self) -> Vec<Arc<dyn PromptComponent>> {
122        Vec::new()
123    }
124
125    fn context_components(&self) -> Vec<Arc<dyn ContextComponent>> {
126        Vec::new()
127    }
128
129    fn tools(&self) -> Vec<Arc<dyn ToolExecutor>> {
130        match self.inner.try_read() {
131            Ok(inner) => inner
132                .tool_defs
133                .iter()
134                .filter_map(|def| McpTool::new(def, self.inner.clone()).ok())
135                .map(|tool| Arc::new(tool) as Arc<dyn ToolExecutor>)
136                .collect(),
137            Err(_) => {
138                warn!("Failed to acquire read lock for MCP tools");
139                Vec::new()
140            }
141        }
142    }
143
144    fn slash_commands(&self) -> Vec<Arc<dyn SlashCommand>> {
145        vec![Arc::new(McpSlashCommand)]
146    }
147}
tycode_core/mcp/mod.rs

tycode_core/mcp/
mod.rs
