vtcode_core/tools/registry/

risk_scorer.rs

//! Risk scoring system for tool execution

use crate::config::constants::tools;
use crate::utils::ansi_codes::{FG_GREEN, FG_MAGENTA, FG_RED, FG_YELLOW};
use serde::{Deserialize, Serialize};

/// Risk level classification for tools
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
pub enum RiskLevel {
    /// Read-only operations with no side effects
    Low,

    /// Operations that create/modify data but within trusted boundaries
    Medium,

    /// Operations with potentially destructive effects or external access
    High,

    /// Operations that could compromise system security
    Critical,
}

impl RiskLevel {
    pub fn as_str(self) -> &'static str {
        match self {
            Self::Low => "low",
            Self::Medium => "medium",
            Self::High => "high",
            Self::Critical => "critical",
        }
    }

    pub fn color_code(self) -> &'static str {
        match self {
            Self::Low => FG_GREEN,
            Self::Medium => FG_YELLOW,
            Self::High => FG_RED,
            Self::Critical => FG_MAGENTA,
        }
    }
}

impl std::fmt::Display for RiskLevel {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(self.as_str())
    }
}

/// Source of the tool (internal, MCP, ACP, etc.)
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ToolSource {
    /// Built-in tools
    Internal,

    /// Model Context Protocol (external)
    Mcp,

    /// Agent Client Protocol (IDE integration)
    Acp,

    /// Other external sources
    External,
}

impl ToolSource {
    /// Get the risk multiplier for this source
    /// MCP/external tools are considered higher risk
    pub fn risk_multiplier(self) -> f32 {
        match self {
            Self::Internal => 1.0,
            Self::Mcp => 1.5,
            Self::Acp => 1.2,
            Self::External => 2.0,
        }
    }
}

/// Workspace trust level
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum WorkspaceTrust {
    Untrusted,
    Partial,
    Trusted,
    FullAuto,
}

impl WorkspaceTrust {
    /// Get the risk reduction multiplier for trusted workspaces
    pub fn risk_reduction(self) -> f32 {
        match self {
            Self::Untrusted => 1.0,
            Self::Partial => 0.8,
            Self::Trusted => 0.6,
            Self::FullAuto => 0.3,
        }
    }
}

/// Context for risk assessment
#[derive(Debug, Clone)]
pub struct ToolRiskContext {
    /// Tool name
    pub tool_name: String,

    /// Source of the tool
    pub source: ToolSource,

    /// Workspace trust level
    pub workspace_trust: WorkspaceTrust,

    /// Number of times this tool has been approved recently
    pub recent_approvals: usize,

    /// Command arguments (if applicable)
    pub command_args: Vec<String>,

    /// Whether this is a write operation
    pub is_write: bool,

    /// Whether this is a potentially destructive operation
    pub is_destructive: bool,

    /// Whether this accesses external network
    pub accesses_network: bool,
}

impl ToolRiskContext {
    /// Create a new risk context
    pub fn new(tool_name: String, source: ToolSource, workspace_trust: WorkspaceTrust) -> Self {
        Self {
            tool_name,
            source,
            workspace_trust,
            recent_approvals: 0,
            command_args: Vec::new(),
            is_write: false,
            is_destructive: false,
            accesses_network: false,
        }
    }

    /// Set command arguments
    pub fn with_args(mut self, args: Vec<String>) -> Self {
        self.command_args = args;
        self
    }

    /// Mark as write operation
    pub fn as_write(mut self) -> Self {
        self.is_write = true;
        self
    }

    /// Mark as potentially destructive
    pub fn as_destructive(mut self) -> Self {
        self.is_destructive = true;
        self
    }

    /// Mark as network-accessing
    pub fn accesses_network(mut self) -> Self {
        self.accesses_network = true;
        self
    }
}

/// Risk scorer for tool execution
pub struct ToolRiskScorer;

