agent_framework_core/client.rs
1//! The chat client trait and the automatic function-invocation loop.
2//!
3//! Rust equivalent of `agent_framework._clients` plus the tool loop from
4//! `_tools.use_function_invocation`.
5
6use async_trait::async_trait;
7use futures::stream::{self, Stream, StreamExt};
8use serde_json::Value;
9use std::collections::HashMap;
10use std::pin::Pin;
11use std::sync::Arc;
12use std::time::{Duration, SystemTime, UNIX_EPOCH};
13use tracing::Instrument;
14
15use crate::error::{Error, Result};
16use crate::middleware::{FunctionInvocationContext, LiveToolList, MiddlewarePipeline, Terminal};
17use crate::tools::{FunctionInvocationConfig, ToolDefinition, ToolKind};
18use crate::types::{
19 ChatOptions, ChatResponse, ChatResponseUpdate, Content, EmbeddingGenerationOptions,
20 FunctionApprovalRequestContent, FunctionApprovalResponseContent, FunctionCallContent,
21 FunctionResultContent, GeneratedEmbeddings, Message, Role, ToolMode, UsageContent,
22 UsageDetails,
23};
24
25/// A boxed stream of streaming chat updates.
26pub type ChatStream = Pin<Box<dyn Stream<Item = Result<ChatResponseUpdate>> + Send>>;
27
28/// The interface every chat client implements.
29///
30/// Implementors provide [`ChatClient::get_response`] and
31/// [`ChatClient::get_streaming_response`]; the framework layers tool invocation
32/// and middleware on top via [`FunctionInvokingChatClient`].
33#[async_trait]
34pub trait ChatClient: Send + Sync {
35 /// Get a complete (non-streaming) response.
36 async fn get_response(
37 &self,
38 messages: Vec<Message>,
39 options: ChatOptions,
40 ) -> Result<ChatResponse>;
41
42 /// Get a streaming response as a sequence of updates.
43 async fn get_streaming_response(
44 &self,
45 messages: Vec<Message>,
46 options: ChatOptions,
47 ) -> Result<ChatStream>;
48
49 /// The default model id for this client, if any.
50 fn model(&self) -> Option<&str> {
51 None
52 }
53}
54
55/// Blanket impl so `Arc<dyn ChatClient>` and wrappers are usable as clients.
56#[async_trait]
57impl<T: ChatClient + ?Sized> ChatClient for Arc<T> {
58 async fn get_response(
59 &self,
60 messages: Vec<Message>,
61 options: ChatOptions,
62 ) -> Result<ChatResponse> {
63 (**self).get_response(messages, options).await
64 }
65 async fn get_streaming_response(
66 &self,
67 messages: Vec<Message>,
68 options: ChatOptions,
69 ) -> Result<ChatStream> {
70 (**self).get_streaming_response(messages, options).await
71 }
72 fn model(&self) -> Option<&str> {
73 (**self).model()
74 }
75}
76
77/// The interface every embedding client implements.
78///
79/// Rust equivalent of upstream's `SupportsGetEmbeddings` protocol /
80/// `BaseEmbeddingClient` (`_clients.py`): generate one embedding per input
81/// string, batched in a single request. Vectors are `Vec<f32>` — see the
82/// note on [`crate::types::Embedding`] about upstream's genericity.
83#[async_trait]
84pub trait EmbeddingClient: Send + Sync {
85 /// Generate embeddings for the given values (one per value, in order).
86 async fn get_embeddings(
87 &self,
88 values: Vec<String>,
89 options: Option<EmbeddingGenerationOptions>,
90 ) -> Result<GeneratedEmbeddings>;
91
92 /// The default embedding model id for this client, if any.
93 fn model(&self) -> Option<&str> {
94 None
95 }
96}
97
98/// Blanket impl so `Arc<dyn EmbeddingClient>` and wrappers are usable as
99/// clients.
100#[async_trait]
101impl<T: EmbeddingClient + ?Sized> EmbeddingClient for Arc<T> {
102 async fn get_embeddings(
103 &self,
104 values: Vec<String>,
105 options: Option<EmbeddingGenerationOptions>,
106 ) -> Result<GeneratedEmbeddings> {
107 (**self).get_embeddings(values, options).await
108 }
109 fn model(&self) -> Option<&str> {
110 (**self).model()
111 }
112}
113
114/// Wraps a [`ChatClient`] to automatically execute local tool calls in a loop,
115/// mirroring `use_function_invocation`.
116pub struct FunctionInvokingChatClient<C: ChatClient> {
117 inner: C,
118 config: FunctionInvocationConfig,
119 /// Governs the `execute_tool` spans this client emits: content capture and
120 /// the GenAI semantic-convention version. Resolved from the environment
121 /// once at construction rather than re-read per tool call, so a caller who
122 /// selects a version explicitly (via [`Self::with_observability_config`])
123 /// gets that version on tool spans too, instead of a trace that mixes
124 /// conventions between its chat and tool spans.
125 observability: crate::observability::ObservabilityConfig,
126 /// Middleware run around every individual tool call (mirrors Python's
127 /// function-middleware pipeline, driven here instead of by a
128 /// `use_function_invocation` decorator).
129 function_middleware: MiddlewarePipeline<FunctionInvocationContext>,
130}
131
132impl<C: ChatClient> FunctionInvokingChatClient<C> {
133 pub fn new(inner: C) -> Self {
134 Self {
135 inner,
136 config: FunctionInvocationConfig::default(),
137 function_middleware: MiddlewarePipeline::default(),
138 observability: crate::observability::ObservabilityConfig::from_env(),
139 }
140 }
141
142 /// Set the [`ObservabilityConfig`](crate::observability::ObservabilityConfig)
143 /// governing this client's `execute_tool` spans. Pass the same config given
144 /// to an [`ObservableChatClient`](crate::observability::ObservableChatClient)
145 /// wrapping the same stack, so one trace reports one semantic-convention
146 /// version throughout.
147 pub fn with_observability_config(
148 mut self,
149 observability: crate::observability::ObservabilityConfig,
150 ) -> Self {
151 self.observability = observability;
152 self
153 }
154
155 /// Override the function-invocation configuration.
156 pub fn with_config(mut self, config: FunctionInvocationConfig) -> Self {
157 self.config = config;
158 self
159 }
160
161 /// Configure the function-invocation middleware pipeline run around every
162 /// tool call: middleware may inspect/rewrite
163 /// [`FunctionInvocationContext::arguments`], short-circuit execution by
164 /// setting [`FunctionInvocationContext::result`] (and either not calling
165 /// `next`, or setting `terminate = true`), or observe a propagated
166 /// execution error by matching on the `Result` returned from their own
167 /// `next.run(...)` call. Replaces any previously configured middleware.
168 pub fn with_function_middleware(
169 mut self,
170 middleware: Vec<Arc<crate::middleware::FunctionMiddleware>>,
171 ) -> Self {
172 self.function_middleware = MiddlewarePipeline::new(middleware);
173 self
174 }
175
176 /// A reference to the wrapped client.
177 pub fn inner(&self) -> &C {
178 &self.inner
179 }
180
181 async fn inner_get_response(
182 &self,
183 messages: Vec<Message>,
184 options: ChatOptions,
185 ) -> Result<ChatResponse> {
186 self.inner.get_response(messages, options).await
187 }
188}
189
190/// Extract the executable tools from the options into a name→tool map.
191fn executable_tools(options: &ChatOptions) -> Vec<ToolDefinition> {
192 options
193 .tools
194 .iter()
195 .filter(|t| t.is_executable())
196 .cloned()
197 .collect()
198}
199
200/// Whether `tool` is a *declaration-only* function tool: a known function with
201/// no local executor. Mirrors Python's `AIFunction.declaration_only`. A call to
202/// such a tool is returned to the caller unexecuted (the frontend-tool pattern
203/// that makes AG-UI client-side tools work). Hosted tools (web search, MCP, …)
204/// are deliberately excluded — they are not function tools and a call whose
205/// name matches none of the local function tools is treated as unknown, not
206/// declaration-only, exactly as Python's `_get_tool_map` omits them.
207fn is_declaration_only(tool: &ToolDefinition) -> bool {
208 tool.kind == ToolKind::Function && tool.executor.is_none()
209}
210
211/// Fold one model call's usage into a running aggregate.
212///
213/// The tool loop issues *several* model calls per logical `get_response`, and
214/// each reports only its own tokens. Without accumulation the returned response
215/// carries the last iteration's usage alone, so a run that called tools five
216/// times under-reports its cost by roughly a factor of five — and because this
217/// port's OTel layer reads `usage_details`, the `gen_ai.usage.*` metrics
218/// under-report with it. Mirrors upstream's `UsageAggregator` (.NET #7539).
219///
220/// `None` means *not reported* rather than zero: an aggregate only carries a
221/// count once some contributor reported one, and stays `None` when no iteration
222/// reported usage at all.
223fn accumulate_usage(aggregate: &mut Option<UsageDetails>, incoming: Option<&UsageDetails>) {
224 let Some(incoming) = incoming else {
225 return;
226 };
227 match aggregate {
228 Some(current) => current.add_assign(incoming),
229 None => *aggregate = Some(incoming.clone()),
230 }
231}
232
233/// The exact rejection payload Python emits for a denied tool call.
234const REJECTION_MESSAGE: &str = "Error: Tool call invocation was rejected by user.";
235
236/// Execute a single requested tool call through the function-middleware
237/// pipeline, with the actual invocation (wrapped in an `execute_tool` span)
238/// as the pipeline's terminal handler.
239///
240/// Returns `(is_error, result)`. `terminate_on_unknown` turns an unknown-tool
241/// call into a hard error (propagated) rather than an error result. Unknown
242/// tools and unparseable arguments are rejected before middleware ever sees
243/// them (there is no function to hand the pipeline in that case); once a
244/// [`FunctionInvocationContext`] is built, middleware can rewrite
245/// `arguments`, short-circuit by setting `result` (without calling `next`, or
246/// with `terminate = true`), or observe an execution error by matching on the
247/// `Result` their own `next.run(...)` call returns. A propagated error is
248/// converted to the same `(true, FunctionResultContent { exception: .. })`
249/// shape the direct-error path used before middleware existed, so
250/// `include_detailed_errors` behaves identically either way.
251///
252/// The single exception is [`Error::MiddlewareFailure`], which is propagated
253/// rather than absorbed: it is the fail-closed signal an enforcement layer
254/// returns to stop the run instead of letting the model retry around a
255/// tool-error result. Because the parallel batch is driven by `try_join_all`,
256/// propagating it also drops (cancels) the sibling calls still in flight.
257/// The ambient state one tool call runs against: everything that comes from
258/// the client and the run rather than from the call itself.
259struct ToolCallEnv<'a> {
260 function_middleware: &'a MiddlewarePipeline<FunctionInvocationContext>,
261 session: Option<&'a crate::session::AgentSession>,
262 live_tools: Option<&'a LiveToolList>,
263 observability: &'a crate::observability::ObservabilityConfig,
264}
265
266async fn execute_tool_call(
267 tool: Option<ToolDefinition>,
268 call: &FunctionCallContent,
269 include_detailed_errors: bool,
270 terminate_on_unknown: bool,
271 env: ToolCallEnv<'_>,
272) -> Result<(bool, FunctionResultContent)> {
273 let ToolCallEnv {
274 function_middleware,
275 session,
276 live_tools,
277 observability,
278 } = env;
279 match tool {
280 None => {
281 if terminate_on_unknown {
282 return Err(Error::tool(format!("unknown tool: {}", call.name)));
283 }
284 Ok((
285 true,
286 FunctionResultContent {
287 call_id: call.call_id.clone(),
288 result: None,
289 exception: Some(format!("tool '{}' not found", call.name)),
290 },
291 ))
292 }
293 Some(def) => {
294 // Reject unparseable arguments rather than silently invoking the tool
295 // with null/default input.
296 let args = match call.parse_arguments() {
297 Ok(m) => Value::Object(m.into_iter().collect()),
298 Err(e) => {
299 let msg = if include_detailed_errors {
300 format!("invalid tool arguments: {e}")
301 } else {
302 "invalid tool arguments".to_string()
303 };
304 return Ok((
305 true,
306 FunctionResultContent {
307 call_id: call.call_id.clone(),
308 result: None,
309 exception: Some(msg),
310 },
311 ));
312 }
313 };
314 let obs_config = observability.clone();
315 let exec = def.executor.as_ref().unwrap().clone();
316 let tool_name = def.name.clone();
317 let description = def.description.clone();
318 let call_id = call.call_id.clone();
319 let terminal: Terminal<FunctionInvocationContext> = Box::new(move |mut ctx| {
320 Box::pin(async move {
321 if ctx.terminate {
322 return Ok(ctx);
323 }
324 let span = crate::observability::tool_span_ex(
325 &tool_name,
326 &call_id,
327 Some(&description),
328 );
329 crate::observability::record_tool_arguments(&span, &ctx.arguments, &obs_config);
330 #[cfg(feature = "otel-metrics")]
331 let started = std::time::Instant::now();
332 let outcome = async {
333 let result = exec.invoke_in_context(ctx.arguments.clone(), &ctx).await;
334 if let Err(e) = &result {
335 crate::observability::record_error(&tracing::Span::current(), e);
336 }
337 result
338 }
339 .instrument(span.clone())
340 .await;
341 #[cfg(feature = "otel-metrics")]
342 crate::observability::metrics::record_function_invocation_duration(
343 &tool_name,
344 started.elapsed(),
345 outcome
346 .as_ref()
347 .err()
348 .map(crate::observability::error_type)
349 .as_deref(),
350 );
351 if let Ok(value) = &outcome {
352 crate::observability::record_tool_result(&span, value, &obs_config);
353 }
354 ctx.result = Some(outcome?);
355 Ok(ctx)
356 }) as crate::tools::BoxFuture<Result<FunctionInvocationContext>>
357 });
358
359 let ctx = FunctionInvocationContext::new(call.name.clone(), args)
360 .with_session(session.cloned())
361 .with_tools(live_tools.cloned());
362 match function_middleware.execute(ctx, terminal).await {
363 Ok(ctx) => Ok((
364 false,
365 FunctionResultContent {
366 call_id: call.call_id.clone(),
367 result: Some(ctx.result.unwrap_or(Value::Null)),
368 exception: None,
369 },
370 )),
371 // The one error the loop does not absorb: middleware that
372 // refuses a call outright (a guardrail, a policy or
373 // authorization gate) needs the run to fail closed rather than
374 // hand the model an error string it can retry around. Every
375 // other error keeps the absorb-and-continue contract below.
376 Err(e) if e.is_middleware_failure() => Err(e),
377 Err(e) => {
378 let msg = if include_detailed_errors {
379 format!("{e}")
380 } else {
381 "tool execution failed".to_string()
382 };
383 Ok((
384 true,
385 FunctionResultContent {
386 call_id: call.call_id.clone(),
387 result: None,
388 exception: Some(msg),
389 },
390 ))
391 }
392 }
393 }
394 }
395}
396
397/// Collect all function-approval responses present in a conversation.
398fn collect_approval_responses(messages: &[Message]) -> Vec<FunctionApprovalResponseContent> {
399 let mut out = Vec::new();
400 for msg in messages {
401 for content in &msg.contents {
402 if let Content::FunctionApprovalResponse(resp) = content {
403 out.push(resp.clone());
404 }
405 }
406 }
407 out
408}
409
410/// Rewrite approval request/response contents in place, mirroring Python's
411/// `_replace_approval_contents_with_results`.
412///
413/// * A [`FunctionApprovalRequestContent`] becomes its embedded
414/// [`FunctionCallContent`], unless an equal call is *outstanding at that
415/// point in the conversation* (a replayed duplicate), in which case the
416/// request is removed instead.
417/// * An approved [`FunctionApprovalResponseContent`] becomes the corresponding
418/// result (correlated strictly by call id) and the message role becomes
419/// `tool`.
420/// * A rejected response becomes a [`FunctionResultContent`] carrying the
421/// rejection payload, and the message role becomes `tool`.
422fn replace_approval_contents_with_results(
423 messages: &mut [Message],
424 approved_results: &HashMap<String, FunctionResultContent>,
425) {
426 /// A call currently awaiting its result. `from_request` marks entries that
427 /// exist because an approval request expanded, so a *real* copy of the same
428 /// call arriving later is recognized as the duplicate instead.
429 struct Outstanding {
430 call: FunctionCallContent,
431 from_request: bool,
432 }
433
434 // One ordered walk. The outstanding-call set is *derived* as the walk
435 // decides each content — a call becomes outstanding when its (kept)
436 // declaration passes by, and stops being outstanding when the content that
437 // answers it passes by. Three earlier revisions instead maintained this
438 // set as separate bookkeeping around a whole-list pre-scan, and every one
439 // was wrong the same way: some site updated the mirror on an event that
440 // does not change what is outstanding, or missed one that does. A
441 // pre-scan is also order-blind — it nets a call against a result that
442 // only arrives *after* the replayed request, expanding the request into a
443 // second declaration — so deriving in order fixes a real timing bug, not
444 // just the structure.
445 let mut outstanding: Vec<Outstanding> = Vec::new();
446 // Nearest-preceding outstanding call sharing the id — the same pairing
447 // rule the compaction module settled on.
448 fn retire(outstanding: &mut Vec<Outstanding>, call_id: &str) {
449 if call_id.is_empty() {
450 return;
451 }
452 if let Some(position) = outstanding
453 .iter()
454 .rposition(|entry| entry.call.call_id == call_id)
455 {
456 outstanding.remove(position);
457 }
458 }
459
460 for msg in messages.iter_mut() {
461 let mut to_remove: Vec<usize> = Vec::new();
462 let mut set_role_tool = false;
463
464 for (idx, content) in msg.contents.iter_mut().enumerate() {
465 match content {
466 Content::FunctionCall(fc) => {
467 if fc.call_id.is_empty() {
468 continue;
469 }
470 // A real call matching an expanded request is the duplicate
471 // now: the expansion already declared it. Keep exactly one,
472 // and mark the survivor as real so it is not itself treated
473 // as expendable.
474 if let Some(position) = outstanding
475 .iter()
476 .position(|entry| entry.from_request && entry.call == *fc)
477 {
478 outstanding[position].from_request = false;
479 to_remove.push(idx);
480 } else {
481 outstanding.push(Outstanding {
482 call: fc.clone(),
483 from_request: false,
484 });
485 }
486 }
487 Content::FunctionResult(fr) => {
488 retire(&mut outstanding, &fr.call_id);
489 }
490 Content::FunctionApprovalRequest(req) => {
491 // Suppressed only for the *same invocation* — ids are
492 // reused, so a request differing in name or arguments is a
493 // fresh call, not a replay. Nothing is consumed on a match:
494 // dropping a request answers no call.
495 if outstanding
496 .iter()
497 .any(|entry| entry.call == req.function_call)
498 {
499 to_remove.push(idx);
500 } else {
501 if !req.function_call.call_id.is_empty() {
502 outstanding.push(Outstanding {
503 call: req.function_call.clone(),
504 from_request: true,
505 });
506 }
507 *content = Content::FunctionCall(req.function_call.clone());
508 }
509 }
510 Content::FunctionApprovalResponse(resp) => {
511 let call_id = resp.function_call.call_id.clone();
512 let mut answered = false;
513 if resp.approved {
514 if let Some(result) = approved_results.get(&call_id) {
515 *content = Content::FunctionResult(result.clone());
516 answered = true;
517 }
518 } else {
519 *content = Content::FunctionResult(FunctionResultContent {
520 call_id: call_id.clone(),
521 result: Some(Value::String(REJECTION_MESSAGE.to_string())),
522 exception: None,
523 });
524 answered = true;
525 }
526 // A response that became a result answers its call, exactly
527 // as a literal result would.
528 if answered {
529 retire(&mut outstanding, &call_id);
530 set_role_tool = true;
531 }
532 }
533 _ => {}
534 }
535 }
536
537 for idx in to_remove.into_iter().rev() {
538 msg.contents.remove(idx);
539 }
540 if set_role_tool {
541 msg.role = Role::tool();
542 }
543 }
544}
545
546#[async_trait]
547impl<C: ChatClient> ChatClient for FunctionInvokingChatClient<C> {
548 async fn get_response(
549 &self,
550 messages: Vec<Message>,
551 mut options: ChatOptions,
552 ) -> Result<ChatResponse> {
553 // After the tool loop settles, auto-populate `ChatResponse.value` from
554 // the final text when a structured `response_format` was requested
555 // (mirrors Python `try_parse_value`). This is the central non-streaming
556 // fill point: it covers a bare `FunctionInvokingChatClient` and every
557 // `Agent` run (whose client is always wrapped in one). The tool
558 // loop is run inside an `async move` block so its interior `return`s
559 // funnel through this single fill/return path.
560 let response_format = options.response_format.clone();
561 let mut response: ChatResponse = async move {
562 self.config.validate()?;
563 // Pop the agent-session side channel before the inner provider
564 // client ever sees the options (mirrors upstream's
565 // `effective_client_kwargs.pop("session")`); it is handed to
566 // invoked tools via `FunctionInvocationContext::session`.
567 let session = options.session.take();
568
569 // Default tool choice to auto when tools are present and unset.
570 if !options.tools.is_empty() && options.tool_choice.is_none() {
571 options.tool_choice = Some(ToolMode::Auto);
572 }
573
574 if executable_tools(&options).is_empty() || !self.config.enabled {
575 return self.inner_get_response(messages, options).await;
576 }
577
578 // The run's live tool list (progressive tool exposure): handed to
579 // every invocation via `FunctionInvocationContext::tools`, and
580 // re-snapshotted into the wire options at the top of every model
581 // iteration — so `add_tools`/`remove_tools` from middleware or
582 // tools take effect on the NEXT iteration, never the in-flight
583 // batch (mirrors upstream `_middleware.py` add_tools/remove_tools
584 // semantics).
585 let live_tools = LiveToolList::new(std::mem::take(&mut options.tools));
586
587 let mut conversation = messages;
588 let mut carried: Vec<Message> = Vec::new();
589 let mut consecutive_errors = 0usize;
590 // Usage summed over every model call this loop makes, applied to
591 // whichever response is returned so the caller sees the cost of the
592 // whole run and not just its final iteration.
593 let mut aggregated_usage: Option<UsageDetails> = None;
594
595 for _ in 0..self.config.max_iterations {
596 options.tools = live_tools.snapshot();
597 let tools = executable_tools(&options);
598 // Process any function-approval responses supplied in the input:
599 // execute the approved calls and splice their results into the
600 // conversation (mirrors Python's `_collect_approval_responses` +
601 // `_replace_approval_contents_with_results`).
602 let approval_responses = collect_approval_responses(&conversation);
603 if !approval_responses.is_empty() {
604 let mut approved_results: HashMap<String, FunctionResultContent> =
605 HashMap::new();
606 let mut had_error = false;
607 for resp in &approval_responses {
608 if !resp.approved {
609 continue;
610 }
611 let call = &resp.function_call;
612 let tool = tools.iter().find(|t| t.name == call.name).cloned();
613 let (is_error, content) = execute_tool_call(
614 tool,
615 call,
616 self.config.include_detailed_errors,
617 self.config.terminate_on_unknown_calls,
618 ToolCallEnv {
619 function_middleware: &self.function_middleware,
620 session: session.as_ref(),
621 live_tools: Some(&live_tools),
622 observability: &self.observability,
623 },
624 )
625 .await?;
626 had_error |= is_error;
627 approved_results.insert(content.call_id.clone(), content);
628 }
629 replace_approval_contents_with_results(&mut conversation, &approved_results);
630 if had_error {
631 consecutive_errors += 1;
632 if consecutive_errors > self.config.max_consecutive_errors_per_request {
633 options.tool_choice = Some(ToolMode::None);
634 }
635 }
636 }
637
638 let response = self
639 .inner_get_response(conversation.clone(), options.clone())
640 .await?;
641 accumulate_usage(&mut aggregated_usage, response.usage_details.as_ref());
642
643 // A call whose result is already present in the same response
644 // was executed by the provider (e.g. Anthropic server-side
645 // web-search/code-execution/MCP `server_tool_use` blocks,
646 // which arrive paired with their `*_tool_result`). Executing
647 // it locally would produce a bogus "tool not found" — only
648 // unresolved calls enter the local tool loop.
649 let resolved_call_ids: std::collections::HashSet<&str> = response
650 .messages
651 .iter()
652 .flat_map(|m| m.contents.iter())
653 .filter_map(Content::as_function_result)
654 .map(|fr| fr.call_id.as_str())
655 .collect();
656 let calls: Vec<_> = response
657 .messages
658 .iter()
659 .flat_map(|m| m.contents.iter())
660 .filter_map(Content::as_function_call)
661 .filter(|fc| !resolved_call_ids.contains(fc.call_id.as_str()))
662 .cloned()
663 .collect();
664
665 if calls.is_empty() {
666 // Prepend the accumulated tool-interaction messages so the final
667 // assistant message stays last.
668 let mut final_resp = response;
669 let mut msgs = std::mem::take(&mut carried);
670 msgs.append(&mut final_resp.messages);
671 final_resp.messages = msgs;
672 final_resp.usage_details = aggregated_usage;
673 return Ok(final_resp);
674 }
675
676 // Human-in-the-loop gate: if *any* requested tool requires approval,
677 // defer *all* calls (matching Python) and return an assistant message
678 // that carries the original calls plus one approval request each.
679 let needs_approval = calls.iter().any(|c| {
680 tools
681 .iter()
682 .find(|t| t.name == c.name)
683 .map(ToolDefinition::requires_approval)
684 .unwrap_or(false)
685 });
686 if needs_approval {
687 let mut resp = response;
688 let approval_contents: Vec<Content> = calls
689 .iter()
690 .map(|c| {
691 Content::FunctionApprovalRequest(FunctionApprovalRequestContent {
692 id: c.call_id.clone(),
693 function_call: c.clone(),
694 })
695 })
696 .collect();
697 if let Some(m) = resp
698 .messages
699 .iter_mut()
700 .rev()
701 .find(|m| m.role == Role::assistant())
702 {
703 m.contents.extend(approval_contents);
704 } else {
705 resp.messages
706 .push(Message::with_contents(Role::assistant(), approval_contents));
707 }
708 let mut msgs = std::mem::take(&mut carried);
709 msgs.append(&mut resp.messages);
710 resp.messages = msgs;
711 resp.usage_details = aggregated_usage;
712 return Ok(resp);
713 }
714
715 // Declaration-only calls: a call targeting a KNOWN tool that has
716 // no local executor (declaration-only — e.g. an AG-UI frontend
717 // tool, or a per-run `additional_tools` entry) terminates the
718 // loop and returns the response with the `FunctionCallContent`
719 // intact, so the caller can execute it. Mirrors Python's
720 // `_try_execute_function_calls` `declaration_only` branch
721 // (`_tools.py:1396-1420`): if *any* requested call is
722 // declaration-only, the whole response is returned unexecuted.
723 // A genuinely unknown tool name is NOT declaration-only and
724 // keeps today's not-found handling in `execute_tool_call`.
725 let has_declaration_only = calls.iter().any(|c| {
726 options
727 .tools
728 .iter()
729 .any(|t| t.name == c.name && is_declaration_only(t))
730 });
731 if has_declaration_only {
732 let mut resp = response;
733 let mut msgs = std::mem::take(&mut carried);
734 msgs.append(&mut resp.messages);
735 resp.messages = msgs;
736 resp.usage_details = aggregated_usage;
737 return Ok(resp);
738 }
739
740 // Record the assistant message(s) that requested the calls.
741 carried.extend(response.messages.iter().cloned());
742 let response_conversation_id = response.conversation_id.clone();
743
744 // Execute all calls concurrently: the model may emit several
745 // parallel tool calls, and I/O-bound tools should not be serialized.
746 let invocations = calls.iter().map(|call| {
747 let tool = tools.iter().find(|t| t.name == call.name).cloned();
748 let call = call.clone();
749 let include_detailed_errors = self.config.include_detailed_errors;
750 let terminate_on_unknown = self.config.terminate_on_unknown_calls;
751 let function_middleware = self.function_middleware.clone();
752 let session = session.clone();
753 let live_tools = live_tools.clone();
754 let observability = self.observability.clone();
755 async move {
756 execute_tool_call(
757 tool,
758 &call,
759 include_detailed_errors,
760 terminate_on_unknown,
761 ToolCallEnv {
762 function_middleware: &function_middleware,
763 session: session.as_ref(),
764 live_tools: Some(&live_tools),
765 observability: &observability,
766 },
767 )
768 .await
769 }
770 });
771
772 let outcomes = futures::future::try_join_all(invocations).await?;
773 let mut result_contents: Vec<Content> = Vec::with_capacity(outcomes.len());
774 let mut had_error = false;
775 for (is_error, content) in outcomes {
776 had_error |= is_error;
777 result_contents.push(Content::FunctionResult(content));
778 }
779
780 if had_error {
781 consecutive_errors += 1;
782 if consecutive_errors > self.config.max_consecutive_errors_per_request {
783 // Give up on tools and let the model answer directly.
784 options.tool_choice = Some(ToolMode::None);
785 }
786 } else {
787 consecutive_errors = 0;
788 }
789
790 let tool_message = Message::with_contents(Role::tool(), result_contents);
791 carried.push(tool_message.clone());
792 match response_conversation_id {
793 // A service-managed client that created (or continued) the
794 // conversation now holds the history server-side. Propagate
795 // its id so the follow-up tool-output submission targets the
796 // right thread — without this, Assistants / Azure AI reject
797 // the submission because `conversation_id` is still `None` —
798 // and send ONLY the new tool results next turn rather than
799 // re-sending the whole history (mirrors Python
800 // `_tools.py:1635-1637, 1695-1699`).
801 Some(cid) => {
802 options.conversation_id = Some(cid);
803 conversation = vec![tool_message];
804 }
805 // Stateless client (e.g. Chat Completions): accumulate and
806 // re-send the full history each turn.
807 None => {
808 conversation.extend(response.messages);
809 conversation.push(tool_message);
810 }
811 }
812 }
813
814 // Failsafe: one final call with tools disabled.
815 options.tool_choice = Some(ToolMode::None);
816 let mut final_resp = self.inner_get_response(conversation, options).await?;
817 accumulate_usage(&mut aggregated_usage, final_resp.usage_details.as_ref());
818 let mut msgs = std::mem::take(&mut carried);
819 msgs.append(&mut final_resp.messages);
820 final_resp.messages = msgs;
821 final_resp.usage_details = aggregated_usage;
822 Ok(final_resp)
823 }
824 .await?;
825 response.try_parse_value(response_format.as_ref());
826 Ok(response)
827 }
828
829 async fn get_streaming_response(
830 &self,
831 messages: Vec<Message>,
832 options: ChatOptions,
833 ) -> Result<ChatStream> {
834 let tools = executable_tools(&options);
835 if tools.is_empty() || !self.config.enabled {
836 return self.inner.get_streaming_response(messages, options).await;
837 }
838 // With tools, run the full loop then stream the aggregated result.
839 // Each message is replayed as its own update with a stable, distinct
840 // `message_id` so that consumers re-aggregating via
841 // `ChatResponse::from_updates` keep the messages separate rather than
842 // merging the tool-call and final assistant messages by role.
843 let response = self.get_response(messages, options).await?;
844 // Response-level metadata must survive the replay so re-aggregation
845 // (and the agent's thread adoption) sees it: ids on every update,
846 // and usage/finish-reason on the final one (usage rides as a
847 // `Content::Usage` item, which `absorb_update` folds into
848 // `usage_details` rather than the message contents — the same shape
849 // providers use for their terminal stream chunk).
850 let conversation_id = response.conversation_id.clone();
851 let response_id = response.response_id.clone();
852 let finish_reason = response.finish_reason.clone();
853 let usage_details = response.usage_details.clone();
854 let last = response.messages.len().saturating_sub(1);
855 // Keep the provider message ids only when they're all present and
856 // distinct; otherwise use positional ids for every message. A service
857 // (e.g. Assistants) can reuse one run id for both the tool-call turn
858 // and the final assistant turn, and `ChatResponse::from_updates` keys
859 // messages by id — a duplicate would merge the final answer into the
860 // tool-call message, ahead of the tool result.
861 let keep_provider_ids = {
862 let mut seen = std::collections::HashSet::new();
863 response.messages.iter().all(|m| {
864 m.message_id
865 .as_ref()
866 .is_some_and(|id| !id.is_empty() && seen.insert(id.as_str()))
867 })
868 };
869 let mut updates: Vec<Result<ChatResponseUpdate>> = response
870 .messages
871 .into_iter()
872 .enumerate()
873 .map(|(i, m)| {
874 let message_id = if keep_provider_ids {
875 m.message_id.clone()
876 } else {
877 Some(format!("replay-{i}"))
878 };
879 let mut contents = m.contents;
880 let is_last = i == last;
881 if is_last {
882 if let Some(usage) = usage_details.clone() {
883 contents.push(Content::Usage(UsageContent { details: usage }));
884 }
885 }
886 Ok(ChatResponseUpdate {
887 contents,
888 role: Some(m.role),
889 author_name: m.author_name,
890 message_id,
891 conversation_id: conversation_id.clone(),
892 response_id: response_id.clone(),
893 finish_reason: is_last.then(|| finish_reason.clone()).flatten(),
894 ..Default::default()
895 })
896 })
897 .collect();
898 // A messageless response (unusual, but possible) still carries its
899 // terminal metadata in one trailing update.
900 if updates.is_empty() && (usage_details.is_some() || finish_reason.is_some()) {
901 let contents = usage_details
902 .map(|u| vec![Content::Usage(UsageContent { details: u })])
903 .unwrap_or_default();
904 updates.push(Ok(ChatResponseUpdate {
905 contents,
906 role: Some(Role::assistant()),
907 conversation_id,
908 response_id,
909 finish_reason,
910 ..Default::default()
911 }));
912 }
913 Ok(stream::iter(updates).boxed())
914 }
915
916 fn model(&self) -> Option<&str> {
917 self.inner.model()
918 }
919}
920
921// ---------------------------------------------------------------------------
922// Retry / backoff layer
923// ---------------------------------------------------------------------------
924
925/// Which errors a [`RetryPolicy`] considers retryable.
926#[derive(Clone)]
927pub enum RetryOn {
928 /// The built-in default predicate (see [`RetryPolicy`] docs for the exact
929 /// rule): retries HTTP `408`/`429`/`5xx` ([`Error::ServiceStatus`]) and
930 /// transport-ish [`Error::Service`] failures (timeouts / connection
931 /// errors). Never retries [`Error::ServiceInvalidAuth`],
932 /// [`Error::ServiceInvalidRequest`], or [`Error::ServiceContentFilter`] —
933 /// authentication/authorization failures, malformed requests, and
934 /// content-filter refusals are non-transient, so retrying would just
935 /// repeat the same rejection.
936 Default,
937 /// A fully custom predicate deciding, per error, whether to retry.
938 Predicate(Arc<dyn Fn(&Error) -> bool + Send + Sync>),
939}
940
941impl std::fmt::Debug for RetryOn {
942 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
943 match self {
944 RetryOn::Default => f.write_str("RetryOn::Default"),
945 RetryOn::Predicate(_) => f.write_str("RetryOn::Predicate(..)"),
946 }
947 }
948}
949
950impl RetryOn {
951 /// A custom retry predicate.
952 pub fn predicate<F>(f: F) -> Self
953 where
954 F: Fn(&Error) -> bool + Send + Sync + 'static,
955 {
956 RetryOn::Predicate(Arc::new(f))
957 }
958
959 fn should_retry(&self, err: &Error) -> bool {
960 match self {
961 RetryOn::Default => default_should_retry(err),
962 RetryOn::Predicate(p) => p(err),
963 }
964 }
965}
966
967/// The default retryability rule used by [`RetryOn::Default`].
968///
969/// Retries when either:
970/// * the error is an [`Error::ServiceStatus`] whose status is `408`
971/// (Request Timeout), `429` (Too Many Requests), or any `5xx`; or
972/// * the error is an [`Error::Service`] whose (lowercased) message contains one
973/// of the transport-failure markers the provider clients emit — `"request
974/// failed"` (the prefix wrapping every `reqwest` send error: DNS, connect,
975/// timeout, reset), `"timed out"`, `"timeout"`, `"connection"`, or `"stream
976/// error"`.
977///
978/// Everything else (4xx other than 408/429, parse errors, tool/workflow errors,
979/// non-transport service errors) is treated as non-retryable. This explicitly
980/// includes [`Error::ServiceInvalidAuth`], [`Error::ServiceInvalidRequest`],
981/// and [`Error::ServiceContentFilter`] — authentication/authorization
982/// failures, malformed requests, and content-filter refusals are
983/// non-transient, so retrying would just repeat the same rejection. None of
984/// the three carry a status via [`Error::status`], so they fall through to
985/// the final `_ => false` below (there's no dedicated match arm for them:
986/// merging one in would just duplicate that `false`, which `clippy` flags as
987/// `match_same_arms`).
988fn default_should_retry(err: &Error) -> bool {
989 if let Some(status) = err.status() {
990 return status == 408 || status == 429 || (500..600).contains(&status);
991 }
992 match err {
993 Error::Service(msg) => {
994 let m = msg.to_lowercase();
995 m.contains("request failed")
996 || m.contains("timed out")
997 || m.contains("timeout")
998 || m.contains("connection")
999 || m.contains("stream error")
1000 }
1001 _ => false,
1002 }
1003}
1004
1005/// Policy controlling [`RetryingChatClient`] backoff.
1006///
1007/// Delays grow exponentially from [`initial_delay`](Self::initial_delay) by
1008/// [`backoff_multiplier`](Self::backoff_multiplier) per attempt, are capped at
1009/// [`max_delay`](Self::max_delay), and are then reduced by up to
1010/// [`jitter`](Self::jitter) (a fraction of the delay). When the failing error
1011/// carries a server `Retry-After` (see [`Error::retry_after`]) that value is
1012/// used instead of the computed backoff (still capped by `max_delay`, and not
1013/// jittered — it is an explicit server instruction).
1014#[derive(Clone, Debug)]
1015pub struct RetryPolicy {
1016 /// Maximum number of *retries* after the initial attempt (default `3`, so
1017 /// up to four total attempts).
1018 pub max_retries: usize,
1019 /// Base delay before the first retry (default `500ms`).
1020 pub initial_delay: Duration,
1021 /// Upper bound on any single delay, also capping a server `Retry-After`
1022 /// (default `30s`).
1023 pub max_delay: Duration,
1024 /// Exponential growth factor applied per retry (default `2.0`).
1025 pub backoff_multiplier: f64,
1026 /// Jitter as a fraction in `0.0..=1.0` (default `0.3`): the computed delay
1027 /// is multiplied by `1 - jitter * r` for a per-attempt random `r` in
1028 /// `[0, 1)`. `0.0` disables jitter (fully deterministic delays).
1029 pub jitter: f64,
1030 /// Which errors to retry (default [`RetryOn::Default`]).
1031 pub retry_on: RetryOn,
1032}
1033
1034impl Default for RetryPolicy {
1035 fn default() -> Self {
1036 Self {
1037 max_retries: 3,
1038 initial_delay: Duration::from_millis(500),
1039 max_delay: Duration::from_secs(30),
1040 backoff_multiplier: 2.0,
1041 jitter: 0.3,
1042 retry_on: RetryOn::Default,
1043 }
1044 }
1045}
1046
1047impl RetryPolicy {
1048 /// A policy with the given retry count and otherwise-default backoff.
1049 pub fn with_max_retries(max_retries: usize) -> Self {
1050 Self {
1051 max_retries,
1052 ..Self::default()
1053 }
1054 }
1055
1056 /// Set the base delay before the first retry.
1057 pub fn initial_delay(mut self, delay: Duration) -> Self {
1058 self.initial_delay = delay;
1059 self
1060 }
1061
1062 /// Set the per-delay cap (also caps a server `Retry-After`).
1063 pub fn max_delay(mut self, delay: Duration) -> Self {
1064 self.max_delay = delay;
1065 self
1066 }
1067
1068 /// Set the exponential growth factor.
1069 pub fn backoff_multiplier(mut self, multiplier: f64) -> Self {
1070 self.backoff_multiplier = multiplier;
1071 self
1072 }
1073
1074 /// Set the jitter fraction (clamped to `0.0..=1.0`).
1075 pub fn jitter(mut self, jitter: f64) -> Self {
1076 self.jitter = jitter.clamp(0.0, 1.0);
1077 self
1078 }
1079
1080 /// Set the retryability rule.
1081 pub fn retry_on(mut self, retry_on: RetryOn) -> Self {
1082 self.retry_on = retry_on;
1083 self
1084 }
1085
1086 /// The delay to wait before a retry, given the 1-based `attempt` number
1087 /// (attempt `1` is the first retry) and the error that triggered it.
1088 fn delay_for(&self, attempt: usize, err: &Error) -> Duration {
1089 // A server-advised `Retry-After` wins over computed backoff (capped by
1090 // `max_delay`, not jittered — it is an explicit instruction).
1091 if let Some(secs) = err.retry_after() {
1092 let capped = secs.min(self.max_delay.as_secs_f64()).max(0.0);
1093 return Duration::from_secs_f64(capped);
1094 }
1095 let exp = self.backoff_multiplier.powi((attempt - 1) as i32);
1096 let base = self.initial_delay.as_secs_f64() * exp;
1097 let capped = base.min(self.max_delay.as_secs_f64());
1098 let jittered = capped * jitter_factor(self.jitter);
1099 Duration::from_secs_f64(jittered.max(0.0))
1100 }
1101}
1102
1103/// A cheap jitter multiplier in `[1 - jitter, 1.0]`, without a `rand`
1104/// dependency: entropy comes from the current wall-clock nanoseconds mixed
1105/// with a process-lifetime counter (so repeated calls within the same
1106/// nanosecond still differ). `jitter <= 0` returns `1.0` (no jitter).
1107fn jitter_factor(jitter: f64) -> f64 {
1108 let jitter = jitter.clamp(0.0, 1.0);
1109 if jitter == 0.0 {
1110 return 1.0;
1111 }
1112 use std::sync::atomic::{AtomicU64, Ordering};
1113 static COUNTER: AtomicU64 = AtomicU64::new(0);
1114 let nanos = SystemTime::now()
1115 .duration_since(UNIX_EPOCH)
1116 .map(|d| d.as_nanos() as u64)
1117 .unwrap_or(0);
1118 let mixed = nanos ^ COUNTER.fetch_add(0x9E37_79B9_7F4A_7C15, Ordering::Relaxed);
1119 // Map to [0, 1) via the top 53 bits (f64 mantissa width).
1120 let r = (mixed >> 11) as f64 / ((1u64 << 53) as f64);
1121 1.0 - jitter * r
1122}
1123
1124/// A [`ChatClient`] decorator that retries transient failures with exponential
1125/// backoff, honoring a server `Retry-After` when present.
1126///
1127/// Wraps any inner [`ChatClient`] and re-issues the request per its
1128/// [`RetryPolicy`]. For streaming, only the *initial connection* is retried:
1129/// if establishing the stream (or its very first item, before anything is
1130/// yielded to the consumer) fails with a retryable error, the connection is
1131/// re-attempted; once the first update flows, later stream errors propagate
1132/// unchanged.
1133///
1134/// ```no_run
1135/// # use std::time::Duration;
1136/// # use agent_framework_core::client::{RetryingChatClient, RetryPolicy};
1137/// # use agent_framework_core::prelude::*;
1138/// # fn demo(inner: impl ChatClient + 'static) {
1139/// let client = RetryingChatClient::new(inner)
1140/// .with_policy(RetryPolicy::with_max_retries(5).initial_delay(Duration::from_millis(200)));
1141/// # let _ = client;
1142/// # }
1143/// ```
1144pub struct RetryingChatClient<C: ChatClient> {
1145 inner: C,
1146 policy: RetryPolicy,
1147}
1148
1149impl<C: ChatClient> RetryingChatClient<C> {
1150 /// Wrap `inner` with the default [`RetryPolicy`].
1151 pub fn new(inner: C) -> Self {
1152 Self {
1153 inner,
1154 policy: RetryPolicy::default(),
1155 }
1156 }
1157
1158 /// Set the retry policy (builder-style).
1159 pub fn with_policy(mut self, policy: RetryPolicy) -> Self {
1160 self.policy = policy;
1161 self
1162 }
1163
1164 /// A reference to the wrapped client.
1165 pub fn inner(&self) -> &C {
1166 &self.inner
1167 }
1168
1169 /// A reference to the active retry policy.
1170 pub fn policy(&self) -> &RetryPolicy {
1171 &self.policy
1172 }
1173
1174 /// Sleep before a retry, emitting a tracing warning describing the attempt.
1175 async fn backoff(&self, attempt: usize, err: &Error) {
1176 let delay = self.policy.delay_for(attempt, err);
1177 tracing::warn!(
1178 attempt,
1179 max_retries = self.policy.max_retries,
1180 delay_ms = delay.as_millis() as u64,
1181 retry_after = err.retry_after(),
1182 status = err.status(),
1183 error = %err,
1184 "retrying chat request after transient error"
1185 );
1186 tokio::time::sleep(delay).await;
1187 }
1188}
1189
1190#[async_trait]
1191impl<C: ChatClient> ChatClient for RetryingChatClient<C> {
1192 async fn get_response(
1193 &self,
1194 messages: Vec<Message>,
1195 options: ChatOptions,
1196 ) -> Result<ChatResponse> {
1197 let mut attempt = 0usize;
1198 loop {
1199 match self
1200 .inner
1201 .get_response(messages.clone(), options.clone())
1202 .await
1203 {
1204 Ok(resp) => return Ok(resp),
1205 Err(e) => {
1206 if attempt >= self.policy.max_retries || !self.policy.retry_on.should_retry(&e)
1207 {
1208 return Err(e);
1209 }
1210 attempt += 1;
1211 self.backoff(attempt, &e).await;
1212 }
1213 }
1214 }
1215 }
1216
1217 async fn get_streaming_response(
1218 &self,
1219 messages: Vec<Message>,
1220 options: ChatOptions,
1221 ) -> Result<ChatStream> {
1222 let mut attempt = 0usize;
1223 loop {
1224 let established = self
1225 .inner
1226 .get_streaming_response(messages.clone(), options.clone())
1227 .await;
1228 match established {
1229 // The stream opened: peek its first item. An error there (with
1230 // nothing yet yielded to the consumer) is still an
1231 // initial-connection failure and is eligible for retry; any Ok
1232 // item — or a non-retryable / retries-exhausted error — is
1233 // handed back with the rest of the stream chained after it.
1234 Ok(mut stream) => match stream.next().await {
1235 Some(Err(e))
1236 if attempt < self.policy.max_retries
1237 && self.policy.retry_on.should_retry(&e) =>
1238 {
1239 attempt += 1;
1240 self.backoff(attempt, &e).await;
1241 continue;
1242 }
1243 Some(first) => {
1244 let head = stream::once(async move { first });
1245 return Ok(head.chain(stream).boxed());
1246 }
1247 None => return Ok(stream::empty().boxed()),
1248 },
1249 // The stream never opened (e.g. a non-success HTTP status).
1250 Err(e) => {
1251 if attempt >= self.policy.max_retries || !self.policy.retry_on.should_retry(&e)
1252 {
1253 return Err(e);
1254 }
1255 attempt += 1;
1256 self.backoff(attempt, &e).await;
1257 }
1258 }
1259 }
1260 }
1261
1262 fn model(&self) -> Option<&str> {
1263 self.inner.model()
1264 }
1265}
1266
1267#[cfg(test)]
1268mod approval_replacement_tests {
1269 use super::*;
1270 use crate::types::{FunctionApprovalRequestContent, FunctionArguments};
1271
1272 fn call(call_id: &str) -> FunctionCallContent {
1273 FunctionCallContent::new(call_id, "get_weather", None)
1274 }
1275
1276 fn approval_request(call_id: &str) -> Message {
1277 Message::with_contents(
1278 Role::assistant(),
1279 vec![Content::FunctionApprovalRequest(
1280 FunctionApprovalRequestContent {
1281 id: format!("req_{call_id}"),
1282 function_call: call(call_id),
1283 },
1284 )],
1285 )
1286 }
1287
1288 fn function_calls(messages: &[Message]) -> Vec<&str> {
1289 messages
1290 .iter()
1291 .flat_map(|m| m.contents.iter())
1292 .filter_map(Content::as_function_call)
1293 .map(|fc| fc.call_id.as_str())
1294 .collect()
1295 }
1296
1297 #[test]
1298 fn an_approval_request_in_a_separate_message_does_not_duplicate_the_call() {
1299 // The round-trip shape: a hosting layer replays the stored function
1300 // call and its approval request as two *separate* assistant messages.
1301 // Deduping per-message never fired, so the request restored a second
1302 // copy of the call and only one copy received a result — the provider
1303 // then rejects the unanswered one.
1304 let mut messages = vec![
1305 Message::with_contents(Role::assistant(), vec![Content::FunctionCall(call("c1"))]),
1306 approval_request("c1"),
1307 ];
1308 replace_approval_contents_with_results(&mut messages, &HashMap::new());
1309 assert_eq!(function_calls(&messages), vec!["c1"]);
1310 }
1311
1312 #[test]
1313 fn two_approval_requests_for_one_call_expand_only_once() {
1314 let mut messages = vec![approval_request("c1"), approval_request("c1")];
1315 replace_approval_contents_with_results(&mut messages, &HashMap::new());
1316 assert_eq!(function_calls(&messages), vec!["c1"]);
1317 }
1318
1319 #[test]
1320 fn a_single_approval_request_still_expands() {
1321 let mut messages = vec![approval_request("c1")];
1322 replace_approval_contents_with_results(&mut messages, &HashMap::new());
1323 assert_eq!(function_calls(&messages), vec!["c1"]);
1324 }
1325
1326 #[test]
1327 fn a_second_unanswered_call_reusing_an_id_still_suppresses_the_request() {
1328 // The completed c1 pair must not mask the *second*, still-unanswered c1
1329 // call: expanding the replayed request here would hand the provider two
1330 // unanswered copies. "Any result exists for this id" cannot tell these
1331 // apart from the reuse-after-completion case above; a call/result count
1332 // can.
1333 let mut messages = vec![
1334 Message::with_contents(Role::assistant(), vec![Content::FunctionCall(call("c1"))]),
1335 Message::with_contents(
1336 Role::tool(),
1337 vec![Content::FunctionResult(FunctionResultContent {
1338 call_id: "c1".into(),
1339 result: Some(Value::String("sunny".into())),
1340 exception: None,
1341 })],
1342 ),
1343 Message::with_contents(Role::assistant(), vec![Content::FunctionCall(call("c1"))]),
1344 approval_request("c1"),
1345 ];
1346 replace_approval_contents_with_results(&mut messages, &HashMap::new());
1347 assert_eq!(function_calls(&messages), vec!["c1", "c1"]);
1348 }
1349
1350 #[test]
1351 fn a_different_invocation_reusing_an_id_is_not_suppressed() {
1352 // Same call_id, different arguments — a distinct invocation. Judging by
1353 // id alone dropped the new call while its approval response still
1354 // executed, attaching that result to the older call on the wire.
1355 let older = FunctionCallContent::new(
1356 "c1",
1357 "get_weather",
1358 Some(FunctionArguments::Raw("{\"city\":\"old\"}".into())),
1359 );
1360 let newer = FunctionCallContent::new(
1361 "c1",
1362 "get_weather",
1363 Some(FunctionArguments::Raw("{\"city\":\"new\"}".into())),
1364 );
1365 let mut messages = vec![
1366 Message::with_contents(Role::assistant(), vec![Content::FunctionCall(older)]),
1367 Message::with_contents(
1368 Role::assistant(),
1369 vec![Content::FunctionApprovalRequest(
1370 FunctionApprovalRequestContent {
1371 id: "req_1".into(),
1372 function_call: newer,
1373 },
1374 )],
1375 ),
1376 ];
1377 replace_approval_contents_with_results(&mut messages, &HashMap::new());
1378
1379 let cities: Vec<String> = messages
1380 .iter()
1381 .flat_map(|m| m.contents.iter())
1382 .filter_map(Content::as_function_call)
1383 .filter_map(|c| {
1384 c.parse_arguments()
1385 .ok()?
1386 .get("city")?
1387 .as_str()
1388 .map(str::to_string)
1389 })
1390 .collect();
1391 assert!(
1392 cities.contains(&"new".to_string()),
1393 "the new invocation must survive, got {cities:?}"
1394 );
1395 }
1396
1397 #[test]
1398 fn two_replayed_requests_for_one_outstanding_call_both_collapse() {
1399 // Removing the *request* answers nothing, so the call stays
1400 // outstanding. Consuming the anchor let the second replayed copy expand
1401 // into a duplicate declaration with only one result to answer it.
1402 let mut messages = vec![
1403 Message::with_contents(Role::assistant(), vec![Content::FunctionCall(call("c1"))]),
1404 approval_request("c1"),
1405 approval_request("c1"),
1406 ];
1407 replace_approval_contents_with_results(&mut messages, &HashMap::new());
1408 assert_eq!(function_calls(&messages), vec!["c1"]);
1409 }
1410
1411 #[test]
1412 fn a_completed_approval_cycle_does_not_suppress_the_next_one() {
1413 // The first request expands and its response answers it. A later cycle
1414 // reusing the same invocation must still expand: a stale outstanding
1415 // entry suppressed it while its own response still converted, leaving
1416 // two results for the old call and no declaration for the new one.
1417 let approved = FunctionApprovalRequestContent {
1418 id: "req_1".into(),
1419 function_call: call("c1"),
1420 }
1421 .create_response(false);
1422 let mut messages = vec![
1423 approval_request("c1"),
1424 Message::with_contents(
1425 Role::assistant(),
1426 vec![Content::FunctionApprovalResponse(approved)],
1427 ),
1428 approval_request("c1"),
1429 ];
1430 replace_approval_contents_with_results(&mut messages, &HashMap::new());
1431 assert_eq!(
1432 function_calls(&messages),
1433 vec!["c1", "c1"],
1434 "the second cycle needs its own declaration"
1435 );
1436 }
1437
1438 #[test]
1439 fn a_result_arriving_after_the_request_still_suppresses_it() {
1440 // Order matters: at the moment the replayed request is seen, the call
1441 // is still unanswered — the result only arrives later in the list. A
1442 // whole-list pre-scan nets the call against that future result,
1443 // concludes nothing is outstanding, and expands the request into a
1444 // second declaration with a single result between them.
1445 let mut messages = vec![
1446 Message::with_contents(Role::assistant(), vec![Content::FunctionCall(call("c1"))]),
1447 approval_request("c1"),
1448 Message::with_contents(
1449 Role::tool(),
1450 vec![Content::FunctionResult(FunctionResultContent {
1451 call_id: "c1".into(),
1452 result: Some(Value::String("sunny".into())),
1453 exception: None,
1454 })],
1455 ),
1456 ];
1457 replace_approval_contents_with_results(&mut messages, &HashMap::new());
1458 assert_eq!(function_calls(&messages), vec!["c1"]);
1459 }
1460
1461 #[test]
1462 fn a_request_replayed_before_its_call_still_collapses_to_one() {
1463 // The replay order is not guaranteed: the approval request can precede
1464 // the stored call. Whichever comes second is the duplicate.
1465 let mut messages = vec![
1466 approval_request("c1"),
1467 Message::with_contents(Role::assistant(), vec![Content::FunctionCall(call("c1"))]),
1468 ];
1469 replace_approval_contents_with_results(&mut messages, &HashMap::new());
1470 assert_eq!(function_calls(&messages), vec!["c1"]);
1471 }
1472
1473 #[test]
1474 fn an_already_answered_call_does_not_suppress_a_reused_call_id() {
1475 // Reusing a call id for a later invocation is supported: the completed
1476 // pair must not suppress the fresh request, which would drop the new
1477 // call and leave its result attached to the old one.
1478 let mut messages = vec![
1479 Message::with_contents(Role::assistant(), vec![Content::FunctionCall(call("c1"))]),
1480 Message::with_contents(
1481 Role::tool(),
1482 vec![Content::FunctionResult(FunctionResultContent {
1483 call_id: "c1".into(),
1484 result: Some(Value::String("sunny".into())),
1485 exception: None,
1486 })],
1487 ),
1488 approval_request("c1"),
1489 ];
1490 replace_approval_contents_with_results(&mut messages, &HashMap::new());
1491 // Both the completed call and the freshly restored one are present.
1492 assert_eq!(function_calls(&messages), vec!["c1", "c1"]);
1493 }
1494}