Skip to main content

actionqueue_core/
budget.rs

1//! Budget domain types for resource consumption tracking.
2//!
3//! These types define the vocabulary for budget allocation and consumption
4//! that Caelum uses to enforce per-thread and per-Vessel resource caps.
5//! ActionQueue enforces the caps; Caelum aggregates and reports consumption.
6
7/// The dimension along which a budget is measured.
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
9#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
10pub enum BudgetDimension {
11    /// Token-count budget (e.g. LLM tokens).
12    Token,
13    /// Cost budget measured in hundredths of a cent (cost-cents).
14    CostCents,
15    /// Wall-clock time budget in seconds.
16    TimeSecs,
17}
18
19impl std::fmt::Display for BudgetDimension {
20    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
21        let name = match self {
22            BudgetDimension::Token => "token",
23            BudgetDimension::CostCents => "cost_cents",
24            BudgetDimension::TimeSecs => "time_secs",
25        };
26        write!(f, "{name}")
27    }
28}
29
30/// A unit of resource consumption reported by a handler after execution.
31///
32/// Handlers attach zero or more `BudgetConsumption` records to their output.
33/// The dispatch loop durably records these and updates the in-memory
34/// `BudgetTracker` after each attempt completes.
35#[derive(Debug, Clone, PartialEq, Eq)]
36#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
37pub struct BudgetConsumption {
38    /// The resource dimension being consumed.
39    pub dimension: BudgetDimension,
40    /// The amount consumed in this attempt.
41    pub amount: u64,
42}
43
44impl BudgetConsumption {
45    /// Creates a new consumption record.
46    pub fn new(dimension: BudgetDimension, amount: u64) -> Self {
47        Self { dimension, amount }
48    }
49}
50
51/// A budget allocation for a single (task, dimension) pair.
52///
53/// Each task may have at most one allocation per `BudgetDimension`.
54/// Validated construction ensures the limit is non-zero.
55#[derive(Debug, Clone, PartialEq, Eq)]
56#[cfg_attr(feature = "serde", derive(serde::Serialize))]
57pub struct BudgetAllocation {
58    dimension: BudgetDimension,
59    limit: u64,
60}
61
62/// Error returned when a `BudgetAllocation` is constructed with invalid values.
63#[derive(Debug, Clone, PartialEq, Eq)]
64pub enum BudgetAllocationError {
65    /// Budget limit must be greater than zero.
66    ZeroLimit,
67}
68
69impl std::fmt::Display for BudgetAllocationError {
70    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
71        match self {
72            BudgetAllocationError::ZeroLimit => write!(f, "budget limit must be greater than zero"),
73        }
74    }
75}
76
77impl std::error::Error for BudgetAllocationError {}
78
79impl BudgetAllocation {
80    /// Creates a new budget allocation with validation.
81    ///
82    /// # Errors
83    ///
84    /// Returns [`BudgetAllocationError::ZeroLimit`] if `limit` is 0.
85    pub fn new(dimension: BudgetDimension, limit: u64) -> Result<Self, BudgetAllocationError> {
86        if limit == 0 {
87            return Err(BudgetAllocationError::ZeroLimit);
88        }
89        Ok(Self { dimension, limit })
90    }
91
92    /// Returns the budget dimension.
93    pub fn dimension(&self) -> BudgetDimension {
94        self.dimension
95    }
96
97    /// Returns the budget limit.
98    pub fn limit(&self) -> u64 {
99        self.limit
100    }
101}
102
103#[cfg(feature = "serde")]
104impl<'de> serde::Deserialize<'de> for BudgetAllocation {
105    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
106    where
107        D: serde::Deserializer<'de>,
108    {
109        #[derive(serde::Deserialize)]
110        struct Wire {
111            dimension: BudgetDimension,
112            limit: u64,
113        }
114
115        let wire = Wire::deserialize(deserializer)?;
116        BudgetAllocation::new(wire.dimension, wire.limit)
117            .map_err(|_| serde::de::Error::custom("budget limit must be greater than zero"))
118    }
119}