agent_code_lib/services/
budget.rs

1//! Cost and token budget enforcement.
2//!
3//! Tracks spending against configurable limits and determines
4//! whether the agent should continue or stop.
5
6/// Budget check result.
7#[derive(Debug, Clone, PartialEq)]
8pub enum BudgetDecision {
9    /// Within budget, continue.
10    Continue,
11    /// Approaching limit, continue with notification.
12    ContinueWithWarning { percent_used: f64, message: String },
13    /// Budget exhausted, stop.
14    Stop { message: String },
15}
16
17/// Configuration for budget limits.
18#[derive(Debug, Clone)]
19pub struct BudgetConfig {
20    /// Maximum USD spend per session (None = unlimited).
21    pub max_cost_usd: Option<f64>,
22    /// Maximum total tokens per session (None = unlimited).
23    pub max_tokens: Option<u64>,
24    /// Warning threshold as fraction of budget (e.g., 0.8 = 80%).
25    pub warning_threshold: f64,
26}
27
28impl Default for BudgetConfig {
29    fn default() -> Self {
30        Self {
31            max_cost_usd: None,
32            max_tokens: None,
33            warning_threshold: 0.8,
34        }
35    }
36}
37
38/// Check whether the current spend is within budget.
39pub fn check_budget(
40    current_cost_usd: f64,
41    current_tokens: u64,
42    config: &BudgetConfig,
43) -> BudgetDecision {
44    // Check cost budget.
45    if let Some(max_cost) = config.max_cost_usd
46        && max_cost > 0.0
47    {
48        let ratio = current_cost_usd / max_cost;
49        if ratio >= 1.0 {
50            return BudgetDecision::Stop {
51                message: format!(
52                    "Cost budget exhausted: ${:.4} / ${:.4}",
53                    current_cost_usd, max_cost
54                ),
55            };
56        }
57        if ratio >= config.warning_threshold {
58            return BudgetDecision::ContinueWithWarning {
59                percent_used: ratio * 100.0,
60                message: format!("Cost at {:.0}% of ${:.4} budget", ratio * 100.0, max_cost),
61            };
62        }
63    }
64
65    // Check token budget.
66    if let Some(max_tokens) = config.max_tokens
67        && max_tokens > 0
68    {
69        let ratio = current_tokens as f64 / max_tokens as f64;
70        if ratio >= 1.0 {
71            return BudgetDecision::Stop {
72                message: format!(
73                    "Token budget exhausted: {} / {} tokens",
74                    current_tokens, max_tokens
75                ),
76            };
77        }
78        if ratio >= config.warning_threshold {
79            return BudgetDecision::ContinueWithWarning {
80                percent_used: ratio * 100.0,
81                message: format!("Tokens at {:.0}% of {} budget", ratio * 100.0, max_tokens),
82            };
83        }
84    }
85
86    BudgetDecision::Continue
87}
88
89/// Continuation check for token budget during multi-turn execution.
90///
91/// After each turn, decide whether to continue based on how many
92/// tokens were consumed and whether progress is diminishing.
93pub fn should_continue_turn(
94    turn_tokens: u64,
95    total_budget: Option<u64>,
96    consecutive_low_progress_turns: u32,
97) -> bool {
98    let Some(budget) = total_budget else {
99        return true; // No budget = always continue.
100    };
101
102    if budget == 0 {
103        return false;
104    }
105
106    // Stop at 90% of budget.
107    if turn_tokens >= (budget as f64 * 0.9) as u64 {
108        return false;
109    }
110
111    // Stop after 3 turns with minimal progress (< 500 tokens each).
112    if consecutive_low_progress_turns >= 3 {
113        return false;
114    }
115
116    true
117}
118
119#[cfg(test)]
120mod tests {
121    use super::*;
122
123    #[test]
124    fn test_within_budget() {
125        let config = BudgetConfig {
126            max_cost_usd: Some(1.0),
127            ..Default::default()
128        };
129        assert_eq!(check_budget(0.5, 0, &config), BudgetDecision::Continue);
130    }
131
132    #[test]
133    fn test_budget_warning() {
134        let config = BudgetConfig {
135            max_cost_usd: Some(1.0),
136            warning_threshold: 0.8,
137            ..Default::default()
138        };
139        match check_budget(0.85, 0, &config) {
140            BudgetDecision::ContinueWithWarning { .. } => {}
141            other => panic!("Expected warning, got: {other:?}"),
142        }
143    }
144
145    #[test]
146    fn test_budget_exhausted() {
147        let config = BudgetConfig {
148            max_cost_usd: Some(1.0),
149            ..Default::default()
150        };
151        match check_budget(1.5, 0, &config) {
152            BudgetDecision::Stop { .. } => {}
153            other => panic!("Expected stop, got: {other:?}"),
154        }
155    }
156
157    #[test]
158    fn test_continuation_logic() {
159        assert!(should_continue_turn(1000, Some(10000), 0));
160        assert!(!should_continue_turn(9500, Some(10000), 0));
161        assert!(!should_continue_turn(1000, Some(10000), 3));
162    }
163}
agent_code_lib/services/budget.rs

agent_code_lib/services/
budget.rs
