Skip to main content

codex_config/
mcp_edit.rs

1use std::collections::BTreeMap;
2use std::fs;
3use std::io::ErrorKind;
4use std::path::Path;
5use std::path::PathBuf;
6
7use tokio::task;
8use toml::Value as TomlValue;
9use toml_edit::DocumentMut;
10use toml_edit::Item as TomlItem;
11use toml_edit::Table as TomlTable;
12use toml_edit::value;
13
14use crate::AppToolApproval;
15use crate::CONFIG_TOML_FILE;
16use crate::McpServerAuth;
17use crate::McpServerConfig;
18use crate::McpServerEnvVar;
19use crate::McpServerTransportConfig;
20
21pub async fn load_global_mcp_servers(
22    codex_home: &Path,
23) -> std::io::Result<BTreeMap<String, McpServerConfig>> {
24    let config_path = codex_home.join(CONFIG_TOML_FILE);
25    let raw = match tokio::fs::read_to_string(&config_path).await {
26        Ok(raw) => raw,
27        Err(err) if err.kind() == ErrorKind::NotFound => return Ok(BTreeMap::new()),
28        Err(err) => return Err(err),
29    };
30    let parsed = toml::from_str::<TomlValue>(&raw)
31        .map_err(|err| std::io::Error::new(ErrorKind::InvalidData, err))?;
32    let Some(servers_value) = parsed.get("mcp_servers") else {
33        return Ok(BTreeMap::new());
34    };
35
36    ensure_no_inline_bearer_tokens(servers_value)?;
37
38    servers_value
39        .clone()
40        .try_into()
41        .map_err(|err| std::io::Error::new(ErrorKind::InvalidData, err))
42}
43
44fn ensure_no_inline_bearer_tokens(value: &TomlValue) -> std::io::Result<()> {
45    let Some(servers_table) = value.as_table() else {
46        return Ok(());
47    };
48
49    for (server_name, server_value) in servers_table {
50        if let Some(server_table) = server_value.as_table()
51            && server_table.contains_key("bearer_token")
52        {
53            let message = format!(
54                "mcp_servers.{server_name} uses unsupported `bearer_token`; set `bearer_token_env_var`."
55            );
56            return Err(std::io::Error::new(ErrorKind::InvalidData, message));
57        }
58    }
59
60    Ok(())
61}
62
63pub struct ConfigEditsBuilder {
64    codex_home: PathBuf,
65    mcp_servers: Option<BTreeMap<String, McpServerConfig>>,
66}
67
68impl ConfigEditsBuilder {
69    pub fn new(codex_home: &Path) -> Self {
70        Self {
71            codex_home: codex_home.to_path_buf(),
72            mcp_servers: None,
73        }
74    }
75
76    pub fn replace_mcp_servers(mut self, servers: &BTreeMap<String, McpServerConfig>) -> Self {
77        self.mcp_servers = Some(servers.clone());
78        self
79    }
80
81    pub async fn apply(self) -> std::io::Result<()> {
82        task::spawn_blocking(move || self.apply_blocking())
83            .await
84            .map_err(|err| {
85                std::io::Error::other(format!("config persistence task panicked: {err}"))
86            })?
87    }
88
89    fn apply_blocking(self) -> std::io::Result<()> {
90        let config_path = self.codex_home.join(CONFIG_TOML_FILE);
91        let mut doc = read_or_create_document(&config_path)?;
92        if let Some(servers) = self.mcp_servers.as_ref() {
93            replace_mcp_servers(&mut doc, servers);
94        }
95        fs::create_dir_all(&self.codex_home)?;
96        fs::write(config_path, doc.to_string())
97    }
98}
99
100fn read_or_create_document(config_path: &Path) -> std::io::Result<DocumentMut> {
101    match fs::read_to_string(config_path) {
102        Ok(raw) => raw
103            .parse::<DocumentMut>()
104            .map_err(|err| std::io::Error::new(ErrorKind::InvalidData, err)),
105        Err(err) if err.kind() == ErrorKind::NotFound => Ok(DocumentMut::new()),
106        Err(err) => Err(err),
107    }
108}
109
110fn replace_mcp_servers(doc: &mut DocumentMut, servers: &BTreeMap<String, McpServerConfig>) {
111    let root = doc.as_table_mut();
112    if servers.is_empty() {
113        root.remove("mcp_servers");
114        return;
115    }
116
117    let mut table = TomlTable::new();
118    table.set_implicit(true);
119    for (name, config) in servers {
120        table.insert(name, serialize_mcp_server(config));
121    }
122    root.insert("mcp_servers", TomlItem::Table(table));
123}
124
125fn serialize_mcp_server(config: &McpServerConfig) -> TomlItem {
126    let mut entry = TomlTable::new();
127    entry.set_implicit(false);
128
129    match &config.transport {
130        McpServerTransportConfig::Stdio {
131            command,
132            args,
133            env,
134            env_vars,
135            cwd,
136        } => {
137            entry["command"] = value(command.clone());
138            if !args.is_empty() {
139                entry["args"] = array_from_strings(args);
140            }
141            if let Some(env) = env
142                && !env.is_empty()
143            {
144                entry["env"] = table_from_pairs(env.iter());
145            }
146            if !env_vars.is_empty() {
147                entry["env_vars"] = array_from_env_vars(env_vars);
148            }
149            if let Some(cwd) = cwd {
150                entry["cwd"] = value(cwd.as_str());
151            }
152        }
153        McpServerTransportConfig::StreamableHttp {
154            url,
155            bearer_token_env_var,
156            http_headers,
157            env_http_headers,
158        } => {
159            entry["url"] = value(url.clone());
160            if let Some(env_var) = bearer_token_env_var {
161                entry["bearer_token_env_var"] = value(env_var.clone());
162            }
163            if let Some(headers) = http_headers
164                && !headers.is_empty()
165            {
166                entry["http_headers"] = table_from_pairs(headers.iter());
167            }
168            if let Some(headers) = env_http_headers
169                && !headers.is_empty()
170            {
171                entry["env_http_headers"] = table_from_pairs(headers.iter());
172            }
173        }
174    }
175
176    if matches!(&config.auth, McpServerAuth::ChatGpt) {
177        entry["auth"] = value("chatgpt");
178    }
179    if !config.enabled {
180        entry["enabled"] = value(false);
181    }
182    if !config.is_local_environment() {
183        entry["environment_id"] = value(config.environment_id.clone());
184    }
185    if config.required {
186        entry["required"] = value(true);
187    }
188    if config.supports_parallel_tool_calls {
189        entry["supports_parallel_tool_calls"] = value(true);
190    }
191    if let Some(timeout) = config.startup_timeout_sec {
192        entry["startup_timeout_sec"] = value(timeout.as_secs_f64());
193    }
194    if let Some(timeout) = config.tool_timeout_sec {
195        entry["tool_timeout_sec"] = value(timeout.as_secs_f64());
196    }
197    if let Some(approval_mode) = config.default_tools_approval_mode {
198        entry["default_tools_approval_mode"] = value(match approval_mode {
199            AppToolApproval::Auto => "auto",
200            AppToolApproval::Prompt => "prompt",
201            AppToolApproval::Writes => "writes",
202            AppToolApproval::Approve => "approve",
203        });
204    }
205    if let Some(enabled_tools) = &config.enabled_tools
206        && !enabled_tools.is_empty()
207    {
208        entry["enabled_tools"] = array_from_strings(enabled_tools);
209    }
210    if let Some(disabled_tools) = &config.disabled_tools
211        && !disabled_tools.is_empty()
212    {
213        entry["disabled_tools"] = array_from_strings(disabled_tools);
214    }
215    if let Some(scopes) = &config.scopes
216        && !scopes.is_empty()
217    {
218        entry["scopes"] = array_from_strings(scopes);
219    }
220    if let Some(oauth) = &config.oauth
221        && let Some(client_id) = &oauth.client_id
222        && !client_id.is_empty()
223    {
224        let mut oauth_table = TomlTable::new();
225        oauth_table.set_implicit(false);
226        oauth_table["client_id"] = value(client_id.clone());
227        entry["oauth"] = TomlItem::Table(oauth_table);
228    }
229    if let Some(resource) = &config.oauth_resource
230        && !resource.is_empty()
231    {
232        entry["oauth_resource"] = value(resource.clone());
233    }
234    if !config.tools.is_empty() {
235        let mut tools = TomlTable::new();
236        tools.set_implicit(false);
237        let mut tool_entries: Vec<_> = config.tools.iter().collect();
238        tool_entries.sort_by_key(|(name, _)| *name);
239        for (name, tool_config) in tool_entries {
240            let mut tool_entry = TomlTable::new();
241            tool_entry.set_implicit(false);
242            if let Some(approval_mode) = tool_config.approval_mode {
243                tool_entry["approval_mode"] = value(match approval_mode {
244                    AppToolApproval::Auto => "auto",
245                    AppToolApproval::Prompt => "prompt",
246                    AppToolApproval::Writes => "writes",
247                    AppToolApproval::Approve => "approve",
248                });
249            }
250            tools.insert(name, TomlItem::Table(tool_entry));
251        }
252        entry.insert("tools", TomlItem::Table(tools));
253    }
254
255    TomlItem::Table(entry)
256}
257
258fn array_from_strings(values: &[String]) -> TomlItem {
259    let mut array = toml_edit::Array::new();
260    for value in values {
261        array.push(value.clone());
262    }
263    TomlItem::Value(array.into())
264}
265
266fn array_from_env_vars(env_vars: &[McpServerEnvVar]) -> TomlItem {
267    let mut array = toml_edit::Array::new();
268    for env_var in env_vars {
269        match env_var {
270            McpServerEnvVar::Name(name) => array.push(name.clone()),
271            McpServerEnvVar::Config { name, source } => {
272                let mut table = toml_edit::InlineTable::new();
273                table.insert("name", name.clone().into());
274                if let Some(source) = source {
275                    table.insert("source", source.clone().into());
276                }
277                array.push(table);
278            }
279        }
280    }
281    TomlItem::Value(array.into())
282}
283
284fn table_from_pairs<'a, I>(pairs: I) -> TomlItem
285where
286    I: IntoIterator<Item = (&'a String, &'a String)>,
287{
288    let mut entries: Vec<_> = pairs.into_iter().collect();
289    entries.sort_by_key(|(key, _)| *key);
290    let mut table = TomlTable::new();
291    table.set_implicit(false);
292    for (key, value_str) in entries {
293        table.insert(key, value(value_str.clone()));
294    }
295    TomlItem::Table(table)
296}
297
298#[cfg(test)]
299#[path = "mcp_edit_tests.rs"]
300mod tests;