lean_ctx/core/a2a/
budget_cascade.rs

1use serde::{Deserialize, Serialize};
2use std::collections::HashSet;
3use std::fmt;
4
5const DEFAULT_CHILD_FRACTION: f64 = 0.5;
6const DEFAULT_MINIMUM_BUDGET: u64 = 1_000;
7const DEFAULT_MAXIMUM_BUDGET: u64 = 500_000;
8const MAX_CASCADE_DEPTH: u32 = 5;
9
10#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
11pub struct BudgetAllocation {
12    pub parent_budget_tokens: u64,
13    pub parent_used_tokens: u64,
14    pub child_fraction: f64,
15    pub minimum_budget: u64,
16    pub maximum_budget: u64,
17}
18
19impl Default for BudgetAllocation {
20    fn default() -> Self {
21        Self {
22            parent_budget_tokens: 0,
23            parent_used_tokens: 0,
24            child_fraction: DEFAULT_CHILD_FRACTION,
25            minimum_budget: DEFAULT_MINIMUM_BUDGET,
26            maximum_budget: DEFAULT_MAXIMUM_BUDGET,
27        }
28    }
29}
30
31#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
32pub struct CascadedBudget {
33    pub allocated_tokens: u64,
34    pub parent_remaining: u64,
35    pub depth: u32,
36    pub lineage: Vec<String>,
37}
38
39#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
40pub enum CascadeError {
41    ZeroAllocation,
42    DepthLimitExceeded { depth: u32, maximum_depth: u32 },
43    LineageCycle { agent_id: String },
44}
45
46impl fmt::Display for CascadeError {
47    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
48        match self {
49            Self::ZeroAllocation => {
50                formatter.write_str("cascaded budget must be greater than zero")
51            }
52            Self::DepthLimitExceeded {
53                depth,
54                maximum_depth,
55            } => write!(
56                formatter,
57                "cascade depth {depth} exceeds maximum depth {maximum_depth}"
58            ),
59            Self::LineageCycle { agent_id } => {
60                write!(
61                    formatter,
62                    "cascade lineage contains cycle at agent {agent_id}"
63                )
64            }
65        }
66    }
67}
68
69impl std::error::Error for CascadeError {}
70
71pub fn cascade_budget(allocation: &BudgetAllocation) -> CascadedBudget {
72    let parent_remaining = allocation
73        .parent_budget_tokens
74        .saturating_sub(allocation.parent_used_tokens);
75    let fractional_budget = (parent_remaining as f64 * allocation.child_fraction) as u64;
76    let allocated_tokens = fractional_budget
77        .max(allocation.minimum_budget)
78        .min(allocation.maximum_budget);
79
80    CascadedBudget {
81        allocated_tokens,
82        parent_remaining,
83        depth: 0,
84        lineage: Vec::new(),
85    }
86}
87
88pub fn validate_cascade(budget: &CascadedBudget) -> Result<(), CascadeError> {
89    if budget.allocated_tokens == 0 {
90        return Err(CascadeError::ZeroAllocation);
91    }
92    if budget.depth > MAX_CASCADE_DEPTH {
93        return Err(CascadeError::DepthLimitExceeded {
94            depth: budget.depth,
95            maximum_depth: MAX_CASCADE_DEPTH,
96        });
97    }
98
99    let mut seen = HashSet::with_capacity(budget.lineage.len());
100    for agent_id in &budget.lineage {
101        if !seen.insert(agent_id) {
102            return Err(CascadeError::LineageCycle {
103                agent_id: agent_id.clone(),
104            });
105        }
106    }
107
108    Ok(())
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114
115    #[test]
116    fn normal_cascade_allocates_default_fraction() {
117        let allocation = BudgetAllocation {
118            parent_budget_tokens: 100_000,
119            parent_used_tokens: 20_000,
120            ..BudgetAllocation::default()
121        };
122
123        let budget = cascade_budget(&allocation);
124
125        assert_eq!(budget.allocated_tokens, 40_000);
126        assert_eq!(budget.parent_remaining, 80_000);
127        assert_eq!(budget.depth, 0);
128        assert!(budget.lineage.is_empty());
129        assert_eq!(validate_cascade(&budget), Ok(()));
130    }
131
132    #[test]
133    fn cascade_saturates_remaining_and_applies_bounds() {
134        let exhausted = BudgetAllocation {
135            parent_budget_tokens: 10,
136            parent_used_tokens: 20,
137            minimum_budget: 1_000,
138            maximum_budget: 5_000,
139            ..BudgetAllocation::default()
140        };
141        let capped = BudgetAllocation {
142            parent_budget_tokens: 2_000_000,
143            parent_used_tokens: 0,
144            child_fraction: 1.0,
145            minimum_budget: 1_000,
146            maximum_budget: 500_000,
147        };
148
149        assert_eq!(cascade_budget(&exhausted).parent_remaining, 0);
150        assert_eq!(cascade_budget(&exhausted).allocated_tokens, 1_000);
151        assert_eq!(cascade_budget(&capped).allocated_tokens, 500_000);
152    }
153
154    #[test]
155    fn validation_rejects_depth_above_limit() {
156        let budget = CascadedBudget {
157            allocated_tokens: 1_000,
158            parent_remaining: 2_000,
159            depth: 6,
160            lineage: vec!["parent".to_string()],
161        };
162
163        assert_eq!(
164            validate_cascade(&budget),
165            Err(CascadeError::DepthLimitExceeded {
166                depth: 6,
167                maximum_depth: 5,
168            })
169        );
170    }
171
172    #[test]
173    fn validation_rejects_zero_budget() {
174        let budget = CascadedBudget {
175            allocated_tokens: 0,
176            parent_remaining: 0,
177            depth: 0,
178            lineage: Vec::new(),
179        };
180
181        assert_eq!(validate_cascade(&budget), Err(CascadeError::ZeroAllocation));
182    }
183
184    #[test]
185    fn validation_rejects_lineage_cycle() {
186        let budget = CascadedBudget {
187            allocated_tokens: 1_000,
188            parent_remaining: 2_000,
189            depth: 2,
190            lineage: vec!["root".to_string(), "child".to_string(), "root".to_string()],
191        };
192
193        assert_eq!(
194            validate_cascade(&budget),
195            Err(CascadeError::LineageCycle {
196                agent_id: "root".to_string(),
197            })
198        );
199    }
200}
lean_ctx/core/a2a/budget_cascade.rs

lean_ctx/core/a2a/
budget_cascade.rs
