Skip to main content

zag_agent/
mcp.rs

1/// Provider-agnostic MCP (Model Context Protocol) server management.
2///
3/// MCP server configs are stored as individual TOML files:
4/// - Global: `~/.zag/mcp/<server-name>.toml`
5/// - Project-scoped: `~/.zag/projects/<sanitized-path>/mcp/<server-name>.toml`
6///
7/// During sync, servers are injected into each provider's native config format
8/// with a `zag-` prefix to avoid collisions with user-managed servers.
9///
10/// Supported providers:
11/// - Claude: `~/.claude.json` under `mcpServers` (JSON)
12/// - Gemini: `~/.gemini/settings.json` under `mcpServers` (JSON)
13/// - Copilot: `~/.copilot/mcp-config.json` under `mcpServers` (JSON)
14/// - Codex: `~/.codex/config.toml` under `[mcp_servers]` (TOML)
15/// - Ollama: No native MCP support
16#[cfg(test)]
17#[path = "mcp_tests.rs"]
18mod tests;
19
20use anyhow::{Context, Result, bail};
21use serde::{Deserialize, Serialize};
22use std::collections::BTreeMap;
23use std::fs;
24use std::path::{Path, PathBuf};
25
26const MCP_PREFIX: &str = "zag-";
27
28// ---------------------------------------------------------------------------
29// Data structures
30// ---------------------------------------------------------------------------
31
32/// An MCP server configuration (one per TOML file).
33#[derive(Debug, Clone, Serialize, Deserialize)]
34pub struct McpServer {
35    /// Human-readable name (also the filename stem).
36    pub name: String,
37    /// Optional description.
38    #[serde(default)]
39    pub description: String,
40    /// Transport type: "stdio" or "http".
41    #[serde(default = "default_transport")]
42    pub transport: String,
43
44    // -- stdio fields --
45    /// Command to start the server (stdio transport).
46    #[serde(default)]
47    pub command: Option<String>,
48    /// Arguments for the command.
49    #[serde(default)]
50    pub args: Vec<String>,
51
52    // -- http fields --
53    /// URL endpoint (http transport).
54    #[serde(default)]
55    pub url: Option<String>,
56    /// Environment variable name containing a bearer token (http transport).
57    #[serde(default)]
58    pub bearer_token_env_var: Option<String>,
59    /// HTTP headers for http transport.
60    #[serde(default)]
61    pub headers: BTreeMap<String, String>,
62
63    // -- shared --
64    /// Environment variables forwarded to the server process.
65    #[serde(default)]
66    pub env: BTreeMap<String, String>,
67}
68
69fn default_transport() -> String {
70    "stdio".to_string()
71}
72
73// ---------------------------------------------------------------------------
74// Directory helpers
75// ---------------------------------------------------------------------------
76
77/// Returns `~/.zag/mcp/`.
78pub fn mcp_dir() -> PathBuf {
79    dirs::home_dir()
80        .unwrap_or_else(|| PathBuf::from("."))
81        .join(".zag")
82        .join("mcp")
83}
84
85/// Returns the project-scoped MCP directory for the given root.
86/// Uses the same sanitization as config: `~/.zag/projects/<sanitized-path>/mcp/`.
87pub fn project_mcp_dir(root: Option<&str>) -> Option<PathBuf> {
88    let base = dirs::home_dir()?.join(".zag");
89
90    let project_dir = if let Some(r) = root {
91        let sanitized = crate::config::Config::sanitize_path(r);
92        base.join("projects").join(sanitized)
93    } else {
94        let current_dir = std::env::current_dir().ok()?;
95        let git_root = find_git_root(&current_dir)?;
96        let sanitized = crate::config::Config::sanitize_path(&git_root.to_string_lossy());
97        base.join("projects").join(sanitized)
98    };
99
100    Some(project_dir.join("mcp"))
101}
102
103fn find_git_root(start_dir: &Path) -> Option<PathBuf> {
104    let output = std::process::Command::new("git")
105        .arg("rev-parse")
106        .arg("--show-toplevel")
107        .current_dir(start_dir)
108        .output()
109        .ok()?;
110    if output.status.success() {
111        let root = String::from_utf8(output.stdout).ok()?;
112        Some(PathBuf::from(root.trim()))
113    } else {
114        None
115    }
116}
117
118/// Returns the provider's native MCP config file path, or `None` if unsupported.
119///
120/// - Claude: `~/.claude.json`
121/// - Gemini: `~/.gemini/settings.json`
122/// - Copilot: `~/.copilot/mcp-config.json`
123/// - Codex: `~/.codex/config.toml`
124pub fn provider_mcp_config_path(provider: &str) -> Option<PathBuf> {
125    let home = dirs::home_dir()?;
126    match provider {
127        "claude" => Some(home.join(".claude.json")),
128        "gemini" => Some(home.join(".gemini").join("settings.json")),
129        "copilot" => Some(home.join(".copilot").join("mcp-config.json")),
130        "codex" => Some(home.join(".codex").join("config.toml")),
131        _ => None,
132    }
133}
134
135/// List of providers that support MCP.
136pub const MCP_PROVIDERS: &[&str] = &["claude", "gemini", "copilot", "codex"];
137
138// ---------------------------------------------------------------------------
139// Loading / saving individual servers
140// ---------------------------------------------------------------------------
141
142/// Parse an MCP server from a TOML file.
143pub fn parse_server(path: &Path) -> Result<McpServer> {
144    let content =
145        fs::read_to_string(path).with_context(|| format!("Failed to read {}", path.display()))?;
146    let server: McpServer = toml::from_str(&content)
147        .with_context(|| format!("Failed to parse MCP server config {}", path.display()))?;
148    Ok(server)
149}
150
151/// Load all MCP servers from a directory. Silently skips invalid files.
152fn load_servers_from(dir: &Path) -> Result<Vec<McpServer>> {
153    if !dir.exists() {
154        return Ok(Vec::new());
155    }
156
157    let mut servers = Vec::new();
158    for entry in fs::read_dir(dir)
159        .with_context(|| format!("Failed to read MCP directory {}", dir.display()))?
160    {
161        let entry = entry?;
162        let path = entry.path();
163        if path.extension().and_then(|e| e.to_str()) != Some("toml") {
164            continue;
165        }
166        match parse_server(&path) {
167            Ok(server) => servers.push(server),
168            Err(e) => {
169                log::warn!("Skipping MCP server at {}: {}", path.display(), e);
170            }
171        }
172    }
173    servers.sort_by(|a, b| a.name.cmp(&b.name));
174    Ok(servers)
175}
176
177/// Load all global MCP servers from `~/.zag/mcp/`.
178pub fn load_global_servers() -> Result<Vec<McpServer>> {
179    load_servers_from(&mcp_dir())
180}
181
182/// Load project-scoped MCP servers. Returns empty vec if not in a project.
183pub fn load_project_servers(root: Option<&str>) -> Result<Vec<McpServer>> {
184    match project_mcp_dir(root) {
185        Some(dir) => load_servers_from(&dir),
186        None => Ok(Vec::new()),
187    }
188}
189
190/// Load all MCP servers (global + project-scoped, project overrides global).
191pub fn load_all_servers(root: Option<&str>) -> Result<Vec<McpServer>> {
192    let mut by_name: BTreeMap<String, McpServer> = BTreeMap::new();
193
194    for server in load_global_servers()? {
195        by_name.insert(server.name.clone(), server);
196    }
197    for server in load_project_servers(root)? {
198        by_name.insert(server.name.clone(), server);
199    }
200
201    Ok(by_name.into_values().collect())
202}
203
204/// List all MCP servers (alias for load_all_servers).
205pub fn list_servers(root: Option<&str>) -> Result<Vec<McpServer>> {
206    load_all_servers(root)
207}
208
209/// Get a single MCP server by name. Checks project-scoped first, then global.
210pub fn get_server(name: &str, root: Option<&str>) -> Result<McpServer> {
211    // Check project-scoped first
212    if let Some(dir) = project_mcp_dir(root) {
213        let path = dir.join(format!("{name}.toml"));
214        if path.exists() {
215            return parse_server(&path);
216        }
217    }
218    // Check global
219    let path = mcp_dir().join(format!("{name}.toml"));
220    if path.exists() {
221        return parse_server(&path);
222    }
223    bail!("MCP server '{name}' not found");
224}
225
226/// Create a new MCP server config file. Returns the path to the new file.
227/// If `project` is true, creates in project-scoped dir; otherwise global.
228pub fn add_server(server: &McpServer, project: bool, root: Option<&str>) -> Result<PathBuf> {
229    let dir = if project {
230        project_mcp_dir(root).context("Not in a project (no git root found)")?
231    } else {
232        mcp_dir()
233    };
234    fs::create_dir_all(&dir)
235        .with_context(|| format!("Failed to create MCP directory {}", dir.display()))?;
236
237    let path = dir.join(format!("{}.toml", server.name));
238    if path.exists() {
239        bail!(
240            "MCP server '{}' already exists at {}",
241            server.name,
242            path.display()
243        );
244    }
245
246    let content =
247        toml::to_string_pretty(server).context("Failed to serialize MCP server config")?;
248    fs::write(&path, content).with_context(|| format!("Failed to write {}", path.display()))?;
249    Ok(path)
250}
251
252/// Remove an MCP server config file and clean up provider configs.
253pub fn remove_server(name: &str, root: Option<&str>) -> Result<()> {
254    let mut found = false;
255
256    // Try project-scoped first
257    if let Some(dir) = project_mcp_dir(root) {
258        let path = dir.join(format!("{name}.toml"));
259        if path.exists() {
260            fs::remove_file(&path)
261                .with_context(|| format!("Failed to remove {}", path.display()))?;
262            found = true;
263        }
264    }
265
266    // Try global
267    let path = mcp_dir().join(format!("{name}.toml"));
268    if path.exists() {
269        fs::remove_file(&path).with_context(|| format!("Failed to remove {}", path.display()))?;
270        found = true;
271    }
272
273    if !found {
274        bail!("MCP server '{name}' not found");
275    }
276
277    // Remove from all provider configs
278    for provider in MCP_PROVIDERS {
279        if let Err(e) = remove_server_from_provider(provider, name) {
280            log::warn!("Failed to clean up {provider} config for '{name}': {e}");
281        }
282    }
283
284    Ok(())
285}
286
287// ---------------------------------------------------------------------------
288// Provider sync: convert McpServer → provider-native format
289// ---------------------------------------------------------------------------
290
291/// Convert an MCP server to a Claude/Gemini/Copilot JSON entry (serde_json::Value).
292fn server_to_json(server: &McpServer, provider: &str) -> serde_json::Value {
293    let mut entry = serde_json::Map::new();
294
295    if server.transport == "stdio" {
296        if let Some(ref cmd) = server.command {
297            // Copilot uses "type": "local", Claude uses "type": "stdio", Gemini omits type
298            match provider {
299                "copilot" => {
300                    entry.insert("type".into(), serde_json::json!("local"));
301                }
302                "claude" => {
303                    entry.insert("type".into(), serde_json::json!("stdio"));
304                }
305                _ => {}
306            }
307            entry.insert("command".into(), serde_json::json!(cmd));
308            if !server.args.is_empty() {
309                entry.insert("args".into(), serde_json::json!(server.args));
310            }
311        }
312    } else if server.transport == "http" {
313        if let Some(ref url) = server.url {
314            match provider {
315                "copilot" => {
316                    entry.insert("type".into(), serde_json::json!("http"));
317                    entry.insert("url".into(), serde_json::json!(url));
318                }
319                "gemini" => {
320                    entry.insert("httpUrl".into(), serde_json::json!(url));
321                }
322                _ => {
323                    entry.insert("type".into(), serde_json::json!("http"));
324                    entry.insert("url".into(), serde_json::json!(url));
325                }
326            }
327        }
328        if !server.headers.is_empty() {
329            entry.insert("headers".into(), serde_json::json!(server.headers));
330        }
331    }
332
333    if !server.env.is_empty() {
334        entry.insert("env".into(), serde_json::json!(server.env));
335    }
336
337    serde_json::Value::Object(entry)
338}
339
340/// Sync all MCP servers into a JSON-based provider's config file.
341/// Only touches entries with the `zag-` prefix.
342fn sync_json_provider(provider: &str, servers: &[McpServer]) -> Result<usize> {
343    let Some(config_path) = provider_mcp_config_path(provider) else {
344        return Ok(0);
345    };
346
347    // Read existing config (or start with empty object)
348    let mut config: serde_json::Value = if config_path.exists() {
349        let content = fs::read_to_string(&config_path)
350            .with_context(|| format!("Failed to read {}", config_path.display()))?;
351        serde_json::from_str(&content)
352            .with_context(|| format!("Failed to parse {}", config_path.display()))?
353    } else {
354        serde_json::json!({})
355    };
356
357    // Ensure mcpServers object exists
358    let mcp_servers = config
359        .as_object_mut()
360        .context("Config is not a JSON object")?
361        .entry("mcpServers")
362        .or_insert_with(|| serde_json::json!({}));
363
364    let mcp_map = mcp_servers
365        .as_object_mut()
366        .context("mcpServers is not a JSON object")?;
367
368    // Remove all existing zag- entries
369    let zag_keys: Vec<String> = mcp_map
370        .keys()
371        .filter(|k| k.starts_with(MCP_PREFIX))
372        .cloned()
373        .collect();
374    for key in &zag_keys {
375        mcp_map.remove(key);
376    }
377
378    // Add current servers with zag- prefix
379    let mut synced = 0;
380    for server in servers {
381        let key = format!("{}{}", MCP_PREFIX, server.name);
382        let value = server_to_json(server, provider);
383        mcp_map.insert(key, value);
384        synced += 1;
385    }
386
387    // Ensure parent directory exists
388    if let Some(parent) = config_path.parent() {
389        fs::create_dir_all(parent)?;
390    }
391
392    // Write back with pretty printing
393    let content = serde_json::to_string_pretty(&config)?;
394    fs::write(&config_path, format!("{content}\n"))
395        .with_context(|| format!("Failed to write {}", config_path.display()))?;
396
397    log::debug!(
398        "Synced {} MCP server(s) to {} at {}",
399        synced,
400        provider,
401        config_path.display()
402    );
403
404    Ok(synced)
405}
406
407/// Sync all MCP servers into Codex's TOML config.
408/// Only touches entries under `[mcp_servers]` with the `zag-` prefix.
409fn sync_codex_provider(servers: &[McpServer]) -> Result<usize> {
410    let Some(config_path) = provider_mcp_config_path("codex") else {
411        return Ok(0);
412    };
413
414    // Read existing config as a TOML table
415    let mut config: toml::Table = if config_path.exists() {
416        let content = fs::read_to_string(&config_path)
417            .with_context(|| format!("Failed to read {}", config_path.display()))?;
418        content
419            .parse::<toml::Table>()
420            .with_context(|| format!("Failed to parse {}", config_path.display()))?
421    } else {
422        toml::Table::new()
423    };
424
425    // Get or create mcp_servers table
426    let mcp_table = config
427        .entry("mcp_servers")
428        .or_insert_with(|| toml::Value::Table(toml::Table::new()))
429        .as_table_mut()
430        .context("mcp_servers is not a TOML table")?;
431
432    // Remove all existing zag- entries
433    let zag_keys: Vec<String> = mcp_table
434        .keys()
435        .filter(|k| k.starts_with(MCP_PREFIX))
436        .cloned()
437        .collect();
438    for key in &zag_keys {
439        mcp_table.remove(key.as_str());
440    }
441
442    // Add current servers with zag- prefix
443    let mut synced = 0;
444    for server in servers {
445        let key = format!("{}{}", MCP_PREFIX, server.name);
446        let mut entry = toml::Table::new();
447
448        if server.transport == "stdio" {
449            if let Some(ref cmd) = server.command {
450                entry.insert("command".into(), toml::Value::String(cmd.clone()));
451            }
452            if !server.args.is_empty() {
453                let args: Vec<toml::Value> = server
454                    .args
455                    .iter()
456                    .map(|a| toml::Value::String(a.clone()))
457                    .collect();
458                entry.insert("args".into(), toml::Value::Array(args));
459            }
460        } else if server.transport == "http" {
461            if let Some(ref url) = server.url {
462                entry.insert("url".into(), toml::Value::String(url.clone()));
463            }
464            if let Some(ref token_var) = server.bearer_token_env_var {
465                entry.insert(
466                    "bearer_token_env_var".into(),
467                    toml::Value::String(token_var.clone()),
468                );
469            }
470        }
471
472        if !server.env.is_empty() {
473            let mut env_table = toml::Table::new();
474            for (k, v) in &server.env {
475                env_table.insert(k.clone(), toml::Value::String(v.clone()));
476            }
477            entry.insert("env".into(), toml::Value::Table(env_table));
478        }
479
480        mcp_table.insert(key, toml::Value::Table(entry));
481        synced += 1;
482    }
483
484    // Ensure parent directory exists
485    if let Some(parent) = config_path.parent() {
486        fs::create_dir_all(parent)?;
487    }
488
489    let content = toml::to_string_pretty(&config)?;
490    fs::write(&config_path, &content)
491        .with_context(|| format!("Failed to write {}", config_path.display()))?;
492
493    log::debug!(
494        "Synced {} MCP server(s) to codex at {}",
495        synced,
496        config_path.display()
497    );
498
499    Ok(synced)
500}
501
502/// Sync MCP servers for a specific provider.
503pub fn sync_servers_for_provider(provider: &str, servers: &[McpServer]) -> Result<usize> {
504    match provider {
505        "claude" | "gemini" | "copilot" => sync_json_provider(provider, servers),
506        "codex" => sync_codex_provider(servers),
507        _ => {
508            log::debug!("Provider '{provider}' does not support MCP servers");
509            Ok(0)
510        }
511    }
512}
513
514/// Remove a single zag-managed server from a provider's config.
515fn remove_server_from_provider(provider: &str, name: &str) -> Result<()> {
516    let Some(config_path) = provider_mcp_config_path(provider) else {
517        return Ok(());
518    };
519    if !config_path.exists() {
520        return Ok(());
521    }
522
523    let key = format!("{MCP_PREFIX}{name}");
524
525    if provider == "codex" {
526        let content = fs::read_to_string(&config_path)?;
527        let mut config: toml::Table = content.parse()?;
528        if let Some(mcp) = config.get_mut("mcp_servers").and_then(|v| v.as_table_mut()) {
529            mcp.remove(&key);
530        }
531        let content = toml::to_string_pretty(&config)?;
532        fs::write(&config_path, &content)?;
533    } else {
534        let content = fs::read_to_string(&config_path)?;
535        let mut config: serde_json::Value = serde_json::from_str(&content)?;
536        if let Some(mcp) = config
537            .as_object_mut()
538            .and_then(|o| o.get_mut("mcpServers"))
539            .and_then(|v| v.as_object_mut())
540        {
541            mcp.remove(&key);
542        }
543        let content = serde_json::to_string_pretty(&config)?;
544        fs::write(&config_path, format!("{content}\n"))?;
545    }
546
547    Ok(())
548}
549
550// ---------------------------------------------------------------------------
551// Import from provider configs
552// ---------------------------------------------------------------------------
553
554/// Import MCP servers from a provider's native config into `~/.zag/mcp/`.
555/// Skips entries prefixed with `zag-` (our own entries).
556/// Returns names of imported servers.
557pub fn import_servers(from_provider: &str) -> Result<Vec<String>> {
558    let Some(config_path) = provider_mcp_config_path(from_provider) else {
559        bail!("Provider '{from_provider}' does not support MCP servers");
560    };
561
562    if !config_path.exists() {
563        bail!(
564            "No MCP config found for '{}' at {}",
565            from_provider,
566            config_path.display()
567        );
568    }
569
570    if from_provider == "codex" {
571        import_from_codex_toml(&config_path)
572    } else {
573        import_from_json(&config_path, from_provider)
574    }
575}
576
577/// Import MCP servers from a JSON config (Claude, Gemini, Copilot).
578fn import_from_json(config_path: &Path, provider: &str) -> Result<Vec<String>> {
579    let content = fs::read_to_string(config_path)?;
580    let config: serde_json::Value = serde_json::from_str(&content)
581        .with_context(|| format!("Failed to parse {}", config_path.display()))?;
582
583    let mcp_servers = match config.get("mcpServers").and_then(|v| v.as_object()) {
584        Some(obj) => obj,
585        None => return Ok(Vec::new()),
586    };
587
588    let dest_dir = mcp_dir();
589    fs::create_dir_all(&dest_dir)?;
590
591    let mut imported = Vec::new();
592
593    for (name, value) in mcp_servers {
594        // Skip our own entries
595        if name.starts_with(MCP_PREFIX) {
596            continue;
597        }
598
599        let dest = dest_dir.join(format!("{name}.toml"));
600        if dest.exists() {
601            log::debug!("Skipping '{name}': already exists in ~/.zag/mcp/");
602            continue;
603        }
604
605        let server = json_entry_to_server(name, value, provider);
606        let content = toml::to_string_pretty(&server).context("Failed to serialize MCP server")?;
607        fs::write(&dest, content).with_context(|| format!("Failed to write {}", dest.display()))?;
608
609        imported.push(name.clone());
610    }
611
612    Ok(imported)
613}
614
615/// Convert a JSON mcpServers entry to our McpServer struct.
616fn json_entry_to_server(name: &str, value: &serde_json::Value, provider: &str) -> McpServer {
617    let obj = value.as_object();
618
619    // Detect transport
620    let transport = if obj.and_then(|o| o.get("url")).is_some()
621        || obj.and_then(|o| o.get("httpUrl")).is_some()
622    {
623        "http".to_string()
624    } else {
625        "stdio".to_string()
626    };
627
628    let command = obj
629        .and_then(|o| o.get("command"))
630        .and_then(|v| v.as_str())
631        .map(|s| s.to_string());
632
633    let args = obj
634        .and_then(|o| o.get("args"))
635        .and_then(|v| v.as_array())
636        .map(|arr| {
637            arr.iter()
638                .filter_map(|v| v.as_str().map(|s| s.to_string()))
639                .collect()
640        })
641        .unwrap_or_default();
642
643    let url = obj
644        .and_then(|o| o.get("url").or_else(|| o.get("httpUrl")))
645        .and_then(|v| v.as_str())
646        .map(|s| s.to_string());
647
648    let env = obj
649        .and_then(|o| o.get("env"))
650        .and_then(|v| v.as_object())
651        .map(|obj| {
652            obj.iter()
653                .filter_map(|(k, v)| v.as_str().map(|s| (k.clone(), s.to_string())))
654                .collect()
655        })
656        .unwrap_or_default();
657
658    let headers = obj
659        .and_then(|o| o.get("headers"))
660        .and_then(|v| v.as_object())
661        .map(|obj| {
662            obj.iter()
663                .filter_map(|(k, v)| v.as_str().map(|s| (k.clone(), s.to_string())))
664                .collect()
665        })
666        .unwrap_or_default();
667
668    let _ = provider; // reserved for provider-specific quirks
669
670    McpServer {
671        name: name.to_string(),
672        description: String::new(),
673        transport,
674        command,
675        args,
676        url,
677        bearer_token_env_var: None,
678        headers,
679        env,
680    }
681}
682
683/// Import MCP servers from Codex TOML config.
684fn import_from_codex_toml(config_path: &Path) -> Result<Vec<String>> {
685    let content = fs::read_to_string(config_path)?;
686    let config: toml::Table = content
687        .parse()
688        .with_context(|| format!("Failed to parse {}", config_path.display()))?;
689
690    let mcp_servers = match config.get("mcp_servers").and_then(|v| v.as_table()) {
691        Some(t) => t,
692        None => return Ok(Vec::new()),
693    };
694
695    let dest_dir = mcp_dir();
696    fs::create_dir_all(&dest_dir)?;
697
698    let mut imported = Vec::new();
699
700    for (name, value) in mcp_servers {
701        if name.starts_with(MCP_PREFIX) {
702            continue;
703        }
704
705        let dest = dest_dir.join(format!("{name}.toml"));
706        if dest.exists() {
707            log::debug!("Skipping '{name}': already exists in ~/.zag/mcp/");
708            continue;
709        }
710
711        let server = toml_entry_to_server(name, value);
712        let content = toml::to_string_pretty(&server)?;
713        fs::write(&dest, content).with_context(|| format!("Failed to write {}", dest.display()))?;
714
715        imported.push(name.clone());
716    }
717
718    Ok(imported)
719}
720
721/// Convert a Codex TOML mcp_servers entry to our McpServer struct.
722fn toml_entry_to_server(name: &str, value: &toml::Value) -> McpServer {
723    let table = value.as_table();
724
725    let transport = if table.and_then(|t| t.get("url")).is_some() {
726        "http".to_string()
727    } else {
728        "stdio".to_string()
729    };
730
731    let command = table
732        .and_then(|t| t.get("command"))
733        .and_then(|v| v.as_str())
734        .map(|s| s.to_string());
735
736    let args = table
737        .and_then(|t| t.get("args"))
738        .and_then(|v| v.as_array())
739        .map(|arr| {
740            arr.iter()
741                .filter_map(|v| v.as_str().map(|s| s.to_string()))
742                .collect()
743        })
744        .unwrap_or_default();
745
746    let url = table
747        .and_then(|t| t.get("url"))
748        .and_then(|v| v.as_str())
749        .map(|s| s.to_string());
750
751    let bearer_token_env_var = table
752        .and_then(|t| t.get("bearer_token_env_var"))
753        .and_then(|v| v.as_str())
754        .map(|s| s.to_string());
755
756    let env = table
757        .and_then(|t| t.get("env"))
758        .and_then(|v| v.as_table())
759        .map(|t| {
760            t.iter()
761                .filter_map(|(k, v)| v.as_str().map(|s| (k.clone(), s.to_string())))
762                .collect()
763        })
764        .unwrap_or_default();
765
766    McpServer {
767        name: name.to_string(),
768        description: String::new(),
769        transport,
770        command,
771        args,
772        url,
773        bearer_token_env_var,
774        headers: BTreeMap::new(),
775        env,
776    }
777}
778
779// ---------------------------------------------------------------------------
780// Orchestration (called automatically before agent sessions)
781// ---------------------------------------------------------------------------
782
783/// Set up MCP servers for the given provider. Called before each agent session.
784pub fn setup_mcp(provider: &str, root: Option<&str>) -> Result<()> {
785    let servers = load_all_servers(root)?;
786    if servers.is_empty() {
787        return Ok(());
788    }
789
790    let synced = sync_servers_for_provider(provider, &servers)?;
791    if synced > 0 {
792        log::info!("Synced {synced} MCP server(s) for {provider}");
793    }
794
795    Ok(())
796}
797
798// ---------------------------------------------------------------------------
799// Testable variants (accept custom directories)
800// ---------------------------------------------------------------------------
801
802/// Load servers from a custom base directory (for testing).
803pub fn load_servers_from_dir(dir: &Path) -> Result<Vec<McpServer>> {
804    load_servers_from(dir)
805}
806
807/// Sync MCP servers to a JSON-format provider config at a custom path (for testing).
808pub fn sync_json_provider_to(
809    provider: &str,
810    servers: &[McpServer],
811    config_path: &Path,
812) -> Result<usize> {
813    let mut config: serde_json::Value = if config_path.exists() {
814        let content = fs::read_to_string(config_path)?;
815        serde_json::from_str(&content)?
816    } else {
817        serde_json::json!({})
818    };
819
820    let mcp_servers = config
821        .as_object_mut()
822        .context("Config is not a JSON object")?
823        .entry("mcpServers")
824        .or_insert_with(|| serde_json::json!({}));
825
826    let mcp_map = mcp_servers
827        .as_object_mut()
828        .context("mcpServers is not a JSON object")?;
829
830    // Remove existing zag- entries
831    let zag_keys: Vec<String> = mcp_map
832        .keys()
833        .filter(|k| k.starts_with(MCP_PREFIX))
834        .cloned()
835        .collect();
836    for key in &zag_keys {
837        mcp_map.remove(key);
838    }
839
840    let mut synced = 0;
841    for server in servers {
842        let key = format!("{}{}", MCP_PREFIX, server.name);
843        let value = server_to_json(server, provider);
844        mcp_map.insert(key, value);
845        synced += 1;
846    }
847
848    let content = serde_json::to_string_pretty(&config)?;
849    fs::write(config_path, format!("{content}\n"))?;
850    Ok(synced)
851}
852
853/// Sync MCP servers to a Codex TOML config at a custom path (for testing).
854pub fn sync_codex_provider_to(servers: &[McpServer], config_path: &Path) -> Result<usize> {
855    let mut config: toml::Table = if config_path.exists() {
856        let content = fs::read_to_string(config_path)?;
857        content.parse()?
858    } else {
859        toml::Table::new()
860    };
861
862    let mcp_table = config
863        .entry("mcp_servers")
864        .or_insert_with(|| toml::Value::Table(toml::Table::new()))
865        .as_table_mut()
866        .context("mcp_servers is not a TOML table")?;
867
868    let zag_keys: Vec<String> = mcp_table
869        .keys()
870        .filter(|k| k.starts_with(MCP_PREFIX))
871        .cloned()
872        .collect();
873    for key in &zag_keys {
874        mcp_table.remove(key.as_str());
875    }
876
877    let mut synced = 0;
878    for server in servers {
879        let key = format!("{}{}", MCP_PREFIX, server.name);
880        let mut entry = toml::Table::new();
881
882        if server.transport == "stdio" {
883            if let Some(ref cmd) = server.command {
884                entry.insert("command".into(), toml::Value::String(cmd.clone()));
885            }
886            if !server.args.is_empty() {
887                let args: Vec<toml::Value> = server
888                    .args
889                    .iter()
890                    .map(|a| toml::Value::String(a.clone()))
891                    .collect();
892                entry.insert("args".into(), toml::Value::Array(args));
893            }
894        } else if server.transport == "http" {
895            if let Some(ref url) = server.url {
896                entry.insert("url".into(), toml::Value::String(url.clone()));
897            }
898            if let Some(ref token_var) = server.bearer_token_env_var {
899                entry.insert(
900                    "bearer_token_env_var".into(),
901                    toml::Value::String(token_var.clone()),
902                );
903            }
904        }
905
906        if !server.env.is_empty() {
907            let mut env_table = toml::Table::new();
908            for (k, v) in &server.env {
909                env_table.insert(k.clone(), toml::Value::String(v.clone()));
910            }
911            entry.insert("env".into(), toml::Value::Table(env_table));
912        }
913
914        mcp_table.insert(key, toml::Value::Table(entry));
915        synced += 1;
916    }
917
918    let content = toml::to_string_pretty(&config)?;
919    fs::write(config_path, &content)?;
920    Ok(synced)
921}
922
923/// Import MCP servers from a JSON config file to a custom destination (for testing).
924pub fn import_from_json_to(
925    config_path: &Path,
926    provider: &str,
927    dest_dir: &Path,
928) -> Result<Vec<String>> {
929    let content = fs::read_to_string(config_path)?;
930    let config: serde_json::Value = serde_json::from_str(&content)?;
931
932    let mcp_servers = match config.get("mcpServers").and_then(|v| v.as_object()) {
933        Some(obj) => obj,
934        None => return Ok(Vec::new()),
935    };
936
937    fs::create_dir_all(dest_dir)?;
938    let mut imported = Vec::new();
939
940    for (name, value) in mcp_servers {
941        if name.starts_with(MCP_PREFIX) {
942            continue;
943        }
944        let dest = dest_dir.join(format!("{name}.toml"));
945        if dest.exists() {
946            continue;
947        }
948        let server = json_entry_to_server(name, value, provider);
949        let content = toml::to_string_pretty(&server)?;
950        fs::write(&dest, content)?;
951        imported.push(name.clone());
952    }
953
954    Ok(imported)
955}
956
957/// Import MCP servers from a Codex TOML config file to a custom destination (for testing).
958pub fn import_from_codex_to(config_path: &Path, dest_dir: &Path) -> Result<Vec<String>> {
959    let content = fs::read_to_string(config_path)?;
960    let config: toml::Table = content.parse()?;
961
962    let mcp_servers = match config.get("mcp_servers").and_then(|v| v.as_table()) {
963        Some(t) => t,
964        None => return Ok(Vec::new()),
965    };
966
967    fs::create_dir_all(dest_dir)?;
968    let mut imported = Vec::new();
969
970    for (name, value) in mcp_servers {
971        if name.starts_with(MCP_PREFIX) {
972            continue;
973        }
974        let dest = dest_dir.join(format!("{name}.toml"));
975        if dest.exists() {
976            continue;
977        }
978        let server = toml_entry_to_server(name, value);
979        let content = toml::to_string_pretty(&server)?;
980        fs::write(&dest, content)?;
981        imported.push(name.clone());
982    }
983
984    Ok(imported)
985}