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