Skip to main content

zagens_runtime_adapters/mcp/
config_io.rs

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    /// Whether the tool is enabled in server config (`enabled_tools` / `disabled_tools`).
27    #[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
161/// Load a single server block from `mcp.json` (for edit UIs).
162pub 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
167/// Remove a server entry and persist.
168/// Returns `true` when removed, `false` when the name was already absent (idempotent).
169pub 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
178/// Replace an existing server block (full document for that key). Errors if the name is missing.
179pub 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
195/// Merge servers (and optionally `timeouts`) from a JSON fragment into `mcp.json`.
196///
197/// Supported shapes:
198/// - Full config: `{ "mcpServers"|"servers": { "name": { … } }, "timeouts"?: … }`
199/// - Server map: `{ "filesystem": { "command": "npx", "args": [] }, … }`
200/// - Single server: `{ "name": "fs", "command": "…", "args": [] }` or `"url"`
201///
202/// Existing servers with the same name are replaced. Returns the number of server entries merged.
203pub 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
337/// Build a manager snapshot from a live pool (shared sidecar pool / hot-reload path).
338pub 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}