1use std::borrow::Cow;
8use std::collections::{BTreeMap, BTreeSet, HashMap};
9use std::future::ready;
10use std::sync::atomic::{AtomicU64, Ordering};
11use std::time::{Duration, SystemTime};
12
13use async_trait::async_trait;
14use futures_util::{StreamExt, stream};
15use http::StatusCode;
16use reqwest::RequestBuilder;
17use reqwest::header::{HeaderMap, RETRY_AFTER};
18use serde_json::{Map, Value, json};
19use switchyard_protocol::{
20 LlmRequest, LlmResponse, LlmResponseChunk, LlmResponseStream, LlmResponseStreamEvent, Metadata,
21 ModelId, Request, Response, RoutedLlmClient,
22};
23use switchyard_translation::{
24 TranslationError, WireFormat, decode_stream, encode_stream_with_extensions,
25};
26use tokio::time::Instant;
27use tracing::Instrument;
28
29use crate::backend::{Backend, openai_url};
30use crate::error::{LlmClientError, Result};
31use crate::metrics;
32use crate::observed_translation;
33use crate::raw::RawResponse;
34
35const ALLOWED_METADATA_HEADERS: &[&str] = &["x-request-id"];
37
38const NON_FORWARDABLE_HEADERS: &[&str] = &[
41 "host",
42 "content-length",
43 "connection",
44 "proxy-connection",
45 "keep-alive",
46 "proxy-authenticate",
47 "proxy-authorization",
48 "te",
49 "trailer",
50 "transfer-encoding",
51 "upgrade",
52 "content-type",
53 "content-encoding",
54 "accept-encoding",
55 "expect",
56];
57
58const INITIAL_RETRY_DELAY: Duration = Duration::from_millis(250);
59const MAX_RETRY_BACKOFF: Duration = Duration::from_secs(2);
60const MAX_RETRY_AFTER: Duration = Duration::from_secs(60);
61
62#[derive(Clone, Debug)]
65pub struct ModelConfig {
66 model_name: ModelId,
67 default_backend: Backend,
68 other_backends: Option<Vec<Backend>>,
69 responses_reasoning: crate::ResponsesReasoningPolicy,
70}
71
72impl ModelConfig {
73 pub fn new(
76 model_name: impl Into<ModelId>,
77 default_backend: Backend,
78 other_backends: Option<Vec<Backend>>,
79 ) -> Self {
80 Self {
81 model_name: model_name.into(),
82 default_backend,
83 other_backends,
84 responses_reasoning: crate::ResponsesReasoningPolicy::default(),
85 }
86 }
87
88 #[must_use]
90 pub fn with_responses_reasoning(mut self, policy: crate::ResponsesReasoningPolicy) -> Self {
91 self.responses_reasoning = policy;
92 self
93 }
94}
95
96#[derive(Clone, Copy, Debug, Eq, PartialEq)]
98pub enum AuxiliaryOperation {
99 AnthropicCountTokens,
101 ResponsesInputTokens,
103 ResponsesCompact,
105}
106
107impl AuxiliaryOperation {
108 const fn wire_format(self) -> WireFormat {
109 match self {
110 Self::AnthropicCountTokens => WireFormat::AnthropicMessages,
111 Self::ResponsesInputTokens | Self::ResponsesCompact => WireFormat::OpenAiResponses,
112 }
113 }
114
115 fn url(self, backend: &Backend) -> String {
116 let suffix = match self {
117 Self::AnthropicCountTokens => return backend.count_tokens_url(),
118 Self::ResponsesInputTokens => "/responses/input_tokens",
119 Self::ResponsesCompact => "/responses/compact",
120 };
121 let base_url = backend.url();
122 openai_url(&base_url, suffix).unwrap_or_else(|error| {
123 tracing::error!(%error, "Unable to build OpenAI auxiliary endpoint URL");
124 base_url
125 })
126 }
127}
128
129pub struct TranslatingLlmClient {
138 model_to_config: HashMap<ModelId, ModelConfig>,
139 client: reqwest::Client,
140 forward_auth_client: reqwest::Client,
141 is_rate_limit_cooldown_enabled: bool,
142 cooldown_epoch: Instant,
143 unavailable_until: HashMap<ModelId, AtomicU64>,
145}
146
147impl TranslatingLlmClient {
148 pub fn new(model_configs: &[ModelConfig]) -> Result<Self> {
151 let mut unavailable_until = HashMap::new();
153 for config in model_configs {
154 for backend in std::iter::once(&config.default_backend)
155 .chain(config.other_backends.iter().flatten())
156 {
157 backend.validate_configured_headers(&config.model_name)?;
158 }
159 unavailable_until.insert(config.model_name.clone(), AtomicU64::new(0));
160 }
161 let build_client = |builder: reqwest::ClientBuilder| {
162 builder.build().map_err(|error| LlmClientError::Transport {
163 source: Box::new(error),
164 })
165 };
166 let client = build_client(reqwest::Client::builder())?;
167 let forward_auth_client =
170 build_client(reqwest::Client::builder().redirect(reqwest::redirect::Policy::none()))?;
171 let model_to_config = model_configs
172 .iter()
173 .map(|config| (config.model_name.clone(), config.clone()))
174 .collect();
175
176 Ok(Self {
177 model_to_config,
178 client,
179 forward_auth_client,
180 is_rate_limit_cooldown_enabled: true,
181 cooldown_epoch: Instant::now(),
182 unavailable_until,
183 })
184 }
185
186 pub fn without_rate_limit_cooldown(mut self) -> Self {
188 self.is_rate_limit_cooldown_enabled = false;
189 self
190 }
191
192 pub fn backend_for(&self, model: &ModelId, format: WireFormat) -> Option<&Backend> {
196 self.model_to_config.get(model).and_then(|config| {
197 if config.default_backend.wire_format() == format {
198 Some(&config.default_backend)
199 } else {
200 config
201 .other_backends
202 .as_ref()
203 .and_then(|backends| backends.iter().find(|b| b.wire_format() == format))
204 }
205 })
206 }
207
208 pub fn supports_auxiliary(&self, model: &ModelId, operation: AuxiliaryOperation) -> bool {
210 self.backend_for(model, operation.wire_format()).is_some()
211 }
212
213 pub async fn get_model_properties(
217 &self,
218 model: &ModelId,
219 format: WireFormat,
220 timeout: Duration,
221 ) -> Result<Value> {
222 let backend =
223 self.backend_for(model, format)
224 .ok_or_else(|| LlmClientError::Configuration {
225 message: format!("model {model} has no backend for {format:?}"),
226 })?;
227 let mut url = reqwest::Url::parse(backend.base_url()).map_err(|error| {
228 LlmClientError::Configuration {
229 message: format!("model {model} has an invalid backend URL: {error}"),
230 }
231 })?;
232 let root = url.path().trim_end_matches('/');
233 let root = root
234 .strip_suffix("/chat/completions")
235 .or_else(|| root.strip_suffix("/responses"))
236 .unwrap_or(root);
237 let path = format!(
238 "{}/props",
239 root.strip_suffix("/v1")
240 .unwrap_or(root)
241 .trim_end_matches('/')
242 );
243 url.set_path(&path);
244 url.set_query(None);
245
246 let builder = self.forward_auth_client.get(url).timeout(timeout);
247 let builder = apply_extra_headers(builder, backend);
248 let response = backend
249 .apply_auth(builder)
250 .send()
251 .await
252 .map_err(convert_reqwest_error)?
253 .error_for_status()
254 .map_err(convert_reqwest_error)?;
255 response
256 .json()
257 .await
258 .map_err(|source| LlmClientError::InvalidResponse {
259 source: Box::new(source),
260 })
261 }
262
263 pub async fn call_auxiliary(
268 &self,
269 model: &ModelId,
270 request: Request,
271 operation: AuxiliaryOperation,
272 ) -> Result<Value> {
273 let wire_format = operation.wire_format();
274 let backend =
275 self.backend_for(model, wire_format)
276 .ok_or_else(|| LlmClientError::Configuration {
277 message: format!("model {model} has no backend for {operation:?}"),
278 })?;
279 let Request {
280 mut llm_request,
281 metadata,
282 ..
283 } = request;
284 llm_request.model = Some(model.to_string());
285 let http_response = self
286 .send_encoded(
287 backend,
288 wire_format,
289 llm_request,
290 metadata.as_ref(),
291 model,
292 UpstreamEndpoint::Auxiliary(operation),
293 )
294 .await?;
295 let EncodedResponse::Buffered { body, .. } = http_response else {
296 return Err(LlmClientError::InvalidRequest {
297 message: "auxiliary endpoints do not support streaming".to_string(),
298 });
299 };
300 serde_json::from_slice(&body).map_err(|error| LlmClientError::InvalidResponse {
301 source: Box::new(error),
302 })
303 }
304
305 async fn send_encoded(
317 &self,
318 backend: &Backend,
319 wire_format: WireFormat,
320 llm_request: LlmRequest,
321 metadata: Option<&Metadata>,
322 model: &ModelId,
323 endpoint: UpstreamEndpoint,
324 ) -> Result<EncodedResponse> {
325 let cooldown = backend.failure_cooldown();
326 let unavailable_until =
327 if matches!(endpoint, UpstreamEndpoint::Completion) && !cooldown.is_zero() {
328 self.unavailable_until.get(model)
329 } else {
330 None
331 };
332 if let Some(until) = unavailable_until
333 && until.load(Ordering::Relaxed) > duration_millis(self.cooldown_epoch.elapsed())
334 {
335 tracing::debug!(model = %model, wire_format = %wire_format, "skipping backend during failure cooldown");
336 return Err(LlmClientError::TemporarilyUnavailable);
337 }
338
339 let mut body = observed_translation::encode_request(&llm_request, wire_format)
340 .map_err(|error| LlmClientError::RequestEncoding(error.to_string()))?;
341 set_json_model(&mut body, model);
345 if matches!(backend, Backend::OpenAiResponses(_)) {
346 sanitize_openai_responses_provider_body(&mut body);
347 }
348 if matches!(backend, Backend::Anthropic(_)) {
351 strip_anthropic_incompatible_fields(&mut body);
352 strip_unsigned_thinking_blocks(&mut body);
353 }
354 omit_configured_body_fields(&mut body, backend.omit_body_fields());
355 merge_extra_body(&mut body, backend.extra_body());
356 if matches!(backend, Backend::OpenAiResponses(_)) {
357 self.model_to_config
358 .get(model)
359 .map(|config| config.responses_reasoning)
360 .unwrap_or_default()
361 .normalize(&mut body);
362 }
363 apply_reasoning_effort(&mut body, backend);
366 if matches!(backend, Backend::OpenAiResponses(_)) {
367 normalize_codex_string_input(&backend.url(), &mut body);
368 }
369 if matches!(backend, Backend::Anthropic(_)) {
370 if nvidia_inference_api_backend(backend) {
371 strip_anthropic_tool_strict(&mut body);
372 }
373 enable_anthropic_prompt_caching(&mut body);
374 }
375 if matches!(backend, Backend::OpenAiChat(_)) {
376 ensure_openai_stream_usage(&mut body);
377 }
378 let streaming = endpoint.allows_streaming()
379 && body.get("stream").and_then(Value::as_bool).unwrap_or(false);
380 let url = endpoint.url(backend);
381 record_gen_ai_request(&url, model, streaming);
382
383 let result = self
384 .send_with_retries(&url, backend, &body, metadata, model, streaming)
385 .await;
386 if let Some(until) = unavailable_until
388 && let Err(error) = &result
389 && is_transient_failure(error)
390 && !((backend.is_forwarding_auth() || !self.is_rate_limit_cooldown_enabled)
391 && matches!(error, LlmClientError::UpstreamHttp { status, .. } if *status == StatusCode::TOO_MANY_REQUESTS))
392 {
393 let deadline = duration_millis(self.cooldown_epoch.elapsed().saturating_add(cooldown));
394 until.fetch_max(deadline, Ordering::Relaxed);
395 }
396 result
398 }
399
400 async fn send_with_retries(
402 &self,
403 url: &str,
404 backend: &Backend,
405 body: &Value,
406 metadata: Option<&Metadata>,
407 model: &ModelId,
408 streaming: bool,
409 ) -> Result<EncodedResponse> {
410 let max_retries = u64::from(backend.max_retries());
411 let max_attempts = max_retries
412 + 1
413 + u64::from(
414 matches!(backend, Backend::OpenAiChat(_))
415 && body.get("chat_template_kwargs").is_some(),
416 );
417 let mut body = Cow::Borrowed(body);
418 let mut template_retry = false;
419 let expires_at = backend
420 .timeout()
421 .map(|timeout| (tokio::time::sleep(timeout).deadline(), timeout));
422 let deadline = async {
423 match expires_at {
424 Some((at, timeout)) => {
425 tokio::time::sleep_until(at).await;
426 timeout
427 }
428 None => std::future::pending().await,
429 }
430 };
431 tokio::pin!(deadline);
432 let mut attempt = 0_u64;
433 loop {
434 if let Some((at, timeout)) = expires_at
435 && tokio::time::Instant::now() >= at
436 {
437 return Err(deadline_error(timeout));
438 }
439 let span = tracing::debug_span!(
440 target: "libsy",
441 "libsy.upstream_attempt",
442 model = %model,
443 wire_format = %backend.wire_format(),
444 attempt = attempt + 1 + u64::from(template_retry),
445 max_attempts,
446 retry = attempt > 0 || template_retry,
447 openinference.span.kind = "CHAIN",
448 outcome = tracing::field::Empty,
449 status_code = tracing::field::Empty,
450 will_retry = tracing::field::Empty,
451 retry_delay_ms = tracing::field::Empty,
452 );
453 let mut attempt_started = false;
454 let result = tokio::select! {
455 biased;
456 timeout = &mut deadline => {
457 if attempt_started {
458 metrics::record_upstream_attempt(None);
459 }
460 span.record("outcome", "error");
461 span.record("will_retry", false);
462 return Err(deadline_error(timeout));
463 }
464 result = async {
465 attempt_started = true;
466 self.send_once(url, backend, body.as_ref(), metadata, model, streaming).await
467 }.instrument(span.clone()) => result,
468 };
469 match result {
471 Ok(response) => {
472 span.record("outcome", "ok");
473 span.record("status_code", response.status());
474 span.record("will_retry", false);
475 if attempt > 0 || template_retry {
476 metrics::record_retry_recovered();
477 }
478 return Ok(match response {
479 EncodedResponse::Streaming {
480 status,
481 chunks,
482 upstream_headers,
483 } => EncodedResponse::Streaming {
484 status,
485 chunks: stream_with_deadline(chunks, expires_at),
486 upstream_headers,
487 },
488 buffered => buffered,
489 });
490 }
491 Err(failure) => {
492 let retry_without_template = matches!(backend, Backend::OpenAiChat(_))
493 && body.get("chat_template_kwargs")
494 == Some(&json!({"enable_thinking": false}))
495 && rejects_template_options(&failure.error);
496 let will_retry =
497 retry_without_template || (attempt < max_retries && failure.is_retryable());
498 span.record("outcome", "error");
499 if let Some(status) = failure.status {
500 span.record("status_code", status.as_u16());
501 }
502 span.record("will_retry", will_retry);
503 if !will_retry {
504 return Err(failure.error);
505 }
506
507 if retry_without_template {
508 if let Some(object) = body.to_mut().as_object_mut() {
509 object.remove("chat_template_kwargs");
510 }
511 template_retry = true;
512 span.record("retry_delay_ms", 0_u64);
513 continue;
514 }
515 let delay = retry_delay(attempt, failure.retry_after);
516 span.record("retry_delay_ms", duration_millis(delay));
517 drop(span);
519 tokio::select! {
520 biased;
521 timeout = &mut deadline => return Err(deadline_error(timeout)),
522 _ = tokio::time::sleep(delay) => {}
523 }
524 attempt += 1;
525 }
526 }
527 }
528 }
529
530 async fn send_once(
532 &self,
533 url: &str,
534 backend: &Backend,
535 body: &Value,
536 metadata: Option<&Metadata>,
537 model: &ModelId,
538 streaming: bool,
539 ) -> std::result::Result<EncodedResponse, AttemptFailure> {
540 let client = if backend.is_forwarding_auth() {
541 &self.forward_auth_client
542 } else {
543 &self.client
544 };
545 let builder = client.post(url).json(body);
546 let builder = forward_metadata_headers(builder, metadata, backend);
547 let builder = backend.apply_forwarded_auth(builder, metadata);
548 let builder = apply_extra_headers(builder, backend);
549 let builder = backend.apply_auth(builder);
550
551 let response = match builder.send().await {
552 Ok(response) => response,
553 Err(error) => {
554 metrics::record_upstream_attempt(None);
555 return Err(AttemptFailure {
556 error: convert_reqwest_error(error),
557 status: None,
558 retry_after: None,
559 });
560 }
561 };
562 let status = response.status();
563 if status.is_success() {
564 if streaming {
565 let upstream_headers = response.headers().clone();
566 let chunks = match prepare_response_stream(response, backend, model).await {
567 Ok(chunks) => chunks,
568 Err(error) => {
569 metrics::record_upstream_attempt(None);
570 return Err(AttemptFailure {
571 error,
572 status: Some(status),
573 retry_after: None,
574 });
575 }
576 };
577 metrics::record_upstream_attempt(Some(status.as_u16()));
578 return Ok(EncodedResponse::Streaming {
579 status: status.as_u16(),
580 chunks,
581 upstream_headers,
582 });
583 }
584 let upstream_headers = response.headers().clone();
585 let body = match response.bytes().await {
586 Ok(body) => body,
587 Err(error) => {
588 metrics::record_upstream_attempt(None);
589 return Err(AttemptFailure {
590 error: convert_reqwest_error(error),
591 status: Some(status),
592 retry_after: None,
593 });
594 }
595 };
596 metrics::record_upstream_attempt(Some(status.as_u16()));
597 return Ok(EncodedResponse::Buffered {
598 status: status.as_u16(),
599 body: body.to_vec(),
600 upstream_headers,
601 });
602 }
603
604 let retry_after = retry_after_delay(response.headers());
605 let body = match response.text().await {
606 Ok(body) => body,
607 Err(error) => {
608 metrics::record_upstream_attempt(None);
609 return Err(AttemptFailure {
610 error: convert_reqwest_error(error),
611 status: Some(status),
612 retry_after,
613 });
614 }
615 };
616 let body = redact_forwarded_headers(body, metadata, backend.is_forwarding_auth());
617 metrics::record_upstream_attempt(Some(status.as_u16()));
618 let error =
619 if status == reqwest::StatusCode::BAD_REQUEST && backend.is_context_overflow(&body) {
620 LlmClientError::ContextWindowExceeded {
621 model: model.clone(),
622 message: body,
623 }
624 } else {
625 LlmClientError::UpstreamHttp { status, body }
626 };
627 Err(AttemptFailure {
628 error,
629 status: Some(status),
630 retry_after,
631 })
632 }
633
634 pub async fn call_rewrite_model(
643 &self,
644 request: Request,
645 model_name: Option<&ModelId>,
646 ) -> Result<Response> {
647 let Request {
648 mut llm_request,
649 metadata,
650 ..
651 } = request;
652
653 let model_id = model_name
654 .cloned()
655 .or_else(|| llm_request.model.map(ModelId::from))
656 .ok_or_else(|| LlmClientError::InvalidRequest {
657 message: "no model given".to_string(),
658 })?;
659 llm_request.model = Some(model_id.to_string());
660
661 let orig_format = metadata.as_ref().and_then(|m| m.wire_format);
662 let wire_format = orig_format.unwrap_or(
663 self.model_to_config
664 .get(&model_id)
665 .map(|config| config.default_backend.wire_format())
666 .ok_or_else(|| LlmClientError::Configuration {
667 message: format!("no backend configured for model {model_id:?}"),
668 })?,
669 );
670 let backend = self.backend_for(&model_id, wire_format).ok_or_else(|| {
671 LlmClientError::Configuration {
672 message: format!("model {model_id:?} has no backend for format {wire_format}"),
673 }
674 })?;
675
676 let http_response = self
677 .send_encoded(
678 backend,
679 wire_format,
680 llm_request,
681 metadata.as_ref(),
682 &model_id,
683 UpstreamEndpoint::Completion,
684 )
685 .await?;
686
687 let (llm_response, upstream_headers) = match http_response {
688 EncodedResponse::Streaming {
689 chunks,
690 upstream_headers,
691 ..
692 } => (LlmResponse::Stream(chunks), upstream_headers),
693 EncodedResponse::Buffered {
694 body,
695 upstream_headers,
696 ..
697 } => {
698 let body = serde_json::from_slice::<Value>(&body).map_err(|error| {
699 LlmClientError::ResponseTranslation(format!("invalid upstream JSON: {error}"))
700 })?;
701 let mut agg =
704 observed_translation::decode_response(&body, wire_format).map_err(|error| {
705 match error {
706 TranslationError::UpstreamFailure { error } => {
707 LlmClientError::UpstreamHttp {
708 status: StatusCode::BAD_GATEWAY,
709 body: redact_forwarded_headers(
710 json!({ "error": error }).to_string(),
711 metadata.as_ref(),
712 backend.is_forwarding_auth(),
713 ),
714 }
715 }
716 error => LlmClientError::ResponseTranslation(error.to_string()),
717 }
718 })?;
719 if agg.usage.provider_cost.is_none() {
721 agg.usage.provider_cost = upstream_headers
722 .get("x-litellm-response-cost")
723 .and_then(|value| value.to_str().ok())
724 .and_then(|value| value.trim().parse::<f64>().ok())
725 .filter(|usd| usd.is_finite() && *usd >= 0.0);
726 }
727 (LlmResponse::Agg(agg), upstream_headers)
728 }
729 };
730
731 Ok(Response {
732 llm_response,
733 metadata,
734 upstream_headers,
735 })
736 }
737
738 pub async fn call_rewrite_model_raw(
756 &self,
757 raw_http_request: Value,
758 http_headers: Option<http::HeaderMap>,
759 model: Option<&ModelId>,
760 wire_format: WireFormat,
761 ) -> Result<RawResponse> {
762 let llm_request = observed_translation::decode_request(wire_format, &raw_http_request)
763 .map_err(|error| LlmClientError::RequestTranslation(error.to_string()))?;
764 let request_extensions = llm_request.extensions.clone();
765 let served_model = model
769 .map(ModelId::to_string)
770 .or_else(|| llm_request.model.clone());
771
772 let request = Request {
773 llm_request,
774 raw_request: None,
775 metadata: Some(Metadata {
776 session_id: None,
777 agent_id: None,
778 task_id: None,
779 correlation_id: None,
780 extra_metadata: None,
781 http_headers,
782 wire_format: None,
783 ..Default::default()
784 }),
785 };
786 let response = self.call_rewrite_model(request, model).await?;
787
788 match response.llm_response {
789 LlmResponse::Agg(agg) => {
790 let body = observed_translation::encode_response(
791 &agg,
792 wire_format,
793 served_model.as_deref(),
794 &request_extensions,
795 )
796 .map_err(|error| LlmClientError::ResponseTranslation(error.to_string()))?;
797 Ok(RawResponse::Buffered(body))
798 }
799 LlmResponse::Stream(chunks) => {
800 let events = encode_stream_with_extensions(
801 chunks,
802 wire_format,
803 served_model,
804 &request_extensions,
805 )?;
806 Ok(RawResponse::Stream(events))
807 }
808 }
809 }
810}
811
812#[async_trait]
813impl RoutedLlmClient for TranslatingLlmClient {
814 async fn call(&self, request: Request) -> Result<Response> {
815 self.call_rewrite_model(request, None).await
816 }
817}
818
819#[derive(Clone, Copy)]
820enum UpstreamEndpoint {
821 Completion,
822 Auxiliary(AuxiliaryOperation),
823}
824
825impl UpstreamEndpoint {
826 fn url(self, backend: &Backend) -> String {
827 match self {
828 UpstreamEndpoint::Completion => backend.url(),
829 UpstreamEndpoint::Auxiliary(operation) => operation.url(backend),
830 }
831 }
832
833 fn allows_streaming(self) -> bool {
834 matches!(self, UpstreamEndpoint::Completion)
835 }
836}
837
838enum EncodedResponse {
839 Buffered {
840 status: u16,
841 body: Vec<u8>,
842 upstream_headers: HeaderMap,
843 },
844 Streaming {
845 status: u16,
846 chunks: LlmResponseStream,
847 upstream_headers: HeaderMap,
848 },
849}
850
851impl EncodedResponse {
852 fn status(&self) -> u16 {
853 match self {
854 EncodedResponse::Buffered { status, .. }
855 | EncodedResponse::Streaming { status, .. } => *status,
856 }
857 }
858}
859
860struct AttemptFailure {
863 error: LlmClientError,
864 status: Option<StatusCode>,
865 retry_after: Option<Duration>,
866}
867
868fn deadline_error(timeout: Duration) -> LlmClientError {
869 LlmClientError::Timeout {
870 source: Box::new(std::io::Error::new(
871 std::io::ErrorKind::TimedOut,
872 format!(
873 "response did not finish within {} ms, retries included",
874 duration_millis(timeout)
875 ),
876 )),
877 }
878}
879
880impl AttemptFailure {
881 fn is_retryable(&self) -> bool {
882 is_transient_failure(&self.error)
883 }
884}
885
886fn is_transient_failure(error: &LlmClientError) -> bool {
887 match error {
888 LlmClientError::Transport { .. } | LlmClientError::Timeout { .. } => true,
889 LlmClientError::UpstreamHttp { status, .. } => {
890 metrics::is_retryable_http_status(status.as_u16())
891 }
892 _ => false,
893 }
894}
895
896fn stream_with_deadline(
897 chunks: LlmResponseStream,
898 deadline: Option<(tokio::time::Instant, Duration)>,
899) -> LlmResponseStream {
900 let Some((at, timeout)) = deadline else {
901 return chunks;
902 };
903 stream::unfold(Some(chunks), move |chunks| async move {
904 let mut chunks = chunks?;
905 let chunk = tokio::select! {
906 biased;
907 _ = tokio::time::sleep_until(at) => return Some((Err(deadline_error(timeout)), None)),
908 chunk = chunks.next() => chunk?,
909 };
910 let next = chunk.is_ok().then_some(chunks);
911 Some((chunk, next))
912 })
913 .boxed()
914}
915
916async fn prepare_response_stream(
917 response: reqwest::Response,
918 backend: &Backend,
919 model: &ModelId,
920) -> Result<LlmResponseStream> {
921 let bytes = response.bytes_stream().map(|chunk| {
922 chunk
923 .map(|bytes| bytes.to_vec())
924 .map_err(convert_reqwest_error)
925 });
926 let mut chunks = decode_stream(bytes, backend.wire_format())?;
927 match chunks.next().await {
928 None => Ok(stream::empty().boxed()),
929 Some(Err(error @ (LlmClientError::Transport { .. } | LlmClientError::Timeout { .. }))) => {
931 Err(error)
932 }
933 Some(first) => {
934 if let Some(message) = first_event_overflow(&first, backend) {
936 return Err(LlmClientError::ContextWindowExceeded {
937 model: model.clone(),
938 message,
939 });
940 }
941 Ok(stream::once(ready(first)).chain(chunks).boxed())
942 }
943 }
944}
945
946fn first_event_overflow(
949 first: &Result<LlmResponseStreamEvent>,
950 backend: &Backend,
951) -> Option<String> {
952 first
953 .as_ref()
954 .ok()?
955 .normalized()
956 .iter()
957 .find_map(|chunk| match chunk {
958 LlmResponseChunk::StreamError { message } if backend.is_context_overflow(message) => {
959 Some(message.clone())
960 }
961 _ => None,
962 })
963}
964
965fn retry_after_delay(headers: &HeaderMap) -> Option<Duration> {
967 let value = headers.get(RETRY_AFTER)?.to_str().ok()?;
968 let delay = if let Ok(seconds) = value.parse::<u64>() {
969 Duration::from_secs(seconds)
970 } else {
971 let retry_at = httpdate::parse_http_date(value).ok()?;
972 retry_at
973 .duration_since(SystemTime::now())
974 .unwrap_or(Duration::ZERO)
975 };
976 Some(delay.min(MAX_RETRY_AFTER))
977}
978
979fn retry_delay(retry_number: u64, retry_after: Option<Duration>) -> Duration {
980 retry_after.unwrap_or_else(|| {
982 let multiplier = 1_u32 << retry_number.min(3);
983 INITIAL_RETRY_DELAY
984 .saturating_mul(multiplier)
985 .min(MAX_RETRY_BACKOFF)
986 })
987}
988
989fn duration_millis(duration: Duration) -> u64 {
990 u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
991}
992
993fn record_gen_ai_request(url: &str, model: &str, streaming: bool) {
994 let span = tracing::Span::current();
995 span.record("gen_ai.request.model", model);
996 if streaming {
997 span.record("gen_ai.request.stream", true);
998 }
999 if let Ok(url) = reqwest::Url::parse(url) {
1000 if let Some(host) = url.host_str() {
1001 span.record("server.address", host);
1002 }
1003 if let Some(port) = url.port_or_known_default() {
1004 span.record("server.port", i64::from(port));
1005 }
1006 }
1007}
1008
1009pub(crate) fn convert_reqwest_error(error: reqwest::Error) -> LlmClientError {
1010 let error = error.without_url();
1013 if error.is_connect() {
1017 LlmClientError::Transport {
1018 source: Box::new(error),
1019 }
1020 } else if error.is_timeout() {
1021 LlmClientError::Timeout {
1022 source: Box::new(error),
1023 }
1024 } else if error.is_builder() {
1025 LlmClientError::Configuration {
1026 message: format!("failed to build upstream request: {error}"),
1027 }
1028 } else {
1029 LlmClientError::Transport {
1030 source: Box::new(error),
1031 }
1032 }
1033}
1034
1035fn forward_metadata_headers(
1037 builder: RequestBuilder,
1038 metadata: Option<&Metadata>,
1039 backend: &Backend,
1040) -> RequestBuilder {
1041 let Some(headers) = metadata.and_then(|metadata| metadata.http_headers.as_ref()) else {
1042 return builder;
1043 };
1044 let mut forwarded = HeaderMap::with_capacity(headers.len());
1045 for (name, value) in headers {
1046 let is_allowed = if backend.is_forwarding_auth() {
1047 !is_non_forwardable_header(name.as_str(), headers)
1048 && !backend.is_provider_owned_header(name.as_str())
1049 } else {
1050 is_allowed_metadata_header(name.as_str())
1051 };
1052 if !is_allowed {
1053 continue;
1054 }
1055 forwarded.append(name, value.clone());
1056 }
1057 builder.headers(forwarded)
1058}
1059
1060fn apply_extra_headers(mut builder: RequestBuilder, backend: &Backend) -> RequestBuilder {
1062 for (name, value) in backend.extra_headers() {
1063 builder = builder.header(name, value);
1064 }
1065 builder
1066}
1067
1068fn set_json_model(body: &mut Value, model: &str) {
1070 if let Value::Object(object) = body {
1071 object.insert("model".to_string(), Value::String(model.to_string()));
1072 }
1073}
1074
1075fn normalize_codex_string_input(url: &str, body: &mut Value) {
1077 let codex = reqwest::Url::parse(url).ok().is_some_and(|url| {
1078 url.path()
1079 .trim_end_matches('/')
1080 .ends_with("/backend-api/codex/responses")
1081 });
1082 if codex && let Some(input @ Value::String(_)) = body.get_mut("input") {
1083 let text = input.take();
1084 *input = json!([{"type": "message", "role": "user", "content": text}]);
1085 }
1086}
1087
1088fn rejects_template_options(error: &LlmClientError) -> bool {
1089 let LlmClientError::UpstreamHttp { status, body } = error else {
1090 return false;
1091 };
1092 if !matches!(
1093 *status,
1094 StatusCode::BAD_REQUEST | StatusCode::UNPROCESSABLE_ENTITY
1095 ) {
1096 return false;
1097 }
1098 let parsed = serde_json::from_str::<Value>(body).ok();
1099 let error = parsed.as_ref().and_then(|value| value.get("error"));
1100 let message = match &parsed {
1101 Some(_) => error
1102 .and_then(|error| error.get("message"))
1103 .and_then(Value::as_str)
1104 .unwrap_or_default(),
1105 None => body,
1106 }
1107 .to_ascii_lowercase();
1108 let field = error
1109 .and_then(|error| error.get("param"))
1110 .and_then(Value::as_str);
1111 let code = error
1112 .and_then(|error| error.get("code"))
1113 .and_then(Value::as_str);
1114 (field == Some("chat_template_kwargs") || message.contains("chat_template_kwargs"))
1115 && (matches!(code, Some("unknown_parameter" | "unsupported_parameter"))
1116 || [
1117 "unrecognized request argument",
1118 "unknown parameter",
1119 "unsupported parameter",
1120 "unexpected keyword argument",
1121 "unknown field",
1122 "unsupported field",
1123 "extra inputs are not permitted",
1124 ]
1125 .iter()
1126 .any(|phrase| message.contains(phrase)))
1127}
1128
1129const CODEX_NAMESPACE_SEPARATOR: &str = "__";
1130
1131fn sanitize_openai_responses_provider_body(body: &mut Value) {
1136 let Value::Object(object) = body else {
1137 return;
1138 };
1139 sanitize_openai_responses_input_for_provider(object.get_mut("input"));
1140 sanitize_openai_responses_tools_for_provider(object.get_mut("tools"));
1141 sanitize_openai_responses_tool_choice_for_provider(object.get_mut("tool_choice"));
1142}
1143
1144fn sanitize_openai_responses_input_for_provider(input: Option<&mut Value>) {
1145 let Some(Value::Array(items)) = input else {
1146 return;
1147 };
1148 for item in items {
1149 let Some(object) = item.as_object_mut() else {
1150 continue;
1151 };
1152 if object.get("type").and_then(Value::as_str) == Some("function_call") {
1153 qualify_responses_function_name(object);
1154 }
1155 }
1156}
1157
1158fn sanitize_openai_responses_tools_for_provider(tools: Option<&mut Value>) {
1159 let Some(Value::Array(tools)) = tools else {
1160 return;
1161 };
1162 let mut flat_tools = Vec::with_capacity(tools.len());
1163 for tool in std::mem::take(tools) {
1164 push_sanitized_openai_responses_tool(&mut flat_tools, tool);
1165 }
1166 *tools = flat_tools;
1167}
1168
1169fn push_sanitized_openai_responses_tool(out: &mut Vec<Value>, tool: Value) {
1170 let Value::Object(mut object) = tool else {
1171 out.push(tool);
1172 return;
1173 };
1174 if object.get("type").and_then(Value::as_str) == Some("namespace") {
1175 let namespace = object
1176 .get("name")
1177 .and_then(Value::as_str)
1178 .unwrap_or_default()
1179 .to_string();
1180 let Some(Value::Array(children)) = object.remove("tools") else {
1181 out.push(Value::Object(object));
1182 return;
1183 };
1184 for child in children {
1185 push_sanitized_namespaced_tool(out, &namespace, child);
1186 }
1187 return;
1188 }
1189 ensure_responses_function_tool_description(&mut object);
1190 out.push(Value::Object(object));
1191}
1192
1193fn push_sanitized_namespaced_tool(out: &mut Vec<Value>, namespace: &str, tool: Value) {
1194 let Value::Object(mut object) = tool else {
1195 out.push(tool);
1196 return;
1197 };
1198 if object.get("type").and_then(Value::as_str) == Some("function") {
1199 qualify_responses_function_name_with_namespace(&mut object, namespace);
1200 ensure_responses_function_tool_description(&mut object);
1201 }
1202 out.push(Value::Object(object));
1203}
1204
1205fn sanitize_openai_responses_tool_choice_for_provider(tool_choice: Option<&mut Value>) {
1206 let Some(Value::Object(object)) = tool_choice else {
1207 return;
1208 };
1209 if object.get("type").and_then(Value::as_str) == Some("function") {
1210 qualify_responses_function_name(object);
1211 }
1212}
1213
1214fn qualify_responses_function_name(object: &mut Map<String, Value>) {
1215 let namespace = object
1216 .remove("namespace")
1217 .and_then(|value| value.as_str().map(ToOwned::to_owned));
1218 let Some(namespace) = namespace.as_deref() else {
1219 return;
1220 };
1221 qualify_responses_function_name_with_namespace(object, namespace);
1222}
1223
1224fn qualify_responses_function_name_with_namespace(
1225 object: &mut Map<String, Value>,
1226 namespace: &str,
1227) {
1228 if namespace.is_empty() {
1229 return;
1230 }
1231 let Some(name) = object.get("name").and_then(Value::as_str) else {
1232 return;
1233 };
1234 let prefix = format!("{namespace}{CODEX_NAMESPACE_SEPARATOR}");
1235 if name.starts_with(&prefix) {
1236 return;
1237 }
1238 object.insert("name".to_string(), Value::String(format!("{prefix}{name}")));
1239}
1240
1241fn ensure_responses_function_tool_description(object: &mut Map<String, Value>) {
1242 if object.get("type").and_then(Value::as_str) == Some("function")
1243 && !matches!(object.get("description"), Some(Value::String(_)))
1244 {
1245 object.insert("description".to_string(), Value::String(String::new()));
1246 }
1247}
1248
1249fn strip_anthropic_incompatible_fields(body: &mut Value) {
1257 if let Value::Object(object) = body {
1258 object.remove("reasoning_effort");
1259 object.remove("context_management");
1260 }
1261}
1262
1263fn nvidia_inference_api_backend(backend: &Backend) -> bool {
1266 reqwest::Url::parse(backend.base_url())
1267 .ok()
1268 .is_some_and(|url| {
1269 url.host_str()
1270 .is_some_and(|host| host.eq_ignore_ascii_case("inference-api.nvidia.com"))
1271 })
1272}
1273
1274fn strip_anthropic_tool_strict(body: &mut Value) {
1275 let Some(tools) = body.get_mut("tools").and_then(Value::as_array_mut) else {
1276 return;
1277 };
1278 for tool in tools.iter_mut().filter_map(Value::as_object_mut) {
1279 tool.remove("strict");
1280 }
1281}
1282
1283fn strip_unsigned_thinking_blocks(body: &mut Value) {
1291 let Value::Object(object) = body else {
1292 return;
1293 };
1294 let Some(Value::Array(messages)) = object.get_mut("messages") else {
1295 return;
1296 };
1297 for message in messages {
1298 strip_unsigned_thinking_from_message(message);
1299 }
1300}
1301
1302fn strip_unsigned_thinking_from_message(message: &mut Value) {
1305 let Value::Object(message) = message else {
1306 return;
1307 };
1308 let Some(Value::Array(blocks)) = message.get("content") else {
1309 return;
1310 };
1311 if !blocks.iter().any(is_unsigned_thinking_block) {
1312 return;
1313 }
1314 let Some(Value::Array(blocks)) = message.get_mut("content") else {
1315 return;
1316 };
1317 blocks.retain(|block| !is_unsigned_thinking_block(block));
1318 if blocks.is_empty() {
1319 message.insert("content".to_string(), Value::String(String::new()));
1320 }
1321}
1322
1323fn is_unsigned_thinking_block(block: &Value) -> bool {
1325 if block.get("type").and_then(Value::as_str) != Some("thinking") {
1326 return false;
1327 }
1328 !matches!(
1329 block.get("signature").and_then(Value::as_str),
1330 Some(signature) if !signature.is_empty()
1331 )
1332}
1333
1334fn apply_reasoning_effort(body: &mut Value, backend: &Backend) {
1338 let Some(effort) = backend.reasoning_effort() else {
1339 return;
1340 };
1341 let Value::Object(object) = body else {
1342 return;
1343 };
1344 match backend {
1345 Backend::OpenAiResponses(_) => {
1346 let reasoning = object
1350 .entry("reasoning".to_string())
1351 .or_insert_with(|| Value::Object(serde_json::Map::new()));
1352 if !reasoning.is_object() {
1353 *reasoning = Value::Object(serde_json::Map::new());
1354 }
1355 if let Value::Object(reasoning) = reasoning {
1356 reasoning.insert("effort".to_string(), Value::String(effort.to_string()));
1357 }
1358 }
1359 Backend::OpenAiChat(_) => {
1360 object.insert(
1362 "reasoning_effort".to_string(),
1363 Value::String(effort.to_string()),
1364 );
1365 }
1366 Backend::Anthropic(_) => {}
1369 }
1370}
1371
1372fn merge_extra_body(body: &mut Value, extra_body: &BTreeMap<String, Value>) {
1374 let Value::Object(object) = body else {
1375 return;
1376 };
1377 for (key, value) in extra_body {
1378 object.entry(key.clone()).or_insert_with(|| value.clone());
1379 }
1380}
1381
1382fn omit_configured_body_fields(body: &mut Value, omit_body_fields: &BTreeSet<String>) {
1383 let Value::Object(object) = body else {
1384 return;
1385 };
1386 for key in omit_body_fields {
1387 object.remove(key);
1388 }
1389}
1390
1391const MAX_CACHE_CONTROL_BLOCKS: usize = 4;
1394
1395fn count_cache_control_blocks(value: &Value) -> usize {
1398 match value {
1399 Value::Object(map) => {
1400 usize::from(map.contains_key("cache_control"))
1401 + map.values().map(count_cache_control_blocks).sum::<usize>()
1402 }
1403 Value::Array(items) => items.iter().map(count_cache_control_blocks).sum(),
1404 _ => 0,
1405 }
1406}
1407
1408fn enable_anthropic_prompt_caching(body: &mut Value) {
1410 if body.get("cache_control").is_some() {
1412 return;
1413 }
1414 if count_cache_control_blocks(body) >= MAX_CACHE_CONTROL_BLOCKS {
1420 return;
1421 }
1422 let Some(content) = body
1423 .get_mut("messages")
1424 .and_then(Value::as_array_mut)
1425 .and_then(|messages| messages.last_mut())
1426 .and_then(|message| message.get_mut("content"))
1427 else {
1428 return;
1429 };
1430 match content {
1431 Value::String(text) => {
1432 *content = serde_json::json!([{
1433 "type": "text",
1434 "text": std::mem::take(text),
1435 "cache_control": {"type": "ephemeral"}
1436 }]);
1437 }
1438 Value::Array(blocks) => {
1439 if let Some(block) = blocks.last_mut().and_then(Value::as_object_mut) {
1440 block
1441 .entry("cache_control".to_string())
1442 .or_insert_with(|| serde_json::json!({"type": "ephemeral"}));
1443 }
1444 }
1445 _ => {}
1446 }
1447}
1448
1449fn ensure_openai_stream_usage(body: &mut Value) {
1451 let Value::Object(object) = body else {
1452 return;
1453 };
1454 if object.get("stream").and_then(Value::as_bool) != Some(true) {
1455 return;
1456 }
1457
1458 match object.get_mut("stream_options") {
1459 Some(Value::Object(options)) => {
1460 options
1461 .entry("include_usage".to_string())
1462 .or_insert(Value::Bool(true));
1463 }
1464 _ => {
1465 let mut options = Map::new();
1466 options.insert("include_usage".to_string(), Value::Bool(true));
1467 object.insert("stream_options".to_string(), Value::Object(options));
1468 }
1469 }
1470}
1471
1472fn is_allowed_metadata_header(name: &str) -> bool {
1473 ALLOWED_METADATA_HEADERS
1474 .iter()
1475 .any(|allowed| name.eq_ignore_ascii_case(allowed))
1476}
1477
1478fn is_non_forwardable_header(name: &str, headers: &HeaderMap) -> bool {
1479 NON_FORWARDABLE_HEADERS
1480 .iter()
1481 .any(|blocked| name.eq_ignore_ascii_case(blocked))
1482 || headers.get_all("connection").iter().any(|value| {
1483 value
1484 .as_bytes()
1485 .split(|byte| *byte == b',')
1486 .any(|option| option.trim_ascii().eq_ignore_ascii_case(name.as_bytes()))
1487 })
1488}
1489
1490fn redact_forwarded_headers(
1492 mut body: String,
1493 metadata: Option<&Metadata>,
1494 is_forwarding_auth: bool,
1495) -> String {
1496 if !is_forwarding_auth {
1497 return body;
1498 }
1499 let Some(headers) = metadata.and_then(|metadata| metadata.http_headers.as_ref()) else {
1500 return body;
1501 };
1502 for (name, value) in headers {
1503 if is_non_forwardable_header(name.as_str(), headers) {
1504 continue;
1505 }
1506 let Ok(value) = value.to_str() else {
1507 continue;
1508 };
1509 if !value.is_empty() {
1510 body = body.replace(value, "[REDACTED]");
1511 }
1512 }
1513 body
1514}
1515
1516#[cfg(test)]
1517mod tests {
1518 use std::collections::BTreeMap;
1519 use std::error::Error;
1520 use std::io::{Read, Write};
1521 use std::sync::Arc;
1522 use std::sync::atomic::{AtomicUsize, Ordering};
1523 use std::thread::JoinHandle;
1524
1525 use serde_json::json;
1526 use switchyard_protocol::{completion_text, text_request};
1527 use wiremock::matchers::{method, path};
1528 use wiremock::{Mock, MockServer, ResponseTemplate};
1529
1530 use super::*;
1531 use crate::backend::HttpBackendConfig;
1532
1533 fn config(base_url: &str) -> HttpBackendConfig {
1534 HttpBackendConfig {
1535 base_url: base_url.to_string(),
1536 api_key: Some("secret".to_string()),
1537 forward_auth: false,
1538 extra_headers: BTreeMap::new(),
1539 extra_body: BTreeMap::new(),
1540 omit_body_fields: BTreeSet::new(),
1541 reasoning_effort: None,
1542 max_retries: 0,
1543 failure_cooldown: Duration::ZERO,
1544 timeout: None,
1545 }
1546 }
1547
1548 #[test]
1549 fn provider_endpoint_urls_preserve_query_and_fragment() {
1550 let invalid_url = "not a URL/";
1551 let backend = Backend::OpenAiResponses(config(invalid_url));
1552 assert!(openai_url(invalid_url, "/responses").is_err());
1553 assert_eq!(backend.url(), invalid_url);
1554 assert_eq!(
1555 AuxiliaryOperation::ResponsesCompact.url(&backend),
1556 invalid_url
1557 );
1558
1559 let anthropic = Backend::Anthropic(config(invalid_url));
1560 assert_eq!(anthropic.url(), invalid_url);
1561 assert_eq!(anthropic.count_tokens_url(), invalid_url);
1562
1563 for path in ["", "/v1/", "/v1/messages"] {
1564 for tail in [
1565 "",
1566 "?api-version=2026-09-01&value=%2F/",
1567 "#fragment?value=/",
1568 ] {
1569 let base_url = format!("https://example.com{path}{tail}");
1570 let backend = Backend::Anthropic(config(&base_url));
1571 assert_eq!(
1572 backend.url(),
1573 format!("https://example.com/v1/messages{tail}")
1574 );
1575 assert_eq!(
1576 AuxiliaryOperation::AnthropicCountTokens.url(&backend),
1577 format!("https://example.com/v1/messages/count_tokens{tail}")
1578 );
1579 }
1580 }
1581
1582 for path in ["/v1", "/v1/responses/", "/v1/chat/completions"] {
1583 for tail in [
1584 "",
1585 "?api-version=2026-09-01&value=%2F/",
1586 "#fragment?value=/",
1587 ] {
1588 let base_url = format!("https://example.com{path}{tail}");
1589 let responses = Backend::OpenAiResponses(config(&base_url));
1590 let chat = Backend::OpenAiChat(config(&base_url));
1591 for (actual, endpoint) in [
1592 (responses.url(), "/responses"),
1593 (chat.url(), "/chat/completions"),
1594 (
1595 AuxiliaryOperation::ResponsesCompact.url(&responses),
1596 "/responses/compact",
1597 ),
1598 (
1599 AuxiliaryOperation::ResponsesInputTokens.url(&responses),
1600 "/responses/input_tokens",
1601 ),
1602 ] {
1603 assert_eq!(actual, format!("https://example.com/v1{endpoint}{tail}"));
1604 }
1605 }
1606 }
1607 }
1608
1609 fn config_with_retries(base_url: &str, max_retries: u32) -> HttpBackendConfig {
1610 HttpBackendConfig {
1611 max_retries,
1612 ..config(base_url)
1613 }
1614 }
1615
1616 fn forwarding_config(base_url: &str) -> HttpBackendConfig {
1617 HttpBackendConfig {
1618 api_key: None,
1619 forward_auth: true,
1620 ..config(base_url)
1621 }
1622 }
1623
1624 fn chat_map(base_url: &str) -> Vec<ModelConfig> {
1626 vec![ModelConfig::new(
1627 "gpt",
1628 Backend::OpenAiChat(config(base_url)),
1629 None,
1630 )]
1631 }
1632
1633 fn responses_map(base_url: &str) -> Vec<ModelConfig> {
1634 vec![ModelConfig::new(
1635 "gpt",
1636 Backend::OpenAiResponses(config(base_url)),
1637 None,
1638 )]
1639 }
1640
1641 fn chat_map_with_extra_body(
1642 base_url: &str,
1643 extra_body: BTreeMap<String, Value>,
1644 ) -> Vec<ModelConfig> {
1645 let mut backend = config(base_url);
1646 backend.extra_body = extra_body;
1647 vec![ModelConfig::new("gpt", Backend::OpenAiChat(backend), None)]
1648 }
1649
1650 fn chat_map_with_omit_body_fields(
1651 base_url: &str,
1652 omit_body_fields: BTreeSet<String>,
1653 ) -> Vec<ModelConfig> {
1654 let mut backend = config(base_url);
1655 backend.omit_body_fields = omit_body_fields;
1656 vec![ModelConfig::new("gpt", Backend::OpenAiChat(backend), None)]
1657 }
1658
1659 fn chat_map_with_effort(base_url: &str, effort: &str) -> Vec<ModelConfig> {
1660 let mut backend = config(base_url);
1661 backend.reasoning_effort = Some(effort.to_string());
1662 vec![ModelConfig::new("gpt", Backend::OpenAiChat(backend), None)]
1663 }
1664
1665 fn responses_map_with_effort(base_url: &str, effort: &str) -> Vec<ModelConfig> {
1666 let mut backend = config(base_url);
1667 backend.reasoning_effort = Some(effort.to_string());
1668 vec![ModelConfig::new(
1669 "gpt",
1670 Backend::OpenAiResponses(backend),
1671 None,
1672 )]
1673 }
1674
1675 fn anthropic_map(base_url: &str) -> Vec<ModelConfig> {
1676 vec![ModelConfig::new(
1677 "claude",
1678 Backend::Anthropic(config(base_url)),
1679 None,
1680 )]
1681 }
1682
1683 fn responses_map_with_drop_policy(base_url: &str) -> Vec<ModelConfig> {
1684 vec![
1685 ModelConfig::new("gpt", Backend::OpenAiResponses(config(base_url)), None)
1686 .with_responses_reasoning(crate::ResponsesReasoningPolicy::Drop),
1687 ]
1688 }
1689
1690 #[tokio::test]
1691 async fn model_properties_uses_configured_auth_and_backend_path()
1692 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
1693 for (base_path, props_path) in [
1694 ("/v1", "/props"),
1695 ("/v1/", "/props"),
1696 ("/v1/chat/completions", "/props"),
1697 ("/v1/responses/", "/props"),
1698 ("/prefix/v1", "/prefix/props"),
1699 ("/prefix/v1/chat/completions", "/prefix/props"),
1700 ("/prefix/v1/responses", "/prefix/props"),
1701 ("/prefix/", "/prefix/props"),
1702 ] {
1703 let server = MockServer::start().await;
1704 Mock::given(method("GET"))
1705 .and(path(props_path))
1706 .and(wiremock::matchers::header("authorization", "Bearer secret"))
1707 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
1708 "default_generation_settings": {"n_ctx": 65_536}
1709 })))
1710 .expect(1)
1711 .mount(&server)
1712 .await;
1713 let client =
1714 TranslatingLlmClient::new(&chat_map(&format!("{}{base_path}", server.uri())))?;
1715 let properties = client
1716 .get_model_properties(
1717 &ModelId::from("gpt"),
1718 WireFormat::OpenAiChat,
1719 Duration::from_secs(1),
1720 )
1721 .await?;
1722 assert_eq!(properties["default_generation_settings"]["n_ctx"], 65_536);
1723 }
1724 Ok(())
1725 }
1726
1727 fn chat_map_with_retries(base_url: &str, max_retries: u32) -> Vec<ModelConfig> {
1728 vec![ModelConfig::new(
1729 "gpt",
1730 Backend::OpenAiChat(config_with_retries(base_url, max_retries)),
1731 None,
1732 )]
1733 }
1734
1735 fn chat_success_response() -> ResponseTemplate {
1736 ResponseTemplate::new(200).set_body_json(json!({
1737 "id": "chatcmpl-1",
1738 "model": "gpt",
1739 "choices": [{
1740 "index": 0,
1741 "message": {"role": "assistant", "content": "recovered"},
1742 "finish_reason": "stop"
1743 }],
1744 "usage": {}
1745 }))
1746 }
1747
1748 fn truncated_response_server(
1749 content_type: &str,
1750 body: &str,
1751 ) -> std::io::Result<(String, JoinHandle<std::io::Result<()>>)> {
1752 response_sequence_server(vec![format!(
1753 "HTTP/1.1 200 OK\r\nContent-Type: {content_type}\r\n\
1754 Content-Length: {}\r\nConnection: close\r\n\r\n{body}",
1755 body.len() + 100
1756 )])
1757 }
1758
1759 fn response_sequence_server(
1760 responses: Vec<String>,
1761 ) -> std::io::Result<(String, JoinHandle<std::io::Result<()>>)> {
1762 let listener = std::net::TcpListener::bind("127.0.0.1:0")?;
1763 let address = listener.local_addr()?;
1764 let handle = std::thread::spawn(move || {
1765 for response in responses {
1766 let (mut stream, _) = listener.accept()?;
1767 let mut request = [0_u8; 1024];
1768 if stream.read(&mut request)? == 0 {
1769 return Err(std::io::Error::new(
1770 std::io::ErrorKind::UnexpectedEof,
1771 "client closed before sending a request",
1772 ));
1773 }
1774 stream.write_all(response.as_bytes())?;
1775 }
1776 Ok(())
1777 });
1778 Ok((format!("http://{address}/v1"), handle))
1779 }
1780
1781 fn raw_chat_success_response() -> String {
1782 let body = r#"{"id":"chatcmpl-1","model":"gpt","choices":[{"index":0,"message":{"role":"assistant","content":"recovered"},"finish_reason":"stop"}],"usage":{}}"#;
1783 format!(
1784 "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n\
1785 Content-Length: {}\r\nConnection: close\r\n\r\n{body}",
1786 body.len()
1787 )
1788 }
1789
1790 fn request_for(model: Option<&str>, stream: bool) -> Request {
1791 let mut llm_request = text_request(model.map(str::to_string), "hi");
1792 llm_request.stream = stream;
1793 Request {
1794 llm_request,
1795 raw_request: None,
1796 metadata: None,
1797 }
1798 }
1799
1800 fn request_with_headers(model: &str, headers: HeaderMap) -> Request {
1801 let mut request = request_for(Some(model), false);
1802 request.metadata = Some(Metadata {
1803 http_headers: Some(headers),
1804 ..Default::default()
1805 });
1806 request
1807 }
1808
1809 #[tokio::test]
1810 async fn transport_errors_drop_the_upstream_url() {
1811 let error = reqwest::Client::new()
1812 .post("http://127.0.0.1:1/v1?key=CANARY")
1813 .send()
1814 .await
1815 .expect_err("closed port");
1816
1817 assert!(!convert_reqwest_error(error).to_string().contains("CANARY"));
1818 }
1819
1820 #[test]
1821 fn nvidia_anthropic_tools_drop_unsupported_strict_field() {
1822 let backend = Backend::Anthropic(config("https://inference-api.nvidia.com"));
1823 assert!(nvidia_inference_api_backend(&backend));
1824 assert!(!nvidia_inference_api_backend(&Backend::Anthropic(config(
1825 "https://inference-api.nvidia.com.example.test"
1826 ))));
1827 let mut body = json!({"tools": [{"name": "a", "input_schema": {}, "strict": true}]});
1828 strip_anthropic_tool_strict(&mut body);
1829 assert!(body["tools"][0].get("strict").is_none());
1830 }
1831
1832 #[test]
1833 fn anthropic_prompt_caching_marks_final_message() {
1834 let caller_managed = json!({
1835 "cache_control": {"type": "ephemeral", "ttl": "1h"},
1836 "messages": [{"role": "user", "content": "hello"}]
1837 });
1838 let mut body = caller_managed.clone();
1839 enable_anthropic_prompt_caching(&mut body);
1840 assert_eq!(body, caller_managed);
1841
1842 let mut body = json!({
1843 "messages": [{"role": "user", "content": "hello"}]
1844 });
1845
1846 enable_anthropic_prompt_caching(&mut body);
1847
1848 assert_eq!(
1849 body["messages"][0]["content"][0]["cache_control"],
1850 json!({"type": "ephemeral"})
1851 );
1852 }
1853
1854 fn cache_control_blocks(value: &Value) -> usize {
1856 count_cache_control_blocks(value)
1857 }
1858
1859 fn body_at_the_cache_control_limit() -> Value {
1862 json!({
1863 "system": [
1864 {"type": "text", "text": "a", "cache_control": {"type": "ephemeral"}},
1865 {"type": "text", "text": "b", "cache_control": {"type": "ephemeral"}}
1866 ],
1867 "tools": [
1868 {"name": "t", "input_schema": {"type": "object"},
1869 "cache_control": {"type": "ephemeral"}}
1870 ],
1871 "messages": [
1872 {"role": "user", "content": [
1873 {"type": "text", "text": "old", "cache_control": {"type": "ephemeral"}}
1874 ]},
1875 {"role": "assistant", "content": [{"type": "text", "text": "ok"}]},
1876 {"role": "user", "content": [{"type": "text", "text": "new"}]}
1877 ]
1878 })
1879 }
1880
1881 #[test]
1882 fn anthropic_prompt_caching_abstains_at_the_cache_control_limit() {
1883 let mut body = body_at_the_cache_control_limit();
1884 assert_eq!(cache_control_blocks(&body), 4);
1885
1886 enable_anthropic_prompt_caching(&mut body);
1887
1888 assert_eq!(cache_control_blocks(&body), 4);
1891 assert!(
1892 body["messages"][2]["content"][0]
1893 .get("cache_control")
1894 .is_none()
1895 );
1896 }
1897
1898 #[test]
1899 fn anthropic_prompt_caching_still_marks_one_below_the_limit() {
1900 let mut body = body_at_the_cache_control_limit();
1901 body["system"].as_array_mut().unwrap().pop();
1903 assert_eq!(cache_control_blocks(&body), 3);
1904
1905 enable_anthropic_prompt_caching(&mut body);
1906
1907 assert_eq!(cache_control_blocks(&body), 4);
1908 assert_eq!(
1909 body["messages"][2]["content"][0]["cache_control"],
1910 json!({"type": "ephemeral"})
1911 );
1912 }
1913
1914 fn request_with_wire_format(model: &str, format: WireFormat) -> Request {
1917 let mut request = request_for(Some(model), false);
1918 request.metadata = Some(Metadata {
1919 session_id: None,
1920 agent_id: None,
1921 task_id: None,
1922 correlation_id: None,
1923 extra_metadata: None,
1924 http_headers: None,
1925 wire_format: Some(format),
1926 ..Default::default()
1927 });
1928 request
1929 }
1930
1931 #[tokio::test]
1932 async fn missing_model_errors()
1933 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
1934 let client = TranslatingLlmClient::new(&[])?;
1935 let Err(error) = client
1936 .call_rewrite_model(request_for(None, false), None)
1937 .await
1938 else {
1939 panic!("expected an error");
1940 };
1941 assert!(matches!(
1942 error,
1943 LlmClientError::InvalidRequest { message } if message == "no model given"
1944 ));
1945 Ok(())
1946 }
1947
1948 #[tokio::test]
1949 async fn unknown_model_errors()
1950 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
1951 let client = TranslatingLlmClient::new(&[])?;
1952 let Err(error) = client
1953 .call_rewrite_model(request_for(Some("gpt"), false), None)
1954 .await
1955 else {
1956 panic!("expected an error");
1957 };
1958 assert!(matches!(
1959 error,
1960 LlmClientError::Configuration { message }
1961 if message.contains("gpt")
1962 ));
1963 Ok(())
1964 }
1965
1966 #[tokio::test]
1967 async fn unknown_model_format_errors()
1968 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
1969 let client = TranslatingLlmClient::new(&chat_map("https://example.test/v1"))?;
1971 let Err(error) = client
1972 .call_rewrite_model(
1973 request_with_wire_format("gpt", WireFormat::AnthropicMessages),
1974 None,
1975 )
1976 .await
1977 else {
1978 panic!("expected an error");
1979 };
1980 assert!(matches!(
1981 error,
1982 LlmClientError::Configuration { message }
1983 if message.contains("gpt")
1984 && message.contains(&WireFormat::AnthropicMessages.to_string())
1985 ));
1986 Ok(())
1987 }
1988
1989 #[test]
1990 fn backend_for_resolves_configured_format()
1991 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
1992 let client = TranslatingLlmClient::new(&chat_map("https://example.test/v1"))?;
1993 assert!(
1995 client
1996 .backend_for(&ModelId::from("gpt"), WireFormat::OpenAiChat)
1997 .is_some()
1998 );
1999 assert!(
2000 client
2001 .backend_for(&ModelId::from("gpt"), WireFormat::AnthropicMessages)
2002 .is_none()
2003 );
2004 assert!(
2005 client
2006 .backend_for(&ModelId::from("missing"), WireFormat::OpenAiChat)
2007 .is_none()
2008 );
2009 Ok(())
2010 }
2011
2012 #[tokio::test]
2013 async fn model_name_arg_wins_over_request_model()
2014 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2015 let client = TranslatingLlmClient::new(&[])?;
2016 let Err(error) = client
2018 .call_rewrite_model(request_for(Some("a"), false), Some(&ModelId::from("b")))
2019 .await
2020 else {
2021 panic!("expected an error");
2022 };
2023 assert!(matches!(
2024 error,
2025 LlmClientError::Configuration { message }
2026 if message.contains("\"b\"")
2027 ));
2028 Ok(())
2029 }
2030
2031 #[tokio::test]
2032 async fn buffered_openai_chat_round_trips()
2033 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2034 let server = MockServer::start().await;
2035 Mock::given(method("POST"))
2036 .and(path("/v1/chat/completions"))
2037 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
2038 "id": "chatcmpl-1",
2039 "model": "gpt",
2040 "choices": [{
2041 "index": 0,
2042 "message": {"role": "assistant", "content": "Hi there"},
2043 "finish_reason": "stop"
2044 }],
2045 "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}
2046 })))
2047 .mount(&server)
2048 .await;
2049
2050 let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
2051
2052 let response = client
2053 .call_rewrite_model(request_for(Some("gpt"), false), None)
2054 .await?;
2055 let agg = response.llm_response.into_agg().await?;
2056 assert_eq!(completion_text(&agg), "Hi there");
2057
2058 Ok(())
2059 }
2060
2061 #[tokio::test]
2062 async fn invalid_json_is_a_response_translation_error()
2063 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2064 let server = MockServer::start().await;
2065 let calls = Arc::new(AtomicUsize::new(0));
2066 let observed_calls = Arc::clone(&calls);
2067 Mock::given(method("POST"))
2068 .respond_with(move |_: &wiremock::Request| {
2069 observed_calls.fetch_add(1, Ordering::SeqCst);
2070 ResponseTemplate::new(200).set_body_raw("not json", "application/json")
2071 })
2072 .mount(&server)
2073 .await;
2074
2075 let client =
2076 TranslatingLlmClient::new(&chat_map_with_retries(&format!("{}/v1", server.uri()), 2))?;
2077 let Err(error) = client
2078 .call_rewrite_model(request_for(Some("gpt"), false), None)
2079 .await
2080 else {
2081 panic!("expected invalid JSON to fail");
2082 };
2083
2084 assert!(matches!(
2085 error,
2086 LlmClientError::ResponseTranslation(message)
2087 if message.contains("invalid upstream JSON")
2088 ));
2089 assert_eq!(calls.load(Ordering::SeqCst), 1);
2090 Ok(())
2091 }
2092
2093 #[tokio::test]
2094 async fn response_body_io_failure_is_a_transport_error()
2095 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2096 let (base_url, server) = truncated_response_server("application/json", "{}")?;
2097 let client = TranslatingLlmClient::new(&chat_map(&base_url))?;
2098 let result = client
2099 .call_rewrite_model(request_for(Some("gpt"), false), None)
2100 .await;
2101 server
2102 .join()
2103 .map_err(|_| std::io::Error::other("response server thread panicked"))??;
2104
2105 let Err(error) = result else {
2106 panic!("expected the truncated response body to fail");
2107 };
2108 let LlmClientError::Transport { source } = error else {
2109 panic!("expected a transport error");
2110 };
2111 let Some(source) = source.downcast_ref::<reqwest::Error>() else {
2112 panic!("expected the reqwest transport source");
2113 };
2114 assert!(source.is_decode());
2115 assert!(
2116 !std::error::Error::source(&source)
2117 .is_some_and(|source| source.is::<serde_json::Error>())
2118 );
2119 Ok(())
2120 }
2121
2122 #[tokio::test]
2123 async fn response_body_transport_failures_are_retried()
2124 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2125 let truncated_responses = [
2126 "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n\
2127 Content-Length: 102\r\nConnection: close\r\n\r\n{}",
2128 "HTTP/1.1 401 Unauthorized\r\nContent-Type: text/plain\r\n\
2129 Content-Length: 100\r\nConnection: close\r\n\r\nbad",
2130 ];
2131
2132 for truncated in truncated_responses {
2133 let (base_url, server) =
2134 response_sequence_server(vec![truncated.to_string(), raw_chat_success_response()])?;
2135 let client = TranslatingLlmClient::new(&chat_map_with_retries(&base_url, 1))?;
2136 let response = client
2137 .call_rewrite_model(request_for(Some("gpt"), false), None)
2138 .await?;
2139 server
2140 .join()
2141 .map_err(|_| std::io::Error::other("response server thread panicked"))??;
2142
2143 assert_eq!(
2144 completion_text(&response.llm_response.into_agg().await?),
2145 "recovered"
2146 );
2147 }
2148 Ok(())
2149 }
2150
2151 #[tokio::test]
2152 async fn pre_first_event_transport_failure_uses_retry_budget()
2153 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2154 let truncated = "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\n\
2155 Content-Length: 100\r\nConnection: close\r\n\r\n";
2156 let body = "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"recovered\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n";
2157 let success = format!(
2158 "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\n\
2159 Content-Length: {}\r\nConnection: close\r\n\r\n{body}",
2160 body.len()
2161 );
2162
2163 for second in [success, truncated.to_string()] {
2164 let is_exhausted = second == truncated;
2165 let (base_url, server) = response_sequence_server(vec![truncated.to_string(), second])?;
2166 let client = TranslatingLlmClient::new(&chat_map_with_retries(&base_url, 1))?;
2167 let result = client
2168 .call_rewrite_model(request_for(Some("gpt"), true), None)
2169 .await;
2170 if is_exhausted {
2171 assert!(matches!(result, Err(LlmClientError::Transport { .. })));
2172 } else {
2173 assert_eq!(
2174 completion_text(&result?.llm_response.into_agg().await?),
2175 "recovered"
2176 );
2177 }
2178 server
2179 .join()
2180 .map_err(|_| std::io::Error::other("response server thread panicked"))??;
2181 }
2182 Ok(())
2183 }
2184
2185 #[tokio::test]
2186 async fn streaming_body_io_failure_preserves_transport_error()
2187 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2188 let body = "data: {\"choices\":[{\"delta\":{\"content\":\"partial\"}}]}\n\n";
2189 let (base_url, server) = truncated_response_server("text/event-stream", body)?;
2190 let client = TranslatingLlmClient::new(&chat_map_with_retries(&base_url, 2))?;
2191 let response = client
2192 .call_rewrite_model(request_for(Some("gpt"), true), None)
2193 .await?;
2194 let result = response.llm_response.into_agg().await;
2195 server
2196 .join()
2197 .map_err(|_| std::io::Error::other("response server thread panicked"))??;
2198
2199 let Err(error) = result else {
2200 panic!("expected the truncated stream body to fail");
2201 };
2202
2203 assert!(matches!(error, LlmClientError::Transport { .. }));
2204 Ok(())
2205 }
2206
2207 #[tokio::test]
2208 async fn rewrites_model_to_resolved_upstream_id()
2209 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2210 let server = MockServer::start().await;
2212 Mock::given(method("POST"))
2213 .and(path("/v1/chat/completions"))
2214 .and(wiremock::matchers::body_partial_json(json!({"model": "gpt"})))
2215 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
2216 "id": "1", "model": "gpt",
2217 "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
2218 "usage": {}
2219 })))
2220 .mount(&server)
2221 .await;
2222
2223 let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
2224 client
2226 .call_rewrite_model(
2227 request_for(Some("switchyard"), false),
2228 Some(&ModelId::from("gpt")),
2229 )
2230 .await?;
2231 Ok(())
2233 }
2234
2235 #[tokio::test]
2238 async fn reasoning_effort_override_replaces_the_callers_effort()
2239 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2240 let server = MockServer::start().await;
2241 Mock::given(method("POST"))
2242 .and(path("/v1/chat/completions"))
2243 .and(wiremock::matchers::body_partial_json(json!({
2244 "model": "gpt",
2245 "reasoning_effort": "max"
2246 })))
2247 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
2248 "id": "1",
2249 "model": "gpt",
2250 "choices": [{
2251 "index": 0,
2252 "message": {"role": "assistant", "content": "ok"},
2253 "finish_reason": "stop"
2254 }],
2255 "usage": {}
2256 })))
2257 .mount(&server)
2258 .await;
2259 let client = TranslatingLlmClient::new(&chat_map_with_effort(
2260 &format!("{}/v1", server.uri()),
2261 "max",
2262 ))?;
2263 client
2264 .call_rewrite_model_raw(
2265 json!({
2266 "model": "client-facing",
2267 "messages": [{"role": "user", "content": "hi"}],
2268 "reasoning_effort": "high"
2269 }),
2270 None,
2271 Some(&ModelId::from("gpt")),
2272 WireFormat::OpenAiChat,
2273 )
2274 .await?;
2275
2276 let server = MockServer::start().await;
2277 Mock::given(method("POST"))
2278 .and(path("/v1/responses"))
2279 .and(wiremock::matchers::body_partial_json(json!({
2280 "model": "gpt",
2281 "reasoning": {"effort": "max", "summary": "auto"}
2282 })))
2283 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
2284 "id": "resp_1",
2285 "object": "response",
2286 "model": "gpt",
2287 "status": "completed",
2288 "output": [{
2289 "type": "message",
2290 "id": "msg_1",
2291 "role": "assistant",
2292 "status": "completed",
2293 "content": [{"type": "output_text", "text": "ok"}]
2294 }],
2295 "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}
2296 })))
2297 .mount(&server)
2298 .await;
2299 let client = TranslatingLlmClient::new(&responses_map_with_effort(
2300 &format!("{}/v1", server.uri()),
2301 "max",
2302 ))?;
2303 client
2304 .call_rewrite_model_raw(
2305 json!({
2306 "model": "client-facing",
2307 "input": [{"role": "user", "content": "hi"}],
2308 "reasoning": {"effort": "high", "summary": "auto"}
2309 }),
2310 None,
2311 Some(&ModelId::from("gpt")),
2312 WireFormat::OpenAiResponses,
2313 )
2314 .await?;
2315 Ok(())
2316 }
2317
2318 #[tokio::test]
2319 async fn template_rejection_retries_once_without_changing_other_fields()
2320 -> std::result::Result<(), Box<dyn Error + Sync + Send>> {
2321 for (message, reject_retry, omit) in [
2322 (
2323 "Unrecognized request argument supplied: chat_template_kwargs",
2324 false,
2325 false,
2326 ),
2327 (
2328 "Unrecognized request argument supplied: chat_template_kwargs",
2329 true,
2330 false,
2331 ),
2332 (
2333 "Invalid temperature for chat_template_kwargs request",
2334 false,
2335 false,
2336 ),
2337 (
2338 "Unrecognized request argument supplied: chat_template_kwargs",
2339 false,
2340 true,
2341 ),
2342 ] {
2343 let server = MockServer::start().await;
2344 Mock::given(method("POST"))
2345 .respond_with(move |request: &wiremock::Request| {
2346 let body: Value = serde_json::from_slice(&request.body).expect("request JSON");
2347 if body.get("chat_template_kwargs").is_some() || reject_retry {
2348 ResponseTemplate::new(400).set_body_json(json!({"error": {"message": message}}))
2349 } else {
2350 ResponseTemplate::new(200).set_body_json(json!({
2351 "id": "1", "model": "gpt", "choices": [{"index": 0,
2352 "message": {"role": "assistant", "content": "yes"}, "finish_reason": "stop"}],
2353 "usage": {}
2354 }))
2355 }
2356 }).mount(&server).await;
2357 let fields = if omit {
2358 BTreeSet::from(["chat_template_kwargs".into()])
2359 } else {
2360 BTreeSet::new()
2361 };
2362 let client = TranslatingLlmClient::new(&chat_map_with_omit_body_fields(
2363 &format!("{}/v1", server.uri()),
2364 fields,
2365 ))?;
2366 let raw = json!({"model": "gpt", "messages": [{"role": "user", "content": "judge"}],
2367 "reasoning_effort": "none", "logprobs": true, "top_logprobs": 8,
2368 "chat_template_kwargs": {"enable_thinking": false}});
2369 let result = client
2370 .call_rewrite_model_raw(
2371 raw.clone(),
2372 None,
2373 Some(&ModelId::from("gpt")),
2374 WireFormat::OpenAiChat,
2375 )
2376 .await;
2377 let rejects_field = message.starts_with("Unrecognized");
2378 assert_eq!(result.is_ok(), omit || (rejects_field && !reject_retry));
2379 let calls = server.received_requests().await.ok_or("missing requests")?;
2380 assert_eq!(calls.len(), if !omit && rejects_field { 2 } else { 1 });
2381 for (index, call) in calls.iter().enumerate() {
2382 let mut expected = raw.clone();
2383 if omit || index > 0 {
2384 expected
2385 .as_object_mut()
2386 .unwrap()
2387 .remove("chat_template_kwargs");
2388 }
2389 assert_eq!(serde_json::from_slice::<Value>(&call.body)?, expected);
2390 }
2391 }
2392 Ok(())
2393 }
2394
2395 #[tokio::test]
2396 async fn template_compatibility_preserves_the_transient_retry_budget()
2397 -> std::result::Result<(), Box<dyn Error + Sync + Send>> {
2398 let server = MockServer::start().await;
2399 let count = Arc::new(AtomicUsize::new(0));
2400 let observed = Arc::clone(&count);
2401 Mock::given(method("POST"))
2402 .respond_with(move |request: &wiremock::Request| {
2403 let body: Value = serde_json::from_slice(&request.body).expect("request JSON");
2404 match observed.fetch_add(1, Ordering::SeqCst) {
2405 0 => {
2406 assert!(body.get("chat_template_kwargs").is_some());
2407 ResponseTemplate::new(422).set_body_json(json!({"error": {
2408 "code": "unsupported_parameter", "param": "chat_template_kwargs"
2409 }}))
2410 }
2411 1 => {
2412 assert!(body.get("chat_template_kwargs").is_none());
2413 ResponseTemplate::new(503).insert_header("retry-after", "0")
2414 }
2415 _ => ResponseTemplate::new(200).set_body_json(json!({
2416 "id": "1", "model": "gpt", "choices": [{"index": 0,
2417 "message": {"role": "assistant", "content": "yes"}, "finish_reason": "stop"}],
2418 "usage": {}
2419 })),
2420 }
2421 }).mount(&server).await;
2422 let client =
2423 TranslatingLlmClient::new(&chat_map_with_retries(&format!("{}/v1", server.uri()), 1))?;
2424 client
2425 .call_rewrite_model_raw(
2426 json!({"model": "gpt", "messages": [{"role": "user", "content": "judge"}],
2427 "chat_template_kwargs": {"enable_thinking": false}}),
2428 None,
2429 Some(&ModelId::from("gpt")),
2430 WireFormat::OpenAiChat,
2431 )
2432 .await?;
2433 assert_eq!(count.load(Ordering::SeqCst), 3);
2434 Ok(())
2435 }
2436
2437 #[tokio::test]
2438 async fn preserved_codex_string_input_is_normalized_at_the_client_boundary()
2439 -> std::result::Result<(), Box<dyn Error + Sync + Send>> {
2440 for root in ["/backend-api/codex", "/v1"] {
2441 for input in [
2442 json!("hello"),
2443 json!([{"role": "user", "content": "hello"}]),
2444 ] {
2445 let server = MockServer::start().await;
2446 Mock::given(method("POST"))
2447 .and(path(format!("{root}/responses")))
2448 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
2449 "id": "resp_1", "object": "response", "model": "gpt", "status": "completed", "output": [], "usage": {}
2450 }))).expect(1).mount(&server).await;
2451 let client =
2452 TranslatingLlmClient::new(&responses_map(&format!("{}{root}", server.uri())))?;
2453 let raw = json!({"model": "client-name", "input": input,
2454 "instructions": "Be brief.", "store": false, "metadata": {"task": "judge"},
2455 "custom_field": {"keep": true}});
2456 client
2457 .call_rewrite_model_raw(
2458 raw.clone(),
2459 None,
2460 Some(&ModelId::from("gpt")),
2461 WireFormat::OpenAiResponses,
2462 )
2463 .await?;
2464 let mut expected = raw;
2465 expected["model"] = json!("gpt");
2466 if root == "/backend-api/codex" && input.is_string() {
2467 expected["input"] =
2468 json!([{"type": "message", "role": "user", "content": input}]);
2469 }
2470 let calls = server.received_requests().await.ok_or("missing requests")?;
2471 assert_eq!(serde_json::from_slice::<Value>(&calls[0].body)?, expected);
2472 }
2473 }
2474 Ok(())
2475 }
2476
2477 #[tokio::test]
2478 async fn extra_body_adds_defaults_without_overriding_the_request()
2479 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2480 let server = MockServer::start().await;
2481 Mock::given(method("POST"))
2482 .and(path("/v1/chat/completions"))
2483 .and(wiremock::matchers::body_partial_json(json!({
2484 "model": "gpt",
2485 "max_tokens": 7,
2486 "service_tier": "priority"
2487 })))
2488 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
2489 "id": "1",
2490 "model": "gpt",
2491 "choices": [{
2492 "index": 0,
2493 "message": {"role": "assistant", "content": "ok"},
2494 "finish_reason": "stop"
2495 }],
2496 "usage": {}
2497 })))
2498 .mount(&server)
2499 .await;
2500
2501 let extra_body = BTreeMap::from([
2502 ("max_tokens".to_string(), json!(999)),
2503 ("service_tier".to_string(), json!("priority")),
2504 ]);
2505 let client = TranslatingLlmClient::new(&chat_map_with_extra_body(
2506 &format!("{}/v1", server.uri()),
2507 extra_body,
2508 ))?;
2509 let raw = json!({
2510 "model": "client-facing",
2511 "messages": [{"role": "user", "content": "hi"}],
2512 "max_tokens": 7
2513 });
2514
2515 client
2516 .call_rewrite_model_raw(
2517 raw,
2518 None,
2519 Some(&ModelId::from("gpt")),
2520 WireFormat::OpenAiChat,
2521 )
2522 .await?;
2523 Ok(())
2524 }
2525
2526 #[tokio::test]
2527 async fn omit_body_fields_strips_a_field_extra_body_cannot_override()
2528 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2529 let server = MockServer::start().await;
2530 Mock::given(method("POST"))
2531 .and(path("/v1/chat/completions"))
2532 .and(|request: &wiremock::Request| {
2533 let body: Value = serde_json::from_slice(&request.body).unwrap_or(Value::Null);
2534 body.get("max_tokens").is_none()
2535 })
2536 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
2537 "id": "1",
2538 "model": "gpt",
2539 "choices": [{
2540 "index": 0,
2541 "message": {"role": "assistant", "content": "ok"},
2542 "finish_reason": "stop"
2543 }],
2544 "usage": {}
2545 })))
2546 .mount(&server)
2547 .await;
2548
2549 let client = TranslatingLlmClient::new(&chat_map_with_omit_body_fields(
2550 &format!("{}/v1", server.uri()),
2551 BTreeSet::from(["max_tokens".to_string()]),
2552 ))?;
2553 let raw = json!({
2554 "model": "client-facing",
2555 "messages": [{"role": "user", "content": "hi"}],
2556 "max_tokens": 7
2557 });
2558
2559 client
2560 .call_rewrite_model_raw(
2561 raw,
2562 None,
2563 Some(&ModelId::from("gpt")),
2564 WireFormat::OpenAiChat,
2565 )
2566 .await?;
2567 Ok(())
2568 }
2569
2570 #[tokio::test]
2574 async fn anthropic_requests_drop_unsigned_thinking_blocks()
2575 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2576 let server = MockServer::start().await;
2577 Mock::given(method("POST"))
2578 .and(path("/v1/messages"))
2579 .and(|request: &wiremock::Request| {
2580 let body: Value = serde_json::from_slice(&request.body).unwrap_or(Value::Null);
2581 let messages = body.get("messages").and_then(Value::as_array).cloned();
2582 let Some(messages) = messages else {
2583 return false;
2584 };
2585 let blocks: Vec<&Value> = messages
2589 .iter()
2590 .filter_map(|message| message.get("content"))
2591 .filter_map(Value::as_array)
2592 .flatten()
2593 .collect();
2594 let thinking: Vec<&&Value> = blocks
2595 .iter()
2596 .filter(|block| block.get("type").and_then(Value::as_str) == Some("thinking"))
2597 .collect();
2598 thinking.len() == 1
2599 && thinking[0].get("signature").and_then(Value::as_str) == Some("sig-abc")
2600 && messages
2601 .iter()
2602 .all(|message| message.get("content") != Some(&json!([])))
2603 })
2604 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
2605 "id": "msg_1",
2606 "type": "message",
2607 "role": "assistant",
2608 "model": "claude",
2609 "content": [{"type": "text", "text": "ok"}],
2610 "stop_reason": "end_turn",
2611 "usage": {"input_tokens": 1, "output_tokens": 1}
2612 })))
2613 .mount(&server)
2614 .await;
2615
2616 let client = TranslatingLlmClient::new(&anthropic_map(&server.uri()))?;
2617 let raw = json!({
2618 "model": "client-facing",
2619 "max_tokens": 7,
2620 "messages": [
2621 {"role": "user", "content": "fix the build"},
2622 {"role": "assistant", "content": [
2623 {"type": "thinking", "thinking": "weak tier reasoning", "signature": ""}
2624 ]},
2625 {"role": "assistant", "content": [
2626 {"type": "thinking", "thinking": "signed reasoning", "signature": "sig-abc"},
2627 {"type": "text", "text": "here goes"}
2628 ]},
2629 {"role": "user", "content": "continue"}
2630 ]
2631 });
2632
2633 client
2634 .call_rewrite_model_raw(
2635 raw,
2636 None,
2637 Some(&ModelId::from("claude")),
2638 WireFormat::AnthropicMessages,
2639 )
2640 .await?;
2641 Ok(())
2642 }
2643
2644 #[tokio::test]
2645 async fn responses_reasoning_policy_preserves_replay_history()
2646 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2647 let history = json!([
2648 {"type": "message", "role": "user", "content": "continue"},
2649 {"type": "reasoning", "encrypted_content": "opaque", "content": [{"type": "reasoning_text", "text": "private"}]},
2650 {"type": "reasoning", "id": "rs_stored", "status": "completed", "summary": []},
2651 {"type": "reasoning", "encrypted_content": null},
2652 {"type": "reasoning", "encrypted_content": ""},
2653 {"type": "reasoning"},
2654 {"type": "function_call", "call_id": "call_1", "name": "shell", "arguments": "{}"},
2655 {"type": "function_call_output", "call_id": "call_1", "output": "ok"}
2656 ]);
2657 for store in [Some(true), Some(false), None] {
2658 for is_drop in [false, true] {
2659 for has_input_override in [false, true] {
2660 let server = MockServer::start().await;
2661 Mock::given(method("POST"))
2662 .and(path("/v1/responses"))
2663 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
2664 "id": "resp_1", "object": "response", "model": "gpt",
2665 "status": "completed", "output": [],
2666 "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}
2667 })))
2668 .expect(1)
2669 .mount(&server)
2670 .await;
2671 let base_url = format!("{}/v1", server.uri());
2672 let mut models = if is_drop {
2673 responses_map_with_drop_policy(&base_url)
2674 } else {
2675 responses_map(&base_url)
2676 };
2677 if has_input_override {
2678 let Backend::OpenAiResponses(backend) = &mut models[0].default_backend
2679 else {
2680 return Err("expected Responses backend".into());
2681 };
2682 backend.omit_body_fields.insert("input".to_string());
2683 backend
2684 .extra_body
2685 .insert("input".to_string(), history.clone());
2686 }
2687 let client = TranslatingLlmClient::new(&models)?;
2688 let mut request = json!({"model": "gpt", "input": history});
2689 if let Some(store) = store {
2690 request["store"] = json!(store);
2691 }
2692 client
2693 .call_rewrite_model_raw(
2694 request,
2695 None,
2696 Some(&ModelId::from("gpt")),
2697 WireFormat::OpenAiResponses,
2698 )
2699 .await?;
2700 let requests = server.received_requests().await.ok_or("missing requests")?;
2701 let body: Value = serde_json::from_slice(&requests[0].body)?;
2702 let mut expected = vec![history[0].clone()];
2703 if !is_drop {
2704 let mut encrypted = history[1].clone();
2705 encrypted["content"] = json!([]);
2706 expected.push(encrypted);
2707 if store != Some(false) {
2708 let mut stored = history[2].clone();
2709 stored["content"] = json!([]);
2710 expected.push(stored);
2711 }
2712 }
2713 expected.extend([history[6].clone(), history[7].clone()]);
2714 assert_eq!(
2715 body["input"],
2716 json!(expected),
2717 "store={store:?}, drop={is_drop}, replace={has_input_override}"
2718 );
2719 }
2720 }
2721 }
2722
2723 Ok(())
2724 }
2725
2726 #[tokio::test]
2727 async fn nvidia_anthropic_requests_strip_strict_from_extra_body()
2728 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2729 let server = MockServer::start().await;
2730 Mock::given(method("POST"))
2731 .and(path("/v1/messages"))
2732 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
2733 "id": "msg_1",
2734 "type": "message",
2735 "role": "assistant",
2736 "model": "claude",
2737 "content": [{"type": "text", "text": "ok"}],
2738 "stop_reason": "end_turn",
2739 "usage": {"input_tokens": 1, "output_tokens": 1}
2740 })))
2741 .expect(1)
2742 .mount(&server)
2743 .await;
2744
2745 let mut backend = config(&format!(
2746 "http://inference-api.nvidia.com:{}",
2747 server.address().port()
2748 ));
2749 backend.extra_body.insert(
2750 "tools".into(),
2751 json!([{
2752 "name": "lookup",
2753 "input_schema": {"type": "object", "properties": {}},
2754 "strict": true
2755 }]),
2756 );
2757 let mut client = TranslatingLlmClient::new(&[ModelConfig::new(
2758 "claude",
2759 Backend::Anthropic(backend),
2760 None,
2761 )])?;
2762 client.client = reqwest::Client::builder()
2763 .no_proxy()
2764 .resolve("inference-api.nvidia.com", *server.address())
2765 .build()?;
2766 client
2767 .call_rewrite_model_raw(
2768 json!({
2769 "model": "claude",
2770 "max_tokens": 7,
2771 "messages": [{"role": "user", "content": "hi"}]
2772 }),
2773 None,
2774 Some(&ModelId::from("claude")),
2775 WireFormat::AnthropicMessages,
2776 )
2777 .await?;
2778
2779 let requests = server.received_requests().await.unwrap();
2780 let body: Value = serde_json::from_slice(&requests[0].body)?;
2781 assert_eq!(body["tools"][0]["name"], "lookup");
2782 assert!(body["tools"][0].get("strict").is_none());
2783 Ok(())
2784 }
2785
2786 #[tokio::test]
2790 async fn anthropic_requests_drop_openai_only_fields()
2791 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2792 let server = MockServer::start().await;
2793 Mock::given(method("POST"))
2794 .and(path("/v1/messages"))
2795 .and(|request: &wiremock::Request| {
2796 let body: Value = serde_json::from_slice(&request.body).unwrap_or(Value::Null);
2797 body.get("context_management").is_none() && body.get("reasoning_effort").is_none()
2798 })
2799 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
2800 "id": "msg_1",
2801 "type": "message",
2802 "role": "assistant",
2803 "model": "claude",
2804 "content": [{"type": "text", "text": "ok"}],
2805 "stop_reason": "end_turn",
2806 "usage": {"input_tokens": 1, "output_tokens": 1}
2807 })))
2808 .mount(&server)
2809 .await;
2810
2811 let client = TranslatingLlmClient::new(&anthropic_map(&server.uri()))?;
2812 let raw = json!({
2813 "model": "client-facing",
2814 "messages": [{"role": "user", "content": "hi"}],
2815 "max_tokens": 7,
2816 "reasoning_effort": "high",
2817 "context_management": {
2818 "edits": [{"type": "clear_thinking_20251015"}]
2819 }
2820 });
2821
2822 client
2823 .call_rewrite_model_raw(
2824 raw,
2825 None,
2826 Some(&ModelId::from("claude")),
2827 WireFormat::AnthropicMessages,
2828 )
2829 .await?;
2830 Ok(())
2831 }
2832
2833 #[tokio::test]
2834 async fn streaming_openai_chat_aggregates()
2835 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2836 let server = MockServer::start().await;
2837 let sse = "data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}\n\n\
2838 data: {\"choices\":[{\"delta\":{\"content\":\" world\"}}]}\n\n\
2839 data: {\"choices\":[],\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":2,\"total_tokens\":3}}\n\n\
2840 data: [DONE]\n\n";
2841 Mock::given(method("POST"))
2842 .and(path("/v1/chat/completions"))
2843 .and(wiremock::matchers::body_partial_json(json!({
2844 "stream": true,
2845 "stream_options": {"include_usage": true}
2846 })))
2847 .respond_with(ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream"))
2848 .mount(&server)
2849 .await;
2850
2851 let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
2852
2853 let response = client
2854 .call_rewrite_model(request_for(Some("gpt"), true), None)
2855 .await?;
2856 assert!(matches!(response.llm_response, LlmResponse::Stream(_)));
2857 let agg = response.llm_response.into_agg().await?;
2858 assert_eq!(completion_text(&agg), "Hello world");
2859 assert_eq!(agg.usage.input_tokens, Some(1));
2860 assert_eq!(agg.usage.output_tokens, Some(2));
2861 assert_eq!(agg.usage.total_tokens, Some(3));
2862 Ok(())
2863 }
2864
2865 #[tokio::test]
2866 async fn streaming_openai_chat_preserves_usage_opt_out()
2867 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2868 let server = MockServer::start().await;
2869 Mock::given(method("POST"))
2870 .and(path("/v1/chat/completions"))
2871 .and(wiremock::matchers::body_partial_json(json!({
2872 "stream": true,
2873 "stream_options": {"include_usage": false}
2874 })))
2875 .respond_with(
2876 ResponseTemplate::new(200).set_body_raw("data: [DONE]\n\n", "text/event-stream"),
2877 )
2878 .mount(&server)
2879 .await;
2880
2881 let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
2882 let raw = json!({
2883 "model": "client-facing",
2884 "messages": [{"role": "user", "content": "hi"}],
2885 "stream": true,
2886 "stream_options": {"include_usage": false}
2887 });
2888
2889 let response = client
2890 .call_rewrite_model_raw(
2891 raw,
2892 None,
2893 Some(&ModelId::from("gpt")),
2894 WireFormat::OpenAiChat,
2895 )
2896 .await?;
2897 assert!(matches!(response, RawResponse::Stream(_)));
2898 Ok(())
2899 }
2900
2901 #[tokio::test]
2902 async fn upstream_500_is_upstream_http()
2903 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2904 let server = MockServer::start().await;
2905 Mock::given(method("POST"))
2906 .respond_with(ResponseTemplate::new(500).set_body_string("boom"))
2907 .mount(&server)
2908 .await;
2909
2910 let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
2911
2912 let Err(error) = client
2913 .call_rewrite_model(request_for(Some("gpt"), false), None)
2914 .await
2915 else {
2916 panic!("expected an error");
2917 };
2918 assert!(matches!(
2919 error,
2920 LlmClientError::UpstreamHttp {
2921 status: StatusCode::INTERNAL_SERVER_ERROR,
2922 ..
2923 }
2924 ));
2925 Ok(())
2926 }
2927
2928 #[tokio::test]
2929 async fn retryable_http_failure_recovers_within_budget()
2930 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2931 let server = MockServer::start().await;
2932 let calls = Arc::new(AtomicUsize::new(0));
2933 let observed_calls = Arc::clone(&calls);
2934 Mock::given(method("POST"))
2935 .respond_with(move |_: &wiremock::Request| {
2936 if observed_calls.fetch_add(1, Ordering::SeqCst) == 0 {
2937 ResponseTemplate::new(503)
2938 .insert_header("retry-after", "0")
2939 .set_body_string("temporarily unavailable")
2940 } else {
2941 chat_success_response()
2942 }
2943 })
2944 .mount(&server)
2945 .await;
2946
2947 let client =
2948 TranslatingLlmClient::new(&chat_map_with_retries(&format!("{}/v1", server.uri()), 1))?;
2949 let response = client
2950 .call_rewrite_model(request_for(Some("gpt"), false), None)
2951 .await?;
2952 let agg = response.llm_response.into_agg().await?;
2953
2954 assert_eq!(completion_text(&agg), "recovered");
2955 assert_eq!(calls.load(Ordering::SeqCst), 2);
2956 Ok(())
2957 }
2958
2959 #[tokio::test]
2960 async fn deterministic_http_failure_is_not_retried()
2961 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2962 let server = MockServer::start().await;
2963 let calls = Arc::new(AtomicUsize::new(0));
2964 let observed_calls = Arc::clone(&calls);
2965 Mock::given(method("POST"))
2966 .respond_with(move |_: &wiremock::Request| {
2967 observed_calls.fetch_add(1, Ordering::SeqCst);
2968 ResponseTemplate::new(401).set_body_string("invalid key")
2969 })
2970 .mount(&server)
2971 .await;
2972
2973 let client =
2974 TranslatingLlmClient::new(&chat_map_with_retries(&format!("{}/v1", server.uri()), 2))?;
2975 let Err(error) = client
2976 .call_rewrite_model(request_for(Some("gpt"), false), None)
2977 .await
2978 else {
2979 panic!("expected an upstream error");
2980 };
2981
2982 assert!(matches!(
2983 error,
2984 LlmClientError::UpstreamHttp {
2985 status: StatusCode::UNAUTHORIZED,
2986 body
2987 } if body == "invalid key"
2988 ));
2989 assert_eq!(calls.load(Ordering::SeqCst), 1);
2990 Ok(())
2991 }
2992
2993 #[tokio::test]
2994 async fn rate_limit_cooldown_depends_on_shared_credentials()
2995 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2996 for forward_auth in [true, false] {
2997 let server = MockServer::start().await;
2998 Mock::given(method("POST"))
2999 .respond_with(ResponseTemplate::new(429).insert_header("retry-after", "0"))
3000 .expect(if forward_auth { 6 } else { 3 })
3001 .mount(&server)
3002 .await;
3003 let backend = HttpBackendConfig {
3004 forward_auth,
3005 failure_cooldown: Duration::from_secs(60),
3006 ..config_with_retries(&format!("{}/v1", server.uri()), 2)
3007 };
3008 let client = TranslatingLlmClient::new(&[ModelConfig::new(
3009 "gpt",
3010 Backend::OpenAiChat(backend),
3011 None,
3012 )])?;
3013 for credential in ["Bearer first-user", "Bearer second-user"] {
3014 let mut headers = http::HeaderMap::new();
3015 headers.insert("authorization", http::HeaderValue::from_static(credential));
3016 let result = client
3017 .call_rewrite_model(request_with_headers("gpt", headers), None)
3018 .await;
3019 if !forward_auth && credential == "Bearer second-user" {
3020 assert!(matches!(
3021 result,
3022 Err(LlmClientError::TemporarilyUnavailable)
3023 ));
3024 } else {
3025 assert!(matches!(
3026 result,
3027 Err(LlmClientError::UpstreamHttp {
3028 status: StatusCode::TOO_MANY_REQUESTS,
3029 ..
3030 })
3031 ));
3032 }
3033 }
3034 }
3035 Ok(())
3036 }
3037
3038 #[tokio::test]
3039 async fn retry_exhaustion_returns_the_final_upstream_error()
3040 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3041 let server = MockServer::start().await;
3042 let calls = Arc::new(AtomicUsize::new(0));
3043 let observed_calls = Arc::clone(&calls);
3044 Mock::given(method("POST"))
3045 .respond_with(move |_: &wiremock::Request| {
3046 let attempt = observed_calls.fetch_add(1, Ordering::SeqCst) + 1;
3047 ResponseTemplate::new(500)
3048 .insert_header("retry-after", "0")
3049 .set_body_string(format!("attempt {attempt}"))
3050 })
3051 .mount(&server)
3052 .await;
3053
3054 let client =
3055 TranslatingLlmClient::new(&chat_map_with_retries(&format!("{}/v1", server.uri()), 2))?;
3056 let Err(error) = client
3057 .call_rewrite_model(request_for(Some("gpt"), false), None)
3058 .await
3059 else {
3060 panic!("expected retry exhaustion");
3061 };
3062
3063 assert!(matches!(
3064 error,
3065 LlmClientError::UpstreamHttp {
3066 status: StatusCode::INTERNAL_SERVER_ERROR,
3067 body
3068 } if body == "attempt 3"
3069 ));
3070 assert_eq!(calls.load(Ordering::SeqCst), 3);
3071 Ok(())
3072 }
3073
3074 #[tokio::test]
3075 async fn deadline_expires_before_send_or_stream_poll()
3076 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3077 let server = MockServer::start().await;
3078 Mock::given(method("POST"))
3079 .respond_with(ResponseTemplate::new(200).set_body_raw(
3080 "data: {\"id\":\"test\",\"model\":\"gpt\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"ok\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n",
3081 "text/event-stream",
3082 ))
3083 .mount(&server)
3084 .await;
3085 for timeout in [Duration::ZERO, Duration::from_millis(200)] {
3086 let mut backend = config(&server.uri());
3087 backend.timeout = Some(timeout);
3088 let client = TranslatingLlmClient::new(&[ModelConfig::new(
3089 "gpt",
3090 Backend::OpenAiChat(backend),
3091 None,
3092 )])?;
3093 let result = client
3094 .call_rewrite_model(request_for(Some("gpt"), true), None)
3095 .await;
3096 if timeout.is_zero() {
3097 assert!(matches!(result, Err(LlmClientError::Timeout { .. })));
3098 assert!(
3099 server
3100 .received_requests()
3101 .await
3102 .expect("recorded requests")
3103 .is_empty()
3104 );
3105 } else {
3106 let LlmResponse::Stream(mut chunks) = result?.llm_response else {
3107 return Err("expected a streaming response".into());
3108 };
3109 tokio::time::sleep(Duration::from_millis(300)).await;
3110 assert!(matches!(
3111 chunks.next().await,
3112 Some(Err(LlmClientError::Timeout { .. }))
3113 ));
3114 assert!(chunks.next().await.is_none());
3115 }
3116 }
3117 Ok(())
3118 }
3119
3120 #[tokio::test]
3121 async fn timeout_is_retried_before_a_response_is_returned()
3122 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3123 let server = MockServer::start().await;
3124 let calls = Arc::new(AtomicUsize::new(0));
3125 let observed_calls = Arc::clone(&calls);
3126 Mock::given(method("POST"))
3127 .respond_with(move |_: &wiremock::Request| {
3128 if observed_calls.fetch_add(1, Ordering::SeqCst) == 0 {
3129 ResponseTemplate::new(200).set_delay(Duration::from_millis(500))
3130 } else {
3131 chat_success_response()
3132 }
3133 })
3134 .mount(&server)
3135 .await;
3136
3137 let mut client =
3138 TranslatingLlmClient::new(&chat_map_with_retries(&format!("{}/v1", server.uri()), 1))?;
3139 client.client = reqwest::Client::builder()
3140 .timeout(Duration::from_millis(100))
3141 .build()?;
3142 let response = client
3143 .call_rewrite_model(request_for(Some("gpt"), false), None)
3144 .await?;
3145
3146 assert_eq!(
3147 completion_text(&response.llm_response.into_agg().await?),
3148 "recovered"
3149 );
3150 assert_eq!(calls.load(Ordering::SeqCst), 2);
3151 Ok(())
3152 }
3153
3154 #[test]
3155 fn retryable_error_classes_are_explicit() {
3156 let transport = AttemptFailure {
3157 error: LlmClientError::Transport {
3158 source: std::io::Error::other("disconnected").into(),
3159 },
3160 status: None,
3161 retry_after: None,
3162 };
3163 assert!(transport.is_retryable());
3164
3165 for status in [
3166 StatusCode::REQUEST_TIMEOUT,
3167 StatusCode::TOO_MANY_REQUESTS,
3168 StatusCode::INTERNAL_SERVER_ERROR,
3169 StatusCode::SERVICE_UNAVAILABLE,
3170 StatusCode::from_u16(599)
3171 .expect("599 is a syntactically valid HTTP code in the http crate"),
3172 ] {
3173 let failure = AttemptFailure {
3174 error: LlmClientError::UpstreamHttp {
3175 status,
3176 body: String::new(),
3177 },
3178 status: Some(status),
3179 retry_after: None,
3180 };
3181 assert!(failure.is_retryable(), "HTTP {status} should retry");
3182 }
3183 for status in [
3184 StatusCode::BAD_REQUEST,
3185 StatusCode::UNAUTHORIZED,
3186 StatusCode::CONFLICT,
3187 StatusCode::from_u16(600)
3188 .expect("600 is a syntactically valid HTTP code in the http crate"),
3189 ] {
3190 let failure = AttemptFailure {
3191 error: LlmClientError::UpstreamHttp {
3192 status,
3193 body: String::new(),
3194 },
3195 status: Some(status),
3196 retry_after: None,
3197 };
3198 assert!(!failure.is_retryable(), "HTTP {status} should fail fast");
3199 }
3200
3201 let configuration = AttemptFailure {
3202 error: LlmClientError::Configuration {
3203 message: "invalid header".to_string(),
3204 },
3205 status: None,
3206 retry_after: None,
3207 };
3208 assert!(!configuration.is_retryable());
3209
3210 let context_window = AttemptFailure {
3211 error: LlmClientError::ContextWindowExceeded {
3212 model: ModelId::from("gpt"),
3213 message: "too long".to_string(),
3214 },
3215 status: Some(StatusCode::BAD_REQUEST),
3216 retry_after: None,
3217 };
3218 assert!(!context_window.is_retryable());
3219 }
3220
3221 #[test]
3222 fn retry_after_supports_seconds_and_http_dates() {
3223 let mut headers = HeaderMap::new();
3224 headers.insert(RETRY_AFTER, reqwest::header::HeaderValue::from_static("3"));
3225 assert_eq!(retry_after_delay(&headers), Some(Duration::from_secs(3)));
3226
3227 let retry_at = SystemTime::now() + Duration::from_secs(2);
3228 let value = httpdate::fmt_http_date(retry_at);
3229 let Ok(value) = reqwest::header::HeaderValue::from_str(&value) else {
3230 panic!("formatted HTTP date should be a valid header");
3231 };
3232 headers.insert(RETRY_AFTER, value);
3233 let Some(delay) = retry_after_delay(&headers) else {
3234 panic!("HTTP date should produce a retry delay");
3235 };
3236 assert!(delay <= Duration::from_secs(2));
3237 }
3238
3239 #[tokio::test]
3240 async fn routed_llm_client_exposes_timeout_variant()
3241 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3242 let server = MockServer::start().await;
3243 Mock::given(method("POST"))
3244 .respond_with(
3245 ResponseTemplate::new(200).set_delay(std::time::Duration::from_millis(100)),
3246 )
3247 .mount(&server)
3248 .await;
3249
3250 let mut client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
3251 client.client = reqwest::Client::builder()
3252 .timeout(std::time::Duration::from_millis(10))
3253 .build()?;
3254
3255 let Err(error) = client.call(request_for(Some("gpt"), false)).await else {
3256 panic!("expected a timeout");
3257 };
3258 let LlmClientError::Timeout { source } = error else {
3259 panic!("expected the protocol timeout variant");
3260 };
3261 let Some(source) = source.downcast_ref::<reqwest::Error>() else {
3262 panic!("expected the reqwest timeout source");
3263 };
3264 assert!(source.is_timeout());
3265 Ok(())
3266 }
3267
3268 #[cfg(target_os = "linux")]
3270 #[tokio::test]
3271 async fn connect_timeout_is_a_transport_error()
3272 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3273 let socket = tokio::net::TcpSocket::new_v4()?;
3276 socket.bind("127.0.0.1:0".parse()?)?;
3277 let listener = socket.listen(0)?;
3278 let addr = listener.local_addr()?;
3279 let _queued: Vec<_> = (0..2)
3280 .map(|_| {
3281 std::net::TcpStream::connect_timeout(&addr, std::time::Duration::from_millis(100))
3282 })
3283 .collect();
3284
3285 let mut client = TranslatingLlmClient::new(&chat_map(&format!("http://{addr}/v1")))?;
3286 client.client = reqwest::Client::builder()
3289 .connect_timeout(std::time::Duration::from_millis(100))
3290 .build()?;
3291
3292 let Err(error) = client.call(request_for(Some("gpt"), false)).await else {
3293 panic!("expected the connect to fail");
3294 };
3295 let LlmClientError::Transport { source } = error else {
3296 panic!("expected Transport, got {error:?}");
3297 };
3298 let Some(source) = source.downcast_ref::<reqwest::Error>() else {
3299 panic!("expected the reqwest connect-timeout source");
3300 };
3301 assert!(source.is_connect());
3302 assert!(source.is_timeout());
3303 Ok(())
3304 }
3305
3306 #[tokio::test]
3307 async fn llama_cpp_context_overflow_400_is_mapped()
3308 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3309 let server = MockServer::start().await;
3310 Mock::given(method("POST"))
3311 .respond_with(ResponseTemplate::new(400).set_body_json(json!({
3312 "error": {
3313 "message": "request (6016 tokens) exceeds the available context size (4096 tokens), try increasing it"
3314 }
3315 })))
3316 .mount(&server)
3317 .await;
3318
3319 let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
3320
3321 let Err(error) = client
3322 .call_rewrite_model(request_for(Some("gpt"), false), None)
3323 .await
3324 else {
3325 panic!("expected an error");
3326 };
3327 assert!(matches!(
3328 error,
3329 LlmClientError::ContextWindowExceeded { model, .. } if model == "gpt"
3330 ));
3331 Ok(())
3332 }
3333
3334 #[tokio::test]
3335 async fn content_policy_400_preserves_upstream_error()
3336 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3337 let server = MockServer::start().await;
3338 let body = json!({
3339 "error": {
3340 "code": "content_policy_violation",
3341 "message": "request blocked by content policy",
3342 "type": "invalid_request_error"
3343 },
3344 "metadata": {"documentation_section": "context window"}
3345 });
3346 Mock::given(method("POST"))
3347 .respond_with(ResponseTemplate::new(400).set_body_json(&body))
3348 .expect(1)
3349 .mount(&server)
3350 .await;
3351 let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
3352 let Err(error) = client
3353 .call_rewrite_model(request_for(Some("gpt"), false), None)
3354 .await
3355 else {
3356 panic!("expected a policy denial");
3357 };
3358 let LlmClientError::UpstreamHttp {
3359 status,
3360 body: actual,
3361 } = error
3362 else {
3363 panic!("expected the upstream HTTP error, got {error:?}");
3364 };
3365 assert_eq!(status, StatusCode::BAD_REQUEST);
3366 assert_eq!(serde_json::from_str::<serde_json::Value>(&actual)?, body);
3367 Ok(())
3368 }
3369
3370 #[tokio::test]
3371 async fn forwards_only_allowlisted_metadata_headers()
3372 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3373 let server = MockServer::start().await;
3374 Mock::given(method("POST"))
3375 .and(wiremock::matchers::header("x-request-id", "abc"))
3376 .and(wiremock::matchers::header("authorization", "Bearer secret"))
3378 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
3379 "id": "1", "model": "gpt",
3380 "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
3381 "usage": {}
3382 })))
3383 .mount(&server)
3384 .await;
3385
3386 let mut headers = http::HeaderMap::new();
3387 headers.insert("x-request-id", http::HeaderValue::from_static("abc"));
3388 headers.insert(
3389 "authorization",
3390 http::HeaderValue::from_static("Bearer client-key"),
3391 );
3392 headers.insert(
3393 "x-goog-api-key",
3394 http::HeaderValue::from_static("client-google-key"),
3395 );
3396 headers.insert(
3397 "x-custom-internal-secret",
3398 http::HeaderValue::from_static("client-custom-key"),
3399 );
3400 let request = request_with_headers("gpt", headers);
3401
3402 let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
3403
3404 client.call_rewrite_model(request, None).await?;
3407 let received = server
3408 .received_requests()
3409 .await
3410 .ok_or("request recording should be enabled")?;
3411 let received = received.first().ok_or("expected one upstream request")?;
3412 assert!(!received.headers.contains_key("x-goog-api-key"));
3413 assert!(!received.headers.contains_key("x-custom-internal-secret"));
3414 Ok(())
3415 }
3416
3417 #[tokio::test]
3418 async fn forwards_application_headers_with_forward_auth()
3419 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3420 let server = MockServer::start().await;
3421 Mock::given(method("POST"))
3422 .and(wiremock::matchers::header(
3423 "authorization",
3424 "Bearer client-key",
3425 ))
3426 .and(wiremock::matchers::header(
3427 "x-goog-api-key",
3428 "client-google-key",
3429 ))
3430 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
3431 "id": "1", "model": "gpt",
3432 "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
3433 "usage": {}
3434 })))
3435 .mount(&server)
3436 .await;
3437
3438 let mut headers = http::HeaderMap::new();
3439 headers.insert(
3440 "authorization",
3441 http::HeaderValue::from_static("Bearer client-key"),
3442 );
3443 headers.insert(
3444 "x-goog-api-key",
3445 http::HeaderValue::from_static("client-google-key"),
3446 );
3447 headers.insert("host", http::HeaderValue::from_static("client.example"));
3448 headers.insert("content-type", http::HeaderValue::from_static("text/plain"));
3449 headers.insert("connection", http::HeaderValue::from_static("x-hop-by-hop"));
3450 headers.insert(
3451 "x-hop-by-hop",
3452 http::HeaderValue::from_static("client-only"),
3453 );
3454 let request = request_with_headers("gpt", headers);
3455 let backend_config = forwarding_config(&format!("{}/v1", server.uri()));
3456 let client = TranslatingLlmClient::new(&[ModelConfig::new(
3457 "gpt",
3458 Backend::OpenAiChat(backend_config),
3459 None,
3460 )])?;
3461
3462 client.call_rewrite_model(request, None).await?;
3463 let received = server
3464 .received_requests()
3465 .await
3466 .ok_or("request recording should be enabled")?;
3467 let received = received.first().ok_or("expected one upstream request")?;
3468 assert_ne!(received.headers["host"], "client.example");
3469 assert_eq!(received.headers["content-type"], "application/json");
3470 assert!(!received.headers.contains_key("connection"));
3471 assert!(!received.headers.contains_key("x-hop-by-hop"));
3472 Ok(())
3473 }
3474
3475 #[tokio::test]
3476 async fn forward_auth_preserves_anthropic_headers_and_redacts_errors()
3477 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3478 let server = MockServer::start().await;
3479 Mock::given(method("POST"))
3480 .and(path("/v1/messages"))
3481 .and(|request: &wiremock::Request| {
3482 request
3483 .headers
3484 .get("anthropic-beta")
3485 .and_then(|value| value.to_str().ok())
3486 == Some("oauth-2025-04-20")
3487 && request.headers.get_all("anthropic-version").iter().count() == 1
3488 && request
3489 .headers
3490 .get("anthropic-version")
3491 .and_then(|value| value.to_str().ok())
3492 == Some("2023-06-01")
3493 })
3494 .respond_with(ResponseTemplate::new(401).set_body_json(json!({
3495 "error": {"message": "rejected client-google-key"}
3496 })))
3497 .mount(&server)
3498 .await;
3499
3500 let mut headers = http::HeaderMap::new();
3501 headers.insert(
3502 "authorization",
3503 http::HeaderValue::from_static("Bearer client-key"),
3504 );
3505 headers.insert(
3506 "x-goog-api-key",
3507 http::HeaderValue::from_static("client-google-key"),
3508 );
3509 headers.insert(
3510 "anthropic-beta",
3511 http::HeaderValue::from_static("oauth-2025-04-20,prompt-caching-2024-07-31"),
3512 );
3513 headers.insert(
3514 "anthropic-version",
3515 http::HeaderValue::from_static("caller-version"),
3516 );
3517 let backend_config = forwarding_config(&server.uri());
3518 let client = TranslatingLlmClient::new(&[ModelConfig::new(
3519 "claude",
3520 Backend::Anthropic(backend_config),
3521 None,
3522 )])?;
3523 let raw = json!({
3524 "model": "claude",
3525 "max_tokens": 8,
3526 "messages": [{"role": "user", "content": "hello"}]
3527 });
3528
3529 let Err(LlmClientError::UpstreamHttp { body, .. }) = client
3530 .call_rewrite_model_raw(
3531 raw,
3532 Some(headers),
3533 Some(&ModelId::from("claude")),
3534 WireFormat::AnthropicMessages,
3535 )
3536 .await
3537 else {
3538 panic!("expected an upstream HTTP error");
3539 };
3540 assert_eq!(body, r#"{"error":{"message":"rejected [REDACTED]"}}"#);
3541 Ok(())
3542 }
3543
3544 #[tokio::test]
3547 async fn routed_llm_client_serves_the_request_model()
3548 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3549 let server = MockServer::start().await;
3550 Mock::given(method("POST"))
3551 .and(path("/v1/chat/completions"))
3552 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
3553 "id": "chatcmpl-1",
3554 "model": "gpt",
3555 "choices": [{
3556 "index": 0,
3557 "message": {"role": "assistant", "content": "routed hi"},
3558 "finish_reason": "stop"
3559 }],
3560 "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}
3561 })))
3562 .mount(&server)
3563 .await;
3564
3565 let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
3566 let response = client.call(request_for(Some("gpt"), false)).await?;
3567 let agg = response.llm_response.into_agg().await?;
3568 assert_eq!(completion_text(&agg), "routed hi");
3569 Ok(())
3570 }
3571
3572 #[tokio::test]
3573 async fn invalid_raw_request_is_a_request_translation_error()
3574 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3575 let client = TranslatingLlmClient::new(&[])?;
3576 let Err(error) = client
3577 .call_rewrite_model_raw(
3578 json!("invalid"),
3579 None,
3580 Some(&ModelId::from("gpt")),
3581 WireFormat::OpenAiChat,
3582 )
3583 .await
3584 else {
3585 panic!("expected request translation to fail");
3586 };
3587
3588 assert!(matches!(
3589 error,
3590 LlmClientError::RequestTranslation(message) if !message.is_empty()
3591 ));
3592 Ok(())
3593 }
3594
3595 #[tokio::test]
3598 async fn call_rewrite_model_raw_round_trips_buffered_json()
3599 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3600 let server = MockServer::start().await;
3601 Mock::given(method("POST"))
3602 .and(path("/v1/chat/completions"))
3603 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
3604 "id": "chatcmpl-1",
3605 "model": "gpt",
3606 "choices": [{
3607 "index": 0,
3608 "message": {"role": "assistant", "content": "Hi there"},
3609 "finish_reason": "stop"
3610 }],
3611 "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}
3612 })))
3613 .mount(&server)
3614 .await;
3615
3616 let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
3617 let raw = json!({
3618 "model": "client-facing",
3619 "messages": [{"role": "user", "content": "hi"}]
3620 });
3621 let RawResponse::Buffered(body) = client
3622 .call_rewrite_model_raw(
3623 raw,
3624 None,
3625 Some(&ModelId::from("gpt")),
3626 WireFormat::OpenAiChat,
3627 )
3628 .await?
3629 else {
3630 panic!("expected a buffered response");
3631 };
3632
3633 assert_eq!(body["choices"][0]["message"]["content"], "Hi there");
3634 assert_eq!(body["model"], "gpt");
3636 Ok(())
3637 }
3638
3639 #[tokio::test]
3642 async fn call_rewrite_model_raw_restores_codex_mcp_namespace()
3643 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3644 let server = MockServer::start().await;
3645 Mock::given(method("POST"))
3646 .and(path("/v1/chat/completions"))
3647 .and(wiremock::matchers::body_partial_json(json!({
3648 "tools": [{
3649 "type": "function",
3650 "function": {"name": "mcp__open_websearch__search"}
3651 }]
3652 })))
3653 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
3654 "id": "chatcmpl-1",
3655 "model": "gpt",
3656 "choices": [{
3657 "index": 0,
3658 "message": {
3659 "role": "assistant",
3660 "content": null,
3661 "tool_calls": [{
3662 "id": "call_1",
3663 "type": "function",
3664 "function": {
3665 "name": "mcp__open_websearch__search",
3666 "arguments": "{\"q\":\"rust\"}"
3667 }
3668 }]
3669 },
3670 "finish_reason": "tool_calls"
3671 }],
3672 "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}
3673 })))
3674 .mount(&server)
3675 .await;
3676
3677 let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
3678 let raw = json!({
3679 "model": "client-facing",
3680 "input": "Search for Rust.",
3681 "tools": [{
3682 "type": "namespace",
3683 "name": "mcp__open_websearch",
3684 "tools": [{
3685 "type": "function",
3686 "name": "search",
3687 "description": "Search the web",
3688 "parameters": {"type": "object", "properties": {"q": {"type": "string"}}}
3689 }]
3690 }]
3691 });
3692
3693 let RawResponse::Buffered(body) = client
3694 .call_rewrite_model_raw(
3695 raw,
3696 None,
3697 Some(&ModelId::from("gpt")),
3698 WireFormat::OpenAiResponses,
3699 )
3700 .await?
3701 else {
3702 panic!("expected a buffered response");
3703 };
3704
3705 assert_eq!(body["output"][0]["type"], "function_call");
3706 assert_eq!(body["output"][0]["name"], "search");
3707 assert_eq!(body["output"][0]["namespace"], "mcp__open_websearch");
3708 assert_eq!(body["output"][0]["arguments"], "{\"q\": \"rust\"}");
3710 Ok(())
3711 }
3712
3713 #[tokio::test]
3714 async fn openai_responses_backend_flattens_codex_namespaces_before_upstream()
3715 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3716 let server = MockServer::start().await;
3717 Mock::given(method("POST"))
3718 .and(path("/v1/responses"))
3719 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
3720 "id": "resp_1",
3721 "model": "gpt",
3722 "object": "response",
3723 "created_at": 0,
3724 "status": "completed",
3725 "output": [{
3726 "type": "function_call",
3727 "call_id": "call_1",
3728 "name": "mcp__open_websearch__search",
3729 "arguments": "{\"q\":\"rust\"}"
3730 }],
3731 "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}
3732 })))
3733 .mount(&server)
3734 .await;
3735
3736 let client = TranslatingLlmClient::new(&responses_map(&format!("{}/v1", server.uri())))?;
3737 let parameters = json!({"type": "object", "properties": {"q": {"type": "string"}}});
3738 let raw = json!({
3739 "model": "client-facing",
3740 "input": [{
3741 "type": "function_call",
3742 "call_id": "call_0",
3743 "name": "search",
3744 "namespace": "mcp__open_websearch",
3745 "arguments": "{}"
3746 }],
3747 "tool_choice": {
3748 "type": "function",
3749 "name": "search",
3750 "namespace": "mcp__open_websearch"
3751 },
3752 "tools": [{
3753 "type": "namespace",
3754 "name": "mcp__open_websearch",
3755 "tools": [{
3756 "type": "function",
3757 "name": "search",
3758 "parameters": parameters.clone()
3759 }]
3760 }]
3761 });
3762
3763 let RawResponse::Buffered(body) = client
3764 .call_rewrite_model_raw(
3765 raw,
3766 None,
3767 Some(&ModelId::from("gpt")),
3768 WireFormat::OpenAiResponses,
3769 )
3770 .await?
3771 else {
3772 panic!("expected a buffered response");
3773 };
3774
3775 let received = server
3776 .received_requests()
3777 .await
3778 .ok_or("request recording should be enabled")?;
3779 let request_body: Value = serde_json::from_slice(&received[0].body)?;
3780 assert_eq!(request_body["model"], "gpt");
3781 assert_eq!(
3782 request_body["tools"],
3783 json!([{
3784 "type": "function",
3785 "name": "mcp__open_websearch__search",
3786 "description": "",
3787 "parameters": parameters
3788 }])
3789 );
3790 assert_eq!(
3791 request_body["tool_choice"],
3792 json!({"type": "function", "name": "mcp__open_websearch__search"})
3793 );
3794 assert_eq!(
3795 request_body["input"][0]["name"],
3796 "mcp__open_websearch__search"
3797 );
3798 assert!(request_body["input"][0].get("namespace").is_none());
3799
3800 assert_eq!(body["output"][0]["type"], "function_call");
3801 assert_eq!(body["output"][0]["name"], "search");
3802 assert_eq!(body["output"][0]["namespace"], "mcp__open_websearch");
3803 Ok(())
3804 }
3805
3806 #[tokio::test]
3812 async fn openai_responses_backend_drops_reasoning_the_provider_cannot_read()
3813 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3814 let server = MockServer::start().await;
3815 Mock::given(method("POST"))
3816 .and(path("/v1/responses"))
3817 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
3818 "id": "resp_1",
3819 "model": "gpt",
3820 "object": "response",
3821 "created_at": 0,
3822 "status": "completed",
3823 "output": [{
3824 "type": "message",
3825 "role": "assistant",
3826 "content": [{"type": "output_text", "text": "ok"}]
3827 }],
3828 "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}
3829 })))
3830 .mount(&server)
3831 .await;
3832
3833 let mut claude_reasoning = json!({
3834 "type": "reasoning",
3835 "id": "rs_b99feb9a-763d-4f7f-b6b2-2c0199b47ded",
3836 "summary": [{"type": "summary_text", "text": "Read the files first."}],
3837 "encrypted_content": null
3838 });
3839 let claude_reply = switchyard_translation::TranslationEngine::default()
3842 .translate_response(
3843 WireFormat::AnthropicMessages,
3844 WireFormat::OpenAiResponses,
3845 &json!({
3846 "id": "msg_1",
3847 "type": "message",
3848 "role": "assistant",
3849 "model": "claude",
3850 "content": [
3851 {"type": "thinking", "thinking": "Plan the edit.", "signature": "sig-abc"},
3852 {"type": "text", "text": "ok"}
3853 ],
3854 "stop_reason": "end_turn",
3855 "usage": {"input_tokens": 1, "output_tokens": 1}
3856 }),
3857 &switchyard_translation::TranslationPolicy::default(),
3858 )?
3859 .body;
3860 let switchyard_reasoning = claude_reply["output"]
3861 .as_array()
3862 .and_then(|output| output.iter().find(|item| item["type"] == "reasoning"))
3863 .cloned()
3864 .ok_or("translated Claude reply should contain a reasoning item")?;
3865 assert!(switchyard_translation::is_anthropic_thinking_envelope(
3866 switchyard_reasoning["encrypted_content"]
3867 .as_str()
3868 .unwrap_or_default()
3869 ));
3870 let mut gpt_reasoning = json!({
3871 "type": "reasoning",
3872 "id": "rs_gpt",
3873 "summary": [],
3874 "encrypted_content": "opaque-encrypted-reasoning"
3875 });
3876 let mut empty_reasoning =
3877 json!({"type": "reasoning", "id": "rs_empty", "summary": [], "encrypted_content": ""});
3878 let user = json!({"type": "message", "role": "user", "content": "fix the build"});
3879 let call = json!({"type": "function_call", "call_id": "call_1", "name": "exec", "arguments": "{}"});
3880 let output = json!({"type": "function_call_output", "call_id": "call_1", "output": "ok"});
3881 let input = json!([
3882 user,
3883 claude_reasoning,
3884 switchyard_reasoning,
3885 call,
3886 output,
3887 gpt_reasoning,
3888 empty_reasoning
3889 ]);
3890
3891 for has_input_override in [false, true] {
3892 let mut models = responses_map(&format!("{}/v1", server.uri()));
3893 if has_input_override {
3894 let Backend::OpenAiResponses(backend) = &mut models[0].default_backend else {
3895 return Err("expected Responses backend".into());
3896 };
3897 backend.omit_body_fields.insert("input".to_string());
3898 backend
3899 .extra_body
3900 .insert("input".to_string(), input.clone());
3901 }
3902 let client = TranslatingLlmClient::new(&models)?;
3903 for store in [json!(false), json!(true)] {
3904 let raw = json!({"model": "client-facing", "store": store, "input": input.clone()});
3905 client
3906 .call_rewrite_model_raw(
3907 raw,
3908 None,
3909 Some(&ModelId::from("gpt")),
3910 WireFormat::OpenAiResponses,
3911 )
3912 .await?;
3913 }
3914 }
3915
3916 let received = server
3917 .received_requests()
3918 .await
3919 .ok_or("request recording should be enabled")?;
3920 let sent: Vec<Value> = received
3921 .iter()
3922 .map(|request| serde_json::from_slice::<Value>(&request.body))
3923 .collect::<std::result::Result<_, _>>()?;
3924 assert_eq!(sent.len(), 4);
3925 for request in &sent {
3926 assert!(
3927 request["input"]
3928 .as_array()
3929 .ok_or("expected input array")?
3930 .iter()
3931 .all(|item| !item["encrypted_content"]
3932 .as_str()
3933 .is_some_and(switchyard_translation::is_anthropic_thinking_envelope)),
3934 "provider cannot read translated Anthropic reasoning"
3935 );
3936 }
3937 for item in [
3938 &mut claude_reasoning,
3939 &mut gpt_reasoning,
3940 &mut empty_reasoning,
3941 ] {
3942 item["content"] = json!([]);
3943 }
3944 for requests in sent.as_chunks::<2>().0 {
3945 assert_eq!(
3947 requests[0]["input"],
3948 json!([user, call, output, gpt_reasoning])
3949 );
3950 assert_eq!(
3953 requests[1]["input"],
3954 json!([
3955 user,
3956 claude_reasoning,
3957 call,
3958 output,
3959 gpt_reasoning,
3960 empty_reasoning
3961 ])
3962 );
3963 }
3964 Ok(())
3965 }
3966
3967 #[tokio::test]
3970 async fn call_rewrite_model_raw_streams_wire_events()
3971 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3972 use futures::TryStreamExt;
3973
3974 let server = MockServer::start().await;
3975 let sse = "data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}\n\n\
3976 data: {\"choices\":[{\"delta\":{\"content\":\" world\"}}]}\n\n\
3977 data: [DONE]\n\n";
3978 Mock::given(method("POST"))
3979 .and(path("/v1/chat/completions"))
3980 .respond_with(ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream"))
3981 .mount(&server)
3982 .await;
3983
3984 let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
3985 let raw = json!({
3986 "model": "client-facing",
3987 "messages": [{"role": "user", "content": "hi"}],
3988 "stream": true
3989 });
3990 let RawResponse::Stream(stream) = client
3991 .call_rewrite_model_raw(
3992 raw,
3993 None,
3994 Some(&ModelId::from("gpt")),
3995 WireFormat::OpenAiChat,
3996 )
3997 .await?
3998 else {
3999 panic!("expected a streamed response");
4000 };
4001
4002 let events: Vec<Value> = stream.try_collect().await?;
4003 assert!(!events.is_empty(), "expected at least one wire event");
4004 let content: String = events
4005 .iter()
4006 .filter_map(|event| event["choices"][0]["delta"]["content"].as_str())
4007 .collect();
4008 assert_eq!(content, "Hello world");
4009 assert!(events.iter().all(|event| event["model"] == "gpt"));
4012 Ok(())
4013 }
4014
4015 #[tokio::test]
4017 async fn call_rewrite_model_raw_forwards_headers()
4018 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
4019 let server = MockServer::start().await;
4020 Mock::given(method("POST"))
4021 .and(wiremock::matchers::header("x-request-id", "abc"))
4022 .and(wiremock::matchers::header("authorization", "Bearer secret"))
4024 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
4025 "id": "1", "model": "gpt",
4026 "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
4027 "usage": {}
4028 })))
4029 .mount(&server)
4030 .await;
4031
4032 let mut headers = http::HeaderMap::new();
4033 headers.insert("x-request-id", http::HeaderValue::from_static("abc"));
4034 headers.insert(
4035 "authorization",
4036 http::HeaderValue::from_static("Bearer client-key"),
4037 );
4038
4039 let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
4040 let raw = json!({"model": "gpt", "messages": [{"role": "user", "content": "hi"}]});
4041 client
4044 .call_rewrite_model_raw(
4045 raw,
4046 Some(headers),
4047 Some(&ModelId::from("gpt")),
4048 WireFormat::OpenAiChat,
4049 )
4050 .await?;
4051 Ok(())
4052 }
4053}