| 1 | use std::time::Duration; |
| 2 | |
| 3 | use serde::{Deserialize, Serialize}; |
| 4 | |
| 5 | #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] |
| 6 | pub struct ExecutionBudget { |
| 7 | pub max_steps: u32, |
| 8 | pub max_tool_calls: u32, |
| 9 | pub max_retries: u32, |
| 10 | pub max_wall_time_ms: u64, |
| 11 | } |
| 12 | |
| 13 | #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] |
| 14 | pub struct BudgetUsage { |
| 15 | pub steps: u32, |
| 16 | pub tool_calls: u32, |
| 17 | pub retries: u32, |
| 18 | pub elapsed_ms: u64, |
| 19 | } |
| 20 | |
| 21 | #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] |
| 22 | #[serde(rename_all = "snake_case")] |
| 23 | pub enum BudgetExhaustion { |
| 24 | Steps, |
| 25 | ToolCalls, |
| 26 | Retries, |
| 27 | WallTime, |
| 28 | } |
| 29 | |
| 30 | impl ExecutionBudget { |
| 31 | pub fn validate(self) -> Result<(), String> { |
| 32 | if self.max_steps == 0 { |
| 33 | return Err("execution budget max_steps must be greater than zero".to_string()); |
| 34 | } |
| 35 | if self.max_tool_calls == 0 { |
| 36 | return Err("execution budget max_tool_calls must be greater than zero".to_string()); |
| 37 | } |
| 38 | if self.max_wall_time_ms == 0 { |
| 39 | return Err("execution budget max_wall_time_ms must be greater than zero".to_string()); |
| 40 | } |
| 41 | Ok(()) |
| 42 | } |
| 43 | |
| 44 | #[must_use] |
| 45 | pub const fn wall_time(self) -> Duration { |
| 46 | Duration::from_millis(self.max_wall_time_ms) |
| 47 | } |
| 48 | |
| 49 | #[must_use] |
| 50 | pub fn exhausted_by(self, usage: BudgetUsage) -> Option<BudgetExhaustion> { |
| 51 | if usage.steps >= self.max_steps { |
| 52 | Some(BudgetExhaustion::Steps) |
| 53 | } else if usage.tool_calls >= self.max_tool_calls { |
| 54 | Some(BudgetExhaustion::ToolCalls) |
| 55 | } else if usage.retries > self.max_retries { |
| 56 | Some(BudgetExhaustion::Retries) |
| 57 | } else if usage.elapsed_ms >= self.max_wall_time_ms { |
| 58 | Some(BudgetExhaustion::WallTime) |
| 59 | } else { |
| 60 | None |
| 61 | } |
| 62 | } |
| 63 | |
| 64 | /// Child work may narrow a parent budget but can never widen it. |
| 65 | #[must_use] |
| 66 | pub fn child_budget(self, requested: Self) -> Self { |
| 67 | Self { |
| 68 | max_steps: self.max_steps.min(requested.max_steps), |
| 69 | max_tool_calls: self.max_tool_calls.min(requested.max_tool_calls), |
| 70 | max_retries: self.max_retries.min(requested.max_retries), |
| 71 | max_wall_time_ms: self.max_wall_time_ms.min(requested.max_wall_time_ms), |
| 72 | } |
| 73 | } |
| 74 | } |
| 75 | |
| 76 | #[cfg(test)] |
| 77 | mod tests { |
| 78 | use super::*; |
| 79 | |
| 80 | #[test] |
| 81 | fn child_cannot_widen_parent_budget() { |
| 82 | let parent = ExecutionBudget { |
| 83 | max_steps: 32, |
| 84 | max_tool_calls: 64, |
| 85 | max_retries: 3, |
| 86 | max_wall_time_ms: 60_000, |
| 87 | }; |
| 88 | let child = parent.child_budget(ExecutionBudget { |
| 89 | max_steps: u32::MAX, |
| 90 | max_tool_calls: u32::MAX, |
| 91 | max_retries: u32::MAX, |
| 92 | max_wall_time_ms: u64::MAX, |
| 93 | }); |
| 94 | assert_eq!(child, parent); |
| 95 | } |
| 96 | |
| 97 | #[test] |
| 98 | fn wall_time_has_a_distinct_exhaustion_reason() { |
| 99 | let budget = ExecutionBudget { |
| 100 | max_steps: 100, |
| 101 | max_tool_calls: 100, |
| 102 | max_retries: 3, |
| 103 | max_wall_time_ms: 10, |
| 104 | }; |
| 105 | assert_eq!( |
| 106 | budget.exhausted_by(BudgetUsage { |
| 107 | steps: 1, |
| 108 | tool_calls: 1, |
| 109 | retries: 0, |
| 110 | elapsed_ms: 10, |
| 111 | }), |
| 112 | Some(BudgetExhaustion::WallTime) |
| 113 | ); |
| 114 | } |
| 115 | } |
| 116 |