Skip to main content

enact_context/
compactor.rs

1//! Context Compaction
2//!
3//! Strategies for reducing context size when approaching limits.
4//!
5//! @see packages/enact-schemas/src/context.schemas.ts
6
7use crate::segment::{ContextPriority, ContextSegment, ContextSegmentType};
8use chrono::{DateTime, Utc};
9use enact_core::kernel::ExecutionId;
10use serde::{Deserialize, Serialize};
11use thiserror::Error;
12
13/// Compaction errors
14#[derive(Debug, Error)]
15pub enum CompactionError {
16    #[error("Nothing to compact")]
17    NothingToCompact,
18
19    #[error("Target token count too low: {0}")]
20    TargetTooLow(usize),
21
22    #[error("Summarization failed: {0}")]
23    SummarizationFailed(String),
24}
25
26/// Available compaction strategies
27///
28/// Matches `compactionStrategyTypeSchema` in @enact/schemas
29#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
30#[serde(rename_all = "snake_case")]
31pub enum CompactionStrategyType {
32    /// Simple truncation (remove oldest)
33    Truncate,
34    /// LLM summarization
35    Summarize,
36    /// Extract key points only
37    ExtractKeyPoints,
38    /// Keep only recent N messages
39    SlidingWindow,
40    /// Keep based on importance scores
41    ImportanceWeighted,
42    /// Combination of strategies
43    Hybrid,
44}
45
46/// Configuration for context compaction
47///
48/// Matches `compactionStrategySchema` in @enact/schemas
49#[derive(Debug, Clone, Serialize, Deserialize)]
50#[serde(rename_all = "camelCase")]
51pub struct CompactionStrategy {
52    /// Strategy type
53    #[serde(rename = "type")]
54    pub strategy_type: CompactionStrategyType,
55
56    /// Target token count after compaction
57    pub target_tokens: usize,
58
59    /// Minimum content to preserve (percentage, 0-100)
60    pub min_preserve_percent: u8,
61
62    /// Segments to compact (in priority order)
63    #[serde(skip_serializing_if = "Option::is_none")]
64    pub segments_to_compact: Option<Vec<ContextSegmentType>>,
65
66    /// Segments to never compact
67    #[serde(skip_serializing_if = "Option::is_none")]
68    pub protected_segments: Option<Vec<ContextSegmentType>>,
69
70    /// For summarize strategy: max summary tokens
71    #[serde(skip_serializing_if = "Option::is_none")]
72    pub summary_max_tokens: Option<usize>,
73
74    /// For sliding_window strategy: window size
75    #[serde(skip_serializing_if = "Option::is_none")]
76    pub window_size: Option<usize>,
77
78    /// For importance_weighted: minimum importance score to keep
79    #[serde(skip_serializing_if = "Option::is_none")]
80    pub min_importance_score: Option<f64>,
81}
82
83impl CompactionStrategy {
84    /// Create a truncation strategy
85    pub fn truncate(target_tokens: usize) -> Self {
86        Self {
87            strategy_type: CompactionStrategyType::Truncate,
88            target_tokens,
89            min_preserve_percent: 20,
90            segments_to_compact: None,
91            protected_segments: Some(vec![
92                ContextSegmentType::System,
93                ContextSegmentType::UserInput,
94            ]),
95            summary_max_tokens: None,
96            window_size: None,
97            min_importance_score: None,
98        }
99    }
100
101    /// Create a sliding window strategy
102    pub fn sliding_window(target_tokens: usize, window_size: usize) -> Self {
103        Self {
104            strategy_type: CompactionStrategyType::SlidingWindow,
105            target_tokens,
106            min_preserve_percent: 20,
107            segments_to_compact: Some(vec![ContextSegmentType::History]),
108            protected_segments: Some(vec![
109                ContextSegmentType::System,
110                ContextSegmentType::UserInput,
111            ]),
112            summary_max_tokens: None,
113            window_size: Some(window_size),
114            min_importance_score: None,
115        }
116    }
117
118    /// Create a summarization strategy
119    pub fn summarize(target_tokens: usize, summary_max_tokens: usize) -> Self {
120        Self {
121            strategy_type: CompactionStrategyType::Summarize,
122            target_tokens,
123            min_preserve_percent: 30,
124            segments_to_compact: Some(vec![
125                ContextSegmentType::History,
126                ContextSegmentType::ToolResults,
127            ]),
128            protected_segments: Some(vec![
129                ContextSegmentType::System,
130                ContextSegmentType::UserInput,
131                ContextSegmentType::Guidance,
132            ]),
133            summary_max_tokens: Some(summary_max_tokens),
134            window_size: None,
135            min_importance_score: None,
136        }
137    }
138
139    /// Check if a segment type is protected
140    pub fn is_protected(&self, segment_type: ContextSegmentType) -> bool {
141        self.protected_segments
142            .as_ref()
143            .map(|p| p.contains(&segment_type))
144            .unwrap_or(false)
145    }
146
147    /// Check if a segment type should be compacted
148    pub fn should_compact(&self, segment_type: ContextSegmentType) -> bool {
149        if self.is_protected(segment_type) {
150            return false;
151        }
152
153        self.segments_to_compact
154            .as_ref()
155            .map(|s| s.contains(&segment_type))
156            .unwrap_or(true) // If not specified, compact all non-protected
157    }
158}
159
160/// Result of a compaction operation
161///
162/// Matches `compactionResultSchema` in @enact/schemas
163#[derive(Debug, Clone, Serialize, Deserialize)]
164#[serde(rename_all = "camelCase")]
165pub struct CompactionResult {
166    /// Execution ID
167    pub execution_id: ExecutionId,
168
169    /// Strategy used
170    pub strategy: CompactionStrategyType,
171
172    /// Tokens before compaction
173    pub tokens_before: usize,
174
175    /// Tokens after compaction
176    pub tokens_after: usize,
177
178    /// Tokens saved
179    pub tokens_saved: usize,
180
181    /// Compression ratio (tokensAfter / tokensBefore)
182    pub compression_ratio: f64,
183
184    /// Number of segments compacted
185    pub segments_compacted: usize,
186
187    /// Duration in milliseconds
188    pub duration_ms: u64,
189
190    /// Whether compaction was successful
191    pub success: bool,
192
193    /// Error message if failed
194    #[serde(skip_serializing_if = "Option::is_none")]
195    pub error: Option<String>,
196
197    /// Timestamp
198    pub compacted_at: DateTime<Utc>,
199}
200
201impl CompactionResult {
202    /// Create a successful result
203    pub fn success(
204        execution_id: ExecutionId,
205        strategy: CompactionStrategyType,
206        tokens_before: usize,
207        tokens_after: usize,
208        segments_compacted: usize,
209        duration_ms: u64,
210    ) -> Self {
211        let tokens_saved = tokens_before.saturating_sub(tokens_after);
212        let compression_ratio = if tokens_before > 0 {
213            tokens_after as f64 / tokens_before as f64
214        } else {
215            1.0
216        };
217
218        Self {
219            execution_id,
220            strategy,
221            tokens_before,
222            tokens_after,
223            tokens_saved,
224            compression_ratio,
225            segments_compacted,
226            duration_ms,
227            success: true,
228            error: None,
229            compacted_at: Utc::now(),
230        }
231    }
232
233    /// Create a failed result
234    pub fn failure(
235        execution_id: ExecutionId,
236        strategy: CompactionStrategyType,
237        tokens_before: usize,
238        error: String,
239        duration_ms: u64,
240    ) -> Self {
241        Self {
242            execution_id,
243            strategy,
244            tokens_before,
245            tokens_after: tokens_before,
246            tokens_saved: 0,
247            compression_ratio: 1.0,
248            segments_compacted: 0,
249            duration_ms,
250            success: false,
251            error: Some(error),
252            compacted_at: Utc::now(),
253        }
254    }
255}
256
257/// Compactor - applies compaction strategies to segments
258pub struct Compactor {
259    strategy: CompactionStrategy,
260}
261
262impl Compactor {
263    /// Create a new compactor with the given strategy
264    pub fn new(strategy: CompactionStrategy) -> Self {
265        Self { strategy }
266    }
267
268    /// Create a truncation compactor
269    pub fn truncate(target_tokens: usize) -> Self {
270        Self::new(CompactionStrategy::truncate(target_tokens))
271    }
272
273    /// Create a sliding window compactor
274    pub fn sliding_window(target_tokens: usize, window_size: usize) -> Self {
275        Self::new(CompactionStrategy::sliding_window(
276            target_tokens,
277            window_size,
278        ))
279    }
280
281    /// Get the strategy
282    pub fn strategy(&self) -> &CompactionStrategy {
283        &self.strategy
284    }
285
286    /// Compact segments using truncation strategy
287    ///
288    /// Removes oldest segments (lowest priority first) until target is reached.
289    pub fn compact_truncate(
290        &self,
291        segments: &mut Vec<ContextSegment>,
292        current_tokens: usize,
293    ) -> Result<usize, CompactionError> {
294        if current_tokens <= self.strategy.target_tokens {
295            return Ok(0);
296        }
297
298        let tokens_to_remove = current_tokens - self.strategy.target_tokens;
299        let mut removed = 0;
300
301        // Sort by priority (ascending) then by sequence (ascending = oldest first)
302        segments.sort_by(|a, b| {
303            a.priority
304                .cmp(&b.priority)
305                .then(a.sequence.cmp(&b.sequence))
306        });
307
308        // Remove lowest priority, oldest segments first
309        let mut i = 0;
310        while i < segments.len() && removed < tokens_to_remove {
311            let segment = &segments[i];
312
313            // Skip protected segments
314            if !segment.compressible || self.strategy.is_protected(segment.segment_type) {
315                i += 1;
316                continue;
317            }
318
319            // Skip critical priority
320            if segment.priority == ContextPriority::Critical {
321                i += 1;
322                continue;
323            }
324
325            removed += segment.token_count;
326            segments.remove(i);
327        }
328
329        Ok(removed)
330    }
331
332    /// Compact using sliding window strategy
333    ///
334    /// Keeps only the most recent N messages in the history.
335    pub fn compact_sliding_window(
336        &self,
337        segments: &mut Vec<ContextSegment>,
338    ) -> Result<usize, CompactionError> {
339        let window_size = self.strategy.window_size.unwrap_or(10);
340
341        // Find history segments
342        let history_indices: Vec<usize> = segments
343            .iter()
344            .enumerate()
345            .filter(|(_, s)| s.segment_type == ContextSegmentType::History)
346            .map(|(i, _)| i)
347            .collect();
348
349        if history_indices.len() <= window_size {
350            return Ok(0);
351        }
352
353        // Remove oldest history segments (keep window_size most recent)
354        let to_remove = history_indices.len() - window_size;
355        let mut removed_tokens = 0;
356
357        // Remove from oldest first (indices are in ascending order)
358        for &idx in history_indices.iter().take(to_remove).rev() {
359            removed_tokens += segments[idx].token_count;
360            segments.remove(idx);
361        }
362
363        Ok(removed_tokens)
364    }
365}
366
367#[cfg(test)]
368mod tests {
369    use super::*;
370
371    #[test]
372    fn test_truncation_strategy() {
373        let strategy = CompactionStrategy::truncate(1000);
374        assert_eq!(strategy.strategy_type, CompactionStrategyType::Truncate);
375        assert!(strategy.is_protected(ContextSegmentType::System));
376        assert!(!strategy.is_protected(ContextSegmentType::History));
377    }
378
379    #[test]
380    fn test_compaction_result() {
381        let exec_id = ExecutionId::new();
382        let result = CompactionResult::success(
383            exec_id,
384            CompactionStrategyType::Truncate,
385            10000,
386            5000,
387            5,
388            100,
389        );
390
391        assert!(result.success);
392        assert_eq!(result.tokens_saved, 5000);
393        assert!((result.compression_ratio - 0.5).abs() < 0.01);
394    }
395}