Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 16 additions & 4 deletions crates/switchyard-runner/src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -207,7 +207,8 @@ impl DeploymentConfig {
}
}

let clients = self.build_clients()?;
let mut provider_api_keys = Vec::new();
let clients = self.build_clients(&mut provider_api_keys)?;
let targets = self.build_targets();
let fallback_base_url = self.fallback_base_url()?;
let mut routes = Vec::with_capacity(self.routes.len());
Expand Down Expand Up @@ -260,11 +261,16 @@ impl DeploymentConfig {
);
routes.push((config.id.clone(), route));
}
let runner = Runner::new(routes).with_fallback_url(fallback_base_url);
let runner = Runner::new(routes)
.with_fallback_url(fallback_base_url)
.with_provider_api_keys(provider_api_keys);
Ok(runner)
}

fn build_clients(&self) -> RunnerResult<BTreeMap<String, Arc<TranslatingLlmClient>>> {
fn build_clients(
&self,
provider_api_keys: &mut Vec<String>,
) -> RunnerResult<BTreeMap<String, Arc<TranslatingLlmClient>>> {
let mut models_by_client = self
.llm_clients
.keys()
Expand All @@ -273,7 +279,13 @@ impl DeploymentConfig {

for (name, client_config) in &self.llm_clients {
validate_value("llm client name", name)?;
build_backend(name, client_config, &BTreeMap::new(), None)?;
let backend = build_backend(name, client_config, &BTreeMap::new(), None)?;
let (Backend::OpenAiChat(config)
| Backend::OpenAiResponses(config)
| Backend::Anthropic(config)) = backend;
if let Some(key) = config.api_key {
provider_api_keys.push(key);
}
}
for (target_name, target) in &self.targets {
let client_config = self.llm_clients.get(&target.llm_client).ok_or_else(|| {
Expand Down
14 changes: 14 additions & 0 deletions crates/switchyard-runner/src/runner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ use crate::{ModelCapabilities, Route, RunnerError};
pub struct Runner {
routes: Vec<(ModelId, Route)>,
fallback_base_url: Option<String>,
provider_api_keys: Vec<String>,
}

/// Borrowed model metadata returned while listing routes.
Expand Down Expand Up @@ -61,9 +62,22 @@ impl Runner {
Self {
routes,
fallback_base_url: None,
provider_api_keys: Vec::new(),
}
}

/// Registers deployment-owned API keys for server response redaction.
/// TOML loading registers these automatically; programmatic hosts must supply them.
pub fn with_provider_api_keys(mut self, keys: Vec<String>) -> Self {
self.provider_api_keys = keys;
self
}

/// Returns deployment-owned secrets for the server's response redactor.
pub fn provider_api_keys(&self) -> &[String] {
&self.provider_api_keys
}

pub(crate) fn with_fallback_url(mut self, fallback_base_url: Option<String>) -> Self {
self.fallback_base_url = fallback_base_url;
self
Expand Down
23 changes: 18 additions & 5 deletions crates/switchyard-server/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
pub mod config;
mod metrics;
mod observability;
mod redaction;
mod response;
mod routing_log;
mod shutdown;
Expand Down Expand Up @@ -157,6 +158,7 @@ struct DecisionLlmClientResponse {
#[derive(Clone)]
pub struct ServerState {
runner: Arc<Runner>,
redactor: Arc<redaction::Redactor>,
fallback_http: reqwest::Client,
metrics: prometheus::Registry,
stats: StatsAccumulator,
Expand Down Expand Up @@ -210,7 +212,9 @@ impl ServerState {
metrics.clone(),
runner.models().map(|model| model.algorithm),
);
let redactor = redaction::Redactor::new(runner.provider_api_keys());
Ok(Self {
redactor: Arc::new(redactor),
runner: Arc::new(runner),
fallback_http,
metrics,
Expand Down Expand Up @@ -518,6 +522,10 @@ fn primary_llm_routes() -> Router<ServerState> {
fn finish_router(router: Router<ServerState>, state: ServerState) -> Router {
router
.layer(DefaultBodyLimit::max(DEFAULT_MAX_REQUEST_BODY_BYTES))
.layer(axum::middleware::from_fn_with_state(
state.clone(),
redaction::redact_response,
))
// `layer` only wraps routes registered before it, so this stays last.
.layer(axum::middleware::from_fn(stamp_request_start))
.with_state(state)
Expand Down Expand Up @@ -1060,11 +1068,16 @@ async fn handle_llm_request(

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()),
};
let mut response = match into_http_response(
response,
wire_format,
response_model,
request_extensions,
Arc::clone(&state.redactor),
) {
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();
Expand Down
221 changes: 221 additions & 0 deletions crates/switchyard-server/src/redaction.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,221 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Removes configured provider credentials from client-facing responses.

use axum::body::{Body, to_bytes};
use axum::extract::{Request, State};
use axum::http::header::{CONTENT_ENCODING, CONTENT_LENGTH, CONTENT_TYPE};
use axum::middleware::Next;
use axum::response::Response;

use crate::{DEFAULT_MAX_REQUEST_BODY_BYTES, ServerState};

#[derive(Default)]
pub(crate) struct Redactor {
raw: Vec<String>,
json: Vec<String>,
}

impl Redactor {
pub(crate) fn new(keys: &[String]) -> Self {
let mut raw: Vec<String> = keys.iter().filter(|key| !key.is_empty()).cloned().collect();
raw.sort_by(|a, b| b.len().cmp(&a.len()).then_with(|| a.cmp(b)));
raw.dedup();
let mut json = Vec::with_capacity(raw.len());
for key in &raw {
let Ok(encoded) = serde_json::to_string(key) else {
unreachable!("serde_json::to_string on a String cannot fail");
Comment thread
grahamking marked this conversation as resolved.
};
json.push(encoded[1..encoded.len() - 1].to_string());
}
json.sort_by_key(|key| std::cmp::Reverse(key.len()));
Self { raw, json }
}

pub(crate) fn json(&self, value: String) -> String {
replace(value, &self.json)
}

pub(crate) fn text(&self, value: String) -> String {
replace(value, &self.raw)
}
}

fn replace(mut value: String, secrets: &[String]) -> String {
for secret in secrets {
if value.contains(secret) {
value = value.replace(secret, "[REDACTED]");
}
}
value
}

pub(crate) async fn redact_response(
State(state): State<ServerState>,
request: Request,
next: Next,
) -> Response {
let mut response = next.run(request).await;
let redactor = &state.redactor;
if redactor.raw.is_empty() {
return response;
}
let is_json = response.headers().get(CONTENT_TYPE).is_some_and(|value| {
value.to_str().is_ok_and(|value| {
value.split(';').next().is_some_and(|mime| {
mime.trim() == "application/json" || mime.trim().ends_with("+json")
})
})
});
let is_encoded = response
.headers()
.get_all(CONTENT_ENCODING)
.iter()
.any(|value| {
value.to_str().map_or(true, |value| {
value
.split(',')
.any(|encoding| !encoding.trim().eq_ignore_ascii_case("identity"))
})
});
for value in response.headers_mut().values_mut() {
if let Ok(text) = value.to_str() {
let sanitized = redactor.text(text.to_string());
if sanitized != text {
// Replacement contains only visible ASCII and cannot invalidate a header.
if let Ok(header) = sanitized.parse() {
*value = header;
}
}
}
}
// SSE bodies are redacted per event before framing, without buffering the stream.
// The fallback proxy passes compressed bodies through without decoding them.
if is_json && !is_encoded {
let (mut parts, body) = response.into_parts();
let body = match to_bytes(body, DEFAULT_MAX_REQUEST_BODY_BYTES).await {
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Ok(bytes) => match String::from_utf8(bytes.to_vec()) {
Ok(json) => Body::from(redactor.json(json)),
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Err(_) => return crate::server_error("Invalid upstream JSON"),
},
Err(_) => return crate::server_error("Unable to read response body"),
};
parts.headers.remove(CONTENT_LENGTH);
response = Response::from_parts(parts, body);
}
response
}

#[cfg(test)]
mod tests {
use std::sync::Arc;

use axum::Json;
use axum::response::IntoResponse;
use axum::routing::get;
use serde_json::json;
use switchyard_runner::Runner;
use switchyard_translation::{LlmStreamError, WireFormat};
use tower::ServiceExt;

use super::*;

#[tokio::test]
async fn responses_redact_provider_keys_without_changing_errors()
-> Result<(), Box<dyn std::error::Error + Send + Sync>> {
let key = "synthetic-\"provider\\key";
let error = json!({"error": {
"message": "upstream failed",
"type": "provider_error",
"debug": format!("Bearer {key}"),
key: [key, "ordinary diagnostics"]
}});
let runner = Runner::new(Vec::new()).with_provider_api_keys(vec![key.to_string()]);
let state = ServerState::from_runner(runner)?;
let buffered = error.clone();
let streamed = error.clone();
let header = axum::http::HeaderValue::from_str(key)?;
// gzip-compressed {"ok":true}.
const GZIP_JSON: &[u8] = &[
31, 139, 8, 0, 0, 0, 0, 0, 2, 3, 171, 86, 202, 207, 86, 178, 42, 41, 42, 77, 173, 5, 0,
144, 95, 212, 167, 11, 0, 0, 0,
];
let router = axum::Router::new()
.route(
"/compressed",
get(|| async {
(
[
(CONTENT_TYPE, "application/json"),
(CONTENT_ENCODING, "gzip"),
(CONTENT_LENGTH, "31"),
],
GZIP_JSON,
)
}),
)
.route("/buffered", get(move || async move { Json(buffered) }))
.route(
"/stream",
get(move |State(state): State<ServerState>| async move {
let events = futures_util::stream::iter([
Ok(json!({"choices": [], "model": "ordinary-model"})),
Err(LlmStreamError::Upstream(streamed)),
]);
let mut response = crate::sse::frame_stream(
Box::pin(events),
WireFormat::OpenAiChat,
Arc::clone(&state.redactor),
)
.into_response();
response.headers_mut().insert("x-upstream-debug", header);
response
}),
);
let app = crate::finish_router(router, state);
let compressed = app
.clone()
.oneshot(Request::builder().uri("/compressed").body(Body::empty())?)
.await?;
assert_eq!(compressed.status(), axum::http::StatusCode::OK);
assert_eq!(compressed.headers()[CONTENT_ENCODING], "gzip");
assert_eq!(compressed.headers()[CONTENT_LENGTH], "31");
assert_eq!(
to_bytes(compressed.into_body(), usize::MAX).await?.as_ref(),
GZIP_JSON
);
for path in ["/buffered", "/stream"] {
let response = app
.clone()
.oneshot(Request::builder().uri(path).body(Body::empty())?)
.await?;
if path == "/stream" {
assert_eq!(response.headers()["x-upstream-debug"], "[REDACTED]");
}
let body =
String::from_utf8(to_bytes(response.into_body(), usize::MAX).await?.to_vec())?;
let error: serde_json::Value = if path == "/stream" {
assert!(body.contains("ordinary-model"));
assert!(!body.contains("[DONE]"));
let data = body
.lines()
.filter_map(|line| line.strip_prefix("data: "))
.next_back()
.ok_or("missing SSE error")?;
serde_json::from_str(data)?
} else {
serde_json::from_str(&body)?
};
assert_eq!(error["error"]["message"], "upstream failed");
assert_eq!(error["error"]["type"], "provider_error");
assert_eq!(error["error"]["debug"], "Bearer [REDACTED]");
assert_eq!(
error["error"]["[REDACTED]"],
json!(["[REDACTED]", "ordinary diagnostics"])
);
assert!(!body.contains("synthetic-"));
}
Ok(())
}
}
4 changes: 3 additions & 1 deletion crates/switchyard-server/src/response.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
//! Response encoding glue for libsy server endpoints.

use std::error::Error;
use std::sync::Arc;

use axum::Json;
use axum::response::{IntoResponse, Response as HttpResponse};
Expand All @@ -24,6 +25,7 @@ pub(crate) fn into_http_response(
target_format: WireFormat,
served_model: Option<String>,
request_extensions: ProviderExtensions,
redactor: Arc<crate::redaction::Redactor>,
) -> Result<HttpResponse, BoxError> {
match response.llm_response {
LlmResponse::Agg(response) => {
Expand All @@ -42,7 +44,7 @@ pub(crate) fn into_http_response(
served_model,
&request_extensions,
)?;
Ok(frame_stream(events, target_format).into_response())
Ok(frame_stream(events, target_format, redactor).into_response())
}
}
}
Loading
Loading