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