1use 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#[derive(Debug, thiserror::Error)]
39pub enum ScenarioError {
40 #[error(transparent)]
42 Prompt(#[from] PromptError),
43 #[error(transparent)]
45 Completion(#[from] CompletionError),
46 #[error(transparent)]
48 Streaming(#[from] StreamingError),
49 #[error(transparent)]
51 Json(#[from] serde_json::Error),
52 #[error(transparent)]
54 Extraction(#[from] crate::extractor::ExtractionError),
55 #[error("{scenario} conformance failed: {details}")]
57 Contract {
58 scenario: &'static str,
60 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
74pub 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 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
119pub 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
158pub 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
191pub 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
208pub 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
233pub 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
273pub 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
293pub 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#[derive(Debug, Clone)]
317pub struct ScenarioReport {
318 pub name: &'static str,
320 pub tool_calls: usize,
322 pub prompt_tokens: u64,
324 pub generated_tokens: u64,
326 pub history_messages: usize,
328 pub duration: Duration,
330 pub response: String,
332}
333
334#[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 text: String,
825 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
958fn 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
971fn 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
990pub 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
1087pub 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
1129pub 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
1184pub 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
1238pub 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
1316pub 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
1385pub 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
1576pub 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
1650pub 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
1724pub 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
1766pub 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
1805pub 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
1899pub 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
1936pub 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
2044pub 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}