Skip to main content

Module hook

Module hook 

Source
Expand description

Event-specific hooks for observing and steering an agent run.

AgentHook replaces the old universal event/action pair with one lifecycle method and one action type per event. Unsupported combinations are therefore rejected by the compiler instead of being interpreted at runtime. Hooks are independent of the agent’s CompletionModel: managed response events carry canonical Rig messages, content, usage, and message IDs. Use the direct completion or streaming APIs when a hook-like integration needs the provider’s typed raw response.

Hooks run in registration order through HookStack. Model selections, tool-call argument rewrites, and tool-result presentation rewrites chain into later hooks; completion-call RequestPatch values accumulate and merge. A ModelTurnAction::Retry or stop action short-circuits the remaining hooks for that event. Nested stacks obey the same rules as flat stacks, including preserving an argument rewrite when an inner stack later skips or stops.

Register observe-only hooks before steering hooks when every observation is required: a steering stop intentionally prevents later observers from running. Tool-result rewrites change the effective presentation sent to the model and recorded as result-content telemetry. The ToolResultEvent::raw_result and its ToolResultEvent::tool_context remain unchanged for policy decisions and execution-outcome metadata. A tool-result stop omits result content from telemetry.

Blocking and streaming agents share model-turn, request, tool-call, and tool-result resolution. Streaming adds text, reasoning, and tool-call delta observations, but shared lifecycle actions have identical semantics on both surfaces. Streamed deltas are provisional until the model turn is accepted; a retry is surfaced as MultiTurnStreamItem::ModelTurnRetried so consumers can discard the rejected turn’s deltas.

§Example

use rig_agent::agent::{
    AgentHook, CompletionResponseEvent, HookContext, ObservationAction,
};

struct ResponseLogger;

impl AgentHook for ResponseLogger {
    async fn on_completion_response(
        &self,
        _ctx: &HookContext,
        event: CompletionResponseEvent<'_>,
    ) -> ObservationAction {
        println!(
            "message {:?}: {:?} ({:?})",
            event.message_id, event.content, event.usage
        );
        ObservationAction::continue_run()
    }
}

§Retrying a completed model turn

A hook can reject a tool-free turn and either reuse the same prompt and preceding history with fresh request preparation, or preserve the rejected response and append corrective feedback. Retries use the run’s existing total model-call budget. A narrower policy limit belongs to the hook and can be stored in the run-scoped Scratchpad:

use std::{collections::HashMap, sync::atomic::{AtomicUsize, Ordering}};
use rig_agent::agent::{AgentHook, HookContext, ModelTurnAction, ModelTurnFinished};
use rig_core::message::AssistantContent;

static NEXT_HOOK_ID: AtomicUsize = AtomicUsize::new(1);

#[derive(Clone, Default)]
struct RetryCounts(HashMap<usize, usize>);

struct RetryOnMarker {
    id: usize,
    max_retries: usize,
}

impl RetryOnMarker {
    fn new(max_retries: usize) -> Self {
        Self {
            id: NEXT_HOOK_ID.fetch_add(1, Ordering::Relaxed),
            max_retries,
        }
    }
}

impl AgentHook for RetryOnMarker {
    async fn on_model_turn_finished(
        &self,
        ctx: &HookContext,
        event: ModelTurnFinished<'_>,
    ) -> ModelTurnAction {
        let rejected = event.content.iter().any(|content| {
            matches!(content, AssistantContent::Text(text) if text.text.contains("RETRY"))
        });
        if !rejected {
            return ModelTurnAction::continue_run();
        }

        let attempt = ctx.scratchpad().update::<RetryCounts, _>(|counts| {
            let attempt = counts.0.entry(self.id).or_default();
            *attempt += 1;
            *attempt
        });
        if attempt <= self.max_retries {
            ModelTurnAction::retry_with_feedback("Return a complete answer.")
        } else {
            ModelTurnAction::stop("response retry limit exceeded")
        }
    }
}

§Retrying a turn the provider cut short

