1use 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 #[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 #[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 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 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 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}