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, HeaderName, HeaderValue, 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(builder: RequestBuilder, backend: &Backend) -> RequestBuilder {
1104 let mut headers = HeaderMap::with_capacity(backend.extra_headers().len());
1105 for (name, value) in backend.extra_headers() {
1107 let name = HeaderName::from_bytes(name.as_bytes()).expect("validated header name");
1108 let value = HeaderValue::from_bytes(value.as_bytes()).expect("validated header value");
1109 headers.insert(name, value);
1110 }
1111 builder.headers(headers)
1112}
1113
1114fn set_json_model(body: &mut Value, model: &str) {
1116 if let Value::Object(object) = body {
1117 object.insert("model".to_string(), Value::String(model.to_string()));
1118 }
1119}
1120
1121fn normalize_codex_string_input(url: &str, body: &mut Value) {
1123 let codex = reqwest::Url::parse(url).ok().is_some_and(|url| {
1124 url.path()
1125 .trim_end_matches('/')
1126 .ends_with("/backend-api/codex/responses")
1127 });
1128 if codex && let Some(input @ Value::String(_)) = body.get_mut("input") {
1129 let text = input.take();
1130 *input = json!([{"type": "message", "role": "user", "content": text}]);
1131 }
1132}
1133
1134fn rejects_template_options(error: &LlmClientError) -> bool {
1135 let LlmClientError::UpstreamHttp { status, body } = error else {
1136 return false;
1137 };
1138 if !matches!(
1139 *status,
1140 StatusCode::BAD_REQUEST | StatusCode::UNPROCESSABLE_ENTITY
1141 ) {
1142 return false;
1143 }
1144 let parsed = serde_json::from_str::<Value>(body).ok();
1145 let error = parsed.as_ref().and_then(|value| value.get("error"));
1146 let message = match &parsed {
1147 Some(_) => error
1148 .and_then(|error| error.get("message"))
1149 .and_then(Value::as_str)
1150 .unwrap_or_default(),
1151 None => body,
1152 }
1153 .to_ascii_lowercase();
1154 let field = error
1155 .and_then(|error| error.get("param"))
1156 .and_then(Value::as_str);
1157 let code = error
1158 .and_then(|error| error.get("code"))
1159 .and_then(Value::as_str);
1160 (field == Some("chat_template_kwargs") || message.contains("chat_template_kwargs"))
1161 && (matches!(code, Some("unknown_parameter" | "unsupported_parameter"))
1162 || [
1163 "unrecognized request argument",
1164 "unknown parameter",
1165 "unsupported parameter",
1166 "unexpected keyword argument",
1167 "unknown field",
1168 "unsupported field",
1169 "extra inputs are not permitted",
1170 ]
1171 .iter()
1172 .any(|phrase| message.contains(phrase)))
1173}
1174
1175const CODEX_NAMESPACE_SEPARATOR: &str = "__";
1176
1177fn sanitize_openai_responses_provider_body(body: &mut Value) {
1182 let Value::Object(object) = body else {
1183 return;
1184 };
1185 sanitize_openai_responses_input_for_provider(object.get_mut("input"));
1186 sanitize_openai_responses_tools_for_provider(object.get_mut("tools"));
1187 sanitize_openai_responses_tool_choice_for_provider(object.get_mut("tool_choice"));
1188}
1189
1190fn sanitize_openai_responses_input_for_provider(input: Option<&mut Value>) {
1191 let Some(Value::Array(items)) = input else {
1192 return;
1193 };
1194 for item in items {
1195 let Some(object) = item.as_object_mut() else {
1196 continue;
1197 };
1198 if object.get("type").and_then(Value::as_str) == Some("function_call") {
1199 qualify_responses_function_name(object);
1200 }
1201 }
1202}
1203
1204fn sanitize_openai_responses_tools_for_provider(tools: Option<&mut Value>) {
1205 let Some(Value::Array(tools)) = tools else {
1206 return;
1207 };
1208 let mut flat_tools = Vec::with_capacity(tools.len());
1209 for tool in std::mem::take(tools) {
1210 push_sanitized_openai_responses_tool(&mut flat_tools, tool);
1211 }
1212 *tools = flat_tools;
1213}
1214
1215fn push_sanitized_openai_responses_tool(out: &mut Vec<Value>, tool: Value) {
1216 let Value::Object(mut object) = tool else {
1217 out.push(tool);
1218 return;
1219 };
1220 if object.get("type").and_then(Value::as_str) == Some("namespace") {
1221 let namespace = object
1222 .get("name")
1223 .and_then(Value::as_str)
1224 .unwrap_or_default()
1225 .to_string();
1226 let Some(Value::Array(children)) = object.remove("tools") else {
1227 out.push(Value::Object(object));
1228 return;
1229 };
1230 for child in children {
1231 push_sanitized_namespaced_tool(out, &namespace, child);
1232 }
1233 return;
1234 }
1235 ensure_responses_function_tool_description(&mut object);
1236 out.push(Value::Object(object));
1237}
1238
1239fn push_sanitized_namespaced_tool(out: &mut Vec<Value>, namespace: &str, tool: Value) {
1240 let Value::Object(mut object) = tool else {
1241 out.push(tool);
1242 return;
1243 };
1244 if object.get("type").and_then(Value::as_str) == Some("function") {
1245 qualify_responses_function_name_with_namespace(&mut object, namespace);
1246 ensure_responses_function_tool_description(&mut object);
1247 }
1248 out.push(Value::Object(object));
1249}
1250
1251fn sanitize_openai_responses_tool_choice_for_provider(tool_choice: Option<&mut Value>) {
1252 let Some(Value::Object(object)) = tool_choice else {
1253 return;
1254 };
1255 if object.get("type").and_then(Value::as_str) == Some("function") {
1256 qualify_responses_function_name(object);
1257 }
1258}
1259
1260fn qualify_responses_function_name(object: &mut Map<String, Value>) {
1261 let namespace = object
1262 .remove("namespace")
1263 .and_then(|value| value.as_str().map(ToOwned::to_owned));
1264 let Some(namespace) = namespace.as_deref() else {
1265 return;
1266 };
1267 qualify_responses_function_name_with_namespace(object, namespace);
1268}
1269
1270fn qualify_responses_function_name_with_namespace(
1271 object: &mut Map<String, Value>,
1272 namespace: &str,
1273) {
1274 if namespace.is_empty() {
1275 return;
1276 }
1277 let Some(name) = object.get("name").and_then(Value::as_str) else {
1278 return;
1279 };
1280 let prefix = format!("{namespace}{CODEX_NAMESPACE_SEPARATOR}");
1281 if name.starts_with(&prefix) {
1282 return;
1283 }
1284 object.insert("name".to_string(), Value::String(format!("{prefix}{name}")));
1285}
1286
1287fn ensure_responses_function_tool_description(object: &mut Map<String, Value>) {
1288 if object.get("type").and_then(Value::as_str) == Some("function")
1289 && !matches!(object.get("description"), Some(Value::String(_)))
1290 {
1291 object.insert("description".to_string(), Value::String(String::new()));
1292 }
1293}
1294
1295fn strip_anthropic_incompatible_fields(body: &mut Value) {
1303 if let Value::Object(object) = body {
1304 object.remove("reasoning_effort");
1305 object.remove("context_management");
1306 if let Some(Value::Array(tools)) = object.get_mut("tools") {
1308 for tool in tools {
1309 if let Value::Object(tool) = tool
1310 && tool.get("strict") == Some(&Value::Bool(false))
1311 {
1312 tool.remove("strict");
1313 }
1314 }
1315 }
1316 }
1317}
1318
1319fn nvidia_inference_api_backend(backend: &Backend) -> bool {
1322 reqwest::Url::parse(backend.base_url())
1323 .ok()
1324 .is_some_and(|url| {
1325 url.host_str()
1326 .is_some_and(|host| host.eq_ignore_ascii_case("inference-api.nvidia.com"))
1327 })
1328}
1329
1330fn strip_anthropic_tool_strict(body: &mut Value) {
1331 let Some(tools) = body.get_mut("tools").and_then(Value::as_array_mut) else {
1332 return;
1333 };
1334 for tool in tools.iter_mut().filter_map(Value::as_object_mut) {
1335 tool.remove("strict");
1336 }
1337}
1338
1339fn strip_unsigned_thinking_blocks(body: &mut Value) {
1347 let Value::Object(object) = body else {
1348 return;
1349 };
1350 let Some(Value::Array(messages)) = object.get_mut("messages") else {
1351 return;
1352 };
1353 for message in messages {
1354 strip_unsigned_thinking_from_message(message);
1355 }
1356}
1357
1358fn strip_unsigned_thinking_from_message(message: &mut Value) {
1361 let Value::Object(message) = message else {
1362 return;
1363 };
1364 let Some(Value::Array(blocks)) = message.get("content") else {
1365 return;
1366 };
1367 if !blocks.iter().any(is_unsigned_thinking_block) {
1368 return;
1369 }
1370 let Some(Value::Array(blocks)) = message.get_mut("content") else {
1371 return;
1372 };
1373 blocks.retain(|block| !is_unsigned_thinking_block(block));
1374 if blocks.is_empty() {
1375 message.insert("content".to_string(), Value::String(String::new()));
1376 }
1377}
1378
1379fn is_unsigned_thinking_block(block: &Value) -> bool {
1381 if block.get("type").and_then(Value::as_str) != Some("thinking") {
1382 return false;
1383 }
1384 !matches!(
1385 block.get("signature").and_then(Value::as_str),
1386 Some(signature) if !signature.is_empty()
1387 )
1388}
1389
1390fn apply_reasoning_effort(body: &mut Value, backend: &Backend) {
1394 let Some(effort) = backend.reasoning_effort() else {
1395 return;
1396 };
1397 let Value::Object(object) = body else {
1398 return;
1399 };
1400 match backend {
1401 Backend::OpenAiResponses(_) => set_nested_effort(object, "reasoning", effort),
1403 Backend::Anthropic(_) => set_nested_effort(object, "output_config", effort),
1404 Backend::OpenAiChat(_) => {
1405 object.insert(
1407 "reasoning_effort".to_string(),
1408 Value::String(effort.to_string()),
1409 );
1410 }
1411 }
1412}
1413
1414fn set_nested_effort(object: &mut serde_json::Map<String, Value>, key: &str, effort: &str) {
1417 let container = object
1418 .entry(key.to_string())
1419 .or_insert_with(|| Value::Object(serde_json::Map::new()));
1420 if !container.is_object() {
1421 *container = Value::Object(serde_json::Map::new());
1422 }
1423 if let Value::Object(container) = container {
1424 container.insert("effort".to_string(), Value::String(effort.to_string()));
1425 }
1426}
1427
1428fn merge_extra_body(body: &mut Value, extra_body: &BTreeMap<String, Value>) {
1430 let Value::Object(object) = body else {
1431 return;
1432 };
1433 for (key, value) in extra_body {
1434 object.entry(key.clone()).or_insert_with(|| value.clone());
1435 }
1436}
1437
1438fn omit_configured_body_fields(body: &mut Value, omit_body_fields: &BTreeSet<String>) {
1439 let Value::Object(object) = body else {
1440 return;
1441 };
1442 for key in omit_body_fields {
1443 object.remove(key);
1444 }
1445}
1446
1447const MAX_CACHE_CONTROL_BLOCKS: usize = 4;
1450
1451fn count_cache_control_blocks(value: &Value) -> usize {
1454 match value {
1455 Value::Object(map) => {
1456 usize::from(map.contains_key("cache_control"))
1457 + map.values().map(count_cache_control_blocks).sum::<usize>()
1458 }
1459 Value::Array(items) => items.iter().map(count_cache_control_blocks).sum(),
1460 _ => 0,
1461 }
1462}
1463
1464fn enable_anthropic_prompt_caching(body: &mut Value, temporary_note_count: usize) {
1466 if count_cache_control_blocks(body) >= MAX_CACHE_CONTROL_BLOCKS {
1469 return;
1470 }
1471 let marker = body
1474 .get("cache_control")
1475 .cloned()
1476 .unwrap_or_else(|| serde_json::json!({"type":"ephemeral"}));
1477 let Some(messages) = body.get_mut("messages").and_then(Value::as_array_mut) else {
1478 return;
1479 };
1480 let mut remaining = temporary_note_count;
1481 for message in messages.iter_mut().rev() {
1483 match message.get_mut("content") {
1484 Some(content @ Value::String(_)) if remaining == 0 => {
1485 let text = std::mem::take(content);
1486 *content = serde_json::json!([{
1487 "type": "text", "text": text, "cache_control": marker
1488 }]);
1489 return;
1490 }
1491 Some(Value::String(_)) => remaining = remaining.saturating_sub(1),
1492 Some(Value::Array(blocks)) => {
1493 if remaining >= blocks.len() {
1494 remaining -= blocks.len();
1495 continue;
1496 }
1497 let index = blocks.len() - remaining - 1;
1498 if let Some(block) = blocks[index].as_object_mut() {
1499 block.entry("cache_control".to_string()).or_insert(marker);
1500 }
1501 return;
1502 }
1503 _ => {}
1504 }
1505 }
1506}
1507
1508fn enable_openai_guidance_caching(
1511 body: &mut Value,
1512 format: WireFormat,
1513 model: &ModelId,
1514 suffix: usize,
1515) {
1516 let name = model.as_str().rsplit('/').next().unwrap_or_default();
1517 let supported = name == "gpt-5.6"
1518 || name.starts_with("gpt-5.6-")
1519 || name == "gpt-6"
1520 || name.starts_with("gpt-6-")
1521 || name.starts_with("gpt-6.");
1522 if suffix == 0
1525 || !supported
1526 || body
1527 .pointer("/prompt_cache_options/mode")
1528 .and_then(Value::as_str)
1529 == Some("explicit")
1530 {
1531 return;
1532 }
1533 let key = if format == WireFormat::OpenAiResponses {
1534 "input"
1535 } else {
1536 "messages"
1537 };
1538 let Some(items) = body.get_mut(key).and_then(Value::as_array_mut) else {
1539 return;
1540 };
1541 let marker_count = items
1542 .iter()
1543 .map(|item| {
1544 ["content", "output"]
1545 .iter()
1546 .filter_map(|key| item.get(key).and_then(Value::as_array))
1547 .flatten()
1548 .filter(|block| block.get("prompt_cache_breakpoint").is_some())
1549 .count()
1550 })
1551 .sum::<usize>();
1552 if marker_count >= 3 {
1554 return;
1555 }
1556 let Some(prefix_len) = items.len().checked_sub(suffix) else {
1557 return;
1558 };
1559 for item in items[..prefix_len].iter_mut().rev() {
1560 let field = if item.get("type").and_then(Value::as_str) == Some("function_call_output") {
1561 "output"
1562 } else if matches!(
1563 item.get("role").and_then(Value::as_str),
1564 Some("user" | "tool" | "developer" | "system")
1565 ) {
1566 "content"
1567 } else {
1568 continue;
1569 };
1570 let Some(content) = item.get_mut(field) else {
1571 continue;
1572 };
1573 if let Value::String(text) = content {
1574 let text = std::mem::take(text);
1575 let kind = if format == WireFormat::OpenAiResponses {
1576 "input_text"
1577 } else {
1578 "text"
1579 };
1580 *content = serde_json::json!([{"type":kind,"text":text}]);
1581 }
1582 if let Some(block) = content
1584 .as_array_mut()
1585 .and_then(|blocks| blocks.last_mut())
1586 .and_then(Value::as_object_mut)
1587 {
1588 block
1589 .entry("prompt_cache_breakpoint".to_string())
1590 .or_insert_with(|| serde_json::json!({"mode":"explicit"}));
1591 return;
1592 }
1593 }
1594}
1595
1596fn ensure_openai_stream_usage(body: &mut Value) {
1598 let Value::Object(object) = body else {
1599 return;
1600 };
1601 if object.get("stream").and_then(Value::as_bool) != Some(true) {
1602 return;
1603 }
1604
1605 match object.get_mut("stream_options") {
1606 Some(Value::Object(options)) => {
1607 options
1608 .entry("include_usage".to_string())
1609 .or_insert(Value::Bool(true));
1610 }
1611 _ => {
1612 let mut options = Map::new();
1613 options.insert("include_usage".to_string(), Value::Bool(true));
1614 object.insert("stream_options".to_string(), Value::Object(options));
1615 }
1616 }
1617}
1618
1619fn is_allowed_metadata_header(name: &str) -> bool {
1620 ALLOWED_METADATA_HEADERS
1621 .iter()
1622 .any(|allowed| name.eq_ignore_ascii_case(allowed))
1623}
1624
1625fn is_non_forwardable_header(name: &str, headers: &HeaderMap) -> bool {
1626 NON_FORWARDABLE_HEADERS
1627 .iter()
1628 .any(|blocked| name.eq_ignore_ascii_case(blocked))
1629 || headers.get_all("connection").iter().any(|value| {
1630 value
1631 .as_bytes()
1632 .split(|byte| *byte == b',')
1633 .any(|option| option.trim_ascii().eq_ignore_ascii_case(name.as_bytes()))
1634 })
1635}
1636
1637fn redact_forwarded_headers(
1639 mut body: String,
1640 metadata: Option<&Metadata>,
1641 is_forwarding_auth: bool,
1642) -> String {
1643 if !is_forwarding_auth {
1644 return body;
1645 }
1646 let Some(headers) = metadata.and_then(|metadata| metadata.http_headers.as_ref()) else {
1647 return body;
1648 };
1649 for (name, value) in headers {
1650 if is_non_forwardable_header(name.as_str(), headers) {
1651 continue;
1652 }
1653 let Ok(value) = value.to_str() else {
1654 continue;
1655 };
1656 if !value.is_empty() {
1657 body = body.replace(value, "[REDACTED]");
1658 }
1659 }
1660 body
1661}
1662
1663#[cfg(test)]
1664mod tests {
1665 use std::collections::BTreeMap;
1666 use std::error::Error;
1667 use std::io::{Read, Write};
1668 use std::sync::Arc;
1669 use std::sync::atomic::{AtomicUsize, Ordering};
1670 use std::thread::JoinHandle;
1671
1672 use serde_json::json;
1673 use switchyard_protocol::{completion_text, text_request};
1674 use wiremock::matchers::{method, path};
1675 use wiremock::{Mock, MockServer, ResponseTemplate};
1676
1677 use super::*;
1678 use crate::backend::HttpBackendConfig;
1679
1680 fn config(base_url: &str) -> HttpBackendConfig {
1681 HttpBackendConfig {
1682 base_url: base_url.to_string(),
1683 api_key: Some("secret".to_string()),
1684 forward_auth: false,
1685 extra_headers: BTreeMap::new(),
1686 extra_body: BTreeMap::new(),
1687 omit_body_fields: BTreeSet::new(),
1688 reasoning_effort: None,
1689 max_retries: 0,
1690 failure_cooldown: Duration::ZERO,
1691 timeout: None,
1692 }
1693 }
1694
1695 #[test]
1696 fn provider_endpoint_urls_preserve_query_and_fragment() {
1697 let invalid_url = "not a URL/";
1698 let backend = Backend::OpenAiResponses(config(invalid_url));
1699 assert!(openai_url(invalid_url, "/responses").is_err());
1700 assert_eq!(backend.url(), invalid_url);
1701 assert_eq!(
1702 AuxiliaryOperation::ResponsesCompact.url(&backend),
1703 invalid_url
1704 );
1705
1706 let anthropic = Backend::Anthropic(config(invalid_url));
1707 assert_eq!(anthropic.url(), invalid_url);
1708 assert_eq!(anthropic.count_tokens_url(), invalid_url);
1709
1710 for path in ["", "/v1/", "/v1/messages"] {
1711 for tail in [
1712 "",
1713 "?api-version=2026-09-01&value=%2F/",
1714 "#fragment?value=/",
1715 ] {
1716 let base_url = format!("https://example.com{path}{tail}");
1717 let backend = Backend::Anthropic(config(&base_url));
1718 assert_eq!(
1719 backend.url(),
1720 format!("https://example.com/v1/messages{tail}")
1721 );
1722 assert_eq!(
1723 AuxiliaryOperation::AnthropicCountTokens.url(&backend),
1724 format!("https://example.com/v1/messages/count_tokens{tail}")
1725 );
1726 }
1727 }
1728
1729 for path in ["/v1", "/v1/responses/", "/v1/chat/completions"] {
1730 for tail in [
1731 "",
1732 "?api-version=2026-09-01&value=%2F/",
1733 "#fragment?value=/",
1734 ] {
1735 let base_url = format!("https://example.com{path}{tail}");
1736 let responses = Backend::OpenAiResponses(config(&base_url));
1737 let chat = Backend::OpenAiChat(config(&base_url));
1738 for (actual, endpoint) in [
1739 (responses.url(), "/responses"),
1740 (chat.url(), "/chat/completions"),
1741 (
1742 AuxiliaryOperation::ResponsesCompact.url(&responses),
1743 "/responses/compact",
1744 ),
1745 (
1746 AuxiliaryOperation::ResponsesInputTokens.url(&responses),
1747 "/responses/input_tokens",
1748 ),
1749 ] {
1750 assert_eq!(actual, format!("https://example.com/v1{endpoint}{tail}"));
1751 }
1752 }
1753 }
1754 }
1755
1756 fn config_with_retries(base_url: &str, max_retries: u32) -> HttpBackendConfig {
1757 HttpBackendConfig {
1758 max_retries,
1759 ..config(base_url)
1760 }
1761 }
1762
1763 fn forwarding_config(base_url: &str) -> HttpBackendConfig {
1764 HttpBackendConfig {
1765 api_key: None,
1766 forward_auth: true,
1767 ..config(base_url)
1768 }
1769 }
1770
1771 fn chat_map(base_url: &str) -> Vec<ModelConfig> {
1773 vec![ModelConfig::new(
1774 "gpt",
1775 Backend::OpenAiChat(config(base_url)),
1776 None,
1777 )]
1778 }
1779
1780 fn responses_map(base_url: &str) -> Vec<ModelConfig> {
1781 vec![ModelConfig::new(
1782 "gpt",
1783 Backend::OpenAiResponses(config(base_url)),
1784 None,
1785 )]
1786 }
1787
1788 fn chat_map_with_extra_body(
1789 base_url: &str,
1790 extra_body: BTreeMap<String, Value>,
1791 ) -> Vec<ModelConfig> {
1792 let mut backend = config(base_url);
1793 backend.extra_body = extra_body;
1794 vec![ModelConfig::new("gpt", Backend::OpenAiChat(backend), None)]
1795 }
1796
1797 fn chat_map_with_omit_body_fields(
1798 base_url: &str,
1799 omit_body_fields: BTreeSet<String>,
1800 ) -> Vec<ModelConfig> {
1801 let mut backend = config(base_url);
1802 backend.omit_body_fields = omit_body_fields;
1803 vec![ModelConfig::new("gpt", Backend::OpenAiChat(backend), None)]
1804 }
1805
1806 fn chat_map_with_effort(base_url: &str, effort: &str) -> Vec<ModelConfig> {
1807 let mut backend = config(base_url);
1808 backend.reasoning_effort = Some(effort.to_string());
1809 vec![ModelConfig::new("gpt", Backend::OpenAiChat(backend), None)]
1810 }
1811
1812 fn responses_map_with_effort(base_url: &str, effort: &str) -> Vec<ModelConfig> {
1813 let mut backend = config(base_url);
1814 backend.reasoning_effort = Some(effort.to_string());
1815 vec![ModelConfig::new(
1816 "gpt",
1817 Backend::OpenAiResponses(backend),
1818 None,
1819 )]
1820 }
1821
1822 fn anthropic_map(base_url: &str) -> Vec<ModelConfig> {
1823 vec![ModelConfig::new(
1824 "claude",
1825 Backend::Anthropic(config(base_url)),
1826 None,
1827 )]
1828 }
1829
1830 #[tokio::test]
1831 async fn model_properties_uses_configured_auth_and_backend_path()
1832 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
1833 for (base_path, props_path) in [
1834 ("/v1", "/props"),
1835 ("/v1/", "/props"),
1836 ("/v1/chat/completions", "/props"),
1837 ("/v1/responses/", "/props"),
1838 ("/prefix/v1", "/prefix/props"),
1839 ("/prefix/v1/chat/completions", "/prefix/props"),
1840 ("/prefix/v1/responses", "/prefix/props"),
1841 ("/prefix/", "/prefix/props"),
1842 ] {
1843 let server = MockServer::start().await;
1844 Mock::given(method("GET"))
1845 .and(path(props_path))
1846 .and(wiremock::matchers::header("authorization", "Bearer secret"))
1847 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
1848 "default_generation_settings": {"n_ctx": 65_536}
1849 })))
1850 .expect(1)
1851 .mount(&server)
1852 .await;
1853 let client =
1854 TranslatingLlmClient::new(&chat_map(&format!("{}{base_path}", server.uri())))?;
1855 let properties = client
1856 .get_model_properties(
1857 &ModelId::from("gpt"),
1858 WireFormat::OpenAiChat,
1859 Duration::from_secs(1),
1860 )
1861 .await?;
1862 assert_eq!(properties["default_generation_settings"]["n_ctx"], 65_536);
1863 }
1864 Ok(())
1865 }
1866
1867 fn chat_map_with_retries(base_url: &str, max_retries: u32) -> Vec<ModelConfig> {
1868 vec![ModelConfig::new(
1869 "gpt",
1870 Backend::OpenAiChat(config_with_retries(base_url, max_retries)),
1871 None,
1872 )]
1873 }
1874
1875 fn chat_success_response() -> ResponseTemplate {
1876 ResponseTemplate::new(200).set_body_json(json!({
1877 "id": "chatcmpl-1",
1878 "model": "gpt",
1879 "choices": [{
1880 "index": 0,
1881 "message": {"role": "assistant", "content": "recovered"},
1882 "finish_reason": "stop"
1883 }],
1884 "usage": {}
1885 }))
1886 }
1887
1888 fn truncated_response_server(
1889 content_type: &str,
1890 body: &str,
1891 ) -> std::io::Result<(String, JoinHandle<std::io::Result<()>>)> {
1892 response_sequence_server(vec![format!(
1893 "HTTP/1.1 200 OK\r\nContent-Type: {content_type}\r\n\
1894 Content-Length: {}\r\nConnection: close\r\n\r\n{body}",
1895 body.len() + 100
1896 )])
1897 }
1898
1899 fn response_sequence_server(
1900 responses: Vec<String>,
1901 ) -> std::io::Result<(String, JoinHandle<std::io::Result<()>>)> {
1902 let listener = std::net::TcpListener::bind("127.0.0.1:0")?;
1903 let address = listener.local_addr()?;
1904 let handle = std::thread::spawn(move || {
1905 for response in responses {
1906 let (mut stream, _) = listener.accept()?;
1907 let mut request = [0_u8; 1024];
1908 if stream.read(&mut request)? == 0 {
1909 return Err(std::io::Error::new(
1910 std::io::ErrorKind::UnexpectedEof,
1911 "client closed before sending a request",
1912 ));
1913 }
1914 stream.write_all(response.as_bytes())?;
1915 }
1916 Ok(())
1917 });
1918 Ok((format!("http://{address}/v1"), handle))
1919 }
1920
1921 fn raw_chat_success_response() -> String {
1922 let body = r#"{"id":"chatcmpl-1","model":"gpt","choices":[{"index":0,"message":{"role":"assistant","content":"recovered"},"finish_reason":"stop"}],"usage":{}}"#;
1923 format!(
1924 "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n\
1925 Content-Length: {}\r\nConnection: close\r\n\r\n{body}",
1926 body.len()
1927 )
1928 }
1929
1930 fn request_for(model: Option<&str>, stream: bool) -> Request {
1931 let mut llm_request = text_request(model.map(str::to_string), "hi");
1932 llm_request.stream = stream;
1933 Request {
1934 llm_request,
1935 raw_request: None,
1936 metadata: None,
1937 }
1938 }
1939
1940 fn request_with_headers(model: &str, headers: HeaderMap) -> Request {
1941 let mut request = request_for(Some(model), false);
1942 request.metadata = Some(Metadata {
1943 http_headers: Some(headers),
1944 ..Default::default()
1945 });
1946 request
1947 }
1948
1949 #[tokio::test]
1950 async fn transport_errors_drop_the_upstream_url() {
1951 let error = reqwest::Client::new()
1952 .post("http://127.0.0.1:1/v1?key=CANARY")
1953 .send()
1954 .await
1955 .expect_err("closed port");
1956
1957 assert!(!convert_reqwest_error(error).to_string().contains("CANARY"));
1958 }
1959
1960 #[tokio::test]
1961 async fn captured_request_build_failure_remains_a_configuration_error() {
1962 let client = TranslatingLlmClient::new(&chat_map("://invalid")).unwrap();
1963 let result =
1964 crate::with_payload_capture(client.call(request_for(Some("gpt"), false))).await;
1965 assert!(matches!(result, Err(LlmClientError::Configuration { .. })));
1966 }
1967
1968 #[test]
1969 fn nvidia_anthropic_tools_drop_unsupported_strict_field() {
1970 let backend = Backend::Anthropic(config("https://inference-api.nvidia.com"));
1971 assert!(nvidia_inference_api_backend(&backend));
1972 assert!(!nvidia_inference_api_backend(&Backend::Anthropic(config(
1973 "https://inference-api.nvidia.com.example.test"
1974 ))));
1975 let mut body = json!({"tools": [{"name": "a", "input_schema": {}, "strict": true}]});
1976 strip_anthropic_tool_strict(&mut body);
1977 assert!(body["tools"][0].get("strict").is_none());
1978 }
1979
1980 #[test]
1981 fn anthropic_prompt_caching_marks_final_message() {
1982 let caller_managed = json!({
1983 "cache_control": {"type": "ephemeral", "ttl": "1h"},
1984 "messages": [{"role": "user", "content": "hello"}]
1985 });
1986 let mut body = caller_managed.clone();
1987 enable_anthropic_prompt_caching(&mut body, 0);
1988 assert_eq!(body["cache_control"], caller_managed["cache_control"]);
1989 assert_eq!(
1990 body["messages"][0]["content"],
1991 json!([{
1992 "type": "text", "text": "hello",
1993 "cache_control": {"type": "ephemeral", "ttl": "1h"}
1994 }])
1995 );
1996
1997 let mut body = json!({
1998 "messages": [{"role": "user", "content": "hello"}]
1999 });
2000
2001 enable_anthropic_prompt_caching(&mut body, 0);
2002
2003 assert_eq!(
2004 body["messages"][0]["content"][0]["cache_control"],
2005 json!({"type": "ephemeral"})
2006 );
2007 }
2008
2009 fn cache_control_blocks(value: &Value) -> usize {
2011 count_cache_control_blocks(value)
2012 }
2013
2014 fn body_at_the_cache_control_limit() -> Value {
2017 json!({
2018 "system": [
2019 {"type": "text", "text": "a", "cache_control": {"type": "ephemeral"}},
2020 {"type": "text", "text": "b", "cache_control": {"type": "ephemeral"}}
2021 ],
2022 "tools": [
2023 {"name": "t", "input_schema": {"type": "object"},
2024 "cache_control": {"type": "ephemeral"}}
2025 ],
2026 "messages": [
2027 {"role": "user", "content": [
2028 {"type": "text", "text": "old", "cache_control": {"type": "ephemeral"}}
2029 ]},
2030 {"role": "assistant", "content": [{"type": "text", "text": "ok"}]},
2031 {"role": "user", "content": [{"type": "text", "text": "new"}]}
2032 ]
2033 })
2034 }
2035
2036 #[test]
2037 fn anthropic_prompt_caching_abstains_at_the_cache_control_limit() {
2038 let mut body = body_at_the_cache_control_limit();
2039 assert_eq!(cache_control_blocks(&body), 4);
2040
2041 enable_anthropic_prompt_caching(&mut body, 0);
2042
2043 assert_eq!(cache_control_blocks(&body), 4);
2046 assert!(
2047 body["messages"][2]["content"][0]
2048 .get("cache_control")
2049 .is_none()
2050 );
2051 }
2052
2053 #[test]
2054 fn anthropic_prompt_caching_still_marks_one_below_the_limit() {
2055 let mut body = body_at_the_cache_control_limit();
2056 body["system"].as_array_mut().unwrap().pop();
2058 assert_eq!(cache_control_blocks(&body), 3);
2059
2060 enable_anthropic_prompt_caching(&mut body, 0);
2061
2062 assert_eq!(cache_control_blocks(&body), 4);
2063 assert_eq!(
2064 body["messages"][2]["content"][0]["cache_control"],
2065 json!({"type": "ephemeral"})
2066 );
2067 }
2068
2069 fn request_with_wire_format(model: &str, format: WireFormat) -> Request {
2072 let mut request = request_for(Some(model), false);
2073 request.metadata = Some(Metadata {
2074 session_id: None,
2075 agent_id: None,
2076 task_id: None,
2077 correlation_id: None,
2078 extra_metadata: None,
2079 http_headers: None,
2080 wire_format: Some(format),
2081 ..Default::default()
2082 });
2083 request
2084 }
2085
2086 #[tokio::test]
2087 async fn missing_model_errors()
2088 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2089 let client = TranslatingLlmClient::new(&[])?;
2090 let Err(error) = client
2091 .call_rewrite_model(request_for(None, false), None)
2092 .await
2093 else {
2094 panic!("expected an error");
2095 };
2096 assert!(matches!(
2097 error,
2098 LlmClientError::InvalidRequest { message } if message == "no model given"
2099 ));
2100 Ok(())
2101 }
2102
2103 #[tokio::test]
2104 async fn unknown_model_errors()
2105 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2106 let client = TranslatingLlmClient::new(&[])?;
2107 let Err(error) = client
2108 .call_rewrite_model(request_for(Some("gpt"), false), None)
2109 .await
2110 else {
2111 panic!("expected an error");
2112 };
2113 assert!(matches!(
2114 error,
2115 LlmClientError::Configuration { message }
2116 if message.contains("gpt")
2117 ));
2118 Ok(())
2119 }
2120
2121 #[tokio::test]
2122 async fn unknown_model_format_errors()
2123 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2124 let client = TranslatingLlmClient::new(&chat_map("https://example.test/v1"))?;
2126 let Err(error) = client
2127 .call_rewrite_model(
2128 request_with_wire_format("gpt", WireFormat::AnthropicMessages),
2129 None,
2130 )
2131 .await
2132 else {
2133 panic!("expected an error");
2134 };
2135 assert!(matches!(
2136 error,
2137 LlmClientError::Configuration { message }
2138 if message.contains("gpt")
2139 && message.contains(&WireFormat::AnthropicMessages.to_string())
2140 ));
2141 Ok(())
2142 }
2143
2144 #[test]
2145 fn backend_for_resolves_configured_format()
2146 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2147 let client = TranslatingLlmClient::new(&chat_map("https://example.test/v1"))?;
2148 assert!(
2150 client
2151 .backend_for(&ModelId::from("gpt"), WireFormat::OpenAiChat)
2152 .is_some()
2153 );
2154 assert!(
2155 client
2156 .backend_for(&ModelId::from("gpt"), WireFormat::AnthropicMessages)
2157 .is_none()
2158 );
2159 assert!(
2160 client
2161 .backend_for(&ModelId::from("missing"), WireFormat::OpenAiChat)
2162 .is_none()
2163 );
2164 Ok(())
2165 }
2166
2167 #[tokio::test]
2168 async fn target_name_arg_wins_over_request_model()
2169 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2170 let client = TranslatingLlmClient::new(&[])?;
2171 let Err(error) = client
2173 .call_rewrite_model(request_for(Some("a"), false), Some(&ModelId::from("b")))
2174 .await
2175 else {
2176 panic!("expected an error");
2177 };
2178 assert!(matches!(
2179 error,
2180 LlmClientError::Configuration { message }
2181 if message.contains("\"b\"")
2182 ));
2183 Ok(())
2184 }
2185
2186 #[tokio::test]
2187 async fn buffered_openai_chat_round_trips()
2188 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2189 let server = MockServer::start().await;
2190 Mock::given(method("POST"))
2191 .and(path("/v1/chat/completions"))
2192 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
2193 "id": "chatcmpl-1",
2194 "model": "gpt",
2195 "choices": [{
2196 "index": 0,
2197 "message": {"role": "assistant", "content": "Hi there"},
2198 "finish_reason": "stop"
2199 }],
2200 "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}
2201 })))
2202 .mount(&server)
2203 .await;
2204
2205 let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
2206
2207 let response = client
2208 .call_rewrite_model(request_for(Some("gpt"), false), None)
2209 .await?;
2210 let agg = response.llm_response.into_agg().await?;
2211 assert_eq!(completion_text(&agg), "Hi there");
2212
2213 Ok(())
2214 }
2215
2216 #[tokio::test]
2217 async fn invalid_json_is_a_response_translation_error()
2218 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2219 let server = MockServer::start().await;
2220 let calls = Arc::new(AtomicUsize::new(0));
2221 let observed_calls = Arc::clone(&calls);
2222 Mock::given(method("POST"))
2223 .respond_with(move |_: &wiremock::Request| {
2224 observed_calls.fetch_add(1, Ordering::SeqCst);
2225 ResponseTemplate::new(200).set_body_raw("not json", "application/json")
2226 })
2227 .mount(&server)
2228 .await;
2229
2230 let client =
2231 TranslatingLlmClient::new(&chat_map_with_retries(&format!("{}/v1", server.uri()), 2))?;
2232 let Err(error) = client
2233 .call_rewrite_model(request_for(Some("gpt"), false), None)
2234 .await
2235 else {
2236 panic!("expected invalid JSON to fail");
2237 };
2238
2239 assert!(matches!(
2240 error,
2241 LlmClientError::ResponseTranslation(message)
2242 if message.contains("invalid upstream JSON")
2243 ));
2244 assert_eq!(calls.load(Ordering::SeqCst), 1);
2245 Ok(())
2246 }
2247
2248 #[tokio::test]
2249 async fn response_body_io_failure_is_a_transport_error()
2250 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2251 let (base_url, server) = truncated_response_server("application/json", "{}")?;
2252 let client = TranslatingLlmClient::new(&chat_map(&base_url))?;
2253 let result = client
2254 .call_rewrite_model(request_for(Some("gpt"), false), None)
2255 .await;
2256 server
2257 .join()
2258 .map_err(|_| std::io::Error::other("response server thread panicked"))??;
2259
2260 let Err(error) = result else {
2261 panic!("expected the truncated response body to fail");
2262 };
2263 let LlmClientError::Transport { source } = error else {
2264 panic!("expected a transport error");
2265 };
2266 let Some(source) = source.downcast_ref::<reqwest::Error>() else {
2267 panic!("expected the reqwest transport source");
2268 };
2269 assert!(source.is_decode());
2270 assert!(
2271 !std::error::Error::source(&source)
2272 .is_some_and(|source| source.is::<serde_json::Error>())
2273 );
2274 Ok(())
2275 }
2276
2277 #[tokio::test]
2278 async fn response_body_transport_failures_are_retried()
2279 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2280 let truncated_responses = [
2281 "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n\
2282 Content-Length: 102\r\nConnection: close\r\n\r\n{}",
2283 "HTTP/1.1 401 Unauthorized\r\nContent-Type: text/plain\r\n\
2284 Content-Length: 100\r\nConnection: close\r\n\r\nbad",
2285 ];
2286
2287 for truncated in truncated_responses {
2288 let (base_url, server) =
2289 response_sequence_server(vec![truncated.to_string(), raw_chat_success_response()])?;
2290 let client = TranslatingLlmClient::new(&chat_map_with_retries(&base_url, 1))?;
2291 let response = client
2292 .call_rewrite_model(request_for(Some("gpt"), false), None)
2293 .await?;
2294 server
2295 .join()
2296 .map_err(|_| std::io::Error::other("response server thread panicked"))??;
2297
2298 assert_eq!(
2299 completion_text(&response.llm_response.into_agg().await?),
2300 "recovered"
2301 );
2302 }
2303 Ok(())
2304 }
2305
2306 #[tokio::test]
2307 async fn pre_first_event_transport_failure_uses_retry_budget()
2308 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2309 let truncated = "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\n\
2310 Content-Length: 100\r\nConnection: close\r\n\r\n";
2311 let body = "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"recovered\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n";
2312 let success = format!(
2313 "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\n\
2314 Content-Length: {}\r\nConnection: close\r\n\r\n{body}",
2315 body.len()
2316 );
2317
2318 for second in [success, truncated.to_string()] {
2319 let is_exhausted = second == truncated;
2320 let (base_url, server) = response_sequence_server(vec![truncated.to_string(), second])?;
2321 let client = TranslatingLlmClient::new(&chat_map_with_retries(&base_url, 1))?;
2322 let result = client
2323 .call_rewrite_model(request_for(Some("gpt"), true), None)
2324 .await;
2325 if is_exhausted {
2326 assert!(matches!(result, Err(LlmClientError::Transport { .. })));
2327 } else {
2328 assert_eq!(
2329 completion_text(&result?.llm_response.into_agg().await?),
2330 "recovered"
2331 );
2332 }
2333 server
2334 .join()
2335 .map_err(|_| std::io::Error::other("response server thread panicked"))??;
2336 }
2337 Ok(())
2338 }
2339
2340 #[tokio::test]
2341 async fn streaming_body_io_failure_preserves_transport_error()
2342 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2343 let body = "data: {\"choices\":[{\"delta\":{\"content\":\"partial\"}}]}\n\n";
2344 let (base_url, server) = truncated_response_server("text/event-stream", body)?;
2345 let client = TranslatingLlmClient::new(&chat_map_with_retries(&base_url, 2))?;
2346 let response = client
2347 .call_rewrite_model(request_for(Some("gpt"), true), None)
2348 .await?;
2349 let result = response.llm_response.into_agg().await;
2350 server
2351 .join()
2352 .map_err(|_| std::io::Error::other("response server thread panicked"))??;
2353
2354 let Err(error) = result else {
2355 panic!("expected the truncated stream body to fail");
2356 };
2357
2358 assert!(matches!(error, LlmClientError::Transport { .. }));
2359 Ok(())
2360 }
2361
2362 #[tokio::test]
2363 async fn rewrites_model_to_resolved_upstream_id()
2364 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2365 let server = MockServer::start().await;
2367 Mock::given(method("POST"))
2368 .and(path("/v1/chat/completions"))
2369 .and(wiremock::matchers::body_partial_json(json!({"model": "gpt"})))
2370 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
2371 "id": "1", "model": "gpt",
2372 "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
2373 "usage": {}
2374 })))
2375 .mount(&server)
2376 .await;
2377
2378 let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
2379 client
2381 .call_rewrite_model(
2382 request_for(Some("switchyard"), false),
2383 Some(&ModelId::from("gpt")),
2384 )
2385 .await?;
2386 Ok(())
2388 }
2389
2390 #[tokio::test]
2393 async fn reasoning_effort_override_replaces_the_callers_effort()
2394 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2395 let server = MockServer::start().await;
2396 Mock::given(method("POST"))
2397 .and(path("/v1/chat/completions"))
2398 .and(wiremock::matchers::body_partial_json(json!({
2399 "model": "gpt",
2400 "reasoning_effort": "max"
2401 })))
2402 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
2403 "id": "1",
2404 "model": "gpt",
2405 "choices": [{
2406 "index": 0,
2407 "message": {"role": "assistant", "content": "ok"},
2408 "finish_reason": "stop"
2409 }],
2410 "usage": {}
2411 })))
2412 .mount(&server)
2413 .await;
2414 let client = TranslatingLlmClient::new(&chat_map_with_effort(
2415 &format!("{}/v1", server.uri()),
2416 "max",
2417 ))?;
2418 client
2419 .call_rewrite_model_raw(
2420 json!({
2421 "model": "client-facing",
2422 "messages": [{"role": "user", "content": "hi"}],
2423 "reasoning_effort": "high"
2424 }),
2425 None,
2426 Some(&ModelId::from("gpt")),
2427 WireFormat::OpenAiChat,
2428 )
2429 .await?;
2430
2431 let server = MockServer::start().await;
2432 Mock::given(method("POST"))
2433 .and(path("/v1/responses"))
2434 .and(wiremock::matchers::body_partial_json(json!({
2435 "model": "gpt",
2436 "reasoning": {"effort": "max", "summary": "auto"}
2437 })))
2438 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
2439 "id": "resp_1",
2440 "object": "response",
2441 "model": "gpt",
2442 "status": "completed",
2443 "output": [{
2444 "type": "message",
2445 "id": "msg_1",
2446 "role": "assistant",
2447 "status": "completed",
2448 "content": [{"type": "output_text", "text": "ok"}]
2449 }],
2450 "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}
2451 })))
2452 .mount(&server)
2453 .await;
2454 let client = TranslatingLlmClient::new(&responses_map_with_effort(
2455 &format!("{}/v1", server.uri()),
2456 "max",
2457 ))?;
2458 client
2459 .call_rewrite_model_raw(
2460 json!({
2461 "model": "client-facing",
2462 "input": [{"role": "user", "content": "hi"}],
2463 "reasoning": {"effort": "high", "summary": "auto"}
2464 }),
2465 None,
2466 Some(&ModelId::from("gpt")),
2467 WireFormat::OpenAiResponses,
2468 )
2469 .await?;
2470 Ok(())
2471 }
2472
2473 #[tokio::test]
2475 async fn anthropic_requests_drop_tool_strict_false()
2476 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2477 let server = MockServer::start().await;
2478 Mock::given(method("POST"))
2479 .and(path("/v1/messages"))
2480 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
2481 "id": "msg_1",
2482 "type": "message",
2483 "role": "assistant",
2484 "model": "claude",
2485 "content": [{"type": "text", "text": "ok"}],
2486 "stop_reason": "end_turn",
2487 "usage": {"input_tokens": 1, "output_tokens": 1}
2488 })))
2489 .expect(2)
2490 .mount(&server)
2491 .await;
2492 let client = TranslatingLlmClient::new(&[ModelConfig::new(
2493 "claude",
2494 Backend::Anthropic(config(&server.uri())),
2495 None,
2496 )])?;
2497 let schema = json!({"type": "object"});
2498 let requests = [
2499 (
2500 WireFormat::OpenAiChat,
2501 json!({
2502 "model": "client-facing",
2503 "messages": [{"role": "user", "content": "hi"}],
2504 "tools": [
2505 {"type": "function", "function": {"name": "loose", "parameters": schema, "strict": false}},
2506 {"type": "function", "function": {"name": "exact", "parameters": schema, "strict": true}}
2507 ]
2508 }),
2509 ),
2510 (
2511 WireFormat::AnthropicMessages,
2512 json!({
2513 "model": "client-facing",
2514 "max_tokens": 8,
2515 "messages": [{"role": "user", "content": "hi"}],
2516 "tools": [
2517 {"name": "loose", "input_schema": schema, "strict": false},
2518 {"name": "exact", "input_schema": schema, "strict": true}
2519 ]
2520 }),
2521 ),
2522 ];
2523 for (format, request) in requests {
2524 client
2525 .call_rewrite_model_raw(request, None, Some(&ModelId::from("claude")), format)
2526 .await?;
2527 }
2528 for received in server.received_requests().await.unwrap_or_default() {
2529 let body: Value = serde_json::from_slice(&received.body)?;
2530 assert!(body["tools"][0].get("strict").is_none(), "{body}");
2531 assert_eq!(body["tools"][1]["strict"], true, "{body}");
2532 }
2533 Ok(())
2534 }
2535
2536 #[tokio::test]
2537 async fn template_rejection_retries_once_without_changing_other_fields()
2538 -> std::result::Result<(), Box<dyn Error + Sync + Send>> {
2539 for (message, reject_retry, omit) in [
2540 (
2541 "Unrecognized request argument supplied: chat_template_kwargs",
2542 false,
2543 false,
2544 ),
2545 (
2546 "Unrecognized request argument supplied: chat_template_kwargs",
2547 true,
2548 false,
2549 ),
2550 (
2551 "Invalid temperature for chat_template_kwargs request",
2552 false,
2553 false,
2554 ),
2555 (
2556 "Unrecognized request argument supplied: chat_template_kwargs",
2557 false,
2558 true,
2559 ),
2560 ] {
2561 let server = MockServer::start().await;
2562 Mock::given(method("POST"))
2563 .respond_with(move |request: &wiremock::Request| {
2564 let body: Value = serde_json::from_slice(&request.body).expect("request JSON");
2565 if body.get("chat_template_kwargs").is_some() || reject_retry {
2566 ResponseTemplate::new(400).set_body_json(json!({"error": {"message": message}}))
2567 } else {
2568 ResponseTemplate::new(200).set_body_json(json!({
2569 "id": "1", "model": "gpt", "choices": [{"index": 0,
2570 "message": {"role": "assistant", "content": "yes"}, "finish_reason": "stop"}],
2571 "usage": {}
2572 }))
2573 }
2574 }).mount(&server).await;
2575 let fields = if omit {
2576 BTreeSet::from(["chat_template_kwargs".into()])
2577 } else {
2578 BTreeSet::new()
2579 };
2580 let client = TranslatingLlmClient::new(&chat_map_with_omit_body_fields(
2581 &format!("{}/v1", server.uri()),
2582 fields,
2583 ))?;
2584 let raw = json!({"model": "gpt", "messages": [{"role": "user", "content": "judge"}],
2585 "reasoning_effort": "none", "logprobs": true, "top_logprobs": 8,
2586 "chat_template_kwargs": {"enable_thinking": false}});
2587 let result = client
2588 .call_rewrite_model_raw(
2589 raw.clone(),
2590 None,
2591 Some(&ModelId::from("gpt")),
2592 WireFormat::OpenAiChat,
2593 )
2594 .await;
2595 let rejects_field = message.starts_with("Unrecognized");
2596 assert_eq!(result.is_ok(), omit || (rejects_field && !reject_retry));
2597 let calls = server.received_requests().await.ok_or("missing requests")?;
2598 assert_eq!(calls.len(), if !omit && rejects_field { 2 } else { 1 });
2599 for (index, call) in calls.iter().enumerate() {
2600 let mut expected = raw.clone();
2601 if omit || index > 0 {
2602 expected
2603 .as_object_mut()
2604 .unwrap()
2605 .remove("chat_template_kwargs");
2606 }
2607 assert_eq!(serde_json::from_slice::<Value>(&call.body)?, expected);
2608 }
2609 }
2610 Ok(())
2611 }
2612
2613 #[tokio::test]
2614 async fn template_compatibility_preserves_the_transient_retry_budget()
2615 -> std::result::Result<(), Box<dyn Error + Sync + Send>> {
2616 let server = MockServer::start().await;
2617 let count = Arc::new(AtomicUsize::new(0));
2618 let observed = Arc::clone(&count);
2619 Mock::given(method("POST"))
2620 .respond_with(move |request: &wiremock::Request| {
2621 let body: Value = serde_json::from_slice(&request.body).expect("request JSON");
2622 match observed.fetch_add(1, Ordering::SeqCst) {
2623 0 => {
2624 assert!(body.get("chat_template_kwargs").is_some());
2625 ResponseTemplate::new(422).set_body_json(json!({"error": {
2626 "code": "unsupported_parameter", "param": "chat_template_kwargs"
2627 }}))
2628 }
2629 1 => {
2630 assert!(body.get("chat_template_kwargs").is_none());
2631 ResponseTemplate::new(503).insert_header("retry-after", "0")
2632 }
2633 _ => ResponseTemplate::new(200).set_body_json(json!({
2634 "id": "1", "model": "gpt", "choices": [{"index": 0,
2635 "message": {"role": "assistant", "content": "yes"}, "finish_reason": "stop"}],
2636 "usage": {}
2637 })),
2638 }
2639 }).mount(&server).await;
2640 let client =
2641 TranslatingLlmClient::new(&chat_map_with_retries(&format!("{}/v1", server.uri()), 1))?;
2642 client
2643 .call_rewrite_model_raw(
2644 json!({"model": "gpt", "messages": [{"role": "user", "content": "judge"}],
2645 "chat_template_kwargs": {"enable_thinking": false}}),
2646 None,
2647 Some(&ModelId::from("gpt")),
2648 WireFormat::OpenAiChat,
2649 )
2650 .await?;
2651 assert_eq!(count.load(Ordering::SeqCst), 3);
2652 Ok(())
2653 }
2654
2655 #[tokio::test]
2656 async fn preserved_codex_string_input_is_normalized_at_the_client_boundary()
2657 -> std::result::Result<(), Box<dyn Error + Sync + Send>> {
2658 for root in ["/backend-api/codex", "/v1"] {
2659 for input in [
2660 json!("hello"),
2661 json!([{"role": "user", "content": "hello"}]),
2662 ] {
2663 let server = MockServer::start().await;
2664 Mock::given(method("POST"))
2665 .and(path(format!("{root}/responses")))
2666 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
2667 "id": "resp_1", "object": "response", "model": "gpt", "status": "completed", "output": [], "usage": {}
2668 }))).expect(1).mount(&server).await;
2669 let client =
2670 TranslatingLlmClient::new(&responses_map(&format!("{}{root}", server.uri())))?;
2671 let raw = json!({"model": "client-name", "input": input,
2672 "instructions": "Be brief.", "store": false, "metadata": {"task": "judge"},
2673 "custom_field": {"keep": true}});
2674 client
2675 .call_rewrite_model_raw(
2676 raw.clone(),
2677 None,
2678 Some(&ModelId::from("gpt")),
2679 WireFormat::OpenAiResponses,
2680 )
2681 .await?;
2682 let mut expected = raw;
2683 expected["model"] = json!("gpt");
2684 if root == "/backend-api/codex" && input.is_string() {
2685 expected["input"] =
2686 json!([{"type": "message", "role": "user", "content": input}]);
2687 }
2688 let calls = server.received_requests().await.ok_or("missing requests")?;
2689 assert_eq!(serde_json::from_slice::<Value>(&calls[0].body)?, expected);
2690 }
2691 }
2692 Ok(())
2693 }
2694
2695 #[tokio::test]
2696 async fn token_counting_does_not_send_the_effort_override()
2697 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2698 let server = MockServer::start().await;
2699 Mock::given(method("POST"))
2700 .and(path("/v1/messages/count_tokens"))
2701 .respond_with(ResponseTemplate::new(200).set_body_json(json!({"input_tokens": 3})))
2702 .expect(1)
2703 .mount(&server)
2704 .await;
2705 let mut backend = config(&server.uri());
2706 backend.reasoning_effort = Some("max".to_string());
2707 let client = TranslatingLlmClient::new(&[ModelConfig::new(
2708 "claude",
2709 Backend::Anthropic(backend),
2710 None,
2711 )])?;
2712 let llm_request = switchyard_translation::decode_request(
2713 WireFormat::AnthropicMessages,
2714 &json!({"model": "claude", "messages": [{"role": "user", "content": "hi"}]}),
2715 )?;
2716 let request = Request {
2717 llm_request,
2718 raw_request: None,
2719 metadata: None,
2720 };
2721 client
2722 .call_auxiliary(
2723 &ModelId::from("claude"),
2724 request,
2725 AuxiliaryOperation::AnthropicCountTokens,
2726 )
2727 .await?;
2728 let received = server.received_requests().await.unwrap_or_default();
2729 let body: Value = serde_json::from_slice(&received[0].body)?;
2730 assert!(body.get("output_config").is_none(), "{body}");
2731 Ok(())
2732 }
2733
2734 #[tokio::test]
2735 async fn reasoning_effort_override_sets_anthropic_output_config_effort()
2736 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2737 let server = MockServer::start().await;
2738 let format = json!({"type": "json_schema", "schema": {"type": "object"}});
2739 Mock::given(method("POST"))
2740 .and(path("/v1/messages"))
2741 .and(wiremock::matchers::body_partial_json(json!({
2742 "model": "claude",
2743 "output_config": {"effort": "max", "format": format}
2744 })))
2745 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
2746 "id": "msg_1",
2747 "type": "message",
2748 "role": "assistant",
2749 "model": "claude",
2750 "content": [{"type": "text", "text": "ok"}],
2751 "stop_reason": "end_turn",
2752 "usage": {"input_tokens": 1, "output_tokens": 1}
2753 })))
2754 .expect(1)
2755 .mount(&server)
2756 .await;
2757 let mut backend = config(&server.uri());
2758 backend.reasoning_effort = Some("max".to_string());
2759 let client = TranslatingLlmClient::new(&[ModelConfig::new(
2760 "claude",
2761 Backend::Anthropic(backend),
2762 None,
2763 )])?;
2764 client
2765 .call_rewrite_model_raw(
2766 json!({
2767 "model": "client-facing",
2768 "max_tokens": 7,
2769 "messages": [{"role": "user", "content": "hi"}],
2770 "output_config": {"effort": "high", "format": format}
2771 }),
2772 None,
2773 Some(&ModelId::from("claude")),
2774 WireFormat::AnthropicMessages,
2775 )
2776 .await?;
2777 Ok(())
2778 }
2779
2780 #[tokio::test]
2781 async fn extra_body_adds_defaults_without_overriding_the_request()
2782 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2783 let server = MockServer::start().await;
2784 Mock::given(method("POST"))
2785 .and(path("/v1/chat/completions"))
2786 .and(wiremock::matchers::body_partial_json(json!({
2787 "model": "gpt",
2788 "max_tokens": 7,
2789 "service_tier": "priority"
2790 })))
2791 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
2792 "id": "1",
2793 "model": "gpt",
2794 "choices": [{
2795 "index": 0,
2796 "message": {"role": "assistant", "content": "ok"},
2797 "finish_reason": "stop"
2798 }],
2799 "usage": {}
2800 })))
2801 .mount(&server)
2802 .await;
2803
2804 let extra_body = BTreeMap::from([
2805 ("max_tokens".to_string(), json!(999)),
2806 ("service_tier".to_string(), json!("priority")),
2807 ]);
2808 let client = TranslatingLlmClient::new(&chat_map_with_extra_body(
2809 &format!("{}/v1", server.uri()),
2810 extra_body,
2811 ))?;
2812 let raw = json!({
2813 "model": "client-facing",
2814 "messages": [{"role": "user", "content": "hi"}],
2815 "max_tokens": 7
2816 });
2817
2818 client
2819 .call_rewrite_model_raw(
2820 raw,
2821 None,
2822 Some(&ModelId::from("gpt")),
2823 WireFormat::OpenAiChat,
2824 )
2825 .await?;
2826 Ok(())
2827 }
2828
2829 #[tokio::test]
2830 async fn omit_body_fields_strips_a_field_extra_body_cannot_override()
2831 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2832 let server = MockServer::start().await;
2833 Mock::given(method("POST"))
2834 .and(path("/v1/chat/completions"))
2835 .and(|request: &wiremock::Request| {
2836 let body: Value = serde_json::from_slice(&request.body).unwrap_or(Value::Null);
2837 body.get("max_tokens").is_none()
2838 })
2839 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
2840 "id": "1",
2841 "model": "gpt",
2842 "choices": [{
2843 "index": 0,
2844 "message": {"role": "assistant", "content": "ok"},
2845 "finish_reason": "stop"
2846 }],
2847 "usage": {}
2848 })))
2849 .mount(&server)
2850 .await;
2851
2852 let client = TranslatingLlmClient::new(&chat_map_with_omit_body_fields(
2853 &format!("{}/v1", server.uri()),
2854 BTreeSet::from(["max_tokens".to_string()]),
2855 ))?;
2856 let raw = json!({
2857 "model": "client-facing",
2858 "messages": [{"role": "user", "content": "hi"}],
2859 "max_tokens": 7
2860 });
2861
2862 client
2863 .call_rewrite_model_raw(
2864 raw,
2865 None,
2866 Some(&ModelId::from("gpt")),
2867 WireFormat::OpenAiChat,
2868 )
2869 .await?;
2870 Ok(())
2871 }
2872
2873 #[tokio::test]
2877 async fn anthropic_requests_drop_unsigned_thinking_blocks()
2878 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2879 let server = MockServer::start().await;
2880 Mock::given(method("POST"))
2881 .and(path("/v1/messages"))
2882 .and(|request: &wiremock::Request| {
2883 let body: Value = serde_json::from_slice(&request.body).unwrap_or(Value::Null);
2884 let messages = body.get("messages").and_then(Value::as_array).cloned();
2885 let Some(messages) = messages else {
2886 return false;
2887 };
2888 let blocks: Vec<&Value> = messages
2892 .iter()
2893 .filter_map(|message| message.get("content"))
2894 .filter_map(Value::as_array)
2895 .flatten()
2896 .collect();
2897 let thinking: Vec<&&Value> = blocks
2898 .iter()
2899 .filter(|block| block.get("type").and_then(Value::as_str) == Some("thinking"))
2900 .collect();
2901 thinking.len() == 1
2902 && thinking[0].get("signature").and_then(Value::as_str) == Some("sig-abc")
2903 && messages
2904 .iter()
2905 .all(|message| message.get("content") != Some(&json!([])))
2906 })
2907 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
2908 "id": "msg_1",
2909 "type": "message",
2910 "role": "assistant",
2911 "model": "claude",
2912 "content": [{"type": "text", "text": "ok"}],
2913 "stop_reason": "end_turn",
2914 "usage": {"input_tokens": 1, "output_tokens": 1}
2915 })))
2916 .mount(&server)
2917 .await;
2918
2919 let client = TranslatingLlmClient::new(&anthropic_map(&server.uri()))?;
2920 let raw = json!({
2921 "model": "client-facing",
2922 "max_tokens": 7,
2923 "messages": [
2924 {"role": "user", "content": "fix the build"},
2925 {"role": "assistant", "content": [
2926 {"type": "thinking", "thinking": "weak tier reasoning", "signature": ""}
2927 ]},
2928 {"role": "assistant", "content": [
2929 {"type": "thinking", "thinking": "signed reasoning", "signature": "sig-abc"},
2930 {"type": "text", "text": "here goes"}
2931 ]},
2932 {"role": "user", "content": "continue"}
2933 ]
2934 });
2935
2936 client
2937 .call_rewrite_model_raw(
2938 raw,
2939 None,
2940 Some(&ModelId::from("claude")),
2941 WireFormat::AnthropicMessages,
2942 )
2943 .await?;
2944 Ok(())
2945 }
2946
2947 #[tokio::test]
2948 async fn responses_reasoning_policy_preserves_replay_history()
2949 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
2950 use crate::ResponsesReasoningPolicy::{Auto, Drop, PreserveEncrypted, PreservePlaintext};
2951 let summary = json!([{"type": "summary_text", "text": "Check the parser."}]);
2952 let history = json!([
2953 {"type": "message", "role": "user", "content": "continue"},
2954 {"type": "reasoning", "encrypted_content": "opaque", "content": [{"type": "reasoning_text", "text": "private"}]},
2955 {"type": "reasoning", "id": "rs_stored", "status": "completed", "summary": []},
2956 {"type": "reasoning", "id": "rs_glm", "summary": summary, "encrypted_content": null},
2958 {"type": "reasoning", "content": [{"type": "reasoning_text", "text": "Run the tests."}], "encrypted_content": ""},
2959 {"type": "reasoning"},
2960 {"type": "function_call", "call_id": "call_1", "name": "shell", "arguments": "{}"},
2961 {"type": "function_call_output", "call_id": "call_1", "output": "ok"}
2962 ]);
2963 let cases = [
2965 (Auto, "openai/gpt-5", PreserveEncrypted),
2966 (Auto, "nvidia/zai-org/glm-5.3-flash", PreservePlaintext),
2967 (
2968 PreserveEncrypted,
2969 "nvidia/zai-org/glm-5.3-flash",
2970 PreserveEncrypted,
2971 ),
2972 (Drop, "openai/gpt-5", Drop),
2973 ];
2974 for store in [Some(true), Some(false), None] {
2975 for (policy, model, behavior) in cases {
2976 for has_input_override in [false, true] {
2977 let server = MockServer::start().await;
2978 Mock::given(method("POST"))
2979 .and(path("/v1/responses"))
2980 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
2981 "id": "resp_1", "object": "response", "model": "gpt",
2982 "status": "completed", "output": [],
2983 "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}
2984 })))
2985 .expect(1)
2986 .mount(&server)
2987 .await;
2988 let mut backend = config(&format!("{}/v1", server.uri()));
2989 if has_input_override {
2990 backend.omit_body_fields.insert("input".to_string());
2991 backend
2992 .extra_body
2993 .insert("input".to_string(), history.clone());
2994 }
2995 let client = TranslatingLlmClient::new(&[ModelConfig::new(
2996 model,
2997 Backend::OpenAiResponses(backend),
2998 None,
2999 )
3000 .with_responses_reasoning(policy)])?;
3001 let mut request = json!({"model": "switchyard", "input": history});
3002 if let Some(store) = store {
3003 request["store"] = json!(store);
3004 }
3005 client
3006 .call_rewrite_model_raw(
3007 request,
3008 None,
3009 Some(&ModelId::from(model)),
3010 WireFormat::OpenAiResponses,
3011 )
3012 .await?;
3013 let requests = server.received_requests().await.ok_or("missing requests")?;
3014 let body: Value = serde_json::from_slice(&requests[0].body)?;
3015 let stored = store != Some(false);
3016 let mut expected = vec![history[0].clone()];
3017 match behavior {
3018 PreserveEncrypted => {
3020 let indexes: &[usize] = if stored { &[1, 2, 3] } else { &[1] };
3021 for &index in indexes {
3022 let mut item = history[index].clone();
3023 item["content"] = json!([]);
3024 expected.push(item);
3025 }
3026 }
3027 PreservePlaintext => {
3029 let mut item = history[1].clone();
3030 item.as_object_mut()
3031 .ok_or("object")?
3032 .remove("encrypted_content");
3033 expected.push(item);
3034 if stored {
3035 expected.push(history[2].clone());
3036 }
3037 expected.extend([history[3].clone(), history[4].clone()]);
3038 }
3039 _ => {}
3040 }
3041 expected.extend([history[6].clone(), history[7].clone()]);
3042 assert_eq!(
3043 body["input"],
3044 json!(expected),
3045 "store={store:?}, policy={policy:?}, model={model}, replace={has_input_override}"
3046 );
3047 }
3048 }
3049 }
3050
3051 Ok(())
3052 }
3053
3054 #[tokio::test]
3055 async fn nvidia_anthropic_requests_strip_strict_from_extra_body()
3056 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3057 let server = MockServer::start().await;
3058 Mock::given(method("POST"))
3059 .and(path("/v1/messages"))
3060 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
3061 "id": "msg_1",
3062 "type": "message",
3063 "role": "assistant",
3064 "model": "claude",
3065 "content": [{"type": "text", "text": "ok"}],
3066 "stop_reason": "end_turn",
3067 "usage": {"input_tokens": 1, "output_tokens": 1}
3068 })))
3069 .expect(1)
3070 .mount(&server)
3071 .await;
3072
3073 let mut backend = config(&format!(
3074 "http://inference-api.nvidia.com:{}",
3075 server.address().port()
3076 ));
3077 backend.extra_body.insert(
3078 "tools".into(),
3079 json!([{
3080 "name": "lookup",
3081 "input_schema": {"type": "object", "properties": {}},
3082 "strict": true
3083 }]),
3084 );
3085 let mut client = TranslatingLlmClient::new(&[ModelConfig::new(
3086 "claude",
3087 Backend::Anthropic(backend),
3088 None,
3089 )])?;
3090 client.client = reqwest::Client::builder()
3091 .no_proxy()
3092 .resolve("inference-api.nvidia.com", *server.address())
3093 .build()?;
3094 client
3095 .call_rewrite_model_raw(
3096 json!({
3097 "model": "claude",
3098 "max_tokens": 7,
3099 "messages": [{"role": "user", "content": "hi"}]
3100 }),
3101 None,
3102 Some(&ModelId::from("claude")),
3103 WireFormat::AnthropicMessages,
3104 )
3105 .await?;
3106
3107 let requests = server.received_requests().await.unwrap();
3108 let body: Value = serde_json::from_slice(&requests[0].body)?;
3109 assert_eq!(body["tools"][0]["name"], "lookup");
3110 assert!(body["tools"][0].get("strict").is_none());
3111 Ok(())
3112 }
3113
3114 #[tokio::test]
3118 async fn anthropic_requests_drop_openai_only_fields()
3119 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3120 let server = MockServer::start().await;
3121 Mock::given(method("POST"))
3122 .and(path("/v1/messages"))
3123 .and(|request: &wiremock::Request| {
3124 let body: Value = serde_json::from_slice(&request.body).unwrap_or(Value::Null);
3125 body.get("context_management").is_none() && body.get("reasoning_effort").is_none()
3126 })
3127 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
3128 "id": "msg_1",
3129 "type": "message",
3130 "role": "assistant",
3131 "model": "claude",
3132 "content": [{"type": "text", "text": "ok"}],
3133 "stop_reason": "end_turn",
3134 "usage": {"input_tokens": 1, "output_tokens": 1}
3135 })))
3136 .mount(&server)
3137 .await;
3138
3139 let client = TranslatingLlmClient::new(&anthropic_map(&server.uri()))?;
3140 let raw = json!({
3141 "model": "client-facing",
3142 "messages": [{"role": "user", "content": "hi"}],
3143 "max_tokens": 7,
3144 "reasoning_effort": "high",
3145 "context_management": {
3146 "edits": [{"type": "clear_thinking_20251015"}]
3147 }
3148 });
3149
3150 client
3151 .call_rewrite_model_raw(
3152 raw,
3153 None,
3154 Some(&ModelId::from("claude")),
3155 WireFormat::AnthropicMessages,
3156 )
3157 .await?;
3158 Ok(())
3159 }
3160
3161 #[tokio::test]
3162 async fn streaming_openai_chat_aggregates()
3163 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3164 let server = MockServer::start().await;
3165 let sse = "data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}\n\n\
3166 data: {\"choices\":[{\"delta\":{\"content\":\" world\"}}]}\n\n\
3167 data: {\"choices\":[],\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":2,\"total_tokens\":3}}\n\n\
3168 data: [DONE]\n\n";
3169 Mock::given(method("POST"))
3170 .and(path("/v1/chat/completions"))
3171 .and(wiremock::matchers::body_partial_json(json!({
3172 "stream": true,
3173 "stream_options": {"include_usage": true}
3174 })))
3175 .respond_with(ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream"))
3176 .mount(&server)
3177 .await;
3178
3179 let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
3180
3181 let response = client
3182 .call_rewrite_model(request_for(Some("gpt"), true), None)
3183 .await?;
3184 assert!(matches!(response.llm_response, LlmResponse::Stream(_)));
3185 let agg = response.llm_response.into_agg().await?;
3186 assert_eq!(completion_text(&agg), "Hello world");
3187 assert_eq!(agg.usage.input_tokens, Some(1));
3188 assert_eq!(agg.usage.output_tokens, Some(2));
3189 assert_eq!(agg.usage.total_tokens, Some(3));
3190 Ok(())
3191 }
3192
3193 #[tokio::test]
3194 async fn streaming_openai_chat_preserves_usage_opt_out()
3195 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3196 let server = MockServer::start().await;
3197 Mock::given(method("POST"))
3198 .and(path("/v1/chat/completions"))
3199 .and(wiremock::matchers::body_partial_json(json!({
3200 "stream": true,
3201 "stream_options": {"include_usage": false}
3202 })))
3203 .respond_with(
3204 ResponseTemplate::new(200).set_body_raw("data: [DONE]\n\n", "text/event-stream"),
3205 )
3206 .mount(&server)
3207 .await;
3208
3209 let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
3210 let raw = json!({
3211 "model": "client-facing",
3212 "messages": [{"role": "user", "content": "hi"}],
3213 "stream": true,
3214 "stream_options": {"include_usage": false}
3215 });
3216
3217 let response = client
3218 .call_rewrite_model_raw(
3219 raw,
3220 None,
3221 Some(&ModelId::from("gpt")),
3222 WireFormat::OpenAiChat,
3223 )
3224 .await?;
3225 assert!(matches!(response, RawResponse::Stream(_)));
3226 Ok(())
3227 }
3228
3229 #[tokio::test]
3230 async fn upstream_500_is_upstream_http()
3231 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3232 let server = MockServer::start().await;
3233 Mock::given(method("POST"))
3234 .respond_with(ResponseTemplate::new(500).set_body_string("boom"))
3235 .mount(&server)
3236 .await;
3237
3238 let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
3239
3240 let Err(error) = client
3241 .call_rewrite_model(request_for(Some("gpt"), false), None)
3242 .await
3243 else {
3244 panic!("expected an error");
3245 };
3246 assert!(matches!(
3247 error,
3248 LlmClientError::UpstreamHttp {
3249 status: StatusCode::INTERNAL_SERVER_ERROR,
3250 ..
3251 }
3252 ));
3253 Ok(())
3254 }
3255
3256 #[tokio::test]
3257 async fn retryable_http_failure_recovers_within_budget()
3258 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3259 let server = MockServer::start().await;
3260 let calls = Arc::new(AtomicUsize::new(0));
3261 let observed_calls = Arc::clone(&calls);
3262 Mock::given(method("POST"))
3263 .respond_with(move |_: &wiremock::Request| {
3264 if observed_calls.fetch_add(1, Ordering::SeqCst) == 0 {
3265 ResponseTemplate::new(503)
3266 .insert_header("retry-after", "0")
3267 .set_body_string("temporarily unavailable")
3268 } else {
3269 chat_success_response()
3270 }
3271 })
3272 .mount(&server)
3273 .await;
3274
3275 let client =
3276 TranslatingLlmClient::new(&chat_map_with_retries(&format!("{}/v1", server.uri()), 1))?;
3277 let response = client
3278 .call_rewrite_model(request_for(Some("gpt"), false), None)
3279 .await?;
3280 let agg = response.llm_response.into_agg().await?;
3281
3282 assert_eq!(completion_text(&agg), "recovered");
3283 assert_eq!(calls.load(Ordering::SeqCst), 2);
3284 Ok(())
3285 }
3286
3287 #[tokio::test]
3288 async fn deterministic_http_failure_is_not_retried()
3289 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3290 let server = MockServer::start().await;
3291 let calls = Arc::new(AtomicUsize::new(0));
3292 let observed_calls = Arc::clone(&calls);
3293 Mock::given(method("POST"))
3294 .respond_with(move |_: &wiremock::Request| {
3295 observed_calls.fetch_add(1, Ordering::SeqCst);
3296 ResponseTemplate::new(401).set_body_string("invalid key")
3297 })
3298 .mount(&server)
3299 .await;
3300
3301 let client =
3302 TranslatingLlmClient::new(&chat_map_with_retries(&format!("{}/v1", server.uri()), 2))?;
3303 let Err(error) = client
3304 .call_rewrite_model(request_for(Some("gpt"), false), None)
3305 .await
3306 else {
3307 panic!("expected an upstream error");
3308 };
3309
3310 assert!(matches!(
3311 error,
3312 LlmClientError::UpstreamHttp {
3313 status: StatusCode::UNAUTHORIZED,
3314 body
3315 } if body == "invalid key"
3316 ));
3317 assert_eq!(calls.load(Ordering::SeqCst), 1);
3318 Ok(())
3319 }
3320
3321 #[tokio::test]
3322 async fn rate_limit_cooldown_depends_on_shared_credentials()
3323 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3324 for forward_auth in [true, false] {
3325 let server = MockServer::start().await;
3326 Mock::given(method("POST"))
3327 .respond_with(ResponseTemplate::new(429).insert_header("retry-after", "0"))
3328 .expect(if forward_auth { 6 } else { 3 })
3329 .mount(&server)
3330 .await;
3331 let backend = HttpBackendConfig {
3332 forward_auth,
3333 failure_cooldown: Duration::from_secs(60),
3334 ..config_with_retries(&format!("{}/v1", server.uri()), 2)
3335 };
3336 let client = TranslatingLlmClient::new(&[ModelConfig::new(
3337 "gpt",
3338 Backend::OpenAiChat(backend),
3339 None,
3340 )])?;
3341 for credential in ["Bearer first-user", "Bearer second-user"] {
3342 let mut headers = http::HeaderMap::new();
3343 headers.insert("authorization", http::HeaderValue::from_static(credential));
3344 let result = client
3345 .call_rewrite_model(request_with_headers("gpt", headers), None)
3346 .await;
3347 if !forward_auth && credential == "Bearer second-user" {
3348 assert!(matches!(
3349 result,
3350 Err(LlmClientError::TemporarilyUnavailable)
3351 ));
3352 } else {
3353 assert!(matches!(
3354 result,
3355 Err(LlmClientError::UpstreamHttp {
3356 status: StatusCode::TOO_MANY_REQUESTS,
3357 ..
3358 })
3359 ));
3360 }
3361 }
3362 }
3363 Ok(())
3364 }
3365
3366 #[tokio::test]
3367 async fn retry_exhaustion_returns_the_final_upstream_error()
3368 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3369 let server = MockServer::start().await;
3370 let calls = Arc::new(AtomicUsize::new(0));
3371 let observed_calls = Arc::clone(&calls);
3372 Mock::given(method("POST"))
3373 .respond_with(move |_: &wiremock::Request| {
3374 let attempt = observed_calls.fetch_add(1, Ordering::SeqCst) + 1;
3375 ResponseTemplate::new(500)
3376 .insert_header("retry-after", "0")
3377 .set_body_string(format!("attempt {attempt}"))
3378 })
3379 .mount(&server)
3380 .await;
3381
3382 let client =
3383 TranslatingLlmClient::new(&chat_map_with_retries(&format!("{}/v1", server.uri()), 2))?;
3384 let Err(error) = client
3385 .call_rewrite_model(request_for(Some("gpt"), false), None)
3386 .await
3387 else {
3388 panic!("expected retry exhaustion");
3389 };
3390
3391 assert!(matches!(
3392 error,
3393 LlmClientError::UpstreamHttp {
3394 status: StatusCode::INTERNAL_SERVER_ERROR,
3395 body
3396 } if body == "attempt 3"
3397 ));
3398 assert_eq!(calls.load(Ordering::SeqCst), 3);
3399 Ok(())
3400 }
3401
3402 #[tokio::test]
3403 async fn deadline_expires_before_send_or_stream_poll()
3404 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3405 let server = MockServer::start().await;
3406 Mock::given(method("POST"))
3407 .respond_with(ResponseTemplate::new(200).set_body_raw(
3408 "data: {\"id\":\"test\",\"model\":\"gpt\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"ok\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n",
3409 "text/event-stream",
3410 ))
3411 .mount(&server)
3412 .await;
3413 for timeout in [Duration::ZERO, Duration::from_millis(200)] {
3414 let mut backend = config(&server.uri());
3415 backend.timeout = Some(timeout);
3416 let client = TranslatingLlmClient::new(&[ModelConfig::new(
3417 "gpt",
3418 Backend::OpenAiChat(backend),
3419 None,
3420 )])?;
3421 let result = client
3422 .call_rewrite_model(request_for(Some("gpt"), true), None)
3423 .await;
3424 if timeout.is_zero() {
3425 assert!(matches!(result, Err(LlmClientError::Timeout { .. })));
3426 assert!(
3427 server
3428 .received_requests()
3429 .await
3430 .expect("recorded requests")
3431 .is_empty()
3432 );
3433 } else {
3434 let LlmResponse::Stream(mut chunks) = result?.llm_response else {
3435 return Err("expected a streaming response".into());
3436 };
3437 tokio::time::sleep(Duration::from_millis(300)).await;
3438 assert!(matches!(
3439 chunks.next().await,
3440 Some(Err(LlmClientError::Timeout { .. }))
3441 ));
3442 assert!(chunks.next().await.is_none());
3443 }
3444 }
3445 Ok(())
3446 }
3447
3448 #[tokio::test]
3449 async fn timeout_is_retried_before_a_response_is_returned()
3450 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3451 let server = MockServer::start().await;
3452 let calls = Arc::new(AtomicUsize::new(0));
3453 let observed_calls = Arc::clone(&calls);
3454 Mock::given(method("POST"))
3455 .respond_with(move |_: &wiremock::Request| {
3456 if observed_calls.fetch_add(1, Ordering::SeqCst) == 0 {
3457 ResponseTemplate::new(200).set_delay(Duration::from_millis(500))
3458 } else {
3459 chat_success_response()
3460 }
3461 })
3462 .mount(&server)
3463 .await;
3464
3465 let mut client =
3466 TranslatingLlmClient::new(&chat_map_with_retries(&format!("{}/v1", server.uri()), 1))?;
3467 client.client = reqwest::Client::builder()
3468 .timeout(Duration::from_millis(100))
3469 .build()?;
3470 let response = client
3471 .call_rewrite_model(request_for(Some("gpt"), false), None)
3472 .await?;
3473
3474 assert_eq!(
3475 completion_text(&response.llm_response.into_agg().await?),
3476 "recovered"
3477 );
3478 assert_eq!(calls.load(Ordering::SeqCst), 2);
3479 Ok(())
3480 }
3481
3482 #[test]
3483 fn retryable_error_classes_are_explicit() {
3484 let transport = AttemptFailure {
3485 error: LlmClientError::Transport {
3486 source: std::io::Error::other("disconnected").into(),
3487 },
3488 status: None,
3489 retry_after: None,
3490 };
3491 assert!(transport.is_retryable());
3492
3493 for status in [
3494 StatusCode::REQUEST_TIMEOUT,
3495 StatusCode::TOO_MANY_REQUESTS,
3496 StatusCode::INTERNAL_SERVER_ERROR,
3497 StatusCode::SERVICE_UNAVAILABLE,
3498 StatusCode::from_u16(599)
3499 .expect("599 is a syntactically valid HTTP code in the http crate"),
3500 ] {
3501 let failure = AttemptFailure {
3502 error: LlmClientError::UpstreamHttp {
3503 status,
3504 body: String::new(),
3505 },
3506 status: Some(status),
3507 retry_after: None,
3508 };
3509 assert!(failure.is_retryable(), "HTTP {status} should retry");
3510 }
3511 for status in [
3512 StatusCode::BAD_REQUEST,
3513 StatusCode::UNAUTHORIZED,
3514 StatusCode::CONFLICT,
3515 StatusCode::from_u16(600)
3516 .expect("600 is a syntactically valid HTTP code in the http crate"),
3517 ] {
3518 let failure = AttemptFailure {
3519 error: LlmClientError::UpstreamHttp {
3520 status,
3521 body: String::new(),
3522 },
3523 status: Some(status),
3524 retry_after: None,
3525 };
3526 assert!(!failure.is_retryable(), "HTTP {status} should fail fast");
3527 }
3528
3529 let configuration = AttemptFailure {
3530 error: LlmClientError::Configuration {
3531 message: "invalid header".to_string(),
3532 },
3533 status: None,
3534 retry_after: None,
3535 };
3536 assert!(!configuration.is_retryable());
3537
3538 let context_window = AttemptFailure {
3539 error: LlmClientError::ContextWindowExceeded {
3540 model: ModelId::from("gpt"),
3541 message: "too long".to_string(),
3542 },
3543 status: Some(StatusCode::BAD_REQUEST),
3544 retry_after: None,
3545 };
3546 assert!(!context_window.is_retryable());
3547 }
3548
3549 #[test]
3550 fn retry_after_supports_seconds_and_http_dates() {
3551 let mut headers = HeaderMap::new();
3552 headers.insert(RETRY_AFTER, reqwest::header::HeaderValue::from_static("3"));
3553 assert_eq!(retry_after_delay(&headers), Some(Duration::from_secs(3)));
3554
3555 let retry_at = SystemTime::now() + Duration::from_secs(2);
3556 let value = httpdate::fmt_http_date(retry_at);
3557 let Ok(value) = reqwest::header::HeaderValue::from_str(&value) else {
3558 panic!("formatted HTTP date should be a valid header");
3559 };
3560 headers.insert(RETRY_AFTER, value);
3561 let Some(delay) = retry_after_delay(&headers) else {
3562 panic!("HTTP date should produce a retry delay");
3563 };
3564 assert!(delay <= Duration::from_secs(2));
3565 }
3566
3567 #[tokio::test]
3568 async fn routed_llm_client_exposes_timeout_variant()
3569 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3570 let server = MockServer::start().await;
3571 Mock::given(method("POST"))
3572 .respond_with(
3573 ResponseTemplate::new(200).set_delay(std::time::Duration::from_millis(100)),
3574 )
3575 .mount(&server)
3576 .await;
3577
3578 let mut client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
3579 client.client = reqwest::Client::builder()
3580 .timeout(std::time::Duration::from_millis(10))
3581 .build()?;
3582
3583 let Err(error) = client.call(request_for(Some("gpt"), false)).await else {
3584 panic!("expected a timeout");
3585 };
3586 let LlmClientError::Timeout { source } = error else {
3587 panic!("expected the protocol timeout variant");
3588 };
3589 let Some(source) = source.downcast_ref::<reqwest::Error>() else {
3590 panic!("expected the reqwest timeout source");
3591 };
3592 assert!(source.is_timeout());
3593 Ok(())
3594 }
3595
3596 #[cfg(target_os = "linux")]
3598 #[tokio::test]
3599 async fn connect_timeout_is_a_transport_error()
3600 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3601 let socket = tokio::net::TcpSocket::new_v4()?;
3604 socket.bind("127.0.0.1:0".parse()?)?;
3605 let listener = socket.listen(0)?;
3606 let addr = listener.local_addr()?;
3607 let _queued: Vec<_> = (0..2)
3608 .map(|_| {
3609 std::net::TcpStream::connect_timeout(&addr, std::time::Duration::from_millis(100))
3610 })
3611 .collect();
3612
3613 let mut client = TranslatingLlmClient::new(&chat_map(&format!("http://{addr}/v1")))?;
3614 client.client = reqwest::Client::builder()
3617 .connect_timeout(std::time::Duration::from_millis(100))
3618 .build()?;
3619
3620 let Err(error) = client.call(request_for(Some("gpt"), false)).await else {
3621 panic!("expected the connect to fail");
3622 };
3623 let LlmClientError::Transport { source } = error else {
3624 panic!("expected Transport, got {error:?}");
3625 };
3626 let Some(source) = source.downcast_ref::<reqwest::Error>() else {
3627 panic!("expected the reqwest connect-timeout source");
3628 };
3629 assert!(source.is_connect());
3630 assert!(source.is_timeout());
3631 Ok(())
3632 }
3633
3634 #[tokio::test]
3635 async fn llama_cpp_context_overflow_400_is_mapped()
3636 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3637 let server = MockServer::start().await;
3638 Mock::given(method("POST"))
3639 .respond_with(ResponseTemplate::new(400).set_body_json(json!({
3640 "error": {
3641 "message": "request (6016 tokens) exceeds the available context size (4096 tokens), try increasing it"
3642 }
3643 })))
3644 .mount(&server)
3645 .await;
3646
3647 let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
3648
3649 let Err(error) = client
3650 .call_rewrite_model(request_for(Some("gpt"), false), None)
3651 .await
3652 else {
3653 panic!("expected an error");
3654 };
3655 assert!(matches!(
3656 error,
3657 LlmClientError::ContextWindowExceeded { model, .. } if model == "gpt"
3658 ));
3659 Ok(())
3660 }
3661
3662 #[tokio::test]
3663 async fn content_policy_400_preserves_upstream_error()
3664 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3665 let server = MockServer::start().await;
3666 let body = json!({
3667 "error": {
3668 "code": "content_policy_violation",
3669 "message": "request blocked by content policy",
3670 "type": "invalid_request_error"
3671 },
3672 "metadata": {"documentation_section": "context window"}
3673 });
3674 Mock::given(method("POST"))
3675 .respond_with(ResponseTemplate::new(400).set_body_json(&body))
3676 .expect(1)
3677 .mount(&server)
3678 .await;
3679 let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
3680 let Err(error) = client
3681 .call_rewrite_model(request_for(Some("gpt"), false), None)
3682 .await
3683 else {
3684 panic!("expected a policy denial");
3685 };
3686 let LlmClientError::UpstreamHttp {
3687 status,
3688 body: actual,
3689 } = error
3690 else {
3691 panic!("expected the upstream HTTP error, got {error:?}");
3692 };
3693 assert_eq!(status, StatusCode::BAD_REQUEST);
3694 assert_eq!(serde_json::from_str::<serde_json::Value>(&actual)?, body);
3695 Ok(())
3696 }
3697
3698 #[tokio::test]
3699 async fn forwards_only_allowlisted_metadata_headers()
3700 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3701 let server = MockServer::start().await;
3702 Mock::given(method("POST"))
3703 .and(wiremock::matchers::header("x-request-id", "abc"))
3704 .and(wiremock::matchers::header("authorization", "Bearer secret"))
3706 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
3707 "id": "1", "model": "gpt",
3708 "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
3709 "usage": {}
3710 })))
3711 .mount(&server)
3712 .await;
3713
3714 let mut headers = http::HeaderMap::new();
3715 headers.insert("x-request-id", http::HeaderValue::from_static("abc"));
3716 headers.insert(
3717 "authorization",
3718 http::HeaderValue::from_static("Bearer client-key"),
3719 );
3720 headers.insert(
3721 "x-goog-api-key",
3722 http::HeaderValue::from_static("client-google-key"),
3723 );
3724 headers.insert(
3725 "x-custom-internal-secret",
3726 http::HeaderValue::from_static("client-custom-key"),
3727 );
3728 let request = request_with_headers("gpt", headers);
3729
3730 let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
3731
3732 client.call_rewrite_model(request, None).await?;
3735 let received = server
3736 .received_requests()
3737 .await
3738 .ok_or("request recording should be enabled")?;
3739 let received = received.first().ok_or("expected one upstream request")?;
3740 assert!(!received.headers.contains_key("x-goog-api-key"));
3741 assert!(!received.headers.contains_key("x-custom-internal-secret"));
3742 Ok(())
3743 }
3744
3745 #[tokio::test]
3746 async fn forwards_application_headers_with_forward_auth()
3747 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3748 let server = MockServer::start().await;
3749 Mock::given(method("POST"))
3750 .and(wiremock::matchers::header(
3751 "authorization",
3752 "Bearer client-key",
3753 ))
3754 .and(wiremock::matchers::header(
3755 "x-goog-api-key",
3756 "client-google-key",
3757 ))
3758 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
3759 "id": "1", "model": "gpt",
3760 "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
3761 "usage": {}
3762 })))
3763 .mount(&server)
3764 .await;
3765
3766 let mut headers = http::HeaderMap::new();
3767 headers.insert(
3768 "authorization",
3769 http::HeaderValue::from_static("Bearer client-key"),
3770 );
3771 headers.insert(
3772 "x-goog-api-key",
3773 http::HeaderValue::from_static("client-google-key"),
3774 );
3775 headers.append("x-custom", http::HeaderValue::from_static("caller-first"));
3776 headers.append("x-custom", http::HeaderValue::from_static("caller-second"));
3777 headers.insert("host", http::HeaderValue::from_static("client.example"));
3778 headers.insert("content-type", http::HeaderValue::from_static("text/plain"));
3779 headers.insert("connection", http::HeaderValue::from_static("x-hop-by-hop"));
3780 headers.insert(
3781 "x-hop-by-hop",
3782 http::HeaderValue::from_static("client-only"),
3783 );
3784 let request = request_with_headers("gpt", headers);
3785 let mut backend_config = forwarding_config(&format!("{}/v1", server.uri()));
3786 backend_config.extra_headers = BTreeMap::from([
3787 ("X-Custom".into(), "configured".into()),
3788 ("X-Configured-Only".into(), "static".into()),
3789 ]);
3790 let client = TranslatingLlmClient::new(&[ModelConfig::new(
3791 "gpt",
3792 Backend::OpenAiChat(backend_config),
3793 None,
3794 )])?;
3795
3796 client.call_rewrite_model(request, None).await?;
3797 let received = server
3798 .received_requests()
3799 .await
3800 .ok_or("request recording should be enabled")?;
3801 let received = received.first().ok_or("expected one upstream request")?;
3802 assert_ne!(received.headers["host"], "client.example");
3803 assert_eq!(received.headers["content-type"], "application/json");
3804 assert!(!received.headers.contains_key("connection"));
3805 assert!(!received.headers.contains_key("x-hop-by-hop"));
3806 assert_eq!(received.headers.get_all("x-custom").iter().count(), 1);
3807 assert_eq!(received.headers["x-custom"], "configured");
3808 assert_eq!(received.headers["x-configured-only"], "static");
3809 Ok(())
3810 }
3811
3812 #[tokio::test]
3813 async fn forward_auth_preserves_anthropic_headers_and_redacts_errors()
3814 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3815 let server = MockServer::start().await;
3816 Mock::given(method("POST"))
3817 .and(path("/v1/messages"))
3818 .and(|request: &wiremock::Request| {
3819 request
3820 .headers
3821 .get("anthropic-beta")
3822 .and_then(|value| value.to_str().ok())
3823 == Some("oauth-2025-04-20")
3824 && request.headers.get_all("anthropic-version").iter().count() == 1
3825 && request
3826 .headers
3827 .get("anthropic-version")
3828 .and_then(|value| value.to_str().ok())
3829 == Some("2023-06-01")
3830 })
3831 .respond_with(ResponseTemplate::new(401).set_body_json(json!({
3832 "error": {"message": "rejected client-google-key"}
3833 })))
3834 .mount(&server)
3835 .await;
3836
3837 let mut headers = http::HeaderMap::new();
3838 headers.insert(
3839 "authorization",
3840 http::HeaderValue::from_static("Bearer client-key"),
3841 );
3842 headers.insert(
3843 "x-goog-api-key",
3844 http::HeaderValue::from_static("client-google-key"),
3845 );
3846 headers.insert(
3847 "anthropic-beta",
3848 http::HeaderValue::from_static("oauth-2025-04-20,prompt-caching-2024-07-31"),
3849 );
3850 headers.insert(
3851 "anthropic-version",
3852 http::HeaderValue::from_static("caller-version"),
3853 );
3854 let backend_config = forwarding_config(&server.uri());
3855 let client = TranslatingLlmClient::new(&[ModelConfig::new(
3856 "claude",
3857 Backend::Anthropic(backend_config),
3858 None,
3859 )])?;
3860 let raw = json!({
3861 "model": "claude",
3862 "max_tokens": 8,
3863 "messages": [{"role": "user", "content": "hello"}]
3864 });
3865
3866 let Err(LlmClientError::UpstreamHttp { body, .. }) = client
3867 .call_rewrite_model_raw(
3868 raw,
3869 Some(headers),
3870 Some(&ModelId::from("claude")),
3871 WireFormat::AnthropicMessages,
3872 )
3873 .await
3874 else {
3875 panic!("expected an upstream HTTP error");
3876 };
3877 assert_eq!(body, r#"{"error":{"message":"rejected [REDACTED]"}}"#);
3878 Ok(())
3879 }
3880
3881 #[tokio::test]
3884 async fn routed_llm_client_serves_the_request_model()
3885 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3886 let server = MockServer::start().await;
3887 Mock::given(method("POST"))
3888 .and(path("/v1/chat/completions"))
3889 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
3890 "id": "chatcmpl-1",
3891 "model": "gpt",
3892 "choices": [{
3893 "index": 0,
3894 "message": {"role": "assistant", "content": "routed hi"},
3895 "finish_reason": "stop"
3896 }],
3897 "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}
3898 })))
3899 .mount(&server)
3900 .await;
3901
3902 let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
3903 let response = client.call(request_for(Some("gpt"), false)).await?;
3904 let agg = response.llm_response.into_agg().await?;
3905 assert_eq!(completion_text(&agg), "routed hi");
3906 Ok(())
3907 }
3908
3909 #[tokio::test]
3910 async fn invalid_raw_request_is_a_request_translation_error()
3911 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3912 let client = TranslatingLlmClient::new(&[])?;
3913 let Err(error) = client
3914 .call_rewrite_model_raw(
3915 json!("invalid"),
3916 None,
3917 Some(&ModelId::from("gpt")),
3918 WireFormat::OpenAiChat,
3919 )
3920 .await
3921 else {
3922 panic!("expected request translation to fail");
3923 };
3924
3925 assert!(matches!(
3926 error,
3927 LlmClientError::RequestTranslation(message) if !message.is_empty()
3928 ));
3929 Ok(())
3930 }
3931
3932 #[tokio::test]
3935 async fn call_rewrite_model_raw_round_trips_buffered_json()
3936 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3937 let server = MockServer::start().await;
3938 Mock::given(method("POST"))
3939 .and(path("/v1/chat/completions"))
3940 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
3941 "id": "chatcmpl-1",
3942 "model": "gpt",
3943 "choices": [{
3944 "index": 0,
3945 "message": {"role": "assistant", "content": "Hi there"},
3946 "finish_reason": "stop"
3947 }],
3948 "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}
3949 })))
3950 .mount(&server)
3951 .await;
3952
3953 let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
3954 let raw = json!({
3955 "model": "client-facing",
3956 "messages": [{"role": "user", "content": "hi"}]
3957 });
3958 let RawResponse::Buffered(body) = client
3959 .call_rewrite_model_raw(
3960 raw,
3961 None,
3962 Some(&ModelId::from("gpt")),
3963 WireFormat::OpenAiChat,
3964 )
3965 .await?
3966 else {
3967 panic!("expected a buffered response");
3968 };
3969
3970 assert_eq!(body["choices"][0]["message"]["content"], "Hi there");
3971 assert_eq!(body["model"], "gpt");
3973 Ok(())
3974 }
3975
3976 #[tokio::test]
3979 async fn call_rewrite_model_raw_restores_codex_mcp_namespace()
3980 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
3981 let server = MockServer::start().await;
3982 Mock::given(method("POST"))
3983 .and(path("/v1/chat/completions"))
3984 .and(wiremock::matchers::body_partial_json(json!({
3985 "tools": [{
3986 "type": "function",
3987 "function": {"name": "mcp__open_websearch__search"}
3988 }]
3989 })))
3990 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
3991 "id": "chatcmpl-1",
3992 "model": "gpt",
3993 "choices": [{
3994 "index": 0,
3995 "message": {
3996 "role": "assistant",
3997 "content": null,
3998 "tool_calls": [{
3999 "id": "call_1",
4000 "type": "function",
4001 "function": {
4002 "name": "mcp__open_websearch__search",
4003 "arguments": "{\"q\":\"rust\"}"
4004 }
4005 }]
4006 },
4007 "finish_reason": "tool_calls"
4008 }],
4009 "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}
4010 })))
4011 .mount(&server)
4012 .await;
4013
4014 let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
4015 let raw = json!({
4016 "model": "client-facing",
4017 "input": "Search for Rust.",
4018 "tools": [{
4019 "type": "namespace",
4020 "name": "mcp__open_websearch",
4021 "tools": [{
4022 "type": "function",
4023 "name": "search",
4024 "description": "Search the web",
4025 "parameters": {"type": "object", "properties": {"q": {"type": "string"}}}
4026 }]
4027 }]
4028 });
4029
4030 let RawResponse::Buffered(body) = client
4031 .call_rewrite_model_raw(
4032 raw,
4033 None,
4034 Some(&ModelId::from("gpt")),
4035 WireFormat::OpenAiResponses,
4036 )
4037 .await?
4038 else {
4039 panic!("expected a buffered response");
4040 };
4041
4042 assert_eq!(body["output"][0]["type"], "function_call");
4043 assert_eq!(body["output"][0]["name"], "search");
4044 assert_eq!(body["output"][0]["namespace"], "mcp__open_websearch");
4045 assert_eq!(body["output"][0]["arguments"], "{\"q\": \"rust\"}");
4047 Ok(())
4048 }
4049
4050 #[tokio::test]
4051 async fn openai_responses_backend_flattens_codex_namespaces_before_upstream()
4052 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
4053 let server = MockServer::start().await;
4054 Mock::given(method("POST"))
4055 .and(path("/v1/responses"))
4056 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
4057 "id": "resp_1",
4058 "model": "gpt",
4059 "object": "response",
4060 "created_at": 0,
4061 "status": "completed",
4062 "output": [{
4063 "type": "function_call",
4064 "call_id": "call_1",
4065 "name": "mcp__open_websearch__search",
4066 "arguments": "{\"q\":\"rust\"}"
4067 }],
4068 "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}
4069 })))
4070 .mount(&server)
4071 .await;
4072
4073 let client = TranslatingLlmClient::new(&responses_map(&format!("{}/v1", server.uri())))?;
4074 let parameters = json!({"type": "object", "properties": {"q": {"type": "string"}}});
4075 let raw = json!({
4076 "model": "client-facing",
4077 "input": [{
4078 "type": "function_call",
4079 "call_id": "call_0",
4080 "name": "search",
4081 "namespace": "mcp__open_websearch",
4082 "arguments": "{}"
4083 }],
4084 "tool_choice": {
4085 "type": "function",
4086 "name": "search",
4087 "namespace": "mcp__open_websearch"
4088 },
4089 "tools": [{
4090 "type": "namespace",
4091 "name": "mcp__open_websearch",
4092 "tools": [{
4093 "type": "function",
4094 "name": "search",
4095 "parameters": parameters.clone()
4096 }]
4097 }]
4098 });
4099
4100 let RawResponse::Buffered(body) = client
4101 .call_rewrite_model_raw(
4102 raw,
4103 None,
4104 Some(&ModelId::from("gpt")),
4105 WireFormat::OpenAiResponses,
4106 )
4107 .await?
4108 else {
4109 panic!("expected a buffered response");
4110 };
4111
4112 let received = server
4113 .received_requests()
4114 .await
4115 .ok_or("request recording should be enabled")?;
4116 let request_body: Value = serde_json::from_slice(&received[0].body)?;
4117 assert_eq!(request_body["model"], "gpt");
4118 assert_eq!(
4119 request_body["tools"],
4120 json!([{
4121 "type": "function",
4122 "name": "mcp__open_websearch__search",
4123 "description": "",
4124 "parameters": parameters
4125 }])
4126 );
4127 assert_eq!(
4128 request_body["tool_choice"],
4129 json!({"type": "function", "name": "mcp__open_websearch__search"})
4130 );
4131 assert_eq!(
4132 request_body["input"][0]["name"],
4133 "mcp__open_websearch__search"
4134 );
4135 assert!(request_body["input"][0].get("namespace").is_none());
4136
4137 assert_eq!(body["output"][0]["type"], "function_call");
4138 assert_eq!(body["output"][0]["name"], "search");
4139 assert_eq!(body["output"][0]["namespace"], "mcp__open_websearch");
4140 Ok(())
4141 }
4142
4143 #[tokio::test]
4150 async fn openai_responses_backend_drops_reasoning_the_provider_cannot_read()
4151 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
4152 let server = MockServer::start().await;
4153 Mock::given(method("POST"))
4154 .and(path("/v1/responses"))
4155 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
4156 "id": "resp_1",
4157 "model": "gpt",
4158 "object": "response",
4159 "created_at": 0,
4160 "status": "completed",
4161 "output": [{
4162 "type": "message",
4163 "role": "assistant",
4164 "content": [{"type": "output_text", "text": "ok"}]
4165 }],
4166 "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}
4167 })))
4168 .mount(&server)
4169 .await;
4170
4171 let mut claude_reasoning = json!({
4172 "type": "reasoning",
4173 "id": "rs_b99feb9a-763d-4f7f-b6b2-2c0199b47ded",
4174 "summary": [{"type": "summary_text", "text": "Read the files first."}],
4175 "encrypted_content": null
4176 });
4177 let claude_reply = switchyard_translation::TranslationEngine::default()
4180 .translate_response(
4181 WireFormat::AnthropicMessages,
4182 WireFormat::OpenAiResponses,
4183 &json!({
4184 "id": "msg_1",
4185 "type": "message",
4186 "role": "assistant",
4187 "model": "claude",
4188 "content": [
4189 {"type": "thinking", "thinking": "Plan the edit.", "signature": "sig-abc"},
4190 {"type": "text", "text": "ok"}
4191 ],
4192 "stop_reason": "end_turn",
4193 "usage": {"input_tokens": 1, "output_tokens": 1}
4194 }),
4195 &switchyard_translation::TranslationPolicy::default(),
4196 )?
4197 .body;
4198 let switchyard_reasoning = claude_reply["output"]
4199 .as_array()
4200 .and_then(|output| output.iter().find(|item| item["type"] == "reasoning"))
4201 .cloned()
4202 .ok_or("translated Claude reply should contain a reasoning item")?;
4203 assert!(switchyard_translation::is_anthropic_thinking_envelope(
4204 switchyard_reasoning["encrypted_content"]
4205 .as_str()
4206 .unwrap_or_default()
4207 ));
4208 let mut gpt_reasoning = json!({
4209 "type": "reasoning",
4210 "id": "rs_gpt",
4211 "summary": [],
4212 "encrypted_content": "opaque-encrypted-reasoning"
4213 });
4214 let mut empty_reasoning =
4215 json!({"type": "reasoning", "id": "rs_empty", "summary": [], "encrypted_content": ""});
4216 let user = json!({"type": "message", "role": "user", "content": "fix the build"});
4217 let call = json!({"type": "function_call", "call_id": "call_1", "name": "exec", "arguments": "{}"});
4218 let output = json!({"type": "function_call_output", "call_id": "call_1", "output": "ok"});
4219 let input = json!([
4220 user,
4221 claude_reasoning,
4222 switchyard_reasoning,
4223 call,
4224 output,
4225 gpt_reasoning,
4226 empty_reasoning
4227 ]);
4228
4229 let models = ["openai/gpt-5", "nvidia/zai-org/glm-5.3-flash"].map(|model| {
4230 ModelConfig::new(
4231 model,
4232 Backend::OpenAiResponses(config(&format!("{}/v1", server.uri()))),
4233 None,
4234 )
4235 });
4236 let client = TranslatingLlmClient::new(&models)?;
4237 for model in ["openai/gpt-5", "nvidia/zai-org/glm-5.3-flash"] {
4238 for store in [json!(false), json!(true)] {
4239 let raw = json!({"model": "client-facing", "store": store, "input": input.clone()});
4240 client
4241 .call_rewrite_model_raw(
4242 raw,
4243 None,
4244 Some(&ModelId::from(model)),
4245 WireFormat::OpenAiResponses,
4246 )
4247 .await?;
4248 }
4249 }
4250
4251 let received = server
4252 .received_requests()
4253 .await
4254 .ok_or("request recording should be enabled")?;
4255 let sent: Vec<Value> = received
4256 .iter()
4257 .map(|request| serde_json::from_slice::<Value>(&request.body))
4258 .collect::<std::result::Result<_, _>>()?;
4259 assert_eq!(sent.len(), 4);
4260 for request in &sent {
4261 assert!(
4262 request["input"]
4263 .as_array()
4264 .ok_or("expected input array")?
4265 .iter()
4266 .all(|item| !item["encrypted_content"]
4267 .as_str()
4268 .is_some_and(switchyard_translation::is_anthropic_thinking_envelope)),
4269 "provider cannot read translated Anthropic reasoning"
4270 );
4271 }
4272
4273 let mut switchyard_plaintext = switchyard_reasoning.clone();
4275 switchyard_plaintext
4276 .as_object_mut()
4277 .ok_or("object")?
4278 .remove("encrypted_content");
4279 assert_eq!(
4280 sent[2]["input"],
4281 json!([user, claude_reasoning, switchyard_plaintext, call, output])
4282 );
4283 assert_eq!(
4284 sent[3]["input"],
4285 json!([
4286 user,
4287 claude_reasoning,
4288 switchyard_plaintext,
4289 call,
4290 output,
4291 empty_reasoning
4292 ])
4293 );
4294
4295 for item in [
4296 &mut claude_reasoning,
4297 &mut gpt_reasoning,
4298 &mut empty_reasoning,
4299 ] {
4300 item["content"] = json!([]);
4301 }
4302 assert_eq!(sent[0]["input"], json!([user, call, output, gpt_reasoning]));
4304 assert_eq!(
4307 sent[1]["input"],
4308 json!([
4309 user,
4310 claude_reasoning,
4311 call,
4312 output,
4313 gpt_reasoning,
4314 empty_reasoning
4315 ])
4316 );
4317 Ok(())
4318 }
4319
4320 #[tokio::test]
4323 async fn call_rewrite_model_raw_streams_wire_events()
4324 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
4325 use futures::TryStreamExt;
4326
4327 let server = MockServer::start().await;
4328 let sse = "data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}\n\n\
4329 data: {\"choices\":[{\"delta\":{\"content\":\" world\"}}]}\n\n\
4330 data: [DONE]\n\n";
4331 Mock::given(method("POST"))
4332 .and(path("/v1/chat/completions"))
4333 .respond_with(ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream"))
4334 .mount(&server)
4335 .await;
4336
4337 let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
4338 let raw = json!({
4339 "model": "client-facing",
4340 "messages": [{"role": "user", "content": "hi"}],
4341 "stream": true
4342 });
4343 let RawResponse::Stream(stream) = client
4344 .call_rewrite_model_raw(
4345 raw,
4346 None,
4347 Some(&ModelId::from("gpt")),
4348 WireFormat::OpenAiChat,
4349 )
4350 .await?
4351 else {
4352 panic!("expected a streamed response");
4353 };
4354
4355 let events: Vec<Value> = stream.try_collect().await?;
4356 assert!(!events.is_empty(), "expected at least one wire event");
4357 let content: String = events
4358 .iter()
4359 .filter_map(|event| event["choices"][0]["delta"]["content"].as_str())
4360 .collect();
4361 assert_eq!(content, "Hello world");
4362 assert!(events.iter().all(|event| event["model"] == "gpt"));
4365 Ok(())
4366 }
4367
4368 #[tokio::test]
4370 async fn call_rewrite_model_raw_forwards_headers()
4371 -> std::result::Result<(), Box<dyn Error + Sync + Send + 'static>> {
4372 let server = MockServer::start().await;
4373 Mock::given(method("POST"))
4374 .and(wiremock::matchers::header("x-request-id", "abc"))
4375 .and(wiremock::matchers::header("authorization", "Bearer secret"))
4377 .respond_with(ResponseTemplate::new(200).set_body_json(json!({
4378 "id": "1", "model": "gpt",
4379 "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
4380 "usage": {}
4381 })))
4382 .mount(&server)
4383 .await;
4384
4385 let mut headers = http::HeaderMap::new();
4386 headers.insert("x-request-id", http::HeaderValue::from_static("abc"));
4387 headers.insert(
4388 "authorization",
4389 http::HeaderValue::from_static("Bearer client-key"),
4390 );
4391
4392 let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri())))?;
4393 let raw = json!({"model": "gpt", "messages": [{"role": "user", "content": "hi"}]});
4394 client
4397 .call_rewrite_model_raw(
4398 raw,
4399 Some(headers),
4400 Some(&ModelId::from("gpt")),
4401 WireFormat::OpenAiChat,
4402 )
4403 .await?;
4404 Ok(())
4405 }
4406}