Skip to main content

vtcode_core/tools/registry/
mod.rs

1//! Tool registry and function declarations
2
3mod approval_recorder;
4mod assembly;
5mod availability_facade;
6mod builder;
7mod builtins;
8mod cache;
9mod catalog_facade;
10mod cgp_facade;
11mod circuit_breaker;
12mod commands_facade;
13mod config_helpers;
14mod dual_output;
15mod error;
16mod execution_facade;
17mod execution_history;
18mod execution_kernel;
19mod execution_request;
20mod executors;
21pub mod file_helpers;
22mod file_monitor_facade;
23mod harness;
24mod harness_facade;
25mod history_facade;
26mod inventory;
27mod inventory_facade;
28mod justification;
29mod justification_extractor;
30pub mod labels;
31mod maintenance;
32mod mcp_facade;
33mod mcp_helpers;
34mod metrics_facade;
35mod optimization_facade;
36mod output_processing;
37mod plan_mode_checks;
38mod plan_mode_facade;
39mod policy;
40mod policy_facade;
41mod progress_facade;
42mod pty;
43mod pty_facade;
44mod registration;
45mod registration_facade;
46mod resiliency;
47mod resiliency_facade;
48mod risk_scorer;
49mod runtime_config_facade;
50mod sandbox_facade;
51mod scheduler_facade;
52mod search_runtime_facade;
53mod shell_policy;
54mod shell_policy_facade;
55mod spooler_facade;
56mod subagent_facade;
57mod telemetry;
58mod timeout;
59mod timeout_category;
60mod timeout_facade;
61mod tool_catalog_facade;
62mod tool_executor_impl;
63mod unified_actions;
64mod utils;
65
66pub use approval_recorder::ApprovalRecorder;
67pub use cgp_facade::CgpRuntimeMode;
68pub use cgp_facade::native_cgp_tool_factory;
69pub use cgp_facade::wrap_registered_native_tool;
70pub use error::{ToolErrorType, ToolExecutionError, classify_error};
71pub use execution_history::{HarnessContextSnapshot, ToolExecutionHistory, ToolExecutionRecord};
72pub use execution_kernel::ToolPreflightOutcome;
73pub use execution_request::{
74    ExecSettlementMode, ExecutionPolicySnapshot, ToolExecutionOutcome, ToolExecutionRequest,
75};
76pub use harness::HarnessContext;
77pub use justification::{ApprovalPattern, JustificationManager, ToolJustification};
78pub use justification_extractor::JustificationExtractor;
79pub use pty::{PtySessionGuard, PtySessionManager};
80pub use registration::{
81    NativeCgpToolFactory, ToolCatalogSource, ToolExecutorFn, ToolHandler, ToolMetadata,
82    ToolRegistration,
83};
84pub use resiliency::{ResiliencyContext, ToolFailureTracker};
85pub use risk_scorer::{RiskLevel, ToolRiskContext, ToolRiskScorer, ToolSource, WorkspaceTrust};
86pub use shell_policy::ShellPolicyChecker;
87pub use telemetry::ToolTelemetryEvent;
88pub use timeout::{
89    AdaptiveTimeoutTuning, ToolLatencyStats, ToolTimeoutCategory, ToolTimeoutPolicy,
90};
91pub use tool_catalog_facade::SessionToolCatalogState;
92pub(crate) use unified_actions::{UnifiedExecAction, UnifiedFileAction, UnifiedSearchAction};
93
94use assembly::ToolAssembly;
95use inventory::ToolInventory;
96use policy::ToolPolicyGateway;
97use utils::normalize_tool_output;
98
99use crate::tools::exec_session::ExecSessionManager;
100use crate::tools::handlers::PlanModeState;
101pub(super) use crate::tools::pty::PtyManager;
102use crate::tools::result::ToolResult as SplitToolResult;
103use crate::tools::safety_gateway::SafetyGateway;
104use parking_lot::Mutex; // Use parking_lot for better performance
105use rustc_hash::FxHashMap;
106use std::sync::Arc;
107
108use crate::mcp::McpClient;
109use crate::subagents::SubagentController;
110use crate::tools::edited_file_monitor::EditedFileMonitor;
111use std::sync::RwLock;
112
113/// Callback for tool progress and output streaming
114pub type ToolProgressCallback = Arc<dyn Fn(&str, &str) + Send + Sync>;
115
116use super::traits::Tool;
117#[cfg(test)]
118use crate::config::types::CapabilityLevel;
119
120/// Default window size for loop detection.
121const DEFAULT_LOOP_DETECT_WINDOW: usize = 5;
122
123#[derive(Clone)]
124pub struct ToolRegistry {
125    inventory: ToolInventory,
126    edited_file_monitor: Arc<EditedFileMonitor>,
127    policy_gateway: Arc<tokio::sync::Mutex<ToolPolicyGateway>>,
128    pty_sessions: PtySessionManager,
129    exec_sessions: ExecSessionManager,
130    mcp_client: Arc<RwLock<Option<Arc<McpClient>>>>,
131    mcp_tool_index: Arc<tokio::sync::RwLock<FxHashMap<String, Vec<String>>>>,
132    mcp_reverse_index: Arc<tokio::sync::RwLock<FxHashMap<String, String>>>,
133    timeout_policy: Arc<RwLock<ToolTimeoutPolicy>>,
134    execution_history: ToolExecutionHistory,
135    harness_context: HarnessContext,
136
137    // Mutable runtime state wrapped for concurrent access
138    resiliency: Arc<Mutex<ResiliencyContext>>,
139
140    /// MP-3: Circuit breaker for MCP client failures
141    mcp_circuit_breaker: Arc<circuit_breaker::McpCircuitBreaker>,
142    /// Shared per-tool circuit breaker state used by the runloop.
143    shared_circuit_breaker: Arc<RwLock<Option<Arc<crate::tools::circuit_breaker::CircuitBreaker>>>>,
144    initialized: Arc<std::sync::atomic::AtomicBool>,
145    // Security & Identity
146    shell_policy: Arc<RwLock<ShellPolicyChecker>>,
147    runtime_sandbox_config: Arc<RwLock<vtcode_config::SandboxConfig>>,
148    agent_type: Arc<RwLock<String>>,
149    // PTY Session Management
150    active_pty_sessions: Arc<RwLock<Option<Arc<std::sync::atomic::AtomicUsize>>>>,
151
152    // Caching
153    cached_available_tools: Arc<RwLock<Option<Vec<String>>>>,
154    /// Callback for streaming tool output and progress
155    progress_callback: Arc<RwLock<Option<ToolProgressCallback>>>,
156    // Performance Observability
157    /// Total tool calls made in current session
158    pub(crate) tool_call_counter: Arc<std::sync::atomic::AtomicU64>,
159    /// Total PTY poll iterations (for monitoring CPU usage)
160    pub(crate) pty_poll_counter: Arc<std::sync::atomic::AtomicU64>,
161    /// Shared metrics collector for reliability and execution observability
162    metrics: Arc<crate::metrics::MetricsCollector>,
163
164    // PERFORMANCE OPTIMIZATIONS - Actually integrated into the real registry
165    /// Memory pool for reducing allocations in hot paths
166    memory_pool: Arc<crate::core::memory_pool::MemoryPool>,
167    /// Hot cache for frequently accessed tools (reduces HashMap lookups)
168    hot_tool_cache: Arc<parking_lot::RwLock<lru::LruCache<String, Arc<dyn Tool>>>>,
169    /// Optimization configuration
170    optimization_config: vtcode_config::OptimizationConfig,
171
172    /// Output spooler for dynamic context discovery (large outputs to files)
173    output_spooler: Arc<super::output_spooler::ToolOutputSpooler>,
174
175    /// Shared Plan Mode state (plan file tracking, active flag) for enter/exit tools
176    /// and read-only enforcement. This is the single source of truth for plan mode;
177    /// use `is_plan_mode()` / `enable_plan_mode()` / `disable_plan_mode()` as accessors.
178    plan_mode_state: PlanModeState,
179    /// Canonical safety gateway shared across registry execution surfaces.
180    safety_gateway: Arc<SafetyGateway>,
181    /// Active CGP runtime mode for wrapping registrations added after startup.
182    cgp_runtime_mode: Arc<RwLock<Option<CgpRuntimeMode>>>,
183    /// Canonical manifest-driven tool assembly used by routing, catalog projections, and policy sync.
184    tool_assembly: Arc<RwLock<ToolAssembly>>,
185    /// Registry-owned tool catalog snapshot cache shared by harnesses.
186    tool_catalog_state: Arc<SessionToolCatalogState>,
187    /// Shared subagent controller when the session enables delegated child agents.
188    subagent_controller: Arc<RwLock<Option<Arc<SubagentController>>>>,
189    /// Session-scoped scheduled prompts for interactive loops and cron tools.
190    session_scheduler: Arc<tokio::sync::Mutex<crate::scheduler::SessionScheduler>>,
191}
192
193#[derive(Debug, Clone, Copy, PartialEq, Eq)]
194pub enum ToolPermissionDecision {
195    Allow,
196    Deny,
197    Prompt,
198}
199
200#[cfg(test)]
201mod tests {
202    use super::*;
203    use crate::config::TimeoutsConfig;
204    use crate::config::ToolDocumentationMode as ConfigToolDocumentationMode;
205    use crate::config::ToolPolicy as ConfigToolPolicy;
206    use crate::config::ToolsConfig;
207    use crate::constants::tools;
208    use crate::tool_policy::ToolPolicy;
209    use crate::tool_policy::ToolPolicyConfig;
210    use crate::tools::handlers::{SessionSurface, SessionToolsConfig, ToolModelCapabilities};
211    use crate::tools::registry::mcp_helpers::normalize_mcp_tool_identifier;
212    use anyhow::Result;
213    use async_trait::async_trait;
214    use futures::future::BoxFuture;
215    use serde_json::Value;
216    use serde_json::json;
217    use std::fs;
218    use std::time::Duration;
219    use tempfile::TempDir;
220
221    const CUSTOM_TOOL_NAME: &str = "custom_test_tool";
222    const SLOW_TIMEOUT_TOOL_NAME: &str = "slow_timeout_test_tool";
223    const REENTRANT_TOOL_NAME: &str = "reentrant_guard_test_tool";
224    const MUTUAL_REENTRANT_TOOL_A: &str = "mutual_reentrant_tool_a";
225    const MUTUAL_REENTRANT_TOOL_B: &str = "mutual_reentrant_tool_b";
226
227    struct CustomEchoTool;
228    struct SlowTimeoutTool;
229
230    #[async_trait]
231    impl Tool for CustomEchoTool {
232        async fn execute(&self, args: Value) -> Result<Value> {
233            Ok(json!({
234                "success": true,
235                "args": args,
236            }))
237        }
238
239        fn name(&self) -> &str {
240            CUSTOM_TOOL_NAME
241        }
242
243        fn description(&self) -> &str {
244            "Custom echo tool for testing"
245        }
246    }
247
248    #[async_trait]
249    impl Tool for SlowTimeoutTool {
250        async fn execute(&self, _args: Value) -> Result<Value> {
251            tokio::time::sleep(Duration::from_millis(1_100)).await;
252            Ok(json!({
253                "ok": true,
254            }))
255        }
256
257        fn name(&self) -> &str {
258            SLOW_TIMEOUT_TOOL_NAME
259        }
260
261        fn description(&self) -> &str {
262            "Tool that intentionally exceeds low timeout ceilings"
263        }
264    }
265
266    fn reentrant_tool_executor<'a>(
267        registry: &'a ToolRegistry,
268        args: Value,
269    ) -> BoxFuture<'a, Result<Value>> {
270        Box::pin(async move { registry.execute_tool_ref(REENTRANT_TOOL_NAME, &args).await })
271    }
272
273    fn mutual_reentrant_tool_a_executor<'a>(
274        registry: &'a ToolRegistry,
275        args: Value,
276    ) -> BoxFuture<'a, Result<Value>> {
277        Box::pin(async move {
278            registry
279                .execute_tool_ref(MUTUAL_REENTRANT_TOOL_B, &args)
280                .await
281        })
282    }
283
284    fn mutual_reentrant_tool_b_executor<'a>(
285        registry: &'a ToolRegistry,
286        args: Value,
287    ) -> BoxFuture<'a, Result<Value>> {
288        Box::pin(async move {
289            registry
290                .execute_tool_ref(MUTUAL_REENTRANT_TOOL_A, &args)
291                .await
292        })
293    }
294
295    #[tokio::test]
296    async fn registers_builtin_tools() -> Result<()> {
297        let temp_dir = TempDir::new()?;
298        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
299        let available = registry.available_tools().await;
300
301        assert!(available.contains(&tools::UNIFIED_SEARCH.to_string()));
302        assert!(available.contains(&tools::UNIFIED_FILE.to_string()));
303        assert!(available.contains(&tools::UNIFIED_EXEC.to_string()));
304        assert!(!available.contains(&tools::READ_FILE.to_string()));
305        assert!(!available.contains(&tools::RUN_PTY_CMD.to_string()));
306        Ok(())
307    }
308
309    #[tokio::test]
310    async fn request_user_input_aliases_are_not_registered() -> Result<()> {
311        let temp_dir = TempDir::new()?;
312        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
313
314        assert!(registry.get_tool(tools::REQUEST_USER_INPUT).is_some());
315        assert!(registry.get_tool(tools::ASK_QUESTIONS).is_none());
316        assert!(registry.get_tool(tools::ASK_USER_QUESTION).is_none());
317
318        Ok(())
319    }
320
321    #[tokio::test]
322    async fn public_tool_projections_stay_in_sync() -> Result<()> {
323        let temp_dir = TempDir::new()?;
324        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
325        let config = SessionToolsConfig::full_public(
326            SessionSurface::Interactive,
327            CapabilityLevel::CodeSearch,
328            ConfigToolDocumentationMode::Full,
329            ToolModelCapabilities::default(),
330        );
331
332        let names = registry
333            .public_tool_names(SessionSurface::Interactive, CapabilityLevel::CodeSearch)
334            .await;
335        let schema_names = registry
336            .schema_entries(config.clone())
337            .await
338            .into_iter()
339            .map(|entry| entry.name)
340            .collect::<Vec<_>>();
341        let declaration_names = registry
342            .function_declarations(config.clone())
343            .await
344            .into_iter()
345            .map(|entry| entry.name)
346            .collect::<Vec<_>>();
347        let mut model_tool_names = registry
348            .model_tools(config)
349            .await
350            .into_iter()
351            .map(|tool| tool.function_name().to_string())
352            .collect::<Vec<_>>();
353
354        model_tool_names.sort();
355
356        assert_eq!(schema_names, names);
357        assert_eq!(declaration_names, names);
358        assert_eq!(model_tool_names, names);
359
360        Ok(())
361    }
362
363    #[tokio::test]
364    async fn public_routing_keeps_aliases_private_and_rebuilds_on_dynamic_updates() -> Result<()> {
365        let temp_dir = TempDir::new()?;
366        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
367        registry.allow_all_tools().await?;
368
369        let test_file = temp_dir.path().join("alias-read.txt");
370        fs::write(&test_file, "via alias\n")?;
371
372        let public_names = registry
373            .public_tool_names(SessionSurface::Interactive, CapabilityLevel::CodeSearch)
374            .await;
375        assert!(!public_names.contains(&tools::READ_FILE.to_string()));
376
377        let read_result = registry
378            .execute_public_tool_ref(
379                tools::READ_FILE,
380                &json!({"path": test_file.to_string_lossy().to_string()}),
381            )
382            .await?;
383        assert_eq!(read_result["success"].as_bool(), Some(true));
384        assert_eq!(read_result["content"].as_str(), Some("via alias"));
385
386        registry
387            .register_tool(
388                ToolRegistration::from_tool_instance(
389                    CUSTOM_TOOL_NAME,
390                    CapabilityLevel::CodeSearch,
391                    CustomEchoTool,
392                )
393                .with_description("Custom echo tool for testing")
394                .with_parameter_schema(json!({
395                    "type": "object",
396                    "properties": {
397                        "input": {"type": "string"}
398                    }
399                }))
400                .with_aliases(["custom_tool_alias"]),
401            )
402            .await?;
403
404        let public_names = registry
405            .public_tool_names(SessionSurface::Interactive, CapabilityLevel::CodeSearch)
406            .await;
407        assert!(public_names.contains(&CUSTOM_TOOL_NAME.to_string()));
408        assert!(!public_names.contains(&"custom_tool_alias".to_string()));
409
410        let schema_names = registry
411            .schema_entries(SessionToolsConfig::full_public(
412                SessionSurface::Interactive,
413                CapabilityLevel::CodeSearch,
414                ConfigToolDocumentationMode::Full,
415                ToolModelCapabilities::default(),
416            ))
417            .await
418            .into_iter()
419            .map(|entry| entry.name)
420            .collect::<Vec<_>>();
421        assert!(schema_names.contains(&CUSTOM_TOOL_NAME.to_string()));
422        assert!(!schema_names.contains(&"custom_tool_alias".to_string()));
423
424        let dynamic_result = registry
425            .execute_public_tool_ref("custom_tool_alias", &json!({"input": "value"}))
426            .await?;
427        assert_eq!(dynamic_result["success"].as_bool(), Some(true));
428
429        registry.unregister_tool(CUSTOM_TOOL_NAME).await?;
430
431        let public_names = registry
432            .public_tool_names(SessionSurface::Interactive, CapabilityLevel::CodeSearch)
433            .await;
434        assert!(!public_names.contains(&CUSTOM_TOOL_NAME.to_string()));
435
436        let err = registry
437            .execute_public_tool_ref("custom_tool_alias", &json!({"input": "value"}))
438            .await
439            .expect_err("alias should be removed with the registration");
440        assert!(err.to_string().contains("Unknown tool"));
441
442        Ok(())
443    }
444
445    #[tokio::test]
446    async fn allows_registering_custom_tools() -> Result<()> {
447        let temp_dir = TempDir::new()?;
448        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
449
450        registry
451            .register_tool(
452                ToolRegistration::from_tool_instance(
453                    CUSTOM_TOOL_NAME,
454                    CapabilityLevel::CodeSearch,
455                    CustomEchoTool,
456                )
457                .with_parameter_schema(json!({
458                    "type": "object",
459                    "properties": {
460                        "input": {"type": "string"}
461                    }
462                })),
463            )
464            .await?;
465
466        registry.allow_all_tools().await?;
467
468        let available = registry.available_tools().await;
469        assert!(available.contains(&CUSTOM_TOOL_NAME.to_string()));
470
471        let response = registry
472            .execute_tool(CUSTOM_TOOL_NAME, json!({"input": "value"}))
473            .await?;
474        assert!(response["success"].as_bool().unwrap_or(false));
475        Ok(())
476    }
477
478    #[tokio::test]
479    async fn dynamic_tool_registration_keeps_policy_catalog_in_sync() -> Result<()> {
480        let temp_dir = TempDir::new()?;
481        let policy_path = temp_dir.path().join("tool-policy.json");
482        let policy_manager =
483            crate::tool_policy::ToolPolicyManager::new_with_config_path(&policy_path).await?;
484        let registry =
485            ToolRegistry::new_with_custom_policy(temp_dir.path().to_path_buf(), policy_manager)
486                .await;
487
488        registry
489            .register_tool(ToolRegistration::from_tool_instance(
490                CUSTOM_TOOL_NAME,
491                CapabilityLevel::CodeSearch,
492                CustomEchoTool,
493            ))
494            .await?;
495
496        let config: ToolPolicyConfig = serde_json::from_str(&fs::read_to_string(&policy_path)?)?;
497        assert!(
498            config
499                .available_tools
500                .contains(&CUSTOM_TOOL_NAME.to_string())
501        );
502
503        registry.unregister_tool(CUSTOM_TOOL_NAME).await?;
504
505        let config: ToolPolicyConfig = serde_json::from_str(&fs::read_to_string(&policy_path)?)?;
506        assert!(
507            !config
508                .available_tools
509                .contains(&CUSTOM_TOOL_NAME.to_string())
510        );
511
512        Ok(())
513    }
514
515    #[tokio::test]
516    async fn executes_prevalidated_tool_path() -> Result<()> {
517        let temp_dir = TempDir::new()?;
518        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
519
520        registry
521            .register_tool(ToolRegistration::from_tool_instance(
522                CUSTOM_TOOL_NAME,
523                CapabilityLevel::CodeSearch,
524                CustomEchoTool,
525            ))
526            .await?;
527        registry.allow_all_tools().await?;
528
529        let args = json!({"input": "value"});
530        let response = registry
531            .execute_tool_ref_prevalidated(CUSTOM_TOOL_NAME, &args)
532            .await?;
533        assert!(response["success"].as_bool().unwrap_or(false));
534
535        Ok(())
536    }
537
538    #[tokio::test]
539    async fn harness_exec_reuses_public_output_normalization() -> Result<()> {
540        let temp_dir = TempDir::new()?;
541        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
542
543        let response = registry
544            .execute_harness_unified_exec(json!({
545                "action": "run",
546                "command": "printf vtcode",
547                "tty": false,
548                "yield_time_ms": 1000
549            }))
550            .await?;
551
552        assert_eq!(response["output"].as_str(), Some("vtcode"));
553        assert!(response.get("stdout").is_none());
554
555        Ok(())
556    }
557
558    #[tokio::test]
559    async fn harness_terminal_runs_retain_completed_sessions_until_close() -> Result<()> {
560        let temp_dir = TempDir::new()?;
561        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
562
563        let response = registry
564            .execute_harness_unified_exec_terminal_run(json!({
565                "action": "run",
566                "command": ["/bin/sh", "-lc", "printf vtcode-terminal"],
567                "tty": true,
568                "yield_time_ms": 200,
569            }))
570            .await?;
571
572        let session_id = response["session_id"]
573            .as_str()
574            .expect("terminal run should expose session_id")
575            .to_string();
576        assert_eq!(response["exit_code"], 0);
577        assert_eq!(response["output"].as_str(), Some("vtcode-terminal"));
578        assert_eq!(
579            registry.harness_exec_session_completed(&session_id).await?,
580            Some(0)
581        );
582
583        registry.close_harness_exec_session(&session_id).await?;
584        registry
585            .harness_exec_session_completed(&session_id)
586            .await
587            .unwrap_err();
588
589        Ok(())
590    }
591
592    fn delayed_exec_args(tty: bool, yield_time_ms: u64) -> Value {
593        json!({
594            "action": "run",
595            "command": ["/bin/sh", "-lc", "printf first && sleep 0.2 && printf second"],
596            "shell": "/bin/sh",
597            "tty": tty,
598            "yield_time_ms": yield_time_ms,
599        })
600    }
601
602    fn long_running_exec_args(tty: bool, yield_time_ms: u64) -> Value {
603        json!({
604            "action": "run",
605            "command": [
606                "/bin/sh",
607                "-lc",
608                "sleep 0.4 && printf second && sleep 0.4 && printf third && sleep 0.4 && printf done"
609            ],
610            "shell": "/bin/sh",
611            "tty": tty,
612            "yield_time_ms": yield_time_ms,
613        })
614    }
615
616    #[tokio::test]
617    async fn prevalidated_exec_mode_settles_noninteractive_run() -> Result<()> {
618        let temp_dir = TempDir::new()?;
619        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
620        registry.allow_all_tools().await?;
621
622        let response = registry
623            .execute_public_tool_ref_prevalidated_with_mode(
624                tools::UNIFIED_EXEC,
625                &delayed_exec_args(false, 50),
626                ExecSettlementMode::SettleNonInteractive,
627            )
628            .await?;
629
630        let output = response["output"]
631            .as_str()
632            .expect("settled exec output should be text");
633        assert!(output.contains("first"));
634        assert!(output.contains("second"));
635        assert_eq!(response["exit_code"], 0);
636        assert!(response.get("next_continue_args").is_none());
637
638        Ok(())
639    }
640
641    #[tokio::test]
642    async fn prevalidated_exec_mode_settles_pipe_poll_until_exit() -> Result<()> {
643        let temp_dir = TempDir::new()?;
644        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
645        registry.allow_all_tools().await?;
646
647        let initial = registry
648            .execute_harness_unified_exec(delayed_exec_args(false, 50))
649            .await?;
650        let session_id = initial["session_id"]
651            .as_str()
652            .expect("partial run should expose session_id")
653            .to_string();
654        let initial_output = initial["output"].as_str().unwrap_or_default().to_string();
655        if initial.get("next_continue_args").is_none() {
656            assert_eq!(initial["exit_code"], 0);
657            assert!(initial_output.contains("first"));
658            assert!(initial_output.contains("second"));
659            return Ok(());
660        }
661        assert!(initial.get("exit_code").is_none());
662
663        let response = registry
664            .execute_public_tool_ref_prevalidated_with_mode(
665                tools::UNIFIED_EXEC,
666                &json!({
667                    "action": "poll",
668                    "session_id": session_id,
669                    "yield_time_ms": 50,
670                }),
671                ExecSettlementMode::SettleNonInteractive,
672            )
673            .await?;
674
675        assert_eq!(response["exit_code"], 0);
676        let settled_output = response["output"]
677            .as_str()
678            .expect("settled poll output should be text");
679        assert!(initial_output.contains("second") || settled_output.contains("second"));
680        assert!(response.get("next_continue_args").is_none());
681
682        Ok(())
683    }
684
685    #[tokio::test]
686    async fn prevalidated_exec_mode_keeps_interactive_runs_manual() -> Result<()> {
687        let temp_dir = TempDir::new()?;
688        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
689        registry.allow_all_tools().await?;
690
691        let response = registry
692            .execute_public_tool_ref_prevalidated_with_mode(
693                tools::UNIFIED_EXEC,
694                &delayed_exec_args(true, 50),
695                ExecSettlementMode::SettleNonInteractive,
696            )
697            .await?;
698
699        assert!(response.get("next_continue_args").is_some());
700        assert!(response.get("exit_code").is_none());
701
702        let session_id = response["session_id"]
703            .as_str()
704            .expect("interactive run should expose session_id")
705            .to_string();
706        registry
707            .execute_harness_unified_exec(json!({
708                "action": "close",
709                "session_id": session_id,
710            }))
711            .await?;
712
713        Ok(())
714    }
715
716    #[tokio::test]
717    async fn unified_exec_run_preserves_requested_session_id_for_follow_up_calls() -> Result<()> {
718        let temp_dir = TempDir::new()?;
719        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
720        registry.allow_all_tools().await?;
721
722        let mut run_args = long_running_exec_args(true, 10);
723        run_args
724            .as_object_mut()
725            .expect("run args should be an object")
726            .insert("session_id".to_string(), json!("check_sh"));
727
728        let initial = registry.execute_harness_unified_exec(run_args).await?;
729        assert_eq!(initial["session_id"], "check_sh");
730        assert_eq!(
731            initial["next_continue_args"],
732            json!({ "session_id": "check_sh" })
733        );
734
735        let response = registry
736            .execute_harness_unified_exec(json!({
737                "action": "poll",
738                "session_id": "check_sh",
739                "yield_time_ms": 10,
740            }))
741            .await?;
742
743        assert!(response.get("output").is_some());
744        assert!(
745            response.get("exit_code").is_some() || response.get("next_continue_args").is_some()
746        );
747
748        if response.get("exit_code").is_none() {
749            registry
750                .execute_harness_unified_exec(json!({
751                    "action": "close",
752                    "session_id": "check_sh",
753                }))
754                .await?;
755        }
756
757        Ok(())
758    }
759
760    #[tokio::test]
761    async fn active_exec_continuations_bypass_identical_call_loop_detection() -> Result<()> {
762        let temp_dir = TempDir::new()?;
763        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
764        registry.allow_all_tools().await?;
765        registry.execution_history.set_loop_detection_limits(5, 2);
766
767        let initial = registry
768            .execute_harness_unified_exec(long_running_exec_args(false, 10))
769            .await?;
770        let session_id = initial["session_id"]
771            .as_str()
772            .expect("partial run should expose session_id")
773            .to_string();
774        let continue_args = json!({
775            "action": "continue",
776            "session_id": session_id,
777            "yield_time_ms": 10,
778        });
779
780        let first = registry
781            .execute_public_tool_ref_prevalidated(tools::UNIFIED_EXEC, &continue_args)
782            .await?;
783        assert_ne!(first.get("loop_detected"), Some(&json!(true)));
784
785        let second = registry
786            .execute_public_tool_ref_prevalidated(tools::UNIFIED_EXEC, &continue_args)
787            .await?;
788        assert_ne!(second.get("loop_detected"), Some(&json!(true)));
789
790        let third = registry
791            .execute_public_tool_ref_prevalidated(tools::UNIFIED_EXEC, &continue_args)
792            .await?;
793        assert_ne!(third.get("loop_detected"), Some(&json!(true)));
794        assert!(
795            third.get("exit_code").is_some() || third.get("next_continue_args").is_some(),
796            "continuation should either remain active or complete cleanly"
797        );
798
799        Ok(())
800    }
801
802    #[tokio::test]
803    async fn unified_exec_accepts_compact_session_alias_for_poll() -> Result<()> {
804        let temp_dir = TempDir::new()?;
805        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
806        registry.allow_all_tools().await?;
807
808        let initial = registry
809            .execute_harness_unified_exec(long_running_exec_args(true, 10))
810            .await?;
811        let session_id = initial["session_id"]
812            .as_str()
813            .expect("partial run should expose session_id")
814            .to_string();
815
816        let response = registry
817            .execute_harness_unified_exec(json!({
818                "s": session_id,
819                "yield_time_ms": 10
820            }))
821            .await?;
822
823        assert!(response.get("output").is_some());
824        assert!(
825            response.get("exit_code").is_some() || response.get("next_continue_args").is_some()
826        );
827
828        Ok(())
829    }
830
831    #[tokio::test]
832    async fn unified_exec_inspect_accepts_compact_session_alias() -> Result<()> {
833        let temp_dir = TempDir::new()?;
834        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
835        registry.allow_all_tools().await?;
836
837        let initial = registry
838            .execute_harness_unified_exec(long_running_exec_args(true, 10))
839            .await?;
840        let session_id = initial["session_id"]
841            .as_str()
842            .expect("partial run should expose session_id")
843            .to_string();
844
845        let response = registry
846            .execute_harness_unified_exec(json!({
847                "action": "inspect",
848                "s": session_id,
849                "head_lines": 1,
850                "tail_lines": 0
851            }))
852            .await?;
853
854        assert_eq!(response["content_type"], "exec_inspect");
855        assert!(response["output"].is_string());
856        assert!(response.get("session_id").is_some());
857
858        Ok(())
859    }
860
861    #[tokio::test]
862    async fn mutating_tools_clear_recent_read_reuse_history() -> Result<()> {
863        let temp_dir = TempDir::new()?;
864        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
865        registry.allow_all_tools().await?;
866        registry.execution_history.set_loop_detection_limits(5, 2);
867
868        let test_file = temp_dir.path().join("test.txt");
869        fs::write(&test_file, "original")?;
870
871        let read_args = json!({
872            "path": test_file.to_string_lossy(),
873            "max_bytes": 1000,
874        });
875        let write_args = json!({
876            "path": test_file.to_string_lossy(),
877            "content": "modified",
878            "mode": "overwrite",
879        });
880
881        let first = registry
882            .execute_tool_ref(tools::READ_FILE, &read_args)
883            .await?;
884        assert_eq!(first["content"], "original");
885
886        let second = registry
887            .execute_tool(tools::READ_FILE, read_args.clone())
888            .await?;
889        assert_eq!(second["content"], "original");
890
891        let write_result = registry.execute_tool(tools::WRITE_FILE, write_args).await?;
892        assert_eq!(write_result["success"], json!(true));
893
894        let after_write = registry.execute_tool(tools::READ_FILE, read_args).await?;
895        assert_eq!(after_write["content"], "modified");
896        assert_ne!(after_write.get("reused_recent_result"), Some(&json!(true)));
897        assert_ne!(after_write.get("loop_detected"), Some(&json!(true)));
898
899        Ok(())
900    }
901
902    #[tokio::test]
903    async fn prevalidated_execution_enforces_plan_mode_guards() -> Result<()> {
904        let temp_dir = TempDir::new()?;
905        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
906        registry.allow_all_tools().await?;
907        registry.enable_plan_mode();
908        registry.plan_mode_state().enable();
909
910        let blocked_path = temp_dir.path().join("blocked.txt");
911        let args = json!({
912            "path": blocked_path.to_string_lossy().to_string(),
913            "content": "should-not-write"
914        });
915
916        let err = registry
917            .execute_tool_ref_prevalidated(tools::WRITE_FILE, &args)
918            .await
919            .expect_err("plan mode should block prevalidated mutating tool call");
920        assert!(err.to_string().contains("plan mode"));
921        assert!(!blocked_path.exists());
922
923        Ok(())
924    }
925
926    #[tokio::test]
927    async fn prevalidated_execution_allows_task_tracker_in_plan_mode() -> Result<()> {
928        let temp_dir = TempDir::new()?;
929        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
930        registry.allow_all_tools().await?;
931        registry.enable_plan_mode();
932        registry.plan_mode_state().enable();
933
934        let plans_dir = temp_dir.path().join(".vtcode").join("plans");
935        fs::create_dir_all(&plans_dir)?;
936        let plan_file = plans_dir.join("adaptive-test.md");
937        fs::write(&plan_file, "# Adaptive test\n")?;
938        registry
939            .plan_mode_state()
940            .set_plan_file(Some(plan_file))
941            .await;
942
943        let args = json!({"action": "create", "items": ["Track step"]});
944
945        let response = registry
946            .execute_tool_ref_prevalidated(tools::TASK_TRACKER, &args)
947            .await
948            .expect("task_tracker should be allowed in plan mode");
949        assert_eq!(response["status"], "created");
950
951        Ok(())
952    }
953
954    #[tokio::test]
955    async fn preflight_rejects_removed_exec_code_alias() -> Result<()> {
956        let temp_dir = TempDir::new()?;
957        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
958
959        let err = registry
960            .preflight_validate_call(
961                "exec_code",
962                &json!({
963                    "command": "echo vtcode"
964                }),
965            )
966            .expect_err("exec_code alias should be rejected");
967        assert!(err.to_string().contains("Unknown tool"));
968
969        Ok(())
970    }
971
972    #[tokio::test]
973    async fn preflight_normalizes_humanized_exec_label_to_unified_exec() -> Result<()> {
974        let temp_dir = TempDir::new()?;
975        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
976
977        let outcome = registry.preflight_validate_call(
978            "Exec code",
979            &json!({
980                "command": "echo vtcode"
981            }),
982        )?;
983        assert_eq!(outcome.normalized_tool_name, tools::UNIFIED_EXEC);
984
985        Ok(())
986    }
987
988    #[tokio::test]
989    async fn preflight_normalizes_execute_code_alias_to_unified_exec() -> Result<()> {
990        let temp_dir = TempDir::new()?;
991        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
992
993        let outcome = registry.preflight_validate_call(
994            tools::EXECUTE_CODE,
995            &json!({
996                "code": "print('vtcode')"
997            }),
998        )?;
999        assert_eq!(outcome.normalized_tool_name, tools::UNIFIED_EXEC);
1000
1001        Ok(())
1002    }
1003
1004    #[tokio::test]
1005    async fn preflight_normalizes_raw_apply_patch_payload_to_input_object() -> Result<()> {
1006        let temp_dir = TempDir::new()?;
1007        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
1008        let patch = "*** Begin Patch\n*** End Patch\n";
1009
1010        let outcome = registry.preflight_validate_call(tools::APPLY_PATCH, &json!(patch))?;
1011
1012        assert_eq!(outcome.normalized_tool_name, tools::APPLY_PATCH);
1013        assert_eq!(outcome.effective_args, json!({ "input": patch }));
1014
1015        Ok(())
1016    }
1017
1018    #[tokio::test]
1019    async fn preflight_normalizes_repo_browser_aliases() -> Result<()> {
1020        let temp_dir = TempDir::new()?;
1021        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
1022
1023        let read_outcome = registry.preflight_validate_call(
1024            "repo_browser.read_file",
1025            &json!({"path": "vtcode-core/src/lib.rs"}),
1026        )?;
1027        assert_eq!(read_outcome.normalized_tool_name, tools::UNIFIED_FILE);
1028
1029        let list_err = registry
1030            .preflight_validate_call(
1031                "repo_browser.list_files",
1032                &json!({"path": "vtcode-core/src"}),
1033            )
1034            .expect_err("repo_browser.list_files alias should be rejected");
1035        assert!(list_err.to_string().contains("Unknown tool"));
1036
1037        Ok(())
1038    }
1039
1040    #[tokio::test]
1041    async fn preflight_prefers_direct_harness_browse_tool_routes() -> Result<()> {
1042        let temp_dir = TempDir::new()?;
1043        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
1044
1045        let read_outcome = registry.preflight_validate_call(
1046            tools::READ_FILE,
1047            &json!({"path": "vtcode-core/src/lib.rs"}),
1048        )?;
1049        assert_eq!(read_outcome.normalized_tool_name, tools::READ_FILE);
1050
1051        let list_outcome = registry.preflight_validate_call(
1052            tools::LIST_FILES,
1053            &json!({"path": "vtcode-core/src", "page": 1, "per_page": 20}),
1054        )?;
1055        assert_eq!(list_outcome.normalized_tool_name, tools::LIST_FILES);
1056
1057        Ok(())
1058    }
1059
1060    #[tokio::test]
1061    async fn preflight_normalizes_plan_mode_force_on_aliases() -> Result<()> {
1062        let temp_dir = TempDir::new()?;
1063        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
1064
1065        let on_outcome = registry.preflight_validate_call("plan_on", &json!({}))?;
1066        assert_eq!(on_outcome.normalized_tool_name, tools::ENTER_PLAN_MODE);
1067
1068        let slash_outcome = registry.preflight_validate_call("/plan", &json!({}))?;
1069        assert_eq!(slash_outcome.normalized_tool_name, tools::ENTER_PLAN_MODE);
1070
1071        Ok(())
1072    }
1073
1074    #[tokio::test]
1075    async fn preflight_normalizes_plan_mode_force_off_aliases() -> Result<()> {
1076        let temp_dir = TempDir::new()?;
1077        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
1078
1079        let off_outcome = registry.preflight_validate_call("mode_edit", &json!({}))?;
1080        assert_eq!(off_outcome.normalized_tool_name, tools::EXIT_PLAN_MODE);
1081
1082        let slash_outcome = registry.preflight_validate_call("/edit", &json!({}))?;
1083        assert_eq!(slash_outcome.normalized_tool_name, tools::EXIT_PLAN_MODE);
1084
1085        Ok(())
1086    }
1087
1088    #[tokio::test]
1089    async fn suggest_fallback_prefers_unified_exec_for_exec_code_alias() -> Result<()> {
1090        let temp_dir = TempDir::new()?;
1091        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
1092
1093        let fallback = registry.suggest_fallback_tool("exec_code").await;
1094        assert_eq!(fallback.as_deref(), Some(tools::UNIFIED_EXEC));
1095
1096        Ok(())
1097    }
1098
1099    #[tokio::test]
1100    async fn suggest_fallback_prefers_unified_exec_for_humanized_exec_label() -> Result<()> {
1101        let temp_dir = TempDir::new()?;
1102        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
1103
1104        let fallback = registry.suggest_fallback_tool("Exec code").await;
1105        assert_eq!(fallback.as_deref(), Some(tools::UNIFIED_EXEC));
1106
1107        Ok(())
1108    }
1109
1110    #[tokio::test]
1111    async fn suggest_fallback_returns_none_for_task_tracker() -> Result<()> {
1112        let temp_dir = TempDir::new()?;
1113        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
1114
1115        let fallback = registry.suggest_fallback_tool(tools::TASK_TRACKER).await;
1116        assert!(fallback.is_none());
1117
1118        Ok(())
1119    }
1120
1121    #[tokio::test]
1122    async fn suggest_fallback_returns_none_for_unknown_tool() -> Result<()> {
1123        let temp_dir = TempDir::new()?;
1124        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
1125
1126        let fallback = registry.suggest_fallback_tool("not_a_real_tool").await;
1127        assert!(fallback.is_none());
1128
1129        Ok(())
1130    }
1131
1132    #[tokio::test]
1133    async fn execute_public_repo_browser_alias_routes_through_public_assembly() -> Result<()> {
1134        let temp_dir = TempDir::new()?;
1135        fs::write(temp_dir.path().join("public-route.txt"), "public route\n")?;
1136        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
1137        registry.allow_all_tools().await?;
1138
1139        let response = registry
1140            .execute_public_tool_ref(
1141                "repo_browser.read_file",
1142                &json!({"path": "public-route.txt"}),
1143            )
1144            .await?;
1145
1146        assert_eq!(response["path"].as_str(), Some("public-route.txt"));
1147        assert!(
1148            response["content"]
1149                .as_str()
1150                .is_some_and(|content| content.contains("public route"))
1151        );
1152
1153        Ok(())
1154    }
1155
1156    #[tokio::test]
1157    async fn set_tool_policy_normalizes_public_aliases() -> Result<()> {
1158        let temp_dir = TempDir::new()?;
1159        let policy_path = temp_dir.path().join("tool-policy.json");
1160        let policy_manager =
1161            crate::tool_policy::ToolPolicyManager::new_with_config_path(&policy_path).await?;
1162        let registry =
1163            ToolRegistry::new_with_custom_policy(temp_dir.path().to_path_buf(), policy_manager)
1164                .await;
1165
1166        registry
1167            .set_tool_policy("Exec code", ToolPolicy::Deny)
1168            .await?;
1169
1170        assert_eq!(
1171            registry.get_tool_policy("Exec code").await,
1172            ToolPolicy::Deny
1173        );
1174        assert_eq!(
1175            registry.get_tool_policy(tools::UNIFIED_EXEC).await,
1176            ToolPolicy::Deny
1177        );
1178
1179        Ok(())
1180    }
1181
1182    #[tokio::test]
1183    async fn apply_config_policies_prefers_explicit_canonical_public_names() -> Result<()> {
1184        let temp_dir = TempDir::new()?;
1185        let policy_path = temp_dir.path().join("tool-policy.json");
1186        let policy_manager =
1187            crate::tool_policy::ToolPolicyManager::new_with_config_path(&policy_path).await?;
1188        let registry =
1189            ToolRegistry::new_with_custom_policy(temp_dir.path().to_path_buf(), policy_manager)
1190                .await;
1191
1192        let mut config = ToolsConfig::default();
1193        config.policies.clear();
1194        config
1195            .policies
1196            .insert(tools::UNIFIED_FILE.to_string(), ConfigToolPolicy::Allow);
1197        config
1198            .policies
1199            .insert(tools::READ_FILE.to_string(), ConfigToolPolicy::Deny);
1200
1201        registry.apply_config_policies(&config).await?;
1202
1203        assert_eq!(
1204            registry.get_tool_policy(tools::UNIFIED_FILE).await,
1205            ToolPolicy::Allow
1206        );
1207        assert_eq!(
1208            registry.get_tool_policy(tools::READ_FILE).await,
1209            ToolPolicy::Allow
1210        );
1211
1212        Ok(())
1213    }
1214
1215    #[tokio::test]
1216    async fn persisted_approval_cache_round_trips_through_registry() -> Result<()> {
1217        let temp_dir = TempDir::new()?;
1218        let policy_path = temp_dir.path().join("tool-policy.json");
1219        let policy_manager =
1220            crate::tool_policy::ToolPolicyManager::new_with_config_path(&policy_path).await?;
1221        let registry =
1222            ToolRegistry::new_with_custom_policy(temp_dir.path().to_path_buf(), policy_manager)
1223                .await;
1224
1225        registry.persist_approval_cache_key("read_file").await?;
1226
1227        assert!(registry.has_persisted_approval("read_file").await);
1228
1229        let manager =
1230            crate::tool_policy::ToolPolicyManager::new_with_config_path(&policy_path).await?;
1231        assert!(manager.has_approval_cache_key("read_file"));
1232
1233        Ok(())
1234    }
1235
1236    #[tokio::test]
1237    async fn public_alias_resolution_stays_consistent_across_execution_preflight_and_policy()
1238    -> Result<()> {
1239        let temp_dir = TempDir::new()?;
1240        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
1241
1242        registry
1243            .register_tool(
1244                ToolRegistration::from_tool_instance(
1245                    CUSTOM_TOOL_NAME,
1246                    CapabilityLevel::CodeSearch,
1247                    CustomEchoTool,
1248                )
1249                .with_description("Custom echo tool for routing parity tests")
1250                .with_parameter_schema(json!({
1251                    "type": "object",
1252                    "properties": {
1253                        "input": {"type": "string"}
1254                    }
1255                }))
1256                .with_permission(ToolPolicy::Allow)
1257                .with_aliases(["custom tool"]),
1258            )
1259            .await?;
1260
1261        let preflight =
1262            registry.preflight_validate_call("Custom Tool", &json!({"input": "value"}))?;
1263        assert_eq!(preflight.normalized_tool_name, CUSTOM_TOOL_NAME);
1264
1265        assert_eq!(
1266            registry.evaluate_tool_policy("Custom Tool").await?,
1267            ToolPermissionDecision::Allow
1268        );
1269
1270        let response = registry
1271            .execute_public_tool_ref("Custom Tool", &json!({"input": "value"}))
1272            .await?;
1273        assert_eq!(response["success"].as_bool(), Some(true));
1274
1275        Ok(())
1276    }
1277
1278    #[tokio::test]
1279    async fn safe_mode_prompt_uses_behavior_metadata() -> Result<()> {
1280        let temp_dir = TempDir::new()?;
1281        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
1282        registry.allow_all_tools().await?;
1283        registry.set_enforce_safe_mode_prompts(true).await;
1284
1285        assert_eq!(
1286            registry.evaluate_tool_policy(tools::UNIFIED_SEARCH).await?,
1287            ToolPermissionDecision::Allow
1288        );
1289        assert_eq!(
1290            registry.evaluate_tool_policy(tools::UNIFIED_EXEC).await?,
1291            ToolPermissionDecision::Prompt
1292        );
1293        assert_eq!(
1294            registry.evaluate_tool_policy(tools::APPLY_PATCH).await?,
1295            ToolPermissionDecision::Prompt
1296        );
1297
1298        Ok(())
1299    }
1300
1301    #[tokio::test]
1302    async fn mcp_policy_paths_resolve_model_visible_aliases() -> Result<()> {
1303        fn noop_executor<'a>(
1304            _registry: &'a ToolRegistry,
1305            _args: Value,
1306        ) -> BoxFuture<'a, Result<Value>> {
1307            Box::pin(async { Ok(json!({"success": true})) })
1308        }
1309
1310        let temp_dir = TempDir::new()?;
1311        let policy_path = temp_dir.path().join("tool-policy.json");
1312        let policy_manager =
1313            crate::tool_policy::ToolPolicyManager::new_with_config_path(&policy_path).await?;
1314        let registry =
1315            ToolRegistry::new_with_custom_policy(temp_dir.path().to_path_buf(), policy_manager)
1316                .await;
1317
1318        let public_name = crate::tools::mcp::model_visible_mcp_tool_name("context7", "search");
1319        registry
1320            .register_tool(
1321                ToolRegistration::new(
1322                    "mcp::context7::search",
1323                    CapabilityLevel::Basic,
1324                    false,
1325                    noop_executor,
1326                )
1327                .with_description("Fake MCP search tool")
1328                .with_parameter_schema(json!({"type": "object"}))
1329                .with_permission(ToolPolicy::Prompt)
1330                .with_aliases([public_name.clone()])
1331                .with_llm_visibility(false),
1332            )
1333            .await?;
1334
1335        registry
1336            .mcp_tool_index
1337            .write()
1338            .await
1339            .insert("context7".to_string(), vec!["search".to_string()]);
1340        registry
1341            .mcp_reverse_index
1342            .write()
1343            .await
1344            .insert("search".to_string(), "context7".to_string());
1345
1346        registry
1347            .persist_mcp_tool_policy(&public_name, ToolPolicy::Allow)
1348            .await?;
1349
1350        let manager =
1351            crate::tool_policy::ToolPolicyManager::new_with_config_path(&policy_path).await?;
1352        assert_eq!(
1353            manager.get_mcp_tool_policy("context7", "search"),
1354            ToolPolicy::Allow
1355        );
1356
1357        assert_eq!(
1358            registry.evaluate_tool_policy(&public_name).await?,
1359            ToolPermissionDecision::Allow
1360        );
1361        assert_eq!(
1362            registry
1363                .evaluate_tool_policy("mcp::context7::search")
1364                .await?,
1365            ToolPermissionDecision::Allow
1366        );
1367
1368        Ok(())
1369    }
1370
1371    #[tokio::test]
1372    async fn apply_patch_alias_executes_without_recursive_reentry() -> Result<()> {
1373        let temp_dir = TempDir::new()?;
1374        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
1375        registry.allow_all_tools().await?;
1376
1377        let patch =
1378            "*** Begin Patch\n*** Add File: patched_via_alias.txt\n+patched\n*** End Patch\n";
1379        let response = registry
1380            .execute_tool(tools::APPLY_PATCH, json!({ "patch": patch }))
1381            .await?;
1382
1383        assert_eq!(response.get("success").and_then(Value::as_bool), Some(true));
1384
1385        let file_contents = fs::read_to_string(temp_dir.path().join("patched_via_alias.txt"))?;
1386        assert_eq!(file_contents, "patched\n");
1387
1388        Ok(())
1389    }
1390
1391    #[tokio::test]
1392    async fn apply_patch_accepts_input_payload() -> Result<()> {
1393        let temp_dir = TempDir::new()?;
1394        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
1395        registry.allow_all_tools().await?;
1396
1397        let patch =
1398            "*** Begin Patch\n*** Add File: patched_via_input.txt\n+patched\n*** End Patch\n";
1399        let response = registry
1400            .execute_tool(tools::APPLY_PATCH, json!({ "input": patch }))
1401            .await?;
1402
1403        assert_eq!(response.get("success").and_then(Value::as_bool), Some(true));
1404
1405        let file_contents = fs::read_to_string(temp_dir.path().join("patched_via_input.txt"))?;
1406        assert_eq!(file_contents, "patched\n");
1407
1408        Ok(())
1409    }
1410
1411    #[tokio::test]
1412    async fn public_apply_patch_accepts_raw_string_payload() -> Result<()> {
1413        let temp_dir = TempDir::new()?;
1414        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
1415        registry.allow_all_tools().await?;
1416
1417        let patch =
1418            "*** Begin Patch\n*** Add File: patched_via_raw_string.txt\n+patched\n*** End Patch\n";
1419        let response = registry
1420            .execute_public_tool_ref(tools::APPLY_PATCH, &json!(patch))
1421            .await?;
1422
1423        assert_eq!(response.get("success").and_then(Value::as_bool), Some(true));
1424
1425        let file_contents = fs::read_to_string(temp_dir.path().join("patched_via_raw_string.txt"))?;
1426        assert_eq!(file_contents, "patched\n");
1427
1428        Ok(())
1429    }
1430
1431    #[tokio::test]
1432    async fn execution_history_records_harness_context() -> Result<()> {
1433        let temp_dir = TempDir::new()?;
1434        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
1435
1436        registry.set_harness_session("session-history");
1437        registry.set_harness_task(Some("task-history".to_owned()));
1438
1439        registry
1440            .register_tool(ToolRegistration::from_tool_instance(
1441                CUSTOM_TOOL_NAME,
1442                CapabilityLevel::CodeSearch,
1443                CustomEchoTool,
1444            ))
1445            .await?;
1446        registry.allow_all_tools().await?;
1447
1448        let args = json!({"input": "value"});
1449        let response = registry
1450            .execute_tool(CUSTOM_TOOL_NAME, args.clone())
1451            .await?;
1452        assert!(response["success"].as_bool().unwrap_or(false));
1453
1454        let records = registry.get_recent_tool_records(1);
1455        let record = records.first().expect("execution record captured");
1456        assert_eq!(record.tool_name, CUSTOM_TOOL_NAME);
1457        assert_eq!(record.context.session_id, "session-history");
1458        assert_eq!(record.context.task_id.as_deref(), Some("task-history"));
1459        assert_eq!(record.args, args);
1460        assert!(record.success);
1461
1462        Ok(())
1463    }
1464
1465    #[tokio::test]
1466    async fn reentrancy_guard_blocks_recursive_tool_loops() -> Result<()> {
1467        let temp_dir = TempDir::new()?;
1468        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
1469
1470        registry
1471            .register_tool(ToolRegistration::new(
1472                REENTRANT_TOOL_NAME,
1473                CapabilityLevel::CodeSearch,
1474                false,
1475                reentrant_tool_executor,
1476            ))
1477            .await?;
1478        registry.allow_all_tools().await?;
1479
1480        let response = registry
1481            .execute_tool(REENTRANT_TOOL_NAME, json!({"input": "loop"}))
1482            .await?;
1483
1484        assert_eq!(
1485            response
1486                .get("reentrant_call_blocked")
1487                .and_then(Value::as_bool),
1488            Some(true)
1489        );
1490        assert_eq!(
1491            response
1492                .pointer("/error/error_type")
1493                .and_then(Value::as_str),
1494            Some("PolicyViolation")
1495        );
1496        assert!(
1497            response
1498                .pointer("/error/message")
1499                .and_then(Value::as_str)
1500                .unwrap_or_default()
1501                .contains("REENTRANCY GUARD")
1502        );
1503
1504        Ok(())
1505    }
1506
1507    #[tokio::test]
1508    async fn reentrancy_guard_blocks_cross_tool_cycles() -> Result<()> {
1509        let temp_dir = TempDir::new()?;
1510        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
1511
1512        registry
1513            .register_tool(ToolRegistration::new(
1514                MUTUAL_REENTRANT_TOOL_A,
1515                CapabilityLevel::CodeSearch,
1516                false,
1517                mutual_reentrant_tool_a_executor,
1518            ))
1519            .await?;
1520        registry
1521            .register_tool(ToolRegistration::new(
1522                MUTUAL_REENTRANT_TOOL_B,
1523                CapabilityLevel::CodeSearch,
1524                false,
1525                mutual_reentrant_tool_b_executor,
1526            ))
1527            .await?;
1528        registry.allow_all_tools().await?;
1529
1530        let response = registry
1531            .execute_tool(MUTUAL_REENTRANT_TOOL_A, json!({"input": "cycle"}))
1532            .await?;
1533
1534        assert_eq!(
1535            response
1536                .get("reentrant_call_blocked")
1537                .and_then(Value::as_bool),
1538            Some(true)
1539        );
1540        assert_eq!(
1541            response
1542                .pointer("/error/error_type")
1543                .and_then(Value::as_str),
1544            Some("PolicyViolation")
1545        );
1546
1547        let stack_trace = response
1548            .get("stack_trace")
1549            .and_then(Value::as_str)
1550            .unwrap_or_default();
1551        assert!(stack_trace.contains(MUTUAL_REENTRANT_TOOL_A));
1552        assert!(stack_trace.contains(MUTUAL_REENTRANT_TOOL_B));
1553
1554        Ok(())
1555    }
1556
1557    #[tokio::test]
1558    async fn full_auto_allowlist_enforced() -> Result<()> {
1559        let temp_dir = TempDir::new()?;
1560        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
1561
1562        registry
1563            .enable_full_auto_mode(&[tools::READ_FILE.to_string()])
1564            .await;
1565
1566        assert!(registry.preflight_tool_permission(tools::READ_FILE).await?);
1567        assert!(
1568            !registry
1569                .preflight_tool_permission(tools::RUN_PTY_CMD)
1570                .await?
1571        );
1572
1573        Ok(())
1574    }
1575
1576    #[test]
1577    fn normalizes_mcp_tool_identifiers() {
1578        assert_eq!(
1579            normalize_mcp_tool_identifier("sequential-thinking"),
1580            "sequentialthinking"
1581        );
1582        assert_eq!(
1583            normalize_mcp_tool_identifier("Context7.Lookup"),
1584            "context7lookup"
1585        );
1586        assert_eq!(normalize_mcp_tool_identifier("alpha_beta"), "alphabeta");
1587    }
1588
1589    #[test]
1590    fn timeout_policy_derives_from_config() {
1591        let config = TimeoutsConfig {
1592            default_ceiling_seconds: 0,
1593            pty_ceiling_seconds: 600,
1594            mcp_ceiling_seconds: 90,
1595            warning_threshold_percent: 75,
1596            ..Default::default()
1597        };
1598
1599        let policy = ToolTimeoutPolicy::from_config(&config);
1600        assert_eq!(policy.ceiling_for(ToolTimeoutCategory::Default), None);
1601        assert_eq!(
1602            policy.ceiling_for(ToolTimeoutCategory::Pty),
1603            Some(Duration::from_secs(600))
1604        );
1605        assert_eq!(
1606            policy.ceiling_for(ToolTimeoutCategory::Mcp),
1607            Some(Duration::from_secs(90))
1608        );
1609        assert!((policy.warning_fraction() - 0.75).abs() < f32::EPSILON);
1610    }
1611
1612    #[tokio::test]
1613    async fn timeout_errors_are_structured_and_track_failures() -> Result<()> {
1614        let temp_dir = TempDir::new()?;
1615        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
1616
1617        registry
1618            .register_tool(ToolRegistration::from_tool_instance(
1619                SLOW_TIMEOUT_TOOL_NAME,
1620                CapabilityLevel::CodeSearch,
1621                SlowTimeoutTool,
1622            ))
1623            .await?;
1624        registry.allow_all_tools().await?;
1625
1626        registry.apply_timeout_policy(&TimeoutsConfig {
1627            default_ceiling_seconds: 1,
1628            pty_ceiling_seconds: 1,
1629            mcp_ceiling_seconds: 1,
1630            ..Default::default()
1631        });
1632
1633        let mut policy = ExecutionPolicySnapshot::default().with_max_retries(4);
1634        policy.retry_base_delay = Duration::from_millis(1);
1635        policy.retry_max_delay = Duration::from_millis(1);
1636        policy.retry_multiplier = 1.0;
1637
1638        let request =
1639            ToolExecutionRequest::new(SLOW_TIMEOUT_TOOL_NAME, json!({})).with_policy(policy);
1640        let outcome = registry.execute_public_tool_request(request).await;
1641
1642        assert!(!outcome.is_success());
1643        assert_eq!(outcome.attempts, 5);
1644
1645        let error = outcome.error.expect("timeout outcome should include error");
1646        assert_eq!(error.tool_name, SLOW_TIMEOUT_TOOL_NAME);
1647        assert!(matches!(error.error_type, ToolErrorType::Timeout));
1648        assert_eq!(error.category, vtcode_commons::ErrorCategory::Timeout);
1649        assert!(error.is_recoverable);
1650        assert!(error.retry_after_ms.is_some());
1651        assert!(
1652            error
1653                .message
1654                .contains("exceeded the standard timeout ceiling")
1655        );
1656        assert_eq!(
1657            error
1658                .debug_context
1659                .as_ref()
1660                .and_then(|ctx| ctx.surface.as_deref()),
1661            Some("tool_registry")
1662        );
1663
1664        let failures = registry.execution_history.get_recent_failures(1);
1665        assert_eq!(failures.len(), 1);
1666        assert_eq!(failures[0].timeout_category.as_deref(), Some("standard"));
1667        assert_eq!(failures[0].effective_timeout_ms, Some(1_000));
1668
1669        let consecutive_failures = registry
1670            .resiliency
1671            .lock()
1672            .failure_trackers
1673            .get(&ToolTimeoutCategory::Default)
1674            .map(|tracker| tracker.consecutive_failures)
1675            .unwrap_or(0);
1676        assert_eq!(consecutive_failures, 5);
1677
1678        Ok(())
1679    }
1680}