engine/

decay.rs

//! Importance Decay Engine for Dakera AI Agent Memory Platform.
//!
//! Background task that periodically decays memory importance scores
//! based on configurable strategies: Exponential, Linear, or StepFunction.
//!
//! On every recall hit, importance gets an access boost.
//! Memories below `min_importance` are auto-deleted.

use std::sync::Arc;

use common::{DecayConfig, DecayStrategy, Memory, MemoryType, Vector};
use serde::{Deserialize, Serialize};
use storage::VectorStorage;
use tracing;

/// Importance Decay Engine that runs as a background task.
pub struct DecayEngine {
    pub config: DecayConfig,
}

/// Configuration loaded from environment variables.
pub struct DecayEngineConfig {
    /// Decay configuration (strategy, half_life, min_importance)
    pub decay_config: DecayConfig,
    /// How often to run decay (in seconds)
    pub interval_secs: u64,
}

impl Default for DecayEngineConfig {
    fn default() -> Self {
        Self {
            decay_config: DecayConfig {
                strategy: DecayStrategy::Exponential,
                half_life_hours: 168.0, // 1 week
                min_importance: 0.01,
            },
            interval_secs: 3600, // 1 hour
        }
    }
}

impl DecayEngineConfig {
    /// Load configuration from environment variables.
    pub fn from_env() -> Self {
        let half_life_hours: f64 = std::env::var("DAKERA_DECAY_HALF_LIFE_HOURS")
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(168.0);

        let min_importance: f32 = std::env::var("DAKERA_DECAY_MIN_IMPORTANCE")
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(0.01);

        let interval_secs: u64 = std::env::var("DAKERA_DECAY_INTERVAL_SECS")
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(3600);

        let strategy_str =
            std::env::var("DAKERA_DECAY_STRATEGY").unwrap_or_else(|_| "exponential".to_string());

        let strategy = match strategy_str.to_lowercase().as_str() {
            "linear" => DecayStrategy::Linear,
            "step" | "stepfunction" | "step_function" => DecayStrategy::StepFunction,
            _ => DecayStrategy::Exponential,
        };

        Self {
            decay_config: DecayConfig {
                strategy,
                half_life_hours,
                min_importance,
            },
            interval_secs,
        }
    }
}

impl DecayEngine {
    /// Create a new DecayEngine with the given configuration.
    pub fn new(config: DecayConfig) -> Self {
        Self { config }
    }

    /// Calculate decayed importance for a single memory.
    ///
    /// Memory type determines base decay speed:
    /// - Working:    3× faster (temporary by design)
    /// - Episodic:   1× normal (events fade naturally)
    /// - Semantic:   0.5× slower (knowledge persists)
    /// - Procedural: 0.3× slower (skills are durable)
    ///
    /// Usage pattern shields from decay (diminishing returns).
    /// Never-accessed memories fade 50% faster.
    pub fn calculate_decay(
        &self,
        current_importance: f32,
        hours_elapsed: f64,
        memory_type: &MemoryType,
        access_count: u32,
    ) -> f32 {
        if hours_elapsed <= 0.0 {
            return current_importance;
        }

        // Memory type determines base decay speed
        let type_multiplier = match memory_type {
            MemoryType::Working => 3.0,
            MemoryType::Episodic => 1.0,
            MemoryType::Semantic => 0.5,
            MemoryType::Procedural => 0.3,
        };

        // Usage pattern shields from decay (diminishing returns)
        let usage_shield = if access_count > 0 {
            1.0 / (1.0 + (access_count as f64 * 0.1))
        } else {
            1.5 // never accessed = 50% faster decay
        };

        let effective_half_life = self.config.half_life_hours / (type_multiplier * usage_shield);

        let decayed = match self.config.strategy {
            DecayStrategy::Exponential => {
                let decay_factor = (0.5_f64).powf(hours_elapsed / effective_half_life);
                current_importance * decay_factor as f32
            }
            DecayStrategy::Linear => {
                let decay_amount = (hours_elapsed / effective_half_life) as f32 * 0.5;
                (current_importance - decay_amount).max(0.0)
            }
            DecayStrategy::StepFunction => {
                let steps = (hours_elapsed / effective_half_life).floor() as u32;
                let decay_factor = (0.5_f32).powi(steps as i32);
                current_importance * decay_factor
            }
        };

        decayed.clamp(0.0, 1.0)
    }

