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}