1use std::sync::{
28 Arc, Mutex,
29 atomic::{AtomicU64, Ordering},
30};
31
32use futures::StreamExt;
33use tracing::{Instrument, info_span, span::Id};
34
35use super::{
36 completion::{Agent, AgentConfig, PreparedCompletionRequest},
37 hook::{
38 AgentHook, CompletionCall, CompletionCallAction,
39 CompletionResponse as CompletionResponseEvent, HookContext, HookStack,
40 InvalidToolCallAction, ModelTurnAction, ModelTurnFinished, ObservationAction, RequestPatch,
41 ToolCall as ToolCallEvent, ToolCallAction, ToolResultAction, ToolResultEvent,
42 },
43 model::ModelHandle,
44 prompt_request::{
45 PromptResponse,
46 streaming::{
47 DriveItem, DriveStream, MultiTurnStreamItem, StreamingError, TurnSource, drive_agent,
48 drive_tool_calls, streaming_error_into_prompt,
49 },
50 tool_result_output,
51 },
52 run::{AgentRun, DEFAULT_OUTPUT_RETRIES, ModelTurn, ModelTurnOutcome, PendingToolCall},
53};
54use rig_core::{
55 memory::ConversationMemory,
56 message::{ToolCall, ToolChoice, UserContent},
57 telemetry::SpanCombinator,
58};
59
60use crate::{
61 completion::{CompletionError, CompletionModel, Document, Message, PromptError, Usage},
62 json_utils,
63 tool::{
64 ToolContext, ToolDispatch, ToolResult,
65 server::{ToolRegistrySnapshot, ToolServerHandle},
66 },
67};
68
69use super::UNKNOWN_AGENT_NAME;
70
71#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
72pub(crate) enum UnhandledInvalidToolCallPolicy {
73 #[default]
74 Fail,
75 IgnoreForExtractor,
76}
77
78macro_rules! build_chat_span {
87 ($runner:expr, $effective_preamble:expr, $name:literal, $operation:literal) => {{
88 let system_instructions = $crate::core::telemetry::system_instructions_json(
89 $effective_preamble,
90 $runner.config.record_telemetry_content,
91 );
92 $crate::core::telemetry::completion_parent_span!(
96 target: "rig::agent_chat",
97 name: $name,
98 operation: $operation,
99 system_instructions: system_instructions.as_deref(),
100 gen_ai.agent.name = $runner.agent_name_or_default(),
101 )
102 }};
103}
104pub(crate) use build_chat_span;
105
106pub(crate) fn observe_action(action: ObservationAction) -> Option<String> {
108 match action {
109 ObservationAction::Continue => None,
110 ObservationAction::Stop(reason) => Some(reason),
111 }
112}
113
114pub(crate) enum ModelTurnDecision {
116 Advance,
118 Retried,
120 Terminate(String),
122}
123
124pub(crate) fn resolve_model_turn_action(
129 run: &mut AgentRun,
130 action: ModelTurnAction,
131) -> Result<ModelTurnDecision, PromptError> {
132 match action {
133 ModelTurnAction::Continue => Ok(ModelTurnDecision::Advance),
134 ModelTurnAction::Retry(request) => {
135 run.retry_model_turn(request)?;
136 Ok(ModelTurnDecision::Retried)
137 }
138 ModelTurnAction::Stop(reason) => Ok(ModelTurnDecision::Terminate(reason)),
139 }
140}
141
142pub struct AgentRunner {
153 pub(crate) config: AgentConfig,
158 pub(crate) prompt: Message,
159 pub(crate) chat_history: Option<Vec<Message>>,
160 pub(crate) max_invalid_tool_call_retries: usize,
161 pub(crate) tool_server_handle: ToolServerHandle,
162 pub(crate) tool_context: ToolContext,
164 pub(crate) output_tool_name: Option<String>,
165 pub(crate) output_tool_description: Option<String>,
166 pub(crate) augment_output_preamble: bool,
167 pub(crate) unhandled_invalid_tool_call_policy: UnhandledInvalidToolCallPolicy,
168 pub(crate) concurrency: usize,
169 pub(crate) error_usage: Option<Arc<Mutex<Usage>>>,
170}
171
172pub(crate) type HistoryAndMemory = (
175 Option<Vec<Message>>,
176 Option<(Arc<dyn ConversationMemory>, String)>,
177);
178
179impl AgentRunner {
180 pub fn from_agent(agent: &Agent, prompt: impl Into<Message>) -> Self {
183 Self {
184 config: agent.config.clone(),
185 prompt: prompt.into(),
186 chat_history: None,
187 max_invalid_tool_call_retries: 0,
188 tool_server_handle: agent.tool_server_handle.clone(),
189 tool_context: ToolContext::new(),
190 output_tool_name: None,
191 output_tool_description: None,
192 augment_output_preamble: true,
193 unhandled_invalid_tool_call_policy: UnhandledInvalidToolCallPolicy::Fail,
194 concurrency: 1,
195 error_usage: None,
196 }
197 }
198
199 pub fn add_hook<H>(mut self, hook: H) -> Self
207 where
208 H: AgentHook + 'static,
209 {
210 self.config.hooks.push(hook);
211 self
212 }
213}
214
215impl AgentRunner {
216 pub fn max_turns(mut self, max_turns: usize) -> Self {
220 self.config.max_turns = max_turns;
221 self
222 }
223
224 pub fn using_model(mut self, model: ModelHandle) -> Self {
231 self.config.model = model;
232 self
233 }
234
235 pub fn using_model_value<M>(self, model: M) -> Self
237 where
238 M: CompletionModel + 'static,
239 {
240 self.using_model(ModelHandle::new(model))
241 }
242
243 pub fn tool_context(mut self, context: ToolContext) -> Self {
245 self.tool_context = context;
246 self
247 }
248
249 pub fn history<I, T>(mut self, history: I) -> Self
252 where
253 I: IntoIterator<Item = T>,
254 T: Into<Message>,
255 {
256 self.chat_history = Some(history.into_iter().map(Into::into).collect());
257 self
258 }
259
260 pub fn preamble(mut self, preamble: impl Into<String>) -> Self {
262 self.config.preamble = Some(preamble.into());
263 self
264 }
265
266 pub fn without_preamble(mut self) -> Self {
268 self.config.preamble = None;
269 self
270 }
271
272 pub fn document(mut self, document: Document) -> Self {
274 self.config.static_context.push(document);
275 self
276 }
277
278 pub fn documents(mut self, documents: impl IntoIterator<Item = Document>) -> Self {
280 self.config.static_context.extend(documents);
281 self
282 }
283
284 pub fn temperature(mut self, temperature: f64) -> Self {
286 self.config.temperature = Some(temperature);
287 self
288 }
289
290 pub fn without_temperature(mut self) -> Self {
292 self.config.temperature = None;
293 self
294 }
295
296 pub fn max_tokens(mut self, max_tokens: u64) -> Self {
298 self.config.max_tokens = Some(max_tokens);
299 self
300 }
301
302 pub fn without_max_tokens(mut self) -> Self {
304 self.config.max_tokens = None;
305 self
306 }
307
308 pub fn merge_additional_params(
314 mut self,
315 params: serde_json::Map<String, serde_json::Value>,
316 ) -> Self {
317 let params = serde_json::Value::Object(params);
318 self.config.additional_params = Some(match self.config.additional_params.take() {
319 Some(baseline) if baseline.is_object() => crate::json_utils::merge(baseline, params),
320 _ => params,
321 });
322 self
323 }
324
325 pub fn replace_additional_params(mut self, params: serde_json::Value) -> Self {
330 self.config.additional_params = Some(params);
331 self
332 }
333
334 pub fn without_additional_params(mut self) -> Self {
337 self.config.additional_params = None;
338 self
339 }
340
341 pub fn tool_choice(mut self, tool_choice: ToolChoice) -> Self {
343 self.config.tool_choice = Some(tool_choice);
344 self
345 }
346
347 pub fn without_tool_choice(mut self) -> Self {
349 self.config.tool_choice = None;
350 self
351 }
352
353 pub(crate) fn output_tool(
355 mut self,
356 name: impl Into<String>,
357 description: impl Into<String>,
358 augment_preamble: bool,
359 ) -> Self {
360 self.output_tool_name = Some(name.into());
361 self.output_tool_description = Some(description.into());
362 self.augment_output_preamble = augment_preamble;
363 self
364 }
365
366 pub(crate) fn ignore_unhandled_invalid_tool_calls(mut self) -> Self {
373 self.unhandled_invalid_tool_call_policy =
374 UnhandledInvalidToolCallPolicy::IgnoreForExtractor;
375 self
376 }
377
378 pub fn record_content_telemetry(mut self, enabled: bool) -> Self {
388 self.config.record_telemetry_content = enabled;
389 self
390 }
391
392 pub fn tool_concurrency(mut self, concurrency: usize) -> Self {
414 self.concurrency = concurrency.max(1);
415 self
416 }
417
418 pub fn conversation(mut self, id: impl Into<String>) -> Self {
420 self.config.conversation_id = Some(id.into());
421 self
422 }
423
424 pub fn without_memory(mut self) -> Self {
426 self.config.memory = None;
427 self.config.conversation_id = None;
428 self
429 }
430
431 pub fn max_invalid_tool_call_retries(mut self, retries: usize) -> Self {
434 self.max_invalid_tool_call_retries = retries;
435 self
436 }
437
438 pub(crate) fn agent_name_or_default(&self) -> &str {
439 self.config.name.as_deref().unwrap_or(UNKNOWN_AGENT_NAME)
440 }
441
442 pub(crate) fn build_run(&self, history_override: Option<Vec<Message>>) -> AgentRun {
447 let run = build_agent_run(
448 self.prompt.clone(),
449 self.config.max_turns,
450 self.max_invalid_tool_call_retries,
451 self.config.output_schema.as_ref(),
452 history_override.or_else(|| self.chat_history.clone()),
453 self.config.tool_choice.clone(),
454 );
455 match &self.output_tool_name {
456 Some(name) => run.with_output_tool_name(name.clone()),
457 None => run,
458 }
459 }
460}
461
462pub(crate) fn build_agent_run(
466 prompt: Message,
467 max_turns: usize,
468 max_invalid_tool_call_retries: usize,
469 output_schema: Option<&schemars::Schema>,
470 history: Option<Vec<Message>>,
471 tool_choice: Option<ToolChoice>,
472) -> AgentRun {
473 let mut run = AgentRun::new(prompt)
474 .max_turns(max_turns)
475 .max_invalid_tool_call_retries(max_invalid_tool_call_retries)
476 .with_output_validation(
477 output_schema.map(|schema| schema.as_value().clone()),
478 DEFAULT_OUTPUT_RETRIES,
479 );
480 if let Some(history) = history {
481 run = run.with_history(history);
482 }
483 if let Some(tool_choice) = tool_choice {
484 run = run.with_tool_choice(tool_choice);
485 }
486 run
487}
488
489pub(crate) fn acquire_agent_span(
497 agent_name: &str,
498 preamble: Option<&str>,
499 record_content: bool,
500) -> (tracing::Span, bool) {
501 if tracing::Span::current().is_disabled() {
502 let system_instructions =
503 rig_core::telemetry::system_instructions_json(preamble, record_content);
504 let span = info_span!(
505 "invoke_agent",
506 gen_ai.operation.name = "invoke_agent",
507 gen_ai.agent.name = agent_name,
508 gen_ai.system_instructions = system_instructions.as_deref(),
509 gen_ai.prompt = tracing::field::Empty,
510 gen_ai.completion = tracing::field::Empty,
511 gen_ai.usage.input_tokens = tracing::field::Empty,
512 gen_ai.usage.output_tokens = tracing::field::Empty,
513 gen_ai.usage.cache_read.input_tokens = tracing::field::Empty,
514 gen_ai.usage.cache_creation.input_tokens = tracing::field::Empty,
515 gen_ai.usage.tool_use_prompt_tokens = tracing::field::Empty,
516 gen_ai.usage.reasoning_tokens = tracing::field::Empty,
517 );
518 (span, true)
519 } else {
520 (tracing::Span::current(), false)
521 }
522}
523
524pub(crate) enum CompletionCallOutcome {
526 Proceed(Option<RequestPatch>),
529 Terminate(String),
531}
532
533pub(crate) async fn resolve_completion_call(
535 hooks: &HookStack,
536 ctx: &HookContext,
537 prompt: &Message,
538 history: &[Message],
539 turn: usize,
540) -> CompletionCallOutcome {
541 match hooks
542 .on_completion_call(
543 ctx,
544 CompletionCall {
545 prompt,
546 history,
547 turn,
548 },
549 )
550 .await
551 {
552 CompletionCallAction::Stop(reason) => CompletionCallOutcome::Terminate(reason),
553 CompletionCallAction::Patch(patch) => CompletionCallOutcome::Proceed(Some(patch)),
554 CompletionCallAction::Continue => CompletionCallOutcome::Proceed(None),
555 }
556}
557
558pub(crate) async fn append_run_messages(
561 memory_handle: Option<&(Arc<dyn ConversationMemory>, String)>,
562 messages: &[Message],
563) {
564 if let Some((memory, id)) = memory_handle
567 && let Err(err) = memory.append(id, messages.to_vec()).await
568 {
569 tracing::warn!(
570 error = %err,
571 conversation_id = %id,
572 "conversation memory append failed; surfacing final response anyway"
573 );
574 }
575}
576
577pub(crate) enum ToolExecution {
579 Executed(Box<ToolCall>),
586 Skipped,
590}
591
592pub(crate) struct ToolCallOutcome {
595 pub content: UserContent,
598 pub execution: ToolExecution,
600}
601
602pub(crate) async fn run_single_tool(
612 runner: &AgentRunner,
613 ctx: &HookContext,
614 tool_snapshot: &ToolRegistrySnapshot,
615 tool_call: &ToolCall,
616 internal_call_id: &str,
617 error_history: &[Message],
618) -> Result<ToolCallOutcome, PromptError> {
619 let hooks = &runner.config.hooks;
620 let tool_context = &runner.tool_context;
621 let record_content = runner.config.record_telemetry_content;
622 let tool_name = &tool_call.function.name;
623 let mut args = json_utils::serialize_json_value(&tool_call.function.arguments);
626
627 let tool_span = tracing::Span::current();
628 tool_span.record("gen_ai.tool.name", tool_name);
629 tool_span.record("gen_ai.tool.call.id", tool_call.id.as_str());
630 if record_content {
631 tool_span.record("gen_ai.tool.call.arguments", &args);
632 }
633
634 let (action, salvaged_rewrite) = hooks
641 .resolve_tool_call(
642 ctx,
643 ToolCallEvent {
644 tool_name,
645 tool_call_id: Some(tool_call.id.as_str()),
646 internal_call_id,
647 args: &args,
648 },
649 )
650 .await;
651
652 if let Some(rewritten) = salvaged_rewrite.as_ref() {
655 args = json_utils::serialize_json_value(rewritten);
656 if record_content {
657 tool_span.record("gen_ai.tool.call.arguments", &args);
658 }
659 tracing::debug!(
660 tool_name = tool_name,
661 "tool-call arguments rewritten by a hook"
662 );
663 }
664
665 let mut skipped: Option<ToolResult> = None;
671 let effective_args: serde_json::Value = match action {
672 ToolCallAction::Stop(reason) => {
673 return Err(PromptError::prompt_cancelled(
674 error_history.to_vec(),
675 reason,
676 ));
677 }
678 ToolCallAction::Skip(reason) => {
679 tracing::info!(tool_name = tool_name, reason = reason, "Tool call rejected");
680 skipped = Some(ToolResult::skipped(reason));
683 salvaged_rewrite.unwrap_or_else(|| tool_call.function.arguments.clone())
686 }
687 ToolCallAction::Rewrite(replacement) => {
688 args = json_utils::serialize_json_value(&replacement);
692 if record_content {
693 tool_span.record("gen_ai.tool.call.arguments", &args);
694 }
695 tracing::debug!(
696 tool_name = tool_name,
697 "tool-call arguments rewritten by a hook"
698 );
699 replacement
700 }
701 ToolCallAction::Run => tool_call.function.arguments.clone(),
702 };
703
704 let (exec, execution, dispatch_context) = match skipped {
708 Some(exec) => (exec, ToolExecution::Skipped, tool_context.for_dispatch()),
709 None => {
710 let mut effective_tool_call = tool_call.clone();
711 effective_tool_call.function.arguments = effective_args;
712 let ToolDispatch {
713 result: exec,
714 context: dispatch_context,
715 } = tool_snapshot.dispatch(tool_name, &args, tool_context).await;
716 (
717 exec,
718 ToolExecution::Executed(Box::new(effective_tool_call)),
719 dispatch_context,
720 )
721 }
722 };
723 let result_action = hooks
726 .on_tool_result(
727 ctx,
728 ToolResultEvent {
729 tool_name,
730 tool_call_id: Some(tool_call.id.as_str()),
731 internal_call_id,
732 args: &args,
733 presentation: exec.output(),
734 raw_result: &exec,
735 tool_context: &dispatch_context,
736 },
737 )
738 .await;
739 record_tool_result(&tool_span, &exec);
743
744 match result_action {
745 ToolResultAction::Stop(reason) => Err(PromptError::prompt_cancelled(
746 error_history.to_vec(),
747 reason,
748 )),
749 ToolResultAction::Rewrite(replacement) => {
750 if record_content {
751 tool_span.record("gen_ai.tool.call.result", replacement.render());
752 }
753 Ok(ToolCallOutcome {
754 content: tool_result_output(
755 tool_call.id.clone(),
756 tool_call.provider.clone(),
757 tool_call.function.name.clone(),
758 replacement,
759 ),
760 execution,
761 })
762 }
763 ToolResultAction::Keep => {
764 if record_content {
765 tool_span.record("gen_ai.tool.call.result", exec.output().render());
766 }
767 let content = tool_result_output(
768 tool_call.id.clone(),
769 tool_call.provider.clone(),
770 tool_call.function.name.clone(),
771 exec.output().clone(),
772 );
773 Ok(ToolCallOutcome { content, execution })
774 }
775 }
776}
777
778fn record_tool_result(span: &tracing::Span, result: &ToolResult) {
779 span.record("gen_ai.tool.call.outcome", result.status_name());
780 if let Some(error) = result.error() {
781 span.record("gen_ai.tool.error.type", error.kind().as_str());
782 }
783}
784
785pub(crate) fn new_execute_tool_span() -> tracing::Span {
791 info_span!(
792 "execute_tool",
793 gen_ai.operation.name = "execute_tool",
794 gen_ai.tool.type = "function",
795 gen_ai.tool.name = tracing::field::Empty,
796 gen_ai.tool.call.id = tracing::field::Empty,
797 gen_ai.tool.call.arguments = tracing::field::Empty,
798 gen_ai.tool.call.result = tracing::field::Empty,
799 gen_ai.tool.call.outcome = tracing::field::Empty,
800 gen_ai.tool.error.type = tracing::field::Empty
801 )
802}
803
804pub(crate) struct UnaryTurnSource {
810 current_span_id: AtomicU64,
819 record_telemetry_content: bool,
820}
821
822impl UnaryTurnSource {
823 pub(crate) fn new(record_telemetry_content: bool) -> Self {
824 Self {
825 current_span_id: AtomicU64::new(0),
826 record_telemetry_content,
827 }
828 }
829
830 fn chain_span(&self, span: tracing::Span) -> tracing::Span {
833 let span = match self.current_span_id.load(Ordering::Relaxed) {
834 0 => span,
835 id => span.follows_from(Id::from_u64(id)).to_owned(),
836 };
837 if let Some(id) = span.id() {
838 self.current_span_id.store(id.into_u64(), Ordering::Relaxed);
839 }
840 span
841 }
842}
843
844impl TurnSource for UnaryTurnSource {
845 fn open_chat_span(
846 &self,
847 runner: &AgentRunner,
848 effective_preamble: Option<&str>,
849 ) -> tracing::Span {
850 let chat_span = build_chat_span!(runner, effective_preamble, "chat", "chat");
851 self.chain_span(chat_span)
852 }
853
854 fn run_model_turn<'a>(
855 &'a mut self,
856 runner: &'a AgentRunner,
857 hook_ctx: &'a HookContext,
858 run: &'a mut AgentRun,
859 prepared: PreparedCompletionRequest,
860 chat_span: tracing::Span,
861 _agent_span: &'a tracing::Span,
862 current_prompt: Message,
863 ) -> DriveStream<'a> {
864 Box::pin(async_stream::stream! {
865 let record_accepted_turn = |run: &AgentRun| {
869 if runner.config.record_telemetry_content
870 && let Some(choice) = run.accepted_turn_choice()
871 {
872 rig_core::telemetry::record_model_output(&chat_span, &choice, true);
873 }
874 };
875
876 let attempt_max_tokens = prepared.max_tokens;
881
882 let resp = match prepared.builder.send().instrument(chat_span.clone()).await {
883 Ok(resp) => resp,
884 Err(err) => {
885 yield Err(StreamingError::from(err));
886 return;
887 }
888 };
889
890 let attempt_finish_reason = resp.finish_reason();
893
894 let mut outcome = match run.model_response(
895 ModelTurn::new(
896 resp.message_id.clone(),
897 resp.choice.clone(),
898 resp.usage,
899 prepared.executable_tool_names,
900 prepared.allowed_tool_names,
901 )
902 .with_identity(
903 resp.response_id.clone(),
904 resp.provider_request_id.clone(),
905 )
906 .with_finish_reason(attempt_finish_reason.clone())
907 .with_raw(resp.raw.clone()),
910 ) {
911 Ok(outcome) => outcome,
912 Err(err) => {
913 yield Err(Box::new(err).into());
914 return;
915 }
916 };
917
918 loop {
919 match outcome {
920 ModelTurnOutcome::NeedsResolution(context) => {
921 let action = runner
922 .config.hooks
923 .on_invalid_tool_call(hook_ctx, &context)
924 .await;
925 let resolution = match action {
926 Some(action) => run.resolve_invalid_tool_call(action),
927 None
928 if runner.unhandled_invalid_tool_call_policy
929 == UnhandledInvalidToolCallPolicy::IgnoreForExtractor =>
930 {
931 run.ignore_invalid_tool_call()
932 }
933 None => run.resolve_invalid_tool_call(InvalidToolCallAction::fail()),
934 };
935 outcome = match resolution {
936 Ok(outcome) => outcome,
937 Err(err) => {
938 yield Err(Box::new(err).into());
939 return;
940 }
941 };
942 }
943 ModelTurnOutcome::TurnRetried => break,
944 ModelTurnOutcome::Continue {
945 response_hook_suppressed,
946 } => {
947 if !response_hook_suppressed {
948 let identity = resp.identity();
959 let attempt_raw = &resp.raw;
960 if let Some(reason) = observe_action(
961 runner
962 .config.hooks
963 .on_completion_response(
964 hook_ctx,
965 CompletionResponseEvent {
966 prompt: ¤t_prompt,
967 content: &resp.choice,
968 usage: resp.usage,
969 message_id: resp.message_id.as_deref(),
970 identity: &identity,
971 raw: attempt_raw,
972 },
973 )
974 .await,
975 ) {
976 record_accepted_turn(run);
977 yield Err(StreamingError::Prompt(Box::new(run.cancel_error(reason))));
978 return;
979 }
980 let action = runner
981 .config.hooks
982 .on_model_turn_finished(
983 hook_ctx,
984 ModelTurnFinished {
985 turn: hook_ctx.turn(),
986 content: &resp.choice,
987 usage: resp.usage,
988 identity: &identity,
989 finish_reason: attempt_finish_reason.as_ref(),
990 max_tokens: attempt_max_tokens,
991 raw: attempt_raw,
992 },
993 )
994 .await;
995 match resolve_model_turn_action(run, action) {
996 Ok(ModelTurnDecision::Advance) => {}
997 Ok(ModelTurnDecision::Retried) => break,
998 Ok(ModelTurnDecision::Terminate(reason)) => {
999 record_accepted_turn(run);
1000 yield Err(StreamingError::Prompt(Box::new(
1001 run.cancel_error(reason),
1002 )));
1003 return;
1004 }
1005 Err(err) => {
1006 yield Err(StreamingError::Prompt(Box::new(err)));
1007 return;
1008 }
1009 }
1010 }
1011
1012 record_accepted_turn(run);
1013 break;
1014 }
1015 }
1016 }
1017 })
1018 }
1019
1020 fn run_tool_calls<'a>(
1021 &'a self,
1022 runner: &'a AgentRunner,
1023 hook_ctx: &'a HookContext,
1024 run: &'a mut AgentRun,
1025 calls: Vec<PendingToolCall>,
1026 tool_snapshot: Arc<ToolRegistrySnapshot>,
1027 ) -> DriveStream<'a> {
1028 drive_tool_calls(
1032 runner,
1033 hook_ctx,
1034 run,
1035 calls,
1036 tool_snapshot,
1037 |span| self.chain_span(span),
1038 false,
1039 )
1040 }
1041
1042 fn record_run_level_telemetry(
1043 &self,
1044 agent_span: &tracing::Span,
1045 response: &PromptResponse,
1046 created_agent_span: bool,
1047 ) {
1048 if created_agent_span {
1053 if self.record_telemetry_content {
1054 agent_span.record("gen_ai.completion", &response.output);
1055 }
1056 agent_span.record_token_usage(&response.usage);
1057 }
1058 }
1059
1060 fn final_item(&self, _response: &PromptResponse) -> Option<MultiTurnStreamItem> {
1061 None
1064 }
1065}
1066
1067impl AgentRunner {
1068 pub(crate) async fn run_with_error_usage(
1069 mut self,
1070 ) -> (Result<PromptResponse, PromptError>, Usage) {
1071 let usage = Arc::new(Mutex::new(Usage::new()));
1072 self.error_usage = Some(usage.clone());
1073 let result = self.run().await;
1074 let observed = result.as_ref().map_or_else(
1075 |_| *usage.lock().unwrap_or_else(|error| error.into_inner()),
1076 |response| response.usage,
1077 );
1078 (result, observed)
1079 }
1080
1081 pub(crate) fn open_agent_span(&self) -> (tracing::Span, bool) {
1084 let (agent_span, created_agent_span) = acquire_agent_span(
1085 self.agent_name_or_default(),
1086 self.config.preamble.as_deref(),
1087 self.config.record_telemetry_content,
1088 );
1089
1090 if self.config.record_telemetry_content
1091 && let Some(text) = self.prompt.rag_text()
1092 {
1093 agent_span.record("gen_ai.prompt", text);
1094 }
1095
1096 (agent_span, created_agent_span)
1097 }
1098
1099 pub(crate) async fn resolve_history_and_memory(
1106 &self,
1107 ) -> Result<HistoryAndMemory, rig_core::memory::MemoryError> {
1108 match &self.chat_history {
1109 Some(_) => Ok((None, None)),
1110 None => match (&self.config.memory, &self.config.conversation_id) {
1111 (Some(memory), Some(id)) => {
1112 let loaded = memory.load(id).await?;
1113 Ok((Some(loaded), Some((memory.clone(), id.clone()))))
1114 }
1115 _ => Ok((None, None)),
1116 },
1117 }
1118 }
1119
1120 pub async fn run(self) -> Result<PromptResponse, PromptError> {
1124 let (agent_span, created_agent_span) = self.open_agent_span();
1125 let (history_override, memory_handle) = self.resolve_history_and_memory().await?;
1126 let run = self.build_run(history_override);
1127
1128 let record_telemetry_content = self.config.record_telemetry_content;
1134 let driver = drive_agent(
1135 self,
1136 UnaryTurnSource::new(record_telemetry_content),
1137 run,
1138 agent_span,
1139 created_agent_span,
1140 memory_handle,
1141 false,
1142 );
1143 futures::pin_mut!(driver);
1144
1145 let mut response = None;
1146 while let Some(item) = driver.next().await {
1147 match item {
1148 Ok(DriveItem::Done(done)) => response = Some(*done),
1149 Ok(DriveItem::Item(_)) => {}
1150 Err(err) => return Err(streaming_error_into_prompt(err)),
1151 }
1152 }
1153
1154 response.ok_or_else(|| {
1156 PromptError::CompletionError(CompletionError::ResponseError(
1157 "agent run ended without producing a final response".to_string(),
1158 ))
1159 })
1160 }
1161}
1162
1163#[cfg(test)]
1164mod tests {
1165 use std::sync::{
1166 Arc, Mutex,
1167 atomic::{AtomicUsize, Ordering},
1168 };
1169
1170 use futures::StreamExt;
1171 use serde_json::json;
1172
1173 use crate::{
1174 agent::{AgentBuilder, AgentHook, HookContext, ToolResultAction, ToolResultEvent},
1175 completion::{CompletionModel, Document},
1176 test_utils::{MockCompletionModel, MockStreamEvent, MockTurn},
1177 tool::{Tool, ToolContext, ToolErrorKind, ToolExecutionError},
1178 };
1179 use rig_core::message::ToolChoice;
1180
1181 struct MetadataFailingTool;
1182
1183 struct SnapshotValue {
1184 value: usize,
1185 clones: Arc<AtomicUsize>,
1186 }
1187
1188 impl Clone for SnapshotValue {
1189 fn clone(&self) -> Self {
1190 self.clones.fetch_add(1, Ordering::SeqCst);
1191 Self {
1192 value: self.value,
1193 clones: self.clones.clone(),
1194 }
1195 }
1196 }
1197
1198 #[derive(Clone, Default)]
1199 struct SnapshotMutatingTool(Arc<Mutex<Vec<usize>>>);
1200
1201 impl Tool for SnapshotMutatingTool {
1202 const NAME: &'static str = "snapshot_mutator";
1203 type Error = rig::tool::ToolExecutionError;
1204 type Args = serde_json::Value;
1205 type Output = String;
1206
1207 fn description(&self) -> String {
1208 "Mutates its per-dispatch context snapshot".into()
1209 }
1210
1211 fn parameters(&self) -> serde_json::Value {
1212 json!({"type": "object", "properties": {}})
1213 }
1214
1215 async fn call(
1216 &self,
1217 context: &mut ToolContext,
1218 _args: Self::Args,
1219 ) -> Result<Self::Output, ToolExecutionError> {
1220 let initial = context.require::<SnapshotValue>()?.value;
1221 self.0.lock().expect("observed values").push(initial);
1222 let updated = {
1223 let value = context
1224 .get_mut::<SnapshotValue>()
1225 .expect("required snapshot value");
1226 value.value += 1;
1227 value.value
1228 };
1229 context.insert_result(updated);
1230 Ok(updated.to_string())
1231 }
1232 }
1233
1234 #[derive(Clone, Default)]
1235 struct SnapshotResults(Arc<Mutex<Vec<usize>>>);
1236
1237 impl AgentHook for SnapshotResults {
1238 async fn on_tool_result(
1239 &self,
1240 _ctx: &HookContext,
1241 event: ToolResultEvent<'_>,
1242 ) -> ToolResultAction {
1243 self.0.lock().expect("result values").push(
1244 *event
1245 .tool_context
1246 .require_result::<usize>()
1247 .expect("per-dispatch result metadata"),
1248 );
1249 ToolResultAction::keep()
1250 }
1251 }
1252
1253 impl Tool for MetadataFailingTool {
1254 const NAME: &'static str = "flaky_tool";
1255 type Error = rig::tool::ToolExecutionError;
1256 type Args = serde_json::Value;
1257 type Output = String;
1258
1259 fn description(&self) -> String {
1260 "Fails after attaching result metadata".into()
1261 }
1262
1263 fn parameters(&self) -> serde_json::Value {
1264 json!({"type": "object", "properties": {}})
1265 }
1266
1267 async fn call(
1268 &self,
1269 context: &mut ToolContext,
1270 _args: Self::Args,
1271 ) -> Result<Self::Output, ToolExecutionError> {
1272 context.insert_result("shared-result-metadata".to_string());
1273 Err(ToolExecutionError::timeout("raw timeout failure"))
1274 }
1275 }
1276
1277 #[derive(Clone, Default)]
1278 struct Results(Arc<Mutex<Vec<(ToolErrorKind, String, String)>>>);
1279
1280 impl AgentHook for Results {
1281 async fn on_tool_result(
1282 &self,
1283 _ctx: &HookContext,
1284 event: ToolResultEvent<'_>,
1285 ) -> ToolResultAction {
1286 if let Some(error) = event.raw_result.error() {
1287 self.0.lock().expect("results").push((
1288 error.kind(),
1289 event.raw_result.output().render(),
1290 event
1291 .tool_context
1292 .result::<String>()
1293 .expect("tool result metadata")
1294 .clone(),
1295 ));
1296 }
1297 ToolResultAction::rewrite("rewritten for model")
1298 }
1299 }
1300
1301 #[test]
1302 fn agent_exposes_read_only_name_and_description() {
1303 let named = AgentBuilder::new(MockCompletionModel::text("done"))
1304 .name("researcher")
1305 .description("Finds evidence")
1306 .build();
1307 assert_eq!(named.name(), Some("researcher"));
1308 assert_eq!(named.description(), Some("Finds evidence"));
1309
1310 let unnamed = AgentBuilder::new(MockCompletionModel::text("done")).build();
1311 assert_eq!(unnamed.name(), None);
1312 assert_eq!(unnamed.description(), None);
1313 }
1314
1315 #[tokio::test]
1316 async fn runner_applies_per_run_request_overrides() {
1317 let model = MockCompletionModel::text("done");
1318 AgentBuilder::new(model.clone())
1319 .preamble("baseline preamble")
1320 .context("baseline document")
1321 .temperature(0.1)
1322 .max_tokens(10)
1323 .additional_params(json!({"baseline": true}))
1324 .build()
1325 .runner("go")
1326 .preamble("run preamble")
1327 .document(Document {
1328 id: "run-one".into(),
1329 text: "first run document".into(),
1330 additional_props: Default::default(),
1331 })
1332 .documents([Document {
1333 id: "run-two".into(),
1334 text: "second run document".into(),
1335 additional_props: Default::default(),
1336 }])
1337 .temperature(0.7)
1338 .max_tokens(42)
1339 .replace_additional_params(json!({"override": true}))
1340 .tool_choice(ToolChoice::None)
1341 .run()
1342 .await
1343 .expect("runner request should succeed");
1344
1345 let requests = model.requests();
1346 let request = requests.first().expect("one request");
1347 assert!(request.chat_history.iter().any(
1348 |message| matches!(message, crate::completion::Message::System { content } if content == "run preamble")
1349 ));
1350 assert!(
1351 request
1352 .documents
1353 .iter()
1354 .any(|document| document.text == "baseline document")
1355 );
1356 assert!(
1357 request
1358 .documents
1359 .iter()
1360 .any(|document| document.id == "run-one")
1361 );
1362 assert!(
1363 request
1364 .documents
1365 .iter()
1366 .any(|document| document.id == "run-two")
1367 );
1368 assert_eq!(request.temperature, Some(0.7));
1369 assert_eq!(request.max_tokens, Some(42));
1370 assert_eq!(request.additional_params, Some(json!({"override": true})));
1371 assert_eq!(request.tool_choice, Some(ToolChoice::None));
1372 }
1373
1374 #[tokio::test]
1375 async fn runner_can_merge_additional_params_into_the_baseline() {
1376 let model = MockCompletionModel::text("done");
1377 AgentBuilder::new(model.clone())
1378 .additional_params(json!({"baseline": true, "winner": "baseline"}))
1379 .build()
1380 .runner("go")
1381 .merge_additional_params(
1382 json!({"override": true, "winner": "runner"})
1383 .as_object()
1384 .expect("object")
1385 .clone(),
1386 )
1387 .run()
1388 .await
1389 .expect("runner request should succeed");
1390
1391 assert_eq!(
1392 model
1393 .requests()
1394 .first()
1395 .expect("one request")
1396 .additional_params,
1397 Some(json!({"baseline": true, "override": true, "winner": "runner"}))
1398 );
1399 }
1400
1401 #[tokio::test]
1402 async fn runner_can_replace_additional_params_wholesale() {
1403 let model = MockCompletionModel::text("done");
1404 AgentBuilder::new(model.clone())
1405 .additional_params(json!({"baseline": true}))
1406 .build()
1407 .runner("go")
1408 .replace_additional_params(json!({"replacement": true}))
1409 .run()
1410 .await
1411 .expect("runner request should succeed");
1412
1413 let requests = model.requests();
1414 let request = requests.first().expect("one request");
1415 assert_eq!(
1416 request.additional_params,
1417 Some(json!({"replacement": true}))
1418 );
1419 }
1420
1421 #[tokio::test]
1422 async fn runner_can_clear_configured_request_defaults() {
1423 let model = MockCompletionModel::text("done");
1424 AgentBuilder::new(model.clone())
1425 .preamble("baseline")
1426 .temperature(0.1)
1427 .max_tokens(10)
1428 .additional_params(json!({"baseline": true}))
1429 .tool_choice(ToolChoice::Required)
1430 .build()
1431 .runner("go")
1432 .without_preamble()
1433 .without_temperature()
1434 .without_max_tokens()
1435 .without_additional_params()
1436 .without_tool_choice()
1437 .run()
1438 .await
1439 .expect("runner request should succeed");
1440
1441 let requests = model.requests();
1442 let request = requests.first().expect("one request");
1443 assert!(
1444 !request
1445 .chat_history
1446 .iter()
1447 .any(|message| matches!(message, crate::completion::Message::System { .. }))
1448 );
1449 assert_eq!(request.temperature, None);
1450 assert_eq!(request.max_tokens, None);
1451 assert_eq!(request.additional_params, None);
1452 assert_eq!(request.tool_choice, None);
1453 }
1454
1455 #[tokio::test]
1456 async fn direct_completion_model_requests_are_intentionally_hook_free() {
1457 #[derive(Clone)]
1458 struct CountCompletionCalls(Arc<AtomicUsize>);
1459
1460 impl AgentHook for CountCompletionCalls {
1461 async fn on_completion_call(
1462 &self,
1463 _ctx: &HookContext,
1464 _event: crate::agent::CompletionCallEvent<'_>,
1465 ) -> crate::agent::CompletionCallAction {
1466 self.0.fetch_add(1, Ordering::SeqCst);
1467 crate::agent::CompletionCallAction::Continue
1468 }
1469 }
1470
1471 let model = MockCompletionModel::text("raw response");
1472 let calls = Arc::new(AtomicUsize::new(0));
1473 let _agent = AgentBuilder::new(model.clone())
1474 .add_hook(CountCompletionCalls(calls.clone()))
1475 .build();
1476
1477 model
1478 .completion_request("raw request")
1479 .send()
1480 .await
1481 .expect("direct model request should succeed");
1482
1483 assert_eq!(calls.load(Ordering::SeqCst), 0);
1484 assert_eq!(model.request_count(), 1);
1485 }
1486
1487 #[tokio::test]
1488 async fn blocking_and_streaming_preserve_raw_failure_while_rewriting_presentation() {
1489 let blocking = Results::default();
1490 let blocking_model = MockCompletionModel::from_turns([
1491 MockTurn::tool_call("tc1", "flaky_tool", json!({})),
1492 MockTurn::text("done"),
1493 ]);
1494 AgentBuilder::new(blocking_model.clone())
1495 .tool(MetadataFailingTool)
1496 .add_hook(blocking.clone())
1497 .build()
1498 .runner("go")
1499 .max_turns(3)
1500 .run()
1501 .await
1502 .expect("blocking run");
1503
1504 let streaming = Results::default();
1505 let streaming_model = MockCompletionModel::from_stream_turns([
1506 vec![
1507 MockStreamEvent::tool_call_name_delta("tc1", "flaky_tool"),
1508 MockStreamEvent::tool_call_arguments_delta("tc1", "{}"),
1509 MockStreamEvent::tool_call("tc1", "flaky_tool", json!({})),
1510 MockStreamEvent::final_response_with_total_tokens(0),
1511 ],
1512 vec![
1513 MockStreamEvent::text("done"),
1514 MockStreamEvent::final_response_with_total_tokens(0),
1515 ],
1516 ]);
1517 let mut stream = AgentBuilder::new(streaming_model.clone())
1518 .tool(MetadataFailingTool)
1519 .add_hook(streaming.clone())
1520 .build()
1521 .runner("go")
1522 .max_turns(3)
1523 .stream()
1524 .await;
1525 while let Some(item) = stream.next().await {
1526 item.expect("stream item");
1527 }
1528
1529 assert_eq!(*blocking.0.lock().unwrap(), *streaming.0.lock().unwrap());
1530 assert_eq!(
1531 *blocking.0.lock().unwrap(),
1532 vec![(
1533 ToolErrorKind::Timeout,
1534 "raw timeout failure".into(),
1535 "shared-result-metadata".into()
1536 )]
1537 );
1538
1539 let blocking_history = serde_json::to_value(
1540 &blocking_model
1541 .requests()
1542 .get(1)
1543 .expect("second blocking request")
1544 .chat_history,
1545 )
1546 .unwrap();
1547 let streaming_history = serde_json::to_value(
1548 &streaming_model
1549 .requests()
1550 .get(1)
1551 .expect("second streaming request")
1552 .chat_history,
1553 )
1554 .unwrap();
1555 assert_eq!(blocking_history, streaming_history);
1556 let history = blocking_history.to_string();
1557 assert!(history.contains("rewritten for model"));
1558 assert!(!history.contains("raw timeout failure"));
1559 }
1560
1561 #[tokio::test]
1562 async fn agent_dispatch_snapshot_clones_once_and_isolates_tool_mutations() {
1563 let clones = Arc::new(AtomicUsize::new(0));
1564 let mut context = ToolContext::new();
1565 context.insert(SnapshotValue {
1566 value: 0,
1567 clones: clones.clone(),
1568 });
1569 let tool = SnapshotMutatingTool::default();
1570 let results = SnapshotResults::default();
1571
1572 AgentBuilder::new(MockCompletionModel::from_turns([
1573 MockTurn::tool_call("tc1", SnapshotMutatingTool::NAME, json!({})),
1574 MockTurn::tool_call("tc2", SnapshotMutatingTool::NAME, json!({})),
1575 MockTurn::text("done"),
1576 ]))
1577 .tool(tool.clone())
1578 .add_hook(results.clone())
1579 .build()
1580 .runner("go")
1581 .tool_context(context)
1582 .max_turns(4)
1583 .run()
1584 .await
1585 .expect("agent run");
1586
1587 assert_eq!(*tool.0.lock().expect("observed values"), vec![0, 0]);
1588 assert_eq!(*results.0.lock().expect("result values"), vec![1, 1]);
1589 assert_eq!(
1590 clones.load(Ordering::SeqCst),
1591 2,
1592 "each of the two agent dispatches should clone inbound context once"
1593 );
1594 }
1595}
1596
1597#[cfg(test)]
1598#[allow(irrefutable_let_patterns, unreachable_patterns)]
1599mod migrated_tests {
1600 use std::collections::HashMap;
1601
1602 use crate::agent::{
1603 CompletionCallAction, CompletionCallEvent, HookStack, InvalidToolCallAction,
1604 InvalidToolCallContext, ModelTurnAction, ModelTurnFinished, ObservationAction,
1605 ReasoningDelta, StreamResponseFinish, TextDelta, ToolCall, ToolCallAction, ToolCallDelta,
1606 ToolResultAction, ToolResultEvent,
1607 };
1608
1609 use std::sync::{
1610 Arc, Mutex,
1611 atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering::SeqCst},
1612 };
1613
1614 use futures::StreamExt;
1615 use serde::Deserialize;
1616 use serde_json::json;
1617 use tokio::sync::{Barrier, Notify};
1618
1619 use crate::agent::AgentBuilder;
1620 use crate::agent::hook::{AgentHook, HookContext, RequestPatch, StepEventKind};
1621 use crate::agent::prompt_request::streaming::{MultiTurnStreamItem, StreamingError};
1622 use crate::agent::run::OutputMode;
1623 use crate::completion::{
1624 CompletionError, CompletionModel, FinishReason, Message, Prompt, PromptError, Usage,
1625 };
1626 use crate::streaming::{StreamedAssistantContent, StreamedUserContent, StreamingPrompt};
1627 use crate::test_utils::{
1628 MockAddTool, MockBarrierTool, MockCompletionModel, MockOperationArgs, MockStreamEvent,
1629 MockSubtractTool, MockToolError, MockTurn, mock_final,
1630 };
1631 use crate::tool::{
1632 Tool, ToolContext, ToolExecutionError, ToolSet,
1633 server::{ToolServer, ToolServerHandle},
1634 };
1635 use rig_core::message::{
1636 AssistantContent, ToolCall as MessageToolCall, ToolChoice, ToolFunction, UserContent,
1637 };
1638 use rig_core::vector_store::{
1639 VectorSearchRequest, VectorStoreError, VectorStoreIndex, request::Filter,
1640 };
1641 use rig_core::wasm_compat::WasmCompatSend;
1642
1643 #[derive(Clone, Default)]
1646 struct RecordingHook {
1647 events: Arc<Mutex<Vec<StepEventKind>>>,
1648 tool_results: Arc<Mutex<Vec<String>>>,
1649 }
1650
1651 impl RecordingHook {
1652 fn shared_events(&self) -> Vec<StepEventKind> {
1656 self.events
1657 .lock()
1658 .expect("events lock")
1659 .iter()
1660 .copied()
1661 .filter(|kind| {
1662 matches!(
1663 kind,
1664 StepEventKind::CompletionCall
1665 | StepEventKind::ToolCall
1666 | StepEventKind::ToolResult
1667 | StepEventKind::InvalidToolCall
1668 )
1669 })
1670 .collect()
1671 }
1672
1673 fn tool_results(&self) -> Vec<String> {
1674 self.tool_results.lock().expect("results lock").clone()
1675 }
1676
1677 fn count(&self, kind: StepEventKind) -> usize {
1680 self.events
1681 .lock()
1682 .expect("events lock")
1683 .iter()
1684 .filter(|recorded| **recorded == kind)
1685 .count()
1686 }
1687 }
1688
1689 impl RecordingHook {
1690 fn record(&self, kind: StepEventKind) {
1691 self.events.lock().expect("events lock").push(kind);
1692 }
1693 }
1694
1695 impl AgentHook for RecordingHook {
1696 async fn on_completion_call(
1697 &self,
1698 _: &HookContext,
1699 _: CompletionCallEvent<'_>,
1700 ) -> CompletionCallAction {
1701 self.record(StepEventKind::CompletionCall);
1702 CompletionCallAction::continue_run()
1703 }
1704 async fn on_completion_response(
1705 &self,
1706 _: &HookContext,
1707 _: crate::agent::hook::CompletionResponse<'_>,
1708 ) -> ObservationAction {
1709 self.record(StepEventKind::CompletionResponse);
1710 ObservationAction::continue_run()
1711 }
1712 async fn on_model_turn_finished(
1713 &self,
1714 _: &HookContext,
1715 _: ModelTurnFinished<'_>,
1716 ) -> ModelTurnAction {
1717 self.record(StepEventKind::ModelTurnFinished);
1718 ModelTurnAction::continue_run()
1719 }
1720 async fn on_invalid_tool_call(
1721 &self,
1722 _: &HookContext,
1723 _: &InvalidToolCallContext,
1724 ) -> Option<InvalidToolCallAction> {
1725 self.record(StepEventKind::InvalidToolCall);
1726 None
1727 }
1728 async fn on_tool_call(&self, _: &HookContext, _: ToolCall<'_>) -> ToolCallAction {
1729 self.record(StepEventKind::ToolCall);
1730 ToolCallAction::run()
1731 }
1732 async fn on_tool_result(
1733 &self,
1734 _: &HookContext,
1735 event: ToolResultEvent<'_>,
1736 ) -> ToolResultAction {
1737 self.record(StepEventKind::ToolResult);
1738 self.tool_results
1739 .lock()
1740 .expect("results lock")
1741 .push(event.presentation.render());
1742 ToolResultAction::keep()
1743 }
1744 async fn on_text_delta(&self, _: &HookContext, _: TextDelta<'_>) -> ObservationAction {
1745 self.record(StepEventKind::TextDelta);
1746 ObservationAction::continue_run()
1747 }
1748 async fn on_reasoning_delta(
1749 &self,
1750 _: &HookContext,
1751 _: ReasoningDelta<'_>,
1752 ) -> ObservationAction {
1753 self.record(StepEventKind::ReasoningDelta);
1754 ObservationAction::continue_run()
1755 }
1756 async fn on_tool_call_delta(
1757 &self,
1758 _: &HookContext,
1759 _: ToolCallDelta<'_>,
1760 ) -> ObservationAction {
1761 self.record(StepEventKind::ToolCallDelta);
1762 ObservationAction::continue_run()
1763 }
1764 async fn on_stream_response_finish(
1765 &self,
1766 _: &HookContext,
1767 _: StreamResponseFinish<'_>,
1768 ) -> ObservationAction {
1769 self.record(StepEventKind::StreamResponseFinish);
1770 ObservationAction::continue_run()
1771 }
1772 }
1773
1774 #[derive(Clone, Debug, PartialEq)]
1775 struct CanonicalResponseSnapshot {
1776 prompt: Message,
1777 content: Vec<AssistantContent>,
1778 usage: Usage,
1779 message_id: Option<String>,
1780 }
1781
1782 #[derive(Clone, Default)]
1783 struct CanonicalResponseHook {
1784 blocking: Arc<Mutex<Vec<CanonicalResponseSnapshot>>>,
1785 streaming: Arc<Mutex<Vec<CanonicalResponseSnapshot>>>,
1786 committed: Arc<Mutex<Vec<Vec<AssistantContent>>>>,
1787 }
1788
1789 impl AgentHook for CanonicalResponseHook {
1790 async fn on_completion_response(
1791 &self,
1792 _ctx: &HookContext,
1793 event: crate::agent::hook::CompletionResponse<'_>,
1794 ) -> ObservationAction {
1795 self.blocking
1796 .lock()
1797 .expect("blocking snapshots")
1798 .push(CanonicalResponseSnapshot {
1799 prompt: event.prompt.clone(),
1800 content: event.content.clone(),
1801 usage: event.usage,
1802 message_id: event.message_id.map(str::to_owned),
1803 });
1804 ObservationAction::continue_run()
1805 }
1806
1807 async fn on_stream_response_finish(
1808 &self,
1809 _ctx: &HookContext,
1810 event: StreamResponseFinish<'_>,
1811 ) -> ObservationAction {
1812 self.streaming
1813 .lock()
1814 .expect("streaming snapshots")
1815 .push(CanonicalResponseSnapshot {
1816 prompt: event.prompt.clone(),
1817 content: event.content.clone(),
1818 usage: event.usage,
1819 message_id: event.message_id.map(str::to_owned),
1820 });
1821 ObservationAction::continue_run()
1822 }
1823
1824 async fn on_model_turn_finished(
1825 &self,
1826 _ctx: &HookContext,
1827 event: ModelTurnFinished<'_>,
1828 ) -> ModelTurnAction {
1829 self.committed
1830 .lock()
1831 .expect("committed snapshots")
1832 .push(event.content.clone());
1833 ModelTurnAction::continue_run()
1834 }
1835 }
1836
1837 #[derive(Clone, Default)]
1838 struct FinishLifecycleHook {
1839 snapshots: Arc<Mutex<Vec<CanonicalResponseSnapshot>>>,
1840 model_turns: Arc<AtomicU32>,
1841 stop: Arc<AtomicBool>,
1842 }
1843
1844 impl FinishLifecycleHook {
1845 fn stopping() -> Self {
1846 let hook = Self::default();
1847 hook.stop.store(true, SeqCst);
1848 hook
1849 }
1850 }
1851
1852 impl AgentHook for FinishLifecycleHook {
1853 async fn on_stream_response_finish(
1854 &self,
1855 _ctx: &HookContext,
1856 event: StreamResponseFinish<'_>,
1857 ) -> ObservationAction {
1858 self.snapshots
1859 .lock()
1860 .expect("finish snapshots")
1861 .push(CanonicalResponseSnapshot {
1862 prompt: event.prompt.clone(),
1863 content: event.content.clone(),
1864 usage: event.usage,
1865 message_id: event.message_id.map(str::to_owned),
1866 });
1867 if self.stop.load(SeqCst) {
1868 ObservationAction::stop("stop at stream EOF")
1869 } else {
1870 ObservationAction::continue_run()
1871 }
1872 }
1873
1874 async fn on_model_turn_finished(
1875 &self,
1876 _ctx: &HookContext,
1877 _event: ModelTurnFinished<'_>,
1878 ) -> ModelTurnAction {
1879 self.model_turns.fetch_add(1, SeqCst);
1880 ModelTurnAction::continue_run()
1881 }
1882 }
1883
1884 fn no_identity() -> &'static rig_core::completion::ResponseIdentity {
1886 static EMPTY: std::sync::OnceLock<rig_core::completion::ResponseIdentity> =
1887 std::sync::OnceLock::new();
1888 EMPTY.get_or_init(Default::default)
1889 }
1890
1891 fn canonical_usage() -> Usage {
1892 Usage {
1893 input_tokens: 11,
1894 output_tokens: 7,
1895 total_tokens: 18,
1896 ..Usage::new()
1897 }
1898 }
1899
1900 #[tokio::test]
1901 async fn blocking_completion_response_hook_receives_canonical_fields() {
1902 let hook = CanonicalResponseHook::default();
1903 let prompt = Message::user("canonical prompt");
1904 AgentBuilder::new(MockCompletionModel::new([MockTurn::text(
1905 "canonical response",
1906 )
1907 .with_usage(canonical_usage())
1908 .with_message_id("msg-canonical")]))
1909 .add_hook(hook.clone())
1910 .build()
1911 .runner(prompt.clone())
1912 .run()
1913 .await
1914 .expect("blocking response");
1915
1916 assert_eq!(
1917 *hook.blocking.lock().expect("blocking snapshots"),
1918 [CanonicalResponseSnapshot {
1919 prompt,
1920 content: vec![AssistantContent::text("canonical response")],
1921 usage: canonical_usage(),
1922 message_id: Some("msg-canonical".to_string()),
1923 }]
1924 );
1925 }
1926
1927 #[tokio::test]
1932 async fn completion_response_hook_and_calls_carry_identity_metadata() {
1933 type IdentityTriple = (Option<String>, Option<String>, Option<String>);
1934
1935 #[derive(Clone, Default)]
1936 struct IdentityHook {
1937 seen: Arc<Mutex<Vec<IdentityTriple>>>,
1938 }
1939
1940 impl AgentHook for IdentityHook {
1941 async fn on_completion_response(
1942 &self,
1943 _ctx: &HookContext,
1944 event: crate::agent::hook::CompletionResponse<'_>,
1945 ) -> ObservationAction {
1946 self.seen.lock().expect("identity snapshots").push((
1947 event.message_id.map(str::to_owned),
1948 event.identity.response_id.clone(),
1949 event.identity.provider_request_id.clone(),
1950 ));
1951 ObservationAction::continue_run()
1952 }
1953 }
1954
1955 let hook = IdentityHook::default();
1956 let response = AgentBuilder::new(MockCompletionModel::new([MockTurn::text("reply")
1957 .with_message_id("msg_1")
1958 .with_response_id("resp_1")
1959 .with_provider_request_id("req_1")]))
1960 .add_hook(hook.clone())
1961 .build()
1962 .runner(Message::user("prompt"))
1963 .run()
1964 .await
1965 .expect("blocking response");
1966
1967 assert_eq!(
1968 *hook.seen.lock().expect("identity snapshots"),
1969 [(
1970 Some("msg_1".to_string()),
1971 Some("resp_1".to_string()),
1972 Some("req_1".to_string()),
1973 )]
1974 );
1975 let call = &response.completion_calls[0];
1976 assert_eq!(call.message_id.as_deref(), Some("msg_1"));
1977 assert_eq!(call.response_id.as_deref(), Some("resp_1"));
1978 assert_eq!(call.provider_request_id.as_deref(), Some("req_1"));
1979 }
1980
1981 #[tokio::test]
1984 async fn absent_identity_metadata_stays_none() {
1985 let response = AgentBuilder::new(MockCompletionModel::new([MockTurn::text("reply")]))
1986 .build()
1987 .runner(Message::user("prompt"))
1988 .run()
1989 .await
1990 .expect("blocking response");
1991
1992 let call = &response.completion_calls[0];
1993 assert_eq!(call.message_id, None);
1994 assert_eq!(call.response_id, None);
1995 assert_eq!(call.provider_request_id, None);
1996 }
1997
1998 #[tokio::test]
2003 async fn failed_attempt_error_carries_its_own_request_id() {
2004 let hook = TurnIdentityHook::default();
2005 let error = AgentBuilder::new(MockCompletionModel::new([
2006 MockTurn::tool_call("tc1", "add", serde_json::json!({"x": 2, "y": 3}))
2007 .with_provider_request_id("req-success-1"),
2008 MockTurn::provider_response_error(
2009 http::StatusCode::TOO_MANY_REQUESTS,
2010 r#"{"error":"rate limited"}"#,
2011 "req-failed-2",
2012 ),
2013 ]))
2014 .tool(crate::test_utils::MockAddTool)
2015 .add_hook(hook.clone())
2016 .build()
2017 .runner(Message::user("add 2 and 3"))
2018 .max_turns(4)
2019 .run()
2020 .await
2021 .expect_err("the second attempt fails");
2022
2023 assert_eq!(
2024 error.provider_request_id(),
2025 Some("req-failed-2"),
2026 "the surfaced error reports the failing attempt's id: {error:?}"
2027 );
2028 let turns = hook.turns.lock().expect("turn identities").clone();
2029 assert_eq!(turns.len(), 1, "only the successful call fired the event");
2030 assert_eq!(
2031 turns[0].provider_request_id.as_deref(),
2032 Some("req-success-1"),
2033 "the success keeps its own id — no cross-attribution"
2034 );
2035 }
2036
2037 #[derive(Clone, Default)]
2041 struct TurnIdentityHook {
2042 turns: Arc<Mutex<Vec<rig_core::completion::ResponseIdentity>>>,
2043 stream_finishes: Arc<Mutex<Vec<rig_core::completion::ResponseIdentity>>>,
2044 }
2045
2046 impl AgentHook for TurnIdentityHook {
2047 async fn on_model_turn_finished(
2048 &self,
2049 _ctx: &HookContext,
2050 event: ModelTurnFinished<'_>,
2051 ) -> ModelTurnAction {
2052 self.turns
2053 .lock()
2054 .expect("turn identities")
2055 .push(event.identity.clone());
2056 ModelTurnAction::continue_run()
2057 }
2058
2059 async fn on_stream_response_finish(
2060 &self,
2061 _ctx: &HookContext,
2062 event: StreamResponseFinish<'_>,
2063 ) -> ObservationAction {
2064 self.stream_finishes
2065 .lock()
2066 .expect("stream finish identities")
2067 .push(event.identity.clone());
2068 ObservationAction::continue_run()
2069 }
2070 }
2071
2072 fn stream_final_with_ids(request_id: &str, response_id: &str) -> MockStreamEvent {
2073 MockStreamEvent::FinalResponse(
2074 rig_core::streaming::StreamFinal::new("mock", Usage::new())
2075 .with_response_id(response_id)
2076 .with_provider_request_id(request_id),
2077 )
2078 }
2079
2080 #[tokio::test]
2083 async fn model_turn_finished_identity_blocking_tool_only_and_text() {
2084 let hook = TurnIdentityHook::default();
2085 let response = AgentBuilder::new(MockCompletionModel::new([
2086 MockTurn::tool_call("tc1", "add", json!({"x": 2, "y": 3}))
2087 .with_provider_request_id("req-turn-1")
2088 .with_response_id("resp-turn-1"),
2089 MockTurn::text("5")
2090 .with_provider_request_id("req-turn-2")
2091 .with_response_id("resp-turn-2"),
2092 ]))
2093 .tool(crate::test_utils::MockAddTool)
2094 .add_hook(hook.clone())
2095 .build()
2096 .runner(Message::user("add 2 and 3"))
2097 .max_turns(3)
2098 .run()
2099 .await
2100 .expect("blocking tool run");
2101
2102 let turns = hook.turns.lock().expect("turn identities").clone();
2103 let request_ids: Vec<_> = turns
2104 .iter()
2105 .map(|identity| identity.provider_request_id.clone())
2106 .collect();
2107 assert_eq!(
2108 request_ids,
2109 [
2110 Some("req-turn-1".to_string()),
2111 Some("req-turn-2".to_string())
2112 ],
2113 "each attempt reports its own transport id, in order"
2114 );
2115 let call_ids: Vec<_> = response
2117 .completion_calls
2118 .iter()
2119 .map(|call| call.provider_request_id.clone())
2120 .collect();
2121 assert_eq!(request_ids, call_ids);
2122 }
2123
2124 #[tokio::test]
2129 async fn model_turn_finished_identity_streamed_tool_only_and_text() {
2130 let hook = TurnIdentityHook::default();
2131 let model = MockCompletionModel::from_stream_turns([
2132 vec![
2133 MockStreamEvent::tool_call_name_delta("tc1", "add"),
2134 MockStreamEvent::tool_call_arguments_delta("tc1", "{\"x\":2,\"y\":3}"),
2135 MockStreamEvent::tool_call("tc1", "add", json!({"x": 2, "y": 3})),
2136 stream_final_with_ids("req-stream-1", "resp-stream-1"),
2137 ],
2138 vec![
2139 MockStreamEvent::text("5"),
2140 stream_final_with_ids("req-stream-2", "resp-stream-2"),
2141 ],
2142 ]);
2143 let mut stream = AgentBuilder::new(model)
2144 .tool(crate::test_utils::MockAddTool)
2145 .add_hook(hook.clone())
2146 .build()
2147 .runner(Message::user("add 2 and 3"))
2148 .max_turns(3)
2149 .stream()
2150 .await;
2151 while let Some(item) = stream.next().await {
2152 item.expect("stream item");
2153 }
2154
2155 let turns = hook.turns.lock().expect("turn identities").clone();
2156 let request_ids: Vec<_> = turns
2157 .iter()
2158 .map(|identity| identity.provider_request_id.clone())
2159 .collect();
2160 assert_eq!(
2161 request_ids,
2162 [
2163 Some("req-stream-1".to_string()),
2164 Some("req-stream-2".to_string())
2165 ],
2166 "streamed tool-only and text turns each carry their own identity"
2167 );
2168 let finishes = hook.stream_finishes.lock().expect("finishes").clone();
2169 assert_eq!(
2170 finishes.len(),
2171 1,
2172 "StreamResponseFinish stays text-turn-scoped; the tool-only turn fires none"
2173 );
2174 assert_eq!(
2175 finishes[0].provider_request_id.as_deref(),
2176 Some("req-stream-2"),
2177 "the text turn's finish event carries that turn's identity"
2178 );
2179 }
2180
2181 #[tokio::test]
2184 async fn model_turn_finished_identity_streamed_reasoning_only() {
2185 let hook = TurnIdentityHook::default();
2186 let model = MockCompletionModel::from_stream_turns([vec![
2187 MockStreamEvent::reasoning("thinking quietly"),
2188 stream_final_with_ids("req-reasoning-only", "resp-reasoning-only"),
2189 ]]);
2190 let mut stream = AgentBuilder::new(model)
2191 .add_hook(hook.clone())
2192 .build()
2193 .runner(Message::user("think"))
2194 .stream()
2195 .await;
2196 while let Some(item) = stream.next().await {
2197 item.expect("stream item");
2198 }
2199
2200 let turns = hook.turns.lock().expect("turn identities").clone();
2201 assert_eq!(turns.len(), 1);
2202 assert_eq!(
2203 turns[0].provider_request_id.as_deref(),
2204 Some("req-reasoning-only")
2205 );
2206 assert_eq!(turns[0].response_id.as_deref(), Some("resp-reasoning-only"));
2207 assert!(
2208 hook.stream_finishes.lock().expect("finishes").is_empty(),
2209 "a reasoning-only turn streams no text, so no StreamResponseFinish"
2210 );
2211 }
2212
2213 #[tokio::test]
2216 async fn retried_turn_reports_the_retried_attempts_own_identity() {
2217 #[derive(Clone, Default)]
2218 struct RetryOnceCapturingIdentity {
2219 seen: Arc<Mutex<Vec<Option<String>>>>,
2220 }
2221
2222 impl AgentHook for RetryOnceCapturingIdentity {
2223 async fn on_model_turn_finished(
2224 &self,
2225 _ctx: &HookContext,
2226 event: ModelTurnFinished<'_>,
2227 ) -> ModelTurnAction {
2228 let mut seen = self.seen.lock().expect("retry identities");
2229 seen.push(event.identity.provider_request_id.clone());
2230 if seen.len() == 1 {
2231 ModelTurnAction::repeat()
2232 } else {
2233 ModelTurnAction::continue_run()
2234 }
2235 }
2236 }
2237
2238 let hook = RetryOnceCapturingIdentity::default();
2239 AgentBuilder::new(MockCompletionModel::new([
2240 MockTurn::text("first attempt").with_provider_request_id("req-attempt-1"),
2241 MockTurn::text("second attempt").with_provider_request_id("req-attempt-2"),
2242 ]))
2243 .add_hook(hook.clone())
2244 .build()
2245 .runner(Message::user("prompt"))
2246 .max_turns(3)
2247 .run()
2248 .await
2249 .expect("retried run");
2250
2251 assert_eq!(
2252 *hook.seen.lock().expect("retry identities"),
2253 [
2254 Some("req-attempt-1".to_string()),
2255 Some("req-attempt-2".to_string())
2256 ],
2257 "each attempt's event carries that attempt's id — no stale leak"
2258 );
2259 }
2260
2261 #[derive(Clone, Default)]
2278 struct RawCaptureHook {
2279 completion_responses: Arc<Mutex<Vec<serde_json::Value>>>,
2280 stream_finishes: Arc<Mutex<Vec<serde_json::Value>>>,
2281 turns: Arc<Mutex<Vec<serde_json::Value>>>,
2282 }
2283
2284 impl RawCaptureHook {
2285 fn completion_responses(&self) -> Vec<serde_json::Value> {
2286 self.completion_responses
2287 .lock()
2288 .expect("completion response raws")
2289 .clone()
2290 }
2291
2292 fn stream_finishes(&self) -> Vec<serde_json::Value> {
2293 self.stream_finishes
2294 .lock()
2295 .expect("stream finish raws")
2296 .clone()
2297 }
2298
2299 fn turns(&self) -> Vec<serde_json::Value> {
2300 self.turns.lock().expect("turn raws").clone()
2301 }
2302 }
2303
2304 impl AgentHook for RawCaptureHook {
2305 async fn on_completion_response(
2306 &self,
2307 _ctx: &HookContext,
2308 event: crate::agent::hook::CompletionResponse<'_>,
2309 ) -> ObservationAction {
2310 self.completion_responses
2311 .lock()
2312 .expect("completion response raws")
2313 .push(event.raw.clone());
2314 ObservationAction::continue_run()
2315 }
2316
2317 async fn on_stream_response_finish(
2318 &self,
2319 _ctx: &HookContext,
2320 event: StreamResponseFinish<'_>,
2321 ) -> ObservationAction {
2322 self.stream_finishes
2323 .lock()
2324 .expect("stream finish raws")
2325 .push(event.raw.clone());
2326 ObservationAction::continue_run()
2327 }
2328
2329 async fn on_model_turn_finished(
2330 &self,
2331 _ctx: &HookContext,
2332 event: ModelTurnFinished<'_>,
2333 ) -> ModelTurnAction {
2334 self.turns
2335 .lock()
2336 .expect("turn raws")
2337 .push(event.raw.clone());
2338 ModelTurnAction::continue_run()
2339 }
2340 }
2341
2342 fn raw_payload(attempt: &str) -> serde_json::Value {
2345 json!({
2346 "id": format!("resp-{attempt}"),
2347 "system_fingerprint": format!("fp-{attempt}"),
2348 "provider_only": attempt,
2349 })
2350 }
2351
2352 fn stream_final_for_attempt(
2356 attempt: &str,
2357 total_tokens: u64,
2358 ) -> rig_core::streaming::StreamFinal {
2359 let mut usage = Usage::new();
2360 usage.total_tokens = total_tokens;
2361 rig_core::streaming::StreamFinal::new("mock", usage)
2362 .with_response_id(format!("resp-{attempt}"))
2363 .with_provider_request_id(format!("req-{attempt}"))
2364 }
2365
2366 fn expected_stream_raw(terminal: &rig_core::streaming::StreamFinal) -> serde_json::Value {
2368 serde_json::to_value(terminal).expect("scripted terminal serializes")
2369 }
2370
2371 fn call_raws(calls: &[crate::agent::CompletionCall]) -> Vec<serde_json::Value> {
2373 calls.iter().map(|call| call.raw.clone()).collect()
2374 }
2375
2376 #[tokio::test]
2380 async fn hook_events_carry_raw_blocking() {
2381 let payload = raw_payload("blocking");
2382
2383 let hook = RawCaptureHook::default();
2384 let response = AgentBuilder::new(MockCompletionModel::new([
2385 MockTurn::text("reply").with_raw(payload.clone())
2386 ]))
2387 .add_hook(hook.clone())
2388 .build()
2389 .prompt("prompt")
2390 .extended_details()
2391 .await
2392 .expect("blocking response");
2393
2394 assert_eq!(hook.completion_responses(), std::slice::from_ref(&payload));
2395 assert_eq!(hook.turns(), std::slice::from_ref(&payload));
2396 assert!(
2397 hook.stream_finishes().is_empty(),
2398 "StreamResponseFinish is a streamed-surface event"
2399 );
2400 assert_eq!(call_raws(&response.completion_calls), [payload]);
2401 }
2402
2403 #[tokio::test]
2408 async fn hook_events_carry_raw_streamed() {
2409 let terminal = stream_final_for_attempt("streamed", 3);
2410 let expected = expected_stream_raw(&terminal);
2411
2412 let hook = RawCaptureHook::default();
2413 let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([vec![
2414 MockStreamEvent::text("reply"),
2415 MockStreamEvent::FinalResponse(terminal),
2416 ]]))
2417 .add_hook(hook.clone())
2418 .build()
2419 .stream_prompt("prompt")
2420 .await;
2421 let mut finals = Vec::new();
2422 let mut final_response = None;
2423 while let Some(item) = stream.next().await {
2424 match item.expect("stream item") {
2425 MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Final(
2426 final_record,
2427 )) => finals.push(final_record.raw.clone()),
2428 MultiTurnStreamItem::FinalResponse(response) => final_response = Some(response),
2429 _ => {}
2430 }
2431 }
2432
2433 assert_eq!(hook.stream_finishes(), std::slice::from_ref(&expected));
2434 assert_eq!(hook.turns(), std::slice::from_ref(&expected));
2435 assert!(
2436 hook.completion_responses().is_empty(),
2437 "CompletionResponse is a blocking-surface event"
2438 );
2439 assert_eq!(finals, std::slice::from_ref(&expected));
2440 let response = final_response.expect("run final response");
2441 assert_eq!(call_raws(&response.completion_calls), [expected]);
2442 }
2443
2444 #[tokio::test]
2449 async fn completion_calls_carry_each_attempts_own_raw_blocking() {
2450 let first = raw_payload("turn-1");
2451 let second = raw_payload("turn-2");
2452 assert_ne!(first, second);
2453
2454 let hook = RawCaptureHook::default();
2455 let response = AgentBuilder::new(MockCompletionModel::new([
2456 MockTurn::tool_call("tc1", "add", json!({"x": 2, "y": 3})).with_raw(first.clone()),
2457 MockTurn::text("5").with_raw(second.clone()),
2458 ]))
2459 .tool(crate::test_utils::MockAddTool)
2460 .add_hook(hook.clone())
2461 .build()
2462 .prompt("add 2 and 3")
2463 .extended_details()
2464 .max_turns(3)
2465 .await
2466 .expect("blocking tool run");
2467
2468 assert_eq!(
2469 call_raws(&response.completion_calls),
2470 [first.clone(), second.clone()],
2471 "each recorded call carries its own attempt's payload"
2472 );
2473 assert_eq!(hook.completion_responses(), [first.clone(), second.clone()]);
2474 assert_eq!(hook.turns(), [first, second]);
2475 }
2476
2477 #[tokio::test]
2484 async fn completion_calls_carry_each_attempts_own_raw_streamed() {
2485 let first_terminal = stream_final_for_attempt("stream-1", 1);
2486 let second_terminal = stream_final_for_attempt("stream-2", 2);
2487 let first = expected_stream_raw(&first_terminal);
2488 let second = expected_stream_raw(&second_terminal);
2489 assert_ne!(first, second);
2490
2491 let hook = RawCaptureHook::default();
2492 let model = MockCompletionModel::from_stream_turns([
2493 vec![
2494 MockStreamEvent::tool_call_name_delta("tc1", "add"),
2495 MockStreamEvent::tool_call_arguments_delta("tc1", "{\"x\":2,\"y\":3}"),
2496 MockStreamEvent::tool_call("tc1", "add", json!({"x": 2, "y": 3})),
2497 MockStreamEvent::FinalResponse(first_terminal),
2498 ],
2499 vec![
2500 MockStreamEvent::text("5"),
2501 MockStreamEvent::FinalResponse(second_terminal),
2502 ],
2503 ]);
2504 let mut stream = AgentBuilder::new(model)
2505 .tool(crate::test_utils::MockAddTool)
2506 .add_hook(hook.clone())
2507 .build()
2508 .stream_prompt("add 2 and 3")
2509 .max_turns(3)
2510 .await;
2511
2512 let mut forwarded_calls = Vec::new();
2513 let mut finals = Vec::new();
2514 let mut final_response = None;
2515 while let Some(item) = stream.next().await {
2516 match item.expect("stream item") {
2517 MultiTurnStreamItem::CompletionCall(call) => forwarded_calls.push(call.raw.clone()),
2518 MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Final(
2519 final_record,
2520 )) => finals.push(final_record.raw.clone()),
2521 MultiTurnStreamItem::FinalResponse(response) => final_response = Some(response),
2522 _ => {}
2523 }
2524 }
2525
2526 let response = final_response.expect("run final response");
2527 assert_eq!(
2528 call_raws(&response.completion_calls),
2529 [first.clone(), second.clone()],
2530 "each recorded call carries its own attempt's terminal record"
2531 );
2532 assert_eq!(
2533 forwarded_calls,
2534 [first.clone(), second.clone()],
2535 "the forwarded CompletionCall items agree with the final record"
2536 );
2537 assert_eq!(hook.turns(), [first, second.clone()]);
2538 assert_eq!(
2539 hook.stream_finishes(),
2540 std::slice::from_ref(&second),
2541 "StreamResponseFinish stays text-turn-scoped and carries that turn's payload"
2542 );
2543 assert_eq!(
2544 finals,
2545 [second],
2546 "the one forwarded Final is the final turn's, carrying its own raw"
2547 );
2548 }
2549
2550 #[derive(Clone, Default)]
2554 struct RetryOnceCapturingRaw {
2555 seen: Arc<Mutex<Vec<serde_json::Value>>>,
2556 }
2557
2558 impl AgentHook for RetryOnceCapturingRaw {
2559 async fn on_model_turn_finished(
2560 &self,
2561 _ctx: &HookContext,
2562 event: ModelTurnFinished<'_>,
2563 ) -> ModelTurnAction {
2564 let mut seen = self.seen.lock().expect("retry raws");
2565 seen.push(event.raw.clone());
2566 if seen.len() == 1 {
2567 ModelTurnAction::repeat()
2568 } else {
2569 ModelTurnAction::continue_run()
2570 }
2571 }
2572 }
2573
2574 #[tokio::test]
2578 async fn retried_turn_records_the_retried_attempts_own_raw_blocking() {
2579 let first = raw_payload("attempt-1");
2580 let second = raw_payload("attempt-2");
2581
2582 let hook = RetryOnceCapturingRaw::default();
2583 let response = AgentBuilder::new(MockCompletionModel::new([
2584 MockTurn::text("first attempt").with_raw(first.clone()),
2585 MockTurn::text("second attempt").with_raw(second.clone()),
2586 ]))
2587 .add_hook(hook.clone())
2588 .build()
2589 .prompt("prompt")
2590 .extended_details()
2591 .max_turns(3)
2592 .await
2593 .expect("retried run");
2594
2595 assert_eq!(response.output, "second attempt");
2596 assert_eq!(
2597 *hook.seen.lock().expect("retry raws"),
2598 [first.clone(), second.clone()],
2599 "each attempt's event carries that attempt's payload — no stale leak"
2600 );
2601 assert_eq!(
2602 call_raws(&response.completion_calls),
2603 [first, second],
2604 "the retried attempt's record carries the retried attempt's payload"
2605 );
2606 }
2607
2608 #[tokio::test]
2613 async fn retried_turn_records_the_retried_attempts_own_raw_streamed() {
2614 let first_terminal = stream_final_for_attempt("attempt-1", 1);
2615 let second_terminal = stream_final_for_attempt("attempt-2", 2);
2616 let first = expected_stream_raw(&first_terminal);
2617 let second = expected_stream_raw(&second_terminal);
2618
2619 let retry = RetryOnceCapturingRaw::default();
2620 let probe = RawCaptureHook::default();
2621 let model = MockCompletionModel::from_stream_turns([
2622 vec![
2623 MockStreamEvent::text("first attempt"),
2624 MockStreamEvent::FinalResponse(first_terminal),
2625 ],
2626 vec![
2627 MockStreamEvent::text("second attempt"),
2628 MockStreamEvent::FinalResponse(second_terminal),
2629 ],
2630 ]);
2631 let mut stream = AgentBuilder::new(model)
2632 .add_hook(probe.clone())
2635 .add_hook(retry.clone())
2636 .build()
2637 .stream_prompt("prompt")
2638 .max_turns(3)
2639 .await;
2640
2641 let mut finals = Vec::new();
2642 let mut final_response = None;
2643 while let Some(item) = stream.next().await {
2644 match item.expect("stream item") {
2645 MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Final(
2646 final_record,
2647 )) => finals.push(final_record.raw.clone()),
2648 MultiTurnStreamItem::FinalResponse(response) => final_response = Some(response),
2649 _ => {}
2650 }
2651 }
2652
2653 let response = final_response.expect("run final response");
2654 assert_eq!(response.output, "second attempt");
2655 assert_eq!(
2656 *retry.seen.lock().expect("retry raws"),
2657 [first.clone(), second.clone()],
2658 "each attempt's event carries that attempt's terminal record"
2659 );
2660 assert_eq!(probe.turns(), [first.clone(), second.clone()]);
2661 assert_eq!(
2662 probe.stream_finishes(),
2663 [first.clone(), second.clone()],
2664 "each attempt's finish event carries its own terminal record"
2665 );
2666 assert_eq!(
2667 call_raws(&response.completion_calls),
2668 [first, second.clone()],
2669 "the retried attempt's record carries the retried attempt's terminal record"
2670 );
2671 assert_eq!(
2672 finals,
2673 [second],
2674 "the rejected attempt's Final is suppressed; the accepted one carries its own raw"
2675 );
2676 }
2677
2678 #[tokio::test]
2679 async fn response_scoped_id_is_not_promoted_into_history() {
2680 let prompt = Message::user("prompt");
2681 let response = AgentBuilder::new(MockCompletionModel::new([
2682 MockTurn::text("reply").with_response_id("chatcmpl-123")
2683 ]))
2684 .build()
2685 .runner(prompt)
2686 .run()
2687 .await
2688 .expect("blocking response");
2689
2690 let messages = response.messages.expect("history enabled");
2691 let assistant_ids: Vec<_> = messages
2692 .iter()
2693 .filter_map(|message| match message {
2694 Message::Assistant { id, .. } => Some(id.clone()),
2695 _ => None,
2696 })
2697 .collect();
2698 assert_eq!(assistant_ids, [None]);
2699 }
2700
2701 #[tokio::test]
2702 async fn message_id_is_promoted_into_history() {
2703 let prompt = Message::user("prompt");
2704 let response = AgentBuilder::new(MockCompletionModel::new([
2705 MockTurn::text("reply").with_message_id("msg_abc")
2706 ]))
2707 .build()
2708 .runner(prompt)
2709 .run()
2710 .await
2711 .expect("blocking response");
2712
2713 let messages = response.messages.expect("history enabled");
2714 let assistant_ids: Vec<_> = messages
2715 .iter()
2716 .filter_map(|message| match message {
2717 Message::Assistant { id, .. } => Some(id.clone()),
2718 _ => None,
2719 })
2720 .collect();
2721 assert_eq!(assistant_ids, [Some("msg_abc".to_string())]);
2722 }
2723
2724 #[tokio::test]
2725 async fn streaming_response_finish_matches_blocking_canonical_fields() {
2726 let prompt = Message::user("canonical prompt");
2727 let blocking_hook = CanonicalResponseHook::default();
2728 AgentBuilder::new(MockCompletionModel::new([MockTurn::text(
2729 "canonical response",
2730 )
2731 .with_usage(canonical_usage())
2732 .with_message_id("msg-canonical")]))
2733 .add_hook(blocking_hook.clone())
2734 .build()
2735 .runner(prompt.clone())
2736 .run()
2737 .await
2738 .expect("blocking response");
2739
2740 let streaming_hook = CanonicalResponseHook::default();
2741 let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([[
2742 MockStreamEvent::text("canonical response"),
2743 MockStreamEvent::final_response(canonical_usage()),
2744 MockStreamEvent::message_id("msg-canonical"),
2745 ]]))
2746 .add_hook(streaming_hook.clone())
2747 .build()
2748 .runner(prompt)
2749 .stream()
2750 .await;
2751 while let Some(item) = stream.next().await {
2752 item.expect("stream item");
2753 }
2754
2755 let blocking = blocking_hook
2756 .blocking
2757 .lock()
2758 .expect("blocking snapshots")
2759 .clone();
2760 let streaming = streaming_hook
2761 .streaming
2762 .lock()
2763 .expect("streaming snapshots")
2764 .clone();
2765 assert_eq!(streaming, blocking);
2766 assert_eq!(streaming[0].usage, canonical_usage());
2767 assert_eq!(streaming[0].message_id.as_deref(), Some("msg-canonical"));
2768 }
2769
2770 #[tokio::test]
2771 async fn streaming_response_finish_without_provider_message_id_reports_none() {
2772 let hook = FinishLifecycleHook::default();
2773 let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([[
2774 MockStreamEvent::text("canonical response"),
2775 MockStreamEvent::final_response(canonical_usage()),
2776 ]]))
2777 .add_hook(hook.clone())
2778 .build()
2779 .runner("canonical prompt")
2780 .stream()
2781 .await;
2782 while let Some(item) = stream.next().await {
2783 item.expect("stream item");
2784 }
2785
2786 let snapshots = hook.snapshots.lock().expect("finish snapshots");
2787 assert_eq!(snapshots.len(), 1);
2788 assert_eq!(snapshots[0].message_id, None);
2789 }
2790
2791 #[tokio::test]
2792 async fn streaming_response_finish_runs_before_buffered_final_is_exposed() {
2793 let hook = FinishLifecycleHook::default();
2794 let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([[
2795 MockStreamEvent::text("canonical response"),
2796 MockStreamEvent::final_response(canonical_usage()),
2797 MockStreamEvent::message_id("msg-after-final"),
2798 ]]))
2799 .add_hook(hook.clone())
2800 .build()
2801 .runner("canonical prompt")
2802 .stream()
2803 .await;
2804 let mut provider_finals = 0;
2805 while let Some(item) = stream.next().await {
2806 if matches!(
2807 item.expect("stream item"),
2808 MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Final(_))
2809 ) {
2810 provider_finals += 1;
2811 let snapshots = hook.snapshots.lock().expect("finish snapshots");
2812 assert_eq!(snapshots.len(), 1, "hook must run before final exposure");
2813 assert_eq!(snapshots[0].message_id.as_deref(), Some("msg-after-final"));
2814 assert_eq!(
2815 hook.model_turns.load(SeqCst),
2816 1,
2817 "the canonical turn hook must accept the turn before final exposure"
2818 );
2819 }
2820 }
2821
2822 assert_eq!(provider_finals, 1);
2823 assert_eq!(hook.snapshots.lock().expect("finish snapshots").len(), 1);
2824 assert_eq!(hook.model_turns.load(SeqCst), 1);
2825 }
2826
2827 #[tokio::test]
2828 async fn streaming_response_finish_stop_suppresses_final_and_turn_commit() {
2829 let hook = FinishLifecycleHook::stopping();
2830 let prompt = Message::user("canonical prompt");
2831 let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([[
2832 MockStreamEvent::text("canonical response"),
2833 MockStreamEvent::final_response(canonical_usage()),
2834 MockStreamEvent::message_id("msg-after-final"),
2835 ]]))
2836 .add_hook(hook.clone())
2837 .build()
2838 .runner(prompt.clone())
2839 .stream()
2840 .await;
2841 let mut saw_provider_final = false;
2842 let mut saw_run_final = false;
2843 let mut error = None;
2844 while let Some(item) = stream.next().await {
2845 match item {
2846 Ok(MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Final(
2847 _,
2848 ))) => saw_provider_final = true,
2849 Ok(MultiTurnStreamItem::FinalResponse(_)) => saw_run_final = true,
2850 Ok(_) => {}
2851 Err(err) => error = Some(err),
2852 }
2853 }
2854
2855 assert!(!saw_provider_final, "the buffered final must remain hidden");
2856 assert!(
2857 !saw_run_final,
2858 "the cancelled run must not produce a response"
2859 );
2860 assert_eq!(hook.snapshots.lock().expect("finish snapshots").len(), 1);
2861 assert_eq!(hook.model_turns.load(SeqCst), 0);
2862 assert!(matches!(
2863 error,
2864 Some(StreamingError::Prompt(error))
2865 if matches!(
2866 error.as_ref(),
2867 PromptError::PromptCancelled { chat_history, reason }
2868 if chat_history == &[prompt] && reason == "stop at stream EOF"
2869 )
2870 ));
2871 }
2872
2873 struct StopCompletedModelTurn;
2874
2875 impl AgentHook for StopCompletedModelTurn {
2876 async fn on_model_turn_finished(
2877 &self,
2878 _ctx: &HookContext,
2879 _event: ModelTurnFinished<'_>,
2880 ) -> ModelTurnAction {
2881 ModelTurnAction::stop("stop completed model turn")
2882 }
2883 }
2884
2885 #[tokio::test]
2886 async fn streaming_model_turn_stop_preserves_completed_provider_final() {
2887 let prompt = Message::user("canonical prompt");
2888 let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([[
2889 MockStreamEvent::text("canonical response"),
2890 MockStreamEvent::final_response(canonical_usage()),
2891 ]]))
2892 .add_hook(StopCompletedModelTurn)
2893 .build()
2894 .runner(prompt.clone())
2895 .stream()
2896 .await;
2897
2898 let mut provider_finals = 0;
2899 let mut saw_retry = false;
2900 let mut saw_run_final = false;
2901 let mut error = None;
2902 while let Some(item) = stream.next().await {
2903 match item {
2904 Ok(MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Final(
2905 _,
2906 ))) => provider_finals += 1,
2907 Ok(MultiTurnStreamItem::ModelTurnRetried { .. }) => saw_retry = true,
2908 Ok(MultiTurnStreamItem::FinalResponse(_)) => saw_run_final = true,
2909 Ok(_) => {}
2910 Err(err) => error = Some(err),
2911 }
2912 }
2913
2914 assert_eq!(provider_finals, 1);
2915 assert!(!saw_retry);
2916 assert!(!saw_run_final);
2917 assert!(matches!(
2918 error,
2919 Some(StreamingError::Prompt(error))
2920 if matches!(
2921 error.as_ref(),
2922 PromptError::PromptCancelled { reason, .. }
2923 if reason == "stop completed model turn"
2924 )
2925 ));
2926 }
2927
2928 #[tokio::test]
2929 async fn provider_error_after_final_suppresses_finish_hook_and_buffered_final() {
2930 let hook = FinishLifecycleHook::default();
2931 let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([[
2932 MockStreamEvent::text("canonical response"),
2933 MockStreamEvent::final_response(canonical_usage()),
2934 MockStreamEvent::error("post-final failure"),
2935 ]]))
2936 .add_hook(hook.clone())
2937 .build()
2938 .runner("canonical prompt")
2939 .stream()
2940 .await;
2941 let mut saw_provider_final = false;
2942 let mut error = None;
2943 while let Some(item) = stream.next().await {
2944 match item {
2945 Ok(MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Final(
2946 _,
2947 ))) => saw_provider_final = true,
2948 Ok(_) => {}
2949 Err(err) => error = Some(err),
2950 }
2951 }
2952
2953 assert!(!saw_provider_final, "the buffered final must remain hidden");
2954 assert!(hook.snapshots.lock().expect("finish snapshots").is_empty());
2955 assert_eq!(hook.model_turns.load(SeqCst), 0);
2956 assert!(matches!(
2957 error,
2958 Some(StreamingError::Completion(CompletionError::ProviderError(message)))
2959 if message == "post-final failure"
2960 ));
2961 }
2962
2963 #[tokio::test]
2964 async fn visible_assistant_items_after_final_are_rejected() {
2965 let cases = [
2966 ("text", MockStreamEvent::text("late text")),
2967 ("reasoning", MockStreamEvent::reasoning("late reasoning")),
2968 (
2969 "reasoning delta",
2970 MockStreamEvent::reasoning_delta("late reasoning"),
2971 ),
2972 (
2973 "tool call",
2974 MockStreamEvent::tool_call("late", "add", json!({"x": 1, "y": 2})),
2975 ),
2976 (
2977 "tool-call delta",
2978 MockStreamEvent::tool_call_name_delta("late", "add"),
2979 ),
2980 ("unknown", MockStreamEvent::unknown(json!({"type": "late"}))),
2981 ];
2982
2983 for (case, visible_item) in cases {
2984 let hook = FinishLifecycleHook::default();
2985 let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([vec![
2986 MockStreamEvent::text("canonical response"),
2987 MockStreamEvent::final_response(canonical_usage()),
2988 visible_item,
2989 ]]))
2990 .add_hook(hook.clone())
2991 .build()
2992 .runner("canonical prompt")
2993 .stream()
2994 .await;
2995 let mut saw_provider_final = false;
2996 let mut error = None;
2997 while let Some(item) = stream.next().await {
2998 match item {
2999 Ok(MultiTurnStreamItem::StreamAssistantItem(
3000 StreamedAssistantContent::Final(_),
3001 )) => saw_provider_final = true,
3002 Ok(_) => {}
3003 Err(err) => error = Some(err),
3004 }
3005 }
3006
3007 assert!(
3008 !saw_provider_final,
3009 "{case}: buffered final must remain hidden"
3010 );
3011 assert!(
3012 hook.snapshots.lock().expect("finish snapshots").is_empty(),
3013 "{case}: finish hook must not run"
3014 );
3015 assert_eq!(hook.model_turns.load(SeqCst), 0, "{case}");
3016 assert!(
3017 matches!(
3018 error,
3019 Some(StreamingError::Completion(CompletionError::ResponseError(ref message)))
3020 if message.contains("visible assistant content after its final response")
3021 ),
3022 "{case}: expected malformed-response error, got {error:?}"
3023 );
3024 }
3025 }
3026
3027 #[tokio::test]
3028 async fn visible_item_after_non_emittable_final_is_rejected() {
3029 let hook = FinishLifecycleHook::default();
3030 let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([[
3031 MockStreamEvent::reasoning("think"),
3032 MockStreamEvent::final_response(canonical_usage()),
3033 MockStreamEvent::text("late text"),
3034 ]]))
3035 .add_hook(hook.clone())
3036 .build()
3037 .runner("canonical prompt")
3038 .stream()
3039 .await;
3040 let mut error = None;
3041 while let Some(item) = stream.next().await {
3042 if let Err(err) = item {
3043 error = Some(err);
3044 }
3045 }
3046
3047 assert!(hook.snapshots.lock().expect("finish snapshots").is_empty());
3048 assert_eq!(hook.model_turns.load(SeqCst), 0);
3049 assert!(matches!(
3050 error,
3051 Some(StreamingError::Completion(CompletionError::ResponseError(message)))
3052 if message.contains("visible assistant content after its final response")
3053 ));
3054 }
3055
3056 #[tokio::test]
3057 async fn streaming_response_finish_normalizes_interleaved_content() {
3058 let hook = CanonicalResponseHook::default();
3059 let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([
3060 vec![
3061 MockStreamEvent::reasoning("think"),
3062 MockStreamEvent::tool_call("tc1", "add", json!({"x": 2, "y": 3})),
3063 MockStreamEvent::text("answer"),
3064 MockStreamEvent::final_response_with_total_tokens(0),
3065 ],
3066 vec![
3067 MockStreamEvent::text("done"),
3068 MockStreamEvent::final_response_with_total_tokens(0),
3069 ],
3070 ]))
3071 .tool(MockAddTool)
3072 .add_hook(hook.clone())
3073 .build()
3074 .runner("go")
3075 .max_turns(3)
3076 .stream()
3077 .await;
3078 while let Some(item) = stream.next().await {
3079 item.expect("stream item");
3080 }
3081
3082 let snapshots = hook.streaming.lock().expect("streaming snapshots");
3083 let committed = hook.committed.lock().expect("committed snapshots");
3084 let kinds = snapshots[0]
3085 .content
3086 .iter()
3087 .map(|content| match content {
3088 AssistantContent::Reasoning(_) => "reasoning",
3089 AssistantContent::Text(_) => "text",
3090 AssistantContent::ToolCall(_) => "tool_call",
3091 _ => "other",
3092 })
3093 .collect::<Vec<_>>();
3094 assert_eq!(kinds, ["reasoning", "text", "tool_call"]);
3095 assert_eq!(
3096 snapshots[0].content, committed[0],
3097 "finish hook and committed turn must share one canonical choice"
3098 );
3099 }
3100
3101 fn blocking_model() -> MockCompletionModel {
3102 MockCompletionModel::from_turns([
3103 MockTurn::tool_call("tc1", "add", json!({"x": 2, "y": 3})),
3104 MockTurn::text("the answer is 5"),
3105 ])
3106 }
3107
3108 fn streaming_model() -> MockCompletionModel {
3113 MockCompletionModel::from_stream_turns([
3114 vec![
3115 MockStreamEvent::tool_call_name_delta("tc1", "add"),
3116 MockStreamEvent::tool_call_arguments_delta("tc1", "{\"x\":2,\"y\":3}"),
3117 MockStreamEvent::tool_call("tc1", "add", json!({"x": 2, "y": 3})),
3118 MockStreamEvent::final_response_with_total_tokens(0),
3119 ],
3120 vec![
3121 MockStreamEvent::text("the answer is 5"),
3122 MockStreamEvent::final_response_with_total_tokens(0),
3123 ],
3124 ])
3125 }
3126
3127 #[tokio::test]
3138 async fn streamed_tool_call_items_share_one_internal_call_id() {
3139 let mut stream = AgentBuilder::new(streaming_model())
3140 .tool(MockAddTool)
3141 .build()
3142 .runner("add 2 and 3")
3143 .max_turns(2)
3144 .stream()
3145 .await;
3146
3147 let mut delta_ids = Vec::new();
3148 let mut completed_ids = Vec::new();
3149 let mut executed_ids = Vec::new();
3150 let mut result_ids = Vec::new();
3151 while let Some(item) = stream.next().await {
3152 match item.expect("stream item") {
3153 MultiTurnStreamItem::StreamAssistantItem(
3154 StreamedAssistantContent::ToolCallDelta {
3155 internal_call_id, ..
3156 },
3157 ) => delta_ids.push(internal_call_id),
3158 MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::ToolCall {
3159 internal_call_id,
3160 ..
3161 }) => completed_ids.push(internal_call_id),
3162 MultiTurnStreamItem::ToolExecutionCommitted {
3163 internal_call_id, ..
3164 } => executed_ids.push(internal_call_id),
3165 MultiTurnStreamItem::StreamUserItem(StreamedUserContent::ToolResult {
3166 internal_call_id,
3167 ..
3168 }) => result_ids.push(internal_call_id),
3169 _ => {}
3170 }
3171 }
3172
3173 assert_eq!(delta_ids.len(), 2, "one name delta and one argument delta");
3174 let correlated = delta_ids.first().expect("a delta id").clone();
3175 assert!(
3176 delta_ids.iter().all(|id| *id == correlated),
3177 "the fragments of one call share one id: {delta_ids:?}"
3178 );
3179 assert_eq!(completed_ids, vec![correlated.clone()]);
3180 assert_eq!(executed_ids, vec![correlated.clone()]);
3181 assert_eq!(result_ids, vec![correlated]);
3182 }
3183
3184 #[tokio::test]
3187 async fn from_agent_preserves_implicit_one_and_explicit_zero_budgets() {
3188 let implicit_model = blocking_model();
3189 let implicit_recorded = implicit_model.clone();
3190 let implicit_agent = AgentBuilder::new(implicit_model).tool(MockAddTool).build();
3191 let implicit_runner = super::AgentRunner::from_agent(&implicit_agent, "add 2 and 3");
3192 assert_eq!(implicit_runner.config.max_turns, 1);
3193
3194 let implicit_err = implicit_runner
3195 .run()
3196 .await
3197 .expect_err("implicit budget should reject the second model call");
3198 assert!(matches!(
3199 implicit_err,
3200 PromptError::MaxTurnsError { max_turns: 1, .. }
3201 ));
3202 assert_eq!(implicit_recorded.request_count(), 1);
3203
3204 let zero_model = MockCompletionModel::text("should not be requested");
3205 let zero_recorded = zero_model.clone();
3206 let zero_agent = AgentBuilder::new(zero_model).default_max_turns(0).build();
3207 let zero_runner = super::AgentRunner::from_agent(&zero_agent, "do not call");
3208 assert_eq!(zero_runner.config.max_turns, 0);
3209
3210 let zero_err = zero_runner
3211 .run()
3212 .await
3213 .expect_err("explicit zero budget should reject the initial model call");
3214 assert!(matches!(
3215 zero_err,
3216 PromptError::MaxTurnsError { max_turns: 0, .. }
3217 ));
3218 assert_eq!(zero_recorded.request_count(), 0);
3219 }
3220
3221 #[tokio::test]
3225 async fn per_run_overrides_do_not_mutate_the_source_agent() {
3226 let agent = AgentBuilder::new(MockCompletionModel::text("ok"))
3227 .name("original")
3228 .preamble("original preamble")
3229 .temperature(0.2)
3230 .default_max_turns(2)
3231 .build();
3232
3233 let overridden = agent
3234 .runner("prompt")
3235 .max_turns(7)
3236 .preamble("overridden preamble")
3237 .temperature(0.9)
3238 .max_tokens(123)
3239 .tool_choice(rig_core::message::ToolChoice::None)
3240 .conversation("per-run-conversation");
3241 assert_eq!(overridden.config.max_turns, 7);
3242 assert_eq!(
3243 overridden.config.preamble.as_deref(),
3244 Some("overridden preamble")
3245 );
3246
3247 assert_eq!(agent.config.max_turns, 2);
3249 assert_eq!(agent.config.preamble.as_deref(), Some("original preamble"));
3250 assert_eq!(agent.config.temperature, Some(0.2));
3251 assert_eq!(agent.config.max_tokens, None);
3252 assert!(agent.config.tool_choice.is_none());
3253 assert!(agent.config.conversation_id.is_none());
3254
3255 let fresh = agent.runner("another prompt");
3257 assert_eq!(fresh.config.max_turns, 2);
3258 assert_eq!(fresh.config.preamble.as_deref(), Some("original preamble"));
3259 assert_eq!(fresh.config.temperature, Some(0.2));
3260 assert!(fresh.config.conversation_id.is_none());
3261 }
3262
3263 #[tokio::test]
3266 async fn prompt_surfaces_reject_second_tool_roundtrip_request_at_budget_one() {
3267 let blocking_model = blocking_model();
3268 let blocking_recorded = blocking_model.clone();
3269 let blocking_agent = AgentBuilder::new(blocking_model).tool(MockAddTool).build();
3270 let blocking_err = blocking_agent
3271 .prompt("add 2 and 3")
3272 .max_turns(1)
3273 .await
3274 .expect_err("blocking prompt should reject request two");
3275 assert!(matches!(
3276 blocking_err,
3277 PromptError::MaxTurnsError { max_turns: 1, .. }
3278 ));
3279 assert_eq!(blocking_recorded.request_count(), 1);
3280
3281 let streaming_model = streaming_model();
3282 let streaming_recorded = streaming_model.clone();
3283 let streaming_agent = AgentBuilder::new(streaming_model).tool(MockAddTool).build();
3284 let mut stream = streaming_agent
3285 .stream_prompt("add 2 and 3")
3286 .max_turns(1)
3287 .await;
3288 let mut streaming_err = None;
3289 while let Some(item) = stream.next().await {
3290 if let Err(err) = item {
3291 streaming_err = Some(err);
3292 break;
3293 }
3294 }
3295 match streaming_err {
3296 Some(StreamingError::Prompt(err)) => assert!(matches!(
3297 *err,
3298 PromptError::MaxTurnsError { max_turns: 1, .. }
3299 )),
3300 other => panic!("expected streaming max-turns error, got {other:?}"),
3301 }
3302 assert_eq!(streaming_recorded.request_count(), 1);
3303 }
3304
3305 #[tokio::test]
3309 async fn run_and_stream_behave_identically_for_a_tool_call() {
3310 let blocking_hook = RecordingHook::default();
3311 let blocking = AgentBuilder::new(blocking_model())
3312 .tool(MockAddTool)
3313 .build()
3314 .runner("add 2 and 3")
3315 .max_turns(2)
3316 .add_hook(blocking_hook.clone())
3317 .run()
3318 .await
3319 .expect("blocking run should succeed");
3320
3321 let streaming_hook = RecordingHook::default();
3324 let mut stream = AgentBuilder::new(streaming_model())
3325 .tool(MockAddTool)
3326 .build()
3327 .runner("add 2 and 3")
3328 .max_turns(2)
3329 .add_hook(streaming_hook.clone())
3330 .stream()
3331 .await;
3332
3333 let mut final_response = None;
3334 while let Some(item) = stream.next().await {
3335 if let Ok(MultiTurnStreamItem::FinalResponse(resp)) =
3336 item.map_err(|err| panic!("stream item errored: {err}"))
3337 {
3338 final_response = Some(resp);
3339 }
3340 }
3341 let final_response = final_response.expect("stream should yield a final response");
3342
3343 assert_eq!(blocking.output, "the answer is 5");
3345 assert_eq!(final_response.output(), blocking.output);
3346
3347 assert_eq!(
3350 blocking_hook.shared_events(),
3351 streaming_hook.shared_events()
3352 );
3353 assert_eq!(
3354 blocking_hook.shared_events(),
3355 vec![
3356 StepEventKind::CompletionCall,
3357 StepEventKind::ToolCall,
3358 StepEventKind::ToolResult,
3359 StepEventKind::CompletionCall,
3360 ]
3361 );
3362
3363 assert_eq!(blocking_hook.tool_results(), streaming_hook.tool_results());
3365 assert_eq!(blocking_hook.tool_results(), vec!["5".to_string()]);
3366
3367 let blocking_messages = blocking.messages.expect("blocking messages");
3369 let streaming_messages = final_response
3370 .messages()
3371 .expect("streaming history")
3372 .to_vec();
3373 assert_eq!(
3374 serde_json::to_value(&blocking_messages).expect("serialize blocking"),
3375 serde_json::to_value(&streaming_messages).expect("serialize streaming"),
3376 );
3377 }
3378
3379 mod structured_tool_results {
3384 use std::sync::{Arc, Mutex};
3385
3386 use futures::StreamExt;
3387 use serde_json::json;
3388
3389 use crate::agent::{
3390 AgentBuilder, AgentHook, HookContext, HookStack, ToolCall, ToolCallAction,
3391 ToolResultAction, ToolResultEvent,
3392 };
3393 use crate::test_utils::{
3394 MockAddTool, MockCompletionModel, MockDeniedTool, MockFailingTool,
3395 MockHandledFailureTool, MockMetadataTool, MockRequestId, MockStreamEvent, MockTurn,
3396 };
3397 use crate::tool::{ToolErrorKind, ToolResult};
3398
3399 #[derive(Clone, Default)]
3402 struct OutcomeHook {
3403 outcomes: Arc<Mutex<Vec<String>>>,
3404 results: Arc<Mutex<Vec<String>>>,
3405 }
3406
3407 impl OutcomeHook {
3408 fn outcomes(&self) -> Vec<String> {
3409 self.outcomes.lock().expect("outcomes").clone()
3410 }
3411
3412 fn results(&self) -> Vec<String> {
3413 self.results.lock().expect("results").clone()
3414 }
3415 }
3416
3417 fn outcome_label(result: &ToolResult) -> String {
3419 if result.is_skipped() {
3420 "skipped".to_string()
3421 } else if result.is_refused() {
3422 "denied".to_string()
3423 } else if let Some(error) = result.error() {
3424 format!("error:{}", error.kind().as_str())
3425 } else {
3426 "success".to_string()
3427 }
3428 }
3429
3430 impl AgentHook for OutcomeHook {
3431 async fn on_tool_result(
3432 &self,
3433 _ctx: &HookContext,
3434 event: ToolResultEvent<'_>,
3435 ) -> ToolResultAction {
3436 if let ToolResultEvent {
3437 presentation,
3438 raw_result,
3439 ..
3440 } = event
3441 {
3442 self.outcomes
3443 .lock()
3444 .expect("outcomes")
3445 .push(outcome_label(raw_result));
3446 self.results
3447 .lock()
3448 .expect("results")
3449 .push(presentation.render());
3450 }
3451 ToolResultAction::keep()
3452 }
3453 }
3454
3455 fn model_one_tool_then_text(tool: &str) -> MockCompletionModel {
3457 MockCompletionModel::from_turns([
3458 MockTurn::tool_call("tc1", tool, json!({})),
3459 MockTurn::text("done"),
3460 ])
3461 }
3462
3463 fn stream_model_one_tool_then_text(tool: &str) -> MockCompletionModel {
3465 MockCompletionModel::from_stream_turns([
3466 vec![
3467 MockStreamEvent::tool_call_name_delta("tc1", tool),
3468 MockStreamEvent::tool_call_arguments_delta("tc1", "{}"),
3469 MockStreamEvent::tool_call("tc1", tool, json!({})),
3470 MockStreamEvent::final_response_with_total_tokens(0),
3471 ],
3472 vec![
3473 MockStreamEvent::text("done"),
3474 MockStreamEvent::final_response_with_total_tokens(0),
3475 ],
3476 ])
3477 }
3478
3479 #[tokio::test]
3482 async fn timeout_failure_surfaces_structured_outcome() {
3483 let hook = OutcomeHook::default();
3484 AgentBuilder::new(model_one_tool_then_text("flaky_tool"))
3485 .tool(MockFailingTool::new(ToolErrorKind::Timeout))
3486 .add_hook(hook.clone())
3487 .build()
3488 .runner("go")
3489 .max_turns(3)
3490 .run()
3491 .await
3492 .expect("run should succeed; a tool timeout is model-visible feedback, not fatal");
3493
3494 assert_eq!(hook.outcomes(), vec!["error:timeout".to_string()]);
3495 assert_eq!(hook.results(), vec!["mock tool call failed".to_string()]);
3497 }
3498
3499 #[tokio::test]
3502 async fn hook_terminates_after_repeated_timeouts() {
3503 #[derive(Clone, Default)]
3504 struct TimeoutCount(usize);
3505
3506 struct TimeoutTerminator;
3507 impl AgentHook for TimeoutTerminator {
3508 async fn on_tool_result(
3509 &self,
3510 ctx: &HookContext,
3511 event: ToolResultEvent<'_>,
3512 ) -> ToolResultAction {
3513 if let ToolResultEvent { raw_result, .. } = event
3514 && raw_result.is_error_kind(ToolErrorKind::Timeout)
3515 {
3516 let count = ctx.scratchpad().update(|c: &mut TimeoutCount| {
3517 c.0 += 1;
3518 c.0
3519 });
3520 if count >= 2 {
3521 return ToolResultAction::stop("aborting after repeated tool timeouts");
3522 }
3523 }
3524 ToolResultAction::keep()
3525 }
3526 }
3527
3528 let observer = OutcomeHook::default();
3529 let err = AgentBuilder::new(MockCompletionModel::from_turns([
3530 MockTurn::tool_call("tc1", "flaky_tool", json!({})),
3531 MockTurn::tool_call("tc2", "flaky_tool", json!({})),
3532 MockTurn::text("unreachable"),
3533 ]))
3534 .tool(MockFailingTool::new(ToolErrorKind::Timeout))
3535 .add_hook(observer.clone())
3537 .add_hook(TimeoutTerminator)
3538 .build()
3539 .runner("go")
3540 .max_turns(5)
3541 .run()
3542 .await
3543 .expect_err("the run must terminate after two timeouts");
3544
3545 assert!(
3546 err.to_string()
3547 .contains("aborting after repeated tool timeouts"),
3548 "unexpected error: {err}"
3549 );
3550 assert_eq!(
3551 observer.outcomes(),
3552 vec!["error:timeout".to_string(), "error:timeout".to_string()],
3553 "both timeout outcomes must be observed before termination"
3554 );
3555 }
3556
3557 #[tokio::test]
3560 async fn not_found_outcome_is_structured_and_non_fatal() {
3561 let hook = OutcomeHook::default();
3562 let status: Arc<Mutex<Option<u16>>> = Arc::new(Mutex::new(None));
3563
3564 struct StatusProbe(Arc<Mutex<Option<u16>>>);
3565 impl AgentHook for StatusProbe {
3566 async fn on_tool_result(
3567 &self,
3568 _ctx: &HookContext,
3569 event: ToolResultEvent<'_>,
3570 ) -> ToolResultAction {
3571 if let Some(error) = event.raw_result.error() {
3572 *self.0.lock().expect("status") = error.http_status();
3573 }
3574 ToolResultAction::keep()
3575 }
3576 }
3577
3578 AgentBuilder::new(model_one_tool_then_text("flaky_tool"))
3579 .tool(MockFailingTool::new(ToolErrorKind::NotFound))
3580 .add_hook(hook.clone())
3581 .add_hook(StatusProbe(status.clone()))
3582 .build()
3583 .runner("go")
3584 .max_turns(3)
3585 .run()
3586 .await
3587 .expect("a 404 must not terminate the run by default");
3588
3589 assert_eq!(hook.outcomes(), vec!["error:not_found".to_string()]);
3590 assert_eq!(
3591 *status.lock().expect("status"),
3592 Some(404),
3593 "the structured failure must carry the HTTP status"
3594 );
3595 }
3596
3597 #[tokio::test]
3600 async fn handled_failure_delivers_model_output_and_error_outcome() {
3601 let hook = OutcomeHook::default();
3602 AgentBuilder::new(model_one_tool_then_text("lookup"))
3603 .tool(MockHandledFailureTool)
3604 .add_hook(hook.clone())
3605 .build()
3606 .runner("go")
3607 .max_turns(3)
3608 .run()
3609 .await
3610 .expect("a handled failure is not fatal");
3611
3612 assert_eq!(hook.outcomes(), vec!["error:not_found".to_string()]);
3613 assert_eq!(
3614 hook.results(),
3615 vec!["no record found for id 42; try a different id".to_string()],
3616 "the tool's model-visible output must survive alongside the error outcome"
3617 );
3618 }
3619
3620 #[tokio::test]
3623 async fn flow_skip_produces_skipped_outcome() {
3624 struct SkipHook;
3625 impl AgentHook for SkipHook {
3626 async fn on_tool_call(
3627 &self,
3628 _ctx: &HookContext,
3629 event: ToolCall<'_>,
3630 ) -> ToolCallAction {
3631 if let ToolCall { .. } = event {
3632 ToolCallAction::skip("not executed (denied by policy); do not retry")
3633 } else {
3634 ToolCallAction::run()
3635 }
3636 }
3637 }
3638
3639 let observer = OutcomeHook::default();
3640 AgentBuilder::new(model_one_tool_then_text("flaky_tool"))
3641 .tool(MockFailingTool::new(ToolErrorKind::Timeout))
3642 .add_hook(SkipHook)
3643 .add_hook(observer.clone())
3644 .build()
3645 .runner("go")
3646 .max_turns(3)
3647 .run()
3648 .await
3649 .expect("run should succeed after skipping the tool");
3650
3651 assert_eq!(observer.outcomes(), vec!["skipped".to_string()]);
3652 assert_eq!(
3653 observer.results(),
3654 vec!["not executed (denied by policy); do not retry".to_string()]
3655 );
3656 }
3657
3658 #[tokio::test]
3663 async fn tool_authored_denial_produces_denied_outcome() {
3664 let hook = OutcomeHook::default();
3665 AgentBuilder::new(model_one_tool_then_text("guarded"))
3666 .tool(MockDeniedTool)
3667 .add_hook(hook.clone())
3668 .build()
3669 .runner("go")
3670 .max_turns(3)
3671 .run()
3672 .await
3673 .expect("a tool-authored denial is not fatal");
3674
3675 assert_eq!(hook.outcomes(), vec!["denied".to_string()]);
3676 assert_eq!(
3677 hook.results(),
3678 vec!["access to this resource is not permitted".to_string()],
3679 "the model still receives the tool's denial message"
3680 );
3681 }
3682
3683 #[tokio::test]
3684 async fn permission_denied_failure_is_not_a_tool_refusal() {
3685 let hook = OutcomeHook::default();
3686 AgentBuilder::new(model_one_tool_then_text("flaky_tool"))
3687 .tool(MockFailingTool::new(ToolErrorKind::PermissionDenied))
3688 .add_hook(hook.clone())
3689 .build()
3690 .runner("go")
3691 .max_turns(3)
3692 .run()
3693 .await
3694 .expect("a permission failure is model-visible feedback, not fatal");
3695
3696 assert_eq!(hook.outcomes(), vec!["error:permission_denied".to_string()]);
3697 assert_eq!(hook.results(), vec!["mock tool call failed".to_string()]);
3698 }
3699
3700 #[tokio::test]
3706 async fn rewrite_args_then_skip_reports_rewritten_args() {
3707 struct RewriteHook;
3709 impl AgentHook for RewriteHook {
3710 async fn on_tool_call(
3711 &self,
3712 _ctx: &HookContext,
3713 event: ToolCall<'_>,
3714 ) -> ToolCallAction {
3715 if let ToolCall { .. } = event {
3716 ToolCallAction::rewrite(json!({ "x": 41, "y": 1 }))
3717 } else {
3718 ToolCallAction::run()
3719 }
3720 }
3721 }
3722 struct SkipHook;
3724 impl AgentHook for SkipHook {
3725 async fn on_tool_call(
3726 &self,
3727 _ctx: &HookContext,
3728 event: ToolCall<'_>,
3729 ) -> ToolCallAction {
3730 if let ToolCall { .. } = event {
3731 ToolCallAction::skip("denied after rewrite")
3732 } else {
3733 ToolCallAction::run()
3734 }
3735 }
3736 }
3737 #[derive(Clone, Default)]
3739 struct ArgsProbe {
3740 args: Arc<Mutex<Option<String>>>,
3741 outcome: Arc<Mutex<Option<String>>>,
3742 }
3743 impl AgentHook for ArgsProbe {
3744 async fn on_tool_result(
3745 &self,
3746 _ctx: &HookContext,
3747 event: ToolResultEvent<'_>,
3748 ) -> ToolResultAction {
3749 if let ToolResultEvent {
3750 args, raw_result, ..
3751 } = event
3752 {
3753 *self.args.lock().expect("args") = Some(args.to_string());
3754 *self.outcome.lock().expect("outcome") = Some(outcome_label(raw_result));
3755 }
3756 ToolResultAction::keep()
3757 }
3758 }
3759
3760 async fn run_surface(streaming: bool) -> (String, String) {
3761 let probe = ArgsProbe::default();
3762 if streaming {
3765 let mut stream = AgentBuilder::new(stream_model_one_tool_then_text("add"))
3766 .tool(MockAddTool)
3767 .add_hook(RewriteHook)
3768 .add_hook(SkipHook)
3769 .add_hook(probe.clone())
3770 .build()
3771 .runner("go")
3772 .max_turns(3)
3773 .stream()
3774 .await;
3775 while let Some(item) = stream.next().await {
3776 if let Err(err) = item {
3777 panic!("stream item errored: {err}");
3778 }
3779 }
3780 } else {
3781 AgentBuilder::new(model_one_tool_then_text("add"))
3782 .tool(MockAddTool)
3783 .add_hook(RewriteHook)
3784 .add_hook(SkipHook)
3785 .add_hook(probe.clone())
3786 .build()
3787 .runner("go")
3788 .max_turns(3)
3789 .run()
3790 .await
3791 .expect("run should succeed after skipping the tool");
3792 }
3793 let args = probe.args.lock().expect("args").clone().expect("args seen");
3794 let outcome = probe
3795 .outcome
3796 .lock()
3797 .expect("outcome")
3798 .clone()
3799 .expect("outcome seen");
3800 (args, outcome)
3801 }
3802
3803 for streaming in [false, true] {
3804 let (args, outcome) = run_surface(streaming).await;
3805 assert_eq!(
3806 outcome, "skipped",
3807 "the skipped tool must produce a Skipped outcome (streaming={streaming})"
3808 );
3809 let parsed: serde_json::Value =
3810 serde_json::from_str(&args).expect("ToolResult args are valid JSON");
3811 assert_eq!(
3812 parsed,
3813 json!({ "x": 41, "y": 1 }),
3814 "the skipped ToolResult must report the rewritten args, not the model's \
3815 original {{}} (streaming={streaming}); got {args}"
3816 );
3817 }
3818 }
3819
3820 #[tokio::test]
3825 async fn nested_hook_stack_rewrite_then_skip_reports_rewritten_args() {
3826 struct RewriteHook;
3827 impl AgentHook for RewriteHook {
3828 async fn on_tool_call(
3829 &self,
3830 _ctx: &HookContext,
3831 event: ToolCall<'_>,
3832 ) -> ToolCallAction {
3833 if let ToolCall { .. } = event {
3834 ToolCallAction::rewrite(json!({ "x": 41, "y": 1 }))
3835 } else {
3836 ToolCallAction::run()
3837 }
3838 }
3839 }
3840 struct SkipHook;
3841 impl AgentHook for SkipHook {
3842 async fn on_tool_call(
3843 &self,
3844 _ctx: &HookContext,
3845 event: ToolCall<'_>,
3846 ) -> ToolCallAction {
3847 if let ToolCall { .. } = event {
3848 ToolCallAction::skip("denied after nested rewrite")
3849 } else {
3850 ToolCallAction::run()
3851 }
3852 }
3853 }
3854 #[derive(Clone, Default)]
3855 struct ArgsProbe {
3856 args: Arc<Mutex<Option<String>>>,
3857 outcome: Arc<Mutex<Option<String>>>,
3858 }
3859 impl AgentHook for ArgsProbe {
3860 async fn on_tool_result(
3861 &self,
3862 _ctx: &HookContext,
3863 event: ToolResultEvent<'_>,
3864 ) -> ToolResultAction {
3865 if let ToolResultEvent {
3866 args, raw_result, ..
3867 } = event
3868 {
3869 *self.args.lock().expect("args") = Some(args.to_string());
3870 *self.outcome.lock().expect("outcome") = Some(outcome_label(raw_result));
3871 }
3872 ToolResultAction::keep()
3873 }
3874 }
3875
3876 fn nested_stack() -> HookStack {
3878 let mut nested = HookStack::new();
3879 nested.push(RewriteHook);
3880 nested.push(SkipHook);
3881 nested
3882 }
3883
3884 for streaming in [false, true] {
3887 let probe = ArgsProbe::default();
3888 if streaming {
3889 let mut stream = AgentBuilder::new(stream_model_one_tool_then_text("add"))
3890 .tool(MockAddTool)
3891 .add_hook(nested_stack())
3892 .add_hook(probe.clone())
3893 .build()
3894 .runner("go")
3895 .max_turns(3)
3896 .stream()
3897 .await;
3898 while let Some(item) = stream.next().await {
3899 if let Err(err) = item {
3900 panic!("stream item errored: {err}");
3901 }
3902 }
3903 } else {
3904 AgentBuilder::new(model_one_tool_then_text("add"))
3905 .tool(MockAddTool)
3906 .add_hook(nested_stack())
3907 .add_hook(probe.clone())
3908 .build()
3909 .runner("go")
3910 .max_turns(3)
3911 .run()
3912 .await
3913 .expect("run should succeed after the nested stack skips the tool");
3914 }
3915
3916 assert_eq!(
3917 probe.outcome.lock().expect("outcome").clone(),
3918 Some("skipped".to_string()),
3919 "streaming={streaming}"
3920 );
3921 let args = probe.args.lock().expect("args").clone().expect("args seen");
3922 let parsed: serde_json::Value =
3923 serde_json::from_str(&args).expect("valid JSON args");
3924 assert_eq!(
3925 parsed,
3926 json!({ "x": 41, "y": 1 }),
3927 "the nested stack's rewrite must survive its skip and reach the ToolResult \
3928 (streaming={streaming}); got {args}"
3929 );
3930 }
3931 }
3932
3933 #[tokio::test]
3936 async fn invalid_args_are_classified_as_invalid_args() {
3937 let hook = OutcomeHook::default();
3938 AgentBuilder::new(MockCompletionModel::from_turns([
3939 MockTurn::tool_call("tc1", "add", json!({ "x": "not-a-number", "y": 1 })),
3941 MockTurn::text("done"),
3942 ]))
3943 .tool(MockAddTool)
3944 .add_hook(hook.clone())
3945 .build()
3946 .runner("go")
3947 .max_turns(3)
3948 .run()
3949 .await
3950 .expect("an invalid-args failure is model-visible feedback, not fatal");
3951
3952 assert_eq!(hook.outcomes(), vec!["error:invalid_args".to_string()]);
3953 }
3954
3955 #[tokio::test]
3958 async fn success_result_metadata_reaches_hook_but_not_model() {
3959 struct MetadataProbe {
3960 seen: Arc<Mutex<Option<String>>>,
3961 model_output: Arc<Mutex<Option<String>>>,
3962 }
3963 impl AgentHook for MetadataProbe {
3964 async fn on_tool_result(
3965 &self,
3966 _ctx: &HookContext,
3967 event: ToolResultEvent<'_>,
3968 ) -> ToolResultAction {
3969 if let ToolResultEvent {
3970 presentation,
3971 tool_context,
3972 ..
3973 } = event
3974 {
3975 *self.seen.lock().expect("seen") = tool_context
3976 .result::<MockRequestId>()
3977 .map(|id| id.0.clone());
3978 *self.model_output.lock().expect("model_output") =
3979 Some(presentation.render());
3980 }
3981 ToolResultAction::keep()
3982 }
3983 }
3984
3985 async fn run_surface(streaming: bool) -> (Option<String>, String) {
3986 let seen: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));
3987 let model_output: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));
3988 let probe = MetadataProbe {
3989 seen: seen.clone(),
3990 model_output: model_output.clone(),
3991 };
3992
3993 if streaming {
3994 let mut stream =
3995 AgentBuilder::new(stream_model_one_tool_then_text("with_meta"))
3996 .tool(MockMetadataTool)
3997 .add_hook(probe)
3998 .build()
3999 .runner("go")
4000 .max_turns(3)
4001 .stream()
4002 .await;
4003 while let Some(item) = stream.next().await {
4004 if let Err(error) = item {
4005 panic!("stream item errored: {error}");
4006 }
4007 }
4008 } else {
4009 AgentBuilder::new(model_one_tool_then_text("with_meta"))
4010 .tool(MockMetadataTool)
4011 .add_hook(probe)
4012 .build()
4013 .runner("go")
4014 .max_turns(3)
4015 .run()
4016 .await
4017 .expect("run should succeed");
4018 }
4019
4020 let seen_value = seen.lock().expect("seen").clone();
4021 let output = model_output
4022 .lock()
4023 .expect("model_output")
4024 .clone()
4025 .expect("output");
4026 (seen_value, output)
4027 }
4028
4029 for streaming in [false, true] {
4030 let (seen, output) = run_surface(streaming).await;
4031 assert_eq!(
4032 seen,
4033 Some("req-7".to_string()),
4034 "the tool's result metadata must reach the hook (streaming={streaming})"
4035 );
4036 assert_eq!(output, "done");
4037 assert!(
4038 !output.contains("req-7"),
4039 "result metadata must never leak into model output (streaming={streaming})"
4040 );
4041 }
4042 }
4043
4044 #[tokio::test]
4048 async fn rewrite_result_does_not_mask_the_structured_outcome() {
4049 struct Redact;
4050 impl AgentHook for Redact {
4051 async fn on_tool_result(
4052 &self,
4053 _ctx: &HookContext,
4054 event: ToolResultEvent<'_>,
4055 ) -> ToolResultAction {
4056 if let ToolResultEvent { .. } = event {
4057 ToolResultAction::rewrite("[REDACTED]")
4058 } else {
4059 ToolResultAction::keep()
4060 }
4061 }
4062 }
4063
4064 let observer = OutcomeHook::default();
4065 AgentBuilder::new(model_one_tool_then_text("flaky_tool"))
4066 .tool(MockFailingTool::new(ToolErrorKind::NotFound))
4067 .add_hook(Redact)
4070 .add_hook(observer.clone())
4071 .build()
4072 .runner("go")
4073 .max_turns(3)
4074 .run()
4075 .await
4076 .expect("run should succeed");
4077
4078 assert_eq!(observer.outcomes(), vec!["error:not_found".to_string()]);
4079 assert_eq!(observer.results(), vec!["[REDACTED]".to_string()]);
4080 }
4081
4082 #[tokio::test]
4085 async fn streaming_and_blocking_outcomes_match() {
4086 let blocking = OutcomeHook::default();
4087 AgentBuilder::new(model_one_tool_then_text("flaky_tool"))
4088 .tool(MockFailingTool::new(ToolErrorKind::Timeout))
4089 .add_hook(blocking.clone())
4090 .build()
4091 .runner("go")
4092 .max_turns(3)
4093 .run()
4094 .await
4095 .expect("blocking run should succeed");
4096
4097 let streaming = OutcomeHook::default();
4098 let mut stream = AgentBuilder::new(stream_model_one_tool_then_text("flaky_tool"))
4099 .tool(MockFailingTool::new(ToolErrorKind::Timeout))
4100 .add_hook(streaming.clone())
4101 .build()
4102 .runner("go")
4103 .max_turns(3)
4104 .stream()
4105 .await;
4106 while let Some(item) = stream.next().await {
4107 if let Err(err) = item {
4108 panic!("stream item errored: {err}");
4109 }
4110 }
4111
4112 assert_eq!(blocking.outcomes(), vec!["error:timeout".to_string()]);
4113 assert_eq!(blocking.outcomes(), streaming.outcomes());
4114 assert_eq!(blocking.results(), streaming.results());
4115 }
4116
4117 #[tokio::test]
4120 async fn concurrent_tools_preserve_order_and_both_outcomes() {
4121 use rig_core::message::{
4122 AssistantContent, ToolCall as MessageToolCall, ToolFunction, UserContent,
4123 };
4124
4125 let turn = MockTurn::from_contents([
4126 AssistantContent::ToolCall(MessageToolCall::from_wire(
4127 "tc_add",
4128 ToolFunction::new("add".to_string(), json!({ "x": 2, "y": 3 })),
4129 )),
4130 AssistantContent::ToolCall(MessageToolCall::from_wire(
4131 "tc_flaky",
4132 ToolFunction::new("flaky_tool".to_string(), json!({})),
4133 )),
4134 ]);
4135
4136 let observer = OutcomeHook::default();
4137 let response = AgentBuilder::new(MockCompletionModel::from_turns([
4138 turn,
4139 MockTurn::text("done"),
4140 ]))
4141 .tool(MockAddTool)
4142 .tool(MockFailingTool::new(ToolErrorKind::Timeout))
4143 .add_hook(observer.clone())
4144 .build()
4145 .runner("go")
4146 .max_turns(3)
4147 .tool_concurrency(2)
4148 .run()
4149 .await
4150 .expect("run should succeed");
4151
4152 let mut outcomes = observer.outcomes();
4154 outcomes.sort();
4155 assert_eq!(
4156 outcomes,
4157 vec!["error:timeout".to_string(), "success".to_string()]
4158 );
4159
4160 let messages = response.messages.expect("messages");
4163 let tool_result_ids: Vec<String> = messages
4164 .iter()
4165 .flat_map(|message| match message {
4166 crate::completion::Message::User { content } => content
4167 .iter()
4168 .filter_map(|c| match c {
4169 UserContent::ToolResult(result) => Some(result.call.to_string()),
4170 _ => None,
4171 })
4172 .collect::<Vec<_>>(),
4173 _ => Vec::new(),
4174 })
4175 .collect();
4176 assert_eq!(
4177 tool_result_ids,
4178 vec!["tc_add".to_string(), "tc_flaky".to_string()],
4179 "tool results must be persisted in call order"
4180 );
4181 }
4182 }
4183
4184 mod span_safety_net {
4190 use std::collections::{HashMap, HashSet};
4191 use std::sync::{Arc, Mutex};
4192
4193 use futures::StreamExt;
4194 use tracing::Instrument;
4195 use tracing::field::{Field, Visit};
4196 use tracing::span::{Attributes, Record};
4197 use tracing::{Id, Subscriber};
4198 use tracing_subscriber::layer::{Context, SubscriberExt};
4199 use tracing_subscriber::{Layer, Registry, registry::LookupSpan};
4200
4201 use crate::agent::{
4202 AgentBuilder, HookContext, MultiTurnStreamItem, ToolResultAction, ToolResultEvent,
4203 };
4204 use crate::completion::{
4205 CompletionError, CompletionModel, CompletionRequest, CompletionResponse, Prompt,
4206 PromptError, Usage,
4207 };
4208 use crate::streaming::StreamedAssistantContent;
4209 use crate::streaming::StreamingCompletionResponse;
4210 use crate::test_utils::{MockAddTool, MockCompletionModel, MockStreamEvent, MockTurn};
4211 use crate::tool::{ToolContext, ToolExecutionError};
4212 use rig_core::telemetry::{CompletionOperation, CompletionSpanBuilder};
4213
4214 use super::{BoundedResponseRetry, StopCompletedModelTurn, TestRetryMode};
4215
4216 #[derive(Clone)]
4217 struct CapturedSpan {
4218 id: u64,
4219 name: String,
4220 target: String,
4221 field_names: HashSet<String>,
4222 u64_fields: HashMap<String, u64>,
4223 string_fields: HashMap<String, Vec<String>>,
4224 }
4225
4226 #[derive(Clone, Default)]
4227 struct Captured {
4228 spans: Arc<Mutex<Vec<CapturedSpan>>>,
4229 follows: Arc<Mutex<Vec<(u64, u64)>>>,
4231 }
4232
4233 impl Captured {
4234 fn insert(&self, id: &Id, name: &str, target: &str) {
4235 self.spans.lock().expect("spans").push(CapturedSpan {
4236 id: id.into_u64(),
4237 name: name.to_string(),
4238 target: target.to_string(),
4239 field_names: HashSet::new(),
4240 u64_fields: HashMap::new(),
4241 string_fields: HashMap::new(),
4242 });
4243 }
4244
4245 fn record(
4246 &self,
4247 id: &Id,
4248 names: HashSet<String>,
4249 u64s: HashMap<String, u64>,
4250 strings: HashMap<String, String>,
4251 ) {
4252 let id = id.into_u64();
4253 if let Ok(mut spans) = self.spans.lock()
4254 && let Some(span) = spans.iter_mut().find(|s| s.id == id)
4255 {
4256 span.field_names.extend(names);
4257 span.u64_fields.extend(u64s);
4258 for (name, value) in strings {
4259 span.string_fields.entry(name).or_default().push(value);
4260 }
4261 }
4262 }
4263
4264 fn follows_from(&self, span: &Id, follows: &Id) {
4265 self.follows
4266 .lock()
4267 .expect("follows")
4268 .push((span.into_u64(), follows.into_u64()));
4269 }
4270
4271 fn clear(&self) {
4272 self.spans.lock().expect("spans").clear();
4273 self.follows.lock().expect("follows").clear();
4274 }
4275
4276 fn snapshot(&self) -> Vec<CapturedSpan> {
4277 self.spans.lock().expect("spans").clone()
4278 }
4279
4280 fn follows_edges(&self) -> Vec<(u64, u64)> {
4281 self.follows.lock().expect("follows").clone()
4282 }
4283 }
4284
4285 struct CaptureLayer {
4286 captured: Captured,
4287 }
4288
4289 impl<S> Layer<S> for CaptureLayer
4290 where
4291 S: Subscriber + for<'l> LookupSpan<'l>,
4292 {
4293 fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, _ctx: Context<'_, S>) {
4294 self.captured
4295 .insert(id, attrs.metadata().name(), attrs.metadata().target());
4296 }
4297
4298 fn on_record(&self, span: &Id, values: &Record<'_>, _ctx: Context<'_, S>) {
4299 let mut visitor = FieldVisitor::default();
4300 values.record(&mut visitor);
4301 self.captured
4302 .record(span, visitor.names, visitor.u64s, visitor.strings);
4303 }
4304
4305 fn on_follows_from(&self, span: &Id, follows: &Id, _ctx: Context<'_, S>) {
4306 self.captured.follows_from(span, follows);
4307 }
4308 }
4309
4310 #[derive(Default)]
4311 struct FieldVisitor {
4312 names: HashSet<String>,
4313 u64s: HashMap<String, u64>,
4314 strings: HashMap<String, String>,
4315 }
4316
4317 impl Visit for FieldVisitor {
4318 fn record_u64(&mut self, field: &Field, value: u64) {
4319 self.names.insert(field.name().to_string());
4320 self.u64s.insert(field.name().to_string(), value);
4321 }
4322
4323 fn record_str(&mut self, field: &Field, value: &str) {
4324 self.names.insert(field.name().to_string());
4325 self.strings
4326 .insert(field.name().to_string(), value.to_string());
4327 }
4328
4329 fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
4330 self.names.insert(field.name().to_string());
4331 self.strings
4332 .insert(field.name().to_string(), format!("{value:?}"));
4333 }
4334 }
4335
4336 fn usage(input: u64, output: u64) -> Usage {
4337 Usage {
4338 input_tokens: input,
4339 output_tokens: output,
4340 ..Usage::new()
4341 }
4342 }
4343
4344 fn tool_then_text_model() -> MockCompletionModel {
4347 MockCompletionModel::from_turns([
4348 MockTurn::tool_call("tc1", "add", serde_json::json!({"x": 2, "y": 3}))
4349 .with_usage(usage(7, 11)),
4350 MockTurn::text("the answer is 5").with_usage(usage(13, 17)),
4351 ])
4352 }
4353
4354 #[derive(Clone)]
4355 struct CompletionTelemetryModel {
4356 inner: MockCompletionModel,
4357 }
4358
4359 impl CompletionModel for CompletionTelemetryModel {
4360 async fn completion(
4361 &self,
4362 request: CompletionRequest,
4363 ) -> Result<CompletionResponse, CompletionError> {
4364 let span = CompletionSpanBuilder::new(
4365 "fixture-provider",
4366 "fixture-model",
4367 CompletionOperation::Chat,
4368 )
4369 .build();
4370 self.inner.completion(request).instrument(span).await
4371 }
4372
4373 async fn stream(
4374 &self,
4375 request: CompletionRequest,
4376 ) -> Result<StreamingCompletionResponse, CompletionError> {
4377 let span = CompletionSpanBuilder::new(
4378 "fixture-provider",
4379 "fixture-model",
4380 CompletionOperation::ChatStreaming,
4381 )
4382 .build();
4383 self.inner.stream(request).instrument(span).await
4384 }
4385 }
4386
4387 async fn warm_blocking_callsites() {
4392 let agent = AgentBuilder::new(tool_then_text_model())
4393 .record_content_telemetry(true)
4394 .tool(MockAddTool)
4395 .build();
4396 let _ = agent.runner("add 2 and 3").max_turns(3).run().await;
4397 }
4398
4399 async fn run_blocking_response_retry_with_content_telemetry() {
4400 AgentBuilder::new(MockCompletionModel::from_turns([
4401 MockTurn::text("rejected"),
4402 MockTurn::text("accepted"),
4403 ]))
4404 .record_content_telemetry(true)
4405 .add_hook(BoundedResponseRetry::new(
4406 "rejected",
4407 1,
4408 TestRetryMode::Repeat,
4409 ))
4410 .build()
4411 .runner("question")
4412 .max_turns(2)
4413 .run()
4414 .await
4415 .expect("blocking retry should succeed");
4416 }
4417
4418 async fn run_streaming_response_retry_with_content_telemetry() {
4419 let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([
4420 [
4421 MockStreamEvent::text("rejected"),
4422 MockStreamEvent::final_response_with_default_usage(),
4423 ],
4424 [
4425 MockStreamEvent::text("accepted"),
4426 MockStreamEvent::final_response_with_default_usage(),
4427 ],
4428 ]))
4429 .record_content_telemetry(true)
4430 .add_hook(BoundedResponseRetry::new(
4431 "rejected",
4432 1,
4433 TestRetryMode::Repeat,
4434 ))
4435 .build()
4436 .runner("question")
4437 .max_turns(2)
4438 .stream()
4439 .await;
4440
4441 let mut saw_final = false;
4442 while let Some(item) = stream.next().await {
4443 if let MultiTurnStreamItem::FinalResponse(response) =
4444 item.expect("streaming retry item")
4445 {
4446 saw_final = true;
4447 assert_eq!(response.output, "accepted");
4448 }
4449 }
4450 assert!(saw_final, "streaming retry should produce a final response");
4451 }
4452
4453 async fn run_blocking_model_turn_stop_with_content_telemetry() {
4454 let error = AgentBuilder::new(MockCompletionModel::from_turns([MockTurn::text(
4455 "stopped blocking response",
4456 )]))
4457 .record_content_telemetry(true)
4458 .add_hook(StopCompletedModelTurn)
4459 .build()
4460 .runner("question")
4461 .run()
4462 .await
4463 .expect_err("blocking model-turn stop should cancel the run");
4464
4465 assert!(matches!(
4466 error,
4467 PromptError::PromptCancelled { reason, .. }
4468 if reason == "stop completed model turn"
4469 ));
4470 }
4471
4472 async fn run_streaming_model_turn_stop_with_content_telemetry() {
4473 let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([[
4474 MockStreamEvent::text("stopped streaming response"),
4475 MockStreamEvent::final_response_with_default_usage(),
4476 ]]))
4477 .record_content_telemetry(true)
4478 .add_hook(StopCompletedModelTurn)
4479 .build()
4480 .runner("question")
4481 .stream()
4482 .await;
4483
4484 let mut provider_finals = 0;
4485 let mut agent_finals = 0;
4486 let mut retries = 0;
4487 let mut errors = 0;
4488 while let Some(item) = stream.next().await {
4489 match item {
4490 Ok(MultiTurnStreamItem::StreamAssistantItem(
4491 StreamedAssistantContent::Final(_),
4492 )) => provider_finals += 1,
4493 Ok(MultiTurnStreamItem::FinalResponse(_)) => agent_finals += 1,
4494 Ok(MultiTurnStreamItem::ModelTurnRetried { .. }) => retries += 1,
4495 Ok(_) => {}
4496 Err(error) => {
4497 errors += 1;
4498 assert!(matches!(
4499 error,
4500 super::StreamingError::Prompt(error)
4501 if matches!(
4502 error.as_ref(),
4503 PromptError::PromptCancelled { reason, .. }
4504 if reason == "stop completed model turn"
4505 )
4506 ));
4507 }
4508 }
4509 }
4510
4511 assert_eq!(provider_finals, 1);
4512 assert_eq!(agent_finals, 0);
4513 assert_eq!(retries, 0);
4514 assert_eq!(errors, 1);
4515 }
4516
4517 #[test]
4524 fn chat_span_declares_the_full_completion_parent_contract() {
4525 use rig_core::telemetry::{
4526 COMPLETION_PARENT_MARKER_FIELD, COMPLETION_PARENT_REQUIRED_FIELDS,
4527 };
4528
4529 let _isolation = crate::test_utils::scoped_tracing_subscriber_guard_blocking();
4530 tracing::subscriber::with_default(Registry::default(), || {
4531 let agent = AgentBuilder::new(MockCompletionModel::text("done"))
4532 .name("contract-agent")
4533 .build();
4534 let runner = agent.runner("hello");
4535 let span = build_chat_span!(runner, None, "chat", "chat");
4536 let Some(metadata) = span.metadata() else {
4537 panic!("chat span was disabled");
4538 };
4539 let declared: HashSet<&str> =
4540 metadata.fields().iter().map(|field| field.name()).collect();
4541 let expected: HashSet<&str> = COMPLETION_PARENT_REQUIRED_FIELDS
4542 .iter()
4543 .copied()
4544 .chain([COMPLETION_PARENT_MARKER_FIELD, "gen_ai.agent.name"])
4545 .collect();
4546 assert_eq!(declared, expected);
4547 assert_eq!(metadata.fields().len(), expected.len());
4551 });
4552 }
4553
4554 #[tokio::test]
4555 async fn response_retry_records_only_accepted_content_on_both_surfaces() {
4556 let _isolation = crate::test_utils::scoped_tracing_subscriber_guard().await;
4557 let captured = Captured::default();
4558 let subscriber = Registry::default().with(CaptureLayer {
4559 captured: captured.clone(),
4560 });
4561 let _default = tracing::subscriber::set_default(subscriber);
4562
4563 run_blocking_response_retry_with_content_telemetry().await;
4566 run_streaming_response_retry_with_content_telemetry().await;
4567 tracing::callsite::rebuild_interest_cache();
4568 captured.clear();
4569
4570 run_blocking_response_retry_with_content_telemetry().await;
4571 let blocking = captured.snapshot();
4572 let blocking_chats = blocking
4573 .iter()
4574 .filter(|span| span.name == "chat")
4575 .collect::<Vec<_>>();
4576 assert_eq!(blocking_chats.len(), 2);
4577 assert!(
4578 blocking_chats
4579 .iter()
4580 .all(|span| span.target == "rig::agent_chat")
4581 );
4582 assert!(
4583 !blocking_chats[0]
4584 .field_names
4585 .contains("gen_ai.output.messages"),
4586 "rejected blocking content must not be recorded as model output"
4587 );
4588 assert!(
4589 blocking_chats[1]
4590 .field_names
4591 .contains("gen_ai.output.messages"),
4592 "accepted blocking content must be recorded as model output"
4593 );
4594 let blocking_output = blocking_chats[1]
4595 .string_fields
4596 .get("gen_ai.output.messages")
4597 .expect("accepted blocking output value");
4598 assert!(
4599 blocking_output
4600 .iter()
4601 .any(|value| value.contains("accepted"))
4602 );
4603 assert!(
4604 blocking_output
4605 .iter()
4606 .all(|value| !value.contains("rejected"))
4607 );
4608 let blocking_completion = blocking
4609 .iter()
4610 .find(|span| span.name == "invoke_agent")
4611 .and_then(|span| span.string_fields.get("gen_ai.completion"))
4612 .expect("accepted blocking run-level completion");
4613 assert_eq!(blocking_completion, &["accepted"]);
4614
4615 captured.clear();
4616 run_streaming_response_retry_with_content_telemetry().await;
4617 let streaming = captured.snapshot();
4618 let streaming_chats = streaming
4619 .iter()
4620 .filter(|span| span.name == "chat_streaming")
4621 .collect::<Vec<_>>();
4622 assert_eq!(streaming_chats.len(), 2);
4623 assert!(
4624 streaming_chats
4625 .iter()
4626 .all(|span| span.target == "rig::agent_chat")
4627 );
4628 assert!(
4629 !streaming_chats[0]
4630 .field_names
4631 .contains("gen_ai.output.messages"),
4632 "rejected streaming content must not be recorded as model output"
4633 );
4634 assert!(
4635 streaming_chats[1]
4636 .field_names
4637 .contains("gen_ai.output.messages"),
4638 "accepted streaming content must be recorded as model output"
4639 );
4640 let streaming_output = streaming_chats[1]
4641 .string_fields
4642 .get("gen_ai.output.messages")
4643 .expect("accepted streaming output value");
4644 assert!(
4645 streaming_output
4646 .iter()
4647 .any(|value| value.contains("accepted"))
4648 );
4649 assert!(
4650 streaming_output
4651 .iter()
4652 .all(|value| !value.contains("rejected"))
4653 );
4654 let streaming_completion = streaming
4655 .iter()
4656 .find(|span| span.name == "invoke_agent")
4657 .and_then(|span| span.string_fields.get("gen_ai.completion"))
4658 .expect("accepted streaming run-level completion");
4659 assert_eq!(streaming_completion, &["accepted"]);
4660 }
4661
4662 #[tokio::test]
4663 async fn model_turn_stop_preserves_completed_content_telemetry_on_both_surfaces() {
4664 let _isolation = crate::test_utils::scoped_tracing_subscriber_guard().await;
4665 let captured = Captured::default();
4666 let subscriber = Registry::default().with(CaptureLayer {
4667 captured: captured.clone(),
4668 });
4669 let _default = tracing::subscriber::set_default(subscriber);
4670
4671 run_blocking_model_turn_stop_with_content_telemetry().await;
4672 run_streaming_model_turn_stop_with_content_telemetry().await;
4673 tracing::callsite::rebuild_interest_cache();
4674 captured.clear();
4675
4676 run_blocking_model_turn_stop_with_content_telemetry().await;
4677 let blocking = captured.snapshot();
4678 let blocking_output = blocking
4679 .iter()
4680 .find(|span| span.name == "chat")
4681 .and_then(|span| span.string_fields.get("gen_ai.output.messages"))
4682 .expect("stopped blocking turn should retain output telemetry");
4683 assert!(
4684 blocking_output
4685 .iter()
4686 .any(|value| value.contains("stopped blocking response"))
4687 );
4688
4689 captured.clear();
4690 run_streaming_model_turn_stop_with_content_telemetry().await;
4691 let streaming = captured.snapshot();
4692 let streaming_output = streaming
4693 .iter()
4694 .find(|span| span.name == "chat_streaming")
4695 .and_then(|span| span.string_fields.get("gen_ai.output.messages"))
4696 .expect("stopped streaming turn should retain output telemetry");
4697 assert!(
4698 streaming_output
4699 .iter()
4700 .any(|value| value.contains("stopped streaming response"))
4701 );
4702 let streaming_completion = streaming
4703 .iter()
4704 .find(|span| span.name == "invoke_agent")
4705 .and_then(|span| span.string_fields.get("gen_ai.completion"))
4706 .expect("stopped streaming turn should retain run-level completion telemetry");
4707 assert_eq!(streaming_completion, &["stopped streaming response"]);
4708 }
4709
4710 #[tokio::test]
4711 async fn run_records_usage_and_chains_chat_spans_on_a_created_agent_span() {
4712 let _isolation = crate::test_utils::scoped_tracing_subscriber_guard().await;
4713 let captured = Captured::default();
4714 let subscriber = Registry::default().with(CaptureLayer {
4715 captured: captured.clone(),
4716 });
4717 let _default = tracing::subscriber::set_default(subscriber);
4718
4719 warm_blocking_callsites().await;
4720 tracing::callsite::rebuild_interest_cache();
4721 captured.clear();
4722
4723 let agent = AgentBuilder::new(tool_then_text_model())
4724 .record_content_telemetry(true)
4725 .tool(MockAddTool)
4726 .build();
4727 let response = agent
4728 .runner("add 2 and 3")
4729 .max_turns(3)
4730 .run()
4731 .await
4732 .expect("blocking run should succeed");
4733 assert_eq!(response.output, "the answer is 5");
4734
4735 let spans = captured.snapshot();
4736
4737 let chat_spans: Vec<&CapturedSpan> =
4739 spans.iter().filter(|s| s.name == "chat").collect();
4740 assert_eq!(chat_spans.len(), 2, "two model turns -> two chat spans");
4741 assert!(
4742 spans.iter().all(|s| s.name != "chat_streaming"),
4743 "blocking driver must not emit chat_streaming spans"
4744 );
4745
4746 let agent_span = spans
4748 .iter()
4749 .find(|s| s.name == "invoke_agent")
4750 .expect("blocking run should create an invoke_agent span");
4751
4752 assert_eq!(
4754 agent_span.u64_fields.get("gen_ai.usage.input_tokens"),
4755 Some(&(7 + 13)),
4756 );
4757 assert_eq!(
4758 agent_span.u64_fields.get("gen_ai.usage.output_tokens"),
4759 Some(&(11 + 17)),
4760 );
4761 assert!(
4762 agent_span.field_names.contains("gen_ai.completion"),
4763 "the created agent span records the final completion text"
4764 );
4765
4766 let tool_span = spans
4771 .iter()
4772 .find(|s| s.name == "execute_tool")
4773 .expect("tool turn should emit an execute_tool span");
4774 let edges = captured.follows_edges();
4775 assert!(
4776 edges.contains(&(tool_span.id, chat_spans[0].id)),
4777 "execute_tool should follow_from the first chat span; edges={edges:?}"
4778 );
4779 assert!(
4780 edges.contains(&(chat_spans[1].id, tool_span.id)),
4781 "the second chat span should follow_from execute_tool; edges={edges:?}"
4782 );
4783 }
4784
4785 #[tokio::test]
4786 async fn classic_completion_parent_is_enriched_without_duplicate_provider_span() {
4787 let _isolation = crate::test_utils::scoped_tracing_subscriber_guard().await;
4788 let captured = Captured::default();
4789 let subscriber = Registry::default().with(CaptureLayer {
4790 captured: captured.clone(),
4791 });
4792 let _default = tracing::subscriber::set_default(subscriber);
4793
4794 let warm = AgentBuilder::new(CompletionTelemetryModel {
4795 inner: MockCompletionModel::text("warm"),
4796 })
4797 .build();
4798 let _ = warm.prompt("warm").await;
4799 tracing::callsite::rebuild_interest_cache();
4800 captured.clear();
4801
4802 let agent = AgentBuilder::new(CompletionTelemetryModel {
4803 inner: MockCompletionModel::text("done"),
4804 })
4805 .build();
4806 let response = agent.prompt("hello").await.expect("prompt should succeed");
4807 assert_eq!(response, "done");
4808
4809 let spans = captured.snapshot();
4810 let chat_spans = spans
4811 .iter()
4812 .filter(|span| span.name == "chat")
4813 .collect::<Vec<_>>();
4814 assert_eq!(chat_spans.len(), 1, "provider telemetry must reuse chat");
4815 assert_eq!(chat_spans[0].target, "rig::agent_chat");
4816 assert!(
4817 spans.iter().all(|span| span.target != "rig::completions"),
4818 "an adopted classic completion parent must not gain a provider child"
4819 );
4820 assert_eq!(
4821 chat_spans[0]
4822 .string_fields
4823 .get("gen_ai.provider.name")
4824 .and_then(|values| values.first())
4825 .map(String::as_str),
4826 Some("fixture-provider")
4827 );
4828 assert_eq!(
4829 chat_spans[0]
4830 .string_fields
4831 .get("gen_ai.request.model")
4832 .and_then(|values| values.first())
4833 .map(String::as_str),
4834 Some("fixture-model")
4835 );
4836 }
4837
4838 #[tokio::test]
4839 async fn run_does_not_record_usage_onto_a_caller_supplied_outer_span() {
4840 let _isolation = crate::test_utils::scoped_tracing_subscriber_guard().await;
4841 let captured = Captured::default();
4842 let subscriber = Registry::default().with(CaptureLayer {
4843 captured: captured.clone(),
4844 });
4845 let _default = tracing::subscriber::set_default(subscriber);
4846
4847 warm_blocking_callsites().await;
4848 tracing::callsite::rebuild_interest_cache();
4849 captured.clear();
4850
4851 let outer = tracing::info_span!(
4855 "outer",
4856 gen_ai.completion = tracing::field::Empty,
4857 gen_ai.usage.input_tokens = tracing::field::Empty,
4858 gen_ai.usage.output_tokens = tracing::field::Empty,
4859 );
4860 async {
4861 let agent = AgentBuilder::new(tool_then_text_model())
4862 .tool(MockAddTool)
4863 .build();
4864 agent
4865 .runner("add 2 and 3")
4866 .max_turns(3)
4867 .run()
4868 .await
4869 .expect("blocking run should succeed");
4870 }
4871 .instrument(outer)
4872 .await;
4873
4874 let spans = captured.snapshot();
4875 assert!(
4877 spans.iter().all(|s| s.name != "invoke_agent"),
4878 "an ambient outer span should be adopted, not wrapped in invoke_agent"
4879 );
4880 let outer_span = spans
4881 .iter()
4882 .find(|s| s.name == "outer")
4883 .expect("outer span should be captured");
4884 assert!(
4885 outer_span
4886 .field_names
4887 .iter()
4888 .all(|name| !name.starts_with("gen_ai.usage.")),
4889 "run-level usage must not be recorded onto a caller-supplied outer span"
4890 );
4891 assert!(
4892 !outer_span.field_names.contains("gen_ai.completion"),
4893 "run-level completion must not be recorded onto a caller-supplied outer span"
4894 );
4895 }
4896
4897 struct RawOutputTool;
4902 impl crate::tool::Tool for RawOutputTool {
4903 const NAME: &'static str = "raw_output";
4904 type Error = rig::tool::ToolExecutionError;
4905 type Args = serde_json::Value;
4906 type Output = String;
4907 fn description(&self) -> String {
4908 "returns a raw output marker".to_string()
4909 }
4910
4911 fn parameters(&self) -> serde_json::Value {
4912 serde_json::json!({ "type": "object", "properties": {} })
4913 }
4914 async fn call(
4915 &self,
4916 _context: &mut ToolContext,
4917 _args: Self::Args,
4918 ) -> Result<Self::Output, ToolExecutionError> {
4919 Ok("RAW_EXECUTION_OUTPUT_42".to_string())
4920 }
4921 }
4922
4923 struct RedactResultHook;
4925 impl crate::agent::AgentHook for RedactResultHook {
4926 async fn on_tool_result(
4927 &self,
4928 _ctx: &HookContext,
4929 event: ToolResultEvent<'_>,
4930 ) -> ToolResultAction {
4931 if let crate::agent::ToolResultEvent { .. } = event {
4932 crate::agent::ToolResultAction::rewrite("[REDACTED]")
4933 } else {
4934 crate::agent::ToolResultAction::keep()
4935 }
4936 }
4937 }
4938
4939 struct StopOnResultHook;
4941 impl crate::agent::AgentHook for StopOnResultHook {
4942 async fn on_tool_result(
4943 &self,
4944 _ctx: &HookContext,
4945 _event: ToolResultEvent<'_>,
4946 ) -> ToolResultAction {
4947 ToolResultAction::stop("stop after raw result")
4948 }
4949 }
4950
4951 #[derive(Default)]
4954 struct ResultValueVisitor {
4955 values: Vec<String>,
4956 }
4957 impl Visit for ResultValueVisitor {
4958 fn record_str(&mut self, field: &Field, value: &str) {
4959 if field.name() == "gen_ai.tool.call.result" {
4960 self.values.push(value.to_string());
4961 }
4962 }
4963 fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
4964 if field.name() == "gen_ai.tool.call.result" {
4965 self.values.push(format!("{value:?}"));
4966 }
4967 }
4968 }
4969
4970 struct ResultValueLayer {
4971 values: Arc<Mutex<Vec<String>>>,
4972 }
4973 impl<S> Layer<S> for ResultValueLayer
4974 where
4975 S: Subscriber + for<'l> LookupSpan<'l>,
4976 {
4977 fn on_record(&self, _span: &Id, values: &Record<'_>, _ctx: Context<'_, S>) {
4978 let mut visitor = ResultValueVisitor::default();
4979 values.record(&mut visitor);
4980 if !visitor.values.is_empty() {
4981 self.values.lock().expect("values").extend(visitor.values);
4982 }
4983 }
4984 }
4985
4986 #[tokio::test]
4989 async fn tool_result_rewrite_redacts_span_output() {
4990 let _isolation = crate::test_utils::scoped_tracing_subscriber_guard().await;
4991 let values: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
4992 let subscriber = Registry::default().with(ResultValueLayer {
4993 values: values.clone(),
4994 });
4995 let _default = tracing::subscriber::set_default(subscriber);
4996
4997 warm_blocking_callsites().await;
5000 tracing::callsite::rebuild_interest_cache();
5001 values.lock().expect("values").clear();
5002
5003 let model = MockCompletionModel::from_turns([
5004 MockTurn::tool_call("tc1", "raw_output", serde_json::json!({})),
5005 MockTurn::text("ok"),
5006 ]);
5007 let response = AgentBuilder::new(model)
5008 .record_content_telemetry(true)
5009 .tool(RawOutputTool)
5010 .add_hook(RedactResultHook)
5011 .build()
5012 .runner("go")
5013 .max_turns(3)
5014 .run()
5015 .await
5016 .expect("run should succeed");
5017 assert_eq!(response.output, "ok");
5018
5019 let captured = values.lock().expect("values").clone();
5020 assert!(
5021 captured.iter().any(|v| v.contains("[REDACTED]")),
5022 "the rewritten presentation must reach telemetry; captured: {captured:?}"
5023 );
5024 assert!(
5025 !captured
5026 .iter()
5027 .any(|v| v.contains("RAW_EXECUTION_OUTPUT_42")),
5028 "the raw tool output must not leak through telemetry; captured: {captured:?}"
5029 );
5030 }
5031
5032 #[tokio::test]
5035 async fn tool_result_stop_omits_span_output() {
5036 let _isolation = crate::test_utils::scoped_tracing_subscriber_guard().await;
5037 let values: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
5038 let subscriber = Registry::default().with(ResultValueLayer {
5039 values: values.clone(),
5040 });
5041 let _default = tracing::subscriber::set_default(subscriber);
5042
5043 warm_blocking_callsites().await;
5044 tracing::callsite::rebuild_interest_cache();
5045 values.lock().expect("values").clear();
5046
5047 let result = AgentBuilder::new(MockCompletionModel::from_turns([MockTurn::tool_call(
5048 "tc1",
5049 "raw_output",
5050 serde_json::json!({}),
5051 )]))
5052 .tool(RawOutputTool)
5053 .add_hook(StopOnResultHook)
5054 .build()
5055 .runner("go")
5056 .max_turns(2)
5057 .run()
5058 .await;
5059 assert!(result.is_err(), "the result hook should stop the run");
5060
5061 let captured = values.lock().expect("values").clone();
5062 assert!(
5063 !captured
5064 .iter()
5065 .any(|value| value.contains("RAW_EXECUTION_OUTPUT_42")),
5066 "a Stop must not leak raw execution telemetry; captured: {captured:?}"
5067 );
5068 }
5069 }
5070
5071 fn tool_call_content(id: &str, args: serde_json::Value) -> AssistantContent {
5072 AssistantContent::ToolCall(MessageToolCall::from_wire(
5076 id,
5077 ToolFunction::new("add".to_string(), args),
5078 ))
5079 }
5080
5081 fn tool_result_text_in_history(messages: &[Message], expected: &str) -> bool {
5085 messages.iter().any(|message| {
5086 matches!(
5087 message,
5088 Message::User { content }
5089 if content.iter().any(|item| matches!(
5090 item,
5091 UserContent::ToolResult(result)
5092 if result.content.iter().any(|c| matches!(
5093 c,
5094 rig_core::message::ToolResultContent::Text(text)
5095 if text.text == expected
5096 ))
5097 ))
5098 )
5099 })
5100 }
5101
5102 fn tool_result_json_in_history(messages: &[Message], expected: &serde_json::Value) -> bool {
5104 messages.iter().any(|message| {
5105 matches!(
5106 message,
5107 Message::User { content }
5108 if content.iter().any(|item| matches!(
5109 item,
5110 UserContent::ToolResult(result)
5111 if result.content.iter().any(|content| matches!(
5112 content,
5113 rig_core::message::ToolResultContent::Json { value }
5114 if value == expected
5115 ))
5116 ))
5117 )
5118 })
5119 }
5120
5121 #[tokio::test]
5126 async fn run_and_stream_same_message_history_for_parallel_tool_calls() {
5127 let blocking_model = MockCompletionModel::from_turns([
5128 MockTurn::from_contents([
5129 tool_call_content("tc1", json!({"x": 2, "y": 3})),
5130 tool_call_content("tc2", json!({"x": 10, "y": 20})),
5131 ]),
5132 MockTurn::text("done"),
5133 ]);
5134 let blocking = AgentBuilder::new(blocking_model)
5135 .tool(MockAddTool)
5136 .build()
5137 .runner("add two pairs")
5138 .max_turns(3)
5139 .tool_concurrency(4)
5140 .run()
5141 .await
5142 .expect("blocking run should succeed");
5143
5144 let streaming_model = MockCompletionModel::from_stream_turns([
5145 vec![
5146 MockStreamEvent::tool_call("tc1", "add", json!({"x": 2, "y": 3})),
5147 MockStreamEvent::tool_call("tc2", "add", json!({"x": 10, "y": 20})),
5148 MockStreamEvent::final_response_with_total_tokens(0),
5149 ],
5150 vec![
5151 MockStreamEvent::text("done"),
5152 MockStreamEvent::final_response_with_total_tokens(0),
5153 ],
5154 ]);
5155 let mut stream = AgentBuilder::new(streaming_model)
5156 .tool(MockAddTool)
5157 .build()
5158 .runner("add two pairs")
5159 .max_turns(3)
5160 .stream()
5161 .await;
5162 let mut final_response = None;
5163 while let Some(item) = stream.next().await {
5164 if let Ok(MultiTurnStreamItem::FinalResponse(resp)) =
5165 item.map_err(|err| panic!("stream item errored: {err}"))
5166 {
5167 final_response = Some(resp);
5168 }
5169 }
5170 let final_response = final_response.expect("stream should yield a final response");
5171
5172 let blocking_messages = blocking.messages.expect("blocking messages");
5173 let streaming_messages = final_response
5174 .messages()
5175 .expect("streaming history")
5176 .to_vec();
5177 assert_eq!(
5178 serde_json::to_value(&blocking_messages).expect("serialize blocking"),
5179 serde_json::to_value(&streaming_messages).expect("serialize streaming"),
5180 );
5181 }
5182
5183 #[derive(Clone)]
5189 struct OutOfOrderTool {
5190 gate: Arc<tokio::sync::Notify>,
5191 order: Arc<AtomicU32>,
5192 }
5193
5194 impl Tool for OutOfOrderTool {
5195 const NAME: &'static str = "add";
5196 type Error = MockToolError;
5197 type Args = MockOperationArgs;
5198 type Output = i32;
5199
5200 fn description(&self) -> String {
5201 MockAddTool.description()
5202 }
5203
5204 fn parameters(&self) -> serde_json::Value {
5205 MockAddTool.parameters()
5206 }
5207
5208 async fn call(
5209 &self,
5210 _context: &mut ToolContext,
5211 _args: Self::Args,
5212 ) -> Result<Self::Output, Self::Error> {
5213 let nth = self.order.fetch_add(1, SeqCst);
5214 if nth == 0 {
5215 self.gate.notified().await;
5217 } else {
5218 self.gate.notify_one();
5220 }
5221 Ok(nth as i32)
5222 }
5223 }
5224
5225 #[tokio::test]
5231 async fn run_preserves_tool_call_order_under_out_of_order_completion() {
5232 let model = MockCompletionModel::from_turns([
5233 MockTurn::from_contents([
5234 tool_call_content("tc1", json!({"x": 1, "y": 0})),
5235 tool_call_content("tc2", json!({"x": 2, "y": 0})),
5236 ]),
5237 MockTurn::text("done"),
5238 ]);
5239 let response = AgentBuilder::new(model)
5240 .tool(OutOfOrderTool {
5241 gate: Arc::new(tokio::sync::Notify::new()),
5242 order: Arc::new(AtomicU32::new(0)),
5243 })
5244 .build()
5245 .runner("go")
5246 .max_turns(3)
5247 .tool_concurrency(4)
5248 .run()
5249 .await
5250 .expect("run should succeed");
5251
5252 let messages = response.messages.expect("messages");
5253 let result_ids: Vec<String> = messages
5254 .iter()
5255 .flat_map(|message| match message {
5256 Message::User { content } => content
5257 .iter()
5258 .filter_map(|item| match item {
5259 UserContent::ToolResult(result) => Some(result.call.to_string()),
5260 _ => None,
5261 })
5262 .collect::<Vec<_>>(),
5263 _ => Vec::new(),
5264 })
5265 .collect();
5266 assert_eq!(result_ids, vec!["tc1".to_string(), "tc2".to_string()]);
5268 }
5269
5270 async fn drive_to_final_response(
5273 mut stream: crate::agent::prompt_request::streaming::StreamingResult,
5274 ) -> crate::agent::prompt_request::PromptResponse {
5275 let mut final_response = None;
5276 while let Some(item) = stream.next().await {
5277 if let MultiTurnStreamItem::FinalResponse(resp) =
5278 item.unwrap_or_else(|err| panic!("stream item errored: {err}"))
5279 {
5280 final_response = Some(resp);
5281 }
5282 }
5283 final_response.expect("stream should yield a final response")
5284 }
5285
5286 fn tool_result_ids(messages: &[Message]) -> Vec<String> {
5288 messages
5289 .iter()
5290 .flat_map(|message| match message {
5291 Message::User { content } => content
5292 .iter()
5293 .filter_map(|item| match item {
5294 UserContent::ToolResult(result) => Some(result.call.to_string()),
5295 _ => None,
5296 })
5297 .collect::<Vec<_>>(),
5298 _ => Vec::new(),
5299 })
5300 .collect()
5301 }
5302
5303 #[tokio::test]
5308 async fn stream_and_run_same_message_history_for_parallel_tool_calls_under_concurrency() {
5309 let blocking_model = MockCompletionModel::from_turns([
5310 MockTurn::from_contents([
5311 tool_call_content("tc1", json!({"x": 2, "y": 3})),
5312 tool_call_content("tc2", json!({"x": 10, "y": 20})),
5313 ]),
5314 MockTurn::text("done"),
5315 ]);
5316 let blocking = AgentBuilder::new(blocking_model)
5317 .tool(MockAddTool)
5318 .build()
5319 .runner("add two pairs")
5320 .max_turns(3)
5321 .tool_concurrency(4)
5322 .run()
5323 .await
5324 .expect("blocking run should succeed");
5325
5326 let streaming_model = MockCompletionModel::from_stream_turns([
5327 vec![
5328 MockStreamEvent::tool_call("tc1", "add", json!({"x": 2, "y": 3})),
5329 MockStreamEvent::tool_call("tc2", "add", json!({"x": 10, "y": 20})),
5330 MockStreamEvent::final_response_with_total_tokens(0),
5331 ],
5332 vec![
5333 MockStreamEvent::text("done"),
5334 MockStreamEvent::final_response_with_total_tokens(0),
5335 ],
5336 ]);
5337 let stream = AgentBuilder::new(streaming_model)
5338 .tool(MockAddTool)
5339 .build()
5340 .runner("add two pairs")
5341 .max_turns(3)
5342 .tool_concurrency(4)
5343 .stream()
5344 .await;
5345 let final_response = drive_to_final_response(stream).await;
5346
5347 let blocking_messages = blocking.messages.expect("blocking messages");
5348 let streaming_messages = final_response
5349 .messages()
5350 .expect("streaming history")
5351 .to_vec();
5352 assert_eq!(
5353 serde_json::to_value(&blocking_messages).expect("serialize blocking"),
5354 serde_json::to_value(&streaming_messages).expect("serialize streaming"),
5355 );
5356 }
5357
5358 #[tokio::test]
5364 async fn stream_preserves_history_order_under_out_of_order_completion() {
5365 let model = MockCompletionModel::from_stream_turns([
5366 vec![
5367 MockStreamEvent::tool_call("tc1", "add", json!({"x": 1, "y": 0})),
5368 MockStreamEvent::tool_call("tc2", "add", json!({"x": 2, "y": 0})),
5369 MockStreamEvent::final_response_with_total_tokens(0),
5370 ],
5371 vec![
5372 MockStreamEvent::text("done"),
5373 MockStreamEvent::final_response_with_total_tokens(0),
5374 ],
5375 ]);
5376 let stream = AgentBuilder::new(model)
5377 .tool(OutOfOrderTool {
5378 gate: Arc::new(tokio::sync::Notify::new()),
5379 order: Arc::new(AtomicU32::new(0)),
5380 })
5381 .build()
5382 .runner("go")
5383 .max_turns(3)
5384 .tool_concurrency(4)
5385 .stream()
5386 .await;
5387 let final_response = tokio::time::timeout(
5390 std::time::Duration::from_secs(5),
5391 drive_to_final_response(stream),
5392 )
5393 .await
5394 .expect("streamed tools must run concurrently, not deadlock on the first call");
5395
5396 let messages = final_response.messages().expect("history").to_vec();
5397 assert_eq!(
5399 tool_result_ids(&messages),
5400 vec!["tc1".to_string(), "tc2".to_string()]
5401 );
5402 }
5403
5404 #[tokio::test]
5409 async fn stream_emits_tool_results_in_call_order_after_batch_settles_under_concurrency() {
5410 let model = MockCompletionModel::from_stream_turns([
5411 vec![
5412 MockStreamEvent::tool_call("tc1", "add", json!({"x": 1, "y": 0})),
5413 MockStreamEvent::tool_call("tc2", "add", json!({"x": 2, "y": 0})),
5414 MockStreamEvent::final_response_with_total_tokens(0),
5415 ],
5416 vec![
5417 MockStreamEvent::text("done"),
5418 MockStreamEvent::final_response_with_total_tokens(0),
5419 ],
5420 ]);
5421 let mut stream = AgentBuilder::new(model)
5422 .tool(OutOfOrderTool {
5423 gate: Arc::new(tokio::sync::Notify::new()),
5424 order: Arc::new(AtomicU32::new(0)),
5425 })
5426 .build()
5427 .runner("go")
5428 .max_turns(3)
5429 .tool_concurrency(4)
5430 .stream()
5431 .await;
5432
5433 let mut streamed_result_ids = Vec::new();
5434 let mut final_response = None;
5435 tokio::time::timeout(std::time::Duration::from_secs(5), async {
5436 while let Some(item) = stream.next().await {
5437 match item.unwrap_or_else(|err| panic!("stream item errored: {err}")) {
5438 MultiTurnStreamItem::StreamUserItem(StreamedUserContent::ToolResult {
5439 tool_result,
5440 ..
5441 }) => streamed_result_ids.push(tool_result.call.into_string()),
5442 MultiTurnStreamItem::FinalResponse(resp) => final_response = Some(resp),
5443 _ => {}
5444 }
5445 }
5446 })
5447 .await
5448 .expect("streamed tools must run concurrently, not deadlock on the first call");
5449
5450 assert_eq!(
5453 streamed_result_ids,
5454 vec!["tc1".to_string(), "tc2".to_string()]
5455 );
5456 let final_response = final_response.expect("stream should yield a final response");
5457 assert_eq!(
5458 tool_result_ids(final_response.messages().expect("history")),
5459 vec!["tc1".to_string(), "tc2".to_string()]
5460 );
5461 }
5462
5463 #[tokio::test]
5469 async fn stream_executes_tools_concurrently_under_concurrency() {
5470 let barrier = Arc::new(tokio::sync::Barrier::new(2));
5471 let model = MockCompletionModel::from_stream_turns([
5472 vec![
5473 MockStreamEvent::tool_call("b1", "barrier_tool", json!({})),
5474 MockStreamEvent::tool_call("b2", "barrier_tool", json!({})),
5475 MockStreamEvent::final_response_with_total_tokens(0),
5476 ],
5477 vec![
5478 MockStreamEvent::text("done"),
5479 MockStreamEvent::final_response_with_total_tokens(0),
5480 ],
5481 ]);
5482 let stream = AgentBuilder::new(model)
5483 .tool(MockBarrierTool::new(barrier))
5484 .build()
5485 .runner("hit the barrier twice")
5486 .max_turns(3)
5487 .tool_concurrency(2)
5488 .stream()
5489 .await;
5490
5491 tokio::time::timeout(
5492 std::time::Duration::from_secs(5),
5493 drive_to_final_response(stream),
5494 )
5495 .await
5496 .expect("streamed tools must run concurrently, not deadlock at the barrier");
5497 }
5498
5499 #[tokio::test]
5506 async fn stream_emits_model_tool_calls_then_atomic_execution_items() {
5507 async fn markers(concurrency: usize) -> Vec<&'static str> {
5508 let model = MockCompletionModel::from_stream_turns([
5509 vec![
5510 MockStreamEvent::tool_call("tc1", "add", json!({"x": 1, "y": 1})),
5511 MockStreamEvent::tool_call("tc2", "add", json!({"x": 2, "y": 2})),
5512 MockStreamEvent::final_response_with_total_tokens(0),
5513 ],
5514 vec![
5515 MockStreamEvent::text("done"),
5516 MockStreamEvent::final_response_with_total_tokens(0),
5517 ],
5518 ]);
5519 let mut stream = AgentBuilder::new(model)
5520 .tool(MockAddTool)
5521 .build()
5522 .runner("add two pairs")
5523 .max_turns(3)
5524 .tool_concurrency(concurrency)
5525 .stream()
5526 .await;
5527 let mut markers = Vec::new();
5528 while let Some(item) = stream.next().await {
5529 match item.unwrap_or_else(|err| panic!("stream item errored: {err}")) {
5530 MultiTurnStreamItem::StreamAssistantItem(
5531 StreamedAssistantContent::ToolCall { .. },
5532 ) => markers.push("model-call"),
5533 MultiTurnStreamItem::ToolExecutionCommitted { .. } => {
5534 markers.push("exec-commit")
5535 }
5536 MultiTurnStreamItem::StreamUserItem(StreamedUserContent::ToolResult {
5537 ..
5538 }) => markers.push("result"),
5539 _ => {}
5540 }
5541 }
5542 markers
5543 }
5544
5545 let expected = vec![
5548 "model-call",
5549 "model-call",
5550 "exec-commit",
5551 "result",
5552 "exec-commit",
5553 "result",
5554 ];
5555 assert_eq!(markers(1).await, expected);
5556 assert_eq!(markers(4).await, expected);
5557 }
5558
5559 struct TerminateAfterSiblingStartedHook {
5563 sibling_started: Arc<tokio::sync::Notify>,
5564 }
5565 impl AgentHook for TerminateAfterSiblingStartedHook {
5566 async fn on_tool_result(
5567 &self,
5568 _ctx: &HookContext,
5569 event: ToolResultEvent<'_>,
5570 ) -> ToolResultAction {
5571 if let ToolResultEvent { args, .. } = event
5572 && serde_json::from_str::<serde_json::Value>(args)
5573 .ok()
5574 .and_then(|v| v.get("x").and_then(serde_json::Value::as_i64))
5575 == Some(1)
5576 {
5577 self.sibling_started.notified().await;
5578 return ToolResultAction::stop("stop after a tool result");
5579 }
5580 ToolResultAction::keep()
5581 }
5582 }
5583
5584 #[derive(Clone)]
5592 struct DrainProbeTool {
5593 started: Arc<AtomicU32>,
5594 completed: Arc<AtomicU32>,
5595 slow_started: Arc<tokio::sync::Notify>,
5596 }
5597
5598 impl Tool for DrainProbeTool {
5599 const NAME: &'static str = "add";
5600 type Error = MockToolError;
5601 type Args = serde_json::Value;
5602 type Output = i32;
5603
5604 fn description(&self) -> String {
5605 MockAddTool.description()
5606 }
5607
5608 fn parameters(&self) -> serde_json::Value {
5609 MockAddTool.parameters()
5610 }
5611
5612 async fn call(
5613 &self,
5614 _context: &mut ToolContext,
5615 args: Self::Args,
5616 ) -> Result<Self::Output, Self::Error> {
5617 self.started.fetch_add(1, SeqCst);
5618 if args.get("x").and_then(serde_json::Value::as_i64) == Some(2) {
5619 self.slow_started.notify_one();
5622 for _ in 0..8 {
5623 tokio::task::yield_now().await;
5624 }
5625 }
5626 self.completed.fetch_add(1, SeqCst);
5627 Ok(0)
5628 }
5629 }
5630
5631 #[tokio::test]
5638 async fn stream_concurrent_tool_result_terminate_drains_in_flight_siblings() {
5639 let started = Arc::new(AtomicU32::new(0));
5640 let completed = Arc::new(AtomicU32::new(0));
5641 let slow_started = Arc::new(tokio::sync::Notify::new());
5642 let model = MockCompletionModel::from_stream_turns([
5643 vec![
5644 MockStreamEvent::tool_call("tc1", "add", json!({"x": 1, "y": 1})),
5645 MockStreamEvent::tool_call("tc2", "add", json!({"x": 2, "y": 2})),
5646 MockStreamEvent::final_response_with_total_tokens(0),
5647 ],
5648 vec![
5649 MockStreamEvent::text("done"),
5650 MockStreamEvent::final_response_with_total_tokens(0),
5651 ],
5652 ]);
5653 let mut stream = AgentBuilder::new(model)
5654 .tool(DrainProbeTool {
5655 started: started.clone(),
5656 completed: completed.clone(),
5657 slow_started: slow_started.clone(),
5658 })
5659 .build()
5660 .runner("add two pairs")
5661 .max_turns(3)
5662 .tool_concurrency(2)
5663 .add_hook(TerminateAfterSiblingStartedHook {
5664 sibling_started: slow_started,
5665 })
5666 .stream()
5667 .await;
5668
5669 let (saw_error, saw_final_response) =
5670 tokio::time::timeout(std::time::Duration::from_secs(5), async move {
5671 let mut saw_error = false;
5672 let mut saw_final_response = false;
5673 while let Some(item) = stream.next().await {
5674 match item {
5675 Ok(MultiTurnStreamItem::FinalResponse(_)) => saw_final_response = true,
5676 Ok(_) => {}
5677 Err(StreamingError::Prompt(_)) => saw_error = true,
5678 Err(other) => panic!("unexpected streaming error: {other}"),
5679 }
5680 }
5681 (saw_error, saw_final_response)
5682 })
5683 .await
5684 .expect("draining the concurrent tools must not hang");
5685
5686 assert!(
5687 saw_error,
5688 "a terminate hook on the concurrent path must surface a StreamingError::Prompt"
5689 );
5690 assert!(
5691 !saw_final_response,
5692 "a terminated run must not yield a final response"
5693 );
5694 assert_eq!(
5697 started.load(SeqCst),
5698 2,
5699 "both tools started (both in flight)"
5700 );
5701 assert_eq!(
5702 completed.load(SeqCst),
5703 2,
5704 "the in-flight sibling must be drained to completion, not cancelled"
5705 );
5706 }
5707
5708 struct OrderedTerminateHook {
5713 gate: Arc<tokio::sync::Notify>,
5714 }
5715
5716 impl AgentHook for OrderedTerminateHook {
5717 async fn on_tool_call(&self, _ctx: &HookContext, event: ToolCall<'_>) -> ToolCallAction {
5718 if let ToolCall { args, .. } = event {
5719 let x = serde_json::from_str::<serde_json::Value>(args)
5720 .ok()
5721 .and_then(|v| v.get("x").and_then(serde_json::Value::as_i64));
5722 match x {
5723 Some(2) => {
5724 self.gate.notify_one();
5725 return ToolCallAction::stop("terminated-by-tc2".to_string());
5726 }
5727 Some(1) => {
5728 self.gate.notified().await;
5729 return ToolCallAction::stop("terminated-by-tc1".to_string());
5730 }
5731 _ => {}
5732 }
5733 }
5734 ToolCallAction::run()
5735 }
5736 }
5737
5738 fn two_terminating_tools_blocking_model() -> MockCompletionModel {
5739 MockCompletionModel::from_turns([
5740 MockTurn::from_contents([
5741 tool_call_content("tc1", json!({"x": 1, "y": 1})),
5742 tool_call_content("tc2", json!({"x": 2, "y": 2})),
5743 ]),
5744 MockTurn::text("unreachable"),
5745 ])
5746 }
5747
5748 fn two_terminating_tools_streaming_model() -> MockCompletionModel {
5749 MockCompletionModel::from_stream_turns([
5750 vec![
5751 MockStreamEvent::tool_call("tc1", "add", json!({"x": 1, "y": 1})),
5752 MockStreamEvent::tool_call("tc2", "add", json!({"x": 2, "y": 2})),
5753 MockStreamEvent::final_response_with_total_tokens(0),
5754 ],
5755 vec![
5756 MockStreamEvent::text("unreachable"),
5757 MockStreamEvent::final_response_with_total_tokens(0),
5758 ],
5759 ])
5760 }
5761
5762 #[tokio::test]
5768 async fn concurrent_simultaneous_tool_terminations_pick_call_order_on_both_drivers() {
5769 let run_err = tokio::time::timeout(
5770 std::time::Duration::from_secs(5),
5771 AgentBuilder::new(two_terminating_tools_blocking_model())
5772 .tool(MockAddTool)
5773 .build()
5774 .runner("go")
5775 .max_turns(3)
5776 .tool_concurrency(2)
5777 .add_hook(OrderedTerminateHook {
5778 gate: Arc::new(tokio::sync::Notify::new()),
5779 })
5780 .run(),
5781 )
5782 .await
5783 .expect("blocking run must not hang")
5784 .expect_err("the run must terminate");
5785
5786 let mut stream = AgentBuilder::new(two_terminating_tools_streaming_model())
5787 .tool(MockAddTool)
5788 .build()
5789 .runner("go")
5790 .max_turns(3)
5791 .tool_concurrency(2)
5792 .add_hook(OrderedTerminateHook {
5793 gate: Arc::new(tokio::sync::Notify::new()),
5794 })
5795 .stream()
5796 .await;
5797
5798 let stream_err = tokio::time::timeout(std::time::Duration::from_secs(5), async move {
5799 while let Some(item) = stream.next().await {
5800 if let Err(err) = item {
5801 return Some(err);
5802 }
5803 }
5804 None
5805 })
5806 .await
5807 .expect("streamed run must not hang")
5808 .expect("the stream must surface a terminate error");
5809
5810 let run_msg = run_err.to_string();
5811 let stream_msg = stream_err.to_string();
5812 assert!(
5813 run_msg.contains("terminated-by-tc1"),
5814 "blocking run should surface the first-called tool's reason, got: {run_msg}"
5815 );
5816 assert!(
5817 stream_msg.contains("terminated-by-tc1"),
5818 "stream should surface the first-called tool's reason, got: {stream_msg}"
5819 );
5820 assert!(
5821 !run_msg.contains("terminated-by-tc2") && !stream_msg.contains("terminated-by-tc2"),
5822 "neither driver should surface the later-completing tool's reason"
5823 );
5824 }
5825
5826 struct TerminateOnFirstToolHook;
5829 impl AgentHook for TerminateOnFirstToolHook {
5830 async fn on_tool_call(&self, _ctx: &HookContext, event: ToolCall<'_>) -> ToolCallAction {
5831 if let ToolCall { args, .. } = event
5832 && serde_json::from_str::<serde_json::Value>(args)
5833 .ok()
5834 .and_then(|v| v.get("x").and_then(serde_json::Value::as_i64))
5835 == Some(1)
5836 {
5837 return ToolCallAction::stop("stop".to_string());
5838 }
5839 ToolCallAction::run()
5840 }
5841 }
5842
5843 #[tokio::test]
5850 async fn default_concurrency_terminate_skips_remaining_tools_on_both_drivers() {
5851 let blocking_calls = Arc::new(AtomicU32::new(0));
5852 AgentBuilder::new(two_terminating_tools_blocking_model())
5853 .tool(CountingAddTool {
5854 calls: blocking_calls.clone(),
5855 })
5856 .build()
5857 .runner("go")
5858 .max_turns(3)
5859 .add_hook(TerminateOnFirstToolHook)
5860 .run()
5861 .await
5862 .expect_err("the run terminates");
5863 assert_eq!(
5864 blocking_calls.load(SeqCst),
5865 0,
5866 "fail-fast: blocking run() must not start the second tool after the first terminates"
5867 );
5868
5869 let streaming_calls = Arc::new(AtomicU32::new(0));
5870 let mut stream = AgentBuilder::new(two_terminating_tools_streaming_model())
5871 .tool(CountingAddTool {
5872 calls: streaming_calls.clone(),
5873 })
5874 .build()
5875 .runner("go")
5876 .max_turns(3)
5877 .add_hook(TerminateOnFirstToolHook)
5878 .stream()
5879 .await;
5880 let mut saw_error = false;
5881 while let Some(item) = stream.next().await {
5882 if let Err(err) = item {
5883 saw_error = true;
5884 assert!(
5885 err.to_string().contains("stop"),
5886 "stream() should surface the terminate reason, got: {err}"
5887 );
5888 break;
5889 }
5890 }
5891 assert!(saw_error, "stream() must surface the terminate error");
5892 assert_eq!(
5893 streaming_calls.load(SeqCst),
5894 0,
5895 "fail-fast: stream() must not start the second tool after the first terminates"
5896 );
5897 }
5898
5899 #[derive(Clone)]
5905 struct RecordingArgsTool {
5906 called: Arc<Mutex<Vec<i64>>>,
5907 sibling_started: Arc<tokio::sync::Notify>,
5908 }
5909
5910 impl Tool for RecordingArgsTool {
5911 const NAME: &'static str = "add";
5912 type Error = MockToolError;
5913 type Args = serde_json::Value;
5914 type Output = i32;
5915
5916 fn description(&self) -> String {
5917 MockAddTool.description()
5918 }
5919
5920 fn parameters(&self) -> serde_json::Value {
5921 MockAddTool.parameters()
5922 }
5923
5924 async fn call(
5925 &self,
5926 _context: &mut ToolContext,
5927 args: Self::Args,
5928 ) -> Result<Self::Output, Self::Error> {
5929 let x = args.get("x").and_then(serde_json::Value::as_i64);
5930 if let Some(x) = x {
5931 self.called.lock().expect("called").push(x);
5932 }
5933 if x == Some(1) {
5934 self.sibling_started.notify_one();
5937 for _ in 0..8 {
5938 tokio::task::yield_now().await;
5939 }
5940 }
5941 Ok(0)
5942 }
5943 }
5944
5945 fn three_tools_first_terminates_streaming_model() -> MockCompletionModel {
5946 MockCompletionModel::from_stream_turns([
5947 vec![
5948 MockStreamEvent::tool_call("tc0", "add", json!({"x": 0, "y": 0})),
5953 MockStreamEvent::tool_call("tc1", "add", json!({"x": 1, "y": 1})),
5954 MockStreamEvent::tool_call("tc2", "add", json!({"x": 2, "y": 2})),
5955 MockStreamEvent::final_response_with_total_tokens(0),
5956 ],
5957 vec![
5958 MockStreamEvent::text("unreachable"),
5959 MockStreamEvent::final_response_with_total_tokens(0),
5960 ],
5961 ])
5962 }
5963
5964 struct TerminateOnArgZeroAfterSiblingHook {
5968 sibling_started: Arc<tokio::sync::Notify>,
5969 }
5970 impl AgentHook for TerminateOnArgZeroAfterSiblingHook {
5971 async fn on_tool_call(&self, _ctx: &HookContext, event: ToolCall<'_>) -> ToolCallAction {
5972 if let ToolCall { args, .. } = event
5973 && serde_json::from_str::<serde_json::Value>(args)
5974 .ok()
5975 .and_then(|v| v.get("x").and_then(serde_json::Value::as_i64))
5976 == Some(0)
5977 {
5978 self.sibling_started.notified().await;
5979 return ToolCallAction::stop("stop");
5980 }
5981 ToolCallAction::run()
5982 }
5983 }
5984
5985 #[tokio::test]
5995 async fn concurrent_terminate_drops_beyond_window_sibling_but_drains_in_flight() {
5996 let called = Arc::new(Mutex::new(Vec::new()));
5997 let sibling_started = Arc::new(tokio::sync::Notify::new());
5998 let mut stream = AgentBuilder::new(three_tools_first_terminates_streaming_model())
5999 .tool(RecordingArgsTool {
6000 called: called.clone(),
6001 sibling_started: sibling_started.clone(),
6002 })
6003 .build()
6004 .runner("go")
6005 .max_turns(3)
6006 .tool_concurrency(2)
6007 .add_hook(TerminateOnArgZeroAfterSiblingHook { sibling_started })
6008 .stream()
6009 .await;
6010
6011 let (saw_error, saw_final) =
6012 tokio::time::timeout(std::time::Duration::from_secs(5), async move {
6013 let mut saw_error = false;
6014 let mut saw_final = false;
6015 while let Some(item) = stream.next().await {
6016 match item {
6017 Ok(MultiTurnStreamItem::FinalResponse(_)) => saw_final = true,
6018 Ok(_) => {}
6019 Err(_) => saw_error = true,
6020 }
6021 }
6022 (saw_error, saw_final)
6023 })
6024 .await
6025 .expect("the concurrent tool drive must not hang");
6026
6027 assert!(saw_error, "the terminated run must surface an error");
6028 assert!(
6029 !saw_final,
6030 "a terminated run must not yield a final response"
6031 );
6032 let called = called.lock().expect("called").clone();
6033 assert!(
6034 called.contains(&1),
6035 "the in-flight sibling (x==1) must be drained to completion; called args: {called:?}"
6036 );
6037 assert!(
6038 !called.contains(&2),
6039 "the not-yet-started sibling beyond the concurrency window (x==2) must be \
6040 dropped, not executed; called args: {called:?}"
6041 );
6042 assert!(
6043 !called.contains(&0),
6044 "the terminator's own body never runs (its ToolCall hook terminated); \
6045 called args: {called:?}"
6046 );
6047 }
6048
6049 #[derive(Clone)]
6053 struct SignalOnRunTool {
6054 a_ran: Arc<AtomicU32>,
6055 a_done: Arc<tokio::sync::Notify>,
6056 }
6057 impl Tool for SignalOnRunTool {
6058 const NAME: &'static str = "add";
6059 type Error = MockToolError;
6060 type Args = serde_json::Value;
6061 type Output = i32;
6062 fn description(&self) -> String {
6063 MockAddTool.description()
6064 }
6065
6066 fn parameters(&self) -> serde_json::Value {
6067 MockAddTool.parameters()
6068 }
6069 async fn call(
6070 &self,
6071 _context: &mut ToolContext,
6072 args: Self::Args,
6073 ) -> Result<Self::Output, Self::Error> {
6074 if args.get("x").and_then(serde_json::Value::as_i64) == Some(1) {
6075 self.a_ran.fetch_add(1, SeqCst);
6076 self.a_done.notify_one();
6077 }
6078 Ok(0)
6079 }
6080 }
6081
6082 struct TerminateAfterSiblingDoneHook {
6086 a_done: Arc<tokio::sync::Notify>,
6087 }
6088 impl AgentHook for TerminateAfterSiblingDoneHook {
6089 async fn on_tool_call(&self, _ctx: &HookContext, event: ToolCall<'_>) -> ToolCallAction {
6090 if let ToolCall { args, .. } = event
6091 && serde_json::from_str::<serde_json::Value>(args)
6092 .ok()
6093 .and_then(|v| v.get("x").and_then(serde_json::Value::as_i64))
6094 == Some(2)
6095 {
6096 self.a_done.notified().await;
6097 return ToolCallAction::stop("stop");
6098 }
6099 ToolCallAction::run()
6100 }
6101 }
6102
6103 #[tokio::test]
6110 async fn concurrent_termination_surfaces_no_execution_items() {
6111 let a_ran = Arc::new(AtomicU32::new(0));
6112 let a_done = Arc::new(tokio::sync::Notify::new());
6113 let model = MockCompletionModel::from_stream_turns([
6114 vec![
6115 MockStreamEvent::tool_call("tc1", "add", json!({"x": 1, "y": 1})),
6116 MockStreamEvent::tool_call("tc2", "add", json!({"x": 2, "y": 2})),
6117 MockStreamEvent::final_response_with_total_tokens(0),
6118 ],
6119 vec![
6120 MockStreamEvent::text("unreachable"),
6121 MockStreamEvent::final_response_with_total_tokens(0),
6122 ],
6123 ]);
6124 let mut stream = AgentBuilder::new(model)
6125 .tool(SignalOnRunTool {
6126 a_ran: a_ran.clone(),
6127 a_done: a_done.clone(),
6128 })
6129 .build()
6130 .runner("go")
6131 .max_turns(3)
6132 .tool_concurrency(2)
6133 .add_hook(TerminateAfterSiblingDoneHook {
6134 a_done: a_done.clone(),
6135 })
6136 .stream()
6137 .await;
6138
6139 let (exec_commits, results, saw_error, saw_final) =
6140 tokio::time::timeout(std::time::Duration::from_secs(5), async move {
6141 let (mut exec_commits, mut results, mut saw_error, mut saw_final) =
6142 (0, 0, false, false);
6143 while let Some(item) = stream.next().await {
6144 match item {
6145 Ok(MultiTurnStreamItem::ToolExecutionCommitted { .. }) => exec_commits += 1,
6146 Ok(MultiTurnStreamItem::StreamUserItem(
6147 StreamedUserContent::ToolResult { .. },
6148 )) => results += 1,
6149 Ok(MultiTurnStreamItem::FinalResponse(_)) => saw_final = true,
6150 Ok(_) => {}
6151 Err(_) => saw_error = true,
6152 }
6153 }
6154 (exec_commits, results, saw_error, saw_final)
6155 })
6156 .await
6157 .expect("the concurrent tool drive must not hang");
6158
6159 assert!(saw_error, "the terminated run must surface an error");
6160 assert!(
6161 !saw_final,
6162 "a terminated run must not yield a final response"
6163 );
6164 assert_eq!(
6165 exec_commits, 0,
6166 "a terminated batch surfaces no ToolExecutionCommitted events"
6167 );
6168 assert_eq!(
6169 results, 0,
6170 "a terminated batch surfaces no successful ToolResult"
6171 );
6172 assert_eq!(
6173 a_ran.load(SeqCst),
6174 1,
6175 "the fast sibling did run (its side effect happened), but its result was suppressed"
6176 );
6177 }
6178
6179 #[tokio::test]
6184 async fn stream_tool_execution_committed_carries_effective_rewritten_args() {
6185 let model = MockCompletionModel::from_stream_turns([
6186 vec![
6187 MockStreamEvent::tool_call("tc1", "add", json!({"x": 2, "y": 3})),
6188 MockStreamEvent::final_response_with_total_tokens(0),
6189 ],
6190 vec![
6191 MockStreamEvent::text("done"),
6192 MockStreamEvent::final_response_with_total_tokens(0),
6193 ],
6194 ]);
6195 let mut stream = AgentBuilder::new(model)
6196 .tool(MockAddTool)
6197 .add_hook(RewriteToolArgsHook(json!({"x": 2, "y": 40})))
6198 .build()
6199 .runner("go")
6200 .max_turns(3)
6201 .stream()
6202 .await;
6203
6204 let mut model_args = None;
6205 let mut exec_args = None;
6206 while let Some(item) = stream.next().await {
6207 match item.unwrap_or_else(|err| panic!("stream item errored: {err}")) {
6208 MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::ToolCall {
6209 tool_call,
6210 ..
6211 }) => model_args = Some(tool_call.function.arguments),
6212 MultiTurnStreamItem::ToolExecutionCommitted { tool_call, .. } => {
6213 exec_args = Some(tool_call.function.arguments)
6214 }
6215 _ => {}
6216 }
6217 }
6218 assert_eq!(
6219 model_args,
6220 Some(json!({"x": 2, "y": 3})),
6221 "the model tool-call event carries the model's original arguments"
6222 );
6223 assert_eq!(
6224 exec_args,
6225 Some(json!({"x": 2, "y": 40})),
6226 "the execution-commit event carries the hook-rewritten (effective) arguments"
6227 );
6228 }
6229
6230 #[tokio::test]
6234 async fn stream_hook_skip_surfaces_result_without_execution_commit() {
6235 struct SkipHook;
6236 impl AgentHook for SkipHook {
6237 async fn on_tool_call(
6238 &self,
6239 _ctx: &HookContext,
6240 event: ToolCall<'_>,
6241 ) -> ToolCallAction {
6242 if let ToolCall { .. } = event {
6243 ToolCallAction::skip("blocked by policy")
6244 } else {
6245 ToolCallAction::run()
6246 }
6247 }
6248 }
6249
6250 let calls = Arc::new(AtomicU32::new(0));
6251 let model = MockCompletionModel::from_stream_turns([
6252 vec![
6253 MockStreamEvent::tool_call("tc1", "add", json!({"x": 1, "y": 2})),
6254 MockStreamEvent::final_response_with_total_tokens(0),
6255 ],
6256 vec![
6257 MockStreamEvent::text("done"),
6258 MockStreamEvent::final_response_with_total_tokens(0),
6259 ],
6260 ]);
6261 let stream = AgentBuilder::new(model)
6262 .tool(CountingAddTool {
6263 calls: calls.clone(),
6264 })
6265 .add_hook(SkipHook)
6266 .build()
6267 .runner("go")
6268 .max_turns(3)
6269 .stream()
6270 .await;
6271
6272 let mut exec_commits = 0;
6273 let mut results = 0;
6274 let mut final_response = None;
6275 let mut stream = stream;
6276 while let Some(item) = stream.next().await {
6277 match item.unwrap_or_else(|err| panic!("stream item errored: {err}")) {
6278 MultiTurnStreamItem::ToolExecutionCommitted { .. } => exec_commits += 1,
6279 MultiTurnStreamItem::StreamUserItem(StreamedUserContent::ToolResult { .. }) => {
6280 results += 1
6281 }
6282 MultiTurnStreamItem::FinalResponse(resp) => final_response = Some(resp),
6283 _ => {}
6284 }
6285 }
6286
6287 assert_eq!(calls.load(SeqCst), 0, "a skipped tool's body never runs");
6288 assert_eq!(
6289 exec_commits, 0,
6290 "a hook-skipped tool produces no execution-commit"
6291 );
6292 assert_eq!(
6293 results, 1,
6294 "the skip result is still surfaced to the consumer"
6295 );
6296 let final_response = final_response.expect("stream should yield a final response");
6297 let history = final_response.messages().expect("history");
6299 assert!(
6300 history.iter().any(|m| serde_json::to_string(m)
6301 .map(|s| s.contains("blocked by policy"))
6302 .unwrap_or(false)),
6303 "the skip result is committed to history"
6304 );
6305 }
6306
6307 #[tokio::test]
6310 async fn required_with_empty_active_tools_errors_locally_without_provider_call() {
6311 struct EmptyActiveToolsHook;
6312 impl AgentHook for EmptyActiveToolsHook {
6313 async fn on_completion_call(
6314 &self,
6315 _ctx: &HookContext,
6316 event: CompletionCallEvent<'_>,
6317 ) -> CompletionCallAction {
6318 if let CompletionCallEvent { .. } = event {
6319 CompletionCallAction::patch(
6320 RequestPatch::new().active_tools(Vec::<String>::new()),
6321 )
6322 } else {
6323 CompletionCallAction::continue_run()
6324 }
6325 }
6326 }
6327
6328 let model = MockCompletionModel::from_turns([MockTurn::text("unreachable")]);
6329 let probe = model.clone();
6330 let err = AgentBuilder::new(model)
6331 .tool(MockAddTool)
6332 .tool_choice(ToolChoice::Required)
6333 .add_hook(EmptyActiveToolsHook)
6334 .build()
6335 .runner("go")
6336 .run()
6337 .await
6338 .expect_err("Required with an empty active_tools filter must fail locally");
6339
6340 assert!(
6341 probe.requests().is_empty(),
6342 "the request must fail locally, with no provider round-trip"
6343 );
6344 let msg = err.to_string();
6345 assert!(
6346 msg.contains("Required"),
6347 "error should mention Required: {msg}"
6348 );
6349 assert!(
6350 msg.contains("active_tools"),
6351 "error should name active_tools: {msg}"
6352 );
6353 }
6354
6355 #[tokio::test]
6358 async fn specific_naming_filtered_out_tool_errors_locally_without_provider_call() {
6359 struct FilterToAddHook;
6360 impl AgentHook for FilterToAddHook {
6361 async fn on_completion_call(
6362 &self,
6363 _ctx: &HookContext,
6364 event: CompletionCallEvent<'_>,
6365 ) -> CompletionCallAction {
6366 if let CompletionCallEvent { .. } = event {
6367 CompletionCallAction::patch(RequestPatch::new().active_tools(["add"]))
6368 } else {
6369 CompletionCallAction::continue_run()
6370 }
6371 }
6372 }
6373
6374 let model = MockCompletionModel::from_turns([MockTurn::text("unreachable")]);
6375 let probe = model.clone();
6376 let err = AgentBuilder::new(model)
6377 .tool(MockAddTool)
6378 .tool(MockSubtractTool)
6379 .tool_choice(ToolChoice::Specific {
6380 function_names: vec!["subtract".to_string()],
6381 })
6382 .add_hook(FilterToAddHook)
6383 .build()
6384 .runner("go")
6385 .run()
6386 .await
6387 .expect_err("Specific naming a filtered-out tool must fail locally");
6388
6389 assert!(
6390 probe.requests().is_empty(),
6391 "the request must fail locally, with no provider round-trip"
6392 );
6393 let msg = err.to_string();
6394 assert!(
6395 msg.contains("subtract"),
6396 "error should name the missing tool: {msg}"
6397 );
6398 assert!(
6399 msg.contains("active_tools"),
6400 "error should name active_tools: {msg}"
6401 );
6402 }
6403
6404 #[tokio::test]
6411 async fn concurrent_tool_execution_stays_within_the_configured_bound() {
6412 #[derive(Clone)]
6413 struct ConcurrencyProbe {
6414 barrier: Arc<Barrier>,
6415 active: Arc<AtomicU32>,
6416 max_active: Arc<AtomicU32>,
6417 }
6418
6419 impl Tool for ConcurrencyProbe {
6420 const NAME: &'static str = "add";
6421 type Error = MockToolError;
6422 type Args = serde_json::Value;
6423 type Output = String;
6424
6425 fn description(&self) -> String {
6426 "concurrency probe".to_string()
6427 }
6428
6429 fn parameters(&self) -> serde_json::Value {
6430 json!({"type": "object", "properties": {}})
6431 }
6432
6433 async fn call(
6434 &self,
6435 _context: &mut ToolContext,
6436 _args: Self::Args,
6437 ) -> Result<Self::Output, Self::Error> {
6438 let now = self.active.fetch_add(1, SeqCst) + 1;
6439 self.max_active.fetch_max(now, SeqCst);
6440 self.barrier.wait().await;
6441 self.active.fetch_sub(1, SeqCst);
6442 Ok("ok".to_string())
6443 }
6444 }
6445
6446 let cap = 2usize;
6447 let probe = ConcurrencyProbe {
6448 barrier: Arc::new(Barrier::new(cap)),
6449 active: Arc::new(AtomicU32::new(0)),
6450 max_active: Arc::new(AtomicU32::new(0)),
6451 };
6452 let max_active = probe.max_active.clone();
6453
6454 let model = MockCompletionModel::from_turns([
6456 MockTurn::from_contents([
6457 tool_call_content("c1", json!({})),
6458 tool_call_content("c2", json!({})),
6459 tool_call_content("c3", json!({})),
6460 tool_call_content("c4", json!({})),
6461 ]),
6462 MockTurn::text("done"),
6463 ]);
6464
6465 let _ = AgentBuilder::new(model)
6466 .tool(probe)
6467 .build()
6468 .runner("probe concurrency")
6469 .max_turns(3)
6470 .tool_concurrency(cap)
6471 .run()
6472 .await
6473 .expect("run should succeed");
6474
6475 let observed = max_active.load(SeqCst);
6476 assert!(
6477 observed > 1,
6478 "tools actually ran concurrently (lower bound): max_active={observed}"
6479 );
6480 assert!(
6481 observed <= cap as u32,
6482 "in-flight never exceeded the configured bound {cap} (upper bound): max_active={observed}"
6483 );
6484 }
6485
6486 #[tokio::test]
6491 async fn tool_concurrency_zero_is_clamped_and_does_not_hang() {
6492 let model = MockCompletionModel::from_turns([
6493 MockTurn::tool_call("tc1", "add", json!({"x": 1, "y": 2})),
6494 MockTurn::text("done"),
6495 ]);
6496 let run = AgentBuilder::new(model)
6497 .tool(MockAddTool)
6498 .build()
6499 .runner("add")
6500 .max_turns(3)
6501 .tool_concurrency(0)
6502 .run();
6503
6504 let response = tokio::time::timeout(std::time::Duration::from_secs(5), run)
6505 .await
6506 .expect("tool_concurrency(0) must clamp to 1, not hang on buffer_unordered(0)")
6507 .expect("run should succeed");
6508 assert_eq!(response.output, "done");
6509 }
6510
6511 #[derive(Clone)]
6513 struct CountingAddTool {
6514 calls: Arc<AtomicU32>,
6515 }
6516 impl Tool for CountingAddTool {
6517 const NAME: &'static str = "add";
6518 type Error = MockToolError;
6519 type Args = MockOperationArgs;
6520 type Output = i32;
6521 fn description(&self) -> String {
6522 MockAddTool.description()
6523 }
6524 fn parameters(&self) -> serde_json::Value {
6525 MockAddTool.parameters()
6526 }
6527 async fn call(
6528 &self,
6529 _context: &mut ToolContext,
6530 args: Self::Args,
6531 ) -> Result<Self::Output, Self::Error> {
6532 self.calls.fetch_add(1, SeqCst);
6533 MockAddTool.call(_context, args).await
6534 }
6535 }
6536
6537 #[derive(Clone, Default)]
6538 struct ToolOnlyHook {
6539 text_delta_calls: Arc<AtomicU32>,
6540 other_calls: Arc<AtomicU32>,
6541 }
6542
6543 impl AgentHook for ToolOnlyHook {
6544 async fn on_text_delta(&self, _: &HookContext, _: TextDelta<'_>) -> ObservationAction {
6545 self.text_delta_calls.fetch_add(1, SeqCst);
6546 ObservationAction::continue_run()
6547 }
6548 async fn on_completion_call(
6549 &self,
6550 _: &HookContext,
6551 _: CompletionCallEvent<'_>,
6552 ) -> CompletionCallAction {
6553 self.other_calls.fetch_add(1, SeqCst);
6554 CompletionCallAction::continue_run()
6555 }
6556 fn observes(&self, kind: StepEventKind) -> bool {
6557 kind != StepEventKind::TextDelta
6558 }
6559 }
6560
6561 #[tokio::test]
6565 async fn observes_gates_text_delta_dispatch() {
6566 let model = MockCompletionModel::from_stream_turns([vec![
6567 MockStreamEvent::text("hel"),
6568 MockStreamEvent::text("lo"),
6569 MockStreamEvent::final_response_with_total_tokens(0),
6570 ]]);
6571 let hook = ToolOnlyHook::default();
6572 let mut stream = AgentBuilder::new(model)
6573 .build()
6574 .runner("hi")
6575 .add_hook(hook.clone())
6576 .stream()
6577 .await;
6578 while stream.next().await.is_some() {}
6579
6580 assert_eq!(
6581 hook.text_delta_calls.load(SeqCst),
6582 0,
6583 "a hook that does not observe TextDelta must not be dispatched for it"
6584 );
6585 assert!(
6586 hook.other_calls.load(SeqCst) > 0,
6587 "the hook should still receive the events it observes"
6588 );
6589 }
6590
6591 struct TerminateOn(StepEventKind);
6594
6595 impl AgentHook for TerminateOn {
6596 async fn on_completion_call(
6597 &self,
6598 _: &HookContext,
6599 _: CompletionCallEvent<'_>,
6600 ) -> CompletionCallAction {
6601 if self.0 == StepEventKind::CompletionCall {
6602 CompletionCallAction::stop("stop here")
6603 } else {
6604 CompletionCallAction::continue_run()
6605 }
6606 }
6607 async fn on_tool_call(&self, _: &HookContext, _: ToolCall<'_>) -> ToolCallAction {
6608 if self.0 == StepEventKind::ToolCall {
6609 ToolCallAction::stop("stop here")
6610 } else {
6611 ToolCallAction::run()
6612 }
6613 }
6614 async fn on_tool_result(
6615 &self,
6616 _: &HookContext,
6617 _: ToolResultEvent<'_>,
6618 ) -> ToolResultAction {
6619 if self.0 == StepEventKind::ToolResult {
6620 ToolResultAction::stop("stop here")
6621 } else {
6622 ToolResultAction::keep()
6623 }
6624 }
6625 }
6626
6627 #[tokio::test]
6631 async fn run_terminates_from_each_shared_event() {
6632 for kind in [
6633 StepEventKind::CompletionCall,
6634 StepEventKind::ToolCall,
6635 StepEventKind::ToolResult,
6636 ] {
6637 let err = AgentBuilder::new(blocking_model())
6638 .tool(MockAddTool)
6639 .build()
6640 .runner("add 2 and 3")
6641 .max_turns(3)
6642 .add_hook(TerminateOn(kind))
6643 .run()
6644 .await
6645 .expect_err(&format!("terminate at {kind:?} must cancel the run"));
6646 assert!(
6647 matches!(err, PromptError::PromptCancelled { .. }),
6648 "terminate at {kind:?} should cancel the run, got {err:?}"
6649 );
6650 }
6651 }
6652
6653 #[tokio::test]
6657 async fn stream_terminates_from_each_shared_event() {
6658 for kind in [
6659 StepEventKind::CompletionCall,
6660 StepEventKind::ToolCall,
6661 StepEventKind::ToolResult,
6662 ] {
6663 let mut stream = AgentBuilder::new(streaming_model())
6664 .tool(MockAddTool)
6665 .build()
6666 .runner("add 2 and 3")
6667 .max_turns(3)
6668 .add_hook(TerminateOn(kind))
6669 .stream()
6670 .await;
6671
6672 let mut saw_error = false;
6673 let mut saw_final = false;
6674 while let Some(item) = stream.next().await {
6675 match item {
6676 Ok(MultiTurnStreamItem::FinalResponse(_)) => saw_final = true,
6677 Err(_) => saw_error = true,
6678 _ => {}
6679 }
6680 }
6681 assert!(saw_error, "terminate at {kind:?} must yield a stream error");
6682 assert!(
6683 !saw_final,
6684 "terminate at {kind:?} must not also produce a final response"
6685 );
6686 }
6687 }
6688
6689 #[tokio::test]
6693 async fn multi_hook_stack_parity_across_run_and_stream() {
6694 let a_block = RecordingHook::default();
6695 let b_block = RecordingHook::default();
6696 let blocking = AgentBuilder::new(blocking_model())
6697 .tool(MockAddTool)
6698 .build()
6699 .runner("add 2 and 3")
6700 .max_turns(3)
6701 .add_hook(a_block.clone())
6702 .add_hook(b_block.clone())
6703 .run()
6704 .await
6705 .expect("blocking run should succeed");
6706
6707 let a_stream = RecordingHook::default();
6708 let b_stream = RecordingHook::default();
6709 let mut stream = AgentBuilder::new(streaming_model())
6710 .tool(MockAddTool)
6711 .build()
6712 .runner("add 2 and 3")
6713 .max_turns(3)
6714 .add_hook(a_stream.clone())
6715 .add_hook(b_stream.clone())
6716 .stream()
6717 .await;
6718 while stream.next().await.is_some() {}
6719
6720 assert_eq!(a_block.shared_events(), b_block.shared_events());
6722 assert_eq!(a_stream.shared_events(), b_stream.shared_events());
6723 assert_eq!(a_block.shared_events(), a_stream.shared_events());
6725 assert_eq!(
6726 a_block.shared_events(),
6727 vec![
6728 StepEventKind::CompletionCall,
6729 StepEventKind::ToolCall,
6730 StepEventKind::ToolResult,
6731 StepEventKind::CompletionCall,
6732 ]
6733 );
6734 assert_eq!(blocking.output, "the answer is 5");
6735 }
6736
6737 struct RepairInvalidToHook(&'static str);
6739
6740 impl AgentHook for RepairInvalidToHook {
6741 async fn on_invalid_tool_call(
6742 &self,
6743 _ctx: &HookContext,
6744 event: &InvalidToolCallContext,
6745 ) -> Option<InvalidToolCallAction> {
6746 Some(if let _ = event {
6747 InvalidToolCallAction::repair(self.0)
6748 } else {
6749 InvalidToolCallAction::fail()
6750 })
6751 }
6752 }
6753
6754 #[derive(Clone)]
6755 struct CaptureAndRepairInvalidHook {
6756 replacement: &'static str,
6757 args: Arc<Mutex<Vec<Option<String>>>>,
6758 }
6759
6760 impl AgentHook for CaptureAndRepairInvalidHook {
6761 async fn on_invalid_tool_call(
6762 &self,
6763 _ctx: &HookContext,
6764 event: &InvalidToolCallContext,
6765 ) -> Option<InvalidToolCallAction> {
6766 self.args
6767 .lock()
6768 .expect("invalid args")
6769 .push(event.args.clone());
6770 Some(InvalidToolCallAction::repair(self.replacement))
6771 }
6772 }
6773
6774 #[tokio::test]
6778 async fn invalid_tool_call_repair_parity_across_run_and_stream() {
6779 let blocking_model = MockCompletionModel::from_turns([
6780 MockTurn::tool_call("tc1", "default_api", json!({"x": 2, "y": 3})),
6781 MockTurn::text("the answer is 5"),
6782 ]);
6783 let blocking_hook = RecordingHook::default();
6784 let blocking = AgentBuilder::new(blocking_model)
6785 .tool(MockAddTool)
6786 .build()
6787 .runner("add 2 and 3")
6788 .max_turns(3)
6789 .add_hook(blocking_hook.clone())
6790 .add_hook(RepairInvalidToHook("add"))
6791 .run()
6792 .await
6793 .expect("blocking run should recover via repair");
6794
6795 let streaming_model = MockCompletionModel::from_stream_turns([
6801 vec![
6802 MockStreamEvent::tool_call("tc1", "default_api", json!({"x": 2, "y": 3})),
6803 MockStreamEvent::final_response_with_total_tokens(0),
6804 ],
6805 vec![
6806 MockStreamEvent::text("the answer is 5"),
6807 MockStreamEvent::final_response_with_total_tokens(0),
6808 ],
6809 ]);
6810 let streaming_hook = RecordingHook::default();
6811 let mut stream = AgentBuilder::new(streaming_model)
6812 .tool(MockAddTool)
6813 .build()
6814 .runner("add 2 and 3")
6815 .max_turns(3)
6816 .add_hook(streaming_hook.clone())
6817 .add_hook(RepairInvalidToHook("add"))
6818 .stream()
6819 .await;
6820 let mut final_response = None;
6821 while let Some(item) = stream.next().await {
6822 if let Ok(MultiTurnStreamItem::FinalResponse(resp)) =
6823 item.map_err(|err| panic!("stream item errored: {err}"))
6824 {
6825 final_response = Some(resp);
6826 }
6827 }
6828 let final_response =
6829 final_response.expect("stream should recover and yield a final response");
6830
6831 assert_eq!(blocking.output, "the answer is 5");
6833 assert_eq!(final_response.output(), blocking.output);
6834
6835 assert_eq!(
6838 blocking_hook.shared_events(),
6839 streaming_hook.shared_events()
6840 );
6841 assert!(
6842 blocking_hook
6843 .shared_events()
6844 .contains(&StepEventKind::InvalidToolCall),
6845 "the hook must observe the invalid tool call"
6846 );
6847 assert_eq!(blocking_hook.tool_results(), streaming_hook.tool_results());
6848 assert_eq!(blocking_hook.tool_results(), vec!["5".to_string()]);
6849
6850 let blocking_messages = blocking.messages.expect("blocking messages");
6852 let streaming_messages = final_response
6853 .messages()
6854 .expect("streaming history")
6855 .to_vec();
6856 assert_eq!(
6857 serde_json::to_value(&blocking_messages).expect("serialize blocking"),
6858 serde_json::to_value(&streaming_messages).expect("serialize streaming"),
6859 );
6860 }
6861
6862 #[tokio::test]
6863 async fn invalid_tool_call_scalar_args_are_canonical_across_run_and_complete_stream() {
6864 let blocking_args = Arc::new(Mutex::new(Vec::new()));
6865 let blocking_hook = RecordingHook::default();
6866 let blocking = AgentBuilder::new(MockCompletionModel::from_turns([
6867 MockTurn::tool_call("tc1", "unknown_echo", json!("payload")),
6868 MockTurn::text("done"),
6869 ]))
6870 .tool(EchoStringArgs)
6871 .build()
6872 .runner("echo a string")
6873 .max_turns(3)
6874 .add_hook(blocking_hook.clone())
6875 .add_hook(CaptureAndRepairInvalidHook {
6876 replacement: EchoStringArgs::NAME,
6877 args: blocking_args.clone(),
6878 })
6879 .run()
6880 .await
6881 .expect("blocking scalar repair should succeed");
6882
6883 let streaming_args = Arc::new(Mutex::new(Vec::new()));
6884 let streaming_hook = RecordingHook::default();
6885 let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([
6886 vec![
6887 MockStreamEvent::tool_call("tc1", "unknown_echo", json!("payload")),
6888 MockStreamEvent::final_response_with_total_tokens(0),
6889 ],
6890 vec![
6891 MockStreamEvent::text("done"),
6892 MockStreamEvent::final_response_with_total_tokens(0),
6893 ],
6894 ]))
6895 .tool(EchoStringArgs)
6896 .build()
6897 .runner("echo a string")
6898 .max_turns(3)
6899 .add_hook(streaming_hook.clone())
6900 .add_hook(CaptureAndRepairInvalidHook {
6901 replacement: EchoStringArgs::NAME,
6902 args: streaming_args.clone(),
6903 })
6904 .stream()
6905 .await;
6906 let mut final_response = None;
6907 while let Some(item) = stream.next().await {
6908 if let MultiTurnStreamItem::FinalResponse(response) =
6909 item.expect("streaming scalar repair should succeed")
6910 {
6911 final_response = Some(response);
6912 }
6913 }
6914 let final_response = final_response.expect("stream should yield a final response");
6915
6916 let canonical_args = vec![Some(serde_json::to_string("payload").unwrap())];
6917 assert_eq!(*blocking_args.lock().unwrap(), canonical_args);
6918 assert_eq!(*streaming_args.lock().unwrap(), canonical_args);
6919 assert_eq!(blocking_hook.tool_results(), vec!["payload"]);
6920 assert_eq!(streaming_hook.tool_results(), vec!["payload"]);
6921 assert_eq!(blocking.output, "done");
6922 assert_eq!(final_response.output(), "done");
6923 assert_eq!(
6924 serde_json::to_value(blocking.messages.expect("blocking history")).unwrap(),
6925 serde_json::to_value(final_response.messages().expect("streaming history")).unwrap()
6926 );
6927 }
6928
6929 #[derive(Clone)]
6951 struct ScriptedToolCall {
6952 id: &'static str,
6953 name: &'static str,
6954 args: serde_json::Value,
6955 }
6956
6957 #[derive(Clone)]
6960 enum ScriptedTurn {
6961 Text(&'static str),
6963 ToolCalls(Vec<ScriptedToolCall>),
6965 }
6966
6967 #[derive(Clone, Copy)]
6973 enum StreamShape {
6974 Complete,
6976 Chunked,
6979 }
6980
6981 impl ScriptedTurn {
6982 fn as_blocking_turn(&self) -> MockTurn {
6983 match self {
6984 ScriptedTurn::Text(text) => MockTurn::text(*text),
6985 ScriptedTurn::ToolCalls(calls) => {
6986 MockTurn::from_contents(calls.iter().map(|call| {
6989 AssistantContent::ToolCall(MessageToolCall::from_wire(
6990 call.id,
6991 ToolFunction::new(call.name.to_string(), call.args.clone()),
6992 ))
6993 }))
6994 }
6995 }
6996 }
6997
6998 fn as_stream_events(&self, shape: StreamShape) -> Vec<MockStreamEvent> {
6999 let mut events = Vec::new();
7000 match self {
7001 ScriptedTurn::Text(text) => events.push(MockStreamEvent::text(*text)),
7002 ScriptedTurn::ToolCalls(calls) => {
7003 for call in calls {
7004 if let StreamShape::Chunked = shape {
7005 let args = serde_json::to_string(&call.args)
7009 .expect("scripted args serialize to json");
7010 events.push(MockStreamEvent::tool_call_name_delta(call.id, call.name));
7011 events.push(MockStreamEvent::tool_call_arguments_delta(call.id, &args));
7012 }
7013 events.push(MockStreamEvent::tool_call(
7014 call.id,
7015 call.name,
7016 call.args.clone(),
7017 ));
7018 }
7019 }
7020 }
7021 events.push(MockStreamEvent::final_response_with_total_tokens(0));
7022 events
7023 }
7024 }
7025
7026 struct ParityOutcome {
7029 output: String,
7030 messages: Vec<Message>,
7031 shared_events: Vec<StepEventKind>,
7032 tool_results: Vec<String>,
7033 }
7034
7035 async fn run_blocking_scenario(prompt: &'static str, turns: &[ScriptedTurn]) -> ParityOutcome {
7036 let model =
7037 MockCompletionModel::from_turns(turns.iter().map(ScriptedTurn::as_blocking_turn));
7038 let hook = RecordingHook::default();
7039 let response = AgentBuilder::new(model)
7040 .tool(MockAddTool)
7041 .build()
7042 .runner(prompt)
7043 .max_turns(8)
7044 .add_hook(hook.clone())
7045 .run()
7046 .await
7047 .expect("blocking scenario should succeed");
7048 ParityOutcome {
7049 output: response.output,
7050 messages: response.messages.expect("blocking messages"),
7051 shared_events: hook.shared_events(),
7052 tool_results: hook.tool_results(),
7053 }
7054 }
7055
7056 async fn run_streaming_scenario(
7057 prompt: &'static str,
7058 turns: &[ScriptedTurn],
7059 shape: StreamShape,
7060 ) -> ParityOutcome {
7061 let model = MockCompletionModel::from_stream_turns(
7062 turns.iter().map(|turn| turn.as_stream_events(shape)),
7063 );
7064 let hook = RecordingHook::default();
7065 let mut stream = AgentBuilder::new(model)
7066 .tool(MockAddTool)
7067 .build()
7068 .runner(prompt)
7069 .max_turns(8)
7070 .add_hook(hook.clone())
7071 .stream()
7072 .await;
7073 let mut final_response = None;
7074 while let Some(item) = stream.next().await {
7075 if let Ok(MultiTurnStreamItem::FinalResponse(resp)) =
7076 item.map_err(|err| panic!("stream item errored: {err}"))
7077 {
7078 final_response = Some(resp);
7079 }
7080 }
7081 let final_response =
7082 final_response.expect("streaming scenario should yield a final response");
7083 ParityOutcome {
7084 output: final_response.output().to_string(),
7085 messages: final_response
7086 .messages()
7087 .expect("streaming history")
7088 .to_vec(),
7089 shared_events: hook.shared_events(),
7090 tool_results: hook.tool_results(),
7091 }
7092 }
7093
7094 fn assert_outcomes_match(blocking: &ParityOutcome, streaming: &ParityOutcome, label: &str) {
7095 assert_eq!(
7096 blocking.output, streaming.output,
7097 "{label}: final output diverged"
7098 );
7099 assert_eq!(
7100 blocking.shared_events, streaming.shared_events,
7101 "{label}: hook event sequence diverged"
7102 );
7103 assert_eq!(
7104 blocking.tool_results, streaming.tool_results,
7105 "{label}: tool-result content diverged"
7106 );
7107 assert_eq!(
7108 serde_json::to_value(&blocking.messages).expect("serialize blocking"),
7109 serde_json::to_value(&streaming.messages).expect("serialize streaming"),
7110 "{label}: message history diverged"
7111 );
7112 }
7113
7114 async fn assert_run_stream_parity(prompt: &'static str, turns: &[ScriptedTurn]) {
7119 let blocking = run_blocking_scenario(prompt, turns).await;
7120 for (shape, label) in [
7121 (StreamShape::Complete, "complete-stream"),
7122 (StreamShape::Chunked, "chunked-stream"),
7123 ] {
7124 let streaming = run_streaming_scenario(prompt, turns, shape).await;
7125 assert_outcomes_match(&blocking, &streaming, label);
7126 }
7127 }
7128
7129 fn add_call(id: &'static str, x: i64, y: i64) -> ScriptedToolCall {
7130 ScriptedToolCall {
7131 id,
7132 name: "add",
7133 args: json!({ "x": x, "y": y }),
7134 }
7135 }
7136
7137 #[tokio::test]
7138 async fn parity_text_only_run() {
7139 assert_run_stream_parity("just say hi", &[ScriptedTurn::Text("hi there")]).await;
7140 }
7141
7142 #[tokio::test]
7143 async fn parity_single_tool_then_text() {
7144 assert_run_stream_parity(
7145 "add 2 and 3",
7146 &[
7147 ScriptedTurn::ToolCalls(vec![add_call("tc1", 2, 3)]),
7148 ScriptedTurn::Text("the answer is 5"),
7149 ],
7150 )
7151 .await;
7152 }
7153
7154 #[tokio::test]
7155 async fn parity_multiple_tools_in_one_turn() {
7156 assert_run_stream_parity(
7157 "add two pairs",
7158 &[
7159 ScriptedTurn::ToolCalls(vec![add_call("tc1", 2, 3), add_call("tc2", 10, 20)]),
7160 ScriptedTurn::Text("done"),
7161 ],
7162 )
7163 .await;
7164 }
7165
7166 #[tokio::test]
7167 async fn parity_multi_turn_sequential_tools() {
7168 assert_run_stream_parity(
7169 "chain two additions",
7170 &[
7171 ScriptedTurn::ToolCalls(vec![add_call("tc1", 1, 1)]),
7172 ScriptedTurn::ToolCalls(vec![add_call("tc2", 2, 2)]),
7173 ScriptedTurn::Text("chained"),
7174 ],
7175 )
7176 .await;
7177 }
7178
7179 struct SkipInvalidHook(&'static str);
7182
7183 impl AgentHook for SkipInvalidHook {
7184 async fn on_invalid_tool_call(
7185 &self,
7186 _ctx: &HookContext,
7187 event: &InvalidToolCallContext,
7188 ) -> Option<InvalidToolCallAction> {
7189 Some(if let _ = event {
7190 InvalidToolCallAction::skip(self.0)
7191 } else {
7192 InvalidToolCallAction::fail()
7193 })
7194 }
7195 }
7196
7197 #[tokio::test]
7202 async fn invalid_tool_call_skip_parity_across_run_and_stream() {
7203 let blocking_model = MockCompletionModel::from_turns([
7204 MockTurn::tool_call("tc1", "default_api", json!({"x": 2, "y": 3})),
7205 MockTurn::text("acknowledged"),
7206 ]);
7207 let blocking_hook = RecordingHook::default();
7208 let blocking = AgentBuilder::new(blocking_model)
7209 .tool(MockAddTool)
7210 .build()
7211 .runner("do the thing")
7212 .max_turns(3)
7213 .add_hook(blocking_hook.clone())
7214 .add_hook(SkipInvalidHook("tool not permitted"))
7215 .run()
7216 .await
7217 .expect("blocking run should recover via skip");
7218
7219 let streaming_model = MockCompletionModel::from_stream_turns([
7222 vec![
7223 MockStreamEvent::tool_call("tc1", "default_api", json!({"x": 2, "y": 3})),
7224 MockStreamEvent::final_response_with_total_tokens(0),
7225 ],
7226 vec![
7227 MockStreamEvent::text("acknowledged"),
7228 MockStreamEvent::final_response_with_total_tokens(0),
7229 ],
7230 ]);
7231 let streaming_hook = RecordingHook::default();
7232 let mut stream = AgentBuilder::new(streaming_model)
7233 .tool(MockAddTool)
7234 .build()
7235 .runner("do the thing")
7236 .max_turns(3)
7237 .add_hook(streaming_hook.clone())
7238 .add_hook(SkipInvalidHook("tool not permitted"))
7239 .stream()
7240 .await;
7241 let mut final_response = None;
7242 while let Some(item) = stream.next().await {
7243 if let Ok(MultiTurnStreamItem::FinalResponse(resp)) =
7244 item.map_err(|err| panic!("stream item errored: {err}"))
7245 {
7246 final_response = Some(resp);
7247 }
7248 }
7249 let final_response =
7250 final_response.expect("stream should recover and yield a final response");
7251
7252 assert_eq!(blocking.output, "acknowledged");
7253 assert_eq!(final_response.output(), blocking.output);
7254 assert_eq!(
7255 blocking_hook.shared_events(),
7256 streaming_hook.shared_events()
7257 );
7258 assert!(
7259 blocking_hook
7260 .shared_events()
7261 .contains(&StepEventKind::InvalidToolCall),
7262 "the hook must observe the invalid tool call"
7263 );
7264
7265 let blocking_messages = blocking.messages.expect("blocking messages");
7266 let streaming_messages = final_response
7267 .messages()
7268 .expect("streaming history")
7269 .to_vec();
7270 assert_eq!(
7271 serde_json::to_value(&blocking_messages).expect("serialize blocking"),
7272 serde_json::to_value(&streaming_messages).expect("serialize streaming"),
7273 );
7274 assert!(
7277 tool_result_text_in_history(&blocking_messages, "tool not permitted"),
7278 "the verbatim invalid-tool skip reason must be the tool result content"
7279 );
7280 }
7281
7282 #[tokio::test]
7289 async fn recovered_turn_suppresses_response_finish_hook_on_both_drivers() {
7290 let blocking_model = MockCompletionModel::from_turns([
7294 MockTurn::from_contents([
7295 AssistantContent::text("let me compute that"),
7296 AssistantContent::ToolCall(MessageToolCall::from_wire(
7297 "tc1",
7298 ToolFunction::new("default_api".to_string(), json!({"x": 2, "y": 3})),
7299 )),
7300 ]),
7301 MockTurn::text("the answer is 5"),
7302 ]);
7303 let blocking_hook = RecordingHook::default();
7304 let blocking = AgentBuilder::new(blocking_model)
7305 .tool(MockAddTool)
7306 .build()
7307 .runner("compute")
7308 .max_turns(3)
7309 .add_hook(blocking_hook.clone())
7310 .add_hook(RepairInvalidToHook("add"))
7311 .run()
7312 .await
7313 .expect("blocking run should recover via repair");
7314
7315 let streaming_model = MockCompletionModel::from_stream_turns([
7316 vec![
7317 MockStreamEvent::text("let me compute that"),
7318 MockStreamEvent::tool_call("tc1", "default_api", json!({"x": 2, "y": 3})),
7319 MockStreamEvent::final_response_with_total_tokens(0),
7320 ],
7321 vec![
7322 MockStreamEvent::text("the answer is 5"),
7323 MockStreamEvent::final_response_with_total_tokens(0),
7324 ],
7325 ]);
7326 let streaming_hook = RecordingHook::default();
7327 let mut stream = AgentBuilder::new(streaming_model)
7328 .tool(MockAddTool)
7329 .build()
7330 .runner("compute")
7331 .max_turns(3)
7332 .add_hook(streaming_hook.clone())
7333 .add_hook(RepairInvalidToHook("add"))
7334 .stream()
7335 .await;
7336 while stream.next().await.is_some() {}
7337
7338 assert_eq!(blocking.output, "the answer is 5");
7340
7341 assert_eq!(
7344 blocking_hook.count(StepEventKind::CompletionResponse),
7345 1,
7346 "the recovered turn must not fire CompletionResponse"
7347 );
7348 assert_eq!(
7351 streaming_hook.count(StepEventKind::StreamResponseFinish),
7352 1,
7353 "the recovered turn must not fire StreamResponseFinish"
7354 );
7355 assert_eq!(
7358 blocking_hook.count(StepEventKind::CompletionResponse),
7359 streaming_hook.count(StepEventKind::StreamResponseFinish),
7360 );
7361
7362 assert_eq!(
7369 blocking_hook.count(StepEventKind::ModelTurnFinished),
7370 1,
7371 "the recovered turn must not fire ModelTurnFinished"
7372 );
7373 assert_eq!(
7374 streaming_hook.count(StepEventKind::ModelTurnFinished),
7375 1,
7376 "the recovered turn must not fire ModelTurnFinished on the streaming surface either"
7377 );
7378 assert_eq!(
7381 blocking_hook.count(StepEventKind::ModelTurnFinished),
7382 streaming_hook.count(StepEventKind::ModelTurnFinished),
7383 );
7384 }
7385
7386 #[tokio::test]
7391 async fn runner_add_hook_appends_to_agent_default_hooks() {
7392 let agent_hook = RecordingHook::default();
7393 let runner_hook = RecordingHook::default();
7394
7395 AgentBuilder::new(blocking_model())
7399 .tool(MockAddTool)
7400 .add_hook(agent_hook.clone())
7401 .build()
7402 .runner("add 2 and 3")
7403 .max_turns(3)
7404 .add_hook(runner_hook.clone())
7405 .run()
7406 .await
7407 .expect("run should succeed");
7408
7409 assert!(
7410 agent_hook.count(StepEventKind::CompletionCall) >= 1,
7411 "the agent-default hook must still observe the run after a runner-level add_hook"
7412 );
7413 assert!(
7414 runner_hook.count(StepEventKind::CompletionCall) >= 1,
7415 "the runner-level hook must also observe the run"
7416 );
7417 assert_eq!(
7420 agent_hook.count(StepEventKind::CompletionCall),
7421 runner_hook.count(StepEventKind::CompletionCall),
7422 "add_hook appends (both hooks observe every turn); it does not replace"
7423 );
7424 }
7425
7426 struct SkipToolCallHook(&'static str);
7428
7429 impl AgentHook for SkipToolCallHook {
7430 async fn on_tool_call(&self, _ctx: &HookContext, event: ToolCall<'_>) -> ToolCallAction {
7431 if let ToolCall { .. } = event {
7432 ToolCallAction::skip(self.0)
7433 } else {
7434 ToolCallAction::run()
7435 }
7436 }
7437 }
7438
7439 #[tokio::test]
7445 async fn valid_tool_call_skip_parity_across_run_and_stream() {
7446 let turns = [
7447 ScriptedTurn::ToolCalls(vec![add_call("tc1", 2, 3)]),
7448 ScriptedTurn::Text("acknowledged"),
7449 ];
7450
7451 let blocking_model =
7452 MockCompletionModel::from_turns(turns.iter().map(ScriptedTurn::as_blocking_turn));
7453 let blocking_hook = RecordingHook::default();
7454 let blocking = AgentBuilder::new(blocking_model)
7455 .tool(MockAddTool)
7456 .build()
7457 .runner("add 2 and 3")
7458 .max_turns(3)
7459 .add_hook(blocking_hook.clone())
7460 .add_hook(SkipToolCallHook("skipped by policy"))
7461 .run()
7462 .await
7463 .expect("blocking run should succeed with a skipped tool call");
7464
7465 let streaming_model = MockCompletionModel::from_stream_turns(
7466 turns
7467 .iter()
7468 .map(|turn| turn.as_stream_events(StreamShape::Complete)),
7469 );
7470 let streaming_hook = RecordingHook::default();
7471 let mut stream = AgentBuilder::new(streaming_model)
7472 .tool(MockAddTool)
7473 .build()
7474 .runner("add 2 and 3")
7475 .max_turns(3)
7476 .add_hook(streaming_hook.clone())
7477 .add_hook(SkipToolCallHook("skipped by policy"))
7478 .stream()
7479 .await;
7480 let mut final_response = None;
7481 while let Some(item) = stream.next().await {
7482 if let Ok(MultiTurnStreamItem::FinalResponse(resp)) =
7483 item.map_err(|err| panic!("stream item errored: {err}"))
7484 {
7485 final_response = Some(resp);
7486 }
7487 }
7488 let final_response = final_response.expect("stream should yield a final response");
7489
7490 assert_eq!(blocking.output, "acknowledged");
7491 assert_eq!(final_response.output(), blocking.output);
7492 assert_eq!(
7493 blocking_hook.shared_events(),
7494 streaming_hook.shared_events()
7495 );
7496 assert_eq!(blocking_hook.tool_results(), streaming_hook.tool_results());
7500 assert_eq!(
7501 blocking_hook.tool_results(),
7502 vec!["skipped by policy".to_string()],
7503 "a skipped tool fires a ToolResult hook with the verbatim skip reason"
7504 );
7505
7506 let blocking_messages = blocking.messages.expect("blocking messages");
7507 let streaming_messages = final_response
7508 .messages()
7509 .expect("streaming history")
7510 .to_vec();
7511 assert_eq!(
7512 serde_json::to_value(&blocking_messages).expect("serialize blocking"),
7513 serde_json::to_value(&streaming_messages).expect("serialize streaming"),
7514 );
7515 assert!(
7518 tool_result_text_in_history(&blocking_messages, "skipped by policy"),
7519 "the verbatim skip reason must be the tool result content in the history"
7520 );
7521 }
7522
7523 struct RewriteToolArgsHook(serde_json::Value);
7527
7528 impl AgentHook for RewriteToolArgsHook {
7529 async fn on_tool_call(&self, _ctx: &HookContext, event: ToolCall<'_>) -> ToolCallAction {
7530 if let ToolCall { .. } = event {
7531 ToolCallAction::rewrite(self.0.clone())
7532 } else {
7533 ToolCallAction::run()
7534 }
7535 }
7536 }
7537
7538 struct EchoStringArgs;
7539
7540 impl Tool for EchoStringArgs {
7541 const NAME: &'static str = "echo_string_args";
7542 type Error = rig::tool::ToolExecutionError;
7543 type Args = String;
7544 type Output = String;
7545
7546 fn description(&self) -> String {
7547 "Echo a JSON string argument".to_string()
7548 }
7549
7550 fn parameters(&self) -> serde_json::Value {
7551 json!({"type": "string"})
7552 }
7553
7554 async fn call(
7555 &self,
7556 _context: &mut ToolContext,
7557 args: Self::Args,
7558 ) -> Result<Self::Output, ToolExecutionError> {
7559 Ok(args)
7560 }
7561 }
7562
7563 #[derive(serde::Deserialize)]
7564 struct FirstGenerationArgs {
7565 old: String,
7566 }
7567
7568 struct FirstGenerationTool(Arc<AtomicU32>);
7569
7570 impl Tool for FirstGenerationTool {
7571 const NAME: &'static str = "generation_pinned";
7572 type Error = rig::tool::ToolExecutionError;
7573 type Args = FirstGenerationArgs;
7574 type Output = String;
7575
7576 fn description(&self) -> String {
7577 "first generation schema".to_string()
7578 }
7579
7580 fn parameters(&self) -> serde_json::Value {
7581 json!({
7582 "type": "object",
7583 "properties": {"old": {"type": "string"}},
7584 "required": ["old"]
7585 })
7586 }
7587
7588 async fn call(
7589 &self,
7590 _context: &mut ToolContext,
7591 args: Self::Args,
7592 ) -> Result<Self::Output, ToolExecutionError> {
7593 self.0.fetch_add(1, SeqCst);
7594 Ok(format!("first:{}", args.old))
7595 }
7596 }
7597
7598 #[derive(serde::Deserialize)]
7599 struct SecondGenerationArgs {
7600 new: String,
7601 }
7602
7603 struct SecondGenerationTool(Arc<AtomicU32>);
7604
7605 impl Tool for SecondGenerationTool {
7606 const NAME: &'static str = FirstGenerationTool::NAME;
7607 type Error = rig::tool::ToolExecutionError;
7608 type Args = SecondGenerationArgs;
7609 type Output = String;
7610
7611 fn description(&self) -> String {
7612 "second generation schema".to_string()
7613 }
7614
7615 fn parameters(&self) -> serde_json::Value {
7616 json!({
7617 "type": "object",
7618 "properties": {"new": {"type": "string"}},
7619 "required": ["new"]
7620 })
7621 }
7622
7623 async fn call(
7624 &self,
7625 _context: &mut ToolContext,
7626 args: Self::Args,
7627 ) -> Result<Self::Output, ToolExecutionError> {
7628 self.0.fetch_add(1, SeqCst);
7629 Ok(format!("second:{}", args.new))
7630 }
7631 }
7632
7633 #[derive(Clone)]
7637 struct PausingCompletionModel {
7638 inner: MockCompletionModel,
7639 request_started: Arc<Notify>,
7640 release_response: Arc<Notify>,
7641 requests: Arc<AtomicU32>,
7642 }
7643
7644 impl PausingCompletionModel {
7645 fn new(inner: MockCompletionModel) -> (Self, Arc<Notify>, Arc<Notify>) {
7646 let request_started = Arc::new(Notify::new());
7647 let release_response = Arc::new(Notify::new());
7648 (
7649 Self {
7650 inner,
7651 request_started: request_started.clone(),
7652 release_response: release_response.clone(),
7653 requests: Arc::new(AtomicU32::new(0)),
7654 },
7655 request_started,
7656 release_response,
7657 )
7658 }
7659
7660 async fn inspect_and_pause(&self, request: &crate::completion::CompletionRequest) {
7661 let request_index = self.requests.fetch_add(1, SeqCst);
7662 let definition = request
7663 .tools
7664 .iter()
7665 .find(|definition| definition.name == FirstGenerationTool::NAME)
7666 .expect("generation tool must be advertised");
7667 if request_index == 0 {
7668 assert_eq!(definition.description, "first generation schema");
7669 self.request_started.notify_one();
7670 self.release_response.notified().await;
7671 } else {
7672 assert_eq!(definition.description, "second generation schema");
7673 }
7674 }
7675 }
7676
7677 impl CompletionModel for PausingCompletionModel {
7678 async fn completion(
7679 &self,
7680 request: crate::completion::CompletionRequest,
7681 ) -> Result<crate::completion::CompletionResponse, crate::completion::CompletionError>
7682 {
7683 self.inspect_and_pause(&request).await;
7684 self.inner.completion(request).await
7685 }
7686
7687 async fn stream(
7688 &self,
7689 request: crate::completion::CompletionRequest,
7690 ) -> Result<crate::streaming::StreamingCompletionResponse, crate::completion::CompletionError>
7691 {
7692 self.inspect_and_pause(&request).await;
7693 self.inner.stream(request).await
7694 }
7695 }
7696
7697 #[test]
7698 fn one_hook_instance_attaches_to_distinct_completion_models() {
7699 #[derive(Clone)]
7700 struct ProviderIndependentHook;
7701
7702 impl AgentHook for ProviderIndependentHook {}
7703
7704 let hook = ProviderIndependentHook;
7705 let _mock_agent = AgentBuilder::new(MockCompletionModel::default())
7706 .add_hook(hook.clone())
7707 .build();
7708 let (other_model, _, _) = PausingCompletionModel::new(MockCompletionModel::default());
7709 let _other_agent = AgentBuilder::new(other_model).add_hook(hook).build();
7710 }
7711
7712 #[tokio::test]
7719 async fn valid_tool_call_rewrite_args_parity_across_run_and_stream() {
7720 let turns = [
7723 ScriptedTurn::ToolCalls(vec![add_call("tc1", 2, 3)]),
7724 ScriptedTurn::Text("acknowledged"),
7725 ];
7726 let replacement = json!({"x": 2, "y": 40});
7727
7728 let blocking_model =
7729 MockCompletionModel::from_turns(turns.iter().map(ScriptedTurn::as_blocking_turn));
7730 let blocking_hook = RecordingHook::default();
7731 let blocking = AgentBuilder::new(blocking_model)
7732 .tool(MockAddTool)
7733 .build()
7734 .runner("add 2 and 3")
7735 .max_turns(3)
7736 .add_hook(blocking_hook.clone())
7737 .add_hook(RewriteToolArgsHook(replacement.clone()))
7738 .run()
7739 .await
7740 .expect("blocking run should succeed with rewritten tool arguments");
7741
7742 let streaming_model = MockCompletionModel::from_stream_turns(
7743 turns
7744 .iter()
7745 .map(|turn| turn.as_stream_events(StreamShape::Complete)),
7746 );
7747 let streaming_hook = RecordingHook::default();
7748 let mut stream = AgentBuilder::new(streaming_model)
7749 .tool(MockAddTool)
7750 .build()
7751 .runner("add 2 and 3")
7752 .max_turns(3)
7753 .add_hook(streaming_hook.clone())
7754 .add_hook(RewriteToolArgsHook(replacement))
7755 .stream()
7756 .await;
7757 let mut final_response = None;
7758 while let Some(item) = stream.next().await {
7759 if let Ok(MultiTurnStreamItem::FinalResponse(resp)) =
7760 item.map_err(|err| panic!("stream item errored: {err}"))
7761 {
7762 final_response = Some(resp);
7763 }
7764 }
7765 let final_response = final_response.expect("stream should yield a final response");
7766
7767 assert_eq!(blocking_hook.tool_results(), vec!["42".to_string()]);
7770 assert_eq!(blocking.output, "acknowledged");
7771 assert_eq!(final_response.output(), blocking.output);
7772 assert_eq!(
7773 blocking_hook.shared_events(),
7774 streaming_hook.shared_events()
7775 );
7776 assert_eq!(blocking_hook.tool_results(), streaming_hook.tool_results());
7777 }
7778
7779 #[tokio::test]
7780 async fn string_tool_call_without_rewrite_is_canonical_across_run_and_stream() {
7781 let turns = [
7782 ScriptedTurn::ToolCalls(vec![ScriptedToolCall {
7783 id: "tc-string",
7784 name: EchoStringArgs::NAME,
7785 args: json!("original"),
7786 }]),
7787 ScriptedTurn::Text("done"),
7788 ];
7789
7790 let blocking_hook = RecordingHook::default();
7791 let blocking = AgentBuilder::new(MockCompletionModel::from_turns(
7792 turns.iter().map(ScriptedTurn::as_blocking_turn),
7793 ))
7794 .tool(EchoStringArgs)
7795 .build()
7796 .runner("echo a string")
7797 .max_turns(3)
7798 .add_hook(blocking_hook.clone())
7799 .run()
7800 .await
7801 .expect("blocking string call should execute");
7802
7803 let streaming_hook = RecordingHook::default();
7804 let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns(
7805 turns
7806 .iter()
7807 .map(|turn| turn.as_stream_events(StreamShape::Complete)),
7808 ))
7809 .tool(EchoStringArgs)
7810 .build()
7811 .runner("echo a string")
7812 .max_turns(3)
7813 .add_hook(streaming_hook.clone())
7814 .stream()
7815 .await;
7816 let mut final_output = None;
7817 while let Some(item) = stream.next().await {
7818 if let MultiTurnStreamItem::FinalResponse(response) =
7819 item.expect("streaming string call should execute")
7820 {
7821 final_output = Some(response.output().to_string());
7822 }
7823 }
7824
7825 assert_eq!(blocking.output, "done");
7826 assert_eq!(final_output.as_deref(), Some("done"));
7827 assert_eq!(blocking_hook.tool_results(), vec!["original"]);
7828 assert_eq!(streaming_hook.tool_results(), vec!["original"]);
7829 }
7830
7831 #[tokio::test]
7832 async fn string_tool_call_rewrite_is_canonical_json_across_run_and_stream() {
7833 let turns = [
7834 ScriptedTurn::ToolCalls(vec![ScriptedToolCall {
7835 id: "tc-string",
7836 name: EchoStringArgs::NAME,
7837 args: json!("original"),
7838 }]),
7839 ScriptedTurn::Text("done"),
7840 ];
7841 let replacement = json!("sanitized");
7842
7843 let blocking_hook = RecordingHook::default();
7844 let blocking = AgentBuilder::new(MockCompletionModel::from_turns(
7845 turns.iter().map(ScriptedTurn::as_blocking_turn),
7846 ))
7847 .tool(EchoStringArgs)
7848 .build()
7849 .runner("echo a string")
7850 .max_turns(3)
7851 .add_hook(blocking_hook.clone())
7852 .add_hook(RewriteToolArgsHook(replacement.clone()))
7853 .run()
7854 .await
7855 .expect("blocking string rewrite should execute");
7856
7857 let streaming_hook = RecordingHook::default();
7858 let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns(
7859 turns
7860 .iter()
7861 .map(|turn| turn.as_stream_events(StreamShape::Complete)),
7862 ))
7863 .tool(EchoStringArgs)
7864 .build()
7865 .runner("echo a string")
7866 .max_turns(3)
7867 .add_hook(streaming_hook.clone())
7868 .add_hook(RewriteToolArgsHook(replacement))
7869 .stream()
7870 .await;
7871 let mut final_output = None;
7872 while let Some(item) = stream.next().await {
7873 if let MultiTurnStreamItem::FinalResponse(response) =
7874 item.expect("streaming string rewrite should execute")
7875 {
7876 final_output = Some(response.output().to_string());
7877 }
7878 }
7879
7880 assert_eq!(blocking.output, "done");
7881 assert_eq!(final_output.as_deref(), Some("done"));
7882 assert_eq!(blocking_hook.tool_results(), vec!["sanitized"]);
7883 assert_eq!(streaming_hook.tool_results(), vec!["sanitized"]);
7884 }
7885
7886 #[tokio::test]
7887 async fn blocking_turn_dispatches_the_registry_generation_it_advertised() {
7888 let first_calls = Arc::new(AtomicU32::new(0));
7889 let second_calls = Arc::new(AtomicU32::new(0));
7890 let handle: ToolServerHandle = ToolServer::new()
7891 .tool(FirstGenerationTool(first_calls.clone()))
7892 .run();
7893 let inner = MockCompletionModel::from_turns([
7894 MockTurn::tool_call(
7895 "tc-generation",
7896 FirstGenerationTool::NAME,
7897 json!({"old": "payload"}),
7898 ),
7899 MockTurn::text("done"),
7900 ]);
7901 let (model, request_started, release_response) = PausingCompletionModel::new(inner);
7902 let runner = AgentBuilder::new(model)
7903 .tool_server_handle(handle.clone())
7904 .build()
7905 .runner("use the generation tool")
7906 .max_turns(3);
7907
7908 let run = runner.run();
7909 let replace = async {
7910 request_started.notified().await;
7911 handle
7912 .add_tool(SecondGenerationTool(second_calls.clone()))
7913 .await;
7914 release_response.notify_one();
7915 };
7916 let (response, ()) = tokio::time::timeout(std::time::Duration::from_secs(2), async {
7917 tokio::join!(run, replace)
7918 })
7919 .await
7920 .expect("in-flight blocking replacement must not hang");
7921 let response = response.expect("blocking run should use its pinned tool generation");
7922
7923 assert_eq!(response.output, "done");
7924 assert_eq!(first_calls.load(SeqCst), 1);
7925 assert_eq!(second_calls.load(SeqCst), 0);
7926 }
7927
7928 #[tokio::test]
7929 async fn streaming_turn_dispatches_the_registry_generation_it_advertised() {
7930 let first_calls = Arc::new(AtomicU32::new(0));
7931 let second_calls = Arc::new(AtomicU32::new(0));
7932 let handle: ToolServerHandle = ToolServer::new()
7933 .tool(FirstGenerationTool(first_calls.clone()))
7934 .run();
7935 let turns = [
7936 ScriptedTurn::ToolCalls(vec![ScriptedToolCall {
7937 id: "tc-generation",
7938 name: FirstGenerationTool::NAME,
7939 args: json!({"old": "payload"}),
7940 }]),
7941 ScriptedTurn::Text("done"),
7942 ];
7943 let inner = MockCompletionModel::from_stream_turns(
7944 turns
7945 .iter()
7946 .map(|turn| turn.as_stream_events(StreamShape::Complete)),
7947 );
7948 let (model, request_started, release_response) = PausingCompletionModel::new(inner);
7949 let runner = AgentBuilder::new(model)
7950 .tool_server_handle(handle.clone())
7951 .build()
7952 .runner("use the generation tool")
7953 .max_turns(3);
7954
7955 let drive = async {
7956 let mut stream = runner.stream().await;
7957 let mut final_output = None;
7958 while let Some(item) = stream.next().await {
7959 if let MultiTurnStreamItem::FinalResponse(response) =
7960 item.expect("streaming run should use its pinned tool generation")
7961 {
7962 final_output = Some(response.output().to_string());
7963 }
7964 }
7965 final_output
7966 };
7967 let replace = async {
7968 request_started.notified().await;
7969 handle
7970 .add_tool(SecondGenerationTool(second_calls.clone()))
7971 .await;
7972 release_response.notify_one();
7973 };
7974 let (final_output, ()) = tokio::time::timeout(std::time::Duration::from_secs(2), async {
7975 tokio::join!(drive, replace)
7976 })
7977 .await
7978 .expect("in-flight streaming replacement must not hang");
7979
7980 assert_eq!(final_output.as_deref(), Some("done"));
7981 assert_eq!(first_calls.load(SeqCst), 1);
7982 assert_eq!(second_calls.load(SeqCst), 0);
7983 }
7984
7985 struct RewriteToolResultHook(&'static str);
7989
7990 impl AgentHook for RewriteToolResultHook {
7991 async fn on_tool_result(
7992 &self,
7993 _ctx: &HookContext,
7994 event: ToolResultEvent<'_>,
7995 ) -> ToolResultAction {
7996 if let ToolResultEvent { .. } = event {
7997 ToolResultAction::rewrite(self.0)
7998 } else {
7999 ToolResultAction::keep()
8000 }
8001 }
8002 }
8003
8004 #[tokio::test]
8011 async fn valid_tool_result_rewrite_parity_across_run_and_stream() {
8012 let turns = [
8015 ScriptedTurn::ToolCalls(vec![add_call("tc1", 2, 3)]),
8016 ScriptedTurn::Text("acknowledged"),
8017 ];
8018
8019 let blocking_model =
8020 MockCompletionModel::from_turns(turns.iter().map(ScriptedTurn::as_blocking_turn));
8021 let blocking_hook = RecordingHook::default();
8022 let blocking = AgentBuilder::new(blocking_model)
8023 .tool(MockAddTool)
8024 .build()
8025 .runner("add 2 and 3")
8026 .max_turns(3)
8027 .add_hook(blocking_hook.clone())
8028 .add_hook(RewriteToolResultHook("redacted-result"))
8029 .run()
8030 .await
8031 .expect("blocking run should succeed with a rewritten tool result");
8032
8033 let streaming_model = MockCompletionModel::from_stream_turns(
8034 turns
8035 .iter()
8036 .map(|turn| turn.as_stream_events(StreamShape::Complete)),
8037 );
8038 let streaming_hook = RecordingHook::default();
8039 let mut stream = AgentBuilder::new(streaming_model)
8040 .tool(MockAddTool)
8041 .build()
8042 .runner("add 2 and 3")
8043 .max_turns(3)
8044 .add_hook(streaming_hook.clone())
8045 .add_hook(RewriteToolResultHook("redacted-result"))
8046 .stream()
8047 .await;
8048 let mut final_response = None;
8049 while let Some(item) = stream.next().await {
8050 if let Ok(MultiTurnStreamItem::FinalResponse(resp)) =
8051 item.map_err(|err| panic!("stream item errored: {err}"))
8052 {
8053 final_response = Some(resp);
8054 }
8055 }
8056 let final_response = final_response.expect("stream should yield a final response");
8057
8058 assert_eq!(blocking.output, "acknowledged");
8059 assert_eq!(final_response.output(), blocking.output);
8060
8061 assert_eq!(blocking_hook.tool_results(), vec!["5".to_string()]);
8064 assert_eq!(blocking_hook.tool_results(), streaming_hook.tool_results());
8065
8066 let blocking_messages = blocking.messages.expect("blocking messages");
8069 let streaming_messages = final_response
8070 .messages()
8071 .expect("streaming history")
8072 .to_vec();
8073 assert_eq!(
8074 serde_json::to_value(&blocking_messages).expect("serialize blocking"),
8075 serde_json::to_value(&streaming_messages).expect("serialize streaming"),
8076 );
8077 assert!(
8078 tool_result_text_in_history(&blocking_messages, "redacted-result"),
8079 "the model-visible tool result must be the hook's replacement"
8080 );
8081 assert!(
8082 !tool_result_text_in_history(&blocking_messages, "5"),
8083 "the tool's original output must not reach the model after a rewrite"
8084 );
8085 }
8086
8087 #[tokio::test]
8093 async fn rewrite_result_is_delivered_verbatim_not_reparsed() {
8094 const IMAGE_JSON: &str = r#"{"type":"image","data":"abc","mimeType":"image/png"}"#;
8095
8096 let turns = [
8097 ScriptedTurn::ToolCalls(vec![add_call("tc1", 2, 3)]),
8098 ScriptedTurn::Text("done"),
8099 ];
8100 let model =
8101 MockCompletionModel::from_turns(turns.iter().map(ScriptedTurn::as_blocking_turn));
8102 let result = AgentBuilder::new(model)
8103 .tool(MockAddTool)
8104 .build()
8105 .runner("add 2 and 3")
8106 .max_turns(3)
8107 .add_hook(RewriteToolResultHook(IMAGE_JSON))
8108 .run()
8109 .await
8110 .expect("run should succeed with a JSON-shaped rewritten result");
8111
8112 let messages = result.messages.expect("messages");
8113 assert!(
8114 tool_result_text_in_history(&messages, IMAGE_JSON),
8115 "the JSON-shaped replacement must reach history verbatim as text, not be \
8116 re-parsed into a structured/image content block"
8117 );
8118 }
8119
8120 struct PatchRequestHook;
8124
8125 impl AgentHook for PatchRequestHook {
8126 async fn on_completion_call(
8127 &self,
8128 _ctx: &HookContext,
8129 event: CompletionCallEvent<'_>,
8130 ) -> CompletionCallAction {
8131 if let CompletionCallEvent { .. } = event {
8132 CompletionCallAction::patch(
8133 RequestPatch::new()
8134 .preamble(OVERRIDE_PREAMBLE)
8135 .temperature(0.25)
8136 .max_tokens(OVERRIDE_MAX_TOKENS)
8137 .tool_choice(ToolChoice::Required)
8138 .active_tools(["add"])
8139 .additional_params(json!({"injected": true})),
8140 )
8141 } else {
8142 CompletionCallAction::continue_run()
8143 }
8144 }
8145 }
8146
8147 const OVERRIDE_PREAMBLE: &str = "overridden: critical-step instructions";
8148 const OVERRIDE_MAX_TOKENS: u64 = 512;
8149
8150 #[tokio::test]
8155 async fn patch_request_parity_across_run_and_stream() {
8156 fn assert_request(req: &crate::completion::CompletionRequest) {
8157 assert_eq!(
8158 req.temperature,
8159 Some(0.25),
8160 "override temperature wins over the agent's 0.9"
8161 );
8162 assert_eq!(
8163 req.max_tokens,
8164 Some(OVERRIDE_MAX_TOKENS),
8165 "override max_tokens wins over the agent's 64"
8166 );
8167 let system = req.chat_history.iter().find_map(|m| match m {
8169 Message::System { content } => Some(content.as_str()),
8170 _ => None,
8171 });
8172 assert_eq!(
8173 system,
8174 Some(OVERRIDE_PREAMBLE),
8175 "override preamble wins over the agent's baseline and is the leading system message"
8176 );
8177 assert!(matches!(req.tool_choice, Some(ToolChoice::Required)));
8178 let tool_names: Vec<&str> = req.tools.iter().map(|t| t.name.as_str()).collect();
8179 assert_eq!(
8180 tool_names,
8181 ["add"],
8182 "active_tools narrows the advertised set to `add` (drops `subtract`)"
8183 );
8184 let params = req.additional_params.as_ref().expect("additional_params");
8187 assert_eq!(params.get("runner").and_then(|v| v.as_str()), Some("keep"));
8188 assert_eq!(params.get("injected").and_then(|v| v.as_bool()), Some(true));
8189 assert!(params.get("baseline").is_none());
8190 }
8191
8192 let blocking_model = MockCompletionModel::from_turns([MockTurn::text("done")]);
8193 let blocking_probe = blocking_model.clone();
8194 let blocking = AgentBuilder::new(blocking_model)
8195 .tool(MockAddTool)
8196 .tool(MockSubtractTool)
8197 .preamble("baseline preamble")
8198 .temperature(0.9)
8199 .max_tokens(64)
8200 .additional_params(json!({"baseline": "keep"}))
8201 .add_hook(PatchRequestHook)
8202 .build()
8203 .runner("go")
8204 .replace_additional_params(json!({"runner": "keep", "injected": false}))
8205 .max_turns(2)
8206 .run()
8207 .await
8208 .expect("blocking run should succeed");
8209 assert_eq!(blocking.output, "done");
8210 let blocking_requests = blocking_probe.requests();
8211 assert_eq!(blocking_requests.len(), 1);
8212 assert_request(&blocking_requests[0]);
8213
8214 let streaming_model = MockCompletionModel::from_stream_turns([
8215 ScriptedTurn::Text("done").as_stream_events(StreamShape::Complete)
8216 ]);
8217 let streaming_probe = streaming_model.clone();
8218 let mut stream = AgentBuilder::new(streaming_model)
8219 .tool(MockAddTool)
8220 .tool(MockSubtractTool)
8221 .preamble("baseline preamble")
8222 .temperature(0.9)
8223 .max_tokens(64)
8224 .additional_params(json!({"baseline": "keep"}))
8225 .add_hook(PatchRequestHook)
8226 .build()
8227 .runner("go")
8228 .replace_additional_params(json!({"runner": "keep", "injected": false}))
8229 .max_turns(2)
8230 .stream()
8231 .await;
8232 while let Some(item) = stream.next().await {
8233 let _ = item.map_err(|err| panic!("stream item errored: {err}"));
8234 }
8235 let streaming_requests = streaming_probe.requests();
8236 assert_eq!(streaming_requests.len(), 1);
8237 assert_request(&streaming_requests[0]);
8238 }
8239
8240 fn hook_doc(id: &str, text: &str) -> crate::completion::Document {
8243 crate::completion::Document {
8244 id: id.to_string(),
8245 text: text.to_string(),
8246 additional_props: Default::default(),
8247 }
8248 }
8249
8250 struct ExtraContextHook {
8252 id: &'static str,
8253 text: &'static str,
8254 }
8255
8256 impl AgentHook for ExtraContextHook {
8257 async fn on_completion_call(
8258 &self,
8259 _ctx: &HookContext,
8260 event: CompletionCallEvent<'_>,
8261 ) -> CompletionCallAction {
8262 if let CompletionCallEvent { .. } = event {
8263 CompletionCallAction::patch(
8264 RequestPatch::new().context(hook_doc(self.id, self.text)),
8265 )
8266 } else {
8267 CompletionCallAction::continue_run()
8268 }
8269 }
8270 }
8271
8272 struct ExtraContextTurnOneHook;
8275
8276 impl AgentHook for ExtraContextTurnOneHook {
8277 async fn on_completion_call(
8278 &self,
8279 _ctx: &HookContext,
8280 event: CompletionCallEvent<'_>,
8281 ) -> CompletionCallAction {
8282 if let CompletionCallEvent { turn, .. } = event
8283 && turn == 1
8284 {
8285 return CompletionCallAction::patch(
8286 RequestPatch::new().context(hook_doc("turn-one", "only turn 1")),
8287 );
8288 }
8289 CompletionCallAction::continue_run()
8290 }
8291 }
8292
8293 #[derive(Clone)]
8294 struct RecordingContextIndex {
8295 id: &'static str,
8296 queries: Arc<Mutex<Vec<(String, u64)>>>,
8297 }
8298
8299 impl VectorStoreIndex for RecordingContextIndex {
8300 type Filter = Filter<serde_json::Value>;
8301
8302 async fn top_n<T: for<'a> Deserialize<'a> + WasmCompatSend>(
8303 &self,
8304 req: VectorSearchRequest,
8305 ) -> Result<Vec<(f64, String, T)>, VectorStoreError> {
8306 self.queries
8307 .lock()
8308 .expect("context query recorder lock")
8309 .push((req.query().to_string(), req.samples()));
8310 let value = serde_json::from_value(json!({ "source": self.id }))?;
8311 Ok(vec![(1.0, self.id.to_string(), value)])
8312 }
8313
8314 async fn top_n_ids(
8315 &self,
8316 _req: VectorSearchRequest,
8317 ) -> Result<Vec<(f64, String)>, VectorStoreError> {
8318 Ok(vec![(1.0, self.id.to_string())])
8319 }
8320 }
8321
8322 struct FailingContextIndex;
8323
8324 impl VectorStoreIndex for FailingContextIndex {
8325 type Filter = Filter<serde_json::Value>;
8326
8327 async fn top_n<T: for<'a> Deserialize<'a> + WasmCompatSend>(
8328 &self,
8329 _req: VectorSearchRequest,
8330 ) -> Result<Vec<(f64, String, T)>, VectorStoreError> {
8331 Err(VectorStoreError::BuilderError(
8332 "context index unavailable".to_string(),
8333 ))
8334 }
8335
8336 async fn top_n_ids(
8337 &self,
8338 _req: VectorSearchRequest,
8339 ) -> Result<Vec<(f64, String)>, VectorStoreError> {
8340 Err(VectorStoreError::BuilderError(
8341 "context index unavailable".to_string(),
8342 ))
8343 }
8344 }
8345
8346 struct QueryRecordingToolIndex {
8347 queries: Arc<Mutex<Vec<String>>>,
8348 }
8349
8350 impl VectorStoreIndex for QueryRecordingToolIndex {
8351 type Filter = Filter<serde_json::Value>;
8352
8353 async fn top_n<T: for<'a> Deserialize<'a> + WasmCompatSend>(
8354 &self,
8355 _req: VectorSearchRequest,
8356 ) -> Result<Vec<(f64, String, T)>, VectorStoreError> {
8357 Ok(Vec::new())
8358 }
8359
8360 async fn top_n_ids(
8361 &self,
8362 req: VectorSearchRequest,
8363 ) -> Result<Vec<(f64, String)>, VectorStoreError> {
8364 self.queries
8365 .lock()
8366 .expect("query recorder lock")
8367 .push(req.query().to_string());
8368 Ok(vec![(1.0, MockAddTool::NAME.to_string())])
8369 }
8370 }
8371
8372 fn one_text_stream_turn(text: &'static str) -> Vec<MockStreamEvent> {
8373 vec![
8374 MockStreamEvent::text(text),
8375 MockStreamEvent::final_response_with_total_tokens(0),
8376 ]
8377 }
8378
8379 #[tokio::test]
8382 async fn extra_context_appears_after_static_context_on_both_surfaces() {
8383 fn assert_docs(req: &crate::completion::CompletionRequest) {
8384 let ids: Vec<&str> = req.documents.iter().map(|d| d.id.as_str()).collect();
8385 let static_pos = ids
8386 .iter()
8387 .position(|id| id.starts_with("static_doc"))
8388 .expect("static context document present");
8389 let extra_pos = ids
8390 .iter()
8391 .position(|id| *id == "hook-doc")
8392 .expect("hook extra_context document present");
8393 assert!(
8394 static_pos < extra_pos,
8395 "static context precedes hook extras: {ids:?}"
8396 );
8397 assert!(
8398 req.documents.iter().any(|d| d.text == "injected"),
8399 "the hook document's text is present"
8400 );
8401 }
8402
8403 let blocking_model = MockCompletionModel::from_turns([MockTurn::text("done")]);
8404 let blocking_probe = blocking_model.clone();
8405 AgentBuilder::new(blocking_model)
8406 .context("static context text")
8407 .add_hook(ExtraContextHook {
8408 id: "hook-doc",
8409 text: "injected",
8410 })
8411 .build()
8412 .runner("go")
8413 .run()
8414 .await
8415 .expect("blocking run should succeed");
8416 assert_docs(blocking_probe.requests().first().expect("one request"));
8417
8418 let streaming_model =
8419 MockCompletionModel::from_stream_turns([one_text_stream_turn("done")]);
8420 let streaming_probe = streaming_model.clone();
8421 let mut stream = AgentBuilder::new(streaming_model)
8422 .context("static context text")
8423 .add_hook(ExtraContextHook {
8424 id: "hook-doc",
8425 text: "injected",
8426 })
8427 .build()
8428 .runner("go")
8429 .stream()
8430 .await;
8431 while let Some(item) = stream.next().await {
8432 let _ = item.map_err(|err| panic!("stream item errored: {err}"));
8433 }
8434 assert_docs(streaming_probe.requests().first().expect("one request"));
8435 }
8436
8437 #[tokio::test]
8439 async fn multiple_hooks_extra_context_append_in_registration_order() {
8440 let model = MockCompletionModel::from_turns([MockTurn::text("done")]);
8441 let probe = model.clone();
8442 AgentBuilder::new(model)
8443 .add_hook(ExtraContextHook {
8444 id: "first",
8445 text: "1",
8446 })
8447 .add_hook(ExtraContextHook {
8448 id: "second",
8449 text: "2",
8450 })
8451 .build()
8452 .runner("go")
8453 .run()
8454 .await
8455 .expect("run should succeed");
8456 let requests = probe.requests();
8457 let req = requests.first().expect("one request");
8458 let ids: Vec<&str> = req.documents.iter().map(|d| d.id.as_str()).collect();
8459 assert_eq!(
8460 ids,
8461 vec!["first", "second"],
8462 "hook extras append in registration order"
8463 );
8464 }
8465
8466 #[tokio::test]
8467 async fn dynamic_context_preserves_query_selection_formatting_and_order_on_both_surfaces() {
8468 fn assert_documents(request: &crate::completion::CompletionRequest) {
8469 let documents = request
8470 .documents
8471 .iter()
8472 .map(|document| (document.id.as_str(), document.text.as_str()))
8473 .collect::<Vec<_>>();
8474 assert_eq!(
8475 documents,
8476 vec![
8477 ("static_doc_0", "static context"),
8478 ("blocking", "{\n \"source\": \"blocking\"\n}"),
8479 ]
8480 );
8481 }
8482
8483 let blocking_queries = Arc::new(Mutex::new(Vec::new()));
8484 let blocking_model = MockCompletionModel::from_turns([MockTurn::text("done")]);
8485 let blocking_probe = blocking_model.clone();
8486 AgentBuilder::new(blocking_model)
8487 .context("static context")
8488 .dynamic_context(
8489 2,
8490 RecordingContextIndex {
8491 id: "blocking",
8492 queries: blocking_queries.clone(),
8493 },
8494 )
8495 .build()
8496 .runner("current blocking query")
8497 .history(vec![Message::user("ignored history query")])
8498 .run()
8499 .await
8500 .expect("blocking dynamic-context run should succeed");
8501 assert_eq!(
8502 *blocking_queries.lock().expect("blocking queries"),
8503 vec![("current blocking query".to_string(), 2)]
8504 );
8505 assert_documents(blocking_probe.requests().first().expect("one request"));
8506
8507 let streaming_queries = Arc::new(Mutex::new(Vec::new()));
8508 let streaming_model =
8509 MockCompletionModel::from_stream_turns([one_text_stream_turn("done")]);
8510 let streaming_probe = streaming_model.clone();
8511 let mut stream = AgentBuilder::new(streaming_model)
8512 .dynamic_context(
8513 3,
8514 RecordingContextIndex {
8515 id: "streaming",
8516 queries: streaming_queries.clone(),
8517 },
8518 )
8519 .build()
8520 .runner(Message::User {
8521 content: vec![UserContent::image_url(
8522 "https://example.com/prompt.png",
8523 None,
8524 None,
8525 )],
8526 })
8527 .history(vec![
8528 Message::user("older history query"),
8529 Message::user("latest history query"),
8530 ])
8531 .stream()
8532 .await;
8533 while let Some(item) = stream.next().await {
8534 item.expect("streaming dynamic-context run should succeed");
8535 }
8536 assert_eq!(
8537 *streaming_queries.lock().expect("streaming queries"),
8538 vec![("latest history query".to_string(), 3)]
8539 );
8540 let streaming_requests = streaming_probe.requests();
8541 let request = streaming_requests.first().expect("one request");
8542 assert_eq!(request.documents.len(), 1);
8543 assert_eq!(request.documents[0].id, "streaming");
8544 assert_eq!(
8545 request.documents[0].text,
8546 "{\n \"source\": \"streaming\"\n}"
8547 );
8548 }
8549
8550 #[tokio::test]
8551 async fn dynamic_context_and_application_hooks_follow_registration_order() {
8552 let queries = Arc::new(Mutex::new(Vec::new()));
8553 let model = MockCompletionModel::from_turns([MockTurn::text("done")]);
8554 let probe = model.clone();
8555 AgentBuilder::new(model)
8556 .context("static")
8557 .add_hook(ExtraContextHook {
8558 id: "before",
8559 text: "before dynamic context",
8560 })
8561 .dynamic_context(
8562 1,
8563 RecordingContextIndex {
8564 id: "first",
8565 queries: queries.clone(),
8566 },
8567 )
8568 .add_hook(ExtraContextHook {
8569 id: "between",
8570 text: "between dynamic contexts",
8571 })
8572 .dynamic_context(
8573 2,
8574 RecordingContextIndex {
8575 id: "second",
8576 queries: queries.clone(),
8577 },
8578 )
8579 .add_hook(ExtraContextHook {
8580 id: "after",
8581 text: "after dynamic context",
8582 })
8583 .build()
8584 .runner("query")
8585 .run()
8586 .await
8587 .expect("run should succeed");
8588
8589 assert_eq!(
8590 probe.requests()[0]
8591 .documents
8592 .iter()
8593 .map(|document| document.id.as_str())
8594 .collect::<Vec<_>>(),
8595 vec![
8596 "static_doc_0",
8597 "before",
8598 "first",
8599 "between",
8600 "second",
8601 "after",
8602 ]
8603 );
8604 assert_eq!(
8605 *queries.lock().expect("context queries"),
8606 vec![("query".to_string(), 1), ("query".to_string(), 2)]
8607 );
8608
8609 let skipped_queries = Arc::new(Mutex::new(Vec::new()));
8610 let error = AgentBuilder::new(MockCompletionModel::from_turns([MockTurn::text("unused")]))
8611 .add_hook(TerminateOn(StepEventKind::CompletionCall))
8612 .dynamic_context(
8613 1,
8614 RecordingContextIndex {
8615 id: "skipped",
8616 queries: skipped_queries.clone(),
8617 },
8618 )
8619 .build()
8620 .runner("query")
8621 .run()
8622 .await
8623 .expect_err("an earlier stop hook should terminate before retrieval");
8624 assert!(matches!(error, PromptError::PromptCancelled { .. }));
8625 assert!(skipped_queries.lock().expect("skipped queries").is_empty());
8626 }
8627
8628 #[tokio::test]
8629 async fn dynamic_context_retrieval_failure_stops_before_provider_io_on_both_surfaces() {
8630 let blocking_model = MockCompletionModel::from_turns([MockTurn::text("unused")]);
8631 let blocking_probe = blocking_model.clone();
8632 let error = AgentBuilder::new(blocking_model)
8633 .dynamic_context(1, FailingContextIndex)
8634 .build()
8635 .runner("retrieve this")
8636 .run()
8637 .await
8638 .expect_err("failed retrieval should stop the run");
8639 assert!(matches!(
8640 error,
8641 PromptError::PromptCancelled { reason, .. }
8642 if reason.contains("context index unavailable")
8643 ));
8644 assert_eq!(blocking_probe.request_count(), 0);
8645
8646 let streaming_model =
8647 MockCompletionModel::from_stream_turns([one_text_stream_turn("unused")]);
8648 let streaming_probe = streaming_model.clone();
8649 let mut stream = AgentBuilder::new(streaming_model)
8650 .dynamic_context(1, FailingContextIndex)
8651 .build()
8652 .runner("retrieve this")
8653 .stream()
8654 .await;
8655 let error = stream
8656 .next()
8657 .await
8658 .expect("stream should report retrieval failure")
8659 .expect_err("failed retrieval should stop the stream");
8660 assert!(matches!(
8661 error,
8662 StreamingError::Prompt(prompt_error)
8663 if matches!(
8664 prompt_error.as_ref(),
8665 PromptError::PromptCancelled { reason, .. }
8666 if reason.contains("context index unavailable")
8667 )
8668 ));
8669 assert_eq!(streaming_probe.request_count(), 0);
8670 }
8671
8672 #[tokio::test]
8673 async fn retrieved_tool_query_selection_is_unchanged_on_both_surfaces() {
8674 let queries = Arc::new(Mutex::new(Vec::new()));
8675 AgentBuilder::new(MockCompletionModel::from_turns([MockTurn::text("done")]))
8676 .retrieved_tools(
8677 1,
8678 QueryRecordingToolIndex {
8679 queries: queries.clone(),
8680 },
8681 ToolSet::from_tools(vec![MockAddTool]),
8682 )
8683 .build()
8684 .runner("blocking retrieval query")
8685 .history(vec![Message::user("blocking history query")])
8686 .run()
8687 .await
8688 .expect("blocking run should succeed");
8689
8690 AgentBuilder::new(MockCompletionModel::from_turns([MockTurn::text("done")]))
8691 .retrieved_tools(
8692 1,
8693 QueryRecordingToolIndex {
8694 queries: queries.clone(),
8695 },
8696 ToolSet::from_tools(vec![MockAddTool]),
8697 )
8698 .build()
8699 .runner(Message::User {
8700 content: vec![UserContent::image_url(
8701 "https://example.com/blocking.png",
8702 None,
8703 None,
8704 )],
8705 })
8706 .history(vec![
8707 Message::user("older blocking history query"),
8708 Message::user("latest blocking history query"),
8709 ])
8710 .run()
8711 .await
8712 .expect("blocking history fallback should succeed");
8713
8714 let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([
8715 one_text_stream_turn("done"),
8716 ]))
8717 .retrieved_tools(
8718 1,
8719 QueryRecordingToolIndex {
8720 queries: queries.clone(),
8721 },
8722 ToolSet::from_tools(vec![MockAddTool]),
8723 )
8724 .build()
8725 .runner("streaming retrieval query")
8726 .history(vec![Message::user("streaming history query")])
8727 .stream()
8728 .await;
8729 while let Some(item) = stream.next().await {
8730 item.expect("stream item should succeed");
8731 }
8732
8733 let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([
8734 one_text_stream_turn("done"),
8735 ]))
8736 .retrieved_tools(
8737 1,
8738 QueryRecordingToolIndex {
8739 queries: queries.clone(),
8740 },
8741 ToolSet::from_tools(vec![MockAddTool]),
8742 )
8743 .build()
8744 .runner(Message::User {
8745 content: vec![UserContent::image_url(
8746 "https://example.com/streaming.png",
8747 None,
8748 None,
8749 )],
8750 })
8751 .history(vec![
8752 Message::user("older streaming history query"),
8753 Message::user("latest streaming history query"),
8754 ])
8755 .stream()
8756 .await;
8757 while let Some(item) = stream.next().await {
8758 item.expect("stream item should succeed");
8759 }
8760
8761 assert_eq!(
8762 *queries.lock().expect("query recorder lock"),
8763 vec![
8764 "blocking retrieval query",
8765 "latest blocking history query",
8766 "streaming retrieval query",
8767 "latest streaming history query",
8768 ]
8769 );
8770 }
8771
8772 #[tokio::test]
8775 async fn extra_context_is_per_turn_non_sticky() {
8776 fn assert_turns(requests: &[crate::completion::CompletionRequest]) {
8777 assert_eq!(requests.len(), 2, "two model turns");
8778 let turn1 = requests.first().expect("turn 1");
8779 let turn2 = requests.get(1).expect("turn 2");
8780 assert!(
8781 turn1.documents.iter().any(|d| d.id == "turn-one"),
8782 "turn 1 carries the injected document"
8783 );
8784 assert!(
8785 turn2.documents.iter().all(|d| d.id != "turn-one"),
8786 "turn 2 does not inherit turn 1's per-turn document"
8787 );
8788 }
8789
8790 let blocking_probe = blocking_model();
8791 let probe = blocking_probe.clone();
8792 AgentBuilder::new(blocking_probe)
8793 .tool(MockAddTool)
8794 .add_hook(ExtraContextTurnOneHook)
8795 .build()
8796 .runner("add 2 and 3")
8797 .max_turns(3)
8798 .run()
8799 .await
8800 .expect("blocking run should succeed");
8801 assert_turns(&probe.requests());
8802
8803 let streaming = streaming_model();
8804 let stream_probe = streaming.clone();
8805 let mut stream = AgentBuilder::new(streaming)
8806 .tool(MockAddTool)
8807 .add_hook(ExtraContextTurnOneHook)
8808 .build()
8809 .runner("add 2 and 3")
8810 .max_turns(3)
8811 .stream()
8812 .await;
8813 while let Some(item) = stream.next().await {
8814 let _ = item.map_err(|err| panic!("stream item errored: {err}"));
8815 }
8816 assert_turns(&stream_probe.requests());
8817 }
8818
8819 #[tokio::test]
8822 async fn history_patch_changes_sent_messages_not_transcript_on_both_surfaces() {
8823 const SENTINEL: &str = "COMPACTED-HISTORY-SENTINEL";
8824
8825 struct HistoryOverrideHook;
8826 impl AgentHook for HistoryOverrideHook {
8827 async fn on_completion_call(
8828 &self,
8829 _ctx: &HookContext,
8830 event: CompletionCallEvent<'_>,
8831 ) -> CompletionCallAction {
8832 if let CompletionCallEvent { .. } = event {
8833 CompletionCallAction::patch(
8834 RequestPatch::new().history([Message::user(SENTINEL)]),
8835 )
8836 } else {
8837 CompletionCallAction::continue_run()
8838 }
8839 }
8840 }
8841
8842 fn request_has_sentinel(req: &crate::completion::CompletionRequest) -> bool {
8843 req.chat_history.iter().any(|m| match m {
8844 Message::User { content } => content
8845 .iter()
8846 .any(|c| matches!(c, UserContent::Text(text) if text.text.contains(SENTINEL))),
8847 _ => false,
8848 })
8849 }
8850
8851 fn messages_have_sentinel(messages: &[Message]) -> bool {
8852 messages.iter().any(|m| match m {
8853 Message::User { content } => content
8854 .iter()
8855 .any(|c| matches!(c, UserContent::Text(text) if text.text.contains(SENTINEL))),
8856 _ => false,
8857 })
8858 }
8859
8860 let blocking_model = MockCompletionModel::from_turns([MockTurn::text("done")]);
8861 let blocking_probe = blocking_model.clone();
8862 let blocking = AgentBuilder::new(blocking_model)
8863 .add_hook(HistoryOverrideHook)
8864 .build()
8865 .runner("real prompt")
8866 .run()
8867 .await
8868 .expect("blocking run should succeed");
8869 assert!(
8870 request_has_sentinel(blocking_probe.requests().first().expect("one request")),
8871 "the overridden history reaches the provider"
8872 );
8873 assert!(
8874 !messages_have_sentinel(blocking.messages.as_deref().unwrap_or_default()),
8875 "the persisted transcript is untouched by the per-turn history override"
8876 );
8877
8878 let streaming_model =
8879 MockCompletionModel::from_stream_turns([one_text_stream_turn("done")]);
8880 let streaming_probe = streaming_model.clone();
8881 let stream = AgentBuilder::new(streaming_model)
8882 .add_hook(HistoryOverrideHook)
8883 .build()
8884 .runner("real prompt")
8885 .stream()
8886 .await;
8887 let final_response = drive_to_final_response(stream).await;
8888 assert!(
8889 request_has_sentinel(streaming_probe.requests().first().expect("one request")),
8890 "the overridden history reaches the provider on the streaming surface too"
8891 );
8892 assert!(
8893 !messages_have_sentinel(final_response.messages().expect("history")),
8894 "the persisted transcript is untouched by the per-turn history override on \
8895 the streaming surface too"
8896 );
8897 }
8898
8899 #[tokio::test]
8902 async fn model_turn_finished_fires_once_per_accepted_turn_including_tool_only() {
8903 let blocking_hook = RecordingHook::default();
8904 AgentBuilder::new(blocking_model())
8905 .tool(MockAddTool)
8906 .add_hook(blocking_hook.clone())
8907 .build()
8908 .runner("add 2 and 3")
8909 .max_turns(3)
8910 .run()
8911 .await
8912 .expect("blocking run should succeed");
8913 assert_eq!(
8914 blocking_hook.count(StepEventKind::ModelTurnFinished),
8915 2,
8916 "one ModelTurnFinished per accepted turn (tool turn + text turn)"
8917 );
8918
8919 let streaming_hook = RecordingHook::default();
8920 let mut stream = AgentBuilder::new(streaming_model())
8921 .tool(MockAddTool)
8922 .add_hook(streaming_hook.clone())
8923 .build()
8924 .runner("add 2 and 3")
8925 .max_turns(3)
8926 .stream()
8927 .await;
8928 while let Some(item) = stream.next().await {
8929 let _ = item.map_err(|err| panic!("stream item errored: {err}"));
8930 }
8931 assert_eq!(
8932 streaming_hook.count(StepEventKind::ModelTurnFinished),
8933 2,
8934 "ModelTurnFinished fires once per turn on the streaming surface too"
8935 );
8936 assert_eq!(
8940 streaming_hook.count(StepEventKind::StreamResponseFinish),
8941 1,
8942 "the tool-only turn fires no StreamResponseFinish"
8943 );
8944 }
8945
8946 #[tokio::test]
8947 async fn reasoning_only_turn_does_not_gain_stream_response_finish() {
8948 let hook = RecordingHook::default();
8949 let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([[
8950 MockStreamEvent::reasoning_delta("think"),
8951 MockStreamEvent::reasoning("think"),
8952 MockStreamEvent::final_response_with_total_tokens(0),
8953 ]]))
8954 .add_hook(hook.clone())
8955 .build()
8956 .runner("reason")
8957 .stream()
8958 .await;
8959 while let Some(item) = stream.next().await {
8960 item.expect("reasoning-only stream item");
8961 }
8962
8963 assert_eq!(
8964 hook.count(StepEventKind::StreamResponseFinish),
8965 0,
8966 "reasoning-only turns must not fire StreamResponseFinish"
8967 );
8968 assert_eq!(
8969 hook.count(StepEventKind::ModelTurnFinished),
8970 1,
8971 "the accepted reasoning-only turn still fires ModelTurnFinished"
8972 );
8973 assert_eq!(
8974 hook.count(StepEventKind::ReasoningDelta),
8975 1,
8976 "the reasoning fragment is observed once and its completed restatement is not a delta"
8977 );
8978 }
8979
8980 #[derive(Clone, Default)]
8982 struct CaptureFirstTurnContent {
8983 kinds: Arc<Mutex<Option<Vec<&'static str>>>>,
8984 }
8985
8986 impl AgentHook for CaptureFirstTurnContent {
8987 async fn on_model_turn_finished(
8988 &self,
8989 _ctx: &HookContext,
8990 event: ModelTurnFinished<'_>,
8991 ) -> ModelTurnAction {
8992 if let ModelTurnFinished { turn, content, .. } = event
8993 && turn == 1
8994 {
8995 let kinds = content
8996 .iter()
8997 .map(|c| match c {
8998 AssistantContent::Reasoning(_) => "reasoning",
8999 AssistantContent::Text(_) => "text",
9000 AssistantContent::ToolCall(_) => "tool_call",
9001 _ => "other",
9002 })
9003 .collect();
9004 *self.kinds.lock().expect("kinds") = Some(kinds);
9005 }
9006 ModelTurnAction::continue_run()
9007 }
9008 }
9009
9010 #[tokio::test]
9017 async fn streaming_model_turn_finished_carries_canonical_committed_content() {
9018 let model = MockCompletionModel::from_stream_turns([
9019 vec![
9020 MockStreamEvent::reasoning("think"),
9021 MockStreamEvent::tool_call("tc1", "add", json!({"x": 2, "y": 3})),
9022 MockStreamEvent::text("answer"),
9023 MockStreamEvent::final_response_with_total_tokens(0),
9024 ],
9025 vec![
9026 MockStreamEvent::text("done"),
9027 MockStreamEvent::final_response_with_total_tokens(0),
9028 ],
9029 ]);
9030 let hook = CaptureFirstTurnContent::default();
9031 let stream = AgentBuilder::new(model)
9032 .tool(MockAddTool)
9033 .add_hook(hook.clone())
9034 .build()
9035 .runner("go")
9036 .max_turns(3)
9037 .stream()
9038 .await;
9039 let _ = drive_to_final_response(stream).await;
9040
9041 assert_eq!(
9042 hook.kinds.lock().expect("kinds").clone(),
9043 Some(vec!["reasoning", "text", "tool_call"]),
9044 "ModelTurnFinished carries the canonical reasoning->text->tool ordering \
9045 from StreamedTurn::finish, not the raw stream.choice emission order"
9046 );
9047 }
9048
9049 #[tokio::test]
9052 async fn chained_rewrites_compose_across_hooks() {
9053 struct SetArg {
9055 key: &'static str,
9056 value: i64,
9057 }
9058 impl AgentHook for SetArg {
9059 async fn on_tool_call(
9060 &self,
9061 _ctx: &HookContext,
9062 event: ToolCall<'_>,
9063 ) -> ToolCallAction {
9064 if let ToolCall { args, .. } = event {
9065 let mut parsed: serde_json::Value =
9066 serde_json::from_str(args).unwrap_or_else(|_| json!({}));
9067 parsed[self.key] = json!(self.value);
9068 ToolCallAction::rewrite(parsed)
9069 } else {
9070 ToolCallAction::run()
9071 }
9072 }
9073 }
9074
9075 struct WrapResult(&'static str);
9077 impl AgentHook for WrapResult {
9078 async fn on_tool_result(
9079 &self,
9080 _ctx: &HookContext,
9081 event: ToolResultEvent<'_>,
9082 ) -> ToolResultAction {
9083 if let ToolResultEvent { presentation, .. } = event {
9084 ToolResultAction::rewrite(format!("{}({})", self.0, presentation.render()))
9085 } else {
9086 ToolResultAction::keep()
9087 }
9088 }
9089 }
9090
9091 let recorder = RecordingHook::default();
9096 let blocking = AgentBuilder::new(blocking_model())
9097 .tool(MockAddTool)
9098 .add_hook(SetArg {
9099 key: "y",
9100 value: 40,
9101 })
9102 .add_hook(SetArg {
9103 key: "x",
9104 value: 100,
9105 })
9106 .add_hook(WrapResult("A"))
9107 .add_hook(WrapResult("B"))
9108 .add_hook(recorder.clone())
9109 .build()
9110 .runner("add 2 and 3")
9111 .max_turns(3)
9112 .run()
9113 .await
9114 .expect("blocking run should succeed");
9115 assert_eq!(blocking.output, "the answer is 5");
9116 assert_eq!(
9117 recorder.tool_results(),
9118 vec!["B(A(140))".to_string()],
9119 "arg rewrites compose (100+40=140) and result rewrites nest B(A(...))"
9120 );
9121
9122 let stream_recorder = RecordingHook::default();
9124 let mut stream = AgentBuilder::new(streaming_model())
9125 .tool(MockAddTool)
9126 .add_hook(SetArg {
9127 key: "y",
9128 value: 40,
9129 })
9130 .add_hook(SetArg {
9131 key: "x",
9132 value: 100,
9133 })
9134 .add_hook(WrapResult("A"))
9135 .add_hook(WrapResult("B"))
9136 .add_hook(stream_recorder.clone())
9137 .build()
9138 .runner("add 2 and 3")
9139 .max_turns(3)
9140 .stream()
9141 .await;
9142 while let Some(item) = stream.next().await {
9143 let _ = item.map_err(|err| panic!("stream item errored: {err}"));
9144 }
9145 assert_eq!(
9146 stream_recorder.tool_results(),
9147 vec!["B(A(140))".to_string()],
9148 "chained rewrites compose identically on the streaming surface"
9149 );
9150 }
9151
9152 #[derive(serde::Deserialize, schemars::JsonSchema)]
9153 #[allow(dead_code)]
9154 struct Answer {
9155 answer: String,
9156 }
9157
9158 struct FinalResultTool;
9162
9163 impl Tool for FinalResultTool {
9164 const NAME: &'static str = "final_result";
9165 type Error = MockToolError;
9166 type Args = serde_json::Value;
9167 type Output = String;
9168
9169 fn description(&self) -> String {
9170 "A real tool sharing the default output-tool name".to_string()
9171 }
9172
9173 fn parameters(&self) -> serde_json::Value {
9174 json!({ "type": "object", "properties": {} })
9175 }
9176
9177 async fn call(
9178 &self,
9179 _context: &mut ToolContext,
9180 _args: Self::Args,
9181 ) -> Result<Self::Output, Self::Error> {
9182 Ok("real final_result output".to_string())
9183 }
9184 }
9185
9186 #[derive(Default)]
9189 struct LateFinalResultIndex {
9190 searches: AtomicU32,
9191 }
9192
9193 impl VectorStoreIndex for LateFinalResultIndex {
9194 type Filter = Filter<serde_json::Value>;
9195
9196 async fn top_n<T: for<'a> Deserialize<'a> + WasmCompatSend>(
9197 &self,
9198 _req: VectorSearchRequest,
9199 ) -> Result<Vec<(f64, String, T)>, VectorStoreError> {
9200 Ok(Vec::new())
9201 }
9202
9203 async fn top_n_ids(
9204 &self,
9205 _req: VectorSearchRequest,
9206 ) -> Result<Vec<(f64, String)>, VectorStoreError> {
9207 if self.searches.fetch_add(1, SeqCst) == 0 {
9208 Ok(Vec::new())
9209 } else {
9210 Ok(vec![(1.0, "final_result".to_string())])
9211 }
9212 }
9213 }
9214
9215 #[derive(Clone)]
9220 struct RegisterLateFinalResultTool {
9221 handle: ToolServerHandle,
9222 second_turn_patch: Option<RequestPatch>,
9223 }
9224
9225 impl AgentHook for RegisterLateFinalResultTool {
9226 async fn on_model_turn_finished(
9227 &self,
9228 ctx: &HookContext,
9229 _event: ModelTurnFinished<'_>,
9230 ) -> ModelTurnAction {
9231 if ctx.turn() == 1 {
9232 self.handle.add_tool(FinalResultTool).await;
9233 }
9234
9235 ModelTurnAction::continue_run()
9236 }
9237
9238 async fn on_completion_call(
9239 &self,
9240 ctx: &HookContext,
9241 _event: CompletionCallEvent<'_>,
9242 ) -> CompletionCallAction {
9243 if ctx.turn() == 2
9244 && let Some(patch) = &self.second_turn_patch
9245 {
9246 return CompletionCallAction::patch(patch.clone());
9247 }
9248
9249 CompletionCallAction::continue_run()
9250 }
9251 }
9252
9253 fn assert_structured_output_collision_error(message: &str) {
9254 assert!(
9255 message.contains("final_result"),
9256 "error should name the conflicting tool: {message}"
9257 );
9258 assert!(
9259 message.contains("structured-output") && message.contains("reserved"),
9260 "error should explain the structured-output reservation: {message}"
9261 );
9262 assert!(
9263 message.contains("rename or remove"),
9264 "error should provide an actionable resolution: {message}"
9265 );
9266 }
9267
9268 #[tokio::test]
9271 async fn initial_output_tool_collision_uses_a_unique_synthetic_name() {
9272 let model = MockCompletionModel::from_turns([
9273 MockTurn::tool_call("real", "final_result", json!({})),
9274 MockTurn::tool_call("output", "final_result_1", json!({ "answer": "done" })),
9275 ]);
9276 let probe = model.clone();
9277 let response = AgentBuilder::new(model)
9278 .tool(FinalResultTool)
9279 .output_schema::<Answer>()
9280 .output_mode(OutputMode::Tool)
9281 .build()
9282 .runner("go")
9283 .max_turns(2)
9284 .run()
9285 .await
9286 .expect("the real tool should dispatch before the unique output tool finalizes");
9287
9288 assert!(response.output.contains("done"));
9289 let requests = probe.requests();
9290 assert_eq!(
9291 requests.len(),
9292 2,
9293 "real-tool dispatch must continue to a second model turn"
9294 );
9295 let tool_names = requests[0]
9296 .tools
9297 .iter()
9298 .map(|tool| tool.name.as_str())
9299 .collect::<Vec<_>>();
9300 assert_eq!(tool_names.len(), 2);
9301 for expected in ["final_result", "final_result_1"] {
9302 assert_eq!(
9303 tool_names.iter().filter(|name| **name == expected).count(),
9304 1,
9305 "the first request should advertise `{expected}` exactly once: {tool_names:?}"
9306 );
9307 }
9308
9309 assert!(
9310 requests[1].chat_history.iter().any(|message| matches!(
9311 message,
9312 Message::User { content }
9313 if content.iter().any(|item| matches!(
9314 item,
9315 UserContent::ToolResult(result)
9316 if result.call == "real"
9317 && result.content.iter().any(|content| matches!(
9318 content,
9319 rig_core::message::ToolResultContent::Text(text)
9320 if text.text == "real final_result output"
9321 ))
9322 ))
9323 )),
9324 "the real `final_result` call must execute normally and its result must reach the follow-up request"
9325 );
9326 }
9327
9328 #[tokio::test]
9333 async fn late_output_tool_collision_fails_before_blocking_provider_for_all_choices() {
9334 let cases = [
9335 ("inherited", None),
9336 (
9337 "required",
9338 Some(RequestPatch::new().tool_choice(ToolChoice::Required)),
9339 ),
9340 (
9341 "none",
9342 Some(RequestPatch::new().tool_choice(ToolChoice::None)),
9343 ),
9344 (
9345 "specific",
9346 Some(RequestPatch::new().tool_choice(ToolChoice::Specific {
9347 function_names: vec!["final_result".to_string()],
9348 })),
9349 ),
9350 ];
9351
9352 for (case, second_turn_patch) in cases {
9353 let handle = ToolServer::new().tool(MockAddTool).run();
9354 let model = MockCompletionModel::from_turns([
9355 MockTurn::tool_call("add-1", "add", json!({ "x": 1, "y": 2 })),
9356 MockTurn::tool_call(
9357 "shadowed",
9358 "final_result",
9359 json!({ "answer": "wrongly finalized" }),
9360 ),
9361 ]);
9362 let probe = model.clone();
9363 let err = AgentBuilder::new(model)
9364 .tool_server_handle(handle.clone())
9365 .output_schema::<Answer>()
9366 .output_mode(OutputMode::Tool)
9367 .add_hook(RegisterLateFinalResultTool {
9368 handle,
9369 second_turn_patch,
9370 })
9371 .build()
9372 .runner("go")
9373 .max_turns(3)
9374 .run()
9375 .await
9376 .unwrap_err();
9377
9378 assert!(
9379 matches!(
9380 &err,
9381 PromptError::CompletionError(CompletionError::RequestError(_))
9382 ),
9383 "{case}: expected a local completion request error, got {err:?}"
9384 );
9385 assert_eq!(
9386 probe.request_count(),
9387 1,
9388 "{case}: the colliding second request must not reach the provider"
9389 );
9390 assert_structured_output_collision_error(&err.to_string());
9391 }
9392 }
9393
9394 #[tokio::test]
9397 async fn late_output_tool_collision_fails_before_streaming_provider() {
9398 let handle = ToolServer::new().tool(MockAddTool).run();
9399 let model = MockCompletionModel::from_stream_turns([
9400 vec![
9401 MockStreamEvent::tool_call("add-1", "add", json!({ "x": 1, "y": 2 })),
9402 MockStreamEvent::final_response_with_total_tokens(0),
9403 ],
9404 vec![
9405 MockStreamEvent::tool_call(
9406 "shadowed",
9407 "final_result",
9408 json!({ "answer": "wrongly finalized" }),
9409 ),
9410 MockStreamEvent::final_response_with_total_tokens(0),
9411 ],
9412 ]);
9413 let probe = model.clone();
9414 let mut stream = AgentBuilder::new(model)
9415 .tool_server_handle(handle.clone())
9416 .output_schema::<Answer>()
9417 .output_mode(OutputMode::Tool)
9418 .add_hook(RegisterLateFinalResultTool {
9419 handle,
9420 second_turn_patch: None,
9421 })
9422 .build()
9423 .runner("go")
9424 .max_turns(3)
9425 .stream()
9426 .await;
9427
9428 let mut collisions = Vec::new();
9429 let mut saw_final_response = false;
9430 while let Some(item) = stream.next().await {
9431 match item {
9432 Err(err) => collisions.push(err),
9433 Ok(MultiTurnStreamItem::FinalResponse(_)) => saw_final_response = true,
9434 Ok(_) => {}
9435 }
9436 }
9437 assert_eq!(
9438 collisions.len(),
9439 1,
9440 "the stream should terminate with exactly one collision error"
9441 );
9442 assert!(
9443 !saw_final_response,
9444 "a collision error must terminate the stream without a final response"
9445 );
9446 let err = collisions.pop().expect("one collision error was asserted");
9447
9448 assert!(
9449 matches!(
9450 &err,
9451 StreamingError::Completion(CompletionError::RequestError(_))
9452 ),
9453 "expected a local streaming completion request error, got {err:?}"
9454 );
9455 assert_eq!(
9456 probe.request_count(),
9457 1,
9458 "the colliding second stream must not reach the provider"
9459 );
9460 assert_structured_output_collision_error(&err.to_string());
9461 }
9462
9463 #[tokio::test]
9467 async fn late_output_tool_collision_is_checked_after_active_tools_filtering() {
9468 let handle = ToolServer::new().tool(MockAddTool).run();
9469 let model = MockCompletionModel::from_turns([
9470 MockTurn::tool_call("add-1", "add", json!({ "x": 1, "y": 2 })),
9471 MockTurn::tool_call("add-2", "add", json!({ "x": 3, "y": 4 })),
9472 MockTurn::tool_call(
9473 "shadowed",
9474 "final_result",
9475 json!({ "answer": "wrongly finalized" }),
9476 ),
9477 ]);
9478 let probe = model.clone();
9479 let err = AgentBuilder::new(model)
9480 .tool_server_handle(handle.clone())
9481 .output_schema::<Answer>()
9482 .output_mode(OutputMode::Tool)
9483 .add_hook(RegisterLateFinalResultTool {
9484 handle,
9485 second_turn_patch: Some(RequestPatch::new().active_tools(["add"])),
9486 })
9487 .build()
9488 .runner("go")
9489 .max_turns(4)
9490 .run()
9491 .await
9492 .expect_err("the exposed third-turn collision should fail locally");
9493
9494 assert_eq!(
9495 probe.request_count(),
9496 2,
9497 "the filtered second turn may run, but the exposed third turn may not"
9498 );
9499 let requests = probe.requests();
9500 let second_turn_names = requests[1]
9501 .tools
9502 .iter()
9503 .map(|tool| tool.name.as_str())
9504 .collect::<Vec<_>>();
9505 assert_eq!(second_turn_names.len(), 2);
9506 for expected in ["add", "final_result"] {
9507 assert_eq!(
9508 second_turn_names
9509 .iter()
9510 .filter(|name| **name == expected)
9511 .count(),
9512 1,
9513 "the second request should advertise `{expected}` exactly once: \
9514 {second_turn_names:?}"
9515 );
9516 }
9517 assert_structured_output_collision_error(&err.to_string());
9518 }
9519
9520 #[tokio::test]
9524 async fn retrieved_output_tool_collision_fails_before_provider_request() {
9525 let mut retrieved_tools = ToolSet::default();
9526 retrieved_tools.add_tool(FinalResultTool);
9527 let handle = ToolServer::new()
9528 .tool(MockAddTool)
9529 .retrieved_tools(1, LateFinalResultIndex::default(), retrieved_tools)
9530 .run();
9531 let model = MockCompletionModel::from_turns([
9532 MockTurn::tool_call("add-1", "add", json!({ "x": 1, "y": 2 })),
9533 MockTurn::tool_call(
9534 "shadowed",
9535 "final_result",
9536 json!({ "answer": "wrongly finalized" }),
9537 ),
9538 ]);
9539 let probe = model.clone();
9540 let err = AgentBuilder::new(model)
9541 .tool_server_handle(handle)
9542 .output_schema::<Answer>()
9543 .output_mode(OutputMode::Tool)
9544 .build()
9545 .runner("go")
9546 .max_turns(3)
9547 .run()
9548 .await
9549 .expect_err("the retrieved second-turn collision should fail locally");
9550
9551 assert!(matches!(
9552 &err,
9553 PromptError::CompletionError(CompletionError::RequestError(_))
9554 ));
9555 assert_eq!(
9556 probe.request_count(),
9557 1,
9558 "the colliding retrieved tool must prevent the second provider request"
9559 );
9560 assert_structured_output_collision_error(&err.to_string());
9561 }
9562
9563 struct ActiveToolsAddOnly;
9566
9567 impl AgentHook for ActiveToolsAddOnly {
9568 async fn on_completion_call(
9569 &self,
9570 _ctx: &HookContext,
9571 event: CompletionCallEvent<'_>,
9572 ) -> CompletionCallAction {
9573 if let CompletionCallEvent { .. } = event {
9574 CompletionCallAction::patch(RequestPatch::new().active_tools(["add"]))
9575 } else {
9576 CompletionCallAction::continue_run()
9577 }
9578 }
9579 }
9580
9581 #[tokio::test]
9590 async fn active_tools_filter_does_not_let_output_tool_collide_with_a_filtered_real_tool() {
9591 let model = MockCompletionModel::from_turns([MockTurn::tool_call(
9597 "out1",
9598 "final_result_1",
9599 json!({ "answer": "done" }),
9600 )]);
9601 let probe = model.clone();
9602 let response = AgentBuilder::new(model)
9603 .tool(MockAddTool)
9604 .tool(FinalResultTool)
9605 .output_schema::<Answer>()
9606 .output_mode(OutputMode::Tool)
9607 .add_hook(ActiveToolsAddOnly)
9608 .build()
9609 .runner("go")
9610 .max_turns(2)
9611 .run()
9612 .await
9613 .expect("run should finalize via the picked output tool `final_result_1`");
9614 assert!(
9615 response.output.contains("done"),
9616 "the intercepted output-tool call should produce the structured result, \
9617 got {:?}",
9618 response.output
9619 );
9620
9621 let requests = probe.requests();
9622 assert!(
9623 !requests.is_empty(),
9624 "the first model request should be captured"
9625 );
9626 let tool_names: Vec<&str> = requests[0].tools.iter().map(|t| t.name.as_str()).collect();
9627 assert!(
9628 tool_names.contains(&"add"),
9629 "active_tools keeps `add` advertised, saw {tool_names:?}"
9630 );
9631 assert!(
9632 tool_names.contains(&"final_result_1"),
9633 "the synthetic output tool must avoid the filtered real `final_result` name, \
9634 saw {tool_names:?}"
9635 );
9636 assert!(
9637 !tool_names.contains(&"final_result"),
9638 "the real `final_result` is filtered out and the output tool must not reuse \
9639 its name, saw {tool_names:?}"
9640 );
9641 }
9642
9643 #[derive(Clone, Default)]
9646 struct CaptureOutputToolInModelTurn {
9647 saw_output_tool_call: Arc<Mutex<bool>>,
9648 }
9649
9650 impl AgentHook for CaptureOutputToolInModelTurn {
9651 async fn on_model_turn_finished(
9652 &self,
9653 _ctx: &HookContext,
9654 event: ModelTurnFinished<'_>,
9655 ) -> ModelTurnAction {
9656 if let ModelTurnFinished { content, .. } = event
9657 && content.iter().any(|c| {
9658 matches!(c, AssistantContent::ToolCall(tc) if tc.function.name == "final_result")
9659 })
9660 {
9661 *self.saw_output_tool_call.lock().expect("lock") = true;
9662 }
9663 ModelTurnAction::continue_run()
9664 }
9665 }
9666
9667 #[tokio::test]
9673 async fn model_turn_finished_content_carries_output_tool_call_in_tool_mode() {
9674 let hook = CaptureOutputToolInModelTurn::default();
9676 let response = AgentBuilder::new(MockCompletionModel::from_turns([MockTurn::tool_call(
9677 "out1",
9678 "final_result",
9679 json!({ "answer": "done" }),
9680 )]))
9681 .output_schema::<Answer>()
9682 .output_mode(OutputMode::Tool)
9683 .add_hook(hook.clone())
9684 .build()
9685 .runner("go")
9686 .max_turns(2)
9687 .run()
9688 .await
9689 .expect("run should finalize via the output tool");
9690 assert!(
9691 *hook.saw_output_tool_call.lock().expect("lock"),
9692 "ModelTurnFinished.content must carry the model-emitted output-tool call (blocking)"
9693 );
9694 assert!(
9695 response.output.contains("done"),
9696 "the run finalizes with the structured output, not the raw tool call: {:?}",
9697 response.output
9698 );
9699
9700 let s_hook = CaptureOutputToolInModelTurn::default();
9702 let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([vec![
9703 MockStreamEvent::tool_call("out1", "final_result", json!({ "answer": "done" })),
9704 MockStreamEvent::final_response_with_total_tokens(0),
9705 ]]))
9706 .output_schema::<Answer>()
9707 .output_mode(OutputMode::Tool)
9708 .add_hook(s_hook.clone())
9709 .build()
9710 .runner("go")
9711 .max_turns(2)
9712 .stream()
9713 .await;
9714 while stream.next().await.is_some() {}
9715 assert!(
9716 *s_hook.saw_output_tool_call.lock().expect("lock"),
9717 "ModelTurnFinished.content must carry the model-emitted output-tool call (streaming)"
9718 );
9719 }
9720
9721 #[tokio::test]
9727 async fn output_tool_finalization_emits_no_complete_tool_call_stream_item() {
9728 let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([vec![
9729 MockStreamEvent::tool_call("out1", "final_result", json!({ "answer": "done" })),
9730 MockStreamEvent::final_response_with_total_tokens(0),
9731 ]]))
9732 .output_schema::<Answer>()
9733 .output_mode(OutputMode::Tool)
9734 .build()
9735 .runner("go")
9736 .max_turns(2)
9737 .stream()
9738 .await;
9739
9740 let mut saw_complete_output_tool_call = false;
9741 let mut final_has_output = false;
9742 while let Some(item) = stream.next().await {
9743 match item.expect("stream item") {
9744 MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::ToolCall {
9745 tool_call,
9746 ..
9747 }) if tool_call.function.name == "final_result" => {
9748 saw_complete_output_tool_call = true;
9749 }
9750 MultiTurnStreamItem::FinalResponse(res) => {
9751 final_has_output = res.output().contains("done");
9752 }
9753 _ => {}
9754 }
9755 }
9756 assert!(
9757 !saw_complete_output_tool_call,
9758 "the output-tool call finalizes the run, so no complete \
9759 StreamAssistantItem::ToolCall item must be emitted for it"
9760 );
9761 assert!(
9762 final_has_output,
9763 "the structured output must be surfaced via the FinalResponse"
9764 );
9765 }
9766
9767 enum Decision {
9776 Approve,
9778 Deny(&'static str),
9780 Edit(serde_json::Value),
9782 Abort(&'static str),
9784 }
9785
9786 #[derive(Clone)]
9790 struct HumanApprovalHook {
9791 decisions: Arc<Mutex<std::collections::VecDeque<Decision>>>,
9792 reviewed: Arc<Mutex<Vec<String>>>,
9793 }
9794
9795 impl HumanApprovalHook {
9796 fn new(decisions: impl IntoIterator<Item = Decision>) -> Self {
9797 Self {
9798 decisions: Arc::new(Mutex::new(decisions.into_iter().collect())),
9799 reviewed: Arc::new(Mutex::new(Vec::new())),
9800 }
9801 }
9802
9803 fn reviewed(&self) -> Vec<String> {
9805 self.reviewed.lock().unwrap().clone()
9806 }
9807 }
9808
9809 impl AgentHook for HumanApprovalHook {
9810 async fn on_tool_call(&self, _ctx: &HookContext, event: ToolCall<'_>) -> ToolCallAction {
9811 let ToolCall {
9812 tool_name, args, ..
9813 } = event
9814 else {
9815 return ToolCallAction::run();
9816 };
9817 self.reviewed
9818 .lock()
9819 .unwrap()
9820 .push(format!("{tool_name}({args})"));
9821 let decision = self.decisions.lock().unwrap().pop_front();
9822 match decision {
9823 Some(Decision::Approve) => ToolCallAction::run(),
9824 Some(Decision::Deny(reason)) => ToolCallAction::skip(reason),
9825 Some(Decision::Edit(args)) => ToolCallAction::rewrite(args),
9826 Some(Decision::Abort(reason)) => ToolCallAction::stop(reason),
9827 None => ToolCallAction::skip("denied: no scripted decision (fail-closed)"),
9830 }
9831 }
9832 }
9833
9834 #[tokio::test]
9840 async fn human_in_the_loop_approve_deny_edit_parity_across_run_and_stream() {
9841 let turns = [
9843 ScriptedTurn::ToolCalls(vec![
9844 add_call("tc1", 2, 3), add_call("tc2", 10, 20), add_call("tc3", 1, 1), ]),
9848 ScriptedTurn::Text("done"),
9849 ];
9850 let denial = "denied by reviewer: amount too large";
9851 let decisions = || {
9852 vec![
9853 Decision::Approve,
9854 Decision::Deny(denial),
9855 Decision::Edit(json!({"x": 1, "y": 100})),
9856 ]
9857 };
9858
9859 let blocking_model =
9860 MockCompletionModel::from_turns(turns.iter().map(ScriptedTurn::as_blocking_turn));
9861 let blocking_recorder = RecordingHook::default();
9862 let blocking_approver = HumanApprovalHook::new(decisions());
9863 let blocking = AgentBuilder::new(blocking_model)
9864 .tool(MockAddTool)
9865 .build()
9866 .runner("carry out the plan")
9867 .max_turns(3)
9868 .add_hook(blocking_recorder.clone())
9869 .add_hook(blocking_approver.clone())
9870 .run()
9871 .await
9872 .expect("blocking HITL run should succeed");
9873
9874 let streaming_model = MockCompletionModel::from_stream_turns(
9875 turns
9876 .iter()
9877 .map(|turn| turn.as_stream_events(StreamShape::Complete)),
9878 );
9879 let streaming_recorder = RecordingHook::default();
9880 let streaming_approver = HumanApprovalHook::new(decisions());
9881 let mut stream = AgentBuilder::new(streaming_model)
9882 .tool(MockAddTool)
9883 .build()
9884 .runner("carry out the plan")
9885 .max_turns(3)
9886 .add_hook(streaming_recorder.clone())
9887 .add_hook(streaming_approver.clone())
9888 .stream()
9889 .await;
9890 let mut final_response = None;
9891 while let Some(item) = stream.next().await {
9892 if let Ok(MultiTurnStreamItem::FinalResponse(resp)) =
9893 item.map_err(|err| panic!("stream item errored: {err}"))
9894 {
9895 final_response = Some(resp);
9896 }
9897 }
9898 let final_response = final_response.expect("stream should yield a final response");
9899
9900 assert_eq!(
9905 blocking_recorder.tool_results(),
9906 vec![
9907 "5".to_string(),
9908 "denied by reviewer: amount too large".to_string(),
9909 "101".to_string()
9910 ]
9911 );
9912 assert_eq!(
9913 blocking_recorder.tool_results(),
9914 streaming_recorder.tool_results()
9915 );
9916
9917 assert!(
9921 !blocking_recorder.tool_results().contains(&"30".to_string()),
9922 "the denied call must not have executed"
9923 );
9924
9925 let reviewed = blocking_approver.reviewed();
9929 assert_eq!(reviewed.len(), 3);
9930 assert_eq!(reviewed, streaming_approver.reviewed());
9931 assert!(
9932 reviewed[0].contains('2') && reviewed[0].contains('3'),
9933 "first reviewed call should be add(2, 3): {reviewed:?}"
9934 );
9935 assert!(
9936 reviewed[1].contains("10") && reviewed[1].contains("20"),
9937 "the denied (second) call should be add(10, 20): {reviewed:?}"
9938 );
9939
9940 assert_eq!(blocking.output, "done");
9941 assert_eq!(final_response.output(), blocking.output);
9942 assert_eq!(
9943 blocking_recorder.shared_events(),
9944 streaming_recorder.shared_events()
9945 );
9946
9947 let blocking_messages = blocking.messages.expect("blocking messages");
9951 let streaming_messages = final_response
9952 .messages()
9953 .expect("streaming history")
9954 .to_vec();
9955 assert_eq!(
9956 serde_json::to_value(&blocking_messages).expect("serialize blocking"),
9957 serde_json::to_value(&streaming_messages).expect("serialize streaming"),
9958 );
9959 assert!(
9960 tool_result_text_in_history(&blocking_messages, denial),
9961 "the denial reason must be the denied call's tool result in the history"
9962 );
9963 assert!(
9964 tool_result_json_in_history(&blocking_messages, &json!(101)),
9965 "the edited call must have executed with the rewritten arguments"
9966 );
9967 }
9968
9969 #[tokio::test]
9973 async fn human_in_the_loop_abort_terminates_the_run() {
9974 let turns = [
9975 ScriptedTurn::ToolCalls(vec![add_call("tc1", 2, 3)]),
9976 ScriptedTurn::Text("unreachable"),
9977 ];
9978 const ABORT_REASON: &str = "aborted by the human reviewer";
9979
9980 let blocking_model =
9982 MockCompletionModel::from_turns(turns.iter().map(ScriptedTurn::as_blocking_turn));
9983 let err = AgentBuilder::new(blocking_model)
9984 .tool(MockAddTool)
9985 .build()
9986 .runner("do the sensitive thing")
9987 .max_turns(3)
9988 .add_hook(HumanApprovalHook::new([Decision::Abort(ABORT_REASON)]))
9989 .run()
9990 .await
9991 .expect_err("an aborted tool call should terminate the blocking run");
9992 assert!(
9993 format!("{err}").contains(ABORT_REASON),
9994 "the abort reason should surface in the blocking error, got: {err}"
9995 );
9996
9997 let streaming_model = MockCompletionModel::from_stream_turns(
10000 turns
10001 .iter()
10002 .map(|turn| turn.as_stream_events(StreamShape::Complete)),
10003 );
10004 let mut stream = AgentBuilder::new(streaming_model)
10005 .tool(MockAddTool)
10006 .build()
10007 .runner("do the sensitive thing")
10008 .max_turns(3)
10009 .add_hook(HumanApprovalHook::new([Decision::Abort(ABORT_REASON)]))
10010 .stream()
10011 .await;
10012 let mut stream_error = None;
10013 while let Some(item) = stream.next().await {
10014 match item {
10015 Err(err) => stream_error = Some(format!("{err}")),
10016 Ok(MultiTurnStreamItem::FinalResponse(resp)) => {
10017 panic!("aborted stream must not finalize, got: {}", resp.output())
10018 }
10019 Ok(_) => {}
10020 }
10021 }
10022 let stream_error = stream_error.expect("an aborted tool call should error the stream");
10023 assert!(
10024 stream_error.contains(ABORT_REASON),
10025 "the abort reason should surface in the streaming error, got: {stream_error}"
10026 );
10027 }
10028
10029 #[derive(Clone)]
10034 struct PolicyHook {
10035 auto_approve: std::collections::HashSet<&'static str>,
10036 evaluated: Arc<Mutex<Vec<String>>>,
10038 cache: Arc<Mutex<std::collections::HashMap<String, bool>>>,
10040 }
10041
10042 impl PolicyHook {
10043 fn new(auto_approve: impl IntoIterator<Item = &'static str>) -> Self {
10044 Self {
10045 auto_approve: auto_approve.into_iter().collect(),
10046 evaluated: Arc::new(Mutex::new(Vec::new())),
10047 cache: Arc::new(Mutex::new(std::collections::HashMap::new())),
10048 }
10049 }
10050
10051 fn evaluated(&self) -> Vec<String> {
10052 self.evaluated.lock().unwrap().clone()
10053 }
10054 }
10055
10056 impl AgentHook for PolicyHook {
10057 async fn on_tool_call(&self, _ctx: &HookContext, event: ToolCall<'_>) -> ToolCallAction {
10058 let ToolCall { tool_name, .. } = event else {
10059 return ToolCallAction::run();
10060 };
10061 let cached = self.cache.lock().unwrap().get(tool_name).copied();
10062 let approved = match cached {
10063 Some(decision) => decision, None => {
10065 self.evaluated.lock().unwrap().push(tool_name.to_string());
10066 let decision = self.auto_approve.contains(tool_name);
10067 self.cache
10068 .lock()
10069 .unwrap()
10070 .insert(tool_name.to_string(), decision);
10071 decision
10072 }
10073 };
10074 if approved {
10075 ToolCallAction::run()
10076 } else {
10077 ToolCallAction::skip(format!("denied by policy: `{tool_name}` not allowed"))
10078 }
10079 }
10080 }
10081
10082 #[tokio::test]
10086 async fn approval_policy_allow_list_with_sticky_decisions() {
10087 let turns = [
10089 ScriptedTurn::ToolCalls(vec![
10090 add_call("c1", 2, 3),
10091 ScriptedToolCall {
10092 id: "c2",
10093 name: "subtract",
10094 args: json!({ "x": 10, "y": 4 }),
10095 },
10096 add_call("c3", 2, 3),
10097 ]),
10098 ScriptedTurn::Text("done"),
10099 ];
10100
10101 let model =
10102 MockCompletionModel::from_turns(turns.iter().map(ScriptedTurn::as_blocking_turn));
10103 let recorder = RecordingHook::default();
10104 let policy = PolicyHook::new(["add"]);
10105 let out = AgentBuilder::new(model)
10106 .tool(MockAddTool)
10107 .tool(MockSubtractTool)
10108 .build()
10109 .runner("go")
10110 .max_turns(3)
10111 .add_hook(recorder.clone())
10112 .add_hook(policy.clone())
10113 .run()
10114 .await
10115 .expect("policy run should succeed");
10116
10117 assert_eq!(out.output, "done");
10118 assert_eq!(
10122 recorder.tool_results(),
10123 vec![
10124 "5".to_string(),
10125 "denied by policy: `subtract` not allowed".to_string(),
10126 "5".to_string()
10127 ]
10128 );
10129 assert_eq!(
10132 policy.evaluated(),
10133 vec!["add".to_string(), "subtract".to_string()]
10134 );
10135 let messages = out.messages.expect("messages");
10136 assert!(
10137 tool_result_text_in_history(&messages, "denied by policy: `subtract` not allowed"),
10138 "the policy denial reason must reach the model as the subtract tool result"
10139 );
10140 }
10141
10142 static NEXT_RESPONSE_RETRY_HOOK_ID: AtomicU64 = AtomicU64::new(1);
10143
10144 #[derive(Clone, Default)]
10145 struct ResponseRetryAttempts(HashMap<u64, usize>);
10146
10147 #[derive(Clone)]
10148 enum TestRetryMode {
10149 Repeat,
10150 Feedback(&'static str),
10151 }
10152
10153 #[derive(Clone)]
10156 struct BoundedResponseRetry {
10157 id: u64,
10158 rejected_text: &'static str,
10159 max_retries: usize,
10160 mode: TestRetryMode,
10161 }
10162
10163 #[derive(Clone, Default)]
10164 struct StatefulCompletionPatch {
10165 calls: Arc<AtomicU32>,
10166 }
10167
10168 impl StatefulCompletionPatch {
10169 fn calls(&self) -> u32 {
10170 self.calls.load(SeqCst)
10171 }
10172 }
10173
10174 impl AgentHook for StatefulCompletionPatch {
10175 async fn on_completion_call(
10176 &self,
10177 _ctx: &HookContext,
10178 _event: crate::agent::CompletionCallEvent<'_>,
10179 ) -> CompletionCallAction {
10180 let call = self.calls.fetch_add(1, SeqCst);
10181 CompletionCallAction::patch(RequestPatch::new().temperature(if call == 0 {
10182 0.1
10183 } else {
10184 0.9
10185 }))
10186 }
10187 }
10188
10189 impl BoundedResponseRetry {
10190 fn new(rejected_text: &'static str, max_retries: usize, mode: TestRetryMode) -> Self {
10191 Self {
10192 id: NEXT_RESPONSE_RETRY_HOOK_ID.fetch_add(1, SeqCst),
10193 rejected_text,
10194 max_retries,
10195 mode,
10196 }
10197 }
10198 }
10199
10200 impl AgentHook for BoundedResponseRetry {
10201 async fn on_model_turn_finished(
10202 &self,
10203 ctx: &HookContext,
10204 event: ModelTurnFinished<'_>,
10205 ) -> ModelTurnAction {
10206 let has_rejected_text = event.content.iter().any(
10207 |content| matches!(content, AssistantContent::Text(text) if text.text == self.rejected_text),
10208 );
10209 let rejected =
10215 has_rejected_text || (self.rejected_text.is_empty() && event.content.is_empty());
10216 if !rejected {
10217 return ModelTurnAction::continue_run();
10218 }
10219
10220 let attempt = ctx
10221 .scratchpad()
10222 .update::<ResponseRetryAttempts, _>(|attempts| {
10223 let attempt = attempts.0.entry(self.id).or_default();
10224 *attempt += 1;
10225 *attempt
10226 });
10227 if attempt > self.max_retries {
10228 return ModelTurnAction::stop(format!(
10229 "response retry limit ({}) exceeded",
10230 self.max_retries
10231 ));
10232 }
10233
10234 match self.mode {
10235 TestRetryMode::Repeat => ModelTurnAction::repeat(),
10236 TestRetryMode::Feedback(feedback) => ModelTurnAction::retry_with_feedback(feedback),
10237 }
10238 }
10239 }
10240
10241 fn retry_usage(input_tokens: u64, output_tokens: u64) -> Usage {
10242 Usage {
10243 input_tokens,
10244 output_tokens,
10245 total_tokens: input_tokens + output_tokens,
10246 ..Usage::new()
10247 }
10248 }
10249
10250 type Termination = (Option<FinishReason>, Option<u64>);
10261
10262 #[derive(Clone, Debug, Default)]
10264 struct TerminationProbe {
10265 observations: Arc<Mutex<Vec<Termination>>>,
10266 }
10267
10268 impl TerminationProbe {
10269 fn observations(&self) -> Vec<Termination> {
10270 self.observations.lock().expect("observations").clone()
10271 }
10272 }
10273
10274 impl AgentHook for TerminationProbe {
10275 async fn on_model_turn_finished(
10276 &self,
10277 _ctx: &HookContext,
10278 event: ModelTurnFinished<'_>,
10279 ) -> ModelTurnAction {
10280 self.observations
10281 .lock()
10282 .expect("observations")
10283 .push((event.finish_reason.cloned(), event.max_tokens));
10284 ModelTurnAction::continue_run()
10285 }
10286 }
10287
10288 #[derive(Clone, Debug, Default)]
10292 struct RetryOnTruncation {
10293 truncated_turns: Arc<AtomicU32>,
10296 }
10297
10298 impl AgentHook for RetryOnTruncation {
10299 async fn on_model_turn_finished(
10300 &self,
10301 _ctx: &HookContext,
10302 event: ModelTurnFinished<'_>,
10303 ) -> ModelTurnAction {
10304 let truncated = event
10305 .finish_reason
10306 .is_some_and(FinishReason::truncated_output);
10307 let has_tool_call = event
10308 .content
10309 .iter()
10310 .any(|content| matches!(content, AssistantContent::ToolCall(_)));
10311 if truncated && !has_tool_call && self.truncated_turns.fetch_add(1, SeqCst) == 0 {
10312 return ModelTurnAction::repeat();
10313 }
10314 ModelTurnAction::continue_run()
10315 }
10316 }
10317
10318 #[derive(Clone, Debug, Default)]
10322 struct EscalatingCap {
10323 calls: Arc<AtomicU32>,
10324 }
10325
10326 impl AgentHook for EscalatingCap {
10327 async fn on_completion_call(
10328 &self,
10329 _ctx: &HookContext,
10330 _event: crate::agent::CompletionCallEvent<'_>,
10331 ) -> CompletionCallAction {
10332 let call = self.calls.fetch_add(1, SeqCst);
10333 CompletionCallAction::patch(RequestPatch::new().max_tokens(if call == 0 {
10334 16
10335 } else {
10336 512
10337 }))
10338 }
10339 }
10340
10341 #[tokio::test]
10344 async fn model_turn_finished_reports_termination_and_effective_max_tokens_blocking() {
10345 let probe = TerminationProbe::default();
10346 let model = MockCompletionModel::from_turns([
10347 MockTurn::text("a partial ans").with_finish_reason(FinishReason::Length)
10348 ]);
10349
10350 AgentBuilder::new(model.clone())
10351 .max_tokens(64)
10352 .add_hook(probe.clone())
10353 .build()
10354 .runner("question")
10355 .run()
10356 .await
10357 .expect("truncated turn is still an answer");
10358
10359 assert_eq!(
10360 probe.observations(),
10361 vec![(Some(FinishReason::Length), Some(64))]
10362 );
10363 assert_eq!(model.requests()[0].max_tokens, Some(64));
10365 }
10366
10367 #[tokio::test]
10369 async fn model_turn_finished_reports_termination_and_effective_max_tokens_streaming() {
10370 let probe = TerminationProbe::default();
10371 let model = MockCompletionModel::from_stream_turns([[
10372 MockStreamEvent::Text("a partial ans".to_string()),
10373 MockStreamEvent::FinalResponse(
10374 mock_final(Usage::new()).with_finish_reason(FinishReason::Length),
10375 ),
10376 ]]);
10377
10378 let mut stream = AgentBuilder::new(model.clone())
10379 .max_tokens(64)
10380 .add_hook(probe.clone())
10381 .build()
10382 .runner("question")
10383 .stream()
10384 .await;
10385 while let Some(item) = stream.next().await {
10386 item.expect("streaming item");
10387 }
10388
10389 assert_eq!(
10390 probe.observations(),
10391 vec![(Some(FinishReason::Length), Some(64))],
10392 "the streaming surface must report the same termination metadata as blocking"
10393 );
10394 assert_eq!(model.requests()[0].max_tokens, Some(64));
10395 }
10396
10397 #[tokio::test]
10401 async fn model_turn_finished_reports_absent_reason_and_absent_cap_as_none() {
10402 let probe = TerminationProbe::default();
10403
10404 AgentBuilder::new(MockCompletionModel::from_turns([MockTurn::text("done")]))
10405 .add_hook(probe.clone())
10406 .build()
10407 .runner("question")
10408 .run()
10409 .await
10410 .expect("run");
10411
10412 assert_eq!(probe.observations(), vec![(None, None)]);
10413 }
10414
10415 #[tokio::test]
10420 async fn model_turn_finished_reports_tool_calls_for_a_mislabelled_tool_turn() {
10421 let probe = TerminationProbe::default();
10422
10423 AgentBuilder::new(MockCompletionModel::from_turns([
10424 MockTurn::tool_call("call-1", "add", json!({ "x": 1, "y": 2 }))
10425 .with_finish_reason(FinishReason::Stop),
10426 MockTurn::text("3"),
10427 ]))
10428 .tool(MockAddTool)
10429 .add_hook(probe.clone())
10430 .build()
10431 .runner("question")
10432 .max_turns(2)
10433 .run()
10434 .await
10435 .expect("tool turn");
10436
10437 assert_eq!(
10438 probe
10439 .observations()
10440 .first()
10441 .map(|(reason, _)| reason.clone()),
10442 Some(Some(FinishReason::ToolCalls)),
10443 "a provider's bare `stop` on a tool turn must not read as a natural stop"
10444 );
10445 }
10446
10447 #[tokio::test]
10452 async fn model_turn_finished_reports_tool_calls_for_a_mislabelled_streamed_tool_turn() {
10453 let probe = TerminationProbe::default();
10454 let model = MockCompletionModel::from_stream_turns([
10455 vec![
10456 MockStreamEvent::tool_call("call-1", "add", json!({ "x": 1, "y": 2 })),
10457 MockStreamEvent::FinalResponse(
10458 mock_final(Usage::new()).with_finish_reason(FinishReason::Stop),
10459 ),
10460 ],
10461 vec![
10462 MockStreamEvent::Text("3".to_string()),
10463 MockStreamEvent::FinalResponse(mock_final(Usage::new())),
10464 ],
10465 ]);
10466
10467 let mut stream = AgentBuilder::new(model)
10468 .tool(MockAddTool)
10469 .add_hook(probe.clone())
10470 .build()
10471 .runner("question")
10472 .max_turns(2)
10473 .stream()
10474 .await;
10475 while let Some(item) = stream.next().await {
10476 item.expect("streaming item");
10477 }
10478
10479 assert_eq!(
10480 probe
10481 .observations()
10482 .first()
10483 .map(|(reason, _)| reason.clone()),
10484 Some(Some(FinishReason::ToolCalls)),
10485 "streaming must reconcile a mislabelled tool turn exactly as blocking does"
10486 );
10487 }
10488
10489 #[tokio::test]
10494 async fn model_turn_finished_passes_an_unmapped_reason_through_verbatim() {
10495 let probe = TerminationProbe::default();
10496
10497 AgentBuilder::new(MockCompletionModel::from_turns([
10498 MockTurn::text("halted").with_finish_reason(FinishReason::Other("guardrail".into()))
10499 ]))
10500 .add_hook(probe.clone())
10501 .build()
10502 .runner("question")
10503 .run()
10504 .await
10505 .expect("run");
10506
10507 assert_eq!(
10508 probe.observations(),
10509 vec![(Some(FinishReason::Other("guardrail".into())), None)]
10510 );
10511 assert!(!FinishReason::Other("guardrail".into()).truncated_output());
10514 }
10515
10516 #[tokio::test]
10521 async fn streaming_retry_reports_the_second_attempts_own_effective_max_tokens() {
10522 let probe = TerminationProbe::default();
10523 let escalating = EscalatingCap::default();
10524 let model = MockCompletionModel::from_stream_turns([
10525 [
10526 MockStreamEvent::Text("rejected".to_string()),
10527 MockStreamEvent::FinalResponse(
10528 mock_final(Usage::new()).with_finish_reason(FinishReason::Length),
10529 ),
10530 ],
10531 [
10532 MockStreamEvent::Text("accepted".to_string()),
10533 MockStreamEvent::FinalResponse(
10534 mock_final(Usage::new()).with_finish_reason(FinishReason::Stop),
10535 ),
10536 ],
10537 ]);
10538
10539 let mut stream = AgentBuilder::new(model.clone())
10540 .max_tokens(64)
10541 .add_hook(escalating.clone())
10542 .add_hook(probe.clone())
10543 .add_hook(BoundedResponseRetry::new(
10544 "rejected",
10545 1,
10546 TestRetryMode::Repeat,
10547 ))
10548 .build()
10549 .runner("question")
10550 .max_turns(2)
10551 .stream()
10552 .await;
10553 while let Some(item) = stream.next().await {
10554 item.expect("streaming item");
10555 }
10556
10557 assert_eq!(
10558 probe.observations(),
10559 vec![
10560 (Some(FinishReason::Length), Some(16)),
10561 (Some(FinishReason::Stop), Some(512)),
10562 ],
10563 "streaming must report each attempt's own post-patch cap, as blocking does"
10564 );
10565 let requests = model.requests();
10566 assert_eq!(requests[0].max_tokens, Some(16));
10567 assert_eq!(requests[1].max_tokens, Some(512));
10568 }
10569
10570 #[tokio::test]
10574 async fn a_portable_hook_can_retry_a_truncated_tool_free_turn() {
10575 let model = MockCompletionModel::from_turns([
10576 MockTurn::text("cut off mid-").with_finish_reason(FinishReason::Length),
10577 MockTurn::text("a complete answer").with_finish_reason(FinishReason::Stop),
10578 ]);
10579 let hook = RetryOnTruncation::default();
10580
10581 let response = AgentBuilder::new(model.clone())
10582 .add_hook(hook.clone())
10583 .build()
10584 .runner("question")
10585 .max_turns(2)
10586 .run()
10587 .await
10588 .expect("the retried turn should answer");
10589
10590 assert_eq!(response.output, "a complete answer");
10591 assert_eq!(model.request_count(), 2, "the truncated turn was retried");
10592 assert_eq!(hook.truncated_turns.load(SeqCst), 1);
10595 }
10596
10597 #[tokio::test]
10601 async fn retry_reports_the_second_attempts_own_effective_max_tokens() {
10602 let probe = TerminationProbe::default();
10603 let escalating = EscalatingCap::default();
10604 let model = MockCompletionModel::from_turns([
10605 MockTurn::text("rejected").with_finish_reason(FinishReason::Length),
10606 MockTurn::text("accepted").with_finish_reason(FinishReason::Stop),
10607 ]);
10608
10609 AgentBuilder::new(model.clone())
10610 .max_tokens(64)
10612 .add_hook(escalating.clone())
10613 .add_hook(probe.clone())
10617 .add_hook(BoundedResponseRetry::new(
10618 "rejected",
10619 1,
10620 TestRetryMode::Repeat,
10621 ))
10622 .build()
10623 .runner("question")
10624 .max_turns(2)
10625 .run()
10626 .await
10627 .expect("repeat should recover");
10628
10629 assert_eq!(
10630 probe.observations(),
10631 vec![
10632 (Some(FinishReason::Length), Some(16)),
10633 (Some(FinishReason::Stop), Some(512)),
10634 ],
10635 "each attempt must report its own post-patch cap, not the agent baseline of 64"
10636 );
10637 let requests = model.requests();
10639 assert_eq!(requests[0].max_tokens, Some(16));
10640 assert_eq!(requests[1].max_tokens, Some(512));
10641 assert_eq!(escalating.calls.load(SeqCst), 2);
10642 }
10643
10644 #[tokio::test]
10645 async fn blocking_model_turn_repeat_preserves_prompt_history_with_fresh_preparation() {
10646 let first_usage = retry_usage(10, 3);
10647 let second_usage = retry_usage(7, 2);
10648 let completion_patch = StatefulCompletionPatch::default();
10649 let model = MockCompletionModel::from_turns([
10650 MockTurn::text("rejected").with_usage(first_usage),
10651 MockTurn::text("accepted").with_usage(second_usage),
10652 ]);
10653 let response = AgentBuilder::new(model.clone())
10654 .add_hook(completion_patch.clone())
10655 .add_hook(BoundedResponseRetry::new(
10656 "rejected",
10657 1,
10658 TestRetryMode::Repeat,
10659 ))
10660 .build()
10661 .runner("question")
10662 .max_turns(2)
10663 .run()
10664 .await
10665 .expect("repeat should recover");
10666
10667 assert_eq!(response.output, "accepted");
10668 assert_eq!(response.usage, first_usage + second_usage);
10669 assert_eq!(response.completion_calls.len(), 2);
10670 let messages = response.messages.expect("response messages");
10671 assert_eq!(
10672 messages,
10673 vec![Message::user("question"), Message::assistant("accepted")]
10674 );
10675
10676 let requests = model.requests();
10677 assert_eq!(requests.len(), 2);
10678 let first = requests[0].chat_history.clone();
10679 let second = requests[1].chat_history.clone();
10680 assert_eq!(first, vec![Message::user("question")]);
10681 assert_eq!(
10682 second, first,
10683 "Repeat must preserve the prompt and preceding history"
10684 );
10685 assert_eq!(requests[0].temperature, Some(0.1));
10686 assert_eq!(requests[1].temperature, Some(0.9));
10687 assert_eq!(completion_patch.calls(), 2);
10688 }
10689
10690 #[tokio::test]
10691 async fn blocking_model_turn_feedback_preserves_rejected_response() {
10692 let model = MockCompletionModel::from_turns([
10693 MockTurn::text("rejected"),
10694 MockTurn::text("accepted"),
10695 ]);
10696 let response = AgentBuilder::new(model.clone())
10697 .add_hook(BoundedResponseRetry::new(
10698 "rejected",
10699 1,
10700 TestRetryMode::Feedback("try another approach"),
10701 ))
10702 .build()
10703 .runner("question")
10704 .max_turns(2)
10705 .run()
10706 .await
10707 .expect("feedback retry should recover");
10708
10709 assert_eq!(response.output, "accepted");
10710 assert_eq!(
10711 response.messages.expect("response messages"),
10712 vec![
10713 Message::user("question"),
10714 Message::assistant("rejected"),
10715 Message::user("try another approach"),
10716 Message::assistant("accepted"),
10717 ]
10718 );
10719 let second_request = &model.requests()[1];
10720 assert_eq!(
10721 second_request.chat_history.clone(),
10722 vec![
10723 Message::user("question"),
10724 Message::assistant("rejected"),
10725 Message::user("try another approach"),
10726 ]
10727 );
10728 }
10729
10730 #[tokio::test]
10731 async fn blocking_empty_feedback_retry_omits_empty_assistant_history() {
10732 let first_usage = retry_usage(5, 1);
10733 let second_usage = retry_usage(7, 2);
10734 let model = MockCompletionModel::from_turns([
10735 MockTurn::text("").with_usage(first_usage),
10736 MockTurn::text("accepted").with_usage(second_usage),
10737 ]);
10738 let response = AgentBuilder::new(model.clone())
10739 .add_hook(BoundedResponseRetry::new(
10740 "",
10741 1,
10742 TestRetryMode::Feedback("provide an answer"),
10743 ))
10744 .build()
10745 .runner("question")
10746 .max_turns(2)
10747 .run()
10748 .await
10749 .expect("feedback retry should recover from an empty turn");
10750
10751 assert_eq!(response.output, "accepted");
10752 assert_eq!(response.usage, first_usage + second_usage);
10753 assert_eq!(response.completion_calls.len(), 2);
10754 assert_eq!(
10755 response.messages.expect("response messages"),
10756 vec![
10757 Message::user("question"),
10758 Message::user("provide an answer"),
10759 Message::assistant("accepted"),
10760 ]
10761 );
10762 assert_eq!(
10763 model.requests()[1].chat_history.clone(),
10764 vec![
10765 Message::user("question"),
10766 Message::user("provide an answer"),
10767 ],
10768 "the retry request must not contain an empty assistant message"
10769 );
10770 }
10771
10772 #[tokio::test]
10773 async fn streaming_model_turn_retry_marks_rollback_and_matches_blocking_accounting() {
10774 let first_usage = retry_usage(10, 3);
10775 let second_usage = retry_usage(7, 2);
10776 let model = MockCompletionModel::from_stream_turns([
10777 [
10778 MockStreamEvent::text("rejected"),
10779 MockStreamEvent::final_response(first_usage),
10780 ],
10781 [
10782 MockStreamEvent::text("accepted"),
10783 MockStreamEvent::final_response(second_usage),
10784 ],
10785 ]);
10786 let mut stream = AgentBuilder::new(model.clone())
10787 .add_hook(BoundedResponseRetry::new(
10788 "rejected",
10789 1,
10790 TestRetryMode::Repeat,
10791 ))
10792 .build()
10793 .runner("question")
10794 .max_turns(2)
10795 .stream()
10796 .await;
10797
10798 let mut retries = Vec::new();
10799 let mut provider_finals = 0;
10800 let mut completion_calls = 0;
10801 let mut final_response = None;
10802 while let Some(item) = stream.next().await {
10803 match item.expect("stream item") {
10804 MultiTurnStreamItem::ModelTurnRetried { turn } => retries.push(turn),
10805 MultiTurnStreamItem::CompletionCall(_) => completion_calls += 1,
10806 MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Final(_)) => {
10807 provider_finals += 1
10808 }
10809 MultiTurnStreamItem::FinalResponse(response) => final_response = Some(response),
10810 _ => {}
10811 }
10812 }
10813
10814 assert_eq!(retries, vec![1]);
10815 assert_eq!(
10816 provider_finals, 1,
10817 "the rejected provider final is suppressed"
10818 );
10819 assert_eq!(completion_calls, 2);
10820 let response = final_response.expect("run final response");
10821 assert_eq!(response.output, "accepted");
10822 assert_eq!(response.usage, first_usage + second_usage);
10823 assert_eq!(response.completion_calls.len(), 2);
10824 assert_eq!(
10825 response.messages.expect("response messages"),
10826 vec![Message::user("question"), Message::assistant("accepted")]
10827 );
10828 assert_eq!(model.requests().len(), 2);
10829 }
10830
10831 #[tokio::test]
10832 async fn streaming_feedback_retry_matches_blocking_history_and_usage() {
10833 let first_usage = retry_usage(5, 2);
10834 let second_usage = retry_usage(8, 4);
10835 let blocking = AgentBuilder::new(MockCompletionModel::from_turns([
10836 MockTurn::text("rejected").with_usage(first_usage),
10837 MockTurn::text("accepted").with_usage(second_usage),
10838 ]))
10839 .add_hook(BoundedResponseRetry::new(
10840 "rejected",
10841 1,
10842 TestRetryMode::Feedback("correct the answer"),
10843 ))
10844 .build()
10845 .runner("question")
10846 .max_turns(2)
10847 .run()
10848 .await
10849 .expect("blocking feedback retry");
10850
10851 let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([
10852 [
10853 MockStreamEvent::text("rejected"),
10854 MockStreamEvent::final_response(first_usage),
10855 ],
10856 [
10857 MockStreamEvent::text("accepted"),
10858 MockStreamEvent::final_response(second_usage),
10859 ],
10860 ]))
10861 .add_hook(BoundedResponseRetry::new(
10862 "rejected",
10863 1,
10864 TestRetryMode::Feedback("correct the answer"),
10865 ))
10866 .build()
10867 .runner("question")
10868 .max_turns(2)
10869 .stream()
10870 .await;
10871 let mut saw_retry = false;
10872 let mut streaming = None;
10873 while let Some(item) = stream.next().await {
10874 match item.expect("stream item") {
10875 MultiTurnStreamItem::ModelTurnRetried { turn: 1 } => saw_retry = true,
10876 MultiTurnStreamItem::FinalResponse(response) => streaming = Some(response),
10877 _ => {}
10878 }
10879 }
10880
10881 let streaming = streaming.expect("streaming final response");
10882 assert!(saw_retry);
10883 assert_eq!(streaming.output, blocking.output);
10884 assert_eq!(streaming.usage, blocking.usage);
10885 let without_raw = |calls: &[crate::agent::CompletionCall]| -> Vec<_> {
10889 calls
10890 .iter()
10891 .cloned()
10892 .map(|call| call.with_raw(serde_json::Value::Null))
10893 .collect()
10894 };
10895 assert_eq!(
10896 without_raw(&streaming.completion_calls),
10897 without_raw(&blocking.completion_calls)
10898 );
10899 assert_eq!(
10900 serde_json::to_value(streaming.messages).expect("streaming history"),
10901 serde_json::to_value(blocking.messages).expect("blocking history")
10902 );
10903 }
10904
10905 #[tokio::test]
10906 async fn streaming_empty_feedback_retry_omits_empty_assistant_history() {
10907 let first_usage = retry_usage(5, 1);
10908 let second_usage = retry_usage(7, 2);
10909 let model = MockCompletionModel::from_stream_turns([
10910 [
10911 MockStreamEvent::text(""),
10912 MockStreamEvent::final_response(first_usage),
10913 ],
10914 [
10915 MockStreamEvent::text("accepted"),
10916 MockStreamEvent::final_response(second_usage),
10917 ],
10918 ]);
10919 let mut stream = AgentBuilder::new(model.clone())
10920 .add_hook(BoundedResponseRetry::new(
10921 "",
10922 1,
10923 TestRetryMode::Feedback("provide an answer"),
10924 ))
10925 .build()
10926 .runner("question")
10927 .max_turns(2)
10928 .stream()
10929 .await;
10930
10931 let mut retries = Vec::new();
10932 let mut provider_finals = 0;
10933 let mut completion_calls = 0;
10934 let mut final_response = None;
10935 while let Some(item) = stream.next().await {
10936 match item.expect("stream item") {
10937 MultiTurnStreamItem::ModelTurnRetried { turn } => retries.push(turn),
10938 MultiTurnStreamItem::CompletionCall(_) => completion_calls += 1,
10939 MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Final(_)) => {
10940 provider_finals += 1;
10941 }
10942 MultiTurnStreamItem::FinalResponse(response) => final_response = Some(response),
10943 _ => {}
10944 }
10945 }
10946
10947 assert_eq!(retries, vec![1]);
10948 assert_eq!(provider_finals, 1, "the rejected final is suppressed");
10949 assert_eq!(completion_calls, 2);
10950 let response = final_response.expect("run final response");
10951 assert_eq!(response.output, "accepted");
10952 assert_eq!(response.usage, first_usage + second_usage);
10953 assert_eq!(response.completion_calls.len(), 2);
10954 assert_eq!(
10955 response.messages.expect("response messages"),
10956 vec![
10957 Message::user("question"),
10958 Message::user("provide an answer"),
10959 Message::assistant("accepted"),
10960 ]
10961 );
10962 assert_eq!(
10963 model.requests()[1].chat_history.clone(),
10964 vec![
10965 Message::user("question"),
10966 Message::user("provide an answer"),
10967 ],
10968 "the retry request must not contain an empty assistant message"
10969 );
10970 }
10971
10972 #[tokio::test]
10973 async fn response_retry_preserves_model_turn_hook_order_across_surfaces() {
10974 let blocking_events = RecordingHook::default();
10975 AgentBuilder::new(MockCompletionModel::from_turns([
10976 MockTurn::text("rejected"),
10977 MockTurn::text("accepted"),
10978 ]))
10979 .add_hook(blocking_events.clone())
10980 .add_hook(BoundedResponseRetry::new(
10981 "rejected",
10982 1,
10983 TestRetryMode::Repeat,
10984 ))
10985 .build()
10986 .runner("question")
10987 .max_turns(2)
10988 .run()
10989 .await
10990 .expect("blocking retry");
10991
10992 let streaming_events = RecordingHook::default();
10993 let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([
10994 [
10995 MockStreamEvent::text("rejected"),
10996 MockStreamEvent::final_response_with_default_usage(),
10997 ],
10998 [
10999 MockStreamEvent::text("accepted"),
11000 MockStreamEvent::final_response_with_default_usage(),
11001 ],
11002 ]))
11003 .add_hook(streaming_events.clone())
11004 .add_hook(BoundedResponseRetry::new(
11005 "rejected",
11006 1,
11007 TestRetryMode::Repeat,
11008 ))
11009 .build()
11010 .runner("question")
11011 .max_turns(2)
11012 .stream()
11013 .await;
11014 while let Some(item) = stream.next().await {
11015 item.expect("streaming retry item");
11016 }
11017
11018 let shared_order = |events: &RecordingHook| {
11019 events
11020 .events
11021 .lock()
11022 .expect("events")
11023 .iter()
11024 .copied()
11025 .filter(|event| {
11026 matches!(
11027 event,
11028 StepEventKind::CompletionCall | StepEventKind::ModelTurnFinished
11029 )
11030 })
11031 .collect::<Vec<_>>()
11032 };
11033 let expected = vec![
11034 StepEventKind::CompletionCall,
11035 StepEventKind::ModelTurnFinished,
11036 StepEventKind::CompletionCall,
11037 StepEventKind::ModelTurnFinished,
11038 ];
11039 assert_eq!(shared_order(&blocking_events), expected);
11040 assert_eq!(shared_order(&streaming_events), expected);
11041
11042 let blocking_order = blocking_events.events.lock().expect("events").clone();
11043 assert_eq!(
11044 blocking_order,
11045 vec![
11046 StepEventKind::CompletionCall,
11047 StepEventKind::CompletionResponse,
11048 StepEventKind::ModelTurnFinished,
11049 StepEventKind::CompletionCall,
11050 StepEventKind::CompletionResponse,
11051 StepEventKind::ModelTurnFinished,
11052 ]
11053 );
11054 let streaming_order = streaming_events.events.lock().expect("events").clone();
11055 assert_eq!(
11056 streaming_order,
11057 vec![
11058 StepEventKind::CompletionCall,
11059 StepEventKind::TextDelta,
11060 StepEventKind::StreamResponseFinish,
11061 StepEventKind::ModelTurnFinished,
11062 StepEventKind::CompletionCall,
11063 StepEventKind::TextDelta,
11064 StepEventKind::StreamResponseFinish,
11065 StepEventKind::ModelTurnFinished,
11066 ]
11067 );
11068 }
11069
11070 #[tokio::test]
11071 async fn streaming_model_turn_retry_respects_max_turns() {
11072 let model = MockCompletionModel::from_stream_turns([[
11073 MockStreamEvent::text("rejected"),
11074 MockStreamEvent::final_response_with_default_usage(),
11075 ]]);
11076 let mut stream = AgentBuilder::new(model)
11077 .add_hook(BoundedResponseRetry::new(
11078 "rejected",
11079 1,
11080 TestRetryMode::Repeat,
11081 ))
11082 .build()
11083 .runner("question")
11084 .max_turns(1)
11085 .stream()
11086 .await;
11087
11088 let mut saw_rollback = false;
11089 let mut error = None;
11090 while let Some(item) = stream.next().await {
11091 match item {
11092 Ok(MultiTurnStreamItem::ModelTurnRetried { turn: 1 }) => saw_rollback = true,
11093 Ok(_) => {}
11094 Err(err) => error = Some(err),
11095 }
11096 }
11097 assert!(saw_rollback);
11098 assert!(matches!(
11099 error,
11100 Some(StreamingError::Prompt(error))
11101 if matches!(error.as_ref(), PromptError::MaxTurnsError { max_turns: 1, .. })
11102 ));
11103 }
11104
11105 struct AlwaysRepeatModelTurn;
11106
11107 impl AgentHook for AlwaysRepeatModelTurn {
11108 async fn on_model_turn_finished(
11109 &self,
11110 _ctx: &HookContext,
11111 _event: ModelTurnFinished<'_>,
11112 ) -> ModelTurnAction {
11113 ModelTurnAction::repeat()
11114 }
11115 }
11116
11117 #[tokio::test]
11118 async fn model_turn_retry_rejects_tool_turn_before_tool_hooks_or_execution() {
11119 let recorder = RecordingHook::default();
11120 let executions = Arc::new(AtomicU32::new(0));
11121 let err = AgentBuilder::new(MockCompletionModel::from_turns([MockTurn::tool_call(
11122 "tc1",
11123 "add",
11124 json!({"x": 1, "y": 2}),
11125 )]))
11126 .tool(CountingAddTool {
11127 calls: executions.clone(),
11128 })
11129 .add_hook(recorder.clone())
11130 .add_hook(AlwaysRepeatModelTurn)
11131 .build()
11132 .runner("add")
11133 .max_turns(2)
11134 .run()
11135 .await
11136 .expect_err("tool-bearing retry must fail closed");
11137
11138 let PromptError::PromptCancelled {
11139 chat_history,
11140 reason,
11141 } = err
11142 else {
11143 panic!("tool-bearing retry should return PromptCancelled");
11144 };
11145 assert!(reason.contains("tool-bearing model turns"));
11146 assert!(reason.contains("tool-call hooks"));
11147 assert_eq!(chat_history, vec![Message::user("add")]);
11148 assert_eq!(recorder.count(StepEventKind::ToolCall), 0);
11149 assert_eq!(recorder.count(StepEventKind::ToolResult), 0);
11150 assert_eq!(executions.load(SeqCst), 0);
11151 }
11152
11153 #[tokio::test]
11154 async fn streaming_model_turn_retry_rejects_tool_turn_without_committed_execution() {
11155 let recorder = RecordingHook::default();
11156 let executions = Arc::new(AtomicU32::new(0));
11157 let mut stream = AgentBuilder::new(MockCompletionModel::from_stream_turns([[
11158 MockStreamEvent::tool_call_name_delta("tc1", "add"),
11159 MockStreamEvent::tool_call_arguments_delta("tc1", r#"{"x":1,"y":2}"#),
11160 MockStreamEvent::tool_call("tc1", "add", json!({"x": 1, "y": 2})),
11161 MockStreamEvent::final_response_with_default_usage(),
11162 ]]))
11163 .tool(CountingAddTool {
11164 calls: executions.clone(),
11165 })
11166 .add_hook(recorder.clone())
11167 .add_hook(AlwaysRepeatModelTurn)
11168 .build()
11169 .runner("add")
11170 .max_turns(2)
11171 .stream()
11172 .await;
11173
11174 let mut execution_commits = 0;
11175 let mut tool_results = 0;
11176 let mut provider_finals = 0;
11177 let mut agent_finals = 0;
11178 let mut retry_markers = 0;
11179 let mut error = None;
11180 while let Some(item) = stream.next().await {
11181 match item {
11182 Ok(MultiTurnStreamItem::ToolExecutionCommitted { .. }) => execution_commits += 1,
11183 Ok(MultiTurnStreamItem::StreamUserItem(StreamedUserContent::ToolResult {
11184 ..
11185 })) => tool_results += 1,
11186 Ok(MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Final(
11187 _,
11188 ))) => provider_finals += 1,
11189 Ok(MultiTurnStreamItem::FinalResponse(_)) => agent_finals += 1,
11190 Ok(MultiTurnStreamItem::ModelTurnRetried { .. }) => retry_markers += 1,
11191 Ok(_) => {}
11192 Err(err) => error = Some(err),
11193 }
11194 }
11195
11196 let Some(StreamingError::Prompt(error)) = error else {
11197 panic!("tool-bearing streaming retry should return PromptCancelled");
11198 };
11199 let PromptError::PromptCancelled {
11200 chat_history,
11201 reason,
11202 } = error.as_ref()
11203 else {
11204 panic!("tool-bearing streaming retry should return PromptCancelled");
11205 };
11206 assert!(reason.contains("tool-bearing model turns"));
11207 assert!(reason.contains("tool-call hooks"));
11208 assert_eq!(chat_history, &[Message::user("add")]);
11209 assert_eq!(execution_commits, 0);
11210 assert_eq!(tool_results, 0);
11211 assert_eq!(provider_finals, 0);
11212 assert_eq!(agent_finals, 0);
11213 assert_eq!(retry_markers, 0);
11214 assert_eq!(recorder.count(StepEventKind::ToolCall), 0);
11215 assert_eq!(recorder.count(StepEventKind::ToolResult), 0);
11216 assert_eq!(executions.load(SeqCst), 0);
11217 }
11218
11219 #[derive(Clone)]
11220 struct BarrierResponseRetry {
11221 inner: BoundedResponseRetry,
11222 barrier: Arc<Barrier>,
11223 }
11224
11225 impl AgentHook for BarrierResponseRetry {
11226 async fn on_model_turn_finished(
11227 &self,
11228 ctx: &HookContext,
11229 event: ModelTurnFinished<'_>,
11230 ) -> ModelTurnAction {
11231 let rejected = event.content.iter().any(
11232 |content| matches!(content, AssistantContent::Text(text) if text.text == "rejected"),
11233 );
11234 if rejected {
11235 self.barrier.wait().await;
11236 }
11237 self.inner.on_model_turn_finished(ctx, event).await
11238 }
11239 }
11240
11241 #[tokio::test]
11242 async fn concurrent_runs_of_same_agent_have_independent_retry_budgets() {
11243 let hook = BarrierResponseRetry {
11244 inner: BoundedResponseRetry::new("rejected", 1, TestRetryMode::Repeat),
11245 barrier: Arc::new(Barrier::new(2)),
11246 };
11247 let agent = AgentBuilder::new(MockCompletionModel::from_turns([
11248 MockTurn::text("rejected"),
11249 MockTurn::text("rejected"),
11250 MockTurn::text("accepted one"),
11251 MockTurn::text("accepted two"),
11252 ]))
11253 .add_hook(hook)
11254 .build();
11255
11256 let first = agent.runner("first").max_turns(2).run();
11257 let second = agent.runner("second").max_turns(2).run();
11258 let (first, second) = tokio::join!(first, second);
11259 let first = first.expect("first run");
11260 let second = second.expect("second run");
11261
11262 let outputs = std::collections::HashSet::from([first.output, second.output]);
11263 assert_eq!(
11264 outputs,
11265 std::collections::HashSet::from([
11266 "accepted one".to_string(),
11267 "accepted two".to_string(),
11268 ])
11269 );
11270 assert_eq!(first.completion_calls.len(), 2);
11271 assert_eq!(second.completion_calls.len(), 2);
11272 }
11273
11274 #[tokio::test]
11275 async fn retry_scratchpad_state_is_isolated_by_run_and_hook_instance() {
11276 let shared_hook = BoundedResponseRetry::new("rejected", 1, TestRetryMode::Repeat);
11277 let first_ctx = HookContext::new(false, None);
11278 let second_ctx = HookContext::new(false, None);
11279 let content = vec![AssistantContent::text("rejected")];
11280 let first_event = ModelTurnFinished {
11281 turn: 1,
11282 content: &content,
11283 usage: Usage::new(),
11284 identity: no_identity(),
11285 finish_reason: None,
11287 max_tokens: None,
11288 raw: &serde_json::Value::Null,
11289 };
11290 let second_event = first_event;
11291
11292 let (first, second) = tokio::join!(
11293 shared_hook.on_model_turn_finished(&first_ctx, first_event),
11294 shared_hook.on_model_turn_finished(&second_ctx, second_event),
11295 );
11296 assert!(matches!(first, ModelTurnAction::Retry(_)));
11297 assert!(matches!(second, ModelTurnAction::Retry(_)));
11298 assert!(matches!(
11299 shared_hook
11300 .on_model_turn_finished(&first_ctx, first_event)
11301 .await,
11302 ModelTurnAction::Stop(_)
11303 ));
11304
11305 let same_run_ctx = HookContext::new(false, None);
11306 let first_hook = BoundedResponseRetry::new("first", 1, TestRetryMode::Repeat);
11307 let second_hook = BoundedResponseRetry::new("second", 1, TestRetryMode::Repeat);
11308 let first_content = vec![AssistantContent::text("first")];
11309 let second_content = vec![AssistantContent::text("second")];
11310 let first_action = first_hook
11311 .on_model_turn_finished(
11312 &same_run_ctx,
11313 ModelTurnFinished {
11314 turn: 1,
11315 content: &first_content,
11316 usage: Usage::new(),
11317 identity: no_identity(),
11318 finish_reason: None,
11319 max_tokens: None,
11320 raw: &serde_json::Value::Null,
11321 },
11322 )
11323 .await;
11324 let second_action = second_hook
11325 .on_model_turn_finished(
11326 &same_run_ctx,
11327 ModelTurnFinished {
11328 turn: 2,
11329 content: &second_content,
11330 usage: Usage::new(),
11331 identity: no_identity(),
11332 finish_reason: None,
11333 max_tokens: None,
11334 raw: &serde_json::Value::Null,
11335 },
11336 )
11337 .await;
11338 assert!(matches!(first_action, ModelTurnAction::Retry(_)));
11339 assert!(matches!(second_action, ModelTurnAction::Retry(_)));
11340 }
11341
11342 #[derive(Clone)]
11343 struct FixedModelTurnAction {
11344 action: ModelTurnAction,
11345 calls: Arc<AtomicU32>,
11346 }
11347
11348 impl AgentHook for FixedModelTurnAction {
11349 async fn on_model_turn_finished(
11350 &self,
11351 _ctx: &HookContext,
11352 _event: ModelTurnFinished<'_>,
11353 ) -> ModelTurnAction {
11354 self.calls.fetch_add(1, SeqCst);
11355 self.action.clone()
11356 }
11357 }
11358
11359 #[tokio::test]
11360 async fn model_turn_action_short_circuits_flat_and_nested_hook_stacks() {
11361 let content = vec![AssistantContent::text("response")];
11362 let event = ModelTurnFinished {
11363 turn: 1,
11364 content: &content,
11365 usage: Usage::new(),
11366 identity: no_identity(),
11367 finish_reason: None,
11369 max_tokens: None,
11370 raw: &serde_json::Value::Null,
11371 };
11372 let ctx = HookContext::new(false, None);
11373
11374 let first_calls = Arc::new(AtomicU32::new(0));
11375 let retry_calls = Arc::new(AtomicU32::new(0));
11376 let skipped_calls = Arc::new(AtomicU32::new(0));
11377 let mut flat = HookStack::new();
11378 flat.push(FixedModelTurnAction {
11379 action: ModelTurnAction::Continue,
11380 calls: first_calls.clone(),
11381 });
11382 flat.push(FixedModelTurnAction {
11383 action: ModelTurnAction::repeat(),
11384 calls: retry_calls.clone(),
11385 });
11386 flat.push(FixedModelTurnAction {
11387 action: ModelTurnAction::stop("unreachable"),
11388 calls: skipped_calls.clone(),
11389 });
11390 assert!(matches!(
11391 flat.on_model_turn_finished(&ctx, event).await,
11392 ModelTurnAction::Retry(_)
11393 ));
11394 assert_eq!(first_calls.load(SeqCst), 1);
11395 assert_eq!(retry_calls.load(SeqCst), 1);
11396 assert_eq!(skipped_calls.load(SeqCst), 0);
11397
11398 let nested_retry_calls = Arc::new(AtomicU32::new(0));
11399 let outer_skipped_calls = Arc::new(AtomicU32::new(0));
11400 let mut nested = HookStack::new();
11401 nested.push(FixedModelTurnAction {
11402 action: ModelTurnAction::retry_with_feedback("fix it"),
11403 calls: nested_retry_calls.clone(),
11404 });
11405 let mut outer = HookStack::new();
11406 outer.push(nested);
11407 outer.push(FixedModelTurnAction {
11408 action: ModelTurnAction::Continue,
11409 calls: outer_skipped_calls.clone(),
11410 });
11411 assert!(matches!(
11412 outer.on_model_turn_finished(&ctx, event).await,
11413 ModelTurnAction::Retry(crate::agent::RetryRequest::Feedback(feedback))
11414 if feedback == "fix it"
11415 ));
11416 assert_eq!(nested_retry_calls.load(SeqCst), 1);
11417 assert_eq!(outer_skipped_calls.load(SeqCst), 0);
11418
11419 let stop_calls = Arc::new(AtomicU32::new(0));
11420 let after_stop_calls = Arc::new(AtomicU32::new(0));
11421 let mut stopping = HookStack::new();
11422 stopping.push(FixedModelTurnAction {
11423 action: ModelTurnAction::stop("stop now"),
11424 calls: stop_calls.clone(),
11425 });
11426 stopping.push(FixedModelTurnAction {
11427 action: ModelTurnAction::Continue,
11428 calls: after_stop_calls.clone(),
11429 });
11430 assert!(matches!(
11431 stopping.on_model_turn_finished(&ctx, event).await,
11432 ModelTurnAction::Stop(reason) if reason == "stop now"
11433 ));
11434 assert_eq!(stop_calls.load(SeqCst), 1);
11435 assert_eq!(after_stop_calls.load(SeqCst), 0);
11436 }
11437}