lean_ctx/core/
work_graph.rs

1//! Bounded Work Graph for multi-agent orchestration (P11 / DIM 4).
2//!
3//! Manages parent/child agent delegation with:
4//! - Budget inheritance (child cannot exceed parent)
5//! - Fan-out limits (max concurrent children)
6//! - Stop conditions (stale, over-budget, redundant)
7//! - Provenance tracking for attribution
8
9use std::collections::BTreeMap;
10
11use crate::core::a2a::budget_cascade::{
12    BudgetAllocation, CascadeError, cascade_budget, validate_cascade,
13};
14use serde::{Deserialize, Serialize};
15
16pub const WORK_GRAPH_SCHEMA_VERSION: u16 = 1;
17const MAX_GRAPH_NODES: usize = 256;
18const MAX_FAN_OUT: usize = 16;
19const MAX_DEPTH: u16 = 8;
20
21#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
22#[serde(rename_all = "snake_case")]
23pub enum NodeStatus {
24    Pending,
25    Active,
26    Completed,
27    Stopped,
28    Failed,
29}
30
31#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
32#[serde(rename_all = "snake_case")]
33pub enum StopReason {
34    BudgetExhausted,
35    Stale,
36    Redundant,
37    ParentStopped,
38    ManualStop,
39    DepthExceeded,
40    FanOutExceeded,
41}
42
43#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
44pub struct WorkNodeBudget {
45    pub tokens_allocated: u64,
46    pub tokens_consumed: u64,
47    pub cost_micros_allocated: u64,
48    pub cost_micros_consumed: u64,
49}
50
51impl WorkNodeBudget {
52    pub fn tokens_remaining(&self) -> u64 {
53        self.tokens_allocated.saturating_sub(self.tokens_consumed)
54    }
55
56    pub fn cost_remaining(&self) -> u64 {
57        self.cost_micros_allocated
58            .saturating_sub(self.cost_micros_consumed)
59    }
60
61    pub fn is_exhausted(&self) -> bool {
62        self.tokens_remaining() == 0 || self.cost_remaining() == 0
63    }
64}
65
66/// Tracks the total budget across an entire delegation chain.
67#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
68pub struct ChainBudget {
69    pub chain_id: String,
70    pub root_budget_tokens: u64,
71    pub total_consumed_tokens: u64,
72    pub total_allocated_tokens: u64,
73    pub depth: u16,
74}
75
76impl ChainBudget {
77    pub fn remaining(&self) -> u64 {
78        self.root_budget_tokens
79            .saturating_sub(self.total_consumed_tokens)
80    }
81
82    pub fn utilization_pct(&self) -> f64 {
83        if self.root_budget_tokens == 0 {
84            return 0.0;
85        }
86
87        (self.total_consumed_tokens as f64 / self.root_budget_tokens as f64) * 100.0
88    }
89}
90
91#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
92pub struct WorkNode {
93    pub node_id: String,
94    pub agent_id: String,
95    pub parent_node_id: Option<String>,
96    pub capsule_ref: String,
97    pub status: NodeStatus,
98    pub budget: WorkNodeBudget,
99    pub depth: u16,
100    pub stop_reason: Option<StopReason>,
101    pub outcome_ref: Option<String>,
102}
103
104/// Bounded, acyclic work graph with enforced fan-out and budget constraints.
105#[derive(Clone, Debug, Serialize, Deserialize)]
106pub struct BoundedWorkGraph {
107    nodes: BTreeMap<String, WorkNode>,
108    children: BTreeMap<String, Vec<String>>,
109    #[serde(default)]
110    chain_budgets: BTreeMap<String, ChainBudget>,
111    #[serde(default)]
112    pending_child_budgets: BTreeMap<String, WorkNodeBudget>,
113    max_fan_out: usize,
114    max_depth: u16,
115}
116
117impl Default for BoundedWorkGraph {
118    fn default() -> Self {
119        Self::new(MAX_FAN_OUT, MAX_DEPTH)
120    }
121}
122
123impl BoundedWorkGraph {
124    #[must_use]
125    pub fn new(max_fan_out: usize, max_depth: u16) -> Self {
126        Self {
127            nodes: BTreeMap::new(),
128            children: BTreeMap::new(),
129            chain_budgets: BTreeMap::new(),
130            pending_child_budgets: BTreeMap::new(),
131            max_fan_out: max_fan_out.clamp(1, MAX_FAN_OUT),
132            max_depth: max_depth.clamp(1, MAX_DEPTH),
133        }
134    }
135
136    /// Add a root node (no parent).
137    pub fn add_root(
138        &mut self,
139        node_id: String,
140        agent_id: String,
141        capsule_ref: String,
142        budget: WorkNodeBudget,
143    ) -> Result<&WorkNode, WorkGraphError> {
144        if self.nodes.len() >= MAX_GRAPH_NODES {
145            return Err(WorkGraphError::CapacityExceeded);
146        }
147        if self.nodes.contains_key(&node_id) {
148            return Err(WorkGraphError::DuplicateNode(node_id));
149        }
150        let node = WorkNode {
151            node_id: node_id.clone(),
152            agent_id,
153            parent_node_id: None,
154            capsule_ref,
155            status: NodeStatus::Active,
156            budget,
157            depth: 0,
158            stop_reason: None,
159            outcome_ref: None,
160        };
161        self.nodes.insert(node_id.clone(), node);
162        self.chain_budgets.insert(
163            node_id.clone(),
164            ChainBudget {
165                chain_id: node_id.clone(),
166                root_budget_tokens: self.nodes[&node_id].budget.tokens_allocated,
167                total_consumed_tokens: 0,
168                total_allocated_tokens: self.nodes[&node_id].budget.tokens_allocated,
169                depth: 0,
170            },
171        );
172        Ok(self.nodes.get(&node_id).unwrap())
173    }
174
175    /// Delegate work to a child node. Validates budget inheritance and fan-out.
176    pub fn delegate(
177        &mut self,
178        parent_node_id: &str,
179        child_node_id: String,
180        child_agent_id: String,
181        capsule_ref: String,
182        child_budget: WorkNodeBudget,
183    ) -> Result<&WorkNode, WorkGraphError> {
184        if self.nodes.len() >= MAX_GRAPH_NODES {
185            return Err(WorkGraphError::CapacityExceeded);
186        }
187        if self.nodes.contains_key(&child_node_id) {
188            return Err(WorkGraphError::DuplicateNode(child_node_id));
189        }
190        let parent = self
191            .nodes
192            .get(parent_node_id)
193            .ok_or_else(|| WorkGraphError::NodeNotFound(parent_node_id.to_string()))?;
194        if parent.status != NodeStatus::Active {
195            return Err(WorkGraphError::ParentNotActive(parent_node_id.to_string()));
196        }
197        let new_depth = parent.depth + 1;
198        if new_depth > self.max_depth {
199            return Err(WorkGraphError::DepthExceeded(self.max_depth));
200        }
201        if child_budget.tokens_allocated > parent.budget.tokens_remaining() {
202            return Err(WorkGraphError::BudgetExceedsParent {
203                child_requested: child_budget.tokens_allocated,
204                parent_remaining: parent.budget.tokens_remaining(),
205            });
206        }
207        if child_budget.cost_micros_allocated > parent.budget.cost_remaining() {
208            return Err(WorkGraphError::BudgetExceedsParent {
209                child_requested: child_budget.cost_micros_allocated,
210                parent_remaining: parent.budget.cost_remaining(),
211            });
212        }
213        let current_children = self.children.get(parent_node_id).map_or(0, Vec::len);
214        if current_children >= self.max_fan_out {
215            return Err(WorkGraphError::FanOutExceeded(self.max_fan_out));
216        }
217        let child_tokens_allocated = child_budget.tokens_allocated;
218        let node = WorkNode {
219            node_id: child_node_id.clone(),
220            agent_id: child_agent_id,
221            parent_node_id: Some(parent_node_id.to_string()),
222            capsule_ref,
223            status: NodeStatus::Active,
224            budget: child_budget,
225            depth: new_depth,
226            stop_reason: None,
227            outcome_ref: None,
228        };
229        self.nodes.insert(child_node_id.clone(), node);
230        self.children
231            .entry(parent_node_id.to_string())
232            .or_default()
233            .push(child_node_id.clone());
234        let chain_id = self
235            .chain_id_for_node(&child_node_id)
236            .expect("delegated child always has a root node");
237        let pending_tokens = self
238            .pending_child_budgets
239            .remove(&child_node_id)
240            .map_or(0, |budget| budget.tokens_allocated);
241        self.record_chain_allocation(&chain_id, pending_tokens, child_tokens_allocated, new_depth);
242        Ok(self.nodes.get(&child_node_id).unwrap())
243    }
244
245    /// Allocates budget for a child node using cascade rules.
246    ///
247    /// The returned budget is reserved for `child_id` until it is passed to
248    /// [`Self::delegate`], so chain allocation is not counted twice.
249    pub fn allocate_child_budget(
250        &mut self,
251        parent_id: &str,
252        child_id: &str,
253        fraction: f64,
254    ) -> Result<WorkNodeBudget, WorkGraphError> {
255        let (parent_budget_tokens, parent_used_tokens, parent_cost_remaining, parent_depth) = {
256            let parent = self
257                .nodes
258                .get(parent_id)
259                .ok_or_else(|| WorkGraphError::NodeNotFound(parent_id.to_string()))?;
260            if parent.status != NodeStatus::Active {
261                return Err(WorkGraphError::ParentNotActive(parent_id.to_string()));
262            }
263            (
264                parent.budget.tokens_allocated,
265                parent.budget.tokens_consumed,
266                parent.budget.cost_remaining(),
267                parent.depth,
268            )
269        };
270
271        let parent_remaining = parent_budget_tokens.saturating_sub(parent_used_tokens);
272        if parent_remaining == 0 {
273            return Err(WorkGraphError::BudgetExceedsParent {
274                child_requested: 1,
275                parent_remaining,
276            });
277        }
278
279        let allocation = BudgetAllocation {
280            parent_budget_tokens,
281            parent_used_tokens,
282            child_fraction: fraction,
283            minimum_budget: 0,
284            maximum_budget: parent_remaining,
285        };
286        let mut cascaded = cascade_budget(&allocation);
287        cascaded.depth = u32::from(parent_depth) + 1;
288        cascaded.lineage = self.node_lineage(parent_id);
289        cascaded.lineage.push(child_id.to_string());
290        validate_cascade(&cascaded)?;
291        if cascaded.allocated_tokens > parent_remaining {
292            return Err(WorkGraphError::BudgetExceedsParent {
293                child_requested: cascaded.allocated_tokens,
294                parent_remaining,
295            });
296        }
297
298        let cost_micros_allocated = if parent_cost_remaining == 0 {
299            0
300        } else {
301            ((parent_cost_remaining as f64 * fraction) as u64)
302                .max(1)
303                .min(parent_cost_remaining)
304        };
305        let child_budget = WorkNodeBudget {
306            tokens_allocated: cascaded.allocated_tokens,
307            tokens_consumed: 0,
308            cost_micros_allocated,
309            cost_micros_consumed: 0,
310        };
311        let chain_id = self
312            .chain_id_for_node(parent_id)
313            .expect("parent node always has a root node");
314        self.pending_child_budgets
315            .insert(child_id.to_string(), child_budget.clone());
316        self.record_chain_allocation(
317            &chain_id,
318            0,
319            child_budget.tokens_allocated,
320            parent_depth + 1,
321        );
322
323        Ok(child_budget)
324    }
325
326    /// Mark a node as completed with an outcome reference.
327    pub fn complete(&mut self, node_id: &str, outcome_ref: String) -> Result<(), WorkGraphError> {
328        let node = self
329            .nodes
330            .get_mut(node_id)
331            .ok_or_else(|| WorkGraphError::NodeNotFound(node_id.to_string()))?;
332        if node.status != NodeStatus::Active {
333            return Err(WorkGraphError::InvalidTransition(node_id.to_string()));
334        }
335        node.status = NodeStatus::Completed;
336        node.outcome_ref = Some(outcome_ref);
337        Ok(())
338    }
339
340    /// Stop a node and all its descendants (cascade).
341    pub fn stop(
342        &mut self,
343        node_id: &str,
344        reason: StopReason,
345    ) -> Result<Vec<String>, WorkGraphError> {
346        if !self.nodes.contains_key(node_id) {
347            return Err(WorkGraphError::NodeNotFound(node_id.to_string()));
348        }
349        let mut stopped = Vec::new();
350        self.stop_recursive(node_id, reason, &mut stopped);
351        Ok(stopped)
352    }
353
354    /// Record token consumption on a node.
355    pub fn consume_budget(
356        &mut self,
357        node_id: &str,
358        tokens: u64,
359        cost_micros: u64,
360    ) -> Result<bool, WorkGraphError> {
361        let exhausted = {
362            let node = self
363                .nodes
364                .get_mut(node_id)
365                .ok_or_else(|| WorkGraphError::NodeNotFound(node_id.to_string()))?;
366            if node.status != NodeStatus::Active {
367                return Err(WorkGraphError::InvalidTransition(node_id.to_string()));
368            }
369            node.budget.tokens_consumed = node.budget.tokens_consumed.saturating_add(tokens);
370            node.budget.cost_micros_consumed =
371                node.budget.cost_micros_consumed.saturating_add(cost_micros);
372            node.budget.is_exhausted()
373        };
374        let chain_id = self
375            .chain_id_for_node(node_id)
376            .expect("node always has a root node");
377        self.ensure_chain_budget(&chain_id);
378        let chain_exhausted = {
379            let chain = self
380                .chain_budgets
381                .get_mut(&chain_id)
382                .expect("chain budget initialized");
383            chain.total_consumed_tokens = chain.total_consumed_tokens.saturating_add(tokens);
384            chain.total_consumed_tokens >= chain.root_budget_tokens
385        };
386
387        if chain_exhausted {
388            self.stop_recursive(&chain_id, StopReason::BudgetExhausted, &mut Vec::new());
389            return Ok(true);
390        }
391        if exhausted {
392            let node = self.nodes.get_mut(node_id).expect("node checked above");
393            node.status = NodeStatus::Stopped;
394            node.stop_reason = Some(StopReason::BudgetExhausted);
395            return Ok(true);
396        }
397
398        Ok(false)
399    }
400
401    /// Records token consumption for a node and updates its chain budget.
402    pub fn consume_tokens(&mut self, node_id: &str, tokens: u64) -> Result<(), WorkGraphError> {
403        self.consume_budget(node_id, tokens, 0)?;
404        Ok(())
405    }
406
407    /// Returns the chain budget for a given node's chain.
408    pub fn chain_budget_for(&self, node_id: &str) -> Option<&ChainBudget> {
409        let chain_id = self.chain_id_for_node(node_id)?;
410        self.chain_budgets.get(&chain_id)
411    }
412
413    /// Returns all chains at or above their utilization threshold.
414    pub fn over_budget_chains(&self, threshold_pct: f64) -> Vec<&ChainBudget> {
415        self.chain_budgets
416            .values()
417            .filter(|budget| budget.utilization_pct() >= threshold_pct)
418            .collect()
419    }
420
421    /// Check all nodes for stop conditions and cascade.
422    pub fn enforce_stop_conditions(&mut self) -> Vec<(String, StopReason)> {
423        let exhausted: Vec<String> = self
424            .nodes
425            .iter()
426            .filter(|(_, n)| n.status == NodeStatus::Active && n.budget.is_exhausted())
427            .map(|(id, _)| id.clone())
428            .collect();
429        let mut stopped = Vec::new();
430        for node_id in exhausted {
431            let mut cascade = Vec::new();
432            self.stop_recursive(&node_id, StopReason::BudgetExhausted, &mut cascade);
433            for id in cascade {
434                stopped.push((id, StopReason::BudgetExhausted));
435            }
436        }
437        stopped
438    }
439
440    pub fn get_node(&self, node_id: &str) -> Option<&WorkNode> {
441        self.nodes.get(node_id)
442    }
443
444    pub fn children_of(&self, node_id: &str) -> &[String] {
445        self.children.get(node_id).map_or(&[], Vec::as_slice)
446    }
447
448    pub fn active_count(&self) -> usize {
449        self.nodes
450            .values()
451            .filter(|n| n.status == NodeStatus::Active)
452            .count()
453    }
454
455    pub fn total_count(&self) -> usize {
456        self.nodes.len()
457    }
458
459    #[allow(clippy::collapsible_if)]
460    fn stop_recursive(&mut self, node_id: &str, reason: StopReason, stopped: &mut Vec<String>) {
461        if let Some(node) = self.nodes.get_mut(node_id) {
462            if matches!(node.status, NodeStatus::Active | NodeStatus::Pending) {
463                node.status = NodeStatus::Stopped;
464                node.stop_reason = Some(reason);
465                stopped.push(node_id.to_string());
466            }
467        }
468        let children: Vec<String> = self.children.get(node_id).cloned().unwrap_or_default();
469        for child_id in children {
470            self.stop_recursive(&child_id, StopReason::ParentStopped, stopped);
471        }
472    }
473
474    fn chain_id_for_node(&self, node_id: &str) -> Option<String> {
475        let mut current_id = node_id;
476        let mut current = self.nodes.get(current_id)?;
477        while let Some(parent_id) = current.parent_node_id.as_deref() {
478            current_id = parent_id;
479            current = self.nodes.get(current_id)?;
480        }
481        Some(current_id.to_string())
482    }
483
484    fn node_lineage(&self, node_id: &str) -> Vec<String> {
485        let mut lineage = Vec::new();
486        let mut current_id = Some(node_id);
487        while let Some(id) = current_id {
488            let Some(node) = self.nodes.get(id) else {
489                break;
490            };
491            lineage.push(id.to_string());
492            current_id = node.parent_node_id.as_deref();
493        }
494        lineage.reverse();
495        lineage
496    }
497
498    fn ensure_chain_budget(&mut self, chain_id: &str) {
499        let root_budget_tokens = self
500            .nodes
501            .get(chain_id)
502            .map_or(0, |node| node.budget.tokens_allocated);
503        self.chain_budgets
504            .entry(chain_id.to_string())
505            .or_insert(ChainBudget {
506                chain_id: chain_id.to_string(),
507                root_budget_tokens,
508                total_consumed_tokens: 0,
509                total_allocated_tokens: root_budget_tokens,
510                depth: 0,
511            });
512    }
513
514    fn record_chain_allocation(
515        &mut self,
516        chain_id: &str,
517        previous_tokens: u64,
518        tokens_allocated: u64,
519        depth: u16,
520    ) {
521        self.ensure_chain_budget(chain_id);
522        let chain = self
523            .chain_budgets
524            .get_mut(chain_id)
525            .expect("chain budget initialized");
526        chain.total_allocated_tokens = chain
527            .total_allocated_tokens
528            .saturating_sub(previous_tokens)
529            .saturating_add(tokens_allocated);
530        chain.depth = chain.depth.max(depth);
531    }
532}
533
534// ─── Errors ──────────────────────────────────────────────────────────────────
535
536#[derive(Debug, thiserror::Error)]
537pub enum WorkGraphError {
538    #[error("graph at capacity ({MAX_GRAPH_NODES} nodes)")]
539    CapacityExceeded,
540    #[error("duplicate node: {0}")]
541    DuplicateNode(String),
542    #[error("node not found: {0}")]
543    NodeNotFound(String),
544    #[error("parent not active: {0}")]
545    ParentNotActive(String),
546    #[error("depth exceeds max {0}")]
547    DepthExceeded(u16),
548    #[error("fan-out exceeds max {0}")]
549    FanOutExceeded(usize),
550    #[error("child budget ({child_requested}) exceeds parent remaining ({parent_remaining})")]
551    BudgetExceedsParent {
552        child_requested: u64,
553        parent_remaining: u64,
554    },
555    #[error("invalid status transition for node: {0}")]
556    InvalidTransition(String),
557    #[error("budget cascade error: {0}")]
558    Cascade(#[from] CascadeError),
559}
560
561// ─── Tests ───────────────────────────────────────────────────────────────────
562
563#[cfg(test)]
564mod tests {
565    use crate::core::work_graph::{
566        BoundedWorkGraph, MAX_DEPTH, MAX_FAN_OUT, NodeStatus, StopReason, WorkGraphError,
567        WorkNodeBudget,
568    };
569
570    fn budget(tokens: u64, cost: u64) -> WorkNodeBudget {
571        WorkNodeBudget {
572            tokens_allocated: tokens,
573            tokens_consumed: 0,
574            cost_micros_allocated: cost,
575            cost_micros_consumed: 0,
576        }
577    }
578
579    #[test]
580    fn basic_delegation_and_budget_inheritance() {
581        let mut g = BoundedWorkGraph::default();
582        g.add_root(
583            "root".into(),
584            "parent-agent".into(),
585            "capsule:abc".into(),
586            budget(1000, 500),
587        )
588        .unwrap();
589        g.delegate(
590            "root",
591            "child-1".into(),
592            "child-agent".into(),
593            "capsule:def".into(),
594            budget(400, 200),
595        )
596        .unwrap();
597        assert_eq!(g.active_count(), 2);
598        assert_eq!(g.children_of("root"), &["child-1"]);
599    }
600
601    #[test]
602    fn child_cannot_exceed_parent_budget() {
603        let mut g = BoundedWorkGraph::default();
604        g.add_root(
605            "root".into(),
606            "a".into(),
607            "capsule:x".into(),
608            budget(100, 50),
609        )
610        .unwrap();
611        assert!(matches!(
612            g.delegate(
613                "root",
614                "c".into(),
615                "b".into(),
616                "capsule:y".into(),
617                budget(200, 30)
618            ),
619            Err(WorkGraphError::BudgetExceedsParent { .. })
620        ));
621    }
622
623    #[test]
624    fn fan_out_limit_enforced() {
625        let mut g = BoundedWorkGraph::new(2, MAX_DEPTH);
626        g.add_root(
627            "root".into(),
628            "a".into(),
629            "capsule:x".into(),
630            budget(1000, 1000),
631        )
632        .unwrap();
633        g.delegate(
634            "root",
635            "c1".into(),
636            "b".into(),
637            "capsule:1".into(),
638            budget(100, 100),
639        )
640        .unwrap();
641        g.delegate(
642            "root",
643            "c2".into(),
644            "b".into(),
645            "capsule:2".into(),
646            budget(100, 100),
647        )
648        .unwrap();
649        assert!(matches!(
650            g.delegate(
651                "root",
652                "c3".into(),
653                "b".into(),
654                "capsule:3".into(),
655                budget(100, 100)
656            ),
657            Err(WorkGraphError::FanOutExceeded(2))
658        ));
659    }
660
661    #[test]
662    fn depth_limit_enforced() {
663        let mut g = BoundedWorkGraph::new(MAX_FAN_OUT, 2);
664        g.add_root(
665            "n0".into(),
666            "a".into(),
667            "capsule:0".into(),
668            budget(1000, 1000),
669        )
670        .unwrap();
671        g.delegate(
672            "n0",
673            "n1".into(),
674            "b".into(),
675            "capsule:1".into(),
676            budget(500, 500),
677        )
678        .unwrap();
679        g.delegate(
680            "n1",
681            "n2".into(),
682            "c".into(),
683            "capsule:2".into(),
684            budget(200, 200),
685        )
686        .unwrap();
687        assert!(matches!(
688            g.delegate(
689                "n2",
690                "n3".into(),
691                "d".into(),
692                "capsule:3".into(),
693                budget(100, 100)
694            ),
695            Err(WorkGraphError::DepthExceeded(2))
696        ));
697    }
698
699    #[test]
700    fn stop_cascades_to_children() {
701        let mut g = BoundedWorkGraph::default();
702        g.add_root(
703            "root".into(),
704            "a".into(),
705            "capsule:r".into(),
706            budget(1000, 1000),
707        )
708        .unwrap();
709        g.delegate(
710            "root",
711            "c1".into(),
712            "b".into(),
713            "capsule:1".into(),
714            budget(300, 300),
715        )
716        .unwrap();
717        g.delegate(
718            "c1",
719            "gc1".into(),
720            "c".into(),
721            "capsule:gc".into(),
722            budget(100, 100),
723        )
724        .unwrap();
725        let stopped = g.stop("c1", StopReason::Stale).unwrap();
726        assert_eq!(stopped, vec!["c1", "gc1"]);
727        assert_eq!(g.get_node("c1").unwrap().status, NodeStatus::Stopped);
728        assert_eq!(
729            g.get_node("gc1").unwrap().stop_reason,
730            Some(StopReason::ParentStopped)
731        );
732    }
733
734    #[test]
735    fn budget_exhaustion_auto_stops() {
736        let mut g = BoundedWorkGraph::default();
737        g.add_root(
738            "root".into(),
739            "a".into(),
740            "capsule:r".into(),
741            budget(100, 100),
742        )
743        .unwrap();
744        let exhausted = g.consume_budget("root", 100, 50).unwrap();
745        assert!(exhausted);
746        assert_eq!(g.get_node("root").unwrap().status, NodeStatus::Stopped);
747    }
748
749    #[test]
750    fn allocate_child_budget_uses_requested_fraction() {
751        let mut g = BoundedWorkGraph::default();
752        g.add_root(
753            "root".into(),
754            "a".into(),
755            "capsule:r".into(),
756            budget(10_000, 5_000),
757        )
758        .unwrap();
759
760        let child_budget = g.allocate_child_budget("root", "child", 0.5).unwrap();
761
762        assert_eq!(child_budget, budget(5_000, 2_500));
763        let chain = g.chain_budget_for("root").unwrap();
764        assert_eq!(chain.total_allocated_tokens, 15_000);
765        assert_eq!(chain.depth, 1);
766    }
767
768    #[test]
769    fn allocate_child_budget_rejects_exhausted_parent() {
770        let mut g = BoundedWorkGraph::default();
771        g.add_root(
772            "root".into(),
773            "a".into(),
774            "capsule:r".into(),
775            budget(100, 100),
776        )
777        .unwrap();
778        g.consume_tokens("root", 100).unwrap();
779
780        assert!(g.allocate_child_budget("root", "child", 0.5).is_err());
781    }
782
783    #[test]
784    fn consume_tokens_updates_node_and_chain() {
785        let mut g = BoundedWorkGraph::default();
786        g.add_root(
787            "root".into(),
788            "a".into(),
789            "capsule:r".into(),
790            budget(1_000, 1_000),
791        )
792        .unwrap();
793
794        g.consume_tokens("root", 250).unwrap();
795
796        assert_eq!(g.get_node("root").unwrap().budget.tokens_consumed, 250);
797        assert_eq!(
798            g.chain_budget_for("root").unwrap().total_consumed_tokens,
799            250
800        );
801    }
802
803    #[test]
804    fn consume_tokens_stops_exhausted_node() {
805        let mut g = BoundedWorkGraph::default();
806        g.add_root(
807            "root".into(),
808            "a".into(),
809            "capsule:r".into(),
810            budget(100, 100),
811        )
812        .unwrap();
813
814        g.consume_tokens("root", 100).unwrap();
815
816        let node = g.get_node("root").unwrap();
817        assert_eq!(node.status, NodeStatus::Stopped);
818        assert_eq!(node.stop_reason, Some(StopReason::BudgetExhausted));
819    }
820
821    #[test]
822    fn over_budget_chains_returns_chains_above_threshold() {
823        let mut g = BoundedWorkGraph::default();
824        g.add_root(
825            "first".into(),
826            "a".into(),
827            "capsule:first".into(),
828            budget(1_000, 1_000),
829        )
830        .unwrap();
831        g.add_root(
832            "second".into(),
833            "b".into(),
834            "capsule:second".into(),
835            budget(1_000, 1_000),
836        )
837        .unwrap();
838        g.consume_tokens("first", 750).unwrap();
839
840        let over_budget = g.over_budget_chains(70.0);
841
842        assert_eq!(over_budget.len(), 1);
843        assert_eq!(over_budget[0].chain_id, "first");
844    }
845
846    #[test]
847    fn complete_sets_outcome() {
848        let mut g = BoundedWorkGraph::default();
849        g.add_root(
850            "root".into(),
851            "a".into(),
852            "capsule:r".into(),
853            budget(1000, 1000),
854        )
855        .unwrap();
856        g.complete("root", "outcome:success".into()).unwrap();
857        let node = g.get_node("root").unwrap();
858        assert_eq!(node.status, NodeStatus::Completed);
859        assert_eq!(node.outcome_ref.as_deref(), Some("outcome:success"));
860    }
861}
lean_ctx/core/work_graph.rs

lean_ctx/core/
work_graph.rs