    /// Calculate access boost for a memory that was just recalled.
    /// Scales with current importance — valuable memories get bigger boosts.
    pub fn access_boost(current_importance: f32) -> f32 {
        let boost = 0.05 + 0.05 * current_importance; // 0.05–0.10 range
        (current_importance + boost).min(1.0)
    }

    /// Apply decay to all memories across all agent namespaces.
    ///
    /// Iterates all namespaces prefixed with `_dakera_agent_`, loads memories,
    /// applies decay based on time since last access, and removes memories
    /// below the minimum importance threshold.
    pub async fn apply_decay(&self, storage: &Arc<dyn VectorStorage>) -> DecayResult {
        let mut result = DecayResult::default();

        // List all namespaces
        let namespaces = match storage.list_namespaces().await {
            Ok(ns) => ns,
            Err(e) => {
                tracing::error!(error = %e, "Failed to list namespaces for decay");
                return result;
            }
        };

        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();

        // Only process Dakera agent namespaces
        for namespace in namespaces {
            if !namespace.starts_with("_dakera_agent_") {
                continue;
            }

            result.namespaces_processed += 1;

            let vectors = match storage.get_all(&namespace).await {
                Ok(v) => v,
                Err(e) => {
                    tracing::warn!(
                        namespace = %namespace,
                        error = %e,
                        "Failed to get vectors for decay"
                    );
                    continue;
                }
            };

            let mut updated_vectors: Vec<Vector> = Vec::new();
            let mut ids_to_delete: Vec<String> = Vec::new();

            for vector in &vectors {
                let memory = match Memory::from_vector(vector) {
                    Some(m) => m,
                    None => continue, // Skip non-memory vectors
                };

                result.memories_processed += 1;

                // Calculate hours since last access
                let hours_elapsed = if now > memory.last_accessed_at {
                    (now - memory.last_accessed_at) as f64 / 3600.0
                } else {
                    0.0
                };

                let new_importance = self.calculate_decay(
                    memory.importance,
                    hours_elapsed,
                    &memory.memory_type,
                    memory.access_count,
                );

                // Check if below minimum importance
                if new_importance < self.config.min_importance {
                    ids_to_delete.push(memory.id.clone());
                    result.memories_deleted += 1;
                    continue;
                }

                // Only update if importance actually changed
                if (new_importance - memory.importance).abs() > 0.001 {
                    let mut updated_memory = memory;
                    updated_memory.importance = new_importance;

                    // Rebuild vector with updated metadata but same embedding
                    let mut updated_vector = vector.clone();
                    updated_vector.metadata = Some(updated_memory.to_vector_metadata());
                    updated_vectors.push(updated_vector);
                    result.memories_decayed += 1;
                }
            }

            // Delete memories below threshold
            if !ids_to_delete.is_empty() {
                if let Err(e) = storage.delete(&namespace, &ids_to_delete).await {
                    tracing::warn!(
                        namespace = %namespace,
                        count = ids_to_delete.len(),
                        error = %e,
                        "Failed to delete expired memories"
                    );
                }
            }

            // Upsert updated memories
            if !updated_vectors.is_empty() {
                if let Err(e) = storage.upsert(&namespace, updated_vectors).await {
                    tracing::warn!(
                        namespace = %namespace,
                        error = %e,
                        "Failed to upsert decayed memories"
                    );
                }
            }
        }

        tracing::info!(
            namespaces_processed = result.namespaces_processed,
            memories_processed = result.memories_processed,
            memories_decayed = result.memories_decayed,
            memories_deleted = result.memories_deleted,
            "Decay cycle completed"
        );

        result
    }

