vv_agent/memory/session/

mod.rs

mod config;
mod entry;
mod parse;
mod prompt;
mod state;
mod storage;

use std::panic::{catch_unwind, AssertUnwindSafe};
use std::path::PathBuf;

use crate::memory::token_utils::estimate_tokens;
use crate::types::Message;

pub use config::{SessionMemoryConfig, SessionMemoryExtractionCallback};
pub use entry::SessionMemoryEntry;
use entry::{entry_key, SESSION_MEMORY_CATEGORIES};
use prompt::{build_extraction_prompt, should_skip_message};
pub use state::SessionMemoryState;

#[derive(Debug, Clone)]
pub struct SessionMemory {
    pub config: SessionMemoryConfig,
    pub state: SessionMemoryState,
    pub(super) workspace: Option<PathBuf>,
    pub(super) storage_scope: Option<String>,
}

impl SessionMemory {
    pub fn new(config: SessionMemoryConfig) -> Self {
        Self::with_workspace(config, None, None)
    }

    pub fn with_workspace(
        config: SessionMemoryConfig,
        workspace: Option<PathBuf>,
        storage_scope: Option<String>,
    ) -> Self {
        Self {
            config,
            state: SessionMemoryState::default(),
            workspace,
            storage_scope: storage_scope.map(|scope| scope.trim().to_string()),
        }
    }

    pub fn should_extract(&self, current_tokens: u64, message_count: usize) -> bool {
        if self.config.extraction_callback.is_none() || current_tokens == 0 || message_count == 0 {
            return false;
        }

        if !self.state.initialized {
            return current_tokens >= self.config.min_tokens_before_extraction
                && message_count >= self.config.min_text_messages;
        }

        let growth_threshold = ((self.config.min_tokens_before_extraction as f64)
            * self.config.growth_ratio)
            .floor()
            .max(1.0) as u64;
        let growth = if current_tokens >= self.state.tokens_at_last_extraction {
            current_tokens - self.state.tokens_at_last_extraction
        } else {
            current_tokens
        };
        growth >= growth_threshold
    }

    pub fn extract(
        &mut self,
        messages: &[Message],
        current_cycle: i32,
        current_tokens: u64,
    ) -> usize {
        let Some(callback) = self.config.extraction_callback.as_ref().cloned() else {
            return 0;
        };
        if messages.is_empty() {
            return 0;
        }

        let start_index = if self.state.last_extracted_message_index >= 0
            && (self.state.last_extracted_message_index as usize) < messages.len()
        {
            self.state.last_extracted_message_index as usize + 1
        } else {
            0
        };
        let new_messages = messages
            .iter()
            .enumerate()
            .filter_map(|(index, message)| {
                (index >= start_index && !should_skip_message(message)).then_some(message)
            })
            .collect::<Vec<_>>();

        if new_messages.is_empty() {
            self.record_extraction(messages.len() as i32 - 1, current_tokens);
            return 0;
        }

        let prompt = build_extraction_prompt(&new_messages);
        let raw_result = catch_unwind(AssertUnwindSafe(|| {
            callback(
                &prompt,
                self.config.extraction_backend.as_deref(),
                self.config.extraction_model.as_deref(),
            )
        }));
        let Some(raw_result) = raw_result.ok().flatten() else {
            return 0;
        };

        let entries = self.parse_extraction_result(&raw_result, current_cycle);
        let merged_count = self.merge_entries(entries);
        self.prune_to_budget();
        self.record_extraction(messages.len() as i32 - 1, current_tokens);
        self.save();
        merged_count
    }

    pub fn render_as_system_context(&self) -> String {
        if self.state.entries.is_empty() {
            return String::new();
        }

        let mut parts = vec!["<Session Memory>".to_string()];
        for category in SESSION_MEMORY_CATEGORIES {
            let entries = self
                .state
                .entries
                .iter()
                .filter(|entry| entry.category == *category)
                .collect::<Vec<_>>();
            if entries.is_empty() {
                continue;
            }
            parts.push(format!("## {category}"));
            for entry in entries {
                parts.push(format!("- {}", entry.content));
            }
        }
        parts.push("</Session Memory>".to_string());
        parts.join("\n")
    }

    pub fn on_compaction(&mut self, current_tokens: Option<u64>) {
        self.state.last_extracted_message_index = -1;
        if let Some(current_tokens) = current_tokens {
            self.state.tokens_at_last_extraction = current_tokens;
            self.state.initialized = true;
        }
        self.save();
    }

    pub fn load(&mut self) {
        let Some(path) = self.storage_path() else {
            return;
        };
        let Ok(content) = std::fs::read_to_string(path) else {
            return;
        };
        let Ok(state) = serde_json::from_str::<SessionMemoryState>(&content) else {
            return;
        };
        self.state = state;
    }

    pub fn save(&self) {
        let Some(path) = self.storage_path() else {
            return;
        };
        let Some(parent) = path.parent() else {
            return;
        };
        if std::fs::create_dir_all(parent).is_err() {
            return;
        }
        let Ok(content) = serde_json::to_string_pretty(&self.state) else {
            return;
        };
        let _ = std::fs::write(path, content);
    }

    pub fn merge_entries(&mut self, entries: Vec<SessionMemoryEntry>) -> usize {
        let mut merged = 0;
        for entry in entries {
            let key = entry_key(&entry);
            if let Some(existing) = self
                .state
                .entries
                .iter_mut()
                .find(|candidate| entry_key(candidate) == key)
            {
                existing.importance = existing.importance.max(entry.importance);
                existing.source_cycle = existing.source_cycle.max(entry.source_cycle);
                continue;
            }
            self.state.entries.push(entry);
            merged += 1;
        }
        merged
    }

    pub fn prune_to_budget(&mut self) {
        if self.config.max_tokens == 0 || self.state.entries.is_empty() {
            return;
        }
        let mut current_tokens =
            estimate_tokens(&self.render_as_system_context(), &self.config.token_model);
        while current_tokens > self.config.max_tokens && !self.state.entries.is_empty() {
            let Some((drop_index, _)) = self
                .state
                .entries
                .iter()
                .enumerate()
                .min_by_key(|(index, entry)| (entry.importance, entry.source_cycle, *index as i32))
            else {
                break;
            };
            self.state.entries.remove(drop_index);
            current_tokens =
                estimate_tokens(&self.render_as_system_context(), &self.config.token_model);
        }
    }

    fn record_extraction(&mut self, last_message_index: i32, current_tokens: u64) {
        self.state.last_extracted_message_index = last_message_index;
        self.state.tokens_at_last_extraction = current_tokens;
        self.state.initialized = true;
    }
}