a3s_code_core/
memory.rs

1//! Memory and learning system for the agent.
2//!
3//! Core types (`MemoryStore`, `MemoryItem`, `MemoryType`, `RelevanceConfig`,
4//! `FileMemoryStore`, `InMemoryStore`) live in `a3s-memory`.
5//!
6//! This module owns `MemoryConfig`, `MemoryStats`, `AgentMemory` (three-tier
7//! session memory), and `MemoryContextProvider` (context injection bridge).
8
9use a3s_memory::{MemoryItem, MemoryStore, MemoryType, PrunePolicy, RelevanceConfig};
10use chrono::{DateTime, Utc};
11use serde::{Deserialize, Serialize};
12use std::collections::VecDeque;
13use std::sync::atomic::{AtomicBool, Ordering};
14use std::sync::Arc;
15use tokio::sync::RwLock;
16
17// ============================================================================
18// Configuration
19// ============================================================================
20
21/// Configuration for the agent memory system (three-tier: working/short-term/long-term)
22#[derive(Debug, Clone, Serialize, Deserialize)]
23#[serde(rename_all = "camelCase")]
24pub struct MemoryConfig {
25    /// Relevance scoring parameters
26    #[serde(default)]
27    pub relevance: RelevanceConfig,
28    /// Maximum short-term memory items (default: 100)
29    #[serde(default = "MemoryConfig::default_max_short_term")]
30    pub max_short_term: usize,
31    /// Maximum working memory items (default: 10)
32    #[serde(default = "MemoryConfig::default_max_working")]
33    pub max_working: usize,
34    /// Automatic pruning policy for long-term storage. `None` disables background pruning.
35    #[serde(default)]
36    pub prune_policy: Option<PrunePolicy>,
37    /// How often the background pruning task runs, in seconds (default: 3600).
38    #[serde(default = "MemoryConfig::default_prune_interval_secs")]
39    pub prune_interval_secs: u64,
40    /// Use an LLM after significant completed turns to distill durable memories
41    /// from the turn transcript.
42    ///
43    /// Enabled by default when memory is configured. The runtime still applies
44    /// a significance gate so trivial turns do not trigger an extraction call.
45    #[serde(
46        default = "MemoryConfig::default_llm_extraction",
47        alias = "llm_extraction"
48    )]
49    pub llm_extraction: bool,
50    /// Maximum durable memories the LLM extractor may write per turn.
51    #[serde(default = "MemoryConfig::default_llm_extraction_max_items")]
52    pub llm_extraction_max_items: usize,
53    /// Maximum transcript characters passed into the LLM memory extractor.
54    #[serde(default = "MemoryConfig::default_llm_extraction_max_input_chars")]
55    pub llm_extraction_max_input_chars: usize,
56}
57
58impl MemoryConfig {
59    fn default_max_short_term() -> usize {
60        100
61    }
62    fn default_max_working() -> usize {
63        10
64    }
65    fn default_prune_interval_secs() -> u64 {
66        3600
67    }
68    fn default_llm_extraction() -> bool {
69        true
70    }
71    fn default_llm_extraction_max_items() -> usize {
72        5
73    }
74    fn default_llm_extraction_max_input_chars() -> usize {
75        8_000
76    }
77}
78
79impl Default for MemoryConfig {
80    fn default() -> Self {
81        Self {
82            relevance: RelevanceConfig::default(),
83            max_short_term: 100,
84            max_working: 10,
85            prune_policy: None,
86            prune_interval_secs: 3600,
87            llm_extraction: true,
88            llm_extraction_max_items: 5,
89            llm_extraction_max_input_chars: 8_000,
90        }
91    }
92}
93
94// ============================================================================
95// Memory Stats
96// ============================================================================
97
98/// Statistics for the three-tier agent memory system
99#[derive(Debug, Clone, Serialize, Deserialize)]
100pub struct MemoryStats {
101    pub long_term_count: usize,
102    pub short_term_count: usize,
103    pub working_count: usize,
104}
105
106// ============================================================================
107// Agent Memory (three-tier: working / short-term / long-term)
108// ============================================================================
109
110/// Three-tier agent memory: working, short-term (session), and long-term (persisted).
111#[derive(Clone)]
112pub struct AgentMemory {
113    /// Long-term memory store
114    pub(crate) store: Arc<dyn MemoryStore>,
115    /// Short-term memory (current session)
116    short_term: Arc<RwLock<VecDeque<MemoryItem>>>,
117    /// Working memory (active context)
118    working: Arc<RwLock<Vec<MemoryItem>>>,
119    pub(crate) max_short_term: usize,
120    pub(crate) max_working: usize,
121    pub(crate) relevance_config: RelevanceConfig,
122    pub(crate) llm_extraction: bool,
123    pub(crate) llm_extraction_max_items: usize,
124    pub(crate) llm_extraction_max_input_chars: usize,
125    llm_extraction_in_flight: Arc<AtomicBool>,
126}
127
128pub(crate) struct MemoryExtractionPermit {
129    in_flight: Arc<AtomicBool>,
130}
131
132impl Drop for MemoryExtractionPermit {
133    fn drop(&mut self) {
134        self.in_flight.store(false, Ordering::Release);
135    }
136}
137
138impl std::fmt::Debug for AgentMemory {
139    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
140        f.debug_struct("AgentMemory")
141            .field("max_short_term", &self.max_short_term)
142            .field("max_working", &self.max_working)
143            .finish()
144    }
145}
146
147impl AgentMemory {
148    /// Create a new agent memory system with default configuration
149    pub fn new(store: Arc<dyn MemoryStore>) -> Self {
150        Self::with_config(store, MemoryConfig::default())
151    }
152
153    /// Create a new agent memory system with custom configuration.
154    ///
155    /// If `config.prune_policy` is `Some`, a background Tokio task is spawned
156    /// that periodically calls `store.prune()` at the configured interval.
157    pub fn with_config(store: Arc<dyn MemoryStore>, config: MemoryConfig) -> Self {
158        if let Some(policy) = config.prune_policy.clone() {
159            let store_for_task = Arc::clone(&store);
160            let interval_secs = config.prune_interval_secs;
161            match tokio::runtime::Handle::try_current() {
162                Ok(handle) => {
163                    handle.spawn(async move {
164                        let mut ticker =
165                            tokio::time::interval(std::time::Duration::from_secs(interval_secs));
166                        ticker.tick().await; // skip the immediate first tick
167                        loop {
168                            ticker.tick().await;
169                            if let Err(e) = store_for_task.prune(&policy).await {
170                                tracing::warn!("memory prune failed: {e}");
171                            }
172                        }
173                    });
174                }
175                Err(_) => {
176                    tracing::warn!(
177                        "memory prune policy configured but no async runtime is available"
178                    );
179                }
180            }
181        }
182
183        Self {
184            store,
185            short_term: Arc::new(RwLock::new(VecDeque::new())),
186            working: Arc::new(RwLock::new(Vec::new())),
187            max_short_term: config.max_short_term,
188            max_working: config.max_working,
189            relevance_config: config.relevance,
190            llm_extraction: config.llm_extraction,
191            llm_extraction_max_items: config.llm_extraction_max_items,
192            llm_extraction_max_input_chars: config.llm_extraction_max_input_chars,
193            llm_extraction_in_flight: Arc::new(AtomicBool::new(false)),
194        }
195    }
196
197    pub(crate) fn score(&self, item: &MemoryItem, now: DateTime<Utc>) -> f32 {
198        let age_days = (now - item.timestamp).num_seconds() as f32 / 86400.0;
199        let decay = (-age_days / self.relevance_config.decay_days).exp();
200        item.importance * self.relevance_config.importance_weight
201            + decay * self.relevance_config.recency_weight
202    }
203
204    /// Store a memory in long-term storage and add to short-term
205    pub async fn remember(&self, item: MemoryItem) -> anyhow::Result<()> {
206        self.remember_item(item).await.map(|_| ())
207    }
208
209    /// Store a memory and return the normalized item that was sent to storage.
210    pub async fn remember_item(&self, item: MemoryItem) -> anyhow::Result<MemoryItem> {
211        self.store.store(item.clone()).await?;
212        let mut short_term = self.short_term.write().await;
213        if let Some(existing) = short_term
214            .iter_mut()
215            .find(|existing| existing.id == item.id)
216        {
217            *existing = item.clone();
218        } else {
219            short_term.push_back(item.clone());
220        }
221        if short_term.len() > self.max_short_term {
222            short_term.pop_front();
223        }
224        Ok(item)
225    }
226
227    /// Remove a memory from long-term storage and session-local memory tiers.
228    pub async fn forget(&self, id: &str) -> anyhow::Result<()> {
229        self.store.delete(id).await?;
230        self.short_term.write().await.retain(|item| item.id != id);
231        self.working.write().await.retain(|item| item.id != id);
232        Ok(())
233    }
234
235    /// Remember a successful pattern
236    pub async fn remember_success(
237        &self,
238        prompt: &str,
239        tools_used: &[String],
240        result: &str,
241    ) -> anyhow::Result<()> {
242        self.remember_success_item(prompt, tools_used, result)
243            .await
244            .map(|_| ())
245    }
246
247    /// Remember a successful pattern and return the stored memory item.
248    pub async fn remember_success_item(
249        &self,
250        prompt: &str,
251        tools_used: &[String],
252        result: &str,
253    ) -> anyhow::Result<MemoryItem> {
254        let content = format!(
255            "Success: {}\nTools: {}\nResult: {}",
256            prompt,
257            tools_used.join(", "),
258            result
259        );
260        let mut item = MemoryItem::new(content)
261            .with_importance(0.8)
262            .with_tag("success")
263            .with_tag("pattern")
264            .with_type(MemoryType::Procedural)
265            .with_metadata("prompt", prompt)
266            .with_metadata("tools", tools_used.join(","));
267        for tool in tools_used {
268            item = item.with_tag(tool.clone());
269        }
270        self.remember_item(item).await
271    }
272
273    /// Remember a failure to avoid repeating
274    pub async fn remember_failure(
275        &self,
276        prompt: &str,
277        error: &str,
278        attempted_tools: &[String],
279    ) -> anyhow::Result<()> {
280        self.remember_failure_item(prompt, error, attempted_tools)
281            .await
282            .map(|_| ())
283    }
284
285    /// Remember a failed pattern and return the stored memory item.
286    pub async fn remember_failure_item(
287        &self,
288        prompt: &str,
289        error: &str,
290        attempted_tools: &[String],
291    ) -> anyhow::Result<MemoryItem> {
292        let content = format!(
293            "Failure: {}\nError: {}\nAttempted tools: {}",
294            prompt,
295            error,
296            attempted_tools.join(", ")
297        );
298        let mut item = MemoryItem::new(content)
299            .with_importance(0.9)
300            .with_tag("failure")
301            .with_tag("avoid")
302            .with_type(MemoryType::Episodic)
303            .with_metadata("prompt", prompt)
304            .with_metadata("error", error);
305        for tool in attempted_tools {
306            item = item.with_tag(tool.clone());
307        }
308        self.remember_item(item).await
309    }
310
311    /// Recall similar past experiences
312    pub async fn recall_similar(
313        &self,
314        prompt: &str,
315        limit: usize,
316    ) -> anyhow::Result<Vec<MemoryItem>> {
317        self.store.search(prompt, limit).await
318    }
319
320    /// Recall by tags
321    pub async fn recall_by_tags(
322        &self,
323        tags: &[String],
324        limit: usize,
325    ) -> anyhow::Result<Vec<MemoryItem>> {
326        self.store.search_by_tags(tags, limit).await
327    }
328
329    /// Get recent memories
330    pub async fn get_recent(&self, limit: usize) -> anyhow::Result<Vec<MemoryItem>> {
331        self.store.get_recent(limit).await
332    }
333
334    /// Add to working memory (auto-trims by relevance if over capacity)
335    pub async fn add_to_working(&self, item: MemoryItem) -> anyhow::Result<()> {
336        let mut working = self.working.write().await;
337        working.push(item);
338        if working.len() > self.max_working {
339            let now = Utc::now();
340            working.sort_by(|a, b| {
341                self.score(b, now)
342                    .partial_cmp(&self.score(a, now))
343                    .unwrap_or(std::cmp::Ordering::Equal)
344            });
345            working.truncate(self.max_working);
346        }
347        Ok(())
348    }
349
350    /// Get working memory
351    pub async fn get_working(&self) -> Vec<MemoryItem> {
352        self.working.read().await.clone()
353    }
354
355    /// Clear working memory
356    pub async fn clear_working(&self) {
357        self.working.write().await.clear();
358    }
359
360    /// Get short-term memory
361    pub async fn get_short_term(&self) -> Vec<MemoryItem> {
362        self.short_term.read().await.iter().cloned().collect()
363    }
364
365    /// Clear short-term memory
366    pub async fn clear_short_term(&self) {
367        self.short_term.write().await.clear();
368    }
369
370    /// Get memory statistics
371    pub async fn stats(&self) -> anyhow::Result<MemoryStats> {
372        Ok(MemoryStats {
373            long_term_count: self.store.count().await?,
374            short_term_count: self.short_term.read().await.len(),
375            working_count: self.working.read().await.len(),
376        })
377    }
378
379    /// Get access to the underlying store
380    pub fn store(&self) -> &Arc<dyn MemoryStore> {
381        &self.store
382    }
383
384    /// Get working memory count
385    pub async fn working_count(&self) -> usize {
386        self.working.read().await.len()
387    }
388
389    /// Get short-term memory count
390    pub async fn short_term_count(&self) -> usize {
391        self.short_term.read().await.len()
392    }
393
394    pub(crate) fn llm_extraction_enabled(&self) -> bool {
395        self.llm_extraction
396    }
397
398    pub(crate) fn llm_extraction_max_items(&self) -> usize {
399        self.llm_extraction_max_items
400    }
401
402    pub(crate) fn llm_extraction_max_input_chars(&self) -> usize {
403        self.llm_extraction_max_input_chars
404    }
405
406    pub(crate) fn try_begin_llm_extraction(&self) -> Option<MemoryExtractionPermit> {
407        self.llm_extraction_in_flight
408            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
409            .ok()
410            .map(|_| MemoryExtractionPermit {
411                in_flight: Arc::clone(&self.llm_extraction_in_flight),
412            })
413    }
414}
415
416// ============================================================================
417// Memory Context Provider
418// ============================================================================
419
420/// Context provider that surfaces past memories as agent context.
421pub struct MemoryContextProvider {
422    memory: AgentMemory,
423}
424
425impl MemoryContextProvider {
426    pub fn new(memory: AgentMemory) -> Self {
427        Self { memory }
428    }
429}
430
431pub(crate) fn memory_items_to_context_result(
432    provider: impl Into<String>,
433    items: Vec<MemoryItem>,
434) -> crate::context::ContextResult {
435    let mut result = crate::context::ContextResult::new(provider);
436    let total = items.len().max(1);
437    for (index, item) in items.into_iter().enumerate() {
438        let supersedes = relation_ids(&item, "supersedes");
439        let conflicts_with = relation_ids(&item, "conflicts_with");
440        let content = memory_context_content(&item, &supersedes, &conflicts_with);
441        let token_count = (content.len() / 4).max(1);
442        let recall_rank_score = 1.0 - (index as f32 / total as f32);
443        let relevance = (item.relevance_score() * 0.35 + recall_rank_score * 0.65).clamp(0.0, 1.0);
444        let context_item = crate::context::ContextItem::new(
445            &item.id,
446            crate::context::ContextType::Memory,
447            content,
448        )
449        .with_relevance(relevance)
450        .with_token_count(token_count)
451        .with_source(format!("memory://{}", item.id))
452        .with_metadata("memory_id", serde_json::json!(item.id))
453        .with_metadata(
454            "memory_type",
455            serde_json::json!(memory_type_label(item.memory_type)),
456        )
457        .with_metadata("tags", serde_json::json!(item.tags))
458        .with_metadata("importance", serde_json::json!(item.importance))
459        .with_provenance("long_term_memory")
460        .with_priority(0.35)
461        .with_trust(0.7)
462        .with_freshness(0.5);
463        let context_item = add_relation_metadata(context_item, "supersedes", supersedes);
464        let context_item = add_relation_metadata(context_item, "conflicts_with", conflicts_with);
465        result.add_item(context_item);
466    }
467    result
468}
469
470fn relation_ids(item: &MemoryItem, key: &str) -> Vec<String> {
471    item.metadata
472        .get(key)
473        .map(|value| {
474            value
475                .split(',')
476                .map(str::trim)
477                .filter(|id| !id.is_empty())
478                .map(ToOwned::to_owned)
479                .collect()
480        })
481        .unwrap_or_default()
482}
483
484fn memory_context_content(
485    item: &MemoryItem,
486    supersedes: &[String],
487    conflicts_with: &[String],
488) -> String {
489    let mut content = item.content.clone();
490    if supersedes.is_empty() && conflicts_with.is_empty() {
491        return content;
492    }
493
494    content.push_str("\n\nMemory relations:");
495    if !supersedes.is_empty() {
496        content.push_str("\n- supersedes: ");
497        content.push_str(&relation_sources(supersedes));
498    }
499    if !conflicts_with.is_empty() {
500        content.push_str("\n- conflicts_with: ");
501        content.push_str(&relation_sources(conflicts_with));
502    }
503    content
504}
505
506fn relation_sources(ids: &[String]) -> String {
507    ids.iter()
508        .map(|id| format!("memory://{id}"))
509        .collect::<Vec<_>>()
510        .join(", ")
511}
512
513fn add_relation_metadata(
514    item: crate::context::ContextItem,
515    key: &str,
516    ids: Vec<String>,
517) -> crate::context::ContextItem {
518    if ids.is_empty() {
519        item
520    } else {
521        item.with_metadata(key, serde_json::json!(ids))
522    }
523}
524
525fn memory_type_label(memory_type: MemoryType) -> &'static str {
526    match memory_type {
527        MemoryType::Episodic => "episodic",
528        MemoryType::Semantic => "semantic",
529        MemoryType::Procedural => "procedural",
530        MemoryType::Working => "working",
531    }
532}
533
534#[async_trait::async_trait]
535impl crate::context::ContextProvider for MemoryContextProvider {
536    fn name(&self) -> &str {
537        "memory"
538    }
539
540    async fn query(
541        &self,
542        query: &crate::context::ContextQuery,
543    ) -> anyhow::Result<crate::context::ContextResult> {
544        let limit = query.max_results.min(5);
545        let items = self.memory.recall_similar(&query.query, limit).await?;
546
547        Ok(memory_items_to_context_result("memory", items))
548    }
549
550    async fn on_turn_complete(
551        &self,
552        _session_id: &str,
553        _prompt: &str,
554        _response: &str,
555    ) -> anyhow::Result<()> {
556        // Memory extraction is owned by the agent loop's gated LLM extractor.
557        // This provider only contributes recalled memories as prompt context.
558        Ok(())
559    }
560}
561
562// ============================================================================
563// Tests
564// ============================================================================
565
566#[cfg(test)]
567mod tests {
568    use super::*;
569    use crate::context::ContextProvider;
570    use a3s_memory::InMemoryStore;
571    use std::sync::Arc;
572
573    #[tokio::test]
574    async fn test_agent_memory_remember_and_recall() {
575        let memory = AgentMemory::new(Arc::new(InMemoryStore::new()));
576        memory
577            .remember_success("create file", &["write".to_string()], "ok")
578            .await
579            .unwrap();
580        memory
581            .remember_failure("delete file", "denied", &["bash".to_string()])
582            .await
583            .unwrap();
584
585        let results = memory.recall_similar("create", 10).await.unwrap();
586        assert!(!results.is_empty());
587
588        let stats = memory.stats().await.unwrap();
589        assert_eq!(stats.long_term_count, 2);
590        assert_eq!(stats.short_term_count, 2);
591    }
592
593    #[tokio::test]
594    async fn test_agent_memory_forget_removes_all_tiers() {
595        let memory = AgentMemory::new(Arc::new(InMemoryStore::new()));
596        let item = memory
597            .remember_item(MemoryItem::new("superseded memory"))
598            .await
599            .unwrap();
600        memory.add_to_working(item.clone()).await.unwrap();
601
602        memory.forget(&item.id).await.unwrap();
603
604        assert_eq!(memory.stats().await.unwrap().long_term_count, 0);
605        assert!(memory.get_short_term().await.is_empty());
606        assert!(memory.get_working().await.is_empty());
607    }
608
609    #[tokio::test]
610    async fn test_agent_memory_uses_canonical_store_item_for_duplicates() {
611        let memory = AgentMemory::new(Arc::new(InMemoryStore::new()));
612        let first = memory
613            .remember_item(
614                MemoryItem::new("Run focused memory extraction tests after parser changes.")
615                    .with_importance(0.3)
616                    .with_tag("memory"),
617            )
618            .await
619            .unwrap();
620
621        let duplicate = memory
622            .remember_item(
623                MemoryItem::new("run focused MEMORY extraction tests after parser changes!")
624                    .with_importance(0.9)
625                    .with_tag("tests"),
626            )
627            .await
628            .unwrap();
629
630        assert_eq!(duplicate.id, first.id);
631        assert_eq!(memory.stats().await.unwrap().long_term_count, 1);
632        let short_term = memory.get_short_term().await;
633        assert_eq!(short_term.len(), 1);
634        assert_eq!(short_term[0].id, first.id);
635        assert_eq!(short_term[0].importance, 0.9);
636        assert!(short_term[0].tags.contains(&"memory".to_string()));
637        assert!(short_term[0].tags.contains(&"tests".to_string()));
638    }
639
640    #[tokio::test]
641    async fn test_agent_memory_working() {
642        let memory = AgentMemory::new(Arc::new(InMemoryStore::new()));
643        memory
644            .add_to_working(MemoryItem::new("task").with_type(MemoryType::Working))
645            .await
646            .unwrap();
647        assert_eq!(memory.working_count().await, 1);
648        memory.clear_working().await;
649        assert_eq!(memory.working_count().await, 0);
650    }
651
652    #[tokio::test]
653    async fn test_agent_memory_working_overflow_trims() {
654        let memory = AgentMemory {
655            store: Arc::new(InMemoryStore::new()),
656            short_term: Arc::new(RwLock::new(VecDeque::new())),
657            working: Arc::new(RwLock::new(Vec::new())),
658            max_short_term: 100,
659            max_working: 3,
660            relevance_config: RelevanceConfig::default(),
661            llm_extraction: false,
662            llm_extraction_max_items: 5,
663            llm_extraction_max_input_chars: 8_000,
664            llm_extraction_in_flight: Arc::new(AtomicBool::new(false)),
665        };
666        for i in 0..5 {
667            memory
668                .add_to_working(
669                    MemoryItem::new(format!("task {i}")).with_importance(i as f32 * 0.2),
670                )
671                .await
672                .unwrap();
673        }
674        assert_eq!(memory.get_working().await.len(), 3);
675    }
676
677    #[tokio::test]
678    async fn test_agent_memory_recall_by_tags() {
679        let memory = AgentMemory::new(Arc::new(InMemoryStore::new()));
680        memory
681            .remember_success("create file", &["write".to_string()], "ok")
682            .await
683            .unwrap();
684        memory
685            .remember_failure("delete file", "denied", &["bash".to_string()])
686            .await
687            .unwrap();
688
689        let successes = memory
690            .recall_by_tags(&["success".to_string()], 10)
691            .await
692            .unwrap();
693        assert_eq!(successes.len(), 1);
694        let failures = memory
695            .recall_by_tags(&["failure".to_string()], 10)
696            .await
697            .unwrap();
698        assert_eq!(failures.len(), 1);
699    }
700
701    #[tokio::test]
702    async fn test_agent_memory_short_term_trim() {
703        let store = Arc::new(InMemoryStore::new());
704        let memory = AgentMemory {
705            store,
706            short_term: Arc::new(RwLock::new(VecDeque::new())),
707            working: Arc::new(RwLock::new(Vec::new())),
708            max_short_term: 3,
709            max_working: 10,
710            relevance_config: RelevanceConfig::default(),
711            llm_extraction: false,
712            llm_extraction_max_items: 5,
713            llm_extraction_max_input_chars: 8_000,
714            llm_extraction_in_flight: Arc::new(AtomicBool::new(false)),
715        };
716        for i in 0..5 {
717            memory
718                .remember(MemoryItem::new(format!("item {i}")))
719                .await
720                .unwrap();
721        }
722        assert_eq!(memory.short_term_count().await, 3);
723    }
724
725    #[tokio::test]
726    async fn test_agent_memory_prune_delegates() {
727        use a3s_memory::PrunePolicy;
728
729        let store = Arc::new(InMemoryStore::new());
730        let memory = AgentMemory::new(store.clone());
731
732        // Insert one old low-importance item directly into the store.
733        let mut old_item = a3s_memory::MemoryItem::new("stale").with_importance(0.2);
734        old_item.timestamp = chrono::Utc::now() - chrono::Duration::days(100);
735        store.store(old_item).await.unwrap();
736
737        assert_eq!(store.count().await.unwrap(), 1);
738
739        // Calling prune on the underlying store via the public accessor works.
740        let policy = PrunePolicy {
741            max_age_days: 90,
742            min_importance_to_keep: 0.5,
743            max_items: 0,
744        };
745        let deleted = memory.store().prune(&policy).await.unwrap();
746        assert_eq!(deleted, 1);
747        assert_eq!(store.count().await.unwrap(), 0);
748    }
749
750    #[test]
751    fn test_agent_memory_score_uses_config() {
752        let config = MemoryConfig {
753            relevance: RelevanceConfig {
754                decay_days: 7.0,
755                importance_weight: 0.9,
756                recency_weight: 0.1,
757            },
758            ..Default::default()
759        };
760        let memory = AgentMemory::with_config(Arc::new(InMemoryStore::new()), config);
761        let item = MemoryItem::new("Test").with_importance(1.0);
762        let score = memory.score(&item, Utc::now());
763        assert!(score > 0.95, "Score was {score}");
764    }
765
766    #[test]
767    fn test_memory_config_partial_deserialize_keeps_llm_extraction_enabled() {
768        let config: MemoryConfig = serde_json::from_str(r#"{"maxShortTerm": 12}"#).unwrap();
769        assert!(config.llm_extraction);
770        assert_eq!(config.max_short_term, 12);
771    }
772
773    #[test]
774    fn test_memory_config_allows_explicit_llm_extraction_disable() {
775        let config: MemoryConfig =
776            serde_json::from_str(r#"{"llmExtraction": false, "maxShortTerm": 12}"#).unwrap();
777        assert!(!config.llm_extraction);
778        assert_eq!(config.max_short_term, 12);
779    }
780
781    #[test]
782    fn test_memory_context_result_includes_relation_context() {
783        let item = MemoryItem::new("Use the file memory store for local sessions.")
784            .with_type(MemoryType::Procedural)
785            .with_tag("consolidated")
786            .with_tag("conflict")
787            .with_metadata("supersedes", "old-preference, old-workflow")
788            .with_metadata("conflicts_with", "legacy-default");
789
790        let result = memory_items_to_context_result("memory", vec![item.clone()]);
791
792        assert_eq!(result.items.len(), 1);
793        let context_item = &result.items[0];
794        assert!(context_item
795            .content
796            .contains("Use the file memory store for local sessions."));
797        assert!(context_item.content.contains("Memory relations:"));
798        assert!(context_item
799            .content
800            .contains("supersedes: memory://old-preference, memory://old-workflow"));
801        assert!(context_item
802            .content
803            .contains("conflicts_with: memory://legacy-default"));
804        assert_eq!(
805            context_item.metadata.get("memory_id"),
806            Some(&serde_json::json!(item.id))
807        );
808        assert_eq!(
809            context_item.metadata.get("memory_type"),
810            Some(&serde_json::json!("procedural"))
811        );
812        assert_eq!(
813            context_item.metadata.get("tags"),
814            Some(&serde_json::json!(["consolidated", "conflict"]))
815        );
816        assert_eq!(
817            context_item.metadata.get("supersedes"),
818            Some(&serde_json::json!(["old-preference", "old-workflow"]))
819        );
820        assert_eq!(
821            context_item.metadata.get("conflicts_with"),
822            Some(&serde_json::json!(["legacy-default"]))
823        );
824        assert_eq!(
825            context_item.token_count,
826            (context_item.content.len() / 4).max(1)
827        );
828    }
829
830    #[test]
831    fn test_memory_context_relevance_preserves_recall_order() {
832        let top_match =
833            MemoryItem::new("Run focused memory extraction tests after parser changes.")
834                .with_importance(0.2)
835                .with_type(MemoryType::Procedural);
836        let generic_high_importance = MemoryItem::new("Remember general memory behavior.")
837            .with_importance(1.0)
838            .with_type(MemoryType::Semantic);
839
840        let result =
841            memory_items_to_context_result("memory", vec![top_match, generic_high_importance]);
842
843        assert_eq!(result.items.len(), 2);
844        assert!(
845            result.items[0].relevance > result.items[1].relevance,
846            "search recall order should remain a strong memory context ranking signal"
847        );
848    }
849
850    #[tokio::test]
851    async fn test_memory_context_provider_does_not_mechanically_store_turns() {
852        let memory = AgentMemory::new(Arc::new(InMemoryStore::new()));
853        let provider = MemoryContextProvider::new(memory.clone());
854
855        provider
856            .on_turn_complete("session-1", "remember nothing", "ok")
857            .await
858            .unwrap();
859
860        assert_eq!(memory.stats().await.unwrap().long_term_count, 0);
861    }
862}
a3s_code_core/memory.rs

a3s_code_core/
memory.rs