ModelTurnFinished::finish_reason and ModelTurnFinished::max_tokens carry a turn’s termination metadata in portable form, so the common “truncated at the cap, so raise it and go again” policy needs no provider types. finish_reason is a normalized FinishReason — anything outside the shared vocabulary arrives as Other in the provider’s own spelling rather than as a natural stop, and None means the provider reported no reason at all. max_tokens is the cap this attempt ran under, after the agent’s configuration, the runner override, and any merged RequestPatch, so the pair below reads its own escalation back on the retried turn:

use std::sync::atomic::{AtomicU64, Ordering};
use rig_agent::agent::{
    AgentHook, CompletionCallAction, CompletionCallEvent, HookContext,
    ModelTurnAction, ModelTurnFinished, RequestPatch,
};
use rig_core::completion::FinishReason;
use rig_core::message::AssistantContent;

/// Doubles the output cap each time a turn is truncated, up to a ceiling.
struct GrowCapOnTruncation {
    cap: AtomicU64,
    ceiling: u64,
}

impl AgentHook for GrowCapOnTruncation {
    /// Every attempt is prepared afresh, so the current cap is applied here
    /// and reported back on that attempt's `ModelTurnFinished`.
    async fn on_completion_call(
        &self,
        _ctx: &HookContext,
        _event: CompletionCallEvent<'_>,
    ) -> CompletionCallAction {
        CompletionCallAction::patch(
            RequestPatch::new().max_tokens(self.cap.load(Ordering::Relaxed)),
        )
    }

    async fn on_model_turn_finished(
        &self,
        _ctx: &HookContext,
        event: ModelTurnFinished<'_>,
    ) -> ModelTurnAction {
        // `truncated_output` covers every reason that means "cut short",
        // so a provider reporting a filter stop retries here too.
        let truncated = event
            .finish_reason
            .is_some_and(FinishReason::truncated_output);
        // Retrying a turn that carries tool calls is rejected, so a policy
        // that might see one has to check before asking.
        let has_tool_call = event
            .content
            .iter()
            .any(|content| matches!(content, AssistantContent::ToolCall(_)));
        // `max_tokens` is this attempt's own cap: growing past the ceiling
        // would be retrying a limit we already know we cannot raise.
        let room = event.max_tokens.is_none_or(|cap| cap < self.ceiling);

        if truncated && !has_tool_call && room {
            let grown = event.max_tokens.map_or(self.ceiling, |cap| {
                cap.saturating_mul(2).min(self.ceiling)
            });
            self.cap.store(grown, Ordering::Relaxed);
            return ModelTurnAction::repeat();
        }
        ModelTurnAction::continue_run()
    }
}

cargo run -p rig-agent --example retry_on_truncation runs this policy against a credential-free scripted model whose output genuinely depends on the cap, on both surfaces.

Structs§

CompletionCall
Completion-call event.
CompletionResponse
Canonical non-streaming completion response event.
HookContext
Run-scoped context supplied to hooks.
HookStack
Ordered composable hook stack.
InvalidToolCallContext
Diagnostics for an invalid model-emitted tool call.
ModelSelection
Model-selection event resolved after completion-call hooks and before request preparation.
ModelTurnFinished
Medium-neutral accepted model-turn event.
ReasoningDelta
Streaming reasoning delta.
RequestPatch
A non-sticky patch applied only to the current turn’s completion request.
RunId
Opaque process-scoped identifier for one agent run.
Scratchpad
Run-scoped typed storage shared by hooks.
StreamResponseFinish
Canonical streaming response-finish event.
TextDelta
Streaming text delta.
ToolCall
Pre-execution tool event.
ToolCallDelta
Streaming tool-call delta.
ToolResultEvent
Post-execution tool event.

Enums§

CompletionCallAction
Action for completion-call hooks.
InvalidToolCallAction
Action for invalid-tool-call hooks and manual invalid-call resolution.
ModelSelectionAction
Action for model-selection hooks.
ModelTurnAction
Action for the medium-neutral ModelTurnFinished event.
ObservationAction
Action for observe-only lifecycle events.
RetryRequest
How an accepted, tool-free model turn should be retried.
StepEventKind
Hook event kind used only as an observation performance hint.
ToolCallAction
Action for pre-tool hooks.
ToolResultAction
Action for post-tool hooks.

Traits§

AgentHook
Per-run lifecycle observer and steerer.