Skip to main content

pmcp_code_mode/
config.rs

1//! Code Mode configuration.
2
3use crate::types::RiskLevel;
4use serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6use std::collections::HashSet;
7
8/// Configuration for Code Mode.
9#[derive(Debug, Clone, Serialize, Deserialize)]
10pub struct CodeModeConfig {
11    /// Whether Code Mode is enabled for this server
12    #[serde(default)]
13    pub enabled: bool,
14
15    // ========================================================================
16    // GraphQL-specific settings
17    // ========================================================================
18    /// Whether to allow mutations (MVP: false)
19    #[serde(default)]
20    pub allow_mutations: bool,
21
22    /// Allowed mutation names (whitelist). If empty and allow_mutations=true, all are allowed.
23    #[serde(default)]
24    pub allowed_mutations: HashSet<String>,
25
26    /// Blocked mutation names (blacklist). Always blocked even if allow_mutations=true.
27    #[serde(default)]
28    pub blocked_mutations: HashSet<String>,
29
30    /// Whether to allow introspection queries
31    #[serde(default)]
32    pub allow_introspection: bool,
33
34    /// Fields that should never be returned (Type.field format) - GraphQL
35    #[serde(default)]
36    pub blocked_fields: HashSet<String>,
37
38    /// Allowed query names (whitelist). If empty and mode is allowlist, none are allowed.
39    #[serde(default)]
40    pub allowed_queries: HashSet<String>,
41
42    /// Blocked query names (blocklist). Always blocked even if reads enabled.
43    #[serde(default)]
44    pub blocked_queries: HashSet<String>,
45
46    // ========================================================================
47    // OpenAPI-specific settings
48    // ========================================================================
49    /// Whether read operations (GET) are enabled (default: true)
50    #[serde(default = "default_true")]
51    pub openapi_reads_enabled: bool,
52
53    /// Whether write operations (POST, PUT, PATCH) are allowed globally
54    #[serde(default)]
55    pub openapi_allow_writes: bool,
56
57    /// Allowed write operations (operationId or "METHOD /path")
58    #[serde(default)]
59    pub openapi_allowed_writes: HashSet<String>,
60
61    /// Blocked write operations
62    #[serde(default)]
63    pub openapi_blocked_writes: HashSet<String>,
64
65    /// Whether delete operations (DELETE) are allowed globally
66    #[serde(default)]
67    pub openapi_allow_deletes: bool,
68
69    /// Allowed delete operations (operationId or "METHOD /path")
70    #[serde(default)]
71    pub openapi_allowed_deletes: HashSet<String>,
72
73    /// Blocked paths (glob patterns like "/admin/*")
74    #[serde(default)]
75    pub openapi_blocked_paths: HashSet<String>,
76
77    /// Fields that are stripped from API responses entirely (no access)
78    #[serde(default)]
79    pub openapi_internal_blocked_fields: HashSet<String>,
80
81    /// Fields that can be used internally but not in script output
82    #[serde(default)]
83    pub openapi_output_blocked_fields: HashSet<String>,
84
85    /// Whether scripts must declare their return type with @returns
86    #[serde(default)]
87    pub openapi_require_output_declaration: bool,
88
89    // ========================================================================
90    // Common settings
91    // ========================================================================
92    /// Action tags to override inferred actions for specific operations.
93    #[serde(default)]
94    pub action_tags: HashMap<String, String>,
95
96    /// Maximum query depth
97    #[serde(default = "default_max_depth")]
98    pub max_depth: u32,
99
100    /// Maximum field count per query
101    #[serde(default = "default_max_field_count")]
102    pub max_field_count: u32,
103
104    /// Maximum estimated query cost
105    #[serde(default = "default_max_cost")]
106    pub max_cost: u32,
107
108    /// Allowed sensitive data categories
109    #[serde(default)]
110    pub allowed_sensitive_categories: HashSet<String>,
111
112    /// Token time-to-live in seconds
113    #[serde(default = "default_token_ttl")]
114    pub token_ttl_seconds: i64,
115
116    /// Risk levels that can be auto-approved without human confirmation
117    #[serde(default = "default_auto_approve_levels")]
118    pub auto_approve_levels: Vec<RiskLevel>,
119
120    /// Maximum query length in characters
121    #[serde(default = "default_max_query_length")]
122    pub max_query_length: usize,
123
124    /// Maximum result rows to return
125    #[serde(default = "default_max_result_rows")]
126    pub max_result_rows: usize,
127
128    /// Query execution timeout in seconds
129    #[serde(default = "default_query_timeout")]
130    pub query_timeout_seconds: u32,
131
132    /// Server ID for token generation
133    #[serde(default)]
134    pub server_id: Option<String>,
135}
136
137impl Default for CodeModeConfig {
138    fn default() -> Self {
139        Self {
140            enabled: false,
141            // GraphQL
142            allow_mutations: false,
143            allowed_mutations: HashSet::new(),
144            blocked_mutations: HashSet::new(),
145            allow_introspection: false,
146            blocked_fields: HashSet::new(),
147            allowed_queries: HashSet::new(),
148            blocked_queries: HashSet::new(),
149            // OpenAPI
150            openapi_reads_enabled: true,
151            openapi_allow_writes: false,
152            openapi_allowed_writes: HashSet::new(),
153            openapi_blocked_writes: HashSet::new(),
154            openapi_allow_deletes: false,
155            openapi_allowed_deletes: HashSet::new(),
156            openapi_blocked_paths: HashSet::new(),
157            openapi_internal_blocked_fields: HashSet::new(),
158            openapi_output_blocked_fields: HashSet::new(),
159            openapi_require_output_declaration: false,
160            // Common
161            action_tags: HashMap::new(),
162            max_depth: default_max_depth(),
163            max_field_count: default_max_field_count(),
164            max_cost: default_max_cost(),
165            allowed_sensitive_categories: HashSet::new(),
166            token_ttl_seconds: default_token_ttl(),
167            auto_approve_levels: default_auto_approve_levels(),
168            max_query_length: default_max_query_length(),
169            max_result_rows: default_max_result_rows(),
170            query_timeout_seconds: default_query_timeout(),
171            server_id: None,
172        }
173    }
174}
175
176impl CodeModeConfig {
177    /// Create a new config with Code Mode enabled.
178    pub fn enabled() -> Self {
179        Self {
180            enabled: true,
181            ..Default::default()
182        }
183    }
184
185    /// Check if a risk level should be auto-approved.
186    pub fn should_auto_approve(&self, risk_level: RiskLevel) -> bool {
187        self.auto_approve_levels.contains(&risk_level)
188    }
189
190    /// Get the server ID, falling back to a default.
191    pub fn server_id(&self) -> &str {
192        self.server_id.as_deref().unwrap_or("unknown")
193    }
194
195    /// Convert to ServerConfigEntity for policy evaluation.
196    pub fn to_server_config_entity(&self) -> crate::policy::ServerConfigEntity {
197        crate::policy::ServerConfigEntity {
198            server_id: self.server_id().to_string(),
199            server_type: "graphql".to_string(),
200            allow_write: self.allow_mutations,
201            allow_delete: self.allow_mutations,
202            allow_admin: self.allow_introspection,
203            allowed_operations: self.allowed_mutations.clone(),
204            blocked_operations: self.blocked_mutations.clone(),
205            max_depth: self.max_depth,
206            max_field_count: self.max_field_count,
207            max_cost: self.max_cost,
208            max_api_calls: 50,
209            blocked_fields: self.blocked_fields.clone(),
210            allowed_sensitive_categories: self.allowed_sensitive_categories.clone(),
211        }
212    }
213
214    /// Convert to OpenAPIServerEntity for policy evaluation (OpenAPI Code Mode).
215    #[cfg(feature = "openapi-code-mode")]
216    pub fn to_openapi_server_entity(&self) -> crate::policy::OpenAPIServerEntity {
217        let mut allowed_operations = self.openapi_allowed_writes.clone();
218        allowed_operations.extend(self.openapi_allowed_deletes.clone());
219
220        let write_mode = if !self.openapi_allow_writes {
221            "deny_all"
222        } else if !self.openapi_allowed_writes.is_empty() {
223            "allowlist"
224        } else if !self.openapi_blocked_writes.is_empty() {
225            "blocklist"
226        } else {
227            "allow_all"
228        };
229
230        crate::policy::OpenAPIServerEntity {
231            server_id: self.server_id().to_string(),
232            server_type: "openapi".to_string(),
233            allow_write: self.openapi_allow_writes,
234            allow_delete: self.openapi_allow_deletes,
235            allow_admin: false,
236            write_mode: write_mode.to_string(),
237            max_depth: self.max_depth,
238            max_cost: self.max_cost,
239            max_api_calls: 50,
240            max_loop_iterations: 100,
241            max_script_length: self.max_query_length as u32,
242            max_nesting_depth: self.max_depth,
243            execution_timeout_seconds: self.query_timeout_seconds,
244            allowed_operations,
245            blocked_operations: self.openapi_blocked_writes.clone(),
246            allowed_methods: HashSet::new(),
247            blocked_methods: HashSet::new(),
248            allowed_path_patterns: HashSet::new(),
249            blocked_path_patterns: self.openapi_blocked_paths.clone(),
250            sensitive_path_patterns: self.openapi_blocked_paths.clone(),
251            auto_approve_read_only: self.openapi_reads_enabled,
252            max_api_calls_for_auto_approve: 10,
253            internal_blocked_fields: self.openapi_internal_blocked_fields.clone(),
254            output_blocked_fields: self.openapi_output_blocked_fields.clone(),
255            require_output_declaration: self.openapi_require_output_declaration,
256        }
257    }
258}
259
260fn default_true() -> bool {
261    true
262}
263
264fn default_token_ttl() -> i64 {
265    300 // 5 minutes
266}
267
268fn default_auto_approve_levels() -> Vec<RiskLevel> {
269    vec![RiskLevel::Low]
270}
271
272fn default_max_query_length() -> usize {
273    10000
274}
275
276fn default_max_result_rows() -> usize {
277    10000
278}
279
280fn default_query_timeout() -> u32 {
281    30
282}
283
284fn default_max_depth() -> u32 {
285    10
286}
287
288fn default_max_field_count() -> u32 {
289    100
290}
291
292fn default_max_cost() -> u32 {
293    1000
294}
295
296#[cfg(test)]
297mod tests {
298    use super::*;
299
300    #[test]
301    fn test_default_config() {
302        let config = CodeModeConfig::default();
303        assert!(!config.enabled);
304        assert!(!config.allow_mutations);
305        assert_eq!(config.token_ttl_seconds, 300);
306        assert_eq!(config.auto_approve_levels, vec![RiskLevel::Low]);
307    }
308
309    #[test]
310    fn test_enabled_config() {
311        let config = CodeModeConfig::enabled();
312        assert!(config.enabled);
313    }
314
315    #[test]
316    fn test_auto_approve() {
317        let config = CodeModeConfig::default();
318        assert!(config.should_auto_approve(RiskLevel::Low));
319        assert!(!config.should_auto_approve(RiskLevel::Medium));
320        assert!(!config.should_auto_approve(RiskLevel::High));
321        assert!(!config.should_auto_approve(RiskLevel::Critical));
322    }
323}