| 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 spendable input ceiling); |
| 13 | //! * **[`PressureLevel`]** — a coarse Low/Medium/High/Critical signal the UI |
| 14 | //! can render without re-deriving thresholds. |
| 15 | //! |
| 16 | //! This module is the budget-math *foundation*. It is intentionally pure (no |
| 17 | //! I/O, no clock, no engine/config types) so it can be unit-tested in isolation. |
| 18 | //! Its consumers live outside it and call in, never the reverse: `route_budget` |
| 19 | //! and `core::engine::context` for [`ContextBudget`], `context_report` for |
| 20 | //! [`PressureLevel`]. |
| 21 | //! |
| 22 | //! ### Why the output reservation is window-dependent |
| 23 | //! |
| 24 | //! The engine's existing input-budget helper |
| 25 | //! (`core::engine::context::context_input_budget_for_window`) computes |
| 26 | //! `window - reserved_output - headroom` and learned the hard way that |
| 27 | //! reserving a large fixed output (262K for V4-class interleaved thinking) on a |
| 28 | //! *small* self-hosted window (e.g. a 256K vLLM deployment) underflows to a |
| 29 | //! negative budget and silently disables every preflight/recovery path. We |
| 30 | //! mirror that lesson here with saturating arithmetic and an output cap that is |
| 31 | //! always clamped to leave at least [`MIN_INPUT_BUDGET_TOKENS`] of input room, |
| 32 | //! so the budget can never collapse to zero on a legitimately sized window. |
| 33 | |
| 34 | // This module IS wired. `ContextBudget` is consumed by `route_budget.rs` and |
| 35 | // `core/engine/context.rs`; `PressureLevel` by `context_report.rs`. It sits on |
| 36 | // the do-not-delete list in AGENTS.md because a blanket `allow(dead_code)` here, |
| 37 | // plus a comment that used to claim the module was "not yet referenced," taught |
| 38 | // several dead-code audits to propose deleting a live file. The allow is now |
| 39 | // per-item on the three genuinely-unused methods, so anything that goes dead |
| 40 | // here shows up as a warning instead of hiding behind a module-wide waiver. |
| 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 remains a window-relative UI signal. The actual compaction trigger |
| 53 | /// is relative to the spendable input ceiling, so these thresholds can diverge |
| 54 | /// when output reservation consumes a substantial part of the window. |
| 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 | #[allow(dead_code)] |
| 128 | pub const fn suggests_compaction(self) -> bool { |
| 129 | matches!(self, PressureLevel::High | PressureLevel::Critical) |
| 130 | } |
| 131 | } |
| 132 | |
| 133 | /// A computed snapshot of how a turn's input sits against a model's context |
| 134 | /// window, plus the derived output cap, compaction trigger, and pressure level. |
| 135 | /// |
| 136 | /// Construct via [`ContextBudget::new`]. All fields are token counts unless the |
| 137 | /// name says otherwise. The struct is `Copy` and holds no borrowed data so it |
| 138 | /// can be cached on UI state cheaply. |
| 139 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] |
| 140 | pub struct ContextBudget { |
| 141 | /// Total context window for the active route (input + output), in tokens. |
| 142 | pub window_tokens: u64, |
| 143 | /// Current estimated input tokens already committed to the turn. |
| 144 | pub input_tokens: u64, |
| 145 | /// Output tokens reserved (and thus the effective output cap) for the turn. |
| 146 | /// Derived from the configured cap, clamped to fit the window while leaving |
| 147 | /// at least [`MIN_INPUT_BUDGET_TOKENS`] of input room. |
| 148 | pub output_cap_tokens: u64, |
| 149 | /// Spendable input ceiling for the turn (`window - output_cap - headroom`, |
| 150 | /// saturating at 0). Compaction percentages use this denominator rather |
| 151 | /// than the raw input-plus-output context window. |
| 152 | pub input_budget_ceiling: u64, |
| 153 | /// Input tokens still available before hitting the reserved boundary |
| 154 | /// (`input_budget_ceiling - input`, saturating at 0). |
| 155 | pub available_input_tokens: u64, |
| 156 | /// Input-token level at or above which compaction should be suggested |
| 157 | /// (`DEFAULT_COMPACTION_TRIGGER_PERCENT` of the 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 output reservation the caller would like |
| 169 | /// (e.g. the engine's `TURN_MAX_OUTPUT_TOKENS`). It is clamped down so it |
| 170 | /// never consumes the headroom or the minimum input budget; on a window |
| 171 | /// too small to hold even the minimum input budget plus headroom, the cap |
| 172 | /// collapses to 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 | let output_cap_tokens = clamp_output_cap(window_tokens, configured_output_cap); |
| 178 | |
| 179 | // Reserve output + safety headroom; whatever remains is spendable input. |
| 180 | let reserved = output_cap_tokens.saturating_add(CONTEXT_HEADROOM_TOKENS); |
| 181 | let input_budget_ceiling = window_tokens.saturating_sub(reserved); |
| 182 | let available_input_tokens = input_budget_ceiling.saturating_sub(input_tokens); |
| 183 | |
| 184 | let compaction_trigger_tokens = |
| 185 | percent_of(input_budget_ceiling, DEFAULT_COMPACTION_TRIGGER_PERCENT); |
| 186 | |
| 187 | let pressure = |
| 188 | PressureLevel::from_usage_percent(usage_percent(window_tokens, input_tokens)); |
| 189 | |
| 190 | ContextBudget { |
| 191 | window_tokens, |
| 192 | input_tokens, |
| 193 | output_cap_tokens, |
| 194 | input_budget_ceiling, |
| 195 | available_input_tokens, |
| 196 | compaction_trigger_tokens, |
| 197 | pressure, |
| 198 | } |
| 199 | } |
| 200 | |
| 201 | /// Derive a compaction trigger from the spendable input ceiling. |
| 202 | /// |
| 203 | /// Applying a clamped percentage to the ceiling guarantees that the |
| 204 | /// trigger cannot sit beyond the input room available to the route. |
| 205 | #[must_use] |
| 206 | pub fn compaction_trigger_for_percent(&self, percent: f64) -> u64 { |
| 207 | percent_of(self.input_budget_ceiling, percent) |
| 208 | } |
| 209 | |
| 210 | /// Fraction of the window currently consumed by input, as a percentage in |
| 211 | /// `0.0..=100.0`. A zero window reports `0.0` rather than dividing by zero. |
| 212 | #[must_use] |
| 213 | pub fn usage_percent(&self) -> f64 { |
| 214 | usage_percent(self.window_tokens, self.input_tokens) |
| 215 | } |
| 216 | |
| 217 | /// Whether current input has reached the compaction trigger and compaction |
| 218 | /// should be suggested. |
| 219 | #[must_use] |
| 220 | #[allow(dead_code)] |
| 221 | pub fn should_compact(&self) -> bool { |
| 222 | self.window_tokens > 0 && self.input_tokens >= self.compaction_trigger_tokens |
| 223 | } |
| 224 | |
| 225 | /// Whether another `additional_input_tokens` of input would fit within the |
| 226 | /// available budget (i.e. not exceed the reserved boundary). |
| 227 | #[must_use] |
| 228 | #[allow(dead_code)] |
| 229 | pub fn fits_additional(&self, additional_input_tokens: u64) -> bool { |
| 230 | additional_input_tokens <= self.available_input_tokens |
| 231 | } |
| 232 | } |
| 233 | |
| 234 | /// Clamp a desired output cap so it fits the window while preserving at least |
| 235 | /// [`MIN_INPUT_BUDGET_TOKENS`] of input room plus [`CONTEXT_HEADROOM_TOKENS`]. |
| 236 | /// |
| 237 | /// On a window too small to hold even that floor, returns whatever room is left |
| 238 | /// after the headroom (possibly zero) rather than underflowing. |
| 239 | fn clamp_output_cap(window_tokens: u64, configured_output_cap: u64) -> u64 { |
| 240 | // The most output we can reserve and still keep the input floor + headroom. |
| 241 | let reserved_floor = MIN_INPUT_BUDGET_TOKENS.saturating_add(CONTEXT_HEADROOM_TOKENS); |
| 242 | let max_output = window_tokens.saturating_sub(reserved_floor); |
| 243 | configured_output_cap.min(max_output) |
| 244 | } |
| 245 | |
| 246 | /// Window usage as a percentage in `0.0..=100.0`. Zero window -> `0.0`. |
| 247 | fn usage_percent(window_tokens: u64, input_tokens: u64) -> f64 { |
| 248 | if window_tokens == 0 { |
| 249 | return 0.0; |
| 250 | } |
| 251 | ((input_tokens as f64 / window_tokens as f64) * 100.0).clamp(0.0, 100.0) |
| 252 | } |
| 253 | |
| 254 | /// `percent`% of `window_tokens`, rounded to the nearest token. Saturates at |
| 255 | /// `u64::MAX` and treats out-of-range percentages by clamping to `0.0..=100.0`. |
| 256 | fn percent_of(window_tokens: u64, percent: f64) -> u64 { |
| 257 | let percent = percent.clamp(0.0, 100.0); |
| 258 | let value = (window_tokens as f64) * (percent / 100.0); |
| 259 | // `as u64` saturates on overflow and floors; add 0.5 to round to nearest. |
| 260 | (value + 0.5) as u64 |
| 261 | } |
| 262 | |
| 263 | #[cfg(test)] |
| 264 | mod tests { |
| 265 | use super::*; |
| 266 | |
| 267 | /// A representative spread of real-world windows: a tight self-hosted |
| 268 | /// deployment, common provider sizes, and a V4-class 1M window. |
| 269 | const WINDOWS: &[u64] = &[8_192, 32_768, 131_072, 262_144, 1_048_576]; |
| 270 | |
| 271 | // -- PressureLevel boundaries ------------------------------------------ |
| 272 | |
| 273 | #[test] |
| 274 | fn pressure_level_boundaries_are_inclusive_lower_bounds() { |
| 275 | assert_eq!(PressureLevel::from_usage_percent(0.0), PressureLevel::Low); |
| 276 | assert_eq!(PressureLevel::from_usage_percent(39.9), PressureLevel::Low); |
| 277 | // 40% is the moderate boundary. |
| 278 | assert_eq!( |
| 279 | PressureLevel::from_usage_percent(40.0), |
| 280 | PressureLevel::Medium |
| 281 | ); |
| 282 | assert_eq!( |
| 283 | PressureLevel::from_usage_percent(74.9), |
| 284 | PressureLevel::Medium |
| 285 | ); |
| 286 | // 75% is the high / compaction boundary. |
| 287 | assert_eq!(PressureLevel::from_usage_percent(75.0), PressureLevel::High); |
| 288 | assert_eq!(PressureLevel::from_usage_percent(89.9), PressureLevel::High); |
| 289 | // 90% is the critical boundary. |
| 290 | assert_eq!( |
| 291 | PressureLevel::from_usage_percent(90.0), |
| 292 | PressureLevel::Critical |
| 293 | ); |
| 294 | assert_eq!( |
| 295 | PressureLevel::from_usage_percent(100.0), |
| 296 | PressureLevel::Critical |
| 297 | ); |
| 298 | } |
| 299 | |
| 300 | #[test] |
| 301 | fn pressure_level_clamps_out_of_range_inputs() { |
| 302 | assert_eq!(PressureLevel::from_usage_percent(-10.0), PressureLevel::Low); |
| 303 | assert_eq!( |
| 304 | PressureLevel::from_usage_percent(150.0), |
| 305 | PressureLevel::Critical |
| 306 | ); |
| 307 | assert_eq!( |
| 308 | PressureLevel::from_usage_percent(f64::INFINITY), |
| 309 | PressureLevel::Critical |
| 310 | ); |
| 311 | } |
| 312 | |
| 313 | #[test] |
| 314 | fn pressure_level_ordering_and_helpers() { |
| 315 | assert!(PressureLevel::Low < PressureLevel::Medium); |
| 316 | assert!(PressureLevel::Medium < PressureLevel::High); |
| 317 | assert!(PressureLevel::High < PressureLevel::Critical); |
| 318 | |
| 319 | assert!(!PressureLevel::Low.suggests_compaction()); |
| 320 | assert!(!PressureLevel::Medium.suggests_compaction()); |
| 321 | assert!(PressureLevel::High.suggests_compaction()); |
| 322 | assert!(PressureLevel::Critical.suggests_compaction()); |
| 323 | |
| 324 | assert_eq!(PressureLevel::Low.label(), "low"); |
| 325 | assert_eq!(PressureLevel::Medium.label(), "moderate"); |
| 326 | assert_eq!(PressureLevel::High.label(), "high"); |
| 327 | assert_eq!(PressureLevel::Critical.label(), "critical"); |
| 328 | } |
| 329 | |
| 330 | // -- Compaction trigger ------------------------------------------------- |
| 331 | |
| 332 | #[test] |
| 333 | fn compaction_trigger_is_three_quarters_of_input_ceiling() { |
| 334 | for &window in WINDOWS { |
| 335 | let budget = ContextBudget::new(window, 0, 64_000); |
| 336 | let expected = percent_of( |
| 337 | budget.input_budget_ceiling, |
| 338 | DEFAULT_COMPACTION_TRIGGER_PERCENT, |
| 339 | ); |
| 340 | assert_eq!( |
| 341 | budget.compaction_trigger_tokens, expected, |
| 342 | "window {window}: trigger should be 75% of the input ceiling" |
| 343 | ); |
| 344 | assert!(budget.compaction_trigger_tokens <= budget.input_budget_ceiling); |
| 345 | } |
| 346 | } |
| 347 | |
| 348 | #[test] |
| 349 | fn should_compact_flips_at_the_trigger() { |
| 350 | let window = 1_048_576; |
| 351 | let cap = 262_144; |
| 352 | let trigger = ContextBudget::new(window, 0, cap).compaction_trigger_tokens; |
| 353 | assert!(trigger > 0); |
| 354 | |
| 355 | let below = ContextBudget::new(window, trigger - 1, cap); |
| 356 | assert!(!below.should_compact()); |
| 357 | |
| 358 | let at = ContextBudget::new(window, trigger, cap); |
| 359 | assert!(at.should_compact()); |
| 360 | |
| 361 | let above = ContextBudget::new(window, trigger + 1, cap); |
| 362 | assert!(above.should_compact()); |
| 363 | } |
| 364 | |
| 365 | #[test] |
| 366 | fn trigger_never_exceeds_input_ceiling() { |
| 367 | const WINDOWS: &[u64] = &[ |
| 368 | 0, 512, 8_192, 65_536, 131_072, 262_144, 500_000, 1_048_576, 2_000_000, |
| 369 | ]; |
| 370 | const CAPS: &[u64] = &[0, 1_024, 64_000, 131_072, 262_144, 1_000_000]; |
| 371 | const INPUTS: &[u64] = &[0, 1_024, 200_000, 2_000_000]; |
| 372 | const PERCENTS: &[f64] = &[10.0, 50.0, 75.0, 80.0, 95.0, 100.0]; |
| 373 | |
| 374 | for &window in WINDOWS { |
| 375 | for &cap in CAPS { |
| 376 | for &input in INPUTS { |
| 377 | let budget = ContextBudget::new(window, input, cap); |
| 378 | assert!(budget.compaction_trigger_tokens <= budget.input_budget_ceiling); |
| 379 | for &percent in PERCENTS { |
| 380 | assert!( |
| 381 | budget.compaction_trigger_for_percent(percent) |
| 382 | <= budget.input_budget_ceiling, |
| 383 | "window={window} cap={cap} input={input} percent={percent}" |
| 384 | ); |
| 385 | } |
| 386 | } |
| 387 | } |
| 388 | } |
| 389 | } |
| 390 | |
| 391 | #[test] |
| 392 | fn zero_window_never_suggests_compaction() { |
| 393 | let budget = ContextBudget::new(0, 0, 64_000); |
| 394 | assert_eq!(budget.compaction_trigger_tokens, 0); |
| 395 | assert_eq!(budget.input_budget_ceiling, 0); |
| 396 | assert!(!budget.should_compact()); |
| 397 | assert_eq!(budget.pressure, PressureLevel::Low); |
| 398 | assert_eq!(budget.available_input_tokens, 0); |
| 399 | assert_eq!(budget.usage_percent(), 0.0); |
| 400 | } |
| 401 | |
| 402 | // -- Output cap clamping & available budget ---------------------------- |
| 403 | |
| 404 | #[test] |
| 405 | fn output_cap_is_preserved_when_window_is_roomy() { |
| 406 | // 1M window, 64K configured cap: cap fits comfortably. |
| 407 | let budget = ContextBudget::new(1_048_576, 0, 64_000); |
| 408 | assert_eq!(budget.output_cap_tokens, 64_000); |
| 409 | // available = window - cap - headroom - input |
| 410 | let expected = 1_048_576 - 64_000 - CONTEXT_HEADROOM_TOKENS; |
| 411 | assert_eq!(budget.available_input_tokens, expected); |
| 412 | } |
| 413 | |
| 414 | #[test] |
| 415 | fn output_cap_is_clamped_to_protect_input_floor_on_small_window() { |
| 416 | // This is the engine's hard-won lesson: a generous output reservation |
| 417 | // on a small window must not underflow the input budget. An 8,192-token |
| 418 | // window with a 262,144-token desired cap must still leave the input |
| 419 | // floor available rather than collapsing to zero or wrapping. |
| 420 | let window = 8_192u64; |
| 421 | let budget = ContextBudget::new(window, 0, 262_144); |
| 422 | |
| 423 | let reserved_floor = MIN_INPUT_BUDGET_TOKENS + CONTEXT_HEADROOM_TOKENS; |
| 424 | let expected_cap = window - reserved_floor; |
| 425 | assert_eq!(budget.output_cap_tokens, expected_cap); |
| 426 | // With zero input committed, the whole remaining budget is available |
| 427 | // and is at least the protected floor. |
| 428 | assert!(budget.available_input_tokens >= MIN_INPUT_BUDGET_TOKENS); |
| 429 | assert_eq!( |
| 430 | budget.available_input_tokens, |
| 431 | window - budget.output_cap_tokens - CONTEXT_HEADROOM_TOKENS |
| 432 | ); |
| 433 | } |
| 434 | |
| 435 | #[test] |
| 436 | fn tiny_window_below_floor_saturates_without_panic() { |
| 437 | // Window smaller than the protected floor: cap collapses to 0 and the |
| 438 | // available budget saturates at 0 instead of underflowing. |
| 439 | let window = 512u64; // < MIN_INPUT_BUDGET_TOKENS + headroom |
| 440 | let budget = ContextBudget::new(window, 100, 262_144); |
| 441 | assert_eq!(budget.output_cap_tokens, 0); |
| 442 | assert_eq!(budget.available_input_tokens, 0); |
| 443 | // Usage still computes a sane percentage. |
| 444 | assert!((budget.usage_percent() - (100.0 / 512.0 * 100.0)).abs() < 1e-9); |
| 445 | } |
| 446 | |
| 447 | #[test] |
| 448 | fn available_budget_saturates_when_input_exceeds_ceiling() { |
| 449 | let window = 131_072u64; |
| 450 | let cap = 32_000u64; |
| 451 | // Commit far more input than the window holds. |
| 452 | let budget = ContextBudget::new(window, window * 2, cap); |
| 453 | assert_eq!(budget.available_input_tokens, 0); |
| 454 | assert!(!budget.fits_additional(1)); |
| 455 | assert!(budget.fits_additional(0)); |
| 456 | // Usage is clamped to 100%. |
| 457 | assert_eq!(budget.usage_percent(), 100.0); |
| 458 | assert_eq!(budget.pressure, PressureLevel::Critical); |
| 459 | } |
| 460 | |
| 461 | #[test] |
| 462 | fn fits_additional_respects_the_reserved_boundary() { |
| 463 | let window = 262_144u64; |
| 464 | let cap = 64_000u64; |
| 465 | let budget = ContextBudget::new(window, 100_000, cap); |
| 466 | let room = budget.available_input_tokens; |
| 467 | assert!(budget.fits_additional(room)); |
| 468 | assert!(!budget.fits_additional(room + 1)); |
| 469 | } |
| 470 | |
| 471 | // -- Usage percent & pressure across window sizes ---------------------- |
| 472 | |
| 473 | #[test] |
| 474 | fn usage_percent_is_proportional_across_window_sizes() { |
| 475 | for &window in WINDOWS { |
| 476 | // Half-full should read ~50% and classify as Medium regardless of |
| 477 | // absolute window size. |
| 478 | let half = window / 2; |
| 479 | let budget = ContextBudget::new(window, half, 64_000); |
| 480 | let pct = budget.usage_percent(); |
| 481 | assert!( |
| 482 | (pct - 50.0).abs() < 0.5, |
| 483 | "window {window}: half-full should be ~50%, got {pct}" |
| 484 | ); |
| 485 | assert_eq!(budget.pressure, PressureLevel::Medium); |
| 486 | } |
| 487 | } |
| 488 | |
| 489 | #[test] |
| 490 | fn pressure_tracks_input_growth_on_a_1m_window() { |
| 491 | let window = 1_048_576u64; |
| 492 | let cap = 262_144u64; |
| 493 | |
| 494 | let low = ContextBudget::new(window, percent_of(window, 10.0), cap); |
| 495 | assert_eq!(low.pressure, PressureLevel::Low); |
| 496 | |
| 497 | let medium = ContextBudget::new(window, percent_of(window, 50.0), cap); |
| 498 | assert_eq!(medium.pressure, PressureLevel::Medium); |
| 499 | |
| 500 | let high = ContextBudget::new(window, percent_of(window, 80.0), cap); |
| 501 | assert_eq!(high.pressure, PressureLevel::High); |
| 502 | assert!(high.should_compact()); |
| 503 | |
| 504 | let critical = ContextBudget::new(window, percent_of(window, 95.0), cap); |
| 505 | assert_eq!(critical.pressure, PressureLevel::Critical); |
| 506 | } |
| 507 | |
| 508 | #[test] |
| 509 | fn snapshot_fields_are_internally_consistent() { |
| 510 | for &window in WINDOWS { |
| 511 | for &input in &[0u64, window / 4, window / 2, window] { |
| 512 | let budget = ContextBudget::new(window, input, 64_000); |
| 513 | // Field mirrors the constructor arguments. |
| 514 | assert_eq!(budget.window_tokens, window); |
| 515 | assert_eq!(budget.input_tokens, input); |
| 516 | // available + input never claims more than the window minus |
| 517 | // reserved output and headroom. |
| 518 | let ceiling = window |
| 519 | .saturating_sub(budget.output_cap_tokens) |
| 520 | .saturating_sub(CONTEXT_HEADROOM_TOKENS); |
| 521 | assert!(budget.available_input_tokens <= ceiling); |
| 522 | // Pressure agrees with the standalone usage percentage. |
| 523 | assert_eq!( |
| 524 | budget.pressure, |
| 525 | PressureLevel::from_usage_percent(budget.usage_percent()) |
| 526 | ); |
| 527 | } |
| 528 | } |
| 529 | } |
| 530 | |
| 531 | #[test] |
| 532 | fn percent_of_rounds_to_nearest_token() { |
| 533 | assert_eq!(percent_of(0, 75.0), 0); |
| 534 | assert_eq!(percent_of(100, 0.0), 0); |
| 535 | assert_eq!(percent_of(100, 100.0), 100); |
| 536 | assert_eq!(percent_of(100, 75.0), 75); |
| 537 | // 3 * 0.75 = 2.25 -> rounds to 2. |
| 538 | assert_eq!(percent_of(3, 75.0), 2); |
| 539 | // 2 * 0.75 = 1.5 -> rounds to 2. |
| 540 | assert_eq!(percent_of(2, 75.0), 2); |
| 541 | } |
| 542 | |
| 543 | #[test] |
| 544 | fn budget_is_copy_and_comparable() { |
| 545 | let a = ContextBudget::new(131_072, 1_000, 32_000); |
| 546 | let b = a; // Copy, not move. |
| 547 | assert_eq!(a, b); |
| 548 | assert_eq!(a.window_tokens, b.window_tokens); |
| 549 | } |
| 550 | } |
| 551 |