1use crate::hooks::{AgentHooks, DefaultHooks, RequestDecision, ResponseDecision};
4use crate::llm::{
5 ChatOutcome, ChatRequest, Content, ContentBlock, LlmProvider, Message, Role, StopReason,
6};
7use crate::types::TokenUsage;
8use anyhow::{Context, Result};
9use async_trait::async_trait;
10use std::fmt::Write;
11use std::sync::Arc;
12
13use super::config::CompactionConfig;
14use super::estimator::TokenEstimator;
15
16const SUMMARY_PREFIX: &str = "[Previous conversation summary]\n\n";
17const 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.";
18const 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";
19const COMPACTION_SUMMARY_PROMPT_SUFFIX: &str =
20 "Provide a concise summary (aim for 500-1000 words):";
21const COMPACT_EMPTY_SUMMARY: &str = "No additional context was available to summarize; the previous messages were already compacted.";
22const SUMMARY_ACKNOWLEDGMENT: &str =
23 "I understand the context from the summary. Let me continue from where we left off.";
24const MAX_TOOL_RESULT_CHARS: usize = 500;
25const TRUNCATED_SUMMARY_MARKER: &str =
26 "\n\n[summary truncated: exceeded the configured summary_max_tokens budget]";
27
28#[async_trait]
32pub trait ContextCompactor: Send + Sync {
33 async fn compact(&self, messages: &[Message]) -> Result<String>;
38
39 fn estimate_tokens(&self, messages: &[Message]) -> usize;
41
42 fn needs_compaction(&self, messages: &[Message]) -> bool;
44
45 async fn compact_history(&self, messages: Vec<Message>) -> Result<CompactionResult>;
50
51 async fn compact_history_with_usage(
63 &self,
64 messages: Vec<Message>,
65 ) -> Result<CompactionResult, FailedCompaction> {
66 self.compact_history(messages)
67 .await
68 .map_err(|error| FailedCompaction {
69 error,
70 llm_usage: TokenUsage::default(),
71 })
72 }
73}
74
75#[derive(Debug, Clone)]
77pub struct CompactionResult {
78 pub messages: Vec<Message>,
80 pub original_count: usize,
82 pub new_count: usize,
84 pub original_tokens: usize,
86 pub new_tokens: usize,
88 pub llm_usage: TokenUsage,
95}
96
97#[derive(Debug)]
105pub struct FailedCompaction {
106 pub error: anyhow::Error,
108 pub llm_usage: TokenUsage,
111}
112
113pub struct LlmContextCompactor<P: LlmProvider + ?Sized, H: AgentHooks = DefaultHooks> {
136 provider: Arc<P>,
137 config: CompactionConfig,
138 hooks: Option<Arc<H>>,
143 system_prompt: String,
144 summary_prompt_prefix: String,
145 summary_prompt_suffix: String,
146}
147
148impl<P: LlmProvider + ?Sized> LlmContextCompactor<P> {
149 #[must_use]
151 pub fn new(provider: Arc<P>, config: CompactionConfig) -> Self {
152 Self {
153 provider,
154 config,
155 hooks: None,
156 system_prompt: COMPACTION_SYSTEM_PROMPT.to_string(),
157 summary_prompt_prefix: COMPACTION_SUMMARY_PROMPT_PREFIX.to_string(),
158 summary_prompt_suffix: COMPACTION_SUMMARY_PROMPT_SUFFIX.to_string(),
159 }
160 }
161
162 #[must_use]
164 pub fn with_defaults(provider: Arc<P>) -> Self {
165 Self::new(provider, CompactionConfig::default())
166 }
167}
168
169impl<P: LlmProvider + ?Sized, H: AgentHooks> LlmContextCompactor<P, H> {
170 #[must_use]
182 pub fn with_guardrail_hooks<H2: AgentHooks>(
183 self,
184 hooks: Arc<H2>,
185 ) -> LlmContextCompactor<P, H2> {
186 LlmContextCompactor {
187 provider: self.provider,
188 config: self.config,
189 hooks: Some(hooks),
190 system_prompt: self.system_prompt,
191 summary_prompt_prefix: self.summary_prompt_prefix,
192 summary_prompt_suffix: self.summary_prompt_suffix,
193 }
194 }
195
196 #[must_use]
198 pub const fn config(&self) -> &CompactionConfig {
199 &self.config
200 }
201
202 #[must_use]
204 pub fn with_prompts(
205 mut self,
206 system_prompt: impl Into<String>,
207 summary_prompt_prefix: impl Into<String>,
208 summary_prompt_suffix: impl Into<String>,
209 ) -> Self {
210 self.system_prompt = system_prompt.into();
211 self.summary_prompt_prefix = summary_prompt_prefix.into();
212 self.summary_prompt_suffix = summary_prompt_suffix.into();
213 self
214 }
215
216 fn extract_summary_text(content: &Content) -> Option<String> {
226 match content {
227 Content::Text(text) => text.strip_prefix(SUMMARY_PREFIX).map(str::to_string),
228 Content::Blocks(blocks) => blocks.iter().find_map(|block| match block {
229 ContentBlock::Text { text } => {
230 text.strip_prefix(SUMMARY_PREFIX).map(str::to_string)
231 }
232 _ => None,
233 }),
234 }
235 }
236
237 fn has_tool_use(content: &Content) -> bool {
239 matches!(
240 content,
241 Content::Blocks(blocks)
242 if blocks
243 .iter()
244 .any(|block| matches!(block, ContentBlock::ToolUse { .. }))
245 )
246 }
247
248 fn has_tool_result(content: &Content) -> bool {
250 matches!(
251 content,
252 Content::Blocks(blocks)
253 if blocks
254 .iter()
255 .any(|block| matches!(block, ContentBlock::ToolResult { .. }))
256 )
257 }
258
259 fn split_point_preserves_tool_pairs(messages: &[Message], mut split_point: usize) -> usize {
272 while split_point > 0 && split_point < messages.len() {
273 let prev = &messages[split_point - 1];
274 let next = &messages[split_point];
275
276 let crosses_tool_pair = prev.role == Role::Assistant
277 && Self::has_tool_use(&prev.content)
278 && next.role == Role::User
279 && Self::has_tool_result(&next.content);
280
281 if crosses_tool_pair {
282 split_point -= 1;
283 continue;
284 }
285
286 break;
287 }
288
289 split_point
290 }
291
292 fn split_point_preserves_tool_pairs_with_cap(
331 messages: &[Message],
332 split_point: usize,
333 max_tokens: usize,
334 ) -> usize {
335 let cap_limit = Self::retain_tail_with_token_cap(messages, split_point, max_tokens);
336 let pair_safe = Self::split_point_preserves_tool_pairs(messages, cap_limit);
337 Self::split_point_skips_leading_orphan(messages, pair_safe)
338 }
339
340 fn split_point_skips_leading_orphan(messages: &[Message], mut split_point: usize) -> usize {
357 while split_point < messages.len() {
358 if Self::leading_message_has_orphan_tool_result(&messages[split_point..]) {
359 split_point = split_point.saturating_add(1);
360 continue;
361 }
362 break;
363 }
364 split_point
365 }
366
367 fn leading_message_has_orphan_tool_result(to_keep: &[Message]) -> bool {
375 let Some(first) = to_keep.first() else {
376 return false;
377 };
378 let Content::Blocks(blocks) = &first.content else {
379 return false;
380 };
381
382 let mut needed: Vec<&str> = Vec::new();
386 for block in blocks {
387 if let ContentBlock::ToolResult { tool_use_id, .. } = block {
388 needed.push(tool_use_id.as_str());
389 }
390 }
391 if needed.is_empty() {
392 return false;
393 }
394
395 let known_ids: std::collections::HashSet<&str> = to_keep
397 .iter()
398 .flat_map(|message| match &message.content {
399 Content::Blocks(blocks) => blocks
400 .iter()
401 .filter_map(|block| match block {
402 ContentBlock::ToolUse { id, .. } => Some(id.as_str()),
403 _ => None,
404 })
405 .collect::<Vec<_>>(),
406 Content::Text(_) => Vec::new(),
407 })
408 .collect();
409
410 needed.iter().any(|id| !known_ids.contains(id))
411 }
412
413 fn retain_tail_with_token_cap(messages: &[Message], start: usize, max_tokens: usize) -> usize {
415 if start >= messages.len() {
416 return messages.len();
417 }
418
419 if max_tokens == 0 {
420 return messages.len();
421 }
422
423 let mut used = 0usize;
424 let mut retained_start = messages.len();
425
426 for idx in (start..messages.len()).rev() {
427 let message_tokens = TokenEstimator::estimate_message(&messages[idx]);
428 if used + message_tokens > max_tokens {
429 break;
430 }
431
432 retained_start = idx;
433 used += message_tokens;
434 }
435
436 retained_start
437 }
438
439 fn format_messages_for_summary<'a>(messages: impl IntoIterator<Item = &'a Message>) -> String {
444 let mut output = String::new();
445
446 for message in messages {
447 let role = match message.role {
448 Role::User => "User",
449 Role::Assistant => "Assistant",
450 };
451
452 let _ = write!(output, "{role}: ");
453
454 match &message.content {
455 Content::Text(text) => {
456 let _ = writeln!(output, "{text}");
457 }
458 Content::Blocks(blocks) => {
459 for block in blocks {
460 match block {
461 ContentBlock::Text { text } => {
462 let _ = writeln!(output, "{text}");
463 }
464 ContentBlock::Thinking { thinking, .. } => {
465 let _ = writeln!(output, "[Thinking: {thinking}]");
467 }
468 ContentBlock::RedactedThinking { .. } => {
469 let _ = writeln!(output, "[Redacted thinking]");
470 }
471 ContentBlock::OpaqueReasoning { .. } => {
472 let _ = writeln!(output, "[Opaque reasoning state omitted]");
477 }
478 ContentBlock::ToolUse { name, input, .. } => {
479 let _ = writeln!(
480 output,
481 "[Called tool: {name} with input: {}]",
482 serde_json::to_string(input).unwrap_or_default()
483 );
484 }
485 ContentBlock::ToolResult {
486 content, is_error, ..
487 } => {
488 let status = if is_error.unwrap_or(false) {
489 "error"
490 } else {
491 "success"
492 };
493 let truncated = if content.chars().count() > MAX_TOOL_RESULT_CHARS {
495 let prefix: String =
496 content.chars().take(MAX_TOOL_RESULT_CHARS).collect();
497 format!("{prefix}... (truncated)")
498 } else {
499 content.clone()
500 };
501 let _ = writeln!(output, "[Tool result ({status}): {truncated}]");
502 }
503 ContentBlock::Image { source } => {
504 let _ = writeln!(output, "[Image: {}]", source.media_type);
505 }
506 ContentBlock::Document { source } => {
507 let _ = writeln!(output, "[Document: {}]", source.media_type);
508 }
509 _ => {
512 let _ = writeln!(output, "[Unrecognized content block]");
513 }
514 }
515 }
516 }
517 }
518 output.push('\n');
519 }
520
521 output
522 }
523
524 fn build_summary_prompt(&self, prior_summaries: &[String], messages_text: &str) -> String {
531 let base = format!(
532 "{}{}{}",
533 self.summary_prompt_prefix, messages_text, self.summary_prompt_suffix
534 );
535
536 if prior_summaries.is_empty() {
537 return base;
538 }
539
540 let prior = prior_summaries.join("\n\n");
541 format!(
542 "Previous summary of earlier conversation. Preserve every fact below \
543 in your new summary so no earlier context is lost:\n{prior}\n\n{base}"
544 )
545 }
546
547 async fn run_summarization(
554 &self,
555 prompt: String,
556 max_tokens: usize,
557 ) -> Result<SummarizationCall, SummarizationFailure> {
558 let mut request = ChatRequest {
559 system: self.system_prompt.clone(),
560 messages: vec![Message::user(prompt)],
561 tools: None,
562 max_tokens: u32::try_from(max_tokens).unwrap_or(u32::MAX),
563 max_tokens_explicit: true,
564 session_id: None,
565 cached_content: None,
566 thinking: None,
567 tool_choice: None,
568 response_format: None,
569 cache: None,
570 };
571
572 if let Some(hooks) = &self.hooks {
576 match hooks.pre_llm_request(&request).await {
577 RequestDecision::Modify(modified) => request = *modified,
578 RequestDecision::Block(reason) => {
579 return Err(SummarizationFailure {
580 error: anyhow::anyhow!(
581 "Summarization request blocked by guardrail: {reason}"
582 ),
583 usage: TokenUsage::default(),
584 });
585 }
586 _ => {}
589 }
590 }
591
592 let outcome = self
593 .provider
594 .chat(request)
595 .await
596 .context("Failed to call LLM for summarization")
597 .map_err(|error| SummarizationFailure {
598 error,
599 usage: TokenUsage::default(),
600 })?;
601
602 match outcome {
603 ChatOutcome::Success(response) => {
604 let usage = TokenUsage {
608 input_tokens: response.usage.input_tokens,
609 output_tokens: response.usage.output_tokens,
610 cached_input_tokens: response.usage.cached_input_tokens,
611 cache_creation_input_tokens: response.usage.cache_creation_input_tokens,
612 };
613 if let Some(hooks) = &self.hooks {
620 match hooks.on_llm_response(&response).await {
621 ResponseDecision::Block(reason) => {
622 return Err(SummarizationFailure {
623 error: anyhow::anyhow!(
624 "Summarization response blocked by guardrail: {reason}"
625 ),
626 usage,
627 });
628 }
629 ResponseDecision::RetryWithFeedback(reason) => {
630 return Err(SummarizationFailure {
631 error: anyhow::anyhow!(
632 "Summarization response rejected by guardrail \
633 (RetryWithFeedback is not retried during compaction): {reason}"
634 ),
635 usage,
636 });
637 }
638 _ => {}
641 }
642 }
643 let truncated = response.stop_reason == Some(StopReason::MaxTokens);
644 let Some(text) = response.first_text().map(String::from) else {
645 return Err(SummarizationFailure {
646 error: anyhow::anyhow!("No text in summarization response"),
647 usage,
648 });
649 };
650 Ok(SummarizationCall {
651 text,
652 truncated,
653 usage,
654 })
655 }
656 ChatOutcome::RateLimited(_) => Err(SummarizationFailure {
657 error: anyhow::anyhow!("Rate limited during summarization"),
658 usage: TokenUsage::default(),
659 }),
660 ChatOutcome::InvalidRequest(msg) => Err(SummarizationFailure {
661 error: anyhow::anyhow!("Invalid request during summarization: {msg}"),
662 usage: TokenUsage::default(),
663 }),
664 ChatOutcome::ServerError(msg) => Err(SummarizationFailure {
665 error: anyhow::anyhow!("Server error during summarization: {msg}"),
666 usage: TokenUsage::default(),
667 }),
668 _ => Err(SummarizationFailure {
671 error: anyhow::anyhow!("Unrecognized provider outcome during summarization"),
672 usage: TokenUsage::default(),
673 }),
674 }
675 }
676
677 async fn summarize_with_usage(
685 &self,
686 messages: &[Message],
687 ) -> Result<(String, TokenUsage), SummarizationFailure> {
688 let mut prior_summaries: Vec<String> = Vec::new();
693 let mut fresh: Vec<&Message> = Vec::new();
694 for message in messages {
695 if let Some(text) = Self::extract_summary_text(&message.content) {
696 if !text.is_empty() {
697 prior_summaries.push(text);
698 }
699 } else {
700 fresh.push(message);
701 }
702 }
703
704 if fresh.is_empty() {
707 if prior_summaries.is_empty() {
708 return Ok((COMPACT_EMPTY_SUMMARY.to_string(), TokenUsage::default()));
709 }
710 return Ok((prior_summaries.join("\n\n"), TokenUsage::default()));
711 }
712
713 let messages_text = Self::format_messages_for_summary(fresh.iter().copied());
714 let prompt = self.build_summary_prompt(&prior_summaries, &messages_text);
715
716 let budget = self.config.summary_max_tokens;
717 let first = self.run_summarization(prompt.clone(), budget).await?;
718 let mut summary = first.text;
719 let mut total_usage = first.usage;
720
721 if first.truncated {
722 log::warn!(
723 "compaction summary hit the max_tokens budget ({budget}); \
724 retrying with a larger budget to avoid silent context loss"
725 );
726 let retry = match self
727 .run_summarization(prompt, budget.saturating_mul(2))
728 .await
729 {
730 Ok(retry) => retry,
731 Err(mut failure) => {
732 failure.usage.add(&total_usage);
735 return Err(failure);
736 }
737 };
738 total_usage.add(&retry.usage);
739 summary = retry.text;
740 if retry.truncated {
741 log::warn!(
742 "compaction summary still truncated after retry; appending a \
743 truncation marker so downstream context loss is visible"
744 );
745 summary.push_str(TRUNCATED_SUMMARY_MARKER);
746 }
747 }
748
749 Ok((summary, total_usage))
750 }
751}
752
753struct SummarizationCall {
755 text: String,
756 truncated: bool,
757 usage: TokenUsage,
758}
759
760struct SummarizationFailure {
764 error: anyhow::Error,
765 usage: TokenUsage,
766}
767
768impl<P: LlmProvider + ?Sized, H: AgentHooks> LlmContextCompactor<P, H> {
769 async fn compact_history_inner(
772 &self,
773 mut messages: Vec<Message>,
774 ) -> Result<CompactionResult, FailedCompaction> {
775 let original_count = messages.len();
776 let original_tokens = self.estimate_tokens(&messages);
777
778 if messages.len() <= self.config.retain_recent {
788 return Ok(CompactionResult {
789 messages,
790 original_count,
791 new_count: original_count,
792 original_tokens,
793 new_tokens: original_tokens,
794 llm_usage: TokenUsage::default(),
795 });
796 }
797
798 let mut split_point = messages.len().saturating_sub(self.config.retain_recent);
800 split_point = Self::split_point_preserves_tool_pairs_with_cap(
801 &messages,
802 split_point,
803 self.config.max_retained_tail_tokens,
804 );
805
806 let to_keep = messages.split_off(split_point);
809 let to_summarize = messages;
810
811 let (summary, llm_usage) =
813 self.summarize_with_usage(&to_summarize)
814 .await
815 .map_err(|failure| FailedCompaction {
816 error: failure.error,
817 llm_usage: failure.usage,
818 })?;
819
820 let mut new_messages = Vec::with_capacity(2 + to_keep.len());
822
823 new_messages.push(Message::user(format!("{SUMMARY_PREFIX}{summary}")));
825
826 if !to_keep.is_empty() {
831 new_messages.push(Message::assistant(SUMMARY_ACKNOWLEDGMENT));
832 }
833
834 new_messages.extend(to_keep);
842
843 let new_count = new_messages.len();
844 let new_tokens = self.estimate_tokens(&new_messages);
845
846 Ok(CompactionResult {
847 messages: new_messages,
848 original_count,
849 new_count,
850 original_tokens,
851 new_tokens,
852 llm_usage,
853 })
854 }
855}
856
857#[async_trait]
858impl<P: LlmProvider + ?Sized, H: AgentHooks> ContextCompactor for LlmContextCompactor<P, H> {
859 async fn compact(&self, messages: &[Message]) -> Result<String> {
860 let (summary, _usage) = self
861 .summarize_with_usage(messages)
862 .await
863 .map_err(|failure| failure.error)?;
864 Ok(summary)
865 }
866
867 fn estimate_tokens(&self, messages: &[Message]) -> usize {
868 TokenEstimator::estimate_history(messages)
869 }
870
871 fn needs_compaction(&self, messages: &[Message]) -> bool {
872 if !self.config.auto_compact {
873 return false;
874 }
875
876 if messages.len() < self.config.min_messages_for_compaction {
877 return false;
878 }
879
880 let estimated_tokens = self.estimate_tokens(messages);
881 estimated_tokens > self.config.threshold_tokens
882 }
883
884 async fn compact_history(&self, messages: Vec<Message>) -> Result<CompactionResult> {
885 self.compact_history_inner(messages)
886 .await
887 .map_err(|failure| failure.error)
888 }
889
890 async fn compact_history_with_usage(
891 &self,
892 messages: Vec<Message>,
893 ) -> Result<CompactionResult, FailedCompaction> {
894 self.compact_history_inner(messages).await
895 }
896}
897
898#[cfg(test)]
899mod tests {
900 use super::*;
901 use crate::llm::{ChatResponse, StopReason, Usage};
902 use anyhow::bail;
903 use std::sync::Mutex;
904
905 struct MockProvider {
906 summary_response: String,
907 requests: Arc<Mutex<Vec<String>>>,
908 echo_input: bool,
911 stop_reason: StopReason,
913 }
914
915 impl MockProvider {
916 fn build(
917 summary: &str,
918 requests: Arc<Mutex<Vec<String>>>,
919 echo_input: bool,
920 stop_reason: StopReason,
921 ) -> Self {
922 Self {
923 summary_response: summary.to_string(),
924 requests,
925 echo_input,
926 stop_reason,
927 }
928 }
929
930 fn new(summary: &str) -> Self {
931 Self::build(
932 summary,
933 Arc::new(Mutex::new(Vec::new())),
934 false,
935 StopReason::EndTurn,
936 )
937 }
938
939 fn new_with_request_log(summary: &str, requests: Arc<Mutex<Vec<String>>>) -> Self {
940 Self::build(summary, requests, false, StopReason::EndTurn)
941 }
942
943 fn new_echo(requests: Arc<Mutex<Vec<String>>>) -> Self {
945 Self::build("", requests, true, StopReason::EndTurn)
946 }
947
948 fn new_truncating(summary: &str, requests: Arc<Mutex<Vec<String>>>) -> Self {
950 Self::build(summary, requests, false, StopReason::MaxTokens)
951 }
952
953 fn user_prompt_of(request: &ChatRequest) -> String {
954 request
955 .messages
956 .iter()
957 .find_map(|message| match &message.content {
958 Content::Text(text) => Some(text.clone()),
959 Content::Blocks(blocks) => {
960 let text = blocks
961 .iter()
962 .filter_map(|block| {
963 if let ContentBlock::Text { text } = block {
964 Some(text.as_str())
965 } else {
966 None
967 }
968 })
969 .collect::<Vec<_>>()
970 .join("\n");
971 if text.is_empty() { None } else { Some(text) }
972 }
973 })
974 .unwrap_or_default()
975 }
976 }
977
978 #[async_trait]
979 impl LlmProvider for MockProvider {
980 async fn chat(&self, request: ChatRequest) -> Result<ChatOutcome> {
981 let user_prompt = Self::user_prompt_of(&request);
982 if let Ok(mut entries) = self.requests.lock() {
983 entries.push(user_prompt.clone());
984 }
985 let text = if self.echo_input {
986 user_prompt
987 } else {
988 self.summary_response.clone()
989 };
990 Ok(ChatOutcome::Success(ChatResponse {
991 id: "test".to_string(),
992 content: vec![ContentBlock::Text { text }],
993 model: "mock".to_string(),
994 stop_reason: Some(self.stop_reason),
995 usage: Usage {
996 served_speed: None,
997 input_tokens: 100,
998 output_tokens: 50,
999 cached_input_tokens: 0,
1000 cache_creation_input_tokens: 0,
1001 },
1002 }))
1003 }
1004
1005 fn model(&self) -> &'static str {
1006 "mock-model"
1007 }
1008
1009 fn provider(&self) -> &'static str {
1010 "mock"
1011 }
1012 }
1013
1014 #[test]
1015 fn test_needs_compaction_below_threshold() {
1016 let provider = Arc::new(MockProvider::new("summary"));
1017 let config = CompactionConfig::default()
1018 .with_threshold_tokens(10_000)
1019 .with_min_messages(5);
1020 let compactor = LlmContextCompactor::new(provider, config);
1021
1022 let messages = vec![
1024 Message::user("Hello"),
1025 Message::assistant("Hi"),
1026 Message::user("How are you?"),
1027 ];
1028
1029 assert!(!compactor.needs_compaction(&messages));
1030 }
1031
1032 #[test]
1033 fn test_needs_compaction_above_threshold() {
1034 let provider = Arc::new(MockProvider::new("summary"));
1035 let config = CompactionConfig::default()
1036 .with_threshold_tokens(50) .with_min_messages(3);
1038 let compactor = LlmContextCompactor::new(provider, config);
1039
1040 let messages = vec![
1042 Message::user("Hello, this is a longer message to test compaction"),
1043 Message::assistant(
1044 "Hi there! This is also a longer response to help trigger compaction",
1045 ),
1046 Message::user("Great, let's continue with even more text here"),
1047 Message::assistant("Absolutely, adding more content to ensure we exceed the threshold"),
1048 ];
1049
1050 assert!(compactor.needs_compaction(&messages));
1051 }
1052
1053 #[test]
1054 fn test_needs_compaction_auto_disabled() {
1055 let provider = Arc::new(MockProvider::new("summary"));
1056 let config = CompactionConfig::default()
1057 .with_threshold_tokens(10) .with_min_messages(1)
1059 .with_auto_compact(false);
1060 let compactor = LlmContextCompactor::new(provider, config);
1061
1062 let messages = vec![
1063 Message::user("Hello, this is a longer message"),
1064 Message::assistant("Response here"),
1065 ];
1066
1067 assert!(!compactor.needs_compaction(&messages));
1068 }
1069
1070 #[test]
1071 fn summary_prompt_redacts_opaque_reasoning_payload() {
1072 let secret = "opaque-secret-that-must-not-enter-the-summary-prompt";
1073 let message = Message::assistant_with_content(vec![ContentBlock::OpaqueReasoning {
1074 provider: "test-provider".to_owned(),
1075 data: serde_json::json!({"encrypted_content": secret}),
1076 }]);
1077
1078 let rendered = LlmContextCompactor::<MockProvider>::format_messages_for_summary([&message]);
1079 assert!(rendered.contains("[Opaque reasoning state omitted]"));
1080 assert!(!rendered.contains(secret));
1081 }
1082
1083 #[tokio::test]
1084 async fn compact_history_preserves_opaque_reasoning_in_retained_tail() -> Result<()> {
1085 let provider = Arc::new(MockProvider::new("older context"));
1086 let config = CompactionConfig::default()
1087 .with_retain_recent(2)
1088 .with_min_messages(3);
1089 let compactor = LlmContextCompactor::new(provider, config);
1090 let opaque_data = serde_json::json!({
1091 "id": "reasoning_1",
1092 "encrypted_content": "ciphertext"
1093 });
1094 let messages = vec![
1095 Message::user("old question"),
1096 Message::assistant("old answer"),
1097 Message::user("current question"),
1098 Message::assistant_with_content(vec![ContentBlock::OpaqueReasoning {
1099 provider: "test-provider".to_owned(),
1100 data: opaque_data.clone(),
1101 }]),
1102 ];
1103
1104 let result = compactor.compact_history(messages).await?;
1105 let retained = result
1106 .messages
1107 .last()
1108 .context("compacted history should retain the newest assistant message")?;
1109 let Content::Blocks(blocks) = &retained.content else {
1110 bail!("retained assistant message should contain blocks");
1111 };
1112 assert!(matches!(
1113 blocks.first(),
1114 Some(ContentBlock::OpaqueReasoning { provider, data })
1115 if provider == "test-provider" && data == &opaque_data
1116 ));
1117 Ok(())
1118 }
1119
1120 #[tokio::test]
1121 async fn compact_history_summarizes_opaque_reasoning_prefix_and_keeps_tail() -> Result<()> {
1122 let requests = Arc::new(Mutex::new(Vec::new()));
1123 let provider = Arc::new(MockProvider::new_with_request_log(
1124 "condensed older context",
1125 Arc::clone(&requests),
1126 ));
1127 let config = CompactionConfig::default()
1128 .with_retain_recent(1)
1129 .with_min_messages(1)
1130 .with_threshold_tokens(1);
1131 let compactor = LlmContextCompactor::new(provider, config);
1132 let opaque_data = serde_json::json!({
1133 "type": "reasoning",
1134 "id": "rs_1",
1135 "encrypted_content": "ciphertext"
1136 });
1137 let messages = vec![
1138 Message::user("older user context"),
1139 Message::assistant_with_content(vec![
1144 ContentBlock::OpaqueReasoning {
1145 provider: "openai-responses".to_owned(),
1146 data: opaque_data,
1147 },
1148 ContentBlock::Text {
1149 text: "older assistant response".to_owned(),
1150 },
1151 ]),
1152 Message::user("newer user context"),
1153 Message::assistant("newer assistant response"),
1154 ];
1155
1156 assert!(
1157 compactor.needs_compaction(&messages),
1158 "opaque reasoning in the summarized prefix must not veto compaction"
1159 );
1160
1161 let result = compactor.compact_history(messages).await?;
1162
1163 assert_eq!(result.messages.len(), 3);
1165 let Content::Text(summary) = &result.messages[0].content else {
1166 bail!("summary should be a text user message");
1167 };
1168 assert!(summary.starts_with(SUMMARY_PREFIX));
1169 assert!(summary.contains("condensed older context"));
1170 assert!(
1171 matches!(&result.messages[2].content, Content::Text(text) if text == "newer assistant response"),
1172 "the retained tail survives verbatim"
1173 );
1174 assert!(
1175 !result.messages.iter().any(|message| matches!(
1176 &message.content,
1177 Content::Blocks(blocks)
1178 if blocks
1179 .iter()
1180 .any(|block| matches!(block, ContentBlock::OpaqueReasoning { .. }))
1181 )),
1182 "the summarized prefix's opaque reasoning is gone from the compacted history"
1183 );
1184
1185 let recorded = requests
1186 .lock()
1187 .map_err(|_| anyhow::anyhow!("request log poisoned"))?;
1188 assert_eq!(recorded.len(), 1, "exactly one summarization call");
1189 assert!(recorded[0].contains("[Opaque reasoning state omitted]"));
1190 assert!(!recorded[0].contains("ciphertext"));
1191 drop(recorded);
1192 Ok(())
1193 }
1194
1195 #[tokio::test]
1196 async fn test_compact_history() -> Result<()> {
1197 let provider = Arc::new(MockProvider::new(
1198 "User asked about Rust programming. Assistant explained ownership, borrowing, and lifetimes.",
1199 ));
1200 let config = CompactionConfig::default()
1201 .with_retain_recent(2)
1202 .with_min_messages(3);
1203 let compactor = LlmContextCompactor::new(provider, config);
1204
1205 let messages = vec![
1207 Message::user(
1208 "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?",
1209 ),
1210 Message::assistant(
1211 "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.",
1212 ),
1213 Message::user(
1214 "Tell me about ownership in detail. How does it work and what are the rules? I want to understand this core concept thoroughly.",
1215 ),
1216 Message::assistant(
1217 "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.",
1218 ),
1219 Message::user("What about borrowing?"), Message::assistant("Borrowing allows references to data without taking ownership."), ];
1222
1223 let result = compactor.compact_history(messages).await?;
1224
1225 assert_eq!(result.new_count, 4);
1227 assert_eq!(result.original_count, 6);
1228
1229 assert!(
1231 result.new_tokens < result.original_tokens,
1232 "Expected fewer tokens after compaction: new={} < original={}",
1233 result.new_tokens,
1234 result.original_tokens
1235 );
1236
1237 if let Content::Text(text) = &result.messages[0].content {
1239 assert!(text.contains("Previous conversation summary"));
1240 }
1241
1242 Ok(())
1243 }
1244
1245 #[tokio::test]
1246 async fn test_compact_history_too_few_messages() -> Result<()> {
1247 let provider = Arc::new(MockProvider::new("summary"));
1248 let config = CompactionConfig::default().with_retain_recent(5);
1249 let compactor = LlmContextCompactor::new(provider, config);
1250
1251 let messages = vec![
1253 Message::user("Hello"),
1254 Message::assistant("Hi"),
1255 Message::user("Bye"),
1256 ];
1257
1258 let result = compactor.compact_history(messages.clone()).await?;
1259
1260 assert_eq!(result.new_count, 3);
1262 assert_eq!(result.messages.len(), 3);
1263
1264 Ok(())
1265 }
1266
1267 #[test]
1268 fn test_format_messages_for_summary() {
1269 let messages = vec![Message::user("Hello"), Message::assistant("Hi there!")];
1270
1271 let formatted = LlmContextCompactor::<MockProvider>::format_messages_for_summary(&messages);
1272
1273 assert!(formatted.contains("User: Hello"));
1274 assert!(formatted.contains("Assistant: Hi there!"));
1275 }
1276
1277 #[test]
1278 fn test_format_messages_for_summary_truncates_tool_results_unicode_safely() {
1279 let long_unicode = "é".repeat(600);
1280
1281 let messages = vec![Message {
1282 role: Role::Assistant,
1283 content: Content::Blocks(vec![ContentBlock::ToolResult {
1284 tool_use_id: "tool-1".to_string(),
1285 content: long_unicode,
1286 is_error: Some(false),
1287 }]),
1288 }];
1289
1290 let formatted = LlmContextCompactor::<MockProvider>::format_messages_for_summary(&messages);
1291
1292 assert!(formatted.contains("... (truncated)"));
1293 }
1294
1295 #[tokio::test]
1296 async fn test_compact_carries_prior_summary_into_request() -> Result<()> {
1297 let requests = Arc::new(Mutex::new(Vec::new()));
1302 let provider = Arc::new(MockProvider::new_with_request_log(
1303 "Fresh summary",
1304 requests.clone(),
1305 ));
1306 let config = CompactionConfig::default().with_min_messages(1);
1307 let compactor = LlmContextCompactor::new(provider, config);
1308
1309 let messages = vec![
1310 Message::user(format!("{SUMMARY_PREFIX}already compacted context")),
1311 Message::assistant("Continue with the next task using this context."),
1312 ];
1313
1314 let summary = compactor.compact(&messages).await?;
1315
1316 let recorded = requests
1317 .lock()
1318 .map_err(|_| anyhow::anyhow!("request log poisoned"))?;
1319 assert_eq!(recorded.len(), 1);
1320 assert_eq!(summary, "Fresh summary");
1323 assert!(recorded[0].contains("Continue with the next task using this context."));
1324 assert!(
1325 recorded[0].contains("already compacted context"),
1326 "prior summary must be carried into the summarization input"
1327 );
1328 drop(recorded);
1329
1330 Ok(())
1331 }
1332
1333 #[tokio::test]
1334 async fn test_compact_history_carries_prior_summary_in_candidate_payload() -> Result<()> {
1335 let requests = Arc::new(Mutex::new(Vec::new()));
1336 let provider = Arc::new(MockProvider::new_with_request_log(
1337 "Fresh history summary",
1338 requests.clone(),
1339 ));
1340 let config = CompactionConfig::default()
1341 .with_retain_recent(2)
1342 .with_min_messages(1);
1343 let compactor = LlmContextCompactor::new(provider, config);
1344
1345 let messages = vec![
1346 Message::user(format!("{SUMMARY_PREFIX}already compacted context")),
1347 Message::assistant("Current turn content from the latest exchange."),
1348 Message::assistant("Recent message that should stay."),
1349 Message::user("Newest note that should stay."),
1350 ];
1351
1352 let result = compactor.compact_history(messages).await?;
1353
1354 let recorded = requests
1355 .lock()
1356 .map_err(|_| anyhow::anyhow!("request log poisoned"))?;
1357 assert_eq!(recorded.len(), 1);
1358 assert!(recorded[0].contains("Current turn content from the latest exchange."));
1359 assert!(
1362 recorded[0].contains("already compacted context"),
1363 "prior summary content must reach the summarizer"
1364 );
1365 drop(recorded);
1366 assert_eq!(result.new_count, 4);
1367
1368 Ok(())
1369 }
1370
1371 #[tokio::test]
1372 async fn test_compact_history_carries_summaries_forward_when_window_has_only_summaries()
1373 -> Result<()> {
1374 let requests = Arc::new(Mutex::new(Vec::new()));
1375 let provider = Arc::new(MockProvider::new_with_request_log(
1376 "This summary should not be used",
1377 requests.clone(),
1378 ));
1379 let config = CompactionConfig::default()
1380 .with_retain_recent(2)
1381 .with_min_messages(1);
1382 let compactor = LlmContextCompactor::new(provider, config);
1383
1384 let messages = vec![
1385 Message::user(format!("{SUMMARY_PREFIX}first prior compacted section")),
1386 Message::assistant(format!("{SUMMARY_PREFIX}second prior compacted section")),
1387 Message::user(format!("{SUMMARY_PREFIX}third prior compacted section")),
1388 Message::assistant("final short note"),
1389 ];
1390
1391 let result = compactor.compact_history(messages).await?;
1392
1393 let recorded = requests
1397 .lock()
1398 .map_err(|_| anyhow::anyhow!("request log poisoned"))?;
1399 assert!(recorded.is_empty());
1400 drop(recorded);
1401 assert_eq!(result.new_count, 4);
1402 assert_eq!(result.messages.len(), 4);
1403
1404 if let Content::Text(text) = &result.messages[0].content {
1405 assert!(
1406 text.contains("first prior compacted section"),
1407 "first prior summary lost"
1408 );
1409 assert!(
1410 text.contains("second prior compacted section"),
1411 "second prior summary lost"
1412 );
1413 assert!(!text.contains(COMPACT_EMPTY_SUMMARY));
1414 } else {
1415 panic!("Expected summary text in first message");
1416 }
1417
1418 Ok(())
1419 }
1420
1421 #[tokio::test]
1422 async fn test_compact_history_preserves_tool_use_tool_result_pairs() -> Result<()> {
1423 let provider = Arc::new(MockProvider::new("Summary of earlier conversation."));
1424 let config = CompactionConfig::default()
1425 .with_retain_recent(2)
1426 .with_min_messages(3);
1427 let compactor = LlmContextCompactor::new(provider, config);
1428
1429 let messages = vec![
1433 Message::user("What files are in the project?"),
1435 Message::assistant("Let me check that for you."),
1437 Message {
1439 role: Role::Assistant,
1440 content: Content::Blocks(vec![ContentBlock::ToolUse {
1441 id: "tool_1".to_string(),
1442 name: "list_files".to_string(),
1443 input: serde_json::json!({}),
1444 thought_signature: None,
1445 }]),
1446 },
1447 Message {
1449 role: Role::User,
1450 content: Content::Blocks(vec![ContentBlock::ToolResult {
1451 tool_use_id: "tool_1".to_string(),
1452 content: "file1.rs\nfile2.rs".to_string(),
1453 is_error: None,
1454 }]),
1455 },
1456 Message::assistant("The project contains file1.rs and file2.rs."),
1458 ];
1459
1460 let result = compactor.compact_history(messages).await?;
1461
1462 assert_eq!(result.new_count, 5);
1466
1467 let kept_assistant = &result.messages[2];
1470 if let Content::Blocks(blocks) = &kept_assistant.content {
1471 assert!(
1472 blocks
1473 .iter()
1474 .any(|b| matches!(b, ContentBlock::ToolUse { .. })),
1475 "Expected assistant tool_use in kept messages"
1476 );
1477 } else {
1478 panic!("Expected Blocks content for assistant tool_use message");
1479 }
1480
1481 let kept_user = &result.messages[3];
1483 if let Content::Blocks(blocks) = &kept_user.content {
1484 assert!(
1485 blocks
1486 .iter()
1487 .any(|b| matches!(b, ContentBlock::ToolResult { .. })),
1488 "Expected user tool_result in kept messages"
1489 );
1490 } else {
1491 panic!("Expected Blocks content for user tool_result message");
1492 }
1493
1494 Ok(())
1495 }
1496
1497 #[tokio::test]
1498 async fn test_compact_history_split_skips_leading_orphan_after_summary_ack() -> Result<()> {
1499 let provider = Arc::new(MockProvider::new("Re-summary."));
1527 let config = CompactionConfig::default()
1528 .with_retain_recent(3)
1529 .with_min_messages(1);
1530 let compactor = LlmContextCompactor::new(provider, config);
1531
1532 let messages = vec![
1533 Message::user(format!("{SUMMARY_PREFIX}Old summary about toolu_X.")),
1534 Message::assistant(SUMMARY_ACKNOWLEDGMENT),
1535 Message {
1536 role: Role::User,
1537 content: Content::Blocks(vec![ContentBlock::ToolResult {
1538 tool_use_id: "toolu_X".to_string(),
1539 content: "result for X".to_string(),
1540 is_error: None,
1541 }]),
1542 },
1543 Message::assistant("Result interpreted."),
1544 Message::user("Now what?"),
1545 ];
1546
1547 let result = compactor.compact_history(messages).await?;
1548
1549 let mut seen_ids: std::collections::HashSet<String> = std::collections::HashSet::new();
1550 for msg in &result.messages {
1551 if let Content::Blocks(blocks) = &msg.content {
1552 for block in blocks {
1553 match block {
1554 ContentBlock::ToolResult { tool_use_id, .. } => {
1555 assert!(
1556 seen_ids.contains(tool_use_id),
1557 "orphan tool_use_id {tool_use_id} survived split selection",
1558 );
1559 }
1560 ContentBlock::ToolUse { id, .. } => {
1561 seen_ids.insert(id.clone());
1562 }
1563 _ => {}
1564 }
1565 }
1566 }
1567 }
1568
1569 Ok(())
1570 }
1571
1572 #[tokio::test]
1573 async fn test_compact_history_keeps_tool_pair_when_immediate_prev_is_text_only() -> Result<()> {
1574 let provider = Arc::new(MockProvider::new("Boundary summary."));
1581 let config = CompactionConfig::default()
1582 .with_retain_recent(2)
1583 .with_min_messages(1);
1584 let compactor = LlmContextCompactor::new(provider, config);
1585
1586 let messages = vec![
1599 Message::user("first turn"),
1600 Message::assistant("text only"),
1601 Message {
1602 role: Role::User,
1603 content: Content::Blocks(vec![ContentBlock::ToolResult {
1604 tool_use_id: "toolu_Y".to_string(),
1605 content: "ancient result".to_string(),
1606 is_error: None,
1607 }]),
1608 },
1609 Message::assistant("then a reply"),
1610 Message::user("ok thanks"),
1611 ];
1612
1613 let result = compactor.compact_history(messages).await?;
1614
1615 let has_tool_result = result.messages.iter().any(|m| {
1619 matches!(
1620 &m.content,
1621 Content::Blocks(blocks)
1622 if blocks.iter().any(|b| matches!(b, ContentBlock::ToolResult { .. }))
1623 )
1624 });
1625 assert!(
1626 !has_tool_result,
1627 "orphan tool_result should have been pushed into to_summarize, not retained",
1628 );
1629
1630 Ok(())
1631 }
1632
1633 #[tokio::test]
1634 async fn test_compact_history_retained_tail_is_token_capped() -> Result<()> {
1635 let provider = Arc::new(MockProvider::new(
1636 "Project summary with a long context and technical context.",
1637 ));
1638 let config = CompactionConfig::default()
1639 .with_retain_recent(8)
1640 .with_min_messages(1)
1641 .with_threshold_tokens(1);
1642 let compactor = LlmContextCompactor::new(provider, config);
1643
1644 let mut messages = Vec::new();
1645
1646 messages.extend((0..6).map(|index| Message::user(format!("pre-compaction noise {index}"))));
1648
1649 messages.extend(
1651 (0..8).map(|index| Message::assistant(format!("kept-{index}: {}", "x".repeat(12_000)))),
1652 );
1653
1654 let result = compactor.compact_history(messages).await?;
1655
1656 let retained_tail = &result.messages[2..];
1658 assert!(retained_tail.len() < 8);
1659
1660 let mut latest_index = -1i32;
1661 let mut all_retained = true;
1662 for message in retained_tail {
1663 if let Content::Text(text) = &message.content {
1664 if let Some(number) = text.split(':').next().and_then(|prefix| {
1665 prefix
1666 .strip_prefix("kept-")
1667 .and_then(|rest| rest.parse::<i32>().ok())
1668 }) {
1669 if number >= 0 {
1670 latest_index = latest_index.max(number);
1671 }
1672 } else {
1673 all_retained = false;
1674 }
1675 } else {
1676 all_retained = false;
1677 }
1678 }
1679
1680 assert!(all_retained);
1681 assert_eq!(latest_index, 7);
1682 assert!(
1683 TokenEstimator::estimate_history(retained_tail)
1684 <= compactor.config().max_retained_tail_tokens
1685 );
1686 assert!(compactor.needs_compaction(&result.messages));
1687
1688 Ok(())
1689 }
1690
1691 #[tokio::test]
1692 async fn test_compact_history_skips_summary_ack_when_retained_tail_is_empty() -> Result<()> {
1693 let provider = Arc::new(MockProvider::new("Summary for oversized user turn."));
1694 let config = CompactionConfig::default()
1695 .with_retain_recent(1)
1696 .with_min_messages(1)
1697 .with_threshold_tokens(1);
1698 let compactor = LlmContextCompactor::new(provider, config);
1699
1700 let messages = vec![
1701 Message::assistant("Earlier assistant context."),
1702 Message::user(format!("oversized-user-turn: {}", "x".repeat(200_000))),
1703 ];
1704
1705 let result = compactor.compact_history(messages).await?;
1706
1707 assert_eq!(result.new_count, 1);
1708 assert_eq!(result.messages.len(), 1);
1709
1710 let only_message = &result.messages[0];
1711 assert_eq!(only_message.role, Role::User);
1712
1713 if let Content::Text(text) = &only_message.content {
1714 assert!(text.contains("Previous conversation summary"));
1715 assert!(!text.contains(SUMMARY_ACKNOWLEDGMENT));
1716 } else {
1717 panic!("Expected summary text when retained tail is empty");
1718 }
1719
1720 Ok(())
1721 }
1722
1723 fn message_contains(message: &Message, needle: &str) -> bool {
1724 match &message.content {
1725 Content::Text(text) => text.contains(needle),
1726 Content::Blocks(blocks) => blocks.iter().any(|block| match block {
1727 ContentBlock::Text { text } => text.contains(needle),
1728 _ => false,
1729 }),
1730 }
1731 }
1732
1733 #[tokio::test]
1734 async fn test_epoch_one_facts_survive_two_compactions() -> Result<()> {
1735 const EPOCH1_FACT: &str = "EPOCH1_FACT: the API key lives in config/secrets.toml";
1741
1742 let requests = Arc::new(Mutex::new(Vec::new()));
1743 let provider = Arc::new(MockProvider::new_echo(requests.clone()));
1744 let config = CompactionConfig::default()
1745 .with_retain_recent(2)
1746 .with_min_messages(1);
1747 let compactor = LlmContextCompactor::new(provider, config);
1748
1749 let epoch1 = vec![
1750 Message::user(EPOCH1_FACT),
1751 Message::assistant("Understood, noted the secrets path."),
1752 Message::user("Now add error handling to main.rs."),
1753 Message::assistant("Added error handling to main.rs."),
1754 Message::user("latest user message one"),
1755 Message::assistant("latest assistant message two"),
1756 ];
1757
1758 let first = compactor.compact_history(epoch1).await?;
1759 assert!(
1760 first
1761 .messages
1762 .iter()
1763 .any(|m| message_contains(m, "EPOCH1_FACT")),
1764 "epoch-1 fact must be captured in the first summary"
1765 );
1766
1767 let mut epoch2 = first.messages;
1769 epoch2.push(Message::user("Another later turn."));
1770 epoch2.push(Message::assistant("Reply to the later turn."));
1771 epoch2.push(Message::user("Final turn a."));
1772 epoch2.push(Message::assistant("Final turn b."));
1773
1774 let second = compactor.compact_history(epoch2).await?;
1775
1776 assert!(
1777 second
1778 .messages
1779 .iter()
1780 .any(|m| message_contains(m, "EPOCH1_FACT")),
1781 "epoch-1 fact must survive the second compaction"
1782 );
1783
1784 let recorded = requests
1787 .lock()
1788 .map_err(|_| anyhow::anyhow!("request log poisoned"))?;
1789 assert!(
1790 recorded.iter().any(|req| req.contains("EPOCH1_FACT")),
1791 "prior summary carrying the epoch-1 fact must reach the summarizer"
1792 );
1793 drop(recorded);
1794
1795 Ok(())
1796 }
1797
1798 #[tokio::test]
1799 async fn test_compact_history_long_tool_chain_respects_token_cap() -> Result<()> {
1800 let provider = Arc::new(MockProvider::new("Summary of the early tool chain."));
1806 let cap = 20_000;
1807 let config = CompactionConfig::default()
1810 .with_retain_recent(18)
1811 .with_min_messages(1)
1812 .with_threshold_tokens(1)
1813 .with_max_retained_tail_tokens(cap);
1814 let compactor = LlmContextCompactor::new(provider, config);
1815
1816 let mut messages = Vec::new();
1819 for i in 0..10 {
1820 messages.push(Message {
1821 role: Role::Assistant,
1822 content: Content::Blocks(vec![ContentBlock::ToolUse {
1823 id: format!("tool_{i}"),
1824 name: "run".to_string(),
1825 input: serde_json::json!({ "arg": "y".repeat(12_000) }),
1826 thought_signature: None,
1827 }]),
1828 });
1829 messages.push(Message {
1830 role: Role::User,
1831 content: Content::Blocks(vec![ContentBlock::ToolResult {
1832 tool_use_id: format!("tool_{i}"),
1833 content: format!("result-{i}: {}", "z".repeat(12_000)),
1834 is_error: None,
1835 }]),
1836 });
1837 }
1838
1839 let full_tokens = TokenEstimator::estimate_history(&messages);
1840 assert!(
1841 full_tokens > cap * 2,
1842 "test setup: full chain must far exceed the cap"
1843 );
1844
1845 let result = compactor.compact_history(messages).await?;
1846
1847 let retained_tail = &result.messages[2..];
1850
1851 let tail_tokens = TokenEstimator::estimate_history(retained_tail);
1852 assert!(
1855 tail_tokens <= cap + 8_000,
1856 "retained tail {tail_tokens} should be bounded by the cap {cap}, not the whole chain"
1857 );
1858 assert!(
1859 retained_tail.len() < 20,
1860 "compaction must have summarized part of the chain"
1861 );
1862
1863 Ok(())
1864 }
1865
1866 #[tokio::test]
1867 async fn test_compact_warns_and_marks_truncated_summary() -> Result<()> {
1868 let requests = Arc::new(Mutex::new(Vec::new()));
1873 let provider = Arc::new(MockProvider::new_truncating(
1874 "partial summary cut off mid-",
1875 requests.clone(),
1876 ));
1877 let config = CompactionConfig::default().with_min_messages(1);
1878 let compactor = LlmContextCompactor::new(provider, config);
1879
1880 let messages = vec![
1881 Message::user("Some content that needs summarizing."),
1882 Message::assistant("More content to summarize here."),
1883 ];
1884
1885 let summary = compactor.compact(&messages).await?;
1886
1887 assert!(
1888 summary.contains("[summary truncated"),
1889 "a persistently truncated summary must carry a truncation marker"
1890 );
1891
1892 let recorded = requests
1894 .lock()
1895 .map_err(|_| anyhow::anyhow!("request log poisoned"))?;
1896 assert_eq!(recorded.len(), 2, "truncation should trigger one retry");
1897 drop(recorded);
1898
1899 Ok(())
1900 }
1901
1902 fn summarizable_messages() -> Vec<Message> {
1905 vec![
1906 Message::user("First question with enough words to summarize meaningfully."),
1907 Message::assistant("First answer, also carrying plenty of prose to compact."),
1908 Message::user("Second question continuing the earlier conversation topic."),
1909 Message::assistant("Second answer expanding on the topic at some length."),
1910 Message::user("Third question?"),
1911 Message::assistant("Third answer."),
1912 ]
1913 }
1914
1915 struct BlockRequestHooks;
1916
1917 #[async_trait]
1918 impl crate::hooks::AgentHooks for BlockRequestHooks {
1919 async fn pre_llm_request(&self, _request: &ChatRequest) -> RequestDecision {
1920 RequestDecision::Block("summaries are not allowed".to_string())
1921 }
1922 }
1923
1924 struct ModifyRequestHooks;
1925
1926 #[async_trait]
1927 impl crate::hooks::AgentHooks for ModifyRequestHooks {
1928 async fn pre_llm_request(&self, request: &ChatRequest) -> RequestDecision {
1929 let mut modified = request.clone();
1930 modified.messages = vec![Message::user("MODIFIED_SUMMARY_PROMPT")];
1931 RequestDecision::Modify(Box::new(modified))
1932 }
1933 }
1934
1935 struct BlockResponseHooks;
1936
1937 #[async_trait]
1938 impl crate::hooks::AgentHooks for BlockResponseHooks {
1939 async fn on_llm_response(&self, _response: &ChatResponse) -> ResponseDecision {
1940 ResponseDecision::Block("summary leaks a secret".to_string())
1941 }
1942 }
1943
1944 struct RetryResponseHooks;
1945
1946 #[async_trait]
1947 impl crate::hooks::AgentHooks for RetryResponseHooks {
1948 async fn on_llm_response(&self, _response: &ChatResponse) -> ResponseDecision {
1949 ResponseDecision::RetryWithFeedback("try harder".to_string())
1950 }
1951 }
1952
1953 #[tokio::test]
1954 async fn blocking_request_hook_aborts_compaction_before_the_llm_call() -> Result<()> {
1955 let requests = Arc::new(Mutex::new(Vec::new()));
1956 let provider = Arc::new(MockProvider::new_with_request_log(
1957 "summary",
1958 requests.clone(),
1959 ));
1960 let config = CompactionConfig::default().with_retain_recent(2);
1961 let compactor = LlmContextCompactor::new(provider, config)
1962 .with_guardrail_hooks(Arc::new(BlockRequestHooks));
1963
1964 let error = match compactor.compact_history(summarizable_messages()).await {
1965 Ok(result) => anyhow::bail!("blocked compaction must not succeed: {result:?}"),
1966 Err(error) => error,
1967 };
1968 assert!(
1969 error.to_string().contains("blocked by guardrail"),
1970 "unexpected error: {error}"
1971 );
1972
1973 let recorded = requests
1974 .lock()
1975 .map_err(|_| anyhow::anyhow!("request log poisoned"))?;
1976 assert!(
1977 recorded.is_empty(),
1978 "a blocked request must never reach the provider"
1979 );
1980 drop(recorded);
1981 Ok(())
1982 }
1983
1984 #[tokio::test]
1985 async fn modify_request_hook_reaches_the_provider() -> Result<()> {
1986 let requests = Arc::new(Mutex::new(Vec::new()));
1987 let provider = Arc::new(MockProvider::new_with_request_log(
1988 "summary",
1989 requests.clone(),
1990 ));
1991 let config = CompactionConfig::default().with_retain_recent(2);
1992 let compactor = LlmContextCompactor::new(provider, config)
1993 .with_guardrail_hooks(Arc::new(ModifyRequestHooks));
1994
1995 let result = compactor.compact_history(summarizable_messages()).await?;
1996 assert!(result.new_count < result.original_count);
1997
1998 let recorded = requests
1999 .lock()
2000 .map_err(|_| anyhow::anyhow!("request log poisoned"))?;
2001 assert_eq!(
2002 recorded.as_slice(),
2003 ["MODIFIED_SUMMARY_PROMPT"],
2004 "the provider must receive the hook-modified request"
2005 );
2006 drop(recorded);
2007 Ok(())
2008 }
2009
2010 #[tokio::test]
2011 async fn blocked_response_aborts_compaction() -> Result<()> {
2012 let provider = Arc::new(MockProvider::new("a summary that leaks the secret"));
2013 let config = CompactionConfig::default().with_retain_recent(2);
2014 let compactor = LlmContextCompactor::new(provider, config)
2015 .with_guardrail_hooks(Arc::new(BlockResponseHooks));
2016
2017 let error = match compactor.compact_history(summarizable_messages()).await {
2018 Ok(result) => anyhow::bail!("blocked summary must not be returned: {result:?}"),
2019 Err(error) => error,
2020 };
2021 assert!(
2022 error.to_string().contains("blocked by guardrail"),
2023 "unexpected error: {error}"
2024 );
2025 Ok(())
2026 }
2027
2028 #[tokio::test]
2029 async fn retry_with_feedback_response_aborts_compaction_without_retrying() -> Result<()> {
2030 let requests = Arc::new(Mutex::new(Vec::new()));
2031 let provider = Arc::new(MockProvider::new_with_request_log(
2032 "summary",
2033 requests.clone(),
2034 ));
2035 let config = CompactionConfig::default().with_retain_recent(2);
2036 let compactor = LlmContextCompactor::new(provider, config)
2037 .with_guardrail_hooks(Arc::new(RetryResponseHooks));
2038
2039 let error = match compactor.compact_history(summarizable_messages()).await {
2040 Ok(result) => anyhow::bail!("rejected summary must not be returned: {result:?}"),
2041 Err(error) => error,
2042 };
2043 assert!(
2044 error.to_string().contains("not retried during compaction"),
2045 "unexpected error: {error}"
2046 );
2047
2048 let recorded = requests
2050 .lock()
2051 .map_err(|_| anyhow::anyhow!("request log poisoned"))?;
2052 assert_eq!(recorded.len(), 1, "RetryWithFeedback must not retry");
2053 drop(recorded);
2054 Ok(())
2055 }
2056
2057 #[tokio::test]
2058 async fn compact_history_reports_summarization_usage() -> Result<()> {
2059 let provider = Arc::new(MockProvider::new("summary"));
2060 let config = CompactionConfig::default().with_retain_recent(2);
2061 let compactor = LlmContextCompactor::new(provider, config);
2062
2063 let result = compactor.compact_history(summarizable_messages()).await?;
2064 assert_eq!(result.llm_usage.input_tokens, 100);
2066 assert_eq!(result.llm_usage.output_tokens, 50);
2067 Ok(())
2068 }
2069}