agentd/supervisor/
budget.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//! Per-run budget: step / token / wall-clock bounds.
3//!
4//! This is the *per-run* budget enforced inside one agent loop (step / token /
5//! wall-clock). The *hierarchical* tree-token accounting — each subagent's
6//! usage rolled up to the tree root and bounded by a tree-wide ceiling — lives
7//! separately in `tree.rs` (`charge_tokens`) and is driven by the reactor
8//! (`KillReason::TreeBudget`); this type is the per-run primitive that records
9//! usage and answers "is a bound hit?".
10
11use crate::agentloop::stop::TerminalStatus;
12use crate::wire::intel::Usage;
13use std::time::Instant;
14
15#[derive(Debug)]
16pub struct Budget {
17    max_steps: u32,
18    max_tokens: u64,
19    deadline: Instant,
20    steps: u32,
21    tokens: u64,
22}
23
24impl Budget {
25    pub fn new(max_steps: u32, max_tokens: u64, deadline: Instant) -> Budget {
26        Budget {
27            max_steps,
28            max_tokens,
29            deadline,
30            steps: 0,
31            tokens: 0,
32        }
33    }
34
35    /// Count one completed loop turn.
36    pub fn record_step(&mut self) {
37        self.steps += 1;
38    }
39
40    /// Add a model call's token usage to the running total.
41    pub fn record_usage(&mut self, usage: Usage) {
42        self.tokens = self.tokens.saturating_add(usage.total());
43    }
44
45    pub fn tokens(&self) -> u64 {
46        self.tokens
47    }
48    pub fn steps(&self) -> u32 {
49        self.steps
50    }
51    /// The configured step ceiling (for an informational `loop.start` field).
52    pub fn max_steps(&self) -> u32 {
53        self.max_steps
54    }
55
56    /// The terminal status for whichever bound is hit, if any. Checked at the
57    /// top of each turn so the loop stops *before* spending more. Deadline is
58    /// checked first (it's the hardest bound).
59    pub fn exceeded(&self) -> Option<TerminalStatus> {
60        if Instant::now() >= self.deadline {
61            Some(TerminalStatus::Deadline)
62        } else if self.steps >= self.max_steps {
63            Some(TerminalStatus::ExhaustedSteps)
64        } else if self.tokens >= self.max_tokens {
65            Some(TerminalStatus::ExhaustedTokens)
66        } else {
67            None
68        }
69    }
70}
71
72#[cfg(test)]
73mod tests {
74    use super::*;
75    use std::time::Duration;
76
77    fn far() -> Instant {
78        Instant::now() + Duration::from_secs(3600)
79    }
80
81    #[test]
82    fn step_bound() {
83        let mut b = Budget::new(2, 1000, far());
84        assert!(b.exceeded().is_none());
85        b.record_step();
86        b.record_step();
87        assert_eq!(b.exceeded(), Some(TerminalStatus::ExhaustedSteps));
88    }
89
90    #[test]
91    fn token_bound() {
92        let mut b = Budget::new(100, 50, far());
93        b.record_usage(Usage {
94            input_tokens: 40,
95            output_tokens: 20,
96        });
97        assert_eq!(b.exceeded(), Some(TerminalStatus::ExhaustedTokens));
98    }
99
100    #[test]
101    fn deadline_bound() {
102        let b = Budget::new(100, 1000, Instant::now() - Duration::from_secs(1));
103        assert_eq!(b.exceeded(), Some(TerminalStatus::Deadline));
104    }
105}
agentd/supervisor/budget.rs

agentd/supervisor/
budget.rs
