Skip to main content

ai_agents_memory/
compacting.rs

1//! CompactingMemory implementation with auto-summarization
2
3use std::sync::Arc;
4
5use async_trait::async_trait;
6use parking_lot::RwLock;
7use serde::{Deserialize, Serialize};
8use tokio::sync::Mutex as AsyncMutex;
9
10use ai_agents_core::{ChatMessage, MemorySnapshot, Result};
11
12use super::Memory;
13use super::context::{CompressResult, ConversationContext, estimate_tokens};
14use super::native::{NativeRetentionInspection, readable_projection};
15use super::summarizer::Summarizer;
16
17fn prefix_at_char_boundary(text: &str, max_chars: usize) -> &str {
18    if max_chars == 0 {
19        return "";
20    }
21
22    match text.char_indices().nth(max_chars) {
23        Some((idx, _)) => &text[..idx],
24        None => text,
25    }
26}
27
28pub struct CompactingMemory {
29    operation_lock: AsyncMutex<()>,
30    summary: RwLock<Option<String>>,
31    messages: RwLock<Vec<ChatMessage>>,
32    summarized_count: RwLock<usize>,
33    config: CompactingMemoryConfig,
34    summarizer: Arc<dyn Summarizer>,
35    compression_history: RwLock<Vec<CompressionEvent>>,
36}
37
38#[derive(Debug, Clone, Serialize, Deserialize)]
39#[serde(deny_unknown_fields)]
40pub struct CompactingMemoryConfig {
41    /// Maximum recent messages retained verbatim, clamped below the compression threshold.
42    #[serde(default = "default_max_recent_messages")]
43    pub max_recent_messages: usize,
44
45    #[serde(default = "default_compress_threshold")]
46    pub compress_threshold: usize,
47
48    #[serde(default = "default_summarize_batch_size")]
49    pub summarize_batch_size: usize,
50
51    // FIXME: unlimited length as default value?
52    #[serde(default = "default_max_summary_length")]
53    pub max_summary_length: usize,
54}
55
56#[derive(Debug, Clone, Serialize, Deserialize)]
57pub struct CompressionEvent {
58    pub timestamp: chrono::DateTime<chrono::Utc>,
59    pub messages_compressed: usize,
60    pub summary_length_before: usize,
61    pub summary_length_after: usize,
62}
63
64fn default_max_recent_messages() -> usize {
65    50
66}
67
68fn default_compress_threshold() -> usize {
69    30
70}
71
72fn default_summarize_batch_size() -> usize {
73    10
74}
75
76fn default_max_summary_length() -> usize {
77    2000
78}
79
80fn protected_recent_count(config: &CompactingMemoryConfig, message_count: usize) -> usize {
81    if config.max_recent_messages < config.compress_threshold {
82        return config.max_recent_messages.min(message_count);
83    }
84
85    let batch_at_threshold = config
86        .summarize_batch_size
87        .max(1)
88        .min(config.compress_threshold);
89    let retention_cap = config.compress_threshold.saturating_sub(batch_at_threshold);
90    config
91        .max_recent_messages
92        .min(retention_cap)
93        .min(message_count)
94}
95
96impl Default for CompactingMemoryConfig {
97    fn default() -> Self {
98        Self {
99            max_recent_messages: default_max_recent_messages(),
100            compress_threshold: default_compress_threshold(),
101            summarize_batch_size: default_summarize_batch_size(),
102            max_summary_length: default_max_summary_length(),
103        }
104    }
105}
106
107impl CompactingMemory {
108    pub fn new(summarizer: Arc<dyn Summarizer>, config: CompactingMemoryConfig) -> Self {
109        Self {
110            operation_lock: AsyncMutex::new(()),
111            summary: RwLock::new(None),
112            messages: RwLock::new(Vec::new()),
113            summarized_count: RwLock::new(0),
114            config,
115            summarizer,
116            compression_history: RwLock::new(Vec::new()),
117        }
118    }
119
120    pub fn with_default_config(summarizer: Arc<dyn Summarizer>) -> Self {
121        Self::new(summarizer, CompactingMemoryConfig::default())
122    }
123
124    pub fn config(&self) -> &CompactingMemoryConfig {
125        &self.config
126    }
127
128    pub fn summary(&self) -> Option<String> {
129        self.summary.read().clone()
130    }
131
132    pub fn summarized_count(&self) -> usize {
133        *self.summarized_count.read()
134    }
135
136    pub fn compression_history(&self) -> Vec<CompressionEvent> {
137        self.compression_history.read().clone()
138    }
139
140    fn record_compression(&self, messages_compressed: usize, before: usize, after: usize) {
141        let event = CompressionEvent {
142            timestamp: chrono::Utc::now(),
143            messages_compressed,
144            summary_length_before: before,
145            summary_length_after: after,
146        };
147        self.compression_history.write().push(event);
148    }
149}
150
151#[async_trait]
152impl ai_agents_core::Memory for CompactingMemory {
153    async fn add_message(&self, message: ChatMessage) -> Result<()> {
154        let _operation = self.operation_lock.lock().await;
155        self.messages.write().push(message);
156        Ok(())
157    }
158
159    async fn get_messages(&self, limit: Option<usize>) -> Result<Vec<ChatMessage>> {
160        let messages = self.messages.read();
161        match limit {
162            Some(n) => {
163                let start = messages.len().saturating_sub(n);
164                Ok(messages[start..].to_vec())
165            }
166            None => Ok(messages.clone()),
167        }
168    }
169
170    async fn clear(&self) -> Result<()> {
171        let _operation = self.operation_lock.lock().await;
172        *self.summary.write() = None;
173        self.messages.write().clear();
174        *self.summarized_count.write() = 0;
175        self.compression_history.write().clear();
176        Ok(())
177    }
178
179    fn len(&self) -> usize {
180        self.messages.read().len()
181    }
182
183    async fn snapshot(&self) -> Result<MemorySnapshot> {
184        let _operation = self.operation_lock.lock().await;
185        let messages = self.messages.read().clone();
186        let summary = self.summary.read().clone();
187        let summarized_count = *self.summarized_count.read();
188
189        let mut snapshot = MemorySnapshot::new(messages).with_summarized_count(summarized_count);
190        if let Some(s) = summary {
191            snapshot = snapshot.with_summary(s);
192        }
193        Ok(snapshot)
194    }
195
196    async fn restore(&self, snapshot: MemorySnapshot) -> Result<()> {
197        let _operation = self.operation_lock.lock().await;
198        *self.messages.write() = snapshot.messages;
199        *self.summary.write() = snapshot.summary;
200        *self.summarized_count.write() = snapshot.summarized_count;
201        self.compression_history.write().clear();
202        Ok(())
203    }
204
205    async fn evict_oldest(&self, count: usize) -> Result<Vec<ChatMessage>> {
206        let _operation = self.operation_lock.lock().await;
207        let mut messages = self.messages.write();
208        let requested = count.min(messages.len());
209        let inspection = NativeRetentionInspection::inspect(&messages)?;
210        let evict_count = if requested == 0 {
211            0
212        } else {
213            inspection
214                .safe_prefix_len_between(requested, messages.len())
215                .ok_or_else(|| {
216                    ai_agents_core::AgentError::MemoryError(
217                        "eviction would split the protected signed native exchange".to_string(),
218                    )
219                })?
220        };
221        let evicted: Vec<ChatMessage> = messages.drain(..evict_count).collect();
222        Ok(evicted)
223    }
224}
225
226#[async_trait]
227impl Memory for CompactingMemory {
228    async fn get_context(&self) -> Result<ConversationContext> {
229        let _operation = self.operation_lock.lock().await;
230        let messages = self.messages.read().clone();
231        let summary = self.summary.read().clone();
232        let summarized_count = *self.summarized_count.read();
233        let total_messages = messages.len() + summarized_count;
234
235        let mut ctx = ConversationContext::with_messages(messages);
236        ctx.total_messages = total_messages;
237
238        if let Some(s) = summary {
239            ctx = ctx.with_summary(s, summarized_count);
240        }
241
242        Ok(ctx)
243    }
244
245    async fn compress(&self, summarizer: Option<&dyn Summarizer>) -> Result<CompressResult> {
246        let _operation = self.operation_lock.lock().await;
247        let message_count = self.messages.read().len();
248
249        if message_count == 0 || message_count < self.config.compress_threshold {
250            return Ok(CompressResult::NotNeeded);
251        }
252
253        let summarizer = summarizer.unwrap_or(self.summarizer.as_ref());
254        let protected_count = protected_recent_count(&self.config, message_count);
255        let configured_compressible_count = message_count - protected_count;
256        let inspection = {
257            let messages = self.messages.read();
258            NativeRetentionInspection::inspect(&messages)?
259        };
260        let compressible_count = inspection
261            .protected_suffix_start()
262            .map_or(configured_compressible_count, |start| {
263                configured_compressible_count.min(start)
264            });
265        if compressible_count == 0 {
266            return Ok(CompressResult::NotNeeded);
267        }
268        let requested_batch = self
269            .config
270            .summarize_batch_size
271            .max(1)
272            .min(compressible_count);
273        // Prefer finishing the group containing the configured batch boundary. If that would
274        // enter the retained suffix, summarize the largest earlier complete group instead.
275        let batch_size = inspection
276            .safe_prefix_len_between(requested_batch, compressible_count)
277            .or_else(|| {
278                let earlier = inspection.safe_prefix_len_at_most(requested_batch);
279                (earlier > 0).then_some(earlier)
280            })
281            .unwrap_or(0);
282        if batch_size == 0 {
283            return Ok(CompressResult::NotNeeded);
284        }
285
286        let original_messages_to_summarize: Vec<ChatMessage> = {
287            let messages = self.messages.read();
288            messages[..batch_size].to_vec()
289        };
290        // Projection happens before the dynamic summarizer boundary so custom implementations
291        // cannot observe or accidentally persist provider replay state.
292        let mut messages_to_summarize = readable_projection(&original_messages_to_summarize)?;
293        let incomplete_exchanges = inspection
294            .incomplete_exchanges()
295            .into_iter()
296            .filter(|(message_index, _)| *message_index < batch_size)
297            .collect::<Vec<_>>();
298        for (message_index, missing_result_ids) in &incomplete_exchanges {
299            let status = serde_json::json!({
300                "native_exchange_status": "incomplete",
301                "missing_result_ids": missing_result_ids,
302                "description": "final execution result was not recorded"
303            });
304            messages_to_summarize[*message_index].content.push('\n');
305            messages_to_summarize[*message_index]
306                .content
307                .push_str(&status.to_string());
308        }
309
310        let new_summary = summarizer.summarize(&messages_to_summarize).await?;
311
312        let summary_before_len = self.summary.read().as_ref().map(|s| s.len()).unwrap_or(0);
313
314        let existing_summary = self.summary.read().clone();
315        let existing_summary_tokens = existing_summary
316            .as_deref()
317            .map(estimate_tokens)
318            .unwrap_or(0);
319        let combined_summary = match existing_summary {
320            Some(existing) => summarizer.merge_summaries(&[existing, new_summary]).await?,
321            None => new_summary,
322        };
323
324        // The framework, rather than a fallible/custom summarizer, owns the durable uncertainty
325        // marker so completed compression cannot erase evidence that a tool result was missing.
326        let incomplete_suffix = incomplete_exchanges
327            .iter()
328            .map(|(_, missing_result_ids)| {
329                serde_json::json!({
330                    "native_exchange_status": "incomplete",
331                    "missing_result_ids": missing_result_ids,
332                    "description": "final execution result was not recorded"
333                })
334                .to_string()
335            })
336            .collect::<Vec<_>>()
337            .join("\n");
338        let suffix = (!incomplete_suffix.is_empty()).then(|| format!("\n{incomplete_suffix}"));
339        let suffix_chars = suffix
340            .as_deref()
341            .map(|suffix| suffix.chars().count())
342            .unwrap_or(0);
343        if suffix_chars > self.config.max_summary_length {
344            return Err(ai_agents_core::AgentError::MemoryError(
345                "summary limit cannot retain incomplete native exchange status".to_string(),
346            ));
347        }
348        let content_limit = self.config.max_summary_length - suffix_chars;
349        let truncated = prefix_at_char_boundary(&combined_summary, content_limit);
350        let mut final_summary = if truncated.len() < combined_summary.len() {
351            truncated.to_string()
352        } else {
353            combined_summary
354        };
355        if let Some(suffix) = suffix {
356            final_summary.push_str(&suffix);
357        }
358
359        let summary_after_len = final_summary.len();
360
361        {
362            let mut messages = self.messages.write();
363            messages.drain(..batch_size);
364        }
365
366        *self.summary.write() = Some(final_summary.clone());
367        *self.summarized_count.write() += batch_size;
368
369        self.record_compression(batch_size, summary_before_len, summary_after_len);
370
371        let tokens_before: u32 = existing_summary_tokens.saturating_add(
372            original_messages_to_summarize
373                .iter()
374                .map(|m| estimate_tokens(&m.content))
375                .sum(),
376        );
377        let tokens_after = estimate_tokens(&final_summary);
378        let tokens_saved = tokens_before.saturating_sub(tokens_after);
379
380        Ok(CompressResult::Compressed {
381            messages_summarized: batch_size,
382            new_summary_length: summary_after_len,
383            tokens_saved,
384        })
385    }
386
387    fn needs_compression(&self) -> bool {
388        let message_count = self.messages.read().len();
389        message_count > 0 && message_count >= self.config.compress_threshold
390    }
391}
392
393#[cfg(test)]
394mod tests {
395    use super::*;
396
397    use crate::summarizer::NoopSummarizer;
398    use ai_agents_core::{
399        AgentError, Memory as CoreMemory, NativeCallBinding, NativeProviderState,
400        NativeProviderTarget, Role, ToolCall, encode_native_tool_call_markers,
401        encode_native_tool_result_marker,
402    };
403    use tokio::time::{Duration, timeout};
404
405    fn make_message(content: &str) -> ChatMessage {
406        ChatMessage {
407            role: Role::User,
408            content: content.to_string(),
409            name: None,
410            timestamp: None,
411        }
412    }
413
414    fn signed_turn(exchange_id: &str) -> Vec<ChatMessage> {
415        let call = ToolCall {
416            id: format!("{exchange_id}-call"),
417            name: "lookup".to_string(),
418            arguments: serde_json::json!({"query":"fixture"}),
419        };
420        let state = NativeProviderState::new(
421            exchange_id,
422            "google",
423            "generateContent",
424            NativeProviderTarget::new("https://example.invalid/v1beta/", "fixture-model").unwrap(),
425            serde_json::json!({
426                "role":"model",
427                "parts":[{
428                    "functionCall":{"name":"lookup","args":{"query":"fixture"}},
429                    "thoughtSignature":"fixture-signature"
430                }]
431            }),
432            vec![NativeCallBinding::new(&call.id, 0).unwrap()],
433        )
434        .unwrap();
435        vec![
436            ChatMessage::user("signed user request"),
437            ChatMessage::assistant(
438                encode_native_tool_call_markers(std::slice::from_ref(&call), Some(&state)).unwrap(),
439            ),
440            ChatMessage::function(
441                "lookup",
442                encode_native_tool_result_marker(&call, serde_json::json!({"ok":true})).unwrap(),
443            ),
444            ChatMessage::assistant("signed turn final response"),
445        ]
446    }
447
448    fn message_contents(messages: &[ChatMessage]) -> Vec<&str> {
449        messages
450            .iter()
451            .map(|message| message.content.as_str())
452            .collect()
453    }
454
455    struct BlockingSummarizer {
456        started: tokio::sync::Notify,
457        release: tokio::sync::Notify,
458        batches: RwLock<Vec<Vec<String>>>,
459    }
460
461    impl BlockingSummarizer {
462        fn new() -> Self {
463            Self {
464                started: tokio::sync::Notify::new(),
465                release: tokio::sync::Notify::new(),
466                batches: RwLock::new(Vec::new()),
467            }
468        }
469
470        async fn wait_until_started(&self) {
471            self.started.notified().await;
472        }
473
474        fn release(&self) {
475            self.release.notify_one();
476        }
477
478        fn batches(&self) -> Vec<Vec<String>> {
479            self.batches.read().clone()
480        }
481    }
482
483    #[async_trait]
484    impl Summarizer for BlockingSummarizer {
485        async fn summarize(&self, messages: &[ChatMessage]) -> Result<String> {
486            let contents: Vec<_> = messages
487                .iter()
488                .map(|message| message.content.clone())
489                .collect();
490            self.batches.write().push(contents.clone());
491            self.started.notify_one();
492            self.release.notified().await;
493            Ok(contents.join(" | "))
494        }
495    }
496
497    struct FailingSummarizer;
498
499    #[async_trait]
500    impl Summarizer for FailingSummarizer {
501        async fn summarize(&self, _messages: &[ChatMessage]) -> Result<String> {
502            Err(AgentError::MemoryError("summary failed".to_string()))
503        }
504    }
505
506    struct FailingMergeSummarizer;
507
508    #[async_trait]
509    impl Summarizer for FailingMergeSummarizer {
510        async fn summarize(&self, _messages: &[ChatMessage]) -> Result<String> {
511            Ok("new summary".to_string())
512        }
513
514        async fn merge_summaries(&self, _summaries: &[String]) -> Result<String> {
515            Err(AgentError::MemoryError("merge failed".to_string()))
516        }
517    }
518
519    #[derive(Default)]
520    struct CapturingSummarizer {
521        batches: RwLock<Vec<Vec<ChatMessage>>>,
522    }
523
524    #[async_trait]
525    impl Summarizer for CapturingSummarizer {
526        async fn summarize(&self, messages: &[ChatMessage]) -> Result<String> {
527            self.batches.write().push(messages.to_vec());
528            Ok(messages
529                .iter()
530                .map(|message| message.content.clone())
531                .collect::<Vec<_>>()
532                .join(" | "))
533        }
534    }
535
536    struct DroppingStatusSummarizer;
537
538    #[async_trait]
539    impl Summarizer for DroppingStatusSummarizer {
540        async fn summarize(&self, _messages: &[ChatMessage]) -> Result<String> {
541            Ok("summary that omitted the native status".to_string())
542        }
543    }
544
545    fn create_test_memory() -> CompactingMemory {
546        let summarizer = Arc::new(NoopSummarizer);
547        let config = CompactingMemoryConfig {
548            max_recent_messages: 3,
549            compress_threshold: 5,
550            summarize_batch_size: 3,
551            max_summary_length: 500,
552        };
553        CompactingMemory::new(summarizer, config)
554    }
555
556    #[tokio::test]
557    async fn test_basic_add_and_get() {
558        let memory = create_test_memory();
559
560        memory.add_message(make_message("Hello")).await.unwrap();
561        memory.add_message(make_message("World")).await.unwrap();
562
563        let messages = memory.get_messages(None).await.unwrap();
564        assert_eq!(messages.len(), 2);
565        assert_eq!(messages[0].content, "Hello");
566        assert_eq!(messages[1].content, "World");
567    }
568
569    #[tokio::test]
570    async fn test_get_messages_with_limit() {
571        let memory = create_test_memory();
572
573        for i in 0..5 {
574            memory
575                .add_message(make_message(&format!("msg{}", i)))
576                .await
577                .unwrap();
578        }
579
580        let messages = memory.get_messages(Some(2)).await.unwrap();
581        assert_eq!(messages.len(), 2);
582        assert_eq!(messages[0].content, "msg3");
583        assert_eq!(messages[1].content, "msg4");
584    }
585
586    #[tokio::test]
587    async fn test_clear() {
588        let memory = create_test_memory();
589
590        memory.add_message(make_message("test")).await.unwrap();
591        assert!(!memory.is_empty());
592
593        memory.clear().await.unwrap();
594        assert!(memory.is_empty());
595        assert!(memory.summary().is_none());
596    }
597
598    #[tokio::test]
599    async fn test_needs_compression() {
600        let memory = create_test_memory();
601
602        for i in 0..4 {
603            memory
604                .add_message(make_message(&format!("msg{}", i)))
605                .await
606                .unwrap();
607        }
608        assert!(!memory.needs_compression());
609
610        memory.add_message(make_message("msg4")).await.unwrap();
611        assert!(memory.needs_compression());
612    }
613
614    #[tokio::test]
615    async fn test_compress_not_needed() {
616        let memory = create_test_memory();
617
618        memory.add_message(make_message("msg1")).await.unwrap();
619        memory.add_message(make_message("msg2")).await.unwrap();
620
621        let result = memory.compress(None).await.unwrap();
622        assert!(matches!(result, CompressResult::NotNeeded));
623    }
624
625    #[tokio::test]
626    async fn test_compress_when_needed() {
627        let memory = create_test_memory();
628
629        for i in 0..6 {
630            memory
631                .add_message(make_message(&format!("message number {}", i)))
632                .await
633                .unwrap();
634        }
635
636        assert!(memory.needs_compression());
637
638        let result = memory.compress(None).await.unwrap();
639
640        if let CompressResult::Compressed {
641            messages_summarized,
642            ..
643        } = result
644        {
645            assert_eq!(messages_summarized, 3);
646        } else {
647            panic!("Expected Compressed result");
648        }
649
650        assert_eq!(memory.len(), 3);
651        assert!(memory.summary().is_some());
652        assert_eq!(memory.summarized_count(), 3);
653    }
654
655    #[tokio::test]
656    async fn test_compress_preserves_configured_recent_tail() {
657        let config = CompactingMemoryConfig {
658            max_recent_messages: 3,
659            compress_threshold: 5,
660            summarize_batch_size: 10,
661            max_summary_length: 500,
662        };
663        let memory = CompactingMemory::new(Arc::new(NoopSummarizer), config);
664
665        for i in 0..7 {
666            memory
667                .add_message(make_message(&format!("msg{}", i)))
668                .await
669                .unwrap();
670        }
671
672        let result = memory.compress(None).await.unwrap();
673        assert!(matches!(
674            result,
675            CompressResult::Compressed {
676                messages_summarized: 4,
677                ..
678            }
679        ));
680        let remaining = memory.get_messages(None).await.unwrap();
681        let contents: Vec<_> = remaining
682            .iter()
683            .map(|message| message.content.as_str())
684            .collect();
685        assert_eq!(contents, vec!["msg4", "msg5", "msg6"]);
686    }
687
688    #[tokio::test]
689    async fn test_multilingual_recent_tail_remains_verbatim() {
690        let config = CompactingMemoryConfig {
691            max_recent_messages: 4,
692            compress_threshold: 6,
693            summarize_batch_size: 3,
694            max_summary_length: 100_000,
695        };
696        let memory = CompactingMemory::new(Arc::new(NoopSummarizer), config);
697        let large_mixed_message = "protected 한국어 日本語 العربية English emoji 🧭".to_string();
698        let source = vec![
699            "Old English context".to_string(),
700            "이전 한국어 문맥".to_string(),
701            "古い日本語の文脈".to_string(),
702            "سياق عربي قديم".to_string(),
703            "पुराना हिंदी संदर्भ".to_string(),
704            "Keep 한국어 and English together".to_string(),
705            "保留する日本語メッセージ".to_string(),
706            "احتفظ بهذه الرسالة العربية".to_string(),
707            large_mixed_message,
708        ];
709        let expected_tail = source[source.len() - 4..].to_vec();
710
711        for content in &source {
712            memory.add_message(make_message(content)).await.unwrap();
713        }
714
715        memory.compress(None).await.unwrap();
716        memory.compress(None).await.unwrap();
717        let remaining = memory.get_messages(None).await.unwrap();
718
719        assert!(!memory.needs_compression());
720        assert_eq!(
721            remaining
722                .iter()
723                .map(|message| message.content.clone())
724                .collect::<Vec<_>>(),
725            expected_tail
726        );
727        assert_eq!(memory.summarized_count() + memory.len(), source.len());
728    }
729
730    #[tokio::test]
731    async fn test_compress_clamps_recent_tail_for_configured_batch() {
732        let config = CompactingMemoryConfig {
733            max_recent_messages: 100,
734            compress_threshold: 5,
735            summarize_batch_size: 10,
736            max_summary_length: 500,
737        };
738        let memory = CompactingMemory::new(Arc::new(NoopSummarizer), config);
739
740        for i in 0..5 {
741            memory
742                .add_message(make_message(&format!("msg{}", i)))
743                .await
744                .unwrap();
745        }
746
747        let result = memory.compress(None).await.unwrap();
748        assert!(matches!(
749            result,
750            CompressResult::Compressed {
751                messages_summarized: 5,
752                ..
753            }
754        ));
755        assert!(memory.get_messages(None).await.unwrap().is_empty());
756        assert!(!memory.needs_compression());
757    }
758
759    #[tokio::test]
760    async fn test_default_config_compresses_full_batches_at_steady_state() {
761        let memory =
762            CompactingMemory::new(Arc::new(NoopSummarizer), CompactingMemoryConfig::default());
763
764        for i in 0..30 {
765            memory
766                .add_message(make_message(&format!("msg{}", i)))
767                .await
768                .unwrap();
769        }
770
771        for round in 0..4 {
772            let result = memory.compress(None).await.unwrap();
773            assert!(matches!(
774                result,
775                CompressResult::Compressed {
776                    messages_summarized: 10,
777                    ..
778                }
779            ));
780            assert_eq!(memory.len(), 20);
781            assert!(!memory.needs_compression());
782
783            if round < 3 {
784                let start = 30 + round * 10;
785                for i in start..start + 10 {
786                    memory
787                        .add_message(make_message(&format!("msg{}", i)))
788                        .await
789                        .unwrap();
790                }
791                assert_eq!(memory.len(), 30);
792                assert!(memory.needs_compression());
793            }
794        }
795
796        assert_eq!(memory.summarized_count(), 40);
797        let remaining = memory.get_messages(None).await.unwrap();
798        assert_eq!(remaining.first().unwrap().content, "msg40");
799        assert_eq!(remaining.last().unwrap().content, "msg59");
800    }
801
802    #[test]
803    fn test_protected_recent_count_edge_cases() {
804        let non_conflicting = CompactingMemoryConfig {
805            max_recent_messages: 25,
806            compress_threshold: 30,
807            summarize_batch_size: 10,
808            max_summary_length: 500,
809        };
810        assert_eq!(protected_recent_count(&non_conflicting, 30), 25);
811
812        let conflicting = CompactingMemoryConfig {
813            max_recent_messages: 50,
814            compress_threshold: 30,
815            summarize_batch_size: 10,
816            max_summary_length: 500,
817        };
818        assert_eq!(protected_recent_count(&conflicting, 30), 20);
819
820        let oversized_batch = CompactingMemoryConfig {
821            max_recent_messages: 5,
822            compress_threshold: 5,
823            summarize_batch_size: 10,
824            max_summary_length: 500,
825        };
826        assert_eq!(protected_recent_count(&oversized_batch, 5), 0);
827
828        let zero_batch = CompactingMemoryConfig {
829            max_recent_messages: 5,
830            compress_threshold: 5,
831            summarize_batch_size: 0,
832            max_summary_length: 500,
833        };
834        assert_eq!(protected_recent_count(&zero_batch, 5), 4);
835
836        let zero_threshold = CompactingMemoryConfig {
837            max_recent_messages: 5,
838            compress_threshold: 0,
839            summarize_batch_size: 10,
840            max_summary_length: 500,
841        };
842        assert_eq!(protected_recent_count(&zero_threshold, 5), 0);
843    }
844
845    #[tokio::test]
846    async fn test_initial_summary_failure_rolls_back_all_accounting() {
847        let config = CompactingMemoryConfig {
848            max_recent_messages: 3,
849            compress_threshold: 5,
850            summarize_batch_size: 2,
851            max_summary_length: 100_000,
852        };
853        let memory = CompactingMemory::new(Arc::new(FailingSummarizer), config);
854        let source = vec![
855            "English before failure".to_string(),
856            "실패 전 한국어".to_string(),
857            "失敗前の日本語".to_string(),
858            "قبل الفشل".to_string(),
859            "large mixed message 한界🙂abc".to_string(),
860        ];
861        for content in &source {
862            memory.add_message(make_message(content)).await.unwrap();
863        }
864
865        let snapshot_before = memory.snapshot().await.unwrap();
866        let context_before = memory.get_context().await.unwrap();
867        let error = memory.compress(None).await.unwrap_err();
868        let snapshot_after = memory.snapshot().await.unwrap();
869        let context_after = memory.get_context().await.unwrap();
870
871        assert!(error.to_string().contains("summary failed"));
872        assert_eq!(
873            message_contents(&snapshot_after.messages),
874            message_contents(&snapshot_before.messages)
875        );
876        assert_eq!(snapshot_after.summary, snapshot_before.summary);
877        assert_eq!(
878            snapshot_after.summarized_count,
879            snapshot_before.summarized_count
880        );
881        assert_eq!(context_after.total_messages, context_before.total_messages);
882        assert_eq!(
883            context_after.summarized_count,
884            context_before.summarized_count
885        );
886        assert_eq!(
887            context_after.estimated_tokens(),
888            context_before.estimated_tokens()
889        );
890        assert!(memory.compression_history().is_empty());
891        assert!(memory.needs_compression());
892
893        let retry = memory.compress(Some(&NoopSummarizer)).await.unwrap();
894        assert!(matches!(
895            retry,
896            CompressResult::Compressed {
897                messages_summarized: 2,
898                ..
899            }
900        ));
901        assert_eq!(memory.summarized_count() + memory.len(), source.len());
902    }
903
904    #[tokio::test]
905    async fn test_compression_failure_is_non_destructive() {
906        let config = CompactingMemoryConfig {
907            max_recent_messages: 2,
908            compress_threshold: 5,
909            summarize_batch_size: 3,
910            max_summary_length: 500,
911        };
912        let memory = CompactingMemory::new(Arc::new(NoopSummarizer), config);
913
914        for i in 0..5 {
915            memory
916                .add_message(make_message(&format!("msg{}", i)))
917                .await
918                .unwrap();
919        }
920        memory.compress(None).await.unwrap();
921        for i in 5..8 {
922            memory
923                .add_message(make_message(&format!("msg{}", i)))
924                .await
925                .unwrap();
926        }
927
928        let messages_before = memory.get_messages(None).await.unwrap();
929        let summary_before = memory.summary();
930        let summarized_count_before = memory.summarized_count();
931        let history_len_before = memory.compression_history().len();
932
933        let error = memory.compress(Some(&FailingMergeSummarizer)).await;
934        assert!(error.is_err());
935        let messages_after = memory.get_messages(None).await.unwrap();
936        let before_contents: Vec<_> = messages_before
937            .iter()
938            .map(|message| message.content.as_str())
939            .collect();
940        let after_contents: Vec<_> = messages_after
941            .iter()
942            .map(|message| message.content.as_str())
943            .collect();
944        assert_eq!(after_contents, before_contents);
945        assert_eq!(memory.summary(), summary_before);
946        assert_eq!(memory.summarized_count(), summarized_count_before);
947        assert_eq!(memory.compression_history().len(), history_len_before);
948    }
949
950    #[tokio::test]
951    async fn test_concurrent_compressions_are_serialized() {
952        let summarizer = Arc::new(BlockingSummarizer::new());
953        let config = CompactingMemoryConfig {
954            max_recent_messages: 2,
955            compress_threshold: 5,
956            summarize_batch_size: 3,
957            max_summary_length: 500,
958        };
959        let memory = Arc::new(CompactingMemory::new(summarizer.clone(), config));
960        for i in 0..5 {
961            memory
962                .add_message(make_message(&format!("msg{}", i)))
963                .await
964                .unwrap();
965        }
966
967        let first_memory = memory.clone();
968        let first = tokio::spawn(async move { first_memory.compress(None).await });
969        summarizer.wait_until_started().await;
970
971        let second_memory = memory.clone();
972        let mut second = tokio::spawn(async move { second_memory.compress(None).await });
973        assert!(
974            timeout(Duration::from_millis(50), &mut second)
975                .await
976                .is_err()
977        );
978
979        summarizer.release();
980        assert!(matches!(
981            first.await.unwrap().unwrap(),
982            CompressResult::Compressed {
983                messages_summarized: 3,
984                ..
985            }
986        ));
987        assert!(matches!(
988            second.await.unwrap().unwrap(),
989            CompressResult::NotNeeded
990        ));
991        assert_eq!(summarizer.batches(), vec![vec!["msg0", "msg1", "msg2"]]);
992        let remaining = memory.get_messages(None).await.unwrap();
993        let contents: Vec<_> = remaining
994            .iter()
995            .map(|message| message.content.as_str())
996            .collect();
997        assert_eq!(contents, vec!["msg3", "msg4"]);
998    }
999
1000    #[tokio::test]
1001    async fn test_compression_serializes_add_message() {
1002        let summarizer = Arc::new(BlockingSummarizer::new());
1003        let config = CompactingMemoryConfig {
1004            max_recent_messages: 2,
1005            compress_threshold: 5,
1006            summarize_batch_size: 3,
1007            max_summary_length: 500,
1008        };
1009        let memory = Arc::new(CompactingMemory::new(summarizer.clone(), config));
1010        for i in 0..5 {
1011            memory
1012                .add_message(make_message(&format!("msg{}", i)))
1013                .await
1014                .unwrap();
1015        }
1016
1017        let compress_memory = memory.clone();
1018        let compress = tokio::spawn(async move { compress_memory.compress(None).await });
1019        summarizer.wait_until_started().await;
1020
1021        let add_memory = memory.clone();
1022        let mut add =
1023            tokio::spawn(async move { add_memory.add_message(make_message("msg5")).await });
1024        assert!(timeout(Duration::from_millis(50), &mut add).await.is_err());
1025
1026        summarizer.release();
1027        compress.await.unwrap().unwrap();
1028        add.await.unwrap().unwrap();
1029        assert_eq!(summarizer.batches(), vec![vec!["msg0", "msg1", "msg2"]]);
1030        let remaining = memory.get_messages(None).await.unwrap();
1031        let contents: Vec<_> = remaining
1032            .iter()
1033            .map(|message| message.content.as_str())
1034            .collect();
1035        assert_eq!(contents, vec!["msg3", "msg4", "msg5"]);
1036    }
1037
1038    #[tokio::test]
1039    async fn test_compression_serializes_eviction() {
1040        let summarizer = Arc::new(BlockingSummarizer::new());
1041        let config = CompactingMemoryConfig {
1042            max_recent_messages: 2,
1043            compress_threshold: 5,
1044            summarize_batch_size: 3,
1045            max_summary_length: 500,
1046        };
1047        let memory = Arc::new(CompactingMemory::new(summarizer.clone(), config));
1048        for i in 0..5 {
1049            memory
1050                .add_message(make_message(&format!("msg{}", i)))
1051                .await
1052                .unwrap();
1053        }
1054
1055        let compress_memory = memory.clone();
1056        let compress = tokio::spawn(async move { compress_memory.compress(None).await });
1057        summarizer.wait_until_started().await;
1058
1059        let evict_memory = memory.clone();
1060        let mut evict = tokio::spawn(async move { evict_memory.evict_oldest(1).await });
1061        assert!(
1062            timeout(Duration::from_millis(50), &mut evict)
1063                .await
1064                .is_err()
1065        );
1066
1067        summarizer.release();
1068        compress.await.unwrap().unwrap();
1069        let evicted = evict.await.unwrap().unwrap();
1070        assert_eq!(evicted.len(), 1);
1071        assert_eq!(evicted[0].content, "msg3");
1072        assert_eq!(summarizer.batches(), vec![vec!["msg0", "msg1", "msg2"]]);
1073        let remaining = memory.get_messages(None).await.unwrap();
1074        assert_eq!(remaining.len(), 1);
1075        assert_eq!(remaining[0].content, "msg4");
1076    }
1077
1078    #[tokio::test]
1079    async fn test_get_context() {
1080        let memory = create_test_memory();
1081
1082        for i in 0..6 {
1083            memory
1084                .add_message(make_message(&format!("msg{}", i)))
1085                .await
1086                .unwrap();
1087        }
1088
1089        memory.compress(None).await.unwrap();
1090
1091        let ctx = memory.get_context().await.unwrap();
1092        assert!(ctx.summary.is_some());
1093        assert_eq!(ctx.messages.len(), 3);
1094        assert_eq!(ctx.summarized_count, 3);
1095    }
1096
1097    #[tokio::test]
1098    async fn test_snapshot_restore() {
1099        let memory = create_test_memory();
1100
1101        memory.add_message(make_message("msg1")).await.unwrap();
1102        memory.add_message(make_message("msg2")).await.unwrap();
1103
1104        let snapshot = memory.snapshot().await.unwrap();
1105        assert_eq!(snapshot.messages.len(), 2);
1106
1107        memory.clear().await.unwrap();
1108        assert!(memory.is_empty());
1109
1110        memory.restore(snapshot).await.unwrap();
1111        let messages = memory.get_messages(None).await.unwrap();
1112        assert_eq!(messages.len(), 2);
1113    }
1114
1115    #[tokio::test]
1116    async fn test_snapshot_restore_preserves_recent_tail() {
1117        let config = CompactingMemoryConfig {
1118            max_recent_messages: 3,
1119            compress_threshold: 5,
1120            summarize_batch_size: 10,
1121            max_summary_length: 500,
1122        };
1123        let memory = CompactingMemory::new(Arc::new(NoopSummarizer), config);
1124
1125        for i in 0..7 {
1126            memory
1127                .add_message(make_message(&format!("msg{}", i)))
1128                .await
1129                .unwrap();
1130        }
1131        memory.compress(None).await.unwrap();
1132        let snapshot = memory.snapshot().await.unwrap();
1133        assert_eq!(snapshot.summarized_count, 4);
1134        let serialized = serde_json::to_string(&snapshot).unwrap();
1135        let persisted: MemorySnapshot = serde_json::from_str(&serialized).unwrap();
1136
1137        memory.clear().await.unwrap();
1138        memory.restore(persisted).await.unwrap();
1139        assert!(memory.summary().is_some());
1140        assert_eq!(memory.summarized_count(), 4);
1141        for i in 7..9 {
1142            memory
1143                .add_message(make_message(&format!("msg{}", i)))
1144                .await
1145                .unwrap();
1146        }
1147        memory.compress(None).await.unwrap();
1148
1149        let remaining = memory.get_messages(None).await.unwrap();
1150        let contents: Vec<_> = remaining
1151            .iter()
1152            .map(|message| message.content.as_str())
1153            .collect();
1154        assert_eq!(contents, vec!["msg6", "msg7", "msg8"]);
1155        let context = memory.get_context().await.unwrap();
1156        assert_eq!(context.summarized_count, 6);
1157        assert_eq!(context.total_messages, 9);
1158    }
1159
1160    #[tokio::test]
1161    async fn test_compression_history() {
1162        let memory = create_test_memory();
1163
1164        for i in 0..6 {
1165            memory
1166                .add_message(make_message(&format!("msg{}", i)))
1167                .await
1168                .unwrap();
1169        }
1170
1171        memory.compress(None).await.unwrap();
1172
1173        let history = memory.compression_history();
1174        assert_eq!(history.len(), 1);
1175        assert_eq!(history[0].messages_compressed, 3);
1176    }
1177
1178    #[test]
1179    fn test_config_default() {
1180        let config = CompactingMemoryConfig::default();
1181        assert_eq!(config.max_recent_messages, 50);
1182        assert_eq!(config.compress_threshold, 30);
1183        assert_eq!(config.summarize_batch_size, 10);
1184        assert_eq!(config.max_summary_length, 2000);
1185    }
1186
1187    #[test]
1188    fn test_config_rejects_unknown_fields() {
1189        let yaml = r#"
1190max_recent_messages: 5
1191compress_thresold: 10
1192"#;
1193        let error = serde_yaml::from_str::<CompactingMemoryConfig>(yaml).unwrap_err();
1194        assert!(
1195            error
1196                .to_string()
1197                .contains("unknown field `compress_thresold`")
1198        );
1199    }
1200
1201    #[tokio::test]
1202    async fn test_evict_oldest() {
1203        let memory = create_test_memory();
1204        for i in 0..5 {
1205            memory
1206                .add_message(make_message(&format!("msg{}", i)))
1207                .await
1208                .unwrap();
1209        }
1210
1211        let evicted = memory.evict_oldest(2).await.unwrap();
1212        assert_eq!(evicted.len(), 2);
1213        assert_eq!(evicted[0].content, "msg0");
1214        assert_eq!(evicted[1].content, "msg1");
1215
1216        let remaining = memory.get_messages(None).await.unwrap();
1217        assert_eq!(remaining.len(), 3);
1218        assert_eq!(remaining[0].content, "msg2");
1219    }
1220
1221    #[tokio::test]
1222    async fn compression_expands_batch_to_complete_signed_past_turn_and_projects_state() {
1223        let summarizer = Arc::new(CapturingSummarizer::default());
1224        let config = CompactingMemoryConfig {
1225            max_recent_messages: 1,
1226            compress_threshold: 5,
1227            summarize_batch_size: 2,
1228            max_summary_length: 100_000,
1229        };
1230        let memory = CompactingMemory::new(summarizer.clone(), config);
1231        for message in signed_turn("compress-past") {
1232            memory.add_message(message).await.unwrap();
1233        }
1234        memory
1235            .add_message(ChatMessage::user("new user turn"))
1236            .await
1237            .unwrap();
1238
1239        let result = memory.compress(None).await.unwrap();
1240
1241        assert!(matches!(
1242            result,
1243            CompressResult::Compressed {
1244                messages_summarized: 4,
1245                ..
1246            }
1247        ));
1248        let remaining = memory.get_messages(None).await.unwrap();
1249        assert_eq!(remaining.len(), 1);
1250        assert_eq!(remaining[0].content, "new user turn");
1251        let batches = summarizer.batches.read();
1252        assert_eq!(batches.len(), 1);
1253        assert_eq!(batches[0].len(), 4);
1254        let projected = batches[0]
1255            .iter()
1256            .map(|message| message.content.as_str())
1257            .collect::<Vec<_>>()
1258            .join("\n");
1259        assert!(!projected.contains("fixture-signature"));
1260        assert!(!projected.contains("_ai_agents_provider_state"));
1261        assert!(projected.contains("native_tool_calls"));
1262        assert!(projected.contains("native_tool_results"));
1263        assert!(!memory.summary().unwrap().contains("fixture-signature"));
1264    }
1265
1266    #[tokio::test]
1267    async fn compression_keeps_latest_signed_turn_intact() {
1268        let config = CompactingMemoryConfig {
1269            max_recent_messages: 0,
1270            compress_threshold: 4,
1271            summarize_batch_size: 2,
1272            max_summary_length: 100_000,
1273        };
1274        let memory = CompactingMemory::new(Arc::new(NoopSummarizer), config);
1275        for message in signed_turn("compress-active") {
1276            memory.add_message(message).await.unwrap();
1277        }
1278
1279        let before = memory.get_messages(None).await.unwrap();
1280        let result = memory.compress(None).await.unwrap();
1281        let after = memory.get_messages(None).await.unwrap();
1282
1283        assert!(matches!(result, CompressResult::NotNeeded));
1284        assert_eq!(after.len(), before.len());
1285        assert!(after[1].content.contains("fixture-signature"));
1286        assert!(memory.summary().is_none());
1287    }
1288
1289    #[tokio::test]
1290    async fn compression_projects_missing_result_for_ended_signed_turn() {
1291        let summarizer = Arc::new(CapturingSummarizer::default());
1292        let config = CompactingMemoryConfig {
1293            max_recent_messages: 1,
1294            compress_threshold: 4,
1295            summarize_batch_size: 2,
1296            max_summary_length: 100_000,
1297        };
1298        let memory = CompactingMemory::new(summarizer.clone(), config);
1299        let mut incomplete = signed_turn("compress-incomplete");
1300        incomplete.remove(2);
1301        for message in incomplete {
1302            memory.add_message(message).await.unwrap();
1303        }
1304        memory
1305            .add_message(ChatMessage::user("new user turn"))
1306            .await
1307            .unwrap();
1308
1309        let result = memory.compress(None).await.unwrap();
1310
1311        assert!(matches!(
1312            result,
1313            CompressResult::Compressed {
1314                messages_summarized: 3,
1315                ..
1316            }
1317        ));
1318        let batches = summarizer.batches.read();
1319        let projected = batches[0]
1320            .iter()
1321            .map(|message| message.content.as_str())
1322            .collect::<Vec<_>>()
1323            .join("\n");
1324        assert!(projected.contains("native_exchange_status"));
1325        assert!(projected.contains("final execution result was not recorded"));
1326        assert!(!projected.contains("fixture-signature"));
1327        assert!(memory.summary().unwrap().contains("native_exchange_status"));
1328    }
1329
1330    #[tokio::test]
1331    async fn framework_retains_incomplete_status_when_summarizer_drops_it() {
1332        let config = CompactingMemoryConfig {
1333            max_recent_messages: 1,
1334            compress_threshold: 4,
1335            summarize_batch_size: 2,
1336            max_summary_length: 100_000,
1337        };
1338        let memory = CompactingMemory::new(Arc::new(DroppingStatusSummarizer), config);
1339        let mut incomplete = signed_turn("compress-dropped-status");
1340        incomplete.remove(2);
1341        for message in incomplete {
1342            memory.add_message(message).await.unwrap();
1343        }
1344        memory
1345            .add_message(ChatMessage::user("new user turn"))
1346            .await
1347            .unwrap();
1348
1349        memory.compress(None).await.unwrap();
1350
1351        let summary = memory.summary().unwrap();
1352        assert!(summary.contains("summary that omitted the native status"));
1353        assert!(summary.contains("native_exchange_status"));
1354        assert!(!summary.contains("fixture-signature"));
1355    }
1356
1357    #[tokio::test]
1358    async fn compacting_eviction_keeps_signed_turn_atomic() {
1359        let memory = create_test_memory();
1360        for message in signed_turn("evict-past") {
1361            memory.add_message(message).await.unwrap();
1362        }
1363        memory
1364            .add_message(ChatMessage::user("new user turn"))
1365            .await
1366            .unwrap();
1367
1368        let evicted = memory.evict_oldest(1).await.unwrap();
1369
1370        assert_eq!(evicted.len(), 4);
1371        let remaining = memory.get_messages(None).await.unwrap();
1372        assert_eq!(remaining.len(), 1);
1373        assert_eq!(remaining[0].content, "new user turn");
1374    }
1375
1376    #[test]
1377    fn test_prefix_at_char_boundary_handles_unicode() {
1378        let text = "계약서 내용을 확인하고 싶어서";
1379        let prefix = prefix_at_char_boundary(text, 5);
1380        assert_eq!(prefix.chars().count(), 5);
1381        assert!(text.starts_with(prefix));
1382    }
1383}