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_top_level_scalar(source: &str, key: &str, value: &str) -> Result<String> {
89 let trailing_newline = source.ends_with('\n');
90 let mut out_lines: Vec<String> = Vec::new();
91 let mut rewrote = false;
92 for line in source.lines() {
93 if !rewrote {
94 let trimmed = line.trim_start();
97 let indent = line.len() - trimmed.len();
98 if indent == 0 && !trimmed.is_empty() && !trimmed.starts_with('#') {
99 if let Some((found_key, _rest)) = trimmed.split_once(':') {
100 if found_key == key {
101 out_lines.push(format!("{key}: {value}"));
102 rewrote = true;
103 continue;
104 }
105 }
106 }
107 }
108 out_lines.push(line.to_string());
109 }
110 if !rewrote {
111 return Err(anyhow!(
112 "set_top_level_scalar: no top-level key `{key}` found in document"
113 ));
114 }
115 let mut joined = out_lines.join("\n");
116 if trailing_newline && !joined.ends_with('\n') {
117 joined.push('\n');
118 }
119 Ok(joined)
126}
127
128pub fn set_nested_mapping(
146 doc: Document,
147 parent_path: &[&str],
148 value_pairs: &[(&str, &str)],
149) -> Result<Document> {
150 if parent_path.is_empty() {
151 return Err(anyhow!("set_nested_mapping: parent_path must not be empty"));
152 }
153 let source = doc.to_string();
154 let edited = splice_nested_mapping(&source, parent_path, value_pairs)?;
155 edited
156 .parse::<Document>()
157 .with_context(|| "re-parse spliced YAML")
158}
159
160fn splice_nested_mapping(
164 source: &str,
165 path: &[&str],
166 value_pairs: &[(&str, &str)],
167) -> Result<String> {
168 let lines: Vec<&str> = source.lines().collect();
169 let trailing_newline = source.ends_with('\n');
170
171 let mut current_indent: usize = 0;
174 let mut search_start: usize = 0;
175 let mut search_end: usize = lines.len();
176 let mut existing_depth: usize = 0;
177 let mut leaf_replace_range: Option<(usize, usize, usize)> = None; for (depth, key) in path.iter().enumerate() {
180 let parent_indent = current_indent;
181 let child_indent_min = if depth == 0 { 0 } else { parent_indent + 1 };
182 match find_key_in_block(&lines, search_start, search_end, key, child_indent_min) {
183 Some((line_idx, key_indent)) => {
184 existing_depth = depth + 1;
185 current_indent = key_indent;
186 let block_end = block_end_after(&lines, line_idx, key_indent);
187 if depth == path.len() - 1 {
188 leaf_replace_range = Some((line_idx, block_end, key_indent));
189 } else {
190 search_start = line_idx + 1;
191 search_end = block_end;
192 }
193 }
194 None => break,
195 }
196 }
197
198 if existing_depth == 0 {
199 return Err(anyhow!(
200 "set_nested_mapping: top-level key `{}` not found",
201 path[0]
202 ));
203 }
204
205 let insert_indent = if existing_depth == path.len() {
207 leaf_replace_range.expect("leaf existed").2
209 } else {
210 current_indent + 2
212 };
213
214 let missing_tail = &path[existing_depth..];
215 let mut block_lines: Vec<String> = Vec::new();
216 let mut indent = insert_indent;
217 for key in missing_tail {
218 block_lines.push(format!("{:indent$}{key}:", "", indent = indent, key = key));
219 indent += 2;
220 }
221 let value_indent = if existing_depth == path.len() {
233 insert_indent + 2
234 } else {
235 indent
236 };
237 if existing_depth == path.len() {
238 block_lines.push(format!(
240 "{:indent$}{key}:",
241 "",
242 indent = insert_indent,
243 key = path[path.len() - 1]
244 ));
245 }
246 for (k, v) in value_pairs {
247 block_lines.push(format!(
248 "{:indent$}{k}: {v}",
249 "",
250 indent = value_indent,
251 k = k,
252 v = v
253 ));
254 }
255
256 let mut out_lines: Vec<String> = lines.iter().map(|s| s.to_string()).collect();
259 if let Some((start, end, _)) = leaf_replace_range {
260 out_lines.splice(start..end, block_lines);
261 } else {
262 out_lines.splice(search_end..search_end, block_lines);
265 }
266
267 let mut joined = out_lines.join("\n");
268 if trailing_newline && !joined.ends_with('\n') {
269 joined.push('\n');
270 }
271 Ok(joined)
272}
273
274fn find_key_in_block(
279 lines: &[&str],
280 start: usize,
281 end: usize,
282 key: &str,
283 min_indent: usize,
284) -> Option<(usize, usize)> {
285 let mut child_indent: Option<usize> = None;
289 for line in lines.iter().take(end).skip(start) {
290 if let Some((indent, _)) = parse_mapping_key_line(line) {
291 if indent >= min_indent {
292 child_indent = Some(child_indent.map_or(indent, |c| c.min(indent)));
293 }
294 }
295 }
296 let child_indent = child_indent?;
297
298 for (i, line) in lines.iter().enumerate().take(end).skip(start) {
300 if let Some((indent, found_key)) = parse_mapping_key_line(line) {
301 if indent == child_indent && found_key == key {
302 return Some((i, indent));
303 }
304 }
305 }
306 None
307}
308
309fn parse_mapping_key_line(line: &str) -> Option<(usize, &str)> {
319 let indent = line.len() - line.trim_start().len();
320 let trimmed = line.trim_start();
321 if trimmed.is_empty() || trimmed.starts_with('#') || trimmed.starts_with('-') {
322 return None;
323 }
324 let colon_idx = trimmed.find(':')?;
325 let key = &trimmed[..colon_idx];
326 if key.is_empty() {
327 return None;
328 }
329 if key.contains(':') {
332 return None;
333 }
334 let after = &trimmed[colon_idx + 1..];
337 if !after.is_empty() && !after.starts_with(char::is_whitespace) {
338 return None;
340 }
341 Some((indent, key))
342}
343
344fn block_end_after(lines: &[&str], key_line: usize, key_indent: usize) -> usize {
350 for (i, line) in lines.iter().enumerate().skip(key_line + 1) {
351 let trimmed = line.trim_start();
352 if trimmed.is_empty() || trimmed.starts_with('#') {
353 continue;
354 }
355 let indent = line.len() - trimmed.len();
356 if indent <= key_indent {
357 return i;
358 }
359 }
360 lines.len()
361}
362
363#[cfg(test)]
364mod tests {
365 use super::*;
366
367 const COMMENTED_FIXTURE: &str = "\
368version: 2
369
370# managers block: each manager is a long-running agent.
371managers:
372 pm:
373 runtime: claude-code # canonical runtime
374 role_prompt: roles/pm.md
375 # interfaces lands here once `teamctl bot setup` runs
376 eng_lead:
377 runtime: claude-code
378 role_prompt: roles/eng_lead.md
379
380# trailing footer
381";
382
383 #[test]
384 fn round_trip_preserves_byte_for_byte() {
385 let dir = tempfile::tempdir().unwrap();
386 let path = dir.path().join("fixture.yaml");
387 fs::write(&path, COMMENTED_FIXTURE).unwrap();
388
389 let doc = load(&path).unwrap();
390 save(&doc, &path).unwrap();
391
392 let after = fs::read_to_string(&path).unwrap();
393 assert_eq!(
394 after, COMMENTED_FIXTURE,
395 "load → save without mutation must be byte-perfect"
396 );
397 }
398
399 #[test]
400 fn mutation_preserves_comments() {
401 let dir = tempfile::tempdir().unwrap();
402 let path = dir.path().join("fixture.yaml");
403 fs::write(&path, COMMENTED_FIXTURE).unwrap();
404
405 let doc = load(&path).unwrap();
406 let doc = set_nested_mapping(
407 doc,
408 &["managers", "pm", "interfaces", "telegram"],
409 &[("bot_token_env", "PM_TOKEN"), ("chat_ids_env", "PM_CHATS")],
410 )
411 .unwrap();
412 save(&doc, &path).unwrap();
413
414 let after = fs::read_to_string(&path).unwrap();
415
416 assert!(
417 after.contains("# managers block: each manager is a long-running agent."),
418 "block comment dropped:\n{after}"
419 );
420 assert!(
421 after.contains("# canonical runtime"),
422 "trailing line comment dropped:\n{after}"
423 );
424 assert!(
425 after.contains("# trailing footer"),
426 "footer comment dropped:\n{after}"
427 );
428 assert!(
429 after.contains(" interfaces:"),
430 "interfaces not properly indented under pm:\n{after}"
431 );
432 assert!(
433 after.contains(" telegram:"),
434 "telegram not properly indented under interfaces:\n{after}"
435 );
436 assert!(
437 after.contains(" bot_token_env: PM_TOKEN"),
438 "leaf not properly indented:\n{after}"
439 );
440 assert!(after.contains(" chat_ids_env: PM_CHATS"));
441
442 let pm_idx = after.find("pm:").expect("pm key");
444 let eng_idx = after.find("eng_lead:").expect("eng_lead key");
445 assert!(pm_idx < eng_idx, "manager key order swapped:\n{after}");
446
447 assert!(
449 after.contains("\n eng_lead:"),
450 "eng_lead boundary broken:\n{after}"
451 );
452 }
453
454 #[test]
458 fn save_does_not_strip_existing_comments() {
459 let dir = tempfile::tempdir().unwrap();
460 let path = dir.path().join("oss-shape.yaml");
461 let fixture = "\
462version: 2
463
464project:
465 id: oss
466 name: OSS Maintainer
467 cwd: ./workspace
468
469# Hub-and-spoke: maintainer is the only manager; workers fan out below.
470managers:
471 maintainer:
472 runtime: claude-code
473 role_prompt: roles/maintainer.md
474 # `teamctl bot setup` writes the interfaces.telegram block here.
475
476workers:
477 bug_fix:
478 runtime: claude-code # workers default to sonnet
479 reports_to: maintainer
480";
481 fs::write(&path, fixture).unwrap();
482
483 let doc = load(&path).unwrap();
484 let doc = set_nested_mapping(
485 doc,
486 &["managers", "maintainer", "interfaces", "telegram"],
487 &[
488 ("bot_token_env", "TEAMCTL_TG_MAINTAINER_TOKEN"),
489 ("chat_ids_env", "TEAMCTL_TG_MAINTAINER_CHATS"),
490 ],
491 )
492 .unwrap();
493 save(&doc, &path).unwrap();
494
495 let after = fs::read_to_string(&path).unwrap();
496 assert!(
497 after.contains(
498 "# Hub-and-spoke: maintainer is the only manager; workers fan out below."
499 ),
500 "block comment dropped — regression class still open:\n{after}"
501 );
502 assert!(
503 after.contains("# `teamctl bot setup` writes the interfaces.telegram block here."),
504 "inline comment dropped:\n{after}"
505 );
506 assert!(
507 after.contains("# workers default to sonnet"),
508 "trailing line comment dropped:\n{after}"
509 );
510 assert!(after.contains(" interfaces:"));
511 assert!(after.contains(" telegram:"));
512 assert!(after.contains(" bot_token_env: TEAMCTL_TG_MAINTAINER_TOKEN"));
513 assert!(after.contains(" chat_ids_env: TEAMCTL_TG_MAINTAINER_CHATS"));
514 }
515
516 #[test]
520 fn idempotent_replace_preserves_siblings() {
521 let dir = tempfile::tempdir().unwrap();
522 let path = dir.path().join("siblings.yaml");
523 let fixture = "\
524version: 2
525managers:
526 pm:
527 runtime: claude-code
528 interfaces:
529 discord:
530 bot_token_env: PM_DISCORD_TOKEN
531 telegram:
532 bot_token_env: OLD_TOKEN
533 chat_ids_env: OLD_CHATS
534";
535 fs::write(&path, fixture).unwrap();
536
537 let doc = load(&path).unwrap();
538 let doc = set_nested_mapping(
539 doc,
540 &["managers", "pm", "interfaces", "telegram"],
541 &[
542 ("bot_token_env", "NEW_TOKEN"),
543 ("chat_ids_env", "NEW_CHATS"),
544 ],
545 )
546 .unwrap();
547 save(&doc, &path).unwrap();
548
549 let after = fs::read_to_string(&path).unwrap();
550 assert_eq!(
551 after.matches("telegram:").count(),
552 1,
553 "duplicate telegram block:\n{after}"
554 );
555 assert_eq!(
556 after.matches("discord:").count(),
557 1,
558 "discord sibling lost:\n{after}"
559 );
560 assert!(
561 after.contains("PM_DISCORD_TOKEN"),
562 "discord adapter contents lost:\n{after}"
563 );
564 assert!(after.contains("NEW_TOKEN"));
565 assert!(after.contains("NEW_CHATS"));
566 assert!(!after.contains("OLD_TOKEN"));
567 assert!(!after.contains("OLD_CHATS"));
568 }
569
570 #[test]
580 fn replace_existing_leaf_nests_values_under_it() {
581 let dir = tempfile::tempdir().unwrap();
582 let path = dir.path().join("prewired.yaml");
583 let fixture = "\
586version: 2
587managers:
588 builder:
589 runtime: claude-code
590 interfaces:
591 telegram:
592 bot_token_env: TEAMCTL_TG_BUILDER_TOKEN
593 chat_ids_env: TEAMCTL_TG_BUILDER_CHATS
594";
595 fs::write(&path, fixture).unwrap();
596
597 let doc = load(&path).unwrap();
598 let doc = set_nested_mapping(
599 doc,
600 &["managers", "builder", "interfaces", "telegram"],
601 &[
602 ("bot_token_env", "TEAMCTL_TG_BUILDER_TOKEN"),
603 ("chat_ids_env", "TEAMCTL_TG_BUILDER_CHATS"),
604 ],
605 )
606 .unwrap();
607 save(&doc, &path).unwrap();
608
609 let after = fs::read_to_string(&path).unwrap();
610 let v: serde_yaml::Value = serde_yaml::from_str(&after)
611 .unwrap_or_else(|e| panic!("re-spliced YAML must parse: {e}\n{after}"));
612 let tg = &v["managers"]["builder"]["interfaces"]["telegram"];
613 assert!(
614 tg.is_mapping(),
615 "telegram leaf must remain a mapping, not be nulled by sibling pairs:\n{after}"
616 );
617 assert_eq!(
618 tg["bot_token_env"].as_str(),
619 Some("TEAMCTL_TG_BUILDER_TOKEN"),
620 "bot_token_env must be nested under telegram:\n{after}"
621 );
622 assert_eq!(
623 tg["chat_ids_env"].as_str(),
624 Some("TEAMCTL_TG_BUILDER_CHATS"),
625 "chat_ids_env must be nested under telegram:\n{after}"
626 );
627 }
628
629 #[test]
632 fn set_top_level_scalar_replaces_integer_with_quoted_string() {
633 let src = "\
636# leading comment
637version: 2
638broker:
639 type: sqlite
640";
641 let edited = set_top_level_scalar(src, "version", "\"2.0.0\"").unwrap();
642 let out = edited;
643 assert!(
644 out.contains("version: \"2.0.0\""),
645 "rewrite missing:\n{out}"
646 );
647 assert!(
648 !out.contains("\nversion: 2\n"),
649 "old literal survived:\n{out}"
650 );
651 assert!(out.contains("# leading comment"));
653 assert!(out.contains("broker:"));
654 assert!(out.contains("type: sqlite"));
655 }
656
657 #[test]
658 fn set_top_level_scalar_is_idempotent() {
659 let src = "version: \"2.0.0\"\nbroker:\n type: sqlite\n";
660 let edited = set_top_level_scalar(src, "version", "\"2.0.0\"").unwrap();
661 let out = edited;
662 assert_eq!(
663 out.matches("version:").count(),
664 1,
665 "no duplicate version line:\n{out}"
666 );
667 assert!(out.contains("version: \"2.0.0\""));
668 }
669
670 #[test]
671 fn set_top_level_scalar_errors_on_missing_key() {
672 let src = "broker:\n type: sqlite\n";
673 let err = set_top_level_scalar(src, "version", "\"2.0.0\"").expect_err("missing key");
674 assert!(
675 err.to_string().contains("no top-level key `version` found"),
676 "error must name the missing key: {err}"
677 );
678 }
679
680 #[test]
681 fn set_top_level_scalar_only_touches_top_level_key() {
682 let src = "\
688version: 2
689nested:
690 version: 99
691 other: ok
692";
693 let edited = set_top_level_scalar(src, "version", "\"2.0.0\"").unwrap();
694 let out = edited;
695 assert!(
696 out.contains("version: \"2.0.0\""),
697 "top-level rewritten:\n{out}"
698 );
699 assert!(
700 out.contains(" version: 99"),
701 "nested version: 99 must be left alone:\n{out}"
702 );
703 }
704}