Skip to main content

rig_agent/agent/
runner.rs

1//! [`AgentRunner`]: the hook-aware driver that turns a sans-IO
2//! [`AgentRun`] into a complete agent loop.
3//!
4//! [`AgentRun`] decides *what* to do next; it
5//! performs no IO and carries no hooks. `AgentRunner` pairs that machine with
6//! the side-effecting concerns — building and sending completion requests,
7//! executing tools, loading/saving conversation memory — and fires an
8//! [`AgentHook`] at every observable point. Both the blocking
9//! [`PromptRequest`](crate::agent::prompt_request::PromptRequest) and the
10//! [`StreamingPromptRequest`](crate::agent::prompt_request::streaming::StreamingPromptRequest)
11//! APIs are thin wrappers over an `AgentRunner`, and you can build one directly
12//! to drive an agent with custom, composable hooks:
13//!
14//! ```rust,no_run
15//! # use rig_agent::Agent;
16//! # use rig_core::completion::CompletionModel;
17//! # async fn example<M: CompletionModel + 'static>(agent: Agent<M>) -> Result<(), Box<dyn std::error::Error>> {
18//! let response = agent
19//!     .runner("What is 2 + 2?")
20//!     .max_turns(3)
21//!     .run()
22//!     .await?;
23//! println!("{}", response.output);
24//! # Ok(())
25//! # }
26//! ```
27
28use std::sync::{
29    Arc, Mutex,
30    atomic::{AtomicU64, Ordering},
31};
32
33use futures::StreamExt;
34use tracing::{Instrument, info_span, span::Id};
35
36use super::{
37    completion::{Agent, PreparedCompletionRequest},
38    hook::{
39        AgentHook, CompletionCall, CompletionCallAction,
40        CompletionResponse as CompletionResponseEvent, HookContext, HookStack,
41        InvalidToolCallAction, ModelTurnAction, ModelTurnFinished, ObservationAction, RequestPatch,
42        ToolCall as ToolCallEvent, ToolCallAction, ToolResultAction, ToolResultEvent,
43    },
44    prompt_request::{
45        PromptResponse,
46        streaming::{
47            DriveItem, DriveStream, MultiTurnStreamItem, StreamingError, TurnSource, drive_agent,
48            drive_tool_calls, record_usage_on_span, streaming_error_into_prompt,
49        },
50        tool_result_output,
51    },
52    run::{
53        AgentRun, DEFAULT_OUTPUT_RETRIES, ModelTurn, ModelTurnOutcome, OutputMode, PendingToolCall,
54    },
55};
56use rig_core::{
57    memory::ConversationMemory,
58    message::{ToolCall, ToolChoice, UserContent},
59};
60
61use crate::{
62    completion::{CompletionError, CompletionModel, Document, Message, PromptError, Usage},
63    json_utils,
64    tool::{
65        ToolContext, ToolDispatch, ToolOutput, ToolResult,
66        server::{ToolRegistrySnapshot, ToolServerHandle},
67    },
68};
69
70use super::UNKNOWN_AGENT_NAME;
71
72#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
73pub(crate) enum UnhandledInvalidToolCallPolicy {
74    #[default]
75    Fail,
76    IgnoreForExtractor,
77}
78
79/// Build the per-turn `chat` span shared by both turn sources.
80///
81/// The span *name* must be a string literal — `tracing` bakes it into static
82/// metadata — so this is a macro parameterized by the name rather than a
83/// function (the two surfaces keep distinct names, `chat` vs `chat_streaming`,
84/// which dashboards split on). The matching operation value is passed with the
85/// name; every other field is identical across the two surfaces, so it lives
86/// here once instead of being copy-pasted into each `TurnSource::open_chat_span`.
87macro_rules! build_chat_span {
88    ($runner:expr, $effective_preamble:expr, $name:literal, $operation:literal) => {{
89        let system_instructions = $crate::core::telemetry::system_instructions_json(
90            $effective_preamble,
91            $runner.record_telemetry_content,
92        );
93        // The core macro is the single source of the completion-parent
94        // contract (marker + required fields); only the agent-specific field
95        // is declared here.
96        $crate::core::telemetry::completion_parent_span!(
97            target: "rig::agent_chat",
98            name: $name,
99            operation: $operation,
100            system_instructions: system_instructions.as_deref(),
101            gen_ai.agent.name = $runner.agent_name_or_default(),
102        )
103    }};
104}
105pub(crate) use build_chat_span;
106
107/// Convert an observe-only action into an optional stop reason.
108pub(crate) fn observe_action(action: ObservationAction) -> Option<String> {
109    match action {
110        ObservationAction::Continue => None,
111        ObservationAction::Stop(reason) => Some(reason),
112    }
113}
114
115/// Resolved outcome of the shared, medium-neutral model-turn hook.
116pub(crate) enum ModelTurnDecision {
117    /// Accept the turn and advance normally.
118    Advance,
119    /// The turn was rejected and the run is ready to issue another model call.
120    Retried,
121    /// Stop the run with the supplied reason.
122    Terminate(String),
123}
124
125/// Apply a model-turn hook action to the sans-IO run.
126///
127/// Both blocking and streaming sources use this resolver so retry history,
128/// tool-turn rejection, and state transitions cannot diverge by medium.
129pub(crate) fn resolve_model_turn_action(
130    run: &mut AgentRun,
131    action: ModelTurnAction,
132) -> Result<ModelTurnDecision, PromptError> {
133    match action {
134        ModelTurnAction::Continue => Ok(ModelTurnDecision::Advance),
135        ModelTurnAction::Retry(request) => {
136            run.retry_model_turn(request)?;
137            Ok(ModelTurnDecision::Retried)
138        }
139        ModelTurnAction::Stop(reason) => Ok(ModelTurnDecision::Terminate(reason)),
140    }
141}
142
143pub(crate) enum ToolCallDecision {
144    Proceed,
145    ProceedWith(serde_json::Value),
146    Skip(String),
147    Terminate(String),
148}
149
150pub(crate) fn tool_call_decision(action: ToolCallAction) -> ToolCallDecision {
151    match action {
152        ToolCallAction::Run => ToolCallDecision::Proceed,
153        ToolCallAction::Rewrite(args) => ToolCallDecision::ProceedWith(args),
154        ToolCallAction::Skip(reason) => ToolCallDecision::Skip(reason),
155        ToolCallAction::Stop(reason) => ToolCallDecision::Terminate(reason),
156    }
157}
158
159pub(crate) enum ToolResultDecision {
160    Keep,
161    Replace(ToolOutput),
162    Terminate(String),
163}
164
165pub(crate) fn tool_result_decision(action: ToolResultAction) -> ToolResultDecision {
166    match action {
167        ToolResultAction::Keep => ToolResultDecision::Keep,
168        ToolResultAction::Rewrite(result) => ToolResultDecision::Replace(result),
169        ToolResultAction::Stop(reason) => ToolResultDecision::Terminate(reason),
170    }
171}
172
173pub(crate) enum CompletionCallDecision {
174    Proceed,
175    Patch(RequestPatch),
176    Terminate(String),
177}
178
179pub(crate) fn completion_call_decision(action: CompletionCallAction) -> CompletionCallDecision {
180    match action {
181        CompletionCallAction::Continue => CompletionCallDecision::Proceed,
182        CompletionCallAction::Patch(patch) => CompletionCallDecision::Patch(patch),
183        CompletionCallAction::Stop(reason) => CompletionCallDecision::Terminate(reason),
184    }
185}
186
187/// A hook-aware driver over [`AgentRun`].
188///
189/// Construct one from an [`Agent`] with [`Agent::runner`], attach hooks with
190/// [`add_hook`](Self::add_hook), then call
191/// [`run`](Self::run) (blocking) or
192/// [`stream`](crate::agent::prompt_request::streaming::StreamingPromptRequest)
193/// (incremental). Hooks are held in a [`HookStack`], an ordered,
194/// runtime-composable list; `run()` and `stream()` share the same loop and fire
195/// the same events, so they behave identically apart from the streamed delta
196/// events the medium adds.
197#[non_exhaustive]
198pub struct AgentRunner<M>
199where
200    M: CompletionModel,
201{
202    pub(crate) prompt: Message,
203    pub(crate) chat_history: Option<Vec<Message>>,
204    pub(crate) max_turns: usize,
205    pub(crate) max_invalid_tool_call_retries: usize,
206    pub(crate) model: Arc<M>,
207    pub(crate) agent_name: Option<String>,
208    pub(crate) preamble: Option<String>,
209    pub(crate) static_context: Vec<Document>,
210    pub(crate) temperature: Option<f64>,
211    pub(crate) max_tokens: Option<u64>,
212    pub(crate) additional_params: Option<serde_json::Value>,
213    pub(crate) record_telemetry_content: bool,
214    pub(crate) tool_server_handle: ToolServerHandle,
215    /// Typed context cloned freshly for every tool dispatch.
216    pub(crate) tool_context: ToolContext,
217    pub(crate) tool_choice: Option<ToolChoice>,
218    pub(crate) output_schema: Option<schemars::Schema>,
219    pub(crate) output_mode: OutputMode,
220    pub(crate) output_tool_name: Option<String>,
221    pub(crate) output_tool_description: Option<String>,
222    pub(crate) augment_output_preamble: bool,
223    pub(crate) unhandled_invalid_tool_call_policy: UnhandledInvalidToolCallPolicy,
224    pub(crate) concurrency: usize,
225    pub(crate) memory: Option<Arc<dyn ConversationMemory>>,
226    pub(crate) conversation_id: Option<String>,
227    pub(crate) hooks: HookStack,
228    pub(crate) error_usage: Option<Arc<Mutex<Usage>>>,
229}
230
231impl<M> AgentRunner<M>
232where
233    M: CompletionModel,
234{
235    /// Build a runner from an agent, seeding it with the agent's default hook
236    /// stack. Prefer [`Agent::runner`].
237    pub fn from_agent(agent: &Agent<M>, prompt: impl Into<Message>) -> Self {
238        Self {
239            prompt: prompt.into(),
240            chat_history: None,
241            max_turns: agent.default_max_turns.unwrap_or(1),
242            max_invalid_tool_call_retries: 0,
243            model: agent.model.clone(),
244            agent_name: agent.name.clone(),
245            preamble: agent.preamble.clone(),
246            static_context: agent.static_context.clone(),
247            temperature: agent.temperature,
248            max_tokens: agent.max_tokens,
249            additional_params: agent.additional_params.clone(),
250            record_telemetry_content: agent.record_telemetry_content,
251            tool_server_handle: agent.tool_server_handle.clone(),
252            tool_context: ToolContext::new(),
253            tool_choice: agent.tool_choice.clone(),
254            output_schema: agent.output_schema.clone(),
255            output_mode: agent.output_mode.clone(),
256            output_tool_name: None,
257            output_tool_description: None,
258            augment_output_preamble: true,
259            unhandled_invalid_tool_call_policy: UnhandledInvalidToolCallPolicy::Fail,
260            concurrency: 1,
261            memory: agent.memory.clone(),
262            conversation_id: agent.default_conversation_id.clone(),
263            hooks: agent.hooks.clone(),
264            error_usage: None,
265        }
266    }
267
268    /// Append a hook to the stack (on top of any the agent already carries).
269    /// Hooks run in registration order; how their results compose is
270    /// event-dependent (`CompletionCall` request patches accumulate and merge,
271    /// `ToolCall`/`ToolResult` rewrites chain, while model-turn steering and
272    /// observe-only/recovery events use their event-specific terminal action). See the
273    /// [`hook`](crate::agent::hook) module docs.
274    pub fn add_hook<H>(mut self, hook: H) -> Self
275    where
276        H: AgentHook + 'static,
277    {
278        self.hooks.push(hook);
279        self
280    }
281}
282
283impl<M> AgentRunner<M>
284where
285    M: CompletionModel,
286{
287    /// Set the total model-call budget, including the initial call and every
288    /// retry or continuation. Zero emits no model calls; one permits only the
289    /// initial call. Exceeding the budget returns [`PromptError::MaxTurnsError`].
290    pub fn max_turns(mut self, max_turns: usize) -> Self {
291        self.max_turns = max_turns;
292        self
293    }
294
295    /// Set the typed context cloned for every tool dispatch in this run.
296    pub fn tool_context(mut self, context: ToolContext) -> Self {
297        self.tool_context = context;
298        self
299    }
300
301    /// Set the chat history preceding the prompt. Passing explicit history
302    /// bypasses conversation memory for this run.
303    pub fn history<I, T>(mut self, history: I) -> Self
304    where
305        I: IntoIterator<Item = T>,
306        T: Into<Message>,
307    {
308        self.chat_history = Some(history.into_iter().map(Into::into).collect());
309        self
310    }
311
312    /// Override the agent preamble for this run.
313    pub fn preamble(mut self, preamble: impl Into<String>) -> Self {
314        self.preamble = Some(preamble.into());
315        self
316    }
317
318    /// Remove the agent's configured preamble for this run.
319    pub fn without_preamble(mut self) -> Self {
320        self.preamble = None;
321        self
322    }
323
324    /// Append one static context document for this run.
325    pub fn document(mut self, document: Document) -> Self {
326        self.static_context.push(document);
327        self
328    }
329
330    /// Append static context documents for this run.
331    pub fn documents(mut self, documents: impl IntoIterator<Item = Document>) -> Self {
332        self.static_context.extend(documents);
333        self
334    }
335
336    /// Override the model temperature for this run.
337    pub fn temperature(mut self, temperature: f64) -> Self {
338        self.temperature = Some(temperature);
339        self
340    }
341
342    /// Remove the agent's configured temperature for this run.
343    pub fn without_temperature(mut self) -> Self {
344        self.temperature = None;
345        self
346    }
347
348    /// Override the maximum completion token count for this run.
349    pub fn max_tokens(mut self, max_tokens: u64) -> Self {
350        self.max_tokens = Some(max_tokens);
351        self
352    }
353
354    /// Remove the agent's configured maximum token count for this run.
355    pub fn without_max_tokens(mut self) -> Self {
356        self.max_tokens = None;
357        self
358    }
359
360    /// Shallow-merge object fields into the provider-specific parameters for
361    /// this run. Later fields win. A non-object baseline is replaced by the
362    /// supplied object. A later completion-call hook patch has final
363    /// precedence: object values shallow-merge, while a non-object on either
364    /// side causes wholesale replacement by the hook value.
365    pub fn merge_additional_params(
366        mut self,
367        params: serde_json::Map<String, serde_json::Value>,
368    ) -> Self {
369        let params = serde_json::Value::Object(params);
370        self.additional_params = Some(match self.additional_params.take() {
371            Some(baseline) if baseline.is_object() => crate::json_utils::merge(baseline, params),
372            _ => params,
373        });
374        self
375    }
376
377    /// Replace all provider-specific parameters for this run. A later
378    /// completion-call hook patch has final precedence: object values
379    /// shallow-merge, while a non-object on either side causes wholesale
380    /// replacement by the hook value.
381    pub fn replace_additional_params(mut self, params: serde_json::Value) -> Self {
382        self.additional_params = Some(params);
383        self
384    }
385
386    /// Remove the agent's configured provider-specific parameters for this run.
387    /// A later completion-call hook may still supply its own parameters.
388    pub fn without_additional_params(mut self) -> Self {
389        self.additional_params = None;
390        self
391    }
392
393    /// Override the tool-choice policy for this run.
394    pub fn tool_choice(mut self, tool_choice: ToolChoice) -> Self {
395        self.tool_choice = Some(tool_choice);
396        self
397    }
398
399    /// Remove the agent's configured tool-choice policy for this run.
400    pub fn without_tool_choice(mut self) -> Self {
401        self.tool_choice = None;
402        self
403    }
404
405    /// Configure the synthetic tool used by an internal Tool-output flow.
406    pub(crate) fn output_tool(
407        mut self,
408        name: impl Into<String>,
409        description: impl Into<String>,
410        augment_preamble: bool,
411    ) -> Self {
412        self.output_tool_name = Some(name.into());
413        self.output_tool_description = Some(description.into());
414        self.augment_output_preamble = augment_preamble;
415        self
416    }
417
418    /// Ignore invalid tool calls when every registered hook declines to act.
419    ///
420    /// This is an internal compatibility policy for extractors, whose legacy
421    /// transport treated every non-`submit` call as irrelevant response
422    /// content. Hooks still receive the invalid-call event first and retain
423    /// full control over recovery or termination.
424    pub(crate) fn ignore_unhandled_invalid_tool_calls(mut self) -> Self {
425        self.unhandled_invalid_tool_call_policy =
426            UnhandledInvalidToolCallPolicy::IgnoreForExtractor;
427        self
428    }
429
430    /// Opt in or out of recording sensitive request, response, and tool content
431    /// on GenAI telemetry spans for this run.
432    ///
433    /// Defaults to the agent's setting, which defaults to `false`. Enabling this
434    /// can expose prompts, retrieved context, tool results, model responses, and
435    /// other sensitive or high-cardinality data through OpenTelemetry span
436    /// attributes, which can increase observability backend storage and query
437    /// costs. Only enable it when content telemetry is acceptable for this run.
438    /// Structural metadata and token usage remain available when disabled.
439    pub fn record_content_telemetry(mut self, enabled: bool) -> Self {
440        self.record_telemetry_content = enabled;
441        self
442    }
443
444    /// Execute up to `concurrency` tools at once (1 by default). Applies to
445    /// **both** the blocking [`run`](Self::run) and the streaming
446    /// [`stream`](Self::stream) paths.
447    ///
448    /// The resulting message history is the same in both paths regardless of
449    /// `concurrency`: final tool results are persisted in tool-call order. At
450    /// the default `concurrency` of 1 the two paths are fully in lock-step; with
451    /// `concurrency > 1` the tools run in parallel, so a `ToolCall`/`ToolResult`
452    /// **hook may fire in completion order** rather than call order — the
453    /// per-tool side effects interleave even though the final history does not.
454    ///
455    /// For the streaming path: the driver emits *all* of a turn's `ToolCall`
456    /// stream items eagerly (in call order) when the model turn commits, then —
457    /// only after the whole tool batch settles successfully — surfaces the
458    /// per-tool `ToolExecutionCommitted` and `ToolResult` stream items in **call
459    /// order** (never completion order), for the tools whose body actually ran.
460    /// The persisted message history is unchanged.
461    ///
462    /// A `concurrency` of 0 is clamped to 1; `0` and `1` both run a turn's tools
463    /// sequentially (the `buffer_unordered` path is used only at `concurrency > 1`).
464    pub fn tool_concurrency(mut self, concurrency: usize) -> Self {
465        self.concurrency = concurrency.max(1);
466        self
467    }
468
469    /// Set the conversation id used to load and persist memory for this run.
470    pub fn conversation(mut self, id: impl Into<String>) -> Self {
471        self.conversation_id = Some(id.into());
472        self
473    }
474
475    /// Disable conversation memory for this run (no load, no save).
476    pub fn without_memory(mut self) -> Self {
477        self.memory = None;
478        self.conversation_id = None;
479        self
480    }
481
482    /// Set the retry budget for invalid tool-call recovery. Invalid tool-call
483    /// retries also consume the total model-call budget.
484    pub fn max_invalid_tool_call_retries(mut self, retries: usize) -> Self {
485        self.max_invalid_tool_call_retries = retries;
486        self
487    }
488
489    pub(crate) fn agent_name_or_default(&self) -> &str {
490        self.agent_name.as_deref().unwrap_or(UNKNOWN_AGENT_NAME)
491    }
492
493    /// Build the sans-IO [`AgentRun`] for this runner's configuration.
494    /// `history_override` replaces the configured chat history (e.g. with
495    /// memory-loaded history). Delegates to [`build_agent_run`] — the single
496    /// construction site shared with the streaming driver.
497    pub(crate) fn build_run(&self, history_override: Option<Vec<Message>>) -> AgentRun {
498        let run = build_agent_run(
499            self.prompt.clone(),
500            self.max_turns,
501            self.max_invalid_tool_call_retries,
502            self.output_schema.as_ref(),
503            history_override.or_else(|| self.chat_history.clone()),
504            self.tool_choice.clone(),
505        );
506        match &self.output_tool_name {
507            Some(name) => run.with_output_tool_name(name.clone()),
508            None => run,
509        }
510    }
511}
512
513/// Construct an [`AgentRun`] from explicit run configuration. The single place a
514/// run is built, so the blocking and streaming drivers configure runs
515/// identically.
516pub(crate) fn build_agent_run(
517    prompt: Message,
518    max_turns: usize,
519    max_invalid_tool_call_retries: usize,
520    output_schema: Option<&schemars::Schema>,
521    history: Option<Vec<Message>>,
522    tool_choice: Option<ToolChoice>,
523) -> AgentRun {
524    let mut run = AgentRun::new(prompt)
525        .max_turns(max_turns)
526        .max_invalid_tool_call_retries(max_invalid_tool_call_retries)
527        .with_output_validation(
528            output_schema.map(|schema| schema.as_value().clone()),
529            DEFAULT_OUTPUT_RETRIES,
530        );
531    if let Some(history) = history {
532        run = run.with_history(history);
533    }
534    if let Some(tool_choice) = tool_choice {
535        run = run.with_tool_choice(tool_choice);
536    }
537    run
538}
539
540/// Build (or adopt) the top-level `invoke_agent` span for a run, shared by the
541/// blocking and streaming drivers so the run-level span shape is defined once.
542///
543/// Returns the span plus whether it was newly created. When the caller is
544/// already inside a span we adopt it and report `false`, so the driver can avoid
545/// recording run-level usage onto a span it does not own (see the
546/// `created_agent_span` guard in both drivers' `Done` handling).
547pub(crate) fn acquire_agent_span(
548    agent_name: &str,
549    preamble: Option<&str>,
550    record_content: bool,
551) -> (tracing::Span, bool) {
552    if tracing::Span::current().is_disabled() {
553        let system_instructions =
554            rig_core::telemetry::system_instructions_json(preamble, record_content);
555        let span = info_span!(
556            "invoke_agent",
557            gen_ai.operation.name = "invoke_agent",
558            gen_ai.agent.name = agent_name,
559            gen_ai.system_instructions = system_instructions.as_deref(),
560            gen_ai.prompt = tracing::field::Empty,
561            gen_ai.completion = tracing::field::Empty,
562            gen_ai.usage.input_tokens = tracing::field::Empty,
563            gen_ai.usage.output_tokens = tracing::field::Empty,
564            gen_ai.usage.cache_read.input_tokens = tracing::field::Empty,
565            gen_ai.usage.cache_creation.input_tokens = tracing::field::Empty,
566            gen_ai.usage.tool_use_prompt_tokens = tracing::field::Empty,
567            gen_ai.usage.reasoning_tokens = tracing::field::Empty,
568        );
569        (span, true)
570    } else {
571        (tracing::Span::current(), false)
572    }
573}
574
575/// Outcome of firing the `CompletionCall` hook for a turn.
576pub(crate) enum CompletionCallOutcome {
577    /// Proceed, optionally applying a per-turn request patch (the merged patch
578    /// from every hook that contributed one).
579    Proceed(Option<RequestPatch>),
580    /// Terminate the run with this reason.
581    Terminate(String),
582}
583
584/// Fire the event-specific completion-call hook for a turn.
585pub(crate) async fn resolve_completion_call(
586    hooks: &HookStack,
587    ctx: &HookContext,
588    prompt: &Message,
589    history: &[Message],
590    turn: usize,
591) -> CompletionCallOutcome {
592    match completion_call_decision(
593        hooks
594            .on_completion_call(
595                ctx,
596                CompletionCall {
597                    prompt,
598                    history,
599                    turn,
600                },
601            )
602            .await,
603    ) {
604        CompletionCallDecision::Terminate(reason) => CompletionCallOutcome::Terminate(reason),
605        CompletionCallDecision::Patch(patch) => CompletionCallOutcome::Proceed(Some(patch)),
606        CompletionCallDecision::Proceed => CompletionCallOutcome::Proceed(None),
607    }
608}
609
610/// Append a finished run's messages to conversation memory, logging and
611/// proceeding on failure. Shared `Done`-arm behavior for both drivers.
612pub(crate) async fn append_run_messages(
613    memory_handle: Option<&(Arc<dyn ConversationMemory>, String)>,
614    messages: &[Message],
615) {
616    // Clone into an owned vec only when there is a backend to append to — the
617    // common no-memory path pays nothing.
618    if let Some((memory, id)) = memory_handle
619        && let Err(err) = memory.append(id, messages.to_vec()).await
620    {
621        tracing::warn!(
622            error = %err,
623            conversation_id = %id,
624            "conversation memory append failed; surfacing final response anyway"
625        );
626    }
627}
628
629/// Whether (and how) a tool call executed, for [`run_single_tool`].
630pub(crate) enum ToolExecution {
631    /// The tool's body ran. Carries the **effective** tool call — the model's
632    /// call with any [`ToolCallAction::Rewrite`] hook
633    /// rewrite applied — so the driver can surface it in the
634    /// [`ToolExecutionCommitted`](crate::agent::prompt_request::streaming::MultiTurnStreamItem::ToolExecutionCommitted)
635    /// event (what actually ran, not the model's original arguments). Boxed to
636    /// keep this enum small (a `ToolCall` is large next to the empty `Skipped`).
637    Executed(Box<ToolCall>),
638    /// A tool-call hook returned [`ToolCallAction::Skip`]: the
639    /// body did not run, so no execution-commit is surfaced — but the skip result
640    /// is still delivered to the model (and surfaced as a `ToolResult`).
641    Skipped,
642}
643
644/// Outcome of [`run_single_tool`]: the tool-result content plus whether the
645/// tool's body ran (and the effective call) or a hook skipped it.
646pub(crate) struct ToolCallOutcome {
647    /// The tool result delivered to the model (a real output, a redacted
648    /// replacement, or a hook skip reason).
649    pub content: UserContent,
650    /// How the call resolved: executed (with the effective tool call) or skipped.
651    pub execution: ToolExecution,
652}
653
654/// Execute a single tool call, firing the `ToolCall` and `ToolResult` hooks and
655/// shaping the result. **Shared by the blocking and streaming drivers** so a
656/// tool call behaves identically in both: same hook events, same fail-closed
657/// skip/terminate handling, and the same result shaping. Hook skips become
658/// [`ToolResult::skipped`], and every result is converted directly into typed
659/// message content through [`tool_result_output`] without reparsing text.
660/// Records `gen_ai.tool.*` on the current span;
661/// `error_history` builds a cancellation error if a hook terminates the run.
662/// Returns whether the tool body executed via [`ToolCallOutcome::execution`].
663pub(crate) async fn run_single_tool<M>(
664    runner: &AgentRunner<M>,
665    ctx: &HookContext,
666    tool_snapshot: &ToolRegistrySnapshot,
667    tool_call: &ToolCall,
668    internal_call_id: &str,
669    error_history: &[Message],
670) -> Result<ToolCallOutcome, PromptError>
671where
672    M: CompletionModel,
673{
674    let hooks = &runner.hooks;
675    let tool_context = &runner.tool_context;
676    let record_content = runner.record_telemetry_content;
677    let tool_name = &tool_call.function.name;
678    // `mut` so a tool-call hook can rewrite the arguments the tool
679    // runs with (the model's emitted arguments are otherwise used verbatim).
680    let mut args = json_utils::serialize_json_value(&tool_call.function.arguments);
681
682    let tool_span = tracing::Span::current();
683    tool_span.record("gen_ai.tool.name", tool_name);
684    tool_span.record("gen_ai.tool.call.id", &tool_call.id);
685    if record_content {
686        tool_span.record("gen_ai.tool.call.arguments", &args);
687    }
688
689    // Resolve the `ToolCall` hook chain. A proceeding chain carries any
690    // `ToolCallAction::Rewrite` in the action itself (→ `ProceedWith`); a chain that a
691    // later hook short-circuits with `Skip`/`Terminate` salvages the accumulated
692    // rewrite into `salvaged_rewrite` so it is *not* lost — the rewritten args
693    // must still be reported on the skipped `ToolResult` and in tracing rather
694    // than leaking the model's original args (see [`HookStack::resolve_tool_call`]).
695    let (action, salvaged_rewrite) = hooks
696        .resolve_tool_call(
697            ctx,
698            ToolCallEvent {
699                tool_name,
700                tool_call_id: tool_call.call_id.as_deref(),
701                internal_call_id,
702                args: &args,
703            },
704        )
705        .await;
706
707    // Apply a salvaged rewrite (short-circuit path only) so `args` — what the
708    // `ToolResult` reports — and the span reflect the effective arguments.
709    if let Some(rewritten) = salvaged_rewrite.as_ref() {
710        args = json_utils::serialize_json_value(rewritten);
711        if record_content {
712            tool_span.record("gen_ai.tool.call.arguments", &args);
713        }
714        tracing::debug!(
715            tool_name = tool_name,
716            "tool-call arguments rewritten by a hook"
717        );
718    }
719
720    // On `Skip` the body does not run and the structured outcome is `Skipped`;
721    // otherwise the tool executes into a structured `ToolResult`.
722    // `effective_args` is what the tool actually ran with (the model's, a hook's
723    // `ToolCallAction::Rewrite` replacement, or a salvaged rewrite) — surfaced in the
724    // execution-commit event so a redaction rewrite does not leak. Unused for a skip.
725    let mut skipped: Option<ToolResult> = None;
726    let effective_args: serde_json::Value = match tool_call_decision(action) {
727        ToolCallDecision::Terminate(reason) => {
728            return Err(PromptError::prompt_cancelled(
729                error_history.to_vec(),
730                reason,
731            ));
732        }
733        ToolCallDecision::Skip(reason) => {
734            tracing::info!(tool_name = tool_name, reason = reason, "Tool call rejected");
735            // Synthetic rejection: `Skipped` outcome, message delivered verbatim.
736            // Still fires the `ToolResult` hook so a policy observes the skip.
737            skipped = Some(ToolResult::skipped(reason));
738            // A skip runs nothing; its effective args are the salvaged rewrite
739            // (if any) so tracing/history stay consistent, though they go unused.
740            salvaged_rewrite.unwrap_or_else(|| tool_call.function.arguments.clone())
741        }
742        ToolCallDecision::ProceedWith(replacement) => {
743            // Proceeding rewrite: re-record the span so the trace, and the
744            // downstream `ToolResult` event, reflect what the tool actually
745            // received rather than what the model emitted.
746            args = json_utils::serialize_json_value(&replacement);
747            if record_content {
748                tool_span.record("gen_ai.tool.call.arguments", &args);
749            }
750            tracing::debug!(
751                tool_name = tool_name,
752                "tool-call arguments rewritten by a hook"
753            );
754            replacement
755        }
756        ToolCallDecision::Proceed => tool_call.function.arguments.clone(),
757    };
758
759    // Resolve the structured execution result and how the call surfaced. A skip
760    // produces no execution-commit event; a real execution carries the effective
761    // tool call (the model's call with any `ToolCallAction::Rewrite` applied).
762    let (exec, execution, dispatch_context) = match skipped {
763        Some(exec) => (exec, ToolExecution::Skipped, tool_context.for_dispatch()),
764        None => {
765            let mut effective_tool_call = tool_call.clone();
766            effective_tool_call.function.arguments = effective_args;
767            let ToolDispatch {
768                result: exec,
769                context: dispatch_context,
770            } = tool_snapshot.dispatch(tool_name, &args, tool_context).await;
771            (
772                exec,
773                ToolExecution::Executed(Box::new(effective_tool_call)),
774                dispatch_context,
775            )
776        }
777    };
778    // Presentation rewrites happen after execution. The raw structured result
779    // and per-dispatch context remain unchanged for every hook.
780    let result_decision = tool_result_decision(
781        hooks
782            .on_tool_result(
783                ctx,
784                ToolResultEvent {
785                    tool_name,
786                    tool_call_id: tool_call.call_id.as_deref(),
787                    internal_call_id,
788                    args: &args,
789                    presentation: exec.output(),
790                    raw_result: &exec,
791                    tool_context: &dispatch_context,
792                },
793            )
794            .await,
795    );
796    // Outcome metadata describes the execution itself, while result content
797    // follows the same presentation policy as the model. This keeps redaction
798    // and stop hooks from leaking raw tool output through telemetry.
799    record_tool_result(&tool_span, &exec);
800
801    match result_decision {
802        ToolResultDecision::Terminate(reason) => Err(PromptError::prompt_cancelled(
803            error_history.to_vec(),
804            reason,
805        )),
806        ToolResultDecision::Replace(replacement) => {
807            if record_content {
808                tool_span.record("gen_ai.tool.call.result", replacement.render());
809            }
810            Ok(ToolCallOutcome {
811                content: tool_result_output(
812                    tool_call.id.clone(),
813                    tool_call.call_id.clone(),
814                    replacement,
815                ),
816                execution,
817            })
818        }
819        ToolResultDecision::Keep => {
820            if record_content {
821                tool_span.record("gen_ai.tool.call.result", exec.output().render());
822            }
823            let content = tool_result_output(
824                tool_call.id.clone(),
825                tool_call.call_id.clone(),
826                exec.output().clone(),
827            );
828            Ok(ToolCallOutcome { content, execution })
829        }
830    }
831}
832
833fn record_tool_result(span: &tracing::Span, result: &ToolResult) {
834    span.record("gen_ai.tool.call.outcome", result.status_name());
835    if let Some(error) = result.error() {
836        span.record("gen_ai.tool.error.type", error.kind().as_str());
837    }
838}
839
840/// Build the per-tool `execute_tool` span carrying the `gen_ai.tool.*` fields
841/// that [`run_single_tool`] records on the current span. Parented to the
842/// contextual current span; the blocking driver additionally chains it via
843/// `follows_from`, while the streaming driver uses it as-is. Shared by both
844/// drivers so the span shape stays defined in one place.
845pub(crate) fn new_execute_tool_span() -> tracing::Span {
846    info_span!(
847        "execute_tool",
848        gen_ai.operation.name = "execute_tool",
849        gen_ai.tool.type = "function",
850        gen_ai.tool.name = tracing::field::Empty,
851        gen_ai.tool.call.id = tracing::field::Empty,
852        gen_ai.tool.call.arguments = tracing::field::Empty,
853        gen_ai.tool.call.result = tracing::field::Empty,
854        gen_ai.tool.call.outcome = tracing::field::Empty,
855        gen_ai.tool.error.type = tracing::field::Empty
856    )
857}
858
859/// [`TurnSource`] for the blocking surface: each turn issues a unary
860/// `model.completion()` request and feeds the whole response into the machine.
861/// Emits no intermediate items (the blocking surface folds the engine to its
862/// final response), but keeps the blocking driver's linear `follows_from` span
863/// chain across chat and tool spans.
864pub(crate) struct UnaryTurnSource {
865    /// Sequences chat and tool spans into a linear `follows_from` chain (the
866    /// streaming surface parents into a tree instead and does not chain).
867    ///
868    /// Atomic rather than `Cell` despite being driven by a single sequential
869    /// task: `run_tool_calls` passes `chain_span` as a closure into
870    /// `drive_tool_calls`, whose returned `DriveStream` is `Send`. That makes the
871    /// closure capture `&self`, so `&UnaryTurnSource` must be `Send`, i.e.
872    /// `UnaryTurnSource: Sync` — which `AtomicU64` provides and `Cell` does not.
873    current_span_id: AtomicU64,
874    record_telemetry_content: bool,
875}
876
877impl UnaryTurnSource {
878    pub(crate) fn new(record_telemetry_content: bool) -> Self {
879        Self {
880            current_span_id: AtomicU64::new(0),
881            record_telemetry_content,
882        }
883    }
884
885    /// Chain `span` onto the previous step's span and record it as the new chain
886    /// head, preserving the blocking driver's linear causal trace.
887    fn chain_span(&self, span: tracing::Span) -> tracing::Span {
888        let span = match self.current_span_id.load(Ordering::Relaxed) {
889            0 => span,
890            id => span.follows_from(Id::from_u64(id)).to_owned(),
891        };
892        if let Some(id) = span.id() {
893            self.current_span_id.store(id.into_u64(), Ordering::Relaxed);
894        }
895        span
896    }
897}
898
899impl<M> TurnSource<M> for UnaryTurnSource
900where
901    M: CompletionModel,
902{
903    type Raw = M::Response;
904
905    fn open_chat_span(
906        &self,
907        runner: &AgentRunner<M>,
908        effective_preamble: Option<&str>,
909    ) -> tracing::Span {
910        let chat_span = build_chat_span!(runner, effective_preamble, "chat", "chat");
911        self.chain_span(chat_span)
912    }
913
914    fn run_model_turn<'a>(
915        &'a mut self,
916        runner: &'a AgentRunner<M>,
917        hook_ctx: &'a HookContext,
918        run: &'a mut AgentRun,
919        prepared: PreparedCompletionRequest<M>,
920        chat_span: tracing::Span,
921        _agent_span: &'a tracing::Span,
922        current_prompt: Message,
923    ) -> DriveStream<'a, M::Response> {
924        Box::pin(async_stream::stream! {
925            let resp = match prepared.builder.send().instrument(chat_span.clone()).await {
926                Ok(resp) => resp,
927                Err(err) => {
928                    yield Err(StreamingError::from(err));
929                    return;
930                }
931            };
932
933            let mut outcome = match run.model_response(ModelTurn::new(
934                resp.message_id.clone(),
935                resp.choice.clone(),
936                resp.usage,
937                prepared.executable_tool_names,
938                prepared.allowed_tool_names,
939            )) {
940                Ok(outcome) => outcome,
941                Err(err) => {
942                    yield Err(Box::new(err).into());
943                    return;
944                }
945            };
946
947            loop {
948                match outcome {
949                    ModelTurnOutcome::NeedsResolution(context) => {
950                        let action = runner
951                            .hooks
952                            .on_invalid_tool_call(hook_ctx, &context)
953                            .await;
954                        let resolution = match action {
955                            Some(action) => run.resolve_invalid_tool_call(action),
956                            None
957                                if runner.unhandled_invalid_tool_call_policy
958                                    == UnhandledInvalidToolCallPolicy::IgnoreForExtractor =>
959                            {
960                                run.ignore_invalid_tool_call()
961                            }
962                            None => run.resolve_invalid_tool_call(InvalidToolCallAction::fail()),
963                        };
964                        outcome = match resolution {
965                            Ok(outcome) => outcome,
966                            Err(err) => {
967                                yield Err(Box::new(err).into());
968                                return;
969                            }
970                        };
971                    }
972                    ModelTurnOutcome::TurnRetried => break,
973                    ModelTurnOutcome::Continue {
974                        response_hook_suppressed,
975                    } => {
976                        if !response_hook_suppressed {
977                            // The response-finish event fires first, then the
978                            // normalized per-turn event. The first observes;
979                            // the second can accept, retry, or stop the canonical
980                            // turn. Both are suppressed for recovered turns.
981                            if let Some(reason) = observe_action(
982                                runner
983                                    .hooks
984                                    .on_completion_response(
985                                        hook_ctx,
986                                        CompletionResponseEvent {
987                                            prompt: &current_prompt,
988                                            content: &resp.choice,
989                                            usage: resp.usage,
990                                            message_id: resp.message_id.as_deref(),
991                                        },
992                                    )
993                                    .await,
994                            ) {
995                                if runner.record_telemetry_content
996                                    && let Some(choice) = run.accepted_turn_choice()
997                                {
998                                    rig_core::telemetry::record_model_output(
999                                        &chat_span, &choice, true,
1000                                    );
1001                                }
1002                                yield Err(StreamingError::Prompt(Box::new(run.cancel_error(reason))));
1003                                return;
1004                            }
1005                            let action = runner
1006                                .hooks
1007                                .on_model_turn_finished(
1008                                    hook_ctx,
1009                                    ModelTurnFinished {
1010                                        turn: hook_ctx.turn(),
1011                                        content: &resp.choice,
1012                                        usage: resp.usage,
1013                                    },
1014                                )
1015                                .await;
1016                            match resolve_model_turn_action(run, action) {
1017                                Ok(ModelTurnDecision::Advance) => {}
1018                                Ok(ModelTurnDecision::Retried) => break,
1019                                Ok(ModelTurnDecision::Terminate(reason)) => {
1020                                    if runner.record_telemetry_content
1021                                        && let Some(choice) = run.accepted_turn_choice()
1022                                    {
1023                                        rig_core::telemetry::record_model_output(
1024                                            &chat_span, &choice, true,
1025                                        );
1026                                    }
1027                                    yield Err(StreamingError::Prompt(Box::new(
1028                                        run.cancel_error(reason),
1029                                    )));
1030                                    return;
1031                                }
1032                                Err(err) => {
1033                                    yield Err(StreamingError::Prompt(Box::new(err)));
1034                                    return;
1035                                }
1036                            }
1037                        }
1038
1039                        if runner.record_telemetry_content
1040                            && let Some(choice) = run.accepted_turn_choice()
1041                        {
1042                            rig_core::telemetry::record_model_output(&chat_span, &choice, true);
1043                        }
1044                        break;
1045                    }
1046                }
1047            }
1048        })
1049    }
1050
1051    fn run_tool_calls<'a>(
1052        &'a self,
1053        runner: &'a AgentRunner<M>,
1054        hook_ctx: &'a HookContext,
1055        run: &'a mut AgentRun,
1056        calls: Vec<PendingToolCall>,
1057        tool_snapshot: Arc<ToolRegistrySnapshot>,
1058    ) -> DriveStream<'a, M::Response> {
1059        // The blocking surface chains tool spans into its linear `follows_from`
1060        // sequence (chat -> tool -> chat), and discards the yielded items, so it
1061        // skips building them.
1062        drive_tool_calls(
1063            runner,
1064            hook_ctx,
1065            run,
1066            calls,
1067            tool_snapshot,
1068            |span| self.chain_span(span),
1069            false,
1070        )
1071    }
1072
1073    fn record_run_level_telemetry(
1074        &self,
1075        agent_span: &tracing::Span,
1076        response: &PromptResponse,
1077        created_agent_span: bool,
1078    ) {
1079        // Record run-level completion + usage onto the agent span, but only when
1080        // we created it — never pollute a caller-supplied outer span. The usage
1081        // fields go through the same recorder the streaming surface uses; the
1082        // blocking surface additionally records the final completion text.
1083        if created_agent_span {
1084            if self.record_telemetry_content {
1085                agent_span.record("gen_ai.completion", &response.output);
1086            }
1087            record_usage_on_span(agent_span, response.usage);
1088        }
1089    }
1090
1091    fn final_item(&self, _response: &PromptResponse) -> Option<MultiTurnStreamItem<M::Response>> {
1092        // The blocking surface folds the engine and discards the final item, so
1093        // building it (an extra full-response clone) is skipped entirely.
1094        None
1095    }
1096}
1097
1098impl<M> AgentRunner<M>
1099where
1100    M: CompletionModel,
1101{
1102    pub(crate) async fn run_with_error_usage(
1103        mut self,
1104    ) -> (Result<PromptResponse, PromptError>, Usage) {
1105        let usage = Arc::new(Mutex::new(Usage::new()));
1106        self.error_usage = Some(usage.clone());
1107        let result = self.run().await;
1108        let observed = result.as_ref().map_or_else(
1109            |_| *usage.lock().unwrap_or_else(|error| error.into_inner()),
1110            |response| response.usage,
1111        );
1112        (result, observed)
1113    }
1114
1115    /// Drive the agent loop to completion, returning the aggregated
1116    /// [`PromptResponse`]. Hooks fire at every observable point; the first hook
1117    /// to terminate cancels the run.
1118    pub async fn run(self) -> Result<PromptResponse, PromptError> {
1119        let (agent_span, created_agent_span) = acquire_agent_span(
1120            self.agent_name_or_default(),
1121            self.preamble.as_deref(),
1122            self.record_telemetry_content,
1123        );
1124
1125        if self.record_telemetry_content
1126            && let Some(text) = self.prompt.rag_text()
1127        {
1128            agent_span.record("gen_ai.prompt", text);
1129        }
1130
1131        // When the caller passes explicit history, memory is fully bypassed for
1132        // this run (no load AND no save). Otherwise, if a memory backend and
1133        // conversation id are both configured, load prior history.
1134        let (history_override, memory_handle) = match &self.chat_history {
1135            Some(_) => (None, None),
1136            None => match (&self.memory, &self.conversation_id) {
1137                (Some(memory), Some(id)) => {
1138                    let loaded = memory.load(id).await?;
1139                    (Some(loaded), Some((memory.clone(), id.clone())))
1140                }
1141                _ => (None, None),
1142            },
1143        };
1144
1145        let run = self.build_run(history_override);
1146
1147        // Fold the shared engine to its final response. The blocking surface
1148        // uses a unary model transport and ignores the intermediate items the
1149        // engine yields; the engine is driven under the caller's ambient span
1150        // (no `instrument`), keeping the agent span detached and the chat/tool
1151        // spans on the blocking `follows_from` chain.
1152        let record_telemetry_content = self.record_telemetry_content;
1153        let driver = drive_agent(
1154            self,
1155            UnaryTurnSource::new(record_telemetry_content),
1156            run,
1157            agent_span,
1158            created_agent_span,
1159            memory_handle,
1160            false,
1161        );
1162        futures::pin_mut!(driver);
1163
1164        let mut response = None;
1165        while let Some(item) = driver.next().await {
1166            match item {
1167                Ok(DriveItem::Done(done)) => response = Some(*done),
1168                Ok(DriveItem::Item(_)) => {}
1169                Err(err) => return Err(streaming_error_into_prompt(err)),
1170            }
1171        }
1172
1173        // The engine yields `Done` unless it errored (handled above).
1174        response.ok_or_else(|| {
1175            PromptError::CompletionError(CompletionError::ResponseError(
1176                "agent run ended without producing a final response".to_string(),
1177            ))
1178        })
1179    }
1180}
1181
1182#[cfg(test)]
1183mod tests {
1184    use std::sync::{
1185        Arc, Mutex,
1186        atomic::{AtomicUsize, Ordering},
1187    };
1188
1189    use futures::StreamExt;
1190    use serde_json::json;
1191
1192    use crate::{
1193        agent::{AgentBuilder, AgentHook, HookContext, ToolResultAction, ToolResultEvent},
1194        completion::{CompletionModel, Document},
1195        test_utils::{MockCompletionModel, MockStreamEvent, MockTurn},
1196        tool::{Tool, ToolContext, ToolErrorKind, ToolExecutionError},
1197    };
1198    use rig_core::message::ToolChoice;
1199
1200    struct MetadataFailingTool;
1201
1202    struct SnapshotValue {
1203        value: usize,
1204        clones: Arc<AtomicUsize>,
1205    }
1206
1207    impl Clone for SnapshotValue {
1208        fn clone(&self) -> Self {
1209            self.clones.fetch_add(1, Ordering::SeqCst);
1210            Self {
1211                value: self.value,
1212                clones: self.clones.clone(),
1213            }
1214        }
1215    }
1216
1217    #[derive(Clone, Default)]
1218    struct SnapshotMutatingTool(Arc<Mutex<Vec<usize>>>);
1219
1220    impl Tool for SnapshotMutatingTool {
1221        const NAME: &'static str = "snapshot_mutator";
1222        type Error = rig::tool::ToolExecutionError;
1223        type Args = serde_json::Value;
1224        type Output = String;
1225
1226        fn description(&self) -> String {
1227            "Mutates its per-dispatch context snapshot".into()
1228        }
1229
1230        fn parameters(&self) -> serde_json::Value {
1231            json!({"type": "object", "properties": {}})
1232        }
1233
1234        async fn call(
1235            &self,
1236            context: &mut ToolContext,
1237            _args: Self::Args,
1238        ) -> Result<Self::Output, ToolExecutionError> {
1239            let initial = context.require::<SnapshotValue>()?.value;
1240            self.0.lock().expect("observed values").push(initial);
1241            let updated = {
1242                let value = context
1243                    .get_mut::<SnapshotValue>()
1244                    .expect("required snapshot value");
1245                value.value += 1;
1246                value.value
1247            };
1248            context.insert_result(updated);
1249            Ok(updated.to_string())
1250        }
1251    }
1252
1253    #[derive(Clone, Default)]
1254    struct SnapshotResults(Arc<Mutex<Vec<usize>>>);
1255
1256    impl AgentHook for SnapshotResults {
1257        async fn on_tool_result(
1258            &self,
1259            _ctx: &HookContext,
1260            event: ToolResultEvent<'_>,
1261        ) -> ToolResultAction {
1262            self.0.lock().expect("result values").push(
1263                *event
1264                    .tool_context
1265                    .require_result::<usize>()
1266                    .expect("per-dispatch result metadata"),
1267            );
1268            ToolResultAction::keep()
1269        }
1270    }
1271
1272    impl Tool for MetadataFailingTool {
1273        const NAME: &'static str = "flaky_tool";
1274        type Error = rig::tool::ToolExecutionError;
1275        type Args = serde_json::Value;
1276        type Output = String;
1277
1278        fn description(&self) -> String {
1279            "Fails after attaching result metadata".into()
1280        }
1281
1282        fn parameters(&self) -> serde_json::Value {
1283            json!({"type": "object", "properties": {}})
1284        }
1285
1286        async fn call(
1287            &self,
1288            context: &mut ToolContext,
1289            _args: Self::Args,
1290        ) -> Result<Self::Output, ToolExecutionError> {
1291            context.insert_result("shared-result-metadata".to_string());
1292            Err(ToolExecutionError::timeout("raw timeout failure"))
1293        }
1294    }
1295
1296    #[derive(Clone, Default)]
1297    struct Results(Arc<Mutex<Vec<(ToolErrorKind, String, String)>>>);
1298
1299    impl AgentHook for Results {
1300        async fn on_tool_result(
1301            &self,
1302            _ctx: &HookContext,
1303            event: ToolResultEvent<'_>,
1304        ) -> ToolResultAction {
1305            if let Some(error) = event.raw_result.error() {
1306                self.0.lock().expect("results").push((
1307                    error.kind(),
1308                    event.raw_result.output().render(),
1309                    event
1310                        .tool_context
1311                        .result::<String>()
1312                        .expect("tool result metadata")
1313                        .clone(),
1314                ));
1315            }
1316            ToolResultAction::rewrite("rewritten for model")
1317        }
1318    }
1319
1320    #[test]
1321    fn agent_exposes_read_only_name_and_description() {
1322        let named = AgentBuilder::new(MockCompletionModel::text("done"))
1323            .name("researcher")
1324            .description("Finds evidence")
1325            .build();
1326        assert_eq!(named.name(), Some("researcher"));
1327        assert_eq!(named.description(), Some("Finds evidence"));
1328
1329        let unnamed = AgentBuilder::new(MockCompletionModel::text("done")).build();
1330        assert_eq!(unnamed.name(), None);
1331        assert_eq!(unnamed.description(), None);
1332    }
1333
1334    #[tokio::test]
1335    async fn runner_applies_per_run_request_overrides() {
1336        let model = MockCompletionModel::text("done");
1337        AgentBuilder::new(model.clone())
1338            .preamble("baseline preamble")
1339            .context("baseline document")
1340            .temperature(0.1)
1341            .max_tokens(10)
1342            .additional_params(json!({"baseline": true}))
1343            .build()
1344            .runner("go")
1345            .preamble("run preamble")
1346            .document(Document {
1347                id: "run-one".into(),
1348                text: "first run document".into(),
1349                additional_props: Default::default(),
1350            })
1351            .documents([Document {
1352                id: "run-two".into(),
1353                text: "second run document".into(),
1354                additional_props: Default::default(),
1355            }])
1356            .temperature(0.7)
1357            .max_tokens(42)
1358            .replace_additional_params(json!({"override": true}))
1359            .tool_choice(ToolChoice::None)
1360            .run()
1361            .await
1362            .expect("runner request should succeed");
1363
1364        let requests = model.requests();
1365        let request = requests.first().expect("one request");
1366        assert!(request.chat_history.iter().any(
1367            |message| matches!(message, crate::completion::Message::System { content } if content == "run preamble")
1368        ));
1369        assert!(
1370            request
1371                .documents
1372                .iter()
1373                .any(|document| document.text == "baseline document")
1374        );
1375        assert!(
1376            request
1377                .documents
1378                .iter()
1379                .any(|document| document.id == "run-one")
1380        );
1381        assert!(
1382            request
1383                .documents
1384                .iter()
1385                .any(|document| document.id == "run-two")
1386        );
1387        assert_eq!(request.temperature, Some(0.7));
1388        assert_eq!(request.max_tokens, Some(42));
1389        assert_eq!(request.additional_params, Some(json!({"override": true})));
1390        assert_eq!(request.tool_choice, Some(ToolChoice::None));
1391    }
1392
1393    #[tokio::test]
1394    async fn runner_can_merge_additional_params_into_the_baseline() {
1395        let model = MockCompletionModel::text("done");
1396        AgentBuilder::new(model.clone())
1397            .additional_params(json!({"baseline": true, "winner": "baseline"}))
1398            .build()
1399            .runner("go")
1400            .merge_additional_params(
1401                json!({"override": true, "winner": "runner"})
1402                    .as_object()
1403                    .expect("object")
1404                    .clone(),
1405            )
1406            .run()
1407            .await
1408            .expect("runner request should succeed");
1409
1410        assert_eq!(
1411            model
1412                .requests()
1413                .first()
1414                .expect("one request")
1415                .additional_params,
1416            Some(json!({"baseline": true, "override": true, "winner": "runner"}))
1417        );
1418    }
1419
1420    #[tokio::test]
1421    async fn runner_can_replace_additional_params_wholesale() {
1422        let model = MockCompletionModel::text("done");
1423        AgentBuilder::new(model.clone())
1424            .additional_params(json!({"baseline": true}))
1425            .build()
1426            .runner("go")
1427            .replace_additional_params(json!({"replacement": true}))
1428            .run()
1429            .await
1430            .expect("runner request should succeed");
1431
1432        let requests = model.requests();
1433        let request = requests.first().expect("one request");
1434        assert_eq!(
1435            request.additional_params,
1436            Some(json!({"replacement": true}))
1437        );
1438    }
1439
1440    #[tokio::test]
1441    async fn runner_can_clear_configured_request_defaults() {
1442        let model = MockCompletionModel::text("done");
1443        AgentBuilder::new(model.clone())
1444            .preamble("baseline")
1445            .temperature(0.1)
1446            .max_tokens(10)
1447            .additional_params(json!({"baseline": true}))
1448            .tool_choice(ToolChoice::Required)
1449            .build()
1450            .runner("go")
1451            .without_preamble()
1452            .without_temperature()
1453            .without_max_tokens()
1454            .without_additional_params()
1455            .without_tool_choice()
1456            .run()
1457            .await
1458            .expect("runner request should succeed");
1459
1460        let requests = model.requests();
1461        let request = requests.first().expect("one request");
1462        assert!(
1463            !request
1464                .chat_history
1465                .iter()
1466                .any(|message| matches!(message, crate::completion::Message::System { .. }))
1467        );
1468        assert_eq!(request.temperature, None);
1469        assert_eq!(request.max_tokens, None);
1470        assert_eq!(request.additional_params, None);
1471        assert_eq!(request.tool_choice, None);
1472    }
1473
1474    #[tokio::test]
1475    async fn direct_completion_model_requests_are_intentionally_hook_free() {
1476        #[derive(Clone)]
1477        struct CountCompletionCalls(Arc<AtomicUsize>);
1478
1479        impl AgentHook for CountCompletionCalls {
1480            async fn on_completion_call(
1481                &self,
1482                _ctx: &HookContext,
1483                _event: crate::agent::CompletionCallEvent<'_>,
1484            ) -> crate::agent::CompletionCallAction {
1485                self.0.fetch_add(1, Ordering::SeqCst);
1486                crate::agent::CompletionCallAction::Continue
1487            }
1488        }
1489
1490        let model = MockCompletionModel::text("raw response");
1491        let calls = Arc::new(AtomicUsize::new(0));
1492        let _agent = AgentBuilder::new(model.clone())
1493            .add_hook(CountCompletionCalls(calls.clone()))
1494            .build();
1495
1496        model
1497            .completion_request("raw request")
1498            .send()
1499            .await
1500            .expect("direct model request should succeed");
1501
1502        assert_eq!(calls.load(Ordering::SeqCst), 0);
1503        assert_eq!(model.request_count(), 1);
1504    }
1505
1506    #[tokio::test]
1507    async fn blocking_and_streaming_preserve_raw_failure_while_rewriting_presentation() {
1508        let blocking = Results::default();
1509        let blocking_model = MockCompletionModel::from_turns([
1510            MockTurn::tool_call("tc1", "flaky_tool", json!({})),
1511            MockTurn::text("done"),
1512        ]);
1513        AgentBuilder::new(blocking_model.clone())
1514            .tool(MetadataFailingTool)
1515            .add_hook(blocking.clone())
1516            .build()
1517            .runner("go")
1518            .max_turns(3)
1519            .run()
1520            .await
1521            .expect("blocking run");
1522
1523        let streaming = Results::default();
1524        let streaming_model = MockCompletionModel::from_stream_turns([
1525            vec![
1526                MockStreamEvent::tool_call_name_delta("tc1", "ic1", "flaky_tool"),
1527                MockStreamEvent::tool_call_arguments_delta("tc1", "ic1", "{}"),
1528                MockStreamEvent::tool_call("tc1", "flaky_tool", json!({})),
1529                MockStreamEvent::final_response_with_total_tokens(0),
1530            ],
1531            vec![
1532                MockStreamEvent::text("done"),
1533                MockStreamEvent::final_response_with_total_tokens(0),
1534            ],
1535        ]);
1536        let mut stream = AgentBuilder::new(streaming_model.clone())
1537            .tool(MetadataFailingTool)
1538            .add_hook(streaming.clone())
1539            .build()
1540            .runner("go")
1541            .max_turns(3)
1542            .stream()
1543            .await;
1544        while let Some(item) = stream.next().await {
1545            item.expect("stream item");
1546        }
1547
1548        assert_eq!(*blocking.0.lock().unwrap(), *streaming.0.lock().unwrap());
1549        assert_eq!(
1550            *blocking.0.lock().unwrap(),
1551            vec![(
1552                ToolErrorKind::Timeout,
1553                "raw timeout failure".into(),
1554                "shared-result-metadata".into()
1555            )]
1556        );
1557
1558        let blocking_history = serde_json::to_value(
1559            &blocking_model
1560                .requests()
1561                .get(1)
1562                .expect("second blocking request")
1563                .chat_history,
1564        )
1565        .unwrap();
1566        let streaming_history = serde_json::to_value(
1567            &streaming_model
1568                .requests()
1569                .get(1)
1570                .expect("second streaming request")
1571                .chat_history,
1572        )
1573        .unwrap();
1574        assert_eq!(blocking_history, streaming_history);
1575        let history = blocking_history.to_string();
1576        assert!(history.contains("rewritten for model"));
1577        assert!(!history.contains("raw timeout failure"));
1578    }
1579
1580    #[tokio::test]
1581    async fn agent_dispatch_snapshot_clones_once_and_isolates_tool_mutations() {
1582        let clones = Arc::new(AtomicUsize::new(0));
1583        let mut context = ToolContext::new();
1584        context.insert(SnapshotValue {
1585            value: 0,
1586            clones: clones.clone(),
1587        });
1588        let tool = SnapshotMutatingTool::default();
1589        let results = SnapshotResults::default();
1590
1591        AgentBuilder::new(MockCompletionModel::from_turns([
1592            MockTurn::tool_call("tc1", SnapshotMutatingTool::NAME, json!({})),
1593            MockTurn::tool_call("tc2", SnapshotMutatingTool::NAME, json!({})),
1594            MockTurn::text("done"),
1595        ]))
1596        .tool(tool.clone())
1597        .add_hook(results.clone())
1598        .build()
1599        .runner("go")
1600        .tool_context(context)
1601        .max_turns(4)
1602        .run()
1603        .await
1604        .expect("agent run");
1605
1606        assert_eq!(*tool.0.lock().expect("observed values"), vec![0, 0]);
1607        assert_eq!(*results.0.lock().expect("result values"), vec![1, 1]);
1608        assert_eq!(
1609            clones.load(Ordering::SeqCst),
1610            2,
1611            "each of the two agent dispatches should clone inbound context once"
1612        );
1613    }
1614}
1615
1616#[cfg(test)]
1617#[allow(irrefutable_let_patterns, unreachable_patterns)]
1618mod migrated_tests {
1619    use std::collections::HashMap;
1620
1621    use crate::agent::{
1622        CompletionCallAction, CompletionCallEvent, HookStack, InvalidToolCallAction,
1623        InvalidToolCallContext, ModelTurnAction, ModelTurnFinished, ObservationAction,
1624        StreamResponseFinish, TextDelta, ToolCall, ToolCallAction, ToolCallDelta, ToolResultAction,
1625        ToolResultEvent,
1626    };
1627
1628    use std::sync::{
1629        Arc, Mutex,
1630        atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering::SeqCst},
1631    };
1632
1633    use futures::StreamExt;
1634    use serde::Deserialize;
1635    use serde_json::json;
1636    use tokio::sync::{Barrier, Notify};
1637
1638    use crate::agent::AgentBuilder;
1639    use crate::agent::hook::{AgentHook, HookContext, RequestPatch, StepEventKind};
1640    use crate::agent::prompt_request::streaming::{MultiTurnStreamItem, StreamingError};
1641    use crate::agent::run::OutputMode;
1642    use crate::completion::{
1643        CompletionError, CompletionModel, Message, Prompt, PromptError, Usage,
1644    };
1645    use crate::streaming::{StreamedAssistantContent, StreamedUserContent, StreamingPrompt};
1646    use crate::test_utils::{
1647        MockAddTool, MockBarrierTool, MockCompletionModel, MockOperationArgs, MockStreamEvent,
1648        MockSubtractTool, MockToolError, MockTurn,
1649    };
1650    use crate::tool::{
1651        Tool, ToolContext, ToolExecutionError, ToolSet,
1652        server::{ToolServer, ToolServerHandle},
1653    };
1654    use rig_core::OneOrMany;
1655    use rig_core::message::{
1656        AssistantContent, ToolCall as MessageToolCall, ToolChoice, ToolFunction, UserContent,
1657    };
1658    use rig_core::vector_store::{
1659        VectorSearchRequest, VectorStoreError, VectorStoreIndex, request::Filter,
1660    };
1661    use rig_core::wasm_compat::WasmCompatSend;
1662
1663    /// Records the kind of every hook event (and every tool-result payload) so a
1664    /// run() and a stream() of the same scenario can be compared.
1665    #[derive(Clone, Default)]
1666    struct RecordingHook {
1667        events: Arc<Mutex<Vec<StepEventKind>>>,
1668        tool_results: Arc<Mutex<Vec<String>>>,
1669    }
1670
1671    impl RecordingHook {
1672        /// Event kinds that should be identical across streaming and
1673        /// non-streaming (excludes the medium-specific delta / response-finish
1674        /// events).
1675        fn shared_events(&self) -> Vec<StepEventKind> {
1676            self.events
1677                .lock()
1678                .expect("events lock")
1679                .iter()
1680                .copied()
1681                .filter(|kind| {
1682                    matches!(
1683                        kind,
1684                        StepEventKind::CompletionCall
1685                            | StepEventKind::ToolCall
1686                            | StepEventKind::ToolResult
1687                            | StepEventKind::InvalidToolCall
1688                    )
1689                })
1690                .collect()
1691        }
1692
1693        fn tool_results(&self) -> Vec<String> {
1694            self.tool_results.lock().expect("results lock").clone()
1695        }
1696
1697        /// Count of a single event kind across the whole run, including the
1698        /// medium-specific response-finish events that `shared_events` excludes.
1699        fn count(&self, kind: StepEventKind) -> usize {
1700            self.events
1701                .lock()
1702                .expect("events lock")
1703                .iter()
1704                .filter(|recorded| **recorded == kind)
1705                .count()
1706        }
1707    }
1708
1709    impl RecordingHook {
1710        fn record(&self, kind: StepEventKind) {
1711            self.events.lock().expect("events lock").push(kind);
1712        }
1713    }
1714
1715    impl AgentHook for RecordingHook {
1716        async fn on_completion_call(
1717            &self,
1718            _: &HookContext,
1719            _: CompletionCallEvent<'_>,
1720        ) -> CompletionCallAction {
1721            self.record(StepEventKind::CompletionCall);
1722            CompletionCallAction::continue_run()
1723        }
1724        async fn on_completion_response(
1725            &self,
1726            _: &HookContext,
1727            _: crate::agent::hook::CompletionResponse<'_>,
1728        ) -> ObservationAction {
1729            self.record(StepEventKind::CompletionResponse);
1730            ObservationAction::continue_run()
1731        }
1732        async fn on_model_turn_finished(
1733            &self,
1734            _: &HookContext,
1735            _: ModelTurnFinished<'_>,
1736        ) -> ModelTurnAction {
1737            self.record(StepEventKind::ModelTurnFinished);
1738            ModelTurnAction::continue_run()
1739        }
1740        async fn on_invalid_tool_call(
1741            &self,
1742            _: &HookContext,
1743            _: &InvalidToolCallContext,
1744        ) -> Option<InvalidToolCallAction> {
1745            self.record(StepEventKind::InvalidToolCall);
1746            None
1747        }
1748        async fn on_tool_call(&self, _: &HookContext, _: ToolCall<'_>) -> ToolCallAction {
1749            self.record(StepEventKind::ToolCall);
1750            ToolCallAction::run()
1751        }
1752        async fn on_tool_result(
1753            &self,
1754            _: &HookContext,
1755            event: ToolResultEvent<'_>,
1756        ) -> ToolResultAction {
1757            self.record(StepEventKind::ToolResult);
1758            self.tool_results
1759                .lock()
1760                .expect("results lock")
1761                .push(event.presentation.render());
1762            ToolResultAction::keep()
1763        }
1764        async fn on_text_delta(&self, _: &HookContext, _: TextDelta<'_>) -> ObservationAction {
1765            self.record(StepEventKind::TextDelta);
1766            ObservationAction::continue_run()
1767        }
1768        async fn on_tool_call_delta(
1769            &self,
1770            _: &HookContext,
1771            _: ToolCallDelta<'_>,
1772        ) -> ObservationAction {
1773            self.record(StepEventKind::ToolCallDelta);
1774            ObservationAction::continue_run()
1775        }
1776        async fn on_stream_response_finish(
1777            &self,
1778            _: &HookContext,
1779            _: StreamResponseFinish<'_>,
1780        ) -> ObservationAction {
1781            self.record(StepEventKind::StreamResponseFinish);
1782            ObservationAction::continue_run()
1783        }
1784    }
1785
1786    #[derive(Clone, Debug, PartialEq)]
1787    struct CanonicalResponseSnapshot {
1788        prompt: Message,
1789        content: OneOrMany<AssistantContent>,
1790        usage: Usage,
1791        message_id: Option<String>,
1792    }
1793
1794    #[derive(Clone, Default)]
1795    struct CanonicalResponseHook {
1796        blocking: Arc<Mutex<Vec<CanonicalResponseSnapshot>>>,
1797        streaming: Arc<Mutex<Vec<CanonicalResponseSnapshot>>>,
1798        committed: Arc<Mutex<Vec<OneOrMany<AssistantContent>>>>,
1799    }
1800
1801    impl AgentHook for CanonicalResponseHook {
1802        async fn on_completion_response(
1803            &self,
1804            _ctx: &HookContext,
1805            event: crate::agent::hook::CompletionResponse<'_>,
1806        ) -> ObservationAction {
1807            self.blocking
1808                .lock()
1809                .expect("blocking snapshots")
1810                .push(CanonicalResponseSnapshot {
1811                    prompt: event.prompt.clone(),
1812                    content: event.content.clone(),
1813                    usage: event.usage,
1814                    message_id: event.message_id.map(str::to_owned),
1815                });
1816            ObservationAction::continue_run()
1817        }
1818
1819        async fn on_stream_response_finish(
1820            &self,
1821            _ctx: &HookContext,
1822            event: StreamResponseFinish<'_>,
1823        ) -> ObservationAction {
1824            self.streaming
1825                .lock()
1826                .expect("streaming snapshots")
1827                .push(CanonicalResponseSnapshot {
1828                    prompt: event.prompt.clone(),
1829                    content: event.content.clone(),
1830                    usage: event.usage,
1831                    message_id: event.message_id.map(str::to_owned),
1832                });
1833            ObservationAction::continue_run()
1834        }
1835
1836        async fn on_model_turn_finished(
1837            &self,
1838            _ctx: &HookContext,
1839            event: ModelTurnFinished<'_>,
1840        ) -> ModelTurnAction {
1841            self.committed
1842                .lock()
1843                .expect("committed snapshots")
1844                .push(event.content.clone());
1845            ModelTurnAction::continue_run()
1846        }
1847    }
1848
1849    #[derive(Clone, Default)]
1850    struct FinishLifecycleHook {
1851        snapshots: Arc<Mutex<Vec<CanonicalResponseSnapshot>>>,
1852        model_turns: Arc<AtomicU32>,
1853        stop: Arc<AtomicBool>,
1854    }
1855
1856    impl FinishLifecycleHook {
1857        fn stopping() -> Self {
1858            let hook = Self::default();
1859            hook.stop.store(true, SeqCst);
1860            hook
1861        }
1862    }
1863
1864    impl AgentHook for FinishLifecycleHook {
1865        async fn on_stream_response_finish(
1866            &self,
1867            _ctx: &HookContext,
1868            event: StreamResponseFinish<'_>,
1869        ) -> ObservationAction {
1870            self.snapshots
1871                .lock()
1872                .expect("finish snapshots")
1873                .push(CanonicalResponseSnapshot {
1874                    prompt: event.prompt.clone(),
1875                    content: event.content.clone(),
1876                    usage: event.usage,
1877                    message_id: event.message_id.map(str::to_owned),
1878                });
1879            if self.stop.load(SeqCst) {
1880                ObservationAction::stop("stop at stream EOF")
1881            } else {
1882                ObservationAction::continue_run()
1883            }
1884        }
1885
1886        async fn on_model_turn_finished(
1887            &self,
1888            _ctx: &HookContext,
1889            _event: ModelTurnFinished<'_>,
1890        ) -> ModelTurnAction {
1891            self.model_turns.fetch_add(1, SeqCst);
1892            ModelTurnAction::continue_run()
1893        }
1894    }
1895
1896    fn canonical_usage() -> Usage {
1897        Usage {
1898            input_tokens: 11,
1899            output_tokens: 7,
1900            total_tokens: 18,
1901            ..Usage::new()
1902        }
1903    }
1904
1905    #[tokio::test]
1906    async fn blocking_completion_response_hook_receives_canonical_fields() {
1907        let hook = CanonicalResponseHook::default();
1908        let prompt = Message::user("canonical prompt");
1909        AgentBuilder::new(MockCompletionModel::new([MockTurn::text(
1910            "canonical response",
1911        )
1912        .with_usage(canonical_usage())
1913        .with_message_id("msg-canonical")]))
1914        .add_hook(hook.clone())
1915        .build()
1916        .runner(prompt.clone())
1917        .run()
1918        .await
1919        .expect("blocking response");
1920
1921        assert_eq!(
1922            *hook.blocking.lock().expect("blocking snapshots"),
1923            [CanonicalResponseSnapshot {
1924                prompt,
1925                content: OneOrMany::one(AssistantContent::text("canonical response")),
1926                usage: canonical_usage(),
1927                message_id: Some("msg-canonical".to_string()),
1928            }]
1929        );
1930    }
1931
1932    #[tokio::test]
1933    async fn streaming_response_finish_matches_blocking_canonical_fields() {
1934        let prompt = Message::user("canonical prompt");
1935        let blocking_hook = CanonicalResponseHook::default();
1936        AgentBuilder::new(MockCompletionModel::new([MockTurn::text(
1937            "canonical response",
1938        )
1939        .with_usage(canonical_usage())
1940        .with_message_id("msg-canonical")]))
1941        .add_hook(blocking_hook.clone())
1942        .build()
1943        .runner(prompt.clone())
1944        .run()
1945        .await
1946        .expect("blocking response");
1947
1948        let streaming_hook = CanonicalResponseHook::default();
1949        let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([[
1950            MockStreamEvent::text("canonical response"),
1951            MockStreamEvent::final_response(canonical_usage()),
1952            MockStreamEvent::message_id("msg-canonical"),
1953        ]]))
1954        .add_hook(streaming_hook.clone())
1955        .build()
1956        .runner(prompt)
1957        .stream()
1958        .await;
1959        while let Some(item) = stream.next().await {
1960            item.expect("stream item");
1961        }
1962
1963        let blocking = blocking_hook
1964            .blocking
1965            .lock()
1966            .expect("blocking snapshots")
1967            .clone();
1968        let streaming = streaming_hook
1969            .streaming
1970            .lock()
1971            .expect("streaming snapshots")
1972            .clone();
1973        assert_eq!(streaming, blocking);
1974        assert_eq!(streaming[0].usage, canonical_usage());
1975        assert_eq!(streaming[0].message_id.as_deref(), Some("msg-canonical"));
1976    }
1977
1978    #[tokio::test]
1979    async fn streaming_response_finish_without_provider_message_id_reports_none() {
1980        let hook = FinishLifecycleHook::default();
1981        let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([[
1982            MockStreamEvent::text("canonical response"),
1983            MockStreamEvent::final_response(canonical_usage()),
1984        ]]))
1985        .add_hook(hook.clone())
1986        .build()
1987        .runner("canonical prompt")
1988        .stream()
1989        .await;
1990        while let Some(item) = stream.next().await {
1991            item.expect("stream item");
1992        }
1993
1994        let snapshots = hook.snapshots.lock().expect("finish snapshots");
1995        assert_eq!(snapshots.len(), 1);
1996        assert_eq!(snapshots[0].message_id, None);
1997    }
1998
1999    #[tokio::test]
2000    async fn streaming_response_finish_runs_before_buffered_final_is_exposed() {
2001        let hook = FinishLifecycleHook::default();
2002        let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([[
2003            MockStreamEvent::text("canonical response"),
2004            MockStreamEvent::final_response(canonical_usage()),
2005            MockStreamEvent::message_id("msg-after-final"),
2006        ]]))
2007        .add_hook(hook.clone())
2008        .build()
2009        .runner("canonical prompt")
2010        .stream()
2011        .await;
2012        let mut provider_finals = 0;
2013        while let Some(item) = stream.next().await {
2014            if matches!(
2015                item.expect("stream item"),
2016                MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Final(_))
2017            ) {
2018                provider_finals += 1;
2019                let snapshots = hook.snapshots.lock().expect("finish snapshots");
2020                assert_eq!(snapshots.len(), 1, "hook must run before final exposure");
2021                assert_eq!(snapshots[0].message_id.as_deref(), Some("msg-after-final"));
2022                assert_eq!(
2023                    hook.model_turns.load(SeqCst),
2024                    1,
2025                    "the canonical turn hook must accept the turn before final exposure"
2026                );
2027            }
2028        }
2029
2030        assert_eq!(provider_finals, 1);
2031        assert_eq!(hook.snapshots.lock().expect("finish snapshots").len(), 1);
2032        assert_eq!(hook.model_turns.load(SeqCst), 1);
2033    }
2034
2035    #[tokio::test]
2036    async fn streaming_response_finish_stop_suppresses_final_and_turn_commit() {
2037        let hook = FinishLifecycleHook::stopping();
2038        let prompt = Message::user("canonical prompt");
2039        let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([[
2040            MockStreamEvent::text("canonical response"),
2041            MockStreamEvent::final_response(canonical_usage()),
2042            MockStreamEvent::message_id("msg-after-final"),
2043        ]]))
2044        .add_hook(hook.clone())
2045        .build()
2046        .runner(prompt.clone())
2047        .stream()
2048        .await;
2049        let mut saw_provider_final = false;
2050        let mut saw_run_final = false;
2051        let mut error = None;
2052        while let Some(item) = stream.next().await {
2053            match item {
2054                Ok(MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Final(
2055                    _,
2056                ))) => saw_provider_final = true,
2057                Ok(MultiTurnStreamItem::FinalResponse(_)) => saw_run_final = true,
2058                Ok(_) => {}
2059                Err(err) => error = Some(err),
2060            }
2061        }
2062
2063        assert!(!saw_provider_final, "the buffered final must remain hidden");
2064        assert!(
2065            !saw_run_final,
2066            "the cancelled run must not produce a response"
2067        );
2068        assert_eq!(hook.snapshots.lock().expect("finish snapshots").len(), 1);
2069        assert_eq!(hook.model_turns.load(SeqCst), 0);
2070        assert!(matches!(
2071            error,
2072            Some(StreamingError::Prompt(error))
2073                if matches!(
2074                    error.as_ref(),
2075                    PromptError::PromptCancelled { chat_history, reason }
2076                        if chat_history == &[prompt] && reason == "stop at stream EOF"
2077                )
2078        ));
2079    }
2080
2081    struct StopCompletedModelTurn;
2082
2083    impl AgentHook for StopCompletedModelTurn {
2084        async fn on_model_turn_finished(
2085            &self,
2086            _ctx: &HookContext,
2087            _event: ModelTurnFinished<'_>,
2088        ) -> ModelTurnAction {
2089            ModelTurnAction::stop("stop completed model turn")
2090        }
2091    }
2092
2093    #[tokio::test]
2094    async fn streaming_model_turn_stop_preserves_completed_provider_final() {
2095        let prompt = Message::user("canonical prompt");
2096        let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([[
2097            MockStreamEvent::text("canonical response"),
2098            MockStreamEvent::final_response(canonical_usage()),
2099        ]]))
2100        .add_hook(StopCompletedModelTurn)
2101        .build()
2102        .runner(prompt.clone())
2103        .stream()
2104        .await;
2105
2106        let mut provider_finals = 0;
2107        let mut saw_retry = false;
2108        let mut saw_run_final = false;
2109        let mut error = None;
2110        while let Some(item) = stream.next().await {
2111            match item {
2112                Ok(MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Final(
2113                    _,
2114                ))) => provider_finals += 1,
2115                Ok(MultiTurnStreamItem::ModelTurnRetried { .. }) => saw_retry = true,
2116                Ok(MultiTurnStreamItem::FinalResponse(_)) => saw_run_final = true,
2117                Ok(_) => {}
2118                Err(err) => error = Some(err),
2119            }
2120        }
2121
2122        assert_eq!(provider_finals, 1);
2123        assert!(!saw_retry);
2124        assert!(!saw_run_final);
2125        assert!(matches!(
2126            error,
2127            Some(StreamingError::Prompt(error))
2128                if matches!(
2129                    error.as_ref(),
2130                    PromptError::PromptCancelled { reason, .. }
2131                        if reason == "stop completed model turn"
2132                )
2133        ));
2134    }
2135
2136    #[tokio::test]
2137    async fn provider_error_after_final_suppresses_finish_hook_and_buffered_final() {
2138        let hook = FinishLifecycleHook::default();
2139        let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([[
2140            MockStreamEvent::text("canonical response"),
2141            MockStreamEvent::final_response(canonical_usage()),
2142            MockStreamEvent::error("post-final failure"),
2143        ]]))
2144        .add_hook(hook.clone())
2145        .build()
2146        .runner("canonical prompt")
2147        .stream()
2148        .await;
2149        let mut saw_provider_final = false;
2150        let mut error = None;
2151        while let Some(item) = stream.next().await {
2152            match item {
2153                Ok(MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Final(
2154                    _,
2155                ))) => saw_provider_final = true,
2156                Ok(_) => {}
2157                Err(err) => error = Some(err),
2158            }
2159        }
2160
2161        assert!(!saw_provider_final, "the buffered final must remain hidden");
2162        assert!(hook.snapshots.lock().expect("finish snapshots").is_empty());
2163        assert_eq!(hook.model_turns.load(SeqCst), 0);
2164        assert!(matches!(
2165            error,
2166            Some(StreamingError::Completion(CompletionError::ProviderError(message)))
2167                if message == "post-final failure"
2168        ));
2169    }
2170
2171    #[tokio::test]
2172    async fn visible_assistant_items_after_final_are_rejected() {
2173        let cases = [
2174            ("text", MockStreamEvent::text("late text")),
2175            ("reasoning", MockStreamEvent::reasoning("late reasoning")),
2176            (
2177                "reasoning delta",
2178                MockStreamEvent::reasoning_delta(None::<String>, "late reasoning"),
2179            ),
2180            (
2181                "tool call",
2182                MockStreamEvent::tool_call("late", "add", json!({"x": 1, "y": 2})),
2183            ),
2184            (
2185                "tool-call delta",
2186                MockStreamEvent::tool_call_name_delta("late", "internal-late", "add"),
2187            ),
2188            ("unknown", MockStreamEvent::unknown(json!({"type": "late"}))),
2189        ];
2190
2191        for (case, visible_item) in cases {
2192            let hook = FinishLifecycleHook::default();
2193            let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([vec![
2194                MockStreamEvent::text("canonical response"),
2195                MockStreamEvent::final_response(canonical_usage()),
2196                visible_item,
2197            ]]))
2198            .add_hook(hook.clone())
2199            .build()
2200            .runner("canonical prompt")
2201            .stream()
2202            .await;
2203            let mut saw_provider_final = false;
2204            let mut error = None;
2205            while let Some(item) = stream.next().await {
2206                match item {
2207                    Ok(MultiTurnStreamItem::StreamAssistantItem(
2208                        StreamedAssistantContent::Final(_),
2209                    )) => saw_provider_final = true,
2210                    Ok(_) => {}
2211                    Err(err) => error = Some(err),
2212                }
2213            }
2214
2215            assert!(
2216                !saw_provider_final,
2217                "{case}: buffered final must remain hidden"
2218            );
2219            assert!(
2220                hook.snapshots.lock().expect("finish snapshots").is_empty(),
2221                "{case}: finish hook must not run"
2222            );
2223            assert_eq!(hook.model_turns.load(SeqCst), 0, "{case}");
2224            assert!(
2225                matches!(
2226                    error,
2227                    Some(StreamingError::Completion(CompletionError::ResponseError(ref message)))
2228                        if message.contains("visible assistant content after its final response")
2229                ),
2230                "{case}: expected malformed-response error, got {error:?}"
2231            );
2232        }
2233    }
2234
2235    #[tokio::test]
2236    async fn visible_item_after_non_emittable_final_is_rejected() {
2237        let hook = FinishLifecycleHook::default();
2238        let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([[
2239            MockStreamEvent::reasoning("think"),
2240            MockStreamEvent::final_response(canonical_usage()),
2241            MockStreamEvent::text("late text"),
2242        ]]))
2243        .add_hook(hook.clone())
2244        .build()
2245        .runner("canonical prompt")
2246        .stream()
2247        .await;
2248        let mut error = None;
2249        while let Some(item) = stream.next().await {
2250            if let Err(err) = item {
2251                error = Some(err);
2252            }
2253        }
2254
2255        assert!(hook.snapshots.lock().expect("finish snapshots").is_empty());
2256        assert_eq!(hook.model_turns.load(SeqCst), 0);
2257        assert!(matches!(
2258            error,
2259            Some(StreamingError::Completion(CompletionError::ResponseError(message)))
2260                if message.contains("visible assistant content after its final response")
2261        ));
2262    }
2263
2264    #[tokio::test]
2265    async fn streaming_response_finish_normalizes_interleaved_content() {
2266        let hook = CanonicalResponseHook::default();
2267        let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([
2268            vec![
2269                MockStreamEvent::reasoning("think"),
2270                MockStreamEvent::tool_call("tc1", "add", json!({"x": 2, "y": 3})),
2271                MockStreamEvent::text("answer"),
2272                MockStreamEvent::final_response_with_total_tokens(0),
2273            ],
2274            vec![
2275                MockStreamEvent::text("done"),
2276                MockStreamEvent::final_response_with_total_tokens(0),
2277            ],
2278        ]))
2279        .tool(MockAddTool)
2280        .add_hook(hook.clone())
2281        .build()
2282        .runner("go")
2283        .max_turns(3)
2284        .stream()
2285        .await;
2286        while let Some(item) = stream.next().await {
2287            item.expect("stream item");
2288        }
2289
2290        let snapshots = hook.streaming.lock().expect("streaming snapshots");
2291        let committed = hook.committed.lock().expect("committed snapshots");
2292        let kinds = snapshots[0]
2293            .content
2294            .iter()
2295            .map(|content| match content {
2296                AssistantContent::Reasoning(_) => "reasoning",
2297                AssistantContent::Text(_) => "text",
2298                AssistantContent::ToolCall(_) => "tool_call",
2299                _ => "other",
2300            })
2301            .collect::<Vec<_>>();
2302        assert_eq!(kinds, ["reasoning", "text", "tool_call"]);
2303        assert_eq!(
2304            snapshots[0].content, committed[0],
2305            "finish hook and committed turn must share one canonical choice"
2306        );
2307    }
2308
2309    fn blocking_model() -> MockCompletionModel {
2310        MockCompletionModel::from_turns([
2311            MockTurn::tool_call("tc1", "add", json!({"x": 2, "y": 3})),
2312            MockTurn::text("the answer is 5"),
2313        ])
2314    }
2315
2316    fn streaming_model() -> MockCompletionModel {
2317        MockCompletionModel::from_stream_turns([
2318            vec![
2319                MockStreamEvent::tool_call_name_delta("tc1", "ic1", "add"),
2320                MockStreamEvent::tool_call_arguments_delta("tc1", "ic1", "{\"x\":2,\"y\":3}"),
2321                MockStreamEvent::tool_call("tc1", "add", json!({"x": 2, "y": 3})),
2322                MockStreamEvent::final_response_with_total_tokens(0),
2323            ],
2324            vec![
2325                MockStreamEvent::text("the answer is 5"),
2326                MockStreamEvent::final_response_with_total_tokens(0),
2327            ],
2328        ])
2329    }
2330
2331    /// `AgentRunner::from_agent` preserves the distinction between an absent
2332    /// agent default (the implicit one-call budget) and an explicit zero budget.
2333    #[tokio::test]
2334    async fn from_agent_preserves_implicit_one_and_explicit_zero_budgets() {
2335        let implicit_model = blocking_model();
2336        let implicit_recorded = implicit_model.clone();
2337        let implicit_agent = AgentBuilder::new(implicit_model).tool(MockAddTool).build();
2338        let implicit_runner = super::AgentRunner::from_agent(&implicit_agent, "add 2 and 3");
2339        assert_eq!(implicit_runner.max_turns, 1);
2340
2341        let implicit_err = implicit_runner
2342            .run()
2343            .await
2344            .expect_err("implicit budget should reject the second model call");
2345        assert!(matches!(
2346            implicit_err,
2347            PromptError::MaxTurnsError { max_turns: 1, .. }
2348        ));
2349        assert_eq!(implicit_recorded.request_count(), 1);
2350
2351        let zero_model = MockCompletionModel::text("should not be requested");
2352        let zero_recorded = zero_model.clone();
2353        let zero_agent = AgentBuilder::new(zero_model).default_max_turns(0).build();
2354        let zero_runner = super::AgentRunner::from_agent(&zero_agent, "do not call");
2355        assert_eq!(zero_runner.max_turns, 0);
2356
2357        let zero_err = zero_runner
2358            .run()
2359            .await
2360            .expect_err("explicit zero budget should reject the initial model call");
2361        assert!(matches!(
2362            zero_err,
2363            PromptError::MaxTurnsError { max_turns: 0, .. }
2364        ));
2365        assert_eq!(zero_recorded.request_count(), 0);
2366    }
2367
2368    /// The public blocking and streaming prompt surfaces enforce the one-call
2369    /// boundary identically after executing a tool-producing first turn.
2370    #[tokio::test]
2371    async fn prompt_surfaces_reject_second_tool_roundtrip_request_at_budget_one() {
2372        let blocking_model = blocking_model();
2373        let blocking_recorded = blocking_model.clone();
2374        let blocking_agent = AgentBuilder::new(blocking_model).tool(MockAddTool).build();
2375        let blocking_err = blocking_agent
2376            .prompt("add 2 and 3")
2377            .max_turns(1)
2378            .await
2379            .expect_err("blocking prompt should reject request two");
2380        assert!(matches!(
2381            blocking_err,
2382            PromptError::MaxTurnsError { max_turns: 1, .. }
2383        ));
2384        assert_eq!(blocking_recorded.request_count(), 1);
2385
2386        let streaming_model = streaming_model();
2387        let streaming_recorded = streaming_model.clone();
2388        let streaming_agent = AgentBuilder::new(streaming_model).tool(MockAddTool).build();
2389        let mut stream = streaming_agent
2390            .stream_prompt("add 2 and 3")
2391            .max_turns(1)
2392            .await;
2393        let mut streaming_err = None;
2394        while let Some(item) = stream.next().await {
2395            if let Err(err) = item {
2396                streaming_err = Some(err);
2397                break;
2398            }
2399        }
2400        match streaming_err {
2401            Some(StreamingError::Prompt(err)) => assert!(matches!(
2402                *err,
2403                PromptError::MaxTurnsError { max_turns: 1, .. }
2404            )),
2405            other => panic!("expected streaming max-turns error, got {other:?}"),
2406        }
2407        assert_eq!(streaming_recorded.request_count(), 1);
2408    }
2409
2410    /// run() and stream() of the same tool-calling scenario produce the same
2411    /// final output, the same final message history, the same tool-result
2412    /// content, and the same medium-independent hook event sequence.
2413    #[tokio::test]
2414    async fn run_and_stream_behave_identically_for_a_tool_call() {
2415        let blocking_hook = RecordingHook::default();
2416        let blocking = AgentBuilder::new(blocking_model())
2417            .tool(MockAddTool)
2418            .build()
2419            .runner("add 2 and 3")
2420            .max_turns(2)
2421            .add_hook(blocking_hook.clone())
2422            .run()
2423            .await
2424            .expect("blocking run should succeed");
2425
2426        // No `.with_history` on either runner — `stream()` must return the final
2427        // history just like `run()` returns `messages`.
2428        let streaming_hook = RecordingHook::default();
2429        let mut stream = AgentBuilder::new(streaming_model())
2430            .tool(MockAddTool)
2431            .build()
2432            .runner("add 2 and 3")
2433            .max_turns(2)
2434            .add_hook(streaming_hook.clone())
2435            .stream()
2436            .await;
2437
2438        let mut final_response = None;
2439        while let Some(item) = stream.next().await {
2440            if let Ok(MultiTurnStreamItem::FinalResponse(resp)) =
2441                item.map_err(|err| panic!("stream item errored: {err}"))
2442            {
2443                final_response = Some(resp);
2444            }
2445        }
2446        let final_response = final_response.expect("stream should yield a final response");
2447
2448        // Same final output.
2449        assert_eq!(blocking.output, "the answer is 5");
2450        assert_eq!(final_response.output(), blocking.output);
2451
2452        // Same medium-independent hook event sequence (model call, tool call,
2453        // tool result, second model call).
2454        assert_eq!(
2455            blocking_hook.shared_events(),
2456            streaming_hook.shared_events()
2457        );
2458        assert_eq!(
2459            blocking_hook.shared_events(),
2460            vec![
2461                StepEventKind::CompletionCall,
2462                StepEventKind::ToolCall,
2463                StepEventKind::ToolResult,
2464                StepEventKind::CompletionCall,
2465            ]
2466        );
2467
2468        // Same tool-result content seen by the hook.
2469        assert_eq!(blocking_hook.tool_results(), streaming_hook.tool_results());
2470        assert_eq!(blocking_hook.tool_results(), vec!["5".to_string()]);
2471
2472        // Same final message history (compared via serialized form to normalize).
2473        let blocking_messages = blocking.messages.expect("blocking messages");
2474        let streaming_messages = final_response
2475            .messages()
2476            .expect("streaming history")
2477            .to_vec();
2478        assert_eq!(
2479            serde_json::to_value(&blocking_messages).expect("serialize blocking"),
2480            serde_json::to_value(&streaming_messages).expect("serialize streaming"),
2481        );
2482    }
2483
2484    /// Structured tool-execution results reach `ToolResultEvent` as machine
2485    /// metadata (error/refusal state plus result context), on both the blocking and streaming paths,
2486    /// so hooks can steer on a classified failure without parsing the result
2487    /// string.
2488    mod structured_tool_results {
2489        use std::sync::{Arc, Mutex};
2490
2491        use futures::StreamExt;
2492        use serde_json::json;
2493
2494        use crate::agent::{
2495            AgentBuilder, AgentHook, HookContext, HookStack, ToolCall, ToolCallAction,
2496            ToolResultAction, ToolResultEvent,
2497        };
2498        use crate::test_utils::{
2499            MockAddTool, MockCompletionModel, MockDeniedTool, MockFailingTool,
2500            MockHandledFailureTool, MockMetadataTool, MockRequestId, MockStreamEvent, MockTurn,
2501        };
2502        use crate::tool::{ToolErrorKind, ToolResult};
2503
2504        /// Records, for every `ToolResult` event, a compact outcome label and the
2505        /// model-visible result string — the machine metadata a policy reads.
2506        #[derive(Clone, Default)]
2507        struct OutcomeHook {
2508            outcomes: Arc<Mutex<Vec<String>>>,
2509            results: Arc<Mutex<Vec<String>>>,
2510        }
2511
2512        impl OutcomeHook {
2513            fn outcomes(&self) -> Vec<String> {
2514                self.outcomes.lock().expect("outcomes").clone()
2515            }
2516
2517            fn results(&self) -> Vec<String> {
2518                self.results.lock().expect("results").clone()
2519            }
2520        }
2521
2522        /// A compact string label for an outcome, e.g. `error:timeout`.
2523        fn outcome_label(result: &ToolResult) -> String {
2524            if result.is_skipped() {
2525                "skipped".to_string()
2526            } else if result.is_refused() {
2527                "denied".to_string()
2528            } else if let Some(error) = result.error() {
2529                format!("error:{}", error.kind().as_str())
2530            } else {
2531                "success".to_string()
2532            }
2533        }
2534
2535        impl AgentHook for OutcomeHook {
2536            async fn on_tool_result(
2537                &self,
2538                _ctx: &HookContext,
2539                event: ToolResultEvent<'_>,
2540            ) -> ToolResultAction {
2541                if let ToolResultEvent {
2542                    presentation,
2543                    raw_result,
2544                    ..
2545                } = event
2546                {
2547                    self.outcomes
2548                        .lock()
2549                        .expect("outcomes")
2550                        .push(outcome_label(raw_result));
2551                    self.results
2552                        .lock()
2553                        .expect("results")
2554                        .push(presentation.render());
2555                }
2556                ToolResultAction::keep()
2557            }
2558        }
2559
2560        /// A blocking model that calls `tool` once, then answers.
2561        fn model_one_tool_then_text(tool: &str) -> MockCompletionModel {
2562            MockCompletionModel::from_turns([
2563                MockTurn::tool_call("tc1", tool, json!({})),
2564                MockTurn::text("done"),
2565            ])
2566        }
2567
2568        /// A streaming model that calls `tool` once, then answers.
2569        fn stream_model_one_tool_then_text(tool: &str) -> MockCompletionModel {
2570            MockCompletionModel::from_stream_turns([
2571                vec![
2572                    MockStreamEvent::tool_call_name_delta("tc1", "ic1", tool),
2573                    MockStreamEvent::tool_call_arguments_delta("tc1", "ic1", "{}"),
2574                    MockStreamEvent::tool_call("tc1", tool, json!({})),
2575                    MockStreamEvent::final_response_with_total_tokens(0),
2576                ],
2577                vec![
2578                    MockStreamEvent::text("done"),
2579                    MockStreamEvent::final_response_with_total_tokens(0),
2580                ],
2581            ])
2582        }
2583
2584        // (1) A `Timeout` failure reaches `ToolResultEvent` as structured
2585        // metadata (not just a string), with the model-visible feedback intact.
2586        #[tokio::test]
2587        async fn timeout_failure_surfaces_structured_outcome() {
2588            let hook = OutcomeHook::default();
2589            AgentBuilder::new(model_one_tool_then_text("flaky_tool"))
2590                .tool(MockFailingTool::new(ToolErrorKind::Timeout))
2591                .add_hook(hook.clone())
2592                .build()
2593                .runner("go")
2594                .max_turns(3)
2595                .run()
2596                .await
2597                .expect("run should succeed; a tool timeout is model-visible feedback, not fatal");
2598
2599            assert_eq!(hook.outcomes(), vec!["error:timeout".to_string()]);
2600            // (4) The model still receives useful text for the handled failure.
2601            assert_eq!(hook.results(), vec!["mock tool call failed".to_string()]);
2602        }
2603
2604        // (2) A hook counts timeout failures in the run scratchpad and terminates
2605        // the run after a threshold — the motivating use case.
2606        #[tokio::test]
2607        async fn hook_terminates_after_repeated_timeouts() {
2608            #[derive(Clone, Default)]
2609            struct TimeoutCount(usize);
2610
2611            struct TimeoutTerminator;
2612            impl AgentHook for TimeoutTerminator {
2613                async fn on_tool_result(
2614                    &self,
2615                    ctx: &HookContext,
2616                    event: ToolResultEvent<'_>,
2617                ) -> ToolResultAction {
2618                    if let ToolResultEvent { raw_result, .. } = event
2619                        && raw_result.is_error_kind(ToolErrorKind::Timeout)
2620                    {
2621                        let count = ctx.scratchpad().update(|c: &mut TimeoutCount| {
2622                            c.0 += 1;
2623                            c.0
2624                        });
2625                        if count >= 2 {
2626                            return ToolResultAction::stop("aborting after repeated tool timeouts");
2627                        }
2628                    }
2629                    ToolResultAction::keep()
2630                }
2631            }
2632
2633            let observer = OutcomeHook::default();
2634            let err = AgentBuilder::new(MockCompletionModel::from_turns([
2635                MockTurn::tool_call("tc1", "flaky_tool", json!({})),
2636                MockTurn::tool_call("tc2", "flaky_tool", json!({})),
2637                MockTurn::text("unreachable"),
2638            ]))
2639            .tool(MockFailingTool::new(ToolErrorKind::Timeout))
2640            // Observer first so it records both timeouts before the terminator fires.
2641            .add_hook(observer.clone())
2642            .add_hook(TimeoutTerminator)
2643            .build()
2644            .runner("go")
2645            .max_turns(5)
2646            .run()
2647            .await
2648            .expect_err("the run must terminate after two timeouts");
2649
2650            assert!(
2651                err.to_string()
2652                    .contains("aborting after repeated tool timeouts"),
2653                "unexpected error: {err}"
2654            );
2655            assert_eq!(
2656                observer.outcomes(),
2657                vec!["error:timeout".to_string(), "error:timeout".to_string()],
2658                "both timeout outcomes must be observed before termination"
2659            );
2660        }
2661
2662        // (3) A not-found (404) failure surfaces as structured `NotFound` metadata
2663        // but does not terminate the run by default — the model may try another path.
2664        #[tokio::test]
2665        async fn not_found_outcome_is_structured_and_non_fatal() {
2666            let hook = OutcomeHook::default();
2667            let status: Arc<Mutex<Option<u16>>> = Arc::new(Mutex::new(None));
2668
2669            struct StatusProbe(Arc<Mutex<Option<u16>>>);
2670            impl AgentHook for StatusProbe {
2671                async fn on_tool_result(
2672                    &self,
2673                    _ctx: &HookContext,
2674                    event: ToolResultEvent<'_>,
2675                ) -> ToolResultAction {
2676                    if let Some(error) = event.raw_result.error() {
2677                        *self.0.lock().expect("status") = error.http_status();
2678                    }
2679                    ToolResultAction::keep()
2680                }
2681            }
2682
2683            AgentBuilder::new(model_one_tool_then_text("flaky_tool"))
2684                .tool(MockFailingTool::new(ToolErrorKind::NotFound))
2685                .add_hook(hook.clone())
2686                .add_hook(StatusProbe(status.clone()))
2687                .build()
2688                .runner("go")
2689                .max_turns(3)
2690                .run()
2691                .await
2692                .expect("a 404 must not terminate the run by default");
2693
2694            assert_eq!(hook.outcomes(), vec!["error:not_found".to_string()]);
2695            assert_eq!(
2696                *status.lock().expect("status"),
2697                Some(404),
2698                "the structured failure must carry the HTTP status"
2699            );
2700        }
2701
2702        // (4) A tool that returns a handled failure via ordinary `Result` shows the
2703        // model useful output while the outcome is a classified error.
2704        #[tokio::test]
2705        async fn handled_failure_delivers_model_output_and_error_outcome() {
2706            let hook = OutcomeHook::default();
2707            AgentBuilder::new(model_one_tool_then_text("lookup"))
2708                .tool(MockHandledFailureTool)
2709                .add_hook(hook.clone())
2710                .build()
2711                .runner("go")
2712                .max_turns(3)
2713                .run()
2714                .await
2715                .expect("a handled failure is not fatal");
2716
2717            assert_eq!(hook.outcomes(), vec!["error:not_found".to_string()]);
2718            assert_eq!(
2719                hook.results(),
2720                vec!["no record found for id 42; try a different id".to_string()],
2721                "the tool's model-visible output must survive alongside the error outcome"
2722            );
2723        }
2724
2725        // (7) `ToolCallAction::Skip` on the tool-call produces a structured `Skipped`
2726        // outcome that the result hook observes.
2727        #[tokio::test]
2728        async fn flow_skip_produces_skipped_outcome() {
2729            struct SkipHook;
2730            impl AgentHook for SkipHook {
2731                async fn on_tool_call(
2732                    &self,
2733                    _ctx: &HookContext,
2734                    event: ToolCall<'_>,
2735                ) -> ToolCallAction {
2736                    if let ToolCall { .. } = event {
2737                        ToolCallAction::skip("not executed (denied by policy); do not retry")
2738                    } else {
2739                        ToolCallAction::run()
2740                    }
2741                }
2742            }
2743
2744            let observer = OutcomeHook::default();
2745            AgentBuilder::new(model_one_tool_then_text("flaky_tool"))
2746                .tool(MockFailingTool::new(ToolErrorKind::Timeout))
2747                .add_hook(SkipHook)
2748                .add_hook(observer.clone())
2749                .build()
2750                .runner("go")
2751                .max_turns(3)
2752                .run()
2753                .await
2754                .expect("run should succeed after skipping the tool");
2755
2756            assert_eq!(observer.outcomes(), vec!["skipped".to_string()]);
2757            assert_eq!(
2758                observer.results(),
2759                vec!["not executed (denied by policy); do not retry".to_string()]
2760            );
2761        }
2762
2763        // A *tool-authored* refusal surfaces as a `Denied`
2764        // outcome — distinct from a hook `ToolCallAction::Skip`, which is `Skipped`. This
2765        // pins the documented `Skipped` vs `Denied` split: `Denied` comes only
2766        // from the tool, never from a hook skip.
2767        #[tokio::test]
2768        async fn tool_authored_denial_produces_denied_outcome() {
2769            let hook = OutcomeHook::default();
2770            AgentBuilder::new(model_one_tool_then_text("guarded"))
2771                .tool(MockDeniedTool)
2772                .add_hook(hook.clone())
2773                .build()
2774                .runner("go")
2775                .max_turns(3)
2776                .run()
2777                .await
2778                .expect("a tool-authored denial is not fatal");
2779
2780            assert_eq!(hook.outcomes(), vec!["denied".to_string()]);
2781            assert_eq!(
2782                hook.results(),
2783                vec!["access to this resource is not permitted".to_string()],
2784                "the model still receives the tool's denial message"
2785            );
2786        }
2787
2788        #[tokio::test]
2789        async fn permission_denied_failure_is_not_a_tool_refusal() {
2790            let hook = OutcomeHook::default();
2791            AgentBuilder::new(model_one_tool_then_text("flaky_tool"))
2792                .tool(MockFailingTool::new(ToolErrorKind::PermissionDenied))
2793                .add_hook(hook.clone())
2794                .build()
2795                .runner("go")
2796                .max_turns(3)
2797                .run()
2798                .await
2799                .expect("a permission failure is model-visible feedback, not fatal");
2800
2801            assert_eq!(hook.outcomes(), vec!["error:permission_denied".to_string()]);
2802            assert_eq!(hook.results(), vec!["mock tool call failed".to_string()]);
2803        }
2804
2805        // A `ToolCallAction::Rewrite` hook followed by a `Skip` hook: the tool must not run,
2806        // the `ToolResult` reports the *rewritten* args (not the model's
2807        // original), and the outcome is `Skipped` — the rewrite (e.g. a
2808        // redaction) is not lost when a later hook short-circuits. Verified on
2809        // both the blocking and streaming surfaces.
2810        #[tokio::test]
2811        async fn rewrite_args_then_skip_reports_rewritten_args() {
2812            // Rewrites the tool args, replacing whatever the model emitted.
2813            struct RewriteHook;
2814            impl AgentHook for RewriteHook {
2815                async fn on_tool_call(
2816                    &self,
2817                    _ctx: &HookContext,
2818                    event: ToolCall<'_>,
2819                ) -> ToolCallAction {
2820                    if let ToolCall { .. } = event {
2821                        ToolCallAction::rewrite(json!({ "x": 41, "y": 1 }))
2822                    } else {
2823                        ToolCallAction::run()
2824                    }
2825                }
2826            }
2827            // Skips *after* the rewrite (registered second).
2828            struct SkipHook;
2829            impl AgentHook for SkipHook {
2830                async fn on_tool_call(
2831                    &self,
2832                    _ctx: &HookContext,
2833                    event: ToolCall<'_>,
2834                ) -> ToolCallAction {
2835                    if let ToolCall { .. } = event {
2836                        ToolCallAction::skip("denied after rewrite")
2837                    } else {
2838                        ToolCallAction::run()
2839                    }
2840                }
2841            }
2842            // Records the args + outcome seen on the `ToolResult` event.
2843            #[derive(Clone, Default)]
2844            struct ArgsProbe {
2845                args: Arc<Mutex<Option<String>>>,
2846                outcome: Arc<Mutex<Option<String>>>,
2847            }
2848            impl AgentHook for ArgsProbe {
2849                async fn on_tool_result(
2850                    &self,
2851                    _ctx: &HookContext,
2852                    event: ToolResultEvent<'_>,
2853                ) -> ToolResultAction {
2854                    if let ToolResultEvent {
2855                        args, raw_result, ..
2856                    } = event
2857                    {
2858                        *self.args.lock().expect("args") = Some(args.to_string());
2859                        *self.outcome.lock().expect("outcome") = Some(outcome_label(raw_result));
2860                    }
2861                    ToolResultAction::keep()
2862                }
2863            }
2864
2865            async fn run_surface(streaming: bool) -> (String, String) {
2866                let probe = ArgsProbe::default();
2867                // The tool must never execute; `MockAddTool` would produce a
2868                // `Success` outcome with result "42" if it (wrongly) ran.
2869                if streaming {
2870                    let mut stream = AgentBuilder::new(stream_model_one_tool_then_text("add"))
2871                        .tool(MockAddTool)
2872                        .add_hook(RewriteHook)
2873                        .add_hook(SkipHook)
2874                        .add_hook(probe.clone())
2875                        .build()
2876                        .runner("go")
2877                        .max_turns(3)
2878                        .stream()
2879                        .await;
2880                    while let Some(item) = stream.next().await {
2881                        if let Err(err) = item {
2882                            panic!("stream item errored: {err}");
2883                        }
2884                    }
2885                } else {
2886                    AgentBuilder::new(model_one_tool_then_text("add"))
2887                        .tool(MockAddTool)
2888                        .add_hook(RewriteHook)
2889                        .add_hook(SkipHook)
2890                        .add_hook(probe.clone())
2891                        .build()
2892                        .runner("go")
2893                        .max_turns(3)
2894                        .run()
2895                        .await
2896                        .expect("run should succeed after skipping the tool");
2897                }
2898                let args = probe.args.lock().expect("args").clone().expect("args seen");
2899                let outcome = probe
2900                    .outcome
2901                    .lock()
2902                    .expect("outcome")
2903                    .clone()
2904                    .expect("outcome seen");
2905                (args, outcome)
2906            }
2907
2908            for streaming in [false, true] {
2909                let (args, outcome) = run_surface(streaming).await;
2910                assert_eq!(
2911                    outcome, "skipped",
2912                    "the skipped tool must produce a Skipped outcome (streaming={streaming})"
2913                );
2914                let parsed: serde_json::Value =
2915                    serde_json::from_str(&args).expect("ToolResult args are valid JSON");
2916                assert_eq!(
2917                    parsed,
2918                    json!({ "x": 41, "y": 1 }),
2919                    "the skipped ToolResult must report the rewritten args, not the model's \
2920                     original {{}} (streaming={streaming}); got {args}"
2921                );
2922            }
2923        }
2924
2925        // End-to-end nesting: a *nested* `HookStack` that rewrites args then skips
2926        // must still report the rewritten args on the skipped `ToolResult` — the
2927        // inner rewrite is not lost behind the inner skip when the stack is added
2928        // as a single composed hook. Guards the nested-composition fix.
2929        #[tokio::test]
2930        async fn nested_hook_stack_rewrite_then_skip_reports_rewritten_args() {
2931            struct RewriteHook;
2932            impl AgentHook for RewriteHook {
2933                async fn on_tool_call(
2934                    &self,
2935                    _ctx: &HookContext,
2936                    event: ToolCall<'_>,
2937                ) -> ToolCallAction {
2938                    if let ToolCall { .. } = event {
2939                        ToolCallAction::rewrite(json!({ "x": 41, "y": 1 }))
2940                    } else {
2941                        ToolCallAction::run()
2942                    }
2943                }
2944            }
2945            struct SkipHook;
2946            impl AgentHook for SkipHook {
2947                async fn on_tool_call(
2948                    &self,
2949                    _ctx: &HookContext,
2950                    event: ToolCall<'_>,
2951                ) -> ToolCallAction {
2952                    if let ToolCall { .. } = event {
2953                        ToolCallAction::skip("denied after nested rewrite")
2954                    } else {
2955                        ToolCallAction::run()
2956                    }
2957                }
2958            }
2959            #[derive(Clone, Default)]
2960            struct ArgsProbe {
2961                args: Arc<Mutex<Option<String>>>,
2962                outcome: Arc<Mutex<Option<String>>>,
2963            }
2964            impl AgentHook for ArgsProbe {
2965                async fn on_tool_result(
2966                    &self,
2967                    _ctx: &HookContext,
2968                    event: ToolResultEvent<'_>,
2969                ) -> ToolResultAction {
2970                    if let ToolResultEvent {
2971                        args, raw_result, ..
2972                    } = event
2973                    {
2974                        *self.args.lock().expect("args") = Some(args.to_string());
2975                        *self.outcome.lock().expect("outcome") = Some(outcome_label(raw_result));
2976                    }
2977                    ToolResultAction::keep()
2978                }
2979            }
2980
2981            // The rewrite + skip live inside a *nested* stack added as one hook.
2982            fn nested_stack() -> HookStack {
2983                let mut nested = HookStack::new();
2984                nested.push(RewriteHook);
2985                nested.push(SkipHook);
2986                nested
2987            }
2988
2989            // Verified on both surfaces: run_single_tool (shared) drives the same
2990            // nested resolution, so blocking and streaming must agree.
2991            for streaming in [false, true] {
2992                let probe = ArgsProbe::default();
2993                if streaming {
2994                    let mut stream = AgentBuilder::new(stream_model_one_tool_then_text("add"))
2995                        .tool(MockAddTool)
2996                        .add_hook(nested_stack())
2997                        .add_hook(probe.clone())
2998                        .build()
2999                        .runner("go")
3000                        .max_turns(3)
3001                        .stream()
3002                        .await;
3003                    while let Some(item) = stream.next().await {
3004                        if let Err(err) = item {
3005                            panic!("stream item errored: {err}");
3006                        }
3007                    }
3008                } else {
3009                    AgentBuilder::new(model_one_tool_then_text("add"))
3010                        .tool(MockAddTool)
3011                        .add_hook(nested_stack())
3012                        .add_hook(probe.clone())
3013                        .build()
3014                        .runner("go")
3015                        .max_turns(3)
3016                        .run()
3017                        .await
3018                        .expect("run should succeed after the nested stack skips the tool");
3019                }
3020
3021                assert_eq!(
3022                    probe.outcome.lock().expect("outcome").clone(),
3023                    Some("skipped".to_string()),
3024                    "streaming={streaming}"
3025                );
3026                let args = probe.args.lock().expect("args").clone().expect("args seen");
3027                let parsed: serde_json::Value =
3028                    serde_json::from_str(&args).expect("valid JSON args");
3029                assert_eq!(
3030                    parsed,
3031                    json!({ "x": 41, "y": 1 }),
3032                    "the nested stack's rewrite must survive its skip and reach the ToolResult \
3033                     (streaming={streaming}); got {args}"
3034                );
3035            }
3036        }
3037
3038        // (8) Invalid JSON arguments are classified as a structured `InvalidArgs`
3039        // failure rather than surfacing as an opaque string.
3040        #[tokio::test]
3041        async fn invalid_args_are_classified_as_invalid_args() {
3042            let hook = OutcomeHook::default();
3043            AgentBuilder::new(MockCompletionModel::from_turns([
3044                // `add` needs integers; a string is a hard parse failure.
3045                MockTurn::tool_call("tc1", "add", json!({ "x": "not-a-number", "y": 1 })),
3046                MockTurn::text("done"),
3047            ]))
3048            .tool(MockAddTool)
3049            .add_hook(hook.clone())
3050            .build()
3051            .runner("go")
3052            .max_turns(3)
3053            .run()
3054            .await
3055            .expect("an invalid-args failure is model-visible feedback, not fatal");
3056
3057            assert_eq!(hook.outcomes(), vec!["error:invalid_args".to_string()]);
3058        }
3059
3060        // Result metadata a tool attaches reaches the hook but never appears in the
3061        // model-visible output on either execution surface.
3062        #[tokio::test]
3063        async fn success_result_metadata_reaches_hook_but_not_model() {
3064            struct MetadataProbe {
3065                seen: Arc<Mutex<Option<String>>>,
3066                model_output: Arc<Mutex<Option<String>>>,
3067            }
3068            impl AgentHook for MetadataProbe {
3069                async fn on_tool_result(
3070                    &self,
3071                    _ctx: &HookContext,
3072                    event: ToolResultEvent<'_>,
3073                ) -> ToolResultAction {
3074                    if let ToolResultEvent {
3075                        presentation,
3076                        tool_context,
3077                        ..
3078                    } = event
3079                    {
3080                        *self.seen.lock().expect("seen") = tool_context
3081                            .result::<MockRequestId>()
3082                            .map(|id| id.0.clone());
3083                        *self.model_output.lock().expect("model_output") =
3084                            Some(presentation.render());
3085                    }
3086                    ToolResultAction::keep()
3087                }
3088            }
3089
3090            async fn run_surface(streaming: bool) -> (Option<String>, String) {
3091                let seen: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));
3092                let model_output: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));
3093                let probe = MetadataProbe {
3094                    seen: seen.clone(),
3095                    model_output: model_output.clone(),
3096                };
3097
3098                if streaming {
3099                    let mut stream =
3100                        AgentBuilder::new(stream_model_one_tool_then_text("with_meta"))
3101                            .tool(MockMetadataTool)
3102                            .add_hook(probe)
3103                            .build()
3104                            .runner("go")
3105                            .max_turns(3)
3106                            .stream()
3107                            .await;
3108                    while let Some(item) = stream.next().await {
3109                        if let Err(error) = item {
3110                            panic!("stream item errored: {error}");
3111                        }
3112                    }
3113                } else {
3114                    AgentBuilder::new(model_one_tool_then_text("with_meta"))
3115                        .tool(MockMetadataTool)
3116                        .add_hook(probe)
3117                        .build()
3118                        .runner("go")
3119                        .max_turns(3)
3120                        .run()
3121                        .await
3122                        .expect("run should succeed");
3123                }
3124
3125                let seen_value = seen.lock().expect("seen").clone();
3126                let output = model_output
3127                    .lock()
3128                    .expect("model_output")
3129                    .clone()
3130                    .expect("output");
3131                (seen_value, output)
3132            }
3133
3134            for streaming in [false, true] {
3135                let (seen, output) = run_surface(streaming).await;
3136                assert_eq!(
3137                    seen,
3138                    Some("req-7".to_string()),
3139                    "the tool's result metadata must reach the hook (streaming={streaming})"
3140                );
3141                assert_eq!(output, "done");
3142                assert!(
3143                    !output.contains("req-7"),
3144                    "result metadata must never leak into model output (streaming={streaming})"
3145                );
3146            }
3147        }
3148
3149        // (6) A `ToolResultAction::Rewrite` hook redacts the model-visible text, but a later
3150        // policy hook still sees the tool's *raw* structured outcome — a rewrite
3151        // changes only what the model sees, not the classification.
3152        #[tokio::test]
3153        async fn rewrite_result_does_not_mask_the_structured_outcome() {
3154            struct Redact;
3155            impl AgentHook for Redact {
3156                async fn on_tool_result(
3157                    &self,
3158                    _ctx: &HookContext,
3159                    event: ToolResultEvent<'_>,
3160                ) -> ToolResultAction {
3161                    if let ToolResultEvent { .. } = event {
3162                        ToolResultAction::rewrite("[REDACTED]")
3163                    } else {
3164                        ToolResultAction::keep()
3165                    }
3166                }
3167            }
3168
3169            let observer = OutcomeHook::default();
3170            AgentBuilder::new(model_one_tool_then_text("flaky_tool"))
3171                .tool(MockFailingTool::new(ToolErrorKind::NotFound))
3172                // Observer AFTER the redactor: it still sees the true outcome, and
3173                // the chained (redacted) model-visible result.
3174                .add_hook(Redact)
3175                .add_hook(observer.clone())
3176                .build()
3177                .runner("go")
3178                .max_turns(3)
3179                .run()
3180                .await
3181                .expect("run should succeed");
3182
3183            assert_eq!(observer.outcomes(), vec!["error:not_found".to_string()]);
3184            assert_eq!(observer.results(), vec!["[REDACTED]".to_string()]);
3185        }
3186
3187        // (9) The blocking and streaming surfaces observe identical structured
3188        // outcomes for the same scenario.
3189        #[tokio::test]
3190        async fn streaming_and_blocking_outcomes_match() {
3191            let blocking = OutcomeHook::default();
3192            AgentBuilder::new(model_one_tool_then_text("flaky_tool"))
3193                .tool(MockFailingTool::new(ToolErrorKind::Timeout))
3194                .add_hook(blocking.clone())
3195                .build()
3196                .runner("go")
3197                .max_turns(3)
3198                .run()
3199                .await
3200                .expect("blocking run should succeed");
3201
3202            let streaming = OutcomeHook::default();
3203            let mut stream = AgentBuilder::new(stream_model_one_tool_then_text("flaky_tool"))
3204                .tool(MockFailingTool::new(ToolErrorKind::Timeout))
3205                .add_hook(streaming.clone())
3206                .build()
3207                .runner("go")
3208                .max_turns(3)
3209                .stream()
3210                .await;
3211            while let Some(item) = stream.next().await {
3212                if let Err(err) = item {
3213                    panic!("stream item errored: {err}");
3214                }
3215            }
3216
3217            assert_eq!(blocking.outcomes(), vec!["error:timeout".to_string()]);
3218            assert_eq!(blocking.outcomes(), streaming.outcomes());
3219            assert_eq!(blocking.results(), streaming.results());
3220        }
3221
3222        // (10) With two tools in one turn at `concurrency > 1`, both structured
3223        // outcomes are observed and the persisted tool results keep call order.
3224        #[tokio::test]
3225        async fn concurrent_tools_preserve_order_and_both_outcomes() {
3226            use rig_core::message::{
3227                AssistantContent, ToolCall as MessageToolCall, ToolFunction, UserContent,
3228            };
3229
3230            let turn = MockTurn::from_contents([
3231                AssistantContent::ToolCall(MessageToolCall::new(
3232                    "tc_add".to_string(),
3233                    ToolFunction::new("add".to_string(), json!({ "x": 2, "y": 3 })),
3234                )),
3235                AssistantContent::ToolCall(MessageToolCall::new(
3236                    "tc_flaky".to_string(),
3237                    ToolFunction::new("flaky_tool".to_string(), json!({})),
3238                )),
3239            ])
3240            .expect("two tool calls");
3241
3242            let observer = OutcomeHook::default();
3243            let response = AgentBuilder::new(MockCompletionModel::from_turns([
3244                turn,
3245                MockTurn::text("done"),
3246            ]))
3247            .tool(MockAddTool)
3248            .tool(MockFailingTool::new(ToolErrorKind::Timeout))
3249            .add_hook(observer.clone())
3250            .build()
3251            .runner("go")
3252            .max_turns(3)
3253            .tool_concurrency(2)
3254            .run()
3255            .await
3256            .expect("run should succeed");
3257
3258            // Hook order may interleave under concurrency, so compare as a set.
3259            let mut outcomes = observer.outcomes();
3260            outcomes.sort();
3261            assert_eq!(
3262                outcomes,
3263                vec!["error:timeout".to_string(), "success".to_string()]
3264            );
3265
3266            // The persisted tool results must keep tool-call order regardless of
3267            // completion timing: `add` (tc_add) before `flaky_tool` (tc_flaky).
3268            let messages = response.messages.expect("messages");
3269            let tool_result_ids: Vec<String> = messages
3270                .iter()
3271                .flat_map(|message| match message {
3272                    crate::completion::Message::User { content } => content
3273                        .iter()
3274                        .filter_map(|c| match c {
3275                            UserContent::ToolResult(result) => Some(result.id.clone()),
3276                            _ => None,
3277                        })
3278                        .collect::<Vec<_>>(),
3279                    _ => Vec::new(),
3280                })
3281                .collect();
3282            assert_eq!(
3283                tool_result_ids,
3284                vec!["tc_add".to_string(), "tc_flaky".to_string()],
3285                "tool results must be persisted in call order"
3286            );
3287        }
3288    }
3289
3290    /// Safety net for the streaming/non-streaming unification: pins the blocking
3291    /// driver's span topology (span name, `invoke_agent` creation, the
3292    /// `follows_from` chain, and `created_agent_span`-gated run-level usage) so a
3293    /// later refactor onto a shared engine cannot silently drift it. The
3294    /// streaming side is already pinned by `assert_stream_usage_recorded_on_chat_spans`.
3295    mod span_safety_net {
3296        use std::collections::{HashMap, HashSet};
3297        use std::sync::{Arc, Mutex};
3298
3299        use futures::StreamExt;
3300        use tracing::Instrument;
3301        use tracing::field::{Field, Visit};
3302        use tracing::span::{Attributes, Record};
3303        use tracing::{Id, Subscriber};
3304        use tracing_subscriber::layer::{Context, SubscriberExt};
3305        use tracing_subscriber::{Layer, Registry, registry::LookupSpan};
3306
3307        use crate::agent::{
3308            AgentBuilder, HookContext, MultiTurnStreamItem, ToolResultAction, ToolResultEvent,
3309        };
3310        use crate::completion::{
3311            CompletionError, CompletionModel, CompletionRequest, CompletionResponse, Prompt,
3312            PromptError, Usage,
3313        };
3314        use crate::streaming::StreamedAssistantContent;
3315        use crate::streaming::StreamingCompletionResponse;
3316        use crate::test_utils::{
3317            MockAddTool, MockCompletionModel, MockResponse, MockStreamEvent, MockTurn,
3318        };
3319        use crate::tool::{ToolContext, ToolExecutionError};
3320        use rig_core::telemetry::{CompletionOperation, CompletionSpanBuilder};
3321
3322        use super::{BoundedResponseRetry, StopCompletedModelTurn, TestRetryMode};
3323
3324        #[derive(Clone)]
3325        struct CapturedSpan {
3326            id: u64,
3327            name: String,
3328            target: String,
3329            field_names: HashSet<String>,
3330            u64_fields: HashMap<String, u64>,
3331            string_fields: HashMap<String, Vec<String>>,
3332        }
3333
3334        #[derive(Clone, Default)]
3335        struct Captured {
3336            spans: Arc<Mutex<Vec<CapturedSpan>>>,
3337            /// `(span, follows_from)` pairs recorded via `Span::follows_from`.
3338            follows: Arc<Mutex<Vec<(u64, u64)>>>,
3339        }
3340
3341        impl Captured {
3342            fn insert(&self, id: &Id, name: &str, target: &str) {
3343                self.spans.lock().expect("spans").push(CapturedSpan {
3344                    id: id.into_u64(),
3345                    name: name.to_string(),
3346                    target: target.to_string(),
3347                    field_names: HashSet::new(),
3348                    u64_fields: HashMap::new(),
3349                    string_fields: HashMap::new(),
3350                });
3351            }
3352
3353            fn record(
3354                &self,
3355                id: &Id,
3356                names: HashSet<String>,
3357                u64s: HashMap<String, u64>,
3358                strings: HashMap<String, String>,
3359            ) {
3360                let id = id.into_u64();
3361                if let Ok(mut spans) = self.spans.lock()
3362                    && let Some(span) = spans.iter_mut().find(|s| s.id == id)
3363                {
3364                    span.field_names.extend(names);
3365                    span.u64_fields.extend(u64s);
3366                    for (name, value) in strings {
3367                        span.string_fields.entry(name).or_default().push(value);
3368                    }
3369                }
3370            }
3371
3372            fn follows_from(&self, span: &Id, follows: &Id) {
3373                self.follows
3374                    .lock()
3375                    .expect("follows")
3376                    .push((span.into_u64(), follows.into_u64()));
3377            }
3378
3379            fn clear(&self) {
3380                self.spans.lock().expect("spans").clear();
3381                self.follows.lock().expect("follows").clear();
3382            }
3383
3384            fn snapshot(&self) -> Vec<CapturedSpan> {
3385                self.spans.lock().expect("spans").clone()
3386            }
3387
3388            fn follows_edges(&self) -> Vec<(u64, u64)> {
3389                self.follows.lock().expect("follows").clone()
3390            }
3391        }
3392
3393        struct CaptureLayer {
3394            captured: Captured,
3395        }
3396
3397        impl<S> Layer<S> for CaptureLayer
3398        where
3399            S: Subscriber + for<'l> LookupSpan<'l>,
3400        {
3401            fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, _ctx: Context<'_, S>) {
3402                self.captured
3403                    .insert(id, attrs.metadata().name(), attrs.metadata().target());
3404            }
3405
3406            fn on_record(&self, span: &Id, values: &Record<'_>, _ctx: Context<'_, S>) {
3407                let mut visitor = FieldVisitor::default();
3408                values.record(&mut visitor);
3409                self.captured
3410                    .record(span, visitor.names, visitor.u64s, visitor.strings);
3411            }
3412
3413            fn on_follows_from(&self, span: &Id, follows: &Id, _ctx: Context<'_, S>) {
3414                self.captured.follows_from(span, follows);
3415            }
3416        }
3417
3418        #[derive(Default)]
3419        struct FieldVisitor {
3420            names: HashSet<String>,
3421            u64s: HashMap<String, u64>,
3422            strings: HashMap<String, String>,
3423        }
3424
3425        impl Visit for FieldVisitor {
3426            fn record_u64(&mut self, field: &Field, value: u64) {
3427                self.names.insert(field.name().to_string());
3428                self.u64s.insert(field.name().to_string(), value);
3429            }
3430
3431            fn record_str(&mut self, field: &Field, value: &str) {
3432                self.names.insert(field.name().to_string());
3433                self.strings
3434                    .insert(field.name().to_string(), value.to_string());
3435            }
3436
3437            fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
3438                self.names.insert(field.name().to_string());
3439                self.strings
3440                    .insert(field.name().to_string(), format!("{value:?}"));
3441            }
3442        }
3443
3444        fn usage(input: u64, output: u64) -> Usage {
3445            Usage {
3446                input_tokens: input,
3447                output_tokens: output,
3448                ..Usage::new()
3449            }
3450        }
3451
3452        /// Two-turn tool scenario: the blocking driver emits chat -> execute_tool
3453        /// -> chat, exercising the `follows_from` chain.
3454        fn tool_then_text_model() -> MockCompletionModel {
3455            MockCompletionModel::from_turns([
3456                MockTurn::tool_call("tc1", "add", serde_json::json!({"x": 2, "y": 3}))
3457                    .with_usage(usage(7, 11)),
3458                MockTurn::text("the answer is 5").with_usage(usage(13, 17)),
3459            ])
3460        }
3461
3462        #[derive(Clone)]
3463        struct CompletionTelemetryModel {
3464            inner: MockCompletionModel,
3465        }
3466
3467        impl CompletionModel for CompletionTelemetryModel {
3468            type Response = MockResponse;
3469            type StreamingResponse = MockResponse;
3470            type Client = ();
3471
3472            fn make(_client: &Self::Client, _model: impl Into<String>) -> Self {
3473                Self {
3474                    inner: MockCompletionModel::default(),
3475                }
3476            }
3477
3478            async fn completion(
3479                &self,
3480                request: CompletionRequest,
3481            ) -> Result<CompletionResponse<Self::Response>, CompletionError> {
3482                let span = CompletionSpanBuilder::new(
3483                    "fixture-provider",
3484                    "fixture-model",
3485                    CompletionOperation::Chat,
3486                )
3487                .build();
3488                self.inner.completion(request).instrument(span).await
3489            }
3490
3491            async fn stream(
3492                &self,
3493                request: CompletionRequest,
3494            ) -> Result<StreamingCompletionResponse<Self::StreamingResponse>, CompletionError>
3495            {
3496                let span = CompletionSpanBuilder::new(
3497                    "fixture-provider",
3498                    "fixture-model",
3499                    CompletionOperation::ChatStreaming,
3500                )
3501                .build();
3502                self.inner.stream(request).instrument(span).await
3503            }
3504        }
3505
3506        /// Register the blocking driver's span callsites against the scoped
3507        /// subscriber before asserting, mirroring the streaming usage test's
3508        /// interest-cache warm-up (a foreign thread without our subscriber can
3509        /// otherwise cache `Interest::never` for these callsites).
3510        async fn warm_blocking_callsites() {
3511            let agent = AgentBuilder::new(tool_then_text_model())
3512                .record_content_telemetry(true)
3513                .tool(MockAddTool)
3514                .build();
3515            let _ = agent.runner("add 2 and 3").max_turns(3).run().await;
3516        }
3517
3518        async fn run_blocking_response_retry_with_content_telemetry() {
3519            AgentBuilder::new(MockCompletionModel::from_turns([
3520                MockTurn::text("rejected"),
3521                MockTurn::text("accepted"),
3522            ]))
3523            .record_content_telemetry(true)
3524            .add_hook(BoundedResponseRetry::new(
3525                "rejected",
3526                1,
3527                TestRetryMode::Repeat,
3528            ))
3529            .build()
3530            .runner("question")
3531            .max_turns(2)
3532            .run()
3533            .await
3534            .expect("blocking retry should succeed");
3535        }
3536
3537        async fn run_streaming_response_retry_with_content_telemetry() {
3538            let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([
3539                [
3540                    MockStreamEvent::text("rejected"),
3541                    MockStreamEvent::final_response_with_default_usage(),
3542                ],
3543                [
3544                    MockStreamEvent::text("accepted"),
3545                    MockStreamEvent::final_response_with_default_usage(),
3546                ],
3547            ]))
3548            .record_content_telemetry(true)
3549            .add_hook(BoundedResponseRetry::new(
3550                "rejected",
3551                1,
3552                TestRetryMode::Repeat,
3553            ))
3554            .build()
3555            .runner("question")
3556            .max_turns(2)
3557            .stream()
3558            .await;
3559
3560            let mut saw_final = false;
3561            while let Some(item) = stream.next().await {
3562                if let MultiTurnStreamItem::FinalResponse(response) =
3563                    item.expect("streaming retry item")
3564                {
3565                    saw_final = true;
3566                    assert_eq!(response.output, "accepted");
3567                }
3568            }
3569            assert!(saw_final, "streaming retry should produce a final response");
3570        }
3571
3572        async fn run_blocking_model_turn_stop_with_content_telemetry() {
3573            let error = AgentBuilder::new(MockCompletionModel::from_turns([MockTurn::text(
3574                "stopped blocking response",
3575            )]))
3576            .record_content_telemetry(true)
3577            .add_hook(StopCompletedModelTurn)
3578            .build()
3579            .runner("question")
3580            .run()
3581            .await
3582            .expect_err("blocking model-turn stop should cancel the run");
3583
3584            assert!(matches!(
3585                error,
3586                PromptError::PromptCancelled { reason, .. }
3587                    if reason == "stop completed model turn"
3588            ));
3589        }
3590
3591        async fn run_streaming_model_turn_stop_with_content_telemetry() {
3592            let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([[
3593                MockStreamEvent::text("stopped streaming response"),
3594                MockStreamEvent::final_response_with_default_usage(),
3595            ]]))
3596            .record_content_telemetry(true)
3597            .add_hook(StopCompletedModelTurn)
3598            .build()
3599            .runner("question")
3600            .stream()
3601            .await;
3602
3603            let mut provider_finals = 0;
3604            let mut agent_finals = 0;
3605            let mut retries = 0;
3606            let mut errors = 0;
3607            while let Some(item) = stream.next().await {
3608                match item {
3609                    Ok(MultiTurnStreamItem::StreamAssistantItem(
3610                        StreamedAssistantContent::Final(_),
3611                    )) => provider_finals += 1,
3612                    Ok(MultiTurnStreamItem::FinalResponse(_)) => agent_finals += 1,
3613                    Ok(MultiTurnStreamItem::ModelTurnRetried { .. }) => retries += 1,
3614                    Ok(_) => {}
3615                    Err(error) => {
3616                        errors += 1;
3617                        assert!(matches!(
3618                            error,
3619                            super::StreamingError::Prompt(error)
3620                                if matches!(
3621                                    error.as_ref(),
3622                                    PromptError::PromptCancelled { reason, .. }
3623                                        if reason == "stop completed model turn"
3624                                )
3625                        ));
3626                    }
3627                }
3628            }
3629
3630            assert_eq!(provider_finals, 1);
3631            assert_eq!(agent_finals, 0);
3632            assert_eq!(retries, 0);
3633            assert_eq!(errors, 1);
3634        }
3635
3636        /// Cross-crate tripwire: the chat span built by `build_chat_span!`
3637        /// must statically declare rig-core's full completion-parent contract
3638        /// (marker + every required field) plus the agent-specific
3639        /// `gen_ai.agent.name`. `Span::record` silently no-ops on undeclared
3640        /// fields, so a missing field here would lose that telemetry on every
3641        /// adopted completion with no error.
3642        #[test]
3643        fn chat_span_declares_the_full_completion_parent_contract() {
3644            use rig_core::telemetry::{
3645                COMPLETION_PARENT_MARKER_FIELD, COMPLETION_PARENT_REQUIRED_FIELDS,
3646            };
3647
3648            let _isolation = crate::test_utils::scoped_tracing_subscriber_guard_blocking();
3649            tracing::subscriber::with_default(Registry::default(), || {
3650                let agent = AgentBuilder::new(MockCompletionModel::text("done"))
3651                    .name("contract-agent")
3652                    .build();
3653                let runner = agent.runner("hello");
3654                let span = build_chat_span!(runner, None, "chat", "chat");
3655                let Some(metadata) = span.metadata() else {
3656                    panic!("chat span was disabled");
3657                };
3658                let declared: HashSet<&str> =
3659                    metadata.fields().iter().map(|field| field.name()).collect();
3660                let expected: HashSet<&str> = COMPLETION_PARENT_REQUIRED_FIELDS
3661                    .iter()
3662                    .copied()
3663                    .chain([COMPLETION_PARENT_MARKER_FIELD, "gen_ai.agent.name"])
3664                    .collect();
3665                assert_eq!(declared, expected);
3666                // Duplicate field names collapse in a `HashSet`, so also pin
3667                // the count: set equality alone cannot catch a field declared
3668                // twice (e.g. an extra colliding with a contract field).
3669                assert_eq!(metadata.fields().len(), expected.len());
3670            });
3671        }
3672
3673        #[tokio::test]
3674        async fn response_retry_records_only_accepted_content_on_both_surfaces() {
3675            let _isolation = crate::test_utils::scoped_tracing_subscriber_guard().await;
3676            let captured = Captured::default();
3677            let subscriber = Registry::default().with(CaptureLayer {
3678                captured: captured.clone(),
3679            });
3680            let _default = tracing::subscriber::set_default(subscriber);
3681
3682            // Register both transport callsites under this subscriber before
3683            // inspecting field recordings.
3684            run_blocking_response_retry_with_content_telemetry().await;
3685            run_streaming_response_retry_with_content_telemetry().await;
3686            tracing::callsite::rebuild_interest_cache();
3687            captured.clear();
3688
3689            run_blocking_response_retry_with_content_telemetry().await;
3690            let blocking = captured.snapshot();
3691            let blocking_chats = blocking
3692                .iter()
3693                .filter(|span| span.name == "chat")
3694                .collect::<Vec<_>>();
3695            assert_eq!(blocking_chats.len(), 2);
3696            assert!(
3697                blocking_chats
3698                    .iter()
3699                    .all(|span| span.target == "rig::agent_chat")
3700            );
3701            assert!(
3702                !blocking_chats[0]
3703                    .field_names
3704                    .contains("gen_ai.output.messages"),
3705                "rejected blocking content must not be recorded as model output"
3706            );
3707            assert!(
3708                blocking_chats[1]
3709                    .field_names
3710                    .contains("gen_ai.output.messages"),
3711                "accepted blocking content must be recorded as model output"
3712            );
3713            let blocking_output = blocking_chats[1]
3714                .string_fields
3715                .get("gen_ai.output.messages")
3716                .expect("accepted blocking output value");
3717            assert!(
3718                blocking_output
3719                    .iter()
3720                    .any(|value| value.contains("accepted"))
3721            );
3722            assert!(
3723                blocking_output
3724                    .iter()
3725                    .all(|value| !value.contains("rejected"))
3726            );
3727            let blocking_completion = blocking
3728                .iter()
3729                .find(|span| span.name == "invoke_agent")
3730                .and_then(|span| span.string_fields.get("gen_ai.completion"))
3731                .expect("accepted blocking run-level completion");
3732            assert_eq!(blocking_completion, &["accepted"]);
3733
3734            captured.clear();
3735            run_streaming_response_retry_with_content_telemetry().await;
3736            let streaming = captured.snapshot();
3737            let streaming_chats = streaming
3738                .iter()
3739                .filter(|span| span.name == "chat_streaming")
3740                .collect::<Vec<_>>();
3741            assert_eq!(streaming_chats.len(), 2);
3742            assert!(
3743                streaming_chats
3744                    .iter()
3745                    .all(|span| span.target == "rig::agent_chat")
3746            );
3747            assert!(
3748                !streaming_chats[0]
3749                    .field_names
3750                    .contains("gen_ai.output.messages"),
3751                "rejected streaming content must not be recorded as model output"
3752            );
3753            assert!(
3754                streaming_chats[1]
3755                    .field_names
3756                    .contains("gen_ai.output.messages"),
3757                "accepted streaming content must be recorded as model output"
3758            );
3759            let streaming_output = streaming_chats[1]
3760                .string_fields
3761                .get("gen_ai.output.messages")
3762                .expect("accepted streaming output value");
3763            assert!(
3764                streaming_output
3765                    .iter()
3766                    .any(|value| value.contains("accepted"))
3767            );
3768            assert!(
3769                streaming_output
3770                    .iter()
3771                    .all(|value| !value.contains("rejected"))
3772            );
3773            let streaming_completion = streaming
3774                .iter()
3775                .find(|span| span.name == "invoke_agent")
3776                .and_then(|span| span.string_fields.get("gen_ai.completion"))
3777                .expect("accepted streaming run-level completion");
3778            assert_eq!(streaming_completion, &["accepted"]);
3779        }
3780
3781        #[tokio::test]
3782        async fn model_turn_stop_preserves_completed_content_telemetry_on_both_surfaces() {
3783            let _isolation = crate::test_utils::scoped_tracing_subscriber_guard().await;
3784            let captured = Captured::default();
3785            let subscriber = Registry::default().with(CaptureLayer {
3786                captured: captured.clone(),
3787            });
3788            let _default = tracing::subscriber::set_default(subscriber);
3789
3790            run_blocking_model_turn_stop_with_content_telemetry().await;
3791            run_streaming_model_turn_stop_with_content_telemetry().await;
3792            tracing::callsite::rebuild_interest_cache();
3793            captured.clear();
3794
3795            run_blocking_model_turn_stop_with_content_telemetry().await;
3796            let blocking = captured.snapshot();
3797            let blocking_output = blocking
3798                .iter()
3799                .find(|span| span.name == "chat")
3800                .and_then(|span| span.string_fields.get("gen_ai.output.messages"))
3801                .expect("stopped blocking turn should retain output telemetry");
3802            assert!(
3803                blocking_output
3804                    .iter()
3805                    .any(|value| value.contains("stopped blocking response"))
3806            );
3807
3808            captured.clear();
3809            run_streaming_model_turn_stop_with_content_telemetry().await;
3810            let streaming = captured.snapshot();
3811            let streaming_output = streaming
3812                .iter()
3813                .find(|span| span.name == "chat_streaming")
3814                .and_then(|span| span.string_fields.get("gen_ai.output.messages"))
3815                .expect("stopped streaming turn should retain output telemetry");
3816            assert!(
3817                streaming_output
3818                    .iter()
3819                    .any(|value| value.contains("stopped streaming response"))
3820            );
3821            let streaming_completion = streaming
3822                .iter()
3823                .find(|span| span.name == "invoke_agent")
3824                .and_then(|span| span.string_fields.get("gen_ai.completion"))
3825                .expect("stopped streaming turn should retain run-level completion telemetry");
3826            assert_eq!(streaming_completion, &["stopped streaming response"]);
3827        }
3828
3829        #[tokio::test]
3830        async fn run_records_usage_and_chains_chat_spans_on_a_created_agent_span() {
3831            let _isolation = crate::test_utils::scoped_tracing_subscriber_guard().await;
3832            let captured = Captured::default();
3833            let subscriber = Registry::default().with(CaptureLayer {
3834                captured: captured.clone(),
3835            });
3836            let _default = tracing::subscriber::set_default(subscriber);
3837
3838            warm_blocking_callsites().await;
3839            tracing::callsite::rebuild_interest_cache();
3840            captured.clear();
3841
3842            let agent = AgentBuilder::new(tool_then_text_model())
3843                .record_content_telemetry(true)
3844                .tool(MockAddTool)
3845                .build();
3846            let response = agent
3847                .runner("add 2 and 3")
3848                .max_turns(3)
3849                .run()
3850                .await
3851                .expect("blocking run should succeed");
3852            assert_eq!(response.output, "the answer is 5");
3853
3854            let spans = captured.snapshot();
3855
3856            // The blocking chat span is named "chat" (NOT "chat_streaming").
3857            let chat_spans: Vec<&CapturedSpan> =
3858                spans.iter().filter(|s| s.name == "chat").collect();
3859            assert_eq!(chat_spans.len(), 2, "two model turns -> two chat spans");
3860            assert!(
3861                spans.iter().all(|s| s.name != "chat_streaming"),
3862                "blocking driver must not emit chat_streaming spans"
3863            );
3864
3865            // A run with no ambient span creates its own invoke_agent span...
3866            let agent_span = spans
3867                .iter()
3868                .find(|s| s.name == "invoke_agent")
3869                .expect("blocking run should create an invoke_agent span");
3870
3871            // ...and records aggregate usage + completion onto it (created_agent_span).
3872            assert_eq!(
3873                agent_span.u64_fields.get("gen_ai.usage.input_tokens"),
3874                Some(&(7 + 13)),
3875            );
3876            assert_eq!(
3877                agent_span.u64_fields.get("gen_ai.usage.output_tokens"),
3878                Some(&(11 + 17)),
3879            );
3880            assert!(
3881                agent_span.field_names.contains("gen_ai.completion"),
3882                "the created agent span records the final completion text"
3883            );
3884
3885            // The blocking driver links chat/tool spans into a linear
3886            // follows_from chain (chat#1 -> execute_tool -> chat#2); the
3887            // streaming driver does not, so this is a blocking-only invariant the
3888            // unification must keep.
3889            let tool_span = spans
3890                .iter()
3891                .find(|s| s.name == "execute_tool")
3892                .expect("tool turn should emit an execute_tool span");
3893            let edges = captured.follows_edges();
3894            assert!(
3895                edges.contains(&(tool_span.id, chat_spans[0].id)),
3896                "execute_tool should follow_from the first chat span; edges={edges:?}"
3897            );
3898            assert!(
3899                edges.contains(&(chat_spans[1].id, tool_span.id)),
3900                "the second chat span should follow_from execute_tool; edges={edges:?}"
3901            );
3902        }
3903
3904        #[tokio::test]
3905        async fn classic_completion_parent_is_enriched_without_duplicate_provider_span() {
3906            let _isolation = crate::test_utils::scoped_tracing_subscriber_guard().await;
3907            let captured = Captured::default();
3908            let subscriber = Registry::default().with(CaptureLayer {
3909                captured: captured.clone(),
3910            });
3911            let _default = tracing::subscriber::set_default(subscriber);
3912
3913            let warm = AgentBuilder::new(CompletionTelemetryModel {
3914                inner: MockCompletionModel::text("warm"),
3915            })
3916            .build();
3917            let _ = warm.prompt("warm").await;
3918            tracing::callsite::rebuild_interest_cache();
3919            captured.clear();
3920
3921            let agent = AgentBuilder::new(CompletionTelemetryModel {
3922                inner: MockCompletionModel::text("done"),
3923            })
3924            .build();
3925            let response = agent.prompt("hello").await.expect("prompt should succeed");
3926            assert_eq!(response, "done");
3927
3928            let spans = captured.snapshot();
3929            let chat_spans = spans
3930                .iter()
3931                .filter(|span| span.name == "chat")
3932                .collect::<Vec<_>>();
3933            assert_eq!(chat_spans.len(), 1, "provider telemetry must reuse chat");
3934            assert_eq!(chat_spans[0].target, "rig::agent_chat");
3935            assert!(
3936                spans.iter().all(|span| span.target != "rig::completions"),
3937                "an adopted classic completion parent must not gain a provider child"
3938            );
3939            assert_eq!(
3940                chat_spans[0]
3941                    .string_fields
3942                    .get("gen_ai.provider.name")
3943                    .and_then(|values| values.first())
3944                    .map(String::as_str),
3945                Some("fixture-provider")
3946            );
3947            assert_eq!(
3948                chat_spans[0]
3949                    .string_fields
3950                    .get("gen_ai.request.model")
3951                    .and_then(|values| values.first())
3952                    .map(String::as_str),
3953                Some("fixture-model")
3954            );
3955        }
3956
3957        #[tokio::test]
3958        async fn run_does_not_record_usage_onto_a_caller_supplied_outer_span() {
3959            let _isolation = crate::test_utils::scoped_tracing_subscriber_guard().await;
3960            let captured = Captured::default();
3961            let subscriber = Registry::default().with(CaptureLayer {
3962                captured: captured.clone(),
3963            });
3964            let _default = tracing::subscriber::set_default(subscriber);
3965
3966            warm_blocking_callsites().await;
3967            tracing::callsite::rebuild_interest_cache();
3968            captured.clear();
3969
3970            // Declare the fields the guard protects so a regression (recording
3971            // onto a caller span) is actually observable rather than a silent
3972            // no-op on an undeclared field.
3973            let outer = tracing::info_span!(
3974                "outer",
3975                gen_ai.completion = tracing::field::Empty,
3976                gen_ai.usage.input_tokens = tracing::field::Empty,
3977                gen_ai.usage.output_tokens = tracing::field::Empty,
3978            );
3979            async {
3980                let agent = AgentBuilder::new(tool_then_text_model())
3981                    .tool(MockAddTool)
3982                    .build();
3983                agent
3984                    .runner("add 2 and 3")
3985                    .max_turns(3)
3986                    .run()
3987                    .await
3988                    .expect("blocking run should succeed");
3989            }
3990            .instrument(outer)
3991            .await;
3992
3993            let spans = captured.snapshot();
3994            // Under an ambient span the driver adopts it; no invoke_agent is created.
3995            assert!(
3996                spans.iter().all(|s| s.name != "invoke_agent"),
3997                "an ambient outer span should be adopted, not wrapped in invoke_agent"
3998            );
3999            let outer_span = spans
4000                .iter()
4001                .find(|s| s.name == "outer")
4002                .expect("outer span should be captured");
4003            assert!(
4004                outer_span
4005                    .field_names
4006                    .iter()
4007                    .all(|name| !name.starts_with("gen_ai.usage.")),
4008                "run-level usage must not be recorded onto a caller-supplied outer span"
4009            );
4010            assert!(
4011                !outer_span.field_names.contains("gen_ai.completion"),
4012                "run-level completion must not be recorded onto a caller-supplied outer span"
4013            );
4014        }
4015
4016        // --- Tool-result rewrites preserve raw policy data and redact telemetry ---
4017
4018        /// A tool that returns a raw marker; a rewrite hook replaces the
4019        /// effective model and telemetry presentation.
4020        struct RawOutputTool;
4021        impl crate::tool::Tool for RawOutputTool {
4022            const NAME: &'static str = "raw_output";
4023            type Error = rig::tool::ToolExecutionError;
4024            type Args = serde_json::Value;
4025            type Output = String;
4026            fn description(&self) -> String {
4027                "returns a raw output marker".to_string()
4028            }
4029
4030            fn parameters(&self) -> serde_json::Value {
4031                serde_json::json!({ "type": "object", "properties": {} })
4032            }
4033            async fn call(
4034                &self,
4035                _context: &mut ToolContext,
4036                _args: Self::Args,
4037            ) -> Result<Self::Output, ToolExecutionError> {
4038                Ok("RAW_EXECUTION_OUTPUT_42".to_string())
4039            }
4040        }
4041
4042        /// Redacts every tool result before the model sees it.
4043        struct RedactResultHook;
4044        impl crate::agent::AgentHook for RedactResultHook {
4045            async fn on_tool_result(
4046                &self,
4047                _ctx: &HookContext,
4048                event: ToolResultEvent<'_>,
4049            ) -> ToolResultAction {
4050                if let crate::agent::ToolResultEvent { .. } = event {
4051                    crate::agent::ToolResultAction::rewrite("[REDACTED]")
4052                } else {
4053                    crate::agent::ToolResultAction::keep()
4054                }
4055            }
4056        }
4057
4058        /// Stops the run after observing a completed tool result.
4059        struct StopOnResultHook;
4060        impl crate::agent::AgentHook for StopOnResultHook {
4061            async fn on_tool_result(
4062                &self,
4063                _ctx: &HookContext,
4064                _event: ToolResultEvent<'_>,
4065            ) -> ToolResultAction {
4066                ToolResultAction::stop("stop after raw result")
4067            }
4068        }
4069
4070        /// Captures every value recorded into the `gen_ai.tool.call.result` span
4071        /// field, so tests can assert telemetry follows result-hook policy.
4072        #[derive(Default)]
4073        struct ResultValueVisitor {
4074            values: Vec<String>,
4075        }
4076        impl Visit for ResultValueVisitor {
4077            fn record_str(&mut self, field: &Field, value: &str) {
4078                if field.name() == "gen_ai.tool.call.result" {
4079                    self.values.push(value.to_string());
4080                }
4081            }
4082            fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
4083                if field.name() == "gen_ai.tool.call.result" {
4084                    self.values.push(format!("{value:?}"));
4085                }
4086            }
4087        }
4088
4089        struct ResultValueLayer {
4090            values: Arc<Mutex<Vec<String>>>,
4091        }
4092        impl<S> Layer<S> for ResultValueLayer
4093        where
4094            S: Subscriber + for<'l> LookupSpan<'l>,
4095        {
4096            fn on_record(&self, _span: &Id, values: &Record<'_>, _ctx: Context<'_, S>) {
4097                let mut visitor = ResultValueVisitor::default();
4098                values.record(&mut visitor);
4099                if !visitor.values.is_empty() {
4100                    self.values.lock().expect("values").extend(visitor.values);
4101                }
4102            }
4103        }
4104
4105        /// A `ToolResult` rewrite applies to both model presentation and
4106        /// telemetry so redaction hooks cannot leak the raw output through spans.
4107        #[tokio::test]
4108        async fn tool_result_rewrite_redacts_span_output() {
4109            let _isolation = crate::test_utils::scoped_tracing_subscriber_guard().await;
4110            let values: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
4111            let subscriber = Registry::default().with(ResultValueLayer {
4112                values: values.clone(),
4113            });
4114            let _default = tracing::subscriber::set_default(subscriber);
4115
4116            // Warm the `execute_tool` result callsite under this subscriber, then
4117            // reset — mirroring the usage tests' interest-cache warm-up.
4118            warm_blocking_callsites().await;
4119            tracing::callsite::rebuild_interest_cache();
4120            values.lock().expect("values").clear();
4121
4122            let model = MockCompletionModel::from_turns([
4123                MockTurn::tool_call("tc1", "raw_output", serde_json::json!({})),
4124                MockTurn::text("ok"),
4125            ]);
4126            let response = AgentBuilder::new(model)
4127                .record_content_telemetry(true)
4128                .tool(RawOutputTool)
4129                .add_hook(RedactResultHook)
4130                .build()
4131                .runner("go")
4132                .max_turns(3)
4133                .run()
4134                .await
4135                .expect("run should succeed");
4136            assert_eq!(response.output, "ok");
4137
4138            let captured = values.lock().expect("values").clone();
4139            assert!(
4140                captured.iter().any(|v| v.contains("[REDACTED]")),
4141                "the rewritten presentation must reach telemetry; captured: {captured:?}"
4142            );
4143            assert!(
4144                !captured
4145                    .iter()
4146                    .any(|v| v.contains("RAW_EXECUTION_OUTPUT_42")),
4147                "the raw tool output must not leak through telemetry; captured: {captured:?}"
4148            );
4149        }
4150
4151        /// Stopping from the result hook retains outcome metadata but omits
4152        /// potentially sensitive result content from telemetry.
4153        #[tokio::test]
4154        async fn tool_result_stop_omits_span_output() {
4155            let _isolation = crate::test_utils::scoped_tracing_subscriber_guard().await;
4156            let values: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
4157            let subscriber = Registry::default().with(ResultValueLayer {
4158                values: values.clone(),
4159            });
4160            let _default = tracing::subscriber::set_default(subscriber);
4161
4162            warm_blocking_callsites().await;
4163            tracing::callsite::rebuild_interest_cache();
4164            values.lock().expect("values").clear();
4165
4166            let result = AgentBuilder::new(MockCompletionModel::from_turns([MockTurn::tool_call(
4167                "tc1",
4168                "raw_output",
4169                serde_json::json!({}),
4170            )]))
4171            .tool(RawOutputTool)
4172            .add_hook(StopOnResultHook)
4173            .build()
4174            .runner("go")
4175            .max_turns(2)
4176            .run()
4177            .await;
4178            assert!(result.is_err(), "the result hook should stop the run");
4179
4180            let captured = values.lock().expect("values").clone();
4181            assert!(
4182                !captured
4183                    .iter()
4184                    .any(|value| value.contains("RAW_EXECUTION_OUTPUT_42")),
4185                "a Stop must not leak raw execution telemetry; captured: {captured:?}"
4186            );
4187        }
4188    }
4189
4190    fn tool_call_content(id: &str, args: serde_json::Value) -> AssistantContent {
4191        AssistantContent::ToolCall(MessageToolCall::new(
4192            id.to_string(),
4193            ToolFunction::new("add".to_string(), args),
4194        ))
4195    }
4196
4197    /// Whether any tool result in `messages` carries `expected` as verbatim text.
4198    /// Used to pin a skip reason's actual value (a reason dropped or altered on
4199    /// both drivers would still satisfy a blocking == streaming equality check).
4200    fn tool_result_text_in_history(messages: &[Message], expected: &str) -> bool {
4201        messages.iter().any(|message| {
4202            matches!(
4203                message,
4204                Message::User { content }
4205                    if content.iter().any(|item| matches!(
4206                        item,
4207                        UserContent::ToolResult(result)
4208                            if result.content.iter().any(|c| matches!(
4209                                c,
4210                                rig_core::message::ToolResultContent::Text(text)
4211                                    if text.text == expected
4212                            ))
4213                    ))
4214            )
4215        })
4216    }
4217
4218    /// Whether any tool result in `messages` carries the exact structured JSON value.
4219    fn tool_result_json_in_history(messages: &[Message], expected: &serde_json::Value) -> bool {
4220        messages.iter().any(|message| {
4221            matches!(
4222                message,
4223                Message::User { content }
4224                    if content.iter().any(|item| matches!(
4225                        item,
4226                        UserContent::ToolResult(result)
4227                            if result.content.iter().any(|content| matches!(
4228                                content,
4229                                rig_core::message::ToolResultContent::Json { value }
4230                                    if value == expected
4231                            ))
4232                    ))
4233            )
4234        })
4235    }
4236
4237    /// Even with `run()` executing tools concurrently, the tool-result order —
4238    /// and so the whole message history — matches the sequential streaming
4239    /// driver. (`run()` runs tools with `buffer_unordered` but writes each result
4240    /// into its original call-index slot, so results still land in call order.)
4241    #[tokio::test]
4242    async fn run_and_stream_same_message_history_for_parallel_tool_calls() {
4243        let blocking_model = MockCompletionModel::from_turns([
4244            MockTurn::from_contents([
4245                tool_call_content("tc1", json!({"x": 2, "y": 3})),
4246                tool_call_content("tc2", json!({"x": 10, "y": 20})),
4247            ])
4248            .expect("two tool calls is a valid turn"),
4249            MockTurn::text("done"),
4250        ]);
4251        let blocking = AgentBuilder::new(blocking_model)
4252            .tool(MockAddTool)
4253            .build()
4254            .runner("add two pairs")
4255            .max_turns(3)
4256            .tool_concurrency(4)
4257            .run()
4258            .await
4259            .expect("blocking run should succeed");
4260
4261        let streaming_model = MockCompletionModel::from_stream_turns([
4262            vec![
4263                MockStreamEvent::tool_call("tc1", "add", json!({"x": 2, "y": 3})),
4264                MockStreamEvent::tool_call("tc2", "add", json!({"x": 10, "y": 20})),
4265                MockStreamEvent::final_response_with_total_tokens(0),
4266            ],
4267            vec![
4268                MockStreamEvent::text("done"),
4269                MockStreamEvent::final_response_with_total_tokens(0),
4270            ],
4271        ]);
4272        let mut stream = AgentBuilder::new(streaming_model)
4273            .tool(MockAddTool)
4274            .build()
4275            .runner("add two pairs")
4276            .max_turns(3)
4277            .stream()
4278            .await;
4279        let mut final_response = None;
4280        while let Some(item) = stream.next().await {
4281            if let Ok(MultiTurnStreamItem::FinalResponse(resp)) =
4282                item.map_err(|err| panic!("stream item errored: {err}"))
4283            {
4284                final_response = Some(resp);
4285            }
4286        }
4287        let final_response = final_response.expect("stream should yield a final response");
4288
4289        let blocking_messages = blocking.messages.expect("blocking messages");
4290        let streaming_messages = final_response
4291            .messages()
4292            .expect("streaming history")
4293            .to_vec();
4294        assert_eq!(
4295            serde_json::to_value(&blocking_messages).expect("serialize blocking"),
4296            serde_json::to_value(&streaming_messages).expect("serialize streaming"),
4297        );
4298    }
4299
4300    /// A tool whose first-*called* invocation completes *after* the second, so
4301    /// `buffer_unordered` yields the results in completion order — yet the
4302    /// persisted history stays in call order because each result is written into
4303    /// its original call-index slot. The first call (in poll/call order) waits on
4304    /// a gate the second call releases.
4305    #[derive(Clone)]
4306    struct OutOfOrderTool {
4307        gate: Arc<tokio::sync::Notify>,
4308        order: Arc<AtomicU32>,
4309    }
4310
4311    impl Tool for OutOfOrderTool {
4312        const NAME: &'static str = "add";
4313        type Error = MockToolError;
4314        type Args = MockOperationArgs;
4315        type Output = i32;
4316
4317        fn description(&self) -> String {
4318            MockAddTool.description()
4319        }
4320
4321        fn parameters(&self) -> serde_json::Value {
4322            MockAddTool.parameters()
4323        }
4324
4325        async fn call(
4326            &self,
4327            _context: &mut ToolContext,
4328            _args: Self::Args,
4329        ) -> Result<Self::Output, Self::Error> {
4330            let nth = self.order.fetch_add(1, SeqCst);
4331            if nth == 0 {
4332                // First call: cannot finish until a later call releases us.
4333                self.gate.notified().await;
4334            } else {
4335                // Later call: finishes immediately and releases the first.
4336                self.gate.notify_one();
4337            }
4338            Ok(nth as i32)
4339        }
4340    }
4341
4342    /// `run()` must persist tool results in tool-call (emission) order even when
4343    /// tools complete out of order under concurrency — it runs them with
4344    /// `buffer_unordered` but reindexes each result into its original call-index
4345    /// slot. (This is what keeps its message history identical to the sequential
4346    /// streaming driver.)
4347    #[tokio::test]
4348    async fn run_preserves_tool_call_order_under_out_of_order_completion() {
4349        let model = MockCompletionModel::from_turns([
4350            MockTurn::from_contents([
4351                tool_call_content("tc1", json!({"x": 1, "y": 0})),
4352                tool_call_content("tc2", json!({"x": 2, "y": 0})),
4353            ])
4354            .expect("two tool calls is a valid turn"),
4355            MockTurn::text("done"),
4356        ]);
4357        let response = AgentBuilder::new(model)
4358            .tool(OutOfOrderTool {
4359                gate: Arc::new(tokio::sync::Notify::new()),
4360                order: Arc::new(AtomicU32::new(0)),
4361            })
4362            .build()
4363            .runner("go")
4364            .max_turns(3)
4365            .tool_concurrency(4)
4366            .run()
4367            .await
4368            .expect("run should succeed");
4369
4370        let messages = response.messages.expect("messages");
4371        let result_ids: Vec<String> = messages
4372            .iter()
4373            .flat_map(|message| match message {
4374                Message::User { content } => content
4375                    .iter()
4376                    .filter_map(|item| match item {
4377                        UserContent::ToolResult(result) => Some(result.id.clone()),
4378                        _ => None,
4379                    })
4380                    .collect::<Vec<_>>(),
4381                _ => Vec::new(),
4382            })
4383            .collect();
4384        // Call order (tc1 then tc2), even though tc2 finished first.
4385        assert_eq!(result_ids, vec!["tc1".to_string(), "tc2".to_string()]);
4386    }
4387
4388    /// Drive a stream to completion, panicking on any stream error, and return
4389    /// its final response.
4390    async fn drive_to_final_response<R: Send + 'static>(
4391        mut stream: crate::agent::prompt_request::streaming::StreamingResult<R>,
4392    ) -> crate::agent::prompt_request::PromptResponse {
4393        let mut final_response = None;
4394        while let Some(item) = stream.next().await {
4395            if let MultiTurnStreamItem::FinalResponse(resp) =
4396                item.unwrap_or_else(|err| panic!("stream item errored: {err}"))
4397            {
4398                final_response = Some(resp);
4399            }
4400        }
4401        final_response.expect("stream should yield a final response")
4402    }
4403
4404    /// Tool-result ids, in history order, across a run's message history.
4405    fn tool_result_ids(messages: &[Message]) -> Vec<String> {
4406        messages
4407            .iter()
4408            .flat_map(|message| match message {
4409                Message::User { content } => content
4410                    .iter()
4411                    .filter_map(|item| match item {
4412                        UserContent::ToolResult(result) => Some(result.id.clone()),
4413                        _ => None,
4414                    })
4415                    .collect::<Vec<_>>(),
4416                _ => Vec::new(),
4417            })
4418            .collect()
4419    }
4420
4421    /// The streaming driver under `tool_concurrency > 1` produces the **same
4422    /// message history** as the blocking driver: streamed results are surfaced in
4423    /// call order after the batch settles, and persisted results stay in tool-call
4424    /// order, so concurrency never reorders the final history.
4425    #[tokio::test]
4426    async fn stream_and_run_same_message_history_for_parallel_tool_calls_under_concurrency() {
4427        let blocking_model = MockCompletionModel::from_turns([
4428            MockTurn::from_contents([
4429                tool_call_content("tc1", json!({"x": 2, "y": 3})),
4430                tool_call_content("tc2", json!({"x": 10, "y": 20})),
4431            ])
4432            .expect("two tool calls is a valid turn"),
4433            MockTurn::text("done"),
4434        ]);
4435        let blocking = AgentBuilder::new(blocking_model)
4436            .tool(MockAddTool)
4437            .build()
4438            .runner("add two pairs")
4439            .max_turns(3)
4440            .tool_concurrency(4)
4441            .run()
4442            .await
4443            .expect("blocking run should succeed");
4444
4445        let streaming_model = MockCompletionModel::from_stream_turns([
4446            vec![
4447                MockStreamEvent::tool_call("tc1", "add", json!({"x": 2, "y": 3})),
4448                MockStreamEvent::tool_call("tc2", "add", json!({"x": 10, "y": 20})),
4449                MockStreamEvent::final_response_with_total_tokens(0),
4450            ],
4451            vec![
4452                MockStreamEvent::text("done"),
4453                MockStreamEvent::final_response_with_total_tokens(0),
4454            ],
4455        ]);
4456        let stream = AgentBuilder::new(streaming_model)
4457            .tool(MockAddTool)
4458            .build()
4459            .runner("add two pairs")
4460            .max_turns(3)
4461            .tool_concurrency(4)
4462            .stream()
4463            .await;
4464        let final_response = drive_to_final_response(stream).await;
4465
4466        let blocking_messages = blocking.messages.expect("blocking messages");
4467        let streaming_messages = final_response
4468            .messages()
4469            .expect("streaming history")
4470            .to_vec();
4471        assert_eq!(
4472            serde_json::to_value(&blocking_messages).expect("serialize blocking"),
4473            serde_json::to_value(&streaming_messages).expect("serialize streaming"),
4474        );
4475    }
4476
4477    /// The streaming driver under concurrency persists tool results in **call
4478    /// order** even when tools complete out of order. `OutOfOrderTool`'s
4479    /// first-called invocation only finishes once the second runs, so this also
4480    /// proves the tools run concurrently: sequential execution would deadlock on
4481    /// the first call.
4482    #[tokio::test]
4483    async fn stream_preserves_history_order_under_out_of_order_completion() {
4484        let model = MockCompletionModel::from_stream_turns([
4485            vec![
4486                MockStreamEvent::tool_call("tc1", "add", json!({"x": 1, "y": 0})),
4487                MockStreamEvent::tool_call("tc2", "add", json!({"x": 2, "y": 0})),
4488                MockStreamEvent::final_response_with_total_tokens(0),
4489            ],
4490            vec![
4491                MockStreamEvent::text("done"),
4492                MockStreamEvent::final_response_with_total_tokens(0),
4493            ],
4494        ]);
4495        let stream = AgentBuilder::new(model)
4496            .tool(OutOfOrderTool {
4497                gate: Arc::new(tokio::sync::Notify::new()),
4498                order: Arc::new(AtomicU32::new(0)),
4499            })
4500            .build()
4501            .runner("go")
4502            .max_turns(3)
4503            .tool_concurrency(4)
4504            .stream()
4505            .await;
4506        // Timeout so a regression to sequential execution fails cleanly instead
4507        // of hanging (the first call only completes once the second runs).
4508        let final_response = tokio::time::timeout(
4509            std::time::Duration::from_secs(5),
4510            drive_to_final_response(stream),
4511        )
4512        .await
4513        .expect("streamed tools must run concurrently, not deadlock on the first call");
4514
4515        let messages = final_response.messages().expect("history").to_vec();
4516        // History stays in call order (tc1 then tc2), even though tc2 finished first.
4517        assert_eq!(
4518            tool_result_ids(&messages),
4519            vec!["tc1".to_string(), "tc2".to_string()]
4520        );
4521    }
4522
4523    /// Under concurrency the streaming driver surfaces tool results **atomically
4524    /// after the whole batch settles**, in **call order** — not as each tool
4525    /// completes. The second call completes first (via the gate), yet its result
4526    /// is still surfaced second, matching persisted history order.
4527    #[tokio::test]
4528    async fn stream_emits_tool_results_in_call_order_after_batch_settles_under_concurrency() {
4529        let model = MockCompletionModel::from_stream_turns([
4530            vec![
4531                MockStreamEvent::tool_call("tc1", "add", json!({"x": 1, "y": 0})),
4532                MockStreamEvent::tool_call("tc2", "add", json!({"x": 2, "y": 0})),
4533                MockStreamEvent::final_response_with_total_tokens(0),
4534            ],
4535            vec![
4536                MockStreamEvent::text("done"),
4537                MockStreamEvent::final_response_with_total_tokens(0),
4538            ],
4539        ]);
4540        let mut stream = AgentBuilder::new(model)
4541            .tool(OutOfOrderTool {
4542                gate: Arc::new(tokio::sync::Notify::new()),
4543                order: Arc::new(AtomicU32::new(0)),
4544            })
4545            .build()
4546            .runner("go")
4547            .max_turns(3)
4548            .tool_concurrency(4)
4549            .stream()
4550            .await;
4551
4552        let mut streamed_result_ids = Vec::new();
4553        let mut final_response = None;
4554        tokio::time::timeout(std::time::Duration::from_secs(5), async {
4555            while let Some(item) = stream.next().await {
4556                match item.unwrap_or_else(|err| panic!("stream item errored: {err}")) {
4557                    MultiTurnStreamItem::StreamUserItem(StreamedUserContent::ToolResult {
4558                        tool_result,
4559                        ..
4560                    }) => streamed_result_ids.push(tool_result.id),
4561                    MultiTurnStreamItem::FinalResponse(resp) => final_response = Some(resp),
4562                    _ => {}
4563                }
4564            }
4565        })
4566        .await
4567        .expect("streamed tools must run concurrently, not deadlock on the first call");
4568
4569        // Call order, even though tc2 completed first — results are surfaced only
4570        // after the whole batch settles.
4571        assert_eq!(
4572            streamed_result_ids,
4573            vec!["tc1".to_string(), "tc2".to_string()]
4574        );
4575        let final_response = final_response.expect("stream should yield a final response");
4576        assert_eq!(
4577            tool_result_ids(final_response.messages().expect("history")),
4578            vec!["tc1".to_string(), "tc2".to_string()]
4579        );
4580    }
4581
4582    /// Two barrier-synchronized tools in one streamed turn finish only if they
4583    /// run concurrently — each waits at the barrier for the other. At
4584    /// `tool_concurrency(2)` the streamed turn completes; sequential execution
4585    /// would block on the first call forever, so the timeout asserts genuine
4586    /// concurrency on the streaming path.
4587    #[tokio::test]
4588    async fn stream_executes_tools_concurrently_under_concurrency() {
4589        let barrier = Arc::new(tokio::sync::Barrier::new(2));
4590        let model = MockCompletionModel::from_stream_turns([
4591            vec![
4592                MockStreamEvent::tool_call("b1", "barrier_tool", json!({})),
4593                MockStreamEvent::tool_call("b2", "barrier_tool", json!({})),
4594                MockStreamEvent::final_response_with_total_tokens(0),
4595            ],
4596            vec![
4597                MockStreamEvent::text("done"),
4598                MockStreamEvent::final_response_with_total_tokens(0),
4599            ],
4600        ]);
4601        let stream = AgentBuilder::new(model)
4602            .tool(MockBarrierTool::new(barrier))
4603            .build()
4604            .runner("hit the barrier twice")
4605            .max_turns(3)
4606            .tool_concurrency(2)
4607            .stream()
4608            .await;
4609
4610        tokio::time::timeout(
4611            std::time::Duration::from_secs(5),
4612            drive_to_final_response(stream),
4613        )
4614        .await
4615        .expect("streamed tools must run concurrently, not deadlock at the barrier");
4616    }
4617
4618    /// The stream-item taxonomy and ordering: the driver emits *all* of a turn's
4619    /// **model** tool-call items ([`StreamedAssistantContent::ToolCall`], one per
4620    /// call the model made) first, then — after the whole tool batch settles —
4621    /// the per-tool **execution** items (`ToolExecutionCommitted` then the
4622    /// `ToolResult`) in call order. This holds identically at every concurrency
4623    /// (the batch is atomic on both the sequential and concurrent paths).
4624    #[tokio::test]
4625    async fn stream_emits_model_tool_calls_then_atomic_execution_items() {
4626        async fn markers(concurrency: usize) -> Vec<&'static str> {
4627            let model = MockCompletionModel::from_stream_turns([
4628                vec![
4629                    MockStreamEvent::tool_call("tc1", "add", json!({"x": 1, "y": 1})),
4630                    MockStreamEvent::tool_call("tc2", "add", json!({"x": 2, "y": 2})),
4631                    MockStreamEvent::final_response_with_total_tokens(0),
4632                ],
4633                vec![
4634                    MockStreamEvent::text("done"),
4635                    MockStreamEvent::final_response_with_total_tokens(0),
4636                ],
4637            ]);
4638            let mut stream = AgentBuilder::new(model)
4639                .tool(MockAddTool)
4640                .build()
4641                .runner("add two pairs")
4642                .max_turns(3)
4643                .tool_concurrency(concurrency)
4644                .stream()
4645                .await;
4646            let mut markers = Vec::new();
4647            while let Some(item) = stream.next().await {
4648                match item.unwrap_or_else(|err| panic!("stream item errored: {err}")) {
4649                    MultiTurnStreamItem::StreamAssistantItem(
4650                        StreamedAssistantContent::ToolCall { .. },
4651                    ) => markers.push("model-call"),
4652                    MultiTurnStreamItem::ToolExecutionCommitted { .. } => {
4653                        markers.push("exec-commit")
4654                    }
4655                    MultiTurnStreamItem::StreamUserItem(StreamedUserContent::ToolResult {
4656                        ..
4657                    }) => markers.push("result"),
4658                    _ => {}
4659                }
4660            }
4661            markers
4662        }
4663
4664        // Both surfaces: all model tool calls first, then per-tool (start, result)
4665        // in call order, surfaced atomically after the batch settles.
4666        let expected = vec![
4667            "model-call",
4668            "model-call",
4669            "exec-commit",
4670            "result",
4671            "exec-commit",
4672            "result",
4673        ];
4674        assert_eq!(markers(1).await, expected);
4675        assert_eq!(markers(4).await, expected);
4676    }
4677
4678    /// Terminates from the `x == 1` tool's result, but only *after* the slow
4679    /// `x == 2` sibling has signalled it started executing — so that sibling is
4680    /// genuinely in flight when the terminate fires (not merely not-yet-started).
4681    struct TerminateAfterSiblingStartedHook {
4682        sibling_started: Arc<tokio::sync::Notify>,
4683    }
4684    impl AgentHook for TerminateAfterSiblingStartedHook {
4685        async fn on_tool_result(
4686            &self,
4687            _ctx: &HookContext,
4688            event: ToolResultEvent<'_>,
4689        ) -> ToolResultAction {
4690            if let ToolResultEvent { args, .. } = event
4691                && serde_json::from_str::<serde_json::Value>(args)
4692                    .ok()
4693                    .and_then(|v| v.get("x").and_then(serde_json::Value::as_i64))
4694                    == Some(1)
4695            {
4696                self.sibling_started.notified().await;
4697                return ToolResultAction::stop("stop after a tool result");
4698            }
4699            ToolResultAction::keep()
4700        }
4701    }
4702
4703    /// A probe tool for the concurrent drain path: records how many calls
4704    /// `started` and `completed`. The `x == 2` call signals it has started, then
4705    /// stays pending across several polls, so it is genuinely in flight — not
4706    /// merely not-yet-started — when the `x == 1` call's result terminates the
4707    /// run. A driver that **drains** the concurrent tool stream polls it to
4708    /// completion (`completed == 2`); one that **cancels** in-flight siblings
4709    /// would drop it mid-poll (`completed == 1`).
4710    #[derive(Clone)]
4711    struct DrainProbeTool {
4712        started: Arc<AtomicU32>,
4713        completed: Arc<AtomicU32>,
4714        slow_started: Arc<tokio::sync::Notify>,
4715    }
4716
4717    impl Tool for DrainProbeTool {
4718        const NAME: &'static str = "add";
4719        type Error = MockToolError;
4720        type Args = serde_json::Value;
4721        type Output = i32;
4722
4723        fn description(&self) -> String {
4724            MockAddTool.description()
4725        }
4726
4727        fn parameters(&self) -> serde_json::Value {
4728            MockAddTool.parameters()
4729        }
4730
4731        async fn call(
4732            &self,
4733            _context: &mut ToolContext,
4734            args: Self::Args,
4735        ) -> Result<Self::Output, Self::Error> {
4736            self.started.fetch_add(1, SeqCst);
4737            if args.get("x").and_then(serde_json::Value::as_i64) == Some(2) {
4738                // Signal that the slow sibling has started, then stay pending so
4739                // it is still executing when the fast call's result terminates.
4740                self.slow_started.notify_one();
4741                for _ in 0..8 {
4742                    tokio::task::yield_now().await;
4743                }
4744            }
4745            self.completed.fetch_add(1, SeqCst);
4746            Ok(0)
4747        }
4748    }
4749
4750    /// On the concurrent path, a terminate surfaces a `StreamingError`, ends the
4751    /// run with no final response, and — for a sibling that is **already in
4752    /// flight** — drains it to completion rather than cancelling it mid-poll (so
4753    /// no detached task is left running and the deterministic terminate reason
4754    /// still surfaces). The `x == 2` sibling signals it started before the
4755    /// `x == 1` result terminates, so `completed == 2` holds only under drain.
4756    #[tokio::test]
4757    async fn stream_concurrent_tool_result_terminate_drains_in_flight_siblings() {
4758        let started = Arc::new(AtomicU32::new(0));
4759        let completed = Arc::new(AtomicU32::new(0));
4760        let slow_started = Arc::new(tokio::sync::Notify::new());
4761        let model = MockCompletionModel::from_stream_turns([
4762            vec![
4763                MockStreamEvent::tool_call("tc1", "add", json!({"x": 1, "y": 1})),
4764                MockStreamEvent::tool_call("tc2", "add", json!({"x": 2, "y": 2})),
4765                MockStreamEvent::final_response_with_total_tokens(0),
4766            ],
4767            vec![
4768                MockStreamEvent::text("done"),
4769                MockStreamEvent::final_response_with_total_tokens(0),
4770            ],
4771        ]);
4772        let mut stream = AgentBuilder::new(model)
4773            .tool(DrainProbeTool {
4774                started: started.clone(),
4775                completed: completed.clone(),
4776                slow_started: slow_started.clone(),
4777            })
4778            .build()
4779            .runner("add two pairs")
4780            .max_turns(3)
4781            .tool_concurrency(2)
4782            .add_hook(TerminateAfterSiblingStartedHook {
4783                sibling_started: slow_started,
4784            })
4785            .stream()
4786            .await;
4787
4788        let (saw_error, saw_final_response) =
4789            tokio::time::timeout(std::time::Duration::from_secs(5), async move {
4790                let mut saw_error = false;
4791                let mut saw_final_response = false;
4792                while let Some(item) = stream.next().await {
4793                    match item {
4794                        Ok(MultiTurnStreamItem::FinalResponse(_)) => saw_final_response = true,
4795                        Ok(_) => {}
4796                        Err(StreamingError::Prompt(_)) => saw_error = true,
4797                        Err(other) => panic!("unexpected streaming error: {other}"),
4798                    }
4799                }
4800                (saw_error, saw_final_response)
4801            })
4802            .await
4803            .expect("draining the concurrent tools must not hang");
4804
4805        assert!(
4806            saw_error,
4807            "a terminate hook on the concurrent path must surface a StreamingError::Prompt"
4808        );
4809        assert!(
4810            !saw_final_response,
4811            "a terminated run must not yield a final response"
4812        );
4813        // The already-in-flight slow sibling is drained to completion, not
4814        // cancelled mid-poll (which would leave `completed == 1`).
4815        assert_eq!(
4816            started.load(SeqCst),
4817            2,
4818            "both tools started (both in flight)"
4819        );
4820        assert_eq!(
4821            completed.load(SeqCst),
4822            2,
4823            "the in-flight sibling must be drained to completion, not cancelled"
4824        );
4825    }
4826
4827    /// A the event-specific stop action from the `ToolCall` event with a reason keyed by the
4828    /// call's `x` arg, forcing the `x == 2` call (tc2) to terminate *before* the
4829    /// `x == 1` call (tc1): tc2 opens the gate after terminating, tc1 awaits it
4830    /// first. So completion order (tc2) differs from call order (tc1).
4831    struct OrderedTerminateHook {
4832        gate: Arc<tokio::sync::Notify>,
4833    }
4834
4835    impl AgentHook for OrderedTerminateHook {
4836        async fn on_tool_call(&self, _ctx: &HookContext, event: ToolCall<'_>) -> ToolCallAction {
4837            if let ToolCall { args, .. } = event {
4838                let x = serde_json::from_str::<serde_json::Value>(args)
4839                    .ok()
4840                    .and_then(|v| v.get("x").and_then(serde_json::Value::as_i64));
4841                match x {
4842                    Some(2) => {
4843                        self.gate.notify_one();
4844                        return ToolCallAction::stop("terminated-by-tc2".to_string());
4845                    }
4846                    Some(1) => {
4847                        self.gate.notified().await;
4848                        return ToolCallAction::stop("terminated-by-tc1".to_string());
4849                    }
4850                    _ => {}
4851                }
4852            }
4853            ToolCallAction::run()
4854        }
4855    }
4856
4857    fn two_terminating_tools_blocking_model() -> MockCompletionModel {
4858        MockCompletionModel::from_turns([
4859            MockTurn::from_contents([
4860                tool_call_content("tc1", json!({"x": 1, "y": 1})),
4861                tool_call_content("tc2", json!({"x": 2, "y": 2})),
4862            ])
4863            .expect("two tool calls is non-empty"),
4864            MockTurn::text("unreachable"),
4865        ])
4866    }
4867
4868    fn two_terminating_tools_streaming_model() -> MockCompletionModel {
4869        MockCompletionModel::from_stream_turns([
4870            vec![
4871                MockStreamEvent::tool_call("tc1", "add", json!({"x": 1, "y": 1})),
4872                MockStreamEvent::tool_call("tc2", "add", json!({"x": 2, "y": 2})),
4873                MockStreamEvent::final_response_with_total_tokens(0),
4874            ],
4875            vec![
4876                MockStreamEvent::text("unreachable"),
4877                MockStreamEvent::final_response_with_total_tokens(0),
4878            ],
4879        ])
4880    }
4881
4882    /// When two tool calls in one turn both terminate the run under
4883    /// `tool_concurrency > 1`, run() and stream() surface the **same** reason —
4884    /// the first-called tool's (call order), not whichever finished first. tc2
4885    /// terminates before tc1, so a completion-order pick would surface tc2's
4886    /// reason and the two drivers would disagree.
4887    #[tokio::test]
4888    async fn concurrent_simultaneous_tool_terminations_pick_call_order_on_both_drivers() {
4889        let run_err = tokio::time::timeout(
4890            std::time::Duration::from_secs(5),
4891            AgentBuilder::new(two_terminating_tools_blocking_model())
4892                .tool(MockAddTool)
4893                .build()
4894                .runner("go")
4895                .max_turns(3)
4896                .tool_concurrency(2)
4897                .add_hook(OrderedTerminateHook {
4898                    gate: Arc::new(tokio::sync::Notify::new()),
4899                })
4900                .run(),
4901        )
4902        .await
4903        .expect("blocking run must not hang")
4904        .expect_err("the run must terminate");
4905
4906        let mut stream = AgentBuilder::new(two_terminating_tools_streaming_model())
4907            .tool(MockAddTool)
4908            .build()
4909            .runner("go")
4910            .max_turns(3)
4911            .tool_concurrency(2)
4912            .add_hook(OrderedTerminateHook {
4913                gate: Arc::new(tokio::sync::Notify::new()),
4914            })
4915            .stream()
4916            .await;
4917
4918        let stream_err = tokio::time::timeout(std::time::Duration::from_secs(5), async move {
4919            while let Some(item) = stream.next().await {
4920                if let Err(err) = item {
4921                    return Some(err);
4922                }
4923            }
4924            None
4925        })
4926        .await
4927        .expect("streamed run must not hang")
4928        .expect("the stream must surface a terminate error");
4929
4930        let run_msg = run_err.to_string();
4931        let stream_msg = stream_err.to_string();
4932        assert!(
4933            run_msg.contains("terminated-by-tc1"),
4934            "blocking run should surface the first-called tool's reason, got: {run_msg}"
4935        );
4936        assert!(
4937            stream_msg.contains("terminated-by-tc1"),
4938            "stream should surface the first-called tool's reason, got: {stream_msg}"
4939        );
4940        assert!(
4941            !run_msg.contains("terminated-by-tc2") && !stream_msg.contains("terminated-by-tc2"),
4942            "neither driver should surface the later-completing tool's reason"
4943        );
4944    }
4945
4946    /// Terminates the run from the `ToolCall` event of the first tool only
4947    /// (`x == 1`), letting any later tool through.
4948    struct TerminateOnFirstToolHook;
4949    impl AgentHook for TerminateOnFirstToolHook {
4950        async fn on_tool_call(&self, _ctx: &HookContext, event: ToolCall<'_>) -> ToolCallAction {
4951            if let ToolCall { args, .. } = event
4952                && serde_json::from_str::<serde_json::Value>(args)
4953                    .ok()
4954                    .and_then(|v| v.get("x").and_then(serde_json::Value::as_i64))
4955                    == Some(1)
4956            {
4957                return ToolCallAction::stop("stop".to_string());
4958            }
4959            ToolCallAction::run()
4960        }
4961    }
4962
4963    /// Fail-fast, lock-step across surfaces: on a multi-tool turn whose first
4964    /// tool's hook terminates the run, the SEQUENTIAL default (`tool_concurrency`
4965    /// == 1) surfaces the terminate immediately and does **not** start the
4966    /// remaining sibling tools — so tool B's side effect never runs. The
4967    /// terminating tool's own body never runs either (its `ToolCall` hook fired
4968    /// first), so `calls == 0` on both drivers, which share the tool driver.
4969    #[tokio::test]
4970    async fn default_concurrency_terminate_skips_remaining_tools_on_both_drivers() {
4971        let blocking_calls = Arc::new(AtomicU32::new(0));
4972        AgentBuilder::new(two_terminating_tools_blocking_model())
4973            .tool(CountingAddTool {
4974                calls: blocking_calls.clone(),
4975            })
4976            .build()
4977            .runner("go")
4978            .max_turns(3)
4979            .add_hook(TerminateOnFirstToolHook)
4980            .run()
4981            .await
4982            .expect_err("the run terminates");
4983        assert_eq!(
4984            blocking_calls.load(SeqCst),
4985            0,
4986            "fail-fast: blocking run() must not start the second tool after the first terminates"
4987        );
4988
4989        let streaming_calls = Arc::new(AtomicU32::new(0));
4990        let mut stream = AgentBuilder::new(two_terminating_tools_streaming_model())
4991            .tool(CountingAddTool {
4992                calls: streaming_calls.clone(),
4993            })
4994            .build()
4995            .runner("go")
4996            .max_turns(3)
4997            .add_hook(TerminateOnFirstToolHook)
4998            .stream()
4999            .await;
5000        let mut saw_error = false;
5001        while let Some(item) = stream.next().await {
5002            if let Err(err) = item {
5003                saw_error = true;
5004                assert!(
5005                    err.to_string().contains("stop"),
5006                    "stream() should surface the terminate reason, got: {err}"
5007                );
5008                break;
5009            }
5010        }
5011        assert!(saw_error, "stream() must surface the terminate error");
5012        assert_eq!(
5013            streaming_calls.load(SeqCst),
5014            0,
5015            "fail-fast: stream() must not start the second tool after the first terminates"
5016        );
5017    }
5018
5019    /// Records the `x` arg of every tool call that reaches its body. The `x == 1`
5020    /// sibling signals it has started (via `sibling_started`) and then stays
5021    /// pending across several polls, so it is genuinely in flight when the
5022    /// terminator (`x == 0`) fires — while a sibling beyond the concurrency
5023    /// window is not yet started and must be dropped.
5024    #[derive(Clone)]
5025    struct RecordingArgsTool {
5026        called: Arc<Mutex<Vec<i64>>>,
5027        sibling_started: Arc<tokio::sync::Notify>,
5028    }
5029
5030    impl Tool for RecordingArgsTool {
5031        const NAME: &'static str = "add";
5032        type Error = MockToolError;
5033        type Args = serde_json::Value;
5034        type Output = i32;
5035
5036        fn description(&self) -> String {
5037            MockAddTool.description()
5038        }
5039
5040        fn parameters(&self) -> serde_json::Value {
5041            MockAddTool.parameters()
5042        }
5043
5044        async fn call(
5045            &self,
5046            _context: &mut ToolContext,
5047            args: Self::Args,
5048        ) -> Result<Self::Output, Self::Error> {
5049            let x = args.get("x").and_then(serde_json::Value::as_i64);
5050            if let Some(x) = x {
5051                self.called.lock().expect("called").push(x);
5052            }
5053            if x == Some(1) {
5054                // Signal that the in-flight sibling has started, then stay pending
5055                // so it is still executing when the terminator fires.
5056                self.sibling_started.notify_one();
5057                for _ in 0..8 {
5058                    tokio::task::yield_now().await;
5059                }
5060            }
5061            Ok(0)
5062        }
5063    }
5064
5065    fn three_tools_first_terminates_streaming_model() -> MockCompletionModel {
5066        MockCompletionModel::from_stream_turns([
5067            vec![
5068                // tc0 (x==0) terminates on its ToolCall hook after the in-flight
5069                // sibling starts; tc1 (x==1) is the in-flight sibling (drains);
5070                // tc2 (x==2) is beyond the concurrency-2 window (not yet started)
5071                // and must be dropped once tc0 terminates.
5072                MockStreamEvent::tool_call("tc0", "add", json!({"x": 0, "y": 0})),
5073                MockStreamEvent::tool_call("tc1", "add", json!({"x": 1, "y": 1})),
5074                MockStreamEvent::tool_call("tc2", "add", json!({"x": 2, "y": 2})),
5075                MockStreamEvent::final_response_with_total_tokens(0),
5076            ],
5077            vec![
5078                MockStreamEvent::text("unreachable"),
5079                MockStreamEvent::final_response_with_total_tokens(0),
5080            ],
5081        ])
5082    }
5083
5084    /// Terminates from the `x == 0` tool's `ToolCall` hook, but only after the
5085    /// `x == 1` sibling has signalled it started executing — so tc1 is genuinely
5086    /// in flight (not merely not-yet-started) when the terminate fires.
5087    struct TerminateOnArgZeroAfterSiblingHook {
5088        sibling_started: Arc<tokio::sync::Notify>,
5089    }
5090    impl AgentHook for TerminateOnArgZeroAfterSiblingHook {
5091        async fn on_tool_call(&self, _ctx: &HookContext, event: ToolCall<'_>) -> ToolCallAction {
5092            if let ToolCall { args, .. } = event
5093                && serde_json::from_str::<serde_json::Value>(args)
5094                    .ok()
5095                    .and_then(|v| v.get("x").and_then(serde_json::Value::as_i64))
5096                    == Some(0)
5097            {
5098                self.sibling_started.notified().await;
5099                return ToolCallAction::stop("stop");
5100            }
5101            ToolCallAction::run()
5102        }
5103    }
5104
5105    /// Concurrent fail-fast: when a tool terminates the turn under
5106    /// `tool_concurrency > 1`, an **already-in-flight** sibling is drained while a
5107    /// sibling **beyond the concurrency window** — not yet started — is dropped.
5108    /// With concurrency 2 and three tools: tc0 (`x == 0`) terminates only after
5109    /// tc1 (`x == 1`) has started, so tc1 is genuinely in flight and drains
5110    /// (`called` contains 1); tc2 (`x == 2`) is pulled only after tc0 frees a slot
5111    /// — by which time the run is terminating — so it is dropped (`called` never
5112    /// contains 2), and tc0's own body never runs (its `ToolCall` hook terminated).
5113    /// The pre-fix run-all-then-decide would have executed tc2 too.
5114    #[tokio::test]
5115    async fn concurrent_terminate_drops_beyond_window_sibling_but_drains_in_flight() {
5116        let called = Arc::new(Mutex::new(Vec::new()));
5117        let sibling_started = Arc::new(tokio::sync::Notify::new());
5118        let mut stream = AgentBuilder::new(three_tools_first_terminates_streaming_model())
5119            .tool(RecordingArgsTool {
5120                called: called.clone(),
5121                sibling_started: sibling_started.clone(),
5122            })
5123            .build()
5124            .runner("go")
5125            .max_turns(3)
5126            .tool_concurrency(2)
5127            .add_hook(TerminateOnArgZeroAfterSiblingHook { sibling_started })
5128            .stream()
5129            .await;
5130
5131        let (saw_error, saw_final) =
5132            tokio::time::timeout(std::time::Duration::from_secs(5), async move {
5133                let mut saw_error = false;
5134                let mut saw_final = false;
5135                while let Some(item) = stream.next().await {
5136                    match item {
5137                        Ok(MultiTurnStreamItem::FinalResponse(_)) => saw_final = true,
5138                        Ok(_) => {}
5139                        Err(_) => saw_error = true,
5140                    }
5141                }
5142                (saw_error, saw_final)
5143            })
5144            .await
5145            .expect("the concurrent tool drive must not hang");
5146
5147        assert!(saw_error, "the terminated run must surface an error");
5148        assert!(
5149            !saw_final,
5150            "a terminated run must not yield a final response"
5151        );
5152        let called = called.lock().expect("called").clone();
5153        assert!(
5154            called.contains(&1),
5155            "the in-flight sibling (x==1) must be drained to completion; called args: {called:?}"
5156        );
5157        assert!(
5158            !called.contains(&2),
5159            "the not-yet-started sibling beyond the concurrency window (x==2) must be \
5160             dropped, not executed; called args: {called:?}"
5161        );
5162        assert!(
5163            !called.contains(&0),
5164            "the terminator's own body never runs (its ToolCall hook terminated); \
5165             called args: {called:?}"
5166        );
5167    }
5168
5169    /// A tool that, for the `x == 1` call, records it ran and signals a gate; the
5170    /// terminating sibling waits on that gate so the `x == 1` call completes
5171    /// *before* the batch terminates.
5172    #[derive(Clone)]
5173    struct SignalOnRunTool {
5174        a_ran: Arc<AtomicU32>,
5175        a_done: Arc<tokio::sync::Notify>,
5176    }
5177    impl Tool for SignalOnRunTool {
5178        const NAME: &'static str = "add";
5179        type Error = MockToolError;
5180        type Args = serde_json::Value;
5181        type Output = i32;
5182        fn description(&self) -> String {
5183            MockAddTool.description()
5184        }
5185
5186        fn parameters(&self) -> serde_json::Value {
5187            MockAddTool.parameters()
5188        }
5189        async fn call(
5190            &self,
5191            _context: &mut ToolContext,
5192            args: Self::Args,
5193        ) -> Result<Self::Output, Self::Error> {
5194            if args.get("x").and_then(serde_json::Value::as_i64) == Some(1) {
5195                self.a_ran.fetch_add(1, SeqCst);
5196                self.a_done.notify_one();
5197            }
5198            Ok(0)
5199        }
5200    }
5201
5202    /// The `x == 2` tool's `ToolCall` hook terminates, but only after the `x == 1`
5203    /// sibling has finished (via the gate), so a *completed* sibling's result is
5204    /// still suppressed by the atomic batch.
5205    struct TerminateAfterSiblingDoneHook {
5206        a_done: Arc<tokio::sync::Notify>,
5207    }
5208    impl AgentHook for TerminateAfterSiblingDoneHook {
5209        async fn on_tool_call(&self, _ctx: &HookContext, event: ToolCall<'_>) -> ToolCallAction {
5210            if let ToolCall { args, .. } = event
5211                && serde_json::from_str::<serde_json::Value>(args)
5212                    .ok()
5213                    .and_then(|v| v.get("x").and_then(serde_json::Value::as_i64))
5214                    == Some(2)
5215            {
5216                self.a_done.notified().await;
5217                return ToolCallAction::stop("stop");
5218            }
5219            ToolCallAction::run()
5220        }
5221    }
5222
5223    /// Atomic concurrent batch: when the batch terminates, even a sibling that
5224    /// completed **successfully** before the terminating sibling produces no
5225    /// `ToolExecutionCommitted` and no `ToolResult` stream item (no orphan
5226    /// execution-commit), and its result is not committed. The `x == 1` tool runs
5227    /// to completion (its side effect happens) and signals; the `x == 2` tool's
5228    /// hook then terminates.
5229    #[tokio::test]
5230    async fn concurrent_termination_surfaces_no_execution_items() {
5231        let a_ran = Arc::new(AtomicU32::new(0));
5232        let a_done = Arc::new(tokio::sync::Notify::new());
5233        let model = MockCompletionModel::from_stream_turns([
5234            vec![
5235                MockStreamEvent::tool_call("tc1", "add", json!({"x": 1, "y": 1})),
5236                MockStreamEvent::tool_call("tc2", "add", json!({"x": 2, "y": 2})),
5237                MockStreamEvent::final_response_with_total_tokens(0),
5238            ],
5239            vec![
5240                MockStreamEvent::text("unreachable"),
5241                MockStreamEvent::final_response_with_total_tokens(0),
5242            ],
5243        ]);
5244        let mut stream = AgentBuilder::new(model)
5245            .tool(SignalOnRunTool {
5246                a_ran: a_ran.clone(),
5247                a_done: a_done.clone(),
5248            })
5249            .build()
5250            .runner("go")
5251            .max_turns(3)
5252            .tool_concurrency(2)
5253            .add_hook(TerminateAfterSiblingDoneHook {
5254                a_done: a_done.clone(),
5255            })
5256            .stream()
5257            .await;
5258
5259        let (exec_commits, results, saw_error, saw_final) =
5260            tokio::time::timeout(std::time::Duration::from_secs(5), async move {
5261                let (mut exec_commits, mut results, mut saw_error, mut saw_final) =
5262                    (0, 0, false, false);
5263                while let Some(item) = stream.next().await {
5264                    match item {
5265                        Ok(MultiTurnStreamItem::ToolExecutionCommitted { .. }) => exec_commits += 1,
5266                        Ok(MultiTurnStreamItem::StreamUserItem(
5267                            StreamedUserContent::ToolResult { .. },
5268                        )) => results += 1,
5269                        Ok(MultiTurnStreamItem::FinalResponse(_)) => saw_final = true,
5270                        Ok(_) => {}
5271                        Err(_) => saw_error = true,
5272                    }
5273                }
5274                (exec_commits, results, saw_error, saw_final)
5275            })
5276            .await
5277            .expect("the concurrent tool drive must not hang");
5278
5279        assert!(saw_error, "the terminated run must surface an error");
5280        assert!(
5281            !saw_final,
5282            "a terminated run must not yield a final response"
5283        );
5284        assert_eq!(
5285            exec_commits, 0,
5286            "a terminated batch surfaces no ToolExecutionCommitted events"
5287        );
5288        assert_eq!(
5289            results, 0,
5290            "a terminated batch surfaces no successful ToolResult"
5291        );
5292        assert_eq!(
5293            a_ran.load(SeqCst),
5294            1,
5295            "the fast sibling did run (its side effect happened), but its result was suppressed"
5296        );
5297    }
5298
5299    /// The model tool-call event carries the model's **original** arguments; the
5300    /// execution-commit event carries the **effective** (hook-rewritten) arguments
5301    /// — so a `ToolCallAction::Rewrite` (e.g. a redaction) is reflected in what
5302    /// actually ran, not leaked as the original.
5303    #[tokio::test]
5304    async fn stream_tool_execution_committed_carries_effective_rewritten_args() {
5305        let model = MockCompletionModel::from_stream_turns([
5306            vec![
5307                MockStreamEvent::tool_call("tc1", "add", json!({"x": 2, "y": 3})),
5308                MockStreamEvent::final_response_with_total_tokens(0),
5309            ],
5310            vec![
5311                MockStreamEvent::text("done"),
5312                MockStreamEvent::final_response_with_total_tokens(0),
5313            ],
5314        ]);
5315        let mut stream = AgentBuilder::new(model)
5316            .tool(MockAddTool)
5317            .add_hook(RewriteToolArgsHook(json!({"x": 2, "y": 40})))
5318            .build()
5319            .runner("go")
5320            .max_turns(3)
5321            .stream()
5322            .await;
5323
5324        let mut model_args = None;
5325        let mut exec_args = None;
5326        while let Some(item) = stream.next().await {
5327            match item.unwrap_or_else(|err| panic!("stream item errored: {err}")) {
5328                MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::ToolCall {
5329                    tool_call,
5330                    ..
5331                }) => model_args = Some(tool_call.function.arguments),
5332                MultiTurnStreamItem::ToolExecutionCommitted { tool_call, .. } => {
5333                    exec_args = Some(tool_call.function.arguments)
5334                }
5335                _ => {}
5336            }
5337        }
5338        assert_eq!(
5339            model_args,
5340            Some(json!({"x": 2, "y": 3})),
5341            "the model tool-call event carries the model's original arguments"
5342        );
5343        assert_eq!(
5344            exec_args,
5345            Some(json!({"x": 2, "y": 40})),
5346            "the execution-commit event carries the hook-rewritten (effective) arguments"
5347        );
5348    }
5349
5350    /// A `ToolCall` hook `ToolCallAction::Skip` surfaces the skip result as a `ToolResult`
5351    /// (the model sees it, and it is committed to history) but produces **no**
5352    /// `ToolExecutionCommitted` — nothing actually ran.
5353    #[tokio::test]
5354    async fn stream_hook_skip_surfaces_result_without_execution_commit() {
5355        struct SkipHook;
5356        impl AgentHook for SkipHook {
5357            async fn on_tool_call(
5358                &self,
5359                _ctx: &HookContext,
5360                event: ToolCall<'_>,
5361            ) -> ToolCallAction {
5362                if let ToolCall { .. } = event {
5363                    ToolCallAction::skip("blocked by policy")
5364                } else {
5365                    ToolCallAction::run()
5366                }
5367            }
5368        }
5369
5370        let calls = Arc::new(AtomicU32::new(0));
5371        let model = MockCompletionModel::from_stream_turns([
5372            vec![
5373                MockStreamEvent::tool_call("tc1", "add", json!({"x": 1, "y": 2})),
5374                MockStreamEvent::final_response_with_total_tokens(0),
5375            ],
5376            vec![
5377                MockStreamEvent::text("done"),
5378                MockStreamEvent::final_response_with_total_tokens(0),
5379            ],
5380        ]);
5381        let stream = AgentBuilder::new(model)
5382            .tool(CountingAddTool {
5383                calls: calls.clone(),
5384            })
5385            .add_hook(SkipHook)
5386            .build()
5387            .runner("go")
5388            .max_turns(3)
5389            .stream()
5390            .await;
5391
5392        let mut exec_commits = 0;
5393        let mut results = 0;
5394        let mut final_response = None;
5395        let mut stream = stream;
5396        while let Some(item) = stream.next().await {
5397            match item.unwrap_or_else(|err| panic!("stream item errored: {err}")) {
5398                MultiTurnStreamItem::ToolExecutionCommitted { .. } => exec_commits += 1,
5399                MultiTurnStreamItem::StreamUserItem(StreamedUserContent::ToolResult { .. }) => {
5400                    results += 1
5401                }
5402                MultiTurnStreamItem::FinalResponse(resp) => final_response = Some(resp),
5403                _ => {}
5404            }
5405        }
5406
5407        assert_eq!(calls.load(SeqCst), 0, "a skipped tool's body never runs");
5408        assert_eq!(
5409            exec_commits, 0,
5410            "a hook-skipped tool produces no execution-commit"
5411        );
5412        assert_eq!(
5413            results, 1,
5414            "the skip result is still surfaced to the consumer"
5415        );
5416        let final_response = final_response.expect("stream should yield a final response");
5417        // The skip result is committed to history (the model sees the reason).
5418        let history = final_response.messages().expect("history");
5419        assert!(
5420            history.iter().any(|m| serde_json::to_string(m)
5421                .map(|s| s.contains("blocked by policy"))
5422                .unwrap_or(false)),
5423            "the skip result is committed to history"
5424        );
5425    }
5426
5427    /// `ToolChoice::Required` + a hook whose `active_tools([])` advertises no tools
5428    /// is a **local** error: the run fails before any provider round-trip.
5429    #[tokio::test]
5430    async fn required_with_empty_active_tools_errors_locally_without_provider_call() {
5431        struct EmptyActiveToolsHook;
5432        impl AgentHook for EmptyActiveToolsHook {
5433            async fn on_completion_call(
5434                &self,
5435                _ctx: &HookContext,
5436                event: CompletionCallEvent<'_>,
5437            ) -> CompletionCallAction {
5438                if let CompletionCallEvent { .. } = event {
5439                    CompletionCallAction::patch(
5440                        RequestPatch::new().active_tools(Vec::<String>::new()),
5441                    )
5442                } else {
5443                    CompletionCallAction::continue_run()
5444                }
5445            }
5446        }
5447
5448        let model = MockCompletionModel::from_turns([MockTurn::text("unreachable")]);
5449        let probe = model.clone();
5450        let err = AgentBuilder::new(model)
5451            .tool(MockAddTool)
5452            .tool_choice(ToolChoice::Required)
5453            .add_hook(EmptyActiveToolsHook)
5454            .build()
5455            .runner("go")
5456            .run()
5457            .await
5458            .expect_err("Required with an empty active_tools filter must fail locally");
5459
5460        assert!(
5461            probe.requests().is_empty(),
5462            "the request must fail locally, with no provider round-trip"
5463        );
5464        let msg = err.to_string();
5465        assert!(
5466            msg.contains("Required"),
5467            "error should mention Required: {msg}"
5468        );
5469        assert!(
5470            msg.contains("active_tools"),
5471            "error should name active_tools: {msg}"
5472        );
5473    }
5474
5475    /// `ToolChoice::Specific` naming a tool that a hook's `active_tools` filtered
5476    /// out is a **local** error naming the filter, before any provider round-trip.
5477    #[tokio::test]
5478    async fn specific_naming_filtered_out_tool_errors_locally_without_provider_call() {
5479        struct FilterToAddHook;
5480        impl AgentHook for FilterToAddHook {
5481            async fn on_completion_call(
5482                &self,
5483                _ctx: &HookContext,
5484                event: CompletionCallEvent<'_>,
5485            ) -> CompletionCallAction {
5486                if let CompletionCallEvent { .. } = event {
5487                    CompletionCallAction::patch(RequestPatch::new().active_tools(["add"]))
5488                } else {
5489                    CompletionCallAction::continue_run()
5490                }
5491            }
5492        }
5493
5494        let model = MockCompletionModel::from_turns([MockTurn::text("unreachable")]);
5495        let probe = model.clone();
5496        let err = AgentBuilder::new(model)
5497            .tool(MockAddTool)
5498            .tool(MockSubtractTool)
5499            .tool_choice(ToolChoice::Specific {
5500                function_names: vec!["subtract".to_string()],
5501            })
5502            .add_hook(FilterToAddHook)
5503            .build()
5504            .runner("go")
5505            .run()
5506            .await
5507            .expect_err("Specific naming a filtered-out tool must fail locally");
5508
5509        assert!(
5510            probe.requests().is_empty(),
5511            "the request must fail locally, with no provider round-trip"
5512        );
5513        let msg = err.to_string();
5514        assert!(
5515            msg.contains("subtract"),
5516            "error should name the missing tool: {msg}"
5517        );
5518        assert!(
5519            msg.contains("active_tools"),
5520            "error should name active_tools: {msg}"
5521        );
5522    }
5523
5524    /// Concurrent tool execution is bounded on *both* sides: real parallelism
5525    /// occurs (lower bound) and the configured `tool_concurrency` cap is never
5526    /// exceeded (upper bound). Four parallel calls run under a cap of two; the
5527    /// barrier is sized to the cap, so it only releases when `cap` calls are in
5528    /// flight together — a serial runtime would deadlock, while an over-eager one
5529    /// (ignoring the cap) would let `max_active` exceed it.
5530    #[tokio::test]
5531    async fn concurrent_tool_execution_stays_within_the_configured_bound() {
5532        #[derive(Clone)]
5533        struct ConcurrencyProbe {
5534            barrier: Arc<Barrier>,
5535            active: Arc<AtomicU32>,
5536            max_active: Arc<AtomicU32>,
5537        }
5538
5539        impl Tool for ConcurrencyProbe {
5540            const NAME: &'static str = "add";
5541            type Error = MockToolError;
5542            type Args = serde_json::Value;
5543            type Output = String;
5544
5545            fn description(&self) -> String {
5546                "concurrency probe".to_string()
5547            }
5548
5549            fn parameters(&self) -> serde_json::Value {
5550                json!({"type": "object", "properties": {}})
5551            }
5552
5553            async fn call(
5554                &self,
5555                _context: &mut ToolContext,
5556                _args: Self::Args,
5557            ) -> Result<Self::Output, Self::Error> {
5558                let now = self.active.fetch_add(1, SeqCst) + 1;
5559                self.max_active.fetch_max(now, SeqCst);
5560                self.barrier.wait().await;
5561                self.active.fetch_sub(1, SeqCst);
5562                Ok("ok".to_string())
5563            }
5564        }
5565
5566        let cap = 2usize;
5567        let probe = ConcurrencyProbe {
5568            barrier: Arc::new(Barrier::new(cap)),
5569            active: Arc::new(AtomicU32::new(0)),
5570            max_active: Arc::new(AtomicU32::new(0)),
5571        };
5572        let max_active = probe.max_active.clone();
5573
5574        // One turn issues four parallel calls to the probe (registered as `add`).
5575        let model = MockCompletionModel::from_turns([
5576            MockTurn::from_contents([
5577                tool_call_content("c1", json!({})),
5578                tool_call_content("c2", json!({})),
5579                tool_call_content("c3", json!({})),
5580                tool_call_content("c4", json!({})),
5581            ])
5582            .expect("four tool calls is a valid turn"),
5583            MockTurn::text("done"),
5584        ]);
5585
5586        let _ = AgentBuilder::new(model)
5587            .tool(probe)
5588            .build()
5589            .runner("probe concurrency")
5590            .max_turns(3)
5591            .tool_concurrency(cap)
5592            .run()
5593            .await
5594            .expect("run should succeed");
5595
5596        let observed = max_active.load(SeqCst);
5597        assert!(
5598            observed > 1,
5599            "tools actually ran concurrently (lower bound): max_active={observed}"
5600        );
5601        assert!(
5602            observed <= cap as u32,
5603            "in-flight never exceeded the configured bound {cap} (upper bound): max_active={observed}"
5604        );
5605    }
5606
5607    /// `tool_concurrency(0)` is clamped to 1 and runs to completion. The timeout
5608    /// guards against a regression that lets `concurrency == 0` reach a
5609    /// `buffer_unordered(0)` (which never makes progress) instead of the
5610    /// sequential `concurrency <= 1` path.
5611    #[tokio::test]
5612    async fn tool_concurrency_zero_is_clamped_and_does_not_hang() {
5613        let model = MockCompletionModel::from_turns([
5614            MockTurn::tool_call("tc1", "add", json!({"x": 1, "y": 2})),
5615            MockTurn::text("done"),
5616        ]);
5617        let run = AgentBuilder::new(model)
5618            .tool(MockAddTool)
5619            .build()
5620            .runner("add")
5621            .max_turns(3)
5622            .tool_concurrency(0)
5623            .run();
5624
5625        let response = tokio::time::timeout(std::time::Duration::from_secs(5), run)
5626            .await
5627            .expect("tool_concurrency(0) must clamp to 1, not hang on buffer_unordered(0)")
5628            .expect("run should succeed");
5629        assert_eq!(response.output, "done");
5630    }
5631
5632    /// A tool that counts how many times it executes.
5633    #[derive(Clone)]
5634    struct CountingAddTool {
5635        calls: Arc<AtomicU32>,
5636    }
5637    impl Tool for CountingAddTool {
5638        const NAME: &'static str = "add";
5639        type Error = MockToolError;
5640        type Args = MockOperationArgs;
5641        type Output = i32;
5642        fn description(&self) -> String {
5643            MockAddTool.description()
5644        }
5645        fn parameters(&self) -> serde_json::Value {
5646            MockAddTool.parameters()
5647        }
5648        async fn call(
5649            &self,
5650            _context: &mut ToolContext,
5651            args: Self::Args,
5652        ) -> Result<Self::Output, Self::Error> {
5653            self.calls.fetch_add(1, SeqCst);
5654            MockAddTool.call(_context, args).await
5655        }
5656    }
5657
5658    #[derive(Clone, Default)]
5659    struct ToolOnlyHook {
5660        text_delta_calls: Arc<AtomicU32>,
5661        other_calls: Arc<AtomicU32>,
5662    }
5663
5664    impl AgentHook for ToolOnlyHook {
5665        async fn on_text_delta(&self, _: &HookContext, _: TextDelta<'_>) -> ObservationAction {
5666            self.text_delta_calls.fetch_add(1, SeqCst);
5667            ObservationAction::continue_run()
5668        }
5669        async fn on_completion_call(
5670            &self,
5671            _: &HookContext,
5672            _: CompletionCallEvent<'_>,
5673        ) -> CompletionCallAction {
5674            self.other_calls.fetch_add(1, SeqCst);
5675            CompletionCallAction::continue_run()
5676        }
5677        fn observes(&self, kind: StepEventKind) -> bool {
5678            kind != StepEventKind::TextDelta
5679        }
5680    }
5681
5682    /// A hook that declares it does not observe text deltas is never dispatched
5683    /// for them (the runner skips building/dispatching that event), but still
5684    /// receives the events it does observe.
5685    #[tokio::test]
5686    async fn observes_gates_text_delta_dispatch() {
5687        let model = MockCompletionModel::from_stream_turns([vec![
5688            MockStreamEvent::text("hel"),
5689            MockStreamEvent::text("lo"),
5690            MockStreamEvent::final_response_with_total_tokens(0),
5691        ]]);
5692        let hook = ToolOnlyHook::default();
5693        let mut stream = AgentBuilder::new(model)
5694            .build()
5695            .runner("hi")
5696            .add_hook(hook.clone())
5697            .stream()
5698            .await;
5699        while stream.next().await.is_some() {}
5700
5701        assert_eq!(
5702            hook.text_delta_calls.load(SeqCst),
5703            0,
5704            "a hook that does not observe TextDelta must not be dispatched for it"
5705        );
5706        assert!(
5707            hook.other_calls.load(SeqCst) > 0,
5708            "the hook should still receive the events it observes"
5709        );
5710    }
5711
5712    /// Terminates the run when it sees a chosen event kind, observing every other
5713    /// event as `Continue`.
5714    struct TerminateOn(StepEventKind);
5715
5716    impl AgentHook for TerminateOn {
5717        async fn on_completion_call(
5718            &self,
5719            _: &HookContext,
5720            _: CompletionCallEvent<'_>,
5721        ) -> CompletionCallAction {
5722            if self.0 == StepEventKind::CompletionCall {
5723                CompletionCallAction::stop("stop here")
5724            } else {
5725                CompletionCallAction::continue_run()
5726            }
5727        }
5728        async fn on_tool_call(&self, _: &HookContext, _: ToolCall<'_>) -> ToolCallAction {
5729            if self.0 == StepEventKind::ToolCall {
5730                ToolCallAction::stop("stop here")
5731            } else {
5732                ToolCallAction::run()
5733            }
5734        }
5735        async fn on_tool_result(
5736            &self,
5737            _: &HookContext,
5738            _: ToolResultEvent<'_>,
5739        ) -> ToolResultAction {
5740            if self.0 == StepEventKind::ToolResult {
5741                ToolResultAction::stop("stop here")
5742            } else {
5743                ToolResultAction::keep()
5744            }
5745        }
5746    }
5747
5748    /// the event-specific stop action cancels the blocking run from *every* shared driver
5749    /// event (model call, model response, tool call, tool result) — none is a
5750    /// silent no-op.
5751    #[tokio::test]
5752    async fn run_terminates_from_each_shared_event() {
5753        for kind in [
5754            StepEventKind::CompletionCall,
5755            StepEventKind::ToolCall,
5756            StepEventKind::ToolResult,
5757        ] {
5758            let err = AgentBuilder::new(blocking_model())
5759                .tool(MockAddTool)
5760                .build()
5761                .runner("add 2 and 3")
5762                .max_turns(3)
5763                .add_hook(TerminateOn(kind))
5764                .run()
5765                .await
5766                .expect_err(&format!("terminate at {kind:?} must cancel the run"));
5767            assert!(
5768                matches!(err, PromptError::PromptCancelled { .. }),
5769                "terminate at {kind:?} should cancel the run, got {err:?}"
5770            );
5771        }
5772    }
5773
5774    /// The same fail-closed termination holds for the streaming driver across the
5775    /// shared events it fires (it surfaces `StreamResponseFinish` instead of
5776    /// `CompletionResponse`): each yields a stream error and no final response.
5777    #[tokio::test]
5778    async fn stream_terminates_from_each_shared_event() {
5779        for kind in [
5780            StepEventKind::CompletionCall,
5781            StepEventKind::ToolCall,
5782            StepEventKind::ToolResult,
5783        ] {
5784            let mut stream = AgentBuilder::new(streaming_model())
5785                .tool(MockAddTool)
5786                .build()
5787                .runner("add 2 and 3")
5788                .max_turns(3)
5789                .add_hook(TerminateOn(kind))
5790                .stream()
5791                .await;
5792
5793            let mut saw_error = false;
5794            let mut saw_final = false;
5795            while let Some(item) = stream.next().await {
5796                match item {
5797                    Ok(MultiTurnStreamItem::FinalResponse(_)) => saw_final = true,
5798                    Err(_) => saw_error = true,
5799                    _ => {}
5800                }
5801            }
5802            assert!(saw_error, "terminate at {kind:?} must yield a stream error");
5803            assert!(
5804                !saw_final,
5805                "terminate at {kind:?} must not also produce a final response"
5806            );
5807        }
5808    }
5809
5810    /// Two hooks pushed onto one stack both observe every event (no short-circuit
5811    /// on `Continue`), and the stack's shared event sequence is identical across
5812    /// the blocking and streaming drivers.
5813    #[tokio::test]
5814    async fn multi_hook_stack_parity_across_run_and_stream() {
5815        let a_block = RecordingHook::default();
5816        let b_block = RecordingHook::default();
5817        let blocking = AgentBuilder::new(blocking_model())
5818            .tool(MockAddTool)
5819            .build()
5820            .runner("add 2 and 3")
5821            .max_turns(3)
5822            .add_hook(a_block.clone())
5823            .add_hook(b_block.clone())
5824            .run()
5825            .await
5826            .expect("blocking run should succeed");
5827
5828        let a_stream = RecordingHook::default();
5829        let b_stream = RecordingHook::default();
5830        let mut stream = AgentBuilder::new(streaming_model())
5831            .tool(MockAddTool)
5832            .build()
5833            .runner("add 2 and 3")
5834            .max_turns(3)
5835            .add_hook(a_stream.clone())
5836            .add_hook(b_stream.clone())
5837            .stream()
5838            .await;
5839        while stream.next().await.is_some() {}
5840
5841        // Both hooks in the stack saw the same events (both ran on every Continue).
5842        assert_eq!(a_block.shared_events(), b_block.shared_events());
5843        assert_eq!(a_stream.shared_events(), b_stream.shared_events());
5844        // The stack's shared event sequence is identical across drivers.
5845        assert_eq!(a_block.shared_events(), a_stream.shared_events());
5846        assert_eq!(
5847            a_block.shared_events(),
5848            vec![
5849                StepEventKind::CompletionCall,
5850                StepEventKind::ToolCall,
5851                StepEventKind::ToolResult,
5852                StepEventKind::CompletionCall,
5853            ]
5854        );
5855        assert_eq!(blocking.output, "the answer is 5");
5856    }
5857
5858    /// Renames an invalid tool call to a known tool; observes everything else.
5859    struct RepairInvalidToHook(&'static str);
5860
5861    impl AgentHook for RepairInvalidToHook {
5862        async fn on_invalid_tool_call(
5863            &self,
5864            _ctx: &HookContext,
5865            event: &InvalidToolCallContext,
5866        ) -> Option<InvalidToolCallAction> {
5867            Some(if let _ = event {
5868                InvalidToolCallAction::repair(self.0)
5869            } else {
5870                InvalidToolCallAction::fail()
5871            })
5872        }
5873    }
5874
5875    #[derive(Clone)]
5876    struct CaptureAndRepairInvalidHook {
5877        replacement: &'static str,
5878        args: Arc<Mutex<Vec<Option<String>>>>,
5879    }
5880
5881    impl AgentHook for CaptureAndRepairInvalidHook {
5882        async fn on_invalid_tool_call(
5883            &self,
5884            _ctx: &HookContext,
5885            event: &InvalidToolCallContext,
5886        ) -> Option<InvalidToolCallAction> {
5887            self.args
5888                .lock()
5889                .expect("invalid args")
5890                .push(event.args.clone());
5891            Some(InvalidToolCallAction::repair(self.replacement))
5892        }
5893    }
5894
5895    /// An invalid tool call repaired by a hook recovers identically under run()
5896    /// and stream(): the renamed tool executes and both drivers reach the same
5897    /// output, tool-result content, and final message history.
5898    #[tokio::test]
5899    async fn invalid_tool_call_repair_parity_across_run_and_stream() {
5900        let blocking_model = MockCompletionModel::from_turns([
5901            MockTurn::tool_call("tc1", "default_api", json!({"x": 2, "y": 3})),
5902            MockTurn::text("the answer is 5"),
5903        ]);
5904        let blocking_hook = RecordingHook::default();
5905        let blocking = AgentBuilder::new(blocking_model)
5906            .tool(MockAddTool)
5907            .build()
5908            .runner("add 2 and 3")
5909            .max_turns(3)
5910            .add_hook(blocking_hook.clone())
5911            .add_hook(RepairInvalidToHook("add"))
5912            .run()
5913            .await
5914            .expect("blocking run should recover via repair");
5915
5916        // Emit the invalid call as a single complete tool call (mirroring the
5917        // blocking model). A provider stream carries one tool call via one
5918        // mechanism — deltas *or* a complete call — so this is the apples-to-
5919        // apples comparison; mixing both would trip the assembler's two
5920        // independent invalid-detection sites and fire the event twice.
5921        let streaming_model = MockCompletionModel::from_stream_turns([
5922            vec![
5923                MockStreamEvent::tool_call("tc1", "default_api", json!({"x": 2, "y": 3})),
5924                MockStreamEvent::final_response_with_total_tokens(0),
5925            ],
5926            vec![
5927                MockStreamEvent::text("the answer is 5"),
5928                MockStreamEvent::final_response_with_total_tokens(0),
5929            ],
5930        ]);
5931        let streaming_hook = RecordingHook::default();
5932        let mut stream = AgentBuilder::new(streaming_model)
5933            .tool(MockAddTool)
5934            .build()
5935            .runner("add 2 and 3")
5936            .max_turns(3)
5937            .add_hook(streaming_hook.clone())
5938            .add_hook(RepairInvalidToHook("add"))
5939            .stream()
5940            .await;
5941        let mut final_response = None;
5942        while let Some(item) = stream.next().await {
5943            if let Ok(MultiTurnStreamItem::FinalResponse(resp)) =
5944                item.map_err(|err| panic!("stream item errored: {err}"))
5945            {
5946                final_response = Some(resp);
5947            }
5948        }
5949        let final_response =
5950            final_response.expect("stream should recover and yield a final response");
5951
5952        // Same recovered output.
5953        assert_eq!(blocking.output, "the answer is 5");
5954        assert_eq!(final_response.output(), blocking.output);
5955
5956        // Both drivers reported the invalid tool call to the hook, then executed
5957        // the repaired tool, so the shared event sequences match.
5958        assert_eq!(
5959            blocking_hook.shared_events(),
5960            streaming_hook.shared_events()
5961        );
5962        assert!(
5963            blocking_hook
5964                .shared_events()
5965                .contains(&StepEventKind::InvalidToolCall),
5966            "the hook must observe the invalid tool call"
5967        );
5968        assert_eq!(blocking_hook.tool_results(), streaming_hook.tool_results());
5969        assert_eq!(blocking_hook.tool_results(), vec!["5".to_string()]);
5970
5971        // Same final message history.
5972        let blocking_messages = blocking.messages.expect("blocking messages");
5973        let streaming_messages = final_response
5974            .messages()
5975            .expect("streaming history")
5976            .to_vec();
5977        assert_eq!(
5978            serde_json::to_value(&blocking_messages).expect("serialize blocking"),
5979            serde_json::to_value(&streaming_messages).expect("serialize streaming"),
5980        );
5981    }
5982
5983    #[tokio::test]
5984    async fn invalid_tool_call_scalar_args_are_canonical_across_run_and_complete_stream() {
5985        let blocking_args = Arc::new(Mutex::new(Vec::new()));
5986        let blocking_hook = RecordingHook::default();
5987        let blocking = AgentBuilder::new(MockCompletionModel::from_turns([
5988            MockTurn::tool_call("tc1", "unknown_echo", json!("payload")),
5989            MockTurn::text("done"),
5990        ]))
5991        .tool(EchoStringArgs)
5992        .build()
5993        .runner("echo a string")
5994        .max_turns(3)
5995        .add_hook(blocking_hook.clone())
5996        .add_hook(CaptureAndRepairInvalidHook {
5997            replacement: EchoStringArgs::NAME,
5998            args: blocking_args.clone(),
5999        })
6000        .run()
6001        .await
6002        .expect("blocking scalar repair should succeed");
6003
6004        let streaming_args = Arc::new(Mutex::new(Vec::new()));
6005        let streaming_hook = RecordingHook::default();
6006        let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([
6007            vec![
6008                MockStreamEvent::tool_call("tc1", "unknown_echo", json!("payload")),
6009                MockStreamEvent::final_response_with_total_tokens(0),
6010            ],
6011            vec![
6012                MockStreamEvent::text("done"),
6013                MockStreamEvent::final_response_with_total_tokens(0),
6014            ],
6015        ]))
6016        .tool(EchoStringArgs)
6017        .build()
6018        .runner("echo a string")
6019        .max_turns(3)
6020        .add_hook(streaming_hook.clone())
6021        .add_hook(CaptureAndRepairInvalidHook {
6022            replacement: EchoStringArgs::NAME,
6023            args: streaming_args.clone(),
6024        })
6025        .stream()
6026        .await;
6027        let mut final_response = None;
6028        while let Some(item) = stream.next().await {
6029            if let MultiTurnStreamItem::FinalResponse(response) =
6030                item.expect("streaming scalar repair should succeed")
6031            {
6032                final_response = Some(response);
6033            }
6034        }
6035        let final_response = final_response.expect("stream should yield a final response");
6036
6037        let canonical_args = vec![Some(serde_json::to_string("payload").unwrap())];
6038        assert_eq!(*blocking_args.lock().unwrap(), canonical_args);
6039        assert_eq!(*streaming_args.lock().unwrap(), canonical_args);
6040        assert_eq!(blocking_hook.tool_results(), vec!["payload"]);
6041        assert_eq!(streaming_hook.tool_results(), vec!["payload"]);
6042        assert_eq!(blocking.output, "done");
6043        assert_eq!(final_response.output(), "done");
6044        assert_eq!(
6045            serde_json::to_value(blocking.messages.expect("blocking history")).unwrap(),
6046            serde_json::to_value(final_response.messages().expect("streaming history")).unwrap()
6047        );
6048    }
6049
6050    // ----------------------------------------------------------------------
6051    // Single-source-of-truth parity harness
6052    // ----------------------------------------------------------------------
6053    //
6054    // `run()` and `stream()` are two implementations of one agent loop; testing
6055    // they agree on the same input is *differential testing*, with each driver
6056    // acting as the other's oracle. The hazard such tests have (and that bit the
6057    // invalid-tool-repair test above) is *fixture drift*: when the blocking
6058    // `MockTurn` list and the streaming `MockStreamEvent` list are hand-written
6059    // separately, they can silently encode different model behavior, so a
6060    // passing test proves nothing.
6061    //
6062    // The fix — the single-source-of-truth / data-driven principle, embodied by
6063    // pydantic-ai's `TestModel` (one scripted response replayed as a stream) and
6064    // litellm's `stream_chunk_builder` (reassemble the stream, compare to the
6065    // whole) — is to derive *both* encodings from one canonical `ScriptedTurn`
6066    // list. The two drivers are then provably fed identical model behavior and
6067    // can be asserted equal on the medium-independent projection (final output,
6068    // message history, tool-result content, shared hook-event sequence).
6069
6070    /// One tool call inside a scripted turn.
6071    #[derive(Clone)]
6072    struct ScriptedToolCall {
6073        id: &'static str,
6074        name: &'static str,
6075        args: serde_json::Value,
6076    }
6077
6078    /// One scripted model turn, described once and rendered into both a blocking
6079    /// `MockTurn` and a streaming `Vec<MockStreamEvent>`.
6080    #[derive(Clone)]
6081    enum ScriptedTurn {
6082        /// A final text answer.
6083        Text(&'static str),
6084        /// One or more tool calls emitted in a single turn.
6085        ToolCalls(Vec<ScriptedToolCall>),
6086    }
6087
6088    /// How a tool call is rendered onto the wire for the streaming driver. Both
6089    /// shapes must yield the *same* canonical turn ("chunked-input invariance",
6090    /// the `tokio-util` `LengthDelimitedCodec` lesson): the assembled message
6091    /// history and tool results may not depend on whether a provider sends a
6092    /// complete tool call or streams it as deltas.
6093    #[derive(Clone, Copy)]
6094    enum StreamShape {
6095        /// One complete tool-call event per call (mirrors the blocking turn).
6096        Complete,
6097        /// Name + argument deltas followed by the complete call, additionally
6098        /// exercising the delta-hook path and the assembler's delta buffering.
6099        Chunked,
6100    }
6101
6102    impl ScriptedTurn {
6103        fn as_blocking_turn(&self) -> MockTurn {
6104            match self {
6105                ScriptedTurn::Text(text) => MockTurn::text(*text),
6106                ScriptedTurn::ToolCalls(calls) => {
6107                    MockTurn::from_contents(calls.iter().map(|call| {
6108                        AssistantContent::ToolCall(MessageToolCall::new(
6109                            call.id.to_string(),
6110                            ToolFunction::new(call.name.to_string(), call.args.clone()),
6111                        ))
6112                    }))
6113                    .expect("a scripted tool-call turn has at least one call")
6114                }
6115            }
6116        }
6117
6118        fn as_stream_events(&self, shape: StreamShape) -> Vec<MockStreamEvent> {
6119            let mut events = Vec::new();
6120            match self {
6121                ScriptedTurn::Text(text) => events.push(MockStreamEvent::text(*text)),
6122                ScriptedTurn::ToolCalls(calls) => {
6123                    for call in calls {
6124                        if let StreamShape::Chunked = shape {
6125                            // Distinct internal id per call; the canonical args
6126                            // still come from the complete event below, so this
6127                            // exercises the delta path without changing the turn.
6128                            let internal = format!("ic-{}", call.id);
6129                            let args = serde_json::to_string(&call.args)
6130                                .expect("scripted args serialize to json");
6131                            events.push(MockStreamEvent::tool_call_name_delta(
6132                                call.id, &internal, call.name,
6133                            ));
6134                            events.push(MockStreamEvent::tool_call_arguments_delta(
6135                                call.id, &internal, &args,
6136                            ));
6137                        }
6138                        events.push(MockStreamEvent::tool_call(
6139                            call.id,
6140                            call.name,
6141                            call.args.clone(),
6142                        ));
6143                    }
6144                }
6145            }
6146            events.push(MockStreamEvent::final_response_with_total_tokens(0));
6147            events
6148        }
6149    }
6150
6151    /// The medium-independent projection of a run that both drivers must agree
6152    /// on.
6153    struct ParityOutcome {
6154        output: String,
6155        messages: Vec<Message>,
6156        shared_events: Vec<StepEventKind>,
6157        tool_results: Vec<String>,
6158    }
6159
6160    async fn run_blocking_scenario(prompt: &'static str, turns: &[ScriptedTurn]) -> ParityOutcome {
6161        let model =
6162            MockCompletionModel::from_turns(turns.iter().map(ScriptedTurn::as_blocking_turn));
6163        let hook = RecordingHook::default();
6164        let response = AgentBuilder::new(model)
6165            .tool(MockAddTool)
6166            .build()
6167            .runner(prompt)
6168            .max_turns(8)
6169            .add_hook(hook.clone())
6170            .run()
6171            .await
6172            .expect("blocking scenario should succeed");
6173        ParityOutcome {
6174            output: response.output,
6175            messages: response.messages.expect("blocking messages"),
6176            shared_events: hook.shared_events(),
6177            tool_results: hook.tool_results(),
6178        }
6179    }
6180
6181    async fn run_streaming_scenario(
6182        prompt: &'static str,
6183        turns: &[ScriptedTurn],
6184        shape: StreamShape,
6185    ) -> ParityOutcome {
6186        let model = MockCompletionModel::from_stream_turns(
6187            turns.iter().map(|turn| turn.as_stream_events(shape)),
6188        );
6189        let hook = RecordingHook::default();
6190        let mut stream = AgentBuilder::new(model)
6191            .tool(MockAddTool)
6192            .build()
6193            .runner(prompt)
6194            .max_turns(8)
6195            .add_hook(hook.clone())
6196            .stream()
6197            .await;
6198        let mut final_response = None;
6199        while let Some(item) = stream.next().await {
6200            if let Ok(MultiTurnStreamItem::FinalResponse(resp)) =
6201                item.map_err(|err| panic!("stream item errored: {err}"))
6202            {
6203                final_response = Some(resp);
6204            }
6205        }
6206        let final_response =
6207            final_response.expect("streaming scenario should yield a final response");
6208        ParityOutcome {
6209            output: final_response.output().to_string(),
6210            messages: final_response
6211                .messages()
6212                .expect("streaming history")
6213                .to_vec(),
6214            shared_events: hook.shared_events(),
6215            tool_results: hook.tool_results(),
6216        }
6217    }
6218
6219    fn assert_outcomes_match(blocking: &ParityOutcome, streaming: &ParityOutcome, label: &str) {
6220        assert_eq!(
6221            blocking.output, streaming.output,
6222            "{label}: final output diverged"
6223        );
6224        assert_eq!(
6225            blocking.shared_events, streaming.shared_events,
6226            "{label}: hook event sequence diverged"
6227        );
6228        assert_eq!(
6229            blocking.tool_results, streaming.tool_results,
6230            "{label}: tool-result content diverged"
6231        );
6232        assert_eq!(
6233            serde_json::to_value(&blocking.messages).expect("serialize blocking"),
6234            serde_json::to_value(&streaming.messages).expect("serialize streaming"),
6235            "{label}: message history diverged"
6236        );
6237    }
6238
6239    /// Drive one canonical scenario through `run()` and through `stream()` in
6240    /// both wire shapes, asserting the medium-independent projection is
6241    /// identical every way. Because both stream shapes are compared against the
6242    /// same blocking outcome, they are also transitively equal to each other.
6243    async fn assert_run_stream_parity(prompt: &'static str, turns: &[ScriptedTurn]) {
6244        let blocking = run_blocking_scenario(prompt, turns).await;
6245        for (shape, label) in [
6246            (StreamShape::Complete, "complete-stream"),
6247            (StreamShape::Chunked, "chunked-stream"),
6248        ] {
6249            let streaming = run_streaming_scenario(prompt, turns, shape).await;
6250            assert_outcomes_match(&blocking, &streaming, label);
6251        }
6252    }
6253
6254    fn add_call(id: &'static str, x: i64, y: i64) -> ScriptedToolCall {
6255        ScriptedToolCall {
6256            id,
6257            name: "add",
6258            args: json!({ "x": x, "y": y }),
6259        }
6260    }
6261
6262    #[tokio::test]
6263    async fn parity_text_only_run() {
6264        assert_run_stream_parity("just say hi", &[ScriptedTurn::Text("hi there")]).await;
6265    }
6266
6267    #[tokio::test]
6268    async fn parity_single_tool_then_text() {
6269        assert_run_stream_parity(
6270            "add 2 and 3",
6271            &[
6272                ScriptedTurn::ToolCalls(vec![add_call("tc1", 2, 3)]),
6273                ScriptedTurn::Text("the answer is 5"),
6274            ],
6275        )
6276        .await;
6277    }
6278
6279    #[tokio::test]
6280    async fn parity_multiple_tools_in_one_turn() {
6281        assert_run_stream_parity(
6282            "add two pairs",
6283            &[
6284                ScriptedTurn::ToolCalls(vec![add_call("tc1", 2, 3), add_call("tc2", 10, 20)]),
6285                ScriptedTurn::Text("done"),
6286            ],
6287        )
6288        .await;
6289    }
6290
6291    #[tokio::test]
6292    async fn parity_multi_turn_sequential_tools() {
6293        assert_run_stream_parity(
6294            "chain two additions",
6295            &[
6296                ScriptedTurn::ToolCalls(vec![add_call("tc1", 1, 1)]),
6297                ScriptedTurn::ToolCalls(vec![add_call("tc2", 2, 2)]),
6298                ScriptedTurn::Text("chained"),
6299            ],
6300        )
6301        .await;
6302    }
6303
6304    /// Skips an invalid tool call (synthetic result, no execution); observes
6305    /// everything else.
6306    struct SkipInvalidHook(&'static str);
6307
6308    impl AgentHook for SkipInvalidHook {
6309        async fn on_invalid_tool_call(
6310            &self,
6311            _ctx: &HookContext,
6312            event: &InvalidToolCallContext,
6313        ) -> Option<InvalidToolCallAction> {
6314            Some(if let _ = event {
6315                InvalidToolCallAction::skip(self.0)
6316            } else {
6317                InvalidToolCallAction::fail()
6318            })
6319        }
6320    }
6321
6322    /// An invalid tool call *skipped* by a hook recovers identically under
6323    /// `run()` and `stream()`: the synthetic skip result enters the history
6324    /// verbatim (it is never re-parsed as tool output) and both drivers reach
6325    /// the same output and message history. Complements the repair-parity test.
6326    #[tokio::test]
6327    async fn invalid_tool_call_skip_parity_across_run_and_stream() {
6328        let blocking_model = MockCompletionModel::from_turns([
6329            MockTurn::tool_call("tc1", "default_api", json!({"x": 2, "y": 3})),
6330            MockTurn::text("acknowledged"),
6331        ]);
6332        let blocking_hook = RecordingHook::default();
6333        let blocking = AgentBuilder::new(blocking_model)
6334            .tool(MockAddTool)
6335            .build()
6336            .runner("do the thing")
6337            .max_turns(3)
6338            .add_hook(blocking_hook.clone())
6339            .add_hook(SkipInvalidHook("tool not permitted"))
6340            .run()
6341            .await
6342            .expect("blocking run should recover via skip");
6343
6344        // Single complete tool call (mirrors the blocking model; see the
6345        // repair-parity test for why deltas are not mixed in here).
6346        let streaming_model = MockCompletionModel::from_stream_turns([
6347            vec![
6348                MockStreamEvent::tool_call("tc1", "default_api", json!({"x": 2, "y": 3})),
6349                MockStreamEvent::final_response_with_total_tokens(0),
6350            ],
6351            vec![
6352                MockStreamEvent::text("acknowledged"),
6353                MockStreamEvent::final_response_with_total_tokens(0),
6354            ],
6355        ]);
6356        let streaming_hook = RecordingHook::default();
6357        let mut stream = AgentBuilder::new(streaming_model)
6358            .tool(MockAddTool)
6359            .build()
6360            .runner("do the thing")
6361            .max_turns(3)
6362            .add_hook(streaming_hook.clone())
6363            .add_hook(SkipInvalidHook("tool not permitted"))
6364            .stream()
6365            .await;
6366        let mut final_response = None;
6367        while let Some(item) = stream.next().await {
6368            if let Ok(MultiTurnStreamItem::FinalResponse(resp)) =
6369                item.map_err(|err| panic!("stream item errored: {err}"))
6370            {
6371                final_response = Some(resp);
6372            }
6373        }
6374        let final_response =
6375            final_response.expect("stream should recover and yield a final response");
6376
6377        assert_eq!(blocking.output, "acknowledged");
6378        assert_eq!(final_response.output(), blocking.output);
6379        assert_eq!(
6380            blocking_hook.shared_events(),
6381            streaming_hook.shared_events()
6382        );
6383        assert!(
6384            blocking_hook
6385                .shared_events()
6386                .contains(&StepEventKind::InvalidToolCall),
6387            "the hook must observe the invalid tool call"
6388        );
6389
6390        let blocking_messages = blocking.messages.expect("blocking messages");
6391        let streaming_messages = final_response
6392            .messages()
6393            .expect("streaming history")
6394            .to_vec();
6395        assert_eq!(
6396            serde_json::to_value(&blocking_messages).expect("serialize blocking"),
6397            serde_json::to_value(&streaming_messages).expect("serialize streaming"),
6398        );
6399        // Pin the actual reason, not just blocking == streaming (see the valid-tool
6400        // skip test): a reason dropped or altered on BOTH paths would still pass.
6401        assert!(
6402            tool_result_text_in_history(&blocking_messages, "tool not permitted"),
6403            "the verbatim invalid-tool skip reason must be the tool result content"
6404        );
6405    }
6406
6407    /// A turn that streams *text and* an invalid tool call, then is repaired, is
6408    /// a recovered turn: its response-finish hook must be suppressed on BOTH
6409    /// drivers — `CompletionResponse` under `run()`, `StreamResponseFinish` under
6410    /// `stream()`. The shared-events parity harness deliberately excludes these
6411    /// medium-specific events, so this asymmetry needs a dedicated assertion (it
6412    /// is the exact event the harness cannot see).
6413    #[tokio::test]
6414    async fn recovered_turn_suppresses_response_finish_hook_on_both_drivers() {
6415        // Turn 1 emits text then an invalid tool call (repaired to "add"); turn 2
6416        // is a plain final-text turn whose response event DOES fire on both
6417        // drivers — so a correct run sees exactly one response-finish event.
6418        let blocking_model = MockCompletionModel::from_turns([
6419            MockTurn::from_contents([
6420                AssistantContent::text("let me compute that"),
6421                AssistantContent::ToolCall(MessageToolCall::new(
6422                    "tc1".to_string(),
6423                    ToolFunction::new("default_api".to_string(), json!({"x": 2, "y": 3})),
6424                )),
6425            ])
6426            .expect("a text + tool-call turn is valid"),
6427            MockTurn::text("the answer is 5"),
6428        ]);
6429        let blocking_hook = RecordingHook::default();
6430        let blocking = AgentBuilder::new(blocking_model)
6431            .tool(MockAddTool)
6432            .build()
6433            .runner("compute")
6434            .max_turns(3)
6435            .add_hook(blocking_hook.clone())
6436            .add_hook(RepairInvalidToHook("add"))
6437            .run()
6438            .await
6439            .expect("blocking run should recover via repair");
6440
6441        let streaming_model = MockCompletionModel::from_stream_turns([
6442            vec![
6443                MockStreamEvent::text("let me compute that"),
6444                MockStreamEvent::tool_call("tc1", "default_api", json!({"x": 2, "y": 3})),
6445                MockStreamEvent::final_response_with_total_tokens(0),
6446            ],
6447            vec![
6448                MockStreamEvent::text("the answer is 5"),
6449                MockStreamEvent::final_response_with_total_tokens(0),
6450            ],
6451        ]);
6452        let streaming_hook = RecordingHook::default();
6453        let mut stream = AgentBuilder::new(streaming_model)
6454            .tool(MockAddTool)
6455            .build()
6456            .runner("compute")
6457            .max_turns(3)
6458            .add_hook(streaming_hook.clone())
6459            .add_hook(RepairInvalidToHook("add"))
6460            .stream()
6461            .await;
6462        while stream.next().await.is_some() {}
6463
6464        // Recovery still reaches the same final answer.
6465        assert_eq!(blocking.output, "the answer is 5");
6466
6467        // Blocking: the recovered turn 1 suppresses `CompletionResponse`; only the
6468        // plain turn 2 fires it.
6469        assert_eq!(
6470            blocking_hook.count(StepEventKind::CompletionResponse),
6471            1,
6472            "the recovered turn must not fire CompletionResponse"
6473        );
6474        // Streaming: the recovered turn 1 must likewise suppress
6475        // `StreamResponseFinish` (without the fix this is 2).
6476        assert_eq!(
6477            streaming_hook.count(StepEventKind::StreamResponseFinish),
6478            1,
6479            "the recovered turn must not fire StreamResponseFinish"
6480        );
6481        // Stated as parity: the count of un-suppressed response-finish events is
6482        // the same across drivers.
6483        assert_eq!(
6484            blocking_hook.count(StepEventKind::CompletionResponse),
6485            streaming_hook.count(StepEventKind::StreamResponseFinish),
6486        );
6487
6488        // The normalized per-turn `ModelTurnFinished` is suppressed on the
6489        // recovered turn 1 on BOTH surfaces too (its own guard, separate from the
6490        // medium-specific response-finish guards above), so only the accepted turn
6491        // 2 fires it — count is 1, not 2, on each driver. Without the suppression
6492        // this would be 2, and a per-turn accounting hook would double-count the
6493        // recovered turn.
6494        assert_eq!(
6495            blocking_hook.count(StepEventKind::ModelTurnFinished),
6496            1,
6497            "the recovered turn must not fire ModelTurnFinished"
6498        );
6499        assert_eq!(
6500            streaming_hook.count(StepEventKind::ModelTurnFinished),
6501            1,
6502            "the recovered turn must not fire ModelTurnFinished on the streaming surface either"
6503        );
6504        // Parity: the normalized per-turn event fires the same number of times on
6505        // both drivers even when a turn is recovered.
6506        assert_eq!(
6507            blocking_hook.count(StepEventKind::ModelTurnFinished),
6508            streaming_hook.count(StepEventKind::ModelTurnFinished),
6509        );
6510    }
6511
6512    /// A prompt/runner-level `add_hook` APPENDS to the agent's default hooks
6513    /// rather than replacing them (the `with_hook` → `add_hook` semantic change):
6514    /// a hook registered on the builder and a hook registered on the runner both
6515    /// observe the same run.
6516    #[tokio::test]
6517    async fn runner_add_hook_appends_to_agent_default_hooks() {
6518        let agent_hook = RecordingHook::default();
6519        let runner_hook = RecordingHook::default();
6520
6521        // `agent_hook` is registered on the builder; `runner_hook` is registered
6522        // on the runner obtained from that agent. `AgentRunner::from_agent` clones
6523        // the agent's hook stack and `add_hook` pushes on top, so both must fire.
6524        AgentBuilder::new(blocking_model())
6525            .tool(MockAddTool)
6526            .add_hook(agent_hook.clone())
6527            .build()
6528            .runner("add 2 and 3")
6529            .max_turns(3)
6530            .add_hook(runner_hook.clone())
6531            .run()
6532            .await
6533            .expect("run should succeed");
6534
6535        assert!(
6536            agent_hook.count(StepEventKind::CompletionCall) >= 1,
6537            "the agent-default hook must still observe the run after a runner-level add_hook"
6538        );
6539        assert!(
6540            runner_hook.count(StepEventKind::CompletionCall) >= 1,
6541            "the runner-level hook must also observe the run"
6542        );
6543        // Both saw the same number of completion calls — the runner-level hook
6544        // appended to the agent stack; it did not replace it.
6545        assert_eq!(
6546            agent_hook.count(StepEventKind::CompletionCall),
6547            runner_hook.count(StepEventKind::CompletionCall),
6548            "add_hook appends (both hooks observe every turn); it does not replace"
6549        );
6550    }
6551
6552    /// Skips a *valid* tool call before execution; observes everything else.
6553    struct SkipToolCallHook(&'static str);
6554
6555    impl AgentHook for SkipToolCallHook {
6556        async fn on_tool_call(&self, _ctx: &HookContext, event: ToolCall<'_>) -> ToolCallAction {
6557            if let ToolCall { .. } = event {
6558                ToolCallAction::skip(self.0)
6559            } else {
6560                ToolCallAction::run()
6561            }
6562        }
6563    }
6564
6565    /// A hook that skips a *valid* tool call (`ToolCallAction::Skip` on `ToolCall`, the
6566    /// honored-action path — distinct from skipping an *invalid* call) recovers
6567    /// identically under `run()` and `stream()`: the synthetic skip result enters
6568    /// the history verbatim without executing the tool, and both drivers reach the
6569    /// same output, tool-result content and message history.
6570    #[tokio::test]
6571    async fn valid_tool_call_skip_parity_across_run_and_stream() {
6572        let turns = [
6573            ScriptedTurn::ToolCalls(vec![add_call("tc1", 2, 3)]),
6574            ScriptedTurn::Text("acknowledged"),
6575        ];
6576
6577        let blocking_model =
6578            MockCompletionModel::from_turns(turns.iter().map(ScriptedTurn::as_blocking_turn));
6579        let blocking_hook = RecordingHook::default();
6580        let blocking = AgentBuilder::new(blocking_model)
6581            .tool(MockAddTool)
6582            .build()
6583            .runner("add 2 and 3")
6584            .max_turns(3)
6585            .add_hook(blocking_hook.clone())
6586            .add_hook(SkipToolCallHook("skipped by policy"))
6587            .run()
6588            .await
6589            .expect("blocking run should succeed with a skipped tool call");
6590
6591        let streaming_model = MockCompletionModel::from_stream_turns(
6592            turns
6593                .iter()
6594                .map(|turn| turn.as_stream_events(StreamShape::Complete)),
6595        );
6596        let streaming_hook = RecordingHook::default();
6597        let mut stream = AgentBuilder::new(streaming_model)
6598            .tool(MockAddTool)
6599            .build()
6600            .runner("add 2 and 3")
6601            .max_turns(3)
6602            .add_hook(streaming_hook.clone())
6603            .add_hook(SkipToolCallHook("skipped by policy"))
6604            .stream()
6605            .await;
6606        let mut final_response = None;
6607        while let Some(item) = stream.next().await {
6608            if let Ok(MultiTurnStreamItem::FinalResponse(resp)) =
6609                item.map_err(|err| panic!("stream item errored: {err}"))
6610            {
6611                final_response = Some(resp);
6612            }
6613        }
6614        let final_response = final_response.expect("stream should yield a final response");
6615
6616        assert_eq!(blocking.output, "acknowledged");
6617        assert_eq!(final_response.output(), blocking.output);
6618        assert_eq!(
6619            blocking_hook.shared_events(),
6620            streaming_hook.shared_events()
6621        );
6622        // A skipped valid tool call fires the `ToolResult` hook carrying a
6623        // structured `Skipped` outcome (the redesign surfaces the skip to result
6624        // hooks), so both drivers record the verbatim skip reason as the result.
6625        assert_eq!(blocking_hook.tool_results(), streaming_hook.tool_results());
6626        assert_eq!(
6627            blocking_hook.tool_results(),
6628            vec!["skipped by policy".to_string()],
6629            "a skipped tool fires a ToolResult hook with the verbatim skip reason"
6630        );
6631
6632        let blocking_messages = blocking.messages.expect("blocking messages");
6633        let streaming_messages = final_response
6634            .messages()
6635            .expect("streaming history")
6636            .to_vec();
6637        assert_eq!(
6638            serde_json::to_value(&blocking_messages).expect("serialize blocking"),
6639            serde_json::to_value(&streaming_messages).expect("serialize streaming"),
6640        );
6641        // Pin the actual reason, not just blocking == streaming: a reason dropped
6642        // or altered on BOTH paths would still satisfy the equality above.
6643        assert!(
6644            tool_result_text_in_history(&blocking_messages, "skipped by policy"),
6645            "the verbatim skip reason must be the tool result content in the history"
6646        );
6647    }
6648
6649    /// A hook that rewrites a valid tool call's arguments (`ToolCallAction::Rewrite` on
6650    /// `ToolCall`) so the tool executes with the replacement instead of what the
6651    /// model emitted.
6652    struct RewriteToolArgsHook(serde_json::Value);
6653
6654    impl AgentHook for RewriteToolArgsHook {
6655        async fn on_tool_call(&self, _ctx: &HookContext, event: ToolCall<'_>) -> ToolCallAction {
6656            if let ToolCall { .. } = event {
6657                ToolCallAction::rewrite(self.0.clone())
6658            } else {
6659                ToolCallAction::run()
6660            }
6661        }
6662    }
6663
6664    struct EchoStringArgs;
6665
6666    impl Tool for EchoStringArgs {
6667        const NAME: &'static str = "echo_string_args";
6668        type Error = rig::tool::ToolExecutionError;
6669        type Args = String;
6670        type Output = String;
6671
6672        fn description(&self) -> String {
6673            "Echo a JSON string argument".to_string()
6674        }
6675
6676        fn parameters(&self) -> serde_json::Value {
6677            json!({"type": "string"})
6678        }
6679
6680        async fn call(
6681            &self,
6682            _context: &mut ToolContext,
6683            args: Self::Args,
6684        ) -> Result<Self::Output, ToolExecutionError> {
6685            Ok(args)
6686        }
6687    }
6688
6689    #[derive(serde::Deserialize)]
6690    struct FirstGenerationArgs {
6691        old: String,
6692    }
6693
6694    struct FirstGenerationTool(Arc<AtomicU32>);
6695
6696    impl Tool for FirstGenerationTool {
6697        const NAME: &'static str = "generation_pinned";
6698        type Error = rig::tool::ToolExecutionError;
6699        type Args = FirstGenerationArgs;
6700        type Output = String;
6701
6702        fn description(&self) -> String {
6703            "first generation schema".to_string()
6704        }
6705
6706        fn parameters(&self) -> serde_json::Value {
6707            json!({
6708                "type": "object",
6709                "properties": {"old": {"type": "string"}},
6710                "required": ["old"]
6711            })
6712        }
6713
6714        async fn call(
6715            &self,
6716            _context: &mut ToolContext,
6717            args: Self::Args,
6718        ) -> Result<Self::Output, ToolExecutionError> {
6719            self.0.fetch_add(1, SeqCst);
6720            Ok(format!("first:{}", args.old))
6721        }
6722    }
6723
6724    #[derive(serde::Deserialize)]
6725    struct SecondGenerationArgs {
6726        new: String,
6727    }
6728
6729    struct SecondGenerationTool(Arc<AtomicU32>);
6730
6731    impl Tool for SecondGenerationTool {
6732        const NAME: &'static str = FirstGenerationTool::NAME;
6733        type Error = rig::tool::ToolExecutionError;
6734        type Args = SecondGenerationArgs;
6735        type Output = String;
6736
6737        fn description(&self) -> String {
6738            "second generation schema".to_string()
6739        }
6740
6741        fn parameters(&self) -> serde_json::Value {
6742            json!({
6743                "type": "object",
6744                "properties": {"new": {"type": "string"}},
6745                "required": ["new"]
6746            })
6747        }
6748
6749        async fn call(
6750            &self,
6751            _context: &mut ToolContext,
6752            args: Self::Args,
6753        ) -> Result<Self::Output, ToolExecutionError> {
6754            self.0.fetch_add(1, SeqCst);
6755            Ok(format!("second:{}", args.new))
6756        }
6757    }
6758
6759    /// Pauses the first provider call after its request has been built. Tests
6760    /// replace the live registry while that request is in flight, then let the
6761    /// model return a call that is valid only for the advertised generation.
6762    #[derive(Clone)]
6763    struct PausingCompletionModel {
6764        inner: MockCompletionModel,
6765        request_started: Arc<Notify>,
6766        release_response: Arc<Notify>,
6767        requests: Arc<AtomicU32>,
6768    }
6769
6770    impl PausingCompletionModel {
6771        fn new(inner: MockCompletionModel) -> (Self, Arc<Notify>, Arc<Notify>) {
6772            let request_started = Arc::new(Notify::new());
6773            let release_response = Arc::new(Notify::new());
6774            (
6775                Self {
6776                    inner,
6777                    request_started: request_started.clone(),
6778                    release_response: release_response.clone(),
6779                    requests: Arc::new(AtomicU32::new(0)),
6780                },
6781                request_started,
6782                release_response,
6783            )
6784        }
6785
6786        async fn inspect_and_pause(&self, request: &crate::completion::CompletionRequest) {
6787            let request_index = self.requests.fetch_add(1, SeqCst);
6788            let definition = request
6789                .tools
6790                .iter()
6791                .find(|definition| definition.name == FirstGenerationTool::NAME)
6792                .expect("generation tool must be advertised");
6793            if request_index == 0 {
6794                assert_eq!(definition.description, "first generation schema");
6795                self.request_started.notify_one();
6796                self.release_response.notified().await;
6797            } else {
6798                assert_eq!(definition.description, "second generation schema");
6799            }
6800        }
6801    }
6802
6803    impl CompletionModel for PausingCompletionModel {
6804        type Response = crate::test_utils::MockResponse;
6805        type StreamingResponse = crate::test_utils::MockResponse;
6806        type Client = ();
6807
6808        fn make(_: &Self::Client, _: impl Into<String>) -> Self {
6809            Self::new(MockCompletionModel::default()).0
6810        }
6811
6812        async fn completion(
6813            &self,
6814            request: crate::completion::CompletionRequest,
6815        ) -> Result<
6816            crate::completion::CompletionResponse<Self::Response>,
6817            crate::completion::CompletionError,
6818        > {
6819            self.inspect_and_pause(&request).await;
6820            self.inner.completion(request).await
6821        }
6822
6823        async fn stream(
6824            &self,
6825            request: crate::completion::CompletionRequest,
6826        ) -> Result<
6827            crate::streaming::StreamingCompletionResponse<Self::StreamingResponse>,
6828            crate::completion::CompletionError,
6829        > {
6830            self.inspect_and_pause(&request).await;
6831            self.inner.stream(request).await
6832        }
6833    }
6834
6835    #[test]
6836    fn one_hook_instance_attaches_to_distinct_completion_models() {
6837        #[derive(Clone)]
6838        struct ProviderIndependentHook;
6839
6840        impl AgentHook for ProviderIndependentHook {}
6841
6842        let hook = ProviderIndependentHook;
6843        let _mock_agent = AgentBuilder::new(MockCompletionModel::default())
6844            .add_hook(hook.clone())
6845            .build();
6846        let (other_model, _, _) = PausingCompletionModel::new(MockCompletionModel::default());
6847        let _other_agent = AgentBuilder::new(other_model).add_hook(hook).build();
6848    }
6849
6850    /// `ToolCallAction::Rewrite` resolves to a `ProceedWith` tool-call decision that
6851    /// carries the replacement arguments, and is named for fail-closed
6852    /// diagnostics.
6853    #[test]
6854    fn rewrite_args_resolves_to_proceed_with_for_tool_call() {
6855        let args = json!({"x": 1, "y": 2});
6856        match super::tool_call_decision(ToolCallAction::rewrite(args.clone())) {
6857            super::ToolCallDecision::ProceedWith(replacement) => assert_eq!(replacement, args),
6858            _ => panic!("ToolCallAction::Rewrite should resolve to ProceedWith"),
6859        }
6860        // The typed convenience builds the same variant as the value constructor.
6861        assert_eq!(
6862            ToolCallAction::try_rewrite(&json!({"x": 1, "y": 2})).expect("serializes"),
6863            ToolCallAction::rewrite(json!({"x": 1, "y": 2})),
6864        );
6865    }
6866
6867    /// A hook that rewrites a *valid* tool call's arguments (`ToolCallAction::Rewrite`
6868    /// on `ToolCall`) is honored identically under `run()` and `stream()`: the
6869    /// tool executes with the replacement, so both drivers observe the same
6870    /// rewritten tool result and reach the same output, tool-result content and
6871    /// message history. Both drivers share `run_single_tool`, so they stay in
6872    /// lock-step.
6873    #[tokio::test]
6874    async fn valid_tool_call_rewrite_args_parity_across_run_and_stream() {
6875        // The model asks to add 2 + 3; the hook rewrites the arguments to 2 + 40,
6876        // so the tool returns 42 rather than 5.
6877        let turns = [
6878            ScriptedTurn::ToolCalls(vec![add_call("tc1", 2, 3)]),
6879            ScriptedTurn::Text("acknowledged"),
6880        ];
6881        let replacement = json!({"x": 2, "y": 40});
6882
6883        let blocking_model =
6884            MockCompletionModel::from_turns(turns.iter().map(ScriptedTurn::as_blocking_turn));
6885        let blocking_hook = RecordingHook::default();
6886        let blocking = AgentBuilder::new(blocking_model)
6887            .tool(MockAddTool)
6888            .build()
6889            .runner("add 2 and 3")
6890            .max_turns(3)
6891            .add_hook(blocking_hook.clone())
6892            .add_hook(RewriteToolArgsHook(replacement.clone()))
6893            .run()
6894            .await
6895            .expect("blocking run should succeed with rewritten tool arguments");
6896
6897        let streaming_model = MockCompletionModel::from_stream_turns(
6898            turns
6899                .iter()
6900                .map(|turn| turn.as_stream_events(StreamShape::Complete)),
6901        );
6902        let streaming_hook = RecordingHook::default();
6903        let mut stream = AgentBuilder::new(streaming_model)
6904            .tool(MockAddTool)
6905            .build()
6906            .runner("add 2 and 3")
6907            .max_turns(3)
6908            .add_hook(streaming_hook.clone())
6909            .add_hook(RewriteToolArgsHook(replacement))
6910            .stream()
6911            .await;
6912        let mut final_response = None;
6913        while let Some(item) = stream.next().await {
6914            if let Ok(MultiTurnStreamItem::FinalResponse(resp)) =
6915                item.map_err(|err| panic!("stream item errored: {err}"))
6916            {
6917                final_response = Some(resp);
6918            }
6919        }
6920        let final_response = final_response.expect("stream should yield a final response");
6921
6922        // The tool ran with the rewritten arguments (2 + 40 = 42), not the
6923        // model's emitted 2 + 3 = 5 — on both drivers.
6924        assert_eq!(blocking_hook.tool_results(), vec!["42".to_string()]);
6925        assert_eq!(blocking.output, "acknowledged");
6926        assert_eq!(final_response.output(), blocking.output);
6927        assert_eq!(
6928            blocking_hook.shared_events(),
6929            streaming_hook.shared_events()
6930        );
6931        assert_eq!(blocking_hook.tool_results(), streaming_hook.tool_results());
6932    }
6933
6934    #[tokio::test]
6935    async fn string_tool_call_without_rewrite_is_canonical_across_run_and_stream() {
6936        let turns = [
6937            ScriptedTurn::ToolCalls(vec![ScriptedToolCall {
6938                id: "tc-string",
6939                name: EchoStringArgs::NAME,
6940                args: json!("original"),
6941            }]),
6942            ScriptedTurn::Text("done"),
6943        ];
6944
6945        let blocking_hook = RecordingHook::default();
6946        let blocking = AgentBuilder::new(MockCompletionModel::from_turns(
6947            turns.iter().map(ScriptedTurn::as_blocking_turn),
6948        ))
6949        .tool(EchoStringArgs)
6950        .build()
6951        .runner("echo a string")
6952        .max_turns(3)
6953        .add_hook(blocking_hook.clone())
6954        .run()
6955        .await
6956        .expect("blocking string call should execute");
6957
6958        let streaming_hook = RecordingHook::default();
6959        let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns(
6960            turns
6961                .iter()
6962                .map(|turn| turn.as_stream_events(StreamShape::Complete)),
6963        ))
6964        .tool(EchoStringArgs)
6965        .build()
6966        .runner("echo a string")
6967        .max_turns(3)
6968        .add_hook(streaming_hook.clone())
6969        .stream()
6970        .await;
6971        let mut final_output = None;
6972        while let Some(item) = stream.next().await {
6973            if let MultiTurnStreamItem::FinalResponse(response) =
6974                item.expect("streaming string call should execute")
6975            {
6976                final_output = Some(response.output().to_string());
6977            }
6978        }
6979
6980        assert_eq!(blocking.output, "done");
6981        assert_eq!(final_output.as_deref(), Some("done"));
6982        assert_eq!(blocking_hook.tool_results(), vec!["original"]);
6983        assert_eq!(streaming_hook.tool_results(), vec!["original"]);
6984    }
6985
6986    #[tokio::test]
6987    async fn string_tool_call_rewrite_is_canonical_json_across_run_and_stream() {
6988        let turns = [
6989            ScriptedTurn::ToolCalls(vec![ScriptedToolCall {
6990                id: "tc-string",
6991                name: EchoStringArgs::NAME,
6992                args: json!("original"),
6993            }]),
6994            ScriptedTurn::Text("done"),
6995        ];
6996        let replacement = json!("sanitized");
6997
6998        let blocking_hook = RecordingHook::default();
6999        let blocking = AgentBuilder::new(MockCompletionModel::from_turns(
7000            turns.iter().map(ScriptedTurn::as_blocking_turn),
7001        ))
7002        .tool(EchoStringArgs)
7003        .build()
7004        .runner("echo a string")
7005        .max_turns(3)
7006        .add_hook(blocking_hook.clone())
7007        .add_hook(RewriteToolArgsHook(replacement.clone()))
7008        .run()
7009        .await
7010        .expect("blocking string rewrite should execute");
7011
7012        let streaming_hook = RecordingHook::default();
7013        let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns(
7014            turns
7015                .iter()
7016                .map(|turn| turn.as_stream_events(StreamShape::Complete)),
7017        ))
7018        .tool(EchoStringArgs)
7019        .build()
7020        .runner("echo a string")
7021        .max_turns(3)
7022        .add_hook(streaming_hook.clone())
7023        .add_hook(RewriteToolArgsHook(replacement))
7024        .stream()
7025        .await;
7026        let mut final_output = None;
7027        while let Some(item) = stream.next().await {
7028            if let MultiTurnStreamItem::FinalResponse(response) =
7029                item.expect("streaming string rewrite should execute")
7030            {
7031                final_output = Some(response.output().to_string());
7032            }
7033        }
7034
7035        assert_eq!(blocking.output, "done");
7036        assert_eq!(final_output.as_deref(), Some("done"));
7037        assert_eq!(blocking_hook.tool_results(), vec!["sanitized"]);
7038        assert_eq!(streaming_hook.tool_results(), vec!["sanitized"]);
7039    }
7040
7041    #[tokio::test]
7042    async fn blocking_turn_dispatches_the_registry_generation_it_advertised() {
7043        let first_calls = Arc::new(AtomicU32::new(0));
7044        let second_calls = Arc::new(AtomicU32::new(0));
7045        let handle: ToolServerHandle = ToolServer::new()
7046            .tool(FirstGenerationTool(first_calls.clone()))
7047            .run();
7048        let inner = MockCompletionModel::from_turns([
7049            MockTurn::tool_call(
7050                "tc-generation",
7051                FirstGenerationTool::NAME,
7052                json!({"old": "payload"}),
7053            ),
7054            MockTurn::text("done"),
7055        ]);
7056        let (model, request_started, release_response) = PausingCompletionModel::new(inner);
7057        let runner = AgentBuilder::new(model)
7058            .tool_server_handle(handle.clone())
7059            .build()
7060            .runner("use the generation tool")
7061            .max_turns(3);
7062
7063        let run = runner.run();
7064        let replace = async {
7065            request_started.notified().await;
7066            handle
7067                .add_tool(SecondGenerationTool(second_calls.clone()))
7068                .await;
7069            release_response.notify_one();
7070        };
7071        let (response, ()) = tokio::time::timeout(std::time::Duration::from_secs(2), async {
7072            tokio::join!(run, replace)
7073        })
7074        .await
7075        .expect("in-flight blocking replacement must not hang");
7076        let response = response.expect("blocking run should use its pinned tool generation");
7077
7078        assert_eq!(response.output, "done");
7079        assert_eq!(first_calls.load(SeqCst), 1);
7080        assert_eq!(second_calls.load(SeqCst), 0);
7081    }
7082
7083    #[tokio::test]
7084    async fn streaming_turn_dispatches_the_registry_generation_it_advertised() {
7085        let first_calls = Arc::new(AtomicU32::new(0));
7086        let second_calls = Arc::new(AtomicU32::new(0));
7087        let handle: ToolServerHandle = ToolServer::new()
7088            .tool(FirstGenerationTool(first_calls.clone()))
7089            .run();
7090        let turns = [
7091            ScriptedTurn::ToolCalls(vec![ScriptedToolCall {
7092                id: "tc-generation",
7093                name: FirstGenerationTool::NAME,
7094                args: json!({"old": "payload"}),
7095            }]),
7096            ScriptedTurn::Text("done"),
7097        ];
7098        let inner = MockCompletionModel::from_stream_turns(
7099            turns
7100                .iter()
7101                .map(|turn| turn.as_stream_events(StreamShape::Complete)),
7102        );
7103        let (model, request_started, release_response) = PausingCompletionModel::new(inner);
7104        let runner = AgentBuilder::new(model)
7105            .tool_server_handle(handle.clone())
7106            .build()
7107            .runner("use the generation tool")
7108            .max_turns(3);
7109
7110        let drive = async {
7111            let mut stream = runner.stream().await;
7112            let mut final_output = None;
7113            while let Some(item) = stream.next().await {
7114                if let MultiTurnStreamItem::FinalResponse(response) =
7115                    item.expect("streaming run should use its pinned tool generation")
7116                {
7117                    final_output = Some(response.output().to_string());
7118                }
7119            }
7120            final_output
7121        };
7122        let replace = async {
7123            request_started.notified().await;
7124            handle
7125                .add_tool(SecondGenerationTool(second_calls.clone()))
7126                .await;
7127            release_response.notify_one();
7128        };
7129        let (final_output, ()) = tokio::time::timeout(std::time::Duration::from_secs(2), async {
7130            tokio::join!(drive, replace)
7131        })
7132        .await
7133        .expect("in-flight streaming replacement must not hang");
7134
7135        assert_eq!(final_output.as_deref(), Some("done"));
7136        assert_eq!(first_calls.load(SeqCst), 1);
7137        assert_eq!(second_calls.load(SeqCst), 0);
7138    }
7139
7140    /// A hook that rewrites a tool's result (`ToolResultAction::Rewrite` on
7141    /// `ToolResult`) so the model sees the replacement instead of the tool's
7142    /// actual output.
7143    struct RewriteToolResultHook(&'static str);
7144
7145    impl AgentHook for RewriteToolResultHook {
7146        async fn on_tool_result(
7147            &self,
7148            _ctx: &HookContext,
7149            event: ToolResultEvent<'_>,
7150        ) -> ToolResultAction {
7151            if let ToolResultEvent { .. } = event {
7152                ToolResultAction::rewrite(self.0)
7153            } else {
7154                ToolResultAction::keep()
7155            }
7156        }
7157    }
7158
7159    /// `ToolResultAction::Rewrite` resolves to a `Replace` tool-result decision carrying
7160    /// the replacement, and is named for fail-closed diagnostics.
7161    #[test]
7162    fn rewrite_result_resolves_to_replace_for_tool_result() {
7163        match super::tool_result_decision(ToolResultAction::rewrite("redacted")) {
7164            super::ToolResultDecision::Replace(result) => {
7165                assert_eq!(result.as_text(), Some("redacted"))
7166            }
7167            _ => panic!("ToolResultAction::Rewrite should resolve to Replace"),
7168        }
7169    }
7170
7171    /// A hook that rewrites a tool's result (`ToolResultAction::Rewrite` on
7172    /// `ToolResult`) is honored identically under `run()` and `stream()`: the
7173    /// model-visible history carries the replacement while the `ToolResult` event
7174    /// still observed the tool's actual output, and both drivers reach the same
7175    /// output and history. Both share `run_single_tool`, so they stay in
7176    /// lock-step.
7177    #[tokio::test]
7178    async fn valid_tool_result_rewrite_parity_across_run_and_stream() {
7179        // The tool computes 2 + 3 = 5; the hook replaces what the model sees with
7180        // "redacted-result".
7181        let turns = [
7182            ScriptedTurn::ToolCalls(vec![add_call("tc1", 2, 3)]),
7183            ScriptedTurn::Text("acknowledged"),
7184        ];
7185
7186        let blocking_model =
7187            MockCompletionModel::from_turns(turns.iter().map(ScriptedTurn::as_blocking_turn));
7188        let blocking_hook = RecordingHook::default();
7189        let blocking = AgentBuilder::new(blocking_model)
7190            .tool(MockAddTool)
7191            .build()
7192            .runner("add 2 and 3")
7193            .max_turns(3)
7194            .add_hook(blocking_hook.clone())
7195            .add_hook(RewriteToolResultHook("redacted-result"))
7196            .run()
7197            .await
7198            .expect("blocking run should succeed with a rewritten tool result");
7199
7200        let streaming_model = MockCompletionModel::from_stream_turns(
7201            turns
7202                .iter()
7203                .map(|turn| turn.as_stream_events(StreamShape::Complete)),
7204        );
7205        let streaming_hook = RecordingHook::default();
7206        let mut stream = AgentBuilder::new(streaming_model)
7207            .tool(MockAddTool)
7208            .build()
7209            .runner("add 2 and 3")
7210            .max_turns(3)
7211            .add_hook(streaming_hook.clone())
7212            .add_hook(RewriteToolResultHook("redacted-result"))
7213            .stream()
7214            .await;
7215        let mut final_response = None;
7216        while let Some(item) = stream.next().await {
7217            if let Ok(MultiTurnStreamItem::FinalResponse(resp)) =
7218                item.map_err(|err| panic!("stream item errored: {err}"))
7219            {
7220                final_response = Some(resp);
7221            }
7222        }
7223        let final_response = final_response.expect("stream should yield a final response");
7224
7225        assert_eq!(blocking.output, "acknowledged");
7226        assert_eq!(final_response.output(), blocking.output);
7227
7228        // The ToolResult event observes the tool's ACTUAL output (5) on both
7229        // drivers — the replacement is applied after the event fires.
7230        assert_eq!(blocking_hook.tool_results(), vec!["5".to_string()]);
7231        assert_eq!(blocking_hook.tool_results(), streaming_hook.tool_results());
7232
7233        // The model-visible history carries the REWRITTEN result, not "5", and is
7234        // byte-identical across drivers.
7235        let blocking_messages = blocking.messages.expect("blocking messages");
7236        let streaming_messages = final_response
7237            .messages()
7238            .expect("streaming history")
7239            .to_vec();
7240        assert_eq!(
7241            serde_json::to_value(&blocking_messages).expect("serialize blocking"),
7242            serde_json::to_value(&streaming_messages).expect("serialize streaming"),
7243        );
7244        assert!(
7245            tool_result_text_in_history(&blocking_messages, "redacted-result"),
7246            "the model-visible tool result must be the hook's replacement"
7247        );
7248        assert!(
7249            !tool_result_text_in_history(&blocking_messages, "5"),
7250            "the tool's original output must not reach the model after a rewrite"
7251        );
7252    }
7253
7254    /// A `ToolResultAction::Rewrite` replacement is delivered to the model verbatim, not
7255    /// re-parsed as structured/multimodal tool output. A JSON-shaped replacement
7256    /// (here, an image payload that `tool_result_output` would turn into an image
7257    /// content block for *real* tool output) reaches history as literal text —
7258    /// so a redaction hook returning JSON cannot be silently restructured.
7259    #[tokio::test]
7260    async fn rewrite_result_is_delivered_verbatim_not_reparsed() {
7261        const IMAGE_JSON: &str = r#"{"type":"image","data":"abc","mimeType":"image/png"}"#;
7262
7263        let turns = [
7264            ScriptedTurn::ToolCalls(vec![add_call("tc1", 2, 3)]),
7265            ScriptedTurn::Text("done"),
7266        ];
7267        let model =
7268            MockCompletionModel::from_turns(turns.iter().map(ScriptedTurn::as_blocking_turn));
7269        let result = AgentBuilder::new(model)
7270            .tool(MockAddTool)
7271            .build()
7272            .runner("add 2 and 3")
7273            .max_turns(3)
7274            .add_hook(RewriteToolResultHook(IMAGE_JSON))
7275            .run()
7276            .await
7277            .expect("run should succeed with a JSON-shaped rewritten result");
7278
7279        let messages = result.messages.expect("messages");
7280        assert!(
7281            tool_result_text_in_history(&messages, IMAGE_JSON),
7282            "the JSON-shaped replacement must reach history verbatim as text, not be \
7283             re-parsed into a structured/image content block"
7284        );
7285    }
7286
7287    /// A hook that patches the model request for the turn (`CompletionCallAction::Patch`
7288    /// on `CompletionCall`): forces tool_choice + temperature, narrows the
7289    /// advertised tools to an allow-list, and injects a passthrough param.
7290    struct PatchRequestHook;
7291
7292    impl AgentHook for PatchRequestHook {
7293        async fn on_completion_call(
7294            &self,
7295            _ctx: &HookContext,
7296            event: CompletionCallEvent<'_>,
7297        ) -> CompletionCallAction {
7298            if let CompletionCallEvent { .. } = event {
7299                CompletionCallAction::patch(
7300                    RequestPatch::new()
7301                        .preamble(OVERRIDE_PREAMBLE)
7302                        .temperature(0.25)
7303                        .max_tokens(OVERRIDE_MAX_TOKENS)
7304                        .tool_choice(ToolChoice::Required)
7305                        .active_tools(["add"])
7306                        .additional_params(json!({"injected": true})),
7307                )
7308            } else {
7309                CompletionCallAction::continue_run()
7310            }
7311        }
7312    }
7313
7314    const OVERRIDE_PREAMBLE: &str = "overridden: critical-step instructions";
7315    const OVERRIDE_MAX_TOKENS: u64 = 512;
7316
7317    /// `CompletionCallAction::Patch` resolves to a `Patch` completion-call decision
7318    /// carrying the patch, and is named for fail-closed diagnostics.
7319    #[test]
7320    fn patch_request_resolves_to_patch_for_completion_call() {
7321        let patch = RequestPatch::new()
7322            .temperature(0.25)
7323            .tool_choice(ToolChoice::Required);
7324        match super::completion_call_decision(CompletionCallAction::patch(patch.clone())) {
7325            super::CompletionCallDecision::Patch(got) => assert_eq!(got, patch),
7326            _ => panic!("PatchRequest should resolve to Patch for a completion call"),
7327        }
7328    }
7329
7330    /// A `CompletionCallAction::Patch` hook patches the request for the turn identically
7331    /// under `run()` and `stream()`: the captured completion request shows the
7332    /// overridden temperature/tool_choice, the merged additional_params, and the
7333    /// tool set narrowed to the allow-list — on both drivers.
7334    #[tokio::test]
7335    async fn patch_request_parity_across_run_and_stream() {
7336        fn assert_request(req: &crate::completion::CompletionRequest) {
7337            assert_eq!(
7338                req.temperature,
7339                Some(0.25),
7340                "override temperature wins over the agent's 0.9"
7341            );
7342            assert_eq!(
7343                req.max_tokens,
7344                Some(OVERRIDE_MAX_TOKENS),
7345                "override max_tokens wins over the agent's 64"
7346            );
7347            // The override preamble wins and is sent as the leading system message.
7348            let system = req.chat_history.iter().find_map(|m| match m {
7349                Message::System { content } => Some(content.as_str()),
7350                _ => None,
7351            });
7352            assert_eq!(
7353                system,
7354                Some(OVERRIDE_PREAMBLE),
7355                "override preamble wins over the agent's baseline and is the leading system message"
7356            );
7357            assert!(matches!(req.tool_choice, Some(ToolChoice::Required)));
7358            let tool_names: Vec<&str> = req.tools.iter().map(|t| t.name.as_str()).collect();
7359            assert_eq!(
7360                tool_names,
7361                ["add"],
7362                "active_tools narrows the advertised set to `add` (drops `subtract`)"
7363            );
7364            // The runner replaces the agent baseline, then the hook shallow-merges
7365            // last and therefore wins conflicts.
7366            let params = req.additional_params.as_ref().expect("additional_params");
7367            assert_eq!(params.get("runner").and_then(|v| v.as_str()), Some("keep"));
7368            assert_eq!(params.get("injected").and_then(|v| v.as_bool()), Some(true));
7369            assert!(params.get("baseline").is_none());
7370        }
7371
7372        let blocking_model = MockCompletionModel::from_turns([MockTurn::text("done")]);
7373        let blocking_probe = blocking_model.clone();
7374        let blocking = AgentBuilder::new(blocking_model)
7375            .tool(MockAddTool)
7376            .tool(MockSubtractTool)
7377            .preamble("baseline preamble")
7378            .temperature(0.9)
7379            .max_tokens(64)
7380            .additional_params(json!({"baseline": "keep"}))
7381            .add_hook(PatchRequestHook)
7382            .build()
7383            .runner("go")
7384            .replace_additional_params(json!({"runner": "keep", "injected": false}))
7385            .max_turns(2)
7386            .run()
7387            .await
7388            .expect("blocking run should succeed");
7389        assert_eq!(blocking.output, "done");
7390        let blocking_requests = blocking_probe.requests();
7391        assert_eq!(blocking_requests.len(), 1);
7392        assert_request(&blocking_requests[0]);
7393
7394        let streaming_model = MockCompletionModel::from_stream_turns([
7395            ScriptedTurn::Text("done").as_stream_events(StreamShape::Complete)
7396        ]);
7397        let streaming_probe = streaming_model.clone();
7398        let mut stream = AgentBuilder::new(streaming_model)
7399            .tool(MockAddTool)
7400            .tool(MockSubtractTool)
7401            .preamble("baseline preamble")
7402            .temperature(0.9)
7403            .max_tokens(64)
7404            .additional_params(json!({"baseline": "keep"}))
7405            .add_hook(PatchRequestHook)
7406            .build()
7407            .runner("go")
7408            .replace_additional_params(json!({"runner": "keep", "injected": false}))
7409            .max_turns(2)
7410            .stream()
7411            .await;
7412        while let Some(item) = stream.next().await {
7413            let _ = item.map_err(|err| panic!("stream item errored: {err}"));
7414        }
7415        let streaming_requests = streaming_probe.requests();
7416        assert_eq!(streaming_requests.len(), 1);
7417        assert_request(&streaming_requests[0]);
7418    }
7419
7420    // --- Hook system v2: extra_context, history view, ModelTurnFinished, chained rewrites ---
7421
7422    fn hook_doc(id: &str, text: &str) -> crate::completion::Document {
7423        crate::completion::Document {
7424            id: id.to_string(),
7425            text: text.to_string(),
7426            additional_props: Default::default(),
7427        }
7428    }
7429
7430    /// Injects one extra context document on every completion call.
7431    struct ExtraContextHook {
7432        id: &'static str,
7433        text: &'static str,
7434    }
7435
7436    impl AgentHook for ExtraContextHook {
7437        async fn on_completion_call(
7438            &self,
7439            _ctx: &HookContext,
7440            event: CompletionCallEvent<'_>,
7441        ) -> CompletionCallAction {
7442            if let CompletionCallEvent { .. } = event {
7443                CompletionCallAction::patch(
7444                    RequestPatch::new().context(hook_doc(self.id, self.text)),
7445                )
7446            } else {
7447                CompletionCallAction::continue_run()
7448            }
7449        }
7450    }
7451
7452    /// Injects an extra context document only on the first turn (to prove
7453    /// per-turn, non-sticky behavior).
7454    struct ExtraContextTurnOneHook;
7455
7456    impl AgentHook for ExtraContextTurnOneHook {
7457        async fn on_completion_call(
7458            &self,
7459            _ctx: &HookContext,
7460            event: CompletionCallEvent<'_>,
7461        ) -> CompletionCallAction {
7462            if let CompletionCallEvent { turn, .. } = event
7463                && turn == 1
7464            {
7465                return CompletionCallAction::patch(
7466                    RequestPatch::new().context(hook_doc("turn-one", "only turn 1")),
7467                );
7468            }
7469            CompletionCallAction::continue_run()
7470        }
7471    }
7472
7473    #[derive(Clone)]
7474    struct RecordingContextIndex {
7475        id: &'static str,
7476        queries: Arc<Mutex<Vec<(String, u64)>>>,
7477    }
7478
7479    impl VectorStoreIndex for RecordingContextIndex {
7480        type Filter = Filter<serde_json::Value>;
7481
7482        async fn top_n<T: for<'a> Deserialize<'a> + WasmCompatSend>(
7483            &self,
7484            req: VectorSearchRequest,
7485        ) -> Result<Vec<(f64, String, T)>, VectorStoreError> {
7486            self.queries
7487                .lock()
7488                .expect("context query recorder lock")
7489                .push((req.query().to_string(), req.samples()));
7490            let value = serde_json::from_value(json!({ "source": self.id }))?;
7491            Ok(vec![(1.0, self.id.to_string(), value)])
7492        }
7493
7494        async fn top_n_ids(
7495            &self,
7496            _req: VectorSearchRequest,
7497        ) -> Result<Vec<(f64, String)>, VectorStoreError> {
7498            Ok(vec![(1.0, self.id.to_string())])
7499        }
7500    }
7501
7502    struct FailingContextIndex;
7503
7504    impl VectorStoreIndex for FailingContextIndex {
7505        type Filter = Filter<serde_json::Value>;
7506
7507        async fn top_n<T: for<'a> Deserialize<'a> + WasmCompatSend>(
7508            &self,
7509            _req: VectorSearchRequest,
7510        ) -> Result<Vec<(f64, String, T)>, VectorStoreError> {
7511            Err(VectorStoreError::BuilderError(
7512                "context index unavailable".to_string(),
7513            ))
7514        }
7515
7516        async fn top_n_ids(
7517            &self,
7518            _req: VectorSearchRequest,
7519        ) -> Result<Vec<(f64, String)>, VectorStoreError> {
7520            Err(VectorStoreError::BuilderError(
7521                "context index unavailable".to_string(),
7522            ))
7523        }
7524    }
7525
7526    struct QueryRecordingToolIndex {
7527        queries: Arc<Mutex<Vec<String>>>,
7528    }
7529
7530    impl VectorStoreIndex for QueryRecordingToolIndex {
7531        type Filter = Filter<serde_json::Value>;
7532
7533        async fn top_n<T: for<'a> Deserialize<'a> + WasmCompatSend>(
7534            &self,
7535            _req: VectorSearchRequest,
7536        ) -> Result<Vec<(f64, String, T)>, VectorStoreError> {
7537            Ok(Vec::new())
7538        }
7539
7540        async fn top_n_ids(
7541            &self,
7542            req: VectorSearchRequest,
7543        ) -> Result<Vec<(f64, String)>, VectorStoreError> {
7544            self.queries
7545                .lock()
7546                .expect("query recorder lock")
7547                .push(req.query().to_string());
7548            Ok(vec![(1.0, MockAddTool::NAME.to_string())])
7549        }
7550    }
7551
7552    fn one_text_stream_turn(text: &'static str) -> Vec<MockStreamEvent> {
7553        vec![
7554            MockStreamEvent::text(text),
7555            MockStreamEvent::final_response_with_total_tokens(0),
7556        ]
7557    }
7558
7559    /// A single hook's `extra_context` document appears in the completion request,
7560    /// after the agent's static context, on both `run()` and `stream()`.
7561    #[tokio::test]
7562    async fn extra_context_appears_after_static_context_on_both_surfaces() {
7563        fn assert_docs(req: &crate::completion::CompletionRequest) {
7564            let ids: Vec<&str> = req.documents.iter().map(|d| d.id.as_str()).collect();
7565            let static_pos = ids
7566                .iter()
7567                .position(|id| id.starts_with("static_doc"))
7568                .expect("static context document present");
7569            let extra_pos = ids
7570                .iter()
7571                .position(|id| *id == "hook-doc")
7572                .expect("hook extra_context document present");
7573            assert!(
7574                static_pos < extra_pos,
7575                "static context precedes hook extras: {ids:?}"
7576            );
7577            assert!(
7578                req.documents.iter().any(|d| d.text == "injected"),
7579                "the hook document's text is present"
7580            );
7581        }
7582
7583        let blocking_model = MockCompletionModel::from_turns([MockTurn::text("done")]);
7584        let blocking_probe = blocking_model.clone();
7585        AgentBuilder::new(blocking_model)
7586            .context("static context text")
7587            .add_hook(ExtraContextHook {
7588                id: "hook-doc",
7589                text: "injected",
7590            })
7591            .build()
7592            .runner("go")
7593            .run()
7594            .await
7595            .expect("blocking run should succeed");
7596        assert_docs(blocking_probe.requests().first().expect("one request"));
7597
7598        let streaming_model =
7599            MockCompletionModel::from_stream_turns([one_text_stream_turn("done")]);
7600        let streaming_probe = streaming_model.clone();
7601        let mut stream = AgentBuilder::new(streaming_model)
7602            .context("static context text")
7603            .add_hook(ExtraContextHook {
7604                id: "hook-doc",
7605                text: "injected",
7606            })
7607            .build()
7608            .runner("go")
7609            .stream()
7610            .await;
7611        while let Some(item) = stream.next().await {
7612            let _ = item.map_err(|err| panic!("stream item errored: {err}"));
7613        }
7614        assert_docs(streaming_probe.requests().first().expect("one request"));
7615    }
7616
7617    /// Two hooks' `extra_context` documents append in registration order.
7618    #[tokio::test]
7619    async fn multiple_hooks_extra_context_append_in_registration_order() {
7620        let model = MockCompletionModel::from_turns([MockTurn::text("done")]);
7621        let probe = model.clone();
7622        AgentBuilder::new(model)
7623            .add_hook(ExtraContextHook {
7624                id: "first",
7625                text: "1",
7626            })
7627            .add_hook(ExtraContextHook {
7628                id: "second",
7629                text: "2",
7630            })
7631            .build()
7632            .runner("go")
7633            .run()
7634            .await
7635            .expect("run should succeed");
7636        let requests = probe.requests();
7637        let req = requests.first().expect("one request");
7638        let ids: Vec<&str> = req.documents.iter().map(|d| d.id.as_str()).collect();
7639        assert_eq!(
7640            ids,
7641            vec!["first", "second"],
7642            "hook extras append in registration order"
7643        );
7644    }
7645
7646    #[tokio::test]
7647    async fn dynamic_context_preserves_query_selection_formatting_and_order_on_both_surfaces() {
7648        fn assert_documents(request: &crate::completion::CompletionRequest) {
7649            let documents = request
7650                .documents
7651                .iter()
7652                .map(|document| (document.id.as_str(), document.text.as_str()))
7653                .collect::<Vec<_>>();
7654            assert_eq!(
7655                documents,
7656                vec![
7657                    ("static_doc_0", "static context"),
7658                    ("blocking", "{\n  \"source\": \"blocking\"\n}"),
7659                ]
7660            );
7661        }
7662
7663        let blocking_queries = Arc::new(Mutex::new(Vec::new()));
7664        let blocking_model = MockCompletionModel::from_turns([MockTurn::text("done")]);
7665        let blocking_probe = blocking_model.clone();
7666        AgentBuilder::new(blocking_model)
7667            .context("static context")
7668            .dynamic_context(
7669                2,
7670                RecordingContextIndex {
7671                    id: "blocking",
7672                    queries: blocking_queries.clone(),
7673                },
7674            )
7675            .build()
7676            .runner("current blocking query")
7677            .history(vec![Message::user("ignored history query")])
7678            .run()
7679            .await
7680            .expect("blocking dynamic-context run should succeed");
7681        assert_eq!(
7682            *blocking_queries.lock().expect("blocking queries"),
7683            vec![("current blocking query".to_string(), 2)]
7684        );
7685        assert_documents(blocking_probe.requests().first().expect("one request"));
7686
7687        let streaming_queries = Arc::new(Mutex::new(Vec::new()));
7688        let streaming_model =
7689            MockCompletionModel::from_stream_turns([one_text_stream_turn("done")]);
7690        let streaming_probe = streaming_model.clone();
7691        let mut stream = AgentBuilder::new(streaming_model)
7692            .dynamic_context(
7693                3,
7694                RecordingContextIndex {
7695                    id: "streaming",
7696                    queries: streaming_queries.clone(),
7697                },
7698            )
7699            .build()
7700            .runner(Message::User {
7701                content: OneOrMany::one(UserContent::image_url(
7702                    "https://example.com/prompt.png",
7703                    None,
7704                    None,
7705                )),
7706            })
7707            .history(vec![
7708                Message::user("older history query"),
7709                Message::user("latest history query"),
7710            ])
7711            .stream()
7712            .await;
7713        while let Some(item) = stream.next().await {
7714            item.expect("streaming dynamic-context run should succeed");
7715        }
7716        assert_eq!(
7717            *streaming_queries.lock().expect("streaming queries"),
7718            vec![("latest history query".to_string(), 3)]
7719        );
7720        let streaming_requests = streaming_probe.requests();
7721        let request = streaming_requests.first().expect("one request");
7722        assert_eq!(request.documents.len(), 1);
7723        assert_eq!(request.documents[0].id, "streaming");
7724        assert_eq!(
7725            request.documents[0].text,
7726            "{\n  \"source\": \"streaming\"\n}"
7727        );
7728    }
7729
7730    #[tokio::test]
7731    async fn dynamic_context_and_application_hooks_follow_registration_order() {
7732        let queries = Arc::new(Mutex::new(Vec::new()));
7733        let model = MockCompletionModel::from_turns([MockTurn::text("done")]);
7734        let probe = model.clone();
7735        AgentBuilder::new(model)
7736            .context("static")
7737            .add_hook(ExtraContextHook {
7738                id: "before",
7739                text: "before dynamic context",
7740            })
7741            .dynamic_context(
7742                1,
7743                RecordingContextIndex {
7744                    id: "first",
7745                    queries: queries.clone(),
7746                },
7747            )
7748            .add_hook(ExtraContextHook {
7749                id: "between",
7750                text: "between dynamic contexts",
7751            })
7752            .dynamic_context(
7753                2,
7754                RecordingContextIndex {
7755                    id: "second",
7756                    queries: queries.clone(),
7757                },
7758            )
7759            .add_hook(ExtraContextHook {
7760                id: "after",
7761                text: "after dynamic context",
7762            })
7763            .build()
7764            .runner("query")
7765            .run()
7766            .await
7767            .expect("run should succeed");
7768
7769        assert_eq!(
7770            probe.requests()[0]
7771                .documents
7772                .iter()
7773                .map(|document| document.id.as_str())
7774                .collect::<Vec<_>>(),
7775            vec![
7776                "static_doc_0",
7777                "before",
7778                "first",
7779                "between",
7780                "second",
7781                "after",
7782            ]
7783        );
7784        assert_eq!(
7785            *queries.lock().expect("context queries"),
7786            vec![("query".to_string(), 1), ("query".to_string(), 2)]
7787        );
7788
7789        let skipped_queries = Arc::new(Mutex::new(Vec::new()));
7790        let error = AgentBuilder::new(MockCompletionModel::from_turns([MockTurn::text("unused")]))
7791            .add_hook(TerminateOn(StepEventKind::CompletionCall))
7792            .dynamic_context(
7793                1,
7794                RecordingContextIndex {
7795                    id: "skipped",
7796                    queries: skipped_queries.clone(),
7797                },
7798            )
7799            .build()
7800            .runner("query")
7801            .run()
7802            .await
7803            .expect_err("an earlier stop hook should terminate before retrieval");
7804        assert!(matches!(error, PromptError::PromptCancelled { .. }));
7805        assert!(skipped_queries.lock().expect("skipped queries").is_empty());
7806    }
7807
7808    #[tokio::test]
7809    async fn dynamic_context_retrieval_failure_stops_before_provider_io_on_both_surfaces() {
7810        let blocking_model = MockCompletionModel::from_turns([MockTurn::text("unused")]);
7811        let blocking_probe = blocking_model.clone();
7812        let error = AgentBuilder::new(blocking_model)
7813            .dynamic_context(1, FailingContextIndex)
7814            .build()
7815            .runner("retrieve this")
7816            .run()
7817            .await
7818            .expect_err("failed retrieval should stop the run");
7819        assert!(matches!(
7820            error,
7821            PromptError::PromptCancelled { reason, .. }
7822                if reason.contains("context index unavailable")
7823        ));
7824        assert_eq!(blocking_probe.request_count(), 0);
7825
7826        let streaming_model =
7827            MockCompletionModel::from_stream_turns([one_text_stream_turn("unused")]);
7828        let streaming_probe = streaming_model.clone();
7829        let mut stream = AgentBuilder::new(streaming_model)
7830            .dynamic_context(1, FailingContextIndex)
7831            .build()
7832            .runner("retrieve this")
7833            .stream()
7834            .await;
7835        let error = stream
7836            .next()
7837            .await
7838            .expect("stream should report retrieval failure")
7839            .expect_err("failed retrieval should stop the stream");
7840        assert!(matches!(
7841            error,
7842            StreamingError::Prompt(prompt_error)
7843                if matches!(
7844                    prompt_error.as_ref(),
7845                    PromptError::PromptCancelled { reason, .. }
7846                        if reason.contains("context index unavailable")
7847                )
7848        ));
7849        assert_eq!(streaming_probe.request_count(), 0);
7850    }
7851
7852    #[tokio::test]
7853    async fn retrieved_tool_query_selection_is_unchanged_on_both_surfaces() {
7854        let queries = Arc::new(Mutex::new(Vec::new()));
7855        AgentBuilder::new(MockCompletionModel::from_turns([MockTurn::text("done")]))
7856            .retrieved_tools(
7857                1,
7858                QueryRecordingToolIndex {
7859                    queries: queries.clone(),
7860                },
7861                ToolSet::from_tools(vec![MockAddTool]),
7862            )
7863            .build()
7864            .runner("blocking retrieval query")
7865            .history(vec![Message::user("blocking history query")])
7866            .run()
7867            .await
7868            .expect("blocking run should succeed");
7869
7870        AgentBuilder::new(MockCompletionModel::from_turns([MockTurn::text("done")]))
7871            .retrieved_tools(
7872                1,
7873                QueryRecordingToolIndex {
7874                    queries: queries.clone(),
7875                },
7876                ToolSet::from_tools(vec![MockAddTool]),
7877            )
7878            .build()
7879            .runner(Message::User {
7880                content: OneOrMany::one(UserContent::image_url(
7881                    "https://example.com/blocking.png",
7882                    None,
7883                    None,
7884                )),
7885            })
7886            .history(vec![
7887                Message::user("older blocking history query"),
7888                Message::user("latest blocking history query"),
7889            ])
7890            .run()
7891            .await
7892            .expect("blocking history fallback should succeed");
7893
7894        let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([
7895            one_text_stream_turn("done"),
7896        ]))
7897        .retrieved_tools(
7898            1,
7899            QueryRecordingToolIndex {
7900                queries: queries.clone(),
7901            },
7902            ToolSet::from_tools(vec![MockAddTool]),
7903        )
7904        .build()
7905        .runner("streaming retrieval query")
7906        .history(vec![Message::user("streaming history query")])
7907        .stream()
7908        .await;
7909        while let Some(item) = stream.next().await {
7910            item.expect("stream item should succeed");
7911        }
7912
7913        let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([
7914            one_text_stream_turn("done"),
7915        ]))
7916        .retrieved_tools(
7917            1,
7918            QueryRecordingToolIndex {
7919                queries: queries.clone(),
7920            },
7921            ToolSet::from_tools(vec![MockAddTool]),
7922        )
7923        .build()
7924        .runner(Message::User {
7925            content: OneOrMany::one(UserContent::image_url(
7926                "https://example.com/streaming.png",
7927                None,
7928                None,
7929            )),
7930        })
7931        .history(vec![
7932            Message::user("older streaming history query"),
7933            Message::user("latest streaming history query"),
7934        ])
7935        .stream()
7936        .await;
7937        while let Some(item) = stream.next().await {
7938            item.expect("stream item should succeed");
7939        }
7940
7941        assert_eq!(
7942            *queries.lock().expect("query recorder lock"),
7943            vec![
7944                "blocking retrieval query",
7945                "latest blocking history query",
7946                "streaming retrieval query",
7947                "latest streaming history query",
7948            ]
7949        );
7950    }
7951
7952    /// A hook's `extra_context` is per-turn and non-sticky: a document injected on
7953    /// turn 1 does not reappear on turn 2. Checked on both surfaces.
7954    #[tokio::test]
7955    async fn extra_context_is_per_turn_non_sticky() {
7956        fn assert_turns(requests: &[crate::completion::CompletionRequest]) {
7957            assert_eq!(requests.len(), 2, "two model turns");
7958            let turn1 = requests.first().expect("turn 1");
7959            let turn2 = requests.get(1).expect("turn 2");
7960            assert!(
7961                turn1.documents.iter().any(|d| d.id == "turn-one"),
7962                "turn 1 carries the injected document"
7963            );
7964            assert!(
7965                turn2.documents.iter().all(|d| d.id != "turn-one"),
7966                "turn 2 does not inherit turn 1's per-turn document"
7967            );
7968        }
7969
7970        let blocking_probe = blocking_model();
7971        let probe = blocking_probe.clone();
7972        AgentBuilder::new(blocking_probe)
7973            .tool(MockAddTool)
7974            .add_hook(ExtraContextTurnOneHook)
7975            .build()
7976            .runner("add 2 and 3")
7977            .max_turns(3)
7978            .run()
7979            .await
7980            .expect("blocking run should succeed");
7981        assert_turns(&probe.requests());
7982
7983        let streaming = streaming_model();
7984        let stream_probe = streaming.clone();
7985        let mut stream = AgentBuilder::new(streaming)
7986            .tool(MockAddTool)
7987            .add_hook(ExtraContextTurnOneHook)
7988            .build()
7989            .runner("add 2 and 3")
7990            .max_turns(3)
7991            .stream()
7992            .await;
7993        while let Some(item) = stream.next().await {
7994            let _ = item.map_err(|err| panic!("stream item errored: {err}"));
7995        }
7996        assert_turns(&stream_probe.requests());
7997    }
7998
7999    /// A hook that overrides `history` changes the messages sent to the provider
8000    /// for the turn without touching the persisted transcript, on both surfaces.
8001    #[tokio::test]
8002    async fn history_patch_changes_sent_messages_not_transcript_on_both_surfaces() {
8003        const SENTINEL: &str = "COMPACTED-HISTORY-SENTINEL";
8004
8005        struct HistoryOverrideHook;
8006        impl AgentHook for HistoryOverrideHook {
8007            async fn on_completion_call(
8008                &self,
8009                _ctx: &HookContext,
8010                event: CompletionCallEvent<'_>,
8011            ) -> CompletionCallAction {
8012                if let CompletionCallEvent { .. } = event {
8013                    CompletionCallAction::patch(
8014                        RequestPatch::new().history([Message::user(SENTINEL)]),
8015                    )
8016                } else {
8017                    CompletionCallAction::continue_run()
8018                }
8019            }
8020        }
8021
8022        fn request_has_sentinel(req: &crate::completion::CompletionRequest) -> bool {
8023            req.chat_history.iter().any(|m| match m {
8024                Message::User { content } => content
8025                    .iter()
8026                    .any(|c| matches!(c, UserContent::Text(text) if text.text.contains(SENTINEL))),
8027                _ => false,
8028            })
8029        }
8030
8031        fn messages_have_sentinel(messages: &[Message]) -> bool {
8032            messages.iter().any(|m| match m {
8033                Message::User { content } => content
8034                    .iter()
8035                    .any(|c| matches!(c, UserContent::Text(text) if text.text.contains(SENTINEL))),
8036                _ => false,
8037            })
8038        }
8039
8040        let blocking_model = MockCompletionModel::from_turns([MockTurn::text("done")]);
8041        let blocking_probe = blocking_model.clone();
8042        let blocking = AgentBuilder::new(blocking_model)
8043            .add_hook(HistoryOverrideHook)
8044            .build()
8045            .runner("real prompt")
8046            .run()
8047            .await
8048            .expect("blocking run should succeed");
8049        assert!(
8050            request_has_sentinel(blocking_probe.requests().first().expect("one request")),
8051            "the overridden history reaches the provider"
8052        );
8053        assert!(
8054            !messages_have_sentinel(blocking.messages.as_deref().unwrap_or_default()),
8055            "the persisted transcript is untouched by the per-turn history override"
8056        );
8057
8058        let streaming_model =
8059            MockCompletionModel::from_stream_turns([one_text_stream_turn("done")]);
8060        let streaming_probe = streaming_model.clone();
8061        let stream = AgentBuilder::new(streaming_model)
8062            .add_hook(HistoryOverrideHook)
8063            .build()
8064            .runner("real prompt")
8065            .stream()
8066            .await;
8067        let final_response = drive_to_final_response(stream).await;
8068        assert!(
8069            request_has_sentinel(streaming_probe.requests().first().expect("one request")),
8070            "the overridden history reaches the provider on the streaming surface too"
8071        );
8072        assert!(
8073            !messages_have_sentinel(final_response.messages().expect("history")),
8074            "the persisted transcript is untouched by the per-turn history override on \
8075             the streaming surface too"
8076        );
8077    }
8078
8079    /// `ModelTurnFinished` fires exactly once per accepted turn on both surfaces,
8080    /// including a streamed tool-only turn that fires no `StreamResponseFinish`.
8081    #[tokio::test]
8082    async fn model_turn_finished_fires_once_per_accepted_turn_including_tool_only() {
8083        let blocking_hook = RecordingHook::default();
8084        AgentBuilder::new(blocking_model())
8085            .tool(MockAddTool)
8086            .add_hook(blocking_hook.clone())
8087            .build()
8088            .runner("add 2 and 3")
8089            .max_turns(3)
8090            .run()
8091            .await
8092            .expect("blocking run should succeed");
8093        assert_eq!(
8094            blocking_hook.count(StepEventKind::ModelTurnFinished),
8095            2,
8096            "one ModelTurnFinished per accepted turn (tool turn + text turn)"
8097        );
8098
8099        let streaming_hook = RecordingHook::default();
8100        let mut stream = AgentBuilder::new(streaming_model())
8101            .tool(MockAddTool)
8102            .add_hook(streaming_hook.clone())
8103            .build()
8104            .runner("add 2 and 3")
8105            .max_turns(3)
8106            .stream()
8107            .await;
8108        while let Some(item) = stream.next().await {
8109            let _ = item.map_err(|err| panic!("stream item errored: {err}"));
8110        }
8111        assert_eq!(
8112            streaming_hook.count(StepEventKind::ModelTurnFinished),
8113            2,
8114            "ModelTurnFinished fires once per turn on the streaming surface too"
8115        );
8116        // The tool-only first turn streams no assistant text, so only the second
8117        // (text) turn fires StreamResponseFinish — proving ModelTurnFinished
8118        // covers the gap.
8119        assert_eq!(
8120            streaming_hook.count(StepEventKind::StreamResponseFinish),
8121            1,
8122            "the tool-only turn fires no StreamResponseFinish"
8123        );
8124    }
8125
8126    #[tokio::test]
8127    async fn reasoning_only_turn_does_not_gain_stream_response_finish() {
8128        let hook = RecordingHook::default();
8129        let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([[
8130            MockStreamEvent::reasoning("think"),
8131            MockStreamEvent::final_response_with_total_tokens(0),
8132        ]]))
8133        .add_hook(hook.clone())
8134        .build()
8135        .runner("reason")
8136        .stream()
8137        .await;
8138        while let Some(item) = stream.next().await {
8139            item.expect("reasoning-only stream item");
8140        }
8141
8142        assert_eq!(
8143            hook.count(StepEventKind::StreamResponseFinish),
8144            0,
8145            "reasoning-only turns must not fire StreamResponseFinish"
8146        );
8147        assert_eq!(
8148            hook.count(StepEventKind::ModelTurnFinished),
8149            1,
8150            "the accepted reasoning-only turn still fires ModelTurnFinished"
8151        );
8152    }
8153
8154    /// Records the content kinds of the first turn's `ModelTurnFinished`.
8155    #[derive(Clone, Default)]
8156    struct CaptureFirstTurnContent {
8157        kinds: Arc<Mutex<Option<Vec<&'static str>>>>,
8158    }
8159
8160    impl AgentHook for CaptureFirstTurnContent {
8161        async fn on_model_turn_finished(
8162            &self,
8163            _ctx: &HookContext,
8164            event: ModelTurnFinished<'_>,
8165        ) -> ModelTurnAction {
8166            if let ModelTurnFinished { turn, content, .. } = event
8167                && turn == 1
8168            {
8169                let kinds = content
8170                    .iter()
8171                    .map(|c| match c {
8172                        AssistantContent::Reasoning(_) => "reasoning",
8173                        AssistantContent::Text(_) => "text",
8174                        AssistantContent::ToolCall(_) => "tool_call",
8175                        _ => "other",
8176                    })
8177                    .collect();
8178                *self.kinds.lock().expect("kinds") = Some(kinds);
8179            }
8180            ModelTurnAction::continue_run()
8181        }
8182    }
8183
8184    /// On the streaming surface, `ModelTurnFinished.content` carries the
8185    /// **canonical** committed content from `StreamedTurn::finish` (reasoning →
8186    /// text → tool calls), not the raw `stream.choice` aggregate. The turn streams
8187    /// reasoning, then a tool call, then text (a non-canonical emission order), so
8188    /// a raw-choice implementation would surface `reasoning, tool_call, text` —
8189    /// the canonical event instead reports `reasoning, text, tool_call`.
8190    #[tokio::test]
8191    async fn streaming_model_turn_finished_carries_canonical_committed_content() {
8192        let model = MockCompletionModel::from_stream_turns([
8193            vec![
8194                MockStreamEvent::reasoning("think"),
8195                MockStreamEvent::tool_call("tc1", "add", json!({"x": 2, "y": 3})),
8196                MockStreamEvent::text("answer"),
8197                MockStreamEvent::final_response_with_total_tokens(0),
8198            ],
8199            vec![
8200                MockStreamEvent::text("done"),
8201                MockStreamEvent::final_response_with_total_tokens(0),
8202            ],
8203        ]);
8204        let hook = CaptureFirstTurnContent::default();
8205        let stream = AgentBuilder::new(model)
8206            .tool(MockAddTool)
8207            .add_hook(hook.clone())
8208            .build()
8209            .runner("go")
8210            .max_turns(3)
8211            .stream()
8212            .await;
8213        let _ = drive_to_final_response(stream).await;
8214
8215        assert_eq!(
8216            hook.kinds.lock().expect("kinds").clone(),
8217            Some(vec!["reasoning", "text", "tool_call"]),
8218            "ModelTurnFinished carries the canonical reasoning->text->tool ordering \
8219             from StreamedTurn::finish, not the raw stream.choice emission order"
8220        );
8221    }
8222
8223    /// `ToolCallAction::Rewrite` and `ToolResultAction::Rewrite` chain across hooks: a later hook observes
8224    /// (and further rewrites) the value produced by earlier hooks.
8225    #[tokio::test]
8226    async fn chained_rewrites_compose_across_hooks() {
8227        /// Sets one key of the tool arguments, preserving the rest.
8228        struct SetArg {
8229            key: &'static str,
8230            value: i64,
8231        }
8232        impl AgentHook for SetArg {
8233            async fn on_tool_call(
8234                &self,
8235                _ctx: &HookContext,
8236                event: ToolCall<'_>,
8237            ) -> ToolCallAction {
8238                if let ToolCall { args, .. } = event {
8239                    let mut parsed: serde_json::Value =
8240                        serde_json::from_str(args).unwrap_or_else(|_| json!({}));
8241                    parsed[self.key] = json!(self.value);
8242                    ToolCallAction::rewrite(parsed)
8243                } else {
8244                    ToolCallAction::run()
8245                }
8246            }
8247        }
8248
8249        /// Wraps the tool result in `label(...)`.
8250        struct WrapResult(&'static str);
8251        impl AgentHook for WrapResult {
8252            async fn on_tool_result(
8253                &self,
8254                _ctx: &HookContext,
8255                event: ToolResultEvent<'_>,
8256            ) -> ToolResultAction {
8257                if let ToolResultEvent { presentation, .. } = event {
8258                    ToolResultAction::rewrite(format!("{}({})", self.0, presentation.render()))
8259                } else {
8260                    ToolResultAction::keep()
8261                }
8262            }
8263        }
8264
8265        // The model asks add(2, 3). SetArg{y:40} then SetArg{x:100} chain, so the
8266        // tool runs with (100, 40) = 140 — proving arg rewrites compose. Then
8267        // WrapResult "A" and "B" chain, and a trailing recorder observes the fully
8268        // chained result "B(A(140))".
8269        let recorder = RecordingHook::default();
8270        let blocking = AgentBuilder::new(blocking_model())
8271            .tool(MockAddTool)
8272            .add_hook(SetArg {
8273                key: "y",
8274                value: 40,
8275            })
8276            .add_hook(SetArg {
8277                key: "x",
8278                value: 100,
8279            })
8280            .add_hook(WrapResult("A"))
8281            .add_hook(WrapResult("B"))
8282            .add_hook(recorder.clone())
8283            .build()
8284            .runner("add 2 and 3")
8285            .max_turns(3)
8286            .run()
8287            .await
8288            .expect("blocking run should succeed");
8289        assert_eq!(blocking.output, "the answer is 5");
8290        assert_eq!(
8291            recorder.tool_results(),
8292            vec!["B(A(140))".to_string()],
8293            "arg rewrites compose (100+40=140) and result rewrites nest B(A(...))"
8294        );
8295
8296        // Same on the streaming surface.
8297        let stream_recorder = RecordingHook::default();
8298        let mut stream = AgentBuilder::new(streaming_model())
8299            .tool(MockAddTool)
8300            .add_hook(SetArg {
8301                key: "y",
8302                value: 40,
8303            })
8304            .add_hook(SetArg {
8305                key: "x",
8306                value: 100,
8307            })
8308            .add_hook(WrapResult("A"))
8309            .add_hook(WrapResult("B"))
8310            .add_hook(stream_recorder.clone())
8311            .build()
8312            .runner("add 2 and 3")
8313            .max_turns(3)
8314            .stream()
8315            .await;
8316        while let Some(item) = stream.next().await {
8317            let _ = item.map_err(|err| panic!("stream item errored: {err}"));
8318        }
8319        assert_eq!(
8320            stream_recorder.tool_results(),
8321            vec!["B(A(140))".to_string()],
8322            "chained rewrites compose identically on the streaming surface"
8323        );
8324    }
8325
8326    #[derive(serde::Deserialize, schemars::JsonSchema)]
8327    #[allow(dead_code)]
8328    struct Answer {
8329        answer: String,
8330    }
8331
8332    /// A real tool whose name equals the default synthetic output-tool name
8333    /// (`final_result`). Used to prove a per-turn `active_tools` filter cannot
8334    /// make the picked output-tool name collide with it.
8335    struct FinalResultTool;
8336
8337    impl Tool for FinalResultTool {
8338        const NAME: &'static str = "final_result";
8339        type Error = MockToolError;
8340        type Args = serde_json::Value;
8341        type Output = String;
8342
8343        fn description(&self) -> String {
8344            "A real tool sharing the default output-tool name".to_string()
8345        }
8346
8347        fn parameters(&self) -> serde_json::Value {
8348            json!({ "type": "object", "properties": {} })
8349        }
8350
8351        async fn call(
8352            &self,
8353            _context: &mut ToolContext,
8354            _args: Self::Args,
8355        ) -> Result<Self::Output, Self::Error> {
8356            Ok("real final_result output".to_string())
8357        }
8358    }
8359
8360    /// Returns no retrieved tool on the first search, then the colliding real
8361    /// `final_result` tool on later searches.
8362    #[derive(Default)]
8363    struct LateFinalResultIndex {
8364        searches: AtomicU32,
8365    }
8366
8367    impl VectorStoreIndex for LateFinalResultIndex {
8368        type Filter = Filter<serde_json::Value>;
8369
8370        async fn top_n<T: for<'a> Deserialize<'a> + WasmCompatSend>(
8371            &self,
8372            _req: VectorSearchRequest,
8373        ) -> Result<Vec<(f64, String, T)>, VectorStoreError> {
8374            Ok(Vec::new())
8375        }
8376
8377        async fn top_n_ids(
8378            &self,
8379            _req: VectorSearchRequest,
8380        ) -> Result<Vec<(f64, String)>, VectorStoreError> {
8381            if self.searches.fetch_add(1, SeqCst) == 0 {
8382                Ok(Vec::new())
8383            } else {
8384                Ok(vec![(1.0, "final_result".to_string())])
8385            }
8386        }
8387    }
8388
8389    /// Registers a real `final_result` tool after the first model turn, once the
8390    /// run has already reserved that name for structured output. An optional
8391    /// second-turn patch lets tests exercise filtering and tool-choice changes
8392    /// without changing the collision source.
8393    #[derive(Clone)]
8394    struct RegisterLateFinalResultTool {
8395        handle: ToolServerHandle,
8396        second_turn_patch: Option<RequestPatch>,
8397    }
8398
8399    impl AgentHook for RegisterLateFinalResultTool {
8400        async fn on_model_turn_finished(
8401            &self,
8402            ctx: &HookContext,
8403            _event: ModelTurnFinished<'_>,
8404        ) -> ModelTurnAction {
8405            if ctx.turn() == 1 {
8406                self.handle.add_tool(FinalResultTool).await;
8407            }
8408
8409            ModelTurnAction::continue_run()
8410        }
8411
8412        async fn on_completion_call(
8413            &self,
8414            ctx: &HookContext,
8415            _event: CompletionCallEvent<'_>,
8416        ) -> CompletionCallAction {
8417            if ctx.turn() == 2
8418                && let Some(patch) = &self.second_turn_patch
8419            {
8420                return CompletionCallAction::patch(patch.clone());
8421            }
8422
8423            CompletionCallAction::continue_run()
8424        }
8425    }
8426
8427    fn assert_structured_output_collision_error(message: &str) {
8428        assert!(
8429            message.contains("final_result"),
8430            "error should name the conflicting tool: {message}"
8431        );
8432        assert!(
8433            message.contains("structured-output") && message.contains("reserved"),
8434            "error should explain the structured-output reservation: {message}"
8435        );
8436        assert!(
8437            message.contains("rename or remove"),
8438            "error should provide an actionable resolution: {message}"
8439        );
8440    }
8441
8442    /// An initially effective real `final_result` keeps normal dispatch while
8443    /// the synthetic structured-output tool is advertised under a unique name.
8444    #[tokio::test]
8445    async fn initial_output_tool_collision_uses_a_unique_synthetic_name() {
8446        let model = MockCompletionModel::from_turns([
8447            MockTurn::tool_call("real", "final_result", json!({})),
8448            MockTurn::tool_call("output", "final_result_1", json!({ "answer": "done" })),
8449        ]);
8450        let probe = model.clone();
8451        let response = AgentBuilder::new(model)
8452            .tool(FinalResultTool)
8453            .output_schema::<Answer>()
8454            .output_mode(OutputMode::Tool)
8455            .build()
8456            .runner("go")
8457            .max_turns(2)
8458            .run()
8459            .await
8460            .expect("the real tool should dispatch before the unique output tool finalizes");
8461
8462        assert!(response.output.contains("done"));
8463        let requests = probe.requests();
8464        assert_eq!(
8465            requests.len(),
8466            2,
8467            "real-tool dispatch must continue to a second model turn"
8468        );
8469        let tool_names = requests[0]
8470            .tools
8471            .iter()
8472            .map(|tool| tool.name.as_str())
8473            .collect::<Vec<_>>();
8474        assert_eq!(tool_names.len(), 2);
8475        for expected in ["final_result", "final_result_1"] {
8476            assert_eq!(
8477                tool_names.iter().filter(|name| **name == expected).count(),
8478                1,
8479                "the first request should advertise `{expected}` exactly once: {tool_names:?}"
8480            );
8481        }
8482
8483        assert!(
8484            requests[1].chat_history.iter().any(|message| matches!(
8485                message,
8486                Message::User { content }
8487                    if content.iter().any(|item| matches!(
8488                        item,
8489                        UserContent::ToolResult(result)
8490                            if result.id == "real"
8491                                && result.content.iter().any(|content| matches!(
8492                                    content,
8493                                    rig_core::message::ToolResultContent::Text(text)
8494                                        if text.text == "real final_result output"
8495                                ))
8496                    ))
8497            )),
8498            "the real `final_result` call must execute normally and its result must reach the follow-up request"
8499        );
8500    }
8501
8502    /// Once Tool output mode has committed a name, a real tool registered under
8503    /// that name must fail the next request locally for every tool-choice shape.
8504    /// Otherwise the provider receives duplicate definitions and the real call
8505    /// is intercepted as final output.
8506    #[tokio::test]
8507    async fn late_output_tool_collision_fails_before_blocking_provider_for_all_choices() {
8508        let cases = [
8509            ("inherited", None),
8510            (
8511                "required",
8512                Some(RequestPatch::new().tool_choice(ToolChoice::Required)),
8513            ),
8514            (
8515                "none",
8516                Some(RequestPatch::new().tool_choice(ToolChoice::None)),
8517            ),
8518            (
8519                "specific",
8520                Some(RequestPatch::new().tool_choice(ToolChoice::Specific {
8521                    function_names: vec!["final_result".to_string()],
8522                })),
8523            ),
8524        ];
8525
8526        for (case, second_turn_patch) in cases {
8527            let handle = ToolServer::new().tool(MockAddTool).run();
8528            let model = MockCompletionModel::from_turns([
8529                MockTurn::tool_call("add-1", "add", json!({ "x": 1, "y": 2 })),
8530                MockTurn::tool_call(
8531                    "shadowed",
8532                    "final_result",
8533                    json!({ "answer": "wrongly finalized" }),
8534                ),
8535            ]);
8536            let probe = model.clone();
8537            let err = AgentBuilder::new(model)
8538                .tool_server_handle(handle.clone())
8539                .output_schema::<Answer>()
8540                .output_mode(OutputMode::Tool)
8541                .add_hook(RegisterLateFinalResultTool {
8542                    handle,
8543                    second_turn_patch,
8544                })
8545                .build()
8546                .runner("go")
8547                .max_turns(3)
8548                .run()
8549                .await
8550                .unwrap_err();
8551
8552            assert!(
8553                matches!(
8554                    &err,
8555                    PromptError::CompletionError(CompletionError::RequestError(_))
8556                ),
8557                "{case}: expected a local completion request error, got {err:?}"
8558            );
8559            assert_eq!(
8560                probe.request_count(),
8561                1,
8562                "{case}: the colliding second request must not reach the provider"
8563            );
8564            assert_structured_output_collision_error(&err.to_string());
8565        }
8566    }
8567
8568    /// The streaming surface uses the same pre-provider collision check as the
8569    /// blocking surface and terminates without starting a second model stream.
8570    #[tokio::test]
8571    async fn late_output_tool_collision_fails_before_streaming_provider() {
8572        let handle = ToolServer::new().tool(MockAddTool).run();
8573        let model = MockCompletionModel::from_stream_turns([
8574            vec![
8575                MockStreamEvent::tool_call("add-1", "add", json!({ "x": 1, "y": 2 })),
8576                MockStreamEvent::final_response_with_total_tokens(0),
8577            ],
8578            vec![
8579                MockStreamEvent::tool_call(
8580                    "shadowed",
8581                    "final_result",
8582                    json!({ "answer": "wrongly finalized" }),
8583                ),
8584                MockStreamEvent::final_response_with_total_tokens(0),
8585            ],
8586        ]);
8587        let probe = model.clone();
8588        let mut stream = AgentBuilder::new(model)
8589            .tool_server_handle(handle.clone())
8590            .output_schema::<Answer>()
8591            .output_mode(OutputMode::Tool)
8592            .add_hook(RegisterLateFinalResultTool {
8593                handle,
8594                second_turn_patch: None,
8595            })
8596            .build()
8597            .runner("go")
8598            .max_turns(3)
8599            .stream()
8600            .await;
8601
8602        let mut collisions = Vec::new();
8603        let mut saw_final_response = false;
8604        while let Some(item) = stream.next().await {
8605            match item {
8606                Err(err) => collisions.push(err),
8607                Ok(MultiTurnStreamItem::FinalResponse(_)) => saw_final_response = true,
8608                Ok(_) => {}
8609            }
8610        }
8611        assert_eq!(
8612            collisions.len(),
8613            1,
8614            "the stream should terminate with exactly one collision error"
8615        );
8616        assert!(
8617            !saw_final_response,
8618            "a collision error must terminate the stream without a final response"
8619        );
8620        let err = collisions.pop().expect("one collision error was asserted");
8621
8622        assert!(
8623            matches!(
8624                &err,
8625                StreamingError::Completion(CompletionError::RequestError(_))
8626            ),
8627            "expected a local streaming completion request error, got {err:?}"
8628        );
8629        assert_eq!(
8630            probe.request_count(),
8631            1,
8632            "the colliding second stream must not reach the provider"
8633        );
8634        assert_structured_output_collision_error(&err.to_string());
8635    }
8636
8637    /// A late colliding tool is harmless while `active_tools` filters it out,
8638    /// but the run must fail as soon as the non-sticky filter lifts and the real
8639    /// tool becomes effective again.
8640    #[tokio::test]
8641    async fn late_output_tool_collision_is_checked_after_active_tools_filtering() {
8642        let handle = ToolServer::new().tool(MockAddTool).run();
8643        let model = MockCompletionModel::from_turns([
8644            MockTurn::tool_call("add-1", "add", json!({ "x": 1, "y": 2 })),
8645            MockTurn::tool_call("add-2", "add", json!({ "x": 3, "y": 4 })),
8646            MockTurn::tool_call(
8647                "shadowed",
8648                "final_result",
8649                json!({ "answer": "wrongly finalized" }),
8650            ),
8651        ]);
8652        let probe = model.clone();
8653        let err = AgentBuilder::new(model)
8654            .tool_server_handle(handle.clone())
8655            .output_schema::<Answer>()
8656            .output_mode(OutputMode::Tool)
8657            .add_hook(RegisterLateFinalResultTool {
8658                handle,
8659                second_turn_patch: Some(RequestPatch::new().active_tools(["add"])),
8660            })
8661            .build()
8662            .runner("go")
8663            .max_turns(4)
8664            .run()
8665            .await
8666            .expect_err("the exposed third-turn collision should fail locally");
8667
8668        assert_eq!(
8669            probe.request_count(),
8670            2,
8671            "the filtered second turn may run, but the exposed third turn may not"
8672        );
8673        let requests = probe.requests();
8674        let second_turn_names = requests[1]
8675            .tools
8676            .iter()
8677            .map(|tool| tool.name.as_str())
8678            .collect::<Vec<_>>();
8679        assert_eq!(second_turn_names.len(), 2);
8680        for expected in ["add", "final_result"] {
8681            assert_eq!(
8682                second_turn_names
8683                    .iter()
8684                    .filter(|name| **name == expected)
8685                    .count(),
8686                1,
8687                "the second request should advertise `{expected}` exactly once: \
8688                 {second_turn_names:?}"
8689            );
8690        }
8691        assert_structured_output_collision_error(&err.to_string());
8692    }
8693
8694    /// Dynamic retrieval shares the same effective per-turn collision check as
8695    /// mutable registration: a name absent on turn one may not shadow the
8696    /// already-reserved output tool when retrieval selects it on turn two.
8697    #[tokio::test]
8698    async fn retrieved_output_tool_collision_fails_before_provider_request() {
8699        let mut retrieved_tools = ToolSet::default();
8700        retrieved_tools.add_tool(FinalResultTool);
8701        let handle = ToolServer::new()
8702            .tool(MockAddTool)
8703            .retrieved_tools(1, LateFinalResultIndex::default(), retrieved_tools)
8704            .run();
8705        let model = MockCompletionModel::from_turns([
8706            MockTurn::tool_call("add-1", "add", json!({ "x": 1, "y": 2 })),
8707            MockTurn::tool_call(
8708                "shadowed",
8709                "final_result",
8710                json!({ "answer": "wrongly finalized" }),
8711            ),
8712        ]);
8713        let probe = model.clone();
8714        let err = AgentBuilder::new(model)
8715            .tool_server_handle(handle)
8716            .output_schema::<Answer>()
8717            .output_mode(OutputMode::Tool)
8718            .build()
8719            .runner("go")
8720            .max_turns(3)
8721            .run()
8722            .await
8723            .expect_err("the retrieved second-turn collision should fail locally");
8724
8725        assert!(matches!(
8726            &err,
8727            PromptError::CompletionError(CompletionError::RequestError(_))
8728        ));
8729        assert_eq!(
8730            probe.request_count(),
8731            1,
8732            "the colliding retrieved tool must prevent the second provider request"
8733        );
8734        assert_structured_output_collision_error(&err.to_string());
8735    }
8736
8737    /// Narrows the advertised tools to `add` for the turn, filtering out the real
8738    /// `final_result` tool.
8739    struct ActiveToolsAddOnly;
8740
8741    impl AgentHook for ActiveToolsAddOnly {
8742        async fn on_completion_call(
8743            &self,
8744            _ctx: &HookContext,
8745            event: CompletionCallEvent<'_>,
8746        ) -> CompletionCallAction {
8747            if let CompletionCallEvent { .. } = event {
8748                CompletionCallAction::patch(RequestPatch::new().active_tools(["add"]))
8749            } else {
8750                CompletionCallAction::continue_run()
8751            }
8752        }
8753    }
8754
8755    /// Regression guard: a per-turn `active_tools` allow-list that filters out a
8756    /// real tool whose name equals the default synthetic output-tool name must not
8757    /// let the picked output-tool name collide with that (filtered) real tool. The
8758    /// name is pinned for the whole run, so picking it against the FULL advertised
8759    /// set — not just this turn's narrowed executable set — keeps it collision-safe
8760    /// once the filter lifts on a later turn. With the bug, the output tool would
8761    /// be named `final_result` (picked against the narrowed `{add}`), colliding
8762    /// with the real `final_result` whenever the filter is gone.
8763    #[tokio::test]
8764    async fn active_tools_filter_does_not_let_output_tool_collide_with_a_filtered_real_tool() {
8765        // The model finalizes by calling the (correctly-picked) output tool, so a
8766        // run on the fixed code completes cleanly in a single turn. Asserting the
8767        // run succeeds also exercises finalization: the model's call to
8768        // `final_result_1` must be intercepted as the output tool, so this fails if
8769        // the picked name and the intercept name ever drift apart.
8770        let model = MockCompletionModel::from_turns([MockTurn::tool_call(
8771            "out1",
8772            "final_result_1",
8773            json!({ "answer": "done" }),
8774        )]);
8775        let probe = model.clone();
8776        let response = AgentBuilder::new(model)
8777            .tool(MockAddTool)
8778            .tool(FinalResultTool)
8779            .output_schema::<Answer>()
8780            .output_mode(OutputMode::Tool)
8781            .add_hook(ActiveToolsAddOnly)
8782            .build()
8783            .runner("go")
8784            .max_turns(2)
8785            .run()
8786            .await
8787            .expect("run should finalize via the picked output tool `final_result_1`");
8788        assert!(
8789            response.output.contains("done"),
8790            "the intercepted output-tool call should produce the structured result, \
8791             got {:?}",
8792            response.output
8793        );
8794
8795        let requests = probe.requests();
8796        assert!(
8797            !requests.is_empty(),
8798            "the first model request should be captured"
8799        );
8800        let tool_names: Vec<&str> = requests[0].tools.iter().map(|t| t.name.as_str()).collect();
8801        assert!(
8802            tool_names.contains(&"add"),
8803            "active_tools keeps `add` advertised, saw {tool_names:?}"
8804        );
8805        assert!(
8806            tool_names.contains(&"final_result_1"),
8807            "the synthetic output tool must avoid the filtered real `final_result` name, \
8808             saw {tool_names:?}"
8809        );
8810        assert!(
8811            !tool_names.contains(&"final_result"),
8812            "the real `final_result` is filtered out and the output tool must not reuse \
8813             its name, saw {tool_names:?}"
8814        );
8815    }
8816
8817    /// Captures whether any `ModelTurnFinished.content` carried a tool call named
8818    /// `final_result` — the model-emitted structured-output output-tool call.
8819    #[derive(Clone, Default)]
8820    struct CaptureOutputToolInModelTurn {
8821        saw_output_tool_call: Arc<Mutex<bool>>,
8822    }
8823
8824    impl AgentHook for CaptureOutputToolInModelTurn {
8825        async fn on_model_turn_finished(
8826            &self,
8827            _ctx: &HookContext,
8828            event: ModelTurnFinished<'_>,
8829        ) -> ModelTurnAction {
8830            if let ModelTurnFinished { content, .. } = event
8831                && content.iter().any(|c| {
8832                    matches!(c, AssistantContent::ToolCall(tc) if tc.function.name == "final_result")
8833                })
8834            {
8835                *self.saw_output_tool_call.lock().expect("lock") = true;
8836            }
8837            ModelTurnAction::continue_run()
8838        }
8839    }
8840
8841    /// `ModelTurnFinished.content` carries the **model-emitted** content — including
8842    /// a structured-output Tool-mode output-tool call — on both surfaces, even though
8843    /// the run persists that turn as assistant text (the structured output) with the
8844    /// tool call dropped. Guards the documented `content` contract: it is the model's
8845    /// committed turn content, not the finalized/persisted content, in Tool mode.
8846    #[tokio::test]
8847    async fn model_turn_finished_content_carries_output_tool_call_in_tool_mode() {
8848        // Blocking surface.
8849        let hook = CaptureOutputToolInModelTurn::default();
8850        let response = AgentBuilder::new(MockCompletionModel::from_turns([MockTurn::tool_call(
8851            "out1",
8852            "final_result",
8853            json!({ "answer": "done" }),
8854        )]))
8855        .output_schema::<Answer>()
8856        .output_mode(OutputMode::Tool)
8857        .add_hook(hook.clone())
8858        .build()
8859        .runner("go")
8860        .max_turns(2)
8861        .run()
8862        .await
8863        .expect("run should finalize via the output tool");
8864        assert!(
8865            *hook.saw_output_tool_call.lock().expect("lock"),
8866            "ModelTurnFinished.content must carry the model-emitted output-tool call (blocking)"
8867        );
8868        assert!(
8869            response.output.contains("done"),
8870            "the run finalizes with the structured output, not the raw tool call: {:?}",
8871            response.output
8872        );
8873
8874        // Streaming surface — same content contract.
8875        let s_hook = CaptureOutputToolInModelTurn::default();
8876        let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([vec![
8877            MockStreamEvent::tool_call("out1", "final_result", json!({ "answer": "done" })),
8878            MockStreamEvent::final_response_with_total_tokens(0),
8879        ]]))
8880        .output_schema::<Answer>()
8881        .output_mode(OutputMode::Tool)
8882        .add_hook(s_hook.clone())
8883        .build()
8884        .runner("go")
8885        .max_turns(2)
8886        .stream()
8887        .await;
8888        while stream.next().await.is_some() {}
8889        assert!(
8890            *s_hook.saw_output_tool_call.lock().expect("lock"),
8891            "ModelTurnFinished.content must carry the model-emitted output-tool call (streaming)"
8892        );
8893    }
8894
8895    /// A structured-output Tool-mode output-tool call finalizes the run directly, so
8896    /// on the streaming surface it is **not** re-emitted as a complete
8897    /// `StreamAssistantItem(StreamedAssistantContent::ToolCall)` item (it bypasses
8898    /// `drive_tool_calls`); its structured result is surfaced in the final `PromptResponse`.
8899    /// Guards the narrowed `StreamAssistantItem` contract.
8900    #[tokio::test]
8901    async fn output_tool_finalization_emits_no_complete_tool_call_stream_item() {
8902        let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([vec![
8903            MockStreamEvent::tool_call("out1", "final_result", json!({ "answer": "done" })),
8904            MockStreamEvent::final_response_with_total_tokens(0),
8905        ]]))
8906        .output_schema::<Answer>()
8907        .output_mode(OutputMode::Tool)
8908        .build()
8909        .runner("go")
8910        .max_turns(2)
8911        .stream()
8912        .await;
8913
8914        let mut saw_complete_output_tool_call = false;
8915        let mut final_has_output = false;
8916        while let Some(item) = stream.next().await {
8917            match item.expect("stream item") {
8918                MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::ToolCall {
8919                    tool_call,
8920                    ..
8921                }) if tool_call.function.name == "final_result" => {
8922                    saw_complete_output_tool_call = true;
8923                }
8924                MultiTurnStreamItem::FinalResponse(res) => {
8925                    final_has_output = res.output().contains("done");
8926                }
8927                _ => {}
8928            }
8929        }
8930        assert!(
8931            !saw_complete_output_tool_call,
8932            "the output-tool call finalizes the run, so no complete \
8933             StreamAssistantItem::ToolCall item must be emitted for it"
8934        );
8935        assert!(
8936            final_has_output,
8937            "the structured output must be surfaced via the FinalResponse"
8938        );
8939    }
8940
8941    // -----------------------------------------------------------------------
8942    // Human-in-the-loop (HITL): one hook gates each tool call behind a human
8943    // decision, mapping approve/deny/edit/abort onto the event-specific actions
8944    // (cont / skip / rewrite_args / terminate). The runnable interactive
8945    // version lives in `examples/agent_with_human_in_the_loop`.
8946    // -----------------------------------------------------------------------
8947
8948    /// A human reviewer's decision for a pending tool call.
8949    enum Decision {
8950        /// Run the tool as the model requested.
8951        Approve,
8952        /// Don't run the tool; feed `reason` back to the model as the result.
8953        Deny(&'static str),
8954        /// Run the tool with these arguments instead of the model's.
8955        Edit(serde_json::Value),
8956        /// Abort the whole run with this reason.
8957        Abort(&'static str),
8958    }
8959
8960    /// Simulates a human reviewer by popping a scripted decision per `ToolCall`
8961    /// and mapping it to the matching event-specific action. A real reviewer would `.await`
8962    /// interactive input here (the hook is async) rather than read a queue.
8963    #[derive(Clone)]
8964    struct HumanApprovalHook {
8965        decisions: Arc<Mutex<std::collections::VecDeque<Decision>>>,
8966        reviewed: Arc<Mutex<Vec<String>>>,
8967    }
8968
8969    impl HumanApprovalHook {
8970        fn new(decisions: impl IntoIterator<Item = Decision>) -> Self {
8971            Self {
8972                decisions: Arc::new(Mutex::new(decisions.into_iter().collect())),
8973                reviewed: Arc::new(Mutex::new(Vec::new())),
8974            }
8975        }
8976
8977        /// `"name(args)"` for each call presented for review, in order.
8978        fn reviewed(&self) -> Vec<String> {
8979            self.reviewed.lock().unwrap().clone()
8980        }
8981    }
8982
8983    impl AgentHook for HumanApprovalHook {
8984        async fn on_tool_call(&self, _ctx: &HookContext, event: ToolCall<'_>) -> ToolCallAction {
8985            let ToolCall {
8986                tool_name, args, ..
8987            } = event
8988            else {
8989                return ToolCallAction::run();
8990            };
8991            self.reviewed
8992                .lock()
8993                .unwrap()
8994                .push(format!("{tool_name}({args})"));
8995            let decision = self.decisions.lock().unwrap().pop_front();
8996            match decision {
8997                Some(Decision::Approve) => ToolCallAction::run(),
8998                Some(Decision::Deny(reason)) => ToolCallAction::skip(reason),
8999                Some(Decision::Edit(args)) => ToolCallAction::rewrite(args),
9000                Some(Decision::Abort(reason)) => ToolCallAction::stop(reason),
9001                // Fail closed if the script is exhausted (it shouldn't be) — deny
9002                // rather than silently approve, matching the example's contract.
9003                None => ToolCallAction::skip("denied: no scripted decision (fail-closed)"),
9004            }
9005        }
9006    }
9007
9008    /// A HITL hook that approves the first tool call, denies the second, and
9009    /// edits the third's arguments behaves identically under `run()` and
9010    /// `stream()`: approved/edited tools execute (and the edit takes effect),
9011    /// the denied tool runs nothing while its reason reaches the model, and the
9012    /// model-visible history is identical across drivers (compared structurally).
9013    #[tokio::test]
9014    async fn human_in_the_loop_approve_deny_edit_parity_across_run_and_stream() {
9015        // One turn issues three tool calls; the reviewer decides each differently.
9016        let turns = [
9017            ScriptedTurn::ToolCalls(vec![
9018                add_call("tc1", 2, 3),   // approve -> runs, 2 + 3 = 5
9019                add_call("tc2", 10, 20), // deny    -> skipped; model sees the reason
9020                add_call("tc3", 1, 1),   // edit    -> runs 1 + 100 = 101, not 1 + 1 = 2
9021            ]),
9022            ScriptedTurn::Text("done"),
9023        ];
9024        let denial = "denied by reviewer: amount too large";
9025        let decisions = || {
9026            vec![
9027                Decision::Approve,
9028                Decision::Deny(denial),
9029                Decision::Edit(json!({"x": 1, "y": 100})),
9030            ]
9031        };
9032
9033        let blocking_model =
9034            MockCompletionModel::from_turns(turns.iter().map(ScriptedTurn::as_blocking_turn));
9035        let blocking_recorder = RecordingHook::default();
9036        let blocking_approver = HumanApprovalHook::new(decisions());
9037        let blocking = AgentBuilder::new(blocking_model)
9038            .tool(MockAddTool)
9039            .build()
9040            .runner("carry out the plan")
9041            .max_turns(3)
9042            .add_hook(blocking_recorder.clone())
9043            .add_hook(blocking_approver.clone())
9044            .run()
9045            .await
9046            .expect("blocking HITL run should succeed");
9047
9048        let streaming_model = MockCompletionModel::from_stream_turns(
9049            turns
9050                .iter()
9051                .map(|turn| turn.as_stream_events(StreamShape::Complete)),
9052        );
9053        let streaming_recorder = RecordingHook::default();
9054        let streaming_approver = HumanApprovalHook::new(decisions());
9055        let mut stream = AgentBuilder::new(streaming_model)
9056            .tool(MockAddTool)
9057            .build()
9058            .runner("carry out the plan")
9059            .max_turns(3)
9060            .add_hook(streaming_recorder.clone())
9061            .add_hook(streaming_approver.clone())
9062            .stream()
9063            .await;
9064        let mut final_response = None;
9065        while let Some(item) = stream.next().await {
9066            if let Ok(MultiTurnStreamItem::FinalResponse(resp)) =
9067                item.map_err(|err| panic!("stream item errored: {err}"))
9068            {
9069                final_response = Some(resp);
9070            }
9071        }
9072        let final_response = final_response.expect("stream should yield a final response");
9073
9074        // Approved (5) and edited (101) tools executed, in call order; the denied
9075        // call executed nothing but now fires a ToolResult carrying its verbatim
9076        // denial reason (structured `Skipped` outcome) — identically on both
9077        // drivers.
9078        assert_eq!(
9079            blocking_recorder.tool_results(),
9080            vec![
9081                "5".to_string(),
9082                "denied by reviewer: amount too large".to_string(),
9083                "101".to_string()
9084            ]
9085        );
9086        assert_eq!(
9087            blocking_recorder.tool_results(),
9088            streaming_recorder.tool_results()
9089        );
9090
9091        // The denied call (10 + 20) never executed, so its result 30 is absent —
9092        // the denial reason stands in its place, ruling out deny being silently
9093        // treated as approve.
9094        assert!(
9095            !blocking_recorder.tool_results().contains(&"30".to_string()),
9096            "the denied call must not have executed"
9097        );
9098
9099        // The reviewer was consulted for all three calls, in order, identically per
9100        // driver — pinning each decision to its call (approve=2+3, deny=10+20,
9101        // edit=the third).
9102        let reviewed = blocking_approver.reviewed();
9103        assert_eq!(reviewed.len(), 3);
9104        assert_eq!(reviewed, streaming_approver.reviewed());
9105        assert!(
9106            reviewed[0].contains('2') && reviewed[0].contains('3'),
9107            "first reviewed call should be add(2, 3): {reviewed:?}"
9108        );
9109        assert!(
9110            reviewed[1].contains("10") && reviewed[1].contains("20"),
9111            "the denied (second) call should be add(10, 20): {reviewed:?}"
9112        );
9113
9114        assert_eq!(blocking.output, "done");
9115        assert_eq!(final_response.output(), blocking.output);
9116        assert_eq!(
9117            blocking_recorder.shared_events(),
9118            streaming_recorder.shared_events()
9119        );
9120
9121        // Model-visible history is identical across drivers (compared structurally
9122        // as serde_json::Value) and carries the denial reason and the edited result
9123        // 101 (not the model's 1 + 1 = 2).
9124        let blocking_messages = blocking.messages.expect("blocking messages");
9125        let streaming_messages = final_response
9126            .messages()
9127            .expect("streaming history")
9128            .to_vec();
9129        assert_eq!(
9130            serde_json::to_value(&blocking_messages).expect("serialize blocking"),
9131            serde_json::to_value(&streaming_messages).expect("serialize streaming"),
9132        );
9133        assert!(
9134            tool_result_text_in_history(&blocking_messages, denial),
9135            "the denial reason must be the denied call's tool result in the history"
9136        );
9137        assert!(
9138            tool_result_json_in_history(&blocking_messages, &json!(101)),
9139            "the edited call must have executed with the rewritten arguments"
9140        );
9141    }
9142
9143    /// A HITL hook that aborts a tool call (`Decision::Abort` -> `ToolCallAction::stop`)
9144    /// stops the run and surfaces the reason as a `PromptCancelled` error — on both
9145    /// the blocking and streaming drivers.
9146    #[tokio::test]
9147    async fn human_in_the_loop_abort_terminates_the_run() {
9148        let turns = [
9149            ScriptedTurn::ToolCalls(vec![add_call("tc1", 2, 3)]),
9150            ScriptedTurn::Text("unreachable"),
9151        ];
9152        const ABORT_REASON: &str = "aborted by the human reviewer";
9153
9154        // Blocking driver: the run resolves to a PromptCancelled error.
9155        let blocking_model =
9156            MockCompletionModel::from_turns(turns.iter().map(ScriptedTurn::as_blocking_turn));
9157        let err = AgentBuilder::new(blocking_model)
9158            .tool(MockAddTool)
9159            .build()
9160            .runner("do the sensitive thing")
9161            .max_turns(3)
9162            .add_hook(HumanApprovalHook::new([Decision::Abort(ABORT_REASON)]))
9163            .run()
9164            .await
9165            .expect_err("an aborted tool call should terminate the blocking run");
9166        assert!(
9167            format!("{err}").contains(ABORT_REASON),
9168            "the abort reason should surface in the blocking error, got: {err}"
9169        );
9170
9171        // Streaming driver: the stream yields an error carrying the same reason and
9172        // never reaches the "unreachable" final text.
9173        let streaming_model = MockCompletionModel::from_stream_turns(
9174            turns
9175                .iter()
9176                .map(|turn| turn.as_stream_events(StreamShape::Complete)),
9177        );
9178        let mut stream = AgentBuilder::new(streaming_model)
9179            .tool(MockAddTool)
9180            .build()
9181            .runner("do the sensitive thing")
9182            .max_turns(3)
9183            .add_hook(HumanApprovalHook::new([Decision::Abort(ABORT_REASON)]))
9184            .stream()
9185            .await;
9186        let mut stream_error = None;
9187        while let Some(item) = stream.next().await {
9188            match item {
9189                Err(err) => stream_error = Some(format!("{err}")),
9190                Ok(MultiTurnStreamItem::FinalResponse(resp)) => {
9191                    panic!("aborted stream must not finalize, got: {}", resp.output())
9192                }
9193                Ok(_) => {}
9194            }
9195        }
9196        let stream_error = stream_error.expect("an aborted tool call should error the stream");
9197        assert!(
9198            stream_error.contains(ABORT_REASON),
9199            "the abort reason should surface in the streaming error, got: {stream_error}"
9200        );
9201    }
9202
9203    /// A non-interactive *policy* HITL hook: auto-approve an allow-list, deny
9204    /// everything else (fail-closed), and cache each decision so a repeated tool
9205    /// is not re-evaluated ("sticky", like the OpenAI Agents SDK's
9206    /// `always_approve`). Backs `examples/agent_with_approval_policy`.
9207    #[derive(Clone)]
9208    struct PolicyHook {
9209        auto_approve: std::collections::HashSet<&'static str>,
9210        /// Tool names the policy actually evaluated (cache misses), in order.
9211        evaluated: Arc<Mutex<Vec<String>>>,
9212        /// Sticky cache of prior decisions, keyed by tool name.
9213        cache: Arc<Mutex<std::collections::HashMap<String, bool>>>,
9214    }
9215
9216    impl PolicyHook {
9217        fn new(auto_approve: impl IntoIterator<Item = &'static str>) -> Self {
9218            Self {
9219                auto_approve: auto_approve.into_iter().collect(),
9220                evaluated: Arc::new(Mutex::new(Vec::new())),
9221                cache: Arc::new(Mutex::new(std::collections::HashMap::new())),
9222            }
9223        }
9224
9225        fn evaluated(&self) -> Vec<String> {
9226            self.evaluated.lock().unwrap().clone()
9227        }
9228    }
9229
9230    impl AgentHook for PolicyHook {
9231        async fn on_tool_call(&self, _ctx: &HookContext, event: ToolCall<'_>) -> ToolCallAction {
9232            let ToolCall { tool_name, .. } = event else {
9233                return ToolCallAction::run();
9234            };
9235            let cached = self.cache.lock().unwrap().get(tool_name).copied();
9236            let approved = match cached {
9237                Some(decision) => decision, // sticky: reuse without re-evaluating
9238                None => {
9239                    self.evaluated.lock().unwrap().push(tool_name.to_string());
9240                    let decision = self.auto_approve.contains(tool_name);
9241                    self.cache
9242                        .lock()
9243                        .unwrap()
9244                        .insert(tool_name.to_string(), decision);
9245                    decision
9246                }
9247            };
9248            if approved {
9249                ToolCallAction::run()
9250            } else {
9251                ToolCallAction::skip(format!("denied by policy: `{tool_name}` not allowed"))
9252            }
9253        }
9254    }
9255
9256    /// The policy hook auto-approves `add` and denies `subtract`, and its decision
9257    /// is sticky: a second `add` call reuses the cached approval instead of being
9258    /// re-evaluated. The denied call never runs and its reason reaches the model.
9259    #[tokio::test]
9260    async fn approval_policy_allow_list_with_sticky_decisions() {
9261        // One turn issues three calls: add, subtract (denied), add again (sticky).
9262        let turns = [
9263            ScriptedTurn::ToolCalls(vec![
9264                add_call("c1", 2, 3),
9265                ScriptedToolCall {
9266                    id: "c2",
9267                    name: "subtract",
9268                    args: json!({ "x": 10, "y": 4 }),
9269                },
9270                add_call("c3", 2, 3),
9271            ]),
9272            ScriptedTurn::Text("done"),
9273        ];
9274
9275        let model =
9276            MockCompletionModel::from_turns(turns.iter().map(ScriptedTurn::as_blocking_turn));
9277        let recorder = RecordingHook::default();
9278        let policy = PolicyHook::new(["add"]);
9279        let out = AgentBuilder::new(model)
9280            .tool(MockAddTool)
9281            .tool(MockSubtractTool)
9282            .build()
9283            .runner("go")
9284            .max_turns(3)
9285            .add_hook(recorder.clone())
9286            .add_hook(policy.clone())
9287            .run()
9288            .await
9289            .expect("policy run should succeed");
9290
9291        assert_eq!(out.output, "done");
9292        // `add` ran twice (auto-approved, then sticky-reused); `subtract` was denied
9293        // and executed nothing, but its denial reason now surfaces as a ToolResult
9294        // (structured `Skipped` outcome) between the two `add` results.
9295        assert_eq!(
9296            recorder.tool_results(),
9297            vec![
9298                "5".to_string(),
9299                "denied by policy: `subtract` not allowed".to_string(),
9300                "5".to_string()
9301            ]
9302        );
9303        // The policy evaluated each distinct tool once; the second `add` reused the
9304        // cached decision rather than being re-evaluated.
9305        assert_eq!(
9306            policy.evaluated(),
9307            vec!["add".to_string(), "subtract".to_string()]
9308        );
9309        let messages = out.messages.expect("messages");
9310        assert!(
9311            tool_result_text_in_history(&messages, "denied by policy: `subtract` not allowed"),
9312            "the policy denial reason must reach the model as the subtract tool result"
9313        );
9314    }
9315
9316    static NEXT_RESPONSE_RETRY_HOOK_ID: AtomicU64 = AtomicU64::new(1);
9317
9318    #[derive(Clone, Default)]
9319    struct ResponseRetryAttempts(HashMap<u64, usize>);
9320
9321    #[derive(Clone)]
9322    enum TestRetryMode {
9323        Repeat,
9324        Feedback(&'static str),
9325    }
9326
9327    /// A policy-owned retry budget. The framework only enforces `max_turns`;
9328    /// this hook stores its narrower limit in the run-scoped scratchpad.
9329    #[derive(Clone)]
9330    struct BoundedResponseRetry {
9331        id: u64,
9332        rejected_text: &'static str,
9333        max_retries: usize,
9334        mode: TestRetryMode,
9335    }
9336
9337    #[derive(Clone, Default)]
9338    struct StatefulCompletionPatch {
9339        calls: Arc<AtomicU32>,
9340    }
9341
9342    impl StatefulCompletionPatch {
9343        fn calls(&self) -> u32 {
9344            self.calls.load(SeqCst)
9345        }
9346    }
9347
9348    impl AgentHook for StatefulCompletionPatch {
9349        async fn on_completion_call(
9350            &self,
9351            _ctx: &HookContext,
9352            _event: crate::agent::CompletionCallEvent<'_>,
9353        ) -> CompletionCallAction {
9354            let call = self.calls.fetch_add(1, SeqCst);
9355            CompletionCallAction::patch(RequestPatch::new().temperature(if call == 0 {
9356                0.1
9357            } else {
9358                0.9
9359            }))
9360        }
9361    }
9362
9363    impl BoundedResponseRetry {
9364        fn new(rejected_text: &'static str, max_retries: usize, mode: TestRetryMode) -> Self {
9365            Self {
9366                id: NEXT_RESPONSE_RETRY_HOOK_ID.fetch_add(1, SeqCst),
9367                rejected_text,
9368                max_retries,
9369                mode,
9370            }
9371        }
9372    }
9373
9374    impl AgentHook for BoundedResponseRetry {
9375        async fn on_model_turn_finished(
9376            &self,
9377            ctx: &HookContext,
9378            event: ModelTurnFinished<'_>,
9379        ) -> ModelTurnAction {
9380            let rejected = event.content.iter().any(
9381                |content| matches!(content, AssistantContent::Text(text) if text.text == self.rejected_text),
9382            );
9383            if !rejected {
9384                return ModelTurnAction::continue_run();
9385            }
9386
9387            let attempt = ctx
9388                .scratchpad()
9389                .update::<ResponseRetryAttempts, _>(|attempts| {
9390                    let attempt = attempts.0.entry(self.id).or_default();
9391                    *attempt += 1;
9392                    *attempt
9393                });
9394            if attempt > self.max_retries {
9395                return ModelTurnAction::stop(format!(
9396                    "response retry limit ({}) exceeded",
9397                    self.max_retries
9398                ));
9399            }
9400
9401            match self.mode {
9402                TestRetryMode::Repeat => ModelTurnAction::repeat(),
9403                TestRetryMode::Feedback(feedback) => ModelTurnAction::retry_with_feedback(feedback),
9404            }
9405        }
9406    }
9407
9408    fn retry_usage(input_tokens: u64, output_tokens: u64) -> Usage {
9409        Usage {
9410            input_tokens,
9411            output_tokens,
9412            total_tokens: input_tokens + output_tokens,
9413            ..Usage::new()
9414        }
9415    }
9416
9417    #[tokio::test]
9418    async fn blocking_model_turn_repeat_preserves_prompt_history_with_fresh_preparation() {
9419        let first_usage = retry_usage(10, 3);
9420        let second_usage = retry_usage(7, 2);
9421        let completion_patch = StatefulCompletionPatch::default();
9422        let model = MockCompletionModel::from_turns([
9423            MockTurn::text("rejected").with_usage(first_usage),
9424            MockTurn::text("accepted").with_usage(second_usage),
9425        ]);
9426        let response = AgentBuilder::new(model.clone())
9427            .add_hook(completion_patch.clone())
9428            .add_hook(BoundedResponseRetry::new(
9429                "rejected",
9430                1,
9431                TestRetryMode::Repeat,
9432            ))
9433            .build()
9434            .runner("question")
9435            .max_turns(2)
9436            .run()
9437            .await
9438            .expect("repeat should recover");
9439
9440        assert_eq!(response.output, "accepted");
9441        assert_eq!(response.usage, first_usage + second_usage);
9442        assert_eq!(response.completion_calls.len(), 2);
9443        let messages = response.messages.expect("response messages");
9444        assert_eq!(
9445            messages,
9446            vec![Message::user("question"), Message::assistant("accepted")]
9447        );
9448
9449        let requests = model.requests();
9450        assert_eq!(requests.len(), 2);
9451        let first = requests[0].chat_history.iter().cloned().collect::<Vec<_>>();
9452        let second = requests[1].chat_history.iter().cloned().collect::<Vec<_>>();
9453        assert_eq!(first, vec![Message::user("question")]);
9454        assert_eq!(
9455            second, first,
9456            "Repeat must preserve the prompt and preceding history"
9457        );
9458        assert_eq!(requests[0].temperature, Some(0.1));
9459        assert_eq!(requests[1].temperature, Some(0.9));
9460        assert_eq!(completion_patch.calls(), 2);
9461    }
9462
9463    #[tokio::test]
9464    async fn blocking_model_turn_feedback_preserves_rejected_response() {
9465        let model = MockCompletionModel::from_turns([
9466            MockTurn::text("rejected"),
9467            MockTurn::text("accepted"),
9468        ]);
9469        let response = AgentBuilder::new(model.clone())
9470            .add_hook(BoundedResponseRetry::new(
9471                "rejected",
9472                1,
9473                TestRetryMode::Feedback("try another approach"),
9474            ))
9475            .build()
9476            .runner("question")
9477            .max_turns(2)
9478            .run()
9479            .await
9480            .expect("feedback retry should recover");
9481
9482        assert_eq!(response.output, "accepted");
9483        assert_eq!(
9484            response.messages.expect("response messages"),
9485            vec![
9486                Message::user("question"),
9487                Message::assistant("rejected"),
9488                Message::user("try another approach"),
9489                Message::assistant("accepted"),
9490            ]
9491        );
9492        let second_request = &model.requests()[1];
9493        assert_eq!(
9494            second_request
9495                .chat_history
9496                .iter()
9497                .cloned()
9498                .collect::<Vec<_>>(),
9499            vec![
9500                Message::user("question"),
9501                Message::assistant("rejected"),
9502                Message::user("try another approach"),
9503            ]
9504        );
9505    }
9506
9507    #[tokio::test]
9508    async fn blocking_empty_feedback_retry_omits_empty_assistant_history() {
9509        let first_usage = retry_usage(5, 1);
9510        let second_usage = retry_usage(7, 2);
9511        let model = MockCompletionModel::from_turns([
9512            MockTurn::text("").with_usage(first_usage),
9513            MockTurn::text("accepted").with_usage(second_usage),
9514        ]);
9515        let response = AgentBuilder::new(model.clone())
9516            .add_hook(BoundedResponseRetry::new(
9517                "",
9518                1,
9519                TestRetryMode::Feedback("provide an answer"),
9520            ))
9521            .build()
9522            .runner("question")
9523            .max_turns(2)
9524            .run()
9525            .await
9526            .expect("feedback retry should recover from an empty turn");
9527
9528        assert_eq!(response.output, "accepted");
9529        assert_eq!(response.usage, first_usage + second_usage);
9530        assert_eq!(response.completion_calls.len(), 2);
9531        assert_eq!(
9532            response.messages.expect("response messages"),
9533            vec![
9534                Message::user("question"),
9535                Message::user("provide an answer"),
9536                Message::assistant("accepted"),
9537            ]
9538        );
9539        assert_eq!(
9540            model.requests()[1]
9541                .chat_history
9542                .iter()
9543                .cloned()
9544                .collect::<Vec<_>>(),
9545            vec![
9546                Message::user("question"),
9547                Message::user("provide an answer"),
9548            ],
9549            "the retry request must not contain an empty assistant message"
9550        );
9551    }
9552
9553    #[tokio::test]
9554    async fn streaming_model_turn_retry_marks_rollback_and_matches_blocking_accounting() {
9555        let first_usage = retry_usage(10, 3);
9556        let second_usage = retry_usage(7, 2);
9557        let model = MockCompletionModel::from_stream_turns([
9558            [
9559                MockStreamEvent::text("rejected"),
9560                MockStreamEvent::final_response(first_usage),
9561            ],
9562            [
9563                MockStreamEvent::text("accepted"),
9564                MockStreamEvent::final_response(second_usage),
9565            ],
9566        ]);
9567        let mut stream = AgentBuilder::new(model.clone())
9568            .add_hook(BoundedResponseRetry::new(
9569                "rejected",
9570                1,
9571                TestRetryMode::Repeat,
9572            ))
9573            .build()
9574            .runner("question")
9575            .max_turns(2)
9576            .stream()
9577            .await;
9578
9579        let mut retries = Vec::new();
9580        let mut provider_finals = 0;
9581        let mut completion_calls = 0;
9582        let mut final_response = None;
9583        while let Some(item) = stream.next().await {
9584            match item.expect("stream item") {
9585                MultiTurnStreamItem::ModelTurnRetried { turn } => retries.push(turn),
9586                MultiTurnStreamItem::CompletionCall(_) => completion_calls += 1,
9587                MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Final(_)) => {
9588                    provider_finals += 1
9589                }
9590                MultiTurnStreamItem::FinalResponse(response) => final_response = Some(response),
9591                _ => {}
9592            }
9593        }
9594
9595        assert_eq!(retries, vec![1]);
9596        assert_eq!(
9597            provider_finals, 1,
9598            "the rejected provider final is suppressed"
9599        );
9600        assert_eq!(completion_calls, 2);
9601        let response = final_response.expect("run final response");
9602        assert_eq!(response.output, "accepted");
9603        assert_eq!(response.usage, first_usage + second_usage);
9604        assert_eq!(response.completion_calls.len(), 2);
9605        assert_eq!(
9606            response.messages.expect("response messages"),
9607            vec![Message::user("question"), Message::assistant("accepted")]
9608        );
9609        assert_eq!(model.requests().len(), 2);
9610    }
9611
9612    #[tokio::test]
9613    async fn streaming_feedback_retry_matches_blocking_history_and_usage() {
9614        let first_usage = retry_usage(5, 2);
9615        let second_usage = retry_usage(8, 4);
9616        let blocking = AgentBuilder::new(MockCompletionModel::from_turns([
9617            MockTurn::text("rejected").with_usage(first_usage),
9618            MockTurn::text("accepted").with_usage(second_usage),
9619        ]))
9620        .add_hook(BoundedResponseRetry::new(
9621            "rejected",
9622            1,
9623            TestRetryMode::Feedback("correct the answer"),
9624        ))
9625        .build()
9626        .runner("question")
9627        .max_turns(2)
9628        .run()
9629        .await
9630        .expect("blocking feedback retry");
9631
9632        let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([
9633            [
9634                MockStreamEvent::text("rejected"),
9635                MockStreamEvent::final_response(first_usage),
9636            ],
9637            [
9638                MockStreamEvent::text("accepted"),
9639                MockStreamEvent::final_response(second_usage),
9640            ],
9641        ]))
9642        .add_hook(BoundedResponseRetry::new(
9643            "rejected",
9644            1,
9645            TestRetryMode::Feedback("correct the answer"),
9646        ))
9647        .build()
9648        .runner("question")
9649        .max_turns(2)
9650        .stream()
9651        .await;
9652        let mut saw_retry = false;
9653        let mut streaming = None;
9654        while let Some(item) = stream.next().await {
9655            match item.expect("stream item") {
9656                MultiTurnStreamItem::ModelTurnRetried { turn: 1 } => saw_retry = true,
9657                MultiTurnStreamItem::FinalResponse(response) => streaming = Some(response),
9658                _ => {}
9659            }
9660        }
9661
9662        let streaming = streaming.expect("streaming final response");
9663        assert!(saw_retry);
9664        assert_eq!(streaming.output, blocking.output);
9665        assert_eq!(streaming.usage, blocking.usage);
9666        assert_eq!(streaming.completion_calls, blocking.completion_calls);
9667        assert_eq!(
9668            serde_json::to_value(streaming.messages).expect("streaming history"),
9669            serde_json::to_value(blocking.messages).expect("blocking history")
9670        );
9671    }
9672
9673    #[tokio::test]
9674    async fn streaming_empty_feedback_retry_omits_empty_assistant_history() {
9675        let first_usage = retry_usage(5, 1);
9676        let second_usage = retry_usage(7, 2);
9677        let model = MockCompletionModel::from_stream_turns([
9678            [
9679                MockStreamEvent::text(""),
9680                MockStreamEvent::final_response(first_usage),
9681            ],
9682            [
9683                MockStreamEvent::text("accepted"),
9684                MockStreamEvent::final_response(second_usage),
9685            ],
9686        ]);
9687        let mut stream = AgentBuilder::new(model.clone())
9688            .add_hook(BoundedResponseRetry::new(
9689                "",
9690                1,
9691                TestRetryMode::Feedback("provide an answer"),
9692            ))
9693            .build()
9694            .runner("question")
9695            .max_turns(2)
9696            .stream()
9697            .await;
9698
9699        let mut retries = Vec::new();
9700        let mut provider_finals = 0;
9701        let mut completion_calls = 0;
9702        let mut final_response = None;
9703        while let Some(item) = stream.next().await {
9704            match item.expect("stream item") {
9705                MultiTurnStreamItem::ModelTurnRetried { turn } => retries.push(turn),
9706                MultiTurnStreamItem::CompletionCall(_) => completion_calls += 1,
9707                MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Final(_)) => {
9708                    provider_finals += 1;
9709                }
9710                MultiTurnStreamItem::FinalResponse(response) => final_response = Some(response),
9711                _ => {}
9712            }
9713        }
9714
9715        assert_eq!(retries, vec![1]);
9716        assert_eq!(provider_finals, 1, "the rejected final is suppressed");
9717        assert_eq!(completion_calls, 2);
9718        let response = final_response.expect("run final response");
9719        assert_eq!(response.output, "accepted");
9720        assert_eq!(response.usage, first_usage + second_usage);
9721        assert_eq!(response.completion_calls.len(), 2);
9722        assert_eq!(
9723            response.messages.expect("response messages"),
9724            vec![
9725                Message::user("question"),
9726                Message::user("provide an answer"),
9727                Message::assistant("accepted"),
9728            ]
9729        );
9730        assert_eq!(
9731            model.requests()[1]
9732                .chat_history
9733                .iter()
9734                .cloned()
9735                .collect::<Vec<_>>(),
9736            vec![
9737                Message::user("question"),
9738                Message::user("provide an answer"),
9739            ],
9740            "the retry request must not contain an empty assistant message"
9741        );
9742    }
9743
9744    #[tokio::test]
9745    async fn response_retry_preserves_model_turn_hook_order_across_surfaces() {
9746        let blocking_events = RecordingHook::default();
9747        AgentBuilder::new(MockCompletionModel::from_turns([
9748            MockTurn::text("rejected"),
9749            MockTurn::text("accepted"),
9750        ]))
9751        .add_hook(blocking_events.clone())
9752        .add_hook(BoundedResponseRetry::new(
9753            "rejected",
9754            1,
9755            TestRetryMode::Repeat,
9756        ))
9757        .build()
9758        .runner("question")
9759        .max_turns(2)
9760        .run()
9761        .await
9762        .expect("blocking retry");
9763
9764        let streaming_events = RecordingHook::default();
9765        let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([
9766            [
9767                MockStreamEvent::text("rejected"),
9768                MockStreamEvent::final_response_with_default_usage(),
9769            ],
9770            [
9771                MockStreamEvent::text("accepted"),
9772                MockStreamEvent::final_response_with_default_usage(),
9773            ],
9774        ]))
9775        .add_hook(streaming_events.clone())
9776        .add_hook(BoundedResponseRetry::new(
9777            "rejected",
9778            1,
9779            TestRetryMode::Repeat,
9780        ))
9781        .build()
9782        .runner("question")
9783        .max_turns(2)
9784        .stream()
9785        .await;
9786        while let Some(item) = stream.next().await {
9787            item.expect("streaming retry item");
9788        }
9789
9790        let shared_order = |events: &RecordingHook| {
9791            events
9792                .events
9793                .lock()
9794                .expect("events")
9795                .iter()
9796                .copied()
9797                .filter(|event| {
9798                    matches!(
9799                        event,
9800                        StepEventKind::CompletionCall | StepEventKind::ModelTurnFinished
9801                    )
9802                })
9803                .collect::<Vec<_>>()
9804        };
9805        let expected = vec![
9806            StepEventKind::CompletionCall,
9807            StepEventKind::ModelTurnFinished,
9808            StepEventKind::CompletionCall,
9809            StepEventKind::ModelTurnFinished,
9810        ];
9811        assert_eq!(shared_order(&blocking_events), expected);
9812        assert_eq!(shared_order(&streaming_events), expected);
9813
9814        let blocking_order = blocking_events.events.lock().expect("events").clone();
9815        assert_eq!(
9816            blocking_order,
9817            vec![
9818                StepEventKind::CompletionCall,
9819                StepEventKind::CompletionResponse,
9820                StepEventKind::ModelTurnFinished,
9821                StepEventKind::CompletionCall,
9822                StepEventKind::CompletionResponse,
9823                StepEventKind::ModelTurnFinished,
9824            ]
9825        );
9826        let streaming_order = streaming_events.events.lock().expect("events").clone();
9827        assert_eq!(
9828            streaming_order,
9829            vec![
9830                StepEventKind::CompletionCall,
9831                StepEventKind::TextDelta,
9832                StepEventKind::StreamResponseFinish,
9833                StepEventKind::ModelTurnFinished,
9834                StepEventKind::CompletionCall,
9835                StepEventKind::TextDelta,
9836                StepEventKind::StreamResponseFinish,
9837                StepEventKind::ModelTurnFinished,
9838            ]
9839        );
9840    }
9841
9842    #[tokio::test]
9843    async fn streaming_model_turn_retry_respects_max_turns() {
9844        let model = MockCompletionModel::from_stream_turns([[
9845            MockStreamEvent::text("rejected"),
9846            MockStreamEvent::final_response_with_default_usage(),
9847        ]]);
9848        let mut stream = AgentBuilder::new(model)
9849            .add_hook(BoundedResponseRetry::new(
9850                "rejected",
9851                1,
9852                TestRetryMode::Repeat,
9853            ))
9854            .build()
9855            .runner("question")
9856            .max_turns(1)
9857            .stream()
9858            .await;
9859
9860        let mut saw_rollback = false;
9861        let mut error = None;
9862        while let Some(item) = stream.next().await {
9863            match item {
9864                Ok(MultiTurnStreamItem::ModelTurnRetried { turn: 1 }) => saw_rollback = true,
9865                Ok(_) => {}
9866                Err(err) => error = Some(err),
9867            }
9868        }
9869        assert!(saw_rollback);
9870        assert!(matches!(
9871            error,
9872            Some(StreamingError::Prompt(error))
9873                if matches!(error.as_ref(), PromptError::MaxTurnsError { max_turns: 1, .. })
9874        ));
9875    }
9876
9877    struct AlwaysRepeatModelTurn;
9878
9879    impl AgentHook for AlwaysRepeatModelTurn {
9880        async fn on_model_turn_finished(
9881            &self,
9882            _ctx: &HookContext,
9883            _event: ModelTurnFinished<'_>,
9884        ) -> ModelTurnAction {
9885            ModelTurnAction::repeat()
9886        }
9887    }
9888
9889    #[tokio::test]
9890    async fn model_turn_retry_rejects_tool_turn_before_tool_hooks_or_execution() {
9891        let recorder = RecordingHook::default();
9892        let executions = Arc::new(AtomicU32::new(0));
9893        let err = AgentBuilder::new(MockCompletionModel::from_turns([MockTurn::tool_call(
9894            "tc1",
9895            "add",
9896            json!({"x": 1, "y": 2}),
9897        )]))
9898        .tool(CountingAddTool {
9899            calls: executions.clone(),
9900        })
9901        .add_hook(recorder.clone())
9902        .add_hook(AlwaysRepeatModelTurn)
9903        .build()
9904        .runner("add")
9905        .max_turns(2)
9906        .run()
9907        .await
9908        .expect_err("tool-bearing retry must fail closed");
9909
9910        let PromptError::PromptCancelled {
9911            chat_history,
9912            reason,
9913        } = err
9914        else {
9915            panic!("tool-bearing retry should return PromptCancelled");
9916        };
9917        assert!(reason.contains("tool-bearing model turns"));
9918        assert!(reason.contains("tool-call hooks"));
9919        assert_eq!(chat_history, vec![Message::user("add")]);
9920        assert_eq!(recorder.count(StepEventKind::ToolCall), 0);
9921        assert_eq!(recorder.count(StepEventKind::ToolResult), 0);
9922        assert_eq!(executions.load(SeqCst), 0);
9923    }
9924
9925    #[tokio::test]
9926    async fn streaming_model_turn_retry_rejects_tool_turn_without_committed_execution() {
9927        let recorder = RecordingHook::default();
9928        let executions = Arc::new(AtomicU32::new(0));
9929        let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([[
9930            MockStreamEvent::tool_call_name_delta("tc1", "ic1", "add"),
9931            MockStreamEvent::tool_call_arguments_delta("tc1", "ic1", r#"{"x":1,"y":2}"#),
9932            MockStreamEvent::tool_call("tc1", "add", json!({"x": 1, "y": 2})),
9933            MockStreamEvent::final_response_with_default_usage(),
9934        ]]))
9935        .tool(CountingAddTool {
9936            calls: executions.clone(),
9937        })
9938        .add_hook(recorder.clone())
9939        .add_hook(AlwaysRepeatModelTurn)
9940        .build()
9941        .runner("add")
9942        .max_turns(2)
9943        .stream()
9944        .await;
9945
9946        let mut execution_commits = 0;
9947        let mut tool_results = 0;
9948        let mut provider_finals = 0;
9949        let mut agent_finals = 0;
9950        let mut retry_markers = 0;
9951        let mut error = None;
9952        while let Some(item) = stream.next().await {
9953            match item {
9954                Ok(MultiTurnStreamItem::ToolExecutionCommitted { .. }) => execution_commits += 1,
9955                Ok(MultiTurnStreamItem::StreamUserItem(StreamedUserContent::ToolResult {
9956                    ..
9957                })) => tool_results += 1,
9958                Ok(MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Final(
9959                    _,
9960                ))) => provider_finals += 1,
9961                Ok(MultiTurnStreamItem::FinalResponse(_)) => agent_finals += 1,
9962                Ok(MultiTurnStreamItem::ModelTurnRetried { .. }) => retry_markers += 1,
9963                Ok(_) => {}
9964                Err(err) => error = Some(err),
9965            }
9966        }
9967
9968        let Some(StreamingError::Prompt(error)) = error else {
9969            panic!("tool-bearing streaming retry should return PromptCancelled");
9970        };
9971        let PromptError::PromptCancelled {
9972            chat_history,
9973            reason,
9974        } = error.as_ref()
9975        else {
9976            panic!("tool-bearing streaming retry should return PromptCancelled");
9977        };
9978        assert!(reason.contains("tool-bearing model turns"));
9979        assert!(reason.contains("tool-call hooks"));
9980        assert_eq!(chat_history, &[Message::user("add")]);
9981        assert_eq!(execution_commits, 0);
9982        assert_eq!(tool_results, 0);
9983        assert_eq!(provider_finals, 0);
9984        assert_eq!(agent_finals, 0);
9985        assert_eq!(retry_markers, 0);
9986        assert_eq!(recorder.count(StepEventKind::ToolCall), 0);
9987        assert_eq!(recorder.count(StepEventKind::ToolResult), 0);
9988        assert_eq!(executions.load(SeqCst), 0);
9989    }
9990
9991    #[derive(Clone)]
9992    struct BarrierResponseRetry {
9993        inner: BoundedResponseRetry,
9994        barrier: Arc<Barrier>,
9995    }
9996
9997    impl AgentHook for BarrierResponseRetry {
9998        async fn on_model_turn_finished(
9999            &self,
10000            ctx: &HookContext,
10001            event: ModelTurnFinished<'_>,
10002        ) -> ModelTurnAction {
10003            let rejected = event.content.iter().any(
10004                |content| matches!(content, AssistantContent::Text(text) if text.text == "rejected"),
10005            );
10006            if rejected {
10007                self.barrier.wait().await;
10008            }
10009            self.inner.on_model_turn_finished(ctx, event).await
10010        }
10011    }
10012
10013    #[tokio::test]
10014    async fn concurrent_runs_of_same_agent_have_independent_retry_budgets() {
10015        let hook = BarrierResponseRetry {
10016            inner: BoundedResponseRetry::new("rejected", 1, TestRetryMode::Repeat),
10017            barrier: Arc::new(Barrier::new(2)),
10018        };
10019        let agent = AgentBuilder::new(MockCompletionModel::from_turns([
10020            MockTurn::text("rejected"),
10021            MockTurn::text("rejected"),
10022            MockTurn::text("accepted one"),
10023            MockTurn::text("accepted two"),
10024        ]))
10025        .add_hook(hook)
10026        .build();
10027
10028        let first = agent.runner("first").max_turns(2).run();
10029        let second = agent.runner("second").max_turns(2).run();
10030        let (first, second) = tokio::join!(first, second);
10031        let first = first.expect("first run");
10032        let second = second.expect("second run");
10033
10034        let outputs = std::collections::HashSet::from([first.output, second.output]);
10035        assert_eq!(
10036            outputs,
10037            std::collections::HashSet::from([
10038                "accepted one".to_string(),
10039                "accepted two".to_string(),
10040            ])
10041        );
10042        assert_eq!(first.completion_calls.len(), 2);
10043        assert_eq!(second.completion_calls.len(), 2);
10044    }
10045
10046    #[tokio::test]
10047    async fn retry_scratchpad_state_is_isolated_by_run_and_hook_instance() {
10048        let shared_hook = BoundedResponseRetry::new("rejected", 1, TestRetryMode::Repeat);
10049        let first_ctx = HookContext::new(false, None);
10050        let second_ctx = HookContext::new(false, None);
10051        let content = OneOrMany::one(AssistantContent::text("rejected"));
10052        let first_event = ModelTurnFinished {
10053            turn: 1,
10054            content: &content,
10055            usage: Usage::new(),
10056        };
10057        let second_event = first_event;
10058
10059        let (first, second) = tokio::join!(
10060            shared_hook.on_model_turn_finished(&first_ctx, first_event),
10061            shared_hook.on_model_turn_finished(&second_ctx, second_event),
10062        );
10063        assert!(matches!(first, ModelTurnAction::Retry(_)));
10064        assert!(matches!(second, ModelTurnAction::Retry(_)));
10065        assert!(matches!(
10066            shared_hook
10067                .on_model_turn_finished(&first_ctx, first_event)
10068                .await,
10069            ModelTurnAction::Stop(_)
10070        ));
10071
10072        let same_run_ctx = HookContext::new(false, None);
10073        let first_hook = BoundedResponseRetry::new("first", 1, TestRetryMode::Repeat);
10074        let second_hook = BoundedResponseRetry::new("second", 1, TestRetryMode::Repeat);
10075        let first_content = OneOrMany::one(AssistantContent::text("first"));
10076        let second_content = OneOrMany::one(AssistantContent::text("second"));
10077        let first_action = first_hook
10078            .on_model_turn_finished(
10079                &same_run_ctx,
10080                ModelTurnFinished {
10081                    turn: 1,
10082                    content: &first_content,
10083                    usage: Usage::new(),
10084                },
10085            )
10086            .await;
10087        let second_action = second_hook
10088            .on_model_turn_finished(
10089                &same_run_ctx,
10090                ModelTurnFinished {
10091                    turn: 2,
10092                    content: &second_content,
10093                    usage: Usage::new(),
10094                },
10095            )
10096            .await;
10097        assert!(matches!(first_action, ModelTurnAction::Retry(_)));
10098        assert!(matches!(second_action, ModelTurnAction::Retry(_)));
10099    }
10100
10101    #[derive(Clone)]
10102    struct FixedModelTurnAction {
10103        action: ModelTurnAction,
10104        calls: Arc<AtomicU32>,
10105    }
10106
10107    impl AgentHook for FixedModelTurnAction {
10108        async fn on_model_turn_finished(
10109            &self,
10110            _ctx: &HookContext,
10111            _event: ModelTurnFinished<'_>,
10112        ) -> ModelTurnAction {
10113            self.calls.fetch_add(1, SeqCst);
10114            self.action.clone()
10115        }
10116    }
10117
10118    #[tokio::test]
10119    async fn model_turn_action_short_circuits_flat_and_nested_hook_stacks() {
10120        let content = OneOrMany::one(AssistantContent::text("response"));
10121        let event = ModelTurnFinished {
10122            turn: 1,
10123            content: &content,
10124            usage: Usage::new(),
10125        };
10126        let ctx = HookContext::new(false, None);
10127
10128        let first_calls = Arc::new(AtomicU32::new(0));
10129        let retry_calls = Arc::new(AtomicU32::new(0));
10130        let skipped_calls = Arc::new(AtomicU32::new(0));
10131        let mut flat = HookStack::new();
10132        flat.push(FixedModelTurnAction {
10133            action: ModelTurnAction::Continue,
10134            calls: first_calls.clone(),
10135        });
10136        flat.push(FixedModelTurnAction {
10137            action: ModelTurnAction::repeat(),
10138            calls: retry_calls.clone(),
10139        });
10140        flat.push(FixedModelTurnAction {
10141            action: ModelTurnAction::stop("unreachable"),
10142            calls: skipped_calls.clone(),
10143        });
10144        assert!(matches!(
10145            flat.on_model_turn_finished(&ctx, event).await,
10146            ModelTurnAction::Retry(_)
10147        ));
10148        assert_eq!(first_calls.load(SeqCst), 1);
10149        assert_eq!(retry_calls.load(SeqCst), 1);
10150        assert_eq!(skipped_calls.load(SeqCst), 0);
10151
10152        let nested_retry_calls = Arc::new(AtomicU32::new(0));
10153        let outer_skipped_calls = Arc::new(AtomicU32::new(0));
10154        let mut nested = HookStack::new();
10155        nested.push(FixedModelTurnAction {
10156            action: ModelTurnAction::retry_with_feedback("fix it"),
10157            calls: nested_retry_calls.clone(),
10158        });
10159        let mut outer = HookStack::new();
10160        outer.push(nested);
10161        outer.push(FixedModelTurnAction {
10162            action: ModelTurnAction::Continue,
10163            calls: outer_skipped_calls.clone(),
10164        });
10165        assert!(matches!(
10166            outer.on_model_turn_finished(&ctx, event).await,
10167            ModelTurnAction::Retry(crate::agent::RetryRequest::Feedback(feedback))
10168                if feedback == "fix it"
10169        ));
10170        assert_eq!(nested_retry_calls.load(SeqCst), 1);
10171        assert_eq!(outer_skipped_calls.load(SeqCst), 0);
10172
10173        let stop_calls = Arc::new(AtomicU32::new(0));
10174        let after_stop_calls = Arc::new(AtomicU32::new(0));
10175        let mut stopping = HookStack::new();
10176        stopping.push(FixedModelTurnAction {
10177            action: ModelTurnAction::stop("stop now"),
10178            calls: stop_calls.clone(),
10179        });
10180        stopping.push(FixedModelTurnAction {
10181            action: ModelTurnAction::Continue,
10182            calls: after_stop_calls.clone(),
10183        });
10184        assert!(matches!(
10185            stopping.on_model_turn_finished(&ctx, event).await,
10186            ModelTurnAction::Stop(reason) if reason == "stop now"
10187        ));
10188        assert_eq!(stop_calls.load(SeqCst), 1);
10189        assert_eq!(after_stop_calls.load(SeqCst), 0);
10190    }
10191}