claude_agent/session/
compact.rs

1//! Context Compaction (Claude Code CLI compatible)
2//!
3//! Summarizes the entire conversation when context exceeds threshold.
4
5use serde::{Deserialize, Serialize};
6
7use super::state::{Session, SessionMessage};
8use super::types::CompactRecord;
9use super::{SessionError, SessionResult};
10use crate::client::DEFAULT_SMALL_MODEL;
11use crate::types::{CompactResult, ContentBlock, DEFAULT_COMPACT_THRESHOLD, Role};
12
13#[derive(Clone, Debug, Serialize, Deserialize)]
14pub struct CompactStrategy {
15    pub enabled: bool,
16    pub threshold_percent: f32,
17    pub summary_model: String,
18    pub max_summary_tokens: u32,
19}
20
21impl Default for CompactStrategy {
22    fn default() -> Self {
23        Self {
24            enabled: true,
25            threshold_percent: DEFAULT_COMPACT_THRESHOLD,
26            summary_model: DEFAULT_SMALL_MODEL.to_string(),
27            max_summary_tokens: 4000,
28        }
29    }
30}
31
32impl CompactStrategy {
33    pub fn disabled() -> Self {
34        Self {
35            enabled: false,
36            ..Default::default()
37        }
38    }
39
40    pub fn with_threshold(mut self, percent: f32) -> Self {
41        self.threshold_percent = percent.clamp(0.5, 0.95);
42        self
43    }
44
45    pub fn with_model(mut self, model: impl Into<String>) -> Self {
46        self.summary_model = model.into();
47        self
48    }
49}
50
51pub struct CompactExecutor {
52    strategy: CompactStrategy,
53}
54
55impl CompactExecutor {
56    pub fn new(strategy: CompactStrategy) -> Self {
57        Self { strategy }
58    }
59
60    pub fn needs_compact(&self, current_tokens: u64, max_tokens: u64) -> bool {
61        if !self.strategy.enabled {
62            return false;
63        }
64        let threshold = (max_tokens as f32 * self.strategy.threshold_percent) as u64;
65        current_tokens >= threshold
66    }
67
68    pub fn prepare_compact(&self, session: &Session) -> SessionResult<PreparedCompact> {
69        if !self.strategy.enabled {
70            return Err(SessionError::Compact {
71                message: "Compact is disabled".to_string(),
72            });
73        }
74
75        let messages = session.get_current_branch();
76        if messages.is_empty() {
77            return Ok(PreparedCompact::NotNeeded);
78        }
79
80        let summary_prompt = self.format_for_summary(&messages);
81
82        Ok(PreparedCompact::Ready {
83            summary_prompt,
84            message_count: messages.len(),
85        })
86    }
87
88    pub fn apply_compact(&self, session: &mut Session, summary: String) -> CompactResult {
89        let original_count = session.messages.len();
90
91        let removed_chars: usize = session
92            .messages
93            .iter()
94            .map(|m| {
95                m.content
96                    .iter()
97                    .filter_map(|c| c.as_text())
98                    .map(|t| t.len())
99                    .sum::<usize>()
100            })
101            .sum();
102        let saved_tokens = (removed_chars / 4) as u64;
103
104        session.messages.clear();
105        session.summary = Some(summary.clone());
106
107        let summary_msg = SessionMessage::user(vec![ContentBlock::text(format!(
108            "[Previous conversation summary]\n\n{}",
109            summary
110        ))])
111        .as_compact_summary();
112
113        session.add_message(summary_msg);
114
115        CompactResult::Compacted {
116            original_count,
117            new_count: 1,
118            saved_tokens: saved_tokens as usize,
119            summary,
120        }
121    }
122
123    pub fn record_compact(&self, session: &mut Session, result: &CompactResult) {
124        if let CompactResult::Compacted {
125            original_count,
126            new_count,
127            saved_tokens,
128            summary,
129        } = result
130        {
131            let record = CompactRecord::new(session.id)
132                .with_counts(*original_count, *new_count)
133                .with_summary(summary.clone())
134                .with_saved_tokens(*saved_tokens);
135            session.record_compact(record);
136        }
137    }
138
139    fn format_for_summary(&self, messages: &[&SessionMessage]) -> String {
140        let mut formatted = String::new();
141        formatted.push_str(COMPACTION_PROMPT);
142        formatted.push_str("\n\n---\n\n");
143
144        for msg in messages {
145            let role = match msg.role {
146                Role::User => "Human",
147                Role::Assistant => "Assistant",
148            };
149
150            formatted.push_str(&format!("**{}**:\n", role));
151
152            for block in &msg.content {
153                if let Some(text) = block.as_text() {
154                    let display_text = if text.len() > 3000 {
155                        format!("{}... [truncated]", &text[..3000])
156                    } else {
157                        text.to_string()
158                    };
159                    formatted.push_str(&display_text);
160                    formatted.push('\n');
161                }
162            }
163            formatted.push('\n');
164        }
165
166        formatted
167    }
168
169    pub fn strategy(&self) -> &CompactStrategy {
170        &self.strategy
171    }
172}
173
174const COMPACTION_PROMPT: &str = r#"Summarize this conversation to continue seamlessly. Preserve:
175
1761. **Original Request**: The core task or question
1772. **Decisions Made**: Architecture, design, approach choices
1783. **Files Modified**: List with brief context
1794. **Code Changes**: Functions, modules modified
1805. **Current Progress**: Completed work and remaining tasks
1816. **Errors & Fixes**: Issues encountered and resolutions
1827. **Key Context**: Constraints, preferences, project structure
183
184Format as structured sections. Be concise but complete."#;
185
186#[derive(Debug)]
187pub enum PreparedCompact {
188    NotNeeded,
189    Ready {
190        summary_prompt: String,
191        message_count: usize,
192    },
193}
194
195#[cfg(test)]
196mod tests {
197    use super::*;
198    use crate::session::state::SessionConfig;
199
200    fn create_test_session(message_count: usize) -> Session {
201        let mut session = Session::new(SessionConfig::default());
202
203        for i in 0..message_count {
204            let content = if i % 2 == 0 {
205                format!("User message {}", i)
206            } else {
207                format!("Assistant response {}", i)
208            };
209
210            let msg = if i % 2 == 0 {
211                SessionMessage::user(vec![ContentBlock::text(content)])
212            } else {
213                SessionMessage::assistant(vec![ContentBlock::text(content)])
214            };
215
216            session.add_message(msg);
217        }
218
219        session
220    }
221
222    #[test]
223    fn test_compact_strategy_default() {
224        let strategy = CompactStrategy::default();
225        assert!(strategy.enabled);
226        assert_eq!(strategy.threshold_percent, 0.8);
227    }
228
229    #[test]
230    fn test_compact_strategy_disabled() {
231        let strategy = CompactStrategy::disabled();
232        assert!(!strategy.enabled);
233    }
234
235    #[test]
236    fn test_needs_compact() {
237        let executor = CompactExecutor::new(CompactStrategy::default().with_threshold(0.8));
238
239        assert!(!executor.needs_compact(70_000, 100_000));
240        assert!(executor.needs_compact(80_000, 100_000));
241        assert!(executor.needs_compact(90_000, 100_000));
242    }
243
244    #[test]
245    fn test_prepare_compact_empty() {
246        let session = Session::new(SessionConfig::default());
247        let executor = CompactExecutor::new(CompactStrategy::default());
248
249        let result = executor.prepare_compact(&session).unwrap();
250        assert!(matches!(result, PreparedCompact::NotNeeded));
251    }
252
253    #[test]
254    fn test_prepare_compact_ready() {
255        let session = create_test_session(10);
256        let executor = CompactExecutor::new(CompactStrategy::default());
257
258        let result = executor.prepare_compact(&session).unwrap();
259
260        match result {
261            PreparedCompact::Ready {
262                summary_prompt,
263                message_count,
264            } => {
265                assert!(summary_prompt.contains("Original Request"));
266                assert_eq!(message_count, 10);
267            }
268            _ => panic!("Expected Ready"),
269        }
270    }
271
272    #[test]
273    fn test_apply_compact() {
274        let mut session = create_test_session(10);
275        let executor = CompactExecutor::new(CompactStrategy::default());
276
277        let result = executor.apply_compact(&mut session, "Test summary".to_string());
278
279        match result {
280            CompactResult::Compacted {
281                original_count,
282                new_count,
283                ..
284            } => {
285                assert_eq!(original_count, 10);
286                assert_eq!(new_count, 1);
287            }
288            _ => panic!("Expected Compacted"),
289        }
290
291        assert!(session.summary.is_some());
292        assert_eq!(session.messages.len(), 1);
293        assert!(session.messages[0].is_compact_summary);
294    }
295}
claude_agent/session/compact.rs

claude_agent/session/
compact.rs
