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