thndrs_agent/
budget.rs

1//! Provider-neutral tool-iteration budgeting for one agent turn.
2
3/// Decision returned by [`ToolIterationBudget`] before a provider request.
4#[derive(Clone, Debug, Eq, PartialEq)]
5pub enum ToolBudgetDecision {
6    /// The next provider request can proceed normally.
7    Continue,
8    /// The segment cap was reached and a provider-visible continuation
9    /// message should be appended before proceeding.
10    ContinueAfterBudgetMessage,
11    /// The full per-turn tool budget has been exhausted.
12    Exhausted {
13        /// Batches in the final segment.
14        segment_iterations: usize,
15        /// All tool batches run in this turn.
16        total_batches: usize,
17        /// Continuations already consumed.
18        continuations_used: usize,
19    },
20}
21
22/// Bounded tool-batch budget for one agent turn.
23#[derive(Clone, Debug, Eq, PartialEq)]
24pub struct ToolIterationBudget {
25    segment_limit: usize,
26    continuation_limit: usize,
27    segment_iterations: usize,
28    total_batches: usize,
29    continuations_used: usize,
30}
31
32impl ToolIterationBudget {
33    /// Create a budget with a segment limit and total continuation limit.
34    pub fn new(segment_limit: usize, continuation_limit: usize) -> Self {
35        Self { segment_limit, continuation_limit, segment_iterations: 0, total_batches: 0, continuations_used: 0 }
36    }
37
38    /// Record a completed provider-requested tool batch.
39    pub fn record_tool_batch(&mut self) {
40        self.segment_iterations = self.segment_iterations.saturating_add(1);
41        self.total_batches = self.total_batches.saturating_add(1);
42    }
43
44    /// Decide whether the next provider request fits the remaining budget.
45    pub fn before_provider_request(&mut self) -> ToolBudgetDecision {
46        if self.segment_iterations < self.segment_limit {
47            return ToolBudgetDecision::Continue;
48        }
49
50        if self.continuations_used < self.continuation_limit {
51            self.segment_iterations = 0;
52            self.continuations_used += 1;
53            return ToolBudgetDecision::ContinueAfterBudgetMessage;
54        }
55
56        ToolBudgetDecision::Exhausted {
57            segment_iterations: self.segment_iterations,
58            total_batches: self.total_batches,
59            continuations_used: self.continuations_used,
60        }
61    }
62
63    /// Return the total number of completed tool batches.
64    pub const fn total_batches(&self) -> usize {
65        self.total_batches
66    }
67
68    /// Return the number of consumed continuation segments.
69    pub const fn continuations_used(&self) -> usize {
70        self.continuations_used
71    }
72}
73
74#[cfg(test)]
75mod tests {
76    use super::*;
77
78    #[test]
79    fn budget_allows_bounded_continuations_before_exhausting() {
80        let mut budget = ToolIterationBudget::new(2, 1);
81        assert_eq!(budget.before_provider_request(), ToolBudgetDecision::Continue);
82        budget.record_tool_batch();
83        budget.record_tool_batch();
84        assert_eq!(
85            budget.before_provider_request(),
86            ToolBudgetDecision::ContinueAfterBudgetMessage
87        );
88        budget.record_tool_batch();
89        budget.record_tool_batch();
90        assert_eq!(
91            budget.before_provider_request(),
92            ToolBudgetDecision::Exhausted { segment_iterations: 2, total_batches: 4, continuations_used: 1 }
93        );
94    }
95}
thndrs_agent/budget.rs

thndrs_agent/
budget.rs
