Skip to main content

ai_agent/services/compact/
compact_warning_hook.rs

1//! Compact warning hook - triggered when compaction is needed
2
3use serde::{Deserialize, Serialize};
4
5/// Warning level for compaction
6#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
7pub enum WarningLevel {
8    /// No warning - everything is fine
9    None,
10    /// Info level - user should be aware
11    Info,
12    /// Warning level - user might want to compact
13    Warning,
14    /// Critical - immediate compaction recommended
15    Critical,
16}
17
18/// Compact warning info
19#[derive(Debug, Clone, Serialize, Deserialize)]
20pub struct CompactWarningInfo {
21    /// Current token count
22    pub current_tokens: u64,
23    /// Maximum tokens allowed
24    pub max_tokens: u64,
25    /// Utilization percentage (0-100)
26    pub utilization: f64,
27    /// Warning level
28    pub level: WarningLevel,
29    /// Suggested action
30    pub suggestion: Option<String>,
31}
32
33/// Get warning level based on utilization
34pub fn get_warning_level(utilization: f64) -> WarningLevel {
35    if utilization >= 0.95 {
36        WarningLevel::Critical
37    } else if utilization >= 0.85 {
38        WarningLevel::Warning
39    } else if utilization >= 0.75 {
40        WarningLevel::Info
41    } else {
42        WarningLevel::None
43    }
44}
45
46/// Create compact warning info
47pub fn create_compact_warning_info(current_tokens: u64, max_tokens: u64) -> CompactWarningInfo {
48    let utilization = if max_tokens > 0 {
49        (current_tokens as f64 / max_tokens as f64) * 100.0
50    } else {
51        0.0
52    };
53
54    let level = get_warning_level(utilization / 100.0);
55
56    let suggestion = match level {
57        WarningLevel::Critical => Some("Run /compact now to avoid losing context".to_string()),
58        WarningLevel::Warning => Some("Consider running /compact soon".to_string()),
59        WarningLevel::Info => Some("You may want to run /compact later".to_string()),
60        WarningLevel::None => None,
61    };
62
63    CompactWarningInfo {
64        current_tokens,
65        max_tokens,
66        utilization,
67        level,
68        suggestion,
69    }
70}
71
72#[cfg(test)]
73mod tests {
74    use super::*;
75
76    #[test]
77    fn test_get_warning_level_critical() {
78        assert_eq!(get_warning_level(0.96), WarningLevel::Critical);
79    }
80
81    #[test]
82    fn test_get_warning_level_warning() {
83        assert_eq!(get_warning_level(0.86), WarningLevel::Warning);
84    }
85
86    #[test]
87    fn test_get_warning_level_info() {
88        assert_eq!(get_warning_level(0.76), WarningLevel::Info);
89    }
90
91    #[test]
92    fn test_get_warning_level_none() {
93        assert_eq!(get_warning_level(0.5), WarningLevel::None);
94    }
95
96    #[test]
97    fn test_create_compact_warning_info() {
98        let info = create_compact_warning_info(80000, 100000);
99        assert_eq!(info.current_tokens, 80000);
100        assert_eq!(info.max_tokens, 100000);
101        assert!((info.utilization - 80.0).abs() < 0.01);
102    }
103}