use serde::Serialize;
pub const HISTORY_FOLD_THRESHOLD: f64 = 0.75;
pub const HISTORY_FOLD_TAIL_FRACTION: f64 = 0.2;
pub const HISTORY_FOLD_AGGRESSIVE_THRESHOLD: f64 = 0.78;
pub const HISTORY_FOLD_AGGRESSIVE_TAIL_FRACTION: f64 = 0.1;
#[cfg(test)]
pub const HISTORY_FOLD_MIN_SAVINGS_FRACTION: f64 = 0.3;
pub const FORCE_SUMMARY_THRESHOLD: f64 = 0.8;
pub const TURN_START_FOLD_THRESHOLD: f64 = 0.9;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
#[serde(rename_all = "kebab-case")]
pub enum PostUsageDecisionKind {
None,
Fold,
ExitWithSummary,
}
#[derive(Debug, Clone, Copy)]
pub struct PostUsageDecision {
pub kind: PostUsageDecisionKind,
#[allow(dead_code)]
pub prompt_tokens: u64,
#[allow(dead_code)]
pub ctx_max: u64,
pub ratio: f64,
pub tail_budget: Option<u64>,
pub aggressive: bool,
}
#[derive(Debug, Clone, Copy)]
pub struct TurnStartEstimate {
pub estimate_tokens: u64,
pub ctx_max: u64,
pub ratio: f64,
}
pub fn decide_after_usage(
prompt_tokens: Option<u64>,
ctx_max: u64,
already_folded_this_turn: bool,
) -> PostUsageDecision {
decide_after_usage_with_threshold(prompt_tokens, ctx_max, already_folded_this_turn, None)
}
pub fn checkpoint_thresholds_for(ctx_max: u64) -> Vec<f64> {
if ctx_max < 25_000 {
return Vec::new();
}
if ctx_max <= 200_000 {
return vec![0.2, 0.4, 0.6, 0.8];
}
if ctx_max <= 500_000 {
return (1..=9).map(|i| i as f64 * 0.1).collect();
}
(1..=18).map(|i| i as f64 * 0.05).collect()
}
#[derive(Debug, Clone)]
pub struct CheckpointSchedule {
thresholds: Vec<f64>,
crossed: Vec<bool>,
}
impl CheckpointSchedule {
pub fn new(ctx_max: u64) -> Self {
let thresholds = checkpoint_thresholds_for(ctx_max);
let crossed = vec![false; thresholds.len()];
Self {
thresholds,
crossed,
}
}
pub fn is_enabled(&self) -> bool {
!self.thresholds.is_empty()
}
pub fn note_usage(&mut self, ratio: f64) -> bool {
let mut newly = false;
for (i, &t) in self.thresholds.iter().enumerate() {
if !self.crossed[i] && ratio >= t {
self.crossed[i] = true;
newly = true;
}
}
newly
}
pub fn reset(&mut self) {
for c in &mut self.crossed {
*c = false;
}
}
}
static FOLD_THRESHOLD_OVERRIDE: std::sync::OnceLock<Option<f64>> = std::sync::OnceLock::new();
pub fn init_fold_threshold(override_fraction: Option<f64>) {
let _ = FOLD_THRESHOLD_OVERRIDE.set(override_fraction);
}
pub const DEFAULT_CONTEXT_TARGET: u64 = 100_000;
const MIN_CONTEXT_TARGET: u64 = 16_000;
static CONTEXT_TARGET: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
pub fn resolve_context_target(configured: Option<u64>) -> u64 {
configured
.map(|t| t.max(MIN_CONTEXT_TARGET))
.unwrap_or(DEFAULT_CONTEXT_TARGET)
}
pub fn init_context_target(target: Option<u64>) {
let _ = CONTEXT_TARGET.set(resolve_context_target(target));
}
pub fn context_target() -> u64 {
*CONTEXT_TARGET.get().unwrap_or(&DEFAULT_CONTEXT_TARGET)
}
pub fn effective_ctx_max(model_window: u64) -> u64 {
model_window.min(context_target())
}
static CONTEXT_WINDOW_OVERRIDE: std::sync::OnceLock<Option<u64>> = std::sync::OnceLock::new();
pub fn init_context_window_override(window: Option<u64>) {
let _ = CONTEXT_WINDOW_OVERRIDE.set(window);
}
pub fn context_window_override() -> Option<u64> {
CONTEXT_WINDOW_OVERRIDE.get().copied().flatten()
}
static INCREMENTAL_CHECKPOINT: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
pub fn init_incremental_checkpoint(enabled: Option<bool>) {
let _ = INCREMENTAL_CHECKPOINT.set(enabled.unwrap_or(true));
}
pub fn incremental_checkpoint_enabled() -> bool {
*INCREMENTAL_CHECKPOINT.get().unwrap_or(&true)
}
static MEMORIES_DIRTY: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
pub fn mark_memories_dirty() {
MEMORIES_DIRTY.store(true, std::sync::atomic::Ordering::Release);
}
pub fn take_memories_dirty() -> bool {
MEMORIES_DIRTY.swap(false, std::sync::atomic::Ordering::AcqRel)
}
pub fn effective_fold_threshold(override_fraction: Option<f64>) -> f64 {
let candidate = override_fraction.or_else(|| FOLD_THRESHOLD_OVERRIDE.get().copied().flatten());
match candidate {
Some(f) if f.is_finite() && (0.3..=HISTORY_FOLD_THRESHOLD).contains(&f) => f,
_ => HISTORY_FOLD_THRESHOLD,
}
}
pub fn decide_after_usage_with_threshold(
prompt_tokens: Option<u64>,
ctx_max: u64,
already_folded_this_turn: bool,
fold_threshold_override: Option<f64>,
) -> PostUsageDecision {
let Some(prompt_tokens) = prompt_tokens else {
return PostUsageDecision {
kind: PostUsageDecisionKind::None,
prompt_tokens: 0,
ctx_max,
ratio: 0.0,
tail_budget: None,
aggressive: false,
};
};
if ctx_max == 0 {
return PostUsageDecision {
kind: PostUsageDecisionKind::None,
prompt_tokens,
ctx_max,
ratio: 0.0,
tail_budget: None,
aggressive: false,
};
}
let ratio = prompt_tokens as f64 / ctx_max as f64;
if ratio > FORCE_SUMMARY_THRESHOLD {
return PostUsageDecision {
kind: PostUsageDecisionKind::ExitWithSummary,
prompt_tokens,
ctx_max,
ratio,
tail_budget: None,
aggressive: false,
};
}
if already_folded_this_turn {
return PostUsageDecision {
kind: PostUsageDecisionKind::None,
prompt_tokens,
ctx_max,
ratio,
tail_budget: None,
aggressive: false,
};
}
if ratio > HISTORY_FOLD_AGGRESSIVE_THRESHOLD {
return PostUsageDecision {
kind: PostUsageDecisionKind::Fold,
prompt_tokens,
ctx_max,
ratio,
tail_budget: Some((ctx_max as f64 * HISTORY_FOLD_AGGRESSIVE_TAIL_FRACTION) as u64),
aggressive: true,
};
}
if ratio > effective_fold_threshold(fold_threshold_override) {
return PostUsageDecision {
kind: PostUsageDecisionKind::Fold,
prompt_tokens,
ctx_max,
ratio,
tail_budget: Some((ctx_max as f64 * HISTORY_FOLD_TAIL_FRACTION) as u64),
aggressive: false,
};
}
PostUsageDecision {
kind: PostUsageDecisionKind::None,
prompt_tokens,
ctx_max,
ratio,
tail_budget: None,
aggressive: false,
}
}
pub fn estimate_turn_start(estimate_tokens: u64, ctx_max: u64) -> TurnStartEstimate {
let ratio = if ctx_max == 0 {
f64::INFINITY
} else {
estimate_tokens as f64 / ctx_max as f64
};
TurnStartEstimate {
estimate_tokens,
ctx_max,
ratio,
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn resolve_context_target_defaults_and_floors() {
assert_eq!(resolve_context_target(None), DEFAULT_CONTEXT_TARGET);
assert_eq!(resolve_context_target(Some(50_000)), 50_000);
assert_eq!(resolve_context_target(Some(1_000)), MIN_CONTEXT_TARGET);
assert_eq!(resolve_context_target(Some(250_000)), 250_000);
}
#[test]
fn capped_budget_folds_within_the_target() {
let budget = DEFAULT_CONTEXT_TARGET;
assert_eq!(
decide_after_usage(Some(74_000), budget, false).kind,
PostUsageDecisionKind::None
);
assert_eq!(
decide_after_usage(Some(76_000), budget, false).kind,
PostUsageDecisionKind::Fold
);
}
#[test]
fn checkpoint_cadence_matches_mimo_by_window() {
assert!(
checkpoint_thresholds_for(8_000).is_empty(),
"tiny window disabled"
);
assert_eq!(checkpoint_thresholds_for(128_000), vec![0.2, 0.4, 0.6, 0.8]);
assert_eq!(checkpoint_thresholds_for(200_000), vec![0.2, 0.4, 0.6, 0.8]);
assert_eq!(checkpoint_thresholds_for(400_000).len(), 9, "10% cadence");
assert_eq!(checkpoint_thresholds_for(1_000_000).len(), 18, "5% cadence");
}
#[test]
fn schedule_fires_each_threshold_once_until_reset() {
let mut s = CheckpointSchedule::new(128_000); assert!(s.is_enabled());
assert!(!s.note_usage(0.1), "below first threshold");
assert!(s.note_usage(0.25), "crossed 20%");
assert!(!s.note_usage(0.30), "no new threshold");
assert!(s.note_usage(0.45), "crossed 40%");
assert!(s.note_usage(0.85), "crossed 60% and 80% together");
assert!(!s.note_usage(0.9), "all crossed");
s.reset();
assert!(s.note_usage(0.25), "fires again after reset");
}
#[test]
fn disabled_schedule_never_fires() {
let mut s = CheckpointSchedule::new(10_000);
assert!(!s.is_enabled());
assert!(!s.note_usage(0.99));
}
#[test]
fn effective_fold_threshold_clamps_override() {
assert_eq!(effective_fold_threshold(Some(0.5)), 0.5);
assert_eq!(effective_fold_threshold(Some(0.3)), 0.3);
assert_eq!(
effective_fold_threshold(Some(0.9)),
HISTORY_FOLD_THRESHOLD,
"above the default is rejected (can't fold later)"
);
assert_eq!(
effective_fold_threshold(Some(0.05)),
HISTORY_FOLD_THRESHOLD,
"below the floor is rejected"
);
assert_eq!(
effective_fold_threshold(Some(f64::NAN)),
HISTORY_FOLD_THRESHOLD
);
assert_eq!(effective_fold_threshold(None), HISTORY_FOLD_THRESHOLD);
}
#[test]
fn lower_override_folds_earlier() {
let d = decide_after_usage_with_threshold(Some(76_800), 128_000, false, None);
assert_eq!(d.kind, PostUsageDecisionKind::None);
let d = decide_after_usage_with_threshold(Some(76_800), 128_000, false, Some(0.5));
assert_eq!(d.kind, PostUsageDecisionKind::Fold);
assert!(!d.aggressive, "still the normal band, not aggressive");
let d = decide_after_usage_with_threshold(Some(110_000), 128_000, false, Some(0.5));
assert_eq!(d.kind, PostUsageDecisionKind::ExitWithSummary);
}
#[test]
fn no_usage_data_returns_none() {
let d = decide_after_usage(None, 128_000, false);
assert_eq!(d.kind, PostUsageDecisionKind::None);
assert_eq!(d.ratio, 0.0);
}
#[test]
fn below_threshold_returns_none() {
let d = decide_after_usage(Some(50_000), 128_000, false);
assert_eq!(d.kind, PostUsageDecisionKind::None);
}
#[test]
fn above_75pct_triggers_fold() {
let d = decide_after_usage(Some(98_000), 128_000, false);
assert_eq!(d.kind, PostUsageDecisionKind::Fold);
assert!(!d.aggressive);
assert_eq!(d.tail_budget, Some(25600));
}
#[test]
fn above_78pct_triggers_aggressive_fold() {
let d = decide_after_usage(Some(101_000), 128_000, false);
assert_eq!(d.kind, PostUsageDecisionKind::Fold);
assert!(d.aggressive);
assert_eq!(d.tail_budget, Some(12800));
}
#[test]
fn above_80pct_triggers_exit_with_summary() {
let d = decide_after_usage(Some(105_000), 128_000, false);
assert_eq!(d.kind, PostUsageDecisionKind::ExitWithSummary);
}
#[test]
fn already_folded_prevents_double_fold() {
let d = decide_after_usage(Some(100_000), 128_000, true);
assert_eq!(d.kind, PostUsageDecisionKind::None);
}
#[test]
fn already_folded_does_not_prevent_exit_with_summary() {
let d = decide_after_usage(Some(105_000), 128_000, true);
assert_eq!(d.kind, PostUsageDecisionKind::ExitWithSummary);
}
#[test]
fn zero_ctx_max_handled_gracefully() {
let d = decide_after_usage(Some(1000), 0, false);
assert_eq!(d.kind, PostUsageDecisionKind::None);
}
#[test]
fn estimate_below_threshold() {
let e = estimate_turn_start(50_000, 128_000);
assert!(e.ratio < TURN_START_FOLD_THRESHOLD);
assert_eq!(e.ctx_max, 128_000);
}
#[test]
fn estimate_above_threshold() {
let e = estimate_turn_start(120_000, 128_000);
assert!(e.ratio > TURN_START_FOLD_THRESHOLD);
}
#[test]
fn estimate_at_boundary() {
let boundary = (128_000.0 * TURN_START_FOLD_THRESHOLD) as u64;
let e = estimate_turn_start(boundary, 128_000);
assert!((e.ratio - TURN_START_FOLD_THRESHOLD).abs() < 0.001);
}
#[test]
fn thresholds_are_strictly_ordered() {
assert!(FORCE_SUMMARY_THRESHOLD > HISTORY_FOLD_AGGRESSIVE_THRESHOLD);
assert!(HISTORY_FOLD_AGGRESSIVE_THRESHOLD > HISTORY_FOLD_THRESHOLD);
assert!(HISTORY_FOLD_THRESHOLD > HISTORY_FOLD_MIN_SAVINGS_FRACTION);
}
#[test]
fn aggressive_tail_is_smaller_than_normal_tail() {
assert!(HISTORY_FOLD_AGGRESSIVE_TAIL_FRACTION < HISTORY_FOLD_TAIL_FRACTION);
}
}