1use std::fs;
38use std::path::Path;
39
40use anyhow::{anyhow, Context, Result};
41
42pub use yaml_edit::path::YamlPath;
43pub use yaml_edit::{Document, Mapping, Sequence};
44
45pub fn load(path: &Path) -> Result<Document> {
51 let raw = fs::read_to_string(path).with_context(|| format!("read {}", path.display()))?;
52 raw.parse::<Document>()
53 .with_context(|| format!("parse {}", path.display()))
54}
55
56pub fn save(doc: &Document, path: &Path) -> Result<()> {
62 fs::write(path, doc.to_string()).with_context(|| format!("write {}", path.display()))?;
63 Ok(())
64}
65
66pub fn set_nested_mapping(
84 doc: Document,
85 parent_path: &[&str],
86 value_pairs: &[(&str, &str)],
87) -> Result<Document> {
88 if parent_path.is_empty() {
89 return Err(anyhow!("set_nested_mapping: parent_path must not be empty"));
90 }
91 let source = doc.to_string();
92 let edited = splice_nested_mapping(&source, parent_path, value_pairs)?;
93 edited
94 .parse::<Document>()
95 .with_context(|| "re-parse spliced YAML")
96}
97
98fn splice_nested_mapping(
102 source: &str,
103 path: &[&str],
104 value_pairs: &[(&str, &str)],
105) -> Result<String> {
106 let lines: Vec<&str> = source.lines().collect();
107 let trailing_newline = source.ends_with('\n');
108
109 let mut current_indent: usize = 0;
112 let mut search_start: usize = 0;
113 let mut search_end: usize = lines.len();
114 let mut existing_depth: usize = 0;
115 let mut leaf_replace_range: Option<(usize, usize, usize)> = None; for (depth, key) in path.iter().enumerate() {
118 let parent_indent = current_indent;
119 let child_indent_min = if depth == 0 { 0 } else { parent_indent + 1 };
120 match find_key_in_block(&lines, search_start, search_end, key, child_indent_min) {
121 Some((line_idx, key_indent)) => {
122 existing_depth = depth + 1;
123 current_indent = key_indent;
124 let block_end = block_end_after(&lines, line_idx, key_indent);
125 if depth == path.len() - 1 {
126 leaf_replace_range = Some((line_idx, block_end, key_indent));
127 } else {
128 search_start = line_idx + 1;
129 search_end = block_end;
130 }
131 }
132 None => break,
133 }
134 }
135
136 if existing_depth == 0 {
137 return Err(anyhow!(
138 "set_nested_mapping: top-level key `{}` not found",
139 path[0]
140 ));
141 }
142
143 let insert_indent = if existing_depth == path.len() {
145 leaf_replace_range.expect("leaf existed").2
147 } else {
148 current_indent + 2
150 };
151
152 let missing_tail = &path[existing_depth..];
153 let mut block_lines: Vec<String> = Vec::new();
154 let mut indent = insert_indent;
155 for key in missing_tail {
156 block_lines.push(format!("{:indent$}{key}:", "", indent = indent, key = key));
157 indent += 2;
158 }
159 let value_indent = if existing_depth == path.len() {
171 insert_indent + 2
172 } else {
173 indent
174 };
175 if existing_depth == path.len() {
176 block_lines.push(format!(
178 "{:indent$}{key}:",
179 "",
180 indent = insert_indent,
181 key = path[path.len() - 1]
182 ));
183 }
184 for (k, v) in value_pairs {
185 block_lines.push(format!(
186 "{:indent$}{k}: {v}",
187 "",
188 indent = value_indent,
189 k = k,
190 v = v
191 ));
192 }
193
194 let mut out_lines: Vec<String> = lines.iter().map(|s| s.to_string()).collect();
197 if let Some((start, end, _)) = leaf_replace_range {
198 out_lines.splice(start..end, block_lines);
199 } else {
200 out_lines.splice(search_end..search_end, block_lines);
203 }
204
205 let mut joined = out_lines.join("\n");
206 if trailing_newline && !joined.ends_with('\n') {
207 joined.push('\n');
208 }
209 Ok(joined)
210}
211
212fn find_key_in_block(
217 lines: &[&str],
218 start: usize,
219 end: usize,
220 key: &str,
221 min_indent: usize,
222) -> Option<(usize, usize)> {
223 let mut child_indent: Option<usize> = None;
227 for line in lines.iter().take(end).skip(start) {
228 if let Some((indent, _)) = parse_mapping_key_line(line) {
229 if indent >= min_indent {
230 child_indent = Some(child_indent.map_or(indent, |c| c.min(indent)));
231 }
232 }
233 }
234 let child_indent = child_indent?;
235
236 for (i, line) in lines.iter().enumerate().take(end).skip(start) {
238 if let Some((indent, found_key)) = parse_mapping_key_line(line) {
239 if indent == child_indent && found_key == key {
240 return Some((i, indent));
241 }
242 }
243 }
244 None
245}
246
247fn parse_mapping_key_line(line: &str) -> Option<(usize, &str)> {
257 let indent = line.len() - line.trim_start().len();
258 let trimmed = line.trim_start();
259 if trimmed.is_empty() || trimmed.starts_with('#') || trimmed.starts_with('-') {
260 return None;
261 }
262 let colon_idx = trimmed.find(':')?;
263 let key = &trimmed[..colon_idx];
264 if key.is_empty() {
265 return None;
266 }
267 if key.contains(':') {
270 return None;
271 }
272 let after = &trimmed[colon_idx + 1..];
275 if !after.is_empty() && !after.starts_with(char::is_whitespace) {
276 return None;
278 }
279 Some((indent, key))
280}
281
282fn block_end_after(lines: &[&str], key_line: usize, key_indent: usize) -> usize {
288 for (i, line) in lines.iter().enumerate().skip(key_line + 1) {
289 let trimmed = line.trim_start();
290 if trimmed.is_empty() || trimmed.starts_with('#') {
291 continue;
292 }
293 let indent = line.len() - trimmed.len();
294 if indent <= key_indent {
295 return i;
296 }
297 }
298 lines.len()
299}
300
301#[cfg(test)]
302mod tests {
303 use super::*;
304
305 const COMMENTED_FIXTURE: &str = "\
306version: 2
307
308# managers block: each manager is a long-running agent.
309managers:
310 pm:
311 runtime: claude-code # canonical runtime
312 role_prompt: roles/pm.md
313 # interfaces lands here once `teamctl bot setup` runs
314 eng_lead:
315 runtime: claude-code
316 role_prompt: roles/eng_lead.md
317
318# trailing footer
319";
320
321 #[test]
322 fn round_trip_preserves_byte_for_byte() {
323 let dir = tempfile::tempdir().unwrap();
324 let path = dir.path().join("fixture.yaml");
325 fs::write(&path, COMMENTED_FIXTURE).unwrap();
326
327 let doc = load(&path).unwrap();
328 save(&doc, &path).unwrap();
329
330 let after = fs::read_to_string(&path).unwrap();
331 assert_eq!(
332 after, COMMENTED_FIXTURE,
333 "load → save without mutation must be byte-perfect"
334 );
335 }
336
337 #[test]
338 fn mutation_preserves_comments() {
339 let dir = tempfile::tempdir().unwrap();
340 let path = dir.path().join("fixture.yaml");
341 fs::write(&path, COMMENTED_FIXTURE).unwrap();
342
343 let doc = load(&path).unwrap();
344 let doc = set_nested_mapping(
345 doc,
346 &["managers", "pm", "interfaces", "telegram"],
347 &[("bot_token_env", "PM_TOKEN"), ("chat_ids_env", "PM_CHATS")],
348 )
349 .unwrap();
350 save(&doc, &path).unwrap();
351
352 let after = fs::read_to_string(&path).unwrap();
353
354 assert!(
355 after.contains("# managers block: each manager is a long-running agent."),
356 "block comment dropped:\n{after}"
357 );
358 assert!(
359 after.contains("# canonical runtime"),
360 "trailing line comment dropped:\n{after}"
361 );
362 assert!(
363 after.contains("# trailing footer"),
364 "footer comment dropped:\n{after}"
365 );
366 assert!(
367 after.contains(" interfaces:"),
368 "interfaces not properly indented under pm:\n{after}"
369 );
370 assert!(
371 after.contains(" telegram:"),
372 "telegram not properly indented under interfaces:\n{after}"
373 );
374 assert!(
375 after.contains(" bot_token_env: PM_TOKEN"),
376 "leaf not properly indented:\n{after}"
377 );
378 assert!(after.contains(" chat_ids_env: PM_CHATS"));
379
380 let pm_idx = after.find("pm:").expect("pm key");
382 let eng_idx = after.find("eng_lead:").expect("eng_lead key");
383 assert!(pm_idx < eng_idx, "manager key order swapped:\n{after}");
384
385 assert!(
387 after.contains("\n eng_lead:"),
388 "eng_lead boundary broken:\n{after}"
389 );
390 }
391
392 #[test]
396 fn save_does_not_strip_existing_comments() {
397 let dir = tempfile::tempdir().unwrap();
398 let path = dir.path().join("oss-shape.yaml");
399 let fixture = "\
400version: 2
401
402project:
403 id: oss
404 name: OSS Maintainer
405 cwd: ./workspace
406
407# Hub-and-spoke: maintainer is the only manager; workers fan out below.
408managers:
409 maintainer:
410 runtime: claude-code
411 role_prompt: roles/maintainer.md
412 # `teamctl bot setup` writes the interfaces.telegram block here.
413
414workers:
415 bug_fix:
416 runtime: claude-code # workers default to sonnet
417 reports_to: maintainer
418";
419 fs::write(&path, fixture).unwrap();
420
421 let doc = load(&path).unwrap();
422 let doc = set_nested_mapping(
423 doc,
424 &["managers", "maintainer", "interfaces", "telegram"],
425 &[
426 ("bot_token_env", "TEAMCTL_TG_MAINTAINER_TOKEN"),
427 ("chat_ids_env", "TEAMCTL_TG_MAINTAINER_CHATS"),
428 ],
429 )
430 .unwrap();
431 save(&doc, &path).unwrap();
432
433 let after = fs::read_to_string(&path).unwrap();
434 assert!(
435 after.contains(
436 "# Hub-and-spoke: maintainer is the only manager; workers fan out below."
437 ),
438 "block comment dropped — regression class still open:\n{after}"
439 );
440 assert!(
441 after.contains("# `teamctl bot setup` writes the interfaces.telegram block here."),
442 "inline comment dropped:\n{after}"
443 );
444 assert!(
445 after.contains("# workers default to sonnet"),
446 "trailing line comment dropped:\n{after}"
447 );
448 assert!(after.contains(" interfaces:"));
449 assert!(after.contains(" telegram:"));
450 assert!(after.contains(" bot_token_env: TEAMCTL_TG_MAINTAINER_TOKEN"));
451 assert!(after.contains(" chat_ids_env: TEAMCTL_TG_MAINTAINER_CHATS"));
452 }
453
454 #[test]
458 fn idempotent_replace_preserves_siblings() {
459 let dir = tempfile::tempdir().unwrap();
460 let path = dir.path().join("siblings.yaml");
461 let fixture = "\
462version: 2
463managers:
464 pm:
465 runtime: claude-code
466 interfaces:
467 discord:
468 bot_token_env: PM_DISCORD_TOKEN
469 telegram:
470 bot_token_env: OLD_TOKEN
471 chat_ids_env: OLD_CHATS
472";
473 fs::write(&path, fixture).unwrap();
474
475 let doc = load(&path).unwrap();
476 let doc = set_nested_mapping(
477 doc,
478 &["managers", "pm", "interfaces", "telegram"],
479 &[
480 ("bot_token_env", "NEW_TOKEN"),
481 ("chat_ids_env", "NEW_CHATS"),
482 ],
483 )
484 .unwrap();
485 save(&doc, &path).unwrap();
486
487 let after = fs::read_to_string(&path).unwrap();
488 assert_eq!(
489 after.matches("telegram:").count(),
490 1,
491 "duplicate telegram block:\n{after}"
492 );
493 assert_eq!(
494 after.matches("discord:").count(),
495 1,
496 "discord sibling lost:\n{after}"
497 );
498 assert!(
499 after.contains("PM_DISCORD_TOKEN"),
500 "discord adapter contents lost:\n{after}"
501 );
502 assert!(after.contains("NEW_TOKEN"));
503 assert!(after.contains("NEW_CHATS"));
504 assert!(!after.contains("OLD_TOKEN"));
505 assert!(!after.contains("OLD_CHATS"));
506 }
507
508 #[test]
518 fn replace_existing_leaf_nests_values_under_it() {
519 let dir = tempfile::tempdir().unwrap();
520 let path = dir.path().join("prewired.yaml");
521 let fixture = "\
524version: 2
525managers:
526 builder:
527 runtime: claude-code
528 interfaces:
529 telegram:
530 bot_token_env: TEAMCTL_TG_BUILDER_TOKEN
531 chat_ids_env: TEAMCTL_TG_BUILDER_CHATS
532";
533 fs::write(&path, fixture).unwrap();
534
535 let doc = load(&path).unwrap();
536 let doc = set_nested_mapping(
537 doc,
538 &["managers", "builder", "interfaces", "telegram"],
539 &[
540 ("bot_token_env", "TEAMCTL_TG_BUILDER_TOKEN"),
541 ("chat_ids_env", "TEAMCTL_TG_BUILDER_CHATS"),
542 ],
543 )
544 .unwrap();
545 save(&doc, &path).unwrap();
546
547 let after = fs::read_to_string(&path).unwrap();
548 let v: serde_yaml::Value = serde_yaml::from_str(&after)
549 .unwrap_or_else(|e| panic!("re-spliced YAML must parse: {e}\n{after}"));
550 let tg = &v["managers"]["builder"]["interfaces"]["telegram"];
551 assert!(
552 tg.is_mapping(),
553 "telegram leaf must remain a mapping, not be nulled by sibling pairs:\n{after}"
554 );
555 assert_eq!(
556 tg["bot_token_env"].as_str(),
557 Some("TEAMCTL_TG_BUILDER_TOKEN"),
558 "bot_token_env must be nested under telegram:\n{after}"
559 );
560 assert_eq!(
561 tg["chat_ids_env"].as_str(),
562 Some("TEAMCTL_TG_BUILDER_CHATS"),
563 "chat_ids_env must be nested under telegram:\n{after}"
564 );
565 }
566}