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