1use ares_types::types::{AgentContext, MemoryFact, Message, Preference, UserMemory};
12use serde::{Deserialize, Serialize};
13use std::collections::HashMap;
14use std::fmt;
15
16pub const DEFAULT_HISTORY_WINDOW: usize = 10;
18
19pub const MAX_FACTS_IN_PROMPT: usize = 20;
21
22pub const MAX_PREFERENCES_IN_PROMPT: usize = 10;
24
25pub fn format_memory_for_prompt(memory: &UserMemory) -> String {
40 let mut parts = Vec::new();
41
42 if !memory.preferences.is_empty() {
44 let prefs: Vec<String> = memory
45 .preferences
46 .iter()
47 .take(MAX_PREFERENCES_IN_PROMPT)
48 .filter(|p| p.confidence >= 0.5) .map(|p| format!("- {}/{}: {}", p.category, p.key, p.value))
50 .collect();
51
52 if !prefs.is_empty() {
53 parts.push(format!("User Preferences:\n{}", prefs.join("\n")));
54 }
55 }
56
57 if !memory.facts.is_empty() {
59 let facts: Vec<String> = memory
60 .facts
61 .iter()
62 .take(MAX_FACTS_IN_PROMPT)
63 .filter(|f| f.confidence >= 0.5) .map(|f| format!("- {}/{}: {}", f.category, f.fact_key, f.fact_value))
65 .collect();
66
67 if !facts.is_empty() {
68 parts.push(format!("Known Facts about User:\n{}", facts.join("\n")));
69 }
70 }
71
72 parts.join("\n\n")
73}
74
75pub fn format_preferences_compact(preferences: &[Preference]) -> String {
80 preferences
81 .iter()
82 .filter(|p| p.confidence >= 0.5)
83 .take(MAX_PREFERENCES_IN_PROMPT)
84 .map(|p| format!("{}: {}", p.key, p.value))
85 .collect::<Vec<_>>()
86 .join(", ")
87}
88
89pub fn truncate_history(history: &[Message], window_size: usize) -> Vec<Message> {
98 if history.len() <= window_size {
99 history.to_vec()
100 } else {
101 history[history.len() - window_size..].to_vec()
102 }
103}
104
105pub fn estimate_tokens(text: &str) -> usize {
114 let words = text.split_whitespace().count();
115 let chars = text.len();
116 let word_estimate = (words as f64 * 1.3) as usize;
118 let char_estimate = chars.div_ceil(4);
119 word_estimate.max(char_estimate).max(1)
121}
122
123pub fn truncate_history_to_tokens(history: &[Message], token_budget: usize) -> Vec<Message> {
134 let mut result: Vec<Message> = Vec::new();
135 let mut total_tokens = 0;
136
137 for msg in history.iter().rev() {
139 let msg_tokens = estimate_tokens(&msg.content);
140 if total_tokens + msg_tokens > token_budget {
141 break;
142 }
143 result.push(msg.clone());
144 total_tokens += msg_tokens;
145 }
146
147 result.reverse();
149 result
150}
151
152pub fn build_context(
164 user_id: String,
165 session_id: String,
166 history: Vec<Message>,
167 memory: Option<UserMemory>,
168 history_window: Option<usize>,
169) -> AgentContext {
170 let window = history_window.unwrap_or(DEFAULT_HISTORY_WINDOW);
171 let truncated_history = truncate_history(&history, window);
172
173 AgentContext {
174 user_id,
175 session_id,
176 conversation_history: truncated_history,
177 user_memory: memory,
178 }
179}
180
181pub fn filter_facts_by_category(facts: &[MemoryFact], category: &str) -> Vec<MemoryFact> {
185 facts
186 .iter()
187 .filter(|f| f.category == category)
188 .cloned()
189 .collect()
190}
191
192pub fn filter_preferences_by_category(
194 preferences: &[Preference],
195 category: &str,
196) -> Vec<Preference> {
197 preferences
198 .iter()
199 .filter(|p| p.category == category)
200 .cloned()
201 .collect()
202}
203
204
205#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
211pub struct MemoryConfig {
212 pub max_sessions: usize,
214 pub session_ttl_secs: u64,
216}
217
218impl Default for MemoryConfig {
219 fn default() -> Self {
220 Self {
221 max_sessions: 128,
222 session_ttl_secs: 3600,
223 }
224 }
225}
226
227#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
229pub struct MemorySession {
230 pub tenant_id: String,
231 pub session_id: String,
232 pub user_id: String,
233 pub payload: serde_json::Value,
234 pub created_at: i64,
235 pub updated_at: i64,
236}
237
238#[derive(Debug, Clone, PartialEq, Eq)]
240pub enum MemoryError {
241 NotFound {
242 tenant_id: String,
243 session_id: String,
244 },
245 CapacityExceeded {
246 max_sessions: usize,
247 },
248 InvalidTenant {
249 expected: String,
250 actual: String,
251 },
252}
253
254impl fmt::Display for MemoryError {
255 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
256 match self {
257 MemoryError::NotFound {
258 tenant_id,
259 session_id,
260 } => write!(
261 f,
262 "session '{session_id}' not found for tenant '{tenant_id}'"
263 ),
264 MemoryError::CapacityExceeded { max_sessions } => {
265 write!(f, "session store capacity exceeded (max {max_sessions})")
266 }
267 MemoryError::InvalidTenant { expected, actual } => write!(
268 f,
269 "invalid tenant: expected '{expected}', got '{actual}'"
270 ),
271 }
272 }
273}
274
275impl std::error::Error for MemoryError {}
276
277pub fn session_key(tenant_id: &str, session_id: &str) -> String {
279 format!("{tenant_id}\x1f{session_id}")
280}
281
282pub fn ttl_expired(created_at: i64, now: i64, ttl_secs: u64) -> bool {
284 ttl_secs > 0 && now.saturating_sub(created_at) > ttl_secs as i64
285}
286
287pub fn lru_get(order: &mut Vec<String>, key: &str) -> bool {
289 if let Some(pos) = order.iter().position(|k| k == key) {
290 let entry = order.remove(pos);
291 order.push(entry);
292 true
293 } else {
294 false
295 }
296}
297
298pub fn lru_put(order: &mut Vec<String>, key: String) {
300 if let Some(pos) = order.iter().position(|k| k == &key) {
301 order.remove(pos);
302 }
303 order.push(key);
304}
305
306pub fn lru_evict(
308 order: &mut Vec<String>,
309 sessions: &mut HashMap<String, MemorySession>,
310 max_sessions: usize,
311) {
312 while order.len() > max_sessions {
313 if let Some(victim) = order.first().cloned() {
314 order.remove(0);
315 sessions.remove(&victim);
316 } else {
317 break;
318 }
319 }
320}
321
322pub fn cleanup_expired(
324 sessions: &mut HashMap<String, MemorySession>,
325 order: &mut Vec<String>,
326 now: i64,
327 ttl_secs: u64,
328) -> usize {
329 if ttl_secs == 0 {
330 return 0;
331 }
332 let expired: Vec<String> = sessions
333 .iter()
334 .filter(|(_, s)| ttl_expired(s.created_at, now, ttl_secs))
335 .map(|(k, _)| k.clone())
336 .collect();
337 for key in &expired {
338 sessions.remove(key);
339 if let Some(pos) = order.iter().position(|k| k == key) {
340 order.remove(pos);
341 }
342 }
343 expired.len()
344}
345
346#[derive(Debug, Clone)]
348pub struct MemoryStore {
349 config: MemoryConfig,
350 sessions: HashMap<String, MemorySession>,
351 lru_order: Vec<String>,
352}
353
354impl MemoryStore {
355 pub fn new(config: MemoryConfig) -> Self {
356 Self {
357 config,
358 sessions: HashMap::new(),
359 lru_order: Vec::new(),
360 }
361 }
362
363 pub fn with_defaults() -> Self {
364 Self::new(MemoryConfig::default())
365 }
366
367 pub fn config(&self) -> &MemoryConfig {
368 &self.config
369 }
370
371 pub fn len(&self) -> usize {
372 self.sessions.len()
373 }
374
375 pub fn is_empty(&self) -> bool {
376 self.sessions.is_empty()
377 }
378
379 pub fn upsert(&mut self, session: MemorySession, now: i64) -> Result<(), MemoryError> {
380 let key = session_key(&session.tenant_id, &session.session_id);
381 lru_put(&mut self.lru_order, key.clone());
382 self.sessions.insert(key, session);
383 lru_evict(
384 &mut self.lru_order,
385 &mut self.sessions,
386 self.config.max_sessions,
387 );
388 if self.sessions.len() > self.config.max_sessions {
389 return Err(MemoryError::CapacityExceeded {
390 max_sessions: self.config.max_sessions,
391 });
392 }
393 let _ = now;
394 Ok(())
395 }
396
397 pub fn put_for_tenant(
398 &mut self,
399 tenant_id: &str,
400 session: MemorySession,
401 now: i64,
402 ) -> Result<(), MemoryError> {
403 if session.tenant_id != tenant_id {
404 return Err(MemoryError::InvalidTenant {
405 expected: tenant_id.to_string(),
406 actual: session.tenant_id.clone(),
407 });
408 }
409 self.upsert(session, now)
410 }
411
412 pub fn get(
413 &mut self,
414 tenant_id: &str,
415 session_id: &str,
416 now: i64,
417 ) -> Result<MemorySession, MemoryError> {
418 let key = session_key(tenant_id, session_id);
419 let Some(session) = self.sessions.get(&key).cloned() else {
420 return Err(MemoryError::NotFound {
421 tenant_id: tenant_id.to_string(),
422 session_id: session_id.to_string(),
423 });
424 };
425 if ttl_expired(session.created_at, now, self.config.session_ttl_secs) {
426 self.remove(tenant_id, session_id);
427 return Err(MemoryError::NotFound {
428 tenant_id: tenant_id.to_string(),
429 session_id: session_id.to_string(),
430 });
431 }
432 lru_get(&mut self.lru_order, &key);
433 Ok(session)
434 }
435
436 pub fn remove(&mut self, tenant_id: &str, session_id: &str) -> bool {
437 let key = session_key(tenant_id, session_id);
438 if self.sessions.remove(&key).is_some() {
439 if let Some(pos) = self.lru_order.iter().position(|k| k == &key) {
440 self.lru_order.remove(pos);
441 }
442 true
443 } else {
444 false
445 }
446 }
447
448 pub fn cleanup(&mut self, now: i64) -> usize {
449 cleanup_expired(
450 &mut self.sessions,
451 &mut self.lru_order,
452 now,
453 self.config.session_ttl_secs,
454 )
455 }
456}
457
458#[cfg(test)]
459mod tests {
460 use super::*;
461 use ares_types::types::MessageRole;
462 use chrono::Utc;
463 use std::collections::HashMap;
464
465 #[test]
466 fn test_format_memory_for_prompt_empty() {
467 let memory = UserMemory {
468 user_id: "test".to_string(),
469 preferences: vec![],
470 facts: vec![],
471 };
472 assert_eq!(format_memory_for_prompt(&memory), "");
473 }
474
475 #[test]
476 fn test_format_memory_for_prompt_with_preferences() {
477 let memory = UserMemory {
478 user_id: "test".to_string(),
479 preferences: vec![Preference {
480 category: "communication".to_string(),
481 key: "style".to_string(),
482 value: "concise".to_string(),
483 confidence: 0.9,
484 }],
485 facts: vec![],
486 };
487 let result = format_memory_for_prompt(&memory);
488 assert!(result.contains("User Preferences:"));
489 assert!(result.contains("communication/style: concise"));
490 }
491
492 #[test]
493 fn test_format_memory_filters_low_confidence() {
494 let memory = UserMemory {
495 user_id: "test".to_string(),
496 preferences: vec![
497 Preference {
498 category: "test".to_string(),
499 key: "high".to_string(),
500 value: "yes".to_string(),
501 confidence: 0.8,
502 },
503 Preference {
504 category: "test".to_string(),
505 key: "low".to_string(),
506 value: "no".to_string(),
507 confidence: 0.3, },
509 ],
510 facts: vec![],
511 };
512 let result = format_memory_for_prompt(&memory);
513 assert!(result.contains("high"));
514 assert!(!result.contains("low"));
515 }
516
517 #[test]
518 fn test_truncate_history() {
519 let history: Vec<Message> = (0..10)
520 .map(|i| Message {
521 role: MessageRole::User,
522 content: format!("Message {}", i),
523 timestamp: Utc::now(),
524 })
525 .collect();
526
527 let truncated = truncate_history(&history, 3);
528 assert_eq!(truncated.len(), 3);
529 assert!(truncated[0].content.contains("7"));
530 assert!(truncated[2].content.contains("9"));
531 }
532
533 #[test]
534 fn test_estimate_tokens() {
535 assert_eq!(estimate_tokens(""), 1); assert_eq!(estimate_tokens("test"), 1);
537 assert_eq!(estimate_tokens("this is a longer test string"), 7);
538 }
539
540 #[test]
541 fn test_format_preferences_compact() {
542 let prefs = vec![
543 Preference {
544 category: "output".to_string(),
545 key: "format".to_string(),
546 value: "markdown".to_string(),
547 confidence: 0.9,
548 },
549 Preference {
550 category: "output".to_string(),
551 key: "length".to_string(),
552 value: "brief".to_string(),
553 confidence: 0.8,
554 },
555 ];
556 let result = format_preferences_compact(&prefs);
557 assert_eq!(result, "format: markdown, length: brief");
558 }
559
560 #[test]
561 fn test_build_context() {
562 let history: Vec<Message> = (0..20)
563 .map(|i| Message {
564 role: MessageRole::User,
565 content: format!("Message {}", i),
566 timestamp: Utc::now(),
567 })
568 .collect();
569
570 let context = build_context(
571 "user1".to_string(),
572 "session1".to_string(),
573 history,
574 None,
575 Some(5),
576 );
577
578 assert_eq!(context.user_id, "user1");
579 assert_eq!(context.session_id, "session1");
580 assert_eq!(context.conversation_history.len(), 5);
581 assert!(context.user_memory.is_none());
582 }
583
584 #[test]
585 fn test_filter_facts_by_category() {
586 let facts = vec![
587 MemoryFact {
588 id: "1".to_string(),
589 user_id: "test".to_string(),
590 category: "work".to_string(),
591 fact_key: "role".to_string(),
592 fact_value: "engineer".to_string(),
593 confidence: 0.9,
594 created_at: Utc::now(),
595 updated_at: Utc::now(),
596 },
597 MemoryFact {
598 id: "2".to_string(),
599 user_id: "test".to_string(),
600 category: "personal".to_string(),
601 fact_key: "hobby".to_string(),
602 fact_value: "reading".to_string(),
603 confidence: 0.8,
604 created_at: Utc::now(),
605 updated_at: Utc::now(),
606 },
607 ];
608
609 let work_facts = filter_facts_by_category(&facts, "work");
610 assert_eq!(work_facts.len(), 1);
611 assert_eq!(work_facts[0].fact_key, "role");
612 }
613
614 #[test]
615 fn test_format_memory_for_prompt_with_facts() {
616 let memory = UserMemory {
617 user_id: "test".to_string(),
618 preferences: vec![],
619 facts: vec![MemoryFact {
620 id: "f1".to_string(),
621 user_id: "test".to_string(),
622 category: "work".to_string(),
623 fact_key: "role".to_string(),
624 fact_value: "engineer".to_string(),
625 confidence: 0.9,
626 created_at: Utc::now(),
627 updated_at: Utc::now(),
628 }],
629 };
630 let result = format_memory_for_prompt(&memory);
631 assert!(result.contains("Known Facts about User:"));
632 assert!(result.contains("work/role: engineer"));
633 }
634
635 #[test]
636 fn test_truncate_history_noop_when_within_window() {
637 let history: Vec<Message> = vec![Message {
638 role: MessageRole::User,
639 content: "only one".to_string(),
640 timestamp: Utc::now(),
641 }];
642 let truncated = truncate_history(&history, 5);
643 assert_eq!(truncated.len(), 1);
644 assert_eq!(truncated[0].content, "only one");
645 }
646
647 #[test]
648 fn test_truncate_history_to_tokens() {
649 let history: Vec<Message> = (0..5)
650 .map(|i| Message {
651 role: MessageRole::User,
652 content: format!("word {}", i),
653 timestamp: Utc::now(),
654 })
655 .collect();
656
657 let truncated = truncate_history_to_tokens(&history, 3);
658 assert!(!truncated.is_empty());
659 assert!(truncated.len() < history.len());
660 assert_eq!(truncated.last().unwrap().content, "word 4");
661 }
662
663 #[test]
664 fn test_filter_preferences_by_category() {
665 let prefs = vec![
666 Preference {
667 category: "output".to_string(),
668 key: "format".to_string(),
669 value: "markdown".to_string(),
670 confidence: 0.9,
671 },
672 Preference {
673 category: "language".to_string(),
674 key: "locale".to_string(),
675 value: "en".to_string(),
676 confidence: 0.9,
677 },
678 ];
679
680 let output = filter_preferences_by_category(&prefs, "output");
681 assert_eq!(output.len(), 1);
682 assert_eq!(output[0].key, "format");
683 }
684
685 #[test]
686 fn test_user_memory_serde_roundtrip() {
687 let memory = UserMemory {
688 user_id: "user-42".to_string(),
689 preferences: vec![Preference {
690 category: "output".to_string(),
691 key: "format".to_string(),
692 value: "markdown".to_string(),
693 confidence: 0.75,
694 }],
695 facts: vec![MemoryFact {
696 id: "fact-1".to_string(),
697 user_id: "user-42".to_string(),
698 category: "work".to_string(),
699 fact_key: "team".to_string(),
700 fact_value: "platform".to_string(),
701 confidence: 0.95,
702 created_at: Utc::now(),
703 updated_at: Utc::now(),
704 }],
705 };
706
707 let json = serde_json::to_string(&memory).expect("serialize UserMemory");
708 let parsed: UserMemory = serde_json::from_str(&json).expect("deserialize UserMemory");
709 assert_eq!(parsed.user_id, "user-42");
710 assert_eq!(parsed.preferences.len(), 1);
711 assert_eq!(parsed.facts[0].fact_key, "team");
712 }
713
714 #[test]
715 fn test_message_serde_roundtrip() {
716 let msg = Message {
717 role: MessageRole::Assistant,
718 content: "hello".to_string(),
719 timestamp: Utc::now(),
720 };
721 let json = serde_json::to_string(&msg).expect("serialize Message");
722 let parsed: Message = serde_json::from_str(&json).expect("deserialize Message");
723 assert_eq!(parsed.content, "hello");
724 assert!(matches!(parsed.role, MessageRole::Assistant));
725 }
726
727 #[test]
728 fn test_format_memory_respects_max_limits() {
729 let prefs: Vec<Preference> = (0..MAX_PREFERENCES_IN_PROMPT + 5)
730 .map(|i| Preference {
731 category: "output".to_string(),
732 key: format!("key-{i}"),
733 value: "v".to_string(),
734 confidence: 0.9,
735 })
736 .collect();
737 let facts: Vec<MemoryFact> = (0..MAX_FACTS_IN_PROMPT + 5)
738 .map(|i| MemoryFact {
739 id: format!("id-{i}"),
740 user_id: "u".to_string(),
741 category: "work".to_string(),
742 fact_key: format!("fact-{i}"),
743 fact_value: "v".to_string(),
744 confidence: 0.9,
745 created_at: Utc::now(),
746 updated_at: Utc::now(),
747 })
748 .collect();
749
750 let memory = UserMemory {
751 user_id: "u".to_string(),
752 preferences: prefs,
753 facts,
754 };
755 let formatted = format_memory_for_prompt(&memory);
756 for i in MAX_PREFERENCES_IN_PROMPT..MAX_PREFERENCES_IN_PROMPT + 5 {
757 assert!(
758 !formatted.contains(&format!("key-{i}")),
759 "preference beyond cap should be omitted"
760 );
761 }
762 for i in MAX_FACTS_IN_PROMPT..MAX_FACTS_IN_PROMPT + 5 {
763 assert!(
764 !formatted.contains(&format!("fact-{i}")),
765 "fact beyond cap should be omitted"
766 );
767 }
768 assert!(formatted.contains("key-0"));
769 assert!(formatted.contains("fact-0"));
770 }
771
772 #[test]
773 fn test_format_memory_filters_low_confidence_facts() {
774 let memory = UserMemory {
775 user_id: "test".to_string(),
776 preferences: vec![],
777 facts: vec![
778 MemoryFact {
779 id: "1".to_string(),
780 user_id: "test".to_string(),
781 category: "work".to_string(),
782 fact_key: "kept".to_string(),
783 fact_value: "yes".to_string(),
784 confidence: 0.9,
785 created_at: Utc::now(),
786 updated_at: Utc::now(),
787 },
788 MemoryFact {
789 id: "2".to_string(),
790 user_id: "test".to_string(),
791 category: "work".to_string(),
792 fact_key: "dropped".to_string(),
793 fact_value: "no".to_string(),
794 confidence: 0.2,
795 created_at: Utc::now(),
796 updated_at: Utc::now(),
797 },
798 ],
799 };
800 let result = format_memory_for_prompt(&memory);
801 assert!(result.contains("kept"));
802 assert!(!result.contains("dropped"));
803 }
804
805 #[test]
806 fn test_truncate_history_zero_window_returns_empty() {
807 let history: Vec<Message> = (0..3)
808 .map(|i| Message {
809 role: MessageRole::User,
810 content: format!("msg {i}"),
811 timestamp: Utc::now(),
812 })
813 .collect();
814 assert!(truncate_history(&history, 0).is_empty());
815 }
816
817 #[test]
818 fn test_truncate_history_to_tokens_zero_budget_returns_empty() {
819 let history: Vec<Message> = vec![Message {
820 role: MessageRole::User,
821 content: "non-empty".to_string(),
822 timestamp: Utc::now(),
823 }];
824 assert!(truncate_history_to_tokens(&history, 0).is_empty());
825 }
826
827 #[test]
828 fn test_build_context_uses_default_history_window() {
829 let history: Vec<Message> = (0..DEFAULT_HISTORY_WINDOW + 5)
830 .map(|i| Message {
831 role: MessageRole::User,
832 content: format!("Message {}", i),
833 timestamp: Utc::now(),
834 })
835 .collect();
836
837 let context = build_context(
838 "user".to_string(),
839 "session".to_string(),
840 history,
841 None,
842 None,
843 );
844 assert_eq!(context.conversation_history.len(), DEFAULT_HISTORY_WINDOW);
845 }
846
847 #[test]
848 fn test_format_preferences_compact_filters_low_confidence() {
849 let prefs = vec![
850 Preference {
851 category: "output".to_string(),
852 key: "keep".to_string(),
853 value: "yes".to_string(),
854 confidence: 0.9,
855 },
856 Preference {
857 category: "output".to_string(),
858 key: "drop".to_string(),
859 value: "no".to_string(),
860 confidence: 0.1,
861 },
862 ];
863 let result = format_preferences_compact(&prefs);
864 assert!(result.contains("keep"));
865 assert!(!result.contains("drop"));
866 }
867
868 #[test]
869 fn test_filter_facts_by_category_no_matches() {
870 let facts = vec![MemoryFact {
871 id: "1".to_string(),
872 user_id: "u".to_string(),
873 category: "work".to_string(),
874 fact_key: "role".to_string(),
875 fact_value: "engineer".to_string(),
876 confidence: 0.9,
877 created_at: Utc::now(),
878 updated_at: Utc::now(),
879 }];
880 assert!(filter_facts_by_category(&facts, "missing").is_empty());
881 }
882
883 #[test]
884 fn test_truncate_history_to_tokens_keeps_recent_within_budget() {
885 let history: Vec<Message> = vec![
886 Message {
887 role: MessageRole::User,
888 content: "a".repeat(40),
889 timestamp: Utc::now(),
890 },
891 Message {
892 role: MessageRole::User,
893 content: "short".to_string(),
894 timestamp: Utc::now(),
895 },
896 ];
897 let truncated = truncate_history_to_tokens(&history, 5);
898 assert_eq!(truncated.len(), 1);
899 assert_eq!(truncated[0].content, "short");
900 }
901 fn sample_session(tenant: &str, session: &str, created_at: i64) -> MemorySession {
906 MemorySession {
907 tenant_id: tenant.into(),
908 session_id: session.into(),
909 user_id: "user-1".into(),
910 payload: serde_json::json!({"turn": 1}),
911 created_at,
912 updated_at: created_at,
913 }
914 }
915
916 #[test]
917 fn memory_config_serde_roundtrip() {
918 let cfg = MemoryConfig {
919 max_sessions: 4,
920 session_ttl_secs: 120,
921 };
922 let json = serde_json::to_string(&cfg).expect("serialize");
923 let back: MemoryConfig = serde_json::from_str(&json).expect("deserialize");
924 assert_eq!(back, cfg);
925 }
926
927 #[test]
928 fn memory_session_serde_roundtrip() {
929 let session = sample_session("tenant-a", "sess-1", 100);
930 let json = serde_json::to_string(&session).expect("serialize");
931 let back: MemorySession = serde_json::from_str(&json).expect("deserialize");
932 assert_eq!(back, session);
933 }
934
935 #[test]
936 fn memory_config_default_values() {
937 let cfg = MemoryConfig::default();
938 assert_eq!(cfg.max_sessions, 128);
939 assert_eq!(cfg.session_ttl_secs, 3600);
940 }
941
942 #[test]
943 fn session_key_includes_tenant_and_session() {
944 let key = session_key("tenant-a", "sess-1");
945 assert!(key.contains("tenant-a"));
946 assert!(key.contains("sess-1"));
947 assert_ne!(session_key("tenant-a", "sess-1"), session_key("tenant-b", "sess-1"));
948 }
949
950 #[test]
951 fn ttl_expired_false_when_within_ttl() {
952 assert!(!ttl_expired(100, 150, 60));
953 }
954
955 #[test]
956 fn ttl_expired_true_when_past_ttl() {
957 assert!(ttl_expired(100, 200, 60));
958 }
959
960 #[test]
961 fn ttl_expired_disabled_when_ttl_zero() {
962 assert!(!ttl_expired(0, 1_000_000, 0));
963 }
964
965 #[test]
966 fn ttl_expired_false_at_exact_boundary() {
967 assert!(!ttl_expired(100, 160, 60));
968 }
969
970 #[test]
971 fn lru_put_appends_new_key() {
972 let mut order = Vec::new();
973 lru_put(&mut order, "a".into());
974 lru_put(&mut order, "b".into());
975 assert_eq!(order, vec!["a", "b"]);
976 }
977
978 #[test]
979 fn lru_put_promotes_existing_key() {
980 let mut order = vec!["a".into(), "b".into(), "c".into()];
981 lru_put(&mut order, "a".into());
982 assert_eq!(order, vec!["b", "c", "a"]);
983 }
984
985 #[test]
986 fn lru_get_promotes_key_to_end() {
987 let mut order = vec!["a".into(), "b".into(), "c".into()];
988 assert!(lru_get(&mut order, "a"));
989 assert_eq!(order, vec!["b", "c", "a"]);
990 }
991
992 #[test]
993 fn lru_get_missing_returns_false() {
994 let mut order = vec!["a".into()];
995 assert!(!lru_get(&mut order, "missing"));
996 }
997
998 #[test]
999 fn lru_evict_drops_least_recently_used() {
1000 let mut order = vec!["old".into(), "mid".into(), "new".into()];
1001 let mut sessions = HashMap::new();
1002 sessions.insert("old".into(), sample_session("t", "old", 1));
1003 sessions.insert("mid".into(), sample_session("t", "mid", 1));
1004 sessions.insert("new".into(), sample_session("t", "new", 1));
1005 lru_evict(&mut order, &mut sessions, 2);
1006 assert_eq!(order, vec!["mid", "new"]);
1007 assert!(!sessions.contains_key("old"));
1008 }
1009
1010 #[test]
1011 fn lru_evict_noop_when_within_capacity() {
1012 let mut order = vec!["a".into()];
1013 let mut sessions = HashMap::new();
1014 sessions.insert("a".into(), sample_session("t", "a", 1));
1015 lru_evict(&mut order, &mut sessions, 2);
1016 assert_eq!(sessions.len(), 1);
1017 }
1018
1019 #[test]
1020 fn cleanup_expired_removes_stale_sessions() {
1021 let mut sessions = HashMap::new();
1022 let mut order = Vec::new();
1023 let key = session_key("t", "stale");
1024 sessions.insert(key.clone(), sample_session("t", "stale", 0));
1025 order.push(key);
1026 assert_eq!(cleanup_expired(&mut sessions, &mut order, 500, 60), 1);
1027 assert!(sessions.is_empty());
1028 }
1029
1030 #[test]
1031 fn cleanup_expired_keeps_fresh_sessions() {
1032 let mut sessions = HashMap::new();
1033 let mut order = Vec::new();
1034 let key = session_key("t", "fresh");
1035 sessions.insert(key.clone(), sample_session("t", "fresh", 400));
1036 order.push(key);
1037 assert_eq!(cleanup_expired(&mut sessions, &mut order, 450, 60), 0);
1038 assert_eq!(sessions.len(), 1);
1039 }
1040
1041 #[test]
1042 fn memory_store_put_and_get_roundtrip() {
1043 let mut store = MemoryStore::new(MemoryConfig {
1044 max_sessions: 4,
1045 session_ttl_secs: 0,
1046 });
1047 let session = sample_session("tenant-a", "sess-1", 10);
1048 store.put_for_tenant("tenant-a", session.clone(), 10).expect("put");
1049 let got = store.get("tenant-a", "sess-1", 10).expect("get");
1050 assert_eq!(got.session_id, "sess-1");
1051 }
1052
1053 #[test]
1054 fn memory_store_get_not_found() {
1055 let mut store = MemoryStore::with_defaults();
1056 let err = store.get("tenant-a", "missing", 0).unwrap_err();
1057 assert!(matches!(err, MemoryError::NotFound { .. }));
1058 }
1059
1060 #[test]
1061 fn memory_store_tenant_isolation() {
1062 let mut store = MemoryStore::new(MemoryConfig {
1063 max_sessions: 4,
1064 session_ttl_secs: 0,
1065 });
1066 store
1067 .put_for_tenant("tenant-a", sample_session("tenant-a", "shared-id", 1), 1)
1068 .expect("put");
1069 assert!(store.get("tenant-b", "shared-id", 1).is_err());
1070 }
1071
1072 #[test]
1073 fn memory_store_invalid_tenant_on_put() {
1074 let mut store = MemoryStore::with_defaults();
1075 let err = store
1076 .put_for_tenant("tenant-a", sample_session("tenant-b", "s", 1), 1)
1077 .unwrap_err();
1078 assert!(matches!(err, MemoryError::InvalidTenant { .. }));
1079 }
1080
1081 #[test]
1082 fn memory_store_capacity_evicts_lru() {
1083 let mut store = MemoryStore::new(MemoryConfig {
1084 max_sessions: 2,
1085 session_ttl_secs: 0,
1086 });
1087 store
1088 .put_for_tenant("t", sample_session("t", "one", 1), 1)
1089 .expect("put one");
1090 store
1091 .put_for_tenant("t", sample_session("t", "two", 2), 2)
1092 .expect("put two");
1093 store
1094 .put_for_tenant("t", sample_session("t", "three", 3), 3)
1095 .expect("put three");
1096 assert_eq!(store.len(), 2);
1097 assert!(store.get("t", "one", 3).is_err());
1098 assert!(store.get("t", "three", 3).is_ok());
1099 }
1100
1101 #[test]
1102 fn memory_store_get_promotes_lru_entry() {
1103 let mut store = MemoryStore::new(MemoryConfig {
1104 max_sessions: 2,
1105 session_ttl_secs: 0,
1106 });
1107 store
1108 .put_for_tenant("t", sample_session("t", "a", 1), 1)
1109 .expect("a");
1110 store
1111 .put_for_tenant("t", sample_session("t", "b", 2), 2)
1112 .expect("b");
1113 store.get("t", "a", 3).expect("touch a");
1114 store
1115 .put_for_tenant("t", sample_session("t", "c", 4), 4)
1116 .expect("c");
1117 assert!(store.get("t", "a", 4).is_ok());
1118 assert!(store.get("t", "b", 4).is_err());
1119 }
1120
1121 #[test]
1122 fn memory_store_get_expired_session_returns_not_found() {
1123 let mut store = MemoryStore::new(MemoryConfig {
1124 max_sessions: 4,
1125 session_ttl_secs: 60,
1126 });
1127 store
1128 .put_for_tenant("t", sample_session("t", "s", 100), 100)
1129 .expect("put");
1130 let err = store.get("t", "s", 200).unwrap_err();
1131 assert!(matches!(err, MemoryError::NotFound { .. }));
1132 }
1133
1134 #[test]
1135 fn memory_store_cleanup_removes_expired() {
1136 let mut store = MemoryStore::new(MemoryConfig {
1137 max_sessions: 4,
1138 session_ttl_secs: 30,
1139 });
1140 store
1141 .put_for_tenant("t", sample_session("t", "s", 10), 10)
1142 .expect("put");
1143 assert_eq!(store.cleanup(100), 1);
1144 assert!(store.is_empty());
1145 }
1146
1147 #[test]
1148 fn memory_store_remove_deletes_session() {
1149 let mut store = MemoryStore::with_defaults();
1150 store
1151 .put_for_tenant("t", sample_session("t", "s", 1), 1)
1152 .expect("put");
1153 assert!(store.remove("t", "s"));
1154 assert!(store.get("t", "s", 1).is_err());
1155 }
1156
1157 #[test]
1158 fn memory_store_remove_missing_is_false() {
1159 let mut store = MemoryStore::with_defaults();
1160 assert!(!store.remove("t", "missing"));
1161 }
1162
1163 #[test]
1164 fn memory_store_clone_preserves_sessions() {
1165 let mut store = MemoryStore::with_defaults();
1166 store
1167 .put_for_tenant("t", sample_session("t", "s", 1), 1)
1168 .expect("put");
1169 let cloned = store.clone();
1170 assert_eq!(cloned.len(), 1);
1171 }
1172
1173 #[test]
1174 fn memory_store_debug_contains_type_name() {
1175 let store = MemoryStore::with_defaults();
1176 assert!(format!("{store:?}").contains("MemoryStore"));
1177 }
1178
1179 #[test]
1180 fn memory_error_display_not_found() {
1181 let msg = MemoryError::NotFound {
1182 tenant_id: "t".into(),
1183 session_id: "s".into(),
1184 }
1185 .to_string();
1186 assert!(msg.contains("t") && msg.contains("s"));
1187 }
1188
1189 #[test]
1190 fn memory_error_display_capacity_exceeded() {
1191 let msg = MemoryError::CapacityExceeded { max_sessions: 2 }.to_string();
1192 assert!(msg.contains("2"));
1193 }
1194
1195 #[test]
1196 fn memory_error_display_invalid_tenant() {
1197 let msg = MemoryError::InvalidTenant {
1198 expected: "a".into(),
1199 actual: "b".into(),
1200 }
1201 .to_string();
1202 assert!(msg.contains("a") && msg.contains("b"));
1203 }
1204
1205 #[test]
1206 fn memory_session_clone_eq() {
1207 let a = sample_session("t", "s", 1);
1208 assert_eq!(a.clone(), a);
1209 }
1210
1211 #[test]
1212 fn memory_config_clone_eq() {
1213 let a = MemoryConfig::default();
1214 assert_eq!(a.clone(), a);
1215 }
1216
1217 #[test]
1218 fn memory_error_clone_eq() {
1219 let err = MemoryError::CapacityExceeded { max_sessions: 1 };
1220 assert_eq!(err.clone(), err);
1221 }
1222
1223 #[test]
1224 fn memory_store_new_is_empty() {
1225 let store = MemoryStore::with_defaults();
1226 assert!(store.is_empty());
1227 }
1228
1229 #[test]
1230 fn memory_store_upsert_updates_existing_payload() {
1231 let mut store = MemoryStore::new(MemoryConfig {
1232 max_sessions: 4,
1233 session_ttl_secs: 0,
1234 });
1235 store
1236 .put_for_tenant("t", sample_session("t", "s", 1), 1)
1237 .expect("first");
1238 let mut updated = sample_session("t", "s", 2);
1239 updated.payload = serde_json::json!({"turn": 2});
1240 store.put_for_tenant("t", updated, 2).expect("update");
1241 let got = store.get("t", "s", 2).expect("get");
1242 assert_eq!(got.payload["turn"], 2);
1243 }
1244
1245 #[test]
1246 fn memory_store_multiple_tenants_do_not_collide() {
1247 let mut store = MemoryStore::new(MemoryConfig {
1248 max_sessions: 8,
1249 session_ttl_secs: 0,
1250 });
1251 store
1252 .put_for_tenant("tenant-a", sample_session("tenant-a", "s", 1), 1)
1253 .expect("a");
1254 store
1255 .put_for_tenant("tenant-b", sample_session("tenant-b", "s", 1), 1)
1256 .expect("b");
1257 assert_eq!(store.len(), 2);
1258 }
1259
1260 #[test]
1261 fn lru_order_least_recently_used_evicted_first() {
1262 let mut order = vec!["first".into(), "second".into()];
1263 let mut sessions = HashMap::new();
1264 sessions.insert("first".into(), sample_session("t", "first", 1));
1265 sessions.insert("second".into(), sample_session("t", "second", 1));
1266 lru_evict(&mut order, &mut sessions, 1);
1267 assert!(!sessions.contains_key("first"));
1268 assert!(sessions.contains_key("second"));
1269 }
1270
1271 #[test]
1272 fn memory_store_config_accessor_returns_config() {
1273 let cfg = MemoryConfig {
1274 max_sessions: 9,
1275 session_ttl_secs: 5,
1276 };
1277 let store = MemoryStore::new(cfg.clone());
1278 assert_eq!(store.config(), &cfg);
1279 }
1280
1281 #[test]
1282 fn cleanup_expired_noop_when_ttl_disabled() {
1283 let mut sessions = HashMap::new();
1284 let mut order = Vec::new();
1285 let key = session_key("t", "s");
1286 sessions.insert(key.clone(), sample_session("t", "s", 0));
1287 order.push(key);
1288 assert_eq!(cleanup_expired(&mut sessions, &mut order, 999, 0), 0);
1289 }
1290
1291 #[test]
1292 fn reexports_memory_constants() {
1293 assert_eq!(DEFAULT_HISTORY_WINDOW, 10);
1294 assert_eq!(MAX_FACTS_IN_PROMPT, 20);
1295 assert_eq!(MAX_PREFERENCES_IN_PROMPT, 10);
1296 }
1297
1298}