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