use std::{fmt, path::PathBuf, sync::Arc, time::Duration};
use mentra::{
error::RuntimeError,
runtime::{
HookDecision, PostExecutionContext, PostExecutionHook, PreExecutionContext,
PreExecutionHook, ResultDecision,
},
tool::ToolResultContent,
};
use serde_json::Value;
use thiserror::Error;
use crate::subprocess::{self, Completion};
use super::{
HookEvent, HookSpec, Interceptor,
chain::{Answer, Chain, Participant},
contract::{HookCall, HookOutcome, HookRequest},
wire::HookResponse,
};
#[derive(Clone)]
pub struct HookRunner {
workspace: PathBuf,
interceptors: Vec<Arc<dyn Interceptor>>,
hooks: Vec<HookSpec>,
report: Arc<dyn Fn(&str) + Send + Sync>,
}
impl HookRunner {
pub fn new(workspace: impl Into<PathBuf>, hooks: Vec<HookSpec>) -> Self {
Self {
workspace: workspace.into(),
interceptors: Vec::new(),
hooks,
report: Arc::new(|message| eprintln!("basis: {message}")),
}
}
pub fn with_interceptor(self, interceptor: impl Interceptor + 'static) -> Self {
Self {
interceptors: {
let mut interceptors = self.interceptors;
interceptors.push(Arc::new(interceptor));
interceptors
},
..self
}
}
pub fn with_reporter(self, report: impl Fn(&str) + Send + Sync + 'static) -> Self {
Self {
report: Arc::new(report),
..self
}
}
pub fn is_empty(&self) -> bool {
self.hooks.is_empty() && self.interceptors.is_empty()
}
pub fn decide(&self, call: &HookCall) -> HookOutcome {
if !self.interceptors.is_empty() {
return HookOutcome::Deny(
"in-process interceptors are registered and cannot be consulted synchronously; \
this call belongs on HookRunner::decide_async"
.to_string(),
);
}
if self.hooks.is_empty() {
return HookOutcome::Allow;
}
match self.consult_hooks(Chain::new(self.request(call))) {
Ok(chain) => chain.outcome(),
Err(outcome) => outcome,
}
}
pub async fn decide_async(&self, call: &HookCall) -> HookOutcome {
if self.is_empty() {
return HookOutcome::Allow;
}
self.consult(HookRequest::from_call(
HookEvent::PreToolUse,
&self.workspace,
call,
))
.await
}
pub async fn review_async(
&self,
call: &HookCall,
output: Value,
is_error: bool,
) -> HookOutcome {
if self.is_empty() {
return HookOutcome::Allow;
}
self.consult(HookRequest::from_result(
&self.workspace,
call,
output,
is_error,
))
.await
}
async fn consult(&self, request: HookRequest) -> HookOutcome {
let chain = match self.consult_interceptors(Chain::new(request)).await {
Ok(chain) => chain,
Err(outcome) => return outcome,
};
if self.hooks.is_empty() {
return chain.outcome();
}
let runner = self.clone();
match tokio::task::spawn_blocking(move || runner.consult_hooks(chain)).await {
Ok(Ok(chain)) => chain.outcome(),
Ok(Err(outcome)) => outcome,
Err(error) => HookOutcome::Deny(format!("hook runner failed: {error}")),
}
}
fn request(&self, call: &HookCall) -> HookRequest {
HookRequest::from_call(HookEvent::PreToolUse, &self.workspace, call)
}
async fn consult_interceptors(&self, chain: Chain) -> Result<Chain, HookOutcome> {
let mut chain = chain;
for interceptor in &self.interceptors {
let answer = match self.ask_interceptor(interceptor, chain.request()).await {
Ok(HookOutcome::Allow) => Answer::Allow,
Ok(HookOutcome::Deny(reason)) => Answer::Deny(Some(reason)),
Ok(HookOutcome::Modify { input, reason }) => Answer::Modify { input, reason },
Ok(HookOutcome::Replace {
output,
is_error,
reason,
}) => Answer::Replace {
output,
is_error: Some(is_error),
reason,
},
Err(failure) => Answer::Broken(failure),
};
chain = chain.advance(
Participant::interceptor(interceptor.name()),
answer,
&*self.report,
)?;
}
Ok(chain)
}
fn consult_hooks(&self, chain: Chain) -> Result<Chain, HookOutcome> {
let mut chain = chain;
for spec in &self.hooks {
if !spec.applies_to(chain.request().event, &chain.request().tool_name) {
continue;
}
let answer = match self.ask(spec, chain.request()) {
Ok(HookResponse::Allow { .. }) => Answer::Allow,
Ok(HookResponse::Deny { reason }) => Answer::Deny(reason),
Ok(HookResponse::Modify { input, reason }) => Answer::Modify { input, reason },
Ok(HookResponse::Replace {
output,
is_error,
reason,
}) => Answer::Replace {
output,
is_error,
reason,
},
Err(failure) => Answer::Broken(failure.to_string()),
};
chain = chain.advance(
Participant::hook(&spec.name, spec.on_failure),
answer,
&*self.report,
)?;
}
Ok(chain)
}
async fn ask_interceptor(
&self,
interceptor: &Arc<dyn Interceptor>,
request: &HookRequest,
) -> Result<HookOutcome, String> {
let interceptor = Arc::clone(interceptor);
let request = request.clone();
match tokio::spawn(async move {
match request.event {
HookEvent::PreToolUse => interceptor.intercept(&request).await,
HookEvent::PostToolUse => interceptor.review(&request).await,
}
})
.await
{
Ok(Ok(outcome)) => Ok(outcome),
Ok(Err(error)) => Err(format!("answered with an error: {error}")),
Err(error) if error.is_panic() => {
Err(format!("panicked: {}", panic_message(error.into_panic())))
}
Err(error) => Err(format!("could not be asked: {error}")),
}
}
fn ask(&self, spec: &HookSpec, request: &HookRequest) -> Result<HookResponse, HookFailure> {
let payload = serde_json::to_string(request).map_err(HookFailure::Payload)?;
let completion = subprocess::execute(
&spec.command,
&self.workspace,
&[],
&payload,
spec.timeout(),
)
.map_err(HookFailure::Spawn)?;
let (code, stdout, stderr) = match completion {
Completion::TimedOut => {
return Err(HookFailure::TimedOut {
timeout: spec.timeout(),
});
}
Completion::Exited {
code,
stdout,
stderr,
} => (code, stdout, stderr),
};
if code != Some(0) {
return Err(HookFailure::Exited {
code: code.map_or_else(|| "a signal".to_string(), |code| format!("code {code}")),
stderr,
});
}
if stdout.trim().is_empty() {
return Err(HookFailure::NoAnswer);
}
serde_json::from_str(&stdout).map_err(|source| HookFailure::Malformed {
output: subprocess::truncated_output(&stdout),
source,
})
}
}
#[async_trait::async_trait]
impl PreExecutionHook for HookRunner {
async fn pre_tool_execution(
&self,
context: &PreExecutionContext,
) -> Result<HookDecision, RuntimeError> {
let call = HookCall::new(
context.agent_id.clone(),
context.tool_name.clone(),
context.tool_call_id.clone(),
context.input_json.clone(),
);
Ok(match self.decide_async(&call).await {
HookOutcome::Allow => HookDecision::Allow,
HookOutcome::Deny(reason) => HookDecision::Deny(reason),
HookOutcome::Modify { input, reason } => match serde_json::to_string(&input) {
Ok(input_json) => HookDecision::Modify { input_json, reason },
Err(error) => HookDecision::Deny(format!(
"a replacement input could not be re-encoded: {error}"
)),
},
HookOutcome::Replace { .. } => HookDecision::Deny(
"a participant replaced the result of a call that has not run yet".to_string(),
),
})
}
}
#[async_trait::async_trait]
impl PostExecutionHook for HookRunner {
async fn post_tool_execution(
&self,
context: &PostExecutionContext,
) -> Result<ResultDecision, RuntimeError> {
let call = HookCall::new(
context.agent_id.clone(),
context.tool_name.clone(),
context.tool_call_id.clone(),
context.input_json.clone(),
);
Ok(
match self
.review_async(&call, as_json(&context.content), context.is_error)
.await
{
HookOutcome::Allow => ResultDecision::Keep,
HookOutcome::Replace {
output, is_error, ..
} => ResultDecision::Replace {
content: as_content(output),
is_error,
},
HookOutcome::Deny(reason) => ResultDecision::Replace {
content: ToolResultContent::text(reason),
is_error: true,
},
HookOutcome::Modify { .. } => ResultDecision::Replace {
content: ToolResultContent::text(
"a participant rewrote the input of a call that had already run",
),
is_error: true,
},
},
)
}
}
fn as_json(content: &ToolResultContent) -> Value {
match content {
ToolResultContent::Text(text) => Value::String(text.clone()),
ToolResultContent::Structured(value) => value.clone(),
}
}
fn as_content(output: Value) -> ToolResultContent {
match output {
Value::String(text) => ToolResultContent::Text(text),
other => ToolResultContent::Structured(other),
}
}
impl fmt::Debug for HookRunner {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("HookRunner")
.field("workspace", &self.workspace)
.field(
"interceptors",
&self
.interceptors
.iter()
.map(|interceptor| interceptor.name())
.collect::<Vec<_>>(),
)
.field(
"hooks",
&self.hooks.iter().map(|spec| &spec.name).collect::<Vec<_>>(),
)
.finish_non_exhaustive()
}
}
fn panic_message(payload: Box<dyn std::any::Any + Send>) -> String {
if let Some(message) = payload.downcast_ref::<&str>() {
return (*message).to_string();
}
if let Some(message) = payload.downcast_ref::<String>() {
return message.clone();
}
"with a payload that is not a message".to_string()
}
#[derive(Debug, Error)]
enum HookFailure {
#[error("could not be started: {0}")]
Spawn(#[source] std::io::Error),
#[error("did not answer within {}ms and was killed", .timeout.as_millis())]
TimedOut { timeout: Duration },
#[error("exited with {code}{}", stderr_tail(.stderr))]
Exited { code: String, stderr: String },
#[error("printed nothing; a hook answers with a JSON decision on stdout")]
NoAnswer,
#[error("printed something that is not a decision ({source}): {output}")]
Malformed {
output: String,
#[source]
source: serde_json::Error,
},
#[error("could not be asked, because the request would not serialize: {0}")]
Payload(#[source] serde_json::Error),
}
fn stderr_tail(stderr: &str) -> String {
let stderr = stderr.trim();
if stderr.is_empty() {
String::new()
} else {
format!(" and said: {stderr}")
}
}
#[cfg(all(test, unix))]
mod tests;