Skip to main content

vtcode_core/prompts/
sections.rs

1use std::fmt::Write as _;
2
3#[derive(Clone, Copy, Debug, PartialEq, Eq)]
4pub(crate) enum SectionBoundaryMode {
5    BracketOnly,
6    BracketOrMarkdown,
7}
8
9pub(crate) fn find_prompt_section_bounds(
10    prompt: &str,
11    section_header: &str,
12    boundary_mode: SectionBoundaryMode,
13) -> Option<(usize, usize)> {
14    fn is_section_header_line(line: &str, boundary_mode: SectionBoundaryMode) -> bool {
15        let trimmed = line.trim();
16        (trimmed.starts_with('[') && trimmed.ends_with(']'))
17            || matches!(boundary_mode, SectionBoundaryMode::BracketOrMarkdown)
18                && (trimmed.starts_with("# ") || trimmed.starts_with("## "))
19    }
20
21    let mut offset = 0usize;
22    let mut section_start = None;
23
24    for line in prompt.split_inclusive('\n') {
25        let trimmed = line.trim();
26        if section_start.is_none() && trimmed == section_header {
27            section_start = Some(offset);
28            offset += line.len();
29            continue;
30        }
31
32        if let Some(start) = section_start
33            && is_section_header_line(line, boundary_mode)
34        {
35            return Some((start, offset));
36        }
37
38        offset += line.len();
39    }
40
41    section_start.map(|start| (start, prompt.len()))
42}
43
44/// Canonical `"\n\n{block}"` append shared by prompt assemblers.
45///
46/// `harness_limits.rs` keeps its own single-`\n` `writeln!` append on purpose:
47/// its section is re-inserted after a full trim and must end with exactly one
48/// newline, so merging the two append shapes would change output.
49pub(crate) fn append_prompt_block(prompt: &mut String, block: &str) {
50    if block.is_empty() {
51        return;
52    }
53
54    if prompt.is_empty() {
55        prompt.push_str(block);
56    } else {
57        let _ = write!(prompt, "\n\n{block}");
58    }
59}
60
61/// Remove every section headed by `section_header`, stopping at the next
62/// header per `boundary_mode`. Callers keep their own mode so bracket-only
63/// sections (harness limits) and markdown-aware sections (active tools,
64/// deferred tools, catalog) preserve their contracts.
65pub(crate) fn remove_prompt_section(prompt: &mut String, section_header: &str, boundary_mode: SectionBoundaryMode) {
66    while let Some((section_start, section_end)) = find_prompt_section_bounds(prompt, section_header, boundary_mode) {
67        prompt.replace_range(section_start..section_end, "");
68    }
69}
70
71/// Strip trailing newlines before re-inserting a rebuilt section.
72pub(crate) fn trim_trailing_newlines(prompt: &mut String) {
73    while prompt.ends_with('\n') {
74        prompt.pop();
75    }
76}
77
78#[cfg(test)]
79mod tests {
80    use super::{SectionBoundaryMode, append_prompt_block, find_prompt_section_bounds, remove_prompt_section};
81
82    #[test]
83    fn markdown_boundary_mode_stops_before_next_markdown_heading() {
84        let prompt = "Base\n## Active Tools\n- a\n## Environment\n- b\n";
85        let bounds = find_prompt_section_bounds(prompt, "## Active Tools", SectionBoundaryMode::BracketOrMarkdown)
86            .expect("section bounds");
87
88        assert_eq!(&prompt[bounds.0..bounds.1], "## Active Tools\n- a\n");
89    }
90
91    #[test]
92    fn bracket_only_mode_ignores_markdown_headings() {
93        let prompt = "Base\n[Harness Limits]\n- a\n# Planning\n- b\n## Environment\n- c\n[Next Section]\n- d\n";
94        let bounds = find_prompt_section_bounds(prompt, "[Harness Limits]", SectionBoundaryMode::BracketOnly)
95            .expect("section bounds");
96
97        assert_eq!(&prompt[bounds.0..bounds.1], "[Harness Limits]\n- a\n# Planning\n- b\n## Environment\n- c\n");
98    }
99
100    #[test]
101    fn active_tools_boundary_preserves_following_planning_contract() {
102        let prompt = "Base\n## Active Tools\n- old\n# PLANNING WORKFLOW (READ-ONLY)\nKeep the plan.\n";
103        let (start, end) =
104            find_prompt_section_bounds(prompt, "## Active Tools", SectionBoundaryMode::BracketOrMarkdown)
105                .expect("section bounds");
106        let mut replaced = prompt.to_string();
107        replaced.replace_range(start..end, "## Active Tools\n- new\n");
108        assert_eq!(replaced, "Base\n## Active Tools\n- new\n# PLANNING WORKFLOW (READ-ONLY)\nKeep the plan.\n");
109    }
110
111    #[test]
112    fn shared_section_helpers_stay_idempotent() {
113        let mut prompt = String::new();
114        append_prompt_block(&mut prompt, "");
115        assert!(prompt.is_empty());
116
117        append_prompt_block(&mut prompt, "## Active Tools\n- a");
118        assert_eq!(prompt, "## Active Tools\n- a");
119
120        append_prompt_block(&mut prompt, "[Deferred Tools]\n- b");
121        assert_eq!(prompt, "## Active Tools\n- a\n\n[Deferred Tools]\n- b");
122
123        remove_prompt_section(&mut prompt, "## Active Tools", SectionBoundaryMode::BracketOrMarkdown);
124        assert_eq!(prompt, "[Deferred Tools]\n- b");
125
126        prompt.push_str("\n\n");
127        super::trim_trailing_newlines(&mut prompt);
128        assert_eq!(prompt, "[Deferred Tools]\n- b");
129    }
130}