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 => {
333            // Single home of the unified-file mapping; mirrors the legacy
334            // file-helper arms above. Keep them in sync.
335            if file_operation_action_is(args, "write") {
336                if let Some(path) = primary_path_arg(args) {
337                    push_file_access_target(&mut targets, path, AccessType::Write);
338                }
339            } else if file_operation_action_is(args, "edit") {
340                if let Some(path) = primary_path_arg(args) {
341                    push_file_access_target(&mut targets, path, AccessType::Modify);
342                }
343            } else if file_operation_action_is(args, "delete") {
344                if let Some(path) = primary_path_arg(args) {
345                    push_file_access_target(&mut targets, path, AccessType::Delete);
346                }
347            } else 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            } else if file_operation_action_is(args, "copy") {
355                if let Some(path) = destination_path_arg(args) {
356                    push_file_access_target(&mut targets, path, AccessType::Write);
357                }
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 prefix = vtcode_commons::formatting::truncate_utf8_prefix(text, PREVIEW_LIMIT);
371        format!("{label} ({} bytes):\n{}{}", text.len(), prefix, if prefix.len() < text.len() { "..." } else { "" })
372    };
373
374    if let Some(content) = args.get("content").and_then(|value| value.as_str()) {
375        return preview_text("Content", content);
376    }
377
378    if let Some(old_str) = args.get("old_str").and_then(|value| value.as_str()) {
379        let new_str = args.get("new_str").and_then(|value| value.as_str()).unwrap_or("");
380        // Bound like the content/patch branches: an unbounded replacement
381        // would bloat the approval prompt and audit log.
382        return format!(
383            "Replace:\n  '{}'\nWith:\n  '{}'",
384            vtcode_commons::formatting::truncate_byte_budget(old_str, PREVIEW_LIMIT, "..."),
385            vtcode_commons::formatting::truncate_byte_budget(new_str, PREVIEW_LIMIT, "...")
386        );
387    }
388
389    if let Ok(Some(patch_input)) = decode_apply_patch_input(args) {
390        return preview_text("Patch", &patch_input.text);
391    }
392
393    "No details provided".to_string()
394}
395
396impl SafetyGateway {
397    /// Create a new safety gateway with default configuration
398    pub fn new() -> Self {
399        Self::with_config(SafetyGatewayConfig::default())
400    }
401
402    /// Create a new safety gateway with custom configuration
403    pub fn with_config(config: SafetyGatewayConfig) -> Self {
404        Self {
405            config: RwLock::new(config),
406            command_policy: None,
407            rate_state: Mutex::new(RateLimiterState::default()),
408            preapproved: Mutex::new(HashSet::new()),
409            dotfile_guardian: None,
410        }
411    }
412
413    /// Set the dotfile guardian for protected file access
414    pub fn with_dotfile_guardian(mut self, guardian: Arc<DotfileGuardian>) -> Self {
415        self.dotfile_guardian = Some(guardian);
416        self
417    }
418
419    /// Create and set a dotfile guardian from configuration
420    pub async fn with_dotfile_protection(mut self, config: DotfileProtectionConfig) -> anyhow::Result<Self> {
421        let guardian = DotfileGuardian::new(config).await?;
422        self.dotfile_guardian = Some(Arc::new(guardian));
423        Ok(self)
424    }
425
426    /// Set the command policy evaluator for shell command checks
427    pub fn with_command_policy(mut self, policy: CommandPolicyEvaluator) -> Self {
428        self.command_policy = Some(policy);
429        self
430    }
431
432    /// Create from commands config
433    pub fn with_commands_config(mut self, config: &CommandsConfig) -> Self {
434        self.command_policy = Some(CommandPolicyEvaluator::from_config(config));
435        self
436    }
437
438    /// Enable or disable planning workflow
439    pub fn set_planning_workflow(&self, active: bool) {
440        self.config.write().planning_active = active;
441    }
442
443    /// Set workspace trust level
444    pub fn set_workspace_trust(&self, trust: WorkspaceTrust) {
445        self.config.write().workspace_trust = trust;
446    }
447
448    /// Update rate limits
449    pub fn set_limits(&self, max_per_turn: usize, max_per_session: usize) {
450        let mut config = self.config.write();
451        config.max_per_turn = max_per_turn;
452        config.max_per_session = max_per_session;
453    }
454
455    /// Update rate-limiter thresholds.
456    pub fn set_rate_limits(&self, rate_limit_per_second: usize, rate_limit_per_minute: Option<usize>) {
457        let mut config = self.config.write();
458        if rate_limit_per_second > 0 {
459            config.rate_limit_per_second = rate_limit_per_second;
460        }
461        config.rate_limit_per_minute = rate_limit_per_minute.filter(|v| *v > 0);
462    }
463
464    /// Enable or disable rate-limit enforcement while preserving counters.
465    pub fn set_rate_limit_enforcement(&self, enabled: bool) {
466        self.config.write().enforce_rate_limits = enabled;
467    }
468
469    /// Increase session limit dynamically
470    pub fn increase_session_limit(&self, increment: usize) {
471        let mut config = self.config.write();
472        let new_max = config.max_per_session.saturating_add(increment);
473        config.max_per_session = new_max;
474        tracing::info!("Session tool limit increased to {}", new_max);
475    }
476
477    pub fn max_per_session(&self) -> usize {
478        self.config.read().max_per_session
479    }
480
481    /// Reset turn counters (call at start of new turn)
482    pub fn start_turn(&self) {
483        let mut state = self.rate_state.lock();
484        state.current_turn_count = 0;
485        state.control_plane_turn_count = 0;
486        state.calls_per_second.clear();
487        state.calls_per_minute.clear();
488    }
489
490    /// Reset all transient tool-budget state for a newly reconstructed
491    /// execution context and restore the limits supplied by the build mode.
492    pub fn reset_for_fresh_execution(&self, max_per_turn: usize, max_per_session: usize) {
493        self.set_limits(max_per_turn, max_per_session);
494        let mut state = self.rate_state.lock();
495        state.current_turn_count = 0;
496        state.control_plane_turn_count = 0;
497        state.session_count = 0;
498        state.calls_per_second.clear();
499        state.calls_per_minute.clear();
500    }
501
502    /// Preapprove a tool for this session
503    pub fn preapprove(&self, tool_name: &str) {
504        let mut preapproved = self.preapproved.lock();
505        preapproved.insert(tool_name.to_string());
506    }
507
508    /// Check if a tool is preapproved
509    pub fn is_preapproved(&self, tool_name: &str) -> bool {
510        let preapproved = self.preapproved.lock();
511        preapproved.contains(tool_name)
512    }
513
514    /// Main entry point: check safety for a tool invocation.
515    ///
516    /// Returns a [`SafetyDecision`] indicating whether execution can proceed.
517    /// Inline-delegating wrapper that returns the inner future directly to
518    /// avoid an extra coroutine state machine (audit section 16).
519    pub fn check_safety<'a>(
520        &'a self,
521        ctx: &'a SafetyContext,
522        tool_name: &'a str,
523        args: &'a Value,
524    ) -> impl Future<Output = SafetyDecision> + 'a {
525        self.check_safety_with_id(ctx, tool_name, args, None)
526    }
527
528    /// Check safety with explicit invocation ID for correlation
529    pub async fn check_safety_with_id(
530        &self,
531        ctx: &SafetyContext,
532        tool_name: &str,
533        args: &Value,
534        invocation_id: Option<ToolInvocationId>,
535    ) -> SafetyDecision {
536        let inv_id = invocation_id.map(|id| id.short()).unwrap_or_else(|| "unknown".to_string());
537
538        tracing::trace!(
539            invocation_id = %inv_id,
540            tool = %tool_name,
541            "SafetyGateway: checking safety"
542        );
543
544        if let Err(err) = self.check_rate_limits(crate::tools::tool_intent::is_turn_budget_exempt_call(tool_name, args))
545        {
546            tracing::warn!(
547                invocation_id = %inv_id,
548                error = %err,
549                "SafetyGateway: rate limit exceeded"
550            );
551            return SafetyDecision::Deny(err.to_string());
552        }
553
554        self.evaluate_non_rate_decision(ctx, tool_name, args, &inv_id).await
555    }
556
557    /// Check safety and atomically reserve a rate-limit slot on success.
558    ///
559    /// This avoids split check/record races by validating rate limits and recording
560    /// execution under a single lock acquisition.
561    pub async fn check_and_record(&self, ctx: &SafetyContext, tool_name: &str, args: &Value) -> SafetyCheckResult {
562        self.check_and_record_with_id(ctx, tool_name, args, None).await
563    }
564
565    /// Check safety with correlation ID and atomically reserve a rate-limit slot.
566    pub async fn check_and_record_with_id(
567        &self,
568        ctx: &SafetyContext,
569        tool_name: &str,
570        args: &Value,
571        invocation_id: Option<ToolInvocationId>,
572    ) -> SafetyCheckResult {
573        let inv_id = invocation_id.map(|id| id.short()).unwrap_or_else(|| "unknown".to_string());
574        tracing::trace!(
575            invocation_id = %inv_id,
576            tool = %tool_name,
577            "SafetyGateway: checking and recording safety"
578        );
579
580        let decision = self.evaluate_non_rate_decision(ctx, tool_name, args, &inv_id).await;
581
582        if decision.is_denied() {
583            return SafetyCheckResult { decision, retry_after: None, violation: None };
584        }
585
586        let control_plane = crate::tools::tool_intent::is_turn_budget_exempt_call(tool_name, args);
587        let now = Instant::now();
588        let mut state = self.rate_state.lock();
589        match self.check_rate_limits_locked(&mut state, now, control_plane) {
590            Ok(()) => {
591                self.record_execution_locked(&mut state, now, control_plane);
592                SafetyCheckResult { decision, retry_after: None, violation: None }
593            }
594            Err(err) => {
595                tracing::warn!(
596                    invocation_id = %inv_id,
597                    error = %err,
598                    "SafetyGateway: rate limit exceeded during atomic reservation"
599                );
600                SafetyCheckResult {
601                    decision: SafetyDecision::Deny(err.to_string()),
602                    retry_after: self.retry_after_for_violation(&err, &state, now),
603                    violation: Some(err),
604                }
605            }
606        }
607    }
608
609    async fn evaluate_non_rate_decision(
610        &self,
611        ctx: &SafetyContext,
612        tool_name: &str,
613        args: &Value,
614        inv_id: &str,
615    ) -> SafetyDecision {
616        // Extract the command text once: the policy check, the risk context,
617        // and the approval justification all read the same value, and each
618        // extraction re-parses the command payload (shell_words split plus
619        // action re-derivation).
620        let command = command_text_for_tool(tool_name, args);
621
622        if let Some(decision) = self.check_dotfile_protection(tool_name, args, &ctx.session_id).await {
623            tracing::info!(
624                invocation_id = %inv_id,
625                tool = %tool_name,
626                "SafetyGateway: dotfile protection triggered"
627            );
628            return decision;
629        }
630
631        if let Some(ref policy) = self.command_policy
632            && let Some(command) = command.as_deref()
633            && !policy.allows_text(command)
634        {
635            let reason = format!("Command '{command}' blocked by policy");
636            tracing::info!(
637                invocation_id = %inv_id,
638                command = %command,
639                "SafetyGateway: command policy denied"
640            );
641            return SafetyDecision::Deny(reason);
642        }
643
644        if self.is_preapproved(tool_name) {
645            tracing::trace!(
646                invocation_id = %inv_id,
647                tool = %tool_name,
648                "SafetyGateway: tool preapproved"
649            );
650            return SafetyDecision::Allow;
651        }
652
653        if ctx.trust_level.can_bypass_approval() {
654            tracing::trace!(
655                invocation_id = %inv_id,
656                tool = %tool_name,
657                trust_level = ?ctx.trust_level,
658                "SafetyGateway: trust level allows bypass"
659            );
660            return SafetyDecision::Allow;
661        }
662
663        // Classify once for the risk context, the justification, and the
664        // destructive gate below: all three read the same tool name and args,
665        // and each re-classification re-parses command payloads.
666        let intent = classify_tool_intent(tool_name, args);
667        let command = command.as_deref();
668        let risk_ctx = self.build_risk_context(tool_name, args, &intent, command);
669        let risk_level = ToolRiskScorer::calculate_risk(&risk_ctx);
670
671        if ToolRiskScorer::requires_justification(risk_level, self.config.read().approval_risk_threshold) {
672            let justification = self.build_approval_justification(tool_name, &risk_level, args, &intent, command);
673            tracing::info!(
674                invocation_id = %inv_id,
675                tool = %tool_name,
676                risk = %risk_level,
677                "SafetyGateway: requires approval"
678            );
679            return SafetyDecision::NeedsApproval(justification);
680        }
681
682        if intent.destructive {
683            let justification = format!("Tool '{tool_name}' is destructive and may modify files or execute commands.");
684            tracing::info!(
685                invocation_id = %inv_id,
686                tool = %tool_name,
687                "SafetyGateway: destructive tool requires approval"
688            );
689            return SafetyDecision::NeedsApproval(justification);
690        }
691
692        SafetyDecision::Allow
693    }
694
695    /// Record that a tool call was executed (for rate limiting)
696    pub fn record_execution(&self) {
697        let mut state = self.rate_state.lock();
698        self.record_execution_locked(&mut state, Instant::now(), false);
699    }
700
701    /// Check rate limits without recording
702    fn check_rate_limits(&self, control_plane: bool) -> Result<(), SafetyError> {
703        let mut state = self.rate_state.lock();
704        self.check_rate_limits_locked(&mut state, Instant::now(), control_plane)
705    }
706
707    fn check_rate_limits_locked(
708        &self,
709        state: &mut RateLimiterState,
710        now: Instant,
711        control_plane: bool,
712    ) -> Result<(), SafetyError> {
713        let config = self.config.read();
714        self.prune_rate_windows(state, now);
715
716        if config.enforce_rate_limits {
717            if state.calls_per_second.len() >= config.rate_limit_per_second {
718                return Err(SafetyError::RateLimitExceeded {
719                    current: state.calls_per_second.len(),
720                    max: config.rate_limit_per_second,
721                    window: "1s",
722                });
723            }
724
725            if let Some(limit) = config.rate_limit_per_minute
726                && state.calls_per_minute.len() >= limit
727            {
728                return Err(SafetyError::RateLimitExceeded {
729                    current: state.calls_per_minute.len(),
730                    max: limit,
731                    window: "60s",
732                });
733            }
734        }
735
736        if control_plane {
737            // Control-plane calls bypass the per-turn work budget but keep
738            // their own bounded cap so the exemption cannot become the
739            // dominant spend path.
740            if state.control_plane_turn_count >= MAX_CONTROL_PLANE_TOOL_CALLS_PER_TURN {
741                return Err(SafetyError::TurnLimitReached { max: MAX_CONTROL_PLANE_TOOL_CALLS_PER_TURN });
742            }
743        } else if state.current_turn_count >= config.max_per_turn {
744            return Err(SafetyError::TurnLimitReached { max: config.max_per_turn });
745        }
746
747        if state.session_count >= config.max_per_session {
748            return Err(SafetyError::SessionLimitReached { max: config.max_per_session });
749        }
750
751        Ok(())
752    }
753
754    fn record_execution_locked(&self, state: &mut RateLimiterState, now: Instant, control_plane: bool) {
755        state.calls_per_second.push_back(now);
756        state.calls_per_minute.push_back(now);
757        state.session_count = state.session_count.saturating_add(1);
758        if control_plane {
759            state.control_plane_turn_count = state.control_plane_turn_count.saturating_add(1);
760        } else {
761            state.current_turn_count = state.current_turn_count.saturating_add(1);
762        }
763    }
764
765    fn prune_rate_windows(&self, state: &mut RateLimiterState, now: Instant) {
766        while let Some(front) = state.calls_per_second.front() {
767            if now.duration_since(*front) > Duration::from_secs(1) {
768                state.calls_per_second.pop_front();
769            } else {
770                break;
771            }
772        }
773        while let Some(front) = state.calls_per_minute.front() {
774            if now.duration_since(*front) > Duration::from_secs(60) {
775                state.calls_per_minute.pop_front();
776            } else {
777                break;
778            }
779        }
780    }
781
782    fn retry_after_for_violation(
783        &self,
784        violation: &SafetyError,
785        state: &RateLimiterState,
786        now: Instant,
787    ) -> Option<Duration> {
788        match violation {
789            SafetyError::RateLimitExceeded { window: "1s", .. } => state
790                .calls_per_second
791                .front()
792                .map(|first| Duration::from_secs(1).saturating_sub(now.duration_since(*first))),
793            SafetyError::RateLimitExceeded { window: "60s", .. } => state
794                .calls_per_minute
795                .front()
796                .map(|first| Duration::from_secs(60).saturating_sub(now.duration_since(*first))),
797            _ => None,
798        }
799    }
800
801    /// Check dotfile protection for file operations.
802    /// Returns Some(SafetyDecision) if dotfile protection applies, None otherwise.
803    async fn check_dotfile_protection(
804        &self,
805        tool_name: &str,
806        args: &Value,
807        session_id: &str,
808    ) -> Option<SafetyDecision> {
809        // Use local guardian if set, otherwise try global guardian
810        let guardian = match self.dotfile_guardian.as_ref() {
811            Some(g) => g.clone(),
812            None => get_global_guardian()?,
813        };
814
815        let file_targets = file_access_targets(tool_name, args);
816        if file_targets.is_empty() {
817            return None;
818        }
819
820        let proposed_changes = proposed_changes_preview(args);
821
822        for target in file_targets {
823            if !guardian.is_protected(Path::new(&target.path)) {
824                continue;
825            }
826
827            let context = AccessContext::new(&target.path, target.access_type, tool_name, session_id)
828                .with_proposed_changes(&proposed_changes);
829
830            match guardian.request_access(&context).await {
831                Ok(ProtectionDecision::Allowed) => continue,
832                Ok(ProtectionDecision::RequiresConfirmation(req)) => {
833                    return Some(SafetyDecision::NeedsApproval(format!(
834                        "DOTFILE PROTECTION\n\n\
835                        File: {}\n\
836                        Operation: {}\n\
837                        Reason: {}\n\n\
838                        Proposed changes:\n{}\n\n\
839                        {}",
840                        req.file_path, req.access_type, req.protection_reason, req.proposed_changes, req.warning
841                    )));
842                }
843                Ok(ProtectionDecision::RequiresSecondaryAuth(req)) => {
844                    return Some(SafetyDecision::NeedsApproval(format!(
845                        "DOTFILE SECONDARY AUTHENTICATION REQUIRED\n\n\
846                        File: {} (whitelisted)\n\
847                        Operation: {}\n\
848                        Reason: {}\n\n\
849                        This file is on the whitelist but requires secondary authentication.\n\n\
850                        Proposed changes:\n{}\n\n\
851                        {}",
852                        req.file_path, req.access_type, req.protection_reason, req.proposed_changes, req.warning
853                    )));
854                }
855                Ok(ProtectionDecision::Blocked(violation)) => {
856                    return Some(SafetyDecision::Deny(format!(
857                        "DOTFILE MODIFICATION BLOCKED\n\n\
858                            File: {}\n\
859                            Reason: {}\n\n\
860                            Suggestion: {}",
861                        violation.file_path, violation.reason, violation.suggestion
862                    )));
863                }
864                Ok(ProtectionDecision::Denied(violation)) => {
865                    return Some(SafetyDecision::Deny(format!(
866                        "DOTFILE ACCESS DENIED\n\n\
867                            File: {}\n\
868                            Reason: {}\n\n\
869                            Suggestion: {}",
870                        violation.file_path, violation.reason, violation.suggestion
871                    )));
872                }
873                Err(e) => {
874                    tracing::error!("Dotfile protection check failed: {}", e);
875                    return Some(SafetyDecision::Deny(format!("Dotfile protection check failed: {e}")));
876                }
877            }
878        }
879
880        None
881    }
882
883    /// Get the dotfile guardian (if configured)
884    pub fn dotfile_guardian(&self) -> Option<&Arc<DotfileGuardian>> {
885        self.dotfile_guardian.as_ref()
886    }
887
888    /// Build risk context from tool name and arguments
889    ///
890    /// `command` is the command extracted once by the caller
891    /// ([`Self::evaluate_non_rate_decision`]); re-deriving it here would
892    /// re-parse the command payload for every risk evaluation.
893    fn build_risk_context(
894        &self,
895        tool_name: &str,
896        args: &Value,
897        intent: &ToolIntent,
898        command: Option<&str>,
899    ) -> ToolRiskContext {
900        let source = if tool_name.starts_with("mcp_") {
901            ToolSource::Mcp
902        } else if tool_name.starts_with("acp_") {
903            ToolSource::Acp
904        } else {
905            ToolSource::Internal
906        };
907
908        let risk_tool_name = action_qualified_policy_name(tool_name, Some(args)).unwrap_or(tool_name);
909
910        let mut ctx = ToolRiskContext::new(risk_tool_name.to_string(), source, self.config.read().workspace_trust);
911
912        // Set flags based on tool type
913        if intent.mutating {
914            ctx = ctx.as_write();
915        }
916        if intent.destructive {
917            ctx = ctx.as_destructive();
918        }
919
920        if ToolRiskScorer::is_network_tool(risk_tool_name) {
921            ctx = ctx.accesses_network();
922        }
923
924        // Extract command args for shell tools
925        if let Some(command) = command {
926            ctx = ctx.with_args(command.split_whitespace().map(String::from).collect());
927        }
928
929        ctx
930    }
931
932    /// Build justification message for approval prompt
933    ///
934    /// `command` is the command extracted once by the caller
935    /// ([`Self::evaluate_non_rate_decision`]); re-deriving it here would
936    /// re-parse the command payload on every approval prompt.
937    fn build_approval_justification(
938        &self,
939        tool_name: &str,
940        risk_level: &RiskLevel,
941        args: &Value,
942        intent: &ToolIntent,
943        command: Option<&str>,
944    ) -> String {
945        let mut parts = Vec::new();
946
947        parts.push(format!("Tool: {tool_name}"));
948        parts.push(format!("Risk level: {risk_level}"));
949
950        if intent.destructive {
951            parts.push("This tool may modify or delete files.".to_string());
952        }
953
954        if let Some(command) = command {
955            parts.push(format!("Command: {command}"));
956        }
957
958        let file_targets = file_access_targets(tool_name, args);
959        if let Some(target) = file_targets.first() {
960            parts.push(format!("Path: {}", target.path.display()));
961            if file_targets.len() > 1 {
962                parts.push(format!("Additional targets: {}", file_targets.len() - 1));
963            }
964        }
965
966        parts.join("\n")
967    }
968
969    /// Get current session statistics
970    pub fn get_stats(&self) -> SafetyStats {
971        let state = self.rate_state.lock();
972        let preapproved = self.preapproved.lock();
973        let config = self.config.read();
974
975        SafetyStats {
976            turn_count: state.current_turn_count,
977            session_count: state.session_count,
978            max_per_turn: config.max_per_turn,
979            max_per_session: config.max_per_session,
980            planning_active: config.planning_active,
981            preapproved_count: preapproved.len(),
982        }
983    }
984}
985
986impl Default for SafetyGateway {
987    fn default() -> Self {
988        Self::new()
989    }
990}
991
992/// Statistics from the safety gateway
993#[derive(Debug, Clone)]
994pub struct SafetyStats {
995    pub turn_count: usize,
996    pub session_count: usize,
997    pub max_per_turn: usize,
998    pub max_per_session: usize,
999    pub planning_active: bool,
1000    pub preapproved_count: usize,
1001}
1002
1003#[cfg(test)]
1004mod tests {
1005    use super::*;
1006    use vtcode_config::core::DotfileProtectionConfig;
1007
1008    fn make_ctx() -> SafetyContext {
1009        SafetyContext::new("test-session")
1010    }
1011
1012    #[tokio::test]
1013    async fn test_allow_read_only_tools() {
1014        let gateway = SafetyGateway::new();
1015        let ctx = make_ctx();
1016
1017        let decision = gateway
1018            .check_safety(&ctx, "read_file", &serde_json::json!({"path": "/tmp/test"}))
1019            .await;
1020
1021        assert!(decision.is_allowed());
1022    }
1023
1024    #[tokio::test]
1025    async fn test_destructive_tool_needs_approval() {
1026        let gateway = SafetyGateway::new();
1027        let ctx = make_ctx();
1028
1029        let decision = gateway
1030            .check_safety(&ctx, "delete_file", &serde_json::json!({"path": "/tmp/test"}))
1031            .await;
1032
1033        assert!(decision.needs_approval());
1034    }
1035
1036    #[tokio::test]
1037    async fn test_planning_workflow_does_not_double_block_mutating() {
1038        let gateway = SafetyGateway::new();
1039        gateway.set_planning_workflow(true);
1040        let ctx = make_ctx();
1041
1042        let decision = gateway
1043            .check_safety(&ctx, "apply_patch", &serde_json::json!({"patch": "..."}))
1044            .await;
1045
1046        assert!(decision.needs_approval());
1047    }
1048
1049    /// The command text is extracted once per check and threaded to the
1050    /// policy, risk, and justification consumers. This pins the observable
1051    /// behavior (identical policy verdict, risk sentence, and justification
1052    /// command line as re-derivation would produce) after the extraction
1053    /// stopped being repeated at each site.
1054    #[tokio::test]
1055    async fn command_extraction_feeds_policy_risk_and_justification_identically() {
1056        let policy = CommandPolicyEvaluator::from_config(&CommandsConfig {
1057            allow_list: vec!["cargo".into()],
1058            ..Default::default()
1059        });
1060        let gateway = SafetyGateway::new().with_command_policy(policy);
1061
1062        // Allowed by policy: command is still classified for risk (mutating
1063        // verifier) and surfaces verbatim in any justification.
1064        let allowed = gateway
1065            .check_safety(&make_ctx(), tools::EXEC_COMMAND, &serde_json::json!({"command": "cargo test"}))
1066            .await;
1067        assert!(allowed.is_allowed());
1068
1069        // Denied by policy: the denial reports the same single command.
1070        let denied = gateway
1071            .check_safety(&make_ctx(), tools::EXEC_COMMAND, &serde_json::json!({"command": "rm -rf /"}))
1072            .await;
1073        let SafetyDecision::Deny(reason) = denied else {
1074            panic!("policy denial expected");
1075        };
1076        assert!(reason.contains("rm -rf /"), "denial must carry the extracted command: {reason}");
1077        assert!(reason.starts_with("Command 'rm -rf /' blocked by policy"), "unexpected denial shape: {reason}");
1078
1079        // Justification path: the risk classifier and the justification see
1080        // the same extracted command (marked unsafe by the gateway config).
1081        gateway.config.write().approval_risk_threshold.clone_from(&RiskLevel::Low);
1082        let justified = gateway
1083            .check_safety(&make_ctx(), tools::EXEC_COMMAND, &serde_json::json!({"command": "cargo publish"}))
1084            .await;
1085        if let SafetyDecision::NeedsApproval(justification) = justified {
1086            assert!(
1087                justification.contains("Command: cargo publish"),
1088                "justification must carry the extracted command: {justification}"
1089            );
1090        }
1091        gateway.config.write().approval_risk_threshold = RiskLevel::High;
1092    }
1093
1094    #[test]
1095    fn build_risk_context_uses_caller_supplied_command() {
1096        let gateway = SafetyGateway::new();
1097        let args = serde_json::json!({"command": "rg --files"});
1098        let intent = classify_tool_intent(tools::EXEC_COMMAND, &args);
1099
1100        let with_command = gateway.build_risk_context(tools::EXEC_COMMAND, &args, &intent, Some("rg --files"));
1101        assert_eq!(with_command.command_args, vec!["rg".to_string(), "--files".to_string()]);
1102
1103        let without_command = gateway.build_risk_context(tools::EXEC_COMMAND, &args, &intent, None);
1104        assert!(without_command.command_args.is_empty());
1105    }
1106
1107    #[tokio::test]
1108    async fn test_preapproved_tools_allowed() {
1109        let gateway = SafetyGateway::new();
1110        gateway.preapprove("delete_file");
1111        let ctx = make_ctx();
1112
1113        let decision = gateway
1114            .check_safety(&ctx, "delete_file", &serde_json::json!({"path": "/tmp/test"}))
1115            .await;
1116
1117        assert!(decision.is_allowed());
1118    }
1119
1120    #[tokio::test]
1121    async fn test_trust_level_bypass() {
1122        let gateway = SafetyGateway::new();
1123        let mut ctx = make_ctx();
1124        ctx.trust_level = SafetyTrustLevel::Full;
1125
1126        let decision = gateway
1127            .check_safety(&ctx, "delete_file", &serde_json::json!({"path": "/tmp/test"}))
1128            .await;
1129
1130        assert!(decision.is_allowed());
1131    }
1132
1133    #[tokio::test]
1134    async fn test_rate_limiting() {
1135        let config = SafetyGatewayConfig { rate_limit_per_second: 2, ..Default::default() };
1136        let gateway = SafetyGateway::with_config(config);
1137        let ctx = make_ctx();
1138
1139        // First two calls should succeed
1140        gateway.record_execution();
1141        gateway.record_execution();
1142
1143        // Third call should be denied
1144        let decision = gateway.check_safety(&ctx, "read_file", &serde_json::json!({})).await;
1145
1146        assert!(decision.is_denied());
1147        assert!(decision.reason().unwrap().contains("Rate limit"));
1148    }
1149
1150    #[test]
1151    fn safety_gateway_defaults_use_shared_tool_call_budget() {
1152        let config = SafetyGatewayConfig::default();
1153
1154        assert_eq!(config.max_per_turn, DEFAULT_MAX_TOOL_CALLS_PER_TURN);
1155        assert_eq!(config.max_per_session, DEFAULT_SAFETY_MAX_TOOL_CALLS_PER_SESSION);
1156    }
1157
1158    #[tokio::test]
1159    async fn test_atomic_check_and_record_rate_limited() {
1160        let config = SafetyGatewayConfig { rate_limit_per_second: 2, ..Default::default() };
1161        let gateway = SafetyGateway::with_config(config);
1162        let ctx = make_ctx();
1163
1164        let first = gateway.check_and_record(&ctx, "read_file", &serde_json::json!({})).await;
1165        assert!(first.decision.is_allowed());
1166
1167        let second = gateway.check_and_record(&ctx, "read_file", &serde_json::json!({})).await;
1168        assert!(second.decision.is_allowed());
1169
1170        let third = gateway.check_and_record(&ctx, "read_file", &serde_json::json!({})).await;
1171        assert!(third.decision.is_denied());
1172        assert!(third.retry_after.is_some());
1173        assert!(matches!(third.violation, Some(SafetyError::RateLimitExceeded { .. })));
1174    }
1175
1176    #[tokio::test]
1177    async fn test_command_policy_enforcement() {
1178        let mut commands_config = CommandsConfig::default();
1179        commands_config.deny_list.push("rm".to_string());
1180
1181        let gateway = SafetyGateway::new().with_commands_config(&commands_config);
1182        let ctx = make_ctx();
1183
1184        let decision = gateway
1185            .check_safety(&ctx, "shell", &serde_json::json!({"command": "rm -rf /"}))
1186            .await;
1187
1188        assert!(decision.is_denied());
1189    }
1190
1191    #[tokio::test]
1192    async fn test_command_session_command_policy_enforcement_with_indexed_args() {
1193        let mut commands_config = CommandsConfig::default();
1194        commands_config.deny_list.push("rm".to_string());
1195
1196        let gateway = SafetyGateway::new().with_commands_config(&commands_config);
1197        let ctx = make_ctx();
1198
1199        let decision = gateway
1200            .check_safety(
1201                &ctx,
1202                tools::UNIFIED_EXEC,
1203                &serde_json::json!({
1204                    "command.0": "rm",
1205                    "command.1": "-rf",
1206                    "command.2": "/"
1207                }),
1208            )
1209            .await;
1210
1211        assert!(decision.is_denied());
1212    }
1213
1214    #[tokio::test]
1215    async fn test_command_session_continue_command_policy_enforcement_with_input() {
1216        let mut commands_config = CommandsConfig::default();
1217        commands_config.deny_list.push("rm".to_string());
1218
1219        let gateway = SafetyGateway::new().with_commands_config(&commands_config);
1220        let ctx = make_ctx();
1221
1222        let decision = gateway
1223            .check_safety(
1224                &ctx,
1225                tools::UNIFIED_EXEC,
1226                &serde_json::json!({
1227                    "action": "continue",
1228                    "session_id": "run-123",
1229                    "input": "rm -rf /\n"
1230                }),
1231            )
1232            .await;
1233
1234        assert!(decision.is_denied());
1235    }
1236
1237    #[tokio::test]
1238    async fn test_send_pty_input_command_policy_enforcement_with_input() {
1239        let mut commands_config = CommandsConfig::default();
1240        commands_config.deny_list.push("rm".to_string());
1241
1242        let gateway = SafetyGateway::new().with_commands_config(&commands_config);
1243        let ctx = make_ctx();
1244
1245        let decision = gateway
1246            .check_safety(
1247                &ctx,
1248                tools::SEND_PTY_INPUT,
1249                &serde_json::json!({
1250                    "session_id": "run-123",
1251                    "input": "rm -rf /\n"
1252                }),
1253            )
1254            .await;
1255
1256        assert!(decision.is_denied());
1257    }
1258
1259    #[tokio::test]
1260    async fn test_apply_patch_dotfile_input_dotfile_protection_requires_approval() {
1261        let guardian = Arc::new(
1262            DotfileGuardian::new(DotfileProtectionConfig {
1263                audit_logging_enabled: false,
1264                create_backups: false,
1265                ..Default::default()
1266            })
1267            .await
1268            .expect("guardian should initialize"),
1269        );
1270        let gateway = SafetyGateway::new().with_dotfile_guardian(guardian);
1271        let ctx = make_ctx();
1272
1273        let decision = gateway
1274            .check_safety(
1275                &ctx,
1276                tools::APPLY_PATCH,
1277                &serde_json::json!({
1278                    "input": "*** Begin Patch\n*** Update File: .gitignore\n@@\n-old\n+new\n*** End Patch\n"
1279                }),
1280            )
1281            .await;
1282
1283        assert!(decision.needs_approval());
1284        assert!(decision.reason().is_some_and(|reason| reason.contains(".gitignore")));
1285    }
1286
1287    #[test]
1288    fn test_patch_file_access_targets_preserve_patch_order() {
1289        let targets = patch_file_access_targets(&serde_json::json!({
1290            "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"
1291        }));
1292
1293        assert_eq!(targets.len(), 2);
1294        assert_eq!(targets[0].path, PathBuf::from("src/main.rs"));
1295        assert_eq!(targets[0].access_type, AccessType::Modify);
1296        assert_eq!(targets[1].path, PathBuf::from(".gitignore"));
1297        assert_eq!(targets[1].access_type, AccessType::Modify);
1298    }
1299
1300    #[test]
1301    fn unified_file_targets_mirror_legacy_file_helpers() {
1302        let write = file_access_targets(tools::UNIFIED_FILE, &serde_json::json!({"action": "write", "path": "a.txt"}));
1303        assert_eq!(write.len(), 1);
1304        assert_eq!(write[0].path, PathBuf::from("a.txt"));
1305        assert_eq!(write[0].access_type, AccessType::Write);
1306
1307        let moved = file_access_targets(
1308            tools::UNIFIED_FILE,
1309            &serde_json::json!({"action": "move", "path": "a.txt", "destination": "b.txt"}),
1310        );
1311        assert_eq!(moved.len(), 2);
1312        assert_eq!(moved[0].access_type, AccessType::Modify);
1313        assert_eq!(moved[1].access_type, AccessType::Write);
1314
1315        let unknown = file_access_targets(tools::UNIFIED_FILE, &serde_json::json!({"action": "read", "path": "a.txt"}));
1316        assert!(unknown.is_empty());
1317    }
1318
1319    #[test]
1320    fn proposed_changes_preview_bounds_old_and_new_str() {
1321        // Short replacements render exactly as before (no truncation
1322        // marker); long ones stay bounded with a marker; a multi-byte
1323        // char on the 500-byte boundary must not panic.
1324        let short = proposed_changes_preview(&serde_json::json!({"old_str": "a", "new_str": "b"}));
1325        assert_eq!(short, "Replace:\n  'a'\nWith:\n  'b'");
1326
1327        let long_old = "a".repeat(600);
1328        let bounded = proposed_changes_preview(&serde_json::json!({"old_str": long_old, "new_str": ""}));
1329        assert!(bounded.contains("..."));
1330        assert!(bounded.len() < 600, "preview must stay bounded: {}", bounded.len());
1331
1332        let boundary_old = format!("{}{}", "a".repeat(499), "日".repeat(10));
1333        let boundary = proposed_changes_preview(&serde_json::json!({"old_str": boundary_old, "new_str": "b"}));
1334        assert!(boundary.contains("..."));
1335        assert!(boundary.ends_with("With:\n  'b'"));
1336    }
1337
1338    #[test]
1339    fn proposed_changes_preview_never_splits_utf8_boundary() {
1340        // 500-byte cut lands mid-`日` (3 bytes): the old `&text[..500]`
1341        // panicked; the canonical prefix rounds down to the boundary.
1342        let content = format!("{}{}", "a".repeat(499), "日".repeat(10));
1343        let preview = proposed_changes_preview(&serde_json::json!({"content": content}));
1344        assert!(preview.contains("Content ("));
1345        assert!(preview.ends_with("..."));
1346        assert!(preview.len() < content.len() + 64);
1347    }
1348
1349    #[tokio::test]
1350    async fn test_apply_patch_patch_field_dotfile_protection_requires_approval() {
1351        let guardian = Arc::new(
1352            DotfileGuardian::new(DotfileProtectionConfig {
1353                audit_logging_enabled: false,
1354                create_backups: false,
1355                ..Default::default()
1356            })
1357            .await
1358            .expect("guardian should initialize"),
1359        );
1360        let gateway = SafetyGateway::new().with_dotfile_guardian(guardian);
1361        let ctx = make_ctx();
1362
1363        let decision = gateway
1364            .check_safety(
1365                &ctx,
1366                tools::APPLY_PATCH,
1367                &serde_json::json!({
1368                    "patch": "*** Begin Patch\n*** Update File: .gitignore\n@@\n-old\n+new\n*** End Patch\n"
1369                }),
1370            )
1371            .await;
1372
1373        assert!(decision.needs_approval());
1374        assert!(decision.reason().is_some_and(|reason| reason.contains(".gitignore")));
1375    }
1376
1377    #[tokio::test]
1378    async fn test_stats_tracking() {
1379        let gateway = SafetyGateway::new();
1380        gateway.preapprove("test_tool");
1381        gateway.record_execution();
1382        gateway.record_execution();
1383
1384        let stats = gateway.get_stats();
1385        assert_eq!(stats.turn_count, 2);
1386        assert_eq!(stats.session_count, 2);
1387        assert_eq!(stats.preapproved_count, 1);
1388    }
1389
1390    #[tokio::test]
1391    async fn test_start_turn_resets_counters() {
1392        let gateway = SafetyGateway::new();
1393
1394        gateway.record_execution();
1395        gateway.record_execution();
1396
1397        let stats_before = gateway.get_stats();
1398        assert_eq!(stats_before.turn_count, 2);
1399
1400        gateway.start_turn();
1401
1402        let stats_after = gateway.get_stats();
1403        assert_eq!(stats_after.turn_count, 0);
1404        assert_eq!(stats_after.session_count, 2); // Session count preserved
1405    }
1406
1407    #[tokio::test]
1408    async fn test_increase_session_limit_updates_limit() {
1409        let gateway = SafetyGateway::new();
1410        gateway.set_limits(10, 1);
1411        let ctx = make_ctx();
1412
1413        let first = gateway.check_and_record(&ctx, "read_file", &serde_json::json!({})).await;
1414        assert!(first.decision.is_allowed());
1415
1416        let second = gateway.check_and_record(&ctx, "read_file", &serde_json::json!({})).await;
1417        assert!(second.decision.is_denied());
1418
1419        gateway.increase_session_limit(1);
1420
1421        let third = gateway.check_and_record(&ctx, "read_file", &serde_json::json!({})).await;
1422        assert!(third.decision.is_allowed());
1423    }
1424
1425    #[tokio::test]
1426    async fn control_plane_wait_does_not_consume_per_turn_budget() {
1427        let config = SafetyGatewayConfig { max_per_turn: 1, ..Default::default() };
1428        let gateway = SafetyGateway::with_config(config);
1429        let ctx = make_ctx();
1430
1431        // The one ordinary slot is consumed by a real work call.
1432        let work = gateway
1433            .check_and_record(&ctx, "read_file", &serde_json::json!({"path": "x.txt"}))
1434            .await;
1435        assert!(work.decision.is_allowed());
1436
1437        // A blocking wait is control-plane: it must not be denied by the
1438        // per-turn counter even though the budget is exhausted.
1439        let wait = gateway
1440            .check_and_record(&ctx, tools::UNIFIED_EXEC, &serde_json::json!({"action": "wait", "session_id": "run-1"}))
1441            .await;
1442        assert!(wait.decision.is_allowed(), "wait must bypass the per-turn work budget");
1443
1444        // Ordinary calls remain denied while the turn budget stays exhausted.
1445        let denied = gateway
1446            .check_and_record(&ctx, "read_file", &serde_json::json!({"path": "y.txt"}))
1447            .await;
1448        assert!(denied.decision.is_denied());
1449    }
1450
1451    #[tokio::test]
1452    async fn control_plane_calls_are_bounded_by_their_own_cap() {
1453        let config = SafetyGatewayConfig {
1454            max_per_turn: 1_000,
1455            enforce_rate_limits: false,
1456            ..Default::default()
1457        };
1458        let gateway = SafetyGateway::with_config(config);
1459        let ctx = make_ctx();
1460
1461        for _ in 0..MAX_CONTROL_PLANE_TOOL_CALLS_PER_TURN {
1462            let call = gateway
1463                .check_and_record(
1464                    &ctx,
1465                    tools::UNIFIED_EXEC,
1466                    &serde_json::json!({"action": "wait", "session_id": "run-1"}),
1467                )
1468                .await;
1469            assert!(call.decision.is_allowed());
1470        }
1471
1472        // The exempt bucket has its own ceiling; the next wait is denied.
1473        let overflow = gateway
1474            .check_and_record(&ctx, tools::UNIFIED_EXEC, &serde_json::json!({"action": "wait", "session_id": "run-1"}))
1475            .await;
1476        assert!(overflow.decision.is_denied());
1477        assert!(overflow.decision.reason().is_some_and(|reason| reason.contains("limit")));
1478    }
1479
1480    #[tokio::test]
1481    async fn control_plane_calls_still_consume_session_and_rate_budget() {
1482        // A tiny session fuse must still trip on control-plane calls.
1483        let config = SafetyGatewayConfig {
1484            max_per_turn: 1_000,
1485            max_per_session: 2,
1486            enforce_rate_limits: false,
1487            ..Default::default()
1488        };
1489        let gateway = SafetyGateway::with_config(config);
1490        let ctx = make_ctx();
1491
1492        for _ in 0..2 {
1493            let call = gateway
1494                .check_and_record(
1495                    &ctx,
1496                    tools::UNIFIED_EXEC,
1497                    &serde_json::json!({"action": "wait", "session_id": "run-1"}),
1498                )
1499                .await;
1500            assert!(call.decision.is_allowed());
1501        }
1502        let third = gateway
1503            .check_and_record(&ctx, tools::UNIFIED_EXEC, &serde_json::json!({"action": "wait", "session_id": "run-1"}))
1504            .await;
1505        assert!(third.decision.is_denied(), "session fuse must bound control-plane calls");
1506    }
1507
1508    #[tokio::test]
1509    async fn start_turn_resets_control_plane_counter() {
1510        let config = SafetyGatewayConfig {
1511            max_per_turn: 1_000,
1512            enforce_rate_limits: false,
1513            ..Default::default()
1514        };
1515        let gateway = SafetyGateway::with_config(config);
1516        let ctx = make_ctx();
1517        let wait_args = serde_json::json!({"action": "wait", "session_id": "run-1"});
1518
1519        for _ in 0..MAX_CONTROL_PLANE_TOOL_CALLS_PER_TURN {
1520            assert!(
1521                gateway
1522                    .check_and_record(&ctx, tools::UNIFIED_EXEC, &wait_args)
1523                    .await
1524                    .decision
1525                    .is_allowed()
1526            );
1527        }
1528        assert!(
1529            gateway
1530                .check_and_record(&ctx, tools::UNIFIED_EXEC, &wait_args)
1531                .await
1532                .decision
1533                .is_denied()
1534        );
1535
1536        gateway.start_turn();
1537
1538        assert!(
1539            gateway
1540                .check_and_record(&ctx, tools::UNIFIED_EXEC, &wait_args)
1541                .await
1542                .decision
1543                .is_allowed(),
1544            "a fresh turn must refill the control-plane bucket"
1545        );
1546    }
1547}
vtcode_core/tools/safety_gateway.rs

vtcode_core/tools/
safety_gateway.rs