    /// Spawn the decay engine as a background tokio task.
    pub fn spawn(
        config: DecayEngineConfig,
        storage: Arc<dyn VectorStorage>,
        metrics: Arc<BackgroundMetrics>,
    ) -> tokio::task::JoinHandle<()> {
        let engine = DecayEngine::new(config.decay_config);
        let interval = std::time::Duration::from_secs(config.interval_secs);

        tokio::spawn(async move {
            tracing::info!(
                strategy = ?engine.config.strategy,
                half_life_hours = engine.config.half_life_hours,
                min_importance = engine.config.min_importance,
                interval_secs = interval.as_secs(),
                "Decay engine started"
            );

            loop {
                tokio::time::sleep(interval).await;
                let result = engine.apply_decay(&storage).await;
                metrics.record_decay(&result);
            }
        })
    }
}

/// Result of a decay cycle.
#[derive(Debug, Default, Clone, Serialize, Deserialize)]
pub struct DecayResult {
    pub namespaces_processed: usize,
    pub memories_processed: usize,
    pub memories_decayed: usize,
    pub memories_deleted: usize,
}

/// Shared metrics for background activity tracking.
///
/// Updated by decay/autopilot spawn loops so the API can expose
/// what background jobs are doing to user memories.
/// Persisted to storage so metrics survive restarts.
#[derive(Debug, Default)]
pub struct BackgroundMetrics {
    inner: std::sync::Mutex<BackgroundMetricsInner>,
    /// Flag: dirty since last persist
    dirty: std::sync::atomic::AtomicBool,
}

/// Max history data points kept (7 days at 1-hour decay interval).
const MAX_HISTORY_POINTS: usize = 168;

#[derive(Debug, Default, Clone, Serialize, Deserialize)]
pub struct BackgroundMetricsInner {
    /// Last decay cycle result
    #[serde(default)]
    pub last_decay: Option<DecayResult>,
    /// Timestamp of last decay run (unix secs)
    #[serde(default)]
    pub last_decay_at: Option<u64>,
    /// Cumulative memories deleted by decay
    #[serde(default)]
    pub total_decay_deleted: u64,
    /// Cumulative memories decayed (importance lowered)
    #[serde(default)]
    pub total_decay_adjusted: u64,

    /// Last dedup cycle result
    #[serde(default)]
    pub last_dedup: Option<DedupResultSnapshot>,
    /// Timestamp of last dedup run (unix secs)
    #[serde(default)]
    pub last_dedup_at: Option<u64>,
    /// Cumulative duplicates removed
    #[serde(default)]
    pub total_dedup_removed: u64,

    /// Last consolidation cycle result
    #[serde(default)]
    pub last_consolidation: Option<ConsolidationResultSnapshot>,
    /// Timestamp of last consolidation run (unix secs)
    #[serde(default)]
    pub last_consolidation_at: Option<u64>,
    /// Cumulative memories consolidated
    #[serde(default)]
    pub total_consolidated: u64,

    /// Historical data points for graphing (ring buffer, newest last)
    #[serde(default)]
    pub history: Vec<ActivityHistoryPoint>,
}

/// A single historical data point for the activity timeline graph.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ActivityHistoryPoint {
    /// Unix timestamp (seconds)
    pub timestamp: u64,
    /// Memories deleted by decay in this cycle
    pub decay_deleted: u64,
    /// Memories adjusted by decay in this cycle
    pub decay_adjusted: u64,
    /// Duplicates removed in this cycle
    pub dedup_removed: u64,
    /// Memories consolidated in this cycle
    pub consolidated: u64,
}

/// Serializable snapshot of a dedup result (avoids coupling to autopilot module).
#[derive(Debug, Default, Clone, Serialize, Deserialize)]
pub struct DedupResultSnapshot {
    pub namespaces_processed: usize,
    pub memories_scanned: usize,
    pub duplicates_removed: usize,
}

/// Serializable snapshot of a consolidation result.
#[derive(Debug, Default, Clone, Serialize, Deserialize)]
pub struct ConsolidationResultSnapshot {
    pub namespaces_processed: usize,
    pub memories_scanned: usize,
    pub clusters_merged: usize,
    pub memories_consolidated: usize,
}

impl BackgroundMetrics {
    pub fn new() -> Self {
        Self::default()
    }

    /// Restore metrics from a previously persisted snapshot.
    pub fn restore(inner: BackgroundMetricsInner) -> Self {
        Self {
            inner: std::sync::Mutex::new(inner),
            dirty: std::sync::atomic::AtomicBool::new(false),
        }
    }

