lean_ctx_workgraph/
budget_cascade.rs

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

lean_ctx_workgraph/
budget_cascade.rs
