1#[derive(Debug, Clone, PartialEq, Eq)]
4pub struct BudgetV1 {
5 pub max_tool_calls: u32,
6 pub max_retries: u32,
7 pub max_turn_millis: u64,
8}
9
10impl Default for BudgetV1 {
11 fn default() -> Self {
12 Self {
13 max_tool_calls: 8,
14 max_retries: 2,
15 max_turn_millis: 30_000,
16 }
17 }
18}
19
20impl BudgetV1 {
21 pub fn allows_tool_call(&self, tool_calls_so_far: u32, requested_tool_calls: u32) -> bool {
22 tool_calls_so_far.saturating_add(requested_tool_calls) <= self.max_tool_calls
23 }
24
25 pub fn allows_retry(&self, retries_so_far: u32) -> bool {
26 retries_so_far < self.max_retries
27 }
28
29 pub fn deadline_exceeded(&self, elapsed_millis: u64) -> bool {
30 elapsed_millis > self.max_turn_millis
31 }
32}
33
34#[cfg(test)]
35mod tests {
36 use super::*;
37
38 #[test]
39 fn tool_call_budget_is_cumulative() {
40 let budget = BudgetV1 {
41 max_tool_calls: 2,
42 max_retries: 1,
43 max_turn_millis: 1000,
44 };
45
46 assert!(budget.allows_tool_call(1, 1));
47 assert!(!budget.allows_tool_call(2, 1));
48 }
49
50 #[test]
51 fn deadline_check_is_explicit() {
52 let budget = BudgetV1 {
53 max_tool_calls: 1,
54 max_retries: 1,
55 max_turn_millis: 10,
56 };
57
58 assert!(!budget.deadline_exceeded(10));
59 assert!(budget.deadline_exceeded(11));
60 }
61}