Skip to main content

ironflow_core/
agent_role.rs

1//! Identity and cost cap of an agent role.
2//!
3//! An agent role (see `ironflow_engine::config::AgentRole`) has a name and
4//! may declare a [`RoleCostCap`]. The role's tasks carry an [`AgentRoleBinding`]
5//! on their [`AgentConfig`](crate::provider::AgentConfig), which the engine
6//! reads to charge every agent step of the role to one shared counter over a
7//! calendar [`CostPeriod`] (UTC).
8//!
9//! # Examples
10//!
11//! ```
12//! use chrono::{TimeZone, Utc};
13//! use ironflow_core::agent_role::{AgentRoleBinding, CostPeriod, RoleCostCap};
14//! use rust_decimal::Decimal;
15//!
16//! let binding = AgentRoleBinding::new("reviewer", Some(RoleCostCap::per_month(Decimal::TEN)));
17//! assert_eq!(binding.period, CostPeriod::Month);
18//!
19//! let now = Utc.with_ymd_and_hms(2026, 7, 26, 15, 30, 0).unwrap();
20//! assert_eq!(CostPeriod::Month.start(now), Utc.with_ymd_and_hms(2026, 7, 1, 0, 0, 0).unwrap());
21//! ```
22
23use chrono::{DateTime, Datelike, Days, Months, TimeZone, Utc};
24use rust_decimal::Decimal;
25use serde::{Deserialize, Serialize};
26use strum::{Display, EnumString, IntoStaticStr};
27
28/// Longest name an agent role may have, in bytes.
29///
30/// # Examples
31///
32/// ```
33/// use ironflow_core::agent_role::MAX_ROLE_NAME_LEN;
34///
35/// assert_eq!(MAX_ROLE_NAME_LEN, 128);
36/// ```
37pub const MAX_ROLE_NAME_LEN: usize = 128;
38
39/// Whether `name` is a valid agent role name.
40///
41/// A name is not blank, has no leading or trailing whitespace, and has at
42/// most [`MAX_ROLE_NAME_LEN`] bytes. The default name of a role, its type
43/// name (`Reviewer`), is valid.
44///
45/// # Examples
46///
47/// ```
48/// use ironflow_core::agent_role::is_valid_role_name;
49///
50/// assert!(is_valid_role_name("code-reviewer.v2"));
51/// assert!(is_valid_role_name("Reviewer"));
52/// assert!(!is_valid_role_name(" reviewer"));
53/// assert!(!is_valid_role_name(""));
54/// ```
55pub fn is_valid_role_name(name: &str) -> bool {
56    !name.trim().is_empty() && name.trim() == name && name.len() <= MAX_ROLE_NAME_LEN
57}
58
59/// Calendar period, in UTC, over which the cost of a role is counted.
60///
61/// # Examples
62///
63/// ```
64/// use std::str::FromStr;
65/// use ironflow_core::agent_role::CostPeriod;
66///
67/// assert_eq!(CostPeriod::from_str("day").ok(), Some(CostPeriod::Day));
68/// assert_eq!(CostPeriod::Month.to_string(), "month");
69/// ```
70#[derive(
71    Debug,
72    Clone,
73    Copy,
74    PartialEq,
75    Eq,
76    Hash,
77    Serialize,
78    Deserialize,
79    Display,
80    EnumString,
81    IntoStaticStr,
82)]
83#[serde(rename_all = "snake_case")]
84#[strum(serialize_all = "snake_case")]
85pub enum CostPeriod {
86    /// From midnight UTC to the next midnight UTC.
87    Day,
88    /// From the 1st of the month at 00:00 UTC to the 1st of the next month.
89    Month,
90}
91
92impl CostPeriod {
93    /// Stable name of the period, as stored and sent over the API.
94    ///
95    /// # Examples
96    ///
97    /// ```
98    /// use ironflow_core::agent_role::CostPeriod;
99    ///
100    /// assert_eq!(CostPeriod::Day.as_str(), "day");
101    /// ```
102    pub fn as_str(self) -> &'static str {
103        self.into()
104    }
105
106    /// Start of the period containing `now`.
107    ///
108    /// # Examples
109    ///
110    /// ```
111    /// use chrono::{TimeZone, Utc};
112    /// use ironflow_core::agent_role::CostPeriod;
113    ///
114    /// let now = Utc.with_ymd_and_hms(2026, 7, 26, 15, 30, 0).unwrap();
115    /// assert_eq!(CostPeriod::Day.start(now), Utc.with_ymd_and_hms(2026, 7, 26, 0, 0, 0).unwrap());
116    /// ```
117    ///
118    /// # Panics
119    ///
120    /// Never panics for a valid [`DateTime<Utc>`]: midnight of an existing day
121    /// is always unambiguous in UTC.
122    pub fn start(self, now: DateTime<Utc>) -> DateTime<Utc> {
123        let day = match self {
124            Self::Day => now.day(),
125            Self::Month => 1,
126        };
127        Utc.with_ymd_and_hms(now.year(), now.month(), day, 0, 0, 0)
128            .single()
129            .expect("midnight UTC of an existing day is always unambiguous")
130    }
131
132    /// Start of the period following the one containing `now`: when the
133    /// counter resets.
134    ///
135    /// # Examples
136    ///
137    /// ```
138    /// use chrono::{TimeZone, Utc};
139    /// use ironflow_core::agent_role::CostPeriod;
140    ///
141    /// let now = Utc.with_ymd_and_hms(2026, 12, 26, 15, 30, 0).unwrap();
142    /// assert_eq!(
143    ///     CostPeriod::Month.next_start(now),
144    ///     Utc.with_ymd_and_hms(2027, 1, 1, 0, 0, 0).unwrap()
145    /// );
146    /// ```
147    ///
148    /// # Panics
149    ///
150    /// Panics only past the last date chrono represents (year 262143).
151    pub fn next_start(self, now: DateTime<Utc>) -> DateTime<Utc> {
152        let start = self.start(now);
153        match self {
154            Self::Day => start
155                .checked_add_days(Days::new(1))
156                .expect("the day after a representable date is representable"),
157            Self::Month => start
158                .checked_add_months(Months::new(1))
159                .expect("the month after a representable date is representable"),
160        }
161    }
162}
163
164/// Cost cap of an agent role over a calendar period.
165///
166/// # Examples
167///
168/// ```
169/// use ironflow_core::agent_role::{CostPeriod, RoleCostCap};
170/// use rust_decimal::Decimal;
171///
172/// let cap = RoleCostCap::per_month(Decimal::new(50, 0));
173/// assert_eq!(cap.limit_usd(), Decimal::new(50, 0));
174/// assert_eq!(cap.period(), CostPeriod::Month);
175/// ```
176#[derive(Debug, Clone, Copy, PartialEq, Eq)]
177pub struct RoleCostCap {
178    limit_usd: Decimal,
179    period: CostPeriod,
180}
181
182impl RoleCostCap {
183    /// Cap of `limit_usd` per calendar day (UTC). Zero disables the role.
184    ///
185    /// # Examples
186    ///
187    /// ```
188    /// use ironflow_core::agent_role::{CostPeriod, RoleCostCap};
189    /// use rust_decimal::Decimal;
190    ///
191    /// assert_eq!(RoleCostCap::per_day(Decimal::ONE).period(), CostPeriod::Day);
192    /// ```
193    ///
194    /// # Panics
195    ///
196    /// Panics if `limit_usd` is negative.
197    pub fn per_day(limit_usd: Decimal) -> Self {
198        Self::new(limit_usd, CostPeriod::Day)
199    }
200
201    /// Cap of `limit_usd` per calendar month (UTC). Zero disables the role.
202    ///
203    /// # Examples
204    ///
205    /// ```
206    /// use ironflow_core::agent_role::{CostPeriod, RoleCostCap};
207    /// use rust_decimal::Decimal;
208    ///
209    /// assert_eq!(RoleCostCap::per_month(Decimal::ONE).period(), CostPeriod::Month);
210    /// ```
211    ///
212    /// # Panics
213    ///
214    /// Panics if `limit_usd` is negative.
215    pub fn per_month(limit_usd: Decimal) -> Self {
216        Self::new(limit_usd, CostPeriod::Month)
217    }
218
219    fn new(limit_usd: Decimal, period: CostPeriod) -> Self {
220        assert!(
221            limit_usd >= Decimal::ZERO,
222            "the cost cap of an agent role must not be negative, got {limit_usd}"
223        );
224        Self { limit_usd, period }
225    }
226
227    /// The cap, in USD.
228    ///
229    /// # Examples
230    ///
231    /// ```
232    /// use ironflow_core::agent_role::RoleCostCap;
233    /// use rust_decimal::Decimal;
234    ///
235    /// assert_eq!(RoleCostCap::per_day(Decimal::TWO).limit_usd(), Decimal::TWO);
236    /// ```
237    pub fn limit_usd(&self) -> Decimal {
238        self.limit_usd
239    }
240
241    /// The period the cap applies to.
242    ///
243    /// # Examples
244    ///
245    /// ```
246    /// use ironflow_core::agent_role::{CostPeriod, RoleCostCap};
247    /// use rust_decimal::Decimal;
248    ///
249    /// assert_eq!(RoleCostCap::per_day(Decimal::TWO).period(), CostPeriod::Day);
250    /// ```
251    pub fn period(&self) -> CostPeriod {
252        self.period
253    }
254}
255
256/// Agent role an invocation is charged to, attached to its
257/// [`AgentConfig`](crate::provider::AgentConfig).
258///
259/// # Examples
260///
261/// ```
262/// use ironflow_core::agent_role::{AgentRoleBinding, CostPeriod};
263///
264/// let binding = AgentRoleBinding::new("reviewer", None);
265/// assert_eq!(binding.name, "reviewer");
266/// assert_eq!(binding.cost_cap_usd, None);
267/// assert_eq!(binding.period, CostPeriod::Month);
268/// ```
269#[derive(Debug, Clone, PartialEq, Eq)]
270pub struct AgentRoleBinding {
271    /// Name of the role, valid per [`is_valid_role_name`].
272    pub name: String,
273    /// Cap declared in code, `None` when the role only counts its cost.
274    pub cost_cap_usd: Option<Decimal>,
275    /// Period of the counter; [`CostPeriod::Month`] when no cap is declared.
276    pub period: CostPeriod,
277}
278
279impl AgentRoleBinding {
280    /// Bind to the role `name`, with its optional code cap.
281    ///
282    /// A role without a code cap still counts its cost over a month, so an
283    /// operator cap set later applies to what was already spent.
284    ///
285    /// # Examples
286    ///
287    /// ```
288    /// use ironflow_core::agent_role::{AgentRoleBinding, CostPeriod, RoleCostCap};
289    /// use rust_decimal::Decimal;
290    ///
291    /// let binding = AgentRoleBinding::new("triage", Some(RoleCostCap::per_day(Decimal::ONE)));
292    /// assert_eq!(binding.cost_cap_usd, Some(Decimal::ONE));
293    /// assert_eq!(binding.period, CostPeriod::Day);
294    /// ```
295    ///
296    /// # Panics
297    ///
298    /// Panics if `name` is not valid per [`is_valid_role_name`].
299    pub fn new(name: &str, cap: Option<RoleCostCap>) -> Self {
300        assert!(
301            is_valid_role_name(name),
302            "invalid agent role name {name:?}: not blank, no leading or trailing whitespace, at most {MAX_ROLE_NAME_LEN} bytes"
303        );
304        Self {
305            name: name.to_string(),
306            cost_cap_usd: cap.map(|c| c.limit_usd()),
307            period: cap.map(|c| c.period()).unwrap_or(CostPeriod::Month),
308        }
309    }
310}
311
312#[cfg(test)]
313mod tests {
314    use std::str::FromStr;
315
316    use serde_json::to_string;
317
318    use super::*;
319
320    fn at(y: i32, m: u32, d: u32, h: u32, min: u32, s: u32) -> DateTime<Utc> {
321        Utc.with_ymd_and_hms(y, m, d, h, min, s).unwrap()
322    }
323
324    #[test]
325    fn role_names_are_validated() {
326        assert!(is_valid_role_name("code-reviewer.v2"));
327        assert!(is_valid_role_name("a"));
328        assert!(is_valid_role_name("Reviewer"));
329        assert!(is_valid_role_name("revi\u{e9}wer"));
330        assert!(is_valid_role_name("code reviewer"));
331        assert!(is_valid_role_name("a/b"));
332        assert!(is_valid_role_name(&"a".repeat(MAX_ROLE_NAME_LEN)));
333        assert!(!is_valid_role_name(&"a".repeat(MAX_ROLE_NAME_LEN + 1)));
334        assert!(!is_valid_role_name(""));
335        assert!(!is_valid_role_name("   "));
336        assert!(!is_valid_role_name(" reviewer"));
337        assert!(!is_valid_role_name("reviewer\n"));
338    }
339
340    #[test]
341    fn day_period_starts_at_midnight_utc() {
342        let day = CostPeriod::Day;
343        assert_eq!(
344            day.start(at(2026, 7, 26, 15, 30, 0)),
345            at(2026, 7, 26, 0, 0, 0)
346        );
347        assert_eq!(
348            day.start(at(2026, 7, 26, 0, 0, 0)),
349            at(2026, 7, 26, 0, 0, 0)
350        );
351        assert_eq!(
352            day.next_start(at(2026, 12, 31, 23, 59, 59)),
353            at(2027, 1, 1, 0, 0, 0)
354        );
355        assert_eq!(
356            day.next_start(at(2028, 2, 28, 12, 0, 0)),
357            at(2028, 2, 29, 0, 0, 0)
358        );
359    }
360
361    #[test]
362    fn month_period_starts_on_the_first_at_midnight_utc() {
363        let month = CostPeriod::Month;
364        assert_eq!(
365            month.start(at(2026, 2, 17, 23, 59, 59)),
366            at(2026, 2, 1, 0, 0, 0)
367        );
368        assert_eq!(
369            month.start(at(2026, 3, 1, 0, 0, 0)),
370            at(2026, 3, 1, 0, 0, 0)
371        );
372        assert_eq!(
373            month.next_start(at(2026, 12, 15, 0, 0, 0)),
374            at(2027, 1, 1, 0, 0, 0)
375        );
376        assert_eq!(
377            month.next_start(at(2028, 2, 29, 10, 0, 0)),
378            at(2028, 3, 1, 0, 0, 0)
379        );
380        assert_eq!(
381            month.next_start(at(2026, 1, 31, 10, 0, 0)),
382            at(2026, 2, 1, 0, 0, 0)
383        );
384    }
385
386    #[test]
387    fn period_names_round_trip() {
388        for period in [CostPeriod::Day, CostPeriod::Month] {
389            assert_eq!(CostPeriod::from_str(period.as_str()).ok(), Some(period));
390            assert_eq!(period.to_string(), period.as_str());
391        }
392        assert!(CostPeriod::from_str("week").is_err());
393        assert_eq!(to_string(&CostPeriod::Day).unwrap(), "\"day\"");
394    }
395
396    #[test]
397    fn zero_cap_is_allowed() {
398        assert_eq!(
399            RoleCostCap::per_day(Decimal::ZERO).limit_usd(),
400            Decimal::ZERO
401        );
402    }
403
404    #[test]
405    #[should_panic(expected = "must not be negative")]
406    fn negative_cap_panics() {
407        let _ = RoleCostCap::per_month(Decimal::NEGATIVE_ONE);
408    }
409
410    #[test]
411    #[should_panic(expected = "invalid agent role name")]
412    fn binding_with_an_invalid_name_panics() {
413        let _ = AgentRoleBinding::new(" reviewer", None);
414    }
415
416    #[test]
417    fn binding_without_cap_counts_over_a_month() {
418        let binding = AgentRoleBinding::new("reviewer", None);
419        assert_eq!(binding.cost_cap_usd, None);
420        assert_eq!(binding.period, CostPeriod::Month);
421    }
422
423    #[test]
424    fn binding_takes_the_cap_and_period() {
425        let binding = AgentRoleBinding::new("reviewer", Some(RoleCostCap::per_day(Decimal::TEN)));
426        assert_eq!(binding.cost_cap_usd, Some(Decimal::TEN));
427        assert_eq!(binding.period, CostPeriod::Day);
428    }
429}