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 } = part
326 {
327 if preserved_refs
328 .iter()
329 .any(|preserved| preserved == call_id.as_str())
330 {
331 continue;
332 }
333 let original_output_tokens = engine.count(output);
334 let tool_name = find_tool_name(call_id, &messages_clone)
335 .unwrap_or_else(|| "unknown".to_string());
336
337 let new_output = if original_output_tokens > 2000 {
338 if let Some(json_excerpt) = extract_json_excerpt(output) {
339 format!(
340 "[tool result: {} | {} | {} tokens]\n{}",
341 call_id, tool_name, original_output_tokens, json_excerpt
342 )
343 } else {
344 let excerpt = excerpt_text(output, 30, 10, engine);
345 format!(
346 "[tool result: {} | {} | {} tokens]\n{}",
347 call_id, tool_name, original_output_tokens, excerpt
348 )
349 }
350 } else {
351 let excerpt = excerpt_text(output, 150, 50, engine);
352 format!(
353 "[tool result: {} | {} | {} tokens]\n{}",
354 call_id, tool_name, original_output_tokens, excerpt
355 )
356 };
357
358 *output = new_output;
359 }
360 }
361 let new_tokens = engine.count_message(msg);
362 msg.token_count = Some(new_tokens);
363 saved += original_tokens.saturating_sub(new_tokens);
364 }
365 }
366
367 partition.token_count = partition.token_count.saturating_sub(saved);
368
369 CompressResult {
371 tokens_saved: saved,
372 prefix_invalidated_at: indices.iter().min().copied(),
373 ..Default::default()
374 }
375 }
376}
377
378pub fn plan_drop_oldest(
385 messages: &[Message],
386 total_tokens: u32,
387 target_tokens: u32,
388 keep: usize,
389 engine: &ContextTokenEngine,
390) -> (usize, u32) {
391 let units = unit_boundaries(messages);
392 let limit = units.len().saturating_sub(keep);
393 let mut saved = 0u32;
394 let mut n = 0usize;
395 for unit in units.iter().take(limit) {
396 if total_tokens.saturating_sub(saved) <= target_tokens {
397 break;
398 }
399 saved += messages[unit.clone()]
400 .iter()
401 .map(|msg| msg.token_count.unwrap_or_else(|| engine.count_message(msg)))
402 .sum::<u32>();
403 n = unit.end;
404 }
405 (n, saved)
406}
407
408pub struct CollapseCompactor;
411
412impl Compressor for CollapseCompactor {
413 fn compress(
414 &self,
415 partitions: &mut ContextPartitions,
416 target_tokens: u32,
417 _max_tokens: u32,
418 preserve_k: usize,
419 engine: &ContextTokenEngine,
420 ) -> CompressResult {
421 let non_history_tokens = partitions
422 .total_tokens(engine)
423 .saturating_sub(partitions.history.token_count);
424 let history_target = target_tokens.saturating_sub(non_history_tokens);
425 let plan = plan_utility_archive(
426 &partitions.history.messages,
427 partitions.history.token_count,
428 history_target,
429 preserve_k,
430 engine,
431 &UtilitySelectionContext {
432 goal: &partitions.task_state.goal,
433 criteria: &partitions.task_state.criteria,
434 preserved_refs: &partitions.task_state.preserved_refs,
435 active_directives: &partitions.task_state.directives,
436 },
437 );
438 if plan.archived_ranges.is_empty() {
439 return CompressResult::default();
440 }
441 let prefix_invalidated_at = plan.archived_ranges.iter().map(|range| range.start).min();
442 let (archived, saved) = apply_utility_plan(&mut partitions.history, &plan);
443
444 CompressResult {
447 tokens_saved: saved,
448 archived,
449 prefix_invalidated_at,
450 ..Default::default()
451 }
452 }
453}
454
455pub struct AutoCompactor;
457
458impl Compressor for AutoCompactor {
459 fn compress(
460 &self,
461 partitions: &mut ContextPartitions,
462 target_tokens: u32,
463 _max_tokens: u32,
464 preserve_k: usize,
465 engine: &ContextTokenEngine,
466 ) -> CompressResult {
467 if partitions.history.messages.is_empty() {
468 return CompressResult::default();
469 }
470 let non_history_tokens = partitions
471 .total_tokens(engine)
472 .saturating_sub(partitions.history.token_count);
473 let history_target = target_tokens.saturating_sub(non_history_tokens);
474 let plan = plan_utility_archive(
475 &partitions.history.messages,
476 partitions.history.token_count,
477 history_target,
478 preserve_k,
479 engine,
480 &UtilitySelectionContext {
481 goal: &partitions.task_state.goal,
482 criteria: &partitions.task_state.criteria,
483 preserved_refs: &partitions.task_state.preserved_refs,
484 active_directives: &partitions.task_state.directives,
485 },
486 );
487 if plan.archived_ranges.is_empty() {
488 return CompressResult::default();
489 }
490 let prefix_invalidated_at = plan.archived_ranges.iter().map(|range| range.start).min();
491 let (archived, saved) = apply_utility_plan(&mut partitions.history, &plan);
492
493 CompressResult {
496 tokens_saved: saved,
497 archived,
498 prefix_invalidated_at,
499 ..Default::default()
500 }
501 }
502}
503
504fn apply_utility_plan(
505 partition: &mut super::partitions::Partition,
506 plan: &super::utility::UtilityArchivePlan,
507) -> (Vec<Message>, u32) {
508 let pairing_was_valid = strict_tool_pairing_is_valid(&partition.messages);
509 let archived_indices = plan
510 .archived_ranges
511 .iter()
512 .flat_map(|range| range.clone())
513 .collect::<std::collections::BTreeSet<_>>();
514 let mut archived = Vec::new();
515 let mut retained = Vec::new();
516 for (index, message) in std::mem::take(&mut partition.messages)
517 .into_iter()
518 .enumerate()
519 {
520 if archived_indices.contains(&index) {
521 archived.push(message);
522 } else {
523 retained.push(message);
524 }
525 }
526 partition.messages = retained;
527 partition.token_count = plan.retained_tokens;
528 debug_assert!(
529 !pairing_was_valid || strict_tool_pairing_is_valid(&partition.messages),
530 "utility selection split a valid tool transaction"
531 );
532 (archived, plan.archived_tokens)
533}
534
535pub struct CompressionPipeline {
537 stages: Vec<(PressureAction, Box<dyn Compressor>)>,
538 preserve_recent_turns: usize,
539}
540
541impl CompressionPipeline {
542 pub fn new(config: &ContextConfig) -> Self {
543 Self {
544 preserve_recent_turns: config.preserve_recent_turns,
545 stages: vec![
546 (
547 PressureAction::SnipCompact,
548 Box::new(SnipCompactor {
549 per_msg_ratio: config.snip_per_msg_ratio,
550 }),
551 ),
552 (PressureAction::MicroCompact, Box::new(MicroCompactor)),
553 (PressureAction::ContextCollapse, Box::new(CollapseCompactor)),
554 (PressureAction::AutoCompact, Box::new(AutoCompactor)),
555 ],
556 }
557 }
558
559 pub fn compress(
560 &self,
561 partitions: &mut ContextPartitions,
562 action: PressureAction,
563 max_tokens: u32,
564 target_tokens: u32,
565 engine: &ContextTokenEngine,
566 ) -> (u32, Option<String>, Vec<Message>, Option<usize>) {
567 if action == PressureAction::None {
568 return (0, None, vec![], None);
569 }
570
571 let mut total_saved = 0;
572 let mut all_archived = vec![];
573 let mut cache_at: Option<usize> = None;
576 let summarizer = super::summarizer::RuleSummarizer;
577
578 for (stage_action, compressor) in &self.stages {
579 if *stage_action <= action {
580 if partitions.total_tokens(engine) <= target_tokens {
581 break;
582 }
583 let res = compressor.compress(
584 partitions,
585 target_tokens,
586 max_tokens,
587 self.preserve_recent_turns,
588 engine,
589 );
590 total_saved += res.tokens_saved;
591 cache_at = [cache_at, res.prefix_invalidated_at]
592 .into_iter()
593 .flatten()
594 .min();
595 all_archived.extend(res.archived);
596 }
597 }
598
599 let summary_budget = if target_tokens == 0 {
611 max_tokens
612 } else {
613 target_tokens
614 };
615 let summary = if all_archived.is_empty() {
616 None
617 } else {
618 let s = summarizer.summarize(&all_archived, action, summary_budget);
619 partitions
620 .task_state
621 .log_compression(action.label(), s.clone());
622 Some(s)
623 };
624
625 (total_saved, summary, all_archived, cache_at)
626 }
627}
628
629#[cfg(test)]
630mod tests {
631 use super::*;
632 use crate::context::config::ContextConfig;
633 use crate::context::partitions::ContextPartitions;
634 use crate::context::token_engine::ContextTokenEngine;
635 use crate::types::message::Message;
636
637 fn engine() -> ContextTokenEngine {
638 ContextTokenEngine::char_approx()
639 }
640 fn config() -> ContextConfig {
641 ContextConfig::default()
642 }
643 const MAX: u32 = 1_000;
644
645 #[test]
646 fn snip_compactor_truncates_oversized_messages() {
647 let cfg = ContextConfig {
648 snip_per_msg_ratio: 0.10,
649 ..Default::default()
650 };
651 let compactor = SnipCompactor {
652 per_msg_ratio: cfg.snip_per_msg_ratio,
653 };
654 let mut ctx = ContextPartitions::new(&cfg);
655 ctx.history.push(Message::user("a".repeat(800)), 200);
656 let result = compactor.compress(&mut ctx, 0, MAX, 0, &engine());
658 assert!(result.tokens_saved > 0);
659 if let Content::Text(ref t) = ctx.history.messages[0].content {
660 assert!(t.contains("… [… 100 tokens omitted …] …"), "got: {t}");
661 }
662 }
663
664 #[test]
665 fn snip_compactor_leaves_small_messages_untouched() {
666 let cfg = ContextConfig {
667 snip_per_msg_ratio: 0.10,
668 ..Default::default()
669 };
670 let compactor = SnipCompactor {
671 per_msg_ratio: cfg.snip_per_msg_ratio,
672 };
673 let mut ctx = ContextPartitions::new(&cfg);
674 ctx.history.push(Message::user("short"), 5);
675 let result = compactor.compress(&mut ctx, 0, MAX, 2, &engine());
676 assert_eq!(result.tokens_saved, 0);
677 }
678
679 #[test]
680 fn micro_compactor_replaces_tool_results_with_measured_placeholder() {
681 use crate::types::message::{ContentPart, Role};
682 use compact_str::CompactString;
683
684 let compactor = MicroCompactor;
685 let mut ctx = ContextPartitions::new(&config());
686 let parts = vec![ContentPart::ToolResult {
687 call_id: CompactString::new("c1"),
688 output: "a".repeat(1200),
689 is_error: false,
690 }];
691 let msg = Message {
692 role: Role::Tool,
693 content: Content::Parts(parts),
694 tool_calls: vec![],
695 token_count: Some(300),
696 };
697 ctx.history.messages.push(msg);
698 ctx.history.token_count = 300;
699
700 let result = compactor.compress(&mut ctx, 0, MAX, 0, &engine());
702 assert!(result.tokens_saved > 0);
703 let Content::Parts(parts) = &ctx.history.messages[0].content else {
704 panic!("tool-result envelope must survive compaction");
705 };
706 let text = parts
707 .iter()
708 .find_map(|part| match part {
709 ContentPart::ToolResult {
710 call_id, output, ..
711 } if call_id.as_str() == "c1" => Some(output.as_str()),
712 _ => None,
713 })
714 .expect("correlated tool result remains present");
715 assert!(
716 text.contains("[tool result: c1 | unknown | 300 tokens]"),
717 "got: {text}"
718 );
719 }
720
721 #[test]
722 fn micro_compactor_recounts_the_entire_mixed_parts_envelope() {
723 use crate::types::message::{ContentPart, Role};
724
725 let compactor = MicroCompactor;
726 let mut ctx = ContextPartitions::new(&config());
727 let msg = Message {
728 role: Role::Tool,
729 content: Content::Parts(vec![
730 ContentPart::Text {
731 text: "metadata that must remain budgeted".into(),
732 },
733 ContentPart::ToolResult {
734 call_id: "c1".into(),
735 output: "a".repeat(1200),
736 is_error: false,
737 },
738 ContentPart::ToolResult {
739 call_id: "c2".into(),
740 output: "b".repeat(1000),
741 is_error: false,
742 },
743 ]),
744 tool_calls: vec![],
745 token_count: None,
746 };
747 let original = engine().count_message(&msg);
748 ctx.history.push(msg, original);
749
750 compactor.compress(&mut ctx, 0, MAX, 0, &engine());
751
752 let message = &ctx.history.messages[0];
753 let recounted = engine().count_message(message);
754 assert_eq!(message.token_count, Some(recounted));
755 assert_eq!(ctx.history.token_count, recounted);
756 let Content::Parts(parts) = &message.content else {
757 panic!("parts preserved")
758 };
759 assert_eq!(
760 parts
761 .iter()
762 .filter(|part| matches!(part, ContentPart::ToolResult { output, .. } if output.starts_with("[tool result:")))
763 .count(),
764 2,
765 "every eligible tool result is excerpted"
766 );
767 }
768
769 #[test]
770 fn collapse_compactor_drops_oldest_to_reach_target() {
771 let compactor = CollapseCompactor;
772 let mut ctx = ContextPartitions::new(&config());
773 for _ in 0..8 {
774 ctx.history.push(Message::user("msg"), 50);
775 }
776 let result = compactor.compress(&mut ctx, 250, MAX, 2, &engine());
777 assert!(result.tokens_saved > 0);
778 assert!(ctx.history.messages.len() < 8);
779 assert!(
782 !result.archived.is_empty(),
783 "drained messages are returned to the pipeline"
784 );
785 assert!(
786 result.summary.is_none(),
787 "compactor no longer self-summarizes"
788 );
789 assert!(
790 ctx.task_state.compression_log.is_empty(),
791 "compactor no longer logs"
792 );
793 }
794
795 #[test]
796 fn collapse_utility_selection_changes_the_archived_units_not_only_the_score() {
797 let compactor = CollapseCompactor;
798 let mut ctx = ContextPartitions::new(&config());
799 ctx.task_state.goal = "ship ORCHID release".into();
800 for (user, assistant) in [
801 (
802 "ORCHID release criterion",
803 "DECISION: retry failure; artifact /work/orchid.json",
804 ),
805 ("routine chatter one", "acknowledged"),
806 ("routine chatter two", "acknowledged"),
807 ("latest request", "working on it"),
808 ] {
809 ctx.history.push(Message::user(user), 40);
810 ctx.history.push(Message::assistant(assistant), 40);
811 }
812
813 let result = compactor.compress(&mut ctx, 160, MAX, 1, &engine());
814 let retained = ctx
815 .history
816 .messages
817 .iter()
818 .filter_map(|message| message.content.as_text())
819 .collect::<Vec<_>>()
820 .join("\n");
821 let archived = result
822 .archived
823 .iter()
824 .filter_map(|message| message.content.as_text())
825 .collect::<Vec<_>>()
826 .join("\n");
827 assert!(retained.contains("ORCHID"));
828 assert!(retained.contains("/work/orchid.json"));
829 assert!(!retained.contains("routine chatter"));
830 assert!(archived.contains("routine chatter one"));
831 assert!(archived.contains("routine chatter two"));
832 assert_eq!(ctx.history.token_count, 160);
833 }
834
835 #[test]
836 fn utility_selector_deducts_fixed_context_before_budgeting_history() {
837 let compactor = CollapseCompactor;
838 let mut ctx = ContextPartitions::new(&config());
839 ctx.system.push(Message::system("fixed"), 600);
840 for index in 0..4 {
841 ctx.history
842 .push(Message::user(format!("unit {index}")), 100);
843 }
844
845 let result = compactor.compress(&mut ctx, 700, MAX, 1, &engine());
846 assert_eq!(result.tokens_saved, 300);
847 assert_eq!(ctx.history.token_count, 100);
848 assert!(ctx.total_tokens(&engine()) <= 700);
849 }
850
851 #[test]
852 fn rule_summarizer_formats_correctly() {
853 use crate::context::summarizer::RuleSummarizer;
854 use crate::types::message::{Content, Message, Role};
855 let summarizer = RuleSummarizer;
856 let mut messages = vec![];
857 messages.push(Message {
858 role: Role::User,
859 content: Content::Text("hello".to_string()),
860 tool_calls: vec![],
861 token_count: Some(5),
862 });
863 messages.push(Message {
864 role: Role::Assistant,
865 content: Content::Text("world".to_string()),
866 tool_calls: vec![],
867 token_count: Some(6),
868 });
869 let summary = summarizer.summarize(&messages, PressureAction::SnipCompact, 100);
870 assert!(summary.contains("[Compressed: snip_compact]"));
871 assert!(summary.contains("archived_messages: 2; archived_tokens: 11"));
872 assert!(summary.contains("constraints:"));
873 }
874
875 #[test]
876 fn micro_compactor_preserves_refs_in_preserved_refs() {
877 use crate::types::message::{ContentPart, Role};
878 use compact_str::CompactString;
879
880 let compactor = MicroCompactor;
881 let mut ctx = ContextPartitions::new(&config());
882 ctx.task_state.preserved_refs = vec!["keep_me".to_string()];
883
884 let parts = vec![ContentPart::ToolResult {
885 call_id: CompactString::new("keep_me"),
886 output: "a".repeat(1200),
887 is_error: false,
888 }];
889 let msg = Message {
890 role: Role::Tool,
891 content: Content::Parts(parts),
892 tool_calls: vec![],
893 token_count: Some(300),
894 };
895 ctx.history.messages.push(msg);
896 ctx.history.token_count = 300;
897
898 let result = compactor.compress(&mut ctx, 0, MAX, 2, &engine());
899 assert_eq!(result.tokens_saved, 0);
901 let text_opt = ctx.history.messages[0].content.as_text();
902 assert!(
903 text_opt.is_none(),
904 "should not be replaced to text placeholder"
905 );
906 }
907
908 #[test]
909 fn auto_compactor_merges_all_except_last_two_turns() {
910 let compactor = AutoCompactor;
911 let mut ctx = ContextPartitions::new(&config());
912 for i in 0..10 {
913 ctx.history.push(Message::user(format!("msg {i}")), 10);
914 }
915 let result = compactor.compress(&mut ctx, 0, MAX, 2, &engine());
916 assert!(result.tokens_saved > 0);
917 assert_eq!(ctx.history.messages.len(), 2); assert!(
921 !result.archived.is_empty(),
922 "drained messages returned to the pipeline"
923 );
924 assert!(
925 result.summary.is_none(),
926 "compactor no longer self-summarizes"
927 );
928 assert!(
929 ctx.task_state.compression_log.is_empty(),
930 "compactor no longer logs"
931 );
932 }
933
934 #[test]
935 fn plan_drop_oldest_respects_target_and_preserve_floor() {
936 let msgs: Vec<Message> = (0..8)
939 .map(|i| {
940 let mut m = Message::user(format!("m{i}"));
941 m.token_count = Some(50);
942 m
943 })
944 .collect();
945 assert_eq!(plan_drop_oldest(&msgs, 400, 250, 2, &engine()), (3, 150));
947 assert_eq!(plan_drop_oldest(&msgs, 400, 0, 2, &engine()), (6, 300));
949 assert_eq!(plan_drop_oldest(&msgs, 400, 500, 2, &engine()), (0, 0));
951 }
952
953 #[test]
954 fn collapse_never_splits_a_tool_transaction() {
955 let mut call = Message::assistant("calling");
956 call.tool_calls.push(crate::types::message::ToolCall {
957 id: "call-1".into(),
958 name: "read".into(),
959 arguments: serde_json::json!({}),
960 });
961 let messages = vec![
962 Message::user("question"),
963 call,
964 Message::tool(vec![ContentPart::ToolResult {
965 call_id: "call-1".into(),
966 output: "ok".into(),
967 is_error: false,
968 }]),
969 Message::assistant("answer"),
970 Message::user("next"),
971 Message::assistant("done"),
972 ]
973 .into_iter()
974 .map(|mut message| {
975 message.token_count = Some(10);
976 message
977 })
978 .collect::<Vec<_>>();
979
980 assert_eq!(plan_drop_oldest(&messages, 60, 30, 1, &engine()), (4, 40));
981 }
982
983 #[test]
984 fn auto_compactor_preserves_the_latest_complete_tool_unit() {
985 let compactor = AutoCompactor;
986 let mut ctx = ContextPartitions::new(&config());
987 ctx.history.push(Message::user("old"), 10);
988 ctx.history.push(Message::assistant("old answer"), 10);
989 let mut call = Message::assistant("calling");
990 call.tool_calls.push(crate::types::message::ToolCall {
991 id: "call-1".into(),
992 name: "read".into(),
993 arguments: serde_json::json!({}),
994 });
995 ctx.history.push(Message::user("question"), 10);
996 ctx.history.push(call, 10);
997 ctx.history.push(
998 Message::tool(vec![ContentPart::ToolResult {
999 call_id: "call-1".into(),
1000 output: "ok".into(),
1001 is_error: false,
1002 }]),
1003 10,
1004 );
1005 ctx.history.push(Message::assistant("answer"), 10);
1006
1007 compactor.compress(&mut ctx, 0, MAX, 1, &engine());
1008
1009 assert_eq!(ctx.history.messages.len(), 4);
1010 assert_eq!(ctx.history.messages[0].content.as_text(), Some("question"));
1011 }
1012
1013 #[test]
1014 fn prefix_keep_yields_without_drop_fallback() {
1015 assert_eq!(prefix_keep_for(6, 2), 2, "len 6 >= 6 → protect oldest 2");
1018 assert_eq!(
1019 prefix_keep_for(5, 2),
1020 0,
1021 "len 5 < 6 → would leave an untouchable message"
1022 );
1023 assert_eq!(prefix_keep_for(3, 2), 0);
1024 assert_eq!(prefix_keep_for(0, 2), 0);
1025 }
1026
1027 #[test]
1028 fn pipeline_reports_accurate_prefix_invalidation() {
1029 let mut cfg = config();
1033 cfg.preserve_recent_turns = 1;
1034 let mut ctx = baseline_partitions();
1035 let (_s, _u, _a, cache_at) = CompressionPipeline::new(&cfg).compress(
1036 &mut ctx,
1037 PressureAction::SnipCompact,
1038 MAX,
1039 500,
1040 &engine(),
1041 );
1042 assert_eq!(
1043 cache_at,
1044 Some(1),
1045 "one protected unit leaves msg 1 as the earliest rewrite"
1046 );
1047
1048 let mut ctx2 = baseline_partitions();
1049 let (_s2, _u2, _a2, cache_at2) = CompressionPipeline::new(&cfg).compress(
1050 &mut ctx2,
1051 PressureAction::AutoCompact,
1052 MAX,
1053 500,
1054 &engine(),
1055 );
1056 assert_eq!(
1057 cache_at2,
1058 Some(0),
1059 "dropping the oldest breaks the cache prefix at 0"
1060 );
1061 }
1062
1063 use crate::types::message::Role;
1072 use compact_str::CompactString;
1073
1074 fn baseline_partitions() -> ContextPartitions {
1077 let cfg = config();
1078 let mut ctx = ContextPartitions::new(&cfg);
1079 ctx.history.push(Message::user("u0 ".repeat(120)), 300);
1081 ctx.history.push(Message::assistant("a0 ".repeat(120)), 300);
1082 ctx.history.messages.push(Message {
1084 role: Role::Tool,
1085 content: Content::Parts(vec![ContentPart::ToolResult {
1086 call_id: CompactString::new("call_1"),
1087 output: serde_json::json!({"rows": 42, "ok": true, "name": "alpha"}).to_string()
1088 + &"-pad".repeat(400),
1089 is_error: false,
1090 }]),
1091 tool_calls: vec![],
1092 token_count: Some(400),
1093 });
1094 ctx.history.token_count += 400;
1095 ctx.history.push(Message::user("u1 ".repeat(120)), 300);
1096 ctx.history.push(Message::assistant("a1 ".repeat(120)), 300);
1097 ctx.history.messages.push(Message {
1098 role: Role::Tool,
1099 content: Content::Parts(vec![ContentPart::ToolResult {
1100 call_id: CompactString::new("call_2"),
1101 output: "y".repeat(1600),
1102 is_error: false,
1103 }]),
1104 tool_calls: vec![],
1105 token_count: Some(400),
1106 });
1107 ctx.history.token_count += 400;
1108 ctx
1109 }
1110
1111 fn run_baseline(action: PressureAction) -> (u32, u32, Option<String>, usize, usize, u32) {
1114 let mut ctx = baseline_partitions();
1115 let before = ctx.total_tokens(&engine());
1116 let mut cfg = config();
1117 cfg.preserve_recent_turns = 1;
1118 let (saved, summary, archived, _cache_at) =
1119 CompressionPipeline::new(&cfg).compress(&mut ctx, action, MAX, 500, &engine());
1120 let archived_len = archived.len();
1121 let msgs_after = ctx.history.messages.len();
1122 let total_after = ctx.total_tokens(&engine());
1123 (
1124 before,
1125 saved,
1126 summary,
1127 archived_len,
1128 msgs_after,
1129 total_after,
1130 )
1131 }
1132
1133 #[test]
1134 fn baseline_snip_only_caps_text_no_archival() {
1135 let (before, saved, summary, archived, msgs, total) =
1139 run_baseline(PressureAction::SnipCompact);
1140 assert_eq!(before, 2000);
1141 assert_eq!(saved, 750, "3 non-prefix oversized messages × 250");
1142 assert_eq!(archived, 0);
1143 assert!(summary.is_none());
1144 assert_eq!(msgs, 6, "snip mutates in place, drops no messages");
1145 assert_eq!(total, 1250);
1146 }
1147
1148 #[test]
1149 fn baseline_micro_excerpts_tool_results() {
1150 let (before, saved, summary, archived, msgs, total) =
1152 run_baseline(PressureAction::MicroCompact);
1153 assert_eq!(before, 2000);
1154 assert_eq!(saved, 1112, "snip(750) + tool-result excerpts(362)");
1155 assert_eq!(archived, 0);
1156 assert!(summary.is_none());
1157 assert_eq!(msgs, 6);
1158 assert_eq!(total, 888);
1159 }
1160
1161 #[test]
1162 fn baseline_collapse_drops_oldest_and_summarizes() {
1163 let (before, saved, summary, archived, msgs, total) =
1166 run_baseline(PressureAction::ContextCollapse);
1167 assert_eq!(before, 2000);
1168 assert_eq!(saved, 1681);
1169 assert_eq!(archived, 3, "drops the complete oldest unit");
1170 assert_eq!(msgs, 3, "one complete recent unit remains");
1171 assert_eq!(total, 319);
1172 let summary = summary.expect("collapse summarizes archived messages");
1173 assert!(
1174 summary.contains("[Compressed: context_collapse]"),
1175 "summary routes the collapse action: {summary}"
1176 );
1177 }
1178
1179 #[test]
1180 fn baseline_auto_attributes_summary_to_auto_compact() {
1181 let (before, saved, summary, archived, msgs, total) =
1187 run_baseline(PressureAction::AutoCompact);
1188 assert_eq!(before, 2000);
1189 assert_eq!(saved, 1681);
1190 assert_eq!(archived, 3);
1191 assert_eq!(msgs, 3);
1192 assert_eq!(total, 319);
1193 let summary = summary.expect("auto-compact summarizes the archived messages");
1194 assert!(
1195 summary.contains("[Compressed: auto_compact]"),
1196 "got: {summary}"
1197 );
1198 }
1199
1200 #[test]
1201 fn baseline_saved_is_monotonic_in_action_level() {
1202 let snip = run_baseline(PressureAction::SnipCompact).1;
1204 let micro = run_baseline(PressureAction::MicroCompact).1;
1205 let collapse = run_baseline(PressureAction::ContextCollapse).1;
1206 let auto = run_baseline(PressureAction::AutoCompact).1;
1207 assert!(snip <= micro, "{snip} <= {micro}");
1208 assert!(micro <= collapse, "{micro} <= {collapse}");
1209 assert!(collapse <= auto, "{collapse} <= {auto}");
1210 }
1211
1212 #[test]
1213 fn pipeline_stops_cascade_when_target_reached() {
1214 let cfg = ContextConfig {
1215 snip_per_msg_ratio: 0.25,
1216 preserve_recent_turns: 0,
1219 ..Default::default()
1220 };
1221 let pipeline = CompressionPipeline::new(&cfg);
1222 let mut ctx = ContextPartitions::new(&cfg);
1223 ctx.history.push(Message::user("a".repeat(3600)), 900);
1224
1225 let (saved, summary, archived, _cache_at) =
1226 pipeline.compress(&mut ctx, PressureAction::AutoCompact, 1_000, 500, &engine());
1227
1228 assert!(saved > 0);
1229 assert!(
1230 summary.is_none(),
1231 "auto compactor should not run after snip reaches target"
1232 );
1233 assert!(
1234 archived.is_empty(),
1235 "heavier archival stages should not run"
1236 );
1237 assert_eq!(ctx.history.messages.len(), 1);
1238 assert!(ctx.total_tokens(&engine()) <= 500);
1239 }
1240
1241 struct Lcg(u64);
1252 impl Lcg {
1253 fn next_u64(&mut self) -> u64 {
1254 self.0 = self
1255 .0
1256 .wrapping_mul(6364136223846793005)
1257 .wrapping_add(1442695040888963407);
1258 self.0 ^ (self.0 >> 33)
1259 }
1260 fn below(&mut self, n: u64) -> u64 {
1261 self.next_u64() % n.max(1)
1262 }
1263 }
1264
1265 fn random_valid_history(rng: &mut Lcg, call_seq: &mut usize) -> Vec<(Message, u32)> {
1269 use crate::types::message::{ContentPart, ToolCall};
1270 let mut out: Vec<(Message, u32)> = Vec::new();
1271 let transactions = 3 + rng.below(10) as usize;
1272 let weight = |rng: &mut Lcg| 10 + rng.below(300) as u32;
1273 for _ in 0..transactions {
1274 if rng.below(10) < 6 {
1275 if rng.below(2) == 0 {
1277 out.push((Message::user(format!("ask {}", call_seq)), weight(rng)));
1278 }
1279 let n_calls = 1 + rng.below(3) as usize;
1280 let ids: Vec<String> = (0..n_calls)
1281 .map(|_| {
1282 *call_seq += 1;
1283 format!("call-{call_seq}")
1284 })
1285 .collect();
1286 let mut assistant = Message::assistant("working");
1287 for id in &ids {
1288 assistant.tool_calls.push(ToolCall {
1289 id: id.clone().into(),
1290 name: "read".into(),
1291 arguments: serde_json::json!({}),
1292 });
1293 }
1294 out.push((assistant, weight(rng)));
1295 if rng.below(2) == 0 {
1297 for id in &ids {
1298 out.push((
1299 Message::tool(vec![ContentPart::ToolResult {
1300 call_id: id.clone().into(),
1301 output: "ok ".repeat(20),
1302 is_error: false,
1303 }]),
1304 weight(rng),
1305 ));
1306 }
1307 } else {
1308 let parts = ids
1309 .iter()
1310 .map(|id| ContentPart::ToolResult {
1311 call_id: id.clone().into(),
1312 output: "ok ".repeat(20),
1313 is_error: false,
1314 })
1315 .collect();
1316 out.push((Message::tool(parts), weight(rng)));
1317 }
1318 if rng.below(2) == 0 {
1319 out.push((Message::assistant("done"), weight(rng)));
1320 }
1321 } else {
1322 out.push((Message::user(format!("plain {}", call_seq)), weight(rng)));
1323 out.push((Message::assistant("reply"), weight(rng)));
1324 }
1325 }
1326 out
1327 }
1328
1329 #[test]
1330 fn compression_and_rendering_preserve_tool_pairing_over_random_transcripts() {
1331 use crate::context::units::strict_tool_pairing_is_valid;
1332 let cfg = config();
1333 let eng = engine();
1334 let pipeline = CompressionPipeline::new(&cfg);
1335 let actions = [
1336 PressureAction::SnipCompact,
1337 PressureAction::MicroCompact,
1338 PressureAction::ContextCollapse,
1339 PressureAction::AutoCompact,
1340 ];
1341
1342 let mut call_seq = 0usize;
1343 for seed in 0..250u64 {
1344 let mut rng = Lcg(seed.wrapping_mul(0x9E3779B97F4A7C15).wrapping_add(1));
1345 let history = random_valid_history(&mut rng, &mut call_seq);
1346 let messages: Vec<Message> = history.iter().map(|(m, _)| m.clone()).collect();
1347 assert!(
1348 strict_tool_pairing_is_valid(&messages),
1349 "generator must emit valid transcripts (seed {seed})"
1350 );
1351
1352 for action in actions {
1353 for &target in &[0u32, 60, 200, 800] {
1354 let mut ctx = ContextPartitions::new(&cfg);
1355 for (message, tokens) in &history {
1356 ctx.history.push(message.clone(), *tokens);
1357 }
1358 pipeline.compress(&mut ctx, action, MAX, target, &eng);
1359 assert!(
1360 strict_tool_pairing_is_valid(&ctx.history.messages),
1361 "compression {action:?} target {target} broke pairing (seed {seed})"
1362 );
1363
1364 for &budget in &[10u32, 150, MAX] {
1366 let rc = super::super::renderer::render(&ctx, budget, &eng, 2);
1367 assert!(
1368 strict_tool_pairing_is_valid(&rc.turns),
1369 "render budget {budget} after {action:?} broke pairing (seed {seed})"
1370 );
1371 }
1372 }
1373 }
1374 }
1375 }
1376}