| 1 | //! Unified context-budget math for the TUI. |
| 2 | //! |
| 3 | //! Given a model's context window, the current input token estimate, and a |
| 4 | //! configured output cap, [`ContextBudget`] derives the four numbers the rest |
| 5 | //! of the app needs to reason about a turn: |
| 6 | //! |
| 7 | //! * **available input budget** — how many input tokens may still be spent |
| 8 | //! after reserving room for the model's output; |
| 9 | //! * **output token cap** — the output reservation actually used to compute |
| 10 | //! that budget (clamped so it never starves the window); |
| 11 | //! * **compaction trigger** — the input-token level at which compaction |
| 12 | //! should be suggested (default: ~75% of the full window, clamped to the |
| 13 | //! spendable input ceiling); |
| 14 | //! * **[`PressureLevel`]** — a coarse Low/Medium/High/Critical signal the UI |
| 15 | //! can render without re-deriving thresholds. |
| 16 | //! |
| 17 | //! This module is the budget-math *foundation*. It is intentionally pure (no |
| 18 | //! I/O, no clock, no engine/config types) so it can be unit-tested in isolation. |
| 19 | //! Its consumers live outside it and call in, never the reverse: `route_budget` |
| 20 | //! and `core::engine::context` for [`ContextBudget`], `context_report` for |
| 21 | //! [`PressureLevel`]. |
| 22 | //! |
| 23 | //! ### Route ceilings stay independent |
| 24 | //! |
| 25 | //! A route can publish a total context window, an output ceiling, and a |
| 26 | //! separate input ceiling. [`ContextBudget::new_with_input_limit`] intersects |
| 27 | //! all three without treating one as a substitute for another: the output |
| 28 | //! reservation is the exact route-effective wire request, the total-window |
| 29 | //! arithmetic remains saturating, and a concrete input limit clamps the final |
| 30 | //! spendable ceiling and compaction trigger. |
| 31 | |
| 32 | // This module IS wired. `ContextBudget` is consumed by `route_budget.rs` and |
| 33 | // `core/engine/context.rs`; `PressureLevel` by `context_report.rs`. It sits on |
| 34 | // the do-not-delete list in AGENTS.md because a blanket `allow(dead_code)` here, |
| 35 | // plus a comment that used to claim the module was "not yet referenced," taught |
| 36 | // several dead-code audits to propose deleting a live file. The suppression is |
| 37 | // now `#[cfg_attr(not(test), expect(dead_code))]` on the three methods unused |
| 38 | // outside tests, so a suppression that stops matching the lint fails the |
| 39 | // build instead of hiding behind a module-wide waiver. Tests call them, so |
| 40 | // a bare `#[expect(dead_code)]` would be unfulfilled under `cfg(test)`. |
| 41 | |
| 42 | /// Fraction of the window, expressed as a percentage, at or above which |
| 43 | /// compaction should be suggested. Mirrors the "high" pressure boundary the |
| 44 | /// existing context report uses for its diagnostic label, rounded up to the |
| 45 | /// conventional three-quarters-full trigger. |
| 46 | pub const DEFAULT_COMPACTION_TRIGGER_PERCENT: f64 = 75.0; |
| 47 | |
| 48 | /// Percentage of the window at or above which pressure is [`PressureLevel::Critical`]. |
| 49 | pub const CRITICAL_PRESSURE_PERCENT: f64 = 90.0; |
| 50 | |
| 51 | /// Percentage of the window at or above which pressure is [`PressureLevel::High`]. |
| 52 | /// Pressure and the requested compaction trigger are window-relative. The |
| 53 | /// trigger is clamped to the spendable input ceiling, so it can fire before |
| 54 | /// this UI boundary when output reservation consumes substantial window space. |
| 55 | pub const HIGH_PRESSURE_PERCENT: f64 = DEFAULT_COMPACTION_TRIGGER_PERCENT; |
| 56 | |
| 57 | /// Percentage of the window at or above which pressure is [`PressureLevel::Medium`]. |
| 58 | /// Matches the "moderate" boundary of the existing diagnostic report. |
| 59 | pub const MEDIUM_PRESSURE_PERCENT: f64 = 40.0; |
| 60 | |
| 61 | /// Safety headroom (tokens) subtracted from the window in addition to the |
| 62 | /// reserved output, to avoid bumping a provider's hard limit. Matches the |
| 63 | /// engine's `CONTEXT_HEADROOM_TOKENS`. |
| 64 | pub const CONTEXT_HEADROOM_TOKENS: u64 = 1_024; |
| 65 | |
| 66 | /// Smallest input budget (tokens) [`ContextBudget`] will report for any window |
| 67 | /// large enough to hold it. The output cap is clamped down as needed to |
| 68 | /// preserve this much input room, so a generous configured output cap can never |
| 69 | /// drive the available input budget to zero on a usable window. |
| 70 | pub const MIN_INPUT_BUDGET_TOKENS: u64 = 1_024; |
| 71 | |
| 72 | /// Coarse, UI-facing description of how full the context window is. |
| 73 | /// |
| 74 | /// Ordered from least to most pressure so the variants can be compared |
| 75 | /// (`level >= PressureLevel::High`) and so the derived `Ord` matches intuition. |
| 76 | #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] |
| 77 | pub enum PressureLevel { |
| 78 | /// Plenty of room; nothing to surface. |
| 79 | Low, |
| 80 | /// Noticeably filling up; informational. |
| 81 | Medium, |
| 82 | /// At or past the compaction trigger; suggest compaction. |
| 83 | High, |
| 84 | /// Near the window limit; compaction/clear is urgent. |
| 85 | Critical, |
| 86 | } |
| 87 | |
| 88 | impl PressureLevel { |
| 89 | /// Classify a window-usage percentage (0.0..=100.0) into a level. |
| 90 | /// |
| 91 | /// Inputs outside the range are clamped, so callers may pass a raw |
| 92 | /// percentage without pre-validating it. |
| 93 | #[must_use] |
| 94 | pub fn from_usage_percent(percent: f64) -> Self { |
| 95 | let percent = percent.clamp(0.0, 100.0); |
| 96 | if percent >= CRITICAL_PRESSURE_PERCENT { |
| 97 | PressureLevel::Critical |
| 98 | } else if percent >= HIGH_PRESSURE_PERCENT { |
| 99 | PressureLevel::High |
| 100 | } else if percent >= MEDIUM_PRESSURE_PERCENT { |
| 101 | PressureLevel::Medium |
| 102 | } else { |
| 103 | PressureLevel::Low |
| 104 | } |
| 105 | } |
| 106 | |
| 107 | /// Lowercase, stable label suitable for status lines and logs. |
| 108 | /// |
| 109 | /// Kept aligned with the existing context-report vocabulary |
| 110 | /// (`low`/`moderate`/`high`/`critical`). |
| 111 | #[must_use] |
| 112 | pub const fn label(self) -> &'static str { |
| 113 | match self { |
| 114 | PressureLevel::Low => "low", |
| 115 | PressureLevel::Medium => "moderate", |
| 116 | PressureLevel::High => "high", |
| 117 | PressureLevel::Critical => "critical", |
| 118 | } |
| 119 | } |
| 120 | |
| 121 | /// Whether this level is at or past the point where compaction should be |
| 122 | /// suggested to the user. |
| 123 | /// |
| 124 | /// Unused by the engine today; kept as the pressure-level counterpart of |
| 125 | /// `ContextBudget::should_compact` so both live next to their thresholds. |
| 126 | #[must_use] |
| 127 | pub const fn suggests_compaction(self) -> bool { |
| 128 | matches!(self, PressureLevel::High | PressureLevel::Critical) |
| 129 | } |
| 130 | } |
| 131 | |
| 132 | /// A computed snapshot of how a turn's input sits against a model's context |
| 133 | /// window, plus the derived output cap, compaction trigger, and pressure level. |
| 134 | /// |
| 135 | /// Construct via [`ContextBudget::new`]. All fields are token counts unless the |
| 136 | /// name says otherwise. The struct is `Copy` and holds no borrowed data so it |
| 137 | /// can be cached on UI state cheaply. |
| 138 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] |
| 139 | pub struct ContextBudget { |
| 140 | /// Total context window for the active route (input + output), in tokens. |
| 141 | pub window_tokens: u64, |
| 142 | /// Current estimated input tokens already committed to the turn. |
| 143 | pub input_tokens: u64, |
| 144 | /// Output tokens reserved (and thus the effective output cap) for the turn. |
| 145 | /// Derived from the configured cap, clamped to fit the window while leaving |
| 146 | /// at least [`MIN_INPUT_BUDGET_TOKENS`] of input room. |
| 147 | pub output_cap_tokens: u64, |
| 148 | /// Spendable input ceiling for the turn (`window - output_cap - headroom`, |
| 149 | /// saturating at 0). Compaction percentages use the full window and clamp |
| 150 | /// their resulting token threshold to this ceiling. |
| 151 | pub input_budget_ceiling: u64, |
| 152 | /// Input tokens still available before hitting the reserved boundary |
| 153 | /// (`input_budget_ceiling - input`, saturating at 0). |
| 154 | pub available_input_tokens: u64, |
| 155 | /// Input-token level at or above which compaction should be suggested |
| 156 | /// (`DEFAULT_COMPACTION_TRIGGER_PERCENT` of the window, clamped to the |
| 157 | /// input budget ceiling). |
| 158 | pub compaction_trigger_tokens: u64, |
| 159 | /// Coarse pressure level derived from window usage. |
| 160 | pub pressure: PressureLevel, |
| 161 | } |
| 162 | |
| 163 | impl ContextBudget { |
| 164 | /// Build a budget snapshot for a route. |
| 165 | /// |
| 166 | /// * `window_tokens` — the route-effective context window (input + output). |
| 167 | /// * `input_tokens` — current estimated input tokens for the turn. |
| 168 | /// * `configured_output_cap` — the route-effective output request the |
| 169 | /// caller needs to reserve. It is clamped down so it never consumes the |
| 170 | /// headroom or the minimum input budget; on a window too small to hold |
| 171 | /// even the minimum input budget plus headroom, the cap collapses to |
| 172 | /// whatever is left (possibly zero). |
| 173 | /// |
| 174 | /// Never panics and never underflows: all arithmetic saturates. |
| 175 | #[must_use] |
| 176 | pub fn new(window_tokens: u64, input_tokens: u64, configured_output_cap: u64) -> Self { |
| 177 | Self::new_with_input_limit(window_tokens, input_tokens, configured_output_cap, None) |
| 178 | } |
| 179 | |
| 180 | /// Build a budget snapshot and intersect it with a provider's independent |
| 181 | /// hard input ceiling when one is published by the resolved route. |
| 182 | #[must_use] |
| 183 | pub fn new_with_input_limit( |
| 184 | window_tokens: u64, |
| 185 | input_tokens: u64, |
| 186 | configured_output_cap: u64, |
| 187 | input_limit_tokens: Option<u64>, |
| 188 | ) -> Self { |
| 189 | let output_cap_tokens = clamp_output_cap(window_tokens, configured_output_cap); |
| 190 | |
| 191 | // Reserve output + safety headroom; whatever remains is spendable input. |
| 192 | let reserved = output_cap_tokens.saturating_add(CONTEXT_HEADROOM_TOKENS); |
| 193 | let window_input_ceiling = window_tokens.saturating_sub(reserved); |
| 194 | let input_budget_ceiling = input_limit_tokens |
| 195 | .filter(|limit| *limit > 0) |
| 196 | .map_or(window_input_ceiling, |limit| { |
| 197 | window_input_ceiling.min(limit) |
| 198 | }); |
| 199 | let available_input_tokens = input_budget_ceiling.saturating_sub(input_tokens); |
| 200 | |
| 201 | let compaction_trigger_tokens = |
| 202 | percent_of(window_tokens, DEFAULT_COMPACTION_TRIGGER_PERCENT).min(input_budget_ceiling); |
| 203 | |
| 204 | let pressure = |
| 205 | PressureLevel::from_usage_percent(usage_percent(window_tokens, input_tokens)); |
| 206 | |
| 207 | ContextBudget { |
| 208 | window_tokens, |
| 209 | input_tokens, |
| 210 | output_cap_tokens, |
| 211 | input_budget_ceiling, |
| 212 | available_input_tokens, |
| 213 | compaction_trigger_tokens, |
| 214 | pressure, |
| 215 | } |
| 216 | } |
| 217 | |
| 218 | /// Derive a compaction trigger from a window percentage. |
| 219 | /// |
| 220 | /// The percentage means what the user-facing context meter says it means: |
| 221 | /// a fraction of the full route window (`80` on a 1M window is 800K input |
| 222 | /// tokens before the route ceiling clamp). One internal clamp applies: |
| 223 | /// |
| 224 | /// ```text |
| 225 | /// trigger = min(window × percent, window − output reservation − headroom) |
| 226 | /// ``` |
| 227 | /// |
| 228 | /// so a late percentage can never push the trigger past the spendable |
| 229 | /// input ceiling and overflow the provider window. Previously the |
| 230 | /// percentage was applied to the ceiling itself, which silently pulled an |
| 231 | /// "80%" setting down to ~60% of the window on routes with large output |
| 232 | /// reservations; UI surfaces should disclose the clamp when it engages |
| 233 | /// instead of redefining what the percentage means. |
| 234 | #[must_use] |
| 235 | pub fn compaction_trigger_for_percent(&self, percent: f64) -> u64 { |
| 236 | percent_of(self.window_tokens, percent).min(self.input_budget_ceiling) |
| 237 | } |
| 238 | |
| 239 | /// Fraction of the window currently consumed by input, as a percentage in |
| 240 | /// `0.0..=100.0`. A zero window reports `0.0` rather than dividing by zero. |
| 241 | #[must_use] |
| 242 | pub fn usage_percent(&self) -> f64 { |
| 243 | usage_percent(self.window_tokens, self.input_tokens) |
| 244 | } |
| 245 | |
| 246 | /// Whether current input has reached the compaction trigger and compaction |
| 247 | /// should be suggested. |
| 248 | #[must_use] |
| 249 | pub fn should_compact(&self) -> bool { |
| 250 | self.window_tokens > 0 && self.input_tokens >= self.compaction_trigger_tokens |
| 251 | } |
| 252 | |
| 253 | /// Whether another `additional_input_tokens` of input would fit within the |
| 254 | /// available budget (i.e. not exceed the reserved boundary). |
| 255 | #[must_use] |
| 256 | pub fn fits_additional(&self, additional_input_tokens: u64) -> bool { |
| 257 | additional_input_tokens <= self.available_input_tokens |
| 258 | } |
| 259 | } |
| 260 | |
| 261 | /// Clamp a desired output cap so it fits the window while preserving at least |
| 262 | /// [`MIN_INPUT_BUDGET_TOKENS`] of input room plus [`CONTEXT_HEADROOM_TOKENS`]. |
| 263 | /// |
| 264 | /// On a window too small to hold even that floor, returns whatever room is left |
| 265 | /// after the headroom (possibly zero) rather than underflowing. |
| 266 | fn clamp_output_cap(window_tokens: u64, configured_output_cap: u64) -> u64 { |
| 267 | // The most output we can reserve and still keep the input floor + headroom. |
| 268 | let reserved_floor = MIN_INPUT_BUDGET_TOKENS.saturating_add(CONTEXT_HEADROOM_TOKENS); |
| 269 | let max_output = window_tokens.saturating_sub(reserved_floor); |
| 270 | configured_output_cap.min(max_output) |
| 271 | } |
| 272 | |
| 273 | /// Window usage as a percentage in `0.0..=100.0`. Zero window -> `0.0`. |
| 274 | fn usage_percent(window_tokens: u64, input_tokens: u64) -> f64 { |
| 275 | if window_tokens == 0 { |
| 276 | return 0.0; |
| 277 | } |
| 278 | ((input_tokens as f64 / window_tokens as f64) * 100.0).clamp(0.0, 100.0) |
| 279 | } |
| 280 | |
| 281 | /// `percent`% of `window_tokens`, rounded to the nearest token. Saturates at |
| 282 | /// `u64::MAX` and treats out-of-range percentages by clamping to `0.0..=100.0`. |
| 283 | fn percent_of(window_tokens: u64, percent: f64) -> u64 { |
| 284 | let percent = percent.clamp(0.0, 100.0); |
| 285 | let value = (window_tokens as f64) * (percent / 100.0); |
| 286 | // `as u64` saturates on overflow and floors; add 0.5 to round to nearest. |
| 287 | (value + 0.5) as u64 |
| 288 | } |
| 289 | |
| 290 | #[cfg(test)] |
| 291 | mod tests { |
| 292 | use super::*; |
| 293 | |
| 294 | /// A representative spread of real-world windows: a tight self-hosted |
| 295 | /// deployment, common provider sizes, and a V4-class 1M window. |
| 296 | const WINDOWS: &[u64] = &[8_192, 32_768, 131_072, 262_144, 1_048_576]; |
| 297 | |
| 298 | // -- PressureLevel boundaries ------------------------------------------ |
| 299 | |
| 300 | #[test] |
| 301 | fn pressure_level_boundaries_are_inclusive_lower_bounds() { |
| 302 | assert_eq!(PressureLevel::from_usage_percent(0.0), PressureLevel::Low); |
| 303 | assert_eq!(PressureLevel::from_usage_percent(39.9), PressureLevel::Low); |
| 304 | // 40% is the moderate boundary. |
| 305 | assert_eq!( |
| 306 | PressureLevel::from_usage_percent(40.0), |
| 307 | PressureLevel::Medium |
| 308 | ); |
| 309 | assert_eq!( |
| 310 | PressureLevel::from_usage_percent(74.9), |
| 311 | PressureLevel::Medium |
| 312 | ); |
| 313 | // 75% is the high / compaction boundary. |
| 314 | assert_eq!(PressureLevel::from_usage_percent(75.0), PressureLevel::High); |
| 315 | assert_eq!(PressureLevel::from_usage_percent(89.9), PressureLevel::High); |
| 316 | // 90% is the critical boundary. |
| 317 | assert_eq!( |
| 318 | PressureLevel::from_usage_percent(90.0), |
| 319 | PressureLevel::Critical |
| 320 | ); |
| 321 | assert_eq!( |
| 322 | PressureLevel::from_usage_percent(100.0), |
| 323 | PressureLevel::Critical |
| 324 | ); |
| 325 | } |
| 326 | |
| 327 | #[test] |
| 328 | fn pressure_level_clamps_out_of_range_inputs() { |
| 329 | assert_eq!(PressureLevel::from_usage_percent(-10.0), PressureLevel::Low); |
| 330 | assert_eq!( |
| 331 | PressureLevel::from_usage_percent(150.0), |
| 332 | PressureLevel::Critical |
| 333 | ); |
| 334 | assert_eq!( |
| 335 | PressureLevel::from_usage_percent(f64::INFINITY), |
| 336 | PressureLevel::Critical |
| 337 | ); |
| 338 | } |
| 339 | |
| 340 | #[test] |
| 341 | fn pressure_level_ordering_and_helpers() { |
| 342 | assert!(PressureLevel::Low < PressureLevel::Medium); |
| 343 | assert!(PressureLevel::Medium < PressureLevel::High); |
| 344 | assert!(PressureLevel::High < PressureLevel::Critical); |
| 345 | |
| 346 | assert!(!PressureLevel::Low.suggests_compaction()); |
| 347 | assert!(!PressureLevel::Medium.suggests_compaction()); |
| 348 | assert!(PressureLevel::High.suggests_compaction()); |
| 349 | assert!(PressureLevel::Critical.suggests_compaction()); |
| 350 | |
| 351 | assert_eq!(PressureLevel::Low.label(), "low"); |
| 352 | assert_eq!(PressureLevel::Medium.label(), "moderate"); |
| 353 | assert_eq!(PressureLevel::High.label(), "high"); |
| 354 | assert_eq!(PressureLevel::Critical.label(), "critical"); |
| 355 | } |
| 356 | |
| 357 | // -- Compaction trigger ------------------------------------------------- |
| 358 | |
| 359 | #[test] |
| 360 | fn compaction_trigger_is_window_percent_clamped_to_input_ceiling() { |
| 361 | for &window in WINDOWS { |
| 362 | let budget = ContextBudget::new(window, 0, 64_000); |
| 363 | let expected = percent_of(window, DEFAULT_COMPACTION_TRIGGER_PERCENT) |
| 364 | .min(budget.input_budget_ceiling); |
| 365 | assert_eq!( |
| 366 | budget.compaction_trigger_tokens, expected, |
| 367 | "window {window}: trigger should be 75% of the window, clamped to the ceiling" |
| 368 | ); |
| 369 | assert!(budget.compaction_trigger_tokens <= budget.input_budget_ceiling); |
| 370 | } |
| 371 | } |
| 372 | |
| 373 | #[test] |
| 374 | fn should_compact_flips_at_the_trigger() { |
| 375 | let window = 1_048_576; |
| 376 | let cap = 262_144; |
| 377 | let trigger = ContextBudget::new(window, 0, cap).compaction_trigger_tokens; |
| 378 | assert!(trigger > 0); |
| 379 | |
| 380 | let below = ContextBudget::new(window, trigger - 1, cap); |
| 381 | assert!(!below.should_compact()); |
| 382 | |
| 383 | let at = ContextBudget::new(window, trigger, cap); |
| 384 | assert!(at.should_compact()); |
| 385 | |
| 386 | let above = ContextBudget::new(window, trigger + 1, cap); |
| 387 | assert!(above.should_compact()); |
| 388 | } |
| 389 | |
| 390 | #[test] |
| 391 | fn trigger_never_exceeds_input_ceiling() { |
| 392 | const WINDOWS: &[u64] = &[ |
| 393 | 0, 512, 8_192, 65_536, 131_072, 262_144, 500_000, 1_048_576, 2_000_000, |
| 394 | ]; |
| 395 | const CAPS: &[u64] = &[0, 1_024, 64_000, 131_072, 262_144, 1_000_000]; |
| 396 | const INPUTS: &[u64] = &[0, 1_024, 200_000, 2_000_000]; |
| 397 | const PERCENTS: &[f64] = &[10.0, 50.0, 75.0, 80.0, 95.0, 100.0]; |
| 398 | |
| 399 | for &window in WINDOWS { |
| 400 | for &cap in CAPS { |
| 401 | for &input in INPUTS { |
| 402 | let budget = ContextBudget::new(window, input, cap); |
| 403 | assert!(budget.compaction_trigger_tokens <= budget.input_budget_ceiling); |
| 404 | for &percent in PERCENTS { |
| 405 | assert!( |
| 406 | budget.compaction_trigger_for_percent(percent) |
| 407 | <= budget.input_budget_ceiling, |
| 408 | "window={window} cap={cap} input={input} percent={percent}" |
| 409 | ); |
| 410 | } |
| 411 | } |
| 412 | } |
| 413 | } |
| 414 | } |
| 415 | |
| 416 | #[test] |
| 417 | fn zero_window_never_suggests_compaction() { |
| 418 | let budget = ContextBudget::new(0, 0, 64_000); |
| 419 | assert_eq!(budget.compaction_trigger_tokens, 0); |
| 420 | assert_eq!(budget.input_budget_ceiling, 0); |
| 421 | assert!(!budget.should_compact()); |
| 422 | assert_eq!(budget.pressure, PressureLevel::Low); |
| 423 | assert_eq!(budget.available_input_tokens, 0); |
| 424 | assert_eq!(budget.usage_percent(), 0.0); |
| 425 | } |
| 426 | |
| 427 | // -- Output cap clamping & available budget ---------------------------- |
| 428 | |
| 429 | #[test] |
| 430 | fn output_cap_is_preserved_when_window_is_roomy() { |
| 431 | // 1M window, 64K configured cap: cap fits comfortably. |
| 432 | let budget = ContextBudget::new(1_048_576, 0, 64_000); |
| 433 | assert_eq!(budget.output_cap_tokens, 64_000); |
| 434 | // available = window - cap - headroom - input |
| 435 | let expected = 1_048_576 - 64_000 - CONTEXT_HEADROOM_TOKENS; |
| 436 | assert_eq!(budget.available_input_tokens, expected); |
| 437 | } |
| 438 | |
| 439 | #[test] |
| 440 | fn output_cap_is_clamped_to_protect_input_floor_on_small_window() { |
| 441 | // This is the engine's hard-won lesson: a generous output reservation |
| 442 | // on a small window must not underflow the input budget. An 8,192-token |
| 443 | // window with a 262,144-token desired cap must still leave the input |
| 444 | // floor available rather than collapsing to zero or wrapping. |
| 445 | let window = 8_192u64; |
| 446 | let budget = ContextBudget::new(window, 0, 262_144); |
| 447 | |
| 448 | let reserved_floor = MIN_INPUT_BUDGET_TOKENS + CONTEXT_HEADROOM_TOKENS; |
| 449 | let expected_cap = window - reserved_floor; |
| 450 | assert_eq!(budget.output_cap_tokens, expected_cap); |
| 451 | // With zero input committed, the whole remaining budget is available |
| 452 | // and is at least the protected floor. |
| 453 | assert!(budget.available_input_tokens >= MIN_INPUT_BUDGET_TOKENS); |
| 454 | assert_eq!( |
| 455 | budget.available_input_tokens, |
| 456 | window - budget.output_cap_tokens - CONTEXT_HEADROOM_TOKENS |
| 457 | ); |
| 458 | } |
| 459 | |
| 460 | #[test] |
| 461 | fn tiny_window_below_floor_saturates_without_panic() { |
| 462 | // Window smaller than the protected floor: cap collapses to 0 and the |
| 463 | // available budget saturates at 0 instead of underflowing. |
| 464 | let window = 512u64; // < MIN_INPUT_BUDGET_TOKENS + headroom |
| 465 | let budget = ContextBudget::new(window, 100, 262_144); |
| 466 | assert_eq!(budget.output_cap_tokens, 0); |
| 467 | assert_eq!(budget.available_input_tokens, 0); |
| 468 | // Usage still computes a sane percentage. |
| 469 | assert!((budget.usage_percent() - (100.0 / 512.0 * 100.0)).abs() < 1e-9); |
| 470 | } |
| 471 | |
| 472 | #[test] |
| 473 | fn available_budget_saturates_when_input_exceeds_ceiling() { |
| 474 | let window = 131_072u64; |
| 475 | let cap = 32_000u64; |
| 476 | // Commit far more input than the window holds. |
| 477 | let budget = ContextBudget::new(window, window * 2, cap); |
| 478 | assert_eq!(budget.available_input_tokens, 0); |
| 479 | assert!(!budget.fits_additional(1)); |
| 480 | assert!(budget.fits_additional(0)); |
| 481 | // Usage is clamped to 100%. |
| 482 | assert_eq!(budget.usage_percent(), 100.0); |
| 483 | assert_eq!(budget.pressure, PressureLevel::Critical); |
| 484 | } |
| 485 | |
| 486 | #[test] |
| 487 | fn fits_additional_respects_the_reserved_boundary() { |
| 488 | let window = 262_144u64; |
| 489 | let cap = 64_000u64; |
| 490 | let budget = ContextBudget::new(window, 100_000, cap); |
| 491 | let room = budget.available_input_tokens; |
| 492 | assert!(budget.fits_additional(room)); |
| 493 | assert!(!budget.fits_additional(room + 1)); |
| 494 | } |
| 495 | |
| 496 | // -- Usage percent & pressure across window sizes ---------------------- |
| 497 | |
| 498 | #[test] |
| 499 | fn usage_percent_is_proportional_across_window_sizes() { |
| 500 | for &window in WINDOWS { |
| 501 | // Half-full should read ~50% and classify as Medium regardless of |
| 502 | // absolute window size. |
| 503 | let half = window / 2; |
| 504 | let budget = ContextBudget::new(window, half, 64_000); |
| 505 | let pct = budget.usage_percent(); |
| 506 | assert!( |
| 507 | (pct - 50.0).abs() < 0.5, |
| 508 | "window {window}: half-full should be ~50%, got {pct}" |
| 509 | ); |
| 510 | assert_eq!(budget.pressure, PressureLevel::Medium); |
| 511 | } |
| 512 | } |
| 513 | |
| 514 | #[test] |
| 515 | fn pressure_tracks_input_growth_on_a_1m_window() { |
| 516 | let window = 1_048_576u64; |
| 517 | let cap = 262_144u64; |
| 518 | |
| 519 | let low = ContextBudget::new(window, percent_of(window, 10.0), cap); |
| 520 | assert_eq!(low.pressure, PressureLevel::Low); |
| 521 | |
| 522 | let medium = ContextBudget::new(window, percent_of(window, 50.0), cap); |
| 523 | assert_eq!(medium.pressure, PressureLevel::Medium); |
| 524 | |
| 525 | let high = ContextBudget::new(window, percent_of(window, 80.0), cap); |
| 526 | assert_eq!(high.pressure, PressureLevel::High); |
| 527 | assert!(high.should_compact()); |
| 528 | |
| 529 | let critical = ContextBudget::new(window, percent_of(window, 95.0), cap); |
| 530 | assert_eq!(critical.pressure, PressureLevel::Critical); |
| 531 | } |
| 532 | |
| 533 | #[test] |
| 534 | fn snapshot_fields_are_internally_consistent() { |
| 535 | for &window in WINDOWS { |
| 536 | for &input in &[0u64, window / 4, window / 2, window] { |
| 537 | let budget = ContextBudget::new(window, input, 64_000); |
| 538 | // Field mirrors the constructor arguments. |
| 539 | assert_eq!(budget.window_tokens, window); |
| 540 | assert_eq!(budget.input_tokens, input); |
| 541 | // available + input never claims more than the window minus |
| 542 | // reserved output and headroom. |
| 543 | let ceiling = window |
| 544 | .saturating_sub(budget.output_cap_tokens) |
| 545 | .saturating_sub(CONTEXT_HEADROOM_TOKENS); |
| 546 | assert!(budget.available_input_tokens <= ceiling); |
| 547 | // Pressure agrees with the standalone usage percentage. |
| 548 | assert_eq!( |
| 549 | budget.pressure, |
| 550 | PressureLevel::from_usage_percent(budget.usage_percent()) |
| 551 | ); |
| 552 | } |
| 553 | } |
| 554 | } |
| 555 | |
| 556 | #[test] |
| 557 | fn percent_of_rounds_to_nearest_token() { |
| 558 | assert_eq!(percent_of(0, 75.0), 0); |
| 559 | assert_eq!(percent_of(100, 0.0), 0); |
| 560 | assert_eq!(percent_of(100, 100.0), 100); |
| 561 | assert_eq!(percent_of(100, 75.0), 75); |
| 562 | // 3 * 0.75 = 2.25 -> rounds to 2. |
| 563 | assert_eq!(percent_of(3, 75.0), 2); |
| 564 | // 2 * 0.75 = 1.5 -> rounds to 2. |
| 565 | assert_eq!(percent_of(2, 75.0), 2); |
| 566 | } |
| 567 | |
| 568 | #[test] |
| 569 | fn budget_is_copy_and_comparable() { |
| 570 | let a = ContextBudget::new(131_072, 1_000, 32_000); |
| 571 | let b = a; // Copy, not move. |
| 572 | assert_eq!(a, b); |
| 573 | assert_eq!(a.window_tokens, b.window_tokens); |
| 574 | } |
| 575 | } |
| 576 |