1use std::collections::HashMap;
2use std::fs;
3use std::path::Path;
4
5use anyhow::{Context, Result};
6
7use crate::network_policy::NetworkPolicyDecider;
8use crate::util::write_atomic;
9
10use super::auth::merge_preserved_secrets;
11use super::config::{McpConfig, McpServerConfig, McpTransportKind};
12use super::pool::McpPool;
13
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub enum McpWriteStatus {
16 Created,
17 Overwritten,
18 SkippedExists,
19}
20
21#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
22pub struct McpDiscoveredItem {
23 pub name: String,
24 pub model_name: String,
25 pub description: Option<String>,
26 #[serde(default = "default_item_enabled")]
28 pub enabled: bool,
29}
30
31fn default_item_enabled() -> bool {
32 true
33}
34
35#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
36pub struct McpServerSnapshot {
37 pub name: String,
38 pub enabled: bool,
39 pub required: bool,
40 pub transport: String,
41 pub command_or_url: String,
42 pub connect_timeout: u64,
43 pub execute_timeout: u64,
44 pub read_timeout: u64,
45 pub connected: bool,
46 pub error: Option<String>,
47 pub tools: Vec<McpDiscoveredItem>,
48 pub resources: Vec<McpDiscoveredItem>,
49 pub prompts: Vec<McpDiscoveredItem>,
50}
51
52#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
53pub struct McpManagerSnapshot {
54 pub config_path: std::path::PathBuf,
55 pub config_exists: bool,
56 pub restart_required: bool,
57 pub servers: Vec<McpServerSnapshot>,
58}
59
60pub fn load_config(path: &Path) -> Result<McpConfig> {
61 if !path.exists() {
62 return Ok(McpConfig::default());
63 }
64 let contents = fs::read_to_string(path)
65 .with_context(|| format!("Failed to read MCP config {}", path.display()))?;
66 serde_json::from_str(&contents)
67 .with_context(|| format!("Failed to parse MCP config {}", path.display()))
68}
69
70pub fn save_config(path: &Path, cfg: &McpConfig) -> Result<()> {
71 if let Some(parent) = path.parent() {
72 fs::create_dir_all(parent).with_context(|| {
73 format!("Failed to create MCP config directory {}", parent.display())
74 })?;
75 }
76 let rendered = serde_json::to_string_pretty(cfg).context("Failed to serialize MCP config")?;
77 write_atomic(path, rendered.as_bytes())
78 .with_context(|| format!("Failed to write MCP config {}", path.display()))?;
79 Ok(())
80}
81
82fn mcp_template_json() -> Result<String> {
83 let mut cfg = McpConfig::default();
84 cfg.servers.insert(
85 "example".to_string(),
86 McpServerConfig {
87 command: Some("node".to_string()),
88 args: vec!["./path/to/your-mcp-server.js".to_string()],
89 env: HashMap::new(),
90 url: None,
91 transport: None,
92 headers: HashMap::new(),
93 auth: None,
94 connect_timeout: None,
95 execute_timeout: None,
96 read_timeout: None,
97 disabled: true,
98 enabled: true,
99 required: false,
100 enabled_tools: Vec::new(),
101 disabled_tools: Vec::new(),
102 },
103 );
104 serde_json::to_string_pretty(&cfg).context("Failed to render MCP template JSON")
105}
106
107pub fn init_config(path: &Path, force: bool) -> Result<McpWriteStatus> {
108 if path.exists() && !force {
109 return Ok(McpWriteStatus::SkippedExists);
110 }
111 let status = if path.exists() {
112 McpWriteStatus::Overwritten
113 } else {
114 McpWriteStatus::Created
115 };
116 if let Some(parent) = path.parent() {
117 fs::create_dir_all(parent).with_context(|| {
118 format!("Failed to create MCP config directory {}", parent.display())
119 })?;
120 }
121 let template = mcp_template_json()?;
122 write_atomic(path, template.as_bytes())
123 .with_context(|| format!("Failed to write MCP config {}", path.display()))?;
124 Ok(status)
125}
126
127pub fn add_server_config(
128 path: &Path,
129 name: String,
130 command: Option<String>,
131 url: Option<String>,
132 args: Vec<String>,
133) -> Result<()> {
134 if command.is_none() && url.is_none() {
135 anyhow::bail!("Provide either a command or URL for MCP server '{name}'.");
136 }
137 let mut cfg = load_config(path)?;
138 cfg.servers.insert(
139 name,
140 McpServerConfig {
141 command,
142 args,
143 env: HashMap::new(),
144 url,
145 transport: None,
146 headers: HashMap::new(),
147 auth: None,
148 connect_timeout: None,
149 execute_timeout: None,
150 read_timeout: None,
151 disabled: false,
152 enabled: true,
153 required: false,
154 enabled_tools: Vec::new(),
155 disabled_tools: Vec::new(),
156 },
157 );
158 save_config(path, &cfg)
159}
160
161pub fn get_server_entry(path: &Path, name: &str) -> Result<Option<McpServerConfig>> {
163 let cfg = load_config(path)?;
164 Ok(cfg.servers.get(name).cloned())
165}
166
167pub fn remove_server_from_config(path: &Path, name: &str) -> Result<bool> {
170 let mut cfg = load_config(path)?;
171 if cfg.servers.remove(name).is_none() {
172 return Ok(false);
173 }
174 save_config(path, &cfg)?;
175 Ok(true)
176}
177
178pub fn replace_server_in_config(path: &Path, name: &str, server: McpServerConfig) -> Result<()> {
180 if server.command.is_none() && server.url.is_none() {
181 anyhow::bail!("MCP server '{name}': provide either `command` or `url`");
182 }
183 let mut cfg = load_config(path)?;
184 let old = cfg
185 .servers
186 .get(name)
187 .cloned()
188 .ok_or_else(|| anyhow::anyhow!("MCP server '{name}' is not configured"))?;
189 let mut server = server;
190 merge_preserved_secrets(&mut server, &old);
191 cfg.servers.insert(name.to_string(), server);
192 save_config(path, &cfg)
193}
194
195pub fn merge_mcp_json_fragment(path: &Path, fragment: &str) -> Result<usize> {
204 let fragment = fragment.trim();
205 if fragment.is_empty() {
206 anyhow::bail!("JSON 不能为空");
207 }
208 let v: serde_json::Value =
209 serde_json::from_str(fragment).context("无效的 JSON:请检查语法(逗号、引号、括号)")?;
210 let obj = v.as_object().context("根节点必须是 JSON 对象 { … }")?;
211
212 let mut cfg = load_config(path)?;
213 let mut merged = 0usize;
214
215 if obj.contains_key("mcpServers") || obj.contains_key("servers") {
216 let partial: McpConfig = serde_json::from_value(v.clone())
217 .context("无法解析为 MCP 配置(检查 mcpServers / servers 字段)")?;
218 if obj.contains_key("timeouts") {
219 cfg.timeouts = partial.timeouts;
220 }
221 if partial.servers.is_empty() && !obj.contains_key("timeouts") {
222 anyhow::bail!("servers 为空:请至少包含一个服务器条目,或同时提供 timeouts");
223 }
224 for (name, sc) in partial.servers {
225 cfg.servers.insert(name, sc);
226 merged += 1;
227 }
228 save_config(path, &cfg)?;
229 return Ok(merged);
230 }
231
232 if obj.contains_key("name") {
233 let name = obj
234 .get("name")
235 .and_then(|x| x.as_str())
236 .context("name 必须是字符串")?;
237 if name.trim().is_empty() {
238 anyhow::bail!("name 不能为空");
239 }
240 let mut inner = obj.clone();
241 inner.remove("name");
242 let server: McpServerConfig = serde_json::from_value(serde_json::Value::Object(inner))
243 .context("服务器字段无效(可与 ~/.zagens/mcp.json 中条目对照)")?;
244 if server.command.is_none() && server.url.is_none() {
245 anyhow::bail!("必须提供 command 或 url");
246 }
247 cfg.servers.insert(name.trim().to_string(), server);
248 merged = 1;
249 save_config(path, &cfg)?;
250 return Ok(merged);
251 }
252
253 let mut timeout_updated = false;
254 let mut incoming: HashMap<String, McpServerConfig> = HashMap::new();
255 for (key, val) in obj {
256 if key == "timeouts" {
257 cfg.timeouts = serde_json::from_value(val.clone()).context("timeouts 格式无效")?;
258 timeout_updated = true;
259 continue;
260 }
261 let server: McpServerConfig = serde_json::from_value(val.clone())
262 .with_context(|| format!("服务器 \"{key}\" 的配置无效"))?;
263 incoming.insert(key.clone(), server);
264 }
265
266 if incoming.is_empty() && !timeout_updated {
267 anyhow::bail!(
268 "未找到服务器条目。可粘贴完整 mcpServers 块,或形如 {{ \"myserver\": {{ \"command\": \"npx\", \"args\": [] }} }}"
269 );
270 }
271 for (k, s) in incoming {
272 cfg.servers.insert(k, s);
273 merged += 1;
274 }
275 save_config(path, &cfg)?;
276 Ok(merged)
277}
278
279pub fn remove_server_config(path: &Path, name: &str) -> Result<()> {
280 let mut cfg = load_config(path)?;
281 if cfg.servers.remove(name).is_none() {
282 anyhow::bail!("MCP server '{name}' not found");
283 }
284 save_config(path, &cfg)
285}
286
287pub fn set_server_enabled(path: &Path, name: &str, enabled: bool) -> Result<()> {
288 let mut cfg = load_config(path)?;
289 let server = cfg
290 .servers
291 .get_mut(name)
292 .ok_or_else(|| anyhow::anyhow!("MCP server '{name}' not found"))?;
293 server.enabled = enabled;
294 server.disabled = !enabled;
295 save_config(path, &cfg)
296}
297
298pub fn manager_snapshot_from_config(
299 path: &Path,
300 restart_required: bool,
301) -> Result<McpManagerSnapshot> {
302 let cfg = load_config(path)?;
303 Ok(snapshot_from_config(
304 path,
305 path.exists(),
306 restart_required,
307 &cfg,
308 None,
309 ))
310}
311
312pub async fn discover_manager_snapshot(
313 path: &Path,
314 network_policy: Option<NetworkPolicyDecider>,
315 restart_required: bool,
316) -> Result<McpManagerSnapshot> {
317 let cfg = load_config(path)?;
318 let mut pool = McpPool::new(cfg.clone());
319 if let Some(policy) = network_policy {
320 pool = pool.with_network_policy(policy);
321 }
322 let errors = pool
323 .connect_all()
324 .await
325 .into_iter()
326 .map(|(name, err)| (name, err.to_string()))
327 .collect::<HashMap<_, _>>();
328 Ok(snapshot_from_config(
329 path,
330 path.exists(),
331 restart_required,
332 &cfg,
333 Some((&pool, &errors)),
334 ))
335}
336
337pub async fn manager_snapshot_from_pool(path: &Path, pool: &mut McpPool) -> McpManagerSnapshot {
339 let cfg = load_config(path).unwrap_or_default();
340 let errors: HashMap<String, String> = pool
341 .connect_all()
342 .await
343 .into_iter()
344 .map(|(name, err)| (name, err.to_string()))
345 .collect();
346 snapshot_from_config(path, path.exists(), false, &cfg, Some((pool, &errors)))
347}
348
349fn snapshot_from_config(
350 path: &Path,
351 config_exists: bool,
352 restart_required: bool,
353 cfg: &McpConfig,
354 discovery: Option<(&McpPool, &HashMap<String, String>)>,
355) -> McpManagerSnapshot {
356 let mut servers = cfg
357 .servers
358 .iter()
359 .map(|(name, server)| {
360 let transport = server
361 .transport_kind()
362 .map_or("unknown", McpTransportKind::as_str);
363 let command_or_url = server.url.clone().unwrap_or_else(|| {
364 let mut command = server
365 .command
366 .clone()
367 .unwrap_or_else(|| "(missing)".to_string());
368 if !server.args.is_empty() {
369 command.push(' ');
370 command.push_str(&server.args.join(" "));
371 }
372 command
373 });
374 let mut snapshot = McpServerSnapshot {
375 name: name.clone(),
376 enabled: server.is_enabled(),
377 required: server.required,
378 transport: transport.to_string(),
379 command_or_url,
380 connect_timeout: server.effective_connect_timeout(&cfg.timeouts),
381 execute_timeout: server.effective_execute_timeout(&cfg.timeouts),
382 read_timeout: server.effective_read_timeout(&cfg.timeouts),
383 connected: false,
384 error: if server.is_enabled() {
385 None
386 } else {
387 Some("disabled".to_string())
388 },
389 tools: Vec::new(),
390 resources: Vec::new(),
391 prompts: Vec::new(),
392 };
393
394 if let Some((pool, errors)) = discovery {
395 if let Some(error) = errors.get(name) {
396 snapshot.error = Some(error.clone());
397 }
398 if let Some(conn) = pool.connections.get(name) {
399 snapshot.connected = conn.is_ready();
400 snapshot.tools = conn
401 .tools()
402 .iter()
403 .map(|tool| McpDiscoveredItem {
404 name: tool.name.clone(),
405 model_name: format!("mcp_{}_{}", name, tool.name),
406 description: tool.description.clone(),
407 enabled: conn.config().is_tool_enabled(&tool.name),
408 })
409 .collect();
410 snapshot.resources =
411 conn.resources()
412 .iter()
413 .map(|resource| McpDiscoveredItem {
414 name: resource.name.clone(),
415 model_name: format!(
416 "mcp_{}_{}",
417 name,
418 resource.name.replace(' ', "_").to_lowercase()
419 ),
420 description: resource.description.clone(),
421 enabled: true,
422 })
423 .chain(conn.resource_templates().iter().map(|template| {
424 McpDiscoveredItem {
425 name: template.name.clone(),
426 model_name: format!(
427 "mcp_{}_{}",
428 name,
429 template.name.replace(' ', "_").to_lowercase()
430 ),
431 description: template.description.clone(),
432 enabled: true,
433 }
434 }))
435 .collect();
436 snapshot.prompts = conn
437 .prompts()
438 .iter()
439 .map(|prompt| McpDiscoveredItem {
440 name: prompt.name.clone(),
441 model_name: format!("mcp_{}_{}", name, prompt.name),
442 description: prompt.description.clone(),
443 enabled: true,
444 })
445 .collect();
446 }
447 }
448
449 snapshot
450 })
451 .collect::<Vec<_>>();
452 servers.sort_by(|a, b| a.name.cmp(&b.name));
453 McpManagerSnapshot {
454 config_path: path.to_path_buf(),
455 config_exists,
456 restart_required,
457 servers,
458 }
459}