Skip to main content

rig_core/agent/
hook.rs

1//! Hooks for observing and steering an agent run.
2//!
3//! A hook is a single [`AgentHook::on_event`] method that the agent loop calls
4//! at every observable point of a run — before each model call, on each model
5//! response, around every tool call, on streamed deltas, and when the model
6//! emits an invalid tool call. Each call receives a [`HookContext`] (run-scoped
7//! identity + a shared scratchpad) and a [`StepEvent`] describing what is
8//! happening, and returns a [`Flow`] that lets the hook observe, patch the
9//! request, skip a tool, terminate the run early, or (for invalid tool calls)
10//! retry/repair/skip recovery.
11//!
12//! Unlike the old multi-method hook trait, a hook implements one method and
13//! matches on the event it cares about — every other event falls through to the
14//! default [`Flow::Continue`]. Hooks compose in a [`HookStack`] that runs several
15//! hooks in registration order.
16//!
17//! # Composition: mergeable patches vs. terminal control actions
18//!
19//! How a [`HookStack`] combines several hooks' [`Flow`] results depends on the
20//! event, and this is the central behavior to understand:
21//!
22//! - **[`StepEvent::CompletionCall`] — accumulate & merge.** Every hook is
23//!   consulted. A hook that returns [`Flow::PatchRequest`] does **not** stop the
24//!   others: patches from all hooks are merged in registration order into one
25//!   effective patch (see [`RequestPatch`] for the per-field merge rules). This
26//!   lets a RAG hook, a tool-policy hook, and a provider-param hook all
27//!   contribute to the same turn. [`Flow::Terminate`] stops the stack and is
28//!   honored; any other (unsupported) flow stops the stack and fails closed,
29//!   discarding the accumulated patch.
30//! - **[`StepEvent::ToolCall`] / [`StepEvent::ToolResult`] — chain.** Every hook
31//!   is consulted; a [`Flow::RewriteArgs`] / [`Flow::RewriteResult`] does not
32//!   stop the others — the rewritten value is threaded into the next hook's
33//!   event, so hook *N* observes the value as rewritten by hooks *1..N-1* and may
34//!   rewrite further (a redaction hook and a truncation hook compose).
35//!   [`Flow::Skip`] / [`Flow::Terminate`] are terminal mid-chain.
36//! - **Every other event — first non-[`Continue`](Flow::Continue) wins.** These
37//!   are observe-only or recovery events ([`CompletionResponse`](StepEvent::CompletionResponse),
38//!   [`ModelTurnFinished`](StepEvent::ModelTurnFinished),
39//!   [`InvalidToolCall`](StepEvent::InvalidToolCall), the streamed deltas): the
40//!   first hook to return a non-[`Continue`](Flow::Continue) result short-circuits
41//!   the rest.
42//!
43//! **Blind merge.** During accumulation/chaining a hook does *not* see earlier
44//! hooks' contributions in its event payload for `CompletionCall` (it sees the
45//! agent baseline); for `ToolCall`/`ToolResult` it *does* see the running
46//! rewritten value. `CompletionCall` patches are declarative with documented
47//! conflict rules, so blind merge is sufficient and keeps [`StepEvent`] `Copy`.
48//!
49//! **Ordering guidance.** Because [`Flow::Terminate`] short-circuits the stack,
50//! register observe-only hooks (telemetry) *before* steering hooks so a later
51//! terminate cannot hide the run from them. A [`HookStack`] pushed *as a hook*
52//! into another stack composes correctly: it returns its own net flow (a merged
53//! patch, a threaded rewrite, or a terminal action) which the outer stack folds
54//! in again — nesting never reintroduces short-circuiting on mergeable results.
55//!
56//! # Why a returned [`Flow`], not a `next()`-style middleware
57//!
58//! A hook returns a typed [`Flow`] rather than receiving a `next` continuation it
59//! must invoke. A `next()`/middleware model — where each layer has to call
60//! `next(ctx)` to let the rest of the chain *and* the wrapped action run — carries
61//! a well-known footgun: forgetting the call silently disables every downstream
62//! hook and the action itself, with no error. The declarative returned-[`Flow`]
63//! model makes that impossible: proceeding is the explicit [`Flow::Continue`], and
64//! any action an event cannot honor is fail-closed (it terminates the run) rather
65//! than silently skipped.
66//!
67//! Hooks are a *driver* concern: they are async, side-effecting and generic over
68//! the model, so they live in the [`AgentRunner`](crate::agent::AgentRunner)
69//! layer rather than inside the sans-IO, serializable
70//! [`AgentRun`](crate::agent::run::AgentRun) state machine.
71//!
72//! # Migrating from `PromptHook`
73//!
74//! The previous eight-method `PromptHook<M>` trait is replaced by the single
75//! [`AgentHook::on_event`] method. Each old method becomes one match arm on a
76//! [`StepEvent`] variant, and the value it used to return becomes the [`Flow`]
77//! you return from that arm (every event you don't care about falls through to
78//! [`Flow::Continue`]). Every `on_event` now also receives a [`HookContext`]
79//! first argument. Attach one or more hooks with `add_hook`.
80//!
81//! | Old `PromptHook` method | [`StepEvent`] variant | [`Flow`] to return |
82//! |---|---|---|
83//! | `on_completion_call` | [`CompletionCall`](StepEvent::CompletionCall) `{ prompt, history, turn }` | [`cont`](Flow::cont) / [`patch_request`](Flow::patch_request) / [`terminate`](Flow::terminate) |
84//! | `on_completion_response` | [`CompletionResponse`](StepEvent::CompletionResponse) `{ prompt, response }` | [`cont`](Flow::cont) / [`terminate`](Flow::terminate) |
85//! | `on_invalid_tool_call` | [`InvalidToolCall`](StepEvent::InvalidToolCall)`(ctx)` | [`fail`](Flow::fail) (default) / [`retry`](Flow::retry) / [`repair`](Flow::repair) / [`skip`](Flow::skip) / [`terminate`](Flow::terminate) |
86//! | `on_tool_call` | [`ToolCall`](StepEvent::ToolCall) `{ tool_name, tool_call_id, internal_call_id, args }` | [`cont`](Flow::cont) / [`rewrite_args`](Flow::rewrite_args) / [`skip`](Flow::skip) / [`terminate`](Flow::terminate) |
87//! | `on_tool_result` | [`ToolResult`](StepEvent::ToolResult) `{ tool_name, .., result, outcome, extensions }` | [`cont`](Flow::cont) / [`rewrite_result`](Flow::rewrite_result) / [`terminate`](Flow::terminate) |
88//! | `on_text_delta` | [`TextDelta`](StepEvent::TextDelta) `{ delta, aggregated }` | [`cont`](Flow::cont) / [`terminate`](Flow::terminate) |
89//! | `on_tool_call_delta` | [`ToolCallDelta`](StepEvent::ToolCallDelta) `{ tool_call_id, internal_call_id, tool_name, delta }` | [`cont`](Flow::cont) / [`terminate`](Flow::terminate) |
90//! | `on_stream_completion_response_finish` | [`StreamResponseFinish`](StepEvent::StreamResponseFinish) `{ prompt, response }` | [`cont`](Flow::cont) / [`terminate`](Flow::terminate) |
91//! | *(new, both surfaces)* | [`ModelTurnFinished`](StepEvent::ModelTurnFinished) `{ turn, content, usage }` | [`cont`](Flow::cont) / [`terminate`](Flow::terminate) |
92//!
93//! Behavioral notes:
94//!
95//! - The invalid-tool-call default is still fail-fast: returning
96//!   [`Flow::Continue`] for [`StepEvent::InvalidToolCall`] is treated as
97//!   [`Flow::fail`], matching the old trait's default `on_invalid_tool_call`.
98//! - A hook opts out of an event by returning [`Flow::cont`] from that arm,
99//!   instead of leaving a trait method unimplemented.
100//! - For per-delta hooks, override [`AgentHook::observes`] to skip the
101//!   high-frequency [`TextDelta`](StepEvent::TextDelta) /
102//!   [`ToolCallDelta`](StepEvent::ToolCallDelta) events you don't consume.
103//!
104//! # Steering on structured tool outcomes
105//!
106//! [`StepEvent::ToolResult`] carries a structured
107//! [`ToolOutcome`] alongside the model-visible
108//! `result`, so a hook can branch on *why* a tool failed — a timeout vs. a 404 —
109//! without parsing strings. The motivating case: abort after repeated timeouts,
110//! but let a not-found flow back to the model as recoverable feedback.
111//!
112//! ```rust,ignore
113//! use rig_core::agent::{AgentHook, Flow, HookContext, StepEvent};
114//! use rig_core::completion::CompletionModel;
115//! use rig_core::tool::ToolFailureKind;
116//!
117//! #[derive(Clone, Default)]
118//! struct TimeoutCount(usize);
119//!
120//! struct OutcomePolicy;
121//!
122//! impl<M: CompletionModel> AgentHook<M> for OutcomePolicy {
123//!     async fn on_event(&self, ctx: &HookContext, event: StepEvent<'_, M>) -> Flow {
124//!         if let StepEvent::ToolResult { outcome, .. } = event {
125//!             // Repeated timeouts abort the run; a 404 does not.
126//!             if outcome.is_error_kind(ToolFailureKind::Timeout) {
127//!                 let count = ctx.scratchpad().update(|c: &mut TimeoutCount| {
128//!                     c.0 += 1;
129//!                     c.0
130//!                 });
131//!                 if count >= 10 {
132//!                     return Flow::terminate("aborting after repeated tool timeouts");
133//!                 }
134//!             }
135//!             // `NotFound` falls through to `Flow::cont`: the model sees the
136//!             // error text and may try another path.
137//!         }
138//!         Flow::cont()
139//!     }
140//! }
141//! ```
142
143use std::sync::Arc;
144use std::sync::atomic::{AtomicUsize, Ordering};
145
146use crate::{
147    OneOrMany,
148    completion::{CompletionModel, Document, Usage},
149    json_utils,
150    message::{AssistantContent, Message, ToolChoice},
151    tool::{ToolCallExtensions, ToolOutcome, ToolResultExtensions},
152    wasm_compat::{WasmBoxedFuture, WasmCompatSend, WasmCompatSync},
153};
154
155/// Opaque, process-scoped identifier for a single agent run.
156///
157/// Minted once when a run's [`HookContext`] is created and stable for the whole
158/// run, so a hook can correlate every event it observes (across turns, tool
159/// calls and streamed deltas) to one run. It is a short URL-safe string from
160/// Rig's internal, non-cryptographic id generator — not globally unique across
161/// process restarts, and not security-sensitive.
162#[derive(Debug, Clone, PartialEq, Eq, Hash)]
163pub struct RunId(String);
164
165impl RunId {
166    /// Mint a fresh run id.
167    pub(crate) fn generate() -> Self {
168        Self(crate::id::generate())
169    }
170
171    /// The id as a string slice.
172    pub fn as_str(&self) -> &str {
173        &self.0
174    }
175}
176
177impl std::fmt::Display for RunId {
178    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
179        f.write_str(&self.0)
180    }
181}
182
183/// A run-scoped, shared scratchpad passed to every hook via [`HookContext`].
184///
185/// A type-map with interior mutability: cooperating hooks read and write typed
186/// values keyed by type, sharing per-run state (a turn counter, a running
187/// budget, a phase flag) without each rolling its own `Arc<Mutex<…>>`. Every
188/// hook in a run receives the same [`HookContext`] by shared reference, so they
189/// all see one scratchpad (and cloning a `Scratchpad` shares its storage too); a
190/// fresh run starts with an empty one.
191///
192/// Hooks receive `&HookContext` (shared), so every accessor here takes `&self`
193/// and mutates through an internal lock. Reads clone the stored value out (the
194/// lock cannot hand out a borrow), so store cheaply-cloneable values.
195///
196/// # Concurrency
197///
198/// Most events are dispatched sequentially within a run, but at
199/// [`tool_concurrency`](crate::agent::AgentRunner::tool_concurrency)` > 1` the
200/// [`ToolCall`](StepEvent::ToolCall) / [`ToolResult`](StepEvent::ToolResult)
201/// hooks for *different* tools in the same turn may run **concurrently**, all
202/// sharing this one scratchpad. Each accessor ([`insert`](Self::insert),
203/// [`update`](Self::update), …) is race-free *per operation* (it holds the lock
204/// for the whole read-modify-write), but the framework imposes **no
205/// deterministic ordering** across those concurrent tool hooks — the order in
206/// which two tools' hooks touch the scratchpad depends on tool completion
207/// timing. Prefer commutative / idempotent state (a counter, a set union), or
208/// key per-tool state by the tool call id / internal call id, rather than
209/// relying on the order of concurrent updates.
210///
211/// # Example
212/// ```
213/// # use rig_core::agent::hook::Scratchpad;
214/// #[derive(Clone, Default)]
215/// struct Calls(u32);
216///
217/// let pad = Scratchpad::default();
218/// pad.update(|c: &mut Calls| c.0 += 1);
219/// assert_eq!(pad.get::<Calls>().map(|c| c.0), Some(1));
220/// ```
221#[derive(Clone, Default)]
222pub struct Scratchpad {
223    // Reuses the tested `ToolCallExtensions` type-map as the storage, wrapped in
224    // a shared lock so `&HookContext` hooks can mutate it. Under
225    // `tool_concurrency > 1` several tools' `ToolCall`/`ToolResult` hooks may
226    // touch this concurrently, so the lock is load-bearing, not decorative.
227    inner: Arc<std::sync::Mutex<ToolCallExtensions>>,
228}
229
230impl Scratchpad {
231    fn lock(&self) -> std::sync::MutexGuard<'_, ToolCallExtensions> {
232        // A poisoned scratchpad (a hook panicked while holding the lock) should
233        // not cascade into cancelling later hooks or the run; recover the guard.
234        self.inner.lock().unwrap_or_else(|e| e.into_inner())
235    }
236
237    /// Insert a typed value, returning the previous value of the same type.
238    pub fn insert<T: Clone + WasmCompatSend + WasmCompatSync + 'static>(
239        &self,
240        val: T,
241    ) -> Option<T> {
242        self.lock().insert(val)
243    }
244
245    /// Get a clone of the stored value of type `T`, if present.
246    pub fn get<T: Clone + WasmCompatSend + WasmCompatSync + 'static>(&self) -> Option<T> {
247        self.lock().get::<T>().cloned()
248    }
249
250    /// Whether a value of type `T` is present.
251    pub fn contains<T: WasmCompatSend + WasmCompatSync + 'static>(&self) -> bool {
252        self.lock().contains::<T>()
253    }
254
255    /// Remove and return the stored value of type `T`, if present.
256    pub fn remove<T: Clone + WasmCompatSend + WasmCompatSync + 'static>(&self) -> Option<T> {
257        self.lock().remove::<T>()
258    }
259
260    /// Read-modify-write the value of type `T` under one lock acquisition,
261    /// starting from [`Default`] when absent. The value is stored back and the
262    /// closure's return value is returned.
263    ///
264    /// The whole read-modify-write is atomic (no lost updates), but at
265    /// `tool_concurrency > 1` it imposes **no ordering** across concurrent tool
266    /// hooks — see the [type-level concurrency note](Scratchpad#concurrency).
267    /// This is the race-free way to bump a counter or accumulate:
268    /// ```
269    /// # use rig_core::agent::hook::Scratchpad;
270    /// # #[derive(Clone, Default)] struct Total(u64);
271    /// # let pad = Scratchpad::default();
272    /// pad.update(|t: &mut Total| t.0 += 10);
273    /// ```
274    pub fn update<T, R>(&self, f: impl FnOnce(&mut T) -> R) -> R
275    where
276        T: Clone + Default + WasmCompatSend + WasmCompatSync + 'static,
277    {
278        let mut guard = self.lock();
279        let mut val = guard.remove::<T>().unwrap_or_default();
280        let out = f(&mut val);
281        guard.insert(val);
282        out
283    }
284}
285
286impl std::fmt::Debug for Scratchpad {
287    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
288        f.debug_struct("Scratchpad")
289            .field("entries", &self.lock().len())
290            .finish()
291    }
292}
293
294/// Run-scoped context passed by shared reference to every [`AgentHook::on_event`]
295/// call.
296///
297/// Carries the run's identity and a shared [`Scratchpad`]. It is a *driver*
298/// construct built once per run by [`AgentRunner`](crate::agent::AgentRunner);
299/// nothing here reaches the sans-IO [`AgentRun`](crate::agent::run::AgentRun)
300/// state machine. Hooks hold it by `&`, so all fields are read via accessors and
301/// run-scoped mutation goes through [`scratchpad`](Self::scratchpad).
302///
303/// One `HookContext` is shared by every hook invocation in a run. At
304/// [`tool_concurrency`](crate::agent::AgentRunner::tool_concurrency)` > 1` the
305/// [`ToolCall`](StepEvent::ToolCall) / [`ToolResult`](StepEvent::ToolResult)
306/// hooks for different tools in a turn can run concurrently against this shared
307/// context — see the [`Scratchpad` concurrency note](Scratchpad#concurrency) for
308/// how to store run-scoped state safely under that concurrency.
309#[derive(Debug)]
310pub struct HookContext {
311    run_id: RunId,
312    // Interior-mutable so the driver can advance it each turn while hooks hold a
313    // shared `&HookContext`; also the reason the context is `Sync`.
314    turn: AtomicUsize,
315    is_streaming: bool,
316    agent_name: Option<String>,
317    scratchpad: Scratchpad,
318}
319
320impl HookContext {
321    /// Build a fresh run-scoped context. `is_streaming` records which surface is
322    /// driving ([`run`](crate::agent::AgentRunner::run) vs.
323    /// [`stream`](crate::agent::AgentRunner::stream)).
324    pub(crate) fn new(is_streaming: bool, agent_name: Option<String>) -> Self {
325        Self {
326            run_id: RunId::generate(),
327            turn: AtomicUsize::new(0),
328            is_streaming,
329            agent_name,
330            scratchpad: Scratchpad::default(),
331        }
332    }
333
334    /// Record the current one-based model-call index (set by the driver before
335    /// each turn), so events that don't carry a turn still see it.
336    pub(crate) fn set_turn(&self, turn: usize) {
337        self.turn.store(turn, Ordering::Relaxed);
338    }
339
340    /// The run's stable identifier.
341    pub fn run_id(&self) -> &RunId {
342        &self.run_id
343    }
344
345    /// The current one-based model-call index (0 before the first turn).
346    pub fn turn(&self) -> usize {
347        self.turn.load(Ordering::Relaxed)
348    }
349
350    /// Whether this run is driven by the streaming surface.
351    pub fn is_streaming(&self) -> bool {
352        self.is_streaming
353    }
354
355    /// The agent's configured name, if any.
356    pub fn agent_name(&self) -> Option<&str> {
357        self.agent_name.as_deref()
358    }
359
360    /// The run-scoped shared scratchpad.
361    pub fn scratchpad(&self) -> &Scratchpad {
362        &self.scratchpad
363    }
364}
365
366// `&HookContext` is borrowed across `.await` points in async hook dispatch, so
367// on native targets `HookContext` must stay `Sync` (and `Send`). This fails to
368// compile if a future change drops the property.
369#[cfg(not(target_family = "wasm"))]
370const _: fn() = || {
371    fn assert_send_sync<T: Send + Sync>() {}
372    assert_send_sync::<HookContext>();
373};
374
375/// Context passed to a hook on a [`StepEvent::InvalidToolCall`] event when the
376/// model emits a tool call that Rig would reject before normal tool-call
377/// handling or execution.
378#[derive(Debug, Clone)]
379#[non_exhaustive]
380pub struct InvalidToolCallContext {
381    /// Tool name emitted by the model.
382    pub tool_name: String,
383    /// Provider-supplied tool call ID, when available.
384    pub tool_call_id: Option<String>,
385    /// Internal Rig call ID, when available.
386    pub internal_call_id: Option<String>,
387    /// JSON arguments emitted for the tool call, when available.
388    pub args: Option<String>,
389    /// Executable Rig tools advertised to the provider for this turn.
390    pub available_tools: Vec<String>,
391    /// Tools allowed by the active [`ToolChoice`] for this turn.
392    pub allowed_tools: Vec<String>,
393    /// Active tool choice for this turn.
394    pub tool_choice: Option<ToolChoice>,
395    /// Diagnostic chat history including the rejected model output when available.
396    pub chat_history: Vec<Message>,
397    /// Whether the rejected call came from the streaming path.
398    pub is_streaming: bool,
399}
400
401/// Recovery action for an invalid tool call, used internally by
402/// [`AgentRun`](crate::agent::run::AgentRun). Hooks express recovery via
403/// [`Flow`]; the [`AgentRunner`](crate::agent::AgentRunner) translates a `Flow`
404/// returned for a [`StepEvent::InvalidToolCall`] into this type.
405#[derive(Debug, Clone, PartialEq, Eq)]
406pub enum InvalidToolCallHookAction {
407    /// Preserve Rig's default fail-fast behavior.
408    Fail,
409    /// Retry the model turn with corrective feedback.
410    Retry { feedback: String },
411    /// Rewrite only the emitted tool name. The repaired name is revalidated
412    /// against registered tools and the current `ToolChoice` before use.
413    Repair { tool_name: String },
414    /// Treat an invalid structured tool call as skipped by returning synthetic
415    /// feedback as its tool result. This does not execute the invalid tool.
416    Skip { reason: String },
417}
418
419impl InvalidToolCallHookAction {
420    /// Preserve Rig's default fail-fast behavior.
421    pub fn fail() -> Self {
422        Self::Fail
423    }
424
425    /// Retry the model turn with corrective feedback.
426    pub fn retry(feedback: impl Into<String>) -> Self {
427        Self::Retry {
428            feedback: feedback.into(),
429        }
430    }
431
432    /// Repair the emitted tool name.
433    pub fn repair(tool_name: impl Into<String>) -> Self {
434        Self::Repair {
435            tool_name: tool_name.into(),
436        }
437    }
438
439    /// Skip the invalid call with a synthetic tool result.
440    pub fn skip(reason: impl Into<String>) -> Self {
441        Self::Skip {
442            reason: reason.into(),
443        }
444    }
445}
446
447/// An observable point in an agent run, passed to [`AgentHook::on_event`].
448///
449/// `StepEvent` borrows everything it carries (it is `Copy`), so a hook may
450/// inspect the event without taking ownership and a [`HookStack`] can forward
451/// the same event to each hook in turn.
452///
453/// The streaming-only variants ([`TextDelta`](StepEvent::TextDelta),
454/// [`ToolCallDelta`](StepEvent::ToolCallDelta) and
455/// [`StreamResponseFinish`](StepEvent::StreamResponseFinish)) are emitted only
456/// by [`AgentRunner::stream`](crate::agent::AgentRunner::stream).
457#[non_exhaustive]
458pub enum StepEvent<'a, M: CompletionModel> {
459    /// Before a completion request is sent to the model. Honors
460    /// [`Flow::Continue`], [`Flow::PatchRequest`] (patch this turn's request) and
461    /// [`Flow::Terminate`]. Across a [`HookStack`], every hook's
462    /// [`PatchRequest`](Flow::PatchRequest) is merged (see the module docs).
463    CompletionCall {
464        /// The prompt message for this turn.
465        prompt: &'a Message,
466        /// The chat history preceding `prompt`.
467        history: &'a [Message],
468        /// One-based index of this model call within the run.
469        turn: usize,
470    },
471    /// After a non-streaming completion response is received. Suppressed for
472    /// turns recovered by invalid tool-call repair, skip, or retry. Honors
473    /// [`Flow::Continue`] and [`Flow::Terminate`]. The medium-specific
474    /// (non-streaming) counterpart of [`ModelTurnFinished`](Self::ModelTurnFinished),
475    /// carrying the raw provider response.
476    CompletionResponse {
477        /// The prompt message for this turn.
478        prompt: &'a Message,
479        /// The model's completion response.
480        response: &'a crate::completion::CompletionResponse<M::Response>,
481    },
482    /// After a model turn is accepted into the run, on **both** surfaces,
483    /// regardless of whether the turn produced text, tool calls, reasoning, or
484    /// mixed content. This is the normalized, medium-neutral counterpart of
485    /// [`CompletionResponse`](Self::CompletionResponse) (non-streaming) and
486    /// [`StreamResponseFinish`](Self::StreamResponseFinish) (streaming) — use it
487    /// for telemetry that must fire once per turn everywhere, including a
488    /// streamed tool-only turn that fires no `StreamResponseFinish`. Suppressed
489    /// for turns recovered by invalid tool-call repair, skip, or retry, and
490    /// fired *after* the medium-specific raw event when one fires. Observe-only:
491    /// honors [`Flow::Continue`] and [`Flow::Terminate`].
492    ModelTurnFinished {
493        /// One-based index of this model call within the run.
494        turn: usize,
495        /// The model's assistant content for this turn — the canonical committed
496        /// model output. For an ordinary turn this is exactly what is recorded
497        /// into the run. On a structured-output Tool-mode turn that finalizes by
498        /// calling the output tool, this is the model-emitted content **including**
499        /// that output-tool call; the run then persists the turn as assistant text
500        /// (the structured output) with the tool call dropped, so the persisted
501        /// message differs from this content.
502        content: &'a OneOrMany<AssistantContent>,
503        /// Token usage for this turn (zeroed if the provider reported none).
504        usage: Usage,
505    },
506    /// The model emitted a tool call that is unknown or disallowed for this
507    /// turn. Honors [`Flow::Fail`] (the default), [`Flow::Retry`],
508    /// [`Flow::Repair`], [`Flow::Skip`] and [`Flow::Terminate`];
509    /// [`Flow::Continue`] is treated as [`Flow::Fail`].
510    InvalidToolCall(&'a InvalidToolCallContext),
511    /// Before a tool is executed. Honors [`Flow::Continue`],
512    /// [`Flow::RewriteArgs`] (execute the tool with rewritten arguments),
513    /// [`Flow::Skip`] (return `reason` as the tool result without executing) and
514    /// [`Flow::Terminate`]. Across a [`HookStack`], [`RewriteArgs`](Flow::RewriteArgs)
515    /// is chained: `args` reflects prior hooks' rewrites (see the module docs).
516    ToolCall {
517        /// Name of the tool about to be called.
518        tool_name: &'a str,
519        /// Provider-supplied tool call ID, when available.
520        tool_call_id: Option<&'a str>,
521        /// Internal Rig call ID correlating this call's events.
522        internal_call_id: &'a str,
523        /// JSON arguments for the call.
524        args: &'a str,
525    },
526    /// After a tool has produced a result (or a [`ToolCall`](Self::ToolCall) hook
527    /// [skipped](Flow::Skip) it). Honors [`Flow::Continue`],
528    /// [`Flow::RewriteResult`] (substitute the result the model sees) and
529    /// [`Flow::Terminate`].
530    ///
531    /// `result` is the model-visible output, and `outcome` / `extensions` are the
532    /// **structured** execution result — the machine-visible half a hook inspects
533    /// without parsing `result`. `outcome` distinguishes success from a classified
534    /// [`ToolFailure`](crate::tool::ToolFailure) (timeout, not-found, …), a
535    /// [`Skipped`](crate::tool::ToolOutcome::Skipped) call, or a
536    /// [`Denied`](crate::tool::ToolOutcome::Denied) one; `extensions` carries
537    /// provider/application metadata the tool attached that is never sent to the
538    /// model.
539    ///
540    /// For the first hook, `result` is the tool's actual output and `outcome` its
541    /// raw structured outcome; across a [`HookStack`],
542    /// [`RewriteResult`](Flow::RewriteResult) is chained so a later hook sees the
543    /// prior hook's replacement in `result`. A rewrite changes only `result` (the
544    /// model-visible text) — `outcome` and `extensions` are the tool's raw
545    /// structured result throughout, so a redaction hook cannot mask the true
546    /// outcome from a later policy hook (see the module docs).
547    ToolResult {
548        /// Name of the tool that was called.
549        tool_name: &'a str,
550        /// Provider-supplied tool call ID, when available.
551        tool_call_id: Option<&'a str>,
552        /// Internal Rig call ID correlating this call's events.
553        internal_call_id: &'a str,
554        /// JSON arguments for the call.
555        args: &'a str,
556        /// The model-visible tool result. Reflects any earlier hook's
557        /// [`RewriteResult`](Flow::RewriteResult); the first hook sees the tool's
558        /// actual output.
559        result: &'a str,
560        /// The structured outcome of the execution (success / classified error /
561        /// skipped / denied). The raw outcome, unaffected by `RewriteResult`.
562        outcome: &'a ToolOutcome,
563        /// Metadata the tool attached to its result, never sent to the model.
564        extensions: &'a ToolResultExtensions,
565    },
566    /// Streaming only: a text delta was received. `aggregated` is the full text
567    /// accumulated for the turn so far. Honors [`Flow::Continue`] and
568    /// [`Flow::Terminate`].
569    TextDelta {
570        /// The newly received text fragment.
571        delta: &'a str,
572        /// All text accumulated for the turn so far.
573        aggregated: &'a str,
574    },
575    /// Streaming only: a tool-call delta was received. `tool_name` is `Some` on
576    /// the first delta for a tool call and `None` on subsequent deltas. Honors
577    /// [`Flow::Continue`] and [`Flow::Terminate`].
578    ToolCallDelta {
579        /// Provider-supplied tool call ID.
580        tool_call_id: &'a str,
581        /// Internal Rig call ID correlating this call's events.
582        internal_call_id: &'a str,
583        /// Tool name, present on the first delta only.
584        tool_name: Option<&'a str>,
585        /// The newly received argument fragment.
586        delta: &'a str,
587    },
588    /// Streaming only: the provider finished streaming a completion response.
589    /// This is the streaming counterpart of [`CompletionResponse`](Self::CompletionResponse)
590    /// and, like it, is suppressed for turns recovered by invalid tool-call
591    /// repair, skip, or retry. Note one medium-specific difference from
592    /// `CompletionResponse`: it fires only on turns that streamed assistant
593    /// **text** — a turn that emits only a tool call (or only reasoning) does
594    /// not fire it. For a per-turn event that fires on *every* turn on both
595    /// surfaces, use [`ModelTurnFinished`](Self::ModelTurnFinished). Honors
596    /// [`Flow::Continue`] and [`Flow::Terminate`].
597    StreamResponseFinish {
598        /// The prompt message for this turn.
599        prompt: &'a Message,
600        /// The provider's final streaming response.
601        response: &'a M::StreamingResponse,
602    },
603}
604
605// `StepEvent` only holds shared references and `Copy` scalars, so it is `Copy`.
606// These are hand-written to avoid `derive` adding a spurious `M: Clone`/`M: Copy`
607// bound (the generic parameter never appears by value).
608impl<M: CompletionModel> Clone for StepEvent<'_, M> {
609    fn clone(&self) -> Self {
610        *self
611    }
612}
613
614impl<M: CompletionModel> Copy for StepEvent<'_, M> {}
615
616/// The discriminant of a [`StepEvent`].
617///
618/// Passed to [`AgentHook::observes`] so a hook can declare which events it cares
619/// about without the runner building the (sometimes expensive) event payload —
620/// in particular the high-frequency streaming [`TextDelta`](StepEventKind::TextDelta)
621/// and [`ToolCallDelta`](StepEventKind::ToolCallDelta) events.
622#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
623#[non_exhaustive]
624pub enum StepEventKind {
625    /// [`StepEvent::CompletionCall`].
626    CompletionCall,
627    /// [`StepEvent::CompletionResponse`].
628    CompletionResponse,
629    /// [`StepEvent::ModelTurnFinished`].
630    ModelTurnFinished,
631    /// [`StepEvent::InvalidToolCall`].
632    InvalidToolCall,
633    /// [`StepEvent::ToolCall`].
634    ToolCall,
635    /// [`StepEvent::ToolResult`].
636    ToolResult,
637    /// [`StepEvent::TextDelta`].
638    TextDelta,
639    /// [`StepEvent::ToolCallDelta`].
640    ToolCallDelta,
641    /// [`StepEvent::StreamResponseFinish`].
642    StreamResponseFinish,
643}
644
645impl<M: CompletionModel> StepEvent<'_, M> {
646    /// The [`StepEventKind`] discriminant of this event.
647    pub fn kind(&self) -> StepEventKind {
648        match self {
649            StepEvent::CompletionCall { .. } => StepEventKind::CompletionCall,
650            StepEvent::CompletionResponse { .. } => StepEventKind::CompletionResponse,
651            StepEvent::ModelTurnFinished { .. } => StepEventKind::ModelTurnFinished,
652            StepEvent::InvalidToolCall(_) => StepEventKind::InvalidToolCall,
653            StepEvent::ToolCall { .. } => StepEventKind::ToolCall,
654            StepEvent::ToolResult { .. } => StepEventKind::ToolResult,
655            StepEvent::TextDelta { .. } => StepEventKind::TextDelta,
656            StepEvent::ToolCallDelta { .. } => StepEventKind::ToolCallDelta,
657            StepEvent::StreamResponseFinish { .. } => StepEventKind::StreamResponseFinish,
658        }
659    }
660}
661
662/// A partial patch over the model request for a single turn, returned by a hook
663/// via [`Flow::PatchRequest`] on a [`StepEvent::CompletionCall`] event.
664///
665/// Every field is optional: a `Some` value overrides the agent's configured
666/// value for this turn, a `None` value inherits it. The patch is **per-turn and
667/// non-sticky** — it never changes the agent's baseline, so the next turn
668/// re-fires [`CompletionCall`](StepEvent::CompletionCall) and resolves from the
669/// baseline again.
670///
671/// # Merge behavior
672///
673/// Two kinds of merge apply. When several hooks in a [`HookStack`] each return a
674/// patch, they are combined **hook ⊕ hook in registration order** with these
675/// per-field rules; the effective patch is then applied **patch → baseline**.
676///
677/// | Field | hook ⊕ hook (registration order) | patch → baseline |
678/// |---|---|---|
679/// | `extra_context` | append (earlier hooks' docs first) | append after static + dynamic context |
680/// | `additional_params` | shallow-merge top-level keys, later hook wins | shallow-merge onto baseline params |
681/// | `preamble` | last writer wins (warns on conflict) | replaces |
682/// | `temperature`, `max_tokens`, `tool_choice` | last writer wins (warns on conflict) | replaces |
683/// | `active_tools` | set **intersection** (warns when empty) | narrows the advertised set |
684/// | `history` | last writer wins (warns on conflict) | replaces the messages sent this turn |
685///
686/// `active_tools` intersects rather than last-writer-wins because it is an
687/// allow-list guardrail: two narrowing hooks must compose as *narrowing*. All
688/// last-writer-wins conflicts emit a `tracing::warn!` so composition stays
689/// debuggable — additive guidance belongs in `extra_context` documents, not in
690/// preamble concatenation.
691///
692/// Build one with the setters:
693///
694/// ```rust,ignore
695/// Flow::patch_request(
696///     RequestPatch::new()
697///         .tool_choice(ToolChoice::Required)
698///         .active_tools(["search"])
699///         .temperature(0.0),
700/// )
701/// ```
702#[derive(Debug, Clone, Default, PartialEq)]
703#[non_exhaustive]
704pub struct RequestPatch {
705    /// Override the system prompt / preamble for this turn.
706    pub preamble: Option<String>,
707    /// Override the sampling temperature for this turn.
708    pub temperature: Option<f64>,
709    /// Override the max output tokens for this turn.
710    pub max_tokens: Option<u64>,
711    /// Override the tool choice for this turn.
712    pub tool_choice: Option<ToolChoice>,
713    /// Restrict the advertised tools to this allow-list (by name) for this turn.
714    /// `Some(vec![])` advertises no executable tools; `None` keeps the full set.
715    pub active_tools: Option<Vec<String>>,
716    /// Provider-passthrough params shallow-merged onto the agent's for this turn.
717    pub additional_params: Option<serde_json::Value>,
718    /// Extra context documents appended (after static and dynamic context) for
719    /// this turn only. The passive-RAG injection point.
720    pub extra_context: Vec<Document>,
721    /// Replace the prior chat history sent to the provider **this turn only**.
722    /// The persisted transcript and the run state are untouched, and RAG's query
723    /// text still derives from the original prompt/history — this changes only
724    /// what messages are sent. `None` sends the real history. The enabling
725    /// primitive for context-window compaction / summarization middleware.
726    pub history: Option<Vec<Message>>,
727}
728
729/// Last-writer-wins merge for a scalar patch field, warning on a real conflict.
730fn merge_last_wins<T>(earlier: Option<T>, later: Option<T>, field: &str) -> Option<T> {
731    match (earlier, later) {
732        (Some(_), Some(l)) => {
733            tracing::warn!(
734                patch_field = field,
735                "two hooks set `{field}` on the same turn; the later hook wins"
736            );
737            Some(l)
738        }
739        (earlier, later) => later.or(earlier),
740    }
741}
742
743impl RequestPatch {
744    /// An empty patch — a no-op, identical to returning [`Flow::cont`].
745    pub fn new() -> Self {
746        Self::default()
747    }
748
749    /// Override the system prompt / preamble for this turn.
750    pub fn preamble(mut self, preamble: impl Into<String>) -> Self {
751        self.preamble = Some(preamble.into());
752        self
753    }
754
755    /// Override the sampling temperature for this turn.
756    pub fn temperature(mut self, temperature: f64) -> Self {
757        self.temperature = Some(temperature);
758        self
759    }
760
761    /// Override the max output tokens for this turn.
762    pub fn max_tokens(mut self, max_tokens: u64) -> Self {
763        self.max_tokens = Some(max_tokens);
764        self
765    }
766
767    /// Override the tool choice for this turn.
768    ///
769    /// Not every provider honors `tool_choice`: some in-core providers (e.g.
770    /// Ollama, Hyperbolic, Mira, Perplexity) ignore it and log a warning, so
771    /// forcing a tool this way is a no-op there. A choice a provider cannot
772    /// represent (e.g. a multi-name [`ToolChoice::Specific`] on Anthropic, which
773    /// forces a single tool) surfaces as a request error rather than being
774    /// silently downgraded.
775    pub fn tool_choice(mut self, tool_choice: ToolChoice) -> Self {
776        self.tool_choice = Some(tool_choice);
777        self
778    }
779
780    /// Restrict the advertised tools to this allow-list (by name) for this turn.
781    ///
782    /// This narrows the executable tool set, so it composes with `tool_choice`:
783    /// if the effective tool choice is a [`ToolChoice::Specific`] naming a tool
784    /// that `active_tools` filters out (e.g. the agent's baseline choice is
785    /// inherited because this patch didn't set its own), the request fails
786    /// closed with a request error rather than silently forcing a dropped tool.
787    /// When narrowing the set, set a compatible `tool_choice` in the same patch.
788    pub fn active_tools<I, S>(mut self, names: I) -> Self
789    where
790        I: IntoIterator<Item = S>,
791        S: Into<String>,
792    {
793        self.active_tools = Some(names.into_iter().map(Into::into).collect());
794        self
795    }
796
797    /// Shallow-merge these provider-passthrough params onto the agent's for this
798    /// turn.
799    pub fn additional_params(mut self, additional_params: serde_json::Value) -> Self {
800        self.additional_params = Some(additional_params);
801        self
802    }
803
804    /// Append extra context documents for this turn (the passive-RAG injection
805    /// point). Documents are appended after the agent's static and dynamic
806    /// (vector-store) context, in the order added.
807    pub fn extra_context<I>(mut self, docs: I) -> Self
808    where
809        I: IntoIterator<Item = Document>,
810    {
811        self.extra_context.extend(docs);
812        self
813    }
814
815    /// Append a single extra context document for this turn.
816    pub fn context(mut self, doc: Document) -> Self {
817        self.extra_context.push(doc);
818        self
819    }
820
821    /// Replace the prior chat history sent to the provider **this turn only**.
822    /// The persisted transcript is untouched; RAG query text still derives from
823    /// the original history. Use for context-window compaction / summarization.
824    pub fn history<I>(mut self, history: I) -> Self
825    where
826        I: IntoIterator<Item = Message>,
827    {
828        self.history = Some(history.into_iter().collect());
829        self
830    }
831
832    /// Whether this patch has no effect (all fields unset).
833    pub(crate) fn is_empty(&self) -> bool {
834        self.preamble.is_none()
835            && self.temperature.is_none()
836            && self.max_tokens.is_none()
837            && self.tool_choice.is_none()
838            && self.active_tools.is_none()
839            && self.additional_params.is_none()
840            && self.extra_context.is_empty()
841            && self.history.is_none()
842    }
843
844    /// Merge a later hook's patch onto this one (registration order: `self` is
845    /// the accumulated earlier hooks, `later` is the next hook). See the struct
846    /// docs for the per-field rules.
847    pub(crate) fn merge(mut self, later: RequestPatch) -> RequestPatch {
848        // extra_context: append (earlier hooks' documents first).
849        self.extra_context.extend(later.extra_context);
850
851        // additional_params: shallow-merge when both are objects (later wins per
852        // key), otherwise the later non-None value wins wholesale — mirroring the
853        // patch → baseline behavior in request assembly.
854        self.additional_params = match (self.additional_params.take(), later.additional_params) {
855            (Some(base), Some(patch)) if base.is_object() && patch.is_object() => {
856                Some(json_utils::merge(base, patch))
857            }
858            (base, patch) => patch.or(base),
859        };
860
861        // Scalars + preamble + history: last writer wins, warn on real conflict.
862        self.preamble = merge_last_wins(self.preamble, later.preamble, "preamble");
863        self.temperature = merge_last_wins(self.temperature, later.temperature, "temperature");
864        self.max_tokens = merge_last_wins(self.max_tokens, later.max_tokens, "max_tokens");
865        self.tool_choice = merge_last_wins(self.tool_choice, later.tool_choice, "tool_choice");
866        self.history = merge_last_wins(self.history, later.history, "history");
867
868        // active_tools: set intersection (two narrowing guardrails compose as
869        // narrowing). One-sided keeps the present allow-list.
870        self.active_tools = match (self.active_tools.take(), later.active_tools) {
871            (Some(earlier), Some(later)) => {
872                let later_set: std::collections::BTreeSet<&String> = later.iter().collect();
873                let intersection: Vec<String> = earlier
874                    .into_iter()
875                    .filter(|name| later_set.contains(name))
876                    .collect();
877                if intersection.is_empty() {
878                    tracing::warn!(
879                        "two hooks' `active_tools` allow-lists have an empty intersection; \
880                         no executable tools will be advertised this turn"
881                    );
882                }
883                Some(intersection)
884            }
885            (earlier, later) => earlier.or(later),
886        };
887
888        self
889    }
890}
891
892/// Control-flow result returned by [`AgentHook::on_event`].
893///
894/// Each [`StepEvent`] honors a specific subset of variants (documented on each
895/// event). The runner is **fail-closed**: an action an event cannot honor never
896/// silently proceeds — it terminates the run with a diagnostic error. In
897/// particular, a blocking action such as [`Flow::Fail`] returned for a
898/// [`StepEvent::ToolCall`] stops the run rather than letting the tool execute.
899/// Returning [`Flow::Continue`] is always the way to "do nothing".
900///
901/// `Flow` is `PartialEq` but not `Eq`, because [`Flow::PatchRequest`] carries a
902/// [`RequestPatch`] whose `temperature` is an `f64`.
903#[derive(Debug, Clone, PartialEq)]
904#[non_exhaustive]
905pub enum Flow {
906    /// Proceed normally.
907    Continue,
908    /// Terminate the agent run early, surfacing `reason`.
909    Terminate {
910        /// Why the run is being terminated.
911        reason: String,
912    },
913    /// Skip the action: for [`StepEvent::ToolCall`], return `reason` as the tool
914    /// result without executing the tool; for [`StepEvent::InvalidToolCall`],
915    /// record `reason` as a synthetic result for the invalid call.
916    Skip {
917        /// The message returned to the model in place of the tool result. It is
918        /// delivered verbatim, so it doubles as a prompt: state that the tool did
919        /// not run and, unless you want the model to try again, tell it not to
920        /// retry — a bare `"denied"` often makes the model re-emit the same call.
921        reason: String,
922    },
923    /// [`StepEvent::ToolCall`] only: rewrite the tool-call arguments, then
924    /// execute the tool with the replacement. This is the steering action for
925    /// guardrails that normalize, clamp, redirect, or inject scoped parameters
926    /// before a tool runs.
927    ///
928    /// The rewritten arguments are what the tool is invoked with, what the
929    /// following [`StepEvent::ToolResult`] reports, and what the
930    /// `gen_ai.tool.call.arguments` span field records.
931    ///
932    /// This rewrites only what the tool *executes against*, not the model's
933    /// transcript: the assistant message that recorded the original tool call is
934    /// unchanged and keeps the model's original arguments. It is therefore an
935    /// execution-args rewrite (inject defaults, clamp a range, redirect a path),
936    /// **not** a history redactor — it does not scrub a value the model already
937    /// emitted from the conversation.
938    ///
939    /// Across a [`HookStack`], rewrites **chain**: the rewritten arguments are
940    /// threaded into the next hook's [`ToolCall`](StepEvent::ToolCall) event, so
941    /// several hooks can each refine the arguments in registration order.
942    RewriteArgs {
943        /// The JSON arguments the tool is invoked with, in place of the ones the
944        /// model emitted.
945        args: serde_json::Value,
946    },
947    /// [`StepEvent::ToolResult`] only: replace the tool's result with this string
948    /// before the model sees it. The post-execution counterpart of
949    /// [`RewriteArgs`](Flow::RewriteArgs) — for guardrails that redact, truncate,
950    /// or normalize a tool's output.
951    ///
952    /// The replacement is what the model receives as the tool result and what the
953    /// `gen_ai.tool.call.result` span field records. As with
954    /// [`RewriteArgs`](Flow::RewriteArgs), this changes only what the model
955    /// *sees*: the tool still ran and produced its real output (which the first
956    /// hook's [`ToolResult`](StepEvent::ToolResult) event observed before this
957    /// replacement is applied). It does not scrub the tool's output from logs.
958    ///
959    /// The replacement is delivered to the model verbatim — it is not re-parsed
960    /// as structured/multimodal tool output, so a JSON-shaped replacement reaches
961    /// the model as literal text.
962    ///
963    /// Across a [`HookStack`], rewrites **chain**: the replacement is threaded
964    /// into the next hook's [`ToolResult`](StepEvent::ToolResult) event, so a
965    /// redaction hook and a truncation hook can compose in registration order.
966    RewriteResult {
967        /// The result delivered to the model in place of the tool's actual
968        /// output.
969        result: String,
970    },
971    /// [`StepEvent::CompletionCall`] only: patch fields of the model request for
972    /// this turn before it is sent. The per-turn request-steering action — for
973    /// hooks that adjust the system prompt, sampling, tool choice, the advertised
974    /// tool set, or inject context documents from run state (force a tool on the
975    /// first turn, lower the temperature on a critical step, add RAG context).
976    ///
977    /// The patch is partial ([`RequestPatch`]): each set field replaces (or, for
978    /// `additional_params`/`extra_context`, merges onto) the agent's configured
979    /// value; unset fields are inherited. It applies to *this turn only* and does
980    /// not change the agent's baseline — the next turn re-fires
981    /// [`CompletionCall`](StepEvent::CompletionCall) and re-resolves from it.
982    ///
983    /// Across a [`HookStack`], patches from all hooks **accumulate** and merge in
984    /// registration order (see [`RequestPatch`] and the module docs); this action
985    /// therefore does *not* short-circuit later hooks.
986    PatchRequest {
987        /// The partial request patch applied to this turn.
988        patch: RequestPatch,
989    },
990    /// [`StepEvent::InvalidToolCall`] only: fail the run fast (the default for
991    /// invalid tool calls).
992    Fail,
993    /// [`StepEvent::InvalidToolCall`] only: retry the model turn with corrective
994    /// feedback.
995    Retry {
996        /// Feedback appended to the conversation before re-prompting.
997        feedback: String,
998    },
999    /// [`StepEvent::InvalidToolCall`] only: rewrite the emitted tool name, which
1000    /// is then revalidated against the allowed tools.
1001    Repair {
1002        /// The corrected tool name.
1003        tool_name: String,
1004    },
1005}
1006
1007impl Flow {
1008    /// Continue the agent loop as normal.
1009    pub fn cont() -> Self {
1010        Self::Continue
1011    }
1012
1013    /// Terminate the agent run early with a reason.
1014    pub fn terminate(reason: impl Into<String>) -> Self {
1015        Self::Terminate {
1016            reason: reason.into(),
1017        }
1018    }
1019
1020    /// Skip the current tool call (or invalid call) with the provided reason.
1021    ///
1022    /// `reason` is delivered to the model verbatim as the tool result, so it
1023    /// doubles as a prompt — tell the model the tool did not run and whether to
1024    /// retry, or it may re-emit the identical call:
1025    ///
1026    /// ```rust,ignore
1027    /// Flow::skip("Not executed (denied by policy). Do not retry unless the user asks.")
1028    /// ```
1029    pub fn skip(reason: impl Into<String>) -> Self {
1030        Self::Skip {
1031            reason: reason.into(),
1032        }
1033    }
1034
1035    /// Rewrite a tool call's arguments, then execute the tool with the
1036    /// replacement (tool calls only).
1037    ///
1038    /// Accepts anything convertible into a [`serde_json::Value`] — most often
1039    /// the [`serde_json::json!`] macro or a value built from the parsed original
1040    /// arguments. To rewrite from a typed value instead, use
1041    /// [`try_rewrite_args`](Flow::try_rewrite_args).
1042    ///
1043    /// ```rust,ignore
1044    /// // Inject a scoped parameter the model never sees, leaving the rest intact.
1045    /// let mut args: serde_json::Value = serde_json::from_str(emitted_args)?;
1046    /// args["account_id"] = serde_json::json!(session.account_id);
1047    /// Flow::rewrite_args(args)
1048    /// ```
1049    pub fn rewrite_args(args: impl Into<serde_json::Value>) -> Self {
1050        Self::RewriteArgs { args: args.into() }
1051    }
1052
1053    /// Rewrite a tool call's arguments from a serializable value (tool calls
1054    /// only), serializing it to JSON.
1055    ///
1056    /// This is the typed convenience over [`rewrite_args`](Flow::rewrite_args)
1057    /// for callers that hold a Rust args struct. It only fails if the value
1058    /// cannot be serialized to JSON; a hook typically maps that error to
1059    /// [`Flow::terminate`]:
1060    ///
1061    /// ```rust,ignore
1062    /// Flow::try_rewrite_args(&new_args).unwrap_or_else(|e| Flow::terminate(e.to_string()))
1063    /// ```
1064    pub fn try_rewrite_args<T: serde::Serialize>(value: &T) -> Result<Self, serde_json::Error> {
1065        Ok(Self::RewriteArgs {
1066            args: serde_json::to_value(value)?,
1067        })
1068    }
1069
1070    /// Replace a tool's result with `result` before the model sees it (tool
1071    /// results only).
1072    ///
1073    /// The post-execution counterpart of [`rewrite_args`](Flow::rewrite_args),
1074    /// for guardrails that redact, truncate, or normalize a tool's output:
1075    ///
1076    /// ```rust,ignore
1077    /// // Redact a secret from the tool output before it reaches the model.
1078    /// Flow::rewrite_result(redact(tool_output))
1079    /// ```
1080    pub fn rewrite_result(result: impl Into<String>) -> Self {
1081        Self::RewriteResult {
1082            result: result.into(),
1083        }
1084    }
1085
1086    /// Patch fields of the model request for this turn (completion calls only).
1087    /// See [`RequestPatch`] for the partial-patch, per-turn, mergeable semantics.
1088    pub fn patch_request(patch: RequestPatch) -> Self {
1089        Self::PatchRequest { patch }
1090    }
1091
1092    /// Fail fast on an invalid tool call (the default).
1093    pub fn fail() -> Self {
1094        Self::Fail
1095    }
1096
1097    /// Retry the model turn with corrective feedback (invalid tool calls only).
1098    ///
1099    /// A common recovery is to let the model self-correct by naming the valid
1100    /// tools, built from the diagnostics in [`InvalidToolCallContext`]:
1101    ///
1102    /// ```rust,ignore
1103    /// // On the `StepEvent::InvalidToolCall(ctx)` arm of `on_event`:
1104    /// Flow::retry(format!(
1105    ///     "`{}` is not a valid tool. Call one of: [{}].",
1106    ///     ctx.tool_name,
1107    ///     ctx.available_tools.join(", "),
1108    /// ))
1109    /// ```
1110    ///
1111    /// Without such a hook the invalid-call default stays fail-closed
1112    /// ([`Flow::Continue`] is treated as [`Flow::fail`]).
1113    pub fn retry(feedback: impl Into<String>) -> Self {
1114        Self::Retry {
1115            feedback: feedback.into(),
1116        }
1117    }
1118
1119    /// Repair the emitted tool name (invalid tool calls only).
1120    pub fn repair(tool_name: impl Into<String>) -> Self {
1121        Self::Repair {
1122            tool_name: tool_name.into(),
1123        }
1124    }
1125}
1126
1127/// A per-run hook that observes and steers an agent run.
1128///
1129/// Implement [`on_event`](AgentHook::on_event) and match on the [`StepEvent`]
1130/// variants you care about; every other event falls through to the default
1131/// [`Flow::Continue`]. Hooks must be cheap to share (`Clone` is not required —
1132/// hooks are held behind an `Arc` once registered).
1133///
1134/// # Example
1135/// ```rust,ignore
1136/// use rig_core::agent::{AgentHook, Flow, HookContext, StepEvent};
1137/// use rig_core::completion::CompletionModel;
1138///
1139/// #[derive(Clone)]
1140/// struct Logger;
1141///
1142/// impl<M: CompletionModel> AgentHook<M> for Logger {
1143///     async fn on_event(&self, ctx: &HookContext, event: StepEvent<'_, M>) -> Flow {
1144///         if let StepEvent::ToolCall { tool_name, args, .. } = event {
1145///             println!("[run {}] calling {tool_name}({args})", ctx.run_id());
1146///         }
1147///         Flow::cont()
1148///     }
1149/// }
1150/// ```
1151pub trait AgentHook<M>: WasmCompatSend + WasmCompatSync
1152where
1153    M: CompletionModel,
1154{
1155    /// Called at every observable point of the agent run (subject to
1156    /// [`observes`](Self::observes)). Receives the run-scoped [`HookContext`] and
1157    /// the [`StepEvent`]. The default implementation is a no-op: it ignores every
1158    /// event and returns [`Flow::Continue`]. It does **not** narrow
1159    /// [`observes`](Self::observes) (which defaults to `true`), so a hook that
1160    /// takes this default is still dispatched every event — override `observes`
1161    /// to skip the high-frequency delta events. (The `()` no-op hook overrides
1162    /// `observes` to `false`, so the runner skips dispatching those delta events
1163    /// to it; it still receives, and returns [`Flow::Continue`] for, every other
1164    /// event.)
1165    fn on_event(
1166        &self,
1167        ctx: &HookContext,
1168        event: StepEvent<'_, M>,
1169    ) -> impl Future<Output = Flow> + WasmCompatSend {
1170        let _ = (ctx, event);
1171        async { Flow::Continue }
1172    }
1173
1174    /// Resolve a [`ToolCall`](StepEvent::ToolCall) for this hook, returning its
1175    /// [`Flow`] plus any tool-argument rewrite that must be **salvaged** when the
1176    /// hook short-circuits — so a nested [`HookStack`] never loses an inner
1177    /// [`Flow::RewriteArgs`] behind a later inner
1178    /// [`Flow::Skip`]/[`Flow::Terminate`].
1179    ///
1180    /// The default — correct for any leaf hook — dispatches the `ToolCall` event
1181    /// to [`on_event`](Self::on_event) and reports **no** salvaged rewrite: a
1182    /// single hook returns exactly one [`Flow`], so it can either rewrite (via
1183    /// [`Flow::RewriteArgs`]) or short-circuit, never both. [`HookStack`]
1184    /// overrides this to compose its members' resolutions, preserving an inner
1185    /// rewrite across a short-circuit. This is an internal composition hook;
1186    /// implementing it is only necessary for a custom composite hook that wraps
1187    /// other hooks and needs the same rewrite-preserving behavior.
1188    #[doc(hidden)]
1189    fn resolve_tool_call(
1190        &self,
1191        ctx: &HookContext,
1192        tool_name: &str,
1193        tool_call_id: Option<&str>,
1194        internal_call_id: &str,
1195        args: &str,
1196    ) -> impl Future<Output = (Flow, Option<serde_json::Value>)> + WasmCompatSend {
1197        async move {
1198            let flow = self
1199                .on_event(
1200                    ctx,
1201                    StepEvent::ToolCall {
1202                        tool_name,
1203                        tool_call_id,
1204                        internal_call_id,
1205                        args,
1206                    },
1207                )
1208                .await;
1209            (flow, None)
1210        }
1211    }
1212
1213    /// Whether this hook observes events of the given [`StepEventKind`].
1214    ///
1215    /// This is a **performance hint for the high-frequency streaming
1216    /// [`TextDelta`](StepEventKind::TextDelta) /
1217    /// [`ToolCallDelta`](StepEventKind::ToolCallDelta) events**, which otherwise
1218    /// cost one boxed future per delta. The runner skips building and
1219    /// dispatching a delta event only when *no* hook in the stack observes it
1220    /// (interest is OR-combined across the stack), so a hook may still be
1221    /// invoked for a delta a sibling observes — `on_event` must therefore stay
1222    /// total (return [`Flow::Continue`] for events it ignores) rather than
1223    /// assume it is only called for observed kinds.
1224    ///
1225    /// Control flow is **never** changed by `observes`: the shared, steering
1226    /// events ([`ToolCall`](StepEventKind::ToolCall),
1227    /// [`InvalidToolCall`](StepEventKind::InvalidToolCall), …) fire identically
1228    /// regardless of this method, so `run()` and `stream()` stay in lock-step.
1229    /// The default observes everything.
1230    fn observes(&self, kind: StepEventKind) -> bool {
1231        let _ = kind;
1232        true
1233    }
1234}
1235
1236/// The no-op hook: observes nothing, never alters control flow.
1237impl<M> AgentHook<M> for ()
1238where
1239    M: CompletionModel,
1240{
1241    /// Observe nothing, so the runner skips building/dispatching the
1242    /// high-frequency streaming delta events (`TextDelta` / `ToolCallDelta`,
1243    /// the only events gated on `observes`) for a `()` hook.
1244    fn observes(&self, _kind: StepEventKind) -> bool {
1245        false
1246    }
1247}
1248
1249/// Object-safe shim over [`AgentHook`] so a [`HookStack`] can hold a
1250/// heterogeneous list of hooks behind `Arc`.
1251trait DynAgentHook<M>: WasmCompatSend + WasmCompatSync
1252where
1253    M: CompletionModel,
1254{
1255    fn on_event_boxed<'a>(
1256        &'a self,
1257        ctx: &'a HookContext,
1258        event: StepEvent<'a, M>,
1259    ) -> WasmBoxedFuture<'a, Flow>
1260    where
1261        M: 'a;
1262
1263    /// Object-safe [`AgentHook::resolve_tool_call`]. Preserves an inner
1264    /// [`Flow::RewriteArgs`] across a short-circuit so nested [`HookStack`]s
1265    /// compose correctly.
1266    fn resolve_tool_call_boxed<'a>(
1267        &'a self,
1268        ctx: &'a HookContext,
1269        tool_name: &'a str,
1270        tool_call_id: Option<&'a str>,
1271        internal_call_id: &'a str,
1272        args: &'a str,
1273    ) -> WasmBoxedFuture<'a, (Flow, Option<serde_json::Value>)>
1274    where
1275        M: 'a;
1276
1277    fn observes_dyn(&self, kind: StepEventKind) -> bool;
1278}
1279
1280impl<M, H> DynAgentHook<M> for H
1281where
1282    M: CompletionModel,
1283    H: AgentHook<M>,
1284{
1285    fn on_event_boxed<'a>(
1286        &'a self,
1287        ctx: &'a HookContext,
1288        event: StepEvent<'a, M>,
1289    ) -> WasmBoxedFuture<'a, Flow>
1290    where
1291        M: 'a,
1292    {
1293        Box::pin(self.on_event(ctx, event))
1294    }
1295
1296    fn resolve_tool_call_boxed<'a>(
1297        &'a self,
1298        ctx: &'a HookContext,
1299        tool_name: &'a str,
1300        tool_call_id: Option<&'a str>,
1301        internal_call_id: &'a str,
1302        args: &'a str,
1303    ) -> WasmBoxedFuture<'a, (Flow, Option<serde_json::Value>)>
1304    where
1305        M: 'a,
1306    {
1307        Box::pin(self.resolve_tool_call(ctx, tool_name, tool_call_id, internal_call_id, args))
1308    }
1309
1310    fn observes_dyn(&self, kind: StepEventKind) -> bool {
1311        self.observes(kind)
1312    }
1313}
1314
1315/// An ordered list of hooks run as one hook.
1316///
1317/// Each hook is consulted in registration order. How their [`Flow`] results
1318/// combine depends on the event (see the [module docs](self)):
1319/// [`CompletionCall`](StepEvent::CompletionCall) patches **accumulate**;
1320/// [`ToolCall`](StepEvent::ToolCall) / [`ToolResult`](StepEvent::ToolResult)
1321/// rewrites **chain**; every other event uses **first non-[`Continue`](Flow::Continue)
1322/// wins**. Because the runner is fail-closed, a non-`Continue` action always
1323/// takes effect or terminates the run — it is never silently ignored. An empty
1324/// stack is the no-op hook and [`observes`](HookStack::observes) nothing, so the
1325/// runner skips event dispatch for it entirely.
1326///
1327/// This is the default hook type carried by an
1328/// [`Agent`](crate::agent::Agent) and an
1329/// [`AgentRunner`](crate::agent::AgentRunner); build one with
1330/// [`add_hook`](crate::agent::AgentRunner::add_hook).
1331pub struct HookStack<M>
1332where
1333    M: CompletionModel,
1334{
1335    hooks: Vec<Arc<dyn DynAgentHook<M>>>,
1336}
1337
1338// Hand-written so the impls do not require `M: Clone`/`M: Default`: `M` only
1339// appears inside `Arc<dyn DynAgentHook<M>>`, never by value.
1340impl<M> Clone for HookStack<M>
1341where
1342    M: CompletionModel,
1343{
1344    fn clone(&self) -> Self {
1345        Self {
1346            hooks: self.hooks.clone(),
1347        }
1348    }
1349}
1350
1351impl<M> Default for HookStack<M>
1352where
1353    M: CompletionModel,
1354{
1355    fn default() -> Self {
1356        Self { hooks: Vec::new() }
1357    }
1358}
1359
1360impl<M> std::fmt::Debug for HookStack<M>
1361where
1362    M: CompletionModel,
1363{
1364    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1365        f.debug_struct("HookStack")
1366            .field("len", &self.hooks.len())
1367            .finish()
1368    }
1369}
1370
1371impl<M> HookStack<M>
1372where
1373    M: CompletionModel,
1374{
1375    /// An empty stack (the no-op hook).
1376    pub fn new() -> Self {
1377        Self::default()
1378    }
1379
1380    /// A stack containing a single hook.
1381    pub fn with<H>(hook: H) -> Self
1382    where
1383        H: AgentHook<M> + 'static,
1384    {
1385        let mut stack = Self::new();
1386        stack.push(hook);
1387        stack
1388    }
1389
1390    /// Append a hook to the end of the stack.
1391    pub fn push<H>(&mut self, hook: H)
1392    where
1393        H: AgentHook<M> + 'static,
1394    {
1395        self.hooks.push(Arc::new(hook));
1396    }
1397
1398    /// Whether the stack contains no hooks.
1399    pub fn is_empty(&self) -> bool {
1400        self.hooks.is_empty()
1401    }
1402
1403    /// Number of hooks in the stack.
1404    pub fn len(&self) -> usize {
1405        self.hooks.len()
1406    }
1407}
1408
1409impl<M> AgentHook<M> for HookStack<M>
1410where
1411    M: CompletionModel,
1412{
1413    /// Compose the stack's members' [`ToolCall`](StepEvent::ToolCall)
1414    /// resolutions, threading tool-arg rewrites through the chain **and**
1415    /// preserving them across a short-circuit — including for a member that is
1416    /// itself a [`HookStack`], which is why members are consulted via
1417    /// [`resolve_tool_call`](AgentHook::resolve_tool_call) rather than
1418    /// [`on_event`](AgentHook::on_event) (the latter can only return a single
1419    /// [`Flow`], losing an inner rewrite behind an inner `Skip`/`Terminate`).
1420    ///
1421    /// When the chain proceeds, any rewrite is carried by the returned [`Flow`]
1422    /// itself ([`RewriteArgs`](Flow::RewriteArgs) for a rewriting chain,
1423    /// [`Continue`](Flow::Continue) otherwise) and the second element is `None`.
1424    /// When a member short-circuits with [`Flow::Skip`] / [`Flow::Terminate`] (or
1425    /// a fail-closed action), that action is returned in the first element while
1426    /// the accumulated rewrite is salvaged into the second element, so the caller
1427    /// (`run_single_tool`) can still report the rewritten args on the resulting
1428    /// [`ToolResult`](StepEvent::ToolResult) event and in tracing rather than
1429    /// leaking the model's original (pre-rewrite) args. The two are therefore
1430    /// mutually exclusive: the [`Flow`] is [`RewriteArgs`](Flow::RewriteArgs) only
1431    /// when the second element is `None`.
1432    async fn resolve_tool_call(
1433        &self,
1434        ctx: &HookContext,
1435        tool_name: &str,
1436        tool_call_id: Option<&str>,
1437        internal_call_id: &str,
1438        args: &str,
1439    ) -> (Flow, Option<serde_json::Value>) {
1440        let mut effective: Option<serde_json::Value> = None;
1441        for hook in &self.hooks {
1442            let rewritten = effective.as_ref().map(json_utils::value_to_json_string);
1443            let args_for_hook = rewritten.as_deref().unwrap_or(args);
1444            let (flow, salvaged) = hook
1445                .resolve_tool_call_boxed(
1446                    ctx,
1447                    tool_name,
1448                    tool_call_id,
1449                    internal_call_id,
1450                    args_for_hook,
1451                )
1452                .await;
1453            // A member (e.g. a nested `HookStack`) may have rewritten the args
1454            // before short-circuiting; adopt that rewrite so it is not lost.
1455            if let Some(rewrite) = salvaged {
1456                effective = Some(rewrite);
1457            }
1458            match flow {
1459                Flow::Continue => {}
1460                Flow::RewriteArgs { args } => effective = Some(args),
1461                // A short-circuit drops the accumulated rewrite from the returned
1462                // flow, so salvage it in the second element for the caller.
1463                other => return (other, effective),
1464            }
1465        }
1466        // The chain proceeded: surface any rewrite through the flow itself.
1467        match effective {
1468            Some(args) => (Flow::RewriteArgs { args }, None),
1469            None => (Flow::Continue, None),
1470        }
1471    }
1472
1473    async fn on_event(&self, ctx: &HookContext, event: StepEvent<'_, M>) -> Flow {
1474        match event {
1475            // Accumulate mergeable request patches from every hook (registration
1476            // order); short-circuit on `Terminate` or any flow the event cannot
1477            // honor (fail-closed downstream, discarding the accumulated patch).
1478            // Hooks see the agent baseline, not earlier hooks' patches (blind
1479            // merge), which keeps `StepEvent` `Copy`.
1480            StepEvent::CompletionCall { .. } => {
1481                let mut merged: Option<RequestPatch> = None;
1482                for hook in &self.hooks {
1483                    match hook.on_event_boxed(ctx, event).await {
1484                        Flow::Continue => {}
1485                        Flow::PatchRequest { patch } => {
1486                            merged = Some(match merged {
1487                                Some(acc) => acc.merge(patch),
1488                                None => patch,
1489                            });
1490                        }
1491                        other => return other,
1492                    }
1493                }
1494                match merged {
1495                    Some(patch) if !patch.is_empty() => Flow::PatchRequest { patch },
1496                    _ => Flow::Continue,
1497                }
1498            }
1499            // Chain tool-arg rewrites: thread the effective arguments through
1500            // each hook so a later hook observes (and may further rewrite) the
1501            // value produced by earlier hooks. A proceeding chain surfaces the
1502            // rewrite as `RewriteArgs`; `Skip`/`Terminate` are terminal and any
1503            // other flow is returned for fail-closed handling. The salvaged
1504            // rewrite (second element) matters only to `run_single_tool`, which
1505            // must report it on a short-circuited `ToolResult`; it is dropped
1506            // here (this result is observe-only).
1507            StepEvent::ToolCall {
1508                tool_name,
1509                tool_call_id,
1510                internal_call_id,
1511                args,
1512            } => {
1513                self.resolve_tool_call(ctx, tool_name, tool_call_id, internal_call_id, args)
1514                    .await
1515                    .0
1516            }
1517            // Chain tool-result rewrites: thread the effective (model-visible)
1518            // result through each hook (the first hook sees the tool's real
1519            // output). The structured `outcome`/`extensions` are the tool's raw
1520            // result and are passed unchanged to every hook — a rewrite alters
1521            // only the model-visible text, never the outcome a later policy sees.
1522            StepEvent::ToolResult {
1523                tool_name,
1524                tool_call_id,
1525                internal_call_id,
1526                args,
1527                result,
1528                outcome,
1529                extensions,
1530            } => {
1531                let mut effective: Option<String> = None;
1532                for hook in &self.hooks {
1533                    let result_for_hook = effective.as_deref().unwrap_or(result);
1534                    let per_hook = StepEvent::ToolResult {
1535                        tool_name,
1536                        tool_call_id,
1537                        internal_call_id,
1538                        args,
1539                        result: result_for_hook,
1540                        outcome,
1541                        extensions,
1542                    };
1543                    match hook.on_event_boxed(ctx, per_hook).await {
1544                        Flow::Continue => {}
1545                        Flow::RewriteResult { result } => effective = Some(result),
1546                        other => return other,
1547                    }
1548                }
1549                match effective {
1550                    Some(result) => Flow::RewriteResult { result },
1551                    None => Flow::Continue,
1552                }
1553            }
1554            // Observe-only / recovery events: first non-`Continue` wins.
1555            _ => {
1556                for hook in &self.hooks {
1557                    match hook.on_event_boxed(ctx, event).await {
1558                        Flow::Continue => {}
1559                        other => return other,
1560                    }
1561                }
1562                Flow::Continue
1563            }
1564        }
1565    }
1566
1567    /// The stack observes an event kind if any of its hooks does (so an empty
1568    /// stack observes nothing).
1569    fn observes(&self, kind: StepEventKind) -> bool {
1570        self.hooks.iter().any(|hook| hook.observes_dyn(kind))
1571    }
1572}
1573
1574#[cfg(test)]
1575mod tests {
1576    use std::sync::{Arc, Mutex};
1577
1578    use super::{
1579        AgentHook, Flow, HookContext, HookStack, RequestPatch, Scratchpad, StepEvent, StepEventKind,
1580    };
1581    use crate::test_utils::MockCompletionModel;
1582
1583    type M = MockCompletionModel;
1584
1585    fn ctx() -> HookContext {
1586        HookContext::new(false, Some("test-agent".to_string()))
1587    }
1588
1589    /// Pushes its label when invoked and returns `Continue` or `Terminate`.
1590    struct Recorder {
1591        label: u32,
1592        log: Arc<Mutex<Vec<u32>>>,
1593        stop: bool,
1594    }
1595
1596    impl AgentHook<M> for Recorder {
1597        async fn on_event(&self, _ctx: &HookContext, _event: StepEvent<'_, M>) -> Flow {
1598            self.log.lock().expect("log").push(self.label);
1599            if self.stop {
1600                Flow::terminate("stop")
1601            } else {
1602                Flow::cont()
1603            }
1604        }
1605    }
1606
1607    /// Observes exactly one event kind (used to probe stack-level `observes`).
1608    struct ObservesOnly(StepEventKind);
1609
1610    impl AgentHook<M> for ObservesOnly {
1611        async fn on_event(&self, _ctx: &HookContext, _event: StepEvent<'_, M>) -> Flow {
1612            Flow::cont()
1613        }
1614
1615        fn observes(&self, kind: StepEventKind) -> bool {
1616            kind == self.0
1617        }
1618    }
1619
1620    /// A hook that returns a fixed patch on `CompletionCall`, and records its
1621    /// label so we can prove every hook ran.
1622    struct Patcher {
1623        label: u32,
1624        log: Arc<Mutex<Vec<u32>>>,
1625        patch: RequestPatch,
1626    }
1627
1628    impl AgentHook<M> for Patcher {
1629        async fn on_event(&self, _ctx: &HookContext, event: StepEvent<'_, M>) -> Flow {
1630            self.log.lock().expect("log").push(self.label);
1631            if matches!(event, StepEvent::CompletionCall { .. }) {
1632                Flow::patch_request(self.patch.clone())
1633            } else {
1634                Flow::cont()
1635            }
1636        }
1637    }
1638
1639    /// A cheap, M-agnostic event to dispatch (no model response required).
1640    fn tool_call_event() -> StepEvent<'static, M> {
1641        StepEvent::ToolCall {
1642            tool_name: "add",
1643            tool_call_id: Some("tc1"),
1644            internal_call_id: "ic1",
1645            args: "{}",
1646        }
1647    }
1648
1649    fn completion_call_event() -> StepEvent<'static, M> {
1650        static PROMPT: std::sync::OnceLock<crate::message::Message> = std::sync::OnceLock::new();
1651        let prompt = PROMPT.get_or_init(|| crate::message::Message::user("hi"));
1652        StepEvent::CompletionCall {
1653            prompt,
1654            history: &[],
1655            turn: 1,
1656        }
1657    }
1658
1659    #[tokio::test]
1660    async fn runs_hooks_in_registration_order_and_consults_all_on_continue() {
1661        let log = Arc::new(Mutex::new(Vec::new()));
1662        let mut stack = HookStack::<M>::with(Recorder {
1663            label: 1,
1664            log: log.clone(),
1665            stop: false,
1666        });
1667        stack.push(Recorder {
1668            label: 2,
1669            log: log.clone(),
1670            stop: false,
1671        });
1672
1673        let flow = stack.on_event(&ctx(), tool_call_event()).await;
1674
1675        assert!(matches!(flow, Flow::Continue));
1676        assert_eq!(*log.lock().expect("log"), vec![1, 2]);
1677    }
1678
1679    #[tokio::test]
1680    async fn first_terminate_short_circuits_on_chained_tool_call() {
1681        // For a tool-call (a chained event), `Terminate` is terminal mid-chain,
1682        // so a later hook must not run once an earlier hook terminates.
1683        let log = Arc::new(Mutex::new(Vec::new()));
1684        let mut stack = HookStack::<M>::with(Recorder {
1685            label: 1,
1686            log: log.clone(),
1687            stop: true,
1688        });
1689        stack.push(Recorder {
1690            label: 2,
1691            log: log.clone(),
1692            stop: false,
1693        });
1694
1695        let flow = stack.on_event(&ctx(), tool_call_event()).await;
1696
1697        assert!(matches!(flow, Flow::Terminate { .. }));
1698        assert_eq!(
1699            *log.lock().expect("log"),
1700            vec![1],
1701            "a later hook must not run after an earlier hook terminates"
1702        );
1703    }
1704
1705    #[tokio::test]
1706    async fn first_terminate_short_circuits_on_observe_only_events() {
1707        // For an observe-only event (the `_ =>` first-non-`Continue`-wins arm,
1708        // here a `TextDelta`), the first hook to terminate must short-circuit the
1709        // rest — this exercises the arm the chained tool-call test does not.
1710        let log = Arc::new(Mutex::new(Vec::new()));
1711        let mut stack = HookStack::<M>::with(Recorder {
1712            label: 1,
1713            log: log.clone(),
1714            stop: true,
1715        });
1716        stack.push(Recorder {
1717            label: 2,
1718            log: log.clone(),
1719            stop: false,
1720        });
1721
1722        let flow = stack
1723            .on_event(
1724                &ctx(),
1725                StepEvent::TextDelta {
1726                    delta: "hi",
1727                    aggregated: "hi",
1728                },
1729            )
1730            .await;
1731
1732        assert!(matches!(flow, Flow::Terminate { .. }));
1733        assert_eq!(
1734            *log.lock().expect("log"),
1735            vec![1],
1736            "a later hook must not run after an earlier hook terminates an observe-only event"
1737        );
1738    }
1739
1740    #[tokio::test]
1741    async fn completion_call_patches_accumulate_and_consult_every_hook() {
1742        // The core composability fix: a patch from hook 1 must NOT skip hook 2.
1743        let log = Arc::new(Mutex::new(Vec::new()));
1744        let mut stack = HookStack::<M>::with(Patcher {
1745            label: 1,
1746            log: log.clone(),
1747            patch: RequestPatch::new().temperature(0.1),
1748        });
1749        stack.push(Patcher {
1750            label: 2,
1751            log: log.clone(),
1752            patch: RequestPatch::new().max_tokens(256),
1753        });
1754
1755        let flow = stack.on_event(&ctx(), completion_call_event()).await;
1756
1757        assert_eq!(
1758            *log.lock().expect("log"),
1759            vec![1, 2],
1760            "both hooks must run; a mergeable patch does not short-circuit"
1761        );
1762        match flow {
1763            Flow::PatchRequest { patch } => {
1764                assert_eq!(patch.temperature, Some(0.1));
1765                assert_eq!(patch.max_tokens, Some(256));
1766            }
1767            other => panic!("expected a merged PatchRequest, got {other:?}"),
1768        }
1769    }
1770
1771    #[tokio::test]
1772    async fn completion_call_terminate_short_circuits_and_discards_patch() {
1773        let log = Arc::new(Mutex::new(Vec::new()));
1774        let mut stack = HookStack::<M>::with(Patcher {
1775            label: 1,
1776            log: log.clone(),
1777            patch: RequestPatch::new().temperature(0.1),
1778        });
1779        // A terminating recorder in the middle.
1780        stack.push(Recorder {
1781            label: 2,
1782            log: log.clone(),
1783            stop: true,
1784        });
1785        stack.push(Patcher {
1786            label: 3,
1787            log: log.clone(),
1788            patch: RequestPatch::new().max_tokens(256),
1789        });
1790
1791        let flow = stack.on_event(&ctx(), completion_call_event()).await;
1792
1793        assert!(matches!(flow, Flow::Terminate { .. }));
1794        assert_eq!(
1795            *log.lock().expect("log"),
1796            vec![1, 2],
1797            "hook 3 must not run after a terminate"
1798        );
1799    }
1800
1801    #[tokio::test]
1802    async fn nested_stack_composes_patches_without_inner_short_circuit() {
1803        // A HookStack pushed as a hook must not reintroduce short-circuiting:
1804        // the inner stack returns its own merged patch, which the outer stack
1805        // merges again.
1806        let log = Arc::new(Mutex::new(Vec::new()));
1807        let mut inner = HookStack::<M>::with(Patcher {
1808            label: 1,
1809            log: log.clone(),
1810            patch: RequestPatch::new().temperature(0.2),
1811        });
1812        inner.push(Patcher {
1813            label: 2,
1814            log: log.clone(),
1815            patch: RequestPatch::new().max_tokens(128),
1816        });
1817
1818        let mut outer = HookStack::<M>::with(inner);
1819        outer.push(Patcher {
1820            label: 3,
1821            log: log.clone(),
1822            patch: RequestPatch::new().preamble("outer"),
1823        });
1824
1825        let flow = outer.on_event(&ctx(), completion_call_event()).await;
1826
1827        assert_eq!(
1828            *log.lock().expect("log"),
1829            vec![1, 2, 3],
1830            "every hook, including both inner-stack hooks, must run"
1831        );
1832        match flow {
1833            Flow::PatchRequest { patch } => {
1834                assert_eq!(patch.temperature, Some(0.2));
1835                assert_eq!(patch.max_tokens, Some(128));
1836                assert_eq!(patch.preamble.as_deref(), Some("outer"));
1837            }
1838            other => panic!("expected a merged PatchRequest, got {other:?}"),
1839        }
1840    }
1841
1842    #[test]
1843    fn stack_observes_is_the_or_of_its_members() {
1844        let mut stack = HookStack::<M>::with(ObservesOnly(StepEventKind::ToolCall));
1845        stack.push(ObservesOnly(StepEventKind::ToolResult));
1846
1847        assert!(stack.observes(StepEventKind::ToolCall));
1848        assert!(stack.observes(StepEventKind::ToolResult));
1849        assert!(
1850            !stack.observes(StepEventKind::TextDelta),
1851            "no member observes TextDelta, so the stack must not either"
1852        );
1853    }
1854
1855    #[tokio::test]
1856    async fn empty_stack_continues_and_observes_nothing() {
1857        let stack = HookStack::<M>::new();
1858
1859        assert!(stack.is_empty());
1860        assert!(!stack.observes(StepEventKind::ToolCall));
1861        assert!(!stack.observes(StepEventKind::TextDelta));
1862        assert!(matches!(
1863            stack.on_event(&ctx(), tool_call_event()).await,
1864            Flow::Continue
1865        ));
1866    }
1867
1868    #[test]
1869    fn unit_hook_observes_no_event_kind() {
1870        // `impl AgentHook for ()` is the no-op hook: it must report interest in
1871        // *no* event kind, so the runner can skip building and dispatching even
1872        // the high-frequency delta events for it. The trait-default `observes`
1873        // returns `true`; `()` deliberately overrides it to `false`, so this is
1874        // the regression guard that the override stays in place.
1875        let all_kinds = [
1876            StepEventKind::CompletionCall,
1877            StepEventKind::CompletionResponse,
1878            StepEventKind::ModelTurnFinished,
1879            StepEventKind::InvalidToolCall,
1880            StepEventKind::ToolCall,
1881            StepEventKind::ToolResult,
1882            StepEventKind::TextDelta,
1883            StepEventKind::ToolCallDelta,
1884            StepEventKind::StreamResponseFinish,
1885        ];
1886        let unit_stack = HookStack::<M>::with(());
1887        for kind in all_kinds {
1888            assert!(
1889                !<() as AgentHook<M>>::observes(&(), kind),
1890                "the `()` no-op hook must not observe {kind:?}"
1891            );
1892            // A stack wrapping only `()` inherits that: it observes nothing, so
1893            // the runner skips delta dispatch for it too.
1894            assert!(
1895                !unit_stack.observes(kind),
1896                "a HookStack::with(()) must not observe {kind:?} either"
1897            );
1898        }
1899    }
1900
1901    // --- RequestPatch merge unit tests ---
1902
1903    #[test]
1904    fn merge_appends_extra_context_in_order() {
1905        let doc = |id: &str| crate::completion::Document {
1906            id: id.to_string(),
1907            text: String::new(),
1908            additional_props: Default::default(),
1909        };
1910        let a = RequestPatch::new().context(doc("a"));
1911        let b = RequestPatch::new().context(doc("b"));
1912        let merged = a.merge(b);
1913        let ids: Vec<&str> = merged.extra_context.iter().map(|d| d.id.as_str()).collect();
1914        assert_eq!(ids, vec!["a", "b"]);
1915    }
1916
1917    #[test]
1918    fn merge_shallow_merges_additional_params_later_wins() {
1919        let a = RequestPatch::new().additional_params(serde_json::json!({"x": 1, "y": 2}));
1920        let b = RequestPatch::new().additional_params(serde_json::json!({"y": 3, "z": 4}));
1921        let merged = a.merge(b);
1922        assert_eq!(
1923            merged.additional_params,
1924            Some(serde_json::json!({"x": 1, "y": 3, "z": 4}))
1925        );
1926    }
1927
1928    #[test]
1929    fn merge_scalar_last_writer_wins() {
1930        let a = RequestPatch::new().temperature(0.1);
1931        let b = RequestPatch::new().temperature(0.9);
1932        assert_eq!(a.merge(b).temperature, Some(0.9));
1933    }
1934
1935    #[test]
1936    fn merge_active_tools_intersects() {
1937        let a = RequestPatch::new().active_tools(["search", "add", "sub"]);
1938        let b = RequestPatch::new().active_tools(["add", "sub", "mul"]);
1939        let merged = a.merge(b);
1940        assert_eq!(
1941            merged.active_tools,
1942            Some(vec!["add".to_string(), "sub".to_string()])
1943        );
1944    }
1945
1946    #[test]
1947    fn merge_active_tools_empty_intersection_yields_empty() {
1948        let a = RequestPatch::new().active_tools(["search"]);
1949        let b = RequestPatch::new().active_tools(["add"]);
1950        let merged = a.merge(b);
1951        assert_eq!(merged.active_tools, Some(vec![]));
1952    }
1953
1954    #[test]
1955    fn merge_one_sided_active_tools_keeps_the_present_list() {
1956        let a = RequestPatch::new().active_tools(["search"]);
1957        let b = RequestPatch::new();
1958        assert_eq!(a.merge(b).active_tools, Some(vec!["search".to_string()]));
1959    }
1960
1961    // --- Scratchpad tests ---
1962
1963    #[test]
1964    fn scratchpad_insert_get_update() {
1965        #[derive(Clone, Default, Debug, PartialEq)]
1966        struct Count(u32);
1967
1968        let pad = Scratchpad::default();
1969        assert_eq!(pad.get::<Count>(), None);
1970        pad.update(|c: &mut Count| c.0 += 1);
1971        pad.update(|c: &mut Count| c.0 += 1);
1972        assert_eq!(pad.get::<Count>(), Some(Count(2)));
1973        assert!(pad.contains::<Count>());
1974        assert_eq!(pad.remove::<Count>(), Some(Count(2)));
1975        assert!(!pad.contains::<Count>());
1976    }
1977
1978    #[test]
1979    fn scratchpad_is_shared_across_clones() {
1980        let pad = Scratchpad::default();
1981        let clone = pad.clone();
1982        pad.insert(7u32);
1983        // The clone shares the same underlying storage.
1984        assert_eq!(clone.get::<u32>(), Some(7));
1985    }
1986
1987    #[test]
1988    fn hook_context_reports_identity_and_turn() {
1989        let ctx = HookContext::new(true, Some("agent".to_string()));
1990        assert!(ctx.is_streaming());
1991        assert_eq!(ctx.agent_name(), Some("agent"));
1992        assert_eq!(ctx.turn(), 0);
1993        ctx.set_turn(3);
1994        assert_eq!(ctx.turn(), 3);
1995        assert!(!ctx.run_id().as_str().is_empty());
1996    }
1997
1998    /// Nested `HookStack` composition of the `ToolCall` chain: a rewrite inside an
1999    /// inner stack must survive a later short-circuit even though the inner stack
2000    /// is dispatched as a single hook. Regression coverage for the bug where
2001    /// `resolve_tool_call` consulted members via `on_event` (one `Flow`, so an
2002    /// inner rewrite was dropped behind an inner `Skip`/`Terminate`).
2003    mod nested_tool_call_resolution {
2004        use super::super::{AgentHook, Flow, HookContext, HookStack, StepEvent};
2005        use super::{M, ctx};
2006        use serde_json::{Value, json};
2007
2008        /// Rewrites the tool args to a fixed value on `ToolCall`.
2009        struct RewriteHook(Value);
2010        impl AgentHook<M> for RewriteHook {
2011            async fn on_event(&self, _ctx: &HookContext, event: StepEvent<'_, M>) -> Flow {
2012                if let StepEvent::ToolCall { .. } = event {
2013                    Flow::rewrite_args(self.0.clone())
2014                } else {
2015                    Flow::cont()
2016                }
2017            }
2018        }
2019
2020        /// Skips on `ToolCall`.
2021        struct SkipHook;
2022        impl AgentHook<M> for SkipHook {
2023            async fn on_event(&self, _ctx: &HookContext, event: StepEvent<'_, M>) -> Flow {
2024                if let StepEvent::ToolCall { .. } = event {
2025                    Flow::skip("denied")
2026                } else {
2027                    Flow::cont()
2028                }
2029            }
2030        }
2031
2032        /// Terminates on `ToolCall`.
2033        struct TerminateHook;
2034        impl AgentHook<M> for TerminateHook {
2035            async fn on_event(&self, _ctx: &HookContext, event: StepEvent<'_, M>) -> Flow {
2036                if let StepEvent::ToolCall { .. } = event {
2037                    Flow::terminate("stop")
2038                } else {
2039                    Flow::cont()
2040                }
2041            }
2042        }
2043
2044        /// Returns `Flow::Fail` on `ToolCall` — not honored there, so it is
2045        /// fail-closed by `run_single_tool`; `resolve_tool_call` returns it verbatim.
2046        struct FailHook;
2047        impl AgentHook<M> for FailHook {
2048            async fn on_event(&self, _ctx: &HookContext, event: StepEvent<'_, M>) -> Flow {
2049                if let StepEvent::ToolCall { .. } = event {
2050                    Flow::fail()
2051                } else {
2052                    Flow::cont()
2053                }
2054            }
2055        }
2056
2057        /// Records the `args` each hook observes on `ToolCall`, to prove the
2058        /// rewritten args are threaded to hooks *after* the rewrite.
2059        #[derive(Clone, Default)]
2060        struct ArgsSpy(std::sync::Arc<std::sync::Mutex<Vec<String>>>);
2061        impl AgentHook<M> for ArgsSpy {
2062            async fn on_event(&self, _ctx: &HookContext, event: StepEvent<'_, M>) -> Flow {
2063                if let StepEvent::ToolCall { args, .. } = event {
2064                    self.0.lock().expect("spy").push(args.to_string());
2065                }
2066                Flow::cont()
2067            }
2068        }
2069
2070        async fn resolve(stack: &HookStack<M>) -> (Flow, Option<Value>) {
2071            stack
2072                .resolve_tool_call(&ctx(), "add", Some("tc1"), "ic1", "{}")
2073                .await
2074        }
2075
2076        #[tokio::test]
2077        async fn nested_rewrite_then_skip_preserves_rewrite() {
2078            // Inner stack: rewrite args, then skip. The rewrite must be salvaged.
2079            let mut inner = HookStack::<M>::new();
2080            inner.push(RewriteHook(json!({ "x": 41 })));
2081            inner.push(SkipHook);
2082
2083            let mut outer = HookStack::<M>::new();
2084            outer.push(inner);
2085
2086            let (flow, salvaged) = resolve(&outer).await;
2087            assert!(matches!(flow, Flow::Skip { .. }), "got {flow:?}");
2088            assert_eq!(
2089                salvaged,
2090                Some(json!({ "x": 41 })),
2091                "the inner rewrite must survive the inner skip through a nested stack"
2092            );
2093        }
2094
2095        #[tokio::test]
2096        async fn nested_rewrite_then_terminate_preserves_rewrite() {
2097            let mut inner = HookStack::<M>::new();
2098            inner.push(RewriteHook(json!({ "x": 7 })));
2099            inner.push(TerminateHook);
2100            let mut outer = HookStack::<M>::new();
2101            outer.push(inner);
2102
2103            let (flow, salvaged) = resolve(&outer).await;
2104            assert!(matches!(flow, Flow::Terminate { .. }), "got {flow:?}");
2105            assert_eq!(salvaged, Some(json!({ "x": 7 })));
2106        }
2107
2108        #[tokio::test]
2109        async fn nested_rewrite_then_fail_closed_preserves_rewrite() {
2110            let mut inner = HookStack::<M>::new();
2111            inner.push(RewriteHook(json!({ "x": 9 })));
2112            inner.push(FailHook);
2113            let mut outer = HookStack::<M>::new();
2114            outer.push(inner);
2115
2116            let (flow, salvaged) = resolve(&outer).await;
2117            // `Fail` is not honored for a tool call, but resolution returns it
2118            // verbatim (run_single_tool fail-closes it); the rewrite still survives.
2119            assert!(matches!(flow, Flow::Fail), "got {flow:?}");
2120            assert_eq!(salvaged, Some(json!({ "x": 9 })));
2121        }
2122
2123        #[tokio::test]
2124        async fn outer_rewrite_then_nested_skip_preserves_outer_rewrite() {
2125            // Outer rewrite, then a nested stack that skips (without its own
2126            // rewrite). The outer rewrite must be salvaged, and the nested stack
2127            // must observe the outer-rewritten args.
2128            let spy = ArgsSpy::default();
2129            let mut inner = HookStack::<M>::new();
2130            inner.push(spy.clone());
2131            inner.push(SkipHook);
2132
2133            let mut outer = HookStack::<M>::new();
2134            outer.push(RewriteHook(json!({ "x": 1, "y": 2 })));
2135            outer.push(inner);
2136
2137            let (flow, salvaged) = resolve(&outer).await;
2138            assert!(matches!(flow, Flow::Skip { .. }), "got {flow:?}");
2139            assert_eq!(salvaged, Some(json!({ "x": 1, "y": 2 })));
2140            // The nested stack saw the outer-rewritten args, not the original `{}`.
2141            assert_eq!(
2142                spy.0.lock().expect("spy").as_slice(),
2143                [serde_json::to_string(&json!({ "x": 1, "y": 2 })).unwrap()],
2144            );
2145        }
2146
2147        #[tokio::test]
2148        async fn deeply_nested_rewrite_then_skip_preserves_rewrite() {
2149            // Three levels: level3 rewrites+skips, wrapped twice.
2150            let mut level3 = HookStack::<M>::new();
2151            level3.push(RewriteHook(json!({ "deep": true })));
2152            level3.push(SkipHook);
2153            let mut level2 = HookStack::<M>::new();
2154            level2.push(level3);
2155            let mut level1 = HookStack::<M>::new();
2156            level1.push(level2);
2157
2158            let (flow, salvaged) = resolve(&level1).await;
2159            assert!(matches!(flow, Flow::Skip { .. }), "got {flow:?}");
2160            assert_eq!(salvaged, Some(json!({ "deep": true })));
2161        }
2162
2163        #[tokio::test]
2164        async fn nested_proceeding_rewrite_surfaces_as_rewrite_args() {
2165            // An inner stack that only rewrites (no short-circuit) surfaces the
2166            // rewrite through the flow itself, with no salvaged second element.
2167            let mut inner = HookStack::<M>::new();
2168            inner.push(RewriteHook(json!({ "x": 5 })));
2169            let mut outer = HookStack::<M>::new();
2170            outer.push(inner);
2171
2172            let (flow, salvaged) = resolve(&outer).await;
2173            assert_eq!(
2174                flow,
2175                Flow::RewriteArgs {
2176                    args: json!({ "x": 5 })
2177                }
2178            );
2179            assert_eq!(salvaged, None);
2180        }
2181    }
2182}