gitdataai/lib/ai/agent/iteration_budget.rs
2026-05-30 01:38:40 +08:00

46 lines
997 B
Rust

#[derive(Clone, Debug)]
pub struct IterationBudget {
pub remaining: usize,
pub hard_limit: usize,
pub grace_call: bool,
pub consumed: usize,
}
impl IterationBudget {
pub fn new(limit: usize) -> Self {
Self {
remaining: limit,
hard_limit: limit,
grace_call: true,
consumed: 0,
}
}
pub fn can_continue(&self) -> bool {
self.remaining > 0 || (self.remaining == 0 && self.grace_call)
}
pub fn consume(&mut self) -> bool {
if self.remaining > 0 {
self.remaining -= 1;
self.consumed += 1;
true
} else if self.grace_call {
self.grace_call = false;
self.consumed += 1;
true
} else {
false
}
}
pub fn exhaust(&mut self) {
self.remaining = 0;
self.grace_call = false;
}
pub const fn total_consumed(&self) -> usize {
self.consumed
}
}