1use crate::high_level::complete;
7use crate::high_level::tokens::estimate as estimate_tokens;
8use crate::{
9 Api, AssistantMessage, ContentBlock, Context, Message, Model, Provider, StreamOptions,
10 TextContent, UserMessage,
11};
12
13fn safe_truncate(s: &str, max_chars: usize) -> String {
15 if s.len() <= max_chars {
16 return s.to_string();
17 }
18 let boundary = s
19 .char_indices()
20 .take_while(|(i, _)| *i <= max_chars)
21 .last()
22 .map(|(i, c)| i + c.len_utf8())
23 .unwrap_or(0);
24 format!("{}...", &s[..boundary])
25}
26
27pub fn generate_branch_summary(messages: &[Message], n: usize) -> String {
32 if messages.is_empty() {
33 return "(empty conversation)".to_string();
34 }
35
36 let last_n: Vec<_> = if n > 0 {
37 messages.iter().rev().take(n).collect()
38 } else {
39 messages.iter().collect()
40 };
41
42 let mut topics = Vec::new();
43 let mut decisions = Vec::new();
44
45 for msg in last_n.iter().rev() {
46 let role = match msg {
47 Message::User(_) => "user",
48 Message::Assistant(_) => "assistant",
49 Message::ToolResult(_) => "tool",
50 };
51 let content = msg.text_content().unwrap_or_default();
52 let preview = safe_truncate(&content, 120);
53
54 if content.contains("created file") || content.contains("edited file") {
56 topics.push("file modifications".to_string());
57 }
58 if content.contains("implemented") || content.contains("added feature") {
59 topics.push("feature implementation".to_string());
60 }
61 if content.contains("decided") || content.contains("chose") || content.contains("agreed") {
62 decisions.push(preview);
63 }
64 if content.contains("search") || content.contains("debug") || content.contains("fix") {
65 topics.push(format!("inquiry/analysis by {}", role));
66 }
67 }
68
69 topics.dedup();
71 decisions.dedup();
72
73 let summary = if topics.is_empty() && decisions.is_empty() {
74 messages
76 .last()
77 .and_then(|m| m.text_content().ok())
78 .map(|c| safe_truncate(&c, 200))
79 .unwrap_or_else(|| "(no content)".to_string())
80 } else {
81 let mut parts = Vec::new();
82 if !topics.is_empty() {
83 parts.push(format!("Topics: {}", topics.join(", ")));
84 }
85 if !decisions.is_empty() {
86 parts.push(format!("Decisions: {}", decisions.join("; ")));
87 }
88 parts.join(" | ")
89 };
90
91 format!("[Branch summary of {} msgs] {}", messages.len(), summary)
92}
93
94use chrono::{DateTime, Utc};
95use serde::{Deserialize, Serialize};
96use std::future::Future;
97use std::pin::Pin;
98use std::sync::Arc;
99use std::time::Duration;
100
101#[derive(Debug, Clone)]
103pub struct CompactionConfig {
104 pub keep_recent: usize,
106 pub max_batch: usize,
108 pub target_ratio: f32,
110 pub summary_max_tokens: usize,
112 pub temperature: f32,
114 pub timeout: Duration,
116 pub custom_instruction: Option<String>,
118}
119
120impl CompactionConfig {
121 pub fn new() -> Self {
123 Self {
124 keep_recent: 4,
125 max_batch: 20,
126 target_ratio: 0.5,
127 summary_max_tokens: 1024,
128 temperature: 0.3,
129 timeout: Duration::from_secs(60),
130 custom_instruction: None,
131 }
132 }
133
134 pub fn with_keep_recent(mut self, count: usize) -> Self {
136 self.keep_recent = count;
137 self
138 }
139
140 pub fn with_max_batch(mut self, count: usize) -> Self {
142 self.max_batch = count;
143 self
144 }
145
146 pub fn with_target_ratio(mut self, ratio: f32) -> Self {
148 self.target_ratio = ratio.clamp(0.1, 0.9);
149 self
150 }
151
152 pub fn with_summary_max_tokens(mut self, tokens: usize) -> Self {
154 self.summary_max_tokens = tokens;
155 self
156 }
157
158 pub fn with_temperature(mut self, temp: f32) -> Self {
160 self.temperature = temp.clamp(0.0, 1.0);
161 self
162 }
163
164 pub fn with_timeout(mut self, timeout: Duration) -> Self {
166 self.timeout = timeout;
167 self
168 }
169
170 pub fn with_custom_instruction(mut self, instruction: impl Into<String>) -> Self {
172 self.custom_instruction = Some(instruction.into());
173 self
174 }
175}
176
177impl Default for CompactionConfig {
178 fn default() -> Self {
179 Self::new()
180 }
181}
182
183#[derive(Debug, Clone, Serialize, Deserialize)]
185pub struct CompactionMetadata {
186 pub original_tokens: usize,
188 pub compacted_tokens: usize,
190 pub messages_compacted: usize,
192 pub messages_kept: usize,
194 pub timestamp: DateTime<Utc>,
196 pub target_ratio: f32,
198 pub actual_ratio: f32,
200 pub success: bool,
202 pub error: Option<String>,
204}
205
206impl CompactionMetadata {
207 pub fn new(
209 original_tokens: usize,
210 compacted_tokens: usize,
211 messages_compacted: usize,
212 messages_kept: usize,
213 target_ratio: f32,
214 ) -> Self {
215 let actual_ratio = if original_tokens > 0 {
216 compacted_tokens as f32 / original_tokens as f32
217 } else {
218 1.0
219 };
220
221 Self {
222 original_tokens,
223 compacted_tokens,
224 messages_compacted,
225 messages_kept,
226 timestamp: Utc::now(),
227 target_ratio,
228 actual_ratio,
229 success: true,
230 error: None,
231 }
232 }
233
234 pub fn failed(
236 original_tokens: usize,
237 messages_compacted: usize,
238 target_ratio: f32,
239 error: impl Into<String>,
240 ) -> Self {
241 Self {
242 original_tokens,
243 compacted_tokens: original_tokens,
244 messages_compacted,
245 messages_kept: 0,
246 timestamp: Utc::now(),
247 target_ratio,
248 actual_ratio: 1.0,
249 success: false,
250 error: Some(error.into()),
251 }
252 }
253
254 pub fn compression_factor(&self) -> f32 {
256 if self.actual_ratio > 0.0 {
257 1.0 - self.actual_ratio
258 } else {
259 0.0
260 }
261 }
262
263 pub fn tokens_saved(&self) -> usize {
265 self.original_tokens.saturating_sub(self.compacted_tokens)
266 }
267}
268
269#[derive(Debug, Clone)]
271pub struct CompactedContext {
272 pub summary: String,
274 pub kept_messages: Vec<Message>,
276 pub compacted_count: usize,
278 pub metadata: CompactionMetadata,
280}
281
282impl CompactedContext {
283 pub fn new(
285 summary: String,
286 kept_messages: Vec<Message>,
287 compacted_count: usize,
288 metadata: CompactionMetadata,
289 ) -> Self {
290 Self {
291 summary,
292 kept_messages,
293 compacted_count,
294 metadata,
295 }
296 }
297
298 pub fn summary(&self) -> &str {
300 &self.summary
301 }
302
303 pub fn kept_count(&self) -> usize {
305 self.kept_messages.len()
306 }
307
308 pub fn compacted_count(&self) -> usize {
310 self.compacted_count
311 }
312
313 pub fn metadata(&self) -> &CompactionMetadata {
315 &self.metadata
316 }
317
318 pub fn is_success(&self) -> bool {
320 self.metadata.success
321 }
322}
323
324#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
326pub enum CompactionStrategy {
327 Disabled,
329 Threshold(f32),
331 EveryNTurns(usize),
333 AbsoluteTokens(usize),
335}
336
337impl CompactionStrategy {
338 pub fn should_compact(
348 &self,
349 context_tokens: usize,
350 context_window: usize,
351 iteration: usize,
352 ) -> bool {
353 match self {
354 CompactionStrategy::Disabled => false,
355 CompactionStrategy::Threshold(threshold) => {
356 if context_window == 0 {
357 return false;
358 }
359 let usage = context_tokens as f32 / context_window as f32;
360 usage >= *threshold
361 }
362 CompactionStrategy::EveryNTurns(n) => iteration > 0 && iteration.is_multiple_of(*n),
363 CompactionStrategy::AbsoluteTokens(max_tokens) => context_tokens >= *max_tokens,
364 }
365 }
366}
367
368impl Default for CompactionStrategy {
369 fn default() -> Self {
370 CompactionStrategy::Threshold(0.8)
371 }
372}
373
374#[derive(Debug, Clone)]
376pub enum CompactionError {
377 LlmError(String),
379 NoMessagesToCompact,
381 TooFewMessages { total: usize, keep_recent: usize },
383 CompactionDisabled,
385 NoContextWindow,
387}
388
389impl std::fmt::Display for CompactionError {
390 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
391 match self {
392 CompactionError::LlmError(msg) => write!(f, "LLM compaction failed: {}", msg),
393 CompactionError::NoMessagesToCompact => write!(f, "No messages to compact"),
394 CompactionError::TooFewMessages { total, keep_recent } => {
395 write!(
396 f,
397 "Not enough messages ({}) to compact (need at least {} for keep_recent)",
398 total,
399 keep_recent + 1
400 )
401 }
402 CompactionError::CompactionDisabled => write!(f, "Compaction is disabled"),
403 CompactionError::NoContextWindow => write!(f, "Context window not configured"),
404 }
405 }
406}
407
408impl std::error::Error for CompactionError {}
409
410pub trait Compactor: Send + Sync {
412 fn compact<'a>(
414 &'a self,
415 messages: &'a [Message],
416 instruction: Option<&'a str>,
417 ) -> Pin<
418 Box<
419 dyn Future<Output = std::result::Result<CompactedContext, CompactionError>> + Send + 'a,
420 >,
421 >;
422
423 fn estimate_tokens(&self, messages: &[Message]) -> usize {
425 messages
426 .iter()
427 .map(|msg| estimate_tokens(&msg.text_content().unwrap_or_default()))
428 .sum()
429 }
430}
431
432pub trait ContextTransformer: Send + Sync {
437 fn transform<'a>(
439 &'a self,
440 context: &'a Context,
441 model: &'a Model,
442 ) -> Pin<Box<dyn Future<Output = Context> + Send + 'a>>;
443}
444
445pub struct NoopContextTransformer;
447
448impl ContextTransformer for NoopContextTransformer {
449 fn transform<'a>(
450 &'a self,
451 context: &'a Context,
452 _model: &'a Model,
453 ) -> Pin<Box<dyn Future<Output = Context> + Send + 'a>> {
454 Box::pin(async move { context.clone() })
455 }
456}
457
458pub struct LlmCompactor {
460 model: Model,
461 _provider: Arc<dyn Provider>,
462 config: CompactionConfig,
463}
464
465impl LlmCompactor {
466 pub fn new(model: Model, provider: Arc<dyn Provider>) -> Self {
468 Self {
469 model,
470 _provider: provider,
471 config: CompactionConfig::new(),
472 }
473 }
474
475 pub fn with_config(
477 model: Model,
478 provider: Arc<dyn Provider>,
479 config: CompactionConfig,
480 ) -> Self {
481 Self {
482 model,
483 _provider: provider,
484 config,
485 }
486 }
487
488 pub fn with_keep_recent(mut self, count: usize) -> Self {
490 self.config.keep_recent = count;
491 self
492 }
493
494 pub fn with_max_batch(mut self, count: usize) -> Self {
496 self.config.max_batch = count;
497 self
498 }
499
500 pub fn with_target_ratio(mut self, ratio: f32) -> Self {
502 self.config.target_ratio = ratio.clamp(0.1, 0.9);
503 self
504 }
505
506 fn build_summarize_prompt(&self, messages: &[Message], instruction: Option<&str>) -> String {
508 let mut prompt = String::new();
509
510 prompt.push_str("Summarize the following conversation concisely. ");
511 prompt.push_str("Capture the key points, decisions, and any ongoing tasks or context.\n\n");
512
513 if let Some(instr) = instruction {
514 prompt.push_str(&format!("Focus areas: {}\n\n", instr));
515 } else if let Some(ref custom_instr) = self.config.custom_instruction {
516 prompt.push_str(&format!("Focus areas: {}\n\n", custom_instr));
517 }
518
519 prompt.push_str("## Conversation to summarize:\n");
520
521 for (i, msg) in messages.iter().enumerate() {
522 let role = match msg {
523 Message::User(_) => "User",
524 Message::Assistant(_) => "Assistant",
525 Message::ToolResult(_) => "Tool",
526 };
527 let content = msg.text_content().unwrap_or_default();
528 let content_preview = safe_truncate(&content, 500);
529 prompt.push_str(&format!("[{} {}]: {}\n", role, i + 1, content_preview));
530 }
531
532 prompt.push_str("\n## Summary:\n");
533 prompt
534 .push_str("Provide a concise summary that captures the essence of this conversation.");
535
536 prompt
537 }
538
539 async fn compact_with_fallback(
541 &self,
542 old_messages: &[Message],
543 recent_messages: &[Message],
544 instruction: Option<&str>,
545 ) -> std::result::Result<CompactedContext, CompactionError> {
546 match self.summarize_with_llm(old_messages, instruction).await {
548 Ok(summary) => {
549 let mut summary_msg =
551 AssistantMessage::new(Api::AnthropicMessages, "compactor", &self.model.id);
552 summary_msg.content = vec![ContentBlock::Text(TextContent::new(format!(
553 "[Previous conversation summarized: {}]",
554 summary
555 )))];
556
557 let mut kept = vec![Message::Assistant(summary_msg)];
559 kept.extend(recent_messages.iter().cloned());
560
561 let original_tokens = self.estimate_tokens(old_messages);
562 let compacted_tokens = self.estimate_tokens(&kept);
563 let kept_len = kept.len();
564
565 Ok(CompactedContext::new(
566 summary,
567 kept,
568 old_messages.len(),
569 CompactionMetadata::new(
570 original_tokens,
571 compacted_tokens,
572 old_messages.len(),
573 kept_len,
574 self.config.target_ratio,
575 ),
576 ))
577 }
578 Err(llm_err) => {
579 self.compact_fallback(old_messages, recent_messages)
581 .await
582 .map_err(|_| CompactionError::LlmError(llm_err.to_string()))
583 }
584 }
585 }
586
587 async fn summarize_with_llm(
589 &self,
590 messages: &[Message],
591 instruction: Option<&str>,
592 ) -> std::result::Result<String, CompactionError> {
593 let prompt = self.build_summarize_prompt(messages, instruction);
594
595 let mut context = Context::new();
596 context.set_system_prompt(
597 "You are a helpful assistant that summarizes conversations concisely.",
598 );
599 context.add_message(Message::User(UserMessage::new(prompt)));
600
601 let options = StreamOptions {
602 temperature: Some(self.config.temperature as f64),
603 max_tokens: Some(self.config.summary_max_tokens),
604 ..Default::default()
605 };
606
607 let summary_message = complete(&self.model, &context, Some(options))
608 .await
609 .map_err(|e| CompactionError::LlmError(e.to_string()))?;
610
611 Ok(summary_message.text_content())
612 }
613
614 async fn compact_fallback(
616 &self,
617 old_messages: &[Message],
618 recent_messages: &[Message],
619 ) -> std::result::Result<CompactedContext, CompactionError> {
620 let mut summary_parts = Vec::new();
622
623 if old_messages.len() > 2 {
624 if let Some(first) = old_messages.first() {
626 let content = first.text_content().unwrap_or_default();
627 let preview = safe_truncate(&content, 200);
628 summary_parts.push(format!("Started discussing: {}", preview));
629 }
630
631 if let Some(last) = old_messages.last() {
633 let content = last.text_content().unwrap_or_default();
634 let preview = safe_truncate(&content, 200);
635 summary_parts.push(format!("Ended with: {}", preview));
636 }
637
638 summary_parts.push(format!(
639 "({} messages omitted)",
640 old_messages.len().saturating_sub(2)
641 ));
642 } else if !old_messages.is_empty() {
643 if let Some(msg) = old_messages.first() {
645 let content = msg.text_content().unwrap_or_default();
646 summary_parts.push(format!("Conversation started: {}", content));
647 }
648 }
649
650 let summary = summary_parts.join(" ");
651
652 let mut summary_msg =
653 AssistantMessage::new(Api::AnthropicMessages, "compactor", &self.model.id);
654 summary_msg.content = vec![ContentBlock::Text(TextContent::new(format!(
655 "[Previous conversation summary: {}]",
656 summary
657 )))];
658
659 let mut kept = vec![Message::Assistant(summary_msg)];
660 kept.extend(recent_messages.iter().cloned());
661
662 let original_tokens = self.estimate_tokens(old_messages);
663 let compacted_tokens = self.estimate_tokens(&kept);
664 let kept_len = kept.len();
665
666 Ok(CompactedContext::new(
667 summary,
668 kept,
669 old_messages.len(),
670 CompactionMetadata::new(
671 original_tokens,
672 compacted_tokens,
673 old_messages.len(),
674 kept_len,
675 self.config.target_ratio,
676 ),
677 ))
678 }
679}
680
681impl Compactor for LlmCompactor {
682 fn compact<'a>(
683 &'a self,
684 messages: &'a [Message],
685 instruction: Option<&'a str>,
686 ) -> Pin<
687 Box<
688 dyn Future<Output = std::result::Result<CompactedContext, CompactionError>> + Send + 'a,
689 >,
690 > {
691 Box::pin(async move {
692 if messages.is_empty() {
694 return Err(CompactionError::NoMessagesToCompact);
695 }
696
697 if messages.len() <= self.config.keep_recent {
698 let original_tokens = self.estimate_tokens(messages);
700 return Ok(CompactedContext::new(
701 String::new(),
702 messages.to_vec(),
703 0,
704 CompactionMetadata::new(
705 original_tokens,
706 original_tokens,
707 0,
708 messages.len(),
709 self.config.target_ratio,
710 ),
711 ));
712 }
713
714 let keep_count = self.config.keep_recent.min(messages.len());
716 let old_messages: Vec<Message> = messages[..messages.len() - keep_count].to_vec();
717 let recent_messages: Vec<Message> = messages[messages.len() - keep_count..].to_vec();
718
719 if old_messages.is_empty() {
720 return Err(CompactionError::NoMessagesToCompact);
721 }
722
723 self.compact_with_fallback(&old_messages, &recent_messages, instruction)
725 .await
726 })
727 }
728}
729
730impl LlmCompactor {
732 pub async fn summarize_branch(
737 &self,
738 messages: &[Message],
739 branch_name: &str,
740 ) -> std::result::Result<String, CompactionError> {
741 if messages.is_empty() {
742 return Ok(format!("Branch '{}' is empty", branch_name));
743 }
744
745 let mut prompt = String::new();
746 prompt.push_str(&format!(
747 "Summarize the conversation branch '{}' concisely. ",
748 branch_name
749 ));
750 prompt.push_str("Focus on: what was discussed, decisions made, and current state.\n\n");
751
752 prompt.push_str("## Branch messages:\n");
753 for (i, msg) in messages.iter().enumerate() {
754 let role = match msg {
755 Message::User(_) => "User",
756 Message::Assistant(_) => "Assistant",
757 Message::ToolResult(_) => "Tool",
758 };
759 let content = msg.text_content().unwrap_or_default();
760 let content_preview = safe_truncate(&content, 300);
761 prompt.push_str(&format!("[{} {}]: {}\n", role, i + 1, content_preview));
762 }
763
764 prompt.push_str("\n## Summary (be concise):\n");
765
766 let mut context = Context::new();
768 context.set_system_prompt(
769 "You are a helpful assistant that summarizes conversation branches. ",
770 );
771 context.add_message(Message::User(UserMessage::new(prompt)));
772
773 let options = StreamOptions {
774 temperature: Some(0.3),
775 max_tokens: Some(512),
776 ..Default::default()
777 };
778
779 let summary_message = complete(&self.model, &context, Some(options))
780 .await
781 .map_err(|e| CompactionError::LlmError(e.to_string()))?;
782
783 Ok(summary_message.text_content())
784 }
785}
786
787pub struct CompactionManager {
789 strategy: CompactionStrategy,
790 compactor: Option<Arc<dyn Compactor>>,
791 context_window: usize,
792 config: CompactionConfig,
793}
794
795impl CompactionManager {
796 pub fn new(strategy: CompactionStrategy, context_window: usize) -> Self {
798 Self {
799 strategy,
800 compactor: None,
801 context_window,
802 config: CompactionConfig::new(),
803 }
804 }
805
806 pub fn with_config(
808 strategy: CompactionStrategy,
809 context_window: usize,
810 config: CompactionConfig,
811 ) -> Self {
812 Self {
813 strategy,
814 compactor: None,
815 context_window,
816 config,
817 }
818 }
819
820 pub fn with_compactor<C: Compactor + 'static>(mut self, compactor: Arc<C>) -> Self {
822 self.compactor = Some(compactor);
823 self
824 }
825
826 pub fn set_compactor(&mut self, compactor: Arc<dyn Compactor>) {
828 self.compactor = Some(compactor);
829 }
830
831 pub fn should_compact(&self, context_tokens: usize, iteration: usize) -> bool {
833 self.strategy
834 .should_compact(context_tokens, self.context_window, iteration)
835 }
836
837 pub fn strategy(&self) -> &CompactionStrategy {
839 &self.strategy
840 }
841
842 pub fn config(&self) -> &CompactionConfig {
844 &self.config
845 }
846
847 pub fn set_config(&mut self, config: CompactionConfig) {
849 self.config = config;
850 }
851
852 pub async fn compact_if_needed(
854 &self,
855 messages: &[Message],
856 instruction: Option<&str>,
857 context_tokens: usize,
858 iteration: usize,
859 ) -> std::result::Result<Option<CompactedContext>, CompactionError> {
860 if !self.should_compact(context_tokens, iteration) {
861 return Ok(None);
862 }
863
864 let compactor = match &self.compactor {
865 Some(c) => c,
866 None => return Err(CompactionError::CompactionDisabled),
867 };
868
869 let result = compactor.compact(messages, instruction).await?;
870 Ok(Some(result))
871 }
872
873 pub async fn compact_now(
875 &self,
876 messages: &[Message],
877 instruction: Option<&str>,
878 ) -> std::result::Result<CompactedContext, CompactionError> {
879 let compactor = match &self.compactor {
880 Some(c) => c,
881 None => return Err(CompactionError::CompactionDisabled),
882 };
883
884 compactor.compact(messages, instruction).await
885 }
886
887 pub fn estimate_tokens(&self, messages: &[Message]) -> usize {
889 messages
890 .iter()
891 .map(|msg| estimate_tokens(&msg.text_content().unwrap_or_default()))
892 .sum()
893 }
894}
895
896impl Default for CompactionManager {
897 fn default() -> Self {
898 Self::new(CompactionStrategy::default(), 128_000)
899 }
900}
901
902#[cfg(test)]
907mod tests {
908 use super::*;
909
910 fn make_user_message(content: &str) -> Message {
912 Message::user(content)
913 }
914
915 fn make_assistant_message(content: &str) -> Message {
917 Message::Assistant({
918 let mut msg = AssistantMessage::new(Api::AnthropicMessages, "test", "test-model");
919 msg.content = vec![ContentBlock::Text(TextContent::new(content))];
920 msg
921 })
922 }
923
924 fn make_test_model() -> Model {
926 Model::new(
927 "test-model",
928 "Test Model",
929 Api::AnthropicMessages,
930 "test",
931 "https://test.example.com",
932 )
933 }
934
935 #[test]
936 fn test_compaction_config_defaults() {
937 let config = CompactionConfig::new();
938 assert_eq!(config.keep_recent, 4);
939 assert_eq!(config.max_batch, 20);
940 assert!((config.target_ratio - 0.5).abs() < 0.001);
941 assert_eq!(config.summary_max_tokens, 1024);
942 assert!((config.temperature - 0.3).abs() < 0.001);
943 }
944
945 #[test]
946 fn test_compaction_config_builder_pattern() {
947 let config = CompactionConfig::new()
948 .with_keep_recent(10)
949 .with_max_batch(30)
950 .with_target_ratio(0.3)
951 .with_temperature(0.5);
952
953 assert_eq!(config.keep_recent, 10);
954 assert_eq!(config.max_batch, 30);
955 assert!((config.target_ratio - 0.3).abs() < 0.001);
956 assert!((config.temperature - 0.5).abs() < 0.001);
957 }
958
959 #[test]
960 fn test_compaction_config_ratio_clamping() {
961 let config = CompactionConfig::new().with_target_ratio(1.5);
963 assert!((config.target_ratio - 0.9).abs() < 0.001);
964
965 let config = CompactionConfig::new().with_target_ratio(-0.5);
967 assert!((config.target_ratio - 0.1).abs() < 0.001);
968 }
969
970 #[test]
971 fn test_compaction_metadata_success() {
972 let metadata = CompactionMetadata::new(
973 1000, 500, 10, 5, 0.5, );
979
980 assert!(metadata.success);
981 assert_eq!(metadata.original_tokens, 1000);
982 assert_eq!(metadata.compacted_tokens, 500);
983 assert_eq!(metadata.messages_compacted, 10);
984 assert_eq!(metadata.messages_kept, 5);
985 assert!((metadata.actual_ratio - 0.5).abs() < 0.001);
986 assert!((metadata.compression_factor() - 0.5).abs() < 0.001);
987 assert_eq!(metadata.tokens_saved(), 500);
988 assert!(metadata.error.is_none());
989 }
990
991 #[test]
992 fn test_compaction_metadata_failure() {
993 let metadata = CompactionError::LlmError("test error".to_string());
994
995 assert!(metadata.to_string().contains("test error"));
997 }
998
999 #[test]
1000 fn test_compaction_metadata_compression_factor() {
1001 let metadata = CompactionMetadata::new(0, 0, 0, 0, 0.5);
1003 assert!((metadata.actual_ratio - 1.0).abs() < 0.001);
1004 assert!((metadata.compression_factor() - 0.0).abs() < 0.001);
1005
1006 let metadata = CompactionMetadata::new(1000, 100, 10, 5, 0.5);
1008 assert!((metadata.compression_factor() - 0.9).abs() < 0.001);
1009 }
1010
1011 #[test]
1012 fn test_compaction_metadata_tokens_saved() {
1013 let metadata = CompactionMetadata::new(1000, 400, 10, 5, 0.5);
1015 assert_eq!(metadata.tokens_saved(), 600);
1016
1017 let metadata = CompactionMetadata::new(1000, 1000, 0, 0, 0.5);
1019 assert_eq!(metadata.tokens_saved(), 0);
1020
1021 let metadata = CompactionMetadata::new(500, 600, 5, 3, 0.5);
1023 assert_eq!(metadata.tokens_saved(), 0); }
1025
1026 #[test]
1027 fn test_compaction_strategy_disabled() {
1028 let strategy = CompactionStrategy::Disabled;
1029 assert!(!strategy.should_compact(100_000, 128_000, 5));
1030 assert!(!strategy.should_compact(120_000, 128_000, 10));
1031 assert!(!strategy.should_compact(0, 128_000, 1));
1032 }
1033
1034 #[test]
1035 fn test_compaction_strategy_threshold() {
1036 let strategy = CompactionStrategy::Threshold(0.8);
1037
1038 assert!(!strategy.should_compact(100_000, 128_000, 1));
1040
1041 assert!(strategy.should_compact(102_400, 128_000, 1));
1043
1044 assert!(strategy.should_compact(120_000, 128_000, 1));
1046
1047 assert!(!strategy.should_compact(100_000, 0, 1));
1049 }
1050
1051 #[test]
1052 fn test_compaction_strategy_every_n_turns() {
1053 let strategy = CompactionStrategy::EveryNTurns(5);
1054
1055 assert!(!strategy.should_compact(0, 128_000, 0));
1057 assert!(!strategy.should_compact(0, 128_000, 3));
1058 assert!(!strategy.should_compact(0, 128_000, 4));
1059
1060 assert!(strategy.should_compact(0, 128_000, 5));
1062 assert!(strategy.should_compact(0, 128_000, 10));
1063 assert!(strategy.should_compact(0, 128_000, 15));
1064
1065 assert!(!strategy.should_compact(0, 128_000, 6));
1067 assert!(!strategy.should_compact(0, 128_000, 9));
1068 }
1069
1070 #[test]
1071 fn test_compaction_strategy_absolute_tokens() {
1072 let strategy = CompactionStrategy::AbsoluteTokens(100_000);
1073
1074 assert!(!strategy.should_compact(50_000, 128_000, 0));
1076 assert!(!strategy.should_compact(99_999, 128_000, 0));
1077
1078 assert!(strategy.should_compact(100_000, 128_000, 0));
1080
1081 assert!(strategy.should_compact(150_000, 128_000, 0));
1083 }
1084
1085 #[test]
1086 fn test_compacted_context_basic() {
1087 let metadata = CompactionMetadata::new(1000, 500, 10, 5, 0.5);
1088 let ctx = CompactedContext::new(
1089 "Test summary".to_string(),
1090 vec![make_user_message("test")],
1091 10,
1092 metadata,
1093 );
1094
1095 assert_eq!(ctx.summary(), "Test summary");
1096 assert_eq!(ctx.kept_count(), 1);
1097 assert_eq!(ctx.compacted_count(), 10);
1098 assert!(ctx.is_success());
1099 assert_eq!(ctx.metadata().tokens_saved(), 500);
1100 }
1101
1102 #[test]
1103 fn test_compacted_context_with_empty_summary() {
1104 let metadata = CompactionMetadata::new(100, 100, 0, 2, 0.5);
1105 let ctx = CompactedContext::new(
1106 String::new(), vec![make_user_message("test1"), make_user_message("test2")],
1108 0,
1109 metadata,
1110 );
1111
1112 assert_eq!(ctx.summary(), "");
1113 assert_eq!(ctx.kept_count(), 2);
1114 assert_eq!(ctx.compacted_count(), 0);
1115 }
1116
1117 #[test]
1118 fn test_llm_compactor_config_builder() {
1119 use crate::providers::OpenAiProvider;
1121 let provider = OpenAiProvider::new();
1122 let model = make_test_model();
1123 let compactor = LlmCompactor::new(model, Arc::new(provider))
1124 .with_keep_recent(6)
1125 .with_max_batch(25)
1126 .with_target_ratio(0.6);
1127
1128 assert!(compactor.config.keep_recent >= 4);
1129 assert!(compactor.config.max_batch >= 20);
1130 }
1131
1132 #[test]
1133 fn test_compaction_error_display() {
1134 let err = CompactionError::NoMessagesToCompact;
1135 assert_eq!(err.to_string(), "No messages to compact");
1136
1137 let err = CompactionError::TooFewMessages {
1138 total: 3,
1139 keep_recent: 5,
1140 };
1141 assert!(err.to_string().contains("3"));
1142 assert!(err.to_string().contains("6"));
1144
1145 let err = CompactionError::CompactionDisabled;
1146 assert_eq!(err.to_string(), "Compaction is disabled");
1147
1148 let err = CompactionError::NoContextWindow;
1149 assert_eq!(err.to_string(), "Context window not configured");
1150
1151 let err = CompactionError::LlmError("API timeout".to_string());
1152 assert!(err.to_string().contains("API timeout"));
1153 }
1154
1155 #[test]
1156 fn test_compaction_manager_default() {
1157 let manager = CompactionManager::default();
1158 assert!(matches!(
1159 manager.strategy(),
1160 CompactionStrategy::Threshold(_)
1161 ));
1162 assert_eq!(manager.config().keep_recent, 4);
1163 }
1164
1165 #[test]
1166 fn test_compaction_manager_with_custom_strategy() {
1167 let strategy = CompactionStrategy::AbsoluteTokens(50_000);
1168 let manager = CompactionManager::new(strategy, 200_000);
1169
1170 assert!(!manager.should_compact(30_000, 0));
1172
1173 assert!(manager.should_compact(60_000, 0));
1175 }
1176
1177 #[test]
1178 fn test_compaction_manager_with_config() {
1179 let config = CompactionConfig::new()
1180 .with_keep_recent(8)
1181 .with_target_ratio(0.4);
1182
1183 let manager =
1184 CompactionManager::with_config(CompactionStrategy::default(), 128_000, config);
1185
1186 assert_eq!(manager.config().keep_recent, 8);
1187 assert!((manager.config().target_ratio - 0.4).abs() < 0.001);
1188 }
1189
1190 #[test]
1191 fn test_compaction_manager_should_compact_integration() {
1192 let manager = CompactionManager::new(CompactionStrategy::Threshold(0.75), 100_000);
1193
1194 assert!(!manager.should_compact(70_000, 0));
1196
1197 assert!(manager.should_compact(75_000, 0));
1199
1200 assert!(manager.should_compact(80_000, 0));
1202 assert!(manager.should_compact(100_000, 0));
1203 }
1204
1205 #[test]
1206 fn test_compaction_manager_no_compactor_set() {
1207 let manager = CompactionManager::new(CompactionStrategy::EveryNTurns(5), 128_000);
1208
1209 assert!(manager.should_compact(0, 5)); }
1213
1214 #[test]
1215 fn test_token_estimation_helper() {
1216 use crate::providers::OpenAiProvider;
1217 let provider = OpenAiProvider::new();
1218 let model = make_test_model();
1219 let compactor = LlmCompactor::new(model, Arc::new(provider));
1220
1221 let messages = vec![
1222 make_user_message("Hello world, this is a test message."),
1223 make_assistant_message("This is a response with some content."),
1224 ];
1225
1226 let tokens = compactor.estimate_tokens(&messages);
1227 assert!(tokens > 0, "Should estimate tokens for messages");
1228 }
1229
1230 #[test]
1231 fn test_compaction_config_custom_instruction() {
1232 let config = CompactionConfig::new()
1233 .with_custom_instruction("Focus on code changes and technical decisions");
1234
1235 assert!(config.custom_instruction.is_some());
1236 assert!(config.custom_instruction.unwrap().contains("code changes"));
1237 }
1238
1239 #[test]
1240 fn test_compaction_metadata_timestamp_is_set() {
1241 let metadata = CompactionMetadata::new(1000, 500, 10, 5, 0.5);
1242 assert!(metadata.timestamp <= Utc::now());
1243 }
1244
1245 #[test]
1246 fn test_compaction_ratio_achievement() {
1247 let metadata = CompactionMetadata::new(1000, 500, 10, 5, 0.5);
1249 assert!((metadata.actual_ratio - 0.5).abs() < 0.001);
1250
1251 let metadata = CompactionMetadata::new(1000, 300, 10, 5, 0.5);
1253 assert!((metadata.actual_ratio - 0.3).abs() < 0.001);
1254 assert!(metadata.compression_factor() > 0.5);
1255
1256 let metadata = CompactionMetadata::new(1000, 700, 10, 5, 0.5);
1258 assert!((metadata.actual_ratio - 0.7).abs() < 0.001);
1259 assert!(metadata.compression_factor() < 0.5);
1260 }
1261
1262 #[test]
1263 fn test_compaction_manager_config_updates() {
1264 let mut manager = CompactionManager::default();
1265
1266 let new_config = CompactionConfig::new()
1267 .with_keep_recent(12)
1268 .with_target_ratio(0.3);
1269
1270 manager.set_config(new_config);
1271
1272 assert_eq!(manager.config().keep_recent, 12);
1273 assert!((manager.config().target_ratio - 0.3).abs() < 0.001);
1274 }
1275
1276 #[test]
1277 fn test_llm_compactor_has_summarize_branch() {
1278 use crate::providers::OpenAiProvider;
1280 let provider = OpenAiProvider::new();
1281 let model = make_test_model();
1282 let compactor = LlmCompactor::new(model, Arc::new(provider));
1283
1284 let messages = vec![
1286 make_user_message("Test message 1"),
1287 make_assistant_message("Test response 1"),
1288 make_user_message("Test message 2"),
1289 ];
1290
1291 let branch_name = "test-branch";
1294 let _future = compactor.summarize_branch(&messages, branch_name);
1296 }
1297
1298 #[test]
1299 fn test_summarize_branch_returns_error_on_llm_failure() {
1300 use crate::providers::OpenAiProvider;
1302 let provider = OpenAiProvider::new();
1303 let model = make_test_model();
1304 let compactor = LlmCompactor::new(model, Arc::new(provider));
1305
1306 let messages: Vec<Message> = vec![];
1308
1309 let _future = compactor.summarize_branch(&messages, "empty-branch");
1312 }
1313}