lean_ctx/core/
session_budget.rs

1//! Session-level cumulative context budget (#1307).
2//!
3//! Tracks total fresh tokens delivered across the session and applies
4//! progressive compression tiers as usage grows:
5//!
6//! | Tier    | Cumulative usage  | Default strategy                    |
7//! |---------|-------------------|-------------------------------------|
8//! | Green   | 0–50%             | No intervention                     |
9//! | Yellow  | 50–75%            | Prefer compressed modes (signatures)|
10//! | Orange  | 75–90%            | Force map/signatures on all reads   |
11//! | Red     | 90–100%           | Reference-only, expand on demand    |
12
13use std::sync::atomic::{AtomicUsize, Ordering};
14
15static CUMULATIVE_TOKENS: AtomicUsize = AtomicUsize::new(0);
16
17/// Session budget pressure tiers.
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub enum PressureTier {
20    Green,
21    Yellow,
22    Orange,
23    Red,
24}
25
26impl PressureTier {
27    /// Suggested read mode downgrade for this pressure tier.
28    pub fn suggested_mode(&self) -> Option<&'static str> {
29        match self {
30            PressureTier::Green => None,
31            PressureTier::Yellow => Some("signatures"),
32            PressureTier::Orange => Some("map"),
33            PressureTier::Red => Some("reference"),
34        }
35    }
36
37    pub fn label(&self) -> &'static str {
38        match self {
39            PressureTier::Green => "green",
40            PressureTier::Yellow => "yellow",
41            PressureTier::Orange => "orange",
42            PressureTier::Red => "red",
43        }
44    }
45}
46
47/// Record `n` fresh tokens delivered in this session.
48pub fn record_delivery(n: usize) {
49    CUMULATIVE_TOKENS.fetch_add(n, Ordering::Relaxed);
50}
51
52/// Current cumulative fresh token count.
53pub fn cumulative_tokens() -> usize {
54    CUMULATIVE_TOKENS.load(Ordering::Relaxed)
55}
56
57/// Current pressure tier based on cumulative usage vs session limit.
58pub fn current_tier(session_limit: usize) -> PressureTier {
59    if session_limit == 0 {
60        return PressureTier::Green;
61    }
62    let used = cumulative_tokens();
63    let ratio = used as f64 / session_limit as f64;
64    if ratio >= 0.90 {
65        PressureTier::Red
66    } else if ratio >= 0.75 {
67        PressureTier::Orange
68    } else if ratio >= 0.50 {
69        PressureTier::Yellow
70    } else {
71        PressureTier::Green
72    }
73}
74
75/// Reset session budget (for testing or session restart).
76pub fn reset() {
77    CUMULATIVE_TOKENS.store(0, Ordering::Relaxed);
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83
84    #[test]
85    fn tier_green_when_below_half() {
86        assert_eq!(current_tier(1000), PressureTier::Green);
87    }
88
89    #[test]
90    fn tier_transitions() {
91        reset();
92        assert_eq!(current_tier(1000), PressureTier::Green);
93
94        record_delivery(500);
95        assert_eq!(current_tier(1000), PressureTier::Yellow);
96
97        record_delivery(250);
98        assert_eq!(current_tier(1000), PressureTier::Orange);
99
100        record_delivery(150);
101        assert_eq!(current_tier(1000), PressureTier::Red);
102
103        reset();
104        assert_eq!(current_tier(1000), PressureTier::Green);
105    }
106
107    #[test]
108    fn tier_green_when_unlimited() {
109        record_delivery(999_999);
110        assert_eq!(current_tier(0), PressureTier::Green);
111        reset();
112    }
113
114    #[test]
115    fn suggested_modes() {
116        assert_eq!(PressureTier::Green.suggested_mode(), None);
117        assert_eq!(PressureTier::Yellow.suggested_mode(), Some("signatures"));
118        assert_eq!(PressureTier::Orange.suggested_mode(), Some("map"));
119        assert_eq!(PressureTier::Red.suggested_mode(), Some("reference"));
120    }
121}
lean_ctx/core/session_budget.rs

lean_ctx/core/
session_budget.rs
