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