vtcode_core/tools/
safety_gateway.rs

1//! Unified Safety Gateway
2//!
3//! Consolidates all safety checking mechanisms into a single gateway:
4//! - Rate limiting (from runloop's tool_call_safety)
5//! - Destructive tool detection
6//! - Command policy enforcement
7//! - Planning workflow restrictions
8//!
9//! This provides consistent safety decisions across all tool execution paths.
10
11use hashbrown::HashSet;
12use parking_lot::{Mutex, RwLock};
13use std::future::Future;
14use std::path::{Path, PathBuf};
15use std::sync::Arc;
16use std::time::{Duration, Instant};
17
18use serde::{Deserialize, Serialize};
19use serde_json::Value;
20use thiserror::Error;
21
22use crate::config::CommandsConfig;
23use crate::config::constants::tools;
24use crate::dotfile_protection::{AccessContext, AccessType, DotfileGuardian, ProtectionDecision, get_global_guardian};
25use crate::tools::apply_patch::{Patch, PatchOperation, decode_apply_patch_input};
26use crate::tools::command_policy::CommandPolicyEvaluator;
27use crate::tools::invocation::ToolInvocationId;
28use crate::tools::rate_limit_config::{tool_calls_per_minute_from_env, tool_calls_per_second_from_env};
29use crate::tools::registry::{RiskLevel, ToolRiskContext, ToolRiskScorer, ToolSource, WorkspaceTrust};
30use crate::tools::tool_intent::{
31    ToolIntent, action_qualified_policy_name, classify_tool_intent, command_session_action_in,
32    command_session_action_is, file_operation_action_is,
33};
34use vtcode_config::constants::tool_limits::{
35    DEFAULT_MAX_TOOL_CALLS_PER_TURN, DEFAULT_SAFETY_MAX_TOOL_CALLS_PER_SESSION, MAX_CONTROL_PLANE_TOOL_CALLS_PER_TURN,
36};
37use vtcode_config::core::DotfileProtectionConfig;
38
39/// Trust level used by the safety gateway for approval bypass decisions.
40#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
41pub enum SafetyTrustLevel {
42    Untrusted,
43    #[default]
44    Standard,
45    Elevated,
46    Full,
47}
48
49impl SafetyTrustLevel {
50    #[inline]
51    pub const fn can_bypass_approval(self) -> bool {
52        matches!(self, Self::Elevated | Self::Full)
53    }
54}
55
56/// Minimal execution context required for safety decisions.
57#[derive(Debug, Clone)]
58pub struct SafetyContext {
59    pub session_id: String,
60    pub trust_level: SafetyTrustLevel,
61}
62
63impl SafetyContext {
64    #[must_use]
65    pub fn new(session_id: impl Into<String>) -> Self {
66        Self {
67            session_id: session_id.into(),
68            trust_level: SafetyTrustLevel::default(),
69        }
70    }
71}
72
73/// Safety decision for a tool invocation
74#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
75pub enum SafetyDecision {
76    /// Tool execution is allowed without user approval
77    Allow,
78    /// Tool execution is denied with a reason
79    Deny(String),
80    /// Tool execution requires user approval with justification
81    NeedsApproval(String),
82}
83
84impl SafetyDecision {
85    /// Whether execution can proceed (Allow only)
86    #[inline]
87    pub fn is_allowed(&self) -> bool {
88        matches!(self, SafetyDecision::Allow)
89    }
90
91    /// Whether execution is blocked (Deny)
92    #[inline]
93    pub fn is_denied(&self) -> bool {
94        matches!(self, SafetyDecision::Deny(_))
95    }
96
97    /// Whether user approval is needed
98    #[inline]
99    pub fn needs_approval(&self) -> bool {
100        matches!(self, SafetyDecision::NeedsApproval(_))
101    }
102
103    /// Get the reason/justification if present
104    pub fn reason(&self) -> Option<&str> {
105        match self {
106            SafetyDecision::Allow => None,
107            SafetyDecision::Deny(reason) | SafetyDecision::NeedsApproval(reason) => Some(reason),
108        }
109    }
110}
111
112/// Errors from safety checks
113#[derive(Debug, Error, Clone)]
114pub enum SafetyError {
115    #[error("Rate limit exceeded: {current} calls in {window} (max: {max})")]
116    RateLimitExceeded {
117        current: usize,
118        max: usize,
119        window: &'static str,
120    },
121    #[error("Per-turn tool limit reached (max: {max})")]
122    TurnLimitReached { max: usize },
123    #[error("Session tool limit reached (max: {max})")]
124    SessionLimitReached { max: usize },
125    #[error("Planning workflow violation: {0}")]
126    PlanningPolicyViolation(String),
127    #[error("Command policy denied: {0}")]
128    CommandPolicyDenied(String),
129    #[error("Dotfile protection violation: {0}")]
130    DotfileProtectionViolation(String),
131}
132
133/// Result of a safety check with optional retry hint metadata.
134#[derive(Debug, Clone)]
135pub struct SafetyCheckResult {
136    /// Final decision for this invocation.
137    pub decision: SafetyDecision,
138    /// Suggested delay before retrying if the decision is a rate-limit denial.
139    pub retry_after: Option<Duration>,
140    /// Structured error when denial is produced by a safety limit.
141    pub violation: Option<SafetyError>,
142}
143
144/// Configuration for the safety gateway
145#[derive(Debug, Clone)]
146pub struct SafetyGatewayConfig {
147    /// Maximum tool calls per turn
148    pub max_per_turn: usize,
149    /// Maximum tool calls per session
150    pub max_per_session: usize,
151    /// Rate limit: calls per second
152    pub rate_limit_per_second: usize,
153    /// Rate limit: calls per minute (optional burst protection)
154    pub rate_limit_per_minute: Option<usize>,
155    /// Whether planning workflow is active (read-only)
156    pub planning_active: bool,
157    /// Workspace trust level
158    pub workspace_trust: WorkspaceTrust,
159    /// Risk threshold for requiring approval
160    pub approval_risk_threshold: RiskLevel,
161    /// Enforce short-window rate limiting (per-second/per-minute).
162    /// Turn/session limits are always enforced.
163    pub enforce_rate_limits: bool,
164}
165
166impl Default for SafetyGatewayConfig {
167    fn default() -> Self {
168        let rate_limit_per_second = tool_calls_per_second_from_env().unwrap_or(5);
169
170        Self {
171            max_per_turn: DEFAULT_MAX_TOOL_CALLS_PER_TURN,
172            max_per_session: DEFAULT_SAFETY_MAX_TOOL_CALLS_PER_SESSION,
173            rate_limit_per_second,
174            rate_limit_per_minute: tool_calls_per_minute_from_env(),
175            planning_active: false,
176            workspace_trust: WorkspaceTrust::Trusted,
177            approval_risk_threshold: RiskLevel::Medium,
178            enforce_rate_limits: true,
179        }
180    }
181}
182
183/// Rate limiter state (shared across async contexts)
184#[derive(Debug, Default)]
185struct RateLimiterState {
186    calls_per_second: std::collections::VecDeque<Instant>,
187    calls_per_minute: std::collections::VecDeque<Instant>,
188    current_turn_count: usize,
189    session_count: usize,
190    /// Count of control-plane calls (`wait`/`inspect`) recorded this turn.
191    /// These are exempt from `current_turn_count` so a long build does not
192    /// consume the budget its follow-up work needs, but they remain bounded
193    /// by their own per-turn cap and the session fuse.
194    control_plane_turn_count: usize,
195}
196
197/// Unified Safety Gateway
198///
199/// Consolidates rate limiting, destructive tool detection, command policy
200/// enforcement, planning workflow restrictions, and dotfile protection into a single
201/// safety decision point.
202pub struct SafetyGateway {
203    /// Configuration
204    config: RwLock<SafetyGatewayConfig>,
205    /// Command policy evaluator (optional, for shell commands)
206    command_policy: Option<CommandPolicyEvaluator>,
207    /// Rate limiter state
208    rate_state: Mutex<RateLimiterState>,
209    /// Preapproved tools for this session
210    preapproved: Mutex<HashSet<String>>,
211    /// Dotfile guardian for protected file access
212    dotfile_guardian: Option<Arc<DotfileGuardian>>,
213}
214
215#[derive(Debug, Clone, PartialEq, Eq, Hash)]
216struct FileAccessTarget {
217    path: PathBuf,
218    access_type: AccessType,
219}
220
221fn primary_path_arg(args: &Value) -> Option<&str> {
222    args.get("path")
223        .and_then(|value| value.as_str())
224        .or_else(|| args.get("file_path").and_then(|value| value.as_str()))
225        .or_else(|| args.get("filepath").and_then(|value| value.as_str()))
226        .or_else(|| args.get("target_path").and_then(|value| value.as_str()))
227}
228
229fn destination_path_arg(args: &Value) -> Option<&str> {
230    args.get("destination").and_then(|value| value.as_str())
231}
232
233fn push_file_access_target(targets: &mut Vec<FileAccessTarget>, path: &str, access_type: AccessType) {
234    let path_str = path.trim();
235    if path_str.is_empty() {
236        return;
237    }
238
239    let path = PathBuf::from(path_str);
240    // For small number of targets, linear search is faster than HashSet.
241    // In large patches, we'll use a local HashSet in patch_file_access_targets.
242    if targets
243        .iter()
244        .any(|existing| existing.path == path && existing.access_type == access_type)
245    {
246        return;
247    }
248
249    targets.push(FileAccessTarget { path, access_type });
250}
251
252fn command_text_for_tool(tool_name: &str, args: &Value) -> Option<String> {
253    match tool_name {
254        tools::EXEC_COMMAND => crate::tools::command_args::command_text(args).ok().flatten(),
255        tools::WRITE_STDIN => crate::tools::command_args::interactive_input_text(args).map(str::to_owned),
256        tools::SHELL | tools::RUN_PTY_CMD => crate::tools::command_args::command_text(args).ok().flatten(),
257        tools::SEND_PTY_INPUT => crate::tools::command_args::interactive_input_text(args).map(str::to_owned),
258        tools::UNIFIED_EXEC if command_session_action_is(args, "run") => {
259            crate::tools::command_args::command_text(args).ok().flatten()
260        }
261        tools::UNIFIED_EXEC if command_session_action_in(args, &["write", "continue"]) => {
262            crate::tools::command_args::interactive_input_text(args).map(str::to_owned)
263        }
264        _ => None,
265    }
266}
267
268fn patch_file_access_targets(args: &Value) -> Vec<FileAccessTarget> {
269    let Ok(Some(patch_input)) = decode_apply_patch_input(args) else {
270        return Vec::new();
271    };
272    let Ok(patch) = Patch::parse(&patch_input.text) else {
273        return Vec::new();
274    };
275
276    let mut targets = Vec::new();
277    for operation in patch.operations() {
278        match operation {
279            PatchOperation::AddFile { path, .. } => {
280                push_file_access_target(&mut targets, path, AccessType::Write);
281            }
282            PatchOperation::DeleteFile { path } => {
283                push_file_access_target(&mut targets, path, AccessType::Delete);
284            }
285            PatchOperation::UpdateFile { path, new_path, .. } => {
286                push_file_access_target(&mut targets, path, AccessType::Modify);
287                if let Some(destination) = new_path.as_deref().filter(|candidate| *candidate != path) {
288                    push_file_access_target(&mut targets, destination, AccessType::Write);
289                }
290            }
291        }
292    }
293
294    targets
295}
296
297fn file_access_targets(tool_name: &str, args: &Value) -> Vec<FileAccessTarget> {
298    let mut targets = Vec::new();
299
300    match tool_name {
301        tools::WRITE_FILE | tools::CREATE_FILE => {
302            if let Some(path) = primary_path_arg(args) {
303                push_file_access_target(&mut targets, path, AccessType::Write);
304            }
305        }
306        tools::EDIT_FILE | "search_replace" => {
307            if let Some(path) = primary_path_arg(args) {
308                push_file_access_target(&mut targets, path, AccessType::Modify);
309            }
310        }
311        tools::DELETE_FILE => {
312            if let Some(path) = primary_path_arg(args) {
313                push_file_access_target(&mut targets, path, AccessType::Delete);
314            }
315        }
316        tools::MOVE_FILE => {
317            if let Some(path) = primary_path_arg(args) {
318                push_file_access_target(&mut targets, path, AccessType::Modify);
319            }
320            if let Some(path) = destination_path_arg(args) {
321                push_file_access_target(&mut targets, path, AccessType::Write);
322            }
323        }
324        tools::COPY_FILE => {
325            if let Some(path) = destination_path_arg(args) {
326                push_file_access_target(&mut targets, path, AccessType::Write);
327            }
328        }
329        tools::APPLY_PATCH => {
330            targets.extend(patch_file_access_targets(args));
331        }
332        tools::UNIFIED_FILE if file_operation_action_is(args, "write") => {
333            if let Some(path) = primary_path_arg(args) {
334                push_file_access_target(&mut targets, path, AccessType::Write);
335            }
336        }
337        tools::UNIFIED_FILE if file_operation_action_is(args, "edit") => {
338            if let Some(path) = primary_path_arg(args) {
339                push_file_access_target(&mut targets, path, AccessType::Modify);
340            }
341        }
342        tools::UNIFIED_FILE if file_operation_action_is(args, "delete") => {
343            if let Some(path) = primary_path_arg(args) {
344                push_file_access_target(&mut targets, path, AccessType::Delete);
345            }
346        }
347        tools::UNIFIED_FILE if file_operation_action_is(args, "move") => {
348            if let Some(path) = primary_path_arg(args) {
349                push_file_access_target(&mut targets, path, AccessType::Modify);
350            }
351            if let Some(path) = destination_path_arg(args) {
352                push_file_access_target(&mut targets, path, AccessType::Write);
353            }
354        }
355        tools::UNIFIED_FILE if file_operation_action_is(args, "copy") => {
356            if let Some(path) = destination_path_arg(args) {
357                push_file_access_target(&mut targets, path, AccessType::Write);
358            }
359        }
360        _ => {}
361    }
362
363    targets
364}
365
366fn proposed_changes_preview(args: &Value) -> String {
367    const PREVIEW_LIMIT: usize = 500;
368
369    let preview_text = |label: &str, text: &str| {
370        let preview_len = text.len().min(PREVIEW_LIMIT);
371        format!(
372            "{label} ({} bytes):\n{}{}",
373            text.len(),
374            &text[..preview_len],
375            if text.len() > preview_len { "..." } else { "" }
376        )
377    };
378
379    if let Some(content) = args.get("content").and_then(|value| value.as_str()) {
380        return preview_text("Content", content);
381    }
382
383    if let Some(old_str) = args.get("old_str").and_then(|value| value.as_str()) {
384        let new_str = args.get("new_str").and_then(|value| value.as_str()).unwrap_or("");
385        return format!("Replace:\n  '{old_str}'\nWith:\n  '{new_str}'");
386    }
387
388    if let Ok(Some(patch_input)) = decode_apply_patch_input(args) {
389        return preview_text("Patch", &patch_input.text);
390    }
391
392    "No details provided".to_string()
393}
394
395impl SafetyGateway {
396    /// Create a new safety gateway with default configuration
397    pub fn new() -> Self {
398        Self::with_config(SafetyGatewayConfig::default())
399    }
400
401    /// Create a new safety gateway with custom configuration
402    pub fn with_config(config: SafetyGatewayConfig) -> Self {
403        Self {
404            config: RwLock::new(config),
405            command_policy: None,
406            rate_state: Mutex::new(RateLimiterState::default()),
407            preapproved: Mutex::new(HashSet::new()),
408            dotfile_guardian: None,
409        }
410    }
411
412    /// Set the dotfile guardian for protected file access
413    pub fn with_dotfile_guardian(mut self, guardian: Arc<DotfileGuardian>) -> Self {
414        self.dotfile_guardian = Some(guardian);
415        self
416    }
417
418    /// Create and set a dotfile guardian from configuration
419    pub async fn with_dotfile_protection(mut self, config: DotfileProtectionConfig) -> anyhow::Result<Self> {
420        let guardian = DotfileGuardian::new(config).await?;
421        self.dotfile_guardian = Some(Arc::new(guardian));
422        Ok(self)
423    }
424
425    /// Set the command policy evaluator for shell command checks
426    pub fn with_command_policy(mut self, policy: CommandPolicyEvaluator) -> Self {
427        self.command_policy = Some(policy);
428        self
429    }
430
431    /// Create from commands config
432    pub fn with_commands_config(mut self, config: &CommandsConfig) -> Self {
433        self.command_policy = Some(CommandPolicyEvaluator::from_config(config));
434        self
435    }
436
437    /// Enable or disable planning workflow
438    pub fn set_planning_workflow(&self, active: bool) {
439        self.config.write().planning_active = active;
440    }
441
442    /// Set workspace trust level
443    pub fn set_workspace_trust(&self, trust: WorkspaceTrust) {
444        self.config.write().workspace_trust = trust;
445    }
446
447    /// Update rate limits
448    pub fn set_limits(&self, max_per_turn: usize, max_per_session: usize) {
449        let mut config = self.config.write();
450        config.max_per_turn = max_per_turn;
451        config.max_per_session = max_per_session;
452    }
453
454    /// Update rate-limiter thresholds.
455    pub fn set_rate_limits(&self, rate_limit_per_second: usize, rate_limit_per_minute: Option<usize>) {
456        let mut config = self.config.write();
457        if rate_limit_per_second > 0 {
458            config.rate_limit_per_second = rate_limit_per_second;
459        }
460        config.rate_limit_per_minute = rate_limit_per_minute.filter(|v| *v > 0);
461    }
462
463    /// Enable or disable rate-limit enforcement while preserving counters.
464    pub fn set_rate_limit_enforcement(&self, enabled: bool) {
465        self.config.write().enforce_rate_limits = enabled;
466    }
467
468    /// Increase session limit dynamically
469    pub fn increase_session_limit(&self, increment: usize) {
470        let mut config = self.config.write();
471        let new_max = config.max_per_session.saturating_add(increment);
472        config.max_per_session = new_max;
473        tracing::info!("Session tool limit increased to {}", new_max);
474    }
475
476    pub fn max_per_session(&self) -> usize {
477        self.config.read().max_per_session
478    }
479
480    /// Reset turn counters (call at start of new turn)
481    pub fn start_turn(&self) {
482        let mut state = self.rate_state.lock();
483        state.current_turn_count = 0;
484        state.control_plane_turn_count = 0;
485        state.calls_per_second.clear();
486        state.calls_per_minute.clear();
487    }
488
489    /// Reset all transient tool-budget state for a newly reconstructed
490    /// execution context and restore the limits supplied by the build mode.
491    pub fn reset_for_fresh_execution(&self, max_per_turn: usize, max_per_session: usize) {
492        self.set_limits(max_per_turn, max_per_session);
493        let mut state = self.rate_state.lock();
494        state.current_turn_count = 0;
495        state.control_plane_turn_count = 0;
496        state.session_count = 0;
497        state.calls_per_second.clear();
498        state.calls_per_minute.clear();
499    }
500
501    /// Preapprove a tool for this session
502    pub fn preapprove(&self, tool_name: &str) {
503        let mut preapproved = self.preapproved.lock();
504        preapproved.insert(tool_name.to_string());
505    }
506
507    /// Check if a tool is preapproved
508    pub fn is_preapproved(&self, tool_name: &str) -> bool {
509        let preapproved = self.preapproved.lock();
510        preapproved.contains(tool_name)
511    }
512
513    /// Main entry point: check safety for a tool invocation.
514    ///
515    /// Returns a [`SafetyDecision`] indicating whether execution can proceed.
516    /// Inline-delegating wrapper that returns the inner future directly to
517    /// avoid an extra coroutine state machine (audit section 16).
518    pub fn check_safety<'a>(
519        &'a self,
520        ctx: &'a SafetyContext,
521        tool_name: &'a str,
522        args: &'a Value,
523    ) -> impl Future<Output = SafetyDecision> + 'a {
524        self.check_safety_with_id(ctx, tool_name, args, None)
525    }
526
527    /// Check safety with explicit invocation ID for correlation
528    pub async fn check_safety_with_id(
529        &self,
530        ctx: &SafetyContext,
531        tool_name: &str,
532        args: &Value,
533        invocation_id: Option<ToolInvocationId>,
534    ) -> SafetyDecision {
535        let inv_id = invocation_id.map(|id| id.short()).unwrap_or_else(|| "unknown".to_string());
536
537        tracing::trace!(
538            invocation_id = %inv_id,
539            tool = %tool_name,
540            "SafetyGateway: checking safety"
541        );
542
543        if let Err(err) = self.check_rate_limits(crate::tools::tool_intent::is_turn_budget_exempt_call(tool_name, args))
544        {
545            tracing::warn!(
546                invocation_id = %inv_id,
547                error = %err,
548                "SafetyGateway: rate limit exceeded"
549            );
550            return SafetyDecision::Deny(err.to_string());
551        }
552
553        self.evaluate_non_rate_decision(ctx, tool_name, args, &inv_id).await
554    }
555
556    /// Check safety and atomically reserve a rate-limit slot on success.
557    ///
558    /// This avoids split check/record races by validating rate limits and recording
559    /// execution under a single lock acquisition.
560    pub async fn check_and_record(&self, ctx: &SafetyContext, tool_name: &str, args: &Value) -> SafetyCheckResult {
561        self.check_and_record_with_id(ctx, tool_name, args, None).await
562    }
563
564    /// Check safety with correlation ID and atomically reserve a rate-limit slot.
565    pub async fn check_and_record_with_id(
566        &self,
567        ctx: &SafetyContext,
568        tool_name: &str,
569        args: &Value,
570        invocation_id: Option<ToolInvocationId>,
571    ) -> SafetyCheckResult {
572        let inv_id = invocation_id.map(|id| id.short()).unwrap_or_else(|| "unknown".to_string());
573        tracing::trace!(
574            invocation_id = %inv_id,
575            tool = %tool_name,
576            "SafetyGateway: checking and recording safety"
577        );
578
579        let decision = self.evaluate_non_rate_decision(ctx, tool_name, args, &inv_id).await;
580
581        if decision.is_denied() {
582            return SafetyCheckResult { decision, retry_after: None, violation: None };
583        }
584
585        let control_plane = crate::tools::tool_intent::is_turn_budget_exempt_call(tool_name, args);
586        let now = Instant::now();
587        let mut state = self.rate_state.lock();
588        match self.check_rate_limits_locked(&mut state, now, control_plane) {
589            Ok(()) => {
590                self.record_execution_locked(&mut state, now, control_plane);
591                SafetyCheckResult { decision, retry_after: None, violation: None }
592            }
593            Err(err) => {
594                tracing::warn!(
595                    invocation_id = %inv_id,
596                    error = %err,
597                    "SafetyGateway: rate limit exceeded during atomic reservation"
598                );
599                SafetyCheckResult {
600                    decision: SafetyDecision::Deny(err.to_string()),
601                    retry_after: self.retry_after_for_violation(&err, &state, now),
602                    violation: Some(err),
603                }
604            }
605        }
606    }
607
608    async fn evaluate_non_rate_decision(
609        &self,
610        ctx: &SafetyContext,
611        tool_name: &str,
612        args: &Value,
613        inv_id: &str,
614    ) -> SafetyDecision {
615        if let Some(decision) = self.check_dotfile_protection(tool_name, args, &ctx.session_id).await {
616            tracing::info!(
617                invocation_id = %inv_id,
618                tool = %tool_name,
619                "SafetyGateway: dotfile protection triggered"
620            );
621            return decision;
622        }
623
624        if let Some(ref policy) = self.command_policy
625            && let Some(command) = command_text_for_tool(tool_name, args)
626            && !policy.allows_text(&command)
627        {
628            let reason = format!("Command '{command}' blocked by policy");
629            tracing::info!(
630                invocation_id = %inv_id,
631                command = %command,
632                "SafetyGateway: command policy denied"
633            );
634            return SafetyDecision::Deny(reason);
635        }
636
637        if self.is_preapproved(tool_name) {
638            tracing::trace!(
639                invocation_id = %inv_id,
640                tool = %tool_name,
641                "SafetyGateway: tool preapproved"
642            );
643            return SafetyDecision::Allow;
644        }
645
646        if ctx.trust_level.can_bypass_approval() {
647            tracing::trace!(
648                invocation_id = %inv_id,
649                tool = %tool_name,
650                trust_level = ?ctx.trust_level,
651                "SafetyGateway: trust level allows bypass"
652            );
653            return SafetyDecision::Allow;
654        }
655
656        // Classify once for the risk context, the justification, and the
657        // destructive gate below: all three read the same tool name and args,
658        // and each re-classification re-parses command payloads.
659        let intent = classify_tool_intent(tool_name, args);
660        let risk_ctx = self.build_risk_context(tool_name, args, &intent);
661        let risk_level = ToolRiskScorer::calculate_risk(&risk_ctx);
662
663        if ToolRiskScorer::requires_justification(risk_level, self.config.read().approval_risk_threshold) {
664            let justification = self.build_approval_justification(tool_name, &risk_level, args, &intent);
665            tracing::info!(
666                invocation_id = %inv_id,
667                tool = %tool_name,
668                risk = %risk_level,
669                "SafetyGateway: requires approval"
670            );
671            return SafetyDecision::NeedsApproval(justification);
672        }
673
674        if intent.destructive {
675            let justification = format!("Tool '{tool_name}' is destructive and may modify files or execute commands.");
676            tracing::info!(
677                invocation_id = %inv_id,
678                tool = %tool_name,
679                "SafetyGateway: destructive tool requires approval"
680            );
681            return SafetyDecision::NeedsApproval(justification);
682        }
683
684        SafetyDecision::Allow
685    }
686
687    /// Record that a tool call was executed (for rate limiting)
688    pub fn record_execution(&self) {
689        let mut state = self.rate_state.lock();
690        self.record_execution_locked(&mut state, Instant::now(), false);
691    }
692
693    /// Check rate limits without recording
694    fn check_rate_limits(&self, control_plane: bool) -> Result<(), SafetyError> {
695        let mut state = self.rate_state.lock();
696        self.check_rate_limits_locked(&mut state, Instant::now(), control_plane)
697    }
698
699    fn check_rate_limits_locked(
700        &self,
701        state: &mut RateLimiterState,
702        now: Instant,
703        control_plane: bool,
704    ) -> Result<(), SafetyError> {
705        let config = self.config.read();
706        self.prune_rate_windows(state, now);
707
708        if config.enforce_rate_limits {
709            if state.calls_per_second.len() >= config.rate_limit_per_second {
710                return Err(SafetyError::RateLimitExceeded {
711                    current: state.calls_per_second.len(),
712                    max: config.rate_limit_per_second,
713                    window: "1s",
714                });
715            }
716
717            if let Some(limit) = config.rate_limit_per_minute
718                && state.calls_per_minute.len() >= limit
719            {
720                return Err(SafetyError::RateLimitExceeded {
721                    current: state.calls_per_minute.len(),
722                    max: limit,
723                    window: "60s",
724                });
725            }
726        }
727
728        if control_plane {
729            // Control-plane calls bypass the per-turn work budget but keep
730            // their own bounded cap so the exemption cannot become the
731            // dominant spend path.
732            if state.control_plane_turn_count >= MAX_CONTROL_PLANE_TOOL_CALLS_PER_TURN {
733                return Err(SafetyError::TurnLimitReached { max: MAX_CONTROL_PLANE_TOOL_CALLS_PER_TURN });
734            }
735        } else if state.current_turn_count >= config.max_per_turn {
736            return Err(SafetyError::TurnLimitReached { max: config.max_per_turn });
737        }
738
739        if state.session_count >= config.max_per_session {
740            return Err(SafetyError::SessionLimitReached { max: config.max_per_session });
741        }
742
743        Ok(())
744    }
745
746    fn record_execution_locked(&self, state: &mut RateLimiterState, now: Instant, control_plane: bool) {
747        state.calls_per_second.push_back(now);
748        state.calls_per_minute.push_back(now);
749        state.session_count = state.session_count.saturating_add(1);
750        if control_plane {
751            state.control_plane_turn_count = state.control_plane_turn_count.saturating_add(1);
752        } else {
753            state.current_turn_count = state.current_turn_count.saturating_add(1);
754        }
755    }
756
757    fn prune_rate_windows(&self, state: &mut RateLimiterState, now: Instant) {
758        while let Some(front) = state.calls_per_second.front() {
759            if now.duration_since(*front) > Duration::from_secs(1) {
760                state.calls_per_second.pop_front();
761            } else {
762                break;
763            }
764        }
765        while let Some(front) = state.calls_per_minute.front() {
766            if now.duration_since(*front) > Duration::from_secs(60) {
767                state.calls_per_minute.pop_front();
768            } else {
769                break;
770            }
771        }
772    }
773
774    fn retry_after_for_violation(
775        &self,
776        violation: &SafetyError,
777        state: &RateLimiterState,
778        now: Instant,
779    ) -> Option<Duration> {
780        match violation {
781            SafetyError::RateLimitExceeded { window: "1s", .. } => state
782                .calls_per_second
783                .front()
784                .map(|first| Duration::from_secs(1).saturating_sub(now.duration_since(*first))),
785            SafetyError::RateLimitExceeded { window: "60s", .. } => state
786                .calls_per_minute
787                .front()
788                .map(|first| Duration::from_secs(60).saturating_sub(now.duration_since(*first))),
789            _ => None,
790        }
791    }
792
793    /// Check dotfile protection for file operations.
794    /// Returns Some(SafetyDecision) if dotfile protection applies, None otherwise.
795    async fn check_dotfile_protection(
796        &self,
797        tool_name: &str,
798        args: &Value,
799        session_id: &str,
800    ) -> Option<SafetyDecision> {
801        // Use local guardian if set, otherwise try global guardian
802        let guardian = match self.dotfile_guardian.as_ref() {
803            Some(g) => g.clone(),
804            None => get_global_guardian()?,
805        };
806
807        let file_targets = file_access_targets(tool_name, args);
808        if file_targets.is_empty() {
809            return None;
810        }
811
812        let proposed_changes = proposed_changes_preview(args);
813
814        for target in file_targets {
815            if !guardian.is_protected(Path::new(&target.path)) {
816                continue;
817            }
818
819            let context = AccessContext::new(&target.path, target.access_type, tool_name, session_id)
820                .with_proposed_changes(&proposed_changes);
821
822            match guardian.request_access(&context).await {
823                Ok(ProtectionDecision::Allowed) => continue,
824                Ok(ProtectionDecision::RequiresConfirmation(req)) => {
825                    return Some(SafetyDecision::NeedsApproval(format!(
826                        "DOTFILE PROTECTION\n\n\
827                        File: {}\n\
828                        Operation: {}\n\
829                        Reason: {}\n\n\
830                        Proposed changes:\n{}\n\n\
831                        {}",
832                        req.file_path, req.access_type, req.protection_reason, req.proposed_changes, req.warning
833                    )));
834                }
835                Ok(ProtectionDecision::RequiresSecondaryAuth(req)) => {
836                    return Some(SafetyDecision::NeedsApproval(format!(
837                        "DOTFILE SECONDARY AUTHENTICATION REQUIRED\n\n\
838                        File: {} (whitelisted)\n\
839                        Operation: {}\n\
840                        Reason: {}\n\n\
841                        This file is on the whitelist but requires secondary authentication.\n\n\
842                        Proposed changes:\n{}\n\n\
843                        {}",
844                        req.file_path, req.access_type, req.protection_reason, req.proposed_changes, req.warning
845                    )));
846                }
847                Ok(ProtectionDecision::Blocked(violation)) => {
848                    return Some(SafetyDecision::Deny(format!(
849                        "DOTFILE MODIFICATION BLOCKED\n\n\
850                            File: {}\n\
851                            Reason: {}\n\n\
852                            Suggestion: {}",
853                        violation.file_path, violation.reason, violation.suggestion
854                    )));
855                }
856                Ok(ProtectionDecision::Denied(violation)) => {
857                    return Some(SafetyDecision::Deny(format!(
858                        "DOTFILE ACCESS DENIED\n\n\
859                            File: {}\n\
860                            Reason: {}\n\n\
861                            Suggestion: {}",
862                        violation.file_path, violation.reason, violation.suggestion
863                    )));
864                }
865                Err(e) => {
866                    tracing::error!("Dotfile protection check failed: {}", e);
867                    return Some(SafetyDecision::Deny(format!("Dotfile protection check failed: {e}")));
868                }
869            }
870        }
871
872        None
873    }
874
875    /// Get the dotfile guardian (if configured)
876    pub fn dotfile_guardian(&self) -> Option<&Arc<DotfileGuardian>> {
877        self.dotfile_guardian.as_ref()
878    }
879
880    /// Build risk context from tool name and arguments
881    fn build_risk_context(&self, tool_name: &str, args: &Value, intent: &ToolIntent) -> ToolRiskContext {
882        let source = if tool_name.starts_with("mcp_") {
883            ToolSource::Mcp
884        } else if tool_name.starts_with("acp_") {
885            ToolSource::Acp
886        } else {
887            ToolSource::Internal
888        };
889
890        let risk_tool_name = action_qualified_policy_name(tool_name, Some(args)).unwrap_or(tool_name);
891
892        let mut ctx = ToolRiskContext::new(risk_tool_name.to_string(), source, self.config.read().workspace_trust);
893
894        // Set flags based on tool type
895        if intent.mutating {
896            ctx = ctx.as_write();
897        }
898        if intent.destructive {
899            ctx = ctx.as_destructive();
900        }
901
902        if ToolRiskScorer::is_network_tool(risk_tool_name) {
903            ctx = ctx.accesses_network();
904        }
905
906        // Extract command args for shell tools
907        if let Some(command) = command_text_for_tool(tool_name, args) {
908            ctx = ctx.with_args(command.split_whitespace().map(String::from).collect());
909        }
910
911        ctx
912    }
913
914    /// Build justification message for approval prompt
915    fn build_approval_justification(
916        &self,
917        tool_name: &str,
918        risk_level: &RiskLevel,
919        args: &Value,
920        intent: &ToolIntent,
921    ) -> String {
922        let mut parts = Vec::new();
923
924        parts.push(format!("Tool: {tool_name}"));
925        parts.push(format!("Risk level: {risk_level}"));
926
927        if intent.destructive {
928            parts.push("This tool may modify or delete files.".to_string());
929        }
930
931        if let Some(command) = command_text_for_tool(tool_name, args) {
932            parts.push(format!("Command: {command}"));
933        }
934
935        let file_targets = file_access_targets(tool_name, args);
936        if let Some(target) = file_targets.first() {
937            parts.push(format!("Path: {}", target.path.display()));
938            if file_targets.len() > 1 {
939                parts.push(format!("Additional targets: {}", file_targets.len() - 1));
940            }
941        }
942
943        parts.join("\n")
944    }
945
946    /// Get current session statistics
947    pub fn get_stats(&self) -> SafetyStats {
948        let state = self.rate_state.lock();
949        let preapproved = self.preapproved.lock();
950        let config = self.config.read();
951
952        SafetyStats {
953            turn_count: state.current_turn_count,
954            session_count: state.session_count,
955            max_per_turn: config.max_per_turn,
956            max_per_session: config.max_per_session,
957            planning_active: config.planning_active,
958            preapproved_count: preapproved.len(),
959        }
960    }
961}
962
963impl Default for SafetyGateway {
964    fn default() -> Self {
965        Self::new()
966    }
967}
968
969/// Statistics from the safety gateway
970#[derive(Debug, Clone)]
971pub struct SafetyStats {
972    pub turn_count: usize,
973    pub session_count: usize,
974    pub max_per_turn: usize,
975    pub max_per_session: usize,
976    pub planning_active: bool,
977    pub preapproved_count: usize,
978}
979
980#[cfg(test)]
981mod tests {
982    use super::*;
983    use vtcode_config::core::DotfileProtectionConfig;
984
985    fn make_ctx() -> SafetyContext {
986        SafetyContext::new("test-session")
987    }
988
989    #[tokio::test]
990    async fn test_allow_read_only_tools() {
991        let gateway = SafetyGateway::new();
992        let ctx = make_ctx();
993
994        let decision = gateway
995            .check_safety(&ctx, "read_file", &serde_json::json!({"path": "/tmp/test"}))
996            .await;
997
998        assert!(decision.is_allowed());
999    }
1000
1001    #[tokio::test]
1002    async fn test_destructive_tool_needs_approval() {
1003        let gateway = SafetyGateway::new();
1004        let ctx = make_ctx();
1005
1006        let decision = gateway
1007            .check_safety(&ctx, "delete_file", &serde_json::json!({"path": "/tmp/test"}))
1008            .await;
1009
1010        assert!(decision.needs_approval());
1011    }
1012
1013    #[tokio::test]
1014    async fn test_planning_workflow_does_not_double_block_mutating() {
1015        let gateway = SafetyGateway::new();
1016        gateway.set_planning_workflow(true);
1017        let ctx = make_ctx();
1018
1019        let decision = gateway
1020            .check_safety(&ctx, "apply_patch", &serde_json::json!({"patch": "..."}))
1021            .await;
1022
1023        assert!(decision.needs_approval());
1024    }
1025
1026    #[tokio::test]
1027    async fn test_preapproved_tools_allowed() {
1028        let gateway = SafetyGateway::new();
1029        gateway.preapprove("delete_file");
1030        let ctx = make_ctx();
1031
1032        let decision = gateway
1033            .check_safety(&ctx, "delete_file", &serde_json::json!({"path": "/tmp/test"}))
1034            .await;
1035
1036        assert!(decision.is_allowed());
1037    }
1038
1039    #[tokio::test]
1040    async fn test_trust_level_bypass() {
1041        let gateway = SafetyGateway::new();
1042        let mut ctx = make_ctx();
1043        ctx.trust_level = SafetyTrustLevel::Full;
1044
1045        let decision = gateway
1046            .check_safety(&ctx, "delete_file", &serde_json::json!({"path": "/tmp/test"}))
1047            .await;
1048
1049        assert!(decision.is_allowed());
1050    }
1051
1052    #[tokio::test]
1053    async fn test_rate_limiting() {
1054        let config = SafetyGatewayConfig { rate_limit_per_second: 2, ..Default::default() };
1055        let gateway = SafetyGateway::with_config(config);
1056        let ctx = make_ctx();
1057
1058        // First two calls should succeed
1059        gateway.record_execution();
1060        gateway.record_execution();
1061
1062        // Third call should be denied
1063        let decision = gateway.check_safety(&ctx, "read_file", &serde_json::json!({})).await;
1064
1065        assert!(decision.is_denied());
1066        assert!(decision.reason().unwrap().contains("Rate limit"));
1067    }
1068
1069    #[test]
1070    fn safety_gateway_defaults_use_shared_tool_call_budget() {
1071        let config = SafetyGatewayConfig::default();
1072
1073        assert_eq!(config.max_per_turn, DEFAULT_MAX_TOOL_CALLS_PER_TURN);
1074        assert_eq!(config.max_per_session, DEFAULT_SAFETY_MAX_TOOL_CALLS_PER_SESSION);
1075    }
1076
1077    #[tokio::test]
1078    async fn test_atomic_check_and_record_rate_limited() {
1079        let config = SafetyGatewayConfig { rate_limit_per_second: 2, ..Default::default() };
1080        let gateway = SafetyGateway::with_config(config);
1081        let ctx = make_ctx();
1082
1083        let first = gateway.check_and_record(&ctx, "read_file", &serde_json::json!({})).await;
1084        assert!(first.decision.is_allowed());
1085
1086        let second = gateway.check_and_record(&ctx, "read_file", &serde_json::json!({})).await;
1087        assert!(second.decision.is_allowed());
1088
1089        let third = gateway.check_and_record(&ctx, "read_file", &serde_json::json!({})).await;
1090        assert!(third.decision.is_denied());
1091        assert!(third.retry_after.is_some());
1092        assert!(matches!(third.violation, Some(SafetyError::RateLimitExceeded { .. })));
1093    }
1094
1095    #[tokio::test]
1096    async fn test_command_policy_enforcement() {
1097        let mut commands_config = CommandsConfig::default();
1098        commands_config.deny_list.push("rm".to_string());
1099
1100        let gateway = SafetyGateway::new().with_commands_config(&commands_config);
1101        let ctx = make_ctx();
1102
1103        let decision = gateway
1104            .check_safety(&ctx, "shell", &serde_json::json!({"command": "rm -rf /"}))
1105            .await;
1106
1107        assert!(decision.is_denied());
1108    }
1109
1110    #[tokio::test]
1111    async fn test_command_session_command_policy_enforcement_with_indexed_args() {
1112        let mut commands_config = CommandsConfig::default();
1113        commands_config.deny_list.push("rm".to_string());
1114
1115        let gateway = SafetyGateway::new().with_commands_config(&commands_config);
1116        let ctx = make_ctx();
1117
1118        let decision = gateway
1119            .check_safety(
1120                &ctx,
1121                tools::UNIFIED_EXEC,
1122                &serde_json::json!({
1123                    "command.0": "rm",
1124                    "command.1": "-rf",
1125                    "command.2": "/"
1126                }),
1127            )
1128            .await;
1129
1130        assert!(decision.is_denied());
1131    }
1132
1133    #[tokio::test]
1134    async fn test_command_session_continue_command_policy_enforcement_with_input() {
1135        let mut commands_config = CommandsConfig::default();
1136        commands_config.deny_list.push("rm".to_string());
1137
1138        let gateway = SafetyGateway::new().with_commands_config(&commands_config);
1139        let ctx = make_ctx();
1140
1141        let decision = gateway
1142            .check_safety(
1143                &ctx,
1144                tools::UNIFIED_EXEC,
1145                &serde_json::json!({
1146                    "action": "continue",
1147                    "session_id": "run-123",
1148                    "input": "rm -rf /\n"
1149                }),
1150            )
1151            .await;
1152
1153        assert!(decision.is_denied());
1154    }
1155
1156    #[tokio::test]
1157    async fn test_send_pty_input_command_policy_enforcement_with_input() {
1158        let mut commands_config = CommandsConfig::default();
1159        commands_config.deny_list.push("rm".to_string());
1160
1161        let gateway = SafetyGateway::new().with_commands_config(&commands_config);
1162        let ctx = make_ctx();
1163
1164        let decision = gateway
1165            .check_safety(
1166                &ctx,
1167                tools::SEND_PTY_INPUT,
1168                &serde_json::json!({
1169                    "session_id": "run-123",
1170                    "input": "rm -rf /\n"
1171                }),
1172            )
1173            .await;
1174
1175        assert!(decision.is_denied());
1176    }
1177
1178    #[tokio::test]
1179    async fn test_apply_patch_dotfile_input_dotfile_protection_requires_approval() {
1180        let guardian = Arc::new(
1181            DotfileGuardian::new(DotfileProtectionConfig {
1182                audit_logging_enabled: false,
1183                create_backups: false,
1184                ..Default::default()
1185            })
1186            .await
1187            .expect("guardian should initialize"),
1188        );
1189        let gateway = SafetyGateway::new().with_dotfile_guardian(guardian);
1190        let ctx = make_ctx();
1191
1192        let decision = gateway
1193            .check_safety(
1194                &ctx,
1195                tools::APPLY_PATCH,
1196                &serde_json::json!({
1197                    "input": "*** Begin Patch\n*** Update File: .gitignore\n@@\n-old\n+new\n*** End Patch\n"
1198                }),
1199            )
1200            .await;
1201
1202        assert!(decision.needs_approval());
1203        assert!(decision.reason().is_some_and(|reason| reason.contains(".gitignore")));
1204    }
1205
1206    #[test]
1207    fn test_patch_file_access_targets_preserve_patch_order() {
1208        let targets = patch_file_access_targets(&serde_json::json!({
1209            "input": "*** Begin Patch\n*** Update File: src/main.rs\n@@\n-old\n+new\n*** Update File: .gitignore\n@@\n-old\n+new\n*** End Patch\n"
1210        }));
1211
1212        assert_eq!(targets.len(), 2);
1213        assert_eq!(targets[0].path, PathBuf::from("src/main.rs"));
1214        assert_eq!(targets[0].access_type, AccessType::Modify);
1215        assert_eq!(targets[1].path, PathBuf::from(".gitignore"));
1216        assert_eq!(targets[1].access_type, AccessType::Modify);
1217    }
1218
1219    #[tokio::test]
1220    async fn test_apply_patch_patch_field_dotfile_protection_requires_approval() {
1221        let guardian = Arc::new(
1222            DotfileGuardian::new(DotfileProtectionConfig {
1223                audit_logging_enabled: false,
1224                create_backups: false,
1225                ..Default::default()
1226            })
1227            .await
1228            .expect("guardian should initialize"),
1229        );
1230        let gateway = SafetyGateway::new().with_dotfile_guardian(guardian);
1231        let ctx = make_ctx();
1232
1233        let decision = gateway
1234            .check_safety(
1235                &ctx,
1236                tools::APPLY_PATCH,
1237                &serde_json::json!({
1238                    "patch": "*** Begin Patch\n*** Update File: .gitignore\n@@\n-old\n+new\n*** End Patch\n"
1239                }),
1240            )
1241            .await;
1242
1243        assert!(decision.needs_approval());
1244        assert!(decision.reason().is_some_and(|reason| reason.contains(".gitignore")));
1245    }
1246
1247    #[tokio::test]
1248    async fn test_stats_tracking() {
1249        let gateway = SafetyGateway::new();
1250        gateway.preapprove("test_tool");
1251        gateway.record_execution();
1252        gateway.record_execution();
1253
1254        let stats = gateway.get_stats();
1255        assert_eq!(stats.turn_count, 2);
1256        assert_eq!(stats.session_count, 2);
1257        assert_eq!(stats.preapproved_count, 1);
1258    }
1259
1260    #[tokio::test]
1261    async fn test_start_turn_resets_counters() {
1262        let gateway = SafetyGateway::new();
1263
1264        gateway.record_execution();
1265        gateway.record_execution();
1266
1267        let stats_before = gateway.get_stats();
1268        assert_eq!(stats_before.turn_count, 2);
1269
1270        gateway.start_turn();
1271
1272        let stats_after = gateway.get_stats();
1273        assert_eq!(stats_after.turn_count, 0);
1274        assert_eq!(stats_after.session_count, 2); // Session count preserved
1275    }
1276
1277    #[tokio::test]
1278    async fn test_increase_session_limit_updates_limit() {
1279        let gateway = SafetyGateway::new();
1280        gateway.set_limits(10, 1);
1281        let ctx = make_ctx();
1282
1283        let first = gateway.check_and_record(&ctx, "read_file", &serde_json::json!({})).await;
1284        assert!(first.decision.is_allowed());
1285
1286        let second = gateway.check_and_record(&ctx, "read_file", &serde_json::json!({})).await;
1287        assert!(second.decision.is_denied());
1288
1289        gateway.increase_session_limit(1);
1290
1291        let third = gateway.check_and_record(&ctx, "read_file", &serde_json::json!({})).await;
1292        assert!(third.decision.is_allowed());
1293    }
1294
1295    #[tokio::test]
1296    async fn control_plane_wait_does_not_consume_per_turn_budget() {
1297        let config = SafetyGatewayConfig { max_per_turn: 1, ..Default::default() };
1298        let gateway = SafetyGateway::with_config(config);
1299        let ctx = make_ctx();
1300
1301        // The one ordinary slot is consumed by a real work call.
1302        let work = gateway
1303            .check_and_record(&ctx, "read_file", &serde_json::json!({"path": "x.txt"}))
1304            .await;
1305        assert!(work.decision.is_allowed());
1306
1307        // A blocking wait is control-plane: it must not be denied by the
1308        // per-turn counter even though the budget is exhausted.
1309        let wait = gateway
1310            .check_and_record(&ctx, tools::UNIFIED_EXEC, &serde_json::json!({"action": "wait", "session_id": "run-1"}))
1311            .await;
1312        assert!(wait.decision.is_allowed(), "wait must bypass the per-turn work budget");
1313
1314        // Ordinary calls remain denied while the turn budget stays exhausted.
1315        let denied = gateway
1316            .check_and_record(&ctx, "read_file", &serde_json::json!({"path": "y.txt"}))
1317            .await;
1318        assert!(denied.decision.is_denied());
1319    }
1320
1321    #[tokio::test]
1322    async fn control_plane_calls_are_bounded_by_their_own_cap() {
1323        let config = SafetyGatewayConfig {
1324            max_per_turn: 1_000,
1325            enforce_rate_limits: false,
1326            ..Default::default()
1327        };
1328        let gateway = SafetyGateway::with_config(config);
1329        let ctx = make_ctx();
1330
1331        for _ in 0..MAX_CONTROL_PLANE_TOOL_CALLS_PER_TURN {
1332            let call = gateway
1333                .check_and_record(
1334                    &ctx,
1335                    tools::UNIFIED_EXEC,
1336                    &serde_json::json!({"action": "wait", "session_id": "run-1"}),
1337                )
1338                .await;
1339            assert!(call.decision.is_allowed());
1340        }
1341
1342        // The exempt bucket has its own ceiling; the next wait is denied.
1343        let overflow = gateway
1344            .check_and_record(&ctx, tools::UNIFIED_EXEC, &serde_json::json!({"action": "wait", "session_id": "run-1"}))
1345            .await;
1346        assert!(overflow.decision.is_denied());
1347        assert!(overflow.decision.reason().is_some_and(|reason| reason.contains("limit")));
1348    }
1349
1350    #[tokio::test]
1351    async fn control_plane_calls_still_consume_session_and_rate_budget() {
1352        // A tiny session fuse must still trip on control-plane calls.
1353        let config = SafetyGatewayConfig {
1354            max_per_turn: 1_000,
1355            max_per_session: 2,
1356            enforce_rate_limits: false,
1357            ..Default::default()
1358        };
1359        let gateway = SafetyGateway::with_config(config);
1360        let ctx = make_ctx();
1361
1362        for _ in 0..2 {
1363            let call = gateway
1364                .check_and_record(
1365                    &ctx,
1366                    tools::UNIFIED_EXEC,
1367                    &serde_json::json!({"action": "wait", "session_id": "run-1"}),
1368                )
1369                .await;
1370            assert!(call.decision.is_allowed());
1371        }
1372        let third = gateway
1373            .check_and_record(&ctx, tools::UNIFIED_EXEC, &serde_json::json!({"action": "wait", "session_id": "run-1"}))
1374            .await;
1375        assert!(third.decision.is_denied(), "session fuse must bound control-plane calls");
1376    }
1377
1378    #[tokio::test]
1379    async fn start_turn_resets_control_plane_counter() {
1380        let config = SafetyGatewayConfig {
1381            max_per_turn: 1_000,
1382            enforce_rate_limits: false,
1383            ..Default::default()
1384        };
1385        let gateway = SafetyGateway::with_config(config);
1386        let ctx = make_ctx();
1387        let wait_args = serde_json::json!({"action": "wait", "session_id": "run-1"});
1388
1389        for _ in 0..MAX_CONTROL_PLANE_TOOL_CALLS_PER_TURN {
1390            assert!(
1391                gateway
1392                    .check_and_record(&ctx, tools::UNIFIED_EXEC, &wait_args)
1393                    .await
1394                    .decision
1395                    .is_allowed()
1396            );
1397        }
1398        assert!(
1399            gateway
1400                .check_and_record(&ctx, tools::UNIFIED_EXEC, &wait_args)
1401                .await
1402                .decision
1403                .is_denied()
1404        );
1405
1406        gateway.start_turn();
1407
1408        assert!(
1409            gateway
1410                .check_and_record(&ctx, tools::UNIFIED_EXEC, &wait_args)
1411                .await
1412                .decision
1413                .is_allowed(),
1414            "a fresh turn must refill the control-plane bucket"
1415        );
1416    }
1417}
vtcode_core/tools/safety_gateway.rs

vtcode_core/tools/
safety_gateway.rs
