use std::fs;
use std::path::Path;
use anyhow::{anyhow, Context, Result};
pub use yaml_edit::path::YamlPath;
pub use yaml_edit::{Document, Mapping, Sequence};
pub fn load(path: &Path) -> Result<Document> {
let raw = fs::read_to_string(path).with_context(|| format!("read {}", path.display()))?;
raw.parse::<Document>()
.with_context(|| format!("parse {}", path.display()))
}
pub fn save(doc: &Document, path: &Path) -> Result<()> {
fs::write(path, doc.to_string()).with_context(|| format!("write {}", path.display()))?;
Ok(())
}
pub fn set_top_level_scalar(source: &str, key: &str, value: &str) -> Result<String> {
let trailing_newline = source.ends_with('\n');
let mut out_lines: Vec<String> = Vec::new();
let mut rewrote = false;
for line in source.lines() {
if !rewrote {
let trimmed = line.trim_start();
let indent = line.len() - trimmed.len();
if indent == 0 && !trimmed.is_empty() && !trimmed.starts_with('#') {
if let Some((found_key, _rest)) = trimmed.split_once(':') {
if found_key == key {
out_lines.push(format!("{key}: {value}"));
rewrote = true;
continue;
}
}
}
}
out_lines.push(line.to_string());
}
if !rewrote {
return Err(anyhow!(
"set_top_level_scalar: no top-level key `{key}` found in document"
));
}
let mut joined = out_lines.join("\n");
if trailing_newline && !joined.ends_with('\n') {
joined.push('\n');
}
Ok(joined)
}
pub fn set_nested_mapping(
doc: Document,
parent_path: &[&str],
value_pairs: &[(&str, &str)],
) -> Result<Document> {
if parent_path.is_empty() {
return Err(anyhow!("set_nested_mapping: parent_path must not be empty"));
}
let source = doc.to_string();
let edited = splice_nested_mapping(&source, parent_path, value_pairs)?;
edited
.parse::<Document>()
.with_context(|| "re-parse spliced YAML")
}
fn splice_nested_mapping(
source: &str,
path: &[&str],
value_pairs: &[(&str, &str)],
) -> Result<String> {
let lines: Vec<&str> = source.lines().collect();
let trailing_newline = source.ends_with('\n');
let mut current_indent: usize = 0;
let mut search_start: usize = 0;
let mut search_end: usize = lines.len();
let mut existing_depth: usize = 0;
let mut leaf_replace_range: Option<(usize, usize, usize)> = None;
for (depth, key) in path.iter().enumerate() {
let parent_indent = current_indent;
let child_indent_min = if depth == 0 { 0 } else { parent_indent + 1 };
match find_key_in_block(&lines, search_start, search_end, key, child_indent_min) {
Some((line_idx, key_indent)) => {
existing_depth = depth + 1;
current_indent = key_indent;
let block_end = block_end_after(&lines, line_idx, key_indent);
if depth == path.len() - 1 {
leaf_replace_range = Some((line_idx, block_end, key_indent));
} else {
search_start = line_idx + 1;
search_end = block_end;
}
}
None => break,
}
}
if existing_depth == 0 {
return Err(anyhow!(
"set_nested_mapping: top-level key `{}` not found",
path[0]
));
}
let insert_indent = if existing_depth == path.len() {
leaf_replace_range.expect("leaf existed").2
} else {
current_indent + 2
};
let missing_tail = &path[existing_depth..];
let mut block_lines: Vec<String> = Vec::new();
let mut indent = insert_indent;
for key in missing_tail {
block_lines.push(format!("{:indent$}{key}:", "", indent = indent, key = key));
indent += 2;
}
let value_indent = if existing_depth == path.len() {
insert_indent + 2
} else {
indent
};
if existing_depth == path.len() {
block_lines.push(format!(
"{:indent$}{key}:",
"",
indent = insert_indent,
key = path[path.len() - 1]
));
}
for (k, v) in value_pairs {
block_lines.push(format!(
"{:indent$}{k}: {v}",
"",
indent = value_indent,
k = k,
v = v
));
}
let mut out_lines: Vec<String> = lines.iter().map(|s| s.to_string()).collect();
if let Some((start, end, _)) = leaf_replace_range {
out_lines.splice(start..end, block_lines);
} else {
out_lines.splice(search_end..search_end, block_lines);
}
let mut joined = out_lines.join("\n");
if trailing_newline && !joined.ends_with('\n') {
joined.push('\n');
}
Ok(joined)
}
fn find_key_in_block(
lines: &[&str],
start: usize,
end: usize,
key: &str,
min_indent: usize,
) -> Option<(usize, usize)> {
let mut child_indent: Option<usize> = None;
for line in lines.iter().take(end).skip(start) {
if let Some((indent, _)) = parse_mapping_key_line(line) {
if indent >= min_indent {
child_indent = Some(child_indent.map_or(indent, |c| c.min(indent)));
}
}
}
let child_indent = child_indent?;
for (i, line) in lines.iter().enumerate().take(end).skip(start) {
if let Some((indent, found_key)) = parse_mapping_key_line(line) {
if indent == child_indent && found_key == key {
return Some((i, indent));
}
}
}
None
}
fn parse_mapping_key_line(line: &str) -> Option<(usize, &str)> {
let indent = line.len() - line.trim_start().len();
let trimmed = line.trim_start();
if trimmed.is_empty() || trimmed.starts_with('#') || trimmed.starts_with('-') {
return None;
}
let colon_idx = trimmed.find(':')?;
let key = &trimmed[..colon_idx];
if key.is_empty() {
return None;
}
if key.contains(':') {
return None;
}
let after = &trimmed[colon_idx + 1..];
if !after.is_empty() && !after.starts_with(char::is_whitespace) {
return None;
}
Some((indent, key))
}
fn block_end_after(lines: &[&str], key_line: usize, key_indent: usize) -> usize {
for (i, line) in lines.iter().enumerate().skip(key_line + 1) {
let trimmed = line.trim_start();
if trimmed.is_empty() || trimmed.starts_with('#') {
continue;
}
let indent = line.len() - trimmed.len();
if indent <= key_indent {
return i;
}
}
lines.len()
}
#[cfg(test)]
mod tests {
use super::*;
const COMMENTED_FIXTURE: &str = "\
version: 2
# managers block: each manager is a long-running agent.
managers:
pm:
runtime: claude-code # canonical runtime
role_prompt: roles/pm.md
# interfaces lands here once `teamctl bot setup` runs
eng_lead:
runtime: claude-code
role_prompt: roles/eng_lead.md
# trailing footer
";
#[test]
fn round_trip_preserves_byte_for_byte() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("fixture.yaml");
fs::write(&path, COMMENTED_FIXTURE).unwrap();
let doc = load(&path).unwrap();
save(&doc, &path).unwrap();
let after = fs::read_to_string(&path).unwrap();
assert_eq!(
after, COMMENTED_FIXTURE,
"load → save without mutation must be byte-perfect"
);
}
#[test]
fn mutation_preserves_comments() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("fixture.yaml");
fs::write(&path, COMMENTED_FIXTURE).unwrap();
let doc = load(&path).unwrap();
let doc = set_nested_mapping(
doc,
&["managers", "pm", "interfaces", "telegram"],
&[("bot_token_env", "PM_TOKEN"), ("chat_ids_env", "PM_CHATS")],
)
.unwrap();
save(&doc, &path).unwrap();
let after = fs::read_to_string(&path).unwrap();
assert!(
after.contains("# managers block: each manager is a long-running agent."),
"block comment dropped:\n{after}"
);
assert!(
after.contains("# canonical runtime"),
"trailing line comment dropped:\n{after}"
);
assert!(
after.contains("# trailing footer"),
"footer comment dropped:\n{after}"
);
assert!(
after.contains(" interfaces:"),
"interfaces not properly indented under pm:\n{after}"
);
assert!(
after.contains(" telegram:"),
"telegram not properly indented under interfaces:\n{after}"
);
assert!(
after.contains(" bot_token_env: PM_TOKEN"),
"leaf not properly indented:\n{after}"
);
assert!(after.contains(" chat_ids_env: PM_CHATS"));
let pm_idx = after.find("pm:").expect("pm key");
let eng_idx = after.find("eng_lead:").expect("eng_lead key");
assert!(pm_idx < eng_idx, "manager key order swapped:\n{after}");
assert!(
after.contains("\n eng_lead:"),
"eng_lead boundary broken:\n{after}"
);
}
#[test]
fn save_does_not_strip_existing_comments() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("oss-shape.yaml");
let fixture = "\
version: 2
project:
id: oss
name: OSS Maintainer
cwd: ./workspace
# Hub-and-spoke: maintainer is the only manager; workers fan out below.
managers:
maintainer:
runtime: claude-code
role_prompt: roles/maintainer.md
# `teamctl bot setup` writes the interfaces.telegram block here.
workers:
bug_fix:
runtime: claude-code # workers default to sonnet
reports_to: maintainer
";
fs::write(&path, fixture).unwrap();
let doc = load(&path).unwrap();
let doc = set_nested_mapping(
doc,
&["managers", "maintainer", "interfaces", "telegram"],
&[
("bot_token_env", "TEAMCTL_TG_MAINTAINER_TOKEN"),
("chat_ids_env", "TEAMCTL_TG_MAINTAINER_CHATS"),
],
)
.unwrap();
save(&doc, &path).unwrap();
let after = fs::read_to_string(&path).unwrap();
assert!(
after.contains(
"# Hub-and-spoke: maintainer is the only manager; workers fan out below."
),
"block comment dropped — regression class still open:\n{after}"
);
assert!(
after.contains("# `teamctl bot setup` writes the interfaces.telegram block here."),
"inline comment dropped:\n{after}"
);
assert!(
after.contains("# workers default to sonnet"),
"trailing line comment dropped:\n{after}"
);
assert!(after.contains(" interfaces:"));
assert!(after.contains(" telegram:"));
assert!(after.contains(" bot_token_env: TEAMCTL_TG_MAINTAINER_TOKEN"));
assert!(after.contains(" chat_ids_env: TEAMCTL_TG_MAINTAINER_CHATS"));
}
#[test]
fn idempotent_replace_preserves_siblings() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("siblings.yaml");
let fixture = "\
version: 2
managers:
pm:
runtime: claude-code
interfaces:
discord:
bot_token_env: PM_DISCORD_TOKEN
telegram:
bot_token_env: OLD_TOKEN
chat_ids_env: OLD_CHATS
";
fs::write(&path, fixture).unwrap();
let doc = load(&path).unwrap();
let doc = set_nested_mapping(
doc,
&["managers", "pm", "interfaces", "telegram"],
&[
("bot_token_env", "NEW_TOKEN"),
("chat_ids_env", "NEW_CHATS"),
],
)
.unwrap();
save(&doc, &path).unwrap();
let after = fs::read_to_string(&path).unwrap();
assert_eq!(
after.matches("telegram:").count(),
1,
"duplicate telegram block:\n{after}"
);
assert_eq!(
after.matches("discord:").count(),
1,
"discord sibling lost:\n{after}"
);
assert!(
after.contains("PM_DISCORD_TOKEN"),
"discord adapter contents lost:\n{after}"
);
assert!(after.contains("NEW_TOKEN"));
assert!(after.contains("NEW_CHATS"));
assert!(!after.contains("OLD_TOKEN"));
assert!(!after.contains("OLD_CHATS"));
}
#[test]
fn replace_existing_leaf_nests_values_under_it() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("prewired.yaml");
let fixture = "\
version: 2
managers:
builder:
runtime: claude-code
interfaces:
telegram:
bot_token_env: TEAMCTL_TG_BUILDER_TOKEN
chat_ids_env: TEAMCTL_TG_BUILDER_CHATS
";
fs::write(&path, fixture).unwrap();
let doc = load(&path).unwrap();
let doc = set_nested_mapping(
doc,
&["managers", "builder", "interfaces", "telegram"],
&[
("bot_token_env", "TEAMCTL_TG_BUILDER_TOKEN"),
("chat_ids_env", "TEAMCTL_TG_BUILDER_CHATS"),
],
)
.unwrap();
save(&doc, &path).unwrap();
let after = fs::read_to_string(&path).unwrap();
let v: serde_yaml::Value = serde_yaml::from_str(&after)
.unwrap_or_else(|e| panic!("re-spliced YAML must parse: {e}\n{after}"));
let tg = &v["managers"]["builder"]["interfaces"]["telegram"];
assert!(
tg.is_mapping(),
"telegram leaf must remain a mapping, not be nulled by sibling pairs:\n{after}"
);
assert_eq!(
tg["bot_token_env"].as_str(),
Some("TEAMCTL_TG_BUILDER_TOKEN"),
"bot_token_env must be nested under telegram:\n{after}"
);
assert_eq!(
tg["chat_ids_env"].as_str(),
Some("TEAMCTL_TG_BUILDER_CHATS"),
"chat_ids_env must be nested under telegram:\n{after}"
);
}
#[test]
fn set_top_level_scalar_replaces_integer_with_quoted_string() {
let src = "\
# leading comment
version: 2
broker:
type: sqlite
";
let edited = set_top_level_scalar(src, "version", "\"2.0.0\"").unwrap();
let out = edited;
assert!(
out.contains("version: \"2.0.0\""),
"rewrite missing:\n{out}"
);
assert!(
!out.contains("\nversion: 2\n"),
"old literal survived:\n{out}"
);
assert!(out.contains("# leading comment"));
assert!(out.contains("broker:"));
assert!(out.contains("type: sqlite"));
}
#[test]
fn set_top_level_scalar_is_idempotent() {
let src = "version: \"2.0.0\"\nbroker:\n type: sqlite\n";
let edited = set_top_level_scalar(src, "version", "\"2.0.0\"").unwrap();
let out = edited;
assert_eq!(
out.matches("version:").count(),
1,
"no duplicate version line:\n{out}"
);
assert!(out.contains("version: \"2.0.0\""));
}
#[test]
fn set_top_level_scalar_errors_on_missing_key() {
let src = "broker:\n type: sqlite\n";
let err = set_top_level_scalar(src, "version", "\"2.0.0\"").expect_err("missing key");
assert!(
err.to_string().contains("no top-level key `version` found"),
"error must name the missing key: {err}"
);
}
#[test]
fn set_top_level_scalar_only_touches_top_level_key() {
let src = "\
version: 2
nested:
version: 99
other: ok
";
let edited = set_top_level_scalar(src, "version", "\"2.0.0\"").unwrap();
let out = edited;
assert!(
out.contains("version: \"2.0.0\""),
"top-level rewritten:\n{out}"
);
assert!(
out.contains(" version: 99"),
"nested version: 99 must be left alone:\n{out}"
);
}
}