Skip to main content

rig_agent/test_utils/
model_conformance.rs

1//! Provider-neutral behavioral scenarios for completion-model conformance.
2//!
3//! These helpers test the model/agent contract only. Provider wire formats,
4//! authentication, HTTP streaming, and cassette matching remain provider-suite
5//! responsibilities.
6
7use std::{
8    collections::BTreeSet,
9    sync::{
10        Arc, Mutex, MutexGuard,
11        atomic::{AtomicUsize, Ordering},
12    },
13    time::{Duration, Instant},
14};
15
16use futures::StreamExt;
17use schemars::JsonSchema;
18use serde::{Deserialize, Serialize};
19
20use crate::{
21    agent::{
22        AgentBuilder, AgentHook, CompletionCallAction, CompletionCallEvent,
23        CompletionResponseEvent, HookContext, InvalidToolCallAction, MultiTurnStreamItem,
24        NoToolConfig, ObservationAction, OutputMode, RequestPatch, StreamingError,
25        ToolCall as ToolCallEvent, ToolCallAction, ToolResultAction, ToolResultEvent,
26        run::{AgentRun, AgentRunStep, ModelTurn, ModelTurnOutcome},
27    },
28    completion::{
29        AssistantContent, CompletionError, CompletionModel, Message, Prompt, PromptError,
30        ToolDefinition,
31    },
32    streaming::StreamingPrompt,
33    tool::{Tool, ToolContext},
34};
35use rig_core::message::{ToolChoice, UserContent};
36
37/// Typed failure from a portable model-conformance scenario.
38#[derive(Debug, thiserror::Error)]
39pub enum ScenarioError {
40    /// A buffered agent run failed.
41    #[error(transparent)]
42    Prompt(#[from] PromptError),
43    /// A direct model completion failed.
44    #[error(transparent)]
45    Completion(#[from] CompletionError),
46    /// A streaming agent run failed.
47    #[error(transparent)]
48    Streaming(#[from] StreamingError),
49    /// Structured content could not be decoded.
50    #[error(transparent)]
51    Json(#[from] serde_json::Error),
52    /// Rig's structured extractor failed.
53    #[error(transparent)]
54    Extraction(#[from] crate::extractor::ExtractionError),
55    /// The model or agent violated the portable behavioral contract.
56    #[error("{scenario} conformance failed: {details}")]
57    Contract {
58        /// Stable scenario name.
59        scenario: &'static str,
60        /// Actionable observation explaining the failure.
61        details: String,
62    },
63}
64
65impl ScenarioError {
66    fn contract(scenario: &'static str, details: impl Into<String>) -> Self {
67        Self::Contract {
68            scenario,
69            details: details.into(),
70        }
71    }
72}
73
74/// Validate the portable diagnostics carried by an unknown or disallowed tool
75/// call failure.
76pub fn validate_unknown_tool_failure(
77    error: &PromptError,
78    expected_tool: &str,
79    expected_allowed_tools: &[&str],
80) -> Result<(), ScenarioError> {
81    const SCENARIO: &str = "unknown_tool_failure";
82    let PromptError::UnknownToolCall {
83        tool_name,
84        allowed_tools,
85        chat_history,
86        ..
87    } = error
88    else {
89        return Err(ScenarioError::contract(
90            SCENARIO,
91            format!("expected UnknownToolCall, observed {error:?}"),
92        ));
93    };
94    let expected_allowed = expected_allowed_tools
95        .iter()
96        .map(|name| (*name).to_string())
97        .collect::<Vec<_>>();
98    // A rejected repair target is never written back into history; diagnostics
99    // retain the model's original call while `tool_name` names the rejected
100    // target. Requiring a call-bearing assistant turn covers both paths.
101    let history_has_call = chat_history.iter().any(|message| {
102        matches!(
103            message,
104            Message::Assistant { content, .. }
105                if content.iter().any(|item| matches!(item, AssistantContent::ToolCall(_)))
106        )
107    });
108    if tool_name != expected_tool || allowed_tools != &expected_allowed || !history_has_call {
109        return Err(ScenarioError::contract(
110            SCENARIO,
111            format!(
112                "tool={tool_name:?}, allowed={allowed_tools:?}, expected_tool={expected_tool:?}, expected_allowed={expected_allowed:?}, history_has_call={history_has_call}, history={chat_history:?}"
113            ),
114        ));
115    }
116    Ok(())
117}
118
119/// Validate cancellation diagnostics, including the exact reason and retained
120/// assistant tool-call history.
121pub fn validate_cancelled_failure(
122    error: &PromptError,
123    expected_reason: &str,
124    expected_tool: &str,
125) -> Result<(), ScenarioError> {
126    const SCENARIO: &str = "cancelled_failure";
127    let PromptError::PromptCancelled {
128        chat_history,
129        reason,
130    } = error
131    else {
132        return Err(ScenarioError::contract(
133            SCENARIO,
134            format!("expected PromptCancelled, observed {error:?}"),
135        ));
136    };
137    let history_has_call = chat_history.iter().any(|message| {
138        matches!(
139            message,
140            Message::Assistant { content, .. }
141                if content.iter().any(|item| matches!(
142                    item,
143                    AssistantContent::ToolCall(call) if call.function.name == expected_tool
144                ))
145        )
146    });
147    if reason != expected_reason || !history_has_call {
148        return Err(ScenarioError::contract(
149            SCENARIO,
150            format!(
151                "reason={reason:?}, expected={expected_reason:?}, history_has_call={history_has_call}, history={chat_history:?}"
152            ),
153        ));
154    }
155    Ok(())
156}
157
158/// Validate max-turn diagnostics, including the exact configured budget and a
159/// retained pending prompt.
160pub fn validate_max_turns_failure(
161    error: &PromptError,
162    expected_max_turns: usize,
163) -> Result<(), ScenarioError> {
164    const SCENARIO: &str = "max_turns_failure";
165    let PromptError::MaxTurnsError {
166        max_turns,
167        chat_history,
168        prompt,
169    } = error
170    else {
171        return Err(ScenarioError::contract(
172            SCENARIO,
173            format!("expected MaxTurnsError, observed {error:?}"),
174        ));
175    };
176    let pending_prompt_retained = matches!(
177        prompt.as_ref(),
178        Message::User { content } if content.iter().next().is_some()
179    );
180    if *max_turns != expected_max_turns || chat_history.is_empty() || !pending_prompt_retained {
181        return Err(ScenarioError::contract(
182            SCENARIO,
183            format!(
184                "max_turns={max_turns}, expected={expected_max_turns}, history={chat_history:?}, pending_prompt_retained={pending_prompt_retained}, pending_prompt={prompt:?}"
185            ),
186        ));
187    }
188    Ok(())
189}
190
191/// Decode a structured-output response with a typed conformance failure that
192/// retains the scenario name and raw response.
193pub fn decode_structured_output<T>(
194    scenario: &'static str,
195    response: &str,
196) -> Result<T, ScenarioError>
197where
198    T: serde::de::DeserializeOwned,
199{
200    serde_json::from_str(response).map_err(|error| {
201        ScenarioError::contract(
202            scenario,
203            format!("structured output did not decode: {error}; response={response:?}"),
204        )
205    })
206}
207
208/// Validate that model-family control markers are absent from user-visible
209/// output and persisted history.
210pub fn validate_protocol_hygiene(
211    scenario: &'static str,
212    visible_output: &str,
213    messages: &[Message],
214    forbidden_markers: &[&str],
215) -> Result<(), ScenarioError> {
216    let serialized = serde_json::to_string(messages)?;
217    let leaked = forbidden_markers
218        .iter()
219        .filter(|marker| visible_output.contains(**marker) || serialized.contains(**marker))
220        .copied()
221        .collect::<Vec<_>>();
222    if !leaked.is_empty() {
223        return Err(ScenarioError::contract(
224            scenario,
225            format!(
226                "protocol markers leaked: {leaked:?}; output={visible_output:?}, history={messages:?}"
227            ),
228        ));
229    }
230    Ok(())
231}
232
233/// Validate that every observed tool invocation contains the expected rewritten
234/// argument fields while allowing unrelated original fields to remain.
235pub fn validate_rewritten_arguments(
236    scenario: &'static str,
237    observations: &[serde_json::Value],
238    expected_fields: &serde_json::Value,
239) -> Result<(), ScenarioError> {
240    let Some(expected) = expected_fields.as_object() else {
241        return Err(ScenarioError::contract(
242            scenario,
243            "expected rewritten fields must be a JSON object",
244        ));
245    };
246    if observations.is_empty() {
247        return Err(ScenarioError::contract(
248            scenario,
249            "the rewritten tool was never invoked",
250        ));
251    }
252    for observation in observations {
253        let Some(actual) = observation.as_object() else {
254            return Err(ScenarioError::contract(
255                scenario,
256                format!("observed rewritten arguments were not an object: {observation:?}"),
257            ));
258        };
259        for (key, expected_value) in expected {
260            if actual.get(key) != Some(expected_value) {
261                return Err(ScenarioError::contract(
262                    scenario,
263                    format!(
264                        "rewritten field {key:?} expected {expected_value:?}, observed {observation:?}"
265                    ),
266                ));
267            }
268        }
269    }
270    Ok(())
271}
272
273/// Validate that a sensitive raw tool result was produced but did not reach the
274/// model's user-visible response after result hooks ran.
275pub fn validate_result_redaction(
276    scenario: &'static str,
277    tool_produced_secret: bool,
278    visible_output: &str,
279    secret: &str,
280) -> Result<(), ScenarioError> {
281    if !tool_produced_secret || visible_output.is_empty() || visible_output.contains(secret) {
282        return Err(ScenarioError::contract(
283            scenario,
284            format!(
285                "produced_secret={tool_produced_secret}, secret_visible={}, output={visible_output:?}",
286                visible_output.contains(secret)
287            ),
288        ));
289    }
290    Ok(())
291}
292
293/// Validate a provider-neutral person extraction and its usage accounting.
294pub fn validate_extraction_fields(
295    scenario: &'static str,
296    first_name: Option<&str>,
297    last_name: Option<&str>,
298    job: Option<&str>,
299    usage: crate::completion::Usage,
300) -> Result<(), ScenarioError> {
301    let fields_match = first_name.is_some_and(|value| value.eq_ignore_ascii_case("Ada"))
302        && last_name.is_some_and(|value| value.eq_ignore_ascii_case("Lovelace"))
303        && job.is_some_and(|value| value.to_ascii_lowercase().contains("mathematician"));
304    if !fields_match || !usage.has_values() {
305        return Err(ScenarioError::contract(
306            scenario,
307            format!(
308                "first_name={first_name:?}, last_name={last_name:?}, job={job:?}, usage={usage:?}"
309            ),
310        ));
311    }
312    Ok(())
313}
314
315/// Summary emitted by a portable model-conformance scenario.
316#[derive(Debug, Clone)]
317pub struct ScenarioReport {
318    /// Stable scenario name.
319    pub name: &'static str,
320    /// Number of model-invoked tool calls observed by the scenario.
321    pub tool_calls: usize,
322    /// Aggregated prompt tokens reported by the model across all turns.
323    pub prompt_tokens: u64,
324    /// Aggregated generated tokens reported by the model across all turns.
325    pub generated_tokens: u64,
326    /// Number of messages retained in the completed run history.
327    pub history_messages: usize,
328    /// End-to-end scenario duration.
329    pub duration: Duration,
330    /// The model's final user-visible response.
331    pub response: String,
332}
333
334/// Error used by the deterministic conformance tools.
335#[derive(Debug, thiserror::Error)]
336#[error("model-conformance tool failed")]
337pub struct ConformanceToolError;
338
339const FORCE_TOOLS_PREAMBLE: &str = "You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing results yourself. Once you have all the tool results you need, reply with the final numeric answer in plain text.";
340const PARALLEL_PROMPT: &str = "Compute 3 + 4 and 10 - 2. You MUST call the add tool and the subtract tool together in your first response, as two parallel function calls, then report both results.";
341const PING_OUTPUT: &str = "pong-crimson-7423";
342const MOTTO_OUTPUT: &str = "steady hands\ncalm waters";
343
344fn lock_recover<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
345    mutex
346        .lock()
347        .unwrap_or_else(|poisoned| poisoned.into_inner())
348}
349
350fn tool_result_values(message: &Message) -> Vec<serde_json::Value> {
351    let Message::User { content } = message else {
352        return Vec::new();
353    };
354    content
355        .iter()
356        .filter_map(|item| match item {
357            UserContent::ToolResult(result) => Some(result),
358            _ => None,
359        })
360        .flat_map(|result| result.content.iter())
361        .filter_map(|content| match content {
362            rig_core::message::ToolResultContent::Text(text) => {
363                Some(serde_json::Value::String(text.text.clone()))
364            }
365            rig_core::message::ToolResultContent::Json { value } => Some(value.clone()),
366            rig_core::message::ToolResultContent::Image(_) => None,
367        })
368        .collect()
369}
370
371fn validate_tool_correlation(
372    scenario: &'static str,
373    messages: &[Message],
374) -> Result<(), ScenarioError> {
375    let mut calls = Vec::new();
376    let mut results = Vec::new();
377    for message in messages {
378        match message {
379            Message::Assistant { content, .. } => {
380                calls.extend(content.iter().filter_map(|item| {
381                    match item {
382                        AssistantContent::ToolCall(call) => Some((
383                            call.id.as_str(),
384                            call.provider
385                                .as_ref()
386                                .map(|provider| provider.call_id.as_str()),
387                        )),
388                        _ => None,
389                    }
390                }));
391            }
392            Message::User { content } => {
393                results.extend(content.iter().filter_map(|item| {
394                    match item {
395                        UserContent::ToolResult(result) => Some((
396                            result.call.as_str(),
397                            result
398                                .provider
399                                .as_ref()
400                                .map(|provider| provider.call_id.as_str()),
401                        )),
402                        _ => None,
403                    }
404                }));
405            }
406            Message::System { .. } => {}
407        }
408    }
409    if calls.is_empty() {
410        return Err(ScenarioError::contract(
411            scenario,
412            format!("history has no assistant tool calls: {messages:?}"),
413        ));
414    }
415    for (id, call_id) in &calls {
416        let matches = results
417            .iter()
418            .filter(|(result_id, result_call_id)| result_id == id && call_id == result_call_id)
419            .count();
420        if matches != 1 {
421            return Err(ScenarioError::contract(
422                scenario,
423                format!(
424                    "tool call id={id:?} call_id={call_id:?} has {matches} correlated results; calls={calls:?}, results={results:?}"
425                ),
426            ));
427        }
428    }
429    if results.len() != calls.len() {
430        return Err(ScenarioError::contract(
431            scenario,
432            format!(
433                "history contains dangling calls or results: calls={calls:?}, results={results:?}"
434            ),
435        ));
436    }
437    Ok(())
438}
439
440#[derive(Debug, Deserialize, JsonSchema)]
441struct OperationArgs {
442    x: i64,
443    y: i64,
444}
445
446#[derive(Clone)]
447struct CountingAdd(Arc<AtomicUsize>);
448
449impl Tool for CountingAdd {
450    const NAME: &'static str = "add";
451    type Error = ConformanceToolError;
452    type Args = OperationArgs;
453    type Output = i64;
454
455    fn description(&self) -> String {
456        "Add x and y together".to_string()
457    }
458
459    fn parameters(&self) -> serde_json::Value {
460        serde_json::json!({
461            "type": "object",
462            "properties": {
463                "x": { "type": "number", "description": "The first operand" },
464                "y": { "type": "number", "description": "The second operand" }
465            },
466            "required": ["x", "y"]
467        })
468    }
469
470    async fn call(
471        &self,
472        _context: &mut ToolContext,
473        args: Self::Args,
474    ) -> Result<Self::Output, Self::Error> {
475        self.0.fetch_add(1, Ordering::SeqCst);
476        Ok(args.x + args.y)
477    }
478}
479
480#[derive(Clone)]
481struct CountingSum(Arc<AtomicUsize>);
482
483impl Tool for CountingSum {
484    const NAME: &'static str = "sum";
485    type Error = ConformanceToolError;
486    type Args = OperationArgs;
487    type Output = i64;
488
489    fn description(&self) -> String {
490        "Add x and y together (alias of add)".to_string()
491    }
492
493    fn parameters(&self) -> serde_json::Value {
494        CountingAdd(Arc::new(AtomicUsize::new(0))).parameters()
495    }
496
497    async fn call(
498        &self,
499        _context: &mut ToolContext,
500        args: Self::Args,
501    ) -> Result<Self::Output, Self::Error> {
502        self.0.fetch_add(1, Ordering::SeqCst);
503        Ok(args.x + args.y)
504    }
505}
506
507#[derive(Clone)]
508struct CountingSubtract(Arc<AtomicUsize>);
509
510impl Tool for CountingSubtract {
511    const NAME: &'static str = "subtract";
512    type Error = ConformanceToolError;
513    type Args = OperationArgs;
514    type Output = i64;
515
516    fn description(&self) -> String {
517        "Subtract y from x (i.e. x - y)".to_string()
518    }
519
520    fn parameters(&self) -> serde_json::Value {
521        serde_json::json!({
522            "type": "object",
523            "properties": {
524                "x": { "type": "number", "description": "The first operand" },
525                "y": { "type": "number", "description": "The second operand" }
526            },
527            "required": ["x", "y"]
528        })
529    }
530
531    async fn call(
532        &self,
533        _context: &mut ToolContext,
534        args: Self::Args,
535    ) -> Result<Self::Output, Self::Error> {
536        self.0.fetch_add(1, Ordering::SeqCst);
537        Ok(args.x - args.y)
538    }
539}
540
541#[derive(Clone)]
542struct RewriteArgument {
543    key: &'static str,
544    value: serde_json::Value,
545}
546
547impl AgentHook for RewriteArgument {
548    async fn on_tool_call(&self, _ctx: &HookContext, event: ToolCallEvent<'_>) -> ToolCallAction {
549        if event.tool_name != CountingAdd::NAME {
550            return ToolCallAction::run();
551        }
552        let Ok(mut arguments) = serde_json::from_str::<serde_json::Value>(event.args) else {
553            return ToolCallAction::run();
554        };
555        let Some(object) = arguments.as_object_mut() else {
556            return ToolCallAction::run();
557        };
558        object.insert(self.key.to_string(), self.value.clone());
559        ToolCallAction::rewrite(arguments)
560    }
561}
562
563#[derive(Clone, Default)]
564struct ObserveArguments(Arc<Mutex<Vec<serde_json::Value>>>);
565
566impl AgentHook for ObserveArguments {
567    async fn on_tool_call(&self, _ctx: &HookContext, event: ToolCallEvent<'_>) -> ToolCallAction {
568        let value = serde_json::from_str(event.args)
569            .unwrap_or_else(|_| serde_json::Value::String(event.args.to_string()));
570        lock_recover(&self.0).push(value);
571        ToolCallAction::run()
572    }
573}
574
575#[derive(Clone)]
576struct ReplaceResult(&'static str);
577
578impl AgentHook for ReplaceResult {
579    async fn on_tool_result(
580        &self,
581        _ctx: &HookContext,
582        event: ToolResultEvent<'_>,
583    ) -> ToolResultAction {
584        if event.tool_name == CountingAdd::NAME {
585            ToolResultAction::rewrite(self.0)
586        } else {
587            ToolResultAction::keep()
588        }
589    }
590}
591
592#[derive(Clone)]
593struct WrapResult;
594
595impl AgentHook for WrapResult {
596    async fn on_tool_result(
597        &self,
598        _ctx: &HookContext,
599        event: ToolResultEvent<'_>,
600    ) -> ToolResultAction {
601        if event.tool_name == CountingAdd::NAME {
602            ToolResultAction::rewrite(format!("[{}]", event.presentation.render()))
603        } else {
604            ToolResultAction::keep()
605        }
606    }
607}
608
609#[derive(Clone)]
610struct FirstTurnPatch(RequestPatch);
611
612impl AgentHook for FirstTurnPatch {
613    async fn on_completion_call(
614        &self,
615        ctx: &HookContext,
616        _event: CompletionCallEvent<'_>,
617    ) -> CompletionCallAction {
618        if ctx.turn() == 1 {
619            CompletionCallAction::patch(self.0.clone())
620        } else {
621            CompletionCallAction::continue_run()
622        }
623    }
624}
625
626#[derive(Clone)]
627struct StopAfterResult(&'static str);
628
629impl AgentHook for StopAfterResult {
630    async fn on_tool_result(
631        &self,
632        _ctx: &HookContext,
633        event: ToolResultEvent<'_>,
634    ) -> ToolResultAction {
635        if event.tool_name == CountingAdd::NAME {
636            ToolResultAction::stop(self.0)
637        } else {
638            ToolResultAction::keep()
639        }
640    }
641}
642
643#[derive(Debug, Deserialize, JsonSchema)]
644struct EmptyArgs {}
645
646#[derive(Clone)]
647struct PingTool(Arc<AtomicUsize>);
648
649impl Tool for PingTool {
650    const NAME: &'static str = "ping";
651    type Error = ConformanceToolError;
652    type Args = EmptyArgs;
653    type Output = String;
654
655    fn description(&self) -> String {
656        "Return the current ping marker. Takes no arguments.".to_string()
657    }
658
659    fn parameters(&self) -> serde_json::Value {
660        serde_json::json!({ "type": "object", "properties": {}, "required": [] })
661    }
662
663    async fn call(
664        &self,
665        _context: &mut ToolContext,
666        _args: Self::Args,
667    ) -> Result<Self::Output, Self::Error> {
668        self.0.fetch_add(1, Ordering::SeqCst);
669        Ok(PING_OUTPUT.to_string())
670    }
671}
672
673#[derive(Clone)]
674struct MottoTool(Arc<AtomicUsize>);
675
676impl Tool for MottoTool {
677    const NAME: &'static str = "fetch_motto";
678    type Error = ConformanceToolError;
679    type Args = EmptyArgs;
680    type Output = String;
681
682    fn description(&self) -> String {
683        "Fetch the two-line workshop motto.".to_string()
684    }
685
686    fn parameters(&self) -> serde_json::Value {
687        serde_json::json!({ "type": "object", "properties": {}, "required": [] })
688    }
689
690    async fn call(
691        &self,
692        _context: &mut ToolContext,
693        _args: Self::Args,
694    ) -> Result<Self::Output, Self::Error> {
695        self.0.fetch_add(1, Ordering::SeqCst);
696        Ok(MOTTO_OUTPUT.to_string())
697    }
698}
699
700#[derive(Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)]
701struct ConfigOutput {
702    service: String,
703    max_retries: u64,
704}
705
706#[derive(Clone)]
707struct ConfigTool(Arc<AtomicUsize>);
708
709impl Tool for ConfigTool {
710    const NAME: &'static str = "fetch_config";
711    type Error = ConformanceToolError;
712    type Args = EmptyArgs;
713    type Output = ConfigOutput;
714
715    fn description(&self) -> String {
716        "Fetch the service configuration object.".to_string()
717    }
718
719    fn parameters(&self) -> serde_json::Value {
720        serde_json::json!({ "type": "object", "properties": {}, "required": [] })
721    }
722
723    async fn call(
724        &self,
725        _context: &mut ToolContext,
726        _args: Self::Args,
727    ) -> Result<Self::Output, Self::Error> {
728        self.0.fetch_add(1, Ordering::SeqCst);
729        Ok(ConfigOutput {
730            service: "cassette-lab".to_string(),
731            max_retries: 3,
732        })
733    }
734}
735
736#[derive(Debug, Clone, Deserialize, Serialize, JsonSchema, PartialEq, Eq)]
737#[serde(rename_all = "snake_case")]
738enum ComplexMode {
739    Careful,
740    Fast,
741}
742
743#[derive(Debug, Clone, Deserialize, Serialize, JsonSchema, PartialEq, Eq)]
744struct ComplexProfile {
745    name: String,
746    tags: Vec<String>,
747}
748
749#[derive(Debug, Clone, Deserialize, Serialize, JsonSchema, PartialEq, Eq)]
750struct ComplexArgs {
751    profile: ComplexProfile,
752    mode: ComplexMode,
753    note: Option<String>,
754    quote: String,
755}
756
757#[derive(Debug, Deserialize, Serialize, JsonSchema)]
758struct ExtractedPerson {
759    #[schemars(required)]
760    first_name: Option<String>,
761    #[schemars(required)]
762    last_name: Option<String>,
763    #[schemars(required)]
764    job: Option<String>,
765}
766
767#[derive(Clone)]
768struct CaptureComplexTool {
769    calls: Arc<AtomicUsize>,
770    captured: Arc<Mutex<Option<ComplexArgs>>>,
771}
772
773impl Tool for CaptureComplexTool {
774    const NAME: &'static str = "store_profile";
775    type Error = ConformanceToolError;
776    type Args = ComplexArgs;
777    type Output = ComplexArgs;
778
779    fn description(&self) -> String {
780        "Store one profile with its nested tags, mode, optional note, and exact quoted text."
781            .to_string()
782    }
783
784    fn parameters(&self) -> serde_json::Value {
785        serde_json::to_value(schemars::schema_for!(ComplexArgs)).unwrap_or_default()
786    }
787
788    async fn call(
789        &self,
790        _context: &mut ToolContext,
791        args: Self::Args,
792    ) -> Result<Self::Output, Self::Error> {
793        self.calls.fetch_add(1, Ordering::SeqCst);
794        *lock_recover(&self.captured) = Some(args.clone());
795        Ok(args)
796    }
797}
798
799fn has_tool_roundtrip(messages: Option<&[Message]>) -> bool {
800    let saw_call = messages.is_some_and(|messages| {
801        messages.iter().any(|message| {
802            matches!(
803                message,
804                Message::Assistant { content, .. }
805                    if content.iter().any(|item| matches!(item, AssistantContent::ToolCall(_)))
806            )
807        })
808    });
809    let saw_result = messages.is_some_and(|messages| {
810        messages.iter().any(|message| {
811            matches!(
812                message,
813                Message::User { content }
814                    if content.iter().any(|item| matches!(item, UserContent::ToolResult(_)))
815            )
816        })
817    });
818    saw_call && saw_result
819}
820
821#[derive(Debug, Deserialize, JsonSchema)]
822struct RepeatArgs {
823    /// The text to repeat.
824    text: String,
825    /// Number of repetitions; defaults to 2 when omitted.
826    times: Option<u32>,
827}
828
829#[derive(Clone)]
830struct RepeatTool {
831    calls: Arc<AtomicUsize>,
832}
833
834impl Tool for RepeatTool {
835    const NAME: &'static str = "repeat_text";
836    type Error = ConformanceToolError;
837    type Args = RepeatArgs;
838    type Output = String;
839
840    fn description(&self) -> String {
841        "Repeat `text`. `times` is optional and defaults to 2.".to_string()
842    }
843
844    fn parameters(&self) -> serde_json::Value {
845        serde_json::to_value(schemars::schema_for!(RepeatArgs)).unwrap_or_default()
846    }
847
848    async fn call(
849        &self,
850        _context: &mut ToolContext,
851        args: Self::Args,
852    ) -> Result<Self::Output, Self::Error> {
853        self.calls.fetch_add(1, Ordering::SeqCst);
854        Ok(vec![args.text.as_str(); args.times.unwrap_or(2) as usize].join(" "))
855    }
856}
857
858#[derive(Debug, Deserialize, JsonSchema)]
859struct BinOpArgs {
860    a: i64,
861    b: i64,
862}
863
864#[derive(Debug, Deserialize, JsonSchema)]
865struct ArithmeticResult {
866    answer: i64,
867    explanation: Option<String>,
868}
869
870#[derive(Clone)]
871struct AddTool(Arc<AtomicUsize>);
872
873impl Tool for AddTool {
874    const NAME: &'static str = "add";
875    type Error = ConformanceToolError;
876    type Args = BinOpArgs;
877    type Output = i64;
878
879    fn description(&self) -> String {
880        "Add two integers a and b.".to_string()
881    }
882
883    fn parameters(&self) -> serde_json::Value {
884        serde_json::to_value(schemars::schema_for!(BinOpArgs)).unwrap_or_default()
885    }
886
887    async fn call(
888        &self,
889        _context: &mut ToolContext,
890        args: Self::Args,
891    ) -> Result<Self::Output, Self::Error> {
892        self.0.fetch_add(1, Ordering::SeqCst);
893        Ok(args.a + args.b)
894    }
895}
896
897#[derive(Clone)]
898struct MultiplyTool(Arc<AtomicUsize>);
899
900impl Tool for MultiplyTool {
901    const NAME: &'static str = "multiply";
902    type Error = ConformanceToolError;
903    type Args = BinOpArgs;
904    type Output = i64;
905
906    fn description(&self) -> String {
907        "Multiply two integers a and b.".to_string()
908    }
909
910    fn parameters(&self) -> serde_json::Value {
911        serde_json::to_value(schemars::schema_for!(BinOpArgs)).unwrap_or_default()
912    }
913
914    async fn call(
915        &self,
916        _context: &mut ToolContext,
917        args: Self::Args,
918    ) -> Result<Self::Output, Self::Error> {
919        self.0.fetch_add(1, Ordering::SeqCst);
920        Ok(args.a * args.b)
921    }
922}
923
924fn report_from_response(
925    name: &'static str,
926    started: Instant,
927    tool_calls: usize,
928    response: crate::agent::PromptResponse,
929) -> Result<ScenarioReport, ScenarioError> {
930    if let Some(messages) = response.messages.as_deref() {
931        validate_protocol_hygiene(
932            name,
933            &response.output,
934            messages,
935            &[
936                "<tool_call>",
937                "</tool_call>",
938                "<tool_response>",
939                "</tool_response>",
940                "<|im_start|>",
941                "<|im_end|>",
942                "<think>",
943                "</think>",
944            ],
945        )?;
946    }
947    Ok(ScenarioReport {
948        name,
949        tool_calls,
950        prompt_tokens: response.usage.input_tokens,
951        generated_tokens: response.usage.output_tokens,
952        history_messages: response.messages.as_ref().map_or(0, Vec::len),
953        duration: started.elapsed(),
954        response: response.output,
955    })
956}
957
958/// Require the extended run's accumulated message history and validate
959/// canonical tool call/result correlation over it.
960fn correlated_messages<'a>(
961    scenario: &'static str,
962    response: &'a crate::agent::PromptResponse,
963) -> Result<&'a [Message], ScenarioError> {
964    let messages = response.messages.as_deref().ok_or_else(|| {
965        ScenarioError::contract(scenario, "extended run omitted accumulated message history")
966    })?;
967    validate_tool_correlation(scenario, messages)?;
968    Ok(messages)
969}
970
971/// [`correlated_messages`], flattened into every tool-result value in history.
972fn correlated_result_values(
973    scenario: &'static str,
974    response: &crate::agent::PromptResponse,
975) -> Result<Vec<serde_json::Value>, ScenarioError> {
976    Ok(correlated_messages(scenario, response)?
977        .iter()
978        .flat_map(tool_result_values)
979        .collect())
980}
981
982fn value_matches_integer(value: &serde_json::Value, expected: i64) -> bool {
983    value.as_i64() == Some(expected)
984        || value
985            .as_str()
986            .and_then(|text| text.trim().parse::<i64>().ok())
987            == Some(expected)
988}
989
990/// Runs two independent calls in one assistant turn and validates canonical
991/// call/result history correlation.
992///
993/// Set `tool_concurrency` to `Some(1)` to prove that serial host execution does
994/// not split or drop the model's parallel call batch.
995pub async fn parallel_tools<M, F>(
996    model: M,
997    configure: F,
998    tool_concurrency: Option<usize>,
999) -> Result<ScenarioReport, ScenarioError>
1000where
1001    M: CompletionModel + 'static,
1002    F: FnOnce(AgentBuilder<NoToolConfig>) -> AgentBuilder<NoToolConfig>,
1003{
1004    let add_calls = Arc::new(AtomicUsize::new(0));
1005    let subtract_calls = Arc::new(AtomicUsize::new(0));
1006    let started = Instant::now();
1007    let agent = configure(AgentBuilder::new(model))
1008        .preamble(FORCE_TOOLS_PREAMBLE)
1009        .temperature(0.0)
1010        .tool(CountingAdd(add_calls.clone()))
1011        .tool(CountingSubtract(subtract_calls.clone()))
1012        .default_max_turns(3)
1013        .build();
1014    let request = agent.prompt(PARALLEL_PROMPT).max_turns(3);
1015    let response = match tool_concurrency {
1016        Some(concurrency) => {
1017            request
1018                .tool_concurrency(concurrency)
1019                .extended_details()
1020                .await?
1021        }
1022        None => request.extended_details().await?,
1023    };
1024    let scenario = if tool_concurrency == Some(1) {
1025        "parallel_tools_serial_execution"
1026    } else {
1027        "parallel_tools"
1028    };
1029    let messages = correlated_messages(scenario, &response)?;
1030
1031    let Some((call_index, calls)) = messages.iter().enumerate().find_map(|(index, message)| {
1032        let Message::Assistant { content, .. } = message else {
1033            return None;
1034        };
1035        let calls = content
1036            .iter()
1037            .filter_map(|item| match item {
1038                AssistantContent::ToolCall(call) => Some(call),
1039                _ => None,
1040            })
1041            .collect::<Vec<_>>();
1042        (calls.len() == 2).then_some((index, calls))
1043    }) else {
1044        return Err(ScenarioError::contract(
1045            scenario,
1046            format!("no assistant turn contained exactly two tool calls: {messages:?}"),
1047        ));
1048    };
1049    let mut names = calls
1050        .iter()
1051        .map(|call| call.function.name.as_str())
1052        .collect::<Vec<_>>();
1053    names.sort_unstable();
1054    if names != ["add", "subtract"] {
1055        return Err(ScenarioError::contract(
1056            scenario,
1057            format!("parallel turn called {names:?}, expected add and subtract"),
1058        ));
1059    }
1060    let results_message = messages.get(call_index + 1).ok_or_else(|| {
1061        ScenarioError::contract(
1062            scenario,
1063            "parallel call turn has no following result message",
1064        )
1065    })?;
1066    let values = tool_result_values(results_message);
1067    if !values.iter().any(|value| value_matches_integer(value, 7))
1068        || !values.iter().any(|value| value_matches_integer(value, 8))
1069        || values.len() != 2
1070    {
1071        return Err(ScenarioError::contract(
1072            scenario,
1073            format!("parallel result message did not contain exactly 7 and 8: {values:?}"),
1074        ));
1075    }
1076    let add = add_calls.load(Ordering::SeqCst);
1077    let subtract = subtract_calls.load(Ordering::SeqCst);
1078    if add != 1 || subtract != 1 {
1079        return Err(ScenarioError::contract(
1080            scenario,
1081            format!("execution counts were add={add}, subtract={subtract}, expected one each"),
1082        ));
1083    }
1084    report_from_response(scenario, started, add + subtract, response)
1085}
1086
1087/// Runs a zero-argument tool and validates verbatim string-result handling.
1088pub async fn zero_argument_tool<M, F>(
1089    model: M,
1090    configure: F,
1091) -> Result<ScenarioReport, ScenarioError>
1092where
1093    M: CompletionModel + 'static,
1094    F: FnOnce(AgentBuilder<NoToolConfig>) -> AgentBuilder<NoToolConfig>,
1095{
1096    const SCENARIO: &str = "zero_argument_tool";
1097    let calls = Arc::new(AtomicUsize::new(0));
1098    let started = Instant::now();
1099    let agent = configure(AgentBuilder::new(model))
1100        .preamble("You must use the provided tools. Report tool outputs exactly as returned.")
1101        .temperature(0.0)
1102        .tool(PingTool(calls.clone()))
1103        .default_max_turns(2)
1104        .build();
1105    let response = agent
1106        .prompt("Call the ping tool, then report the exact marker it returns.")
1107        .max_turns(2)
1108        .extended_details()
1109        .await?;
1110    let values = correlated_result_values(SCENARIO, &response)?;
1111    if calls.load(Ordering::SeqCst) != 1
1112        || !values
1113            .iter()
1114            .any(|value| value.as_str() == Some(PING_OUTPUT))
1115        || !response.output.contains(PING_OUTPUT)
1116    {
1117        return Err(ScenarioError::contract(
1118            SCENARIO,
1119            format!(
1120                "calls={}, results={values:?}, response={:?}",
1121                calls.load(Ordering::SeqCst),
1122                response.output
1123            ),
1124        ));
1125    }
1126    report_from_response(SCENARIO, started, 1, response)
1127}
1128
1129/// Runs string- and JSON-returning tools and validates that neither output is
1130/// double encoded.
1131pub async fn tool_output_serialization<M, F>(
1132    model: M,
1133    configure: F,
1134) -> Result<ScenarioReport, ScenarioError>
1135where
1136    M: CompletionModel + 'static,
1137    F: FnOnce(AgentBuilder<NoToolConfig>) -> AgentBuilder<NoToolConfig>,
1138{
1139    const SCENARIO: &str = "tool_output_serialization";
1140    let started = Instant::now();
1141    let motto_calls = Arc::new(AtomicUsize::new(0));
1142    let config_calls = Arc::new(AtomicUsize::new(0));
1143    let agent = configure(AgentBuilder::new(model))
1144        .preamble("You must use the provided tools before answering.")
1145        .temperature(0.0)
1146        .tool(MottoTool(motto_calls.clone()))
1147        .tool(ConfigTool(config_calls.clone()))
1148        .default_max_turns(3)
1149        .build();
1150    let response = agent
1151        .prompt("Call fetch_motto and fetch_config, then summarize both outputs in one sentence.")
1152        .max_turns(3)
1153        .extended_details()
1154        .await?;
1155    let values = correlated_result_values(SCENARIO, &response)?;
1156    let expected_config = serde_json::to_value(ConfigOutput {
1157        service: "cassette-lab".to_string(),
1158        max_retries: 3,
1159    })?;
1160    let motto_ok = values
1161        .iter()
1162        .any(|value| value.as_str() == Some(MOTTO_OUTPUT));
1163    let config_ok = values.iter().any(|value| {
1164        value == &expected_config
1165            || value
1166                .as_str()
1167                .and_then(|text| serde_json::from_str::<serde_json::Value>(text).ok())
1168                .as_ref()
1169                == Some(&expected_config)
1170    });
1171    let motto_count = motto_calls.load(Ordering::SeqCst);
1172    let config_count = config_calls.load(Ordering::SeqCst);
1173    if !motto_ok || !config_ok || motto_count != 1 || config_count != 1 {
1174        return Err(ScenarioError::contract(
1175            SCENARIO,
1176            format!(
1177                "expected one verbatim motto and one semantic config JSON; motto_calls={motto_count}, config_calls={config_count}, values={values:?}"
1178            ),
1179        ));
1180    }
1181    report_from_response(SCENARIO, started, motto_count + config_count, response)
1182}
1183
1184/// Runs a nested, escaped, Unicode-bearing argument payload through typed tool
1185/// deserialization and validates the exact semantic value received by the tool.
1186pub async fn complex_tool_arguments<M, F>(
1187    model: M,
1188    configure: F,
1189) -> Result<ScenarioReport, ScenarioError>
1190where
1191    M: CompletionModel + 'static,
1192    F: FnOnce(AgentBuilder<NoToolConfig>) -> AgentBuilder<NoToolConfig>,
1193{
1194    const SCENARIO: &str = "complex_tool_arguments";
1195    let expected = ComplexArgs {
1196        profile: ComplexProfile {
1197            name: "Zoë \"Z\"".to_string(),
1198            tags: vec!["rust".to_string(), "東京".to_string()],
1199        },
1200        mode: ComplexMode::Careful,
1201        note: Some("line one\nline two".to_string()),
1202        quote: "path C:\\tmp and \"quoted\"".to_string(),
1203    };
1204    let calls = Arc::new(AtomicUsize::new(0));
1205    let captured = Arc::new(Mutex::new(None));
1206    let started = Instant::now();
1207    let agent = configure(AgentBuilder::new(model))
1208        .preamble("Use store_profile exactly once with every value supplied by the user.")
1209        .temperature(0.0)
1210        .tool(CaptureComplexTool {
1211            calls: calls.clone(),
1212            captured: captured.clone(),
1213        })
1214        .default_max_turns(3)
1215        .build();
1216    let response = agent
1217        .prompt(
1218            "Call store_profile with profile.name exactly `Zoë \\\"Z\\\"`, profile.tags exactly [`rust`, `東京`], mode `careful`, note containing the two lines `line one` and `line two` separated by a newline, and quote exactly `path C:\\\\tmp and \\\"quoted\\\"`. Then confirm it was stored.",
1219        )
1220        .max_turns(3)
1221        .extended_details()
1222        .await?;
1223    let observed = lock_recover(&captured).clone();
1224    if calls.load(Ordering::SeqCst) != 1 || observed.as_ref() != Some(&expected) {
1225        return Err(ScenarioError::contract(
1226            SCENARIO,
1227            format!(
1228                "calls={}, expected={expected:?}, observed={observed:?}, response={:?}",
1229                calls.load(Ordering::SeqCst),
1230                response.output
1231            ),
1232        ));
1233    }
1234    correlated_messages(SCENARIO, &response)?;
1235    report_from_response(SCENARIO, started, 1, response)
1236}
1237
1238/// Runs the same deterministic text request through buffered and raw streaming
1239/// completion surfaces and validates equivalent visible content and usage.
1240pub async fn buffered_streaming_text_parity<M>(model: M) -> Result<ScenarioReport, ScenarioError>
1241where
1242    M: CompletionModel + Clone + 'static,
1243{
1244    const SCENARIO: &str = "buffered_streaming_text_parity";
1245    const PROMPT: &str = "Answer with exactly the single word Paris.";
1246    let started = Instant::now();
1247    let request = || {
1248        model
1249            .completion_request(PROMPT)
1250            .temperature(0.0)
1251            .max_tokens(32)
1252            .build()
1253    };
1254    let buffered = model.completion(request()).await?;
1255    let buffered_text = buffered
1256        .choice
1257        .iter()
1258        .filter_map(|item| match item {
1259            AssistantContent::Text(text) => Some(text.text.as_str()),
1260            _ => None,
1261        })
1262        .collect::<String>();
1263
1264    let mut stream = model.stream(request()).await?;
1265    let mut streamed_text = String::new();
1266    let mut streamed_usage = None;
1267    while let Some(item) = stream.next().await {
1268        match item? {
1269            crate::streaming::StreamedAssistantContent::Text(text) => {
1270                streamed_text.push_str(&text.text);
1271            }
1272            crate::streaming::StreamedAssistantContent::Final(response) => {
1273                streamed_usage = Some(response.usage);
1274            }
1275            crate::streaming::StreamedAssistantContent::ToolCall { .. }
1276            | crate::streaming::StreamedAssistantContent::ToolCallDelta { .. }
1277            | crate::streaming::StreamedAssistantContent::Reasoning { .. }
1278            | crate::streaming::StreamedAssistantContent::ReasoningDelta { .. }
1279            | crate::streaming::StreamedAssistantContent::Unknown(_) => {}
1280        }
1281    }
1282    let usage = streamed_usage.ok_or_else(|| {
1283        ScenarioError::contract(SCENARIO, "raw stream omitted its final response metadata")
1284    })?;
1285    let normalize = |text: &str| {
1286        text.trim()
1287            .trim_matches(|character: char| !character.is_alphanumeric())
1288            .to_string()
1289    };
1290    let buffered_answer = normalize(&buffered_text);
1291    let streamed_answer = normalize(&streamed_text);
1292    if !buffered_answer.eq_ignore_ascii_case("Paris")
1293        || !streamed_answer.eq_ignore_ascii_case("Paris")
1294        || !buffered.usage.has_values()
1295        || !usage.has_values()
1296    {
1297        return Err(ScenarioError::contract(
1298            SCENARIO,
1299            format!(
1300                "buffered={buffered_text:?}, streamed={streamed_text:?}, buffered_usage={:?}, streamed_usage={usage:?}",
1301                buffered.usage
1302            ),
1303        ));
1304    }
1305    Ok(ScenarioReport {
1306        name: SCENARIO,
1307        tool_calls: 0,
1308        prompt_tokens: usage.input_tokens,
1309        generated_tokens: usage.output_tokens,
1310        history_messages: 0,
1311        duration: started.elapsed(),
1312        response: streamed_text,
1313    })
1314}
1315
1316/// Runs Rig's structured extractor and validates both the extracted semantic
1317/// fields and accumulated usage.
1318pub async fn structured_extraction<M>(model: M) -> Result<ScenarioReport, ScenarioError>
1319where
1320    M: CompletionModel + 'static,
1321{
1322    const SCENARIO: &str = "structured_extraction";
1323    const INPUT: &str = "Hello, my name is Ada Lovelace and I work as a mathematician.";
1324    let started = Instant::now();
1325    let response = crate::extractor::ExtractorBuilder::<ExtractedPerson>::new(model)
1326        .max_tokens(384)
1327        .retries(0)
1328        .build()
1329        .extract_with_usage(INPUT)
1330        .await?;
1331    validate_extraction_fields(
1332        SCENARIO,
1333        response.data.first_name.as_deref(),
1334        response.data.last_name.as_deref(),
1335        response.data.job.as_deref(),
1336        response.usage,
1337    )?;
1338    Ok(ScenarioReport {
1339        name: SCENARIO,
1340        tool_calls: 1,
1341        prompt_tokens: response.usage.input_tokens,
1342        generated_tokens: response.usage.output_tokens,
1343        history_messages: 0,
1344        duration: started.elapsed(),
1345        response: format!(
1346            "{} {} — {}",
1347            response.data.first_name.as_deref().unwrap_or_default(),
1348            response.data.last_name.as_deref().unwrap_or_default(),
1349            response.data.job.as_deref().unwrap_or_default()
1350        ),
1351    })
1352}
1353
1354fn restricted_recovery_run(
1355    prompt: &str,
1356    turn: ModelTurn,
1357    retries: usize,
1358) -> Result<AgentRun, ScenarioError> {
1359    const SCENARIO: &str = "invalid_tool_recovery";
1360    let mut run = AgentRun::new(prompt)
1361        .max_turns(2)
1362        .max_invalid_tool_call_retries(retries);
1363    if !matches!(run.next_step()?, AgentRunStep::CallModel { .. }) {
1364        return Err(ScenarioError::contract(
1365            SCENARIO,
1366            "fresh AgentRun did not request a model turn",
1367        ));
1368    }
1369    let outcome = run.model_response(turn)?;
1370    let ModelTurnOutcome::NeedsResolution(context) = outcome else {
1371        return Err(ScenarioError::contract(
1372            SCENARIO,
1373            format!("disallowed tool call did not require resolution: {outcome:?}"),
1374        ));
1375    };
1376    if context.tool_name != CountingAdd::NAME {
1377        return Err(ScenarioError::contract(
1378            SCENARIO,
1379            format!("expected rejected add call, observed {context:?}"),
1380        ));
1381    }
1382    Ok(run)
1383}
1384
1385/// Uses one real model turn to exercise fail-fast, retry exhaustion, repair,
1386/// rejected repair, and skip handling without executing a disallowed call.
1387pub async fn invalid_tool_recovery<M, F>(
1388    model: M,
1389    configure: F,
1390) -> Result<ScenarioReport, ScenarioError>
1391where
1392    M: CompletionModel + 'static,
1393    F: FnOnce(AgentBuilder<NoToolConfig>) -> AgentBuilder<NoToolConfig>,
1394{
1395    const SCENARIO: &str = "invalid_tool_recovery";
1396    const PROMPT: &str = "Call the add tool exactly once with x=2 and y=3. Do not call sum.";
1397    let started = Instant::now();
1398    let add_calls = Arc::new(AtomicUsize::new(0));
1399    let sum_calls = Arc::new(AtomicUsize::new(0));
1400    let agent = configure(AgentBuilder::new(model))
1401        .preamble(FORCE_TOOLS_PREAMBLE)
1402        .temperature(0.0)
1403        .tool(CountingAdd(add_calls.clone()))
1404        .tool(CountingSum(sum_calls.clone()))
1405        .tool_choice(ToolChoice::Required)
1406        .build();
1407    #[derive(Clone)]
1408    struct CaptureTurn(Arc<Mutex<Option<ModelTurn>>>);
1409
1410    impl AgentHook for CaptureTurn {
1411        async fn on_completion_response(
1412            &self,
1413            _ctx: &HookContext,
1414            event: CompletionResponseEvent<'_>,
1415        ) -> ObservationAction {
1416            *lock_recover(&self.0) = Some(ModelTurn::new(
1417                event.message_id.map(str::to_owned),
1418                event.content.clone(),
1419                event.usage,
1420                BTreeSet::new(),
1421                BTreeSet::new(),
1422            ));
1423            ObservationAction::stop("captured conformance model turn")
1424        }
1425    }
1426
1427    let captured = Arc::new(Mutex::new(None));
1428    let stopped = agent
1429        .runner(PROMPT)
1430        .add_hook(CaptureTurn(captured.clone()))
1431        .run()
1432        .await;
1433    if !matches!(stopped, Err(PromptError::PromptCancelled { .. })) {
1434        return Err(ScenarioError::contract(
1435            SCENARIO,
1436            format!("capture hook did not stop after the model response: {stopped:?}"),
1437        ));
1438    }
1439    let response = lock_recover(&captured).take().ok_or_else(|| {
1440        ScenarioError::contract(SCENARIO, "capture hook observed no model response")
1441    })?;
1442    let emitted = response
1443        .choice
1444        .iter()
1445        .filter(|item| {
1446            matches!(item, AssistantContent::ToolCall(call) if call.function.name == CountingAdd::NAME)
1447        })
1448        .count();
1449    if emitted != 1 {
1450        return Err(ScenarioError::contract(
1451            SCENARIO,
1452            format!(
1453                "model emitted {emitted} add calls, response={:?}",
1454                response.choice
1455            ),
1456        ));
1457    }
1458    let executable = BTreeSet::from([CountingAdd::NAME.to_string(), CountingSum::NAME.to_string()]);
1459    let allowed = BTreeSet::from([CountingSum::NAME.to_string()]);
1460    let turn = ModelTurn::new(
1461        response.message_id,
1462        response.choice,
1463        response.usage,
1464        executable,
1465        allowed,
1466    );
1467
1468    let mut fail = restricted_recovery_run(PROMPT, turn.clone(), 0)?;
1469    let error = match fail.resolve_invalid_tool_call(InvalidToolCallAction::fail()) {
1470        Err(error) => error,
1471        Ok(outcome) => {
1472            return Err(ScenarioError::contract(
1473                SCENARIO,
1474                format!("fail action unexpectedly returned {outcome:?}"),
1475            ));
1476        }
1477    };
1478    validate_unknown_tool_failure(&error, CountingAdd::NAME, &[CountingSum::NAME])?;
1479
1480    let mut retry = restricted_recovery_run(PROMPT, turn.clone(), 0)?;
1481    let error = match retry
1482        .resolve_invalid_tool_call(InvalidToolCallAction::retry("choose an allowed tool"))
1483    {
1484        Err(error) => error,
1485        Ok(outcome) => {
1486            return Err(ScenarioError::contract(
1487                SCENARIO,
1488                format!("exhausted retry unexpectedly returned {outcome:?}"),
1489            ));
1490        }
1491    };
1492    validate_unknown_tool_failure(&error, CountingAdd::NAME, &[CountingSum::NAME])?;
1493
1494    let mut rejected_repair = restricted_recovery_run(PROMPT, turn.clone(), 0)?;
1495    let error =
1496        match rejected_repair.resolve_invalid_tool_call(InvalidToolCallAction::repair("missing")) {
1497            Err(error) => error,
1498            Ok(outcome) => {
1499                return Err(ScenarioError::contract(
1500                    SCENARIO,
1501                    format!("disallowed repair unexpectedly returned {outcome:?}"),
1502                ));
1503            }
1504        };
1505    validate_unknown_tool_failure(&error, "missing", &[CountingSum::NAME])?;
1506
1507    let mut repaired = restricted_recovery_run(PROMPT, turn.clone(), 0)?;
1508    if !matches!(
1509        repaired.resolve_invalid_tool_call(InvalidToolCallAction::repair(CountingSum::NAME))?,
1510        ModelTurnOutcome::Continue { .. }
1511    ) {
1512        return Err(ScenarioError::contract(
1513            SCENARIO,
1514            "valid repair did not continue",
1515        ));
1516    }
1517    let AgentRunStep::CallTools { calls } = repaired.next_step()? else {
1518        return Err(ScenarioError::contract(
1519            SCENARIO,
1520            "valid repair did not produce pending tool execution",
1521        ));
1522    };
1523    let repaired_call = calls.first();
1524    if calls.len() != 1
1525        || !repaired_call.is_some_and(|call| {
1526            call.tool_call.function.name == CountingSum::NAME && call.preresolved_result.is_none()
1527        })
1528    {
1529        return Err(ScenarioError::contract(
1530            SCENARIO,
1531            format!("repaired pending calls were incorrect: {calls:?}"),
1532        ));
1533    }
1534
1535    let mut skipped = restricted_recovery_run(PROMPT, turn.clone(), 0)?;
1536    if !matches!(
1537        skipped.resolve_invalid_tool_call(InvalidToolCallAction::skip("disabled for this turn"))?,
1538        ModelTurnOutcome::Continue { .. }
1539    ) {
1540        return Err(ScenarioError::contract(SCENARIO, "skip did not continue"));
1541    }
1542    let AgentRunStep::CallTools { calls } = skipped.next_step()? else {
1543        return Err(ScenarioError::contract(
1544            SCENARIO,
1545            "skip did not produce a pre-resolved pending call",
1546        ));
1547    };
1548    let skipped_is_preresolved = match calls.first() {
1549        Some(call) => call.preresolved_result.is_some(),
1550        None => false,
1551    };
1552    if calls.len() != 1 || !skipped_is_preresolved {
1553        return Err(ScenarioError::contract(
1554            SCENARIO,
1555            format!("skipped pending calls were incorrect: {calls:?}"),
1556        ));
1557    }
1558    if add_calls.load(Ordering::SeqCst) != 0 || sum_calls.load(Ordering::SeqCst) != 0 {
1559        return Err(ScenarioError::contract(
1560            SCENARIO,
1561            "recovery scenario executed a tool body",
1562        ));
1563    }
1564
1565    Ok(ScenarioReport {
1566        name: SCENARIO,
1567        tool_calls: emitted,
1568        prompt_tokens: turn.usage.input_tokens,
1569        generated_tokens: turn.usage.output_tokens,
1570        history_messages: 2,
1571        duration: started.elapsed(),
1572        response: "fail, retry, repair, rejected repair, and skip passed".to_string(),
1573    })
1574}
1575
1576/// Exercises chained argument/result rewrites and a first-turn-only request
1577/// patch. Completion proves `tool_choice=Required` did not leak to turn two.
1578pub async fn hook_rewrites_and_request_patch<M, F>(
1579    model: M,
1580    configure: F,
1581) -> Result<ScenarioReport, ScenarioError>
1582where
1583    M: CompletionModel + 'static,
1584    F: FnOnce(AgentBuilder<NoToolConfig>) -> AgentBuilder<NoToolConfig>,
1585{
1586    const SCENARIO: &str = "hook_rewrites_and_request_patch";
1587    let started = Instant::now();
1588    let calls = Arc::new(AtomicUsize::new(0));
1589    let observed = ObserveArguments::default();
1590    let observed_probe = observed.clone();
1591    let agent = configure(AgentBuilder::new(model))
1592        .preamble("Use add for arithmetic and report only the tool result.")
1593        .temperature(0.0)
1594        .tool(CountingAdd(calls.clone()))
1595        .default_max_turns(3)
1596        .build();
1597    let response = agent
1598        .prompt("Use add once for x=1 and y=1, then report what the tool returns.")
1599        .max_turns(3)
1600        .add_hook(FirstTurnPatch(
1601            RequestPatch::new()
1602                .active_tools([CountingAdd::NAME])
1603                .tool_choice(ToolChoice::Required),
1604        ))
1605        .add_hook(RewriteArgument {
1606            key: "x",
1607            value: serde_json::json!(7),
1608        })
1609        .add_hook(RewriteArgument {
1610            key: "y",
1611            value: serde_json::json!(8),
1612        })
1613        .add_hook(observed)
1614        .add_hook(ReplaceResult("portable-redacted"))
1615        .add_hook(WrapResult)
1616        .extended_details()
1617        .await?;
1618    let observations = lock_recover(&observed_probe.0).clone();
1619    validate_rewritten_arguments(
1620        SCENARIO,
1621        &observations,
1622        &serde_json::json!({ "x": 7, "y": 8 }),
1623    )?;
1624    let messages = response.messages.as_deref().ok_or_else(|| {
1625        ScenarioError::contract(SCENARIO, "extended hook run omitted message history")
1626    })?;
1627    let results = messages
1628        .iter()
1629        .flat_map(tool_result_values)
1630        .collect::<Vec<_>>();
1631    if calls.load(Ordering::SeqCst) != 1
1632        || !results
1633            .iter()
1634            .any(|value| value == &serde_json::json!("[portable-redacted]"))
1635        || response.completion_calls.len() != 2
1636    {
1637        return Err(ScenarioError::contract(
1638            SCENARIO,
1639            format!(
1640                "calls={}, completion_calls={}, results={results:?}, output={:?}",
1641                calls.load(Ordering::SeqCst),
1642                response.completion_calls.len(),
1643                response.output
1644            ),
1645        ));
1646    }
1647    report_from_response(SCENARIO, started, 1, response)
1648}
1649
1650/// Exercises post-execution cancellation and max-turn diagnostics through the
1651/// public agent driver using real model-emitted calls.
1652pub async fn cancellation_and_max_turns<M, F>(
1653    model: M,
1654    configure: F,
1655) -> Result<ScenarioReport, ScenarioError>
1656where
1657    M: CompletionModel + Clone + 'static,
1658    F: Fn(AgentBuilder<NoToolConfig>) -> AgentBuilder<NoToolConfig>,
1659{
1660    const SCENARIO: &str = "cancellation_and_max_turns";
1661    const REASON: &str = "portable result veto";
1662    let started = Instant::now();
1663    let cancelled_calls = Arc::new(AtomicUsize::new(0));
1664    let cancelled_agent = configure(AgentBuilder::new(model.clone()))
1665        .preamble("Use add for arithmetic; never calculate by hand.")
1666        .temperature(0.0)
1667        .tool(CountingAdd(cancelled_calls.clone()))
1668        .build();
1669    let cancelled = match cancelled_agent
1670        .prompt("Use add once to compute x=20 plus y=22.")
1671        .max_turns(2)
1672        .add_hook(StopAfterResult(REASON))
1673        .await
1674    {
1675        Err(error) => error,
1676        Ok(output) => {
1677            return Err(ScenarioError::contract(
1678                SCENARIO,
1679                format!("result cancellation unexpectedly completed: {output:?}"),
1680            ));
1681        }
1682    };
1683    validate_cancelled_failure(&cancelled, REASON, CountingAdd::NAME)?;
1684
1685    let max_turn_calls = Arc::new(AtomicUsize::new(0));
1686    let max_turn_agent = configure(AgentBuilder::new(model))
1687        .preamble("Use add for arithmetic; never calculate by hand.")
1688        .temperature(0.0)
1689        .tool(CountingAdd(max_turn_calls.clone()))
1690        .build();
1691    let max_turn = match max_turn_agent
1692        .prompt("Use add once to compute x=20 plus y=22, then report the result.")
1693        .max_turns(1)
1694        .await
1695    {
1696        Err(error) => error,
1697        Ok(output) => {
1698            return Err(ScenarioError::contract(
1699                SCENARIO,
1700                format!("one-turn budget unexpectedly completed: {output:?}"),
1701            ));
1702        }
1703    };
1704    validate_max_turns_failure(&max_turn, 1)?;
1705    let cancelled_count = cancelled_calls.load(Ordering::SeqCst);
1706    let max_turn_count = max_turn_calls.load(Ordering::SeqCst);
1707    if cancelled_count != 1 || max_turn_count != 1 {
1708        return Err(ScenarioError::contract(
1709            SCENARIO,
1710            format!("cancelled executions={cancelled_count}, max-turn executions={max_turn_count}"),
1711        ));
1712    }
1713    Ok(ScenarioReport {
1714        name: SCENARIO,
1715        tool_calls: cancelled_count + max_turn_count,
1716        prompt_tokens: 0,
1717        generated_tokens: 0,
1718        history_messages: 4,
1719        duration: started.elapsed(),
1720        response: "post-result cancellation and max-turn diagnostics passed".to_string(),
1721    })
1722}
1723
1724/// Runs the portable optional-argument tool scenario.
1725///
1726/// `configure` is deliberately outside the scenario so a provider suite can
1727/// attach transport-only settings without putting them into the shared model
1728/// contract.
1729pub async fn optional_argument<M, F>(
1730    model: M,
1731    configure: F,
1732) -> Result<ScenarioReport, ScenarioError>
1733where
1734    M: CompletionModel + 'static,
1735    F: FnOnce(AgentBuilder<NoToolConfig>) -> AgentBuilder<NoToolConfig>,
1736{
1737    let calls = Arc::new(AtomicUsize::new(0));
1738    let started = Instant::now();
1739    let agent = configure(AgentBuilder::new(model))
1740        .preamble("Use the repeat_text tool whenever asked to repeat text.")
1741        .tool(RepeatTool {
1742            calls: calls.clone(),
1743        })
1744        .default_max_turns(4)
1745        .build();
1746    let result = agent
1747        .prompt(
1748            "Use the repeat_text tool to repeat the word \"banana\" 3 times, then show me the exact result.",
1749        )
1750        .extended_details()
1751        .await?;
1752    let response = result.output.clone();
1753    let tool_calls = calls.load(Ordering::SeqCst);
1754    if tool_calls == 0
1755        || response.matches("banana").count() < 1
1756        || !has_tool_roundtrip(result.messages.as_deref())
1757    {
1758        return Err(ScenarioError::contract(
1759            "optional_argument",
1760            format!("calls={tool_calls}, response={response:?}"),
1761        ));
1762    }
1763    report_from_response("optional_argument", started, tool_calls, result)
1764}
1765
1766/// Runs a portable two-tool sequential arithmetic scenario.
1767pub async fn sequential_tools<M, F>(model: M, configure: F) -> Result<ScenarioReport, ScenarioError>
1768where
1769    M: CompletionModel + 'static,
1770    F: FnOnce(AgentBuilder<NoToolConfig>) -> AgentBuilder<NoToolConfig>,
1771{
1772    let add_calls = Arc::new(AtomicUsize::new(0));
1773    let multiply_calls = Arc::new(AtomicUsize::new(0));
1774    let started = Instant::now();
1775    let agent = configure(AgentBuilder::new(model))
1776        .preamble(
1777            "You are a calculator. Use the add and multiply tools for arithmetic; never compute by hand.",
1778        )
1779        .tool(AddTool(add_calls.clone()))
1780        .tool(MultiplyTool(multiply_calls.clone()))
1781        .default_max_turns(6)
1782        .build();
1783    let result = agent
1784        .prompt(
1785            "Compute (4 + 6) * 2. First call the add tool, then call the multiply tool on the result. Tell me the final number.",
1786        )
1787        .extended_details()
1788        .await?;
1789    let response = result.output.clone();
1790    let add = add_calls.load(Ordering::SeqCst);
1791    let multiply = multiply_calls.load(Ordering::SeqCst);
1792    if add == 0
1793        || multiply == 0
1794        || !response.contains("20")
1795        || !has_tool_roundtrip(result.messages.as_deref())
1796    {
1797        return Err(ScenarioError::contract(
1798            "sequential_tools",
1799            format!("add={add}, multiply={multiply}, response={response:?}"),
1800        ));
1801    }
1802    report_from_response("sequential_tools", started, add + multiply, result)
1803}
1804
1805/// Runs a tool through Rig's multi-turn streaming agent driver.
1806pub async fn streaming_tool<M, F>(model: M, configure: F) -> Result<ScenarioReport, ScenarioError>
1807where
1808    M: CompletionModel + 'static,
1809    F: FnOnce(AgentBuilder<NoToolConfig>) -> AgentBuilder<NoToolConfig>,
1810{
1811    let calls = Arc::new(AtomicUsize::new(0));
1812    let started = Instant::now();
1813    let agent = configure(AgentBuilder::new(model))
1814        .preamble("Use the add tool for arithmetic; do not calculate by hand.")
1815        .tool(AddTool(calls.clone()))
1816        .default_max_turns(4)
1817        .build();
1818    let mut stream = agent
1819        .stream_prompt("Use add to calculate 17 + 25, then state the final number.")
1820        .max_turns(4)
1821        .await;
1822    let mut final_response = None;
1823    let mut final_count = 0_usize;
1824    let mut completion_usage = crate::completion::Usage::new();
1825    let mut streamed_call_ids = Vec::new();
1826    let mut streamed_result_ids = Vec::new();
1827    while let Some(item) = stream.next().await {
1828        match item? {
1829            MultiTurnStreamItem::StreamAssistantItem(
1830                crate::streaming::StreamedAssistantContent::ToolCall {
1831                    internal_call_id, ..
1832                },
1833            ) => streamed_call_ids.push(internal_call_id),
1834            MultiTurnStreamItem::StreamUserItem(
1835                crate::streaming::StreamedUserContent::ToolResult {
1836                    internal_call_id, ..
1837                },
1838            ) => streamed_result_ids.push(internal_call_id),
1839            MultiTurnStreamItem::CompletionCall(call) => completion_usage += call.usage,
1840            MultiTurnStreamItem::FinalResponse(response) => {
1841                final_count += 1;
1842                final_response = Some(response);
1843            }
1844            MultiTurnStreamItem::StreamAssistantItem(_)
1845            | MultiTurnStreamItem::ToolExecutionCommitted { .. }
1846            | MultiTurnStreamItem::ModelTurnRetried { .. } => {}
1847        }
1848    }
1849    let result = final_response.ok_or_else(|| {
1850        ScenarioError::contract("streaming_tool", "stream produced no final response")
1851    })?;
1852    let response = result.output.clone();
1853    let history_messages = result.messages.as_ref().map_or(0, Vec::len);
1854    let tool_calls = calls.load(Ordering::SeqCst);
1855    streamed_call_ids.sort();
1856    streamed_result_ids.sort();
1857    let correlated_stream =
1858        !streamed_call_ids.is_empty() && streamed_call_ids == streamed_result_ids;
1859    let correlated_history = result
1860        .messages
1861        .as_deref()
1862        .is_some_and(|messages| validate_tool_correlation("streaming_tool", messages).is_ok());
1863    if tool_calls == 0
1864        || !response.contains("42")
1865        || history_messages < 4
1866        || final_count != 1
1867        || !correlated_stream
1868        || !correlated_history
1869        || completion_usage != result.usage
1870        || result.completion_calls.is_empty()
1871    {
1872        return Err(ScenarioError::contract(
1873            "streaming_tool",
1874            format!(
1875                "calls={tool_calls}, final_count={final_count}, streamed_call_ids={streamed_call_ids:?}, streamed_result_ids={streamed_result_ids:?}, completion_usage={completion_usage:?}, final_usage={:?}, history_messages={history_messages}, response={response:?}",
1876                result.usage
1877            ),
1878        ));
1879    }
1880    if let Some(messages) = result.messages.as_deref() {
1881        validate_protocol_hygiene(
1882            "streaming_tool",
1883            &response,
1884            messages,
1885            &["<tool_call>", "</tool_call>", "<think>", "</think>"],
1886        )?;
1887    }
1888    Ok(ScenarioReport {
1889        name: "streaming_tool",
1890        tool_calls,
1891        prompt_tokens: result.usage.input_tokens,
1892        generated_tokens: result.usage.output_tokens,
1893        history_messages,
1894        duration: started.elapsed(),
1895        response,
1896    })
1897}
1898
1899/// Runs a normal tool followed by Rig's synthetic structured-output tool.
1900pub async fn structured_after_tool<M, F>(
1901    model: M,
1902    configure: F,
1903) -> Result<ScenarioReport, ScenarioError>
1904where
1905    M: CompletionModel + 'static,
1906    F: FnOnce(AgentBuilder<NoToolConfig>) -> AgentBuilder<NoToolConfig>,
1907{
1908    let calls = Arc::new(AtomicUsize::new(0));
1909    let started = Instant::now();
1910    let agent = configure(AgentBuilder::new(model))
1911        .preamble(
1912            "Use add for arithmetic, then finish by calling the structured output tool exactly once.",
1913        )
1914        .output_schema::<ArithmeticResult>()
1915        .output_mode(OutputMode::Tool)
1916        .tool(AddTool(calls.clone()))
1917        .default_max_turns(5)
1918        .build();
1919    let result = agent
1920        .prompt("Use add to calculate 19 + 23. Return answer=42 and a short optional explanation.")
1921        .extended_details()
1922        .await?;
1923    let response = result.output.clone();
1924    let parsed: ArithmeticResult = serde_json::from_str(&response)?;
1925    let tool_calls = calls.load(Ordering::SeqCst);
1926    if tool_calls == 0 || parsed.answer != 42 || !has_tool_roundtrip(result.messages.as_deref()) {
1927        return Err(ScenarioError::contract(
1928            "structured_after_tool",
1929            format!("calls={tool_calls}, response={response:?}"),
1930        ));
1931    }
1932    let _ = parsed.explanation;
1933    report_from_response("structured_after_tool", started, tool_calls + 1, result)
1934}
1935
1936/// Runs all portable tool-choice modes directly against a completion model.
1937pub async fn tool_choice_modes<M>(model: M) -> Result<ScenarioReport, ScenarioError>
1938where
1939    M: CompletionModel + Clone + 'static,
1940{
1941    let definition = |name: &str| ToolDefinition {
1942        name: name.to_string(),
1943        description: format!("Return the supplied integer using {name}."),
1944        parameters: serde_json::json!({
1945            "type": "object",
1946            "properties": {"value": {"type": "integer"}},
1947            "required": ["value"]
1948        }),
1949    };
1950    let tools = vec![definition("alpha"), definition("beta")];
1951    let started = Instant::now();
1952    let none = model
1953        .completion(
1954            model
1955                .completion_request("Answer with only the number 4. Do not call a function.")
1956                .tools(tools.clone())
1957                .tool_choice(ToolChoice::None)
1958                .temperature(0.0)
1959                .max_tokens(64)
1960                .build(),
1961        )
1962        .await?;
1963    if none
1964        .choice
1965        .iter()
1966        .any(|item| matches!(item, AssistantContent::ToolCall(_)))
1967    {
1968        return Err(ScenarioError::contract(
1969            "tool_choice_modes",
1970            "tool_choice none emitted a tool call",
1971        ));
1972    }
1973
1974    let required = model
1975        .completion(
1976            model
1977                .completion_request("Call alpha with value 7.")
1978                .tools(tools.clone())
1979                .tool_choice(ToolChoice::Required)
1980                .temperature(0.0)
1981                .max_tokens(96)
1982                .build(),
1983        )
1984        .await?;
1985    let required_calls = required
1986        .choice
1987        .iter()
1988        .filter(|item| matches!(item, AssistantContent::ToolCall(_)))
1989        .count();
1990    if required_calls == 0 {
1991        return Err(ScenarioError::contract(
1992            "tool_choice_modes",
1993            "tool_choice required emitted no tool call",
1994        ));
1995    }
1996
1997    let specific = model
1998        .completion(
1999            model
2000                .completion_request("Call beta with value 9.")
2001                .tools(tools)
2002                .tool_choice(ToolChoice::Specific {
2003                    function_names: vec!["beta".to_string()],
2004                })
2005                .temperature(0.0)
2006                .max_tokens(96)
2007                .build(),
2008        )
2009        .await?;
2010    let specific_calls = specific
2011        .choice
2012        .iter()
2013        .filter_map(|item| match item {
2014            AssistantContent::ToolCall(call) => Some(call),
2015            _ => None,
2016        })
2017        .collect::<Vec<_>>();
2018    if specific_calls.is_empty()
2019        || specific_calls
2020            .iter()
2021            .any(|call| call.function.name != "beta")
2022    {
2023        return Err(ScenarioError::contract(
2024            "tool_choice_modes",
2025            "specific tool choice did not select only beta",
2026        ));
2027    }
2028
2029    Ok(ScenarioReport {
2030        name: "tool_choice_modes",
2031        tool_calls: required_calls + specific_calls.len(),
2032        prompt_tokens: none.usage.input_tokens
2033            + required.usage.input_tokens
2034            + specific.usage.input_tokens,
2035        generated_tokens: none.usage.output_tokens
2036            + required.usage.output_tokens
2037            + specific.usage.output_tokens,
2038        history_messages: 0,
2039        duration: started.elapsed(),
2040        response: "none, required, and specific modes passed".to_string(),
2041    })
2042}
2043
2044/// Runs a streamed real-tool turn followed by Rig's synthetic output tool.
2045pub async fn streaming_structured_after_tool<M, F>(
2046    model: M,
2047    configure: F,
2048) -> Result<ScenarioReport, ScenarioError>
2049where
2050    M: CompletionModel + 'static,
2051    F: FnOnce(AgentBuilder<NoToolConfig>) -> AgentBuilder<NoToolConfig>,
2052{
2053    let calls = Arc::new(AtomicUsize::new(0));
2054    let started = Instant::now();
2055    let agent = configure(AgentBuilder::new(model))
2056        .preamble(
2057            "Use add for arithmetic, then finish by calling the structured output tool exactly once.",
2058        )
2059        .output_schema::<ArithmeticResult>()
2060        .output_mode(OutputMode::Tool)
2061        .tool(AddTool(calls.clone()))
2062        .default_max_turns(5)
2063        .build();
2064    let mut stream = agent
2065        .stream_prompt(
2066            "Use add to calculate 19 + 23. Return answer=42 and a short optional explanation.",
2067        )
2068        .max_turns(5)
2069        .await;
2070    let mut final_response = None;
2071    let mut final_count = 0_usize;
2072    while let Some(item) = stream.next().await {
2073        if let MultiTurnStreamItem::FinalResponse(response) = item? {
2074            final_count += 1;
2075            final_response = Some(response);
2076        }
2077    }
2078    let result = final_response.ok_or_else(|| {
2079        ScenarioError::contract(
2080            "streaming_structured_after_tool",
2081            "stream produced no final response",
2082        )
2083    })?;
2084    let parsed: ArithmeticResult = serde_json::from_str(&result.output)?;
2085    let calls = calls.load(Ordering::SeqCst);
2086    if calls == 0
2087        || final_count != 1
2088        || parsed.answer != 42
2089        || !has_tool_roundtrip(result.messages.as_deref())
2090    {
2091        return Err(ScenarioError::contract(
2092            "streaming_structured_after_tool",
2093            format!(
2094                "calls={calls}, final_count={final_count}, response={:?}",
2095                result.output
2096            ),
2097        ));
2098    }
2099    report_from_response(
2100        "streaming_structured_after_tool",
2101        started,
2102        calls + 1,
2103        result,
2104    )
2105}
2106
2107#[cfg(test)]
2108mod tests {
2109    use super::*;
2110    use crate::{
2111        completion::Usage,
2112        test_utils::{MockCompletionModel, MockStreamEvent, MockTurn, mock_final},
2113    };
2114    use rig_core::message::{ToolCall, ToolFunction};
2115
2116    fn tool_call(id: &str, name: &str, arguments: serde_json::Value) -> AssistantContent {
2117        AssistantContent::ToolCall(ToolCall::from_wire(
2118            id,
2119            ToolFunction::new(name.to_string(), arguments),
2120        ))
2121    }
2122
2123    fn usage(input: u64, output: u64) -> Usage {
2124        Usage {
2125            input_tokens: input,
2126            output_tokens: output,
2127            total_tokens: input + output,
2128            ..Usage::new()
2129        }
2130    }
2131
2132    fn fixture_contract(condition: bool, details: &str) -> Result<(), ScenarioError> {
2133        if condition {
2134            Ok(())
2135        } else {
2136            Err(ScenarioError::contract("test_fixture", details))
2137        }
2138    }
2139
2140    #[tokio::test]
2141    async fn parallel_contract_validates_batch_and_correlation() -> Result<(), ScenarioError> {
2142        let first = MockTurn::from_contents([
2143            tool_call("call_add", "add", serde_json::json!({"x": 3, "y": 4})),
2144            tool_call(
2145                "call_subtract",
2146                "subtract",
2147                serde_json::json!({"x": 10, "y": 2}),
2148            ),
2149        ]);
2150        let report = parallel_tools(
2151            MockCompletionModel::new([first, MockTurn::text("7 and 8")]),
2152            |builder| builder,
2153            Some(1),
2154        )
2155        .await?;
2156        fixture_contract(report.tool_calls == 2, "parallel tool-call count")?;
2157        fixture_contract(report.history_messages >= 4, "parallel history length")?;
2158        Ok(())
2159    }
2160
2161    #[tokio::test]
2162    async fn zero_argument_and_output_serialization_contracts_pass() -> Result<(), ScenarioError> {
2163        let zero = zero_argument_tool(
2164            MockCompletionModel::new([
2165                MockTurn::tool_call("ping_call", "ping", serde_json::json!({})),
2166                MockTurn::text(PING_OUTPUT),
2167            ]),
2168            |builder| builder,
2169        )
2170        .await?;
2171        fixture_contract(zero.tool_calls == 1, "zero-argument call count")?;
2172
2173        let first = MockTurn::from_contents([
2174            tool_call("motto_call", "fetch_motto", serde_json::json!({})),
2175            tool_call("config_call", "fetch_config", serde_json::json!({})),
2176        ]);
2177        let serialized = tool_output_serialization(
2178            MockCompletionModel::new([first, MockTurn::text("summary")]),
2179            |builder| builder,
2180        )
2181        .await?;
2182        fixture_contract(serialized.tool_calls == 2, "serialized-output call count")?;
2183        Ok(())
2184    }
2185
2186    #[tokio::test]
2187    async fn complex_arguments_preserve_nested_unicode_and_escapes() -> Result<(), ScenarioError> {
2188        let arguments = serde_json::json!({
2189            "profile": {"name": "Zoë \"Z\"", "tags": ["rust", "東京"]},
2190            "mode": "careful",
2191            "note": "line one\nline two",
2192            "quote": "path C:\\tmp and \"quoted\""
2193        });
2194        let report = complex_tool_arguments(
2195            MockCompletionModel::new([
2196                MockTurn::tool_call("profile_call", "store_profile", arguments),
2197                MockTurn::text("stored"),
2198            ]),
2199            |builder| builder,
2200        )
2201        .await?;
2202        fixture_contract(report.tool_calls == 1, "complex-argument call count")?;
2203        Ok(())
2204    }
2205
2206    #[tokio::test]
2207    async fn extraction_contract_requires_fields_and_usage() -> Result<(), ScenarioError> {
2208        let report = structured_extraction(MockCompletionModel::new([MockTurn::tool_call(
2209            "submit_call",
2210            "submit",
2211            serde_json::json!({
2212                "first_name": "Ada",
2213                "last_name": "Lovelace",
2214                "job": "mathematician"
2215            }),
2216        )
2217        .with_usage(usage(20, 5))]))
2218        .await?;
2219        fixture_contract(report.prompt_tokens == 20, "extraction input usage")?;
2220        fixture_contract(report.generated_tokens == 5, "extraction output usage")?;
2221        Ok(())
2222    }
2223
2224    #[tokio::test]
2225    async fn streaming_contract_checks_events_history_and_usage() -> Result<(), ScenarioError> {
2226        let model = MockCompletionModel::from_stream_turns([
2227            vec![
2228                MockStreamEvent::tool_call(
2229                    "add_call",
2230                    "add",
2231                    serde_json::json!({"a": 17, "b": 25}),
2232                ),
2233                MockStreamEvent::FinalResponse(mock_final(usage(10, 2))),
2234            ],
2235            vec![
2236                MockStreamEvent::text("42"),
2237                MockStreamEvent::FinalResponse(mock_final(usage(14, 1))),
2238            ],
2239        ]);
2240        let report = streaming_tool(model, |builder| builder).await?;
2241        fixture_contract(report.prompt_tokens == 24, "streaming input usage")?;
2242        fixture_contract(report.generated_tokens == 3, "streaming output usage")?;
2243        Ok(())
2244    }
2245
2246    #[tokio::test]
2247    async fn invalid_recovery_paths_do_not_execute_tools() -> Result<(), ScenarioError> {
2248        let report = invalid_tool_recovery(
2249            MockCompletionModel::new([MockTurn::tool_call(
2250                "invalid-add",
2251                "add",
2252                serde_json::json!({ "x": 2, "y": 3 }),
2253            )]),
2254            |builder| builder,
2255        )
2256        .await?;
2257        fixture_contract(report.tool_calls == 1, "recovery source call count")?;
2258        Ok(())
2259    }
2260
2261    #[tokio::test]
2262    async fn hook_rewrites_chain_and_request_patch_is_turn_local() -> Result<(), ScenarioError> {
2263        let report = hook_rewrites_and_request_patch(
2264            MockCompletionModel::new([
2265                MockTurn::tool_call("hook-add", "add", serde_json::json!({ "x": 1, "y": 1 })),
2266                MockTurn::text("[portable-redacted]"),
2267            ]),
2268            |builder| builder,
2269        )
2270        .await?;
2271        fixture_contract(report.tool_calls == 1, "hook execution count")?;
2272        Ok(())
2273    }
2274
2275    #[tokio::test]
2276    async fn cancellation_and_max_turn_controls_retain_diagnostics() -> Result<(), ScenarioError> {
2277        let report = cancellation_and_max_turns(
2278            MockCompletionModel::new([
2279                MockTurn::tool_call("cancel-add", "add", serde_json::json!({ "x": 20, "y": 22 })),
2280                MockTurn::tool_call("budget-add", "add", serde_json::json!({ "x": 20, "y": 22 })),
2281            ]),
2282            |builder| builder,
2283        )
2284        .await?;
2285        fixture_contract(report.tool_calls == 2, "run-control execution count")?;
2286        Ok(())
2287    }
2288
2289    #[test]
2290    fn typed_validators_reject_bad_structured_output_and_protocol_leaks() {
2291        let invalid = decode_structured_output::<ConfigOutput>("invalid_json", "not json");
2292        assert!(matches!(invalid, Err(ScenarioError::Contract { .. })));
2293
2294        let messages = vec![Message::Assistant {
2295            id: None,
2296            content: vec![AssistantContent::text("visible <tool_call>")],
2297        }];
2298        let hygiene = validate_protocol_hygiene(
2299            "protocol_hygiene",
2300            "visible <tool_call>",
2301            &messages,
2302            &["<tool_call>"],
2303        );
2304        assert!(matches!(hygiene, Err(ScenarioError::Contract { .. })));
2305    }
2306
2307    #[test]
2308    fn invalid_tool_diagnostics_require_rejected_call_history() {
2309        let history = vec![Message::Assistant {
2310            id: None,
2311            content: vec![tool_call(
2312                "bad_call",
2313                "missing",
2314                serde_json::json!({"value": 1}),
2315            )],
2316        }];
2317        let error = PromptError::UnknownToolCall {
2318            tool_name: "missing".to_string(),
2319            available_tools: vec!["add".to_string()],
2320            allowed_tools: Vec::new(),
2321            chat_history: Box::new(history),
2322        };
2323        assert!(validate_unknown_tool_failure(&error, "missing", &[]).is_ok());
2324        assert!(validate_unknown_tool_failure(&error, "other", &[]).is_err());
2325    }
2326}