From 4f4f1f98a30c27be16dd191c3b898337fe4786f4 Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Thu, 5 Mar 2026 02:49:06 -0500 Subject: [PATCH] Simplify hook implementation: deduplicate LLM setup, fix async I/O, cache HTTP clients MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Extract shared prompt/agent hook setup (model resolution, system prompt, user message, timeout wrapper) into reusable helpers - Fix blocking I/O: std::process::Command → tokio::process::Command for host-mode hook execution - Cache reqwest::Client per TLS mode via OnceLock instead of rebuilding per HTTP hook call - Return Cow from resolved_hook_type() to avoid cloning HookType on every call - Rename command_executor → executor in HookRunner Co-Authored-By: Claude Opus 4.6 --- crates/arc-workflows/src/hook/config.rs | 46 ++--- crates/arc-workflows/src/hook/executor.rs | 213 +++++++++++++++------- crates/arc-workflows/src/hook/runner.rs | 14 +- crates/arc-workflows/tests/integration.rs | 10 +- 4 files changed, 188 insertions(+), 95 deletions(-) diff --git a/crates/arc-workflows/src/hook/config.rs b/crates/arc-workflows/src/hook/config.rs index ce0887422..7d879e1ce 100644 --- a/crates/arc-workflows/src/hook/config.rs +++ b/crates/arc-workflows/src/hook/config.rs @@ -1,3 +1,5 @@ +use std::borrow::Cow; + use serde::Deserialize; use super::types::HookEvent; @@ -64,13 +66,13 @@ pub struct HookDefinition { impl HookDefinition { /// Resolve the effective hook type: explicit `hook_type` wins, then `command` /// shorthand, then error. - pub fn resolved_hook_type(&self) -> Option { + pub fn resolved_hook_type(&self) -> Option> { if let Some(ref ht) = self.hook_type { - return Some(ht.clone()); + return Some(Cow::Borrowed(ht)); } self.command .as_ref() - .map(|cmd| HookType::Command { command: cmd.clone() }) + .map(|cmd| Cow::Owned(HookType::Command { command: cmd.clone() })) } /// Whether this hook is blocking for its event. @@ -87,7 +89,7 @@ impl HookDefinition { if let Some(ms) = self.timeout_ms { return std::time::Duration::from_millis(ms); } - let default_ms = match self.resolved_hook_type() { + let default_ms = match self.resolved_hook_type().as_deref() { Some(HookType::Prompt { .. }) => 30_000, _ => 60_000, }; @@ -107,7 +109,7 @@ impl HookDefinition { return n.clone(); } let event_str = self.event.to_string(); - match self.resolved_hook_type() { + match self.resolved_hook_type().as_deref() { Some(HookType::Command { ref command }) => { let short = &command[..arc_agent::floor_char_boundary(command, 20)]; format!("{event_str}:{short}") @@ -178,7 +180,7 @@ command = "./scripts/pre-check.sh" assert_eq!(hook.event, HookEvent::StageStart); assert_eq!(hook.command.as_deref(), Some("./scripts/pre-check.sh")); let resolved = hook.resolved_hook_type().unwrap(); - assert!(matches!(resolved, HookType::Command { command } if command == "./scripts/pre-check.sh")); + assert!(matches!(&*resolved, HookType::Command { command } if command == "./scripts/pre-check.sh")); } #[test] @@ -207,7 +209,7 @@ url = "https://hooks.example.com/done" let config: HookConfig = toml::from_str(toml).unwrap(); let hook = &config.hooks[0]; assert!(matches!( - hook.resolved_hook_type(), + hook.resolved_hook_type().as_deref(), Some(HookType::Http { url, .. }) if url == "https://hooks.example.com/done" )); } @@ -226,7 +228,7 @@ Authorization = "Bearer $API_KEY" "#; let config: HookConfig = toml::from_str(toml).unwrap(); let hook = &config.hooks[0]; - match hook.resolved_hook_type().unwrap() { + match &*hook.resolved_hook_type().unwrap() { HookType::Http { url, headers, @@ -234,9 +236,9 @@ Authorization = "Bearer $API_KEY" .. } => { assert_eq!(url, "https://hooks.example.com/start"); - assert_eq!(allowed_env_vars, vec!["API_KEY", "SECRET"]); + assert_eq!(allowed_env_vars, &["API_KEY", "SECRET"]); assert_eq!( - headers.unwrap().get("Authorization").unwrap(), + headers.as_ref().unwrap().get("Authorization").unwrap(), "Bearer $API_KEY" ); } @@ -254,7 +256,7 @@ url = "https://hooks.example.com/done" "#; let config: HookConfig = toml::from_str(toml).unwrap(); let hook = &config.hooks[0]; - match hook.resolved_hook_type().unwrap() { + match &*hook.resolved_hook_type().unwrap() { HookType::Http { allowed_env_vars, .. } => { @@ -455,8 +457,8 @@ url = "https://hooks.example.com/done" "#; let config: HookConfig = toml::from_str(toml).unwrap(); let hook = &config.hooks[0]; - match hook.resolved_hook_type().unwrap() { - HookType::Http { tls, .. } => assert_eq!(tls, TlsMode::Verify), + match &*hook.resolved_hook_type().unwrap() { + HookType::Http { tls, .. } => assert_eq!(*tls, TlsMode::Verify), _ => panic!("expected Http hook type"), } } @@ -472,8 +474,8 @@ tls = "no_verify" "#; let config: HookConfig = toml::from_str(toml).unwrap(); let hook = &config.hooks[0]; - match hook.resolved_hook_type().unwrap() { - HookType::Http { tls, .. } => assert_eq!(tls, TlsMode::NoVerify), + match &*hook.resolved_hook_type().unwrap() { + HookType::Http { tls, .. } => assert_eq!(*tls, TlsMode::NoVerify), _ => panic!("expected Http hook type"), } } @@ -489,8 +491,8 @@ tls = "off" "#; let config: HookConfig = toml::from_str(toml).unwrap(); let hook = &config.hooks[0]; - match hook.resolved_hook_type().unwrap() { - HookType::Http { tls, .. } => assert_eq!(tls, TlsMode::Off), + match &*hook.resolved_hook_type().unwrap() { + HookType::Http { tls, .. } => assert_eq!(*tls, TlsMode::Off), _ => panic!("expected Http hook type"), } } @@ -507,9 +509,9 @@ model = "haiku" let config: HookConfig = toml::from_str(toml).unwrap(); let hook = &config.hooks[0]; assert!(matches!( - hook.resolved_hook_type(), + hook.resolved_hook_type().as_deref(), Some(HookType::Prompt { prompt, model }) - if prompt == "Should this stage proceed?" && model == Some("haiku".into()) + if prompt == "Should this stage proceed?" && *model == Some("haiku".into()) )); } @@ -526,11 +528,11 @@ max_tool_rounds = 10 let config: HookConfig = toml::from_str(toml).unwrap(); let hook = &config.hooks[0]; assert!(matches!( - hook.resolved_hook_type(), + hook.resolved_hook_type().as_deref(), Some(HookType::Agent { prompt, model, max_tool_rounds }) if prompt == "Verify tests pass." - && model == Some("sonnet".into()) - && max_tool_rounds == Some(10) + && *model == Some("sonnet".into()) + && *max_tool_rounds == Some(10) )); } diff --git a/crates/arc-workflows/src/hook/executor.rs b/crates/arc-workflows/src/hook/executor.rs index e84648744..78e392687 100644 --- a/crates/arc-workflows/src/hook/executor.rs +++ b/crates/arc-workflows/src/hook/executor.rs @@ -1,3 +1,4 @@ +use std::borrow::Cow; use std::collections::HashMap; use std::path::Path; use std::sync::Arc; @@ -10,6 +11,8 @@ use arc_agent::Sandbox; use super::config::{HookDefinition, HookType, TlsMode}; use super::types::{HookContext, HookDecision, HookResult, PromptHookResponse}; +const HOOK_EVALUATOR_SYSTEM_PROMPT: &str = "You are a hook evaluator for a workflow engine. Given context about a workflow event, evaluate the condition and respond with JSON: {\"ok\": true} or {\"ok\": false, \"reason\": \"...\"}. Respond ONLY with valid JSON."; + /// Trait for executing hooks via different transports. #[async_trait] pub trait HookExecutor: Send + Sync { @@ -130,7 +133,7 @@ impl HookExecutorImpl { }, } } else { - let mut cmd = std::process::Command::new("sh"); + let mut cmd = tokio::process::Command::new("sh"); cmd.arg("-c").arg(command); if let Some(wd) = work_dir { cmd.current_dir(wd); @@ -145,10 +148,10 @@ impl HookExecutorImpl { match cmd.spawn() { Ok(mut child) => { if let Some(mut stdin) = child.stdin.take() { - use std::io::Write; - let _ = stdin.write_all(context_json.as_bytes()); + use tokio::io::AsyncWriteExt; + let _ = stdin.write_all(context_json.as_bytes()).await; } - match child.wait_with_output() { + match child.wait_with_output().await { Ok(output) => { let exit_code = output.status.code().unwrap_or(1); let stdout = String::from_utf8_lossy(&output.stdout); @@ -198,6 +201,41 @@ impl HookExecutorImpl { } } + /// Resolve a model alias (e.g. "haiku") to a concrete model ID. + fn resolve_model(model: &Option) -> String { + let model_id = model.as_deref().unwrap_or("haiku"); + let model_info = arc_llm::catalog::get_model_info(model_id); + model_info + .as_ref() + .map_or(model_id, |m| m.id.as_str()) + .to_string() + } + + /// Build the user message for prompt/agent hooks. + fn build_hook_user_message(prompt: &str, context: &HookContext) -> String { + let context_json = serde_json::to_string(context).unwrap_or_default(); + format!("Hook prompt: {prompt}\n\nEvent context:\n{context_json}") + } + + /// Execute an LLM hook with a timeout, failing open on error or timeout. + async fn execute_llm_with_timeout( + timeout: std::time::Duration, + hook_kind: &str, + f: F, + ) -> HookDecision + where + F: FnOnce() -> Fut, + Fut: std::future::Future, + { + match tokio::time::timeout(timeout, f()).await { + Ok(decision) => decision, + Err(_) => { + tracing::warn!("{hook_kind} hook timed out, proceeding"); + HookDecision::Proceed + } + } + } + /// Execute a prompt hook: single-turn LLM call returning ok/block. async fn execute_prompt( prompt: &str, @@ -205,17 +243,12 @@ impl HookExecutorImpl { context: &HookContext, timeout: std::time::Duration, ) -> HookDecision { - let result = tokio::time::timeout(timeout, async { - let model_id = model.as_deref().unwrap_or("haiku"); - let model_info = arc_llm::catalog::get_model_info(model_id); - let resolved_model = model_info.as_ref().map_or(model_id, |m| m.id.as_str()); + let resolved_model = Self::resolve_model(model); + let user_msg = Self::build_hook_user_message(prompt, context); - let context_json = serde_json::to_string(context).unwrap_or_default(); - let system = "You are a hook evaluator for a workflow engine. Given context about a workflow event, evaluate the condition and respond with JSON: {\"ok\": true} or {\"ok\": false, \"reason\": \"...\"}. Respond ONLY with valid JSON."; - let user_msg = format!("Hook prompt: {prompt}\n\nEvent context:\n{context_json}"); - - let params = arc_llm::generate::GenerateParams::new(resolved_model) - .system(system) + Self::execute_llm_with_timeout(timeout, "prompt", || async move { + let params = arc_llm::generate::GenerateParams::new(&resolved_model) + .system(HOOK_EVALUATOR_SYSTEM_PROMPT) .prompt(user_msg) .max_tokens(1024); @@ -227,15 +260,7 @@ impl HookExecutorImpl { } } }) - .await; - - match result { - Ok(decision) => decision, - Err(_) => { - tracing::warn!("prompt hook timed out, proceeding"); - HookDecision::Proceed - } - } + .await } /// Execute an agent hook: multi-turn LLM call with sandbox tool access. @@ -251,11 +276,10 @@ impl HookExecutorImpl { sandbox: Arc, timeout: std::time::Duration, ) -> HookDecision { - let result = tokio::time::timeout(timeout, async { - let model_id = model.as_deref().unwrap_or("haiku"); - let model_info = arc_llm::catalog::get_model_info(model_id); - let resolved_model = model_info.as_ref().map_or(model_id, |m| m.id.as_str()); + let resolved_model = Self::resolve_model(model); + let user_msg = Self::build_hook_user_message(prompt, context); + Self::execute_llm_with_timeout(timeout, "agent", || async move { let client = match arc_llm::client::Client::from_env().await { Ok(c) => c, Err(e) => { @@ -264,18 +288,13 @@ impl HookExecutorImpl { } }; - let context_json = serde_json::to_string(context).unwrap_or_default(); - let system = "You are a hook evaluator for a workflow engine. Given context about a workflow event, evaluate the condition and respond with JSON: {\"ok\": true} or {\"ok\": false, \"reason\": \"...\"}. Respond ONLY with valid JSON."; - let user_msg = format!("Hook prompt: {prompt}\n\nEvent context:\n{context_json}"); - - // Reuse the same tool registry as normal agent sessions. let config = arc_agent::SessionConfig::default(); let mut registry = arc_agent::ToolRegistry::new(); arc_agent::register_core_tools(&mut registry, &config, None); let tool_defs = registry.definitions(); let mut messages = vec![ - arc_llm::types::Message::system(system), + arc_llm::types::Message::system(HOOK_EVALUATOR_SYSTEM_PROMPT), arc_llm::types::Message::user(user_msg), ]; @@ -284,7 +303,7 @@ impl HookExecutorImpl { for _ in 0..rounds { let request = arc_llm::types::Request { - model: resolved_model.to_string(), + model: resolved_model.clone(), messages: messages.clone(), provider: None, tools: Some(tool_defs.clone()), @@ -312,10 +331,8 @@ impl HookExecutorImpl { return Self::parse_prompt_response(&response.text()); } - // Append assistant message with tool calls messages.push(response.message.clone()); - // Execute each tool call via the shared registry for tc in &tool_calls { let tool = registry.get(&tc.name).cloned(); let ctx = arc_agent::tool_registry::ToolContext { @@ -346,20 +363,22 @@ impl HookExecutorImpl { tracing::warn!("agent hook exhausted max tool rounds, proceeding"); HookDecision::Proceed }) - .await; + .await + } - match result { - Ok(decision) => decision, - Err(_) => { - tracing::warn!("agent hook timed out, proceeding"); - HookDecision::Proceed - } - } + /// Build a reqwest client for the given TLS mode. + fn build_http_client(tls: &TlsMode) -> reqwest::Client { + let accept_invalid = matches!(tls, TlsMode::NoVerify | TlsMode::Off); + reqwest::Client::builder() + .danger_accept_invalid_certs(accept_invalid) + .build() + .unwrap_or_default() } /// Execute an HTTP hook: POST context JSON and parse the response. /// Fail-open: non-2xx and connection errors return `Proceed`. async fn execute_http( + client: &reqwest::Client, url: &str, headers: &Option>, allowed_env_vars: &[String], @@ -381,14 +400,7 @@ impl HookExecutorImpl { TlsMode::Off => {} } - let accept_invalid = matches!(tls, TlsMode::NoVerify | TlsMode::Off); - let client = reqwest::Client::builder() - .timeout(timeout) - .danger_accept_invalid_certs(accept_invalid) - .build() - .unwrap_or_default(); - - let mut request = client.post(url).json(context); + let mut request = client.post(url).timeout(timeout).json(context); if let Some(hdrs) = headers { for (key, value) in hdrs { @@ -436,6 +448,37 @@ impl HookExecutorImpl { } } +/// Cached HTTP clients keyed by TLS mode. +struct HttpClientCache { + verify: reqwest::Client, + no_verify: reqwest::Client, + off: reqwest::Client, +} + +impl HttpClientCache { + fn new() -> Self { + Self { + verify: HookExecutorImpl::build_http_client(&TlsMode::Verify), + no_verify: HookExecutorImpl::build_http_client(&TlsMode::NoVerify), + off: HookExecutorImpl::build_http_client(&TlsMode::Off), + } + } + + fn get(&self, tls: &TlsMode) -> &reqwest::Client { + match tls { + TlsMode::Verify => &self.verify, + TlsMode::NoVerify => &self.no_verify, + TlsMode::Off => &self.off, + } + } +} + +impl Default for HttpClientCache { + fn default() -> Self { + Self::new() + } +} + #[async_trait] impl HookExecutor for HookExecutorImpl { async fn execute( @@ -445,32 +488,60 @@ impl HookExecutor for HookExecutorImpl { sandbox: Arc, work_dir: Option<&Path>, ) -> HookResult { + use std::sync::OnceLock; + static HTTP_CLIENTS: OnceLock = OnceLock::new(); + let start = Instant::now(); let decision = match definition.resolved_hook_type() { - Some(HookType::Command { ref command }) => { + Some(Cow::Borrowed(HookType::Command { ref command }) + | Cow::Owned(HookType::Command { ref command })) => { Self::execute_command(definition, command, context, &sandbox, work_dir).await } - Some(HookType::Http { + Some(Cow::Borrowed(HookType::Http { ref url, ref headers, ref allowed_env_vars, ref tls, - }) => { - Self::execute_http(url, headers, allowed_env_vars, tls, context, definition.timeout()) - .await + }) + | Cow::Owned(HookType::Http { + ref url, + ref headers, + ref allowed_env_vars, + ref tls, + })) => { + let clients = HTTP_CLIENTS.get_or_init(HttpClientCache::new); + Self::execute_http( + clients.get(tls), + url, + headers, + allowed_env_vars, + tls, + context, + definition.timeout(), + ) + .await } - Some(HookType::Prompt { + Some(Cow::Borrowed(HookType::Prompt { ref prompt, ref model, - }) => { + }) + | Cow::Owned(HookType::Prompt { + ref prompt, + ref model, + })) => { Self::execute_prompt(prompt, model, context, definition.timeout()).await } - Some(HookType::Agent { + Some(Cow::Borrowed(HookType::Agent { ref prompt, ref model, ref max_tool_rounds, - }) => { + }) + | Cow::Owned(HookType::Agent { + ref prompt, + ref model, + ref max_tool_rounds, + })) => { Self::execute_agent( prompt, model, @@ -510,6 +581,10 @@ mod tests { Arc::new(arc_agent::LocalSandbox::new(std::env::current_dir().unwrap())) } + fn test_http_client() -> reqwest::Client { + HookExecutorImpl::build_http_client(&TlsMode::Off) + } + fn make_definition(command: &str) -> HookDefinition { HookDefinition { name: Some("test-hook".into()), @@ -801,7 +876,9 @@ mod tests { .create_async() .await; + let client = test_http_client(); let decision = HookExecutorImpl::execute_http( + &client, &format!("{}/hook", server.url()), &None, &[], @@ -830,7 +907,9 @@ mod tests { .create_async() .await; + let client = test_http_client(); let decision = HookExecutorImpl::execute_http( + &client, &format!("{}/hook", server.url()), &None, &[], @@ -854,7 +933,9 @@ mod tests { .create_async() .await; + let client = test_http_client(); let decision = HookExecutorImpl::execute_http( + &client, &format!("{}/hook", server.url()), &None, &[], @@ -870,7 +951,9 @@ mod tests { #[tokio::test] async fn http_hook_connection_failure_returns_proceed() { + let client = test_http_client(); let decision = HookExecutorImpl::execute_http( + &client, "http://127.0.0.1:1", &None, &[], @@ -900,7 +983,9 @@ mod tests { ("Authorization".to_string(), "Bearer $ARC_TEST_TOKEN".to_string()), ]); + let client = test_http_client(); let decision = HookExecutorImpl::execute_http( + &client, &format!("{}/hook", server.url()), &Some(headers), &["ARC_TEST_TOKEN".to_string()], @@ -919,7 +1004,9 @@ mod tests { #[tokio::test] async fn http_hook_rejects_http_url_when_tls_verify() { + let client = test_http_client(); let decision = HookExecutorImpl::execute_http( + &client, "http://example.com/hook", &None, &[], @@ -934,7 +1021,9 @@ mod tests { #[tokio::test] async fn http_hook_rejects_http_url_when_tls_no_verify() { + let client = test_http_client(); let decision = HookExecutorImpl::execute_http( + &client, "http://example.com/hook", &None, &[], @@ -957,7 +1046,9 @@ mod tests { .create_async() .await; + let client = test_http_client(); let decision = HookExecutorImpl::execute_http( + &client, &format!("{}/hook", server.url()), &None, &[], diff --git a/crates/arc-workflows/src/hook/runner.rs b/crates/arc-workflows/src/hook/runner.rs index d435080e4..dcdc5e166 100644 --- a/crates/arc-workflows/src/hook/runner.rs +++ b/crates/arc-workflows/src/hook/runner.rs @@ -11,7 +11,7 @@ use super::types::{HookContext, HookDecision}; /// Central orchestrator: filters matching hooks, executes them, merges decisions. pub struct HookRunner { config: HookConfig, - command_executor: Arc, + executor: Arc, /// Pre-compiled regexes keyed by matcher pattern string. compiled_matchers: HashMap, } @@ -22,7 +22,7 @@ impl HookRunner { let compiled_matchers = Self::compile_matchers(&config); Self { config, - command_executor: Arc::new(HookExecutorImpl), + executor: Arc::new(HookExecutorImpl), compiled_matchers, } } @@ -33,7 +33,7 @@ impl HookRunner { let compiled_matchers = Self::compile_matchers(&config); Self { config, - command_executor: executor, + executor: executor, compiled_matchers, } } @@ -136,7 +136,7 @@ impl HookRunner { "Executing hook" ); let result = self - .command_executor + .executor .execute(hook, context, sandbox.clone(), work_dir) .await; tracing::debug!( @@ -184,7 +184,7 @@ impl HookRunner { "Executing hook" ); let result = self - .command_executor + .executor .execute(hook, context, sandbox.clone(), work_dir) .await; tracing::debug!( @@ -398,7 +398,7 @@ mod tests { } #[tokio::test] - async fn command_executor_integration_success() { + async fn executor_integration_success() { let config = HookConfig { hooks: vec![{ let mut h = make_hook(HookEvent::RunStart, "echo-hook"); @@ -414,7 +414,7 @@ mod tests { } #[tokio::test] - async fn command_executor_integration_block() { + async fn executor_integration_block() { let config = HookConfig { hooks: vec![{ let mut h = make_hook(HookEvent::RunStart, "fail-hook"); diff --git a/crates/arc-workflows/tests/integration.rs b/crates/arc-workflows/tests/integration.rs index 92154cc7f..9026406f7 100644 --- a/crates/arc-workflows/tests/integration.rs +++ b/crates/arc-workflows/tests/integration.rs @@ -7683,20 +7683,20 @@ timeout_ms = 120000 // Prompt hook assert_eq!(cfg.hooks[0].event, arc_workflows::hook::HookEvent::StageStart); assert!(matches!( - cfg.hooks[0].resolved_hook_type(), + cfg.hooks[0].resolved_hook_type().as_deref(), Some(arc_workflows::hook::HookType::Prompt { prompt, model }) - if prompt == "Should this stage proceed?" && model == Some("haiku".into()) + if prompt == "Should this stage proceed?" && *model == Some("haiku".into()) )); assert_eq!(cfg.hooks[0].timeout(), std::time::Duration::from_millis(30000)); // Agent hook assert_eq!(cfg.hooks[1].event, arc_workflows::hook::HookEvent::RunComplete); assert!(matches!( - cfg.hooks[1].resolved_hook_type(), + cfg.hooks[1].resolved_hook_type().as_deref(), Some(arc_workflows::hook::HookType::Agent { prompt, model, max_tool_rounds }) if prompt == "Verify all tests pass." - && model == Some("sonnet".into()) - && max_tool_rounds == Some(10) + && *model == Some("sonnet".into()) + && *max_tool_rounds == Some(10) )); assert_eq!(cfg.hooks[1].timeout(), std::time::Duration::from_millis(120000)); }