1use anyhow::Result;
8
9use crate::providers::{ContentBlock, Message, MessageContent, Provider, Role};
10
11use super::compressor::{compress_messages, estimate_total_tokens};
12use super::config::{CompressionConfig, CircuitBreakerState, ThresholdLevel,
13 TIME_BASED_MC_CLEARED_MESSAGE};
14use super::dependency::DependencyBuilder;
15use super::phase_detector::PhaseDetector;
16use super::scorer::Scorer;
17use super::summarizer::Summarizer;
18use super::tool_compressor::ToolCompressor;
19use super::types::{
20 AiCompressionMode, CompressionThresholds, DependencyGraph,
21 ScoredMessage, CompressionStrategy,
22};
23
24pub struct CompressionPipeline {
26 config: CompressionConfig,
28 scorer: Scorer,
30 tool_compressor: ToolCompressor,
32 circuit_breaker: CircuitBreakerState,
34}
35
36pub struct CompressionOutcome {
38 pub messages: Vec<Message>,
40 pub threshold_level: ThresholdLevel,
42 pub percent_left: u32,
44 pub success: bool,
46 pub error: Option<String>,
48 pub circuit_breaker_tripped: bool,
50}
51
52#[derive(Debug, Clone)]
54pub enum ValidationError {
55 OrphanedToolResult { tool_use_id: String, index: usize },
57 OrphanedToolUse { tool_use_id: String, index: usize },
59 MissingFirstMessage,
61 OrderViolation { expected_role: Role, actual_role: Role, index: usize },
63}
64
65impl CompressionPipeline {
66 pub fn new_rule_only(config: CompressionConfig) -> Self {
68 let thresholds = CompressionThresholds::default();
69 Self {
70 config,
71 scorer: Scorer::new_rule_only(),
72 tool_compressor: ToolCompressor::new_truncate_only(thresholds),
73 circuit_breaker: CircuitBreakerState::new(),
74 }
75 }
76
77 pub fn new_with_ai(
79 config: CompressionConfig,
80 fast_model: Box<dyn Provider>,
81 ) -> Self {
82 let thresholds = CompressionThresholds::default();
83 let summarizer = Summarizer::new(fast_model.clone());
84
85 Self {
86 config,
87 scorer: Scorer::new_with_ai(fast_model),
88 tool_compressor: ToolCompressor::new_with_ai(summarizer, thresholds),
89 circuit_breaker: CircuitBreakerState::new(),
90 }
91 }
92
93 pub fn new_with_full_ai(
95 config: CompressionConfig,
96 fast_model: Box<dyn Provider>,
97 main_model: Box<dyn Provider>,
98 ) -> Self {
99 let thresholds = CompressionThresholds::default();
100 let summarizer = Summarizer::new_with_main(fast_model.clone(), main_model);
101
102 Self {
103 config,
104 scorer: Scorer::new_with_ai(fast_model),
105 tool_compressor: ToolCompressor::new_with_ai(summarizer, thresholds),
106 circuit_breaker: CircuitBreakerState::new(),
107 }
108 }
109
110 pub fn should_compress(
112 &self,
113 token_usage: u32,
114 context_window: u32,
115 ) -> (bool, ThresholdLevel) {
116 if self.circuit_breaker.should_skip() {
118 return (false, ThresholdLevel::Blocking);
119 }
120
121 let (level, _) = CompressionConfig::calculate_threshold_level(token_usage, context_window);
122
123 let should_compress = level != ThresholdLevel::Normal;
124 (should_compress, level)
125 }
126
127 pub fn should_time_based_clear(messages: &[Message]) -> bool {
130 let last_assistant = messages.iter().rev().find(|m| m.role == Role::Assistant);
131
132 if let Some(_msg) = last_assistant {
133 let messages_since = messages.iter().rev().take_while(|m| m.role != Role::Assistant).count();
136 messages_since > 10
138 } else {
139 false
140 }
141 }
142
143 pub fn time_based_microcompact(messages: &[Message]) -> Vec<Message> {
145 messages.iter().map(|msg| {
146 if msg.role != Role::Tool {
147 return msg.clone();
148 }
149
150 match &msg.content {
152 MessageContent::Blocks(blocks) => {
153 let new_blocks: Vec<ContentBlock> = blocks.iter().map(|b| {
154 if let ContentBlock::ToolResult { tool_use_id, content } = b {
155 if content.len() > 500 && content != TIME_BASED_MC_CLEARED_MESSAGE {
157 ContentBlock::ToolResult {
158 tool_use_id: tool_use_id.clone(),
159 content: TIME_BASED_MC_CLEARED_MESSAGE.to_string(),
160 }
161 } else {
162 b.clone()
163 }
164 } else {
165 b.clone()
166 }
167 }).collect();
168 Message {
169 role: msg.role.clone(),
170 content: MessageContent::Blocks(new_blocks),
171 }
172 }
173 _ => msg.clone(),
174 }
175 }).collect()
176 }
177
178 pub fn strip_thinking(messages: &[Message]) -> Vec<Message> {
181 messages.iter().map(|msg| {
182 match &msg.content {
183 MessageContent::Blocks(blocks) => {
184 let new_blocks: Vec<ContentBlock> = blocks.iter()
185 .filter(|b| {
186 !matches!(b, ContentBlock::Thinking { .. })
188 })
189 .cloned()
190 .collect();
191 Message {
192 role: msg.role.clone(),
193 content: MessageContent::Blocks(new_blocks),
194 }
195 }
196 _ => msg.clone(),
197 }
198 }).collect()
199 }
200
201 const COMPACTABLE_TOOLS: &[&str] = &[
204 "bash", "read", "glob", "grep", "ls", "edit", "write",
205 "notebook_edit", "web_fetch", "web_search",
206 ];
207
208 pub fn is_compactable_tool(tool_name: &str) -> bool {
210 Self::COMPACTABLE_TOOLS.contains(&tool_name)
211 }
212
213 pub fn clear_tool_results(messages: &[Message], _tool_names: &[&str]) -> Vec<Message> {
215 messages.iter().map(|msg| {
216 if msg.role != Role::Tool {
217 return msg.clone();
218 }
219
220 match &msg.content {
221 MessageContent::Blocks(blocks) => {
222 let new_blocks: Vec<ContentBlock> = blocks.iter().map(|b| {
223 if let ContentBlock::ToolResult { tool_use_id, content } = b {
224 if content.len() > 500 && content != TIME_BASED_MC_CLEARED_MESSAGE {
227 ContentBlock::ToolResult {
228 tool_use_id: tool_use_id.clone(),
229 content: TIME_BASED_MC_CLEARED_MESSAGE.to_string(),
230 }
231 } else {
232 b.clone()
233 }
234 } else {
235 b.clone()
236 }
237 }).collect();
238 Message {
239 role: msg.role.clone(),
240 content: MessageContent::Blocks(new_blocks),
241 }
242 }
243 _ => msg.clone(),
244 }
245 }).collect()
246 }
247
248 pub fn full_microcompact(messages: &[Message]) -> Vec<Message> {
250 let no_thinking = Self::strip_thinking(messages);
252 Self::time_based_microcompact(&no_thinking)
254 }
255
256 pub fn validate_compression(messages: &[Message], _original_deps: &DependencyGraph) -> Vec<ValidationError> {
262 let mut errors = Vec::new();
263
264 if messages.is_empty() {
266 errors.push(ValidationError::MissingFirstMessage);
267 return errors;
268 }
269
270 let new_deps = DependencyBuilder::build(messages);
272
273 for (idx, msg) in messages.iter().enumerate() {
275 if msg.role == Role::Tool
276 && let MessageContent::Blocks(blocks) = &msg.content {
277 for block in blocks {
278 if let ContentBlock::ToolResult { tool_use_id, .. } = block {
279 let has_tool_use = messages.iter().any(|m| {
281 if let MessageContent::Blocks(bs) = &m.content {
282 bs.iter().any(|b| {
283 if let ContentBlock::ToolUse { id, .. } = b {
284 id == tool_use_id
285 } else {
286 false
287 }
288 })
289 } else {
290 false
291 }
292 });
293
294 if !has_tool_use {
295 errors.push(ValidationError::OrphanedToolResult {
296 tool_use_id: tool_use_id.clone(),
297 index: idx,
298 });
299 }
300 }
301 }
302 }
303 }
304
305 for (idx, msg) in messages.iter().enumerate() {
307 if let MessageContent::Blocks(blocks) = &msg.content {
308 for block in blocks {
309 if let ContentBlock::ToolUse { id, .. } = block {
310 let has_tool_result = messages.iter().any(|m| {
312 if m.role == Role::Tool {
313 if let MessageContent::Blocks(bs) = &m.content {
314 bs.iter().any(|b| {
315 if let ContentBlock::ToolResult { tool_use_id, .. } = b {
316 tool_use_id == id
317 } else {
318 false
319 }
320 })
321 } else {
322 false
323 }
324 } else {
325 false
326 }
327 });
328
329 if !has_tool_result {
330 errors.push(ValidationError::OrphanedToolUse {
331 tool_use_id: id.clone(),
332 index: idx,
333 });
334 }
335 }
336 }
337 }
338 }
339
340 for dep in &new_deps.dependencies {
342 if dep.tool_use_idx >= messages.len() {
343 errors.push(ValidationError::OrphanedToolUse {
344 tool_use_id: dep.tool_name.clone(),
345 index: dep.tool_use_idx,
346 });
347 }
348 if dep.tool_result_idx >= messages.len() {
349 errors.push(ValidationError::OrphanedToolResult {
350 tool_use_id: dep.tool_name.clone(),
351 index: dep.tool_result_idx,
352 });
353 }
354 }
355
356 errors
357 }
358
359 pub fn is_valid_compression(messages: &[Message], original_deps: &DependencyGraph) -> bool {
361 Self::validate_compression(messages, original_deps).is_empty()
362 }
363
364 pub async fn execute(
366 &mut self,
367 messages: &[Message],
368 ai_mode: AiCompressionMode,
369 token_usage: u32,
370 context_window: u32,
371 ) -> Result<CompressionOutcome> {
372 if self.circuit_breaker.should_skip() {
374 return Ok(CompressionOutcome {
375 messages: messages.to_vec(),
376 threshold_level: ThresholdLevel::Blocking,
377 percent_left: 0,
378 success: false,
379 error: Some("Circuit breaker tripped - too many consecutive failures".to_string()),
380 circuit_breaker_tripped: true,
381 });
382 }
383
384 if messages.len() <= self.config.min_preserve_messages {
385 let (level, percent) = CompressionConfig::calculate_threshold_level(token_usage, context_window);
386 return Ok(CompressionOutcome {
387 messages: messages.to_vec(),
388 threshold_level: level,
389 percent_left: percent,
390 success: true,
391 error: None,
392 circuit_breaker_tripped: false,
393 });
394 }
395
396 let pre_processed = if Self::should_time_based_clear(messages) {
398 Self::time_based_microcompact(messages)
399 } else {
400 messages.to_vec()
401 };
402
403 let phase = PhaseDetector::detect(&pre_processed);
405 let weights = phase.default_weights();
406 let deps = DependencyBuilder::build(&pre_processed);
407
408 let scored = self.scorer.score_all(&pre_processed, &weights, &deps, ai_mode).await?;
410
411 let compressed = self.tool_compressor.compress_results(&pre_processed, ai_mode).await?;
413
414 let target_count = calculate_target_count(pre_processed.len(), &self.config);
416 let selected = self.select_messages(scored, &deps, target_count, &compressed);
417
418 let final_messages = self.ensure_dependency_integrity(selected, &deps, &pre_processed);
420
421 self.circuit_breaker.record_success();
423
424 let post_tokens = estimate_total_tokens(&final_messages);
426 let (level, percent) = CompressionConfig::calculate_threshold_level(post_tokens, context_window);
427
428 Ok(CompressionOutcome {
429 messages: final_messages,
430 threshold_level: level,
431 percent_left: percent,
432 success: true,
433 error: None,
434 circuit_breaker_tripped: false,
435 })
436 }
437
438 pub async fn execute_with_circuit_breaker(
440 &mut self,
441 messages: &[Message],
442 ai_mode: AiCompressionMode,
443 token_usage: u32,
444 context_window: u32,
445 ) -> Result<CompressionOutcome> {
446 let result = self.execute(messages, ai_mode, token_usage, context_window).await;
447
448 match result {
449 Ok(res) => Ok(res),
450 Err(e) => {
451 let tripped = self.circuit_breaker.record_failure();
453
454 let (level, percent) = CompressionConfig::calculate_threshold_level(token_usage, context_window);
455
456 Ok(CompressionOutcome {
457 messages: messages.to_vec(),
458 threshold_level: level,
459 percent_left: percent,
460 success: false,
461 error: Some(e.to_string()),
462 circuit_breaker_tripped: tripped,
463 })
464 }
465 }
466 }
467
468 pub fn execute_sync(&self, messages: &[Message]) -> Result<Vec<Message>> {
470 compress_messages(messages, CompressionStrategy::BiasBased, &self.config)
472 }
473
474 fn select_messages(
476 &self,
477 scored: Vec<ScoredMessage>,
478 deps: &DependencyGraph,
479 target_count: usize,
480 compressed_messages: &[Message],
481 ) -> Vec<Message> {
482 let mut sorted = scored;
484 sorted.sort_by(|a, b| b.final_score.partial_cmp(&a.final_score).unwrap());
485
486 let mut preserve_indices: std::collections::HashSet<usize> = std::collections::HashSet::new();
488
489 for sm in sorted.iter().take(target_count) {
491 preserve_indices.insert(sm.index);
492
493 for pair_idx in deps.get_pair_indices(sm.index) {
495 preserve_indices.insert(pair_idx);
496 }
497 }
498
499 let selected: Vec<Message> = preserve_indices
501 .iter()
502 .filter_map(|idx| compressed_messages.get(*idx).cloned())
503 .collect();
504
505 selected
506 }
507
508 fn ensure_dependency_integrity(
510 &self,
511 selected: Vec<Message>,
512 _deps: &DependencyGraph,
513 _original: &[Message],
514 ) -> Vec<Message> {
515 selected
518 }
519
520 pub fn score_only(&self, messages: &[Message]) -> Vec<ScoredMessage> {
522 let phase = PhaseDetector::detect(messages);
523 let weights = phase.default_weights();
524 let deps = DependencyBuilder::build(messages);
525
526 let mut scored: Vec<ScoredMessage> = Vec::new();
528 for (idx, msg) in messages.iter().enumerate() {
529 let base_score = super::scorer::score_by_rules(msg, idx, &weights);
530 scored.push(ScoredMessage::new(idx, msg.clone(), base_score));
531 }
532
533 let bonus = weights.dependency_pair_bonus;
535 for dep in &deps.dependencies {
536 if let Some(sm) = scored.get_mut(dep.tool_use_idx) {
537 sm.with_dependency_bonus(bonus);
538 }
539 if let Some(sm) = scored.get_mut(dep.tool_result_idx) {
540 sm.with_dependency_bonus(bonus);
541 }
542 }
543
544 scored
545 }
546}
547
548fn calculate_target_count(total: usize, config: &CompressionConfig) -> usize {
550 let target = (total as f64 * config.target_ratio) as usize;
551 target.max(config.min_preserve_messages)
552}
553
554pub fn compress_with_pipeline(
556 messages: &[Message],
557 config: &CompressionConfig,
558 ai_mode: AiCompressionMode,
559 fast_model: Option<Box<dyn Provider>>,
560) -> Result<Vec<Message>> {
561 let pipeline = match (ai_mode, fast_model) {
563 (AiCompressionMode::None, _) => CompressionPipeline::new_rule_only(config.clone()),
564 (AiCompressionMode::Light | AiCompressionMode::Deep, Some(model)) => {
565 CompressionPipeline::new_with_ai(config.clone(), model)
566 }
567 _ => CompressionPipeline::new_rule_only(config.clone()),
568 };
569
570 pipeline.execute_sync(messages)
572}
573
574#[cfg(test)]
575mod tests {
576 use super::*;
577 use crate::providers::{MessageContent, Role};
578
579 #[test]
580 fn test_pipeline_new_rule_only() {
581 let config = CompressionConfig::default();
582 let pipeline = CompressionPipeline::new_rule_only(config);
583 let messages = vec![
585 Message {
586 role: Role::User,
587 content: MessageContent::Text("Test".to_string()),
588 },
589 ];
590 let result = pipeline.execute_sync(&messages);
591 assert!(result.is_ok());
592 }
593
594 #[test]
595 fn test_calculate_target_count() {
596 let config = CompressionConfig::default();
597 let total = 100;
598 let target = calculate_target_count(total, &config);
599 assert!(target >= config.min_preserve_messages);
600 assert!(target < total);
601 }
602
603 #[test]
604 fn test_score_only() {
605 let config = CompressionConfig::default();
606 let pipeline = CompressionPipeline::new_rule_only(config);
607
608 let messages = vec![
609 Message {
610 role: Role::User,
611 content: MessageContent::Text("Hello".to_string()),
612 },
613 Message {
614 role: Role::Assistant,
615 content: MessageContent::Text("Hi".to_string()),
616 },
617 ];
618
619 let scored = pipeline.score_only(&messages);
620 assert_eq!(scored.len(), 2);
621 assert!(scored[0].final_score > scored[1].final_score); }
623
624 #[test]
625 fn test_execute_sync_small() {
626 let config = CompressionConfig::default();
627 let pipeline = CompressionPipeline::new_rule_only(config);
628
629 let messages = vec![
630 Message {
631 role: Role::User,
632 content: MessageContent::Text("Hello".to_string()),
633 },
634 ];
635
636 let result = pipeline.execute_sync(&messages).unwrap();
637 assert_eq!(result.len(), 1); }
639
640 #[test]
641 fn test_time_based_microcompact() {
642 let messages = vec![
643 Message {
644 role: Role::Tool,
645 content: MessageContent::Blocks(vec![
646 ContentBlock::ToolResult {
647 tool_use_id: "tool_1".to_string(),
648 content: "This is a very long tool result content that should be cleared...".repeat(20),
649 },
650 ]),
651 },
652 Message {
653 role: Role::Tool,
654 content: MessageContent::Blocks(vec![
655 ContentBlock::ToolResult {
656 tool_use_id: "tool_2".to_string(),
657 content: "Short content".to_string(),
658 },
659 ]),
660 },
661 ];
662
663 let compacted = CompressionPipeline::time_based_microcompact(&messages);
664
665 if let MessageContent::Blocks(blocks) = &compacted[0].content {
667 if let ContentBlock::ToolResult { content, .. } = &blocks[0] {
668 assert_eq!(content, TIME_BASED_MC_CLEARED_MESSAGE);
669 }
670 }
671
672 if let MessageContent::Blocks(blocks) = &compacted[1].content {
674 if let ContentBlock::ToolResult { content, .. } = &blocks[0] {
675 assert_eq!(content, "Short content");
676 }
677 }
678 }
679
680 #[test]
681 fn test_strip_thinking() {
682 let messages = vec![
683 Message {
684 role: Role::Assistant,
685 content: MessageContent::Blocks(vec![
686 ContentBlock::Text { text: "Response".to_string() },
687 ContentBlock::Thinking { thinking: "Long thinking process...".to_string(), signature: None },
688 ]),
689 },
690 ];
691
692 let stripped = CompressionPipeline::strip_thinking(&messages);
693
694 if let MessageContent::Blocks(blocks) = &stripped[0].content {
696 assert_eq!(blocks.len(), 1);
697 assert!(matches!(&blocks[0], ContentBlock::Text { .. }));
698 }
699 }
700
701 #[test]
702 fn test_is_compactable_tool() {
703 assert!(CompressionPipeline::is_compactable_tool("bash"));
704 assert!(CompressionPipeline::is_compactable_tool("read"));
705 assert!(CompressionPipeline::is_compactable_tool("glob"));
706 assert!(!CompressionPipeline::is_compactable_tool("unknown_tool"));
707 }
708
709 #[test]
710 fn test_should_time_based_clear() {
711 let mut many_messages: Vec<Message> = vec![
713 Message {
714 role: Role::Assistant,
715 content: MessageContent::Text("response".to_string()),
716 },
717 ];
718 for i in 0..15 {
720 many_messages.push(Message {
721 role: if i % 2 == 0 { Role::User } else { Role::Tool },
722 content: MessageContent::Text("content".to_string()),
723 });
724 }
725
726 assert!(CompressionPipeline::should_time_based_clear(&many_messages));
727
728 let few_messages = vec![
730 Message {
731 role: Role::Assistant,
732 content: MessageContent::Text("response".to_string()),
733 },
734 Message {
735 role: Role::User,
736 content: MessageContent::Text("follow-up".to_string()),
737 },
738 ];
739
740 assert!(!CompressionPipeline::should_time_based_clear(&few_messages));
741 }
742
743 #[test]
744 fn test_validate_compression_valid() {
745 let messages = vec![
746 Message {
747 role: Role::User,
748 content: MessageContent::Text("Request".to_string()),
749 },
750 Message {
751 role: Role::Assistant,
752 content: MessageContent::Blocks(vec![
753 ContentBlock::ToolUse {
754 id: "tool_1".to_string(),
755 name: "read".to_string(),
756 input: serde_json::json!({"path": "test.txt"}),
757 },
758 ]),
759 },
760 Message {
761 role: Role::Tool,
762 content: MessageContent::Blocks(vec![
763 ContentBlock::ToolResult {
764 tool_use_id: "tool_1".to_string(),
765 content: "File content".to_string(),
766 },
767 ]),
768 },
769 ];
770
771 let deps = DependencyBuilder::build(&messages);
772 let errors = CompressionPipeline::validate_compression(&messages, &deps);
773 assert!(errors.is_empty());
774 }
775
776 #[test]
777 fn test_validate_compression_orphaned_tool_result() {
778 let messages = vec![
779 Message {
780 role: Role::User,
781 content: MessageContent::Text("Request".to_string()),
782 },
783 Message {
784 role: Role::Tool,
785 content: MessageContent::Blocks(vec![
786 ContentBlock::ToolResult {
787 tool_use_id: "tool_missing".to_string(),
788 content: "Orphaned result".to_string(),
789 },
790 ]),
791 },
792 ];
793
794 let deps = DependencyBuilder::build(&messages);
795 let errors = CompressionPipeline::validate_compression(&messages, &deps);
796 assert!(!errors.is_empty());
797 assert!(errors.iter().any(|e| matches!(e, ValidationError::OrphanedToolResult { .. })));
798 }
799
800 #[test]
801 fn test_validate_compression_empty() {
802 let messages: Vec<Message> = vec![];
803 let deps = DependencyBuilder::build(&messages);
804 let errors = CompressionPipeline::validate_compression(&messages, &deps);
805 assert!(!errors.is_empty());
806 assert!(errors.iter().any(|e| matches!(e, ValidationError::MissingFirstMessage)));
807 }
808}