use super::*;
#[tokio::test]
async fn token_budget_prunes_oldest_when_over_budget() {
let mut tb = TokenBudget::new(512, 0.99);
let budget = 506;
let big = "x".repeat(200); tb.sync_turn(&big, &big, &metrics_with_input(200)).await;
assert_eq!(tb.history.len(), 1);
tb.sync_turn(&big, &big, &metrics_with_input(520)).await;
assert!(tb.used_tokens() <= budget, "used = {}", tb.used_tokens());
assert_eq!(tb.history.len(), 1, "the oldest turn must have been pruned");
}
#[tokio::test]
async fn token_budget_uses_metrics_when_available() {
let mut tb = TokenBudget::new(1000, 1.0);
let mut m = dummy_metrics();
m.usage = Some(crate::metrics::TokenUsage {
input_tokens: 30,
output_tokens: 20,
});
tb.sync_turn("q", "a", &m).await;
assert_eq!(tb.used_tokens(), 31); }
#[tokio::test]
async fn token_budget_no_runaway_drift_b3_regression() {
let mut tb = TokenBudget::new(8_192, 0.80); let mut old_running_sum: u64 = 0;
let turns = 20u32;
for i in 0..turns {
let input = 2_582 + (4_748 - 2_582) * i / (turns - 1);
old_running_sum += u64::from(input) + 20;
tb.sync_turn("reply ok", "ok", &metrics_with_input(input))
.await;
}
let used = u64::from(tb.used_tokens());
assert!(
(4_748..4_800).contains(&used),
"used must track the last real prompt, got {used}"
);
assert_eq!(tb.history.len(), turns as usize, "no spurious pruning");
assert_eq!(tb.pruned_count, 0);
assert!(
used * 5 < old_running_sum,
"anchored used ({used}) must be at least 5× below the old \
running sum ({old_running_sum})"
);
}
fn turn_with_tokens(
user: &str,
assistant: &str,
tokens_in: Option<u32>,
tokens_out: Option<u32>,
) -> crate::ConversationTurn {
let mut t = crate::ConversationTurn::new(user, assistant);
t.tokens_in = tokens_in;
t.tokens_out = tokens_out;
t
}
#[tokio::test]
async fn summarizing_no_runaway_compression_b3_regression() {
let mut s = Summarizing::new(8_192) .with_summarizer(stub_summarizer("SUMMARY"));
for i in 0..20u32 {
let input = 2_582 + (4_748 - 2_582) * i / 19;
s.sync_turn("reply ok", "ok", &metrics_with_input(input))
.await;
}
assert_eq!(
s.compress_count, 0,
"prompts never exceeded the budget — compression must not fire"
);
assert!(s.prev_summary.is_empty());
assert_eq!(s.history.len(), 20);
}
#[test]
fn token_budget_restore_anchors_on_column_tokens() {
let mut tb = TokenBudget::new(100_000, 0.80);
let turns = vec![
turn_with_tokens("u1", "a1", Some(900), Some(10)),
turn_with_tokens("u2", "aa", Some(1_000), Some(12)),
];
tb.restore_turns(&turns);
assert_eq!(tb.last_prompt_tokens, Some(1_000));
assert_eq!(tb.used_tokens(), 1_001);
}
#[test]
fn token_budget_restore_null_columns_fall_back_to_estimate() {
let mut tb = TokenBudget::new(100_000, 0.80);
let text = "x".repeat(40); let turns = vec![
crate::ConversationTurn::new(text.clone(), text.clone()),
crate::ConversationTurn::new(text.clone(), text.clone()),
];
tb.restore_turns(&turns);
assert_eq!(
tb.last_prompt_tokens, None,
"no measurement in the store → no anchor, only estimates"
);
assert_eq!(tb.used_tokens(), 40, "2 turns × (10 + 10) est tokens");
}
#[test]
fn token_budget_restore_unmeasured_tail_extends_delta() {
let mut tb = TokenBudget::new(100_000, 0.80);
let turns = vec![
turn_with_tokens("u1", "aaaa", Some(1_000), Some(8)), crate::ConversationTurn::new("xxxx", "yyyy"), ];
tb.restore_turns(&turns);
assert_eq!(tb.last_prompt_tokens, Some(1_000));
assert_eq!(tb.used_tokens(), 1_003);
}
#[test]
fn summarizing_restore_anchors_on_column_tokens() {
let mut s = Summarizing::new(100_000);
let turns = vec![
turn_with_tokens("u1", "a1", Some(700), Some(9)),
turn_with_tokens("u2", "aaaa", Some(2_000), Some(11)),
];
s.restore_turns(&turns);
assert_eq!(s.last_prompt_tokens, Some(2_000));
assert_eq!(s.used_tokens(), 2_001); }
#[test]
fn summarizing_restore_ignores_measurements_before_the_cut() {
let marked = format!(
"{}\nolder work compressed\n{}",
crate::agentic::SUMMARY_PREFIX,
crate::agentic::SUMMARY_END_MARKER
);
let mut s = Summarizing::new(100_000);
let turns = vec![
turn_with_tokens("old", "old", Some(50_000), Some(10)),
crate::ConversationTurn::new(marked, ""),
crate::ConversationTurn::new("after", "compaction"), ];
s.restore_turns(&turns);
assert_eq!(
s.last_prompt_tokens, None,
"a pre-compression measurement must not anchor the cut working set"
);
}
#[tokio::test]
async fn token_budget_does_not_prune_within_budget() {
let mut tb = TokenBudget::new(200, 1.0); let mut m = dummy_metrics();
m.usage = Some(crate::metrics::TokenUsage {
input_tokens: 50,
output_tokens: 50,
});
tb.sync_turn("q", "a", &m).await; assert_eq!(tb.history.len(), 1);
}
#[tokio::test]
async fn token_budget_estimate_fallback_counts_each_turn_once() {
let mut tb = TokenBudget::new(1000, 1.0);
let mut m = dummy_metrics();
m.usage = None; let text = "x".repeat(40); tb.sync_turn(&text, &text, &m).await;
tb.sync_turn(&text, &text, &m).await;
assert_eq!(tb.used_tokens(), 40, "2 turns × (10 + 10) est tokens");
}