Skip to main content

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}