battlecommand_forge/
context.rs

1//! Context compaction — prevents token limit crashes.
2//! Ported from battleclaw-v2 context_compact.rs.
3
4const MAX_CONTEXT_CHARS: usize = 120_000; // ~30K tokens
5const COMPACT_THRESHOLD: f64 = 0.95;
6const TARGET_AFTER_COMPACT: f64 = 0.60;
7
8#[derive(Debug, Clone)]
9pub struct ContextMessage {
10    pub role: String,
11    pub content: String,
12    pub compactable: bool,
13}
14
15pub struct ContextManager {
16    messages: Vec<ContextMessage>,
17    total_chars: usize,
18}
19
20impl Default for ContextManager {
21    fn default() -> Self {
22        Self::new()
23    }
24}
25
26impl ContextManager {
27    pub fn new() -> Self {
28        Self {
29            messages: Vec::new(),
30            total_chars: 0,
31        }
32    }
33
34    pub fn add(&mut self, role: &str, content: &str, compactable: bool) {
35        self.total_chars += content.len();
36        self.messages.push(ContextMessage {
37            role: role.to_string(),
38            content: content.to_string(),
39            compactable,
40        });
41        if self.usage_ratio() >= COMPACT_THRESHOLD {
42            self.compact();
43        }
44    }
45
46    pub fn usage_ratio(&self) -> f64 {
47        self.total_chars as f64 / MAX_CONTEXT_CHARS as f64
48    }
49
50    pub fn usage_percent(&self) -> u32 {
51        (self.usage_ratio() * 100.0) as u32
52    }
53
54    pub fn to_string(&self) -> String {
55        self.messages
56            .iter()
57            .map(|m| format!("{}: {}", m.role, m.content))
58            .collect::<Vec<_>>()
59            .join("\n\n")
60    }
61
62    pub fn len(&self) -> usize {
63        self.messages.len()
64    }
65
66    pub fn compact(&mut self) {
67        let target_chars = (MAX_CONTEXT_CHARS as f64 * TARGET_AFTER_COMPACT) as usize;
68
69        // Phase 1: Truncate long compactable messages
70        for msg in &mut self.messages {
71            if msg.compactable && msg.content.len() > 500 {
72                let end = msg.content.len().min(200);
73                let truncated = msg.content.len() - end;
74                msg.content = format!("{}...[truncated {} chars]", &msg.content[..end], truncated);
75            }
76        }
77        self.recalc();
78        if self.total_chars <= target_chars {
79            return;
80        }
81
82        // Phase 2: Drop old compactable messages (keep last 20)
83        if self.messages.len() > 20 {
84            let to_remove = self.messages.len() - 20;
85            let summary = format!(
86                "[Compacted {} earlier messages at {}% capacity]",
87                to_remove,
88                self.usage_percent()
89            );
90            self.messages.drain(..to_remove);
91            self.messages.insert(
92                0,
93                ContextMessage {
94                    role: "system".into(),
95                    content: summary,
96                    compactable: false,
97                },
98            );
99        }
100        self.recalc();
101    }
102
103    fn recalc(&mut self) {
104        self.total_chars = self.messages.iter().map(|m| m.content.len()).sum();
105    }
106}
107
108#[cfg(test)]
109mod tests {
110    use super::*;
111
112    #[test]
113    fn test_basic_add() {
114        let mut cm = ContextManager::new();
115        cm.add("user", "hello", false);
116        assert_eq!(cm.len(), 1);
117        assert_eq!(cm.total_chars, 5);
118    }
119
120    #[test]
121    fn test_auto_compact() {
122        let mut cm = ContextManager::new();
123        // Add enough to trigger compaction
124        for _ in 0..200 {
125            cm.add("user", &"x".repeat(1000), true);
126        }
127        // Should have compacted
128        assert!(cm.total_chars < MAX_CONTEXT_CHARS);
129    }
130
131    #[test]
132    fn test_usage_ratio() {
133        let mut cm = ContextManager::new();
134        cm.add("user", &"a".repeat(60_000), false);
135        assert!((cm.usage_ratio() - 0.5).abs() < 0.01);
136    }
137}
battlecommand_forge/context.rs

battlecommand_forge/
context.rs
