Simplify hook implementation: deduplicate LLM setup, fix async I/O, cache HTTP clients

- 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 <noreply@anthropic.com>
This commit is contained in:
Bryan Helmkamp 2026-03-05 02:49:06 -05:00
parent c7e7b30906
commit 4f4f1f98a3
4 changed files with 188 additions and 95 deletions

View file

@ -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<HookType> {
pub fn resolved_hook_type(&self) -> Option<Cow<'_, HookType>> {
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)
));
}

View file

@ -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>) -> 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<F, Fut>(
timeout: std::time::Duration,
hook_kind: &str,
f: F,
) -> HookDecision
where
F: FnOnce() -> Fut,
Fut: std::future::Future<Output = HookDecision>,
{
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<dyn Sandbox>,
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<HashMap<String, String>>,
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<dyn Sandbox>,
work_dir: Option<&Path>,
) -> HookResult {
use std::sync::OnceLock;
static HTTP_CLIENTS: OnceLock<HttpClientCache> = 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,
&[],

View file

@ -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<dyn HookExecutor>,
executor: Arc<dyn HookExecutor>,
/// Pre-compiled regexes keyed by matcher pattern string.
compiled_matchers: HashMap<String, regex::Regex>,
}
@ -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");

View file

@ -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));
}