1use super::config::ContextConfig;
2use super::partitions::ContextPartitions;
3use super::pressure::PressureAction;
4use super::token_engine::ContextTokenEngine;
5use super::units::{strict_tool_pairing_is_valid, unit_boundaries};
6use super::utility::{UtilitySelectionContext, plan_utility_archive};
7use crate::types::message::{Content, ContentPart, Message};
8
9#[derive(Default)]
11pub struct CompressResult {
12 pub tokens_saved: u32,
14 pub summary: Option<String>,
16 pub archived: Vec<Message>,
18 pub prefix_invalidated_at: Option<usize>,
22}
23
24pub trait Compressor: Send + Sync {
26 fn compress(
27 &self,
28 partitions: &mut ContextPartitions,
29 target_tokens: u32,
30 max_tokens: u32,
31 preserve_k: usize,
32 engine: &ContextTokenEngine,
33 ) -> CompressResult;
34}
35
36pub struct SnipCompactor {
38 pub per_msg_ratio: f64,
39}
40
41impl Compressor for SnipCompactor {
42 fn compress(
43 &self,
44 partitions: &mut ContextPartitions,
45 _target_tokens: u32,
46 max_tokens: u32,
47 preserve_k: usize,
48 engine: &ContextTokenEngine,
49 ) -> CompressResult {
50 let per_msg_limit = ((max_tokens as f64 * self.per_msg_ratio) as u32).max(50);
51 let mut saved = 0u32;
52 let partition = &mut partitions.history;
53 let prefix_keep = prefix_keep_for(partition.messages.len(), preserve_k);
59 let indices =
60 oversized_text_message_indices(&partition.messages, per_msg_limit, prefix_keep, engine);
61
62 for &i in &indices {
63 let msg = &mut partition.messages[i];
64 let original_tokens = msg.token_count.unwrap_or_else(|| engine.count_message(msg));
65 let head_limit = per_msg_limit / 2;
66 let tail_limit = per_msg_limit.saturating_sub(head_limit);
67 let snipped = if let Content::Text(ref t) = msg.content {
70 Some(excerpt_text_with_total(
71 t,
72 head_limit,
73 tail_limit,
74 engine,
75 original_tokens,
76 ))
77 } else {
78 None
79 };
80 if let Some(text) = snipped {
81 msg.content = Content::Text(text);
82 msg.token_count = Some(per_msg_limit);
83 saved += original_tokens.saturating_sub(per_msg_limit);
84 }
85 }
86
87 partition.token_count = partition.token_count.saturating_sub(saved);
88
89 CompressResult {
92 tokens_saved: saved,
93 prefix_invalidated_at: indices.iter().min().copied(),
94 ..Default::default()
95 }
96 }
97}
98
99fn prefix_keep_for(len: usize, preserve_k: usize) -> usize {
109 if len >= preserve_k.saturating_mul(3) {
110 preserve_k
111 } else {
112 0
113 }
114}
115
116fn oversized_text_message_indices(
117 messages: &[Message],
118 per_msg_limit: u32,
119 prefix_keep: usize,
120 engine: &ContextTokenEngine,
121) -> Vec<usize> {
122 messages
123 .iter()
124 .enumerate()
125 .filter(|(i, msg)| {
126 if *i < prefix_keep {
130 return false;
131 }
132 if !matches!(msg.content, Content::Text(_)) {
133 return false;
134 }
135 let toks = msg.token_count.unwrap_or_else(|| engine.count_message(msg));
136 toks > per_msg_limit && toks > 10
137 })
138 .map(|(i, _)| i)
139 .collect()
140}
141
142fn extract_json_excerpt(output: &str) -> Option<String> {
144 let val: serde_json::Value = serde_json::from_str(output).ok()?;
145 match val {
146 serde_json::Value::Object(map) => {
147 let mut summary_parts = Vec::new();
148 let mut keys = Vec::new();
149 for (k, v) in &map {
150 keys.push(k.as_str());
151 if v.is_number() || v.is_boolean() {
152 summary_parts.push(format!("{}: {}", k, v));
153 } else if let Some(s) = v.as_str() {
154 if s.len() <= 50 {
155 summary_parts.push(format!("{}: \"{}\"", k, s));
156 }
157 }
158 }
159 Some(format!(
160 "JSON Keys: [{}]\nJSON Fields: {{{}}}",
161 keys.join(", "),
162 summary_parts.join(", ")
163 ))
164 }
165 serde_json::Value::Array(arr) => {
166 if arr.is_empty() {
167 return Some("JSON Array: []".to_string());
168 }
169 let mut headers = Vec::new();
170 if let Some(serde_json::Value::Object(first_map)) = arr.first() {
171 for k in first_map.keys() {
172 headers.push(k.as_str());
173 }
174 }
175 let len = arr.len();
176 Some(format!(
177 "JSON Array: {} items. Keys: [{}]",
178 len,
179 headers.join(", ")
180 ))
181 }
182 _ => None,
183 }
184}
185
186fn excerpt_text(
188 text: &str,
189 head_tokens: u32,
190 tail_tokens: u32,
191 engine: &ContextTokenEngine,
192) -> String {
193 excerpt_text_with_total(text, head_tokens, tail_tokens, engine, engine.count(text))
194}
195
196fn excerpt_text_with_total(
199 text: &str,
200 head_tokens: u32,
201 tail_tokens: u32,
202 engine: &ContextTokenEngine,
203 total_tokens: u32,
204) -> String {
205 if total_tokens <= head_tokens + tail_tokens {
206 return text.to_string();
207 }
208 let head = engine.truncate(text, head_tokens);
209
210 let chars: Vec<char> = text.chars().collect();
211 let mut low = head.chars().count();
212 let mut high = chars.len();
213 let mut suffix_start = chars.len();
214 while low <= high {
215 let mid = (low + high) / 2;
216 if mid >= chars.len() {
217 break;
218 }
219 let candidate: String = chars[mid..].iter().collect();
220 let tokens = engine.count(&candidate);
221 if tokens <= tail_tokens {
222 suffix_start = mid;
223 if mid == 0 {
224 break;
225 }
226 high = mid - 1;
227 } else {
228 low = mid + 1;
229 }
230 }
231 let tail: String = chars[suffix_start..].iter().collect();
232 let remaining = total_tokens
233 .saturating_sub(head_tokens)
234 .saturating_sub(tail_tokens);
235 format!("{}… [… {} tokens omitted …] …{}", head, remaining, tail)
236}
237
238fn excerptable_tool_result_indices(
243 messages: &[Message],
244 preserved_refs: &[String],
245 prefix_keep: usize,
246 engine: &ContextTokenEngine,
247) -> Vec<usize> {
248 messages
249 .iter()
250 .enumerate()
251 .filter_map(|(i, msg)| {
252 if i < prefix_keep {
255 return None;
256 }
257 let toks = msg.token_count.unwrap_or_else(|| engine.count_message(msg));
258 if toks < 200 {
259 return None;
260 }
261 let Content::Parts(parts) = &msg.content else {
262 return None;
263 };
264 parts
265 .iter()
266 .any(|part| {
267 matches!(
268 part,
269 ContentPart::ToolResult { call_id, .. }
270 if !preserved_refs.iter().any(|preserved| preserved == call_id.as_str())
271 )
272 })
273 .then_some(i)
274 })
275 .collect()
276}
277
278pub struct MicroCompactor;
281
282impl Compressor for MicroCompactor {
283 fn compress(
284 &self,
285 partitions: &mut ContextPartitions,
286 _target_tokens: u32,
287 _max_tokens: u32,
288 preserve_k: usize,
289 engine: &ContextTokenEngine,
290 ) -> CompressResult {
291 let find_tool_name = |call_id: &str, msgs: &[Message]| -> Option<String> {
292 for m in msgs {
293 for tc in &m.tool_calls {
294 if tc.id == call_id {
295 return Some(tc.name.to_string());
296 }
297 }
298 }
299 None
300 };
301
302 let prefix_keep = prefix_keep_for(partitions.history.messages.len(), preserve_k);
305 let indices = excerptable_tool_result_indices(
306 &partitions.history.messages,
307 &partitions.task_state.preserved_refs,
308 prefix_keep,
309 engine,
310 );
311 let messages_clone = partitions.history.messages.clone();
312 let preserved_refs = partitions.task_state.preserved_refs.clone();
313 let partition = &mut partitions.history;
314 let mut saved = 0u32;
315
316 for &i in &indices {
317 let msg = &mut partition.messages[i];
318 let original_tokens = msg.token_count.unwrap_or_else(|| engine.count_message(msg));
319 if let Content::Parts(ref mut parts) = msg.content {
320 for part in parts.iter_mut() {
321 if let ContentPart::ToolResult {
322 call_id,
323 output,
324 is_error: _,
325 durable_content,
326 ..
327 } = part
328 {
329 if preserved_refs
330 .iter()
331 .any(|preserved| preserved == call_id.as_str())
332 {
333 continue;
334 }
335 let original_output_tokens = engine.count(output);
336 let tool_name = find_tool_name(call_id, &messages_clone)
337 .unwrap_or_else(|| "unknown".to_string());
338
339 let new_output = if original_output_tokens > 2000 {
340 if let Some(json_excerpt) = extract_json_excerpt(output) {
341 format!(
342 "[tool result: {} | {} | {} tokens]\n{}",
343 call_id, tool_name, original_output_tokens, json_excerpt
344 )
345 } else {
346 let excerpt = excerpt_text(output, 30, 10, engine);
347 format!(
348 "[tool result: {} | {} | {} tokens]\n{}",
349 call_id, tool_name, original_output_tokens, excerpt
350 )
351 }
352 } else {
353 let excerpt = excerpt_text(output, 150, 50, engine);
354 format!(
355 "[tool result: {} | {} | {} tokens]\n{}",
356 call_id, tool_name, original_output_tokens, excerpt
357 )
358 };
359
360 *output = new_output;
361 *durable_content = None;
365 }
366 }
367 let new_tokens = engine.count_message(msg);
368 msg.token_count = Some(new_tokens);
369 saved += original_tokens.saturating_sub(new_tokens);
370 }
371 }
372
373 partition.token_count = partition.token_count.saturating_sub(saved);
374
375 CompressResult {
377 tokens_saved: saved,
378 prefix_invalidated_at: indices.iter().min().copied(),
379 ..Default::default()
380 }
381 }
382}
383
384pub fn plan_drop_oldest(
391 messages: &[Message],
392 total_tokens: u32,
393 target_tokens: u32,
394 keep: usize,
395 engine: &ContextTokenEngine,
396) -> (usize, u32) {
397 let units = unit_boundaries(messages);
398 let limit = units.len().saturating_sub(keep);
399 let mut saved = 0u32;
400 let mut n = 0usize;
401 for unit in units.iter().take(limit) {
402 if total_tokens.saturating_sub(saved) <= target_tokens {
403 break;
404 }
405 saved += messages[unit.clone()]
406 .iter()
407 .map(|msg| msg.token_count.unwrap_or_else(|| engine.count_message(msg)))
408 .sum::<u32>();
409 n = unit.end;
410 }
411 (n, saved)
412}
413
414pub struct CollapseCompactor;
417
418impl Compressor for CollapseCompactor {
419 fn compress(
420 &self,
421 partitions: &mut ContextPartitions,
422 target_tokens: u32,
423 _max_tokens: u32,
424 preserve_k: usize,
425 engine: &ContextTokenEngine,
426 ) -> CompressResult {
427 let non_history_tokens = partitions
428 .total_tokens(engine)
429 .saturating_sub(partitions.history.token_count);
430 let history_target = target_tokens.saturating_sub(non_history_tokens);
431 let plan = plan_utility_archive(
432 &partitions.history.messages,
433 partitions.history.token_count,
434 history_target,
435 preserve_k,
436 engine,
437 &UtilitySelectionContext {
438 goal: &partitions.task_state.goal,
439 criteria: &partitions.task_state.criteria,
440 preserved_refs: &partitions.task_state.preserved_refs,
441 active_directives: &partitions.task_state.directives,
442 },
443 );
444 if plan.archived_ranges.is_empty() {
445 return CompressResult::default();
446 }
447 let prefix_invalidated_at = plan.archived_ranges.iter().map(|range| range.start).min();
448 let (archived, saved) = apply_utility_plan(&mut partitions.history, &plan);
449
450 CompressResult {
453 tokens_saved: saved,
454 archived,
455 prefix_invalidated_at,
456 ..Default::default()
457 }
458 }
459}
460
461pub struct AutoCompactor;
463
464impl Compressor for AutoCompactor {
465 fn compress(
466 &self,
467 partitions: &mut ContextPartitions,
468 target_tokens: u32,
469 _max_tokens: u32,
470 preserve_k: usize,
471 engine: &ContextTokenEngine,
472 ) -> CompressResult {
473 if partitions.history.messages.is_empty() {
474 return CompressResult::default();
475 }
476 let non_history_tokens = partitions
477 .total_tokens(engine)
478 .saturating_sub(partitions.history.token_count);
479 let history_target = target_tokens.saturating_sub(non_history_tokens);
480 let plan = plan_utility_archive(
481 &partitions.history.messages,
482 partitions.history.token_count,
483 history_target,
484 preserve_k,
485 engine,
486 &UtilitySelectionContext {
487 goal: &partitions.task_state.goal,
488 criteria: &partitions.task_state.criteria,
489 preserved_refs: &partitions.task_state.preserved_refs,
490 active_directives: &partitions.task_state.directives,
491 },
492 );
493 if plan.archived_ranges.is_empty() {
494 return CompressResult::default();
495 }
496 let prefix_invalidated_at = plan.archived_ranges.iter().map(|range| range.start).min();
497 let (archived, saved) = apply_utility_plan(&mut partitions.history, &plan);
498
499 CompressResult {
502 tokens_saved: saved,
503 archived,
504 prefix_invalidated_at,
505 ..Default::default()
506 }
507 }
508}
509
510fn apply_utility_plan(
511 partition: &mut super::partitions::Partition,
512 plan: &super::utility::UtilityArchivePlan,
513) -> (Vec<Message>, u32) {
514 let pairing_was_valid = strict_tool_pairing_is_valid(&partition.messages);
515 let archived_indices = plan
516 .archived_ranges
517 .iter()
518 .flat_map(|range| range.clone())
519 .collect::<std::collections::BTreeSet<_>>();
520 let mut archived = Vec::new();
521 let mut retained = Vec::new();
522 for (index, message) in std::mem::take(&mut partition.messages)
523 .into_iter()
524 .enumerate()
525 {
526 if archived_indices.contains(&index) {
527 archived.push(message);
528 } else {
529 retained.push(message);
530 }
531 }
532 partition.messages = retained;
533 partition.token_count = plan.retained_tokens;
534 debug_assert!(
535 !pairing_was_valid || strict_tool_pairing_is_valid(&partition.messages),
536 "utility selection split a valid tool transaction"
537 );
538 (archived, plan.archived_tokens)
539}
540
541pub struct CompressionPipeline {
543 stages: Vec<(PressureAction, Box<dyn Compressor>)>,
544 preserve_recent_turns: usize,
545}
546
547impl CompressionPipeline {
548 pub fn new(config: &ContextConfig) -> Self {
549 Self {
550 preserve_recent_turns: config.preserve_recent_turns,
551 stages: vec![
552 (
553 PressureAction::SnipCompact,
554 Box::new(SnipCompactor {
555 per_msg_ratio: config.snip_per_msg_ratio,
556 }),
557 ),
558 (PressureAction::MicroCompact, Box::new(MicroCompactor)),
559 (PressureAction::ContextCollapse, Box::new(CollapseCompactor)),
560 (PressureAction::AutoCompact, Box::new(AutoCompactor)),
561 ],
562 }
563 }
564
565 pub fn compress(
566 &self,
567 partitions: &mut ContextPartitions,
568 action: PressureAction,
569 max_tokens: u32,
570 target_tokens: u32,
571 engine: &ContextTokenEngine,
572 ) -> (u32, Option<String>, Vec<Message>, Option<usize>) {
573 if action == PressureAction::None {
574 return (0, None, vec![], None);
575 }
576
577 let mut total_saved = 0;
578 let mut all_archived = vec![];
579 let mut cache_at: Option<usize> = None;
582 let summarizer = super::summarizer::RuleSummarizer;
583
584 for (stage_action, compressor) in &self.stages {
585 if *stage_action <= action {
586 if partitions.total_tokens(engine) <= target_tokens {
587 break;
588 }
589 let res = compressor.compress(
590 partitions,
591 target_tokens,
592 max_tokens,
593 self.preserve_recent_turns,
594 engine,
595 );
596 total_saved += res.tokens_saved;
597 cache_at = [cache_at, res.prefix_invalidated_at]
598 .into_iter()
599 .flatten()
600 .min();
601 all_archived.extend(res.archived);
602 }
603 }
604
605 let summary_budget = if target_tokens == 0 {
617 max_tokens
618 } else {
619 target_tokens
620 };
621 let summary = if all_archived.is_empty() {
622 None
623 } else {
624 let s = summarizer.summarize(&all_archived, action, summary_budget);
625 partitions
626 .task_state
627 .log_compression(action.label(), s.clone());
628 Some(s)
629 };
630
631 (total_saved, summary, all_archived, cache_at)
632 }
633}
634
635#[cfg(test)]
636mod tests {
637 use super::*;
638 use crate::context::config::ContextConfig;
639 use crate::context::partitions::ContextPartitions;
640 use crate::context::token_engine::ContextTokenEngine;
641 use crate::types::message::Message;
642
643 fn engine() -> ContextTokenEngine {
644 ContextTokenEngine::char_approx()
645 }
646 fn config() -> ContextConfig {
647 ContextConfig::default()
648 }
649 const MAX: u32 = 1_000;
650
651 #[test]
652 fn snip_compactor_truncates_oversized_messages() {
653 let cfg = ContextConfig {
654 snip_per_msg_ratio: 0.10,
655 ..Default::default()
656 };
657 let compactor = SnipCompactor {
658 per_msg_ratio: cfg.snip_per_msg_ratio,
659 };
660 let mut ctx = ContextPartitions::new(&cfg);
661 ctx.history.push(Message::user("a".repeat(800)), 200);
662 let result = compactor.compress(&mut ctx, 0, MAX, 0, &engine());
664 assert!(result.tokens_saved > 0);
665 if let Content::Text(ref t) = ctx.history.messages[0].content {
666 assert!(t.contains("… [… 100 tokens omitted …] …"), "got: {t}");
667 }
668 }
669
670 #[test]
671 fn snip_compactor_leaves_small_messages_untouched() {
672 let cfg = ContextConfig {
673 snip_per_msg_ratio: 0.10,
674 ..Default::default()
675 };
676 let compactor = SnipCompactor {
677 per_msg_ratio: cfg.snip_per_msg_ratio,
678 };
679 let mut ctx = ContextPartitions::new(&cfg);
680 ctx.history.push(Message::user("short"), 5);
681 let result = compactor.compress(&mut ctx, 0, MAX, 2, &engine());
682 assert_eq!(result.tokens_saved, 0);
683 }
684
685 #[test]
686 fn micro_compactor_replaces_tool_results_with_measured_placeholder() {
687 use crate::types::message::{ContentPart, Role};
688 use compact_str::CompactString;
689
690 let compactor = MicroCompactor;
691 let mut ctx = ContextPartitions::new(&config());
692 let parts = vec![ContentPart::ToolResult {
693 call_id: CompactString::new("c1"),
694 output: "a".repeat(1200),
695 is_error: false,
696 durable_content: None,
697 }];
698 let msg = Message {
699 role: Role::Tool,
700 content: Content::Parts(parts),
701 tool_calls: vec![],
702 token_count: Some(300),
703 };
704 ctx.history.messages.push(msg);
705 ctx.history.token_count = 300;
706
707 let result = compactor.compress(&mut ctx, 0, MAX, 0, &engine());
709 assert!(result.tokens_saved > 0);
710 let Content::Parts(parts) = &ctx.history.messages[0].content else {
711 panic!("tool-result envelope must survive compaction");
712 };
713 let text = parts
714 .iter()
715 .find_map(|part| match part {
716 ContentPart::ToolResult {
717 call_id, output, ..
718 } if call_id.as_str() == "c1" => Some(output.as_str()),
719 _ => None,
720 })
721 .expect("correlated tool result remains present");
722 assert!(
723 text.contains("[tool result: c1 | unknown | 300 tokens]"),
724 "got: {text}"
725 );
726 }
727
728 #[test]
729 fn micro_compactor_recounts_the_entire_mixed_parts_envelope() {
730 use crate::types::message::{ContentPart, Role};
731
732 let compactor = MicroCompactor;
733 let mut ctx = ContextPartitions::new(&config());
734 let msg = Message {
735 role: Role::Tool,
736 content: Content::Parts(vec![
737 ContentPart::Text {
738 text: "metadata that must remain budgeted".into(),
739 },
740 ContentPart::ToolResult {
741 call_id: "c1".into(),
742 output: "a".repeat(1200),
743 is_error: false,
744 durable_content: None,
745 },
746 ContentPart::ToolResult {
747 call_id: "c2".into(),
748 output: "b".repeat(1000),
749 is_error: false,
750 durable_content: None,
751 },
752 ]),
753 tool_calls: vec![],
754 token_count: None,
755 };
756 let original = engine().count_message(&msg);
757 ctx.history.push(msg, original);
758
759 compactor.compress(&mut ctx, 0, MAX, 0, &engine());
760
761 let message = &ctx.history.messages[0];
762 let recounted = engine().count_message(message);
763 assert_eq!(message.token_count, Some(recounted));
764 assert_eq!(ctx.history.token_count, recounted);
765 let Content::Parts(parts) = &message.content else {
766 panic!("parts preserved")
767 };
768 assert_eq!(
769 parts
770 .iter()
771 .filter(|part| matches!(part, ContentPart::ToolResult { output, .. } if output.starts_with("[tool result:")))
772 .count(),
773 2,
774 "every eligible tool result is excerpted"
775 );
776 }
777
778 #[test]
779 fn collapse_compactor_drops_oldest_to_reach_target() {
780 let compactor = CollapseCompactor;
781 let mut ctx = ContextPartitions::new(&config());
782 for _ in 0..8 {
783 ctx.history.push(Message::user("msg"), 50);
784 }
785 let result = compactor.compress(&mut ctx, 250, MAX, 2, &engine());
786 assert!(result.tokens_saved > 0);
787 assert!(ctx.history.messages.len() < 8);
788 assert!(
791 !result.archived.is_empty(),
792 "drained messages are returned to the pipeline"
793 );
794 assert!(
795 result.summary.is_none(),
796 "compactor no longer self-summarizes"
797 );
798 assert!(
799 ctx.task_state.compression_log.is_empty(),
800 "compactor no longer logs"
801 );
802 }
803
804 #[test]
805 fn collapse_utility_selection_changes_the_archived_units_not_only_the_score() {
806 let compactor = CollapseCompactor;
807 let mut ctx = ContextPartitions::new(&config());
808 ctx.task_state.goal = "ship ORCHID release".into();
809 for (user, assistant) in [
810 (
811 "ORCHID release criterion",
812 "DECISION: retry failure; artifact /work/orchid.json",
813 ),
814 ("routine chatter one", "acknowledged"),
815 ("routine chatter two", "acknowledged"),
816 ("latest request", "working on it"),
817 ] {
818 ctx.history.push(Message::user(user), 40);
819 ctx.history.push(Message::assistant(assistant), 40);
820 }
821
822 let result = compactor.compress(&mut ctx, 160, MAX, 1, &engine());
823 let retained = ctx
824 .history
825 .messages
826 .iter()
827 .filter_map(|message| message.content.as_text())
828 .collect::<Vec<_>>()
829 .join("\n");
830 let archived = result
831 .archived
832 .iter()
833 .filter_map(|message| message.content.as_text())
834 .collect::<Vec<_>>()
835 .join("\n");
836 assert!(retained.contains("ORCHID"));
837 assert!(retained.contains("/work/orchid.json"));
838 assert!(!retained.contains("routine chatter"));
839 assert!(archived.contains("routine chatter one"));
840 assert!(archived.contains("routine chatter two"));
841 assert_eq!(ctx.history.token_count, 160);
842 }
843
844 #[test]
845 fn utility_selector_deducts_fixed_context_before_budgeting_history() {
846 let compactor = CollapseCompactor;
847 let mut ctx = ContextPartitions::new(&config());
848 ctx.system.push(Message::system("fixed"), 600);
849 for index in 0..4 {
850 ctx.history
851 .push(Message::user(format!("unit {index}")), 100);
852 }
853
854 let result = compactor.compress(&mut ctx, 700, MAX, 1, &engine());
855 assert_eq!(result.tokens_saved, 300);
856 assert_eq!(ctx.history.token_count, 100);
857 assert!(ctx.total_tokens(&engine()) <= 700);
858 }
859
860 #[test]
861 fn rule_summarizer_formats_correctly() {
862 use crate::context::summarizer::RuleSummarizer;
863 use crate::types::message::{Content, Message, Role};
864 let summarizer = RuleSummarizer;
865 let mut messages = vec![];
866 messages.push(Message {
867 role: Role::User,
868 content: Content::Text("hello".to_string()),
869 tool_calls: vec![],
870 token_count: Some(5),
871 });
872 messages.push(Message {
873 role: Role::Assistant,
874 content: Content::Text("world".to_string()),
875 tool_calls: vec![],
876 token_count: Some(6),
877 });
878 let summary = summarizer.summarize(&messages, PressureAction::SnipCompact, 100);
879 assert!(summary.contains("[Compressed: snip_compact]"));
880 assert!(summary.contains("archived_messages: 2; archived_tokens: 11"));
881 assert!(summary.contains("constraints:"));
882 }
883
884 #[test]
885 fn micro_compactor_preserves_refs_in_preserved_refs() {
886 use crate::types::message::{ContentPart, Role};
887 use compact_str::CompactString;
888
889 let compactor = MicroCompactor;
890 let mut ctx = ContextPartitions::new(&config());
891 ctx.task_state.preserved_refs = vec!["keep_me".to_string()];
892
893 let parts = vec![ContentPart::ToolResult {
894 call_id: CompactString::new("keep_me"),
895 output: "a".repeat(1200),
896 is_error: false,
897 durable_content: None,
898 }];
899 let msg = Message {
900 role: Role::Tool,
901 content: Content::Parts(parts),
902 tool_calls: vec![],
903 token_count: Some(300),
904 };
905 ctx.history.messages.push(msg);
906 ctx.history.token_count = 300;
907
908 let result = compactor.compress(&mut ctx, 0, MAX, 2, &engine());
909 assert_eq!(result.tokens_saved, 0);
911 let text_opt = ctx.history.messages[0].content.as_text();
912 assert!(
913 text_opt.is_none(),
914 "should not be replaced to text placeholder"
915 );
916 }
917
918 #[test]
919 fn auto_compactor_merges_all_except_last_two_turns() {
920 let compactor = AutoCompactor;
921 let mut ctx = ContextPartitions::new(&config());
922 for i in 0..10 {
923 ctx.history.push(Message::user(format!("msg {i}")), 10);
924 }
925 let result = compactor.compress(&mut ctx, 0, MAX, 2, &engine());
926 assert!(result.tokens_saved > 0);
927 assert_eq!(ctx.history.messages.len(), 2); assert!(
931 !result.archived.is_empty(),
932 "drained messages returned to the pipeline"
933 );
934 assert!(
935 result.summary.is_none(),
936 "compactor no longer self-summarizes"
937 );
938 assert!(
939 ctx.task_state.compression_log.is_empty(),
940 "compactor no longer logs"
941 );
942 }
943
944 #[test]
945 fn plan_drop_oldest_respects_target_and_preserve_floor() {
946 let msgs: Vec<Message> = (0..8)
949 .map(|i| {
950 let mut m = Message::user(format!("m{i}"));
951 m.token_count = Some(50);
952 m
953 })
954 .collect();
955 assert_eq!(plan_drop_oldest(&msgs, 400, 250, 2, &engine()), (3, 150));
957 assert_eq!(plan_drop_oldest(&msgs, 400, 0, 2, &engine()), (6, 300));
959 assert_eq!(plan_drop_oldest(&msgs, 400, 500, 2, &engine()), (0, 0));
961 }
962
963 #[test]
964 fn collapse_never_splits_a_tool_transaction() {
965 let mut call = Message::assistant("calling");
966 call.tool_calls.push(crate::types::message::ToolCall {
967 id: "call-1".into(),
968 name: "read".into(),
969 arguments: serde_json::json!({}),
970 });
971 let messages = vec![
972 Message::user("question"),
973 call,
974 Message::tool(vec![ContentPart::ToolResult {
975 call_id: "call-1".into(),
976 output: "ok".into(),
977 is_error: false,
978 durable_content: None,
979 }]),
980 Message::assistant("answer"),
981 Message::user("next"),
982 Message::assistant("done"),
983 ]
984 .into_iter()
985 .map(|mut message| {
986 message.token_count = Some(10);
987 message
988 })
989 .collect::<Vec<_>>();
990
991 assert_eq!(plan_drop_oldest(&messages, 60, 30, 1, &engine()), (4, 40));
992 }
993
994 #[test]
995 fn auto_compactor_preserves_the_latest_complete_tool_unit() {
996 let compactor = AutoCompactor;
997 let mut ctx = ContextPartitions::new(&config());
998 ctx.history.push(Message::user("old"), 10);
999 ctx.history.push(Message::assistant("old answer"), 10);
1000 let mut call = Message::assistant("calling");
1001 call.tool_calls.push(crate::types::message::ToolCall {
1002 id: "call-1".into(),
1003 name: "read".into(),
1004 arguments: serde_json::json!({}),
1005 });
1006 ctx.history.push(Message::user("question"), 10);
1007 ctx.history.push(call, 10);
1008 ctx.history.push(
1009 Message::tool(vec![ContentPart::ToolResult {
1010 call_id: "call-1".into(),
1011 output: "ok".into(),
1012 is_error: false,
1013 durable_content: None,
1014 }]),
1015 10,
1016 );
1017 ctx.history.push(Message::assistant("answer"), 10);
1018
1019 compactor.compress(&mut ctx, 0, MAX, 1, &engine());
1020
1021 assert_eq!(ctx.history.messages.len(), 4);
1022 assert_eq!(ctx.history.messages[0].content.as_text(), Some("question"));
1023 }
1024
1025 #[test]
1026 fn prefix_keep_yields_without_drop_fallback() {
1027 assert_eq!(prefix_keep_for(6, 2), 2, "len 6 >= 6 → protect oldest 2");
1030 assert_eq!(
1031 prefix_keep_for(5, 2),
1032 0,
1033 "len 5 < 6 → would leave an untouchable message"
1034 );
1035 assert_eq!(prefix_keep_for(3, 2), 0);
1036 assert_eq!(prefix_keep_for(0, 2), 0);
1037 }
1038
1039 #[test]
1040 fn pipeline_reports_accurate_prefix_invalidation() {
1041 let mut cfg = config();
1045 cfg.preserve_recent_turns = 1;
1046 let mut ctx = baseline_partitions();
1047 let (_s, _u, _a, cache_at) = CompressionPipeline::new(&cfg).compress(
1048 &mut ctx,
1049 PressureAction::SnipCompact,
1050 MAX,
1051 500,
1052 &engine(),
1053 );
1054 assert_eq!(
1055 cache_at,
1056 Some(1),
1057 "one protected unit leaves msg 1 as the earliest rewrite"
1058 );
1059
1060 let mut ctx2 = baseline_partitions();
1061 let (_s2, _u2, _a2, cache_at2) = CompressionPipeline::new(&cfg).compress(
1062 &mut ctx2,
1063 PressureAction::AutoCompact,
1064 MAX,
1065 500,
1066 &engine(),
1067 );
1068 assert_eq!(
1069 cache_at2,
1070 Some(0),
1071 "dropping the oldest breaks the cache prefix at 0"
1072 );
1073 }
1074
1075 use crate::types::message::Role;
1084 use compact_str::CompactString;
1085
1086 fn baseline_partitions() -> ContextPartitions {
1089 let cfg = config();
1090 let mut ctx = ContextPartitions::new(&cfg);
1091 ctx.history.push(Message::user("u0 ".repeat(120)), 300);
1093 ctx.history.push(Message::assistant("a0 ".repeat(120)), 300);
1094 ctx.history.messages.push(Message {
1096 role: Role::Tool,
1097 content: Content::Parts(vec![ContentPart::ToolResult {
1098 call_id: CompactString::new("call_1"),
1099 output: serde_json::json!({"rows": 42, "ok": true, "name": "alpha"}).to_string()
1100 + &"-pad".repeat(400),
1101 is_error: false,
1102 durable_content: None,
1103 }]),
1104 tool_calls: vec![],
1105 token_count: Some(400),
1106 });
1107 ctx.history.token_count += 400;
1108 ctx.history.push(Message::user("u1 ".repeat(120)), 300);
1109 ctx.history.push(Message::assistant("a1 ".repeat(120)), 300);
1110 ctx.history.messages.push(Message {
1111 role: Role::Tool,
1112 content: Content::Parts(vec![ContentPart::ToolResult {
1113 call_id: CompactString::new("call_2"),
1114 output: "y".repeat(1600),
1115 is_error: false,
1116 durable_content: None,
1117 }]),
1118 tool_calls: vec![],
1119 token_count: Some(400),
1120 });
1121 ctx.history.token_count += 400;
1122 ctx
1123 }
1124
1125 fn run_baseline(action: PressureAction) -> (u32, u32, Option<String>, usize, usize, u32) {
1128 let mut ctx = baseline_partitions();
1129 let before = ctx.total_tokens(&engine());
1130 let mut cfg = config();
1131 cfg.preserve_recent_turns = 1;
1132 let (saved, summary, archived, _cache_at) =
1133 CompressionPipeline::new(&cfg).compress(&mut ctx, action, MAX, 500, &engine());
1134 let archived_len = archived.len();
1135 let msgs_after = ctx.history.messages.len();
1136 let total_after = ctx.total_tokens(&engine());
1137 (
1138 before,
1139 saved,
1140 summary,
1141 archived_len,
1142 msgs_after,
1143 total_after,
1144 )
1145 }
1146
1147 #[test]
1148 fn baseline_snip_only_caps_text_no_archival() {
1149 let (before, saved, summary, archived, msgs, total) =
1153 run_baseline(PressureAction::SnipCompact);
1154 assert_eq!(before, 2000);
1155 assert_eq!(saved, 750, "3 non-prefix oversized messages × 250");
1156 assert_eq!(archived, 0);
1157 assert!(summary.is_none());
1158 assert_eq!(msgs, 6, "snip mutates in place, drops no messages");
1159 assert_eq!(total, 1250);
1160 }
1161
1162 #[test]
1163 fn baseline_micro_excerpts_tool_results() {
1164 let (before, saved, summary, archived, msgs, total) =
1166 run_baseline(PressureAction::MicroCompact);
1167 assert_eq!(before, 2000);
1168 assert_eq!(saved, 1112, "snip(750) + tool-result excerpts(362)");
1169 assert_eq!(archived, 0);
1170 assert!(summary.is_none());
1171 assert_eq!(msgs, 6);
1172 assert_eq!(total, 888);
1173 }
1174
1175 #[test]
1176 fn baseline_collapse_drops_oldest_and_summarizes() {
1177 let (before, saved, summary, archived, msgs, total) =
1180 run_baseline(PressureAction::ContextCollapse);
1181 assert_eq!(before, 2000);
1182 assert_eq!(saved, 1681);
1183 assert_eq!(archived, 3, "drops the complete oldest unit");
1184 assert_eq!(msgs, 3, "one complete recent unit remains");
1185 assert_eq!(total, 319);
1186 let summary = summary.expect("collapse summarizes archived messages");
1187 assert!(
1188 summary.contains("[Compressed: context_collapse]"),
1189 "summary routes the collapse action: {summary}"
1190 );
1191 }
1192
1193 #[test]
1194 fn baseline_auto_attributes_summary_to_auto_compact() {
1195 let (before, saved, summary, archived, msgs, total) =
1201 run_baseline(PressureAction::AutoCompact);
1202 assert_eq!(before, 2000);
1203 assert_eq!(saved, 1681);
1204 assert_eq!(archived, 3);
1205 assert_eq!(msgs, 3);
1206 assert_eq!(total, 319);
1207 let summary = summary.expect("auto-compact summarizes the archived messages");
1208 assert!(
1209 summary.contains("[Compressed: auto_compact]"),
1210 "got: {summary}"
1211 );
1212 }
1213
1214 #[test]
1215 fn baseline_saved_is_monotonic_in_action_level() {
1216 let snip = run_baseline(PressureAction::SnipCompact).1;
1218 let micro = run_baseline(PressureAction::MicroCompact).1;
1219 let collapse = run_baseline(PressureAction::ContextCollapse).1;
1220 let auto = run_baseline(PressureAction::AutoCompact).1;
1221 assert!(snip <= micro, "{snip} <= {micro}");
1222 assert!(micro <= collapse, "{micro} <= {collapse}");
1223 assert!(collapse <= auto, "{collapse} <= {auto}");
1224 }
1225
1226 #[test]
1227 fn pipeline_stops_cascade_when_target_reached() {
1228 let cfg = ContextConfig {
1229 snip_per_msg_ratio: 0.25,
1230 preserve_recent_turns: 0,
1233 ..Default::default()
1234 };
1235 let pipeline = CompressionPipeline::new(&cfg);
1236 let mut ctx = ContextPartitions::new(&cfg);
1237 ctx.history.push(Message::user("a".repeat(3600)), 900);
1238
1239 let (saved, summary, archived, _cache_at) =
1240 pipeline.compress(&mut ctx, PressureAction::AutoCompact, 1_000, 500, &engine());
1241
1242 assert!(saved > 0);
1243 assert!(
1244 summary.is_none(),
1245 "auto compactor should not run after snip reaches target"
1246 );
1247 assert!(
1248 archived.is_empty(),
1249 "heavier archival stages should not run"
1250 );
1251 assert_eq!(ctx.history.messages.len(), 1);
1252 assert!(ctx.total_tokens(&engine()) <= 500);
1253 }
1254
1255 struct Lcg(u64);
1266 impl Lcg {
1267 fn next_u64(&mut self) -> u64 {
1268 self.0 = self
1269 .0
1270 .wrapping_mul(6364136223846793005)
1271 .wrapping_add(1442695040888963407);
1272 self.0 ^ (self.0 >> 33)
1273 }
1274 fn below(&mut self, n: u64) -> u64 {
1275 self.next_u64() % n.max(1)
1276 }
1277 }
1278
1279 fn random_valid_history(rng: &mut Lcg, call_seq: &mut usize) -> Vec<(Message, u32)> {
1283 use crate::types::message::{ContentPart, ToolCall};
1284 let mut out: Vec<(Message, u32)> = Vec::new();
1285 let transactions = 3 + rng.below(10) as usize;
1286 let weight = |rng: &mut Lcg| 10 + rng.below(300) as u32;
1287 for _ in 0..transactions {
1288 if rng.below(10) < 6 {
1289 if rng.below(2) == 0 {
1291 out.push((Message::user(format!("ask {}", call_seq)), weight(rng)));
1292 }
1293 let n_calls = 1 + rng.below(3) as usize;
1294 let ids: Vec<String> = (0..n_calls)
1295 .map(|_| {
1296 *call_seq += 1;
1297 format!("call-{call_seq}")
1298 })
1299 .collect();
1300 let mut assistant = Message::assistant("working");
1301 for id in &ids {
1302 assistant.tool_calls.push(ToolCall {
1303 id: id.clone().into(),
1304 name: "read".into(),
1305 arguments: serde_json::json!({}),
1306 });
1307 }
1308 out.push((assistant, weight(rng)));
1309 if rng.below(2) == 0 {
1311 for id in &ids {
1312 out.push((
1313 Message::tool(vec![ContentPart::ToolResult {
1314 call_id: id.clone().into(),
1315 output: "ok ".repeat(20),
1316 is_error: false,
1317 durable_content: None,
1318 }]),
1319 weight(rng),
1320 ));
1321 }
1322 } else {
1323 let parts = ids
1324 .iter()
1325 .map(|id| ContentPart::ToolResult {
1326 call_id: id.clone().into(),
1327 output: "ok ".repeat(20),
1328 is_error: false,
1329 durable_content: None,
1330 })
1331 .collect();
1332 out.push((Message::tool(parts), weight(rng)));
1333 }
1334 if rng.below(2) == 0 {
1335 out.push((Message::assistant("done"), weight(rng)));
1336 }
1337 } else {
1338 out.push((Message::user(format!("plain {}", call_seq)), weight(rng)));
1339 out.push((Message::assistant("reply"), weight(rng)));
1340 }
1341 }
1342 out
1343 }
1344
1345 #[test]
1346 fn compression_and_rendering_preserve_tool_pairing_over_random_transcripts() {
1347 use crate::context::units::strict_tool_pairing_is_valid;
1348 let cfg = config();
1349 let eng = engine();
1350 let pipeline = CompressionPipeline::new(&cfg);
1351 let actions = [
1352 PressureAction::SnipCompact,
1353 PressureAction::MicroCompact,
1354 PressureAction::ContextCollapse,
1355 PressureAction::AutoCompact,
1356 ];
1357
1358 let mut call_seq = 0usize;
1359 for seed in 0..250u64 {
1360 let mut rng = Lcg(seed.wrapping_mul(0x9E3779B97F4A7C15).wrapping_add(1));
1361 let history = random_valid_history(&mut rng, &mut call_seq);
1362 let messages: Vec<Message> = history.iter().map(|(m, _)| m.clone()).collect();
1363 assert!(
1364 strict_tool_pairing_is_valid(&messages),
1365 "generator must emit valid transcripts (seed {seed})"
1366 );
1367
1368 for action in actions {
1369 for &target in &[0u32, 60, 200, 800] {
1370 let mut ctx = ContextPartitions::new(&cfg);
1371 for (message, tokens) in &history {
1372 ctx.history.push(message.clone(), *tokens);
1373 }
1374 pipeline.compress(&mut ctx, action, MAX, target, &eng);
1375 assert!(
1376 strict_tool_pairing_is_valid(&ctx.history.messages),
1377 "compression {action:?} target {target} broke pairing (seed {seed})"
1378 );
1379
1380 for &budget in &[10u32, 150, MAX] {
1382 let rc = super::super::renderer::render(&ctx, budget, &eng, 2);
1383 assert!(
1384 strict_tool_pairing_is_valid(&rc.turns),
1385 "render budget {budget} after {action:?} broke pairing (seed {seed})"
1386 );
1387 }
1388 }
1389 }
1390 }
1391 }
1392}