chore: persist comments in edit (#7931)

This PR makes sure that inline comment is preserved for mcp server
config and arbitrary key/value setPath config.

---------

Co-authored-by: celia-oai <celia@openai.com>
This commit is contained in:
jif-oai 2025-12-16 01:05:49 +01:00 committed by GitHub
parent f754b19e80
commit 3ee5c40261
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 424 additions and 9 deletions

View file

@ -90,7 +90,7 @@ mod document_helpers {
}
}
pub(super) fn serialize_mcp_server(config: &McpServerConfig) -> TomlItem {
fn serialize_mcp_server_table(config: &McpServerConfig) -> TomlTable {
let mut entry = TomlTable::new();
entry.set_implicit(false);
@ -161,7 +161,29 @@ mod document_helpers {
entry["disabled_tools"] = array_from_iter(disabled_tools.iter().cloned());
}
TomlItem::Table(entry)
entry
}
pub(super) fn serialize_mcp_server(config: &McpServerConfig) -> TomlItem {
TomlItem::Table(serialize_mcp_server_table(config))
}
pub(super) fn serialize_mcp_server_inline(config: &McpServerConfig) -> InlineTable {
serialize_mcp_server_table(config).into_inline_table()
}
pub(super) fn merge_inline_table(existing: &mut InlineTable, replacement: InlineTable) {
existing.retain(|key, _| replacement.get(key).is_some());
for (key, value) in replacement.iter() {
if let Some(existing_value) = existing.get_mut(key) {
let mut updated_value = value.clone();
*updated_value.decor_mut() = existing_value.decor().clone();
*existing_value = updated_value;
} else {
existing.insert(key.to_string(), value.clone());
}
}
}
fn table_from_inline(inline: &InlineTable) -> TomlTable {
@ -317,15 +339,52 @@ impl ConfigDocument {
return self.clear(Scope::Global, &["mcp_servers"]);
}
let mut table = TomlTable::new();
table.set_implicit(true);
for (name, config) in servers {
table.insert(name, document_helpers::serialize_mcp_server(config));
let root = self.doc.as_table_mut();
if !root.contains_key("mcp_servers") {
root.insert(
"mcp_servers",
TomlItem::Table(document_helpers::new_implicit_table()),
);
}
let item = TomlItem::Table(table);
self.write_value(Scope::Global, &["mcp_servers"], item)
let Some(item) = root.get_mut("mcp_servers") else {
return false;
};
if document_helpers::ensure_table_for_write(item).is_none() {
*item = TomlItem::Table(document_helpers::new_implicit_table());
}
let Some(table) = item.as_table_mut() else {
return false;
};
let keys_to_remove: Vec<String> = table
.iter()
.map(|(key, _)| key.to_string())
.filter(|key| !servers.contains_key(key.as_str()))
.collect();
for key in keys_to_remove {
table.remove(&key);
}
for (name, config) in servers {
if let Some(existing) = table.get_mut(name.as_str()) {
if let TomlItem::Value(value) = existing
&& let Some(inline) = value.as_inline_table_mut()
{
let replacement = document_helpers::serialize_mcp_server_inline(config);
document_helpers::merge_inline_table(inline, replacement);
} else {
*existing = document_helpers::serialize_mcp_server(config);
}
} else {
table.insert(name, document_helpers::serialize_mcp_server(config));
}
}
true
}
fn scoped_segments(&self, scope: Scope, segments: &[&str]) -> Vec<String> {
@ -357,6 +416,10 @@ impl ConfigDocument {
return false;
};
let mut value = value;
if let Some(existing) = parent.get(last) {
Self::preserve_decor(existing, &mut value);
}
parent[last] = value;
true
}
@ -398,6 +461,37 @@ impl ConfigDocument {
Some(current)
}
fn preserve_decor(existing: &TomlItem, replacement: &mut TomlItem) {
match (existing, replacement) {
(TomlItem::Table(existing_table), TomlItem::Table(replacement_table)) => {
replacement_table
.decor_mut()
.clone_from(existing_table.decor());
for (key, existing_item) in existing_table.iter() {
if let (Some(existing_key), Some(mut replacement_key)) =
(existing_table.key(key), replacement_table.key_mut(key))
{
replacement_key
.leaf_decor_mut()
.clone_from(existing_key.leaf_decor());
replacement_key
.dotted_decor_mut()
.clone_from(existing_key.dotted_decor());
}
if let Some(replacement_item) = replacement_table.get_mut(key) {
Self::preserve_decor(existing_item, replacement_item);
}
}
}
(TomlItem::Value(existing_value), TomlItem::Value(replacement_value)) => {
replacement_value
.decor_mut()
.clone_from(existing_value.decor());
}
_ => {}
}
}
}
/// Persist edits using a blocking strategy.
@ -691,6 +785,68 @@ profiles = { fast = { model = "gpt-4o", sandbox_mode = "strict" } }
);
}
#[test]
fn batch_write_table_upsert_preserves_inline_comments() {
let tmp = tempdir().expect("tmpdir");
let codex_home = tmp.path();
let original = r#"approval_policy = "never"
[mcp_servers.linear]
name = "linear"
# ok
url = "https://linear.example"
[mcp_servers.linear.http_headers]
foo = "bar"
[sandbox_workspace_write]
# ok 3
network_access = false
"#;
std::fs::write(codex_home.join(CONFIG_TOML_FILE), original).expect("seed config");
apply_blocking(
codex_home,
None,
&[
ConfigEdit::SetPath {
segments: vec![
"mcp_servers".to_string(),
"linear".to_string(),
"url".to_string(),
],
value: value("https://linear.example/v2"),
},
ConfigEdit::SetPath {
segments: vec![
"sandbox_workspace_write".to_string(),
"network_access".to_string(),
],
value: value(true),
},
],
)
.expect("apply");
let updated =
std::fs::read_to_string(codex_home.join(CONFIG_TOML_FILE)).expect("read config");
let expected = r#"approval_policy = "never"
[mcp_servers.linear]
name = "linear"
# ok
url = "https://linear.example/v2"
[mcp_servers.linear.http_headers]
foo = "bar"
[sandbox_workspace_write]
# ok 3
network_access = true
"#;
assert_eq!(updated, expected);
}
#[test]
fn blocking_clear_model_removes_inline_table_entry() {
let tmp = tempdir().expect("tmpdir");
@ -1028,6 +1184,178 @@ B = \"2\"
assert_eq!(raw, expected);
}
#[test]
fn blocking_replace_mcp_servers_preserves_inline_comments() {
let tmp = tempdir().expect("tmpdir");
let codex_home = tmp.path();
std::fs::write(
codex_home.join(CONFIG_TOML_FILE),
r#"[mcp_servers]
# keep me
foo = { command = "cmd" }
"#,
)
.expect("seed");
let mut servers = BTreeMap::new();
servers.insert(
"foo".to_string(),
McpServerConfig {
transport: McpServerTransportConfig::Stdio {
command: "cmd".to_string(),
args: Vec::new(),
env: None,
env_vars: Vec::new(),
cwd: None,
},
enabled: true,
startup_timeout_sec: None,
tool_timeout_sec: None,
enabled_tools: None,
disabled_tools: None,
},
);
apply_blocking(codex_home, None, &[ConfigEdit::ReplaceMcpServers(servers)])
.expect("persist");
let contents =
std::fs::read_to_string(codex_home.join(CONFIG_TOML_FILE)).expect("read config");
let expected = r#"[mcp_servers]
# keep me
foo = { command = "cmd" }
"#;
assert_eq!(contents, expected);
}
#[test]
fn blocking_replace_mcp_servers_preserves_inline_comment_suffix() {
let tmp = tempdir().expect("tmpdir");
let codex_home = tmp.path();
std::fs::write(
codex_home.join(CONFIG_TOML_FILE),
r#"[mcp_servers]
foo = { command = "cmd" } # keep me
"#,
)
.expect("seed");
let mut servers = BTreeMap::new();
servers.insert(
"foo".to_string(),
McpServerConfig {
transport: McpServerTransportConfig::Stdio {
command: "cmd".to_string(),
args: Vec::new(),
env: None,
env_vars: Vec::new(),
cwd: None,
},
enabled: false,
startup_timeout_sec: None,
tool_timeout_sec: None,
enabled_tools: None,
disabled_tools: None,
},
);
apply_blocking(codex_home, None, &[ConfigEdit::ReplaceMcpServers(servers)])
.expect("persist");
let contents =
std::fs::read_to_string(codex_home.join(CONFIG_TOML_FILE)).expect("read config");
let expected = r#"[mcp_servers]
foo = { command = "cmd" , enabled = false } # keep me
"#;
assert_eq!(contents, expected);
}
#[test]
fn blocking_replace_mcp_servers_preserves_inline_comment_after_removing_keys() {
let tmp = tempdir().expect("tmpdir");
let codex_home = tmp.path();
std::fs::write(
codex_home.join(CONFIG_TOML_FILE),
r#"[mcp_servers]
foo = { command = "cmd", args = ["--flag"] } # keep me
"#,
)
.expect("seed");
let mut servers = BTreeMap::new();
servers.insert(
"foo".to_string(),
McpServerConfig {
transport: McpServerTransportConfig::Stdio {
command: "cmd".to_string(),
args: Vec::new(),
env: None,
env_vars: Vec::new(),
cwd: None,
},
enabled: true,
startup_timeout_sec: None,
tool_timeout_sec: None,
enabled_tools: None,
disabled_tools: None,
},
);
apply_blocking(codex_home, None, &[ConfigEdit::ReplaceMcpServers(servers)])
.expect("persist");
let contents =
std::fs::read_to_string(codex_home.join(CONFIG_TOML_FILE)).expect("read config");
let expected = r#"[mcp_servers]
foo = { command = "cmd"} # keep me
"#;
assert_eq!(contents, expected);
}
#[test]
fn blocking_replace_mcp_servers_preserves_inline_comment_prefix_on_update() {
let tmp = tempdir().expect("tmpdir");
let codex_home = tmp.path();
std::fs::write(
codex_home.join(CONFIG_TOML_FILE),
r#"[mcp_servers]
# keep me
foo = { command = "cmd" }
"#,
)
.expect("seed");
let mut servers = BTreeMap::new();
servers.insert(
"foo".to_string(),
McpServerConfig {
transport: McpServerTransportConfig::Stdio {
command: "cmd".to_string(),
args: Vec::new(),
env: None,
env_vars: Vec::new(),
cwd: None,
},
enabled: false,
startup_timeout_sec: None,
tool_timeout_sec: None,
enabled_tools: None,
disabled_tools: None,
},
);
apply_blocking(codex_home, None, &[ConfigEdit::ReplaceMcpServers(servers)])
.expect("persist");
let contents =
std::fs::read_to_string(codex_home.join(CONFIG_TOML_FILE)).expect("read config");
let expected = r#"[mcp_servers]
# keep me
foo = { command = "cmd" , enabled = false }
"#;
assert_eq!(contents, expected);
}
#[test]
fn blocking_clear_path_noop_when_missing() {
let tmp = tempdir().expect("tmpdir");

View file

@ -932,4 +932,91 @@ remote_compaction = true
assert_eq!(overridden.overriding_layer.name, ConfigLayerName::System);
assert_eq!(overridden.effective_value, serde_json::json!("never"));
}
#[tokio::test]
async fn upsert_merges_tables_replace_overwrites() -> Result<()> {
let tmp = tempdir().expect("tempdir");
let path = tmp.path().join(CONFIG_TOML_FILE);
let base = r#"[mcp_servers.linear]
bearer_token_env_var = "TOKEN"
name = "linear"
url = "https://linear.example"
[mcp_servers.linear.env_http_headers]
existing = "keep"
[mcp_servers.linear.http_headers]
alpha = "a"
"#;
let overlay = serde_json::json!({
"bearer_token_env_var": "NEW_TOKEN",
"http_headers": {
"alpha": "updated",
"beta": "b"
},
"name": "linear",
"url": "https://linear.example"
});
std::fs::write(&path, base)?;
let service = ConfigService::new(tmp.path().to_path_buf(), vec![]);
service
.write_value(ConfigValueWriteParams {
file_path: Some(path.display().to_string()),
key_path: "mcp_servers.linear".to_string(),
value: overlay.clone(),
merge_strategy: MergeStrategy::Upsert,
expected_version: None,
})
.await
.expect("upsert succeeds");
let upserted: TomlValue = toml::from_str(&std::fs::read_to_string(&path)?)?;
let expected_upsert: TomlValue = toml::from_str(
r#"[mcp_servers.linear]
bearer_token_env_var = "NEW_TOKEN"
name = "linear"
url = "https://linear.example"
[mcp_servers.linear.env_http_headers]
existing = "keep"
[mcp_servers.linear.http_headers]
alpha = "updated"
beta = "b"
"#,
)?;
assert_eq!(upserted, expected_upsert);
std::fs::write(&path, base)?;
service
.write_value(ConfigValueWriteParams {
file_path: Some(path.display().to_string()),
key_path: "mcp_servers.linear".to_string(),
value: overlay,
merge_strategy: MergeStrategy::Replace,
expected_version: None,
})
.await
.expect("replace succeeds");
let replaced: TomlValue = toml::from_str(&std::fs::read_to_string(&path)?)?;
let expected_replace: TomlValue = toml::from_str(
r#"[mcp_servers.linear]
bearer_token_env_var = "NEW_TOKEN"
name = "linear"
url = "https://linear.example"
[mcp_servers.linear.http_headers]
alpha = "updated"
beta = "b"
"#,
)?;
assert_eq!(replaced, expected_replace);
Ok(())
}
}