matrixcode_core/agent/context/

manager.rs

//! Agent context management.
//!
//! This module manages the context of the Agent, including:
//! - Skills and project overview
//! - Memory summary
//! - System prompt building
//! - CodeGraph tool management
//!
//! By extracting context into a dedicated struct, we enable:
//! - Clear separation between context and state
//! - Easier testing of prompt building
//! - Better management of project-specific information

use std::path::PathBuf;

use crate::prompt::PromptProfile;
use crate::skills::Skill;

/// Agent context information.
///
/// Manages all context-related information that influences agent behavior.
/// This includes project overview, memory, skills, and system prompts.
pub struct AgentContext {
    /// Skills available to the agent.
    skills: Vec<Skill>,

    /// Project overview (generated by /init).
    project_overview: Option<String>,

    /// Memory summary from previous conversations.
    memory_summary: Option<String>,

    /// Project root path (for CodeGraph and file operations).
    project_path: Option<PathBuf>,

    /// System prompt sent to LLM.
    system_prompt: String,

    /// Prompt profile (defines agent personality and rules).
    profile: PromptProfile,
}

impl AgentContext {
    /// Create a new context with given profile.
    pub fn new(profile: PromptProfile) -> Self {
        let system_prompt = crate::prompt::build_system_prompt(
            &profile,
            &[], // No skills initially
            None,
            None,
        );

        Self {
            skills: Vec::new(),
            project_overview: None,
            memory_summary: None,
            project_path: None,
            system_prompt,
            profile,
        }
    }

    /// Create a new context with skills and overview.
    pub fn with_context(
        profile: PromptProfile,
        skills: Vec<Skill>,
        project_overview: Option<String>,
        memory_summary: Option<String>,
        project_path: Option<PathBuf>,
    ) -> Self {
        let system_prompt = crate::prompt::build_system_prompt(
            &profile,
            &skills,
            project_overview.as_deref(),
            memory_summary.as_deref(),
        );

        Self {
            skills,
            project_overview,
            memory_summary,
            project_path,
            system_prompt,
            profile,
        }
    }

    /// Get reference to skills.
    pub fn skills(&self) -> &[Skill] {
        &self.skills
    }

    /// Set skills (rebuilds system prompt).
    pub fn set_skills(&mut self, skills: Vec<Skill>) {
        self.skills = skills;
        self.rebuild_system_prompt();
    }

    /// Get project overview.
    pub fn project_overview(&self) -> Option<&str> {
        self.project_overview.as_deref()
    }

    /// Set project overview (rebuilds system prompt).
    pub fn set_project_overview(&mut self, overview: Option<String>) {
        self.project_overview = overview;
        self.rebuild_system_prompt();
    }

    /// Get memory summary.
    pub fn memory_summary(&self) -> Option<&str> {
        self.memory_summary.as_deref()
    }

    /// Update memory summary (rebuilds system prompt).
    pub fn update_memory(&mut self, summary: Option<String>) {
        self.memory_summary = summary;
        self.rebuild_system_prompt();
    }

    /// Get project path.
    pub fn project_path(&self) -> Option<&PathBuf> {
        self.project_path.as_ref()
    }

    /// Set project path (may trigger CodeGraph injection).
    pub fn set_project_path(&mut self, path: Option<PathBuf>) {
        self.project_path = path;
        // Note: CodeGraph injection will be handled in Phase 4
        // when we integrate with tool management
    }

    /// Get system prompt.
    pub fn system_prompt(&self) -> &str {
        &self.system_prompt
    }

    /// Get prompt profile.
    pub fn profile(&self) -> &PromptProfile {
        &self.profile
    }

    /// Rebuild system prompt with current context.
    fn rebuild_system_prompt(&mut self) {
        self.system_prompt = crate::prompt::build_system_prompt(
            &self.profile,
            &self.skills,
            self.project_overview.as_deref(),
            self.memory_summary.as_deref(),
        );
    }

    /// Set system prompt directly (for workflows with project_path).
    pub fn set_system_prompt(&mut self, prompt: String) {
        self.system_prompt = prompt;
    }

    /// Rebuild system prompt with workflows (includes project_path for CodeGraph).
    pub fn rebuild_system_prompt_with_workflows(&mut self, project_path: Option<PathBuf>) {
        self.system_prompt = crate::prompt::build_system_prompt_with_workflows(
            &self.profile,
            &self.skills,
            self.project_overview.as_deref(),
            self.memory_summary.as_deref(),
            project_path.as_ref(),
            None, // LSP servers not available in agent context
        );
    }

