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 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 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
113pub type ToolProgressCallback = Arc<dyn Fn(&str, &str) + Send + Sync>;
115
116use super::traits::Tool;
117#[cfg(test)]
118use crate::config::types::CapabilityLevel;
119
120const 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 resiliency: Arc<Mutex<ResiliencyContext>>,
139
140 mcp_circuit_breaker: Arc<circuit_breaker::McpCircuitBreaker>,
142 shared_circuit_breaker: Arc<RwLock<Option<Arc<crate::tools::circuit_breaker::CircuitBreaker>>>>,
144 initialized: Arc<std::sync::atomic::AtomicBool>,
145 shell_policy: Arc<RwLock<ShellPolicyChecker>>,
147 runtime_sandbox_config: Arc<RwLock<vtcode_config::SandboxConfig>>,
148 agent_type: Arc<RwLock<String>>,
149 active_pty_sessions: Arc<RwLock<Option<Arc<std::sync::atomic::AtomicUsize>>>>,
151
152 cached_available_tools: Arc<RwLock<Option<Vec<String>>>>,
154 progress_callback: Arc<RwLock<Option<ToolProgressCallback>>>,
156 pub(crate) tool_call_counter: Arc<std::sync::atomic::AtomicU64>,
159 pub(crate) pty_poll_counter: Arc<std::sync::atomic::AtomicU64>,
161 metrics: Arc<crate::metrics::MetricsCollector>,
163
164 memory_pool: Arc<crate::core::memory_pool::MemoryPool>,
167 hot_tool_cache: Arc<parking_lot::RwLock<lru::LruCache<String, Arc<dyn Tool>>>>,
169 optimization_config: vtcode_config::OptimizationConfig,
171
172 output_spooler: Arc<super::output_spooler::ToolOutputSpooler>,
174
175 plan_mode_state: PlanModeState,
179 safety_gateway: Arc<SafetyGateway>,
181 cgp_runtime_mode: Arc<RwLock<Option<CgpRuntimeMode>>>,
183 tool_assembly: Arc<RwLock<ToolAssembly>>,
185 tool_catalog_state: Arc<SessionToolCatalogState>,
187 subagent_controller: Arc<RwLock<Option<Arc<SubagentController>>>>,
189 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}