Skip to main content

agent_air_runtime/controller/session/
compactor.rs

1use std::collections::{HashMap, HashSet};
2use std::future::Future;
3use std::pin::Pin;
4use std::time::Duration;
5
6use thiserror::Error;
7
8use crate::client::LLMClient;
9use crate::client::models::{Message as LLMMessage, MessageOptions};
10
11use crate::controller::types::{ContentBlock, Message, TextBlock, TurnId, UserMessage};
12
13/// Error type for compactor configuration.
14#[derive(Error, Debug)]
15pub enum CompactorConfigError {
16    /// Invalid threshold value.
17    #[error("Invalid threshold {0}: must be between 0.0 and 1.0 (exclusive)")]
18    InvalidThreshold(f64),
19
20    /// Invalid keep_recent_turns value.
21    #[error("Invalid keep_recent_turns {0}: must be at least 1")]
22    InvalidKeepRecentTurns(usize),
23}
24
25/// Defines how tool results are handled during compaction.
26#[derive(Debug, Clone, Copy, PartialEq, Eq)]
27pub enum ToolCompaction {
28    /// Replace tool results with summaries (uses CompactSummary if available).
29    Summarize,
30    /// Replace tool result content with a redaction notice.
31    Redact,
32}
33
34impl std::fmt::Display for ToolCompaction {
35    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
36        match self {
37            ToolCompaction::Summarize => write!(f, "summarize"),
38            ToolCompaction::Redact => write!(f, "redact"),
39        }
40    }
41}
42
43/// Result of a compaction operation.
44#[derive(Debug, Clone, Default)]
45pub struct CompactionResult {
46    /// Number of tool results that were summarized.
47    pub tool_results_summarized: usize,
48    /// Number of tool results that were redacted.
49    pub tool_results_redacted: usize,
50    /// Number of turns that were compacted.
51    pub turns_compacted: usize,
52}
53
54impl CompactionResult {
55    /// Total number of tool results compacted.
56    pub fn total_compacted(&self) -> usize {
57        self.tool_results_summarized + self.tool_results_redacted
58    }
59}
60
61/// Error type for async compaction operations.
62#[derive(Debug)]
63pub enum CompactionError {
64    /// LLM call failed during summarization.
65    LLMError(String),
66    /// Timeout during summarization.
67    Timeout,
68    /// Configuration error.
69    ConfigError(String),
70}
71
72impl std::fmt::Display for CompactionError {
73    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
74        match self {
75            CompactionError::LLMError(msg) => write!(f, "LLM error: {}", msg),
76            CompactionError::Timeout => write!(f, "compaction timed out"),
77            CompactionError::ConfigError(msg) => write!(f, "config error: {}", msg),
78        }
79    }
80}
81
82impl std::error::Error for CompactionError {}
83
84/// Trait for conversation compaction strategies.
85/// Implementations decide when and how to compact conversation history
86/// to reduce token usage while preserving important context.
87pub trait Compactor: Send + Sync {
88    /// Returns true if compaction should occur before this LLM call.
89    ///
90    /// # Arguments
91    /// * `context_used` - Current number of tokens in the conversation context
92    /// * `context_limit` - Maximum context size for the model
93    fn should_compact(&self, context_used: i64, context_limit: i32) -> bool;
94
95    /// Performs synchronous compaction on the conversation.
96    /// Returns the compaction statistics.
97    ///
98    /// # Arguments
99    /// * `conversation` - The conversation messages to compact
100    /// * `compact_summaries` - Map of tool_use_id to pre-computed compact summaries
101    fn compact(
102        &self,
103        conversation: &mut Vec<Message>,
104        compact_summaries: &HashMap<String, String>,
105    ) -> CompactionResult;
106
107    /// Returns true if this compactor requires async compaction.
108    /// When true, the caller should use `AsyncCompactor::compact_async()` instead of `compact()`.
109    fn is_async(&self) -> bool {
110        false
111    }
112}
113
114/// A boxed future for async compaction results.
115pub type CompactAsyncFuture<'a> = Pin<
116    Box<dyn Future<Output = Result<(Vec<Message>, CompactionResult), CompactionError>> + Send + 'a>,
117>;
118
119/// Trait for async compaction strategies that require LLM calls.
120/// Implementors must also implement `Compactor` with `is_async() -> true`.
121pub trait AsyncCompactor: Compactor {
122    /// Performs async compaction that may involve LLM calls.
123    /// Returns the new conversation and compaction statistics.
124    ///
125    /// Unlike sync `compact()` which modifies in place, this returns a new conversation
126    /// because async compaction may replace multiple messages with a single summary.
127    fn compact_async<'a>(
128        &'a self,
129        conversation: Vec<Message>,
130        compact_summaries: &'a HashMap<String, String>,
131    ) -> CompactAsyncFuture<'a>;
132}
133
134/// Compacts when context usage exceeds a threshold.
135/// It identifies turns to keep (most recent N) and applies the configured
136/// compaction strategy to older tool results.
137pub struct ThresholdCompactor {
138    /// Context utilization ratio (0.0-1.0) that triggers compaction.
139    /// For example, 0.75 means compact when 75% of context is used.
140    threshold: f64,
141
142    /// Number of recent turns to preserve during compaction.
143    /// A turn includes user message, assistant response, tool calls, and tool results.
144    keep_recent_turns: usize,
145
146    /// Strategy for handling old tool results.
147    tool_compaction: ToolCompaction,
148}
149
150impl ThresholdCompactor {
151    /// Creates a new threshold compactor.
152    ///
153    /// # Arguments
154    /// * `threshold` - Context utilization ratio (0.0-1.0, exclusive)
155    /// * `keep_recent_turns` - Number of recent turns to preserve
156    /// * `tool_compaction` - Strategy for handling old tool results
157    ///
158    /// # Returns
159    /// Error if parameters are invalid.
160    pub fn new(
161        threshold: f64,
162        keep_recent_turns: usize,
163        tool_compaction: ToolCompaction,
164    ) -> Result<Self, CompactorConfigError> {
165        if threshold <= 0.0 || threshold >= 1.0 {
166            return Err(CompactorConfigError::InvalidThreshold(threshold));
167        }
168
169        Ok(Self {
170            threshold,
171            keep_recent_turns,
172            tool_compaction,
173        })
174    }
175
176    /// Extract unique turn IDs in order of first appearance.
177    fn unique_turn_ids(&self, conversation: &[Message]) -> Vec<TurnId> {
178        let mut seen = HashSet::new();
179        let mut ids = Vec::new();
180
181        for msg in conversation {
182            let turn_id = msg.turn_id();
183            if seen.insert(turn_id.clone()) {
184                ids.push(turn_id.clone());
185            }
186        }
187
188        ids
189    }
190
191    /// Compact tool results in a single message.
192    fn compact_message(
193        &self,
194        msg: &mut Message,
195        compact_summaries: &HashMap<String, String>,
196    ) -> (usize, usize) {
197        let mut summarized = 0;
198        let mut redacted = 0;
199
200        for block in msg.content_mut() {
201            if let ContentBlock::ToolResult(tool_result) = block {
202                match self.tool_compaction {
203                    ToolCompaction::Summarize => {
204                        // Use pre-computed summary if available
205                        if let Some(summary) = compact_summaries.get(&tool_result.tool_use_id) {
206                            tool_result.content = summary.clone();
207                            summarized += 1;
208                            tracing::debug!(
209                                tool_use_id = %tool_result.tool_use_id,
210                                "Tool result summarized"
211                            );
212                        }
213                        // If no summary available, keep original content
214                    }
215                    ToolCompaction::Redact => {
216                        tool_result.content =
217                            "[Tool result redacted during compaction]".to_string();
218                        redacted += 1;
219                        tracing::debug!(
220                            tool_use_id = %tool_result.tool_use_id,
221                            "Tool result redacted"
222                        );
223                    }
224                }
225            }
226        }
227
228        (summarized, redacted)
229    }
230}
231
232impl Compactor for ThresholdCompactor {
233    fn should_compact(&self, context_used: i64, context_limit: i32) -> bool {
234        if context_limit == 0 {
235            return false;
236        }
237
238        let utilization = context_used as f64 / context_limit as f64;
239        let should_compact = utilization > self.threshold;
240
241        if should_compact {
242            tracing::info!(
243                utilization = utilization,
244                threshold = self.threshold,
245                context_used,
246                context_limit,
247                "Compaction triggered - context utilization exceeded threshold"
248            );
249        }
250
251        should_compact
252    }
253
254    fn compact(
255        &self,
256        conversation: &mut Vec<Message>,
257        compact_summaries: &HashMap<String, String>,
258    ) -> CompactionResult {
259        if conversation.is_empty() {
260            return CompactionResult::default();
261        }
262
263        // Find unique turn IDs in order
264        let turn_ids = self.unique_turn_ids(conversation);
265
266        // If fewer turns than keep_recent_turns, nothing to compact
267        if turn_ids.len() <= self.keep_recent_turns {
268            tracing::debug!(
269                total_turns = turn_ids.len(),
270                keep_recent = self.keep_recent_turns,
271                "Skipping compaction - not enough turns"
272            );
273            return CompactionResult::default();
274        }
275
276        // Identify turns to keep (most recent N)
277        let start_idx = turn_ids.len() - self.keep_recent_turns;
278        let turns_to_keep: HashSet<_> = turn_ids[start_idx..].iter().cloned().collect();
279        let turns_compacted = start_idx;
280
281        tracing::info!(
282            total_turns = turn_ids.len(),
283            keep_recent = self.keep_recent_turns,
284            compacting_turns = turns_compacted,
285            tool_compaction_strategy = %self.tool_compaction,
286            "Starting conversation compaction"
287        );
288
289        // Compact older messages
290        let mut total_summarized = 0;
291        let mut total_redacted = 0;
292
293        for msg in conversation.iter_mut() {
294            let turn_id = msg.turn_id();
295            if turns_to_keep.contains(turn_id) {
296                continue; // Keep recent turns intact
297            }
298
299            let (summarized, redacted) = self.compact_message(msg, compact_summaries);
300            total_summarized += summarized;
301            total_redacted += redacted;
302        }
303
304        tracing::info!(
305            tool_results_summarized = total_summarized,
306            tool_results_redacted = total_redacted,
307            turns_compacted,
308            "Conversation compaction completed"
309        );
310
311        CompactionResult {
312            tool_results_summarized: total_summarized,
313            tool_results_redacted: total_redacted,
314            turns_compacted,
315        }
316    }
317}
318
319// ============================================================================
320// LLM-Based Compaction
321// ============================================================================
322
323/// Default system prompt for conversation summarization.
324pub const DEFAULT_SUMMARY_SYSTEM_PROMPT: &str = r#"You are a conversation summarizer. Your task is to create a concise summary of the conversation history provided.
325
326Guidelines:
327- Capture the key topics discussed, decisions made, and important context
328- Preserve any technical details, file paths, code snippets, or specific values that would be needed to continue the conversation
329- Include the user's original goals and any progress made toward them
330- Note any pending tasks or unresolved questions
331- Keep the summary focused and actionable
332- Format the summary as a narrative that provides context for continuing the conversation
333
334Respond with only the summary, no additional commentary."#;
335
336/// Default maximum tokens for summary responses.
337pub const DEFAULT_MAX_SUMMARY_TOKENS: i64 = 2048;
338
339/// Default timeout for summarization LLM calls (60 seconds).
340pub const DEFAULT_SUMMARY_TIMEOUT: Duration = Duration::from_secs(60);
341
342/// Configuration for LLM-based conversation compaction.
343#[derive(Debug, Clone)]
344pub struct LLMCompactorConfig {
345    /// Context utilization ratio (0.0-1.0) that triggers compaction.
346    pub threshold: f64,
347
348    /// Number of recent turns to preserve during compaction.
349    pub keep_recent_turns: usize,
350
351    /// System prompt for summarization. Uses DEFAULT_SUMMARY_SYSTEM_PROMPT if None.
352    pub summary_system_prompt: Option<String>,
353
354    /// Maximum tokens for summary response. Uses DEFAULT_MAX_SUMMARY_TOKENS if None.
355    pub max_summary_tokens: Option<i64>,
356
357    /// Timeout for summarization call. Uses DEFAULT_SUMMARY_TIMEOUT if None.
358    pub summary_timeout: Option<Duration>,
359}
360
361impl LLMCompactorConfig {
362    /// Creates a new LLM compactor config with default optional values.
363    pub fn new(threshold: f64, keep_recent_turns: usize) -> Self {
364        Self {
365            threshold,
366            keep_recent_turns,
367            summary_system_prompt: None,
368            max_summary_tokens: None,
369            summary_timeout: None,
370        }
371    }
372
373    /// Validates the configuration.
374    pub fn validate(&self) -> Result<(), CompactorConfigError> {
375        if self.threshold <= 0.0 || self.threshold >= 1.0 {
376            return Err(CompactorConfigError::InvalidThreshold(self.threshold));
377        }
378        Ok(())
379    }
380
381    /// Returns the system prompt to use (config value or default).
382    pub fn system_prompt(&self) -> &str {
383        self.summary_system_prompt
384            .as_deref()
385            .unwrap_or(DEFAULT_SUMMARY_SYSTEM_PROMPT)
386    }
387
388    /// Returns the max tokens to use (config value or default).
389    pub fn max_tokens(&self) -> i64 {
390        self.max_summary_tokens
391            .unwrap_or(DEFAULT_MAX_SUMMARY_TOKENS)
392    }
393
394    /// Returns the timeout to use (config value or default).
395    pub fn timeout(&self) -> Duration {
396        self.summary_timeout.unwrap_or(DEFAULT_SUMMARY_TIMEOUT)
397    }
398}
399
400impl Default for LLMCompactorConfig {
401    fn default() -> Self {
402        Self::new(0.75, 5)
403    }
404}
405
406/// LLM-based compactor that summarizes older conversation using an LLM.
407/// It replaces older messages with a single summary message while preserving recent turns.
408pub struct LLMCompactor {
409    /// LLMClient client for making summarization LLM calls.
410    client: LLMClient,
411
412    /// Configuration for compaction behavior.
413    config: LLMCompactorConfig,
414}
415
416impl LLMCompactor {
417    /// Creates a new LLM compactor.
418    ///
419    /// # Arguments
420    /// * `client` - LLMClient client for making LLM calls
421    /// * `config` - Configuration for compaction behavior
422    ///
423    /// # Returns
424    /// Error if configuration is invalid.
425    pub fn new(
426        client: LLMClient,
427        config: LLMCompactorConfig,
428    ) -> Result<Self, CompactorConfigError> {
429        config.validate()?;
430
431        tracing::info!(
432            threshold = config.threshold,
433            keep_recent_turns = config.keep_recent_turns,
434            max_summary_tokens = config.max_tokens(),
435            "LLM compactor created"
436        );
437
438        Ok(Self { client, config })
439    }
440
441    /// Extract unique turn IDs in order of first appearance.
442    fn unique_turn_ids(&self, conversation: &[Message]) -> Vec<TurnId> {
443        let mut seen = HashSet::new();
444        let mut ids = Vec::new();
445
446        for msg in conversation {
447            let turn_id = msg.turn_id();
448            if seen.insert(turn_id.clone()) {
449                ids.push(turn_id.clone());
450            }
451        }
452
453        ids
454    }
455
456    /// Format messages for LLM summarization.
457    fn format_messages_for_summary(&self, messages: &[Message]) -> String {
458        let mut builder = String::new();
459
460        for msg in messages {
461            // Add role label
462            if msg.is_user() {
463                builder.push_str("User: ");
464            } else {
465                builder.push_str("Assistant: ");
466            }
467
468            // Format content blocks
469            for block in msg.content() {
470                match block {
471                    ContentBlock::Text(text) => {
472                        builder.push_str(&text.text);
473                    }
474                    ContentBlock::ToolUse(tool_use) => {
475                        builder.push_str(&format!(
476                            "[Called tool: {} with input: {:?}]",
477                            tool_use.name, tool_use.input
478                        ));
479                    }
480                    ContentBlock::ToolResult(tool_result) => {
481                        let content = truncate_content(&tool_result.content, 1000);
482                        if tool_result.is_error {
483                            builder.push_str(&format!("[Tool error: {}]", content));
484                        } else {
485                            builder.push_str(&format!("[Tool result: {}]", content));
486                        }
487                    }
488                }
489            }
490            builder.push_str("\n\n");
491        }
492
493        builder
494    }
495
496    /// Create a summary message from LLM output.
497    fn create_summary_message(&self, summary: &str, session_id: &str) -> Message {
498        // Use turn number 0 to ensure it comes before all other turns
499        let turn_id = TurnId::new_user_turn(0);
500
501        // Use std::time for timestamps
502        let now = std::time::SystemTime::now()
503            .duration_since(std::time::UNIX_EPOCH)
504            .unwrap_or_default();
505
506        Message::User(UserMessage {
507            id: format!("summary-{}", now.as_nanos()),
508            session_id: session_id.to_string(),
509            turn_id,
510            created_at: now.as_secs() as i64,
511            content: vec![ContentBlock::Text(TextBlock {
512                text: format!("[Previous conversation summary]:\n\n{}", summary),
513            })],
514        })
515    }
516
517    /// Get session ID from conversation (uses first message's session_id).
518    fn get_session_id(&self, conversation: &[Message]) -> String {
519        conversation
520            .first()
521            .map(|m| m.session_id().to_string())
522            .unwrap_or_default()
523    }
524}
525
526impl Compactor for LLMCompactor {
527    fn should_compact(&self, context_used: i64, context_limit: i32) -> bool {
528        if context_limit == 0 {
529            return false;
530        }
531
532        let utilization = context_used as f64 / context_limit as f64;
533        let should_compact = utilization > self.config.threshold;
534
535        if should_compact {
536            tracing::info!(
537                utilization = utilization,
538                threshold = self.config.threshold,
539                context_used,
540                context_limit,
541                "LLM compaction triggered - context utilization exceeded threshold"
542            );
543        }
544
545        should_compact
546    }
547
548    fn compact(
549        &self,
550        _conversation: &mut Vec<Message>,
551        _compact_summaries: &HashMap<String, String>,
552    ) -> CompactionResult {
553        // LLMCompactor requires async - this should not be called
554        tracing::warn!("LLMCompactor::compact() called - use compact_async() instead");
555        CompactionResult::default()
556    }
557
558    fn is_async(&self) -> bool {
559        true
560    }
561}
562
563impl AsyncCompactor for LLMCompactor {
564    fn compact_async<'a>(
565        &'a self,
566        conversation: Vec<Message>,
567        _compact_summaries: &'a HashMap<String, String>,
568    ) -> CompactAsyncFuture<'a> {
569        Box::pin(async move {
570            if conversation.is_empty() {
571                return Ok((conversation, CompactionResult::default()));
572            }
573
574            // Find unique turn IDs in order
575            let turn_ids = self.unique_turn_ids(&conversation);
576
577            // If fewer turns than keep_recent_turns, nothing to compact
578            if turn_ids.len() <= self.config.keep_recent_turns {
579                tracing::debug!(
580                    total_turns = turn_ids.len(),
581                    keep_recent = self.config.keep_recent_turns,
582                    "Skipping LLM compaction - not enough turns"
583                );
584                return Ok((conversation, CompactionResult::default()));
585            }
586
587            // Separate messages into old (to summarize) and recent (to keep)
588            let start_idx = turn_ids.len() - self.config.keep_recent_turns;
589            let turns_to_keep: HashSet<_> = turn_ids[start_idx..].iter().cloned().collect();
590
591            let mut old_messages = Vec::new();
592            let mut recent_messages = Vec::new();
593
594            for msg in conversation {
595                if turns_to_keep.contains(msg.turn_id()) {
596                    recent_messages.push(msg);
597                } else {
598                    old_messages.push(msg);
599                }
600            }
601
602            if old_messages.is_empty() {
603                return Ok((recent_messages, CompactionResult::default()));
604            }
605
606            let session_id = self.get_session_id(&old_messages);
607
608            // Format old messages for summarization
609            let formatted_conversation = self.format_messages_for_summary(&old_messages);
610
611            tracing::info!(
612                total_turns = turn_ids.len(),
613                turns_to_summarize = start_idx,
614                turns_to_keep = self.config.keep_recent_turns,
615                messages_to_summarize = old_messages.len(),
616                formatted_length = formatted_conversation.len(),
617                "Starting LLM conversation compaction"
618            );
619
620            // Call LLM for summarization with timeout
621            let options = MessageOptions {
622                max_tokens: Some(self.config.max_tokens() as u32),
623                ..Default::default()
624            };
625
626            // Create messages for summarization
627            let llm_messages = vec![
628                LLMMessage::system(self.config.system_prompt()),
629                LLMMessage::user(formatted_conversation),
630            ];
631
632            // Make LLM call with timeout
633            let result = tokio::time::timeout(
634                self.config.timeout(),
635                self.client.send_message(&llm_messages, &options),
636            )
637            .await;
638
639            let response = match result {
640                Ok(Ok(msg)) => msg,
641                Ok(Err(e)) => {
642                    tracing::error!(error = %e, "LLM compaction failed");
643                    return Err(CompactionError::LLMError(e.to_string()));
644                }
645                Err(_) => {
646                    tracing::error!("LLM compaction timed out");
647                    return Err(CompactionError::Timeout);
648                }
649            };
650
651            // Extract text from response
652            let summary_text = response
653                .content
654                .iter()
655                .filter_map(|c| {
656                    if let crate::client::models::Content::Text(t) = c {
657                        Some(t.as_str())
658                    } else {
659                        None
660                    }
661                })
662                .collect::<Vec<_>>()
663                .join("");
664
665            // Create summary message
666            let summary_message = self.create_summary_message(&summary_text, &session_id);
667
668            // Build new conversation: summary + recent messages
669            let mut new_conversation = Vec::with_capacity(1 + recent_messages.len());
670            new_conversation.push(summary_message);
671            new_conversation.extend(recent_messages);
672
673            let result = CompactionResult {
674                tool_results_summarized: 0,
675                tool_results_redacted: 0,
676                turns_compacted: start_idx,
677            };
678
679            tracing::info!(
680                original_messages = old_messages.len() + result.turns_compacted,
681                new_messages = new_conversation.len(),
682                summary_length = summary_text.len(),
683                turns_compacted = result.turns_compacted,
684                "LLM compaction completed"
685            );
686
687            Ok((new_conversation, result))
688        })
689    }
690}
691
692/// Truncate content to a maximum length, adding ellipsis if needed.
693fn truncate_content(content: &str, max_len: usize) -> String {
694    if content.len() <= max_len {
695        content.to_string()
696    } else {
697        format!("{}...", &content[..max_len.saturating_sub(3)])
698    }
699}
700
701#[cfg(test)]
702mod tests {
703    use super::*;
704    use crate::controller::types::{AssistantMessage, UserMessage};
705
706    fn make_user_message(turn_id: TurnId) -> Message {
707        Message::User(UserMessage {
708            id: format!("msg_{}", turn_id),
709            session_id: "test_session".to_string(),
710            turn_id,
711            created_at: 0,
712            content: vec![ContentBlock::Text(crate::controller::types::TextBlock {
713                text: "test".to_string(),
714            })],
715        })
716    }
717
718    fn make_assistant_message(turn_id: TurnId) -> Message {
719        Message::Assistant(AssistantMessage {
720            id: format!("msg_{}", turn_id),
721            session_id: "test_session".to_string(),
722            turn_id,
723            parent_id: String::new(),
724            created_at: 0,
725            completed_at: None,
726            model_id: "test_model".to_string(),
727            provider_id: "test_provider".to_string(),
728            input_tokens: 0,
729            output_tokens: 0,
730            cache_read_tokens: 0,
731            cache_write_tokens: 0,
732            finish_reason: None,
733            error: None,
734            content: vec![ContentBlock::Text(crate::controller::types::TextBlock {
735                text: "test".to_string(),
736            })],
737        })
738    }
739
740    fn make_tool_result_message(tool_use_id: &str, content: &str, turn_id: TurnId) -> Message {
741        Message::User(UserMessage {
742            id: format!("msg_{}", turn_id),
743            session_id: "test_session".to_string(),
744            turn_id,
745            created_at: 0,
746            content: vec![ContentBlock::ToolResult(
747                crate::controller::types::ToolResultBlock {
748                    tool_use_id: tool_use_id.to_string(),
749                    content: content.to_string(),
750                    is_error: false,
751                    compact_summary: None,
752                },
753            )],
754        })
755    }
756
757    #[test]
758    fn test_threshold_compactor_creation() {
759        // Valid threshold
760        let compactor = ThresholdCompactor::new(0.75, 3, ToolCompaction::Redact);
761        assert!(compactor.is_ok());
762
763        // Invalid threshold (too low)
764        let compactor = ThresholdCompactor::new(0.0, 3, ToolCompaction::Redact);
765        assert!(compactor.is_err());
766
767        // Invalid threshold (too high)
768        let compactor = ThresholdCompactor::new(1.0, 3, ToolCompaction::Redact);
769        assert!(compactor.is_err());
770    }
771
772    #[test]
773    fn test_should_compact() {
774        let compactor = ThresholdCompactor::new(0.75, 3, ToolCompaction::Redact).unwrap();
775
776        // Below threshold - don't compact
777        assert!(!compactor.should_compact(7000, 10000));
778
779        // Above threshold - compact
780        assert!(compactor.should_compact(8000, 10000));
781
782        // Zero context limit - don't compact
783        assert!(!compactor.should_compact(8000, 0));
784    }
785
786    #[test]
787    fn test_compact_not_enough_turns() {
788        let compactor = ThresholdCompactor::new(0.75, 3, ToolCompaction::Redact).unwrap();
789
790        let mut conversation = vec![
791            make_user_message(TurnId::new_user_turn(1)),
792            make_assistant_message(TurnId::new_assistant_turn(1)),
793        ];
794
795        let summaries = std::collections::HashMap::new();
796        let result = compactor.compact(&mut conversation, &summaries);
797
798        // Only 2 turns, but we keep 3, so nothing compacted
799        assert_eq!(result.turns_compacted, 0);
800    }
801
802    #[test]
803    fn test_compact_redacts_old_tool_results() {
804        // keep_recent_turns=2 keeps u2 and a2, compacts u1 and a1
805        let compactor = ThresholdCompactor::new(0.75, 2, ToolCompaction::Redact).unwrap();
806
807        let mut conversation = vec![
808            // Turn 1 - old, should be compacted
809            make_tool_result_message("tool_1", "old result", TurnId::new_user_turn(1)),
810            make_assistant_message(TurnId::new_assistant_turn(1)),
811            // Turn 2 - recent, should be kept
812            make_tool_result_message("tool_2", "new result", TurnId::new_user_turn(2)),
813            make_assistant_message(TurnId::new_assistant_turn(2)),
814        ];
815
816        let summaries = std::collections::HashMap::new();
817        let result = compactor.compact(&mut conversation, &summaries);
818
819        assert_eq!(result.tool_results_redacted, 1);
820        assert_eq!(result.turns_compacted, 2); // u1 and a1 compacted
821
822        // Check that old tool result was redacted
823        if let ContentBlock::ToolResult(tr) = &conversation[0].content()[0] {
824            assert!(tr.content.contains("redacted"));
825        } else {
826            panic!("Expected ToolResult");
827        }
828
829        // Check that new tool result was kept
830        if let ContentBlock::ToolResult(tr) = &conversation[2].content()[0] {
831            assert_eq!(tr.content, "new result");
832        } else {
833            panic!("Expected ToolResult");
834        }
835    }
836
837    #[test]
838    fn test_compact_summarizes_with_summary() {
839        // keep_recent_turns=2 keeps u2 and a2, compacts u1 and a1
840        let compactor = ThresholdCompactor::new(0.75, 2, ToolCompaction::Summarize).unwrap();
841
842        let mut conversation = vec![
843            make_tool_result_message("tool_1", "very long result", TurnId::new_user_turn(1)),
844            make_assistant_message(TurnId::new_assistant_turn(1)),
845            make_user_message(TurnId::new_user_turn(2)),
846            make_assistant_message(TurnId::new_assistant_turn(2)),
847        ];
848
849        let mut summaries = std::collections::HashMap::new();
850        summaries.insert("tool_1".to_string(), "[summary]".to_string());
851
852        let result = compactor.compact(&mut conversation, &summaries);
853
854        assert_eq!(result.tool_results_summarized, 1);
855
856        // Check that tool result was summarized
857        if let ContentBlock::ToolResult(tr) = &conversation[0].content()[0] {
858            assert_eq!(tr.content, "[summary]");
859        } else {
860            panic!("Expected ToolResult");
861        }
862    }
863
864    // ============================================================================
865    // LLMCompactorConfig Tests
866    // ============================================================================
867
868    #[test]
869    fn test_llm_compactor_config_creation() {
870        let config = LLMCompactorConfig::new(0.75, 5);
871        assert_eq!(config.threshold, 0.75);
872        assert_eq!(config.keep_recent_turns, 5);
873        assert!(config.summary_system_prompt.is_none());
874        assert!(config.max_summary_tokens.is_none());
875        assert!(config.summary_timeout.is_none());
876    }
877
878    #[test]
879    fn test_llm_compactor_config_validation() {
880        // Valid config
881        let config = LLMCompactorConfig::new(0.75, 5);
882        assert!(config.validate().is_ok());
883
884        // Invalid threshold (too low)
885        let config = LLMCompactorConfig::new(0.0, 5);
886        assert!(config.validate().is_err());
887
888        // Invalid threshold (too high)
889        let config = LLMCompactorConfig::new(1.0, 5);
890        assert!(config.validate().is_err());
891
892        // Edge case: threshold just above 0
893        let config = LLMCompactorConfig::new(0.01, 5);
894        assert!(config.validate().is_ok());
895
896        // Edge case: threshold just below 1
897        let config = LLMCompactorConfig::new(0.99, 5);
898        assert!(config.validate().is_ok());
899    }
900
901    #[test]
902    fn test_llm_compactor_config_defaults() {
903        let config = LLMCompactorConfig::default();
904        assert_eq!(config.threshold, 0.75);
905        assert_eq!(config.keep_recent_turns, 5);
906
907        // Check that getter methods return defaults
908        assert_eq!(config.system_prompt(), DEFAULT_SUMMARY_SYSTEM_PROMPT);
909        assert_eq!(config.max_tokens(), DEFAULT_MAX_SUMMARY_TOKENS);
910        assert_eq!(config.timeout(), DEFAULT_SUMMARY_TIMEOUT);
911    }
912
913    #[test]
914    fn test_llm_compactor_config_custom_values() {
915        let config = LLMCompactorConfig {
916            threshold: 0.8,
917            keep_recent_turns: 3,
918            summary_system_prompt: Some("Custom prompt".to_string()),
919            max_summary_tokens: Some(4096),
920            summary_timeout: Some(Duration::from_secs(120)),
921        };
922
923        assert_eq!(config.system_prompt(), "Custom prompt");
924        assert_eq!(config.max_tokens(), 4096);
925        assert_eq!(config.timeout(), Duration::from_secs(120));
926    }
927
928    #[test]
929    fn test_truncate_content() {
930        // Short content - no truncation
931        assert_eq!(truncate_content("hello", 10), "hello");
932
933        // Exact length - no truncation
934        assert_eq!(truncate_content("hello", 5), "hello");
935
936        // Long content - truncated with ellipsis
937        assert_eq!(truncate_content("hello world", 8), "hello...");
938
939        // Very short max - edge case
940        assert_eq!(truncate_content("hello", 3), "...");
941    }
942}