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