actionqueue_budget/
gate.rs

1//! Pre-dispatch budget gate.
2//!
3//! The gate wraps the [`BudgetTracker`] and provides the eligibility check
4//! used by the dispatch loop's step 7 (run selection). A task with any
5//! exhausted budget dimension is blocked from dispatch.
6
7use actionqueue_core::budget::BudgetDimension;
8use actionqueue_core::ids::TaskId;
9use tracing;
10
11use crate::tracker::BudgetTracker;
12
13/// Pre-dispatch eligibility gate backed by the budget tracker.
14///
15/// All checks are O(dimensions) per task — linear in the number of budget
16/// dimensions allocated to the task, typically very small (1-3).
17#[derive(Debug)]
18pub struct BudgetGate<'a> {
19    tracker: &'a BudgetTracker,
20}
21
22impl<'a> BudgetGate<'a> {
23    /// Creates a gate that borrows the given tracker.
24    pub fn new(tracker: &'a BudgetTracker) -> Self {
25        Self { tracker }
26    }
27
28    /// Returns `true` if no budget dimension is exhausted for this task.
29    pub fn can_dispatch(&self, task_id: TaskId) -> bool {
30        let allowed = !self.tracker.is_any_exhausted(task_id);
31        if !allowed {
32            tracing::debug!(%task_id, "dispatch blocked by exhausted budget");
33        }
34        allowed
35    }
36
37    /// Returns the (dimension, pct) pairs that have crossed the given threshold.
38    ///
39    /// Used by the dispatch loop to fire `BudgetThresholdCrossed` events.
40    pub fn check_threshold(
41        &self,
42        task_id: TaskId,
43        threshold_pct: u8,
44    ) -> Vec<(BudgetDimension, u8)> {
45        let dimensions =
46            [BudgetDimension::Token, BudgetDimension::CostCents, BudgetDimension::TimeSecs];
47        dimensions
48            .into_iter()
49            .filter_map(|dim| {
50                let pct = self.tracker.threshold_pct(task_id, dim)?;
51                if pct >= threshold_pct {
52                    Some((dim, pct))
53                } else {
54                    None
55                }
56            })
57            .collect()
58    }
59}
60
61#[cfg(test)]
62mod tests {
63    use actionqueue_core::budget::BudgetDimension;
64    use actionqueue_core::ids::TaskId;
65
66    use super::BudgetGate;
67    use crate::tracker::BudgetTracker;
68
69    #[test]
70    fn can_dispatch_when_no_budget_allocated() {
71        let tracker = BudgetTracker::new();
72        let gate = BudgetGate::new(&tracker);
73        let task = TaskId::new();
74        assert!(gate.can_dispatch(task));
75    }
76
77    #[test]
78    fn can_dispatch_within_budget() {
79        let mut tracker = BudgetTracker::new();
80        let task = TaskId::new();
81        tracker.allocate(task, BudgetDimension::Token, 1000);
82        tracker.consume(task, BudgetDimension::Token, 500);
83        let gate = BudgetGate::new(&tracker);
84        assert!(gate.can_dispatch(task));
85    }
86
87    #[test]
88    fn cannot_dispatch_when_exhausted() {
89        let mut tracker = BudgetTracker::new();
90        let task = TaskId::new();
91        tracker.allocate(task, BudgetDimension::Token, 100);
92        tracker.consume(task, BudgetDimension::Token, 100);
93        let gate = BudgetGate::new(&tracker);
94        assert!(!gate.can_dispatch(task));
95    }
96
97    #[test]
98    fn check_threshold_at_zero_returns_all_allocated() {
99        let mut tracker = BudgetTracker::new();
100        let task = TaskId::new();
101        tracker.allocate(task, BudgetDimension::Token, 100);
102        tracker.allocate(task, BudgetDimension::CostCents, 50);
103        let gate = BudgetGate::new(&tracker);
104        // threshold_pct=0 means everything at or above 0% should be returned
105        let crossing = gate.check_threshold(task, 0);
106        assert_eq!(crossing.len(), 2);
107    }
108
109    #[test]
110    fn check_threshold_at_100_only_returns_exhausted() {
111        let mut tracker = BudgetTracker::new();
112        let task = TaskId::new();
113        tracker.allocate(task, BudgetDimension::Token, 100);
114        tracker.consume(task, BudgetDimension::Token, 100);
115        tracker.allocate(task, BudgetDimension::CostCents, 100);
116        tracker.consume(task, BudgetDimension::CostCents, 50);
117        let gate = BudgetGate::new(&tracker);
118        let crossing = gate.check_threshold(task, 100);
119        assert_eq!(crossing.len(), 1);
120        assert_eq!(crossing[0].0, BudgetDimension::Token);
121    }
122
123    #[test]
124    fn check_threshold_no_allocations_returns_empty() {
125        let tracker = BudgetTracker::new();
126        let gate = BudgetGate::new(&tracker);
127        let task = TaskId::new();
128        let crossing = gate.check_threshold(task, 50);
129        assert!(crossing.is_empty());
130    }
131
132    #[test]
133    fn check_threshold_multiple_dimensions_crossing() {
134        let mut tracker = BudgetTracker::new();
135        let task = TaskId::new();
136        tracker.allocate(task, BudgetDimension::Token, 100);
137        tracker.consume(task, BudgetDimension::Token, 90); // 90%
138        tracker.allocate(task, BudgetDimension::CostCents, 200);
139        tracker.consume(task, BudgetDimension::CostCents, 160); // 80%
140        tracker.allocate(task, BudgetDimension::TimeSecs, 50);
141        tracker.consume(task, BudgetDimension::TimeSecs, 20); // 40%
142        let gate = BudgetGate::new(&tracker);
143        let crossing = gate.check_threshold(task, 80);
144        assert_eq!(crossing.len(), 2);
145    }
146
147    #[test]
148    fn check_threshold_returns_crossing_dimensions() {
149        let mut tracker = BudgetTracker::new();
150        let task = TaskId::new();
151        tracker.allocate(task, BudgetDimension::Token, 100);
152        tracker.consume(task, BudgetDimension::Token, 85);
153        let gate = BudgetGate::new(&tracker);
154        let crossing = gate.check_threshold(task, 80);
155        assert_eq!(crossing.len(), 1);
156        assert_eq!(crossing[0].0, BudgetDimension::Token);
157    }
158}
actionqueue_budget/gate.rs

actionqueue_budget/
gate.rs