    /// Clear all context (reset to initial state).
    pub fn clear(&mut self) {
        self.skills.clear();
        self.project_overview = None;
        self.memory_summary = None;
        self.project_path = None;
        self.rebuild_system_prompt();
    }
}

impl Default for AgentContext {
    fn default() -> Self {
        Self::new(PromptProfile::default())
    }
}

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

    #[test]
    fn test_context_new_creates_empty_context() {
        let context = AgentContext::new(PromptProfile::default());

        assert_eq!(context.skills().len(), 0);
        assert!(context.project_overview().is_none());
        assert!(context.memory_summary().is_none());
        assert!(context.project_path().is_none());
        assert!(!context.system_prompt().is_empty(), "system prompt should not be empty");
    }

    #[test]
    fn test_context_set_skills_updates_prompt() {
        let mut context = AgentContext::new(PromptProfile::default());
        let initial_prompt = context.system_prompt().to_string();

        // Add a skill
        let skills = vec![Skill {
            name: "test_skill".to_string(),
            description: "Test skill description".to_string(),
            trigger: Some("/test".to_string()),
            skill_type: crate::skills::SkillType::Flexible,
            priority: crate::skills::SkillPriority::Implementation,
            mandatory: false,
            dir: PathBuf::from("/skills/test_skill"),
            body: "Test skill content".to_string(),
            source_file: PathBuf::from("/skills/test_skill/SKILL.md"),
        }];
        context.set_skills(skills);

        // System prompt should change
        assert_ne!(context.system_prompt(), initial_prompt, "prompt should change when skills added");
        assert!(context.system_prompt().contains("test_skill"), "prompt should include skill name");
        assert!(context.system_prompt().contains("Test skill description"), "prompt should include skill description");
    }

    #[test]
    fn test_context_update_memory_rebuilds_prompt() {
        let mut context = AgentContext::new(PromptProfile::default());
        let initial_prompt = context.system_prompt().to_string();

        // Update memory
        context.update_memory(Some("Previous conversation summary".to_string()));

        // System prompt should change
        assert_ne!(context.system_prompt(), initial_prompt, "prompt should change when memory updated");
        assert!(context.system_prompt().contains("Previous conversation summary"), "prompt should include memory");
    }

    #[test]
    fn test_context_set_project_overview() {
        let mut context = AgentContext::new(PromptProfile::default());

        context.set_project_overview(Some("Project overview text".to_string()));

        assert_eq!(context.project_overview(), Some("Project overview text"));
    }

    #[test]
    fn test_context_clear_resets_state() {
        let mut context = AgentContext::new(PromptProfile::default());

        // Add some context
        context.set_skills(vec![Skill {
            name: "skill".to_string(),
            description: "description".to_string(),
            trigger: Some("trigger".to_string()),
            skill_type: crate::skills::SkillType::Flexible,
            priority: crate::skills::SkillPriority::Implementation,
            mandatory: false,
            dir: PathBuf::from("/skills/skill"),
            body: "content".to_string(),
            source_file: PathBuf::from("/skills/skill/SKILL.md"),
        }]);
        context.set_project_overview(Some("overview".to_string()));
        context.update_memory(Some("memory".to_string()));
        context.set_project_path(Some(PathBuf::from("/path")));

        // Clear
        context.clear();

        // Verify all cleared
        assert_eq!(context.skills().len(), 0);
        assert!(context.project_overview().is_none());
        assert!(context.memory_summary().is_none());
        assert!(context.project_path().is_none());
    }

    #[test]
    fn test_context_with_context_creates_full_context() {
        let skills = vec![Skill {
            name: "skill".to_string(),
            description: "skill description".to_string(),
            trigger: Some("trigger".to_string()),
            skill_type: crate::skills::SkillType::Flexible,
            priority: crate::skills::SkillPriority::Implementation,
            mandatory: false,
            dir: PathBuf::from("/skills/skill"),
            body: "skill content".to_string(),
            source_file: PathBuf::from("/skills/skill/SKILL.md"),
        }];
        let overview = Some("Project overview".to_string());
        let memory = Some("Memory summary".to_string());
        let path = Some(PathBuf::from("/project"));

        let context = AgentContext::with_context(
            PromptProfile::default(),
            skills.clone(),
            overview.clone(),
            memory.clone(),
            path.clone(),
        );

        assert_eq!(context.skills(), &skills);
        assert_eq!(context.project_overview(), overview.as_deref());
        assert_eq!(context.memory_summary(), memory.as_deref());
        assert_eq!(context.project_path(), path.as_ref());
    }
}