Skip to main content

vtcode_core/prompts/
harness_limits.rs

1use crate::prompts::sections::{SectionBoundaryMode, find_prompt_section_bounds};
2use std::fmt::Write as _;
3
4/// Keep prompt guidance aligned with the selected user-visible harness limits.
5///
6/// This operation is idempotent: existing `[Harness Limits]` sections are removed before
7/// inserting the current values.
8pub fn upsert_harness_limits_section(
9    prompt: &mut String,
10    max_tool_calls_per_turn: usize,
11    max_tool_wall_clock_secs: u64,
12    max_tool_retries: u32,
13) {
14    let max_tool_calls_label = if max_tool_calls_per_turn == 0 {
15        "unlimited".to_string()
16    } else {
17        max_tool_calls_per_turn.to_string()
18    };
19
20    while let Some((section_start, section_end)) =
21        find_prompt_section_bounds(prompt, "[Harness Limits]", SectionBoundaryMode::BracketOnly)
22    {
23        prompt.replace_range(section_start..section_end, "");
24    }
25
26    while prompt.ends_with('\n') {
27        prompt.pop();
28    }
29
30    if prompt.is_empty() {
31        let _ = writeln!(
32            prompt,
33            "[Harness Limits]\n- max_tool_calls_per_turn: {}\n- max_tool_wall_clock_secs: {}\n- max_tool_retries: {}",
34            max_tool_calls_label, max_tool_wall_clock_secs, max_tool_retries
35        );
36    } else {
37        let _ = writeln!(
38            prompt,
39            "\n[Harness Limits]\n- max_tool_calls_per_turn: {}\n- max_tool_wall_clock_secs: {}\n- max_tool_retries: {}",
40            max_tool_calls_label, max_tool_wall_clock_secs, max_tool_retries
41        );
42    }
43}
44
45#[cfg(test)]
46mod tests {
47    use super::upsert_harness_limits_section;
48
49    #[test]
50    fn upsert_harness_limits_adds_single_section() {
51        let mut prompt = "Base prompt".to_string();
52
53        upsert_harness_limits_section(&mut prompt, 12, 180, 2);
54
55        assert_eq!(prompt.matches("[Harness Limits]").count(), 1);
56        assert!(prompt.contains("- max_tool_calls_per_turn: 12"));
57        assert!(prompt.contains("- max_tool_wall_clock_secs: 180"));
58        assert!(prompt.contains("- max_tool_retries: 2"));
59    }
60
61    #[test]
62    fn upsert_harness_limits_replaces_existing_values() {
63        let mut prompt = "Base prompt\n[Harness Limits]\n- max_tool_calls_per_turn: 3\n- max_tool_wall_clock_secs: 60\n- max_tool_retries: 1\n".to_string();
64
65        upsert_harness_limits_section(&mut prompt, 9, 240, 4);
66
67        assert_eq!(prompt.matches("[Harness Limits]").count(), 1);
68        assert!(prompt.contains("- max_tool_calls_per_turn: 9"));
69        assert!(prompt.contains("- max_tool_wall_clock_secs: 240"));
70        assert!(prompt.contains("- max_tool_retries: 4"));
71        assert!(!prompt.contains("- max_tool_calls_per_turn: 3"));
72    }
73
74    #[test]
75    fn upsert_harness_limits_preserves_trailing_prompt_sections() {
76        let mut prompt = "Base prompt\n[Harness Limits]\n- max_tool_calls_per_turn: 3\n- max_tool_wall_clock_secs: 60\n- max_tool_retries: 1\n[Additional Context]\nKeep this section".to_string();
77
78        upsert_harness_limits_section(&mut prompt, 11, 90, 3);
79
80        assert_eq!(prompt.matches("[Harness Limits]").count(), 1);
81        assert!(prompt.contains("[Additional Context]\nKeep this section"));
82        assert!(prompt.ends_with("- max_tool_retries: 3\n"));
83    }
84
85    #[test]
86    fn upsert_harness_limits_replaces_indented_section_header() {
87        let mut prompt = "Base prompt\n  [Harness Limits]\n- max_tool_calls_per_turn: 1\n- max_tool_wall_clock_secs: 1\n- max_tool_retries: 1\n".to_string();
88
89        upsert_harness_limits_section(&mut prompt, 5, 30, 2);
90
91        assert_eq!(prompt.matches("[Harness Limits]").count(), 1);
92        assert!(prompt.contains("- max_tool_calls_per_turn: 5"));
93        assert!(!prompt.contains("- max_tool_calls_per_turn: 1"));
94    }
95
96    #[test]
97    fn upsert_harness_limits_removes_duplicate_sections() {
98        let mut prompt = "Base prompt\n[Harness Limits]\n- max_tool_calls_per_turn: 2\n- max_tool_wall_clock_secs: 10\n- max_tool_retries: 1\n[Other]\nkeep\n[Harness Limits]\n- max_tool_calls_per_turn: 3\n- max_tool_wall_clock_secs: 20\n- max_tool_retries: 2\n".to_string();
99
100        upsert_harness_limits_section(&mut prompt, 7, 70, 3);
101
102        assert_eq!(prompt.matches("[Harness Limits]").count(), 1);
103        assert!(prompt.contains("- max_tool_calls_per_turn: 7"));
104        assert!(prompt.contains("[Other]\nkeep"));
105    }
106
107    #[test]
108    fn upsert_harness_limits_renders_unlimited_when_tool_cap_disabled() {
109        let mut prompt = "Base prompt".to_string();
110
111        upsert_harness_limits_section(&mut prompt, 0, 600, 2);
112
113        assert!(prompt.contains("- max_tool_calls_per_turn: unlimited"));
114        assert!(prompt.contains("- max_tool_wall_clock_secs: 600"));
115        assert!(prompt.contains("- max_tool_retries: 2"));
116    }
117}