From 242899d4b6dffb6836a4a4de43a5ddc9f4dd3c51 Mon Sep 17 00:00:00 2001 From: Paulo Aboim Pinto Date: Thu, 28 May 2026 00:47:37 +0200 Subject: [PATCH] feat: run ToolCallBefore hooks before tool execution --- crates/tui/src/core/engine.rs | 9 ++ crates/tui/src/core/engine/turn_loop.rs | 140 ++++++++++++++++++++++++ crates/tui/src/core/ops.rs | 3 + crates/tui/src/main.rs | 2 + crates/tui/src/runtime_threads.rs | 2 + crates/tui/src/tui/ui.rs | 2 + 6 files changed, 158 insertions(+) diff --git a/crates/tui/src/core/engine.rs b/crates/tui/src/core/engine.rs index 5813b538..2f920b3f 100644 --- a/crates/tui/src/core/engine.rs +++ b/crates/tui/src/core/engine.rs @@ -166,6 +166,9 @@ pub struct EngineConfig { /// Tool restriction from custom slash command frontmatter. /// `None` means the current turn may use the normal tool set. pub allowed_tools: Option>, + /// Hook executor for control-plane hooks. + /// `ToolCallBefore` hooks may deny a tool call with exit code 2. + pub hook_executor: Option>, /// Resolved BCP-47 locale tag (e.g. `"en"`, `"zh-Hans"`, `"ja"`) /// for the `## Environment` block in the system prompt. The /// caller resolves this from `Settings` once at engine @@ -237,6 +240,7 @@ impl Default for EngineConfig { strict_tool_mode: false, goal_objective: None, allowed_tools: None, + hook_executor: None, locale_tag: "en".to_string(), workshop: None, search_provider: crate::config::SearchProvider::default(), @@ -650,6 +654,7 @@ impl Engine { translation_enabled, show_thinking, allowed_tools, + hook_executor, } => { self.handle_send_message( content, @@ -666,6 +671,7 @@ impl Engine { translation_enabled, show_thinking, allowed_tools, + hook_executor, ) .await; } @@ -884,6 +890,7 @@ impl Engine { self.config.translation_enabled, self.config.show_thinking, self.config.allowed_tools.clone(), + self.config.hook_executor.clone(), ) .await; } @@ -1008,6 +1015,7 @@ impl Engine { translation_enabled: bool, show_thinking: bool, allowed_tools: Option>, + hook_executor: Option>, ) { // Reset cancel token for fresh turn (in case previous was cancelled) self.reset_cancel_token(); @@ -1114,6 +1122,7 @@ impl Engine { ); } self.config.allowed_tools = allowed_tools; + self.config.hook_executor = hook_executor; self.session.reasoning_effort = reasoning_effort; self.session.reasoning_effort_auto = reasoning_effort_auto; self.session.auto_model = auto_model; diff --git a/crates/tui/src/core/engine/turn_loop.rs b/crates/tui/src/core/engine/turn_loop.rs index 9a3245e1..9d0bdd91 100644 --- a/crates/tui/src/core/engine/turn_loop.rs +++ b/crates/tui/src/core/engine/turn_loop.rs @@ -1261,6 +1261,45 @@ impl Engine { ))); } + if blocked_error.is_none() + && let Some(hook_executor) = self.config.hook_executor.as_ref() + && hook_executor.has_hooks_for_event(crate::hooks::HookEvent::ToolCallBefore) + { + let hook_context = crate::hooks::HookContext::new() + .with_tool_name(&tool_name) + .with_tool_args(&tool_input) + .with_mode(&format!("{mode:?}")) + .with_workspace(self.session.workspace.clone()) + .with_model(&self.config.model) + .with_session_id(&self.session.id); + let hook_results = hook_executor + .execute(crate::hooks::HookEvent::ToolCallBefore, &hook_context); + if let Some(denial) = hook_results + .iter() + .find(|result| result.exit_code == Some(2)) + { + let reason = denial + .stdout + .trim() + .lines() + .next() + .filter(|line| !line.is_empty()) + .or_else(|| { + denial + .stderr + .trim() + .lines() + .next() + .filter(|line| !line.is_empty()) + }) + .or(denial.error.as_deref()) + .unwrap_or("ToolCallBefore hook denied tool execution"); + blocked_error = Some(ToolError::permission_denied(format!( + "ToolCallBefore hook denied tool '{tool_name}': {reason}" + ))); + } + } + if !caller_allowed_for_tool(tool_caller.as_ref(), tool_def) { blocked_error = Some(ToolError::permission_denied(format!( "Tool '{tool_name}' does not allow caller '{}'", @@ -2514,4 +2553,105 @@ mod tests { let allowed = vec!["read_file".to_string()]; assert!(command_allows_tool(Some(&allowed), &tool_name)); } + + #[test] + fn hook_gate_denies_with_exit_code_2() { + use crate::hooks::{Hook, HookContext, HookEvent, HookExecutor, HooksConfig}; + + let deny_cmd = if cfg!(windows) { "exit /b 2" } else { "exit 2" }; + let config = HooksConfig { + enabled: true, + hooks: vec![Hook::new(HookEvent::ToolCallBefore, deny_cmd)], + ..HooksConfig::default() + }; + let executor = HookExecutor::new(config, std::path::PathBuf::from(".")); + let ctx = HookContext::new() + .with_tool_name("exec_shell") + .with_tool_args(&serde_json::json!({})); + let results = executor.execute(HookEvent::ToolCallBefore, &ctx); + + assert_eq!(results.len(), 1); + assert_eq!(results[0].exit_code, Some(2)); + } + + #[test] + fn hook_gate_allows_with_exit_code_0() { + use crate::hooks::{Hook, HookContext, HookEvent, HookExecutor, HooksConfig}; + + let allow_cmd = if cfg!(windows) { "exit /b 0" } else { "exit 0" }; + let config = HooksConfig { + enabled: true, + hooks: vec![Hook::new(HookEvent::ToolCallBefore, allow_cmd)], + ..HooksConfig::default() + }; + let executor = HookExecutor::new(config, std::path::PathBuf::from(".")); + let ctx = HookContext::new() + .with_tool_name("read_file") + .with_tool_args(&serde_json::json!({})); + let results = executor.execute(HookEvent::ToolCallBefore, &ctx); + + assert_eq!(results.len(), 1); + assert_eq!(results[0].exit_code, Some(0)); + assert!(results[0].success); + } + + #[test] + fn hook_gate_failure_exit_code_1_is_not_denial() { + use crate::hooks::{Hook, HookContext, HookEvent, HookExecutor, HooksConfig}; + + let fail_cmd = if cfg!(windows) { "exit /b 1" } else { "exit 1" }; + let config = HooksConfig { + enabled: true, + hooks: vec![Hook::new(HookEvent::ToolCallBefore, fail_cmd)], + ..HooksConfig::default() + }; + let executor = HookExecutor::new(config, std::path::PathBuf::from(".")); + let ctx = HookContext::new() + .with_tool_name("write_file") + .with_tool_args(&serde_json::json!({})); + let results = executor.execute(HookEvent::ToolCallBefore, &ctx); + + assert_eq!(results.len(), 1); + assert_eq!(results[0].exit_code, Some(1)); + assert_ne!(results[0].exit_code, Some(2)); + } + + #[test] + fn hook_gate_no_hooks_returns_no_results() { + use crate::hooks::{HookContext, HookEvent, HookExecutor, HooksConfig}; + + let config = HooksConfig { + enabled: true, + hooks: vec![], + ..HooksConfig::default() + }; + let executor = HookExecutor::new(config, std::path::PathBuf::from(".")); + let ctx = HookContext::new().with_tool_name("grep_files"); + let results = executor.execute(HookEvent::ToolCallBefore, &ctx); + + assert!(results.is_empty()); + } + + #[test] + fn hook_gate_denial_reason_can_come_from_stdout() { + use crate::hooks::{Hook, HookContext, HookEvent, HookExecutor, HooksConfig}; + + let deny_cmd = if cfg!(windows) { + "echo Tool blocked by security policy & exit /b 2" + } else { + "echo 'Tool blocked by security policy' && exit 2" + }; + let config = HooksConfig { + enabled: true, + hooks: vec![Hook::new(HookEvent::ToolCallBefore, deny_cmd)], + ..HooksConfig::default() + }; + let executor = HookExecutor::new(config, std::path::PathBuf::from(".")); + let ctx = HookContext::new().with_tool_name("exec_shell"); + let results = executor.execute(HookEvent::ToolCallBefore, &ctx); + + assert_eq!(results.len(), 1); + assert_eq!(results[0].exit_code, Some(2)); + assert!(results[0].stdout.contains("security")); + } } diff --git a/crates/tui/src/core/ops.rs b/crates/tui/src/core/ops.rs index 87f47945..df6f0aa0 100644 --- a/crates/tui/src/core/ops.rs +++ b/crates/tui/src/core/ops.rs @@ -35,6 +35,9 @@ pub enum Op { /// Tool restriction from custom slash command frontmatter. /// `None` means the current turn may use the normal tool set. allowed_tools: Option>, + /// Hook executor for control-plane hooks. + /// `ToolCallBefore` hooks may deny a tool call with exit code 2. + hook_executor: Option>, }, /// Cancel the current request diff --git a/crates/tui/src/main.rs b/crates/tui/src/main.rs index 9feaaac4..47468b73 100644 --- a/crates/tui/src/main.rs +++ b/crates/tui/src/main.rs @@ -5379,6 +5379,7 @@ async fn run_exec_agent( strict_tool_mode: config.strict_tool_mode.unwrap_or(false), goal_objective: None, allowed_tools: None, + hook_executor: None, locale_tag: crate::localization::resolve_locale(&settings.locale) .tag() .to_string(), @@ -5435,6 +5436,7 @@ async fn run_exec_agent( model: effective_model.clone(), goal_objective: None, allowed_tools: None, + hook_executor: None, reasoning_effort: effective_reasoning_effort, reasoning_effort_auto: auto_model, auto_model, diff --git a/crates/tui/src/runtime_threads.rs b/crates/tui/src/runtime_threads.rs index 51f79922..39c009ed 100644 --- a/crates/tui/src/runtime_threads.rs +++ b/crates/tui/src/runtime_threads.rs @@ -1654,6 +1654,7 @@ impl RuntimeThreadManager { translation_enabled: false, show_thinking, allowed_tools: None, + hook_executor: None, approval_mode: if auto_approve { crate::tui::approval::ApprovalMode::Auto } else { @@ -2020,6 +2021,7 @@ impl RuntimeThreadManager { strict_tool_mode: self.config.strict_tool_mode.unwrap_or(false), goal_objective: None, allowed_tools: None, + hook_executor: None, locale_tag: crate::localization::resolve_locale(&settings.locale) .tag() .to_string(), diff --git a/crates/tui/src/tui/ui.rs b/crates/tui/src/tui/ui.rs index e92a2a05..d0dc6406 100644 --- a/crates/tui/src/tui/ui.rs +++ b/crates/tui/src/tui/ui.rs @@ -754,6 +754,7 @@ fn build_engine_config(app: &App, config: &Config) -> EngineConfig { ), max_spawn_depth: crate::tools::subagent::DEFAULT_MAX_SPAWN_DEPTH, allowed_tools: app.active_allowed_tools.clone(), + hook_executor: Some(std::sync::Arc::new(app.hooks.clone())), network_policy: config.network.clone().map(|toml_cfg| { crate::network_policy::NetworkPolicyDecider::with_default_audit(toml_cfg.into_runtime()) }), @@ -4706,6 +4707,7 @@ async fn dispatch_user_message( translation_enabled: app.translation_enabled, show_thinking: app.show_thinking, allowed_tools: app.active_allowed_tools.clone(), + hook_executor: Some(std::sync::Arc::new(app.hooks.clone())), }) .await {