lc_agents/executor/
budget.rs

1// lc-agents/src/executor/budget.rs
2//! Budget gates (§4.2): `BudgetConfig` hard-limit configuration + `BudgetExceeded`
3//! over-limit details + three gate-check functions (shared by the invoke / stream paths).
4//!
5//! The `AgentExecutor` control loop is all-off by default (`None` field = unlimited) and
6//! existing behavior is unchanged; once enabled via `.with_budget(BudgetConfig { .. })`,
7//! hitting any limit returns [`super::AgentError::BudgetExceeded`], letting the caller
8//! distinguish a budget stop from the model not converging. The gate functions are called
9//! with the same semantics in `run_agent_loop_from` (invoke) and `stream`, so the two
10//! paths cannot diverge.
11
12use super::AgentError;
13use crate::metrics::AgentMetrics;
14use std::time::{Duration, Instant};
15
16/// Hard budget for the agent control loop. A `None` field means that item is unlimited.
17#[derive(Debug, Clone, Default)]
18pub struct BudgetConfig {
19    /// Cumulative tool-call cap (parallel calls included; stops when exceeded).
20    pub max_tool_calls: Option<usize>,
21    /// Cumulative LLM output-token cap (reads `AgentMetrics.total_tokens`; has no effect
22    /// when the agent does not report tokens).
23    pub max_tokens: Option<usize>,
24    /// Loop wall-clock cap (timed from `run_agent_loop`).
25    pub max_duration: Option<Duration>,
26    /// Iteration cap (tightens `AgentExecutor::max_iterations`; hitting it returns an error
27    /// instead of the placeholder return path used at the iteration limit).
28    pub max_iterations: Option<usize>,
29}
30
31/// Details of a budget over-limit.
32#[derive(Debug, Clone)]
33pub enum BudgetExceeded {
34    /// Cumulative tool-call count exceeded.
35    ToolCalls {
36        /// Configured limit.
37        limit: usize,
38        /// Actual cumulative count at trigger time.
39        actual: usize,
40    },
41    /// Cumulative LLM output-token count exceeded.
42    Tokens {
43        /// Configured limit.
44        limit: usize,
45        /// Actual cumulative tokens at trigger time.
46        actual: usize,
47    },
48    /// Loop wall-clock duration exceeded.
49    Duration {
50        /// Configured limit.
51        limit: Duration,
52        /// Actual elapsed time at trigger time.
53        elapsed: Duration,
54    },
55    /// Iteration count exceeded.
56    Iterations {
57        /// Effective limit (already `min`'d with `AgentExecutor::max_iterations`).
58        limit: usize,
59    },
60}
61
62/// Budget gate (§4.2): iteration-level check (iteration count + wall-clock). Returns an
63/// error when a limit is exceeded.
64///
65/// `max_iterations` uses `min(self.max_iterations, budget.max_iterations)` as the
66/// effective limit — when the budget is tighter than the default it hard-stops on
67/// exceeding; when looser, `max_iterations` backs it up without changing the original
68/// placeholder return path. Shared by invoke / stream.
69pub(crate) fn budget_iteration_gate(
70    budget: Option<&BudgetConfig>,
71    max_iterations: usize,
72    iteration: usize,
73    loop_start: Instant,
74) -> Option<AgentError> {
75    let budget = budget?;
76    if let Some(limit) = budget.max_iterations {
77        let effective = limit.min(max_iterations);
78        if iteration >= effective {
79            return Some(AgentError::BudgetExceeded(BudgetExceeded::Iterations {
80                limit: effective,
81            }));
82        }
83    }
84    if let Some(limit) = budget.max_duration {
85        let elapsed = loop_start.elapsed();
86        if elapsed >= limit {
87            return Some(AgentError::BudgetExceeded(BudgetExceeded::Duration {
88                limit,
89                elapsed,
90            }));
91        }
92    }
93    None
94}
95
96/// Budget gate (§4.2): cumulative-token check after an LLM call. No effect when the
97/// agent does not report tokens.
98pub(crate) fn budget_token_gate(
99    budget: Option<&BudgetConfig>,
100    metrics: &AgentMetrics,
101) -> Option<AgentError> {
102    let budget = budget?;
103    let limit = budget.max_tokens?;
104    let actual = metrics.total_tokens.unwrap_or(0);
105    if actual >= limit {
106        return Some(AgentError::BudgetExceeded(BudgetExceeded::Tokens {
107            limit,
108            actual,
109        }));
110    }
111    None
112}
113
114/// Budget gate (§4.2): checks cumulative call count and wall-clock before a tool runs.
115///
116/// `metrics.tool_calls` is already incremented, so the check uses `> limit` — allowing
117/// exactly `limit` tool executions, with the `limit + 1`-th triggering the hard stop.
118pub(crate) fn budget_tool_gate(
119    budget: Option<&BudgetConfig>,
120    metrics: &AgentMetrics,
121    loop_start: Instant,
122) -> Option<AgentError> {
123    let budget = budget?;
124    if let Some(limit) = budget.max_tool_calls {
125        if metrics.tool_calls > limit {
126            return Some(AgentError::BudgetExceeded(BudgetExceeded::ToolCalls {
127                limit,
128                actual: metrics.tool_calls,
129            }));
130        }
131    }
132    if let Some(limit) = budget.max_duration {
133        let elapsed = loop_start.elapsed();
134        if elapsed >= limit {
135            return Some(AgentError::BudgetExceeded(BudgetExceeded::Duration {
136                limit,
137                elapsed,
138            }));
139        }
140    }
141    None
142}
lc_agents/executor/budget.rs

lc_agents/executor/
budget.rs
