From 435f6e33accbb4a7f90503676d892a8d6a6ab975 Mon Sep 17 00:00:00 2001 From: Lars van der Zande Date: Fri, 28 Aug 2026 12:00:37 -0500 Subject: [PATCH] feat(server): forward upstream response headers The LLM client records the upstream HTTP response headers on both the buffered and the streaming capture path, and the server replays an allowlisted subset downstream: W3C tracing, x-request-id, Anthropic's request-id, openai-processing-ms, and the anthropic-ratelimit-, x-ratelimit-, and x-upstream- namespaces. Body description, hop-by-hop, cookie, and Switchyard-owned headers never forward, and a header the server writes itself always beats an upstream echo of the same name. Source-breaking for Rust callers: `switchyard_protocol::Response` gains a required `upstream_headers` field, so struct-literal constructions need `upstream_headers: Default::default()`. No wire or Python-surface change. Signed-off-by: Lars van der Zande --- CHANGELOG.md | 14 ++ crates/libsy-llm-client/src/client.rs | 25 ++- crates/libsy-llm-client/src/run.rs | 2 + .../libsy-llm-client/tests/observability.rs | 9 + .../src/algorithms/advisor_gate/tests.rs | 22 ++- .../libsy/src/algorithms/advisor_gate/turn.rs | 5 + crates/libsy/src/algorithms/escalation.rs | 3 + crates/libsy/src/algorithms/llm_class.rs | 2 + crates/libsy/src/algorithms/noop.rs | 1 + crates/libsy/src/algorithms/util/llm_judge.rs | 3 + crates/libsy/src/core/algorithm.rs | 3 + crates/libsy/src/core/testing.rs | 1 + crates/prefill-router/tests/unit/algorithm.rs | 1 + crates/protocol/src/envelope.rs | 7 +- .../src/runtime.rs | 4 + crates/switchyard-py/src/libsy_bindings.rs | 2 + crates/switchyard-runner/tests/route.rs | 2 + crates/switchyard-server/src/lib.rs | 32 +++- crates/switchyard-server/src/usage_metrics.rs | 3 + crates/switchyard-server/tests/server.rs | 165 ++++++++++++++++++ 20 files changed, 297 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 76884e283..dc4d0f685 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -69,6 +69,15 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). - **Hierarchical routing** — libsy adds hierarchical routing, with stages delegating to their own sub-router; a hierarchical stage router that carries its own judge is rejected. (#533) +- **Upstream response headers forwarded** — the LLM client records the + upstream HTTP response headers on both the buffered and the streaming path, + and `switchyard-server` replays an allowlisted subset to the downstream + client: W3C tracing (`traceparent`, `tracestate`, `baggage`), + `x-request-id`, Anthropic's `request-id`, `openai-processing-ms`, and the + `anthropic-ratelimit-`, `x-ratelimit-`, and `x-upstream-` namespaces. Body + description, hop-by-hop, cookie, and Switchyard-owned headers are never + forwarded, and a header Switchyard writes itself always beats an upstream + echo of the same name. (#571) ### Changed @@ -83,6 +92,11 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). - **LiteLLM integration replaced by a routing plugin** — the client integration becomes a routing plugin, and its example moves out of `experimental`. (#532) +- **`Response` gains a required `upstream_headers` field** *(source-breaking)* + — `switchyard_protocol::Response` carries the upstream HTTP headers, so Rust + callers that construct a `Response` with a struct literal must add + `upstream_headers: Default::default()`. Field access and every other use are + unaffected, and there is no wire or Python-surface change. (#571) ### Removed diff --git a/crates/libsy-llm-client/src/client.rs b/crates/libsy-llm-client/src/client.rs index 371fc586f..b2ff28934 100644 --- a/crates/libsy-llm-client/src/client.rs +++ b/crates/libsy-llm-client/src/client.rs @@ -358,6 +358,7 @@ impl TranslatingLlmClient { metrics::record_upstream_attempt(Some(status.as_u16())); return Ok(EncodedResponse::Streaming(response)); } + let upstream_headers = response.headers().clone(); let body = match response.bytes().await { Ok(body) => body, Err(error) => { @@ -373,6 +374,7 @@ impl TranslatingLlmClient { return Ok(EncodedResponse::Buffered { status: status.as_u16(), body: body.to_vec(), + upstream_headers, }); } @@ -459,10 +461,11 @@ impl TranslatingLlmClient { ) .await?; - let llm_response = match http_response { + let (llm_response, upstream_headers) = match http_response { EncodedResponse::Streaming(http_response) => { // Adapt the reqwest body stream to plain bytes; the SSE-decode itself is // transport-agnostic and lives in `switchyard-translation`. + let upstream_headers = http_response.headers().clone(); let bytes = http_response.bytes_stream().map(|chunk| { chunk.map(|bytes| bytes.to_vec()).map_err(|error| { if error.is_timeout() { @@ -481,7 +484,7 @@ impl TranslatingLlmClient { // error event on an HTTP 200. Classify the first event before returning // the stream: nothing has reached the caller yet, so an overflow can // still fail the call and let routing try the next candidate. - match chunks.next().await { + let llm_response = match chunks.next().await { None => LlmResponse::Stream(stream::empty().boxed()), Some(first) => { if let Some(message) = first_event_overflow(&first, backend) { @@ -492,21 +495,27 @@ impl TranslatingLlmClient { } LlmResponse::Stream(stream::once(ready(first)).chain(chunks).boxed()) } - } + }; + (llm_response, upstream_headers) } - EncodedResponse::Buffered { body, .. } => { + EncodedResponse::Buffered { + body, + upstream_headers, + .. + } => { let body = serde_json::from_slice::(&body).map_err(|error| { LlmClientError::ResponseTranslation(format!("invalid upstream JSON: {error}")) })?; let agg = decode_aggregated_response(&body, wire_format) .map_err(|error| LlmClientError::ResponseTranslation(error.to_string()))?; - LlmResponse::Agg(agg) + (LlmResponse::Agg(agg), upstream_headers) } }; Ok(Response { llm_response, metadata, + upstream_headers, }) } @@ -608,7 +617,11 @@ impl UpstreamEndpoint { } enum EncodedResponse { - Buffered { status: u16, body: Vec }, + Buffered { + status: u16, + body: Vec, + upstream_headers: HeaderMap, + }, Streaming(reqwest::Response), } diff --git a/crates/libsy-llm-client/src/run.rs b/crates/libsy-llm-client/src/run.rs index 8103a3bb6..1992aff8b 100644 --- a/crates/libsy-llm-client/src/run.rs +++ b/crates/libsy-llm-client/src/run.rs @@ -558,6 +558,7 @@ mod tests { Ok(Response { llm_response: LlmResponse::Agg(text_response(Some(model.to_string()), model)), metadata: None, + upstream_headers: http::HeaderMap::new(), }) } } @@ -573,6 +574,7 @@ mod tests { .boxed(), ), metadata: None, + upstream_headers: http::HeaderMap::new(), } } diff --git a/crates/libsy-llm-client/tests/observability.rs b/crates/libsy-llm-client/tests/observability.rs index 4d3921afa..d07d17147 100644 --- a/crates/libsy-llm-client/tests/observability.rs +++ b/crates/libsy-llm-client/tests/observability.rs @@ -380,6 +380,7 @@ impl RoutedLlmClient for AffinityFallbackClient { r#"{"crux":"bounded task","primary_rule":"SUP-1","capability_boundary":"supported","p_solve":0.9}"#, )), metadata: None, + upstream_headers: http::HeaderMap::new(), }); } if model == "affinity-fallback-weak" && !self.efficient_available.load(Ordering::Relaxed) { @@ -391,6 +392,7 @@ impl RoutedLlmClient for AffinityFallbackClient { Ok(Response { llm_response: LlmResponse::Agg(text_response(Some(model.to_string()), "answer")), metadata: None, + upstream_headers: http::HeaderMap::new(), }) } } @@ -409,6 +411,7 @@ impl RoutedLlmClient for ClassifierClient { Ok(Response { llm_response: LlmResponse::Agg(text_response(Some(model_id.to_string()), completion)), metadata: None, + upstream_headers: http::HeaderMap::new(), }) } } @@ -438,6 +441,7 @@ impl RoutedLlmClient for JudgeClient { "routed response", )), metadata: None, + upstream_headers: http::HeaderMap::new(), }); } match &self.outcome { @@ -448,6 +452,7 @@ impl RoutedLlmClient for JudgeClient { JudgeOutcome::Reply(text) => Ok(Response { llm_response: LlmResponse::Agg(text_response(None, *text)), metadata: None, + upstream_headers: http::HeaderMap::new(), }), JudgeOutcome::StreamDecodeFailure => Ok(Response { llm_response: LlmResponse::Stream( @@ -459,6 +464,7 @@ impl RoutedLlmClient for JudgeClient { .boxed(), ), metadata: None, + upstream_headers: http::HeaderMap::new(), }), } } @@ -480,6 +486,7 @@ impl RoutedLlmClient for UsageClient { Ok(Response { llm_response: LlmResponse::Agg(response), metadata: None, + upstream_headers: http::HeaderMap::new(), }) } } @@ -1104,6 +1111,7 @@ impl RoutedLlmClient for StreamingUsageClient { Ok(Response { llm_response: LlmResponse::Stream(Box::pin(futures::stream::iter(chunks))), metadata: None, + upstream_headers: http::HeaderMap::new(), }) } } @@ -1539,6 +1547,7 @@ async fn in_flight_gauge_reads_a_run_parked_on_an_unanswered_routing_call() call.respond(Ok(Response { llm_response: LlmResponse::Agg(text_response(Some(MODEL.to_string()), "answer")), metadata: None, + upstream_headers: http::HeaderMap::new(), }))?; while stream.next().await.is_some() {} diff --git a/crates/libsy/src/algorithms/advisor_gate/tests.rs b/crates/libsy/src/algorithms/advisor_gate/tests.rs index 464396f0e..0fb0150a3 100644 --- a/crates/libsy/src/algorithms/advisor_gate/tests.rs +++ b/crates/libsy/src/algorithms/advisor_gate/tests.rs @@ -77,6 +77,7 @@ fn tool_call_turn() -> Response { ..AggLlmResponse::default() }), metadata: None, + upstream_headers: http::HeaderMap::new(), } } @@ -93,6 +94,7 @@ fn tool_use_stop_turn() -> Response { ..AggLlmResponse::default() }), metadata: None, + upstream_headers: http::HeaderMap::new(), } } @@ -111,6 +113,7 @@ fn reasoning_only_turn() -> Response { ..AggLlmResponse::default() }), metadata: None, + upstream_headers: http::HeaderMap::new(), } } @@ -125,6 +128,7 @@ fn empty_turn() -> Response { ..AggLlmResponse::default() }), metadata: None, + upstream_headers: http::HeaderMap::new(), } } @@ -134,6 +138,7 @@ fn streamed(events: Vec) -> Response { events.into_iter().map(Ok), ))), metadata: None, + upstream_headers: http::HeaderMap::new(), } } @@ -269,7 +274,14 @@ async fn tool_call_turn_replays_without_review() { async fn approved_terminal_turn_returns_buffered_body() { let script = Script::new(); let gate = gate(AdvisorGateConfig::default()); - let serve = script.serve("APPROVE", |_| reply("all done")); + let serve = script.serve("APPROVE", |_| { + let mut response = reply("all done"); + response.upstream_headers.insert( + "x-upstream-trace", + "trace-123".parse().expect("valid test header"), + ); + response + }); let (selected_model, response) = test_drive(gate, task_request(), serve) .await .expect("routes"); @@ -277,6 +289,13 @@ async fn approved_terminal_turn_returns_buffered_body() { script.models(), vec![EXECUTOR.to_string(), ADVISOR.to_string()] ); + assert_eq!( + response + .upstream_headers + .get("x-upstream-trace") + .and_then(|value| value.to_str().ok()), + Some("trace-123") + ); assert_eq!(completion_text(&agg_of(response).await), "all done"); assert_eq!(selected_model, EXECUTOR); } @@ -752,6 +771,7 @@ async fn pattern_trigger_matches_on_tool_call_turns() { ..AggLlmResponse::default() }), metadata: None, + upstream_headers: http::HeaderMap::new(), }; let serve = script.serve("APPROVE", { let turn = parking_lot::Mutex::new(Some(turn)); diff --git a/crates/libsy/src/algorithms/advisor_gate/turn.rs b/crates/libsy/src/algorithms/advisor_gate/turn.rs index 18dc73c58..62964caea 100644 --- a/crates/libsy/src/algorithms/advisor_gate/turn.rs +++ b/crates/libsy/src/algorithms/advisor_gate/turn.rs @@ -25,6 +25,7 @@ pub(super) struct GatedTurn { /// response, its own preservation intact. pub(super) agg: AggLlmResponse, pub(super) metadata: Option, + pub(super) upstream_headers: http::HeaderMap, } impl GatedTurn { @@ -40,6 +41,7 @@ impl GatedTurn { Response { llm_response, metadata: self.metadata, + upstream_headers: self.upstream_headers, } } } @@ -50,11 +52,13 @@ impl GatedTurn { /// turn fails whole. pub(super) async fn buffer_turn(executor: &str, response: Response) -> Result { let metadata = response.metadata; + let upstream_headers = response.upstream_headers; match response.llm_response { LlmResponse::Agg(agg) => Ok(GatedTurn { events: None, agg, metadata, + upstream_headers, }), LlmResponse::Stream(mut stream) => { let mut events = Vec::new(); @@ -88,6 +92,7 @@ pub(super) async fn buffer_turn(executor: &str, response: Response) -> Result for EscalationClassifier { LlmResponse::Agg(agg) }, metadata: efficient_response.metadata, + upstream_headers: efficient_response.upstream_headers, }; let (classification, _) = self @@ -211,6 +212,7 @@ mod tests { Ok(Response { llm_response: LlmResponse::Agg(text_response(None, queue.take())), metadata: request.metadata, + upstream_headers: http::HeaderMap::new(), }) } } @@ -246,6 +248,7 @@ mod tests { Err(error), ]))), metadata: None, + upstream_headers: http::HeaderMap::new(), } } diff --git a/crates/libsy/src/algorithms/llm_class.rs b/crates/libsy/src/algorithms/llm_class.rs index 599e9f9e0..bbf506ec2 100644 --- a/crates/libsy/src/algorithms/llm_class.rs +++ b/crates/libsy/src/algorithms/llm_class.rs @@ -929,6 +929,7 @@ mod tests { Ok(Response { llm_response: LlmResponse::Agg(text_response(None, completion)), metadata: request.metadata, + upstream_headers: http::HeaderMap::new(), }) } } @@ -947,6 +948,7 @@ mod tests { Ok(Response { llm_response: LlmResponse::Agg(text_response(None, format!("answer from {model}"))), metadata: request.metadata, + upstream_headers: http::HeaderMap::new(), }) } } diff --git a/crates/libsy/src/algorithms/noop.rs b/crates/libsy/src/algorithms/noop.rs index 8a9e42c47..8ea0defea 100644 --- a/crates/libsy/src/algorithms/noop.rs +++ b/crates/libsy/src/algorithms/noop.rs @@ -48,6 +48,7 @@ impl Algorithm for Noop { let response = Response { llm_response, metadata: request.metadata.clone(), + upstream_headers: http::HeaderMap::new(), }; Ok(RoutingOutcome::answered(model_id, request, response)) } diff --git a/crates/libsy/src/algorithms/util/llm_judge.rs b/crates/libsy/src/algorithms/util/llm_judge.rs index 0210bda27..70ef24bcf 100644 --- a/crates/libsy/src/algorithms/util/llm_judge.rs +++ b/crates/libsy/src/algorithms/util/llm_judge.rs @@ -492,6 +492,7 @@ mod tests { Response { llm_response: LlmResponse::Agg(text_response(None, completion)), metadata: None, + upstream_headers: http::HeaderMap::new(), } } @@ -501,6 +502,7 @@ mod tests { futures::stream::iter(chunks.into_iter().map(|chunk| Ok(chunk.into()))).boxed(), ), metadata: None, + upstream_headers: http::HeaderMap::new(), } } @@ -514,6 +516,7 @@ mod tests { Response { llm_response: LlmResponse::Stream(items.boxed()), metadata: None, + upstream_headers: http::HeaderMap::new(), } } diff --git a/crates/libsy/src/core/algorithm.rs b/crates/libsy/src/core/algorithm.rs index 9c89a4553..d2d0c9a44 100644 --- a/crates/libsy/src/core/algorithm.rs +++ b/crates/libsy/src/core/algorithm.rs @@ -493,6 +493,7 @@ mod tests { Response { llm_response: LlmResponse::Agg(text_response(None, "existing")), metadata: None, + upstream_headers: http::HeaderMap::new(), }, ); @@ -630,6 +631,7 @@ mod tests { Ok(Response { llm_response: LlmResponse::Stream(stream), metadata: None, + upstream_headers: http::HeaderMap::new(), }) } }; @@ -714,6 +716,7 @@ mod tests { "fulfilled".to_string(), )), metadata: None, + upstream_headers: http::HeaderMap::new(), }))?; } Step::Done(outcome) => { diff --git a/crates/libsy/src/core/testing.rs b/crates/libsy/src/core/testing.rs index aeb570988..7d7e9e17e 100644 --- a/crates/libsy/src/core/testing.rs +++ b/crates/libsy/src/core/testing.rs @@ -87,5 +87,6 @@ pub(crate) fn reply(completion: impl Into) -> Response { Response { llm_response: LlmResponse::Agg(text_response(None, completion.into())), metadata: None, + upstream_headers: http::HeaderMap::new(), } } diff --git a/crates/prefill-router/tests/unit/algorithm.rs b/crates/prefill-router/tests/unit/algorithm.rs index 9e71759f2..29004b9e3 100644 --- a/crates/prefill-router/tests/unit/algorithm.rs +++ b/crates/prefill-router/tests/unit/algorithm.rs @@ -72,6 +72,7 @@ async fn selected(route: Arc, request: Request) -> libsy::Result< switchyard_protocol::text_response(None, "unused"), ), metadata: None, + upstream_headers: Default::default(), })) }) .await?; diff --git a/crates/protocol/src/envelope.rs b/crates/protocol/src/envelope.rs index 5c50e3341..7347ee837 100644 --- a/crates/protocol/src/envelope.rs +++ b/crates/protocol/src/envelope.rs @@ -3,8 +3,8 @@ //! The request/response envelope: the normalized [`LlmRequest`]/[`LlmResponse`] paired //! with the original provider payload and correlation [`Metadata`]. - use crate::{LlmRequest, LlmResponse, Metadata, ModelId}; +use http::HeaderMap; /// A request an algorithm routes: the normalized [`LlmRequest`] plus optional /// host-owned raw data and correlation [`Metadata`]. @@ -39,6 +39,10 @@ pub struct Response { pub llm_response: LlmResponse, /// Correlation metadata carried through the response. pub metadata: Option, + /// Upstream HTTP response headers preserved from the LLM backend (or proxy). + /// Populated by the LLM client; consumers (e.g. switchyard-server) may forward + /// these to the downstream client for observability. + pub upstream_headers: HeaderMap, } impl Response { @@ -74,6 +78,7 @@ mod tests { let mut response = Response { llm_response: LlmResponse::Agg(text_response(None, "answer")), metadata: None, + upstream_headers: HeaderMap::new(), }; assert_eq!(response.served_model(), None); diff --git a/crates/switchyard-nemo-relay-plugin/src/runtime.rs b/crates/switchyard-nemo-relay-plugin/src/runtime.rs index 7b5a12fb9..974d530b7 100644 --- a/crates/switchyard-nemo-relay-plugin/src/runtime.rs +++ b/crates/switchyard-nemo-relay-plugin/src/runtime.rs @@ -1189,6 +1189,7 @@ mod tests { served_model: Some(ModelId::from("selected-target")), ..Default::default() }), + upstream_headers: http::HeaderMap::new(), }; let captured = Arc::new(Mutex::new(Vec::new())); let emitted = Arc::clone(&captured); @@ -1241,6 +1242,7 @@ mod tests { served_model: Some(ModelId::from("selected-target")), ..Default::default() }), + upstream_headers: http::HeaderMap::new(), }; let captured = Arc::new(Mutex::new(Vec::new())); let emitted = Arc::clone(&captured); @@ -1275,6 +1277,7 @@ mod tests { served_model: Some(ModelId::from("strong")), ..Default::default() }), + upstream_headers: http::HeaderMap::new(), }; let captured = Arc::new(Mutex::new(Vec::new())); let emitted = Arc::clone(&captured); @@ -1333,6 +1336,7 @@ mod tests { served_model: Some(ModelId::from("selected-target")), ..Default::default() }), + upstream_headers: http::HeaderMap::new(), }; let captured = Arc::new(Mutex::new(Vec::new())); let emitted = Arc::clone(&captured); diff --git a/crates/switchyard-py/src/libsy_bindings.rs b/crates/switchyard-py/src/libsy_bindings.rs index 0571a3e12..4147ac5c1 100644 --- a/crates/switchyard-py/src/libsy_bindings.rs +++ b/crates/switchyard-py/src/libsy_bindings.rs @@ -7,6 +7,7 @@ use std::collections::HashMap; use std::sync::Arc; use futures::StreamExt; +use http::HeaderMap; use http::header::{HeaderName, HeaderValue}; use pyo3::exceptions::{PyBaseException, PyStopAsyncIteration, PyTypeError, PyValueError}; use pyo3::prelude::*; @@ -488,6 +489,7 @@ impl PyModelCall { call.respond(Ok(Response { llm_response, metadata, + upstream_headers: HeaderMap::new(), })) .map_err(py_libsy_error) } diff --git a/crates/switchyard-runner/tests/route.rs b/crates/switchyard-runner/tests/route.rs index df03c147d..70abbf1c8 100644 --- a/crates/switchyard-runner/tests/route.rs +++ b/crates/switchyard-runner/tests/route.rs @@ -25,6 +25,7 @@ impl RoutedLlmClient for StubClient { "plugin response", )), metadata: None, + upstream_headers: Default::default(), }) } } @@ -118,6 +119,7 @@ impl RoutedLlmClient for LazyStreamClient { Ok(Response { llm_response: LlmResponse::Stream(stream), metadata: None, + upstream_headers: Default::default(), }) } } diff --git a/crates/switchyard-server/src/lib.rs b/crates/switchyard-server/src/lib.rs index 8b6f06de8..66fefaabc 100644 --- a/crates/switchyard-server/src/lib.rs +++ b/crates/switchyard-server/src/lib.rs @@ -63,7 +63,27 @@ pub const DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(30); pub const DEFAULT_MAX_REQUEST_BODY_BYTES: usize = 32 * 1024 * 1024; const HEADER_SELECTED_MODEL: &str = "x-model-router-selected-model"; +const FORWARDED_UPSTREAM_HEADERS: &[&str] = &[ + "baggage", + "openai-processing-ms", + // Anthropic spells its correlation id without the `x-` prefix. + "request-id", + "traceparent", + "tracestate", + "x-request-id", +]; +const FORWARDED_UPSTREAM_HEADER_PREFIXES: &[&str] = + &["anthropic-ratelimit-", "x-ratelimit-", "x-upstream-"]; const MAX_ROUTING_HEADER_VALUE_LEN: usize = 512; + +/// Whether an upstream header is safe and useful to expose downstream. +fn should_forward_upstream_header(name: &HeaderName) -> bool { + let name = name.as_str(); + FORWARDED_UPSTREAM_HEADERS.contains(&name) + || FORWARDED_UPSTREAM_HEADER_PREFIXES + .iter() + .any(|prefix| name.starts_with(prefix)) +} /// Non-standard status used only in logs and metrics for a request whose /// downstream client disconnected before any response was written. const CLIENT_CLOSED_REQUEST: u16 = 499; @@ -1042,7 +1062,7 @@ async fn handle_llm_request( // The response carries the candidate that actually served it. Fall back to the routing // selection for algorithms that return a response without an offloaded model call. let served_model = response.served_model().cloned().or(Some(selected_model)); - let response = if let Some(served_model) = served_model.as_ref() { + let mut response = if let Some(served_model) = served_model.as_ref() { let cache_eligible = cache_probe .as_ref() .map(|probe| state.stats.prefix_eligibility(served_model, probe)) @@ -1059,12 +1079,22 @@ async fn handle_llm_request( response }; + let upstream_headers = std::mem::take(&mut response.upstream_headers); let response_model = served_model.as_ref().map(ToString::to_string); let mut response = match into_http_response(response, wire_format, response_model, request_extensions) { Ok(response) => response, Err(error) => return server_error(error.to_string()), }; + // Forward upstream headers before Switchyard writes its own so any header + // this server emits always overrides an upstream echo of the same name. + let response_headers = response.headers_mut(); + for (name, value) in upstream_headers.iter() { + if !should_forward_upstream_header(name) { + continue; + } + response_headers.append(name.clone(), value.clone()); + } if let Some(served_model) = served_model.as_ref() { attach_routing_headers(&mut response, served_model.as_str()); } diff --git a/crates/switchyard-server/src/usage_metrics.rs b/crates/switchyard-server/src/usage_metrics.rs index 754df811e..bb4e21ad8 100644 --- a/crates/switchyard-server/src/usage_metrics.rs +++ b/crates/switchyard-server/src/usage_metrics.rs @@ -25,6 +25,7 @@ pub(crate) fn observe( let Response { llm_response, metadata, + upstream_headers, } = response; let model = model.to_string(); @@ -99,6 +100,7 @@ pub(crate) fn observe( Response { llm_response, metadata, + upstream_headers, } } @@ -213,6 +215,7 @@ mod tests { let response = Response { llm_response: LlmResponse::Stream(Box::pin(source)), metadata: None, + upstream_headers: http::HeaderMap::new(), }; let stats = StatsAccumulator::default(); let observed = observe( diff --git a/crates/switchyard-server/tests/server.rs b/crates/switchyard-server/tests/server.rs index 43e6abf7f..04189bfbe 100644 --- a/crates/switchyard-server/tests/server.rs +++ b/crates/switchyard-server/tests/server.rs @@ -156,6 +156,69 @@ async fn upstream_chat( ) .into_response(); } + if prompt == "upstream-headers" { + // Both the buffered and the streamed reply echo the same set, so the two + // capture paths are compared against one expectation. + const UPSTREAM_HEADER_ECHO: [(&str, &str); 5] = [ + ("x-upstream-trace", "trace-123"), + ("x-request-id", "req-42"), + ("request-id", "req_anthropic_42"), + ("x-model-router-selected-model", "model/upstream-echo"), + ("x-switchyard-session-id", "spoofed-by-upstream"), + ]; + // Streaming captures the headers off the response head, before any body + // arrives, so the streamed variant exercises a different capture branch. + let mut response = if body["stream"].as_bool() == Some(true) { + let events = [ + json!({"id": "chatcmpl-headers", "model": model, "choices": [{"index": 0, "delta": {"role": "assistant"}}]}).to_string(), + json!({"id": "chatcmpl-headers", "model": model, "choices": [{"index": 0, "delta": {"content": "ok"}}]}).to_string(), + json!({"id": "chatcmpl-headers", "model": model, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12}}).to_string(), + "[DONE]".to_string(), + ]; + let stream = futures_util::stream::iter( + events + .into_iter() + .map(|data| Ok::(Event::default().data(data))), + ); + let mut response = Sse::new(stream).into_response(); + let headers = response.headers_mut(); + for (name, value) in UPSTREAM_HEADER_ECHO { + headers.append(name, HeaderValue::from_static(value)); + } + response + } else { + ( + UPSTREAM_HEADER_ECHO, + Json(json!({ + "id": "chatcmpl-test", + "object": "chat.completion", + "model": model, + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop" + }], + "usage": {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12} + })), + ) + .into_response() + }; + let headers = response.headers_mut(); + headers.append("x-upstream-trace", HeaderValue::from_static("trace-456")); + headers.append( + "link", + HeaderValue::from_static("; rel=next"), + ); + headers.append( + "link", + HeaderValue::from_static("; rel=prev"), + ); + headers.append( + "set-cookie", + HeaderValue::from_static("session=upstream; HttpOnly"), + ); + return response; + } if body["stream"].as_bool() == Some(true) { // Streamed tool call, for the namespace-on-every-event assertions. The // model calls a tool by the name it was given, so echo that name back. @@ -3596,6 +3659,108 @@ async fn responses_round_trips_codex_tool_namespaces() -> TestResult { Ok(()) } +/// Allowed upstream response headers ride through to the client, while body, cookie, +/// and Switchyard-owned headers do not; a header this server writes always beats an +/// upstream echo of the same name. +#[tokio::test] +async fn upstream_headers_forward_but_switchyard_writes_win() -> TestResult { + let (_upstream, app) = test_app(&[(ROUTE_MODEL, &["model/a"])]).await?; + let body = json!({ + "model": ROUTE_MODEL, + "messages": [{"role": "user", "content": "upstream-headers"}] + }); + let response = send(&app, "POST", "/v1/chat/completions", Some(body)).await?; + assert_eq!(response.status, StatusCode::OK); + + // Observability headers survive the proxy hop. + let traces = response + .headers + .get_all("x-upstream-trace") + .iter() + .map(|value| value.to_str()) + .collect::, _>>()?; + assert_eq!(traces, ["trace-123", "trace-456"]); + assert_eq!( + response + .headers + .get("x-request-id") + .and_then(|value| value.to_str().ok()), + Some("req-42") + ); + // Anthropic spells its correlation id without the `x-` prefix. + assert_eq!( + response + .headers + .get("request-id") + .and_then(|value| value.to_str().ok()), + Some("req_anthropic_42") + ); + assert!(!response.headers.contains_key("link")); + + // Upstream cookies must never become Switchyard-origin cookies. + assert!(!response.headers.contains_key("set-cookie")); + + // Switchyard's own namespace never forwards from upstream. + assert!(!response.headers.contains_key("x-switchyard-session-id")); + + // …and Switchyard's routing write beats the upstream echo. + assert_eq!( + response + .headers + .get("x-model-router-selected-model") + .and_then(|value| value.to_str().ok()), + Some("model/a") + ); + Ok(()) +} + +/// A streamed reply captures its headers off the response head, on a branch the +/// buffered path never touches, so the same contract is asserted there too. +#[tokio::test] +async fn upstream_headers_forward_on_streaming_responses() -> TestResult { + let (_upstream, app) = test_app(&[(ROUTE_MODEL, &["model/a"])]).await?; + let body = json!({ + "model": ROUTE_MODEL, + "stream": true, + "messages": [{"role": "user", "content": "upstream-headers"}] + }); + let response = send(&app, "POST", "/v1/chat/completions", Some(body)).await?; + assert_eq!(response.status, StatusCode::OK); + + let traces = response + .headers + .get_all("x-upstream-trace") + .iter() + .map(|value| value.to_str()) + .collect::, _>>()?; + assert_eq!(traces, ["trace-123", "trace-456"]); + assert_eq!( + response + .headers + .get("x-request-id") + .and_then(|value| value.to_str().ok()), + Some("req-42") + ); + assert_eq!( + response + .headers + .get("request-id") + .and_then(|value| value.to_str().ok()), + Some("req_anthropic_42") + ); + assert!(!response.headers.contains_key("link")); + assert!(!response.headers.contains_key("set-cookie")); + assert!(!response.headers.contains_key("x-switchyard-session-id")); + assert_eq!( + response + .headers + .get("x-model-router-selected-model") + .and_then(|value| value.to_str().ok()), + Some("model/a") + ); + Ok(()) +} + // Verifies a route declaring `vision = true` advertises image input, and that an // undeclared route still fails closed to text-only. //