ironflow_engine/
budget.rs

1//! Cost guardrails: per-run cost cap and global monthly quota.
2//!
3//! `max_budget_usd` on an [`AgentConfig`](ironflow_core::provider::AgentConfig)
4//! only caps a single agent invocation. A workflow with a loop, a dynamic
5//! `ctx.parallel()`, or a chain of sub-workflows can chain dozens of agent
6//! calls with no cumulative bound. [`BudgetConfig`] closes that gap:
7//!
8//! - **Per-run cap** -- checked before every agent step. Crossing it cancels
9//!   the run *before* the step is launched, so no spend happens.
10//! - **Monthly quota** -- checked when a new run is created. Crossing it
11//!   refuses the creation; runs already in flight are untouched.
12//!
13//! # Examples
14//!
15//! ```
16//! use ironflow_engine::budget::BudgetConfig;
17//! use rust_decimal::Decimal;
18//!
19//! let config = BudgetConfig::new()
20//!     .default_run_max_cost_usd(Decimal::new(500, 2))
21//!     .monthly_cost_limit_usd(Decimal::new(20000, 2));
22//!
23//! assert_eq!(config.default_run_max_cost_usd, Some(Decimal::new(500, 2)));
24//! ```
25
26use std::env;
27use std::str::FromStr;
28
29use chrono::{DateTime, Datelike, TimeZone, Utc};
30use rust_decimal::Decimal;
31use rust_decimal::prelude::FromPrimitive;
32use tracing::warn;
33
34/// Environment variable holding the server-wide default per-run cost cap.
35pub const DEFAULT_RUN_MAX_COST_ENV: &str = "IRONFLOW_DEFAULT_RUN_MAX_COST_USD";
36
37/// Environment variable holding the global monthly cost quota.
38pub const MONTHLY_COST_LIMIT_ENV: &str = "IRONFLOW_MONTHLY_COST_LIMIT_USD";
39
40/// Server-level cost guardrails.
41///
42/// Both fields default to `None`, which disables the corresponding check and
43/// preserves the pre-existing behaviour exactly.
44///
45/// # Examples
46///
47/// ```
48/// use ironflow_engine::budget::BudgetConfig;
49///
50/// let config = BudgetConfig::new();
51/// assert!(config.default_run_max_cost_usd.is_none());
52/// assert!(config.monthly_cost_limit_usd.is_none());
53/// ```
54#[derive(Debug, Clone, Default, PartialEq, Eq)]
55pub struct BudgetConfig {
56    /// Default cost cap applied to a run when neither the creation request nor
57    /// the handler declares one. `None` means no default cap.
58    pub default_run_max_cost_usd: Option<Decimal>,
59    /// Global quota for the current calendar month (UTC). `None` disables the
60    /// monthly check.
61    pub monthly_cost_limit_usd: Option<Decimal>,
62}
63
64impl BudgetConfig {
65    /// Create a configuration with both guardrails disabled.
66    ///
67    /// # Examples
68    ///
69    /// ```
70    /// use ironflow_engine::budget::BudgetConfig;
71    ///
72    /// assert_eq!(BudgetConfig::new(), BudgetConfig::default());
73    /// ```
74    pub fn new() -> Self {
75        Self::default()
76    }
77
78    /// Set the server-wide default per-run cost cap.
79    ///
80    /// # Examples
81    ///
82    /// ```
83    /// use ironflow_engine::budget::BudgetConfig;
84    /// use rust_decimal::Decimal;
85    ///
86    /// let config = BudgetConfig::new().default_run_max_cost_usd(Decimal::new(150, 2));
87    /// assert_eq!(config.default_run_max_cost_usd, Some(Decimal::new(150, 2)));
88    /// ```
89    pub fn default_run_max_cost_usd(mut self, cap: Decimal) -> Self {
90        self.default_run_max_cost_usd = Some(cap);
91        self
92    }
93
94    /// Set the global monthly cost quota.
95    ///
96    /// # Examples
97    ///
98    /// ```
99    /// use ironflow_engine::budget::BudgetConfig;
100    /// use rust_decimal::Decimal;
101    ///
102    /// let config = BudgetConfig::new().monthly_cost_limit_usd(Decimal::new(10000, 2));
103    /// assert_eq!(config.monthly_cost_limit_usd, Some(Decimal::new(10000, 2)));
104    /// ```
105    pub fn monthly_cost_limit_usd(mut self, limit: Decimal) -> Self {
106        self.monthly_cost_limit_usd = Some(limit);
107        self
108    }
109
110    /// Load the configuration from the environment.
111    ///
112    /// Reads [`DEFAULT_RUN_MAX_COST_ENV`] and [`MONTHLY_COST_LIMIT_ENV`]. A
113    /// variable that is unset, empty, unparseable, or negative is ignored (the
114    /// corresponding guardrail stays disabled) and logged at `WARN`. Loading
115    /// never fails, so a malformed value can never prevent the server from
116    /// starting.
117    ///
118    /// # Examples
119    ///
120    /// ```
121    /// use ironflow_engine::budget::BudgetConfig;
122    ///
123    /// let config = BudgetConfig::from_env();
124    /// # let _ = config;
125    /// ```
126    pub fn from_env() -> Self {
127        Self {
128            default_run_max_cost_usd: read_decimal_env(DEFAULT_RUN_MAX_COST_ENV),
129            monthly_cost_limit_usd: read_decimal_env(MONTHLY_COST_LIMIT_ENV),
130        }
131    }
132
133    /// Resolve the cost cap of a run being created.
134    ///
135    /// Priority, strongest first: the value supplied at creation time, then the
136    /// handler default, then the server default. `None` at every level means
137    /// the run has no cap.
138    ///
139    /// # Examples
140    ///
141    /// ```
142    /// use ironflow_engine::budget::BudgetConfig;
143    /// use rust_decimal::Decimal;
144    ///
145    /// let server = Decimal::new(100, 2);
146    /// let handler = Decimal::new(200, 2);
147    /// let requested = Decimal::new(300, 2);
148    /// let config = BudgetConfig::new().default_run_max_cost_usd(server);
149    ///
150    /// assert_eq!(config.resolve_run_cap(Some(requested), Some(handler)), Some(requested));
151    /// assert_eq!(config.resolve_run_cap(None, Some(handler)), Some(handler));
152    /// assert_eq!(config.resolve_run_cap(None, None), Some(server));
153    /// ```
154    pub fn resolve_run_cap(
155        &self,
156        requested: Option<Decimal>,
157        handler_default: Option<Decimal>,
158    ) -> Option<Decimal> {
159        requested
160            .or(handler_default)
161            .or(self.default_run_max_cost_usd)
162    }
163}
164
165/// Read a non-negative [`Decimal`] from an environment variable.
166///
167/// Returns `None` when the variable is unset, empty, unparseable, or negative.
168fn read_decimal_env(name: &str) -> Option<Decimal> {
169    let raw = env::var(name).ok()?;
170    let trimmed = raw.trim();
171    if trimmed.is_empty() {
172        return None;
173    }
174
175    match Decimal::from_str(trimmed) {
176        Ok(value) if value >= Decimal::ZERO => Some(value),
177        Ok(value) => {
178            warn!(env = name, value = %value, "ignoring negative budget limit");
179            None
180        }
181        Err(e) => {
182            warn!(env = name, value = trimmed, error = %e, "ignoring unparseable budget limit");
183            None
184        }
185    }
186}
187
188/// Convert an agent step budget expressed as `f64` into a [`Decimal`].
189///
190/// A missing budget counts as zero: the run cap check then reduces to
191/// "has the run already spent more than its cap?". A `NaN` or infinite value
192/// also counts as zero rather than poisoning the comparison.
193///
194/// # Examples
195///
196/// ```
197/// use ironflow_engine::budget::step_budget_usd;
198/// use rust_decimal::Decimal;
199///
200/// assert_eq!(step_budget_usd(Some(0.25)), Decimal::new(25, 2));
201/// assert_eq!(step_budget_usd(None), Decimal::ZERO);
202/// assert_eq!(step_budget_usd(Some(f64::NAN)), Decimal::ZERO);
203/// ```
204pub fn step_budget_usd(max_budget_usd: Option<f64>) -> Decimal {
205    max_budget_usd
206        .and_then(Decimal::from_f64)
207        .unwrap_or(Decimal::ZERO)
208}
209
210/// Start of the current calendar month, at 00:00:00 UTC.
211///
212/// Used as the lower bound of the monthly quota window.
213///
214/// # Examples
215///
216/// ```
217/// use chrono::{Datelike, TimeZone, Timelike, Utc};
218/// use ironflow_engine::budget::month_start;
219///
220/// let start = month_start(Utc.with_ymd_and_hms(2026, 7, 26, 15, 30, 0).unwrap());
221/// assert_eq!(start.year(), 2026);
222/// assert_eq!(start.month(), 7);
223/// assert_eq!(start.day(), 1);
224/// assert_eq!(start.hour(), 0);
225/// ```
226///
227/// # Panics
228///
229/// Never panics for a valid [`DateTime<Utc>`]: day 1 at midnight always exists
230/// for any year/month pair reachable from an existing timestamp.
231pub fn month_start(now: DateTime<Utc>) -> DateTime<Utc> {
232    Utc.with_ymd_and_hms(now.year(), now.month(), 1, 0, 0, 0)
233        .single()
234        .expect("first day of month at midnight UTC is always unambiguous")
235}
236
237#[cfg(test)]
238mod tests {
239    use chrono::Timelike;
240
241    use super::*;
242
243    #[test]
244    fn new_disables_both_guardrails() {
245        let config = BudgetConfig::new();
246        assert!(config.default_run_max_cost_usd.is_none());
247        assert!(config.monthly_cost_limit_usd.is_none());
248    }
249
250    #[test]
251    fn resolve_run_cap_prefers_request_then_handler_then_server() {
252        let config = BudgetConfig::new().default_run_max_cost_usd(Decimal::ONE);
253
254        assert_eq!(
255            config.resolve_run_cap(Some(Decimal::TEN), Some(Decimal::TWO)),
256            Some(Decimal::TEN)
257        );
258        assert_eq!(
259            config.resolve_run_cap(None, Some(Decimal::TWO)),
260            Some(Decimal::TWO)
261        );
262        assert_eq!(config.resolve_run_cap(None, None), Some(Decimal::ONE));
263    }
264
265    #[test]
266    fn resolve_run_cap_without_server_default_is_none() {
267        let config = BudgetConfig::new();
268        assert_eq!(config.resolve_run_cap(None, None), None);
269    }
270
271    #[test]
272    fn resolve_run_cap_accepts_explicit_zero() {
273        let config = BudgetConfig::new().default_run_max_cost_usd(Decimal::TEN);
274        assert_eq!(
275            config.resolve_run_cap(Some(Decimal::ZERO), None),
276            Some(Decimal::ZERO)
277        );
278    }
279
280    #[test]
281    fn step_budget_usd_maps_missing_and_invalid_to_zero() {
282        assert_eq!(step_budget_usd(None), Decimal::ZERO);
283        assert_eq!(step_budget_usd(Some(f64::NAN)), Decimal::ZERO);
284        assert_eq!(step_budget_usd(Some(f64::INFINITY)), Decimal::ZERO);
285        assert_eq!(step_budget_usd(Some(f64::NEG_INFINITY)), Decimal::ZERO);
286    }
287
288    #[test]
289    fn step_budget_usd_converts_finite_values() {
290        assert_eq!(step_budget_usd(Some(0.0)), Decimal::ZERO);
291        assert_eq!(step_budget_usd(Some(0.25)), Decimal::new(25, 2));
292        assert_eq!(step_budget_usd(Some(1.5)), Decimal::new(15, 1));
293    }
294
295    #[test]
296    fn month_start_truncates_to_first_day_midnight() {
297        let now = Utc.with_ymd_and_hms(2026, 2, 17, 23, 59, 59).unwrap();
298        let start = month_start(now);
299
300        assert_eq!(start.year(), 2026);
301        assert_eq!(start.month(), 2);
302        assert_eq!(start.day(), 1);
303        assert_eq!(start.hour(), 0);
304        assert_eq!(start.minute(), 0);
305        assert_eq!(start.second(), 0);
306    }
307
308    #[test]
309    fn month_start_is_idempotent() {
310        let now = Utc.with_ymd_and_hms(2026, 12, 1, 0, 0, 0).unwrap();
311        assert_eq!(month_start(month_start(now)), month_start(now));
312    }
313
314    #[test]
315    fn read_decimal_env_rejects_unset_empty_negative_and_garbage() {
316        // SAFETY: single-threaded test, variables are scoped to this test name.
317        unsafe {
318            env::remove_var("IRONFLOW_TEST_BUDGET_UNSET");
319            env::set_var("IRONFLOW_TEST_BUDGET_EMPTY", "   ");
320            env::set_var("IRONFLOW_TEST_BUDGET_NEGATIVE", "-1.5");
321            env::set_var("IRONFLOW_TEST_BUDGET_GARBAGE", "five dollars");
322            env::set_var("IRONFLOW_TEST_BUDGET_VALID", " 2.50 ");
323        }
324
325        assert_eq!(read_decimal_env("IRONFLOW_TEST_BUDGET_UNSET"), None);
326        assert_eq!(read_decimal_env("IRONFLOW_TEST_BUDGET_EMPTY"), None);
327        assert_eq!(read_decimal_env("IRONFLOW_TEST_BUDGET_NEGATIVE"), None);
328        assert_eq!(read_decimal_env("IRONFLOW_TEST_BUDGET_GARBAGE"), None);
329        assert_eq!(
330            read_decimal_env("IRONFLOW_TEST_BUDGET_VALID"),
331            Some(Decimal::new(250, 2))
332        );
333
334        // SAFETY: single-threaded test cleanup.
335        unsafe {
336            env::remove_var("IRONFLOW_TEST_BUDGET_EMPTY");
337            env::remove_var("IRONFLOW_TEST_BUDGET_NEGATIVE");
338            env::remove_var("IRONFLOW_TEST_BUDGET_GARBAGE");
339            env::remove_var("IRONFLOW_TEST_BUDGET_VALID");
340        }
341    }
342
343    #[test]
344    fn read_decimal_env_accepts_zero() {
345        // SAFETY: single-threaded test.
346        unsafe { env::set_var("IRONFLOW_TEST_BUDGET_ZERO", "0") };
347        assert_eq!(
348            read_decimal_env("IRONFLOW_TEST_BUDGET_ZERO"),
349            Some(Decimal::ZERO)
350        );
351        // SAFETY: single-threaded test cleanup.
352        unsafe { env::remove_var("IRONFLOW_TEST_BUDGET_ZERO") };
353    }
354}
ironflow_engine/budget.rs

ironflow_engine/
budget.rs
