vtcode_core/tools/registry/

justification.rs

/// Tool Justification System
///
/// Captures agent reasoning before high-risk tool execution to improve approval UX
/// and enable learning of approval patterns.
use crate::tools::registry::risk_scorer::RiskLevel;
use crate::utils::file_utils::{
    ensure_dir_exists_sync, read_file_with_context_sync, write_file_with_context_sync,
};
use anyhow::Result;
use hashbrown::HashMap;
use serde::{Deserialize, Serialize};
use std::path::PathBuf;

/// Justification provided by the agent for executing a high-risk tool
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolJustification {
    /// Tool being justified
    pub tool_name: String,
    /// Brief explanation from the agent
    pub reason: String,
    /// Expected outcome of tool execution
    pub expected_outcome: Option<String>,
    /// Risk level that triggered justification
    pub risk_level: String,
    /// Timestamp when justification was provided
    pub timestamp: String,
}

impl ToolJustification {
    /// Create a new tool justification
    pub fn new(
        tool_name: impl Into<String>,
        reason: impl Into<String>,
        risk_level: &RiskLevel,
    ) -> Self {
        Self {
            tool_name: tool_name.into(),
            reason: reason.into(),
            expected_outcome: None,
            risk_level: format!("{:?}", risk_level),
            timestamp: chrono::Local::now().to_rfc3339(),
        }
    }

    /// Add expected outcome to justification
    pub fn with_outcome(mut self, outcome: impl Into<String>) -> Self {
        self.expected_outcome = Some(outcome.into());
        self
    }

    /// Format justification for display in approval dialog
    pub fn format_for_dialog(&self) -> Vec<String> {
        let mut lines = vec![];

        lines.push(String::new());
        lines.push("Agent Reasoning:".to_owned());

        // Wrap reason text if needed - iterate directly without collecting
        for line in self.reason.lines() {
            let wrapped = textwrap::fill(&format!("  {line}"), 78);
            for wrapped_line in wrapped.lines() {
                lines.push(wrapped_line.to_owned());
            }
        }

        if let Some(outcome) = &self.expected_outcome {
            lines.push(String::new());
            lines.push("Expected Outcome:".to_owned());
            let wrapped = textwrap::fill(&format!("  {outcome}"), 78);
            for wrapped_line in wrapped.lines() {
                lines.push(wrapped_line.to_owned());
            }
        }

        lines.push(String::new());
        lines.push(format!("Risk Level: {}", self.risk_level));

        lines
    }
}

/// Tracks approval patterns to learn from user decisions
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct ApprovalPattern {
    /// Stable approval key used for lookup and persistence
    pub tool_name: String,
    /// Human-readable label for prompts and summaries
    #[serde(default)]
    pub display_name: Option<String>,
    /// Number of times user approved
    pub approve_count: u32,
    /// Number of times user denied
    pub deny_count: u32,
    /// Last decision (true = approve, false = deny)
    pub last_decision: Option<bool>,
    /// Most recent reason (if available)
    pub recent_reason: Option<String>,
}

impl ApprovalPattern {
    /// Compute approval rate (0.0 to 1.0)
    pub fn approval_rate(&self) -> f32 {
        let total = self.approve_count + self.deny_count;
        if total == 0 {
            0.0
        } else {
            self.approve_count as f32 / total as f32
        }
    }

    /// Check if this tool has high approval rate (>80%)
    pub fn has_high_approval_rate(&self) -> bool {
        self.approval_count() >= 3 && self.approval_rate() > 0.8
    }

    /// Return approval count
    pub fn approval_count(&self) -> u32 {
        self.approve_count
    }

    pub fn display_name<'a>(&'a self, fallback: &'a str) -> &'a str {
        self.display_name.as_deref().unwrap_or(fallback)
    }
}

