1use super::{
14 Capability, CapabilityLocalization, CapabilityStatus, ModelViewContext, ModelViewProvider,
15};
16use crate::events::TokenUsage;
17use crate::message::{ContentPart, Message, MessageRole};
18use crate::message_filter::MessageFilterProvider;
19use serde::{Deserialize, Serialize};
20use std::collections::{HashMap, HashSet};
21use std::sync::Arc;
22
23pub const COMPACTION_CAPABILITY_ID: &str = "compaction";
25const MAX_RELATED_RECENT_READ_RESULTS: usize = 4;
26
27#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
29#[serde(rename_all = "snake_case")]
30pub enum CompactionStrategy {
31 #[default]
33 Auto,
34 Native,
36 ObservationMasking,
38 Summarization,
40}
41
42impl std::fmt::Display for CompactionStrategy {
43 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
44 match self {
45 Self::Auto => write!(f, "auto"),
46 Self::Native => write!(f, "native"),
47 Self::ObservationMasking => write!(f, "observation_masking"),
48 Self::Summarization => write!(f, "summarization"),
49 }
50 }
51}
52
53#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
55#[serde(rename_all = "snake_case")]
56pub enum MaskingSummaryFormat {
57 #[default]
59 OneLine,
60 HeadTail,
62}
63
64#[derive(Debug, Clone, Serialize, Deserialize)]
66pub struct ObservationMaskingConfig {
67 #[serde(default = "default_keep_recent_tool_outputs")]
69 pub keep_recent_tool_outputs: usize,
70
71 #[serde(default)]
73 pub summary_format: MaskingSummaryFormat,
74}
75
76impl Default for ObservationMaskingConfig {
77 fn default() -> Self {
78 Self {
79 keep_recent_tool_outputs: default_keep_recent_tool_outputs(),
80 summary_format: MaskingSummaryFormat::default(),
81 }
82 }
83}
84
85fn default_keep_recent_tool_outputs() -> usize {
86 2
90}
91
92#[derive(Debug, Clone, Serialize, Deserialize)]
98pub struct CostControlConfig {
99 #[serde(default = "default_cost_control_enabled")]
101 pub enabled: bool,
102
103 #[serde(default = "default_cost_control_keep_recent_tool_results")]
105 pub keep_recent_tool_results: usize,
106
107 #[serde(default = "default_cost_control_mask_after_tool_results")]
109 pub mask_after_tool_results: usize,
110
111 #[serde(default = "default_cost_control_max_live_tool_result_bytes")]
113 pub max_live_tool_result_bytes: usize,
114
115 #[serde(default = "default_cost_control_max_uncached_input_tokens")]
117 pub max_uncached_input_tokens: u32,
118
119 #[serde(default = "default_cost_control_min_cache_read_ratio")]
121 pub min_cache_read_ratio: f32,
122}
123
124impl Default for CostControlConfig {
125 fn default() -> Self {
126 Self {
127 enabled: default_cost_control_enabled(),
128 keep_recent_tool_results: default_cost_control_keep_recent_tool_results(),
129 mask_after_tool_results: default_cost_control_mask_after_tool_results(),
130 max_live_tool_result_bytes: default_cost_control_max_live_tool_result_bytes(),
131 max_uncached_input_tokens: default_cost_control_max_uncached_input_tokens(),
132 min_cache_read_ratio: default_cost_control_min_cache_read_ratio(),
133 }
134 }
135}
136
137fn default_cost_control_enabled() -> bool {
138 true
139}
140
141fn default_cost_control_keep_recent_tool_results() -> usize {
142 2
143}
144
145fn default_cost_control_mask_after_tool_results() -> usize {
146 4
147}
148
149fn default_cost_control_max_live_tool_result_bytes() -> usize {
150 24 * 1024
151}
152
153fn default_cost_control_max_uncached_input_tokens() -> u32 {
154 100_000
155}
156
157fn default_cost_control_min_cache_read_ratio() -> f32 {
158 0.35
159}
160
161#[derive(Debug, Clone, Serialize, Deserialize)]
163pub struct SummarizationConfig {
164 #[serde(default)]
166 pub model: Option<String>,
167
168 #[serde(default = "default_preserve")]
170 pub preserve: Vec<String>,
171
172 #[serde(default)]
174 pub instructions: Option<String>,
175}
176
177impl Default for SummarizationConfig {
178 fn default() -> Self {
179 Self {
180 model: None,
181 preserve: default_preserve(),
182 instructions: None,
183 }
184 }
185}
186
187fn default_preserve() -> Vec<String> {
188 vec![
189 "decisions".to_string(),
190 "files_modified".to_string(),
191 "errors".to_string(),
192 "current_plan".to_string(),
193 "skill_instructions".to_string(),
194 ]
195}
196
197#[derive(Debug, Clone, Serialize, Deserialize)]
204pub struct CompactionConfig {
205 #[serde(default)]
207 pub strategy: CompactionStrategy,
208
209 #[serde(default = "default_proactive")]
211 pub proactive: bool,
212
213 #[serde(default = "default_budget_percent")]
215 pub budget_percent: f32,
216
217 #[serde(default)]
219 pub observation_masking: ObservationMaskingConfig,
220
221 #[serde(default)]
223 pub summarization: SummarizationConfig,
224
225 #[serde(default)]
227 pub memory_tiers: HierarchicalMemoryConfig,
228
229 #[serde(default)]
231 pub cost_control: CostControlConfig,
232}
233
234impl Default for CompactionConfig {
235 fn default() -> Self {
236 Self {
237 strategy: CompactionStrategy::default(),
238 proactive: default_proactive(),
239 budget_percent: default_budget_percent(),
240 observation_masking: ObservationMaskingConfig::default(),
241 summarization: SummarizationConfig::default(),
242 memory_tiers: HierarchicalMemoryConfig::default(),
243 cost_control: CostControlConfig::default(),
244 }
245 }
246}
247
248fn default_proactive() -> bool {
249 true
250}
251
252fn default_budget_percent() -> f32 {
253 0.85
254}
255
256impl CompactionConfig {
257 pub fn from_json(value: &serde_json::Value) -> Self {
259 serde_json::from_value(value.clone()).unwrap_or_default()
260 }
261}
262
263pub struct CompactionCapability;
265
266impl Capability for CompactionCapability {
267 fn id(&self) -> &str {
268 COMPACTION_CAPABILITY_ID
269 }
270
271 fn name(&self) -> &str {
272 "Compaction"
273 }
274
275 fn description(&self) -> &str {
276 r#"Configurable context compaction when conversations exceed LLM context windows.
277
278Choose between native provider compaction (e.g., OpenAI /responses/compact), observation masking (strip old tool outputs), or LLM summarization. The `auto` strategy cascades through all available options."#
279 }
280
281 fn status(&self) -> CapabilityStatus {
282 CapabilityStatus::Available
283 }
284
285 fn icon(&self) -> Option<&str> {
286 Some("shrink")
287 }
288
289 fn category(&self) -> Option<&str> {
290 Some("Optimization")
291 }
292
293 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
294 Some(Arc::new(CompactionFilterProvider))
295 }
296
297 fn model_view_provider(&self) -> Option<Arc<dyn ModelViewProvider>> {
298 Some(Arc::new(CompactionModelViewProvider))
299 }
300
301 fn config_schema(&self) -> Option<serde_json::Value> {
308 Some(serde_json::json!({
309 "type": "object",
310 "properties": {
311 "strategy": {
312 "type": "string",
313 "title": "Strategy",
314 "description": "Compaction strategy used when the conversation approaches the context window.",
315 "oneOf": [
316 { "const": "auto", "title": "Automatic" },
317 { "const": "native", "title": "Provider-native" },
318 { "const": "observation_masking", "title": "Observation masking" },
319 { "const": "summarization", "title": "LLM summarization" }
320 ],
321 "default": "auto"
322 },
323 "proactive": {
324 "type": "boolean",
325 "title": "Proactive compaction",
326 "description": "Compact at the budget threshold instead of waiting for a request-too-large error.",
327 "default": true
328 },
329 "budget_percent": {
330 "type": "number",
331 "title": "Context budget threshold",
332 "description": "Fraction of the model context window at which proactive compaction triggers.",
333 "minimum": 0.1,
334 "maximum": 1.0,
335 "default": 0.85
339 }
340 }
341 }))
342 }
343
344 fn validate_config(&self, config: &serde_json::Value) -> Result<(), String> {
345 if config.is_null() {
346 return Ok(());
347 }
348 let typed: CompactionConfig = serde_json::from_value(config.clone())
349 .map_err(|e| format!("invalid compaction config: {e}"))?;
350 if !(0.1..=1.0).contains(&typed.budget_percent) {
351 return Err(format!(
352 "budget_percent must be between 0.1 and 1.0, got {}",
353 typed.budget_percent
354 ));
355 }
356 Ok(())
357 }
358
359 fn localizations(&self) -> Vec<CapabilityLocalization> {
360 vec![
361 CapabilityLocalization {
362 locale: "en",
363 name: None,
364 description: None,
365 config_description: Some(
366 "Controls the compaction strategy, proactive triggering, and the \
367 context-budget threshold.",
368 ),
369 config_overlay: None,
370 },
371 CapabilityLocalization {
372 locale: "uk",
373 name: Some("Ущільнення контексту"),
374 description: Some(
375 "Налаштовуване ущільнення контексту, коли розмова перевищує контекстне \
376 вікно LLM. Доступні стратегії: нативне ущільнення провайдера, маскування \
377 результатів інструментів і підсумовування через LLM; стратегія auto \
378 перебирає всі доступні варіанти.",
379 ),
380 config_description: Some(
381 "Визначає стратегію ущільнення контексту, проактивний запуск і поріг \
382 бюджету контексту.",
383 ),
384 config_overlay: Some(serde_json::json!({
385 "properties": {
386 "strategy": {
387 "title": "Стратегія",
388 "description": "Стратегія ущільнення, коли розмова наближається до межі контекстного вікна.",
389 "enum_labels": {
390 "auto": "Автоматично",
391 "native": "Нативна (провайдер)",
392 "observation_masking": "Маскування результатів інструментів",
393 "summarization": "Підсумовування через LLM"
394 }
395 },
396 "proactive": {
397 "title": "Проактивне ущільнення",
398 "description": "Ущільнювати контекст при досягненні порогу бюджету, а не лише після помилки про завеликий запит."
399 },
400 "budget_percent": {
401 "title": "Поріг бюджету контексту",
402 "description": "Частка контекстного вікна моделі, після якої запускається проактивне ущільнення."
403 }
404 }
405 })),
406 },
407 ]
408 }
409}
410
411struct CompactionModelViewProvider;
412
413impl ModelViewProvider for CompactionModelViewProvider {
414 fn apply_model_view(
415 &self,
416 messages: Vec<Message>,
417 config: &serde_json::Value,
418 context: &ModelViewContext<'_>,
419 ) -> Vec<Message> {
420 let config = CompactionConfig::from_json(config);
421 let masking = build_model_view_messages_owned(messages, &config, context.prior_usage);
422 if masking.masked_count > 0 {
423 tracing::info!(
424 session_id = %context.session_id,
425 masked_count = masking.masked_count,
426 tool_result_bytes_before = masking.tool_result_bytes_before,
427 tool_result_bytes_after = masking.tool_result_bytes_after,
428 "CompactionCapability: masked stale tool results for model view"
429 );
430 }
431 masking.messages
432 }
433
434 fn priority(&self) -> i32 {
435 50
436 }
437}
438
439struct CompactionFilterProvider;
449
450impl MessageFilterProvider for CompactionFilterProvider {
451 fn apply_filters(
452 &self,
453 _query: &mut crate::message_filter::MessageQuery,
454 _config: &serde_json::Value,
455 ) {
456 }
464
465 fn priority(&self) -> i32 {
466 50 }
468}
469
470pub fn estimate_tokens(msg: &LlmMessage) -> usize {
479 let text_len = match &msg.content {
480 LlmMessageContent::Text(t) => t.len(),
481 LlmMessageContent::Parts(parts) => parts
482 .iter()
483 .map(|p| match p {
484 LlmContentPart::Text { text } => text.len(),
485 _ => 50, })
487 .sum(),
488 };
489
490 let tool_call_len = msg
492 .tool_calls
493 .as_ref()
494 .map(|calls| {
495 calls
496 .iter()
497 .map(|tc| tc.name.len() + tc.arguments.to_string().len() + 20)
498 .sum::<usize>()
499 })
500 .unwrap_or(0);
501
502 (text_len + tool_call_len) / 4
503}
504
505pub fn estimate_total_tokens(messages: &[LlmMessage]) -> usize {
507 messages.iter().map(estimate_tokens).sum()
508}
509
510pub fn should_compact_proactively(
515 messages: &[LlmMessage],
516 config: &CompactionConfig,
517 context_window_tokens: usize,
518) -> bool {
519 if !config.proactive {
520 return false;
521 }
522 let budget = (context_window_tokens as f32 * config.budget_percent) as usize;
523 let estimated = estimate_total_tokens(messages);
524 estimated > budget
525}
526
527pub fn aggressive_trim(
537 messages: &[LlmMessage],
538 target_tokens: usize,
539 has_system_prompt: bool,
540) -> Vec<LlmMessage> {
541 let mut result = Vec::new();
542 let mut token_budget = target_tokens;
543
544 let start_idx = if has_system_prompt && !messages.is_empty() {
546 let sys_tokens = estimate_tokens(&messages[0]);
547 if sys_tokens < token_budget {
548 result.push(messages[0].clone());
549 token_budget -= sys_tokens;
550 }
551 1
552 } else {
553 0
554 };
555
556 let conversation = &messages[start_idx..];
557
558 let mut protected_indices: std::collections::HashSet<usize> = conversation
561 .iter()
562 .enumerate()
563 .filter(|(_, m)| {
564 is_protected_tool_result(conversation, m) || is_protected_tool_call_message(m)
565 })
566 .map(|(i, _)| i)
567 .collect();
568
569 if !conversation.is_empty() {
577 protected_indices.insert(0);
578 }
579
580 let mut protected_budget: usize = 0;
581 for &idx in &protected_indices {
582 protected_budget += estimate_tokens(&conversation[idx]);
583 }
584
585 if protected_budget > token_budget {
588 let mut protected_with_indices: Vec<(usize, LlmMessage)> = protected_indices
589 .iter()
590 .map(|&idx| (idx, conversation[idx].clone()))
591 .collect();
592 protected_with_indices.sort_by_key(|(i, _)| *i);
593
594 let mut remaining = token_budget;
595 let mut kept: Vec<(usize, LlmMessage)> = Vec::new();
596 for (idx, msg) in protected_with_indices.into_iter().rev() {
597 let t = estimate_tokens(&msg);
598 if t <= remaining {
599 kept.push((idx, msg));
600 remaining -= t;
601 }
602 }
603 kept.sort_by_key(|(i, _)| *i);
604 result.extend(kept.into_iter().map(|(_, m)| m));
605 return crate::tool_call_integrity::retain_complete_llm_tool_exchanges(result);
606 }
607
608 token_budget -= protected_budget;
609
610 let mut keep_from_end = Vec::new();
612 for (i, msg) in conversation.iter().enumerate().rev() {
613 if protected_indices.contains(&i) {
614 continue; }
616 let msg_tokens = estimate_tokens(msg);
617 if msg_tokens <= token_budget {
618 keep_from_end.push((i, msg.clone()));
619 token_budget -= msg_tokens;
620 } else {
621 break;
622 }
623 }
624
625 let mut all_kept: Vec<(usize, LlmMessage)> = Vec::new();
627 for &idx in &protected_indices {
628 all_kept.push((idx, conversation[idx].clone()));
629 }
630 all_kept.extend(keep_from_end);
631 all_kept.sort_by_key(|(i, _)| *i);
632
633 result.extend(all_kept.into_iter().map(|(_, m)| m));
634 crate::tool_call_integrity::retain_complete_llm_tool_exchanges(result)
635}
636
637#[derive(Debug, Clone, Default, Serialize, Deserialize)]
643pub struct SessionCompactionMetrics {
644 pub compaction_count: u32,
646 pub total_messages_saved: u64,
648 pub strategy_counts: HashMap<String, u32>,
650 pub total_duration_ms: u64,
652}
653
654impl SessionCompactionMetrics {
655 pub fn record(
657 &mut self,
658 strategy_used: &str,
659 messages_before: usize,
660 messages_after: usize,
661 duration_ms: u64,
662 ) {
663 self.compaction_count += 1;
664 self.total_messages_saved += (messages_before.saturating_sub(messages_after)) as u64;
665 self.total_duration_ms += duration_ms;
666
667 for strategy in strategy_used.split('+') {
668 *self
669 .strategy_counts
670 .entry(strategy.to_string())
671 .or_insert(0) += 1;
672 }
673 }
674}
675
676#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
682#[serde(rename_all = "snake_case")]
683pub enum MemoryTier {
684 Hot,
686 Warm,
688 Cold,
690}
691
692#[derive(Debug, Clone, Serialize, Deserialize)]
694pub struct HierarchicalMemoryConfig {
695 #[serde(default = "default_hot_messages")]
697 pub hot_messages: usize,
698 #[serde(default = "default_warm_messages")]
700 pub warm_messages: usize,
701 }
703
704impl Default for HierarchicalMemoryConfig {
705 fn default() -> Self {
706 Self {
707 hot_messages: default_hot_messages(),
708 warm_messages: default_warm_messages(),
709 }
710 }
711}
712
713fn default_hot_messages() -> usize {
714 20
715}
716
717fn default_warm_messages() -> usize {
718 100
719}
720
721pub fn classify_memory_tiers<'a>(
725 messages: &'a [LlmMessage],
726 config: &HierarchicalMemoryConfig,
727) -> Vec<(MemoryTier, &'a LlmMessage)> {
728 let len = messages.len();
729 messages
730 .iter()
731 .enumerate()
732 .map(|(i, msg)| {
733 let from_end = len - 1 - i;
734 let tier = if from_end < config.hot_messages {
735 MemoryTier::Hot
736 } else if from_end < config.hot_messages + config.warm_messages {
737 MemoryTier::Warm
738 } else {
739 MemoryTier::Cold
740 };
741 (tier, msg)
742 })
743 .collect()
744}
745
746pub fn apply_hierarchical_memory(
754 messages: &[LlmMessage],
755 config: &HierarchicalMemoryConfig,
756 masking_config: &ObservationMaskingConfig,
757 cold_summary: Option<&str>,
758) -> Vec<LlmMessage> {
759 let len = messages.len();
760 let hot_start = len.saturating_sub(config.hot_messages);
761 let warm_start = hot_start.saturating_sub(config.warm_messages);
762
763 let mut result = Vec::new();
764
765 if warm_start > 0 {
767 let cold_msgs = &messages[..warm_start];
769 let protected_cold: Vec<LlmMessage> = cold_msgs
770 .iter()
771 .filter(|m| is_protected_tool_result(cold_msgs, m) || is_protected_tool_call_message(m))
772 .cloned()
773 .collect();
774
775 if let Some(summary) = cold_summary {
776 result.push(build_summary_message(summary));
777 }
778
779 result.extend(protected_cold);
781 }
782
783 if warm_start < hot_start {
787 let warm_msgs = &messages[warm_start..hot_start];
788
789 let protected_call_ids: std::collections::HashSet<String> = warm_msgs
791 .iter()
792 .filter(|m| is_protected_tool_result(messages, m))
793 .filter_map(|m| m.tool_call_id.clone())
794 .collect();
795
796 let masked = apply_observation_masking_with_protected(
797 warm_msgs,
798 masking_config,
799 &protected_call_ids,
800 );
801 result.extend(masked.messages);
802 }
803
804 if hot_start < len {
806 result.extend_from_slice(&messages[hot_start..]);
807 }
808
809 crate::tool_call_integrity::retain_complete_llm_tool_exchanges(result)
810}
811
812use crate::driver_registry::{LlmContentPart, LlmMessage, LlmMessageContent, LlmMessageRole};
817
818const PROTECTED_TOOL_NAMES: &[&str] = &["activate_skill"];
826
827fn is_protected_tool_result(messages: &[LlmMessage], tool_msg: &LlmMessage) -> bool {
832 if tool_msg.role != LlmMessageRole::Tool {
833 return false;
834 }
835 let tool_name = find_tool_call_name(messages, tool_msg);
836 PROTECTED_TOOL_NAMES.contains(&tool_name.as_str())
837}
838
839fn is_protected_tool_call_message(msg: &LlmMessage) -> bool {
843 if msg.role != LlmMessageRole::Assistant {
844 return false;
845 }
846 msg.tool_calls.as_ref().is_some_and(|calls| {
847 calls
848 .iter()
849 .any(|tc| PROTECTED_TOOL_NAMES.contains(&tc.name.as_str()))
850 })
851}
852
853#[derive(Debug)]
859pub struct ObservationMaskingResult {
860 pub messages: Vec<LlmMessage>,
862 pub masked_count: usize,
864}
865
866pub fn apply_observation_masking(
874 messages: &[LlmMessage],
875 config: &ObservationMaskingConfig,
876) -> ObservationMaskingResult {
877 apply_observation_masking_with_protected(messages, config, &std::collections::HashSet::new())
878}
879
880#[derive(Debug)]
882pub struct CostControlMaskingResult {
883 pub messages: Vec<Message>,
885 pub masked_count: usize,
887 pub tool_result_bytes_before: usize,
889 pub tool_result_bytes_after: usize,
891}
892
893pub fn build_model_view_messages(
899 stored_messages: &[Message],
900 compaction_config: &CompactionConfig,
901 prior_usage: Option<&TokenUsage>,
902) -> CostControlMaskingResult {
903 apply_cost_control_masking(stored_messages, compaction_config, prior_usage)
904}
905
906pub fn build_model_view_messages_owned(
910 stored_messages: Vec<Message>,
911 compaction_config: &CompactionConfig,
912 prior_usage: Option<&TokenUsage>,
913) -> CostControlMaskingResult {
914 apply_cost_control_masking_owned(stored_messages, compaction_config, prior_usage)
915}
916
917pub fn apply_cost_control_masking(
926 messages: &[Message],
927 config: &CompactionConfig,
928 prior_usage: Option<&TokenUsage>,
929) -> CostControlMaskingResult {
930 apply_cost_control_masking_owned(messages.to_vec(), config, prior_usage)
931}
932
933fn apply_cost_control_masking_owned(
934 messages: Vec<Message>,
935 config: &CompactionConfig,
936 prior_usage: Option<&TokenUsage>,
937) -> CostControlMaskingResult {
938 let cost_config = &config.cost_control;
939 let tool_indices: Vec<usize> = messages
940 .iter()
941 .enumerate()
942 .filter(|(_, message)| {
943 message.role == MessageRole::ToolResult
944 && !is_protected_message_tool_result(&messages, message)
945 })
946 .map(|(index, _)| index)
947 .collect();
948 let tool_result_bytes_before = tool_indices
949 .iter()
950 .map(|index| message_tool_result_len(&messages[*index]))
951 .sum();
952
953 if !cost_config.enabled
954 || tool_indices.len() <= cost_config.keep_recent_tool_results
955 || !should_apply_cost_control_masking(
956 tool_indices.len(),
957 tool_result_bytes_before,
958 cost_config,
959 prior_usage,
960 )
961 {
962 return CostControlMaskingResult {
963 messages,
964 masked_count: 0,
965 tool_result_bytes_before,
966 tool_result_bytes_after: tool_result_bytes_before,
967 };
968 }
969
970 let keep_recent = cost_config.keep_recent_tool_results;
971 let to_mask_count = tool_indices.len().saturating_sub(keep_recent);
972 let related_recent_reads =
973 related_recent_paginated_read_results(&messages, &tool_indices, keep_recent);
974 let indices_to_mask: HashSet<usize> = tool_indices[..to_mask_count]
975 .iter()
976 .copied()
977 .filter(|index| !related_recent_reads.contains(index))
978 .collect();
979 let tool_names: std::collections::HashMap<usize, String> = indices_to_mask
980 .iter()
981 .map(|index| {
982 (
983 *index,
984 find_message_tool_call_name(&messages, &messages[*index]),
985 )
986 })
987 .collect();
988
989 let mut masked_count = 0;
990 let mut masked_messages = Vec::with_capacity(messages.len());
991 for (index, message) in messages.into_iter().enumerate() {
992 if let Some(tool_name) = tool_names.get(&index) {
993 masked_messages.push(mask_tool_result_message(&message, tool_name));
994 masked_count += 1;
995 } else {
996 masked_messages.push(message);
997 }
998 }
999
1000 let tool_result_bytes_after = masked_messages
1001 .iter()
1002 .filter(|message| message.role == MessageRole::ToolResult)
1003 .map(message_tool_result_len)
1004 .sum();
1005
1006 CostControlMaskingResult {
1007 messages: masked_messages,
1008 masked_count,
1009 tool_result_bytes_before,
1010 tool_result_bytes_after,
1011 }
1012}
1013
1014fn should_apply_cost_control_masking(
1015 tool_result_count: usize,
1016 tool_result_bytes: usize,
1017 config: &CostControlConfig,
1018 prior_usage: Option<&TokenUsage>,
1019) -> bool {
1020 if tool_result_count >= config.mask_after_tool_results {
1021 return true;
1022 }
1023 if tool_result_bytes >= config.max_live_tool_result_bytes {
1024 return true;
1025 }
1026 let Some(usage) = prior_usage else {
1027 return false;
1028 };
1029 let cache_read = usage.cache_read_tokens.unwrap_or(0);
1033 let cache_creation = usage.cache_creation_tokens.unwrap_or(0);
1034 if usage.input_tokens >= config.max_uncached_input_tokens {
1035 return true;
1036 }
1037 let total_prompt = usage.input_tokens + cache_read + cache_creation;
1038 total_prompt > 0 && (cache_read as f32 / total_prompt as f32) < config.min_cache_read_ratio
1039}
1040
1041fn is_protected_message_tool_result(messages: &[Message], tool_msg: &Message) -> bool {
1042 if tool_msg.role != MessageRole::ToolResult {
1043 return false;
1044 }
1045 let tool_name = find_message_tool_call_name(messages, tool_msg);
1046 PROTECTED_TOOL_NAMES.contains(&tool_name.as_str())
1047}
1048
1049#[derive(Debug, Clone, PartialEq, Eq, Hash)]
1050struct ReadResultKey {
1051 tool_name: String,
1052 path: String,
1053 content_hash: String,
1054}
1055
1056fn related_recent_paginated_read_results(
1057 messages: &[Message],
1058 tool_indices: &[usize],
1059 keep_recent: usize,
1060) -> HashSet<usize> {
1061 if keep_recent == 0 || tool_indices.len() <= keep_recent {
1062 return HashSet::new();
1063 }
1064
1065 let keep_start = tool_indices.len().saturating_sub(keep_recent);
1066 let recent_keys: HashSet<ReadResultKey> = tool_indices[keep_start..]
1067 .iter()
1068 .filter_map(|index| paginated_read_result_key(messages, &messages[*index]))
1069 .collect();
1070 if recent_keys.is_empty() {
1071 return HashSet::new();
1072 }
1073
1074 let mut protected = HashSet::new();
1075 for index in tool_indices[..keep_start].iter().rev() {
1076 let Some(key) = paginated_read_result_key(messages, &messages[*index]) else {
1077 break;
1078 };
1079 if !recent_keys.contains(&key) {
1080 break;
1081 }
1082 protected.insert(*index);
1083 if protected.len() >= MAX_RELATED_RECENT_READ_RESULTS {
1084 break;
1085 }
1086 }
1087 protected
1088}
1089
1090fn paginated_read_result_key(messages: &[Message], tool_msg: &Message) -> Option<ReadResultKey> {
1091 let tool_name = find_message_tool_call_name(messages, tool_msg);
1092 if !is_read_file_tool_name(&tool_name) {
1093 return None;
1094 }
1095 let value = tool_msg.tool_result_content()?.result.as_ref()?;
1096 let object = value.as_object()?;
1097 object.get("lines_shown").and_then(|v| v.as_object())?;
1098 Some(ReadResultKey {
1099 tool_name,
1100 path: object.get("path")?.as_str()?.to_string(),
1101 content_hash: object.get("content_hash")?.as_str()?.to_string(),
1102 })
1103}
1104
1105fn is_read_file_tool_name(tool_name: &str) -> bool {
1106 matches!(
1107 tool_name,
1108 "read_file"
1109 | "daytona_read_file"
1110 | "sandbox_read_file"
1111 | "e2b_read_file"
1112 | "docker_read_file"
1113 | "deno_read_file"
1114 | "sprites_read_file"
1115 | "read_github_file"
1116 )
1117}
1118
1119fn find_message_tool_call_name(messages: &[Message], tool_msg: &Message) -> String {
1120 let Some(call_id) = tool_msg.tool_call_id() else {
1121 return "unknown_tool".to_string();
1122 };
1123
1124 for msg in messages.iter().rev() {
1125 if msg.role != MessageRole::Agent {
1126 continue;
1127 }
1128 for tool_call in msg.tool_calls() {
1129 if tool_call.id == call_id {
1130 return tool_call.name.clone();
1131 }
1132 }
1133 }
1134
1135 "unknown_tool".to_string()
1136}
1137
1138fn message_tool_result_len(message: &Message) -> usize {
1139 let Some(result) = message.tool_result_content() else {
1140 return 0;
1141 };
1142 result
1143 .result
1144 .as_ref()
1145 .map(estimate_json_value_len)
1146 .unwrap_or(0)
1147 + result.error.as_ref().map_or(0, String::len)
1148}
1149
1150fn mask_tool_result_message(message: &Message, tool_name: &str) -> Message {
1151 let Some(result) = message.tool_result_content() else {
1152 return message.clone();
1153 };
1154 let summary = summarize_tool_result(tool_name, result.result.as_ref(), result.error.as_ref());
1155 let was_error = result.error.is_some();
1156 let mut masked = message.clone();
1157 for part in &mut masked.content {
1158 if let ContentPart::ToolResult(tool_result) = part {
1159 if was_error {
1160 tool_result.result = None;
1161 tool_result.error = Some(summary);
1162 } else {
1163 tool_result.result = Some(serde_json::json!({
1164 "masked": true,
1165 "summary": summary,
1166 }));
1167 tool_result.error = None;
1168 }
1169 break;
1170 }
1171 }
1172 masked
1173}
1174
1175fn summarize_tool_result(
1176 tool_name: &str,
1177 result: Option<&serde_json::Value>,
1178 error: Option<&String>,
1179) -> String {
1180 if let Some(error) = error {
1181 return format!("[{tool_name} error: {}]", truncate_inline(error, 160));
1182 }
1183 let Some(value) = result else {
1184 return format!("[{tool_name} returned no result]");
1185 };
1186 let Some(object) = value.as_object() else {
1187 return format!(
1188 "[{tool_name} -> {}, {} bytes]",
1189 value_kind(value),
1190 estimate_json_value_len(value)
1191 );
1192 };
1193
1194 match tool_name {
1195 tool_name if is_read_file_tool_name(tool_name) => {
1196 summarize_read_file_result(tool_name, object, value)
1197 }
1198 "bash" | "daytona_exec" | "sandbox_exec" | "e2b_exec" | "docker_exec" | "deno_exec" => {
1199 summarize_exec_result(tool_name, object, value)
1200 }
1201 "list_directory" => summarize_list_directory_result(tool_name, object, value),
1202 "grep_files" => summarize_grep_files_result(tool_name, object, value),
1203 _ => summarize_generic_tool_result(tool_name, object, value),
1204 }
1205}
1206
1207fn summarize_read_file_result(
1208 tool_name: &str,
1209 object: &serde_json::Map<String, serde_json::Value>,
1210 value: &serde_json::Value,
1211) -> String {
1212 let path = object
1213 .get("path")
1214 .and_then(|v| v.as_str())
1215 .unwrap_or("(unknown path)");
1216 let lines = object.get("lines_shown").and_then(|v| v.as_object());
1217 let line_range = lines
1218 .and_then(|lines| {
1219 let start = lines.get("start")?.as_u64()?;
1220 let end = lines.get("end")?.as_u64()?;
1221 Some(format!(" lines {start}-{end}"))
1222 })
1223 .unwrap_or_default();
1224 let total_lines = object
1225 .get("total_lines")
1226 .and_then(|v| v.as_u64())
1227 .map(|lines| format!(", total_lines={lines}"))
1228 .unwrap_or_default();
1229 let next_offset = object
1230 .get("truncation")
1231 .and_then(|v| v.as_object())
1232 .and_then(|truncation| truncation.get("next_offset"))
1233 .and_then(|v| v.as_u64())
1234 .map(|offset| format!(", next_offset={offset}"))
1235 .unwrap_or_default();
1236 let hash = object
1237 .get("content_hash")
1238 .and_then(|v| v.as_str())
1239 .map(|hash| format!(", hash={hash}"))
1240 .unwrap_or_default();
1241 let truncated = object
1242 .get("truncated")
1243 .and_then(|v| v.as_bool())
1244 .unwrap_or(false);
1245
1246 format!(
1247 "[{tool_name} {path}{line_range}, {} bytes, truncated={truncated}{total_lines}{next_offset}{hash}]",
1248 estimate_json_value_len(value)
1249 )
1250}
1251
1252fn summarize_exec_result(
1253 tool_name: &str,
1254 object: &serde_json::Map<String, serde_json::Value>,
1255 value: &serde_json::Value,
1256) -> String {
1257 let exit = object
1258 .get("exit_code")
1259 .and_then(|v| v.as_i64())
1260 .map(|code| format!(" exit={code}"))
1261 .unwrap_or_default();
1262 let stdout_len = object
1263 .get("stdout")
1264 .and_then(|v| v.as_str())
1265 .map(|stdout| stdout.len())
1266 .unwrap_or(0);
1267 let stderr_len = object
1268 .get("stderr")
1269 .and_then(|v| v.as_str())
1270 .map(|stderr| stderr.len())
1271 .unwrap_or(0);
1272 let full_output = object
1273 .get("full_output")
1274 .and_then(|v| v.as_str())
1275 .map(|path| format!(", full_output={path}"))
1276 .unwrap_or_default();
1277 let total_lines = object
1278 .get("total_lines")
1279 .and_then(|v| v.as_u64())
1280 .map(|lines| format!(", total_lines={lines}"))
1281 .unwrap_or_default();
1282
1283 format!(
1284 "[{tool_name}{exit}, stdout={} bytes, stderr={} bytes, result={} bytes{full_output}{total_lines}]",
1285 stdout_len,
1286 stderr_len,
1287 estimate_json_value_len(value)
1288 )
1289}
1290
1291fn summarize_list_directory_result(
1292 tool_name: &str,
1293 object: &serde_json::Map<String, serde_json::Value>,
1294 value: &serde_json::Value,
1295) -> String {
1296 let path = object
1297 .get("path")
1298 .and_then(|v| v.as_str())
1299 .unwrap_or("(unknown path)");
1300 let count = object
1301 .get("count")
1302 .and_then(|v| v.as_u64())
1303 .or_else(|| {
1304 object
1305 .get("entries")
1306 .and_then(|v| v.as_array())
1307 .map(|v| v.len() as u64)
1308 })
1309 .unwrap_or(0);
1310 format!(
1311 "[{tool_name} {path}, {count} entries, {} bytes]",
1312 estimate_json_value_len(value)
1313 )
1314}
1315
1316fn summarize_grep_files_result(
1317 tool_name: &str,
1318 object: &serde_json::Map<String, serde_json::Value>,
1319 value: &serde_json::Value,
1320) -> String {
1321 let pattern = object
1322 .get("pattern")
1323 .and_then(|v| v.as_str())
1324 .map(|pattern| format!(" pattern={:?}", truncate_inline(pattern, 80)))
1325 .unwrap_or_default();
1326 let match_count = object
1327 .get("match_count")
1328 .and_then(|v| v.as_u64())
1329 .unwrap_or(0);
1330 format!(
1331 "[{tool_name}{pattern}, matches={match_count}, {} bytes]",
1332 estimate_json_value_len(value)
1333 )
1334}
1335
1336fn summarize_generic_tool_result(
1337 tool_name: &str,
1338 object: &serde_json::Map<String, serde_json::Value>,
1339 value: &serde_json::Value,
1340) -> String {
1341 let keys = object.keys().take(5).cloned().collect::<Vec<_>>().join(",");
1342 format!(
1343 "[{tool_name} result, {} bytes, keys={keys}]",
1344 estimate_json_value_len(value)
1345 )
1346}
1347
1348fn value_kind(value: &serde_json::Value) -> &'static str {
1349 match value {
1350 serde_json::Value::Null => "null",
1351 serde_json::Value::Bool(_) => "bool",
1352 serde_json::Value::Number(_) => "number",
1353 serde_json::Value::String(_) => "string",
1354 serde_json::Value::Array(_) => "array",
1355 serde_json::Value::Object(_) => "object",
1356 }
1357}
1358
1359fn estimate_json_value_len(value: &serde_json::Value) -> usize {
1360 let mut writer = CountingWriter::default();
1361 serde_json::to_writer(&mut writer, value)
1362 .map(|_| writer.bytes)
1363 .unwrap_or(0)
1364}
1365
1366#[derive(Default)]
1367struct CountingWriter {
1368 bytes: usize,
1369}
1370
1371impl std::io::Write for CountingWriter {
1372 fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
1373 self.bytes += buf.len();
1374 Ok(buf.len())
1375 }
1376
1377 fn flush(&mut self) -> std::io::Result<()> {
1378 Ok(())
1379 }
1380}
1381
1382fn truncate_inline(text: &str, max_chars: usize) -> String {
1383 if text.chars().count() <= max_chars {
1384 return text.to_string();
1385 }
1386 let mut truncated = text.chars().take(max_chars).collect::<String>();
1387 truncated.push_str("...");
1388 truncated
1389}
1390
1391fn apply_observation_masking_with_protected(
1395 messages: &[LlmMessage],
1396 config: &ObservationMaskingConfig,
1397 extra_protected_call_ids: &std::collections::HashSet<String>,
1398) -> ObservationMaskingResult {
1399 let tool_indices: Vec<usize> = messages
1401 .iter()
1402 .enumerate()
1403 .filter(|(_, m)| {
1404 m.role == LlmMessageRole::Tool
1405 && !is_protected_tool_result(messages, m)
1406 && !m
1407 .tool_call_id
1408 .as_ref()
1409 .is_some_and(|id| extra_protected_call_ids.contains(id))
1410 })
1411 .map(|(i, _)| i)
1412 .collect();
1413
1414 if tool_indices.len() <= config.keep_recent_tool_outputs {
1415 return ObservationMaskingResult {
1416 messages: messages.to_vec(),
1417 masked_count: 0,
1418 };
1419 }
1420
1421 let to_mask_count = tool_indices.len() - config.keep_recent_tool_outputs;
1422 let indices_to_mask: std::collections::HashSet<usize> =
1423 tool_indices[..to_mask_count].iter().copied().collect();
1424
1425 let mut result = Vec::with_capacity(messages.len());
1426 let mut masked_count = 0;
1427
1428 for (i, msg) in messages.iter().enumerate() {
1429 if indices_to_mask.contains(&i) {
1430 let tool_name = find_tool_call_name(messages, msg);
1431 let summary = match config.summary_format {
1432 MaskingSummaryFormat::OneLine => format_one_line_summary(&tool_name, &msg.content),
1433 MaskingSummaryFormat::HeadTail => format_head_tail_summary(&msg.content),
1434 };
1435 result.push(LlmMessage {
1436 role: LlmMessageRole::Tool,
1437 content: LlmMessageContent::Text(summary),
1438 tool_calls: msg.tool_calls.clone(),
1439 tool_call_id: msg.tool_call_id.clone(),
1440 phase: msg.phase,
1441 thinking: None,
1442 thinking_signature: None,
1443 });
1444 masked_count += 1;
1445 } else {
1446 result.push(msg.clone());
1447 }
1448 }
1449
1450 ObservationMaskingResult {
1451 messages: result,
1452 masked_count,
1453 }
1454}
1455
1456fn find_tool_call_name(messages: &[LlmMessage], tool_msg: &LlmMessage) -> String {
1458 let Some(ref call_id) = tool_msg.tool_call_id else {
1459 return "unknown_tool".to_string();
1460 };
1461
1462 for msg in messages.iter().rev() {
1463 if msg.role == LlmMessageRole::Assistant
1464 && let Some(ref tool_calls) = msg.tool_calls
1465 {
1466 for tc in tool_calls {
1467 if tc.id == *call_id {
1468 return tc.name.clone();
1469 }
1470 }
1471 }
1472 }
1473
1474 "unknown_tool".to_string()
1475}
1476
1477fn extract_text(content: &LlmMessageContent) -> String {
1478 match content {
1479 LlmMessageContent::Text(t) => t.clone(),
1480 LlmMessageContent::Parts(parts) => parts
1481 .iter()
1482 .filter_map(|p| {
1483 if let LlmContentPart::Text { text } = p {
1484 Some(text.clone())
1485 } else {
1486 None
1487 }
1488 })
1489 .collect::<Vec<_>>()
1490 .join(" "),
1491 }
1492}
1493
1494fn format_one_line_summary(tool_name: &str, content: &LlmMessageContent) -> String {
1495 let text = extract_text(content);
1496 let line_count = text.lines().count();
1497 let byte_len = text.len();
1498
1499 if byte_len <= 100 {
1500 format!("[{tool_name} → {text}]")
1501 } else {
1502 format!("[{tool_name} → {line_count} lines, {byte_len} bytes]")
1503 }
1504}
1505
1506fn format_head_tail_summary(content: &LlmMessageContent) -> String {
1507 let text = extract_text(content);
1508 let lines: Vec<&str> = text.lines().collect();
1509
1510 if lines.len() <= 6 {
1511 return text;
1512 }
1513
1514 let head: Vec<&str> = lines[..3].to_vec();
1515 let tail: Vec<&str> = lines[lines.len() - 3..].to_vec();
1516
1517 format!(
1518 "{}\n... ({} lines omitted) ...\n{}",
1519 head.join("\n"),
1520 lines.len() - 6,
1521 tail.join("\n")
1522 )
1523}
1524
1525pub fn build_summarization_prompt(config: &SummarizationConfig) -> String {
1531 let preserve_items = if config.preserve.is_empty() {
1532 default_preserve()
1533 } else {
1534 config.preserve.clone()
1535 };
1536
1537 let preserve_list = preserve_items
1538 .iter()
1539 .map(|item| format!("- {item}"))
1540 .collect::<Vec<_>>()
1541 .join("\n");
1542
1543 let custom_instructions = config
1544 .instructions
1545 .as_deref()
1546 .map(|instr| format!("\n- {instr}"))
1547 .unwrap_or_default();
1548
1549 format!(
1550 r#"<task>
1551Summarize the following conversation history. The summary replaces these
1552messages in the agent's context window — it must contain everything the
1553agent needs to continue working.
1554</task>
1555
1556<preserve>
1557{preserve_list}{custom_instructions}
1558</preserve>
1559
1560<format>
1561Produce a structured summary. Use sections. Be concise but complete.
1562Do not include tool output verbatim — reference files by path.
1563IMPORTANT: Any activate_skill tool results contain durable skill instructions.
1564Include them verbatim in a dedicated "Active Skills" section — do not summarize
1565or paraphrase skill instructions.
1566</format>"#
1567 )
1568}
1569
1570pub fn format_messages_for_summarization(messages: &[LlmMessage]) -> String {
1572 let mut parts = Vec::new();
1573 for msg in messages {
1574 let role = match msg.role {
1575 LlmMessageRole::System => "system",
1576 LlmMessageRole::User => "user",
1577 LlmMessageRole::Assistant => "assistant",
1578 LlmMessageRole::Tool => "tool",
1579 };
1580
1581 let content = extract_text(&msg.content);
1582
1583 let is_protected = is_protected_tool_result(messages, msg);
1586
1587 let truncated = if !is_protected && content.len() > 2000 {
1589 let safe_prefix = truncate_at_char_boundary(&content, 2000);
1590 format!(
1591 "{}... [truncated, {} chars total]",
1592 safe_prefix,
1593 content.len()
1594 )
1595 } else {
1596 content
1597 };
1598
1599 parts.push(format!("[{role}]: {truncated}"));
1600 }
1601 parts.join("\n\n")
1602}
1603
1604fn truncate_at_char_boundary(content: &str, max_bytes: usize) -> &str {
1605 if content.len() <= max_bytes {
1606 return content;
1607 }
1608
1609 if content.is_char_boundary(max_bytes) {
1610 return &content[..max_bytes];
1611 }
1612
1613 let mut end = max_bytes;
1614 while end > 0 && !content.is_char_boundary(end) {
1615 end -= 1;
1616 }
1617
1618 &content[..end]
1619}
1620
1621pub fn build_summary_message(summary_text: &str) -> LlmMessage {
1623 LlmMessage {
1624 role: LlmMessageRole::System,
1625 content: LlmMessageContent::Text(format!(
1626 "[CONVERSATION_SUMMARY]\n{summary_text}\n[/CONVERSATION_SUMMARY]"
1627 )),
1628 tool_calls: None,
1629 tool_call_id: None,
1630 phase: None,
1631 thinking: None,
1632 thinking_signature: None,
1633 }
1634}
1635
1636pub fn compose_summary_with_recent(
1642 system_message: Option<LlmMessage>,
1643 summary_text: &str,
1644 recent_messages: &[LlmMessage],
1645) -> Vec<LlmMessage> {
1646 let mut messages = Vec::with_capacity(recent_messages.len() + 2);
1647 if let Some(system_message) = system_message {
1648 messages.push(system_message);
1649 }
1650 messages.push(build_summary_message(summary_text));
1651 messages.extend_from_slice(recent_messages);
1652 crate::tool_call_integrity::retain_complete_llm_tool_exchanges(messages)
1653}
1654
1655#[derive(Debug, Clone, Serialize, Deserialize)]
1661pub struct CompactionStep {
1662 pub strategy: String,
1664 pub messages_after: usize,
1666 pub duration_ms: u64,
1668}
1669
1670#[cfg(test)]
1675mod tests {
1676 use super::*;
1677 use crate::tool_types::ToolCall;
1678 use serde_json::json;
1679
1680 fn assert_complete_tool_exchanges(messages: &[LlmMessage]) {
1681 let calls: std::collections::HashSet<_> = messages
1682 .iter()
1683 .flat_map(|message| message.tool_calls.iter().flatten())
1684 .map(|call| call.id.as_str())
1685 .collect();
1686 let results: std::collections::HashSet<_> = messages
1687 .iter()
1688 .filter_map(|message| message.tool_call_id.as_deref())
1689 .collect();
1690 assert_eq!(calls, results, "tool calls and results must remain atomic");
1691 }
1692
1693 fn make_user_msg(text: &str) -> LlmMessage {
1694 LlmMessage {
1695 role: LlmMessageRole::User,
1696 content: LlmMessageContent::Text(text.to_string()),
1697 tool_calls: None,
1698 tool_call_id: None,
1699 phase: None,
1700 thinking: None,
1701 thinking_signature: None,
1702 }
1703 }
1704
1705 fn make_assistant_msg(text: &str) -> LlmMessage {
1706 LlmMessage {
1707 role: LlmMessageRole::Assistant,
1708 content: LlmMessageContent::Text(text.to_string()),
1709 tool_calls: None,
1710 tool_call_id: None,
1711 phase: None,
1712 thinking: None,
1713 thinking_signature: None,
1714 }
1715 }
1716
1717 fn make_assistant_with_tool_call(call_id: &str, tool_name: &str) -> LlmMessage {
1718 LlmMessage {
1719 role: LlmMessageRole::Assistant,
1720 content: LlmMessageContent::Text(String::new()),
1721 tool_calls: Some(vec![ToolCall {
1722 id: call_id.to_string(),
1723 name: tool_name.to_string(),
1724 arguments: json!({"path": "src/main.rs"}),
1725 }]),
1726 tool_call_id: None,
1727 phase: None,
1728 thinking: None,
1729 thinking_signature: None,
1730 }
1731 }
1732
1733 fn make_assistant_with_tool_calls(calls: &[(&str, &str)]) -> LlmMessage {
1734 LlmMessage {
1735 role: LlmMessageRole::Assistant,
1736 content: LlmMessageContent::Text(String::new()),
1737 tool_calls: Some(
1738 calls
1739 .iter()
1740 .map(|(call_id, tool_name)| ToolCall {
1741 id: (*call_id).to_string(),
1742 name: (*tool_name).to_string(),
1743 arguments: json!({}),
1744 })
1745 .collect(),
1746 ),
1747 tool_call_id: None,
1748 phase: None,
1749 thinking: None,
1750 thinking_signature: None,
1751 }
1752 }
1753
1754 fn make_tool_result(call_id: &str, output: &str) -> LlmMessage {
1755 LlmMessage {
1756 role: LlmMessageRole::Tool,
1757 content: LlmMessageContent::Text(output.to_string()),
1758 tool_calls: None,
1759 tool_call_id: Some(call_id.to_string()),
1760 phase: None,
1761 thinking: None,
1762 thinking_signature: None,
1763 }
1764 }
1765
1766 #[test]
1771 fn test_capability_metadata() {
1772 let cap = CompactionCapability;
1773 assert_eq!(cap.id(), COMPACTION_CAPABILITY_ID);
1774 assert_eq!(cap.name(), "Compaction");
1775 assert_eq!(cap.status(), CapabilityStatus::Available);
1776 assert_eq!(cap.category(), Some("Optimization"));
1777 assert!(cap.tools().is_empty());
1778 assert!(cap.message_filter_provider().is_some());
1779 }
1780
1781 #[test]
1782 fn test_config_schema_and_validate_config() {
1783 let cap = CompactionCapability;
1784
1785 let schema = cap.config_schema().expect("config schema");
1786 assert_eq!(schema["type"], "object");
1787 assert!(schema["properties"]["strategy"].is_object());
1789 assert!(schema["properties"]["proactive"].is_object());
1790 assert!(schema["properties"]["budget_percent"].is_object());
1791 assert!(schema["properties"].get("observation_masking").is_none());
1792 assert!(schema["properties"].get("cost_control").is_none());
1793
1794 assert!(cap.validate_config(&serde_json::Value::Null).is_ok());
1796 assert!(
1797 cap.validate_config(&json!({
1798 "strategy": "native",
1799 "proactive": false,
1800 "budget_percent": 0.9
1801 }))
1802 .is_ok()
1803 );
1804 assert!(
1806 cap.validate_config(&json!({
1807 "strategy": "observation_masking",
1808 "observation_masking": { "keep_recent_tool_outputs": 4 },
1809 "cost_control": { "enabled": false }
1810 }))
1811 .is_ok()
1812 );
1813
1814 assert!(cap.validate_config(&json!({"strategy": "bogus"})).is_err());
1816 let err = cap
1817 .validate_config(&json!({"budget_percent": 5.0}))
1818 .unwrap_err();
1819 assert!(err.contains("budget_percent"));
1820 }
1821
1822 #[test]
1823 fn test_localizations_resolve_uk() {
1824 let cap = CompactionCapability;
1825 assert_eq!(cap.localized_name(Some("uk-UA")), "Ущільнення контексту");
1826 assert!(cap.describe_schema(None).is_some());
1827 }
1828
1829 #[test]
1830 fn test_default_config() {
1831 let config = CompactionConfig::default();
1832 assert_eq!(config.strategy, CompactionStrategy::Auto);
1833 assert!(config.proactive);
1834 assert!((config.budget_percent - 0.85).abs() < f32::EPSILON);
1835 assert_eq!(config.observation_masking.keep_recent_tool_outputs, 2);
1836 assert_eq!(
1837 config.observation_masking.summary_format,
1838 MaskingSummaryFormat::OneLine
1839 );
1840 assert!(config.summarization.model.is_none());
1841 assert_eq!(config.summarization.preserve.len(), 5);
1842 assert!(config.summarization.instructions.is_none());
1843 assert!(config.cost_control.enabled);
1844 assert_eq!(config.cost_control.keep_recent_tool_results, 2);
1845 }
1846
1847 #[test]
1848 fn test_config_from_empty_json() {
1849 let config = CompactionConfig::from_json(&json!({}));
1850 assert_eq!(config.strategy, CompactionStrategy::Auto);
1851 assert!(config.proactive);
1852 }
1853
1854 #[test]
1855 fn test_config_native_only() {
1856 let config = CompactionConfig::from_json(&json!({"strategy": "native"}));
1857 assert_eq!(config.strategy, CompactionStrategy::Native);
1858 assert!(config.proactive);
1859 }
1860
1861 #[test]
1862 fn test_config_observation_masking_with_custom_settings() {
1863 let config = CompactionConfig::from_json(&json!({
1864 "strategy": "observation_masking",
1865 "proactive": false,
1866 "observation_masking": {
1867 "keep_recent_tool_outputs": 10,
1868 "summary_format": "head_tail"
1869 }
1870 }));
1871 assert_eq!(config.strategy, CompactionStrategy::ObservationMasking);
1872 assert!(!config.proactive);
1873 assert_eq!(config.observation_masking.keep_recent_tool_outputs, 10);
1874 assert_eq!(
1875 config.observation_masking.summary_format,
1876 MaskingSummaryFormat::HeadTail
1877 );
1878 }
1879
1880 #[test]
1881 fn test_config_cost_control_with_custom_settings() {
1882 let config = CompactionConfig::from_json(&json!({
1883 "cost_control": {
1884 "enabled": true,
1885 "keep_recent_tool_results": 1,
1886 "mask_after_tool_results": 2,
1887 "max_live_tool_result_bytes": 4096,
1888 "max_uncached_input_tokens": 50000,
1889 "min_cache_read_ratio": 0.5
1890 }
1891 }));
1892
1893 assert!(config.cost_control.enabled);
1894 assert_eq!(config.cost_control.keep_recent_tool_results, 1);
1895 assert_eq!(config.cost_control.mask_after_tool_results, 2);
1896 assert_eq!(config.cost_control.max_live_tool_result_bytes, 4096);
1897 assert_eq!(config.cost_control.max_uncached_input_tokens, 50000);
1898 assert!((config.cost_control.min_cache_read_ratio - 0.5).abs() < f32::EPSILON);
1899 }
1900
1901 #[test]
1902 fn test_config_summarization_with_custom_model() {
1903 let config = CompactionConfig::from_json(&json!({
1904 "strategy": "summarization",
1905 "summarization": {
1906 "model": "claude-haiku-4-5-20251001",
1907 "instructions": "Focus on API decisions",
1908 "preserve": ["decisions", "errors"]
1909 }
1910 }));
1911 assert_eq!(config.strategy, CompactionStrategy::Summarization);
1912 assert_eq!(
1913 config.summarization.model.as_deref(),
1914 Some("claude-haiku-4-5-20251001")
1915 );
1916 assert_eq!(
1917 config.summarization.instructions.as_deref(),
1918 Some("Focus on API decisions")
1919 );
1920 assert_eq!(config.summarization.preserve.len(), 2);
1921 }
1922
1923 fn make_message_tool_turn(
1924 call_id: &str,
1925 tool_name: &str,
1926 result: serde_json::Value,
1927 ) -> Vec<Message> {
1928 vec![
1929 Message::assistant_with_tools(
1930 "",
1931 vec![ToolCall {
1932 id: call_id.to_string(),
1933 name: tool_name.to_string(),
1934 arguments: json!({"path": "/workspace/src/lib.rs"}),
1935 }],
1936 ),
1937 Message::tool_result(call_id, Some(result), None),
1938 ]
1939 }
1940
1941 #[test]
1942 fn test_cost_control_masks_old_read_file_results() {
1943 let mut messages = vec![Message::user("inspect files")];
1944 for index in 0..5 {
1945 messages.extend(make_message_tool_turn(
1946 &format!("call_{index}"),
1947 "read_file",
1948 json!({
1949 "path": "/workspace/src/lib.rs",
1950 "content": format!("{}{}", "line\n".repeat(400), index),
1951 "total_lines": 900,
1952 "lines_shown": {"start": 1, "end": 400},
1953 "truncated": true,
1954 "content_hash": format!("sha256:{index}"),
1955 "truncation": {"truncated": true, "next_offset": 400, "reason": "line_cap"}
1956 }),
1957 ));
1958 }
1959
1960 let config = CompactionConfig::from_json(&json!({
1961 "cost_control": {
1962 "keep_recent_tool_results": 2,
1963 "mask_after_tool_results": 4
1964 }
1965 }));
1966 let result = apply_cost_control_masking(&messages, &config, None);
1967
1968 assert_eq!(result.masked_count, 3);
1969 assert!(result.tool_result_bytes_after < result.tool_result_bytes_before);
1970
1971 let first_tool = result.messages[2].tool_result_content().unwrap();
1972 let masked = first_tool.result.as_ref().unwrap();
1973 assert_eq!(masked["masked"], true);
1974 let summary = masked["summary"].as_str().unwrap();
1975 assert!(summary.contains("read_file"));
1976 assert!(summary.contains("/workspace/src/lib.rs"));
1977 assert!(summary.contains("lines 1-400"));
1978 assert!(summary.contains("next_offset=400"));
1979 assert!(!summary.contains("line\nline"));
1980
1981 let last_tool = result
1982 .messages
1983 .last()
1984 .unwrap()
1985 .tool_result_content()
1986 .unwrap();
1987 assert!(last_tool.result.as_ref().unwrap().get("content").is_some());
1988 }
1989
1990 #[test]
1991 fn test_cost_control_keeps_recent_paginated_read_group() {
1992 let mut messages = vec![Message::user("inspect saved output")];
1993 messages.extend(make_message_tool_turn(
1994 "call_bash",
1995 "bash",
1996 json!({
1997 "stdout": "old command output",
1998 "stderr": "",
1999 "exit_code": 0,
2000 "success": true
2001 }),
2002 ));
2003 messages.extend(make_message_tool_turn(
2004 "call_read_first",
2005 "read_file",
2006 json!({
2007 "path": "/workspace/outputs/call_123.stdout",
2008 "content": "first page\n".repeat(200),
2009 "total_lines": 400,
2010 "lines_shown": {"start": 1, "end": 200},
2011 "truncated": true,
2012 "content_hash": "sha256:same-output",
2013 "truncation": {"truncated": true, "next_offset": 200, "reason": "line_cap"}
2014 }),
2015 ));
2016 messages.extend(make_message_tool_turn(
2017 "call_read_second",
2018 "read_file",
2019 json!({
2020 "path": "/workspace/outputs/call_123.stdout",
2021 "content": "second page\n".repeat(200),
2022 "total_lines": 400,
2023 "lines_shown": {"start": 201, "end": 400},
2024 "truncated": false,
2025 "content_hash": "sha256:same-output"
2026 }),
2027 ));
2028
2029 let config = CompactionConfig::from_json(&json!({
2030 "cost_control": {
2031 "keep_recent_tool_results": 1,
2032 "mask_after_tool_results": 2
2033 }
2034 }));
2035 let result = build_model_view_messages(&messages, &config, None);
2036
2037 assert_eq!(result.masked_count, 1);
2038 let bash_result = result.messages[2].tool_result_content().unwrap();
2039 assert_eq!(bash_result.result.as_ref().unwrap()["masked"], true);
2040 let first_page = result.messages[4].tool_result_content().unwrap();
2041 assert!(first_page.result.as_ref().unwrap().get("content").is_some());
2042 let second_page = result.messages[6].tool_result_content().unwrap();
2043 assert!(
2044 second_page
2045 .result
2046 .as_ref()
2047 .unwrap()
2048 .get("content")
2049 .is_some()
2050 );
2051 }
2052
2053 #[test]
2054 fn test_model_view_masks_with_compaction_config() {
2055 let mut messages = vec![Message::user("inspect files repeatedly")];
2056 for index in 0..9 {
2057 messages.extend(make_message_tool_turn(
2058 &format!("call_{index}"),
2059 "read_file",
2060 json!({
2061 "path": "/workspace/session_019e4c9dd1b17021af70ad3227361b16.jsonl",
2062 "content": format!("{}{}", "large transcript line\n".repeat(1000), index),
2063 "total_lines": 1000,
2064 "lines_shown": {"start": 1, "end": 1000},
2065 "truncated": false,
2066 "content_hash": format!("sha256:{index}")
2067 }),
2068 ));
2069 }
2070
2071 let config = CompactionConfig::default();
2072 let result = build_model_view_messages(&messages, &config, None);
2073
2074 assert_eq!(result.masked_count, 7);
2075 assert!(result.tool_result_bytes_after < result.tool_result_bytes_before / 4);
2076 let first_tool = result.messages[2].tool_result_content().unwrap();
2077 let masked = first_tool.result.as_ref().unwrap();
2078 assert_eq!(masked["masked"], true);
2079 assert!(masked["summary"].as_str().unwrap().contains("read_file"));
2080 let last_tool = result
2081 .messages
2082 .last()
2083 .unwrap()
2084 .tool_result_content()
2085 .unwrap();
2086 assert!(last_tool.result.as_ref().unwrap().get("content").is_some());
2087 }
2088
2089 #[test]
2090 fn test_compaction_capability_contributes_model_view_provider() {
2091 let mut messages = vec![Message::user("inspect files repeatedly")];
2092 for index in 0..9 {
2093 messages.extend(make_message_tool_turn(
2094 &format!("call_{index}"),
2095 "read_file",
2096 json!({
2097 "path": "/workspace/src/lib.rs",
2098 "content": format!("{}{}", "large file line\n".repeat(1000), index),
2099 "total_lines": 1000,
2100 "lines_shown": {"start": 1, "end": 1000},
2101 "truncated": false
2102 }),
2103 ));
2104 }
2105
2106 let capability = CompactionCapability;
2107 let provider = capability.model_view_provider().unwrap();
2108 let context = ModelViewContext {
2109 session_id: crate::typed_id::SessionId::new(),
2110 prior_usage: None,
2111 };
2112 let result = provider.apply_model_view(messages, &json!({}), &context);
2113
2114 let first_tool = result[2].tool_result_content().unwrap();
2115 assert_eq!(first_tool.result.as_ref().unwrap()["masked"], true);
2116 let last_tool = result.last().unwrap().tool_result_content().unwrap();
2117 assert!(last_tool.result.as_ref().unwrap().get("content").is_some());
2118 }
2119
2120 #[test]
2121 fn test_model_view_respects_disabled_cost_control_config() {
2122 let mut messages = vec![Message::user("inspect files repeatedly")];
2123 for index in 0..5 {
2124 messages.extend(make_message_tool_turn(
2125 &format!("call_{index}"),
2126 "read_file",
2127 json!({
2128 "path": "/workspace/src/lib.rs",
2129 "content": "line\n".repeat(400),
2130 "total_lines": 400,
2131 "lines_shown": {"start": 1, "end": 400},
2132 "truncated": false
2133 }),
2134 ));
2135 }
2136
2137 let config = CompactionConfig::from_json(&json!({
2138 "cost_control": {
2139 "enabled": false,
2140 "keep_recent_tool_results": 1,
2141 "mask_after_tool_results": 2
2142 }
2143 }));
2144 let result = build_model_view_messages(&messages, &config, None);
2145
2146 assert_eq!(result.masked_count, 0);
2147 assert_eq!(
2148 result.tool_result_bytes_after,
2149 result.tool_result_bytes_before
2150 );
2151 }
2152
2153 #[test]
2154 fn test_cost_control_uses_prior_usage_signal() {
2155 let mut messages = vec![Message::user("run commands")];
2156 for index in 0..3 {
2157 messages.extend(make_message_tool_turn(
2158 &format!("call_{index}"),
2159 "bash",
2160 json!({
2161 "stdout": "small output",
2162 "stderr": "",
2163 "exit_code": 0,
2164 "success": true
2165 }),
2166 ));
2167 }
2168
2169 let config = CompactionConfig::from_json(&json!({
2170 "cost_control": {
2171 "keep_recent_tool_results": 1,
2172 "mask_after_tool_results": 99,
2173 "max_live_tool_result_bytes": 999999,
2174 "max_uncached_input_tokens": 1000
2175 }
2176 }));
2177 let usage = TokenUsage::with_cache(10_000, 100, Some(0), None);
2178 let result = apply_cost_control_masking(&messages, &config, Some(&usage));
2179
2180 assert_eq!(result.masked_count, 2);
2181 let first_tool = result.messages[2].tool_result_content().unwrap();
2182 let summary = first_tool.result.as_ref().unwrap()["summary"]
2183 .as_str()
2184 .unwrap();
2185 assert!(summary.contains("bash exit=0"));
2186 }
2187
2188 #[test]
2189 fn test_cost_control_masking_uses_disjoint_cache_buckets() {
2190 let config = CostControlConfig {
2193 mask_after_tool_results: usize::MAX,
2194 max_live_tool_result_bytes: usize::MAX,
2195 max_uncached_input_tokens: 50_000,
2196 min_cache_read_ratio: 0.35,
2197 ..CostControlConfig::default()
2198 };
2199
2200 let well_cached = TokenUsage::with_cache(20_000, 100, Some(180_000), None);
2203 assert!(!should_apply_cost_control_masking(
2204 0,
2205 0,
2206 &config,
2207 Some(&well_cached)
2208 ));
2209
2210 let cache_poor = TokenUsage::with_cache(20_000, 100, Some(2_000), None);
2212 assert!(should_apply_cost_control_masking(
2213 0,
2214 0,
2215 &config,
2216 Some(&cache_poor)
2217 ));
2218
2219 let heavy_uncached = TokenUsage::with_cache(60_000, 100, Some(200_000), None);
2221 assert!(should_apply_cost_control_masking(
2222 0,
2223 0,
2224 &config,
2225 Some(&heavy_uncached)
2226 ));
2227 }
2228
2229 #[test]
2230 fn test_model_view_uses_provider_cache_signal_from_compaction_config() {
2231 let mut messages = vec![Message::user("run commands")];
2232 for index in 0..3 {
2233 messages.extend(make_message_tool_turn(
2234 &format!("call_{index}"),
2235 "bash",
2236 json!({
2237 "stdout": "small output",
2238 "stderr": "",
2239 "exit_code": 0,
2240 "success": true
2241 }),
2242 ));
2243 }
2244 let usage = TokenUsage::with_cache(150_000, 100, Some(0), None);
2245
2246 let config = CompactionConfig::default();
2247 let result = build_model_view_messages(&messages, &config, Some(&usage));
2248
2249 assert_eq!(result.masked_count, 1);
2250 let first_tool = result.messages[2].tool_result_content().unwrap();
2251 assert_eq!(first_tool.result.as_ref().unwrap()["masked"], true);
2252 }
2253
2254 #[test]
2255 fn test_config_falls_back_to_defaults_for_invalid_json() {
2256 let config = CompactionConfig::from_json(&json!({
2257 "strategy": "nonexistent_strategy",
2258 "budget_percent": "not-a-number"
2259 }));
2260 assert_eq!(config.strategy, CompactionStrategy::Auto);
2261 assert!(config.proactive);
2262 }
2263
2264 #[test]
2265 fn test_config_partial_override() {
2266 let config = CompactionConfig::from_json(&json!({
2267 "budget_percent": 0.7,
2268 "observation_masking": {
2269 "keep_recent_tool_outputs": 3
2270 }
2271 }));
2272 assert_eq!(config.strategy, CompactionStrategy::Auto);
2273 assert!(config.proactive);
2274 assert!((config.budget_percent - 0.7).abs() < f32::EPSILON);
2275 assert_eq!(config.observation_masking.keep_recent_tool_outputs, 3);
2276 assert_eq!(
2277 config.observation_masking.summary_format,
2278 MaskingSummaryFormat::OneLine
2279 );
2280 }
2281
2282 #[test]
2283 fn test_strategy_serialization_roundtrip() {
2284 for strategy in [
2285 CompactionStrategy::Auto,
2286 CompactionStrategy::Native,
2287 CompactionStrategy::ObservationMasking,
2288 CompactionStrategy::Summarization,
2289 ] {
2290 let json = serde_json::to_value(strategy).unwrap();
2291 let deserialized: CompactionStrategy = serde_json::from_value(json).unwrap();
2292 assert_eq!(strategy, deserialized);
2293 }
2294 }
2295
2296 #[test]
2297 fn test_strategy_display() {
2298 assert_eq!(CompactionStrategy::Auto.to_string(), "auto");
2299 assert_eq!(CompactionStrategy::Native.to_string(), "native");
2300 assert_eq!(
2301 CompactionStrategy::ObservationMasking.to_string(),
2302 "observation_masking"
2303 );
2304 assert_eq!(
2305 CompactionStrategy::Summarization.to_string(),
2306 "summarization"
2307 );
2308 }
2309
2310 #[test]
2311 fn test_masking_format_serialization_roundtrip() {
2312 for format in [
2313 MaskingSummaryFormat::OneLine,
2314 MaskingSummaryFormat::HeadTail,
2315 ] {
2316 let json = serde_json::to_value(format).unwrap();
2317 let deserialized: MaskingSummaryFormat = serde_json::from_value(json).unwrap();
2318 assert_eq!(format, deserialized);
2319 }
2320 }
2321
2322 #[test]
2323 fn test_budget_percent_boundary_values() {
2324 let config = CompactionConfig::from_json(&json!({"budget_percent": 0.1}));
2325 assert!((config.budget_percent - 0.1).abs() < f32::EPSILON);
2326
2327 let config = CompactionConfig::from_json(&json!({"budget_percent": 0.99}));
2328 assert!((config.budget_percent - 0.99).abs() < f32::EPSILON);
2329 }
2330
2331 #[test]
2332 fn test_keep_recent_tool_outputs_zero() {
2333 let config = CompactionConfig::from_json(&json!({
2334 "observation_masking": {"keep_recent_tool_outputs": 0}
2335 }));
2336 assert_eq!(config.observation_masking.keep_recent_tool_outputs, 0);
2337 }
2338
2339 #[test]
2344 fn test_masking_no_tool_messages() {
2345 let messages = vec![make_user_msg("hello"), make_assistant_msg("hi")];
2346 let config = ObservationMaskingConfig::default();
2347 let result = apply_observation_masking(&messages, &config);
2348 assert_eq!(result.masked_count, 0);
2349 assert_eq!(result.messages.len(), 2);
2350 }
2351
2352 #[test]
2353 fn test_masking_fewer_than_keep_recent() {
2354 let messages = vec![
2355 make_user_msg("read file"),
2356 make_assistant_with_tool_call("call_1", "read_file"),
2357 make_tool_result("call_1", "file contents"),
2358 make_assistant_msg("done"),
2359 ];
2360 let config = ObservationMaskingConfig {
2361 keep_recent_tool_outputs: 5,
2362 summary_format: MaskingSummaryFormat::OneLine,
2363 };
2364 let result = apply_observation_masking(&messages, &config);
2365 assert_eq!(result.masked_count, 0);
2366 }
2367
2368 #[test]
2369 fn test_masking_masks_old_outputs() {
2370 let messages = vec![
2371 make_user_msg("start"),
2372 make_assistant_with_tool_call("call_1", "read_file"),
2373 make_tool_result(
2374 "call_1",
2375 "old file contents that are very long and should be masked by the observation masking strategy because it exceeds 100 chars",
2376 ),
2377 make_assistant_msg("got it"),
2378 make_user_msg("next"),
2379 make_assistant_with_tool_call("call_2", "search"),
2380 make_tool_result("call_2", "search results"),
2381 make_assistant_msg("found it"),
2382 make_user_msg("more"),
2383 make_assistant_with_tool_call("call_3", "bash"),
2384 make_tool_result("call_3", "command output"),
2385 ];
2386
2387 let config = ObservationMaskingConfig {
2388 keep_recent_tool_outputs: 2,
2389 summary_format: MaskingSummaryFormat::OneLine,
2390 };
2391 let result = apply_observation_masking(&messages, &config);
2392
2393 assert_eq!(result.masked_count, 1);
2394
2395 let masked = &result.messages[2];
2397 assert_eq!(masked.role, LlmMessageRole::Tool);
2398 let text = extract_text(&masked.content);
2399 assert!(
2400 text.starts_with('['),
2401 "Expected masked summary, got: {text}"
2402 );
2403 assert!(text.contains("read_file"), "Expected tool name: {text}");
2404
2405 assert_eq!(extract_text(&result.messages[6].content), "search results");
2407 assert_eq!(extract_text(&result.messages[10].content), "command output");
2408 }
2409
2410 #[test]
2411 fn test_masking_preserves_tool_call_id() {
2412 let messages = vec![
2413 make_assistant_with_tool_call("call_1", "read_file"),
2414 make_tool_result("call_1", "content"),
2415 make_assistant_with_tool_call("call_2", "bash"),
2416 make_tool_result("call_2", "output"),
2417 ];
2418
2419 let config = ObservationMaskingConfig {
2420 keep_recent_tool_outputs: 1,
2421 summary_format: MaskingSummaryFormat::OneLine,
2422 };
2423 let result = apply_observation_masking(&messages, &config);
2424 assert_eq!(result.messages[1].tool_call_id, Some("call_1".to_string()));
2425 }
2426
2427 #[test]
2428 fn test_masking_head_tail_format() {
2429 let long_output = (0..20)
2430 .map(|i| format!("line {i}"))
2431 .collect::<Vec<_>>()
2432 .join("\n");
2433
2434 let messages = vec![
2435 make_assistant_with_tool_call("call_1", "bash"),
2436 make_tool_result("call_1", &long_output),
2437 make_assistant_with_tool_call("call_2", "bash"),
2438 make_tool_result("call_2", "recent output"),
2439 ];
2440
2441 let config = ObservationMaskingConfig {
2442 keep_recent_tool_outputs: 1,
2443 summary_format: MaskingSummaryFormat::HeadTail,
2444 };
2445 let result = apply_observation_masking(&messages, &config);
2446
2447 let text = extract_text(&result.messages[1].content);
2448 assert!(text.contains("line 0"), "Should contain first lines");
2449 assert!(text.contains("line 19"), "Should contain last lines");
2450 assert!(text.contains("lines omitted"), "Should indicate omissions");
2451 }
2452
2453 #[test]
2454 fn test_masking_short_output_inline() {
2455 let messages = vec![
2456 make_assistant_with_tool_call("call_1", "get_time"),
2457 make_tool_result("call_1", "2024-01-01"),
2458 make_assistant_with_tool_call("call_2", "bash"),
2459 make_tool_result("call_2", "ok"),
2460 ];
2461
2462 let config = ObservationMaskingConfig {
2463 keep_recent_tool_outputs: 1,
2464 summary_format: MaskingSummaryFormat::OneLine,
2465 };
2466 let result = apply_observation_masking(&messages, &config);
2467 let text = extract_text(&result.messages[1].content);
2468 assert!(text.contains("2024-01-01"), "Short output included: {text}");
2469 }
2470
2471 #[test]
2472 fn test_masking_all_when_keep_zero() {
2473 let messages = vec![
2474 make_assistant_with_tool_call("call_1", "a"),
2475 make_tool_result("call_1", "output1"),
2476 make_assistant_with_tool_call("call_2", "b"),
2477 make_tool_result("call_2", "output2"),
2478 ];
2479
2480 let config = ObservationMaskingConfig {
2481 keep_recent_tool_outputs: 0,
2482 summary_format: MaskingSummaryFormat::OneLine,
2483 };
2484 let result = apply_observation_masking(&messages, &config);
2485 assert_eq!(result.masked_count, 2);
2486 }
2487
2488 #[test]
2489 fn test_masking_empty_messages() {
2490 let result = apply_observation_masking(&[], &ObservationMaskingConfig::default());
2491 assert_eq!(result.masked_count, 0);
2492 assert!(result.messages.is_empty());
2493 }
2494
2495 #[test]
2496 fn test_masking_preserves_message_count() {
2497 let messages = vec![
2498 make_user_msg("start"),
2499 make_assistant_with_tool_call("c1", "read_file"),
2500 make_tool_result("c1", "content 1"),
2501 make_assistant_msg("ok"),
2502 make_user_msg("next"),
2503 make_assistant_with_tool_call("c2", "bash"),
2504 make_tool_result("c2", "content 2"),
2505 make_assistant_msg("done"),
2506 ];
2507
2508 let config = ObservationMaskingConfig {
2509 keep_recent_tool_outputs: 1,
2510 summary_format: MaskingSummaryFormat::OneLine,
2511 };
2512 let result = apply_observation_masking(&messages, &config);
2513 assert_eq!(result.messages.len(), messages.len());
2514 }
2515
2516 #[test]
2517 fn test_masking_unknown_tool_call_id() {
2518 let messages = vec![
2519 make_tool_result("orphan", "some output"),
2520 make_assistant_with_tool_call("call_2", "bash"),
2521 make_tool_result("call_2", "recent"),
2522 ];
2523
2524 let config = ObservationMaskingConfig {
2525 keep_recent_tool_outputs: 1,
2526 summary_format: MaskingSummaryFormat::OneLine,
2527 };
2528 let result = apply_observation_masking(&messages, &config);
2529 assert_eq!(result.masked_count, 1);
2530 let text = extract_text(&result.messages[0].content);
2531 assert!(text.contains("unknown_tool"), "Fallback name: {text}");
2532 }
2533
2534 #[test]
2535 fn test_masking_many_tool_calls_keeps_exactly_n() {
2536 let mut messages = Vec::new();
2537 for i in 0..10 {
2538 let id = format!("call_{i}");
2539 messages.push(make_assistant_with_tool_call(&id, &format!("tool_{i}")));
2540 messages.push(make_tool_result(&id, &format!("output {i}")));
2541 }
2542
2543 let config = ObservationMaskingConfig {
2544 keep_recent_tool_outputs: 3,
2545 summary_format: MaskingSummaryFormat::OneLine,
2546 };
2547 let result = apply_observation_masking(&messages, &config);
2548 assert_eq!(result.masked_count, 7);
2549
2550 assert_eq!(extract_text(&result.messages[15].content), "output 7");
2552 assert_eq!(extract_text(&result.messages[17].content), "output 8");
2553 assert_eq!(extract_text(&result.messages[19].content), "output 9");
2554 }
2555
2556 #[test]
2561 fn test_summarization_prompt_default() {
2562 let config = SummarizationConfig::default();
2563 let prompt = build_summarization_prompt(&config);
2564 assert!(prompt.contains("<task>"));
2565 assert!(prompt.contains("decisions"));
2566 assert!(prompt.contains("files_modified"));
2567 assert!(prompt.contains("errors"));
2568 assert!(prompt.contains("current_plan"));
2569 }
2570
2571 #[test]
2572 fn test_summarization_prompt_custom_instructions() {
2573 let config = SummarizationConfig {
2574 instructions: Some("Focus on API changes".to_string()),
2575 ..Default::default()
2576 };
2577 let prompt = build_summarization_prompt(&config);
2578 assert!(prompt.contains("Focus on API changes"));
2579 }
2580
2581 #[test]
2582 fn test_summarization_prompt_custom_preserve() {
2583 let config = SummarizationConfig {
2584 preserve: vec!["auth_tokens".to_string(), "database_schema".to_string()],
2585 ..Default::default()
2586 };
2587 let prompt = build_summarization_prompt(&config);
2588 assert!(prompt.contains("auth_tokens"));
2589 assert!(prompt.contains("database_schema"));
2590 assert!(!prompt.contains("decisions"));
2591 }
2592
2593 #[test]
2594 fn test_summarization_prompt_empty_preserve_uses_defaults() {
2595 let config = SummarizationConfig {
2596 preserve: vec![],
2597 ..Default::default()
2598 };
2599 let prompt = build_summarization_prompt(&config);
2600 assert!(prompt.contains("decisions"));
2601 }
2602
2603 #[test]
2604 fn test_format_messages_for_summarization() {
2605 let messages = vec![
2606 make_user_msg("What is 2+2?"),
2607 make_assistant_msg("The answer is 4."),
2608 ];
2609 let formatted = format_messages_for_summarization(&messages);
2610 assert!(formatted.contains("[user]: What is 2+2?"));
2611 assert!(formatted.contains("[assistant]: The answer is 4."));
2612 }
2613
2614 #[test]
2615 fn test_format_messages_truncates_long_content() {
2616 let long_content = "x".repeat(5000);
2617 let messages = vec![make_user_msg(&long_content)];
2618 let formatted = format_messages_for_summarization(&messages);
2619 assert!(formatted.contains("truncated"));
2620 assert!(formatted.len() < long_content.len());
2621 }
2622
2623 #[test]
2624 fn test_format_messages_truncates_utf8_without_panic() {
2625 let multibyte = "é".repeat(1001); let messages = vec![make_user_msg(&multibyte)];
2627 let formatted = format_messages_for_summarization(&messages);
2628 assert!(formatted.contains("truncated"));
2629 assert!(formatted.contains("[truncated, 2002 chars total]"));
2630 }
2631
2632 #[test]
2633 fn test_build_summary_message() {
2634 let msg = build_summary_message("The user asked about APIs.");
2635 assert_eq!(msg.role, LlmMessageRole::System);
2636 let text = extract_text(&msg.content);
2637 assert!(text.contains("[CONVERSATION_SUMMARY]"));
2638 assert!(text.contains("The user asked about APIs."));
2639 assert!(text.contains("[/CONVERSATION_SUMMARY]"));
2640 }
2641
2642 #[test]
2647 fn test_head_tail_short_content_unchanged() {
2648 let content = LlmMessageContent::Text("line1\nline2\nline3".to_string());
2649 assert_eq!(format_head_tail_summary(&content), "line1\nline2\nline3");
2650 }
2651
2652 #[test]
2653 fn test_head_tail_exactly_six_lines() {
2654 let content = LlmMessageContent::Text("1\n2\n3\n4\n5\n6".to_string());
2655 assert_eq!(format_head_tail_summary(&content), "1\n2\n3\n4\n5\n6");
2656 }
2657
2658 #[test]
2659 fn test_head_tail_seven_lines() {
2660 let content = LlmMessageContent::Text("1\n2\n3\n4\n5\n6\n7".to_string());
2661 let result = format_head_tail_summary(&content);
2662 assert!(result.contains("1\n2\n3"));
2663 assert!(result.contains("5\n6\n7"));
2664 assert!(result.contains("1 lines omitted"));
2665 }
2666
2667 #[test]
2672 fn test_one_line_empty_output() {
2673 let result = format_one_line_summary("bash", &LlmMessageContent::Text(String::new()));
2674 assert_eq!(result, "[bash → ]");
2675 }
2676
2677 #[test]
2678 fn test_one_line_exactly_100_chars() {
2679 let text = "x".repeat(100);
2680 let result = format_one_line_summary("bash", &LlmMessageContent::Text(text.clone()));
2681 assert!(result.contains(&text));
2682 }
2683
2684 #[test]
2685 fn test_one_line_101_chars_summarized() {
2686 let text = "x".repeat(101);
2687 let result = format_one_line_summary("bash", &LlmMessageContent::Text(text));
2688 assert!(result.contains("lines"));
2689 assert!(result.contains("bytes"));
2690 }
2691
2692 #[test]
2693 fn test_one_line_multipart_content() {
2694 let content = LlmMessageContent::Parts(vec![
2695 LlmContentPart::Text {
2696 text: "part1".to_string(),
2697 },
2698 LlmContentPart::Text {
2699 text: "part2".to_string(),
2700 },
2701 ]);
2702 let result = format_one_line_summary("tool", &content);
2703 assert!(result.contains("part1"));
2704 assert!(result.contains("part2"));
2705 }
2706
2707 #[test]
2712 fn test_compaction_step_serialization() {
2713 let step = CompactionStep {
2714 strategy: "observation_masking".to_string(),
2715 messages_after: 42,
2716 duration_ms: 12,
2717 };
2718 let json = serde_json::to_value(&step).unwrap();
2719 assert_eq!(json["strategy"], "observation_masking");
2720 assert_eq!(json["messages_after"], 42);
2721 assert_eq!(json["duration_ms"], 12);
2722 }
2723
2724 #[test]
2729 fn test_estimate_tokens_text() {
2730 let msg = make_user_msg("hello world"); let tokens = estimate_tokens(&msg);
2732 assert_eq!(tokens, 11 / 4);
2733 }
2734
2735 #[test]
2736 fn test_estimate_tokens_empty() {
2737 let msg = make_user_msg("");
2738 assert_eq!(estimate_tokens(&msg), 0);
2739 }
2740
2741 #[test]
2742 fn test_estimate_total_tokens() {
2743 let messages = vec![
2744 make_user_msg("a".repeat(400).as_str()), make_assistant_msg("b".repeat(200).as_str()), ];
2747 assert_eq!(estimate_total_tokens(&messages), 150);
2748 }
2749
2750 #[test]
2751 fn test_estimate_tokens_with_tool_calls() {
2752 let msg = make_assistant_with_tool_call("call_1", "read_file");
2753 let tokens = estimate_tokens(&msg);
2754 assert!(tokens > 0, "Tool call should contribute tokens");
2755 }
2756
2757 #[test]
2762 fn test_should_compact_proactively_under_budget() {
2763 let messages = vec![make_user_msg("short")];
2764 let config = CompactionConfig::default(); assert!(!should_compact_proactively(&messages, &config, 128_000));
2766 }
2767
2768 #[test]
2769 fn test_should_compact_proactively_over_budget() {
2770 let big_text = "x".repeat(4000); let messages = vec![make_user_msg(&big_text)];
2773 let config = CompactionConfig::default();
2774 assert!(should_compact_proactively(&messages, &config, 1000));
2775 }
2776
2777 #[test]
2778 fn test_should_compact_proactively_disabled() {
2779 let big_text = "x".repeat(4000);
2780 let messages = vec![make_user_msg(&big_text)];
2781 let config = CompactionConfig {
2782 proactive: false,
2783 ..Default::default()
2784 };
2785 assert!(!should_compact_proactively(&messages, &config, 1000));
2786 }
2787
2788 #[test]
2793 fn test_aggressive_trim_keeps_newest() {
2794 let messages = vec![
2796 make_user_msg(&"s".repeat(400)), make_user_msg(&"a".repeat(400)), make_assistant_msg(&"b".repeat(400)), make_user_msg(&"c".repeat(400)), make_assistant_msg(&"d".repeat(400)), ];
2802 let target_tokens = 300;
2804 let result = aggressive_trim(&messages, target_tokens, true);
2805 assert!(
2806 result.len() < messages.len(),
2807 "Expected trim, got {} messages",
2808 result.len()
2809 );
2810 assert_eq!(result[0].role, LlmMessageRole::User);
2812 }
2813
2814 #[test]
2815 fn test_aggressive_trim_empty() {
2816 let result = aggressive_trim(&[], 100, false);
2817 assert!(result.is_empty());
2818 }
2819
2820 #[test]
2821 fn test_aggressive_trim_anchors_first_conversation_message() {
2822 let messages = vec![
2826 make_user_msg("sys"), make_user_msg(&"TASK ".repeat(80)), make_assistant_msg(&"x".repeat(400)), make_user_msg(&"c".repeat(400)), make_assistant_msg(&"d".repeat(400)), ];
2832 let result = aggressive_trim(&messages, 250, true);
2833
2834 assert!(
2835 result.len() < messages.len(),
2836 "expected a trim, got {} messages",
2837 result.len()
2838 );
2839 let kept_task = result.iter().any(|m| match &m.content {
2840 LlmMessageContent::Text(t) => t.contains("TASK"),
2841 _ => false,
2842 });
2843 assert!(kept_task, "the original task must be anchored, not dropped");
2844 }
2845
2846 #[test]
2847 fn test_aggressive_trim_everything_fits() {
2848 let messages = vec![make_user_msg("hi"), make_assistant_msg("hello")];
2849 let result = aggressive_trim(&messages, 100_000, false);
2850 assert_eq!(result.len(), 2);
2851 }
2852
2853 #[test]
2858 fn test_session_metrics_record() {
2859 let mut metrics = SessionCompactionMetrics::default();
2860 metrics.record("observation_masking+native", 100, 50, 200);
2861
2862 assert_eq!(metrics.compaction_count, 1);
2863 assert_eq!(metrics.total_messages_saved, 50);
2864 assert_eq!(metrics.total_duration_ms, 200);
2865 assert_eq!(metrics.strategy_counts["observation_masking"], 1);
2866 assert_eq!(metrics.strategy_counts["native"], 1);
2867 }
2868
2869 #[test]
2870 fn test_session_metrics_accumulate() {
2871 let mut metrics = SessionCompactionMetrics::default();
2872 metrics.record("observation_masking", 100, 80, 10);
2873 metrics.record("summarization", 80, 40, 500);
2874
2875 assert_eq!(metrics.compaction_count, 2);
2876 assert_eq!(metrics.total_messages_saved, 60);
2877 assert_eq!(metrics.total_duration_ms, 510);
2878 assert_eq!(metrics.strategy_counts["observation_masking"], 1);
2879 assert_eq!(metrics.strategy_counts["summarization"], 1);
2880 }
2881
2882 #[test]
2883 fn test_session_metrics_serialization() {
2884 let mut metrics = SessionCompactionMetrics::default();
2885 metrics.record("auto", 50, 30, 100);
2886 let json = serde_json::to_value(&metrics).unwrap();
2887 assert_eq!(json["compaction_count"], 1);
2888 assert_eq!(json["total_messages_saved"], 20);
2889 }
2890
2891 #[test]
2896 fn test_classify_memory_tiers_basic() {
2897 let messages: Vec<LlmMessage> = (0..30)
2898 .map(|i| make_user_msg(&format!("msg {i}")))
2899 .collect();
2900
2901 let config = HierarchicalMemoryConfig {
2902 hot_messages: 5,
2903 warm_messages: 10,
2904 };
2905
2906 let classified = classify_memory_tiers(&messages, &config);
2907 assert_eq!(classified.len(), 30);
2908
2909 assert_eq!(classified[29].0, MemoryTier::Hot);
2911 assert_eq!(classified[25].0, MemoryTier::Hot);
2912
2913 assert_eq!(classified[24].0, MemoryTier::Warm);
2915 assert_eq!(classified[15].0, MemoryTier::Warm);
2916
2917 assert_eq!(classified[14].0, MemoryTier::Cold);
2919 assert_eq!(classified[0].0, MemoryTier::Cold);
2920 }
2921
2922 #[test]
2923 fn test_classify_memory_tiers_all_hot() {
2924 let messages: Vec<LlmMessage> =
2925 (0..3).map(|i| make_user_msg(&format!("msg {i}"))).collect();
2926
2927 let config = HierarchicalMemoryConfig::default(); let classified = classify_memory_tiers(&messages, &config);
2930 assert!(classified.iter().all(|(tier, _)| *tier == MemoryTier::Hot));
2931 }
2932
2933 #[test]
2934 fn test_apply_hierarchical_memory_basic() {
2935 let mut messages = Vec::new();
2936
2937 for i in 0..5 {
2939 let id = format!("old_{i}");
2940 messages.push(make_assistant_with_tool_call(&id, "read_file"));
2941 messages.push(make_tool_result(&id, &format!("old content {i}")));
2942 }
2943
2944 for i in 0..3 {
2946 let id = format!("mid_{i}");
2947 messages.push(make_assistant_with_tool_call(&id, "bash"));
2948 messages.push(make_tool_result(&id, &format!("mid output {i}")));
2949 }
2950
2951 messages.push(make_user_msg("what now?"));
2953 messages.push(make_assistant_msg("let me check"));
2954
2955 let config = HierarchicalMemoryConfig {
2956 hot_messages: 2,
2957 warm_messages: 6,
2958 };
2959 let masking_config = ObservationMaskingConfig::default();
2960
2961 let result = apply_hierarchical_memory(
2962 &messages,
2963 &config,
2964 &masking_config,
2965 Some("Summary of old work"),
2966 );
2967
2968 assert!(result.len() <= 9);
2970 let first_text = extract_text(&result[0].content);
2972 assert!(first_text.contains("CONVERSATION_SUMMARY"));
2973 let last = extract_text(&result[result.len() - 1].content);
2975 assert!(last.contains("let me check"));
2976 }
2977
2978 #[test]
2979 fn test_apply_hierarchical_memory_no_cold() {
2980 let messages = vec![make_user_msg("hello"), make_assistant_msg("hi")];
2981
2982 let config = HierarchicalMemoryConfig {
2983 hot_messages: 5,
2984 warm_messages: 5,
2985 };
2986
2987 let result = apply_hierarchical_memory(
2988 &messages,
2989 &config,
2990 &ObservationMaskingConfig::default(),
2991 None,
2992 );
2993 assert_eq!(result.len(), 2);
2995 }
2996
2997 #[test]
2998 fn test_memory_tier_config_from_json() {
2999 let config: HierarchicalMemoryConfig = serde_json::from_value(json!({
3000 "hot_messages": 10,
3001 "warm_messages": 50
3002 }))
3003 .unwrap();
3004 assert_eq!(config.hot_messages, 10);
3005 assert_eq!(config.warm_messages, 50);
3006 }
3007
3008 #[test]
3009 fn test_memory_tier_config_defaults() {
3010 let config = HierarchicalMemoryConfig::default();
3011 assert_eq!(config.hot_messages, 20);
3012 assert_eq!(config.warm_messages, 100);
3013 }
3014
3015 #[test]
3016 fn test_compaction_config_with_memory_tiers() {
3017 let config = CompactionConfig::from_json(&json!({
3018 "strategy": "auto",
3019 "memory_tiers": {
3020 "hot_messages": 15,
3021 "warm_messages": 80
3022 }
3023 }));
3024 assert_eq!(config.memory_tiers.hot_messages, 15);
3025 assert_eq!(config.memory_tiers.warm_messages, 80);
3026 }
3027
3028 #[test]
3029 fn test_memory_tier_serialization() {
3030 assert_eq!(serde_json::to_value(MemoryTier::Hot).unwrap(), json!("hot"));
3031 assert_eq!(
3032 serde_json::to_value(MemoryTier::Warm).unwrap(),
3033 json!("warm")
3034 );
3035 assert_eq!(
3036 serde_json::to_value(MemoryTier::Cold).unwrap(),
3037 json!("cold")
3038 );
3039 }
3040
3041 #[test]
3046 fn test_masking_skips_activate_skill_results() {
3047 let messages = vec![
3050 make_assistant_with_tool_call("call_skill", "activate_skill"),
3051 make_tool_result(
3052 "call_skill",
3053 "You are a code review agent. Follow these instructions...",
3054 ),
3055 make_assistant_msg("Skill activated"),
3056 make_assistant_with_tool_call("call_read", "read_file"),
3057 make_tool_result(
3058 "call_read",
3059 "file contents that are long enough to be masked by observation masking because they exceed one hundred characters easily",
3060 ),
3061 make_assistant_msg("got it"),
3062 make_assistant_with_tool_call("call_bash", "bash"),
3063 make_tool_result("call_bash", "command output"),
3064 ];
3065
3066 let config = ObservationMaskingConfig {
3067 keep_recent_tool_outputs: 1,
3068 summary_format: MaskingSummaryFormat::OneLine,
3069 };
3070 let result = apply_observation_masking(&messages, &config);
3071
3072 assert_eq!(
3074 extract_text(&result.messages[1].content),
3075 "You are a code review agent. Follow these instructions..."
3076 );
3077 assert!(extract_text(&result.messages[4].content).starts_with('['));
3079 assert_eq!(extract_text(&result.messages[7].content), "command output");
3081 assert_eq!(result.masked_count, 1);
3082 }
3083
3084 #[test]
3085 fn test_masking_all_activate_skill_exempt_from_count() {
3086 let messages = vec![
3089 make_assistant_with_tool_call("s1", "activate_skill"),
3090 make_tool_result("s1", "Skill 1 instructions"),
3091 make_assistant_with_tool_call("s2", "activate_skill"),
3092 make_tool_result("s2", "Skill 2 instructions"),
3093 make_assistant_with_tool_call("c1", "bash"),
3094 make_tool_result("c1", "output"),
3095 ];
3096
3097 let config = ObservationMaskingConfig {
3098 keep_recent_tool_outputs: 0,
3099 summary_format: MaskingSummaryFormat::OneLine,
3100 };
3101 let result = apply_observation_masking(&messages, &config);
3102
3103 assert_eq!(result.masked_count, 1);
3104 assert_eq!(
3106 extract_text(&result.messages[1].content),
3107 "Skill 1 instructions"
3108 );
3109 assert_eq!(
3110 extract_text(&result.messages[3].content),
3111 "Skill 2 instructions"
3112 );
3113 assert_complete_tool_exchanges(&result.messages);
3114 }
3115
3116 #[test]
3117 fn test_aggressive_trim_preserves_skill_messages() {
3118 let messages = vec![
3121 make_user_msg(&"s".repeat(400)), make_assistant_with_tool_call("skill1", "activate_skill"),
3123 make_tool_result("skill1", "Important skill instructions"),
3124 make_user_msg(&"a".repeat(400)), make_assistant_msg(&"b".repeat(400)), make_user_msg(&"c".repeat(400)), make_assistant_msg(&"d".repeat(400)), ];
3129
3130 let target_tokens = 400;
3133 let result = aggressive_trim(&messages, target_tokens, true);
3134
3135 let has_skill_result = result.iter().any(|m| {
3137 m.role == LlmMessageRole::Tool
3138 && extract_text(&m.content) == "Important skill instructions"
3139 });
3140 assert!(
3141 has_skill_result,
3142 "Skill tool result must survive aggressive trim"
3143 );
3144
3145 let has_skill_call = result.iter().any(|m| {
3146 m.tool_calls
3147 .as_ref()
3148 .is_some_and(|calls| calls.iter().any(|tc| tc.name == "activate_skill"))
3149 });
3150 assert!(
3151 has_skill_call,
3152 "Skill tool call must survive aggressive trim"
3153 );
3154 }
3155
3156 #[test]
3157 fn aggressive_trim_keeps_parallel_tool_calls_atomic() {
3158 let messages = vec![
3159 make_user_msg("system"),
3160 make_user_msg("original task"),
3161 make_assistant_with_tool_calls(&[
3162 ("call_skill", "activate_skill"),
3163 ("call_bash", "bash"),
3164 ]),
3165 make_tool_result("call_skill", "Important skill instructions"),
3166 make_tool_result("call_bash", &"output ".repeat(100)),
3167 make_user_msg(&"recent user ".repeat(40)),
3168 make_assistant_msg(&"recent answer ".repeat(40)),
3169 ];
3170
3171 let result = aggressive_trim(&messages, 300, true);
3172 let visible_calls: Vec<_> = result
3173 .iter()
3174 .flat_map(|message| message.tool_calls.iter().flatten())
3175 .map(|call| call.id.as_str())
3176 .collect();
3177 let visible_results: Vec<_> = result
3178 .iter()
3179 .filter_map(|message| message.tool_call_id.as_deref())
3180 .collect();
3181
3182 assert!(visible_calls.contains(&"call_skill"));
3183 assert!(visible_results.contains(&"call_skill"));
3184 assert_eq!(
3185 visible_calls, visible_results,
3186 "reduction must not expose a call without its result"
3187 );
3188 assert_complete_tool_exchanges(&result);
3189 }
3190
3191 #[test]
3192 fn hierarchical_memory_prunes_only_evicted_calls_from_a_protected_parallel_batch() {
3193 let messages = vec![
3194 make_assistant_with_tool_calls(&[
3195 ("call_skill", "activate_skill"),
3196 ("call_bash", "bash"),
3197 ]),
3198 make_tool_result("call_skill", "Important skill instructions"),
3199 make_tool_result("call_bash", "ordinary output"),
3200 make_user_msg("recent question"),
3201 make_assistant_msg("recent answer"),
3202 ];
3203 let result = apply_hierarchical_memory(
3204 &messages,
3205 &HierarchicalMemoryConfig {
3206 hot_messages: 2,
3207 warm_messages: 0,
3208 },
3209 &ObservationMaskingConfig::default(),
3210 Some("Summary of old work"),
3211 );
3212
3213 assert_complete_tool_exchanges(&result);
3214 let calls: Vec<_> = result
3215 .iter()
3216 .flat_map(|message| message.tool_calls.iter().flatten())
3217 .map(|call| call.id.as_str())
3218 .collect();
3219 assert_eq!(calls, vec!["call_skill"]);
3220 }
3221
3222 #[test]
3223 fn summary_boundary_drops_a_recent_result_whose_call_was_summarized() {
3224 let recent = vec![
3225 make_tool_result("call_old", "old result"),
3226 make_user_msg("continue"),
3227 ];
3228 let result = compose_summary_with_recent(None, "Earlier work", &recent);
3229
3230 assert_complete_tool_exchanges(&result);
3231 assert!(result.iter().all(|message| message.tool_call_id.is_none()));
3232 assert!(
3233 result
3234 .iter()
3235 .any(|message| extract_text(&message.content) == "continue")
3236 );
3237 }
3238
3239 #[test]
3240 fn test_hierarchical_memory_rescues_skill_from_cold_tier() {
3241 let mut messages = Vec::new();
3242
3243 messages.push(make_assistant_with_tool_call("skill1", "activate_skill"));
3245 messages.push(make_tool_result(
3246 "skill1",
3247 "You must always validate input.",
3248 ));
3249 for i in 0..8 {
3250 let id = format!("old_{i}");
3251 messages.push(make_assistant_with_tool_call(&id, "read_file"));
3252 messages.push(make_tool_result(&id, &format!("old content {i}")));
3253 }
3254
3255 for i in 0..3 {
3257 let id = format!("mid_{i}");
3258 messages.push(make_assistant_with_tool_call(&id, "bash"));
3259 messages.push(make_tool_result(&id, &format!("mid output {i}")));
3260 }
3261
3262 messages.push(make_user_msg("what now?"));
3264 messages.push(make_assistant_msg("let me check"));
3265
3266 let config = HierarchicalMemoryConfig {
3267 hot_messages: 2,
3268 warm_messages: 6,
3269 };
3270 let masking_config = ObservationMaskingConfig::default();
3271
3272 let result = apply_hierarchical_memory(
3273 &messages,
3274 &config,
3275 &masking_config,
3276 Some("Summary of old work"),
3277 );
3278
3279 let has_skill_instructions = result
3281 .iter()
3282 .any(|m| extract_text(&m.content).contains("You must always validate input."));
3283 assert!(
3284 has_skill_instructions,
3285 "Skill instructions from cold tier must be rescued into output"
3286 );
3287
3288 assert!(extract_text(&result[0].content).contains("CONVERSATION_SUMMARY"));
3290 }
3291
3292 #[test]
3293 fn test_is_protected_tool_result_detection() {
3294 let messages = vec![
3295 make_assistant_with_tool_call("s1", "activate_skill"),
3296 make_tool_result("s1", "skill content"),
3297 make_assistant_with_tool_call("r1", "read_file"),
3298 make_tool_result("r1", "file content"),
3299 ];
3300
3301 assert!(is_protected_tool_result(&messages, &messages[1]));
3303 assert!(!is_protected_tool_result(&messages, &messages[3]));
3305 assert!(!is_protected_tool_result(&messages, &messages[0]));
3307 }
3308
3309 #[test]
3310 fn test_is_protected_tool_call_message_detection() {
3311 let skill_call = make_assistant_with_tool_call("s1", "activate_skill");
3312 let regular_call = make_assistant_with_tool_call("r1", "read_file");
3313 let user_msg = make_user_msg("hello");
3314
3315 assert!(is_protected_tool_call_message(&skill_call));
3316 assert!(!is_protected_tool_call_message(®ular_call));
3317 assert!(!is_protected_tool_call_message(&user_msg));
3318 }
3319
3320 #[test]
3321 fn test_default_preserve_includes_skill_instructions() {
3322 let config = SummarizationConfig::default();
3323 assert!(
3324 config.preserve.contains(&"skill_instructions".to_string()),
3325 "Default preserve list must include skill_instructions"
3326 );
3327 }
3328
3329 #[test]
3330 fn test_summarization_prompt_mentions_skill_protection() {
3331 let config = SummarizationConfig::default();
3332 let prompt = build_summarization_prompt(&config);
3333 assert!(
3334 prompt.contains("activate_skill"),
3335 "Summarization prompt must instruct LLM to preserve skill content"
3336 );
3337 }
3338
3339 #[test]
3340 fn test_aggressive_trim_protected_exceed_budget() {
3341 let messages = vec![
3344 make_user_msg(&"s".repeat(400)), make_assistant_with_tool_call("skill1", "activate_skill"), make_tool_result("skill1", &"x".repeat(800)), make_assistant_with_tool_call("skill2", "activate_skill"), make_tool_result("skill2", &"y".repeat(800)), make_user_msg(&"z".repeat(400)), ];
3351
3352 let result = aggressive_trim(&messages, 200, true);
3354
3355 let has_non_protected = result
3357 .iter()
3358 .any(|m| m.role == LlmMessageRole::User && extract_text(&m.content).contains('z'));
3359 assert!(
3360 !has_non_protected,
3361 "Non-protected messages must be dropped when protected exceed budget"
3362 );
3363 }
3364
3365 #[test]
3366 fn test_format_messages_no_truncate_protected_tool_result() {
3367 let long_instructions = "a".repeat(5000);
3369 let messages = vec![
3370 make_assistant_with_tool_call("s1", "activate_skill"),
3371 make_tool_result("s1", &long_instructions),
3372 make_assistant_with_tool_call("r1", "read_file"),
3373 make_tool_result("r1", &"b".repeat(5000)),
3374 ];
3375
3376 let formatted = format_messages_for_summarization(&messages);
3377
3378 assert!(
3380 formatted.contains(&long_instructions),
3381 "Protected tool result must not be truncated"
3382 );
3383 assert!(
3385 formatted.contains("[truncated, 5000 chars total]"),
3386 "Non-protected tool result should be truncated"
3387 );
3388 }
3389
3390 #[test]
3391 fn test_hierarchical_memory_cross_tier_boundary_protection() {
3392 let mut messages = Vec::new();
3395
3396 messages.push(make_assistant_with_tool_call("skill1", "activate_skill"));
3398 for i in 0..9 {
3399 let id = format!("cold_{i}");
3400 messages.push(make_assistant_with_tool_call(&id, "read_file"));
3401 messages.push(make_tool_result(&id, &format!("cold content {i}")));
3402 }
3403
3404 messages.push(make_tool_result(
3406 "skill1",
3407 "Cross-tier skill instructions that must survive",
3408 ));
3409 for i in 0..2 {
3410 let id = format!("warm_{i}");
3411 messages.push(make_assistant_with_tool_call(&id, "bash"));
3412 messages.push(make_tool_result(&id, &format!("warm output {i}")));
3413 }
3414
3415 messages.push(make_user_msg("continue"));
3417 messages.push(make_assistant_msg("ok"));
3418
3419 let config = HierarchicalMemoryConfig {
3420 hot_messages: 2,
3421 warm_messages: 5, };
3423 let masking_config = ObservationMaskingConfig {
3424 keep_recent_tool_outputs: 0,
3425 summary_format: MaskingSummaryFormat::OneLine,
3426 };
3427
3428 let result = apply_hierarchical_memory(&messages, &config, &masking_config, None);
3429
3430 let has_skill_instructions = result.iter().any(|m| {
3431 extract_text(&m.content).contains("Cross-tier skill instructions that must survive")
3432 });
3433 assert!(
3434 has_skill_instructions,
3435 "Skill result in warm tier with call in cold tier must be protected"
3436 );
3437 }
3438}