harness_loop_engine/
budget.rs

1//! Per-loop token & iteration budgets.
2//!
3//! Loop engineering's hard-won lesson: *token consumption explodes with
4//! sub-agents and long-running loops*. A loop that runs every five minutes
5//! forever will, without a ceiling, spend without bound. `TokenBudget` is
6//! that ceiling, enforced per round; `BudgetState` accumulates spend across
7//! the maker and checker sub-agents and reports when a limit is crossed.
8
9use harness_core::Usage;
10use serde::{Deserialize, Serialize};
11
12/// A declarative spend ceiling for a single round of a loop.
13///
14/// `None` on a field means "no limit on this axis". `max_iters_per_round`
15/// caps how many tool-using iterations each sub-agent may take and is the
16/// one limit that is always present (sub-agents need a finite iteration
17/// budget regardless).
18#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
19pub struct TokenBudget {
20    /// Max input (prompt) tokens summed across all sub-agents in a round.
21    pub max_input_tokens: Option<u64>,
22    /// Max output (completion) tokens summed across the round.
23    pub max_output_tokens: Option<u64>,
24    /// Max total (input + output) tokens for the round.
25    pub max_total_tokens: Option<u64>,
26    /// Iteration cap handed to each sub-agent's loop.
27    pub max_iters_per_round: u32,
28}
29
30impl Default for TokenBudget {
31    fn default() -> Self {
32        Self {
33            max_input_tokens: None,
34            max_output_tokens: None,
35            max_total_tokens: None,
36            max_iters_per_round: 12,
37        }
38    }
39}
40
41impl TokenBudget {
42    /// A budget with only an iteration cap (no token ceilings).
43    pub fn iters(max_iters_per_round: u32) -> Self {
44        Self {
45            max_iters_per_round,
46            ..Default::default()
47        }
48    }
49
50    pub fn with_max_total_tokens(mut self, n: u64) -> Self {
51        self.max_total_tokens = Some(n);
52        self
53    }
54    pub fn with_max_input_tokens(mut self, n: u64) -> Self {
55        self.max_input_tokens = Some(n);
56        self
57    }
58    pub fn with_max_output_tokens(mut self, n: u64) -> Self {
59        self.max_output_tokens = Some(n);
60        self
61    }
62    pub fn with_max_iters_per_round(mut self, n: u32) -> Self {
63        self.max_iters_per_round = n;
64        self
65    }
66}
67
68/// Which ceiling a round crossed.
69#[derive(Debug, Clone, Copy, PartialEq, Eq)]
70pub enum BudgetLimit {
71    Input,
72    Output,
73    Total,
74}
75
76impl BudgetLimit {
77    pub fn label(self) -> &'static str {
78        match self {
79            BudgetLimit::Input => "input-tokens",
80            BudgetLimit::Output => "output-tokens",
81            BudgetLimit::Total => "total-tokens",
82        }
83    }
84}
85
86/// Running tally of spend within a round, checked against a [`TokenBudget`].
87#[derive(Debug, Clone, Copy)]
88pub struct BudgetState {
89    budget: TokenBudget,
90    pub input_tokens: u64,
91    pub output_tokens: u64,
92}
93
94impl BudgetState {
95    pub fn new(budget: TokenBudget) -> Self {
96        Self {
97            budget,
98            input_tokens: 0,
99            output_tokens: 0,
100        }
101    }
102
103    /// Fold one sub-agent's usage into the tally.
104    pub fn add(&mut self, usage: &Usage) {
105        self.input_tokens += usage.input_tokens as u64;
106        self.output_tokens += usage.output_tokens as u64;
107    }
108
109    pub fn total_tokens(&self) -> u64 {
110        self.input_tokens + self.output_tokens
111    }
112
113    /// The iteration cap each sub-agent should run under.
114    pub fn max_iters(&self) -> u32 {
115        self.budget.max_iters_per_round
116    }
117
118    /// Returns the first limit that has been crossed, if any. The engine
119    /// checks this after each sub-agent so it can stop before spawning the
120    /// next one.
121    pub fn exceeded(&self) -> Option<BudgetLimit> {
122        if let Some(m) = self.budget.max_input_tokens
123            && self.input_tokens > m
124        {
125            return Some(BudgetLimit::Input);
126        }
127        if let Some(m) = self.budget.max_output_tokens
128            && self.output_tokens > m
129        {
130            return Some(BudgetLimit::Output);
131        }
132        if let Some(m) = self.budget.max_total_tokens
133            && self.total_tokens() > m
134        {
135            return Some(BudgetLimit::Total);
136        }
137        None
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144
145    fn usage(input: u32, output: u32) -> Usage {
146        Usage {
147            input_tokens: input,
148            output_tokens: output,
149            cached_input_tokens: 0,
150        }
151    }
152
153    #[test]
154    fn no_limits_never_exceeds() {
155        let mut s = BudgetState::new(TokenBudget::iters(8));
156        s.add(&usage(1_000_000, 1_000_000));
157        assert!(s.exceeded().is_none());
158        assert_eq!(s.max_iters(), 8);
159    }
160
161    #[test]
162    fn total_limit_trips() {
163        let mut s = BudgetState::new(TokenBudget::iters(8).with_max_total_tokens(100));
164        s.add(&usage(60, 30)); // 90 — under
165        assert!(s.exceeded().is_none());
166        s.add(&usage(20, 0)); // 110 — over
167        assert_eq!(s.exceeded(), Some(BudgetLimit::Total));
168    }
169
170    #[test]
171    fn input_and_output_limits_trip_independently() {
172        let mut s = BudgetState::new(
173            TokenBudget::iters(8)
174                .with_max_input_tokens(50)
175                .with_max_output_tokens(50),
176        );
177        s.add(&usage(51, 1));
178        assert_eq!(s.exceeded(), Some(BudgetLimit::Input));
179
180        let mut s2 = BudgetState::new(TokenBudget::iters(8).with_max_output_tokens(50));
181        s2.add(&usage(1, 51));
182        assert_eq!(s2.exceeded(), Some(BudgetLimit::Output));
183    }
184}
harness_loop_engine/budget.rs

harness_loop_engine/
budget.rs
