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