winx-code-agent 0.2.343

//! Per-principal MCP tool catalog policies.

use serde::Deserialize;

use crate::errors::{Result, WinxError};
use crate::tool_registry::ToolKind;

pub use crate::tool_registry::ALL_TOOL_NAMES;

/// Curated catalog shapes for common MCP clients and workflows.
#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, clap::ValueEnum)]
#[serde(rename_all = "kebab-case")]
pub enum ToolProfile {
    /// Advertise and permit every Winx tool.
    #[default]
    Full,
    /// Core repository exploration and editing without image or handoff helpers.
    Coding,
    /// Filesystem read/navigation tools only; shell and edit calls are rejected.
    ReadOnly,
    /// Durable terminal access with the smallest useful catalog.
    Terminal,
}

impl ToolProfile {
    const fn mask(self) -> u16 {
        match self {
            Self::Full => {
                ToolKind::Initialize.bit()
                    | ToolKind::BashCommand.bit()
                    | ToolKind::ReadFiles.bit()
                    | ToolKind::FileWriteOrEdit.bit()
                    | ToolKind::MultiFileEdit.bit()
                    | ToolKind::VerifyEdit.bit()
                    | ToolKind::UndoEdit.bit()
                    | ToolKind::ContextSave.bit()
                    | ToolKind::ReadImage.bit()
                    | ToolKind::CodeMap.bit()
                    | ToolKind::ApplyPatch.bit()
            }
            Self::Coding => {
                ToolKind::Initialize.bit()
                    | ToolKind::BashCommand.bit()
                    | ToolKind::ReadFiles.bit()
                    | ToolKind::FileWriteOrEdit.bit()
                    | ToolKind::MultiFileEdit.bit()
                    | ToolKind::VerifyEdit.bit()
                    | ToolKind::UndoEdit.bit()
                    | ToolKind::CodeMap.bit()
                    | ToolKind::ApplyPatch.bit()
            }
            Self::ReadOnly => {
                ToolKind::Initialize.bit()
                    | ToolKind::ReadFiles.bit()
                    | ToolKind::ReadImage.bit()
                    | ToolKind::CodeMap.bit()
            }
            Self::Terminal => ToolKind::Initialize.bit() | ToolKind::BashCommand.bit(),
        }
    }
}

/// Fast, immutable allowlist used both for catalog filtering and call routing.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct ToolPolicy {
    mask: u16,
}

impl Default for ToolPolicy {
    fn default() -> Self {
        Self::from_profile(ToolProfile::Full)
    }
}

impl ToolPolicy {
    /// Build a policy from one curated profile.
    pub const fn from_profile(profile: ToolProfile) -> Self {
        Self { mask: profile.mask() }
    }

    /// Build an exact allowlist. Names are case-sensitive and duplicates are harmless.
    pub fn from_allowed_tools<I, S>(names: I) -> Result<Self>
    where
        I: IntoIterator<Item = S>,
        S: AsRef<str>,
    {
        let mut mask = 0_u16;
        let mut saw_name = false;
        for name in names {
            saw_name = true;
            let name = name.as_ref();
            let kind = ToolKind::parse(name).ok_or_else(|| {
                WinxError::ConfigurationError(format!(
                    "unknown MCP tool {name:?}; expected one of {}",
                    ALL_TOOL_NAMES.join(", ")
                ))
            })?;
            mask |= kind.bit();
        }
        if !saw_name {
            return Err(WinxError::ConfigurationError(
                "MCP tool allowlist cannot be empty".to_string(),
            ));
        }
        if mask & ToolKind::VerifyEdit.bit() != 0 && mask & ToolKind::BashCommand.bit() == 0 {
            return Err(WinxError::ConfigurationError(
                "VerifyEdit requires BashCommand in the same MCP tool allowlist".to_string(),
            ));
        }
        Ok(Self { mask })
    }

    /// Use an explicit allowlist when present, otherwise use the selected profile.
    pub fn resolve(profile: ToolProfile, allowed_tools: Option<&[String]>) -> Result<Self> {
        allowed_tools.map_or_else(
            || Ok(Self::from_profile(profile)),
            |names| Self::from_allowed_tools(names.iter()),
        )
    }

    /// Whether a tool may be advertised and called.
    pub fn allows(self, name: &str) -> bool {
        ToolKind::parse(name).is_some_and(|kind| self.allows_kind(kind))
    }

    /// Typed policy check used after the wire name has been parsed once.
    pub const fn allows_kind(self, kind: ToolKind) -> bool {
        self.mask & kind.bit() != 0
    }

    /// Allowed names in stable catalog order.
    pub fn names(self) -> impl Iterator<Item = &'static str> {
        ToolKind::ALL.into_iter().filter(move |kind| self.allows_kind(*kind)).map(ToolKind::as_str)
    }

    /// Number of advertised tools.
    pub fn len(self) -> usize {
        self.mask.count_ones() as usize
    }

    /// Whether this policy advertises no tools.
    pub fn is_empty(self) -> bool {
        self.mask == 0
    }
}

#[cfg(test)]
mod tests {
    #![allow(clippy::expect_used)]

    use super::{ToolPolicy, ToolProfile, ALL_TOOL_NAMES};

    #[test]
    fn profiles_have_stable_expected_catalogs() {
        assert_eq!(ToolPolicy::default().names().collect::<Vec<_>>(), ALL_TOOL_NAMES);
        assert_eq!(
            ToolPolicy::from_profile(ToolProfile::Coding).names().collect::<Vec<_>>(),
            vec![
                "Initialize",
                "BashCommand",
                "ReadFiles",
                "FileWriteOrEdit",
                "MultiFileEdit",
                "VerifyEdit",
                "UndoEdit",
                "CodeMap",
                "ApplyPatch",
            ]
        );
        assert_eq!(
            ToolPolicy::from_profile(ToolProfile::ReadOnly).names().collect::<Vec<_>>(),
            vec!["Initialize", "ReadFiles", "ReadImage", "CodeMap"]
        );
        assert_eq!(
            ToolPolicy::from_profile(ToolProfile::Terminal).names().collect::<Vec<_>>(),
            vec!["Initialize", "BashCommand"]
        );
    }

    #[test]
    fn explicit_allowlist_is_exact_deduplicated_and_validated() {
        let policy = ToolPolicy::from_allowed_tools(["ReadFiles", "Initialize", "ReadFiles"])
            .expect("valid policy");
        assert_eq!(policy.names().collect::<Vec<_>>(), vec!["Initialize", "ReadFiles"]);
        assert!(ToolPolicy::from_allowed_tools(["readfiles"]).is_err());
        assert!(ToolPolicy::from_allowed_tools(Vec::<String>::new()).is_err());
        assert!(ToolPolicy::from_allowed_tools(["VerifyEdit"]).is_err());
        assert!(ToolPolicy::from_allowed_tools(["VerifyEdit", "BashCommand"]).is_ok());
    }
}