use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use crate::{
OneOrMany,
completion::{CompletionModel, Document, Usage},
json_utils,
message::{AssistantContent, Message, ToolChoice},
tool::{ToolCallExtensions, ToolOutcome, ToolResultExtensions},
wasm_compat::{WasmBoxedFuture, WasmCompatSend, WasmCompatSync},
};
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct RunId(String);
impl RunId {
pub(crate) fn generate() -> Self {
Self(crate::id::generate())
}
pub fn as_str(&self) -> &str {
&self.0
}
}
impl std::fmt::Display for RunId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(&self.0)
}
}
#[derive(Clone, Default)]
pub struct Scratchpad {
inner: Arc<std::sync::Mutex<ToolCallExtensions>>,
}
impl Scratchpad {
fn lock(&self) -> std::sync::MutexGuard<'_, ToolCallExtensions> {
self.inner.lock().unwrap_or_else(|e| e.into_inner())
}
pub fn insert<T: Clone + WasmCompatSend + WasmCompatSync + 'static>(
&self,
val: T,
) -> Option<T> {
self.lock().insert(val)
}
pub fn get<T: Clone + WasmCompatSend + WasmCompatSync + 'static>(&self) -> Option<T> {
self.lock().get::<T>().cloned()
}
pub fn contains<T: WasmCompatSend + WasmCompatSync + 'static>(&self) -> bool {
self.lock().contains::<T>()
}
pub fn remove<T: Clone + WasmCompatSend + WasmCompatSync + 'static>(&self) -> Option<T> {
self.lock().remove::<T>()
}
pub fn update<T, R>(&self, f: impl FnOnce(&mut T) -> R) -> R
where
T: Clone + Default + WasmCompatSend + WasmCompatSync + 'static,
{
let mut guard = self.lock();
let mut val = guard.remove::<T>().unwrap_or_default();
let out = f(&mut val);
guard.insert(val);
out
}
}
impl std::fmt::Debug for Scratchpad {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Scratchpad")
.field("entries", &self.lock().len())
.finish()
}
}
#[derive(Debug)]
pub struct HookContext {
run_id: RunId,
turn: AtomicUsize,
is_streaming: bool,
agent_name: Option<String>,
scratchpad: Scratchpad,
}
impl HookContext {
pub(crate) fn new(is_streaming: bool, agent_name: Option<String>) -> Self {
Self {
run_id: RunId::generate(),
turn: AtomicUsize::new(0),
is_streaming,
agent_name,
scratchpad: Scratchpad::default(),
}
}
pub(crate) fn set_turn(&self, turn: usize) {
self.turn.store(turn, Ordering::Relaxed);
}
pub fn run_id(&self) -> &RunId {
&self.run_id
}
pub fn turn(&self) -> usize {
self.turn.load(Ordering::Relaxed)
}
pub fn is_streaming(&self) -> bool {
self.is_streaming
}
pub fn agent_name(&self) -> Option<&str> {
self.agent_name.as_deref()
}
pub fn scratchpad(&self) -> &Scratchpad {
&self.scratchpad
}
}
#[cfg(not(target_family = "wasm"))]
const _: fn() = || {
fn assert_send_sync<T: Send + Sync>() {}
assert_send_sync::<HookContext>();
};
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct InvalidToolCallContext {
pub tool_name: String,
pub tool_call_id: Option<String>,
pub internal_call_id: Option<String>,
pub args: Option<String>,
pub available_tools: Vec<String>,
pub allowed_tools: Vec<String>,
pub tool_choice: Option<ToolChoice>,
pub chat_history: Vec<Message>,
pub is_streaming: bool,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum InvalidToolCallHookAction {
Fail,
Retry { feedback: String },
Repair { tool_name: String },
Skip { reason: String },
}
impl InvalidToolCallHookAction {
pub fn fail() -> Self {
Self::Fail
}
pub fn retry(feedback: impl Into<String>) -> Self {
Self::Retry {
feedback: feedback.into(),
}
}
pub fn repair(tool_name: impl Into<String>) -> Self {
Self::Repair {
tool_name: tool_name.into(),
}
}
pub fn skip(reason: impl Into<String>) -> Self {
Self::Skip {
reason: reason.into(),
}
}
}
#[non_exhaustive]
pub enum StepEvent<'a, M: CompletionModel> {
CompletionCall {
prompt: &'a Message,
history: &'a [Message],
turn: usize,
},
CompletionResponse {
prompt: &'a Message,
response: &'a crate::completion::CompletionResponse<M::Response>,
},
ModelTurnFinished {
turn: usize,
content: &'a OneOrMany<AssistantContent>,
usage: Usage,
},
InvalidToolCall(&'a InvalidToolCallContext),
ToolCall {
tool_name: &'a str,
tool_call_id: Option<&'a str>,
internal_call_id: &'a str,
args: &'a str,
},
ToolResult {
tool_name: &'a str,
tool_call_id: Option<&'a str>,
internal_call_id: &'a str,
args: &'a str,
result: &'a str,
outcome: &'a ToolOutcome,
extensions: &'a ToolResultExtensions,
},
TextDelta {
delta: &'a str,
aggregated: &'a str,
},
ToolCallDelta {
tool_call_id: &'a str,
internal_call_id: &'a str,
tool_name: Option<&'a str>,
delta: &'a str,
},
StreamResponseFinish {
prompt: &'a Message,
response: &'a M::StreamingResponse,
},
}
impl<M: CompletionModel> Clone for StepEvent<'_, M> {
fn clone(&self) -> Self {
*self
}
}
impl<M: CompletionModel> Copy for StepEvent<'_, M> {}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum StepEventKind {
CompletionCall,
CompletionResponse,
ModelTurnFinished,
InvalidToolCall,
ToolCall,
ToolResult,
TextDelta,
ToolCallDelta,
StreamResponseFinish,
}
impl<M: CompletionModel> StepEvent<'_, M> {
pub fn kind(&self) -> StepEventKind {
match self {
StepEvent::CompletionCall { .. } => StepEventKind::CompletionCall,
StepEvent::CompletionResponse { .. } => StepEventKind::CompletionResponse,
StepEvent::ModelTurnFinished { .. } => StepEventKind::ModelTurnFinished,
StepEvent::InvalidToolCall(_) => StepEventKind::InvalidToolCall,
StepEvent::ToolCall { .. } => StepEventKind::ToolCall,
StepEvent::ToolResult { .. } => StepEventKind::ToolResult,
StepEvent::TextDelta { .. } => StepEventKind::TextDelta,
StepEvent::ToolCallDelta { .. } => StepEventKind::ToolCallDelta,
StepEvent::StreamResponseFinish { .. } => StepEventKind::StreamResponseFinish,
}
}
}
#[derive(Debug, Clone, Default, PartialEq)]
#[non_exhaustive]
pub struct RequestPatch {
pub preamble: Option<String>,
pub temperature: Option<f64>,
pub max_tokens: Option<u64>,
pub tool_choice: Option<ToolChoice>,
pub active_tools: Option<Vec<String>>,
pub additional_params: Option<serde_json::Value>,
pub extra_context: Vec<Document>,
pub history: Option<Vec<Message>>,
}
fn merge_last_wins<T>(earlier: Option<T>, later: Option<T>, field: &str) -> Option<T> {
match (earlier, later) {
(Some(_), Some(l)) => {
tracing::warn!(
patch_field = field,
"two hooks set `{field}` on the same turn; the later hook wins"
);
Some(l)
}
(earlier, later) => later.or(earlier),
}
}
impl RequestPatch {
pub fn new() -> Self {
Self::default()
}
pub fn preamble(mut self, preamble: impl Into<String>) -> Self {
self.preamble = Some(preamble.into());
self
}
pub fn temperature(mut self, temperature: f64) -> Self {
self.temperature = Some(temperature);
self
}
pub fn max_tokens(mut self, max_tokens: u64) -> Self {
self.max_tokens = Some(max_tokens);
self
}
pub fn tool_choice(mut self, tool_choice: ToolChoice) -> Self {
self.tool_choice = Some(tool_choice);
self
}
pub fn active_tools<I, S>(mut self, names: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.active_tools = Some(names.into_iter().map(Into::into).collect());
self
}
pub fn additional_params(mut self, additional_params: serde_json::Value) -> Self {
self.additional_params = Some(additional_params);
self
}
pub fn extra_context<I>(mut self, docs: I) -> Self
where
I: IntoIterator<Item = Document>,
{
self.extra_context.extend(docs);
self
}
pub fn context(mut self, doc: Document) -> Self {
self.extra_context.push(doc);
self
}
pub fn history<I>(mut self, history: I) -> Self
where
I: IntoIterator<Item = Message>,
{
self.history = Some(history.into_iter().collect());
self
}
pub(crate) fn is_empty(&self) -> bool {
self.preamble.is_none()
&& self.temperature.is_none()
&& self.max_tokens.is_none()
&& self.tool_choice.is_none()
&& self.active_tools.is_none()
&& self.additional_params.is_none()
&& self.extra_context.is_empty()
&& self.history.is_none()
}
pub(crate) fn merge(mut self, later: RequestPatch) -> RequestPatch {
self.extra_context.extend(later.extra_context);
self.additional_params = match (self.additional_params.take(), later.additional_params) {
(Some(base), Some(patch)) if base.is_object() && patch.is_object() => {
Some(json_utils::merge(base, patch))
}
(base, patch) => patch.or(base),
};
self.preamble = merge_last_wins(self.preamble, later.preamble, "preamble");
self.temperature = merge_last_wins(self.temperature, later.temperature, "temperature");
self.max_tokens = merge_last_wins(self.max_tokens, later.max_tokens, "max_tokens");
self.tool_choice = merge_last_wins(self.tool_choice, later.tool_choice, "tool_choice");
self.history = merge_last_wins(self.history, later.history, "history");
self.active_tools = match (self.active_tools.take(), later.active_tools) {
(Some(earlier), Some(later)) => {
let later_set: std::collections::BTreeSet<&String> = later.iter().collect();
let intersection: Vec<String> = earlier
.into_iter()
.filter(|name| later_set.contains(name))
.collect();
if intersection.is_empty() {
tracing::warn!(
"two hooks' `active_tools` allow-lists have an empty intersection; \
no executable tools will be advertised this turn"
);
}
Some(intersection)
}
(earlier, later) => earlier.or(later),
};
self
}
}
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
pub enum Flow {
Continue,
Terminate {
reason: String,
},
Skip {
reason: String,
},
RewriteArgs {
args: serde_json::Value,
},
RewriteResult {
result: String,
},
PatchRequest {
patch: RequestPatch,
},
Fail,
Retry {
feedback: String,
},
Repair {
tool_name: String,
},
}
impl Flow {
pub fn cont() -> Self {
Self::Continue
}
pub fn terminate(reason: impl Into<String>) -> Self {
Self::Terminate {
reason: reason.into(),
}
}
pub fn skip(reason: impl Into<String>) -> Self {
Self::Skip {
reason: reason.into(),
}
}
pub fn rewrite_args(args: impl Into<serde_json::Value>) -> Self {
Self::RewriteArgs { args: args.into() }
}
pub fn try_rewrite_args<T: serde::Serialize>(value: &T) -> Result<Self, serde_json::Error> {
Ok(Self::RewriteArgs {
args: serde_json::to_value(value)?,
})
}
pub fn rewrite_result(result: impl Into<String>) -> Self {
Self::RewriteResult {
result: result.into(),
}
}
pub fn patch_request(patch: RequestPatch) -> Self {
Self::PatchRequest { patch }
}
pub fn fail() -> Self {
Self::Fail
}
pub fn retry(feedback: impl Into<String>) -> Self {
Self::Retry {
feedback: feedback.into(),
}
}
pub fn repair(tool_name: impl Into<String>) -> Self {
Self::Repair {
tool_name: tool_name.into(),
}
}
}
pub trait AgentHook<M>: WasmCompatSend + WasmCompatSync
where
M: CompletionModel,
{
fn on_event(
&self,
ctx: &HookContext,
event: StepEvent<'_, M>,
) -> impl Future<Output = Flow> + WasmCompatSend {
let _ = (ctx, event);
async { Flow::Continue }
}
#[doc(hidden)]
fn resolve_tool_call(
&self,
ctx: &HookContext,
tool_name: &str,
tool_call_id: Option<&str>,
internal_call_id: &str,
args: &str,
) -> impl Future<Output = (Flow, Option<serde_json::Value>)> + WasmCompatSend {
async move {
let flow = self
.on_event(
ctx,
StepEvent::ToolCall {
tool_name,
tool_call_id,
internal_call_id,
args,
},
)
.await;
(flow, None)
}
}
fn observes(&self, kind: StepEventKind) -> bool {
let _ = kind;
true
}
}
impl<M> AgentHook<M> for ()
where
M: CompletionModel,
{
fn observes(&self, _kind: StepEventKind) -> bool {
false
}
}
trait DynAgentHook<M>: WasmCompatSend + WasmCompatSync
where
M: CompletionModel,
{
fn on_event_boxed<'a>(
&'a self,
ctx: &'a HookContext,
event: StepEvent<'a, M>,
) -> WasmBoxedFuture<'a, Flow>
where
M: 'a;
fn resolve_tool_call_boxed<'a>(
&'a self,
ctx: &'a HookContext,
tool_name: &'a str,
tool_call_id: Option<&'a str>,
internal_call_id: &'a str,
args: &'a str,
) -> WasmBoxedFuture<'a, (Flow, Option<serde_json::Value>)>
where
M: 'a;
fn observes_dyn(&self, kind: StepEventKind) -> bool;
}
impl<M, H> DynAgentHook<M> for H
where
M: CompletionModel,
H: AgentHook<M>,
{
fn on_event_boxed<'a>(
&'a self,
ctx: &'a HookContext,
event: StepEvent<'a, M>,
) -> WasmBoxedFuture<'a, Flow>
where
M: 'a,
{
Box::pin(self.on_event(ctx, event))
}
fn resolve_tool_call_boxed<'a>(
&'a self,
ctx: &'a HookContext,
tool_name: &'a str,
tool_call_id: Option<&'a str>,
internal_call_id: &'a str,
args: &'a str,
) -> WasmBoxedFuture<'a, (Flow, Option<serde_json::Value>)>
where
M: 'a,
{
Box::pin(self.resolve_tool_call(ctx, tool_name, tool_call_id, internal_call_id, args))
}
fn observes_dyn(&self, kind: StepEventKind) -> bool {
self.observes(kind)
}
}
pub struct HookStack<M>
where
M: CompletionModel,
{
hooks: Vec<Arc<dyn DynAgentHook<M>>>,
}
impl<M> Clone for HookStack<M>
where
M: CompletionModel,
{
fn clone(&self) -> Self {
Self {
hooks: self.hooks.clone(),
}
}
}
impl<M> Default for HookStack<M>
where
M: CompletionModel,
{
fn default() -> Self {
Self { hooks: Vec::new() }
}
}
impl<M> std::fmt::Debug for HookStack<M>
where
M: CompletionModel,
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("HookStack")
.field("len", &self.hooks.len())
.finish()
}
}
impl<M> HookStack<M>
where
M: CompletionModel,
{
pub fn new() -> Self {
Self::default()
}
pub fn with<H>(hook: H) -> Self
where
H: AgentHook<M> + 'static,
{
let mut stack = Self::new();
stack.push(hook);
stack
}
pub fn push<H>(&mut self, hook: H)
where
H: AgentHook<M> + 'static,
{
self.hooks.push(Arc::new(hook));
}
pub fn is_empty(&self) -> bool {
self.hooks.is_empty()
}
pub fn len(&self) -> usize {
self.hooks.len()
}
}
impl<M> AgentHook<M> for HookStack<M>
where
M: CompletionModel,
{
async fn resolve_tool_call(
&self,
ctx: &HookContext,
tool_name: &str,
tool_call_id: Option<&str>,
internal_call_id: &str,
args: &str,
) -> (Flow, Option<serde_json::Value>) {
let mut effective: Option<serde_json::Value> = None;
for hook in &self.hooks {
let rewritten = effective.as_ref().map(json_utils::value_to_json_string);
let args_for_hook = rewritten.as_deref().unwrap_or(args);
let (flow, salvaged) = hook
.resolve_tool_call_boxed(
ctx,
tool_name,
tool_call_id,
internal_call_id,
args_for_hook,
)
.await;
if let Some(rewrite) = salvaged {
effective = Some(rewrite);
}
match flow {
Flow::Continue => {}
Flow::RewriteArgs { args } => effective = Some(args),
other => return (other, effective),
}
}
match effective {
Some(args) => (Flow::RewriteArgs { args }, None),
None => (Flow::Continue, None),
}
}
async fn on_event(&self, ctx: &HookContext, event: StepEvent<'_, M>) -> Flow {
match event {
StepEvent::CompletionCall { .. } => {
let mut merged: Option<RequestPatch> = None;
for hook in &self.hooks {
match hook.on_event_boxed(ctx, event).await {
Flow::Continue => {}
Flow::PatchRequest { patch } => {
merged = Some(match merged {
Some(acc) => acc.merge(patch),
None => patch,
});
}
other => return other,
}
}
match merged {
Some(patch) if !patch.is_empty() => Flow::PatchRequest { patch },
_ => Flow::Continue,
}
}
StepEvent::ToolCall {
tool_name,
tool_call_id,
internal_call_id,
args,
} => {
self.resolve_tool_call(ctx, tool_name, tool_call_id, internal_call_id, args)
.await
.0
}
StepEvent::ToolResult {
tool_name,
tool_call_id,
internal_call_id,
args,
result,
outcome,
extensions,
} => {
let mut effective: Option<String> = None;
for hook in &self.hooks {
let result_for_hook = effective.as_deref().unwrap_or(result);
let per_hook = StepEvent::ToolResult {
tool_name,
tool_call_id,
internal_call_id,
args,
result: result_for_hook,
outcome,
extensions,
};
match hook.on_event_boxed(ctx, per_hook).await {
Flow::Continue => {}
Flow::RewriteResult { result } => effective = Some(result),
other => return other,
}
}
match effective {
Some(result) => Flow::RewriteResult { result },
None => Flow::Continue,
}
}
_ => {
for hook in &self.hooks {
match hook.on_event_boxed(ctx, event).await {
Flow::Continue => {}
other => return other,
}
}
Flow::Continue
}
}
}
fn observes(&self, kind: StepEventKind) -> bool {
self.hooks.iter().any(|hook| hook.observes_dyn(kind))
}
}
#[cfg(test)]
mod tests {
use std::sync::{Arc, Mutex};
use super::{
AgentHook, Flow, HookContext, HookStack, RequestPatch, Scratchpad, StepEvent, StepEventKind,
};
use crate::test_utils::MockCompletionModel;
type M = MockCompletionModel;
fn ctx() -> HookContext {
HookContext::new(false, Some("test-agent".to_string()))
}
struct Recorder {
label: u32,
log: Arc<Mutex<Vec<u32>>>,
stop: bool,
}
impl AgentHook<M> for Recorder {
async fn on_event(&self, _ctx: &HookContext, _event: StepEvent<'_, M>) -> Flow {
self.log.lock().expect("log").push(self.label);
if self.stop {
Flow::terminate("stop")
} else {
Flow::cont()
}
}
}
struct ObservesOnly(StepEventKind);
impl AgentHook<M> for ObservesOnly {
async fn on_event(&self, _ctx: &HookContext, _event: StepEvent<'_, M>) -> Flow {
Flow::cont()
}
fn observes(&self, kind: StepEventKind) -> bool {
kind == self.0
}
}
struct Patcher {
label: u32,
log: Arc<Mutex<Vec<u32>>>,
patch: RequestPatch,
}
impl AgentHook<M> for Patcher {
async fn on_event(&self, _ctx: &HookContext, event: StepEvent<'_, M>) -> Flow {
self.log.lock().expect("log").push(self.label);
if matches!(event, StepEvent::CompletionCall { .. }) {
Flow::patch_request(self.patch.clone())
} else {
Flow::cont()
}
}
}
fn tool_call_event() -> StepEvent<'static, M> {
StepEvent::ToolCall {
tool_name: "add",
tool_call_id: Some("tc1"),
internal_call_id: "ic1",
args: "{}",
}
}
fn completion_call_event() -> StepEvent<'static, M> {
static PROMPT: std::sync::OnceLock<crate::message::Message> = std::sync::OnceLock::new();
let prompt = PROMPT.get_or_init(|| crate::message::Message::user("hi"));
StepEvent::CompletionCall {
prompt,
history: &[],
turn: 1,
}
}
#[tokio::test]
async fn runs_hooks_in_registration_order_and_consults_all_on_continue() {
let log = Arc::new(Mutex::new(Vec::new()));
let mut stack = HookStack::<M>::with(Recorder {
label: 1,
log: log.clone(),
stop: false,
});
stack.push(Recorder {
label: 2,
log: log.clone(),
stop: false,
});
let flow = stack.on_event(&ctx(), tool_call_event()).await;
assert!(matches!(flow, Flow::Continue));
assert_eq!(*log.lock().expect("log"), vec![1, 2]);
}
#[tokio::test]
async fn first_terminate_short_circuits_on_chained_tool_call() {
let log = Arc::new(Mutex::new(Vec::new()));
let mut stack = HookStack::<M>::with(Recorder {
label: 1,
log: log.clone(),
stop: true,
});
stack.push(Recorder {
label: 2,
log: log.clone(),
stop: false,
});
let flow = stack.on_event(&ctx(), tool_call_event()).await;
assert!(matches!(flow, Flow::Terminate { .. }));
assert_eq!(
*log.lock().expect("log"),
vec![1],
"a later hook must not run after an earlier hook terminates"
);
}
#[tokio::test]
async fn first_terminate_short_circuits_on_observe_only_events() {
let log = Arc::new(Mutex::new(Vec::new()));
let mut stack = HookStack::<M>::with(Recorder {
label: 1,
log: log.clone(),
stop: true,
});
stack.push(Recorder {
label: 2,
log: log.clone(),
stop: false,
});
let flow = stack
.on_event(
&ctx(),
StepEvent::TextDelta {
delta: "hi",
aggregated: "hi",
},
)
.await;
assert!(matches!(flow, Flow::Terminate { .. }));
assert_eq!(
*log.lock().expect("log"),
vec![1],
"a later hook must not run after an earlier hook terminates an observe-only event"
);
}
#[tokio::test]
async fn completion_call_patches_accumulate_and_consult_every_hook() {
let log = Arc::new(Mutex::new(Vec::new()));
let mut stack = HookStack::<M>::with(Patcher {
label: 1,
log: log.clone(),
patch: RequestPatch::new().temperature(0.1),
});
stack.push(Patcher {
label: 2,
log: log.clone(),
patch: RequestPatch::new().max_tokens(256),
});
let flow = stack.on_event(&ctx(), completion_call_event()).await;
assert_eq!(
*log.lock().expect("log"),
vec![1, 2],
"both hooks must run; a mergeable patch does not short-circuit"
);
match flow {
Flow::PatchRequest { patch } => {
assert_eq!(patch.temperature, Some(0.1));
assert_eq!(patch.max_tokens, Some(256));
}
other => panic!("expected a merged PatchRequest, got {other:?}"),
}
}
#[tokio::test]
async fn completion_call_terminate_short_circuits_and_discards_patch() {
let log = Arc::new(Mutex::new(Vec::new()));
let mut stack = HookStack::<M>::with(Patcher {
label: 1,
log: log.clone(),
patch: RequestPatch::new().temperature(0.1),
});
stack.push(Recorder {
label: 2,
log: log.clone(),
stop: true,
});
stack.push(Patcher {
label: 3,
log: log.clone(),
patch: RequestPatch::new().max_tokens(256),
});
let flow = stack.on_event(&ctx(), completion_call_event()).await;
assert!(matches!(flow, Flow::Terminate { .. }));
assert_eq!(
*log.lock().expect("log"),
vec![1, 2],
"hook 3 must not run after a terminate"
);
}
#[tokio::test]
async fn nested_stack_composes_patches_without_inner_short_circuit() {
let log = Arc::new(Mutex::new(Vec::new()));
let mut inner = HookStack::<M>::with(Patcher {
label: 1,
log: log.clone(),
patch: RequestPatch::new().temperature(0.2),
});
inner.push(Patcher {
label: 2,
log: log.clone(),
patch: RequestPatch::new().max_tokens(128),
});
let mut outer = HookStack::<M>::with(inner);
outer.push(Patcher {
label: 3,
log: log.clone(),
patch: RequestPatch::new().preamble("outer"),
});
let flow = outer.on_event(&ctx(), completion_call_event()).await;
assert_eq!(
*log.lock().expect("log"),
vec![1, 2, 3],
"every hook, including both inner-stack hooks, must run"
);
match flow {
Flow::PatchRequest { patch } => {
assert_eq!(patch.temperature, Some(0.2));
assert_eq!(patch.max_tokens, Some(128));
assert_eq!(patch.preamble.as_deref(), Some("outer"));
}
other => panic!("expected a merged PatchRequest, got {other:?}"),
}
}
#[test]
fn stack_observes_is_the_or_of_its_members() {
let mut stack = HookStack::<M>::with(ObservesOnly(StepEventKind::ToolCall));
stack.push(ObservesOnly(StepEventKind::ToolResult));
assert!(stack.observes(StepEventKind::ToolCall));
assert!(stack.observes(StepEventKind::ToolResult));
assert!(
!stack.observes(StepEventKind::TextDelta),
"no member observes TextDelta, so the stack must not either"
);
}
#[tokio::test]
async fn empty_stack_continues_and_observes_nothing() {
let stack = HookStack::<M>::new();
assert!(stack.is_empty());
assert!(!stack.observes(StepEventKind::ToolCall));
assert!(!stack.observes(StepEventKind::TextDelta));
assert!(matches!(
stack.on_event(&ctx(), tool_call_event()).await,
Flow::Continue
));
}
#[test]
fn unit_hook_observes_no_event_kind() {
let all_kinds = [
StepEventKind::CompletionCall,
StepEventKind::CompletionResponse,
StepEventKind::ModelTurnFinished,
StepEventKind::InvalidToolCall,
StepEventKind::ToolCall,
StepEventKind::ToolResult,
StepEventKind::TextDelta,
StepEventKind::ToolCallDelta,
StepEventKind::StreamResponseFinish,
];
let unit_stack = HookStack::<M>::with(());
for kind in all_kinds {
assert!(
!<() as AgentHook<M>>::observes(&(), kind),
"the `()` no-op hook must not observe {kind:?}"
);
assert!(
!unit_stack.observes(kind),
"a HookStack::with(()) must not observe {kind:?} either"
);
}
}
#[test]
fn merge_appends_extra_context_in_order() {
let doc = |id: &str| crate::completion::Document {
id: id.to_string(),
text: String::new(),
additional_props: Default::default(),
};
let a = RequestPatch::new().context(doc("a"));
let b = RequestPatch::new().context(doc("b"));
let merged = a.merge(b);
let ids: Vec<&str> = merged.extra_context.iter().map(|d| d.id.as_str()).collect();
assert_eq!(ids, vec!["a", "b"]);
}
#[test]
fn merge_shallow_merges_additional_params_later_wins() {
let a = RequestPatch::new().additional_params(serde_json::json!({"x": 1, "y": 2}));
let b = RequestPatch::new().additional_params(serde_json::json!({"y": 3, "z": 4}));
let merged = a.merge(b);
assert_eq!(
merged.additional_params,
Some(serde_json::json!({"x": 1, "y": 3, "z": 4}))
);
}
#[test]
fn merge_scalar_last_writer_wins() {
let a = RequestPatch::new().temperature(0.1);
let b = RequestPatch::new().temperature(0.9);
assert_eq!(a.merge(b).temperature, Some(0.9));
}
#[test]
fn merge_active_tools_intersects() {
let a = RequestPatch::new().active_tools(["search", "add", "sub"]);
let b = RequestPatch::new().active_tools(["add", "sub", "mul"]);
let merged = a.merge(b);
assert_eq!(
merged.active_tools,
Some(vec!["add".to_string(), "sub".to_string()])
);
}
#[test]
fn merge_active_tools_empty_intersection_yields_empty() {
let a = RequestPatch::new().active_tools(["search"]);
let b = RequestPatch::new().active_tools(["add"]);
let merged = a.merge(b);
assert_eq!(merged.active_tools, Some(vec![]));
}
#[test]
fn merge_one_sided_active_tools_keeps_the_present_list() {
let a = RequestPatch::new().active_tools(["search"]);
let b = RequestPatch::new();
assert_eq!(a.merge(b).active_tools, Some(vec!["search".to_string()]));
}
#[test]
fn scratchpad_insert_get_update() {
#[derive(Clone, Default, Debug, PartialEq)]
struct Count(u32);
let pad = Scratchpad::default();
assert_eq!(pad.get::<Count>(), None);
pad.update(|c: &mut Count| c.0 += 1);
pad.update(|c: &mut Count| c.0 += 1);
assert_eq!(pad.get::<Count>(), Some(Count(2)));
assert!(pad.contains::<Count>());
assert_eq!(pad.remove::<Count>(), Some(Count(2)));
assert!(!pad.contains::<Count>());
}
#[test]
fn scratchpad_is_shared_across_clones() {
let pad = Scratchpad::default();
let clone = pad.clone();
pad.insert(7u32);
assert_eq!(clone.get::<u32>(), Some(7));
}
#[test]
fn hook_context_reports_identity_and_turn() {
let ctx = HookContext::new(true, Some("agent".to_string()));
assert!(ctx.is_streaming());
assert_eq!(ctx.agent_name(), Some("agent"));
assert_eq!(ctx.turn(), 0);
ctx.set_turn(3);
assert_eq!(ctx.turn(), 3);
assert!(!ctx.run_id().as_str().is_empty());
}
mod nested_tool_call_resolution {
use super::super::{AgentHook, Flow, HookContext, HookStack, StepEvent};
use super::{M, ctx};
use serde_json::{Value, json};
struct RewriteHook(Value);
impl AgentHook<M> for RewriteHook {
async fn on_event(&self, _ctx: &HookContext, event: StepEvent<'_, M>) -> Flow {
if let StepEvent::ToolCall { .. } = event {
Flow::rewrite_args(self.0.clone())
} else {
Flow::cont()
}
}
}
struct SkipHook;
impl AgentHook<M> for SkipHook {
async fn on_event(&self, _ctx: &HookContext, event: StepEvent<'_, M>) -> Flow {
if let StepEvent::ToolCall { .. } = event {
Flow::skip("denied")
} else {
Flow::cont()
}
}
}
struct TerminateHook;
impl AgentHook<M> for TerminateHook {
async fn on_event(&self, _ctx: &HookContext, event: StepEvent<'_, M>) -> Flow {
if let StepEvent::ToolCall { .. } = event {
Flow::terminate("stop")
} else {
Flow::cont()
}
}
}
struct FailHook;
impl AgentHook<M> for FailHook {
async fn on_event(&self, _ctx: &HookContext, event: StepEvent<'_, M>) -> Flow {
if let StepEvent::ToolCall { .. } = event {
Flow::fail()
} else {
Flow::cont()
}
}
}
#[derive(Clone, Default)]
struct ArgsSpy(std::sync::Arc<std::sync::Mutex<Vec<String>>>);
impl AgentHook<M> for ArgsSpy {
async fn on_event(&self, _ctx: &HookContext, event: StepEvent<'_, M>) -> Flow {
if let StepEvent::ToolCall { args, .. } = event {
self.0.lock().expect("spy").push(args.to_string());
}
Flow::cont()
}
}
async fn resolve(stack: &HookStack<M>) -> (Flow, Option<Value>) {
stack
.resolve_tool_call(&ctx(), "add", Some("tc1"), "ic1", "{}")
.await
}
#[tokio::test]
async fn nested_rewrite_then_skip_preserves_rewrite() {
let mut inner = HookStack::<M>::new();
inner.push(RewriteHook(json!({ "x": 41 })));
inner.push(SkipHook);
let mut outer = HookStack::<M>::new();
outer.push(inner);
let (flow, salvaged) = resolve(&outer).await;
assert!(matches!(flow, Flow::Skip { .. }), "got {flow:?}");
assert_eq!(
salvaged,
Some(json!({ "x": 41 })),
"the inner rewrite must survive the inner skip through a nested stack"
);
}
#[tokio::test]
async fn nested_rewrite_then_terminate_preserves_rewrite() {
let mut inner = HookStack::<M>::new();
inner.push(RewriteHook(json!({ "x": 7 })));
inner.push(TerminateHook);
let mut outer = HookStack::<M>::new();
outer.push(inner);
let (flow, salvaged) = resolve(&outer).await;
assert!(matches!(flow, Flow::Terminate { .. }), "got {flow:?}");
assert_eq!(salvaged, Some(json!({ "x": 7 })));
}
#[tokio::test]
async fn nested_rewrite_then_fail_closed_preserves_rewrite() {
let mut inner = HookStack::<M>::new();
inner.push(RewriteHook(json!({ "x": 9 })));
inner.push(FailHook);
let mut outer = HookStack::<M>::new();
outer.push(inner);
let (flow, salvaged) = resolve(&outer).await;
assert!(matches!(flow, Flow::Fail), "got {flow:?}");
assert_eq!(salvaged, Some(json!({ "x": 9 })));
}
#[tokio::test]
async fn outer_rewrite_then_nested_skip_preserves_outer_rewrite() {
let spy = ArgsSpy::default();
let mut inner = HookStack::<M>::new();
inner.push(spy.clone());
inner.push(SkipHook);
let mut outer = HookStack::<M>::new();
outer.push(RewriteHook(json!({ "x": 1, "y": 2 })));
outer.push(inner);
let (flow, salvaged) = resolve(&outer).await;
assert!(matches!(flow, Flow::Skip { .. }), "got {flow:?}");
assert_eq!(salvaged, Some(json!({ "x": 1, "y": 2 })));
assert_eq!(
spy.0.lock().expect("spy").as_slice(),
[serde_json::to_string(&json!({ "x": 1, "y": 2 })).unwrap()],
);
}
#[tokio::test]
async fn deeply_nested_rewrite_then_skip_preserves_rewrite() {
let mut level3 = HookStack::<M>::new();
level3.push(RewriteHook(json!({ "deep": true })));
level3.push(SkipHook);
let mut level2 = HookStack::<M>::new();
level2.push(level3);
let mut level1 = HookStack::<M>::new();
level1.push(level2);
let (flow, salvaged) = resolve(&level1).await;
assert!(matches!(flow, Flow::Skip { .. }), "got {flow:?}");
assert_eq!(salvaged, Some(json!({ "deep": true })));
}
#[tokio::test]
async fn nested_proceeding_rewrite_surfaces_as_rewrite_args() {
let mut inner = HookStack::<M>::new();
inner.push(RewriteHook(json!({ "x": 5 })));
let mut outer = HookStack::<M>::new();
outer.push(inner);
let (flow, salvaged) = resolve(&outer).await;
assert_eq!(
flow,
Flow::RewriteArgs {
args: json!({ "x": 5 })
}
);
assert_eq!(salvaged, None);
}
}
}