    /// Whether metrics changed since last persist.
    pub fn is_dirty(&self) -> bool {
        self.dirty.load(std::sync::atomic::Ordering::Relaxed)
    }

    /// Clear the dirty flag after a successful persist.
    pub fn clear_dirty(&self) {
        self.dirty
            .store(false, std::sync::atomic::Ordering::Relaxed);
    }

    /// Record a decay cycle result.
    pub fn record_decay(&self, result: &DecayResult) {
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        inner.total_decay_deleted += result.memories_deleted as u64;
        inner.total_decay_adjusted += result.memories_decayed as u64;
        inner.last_decay = Some(result.clone());
        inner.last_decay_at = Some(now);
        // Append history point
        push_history(
            &mut inner.history,
            ActivityHistoryPoint {
                timestamp: now,
                decay_deleted: result.memories_deleted as u64,
                decay_adjusted: result.memories_decayed as u64,
                dedup_removed: 0,
                consolidated: 0,
            },
        );
        self.dirty.store(true, std::sync::atomic::Ordering::Relaxed);
    }

    /// Record a dedup cycle result.
    pub fn record_dedup(
        &self,
        namespaces_processed: usize,
        memories_scanned: usize,
        duplicates_removed: usize,
    ) {
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        inner.total_dedup_removed += duplicates_removed as u64;
        inner.last_dedup = Some(DedupResultSnapshot {
            namespaces_processed,
            memories_scanned,
            duplicates_removed,
        });
        inner.last_dedup_at = Some(now);
        push_history(
            &mut inner.history,
            ActivityHistoryPoint {
                timestamp: now,
                decay_deleted: 0,
                decay_adjusted: 0,
                dedup_removed: duplicates_removed as u64,
                consolidated: 0,
            },
        );
        self.dirty.store(true, std::sync::atomic::Ordering::Relaxed);
    }

    /// Record a consolidation cycle result.
    pub fn record_consolidation(
        &self,
        namespaces_processed: usize,
        memories_scanned: usize,
        clusters_merged: usize,
        memories_consolidated: usize,
    ) {
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        inner.total_consolidated += memories_consolidated as u64;
        inner.last_consolidation = Some(ConsolidationResultSnapshot {
            namespaces_processed,
            memories_scanned,
            clusters_merged,
            memories_consolidated,
        });
        inner.last_consolidation_at = Some(now);
        push_history(
            &mut inner.history,
            ActivityHistoryPoint {
                timestamp: now,
                decay_deleted: 0,
                decay_adjusted: 0,
                dedup_removed: 0,
                consolidated: memories_consolidated as u64,
            },
        );
        self.dirty.store(true, std::sync::atomic::Ordering::Relaxed);
    }

    /// Replace the inner state with a restored snapshot (used on startup).
    pub fn restore_into(&self, restored: BackgroundMetricsInner) {
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        *inner = restored;
        // Don't set dirty — this was just loaded from storage
    }

    /// Get a snapshot of all metrics for API response.
    pub fn snapshot(&self) -> BackgroundMetricsInner {
        self.inner.lock().unwrap_or_else(|e| e.into_inner()).clone()
    }
}

