1use crate::llm::{
4 ChatOutcome, ChatRequest, Content, ContentBlock, LlmProvider, Message, Role, StopReason,
5};
6use anyhow::{Context, Result, bail};
7use async_trait::async_trait;
8use std::fmt::Write;
9use std::sync::Arc;
10
11use super::config::CompactionConfig;
12use super::estimator::TokenEstimator;
13
14const SUMMARY_PREFIX: &str = "[Previous conversation summary]\n\n";
15const COMPACTION_SYSTEM_PROMPT: &str = "You are a precise summarizer. Your task is to create concise but complete summaries of conversations, preserving all technical details needed to continue the work.";
16const COMPACTION_SUMMARY_PROMPT_PREFIX: &str = "Summarize this conversation concisely, preserving:\n- Key decisions and conclusions reached\n- Important file paths, code changes, and technical details\n- Current task context and what has been accomplished\n- Any pending items, errors encountered, or next steps\n\nBe specific about technical details (file names, function names, error messages) as these\nare critical for continuing the work.\n\nConversation:\n";
17const COMPACTION_SUMMARY_PROMPT_SUFFIX: &str =
18 "Provide a concise summary (aim for 500-1000 words):";
19const COMPACT_EMPTY_SUMMARY: &str = "No additional context was available to summarize; the previous messages were already compacted.";
20const SUMMARY_ACKNOWLEDGMENT: &str =
21 "I understand the context from the summary. Let me continue from where we left off.";
22const MAX_TOOL_RESULT_CHARS: usize = 500;
23const TRUNCATED_SUMMARY_MARKER: &str =
24 "\n\n[summary truncated: exceeded the configured summary_max_tokens budget]";
25
26#[async_trait]
30pub trait ContextCompactor: Send + Sync {
31 async fn compact(&self, messages: &[Message]) -> Result<String>;
36
37 fn estimate_tokens(&self, messages: &[Message]) -> usize;
39
40 fn needs_compaction(&self, messages: &[Message]) -> bool;
42
43 async fn compact_history(&self, messages: Vec<Message>) -> Result<CompactionResult>;
48}
49
50#[derive(Debug, Clone)]
52pub struct CompactionResult {
53 pub messages: Vec<Message>,
55 pub original_count: usize,
57 pub new_count: usize,
59 pub original_tokens: usize,
61 pub new_tokens: usize,
63}
64
65pub struct LlmContextCompactor<P: LlmProvider + ?Sized> {
76 provider: Arc<P>,
77 config: CompactionConfig,
78 system_prompt: String,
79 summary_prompt_prefix: String,
80 summary_prompt_suffix: String,
81}
82
83impl<P: LlmProvider + ?Sized> LlmContextCompactor<P> {
84 #[must_use]
86 pub fn new(provider: Arc<P>, config: CompactionConfig) -> Self {
87 Self {
88 provider,
89 config,
90 system_prompt: COMPACTION_SYSTEM_PROMPT.to_string(),
91 summary_prompt_prefix: COMPACTION_SUMMARY_PROMPT_PREFIX.to_string(),
92 summary_prompt_suffix: COMPACTION_SUMMARY_PROMPT_SUFFIX.to_string(),
93 }
94 }
95
96 #[must_use]
98 pub fn with_defaults(provider: Arc<P>) -> Self {
99 Self::new(provider, CompactionConfig::default())
100 }
101
102 #[must_use]
104 pub const fn config(&self) -> &CompactionConfig {
105 &self.config
106 }
107
108 #[must_use]
110 pub fn with_prompts(
111 mut self,
112 system_prompt: impl Into<String>,
113 summary_prompt_prefix: impl Into<String>,
114 summary_prompt_suffix: impl Into<String>,
115 ) -> Self {
116 self.system_prompt = system_prompt.into();
117 self.summary_prompt_prefix = summary_prompt_prefix.into();
118 self.summary_prompt_suffix = summary_prompt_suffix.into();
119 self
120 }
121
122 fn extract_summary_text(content: &Content) -> Option<String> {
132 match content {
133 Content::Text(text) => text.strip_prefix(SUMMARY_PREFIX).map(str::to_string),
134 Content::Blocks(blocks) => blocks.iter().find_map(|block| match block {
135 ContentBlock::Text { text } => {
136 text.strip_prefix(SUMMARY_PREFIX).map(str::to_string)
137 }
138 _ => None,
139 }),
140 }
141 }
142
143 fn has_tool_use(content: &Content) -> bool {
145 matches!(
146 content,
147 Content::Blocks(blocks)
148 if blocks
149 .iter()
150 .any(|block| matches!(block, ContentBlock::ToolUse { .. }))
151 )
152 }
153
154 fn has_tool_result(content: &Content) -> bool {
156 matches!(
157 content,
158 Content::Blocks(blocks)
159 if blocks
160 .iter()
161 .any(|block| matches!(block, ContentBlock::ToolResult { .. }))
162 )
163 }
164
165 fn has_opaque_reasoning(messages: &[Message]) -> bool {
174 messages.iter().any(|message| {
175 matches!(
176 &message.content,
177 Content::Blocks(blocks)
178 if blocks
179 .iter()
180 .any(|block| matches!(block, ContentBlock::OpaqueReasoning { .. }))
181 )
182 })
183 }
184
185 fn split_point_preserves_tool_pairs(messages: &[Message], mut split_point: usize) -> usize {
198 while split_point > 0 && split_point < messages.len() {
199 let prev = &messages[split_point - 1];
200 let next = &messages[split_point];
201
202 let crosses_tool_pair = prev.role == Role::Assistant
203 && Self::has_tool_use(&prev.content)
204 && next.role == Role::User
205 && Self::has_tool_result(&next.content);
206
207 if crosses_tool_pair {
208 split_point -= 1;
209 continue;
210 }
211
212 break;
213 }
214
215 split_point
216 }
217
218 fn split_point_preserves_tool_pairs_with_cap(
257 messages: &[Message],
258 split_point: usize,
259 max_tokens: usize,
260 ) -> usize {
261 let cap_limit = Self::retain_tail_with_token_cap(messages, split_point, max_tokens);
262 let pair_safe = Self::split_point_preserves_tool_pairs(messages, cap_limit);
263 Self::split_point_skips_leading_orphan(messages, pair_safe)
264 }
265
266 fn split_point_skips_leading_orphan(messages: &[Message], mut split_point: usize) -> usize {
283 while split_point < messages.len() {
284 if Self::leading_message_has_orphan_tool_result(&messages[split_point..]) {
285 split_point = split_point.saturating_add(1);
286 continue;
287 }
288 break;
289 }
290 split_point
291 }
292
293 fn leading_message_has_orphan_tool_result(to_keep: &[Message]) -> bool {
301 let Some(first) = to_keep.first() else {
302 return false;
303 };
304 let Content::Blocks(blocks) = &first.content else {
305 return false;
306 };
307
308 let mut needed: Vec<&str> = Vec::new();
312 for block in blocks {
313 if let ContentBlock::ToolResult { tool_use_id, .. } = block {
314 needed.push(tool_use_id.as_str());
315 }
316 }
317 if needed.is_empty() {
318 return false;
319 }
320
321 let known_ids: std::collections::HashSet<&str> = to_keep
323 .iter()
324 .flat_map(|message| match &message.content {
325 Content::Blocks(blocks) => blocks
326 .iter()
327 .filter_map(|block| match block {
328 ContentBlock::ToolUse { id, .. } => Some(id.as_str()),
329 _ => None,
330 })
331 .collect::<Vec<_>>(),
332 Content::Text(_) => Vec::new(),
333 })
334 .collect();
335
336 needed.iter().any(|id| !known_ids.contains(id))
337 }
338
339 fn retain_tail_with_token_cap(messages: &[Message], start: usize, max_tokens: usize) -> usize {
341 if start >= messages.len() {
342 return messages.len();
343 }
344
345 if max_tokens == 0 {
346 return messages.len();
347 }
348
349 let mut used = 0usize;
350 let mut retained_start = messages.len();
351
352 for idx in (start..messages.len()).rev() {
353 let message_tokens = TokenEstimator::estimate_message(&messages[idx]);
354 if used + message_tokens > max_tokens {
355 break;
356 }
357
358 retained_start = idx;
359 used += message_tokens;
360 }
361
362 retained_start
363 }
364
365 fn format_messages_for_summary<'a>(messages: impl IntoIterator<Item = &'a Message>) -> String {
370 let mut output = String::new();
371
372 for message in messages {
373 let role = match message.role {
374 Role::User => "User",
375 Role::Assistant => "Assistant",
376 };
377
378 let _ = write!(output, "{role}: ");
379
380 match &message.content {
381 Content::Text(text) => {
382 let _ = writeln!(output, "{text}");
383 }
384 Content::Blocks(blocks) => {
385 for block in blocks {
386 match block {
387 ContentBlock::Text { text } => {
388 let _ = writeln!(output, "{text}");
389 }
390 ContentBlock::Thinking { thinking, .. } => {
391 let _ = writeln!(output, "[Thinking: {thinking}]");
393 }
394 ContentBlock::RedactedThinking { .. } => {
395 let _ = writeln!(output, "[Redacted thinking]");
396 }
397 ContentBlock::OpaqueReasoning { .. } => {
398 let _ = writeln!(output, "[Opaque reasoning state omitted]");
403 }
404 ContentBlock::ToolUse { name, input, .. } => {
405 let _ = writeln!(
406 output,
407 "[Called tool: {name} with input: {}]",
408 serde_json::to_string(input).unwrap_or_default()
409 );
410 }
411 ContentBlock::ToolResult {
412 content, is_error, ..
413 } => {
414 let status = if is_error.unwrap_or(false) {
415 "error"
416 } else {
417 "success"
418 };
419 let truncated = if content.chars().count() > MAX_TOOL_RESULT_CHARS {
421 let prefix: String =
422 content.chars().take(MAX_TOOL_RESULT_CHARS).collect();
423 format!("{prefix}... (truncated)")
424 } else {
425 content.clone()
426 };
427 let _ = writeln!(output, "[Tool result ({status}): {truncated}]");
428 }
429 ContentBlock::Image { source } => {
430 let _ = writeln!(output, "[Image: {}]", source.media_type);
431 }
432 ContentBlock::Document { source } => {
433 let _ = writeln!(output, "[Document: {}]", source.media_type);
434 }
435 _ => {
438 let _ = writeln!(output, "[Unrecognized content block]");
439 }
440 }
441 }
442 }
443 }
444 output.push('\n');
445 }
446
447 output
448 }
449
450 fn build_summary_prompt(&self, prior_summaries: &[String], messages_text: &str) -> String {
457 let base = format!(
458 "{}{}{}",
459 self.summary_prompt_prefix, messages_text, self.summary_prompt_suffix
460 );
461
462 if prior_summaries.is_empty() {
463 return base;
464 }
465
466 let prior = prior_summaries.join("\n\n");
467 format!(
468 "Previous summary of earlier conversation. Preserve every fact below \
469 in your new summary so no earlier context is lost:\n{prior}\n\n{base}"
470 )
471 }
472
473 async fn run_summarization(&self, prompt: String, max_tokens: usize) -> Result<(String, bool)> {
476 let request = ChatRequest {
477 system: self.system_prompt.clone(),
478 messages: vec![Message::user(prompt)],
479 tools: None,
480 max_tokens: u32::try_from(max_tokens).unwrap_or(u32::MAX),
481 max_tokens_explicit: true,
482 session_id: None,
483 cached_content: None,
484 thinking: None,
485 tool_choice: None,
486 response_format: None,
487 cache: None,
488 };
489
490 let outcome = self
491 .provider
492 .chat(request)
493 .await
494 .context("Failed to call LLM for summarization")?;
495
496 match outcome {
497 ChatOutcome::Success(response) => {
498 let truncated = response.stop_reason == Some(StopReason::MaxTokens);
499 let text = response
500 .first_text()
501 .map(String::from)
502 .context("No text in summarization response")?;
503 Ok((text, truncated))
504 }
505 ChatOutcome::RateLimited(_) => {
506 bail!("Rate limited during summarization")
507 }
508 ChatOutcome::InvalidRequest(msg) => {
509 bail!("Invalid request during summarization: {msg}")
510 }
511 ChatOutcome::ServerError(msg) => {
512 bail!("Server error during summarization: {msg}")
513 }
514 _ => {
517 bail!("Unrecognized provider outcome during summarization")
518 }
519 }
520 }
521}
522
523#[async_trait]
524impl<P: LlmProvider + ?Sized> ContextCompactor for LlmContextCompactor<P> {
525 async fn compact(&self, messages: &[Message]) -> Result<String> {
526 let mut prior_summaries: Vec<String> = Vec::new();
531 let mut fresh: Vec<&Message> = Vec::new();
532 for message in messages {
533 if let Some(text) = Self::extract_summary_text(&message.content) {
534 if !text.is_empty() {
535 prior_summaries.push(text);
536 }
537 } else {
538 fresh.push(message);
539 }
540 }
541
542 if fresh.is_empty() {
545 if prior_summaries.is_empty() {
546 return Ok(COMPACT_EMPTY_SUMMARY.to_string());
547 }
548 return Ok(prior_summaries.join("\n\n"));
549 }
550
551 let messages_text = Self::format_messages_for_summary(fresh.iter().copied());
552 let prompt = self.build_summary_prompt(&prior_summaries, &messages_text);
553
554 let budget = self.config.summary_max_tokens;
555 let (mut summary, truncated) = self.run_summarization(prompt.clone(), budget).await?;
556
557 if truncated {
558 log::warn!(
559 "compaction summary hit the max_tokens budget ({budget}); \
560 retrying with a larger budget to avoid silent context loss"
561 );
562 let (retry_summary, still_truncated) = self
563 .run_summarization(prompt, budget.saturating_mul(2))
564 .await?;
565 summary = retry_summary;
566 if still_truncated {
567 log::warn!(
568 "compaction summary still truncated after retry; appending a \
569 truncation marker so downstream context loss is visible"
570 );
571 summary.push_str(TRUNCATED_SUMMARY_MARKER);
572 }
573 }
574
575 Ok(summary)
576 }
577
578 fn estimate_tokens(&self, messages: &[Message]) -> usize {
579 TokenEstimator::estimate_history(messages)
580 }
581
582 fn needs_compaction(&self, messages: &[Message]) -> bool {
583 if !self.config.auto_compact {
584 return false;
585 }
586
587 if Self::has_opaque_reasoning(messages) {
592 return false;
593 }
594
595 if messages.len() < self.config.min_messages_for_compaction {
596 return false;
597 }
598
599 let estimated_tokens = self.estimate_tokens(messages);
600 estimated_tokens > self.config.threshold_tokens
601 }
602
603 async fn compact_history(&self, mut messages: Vec<Message>) -> Result<CompactionResult> {
604 let original_count = messages.len();
605 let original_tokens = self.estimate_tokens(&messages);
606
607 if Self::has_opaque_reasoning(&messages) {
612 log::debug!(
613 "skipping generic context compaction for history with opaque provider reasoning state"
614 );
615 return Ok(CompactionResult {
616 messages,
617 original_count,
618 new_count: original_count,
619 original_tokens,
620 new_tokens: original_tokens,
621 });
622 }
623
624 if messages.len() <= self.config.retain_recent {
626 return Ok(CompactionResult {
627 messages,
628 original_count,
629 new_count: original_count,
630 original_tokens,
631 new_tokens: original_tokens,
632 });
633 }
634
635 let mut split_point = messages.len().saturating_sub(self.config.retain_recent);
637 split_point = Self::split_point_preserves_tool_pairs_with_cap(
638 &messages,
639 split_point,
640 self.config.max_retained_tail_tokens,
641 );
642
643 let to_keep = messages.split_off(split_point);
646 let to_summarize = messages;
647
648 let summary = self.compact(&to_summarize).await?;
650
651 let mut new_messages = Vec::with_capacity(2 + to_keep.len());
653
654 new_messages.push(Message::user(format!("{SUMMARY_PREFIX}{summary}")));
656
657 if !to_keep.is_empty() {
662 new_messages.push(Message::assistant(SUMMARY_ACKNOWLEDGMENT));
663 }
664
665 new_messages.extend(to_keep);
673
674 let new_count = new_messages.len();
675 let new_tokens = self.estimate_tokens(&new_messages);
676
677 Ok(CompactionResult {
678 messages: new_messages,
679 original_count,
680 new_count,
681 original_tokens,
682 new_tokens,
683 })
684 }
685}
686
687#[cfg(test)]
688mod tests {
689 use super::*;
690 use crate::llm::{ChatResponse, StopReason, Usage};
691 use std::sync::Mutex;
692
693 struct MockProvider {
694 summary_response: String,
695 requests: Arc<Mutex<Vec<String>>>,
696 echo_input: bool,
699 stop_reason: StopReason,
701 }
702
703 impl MockProvider {
704 fn build(
705 summary: &str,
706 requests: Arc<Mutex<Vec<String>>>,
707 echo_input: bool,
708 stop_reason: StopReason,
709 ) -> Self {
710 Self {
711 summary_response: summary.to_string(),
712 requests,
713 echo_input,
714 stop_reason,
715 }
716 }
717
718 fn new(summary: &str) -> Self {
719 Self::build(
720 summary,
721 Arc::new(Mutex::new(Vec::new())),
722 false,
723 StopReason::EndTurn,
724 )
725 }
726
727 fn new_with_request_log(summary: &str, requests: Arc<Mutex<Vec<String>>>) -> Self {
728 Self::build(summary, requests, false, StopReason::EndTurn)
729 }
730
731 fn new_echo(requests: Arc<Mutex<Vec<String>>>) -> Self {
733 Self::build("", requests, true, StopReason::EndTurn)
734 }
735
736 fn new_truncating(summary: &str, requests: Arc<Mutex<Vec<String>>>) -> Self {
738 Self::build(summary, requests, false, StopReason::MaxTokens)
739 }
740
741 fn user_prompt_of(request: &ChatRequest) -> String {
742 request
743 .messages
744 .iter()
745 .find_map(|message| match &message.content {
746 Content::Text(text) => Some(text.clone()),
747 Content::Blocks(blocks) => {
748 let text = blocks
749 .iter()
750 .filter_map(|block| {
751 if let ContentBlock::Text { text } = block {
752 Some(text.as_str())
753 } else {
754 None
755 }
756 })
757 .collect::<Vec<_>>()
758 .join("\n");
759 if text.is_empty() { None } else { Some(text) }
760 }
761 })
762 .unwrap_or_default()
763 }
764 }
765
766 #[async_trait]
767 impl LlmProvider for MockProvider {
768 async fn chat(&self, request: ChatRequest) -> Result<ChatOutcome> {
769 let user_prompt = Self::user_prompt_of(&request);
770 if let Ok(mut entries) = self.requests.lock() {
771 entries.push(user_prompt.clone());
772 }
773 let text = if self.echo_input {
774 user_prompt
775 } else {
776 self.summary_response.clone()
777 };
778 Ok(ChatOutcome::Success(ChatResponse {
779 id: "test".to_string(),
780 content: vec![ContentBlock::Text { text }],
781 model: "mock".to_string(),
782 stop_reason: Some(self.stop_reason),
783 usage: Usage {
784 input_tokens: 100,
785 output_tokens: 50,
786 cached_input_tokens: 0,
787 cache_creation_input_tokens: 0,
788 },
789 }))
790 }
791
792 fn model(&self) -> &'static str {
793 "mock-model"
794 }
795
796 fn provider(&self) -> &'static str {
797 "mock"
798 }
799 }
800
801 #[test]
802 fn test_needs_compaction_below_threshold() {
803 let provider = Arc::new(MockProvider::new("summary"));
804 let config = CompactionConfig::default()
805 .with_threshold_tokens(10_000)
806 .with_min_messages(5);
807 let compactor = LlmContextCompactor::new(provider, config);
808
809 let messages = vec![
811 Message::user("Hello"),
812 Message::assistant("Hi"),
813 Message::user("How are you?"),
814 ];
815
816 assert!(!compactor.needs_compaction(&messages));
817 }
818
819 #[test]
820 fn test_needs_compaction_above_threshold() {
821 let provider = Arc::new(MockProvider::new("summary"));
822 let config = CompactionConfig::default()
823 .with_threshold_tokens(50) .with_min_messages(3);
825 let compactor = LlmContextCompactor::new(provider, config);
826
827 let messages = vec![
829 Message::user("Hello, this is a longer message to test compaction"),
830 Message::assistant(
831 "Hi there! This is also a longer response to help trigger compaction",
832 ),
833 Message::user("Great, let's continue with even more text here"),
834 Message::assistant("Absolutely, adding more content to ensure we exceed the threshold"),
835 ];
836
837 assert!(compactor.needs_compaction(&messages));
838 }
839
840 #[test]
841 fn test_needs_compaction_auto_disabled() {
842 let provider = Arc::new(MockProvider::new("summary"));
843 let config = CompactionConfig::default()
844 .with_threshold_tokens(10) .with_min_messages(1)
846 .with_auto_compact(false);
847 let compactor = LlmContextCompactor::new(provider, config);
848
849 let messages = vec![
850 Message::user("Hello, this is a longer message"),
851 Message::assistant("Response here"),
852 ];
853
854 assert!(!compactor.needs_compaction(&messages));
855 }
856
857 #[test]
858 fn summary_prompt_redacts_opaque_reasoning_payload() {
859 let secret = "opaque-secret-that-must-not-enter-the-summary-prompt";
860 let message = Message::assistant_with_content(vec![ContentBlock::OpaqueReasoning {
861 provider: "test-provider".to_owned(),
862 data: serde_json::json!({"encrypted_content": secret}),
863 }]);
864
865 let rendered = LlmContextCompactor::<MockProvider>::format_messages_for_summary([&message]);
866 assert!(rendered.contains("[Opaque reasoning state omitted]"));
867 assert!(!rendered.contains(secret));
868 }
869
870 #[tokio::test]
871 async fn compact_history_preserves_opaque_reasoning_in_retained_tail() -> Result<()> {
872 let provider = Arc::new(MockProvider::new("older context"));
873 let config = CompactionConfig::default()
874 .with_retain_recent(2)
875 .with_min_messages(3);
876 let compactor = LlmContextCompactor::new(provider, config);
877 let opaque_data = serde_json::json!({
878 "id": "reasoning_1",
879 "encrypted_content": "ciphertext"
880 });
881 let messages = vec![
882 Message::user("old question"),
883 Message::assistant("old answer"),
884 Message::user("current question"),
885 Message::assistant_with_content(vec![ContentBlock::OpaqueReasoning {
886 provider: "test-provider".to_owned(),
887 data: opaque_data.clone(),
888 }]),
889 ];
890
891 let result = compactor.compact_history(messages).await?;
892 let retained = result
893 .messages
894 .last()
895 .context("compacted history should retain the newest assistant message")?;
896 let Content::Blocks(blocks) = &retained.content else {
897 bail!("retained assistant message should contain blocks");
898 };
899 assert!(matches!(
900 blocks.first(),
901 Some(ContentBlock::OpaqueReasoning { provider, data })
902 if provider == "test-provider" && data == &opaque_data
903 ));
904 Ok(())
905 }
906
907 #[tokio::test]
908 async fn compact_history_bypasses_opaque_reasoning_before_the_split() -> Result<()> {
909 let requests = Arc::new(Mutex::new(Vec::new()));
910 let provider = Arc::new(MockProvider::new_with_request_log(
911 "summary that must not be requested",
912 Arc::clone(&requests),
913 ));
914 let config = CompactionConfig::default()
915 .with_retain_recent(1)
916 .with_min_messages(1)
917 .with_threshold_tokens(1);
918 let compactor = LlmContextCompactor::new(provider, config);
919 let opaque_data = serde_json::json!({
920 "type": "reasoning",
921 "id": "rs_1",
922 "encrypted_content": "ciphertext"
923 });
924 let messages = vec![
925 Message::user("older user context"),
926 Message::assistant_with_content(vec![
929 ContentBlock::OpaqueReasoning {
930 provider: "openai-responses".to_owned(),
931 data: opaque_data,
932 },
933 ContentBlock::Text {
934 text: "older assistant response".to_owned(),
935 },
936 ]),
937 Message::user("newer user context"),
938 Message::assistant("newer assistant response"),
939 ];
940 let serialized_before = serde_json::to_value(&messages)?;
941
942 assert!(
943 !compactor.needs_compaction(&messages),
944 "generic compaction must not be scheduled for opaque replay state"
945 );
946
947 let result = compactor.compact_history(messages).await?;
948
949 assert_eq!(serde_json::to_value(&result.messages)?, serialized_before);
950 assert_eq!(result.new_count, result.original_count);
951 assert_eq!(result.new_tokens, result.original_tokens);
952 let recorded = requests
953 .lock()
954 .map_err(|_| anyhow::anyhow!("request log poisoned"))?;
955 assert!(
956 recorded.is_empty(),
957 "opaque replay state must not be sent to the generic summarizer"
958 );
959 drop(recorded);
960 Ok(())
961 }
962
963 #[tokio::test]
964 async fn test_compact_history() -> Result<()> {
965 let provider = Arc::new(MockProvider::new(
966 "User asked about Rust programming. Assistant explained ownership, borrowing, and lifetimes.",
967 ));
968 let config = CompactionConfig::default()
969 .with_retain_recent(2)
970 .with_min_messages(3);
971 let compactor = LlmContextCompactor::new(provider, config);
972
973 let messages = vec![
975 Message::user(
976 "What is Rust? I've heard it's a systems programming language but I don't know much about it. Can you explain the key features and why people are excited about it?",
977 ),
978 Message::assistant(
979 "Rust is a systems programming language focused on safety, speed, and concurrency. It achieves memory safety without garbage collection through its ownership system. The key features include zero-cost abstractions, guaranteed memory safety, threads without data races, and minimal runtime.",
980 ),
981 Message::user(
982 "Tell me about ownership in detail. How does it work and what are the rules? I want to understand this core concept thoroughly.",
983 ),
984 Message::assistant(
985 "Ownership is Rust's central feature with three rules: each value has one owner, only one owner at a time, and the value is dropped when owner goes out of scope. This system prevents memory leaks, double frees, and dangling pointers at compile time.",
986 ),
987 Message::user("What about borrowing?"), Message::assistant("Borrowing allows references to data without taking ownership."), ];
990
991 let result = compactor.compact_history(messages).await?;
992
993 assert_eq!(result.new_count, 4);
995 assert_eq!(result.original_count, 6);
996
997 assert!(
999 result.new_tokens < result.original_tokens,
1000 "Expected fewer tokens after compaction: new={} < original={}",
1001 result.new_tokens,
1002 result.original_tokens
1003 );
1004
1005 if let Content::Text(text) = &result.messages[0].content {
1007 assert!(text.contains("Previous conversation summary"));
1008 }
1009
1010 Ok(())
1011 }
1012
1013 #[tokio::test]
1014 async fn test_compact_history_too_few_messages() -> Result<()> {
1015 let provider = Arc::new(MockProvider::new("summary"));
1016 let config = CompactionConfig::default().with_retain_recent(5);
1017 let compactor = LlmContextCompactor::new(provider, config);
1018
1019 let messages = vec![
1021 Message::user("Hello"),
1022 Message::assistant("Hi"),
1023 Message::user("Bye"),
1024 ];
1025
1026 let result = compactor.compact_history(messages.clone()).await?;
1027
1028 assert_eq!(result.new_count, 3);
1030 assert_eq!(result.messages.len(), 3);
1031
1032 Ok(())
1033 }
1034
1035 #[test]
1036 fn test_format_messages_for_summary() {
1037 let messages = vec![Message::user("Hello"), Message::assistant("Hi there!")];
1038
1039 let formatted = LlmContextCompactor::<MockProvider>::format_messages_for_summary(&messages);
1040
1041 assert!(formatted.contains("User: Hello"));
1042 assert!(formatted.contains("Assistant: Hi there!"));
1043 }
1044
1045 #[test]
1046 fn test_format_messages_for_summary_truncates_tool_results_unicode_safely() {
1047 let long_unicode = "é".repeat(600);
1048
1049 let messages = vec![Message {
1050 role: Role::Assistant,
1051 content: Content::Blocks(vec![ContentBlock::ToolResult {
1052 tool_use_id: "tool-1".to_string(),
1053 content: long_unicode,
1054 is_error: Some(false),
1055 }]),
1056 }];
1057
1058 let formatted = LlmContextCompactor::<MockProvider>::format_messages_for_summary(&messages);
1059
1060 assert!(formatted.contains("... (truncated)"));
1061 }
1062
1063 #[tokio::test]
1064 async fn test_compact_carries_prior_summary_into_request() -> Result<()> {
1065 let requests = Arc::new(Mutex::new(Vec::new()));
1070 let provider = Arc::new(MockProvider::new_with_request_log(
1071 "Fresh summary",
1072 requests.clone(),
1073 ));
1074 let config = CompactionConfig::default().with_min_messages(1);
1075 let compactor = LlmContextCompactor::new(provider, config);
1076
1077 let messages = vec![
1078 Message::user(format!("{SUMMARY_PREFIX}already compacted context")),
1079 Message::assistant("Continue with the next task using this context."),
1080 ];
1081
1082 let summary = compactor.compact(&messages).await?;
1083
1084 let recorded = requests
1085 .lock()
1086 .map_err(|_| anyhow::anyhow!("request log poisoned"))?;
1087 assert_eq!(recorded.len(), 1);
1088 assert_eq!(summary, "Fresh summary");
1091 assert!(recorded[0].contains("Continue with the next task using this context."));
1092 assert!(
1093 recorded[0].contains("already compacted context"),
1094 "prior summary must be carried into the summarization input"
1095 );
1096 drop(recorded);
1097
1098 Ok(())
1099 }
1100
1101 #[tokio::test]
1102 async fn test_compact_history_carries_prior_summary_in_candidate_payload() -> Result<()> {
1103 let requests = Arc::new(Mutex::new(Vec::new()));
1104 let provider = Arc::new(MockProvider::new_with_request_log(
1105 "Fresh history summary",
1106 requests.clone(),
1107 ));
1108 let config = CompactionConfig::default()
1109 .with_retain_recent(2)
1110 .with_min_messages(1);
1111 let compactor = LlmContextCompactor::new(provider, config);
1112
1113 let messages = vec![
1114 Message::user(format!("{SUMMARY_PREFIX}already compacted context")),
1115 Message::assistant("Current turn content from the latest exchange."),
1116 Message::assistant("Recent message that should stay."),
1117 Message::user("Newest note that should stay."),
1118 ];
1119
1120 let result = compactor.compact_history(messages).await?;
1121
1122 let recorded = requests
1123 .lock()
1124 .map_err(|_| anyhow::anyhow!("request log poisoned"))?;
1125 assert_eq!(recorded.len(), 1);
1126 assert!(recorded[0].contains("Current turn content from the latest exchange."));
1127 assert!(
1130 recorded[0].contains("already compacted context"),
1131 "prior summary content must reach the summarizer"
1132 );
1133 drop(recorded);
1134 assert_eq!(result.new_count, 4);
1135
1136 Ok(())
1137 }
1138
1139 #[tokio::test]
1140 async fn test_compact_history_carries_summaries_forward_when_window_has_only_summaries()
1141 -> Result<()> {
1142 let requests = Arc::new(Mutex::new(Vec::new()));
1143 let provider = Arc::new(MockProvider::new_with_request_log(
1144 "This summary should not be used",
1145 requests.clone(),
1146 ));
1147 let config = CompactionConfig::default()
1148 .with_retain_recent(2)
1149 .with_min_messages(1);
1150 let compactor = LlmContextCompactor::new(provider, config);
1151
1152 let messages = vec![
1153 Message::user(format!("{SUMMARY_PREFIX}first prior compacted section")),
1154 Message::assistant(format!("{SUMMARY_PREFIX}second prior compacted section")),
1155 Message::user(format!("{SUMMARY_PREFIX}third prior compacted section")),
1156 Message::assistant("final short note"),
1157 ];
1158
1159 let result = compactor.compact_history(messages).await?;
1160
1161 let recorded = requests
1165 .lock()
1166 .map_err(|_| anyhow::anyhow!("request log poisoned"))?;
1167 assert!(recorded.is_empty());
1168 drop(recorded);
1169 assert_eq!(result.new_count, 4);
1170 assert_eq!(result.messages.len(), 4);
1171
1172 if let Content::Text(text) = &result.messages[0].content {
1173 assert!(
1174 text.contains("first prior compacted section"),
1175 "first prior summary lost"
1176 );
1177 assert!(
1178 text.contains("second prior compacted section"),
1179 "second prior summary lost"
1180 );
1181 assert!(!text.contains(COMPACT_EMPTY_SUMMARY));
1182 } else {
1183 panic!("Expected summary text in first message");
1184 }
1185
1186 Ok(())
1187 }
1188
1189 #[tokio::test]
1190 async fn test_compact_history_preserves_tool_use_tool_result_pairs() -> Result<()> {
1191 let provider = Arc::new(MockProvider::new("Summary of earlier conversation."));
1192 let config = CompactionConfig::default()
1193 .with_retain_recent(2)
1194 .with_min_messages(3);
1195 let compactor = LlmContextCompactor::new(provider, config);
1196
1197 let messages = vec![
1201 Message::user("What files are in the project?"),
1203 Message::assistant("Let me check that for you."),
1205 Message {
1207 role: Role::Assistant,
1208 content: Content::Blocks(vec![ContentBlock::ToolUse {
1209 id: "tool_1".to_string(),
1210 name: "list_files".to_string(),
1211 input: serde_json::json!({}),
1212 thought_signature: None,
1213 }]),
1214 },
1215 Message {
1217 role: Role::User,
1218 content: Content::Blocks(vec![ContentBlock::ToolResult {
1219 tool_use_id: "tool_1".to_string(),
1220 content: "file1.rs\nfile2.rs".to_string(),
1221 is_error: None,
1222 }]),
1223 },
1224 Message::assistant("The project contains file1.rs and file2.rs."),
1226 ];
1227
1228 let result = compactor.compact_history(messages).await?;
1229
1230 assert_eq!(result.new_count, 5);
1234
1235 let kept_assistant = &result.messages[2];
1238 if let Content::Blocks(blocks) = &kept_assistant.content {
1239 assert!(
1240 blocks
1241 .iter()
1242 .any(|b| matches!(b, ContentBlock::ToolUse { .. })),
1243 "Expected assistant tool_use in kept messages"
1244 );
1245 } else {
1246 panic!("Expected Blocks content for assistant tool_use message");
1247 }
1248
1249 let kept_user = &result.messages[3];
1251 if let Content::Blocks(blocks) = &kept_user.content {
1252 assert!(
1253 blocks
1254 .iter()
1255 .any(|b| matches!(b, ContentBlock::ToolResult { .. })),
1256 "Expected user tool_result in kept messages"
1257 );
1258 } else {
1259 panic!("Expected Blocks content for user tool_result message");
1260 }
1261
1262 Ok(())
1263 }
1264
1265 #[tokio::test]
1266 async fn test_compact_history_split_skips_leading_orphan_after_summary_ack() -> Result<()> {
1267 let provider = Arc::new(MockProvider::new("Re-summary."));
1295 let config = CompactionConfig::default()
1296 .with_retain_recent(3)
1297 .with_min_messages(1);
1298 let compactor = LlmContextCompactor::new(provider, config);
1299
1300 let messages = vec![
1301 Message::user(format!("{SUMMARY_PREFIX}Old summary about toolu_X.")),
1302 Message::assistant(SUMMARY_ACKNOWLEDGMENT),
1303 Message {
1304 role: Role::User,
1305 content: Content::Blocks(vec![ContentBlock::ToolResult {
1306 tool_use_id: "toolu_X".to_string(),
1307 content: "result for X".to_string(),
1308 is_error: None,
1309 }]),
1310 },
1311 Message::assistant("Result interpreted."),
1312 Message::user("Now what?"),
1313 ];
1314
1315 let result = compactor.compact_history(messages).await?;
1316
1317 let mut seen_ids: std::collections::HashSet<String> = std::collections::HashSet::new();
1318 for msg in &result.messages {
1319 if let Content::Blocks(blocks) = &msg.content {
1320 for block in blocks {
1321 match block {
1322 ContentBlock::ToolResult { tool_use_id, .. } => {
1323 assert!(
1324 seen_ids.contains(tool_use_id),
1325 "orphan tool_use_id {tool_use_id} survived split selection",
1326 );
1327 }
1328 ContentBlock::ToolUse { id, .. } => {
1329 seen_ids.insert(id.clone());
1330 }
1331 _ => {}
1332 }
1333 }
1334 }
1335 }
1336
1337 Ok(())
1338 }
1339
1340 #[tokio::test]
1341 async fn test_compact_history_keeps_tool_pair_when_immediate_prev_is_text_only() -> Result<()> {
1342 let provider = Arc::new(MockProvider::new("Boundary summary."));
1349 let config = CompactionConfig::default()
1350 .with_retain_recent(2)
1351 .with_min_messages(1);
1352 let compactor = LlmContextCompactor::new(provider, config);
1353
1354 let messages = vec![
1367 Message::user("first turn"),
1368 Message::assistant("text only"),
1369 Message {
1370 role: Role::User,
1371 content: Content::Blocks(vec![ContentBlock::ToolResult {
1372 tool_use_id: "toolu_Y".to_string(),
1373 content: "ancient result".to_string(),
1374 is_error: None,
1375 }]),
1376 },
1377 Message::assistant("then a reply"),
1378 Message::user("ok thanks"),
1379 ];
1380
1381 let result = compactor.compact_history(messages).await?;
1382
1383 let has_tool_result = result.messages.iter().any(|m| {
1387 matches!(
1388 &m.content,
1389 Content::Blocks(blocks)
1390 if blocks.iter().any(|b| matches!(b, ContentBlock::ToolResult { .. }))
1391 )
1392 });
1393 assert!(
1394 !has_tool_result,
1395 "orphan tool_result should have been pushed into to_summarize, not retained",
1396 );
1397
1398 Ok(())
1399 }
1400
1401 #[tokio::test]
1402 async fn test_compact_history_retained_tail_is_token_capped() -> Result<()> {
1403 let provider = Arc::new(MockProvider::new(
1404 "Project summary with a long context and technical context.",
1405 ));
1406 let config = CompactionConfig::default()
1407 .with_retain_recent(8)
1408 .with_min_messages(1)
1409 .with_threshold_tokens(1);
1410 let compactor = LlmContextCompactor::new(provider, config);
1411
1412 let mut messages = Vec::new();
1413
1414 messages.extend((0..6).map(|index| Message::user(format!("pre-compaction noise {index}"))));
1416
1417 messages.extend(
1419 (0..8).map(|index| Message::assistant(format!("kept-{index}: {}", "x".repeat(12_000)))),
1420 );
1421
1422 let result = compactor.compact_history(messages).await?;
1423
1424 let retained_tail = &result.messages[2..];
1426 assert!(retained_tail.len() < 8);
1427
1428 let mut latest_index = -1i32;
1429 let mut all_retained = true;
1430 for message in retained_tail {
1431 if let Content::Text(text) = &message.content {
1432 if let Some(number) = text.split(':').next().and_then(|prefix| {
1433 prefix
1434 .strip_prefix("kept-")
1435 .and_then(|rest| rest.parse::<i32>().ok())
1436 }) {
1437 if number >= 0 {
1438 latest_index = latest_index.max(number);
1439 }
1440 } else {
1441 all_retained = false;
1442 }
1443 } else {
1444 all_retained = false;
1445 }
1446 }
1447
1448 assert!(all_retained);
1449 assert_eq!(latest_index, 7);
1450 assert!(
1451 TokenEstimator::estimate_history(retained_tail)
1452 <= compactor.config().max_retained_tail_tokens
1453 );
1454 assert!(compactor.needs_compaction(&result.messages));
1455
1456 Ok(())
1457 }
1458
1459 #[tokio::test]
1460 async fn test_compact_history_skips_summary_ack_when_retained_tail_is_empty() -> Result<()> {
1461 let provider = Arc::new(MockProvider::new("Summary for oversized user turn."));
1462 let config = CompactionConfig::default()
1463 .with_retain_recent(1)
1464 .with_min_messages(1)
1465 .with_threshold_tokens(1);
1466 let compactor = LlmContextCompactor::new(provider, config);
1467
1468 let messages = vec![
1469 Message::assistant("Earlier assistant context."),
1470 Message::user(format!("oversized-user-turn: {}", "x".repeat(200_000))),
1471 ];
1472
1473 let result = compactor.compact_history(messages).await?;
1474
1475 assert_eq!(result.new_count, 1);
1476 assert_eq!(result.messages.len(), 1);
1477
1478 let only_message = &result.messages[0];
1479 assert_eq!(only_message.role, Role::User);
1480
1481 if let Content::Text(text) = &only_message.content {
1482 assert!(text.contains("Previous conversation summary"));
1483 assert!(!text.contains(SUMMARY_ACKNOWLEDGMENT));
1484 } else {
1485 panic!("Expected summary text when retained tail is empty");
1486 }
1487
1488 Ok(())
1489 }
1490
1491 fn message_contains(message: &Message, needle: &str) -> bool {
1492 match &message.content {
1493 Content::Text(text) => text.contains(needle),
1494 Content::Blocks(blocks) => blocks.iter().any(|block| match block {
1495 ContentBlock::Text { text } => text.contains(needle),
1496 _ => false,
1497 }),
1498 }
1499 }
1500
1501 #[tokio::test]
1502 async fn test_epoch_one_facts_survive_two_compactions() -> Result<()> {
1503 const EPOCH1_FACT: &str = "EPOCH1_FACT: the API key lives in config/secrets.toml";
1509
1510 let requests = Arc::new(Mutex::new(Vec::new()));
1511 let provider = Arc::new(MockProvider::new_echo(requests.clone()));
1512 let config = CompactionConfig::default()
1513 .with_retain_recent(2)
1514 .with_min_messages(1);
1515 let compactor = LlmContextCompactor::new(provider, config);
1516
1517 let epoch1 = vec![
1518 Message::user(EPOCH1_FACT),
1519 Message::assistant("Understood, noted the secrets path."),
1520 Message::user("Now add error handling to main.rs."),
1521 Message::assistant("Added error handling to main.rs."),
1522 Message::user("latest user message one"),
1523 Message::assistant("latest assistant message two"),
1524 ];
1525
1526 let first = compactor.compact_history(epoch1).await?;
1527 assert!(
1528 first
1529 .messages
1530 .iter()
1531 .any(|m| message_contains(m, "EPOCH1_FACT")),
1532 "epoch-1 fact must be captured in the first summary"
1533 );
1534
1535 let mut epoch2 = first.messages;
1537 epoch2.push(Message::user("Another later turn."));
1538 epoch2.push(Message::assistant("Reply to the later turn."));
1539 epoch2.push(Message::user("Final turn a."));
1540 epoch2.push(Message::assistant("Final turn b."));
1541
1542 let second = compactor.compact_history(epoch2).await?;
1543
1544 assert!(
1545 second
1546 .messages
1547 .iter()
1548 .any(|m| message_contains(m, "EPOCH1_FACT")),
1549 "epoch-1 fact must survive the second compaction"
1550 );
1551
1552 let recorded = requests
1555 .lock()
1556 .map_err(|_| anyhow::anyhow!("request log poisoned"))?;
1557 assert!(
1558 recorded.iter().any(|req| req.contains("EPOCH1_FACT")),
1559 "prior summary carrying the epoch-1 fact must reach the summarizer"
1560 );
1561 drop(recorded);
1562
1563 Ok(())
1564 }
1565
1566 #[tokio::test]
1567 async fn test_compact_history_long_tool_chain_respects_token_cap() -> Result<()> {
1568 let provider = Arc::new(MockProvider::new("Summary of the early tool chain."));
1574 let cap = 20_000;
1575 let config = CompactionConfig::default()
1578 .with_retain_recent(18)
1579 .with_min_messages(1)
1580 .with_threshold_tokens(1)
1581 .with_max_retained_tail_tokens(cap);
1582 let compactor = LlmContextCompactor::new(provider, config);
1583
1584 let mut messages = Vec::new();
1587 for i in 0..10 {
1588 messages.push(Message {
1589 role: Role::Assistant,
1590 content: Content::Blocks(vec![ContentBlock::ToolUse {
1591 id: format!("tool_{i}"),
1592 name: "run".to_string(),
1593 input: serde_json::json!({ "arg": "y".repeat(12_000) }),
1594 thought_signature: None,
1595 }]),
1596 });
1597 messages.push(Message {
1598 role: Role::User,
1599 content: Content::Blocks(vec![ContentBlock::ToolResult {
1600 tool_use_id: format!("tool_{i}"),
1601 content: format!("result-{i}: {}", "z".repeat(12_000)),
1602 is_error: None,
1603 }]),
1604 });
1605 }
1606
1607 let full_tokens = TokenEstimator::estimate_history(&messages);
1608 assert!(
1609 full_tokens > cap * 2,
1610 "test setup: full chain must far exceed the cap"
1611 );
1612
1613 let result = compactor.compact_history(messages).await?;
1614
1615 let retained_tail = &result.messages[2..];
1618
1619 let tail_tokens = TokenEstimator::estimate_history(retained_tail);
1620 assert!(
1623 tail_tokens <= cap + 8_000,
1624 "retained tail {tail_tokens} should be bounded by the cap {cap}, not the whole chain"
1625 );
1626 assert!(
1627 retained_tail.len() < 20,
1628 "compaction must have summarized part of the chain"
1629 );
1630
1631 Ok(())
1632 }
1633
1634 #[tokio::test]
1635 async fn test_compact_warns_and_marks_truncated_summary() -> Result<()> {
1636 let requests = Arc::new(Mutex::new(Vec::new()));
1641 let provider = Arc::new(MockProvider::new_truncating(
1642 "partial summary cut off mid-",
1643 requests.clone(),
1644 ));
1645 let config = CompactionConfig::default().with_min_messages(1);
1646 let compactor = LlmContextCompactor::new(provider, config);
1647
1648 let messages = vec![
1649 Message::user("Some content that needs summarizing."),
1650 Message::assistant("More content to summarize here."),
1651 ];
1652
1653 let summary = compactor.compact(&messages).await?;
1654
1655 assert!(
1656 summary.contains("[summary truncated"),
1657 "a persistently truncated summary must carry a truncation marker"
1658 );
1659
1660 let recorded = requests
1662 .lock()
1663 .map_err(|_| anyhow::anyhow!("request log poisoned"))?;
1664 assert_eq!(recorded.len(), 2, "truncation should trigger one retry");
1665 drop(recorded);
1666
1667 Ok(())
1668 }
1669}