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 input_tokens: 100,
997 output_tokens: 50,
998 cached_input_tokens: 0,
999 cache_creation_input_tokens: 0,
1000 },
1001 }))
1002 }
1003
1004 fn model(&self) -> &'static str {
1005 "mock-model"
1006 }
1007
1008 fn provider(&self) -> &'static str {
1009 "mock"
1010 }
1011 }
1012
1013 #[test]
1014 fn test_needs_compaction_below_threshold() {
1015 let provider = Arc::new(MockProvider::new("summary"));
1016 let config = CompactionConfig::default()
1017 .with_threshold_tokens(10_000)
1018 .with_min_messages(5);
1019 let compactor = LlmContextCompactor::new(provider, config);
1020
1021 let messages = vec![
1023 Message::user("Hello"),
1024 Message::assistant("Hi"),
1025 Message::user("How are you?"),
1026 ];
1027
1028 assert!(!compactor.needs_compaction(&messages));
1029 }
1030
1031 #[test]
1032 fn test_needs_compaction_above_threshold() {
1033 let provider = Arc::new(MockProvider::new("summary"));
1034 let config = CompactionConfig::default()
1035 .with_threshold_tokens(50) .with_min_messages(3);
1037 let compactor = LlmContextCompactor::new(provider, config);
1038
1039 let messages = vec![
1041 Message::user("Hello, this is a longer message to test compaction"),
1042 Message::assistant(
1043 "Hi there! This is also a longer response to help trigger compaction",
1044 ),
1045 Message::user("Great, let's continue with even more text here"),
1046 Message::assistant("Absolutely, adding more content to ensure we exceed the threshold"),
1047 ];
1048
1049 assert!(compactor.needs_compaction(&messages));
1050 }
1051
1052 #[test]
1053 fn test_needs_compaction_auto_disabled() {
1054 let provider = Arc::new(MockProvider::new("summary"));
1055 let config = CompactionConfig::default()
1056 .with_threshold_tokens(10) .with_min_messages(1)
1058 .with_auto_compact(false);
1059 let compactor = LlmContextCompactor::new(provider, config);
1060
1061 let messages = vec![
1062 Message::user("Hello, this is a longer message"),
1063 Message::assistant("Response here"),
1064 ];
1065
1066 assert!(!compactor.needs_compaction(&messages));
1067 }
1068
1069 #[test]
1070 fn summary_prompt_redacts_opaque_reasoning_payload() {
1071 let secret = "opaque-secret-that-must-not-enter-the-summary-prompt";
1072 let message = Message::assistant_with_content(vec![ContentBlock::OpaqueReasoning {
1073 provider: "test-provider".to_owned(),
1074 data: serde_json::json!({"encrypted_content": secret}),
1075 }]);
1076
1077 let rendered = LlmContextCompactor::<MockProvider>::format_messages_for_summary([&message]);
1078 assert!(rendered.contains("[Opaque reasoning state omitted]"));
1079 assert!(!rendered.contains(secret));
1080 }
1081
1082 #[tokio::test]
1083 async fn compact_history_preserves_opaque_reasoning_in_retained_tail() -> Result<()> {
1084 let provider = Arc::new(MockProvider::new("older context"));
1085 let config = CompactionConfig::default()
1086 .with_retain_recent(2)
1087 .with_min_messages(3);
1088 let compactor = LlmContextCompactor::new(provider, config);
1089 let opaque_data = serde_json::json!({
1090 "id": "reasoning_1",
1091 "encrypted_content": "ciphertext"
1092 });
1093 let messages = vec![
1094 Message::user("old question"),
1095 Message::assistant("old answer"),
1096 Message::user("current question"),
1097 Message::assistant_with_content(vec![ContentBlock::OpaqueReasoning {
1098 provider: "test-provider".to_owned(),
1099 data: opaque_data.clone(),
1100 }]),
1101 ];
1102
1103 let result = compactor.compact_history(messages).await?;
1104 let retained = result
1105 .messages
1106 .last()
1107 .context("compacted history should retain the newest assistant message")?;
1108 let Content::Blocks(blocks) = &retained.content else {
1109 bail!("retained assistant message should contain blocks");
1110 };
1111 assert!(matches!(
1112 blocks.first(),
1113 Some(ContentBlock::OpaqueReasoning { provider, data })
1114 if provider == "test-provider" && data == &opaque_data
1115 ));
1116 Ok(())
1117 }
1118
1119 #[tokio::test]
1120 async fn compact_history_summarizes_opaque_reasoning_prefix_and_keeps_tail() -> Result<()> {
1121 let requests = Arc::new(Mutex::new(Vec::new()));
1122 let provider = Arc::new(MockProvider::new_with_request_log(
1123 "condensed older context",
1124 Arc::clone(&requests),
1125 ));
1126 let config = CompactionConfig::default()
1127 .with_retain_recent(1)
1128 .with_min_messages(1)
1129 .with_threshold_tokens(1);
1130 let compactor = LlmContextCompactor::new(provider, config);
1131 let opaque_data = serde_json::json!({
1132 "type": "reasoning",
1133 "id": "rs_1",
1134 "encrypted_content": "ciphertext"
1135 });
1136 let messages = vec![
1137 Message::user("older user context"),
1138 Message::assistant_with_content(vec![
1143 ContentBlock::OpaqueReasoning {
1144 provider: "openai-responses".to_owned(),
1145 data: opaque_data,
1146 },
1147 ContentBlock::Text {
1148 text: "older assistant response".to_owned(),
1149 },
1150 ]),
1151 Message::user("newer user context"),
1152 Message::assistant("newer assistant response"),
1153 ];
1154
1155 assert!(
1156 compactor.needs_compaction(&messages),
1157 "opaque reasoning in the summarized prefix must not veto compaction"
1158 );
1159
1160 let result = compactor.compact_history(messages).await?;
1161
1162 assert_eq!(result.messages.len(), 3);
1164 let Content::Text(summary) = &result.messages[0].content else {
1165 bail!("summary should be a text user message");
1166 };
1167 assert!(summary.starts_with(SUMMARY_PREFIX));
1168 assert!(summary.contains("condensed older context"));
1169 assert!(
1170 matches!(&result.messages[2].content, Content::Text(text) if text == "newer assistant response"),
1171 "the retained tail survives verbatim"
1172 );
1173 assert!(
1174 !result.messages.iter().any(|message| matches!(
1175 &message.content,
1176 Content::Blocks(blocks)
1177 if blocks
1178 .iter()
1179 .any(|block| matches!(block, ContentBlock::OpaqueReasoning { .. }))
1180 )),
1181 "the summarized prefix's opaque reasoning is gone from the compacted history"
1182 );
1183
1184 let recorded = requests
1185 .lock()
1186 .map_err(|_| anyhow::anyhow!("request log poisoned"))?;
1187 assert_eq!(recorded.len(), 1, "exactly one summarization call");
1188 assert!(recorded[0].contains("[Opaque reasoning state omitted]"));
1189 assert!(!recorded[0].contains("ciphertext"));
1190 drop(recorded);
1191 Ok(())
1192 }
1193
1194 #[tokio::test]
1195 async fn test_compact_history() -> Result<()> {
1196 let provider = Arc::new(MockProvider::new(
1197 "User asked about Rust programming. Assistant explained ownership, borrowing, and lifetimes.",
1198 ));
1199 let config = CompactionConfig::default()
1200 .with_retain_recent(2)
1201 .with_min_messages(3);
1202 let compactor = LlmContextCompactor::new(provider, config);
1203
1204 let messages = vec![
1206 Message::user(
1207 "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?",
1208 ),
1209 Message::assistant(
1210 "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.",
1211 ),
1212 Message::user(
1213 "Tell me about ownership in detail. How does it work and what are the rules? I want to understand this core concept thoroughly.",
1214 ),
1215 Message::assistant(
1216 "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.",
1217 ),
1218 Message::user("What about borrowing?"), Message::assistant("Borrowing allows references to data without taking ownership."), ];
1221
1222 let result = compactor.compact_history(messages).await?;
1223
1224 assert_eq!(result.new_count, 4);
1226 assert_eq!(result.original_count, 6);
1227
1228 assert!(
1230 result.new_tokens < result.original_tokens,
1231 "Expected fewer tokens after compaction: new={} < original={}",
1232 result.new_tokens,
1233 result.original_tokens
1234 );
1235
1236 if let Content::Text(text) = &result.messages[0].content {
1238 assert!(text.contains("Previous conversation summary"));
1239 }
1240
1241 Ok(())
1242 }
1243
1244 #[tokio::test]
1245 async fn test_compact_history_too_few_messages() -> Result<()> {
1246 let provider = Arc::new(MockProvider::new("summary"));
1247 let config = CompactionConfig::default().with_retain_recent(5);
1248 let compactor = LlmContextCompactor::new(provider, config);
1249
1250 let messages = vec![
1252 Message::user("Hello"),
1253 Message::assistant("Hi"),
1254 Message::user("Bye"),
1255 ];
1256
1257 let result = compactor.compact_history(messages.clone()).await?;
1258
1259 assert_eq!(result.new_count, 3);
1261 assert_eq!(result.messages.len(), 3);
1262
1263 Ok(())
1264 }
1265
1266 #[test]
1267 fn test_format_messages_for_summary() {
1268 let messages = vec![Message::user("Hello"), Message::assistant("Hi there!")];
1269
1270 let formatted = LlmContextCompactor::<MockProvider>::format_messages_for_summary(&messages);
1271
1272 assert!(formatted.contains("User: Hello"));
1273 assert!(formatted.contains("Assistant: Hi there!"));
1274 }
1275
1276 #[test]
1277 fn test_format_messages_for_summary_truncates_tool_results_unicode_safely() {
1278 let long_unicode = "é".repeat(600);
1279
1280 let messages = vec![Message {
1281 role: Role::Assistant,
1282 content: Content::Blocks(vec![ContentBlock::ToolResult {
1283 tool_use_id: "tool-1".to_string(),
1284 content: long_unicode,
1285 is_error: Some(false),
1286 }]),
1287 }];
1288
1289 let formatted = LlmContextCompactor::<MockProvider>::format_messages_for_summary(&messages);
1290
1291 assert!(formatted.contains("... (truncated)"));
1292 }
1293
1294 #[tokio::test]
1295 async fn test_compact_carries_prior_summary_into_request() -> Result<()> {
1296 let requests = Arc::new(Mutex::new(Vec::new()));
1301 let provider = Arc::new(MockProvider::new_with_request_log(
1302 "Fresh summary",
1303 requests.clone(),
1304 ));
1305 let config = CompactionConfig::default().with_min_messages(1);
1306 let compactor = LlmContextCompactor::new(provider, config);
1307
1308 let messages = vec![
1309 Message::user(format!("{SUMMARY_PREFIX}already compacted context")),
1310 Message::assistant("Continue with the next task using this context."),
1311 ];
1312
1313 let summary = compactor.compact(&messages).await?;
1314
1315 let recorded = requests
1316 .lock()
1317 .map_err(|_| anyhow::anyhow!("request log poisoned"))?;
1318 assert_eq!(recorded.len(), 1);
1319 assert_eq!(summary, "Fresh summary");
1322 assert!(recorded[0].contains("Continue with the next task using this context."));
1323 assert!(
1324 recorded[0].contains("already compacted context"),
1325 "prior summary must be carried into the summarization input"
1326 );
1327 drop(recorded);
1328
1329 Ok(())
1330 }
1331
1332 #[tokio::test]
1333 async fn test_compact_history_carries_prior_summary_in_candidate_payload() -> Result<()> {
1334 let requests = Arc::new(Mutex::new(Vec::new()));
1335 let provider = Arc::new(MockProvider::new_with_request_log(
1336 "Fresh history summary",
1337 requests.clone(),
1338 ));
1339 let config = CompactionConfig::default()
1340 .with_retain_recent(2)
1341 .with_min_messages(1);
1342 let compactor = LlmContextCompactor::new(provider, config);
1343
1344 let messages = vec![
1345 Message::user(format!("{SUMMARY_PREFIX}already compacted context")),
1346 Message::assistant("Current turn content from the latest exchange."),
1347 Message::assistant("Recent message that should stay."),
1348 Message::user("Newest note that should stay."),
1349 ];
1350
1351 let result = compactor.compact_history(messages).await?;
1352
1353 let recorded = requests
1354 .lock()
1355 .map_err(|_| anyhow::anyhow!("request log poisoned"))?;
1356 assert_eq!(recorded.len(), 1);
1357 assert!(recorded[0].contains("Current turn content from the latest exchange."));
1358 assert!(
1361 recorded[0].contains("already compacted context"),
1362 "prior summary content must reach the summarizer"
1363 );
1364 drop(recorded);
1365 assert_eq!(result.new_count, 4);
1366
1367 Ok(())
1368 }
1369
1370 #[tokio::test]
1371 async fn test_compact_history_carries_summaries_forward_when_window_has_only_summaries()
1372 -> Result<()> {
1373 let requests = Arc::new(Mutex::new(Vec::new()));
1374 let provider = Arc::new(MockProvider::new_with_request_log(
1375 "This summary should not be used",
1376 requests.clone(),
1377 ));
1378 let config = CompactionConfig::default()
1379 .with_retain_recent(2)
1380 .with_min_messages(1);
1381 let compactor = LlmContextCompactor::new(provider, config);
1382
1383 let messages = vec![
1384 Message::user(format!("{SUMMARY_PREFIX}first prior compacted section")),
1385 Message::assistant(format!("{SUMMARY_PREFIX}second prior compacted section")),
1386 Message::user(format!("{SUMMARY_PREFIX}third prior compacted section")),
1387 Message::assistant("final short note"),
1388 ];
1389
1390 let result = compactor.compact_history(messages).await?;
1391
1392 let recorded = requests
1396 .lock()
1397 .map_err(|_| anyhow::anyhow!("request log poisoned"))?;
1398 assert!(recorded.is_empty());
1399 drop(recorded);
1400 assert_eq!(result.new_count, 4);
1401 assert_eq!(result.messages.len(), 4);
1402
1403 if let Content::Text(text) = &result.messages[0].content {
1404 assert!(
1405 text.contains("first prior compacted section"),
1406 "first prior summary lost"
1407 );
1408 assert!(
1409 text.contains("second prior compacted section"),
1410 "second prior summary lost"
1411 );
1412 assert!(!text.contains(COMPACT_EMPTY_SUMMARY));
1413 } else {
1414 panic!("Expected summary text in first message");
1415 }
1416
1417 Ok(())
1418 }
1419
1420 #[tokio::test]
1421 async fn test_compact_history_preserves_tool_use_tool_result_pairs() -> Result<()> {
1422 let provider = Arc::new(MockProvider::new("Summary of earlier conversation."));
1423 let config = CompactionConfig::default()
1424 .with_retain_recent(2)
1425 .with_min_messages(3);
1426 let compactor = LlmContextCompactor::new(provider, config);
1427
1428 let messages = vec![
1432 Message::user("What files are in the project?"),
1434 Message::assistant("Let me check that for you."),
1436 Message {
1438 role: Role::Assistant,
1439 content: Content::Blocks(vec![ContentBlock::ToolUse {
1440 id: "tool_1".to_string(),
1441 name: "list_files".to_string(),
1442 input: serde_json::json!({}),
1443 thought_signature: None,
1444 }]),
1445 },
1446 Message {
1448 role: Role::User,
1449 content: Content::Blocks(vec![ContentBlock::ToolResult {
1450 tool_use_id: "tool_1".to_string(),
1451 content: "file1.rs\nfile2.rs".to_string(),
1452 is_error: None,
1453 }]),
1454 },
1455 Message::assistant("The project contains file1.rs and file2.rs."),
1457 ];
1458
1459 let result = compactor.compact_history(messages).await?;
1460
1461 assert_eq!(result.new_count, 5);
1465
1466 let kept_assistant = &result.messages[2];
1469 if let Content::Blocks(blocks) = &kept_assistant.content {
1470 assert!(
1471 blocks
1472 .iter()
1473 .any(|b| matches!(b, ContentBlock::ToolUse { .. })),
1474 "Expected assistant tool_use in kept messages"
1475 );
1476 } else {
1477 panic!("Expected Blocks content for assistant tool_use message");
1478 }
1479
1480 let kept_user = &result.messages[3];
1482 if let Content::Blocks(blocks) = &kept_user.content {
1483 assert!(
1484 blocks
1485 .iter()
1486 .any(|b| matches!(b, ContentBlock::ToolResult { .. })),
1487 "Expected user tool_result in kept messages"
1488 );
1489 } else {
1490 panic!("Expected Blocks content for user tool_result message");
1491 }
1492
1493 Ok(())
1494 }
1495
1496 #[tokio::test]
1497 async fn test_compact_history_split_skips_leading_orphan_after_summary_ack() -> Result<()> {
1498 let provider = Arc::new(MockProvider::new("Re-summary."));
1526 let config = CompactionConfig::default()
1527 .with_retain_recent(3)
1528 .with_min_messages(1);
1529 let compactor = LlmContextCompactor::new(provider, config);
1530
1531 let messages = vec![
1532 Message::user(format!("{SUMMARY_PREFIX}Old summary about toolu_X.")),
1533 Message::assistant(SUMMARY_ACKNOWLEDGMENT),
1534 Message {
1535 role: Role::User,
1536 content: Content::Blocks(vec![ContentBlock::ToolResult {
1537 tool_use_id: "toolu_X".to_string(),
1538 content: "result for X".to_string(),
1539 is_error: None,
1540 }]),
1541 },
1542 Message::assistant("Result interpreted."),
1543 Message::user("Now what?"),
1544 ];
1545
1546 let result = compactor.compact_history(messages).await?;
1547
1548 let mut seen_ids: std::collections::HashSet<String> = std::collections::HashSet::new();
1549 for msg in &result.messages {
1550 if let Content::Blocks(blocks) = &msg.content {
1551 for block in blocks {
1552 match block {
1553 ContentBlock::ToolResult { tool_use_id, .. } => {
1554 assert!(
1555 seen_ids.contains(tool_use_id),
1556 "orphan tool_use_id {tool_use_id} survived split selection",
1557 );
1558 }
1559 ContentBlock::ToolUse { id, .. } => {
1560 seen_ids.insert(id.clone());
1561 }
1562 _ => {}
1563 }
1564 }
1565 }
1566 }
1567
1568 Ok(())
1569 }
1570
1571 #[tokio::test]
1572 async fn test_compact_history_keeps_tool_pair_when_immediate_prev_is_text_only() -> Result<()> {
1573 let provider = Arc::new(MockProvider::new("Boundary summary."));
1580 let config = CompactionConfig::default()
1581 .with_retain_recent(2)
1582 .with_min_messages(1);
1583 let compactor = LlmContextCompactor::new(provider, config);
1584
1585 let messages = vec![
1598 Message::user("first turn"),
1599 Message::assistant("text only"),
1600 Message {
1601 role: Role::User,
1602 content: Content::Blocks(vec![ContentBlock::ToolResult {
1603 tool_use_id: "toolu_Y".to_string(),
1604 content: "ancient result".to_string(),
1605 is_error: None,
1606 }]),
1607 },
1608 Message::assistant("then a reply"),
1609 Message::user("ok thanks"),
1610 ];
1611
1612 let result = compactor.compact_history(messages).await?;
1613
1614 let has_tool_result = result.messages.iter().any(|m| {
1618 matches!(
1619 &m.content,
1620 Content::Blocks(blocks)
1621 if blocks.iter().any(|b| matches!(b, ContentBlock::ToolResult { .. }))
1622 )
1623 });
1624 assert!(
1625 !has_tool_result,
1626 "orphan tool_result should have been pushed into to_summarize, not retained",
1627 );
1628
1629 Ok(())
1630 }
1631
1632 #[tokio::test]
1633 async fn test_compact_history_retained_tail_is_token_capped() -> Result<()> {
1634 let provider = Arc::new(MockProvider::new(
1635 "Project summary with a long context and technical context.",
1636 ));
1637 let config = CompactionConfig::default()
1638 .with_retain_recent(8)
1639 .with_min_messages(1)
1640 .with_threshold_tokens(1);
1641 let compactor = LlmContextCompactor::new(provider, config);
1642
1643 let mut messages = Vec::new();
1644
1645 messages.extend((0..6).map(|index| Message::user(format!("pre-compaction noise {index}"))));
1647
1648 messages.extend(
1650 (0..8).map(|index| Message::assistant(format!("kept-{index}: {}", "x".repeat(12_000)))),
1651 );
1652
1653 let result = compactor.compact_history(messages).await?;
1654
1655 let retained_tail = &result.messages[2..];
1657 assert!(retained_tail.len() < 8);
1658
1659 let mut latest_index = -1i32;
1660 let mut all_retained = true;
1661 for message in retained_tail {
1662 if let Content::Text(text) = &message.content {
1663 if let Some(number) = text.split(':').next().and_then(|prefix| {
1664 prefix
1665 .strip_prefix("kept-")
1666 .and_then(|rest| rest.parse::<i32>().ok())
1667 }) {
1668 if number >= 0 {
1669 latest_index = latest_index.max(number);
1670 }
1671 } else {
1672 all_retained = false;
1673 }
1674 } else {
1675 all_retained = false;
1676 }
1677 }
1678
1679 assert!(all_retained);
1680 assert_eq!(latest_index, 7);
1681 assert!(
1682 TokenEstimator::estimate_history(retained_tail)
1683 <= compactor.config().max_retained_tail_tokens
1684 );
1685 assert!(compactor.needs_compaction(&result.messages));
1686
1687 Ok(())
1688 }
1689
1690 #[tokio::test]
1691 async fn test_compact_history_skips_summary_ack_when_retained_tail_is_empty() -> Result<()> {
1692 let provider = Arc::new(MockProvider::new("Summary for oversized user turn."));
1693 let config = CompactionConfig::default()
1694 .with_retain_recent(1)
1695 .with_min_messages(1)
1696 .with_threshold_tokens(1);
1697 let compactor = LlmContextCompactor::new(provider, config);
1698
1699 let messages = vec![
1700 Message::assistant("Earlier assistant context."),
1701 Message::user(format!("oversized-user-turn: {}", "x".repeat(200_000))),
1702 ];
1703
1704 let result = compactor.compact_history(messages).await?;
1705
1706 assert_eq!(result.new_count, 1);
1707 assert_eq!(result.messages.len(), 1);
1708
1709 let only_message = &result.messages[0];
1710 assert_eq!(only_message.role, Role::User);
1711
1712 if let Content::Text(text) = &only_message.content {
1713 assert!(text.contains("Previous conversation summary"));
1714 assert!(!text.contains(SUMMARY_ACKNOWLEDGMENT));
1715 } else {
1716 panic!("Expected summary text when retained tail is empty");
1717 }
1718
1719 Ok(())
1720 }
1721
1722 fn message_contains(message: &Message, needle: &str) -> bool {
1723 match &message.content {
1724 Content::Text(text) => text.contains(needle),
1725 Content::Blocks(blocks) => blocks.iter().any(|block| match block {
1726 ContentBlock::Text { text } => text.contains(needle),
1727 _ => false,
1728 }),
1729 }
1730 }
1731
1732 #[tokio::test]
1733 async fn test_epoch_one_facts_survive_two_compactions() -> Result<()> {
1734 const EPOCH1_FACT: &str = "EPOCH1_FACT: the API key lives in config/secrets.toml";
1740
1741 let requests = Arc::new(Mutex::new(Vec::new()));
1742 let provider = Arc::new(MockProvider::new_echo(requests.clone()));
1743 let config = CompactionConfig::default()
1744 .with_retain_recent(2)
1745 .with_min_messages(1);
1746 let compactor = LlmContextCompactor::new(provider, config);
1747
1748 let epoch1 = vec![
1749 Message::user(EPOCH1_FACT),
1750 Message::assistant("Understood, noted the secrets path."),
1751 Message::user("Now add error handling to main.rs."),
1752 Message::assistant("Added error handling to main.rs."),
1753 Message::user("latest user message one"),
1754 Message::assistant("latest assistant message two"),
1755 ];
1756
1757 let first = compactor.compact_history(epoch1).await?;
1758 assert!(
1759 first
1760 .messages
1761 .iter()
1762 .any(|m| message_contains(m, "EPOCH1_FACT")),
1763 "epoch-1 fact must be captured in the first summary"
1764 );
1765
1766 let mut epoch2 = first.messages;
1768 epoch2.push(Message::user("Another later turn."));
1769 epoch2.push(Message::assistant("Reply to the later turn."));
1770 epoch2.push(Message::user("Final turn a."));
1771 epoch2.push(Message::assistant("Final turn b."));
1772
1773 let second = compactor.compact_history(epoch2).await?;
1774
1775 assert!(
1776 second
1777 .messages
1778 .iter()
1779 .any(|m| message_contains(m, "EPOCH1_FACT")),
1780 "epoch-1 fact must survive the second compaction"
1781 );
1782
1783 let recorded = requests
1786 .lock()
1787 .map_err(|_| anyhow::anyhow!("request log poisoned"))?;
1788 assert!(
1789 recorded.iter().any(|req| req.contains("EPOCH1_FACT")),
1790 "prior summary carrying the epoch-1 fact must reach the summarizer"
1791 );
1792 drop(recorded);
1793
1794 Ok(())
1795 }
1796
1797 #[tokio::test]
1798 async fn test_compact_history_long_tool_chain_respects_token_cap() -> Result<()> {
1799 let provider = Arc::new(MockProvider::new("Summary of the early tool chain."));
1805 let cap = 20_000;
1806 let config = CompactionConfig::default()
1809 .with_retain_recent(18)
1810 .with_min_messages(1)
1811 .with_threshold_tokens(1)
1812 .with_max_retained_tail_tokens(cap);
1813 let compactor = LlmContextCompactor::new(provider, config);
1814
1815 let mut messages = Vec::new();
1818 for i in 0..10 {
1819 messages.push(Message {
1820 role: Role::Assistant,
1821 content: Content::Blocks(vec![ContentBlock::ToolUse {
1822 id: format!("tool_{i}"),
1823 name: "run".to_string(),
1824 input: serde_json::json!({ "arg": "y".repeat(12_000) }),
1825 thought_signature: None,
1826 }]),
1827 });
1828 messages.push(Message {
1829 role: Role::User,
1830 content: Content::Blocks(vec![ContentBlock::ToolResult {
1831 tool_use_id: format!("tool_{i}"),
1832 content: format!("result-{i}: {}", "z".repeat(12_000)),
1833 is_error: None,
1834 }]),
1835 });
1836 }
1837
1838 let full_tokens = TokenEstimator::estimate_history(&messages);
1839 assert!(
1840 full_tokens > cap * 2,
1841 "test setup: full chain must far exceed the cap"
1842 );
1843
1844 let result = compactor.compact_history(messages).await?;
1845
1846 let retained_tail = &result.messages[2..];
1849
1850 let tail_tokens = TokenEstimator::estimate_history(retained_tail);
1851 assert!(
1854 tail_tokens <= cap + 8_000,
1855 "retained tail {tail_tokens} should be bounded by the cap {cap}, not the whole chain"
1856 );
1857 assert!(
1858 retained_tail.len() < 20,
1859 "compaction must have summarized part of the chain"
1860 );
1861
1862 Ok(())
1863 }
1864
1865 #[tokio::test]
1866 async fn test_compact_warns_and_marks_truncated_summary() -> Result<()> {
1867 let requests = Arc::new(Mutex::new(Vec::new()));
1872 let provider = Arc::new(MockProvider::new_truncating(
1873 "partial summary cut off mid-",
1874 requests.clone(),
1875 ));
1876 let config = CompactionConfig::default().with_min_messages(1);
1877 let compactor = LlmContextCompactor::new(provider, config);
1878
1879 let messages = vec![
1880 Message::user("Some content that needs summarizing."),
1881 Message::assistant("More content to summarize here."),
1882 ];
1883
1884 let summary = compactor.compact(&messages).await?;
1885
1886 assert!(
1887 summary.contains("[summary truncated"),
1888 "a persistently truncated summary must carry a truncation marker"
1889 );
1890
1891 let recorded = requests
1893 .lock()
1894 .map_err(|_| anyhow::anyhow!("request log poisoned"))?;
1895 assert_eq!(recorded.len(), 2, "truncation should trigger one retry");
1896 drop(recorded);
1897
1898 Ok(())
1899 }
1900
1901 fn summarizable_messages() -> Vec<Message> {
1904 vec![
1905 Message::user("First question with enough words to summarize meaningfully."),
1906 Message::assistant("First answer, also carrying plenty of prose to compact."),
1907 Message::user("Second question continuing the earlier conversation topic."),
1908 Message::assistant("Second answer expanding on the topic at some length."),
1909 Message::user("Third question?"),
1910 Message::assistant("Third answer."),
1911 ]
1912 }
1913
1914 struct BlockRequestHooks;
1915
1916 #[async_trait]
1917 impl crate::hooks::AgentHooks for BlockRequestHooks {
1918 async fn pre_llm_request(&self, _request: &ChatRequest) -> RequestDecision {
1919 RequestDecision::Block("summaries are not allowed".to_string())
1920 }
1921 }
1922
1923 struct ModifyRequestHooks;
1924
1925 #[async_trait]
1926 impl crate::hooks::AgentHooks for ModifyRequestHooks {
1927 async fn pre_llm_request(&self, request: &ChatRequest) -> RequestDecision {
1928 let mut modified = request.clone();
1929 modified.messages = vec![Message::user("MODIFIED_SUMMARY_PROMPT")];
1930 RequestDecision::Modify(Box::new(modified))
1931 }
1932 }
1933
1934 struct BlockResponseHooks;
1935
1936 #[async_trait]
1937 impl crate::hooks::AgentHooks for BlockResponseHooks {
1938 async fn on_llm_response(&self, _response: &ChatResponse) -> ResponseDecision {
1939 ResponseDecision::Block("summary leaks a secret".to_string())
1940 }
1941 }
1942
1943 struct RetryResponseHooks;
1944
1945 #[async_trait]
1946 impl crate::hooks::AgentHooks for RetryResponseHooks {
1947 async fn on_llm_response(&self, _response: &ChatResponse) -> ResponseDecision {
1948 ResponseDecision::RetryWithFeedback("try harder".to_string())
1949 }
1950 }
1951
1952 #[tokio::test]
1953 async fn blocking_request_hook_aborts_compaction_before_the_llm_call() -> Result<()> {
1954 let requests = Arc::new(Mutex::new(Vec::new()));
1955 let provider = Arc::new(MockProvider::new_with_request_log(
1956 "summary",
1957 requests.clone(),
1958 ));
1959 let config = CompactionConfig::default().with_retain_recent(2);
1960 let compactor = LlmContextCompactor::new(provider, config)
1961 .with_guardrail_hooks(Arc::new(BlockRequestHooks));
1962
1963 let error = match compactor.compact_history(summarizable_messages()).await {
1964 Ok(result) => anyhow::bail!("blocked compaction must not succeed: {result:?}"),
1965 Err(error) => error,
1966 };
1967 assert!(
1968 error.to_string().contains("blocked by guardrail"),
1969 "unexpected error: {error}"
1970 );
1971
1972 let recorded = requests
1973 .lock()
1974 .map_err(|_| anyhow::anyhow!("request log poisoned"))?;
1975 assert!(
1976 recorded.is_empty(),
1977 "a blocked request must never reach the provider"
1978 );
1979 drop(recorded);
1980 Ok(())
1981 }
1982
1983 #[tokio::test]
1984 async fn modify_request_hook_reaches_the_provider() -> Result<()> {
1985 let requests = Arc::new(Mutex::new(Vec::new()));
1986 let provider = Arc::new(MockProvider::new_with_request_log(
1987 "summary",
1988 requests.clone(),
1989 ));
1990 let config = CompactionConfig::default().with_retain_recent(2);
1991 let compactor = LlmContextCompactor::new(provider, config)
1992 .with_guardrail_hooks(Arc::new(ModifyRequestHooks));
1993
1994 let result = compactor.compact_history(summarizable_messages()).await?;
1995 assert!(result.new_count < result.original_count);
1996
1997 let recorded = requests
1998 .lock()
1999 .map_err(|_| anyhow::anyhow!("request log poisoned"))?;
2000 assert_eq!(
2001 recorded.as_slice(),
2002 ["MODIFIED_SUMMARY_PROMPT"],
2003 "the provider must receive the hook-modified request"
2004 );
2005 drop(recorded);
2006 Ok(())
2007 }
2008
2009 #[tokio::test]
2010 async fn blocked_response_aborts_compaction() -> Result<()> {
2011 let provider = Arc::new(MockProvider::new("a summary that leaks the secret"));
2012 let config = CompactionConfig::default().with_retain_recent(2);
2013 let compactor = LlmContextCompactor::new(provider, config)
2014 .with_guardrail_hooks(Arc::new(BlockResponseHooks));
2015
2016 let error = match compactor.compact_history(summarizable_messages()).await {
2017 Ok(result) => anyhow::bail!("blocked summary must not be returned: {result:?}"),
2018 Err(error) => error,
2019 };
2020 assert!(
2021 error.to_string().contains("blocked by guardrail"),
2022 "unexpected error: {error}"
2023 );
2024 Ok(())
2025 }
2026
2027 #[tokio::test]
2028 async fn retry_with_feedback_response_aborts_compaction_without_retrying() -> Result<()> {
2029 let requests = Arc::new(Mutex::new(Vec::new()));
2030 let provider = Arc::new(MockProvider::new_with_request_log(
2031 "summary",
2032 requests.clone(),
2033 ));
2034 let config = CompactionConfig::default().with_retain_recent(2);
2035 let compactor = LlmContextCompactor::new(provider, config)
2036 .with_guardrail_hooks(Arc::new(RetryResponseHooks));
2037
2038 let error = match compactor.compact_history(summarizable_messages()).await {
2039 Ok(result) => anyhow::bail!("rejected summary must not be returned: {result:?}"),
2040 Err(error) => error,
2041 };
2042 assert!(
2043 error.to_string().contains("not retried during compaction"),
2044 "unexpected error: {error}"
2045 );
2046
2047 let recorded = requests
2049 .lock()
2050 .map_err(|_| anyhow::anyhow!("request log poisoned"))?;
2051 assert_eq!(recorded.len(), 1, "RetryWithFeedback must not retry");
2052 drop(recorded);
2053 Ok(())
2054 }
2055
2056 #[tokio::test]
2057 async fn compact_history_reports_summarization_usage() -> Result<()> {
2058 let provider = Arc::new(MockProvider::new("summary"));
2059 let config = CompactionConfig::default().with_retain_recent(2);
2060 let compactor = LlmContextCompactor::new(provider, config);
2061
2062 let result = compactor.compact_history(summarizable_messages()).await?;
2063 assert_eq!(result.llm_usage.input_tokens, 100);
2065 assert_eq!(result.llm_usage.output_tokens, 50);
2066 Ok(())
2067 }
2068}