use serde::{Deserialize, Serialize};
use crate::agent::AgentModelRef;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct CompactionConfig {
pub auto_enabled: bool,
pub trigger_mode: CompactionTriggerMode,
pub threshold: CompactionThreshold,
pub reserve_tokens: u64,
pub preserve: PreserveConfig,
pub prune: PruneConfig,
pub strategy: CompactionStrategy,
pub auto_continue: bool,
pub max_consecutive_failures: u32,
pub model: Option<AgentModelRef>,
pub summary_max_tokens: Option<u32>,
}
impl Default for CompactionConfig {
fn default() -> Self {
Self {
auto_enabled: true,
trigger_mode: CompactionTriggerMode::default(),
threshold: CompactionThreshold::default(),
reserve_tokens: 16_384,
preserve: PreserveConfig::default(),
prune: PruneConfig::default(),
strategy: CompactionStrategy::default(),
auto_continue: true,
max_consecutive_failures: 3,
model: None,
summary_max_tokens: Some(20_000),
}
}
}
impl CompactionConfig {
pub fn with_auto_enabled(mut self, enabled: bool) -> Self {
self.auto_enabled = enabled;
self
}
pub fn with_trigger_mode(mut self, mode: CompactionTriggerMode) -> Self {
self.trigger_mode = mode;
self
}
pub fn with_threshold(mut self, threshold: CompactionThreshold) -> Self {
self.threshold = threshold;
self
}
pub fn with_reserve_tokens(mut self, tokens: u64) -> Self {
self.reserve_tokens = tokens;
self
}
pub fn with_preserve(mut self, preserve: PreserveConfig) -> Self {
self.preserve = preserve;
self
}
pub fn with_prune(mut self, prune: PruneConfig) -> Self {
self.prune = prune;
self
}
pub fn with_strategy(mut self, strategy: CompactionStrategy) -> Self {
self.strategy = strategy;
self
}
pub fn with_auto_continue(mut self, auto_continue: bool) -> Self {
self.auto_continue = auto_continue;
self
}
pub fn with_max_consecutive_failures(mut self, max: u32) -> Self {
self.max_consecutive_failures = max;
self
}
pub fn with_model(mut self, model: AgentModelRef) -> Self {
self.model = Some(model);
self
}
pub fn with_summary_max_tokens(mut self, tokens: u32) -> Self {
self.summary_max_tokens = Some(tokens);
self
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum CompactionTriggerMode {
#[default]
Threshold,
Overflow,
}
#[derive(Debug, Clone, Copy, PartialEq, Default, Serialize, Deserialize)]
pub struct CompactionThreshold {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub tokens: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub ratio: Option<f64>,
}
impl CompactionThreshold {
pub fn fixed(tokens: u64) -> Self {
Self {
tokens: Some(tokens),
ratio: None,
}
}
pub fn ratio(ratio: f64) -> Self {
Self {
tokens: None,
ratio: Some(ratio),
}
}
pub fn resolve(&self, context_window: u64, reserve_tokens: u64) -> u64 {
if let Some(tokens) = self.tokens {
return tokens.clamp(1, context_window.saturating_sub(1));
}
if let Some(ratio) = self.ratio {
let clamped = ratio.clamp(0.01, 0.99);
return (context_window as f64 * clamped) as u64;
}
let effective_reserve = reserve_tokens.max((context_window as f64 * 0.15) as u64);
context_window.saturating_sub(effective_reserve)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct PreserveConfig {
pub recent_turns: u32,
pub recent_tokens: Option<u64>,
}
impl Default for PreserveConfig {
fn default() -> Self {
Self {
recent_turns: 2,
recent_tokens: None, }
}
}
impl PreserveConfig {
pub fn new(recent_turns: u32, recent_tokens: u64) -> Self {
Self {
recent_turns,
recent_tokens: Some(recent_tokens),
}
}
pub fn with_recent_turns(mut self, turns: u32) -> Self {
self.recent_turns = turns;
self
}
pub fn with_recent_tokens(mut self, tokens: u64) -> Self {
self.recent_tokens = Some(tokens);
self
}
pub fn resolve_recent_tokens(&self, usable_tokens: u64, min: u64, max: u64) -> u64 {
self.recent_tokens.unwrap_or_else(|| {
let auto = (usable_tokens as f64 * 0.25) as u64;
auto.clamp(min, max)
})
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PruneConfig {
pub enabled: bool,
pub protect_tokens: u64,
pub minimum_tokens: u64,
pub tool_output_max_chars: usize,
#[serde(default)]
pub protected_tools: Vec<String>,
}
impl Default for PruneConfig {
fn default() -> Self {
Self {
enabled: true,
protect_tokens: 40_000,
minimum_tokens: 20_000,
tool_output_max_chars: 2_000,
protected_tools: Vec::new(),
}
}
}
impl PruneConfig {
pub fn with_enabled(mut self, enabled: bool) -> Self {
self.enabled = enabled;
self
}
pub fn with_protect_tokens(mut self, tokens: u64) -> Self {
self.protect_tokens = tokens;
self
}
pub fn with_minimum_tokens(mut self, tokens: u64) -> Self {
self.minimum_tokens = tokens;
self
}
pub fn with_tool_output_max_chars(mut self, chars: usize) -> Self {
self.tool_output_max_chars = chars;
self
}
pub fn add_protected_tool(mut self, tool: impl Into<String>) -> Self {
self.protected_tools.push(tool.into());
self
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum CompactionStrategy {
#[default]
Summarize,
Handoff,
}
impl CompactionStrategy {
pub fn is_summarize(&self) -> bool {
matches!(self, Self::Summarize)
}
pub fn is_handoff(&self) -> bool {
matches!(self, Self::Handoff)
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct CompactionResult {
pub summary: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub short_summary: Option<String>,
pub trigger: CompactTrigger,
pub tokens_before: u64,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub tokens_after: Option<u64>,
pub messages_compacted: usize,
pub messages_kept: usize,
pub success: bool,
}
impl CompactionResult {
pub fn tokens_saved(&self) -> Option<u64> {
self.tokens_after
.map(|after| self.tokens_before.saturating_sub(after))
}
pub fn compression_ratio(&self) -> Option<f64> {
self.tokens_after.map(|after| {
if self.tokens_before == 0 {
return 0.0;
}
1.0 - (after as f64 / self.tokens_before as f64)
})
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum CompactTrigger {
Auto,
Manual,
Overflow,
Idle,
}
impl CompactTrigger {
pub fn is_auto(&self) -> bool {
!matches!(self, Self::Manual)
}
pub fn is_manual(&self) -> bool {
matches!(self, Self::Manual)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
#[serde(tag = "level", rename_all = "snake_case")]
pub enum TokenBudgetState {
Normal {
percent_remaining: f64,
},
Warning {
percent_remaining: f64,
},
Error {
percent_remaining: f64,
},
Blocking,
}
impl TokenBudgetState {
pub fn from_usage(
used_tokens: u64,
context_window: u64,
auto_compact_buffer: u64,
) -> Self {
if context_window == 0 {
return Self::Blocking;
}
let percent_remaining = 1.0 - (used_tokens as f64 / context_window as f64);
if used_tokens >= context_window {
return Self::Blocking;
}
let auto_threshold = context_window.saturating_sub(auto_compact_buffer);
let error_threshold = context_window.saturating_sub(20_000);
let warning_threshold = auto_threshold.saturating_sub(20_000);
if used_tokens >= error_threshold {
Self::Error { percent_remaining }
} else if used_tokens >= warning_threshold {
Self::Warning { percent_remaining }
} else {
Self::Normal { percent_remaining }
}
}
pub fn should_auto_compact(&self) -> bool {
matches!(self, Self::Error { .. } | Self::Blocking)
}
pub fn is_blocking(&self) -> bool {
matches!(self, Self::Blocking)
}
pub fn is_warning_or_worse(&self) -> bool {
!matches!(self, Self::Normal { .. })
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_default_config() {
let config = CompactionConfig::default();
assert!(config.auto_enabled);
assert_eq!(config.trigger_mode, CompactionTriggerMode::Threshold);
assert_eq!(config.reserve_tokens, 16_384);
assert_eq!(config.preserve.recent_turns, 2);
assert!(config.prune.enabled);
assert_eq!(config.strategy, CompactionStrategy::Summarize);
assert!(config.auto_continue);
assert_eq!(config.max_consecutive_failures, 3);
assert!(config.model.is_none());
assert_eq!(config.summary_max_tokens, Some(20_000));
}
#[test]
fn test_config_builder() {
let config = CompactionConfig::default()
.with_auto_enabled(false)
.with_trigger_mode(CompactionTriggerMode::Overflow)
.with_reserve_tokens(20_000)
.with_auto_continue(false)
.with_max_consecutive_failures(5);
assert!(!config.auto_enabled);
assert_eq!(config.trigger_mode, CompactionTriggerMode::Overflow);
assert_eq!(config.reserve_tokens, 20_000);
assert!(!config.auto_continue);
assert_eq!(config.max_consecutive_failures, 5);
}
#[test]
fn test_config_serde_roundtrip() {
let config = CompactionConfig::default()
.with_strategy(CompactionStrategy::Handoff)
.with_prune(PruneConfig::default().with_enabled(false));
let json = serde_json::to_string(&config).unwrap();
let restored: CompactionConfig = serde_json::from_str(&json).unwrap();
assert_eq!(config, restored);
}
#[test]
fn test_threshold_fixed() {
let t = CompactionThreshold::fixed(150_000);
assert_eq!(t.resolve(200_000, 16_384), 150_000);
}
#[test]
fn test_threshold_fixed_clamp() {
let t = CompactionThreshold::fixed(300_000);
assert_eq!(t.resolve(200_000, 16_384), 199_999);
}
#[test]
fn test_threshold_ratio() {
let t = CompactionThreshold::ratio(0.85);
assert_eq!(t.resolve(200_000, 16_384), 170_000);
}
#[test]
fn test_threshold_fallback() {
let t = CompactionThreshold::default();
assert_eq!(t.resolve(200_000, 16_384), 170_000);
}
#[test]
fn test_threshold_fallback_small_window() {
let t = CompactionThreshold::default();
assert_eq!(t.resolve(50_000, 16_384), 33_616);
}
#[test]
fn test_preserve_default() {
let p = PreserveConfig::default();
assert_eq!(p.recent_turns, 2);
assert!(p.recent_tokens.is_none());
}
#[test]
fn test_preserve_resolve_auto() {
let p = PreserveConfig::default();
assert_eq!(p.resolve_recent_tokens(100_000, 2_000, 8_000), 8_000);
assert_eq!(p.resolve_recent_tokens(4_000, 2_000, 8_000), 2_000);
assert_eq!(p.resolve_recent_tokens(20_000, 2_000, 8_000), 5_000);
}
#[test]
fn test_preserve_resolve_explicit() {
let p = PreserveConfig::new(3, 15_000);
assert_eq!(p.resolve_recent_tokens(100_000, 2_000, 8_000), 15_000);
}
#[test]
fn test_prune_default() {
let p = PruneConfig::default();
assert!(p.enabled);
assert_eq!(p.protect_tokens, 40_000);
assert_eq!(p.minimum_tokens, 20_000);
assert_eq!(p.tool_output_max_chars, 2_000);
assert!(p.protected_tools.is_empty());
}
#[test]
fn test_prune_builder() {
let p = PruneConfig::default()
.with_enabled(false)
.with_protect_tokens(50_000)
.add_protected_tool("skill")
.add_protected_tool("memory");
assert!(!p.enabled);
assert_eq!(p.protect_tokens, 50_000);
assert_eq!(p.protected_tools, vec!["skill", "memory"]);
}
#[test]
fn test_strategy_predicates() {
assert!(CompactionStrategy::Summarize.is_summarize());
assert!(!CompactionStrategy::Summarize.is_handoff());
assert!(CompactionStrategy::Handoff.is_handoff());
assert!(!CompactionStrategy::Handoff.is_summarize());
}
#[test]
fn test_strategy_serde() {
let json = serde_json::to_string(&CompactionStrategy::Handoff).unwrap();
assert_eq!(json, r#""handoff""#);
let restored: CompactionStrategy = serde_json::from_str(&json).unwrap();
assert_eq!(restored, CompactionStrategy::Handoff);
}
#[test]
fn test_trigger_is_auto() {
assert!(CompactTrigger::Auto.is_auto());
assert!(CompactTrigger::Overflow.is_auto());
assert!(CompactTrigger::Idle.is_auto());
assert!(!CompactTrigger::Manual.is_auto());
}
#[test]
fn test_trigger_serde() {
for trigger in [
CompactTrigger::Auto,
CompactTrigger::Manual,
CompactTrigger::Overflow,
CompactTrigger::Idle,
] {
let json = serde_json::to_string(&trigger).unwrap();
let restored: CompactTrigger = serde_json::from_str(&json).unwrap();
assert_eq!(trigger, restored);
}
}
#[test]
fn test_result_tokens_saved() {
let result = CompactionResult {
summary: "test".into(),
short_summary: None,
trigger: CompactTrigger::Auto,
tokens_before: 150_000,
tokens_after: Some(5_000),
messages_compacted: 40,
messages_kept: 8,
success: true,
};
assert_eq!(result.tokens_saved(), Some(145_000));
}
#[test]
fn test_result_compression_ratio() {
let result = CompactionResult {
summary: "test".into(),
short_summary: None,
trigger: CompactTrigger::Auto,
tokens_before: 100_000,
tokens_after: Some(10_000),
messages_compacted: 30,
messages_kept: 5,
success: true,
};
let ratio = result.compression_ratio().unwrap();
assert!((ratio - 0.9).abs() < 0.001);
}
#[test]
fn test_result_no_tokens_after() {
let result = CompactionResult {
summary: "test".into(),
short_summary: None,
trigger: CompactTrigger::Manual,
tokens_before: 100_000,
tokens_after: None,
messages_compacted: 20,
messages_kept: 5,
success: true,
};
assert!(result.tokens_saved().is_none());
assert!(result.compression_ratio().is_none());
}
#[test]
fn test_result_serde_roundtrip() {
let result = CompactionResult {
summary: "The user asked about Rust ownership...".into(),
short_summary: Some("Discussed Rust ownership".into()),
trigger: CompactTrigger::Overflow,
tokens_before: 180_000,
tokens_after: Some(8_000),
messages_compacted: 50,
messages_kept: 6,
success: true,
};
let json = serde_json::to_string(&result).unwrap();
let restored: CompactionResult = serde_json::from_str(&json).unwrap();
assert_eq!(result, restored);
}
#[test]
fn test_budget_state_normal() {
let state = TokenBudgetState::from_usage(50_000, 200_000, 13_000);
assert!(matches!(state, TokenBudgetState::Normal { .. }));
assert!(!state.should_auto_compact());
assert!(!state.is_blocking());
assert!(!state.is_warning_or_worse());
}
#[test]
fn test_budget_state_warning() {
let state = TokenBudgetState::from_usage(170_000, 200_000, 13_000);
assert!(matches!(state, TokenBudgetState::Warning { .. }));
assert!(!state.should_auto_compact());
assert!(state.is_warning_or_worse());
}
#[test]
fn test_budget_state_error() {
let state = TokenBudgetState::from_usage(185_000, 200_000, 13_000);
assert!(matches!(state, TokenBudgetState::Error { .. }));
assert!(state.should_auto_compact());
assert!(state.is_warning_or_worse());
}
#[test]
fn test_budget_state_blocking() {
let state = TokenBudgetState::from_usage(200_000, 200_000, 13_000);
assert!(matches!(state, TokenBudgetState::Blocking));
assert!(state.should_auto_compact());
assert!(state.is_blocking());
}
#[test]
fn test_budget_state_zero_window() {
let state = TokenBudgetState::from_usage(0, 0, 0);
assert!(matches!(state, TokenBudgetState::Blocking));
}
#[test]
fn test_budget_state_serde_roundtrip() {
let states = vec![
TokenBudgetState::Normal {
percent_remaining: 0.75,
},
TokenBudgetState::Warning {
percent_remaining: 0.15,
},
TokenBudgetState::Error {
percent_remaining: 0.05,
},
TokenBudgetState::Blocking,
];
for state in states {
let json = serde_json::to_string(&state).unwrap();
let restored: TokenBudgetState = serde_json::from_str(&json).unwrap();
assert_eq!(state, restored);
}
}
}