impl ToolRiskScorer {
    /// Calculate risk level for a tool
    pub fn calculate_risk(ctx: &ToolRiskContext) -> RiskLevel {
        let mut base_score = Self::base_risk_for_tool(&ctx.tool_name);

        // Apply modifiers
        if ctx.is_destructive {
            base_score += 30;
        }
        if ctx.is_write {
            base_score += 15;
        }
        if ctx.accesses_network {
            base_score += 10;
        }

        // Apply source multiplier
        #[allow(
            clippy::cast_sign_loss,
            clippy::let_and_return,
            reason = "Intentional compatibility, platform, or test-only suppression."
        )]
        let adjusted = ((base_score as f32 * ctx.source.risk_multiplier()).max(0.0)) as u32;
        base_score = adjusted;

        // Apply trust reduction
        #[allow(
            clippy::cast_sign_loss,
            clippy::let_and_return,
            reason = "Intentional compatibility, platform, or test-only suppression."
        )]
        let adjusted = ((base_score as f32 * ctx.workspace_trust.risk_reduction()).max(0.0)) as u32;
        base_score = adjusted;

        // Approval history reduces risk (diminishing returns)
        let approval_reduction = ctx.recent_approvals.min(3) as u32 * 5;
        base_score = base_score.saturating_sub(approval_reduction);

        // Convert to risk level
        match base_score {
            0..=25 => RiskLevel::Low,
            26..=50 => RiskLevel::Medium,
            51..=75 => RiskLevel::High,
            _ => RiskLevel::Critical,
        }
    }

    /// Whether a tool name performs outbound network access.
    ///
    /// Centralizes the network-tool set so policy auto-approval, risk scoring,
    /// and the safety gateway stay in agreement.
    pub fn is_network_tool(tool_name: &str) -> bool {
        matches!(
            tool_name,
            tools::WEB_SEARCH
                | tools::WEB_FETCH
                | tools::FETCH_URL
                | tools::DEFUDDLE_FETCH
                | tools::MCP_CONNECT_SERVER
                | tools::MCP_DISCONNECT_SERVER
        ) || matches!(tool_name, "mcp:connect" | "mcp:disconnect")
    }

    /// Determine if justification is required
    pub fn requires_justification(risk: RiskLevel, threshold: RiskLevel) -> bool {
        risk >= threshold
    }

    /// Base risk score for common tools
    fn base_risk_for_tool(tool_name: &str) -> u32 {
        match tool_name {
            // Read-only tools (base: 0)
            tools::READ_FILE
            | tools::CODE_SEARCH
            | tools::MCP_SEARCH_TOOLS
            | tools::MCP_GET_TOOL_DETAILS
            | tools::MCP_LIST_SERVERS => 0,

            // Safe metadata tools (base: 5)
            "file_info" | "status" | "logs" => 5,

            // Write tools (base: 20)
            tools::WRITE_FILE | tools::EDIT_FILE | tools::CREATE_FILE => 20,

            // Potentially risky write operations (base: 25)
            tools::APPLY_PATCH | tools::DELETE_FILE => 25,

            // Skill loading can cross from untrusted model-selected content
            // into executable or native-backed implementations. Keep it out
            // of the low-risk auto-approval path even in trusted workspaces.
            tools::LOAD_SKILL => 60,

            // PTY/interactive commands (base: 35)
            tools::CREATE_PTY_SESSION | tools::RUN_PTY_CMD | tools::SEND_PTY_INPUT | tools::UNIFIED_EXEC => 35,

            // Network operations (base: 40)
            tools::WEB_SEARCH
            | tools::WEB_FETCH
            | tools::FETCH_URL
            | tools::DEFUDDLE_FETCH
            | tools::MCP_CONNECT_SERVER
            | tools::MCP_DISCONNECT_SERVER
            | "mcp:connect"
            | "mcp:disconnect" => 40,

            // MCP tools (default to medium risk)
            _ if tool_name.starts_with("mcp_") => 30,

            // Unknown tools default to medium-high risk
            _ => 35,
        }
    }
}

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

    #[test]
    fn test_risk_level_ordering() {
        assert!(RiskLevel::Low < RiskLevel::Medium);
        assert!(RiskLevel::Medium < RiskLevel::High);
        assert!(RiskLevel::High < RiskLevel::Critical);
    }

    #[test]
    fn test_risk_calculation() {
        // Read-only operation in trusted workspace
        let ctx = ToolRiskContext::new(tools::READ_FILE.to_string(), ToolSource::Internal, WorkspaceTrust::Trusted);
        let risk = ToolRiskScorer::calculate_risk(&ctx);
        assert_eq!(risk, RiskLevel::Low);

        // Write operation in untrusted workspace
        let ctx = ToolRiskContext::new(tools::WRITE_FILE.to_string(), ToolSource::External, WorkspaceTrust::Untrusted)
            .as_write();
        let risk = ToolRiskScorer::calculate_risk(&ctx);
        assert!(risk >= RiskLevel::High);
    }

    #[test]
    fn test_network_tools_are_classified() {
        assert!(ToolRiskScorer::is_network_tool(tools::WEB_FETCH));
        assert!(ToolRiskScorer::is_network_tool(tools::WEB_SEARCH));
        assert!(ToolRiskScorer::is_network_tool(tools::FETCH_URL));
        assert!(ToolRiskScorer::is_network_tool(tools::DEFUDDLE_FETCH));
        assert!(ToolRiskScorer::is_network_tool("mcp:connect"));
        assert!(ToolRiskScorer::is_network_tool("mcp:disconnect"));
        assert!(!ToolRiskScorer::is_network_tool(tools::CODE_SEARCH));
        assert!(!ToolRiskScorer::is_network_tool(tools::READ_FILE));
    }

    #[test]
    fn test_network_fetch_not_low_risk_even_when_trusted() {
        // Mirrors the policy auto-approval path: trusted workspace + network
        // flag. Must stay above Low so HITL approval is required.
        for tool in [tools::WEB_FETCH, "mcp:connect", "mcp:disconnect"] {
            let ctx = ToolRiskContext::new(tool.to_string(), ToolSource::Internal, WorkspaceTrust::Trusted)
                .accesses_network();
            let risk = ToolRiskScorer::calculate_risk(&ctx);
            assert!(
                risk > RiskLevel::Low,
                "network tool '{tool}' should not be auto-approved as low risk, got {risk:?}"
            );
        }
    }

    #[test]
    fn test_skill_loading_is_not_low_risk_even_when_trusted() {
        let ctx = ToolRiskContext::new(tools::LOAD_SKILL.to_string(), ToolSource::Internal, WorkspaceTrust::Trusted);
        let risk = ToolRiskScorer::calculate_risk(&ctx);

        assert!(risk > RiskLevel::Low, "skill loading should require approval, got {risk:?}");
    }

    #[test]
    fn test_approval_history_reduces_risk() {
        let mut ctx =
            ToolRiskContext::new(tools::RUN_PTY_CMD.to_string(), ToolSource::Internal, WorkspaceTrust::Untrusted);

        let risk_before = ToolRiskScorer::calculate_risk(&ctx);

        ctx.recent_approvals = 3;
        let risk_after = ToolRiskScorer::calculate_risk(&ctx);

        assert!(risk_after <= risk_before);
    }

    #[test]
    fn test_source_multiplier() {
        let base = ToolRiskContext::new("mcp_tool".to_string(), ToolSource::Internal, WorkspaceTrust::Trusted);
        let base_risk = ToolRiskScorer::calculate_risk(&base);

        let mcp = ToolRiskContext::new("mcp_tool".to_string(), ToolSource::Mcp, WorkspaceTrust::Trusted);
        let mcp_risk = ToolRiskScorer::calculate_risk(&mcp);

        // MCP tool should have higher risk
        assert!(mcp_risk > base_risk || mcp_risk == RiskLevel::Critical);
    }

    #[test]
    fn test_requires_justification() {
        assert!(ToolRiskScorer::requires_justification(RiskLevel::High, RiskLevel::High));
        assert!(!ToolRiskScorer::requires_justification(RiskLevel::Medium, RiskLevel::High));
    }
}