codetether_agent/session/helper/
cost_guard.rs

1//! Cost-budget enforcement for the agentic loop.
2//!
3//! Checked before every provider request. Reads [`CostGuardrails`] from
4//! env vars (`CODETETHER_COST_WARN_USD`, `CODETETHER_COST_LIMIT_USD`)
5//! and compares against the running session cost estimate from
6//! [`crate::provider::pricing::session_cost_usd`].
7
8use std::sync::atomic::{AtomicBool, Ordering};
9
10use crate::config::guardrails::CostGuardrails;
11use crate::provider::pricing::session_cost_usd;
12
13/// One-shot latch so we don't spam the log every loop iteration once the
14/// warn threshold is crossed.
15static WARNED: AtomicBool = AtomicBool::new(false);
16
17/// Result of a budget check.
18#[derive(Debug, Clone, Copy)]
19pub enum CostGuardStatus {
20    Ok,
21    /// Warn threshold crossed this call (one-shot).
22    Warned {
23        spent_usd: f64,
24        threshold_usd: f64,
25    },
26    /// Hard limit crossed.
27    Block {
28        spent_usd: f64,
29        limit_usd: f64,
30    },
31}
32
33fn check() -> CostGuardStatus {
34    let g = CostGuardrails::from_env();
35    let spent = session_cost_usd();
36
37    if let Some(limit) = g.hard_limit_usd
38        && spent >= limit
39    {
40        return CostGuardStatus::Block {
41            spent_usd: spent,
42            limit_usd: limit,
43        };
44    }
45    if let Some(warn) = g.warn_usd
46        && spent >= warn
47        && !WARNED.swap(true, Ordering::Relaxed)
48    {
49        return CostGuardStatus::Warned {
50            spent_usd: spent,
51            threshold_usd: warn,
52        };
53    }
54    CostGuardStatus::Ok
55}
56
57/// Enforce the cost budget before sending a provider request.
58///
59/// Returns `Err` if the hard limit has been reached. Logs a one-shot
60/// warning the first time the warn threshold is crossed. Returns `Ok(())`
61/// in the no-limits-configured case.
62pub fn enforce_cost_budget() -> anyhow::Result<()> {
63    match check() {
64        CostGuardStatus::Block {
65            spent_usd,
66            limit_usd,
67        } => {
68            anyhow::bail!(
69                "Cost guardrail tripped: session has spent ~${:.2} which meets/exceeds the \
70                 hard limit of ${:.2}. Raise CODETETHER_COST_LIMIT_USD (or \
71                 `[guardrails] hard_limit_usd` in config) to continue.",
72                spent_usd,
73                limit_usd
74            )
75        }
76        CostGuardStatus::Warned {
77            spent_usd,
78            threshold_usd,
79        } => {
80            tracing::warn!(
81                spent_usd,
82                threshold_usd,
83                "Cost guardrail warn threshold reached; set CODETETHER_COST_LIMIT_USD to cap spend"
84            );
85            Ok(())
86        }
87        CostGuardStatus::Ok => Ok(()),
88    }
89}
90
91/// Non-mutating probe used by UI widgets to color the cost badge.
92pub fn cost_guard_level() -> CostGuardLevel {
93    let g = CostGuardrails::from_env();
94    let spent = session_cost_usd();
95    if let Some(limit) = g.hard_limit_usd
96        && spent >= limit
97    {
98        return CostGuardLevel::OverLimit;
99    }
100    if let Some(warn) = g.warn_usd
101        && spent >= warn
102    {
103        return CostGuardLevel::OverWarn;
104    }
105    CostGuardLevel::Ok
106}
107
108/// UI-facing severity of the current session's spend.
109#[derive(Debug, Clone, Copy, PartialEq, Eq)]
110pub enum CostGuardLevel {
111    Ok,
112    OverWarn,
113    OverLimit,
114}
codetether_agent/session/helper/cost_guard.rs

codetether_agent/session/helper/
cost_guard.rs
