lc_core/router_llm/
budget.rs

1//! Router-level spend/token budget circuit breaker (B13, 0.22.4).
2//!
3//! A [`RouterBudget`] is shared (`Arc`) by every call through one
4//! [`super::RouterLLM`]. Before a candidate model is attempted the router
5//! asks [`RouterBudget::precheck`] whether the **projected** spend (the
6//! prompt priced through the slot's [`ModelPrice`](crate::cost::ModelPrice))
7//! still fits; after a successful call it records the model-reported
8//! [`TokenUsage`](crate::language_models::TokenUsage), which is what actually
9//! trips the breaker once cumulative spend crosses the cap.
10//!
11//! Crucially, the breaker does not make the router fail: it makes the router
12//! **skip** the candidate that would exceed the budget and continue down the
13//! fallback chain. A slot without a price (a local / free model) projects
14//! zero cost, so after the paid tier trips, traffic rolls over to free
15//! fallbacks — the mainstream "cost guard" pattern. Only when every
16//! candidate is skipped does the caller see
17//! [`RouterError::BudgetExceeded`](super::RouterError::BudgetExceeded).
18//!
19//! All state lives in atomics, so the precheck is a cheap synchronous read
20//! usable on both the chat and streaming paths.
21
22use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
23use std::sync::Arc;
24
25/// Floating-point tolerance for "already exactly at the limit" comparisons.
26const COST_EPSILON: f64 = 1e-9;
27
28/// Which dimension of a [`RouterBudget`] was exceeded.
29#[derive(Debug, Clone, Copy, PartialEq, Eq)]
30pub enum BudgetKind {
31    /// Cumulative USD spend crossed the configured cost cap.
32    CostUsd,
33    /// Cumulative token usage crossed the configured token cap.
34    Tokens,
35}
36
37impl std::fmt::Display for BudgetKind {
38    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
39        match self {
40            BudgetKind::CostUsd => f.write_str("cost budget (USD)"),
41            BudgetKind::Tokens => f.write_str("token budget"),
42        }
43    }
44}
45
46/// A circuit breaker over cumulative spend and/or token usage.
47///
48/// Construct with one of the constructors, wrap in an `Arc`, and attach via
49/// [`super::RouterLLM::with_budget`]. Cloning is not supported; share by
50/// `Arc` so every call of one run observes the same totals.
51pub struct RouterBudget {
52    max_cost_usd: Option<f64>,
53    max_tokens: Option<u64>,
54    /// Cumulative spend as `f64` bits, updated with a CAS loop.
55    spent_bits: AtomicU64,
56    tokens: AtomicUsize,
57    trips: AtomicUsize,
58}
59
60impl std::fmt::Debug for RouterBudget {
61    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
62        f.debug_struct("RouterBudget")
63            .field("max_cost_usd", &self.max_cost_usd)
64            .field("max_tokens", &self.max_tokens)
65            .field("spent_usd", &self.spent_usd())
66            .field("tokens", &self.tokens())
67            .field("trips", &self.trips())
68            .finish()
69    }
70}
71
72impl RouterBudget {
73    /// Only a USD spend cap.
74    pub fn with_cost_limit(max_cost_usd: f64) -> Self {
75        Self::new(Some(max_cost_usd), None)
76    }
77
78    /// Only a cumulative token cap.
79    pub fn with_token_limit(max_tokens: u64) -> Self {
80        Self::new(None, Some(max_tokens))
81    }
82
83    /// Both a USD cap and a token cap; either one tripping skips the slot.
84    pub fn with_cost_and_token_limits(max_cost_usd: f64, max_tokens: u64) -> Arc<Self> {
85        Arc::new(Self::new(Some(max_cost_usd), Some(max_tokens)))
86    }
87
88    fn new(max_cost_usd: Option<f64>, max_tokens: Option<u64>) -> Self {
89        Self {
90            max_cost_usd,
91            max_tokens,
92            spent_bits: AtomicU64::new(0.0f64.to_bits()),
93            tokens: AtomicUsize::new(0),
94            trips: AtomicUsize::new(0),
95        }
96    }
97
98    /// Configured USD cap, if any.
99    pub fn max_cost_usd(&self) -> Option<f64> {
100        self.max_cost_usd
101    }
102
103    /// Configured token cap, if any.
104    pub fn max_tokens(&self) -> Option<u64> {
105        self.max_tokens
106    }
107
108    /// Cumulative recorded USD spend.
109    pub fn spent_usd(&self) -> f64 {
110        f64::from_bits(self.spent_bits.load(Ordering::Acquire))
111    }
112
113    /// Cumulative recorded token usage.
114    pub fn tokens(&self) -> u64 {
115        self.tokens.load(Ordering::Acquire) as u64
116    }
117
118    /// Number of times a precheck or a post-call record crossed a cap.
119    pub fn trips(&self) -> usize {
120        self.trips.load(Ordering::Acquire)
121    }
122
123    /// Whether either cap is already crossed.
124    pub fn is_tripped(&self) -> bool {
125        if let Some(limit) = self.max_cost_usd {
126            if self.spent_usd() > limit + COST_EPSILON {
127                return true;
128            }
129        }
130        if let Some(limit) = self.max_tokens {
131            if self.tokens() > limit {
132                return true;
133            }
134        }
135        false
136    }
137
138    /// Projects one pending call and returns the dimension that would be
139    /// exceeded, if any.
140    ///
141    /// `projected_cost_usd` is the prompt priced against the slot's price
142    /// (`0.0` for free/unpriced slots); `projected_tokens` is the prompt-token
143    /// estimate. A projected **zero-cost** call always passes the *cost*
144    /// dimension even while the breaker is tripped — that is what lets free
145    /// fallbacks keep serving. The token dimension is price-agnostic.
146    pub fn precheck(
147        &self,
148        projected_cost_usd: f64,
149        projected_tokens: u64,
150    ) -> Result<(), BudgetExceeded> {
151        if let Some(limit) = self.max_cost_usd {
152            if projected_cost_usd > 0.0
153                && self.spent_usd() + projected_cost_usd > limit + COST_EPSILON
154            {
155                self.trips.fetch_add(1, Ordering::AcqRel);
156                return Err(BudgetExceeded {
157                    kind: BudgetKind::CostUsd,
158                    used: self.spent_usd(),
159                    limit,
160                });
161            }
162        }
163        if let Some(limit) = self.max_tokens {
164            let used = self.tokens();
165            if used + projected_tokens > limit {
166                self.trips.fetch_add(1, Ordering::AcqRel);
167                return Err(BudgetExceeded {
168                    kind: BudgetKind::Tokens,
169                    used: used as f64,
170                    limit: limit as f64,
171                });
172            }
173        }
174        Ok(())
175    }
176
177    /// Records the measured cost/tokens of one finished call.
178    ///
179    /// Returns the first dimension this record pushed over its cap (the
180    /// breaker latches — later prechecks keep skipping paid candidates).
181    pub fn record(&self, cost_usd: f64, total_tokens: u64) -> Option<BudgetExceeded> {
182        if cost_usd != 0.0 {
183            let mut cur = self.spent_bits.load(Ordering::Acquire);
184            loop {
185                let next = f64::from_bits(cur) + cost_usd;
186                match self.spent_bits.compare_exchange(
187                    cur,
188                    next.to_bits(),
189                    Ordering::AcqRel,
190                    Ordering::Acquire,
191                ) {
192                    Ok(_) => break,
193                    Err(actual) => cur = actual,
194                }
195            }
196        }
197        if total_tokens != 0 {
198            self.tokens
199                .fetch_add(total_tokens as usize, Ordering::AcqRel);
200        }
201
202        if let Some(limit) = self.max_cost_usd {
203            if self.spent_usd() > limit + COST_EPSILON {
204                self.trips.fetch_add(1, Ordering::AcqRel);
205                return Some(BudgetExceeded {
206                    kind: BudgetKind::CostUsd,
207                    used: self.spent_usd(),
208                    limit,
209                });
210            }
211        }
212        if let Some(limit) = self.max_tokens {
213            if self.tokens() > limit {
214                self.trips.fetch_add(1, Ordering::AcqRel);
215                return Some(BudgetExceeded {
216                    kind: BudgetKind::Tokens,
217                    used: self.tokens() as f64,
218                    limit: limit as f64,
219                });
220            }
221        }
222        None
223    }
224
225    /// Zeroes accumulated spend/tokens/trip count (start a new run reusing
226    /// the same budget configuration).
227    pub fn reset(&self) {
228        self.spent_bits.store(0.0f64.to_bits(), Ordering::Release);
229        self.tokens.store(0, Ordering::Release);
230        self.trips.store(0, Ordering::Release);
231    }
232}
233
234/// Snapshot carried by [`RouterError::BudgetExceeded`](super::RouterError::BudgetExceeded).
235#[derive(Debug, Clone, Copy, PartialEq)]
236pub struct BudgetExceeded {
237    /// Which dimension tripped.
238    pub kind: BudgetKind,
239    /// Used amount at the moment of the trip (USD or tokens, per `kind`).
240    pub used: f64,
241    /// Configured limit in the same unit.
242    pub limit: f64,
243}
244
245impl std::fmt::Display for BudgetExceeded {
246    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
247        write!(
248            f,
249            "{} exceeded: used {:.6}, limit {:.6}",
250            self.kind, self.used, self.limit
251        )
252    }
253}
254
255impl std::error::Error for BudgetExceeded {}
256
257#[cfg(test)]
258mod tests {
259    use super::*;
260
261    #[test]
262    fn allows_within_cost_cap_and_blocks_projected_overrun() {
263        let b = RouterBudget::with_cost_limit(1.0);
264        b.precheck(0.4, 10).unwrap();
265        b.record(0.4, 10);
266        assert_eq!(b.spent_usd(), 0.4);
267        // 0.4 + 0.7 > 1.0 → blocked.
268        let e = b.precheck(0.7, 0).unwrap_err();
269        assert_eq!(e.kind, BudgetKind::CostUsd);
270        assert_eq!(e.used, 0.4);
271        assert_eq!(e.limit, 1.0);
272        assert!(b.trips() >= 1);
273    }
274
275    #[test]
276    fn free_call_passes_cost_dimension_even_when_tripped() {
277        let b = RouterBudget::with_cost_limit(1.0);
278        b.record(2.0, 0);
279        assert!(b.is_tripped());
280        // A zero-cost fallback must still be reachable.
281        b.precheck(0.0, 0).unwrap();
282        // But a paid call stays blocked.
283        assert!(b.precheck(0.01, 0).is_err());
284    }
285
286    #[test]
287    fn token_cap_counts_estimates_independent_of_price() {
288        let b = RouterBudget::with_token_limit(100);
289        b.precheck(0.0, 60).unwrap();
290        b.record(0.0, 60);
291        let e = b.precheck(0.0, 50).unwrap_err();
292        assert_eq!(e.kind, BudgetKind::Tokens);
293        assert_eq!(e.used, 60.0);
294        assert_eq!(e.limit, 100.0);
295    }
296
297    #[test]
298    fn record_latches_breaker_on_overshoot() {
299        let b = RouterBudget::with_cost_limit(1.0);
300        // Projected zero output cost, but the measured call overshoots.
301        b.precheck(0.9, 0).unwrap();
302        let trip = b.record(1.5, 100).expect("should trip on record");
303        assert_eq!(trip.kind, BudgetKind::CostUsd);
304        assert!(b.is_tripped());
305    }
306
307    #[test]
308    fn concurrent_record_sums_without_losing_updates() {
309        let b = Arc::new(RouterBudget::with_cost_limit(f64::INFINITY));
310        let mut handles = Vec::new();
311        for _ in 0..8 {
312            let b = b.clone();
313            handles.push(std::thread::spawn(move || {
314                for _ in 0..1000 {
315                    b.record(0.001, 1);
316                }
317            }));
318        }
319        for h in handles {
320            h.join().unwrap();
321        }
322        assert!((b.spent_usd() - 8.0).abs() < 1e-9);
323        assert_eq!(b.tokens(), 8000);
324    }
325
326    #[test]
327    fn reset_clears_totals_and_trip() {
328        let b = RouterBudget::with_cost_limit(1.0);
329        b.record(2.0, 50);
330        assert!(b.is_tripped());
331        b.reset();
332        assert!(!b.is_tripped());
333        assert_eq!(b.spent_usd(), 0.0);
334        assert_eq!(b.tokens(), 0);
335        assert_eq!(b.trips(), 0);
336    }
337}
lc_core/router_llm/budget.rs

lc_core/router_llm/
budget.rs
