winx_code_agent/
tool_policy.rs1use serde::Deserialize;
4
5use crate::errors::{Result, WinxError};
6
7pub 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#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, clap::ValueEnum)]
32#[serde(rename_all = "kebab-case")]
33pub enum ToolProfile {
34 #[default]
36 Full,
37 Coding,
39 ReadOnly,
41 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#[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 pub const fn from_profile(profile: ToolProfile) -> Self {
89 Self { mask: profile.mask() }
90 }
91
92 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 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 pub fn allows(self, name: &str) -> bool {
129 tool_bit(name).is_some_and(|bit| self.mask & bit != 0)
130 }
131
132 pub fn names(self) -> impl Iterator<Item = &'static str> {
134 ALL_TOOL_NAMES.into_iter().filter(move |name| self.allows(name))
135 }
136
137 pub fn len(self) -> usize {
139 self.mask.count_ones() as usize
140 }
141
142 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}