1use crate::core::config::{ContextManagementConfig, ContextManagementMode};
16
17#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
19pub enum WorkPhase {
20 #[default]
21 Unknown,
22 Execute,
24 WrapUp,
26 Plan,
28 NewTask,
29}
30
31impl WorkPhase {
32 pub fn parse(value: &str) -> Option<Self> {
33 match value.trim().to_ascii_lowercase().as_str() {
34 "unknown" => Some(Self::Unknown),
35 "execute" => Some(Self::Execute),
36 "wrap_up" | "wrapup" | "wrap-up" => Some(Self::WrapUp),
37 "plan" => Some(Self::Plan),
38 "new_task" | "newtask" | "new-task" => Some(Self::NewTask),
39 _ => None,
40 }
41 }
42
43 pub fn as_str(self) -> &'static str {
44 match self {
45 Self::Unknown => "unknown",
46 Self::Execute => "execute",
47 Self::WrapUp => "wrap_up",
48 Self::Plan => "plan",
49 Self::NewTask => "new_task",
50 }
51 }
52}
53
54#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
57pub struct ContextState {
58 phase: WorkPhase,
59 plan_pending_execution: bool,
60 boundary_pending: bool,
61 rollover_due: bool,
62 rollover_required: bool,
63 finish_rounds_used: u32,
64 unproductive_rollover: Option<(u64, u8)>,
67}
68
69impl ContextState {
70 pub fn phase(&self) -> WorkPhase {
71 self.phase
72 }
73
74 pub fn finish_rounds_used(&self) -> u32 {
75 self.finish_rounds_used
76 }
77
78 pub fn report_phase(&mut self, phase: WorkPhase) {
81 match phase {
82 WorkPhase::Plan => self.plan_pending_execution = true,
83 WorkPhase::Execute if self.plan_pending_execution => {
84 self.boundary_pending = true;
85 self.plan_pending_execution = false;
86 }
87 WorkPhase::NewTask => self.boundary_pending = true,
88 _ => {}
89 }
90 self.phase = phase;
91 }
92
93 pub fn defer_unproductive_rollover(&mut self, used_tokens: u64) {
97 self.unproductive_rollover = Some((used_tokens, 0));
98 }
99
100 pub fn reset(&mut self) {
102 *self = Self::default();
103 }
104}
105
106#[derive(Debug, Clone, Copy, PartialEq, Eq)]
108pub struct ContextBudget {
109 pub context_window_tokens: u64,
110 pub used_tokens: u64,
112 pub hard_remaining_tokens: u64,
115 pub required_next_round_tokens: u64,
118}
119
120#[derive(Debug, Clone, Copy, PartialEq, Eq)]
121pub struct ContextThresholds {
122 pub pressure_tokens: u64,
123 pub rollover_tokens: u64,
124}
125
126impl ContextThresholds {
127 pub fn resolve(
132 config: &ContextManagementConfig,
133 context_window_tokens: u64,
134 ) -> Result<Self, &'static str> {
135 config.validate_for_window(context_window_tokens)?;
136 let window = context_window_tokens;
137 let (pressure, rollover) = if window <= 200_000 {
138 ((window * 7 / 10).max(1), (window * 9 / 10).max(2))
139 } else if window < 1_000_000 {
140 let extra = window - 200_000;
141 (
142 140_000 + extra * 110_000 / 800_000,
143 180_000 + extra * 220_000 / 800_000,
144 )
145 } else {
146 (250_000, 400_000)
147 };
148 let (pressure_tokens, rollover_tokens) =
149 match (config.pressure_tokens, config.rollover_tokens) {
150 (Some(p), Some(r)) => (p, r),
151 (Some(p), None) => (p, rollover.max(p + 1)),
152 (None, Some(r)) => (pressure.min(r - 1), r),
153 (None, None) => (pressure, rollover),
154 };
155 Ok(Self {
156 pressure_tokens,
157 rollover_tokens,
158 })
159 }
160}
161
162#[derive(Debug, Clone, Copy, PartialEq, Eq)]
163pub enum ContextBand {
164 Normal,
165 Pressure,
166 Rollover,
167 HardLimit,
168}
169
170#[derive(Debug, Clone, Copy, PartialEq, Eq)]
171pub enum ContextAction {
172 Disabled,
173 Continue,
174 Advisory,
176 FinishBounded {
178 rounds_remaining: u32,
179 },
180 Rollover,
182 HardStop,
185}
186
187#[derive(Debug, Clone, Copy, PartialEq, Eq)]
188pub enum ContextReason {
189 ModeOff,
190 BelowPressure,
191 PrepareForRollover,
192 TaskBoundary,
193 RolloverThreshold,
194 BoundedFinish,
195 FinishAllowanceExhausted,
196 RolloverPending,
197 UnproductiveRollover,
198 HardCapacity,
199 InvalidConfig,
200}
201
202#[derive(Debug, Clone, Copy, PartialEq, Eq)]
203pub struct ContextAssessment {
204 pub action: ContextAction,
205 pub band: ContextBand,
206 pub reason: ContextReason,
207 pub thresholds: Option<ContextThresholds>,
209 pub next_state: ContextState,
212}
213
214pub fn assess_context(
218 config: &ContextManagementConfig,
219 state: &ContextState,
220 budget: ContextBudget,
221) -> ContextAssessment {
222 let mut assessment = ContextAssessment {
223 action: ContextAction::HardStop,
224 band: ContextBand::HardLimit,
225 reason: ContextReason::HardCapacity,
226 thresholds: None,
227 next_state: *state,
228 };
229 let remaining = budget.hard_remaining_tokens.min(
230 budget
231 .context_window_tokens
232 .saturating_sub(budget.used_tokens),
233 );
234 let required = budget.required_next_round_tokens.max(config.reserve_tokens);
235 if remaining == 0 || remaining < required {
236 return assessment;
237 }
238
239 assessment.next_state.boundary_pending = false;
242 if config.mode == ContextManagementMode::Off {
243 assessment.action = ContextAction::Disabled;
244 assessment.band = ContextBand::Normal;
245 assessment.reason = ContextReason::ModeOff;
246 return assessment;
247 }
248 let Ok(thresholds) = ContextThresholds::resolve(config, budget.context_window_tokens) else {
249 assessment.reason = ContextReason::InvalidConfig;
250 return assessment;
251 };
252 assessment.thresholds = Some(thresholds);
253 assessment.band = if budget.used_tokens >= thresholds.rollover_tokens {
254 ContextBand::Rollover
255 } else if budget.used_tokens >= thresholds.pressure_tokens {
256 ContextBand::Pressure
257 } else {
258 ContextBand::Normal
259 };
260
261 if let Some((used, skipped)) = state.unproductive_rollover {
266 if budget.used_tokens.abs_diff(used) < 8_192 && skipped < 4 {
267 assessment.next_state.unproductive_rollover = Some((used, skipped + 1));
268 assessment.action = ContextAction::Advisory;
269 assessment.reason = ContextReason::UnproductiveRollover;
270 return assessment;
271 }
272 assessment.next_state.unproductive_rollover = None;
273 }
274
275 if state.rollover_required {
276 assessment.action = ContextAction::Rollover;
277 assessment.reason = ContextReason::RolloverPending;
278 return assessment;
279 }
280 if budget.used_tokens >= thresholds.pressure_tokens
281 && (state.boundary_pending || state.phase == WorkPhase::NewTask)
282 {
283 assessment.next_state.rollover_required = true;
284 assessment.action = ContextAction::Rollover;
285 assessment.reason = ContextReason::TaskBoundary;
286 return assessment;
287 }
288 if assessment.band == ContextBand::Rollover || state.rollover_due {
289 assessment.next_state.rollover_due = true;
292 if matches!(state.phase, WorkPhase::Plan | WorkPhase::WrapUp)
293 && state.finish_rounds_used < config.finish_rounds
294 {
295 assessment.next_state.finish_rounds_used += 1;
296 assessment.action = ContextAction::FinishBounded {
297 rounds_remaining: config.finish_rounds - assessment.next_state.finish_rounds_used,
298 };
299 assessment.reason = ContextReason::BoundedFinish;
300 } else {
301 assessment.next_state.rollover_required = true;
302 assessment.action = ContextAction::Rollover;
303 assessment.reason = if state.finish_rounds_used >= config.finish_rounds {
304 ContextReason::FinishAllowanceExhausted
305 } else {
306 ContextReason::RolloverThreshold
307 };
308 }
309 } else if assessment.band == ContextBand::Pressure {
310 assessment.action = ContextAction::Advisory;
311 assessment.reason = ContextReason::PrepareForRollover;
312 } else {
313 assessment.action = ContextAction::Continue;
314 assessment.reason = ContextReason::BelowPressure;
315 }
316 assessment
317}
318
319#[cfg(test)]
320mod tests {
321 use super::*;
322
323 fn auto_config() -> ContextManagementConfig {
324 ContextManagementConfig {
325 mode: ContextManagementMode::Auto,
326 ..Default::default()
327 }
328 }
329
330 fn budget(window: u64, used: u64) -> ContextBudget {
331 ContextBudget {
332 context_window_tokens: window,
333 used_tokens: used,
334 hard_remaining_tokens: window.saturating_sub(used),
335 required_next_round_tokens: 8_000,
336 }
337 }
338
339 #[test]
340 fn unproductive_rollover_defers_only_soft_admission_and_retries_after_progress() {
341 let config = auto_config();
342 let mut state = ContextState::default();
343 state.report_phase(WorkPhase::NewTask);
344 state = assess_context(&config, &state, budget(1_000_000, 339_058)).next_state;
345 state.defer_unproductive_rollover(339_058);
346 for _ in 0..4 {
347 state.report_phase(WorkPhase::NewTask); let next = assess_context(&config, &state, budget(1_000_000, 339_100));
349 assert_eq!(next.action, ContextAction::Advisory);
350 assert_eq!(next.reason, ContextReason::UnproductiveRollover);
351 state = next.next_state;
352 }
353 assert_eq!(
354 assess_context(&config, &state, budget(1_000_000, 339_100)).action,
355 ContextAction::Rollover
356 );
357 for changed in [329_000, 350_000] {
358 state.defer_unproductive_rollover(339_058);
359 assert_eq!(
360 assess_context(&config, &state, budget(1_000_000, changed)).action,
361 ContextAction::Rollover
362 );
363 }
364 state.reset();
365 assert_eq!(state.unproductive_rollover, None);
366 assert_eq!(state.phase(), WorkPhase::Unknown);
367 }
368
369 #[test]
370 fn unproductive_cooldown_cannot_bypass_hard_budget_or_renew_finish_allowance() {
371 let config = auto_config();
372 let mut state = ContextState::default();
373 state.report_phase(WorkPhase::Plan);
374 for _ in 0..config.finish_rounds {
375 state = assess_context(&config, &state, budget(1_000_000, 400_000)).next_state;
376 }
377 state = assess_context(&config, &state, budget(1_000_000, 400_000)).next_state;
378 state.defer_unproductive_rollover(400_000);
379 state.report_phase(WorkPhase::WrapUp);
380 let off = ContextManagementConfig {
381 mode: ContextManagementMode::Off,
382 ..config
383 };
384 state = assess_context(&off, &state, budget(1_000_000, 400_000)).next_state;
385 assert_eq!(state.unproductive_rollover, Some((400_000, 0)));
386 assert_eq!(state.finish_rounds_used(), config.finish_rounds);
387 let denied = assess_context(
388 &config,
389 &state,
390 ContextBudget {
391 hard_remaining_tokens: 0,
392 ..budget(1_000_000, 400_000)
393 },
394 );
395 assert_eq!(denied.action, ContextAction::HardStop);
396 assert_eq!(denied.next_state, state);
397 assert_eq!(
398 assess_context(&config, &state, budget(410_000, 400_000)).action,
399 ContextAction::HardStop
400 );
401 let allowed = assess_context(&config, &state, budget(1_000_000, 400_000));
402 assert_eq!(allowed.reason, ContextReason::UnproductiveRollover);
403 assert_eq!(
404 allowed.next_state.finish_rounds_used(),
405 config.finish_rounds
406 );
407 }
408
409 #[test]
410 fn context_threshold_anchors_and_pressure_are_advisory() {
411 let config = auto_config();
412 for (window, pressure, rollover) in [
413 (200_000, 140_000, 180_000),
414 (1_000_000, 250_000, 400_000),
415 (2_000_000, 250_000, 400_000),
416 ] {
417 assert_eq!(
418 ContextThresholds::resolve(&config, window).unwrap(),
419 ContextThresholds {
420 pressure_tokens: pressure,
421 rollover_tokens: rollover
422 }
423 );
424 let mut state = ContextState::default();
425 state.report_phase(WorkPhase::Execute);
426 assert_eq!(
427 assess_context(&config, &state, budget(window, pressure - 1)).action,
428 ContextAction::Continue
429 );
430 let at_pressure = assess_context(&config, &state, budget(window, pressure));
431 assert_eq!(at_pressure.action, ContextAction::Advisory);
432 assert_eq!(at_pressure.band, ContextBand::Pressure);
433 assert_eq!(
434 assess_context(&config, &state, budget(window, rollover - 1)).action,
435 ContextAction::Advisory
436 );
437 assert_eq!(
438 assess_context(&config, &state, budget(window, rollover)).action,
439 ContextAction::Rollover
440 );
441 }
442 }
443
444 #[test]
445 fn context_350k_plan_finishes_before_execute_requests_rollover() {
446 let config = auto_config();
447 let current_budget = budget(1_000_000, 350_000);
448 for phase in [WorkPhase::Execute, WorkPhase::WrapUp, WorkPhase::Plan] {
449 let mut state = ContextState::default();
450 state.report_phase(phase);
451 assert_eq!(
452 assess_context(&config, &state, current_budget).action,
453 ContextAction::Advisory
454 );
455 }
456 let mut state = ContextState::default();
457 state.report_phase(WorkPhase::Plan);
458 state = assess_context(&config, &state, current_budget).next_state;
459 state.report_phase(WorkPhase::Unknown);
460 state.report_phase(WorkPhase::WrapUp);
461 state.report_phase(WorkPhase::Execute);
462 state.report_phase(WorkPhase::Plan); let assessment = assess_context(&config, &state, current_budget);
464 assert_eq!(assessment.action, ContextAction::Rollover);
465 assert_eq!(assessment.reason, ContextReason::TaskBoundary);
466 state = assessment.next_state;
467 state.report_phase(WorkPhase::WrapUp);
468 assert_eq!(
469 assess_context(&config, &state, current_budget).action,
470 ContextAction::Rollover
471 );
472 state.reset();
473 assert_eq!(
474 assess_context(&config, &state, budget(1_000_000, 0)).action,
475 ContextAction::Continue
476 );
477 }
478
479 #[test]
480 fn context_new_task_rolls_over_at_pressure_not_before() {
481 let config = auto_config();
482 let mut state = ContextState::default();
483 state.report_phase(WorkPhase::NewTask);
484 assert_eq!(
485 assess_context(&config, &state, budget(1_000_000, 249_999)).action,
486 ContextAction::Continue
487 );
488 assert_eq!(
489 assess_context(&config, &state, budget(1_000_000, 250_000)).action,
490 ContextAction::Rollover
491 );
492 }
493
494 #[test]
495 fn context_low_pressure_boundary_is_consumed_at_admission() {
496 let config = auto_config();
497 let mut state = ContextState::default();
498 state.report_phase(WorkPhase::Plan);
499 state.report_phase(WorkPhase::Execute);
500 state = assess_context(&config, &state, budget(1_000_000, 100_000)).next_state;
501 assert_eq!(
502 assess_context(&config, &state, budget(1_000_000, 350_000)).action,
503 ContextAction::Advisory
504 );
505 }
506
507 #[test]
508 fn context_400k_finish_is_bounded_across_reports_and_estimate_decreases() {
509 let config = auto_config();
510 for phase in [WorkPhase::Plan, WorkPhase::WrapUp] {
511 let mut state = ContextState::default();
512 state.report_phase(phase);
513 let first = assess_context(&config, &state, budget(1_000_000, 400_000));
514 assert_eq!(
515 first.action,
516 ContextAction::FinishBounded {
517 rounds_remaining: 1
518 }
519 );
520 assert_eq!(state.finish_rounds_used(), 0); assert_eq!(
522 first,
523 assess_context(&config, &state, budget(1_000_000, 400_000))
524 );
525 state = first.next_state;
526 state.report_phase(WorkPhase::Unknown);
527 state.report_phase(phase);
528 let second = assess_context(&config, &state, budget(1_000_000, 350_000));
529 assert_eq!(
530 second.action,
531 ContextAction::FinishBounded {
532 rounds_remaining: 0
533 }
534 );
535 state = second.next_state;
536 assert_eq!(state.finish_rounds_used(), 2);
537 state.report_phase(phase);
538 let exhausted = assess_context(&config, &state, budget(1_000_000, 100_000));
539 assert_eq!(exhausted.action, ContextAction::Rollover);
540 assert_eq!(exhausted.reason, ContextReason::FinishAllowanceExhausted);
541 state = exhausted.next_state;
542 state.report_phase(WorkPhase::Plan);
543 assert_eq!(
544 assess_context(&config, &state, budget(1_000_000, 100_000)).action,
545 ContextAction::Rollover
546 );
547 state.reset();
548 assert_eq!(state.finish_rounds_used(), 0);
549 assert_eq!(state.phase(), WorkPhase::Unknown);
550 }
551 }
552
553 #[test]
554 fn context_finish_override_is_bounded_including_zero() {
555 for limit in [0, 1, ContextManagementConfig::MAX_FINISH_ROUNDS] {
556 let config = ContextManagementConfig {
557 finish_rounds: limit,
558 ..auto_config()
559 };
560 let mut state = ContextState::default();
561 state.report_phase(WorkPhase::WrapUp);
562 for used in 0..limit {
563 let assessment = assess_context(&config, &state, budget(1_000_000, 400_000));
564 assert_eq!(
565 assessment.action,
566 ContextAction::FinishBounded {
567 rounds_remaining: limit - used - 1
568 }
569 );
570 state = assessment.next_state;
571 }
572 assert_eq!(
573 assess_context(&config, &state, budget(1_000_000, 400_000)).action,
574 ContextAction::Rollover
575 );
576 }
577 }
578
579 #[test]
580 fn context_off_is_default_and_mode_toggles_do_not_replenish_finish() {
581 let off = ContextManagementConfig::default();
582 let mut state = ContextState::default();
583 state.report_phase(WorkPhase::WrapUp);
584 assert_eq!(
585 assess_context(&off, &state, budget(1_000_000, 400_000)).action,
586 ContextAction::Disabled
587 );
588 state = assess_context(&auto_config(), &state, budget(1_000_000, 400_000)).next_state;
589 state = assess_context(&off, &state, budget(1_000_000, 410_000)).next_state;
590 assert_eq!(state.finish_rounds_used(), 1);
591 let second = assess_context(&auto_config(), &state, budget(1_000_000, 420_000));
592 assert_eq!(
593 second.action,
594 ContextAction::FinishBounded {
595 rounds_remaining: 0
596 }
597 );
598 assert_eq!(
599 assess_context(
600 &auto_config(),
601 &second.next_state,
602 budget(1_000_000, 420_000)
603 )
604 .action,
605 ContextAction::Rollover
606 );
607 }
608
609 #[test]
610 fn context_hard_capacity_wins_for_every_phase_and_mode() {
611 for mode in [ContextManagementMode::Off, ContextManagementMode::Auto] {
612 let config = ContextManagementConfig {
613 mode,
614 ..Default::default()
615 };
616 for phase in [
617 WorkPhase::Unknown,
618 WorkPhase::Execute,
619 WorkPhase::WrapUp,
620 WorkPhase::Plan,
621 WorkPhase::NewTask,
622 ] {
623 let mut state = ContextState::default();
624 state.report_phase(phase);
625 for hard_budget in [
626 ContextBudget {
627 hard_remaining_tokens: 15_999,
628 ..budget(1_000_000, 350_000)
629 },
630 ContextBudget {
631 required_next_round_tokens: 650_001,
632 ..budget(1_000_000, 350_000)
633 },
634 ContextBudget {
635 hard_remaining_tokens: u64::MAX,
636 ..budget(200_000, 190_000)
637 },
638 budget(200_000, 200_001),
639 budget(0, 0),
640 ] {
641 let assessment = assess_context(&config, &state, hard_budget);
642 assert_eq!(
643 assessment.action,
644 ContextAction::HardStop,
645 "{mode:?} {phase:?} {hard_budget:?}"
646 );
647 assert_eq!(assessment.band, ContextBand::HardLimit);
648 assert_eq!(assessment.reason, ContextReason::HardCapacity);
649 assert_eq!(assessment.next_state, state);
650 }
651 }
652 }
653 }
654
655 #[test]
656 fn context_exact_reserve_is_allowed_and_next_round_can_require_more() {
657 let config = auto_config();
658 let state = ContextState::default();
659 let exact = ContextBudget {
660 hard_remaining_tokens: 16_000,
661 ..budget(200_000, 140_000)
662 };
663 assert_eq!(
664 assess_context(&config, &state, exact).action,
665 ContextAction::Advisory
666 );
667 let larger = ContextBudget {
668 required_next_round_tokens: 16_001,
669 ..exact
670 };
671 assert_eq!(
672 assess_context(&config, &state, larger).action,
673 ContextAction::HardStop
674 );
675 }
676
677 #[test]
678 fn context_threshold_overrides_preserve_explicit_values() {
679 for (pressure, rollover, expected) in [
680 (Some(350_000), Some(500_000), (350_000, 500_000)),
681 (Some(450_000), None, (450_000, 450_001)),
682 (None, Some(200_000), (199_999, 200_000)),
683 ] {
684 let config = ContextManagementConfig {
685 pressure_tokens: pressure,
686 rollover_tokens: rollover,
687 ..auto_config()
688 };
689 let thresholds = ContextThresholds::resolve(&config, 1_000_000).unwrap();
690 assert_eq!(
691 (thresholds.pressure_tokens, thresholds.rollover_tokens),
692 expected
693 );
694 }
695 assert!(ContextThresholds::resolve(&auto_config(), u64::MAX).is_ok());
696 assert_eq!(
697 ContextThresholds::resolve(&auto_config(), 100_000)
698 .unwrap()
699 .pressure_tokens,
700 70_000
701 );
702 assert_eq!(
703 ContextThresholds::resolve(&auto_config(), 600_000)
704 .unwrap()
705 .rollover_tokens,
706 290_000
707 );
708 }
709
710 #[test]
711 fn context_invalid_programmatic_config_fails_closed() {
712 for config in [
713 ContextManagementConfig {
714 pressure_tokens: Some(0),
715 ..auto_config()
716 },
717 ContextManagementConfig {
718 pressure_tokens: Some(u64::MAX),
719 ..auto_config()
720 },
721 ContextManagementConfig {
722 rollover_tokens: Some(1_000_001),
723 ..auto_config()
724 },
725 ContextManagementConfig {
726 finish_rounds: 9,
727 ..auto_config()
728 },
729 ContextManagementConfig {
730 reserve_tokens: 0,
731 ..auto_config()
732 },
733 ] {
734 let assessment = assess_context(
735 &config,
736 &ContextState::default(),
737 budget(1_000_000, 350_000),
738 );
739 assert_eq!(assessment.action, ContextAction::HardStop);
740 assert_eq!(assessment.reason, ContextReason::InvalidConfig);
741 }
742 }
743
744 #[test]
745 fn context_work_phase_roundtrips_and_rejects_unknown_reports() {
746 for phase in [
747 WorkPhase::Unknown,
748 WorkPhase::Execute,
749 WorkPhase::WrapUp,
750 WorkPhase::Plan,
751 WorkPhase::NewTask,
752 ] {
753 assert_eq!(WorkPhase::parse(phase.as_str()), Some(phase));
754 }
755 assert_eq!(WorkPhase::parse(" PLAN "), Some(WorkPhase::Plan));
756 assert_eq!(WorkPhase::parse("wrap-up"), Some(WorkPhase::WrapUp));
757 assert_eq!(WorkPhase::parse("new-task"), Some(WorkPhase::NewTask));
758 assert_eq!(WorkPhase::parse("finish_forever"), None);
759 }
760}