Skip to main content

switchyard_llm_client/
client.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! [`TranslatingLlmClient`]: encode a neutral
5//! request, call the configured backend over HTTP, decode the neutral response.
6
7use std::collections::{BTreeMap, BTreeSet, HashMap};
8use std::future::ready;
9use std::sync::atomic::{AtomicU64, Ordering};
10use std::time::{Duration, SystemTime};
11
12use async_trait::async_trait;
13use futures_util::{StreamExt, stream};
14use http::StatusCode;
15use reqwest::RequestBuilder;
16use reqwest::header::{HeaderMap, RETRY_AFTER};
17use serde_json::{Map, Value, json};
18use switchyard_protocol::{
19    LlmRequest, LlmResponse, LlmResponseChunk, LlmResponseStream, LlmResponseStreamEvent, Metadata,
20    ModelId, Request, Response, RoutedLlmClient,
21};
22use switchyard_translation::{
23    TranslationError, WireFormat, decode_stream, encode_stream_with_extensions,
24};
25use tokio::time::Instant;
26use tracing::Instrument;
27
28use crate::backend::{Backend, openai_url};
29use crate::error::{LlmClientError, Result};
30use crate::metrics;
31use crate::observed_translation;
32use crate::raw::RawResponse;
33
34// Caller headers safe to send when caller auth forwarding is disabled.
35const ALLOWED_METADATA_HEADERS: &[&str] = &["x-request-id"];
36
37// Headers tied to the inbound connection, destination, or body. The HTTP client
38// must rebuild these for the upstream request, even when forwarding auth.
39const NON_FORWARDABLE_HEADERS: &[&str] = &[
40    "host",
41    "content-length",
42    "connection",
43    "proxy-connection",
44    "keep-alive",
45    "proxy-authenticate",
46    "proxy-authorization",
47    "te",
48    "trailer",
49    "transfer-encoding",
50    "upgrade",
51    "content-type",
52    "content-encoding",
53    "accept-encoding",
54    "expect",
55];
56
57const INITIAL_RETRY_DELAY: Duration = Duration::from_millis(250);
58const MAX_RETRY_BACKOFF: Duration = Duration::from_secs(2);
59const MAX_RETRY_AFTER: Duration = Duration::from_secs(60);
60
61/// How one model is served: the `default_backend` used when the request does not
62/// pin a wire format, plus any `other_backends` reachable over additional formats.
63#[derive(Clone, Debug)]
64pub struct ModelConfig {
65    model_name: ModelId,
66    default_backend: Backend,
67    other_backends: Option<Vec<Backend>>,
68}
69
70impl ModelConfig {
71    /// A model named `model_name` served by `default_backend`, optionally reachable
72    /// over additional wire formats via `other_backends`.
73    pub fn new(
74        model_name: impl Into<ModelId>,
75        default_backend: Backend,
76        other_backends: Option<Vec<Backend>>,
77    ) -> Self {
78        Self {
79            model_name: model_name.into(),
80            default_backend,
81            other_backends,
82        }
83    }
84}
85
86/// A model-bearing provider operation outside the normal completion endpoint.
87#[derive(Clone, Copy, Debug, Eq, PartialEq)]
88pub enum AuxiliaryOperation {
89    /// Anthropic Messages input-token counting.
90    AnthropicCountTokens,
91    /// OpenAI Responses input-token counting.
92    ResponsesInputTokens,
93    /// OpenAI Responses compaction.
94    ResponsesCompact,
95}
96
97impl AuxiliaryOperation {
98    const fn wire_format(self) -> WireFormat {
99        match self {
100            Self::AnthropicCountTokens => WireFormat::AnthropicMessages,
101            Self::ResponsesInputTokens | Self::ResponsesCompact => WireFormat::OpenAiResponses,
102        }
103    }
104
105    fn url(self, backend: &Backend) -> String {
106        let suffix = match self {
107            Self::AnthropicCountTokens => return backend.count_tokens_url(),
108            Self::ResponsesInputTokens => "/responses/input_tokens",
109            Self::ResponsesCompact => "/responses/compact",
110        };
111        let base_url = backend.url();
112        openai_url(&base_url, suffix).unwrap_or_else(|error| {
113            tracing::error!(%error, "Unable to build OpenAI auxiliary endpoint URL");
114            base_url
115        })
116    }
117}
118
119/// A client that dispatches neutral-IR requests to per-model HTTP backends.
120///
121/// Construct it with a list of [`ModelConfig`]s — one per model, each naming a
122/// default [`Backend`] and any additional per-format backends. Each call resolves
123/// the model and wire format, encodes the request to that backend's wire format,
124/// applies auth and forwarded headers, sends the HTTP request with a shared
125/// [`reqwest::Client`], and decodes the response back to the neutral IR (buffered
126/// or streamed).
127pub struct TranslatingLlmClient {
128    model_to_config: HashMap<ModelId, ModelConfig>,
129    client: reqwest::Client,
130    forward_auth_client: reqwest::Client,
131    is_rate_limit_cooldown_enabled: bool,
132    cooldown_epoch: Instant,
133    // Milliseconds since cooldown_epoch. Each atomic is independent, so relaxed ordering suffices.
134    unavailable_until: HashMap<ModelId, AtomicU64>,
135}
136
137impl TranslatingLlmClient {
138    /// Builds a client over the given [`ModelConfig`]s, with a fresh shared HTTP
139    /// client and the built-in translation codecs.
140    pub fn new(model_configs: &[ModelConfig]) -> Result<Self> {
141        // Pre-populate so concurrent calls update only atomics and keep map access lock-free.
142        let mut unavailable_until = HashMap::new();
143        for config in model_configs {
144            for backend in std::iter::once(&config.default_backend)
145                .chain(config.other_backends.iter().flatten())
146            {
147                backend.validate_configured_headers(&config.model_name)?;
148            }
149            unavailable_until.insert(config.model_name.clone(), AtomicU64::new(0));
150        }
151        let build_client = |builder: reqwest::ClientBuilder| {
152            builder.build().map_err(|error| LlmClientError::Transport {
153                source: Box::new(error),
154            })
155        };
156        let client = build_client(reqwest::Client::builder())?;
157        // A redirect could move provider-specific headers to another origin.
158        // Forwarded credentials are sent only to the configured URL.
159        let forward_auth_client =
160            build_client(reqwest::Client::builder().redirect(reqwest::redirect::Policy::none()))?;
161        let model_to_config = model_configs
162            .iter()
163            .map(|config| (config.model_name.clone(), config.clone()))
164            .collect();
165
166        Ok(Self {
167            model_to_config,
168            client,
169            forward_auth_client,
170            is_rate_limit_cooldown_enabled: true,
171            cooldown_epoch: Instant::now(),
172            unavailable_until,
173        })
174    }
175
176    /// Disable HTTP 429 cooldown so subsequent calls can receive the upstream error.
177    pub fn without_rate_limit_cooldown(mut self) -> Self {
178        self.is_rate_limit_cooldown_enabled = false;
179        self
180    }
181
182    /// The backend serving `model` over `format` — the default backend when its
183    /// format matches, otherwise a matching entry in `other_backends`; `None` when
184    /// the model is unknown or has no backend for `format`.
185    pub fn backend_for(&self, model: &ModelId, format: WireFormat) -> Option<&Backend> {
186        self.model_to_config.get(model).and_then(|config| {
187            if config.default_backend.wire_format() == format {
188                Some(&config.default_backend)
189            } else {
190                config
191                    .other_backends
192                    .as_ref()
193                    .and_then(|backends| backends.iter().find(|b| b.wire_format() == format))
194            }
195        })
196    }
197
198    /// Whether `model` has a backend for `operation`.
199    pub fn supports_auxiliary(&self, model: &ModelId, operation: AuxiliaryOperation) -> bool {
200        self.backend_for(model, operation.wire_format()).is_some()
201    }
202
203    /// Reads llama.cpp model properties through the backend's configured auth and headers.
204    ///
205    /// Redirects are disabled so credentials cannot move to another origin.
206    pub async fn get_model_properties(
207        &self,
208        model: &ModelId,
209        format: WireFormat,
210        timeout: Duration,
211    ) -> Result<Value> {
212        let backend =
213            self.backend_for(model, format)
214                .ok_or_else(|| LlmClientError::Configuration {
215                    message: format!("model {model} has no backend for {format:?}"),
216                })?;
217        let mut url = reqwest::Url::parse(backend.base_url()).map_err(|error| {
218            LlmClientError::Configuration {
219                message: format!("model {model} has an invalid backend URL: {error}"),
220            }
221        })?;
222        let root = url.path().trim_end_matches('/');
223        let root = root
224            .strip_suffix("/chat/completions")
225            .or_else(|| root.strip_suffix("/responses"))
226            .unwrap_or(root);
227        let path = format!(
228            "{}/props",
229            root.strip_suffix("/v1")
230                .unwrap_or(root)
231                .trim_end_matches('/')
232        );
233        url.set_path(&path);
234        url.set_query(None);
235
236        let builder = self.forward_auth_client.get(url).timeout(timeout);
237        let builder = apply_extra_headers(builder, backend);
238        let response = backend
239            .apply_auth(builder)
240            .send()
241            .await
242            .map_err(convert_reqwest_error)?
243            .error_for_status()
244            .map_err(convert_reqwest_error)?;
245        response
246            .json()
247            .await
248            .map_err(|source| LlmClientError::InvalidResponse {
249                source: Box::new(source),
250            })
251    }
252
253    /// Calls a model-bearing auxiliary provider operation.
254    ///
255    /// Returns an error when the model has no compatible backend or the upstream
256    /// request fails or returns invalid JSON.
257    pub async fn call_auxiliary(
258        &self,
259        model: &ModelId,
260        request: Request,
261        operation: AuxiliaryOperation,
262    ) -> Result<Value> {
263        let wire_format = operation.wire_format();
264        let backend =
265            self.backend_for(model, wire_format)
266                .ok_or_else(|| LlmClientError::Configuration {
267                    message: format!("model {model} has no backend for {operation:?}"),
268                })?;
269        let Request {
270            mut llm_request,
271            metadata,
272            ..
273        } = request;
274        llm_request.model = Some(model.to_string());
275        let http_response = self
276            .send_encoded(
277                backend,
278                wire_format,
279                llm_request,
280                metadata.as_ref(),
281                model,
282                UpstreamEndpoint::Auxiliary(operation),
283            )
284            .await?;
285        let EncodedResponse::Buffered { body, .. } = http_response else {
286            return Err(LlmClientError::InvalidRequest {
287                message: "auxiliary endpoints do not support streaming".to_string(),
288            });
289        };
290        serde_json::from_slice(&body).map_err(|error| LlmClientError::InvalidResponse {
291            source: Box::new(error),
292        })
293    }
294
295    /// Encode `llm_request` for `wire_format`, POST it to `url` with the request's
296    /// forwarded headers plus the backend's static headers and auth, and return the
297    /// successful upstream response. A
298    /// buffered response is fully collected within the retry boundary; a streamed
299    /// response has its first decoded event checked within that boundary. The backend's
300    /// deadline, when set, spans every attempt, retry delay, and subsequent stream read. A
301    /// non-success status maps to a typed error — a 400 is classified as a context-window
302    /// overflow via the backend's provider rules. Shared by
303    /// [`call_rewrite_model`](Self::call_rewrite_model) (which POSTs to the
304    /// backend's completion URL and decodes a response) and
305    /// the model-bearing auxiliary operations, which return raw JSON.
306    async fn send_encoded(
307        &self,
308        backend: &Backend,
309        wire_format: WireFormat,
310        llm_request: LlmRequest,
311        metadata: Option<&Metadata>,
312        model: &ModelId,
313        endpoint: UpstreamEndpoint,
314    ) -> Result<EncodedResponse> {
315        let cooldown = backend.failure_cooldown();
316        let unavailable_until =
317            if matches!(endpoint, UpstreamEndpoint::Completion) && !cooldown.is_zero() {
318                self.unavailable_until.get(model)
319            } else {
320                None
321            };
322        if let Some(until) = unavailable_until
323            && until.load(Ordering::Relaxed) > duration_millis(self.cooldown_epoch.elapsed())
324        {
325            tracing::debug!(model = %model, wire_format = %wire_format, "skipping backend during failure cooldown");
326            return Err(LlmClientError::TemporarilyUnavailable);
327        }
328
329        let mut body = observed_translation::encode_request(&llm_request, wire_format)
330            .map_err(|error| LlmClientError::RequestEncoding(error.to_string()))?;
331        // `encode_request` round-trips a preserved same-format body verbatim,
332        // which keeps the caller's original `model`; force the resolved model so
333        // the upstream always sees the target id.
334        set_json_model(&mut body, model);
335        if matches!(backend, Backend::OpenAiResponses(_)) {
336            sanitize_openai_responses_provider_body(&mut body);
337        }
338        // Strip before `merge_extra_body` so a target can reinstate either field
339        // deliberately via `extra_body`.
340        if matches!(backend, Backend::Anthropic(_)) {
341            strip_anthropic_incompatible_fields(&mut body);
342            strip_unsigned_thinking_blocks(&mut body);
343        }
344        omit_configured_body_fields(&mut body, backend.omit_body_fields());
345        merge_extra_body(&mut body, backend.extra_body());
346        // After the merge on purpose: the effort override must win over both the caller's
347        // value and any `reasoning` default a target set through `extra_body`.
348        apply_reasoning_effort(&mut body, backend);
349        if matches!(backend, Backend::Anthropic(_)) {
350            if nvidia_inference_api_backend(backend) {
351                strip_anthropic_tool_strict(&mut body);
352            }
353            enable_anthropic_prompt_caching(&mut body);
354        }
355        if matches!(backend, Backend::OpenAiChat(_)) {
356            ensure_openai_stream_usage(&mut body);
357        }
358        let streaming = endpoint.allows_streaming()
359            && body.get("stream").and_then(Value::as_bool).unwrap_or(false);
360        let url = endpoint.url(backend);
361        record_gen_ai_request(&url, model, streaming);
362
363        let result = self
364            .send_with_retries(&url, backend, &body, metadata, model, streaming)
365            .await;
366        // Keep quota failures request-local when forwarding auth or returning 429 to callers.
367        if let Some(until) = unavailable_until
368            && let Err(error) = &result
369            && is_transient_failure(error)
370            && !((backend.is_forwarding_auth() || !self.is_rate_limit_cooldown_enabled)
371                && matches!(error, LlmClientError::UpstreamHttp { status, .. } if *status == StatusCode::TOO_MANY_REQUESTS))
372        {
373            let deadline = duration_millis(self.cooldown_epoch.elapsed().saturating_add(cooldown));
374            until.fetch_max(deadline, Ordering::Relaxed);
375        }
376        // Let cooldown expire naturally; a successful in-flight call may overlap a newer failure.
377        result
378    }
379
380    // Sends the encoded body, retrying retryable failures within the backend's retry budget.
381    async fn send_with_retries(
382        &self,
383        url: &str,
384        backend: &Backend,
385        body: &Value,
386        metadata: Option<&Metadata>,
387        model: &ModelId,
388        streaming: bool,
389    ) -> Result<EncodedResponse> {
390        let max_retries = u64::from(backend.max_retries());
391        let max_attempts = max_retries + 1;
392        let expires_at = backend
393            .timeout()
394            .map(|timeout| (tokio::time::sleep(timeout).deadline(), timeout));
395        let deadline = async {
396            match expires_at {
397                Some((at, timeout)) => {
398                    tokio::time::sleep_until(at).await;
399                    timeout
400                }
401                None => std::future::pending().await,
402            }
403        };
404        tokio::pin!(deadline);
405        let mut attempt = 0_u64;
406        loop {
407            if let Some((at, timeout)) = expires_at
408                && tokio::time::Instant::now() >= at
409            {
410                return Err(deadline_error(timeout));
411            }
412            let span = tracing::debug_span!(
413                target: "libsy",
414                "libsy.upstream_attempt",
415                model = %model,
416                wire_format = %backend.wire_format(),
417                attempt = attempt + 1,
418                max_attempts,
419                retry = attempt > 0,
420                openinference.span.kind = "CHAIN",
421                outcome = tracing::field::Empty,
422                status_code = tracing::field::Empty,
423                will_retry = tracing::field::Empty,
424                retry_delay_ms = tracing::field::Empty,
425            );
426            let mut attempt_started = false;
427            let result = tokio::select! {
428                biased;
429                timeout = &mut deadline => {
430                    if attempt_started {
431                        metrics::record_upstream_attempt(None);
432                    }
433                    span.record("outcome", "error");
434                    span.record("will_retry", false);
435                    return Err(deadline_error(timeout));
436                }
437                result = async {
438                    attempt_started = true;
439                    self.send_once(url, backend, body, metadata, model, streaming).await
440                }.instrument(span.clone()) => result,
441            };
442            // The retained handle updates this same attempt span with its outcome.
443            match result {
444                Ok(response) => {
445                    span.record("outcome", "ok");
446                    span.record("status_code", response.status());
447                    span.record("will_retry", false);
448                    if attempt > 0 {
449                        metrics::record_retry_recovered();
450                    }
451                    return Ok(match response {
452                        EncodedResponse::Streaming {
453                            status,
454                            chunks,
455                            upstream_headers,
456                        } => EncodedResponse::Streaming {
457                            status,
458                            chunks: stream_with_deadline(chunks, expires_at),
459                            upstream_headers,
460                        },
461                        buffered => buffered,
462                    });
463                }
464                Err(failure) => {
465                    let will_retry = attempt < max_retries && failure.is_retryable();
466                    span.record("outcome", "error");
467                    if let Some(status) = failure.status {
468                        span.record("status_code", status.as_u16());
469                    }
470                    span.record("will_retry", will_retry);
471                    if !will_retry {
472                        return Err(failure.error);
473                    }
474
475                    let delay = retry_delay(attempt, failure.retry_after);
476                    span.record("retry_delay_ms", duration_millis(delay));
477                    // Close the attempt span before sleeping so backoff is not attempt latency.
478                    drop(span);
479                    tokio::select! {
480                        biased;
481                        timeout = &mut deadline => return Err(deadline_error(timeout)),
482                        _ = tokio::time::sleep(delay) => {}
483                    }
484                    attempt += 1;
485                }
486            }
487        }
488    }
489
490    // Performs one HTTP attempt and retains the retry metadata alongside any error.
491    async fn send_once(
492        &self,
493        url: &str,
494        backend: &Backend,
495        body: &Value,
496        metadata: Option<&Metadata>,
497        model: &ModelId,
498        streaming: bool,
499    ) -> std::result::Result<EncodedResponse, AttemptFailure> {
500        let client = if backend.is_forwarding_auth() {
501            &self.forward_auth_client
502        } else {
503            &self.client
504        };
505        let builder = client.post(url).json(body);
506        let builder = forward_metadata_headers(builder, metadata, backend);
507        let builder = backend.apply_forwarded_auth(builder, metadata);
508        let builder = apply_extra_headers(builder, backend);
509        let builder = backend.apply_auth(builder);
510
511        let response = match builder.send().await {
512            Ok(response) => response,
513            Err(error) => {
514                metrics::record_upstream_attempt(None);
515                return Err(AttemptFailure {
516                    error: convert_reqwest_error(error),
517                    status: None,
518                    retry_after: None,
519                });
520            }
521        };
522        let status = response.status();
523        if status.is_success() {
524            if streaming {
525                let upstream_headers = response.headers().clone();
526                let chunks = match prepare_response_stream(response, backend, model).await {
527                    Ok(chunks) => chunks,
528                    Err(error) => {
529                        metrics::record_upstream_attempt(None);
530                        return Err(AttemptFailure {
531                            error,
532                            status: Some(status),
533                            retry_after: None,
534                        });
535                    }
536                };
537                metrics::record_upstream_attempt(Some(status.as_u16()));
538                return Ok(EncodedResponse::Streaming {
539                    status: status.as_u16(),
540                    chunks,
541                    upstream_headers,
542                });
543            }
544            let upstream_headers = response.headers().clone();
545            let body = match response.bytes().await {
546                Ok(body) => body,
547                Err(error) => {
548                    metrics::record_upstream_attempt(None);
549                    return Err(AttemptFailure {
550                        error: convert_reqwest_error(error),
551                        status: Some(status),
552                        retry_after: None,
553                    });
554                }
555            };
556            metrics::record_upstream_attempt(Some(status.as_u16()));
557            return Ok(EncodedResponse::Buffered {
558                status: status.as_u16(),
559                body: body.to_vec(),
560                upstream_headers,
561            });
562        }
563
564        let retry_after = retry_after_delay(response.headers());
565        let body = match response.text().await {
566            Ok(body) => body,
567            Err(error) => {
568                metrics::record_upstream_attempt(None);
569                return Err(AttemptFailure {
570                    error: convert_reqwest_error(error),
571                    status: Some(status),
572                    retry_after,
573                });
574            }
575        };
576        let body = redact_forwarded_headers(body, metadata, backend.is_forwarding_auth());
577        metrics::record_upstream_attempt(Some(status.as_u16()));
578        let error =
579            if status == reqwest::StatusCode::BAD_REQUEST && backend.is_context_overflow(&body) {
580                LlmClientError::ContextWindowExceeded {
581                    model: model.clone(),
582                    message: body,
583                }
584            } else {
585                LlmClientError::UpstreamHttp { status, body }
586            };
587        Err(AttemptFailure {
588            error,
589            status: Some(status),
590            retry_after,
591        })
592    }
593
594    /// Calls the backend for `model_name` (or the request's own model), over the
595    /// wire format the request pins in its metadata (else the model's default
596    /// backend), and returns the neutral response.
597    ///
598    /// Resolution: `model_name` wins over `request.llm_request.model`; the
599    /// resolved name is both the outer map key and the model id written into the
600    /// request before translation. Missing models are invalid requests; unknown
601    /// models or wire formats are configuration errors.
602    pub async fn call_rewrite_model(
603        &self,
604        request: Request,
605        model_name: Option<&ModelId>,
606    ) -> Result<Response> {
607        let Request {
608            mut llm_request,
609            metadata,
610            ..
611        } = request;
612
613        let model_id = model_name
614            .cloned()
615            .or_else(|| llm_request.model.map(ModelId::from))
616            .ok_or_else(|| LlmClientError::InvalidRequest {
617                message: "no model given".to_string(),
618            })?;
619        llm_request.model = Some(model_id.to_string());
620
621        let orig_format = metadata.as_ref().and_then(|m| m.wire_format);
622        let wire_format = orig_format.unwrap_or(
623            self.model_to_config
624                .get(&model_id)
625                .map(|config| config.default_backend.wire_format())
626                .ok_or_else(|| LlmClientError::Configuration {
627                    message: format!("no backend configured for model {model_id:?}"),
628                })?,
629        );
630        let backend = self.backend_for(&model_id, wire_format).ok_or_else(|| {
631            LlmClientError::Configuration {
632                message: format!("model {model_id:?} has no backend for format {wire_format}"),
633            }
634        })?;
635
636        let http_response = self
637            .send_encoded(
638                backend,
639                wire_format,
640                llm_request,
641                metadata.as_ref(),
642                &model_id,
643                UpstreamEndpoint::Completion,
644            )
645            .await?;
646
647        let (llm_response, upstream_headers) = match http_response {
648            EncodedResponse::Streaming {
649                chunks,
650                upstream_headers,
651                ..
652            } => (LlmResponse::Stream(chunks), upstream_headers),
653            EncodedResponse::Buffered {
654                body,
655                upstream_headers,
656                ..
657            } => {
658                let body = serde_json::from_slice::<Value>(&body).map_err(|error| {
659                    LlmClientError::ResponseTranslation(format!("invalid upstream JSON: {error}"))
660                })?;
661                // Map a provider's failed generation to 502, even under HTTP 200.
662                // Redact forwarded credentials before returning the provider error.
663                let mut agg =
664                    observed_translation::decode_response(&body, wire_format).map_err(|error| {
665                        match error {
666                            TranslationError::UpstreamFailure { error } => {
667                                LlmClientError::UpstreamHttp {
668                                    status: StatusCode::BAD_GATEWAY,
669                                    body: redact_forwarded_headers(
670                                        json!({ "error": error }).to_string(),
671                                        metadata.as_ref(),
672                                        backend.is_forwarding_auth(),
673                                    ),
674                                }
675                            }
676                            error => LlmClientError::ResponseTranslation(error.to_string()),
677                        }
678                    })?;
679                // The gateway may report a buffered charge only in its response header.
680                if agg.usage.provider_cost.is_none() {
681                    agg.usage.provider_cost = upstream_headers
682                        .get("x-litellm-response-cost")
683                        .and_then(|value| value.to_str().ok())
684                        .and_then(|value| value.trim().parse::<f64>().ok())
685                        .filter(|usd| usd.is_finite() && *usd >= 0.0);
686                }
687                (LlmResponse::Agg(agg), upstream_headers)
688            }
689        };
690
691        Ok(Response {
692            llm_response,
693            metadata,
694            upstream_headers,
695        })
696    }
697
698    /// The whole decode → call → encode path a wire endpoint needs, in one call.
699    ///
700    /// Decodes `raw_http_request` from `wire_format` to the neutral IR, serves it via
701    /// [`call_rewrite_model`](Self::call_rewrite_model) — the *upstream* wire format is
702    /// resolved there from the model's backend, independently of `wire_format` — then
703    /// encodes the neutral response back into `wire_format`. The result is a buffered
704    /// [`RawResponse::Buffered`] JSON body or a streamed [`RawResponse::Stream`] of
705    /// wire events (the caller frames the stream as SSE). The response's `model` is
706    /// restamped with the model that actually served the call, so the body names the
707    /// model that answered rather than the route the caller addressed.
708    ///
709    /// `http_headers` are carried through as the request's
710    /// [`Metadata::http_headers`]. Backends with `forward_auth` disabled forward only
711    /// allowed metadata headers; `forward_auth` backends forward all application
712    /// headers. The forwarding backends in one route must use one credential family
713    /// unless they all use the same scheme, host, and port. Transport headers are always
714    /// rebuilt. Pass `None` to forward nothing.
715    pub async fn call_rewrite_model_raw(
716        &self,
717        raw_http_request: Value,
718        http_headers: Option<http::HeaderMap>,
719        model: Option<&ModelId>,
720        wire_format: WireFormat,
721    ) -> Result<RawResponse> {
722        let llm_request = observed_translation::decode_request(wire_format, &raw_http_request)
723            .map_err(|error| LlmClientError::RequestTranslation(error.to_string()))?;
724        let request_extensions = llm_request.extensions.clone();
725        // The model that serves the call — the rewrite target when the caller pinned
726        // one, else the request's own model. Mirrors `call_rewrite_model`'s own
727        // resolution so the response names whoever answered.
728        let served_model = model
729            .map(ModelId::to_string)
730            .or_else(|| llm_request.model.clone());
731
732        let request = Request {
733            llm_request,
734            raw_request: None,
735            metadata: Some(Metadata {
736                session_id: None,
737                agent_id: None,
738                task_id: None,
739                correlation_id: None,
740                extra_metadata: None,
741                http_headers,
742                wire_format: None,
743                ..Default::default()
744            }),
745        };
746        let response = self.call_rewrite_model(request, model).await?;
747
748        match response.llm_response {
749            LlmResponse::Agg(agg) => {
750                let body = observed_translation::encode_response(
751                    &agg,
752                    wire_format,
753                    served_model.as_deref(),
754                    &request_extensions,
755                )
756                .map_err(|error| LlmClientError::ResponseTranslation(error.to_string()))?;
757                Ok(RawResponse::Buffered(body))
758            }
759            LlmResponse::Stream(chunks) => {
760                let events = encode_stream_with_extensions(
761                    chunks,
762                    wire_format,
763                    served_model,
764                    &request_extensions,
765                )?;
766                Ok(RawResponse::Stream(events))
767            }
768        }
769    }
770}
771
772#[async_trait]
773impl RoutedLlmClient for TranslatingLlmClient {
774    async fn call(&self, request: Request) -> Result<Response> {
775        self.call_rewrite_model(request, None).await
776    }
777}
778
779#[derive(Clone, Copy)]
780enum UpstreamEndpoint {
781    Completion,
782    Auxiliary(AuxiliaryOperation),
783}
784
785impl UpstreamEndpoint {
786    fn url(self, backend: &Backend) -> String {
787        match self {
788            UpstreamEndpoint::Completion => backend.url(),
789            UpstreamEndpoint::Auxiliary(operation) => operation.url(backend),
790        }
791    }
792
793    fn allows_streaming(self) -> bool {
794        matches!(self, UpstreamEndpoint::Completion)
795    }
796}
797
798enum EncodedResponse {
799    Buffered {
800        status: u16,
801        body: Vec<u8>,
802        upstream_headers: HeaderMap,
803    },
804    Streaming {
805        status: u16,
806        chunks: LlmResponseStream,
807        upstream_headers: HeaderMap,
808    },
809}
810
811impl EncodedResponse {
812    fn status(&self) -> u16 {
813        match self {
814            EncodedResponse::Buffered { status, .. }
815            | EncodedResponse::Streaming { status, .. } => *status,
816        }
817    }
818}
819
820// The typed error decides retry eligibility; status and Retry-After feed
821// attempt telemetry and delay selection.
822struct AttemptFailure {
823    error: LlmClientError,
824    status: Option<StatusCode>,
825    retry_after: Option<Duration>,
826}
827
828fn deadline_error(timeout: Duration) -> LlmClientError {
829    LlmClientError::Timeout {
830        source: Box::new(std::io::Error::new(
831            std::io::ErrorKind::TimedOut,
832            format!(
833                "response did not finish within {} ms, retries included",
834                duration_millis(timeout)
835            ),
836        )),
837    }
838}
839
840impl AttemptFailure {
841    fn is_retryable(&self) -> bool {
842        is_transient_failure(&self.error)
843    }
844}
845
846fn is_transient_failure(error: &LlmClientError) -> bool {
847    match error {
848        LlmClientError::Transport { .. } | LlmClientError::Timeout { .. } => true,
849        LlmClientError::UpstreamHttp { status, .. } => {
850            metrics::is_retryable_http_status(status.as_u16())
851        }
852        _ => false,
853    }
854}
855
856fn stream_with_deadline(
857    chunks: LlmResponseStream,
858    deadline: Option<(tokio::time::Instant, Duration)>,
859) -> LlmResponseStream {
860    let Some((at, timeout)) = deadline else {
861        return chunks;
862    };
863    stream::unfold(Some(chunks), move |chunks| async move {
864        let mut chunks = chunks?;
865        let chunk = tokio::select! {
866            biased;
867            _ = tokio::time::sleep_until(at) => return Some((Err(deadline_error(timeout)), None)),
868            chunk = chunks.next() => chunk?,
869        };
870        let next = chunk.is_ok().then_some(chunks);
871        Some((chunk, next))
872    })
873    .boxed()
874}
875
876async fn prepare_response_stream(
877    response: reqwest::Response,
878    backend: &Backend,
879    model: &ModelId,
880) -> Result<LlmResponseStream> {
881    let bytes = response.bytes_stream().map(|chunk| {
882        chunk
883            .map(|bytes| bytes.to_vec())
884            .map_err(convert_reqwest_error)
885    });
886    let mut chunks = decode_stream(bytes, backend.wire_format())?;
887    match chunks.next().await {
888        None => Ok(stream::empty().boxed()),
889        // Nothing has reached the caller, so transport failures can still be retried.
890        Some(Err(error @ (LlmClientError::Transport { .. } | LlmClientError::Timeout { .. }))) => {
891            Err(error)
892        }
893        Some(first) => {
894            // An in-band context overflow skips retries and advances to another candidate.
895            if let Some(message) = first_event_overflow(&first, backend) {
896                return Err(LlmClientError::ContextWindowExceeded {
897                    model: model.clone(),
898                    message,
899                });
900            }
901            Ok(stream::once(ready(first)).chain(chunks).boxed())
902        }
903    }
904}
905
906// The overflow message when a stream's first event is an in-band provider rejection
907// of the whole request, rather than the start of a response.
908fn first_event_overflow(
909    first: &Result<LlmResponseStreamEvent>,
910    backend: &Backend,
911) -> Option<String> {
912    first
913        .as_ref()
914        .ok()?
915        .normalized()
916        .iter()
917        .find_map(|chunk| match chunk {
918            LlmResponseChunk::StreamError { message } if backend.is_context_overflow(message) => {
919                Some(message.clone())
920            }
921            _ => None,
922        })
923}
924
925// Uses Retry-After when supplied, capped so an upstream cannot stall a request indefinitely.
926fn retry_after_delay(headers: &HeaderMap) -> Option<Duration> {
927    let value = headers.get(RETRY_AFTER)?.to_str().ok()?;
928    let delay = if let Ok(seconds) = value.parse::<u64>() {
929        Duration::from_secs(seconds)
930    } else {
931        let retry_at = httpdate::parse_http_date(value).ok()?;
932        retry_at
933            .duration_since(SystemTime::now())
934            .unwrap_or(Duration::ZERO)
935    };
936    Some(delay.min(MAX_RETRY_AFTER))
937}
938
939fn retry_delay(retry_number: u64, retry_after: Option<Duration>) -> Duration {
940    // Retry-After wins; otherwise double 250 ms up to the two-second cap.
941    retry_after.unwrap_or_else(|| {
942        let multiplier = 1_u32 << retry_number.min(3);
943        INITIAL_RETRY_DELAY
944            .saturating_mul(multiplier)
945            .min(MAX_RETRY_BACKOFF)
946    })
947}
948
949fn duration_millis(duration: Duration) -> u64 {
950    u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
951}
952
953fn record_gen_ai_request(url: &str, model: &str, streaming: bool) {
954    let span = tracing::Span::current();
955    span.record("gen_ai.request.model", model);
956    if streaming {
957        span.record("gen_ai.request.stream", true);
958    }
959    if let Ok(url) = reqwest::Url::parse(url) {
960        if let Some(host) = url.host_str() {
961            span.record("server.address", host);
962        }
963        if let Some(port) = url.port_or_known_default() {
964            span.record("server.port", i64::from(port));
965        }
966    }
967}
968
969pub(crate) fn convert_reqwest_error(error: reqwest::Error) -> LlmClientError {
970    // Reqwest labels truncated or otherwise unreadable response bodies as decode
971    // errors, so distinguish them from serde JSON failures at the call site.
972    let error = error.without_url();
973    if error.is_timeout() {
974        LlmClientError::Timeout {
975            source: Box::new(error),
976        }
977    } else if error.is_builder() {
978        LlmClientError::Configuration {
979            message: format!("failed to build upstream request: {error}"),
980        }
981    } else {
982        LlmClientError::Transport {
983            source: Box::new(error),
984        }
985    }
986}
987
988// Forwards safe metadata headers, or all application headers when forwarding auth.
989fn forward_metadata_headers(
990    builder: RequestBuilder,
991    metadata: Option<&Metadata>,
992    backend: &Backend,
993) -> RequestBuilder {
994    let Some(headers) = metadata.and_then(|metadata| metadata.http_headers.as_ref()) else {
995        return builder;
996    };
997    let mut forwarded = HeaderMap::with_capacity(headers.len());
998    for (name, value) in headers {
999        let is_allowed = if backend.is_forwarding_auth() {
1000            !is_non_forwardable_header(name.as_str(), headers)
1001                && !backend.is_provider_owned_header(name.as_str())
1002        } else {
1003            is_allowed_metadata_header(name.as_str())
1004        };
1005        if !is_allowed {
1006            continue;
1007        }
1008        forwarded.append(name, value.clone());
1009    }
1010    builder.headers(forwarded)
1011}
1012
1013// Adds the backend's custom per-call headers.
1014fn apply_extra_headers(mut builder: RequestBuilder, backend: &Backend) -> RequestBuilder {
1015    for (name, value) in backend.extra_headers() {
1016        builder = builder.header(name, value);
1017    }
1018    builder
1019}
1020
1021// Overwrites the outbound body's `model` field with the resolved model id.
1022fn set_json_model(body: &mut Value, model: &str) {
1023    if let Value::Object(object) = body {
1024        object.insert("model".to_string(), Value::String(model.to_string()));
1025    }
1026}
1027
1028const CODEX_NAMESPACE_SEPARATOR: &str = "__";
1029
1030/// Rewrites a Responses body into the shape OpenAI-compatible providers accept.
1031///
1032/// Codex extends Responses with namespace containers and namespaced function calls, and
1033/// OpenAI-compatible providers expect a flat Responses tool namespace. The input also loses
1034/// reasoning items that the provider cannot read.
1035fn sanitize_openai_responses_provider_body(body: &mut Value) {
1036    let Value::Object(object) = body else {
1037        return;
1038    };
1039    let store_off = object.get("store") == Some(&Value::Bool(false));
1040    sanitize_openai_responses_input_for_provider(object.get_mut("input"), store_off);
1041    sanitize_openai_responses_tools_for_provider(object.get_mut("tools"));
1042    sanitize_openai_responses_tool_choice_for_provider(object.get_mut("tool_choice"));
1043}
1044
1045/// Removes reasoning items that the provider cannot read and flattens namespaced function
1046/// calls in the Responses `input` array.
1047fn sanitize_openai_responses_input_for_provider(input: Option<&mut Value>, store_off: bool) {
1048    let Some(Value::Array(items)) = input else {
1049        return;
1050    };
1051    items.retain(|item| !is_unreadable_reasoning_item(item, store_off));
1052    for item in items {
1053        let Some(object) = item.as_object_mut() else {
1054            continue;
1055        };
1056        if object.get("type").and_then(Value::as_str) == Some("function_call") {
1057            qualify_responses_function_name(object);
1058        }
1059    }
1060}
1061
1062/// Returns true for a reasoning item that an OpenAI provider would reject with 400.
1063///
1064/// A router can serve earlier steps of a turn from Claude and a later step from an OpenAI
1065/// model, and the client replays Claude's reasoning items on every step. Two kinds fail:
1066///
1067/// - An item without `encrypted_content` (absent, null, or empty) when the request sets
1068///   `store: false`. OpenAI treats the item's ID as a reference to stored reasoning and finds
1069///   nothing. With `store` on, the ID can refer to stored reasoning, so the item is kept.
1070/// - An item whose `encrypted_content` holds Anthropic thinking that Switchyard packed for a
1071///   Responses client. OpenAI cannot decrypt it, whatever `store` is set to.
1072fn is_unreadable_reasoning_item(item: &Value, store_off: bool) -> bool {
1073    if item.get("type").and_then(Value::as_str) != Some("reasoning") {
1074        return false;
1075    }
1076    match item.get("encrypted_content").and_then(Value::as_str) {
1077        Some(encrypted) if !encrypted.is_empty() => {
1078            switchyard_translation::is_anthropic_thinking_envelope(encrypted)
1079        }
1080        _ => store_off,
1081    }
1082}
1083
1084fn sanitize_openai_responses_tools_for_provider(tools: Option<&mut Value>) {
1085    let Some(Value::Array(tools)) = tools else {
1086        return;
1087    };
1088    let mut flat_tools = Vec::with_capacity(tools.len());
1089    for tool in std::mem::take(tools) {
1090        push_sanitized_openai_responses_tool(&mut flat_tools, tool);
1091    }
1092    *tools = flat_tools;
1093}
1094
1095fn push_sanitized_openai_responses_tool(out: &mut Vec<Value>, tool: Value) {
1096    let Value::Object(mut object) = tool else {
1097        out.push(tool);
1098        return;
1099    };
1100    if object.get("type").and_then(Value::as_str) == Some("namespace") {
1101        let namespace = object
1102            .get("name")
1103            .and_then(Value::as_str)
1104            .unwrap_or_default()
1105            .to_string();
1106        let Some(Value::Array(children)) = object.remove("tools") else {
1107            out.push(Value::Object(object));
1108            return;
1109        };
1110        for child in children {
1111            push_sanitized_namespaced_tool(out, &namespace, child);
1112        }
1113        return;
1114    }
1115    ensure_responses_function_tool_description(&mut object);
1116    out.push(Value::Object(object));
1117}
1118
1119fn push_sanitized_namespaced_tool(out: &mut Vec<Value>, namespace: &str, tool: Value) {
1120    let Value::Object(mut object) = tool else {
1121        out.push(tool);
1122        return;
1123    };
1124    if object.get("type").and_then(Value::as_str) == Some("function") {
1125        qualify_responses_function_name_with_namespace(&mut object, namespace);
1126        ensure_responses_function_tool_description(&mut object);
1127    }
1128    out.push(Value::Object(object));
1129}
1130
1131fn sanitize_openai_responses_tool_choice_for_provider(tool_choice: Option<&mut Value>) {
1132    let Some(Value::Object(object)) = tool_choice else {
1133        return;
1134    };
1135    if object.get("type").and_then(Value::as_str) == Some("function") {
1136        qualify_responses_function_name(object);
1137    }
1138}
1139
1140fn qualify_responses_function_name(object: &mut Map<String, Value>) {
1141    let namespace = object
1142        .remove("namespace")
1143        .and_then(|value| value.as_str().map(ToOwned::to_owned));
1144    let Some(namespace) = namespace.as_deref() else {
1145        return;
1146    };
1147    qualify_responses_function_name_with_namespace(object, namespace);
1148}
1149
1150fn qualify_responses_function_name_with_namespace(
1151    object: &mut Map<String, Value>,
1152    namespace: &str,
1153) {
1154    if namespace.is_empty() {
1155        return;
1156    }
1157    let Some(name) = object.get("name").and_then(Value::as_str) else {
1158        return;
1159    };
1160    let prefix = format!("{namespace}{CODEX_NAMESPACE_SEPARATOR}");
1161    if name.starts_with(&prefix) {
1162        return;
1163    }
1164    object.insert("name".to_string(), Value::String(format!("{prefix}{name}")));
1165}
1166
1167fn ensure_responses_function_tool_description(object: &mut Map<String, Value>) {
1168    if object.get("type").and_then(Value::as_str) == Some("function")
1169        && !matches!(object.get("description"), Some(Value::String(_)))
1170    {
1171        object.insert("description".to_string(), Value::String(String::new()));
1172    }
1173}
1174
1175// Drops fields accepted by OpenAI-like APIs but rejected by Anthropic Messages.
1176//
1177// A router can serve earlier turns of a session from an OpenAI-format target and
1178// later turns from an Anthropic one. Clients such as Claude Code send
1179// `context_management` on every turn, so the Anthropic leg must strip it or the
1180// upstream rejects the request (for example `clear_thinking_20251015` strategy
1181// requires `thinking` to be enabled or adaptive).
1182fn strip_anthropic_incompatible_fields(body: &mut Value) {
1183    if let Value::Object(object) = body {
1184        object.remove("reasoning_effort");
1185        object.remove("context_management");
1186    }
1187}
1188
1189// NVIDIA Inference API's Anthropic-compatible schema rejects the strict tool-use
1190// extension that native Anthropic accepts.
1191fn nvidia_inference_api_backend(backend: &Backend) -> bool {
1192    reqwest::Url::parse(backend.base_url())
1193        .ok()
1194        .is_some_and(|url| {
1195            url.host_str()
1196                .is_some_and(|host| host.eq_ignore_ascii_case("inference-api.nvidia.com"))
1197        })
1198}
1199
1200fn strip_anthropic_tool_strict(body: &mut Value) {
1201    let Some(tools) = body.get_mut("tools").and_then(Value::as_array_mut) else {
1202        return;
1203    };
1204    for tool in tools.iter_mut().filter_map(Value::as_object_mut) {
1205        tool.remove("strict");
1206    }
1207}
1208
1209// Removes replayed `thinking` blocks that carry no signature.
1210//
1211// Anthropic requires signed thinking blocks on replay. A router can serve earlier
1212// turns of a session from an OpenAI-format target whose thinking blocks are
1213// unsigned, so the Anthropic leg must drop them or the upstream rejects the
1214// request. Bedrock enforces this (surfacing as a SigV4 signature mismatch) where
1215// Azure-hosted Anthropic currently does not.
1216fn strip_unsigned_thinking_blocks(body: &mut Value) {
1217    let Value::Object(object) = body else {
1218        return;
1219    };
1220    let Some(Value::Array(messages)) = object.get_mut("messages") else {
1221        return;
1222    };
1223    for message in messages {
1224        strip_unsigned_thinking_from_message(message);
1225    }
1226}
1227
1228// Drops unsigned thinking blocks from one message, collapsing content that ends
1229// up empty to an empty string so the message stays valid.
1230fn strip_unsigned_thinking_from_message(message: &mut Value) {
1231    let Value::Object(message) = message else {
1232        return;
1233    };
1234    let Some(Value::Array(blocks)) = message.get("content") else {
1235        return;
1236    };
1237    if !blocks.iter().any(is_unsigned_thinking_block) {
1238        return;
1239    }
1240    let Some(Value::Array(blocks)) = message.get_mut("content") else {
1241        return;
1242    };
1243    blocks.retain(|block| !is_unsigned_thinking_block(block));
1244    if blocks.is_empty() {
1245        message.insert("content".to_string(), Value::String(String::new()));
1246    }
1247}
1248
1249// A thinking block is unsigned when `signature` is absent or empty.
1250fn is_unsigned_thinking_block(block: &Value) -> bool {
1251    if block.get("type").and_then(Value::as_str) != Some("thinking") {
1252        return false;
1253    }
1254    !matches!(
1255        block.get("signature").and_then(Value::as_str),
1256        Some(signature) if !signature.is_empty()
1257    )
1258}
1259
1260// Forces the target's configured reasoning effort onto the outbound body, replacing the
1261// caller's value. Unlike `extra_body`, this is an override: a route that sends one model at a
1262// higher effort than the client asked for is the point of the setting.
1263fn apply_reasoning_effort(body: &mut Value, backend: &Backend) {
1264    let Some(effort) = backend.reasoning_effort() else {
1265        return;
1266    };
1267    let Value::Object(object) = body else {
1268        return;
1269    };
1270    match backend {
1271        Backend::OpenAiResponses(_) => {
1272            // Responses nests effort under `reasoning` next to fields the caller may have set
1273            // (`summary`, for example), so only the `effort` key is replaced. A `reasoning`
1274            // value that is not an object is malformed and is replaced whole.
1275            let reasoning = object
1276                .entry("reasoning".to_string())
1277                .or_insert_with(|| Value::Object(serde_json::Map::new()));
1278            if !reasoning.is_object() {
1279                *reasoning = Value::Object(serde_json::Map::new());
1280            }
1281            if let Value::Object(reasoning) = reasoning {
1282                reasoning.insert("effort".to_string(), Value::String(effort.to_string()));
1283            }
1284        }
1285        Backend::OpenAiChat(_) => {
1286            // Chat Completions takes effort as a top-level field.
1287            object.insert(
1288                "reasoning_effort".to_string(),
1289                Value::String(effort.to_string()),
1290            );
1291        }
1292        // Anthropic has no effort field (thinking is a token budget); the runner rejects the
1293        // setting on Anthropic clients at load time, so this arm is unreachable in practice.
1294        Backend::Anthropic(_) => {}
1295    }
1296}
1297
1298// Applies target defaults without overriding fields supplied by the caller.
1299fn merge_extra_body(body: &mut Value, extra_body: &BTreeMap<String, Value>) {
1300    let Value::Object(object) = body else {
1301        return;
1302    };
1303    for (key, value) in extra_body {
1304        object.entry(key.clone()).or_insert_with(|| value.clone());
1305    }
1306}
1307
1308fn omit_configured_body_fields(body: &mut Value, omit_body_fields: &BTreeSet<String>) {
1309    let Value::Object(object) = body else {
1310        return;
1311    };
1312    for key in omit_body_fields {
1313        object.remove(key);
1314    }
1315}
1316
1317// Anthropic and Bedrock both cap a request at four blocks carrying
1318// `cache_control`, counting tools, system blocks and message blocks together.
1319const MAX_CACHE_CONTROL_BLOCKS: usize = 4;
1320
1321// Counts the blocks upstream will see carrying a `cache_control` marker. The
1322// marker's own value holds no such key, so it is never counted twice.
1323fn count_cache_control_blocks(value: &Value) -> usize {
1324    match value {
1325        Value::Object(map) => {
1326            usize::from(map.contains_key("cache_control"))
1327                + map.values().map(count_cache_control_blocks).sum::<usize>()
1328        }
1329        Value::Array(items) => items.iter().map(count_cache_control_blocks).sum(),
1330        _ => 0,
1331    }
1332}
1333
1334// Marks the final message content block as the Anthropic prompt-cache breakpoint.
1335fn enable_anthropic_prompt_caching(body: &mut Value) {
1336    // Top-level cache control manages the final breakpoint and its TTL.
1337    if body.get("cache_control").is_some() {
1338        return;
1339    }
1340    // Abstain once the caller has spent the budget itself. Adding a fifth marker
1341    // turns a request that was valid on arrival into an upstream HTTP 400, and a
1342    // caller that placed four breakpoints deliberately needs them more than we
1343    // need a fifth. When the final block is already marked this returns early
1344    // and changes nothing, which is what the insert below would have done.
1345    if count_cache_control_blocks(body) >= MAX_CACHE_CONTROL_BLOCKS {
1346        return;
1347    }
1348    let Some(content) = body
1349        .get_mut("messages")
1350        .and_then(Value::as_array_mut)
1351        .and_then(|messages| messages.last_mut())
1352        .and_then(|message| message.get_mut("content"))
1353    else {
1354        return;
1355    };
1356    match content {
1357        Value::String(text) => {
1358            *content = serde_json::json!([{
1359                "type": "text",
1360                "text": std::mem::take(text),
1361                "cache_control": {"type": "ephemeral"}
1362            }]);
1363        }
1364        Value::Array(blocks) => {
1365            if let Some(block) = blocks.last_mut().and_then(Value::as_object_mut) {
1366                block
1367                    .entry("cache_control".to_string())
1368                    .or_insert_with(|| serde_json::json!({"type": "ephemeral"}));
1369            }
1370        }
1371        _ => {}
1372    }
1373}
1374
1375// Requests streamed Chat usage by default while preserving an explicit caller choice.
1376fn ensure_openai_stream_usage(body: &mut Value) {
1377    let Value::Object(object) = body else {
1378        return;
1379    };
1380    if object.get("stream").and_then(Value::as_bool) != Some(true) {
1381        return;
1382    }
1383
1384    match object.get_mut("stream_options") {
1385        Some(Value::Object(options)) => {
1386            options
1387                .entry("include_usage".to_string())
1388                .or_insert(Value::Bool(true));
1389        }
1390        _ => {
1391            let mut options = Map::new();
1392            options.insert("include_usage".to_string(), Value::Bool(true));
1393            object.insert("stream_options".to_string(), Value::Object(options));
1394        }
1395    }
1396}
1397
1398fn is_allowed_metadata_header(name: &str) -> bool {
1399    ALLOWED_METADATA_HEADERS
1400        .iter()
1401        .any(|allowed| name.eq_ignore_ascii_case(allowed))
1402}
1403
1404fn is_non_forwardable_header(name: &str, headers: &HeaderMap) -> bool {
1405    NON_FORWARDABLE_HEADERS
1406        .iter()
1407        .any(|blocked| name.eq_ignore_ascii_case(blocked))
1408        || headers.get_all("connection").iter().any(|value| {
1409            value
1410                .as_bytes()
1411                .split(|byte| *byte == b',')
1412                .any(|option| option.trim_ascii().eq_ignore_ascii_case(name.as_bytes()))
1413        })
1414}
1415
1416// Unknown application headers can carry credentials when auth forwarding is enabled.
1417fn redact_forwarded_headers(
1418    mut body: String,
1419    metadata: Option<&Metadata>,
1420    is_forwarding_auth: bool,
1421) -> String {
1422    if !is_forwarding_auth {
1423        return body;
1424    }
1425    let Some(headers) = metadata.and_then(|metadata| metadata.http_headers.as_ref()) else {
1426        return body;
1427    };
1428    for (name, value) in headers {
1429        if is_non_forwardable_header(name.as_str(), headers) {
1430            continue;
1431        }
1432        let Ok(value) = value.to_str() else {
1433            continue;
1434        };
1435        if !value.is_empty() {
1436            body = body.replace(value, "[REDACTED]");
1437        }
1438    }
1439    body
1440}
1441
1442#[cfg(test)]
1443mod tests {
1444    use std::collections::BTreeMap;
1445    use std::error::Error;
1446    use std::io::{Read, Write};
1447    use std::sync::Arc;
1448    use std::sync::atomic::{AtomicUsize, Ordering};
1449    use std::thread::JoinHandle;
1450
1451    use serde_json::json;
1452    use switchyard_protocol::{completion_text, text_request};
1453    use wiremock::matchers::{method, path};
1454    use wiremock::{Mock, MockServer, ResponseTemplate};
1455
1456    use super::*;
1457    use crate::backend::HttpBackendConfig;
1458
1459    fn config(base_url: &str) -> HttpBackendConfig {
1460        HttpBackendConfig {
1461            base_url: base_url.to_string(),
1462            api_key: Some("secret".to_string()),
1463            forward_auth: false,
1464            extra_headers: BTreeMap::new(),
1465            extra_body: BTreeMap::new(),
1466            omit_body_fields: BTreeSet::new(),
1467            reasoning_effort: None,
1468            max_retries: 0,
1469            failure_cooldown: Duration::ZERO,
1470            timeout: None,
1471        }
1472    }
1473
1474    #[test]
1475    fn provider_endpoint_urls_preserve_query_and_fragment() {
1476        let invalid_url = "not a URL/";
1477        let backend = Backend::OpenAiResponses(config(invalid_url));
1478        assert!(openai_url(invalid_url, "/responses").is_err());
1479        assert_eq!(backend.url(), invalid_url);
1480        assert_eq!(
1481            AuxiliaryOperation::ResponsesCompact.url(&backend),
1482            invalid_url
1483        );
1484
1485        let anthropic = Backend::Anthropic(config(invalid_url));
1486        assert_eq!(anthropic.url(), invalid_url);
1487        assert_eq!(anthropic.count_tokens_url(), invalid_url);
1488
1489        for path in ["", "/v1/", "/v1/messages"] {
1490            for tail in [
1491                "",
1492                "?api-version=2026-09-01&value=%2F/",
1493                "#fragment?value=/",
1494            ] {
1495                let base_url = format!("https://example.com{path}{tail}");
1496                let backend = Backend::Anthropic(config(&base_url));
1497                assert_eq!(
1498                    backend.url(),
1499                    format!("https://example.com/v1/messages{tail}")
1500                );
1501                assert_eq!(
1502                    AuxiliaryOperation::AnthropicCountTokens.url(&backend),
1503                    format!("https://example.com/v1/messages/count_tokens{tail}")
1504                );
1505            }
1506        }
1507
1508        for path in ["/v1", "/v1/responses/", "/v1/chat/completions"] {
1509            for tail in [
1510                "",
1511                "?api-version=2026-09-01&value=%2F/",
1512                "#fragment?value=/",
1513            ] {
1514                let base_url = format!("https://example.com{path}{tail}");
1515                let responses = Backend::OpenAiResponses(config(&base_url));
1516                let chat = Backend::OpenAiChat(config(&base_url));
1517                for (actual, endpoint) in [
1518                    (responses.url(), "/responses"),
1519                    (chat.url(), "/chat/completions"),
1520                    (
1521                        AuxiliaryOperation::ResponsesCompact.url(&responses),
1522                        "/responses/compact",
1523                    ),
1524                    (
1525                        AuxiliaryOperation::ResponsesInputTokens.url(&responses),
1526                        "/responses/input_tokens",
1527                    ),
1528                ] {
1529                    assert_eq!(actual, format!("https://example.com/v1{endpoint}{tail}"));
1530                }
1531            }
1532        }
1533    }
1534
1535    fn config_with_retries(base_url: &str, max_retries: u32) -> HttpBackendConfig {
1536        HttpBackendConfig {
1537            max_retries,
1538            ..config(base_url)
1539        }
1540    }
1541
1542    fn forwarding_config(base_url: &str) -> HttpBackendConfig {
1543        HttpBackendConfig {
1544            api_key: None,
1545            forward_auth: true,
1546            ..config(base_url)
1547        }
1548    }
1549
1550    // A one-model config list: "gpt" served over OpenAI Chat at base_url.
1551    fn chat_map(base_url: &str) -> Vec<ModelConfig> {
1552        vec![ModelConfig::new(
1553            "gpt",
1554            Backend::OpenAiChat(config(base_url)),
1555            None,
1556        )]
1557    }
1558
1559    fn responses_map(base_url: &str) -> Vec<ModelConfig> {
1560        vec![ModelConfig::new(
1561            "gpt",
1562            Backend::OpenAiResponses(config(base_url)),
1563            None,
1564        )]
1565    }
1566
1567    fn chat_map_with_extra_body(
1568        base_url: &str,
1569        extra_body: BTreeMap<String, Value>,
1570    ) -> Vec<ModelConfig> {
1571        let mut backend = config(base_url);
1572        backend.extra_body = extra_body;
1573        vec![ModelConfig::new("gpt", Backend::OpenAiChat(backend), None)]
1574    }
1575
1576    fn chat_map_with_omit_body_fields(
1577        base_url: &str,
1578        omit_body_fields: BTreeSet<String>,
1579    ) -> Vec<ModelConfig> {
1580        let mut backend = config(base_url);
1581        backend.omit_body_fields = omit_body_fields;
1582        vec![ModelConfig::new("gpt", Backend::OpenAiChat(backend), None)]
1583    }
1584
1585    fn chat_map_with_effort(base_url: &str, effort: &str) -> Vec<ModelConfig> {
1586        let mut backend = config(base_url);
1587        backend.reasoning_effort = Some(effort.to_string());
1588        vec![ModelConfig::new("gpt", Backend::OpenAiChat(backend), None)]
1589    }
1590
1591    fn responses_map_with_effort(base_url: &str, effort: &str) -> Vec<ModelConfig> {
1592        let mut backend = config(base_url);
1593        backend.reasoning_effort = Some(effort.to_string());
1594        vec![ModelConfig::new(
1595            "gpt",
1596            Backend::OpenAiResponses(backend),
1597            None,
1598        )]
1599    }
1600
1601    fn anthropic_map(base_url: &str) -> Vec<ModelConfig> {
1602        vec![ModelConfig::new(
1603            "claude",
1604            Backend::Anthropic(config(base_url)),
1605            None,
1606        )]
1607    }
1608
1609    #[tokio::test]
1610    async fn model_properties_uses_configured_auth_and_backend_path()
1611    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
1612        for (base_path, props_path) in [
1613            ("/v1", "/props"),
1614            ("/v1/", "/props"),
1615            ("/v1/chat/completions", "/props"),
1616            ("/v1/responses/", "/props"),
1617            ("/prefix/v1", "/prefix/props"),
1618            ("/prefix/v1/chat/completions", "/prefix/props"),
1619            ("/prefix/v1/responses", "/prefix/props"),
1620            ("/prefix/", "/prefix/props"),
1621        ] {
1622            let server = MockServer::start().await;
1623            Mock::given(method("GET"))
1624                .and(path(props_path))
1625                .and(wiremock::matchers::header("authorization", "Bearer secret"))
1626                .respond_with(ResponseTemplate::new(200).set_body_json(json!({
1627                    "default_generation_settings": {"n_ctx": 65_536}
1628                })))
1629                .expect(1)
1630                .mount(&server)
1631                .await;
1632            let client =
1633                TranslatingLlmClient::new(&chat_map(&format!("{}{base_path}", server.uri())))?;
1634            let properties = client
1635                .get_model_properties(
1636                    &ModelId::from("gpt"),
1637                    WireFormat::OpenAiChat,
1638                    Duration::from_secs(1),
1639                )
1640                .await?;
1641            assert_eq!(properties["default_generation_settings"]["n_ctx"], 65_536);
1642        }
1643        Ok(())
1644    }
1645
1646    fn chat_map_with_retries(base_url: &str, max_retries: u32) -> Vec<ModelConfig> {
1647        vec![ModelConfig::new(
1648            "gpt",
1649            Backend::OpenAiChat(config_with_retries(base_url, max_retries)),
1650            None,
1651        )]
1652    }
1653
1654    fn chat_success_response() -> ResponseTemplate {
1655        ResponseTemplate::new(200).set_body_json(json!({
1656            "id": "chatcmpl-1",
1657            "model": "gpt",
1658            "choices": [{
1659                "index": 0,
1660                "message": {"role": "assistant", "content": "recovered"},
1661                "finish_reason": "stop"
1662            }],
1663            "usage": {}
1664        }))
1665    }
1666
1667    fn truncated_response_server(
1668        content_type: &str,
1669        body: &str,
1670    ) -> std::io::Result<(String, JoinHandle<std::io::Result<()>>)> {
1671        response_sequence_server(vec![format!(
1672            "HTTP/1.1 200 OK\r\nContent-Type: {content_type}\r\n\
1673             Content-Length: {}\r\nConnection: close\r\n\r\n{body}",
1674            body.len() + 100
1675        )])
1676    }
1677
1678    fn response_sequence_server(
1679        responses: Vec<String>,
1680    ) -> std::io::Result<(String, JoinHandle<std::io::Result<()>>)> {
1681        let listener = std::net::TcpListener::bind("127.0.0.1:0")?;
1682        let address = listener.local_addr()?;
1683        let handle = std::thread::spawn(move || {
1684            for response in responses {
1685                let (mut stream, _) = listener.accept()?;
1686                let mut request = [0_u8; 1024];
1687                if stream.read(&mut request)? == 0 {
1688                    return Err(std::io::Error::new(
1689                        std::io::ErrorKind::UnexpectedEof,
1690                        "client closed before sending a request",
1691                    ));
1692                }
1693                stream.write_all(response.as_bytes())?;
1694            }
1695            Ok(())
1696        });
1697        Ok((format!("http://{address}/v1"), handle))
1698    }
1699
1700    fn raw_chat_success_response() -> String {
1701        let body = r#"{"id":"chatcmpl-1","model":"gpt","choices":[{"index":0,"message":{"role":"assistant","content":"recovered"},"finish_reason":"stop"}],"usage":{}}"#;
1702        format!(
1703            "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n\
1704             Content-Length: {}\r\nConnection: close\r\n\r\n{body}",
1705            body.len()
1706        )
1707    }
1708
1709    fn request_for(model: Option<&str>, stream: bool) -> Request {
1710        let mut llm_request = text_request(model.map(str::to_string), "hi");
1711        llm_request.stream = stream;
1712        Request {
1713            llm_request,
1714            raw_request: None,
1715            metadata: None,
1716        }
1717    }
1718
1719    fn request_with_headers(model: &str, headers: HeaderMap) -> Request {
1720        let mut request = request_for(Some(model), false);
1721        request.metadata = Some(Metadata {
1722            http_headers: Some(headers),
1723            ..Default::default()
1724        });
1725        request
1726    }
1727
1728    #[tokio::test]
1729    async fn transport_errors_drop_the_upstream_url() {
1730        let error = reqwest::Client::new()
1731            .post("http://127.0.0.1:1/v1?key=CANARY")
1732            .send()
1733            .await
1734            .expect_err("closed port");
1735
1736        assert!(!convert_reqwest_error(error).to_string().contains("CANARY"));
1737    }
1738
1739    #[test]
1740    fn nvidia_anthropic_tools_drop_unsupported_strict_field() {
1741        let backend = Backend::Anthropic(config("https://inference-api.nvidia.com"));
1742        assert!(nvidia_inference_api_backend(&backend));
1743        assert!(!nvidia_inference_api_backend(&Backend::Anthropic(config(
1744            "https://inference-api.nvidia.com.example.test"
1745        ))));
1746        let mut body = json!({"tools": [{"name": "a", "input_schema": {}, "strict": true}]});
1747        strip_anthropic_tool_strict(&mut body);
1748        assert!(body["tools"][0].get("strict").is_none());
1749    }
1750
1751    #[test]
1752    fn anthropic_prompt_caching_marks_final_message() {
1753        let caller_managed = json!({
1754            "cache_control": {"type": "ephemeral", "ttl": "1h"},
1755            "messages": [{"role": "user", "content": "hello"}]
1756        });
1757        let mut body = caller_managed.clone();
1758        enable_anthropic_prompt_caching(&mut body);
1759        assert_eq!(body, caller_managed);
1760
1761        let mut body = json!({
1762            "messages": [{"role": "user", "content": "hello"}]
1763        });
1764
1765        enable_anthropic_prompt_caching(&mut body);
1766
1767        assert_eq!(
1768            body["messages"][0]["content"][0]["cache_control"],
1769            json!({"type": "ephemeral"})
1770        );
1771    }
1772
1773    // Counts the blocks upstream will see carrying a `cache_control` marker.
1774    fn cache_control_blocks(value: &Value) -> usize {
1775        count_cache_control_blocks(value)
1776    }
1777
1778    // A body whose four breakpoints are all spent elsewhere: the caller manages
1779    // its own caching and left the final block unmarked on purpose.
1780    fn body_at_the_cache_control_limit() -> Value {
1781        json!({
1782            "system": [
1783                {"type": "text", "text": "a", "cache_control": {"type": "ephemeral"}},
1784                {"type": "text", "text": "b", "cache_control": {"type": "ephemeral"}}
1785            ],
1786            "tools": [
1787                {"name": "t", "input_schema": {"type": "object"},
1788                 "cache_control": {"type": "ephemeral"}}
1789            ],
1790            "messages": [
1791                {"role": "user", "content": [
1792                    {"type": "text", "text": "old", "cache_control": {"type": "ephemeral"}}
1793                ]},
1794                {"role": "assistant", "content": [{"type": "text", "text": "ok"}]},
1795                {"role": "user", "content": [{"type": "text", "text": "new"}]}
1796            ]
1797        })
1798    }
1799
1800    #[test]
1801    fn anthropic_prompt_caching_abstains_at_the_cache_control_limit() {
1802        let mut body = body_at_the_cache_control_limit();
1803        assert_eq!(cache_control_blocks(&body), 4);
1804
1805        enable_anthropic_prompt_caching(&mut body);
1806
1807        // Anthropic and Bedrock both reject a fifth marker with HTTP 400, which
1808        // would fail a request that was valid before it reached us.
1809        assert_eq!(cache_control_blocks(&body), 4);
1810        assert!(
1811            body["messages"][2]["content"][0]
1812                .get("cache_control")
1813                .is_none()
1814        );
1815    }
1816
1817    #[test]
1818    fn anthropic_prompt_caching_still_marks_one_below_the_limit() {
1819        let mut body = body_at_the_cache_control_limit();
1820        // Free one breakpoint, so there is room for ours.
1821        body["system"].as_array_mut().unwrap().pop();
1822        assert_eq!(cache_control_blocks(&body), 3);
1823
1824        enable_anthropic_prompt_caching(&mut body);
1825
1826        assert_eq!(cache_control_blocks(&body), 4);
1827        assert_eq!(
1828            body["messages"][2]["content"][0]["cache_control"],
1829            json!({"type": "ephemeral"})
1830        );
1831    }
1832
1833    // A request that pins `format` in its metadata, so the client resolves that
1834    // wire format instead of the model's default backend.
1835    fn request_with_wire_format(model: &str, format: WireFormat) -> Request {
1836        let mut request = request_for(Some(model), false);
1837        request.metadata = Some(Metadata {
1838            session_id: None,
1839            agent_id: None,
1840            task_id: None,
1841            correlation_id: None,
1842            extra_metadata: None,
1843            http_headers: None,
1844            wire_format: Some(format),
1845            ..Default::default()
1846        });
1847        request
1848    }
1849
1850    #[tokio::test]
1851    async fn missing_model_errors()
1852    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
1853        let client = TranslatingLlmClient::new(&[])?;
1854        let Err(error) = client
1855            .call_rewrite_model(request_for(None, false), None)
1856            .await
1857        else {
1858            panic!("expected an error");
1859        };
1860        assert!(matches!(
1861            error,
1862            LlmClientError::InvalidRequest { message } if message == "no model given"
1863        ));
1864        Ok(())
1865    }
1866
1867    #[tokio::test]
1868    async fn unknown_model_errors()
1869    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
1870        let client = TranslatingLlmClient::new(&[])?;
1871        let Err(error) = client
1872            .call_rewrite_model(request_for(Some("gpt"), false), None)
1873            .await
1874        else {
1875            panic!("expected an error");
1876        };
1877        assert!(matches!(
1878            error,
1879            LlmClientError::Configuration { message }
1880                if message.contains("gpt")
1881        ));
1882        Ok(())
1883    }
1884
1885    #[tokio::test]
1886    async fn unknown_model_format_errors()
1887    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
1888        // "gpt" exists but only over OpenAI Chat; the request pins Anthropic.
1889        let client = TranslatingLlmClient::new(&chat_map("https://example.test/v1"))?;
1890        let Err(error) = client
1891            .call_rewrite_model(
1892                request_with_wire_format("gpt", WireFormat::AnthropicMessages),
1893                None,
1894            )
1895            .await
1896        else {
1897            panic!("expected an error");
1898        };
1899        assert!(matches!(
1900            error,
1901            LlmClientError::Configuration { message }
1902                if message.contains("gpt")
1903                    && message.contains(&WireFormat::AnthropicMessages.to_string())
1904        ));
1905        Ok(())
1906    }
1907
1908    #[test]
1909    fn backend_for_resolves_configured_format()
1910    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
1911        let client = TranslatingLlmClient::new(&chat_map("https://example.test/v1"))?;
1912        // "gpt" is served over OpenAI Chat only; other formats and models miss.
1913        assert!(
1914            client
1915                .backend_for(&ModelId::from("gpt"), WireFormat::OpenAiChat)
1916                .is_some()
1917        );
1918        assert!(
1919            client
1920                .backend_for(&ModelId::from("gpt"), WireFormat::AnthropicMessages)
1921                .is_none()
1922        );
1923        assert!(
1924            client
1925                .backend_for(&ModelId::from("missing"), WireFormat::OpenAiChat)
1926                .is_none()
1927        );
1928        Ok(())
1929    }
1930
1931    #[tokio::test]
1932    async fn model_name_arg_wins_over_request_model()
1933    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
1934        let client = TranslatingLlmClient::new(&[])?;
1935        // Arg "b" is looked up (and reported), not the request's "a".
1936        let Err(error) = client
1937            .call_rewrite_model(request_for(Some("a"), false), Some(&ModelId::from("b")))
1938            .await
1939        else {
1940            panic!("expected an error");
1941        };
1942        assert!(matches!(
1943            error,
1944            LlmClientError::Configuration { message }
1945                if message.contains("\"b\"")
1946        ));
1947        Ok(())
1948    }
1949
1950    #[tokio::test]
1951    async fn buffered_openai_chat_round_trips()
1952    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
1953        let server = MockServer::start().await;
1954        Mock::given(method("POST"))
1955            .and(path("/v1/chat/completions"))
1956            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
1957                "id": "chatcmpl-1",
1958                "model": "gpt",
1959                "choices": [{
1960                    "index": 0,
1961                    "message": {"role": "assistant", "content": "Hi there"},
1962                    "finish_reason": "stop"
1963                }],
1964                "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}
1965            })))
1966            .mount(&server)
1967            .await;
1968
1969        let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
1970
1971        let response = client
1972            .call_rewrite_model(request_for(Some("gpt"), false), None)
1973            .await?;
1974        let agg = response.llm_response.into_agg().await?;
1975        assert_eq!(completion_text(&agg), "Hi there");
1976
1977        Ok(())
1978    }
1979
1980    #[tokio::test]
1981    async fn invalid_json_is_a_response_translation_error()
1982    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
1983        let server = MockServer::start().await;
1984        let calls = Arc::new(AtomicUsize::new(0));
1985        let observed_calls = Arc::clone(&calls);
1986        Mock::given(method("POST"))
1987            .respond_with(move |_: &wiremock::Request| {
1988                observed_calls.fetch_add(1, Ordering::SeqCst);
1989                ResponseTemplate::new(200).set_body_raw("not json", "application/json")
1990            })
1991            .mount(&server)
1992            .await;
1993
1994        let client =
1995            TranslatingLlmClient::new(&chat_map_with_retries(&format!("{}/v1", server.uri()), 2))?;
1996        let Err(error) = client
1997            .call_rewrite_model(request_for(Some("gpt"), false), None)
1998            .await
1999        else {
2000            panic!("expected invalid JSON to fail");
2001        };
2002
2003        assert!(matches!(
2004            error,
2005            LlmClientError::ResponseTranslation(message)
2006                if message.contains("invalid upstream JSON")
2007        ));
2008        assert_eq!(calls.load(Ordering::SeqCst), 1);
2009        Ok(())
2010    }
2011
2012    #[tokio::test]
2013    async fn response_body_io_failure_is_a_transport_error()
2014    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2015        let (base_url, server) = truncated_response_server("application/json", "{}")?;
2016        let client = TranslatingLlmClient::new(&chat_map(&base_url))?;
2017        let result = client
2018            .call_rewrite_model(request_for(Some("gpt"), false), None)
2019            .await;
2020        server
2021            .join()
2022            .map_err(|_| std::io::Error::other("response server thread panicked"))??;
2023
2024        let Err(error) = result else {
2025            panic!("expected the truncated response body to fail");
2026        };
2027        let LlmClientError::Transport { source } = error else {
2028            panic!("expected a transport error");
2029        };
2030        let Some(source) = source.downcast_ref::<reqwest::Error>() else {
2031            panic!("expected the reqwest transport source");
2032        };
2033        assert!(source.is_decode());
2034        assert!(
2035            !std::error::Error::source(&source)
2036                .is_some_and(|source| source.is::<serde_json::Error>())
2037        );
2038        Ok(())
2039    }
2040
2041    #[tokio::test]
2042    async fn response_body_transport_failures_are_retried()
2043    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2044        let truncated_responses = [
2045            "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n\
2046             Content-Length: 102\r\nConnection: close\r\n\r\n{}",
2047            "HTTP/1.1 401 Unauthorized\r\nContent-Type: text/plain\r\n\
2048             Content-Length: 100\r\nConnection: close\r\n\r\nbad",
2049        ];
2050
2051        for truncated in truncated_responses {
2052            let (base_url, server) =
2053                response_sequence_server(vec![truncated.to_string(), raw_chat_success_response()])?;
2054            let client = TranslatingLlmClient::new(&chat_map_with_retries(&base_url, 1))?;
2055            let response = client
2056                .call_rewrite_model(request_for(Some("gpt"), false), None)
2057                .await?;
2058            server
2059                .join()
2060                .map_err(|_| std::io::Error::other("response server thread panicked"))??;
2061
2062            assert_eq!(
2063                completion_text(&response.llm_response.into_agg().await?),
2064                "recovered"
2065            );
2066        }
2067        Ok(())
2068    }
2069
2070    #[tokio::test]
2071    async fn pre_first_event_transport_failure_uses_retry_budget()
2072    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2073        let truncated = "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\n\
2074                         Content-Length: 100\r\nConnection: close\r\n\r\n";
2075        let body = "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"recovered\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n";
2076        let success = format!(
2077            "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\n\
2078             Content-Length: {}\r\nConnection: close\r\n\r\n{body}",
2079            body.len()
2080        );
2081
2082        for second in [success, truncated.to_string()] {
2083            let is_exhausted = second == truncated;
2084            let (base_url, server) = response_sequence_server(vec![truncated.to_string(), second])?;
2085            let client = TranslatingLlmClient::new(&chat_map_with_retries(&base_url, 1))?;
2086            let result = client
2087                .call_rewrite_model(request_for(Some("gpt"), true), None)
2088                .await;
2089            if is_exhausted {
2090                assert!(matches!(result, Err(LlmClientError::Transport { .. })));
2091            } else {
2092                assert_eq!(
2093                    completion_text(&result?.llm_response.into_agg().await?),
2094                    "recovered"
2095                );
2096            }
2097            server
2098                .join()
2099                .map_err(|_| std::io::Error::other("response server thread panicked"))??;
2100        }
2101        Ok(())
2102    }
2103
2104    #[tokio::test]
2105    async fn streaming_body_io_failure_preserves_transport_error()
2106    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2107        let body = "data: {\"choices\":[{\"delta\":{\"content\":\"partial\"}}]}\n\n";
2108        let (base_url, server) = truncated_response_server("text/event-stream", body)?;
2109        let client = TranslatingLlmClient::new(&chat_map_with_retries(&base_url, 2))?;
2110        let response = client
2111            .call_rewrite_model(request_for(Some("gpt"), true), None)
2112            .await?;
2113        let result = response.llm_response.into_agg().await;
2114        server
2115            .join()
2116            .map_err(|_| std::io::Error::other("response server thread panicked"))??;
2117
2118        let Err(error) = result else {
2119            panic!("expected the truncated stream body to fail");
2120        };
2121
2122        assert!(matches!(error, LlmClientError::Transport { .. }));
2123        Ok(())
2124    }
2125
2126    #[tokio::test]
2127    async fn rewrites_model_to_resolved_upstream_id()
2128    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2129        // Inbound body says "switchyard"; the upstream must receive "gpt".
2130        let server = MockServer::start().await;
2131        Mock::given(method("POST"))
2132            .and(path("/v1/chat/completions"))
2133            .and(wiremock::matchers::body_partial_json(json!({"model": "gpt"})))
2134            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
2135                "id": "1", "model": "gpt",
2136                "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
2137                "usage": {}
2138            })))
2139            .mount(&server)
2140            .await;
2141
2142        let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
2143        // Inbound model differs from the map key / resolved model.
2144        client
2145            .call_rewrite_model(
2146                request_for(Some("switchyard"), false),
2147                Some(&ModelId::from("gpt")),
2148            )
2149            .await?;
2150        // The body_partial_json matcher asserts the upstream saw model "gpt".
2151        Ok(())
2152    }
2153
2154    /// A configured reasoning effort replaces the caller's value on both OpenAI wire formats,
2155    /// which `extra_body` (defaults only) cannot do.
2156    #[tokio::test]
2157    async fn reasoning_effort_override_replaces_the_callers_effort()
2158    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2159        let server = MockServer::start().await;
2160        Mock::given(method("POST"))
2161            .and(path("/v1/chat/completions"))
2162            .and(wiremock::matchers::body_partial_json(json!({
2163                "model": "gpt",
2164                "reasoning_effort": "max"
2165            })))
2166            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
2167                "id": "1",
2168                "model": "gpt",
2169                "choices": [{
2170                    "index": 0,
2171                    "message": {"role": "assistant", "content": "ok"},
2172                    "finish_reason": "stop"
2173                }],
2174                "usage": {}
2175            })))
2176            .mount(&server)
2177            .await;
2178        let client = TranslatingLlmClient::new(&chat_map_with_effort(
2179            &format!("{}/v1", server.uri()),
2180            "max",
2181        ))?;
2182        client
2183            .call_rewrite_model_raw(
2184                json!({
2185                    "model": "client-facing",
2186                    "messages": [{"role": "user", "content": "hi"}],
2187                    "reasoning_effort": "high"
2188                }),
2189                None,
2190                Some(&ModelId::from("gpt")),
2191                WireFormat::OpenAiChat,
2192            )
2193            .await?;
2194
2195        let server = MockServer::start().await;
2196        Mock::given(method("POST"))
2197            .and(path("/v1/responses"))
2198            .and(wiremock::matchers::body_partial_json(json!({
2199                "model": "gpt",
2200                "reasoning": {"effort": "max", "summary": "auto"}
2201            })))
2202            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
2203                "id": "resp_1",
2204                "object": "response",
2205                "model": "gpt",
2206                "status": "completed",
2207                "output": [{
2208                    "type": "message",
2209                    "id": "msg_1",
2210                    "role": "assistant",
2211                    "status": "completed",
2212                    "content": [{"type": "output_text", "text": "ok"}]
2213                }],
2214                "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}
2215            })))
2216            .mount(&server)
2217            .await;
2218        let client = TranslatingLlmClient::new(&responses_map_with_effort(
2219            &format!("{}/v1", server.uri()),
2220            "max",
2221        ))?;
2222        client
2223            .call_rewrite_model_raw(
2224                json!({
2225                    "model": "client-facing",
2226                    "input": [{"role": "user", "content": "hi"}],
2227                    "reasoning": {"effort": "high", "summary": "auto"}
2228                }),
2229                None,
2230                Some(&ModelId::from("gpt")),
2231                WireFormat::OpenAiResponses,
2232            )
2233            .await?;
2234        Ok(())
2235    }
2236
2237    #[tokio::test]
2238    async fn extra_body_adds_defaults_without_overriding_the_request()
2239    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2240        let server = MockServer::start().await;
2241        Mock::given(method("POST"))
2242            .and(path("/v1/chat/completions"))
2243            .and(wiremock::matchers::body_partial_json(json!({
2244                "model": "gpt",
2245                "max_tokens": 7,
2246                "service_tier": "priority"
2247            })))
2248            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
2249                "id": "1",
2250                "model": "gpt",
2251                "choices": [{
2252                    "index": 0,
2253                    "message": {"role": "assistant", "content": "ok"},
2254                    "finish_reason": "stop"
2255                }],
2256                "usage": {}
2257            })))
2258            .mount(&server)
2259            .await;
2260
2261        let extra_body = BTreeMap::from([
2262            ("max_tokens".to_string(), json!(999)),
2263            ("service_tier".to_string(), json!("priority")),
2264        ]);
2265        let client = TranslatingLlmClient::new(&chat_map_with_extra_body(
2266            &format!("{}/v1", server.uri()),
2267            extra_body,
2268        ))?;
2269        let raw = json!({
2270            "model": "client-facing",
2271            "messages": [{"role": "user", "content": "hi"}],
2272            "max_tokens": 7
2273        });
2274
2275        client
2276            .call_rewrite_model_raw(
2277                raw,
2278                None,
2279                Some(&ModelId::from("gpt")),
2280                WireFormat::OpenAiChat,
2281            )
2282            .await?;
2283        Ok(())
2284    }
2285
2286    #[tokio::test]
2287    async fn omit_body_fields_strips_a_field_extra_body_cannot_override()
2288    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2289        let server = MockServer::start().await;
2290        Mock::given(method("POST"))
2291            .and(path("/v1/chat/completions"))
2292            .and(|request: &wiremock::Request| {
2293                let body: Value = serde_json::from_slice(&request.body).unwrap_or(Value::Null);
2294                body.get("max_tokens").is_none()
2295            })
2296            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
2297                "id": "1",
2298                "model": "gpt",
2299                "choices": [{
2300                    "index": 0,
2301                    "message": {"role": "assistant", "content": "ok"},
2302                    "finish_reason": "stop"
2303                }],
2304                "usage": {}
2305            })))
2306            .mount(&server)
2307            .await;
2308
2309        let client = TranslatingLlmClient::new(&chat_map_with_omit_body_fields(
2310            &format!("{}/v1", server.uri()),
2311            BTreeSet::from(["max_tokens".to_string()]),
2312        ))?;
2313        let raw = json!({
2314            "model": "client-facing",
2315            "messages": [{"role": "user", "content": "hi"}],
2316            "max_tokens": 7
2317        });
2318
2319        client
2320            .call_rewrite_model_raw(
2321                raw,
2322                None,
2323                Some(&ModelId::from("gpt")),
2324                WireFormat::OpenAiChat,
2325            )
2326            .await?;
2327        Ok(())
2328    }
2329
2330    // A weak OpenAI-format tier emits thinking blocks with no signature. Replaying
2331    // them to Anthropic is rejected (Bedrock reports it as a SigV4 mismatch), so
2332    // the Anthropic leg must drop them while keeping signed ones.
2333    #[tokio::test]
2334    async fn anthropic_requests_drop_unsigned_thinking_blocks()
2335    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2336        let server = MockServer::start().await;
2337        Mock::given(method("POST"))
2338            .and(path("/v1/messages"))
2339            .and(|request: &wiremock::Request| {
2340                let body: Value = serde_json::from_slice(&request.body).unwrap_or(Value::Null);
2341                let messages = body.get("messages").and_then(Value::as_array).cloned();
2342                let Some(messages) = messages else {
2343                    return false;
2344                };
2345                // The unsigned block is gone, the signed one survives, and the
2346                // message whose only block was unsigned is not left with an empty
2347                // content array.
2348                let blocks: Vec<&Value> = messages
2349                    .iter()
2350                    .filter_map(|message| message.get("content"))
2351                    .filter_map(Value::as_array)
2352                    .flatten()
2353                    .collect();
2354                let thinking: Vec<&&Value> = blocks
2355                    .iter()
2356                    .filter(|block| block.get("type").and_then(Value::as_str) == Some("thinking"))
2357                    .collect();
2358                thinking.len() == 1
2359                    && thinking[0].get("signature").and_then(Value::as_str) == Some("sig-abc")
2360                    && messages
2361                        .iter()
2362                        .all(|message| message.get("content") != Some(&json!([])))
2363            })
2364            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
2365                "id": "msg_1",
2366                "type": "message",
2367                "role": "assistant",
2368                "model": "claude",
2369                "content": [{"type": "text", "text": "ok"}],
2370                "stop_reason": "end_turn",
2371                "usage": {"input_tokens": 1, "output_tokens": 1}
2372            })))
2373            .mount(&server)
2374            .await;
2375
2376        let client = TranslatingLlmClient::new(&anthropic_map(&server.uri()))?;
2377        let raw = json!({
2378            "model": "client-facing",
2379            "max_tokens": 7,
2380            "messages": [
2381                {"role": "user", "content": "fix the build"},
2382                {"role": "assistant", "content": [
2383                    {"type": "thinking", "thinking": "weak tier reasoning", "signature": ""}
2384                ]},
2385                {"role": "assistant", "content": [
2386                    {"type": "thinking", "thinking": "signed reasoning", "signature": "sig-abc"},
2387                    {"type": "text", "text": "here goes"}
2388                ]},
2389                {"role": "user", "content": "continue"}
2390            ]
2391        });
2392
2393        client
2394            .call_rewrite_model_raw(
2395                raw,
2396                None,
2397                Some(&ModelId::from("claude")),
2398                WireFormat::AnthropicMessages,
2399            )
2400            .await?;
2401        Ok(())
2402    }
2403
2404    #[tokio::test]
2405    async fn nvidia_anthropic_requests_strip_strict_from_extra_body()
2406    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2407        let server = MockServer::start().await;
2408        Mock::given(method("POST"))
2409            .and(path("/v1/messages"))
2410            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
2411                "id": "msg_1",
2412                "type": "message",
2413                "role": "assistant",
2414                "model": "claude",
2415                "content": [{"type": "text", "text": "ok"}],
2416                "stop_reason": "end_turn",
2417                "usage": {"input_tokens": 1, "output_tokens": 1}
2418            })))
2419            .expect(1)
2420            .mount(&server)
2421            .await;
2422
2423        let mut backend = config(&format!(
2424            "http://inference-api.nvidia.com:{}",
2425            server.address().port()
2426        ));
2427        backend.extra_body.insert(
2428            "tools".into(),
2429            json!([{
2430                "name": "lookup",
2431                "input_schema": {"type": "object", "properties": {}},
2432                "strict": true
2433            }]),
2434        );
2435        let mut client = TranslatingLlmClient::new(&[ModelConfig::new(
2436            "claude",
2437            Backend::Anthropic(backend),
2438            None,
2439        )])?;
2440        client.client = reqwest::Client::builder()
2441            .no_proxy()
2442            .resolve("inference-api.nvidia.com", *server.address())
2443            .build()?;
2444        client
2445            .call_rewrite_model_raw(
2446                json!({
2447                    "model": "claude",
2448                    "max_tokens": 7,
2449                    "messages": [{"role": "user", "content": "hi"}]
2450                }),
2451                None,
2452                Some(&ModelId::from("claude")),
2453                WireFormat::AnthropicMessages,
2454            )
2455            .await?;
2456
2457        let requests = server.received_requests().await.unwrap();
2458        let body: Value = serde_json::from_slice(&requests[0].body)?;
2459        assert_eq!(body["tools"][0]["name"], "lookup");
2460        assert!(body["tools"][0].get("strict").is_none());
2461        Ok(())
2462    }
2463
2464    // A router can serve earlier turns from an OpenAI target and later turns from
2465    // an Anthropic one, so the Anthropic leg must drop OpenAI-only fields the
2466    // caller keeps sending or the upstream rejects the whole request.
2467    #[tokio::test]
2468    async fn anthropic_requests_drop_openai_only_fields()
2469    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2470        let server = MockServer::start().await;
2471        Mock::given(method("POST"))
2472            .and(path("/v1/messages"))
2473            .and(|request: &wiremock::Request| {
2474                let body: Value = serde_json::from_slice(&request.body).unwrap_or(Value::Null);
2475                body.get("context_management").is_none() && body.get("reasoning_effort").is_none()
2476            })
2477            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
2478                "id": "msg_1",
2479                "type": "message",
2480                "role": "assistant",
2481                "model": "claude",
2482                "content": [{"type": "text", "text": "ok"}],
2483                "stop_reason": "end_turn",
2484                "usage": {"input_tokens": 1, "output_tokens": 1}
2485            })))
2486            .mount(&server)
2487            .await;
2488
2489        let client = TranslatingLlmClient::new(&anthropic_map(&server.uri()))?;
2490        let raw = json!({
2491            "model": "client-facing",
2492            "messages": [{"role": "user", "content": "hi"}],
2493            "max_tokens": 7,
2494            "reasoning_effort": "high",
2495            "context_management": {
2496                "edits": [{"type": "clear_thinking_20251015"}]
2497            }
2498        });
2499
2500        client
2501            .call_rewrite_model_raw(
2502                raw,
2503                None,
2504                Some(&ModelId::from("claude")),
2505                WireFormat::AnthropicMessages,
2506            )
2507            .await?;
2508        Ok(())
2509    }
2510
2511    #[tokio::test]
2512    async fn streaming_openai_chat_aggregates()
2513    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2514        let server = MockServer::start().await;
2515        let sse = "data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}\n\n\
2516             data: {\"choices\":[{\"delta\":{\"content\":\" world\"}}]}\n\n\
2517             data: {\"choices\":[],\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":2,\"total_tokens\":3}}\n\n\
2518             data: [DONE]\n\n";
2519        Mock::given(method("POST"))
2520            .and(path("/v1/chat/completions"))
2521            .and(wiremock::matchers::body_partial_json(json!({
2522                "stream": true,
2523                "stream_options": {"include_usage": true}
2524            })))
2525            .respond_with(ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream"))
2526            .mount(&server)
2527            .await;
2528
2529        let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
2530
2531        let response = client
2532            .call_rewrite_model(request_for(Some("gpt"), true), None)
2533            .await?;
2534        assert!(matches!(response.llm_response, LlmResponse::Stream(_)));
2535        let agg = response.llm_response.into_agg().await?;
2536        assert_eq!(completion_text(&agg), "Hello world");
2537        assert_eq!(agg.usage.input_tokens, Some(1));
2538        assert_eq!(agg.usage.output_tokens, Some(2));
2539        assert_eq!(agg.usage.total_tokens, Some(3));
2540        Ok(())
2541    }
2542
2543    #[tokio::test]
2544    async fn streaming_openai_chat_preserves_usage_opt_out()
2545    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2546        let server = MockServer::start().await;
2547        Mock::given(method("POST"))
2548            .and(path("/v1/chat/completions"))
2549            .and(wiremock::matchers::body_partial_json(json!({
2550                "stream": true,
2551                "stream_options": {"include_usage": false}
2552            })))
2553            .respond_with(
2554                ResponseTemplate::new(200).set_body_raw("data: [DONE]\n\n", "text/event-stream"),
2555            )
2556            .mount(&server)
2557            .await;
2558
2559        let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
2560        let raw = json!({
2561            "model": "client-facing",
2562            "messages": [{"role": "user", "content": "hi"}],
2563            "stream": true,
2564            "stream_options": {"include_usage": false}
2565        });
2566
2567        let response = client
2568            .call_rewrite_model_raw(
2569                raw,
2570                None,
2571                Some(&ModelId::from("gpt")),
2572                WireFormat::OpenAiChat,
2573            )
2574            .await?;
2575        assert!(matches!(response, RawResponse::Stream(_)));
2576        Ok(())
2577    }
2578
2579    #[tokio::test]
2580    async fn upstream_500_is_upstream_http()
2581    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2582        let server = MockServer::start().await;
2583        Mock::given(method("POST"))
2584            .respond_with(ResponseTemplate::new(500).set_body_string("boom"))
2585            .mount(&server)
2586            .await;
2587
2588        let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
2589
2590        let Err(error) = client
2591            .call_rewrite_model(request_for(Some("gpt"), false), None)
2592            .await
2593        else {
2594            panic!("expected an error");
2595        };
2596        assert!(matches!(
2597            error,
2598            LlmClientError::UpstreamHttp {
2599                status: StatusCode::INTERNAL_SERVER_ERROR,
2600                ..
2601            }
2602        ));
2603        Ok(())
2604    }
2605
2606    #[tokio::test]
2607    async fn retryable_http_failure_recovers_within_budget()
2608    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2609        let server = MockServer::start().await;
2610        let calls = Arc::new(AtomicUsize::new(0));
2611        let observed_calls = Arc::clone(&calls);
2612        Mock::given(method("POST"))
2613            .respond_with(move |_: &wiremock::Request| {
2614                if observed_calls.fetch_add(1, Ordering::SeqCst) == 0 {
2615                    ResponseTemplate::new(503)
2616                        .insert_header("retry-after", "0")
2617                        .set_body_string("temporarily unavailable")
2618                } else {
2619                    chat_success_response()
2620                }
2621            })
2622            .mount(&server)
2623            .await;
2624
2625        let client =
2626            TranslatingLlmClient::new(&chat_map_with_retries(&format!("{}/v1", server.uri()), 1))?;
2627        let response = client
2628            .call_rewrite_model(request_for(Some("gpt"), false), None)
2629            .await?;
2630        let agg = response.llm_response.into_agg().await?;
2631
2632        assert_eq!(completion_text(&agg), "recovered");
2633        assert_eq!(calls.load(Ordering::SeqCst), 2);
2634        Ok(())
2635    }
2636
2637    #[tokio::test]
2638    async fn deterministic_http_failure_is_not_retried()
2639    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2640        let server = MockServer::start().await;
2641        let calls = Arc::new(AtomicUsize::new(0));
2642        let observed_calls = Arc::clone(&calls);
2643        Mock::given(method("POST"))
2644            .respond_with(move |_: &wiremock::Request| {
2645                observed_calls.fetch_add(1, Ordering::SeqCst);
2646                ResponseTemplate::new(401).set_body_string("invalid key")
2647            })
2648            .mount(&server)
2649            .await;
2650
2651        let client =
2652            TranslatingLlmClient::new(&chat_map_with_retries(&format!("{}/v1", server.uri()), 2))?;
2653        let Err(error) = client
2654            .call_rewrite_model(request_for(Some("gpt"), false), None)
2655            .await
2656        else {
2657            panic!("expected an upstream error");
2658        };
2659
2660        assert!(matches!(
2661            error,
2662            LlmClientError::UpstreamHttp {
2663                status: StatusCode::UNAUTHORIZED,
2664                body
2665            } if body == "invalid key"
2666        ));
2667        assert_eq!(calls.load(Ordering::SeqCst), 1);
2668        Ok(())
2669    }
2670
2671    #[tokio::test]
2672    async fn rate_limit_cooldown_depends_on_shared_credentials()
2673    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2674        for forward_auth in [true, false] {
2675            let server = MockServer::start().await;
2676            Mock::given(method("POST"))
2677                .respond_with(ResponseTemplate::new(429).insert_header("retry-after", "0"))
2678                .expect(if forward_auth { 6 } else { 3 })
2679                .mount(&server)
2680                .await;
2681            let backend = HttpBackendConfig {
2682                forward_auth,
2683                failure_cooldown: Duration::from_secs(60),
2684                ..config_with_retries(&format!("{}/v1", server.uri()), 2)
2685            };
2686            let client = TranslatingLlmClient::new(&[ModelConfig::new(
2687                "gpt",
2688                Backend::OpenAiChat(backend),
2689                None,
2690            )])?;
2691            for credential in ["Bearer first-user", "Bearer second-user"] {
2692                let mut headers = http::HeaderMap::new();
2693                headers.insert("authorization", http::HeaderValue::from_static(credential));
2694                let result = client
2695                    .call_rewrite_model(request_with_headers("gpt", headers), None)
2696                    .await;
2697                if !forward_auth && credential == "Bearer second-user" {
2698                    assert!(matches!(
2699                        result,
2700                        Err(LlmClientError::TemporarilyUnavailable)
2701                    ));
2702                } else {
2703                    assert!(matches!(
2704                        result,
2705                        Err(LlmClientError::UpstreamHttp {
2706                            status: StatusCode::TOO_MANY_REQUESTS,
2707                            ..
2708                        })
2709                    ));
2710                }
2711            }
2712        }
2713        Ok(())
2714    }
2715
2716    #[tokio::test]
2717    async fn retry_exhaustion_returns_the_final_upstream_error()
2718    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2719        let server = MockServer::start().await;
2720        let calls = Arc::new(AtomicUsize::new(0));
2721        let observed_calls = Arc::clone(&calls);
2722        Mock::given(method("POST"))
2723            .respond_with(move |_: &wiremock::Request| {
2724                let attempt = observed_calls.fetch_add(1, Ordering::SeqCst) + 1;
2725                ResponseTemplate::new(500)
2726                    .insert_header("retry-after", "0")
2727                    .set_body_string(format!("attempt {attempt}"))
2728            })
2729            .mount(&server)
2730            .await;
2731
2732        let client =
2733            TranslatingLlmClient::new(&chat_map_with_retries(&format!("{}/v1", server.uri()), 2))?;
2734        let Err(error) = client
2735            .call_rewrite_model(request_for(Some("gpt"), false), None)
2736            .await
2737        else {
2738            panic!("expected retry exhaustion");
2739        };
2740
2741        assert!(matches!(
2742            error,
2743            LlmClientError::UpstreamHttp {
2744                status: StatusCode::INTERNAL_SERVER_ERROR,
2745                body
2746            } if body == "attempt 3"
2747        ));
2748        assert_eq!(calls.load(Ordering::SeqCst), 3);
2749        Ok(())
2750    }
2751
2752    #[tokio::test]
2753    async fn deadline_expires_before_send_or_stream_poll()
2754    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2755        let server = MockServer::start().await;
2756        Mock::given(method("POST"))
2757            .respond_with(ResponseTemplate::new(200).set_body_raw(
2758                "data: {\"id\":\"test\",\"model\":\"gpt\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"ok\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n",
2759                "text/event-stream",
2760            ))
2761            .mount(&server)
2762            .await;
2763        for timeout in [Duration::ZERO, Duration::from_millis(200)] {
2764            let mut backend = config(&server.uri());
2765            backend.timeout = Some(timeout);
2766            let client = TranslatingLlmClient::new(&[ModelConfig::new(
2767                "gpt",
2768                Backend::OpenAiChat(backend),
2769                None,
2770            )])?;
2771            let result = client
2772                .call_rewrite_model(request_for(Some("gpt"), true), None)
2773                .await;
2774            if timeout.is_zero() {
2775                assert!(matches!(result, Err(LlmClientError::Timeout { .. })));
2776                assert!(
2777                    server
2778                        .received_requests()
2779                        .await
2780                        .expect("recorded requests")
2781                        .is_empty()
2782                );
2783            } else {
2784                let LlmResponse::Stream(mut chunks) = result?.llm_response else {
2785                    return Err("expected a streaming response".into());
2786                };
2787                tokio::time::sleep(Duration::from_millis(300)).await;
2788                assert!(matches!(
2789                    chunks.next().await,
2790                    Some(Err(LlmClientError::Timeout { .. }))
2791                ));
2792                assert!(chunks.next().await.is_none());
2793            }
2794        }
2795        Ok(())
2796    }
2797
2798    #[tokio::test]
2799    async fn timeout_is_retried_before_a_response_is_returned()
2800    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2801        let server = MockServer::start().await;
2802        let calls = Arc::new(AtomicUsize::new(0));
2803        let observed_calls = Arc::clone(&calls);
2804        Mock::given(method("POST"))
2805            .respond_with(move |_: &wiremock::Request| {
2806                if observed_calls.fetch_add(1, Ordering::SeqCst) == 0 {
2807                    ResponseTemplate::new(200).set_delay(Duration::from_millis(500))
2808                } else {
2809                    chat_success_response()
2810                }
2811            })
2812            .mount(&server)
2813            .await;
2814
2815        let mut client =
2816            TranslatingLlmClient::new(&chat_map_with_retries(&format!("{}/v1", server.uri()), 1))?;
2817        client.client = reqwest::Client::builder()
2818            .timeout(Duration::from_millis(100))
2819            .build()?;
2820        let response = client
2821            .call_rewrite_model(request_for(Some("gpt"), false), None)
2822            .await?;
2823
2824        assert_eq!(
2825            completion_text(&response.llm_response.into_agg().await?),
2826            "recovered"
2827        );
2828        assert_eq!(calls.load(Ordering::SeqCst), 2);
2829        Ok(())
2830    }
2831
2832    #[test]
2833    fn retryable_error_classes_are_explicit() {
2834        let transport = AttemptFailure {
2835            error: LlmClientError::Transport {
2836                source: std::io::Error::other("disconnected").into(),
2837            },
2838            status: None,
2839            retry_after: None,
2840        };
2841        assert!(transport.is_retryable());
2842
2843        for status in [
2844            StatusCode::REQUEST_TIMEOUT,
2845            StatusCode::TOO_MANY_REQUESTS,
2846            StatusCode::INTERNAL_SERVER_ERROR,
2847            StatusCode::SERVICE_UNAVAILABLE,
2848            StatusCode::from_u16(599)
2849                .expect("599 is a syntactically valid HTTP code in the http crate"),
2850        ] {
2851            let failure = AttemptFailure {
2852                error: LlmClientError::UpstreamHttp {
2853                    status,
2854                    body: String::new(),
2855                },
2856                status: Some(status),
2857                retry_after: None,
2858            };
2859            assert!(failure.is_retryable(), "HTTP {status} should retry");
2860        }
2861        for status in [
2862            StatusCode::BAD_REQUEST,
2863            StatusCode::UNAUTHORIZED,
2864            StatusCode::CONFLICT,
2865            StatusCode::from_u16(600)
2866                .expect("600 is a syntactically valid HTTP code in the http crate"),
2867        ] {
2868            let failure = AttemptFailure {
2869                error: LlmClientError::UpstreamHttp {
2870                    status,
2871                    body: String::new(),
2872                },
2873                status: Some(status),
2874                retry_after: None,
2875            };
2876            assert!(!failure.is_retryable(), "HTTP {status} should fail fast");
2877        }
2878
2879        let configuration = AttemptFailure {
2880            error: LlmClientError::Configuration {
2881                message: "invalid header".to_string(),
2882            },
2883            status: None,
2884            retry_after: None,
2885        };
2886        assert!(!configuration.is_retryable());
2887
2888        let context_window = AttemptFailure {
2889            error: LlmClientError::ContextWindowExceeded {
2890                model: ModelId::from("gpt"),
2891                message: "too long".to_string(),
2892            },
2893            status: Some(StatusCode::BAD_REQUEST),
2894            retry_after: None,
2895        };
2896        assert!(!context_window.is_retryable());
2897    }
2898
2899    #[test]
2900    fn retry_after_supports_seconds_and_http_dates() {
2901        let mut headers = HeaderMap::new();
2902        headers.insert(RETRY_AFTER, reqwest::header::HeaderValue::from_static("3"));
2903        assert_eq!(retry_after_delay(&headers), Some(Duration::from_secs(3)));
2904
2905        let retry_at = SystemTime::now() + Duration::from_secs(2);
2906        let value = httpdate::fmt_http_date(retry_at);
2907        let Ok(value) = reqwest::header::HeaderValue::from_str(&value) else {
2908            panic!("formatted HTTP date should be a valid header");
2909        };
2910        headers.insert(RETRY_AFTER, value);
2911        let Some(delay) = retry_after_delay(&headers) else {
2912            panic!("HTTP date should produce a retry delay");
2913        };
2914        assert!(delay <= Duration::from_secs(2));
2915    }
2916
2917    #[tokio::test]
2918    async fn routed_llm_client_exposes_timeout_variant()
2919    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2920        let server = MockServer::start().await;
2921        Mock::given(method("POST"))
2922            .respond_with(
2923                ResponseTemplate::new(200).set_delay(std::time::Duration::from_millis(100)),
2924            )
2925            .mount(&server)
2926            .await;
2927
2928        let mut client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
2929        client.client = reqwest::Client::builder()
2930            .timeout(std::time::Duration::from_millis(10))
2931            .build()?;
2932
2933        let Err(error) = client.call(request_for(Some("gpt"), false)).await else {
2934            panic!("expected a timeout");
2935        };
2936        let LlmClientError::Timeout { source } = error else {
2937            panic!("expected the protocol timeout variant");
2938        };
2939        let Some(source) = source.downcast_ref::<reqwest::Error>() else {
2940            panic!("expected the reqwest timeout source");
2941        };
2942        assert!(source.is_timeout());
2943        Ok(())
2944    }
2945
2946    #[tokio::test]
2947    async fn llama_cpp_context_overflow_400_is_mapped()
2948    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2949        let server = MockServer::start().await;
2950        Mock::given(method("POST"))
2951            .respond_with(ResponseTemplate::new(400).set_body_json(json!({
2952                "error": {
2953                    "message": "request (6016 tokens) exceeds the available context size (4096 tokens), try increasing it"
2954                }
2955            })))
2956            .mount(&server)
2957            .await;
2958
2959        let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
2960
2961        let Err(error) = client
2962            .call_rewrite_model(request_for(Some("gpt"), false), None)
2963            .await
2964        else {
2965            panic!("expected an error");
2966        };
2967        assert!(matches!(
2968            error,
2969            LlmClientError::ContextWindowExceeded { model, .. } if model == "gpt"
2970        ));
2971        Ok(())
2972    }
2973
2974    #[tokio::test]
2975    async fn content_policy_400_preserves_upstream_error()
2976    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2977        let server = MockServer::start().await;
2978        let body = json!({
2979            "error": {
2980                "code": "content_policy_violation",
2981                "message": "request blocked by content policy",
2982                "type": "invalid_request_error"
2983            },
2984            "metadata": {"documentation_section": "context window"}
2985        });
2986        Mock::given(method("POST"))
2987            .respond_with(ResponseTemplate::new(400).set_body_json(&body))
2988            .expect(1)
2989            .mount(&server)
2990            .await;
2991        let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
2992        let Err(error) = client
2993            .call_rewrite_model(request_for(Some("gpt"), false), None)
2994            .await
2995        else {
2996            panic!("expected a policy denial");
2997        };
2998        let LlmClientError::UpstreamHttp {
2999            status,
3000            body: actual,
3001        } = error
3002        else {
3003            panic!("expected the upstream HTTP error, got {error:?}");
3004        };
3005        assert_eq!(status, StatusCode::BAD_REQUEST);
3006        assert_eq!(serde_json::from_str::<serde_json::Value>(&actual)?, body);
3007        Ok(())
3008    }
3009
3010    #[tokio::test]
3011    async fn forwards_only_allowlisted_metadata_headers()
3012    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3013        let server = MockServer::start().await;
3014        Mock::given(method("POST"))
3015            .and(wiremock::matchers::header("x-request-id", "abc"))
3016            // A forwarded Authorization must NOT override the backend's bearer key.
3017            .and(wiremock::matchers::header("authorization", "Bearer secret"))
3018            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
3019                "id": "1", "model": "gpt",
3020                "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
3021                "usage": {}
3022            })))
3023            .mount(&server)
3024            .await;
3025
3026        let mut headers = http::HeaderMap::new();
3027        headers.insert("x-request-id", http::HeaderValue::from_static("abc"));
3028        headers.insert(
3029            "authorization",
3030            http::HeaderValue::from_static("Bearer client-key"),
3031        );
3032        headers.insert(
3033            "x-goog-api-key",
3034            http::HeaderValue::from_static("client-google-key"),
3035        );
3036        headers.insert(
3037            "x-custom-internal-secret",
3038            http::HeaderValue::from_static("client-custom-key"),
3039        );
3040        let request = request_with_headers("gpt", headers);
3041
3042        let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
3043
3044        // The allowed request id survives and caller authorization cannot replace
3045        // the backend's credential.
3046        client.call_rewrite_model(request, None).await?;
3047        let received = server
3048            .received_requests()
3049            .await
3050            .ok_or("request recording should be enabled")?;
3051        let received = received.first().ok_or("expected one upstream request")?;
3052        assert!(!received.headers.contains_key("x-goog-api-key"));
3053        assert!(!received.headers.contains_key("x-custom-internal-secret"));
3054        Ok(())
3055    }
3056
3057    #[tokio::test]
3058    async fn forwards_application_headers_with_forward_auth()
3059    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3060        let server = MockServer::start().await;
3061        Mock::given(method("POST"))
3062            .and(wiremock::matchers::header(
3063                "authorization",
3064                "Bearer client-key",
3065            ))
3066            .and(wiremock::matchers::header(
3067                "x-goog-api-key",
3068                "client-google-key",
3069            ))
3070            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
3071                "id": "1", "model": "gpt",
3072                "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
3073                "usage": {}
3074            })))
3075            .mount(&server)
3076            .await;
3077
3078        let mut headers = http::HeaderMap::new();
3079        headers.insert(
3080            "authorization",
3081            http::HeaderValue::from_static("Bearer client-key"),
3082        );
3083        headers.insert(
3084            "x-goog-api-key",
3085            http::HeaderValue::from_static("client-google-key"),
3086        );
3087        headers.insert("host", http::HeaderValue::from_static("client.example"));
3088        headers.insert("content-type", http::HeaderValue::from_static("text/plain"));
3089        headers.insert("connection", http::HeaderValue::from_static("x-hop-by-hop"));
3090        headers.insert(
3091            "x-hop-by-hop",
3092            http::HeaderValue::from_static("client-only"),
3093        );
3094        let request = request_with_headers("gpt", headers);
3095        let backend_config = forwarding_config(&format!("{}/v1", server.uri()));
3096        let client = TranslatingLlmClient::new(&[ModelConfig::new(
3097            "gpt",
3098            Backend::OpenAiChat(backend_config),
3099            None,
3100        )])?;
3101
3102        client.call_rewrite_model(request, None).await?;
3103        let received = server
3104            .received_requests()
3105            .await
3106            .ok_or("request recording should be enabled")?;
3107        let received = received.first().ok_or("expected one upstream request")?;
3108        assert_ne!(received.headers["host"], "client.example");
3109        assert_eq!(received.headers["content-type"], "application/json");
3110        assert!(!received.headers.contains_key("connection"));
3111        assert!(!received.headers.contains_key("x-hop-by-hop"));
3112        Ok(())
3113    }
3114
3115    #[tokio::test]
3116    async fn forward_auth_preserves_anthropic_headers_and_redacts_errors()
3117    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3118        let server = MockServer::start().await;
3119        Mock::given(method("POST"))
3120            .and(path("/v1/messages"))
3121            .and(|request: &wiremock::Request| {
3122                request
3123                    .headers
3124                    .get("anthropic-beta")
3125                    .and_then(|value| value.to_str().ok())
3126                    == Some("oauth-2025-04-20")
3127                    && request.headers.get_all("anthropic-version").iter().count() == 1
3128                    && request
3129                        .headers
3130                        .get("anthropic-version")
3131                        .and_then(|value| value.to_str().ok())
3132                        == Some("2023-06-01")
3133            })
3134            .respond_with(ResponseTemplate::new(401).set_body_json(json!({
3135                "error": {"message": "rejected client-google-key"}
3136            })))
3137            .mount(&server)
3138            .await;
3139
3140        let mut headers = http::HeaderMap::new();
3141        headers.insert(
3142            "authorization",
3143            http::HeaderValue::from_static("Bearer client-key"),
3144        );
3145        headers.insert(
3146            "x-goog-api-key",
3147            http::HeaderValue::from_static("client-google-key"),
3148        );
3149        headers.insert(
3150            "anthropic-beta",
3151            http::HeaderValue::from_static("oauth-2025-04-20,prompt-caching-2024-07-31"),
3152        );
3153        headers.insert(
3154            "anthropic-version",
3155            http::HeaderValue::from_static("caller-version"),
3156        );
3157        let backend_config = forwarding_config(&server.uri());
3158        let client = TranslatingLlmClient::new(&[ModelConfig::new(
3159            "claude",
3160            Backend::Anthropic(backend_config),
3161            None,
3162        )])?;
3163        let raw = json!({
3164            "model": "claude",
3165            "max_tokens": 8,
3166            "messages": [{"role": "user", "content": "hello"}]
3167        });
3168
3169        let Err(LlmClientError::UpstreamHttp { body, .. }) = client
3170            .call_rewrite_model_raw(
3171                raw,
3172                Some(headers),
3173                Some(&ModelId::from("claude")),
3174                WireFormat::AnthropicMessages,
3175            )
3176            .await
3177        else {
3178            panic!("expected an upstream HTTP error");
3179        };
3180        assert_eq!(body, r#"{"error":{"message":"rejected [REDACTED]"}}"#);
3181        Ok(())
3182    }
3183
3184    // Exercises the `RoutedLlmClient` impl: `call` uses the model already materialized in the
3185    // request and round-trips a buffered response.
3186    #[tokio::test]
3187    async fn routed_llm_client_serves_the_request_model()
3188    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3189        let server = MockServer::start().await;
3190        Mock::given(method("POST"))
3191            .and(path("/v1/chat/completions"))
3192            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
3193                "id": "chatcmpl-1",
3194                "model": "gpt",
3195                "choices": [{
3196                    "index": 0,
3197                    "message": {"role": "assistant", "content": "routed hi"},
3198                    "finish_reason": "stop"
3199                }],
3200                "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}
3201            })))
3202            .mount(&server)
3203            .await;
3204
3205        let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
3206        let response = client.call(request_for(Some("gpt"), false)).await?;
3207        let agg = response.llm_response.into_agg().await?;
3208        assert_eq!(completion_text(&agg), "routed hi");
3209        Ok(())
3210    }
3211
3212    #[tokio::test]
3213    async fn invalid_raw_request_is_a_request_translation_error()
3214    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3215        let client = TranslatingLlmClient::new(&[])?;
3216        let Err(error) = client
3217            .call_rewrite_model_raw(
3218                json!("invalid"),
3219                None,
3220                Some(&ModelId::from("gpt")),
3221                WireFormat::OpenAiChat,
3222            )
3223            .await
3224        else {
3225            panic!("expected request translation to fail");
3226        };
3227
3228        assert!(matches!(
3229            error,
3230            LlmClientError::RequestTranslation(message) if !message.is_empty()
3231        ));
3232        Ok(())
3233    }
3234
3235    // Raw path, buffered: decode an OpenAI Chat body -> call -> encode back to OpenAI
3236    // Chat JSON, with the served `model` restamped over the id the caller addressed.
3237    #[tokio::test]
3238    async fn call_rewrite_model_raw_round_trips_buffered_json()
3239    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3240        let server = MockServer::start().await;
3241        Mock::given(method("POST"))
3242            .and(path("/v1/chat/completions"))
3243            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
3244                "id": "chatcmpl-1",
3245                "model": "gpt",
3246                "choices": [{
3247                    "index": 0,
3248                    "message": {"role": "assistant", "content": "Hi there"},
3249                    "finish_reason": "stop"
3250                }],
3251                "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}
3252            })))
3253            .mount(&server)
3254            .await;
3255
3256        let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
3257        let raw = json!({
3258            "model": "client-facing",
3259            "messages": [{"role": "user", "content": "hi"}]
3260        });
3261        let RawResponse::Buffered(body) = client
3262            .call_rewrite_model_raw(
3263                raw,
3264                None,
3265                Some(&ModelId::from("gpt")),
3266                WireFormat::OpenAiChat,
3267            )
3268            .await?
3269        else {
3270            panic!("expected a buffered response");
3271        };
3272
3273        assert_eq!(body["choices"][0]["message"]["content"], "Hi there");
3274        // The client sees the model that answered, not the "client-facing" route id.
3275        assert_eq!(body["model"], "gpt");
3276        Ok(())
3277    }
3278
3279    // A Codex namespace is folded into the upstream tool name, then split back
3280    // into name and namespace on the Responses call that returns to Codex.
3281    #[tokio::test]
3282    async fn call_rewrite_model_raw_restores_codex_mcp_namespace()
3283    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3284        let server = MockServer::start().await;
3285        Mock::given(method("POST"))
3286            .and(path("/v1/chat/completions"))
3287            .and(wiremock::matchers::body_partial_json(json!({
3288                "tools": [{
3289                    "type": "function",
3290                    "function": {"name": "mcp__open_websearch__search"}
3291                }]
3292            })))
3293            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
3294                "id": "chatcmpl-1",
3295                "model": "gpt",
3296                "choices": [{
3297                    "index": 0,
3298                    "message": {
3299                        "role": "assistant",
3300                        "content": null,
3301                        "tool_calls": [{
3302                            "id": "call_1",
3303                            "type": "function",
3304                            "function": {
3305                                "name": "mcp__open_websearch__search",
3306                                "arguments": "{\"q\":\"rust\"}"
3307                            }
3308                        }]
3309                    },
3310                    "finish_reason": "tool_calls"
3311                }],
3312                "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}
3313            })))
3314            .mount(&server)
3315            .await;
3316
3317        let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
3318        let raw = json!({
3319            "model": "client-facing",
3320            "input": "Search for Rust.",
3321            "tools": [{
3322                "type": "namespace",
3323                "name": "mcp__open_websearch",
3324                "tools": [{
3325                    "type": "function",
3326                    "name": "search",
3327                    "description": "Search the web",
3328                    "parameters": {"type": "object", "properties": {"q": {"type": "string"}}}
3329                }]
3330            }]
3331        });
3332
3333        let RawResponse::Buffered(body) = client
3334            .call_rewrite_model_raw(
3335                raw,
3336                None,
3337                Some(&ModelId::from("gpt")),
3338                WireFormat::OpenAiResponses,
3339            )
3340            .await?
3341        else {
3342            panic!("expected a buffered response");
3343        };
3344
3345        assert_eq!(body["output"][0]["type"], "function_call");
3346        assert_eq!(body["output"][0]["name"], "search");
3347        assert_eq!(body["output"][0]["namespace"], "mcp__open_websearch");
3348        // Arguments are parsed and re-serialized, so the spacing is normalized.
3349        assert_eq!(body["output"][0]["arguments"], "{\"q\": \"rust\"}");
3350        Ok(())
3351    }
3352
3353    #[tokio::test]
3354    async fn openai_responses_backend_flattens_codex_namespaces_before_upstream()
3355    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3356        let server = MockServer::start().await;
3357        Mock::given(method("POST"))
3358            .and(path("/v1/responses"))
3359            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
3360                "id": "resp_1",
3361                "model": "gpt",
3362                "object": "response",
3363                "created_at": 0,
3364                "status": "completed",
3365                "output": [{
3366                    "type": "function_call",
3367                    "call_id": "call_1",
3368                    "name": "mcp__open_websearch__search",
3369                    "arguments": "{\"q\":\"rust\"}"
3370                }],
3371                "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}
3372            })))
3373            .mount(&server)
3374            .await;
3375
3376        let client = TranslatingLlmClient::new(&responses_map(&format!("{}/v1", server.uri())))?;
3377        let parameters = json!({"type": "object", "properties": {"q": {"type": "string"}}});
3378        let raw = json!({
3379            "model": "client-facing",
3380            "input": [{
3381                "type": "function_call",
3382                "call_id": "call_0",
3383                "name": "search",
3384                "namespace": "mcp__open_websearch",
3385                "arguments": "{}"
3386            }],
3387            "tool_choice": {
3388                "type": "function",
3389                "name": "search",
3390                "namespace": "mcp__open_websearch"
3391            },
3392            "tools": [{
3393                "type": "namespace",
3394                "name": "mcp__open_websearch",
3395                "tools": [{
3396                    "type": "function",
3397                    "name": "search",
3398                    "parameters": parameters.clone()
3399                }]
3400            }]
3401        });
3402
3403        let RawResponse::Buffered(body) = client
3404            .call_rewrite_model_raw(
3405                raw,
3406                None,
3407                Some(&ModelId::from("gpt")),
3408                WireFormat::OpenAiResponses,
3409            )
3410            .await?
3411        else {
3412            panic!("expected a buffered response");
3413        };
3414
3415        let received = server
3416            .received_requests()
3417            .await
3418            .ok_or("request recording should be enabled")?;
3419        let request_body: Value = serde_json::from_slice(&received[0].body)?;
3420        assert_eq!(request_body["model"], "gpt");
3421        assert_eq!(
3422            request_body["tools"],
3423            json!([{
3424                "type": "function",
3425                "name": "mcp__open_websearch__search",
3426                "description": "",
3427                "parameters": parameters
3428            }])
3429        );
3430        assert_eq!(
3431            request_body["tool_choice"],
3432            json!({"type": "function", "name": "mcp__open_websearch__search"})
3433        );
3434        assert_eq!(
3435            request_body["input"][0]["name"],
3436            "mcp__open_websearch__search"
3437        );
3438        assert!(request_body["input"][0].get("namespace").is_none());
3439
3440        assert_eq!(body["output"][0]["type"], "function_call");
3441        assert_eq!(body["output"][0]["name"], "search");
3442        assert_eq!(body["output"][0]["namespace"], "mcp__open_websearch");
3443        Ok(())
3444    }
3445
3446    // A stage route can answer early steps of a Codex turn from Claude and a later step from
3447    // GPT, and Codex replays Claude's reasoning items on every step. OpenAI returns 400 for two
3448    // kinds: items with an ID but no `encrypted_content` when the request sets `store: false`,
3449    // and items whose `encrypted_content` holds Anthropic thinking that Switchyard packed for
3450    // the client. The Responses leg must drop both and keep GPT's own encrypted reasoning.
3451    #[tokio::test]
3452    async fn openai_responses_backend_drops_reasoning_the_provider_cannot_read()
3453    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3454        let server = MockServer::start().await;
3455        Mock::given(method("POST"))
3456            .and(path("/v1/responses"))
3457            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
3458                "id": "resp_1",
3459                "model": "gpt",
3460                "object": "response",
3461                "created_at": 0,
3462                "status": "completed",
3463                "output": [{
3464                    "type": "message",
3465                    "role": "assistant",
3466                    "content": [{"type": "output_text", "text": "ok"}]
3467                }],
3468                "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}
3469            })))
3470            .mount(&server)
3471            .await;
3472
3473        let client = TranslatingLlmClient::new(&responses_map(&format!("{}/v1", server.uri())))?;
3474        let claude_reasoning = json!({
3475            "type": "reasoning",
3476            "id": "rs_b99feb9a-763d-4f7f-b6b2-2c0199b47ded",
3477            "summary": [{"type": "summary_text", "text": "Read the files first."}],
3478            "encrypted_content": null
3479        });
3480        // Build the item the way Switchyard does when it returns a Claude reply to a Responses
3481        // client, so the test follows any change to that format.
3482        let claude_reply = switchyard_translation::TranslationEngine::default()
3483            .translate_response(
3484                WireFormat::AnthropicMessages,
3485                WireFormat::OpenAiResponses,
3486                &json!({
3487                    "id": "msg_1",
3488                    "type": "message",
3489                    "role": "assistant",
3490                    "model": "claude",
3491                    "content": [
3492                        {"type": "thinking", "thinking": "Plan the edit.", "signature": "sig-abc"},
3493                        {"type": "text", "text": "ok"}
3494                    ],
3495                    "stop_reason": "end_turn",
3496                    "usage": {"input_tokens": 1, "output_tokens": 1}
3497                }),
3498                &switchyard_translation::TranslationPolicy::default(),
3499            )?
3500            .body;
3501        let switchyard_reasoning = claude_reply["output"]
3502            .as_array()
3503            .and_then(|output| output.iter().find(|item| item["type"] == "reasoning"))
3504            .cloned()
3505            .ok_or("translated Claude reply should contain a reasoning item")?;
3506        assert!(switchyard_translation::is_anthropic_thinking_envelope(
3507            switchyard_reasoning["encrypted_content"]
3508                .as_str()
3509                .unwrap_or_default()
3510        ));
3511        let gpt_reasoning = json!({
3512            "type": "reasoning",
3513            "id": "rs_gpt",
3514            "summary": [],
3515            "encrypted_content": "opaque-encrypted-reasoning"
3516        });
3517        let empty_reasoning =
3518            json!({"type": "reasoning", "id": "rs_empty", "summary": [], "encrypted_content": ""});
3519        let user = json!({"type": "message", "role": "user", "content": "fix the build"});
3520        let call = json!({"type": "function_call", "call_id": "call_1", "name": "exec", "arguments": "{}"});
3521        let output = json!({"type": "function_call_output", "call_id": "call_1", "output": "ok"});
3522        let input = json!([
3523            user,
3524            claude_reasoning,
3525            switchyard_reasoning,
3526            call,
3527            output,
3528            gpt_reasoning,
3529            empty_reasoning
3530        ]);
3531
3532        for store in [json!(false), json!(true)] {
3533            let raw = json!({"model": "client-facing", "store": store, "input": input.clone()});
3534            client
3535                .call_rewrite_model_raw(
3536                    raw,
3537                    None,
3538                    Some(&ModelId::from("gpt")),
3539                    WireFormat::OpenAiResponses,
3540                )
3541                .await?;
3542        }
3543
3544        let received = server
3545            .received_requests()
3546            .await
3547            .ok_or("request recording should be enabled")?;
3548        let sent: Vec<Value> = received
3549            .iter()
3550            .map(|request| serde_json::from_slice::<Value>(&request.body))
3551            .collect::<std::result::Result<_, _>>()?;
3552        // With `store: false`, only GPT's own encrypted reasoning survives.
3553        assert_eq!(sent[0]["input"], json!([user, call, output, gpt_reasoning]));
3554        // With `store: true`, IDs can refer to stored reasoning, but OpenAI still cannot
3555        // decrypt the Switchyard item.
3556        assert_eq!(
3557            sent[1]["input"],
3558            json!([
3559                user,
3560                claude_reasoning,
3561                call,
3562                output,
3563                gpt_reasoning,
3564                empty_reasoning
3565            ])
3566        );
3567        Ok(())
3568    }
3569
3570    // Raw path, streaming: an inbound `stream: true` request yields an unframed stream
3571    // of OpenAI Chat chunk objects whose deltas reassemble the completion.
3572    #[tokio::test]
3573    async fn call_rewrite_model_raw_streams_wire_events()
3574    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3575        use futures::TryStreamExt;
3576
3577        let server = MockServer::start().await;
3578        let sse = "data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}\n\n\
3579             data: {\"choices\":[{\"delta\":{\"content\":\" world\"}}]}\n\n\
3580             data: [DONE]\n\n";
3581        Mock::given(method("POST"))
3582            .and(path("/v1/chat/completions"))
3583            .respond_with(ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream"))
3584            .mount(&server)
3585            .await;
3586
3587        let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
3588        let raw = json!({
3589            "model": "client-facing",
3590            "messages": [{"role": "user", "content": "hi"}],
3591            "stream": true
3592        });
3593        let RawResponse::Stream(stream) = client
3594            .call_rewrite_model_raw(
3595                raw,
3596                None,
3597                Some(&ModelId::from("gpt")),
3598                WireFormat::OpenAiChat,
3599            )
3600            .await?
3601        else {
3602            panic!("expected a streamed response");
3603        };
3604
3605        let events: Vec<Value> = stream.try_collect().await?;
3606        assert!(!events.is_empty(), "expected at least one wire event");
3607        let content: String = events
3608            .iter()
3609            .filter_map(|event| event["choices"][0]["delta"]["content"].as_str())
3610            .collect();
3611        assert_eq!(content, "Hello world");
3612        // The mock frames carry no `model`, so every chunk's model comes from the
3613        // served id rather than the "unknown" fallback or the caller's route id.
3614        assert!(events.iter().all(|event| event["model"] == "gpt"));
3615        Ok(())
3616    }
3617
3618    // Raw path forwards allowed caller headers to the upstream.
3619    #[tokio::test]
3620    async fn call_rewrite_model_raw_forwards_headers()
3621    -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3622        let server = MockServer::start().await;
3623        Mock::given(method("POST"))
3624            .and(wiremock::matchers::header("x-request-id", "abc"))
3625            // A forwarded authorization must NOT override the backend's bearer key.
3626            .and(wiremock::matchers::header("authorization", "Bearer secret"))
3627            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
3628                "id": "1", "model": "gpt",
3629                "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
3630                "usage": {}
3631            })))
3632            .mount(&server)
3633            .await;
3634
3635        let mut headers = http::HeaderMap::new();
3636        headers.insert("x-request-id", http::HeaderValue::from_static("abc"));
3637        headers.insert(
3638            "authorization",
3639            http::HeaderValue::from_static("Bearer client-key"),
3640        );
3641
3642        let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
3643        let raw = json!({"model": "gpt", "messages": [{"role": "user", "content": "hi"}]});
3644        // The allowed request id survives and caller authorization cannot replace
3645        // the backend's credential.
3646        client
3647            .call_rewrite_model_raw(
3648                raw,
3649                Some(headers),
3650                Some(&ModelId::from("gpt")),
3651                WireFormat::OpenAiChat,
3652            )
3653            .await?;
3654        Ok(())
3655    }
3656}