// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 //! [`TranslatingLlmClient`]: encode a neutral //! request, call the configured backend over HTTP, decode the neutral response. use std::borrow::Cow; use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::future::ready; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{Duration, SystemTime}; use async_trait::async_trait; use futures_util::{StreamExt, stream}; use http::StatusCode; use reqwest::RequestBuilder; use reqwest::header::{HeaderMap, RETRY_AFTER}; use serde_json::{Map, Value, json}; use switchyard_protocol::codex_namespaces::NAMESPACE_SEPARATOR; use switchyard_protocol::{ LlmRequest, LlmResponse, LlmResponseChunk, LlmResponseStream, LlmResponseStreamEvent, Metadata, ModelId, Request, Response, RoutedLlmClient, }; use switchyard_translation::{ TranslationError, WireFormat, decode_stream, encode_stream_with_extensions, }; use tokio::time::Instant; use tracing::Instrument; use crate::backend::{Backend, openai_url}; use crate::error::{LlmClientError, Result}; use crate::metrics; use crate::observed_translation; use crate::raw::RawResponse; // Caller headers safe to send when caller auth forwarding is disabled. const ALLOWED_METADATA_HEADERS: &[&str] = &["x-request-id"]; // Headers tied to the inbound connection, destination, or body. The HTTP client // must rebuild these for the upstream request, even when forwarding auth. const NON_FORWARDABLE_HEADERS: &[&str] = &[ "host", "content-length", "connection", "proxy-connection", "keep-alive", "proxy-authenticate", "proxy-authorization", "te", "trailer", "transfer-encoding", "upgrade", "content-type", "content-encoding", "accept-encoding", "expect", ]; const INITIAL_RETRY_DELAY: Duration = Duration::from_millis(250); const MAX_RETRY_BACKOFF: Duration = Duration::from_secs(2); const MAX_RETRY_AFTER: Duration = Duration::from_secs(60); /// A named target's provider model and backends. #[derive(Clone, Debug)] pub struct ModelConfig { target_name: ModelId, provider_model: Option, default_backend: Backend, other_backends: Option>, responses_reasoning: crate::ResponsesReasoningPolicy, } impl ModelConfig { /// Configures a target. Its name is also the provider model ID unless overridden. pub fn new( target_name: impl Into, default_backend: Backend, other_backends: Option>, ) -> Self { Self { target_name: target_name.into(), provider_model: None, default_backend, other_backends, responses_reasoning: crate::ResponsesReasoningPolicy::default(), } } /// Sets how Responses reasoning items are replayed to this target. #[must_use] pub fn with_responses_reasoning(mut self, policy: crate::ResponsesReasoningPolicy) -> Self { self.responses_reasoning = policy; self } /// Sends `provider_model` upstream in place of `target_name`. The client is still /// called with `target_name`, so several entries can serve one provider model with /// different backend settings. pub fn with_provider_model(mut self, provider_model: impl Into) -> Self { self.provider_model = Some(provider_model.into()); self } } /// A model-bearing provider operation outside the normal completion endpoint. #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum AuxiliaryOperation { /// Anthropic Messages input-token counting. AnthropicCountTokens, /// OpenAI Responses input-token counting. ResponsesInputTokens, /// OpenAI Responses compaction. ResponsesCompact, } impl AuxiliaryOperation { const fn wire_format(self) -> WireFormat { match self { Self::AnthropicCountTokens => WireFormat::AnthropicMessages, Self::ResponsesInputTokens | Self::ResponsesCompact => WireFormat::OpenAiResponses, } } fn url(self, backend: &Backend) -> String { let suffix = match self { Self::AnthropicCountTokens => return backend.count_tokens_url(), Self::ResponsesInputTokens => "/responses/input_tokens", Self::ResponsesCompact => "/responses/compact", }; let base_url = backend.url(); openai_url(&base_url, suffix).unwrap_or_else(|error| { tracing::error!(%error, "Unable to build OpenAI auxiliary endpoint URL"); base_url }) } } /// A client that dispatches neutral-IR requests to named targets. /// /// Each [`ModelConfig`] defines a target's provider model and backends. /// Each call resolves the target and wire format, encodes the request, /// applies auth and forwarded headers, sends the HTTP request with a shared /// [`reqwest::Client`], and decodes the response back to the neutral IR (buffered /// or streamed). pub struct TranslatingLlmClient { configs_by_target: HashMap, client: reqwest::Client, forward_auth_client: reqwest::Client, is_rate_limit_cooldown_enabled: bool, cooldown_epoch: Instant, // Milliseconds since cooldown_epoch. Each atomic is independent, so relaxed ordering suffices. unavailable_until: HashMap, } impl TranslatingLlmClient { /// Builds a client over the given [`ModelConfig`]s, with a fresh shared HTTP /// client and the built-in translation codecs. pub fn new(model_configs: &[ModelConfig]) -> Result { // Pre-populate so concurrent calls update only atomics and keep map access lock-free. let mut unavailable_until = HashMap::new(); for config in model_configs { for backend in std::iter::once(&config.default_backend) .chain(config.other_backends.iter().flatten()) { backend.validate_configured_headers(&config.target_name)?; } unavailable_until.insert(config.target_name.clone(), AtomicU64::new(0)); } let build_client = |builder: reqwest::ClientBuilder| { builder.build().map_err(|error| LlmClientError::Transport { source: Box::new(error), }) }; let client = build_client(reqwest::Client::builder())?; // A redirect could move provider-specific headers to another origin. // Forwarded credentials are sent only to the configured URL. let forward_auth_client = build_client(reqwest::Client::builder().redirect(reqwest::redirect::Policy::none()))?; let configs_by_target = model_configs .iter() .map(|config| (config.target_name.clone(), config.clone())) .collect(); Ok(Self { configs_by_target, client, forward_auth_client, is_rate_limit_cooldown_enabled: true, cooldown_epoch: Instant::now(), unavailable_until, }) } /// Disable HTTP 429 cooldown so subsequent calls can receive the upstream error. pub fn without_rate_limit_cooldown(mut self) -> Self { self.is_rate_limit_cooldown_enabled = false; self } /// The target's backend for `format`, or `None` if no matching backend exists. pub fn backend_for(&self, target_name: &ModelId, format: WireFormat) -> Option<&Backend> { self.configs_by_target.get(target_name).and_then(|config| { if config.default_backend.wire_format() == format { Some(&config.default_backend) } else { config .other_backends .as_ref() .and_then(|backends| backends.iter().find(|b| b.wire_format() == format)) } }) } /// The target's provider model ID, falling back to its name when none is configured. pub fn provider_model<'a>(&'a self, target_name: &'a ModelId) -> &'a ModelId { self.configs_by_target .get(target_name) .and_then(|config| config.provider_model.as_ref()) .unwrap_or(target_name) } /// Whether the target has a backend for `operation`. pub fn supports_auxiliary(&self, target_name: &ModelId, operation: AuxiliaryOperation) -> bool { self.backend_for(target_name, operation.wire_format()) .is_some() } /// Reads llama.cpp model properties through the backend's configured auth and headers. /// /// Redirects are disabled so credentials cannot move to another origin. pub async fn get_model_properties( &self, target_name: &ModelId, format: WireFormat, timeout: Duration, ) -> Result { let backend = self.backend_for(target_name, format) .ok_or_else(|| LlmClientError::Configuration { message: format!("model {target_name} has no backend for {format:?}"), })?; let mut url = reqwest::Url::parse(backend.base_url()).map_err(|error| { LlmClientError::Configuration { message: format!("model {target_name} has an invalid backend URL: {error}"), } })?; let root = url.path().trim_end_matches('/'); let root = root .strip_suffix("/chat/completions") .or_else(|| root.strip_suffix("/responses")) .unwrap_or(root); let path = format!( "{}/props", root.strip_suffix("/v1") .unwrap_or(root) .trim_end_matches('/') ); url.set_path(&path); url.set_query(None); let builder = self.forward_auth_client.get(url).timeout(timeout); let builder = apply_extra_headers(builder, backend); let response = backend .apply_auth(builder) .send() .await .map_err(convert_reqwest_error)? .error_for_status() .map_err(convert_reqwest_error)?; response .json() .await .map_err(|source| LlmClientError::InvalidResponse { source: Box::new(source), }) } /// Calls a model-bearing auxiliary provider operation. /// /// Returns an error when the model has no compatible backend or the upstream /// request fails or returns invalid JSON. pub async fn call_auxiliary( &self, target_name: &ModelId, request: Request, operation: AuxiliaryOperation, ) -> Result { let wire_format = operation.wire_format(); let backend = self.backend_for(target_name, wire_format).ok_or_else(|| { LlmClientError::Configuration { message: format!("model {target_name} has no backend for {operation:?}"), } })?; let Request { mut llm_request, metadata, .. } = request; llm_request.model = Some(self.provider_model(target_name).to_string()); let http_response = self .send_encoded( backend, wire_format, llm_request, metadata.as_ref(), target_name, UpstreamEndpoint::Auxiliary(operation), ) .await?; let EncodedResponse::Buffered { body, .. } = http_response else { return Err(LlmClientError::InvalidRequest { message: "auxiliary endpoints do not support streaming".to_string(), }); }; serde_json::from_slice(&body).map_err(|error| LlmClientError::InvalidResponse { source: Box::new(error), }) } /// Encode `llm_request` for `wire_format`, POST it to `url` with the request's /// forwarded headers plus the backend's static headers and auth, and return the /// successful upstream response. A /// buffered response is fully collected within the retry boundary; a streamed /// response has its first decoded event checked within that boundary. The backend's /// deadline, when set, spans every attempt, retry delay, and subsequent stream read. A /// non-success status maps to a typed error — a 400 is classified as a context-window /// overflow via the backend's provider rules. Shared by /// [`call_rewrite_model`](Self::call_rewrite_model) (which POSTs to the /// backend's completion URL and decodes a response) and /// the model-bearing auxiliary operations, which return raw JSON. async fn send_encoded( &self, backend: &Backend, wire_format: WireFormat, llm_request: LlmRequest, metadata: Option<&Metadata>, target_name: &ModelId, endpoint: UpstreamEndpoint, ) -> Result { let cooldown = backend.failure_cooldown(); let unavailable_until = if matches!(endpoint, UpstreamEndpoint::Completion) && !cooldown.is_zero() { self.unavailable_until.get(target_name) } else { None }; if let Some(until) = unavailable_until && until.load(Ordering::Relaxed) > duration_millis(self.cooldown_epoch.elapsed()) { tracing::debug!(model = %target_name, wire_format = %wire_format, "skipping backend during failure cooldown"); return Err(LlmClientError::TemporarilyUnavailable); } let mut body = observed_translation::encode_request(&llm_request, wire_format) .map_err(|error| LlmClientError::RequestEncoding(error.to_string()))?; // `encode_request` round-trips a preserved same-format body verbatim, // which keeps the caller's original `model`. Set the provider model explicitly. let provider_model_id = self.provider_model(target_name); set_json_model(&mut body, provider_model_id); if matches!(backend, Backend::OpenAiResponses(_)) { sanitize_openai_responses_provider_body(&mut body); } // Strip before `merge_extra_body` so a target can reinstate either field // deliberately via `extra_body`. if matches!(backend, Backend::Anthropic(_)) { strip_anthropic_incompatible_fields(&mut body); strip_unsigned_thinking_blocks(&mut body); } omit_configured_body_fields(&mut body, backend.omit_body_fields()); merge_extra_body(&mut body, backend.extra_body()); if matches!(backend, Backend::OpenAiResponses(_)) { let outbound_model_id = body .get("model") .and_then(Value::as_str) .unwrap_or(provider_model_id) .to_owned(); self.configs_by_target .get(target_name) .map(|config| config.responses_reasoning) .unwrap_or_default() .normalize(&mut body, backend.base_url(), &outbound_model_id); } // After the merge on purpose: the effort override must win over both the caller's // value and any `reasoning` default a target set through `extra_body`. Token counting // does not depend on effort, and Anthropic's count endpoint may reject the field. if matches!(endpoint, UpstreamEndpoint::Completion) { apply_reasoning_effort(&mut body, backend); } if matches!(backend, Backend::OpenAiResponses(_)) { normalize_codex_string_input(&backend.url(), &mut body); } if matches!(backend, Backend::Anthropic(_)) { if nvidia_inference_api_backend(backend) { strip_anthropic_tool_strict(&mut body); } enable_anthropic_prompt_caching(&mut body, llm_request.temporary_note_count); } if matches!( backend, Backend::OpenAiChat(_) | Backend::OpenAiResponses(_) ) { enable_openai_guidance_caching( &mut body, wire_format, provider_model_id, llm_request.temporary_note_count, ); } if matches!(backend, Backend::OpenAiChat(_)) { ensure_openai_stream_usage(&mut body); } let streaming = endpoint.allows_streaming() && body.get("stream").and_then(Value::as_bool).unwrap_or(false); let url = endpoint.url(backend); record_gen_ai_request(&url, provider_model_id, streaming); let result = self .send_with_retries(&url, backend, &body, metadata, target_name, streaming) .await; // Keep quota failures request-local when forwarding auth or returning 429 to callers. if let Some(until) = unavailable_until && let Err(error) = &result && is_transient_failure(error) && !((backend.is_forwarding_auth() || !self.is_rate_limit_cooldown_enabled) && matches!(error, LlmClientError::UpstreamHttp { status, .. } if *status == StatusCode::TOO_MANY_REQUESTS)) { let deadline = duration_millis(self.cooldown_epoch.elapsed().saturating_add(cooldown)); until.fetch_max(deadline, Ordering::Relaxed); } // Let cooldown expire naturally; a successful in-flight call may overlap a newer failure. result } // Sends the encoded body, retrying retryable failures within the backend's retry budget. async fn send_with_retries( &self, url: &str, backend: &Backend, body: &Value, metadata: Option<&Metadata>, target_name: &ModelId, streaming: bool, ) -> Result { let max_retries = u64::from(backend.max_retries()); let max_attempts = max_retries + 1 + u64::from( matches!(backend, Backend::OpenAiChat(_)) && body.get("chat_template_kwargs").is_some(), ); let mut body = Cow::Borrowed(body); let mut template_retry = false; let expires_at = backend .timeout() .map(|timeout| (tokio::time::sleep(timeout).deadline(), timeout)); let deadline = async { match expires_at { Some((at, timeout)) => { tokio::time::sleep_until(at).await; timeout } None => std::future::pending().await, } }; tokio::pin!(deadline); let mut attempt = 0_u64; loop { if let Some((at, timeout)) = expires_at && tokio::time::Instant::now() >= at { return Err(deadline_error(timeout)); } let span = tracing::debug_span!( target: "libsy", "libsy.upstream_attempt", model = %target_name, wire_format = %backend.wire_format(), attempt = attempt + 1 + u64::from(template_retry), max_attempts, retry = attempt > 0 || template_retry, openinference.span.kind = "CHAIN", outcome = tracing::field::Empty, status_code = tracing::field::Empty, will_retry = tracing::field::Empty, retry_delay_ms = tracing::field::Empty, ); let mut attempt_started = false; let result = tokio::select! { biased; timeout = &mut deadline => { if attempt_started { metrics::record_upstream_attempt(None); } span.record("outcome", "error"); span.record("will_retry", false); return Err(deadline_error(timeout)); } result = async { attempt_started = true; self.send_once(url, backend, body.as_ref(), metadata, target_name, streaming).await }.instrument(span.clone()) => result, }; // The retained handle updates this same attempt span with its outcome. match result { Ok(response) => { span.record("outcome", "ok"); span.record("status_code", response.status()); span.record("will_retry", false); if attempt > 0 || template_retry { metrics::record_retry_recovered(); } return Ok(match response { EncodedResponse::Streaming { status, chunks, upstream_headers, } => EncodedResponse::Streaming { status, chunks: stream_with_deadline(chunks, expires_at), upstream_headers, }, buffered => buffered, }); } Err(failure) => { let retry_without_template = matches!(backend, Backend::OpenAiChat(_)) && body.get("chat_template_kwargs") == Some(&json!({"enable_thinking": false})) && rejects_template_options(&failure.error); let will_retry = retry_without_template || (attempt < max_retries && failure.is_retryable()); span.record("outcome", "error"); if let Some(status) = failure.status { span.record("status_code", status.as_u16()); } span.record("will_retry", will_retry); if !will_retry { return Err(failure.error); } if retry_without_template { if let Some(object) = body.to_mut().as_object_mut() { object.remove("chat_template_kwargs"); } template_retry = true; span.record("retry_delay_ms", 0_u64); continue; } let delay = retry_delay(attempt, failure.retry_after); span.record("retry_delay_ms", duration_millis(delay)); // Close the attempt span before sleeping so backoff is not attempt latency. drop(span); tokio::select! { biased; timeout = &mut deadline => return Err(deadline_error(timeout)), _ = tokio::time::sleep(delay) => {} } attempt += 1; } } } } // Performs one HTTP attempt and retains the retry metadata alongside any error. async fn send_once( &self, url: &str, backend: &Backend, body: &Value, metadata: Option<&Metadata>, target_name: &ModelId, streaming: bool, ) -> std::result::Result { let client = if backend.is_forwarding_auth() { &self.forward_auth_client } else { &self.client }; let builder = client.post(url).json(body); let builder = forward_metadata_headers(builder, metadata, backend); let builder = backend.apply_forwarded_auth(builder, metadata); let builder = apply_extra_headers(builder, backend); let builder = backend.apply_auth(builder); let response = match async { let request = builder.build()?; if let Some(body) = request.body().and_then(reqwest::Body::as_bytes) { crate::payload_capture::record("request", || String::from_utf8_lossy(body)); } client.execute(request).await } .await { Ok(response) => response, Err(error) => { metrics::record_upstream_attempt(None); return Err(AttemptFailure { error: convert_reqwest_error(error), status: None, retry_after: None, }); } }; let status = response.status(); if status.is_success() { if streaming { let upstream_headers = response.headers().clone(); let chunks = match prepare_response_stream(response, backend, target_name).await { Ok(chunks) => chunks, Err(error) => { metrics::record_upstream_attempt(None); return Err(AttemptFailure { error, status: Some(status), retry_after: None, }); } }; metrics::record_upstream_attempt(Some(status.as_u16())); return Ok(EncodedResponse::Streaming { status: status.as_u16(), chunks, upstream_headers, }); } let upstream_headers = response.headers().clone(); let body = match response.bytes().await { Ok(body) => body, Err(error) => { metrics::record_upstream_attempt(None); return Err(AttemptFailure { error: convert_reqwest_error(error), status: Some(status), retry_after: None, }); } }; metrics::record_upstream_attempt(Some(status.as_u16())); crate::payload_capture::record("response", || String::from_utf8_lossy(&body)); return Ok(EncodedResponse::Buffered { status: status.as_u16(), body: body.to_vec(), upstream_headers, }); } let retry_after = retry_after_delay(response.headers()); let body = match response.text().await { Ok(body) => body, Err(error) => { metrics::record_upstream_attempt(None); return Err(AttemptFailure { error: convert_reqwest_error(error), status: Some(status), retry_after, }); } }; let body = redact_forwarded_headers(body, metadata, backend.is_forwarding_auth()); metrics::record_upstream_attempt(Some(status.as_u16())); let error = if status == reqwest::StatusCode::BAD_REQUEST && backend.is_context_overflow(&body) { LlmClientError::ContextWindowExceeded { model: target_name.clone(), message: body, } } else { LlmClientError::UpstreamHttp { status, body } }; Err(AttemptFailure { error, status: Some(status), retry_after, }) } /// Calls the selected target using its configured provider model and backend. /// `target_name` takes precedence over the request's model field. /// Request metadata selects the wire format, falling back to the target's default. pub async fn call_rewrite_model( &self, request: Request, target_name: Option<&ModelId>, ) -> Result { let Request { mut llm_request, metadata, .. } = request; let selected_target = target_name .cloned() .or_else(|| llm_request.model.map(ModelId::from)) .ok_or_else(|| LlmClientError::InvalidRequest { message: "no model given".to_string(), })?; llm_request.model = Some(self.provider_model(&selected_target).to_string()); let orig_format = metadata.as_ref().and_then(|m| m.wire_format); let wire_format = orig_format.unwrap_or( self.configs_by_target .get(&selected_target) .map(|config| config.default_backend.wire_format()) .ok_or_else(|| LlmClientError::Configuration { message: format!("no backend configured for model {selected_target:?}"), })?, ); let backend = self .backend_for(&selected_target, wire_format) .ok_or_else(|| LlmClientError::Configuration { message: format!( "model {selected_target:?} has no backend for format {wire_format}" ), })?; let http_response = self .send_encoded( backend, wire_format, llm_request, metadata.as_ref(), &selected_target, UpstreamEndpoint::Completion, ) .await?; let (llm_response, upstream_headers) = match http_response { EncodedResponse::Streaming { chunks, upstream_headers, .. } => (LlmResponse::Stream(chunks), upstream_headers), EncodedResponse::Buffered { body, upstream_headers, .. } => { let body = serde_json::from_slice::(&body).map_err(|error| { LlmClientError::ResponseTranslation(format!("invalid upstream JSON: {error}")) })?; // Map a provider's failed generation to 502, even under HTTP 200. // Redact forwarded credentials before returning the provider error. let mut agg = observed_translation::decode_response(&body, wire_format).map_err(|error| { match error { TranslationError::UpstreamFailure { error } => { LlmClientError::UpstreamHttp { status: StatusCode::BAD_GATEWAY, body: redact_forwarded_headers( json!({ "error": error }).to_string(), metadata.as_ref(), backend.is_forwarding_auth(), ), } } error => LlmClientError::ResponseTranslation(error.to_string()), } })?; // The gateway may report a buffered charge only in its response header. if agg.usage.provider_cost.is_none() { agg.usage.provider_cost = upstream_headers .get("x-litellm-response-cost") .and_then(|value| value.to_str().ok()) .and_then(|value| value.trim().parse::().ok()) .filter(|usd| usd.is_finite() && *usd >= 0.0); } (LlmResponse::Agg(agg), upstream_headers) } }; Ok(Response { llm_response, metadata, upstream_headers, }) } /// The whole decode → call → encode path a wire endpoint needs, in one call. /// /// Decodes `raw_http_request` from `wire_format` to the neutral IR, serves it via /// [`call_rewrite_model`](Self::call_rewrite_model) — the *upstream* wire format is /// resolved there from the model's backend, independently of `wire_format` — then /// encodes the neutral response back into `wire_format`. The result is a buffered /// [`RawResponse::Buffered`] JSON body or a streamed [`RawResponse::Stream`] of /// wire events (the caller frames the stream as SSE). The response's `model` is /// restamped with the model that actually served the call, so the body names the /// model that answered rather than the route the caller addressed. /// /// `http_headers` are carried through as the request's /// [`Metadata::http_headers`]. Backends with `forward_auth` disabled forward only /// allowed metadata headers; `forward_auth` backends forward all application /// headers. The forwarding backends in one route must use one credential family /// unless they all use the same scheme, host, and port. Transport headers are always /// rebuilt. Pass `None` to forward nothing. pub async fn call_rewrite_model_raw( &self, raw_http_request: Value, http_headers: Option, target_name: Option<&ModelId>, wire_format: WireFormat, ) -> Result { let llm_request = observed_translation::decode_request(wire_format, &raw_http_request) .map_err(|error| LlmClientError::RequestTranslation(error.to_string()))?; let request_extensions = llm_request.extensions.clone(); // Report the provider model for the same target selected by `call_rewrite_model`. let served_model = target_name .cloned() .or_else(|| llm_request.model.clone().map(ModelId::from)) .map(|target_name| self.provider_model(&target_name).to_string()); let request = Request { llm_request, raw_request: None, metadata: Some(Metadata { session_id: None, agent_id: None, task_id: None, correlation_id: None, extra_metadata: None, http_headers, wire_format: None, ..Default::default() }), }; let response = self.call_rewrite_model(request, target_name).await?; match response.llm_response { LlmResponse::Agg(agg) => { let body = observed_translation::encode_response( &agg, wire_format, served_model.as_deref(), &request_extensions, ) .map_err(|error| LlmClientError::ResponseTranslation(error.to_string()))?; Ok(RawResponse::Buffered(body)) } LlmResponse::Stream(chunks) => { let events = encode_stream_with_extensions( chunks, wire_format, served_model, &request_extensions, )?; Ok(RawResponse::Stream(events)) } } } } #[async_trait] impl RoutedLlmClient for TranslatingLlmClient { fn provider_model<'a>(&'a self, target: &'a ModelId) -> &'a ModelId { TranslatingLlmClient::provider_model(self, target) } async fn call(&self, request: Request) -> Result { self.call_rewrite_model(request, None).await } } #[derive(Clone, Copy)] enum UpstreamEndpoint { Completion, Auxiliary(AuxiliaryOperation), } impl UpstreamEndpoint { fn url(self, backend: &Backend) -> String { match self { UpstreamEndpoint::Completion => backend.url(), UpstreamEndpoint::Auxiliary(operation) => operation.url(backend), } } fn allows_streaming(self) -> bool { matches!(self, UpstreamEndpoint::Completion) } } enum EncodedResponse { Buffered { status: u16, body: Vec, upstream_headers: HeaderMap, }, Streaming { status: u16, chunks: LlmResponseStream, upstream_headers: HeaderMap, }, } impl EncodedResponse { fn status(&self) -> u16 { match self { EncodedResponse::Buffered { status, .. } | EncodedResponse::Streaming { status, .. } => *status, } } } // The typed error decides retry eligibility; status and Retry-After feed // attempt telemetry and delay selection. struct AttemptFailure { error: LlmClientError, status: Option, retry_after: Option, } fn deadline_error(timeout: Duration) -> LlmClientError { LlmClientError::Timeout { source: Box::new(std::io::Error::new( std::io::ErrorKind::TimedOut, format!( "response did not finish within {} ms, retries included", duration_millis(timeout) ), )), } } impl AttemptFailure { fn is_retryable(&self) -> bool { is_transient_failure(&self.error) } } fn is_transient_failure(error: &LlmClientError) -> bool { match error { LlmClientError::Transport { .. } | LlmClientError::Timeout { .. } => true, LlmClientError::UpstreamHttp { status, .. } => { metrics::is_retryable_http_status(status.as_u16()) } _ => false, } } fn stream_with_deadline( chunks: LlmResponseStream, deadline: Option<(tokio::time::Instant, Duration)>, ) -> LlmResponseStream { let Some((at, timeout)) = deadline else { return chunks; }; stream::unfold(Some(chunks), move |chunks| async move { let mut chunks = chunks?; let chunk = tokio::select! { biased; _ = tokio::time::sleep_until(at) => return Some((Err(deadline_error(timeout)), None)), chunk = chunks.next() => chunk?, }; let next = chunk.is_ok().then_some(chunks); Some((chunk, next)) }) .boxed() } async fn prepare_response_stream( response: reqwest::Response, backend: &Backend, target_name: &ModelId, ) -> Result { let bytes = response.bytes_stream().map(|chunk| { chunk .map(|bytes| bytes.to_vec()) .map_err(convert_reqwest_error) }); let mut chunks = decode_stream(bytes, backend.wire_format())?; match chunks.next().await { None => Ok(stream::empty().boxed()), // Nothing has reached the caller, so transport failures can still be retried. Some(Err(error @ (LlmClientError::Transport { .. } | LlmClientError::Timeout { .. }))) => { Err(error) } Some(first) => { // An in-band context overflow skips retries and advances to another candidate. if let Some(message) = first_event_overflow(&first, backend) { return Err(LlmClientError::ContextWindowExceeded { model: target_name.clone(), message, }); } Ok(stream::once(ready(first)).chain(chunks).boxed()) } } } // The overflow message when a stream's first event is an in-band provider rejection // of the whole request, rather than the start of a response. fn first_event_overflow( first: &Result, backend: &Backend, ) -> Option { first .as_ref() .ok()? .normalized() .iter() .find_map(|chunk| match chunk { LlmResponseChunk::StreamError { message } if backend.is_context_overflow(message) => { Some(message.clone()) } _ => None, }) } // Uses Retry-After when supplied, capped so an upstream cannot stall a request indefinitely. fn retry_after_delay(headers: &HeaderMap) -> Option { let value = headers.get(RETRY_AFTER)?.to_str().ok()?; let delay = if let Ok(seconds) = value.parse::() { Duration::from_secs(seconds) } else { let retry_at = httpdate::parse_http_date(value).ok()?; retry_at .duration_since(SystemTime::now()) .unwrap_or(Duration::ZERO) }; Some(delay.min(MAX_RETRY_AFTER)) } fn retry_delay(retry_number: u64, retry_after: Option) -> Duration { // Retry-After wins; otherwise double 250 ms up to the two-second cap. retry_after.unwrap_or_else(|| { let multiplier = 1_u32 << retry_number.min(3); INITIAL_RETRY_DELAY .saturating_mul(multiplier) .min(MAX_RETRY_BACKOFF) }) } fn duration_millis(duration: Duration) -> u64 { u64::try_from(duration.as_millis()).unwrap_or(u64::MAX) } fn record_gen_ai_request(url: &str, model: &str, streaming: bool) { let span = tracing::Span::current(); span.record("gen_ai.request.model", model); if streaming { span.record("gen_ai.request.stream", true); } if let Ok(url) = reqwest::Url::parse(url) { if let Some(host) = url.host_str() { span.record("server.address", host); } if let Some(port) = url.port_or_known_default() { span.record("server.port", i64::from(port)); } } } pub(crate) fn convert_reqwest_error(error: reqwest::Error) -> LlmClientError { // Reqwest labels truncated or otherwise unreadable response bodies as decode // errors, so distinguish them from serde JSON failures at the call site. let error = error.without_url(); // reqwest reports a connect-phase timeout as both `is_connect()` and // `is_timeout()`. Check connect first: no request was sent on the failed // connection, so it is a transport failure like a refused connection. if error.is_connect() { LlmClientError::Transport { source: Box::new(error), } } else if error.is_timeout() { LlmClientError::Timeout { source: Box::new(error), } } else if error.is_builder() { LlmClientError::Configuration { message: format!("failed to build upstream request: {error}"), } } else { LlmClientError::Transport { source: Box::new(error), } } } // Forwards safe metadata headers, or all application headers when forwarding auth. fn forward_metadata_headers( builder: RequestBuilder, metadata: Option<&Metadata>, backend: &Backend, ) -> RequestBuilder { let Some(headers) = metadata.and_then(|metadata| metadata.http_headers.as_ref()) else { return builder; }; let mut forwarded = HeaderMap::with_capacity(headers.len()); for (name, value) in headers { let is_allowed = if backend.is_forwarding_auth() { !is_non_forwardable_header(name.as_str(), headers) && !backend.is_provider_owned_header(name.as_str()) } else { is_allowed_metadata_header(name.as_str()) }; if !is_allowed { continue; } forwarded.append(name, value.clone()); } builder.headers(forwarded) } // Adds the backend's custom per-call headers. fn apply_extra_headers(mut builder: RequestBuilder, backend: &Backend) -> RequestBuilder { for (name, value) in backend.extra_headers() { builder = builder.header(name, value); } builder } // Overwrites the outbound body's `model` field with the resolved model id. fn set_json_model(body: &mut Value, model: &str) { if let Value::Object(object) = body { object.insert("model".to_string(), Value::String(model.to_string())); } } // The Codex Responses endpoint requires item-list input even on the preserved-body path. fn normalize_codex_string_input(url: &str, body: &mut Value) { let codex = reqwest::Url::parse(url).ok().is_some_and(|url| { url.path() .trim_end_matches('/') .ends_with("/backend-api/codex/responses") }); if codex && let Some(input @ Value::String(_)) = body.get_mut("input") { let text = input.take(); *input = json!([{"type": "message", "role": "user", "content": text}]); } } fn rejects_template_options(error: &LlmClientError) -> bool { let LlmClientError::UpstreamHttp { status, body } = error else { return false; }; if !matches!( *status, StatusCode::BAD_REQUEST | StatusCode::UNPROCESSABLE_ENTITY ) { return false; } let parsed = serde_json::from_str::(body).ok(); let error = parsed.as_ref().and_then(|value| value.get("error")); let message = match &parsed { Some(_) => error .and_then(|error| error.get("message")) .and_then(Value::as_str) .unwrap_or_default(), None => body, } .to_ascii_lowercase(); let field = error .and_then(|error| error.get("param")) .and_then(Value::as_str); let code = error .and_then(|error| error.get("code")) .and_then(Value::as_str); (field == Some("chat_template_kwargs") || message.contains("chat_template_kwargs")) && (matches!(code, Some("unknown_parameter" | "unsupported_parameter")) || [ "unrecognized request argument", "unknown parameter", "unsupported parameter", "unexpected keyword argument", "unknown field", "unsupported field", "extra inputs are not permitted", ] .iter() .any(|phrase| message.contains(phrase))) } /// Rewrites a Responses body into the shape OpenAI-compatible providers accept. /// /// Codex extends Responses with namespace containers and namespaced function calls, and /// OpenAI-compatible providers expect a flat Responses tool namespace. fn sanitize_openai_responses_provider_body(body: &mut Value) { let Value::Object(object) = body else { return; }; sanitize_openai_responses_input_for_provider(object.get_mut("input")); sanitize_openai_responses_tools_for_provider(object.get_mut("tools")); sanitize_openai_responses_tool_choice_for_provider(object.get_mut("tool_choice")); } fn sanitize_openai_responses_input_for_provider(input: Option<&mut Value>) { let Some(Value::Array(items)) = input else { return; }; for item in items { let Some(object) = item.as_object_mut() else { continue; }; if object.get("type").and_then(Value::as_str) == Some("function_call") { qualify_responses_function_name(object); } } } fn sanitize_openai_responses_tools_for_provider(tools: Option<&mut Value>) { let Some(Value::Array(tools)) = tools else { return; }; let mut flat_tools = Vec::with_capacity(tools.len()); for tool in std::mem::take(tools) { push_sanitized_openai_responses_tool(&mut flat_tools, tool); } *tools = flat_tools; } fn push_sanitized_openai_responses_tool(out: &mut Vec, tool: Value) { let Value::Object(mut object) = tool else { out.push(tool); return; }; if object.get("type").and_then(Value::as_str) == Some("namespace") { let namespace = object .get("name") .and_then(Value::as_str) .unwrap_or_default() .to_string(); let Some(Value::Array(children)) = object.remove("tools") else { out.push(Value::Object(object)); return; }; for child in children { push_sanitized_namespaced_tool(out, &namespace, child); } return; } ensure_responses_function_tool_description(&mut object); out.push(Value::Object(object)); } fn push_sanitized_namespaced_tool(out: &mut Vec, namespace: &str, tool: Value) { let Value::Object(mut object) = tool else { out.push(tool); return; }; if object.get("type").and_then(Value::as_str) == Some("function") { qualify_responses_function_name_with_namespace(&mut object, namespace); ensure_responses_function_tool_description(&mut object); } out.push(Value::Object(object)); } fn sanitize_openai_responses_tool_choice_for_provider(tool_choice: Option<&mut Value>) { let Some(Value::Object(object)) = tool_choice else { return; }; if object.get("type").and_then(Value::as_str) == Some("function") { qualify_responses_function_name(object); } } fn qualify_responses_function_name(object: &mut Map) { let namespace = object .remove("namespace") .and_then(|value| value.as_str().map(ToOwned::to_owned)); let Some(namespace) = namespace.as_deref() else { return; }; qualify_responses_function_name_with_namespace(object, namespace); } fn qualify_responses_function_name_with_namespace( object: &mut Map, namespace: &str, ) { if namespace.is_empty() { return; } let Some(name) = object.get("name").and_then(Value::as_str) else { return; }; let prefix = format!("{namespace}{NAMESPACE_SEPARATOR}"); if name.starts_with(&prefix) { return; } object.insert("name".to_string(), Value::String(format!("{prefix}{name}"))); } fn ensure_responses_function_tool_description(object: &mut Map) { if object.get("type").and_then(Value::as_str) == Some("function") && !matches!(object.get("description"), Some(Value::String(_))) { object.insert("description".to_string(), Value::String(String::new())); } } // Drops fields accepted by OpenAI-like APIs but rejected by Anthropic Messages. // // A router can serve earlier turns of a session from an OpenAI-format target and // later turns from an Anthropic one. Clients such as Claude Code send // `context_management` on every turn, so the Anthropic leg must strip it or the // upstream rejects the request (for example `clear_thinking_20251015` strategy // requires `thinking` to be enabled or adaptive). fn strip_anthropic_incompatible_fields(body: &mut Value) { if let Value::Object(object) = body { object.remove("reasoning_effort"); object.remove("context_management"); // `strict: false` is the default, and Bedrock rejects the field outright. if let Some(Value::Array(tools)) = object.get_mut("tools") { for tool in tools { if let Value::Object(tool) = tool && tool.get("strict") == Some(&Value::Bool(false)) { tool.remove("strict"); } } } } } // NVIDIA Inference API's Anthropic-compatible schema rejects the strict tool-use // extension that native Anthropic accepts. fn nvidia_inference_api_backend(backend: &Backend) -> bool { reqwest::Url::parse(backend.base_url()) .ok() .is_some_and(|url| { url.host_str() .is_some_and(|host| host.eq_ignore_ascii_case("inference-api.nvidia.com")) }) } fn strip_anthropic_tool_strict(body: &mut Value) { let Some(tools) = body.get_mut("tools").and_then(Value::as_array_mut) else { return; }; for tool in tools.iter_mut().filter_map(Value::as_object_mut) { tool.remove("strict"); } } // Removes replayed `thinking` blocks that carry no signature. // // Anthropic requires signed thinking blocks on replay. A router can serve earlier // turns of a session from an OpenAI-format target whose thinking blocks are // unsigned, so the Anthropic leg must drop them or the upstream rejects the // request. Bedrock enforces this (surfacing as a SigV4 signature mismatch) where // Azure-hosted Anthropic currently does not. fn strip_unsigned_thinking_blocks(body: &mut Value) { let Value::Object(object) = body else { return; }; let Some(Value::Array(messages)) = object.get_mut("messages") else { return; }; for message in messages { strip_unsigned_thinking_from_message(message); } } // Drops unsigned thinking blocks from one message, collapsing content that ends // up empty to an empty string so the message stays valid. fn strip_unsigned_thinking_from_message(message: &mut Value) { let Value::Object(message) = message else { return; }; let Some(Value::Array(blocks)) = message.get("content") else { return; }; if !blocks.iter().any(is_unsigned_thinking_block) { return; } let Some(Value::Array(blocks)) = message.get_mut("content") else { return; }; blocks.retain(|block| !is_unsigned_thinking_block(block)); if blocks.is_empty() { message.insert("content".to_string(), Value::String(String::new())); } } // A thinking block is unsigned when `signature` is absent or empty. fn is_unsigned_thinking_block(block: &Value) -> bool { if block.get("type").and_then(Value::as_str) != Some("thinking") { return false; } !matches!( block.get("signature").and_then(Value::as_str), Some(signature) if !signature.is_empty() ) } // Forces the target's configured reasoning effort onto the outbound body, replacing the // caller's value. Unlike `extra_body`, this is an override: a route that sends one model at a // higher effort than the client asked for is the point of the setting. fn apply_reasoning_effort(body: &mut Value, backend: &Backend) { let Some(effort) = backend.reasoning_effort() else { return; }; let Value::Object(object) = body else { return; }; match backend { // Responses nests effort under `reasoning`, Anthropic under `output_config`. Backend::OpenAiResponses(_) => set_nested_effort(object, "reasoning", effort), Backend::Anthropic(_) => set_nested_effort(object, "output_config", effort), Backend::OpenAiChat(_) => { // Chat Completions takes effort as a top-level field. object.insert( "reasoning_effort".to_string(), Value::String(effort.to_string()), ); } } } // Replaces only `effort` inside `key`, keeping fields the caller set next to it (`summary` or // `format`, for example). A non-object value there is malformed and is replaced whole. fn set_nested_effort(object: &mut serde_json::Map, key: &str, effort: &str) { let container = object .entry(key.to_string()) .or_insert_with(|| Value::Object(serde_json::Map::new())); if !container.is_object() { *container = Value::Object(serde_json::Map::new()); } if let Value::Object(container) = container { container.insert("effort".to_string(), Value::String(effort.to_string())); } } // Applies target defaults without overriding fields supplied by the caller. fn merge_extra_body(body: &mut Value, extra_body: &BTreeMap) { let Value::Object(object) = body else { return; }; for (key, value) in extra_body { object.entry(key.clone()).or_insert_with(|| value.clone()); } } fn omit_configured_body_fields(body: &mut Value, omit_body_fields: &BTreeSet) { let Value::Object(object) = body else { return; }; for key in omit_body_fields { object.remove(key); } } // Anthropic and Bedrock both cap a request at four blocks carrying // `cache_control`, counting tools, system blocks and message blocks together. const MAX_CACHE_CONTROL_BLOCKS: usize = 4; // Counts the blocks upstream will see carrying a `cache_control` marker. The // marker's own value holds no such key, so it is never counted twice. fn count_cache_control_blocks(value: &Value) -> usize { match value { Value::Object(map) => { usize::from(map.contains_key("cache_control")) + map.values().map(count_cache_control_blocks).sum::() } Value::Array(items) => items.iter().map(count_cache_control_blocks).sum(), _ => 0, } } // This helper marks the final caller content block before any temporary notes, if policy permits. fn enable_anthropic_prompt_caching(body: &mut Value, temporary_note_count: usize) { // Anthropic rejects a fifth cache marker with HTTP 400. // This check preserves the caller's markers when the request has reached that limit. if count_cache_control_blocks(body) >= MAX_CACHE_CONTROL_BLOCKS { return; } // Requests with only a top-level `1h` policy can return HTTP 400. // New block markers match the caller's TTL on turns with or without notes. let marker = body .get("cache_control") .cloned() .unwrap_or_else(|| serde_json::json!({"type":"ephemeral"})); let Some(messages) = body.get_mut("messages").and_then(Value::as_array_mut) else { return; }; let mut remaining = temporary_note_count; // Temporary guidance is absent next turn, so the breakpoint must precede it. for message in messages.iter_mut().rev() { match message.get_mut("content") { Some(content @ Value::String(_)) if remaining == 0 => { let text = std::mem::take(content); *content = serde_json::json!([{ "type": "text", "text": text, "cache_control": marker }]); return; } Some(Value::String(_)) => remaining = remaining.saturating_sub(1), Some(Value::Array(blocks)) => { if remaining >= blocks.len() { remaining -= blocks.len(); continue; } let index = blocks.len() - remaining - 1; if let Some(block) = blocks[index].as_object_mut() { block.entry("cache_control".to_string()).or_insert(marker); } return; } _ => {} } } } // GPT-5.6 and GPT-6 cache writes end at selected message or content boundaries. // A temporary suffix needs a separate cache entry before it so the next request can reuse the prefix. fn enable_openai_guidance_caching( body: &mut Value, format: WireFormat, model: &ModelId, suffix: usize, ) { let name = model.as_str().rsplit('/').next().unwrap_or_default(); let supported = name == "gpt-5.6" || name.starts_with("gpt-5.6-") || name == "gpt-6" || name.starts_with("gpt-6-") || name.starts_with("gpt-6."); // Explicit-only callers choose their write locations. // This helper adds markers only for names recognized by its model-name check. if suffix == 0 || !supported || body .pointer("/prompt_cache_options/mode") .and_then(Value::as_str) == Some("explicit") { return; } let key = if format == WireFormat::OpenAiResponses { "input" } else { "messages" }; let Some(items) = body.get_mut(key).and_then(Value::as_array_mut) else { return; }; let marker_count = items .iter() .map(|item| { ["content", "output"] .iter() .filter_map(|key| item.get(key).and_then(Value::as_array)) .flatten() .filter(|block| block.get("prompt_cache_breakpoint").is_some()) .count() }) .sum::(); // Implicit mode reserves one of the four write slots for the provider's final breakpoint. if marker_count >= 3 { return; } let Some(prefix_len) = items.len().checked_sub(suffix) else { return; }; for item in items[..prefix_len].iter_mut().rev() { let field = if item.get("type").and_then(Value::as_str) == Some("function_call_output") { "output" } else if matches!( item.get("role").and_then(Value::as_str), Some("user" | "tool" | "developer" | "system") ) { "content" } else { continue; }; let Some(content) = item.get_mut(field) else { continue; }; if let Value::String(text) = content { let text = std::mem::take(text); let kind = if format == WireFormat::OpenAiResponses { "input_text" } else { "text" }; *content = serde_json::json!([{"type":kind,"text":text}]); } // Empty content arrays have no block to mark, so the search continues through earlier items. if let Some(block) = content .as_array_mut() .and_then(|blocks| blocks.last_mut()) .and_then(Value::as_object_mut) { block .entry("prompt_cache_breakpoint".to_string()) .or_insert_with(|| serde_json::json!({"mode":"explicit"})); return; } } } // Requests streamed Chat usage by default while preserving an explicit caller choice. fn ensure_openai_stream_usage(body: &mut Value) { let Value::Object(object) = body else { return; }; if object.get("stream").and_then(Value::as_bool) != Some(true) { return; } match object.get_mut("stream_options") { Some(Value::Object(options)) => { options .entry("include_usage".to_string()) .or_insert(Value::Bool(true)); } _ => { let mut options = Map::new(); options.insert("include_usage".to_string(), Value::Bool(true)); object.insert("stream_options".to_string(), Value::Object(options)); } } } fn is_allowed_metadata_header(name: &str) -> bool { ALLOWED_METADATA_HEADERS .iter() .any(|allowed| name.eq_ignore_ascii_case(allowed)) } fn is_non_forwardable_header(name: &str, headers: &HeaderMap) -> bool { NON_FORWARDABLE_HEADERS .iter() .any(|blocked| name.eq_ignore_ascii_case(blocked)) || headers.get_all("connection").iter().any(|value| { value .as_bytes() .split(|byte| *byte == b',') .any(|option| option.trim_ascii().eq_ignore_ascii_case(name.as_bytes())) }) } // Unknown application headers can carry credentials when auth forwarding is enabled. fn redact_forwarded_headers( mut body: String, metadata: Option<&Metadata>, is_forwarding_auth: bool, ) -> String { if !is_forwarding_auth { return body; } let Some(headers) = metadata.and_then(|metadata| metadata.http_headers.as_ref()) else { return body; }; for (name, value) in headers { if is_non_forwardable_header(name.as_str(), headers) { continue; } let Ok(value) = value.to_str() else { continue; }; if !value.is_empty() { body = body.replace(value, "[REDACTED]"); } } body } #[cfg(test)] mod tests { use std::collections::BTreeMap; use std::error::Error; use std::io::{Read, Write}; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; use std::thread::JoinHandle; use serde_json::json; use switchyard_protocol::{completion_text, text_request}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; use super::*; use crate::backend::HttpBackendConfig; fn config(base_url: &str) -> HttpBackendConfig { HttpBackendConfig { base_url: base_url.to_string(), api_key: Some("secret".to_string()), forward_auth: false, extra_headers: BTreeMap::new(), extra_body: BTreeMap::new(), omit_body_fields: BTreeSet::new(), reasoning_effort: None, max_retries: 0, failure_cooldown: Duration::ZERO, timeout: None, } } #[test] fn provider_endpoint_urls_preserve_query_and_fragment() { let invalid_url = "not a URL/"; let backend = Backend::OpenAiResponses(config(invalid_url)); assert!(openai_url(invalid_url, "/responses").is_err()); assert_eq!(backend.url(), invalid_url); assert_eq!( AuxiliaryOperation::ResponsesCompact.url(&backend), invalid_url ); let anthropic = Backend::Anthropic(config(invalid_url)); assert_eq!(anthropic.url(), invalid_url); assert_eq!(anthropic.count_tokens_url(), invalid_url); for path in ["", "/v1/", "/v1/messages"] { for tail in [ "", "?api-version=2026-09-01&value=%2F/", "#fragment?value=/", ] { let base_url = format!("https://example.com{path}{tail}"); let backend = Backend::Anthropic(config(&base_url)); assert_eq!( backend.url(), format!("https://example.com/v1/messages{tail}") ); assert_eq!( AuxiliaryOperation::AnthropicCountTokens.url(&backend), format!("https://example.com/v1/messages/count_tokens{tail}") ); } } for path in ["/v1", "/v1/responses/", "/v1/chat/completions"] { for tail in [ "", "?api-version=2026-09-01&value=%2F/", "#fragment?value=/", ] { let base_url = format!("https://example.com{path}{tail}"); let responses = Backend::OpenAiResponses(config(&base_url)); let chat = Backend::OpenAiChat(config(&base_url)); for (actual, endpoint) in [ (responses.url(), "/responses"), (chat.url(), "/chat/completions"), ( AuxiliaryOperation::ResponsesCompact.url(&responses), "/responses/compact", ), ( AuxiliaryOperation::ResponsesInputTokens.url(&responses), "/responses/input_tokens", ), ] { assert_eq!(actual, format!("https://example.com/v1{endpoint}{tail}")); } } } } fn config_with_retries(base_url: &str, max_retries: u32) -> HttpBackendConfig { HttpBackendConfig { max_retries, ..config(base_url) } } fn forwarding_config(base_url: &str) -> HttpBackendConfig { HttpBackendConfig { api_key: None, forward_auth: true, ..config(base_url) } } // A one-model config list: "gpt" served over OpenAI Chat at base_url. fn chat_map(base_url: &str) -> Vec { vec![ModelConfig::new( "gpt", Backend::OpenAiChat(config(base_url)), None, )] } fn responses_map(base_url: &str) -> Vec { vec![ModelConfig::new( "gpt", Backend::OpenAiResponses(config(base_url)), None, )] } fn chat_map_with_extra_body( base_url: &str, extra_body: BTreeMap, ) -> Vec { let mut backend = config(base_url); backend.extra_body = extra_body; vec![ModelConfig::new("gpt", Backend::OpenAiChat(backend), None)] } fn chat_map_with_omit_body_fields( base_url: &str, omit_body_fields: BTreeSet, ) -> Vec { let mut backend = config(base_url); backend.omit_body_fields = omit_body_fields; vec![ModelConfig::new("gpt", Backend::OpenAiChat(backend), None)] } fn chat_map_with_effort(base_url: &str, effort: &str) -> Vec { let mut backend = config(base_url); backend.reasoning_effort = Some(effort.to_string()); vec![ModelConfig::new("gpt", Backend::OpenAiChat(backend), None)] } fn responses_map_with_effort(base_url: &str, effort: &str) -> Vec { let mut backend = config(base_url); backend.reasoning_effort = Some(effort.to_string()); vec![ModelConfig::new( "gpt", Backend::OpenAiResponses(backend), None, )] } fn anthropic_map(base_url: &str) -> Vec { vec![ModelConfig::new( "claude", Backend::Anthropic(config(base_url)), None, )] } #[tokio::test] async fn model_properties_uses_configured_auth_and_backend_path() -> std::result::Result<(), Box> { for (base_path, props_path) in [ ("/v1", "/props"), ("/v1/", "/props"), ("/v1/chat/completions", "/props"), ("/v1/responses/", "/props"), ("/prefix/v1", "/prefix/props"), ("/prefix/v1/chat/completions", "/prefix/props"), ("/prefix/v1/responses", "/prefix/props"), ("/prefix/", "/prefix/props"), ] { let server = MockServer::start().await; Mock::given(method("GET")) .and(path(props_path)) .and(wiremock::matchers::header("authorization", "Bearer secret")) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "default_generation_settings": {"n_ctx": 65_536} }))) .expect(1) .mount(&server) .await; let client = TranslatingLlmClient::new(&chat_map(&format!("{}{base_path}", server.uri())))?; let properties = client .get_model_properties( &ModelId::from("gpt"), WireFormat::OpenAiChat, Duration::from_secs(1), ) .await?; assert_eq!(properties["default_generation_settings"]["n_ctx"], 65_536); } Ok(()) } fn chat_map_with_retries(base_url: &str, max_retries: u32) -> Vec { vec![ModelConfig::new( "gpt", Backend::OpenAiChat(config_with_retries(base_url, max_retries)), None, )] } fn chat_success_response() -> ResponseTemplate { ResponseTemplate::new(200).set_body_json(json!({ "id": "chatcmpl-1", "model": "gpt", "choices": [{ "index": 0, "message": {"role": "assistant", "content": "recovered"}, "finish_reason": "stop" }], "usage": {} })) } fn truncated_response_server( content_type: &str, body: &str, ) -> std::io::Result<(String, JoinHandle>)> { response_sequence_server(vec![format!( "HTTP/1.1 200 OK\r\nContent-Type: {content_type}\r\n\ Content-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len() + 100 )]) } fn response_sequence_server( responses: Vec, ) -> std::io::Result<(String, JoinHandle>)> { let listener = std::net::TcpListener::bind("127.0.0.1:0")?; let address = listener.local_addr()?; let handle = std::thread::spawn(move || { for response in responses { let (mut stream, _) = listener.accept()?; let mut request = [0_u8; 1024]; if stream.read(&mut request)? == 0 { return Err(std::io::Error::new( std::io::ErrorKind::UnexpectedEof, "client closed before sending a request", )); } stream.write_all(response.as_bytes())?; } Ok(()) }); Ok((format!("http://{address}/v1"), handle)) } fn raw_chat_success_response() -> String { let body = r#"{"id":"chatcmpl-1","model":"gpt","choices":[{"index":0,"message":{"role":"assistant","content":"recovered"},"finish_reason":"stop"}],"usage":{}}"#; format!( "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n\ Content-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len() ) } fn request_for(model: Option<&str>, stream: bool) -> Request { let mut llm_request = text_request(model.map(str::to_string), "hi"); llm_request.stream = stream; Request { llm_request, raw_request: None, metadata: None, } } fn request_with_headers(model: &str, headers: HeaderMap) -> Request { let mut request = request_for(Some(model), false); request.metadata = Some(Metadata { http_headers: Some(headers), ..Default::default() }); request } #[tokio::test] async fn transport_errors_drop_the_upstream_url() { let error = reqwest::Client::new() .post("http://127.0.0.1:1/v1?key=CANARY") .send() .await .expect_err("closed port"); assert!(!convert_reqwest_error(error).to_string().contains("CANARY")); } #[tokio::test] async fn captured_request_build_failure_remains_a_configuration_error() { let client = TranslatingLlmClient::new(&chat_map("://invalid")).unwrap(); let result = crate::with_payload_capture(client.call(request_for(Some("gpt"), false))).await; assert!(matches!(result, Err(LlmClientError::Configuration { .. }))); } #[test] fn nvidia_anthropic_tools_drop_unsupported_strict_field() { let backend = Backend::Anthropic(config("https://inference-api.nvidia.com")); assert!(nvidia_inference_api_backend(&backend)); assert!(!nvidia_inference_api_backend(&Backend::Anthropic(config( "https://inference-api.nvidia.com.example.test" )))); let mut body = json!({"tools": [{"name": "a", "input_schema": {}, "strict": true}]}); strip_anthropic_tool_strict(&mut body); assert!(body["tools"][0].get("strict").is_none()); } #[test] fn anthropic_prompt_caching_marks_final_message() { let caller_managed = json!({ "cache_control": {"type": "ephemeral", "ttl": "1h"}, "messages": [{"role": "user", "content": "hello"}] }); let mut body = caller_managed.clone(); enable_anthropic_prompt_caching(&mut body, 0); assert_eq!(body["cache_control"], caller_managed["cache_control"]); assert_eq!( body["messages"][0]["content"], json!([{ "type": "text", "text": "hello", "cache_control": {"type": "ephemeral", "ttl": "1h"} }]) ); let mut body = json!({ "messages": [{"role": "user", "content": "hello"}] }); enable_anthropic_prompt_caching(&mut body, 0); assert_eq!( body["messages"][0]["content"][0]["cache_control"], json!({"type": "ephemeral"}) ); } // Counts the blocks upstream will see carrying a `cache_control` marker. fn cache_control_blocks(value: &Value) -> usize { count_cache_control_blocks(value) } // A body whose four breakpoints are all spent elsewhere: the caller manages // its own caching and left the final block unmarked on purpose. fn body_at_the_cache_control_limit() -> Value { json!({ "system": [ {"type": "text", "text": "a", "cache_control": {"type": "ephemeral"}}, {"type": "text", "text": "b", "cache_control": {"type": "ephemeral"}} ], "tools": [ {"name": "t", "input_schema": {"type": "object"}, "cache_control": {"type": "ephemeral"}} ], "messages": [ {"role": "user", "content": [ {"type": "text", "text": "old", "cache_control": {"type": "ephemeral"}} ]}, {"role": "assistant", "content": [{"type": "text", "text": "ok"}]}, {"role": "user", "content": [{"type": "text", "text": "new"}]} ] }) } #[test] fn anthropic_prompt_caching_abstains_at_the_cache_control_limit() { let mut body = body_at_the_cache_control_limit(); assert_eq!(cache_control_blocks(&body), 4); enable_anthropic_prompt_caching(&mut body, 0); // Anthropic and Bedrock both reject a fifth marker with HTTP 400, which // would fail a request that was valid before it reached us. assert_eq!(cache_control_blocks(&body), 4); assert!( body["messages"][2]["content"][0] .get("cache_control") .is_none() ); } #[test] fn anthropic_prompt_caching_still_marks_one_below_the_limit() { let mut body = body_at_the_cache_control_limit(); // Free one breakpoint, so there is room for ours. body["system"].as_array_mut().unwrap().pop(); assert_eq!(cache_control_blocks(&body), 3); enable_anthropic_prompt_caching(&mut body, 0); assert_eq!(cache_control_blocks(&body), 4); assert_eq!( body["messages"][2]["content"][0]["cache_control"], json!({"type": "ephemeral"}) ); } // A request that pins `format` in its metadata, so the client resolves that // wire format instead of the model's default backend. fn request_with_wire_format(model: &str, format: WireFormat) -> Request { let mut request = request_for(Some(model), false); request.metadata = Some(Metadata { session_id: None, agent_id: None, task_id: None, correlation_id: None, extra_metadata: None, http_headers: None, wire_format: Some(format), ..Default::default() }); request } #[tokio::test] async fn missing_model_errors() -> std::result::Result<(), Box> { let client = TranslatingLlmClient::new(&[])?; let Err(error) = client .call_rewrite_model(request_for(None, false), None) .await else { panic!("expected an error"); }; assert!(matches!( error, LlmClientError::InvalidRequest { message } if message == "no model given" )); Ok(()) } #[tokio::test] async fn unknown_model_errors() -> std::result::Result<(), Box> { let client = TranslatingLlmClient::new(&[])?; let Err(error) = client .call_rewrite_model(request_for(Some("gpt"), false), None) .await else { panic!("expected an error"); }; assert!(matches!( error, LlmClientError::Configuration { message } if message.contains("gpt") )); Ok(()) } #[tokio::test] async fn unknown_model_format_errors() -> std::result::Result<(), Box> { // "gpt" exists but only over OpenAI Chat; the request pins Anthropic. let client = TranslatingLlmClient::new(&chat_map("https://example.test/v1"))?; let Err(error) = client .call_rewrite_model( request_with_wire_format("gpt", WireFormat::AnthropicMessages), None, ) .await else { panic!("expected an error"); }; assert!(matches!( error, LlmClientError::Configuration { message } if message.contains("gpt") && message.contains(&WireFormat::AnthropicMessages.to_string()) )); Ok(()) } #[test] fn backend_for_resolves_configured_format() -> std::result::Result<(), Box> { let client = TranslatingLlmClient::new(&chat_map("https://example.test/v1"))?; // "gpt" is served over OpenAI Chat only; other formats and models miss. assert!( client .backend_for(&ModelId::from("gpt"), WireFormat::OpenAiChat) .is_some() ); assert!( client .backend_for(&ModelId::from("gpt"), WireFormat::AnthropicMessages) .is_none() ); assert!( client .backend_for(&ModelId::from("missing"), WireFormat::OpenAiChat) .is_none() ); Ok(()) } #[tokio::test] async fn target_name_arg_wins_over_request_model() -> std::result::Result<(), Box> { let client = TranslatingLlmClient::new(&[])?; // Arg "b" is looked up (and reported), not the request's "a". let Err(error) = client .call_rewrite_model(request_for(Some("a"), false), Some(&ModelId::from("b"))) .await else { panic!("expected an error"); }; assert!(matches!( error, LlmClientError::Configuration { message } if message.contains("\"b\"") )); Ok(()) } #[tokio::test] async fn buffered_openai_chat_round_trips() -> std::result::Result<(), Box> { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/v1/chat/completions")) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "id": "chatcmpl-1", "model": "gpt", "choices": [{ "index": 0, "message": {"role": "assistant", "content": "Hi there"}, "finish_reason": "stop" }], "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3} }))) .mount(&server) .await; let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?; let response = client .call_rewrite_model(request_for(Some("gpt"), false), None) .await?; let agg = response.llm_response.into_agg().await?; assert_eq!(completion_text(&agg), "Hi there"); Ok(()) } #[tokio::test] async fn invalid_json_is_a_response_translation_error() -> std::result::Result<(), Box> { let server = MockServer::start().await; let calls = Arc::new(AtomicUsize::new(0)); let observed_calls = Arc::clone(&calls); Mock::given(method("POST")) .respond_with(move |_: &wiremock::Request| { observed_calls.fetch_add(1, Ordering::SeqCst); ResponseTemplate::new(200).set_body_raw("not json", "application/json") }) .mount(&server) .await; let client = TranslatingLlmClient::new(&chat_map_with_retries(&format!("{}/v1", server.uri()), 2))?; let Err(error) = client .call_rewrite_model(request_for(Some("gpt"), false), None) .await else { panic!("expected invalid JSON to fail"); }; assert!(matches!( error, LlmClientError::ResponseTranslation(message) if message.contains("invalid upstream JSON") )); assert_eq!(calls.load(Ordering::SeqCst), 1); Ok(()) } #[tokio::test] async fn response_body_io_failure_is_a_transport_error() -> std::result::Result<(), Box> { let (base_url, server) = truncated_response_server("application/json", "{}")?; let client = TranslatingLlmClient::new(&chat_map(&base_url))?; let result = client .call_rewrite_model(request_for(Some("gpt"), false), None) .await; server .join() .map_err(|_| std::io::Error::other("response server thread panicked"))??; let Err(error) = result else { panic!("expected the truncated response body to fail"); }; let LlmClientError::Transport { source } = error else { panic!("expected a transport error"); }; let Some(source) = source.downcast_ref::() else { panic!("expected the reqwest transport source"); }; assert!(source.is_decode()); assert!( !std::error::Error::source(&source) .is_some_and(|source| source.is::()) ); Ok(()) } #[tokio::test] async fn response_body_transport_failures_are_retried() -> std::result::Result<(), Box> { let truncated_responses = [ "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n\ Content-Length: 102\r\nConnection: close\r\n\r\n{}", "HTTP/1.1 401 Unauthorized\r\nContent-Type: text/plain\r\n\ Content-Length: 100\r\nConnection: close\r\n\r\nbad", ]; for truncated in truncated_responses { let (base_url, server) = response_sequence_server(vec![truncated.to_string(), raw_chat_success_response()])?; let client = TranslatingLlmClient::new(&chat_map_with_retries(&base_url, 1))?; let response = client .call_rewrite_model(request_for(Some("gpt"), false), None) .await?; server .join() .map_err(|_| std::io::Error::other("response server thread panicked"))??; assert_eq!( completion_text(&response.llm_response.into_agg().await?), "recovered" ); } Ok(()) } #[tokio::test] async fn pre_first_event_transport_failure_uses_retry_budget() -> std::result::Result<(), Box> { let truncated = "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\n\ Content-Length: 100\r\nConnection: close\r\n\r\n"; let body = "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"recovered\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n"; let success = format!( "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\n\ Content-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len() ); for second in [success, truncated.to_string()] { let is_exhausted = second == truncated; let (base_url, server) = response_sequence_server(vec![truncated.to_string(), second])?; let client = TranslatingLlmClient::new(&chat_map_with_retries(&base_url, 1))?; let result = client .call_rewrite_model(request_for(Some("gpt"), true), None) .await; if is_exhausted { assert!(matches!(result, Err(LlmClientError::Transport { .. }))); } else { assert_eq!( completion_text(&result?.llm_response.into_agg().await?), "recovered" ); } server .join() .map_err(|_| std::io::Error::other("response server thread panicked"))??; } Ok(()) } #[tokio::test] async fn streaming_body_io_failure_preserves_transport_error() -> std::result::Result<(), Box> { let body = "data: {\"choices\":[{\"delta\":{\"content\":\"partial\"}}]}\n\n"; let (base_url, server) = truncated_response_server("text/event-stream", body)?; let client = TranslatingLlmClient::new(&chat_map_with_retries(&base_url, 2))?; let response = client .call_rewrite_model(request_for(Some("gpt"), true), None) .await?; let result = response.llm_response.into_agg().await; server .join() .map_err(|_| std::io::Error::other("response server thread panicked"))??; let Err(error) = result else { panic!("expected the truncated stream body to fail"); }; assert!(matches!(error, LlmClientError::Transport { .. })); Ok(()) } #[tokio::test] async fn rewrites_model_to_resolved_upstream_id() -> std::result::Result<(), Box> { // Inbound body says "switchyard"; the upstream must receive "gpt". let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/v1/chat/completions")) .and(wiremock::matchers::body_partial_json(json!({"model": "gpt"}))) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "id": "1", "model": "gpt", "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], "usage": {} }))) .mount(&server) .await; let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?; // Inbound model differs from the map key / resolved model. client .call_rewrite_model( request_for(Some("switchyard"), false), Some(&ModelId::from("gpt")), ) .await?; // The body_partial_json matcher asserts the upstream saw model "gpt". Ok(()) } /// A configured reasoning effort replaces the caller's value on both OpenAI wire formats, /// which `extra_body` (defaults only) cannot do. #[tokio::test] async fn reasoning_effort_override_replaces_the_callers_effort() -> std::result::Result<(), Box> { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/v1/chat/completions")) .and(wiremock::matchers::body_partial_json(json!({ "model": "gpt", "reasoning_effort": "max" }))) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "id": "1", "model": "gpt", "choices": [{ "index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop" }], "usage": {} }))) .mount(&server) .await; let client = TranslatingLlmClient::new(&chat_map_with_effort( &format!("{}/v1", server.uri()), "max", ))?; client .call_rewrite_model_raw( json!({ "model": "client-facing", "messages": [{"role": "user", "content": "hi"}], "reasoning_effort": "high" }), None, Some(&ModelId::from("gpt")), WireFormat::OpenAiChat, ) .await?; let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/v1/responses")) .and(wiremock::matchers::body_partial_json(json!({ "model": "gpt", "reasoning": {"effort": "max", "summary": "auto"} }))) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "id": "resp_1", "object": "response", "model": "gpt", "status": "completed", "output": [{ "type": "message", "id": "msg_1", "role": "assistant", "status": "completed", "content": [{"type": "output_text", "text": "ok"}] }], "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2} }))) .mount(&server) .await; let client = TranslatingLlmClient::new(&responses_map_with_effort( &format!("{}/v1", server.uri()), "max", ))?; client .call_rewrite_model_raw( json!({ "model": "client-facing", "input": [{"role": "user", "content": "hi"}], "reasoning": {"effort": "high", "summary": "auto"} }), None, Some(&ModelId::from("gpt")), WireFormat::OpenAiResponses, ) .await?; Ok(()) } // Covers both a translated request and an Anthropic caller's body sent as is. #[tokio::test] async fn anthropic_requests_drop_tool_strict_false() -> std::result::Result<(), Box> { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/v1/messages")) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "id": "msg_1", "type": "message", "role": "assistant", "model": "claude", "content": [{"type": "text", "text": "ok"}], "stop_reason": "end_turn", "usage": {"input_tokens": 1, "output_tokens": 1} }))) .expect(2) .mount(&server) .await; let client = TranslatingLlmClient::new(&[ModelConfig::new( "claude", Backend::Anthropic(config(&server.uri())), None, )])?; let schema = json!({"type": "object"}); let requests = [ ( WireFormat::OpenAiChat, json!({ "model": "client-facing", "messages": [{"role": "user", "content": "hi"}], "tools": [ {"type": "function", "function": {"name": "loose", "parameters": schema, "strict": false}}, {"type": "function", "function": {"name": "exact", "parameters": schema, "strict": true}} ] }), ), ( WireFormat::AnthropicMessages, json!({ "model": "client-facing", "max_tokens": 8, "messages": [{"role": "user", "content": "hi"}], "tools": [ {"name": "loose", "input_schema": schema, "strict": false}, {"name": "exact", "input_schema": schema, "strict": true} ] }), ), ]; for (format, request) in requests { client .call_rewrite_model_raw(request, None, Some(&ModelId::from("claude")), format) .await?; } for received in server.received_requests().await.unwrap_or_default() { let body: Value = serde_json::from_slice(&received.body)?; assert!(body["tools"][0].get("strict").is_none(), "{body}"); assert_eq!(body["tools"][1]["strict"], true, "{body}"); } Ok(()) } #[tokio::test] async fn template_rejection_retries_once_without_changing_other_fields() -> std::result::Result<(), Box> { for (message, reject_retry, omit) in [ ( "Unrecognized request argument supplied: chat_template_kwargs", false, false, ), ( "Unrecognized request argument supplied: chat_template_kwargs", true, false, ), ( "Invalid temperature for chat_template_kwargs request", false, false, ), ( "Unrecognized request argument supplied: chat_template_kwargs", false, true, ), ] { let server = MockServer::start().await; Mock::given(method("POST")) .respond_with(move |request: &wiremock::Request| { let body: Value = serde_json::from_slice(&request.body).expect("request JSON"); if body.get("chat_template_kwargs").is_some() || reject_retry { ResponseTemplate::new(400).set_body_json(json!({"error": {"message": message}})) } else { ResponseTemplate::new(200).set_body_json(json!({ "id": "1", "model": "gpt", "choices": [{"index": 0, "message": {"role": "assistant", "content": "yes"}, "finish_reason": "stop"}], "usage": {} })) } }).mount(&server).await; let fields = if omit { BTreeSet::from(["chat_template_kwargs".into()]) } else { BTreeSet::new() }; let client = TranslatingLlmClient::new(&chat_map_with_omit_body_fields( &format!("{}/v1", server.uri()), fields, ))?; let raw = json!({"model": "gpt", "messages": [{"role": "user", "content": "judge"}], "reasoning_effort": "none", "logprobs": true, "top_logprobs": 8, "chat_template_kwargs": {"enable_thinking": false}}); let result = client .call_rewrite_model_raw( raw.clone(), None, Some(&ModelId::from("gpt")), WireFormat::OpenAiChat, ) .await; let rejects_field = message.starts_with("Unrecognized"); assert_eq!(result.is_ok(), omit || (rejects_field && !reject_retry)); let calls = server.received_requests().await.ok_or("missing requests")?; assert_eq!(calls.len(), if !omit && rejects_field { 2 } else { 1 }); for (index, call) in calls.iter().enumerate() { let mut expected = raw.clone(); if omit || index > 0 { expected .as_object_mut() .unwrap() .remove("chat_template_kwargs"); } assert_eq!(serde_json::from_slice::(&call.body)?, expected); } } Ok(()) } #[tokio::test] async fn template_compatibility_preserves_the_transient_retry_budget() -> std::result::Result<(), Box> { let server = MockServer::start().await; let count = Arc::new(AtomicUsize::new(0)); let observed = Arc::clone(&count); Mock::given(method("POST")) .respond_with(move |request: &wiremock::Request| { let body: Value = serde_json::from_slice(&request.body).expect("request JSON"); match observed.fetch_add(1, Ordering::SeqCst) { 0 => { assert!(body.get("chat_template_kwargs").is_some()); ResponseTemplate::new(422).set_body_json(json!({"error": { "code": "unsupported_parameter", "param": "chat_template_kwargs" }})) } 1 => { assert!(body.get("chat_template_kwargs").is_none()); ResponseTemplate::new(503).insert_header("retry-after", "0") } _ => ResponseTemplate::new(200).set_body_json(json!({ "id": "1", "model": "gpt", "choices": [{"index": 0, "message": {"role": "assistant", "content": "yes"}, "finish_reason": "stop"}], "usage": {} })), } }).mount(&server).await; let client = TranslatingLlmClient::new(&chat_map_with_retries(&format!("{}/v1", server.uri()), 1))?; client .call_rewrite_model_raw( json!({"model": "gpt", "messages": [{"role": "user", "content": "judge"}], "chat_template_kwargs": {"enable_thinking": false}}), None, Some(&ModelId::from("gpt")), WireFormat::OpenAiChat, ) .await?; assert_eq!(count.load(Ordering::SeqCst), 3); Ok(()) } #[tokio::test] async fn preserved_codex_string_input_is_normalized_at_the_client_boundary() -> std::result::Result<(), Box> { for root in ["/backend-api/codex", "/v1"] { for input in [ json!("hello"), json!([{"role": "user", "content": "hello"}]), ] { let server = MockServer::start().await; Mock::given(method("POST")) .and(path(format!("{root}/responses"))) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "id": "resp_1", "object": "response", "model": "gpt", "status": "completed", "output": [], "usage": {} }))).expect(1).mount(&server).await; let client = TranslatingLlmClient::new(&responses_map(&format!("{}{root}", server.uri())))?; let raw = json!({"model": "client-name", "input": input, "instructions": "Be brief.", "store": false, "metadata": {"task": "judge"}, "custom_field": {"keep": true}}); client .call_rewrite_model_raw( raw.clone(), None, Some(&ModelId::from("gpt")), WireFormat::OpenAiResponses, ) .await?; let mut expected = raw; expected["model"] = json!("gpt"); if root == "/backend-api/codex" && input.is_string() { expected["input"] = json!([{"type": "message", "role": "user", "content": input}]); } let calls = server.received_requests().await.ok_or("missing requests")?; assert_eq!(serde_json::from_slice::(&calls[0].body)?, expected); } } Ok(()) } #[tokio::test] async fn token_counting_does_not_send_the_effort_override() -> std::result::Result<(), Box> { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/v1/messages/count_tokens")) .respond_with(ResponseTemplate::new(200).set_body_json(json!({"input_tokens": 3}))) .expect(1) .mount(&server) .await; let mut backend = config(&server.uri()); backend.reasoning_effort = Some("max".to_string()); let client = TranslatingLlmClient::new(&[ModelConfig::new( "claude", Backend::Anthropic(backend), None, )])?; let llm_request = switchyard_translation::decode_request( WireFormat::AnthropicMessages, &json!({"model": "claude", "messages": [{"role": "user", "content": "hi"}]}), )?; let request = Request { llm_request, raw_request: None, metadata: None, }; client .call_auxiliary( &ModelId::from("claude"), request, AuxiliaryOperation::AnthropicCountTokens, ) .await?; let received = server.received_requests().await.unwrap_or_default(); let body: Value = serde_json::from_slice(&received[0].body)?; assert!(body.get("output_config").is_none(), "{body}"); Ok(()) } #[tokio::test] async fn reasoning_effort_override_sets_anthropic_output_config_effort() -> std::result::Result<(), Box> { let server = MockServer::start().await; let format = json!({"type": "json_schema", "schema": {"type": "object"}}); Mock::given(method("POST")) .and(path("/v1/messages")) .and(wiremock::matchers::body_partial_json(json!({ "model": "claude", "output_config": {"effort": "max", "format": format} }))) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "id": "msg_1", "type": "message", "role": "assistant", "model": "claude", "content": [{"type": "text", "text": "ok"}], "stop_reason": "end_turn", "usage": {"input_tokens": 1, "output_tokens": 1} }))) .expect(1) .mount(&server) .await; let mut backend = config(&server.uri()); backend.reasoning_effort = Some("max".to_string()); let client = TranslatingLlmClient::new(&[ModelConfig::new( "claude", Backend::Anthropic(backend), None, )])?; client .call_rewrite_model_raw( json!({ "model": "client-facing", "max_tokens": 7, "messages": [{"role": "user", "content": "hi"}], "output_config": {"effort": "high", "format": format} }), None, Some(&ModelId::from("claude")), WireFormat::AnthropicMessages, ) .await?; Ok(()) } #[tokio::test] async fn extra_body_adds_defaults_without_overriding_the_request() -> std::result::Result<(), Box> { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/v1/chat/completions")) .and(wiremock::matchers::body_partial_json(json!({ "model": "gpt", "max_tokens": 7, "service_tier": "priority" }))) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "id": "1", "model": "gpt", "choices": [{ "index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop" }], "usage": {} }))) .mount(&server) .await; let extra_body = BTreeMap::from([ ("max_tokens".to_string(), json!(999)), ("service_tier".to_string(), json!("priority")), ]); let client = TranslatingLlmClient::new(&chat_map_with_extra_body( &format!("{}/v1", server.uri()), extra_body, ))?; let raw = json!({ "model": "client-facing", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 7 }); client .call_rewrite_model_raw( raw, None, Some(&ModelId::from("gpt")), WireFormat::OpenAiChat, ) .await?; Ok(()) } #[tokio::test] async fn omit_body_fields_strips_a_field_extra_body_cannot_override() -> std::result::Result<(), Box> { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/v1/chat/completions")) .and(|request: &wiremock::Request| { let body: Value = serde_json::from_slice(&request.body).unwrap_or(Value::Null); body.get("max_tokens").is_none() }) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "id": "1", "model": "gpt", "choices": [{ "index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop" }], "usage": {} }))) .mount(&server) .await; let client = TranslatingLlmClient::new(&chat_map_with_omit_body_fields( &format!("{}/v1", server.uri()), BTreeSet::from(["max_tokens".to_string()]), ))?; let raw = json!({ "model": "client-facing", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 7 }); client .call_rewrite_model_raw( raw, None, Some(&ModelId::from("gpt")), WireFormat::OpenAiChat, ) .await?; Ok(()) } // A weak OpenAI-format tier emits thinking blocks with no signature. Replaying // them to Anthropic is rejected (Bedrock reports it as a SigV4 mismatch), so // the Anthropic leg must drop them while keeping signed ones. #[tokio::test] async fn anthropic_requests_drop_unsigned_thinking_blocks() -> std::result::Result<(), Box> { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/v1/messages")) .and(|request: &wiremock::Request| { let body: Value = serde_json::from_slice(&request.body).unwrap_or(Value::Null); let messages = body.get("messages").and_then(Value::as_array).cloned(); let Some(messages) = messages else { return false; }; // The unsigned block is gone, the signed one survives, and the // message whose only block was unsigned is not left with an empty // content array. let blocks: Vec<&Value> = messages .iter() .filter_map(|message| message.get("content")) .filter_map(Value::as_array) .flatten() .collect(); let thinking: Vec<&&Value> = blocks .iter() .filter(|block| block.get("type").and_then(Value::as_str) == Some("thinking")) .collect(); thinking.len() == 1 && thinking[0].get("signature").and_then(Value::as_str) == Some("sig-abc") && messages .iter() .all(|message| message.get("content") != Some(&json!([]))) }) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "id": "msg_1", "type": "message", "role": "assistant", "model": "claude", "content": [{"type": "text", "text": "ok"}], "stop_reason": "end_turn", "usage": {"input_tokens": 1, "output_tokens": 1} }))) .mount(&server) .await; let client = TranslatingLlmClient::new(&anthropic_map(&server.uri()))?; let raw = json!({ "model": "client-facing", "max_tokens": 7, "messages": [ {"role": "user", "content": "fix the build"}, {"role": "assistant", "content": [ {"type": "thinking", "thinking": "weak tier reasoning", "signature": ""} ]}, {"role": "assistant", "content": [ {"type": "thinking", "thinking": "signed reasoning", "signature": "sig-abc"}, {"type": "text", "text": "here goes"} ]}, {"role": "user", "content": "continue"} ] }); client .call_rewrite_model_raw( raw, None, Some(&ModelId::from("claude")), WireFormat::AnthropicMessages, ) .await?; Ok(()) } #[tokio::test] async fn responses_reasoning_policy_preserves_replay_history() -> std::result::Result<(), Box> { use crate::ResponsesReasoningPolicy::{Auto, Drop, PreserveEncrypted, PreservePlaintext}; let summary = json!([{"type": "summary_text", "text": "Check the parser."}]); let history = json!([ {"type": "message", "role": "user", "content": "continue"}, {"type": "reasoning", "encrypted_content": "opaque", "content": [{"type": "reasoning_text", "text": "private"}]}, {"type": "reasoning", "id": "rs_stored", "status": "completed", "summary": []}, // GLM's plaintext forms, as Codex replays them. {"type": "reasoning", "id": "rs_glm", "summary": summary, "encrypted_content": null}, {"type": "reasoning", "content": [{"type": "reasoning_text", "text": "Run the tests."}], "encrypted_content": ""}, {"type": "reasoning"}, {"type": "function_call", "call_id": "call_1", "name": "shell", "arguments": "{}"}, {"type": "function_call_output", "call_id": "call_1", "output": "ok"} ]); // An explicit policy wins over the target. `Auto` follows the target. let cases = [ (Auto, "openai/gpt-5", PreserveEncrypted), (Auto, "nvidia/zai-org/glm-5.3-flash", PreservePlaintext), ( PreserveEncrypted, "nvidia/zai-org/glm-5.3-flash", PreserveEncrypted, ), (Drop, "openai/gpt-5", Drop), ]; for store in [Some(true), Some(false), None] { for (policy, model, behavior) in cases { for has_input_override in [false, true] { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/v1/responses")) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "id": "resp_1", "object": "response", "model": "gpt", "status": "completed", "output": [], "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2} }))) .expect(1) .mount(&server) .await; let mut backend = config(&format!("{}/v1", server.uri())); if has_input_override { backend.omit_body_fields.insert("input".to_string()); backend .extra_body .insert("input".to_string(), history.clone()); } let client = TranslatingLlmClient::new(&[ModelConfig::new( model, Backend::OpenAiResponses(backend), None, ) .with_responses_reasoning(policy)])?; let mut request = json!({"model": "switchyard", "input": history}); if let Some(store) = store { request["store"] = json!(store); } client .call_rewrite_model_raw( request, None, Some(&ModelId::from(model)), WireFormat::OpenAiResponses, ) .await?; let requests = server.received_requests().await.ok_or("missing requests")?; let body: Value = serde_json::from_slice(&requests[0].body)?; let stored = store != Some(false); let mut expected = vec![history[0].clone()]; match behavior { // OpenAI reads its own payload and stored IDs, but rejects plaintext. PreserveEncrypted => { let indexes: &[usize] = if stored { &[1, 2, 3] } else { &[1] }; for &index in indexes { let mut item = history[index].clone(); item["content"] = json!([]); expected.push(item); } } // Other targets keep plaintext and lose payloads of unknown origin. PreservePlaintext => { let mut item = history[1].clone(); item.as_object_mut() .ok_or("object")? .remove("encrypted_content"); expected.push(item); if stored { expected.push(history[2].clone()); } expected.extend([history[3].clone(), history[4].clone()]); } _ => {} } expected.extend([history[6].clone(), history[7].clone()]); assert_eq!( body["input"], json!(expected), "store={store:?}, policy={policy:?}, model={model}, replace={has_input_override}" ); } } } Ok(()) } #[tokio::test] async fn nvidia_anthropic_requests_strip_strict_from_extra_body() -> std::result::Result<(), Box> { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/v1/messages")) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "id": "msg_1", "type": "message", "role": "assistant", "model": "claude", "content": [{"type": "text", "text": "ok"}], "stop_reason": "end_turn", "usage": {"input_tokens": 1, "output_tokens": 1} }))) .expect(1) .mount(&server) .await; let mut backend = config(&format!( "http://inference-api.nvidia.com:{}", server.address().port() )); backend.extra_body.insert( "tools".into(), json!([{ "name": "lookup", "input_schema": {"type": "object", "properties": {}}, "strict": true }]), ); let mut client = TranslatingLlmClient::new(&[ModelConfig::new( "claude", Backend::Anthropic(backend), None, )])?; client.client = reqwest::Client::builder() .no_proxy() .resolve("inference-api.nvidia.com", *server.address()) .build()?; client .call_rewrite_model_raw( json!({ "model": "claude", "max_tokens": 7, "messages": [{"role": "user", "content": "hi"}] }), None, Some(&ModelId::from("claude")), WireFormat::AnthropicMessages, ) .await?; let requests = server.received_requests().await.unwrap(); let body: Value = serde_json::from_slice(&requests[0].body)?; assert_eq!(body["tools"][0]["name"], "lookup"); assert!(body["tools"][0].get("strict").is_none()); Ok(()) } // A router can serve earlier turns from an OpenAI target and later turns from // an Anthropic one, so the Anthropic leg must drop OpenAI-only fields the // caller keeps sending or the upstream rejects the whole request. #[tokio::test] async fn anthropic_requests_drop_openai_only_fields() -> std::result::Result<(), Box> { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/v1/messages")) .and(|request: &wiremock::Request| { let body: Value = serde_json::from_slice(&request.body).unwrap_or(Value::Null); body.get("context_management").is_none() && body.get("reasoning_effort").is_none() }) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "id": "msg_1", "type": "message", "role": "assistant", "model": "claude", "content": [{"type": "text", "text": "ok"}], "stop_reason": "end_turn", "usage": {"input_tokens": 1, "output_tokens": 1} }))) .mount(&server) .await; let client = TranslatingLlmClient::new(&anthropic_map(&server.uri()))?; let raw = json!({ "model": "client-facing", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 7, "reasoning_effort": "high", "context_management": { "edits": [{"type": "clear_thinking_20251015"}] } }); client .call_rewrite_model_raw( raw, None, Some(&ModelId::from("claude")), WireFormat::AnthropicMessages, ) .await?; Ok(()) } #[tokio::test] async fn streaming_openai_chat_aggregates() -> std::result::Result<(), Box> { let server = MockServer::start().await; let sse = "data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}\n\n\ data: {\"choices\":[{\"delta\":{\"content\":\" world\"}}]}\n\n\ data: {\"choices\":[],\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":2,\"total_tokens\":3}}\n\n\ data: [DONE]\n\n"; Mock::given(method("POST")) .and(path("/v1/chat/completions")) .and(wiremock::matchers::body_partial_json(json!({ "stream": true, "stream_options": {"include_usage": true} }))) .respond_with(ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream")) .mount(&server) .await; let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?; let response = client .call_rewrite_model(request_for(Some("gpt"), true), None) .await?; assert!(matches!(response.llm_response, LlmResponse::Stream(_))); let agg = response.llm_response.into_agg().await?; assert_eq!(completion_text(&agg), "Hello world"); assert_eq!(agg.usage.input_tokens, Some(1)); assert_eq!(agg.usage.output_tokens, Some(2)); assert_eq!(agg.usage.total_tokens, Some(3)); Ok(()) } #[tokio::test] async fn streaming_openai_chat_preserves_usage_opt_out() -> std::result::Result<(), Box> { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/v1/chat/completions")) .and(wiremock::matchers::body_partial_json(json!({ "stream": true, "stream_options": {"include_usage": false} }))) .respond_with( ResponseTemplate::new(200).set_body_raw("data: [DONE]\n\n", "text/event-stream"), ) .mount(&server) .await; let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?; let raw = json!({ "model": "client-facing", "messages": [{"role": "user", "content": "hi"}], "stream": true, "stream_options": {"include_usage": false} }); let response = client .call_rewrite_model_raw( raw, None, Some(&ModelId::from("gpt")), WireFormat::OpenAiChat, ) .await?; assert!(matches!(response, RawResponse::Stream(_))); Ok(()) } #[tokio::test] async fn upstream_500_is_upstream_http() -> std::result::Result<(), Box> { let server = MockServer::start().await; Mock::given(method("POST")) .respond_with(ResponseTemplate::new(500).set_body_string("boom")) .mount(&server) .await; let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?; let Err(error) = client .call_rewrite_model(request_for(Some("gpt"), false), None) .await else { panic!("expected an error"); }; assert!(matches!( error, LlmClientError::UpstreamHttp { status: StatusCode::INTERNAL_SERVER_ERROR, .. } )); Ok(()) } #[tokio::test] async fn retryable_http_failure_recovers_within_budget() -> std::result::Result<(), Box> { let server = MockServer::start().await; let calls = Arc::new(AtomicUsize::new(0)); let observed_calls = Arc::clone(&calls); Mock::given(method("POST")) .respond_with(move |_: &wiremock::Request| { if observed_calls.fetch_add(1, Ordering::SeqCst) == 0 { ResponseTemplate::new(503) .insert_header("retry-after", "0") .set_body_string("temporarily unavailable") } else { chat_success_response() } }) .mount(&server) .await; let client = TranslatingLlmClient::new(&chat_map_with_retries(&format!("{}/v1", server.uri()), 1))?; let response = client .call_rewrite_model(request_for(Some("gpt"), false), None) .await?; let agg = response.llm_response.into_agg().await?; assert_eq!(completion_text(&agg), "recovered"); assert_eq!(calls.load(Ordering::SeqCst), 2); Ok(()) } #[tokio::test] async fn deterministic_http_failure_is_not_retried() -> std::result::Result<(), Box> { let server = MockServer::start().await; let calls = Arc::new(AtomicUsize::new(0)); let observed_calls = Arc::clone(&calls); Mock::given(method("POST")) .respond_with(move |_: &wiremock::Request| { observed_calls.fetch_add(1, Ordering::SeqCst); ResponseTemplate::new(401).set_body_string("invalid key") }) .mount(&server) .await; let client = TranslatingLlmClient::new(&chat_map_with_retries(&format!("{}/v1", server.uri()), 2))?; let Err(error) = client .call_rewrite_model(request_for(Some("gpt"), false), None) .await else { panic!("expected an upstream error"); }; assert!(matches!( error, LlmClientError::UpstreamHttp { status: StatusCode::UNAUTHORIZED, body } if body == "invalid key" )); assert_eq!(calls.load(Ordering::SeqCst), 1); Ok(()) } #[tokio::test] async fn rate_limit_cooldown_depends_on_shared_credentials() -> std::result::Result<(), Box> { for forward_auth in [true, false] { let server = MockServer::start().await; Mock::given(method("POST")) .respond_with(ResponseTemplate::new(429).insert_header("retry-after", "0")) .expect(if forward_auth { 6 } else { 3 }) .mount(&server) .await; let backend = HttpBackendConfig { forward_auth, failure_cooldown: Duration::from_secs(60), ..config_with_retries(&format!("{}/v1", server.uri()), 2) }; let client = TranslatingLlmClient::new(&[ModelConfig::new( "gpt", Backend::OpenAiChat(backend), None, )])?; for credential in ["Bearer first-user", "Bearer second-user"] { let mut headers = http::HeaderMap::new(); headers.insert("authorization", http::HeaderValue::from_static(credential)); let result = client .call_rewrite_model(request_with_headers("gpt", headers), None) .await; if !forward_auth && credential == "Bearer second-user" { assert!(matches!( result, Err(LlmClientError::TemporarilyUnavailable) )); } else { assert!(matches!( result, Err(LlmClientError::UpstreamHttp { status: StatusCode::TOO_MANY_REQUESTS, .. }) )); } } } Ok(()) } #[tokio::test] async fn retry_exhaustion_returns_the_final_upstream_error() -> std::result::Result<(), Box> { let server = MockServer::start().await; let calls = Arc::new(AtomicUsize::new(0)); let observed_calls = Arc::clone(&calls); Mock::given(method("POST")) .respond_with(move |_: &wiremock::Request| { let attempt = observed_calls.fetch_add(1, Ordering::SeqCst) + 1; ResponseTemplate::new(500) .insert_header("retry-after", "0") .set_body_string(format!("attempt {attempt}")) }) .mount(&server) .await; let client = TranslatingLlmClient::new(&chat_map_with_retries(&format!("{}/v1", server.uri()), 2))?; let Err(error) = client .call_rewrite_model(request_for(Some("gpt"), false), None) .await else { panic!("expected retry exhaustion"); }; assert!(matches!( error, LlmClientError::UpstreamHttp { status: StatusCode::INTERNAL_SERVER_ERROR, body } if body == "attempt 3" )); assert_eq!(calls.load(Ordering::SeqCst), 3); Ok(()) } #[tokio::test] async fn deadline_expires_before_send_or_stream_poll() -> std::result::Result<(), Box> { let server = MockServer::start().await; Mock::given(method("POST")) .respond_with(ResponseTemplate::new(200).set_body_raw( "data: {\"id\":\"test\",\"model\":\"gpt\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"ok\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n", "text/event-stream", )) .mount(&server) .await; for timeout in [Duration::ZERO, Duration::from_millis(200)] { let mut backend = config(&server.uri()); backend.timeout = Some(timeout); let client = TranslatingLlmClient::new(&[ModelConfig::new( "gpt", Backend::OpenAiChat(backend), None, )])?; let result = client .call_rewrite_model(request_for(Some("gpt"), true), None) .await; if timeout.is_zero() { assert!(matches!(result, Err(LlmClientError::Timeout { .. }))); assert!( server .received_requests() .await .expect("recorded requests") .is_empty() ); } else { let LlmResponse::Stream(mut chunks) = result?.llm_response else { return Err("expected a streaming response".into()); }; tokio::time::sleep(Duration::from_millis(300)).await; assert!(matches!( chunks.next().await, Some(Err(LlmClientError::Timeout { .. })) )); assert!(chunks.next().await.is_none()); } } Ok(()) } #[tokio::test] async fn timeout_is_retried_before_a_response_is_returned() -> std::result::Result<(), Box> { let server = MockServer::start().await; let calls = Arc::new(AtomicUsize::new(0)); let observed_calls = Arc::clone(&calls); Mock::given(method("POST")) .respond_with(move |_: &wiremock::Request| { if observed_calls.fetch_add(1, Ordering::SeqCst) == 0 { ResponseTemplate::new(200).set_delay(Duration::from_millis(500)) } else { chat_success_response() } }) .mount(&server) .await; let mut client = TranslatingLlmClient::new(&chat_map_with_retries(&format!("{}/v1", server.uri()), 1))?; client.client = reqwest::Client::builder() .timeout(Duration::from_millis(100)) .build()?; let response = client .call_rewrite_model(request_for(Some("gpt"), false), None) .await?; assert_eq!( completion_text(&response.llm_response.into_agg().await?), "recovered" ); assert_eq!(calls.load(Ordering::SeqCst), 2); Ok(()) } #[test] fn retryable_error_classes_are_explicit() { let transport = AttemptFailure { error: LlmClientError::Transport { source: std::io::Error::other("disconnected").into(), }, status: None, retry_after: None, }; assert!(transport.is_retryable()); for status in [ StatusCode::REQUEST_TIMEOUT, StatusCode::TOO_MANY_REQUESTS, StatusCode::INTERNAL_SERVER_ERROR, StatusCode::SERVICE_UNAVAILABLE, StatusCode::from_u16(599) .expect("599 is a syntactically valid HTTP code in the http crate"), ] { let failure = AttemptFailure { error: LlmClientError::UpstreamHttp { status, body: String::new(), }, status: Some(status), retry_after: None, }; assert!(failure.is_retryable(), "HTTP {status} should retry"); } for status in [ StatusCode::BAD_REQUEST, StatusCode::UNAUTHORIZED, StatusCode::CONFLICT, StatusCode::from_u16(600) .expect("600 is a syntactically valid HTTP code in the http crate"), ] { let failure = AttemptFailure { error: LlmClientError::UpstreamHttp { status, body: String::new(), }, status: Some(status), retry_after: None, }; assert!(!failure.is_retryable(), "HTTP {status} should fail fast"); } let configuration = AttemptFailure { error: LlmClientError::Configuration { message: "invalid header".to_string(), }, status: None, retry_after: None, }; assert!(!configuration.is_retryable()); let context_window = AttemptFailure { error: LlmClientError::ContextWindowExceeded { model: ModelId::from("gpt"), message: "too long".to_string(), }, status: Some(StatusCode::BAD_REQUEST), retry_after: None, }; assert!(!context_window.is_retryable()); } #[test] fn retry_after_supports_seconds_and_http_dates() { let mut headers = HeaderMap::new(); headers.insert(RETRY_AFTER, reqwest::header::HeaderValue::from_static("3")); assert_eq!(retry_after_delay(&headers), Some(Duration::from_secs(3))); let retry_at = SystemTime::now() + Duration::from_secs(2); let value = httpdate::fmt_http_date(retry_at); let Ok(value) = reqwest::header::HeaderValue::from_str(&value) else { panic!("formatted HTTP date should be a valid header"); }; headers.insert(RETRY_AFTER, value); let Some(delay) = retry_after_delay(&headers) else { panic!("HTTP date should produce a retry delay"); }; assert!(delay <= Duration::from_secs(2)); } #[tokio::test] async fn routed_llm_client_exposes_timeout_variant() -> std::result::Result<(), Box> { let server = MockServer::start().await; Mock::given(method("POST")) .respond_with( ResponseTemplate::new(200).set_delay(std::time::Duration::from_millis(100)), ) .mount(&server) .await; let mut client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?; client.client = reqwest::Client::builder() .timeout(std::time::Duration::from_millis(10)) .build()?; let Err(error) = client.call(request_for(Some("gpt"), false)).await else { panic!("expected a timeout"); }; let LlmClientError::Timeout { source } = error else { panic!("expected the protocol timeout variant"); }; let Some(source) = source.downcast_ref::() else { panic!("expected the reqwest timeout source"); }; assert!(source.is_timeout()); Ok(()) } // A connect that never completes sends nothing, like a refused connection. #[cfg(target_os = "linux")] #[tokio::test] async fn connect_timeout_is_a_transport_error() -> std::result::Result<(), Box> { // Linux drops SYNs once a listener's accept queue is full, so later // connects hang like an upstream behind a dropping firewall. let socket = tokio::net::TcpSocket::new_v4()?; socket.bind("127.0.0.1:0".parse()?)?; let listener = socket.listen(0)?; let addr = listener.local_addr()?; let _queued: Vec<_> = (0..2) .map(|_| { std::net::TcpStream::connect_timeout(&addr, std::time::Duration::from_millis(100)) }) .collect(); let mut client = TranslatingLlmClient::new(&chat_map(&format!("http://{addr}/v1")))?; // Stands in for the OS connect timeout (30 s on Linux via reqwest's default // TCP_USER_TIMEOUT); both surface as `is_connect() && is_timeout()`. client.client = reqwest::Client::builder() .connect_timeout(std::time::Duration::from_millis(100)) .build()?; let Err(error) = client.call(request_for(Some("gpt"), false)).await else { panic!("expected the connect to fail"); }; let LlmClientError::Transport { source } = error else { panic!("expected Transport, got {error:?}"); }; let Some(source) = source.downcast_ref::() else { panic!("expected the reqwest connect-timeout source"); }; assert!(source.is_connect()); assert!(source.is_timeout()); Ok(()) } #[tokio::test] async fn llama_cpp_context_overflow_400_is_mapped() -> std::result::Result<(), Box> { let server = MockServer::start().await; Mock::given(method("POST")) .respond_with(ResponseTemplate::new(400).set_body_json(json!({ "error": { "message": "request (6016 tokens) exceeds the available context size (4096 tokens), try increasing it" } }))) .mount(&server) .await; let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?; let Err(error) = client .call_rewrite_model(request_for(Some("gpt"), false), None) .await else { panic!("expected an error"); }; assert!(matches!( error, LlmClientError::ContextWindowExceeded { model, .. } if model == "gpt" )); Ok(()) } #[tokio::test] async fn content_policy_400_preserves_upstream_error() -> std::result::Result<(), Box> { let server = MockServer::start().await; let body = json!({ "error": { "code": "content_policy_violation", "message": "request blocked by content policy", "type": "invalid_request_error" }, "metadata": {"documentation_section": "context window"} }); Mock::given(method("POST")) .respond_with(ResponseTemplate::new(400).set_body_json(&body)) .expect(1) .mount(&server) .await; let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?; let Err(error) = client .call_rewrite_model(request_for(Some("gpt"), false), None) .await else { panic!("expected a policy denial"); }; let LlmClientError::UpstreamHttp { status, body: actual, } = error else { panic!("expected the upstream HTTP error, got {error:?}"); }; assert_eq!(status, StatusCode::BAD_REQUEST); assert_eq!(serde_json::from_str::(&actual)?, body); Ok(()) } #[tokio::test] async fn forwards_only_allowlisted_metadata_headers() -> std::result::Result<(), Box> { let server = MockServer::start().await; Mock::given(method("POST")) .and(wiremock::matchers::header("x-request-id", "abc")) // A forwarded Authorization must NOT override the backend's bearer key. .and(wiremock::matchers::header("authorization", "Bearer secret")) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "id": "1", "model": "gpt", "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], "usage": {} }))) .mount(&server) .await; let mut headers = http::HeaderMap::new(); headers.insert("x-request-id", http::HeaderValue::from_static("abc")); headers.insert( "authorization", http::HeaderValue::from_static("Bearer client-key"), ); headers.insert( "x-goog-api-key", http::HeaderValue::from_static("client-google-key"), ); headers.insert( "x-custom-internal-secret", http::HeaderValue::from_static("client-custom-key"), ); let request = request_with_headers("gpt", headers); let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?; // The allowed request id survives and caller authorization cannot replace // the backend's credential. client.call_rewrite_model(request, None).await?; let received = server .received_requests() .await .ok_or("request recording should be enabled")?; let received = received.first().ok_or("expected one upstream request")?; assert!(!received.headers.contains_key("x-goog-api-key")); assert!(!received.headers.contains_key("x-custom-internal-secret")); Ok(()) } #[tokio::test] async fn forwards_application_headers_with_forward_auth() -> std::result::Result<(), Box> { let server = MockServer::start().await; Mock::given(method("POST")) .and(wiremock::matchers::header( "authorization", "Bearer client-key", )) .and(wiremock::matchers::header( "x-goog-api-key", "client-google-key", )) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "id": "1", "model": "gpt", "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], "usage": {} }))) .mount(&server) .await; let mut headers = http::HeaderMap::new(); headers.insert( "authorization", http::HeaderValue::from_static("Bearer client-key"), ); headers.insert( "x-goog-api-key", http::HeaderValue::from_static("client-google-key"), ); headers.insert("host", http::HeaderValue::from_static("client.example")); headers.insert("content-type", http::HeaderValue::from_static("text/plain")); headers.insert("connection", http::HeaderValue::from_static("x-hop-by-hop")); headers.insert( "x-hop-by-hop", http::HeaderValue::from_static("client-only"), ); let request = request_with_headers("gpt", headers); let backend_config = forwarding_config(&format!("{}/v1", server.uri())); let client = TranslatingLlmClient::new(&[ModelConfig::new( "gpt", Backend::OpenAiChat(backend_config), None, )])?; client.call_rewrite_model(request, None).await?; let received = server .received_requests() .await .ok_or("request recording should be enabled")?; let received = received.first().ok_or("expected one upstream request")?; assert_ne!(received.headers["host"], "client.example"); assert_eq!(received.headers["content-type"], "application/json"); assert!(!received.headers.contains_key("connection")); assert!(!received.headers.contains_key("x-hop-by-hop")); Ok(()) } #[tokio::test] async fn forward_auth_preserves_anthropic_headers_and_redacts_errors() -> std::result::Result<(), Box> { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/v1/messages")) .and(|request: &wiremock::Request| { request .headers .get("anthropic-beta") .and_then(|value| value.to_str().ok()) == Some("oauth-2025-04-20") && request.headers.get_all("anthropic-version").iter().count() == 1 && request .headers .get("anthropic-version") .and_then(|value| value.to_str().ok()) == Some("2023-06-01") }) .respond_with(ResponseTemplate::new(401).set_body_json(json!({ "error": {"message": "rejected client-google-key"} }))) .mount(&server) .await; let mut headers = http::HeaderMap::new(); headers.insert( "authorization", http::HeaderValue::from_static("Bearer client-key"), ); headers.insert( "x-goog-api-key", http::HeaderValue::from_static("client-google-key"), ); headers.insert( "anthropic-beta", http::HeaderValue::from_static("oauth-2025-04-20,prompt-caching-2024-07-31"), ); headers.insert( "anthropic-version", http::HeaderValue::from_static("caller-version"), ); let backend_config = forwarding_config(&server.uri()); let client = TranslatingLlmClient::new(&[ModelConfig::new( "claude", Backend::Anthropic(backend_config), None, )])?; let raw = json!({ "model": "claude", "max_tokens": 8, "messages": [{"role": "user", "content": "hello"}] }); let Err(LlmClientError::UpstreamHttp { body, .. }) = client .call_rewrite_model_raw( raw, Some(headers), Some(&ModelId::from("claude")), WireFormat::AnthropicMessages, ) .await else { panic!("expected an upstream HTTP error"); }; assert_eq!(body, r#"{"error":{"message":"rejected [REDACTED]"}}"#); Ok(()) } // Exercises the `RoutedLlmClient` impl: `call` uses the model already materialized in the // request and round-trips a buffered response. #[tokio::test] async fn routed_llm_client_serves_the_request_model() -> std::result::Result<(), Box> { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/v1/chat/completions")) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "id": "chatcmpl-1", "model": "gpt", "choices": [{ "index": 0, "message": {"role": "assistant", "content": "routed hi"}, "finish_reason": "stop" }], "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3} }))) .mount(&server) .await; let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?; let response = client.call(request_for(Some("gpt"), false)).await?; let agg = response.llm_response.into_agg().await?; assert_eq!(completion_text(&agg), "routed hi"); Ok(()) } #[tokio::test] async fn invalid_raw_request_is_a_request_translation_error() -> std::result::Result<(), Box> { let client = TranslatingLlmClient::new(&[])?; let Err(error) = client .call_rewrite_model_raw( json!("invalid"), None, Some(&ModelId::from("gpt")), WireFormat::OpenAiChat, ) .await else { panic!("expected request translation to fail"); }; assert!(matches!( error, LlmClientError::RequestTranslation(message) if !message.is_empty() )); Ok(()) } // Raw path, buffered: decode an OpenAI Chat body -> call -> encode back to OpenAI // Chat JSON, with the served `model` restamped over the id the caller addressed. #[tokio::test] async fn call_rewrite_model_raw_round_trips_buffered_json() -> std::result::Result<(), Box> { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/v1/chat/completions")) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "id": "chatcmpl-1", "model": "gpt", "choices": [{ "index": 0, "message": {"role": "assistant", "content": "Hi there"}, "finish_reason": "stop" }], "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3} }))) .mount(&server) .await; let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?; let raw = json!({ "model": "client-facing", "messages": [{"role": "user", "content": "hi"}] }); let RawResponse::Buffered(body) = client .call_rewrite_model_raw( raw, None, Some(&ModelId::from("gpt")), WireFormat::OpenAiChat, ) .await? else { panic!("expected a buffered response"); }; assert_eq!(body["choices"][0]["message"]["content"], "Hi there"); // The client sees the model that answered, not the "client-facing" route id. assert_eq!(body["model"], "gpt"); Ok(()) } // A Codex namespace is folded into the upstream tool name, then split back // into name and namespace on the Responses call that returns to Codex. #[tokio::test] async fn call_rewrite_model_raw_restores_codex_mcp_namespace() -> std::result::Result<(), Box> { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/v1/chat/completions")) .and(wiremock::matchers::body_partial_json(json!({ "tools": [{ "type": "function", "function": {"name": "mcp__open_websearch_sy_ns_search"} }] }))) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "id": "chatcmpl-1", "model": "gpt", "choices": [{ "index": 0, "message": { "role": "assistant", "content": null, "tool_calls": [{ "id": "call_1", "type": "function", "function": { "name": "mcp__open_websearch_sy_ns_search", "arguments": "{\"q\":\"rust\"}" } }] }, "finish_reason": "tool_calls" }], "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2} }))) .mount(&server) .await; let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?; let raw = json!({ "model": "client-facing", "input": "Search for Rust.", "tools": [{ "type": "namespace", "name": "mcp__open_websearch", "tools": [{ "type": "function", "name": "search", "description": "Search the web", "parameters": {"type": "object", "properties": {"q": {"type": "string"}}} }] }] }); let RawResponse::Buffered(body) = client .call_rewrite_model_raw( raw, None, Some(&ModelId::from("gpt")), WireFormat::OpenAiResponses, ) .await? else { panic!("expected a buffered response"); }; assert_eq!(body["output"][0]["type"], "function_call"); assert_eq!(body["output"][0]["name"], "search"); assert_eq!(body["output"][0]["namespace"], "mcp__open_websearch"); // Arguments are parsed and re-serialized, so the spacing is normalized. assert_eq!(body["output"][0]["arguments"], "{\"q\": \"rust\"}"); Ok(()) } #[tokio::test] async fn openai_responses_backend_flattens_codex_namespaces_before_upstream() -> std::result::Result<(), Box> { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/v1/responses")) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "id": "resp_1", "model": "gpt", "object": "response", "created_at": 0, "status": "completed", "output": [{ "type": "function_call", "call_id": "call_1", "name": "mcp__open_websearch_sy_ns_search", "arguments": "{\"q\":\"rust\"}" }], "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2} }))) .mount(&server) .await; let client = TranslatingLlmClient::new(&responses_map(&format!("{}/v1", server.uri())))?; let parameters = json!({"type": "object", "properties": {"q": {"type": "string"}}}); let raw = json!({ "model": "client-facing", "input": [{ "type": "function_call", "call_id": "call_0", "name": "search", "namespace": "mcp__open_websearch", "arguments": "{}" }], "tool_choice": { "type": "function", "name": "search", "namespace": "mcp__open_websearch" }, "tools": [{ "type": "namespace", "name": "mcp__open_websearch", "tools": [{ "type": "function", "name": "search", "parameters": parameters.clone() }] }] }); let RawResponse::Buffered(body) = client .call_rewrite_model_raw( raw, None, Some(&ModelId::from("gpt")), WireFormat::OpenAiResponses, ) .await? else { panic!("expected a buffered response"); }; let received = server .received_requests() .await .ok_or("request recording should be enabled")?; let request_body: Value = serde_json::from_slice(&received[0].body)?; assert_eq!(request_body["model"], "gpt"); assert_eq!( request_body["tools"], json!([{ "type": "function", "name": "mcp__open_websearch_sy_ns_search", "description": "", "parameters": parameters }]) ); assert_eq!( request_body["tool_choice"], json!({"type": "function", "name": "mcp__open_websearch_sy_ns_search"}) ); assert_eq!( request_body["input"][0]["name"], "mcp__open_websearch_sy_ns_search" ); assert!(request_body["input"][0].get("namespace").is_none()); assert_eq!(body["output"][0]["type"], "function_call"); assert_eq!(body["output"][0]["name"], "search"); assert_eq!(body["output"][0]["namespace"], "mcp__open_websearch"); Ok(()) } // A stage route can answer early steps of a Codex turn from Claude and a later step from // GPT, and Codex replays Claude's reasoning items on every step. OpenAI returns 400 for two // kinds: items with an ID but no `encrypted_content` when the request sets `store: false`, // and items whose `encrypted_content` holds Anthropic thinking that Switchyard packed for // the client. The Responses leg must drop both and keep GPT's own encrypted reasoning. // A non-OpenAI target instead keeps the plaintext and drops both payloads. #[tokio::test] async fn openai_responses_backend_drops_reasoning_the_provider_cannot_read() -> std::result::Result<(), Box> { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/v1/responses")) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "id": "resp_1", "model": "gpt", "object": "response", "created_at": 0, "status": "completed", "output": [{ "type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "ok"}] }], "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2} }))) .mount(&server) .await; let mut claude_reasoning = json!({ "type": "reasoning", "id": "rs_b99feb9a-763d-4f7f-b6b2-2c0199b47ded", "summary": [{"type": "summary_text", "text": "Read the files first."}], "encrypted_content": null }); // Build the item the way Switchyard does when it returns a Claude reply to a Responses // client, so the test follows any change to that format. let claude_reply = switchyard_translation::TranslationEngine::default() .translate_response( WireFormat::AnthropicMessages, WireFormat::OpenAiResponses, &json!({ "id": "msg_1", "type": "message", "role": "assistant", "model": "claude", "content": [ {"type": "thinking", "thinking": "Plan the edit.", "signature": "sig-abc"}, {"type": "text", "text": "ok"} ], "stop_reason": "end_turn", "usage": {"input_tokens": 1, "output_tokens": 1} }), &switchyard_translation::TranslationPolicy::default(), )? .body; let switchyard_reasoning = claude_reply["output"] .as_array() .and_then(|output| output.iter().find(|item| item["type"] == "reasoning")) .cloned() .ok_or("translated Claude reply should contain a reasoning item")?; assert!(switchyard_translation::is_anthropic_thinking_envelope( switchyard_reasoning["encrypted_content"] .as_str() .unwrap_or_default() )); let mut gpt_reasoning = json!({ "type": "reasoning", "id": "rs_gpt", "summary": [], "encrypted_content": "opaque-encrypted-reasoning" }); let mut empty_reasoning = json!({"type": "reasoning", "id": "rs_empty", "summary": [], "encrypted_content": ""}); let user = json!({"type": "message", "role": "user", "content": "fix the build"}); let call = json!({"type": "function_call", "call_id": "call_1", "name": "exec", "arguments": "{}"}); let output = json!({"type": "function_call_output", "call_id": "call_1", "output": "ok"}); let input = json!([ user, claude_reasoning, switchyard_reasoning, call, output, gpt_reasoning, empty_reasoning ]); let models = ["openai/gpt-5", "nvidia/zai-org/glm-5.3-flash"].map(|model| { ModelConfig::new( model, Backend::OpenAiResponses(config(&format!("{}/v1", server.uri()))), None, ) }); let client = TranslatingLlmClient::new(&models)?; for model in ["openai/gpt-5", "nvidia/zai-org/glm-5.3-flash"] { for store in [json!(false), json!(true)] { let raw = json!({"model": "client-facing", "store": store, "input": input.clone()}); client .call_rewrite_model_raw( raw, None, Some(&ModelId::from(model)), WireFormat::OpenAiResponses, ) .await?; } } let received = server .received_requests() .await .ok_or("request recording should be enabled")?; let sent: Vec = received .iter() .map(|request| serde_json::from_slice::(&request.body)) .collect::>()?; assert_eq!(sent.len(), 4); for request in &sent { assert!( request["input"] .as_array() .ok_or("expected input array")? .iter() .all(|item| !item["encrypted_content"] .as_str() .is_some_and(switchyard_translation::is_anthropic_thinking_envelope)), "provider cannot read translated Anthropic reasoning" ); } // GLM keeps the plaintext beside the Anthropic envelope and drops GPT's payload. let mut switchyard_plaintext = switchyard_reasoning.clone(); switchyard_plaintext .as_object_mut() .ok_or("object")? .remove("encrypted_content"); assert_eq!( sent[2]["input"], json!([user, claude_reasoning, switchyard_plaintext, call, output]) ); assert_eq!( sent[3]["input"], json!([ user, claude_reasoning, switchyard_plaintext, call, output, empty_reasoning ]) ); for item in [ &mut claude_reasoning, &mut gpt_reasoning, &mut empty_reasoning, ] { item["content"] = json!([]); } // With `store: false`, only GPT's own encrypted reasoning survives. assert_eq!(sent[0]["input"], json!([user, call, output, gpt_reasoning])); // With `store: true`, IDs can refer to stored reasoning, but OpenAI still cannot // decrypt the Switchyard item. assert_eq!( sent[1]["input"], json!([ user, claude_reasoning, call, output, gpt_reasoning, empty_reasoning ]) ); Ok(()) } // Raw path, streaming: an inbound `stream: true` request yields an unframed stream // of OpenAI Chat chunk objects whose deltas reassemble the completion. #[tokio::test] async fn call_rewrite_model_raw_streams_wire_events() -> std::result::Result<(), Box> { use futures::TryStreamExt; let server = MockServer::start().await; let sse = "data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}\n\n\ data: {\"choices\":[{\"delta\":{\"content\":\" world\"}}]}\n\n\ data: [DONE]\n\n"; Mock::given(method("POST")) .and(path("/v1/chat/completions")) .respond_with(ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream")) .mount(&server) .await; let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?; let raw = json!({ "model": "client-facing", "messages": [{"role": "user", "content": "hi"}], "stream": true }); let RawResponse::Stream(stream) = client .call_rewrite_model_raw( raw, None, Some(&ModelId::from("gpt")), WireFormat::OpenAiChat, ) .await? else { panic!("expected a streamed response"); }; let events: Vec = stream.try_collect().await?; assert!(!events.is_empty(), "expected at least one wire event"); let content: String = events .iter() .filter_map(|event| event["choices"][0]["delta"]["content"].as_str()) .collect(); assert_eq!(content, "Hello world"); // The mock frames carry no `model`, so every chunk's model comes from the // served id rather than the "unknown" fallback or the caller's route id. assert!(events.iter().all(|event| event["model"] == "gpt")); Ok(()) } // Raw path forwards allowed caller headers to the upstream. #[tokio::test] async fn call_rewrite_model_raw_forwards_headers() -> std::result::Result<(), Box> { let server = MockServer::start().await; Mock::given(method("POST")) .and(wiremock::matchers::header("x-request-id", "abc")) // A forwarded authorization must NOT override the backend's bearer key. .and(wiremock::matchers::header("authorization", "Bearer secret")) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "id": "1", "model": "gpt", "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], "usage": {} }))) .mount(&server) .await; let mut headers = http::HeaderMap::new(); headers.insert("x-request-id", http::HeaderValue::from_static("abc")); headers.insert( "authorization", http::HeaderValue::from_static("Bearer client-key"), ); let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?; let raw = json!({"model": "gpt", "messages": [{"role": "user", "content": "hi"}]}); // The allowed request id survives and caller authorization cannot replace // the backend's credential. client .call_rewrite_model_raw( raw, Some(headers), Some(&ModelId::from("gpt")), WireFormat::OpenAiChat, ) .await?; Ok(()) } }