Skip to main content

winx_code_agent/
tool_policy.rs

1//! Per-principal MCP tool catalog policies.
2
3use serde::Deserialize;
4
5use crate::errors::{Result, WinxError};
6
7/// Canonical tool names in the order advertised by `tools/list`.
8pub const ALL_TOOL_NAMES: [&str; 9] = [
9    "Initialize",
10    "BashCommand",
11    "ReadFiles",
12    "FileWriteOrEdit",
13    "MultiFileEdit",
14    "UndoEdit",
15    "ContextSave",
16    "ReadImage",
17    "CodeMap",
18];
19
20const INITIALIZE: u16 = 1 << 0;
21const BASH_COMMAND: u16 = 1 << 1;
22const READ_FILES: u16 = 1 << 2;
23const FILE_WRITE_OR_EDIT: u16 = 1 << 3;
24const MULTI_FILE_EDIT: u16 = 1 << 4;
25const UNDO_EDIT: u16 = 1 << 5;
26const CONTEXT_SAVE: u16 = 1 << 6;
27const READ_IMAGE: u16 = 1 << 7;
28const CODE_MAP: u16 = 1 << 8;
29
30/// Curated catalog shapes for common MCP clients and workflows.
31#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, clap::ValueEnum)]
32#[serde(rename_all = "kebab-case")]
33pub enum ToolProfile {
34    /// Advertise and permit every Winx tool.
35    #[default]
36    Full,
37    /// Core repository exploration and editing without image or handoff helpers.
38    Coding,
39    /// Filesystem read/navigation tools only; shell and edit calls are rejected.
40    ReadOnly,
41    /// Durable terminal access with the smallest useful catalog.
42    Terminal,
43}
44
45impl ToolProfile {
46    const fn mask(self) -> u16 {
47        match self {
48            Self::Full => {
49                INITIALIZE
50                    | BASH_COMMAND
51                    | READ_FILES
52                    | FILE_WRITE_OR_EDIT
53                    | MULTI_FILE_EDIT
54                    | UNDO_EDIT
55                    | CONTEXT_SAVE
56                    | READ_IMAGE
57                    | CODE_MAP
58            }
59            Self::Coding => {
60                INITIALIZE
61                    | BASH_COMMAND
62                    | READ_FILES
63                    | FILE_WRITE_OR_EDIT
64                    | MULTI_FILE_EDIT
65                    | UNDO_EDIT
66                    | CODE_MAP
67            }
68            Self::ReadOnly => INITIALIZE | READ_FILES | READ_IMAGE | CODE_MAP,
69            Self::Terminal => INITIALIZE | BASH_COMMAND,
70        }
71    }
72}
73
74/// Fast, immutable allowlist used both for catalog filtering and call routing.
75#[derive(Clone, Copy, Debug, Eq, PartialEq)]
76pub struct ToolPolicy {
77    mask: u16,
78}
79
80impl Default for ToolPolicy {
81    fn default() -> Self {
82        Self::from_profile(ToolProfile::Full)
83    }
84}
85
86impl ToolPolicy {
87    /// Build a policy from one curated profile.
88    pub const fn from_profile(profile: ToolProfile) -> Self {
89        Self { mask: profile.mask() }
90    }
91
92    /// Build an exact allowlist. Names are case-sensitive and duplicates are harmless.
93    pub fn from_allowed_tools<I, S>(names: I) -> Result<Self>
94    where
95        I: IntoIterator<Item = S>,
96        S: AsRef<str>,
97    {
98        let mut mask = 0_u16;
99        let mut saw_name = false;
100        for name in names {
101            saw_name = true;
102            let name = name.as_ref();
103            let bit = tool_bit(name).ok_or_else(|| {
104                WinxError::ConfigurationError(format!(
105                    "unknown MCP tool {name:?}; expected one of {}",
106                    ALL_TOOL_NAMES.join(", ")
107                ))
108            })?;
109            mask |= bit;
110        }
111        if !saw_name {
112            return Err(WinxError::ConfigurationError(
113                "MCP tool allowlist cannot be empty".to_string(),
114            ));
115        }
116        Ok(Self { mask })
117    }
118
119    /// Use an explicit allowlist when present, otherwise use the selected profile.
120    pub fn resolve(profile: ToolProfile, allowed_tools: Option<&[String]>) -> Result<Self> {
121        allowed_tools.map_or_else(
122            || Ok(Self::from_profile(profile)),
123            |names| Self::from_allowed_tools(names.iter()),
124        )
125    }
126
127    /// Whether a tool may be advertised and called.
128    pub fn allows(self, name: &str) -> bool {
129        tool_bit(name).is_some_and(|bit| self.mask & bit != 0)
130    }
131
132    /// Allowed names in stable catalog order.
133    pub fn names(self) -> impl Iterator<Item = &'static str> {
134        ALL_TOOL_NAMES.into_iter().filter(move |name| self.allows(name))
135    }
136
137    /// Number of advertised tools.
138    pub fn len(self) -> usize {
139        self.mask.count_ones() as usize
140    }
141
142    /// Whether this policy advertises no tools.
143    pub fn is_empty(self) -> bool {
144        self.mask == 0
145    }
146}
147
148fn tool_bit(name: &str) -> Option<u16> {
149    ALL_TOOL_NAMES.iter().position(|candidate| *candidate == name).map(|index| 1 << index)
150}
151
152#[cfg(test)]
153mod tests {
154    #![allow(clippy::expect_used)]
155
156    use super::{ToolPolicy, ToolProfile, ALL_TOOL_NAMES};
157
158    #[test]
159    fn profiles_have_stable_expected_catalogs() {
160        assert_eq!(ToolPolicy::default().names().collect::<Vec<_>>(), ALL_TOOL_NAMES);
161        assert_eq!(
162            ToolPolicy::from_profile(ToolProfile::Coding).names().collect::<Vec<_>>(),
163            vec![
164                "Initialize",
165                "BashCommand",
166                "ReadFiles",
167                "FileWriteOrEdit",
168                "MultiFileEdit",
169                "UndoEdit",
170                "CodeMap",
171            ]
172        );
173        assert_eq!(
174            ToolPolicy::from_profile(ToolProfile::ReadOnly).names().collect::<Vec<_>>(),
175            vec!["Initialize", "ReadFiles", "ReadImage", "CodeMap"]
176        );
177        assert_eq!(
178            ToolPolicy::from_profile(ToolProfile::Terminal).names().collect::<Vec<_>>(),
179            vec!["Initialize", "BashCommand"]
180        );
181    }
182
183    #[test]
184    fn explicit_allowlist_is_exact_deduplicated_and_validated() {
185        let policy = ToolPolicy::from_allowed_tools(["ReadFiles", "Initialize", "ReadFiles"])
186            .expect("valid policy");
187        assert_eq!(policy.names().collect::<Vec<_>>(), vec!["Initialize", "ReadFiles"]);
188        assert!(ToolPolicy::from_allowed_tools(["readfiles"]).is_err());
189        assert!(ToolPolicy::from_allowed_tools(Vec::<String>::new()).is_err());
190    }
191}