Skip to main content

tycode_core/mcp/
command.rs

1use std::collections::HashMap;
2
3use chrono::Utc;
4
5use crate::chat::actor::ActorState;
6use crate::chat::events::{ChatMessage, MessageSender};
7use crate::module::SlashCommand;
8use crate::settings::config::McpServerConfig;
9
10pub struct McpSlashCommand;
11
12#[async_trait::async_trait(?Send)]
13impl SlashCommand for McpSlashCommand {
14    fn name(&self) -> &'static str {
15        "mcp"
16    }
17
18    fn description(&self) -> &'static str {
19        "Manage MCP server configurations"
20    }
21
22    fn usage(&self) -> &'static str {
23        "/mcp [add|remove] [args...]"
24    }
25
26    async fn execute(&self, state: &mut ActorState, args: &[&str]) -> Vec<ChatMessage> {
27        let parts: Vec<String> = std::iter::once("mcp".to_string())
28            .chain(args.iter().map(|s| s.to_string()))
29            .collect();
30        handle_mcp_command(state, &parts).await
31    }
32}
33
34fn create_message(content: String, sender: MessageSender) -> ChatMessage {
35    ChatMessage {
36        content,
37        sender,
38        timestamp: Utc::now().timestamp_millis() as u64,
39        reasoning: None,
40        tool_calls: Vec::new(),
41        model_info: None,
42        token_usage: None,
43        images: vec![],
44    }
45}
46
47async fn handle_mcp_command(state: &mut ActorState, parts: &[String]) -> Vec<ChatMessage> {
48    if parts.len() < 2 {
49        let settings = state.settings.settings();
50        if settings.mcp_servers.is_empty() {
51            return vec![create_message(
52                "No MCP servers configured. Use `/mcp add <name> <command> [--args \"args...\"] [--env \"KEY=VALUE\"]` to add one.".to_string(),
53                MessageSender::System,
54            )];
55        }
56
57        let mut message = String::from("Configured MCP servers:\n\n");
58        for (name, config) in &settings.mcp_servers {
59            message.push_str(&format!(
60                "  {}:\n    Command: {}\n    Args: {}\n    Env: {}\n\n",
61                name,
62                config.command,
63                if config.args.is_empty() {
64                    "<none>".to_string()
65                } else {
66                    config.args.join(" ")
67                },
68                if config.env.is_empty() {
69                    "<none>".to_string()
70                } else {
71                    config
72                        .env
73                        .iter()
74                        .map(|(k, v)| format!("{}={}", k, v))
75                        .collect::<Vec<_>>()
76                        .join(", ")
77                }
78            ));
79        }
80        return vec![create_message(message, MessageSender::System)];
81    }
82
83    match parts[1].as_str() {
84        "add" => handle_mcp_add_command(state, parts).await,
85        "remove" => handle_mcp_remove_command(state, parts).await,
86        _ => vec![create_message(
87            "Usage: /mcp [add|remove] [args...]. Use `/mcp` to list all servers.".to_string(),
88            MessageSender::Error,
89        )],
90    }
91}
92
93fn parse_mcp_args_value(parts: &[String], i: usize) -> Result<Vec<String>, String> {
94    let args_str = parts.get(i + 1).ok_or("--args requires a value")?;
95    Ok(args_str.split_whitespace().map(|s| s.to_string()).collect())
96}
97
98fn parse_mcp_env_var(parts: &[String], i: usize) -> Result<(String, String), String> {
99    let env_str = parts
100        .get(i + 1)
101        .ok_or("--env requires a value in format KEY=VALUE")?;
102    let eq_pos = env_str
103        .find('=')
104        .ok_or("Environment variable must be in format KEY=VALUE")?;
105    let key = env_str[..eq_pos].to_string();
106    if key.is_empty() {
107        return Err("Environment variable key cannot be empty".to_string());
108    }
109    let value = env_str[eq_pos + 1..].to_string();
110    Ok((key, value))
111}
112
113fn process_mcp_optional_args(parts: &[String], config: &mut McpServerConfig) -> Result<(), String> {
114    let mut i = 4;
115    while i < parts.len() {
116        match parts[i].as_str() {
117            "--args" => {
118                config.args = parse_mcp_args_value(parts, i)?;
119                i += 2;
120            }
121            "--env" => {
122                let (key, value) = parse_mcp_env_var(parts, i)?;
123                config.env.insert(key, value);
124                i += 2;
125            }
126            arg => return Err(format!("Unknown argument: {}", arg)),
127        }
128    }
129    Ok(())
130}
131
132async fn handle_mcp_add_command(state: &mut ActorState, parts: &[String]) -> Vec<ChatMessage> {
133    if parts.len() < 4 {
134        return vec![create_message(
135            "Usage: /mcp add <name> <command> [--args \"args...\"] [--env \"KEY=VALUE\"]"
136                .to_string(),
137            MessageSender::Error,
138        )];
139    }
140
141    let name = parts[2].trim().to_string();
142    let command = parts[3].trim().to_string();
143
144    if name.is_empty() {
145        return vec![create_message(
146            "Server name cannot be empty".to_string(),
147            MessageSender::Error,
148        )];
149    }
150
151    if command.is_empty() {
152        return vec![create_message(
153            "Command path cannot be empty".to_string(),
154            MessageSender::Error,
155        )];
156    }
157
158    let mut config = McpServerConfig {
159        command,
160        args: Vec::new(),
161        env: HashMap::new(),
162    };
163
164    if let Err(e) = process_mcp_optional_args(parts, &mut config) {
165        return vec![create_message(e, MessageSender::Error)];
166    }
167
168    let current_settings = state.settings.settings();
169    let replacing = current_settings.mcp_servers.contains_key(&name);
170
171    state.settings.update_setting(|settings| {
172        settings.mcp_servers.insert(name.clone(), config.clone());
173    });
174
175    if let Err(e) = state.settings.save() {
176        return vec![create_message(
177            format!("MCP server updated for this session but failed to save settings: {e:?}"),
178            MessageSender::Error,
179        )];
180    }
181
182    let connection_status = match state.mcp_manager.add_server(name.clone(), config).await {
183        Ok(()) => "\nServer connected successfully.".to_string(),
184        Err(e) => format!(
185            "\nWarning: Failed to connect to server: {e:?}. Server will be retried on next session."
186        ),
187    };
188
189    let response = if replacing {
190        format!("Updated MCP server '{name}'")
191    } else {
192        format!("Added MCP server '{name}'")
193    };
194
195    vec![create_message(
196        format!(
197            "{}{}\n\nSettings saved to disk. The MCP server configuration is now persistent across sessions.",
198            response, connection_status
199        ),
200        MessageSender::System,
201    )]
202}
203
204async fn handle_mcp_remove_command(state: &mut ActorState, parts: &[String]) -> Vec<ChatMessage> {
205    if parts.len() < 3 {
206        return vec![create_message(
207            "Usage: /mcp remove <name>".to_string(),
208            MessageSender::Error,
209        )];
210    }
211
212    let name = parts[2].trim();
213
214    if name.is_empty() {
215        return vec![create_message(
216            "Server name cannot be empty".to_string(),
217            MessageSender::Error,
218        )];
219    }
220
221    let current_settings = state.settings.settings();
222    if !current_settings.mcp_servers.contains_key(name) {
223        return vec![create_message(
224            format!("MCP server '{name}' not found"),
225            MessageSender::Error,
226        )];
227    }
228
229    state.settings.update_setting(|settings| {
230        settings.mcp_servers.remove(name);
231    });
232
233    if let Err(e) = state.settings.save() {
234        return vec![create_message(
235            format!("MCP server removed for this session but failed to save settings: {e:?}"),
236            MessageSender::Error,
237        )];
238    }
239
240    vec![create_message(
241        format!(
242            "Removed MCP server '{name}'\n\nSettings saved to disk. The MCP server configuration is now persistent across sessions."
243        ),
244        MessageSender::System,
245    )]
246}