/// Merge an on-disk pattern into the in-memory entry by taking the max of
/// counters and preferring any non-`None` metadata from disk. Conservative:
/// undercount is safer (more prompts) than overcount (fewer prompts → risk).
fn merge_pattern_from_disk(local: &mut ApprovalPattern, disk: &ApprovalPattern) {
    local.approve_count = local.approve_count.max(disk.approve_count);
    local.deny_count = local.deny_count.max(disk.deny_count);
    if disk.display_name.is_some() {
        local.display_name = disk.display_name.clone();
    }
    if disk.last_decision.is_some() {
        local.last_decision = disk.last_decision;
    }
    if disk.recent_reason.is_some() {
        local.recent_reason = disk.recent_reason.clone();
    }
}

/// Manager for approval pattern learning and justifications
pub struct JustificationManager {
    cache_dir: PathBuf,
    patterns: std::sync::Arc<std::sync::Mutex<HashMap<String, ApprovalPattern>>>,
}

impl JustificationManager {
    /// Create a new justification manager
    pub fn new(cache_dir: PathBuf) -> Self {
        let patterns = std::sync::Arc::new(std::sync::Mutex::new(HashMap::new()));
        let manager = Self {
            cache_dir,
            patterns,
        };

        // Try to load existing patterns
        let _ = manager.load_patterns();

        manager
    }

    /// Load approval patterns from disk and merge into the in-memory map.
    ///
    /// Merging (rather than replacing) keeps in-memory increments that have not
    /// yet been flushed to disk — important when refreshing right before an
    /// auto-approval check while a concurrent vtcode session may have written
    /// newer counts to the same file.
    fn load_patterns(&self) -> Result<()> {
        let patterns_file = self.cache_dir.join("approval_patterns.json");
        if !patterns_file.exists() {
            return Ok(());
        }

        let content = read_file_with_context_sync(&patterns_file, "approval patterns cache")?;
        let loaded_patterns: HashMap<String, ApprovalPattern> = serde_json::from_str(&content)?;

        let mut patterns = self
            .patterns
            .lock()
            .map_err(|e| anyhow::anyhow!("Failed to lock patterns: {}", e))?;

        for (key, disk) in loaded_patterns {
            patterns
                .entry(key)
                .and_modify(|local| merge_pattern_from_disk(local, &disk))
                .or_insert(disk);
        }

        Ok(())
    }

    /// Re-read patterns from disk, merging with any in-memory state.
    pub fn refresh_patterns(&self) -> Result<()> {
        self.load_patterns()
    }

    /// Get approval pattern for a key
    pub fn get_pattern(&self, approval_key: &str) -> Option<ApprovalPattern> {
        if let Ok(patterns) = self.patterns.lock() {
            patterns.get(approval_key).cloned()
        } else {
            None
        }
    }

    /// Record user approval decision
    pub fn record_decision(
        &self,
        approval_key: &str,
        display_name: Option<&str>,
        approved: bool,
        reason: Option<String>,
    ) {
        let should_persist = if let Ok(mut patterns) = self.patterns.lock() {
            let pattern =
                patterns
                    .entry(approval_key.to_owned())
                    .or_insert_with(|| ApprovalPattern {
                        tool_name: approval_key.to_owned(),
                        display_name: display_name.map(str::to_owned),
                        approve_count: 0,
                        deny_count: 0,
                        last_decision: None,
                        recent_reason: None,
                    });

            if let Some(display_name) = display_name {
                pattern.display_name = Some(display_name.to_owned());
            }

            if approved {
                pattern.approve_count += 1;
            } else {
                pattern.deny_count += 1;
            }

            pattern.last_decision = Some(approved);
            pattern.recent_reason = reason;
            true
        } else {
            false
        };

        // Persist to disk after releasing the lock.
        if should_persist {
            let _ = self.persist_patterns();
        }
    }

