zclaw_openfang/crates/zclaw-runtime/src/middleware/memory.rs

//! Memory middleware — unified pre/post hooks for memory retrieval and extraction.
//!
//! This middleware unifies the memory lifecycle:
//! - `before_completion`: retrieves relevant memories and injects them into the system prompt
//! - `after_completion`: extracts learnings from the conversation and stores them
//!
//! It replaces both the inline `GrowthIntegration` calls in `AgentLoop` and the
//! `intelligence_hooks` calls in the Tauri desktop layer.

use async_trait::async_trait;
use zclaw_types::Result;
use crate::growth::GrowthIntegration;
use crate::middleware::{AgentMiddleware, MiddlewareContext, MiddlewareDecision};
use crate::middleware::evolution::EvolutionMiddleware;

/// Middleware that handles memory retrieval (pre-completion) and extraction (post-completion).
///
/// Wraps `GrowthIntegration` and delegates:
/// - `before_completion` → `enhance_prompt()` for memory injection
/// - `after_completion` → `extract_combined()` for memory extraction + evolution check
pub struct MemoryMiddleware {
    growth: GrowthIntegration,
    /// Shared EvolutionMiddleware for pushing evolution suggestions
    evolution_mw: Option<std::sync::Arc<EvolutionMiddleware>>,
    /// Minimum seconds between extractions for the same agent (debounce).
    debounce_secs: u64,
    /// Timestamp of last extraction per agent (for debouncing).
    last_extraction: std::sync::Mutex<std::collections::HashMap<String, std::time::Instant>>,
}

impl MemoryMiddleware {
    pub fn new(growth: GrowthIntegration) -> Self {
        Self {
            growth,
            evolution_mw: None,
            debounce_secs: 30,
            last_extraction: std::sync::Mutex::new(std::collections::HashMap::new()),
        }
    }

    /// Attach a shared EvolutionMiddleware for pushing evolution suggestions.
    pub fn with_evolution(mut self, mw: std::sync::Arc<EvolutionMiddleware>) -> Self {
        self.evolution_mw = Some(mw);
        self
    }

    /// Set the debounce interval in seconds.
    pub fn with_debounce_secs(mut self, secs: u64) -> Self {
        self.debounce_secs = secs;
        self
    }

    /// Check if enough time has passed since the last extraction for this agent.
    fn should_extract(&self, agent_id: &str) -> bool {
        let now = std::time::Instant::now();
        let mut map = self.last_extraction.lock().unwrap_or_else(|e| e.into_inner());
        if let Some(last) = map.get(agent_id) {
            if now.duration_since(*last).as_secs() < self.debounce_secs {
                return false;
            }
        }
        map.insert(agent_id.to_string(), now);
        true
    }

    /// Check for evolvable patterns and push suggestions to EvolutionMiddleware.
    async fn check_and_push_evolution(&self, agent_id: &zclaw_types::AgentId) {
        let evolution_mw = match &self.evolution_mw {
            Some(mw) => mw,
            None => return,
        };

        match self.growth.check_evolution(agent_id).await {
            Ok(patterns) if !patterns.is_empty() => {
                for pattern in &patterns {
                    let trigger = pattern
                        .common_steps
                        .first()
                        .cloned()
                        .unwrap_or_else(|| pattern.pain_pattern.clone());
                    evolution_mw.add_pending(
                        crate::middleware::evolution::PendingEvolution {
                            pattern_name: pattern.pain_pattern.clone(),
                            trigger_suggestion: trigger,
                            description: format!(
                                "基于 {} 次重复经验，自动固化技能",
                                pattern.total_reuse
                            ),
                        },
                    ).await;
                }
                tracing::info!(
                    "[MemoryMiddleware] Pushed {} evolution candidates for agent {}",
                    patterns.len(),
                    agent_id
                );
            }
            Ok(_) => {
                tracing::debug!("[MemoryMiddleware] No evolvable patterns found");
            }
            Err(e) => {
                tracing::debug!(
                    "[MemoryMiddleware] Evolution check failed (non-fatal): {}", e
                );
            }
        }
    }
}

#[async_trait]
impl AgentMiddleware for MemoryMiddleware {
    fn name(&self) -> &str { "memory" }
    fn priority(&self) -> i32 { 150 }

    async fn before_completion(&self, ctx: &mut MiddlewareContext) -> Result<MiddlewareDecision> {
        tracing::debug!(
            "[MemoryMiddleware] before_completion for query: {:?}",
            ctx.user_input.chars().take(50).collect::<String>()
        );

        let base = &ctx.system_prompt;
        match self.growth.enhance_prompt(&ctx.agent_id, base, &ctx.user_input).await {
            Ok(enhanced) => {
                if enhanced != *base {
                    tracing::info!(
                        "[MemoryMiddleware] Injected memories into system prompt for agent {}",
                        ctx.agent_id
                    );
                    ctx.system_prompt = enhanced;
                } else {
                    tracing::debug!(
                        "[MemoryMiddleware] No relevant memories found for query: {:?}",
                        ctx.user_input.chars().take(50).collect::<String>()
                    );
                }
                Ok(MiddlewareDecision::Continue)
            }
            Err(e) => {
                tracing::warn!(
                    "[MemoryMiddleware] Memory retrieval failed (non-fatal): {}",
                    e
                );
                Ok(MiddlewareDecision::Continue)
            }
        }
    }

    async fn after_completion(&self, ctx: &MiddlewareContext) -> Result<()> {
        let agent_key = ctx.agent_id.to_string();
        if !self.should_extract(&agent_key) {
            tracing::debug!(
                "[MemoryMiddleware] Skipping extraction for agent {} (debounced)",
                agent_key
            );
            return Ok(());
        }

        if ctx.messages.is_empty() {
            return Ok(());
        }

        match self.growth.extract_combined(
            &ctx.agent_id,
            &ctx.messages,
            &ctx.session_id,
        ).await {
            Ok(Some((mem_count, facts))) => {
                tracing::info!(
                    "[MemoryMiddleware] Extracted {} memories + {} structured facts for agent {}",
                    mem_count,
                    facts.len(),
                    agent_key
                );

                // Check for evolvable patterns after successful extraction
                self.check_and_push_evolution(&ctx.agent_id).await;
            }
            Ok(None) => {
                tracing::debug!("[MemoryMiddleware] No memories or facts extracted");
            }
            Err(e) => {
                tracing::warn!("[MemoryMiddleware] Combined extraction failed: {}", e);
            }
        }

        Ok(())
    }
}