use anyhow::Result;
use chrono::Utc;
use std::sync::Arc;
use tracing::info;
use crate::importance::ImportanceCalculator;
use crate::storage::MemoryStorage;
use mr_common::MemoryConfig;
pub struct LifecycleManager {
storage: Arc<dyn MemoryStorage>,
calculator: ImportanceCalculator,
config: MemoryConfig,
}
impl LifecycleManager {
pub fn new(
storage: Arc<dyn MemoryStorage>,
calculator: ImportanceCalculator,
config: MemoryConfig,
) -> Self {
Self {
storage,
calculator,
config,
}
}
pub async fn recalculate_importance(&self) -> Result<()> {
info!("Recalculating importance for all memories");
let memories = self.storage.list(1000).await?;
for memory in memories {
let new_importance = self.calculator.calculate(&memory);
if memory.importance != new_importance {
let mut updated = memory;
updated.importance = new_importance;
self.storage.update(&updated).await?;
}
}
info!("Importance recalculation completed");
Ok(())
}
pub async fn cleanup_cycle(&self) -> Result<()> {
info!("Starting cleanup cycle");
self.cleanup_deleted().await?;
self.cleanup_low_importance().await?;
info!("Cleanup cycle completed");
Ok(())
}
async fn cleanup_deleted(&self) -> Result<()> {
let deleted = self.storage.list_deleted().await?;
let now = Utc::now();
for memory in deleted {
if let Some(deleted_at) = memory.deleted_at {
let days_since_delete = (now - deleted_at).num_days();
if days_since_delete > self.config.soft_delete_recovery_days as i64 {
info!(
"Hard deleting memory {} (deleted {} days ago)",
memory.id, days_since_delete
);
self.storage.delete(&memory.id).await?;
}
}
}
Ok(())
}
async fn cleanup_low_importance(&self) -> Result<()> {
let low_importance = self
.storage
.list_by_importance(0.0, self.config.hard_delete_importance)
.await?;
let now = Utc::now();
for memory in low_importance {
let days_inactive = (now - memory.last_accessed).num_days();
if days_inactive > self.config.hard_delete_inactive_days as i64 {
info!(
"Deleting low importance memory {} (inactive {} days)",
memory.id, days_inactive
);
self.storage.delete(&memory.id).await?;
}
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::storage::{MemoryStore, RocksDBStore};
use mr_common::{Memory, MemoryType};
use tempfile::tempdir;
#[tokio::test]
async fn test_recalculate_importance() {
let dir = tempdir().unwrap();
let rocksdb = RocksDBStore::open(dir.path()).unwrap();
let storage = Arc::new(MemoryStore::new(rocksdb));
let memory = Memory::new("test".to_string(), MemoryType::Knowledge);
storage.save(&memory).await.unwrap();
let calc = ImportanceCalculator::default();
let config = MemoryConfig::default();
let manager = LifecycleManager::new(storage.clone(), calc, config);
manager.recalculate_importance().await.unwrap();
let retrieved = storage.get(&memory.id).await.unwrap().unwrap();
assert!(retrieved.importance > 0.0);
}
#[tokio::test]
async fn test_cleanup_soft_deleted_memories() {
let dir = tempdir().unwrap();
let rocksdb = RocksDBStore::open(dir.path()).unwrap();
let storage = Arc::new(MemoryStore::new(rocksdb));
let memory = Memory::new("to be deleted".to_string(), MemoryType::Knowledge);
storage.save(&memory).await.unwrap();
storage.delete(&memory.id).await.unwrap();
let mut deleted_mem = storage.get(&memory.id).await.unwrap().unwrap();
deleted_mem.deleted_at = Some(Utc::now() - chrono::Duration::days(2));
storage.update(&deleted_mem).await.unwrap();
let config = MemoryConfig {
soft_delete_recovery_days: 1,
..Default::default()
};
let calc = ImportanceCalculator::default();
let manager = LifecycleManager::new(storage.clone(), calc, config);
manager.cleanup_cycle().await.unwrap();
let result = storage.get(&memory.id).await.unwrap();
assert!(
result.is_none(),
"Memory should be hard deleted after cleanup"
);
}
#[tokio::test]
async fn test_cleanup_does_not_affect_active_memories() {
let dir = tempdir().unwrap();
let rocksdb = RocksDBStore::open(dir.path()).unwrap();
let storage = Arc::new(MemoryStore::new(rocksdb));
let active_memory = Memory::new("active memory".to_string(), MemoryType::Knowledge)
.with_tags(vec!["critical".to_string()]);
storage.save(&active_memory).await.unwrap();
let calc = ImportanceCalculator::default();
let config = MemoryConfig::default();
let manager = LifecycleManager::new(storage.clone(), calc, config);
manager.cleanup_cycle().await.unwrap();
let result = storage.get(&active_memory.id).await.unwrap();
assert!(result.is_some());
}
#[tokio::test]
async fn test_cleanup_low_importance_inactive_memories() {
let dir = tempdir().unwrap();
let rocksdb = RocksDBStore::open(dir.path()).unwrap();
let storage = Arc::new(MemoryStore::new(rocksdb));
let mut memory = Memory::new("low importance".to_string(), MemoryType::Conversation);
memory.importance = 0.05;
memory.last_accessed = Utc::now() - chrono::Duration::days(35);
storage.save(&memory).await.unwrap();
let low_importance = storage.list_by_importance(0.0, 0.1).await.unwrap();
assert!(
!low_importance.is_empty(),
"Memory should be in low importance list"
);
let config = MemoryConfig {
hard_delete_importance: 0.1,
hard_delete_inactive_days: 30,
..Default::default()
};
let calc = ImportanceCalculator::default();
let manager = LifecycleManager::new(storage.clone(), calc, config);
manager.cleanup_cycle().await.unwrap();
let result = storage.get(&memory.id).await.unwrap();
assert!(
result.is_some_and(|m| m.is_deleted),
"Memory should be soft deleted"
);
}
}