/// Append a history point, capping at MAX_HISTORY_POINTS.
fn push_history(history: &mut Vec<ActivityHistoryPoint>, point: ActivityHistoryPoint) {
    history.push(point);
    if history.len() > MAX_HISTORY_POINTS {
        let excess = history.len() - MAX_HISTORY_POINTS;
        history.drain(..excess);
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn make_engine(strategy: DecayStrategy, half_life: f64) -> DecayEngine {
        DecayEngine::new(DecayConfig {
            strategy,
            half_life_hours: half_life,
            min_importance: 0.01,
        })
    }

    // Default args: Episodic type, 0 access count (same as old behavior with multiplier=1.0, shield=1.5)
    const EPISODIC: MemoryType = MemoryType::Episodic;

    #[test]
    fn test_exponential_decay_at_half_life_episodic_no_access() {
        let engine = make_engine(DecayStrategy::Exponential, 168.0);
        // Effective half-life = 168 / (1.0 * 1.5) = 112h (faster for never-accessed)
        let result = engine.calculate_decay(1.0, 112.0, &EPISODIC, 0);
        assert!((result - 0.5).abs() < 0.01, "Expected ~0.5, got {}", result);
    }

    #[test]
    fn test_exponential_decay_zero_time() {
        let engine = make_engine(DecayStrategy::Exponential, 168.0);
        let result = engine.calculate_decay(0.8, 0.0, &EPISODIC, 0);
        assert!((result - 0.8).abs() < 0.001);
    }

    #[test]
    fn test_linear_decay_floors_at_zero() {
        let engine = make_engine(DecayStrategy::Linear, 168.0);
        let result = engine.calculate_decay(0.3, 168.0, &EPISODIC, 0);
        assert!(result >= 0.0, "Should not go below 0, got {}", result);
    }

    #[test]
    fn test_procedural_decays_slower_than_working() {
        let engine = make_engine(DecayStrategy::Exponential, 168.0);
        let working = engine.calculate_decay(1.0, 168.0, &MemoryType::Working, 0);
        let procedural = engine.calculate_decay(1.0, 168.0, &MemoryType::Procedural, 0);
        assert!(
            procedural > working,
            "Procedural ({}) should decay slower than Working ({})",
            procedural,
            working
        );
    }

    #[test]
    fn test_high_access_count_decays_slower() {
        let engine = make_engine(DecayStrategy::Exponential, 168.0);
        let no_access = engine.calculate_decay(1.0, 168.0, &EPISODIC, 0);
        let high_access = engine.calculate_decay(1.0, 168.0, &EPISODIC, 10);
        assert!(
            high_access > no_access,
            "High access ({}) should decay slower than no access ({})",
            high_access,
            no_access
        );
    }

    #[test]
    fn test_semantic_decays_slower_than_episodic() {
        let engine = make_engine(DecayStrategy::Exponential, 168.0);
        let episodic = engine.calculate_decay(1.0, 168.0, &EPISODIC, 5);
        let semantic = engine.calculate_decay(1.0, 168.0, &MemoryType::Semantic, 5);
        assert!(
            semantic > episodic,
            "Semantic ({}) should decay slower than Episodic ({})",
            semantic,
            episodic
        );
    }

    #[test]
    fn test_access_boost_scales_with_importance() {
        let low = DecayEngine::access_boost(0.2);
        let high = DecayEngine::access_boost(0.8);
        let boost_low = low - 0.2;
        let boost_high = high - 0.8;
        assert!(
            boost_high > boost_low,
            "High-importance boost ({}) should be larger than low-importance boost ({})",
            boost_high,
            boost_low
        );
        // Verify range: 0.05 + 0.05*importance
        assert!((boost_low - (0.05 + 0.05 * 0.2)).abs() < 0.001);
        assert!((boost_high - (0.05 + 0.05 * 0.8)).abs() < 0.001);
    }

    #[test]
    fn test_access_boost_caps_at_one() {
        assert!((DecayEngine::access_boost(1.0) - 1.0).abs() < 0.001);
        assert!((DecayEngine::access_boost(0.96) - 1.0).abs() < 0.001);
    }

    #[test]
    fn test_decay_clamps_to_range() {
        let engine = make_engine(DecayStrategy::Exponential, 1.0);
        let result = engine.calculate_decay(0.001, 100.0, &EPISODIC, 0);
        assert!(result >= 0.0 && result <= 1.0);
    }

    #[test]
    fn test_step_function_decay() {
        let engine = make_engine(DecayStrategy::StepFunction, 168.0);
        // Effective half-life for Episodic+0 access = 168/1.5 = 112h
        let eff_hl = 168.0 / 1.5;

        // Before first step - no decay
        let result = engine.calculate_decay(1.0, eff_hl * 0.5, &EPISODIC, 0);
        assert!((result - 1.0).abs() < 0.001);

        // At first step
        let result = engine.calculate_decay(1.0, eff_hl, &EPISODIC, 0);
        assert!((result - 0.5).abs() < 0.001);
    }
}