    /// Persist patterns to disk
    fn persist_patterns(&self) -> Result<()> {
        ensure_dir_exists_sync(&self.cache_dir)?;
        let patterns_file = self.cache_dir.join("approval_patterns.json");
        let patterns = self
            .patterns
            .lock()
            .map_err(|e| anyhow::anyhow!("Failed to lock patterns: {}", e))?;
        let content = serde_json::to_string_pretty(&*patterns)?;
        write_file_with_context_sync(&patterns_file, &content, "approval patterns cache")?;
        Ok(())
    }

    /// Get learning summary for a key
    pub fn get_learning_summary(&self, approval_key: &str) -> Option<String> {
        let pattern = self.get_pattern(approval_key)?;

        if pattern.approval_count() == 0 {
            return None;
        }

        Some(format!(
            "Approved {} of {} times ({:.0}%)",
            pattern.approve_count,
            pattern.approve_count + pattern.deny_count,
            pattern.approval_rate() * 100.0
        ))
    }
}

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

    #[test]
    fn test_tool_justification_creation() {
        let just = ToolJustification::new(
            "read_file",
            "Need to understand code structure",
            &RiskLevel::Low,
        )
        .with_outcome("Will analyze the AST to provide better context");

        assert_eq!(just.tool_name, "read_file");
        assert!(just.reason.contains("understand"));
        assert!(just.expected_outcome.is_some());
    }

    #[test]
    fn test_justification_formatting() {
        let just = ToolJustification::new(
            "run_command",
            "Execute build to check for compilation errors",
            &RiskLevel::High,
        )
        .with_outcome("Will produce build output for analysis");

        let formatted = just.format_for_dialog();
        assert!(formatted.iter().any(|l| l.contains("Agent Reasoning")));
        assert!(formatted.iter().any(|l| l.contains("Expected Outcome")));
        assert!(formatted.iter().any(|l| l.contains("Risk Level")));
    }

    #[test]
    fn test_approval_pattern_calculation() {
        let mut pattern = ApprovalPattern {
            tool_name: "read_file".to_owned(),
            display_name: None,
            approve_count: 9,
            deny_count: 1,
            last_decision: Some(true),
            recent_reason: None,
        };

        assert_eq!(pattern.approval_rate(), 0.9);
        assert!(pattern.has_high_approval_rate());

        pattern.approve_count = 3;
        pattern.deny_count = 7;
        assert!(!pattern.has_high_approval_rate()); // < 0.8 rate
    }

    #[test]
    fn test_justification_manager_basic() {
        let temp_dir = std::env::temp_dir().join(format!("vtcode_test_{}", std::process::id()));
        let manager = JustificationManager::new(temp_dir.clone());

        manager.record_decision("read_file", Some("Read File"), true, None);
        manager.record_decision("read_file", Some("Read File"), true, None);
        manager.record_decision("read_file", Some("Read File"), false, None);

        let pattern = manager.get_pattern("read_file").unwrap();
        assert_eq!(pattern.approve_count, 2);
        assert_eq!(pattern.deny_count, 1);
        assert_eq!(pattern.approval_rate(), 2.0 / 3.0);
        assert_eq!(pattern.display_name.as_deref(), Some("Read File"));

        // Cleanup
        let _ = std::fs::remove_dir_all(&temp_dir);
    }

    #[test]
    fn test_justification_manager_preserves_new_display_name() {
        let temp_dir = std::env::temp_dir().join(format!("vtcode_test_{}", std::process::id()));
        let manager = JustificationManager::new(temp_dir.clone());

        manager.record_decision("shell:key", Some("command `cargo test`"), true, None);
        manager.record_decision(
            "shell:key",
            Some("commands starting with `cargo`"),
            true,
            None,
        );

        let pattern = manager.get_pattern("shell:key").unwrap();
        assert_eq!(
            pattern.display_name.as_deref(),
            Some("commands starting with `cargo`")
        );

        // Cleanup
        let _ = std::fs::remove_dir_all(&temp_dir);
    }
}