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}