From db9f11ed59b2741a455df72f4f90c8204b5dc0fb Mon Sep 17 00:00:00 2001 From: nachiketb Date: Fri, 28 Aug 2026 11:03:32 -0700 Subject: [PATCH 1/5] feat(server): add fallback client proxy Signed-off-by: nachiketb --- Cargo.lock | 1 + crates/libsy-llm-client/src/backend.rs | 50 ++------ crates/libsy-llm-client/src/client.rs | 106 +++++----------- crates/switchyard-runner/src/algorithm.rs | 2 +- crates/switchyard-runner/src/config.rs | 58 ++++----- crates/switchyard-runner/src/failure.rs | 10 +- crates/switchyard-runner/src/lib.rs | 4 +- crates/switchyard-runner/src/route.rs | 28 +--- crates/switchyard-runner/src/runner.rs | 46 ++++++- crates/switchyard-runner/tests/route.rs | 1 - crates/switchyard-server/Cargo.toml | 1 + crates/switchyard-server/README.md | 7 +- crates/switchyard-server/src/lib.rs | 77 ++++------- crates/switchyard-server/tests/server.rs | 148 +++++----------------- docs/reference/toml_schema.md | 5 + 15 files changed, 187 insertions(+), 357 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index fd4d9b6bf..89b712d9d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2392,6 +2392,7 @@ dependencies = [ "opentelemetry_sdk", "parking_lot", "prometheus", + "reqwest", "rustls", "serde", "serde_json", diff --git a/crates/libsy-llm-client/src/backend.rs b/crates/libsy-llm-client/src/backend.rs index fab446095..15ff205f1 100644 --- a/crates/libsy-llm-client/src/backend.rs +++ b/crates/libsy-llm-client/src/backend.rs @@ -255,17 +255,19 @@ impl Backend { self.config().max_retries } - /// Whether this backend speaks the Anthropic Messages wire format — the only - /// one with a `count_tokens` endpoint. - pub fn is_anthropic(&self) -> bool { - matches!(self, Backend::Anthropic(_)) - } - - /// The upstream `/v1/messages/count_tokens` URL, derived from the same base - /// URL join as [`url`](Self::url). - pub fn count_tokens_url(&self) -> String { + /// Resolves an unmatched provider path against this backend's API root. + pub(crate) fn forwarding_url(&self, path_and_query: &str) -> String { let base_url = self.config().base_url.trim_end_matches('/'); - format!("{}/count_tokens", anthropic_url(base_url)) + let root = [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages", + "/v1", + ] + .iter() + .find_map(|suffix| base_url.strip_suffix(suffix)) + .unwrap_or(base_url); + format!("{root}{path_and_query}") } /// Whether an upstream 400 `body` looks like a context-window overflow for @@ -391,34 +393,6 @@ mod tests { ); } - #[test] - fn count_tokens_url_joins_every_base_url_shape() { - assert_eq!( - Backend::Anthropic(config("https://host")).count_tokens_url(), - "https://host/v1/messages/count_tokens" - ); - assert_eq!( - Backend::Anthropic(config("https://host/v1")).count_tokens_url(), - "https://host/v1/messages/count_tokens" - ); - assert_eq!( - Backend::Anthropic(config("https://host/v1/messages")).count_tokens_url(), - "https://host/v1/messages/count_tokens" - ); - // Trailing slash is trimmed before the join. - assert_eq!( - Backend::Anthropic(config("https://host/v1/")).count_tokens_url(), - "https://host/v1/messages/count_tokens" - ); - } - - #[test] - fn only_anthropic_backend_is_anthropic() { - assert!(Backend::Anthropic(config("x")).is_anthropic()); - assert!(!Backend::OpenAiChat(config("x")).is_anthropic()); - assert!(!Backend::OpenAiResponses(config("x")).is_anthropic()); - } - #[test] fn wire_format_matches_variant() { assert_eq!( diff --git a/crates/libsy-llm-client/src/client.rs b/crates/libsy-llm-client/src/client.rs index 4a517936f..3b10dc32a 100644 --- a/crates/libsy-llm-client/src/client.rs +++ b/crates/libsy-llm-client/src/client.rs @@ -10,9 +10,9 @@ use std::time::{Duration, SystemTime}; use async_trait::async_trait; use futures_util::{StreamExt, stream}; -use http::StatusCode; +use http::{Method, StatusCode}; use reqwest::RequestBuilder; -use reqwest::header::{HeaderMap, RETRY_AFTER}; +use reqwest::header::{CONTENT_LENGTH, CONTENT_TYPE, HeaderMap, RETRY_AFTER}; use serde_json::{Map, Value}; use switchyard_protocol::{ LlmRequest, LlmResponse, LlmResponseChunk, LlmResponseStreamEvent, Metadata, ModelId, Request, @@ -146,49 +146,35 @@ impl TranslatingLlmClient { }) } - /// Whether `model` has an Anthropic backend that supports token counting. - pub fn supports_count_tokens(&self, model: &ModelId) -> bool { - self.backend_for(model, WireFormat::AnthropicMessages) - .is_some() - } - - /// Counts input tokens with `model`'s Anthropic backend. - /// - /// Returns an error when the model has no Anthropic backend or the upstream - /// request fails or returns invalid JSON. - pub async fn count_tokens(&self, model: &ModelId, request: Request) -> Result { - let backend = self - .backend_for(model, WireFormat::AnthropicMessages) - .ok_or_else(|| LlmClientError::Configuration { - message: format!("model {model} has no Anthropic backend for count_tokens"), - })?; - let Request { - mut llm_request, - metadata, - .. - } = request; - llm_request.model = Some(model.to_string()); - let http_response = self - .send_encoded( - backend, - WireFormat::AnthropicMessages, - llm_request, - metadata.as_ref(), - model, - UpstreamEndpoint::CountTokens, - ) - .await?; - let body = match http_response { - EncodedResponse::Buffered { body, .. } => body, - EncodedResponse::Streaming(_) => { - return Err(LlmClientError::InvalidRequest { - message: "count_tokens does not support streaming requests".to_string(), - }); - } + /// Forwards a provider-native request through `backend` without translation. + pub async fn forward( + &self, + backend: &Backend, + method: Method, + path_and_query: &str, + body: reqwest::Body, + metadata: Option<&Metadata>, + ) -> Result { + let client = if backend.is_forwarding_auth() { + &self.forward_auth_client + } else { + &self.client }; - serde_json::from_slice(&body).map_err(|error| LlmClientError::InvalidResponse { - source: Box::new(error), - }) + let mut builder = client + .request(method, backend.forwarding_url(path_and_query)) + .body(body); + builder = forward_metadata_headers(builder, metadata); + if let Some(headers) = metadata.and_then(|metadata| metadata.http_headers.as_ref()) { + for name in [CONTENT_TYPE, CONTENT_LENGTH] { + if let Some(value) = headers.get(&name) { + builder = builder.header(name, value); + } + } + } + builder = backend.apply_forwarded_auth(builder, metadata); + builder = apply_extra_headers(builder, backend); + builder = backend.apply_auth(builder); + builder.send().await.map_err(convert_reqwest_error) } /// Encode `llm_request` for `wire_format`, POST it to `url` with the request's @@ -198,10 +184,8 @@ impl TranslatingLlmClient { /// response is returned as soon as its successful headers arrive. A non-success /// status maps to a typed error — a 400 is classified as a context-window /// overflow via the backend's provider rules. Shared by - /// [`call_rewrite_model`](Self::call_rewrite_model) (which POSTs to the - /// backend's completion URL and decodes a response) and - /// [`count_tokens`](Self::count_tokens) (which POSTs to the `count_tokens` - /// URL and returns the raw JSON). + /// [`call_rewrite_model`](Self::call_rewrite_model), which POSTs to the + /// backend's completion URL and decodes a response. async fn send_encoded( &self, backend: &Backend, @@ -209,7 +193,6 @@ impl TranslatingLlmClient { llm_request: LlmRequest, metadata: Option<&Metadata>, model: &ModelId, - endpoint: UpstreamEndpoint, ) -> Result { let mut body = encode_request(&llm_request, wire_format) .map_err(|error| LlmClientError::RequestEncoding(error.to_string()))?; @@ -230,9 +213,8 @@ impl TranslatingLlmClient { if matches!(backend, Backend::OpenAiChat(_)) { ensure_openai_stream_usage(&mut body); } - let streaming = endpoint.allows_streaming() - && body.get("stream").and_then(Value::as_bool).unwrap_or(false); - let url = endpoint.url(backend); + let streaming = body.get("stream").and_then(Value::as_bool).unwrap_or(false); + let url = backend.url(); record_gen_ai_request(&url, model, streaming); let max_retries = u64::from(backend.max_retries()); @@ -426,7 +408,6 @@ impl TranslatingLlmClient { llm_request, metadata.as_ref(), &model_id, - UpstreamEndpoint::Completion, ) .await?; @@ -559,25 +540,6 @@ impl RoutedLlmClient for TranslatingLlmClient { } } -#[derive(Clone, Copy)] -enum UpstreamEndpoint { - Completion, - CountTokens, -} - -impl UpstreamEndpoint { - fn url(self, backend: &Backend) -> String { - match self { - UpstreamEndpoint::Completion => backend.url(), - UpstreamEndpoint::CountTokens => backend.count_tokens_url(), - } - } - - fn allows_streaming(self) -> bool { - matches!(self, UpstreamEndpoint::Completion) - } -} - enum EncodedResponse { Buffered { status: u16, body: Vec }, Streaming(reqwest::Response), diff --git a/crates/switchyard-runner/src/algorithm.rs b/crates/switchyard-runner/src/algorithm.rs index a601ab94a..243c56b60 100644 --- a/crates/switchyard-runner/src/algorithm.rs +++ b/crates/switchyard-runner/src/algorithm.rs @@ -467,7 +467,7 @@ impl AlgorithmSpec { names } // The advisor is judge-only: reviews go through its own client, - // so it is not a completion (or count_tokens) destination. + // so it is not a completion destination. Self::Advisor { executor_target, .. } => vec![executor_target], diff --git a/crates/switchyard-runner/src/config.rs b/crates/switchyard-runner/src/config.rs index d4887fab1..12b35d02e 100644 --- a/crates/switchyard-runner/src/config.rs +++ b/crates/switchyard-runner/src/config.rs @@ -17,9 +17,9 @@ use switchyard_llm_client::{ }; use switchyard_protocol::{ModelId, RoutedLlmClient, WireFormat}; +use crate::runner::FallbackClient; use crate::{ - AlgorithmSpec, CallerAuthKind, CountTokensTarget, DecisionTarget, ModelCapabilities, Route, - Runner, RunnerError, + AlgorithmSpec, CallerAuthKind, DecisionTarget, ModelCapabilities, Route, Runner, RunnerError, }; const SUPPORTED_SCHEMA_VERSION: u32 = 1; @@ -54,6 +54,7 @@ pub(crate) fn runner_from_toml(source: &str) -> RunnerResult { #[serde(deny_unknown_fields)] pub(crate) struct DeploymentConfig { schema_version: u32, + fallback_client: Option, #[serde(default)] llm_clients: BTreeMap, targets: BTreeMap, @@ -165,6 +166,7 @@ impl DeploymentConfig { let clients = self.build_clients()?; let targets = self.build_targets(); + let fallback_client = self.build_fallback_client(&clients)?; let mut routes = Vec::with_capacity(self.routes.len()); for (route_name, config) in &self.routes { validate_value("route name", route_name)?; @@ -188,7 +190,6 @@ impl DeploymentConfig { .map_err(|error| RunnerError::configuration_source(error.to_string(), error))?; let (route_clients, caller_auth) = self.build_route_clients(route_name, config, &clients)?; - let count_tokens_target = self.build_count_tokens_target(config, &clients); let decision_targets = config .routing_target_names() .into_iter() @@ -199,12 +200,11 @@ impl DeploymentConfig { route_clients, caller_auth, capabilities, - count_tokens_target, decision_targets, ); routes.push((config.id.clone(), route)); } - Ok(Runner::new(routes)) + Ok(Runner::new(routes).with_fallback_client(fallback_client)) } fn build_clients(&self) -> RunnerResult>> { @@ -291,42 +291,28 @@ impl DeploymentConfig { Ok((ClientRouter::new(by_model), caller_auth)) } - fn build_count_tokens_target( + fn build_fallback_client( &self, - route: &RouteConfig, clients: &BTreeMap>, - ) -> Option { - route - .routing_target_names() - .into_iter() - .enumerate() - .filter_map(|(index, name)| { - let target = self.targets.get(name)?; - let client = clients.get(&target.llm_client)?; - client.supports_count_tokens(&target.id).then_some(( - count_tokens_priority(name, &target.id), - index, - target, - client, - )) - }) - .min_by_key(|(priority, index, _, _)| (*priority, *index)) - .map(|(_, _, target, client)| CountTokensTarget { - model: target.id.clone(), - client: client.clone(), - }) + ) -> RunnerResult> { + let Some(name) = &self.fallback_client else { + return Ok(None); + }; + let config = self.llm_clients.get(name).ok_or_else(|| { + RunnerError::configuration(format!( + "fallback_client references unknown llm client {name}" + )) + })?; + let client = clients.get(name).ok_or_else(|| { + RunnerError::configuration("validated fallback client was not initialized") + })?; + Ok(Some(FallbackClient { + backend: build_backend(name, config, &BTreeMap::new())?, + client: client.clone(), + })) } } -fn count_tokens_priority(target_name: &str, model_id: &ModelId) -> usize { - let target_name = target_name.to_ascii_lowercase(); - let model_id = model_id.to_ascii_lowercase(); - ["opus", "sonnet", "haiku"] - .iter() - .position(|hint| target_name.contains(hint) || model_id.contains(hint)) - .unwrap_or(3) -} - /// A client endpoint, parsed when the config loads rather than checked afterwards. /// /// Holding a `HttpBaseUrl` is proof the value is an absolute HTTP(S) URL, so no diff --git a/crates/switchyard-runner/src/failure.rs b/crates/switchyard-runner/src/failure.rs index 4e2901b7a..7dd470181 100644 --- a/crates/switchyard-runner/src/failure.rs +++ b/crates/switchyard-runner/src/failure.rs @@ -98,9 +98,7 @@ impl RunnerError { None, None, ), - Self::UnknownRouteModel(_) - | Self::IncompatibleCallerFormat(_) - | Self::CountTokensUnsupported => summary( + Self::UnknownRouteModel(_) | Self::IncompatibleCallerFormat(_) => summary( RouteErrorKind::InvalidRequest, RouteErrorPhase::BeforeResponse, None, @@ -328,11 +326,5 @@ mod tests { configuration.execution_error_summary().kind, RouteErrorKind::Configuration )); - - let unsupported = RunnerError::CountTokensUnsupported; - assert!(matches!( - unsupported.execution_error_summary().kind, - RouteErrorKind::InvalidRequest - )); } } diff --git a/crates/switchyard-runner/src/lib.rs b/crates/switchyard-runner/src/lib.rs index 75543c2fe..850e513dc 100644 --- a/crates/switchyard-runner/src/lib.rs +++ b/crates/switchyard-runner/src/lib.rs @@ -14,7 +14,5 @@ pub use algorithm::{ ClassifierPolicyConfig, LlmClassifierRouteConfig, StageClassifierConfig, SubagentRouteConfig, }; pub use failure::{RouteErrorKind, RouteErrorPhase, RouteErrorSummary, stream_error_summary}; -pub use route::{ - CallerAuthKind, CountTokensTarget, ModelCapabilities, Route, RunOutput, RunnerError, -}; +pub use route::{CallerAuthKind, ModelCapabilities, Route, RunOutput, RunnerError}; pub use runner::{DecisionDescription, DecisionTarget, ModelInfo, Runner}; diff --git a/crates/switchyard-runner/src/route.rs b/crates/switchyard-runner/src/route.rs index dc4db13ee..451f365c1 100644 --- a/crates/switchyard-runner/src/route.rs +++ b/crates/switchyard-runner/src/route.rs @@ -7,8 +7,7 @@ use std::error::Error; use std::sync::Arc; use libsy::{Algorithm, CallModel, LibsyError, RoutingOutcome, drive}; -use serde_json::Value; -use switchyard_llm_client::{ClientRouter, RunObserver, TranslatingLlmClient}; +use switchyard_llm_client::{ClientRouter, RunObserver}; use switchyard_protocol::{LlmClientError, ModelId, Request, Response, WireFormat}; use thiserror::Error; @@ -56,13 +55,6 @@ impl CallerAuthKind { } } -/// Exact upstream model used for Anthropic token counting. -#[derive(Clone)] -pub struct CountTokensTarget { - pub model: ModelId, - pub client: Arc, -} - /// Error returned while loading or executing configured routes. #[derive(Debug, Error)] pub enum RunnerError { @@ -76,8 +68,6 @@ pub enum RunnerError { UnknownRouteModel(String), #[error("caller format is incompatible with {} credentials", .0.as_str())] IncompatibleCallerFormat(CallerAuthKind), - #[error("route has no Anthropic target for token counting")] - CountTokensUnsupported, #[error(transparent)] Algorithm(#[from] LibsyError), #[error(transparent)] @@ -112,7 +102,6 @@ pub struct Route { clients: ClientRouter, caller_auth: Option, capabilities: ModelCapabilities, - count_tokens_target: Option, decision_targets: Vec, } @@ -129,7 +118,6 @@ impl Route { clients: ClientRouter, caller_auth: Option, capabilities: ModelCapabilities, - count_tokens_target: Option, decision_targets: Vec, ) -> Self { Self { @@ -137,7 +125,6 @@ impl Route { clients, caller_auth, capabilities, - count_tokens_target, decision_targets, } } @@ -202,19 +189,6 @@ impl Route { .await .map_err(Into::into) } - - /// Counts tokens using the configured Anthropic-capable target. - pub async fn count_tokens(&self, request: Request) -> Result { - let target = self - .count_tokens_target - .as_ref() - .ok_or(RunnerError::CountTokensUnsupported)?; - target - .client - .count_tokens(&target.model, request) - .await - .map_err(Into::into) - } } async fn serve_decision_dependency(clients: ClientRouter, call: CallModel) -> libsy::Result<()> { diff --git a/crates/switchyard-runner/src/runner.rs b/crates/switchyard-runner/src/runner.rs index 93dce5e76..c4055cf4b 100644 --- a/crates/switchyard-runner/src/runner.rs +++ b/crates/switchyard-runner/src/runner.rs @@ -5,10 +5,13 @@ use std::collections::BTreeMap; use std::path::Path; +use std::sync::Arc; use libsy::RoutingOutcome; +use reqwest::{Body, Method, Response}; use serde_json::Value; -use switchyard_protocol::{ModelId, WireFormat}; +use switchyard_llm_client::{Backend, TranslatingLlmClient}; +use switchyard_protocol::{Metadata, ModelId, WireFormat}; use crate::config; use crate::{ModelCapabilities, Route, RunnerError}; @@ -16,6 +19,12 @@ use crate::{ModelCapabilities, Route, RunnerError}; /// Immutable named route table. pub struct Runner { routes: Vec<(ModelId, Route)>, + fallback_client: Option, +} + +pub(crate) struct FallbackClient { + pub backend: Backend, + pub client: Arc, } /// Borrowed model metadata returned while listing routes. @@ -57,7 +66,15 @@ impl Runner { /// Builds a runner from named routes in caller-provided order. /// Pre-condition: There must be at least one route. pub fn new(routes: Vec<(ModelId, Route)>) -> Self { - Self { routes } + Self { + routes, + fallback_client: None, + } + } + + pub(crate) fn with_fallback_client(mut self, fallback_client: Option) -> Self { + self.fallback_client = fallback_client; + self } /// Returns the route registered for a model. @@ -77,6 +94,31 @@ impl Runner { }) } + /// Forwards an unmatched request through the configured fallback client. + pub async fn forward_fallback( + &self, + method: Method, + path_and_query: &str, + body: Body, + metadata: Metadata, + ) -> Result, RunnerError> { + let Some(fallback) = &self.fallback_client else { + return Ok(None); + }; + fallback + .client + .forward( + &fallback.backend, + method, + path_and_query, + body, + Some(&metadata), + ) + .await + .map(Some) + .map_err(Into::into) + } + /// Resolves an outcome to configured target names and non-secret client settings. pub fn describe_decision( &self, diff --git a/crates/switchyard-runner/tests/route.rs b/crates/switchyard-runner/tests/route.rs index 71ce57c7e..9442004ba 100644 --- a/crates/switchyard-runner/tests/route.rs +++ b/crates/switchyard-runner/tests/route.rs @@ -51,7 +51,6 @@ fn plugin_route(client: Arc) -> Route { clients, None, ModelCapabilities::default(), - None, Vec::new(), ) } diff --git a/crates/switchyard-server/Cargo.toml b/crates/switchyard-server/Cargo.toml index 0c4b47e24..29cea432f 100644 --- a/crates/switchyard-server/Cargo.toml +++ b/crates/switchyard-server/Cargo.toml @@ -32,6 +32,7 @@ opentelemetry-prometheus = "0.32" opentelemetry_sdk = { version = "0.32", default-features = false, features = ["metrics", "trace"] } parking_lot.workspace = true prometheus = "0.14" +reqwest.workspace = true serde.workspace = true switchyard-llm-client.workspace = true switchyard-protocol.workspace = true diff --git a/crates/switchyard-server/README.md b/crates/switchyard-server/README.md index 2ff1a7a5a..7e6426260 100644 --- a/crates/switchyard-server/README.md +++ b/crates/switchyard-server/README.md @@ -129,7 +129,7 @@ documented in [Stage-Router Routing](../../docs/routing_algorithms/stage_router_ | `POST` | `/v1/messages` | Anthropic Messages | | `POST` | `/v1/responses` | OpenAI Responses | | `POST` | `/v1/decision` | Resolve selected and fallback targets without a post-routing answer call | -| `POST` | `/v1/messages/count_tokens` | Token count from a route's Anthropic target | +| `ANY` | Any unmatched path | Raw forward through the optional `fallback_client` | | `GET` | `/v1/models` | Routes served by this deployment | | `GET` | `/v1/stats` | Per-model usage plus curated algorithm stats | | `POST` | `/v1/stats/reset` | Clear accumulated stats | @@ -140,6 +140,11 @@ Requests name a route by its `id`, so `POST /v1/chat/completions` with `"model": routes through the `[routes.general]` entry above. Any of the three request formats can address any route, and the server translates between them. +Set the top-level `fallback_client` to an entry under `[llm_clients]` to proxy any otherwise +unmatched method and path through that client. The fallback client does not need a target. Requests +and responses are forwarded without translation, including paths, query strings, bodies, and model +identifiers. Without `fallback_client`, unmatched paths return `404`. + `POST /v1/decision` accepts `{"input_format": "openai_chat", "request": {...}}`, where the nested request names the route in `model`. It executes required classifier or judge calls, then returns the selected target and ordered fallbacks with their model, format, base URL, and diff --git a/crates/switchyard-server/src/lib.rs b/crates/switchyard-server/src/lib.rs index 77150399c..a382890d1 100644 --- a/crates/switchyard-server/src/lib.rs +++ b/crates/switchyard-server/src/lib.rs @@ -23,6 +23,7 @@ use std::path::PathBuf; use std::sync::Arc; use std::time::{Duration, Instant}; +use axum::body::Body; use axum::extract::rejection::{JsonRejection, QueryRejection}; use axum::extract::{DefaultBodyLimit, Query, Request as HttpRequest, State}; use axum::http::header::CONTENT_TYPE; @@ -189,7 +190,6 @@ impl ServerState { clients, None, ModelCapabilities::default(), - None, Vec::new(), ), ) @@ -478,7 +478,6 @@ pub fn build_switchyard_router(state: ServerState) -> Router { .route("/v1/messages", post(anthropic_messages)) .route("/v1/responses", post(openai_responses)) .route("/v1/decision", post(decision)) - .route("/v1/messages/count_tokens", post(anthropic_count_tokens)) .route("/v1/models", get(models)) .route("/v1/stats", get(get_stats)) .route("/v1/stats/reset", post(reset_stats)) @@ -488,7 +487,7 @@ pub fn build_switchyard_router(state: ServerState) -> Router { router = router.route("/v1/routing/session-stats", get(get_session_stats)); } router - .fallback(not_found) + .fallback(proxy_unmatched) .layer(DefaultBodyLimit::max(DEFAULT_MAX_REQUEST_BODY_BYTES)) // `layer` only wraps routes registered before it, so this stays last. .layer(axum::middleware::from_fn(stamp_request_start)) @@ -539,6 +538,29 @@ async fn openai_responses( handle_endpoint(state, started, headers, body, WireFormat::OpenAiResponses).await } +// Forwards unmatched requests through the deployment's optional fallback client. +async fn proxy_unmatched(State(state): State, request: HttpRequest) -> Response { + let (parts, body) = request.into_parts(); + let path_and_query = parts + .uri + .path_and_query() + .map_or_else(|| parts.uri.path().to_string(), ToString::to_string); + let metadata = metadata_from_headers(parts.headers); + let body = reqwest::Body::wrap_stream(body.into_data_stream()); + match state + .runner + .forward_fallback(parts.method, &path_and_query, body, metadata) + .await + { + Ok(Some(response)) => { + let response: http::Response = response.into(); + response.map(Body::new) + } + Ok(None) => not_found().await, + Err(error) => runner_error(error), + } +} + /// One provider request submitted for a routing decision without an answer-model call. #[derive(Deserialize)] #[serde(deny_unknown_fields)] @@ -613,45 +635,6 @@ async fn decision( } } -/// Anthropic token counting against the route's explicitly configured target. -async fn anthropic_count_tokens( - State(state): State, - headers: HeaderMap, - body: std::result::Result, JsonRejection>, -) -> Response { - let body = match llm_json_body(body) { - Ok(body) => body, - Err((status, message)) => { - return anthropic_error_response(invalid_body_error(status, message)); - } - }; - let (route, request) = match resolve_route( - &state, - metadata_from_headers(headers), - body, - WireFormat::AnthropicMessages, - ) { - Ok(resolved) => resolved, - Err(response) => return anthropic_error_response(response), - }; - anthropic_error_response(match route.count_tokens(request).await { - Ok(payload) => (StatusCode::OK, Json(payload)).into_response(), - Err(RunnerError::CountTokensUnsupported) => error_response( - StatusCode::BAD_REQUEST, - "route has no Anthropic target for token counting", - "invalid_request_error", - "count_tokens_unsupported", - ), - Err(RunnerError::Client(error)) => count_tokens_error(error), - Err(error) => server_error(error.to_string()), - }) -} - -/// Maps a token-count client failure with the same policy as a routed client call. -fn count_tokens_error(error: LlmClientError) -> Response { - client_error(&error) -} - async fn handle_endpoint( state: ServerState, started: RequestStart, @@ -739,8 +722,8 @@ fn llm_json_body( } /// Decode `body`, resolve the route named by its `model`, and build the -/// [`Request`]. Shared by the completion and `count_tokens` handlers. Returns -/// the resolved route and the built request — or an error [`Response`] +/// [`Request`]. Shared by the completion handlers. Returns the resolved route +/// and the built request — or an error [`Response`] /// (invalid body, empty `model` → 400, unknown route → 404). // Both callers immediately return the `Err(Response)` as the HTTP response, so // the large error type is intentional, not propagated up a call stack. @@ -1128,10 +1111,6 @@ fn render_error_response(response: Response, wire_format: WireFormat) -> Respons error.into_response(wire_format) } -fn anthropic_error_response(response: Response) -> Response { - render_error_response(response, WireFormat::AnthropicMessages) -} - fn anthropic_error_type(status: StatusCode) -> &'static str { match status { StatusCode::BAD_REQUEST => "invalid_request_error", @@ -1449,7 +1428,7 @@ fn endpoint_listing(has_routing_log: bool) -> String { " POST /v1/chat/completions OpenAI Chat Completions", " POST /v1/messages Anthropic Messages", " POST /v1/responses OpenAI Responses", - " POST /v1/messages/count_tokens", + " ANY unmatched paths optional fallback client", " GET /v1/models configured routes", " GET /v1/stats routing stats", " POST /v1/stats/reset", diff --git a/crates/switchyard-server/tests/server.rs b/crates/switchyard-server/tests/server.rs index 32290ea36..0343ca820 100644 --- a/crates/switchyard-server/tests/server.rs +++ b/crates/switchyard-server/tests/server.rs @@ -57,7 +57,7 @@ impl MockUpstream { post(upstream_responses_requires_forwarded_auth), ) .route("/capture", post(upstream_redirect_capture)) - .route("/v1/messages/count_tokens", post(upstream_count_tokens)) + .route("/future/provider/endpoint", post(upstream_fallback)) .layer(DefaultBodyLimit::disable()) .with_state(Arc::clone(&calls)); let listener = TcpListener::bind("127.0.0.1:0").await?; @@ -422,7 +422,7 @@ async fn upstream_redirect_capture( StatusCode::OK.into_response() } -async fn upstream_count_tokens( +async fn upstream_fallback( State(calls): State>>>, Json(body): Json, ) -> HttpResponse { @@ -1713,28 +1713,29 @@ response_format_type = "json_object" } #[tokio::test] -async fn count_tokens_forwards_to_configured_anthropic_target() -> TestResult { +async fn fallback_client_forwards_unmatched_requests_and_is_optional() -> TestResult { let upstream = MockUpstream::start().await?; let state = load_test_config(&format!( r#" schema_version = 1 +fallback_client = "fallback" -[llm_clients.claude] -format = "anthropic_messages" +[llm_clients.fallback] +format = "openai_responses" base_url = "{base_url}" -[targets.strong] -id = "real/opus" -llm_client = "claude" +[llm_clients.routed] +format = "openai_chat" +base_url = "{base_url}" -[targets.other] -id = "real/sonnet" -llm_client = "claude" +[targets.weak] +id = "model/weak" +llm_client = "routed" [routes.random] id = "switchyard/random" -type = "random" -targets = ["other", "strong"] +type = "passthrough" +target = "weak" "#, base_url = upstream.base_url ))?; @@ -1743,20 +1744,29 @@ targets = ["other", "strong"] let response = send( &app, "POST", - "/v1/messages/count_tokens", + "/future/provider/endpoint?mode=raw", Some(json!({ - "model": "switchyard/random", - "messages": [{"role": "user", "content": "hi"}] + "model": "provider/model", + "provider_field": {"nested": true} })), ) .await?; assert_eq!(response.status, StatusCode::OK); assert_eq!(response.json()?["input_tokens"], 7); - let calls = upstream.calls.lock().await; - assert_eq!(calls.len(), 1); - // The inbound route name is rewritten to the real upstream model. - assert_eq!(calls[0]["model"], "real/opus"); + assert_eq!( + calls.as_slice(), + &[json!({ + "model": "provider/model", + "provider_field": {"nested": true} + })] + ); + drop(calls); + + let state = random_state(&upstream.base_url, &[(ROUTE_MODEL, &["model/weak"])])?; + let app = build_switchyard_router(state); + let response = send(&app, "POST", "/future/provider/endpoint", None).await?; + assert_eq!(response.status, StatusCode::NOT_FOUND); Ok(()) } @@ -1896,55 +1906,6 @@ target = "openai" Ok(()) } -#[tokio::test] -async fn count_tokens_without_anthropic_target_returns_bad_request() -> TestResult { - let upstream = MockUpstream::start().await?; - let state = load_test_config(&format!( - r#" -schema_version = 1 - -[llm_clients.upstream] -format = "openai_chat" -base_url = "{base_url}" - -[targets.weak] -id = "model/weak" -llm_client = "upstream" - -[routes.random] -id = "switchyard/random" -type = "random" -targets = ["weak"] -"#, - base_url = upstream.base_url - ))?; - let app = build_switchyard_router(state); - - let response = send( - &app, - "POST", - "/v1/messages/count_tokens", - Some(json!({ - "model": "switchyard/random", - "messages": [{"role": "user", "content": "hi"}] - })), - ) - .await?; - assert_eq!(response.status, StatusCode::BAD_REQUEST); - // This route has no Anthropic-format target. - assert_eq!( - response.json()?, - json!({ - "type": "error", - "error": { - "type": "invalid_request_error", - "message": "route has no Anthropic target for token counting" - } - }) - ); - Ok(()) -} - #[tokio::test] async fn routes_dispatch_and_discovery_endpoints_are_stable() -> TestResult { let (upstream, app) = test_app(&[ @@ -3069,55 +3030,6 @@ async fn advisor_route_routing_log_records_classifier_tier() -> TestResult { Ok(()) } -#[tokio::test] -async fn advisor_route_count_tokens_uses_executor() -> TestResult { - let upstream = MockUpstream::start().await?; - let state = load_test_config(&format!( - r#" -schema_version = 1 - -[llm_clients.claude] -format = "anthropic_messages" -base_url = "{base_url}" - -[targets.executor] -id = "model/executor" -llm_client = "claude" - -[targets.advisor] -id = "model/advisor" -llm_client = "claude" - -[routes.gated] -id = "switchyard/advisor" -type = "advisor" -executor_target = "executor" -advisor_target = "advisor" -"#, - base_url = upstream.base_url - ))?; - let app = build_switchyard_router(state); - - let response = send( - &app, - "POST", - "/v1/messages/count_tokens", - Some(json!({ - "model": "switchyard/advisor", - "messages": [{"role": "user", "content": "hi"}] - })), - ) - .await?; - assert_eq!(response.status, StatusCode::OK); - assert_eq!(response.json()?["input_tokens"], 7); - // The executor is the route's only completion target, so it backs - // count_tokens; the judge-only advisor never does. - let calls = upstream.calls.lock().await; - assert_eq!(calls.len(), 1); - assert_eq!(calls[0]["model"], "model/executor"); - Ok(()) -} - /// An advisor deployment whose reviewer client never retries, so a down /// advisor hits fail-open after a single attempt (the documented deployment /// posture for the advisor tier). diff --git a/docs/reference/toml_schema.md b/docs/reference/toml_schema.md index f4f2f7742..430d0f1fe 100644 --- a/docs/reference/toml_schema.md +++ b/docs/reference/toml_schema.md @@ -32,6 +32,11 @@ target = "strong" `schema_version` must be `1`. Table names under `llm_clients`, `targets`, and `routes` are local references; clients send the route's `id` as the model name. +The optional top-level `fallback_client` names an entry under `[llm_clients]`. When set, any HTTP +method and path not implemented by Switchyard is forwarded through that client without translating +the path, query string, body, response, or model identifier. The fallback client does not need a +target. When omitted, unmatched paths return `404`. + `schema_version`, `[targets]`, and `[routes]` must all be present, even when a route reaches no upstream. A file without a `[targets]` table is rejected with `missing field targets`; an empty `[targets]` table satisfies it. From ae4fc0f7234ad486a1ff0ea2863efe11e8048553 Mon Sep 17 00:00:00 2001 From: nachiketb Date: Fri, 28 Aug 2026 13:24:13 -0700 Subject: [PATCH 2/5] refactor(server): isolate fallback proxy Signed-off-by: nachiketb --- crates/libsy-llm-client/src/backend.rs | 15 ----- crates/libsy-llm-client/src/client.rs | 35 +---------- crates/switchyard-runner/src/config.rs | 19 ++---- crates/switchyard-runner/src/runner.rs | 44 +++----------- crates/switchyard-server/README.md | 4 +- crates/switchyard-server/src/lib.rs | 77 +++++++++++++++++++++--- crates/switchyard-server/tests/server.rs | 55 +++++++++++++++-- docs/reference/toml_schema.md | 4 +- 8 files changed, 138 insertions(+), 115 deletions(-) diff --git a/crates/libsy-llm-client/src/backend.rs b/crates/libsy-llm-client/src/backend.rs index 15ff205f1..578a2275d 100644 --- a/crates/libsy-llm-client/src/backend.rs +++ b/crates/libsy-llm-client/src/backend.rs @@ -255,21 +255,6 @@ impl Backend { self.config().max_retries } - /// Resolves an unmatched provider path against this backend's API root. - pub(crate) fn forwarding_url(&self, path_and_query: &str) -> String { - let base_url = self.config().base_url.trim_end_matches('/'); - let root = [ - "/v1/chat/completions", - "/v1/responses", - "/v1/messages", - "/v1", - ] - .iter() - .find_map(|suffix| base_url.strip_suffix(suffix)) - .unwrap_or(base_url); - format!("{root}{path_and_query}") - } - /// Whether an upstream 400 `body` looks like a context-window overflow for /// this backend's provider. pub(crate) fn is_context_overflow(&self, body: &str) -> bool { diff --git a/crates/libsy-llm-client/src/client.rs b/crates/libsy-llm-client/src/client.rs index 3b10dc32a..b8022a3db 100644 --- a/crates/libsy-llm-client/src/client.rs +++ b/crates/libsy-llm-client/src/client.rs @@ -10,9 +10,9 @@ use std::time::{Duration, SystemTime}; use async_trait::async_trait; use futures_util::{StreamExt, stream}; -use http::{Method, StatusCode}; +use http::StatusCode; use reqwest::RequestBuilder; -use reqwest::header::{CONTENT_LENGTH, CONTENT_TYPE, HeaderMap, RETRY_AFTER}; +use reqwest::header::{HeaderMap, RETRY_AFTER}; use serde_json::{Map, Value}; use switchyard_protocol::{ LlmRequest, LlmResponse, LlmResponseChunk, LlmResponseStreamEvent, Metadata, ModelId, Request, @@ -146,37 +146,6 @@ impl TranslatingLlmClient { }) } - /// Forwards a provider-native request through `backend` without translation. - pub async fn forward( - &self, - backend: &Backend, - method: Method, - path_and_query: &str, - body: reqwest::Body, - metadata: Option<&Metadata>, - ) -> Result { - let client = if backend.is_forwarding_auth() { - &self.forward_auth_client - } else { - &self.client - }; - let mut builder = client - .request(method, backend.forwarding_url(path_and_query)) - .body(body); - builder = forward_metadata_headers(builder, metadata); - if let Some(headers) = metadata.and_then(|metadata| metadata.http_headers.as_ref()) { - for name in [CONTENT_TYPE, CONTENT_LENGTH] { - if let Some(value) = headers.get(&name) { - builder = builder.header(name, value); - } - } - } - builder = backend.apply_forwarded_auth(builder, metadata); - builder = apply_extra_headers(builder, backend); - builder = backend.apply_auth(builder); - builder.send().await.map_err(convert_reqwest_error) - } - /// Encode `llm_request` for `wire_format`, POST it to `url` with the request's /// forwarded headers plus the backend's static headers and auth, and return the /// successful upstream response. A diff --git a/crates/switchyard-runner/src/config.rs b/crates/switchyard-runner/src/config.rs index 12b35d02e..e156862df 100644 --- a/crates/switchyard-runner/src/config.rs +++ b/crates/switchyard-runner/src/config.rs @@ -17,7 +17,6 @@ use switchyard_llm_client::{ }; use switchyard_protocol::{ModelId, RoutedLlmClient, WireFormat}; -use crate::runner::FallbackClient; use crate::{ AlgorithmSpec, CallerAuthKind, DecisionTarget, ModelCapabilities, Route, Runner, RunnerError, }; @@ -166,7 +165,7 @@ impl DeploymentConfig { let clients = self.build_clients()?; let targets = self.build_targets(); - let fallback_client = self.build_fallback_client(&clients)?; + let fallback_base_url = self.fallback_base_url()?; let mut routes = Vec::with_capacity(self.routes.len()); for (route_name, config) in &self.routes { validate_value("route name", route_name)?; @@ -204,7 +203,8 @@ impl DeploymentConfig { ); routes.push((config.id.clone(), route)); } - Ok(Runner::new(routes).with_fallback_client(fallback_client)) + let runner = Runner::new(routes).with_fallback_url(fallback_base_url); + Ok(runner) } fn build_clients(&self) -> RunnerResult>> { @@ -291,10 +291,7 @@ impl DeploymentConfig { Ok((ClientRouter::new(by_model), caller_auth)) } - fn build_fallback_client( - &self, - clients: &BTreeMap>, - ) -> RunnerResult> { + fn fallback_base_url(&self) -> RunnerResult> { let Some(name) = &self.fallback_client else { return Ok(None); }; @@ -303,13 +300,7 @@ impl DeploymentConfig { "fallback_client references unknown llm client {name}" )) })?; - let client = clients.get(name).ok_or_else(|| { - RunnerError::configuration("validated fallback client was not initialized") - })?; - Ok(Some(FallbackClient { - backend: build_backend(name, config, &BTreeMap::new())?, - client: client.clone(), - })) + Ok(Some(config.base_url.as_str().to_string())) } } diff --git a/crates/switchyard-runner/src/runner.rs b/crates/switchyard-runner/src/runner.rs index c4055cf4b..bcd175344 100644 --- a/crates/switchyard-runner/src/runner.rs +++ b/crates/switchyard-runner/src/runner.rs @@ -5,13 +5,10 @@ use std::collections::BTreeMap; use std::path::Path; -use std::sync::Arc; use libsy::RoutingOutcome; -use reqwest::{Body, Method, Response}; use serde_json::Value; -use switchyard_llm_client::{Backend, TranslatingLlmClient}; -use switchyard_protocol::{Metadata, ModelId, WireFormat}; +use switchyard_protocol::{ModelId, WireFormat}; use crate::config; use crate::{ModelCapabilities, Route, RunnerError}; @@ -19,12 +16,7 @@ use crate::{ModelCapabilities, Route, RunnerError}; /// Immutable named route table. pub struct Runner { routes: Vec<(ModelId, Route)>, - fallback_client: Option, -} - -pub(crate) struct FallbackClient { - pub backend: Backend, - pub client: Arc, + fallback_base_url: Option, } /// Borrowed model metadata returned while listing routes. @@ -68,12 +60,12 @@ impl Runner { pub fn new(routes: Vec<(ModelId, Route)>) -> Self { Self { routes, - fallback_client: None, + fallback_base_url: None, } } - pub(crate) fn with_fallback_client(mut self, fallback_client: Option) -> Self { - self.fallback_client = fallback_client; + pub(crate) fn with_fallback_url(mut self, fallback_base_url: Option) -> Self { + self.fallback_base_url = fallback_base_url; self } @@ -94,29 +86,9 @@ impl Runner { }) } - /// Forwards an unmatched request through the configured fallback client. - pub async fn forward_fallback( - &self, - method: Method, - path_and_query: &str, - body: Body, - metadata: Metadata, - ) -> Result, RunnerError> { - let Some(fallback) = &self.fallback_client else { - return Ok(None); - }; - fallback - .client - .forward( - &fallback.backend, - method, - path_and_query, - body, - Some(&metadata), - ) - .await - .map(Some) - .map_err(Into::into) + /// Returns the validated API root used for unmatched HTTP requests. + pub fn fallback_base_url(&self) -> Option<&str> { + self.fallback_base_url.as_deref() } /// Resolves an outcome to configured target names and non-secret client settings. diff --git a/crates/switchyard-server/README.md b/crates/switchyard-server/README.md index 7e6426260..8171dfbd2 100644 --- a/crates/switchyard-server/README.md +++ b/crates/switchyard-server/README.md @@ -143,7 +143,9 @@ route, and the server translates between them. Set the top-level `fallback_client` to an entry under `[llm_clients]` to proxy any otherwise unmatched method and path through that client. The fallback client does not need a target. Requests and responses are forwarded without translation, including paths, query strings, bodies, and model -identifiers. Without `fallback_client`, unmatched paths return `404`. +identifiers. The caller's end-to-end headers, including authorization, are forwarded; hop-by-hop +headers are removed, and the client's configured API key, extra headers, format, and retry policy +are not applied. Without `fallback_client`, unmatched paths return `404`. `POST /v1/decision` accepts `{"input_format": "openai_chat", "request": {...}}`, where the nested request names the route in `model`. It executes required classifier or judge calls, then diff --git a/crates/switchyard-server/src/lib.rs b/crates/switchyard-server/src/lib.rs index a382890d1..c2c0dce3b 100644 --- a/crates/switchyard-server/src/lib.rs +++ b/crates/switchyard-server/src/lib.rs @@ -137,6 +137,7 @@ struct DecisionLlmClientResponse { #[derive(Clone)] pub struct ServerState { runner: Arc, + fallback_http: reqwest::Client, metrics: prometheus::Registry, stats: StatsAccumulator, routing_log: Option, @@ -202,12 +203,17 @@ impl ServerState { /// Creates HTTP-server state around an already configured runner. pub fn from_runner(runner: Runner) -> ServerResult { let metrics = metrics::registry().map_err(ServerError::new)?; + let fallback_http = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .build() + .map_err(|error| ServerError::new(error.to_string()))?; let stats = StatsAccumulator::new( metrics.clone(), runner.models().map(|model| model.algorithm), ); Ok(Self { runner: Arc::new(runner), + fallback_http, metrics, stats, routing_log: None, @@ -538,26 +544,79 @@ async fn openai_responses( handle_endpoint(state, started, headers, body, WireFormat::OpenAiResponses).await } -// Forwards unmatched requests through the deployment's optional fallback client. +// Forwards unmatched requests unchanged to the configured API root. async fn proxy_unmatched(State(state): State, request: HttpRequest) -> Response { - let (parts, body) = request.into_parts(); + let Some(base_url) = state.runner.fallback_base_url() else { + return not_found().await; + }; + let (mut parts, body) = request.into_parts(); let path_and_query = parts .uri .path_and_query() .map_or_else(|| parts.uri.path().to_string(), ToString::to_string); - let metadata = metadata_from_headers(parts.headers); + strip_hop_by_hop_headers(&mut parts.headers); + parts.headers.remove(axum::http::header::HOST); let body = reqwest::Body::wrap_stream(body.into_data_stream()); match state - .runner - .forward_fallback(parts.method, &path_and_query, body, metadata) + .fallback_http + .request(parts.method, fallback_url(base_url, &path_and_query)) + .headers(parts.headers) + .body(body) + .send() .await { - Ok(Some(response)) => { + Ok(response) => { let response: http::Response = response.into(); - response.map(Body::new) + let (mut parts, body) = response.into_parts(); + strip_hop_by_hop_headers(&mut parts.headers); + Response::from_parts(parts, Body::new(body)) } - Ok(None) => not_found().await, - Err(error) => runner_error(error), + Err(error) => error_response( + StatusCode::BAD_GATEWAY, + error.to_string(), + "upstream_error", + "upstream_error", + ), + } +} + +fn fallback_url(base_url: &str, path_and_query: &str) -> String { + let base_url = base_url.trim_end_matches('/'); + let root = [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages", + "/v1", + ] + .iter() + .find_map(|suffix| base_url.strip_suffix(suffix)) + .unwrap_or(base_url); + format!("{root}{path_and_query}") +} + +fn strip_hop_by_hop_headers(headers: &mut HeaderMap) { + let connection_headers = headers + .get_all(axum::http::header::CONNECTION) + .iter() + .filter_map(|value| value.to_str().ok()) + .flat_map(|value| value.split(',')) + .filter_map(|name| HeaderName::from_bytes(name.trim().as_bytes()).ok()) + .collect::>(); + for name in connection_headers { + headers.remove(name); + } + for name in [ + "connection", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "proxy-connection", + "te", + "trailer", + "transfer-encoding", + "upgrade", + ] { + headers.remove(name); } } diff --git a/crates/switchyard-server/tests/server.rs b/crates/switchyard-server/tests/server.rs index 0343ca820..0d13a8407 100644 --- a/crates/switchyard-server/tests/server.rs +++ b/crates/switchyard-server/tests/server.rs @@ -11,7 +11,7 @@ use std::sync::Arc; use axum::body::{Body, Bytes}; use axum::extract::{DefaultBodyLimit, State}; -use axum::http::{HeaderMap, Request as HttpRequest, StatusCode}; +use axum::http::{HeaderMap, HeaderValue, Request as HttpRequest, StatusCode}; use axum::response::sse::{Event, Sse}; use axum::response::{IntoResponse, Response as HttpResponse}; use axum::routing::post; @@ -424,10 +424,29 @@ async fn upstream_redirect_capture( async fn upstream_fallback( State(calls): State>>>, + headers: HeaderMap, Json(body): Json, ) -> HttpResponse { - calls.lock().await.push(body.clone()); - Json(json!({"input_tokens": 7})).into_response() + calls.lock().await.push(json!({ + "body": body, + "authorization": headers.get("authorization").and_then(|value| value.to_str().ok()), + "end_to_end": headers.get("x-end-to-end").and_then(|value| value.to_str().ok()), + "configured_secret": headers.contains_key("x-configured-secret"), + "connection": headers.contains_key("connection"), + "connection_nominated": headers.contains_key("x-remove-me") + })); + let mut response = Json(json!({"input_tokens": 7})).into_response(); + response + .headers_mut() + .insert("connection", HeaderValue::from_static("x-upstream-hop")); + response + .headers_mut() + .insert("x-upstream-hop", HeaderValue::from_static("remove")); + response.headers_mut().insert( + "x-end-to-end-response", + HeaderValue::from_static("preserve"), + ); + response } fn random_state(base_url: &str, routes: &[(&str, &[&str])]) -> TestResult { @@ -1723,6 +1742,7 @@ fallback_client = "fallback" [llm_clients.fallback] format = "openai_responses" base_url = "{base_url}" +extra_headers = {{ "x-configured-secret" = "must-not-forward" }} [llm_clients.routed] format = "openai_chat" @@ -1741,7 +1761,7 @@ target = "weak" ))?; let app = build_switchyard_router(state); - let response = send( + let response = send_with_headers( &app, "POST", "/future/provider/endpoint?mode=raw", @@ -1749,16 +1769,39 @@ target = "weak" "model": "provider/model", "provider_field": {"nested": true} })), + &[ + ("authorization", "Bearer caller-key"), + ("connection", "keep-alive, x-remove-me"), + ("keep-alive", "timeout=5"), + ("x-remove-me", "remove"), + ("x-end-to-end", "preserve"), + ], ) .await?; assert_eq!(response.status, StatusCode::OK); assert_eq!(response.json()?["input_tokens"], 7); + assert!(!response.headers.contains_key("connection")); + assert!(!response.headers.contains_key("x-upstream-hop")); + assert_eq!( + response + .headers + .get("x-end-to-end-response") + .and_then(|value| value.to_str().ok()), + Some("preserve") + ); let calls = upstream.calls.lock().await; assert_eq!( calls.as_slice(), &[json!({ - "model": "provider/model", - "provider_field": {"nested": true} + "body": { + "model": "provider/model", + "provider_field": {"nested": true} + }, + "authorization": "Bearer caller-key", + "end_to_end": "preserve", + "configured_secret": false, + "connection": false, + "connection_nominated": false })] ); drop(calls); diff --git a/docs/reference/toml_schema.md b/docs/reference/toml_schema.md index 430d0f1fe..966edd978 100644 --- a/docs/reference/toml_schema.md +++ b/docs/reference/toml_schema.md @@ -35,7 +35,9 @@ target = "strong" The optional top-level `fallback_client` names an entry under `[llm_clients]`. When set, any HTTP method and path not implemented by Switchyard is forwarded through that client without translating the path, query string, body, response, or model identifier. The fallback client does not need a -target. When omitted, unmatched paths return `404`. +target. Only its `base_url` is used: caller end-to-end headers are forwarded, hop-by-hop headers are +removed, and configured API keys, extra headers, format, and retries are not applied. When omitted, +unmatched paths return `404`. `schema_version`, `[targets]`, and `[routes]` must all be present, even when a route reaches no upstream. A file without a `[targets]` table is rejected with From 5baeb5b2a2b2776303e76b2e100a5ac024975734 Mon Sep 17 00:00:00 2001 From: nachiketb Date: Fri, 28 Aug 2026 15:21:51 -0700 Subject: [PATCH 3/5] feat(server): route model-bearing auxiliary endpoints Signed-off-by: nachiketb --- crates/libsy-llm-client/src/backend.rs | 13 ++ crates/libsy-llm-client/src/client.rs | 135 +++++++++++++++++++- crates/switchyard-runner/src/algorithm.rs | 2 +- crates/switchyard-runner/src/config.rs | 59 ++++++++- crates/switchyard-runner/src/failure.rs | 11 +- crates/switchyard-runner/src/lib.rs | 5 +- crates/switchyard-runner/src/route.rs | 66 +++++++++- crates/switchyard-runner/tests/route.rs | 2 + crates/switchyard-server/README.md | 6 + crates/switchyard-server/src/lib.rs | 131 ++++++++++++++++++- crates/switchyard-server/tests/server.rs | 147 +++++++++++++++++++++- 11 files changed, 565 insertions(+), 12 deletions(-) diff --git a/crates/libsy-llm-client/src/backend.rs b/crates/libsy-llm-client/src/backend.rs index 578a2275d..a526712fc 100644 --- a/crates/libsy-llm-client/src/backend.rs +++ b/crates/libsy-llm-client/src/backend.rs @@ -255,6 +255,19 @@ impl Backend { self.config().max_retries } + /// Whether this backend speaks the Anthropic Messages wire format — the only + /// one with a `count_tokens` endpoint. + pub fn is_anthropic(&self) -> bool { + matches!(self, Backend::Anthropic(_)) + } + + /// The upstream `/v1/messages/count_tokens` URL, derived from the same base + /// URL join as [`url`](Self::url). + pub fn count_tokens_url(&self) -> String { + let base_url = self.config().base_url.trim_end_matches('/'); + format!("{}/count_tokens", anthropic_url(base_url)) + } + /// Whether an upstream 400 `body` looks like a context-window overflow for /// this backend's provider. pub(crate) fn is_context_overflow(&self, body: &str) -> bool { diff --git a/crates/libsy-llm-client/src/client.rs b/crates/libsy-llm-client/src/client.rs index b8022a3db..a19715c38 100644 --- a/crates/libsy-llm-client/src/client.rs +++ b/crates/libsy-llm-client/src/client.rs @@ -146,6 +146,106 @@ impl TranslatingLlmClient { }) } + /// Whether `model` has an Anthropic backend that supports token counting. + pub fn supports_count_tokens(&self, model: &ModelId) -> bool { + self.backend_for(model, WireFormat::AnthropicMessages) + .is_some() + } + + /// Whether `model` has an OpenAI Responses backend for auxiliary operations. + pub fn supports_responses_auxiliary(&self, model: &ModelId) -> bool { + self.backend_for(model, WireFormat::OpenAiResponses) + .is_some() + } + + /// Counts input tokens with `model`'s Anthropic backend. + /// + /// Returns an error when the model has no Anthropic backend or the upstream + /// request fails or returns invalid JSON. + pub async fn count_tokens(&self, model: &ModelId, request: Request) -> Result { + let backend = self + .backend_for(model, WireFormat::AnthropicMessages) + .ok_or_else(|| LlmClientError::Configuration { + message: format!("model {model} has no Anthropic backend for count_tokens"), + })?; + let Request { + mut llm_request, + metadata, + .. + } = request; + llm_request.model = Some(model.to_string()); + let http_response = self + .send_encoded( + backend, + WireFormat::AnthropicMessages, + llm_request, + metadata.as_ref(), + model, + UpstreamEndpoint::CountTokens, + ) + .await?; + let body = match http_response { + EncodedResponse::Buffered { body, .. } => body, + EncodedResponse::Streaming(_) => { + return Err(LlmClientError::InvalidRequest { + message: "count_tokens does not support streaming requests".to_string(), + }); + } + }; + serde_json::from_slice(&body).map_err(|error| LlmClientError::InvalidResponse { + source: Box::new(error), + }) + } + + /// Counts input tokens with `model`'s OpenAI Responses backend. + pub async fn responses_input_tokens(&self, model: &ModelId, request: Request) -> Result { + self.responses_auxiliary(model, request, UpstreamEndpoint::ResponsesInputTokens) + .await + } + + /// Compacts a request with `model`'s OpenAI Responses backend. + pub async fn responses_compact(&self, model: &ModelId, request: Request) -> Result { + self.responses_auxiliary(model, request, UpstreamEndpoint::ResponsesCompact) + .await + } + + async fn responses_auxiliary( + &self, + model: &ModelId, + request: Request, + endpoint: UpstreamEndpoint, + ) -> Result { + let backend = self + .backend_for(model, WireFormat::OpenAiResponses) + .ok_or_else(|| LlmClientError::Configuration { + message: format!("model {model} has no OpenAI Responses backend"), + })?; + let Request { + mut llm_request, + metadata, + .. + } = request; + llm_request.model = Some(model.to_string()); + let http_response = self + .send_encoded( + backend, + WireFormat::OpenAiResponses, + llm_request, + metadata.as_ref(), + model, + endpoint, + ) + .await?; + let EncodedResponse::Buffered { body, .. } = http_response else { + return Err(LlmClientError::InvalidRequest { + message: "Responses auxiliary endpoints do not support streaming".to_string(), + }); + }; + serde_json::from_slice(&body).map_err(|error| LlmClientError::InvalidResponse { + source: Box::new(error), + }) + } + /// Encode `llm_request` for `wire_format`, POST it to `url` with the request's /// forwarded headers plus the backend's static headers and auth, and return the /// successful upstream response. A @@ -153,8 +253,9 @@ impl TranslatingLlmClient { /// response is returned as soon as its successful headers arrive. A non-success /// status maps to a typed error — a 400 is classified as a context-window /// overflow via the backend's provider rules. Shared by - /// [`call_rewrite_model`](Self::call_rewrite_model), which POSTs to the - /// backend's completion URL and decodes a response. + /// [`call_rewrite_model`](Self::call_rewrite_model) (which POSTs to the + /// backend's completion URL and decodes a response) and + /// the model-bearing auxiliary operations, which return raw JSON. async fn send_encoded( &self, backend: &Backend, @@ -162,6 +263,7 @@ impl TranslatingLlmClient { llm_request: LlmRequest, metadata: Option<&Metadata>, model: &ModelId, + endpoint: UpstreamEndpoint, ) -> Result { let mut body = encode_request(&llm_request, wire_format) .map_err(|error| LlmClientError::RequestEncoding(error.to_string()))?; @@ -182,8 +284,9 @@ impl TranslatingLlmClient { if matches!(backend, Backend::OpenAiChat(_)) { ensure_openai_stream_usage(&mut body); } - let streaming = body.get("stream").and_then(Value::as_bool).unwrap_or(false); - let url = backend.url(); + let streaming = endpoint.allows_streaming() + && body.get("stream").and_then(Value::as_bool).unwrap_or(false); + let url = endpoint.url(backend); record_gen_ai_request(&url, model, streaming); let max_retries = u64::from(backend.max_retries()); @@ -377,6 +480,7 @@ impl TranslatingLlmClient { llm_request, metadata.as_ref(), &model_id, + UpstreamEndpoint::Completion, ) .await?; @@ -509,6 +613,29 @@ impl RoutedLlmClient for TranslatingLlmClient { } } +#[derive(Clone, Copy)] +enum UpstreamEndpoint { + Completion, + CountTokens, + ResponsesInputTokens, + ResponsesCompact, +} + +impl UpstreamEndpoint { + fn url(self, backend: &Backend) -> String { + match self { + UpstreamEndpoint::Completion => backend.url(), + UpstreamEndpoint::CountTokens => backend.count_tokens_url(), + UpstreamEndpoint::ResponsesInputTokens => format!("{}/input_tokens", backend.url()), + UpstreamEndpoint::ResponsesCompact => format!("{}/compact", backend.url()), + } + } + + fn allows_streaming(self) -> bool { + matches!(self, UpstreamEndpoint::Completion) + } +} + enum EncodedResponse { Buffered { status: u16, body: Vec }, Streaming(reqwest::Response), diff --git a/crates/switchyard-runner/src/algorithm.rs b/crates/switchyard-runner/src/algorithm.rs index 243c56b60..a601ab94a 100644 --- a/crates/switchyard-runner/src/algorithm.rs +++ b/crates/switchyard-runner/src/algorithm.rs @@ -467,7 +467,7 @@ impl AlgorithmSpec { names } // The advisor is judge-only: reviews go through its own client, - // so it is not a completion destination. + // so it is not a completion (or count_tokens) destination. Self::Advisor { executor_target, .. } => vec![executor_target], diff --git a/crates/switchyard-runner/src/config.rs b/crates/switchyard-runner/src/config.rs index e156862df..ddf8a9bf3 100644 --- a/crates/switchyard-runner/src/config.rs +++ b/crates/switchyard-runner/src/config.rs @@ -18,7 +18,8 @@ use switchyard_llm_client::{ use switchyard_protocol::{ModelId, RoutedLlmClient, WireFormat}; use crate::{ - AlgorithmSpec, CallerAuthKind, DecisionTarget, ModelCapabilities, Route, Runner, RunnerError, + AlgorithmSpec, CallerAuthKind, CountTokensTarget, DecisionTarget, ModelCapabilities, + ResponsesTarget, Route, Runner, RunnerError, }; const SUPPORTED_SCHEMA_VERSION: u32 = 1; @@ -189,6 +190,8 @@ impl DeploymentConfig { .map_err(|error| RunnerError::configuration_source(error.to_string(), error))?; let (route_clients, caller_auth) = self.build_route_clients(route_name, config, &clients)?; + let count_tokens_target = self.build_count_tokens_target(config, &clients); + let responses_target = self.build_responses_target(config, &clients); let decision_targets = config .routing_target_names() .into_iter() @@ -199,6 +202,8 @@ impl DeploymentConfig { route_clients, caller_auth, capabilities, + count_tokens_target, + responses_target, decision_targets, ); routes.push((config.id.clone(), route)); @@ -302,6 +307,58 @@ impl DeploymentConfig { })?; Ok(Some(config.base_url.as_str().to_string())) } + + fn build_count_tokens_target( + &self, + route: &RouteConfig, + clients: &BTreeMap>, + ) -> Option { + route + .routing_target_names() + .into_iter() + .enumerate() + .filter_map(|(index, name)| { + let target = self.targets.get(name)?; + let client = clients.get(&target.llm_client)?; + client.supports_count_tokens(&target.id).then_some(( + count_tokens_priority(name, &target.id), + index, + target, + client, + )) + }) + .min_by_key(|(priority, index, _, _)| (*priority, *index)) + .map(|(_, _, target, client)| CountTokensTarget { + model: target.id.clone(), + client: client.clone(), + }) + } + + fn build_responses_target( + &self, + route: &RouteConfig, + clients: &BTreeMap>, + ) -> Option { + route.routing_target_names().into_iter().find_map(|name| { + let target = self.targets.get(name)?; + let client = clients.get(&target.llm_client)?; + client + .supports_responses_auxiliary(&target.id) + .then(|| ResponsesTarget { + model: target.id.clone(), + client: client.clone(), + }) + }) + } +} + +fn count_tokens_priority(target_name: &str, model_id: &ModelId) -> usize { + let target_name = target_name.to_ascii_lowercase(); + let model_id = model_id.to_ascii_lowercase(); + ["opus", "sonnet", "haiku"] + .iter() + .position(|hint| target_name.contains(hint) || model_id.contains(hint)) + .unwrap_or(3) } /// A client endpoint, parsed when the config loads rather than checked afterwards. diff --git a/crates/switchyard-runner/src/failure.rs b/crates/switchyard-runner/src/failure.rs index 7dd470181..598d8c31b 100644 --- a/crates/switchyard-runner/src/failure.rs +++ b/crates/switchyard-runner/src/failure.rs @@ -98,7 +98,10 @@ impl RunnerError { None, None, ), - Self::UnknownRouteModel(_) | Self::IncompatibleCallerFormat(_) => summary( + Self::UnknownRouteModel(_) + | Self::IncompatibleCallerFormat(_) + | Self::CountTokensUnsupported + | Self::ResponsesAuxiliaryUnsupported => summary( RouteErrorKind::InvalidRequest, RouteErrorPhase::BeforeResponse, None, @@ -326,5 +329,11 @@ mod tests { configuration.execution_error_summary().kind, RouteErrorKind::Configuration )); + + let unsupported = RunnerError::CountTokensUnsupported; + assert!(matches!( + unsupported.execution_error_summary().kind, + RouteErrorKind::InvalidRequest + )); } } diff --git a/crates/switchyard-runner/src/lib.rs b/crates/switchyard-runner/src/lib.rs index 850e513dc..5e0cbb620 100644 --- a/crates/switchyard-runner/src/lib.rs +++ b/crates/switchyard-runner/src/lib.rs @@ -14,5 +14,8 @@ pub use algorithm::{ ClassifierPolicyConfig, LlmClassifierRouteConfig, StageClassifierConfig, SubagentRouteConfig, }; pub use failure::{RouteErrorKind, RouteErrorPhase, RouteErrorSummary, stream_error_summary}; -pub use route::{CallerAuthKind, ModelCapabilities, Route, RunOutput, RunnerError}; +pub use route::{ + CallerAuthKind, CountTokensTarget, ModelCapabilities, ResponsesTarget, Route, RunOutput, + RunnerError, +}; pub use runner::{DecisionDescription, DecisionTarget, ModelInfo, Runner}; diff --git a/crates/switchyard-runner/src/route.rs b/crates/switchyard-runner/src/route.rs index 451f365c1..4e5b683a1 100644 --- a/crates/switchyard-runner/src/route.rs +++ b/crates/switchyard-runner/src/route.rs @@ -7,7 +7,8 @@ use std::error::Error; use std::sync::Arc; use libsy::{Algorithm, CallModel, LibsyError, RoutingOutcome, drive}; -use switchyard_llm_client::{ClientRouter, RunObserver}; +use serde_json::Value; +use switchyard_llm_client::{ClientRouter, RunObserver, TranslatingLlmClient}; use switchyard_protocol::{LlmClientError, ModelId, Request, Response, WireFormat}; use thiserror::Error; @@ -55,6 +56,20 @@ impl CallerAuthKind { } } +/// Exact upstream model used for Anthropic token counting. +#[derive(Clone)] +pub struct CountTokensTarget { + pub model: ModelId, + pub client: Arc, +} + +/// Exact upstream model used for OpenAI Responses auxiliary operations. +#[derive(Clone)] +pub struct ResponsesTarget { + pub model: ModelId, + pub client: Arc, +} + /// Error returned while loading or executing configured routes. #[derive(Debug, Error)] pub enum RunnerError { @@ -68,6 +83,10 @@ pub enum RunnerError { UnknownRouteModel(String), #[error("caller format is incompatible with {} credentials", .0.as_str())] IncompatibleCallerFormat(CallerAuthKind), + #[error("route has no Anthropic target for token counting")] + CountTokensUnsupported, + #[error("route has no OpenAI Responses target")] + ResponsesAuxiliaryUnsupported, #[error(transparent)] Algorithm(#[from] LibsyError), #[error(transparent)] @@ -102,6 +121,8 @@ pub struct Route { clients: ClientRouter, caller_auth: Option, capabilities: ModelCapabilities, + count_tokens_target: Option, + responses_target: Option, decision_targets: Vec, } @@ -118,6 +139,8 @@ impl Route { clients: ClientRouter, caller_auth: Option, capabilities: ModelCapabilities, + count_tokens_target: Option, + responses_target: Option, decision_targets: Vec, ) -> Self { Self { @@ -125,6 +148,8 @@ impl Route { clients, caller_auth, capabilities, + count_tokens_target, + responses_target, decision_targets, } } @@ -189,6 +214,45 @@ impl Route { .await .map_err(Into::into) } + + /// Counts tokens using the configured Anthropic-capable target. + pub async fn count_tokens(&self, request: Request) -> Result { + let target = self + .count_tokens_target + .as_ref() + .ok_or(RunnerError::CountTokensUnsupported)?; + target + .client + .count_tokens(&target.model, request) + .await + .map_err(Into::into) + } + + /// Counts tokens using the configured OpenAI Responses-capable target. + pub async fn responses_input_tokens(&self, request: Request) -> Result { + let target = self + .responses_target + .as_ref() + .ok_or(RunnerError::ResponsesAuxiliaryUnsupported)?; + target + .client + .responses_input_tokens(&target.model, request) + .await + .map_err(Into::into) + } + + /// Compacts a request using the configured OpenAI Responses-capable target. + pub async fn responses_compact(&self, request: Request) -> Result { + let target = self + .responses_target + .as_ref() + .ok_or(RunnerError::ResponsesAuxiliaryUnsupported)?; + target + .client + .responses_compact(&target.model, request) + .await + .map_err(Into::into) + } } async fn serve_decision_dependency(clients: ClientRouter, call: CallModel) -> libsy::Result<()> { diff --git a/crates/switchyard-runner/tests/route.rs b/crates/switchyard-runner/tests/route.rs index 9442004ba..b6b1d1827 100644 --- a/crates/switchyard-runner/tests/route.rs +++ b/crates/switchyard-runner/tests/route.rs @@ -51,6 +51,8 @@ fn plugin_route(client: Arc) -> Route { clients, None, ModelCapabilities::default(), + None, + None, Vec::new(), ) } diff --git a/crates/switchyard-server/README.md b/crates/switchyard-server/README.md index 8171dfbd2..3f59fd0af 100644 --- a/crates/switchyard-server/README.md +++ b/crates/switchyard-server/README.md @@ -129,6 +129,9 @@ documented in [Stage-Router Routing](../../docs/routing_algorithms/stage_router_ | `POST` | `/v1/messages` | Anthropic Messages | | `POST` | `/v1/responses` | OpenAI Responses | | `POST` | `/v1/decision` | Resolve selected and fallback targets without a post-routing answer call | +| `POST` | `/v1/messages/count_tokens` | Token count from a route's Anthropic target | +| `POST` | `/v1/responses/input_tokens` | Token count from a route's OpenAI Responses target | +| `POST` | `/v1/responses/compact` | Compaction through a route's OpenAI Responses target | | `ANY` | Any unmatched path | Raw forward through the optional `fallback_client` | | `GET` | `/v1/models` | Routes served by this deployment | | `GET` | `/v1/stats` | Per-model usage plus curated algorithm stats | @@ -147,6 +150,9 @@ identifiers. The caller's end-to-end headers, including authorization, are forwa headers are removed, and the client's configured API key, extra headers, format, and retry policy are not applied. Without `fallback_client`, unmatched paths return `404`. +The three model-bearing auxiliary endpoints above remain routed operations: they resolve the route +name to a compatible target, rewrite the upstream model ID, and use that target's configured client. + `POST /v1/decision` accepts `{"input_format": "openai_chat", "request": {...}}`, where the nested request names the route in `model`. It executes required classifier or judge calls, then returns the selected target and ordered fallbacks with their model, format, base URL, and diff --git a/crates/switchyard-server/src/lib.rs b/crates/switchyard-server/src/lib.rs index c2c0dce3b..120d1ee58 100644 --- a/crates/switchyard-server/src/lib.rs +++ b/crates/switchyard-server/src/lib.rs @@ -191,6 +191,8 @@ impl ServerState { clients, None, ModelCapabilities::default(), + None, + None, Vec::new(), ), ) @@ -484,6 +486,12 @@ pub fn build_switchyard_router(state: ServerState) -> Router { .route("/v1/messages", post(anthropic_messages)) .route("/v1/responses", post(openai_responses)) .route("/v1/decision", post(decision)) + .route("/v1/messages/count_tokens", post(anthropic_count_tokens)) + .route( + "/v1/responses/input_tokens", + post(openai_responses_input_tokens), + ) + .route("/v1/responses/compact", post(openai_responses_compact)) .route("/v1/models", get(models)) .route("/v1/stats", get(get_stats)) .route("/v1/stats/reset", post(reset_stats)) @@ -694,6 +702,118 @@ async fn decision( } } +/// Anthropic token counting against the route's explicitly configured target. +async fn anthropic_count_tokens( + State(state): State, + headers: HeaderMap, + body: std::result::Result, JsonRejection>, +) -> Response { + let body = match llm_json_body(body) { + Ok(body) => body, + Err((status, message)) => { + return anthropic_error_response(invalid_body_error(status, message)); + } + }; + let (route, request) = match resolve_route( + &state, + metadata_from_headers(headers), + body, + WireFormat::AnthropicMessages, + ) { + Ok(resolved) => resolved, + Err(response) => return anthropic_error_response(response), + }; + anthropic_error_response(match route.count_tokens(request).await { + Ok(payload) => (StatusCode::OK, Json(payload)).into_response(), + Err(RunnerError::CountTokensUnsupported) => error_response( + StatusCode::BAD_REQUEST, + "route has no Anthropic target for token counting", + "invalid_request_error", + "count_tokens_unsupported", + ), + Err(RunnerError::Client(error)) => count_tokens_error(error), + Err(error) => server_error(error.to_string()), + }) +} + +/// Maps a token-count client failure with the same policy as a routed client call. +fn count_tokens_error(error: LlmClientError) -> Response { + client_error(&error) +} + +#[derive(Clone, Copy)] +enum ResponsesAuxiliaryEndpoint { + InputTokens, + Compact, +} + +async fn openai_responses_input_tokens( + State(state): State, + headers: HeaderMap, + body: std::result::Result, JsonRejection>, +) -> Response { + openai_responses_auxiliary( + state, + headers, + body, + ResponsesAuxiliaryEndpoint::InputTokens, + ) + .await +} + +async fn openai_responses_compact( + State(state): State, + headers: HeaderMap, + body: std::result::Result, JsonRejection>, +) -> Response { + openai_responses_auxiliary(state, headers, body, ResponsesAuxiliaryEndpoint::Compact).await +} + +// Resolves route aliases and configured credentials before calling a model-bearing Responses API. +async fn openai_responses_auxiliary( + state: ServerState, + headers: HeaderMap, + body: std::result::Result, JsonRejection>, + endpoint: ResponsesAuxiliaryEndpoint, +) -> Response { + let body = match llm_json_body(body) { + Ok(body) => body, + Err((status, message)) => { + return render_error_response( + invalid_body_error(status, message), + WireFormat::OpenAiResponses, + ); + } + }; + let (route, request) = match resolve_route( + &state, + metadata_from_headers(headers), + body, + WireFormat::OpenAiResponses, + ) { + Ok(resolved) => resolved, + Err(response) => return render_error_response(response, WireFormat::OpenAiResponses), + }; + let result = match endpoint { + ResponsesAuxiliaryEndpoint::InputTokens => route.responses_input_tokens(request).await, + ResponsesAuxiliaryEndpoint::Compact => route.responses_compact(request).await, + }; + render_error_response( + match result { + Ok(payload) => (StatusCode::OK, Json(payload)).into_response(), + Err(RunnerError::ResponsesAuxiliaryUnsupported) => error_response( + StatusCode::BAD_REQUEST, + "route has no OpenAI Responses target", + "invalid_request_error", + "responses_auxiliary_unsupported", + ), + Err(RunnerError::Client(error)) => client_error(&error), + Err(error) => server_error(error.to_string()), + }, + WireFormat::OpenAiResponses, + ) +} + async fn handle_endpoint( state: ServerState, started: RequestStart, @@ -781,8 +901,8 @@ fn llm_json_body( } /// Decode `body`, resolve the route named by its `model`, and build the -/// [`Request`]. Shared by the completion handlers. Returns the resolved route -/// and the built request — or an error [`Response`] +/// [`Request`]. Shared by the completion and `count_tokens` handlers. Returns +/// the resolved route and the built request — or an error [`Response`] /// (invalid body, empty `model` → 400, unknown route → 404). // Both callers immediately return the `Err(Response)` as the HTTP response, so // the large error type is intentional, not propagated up a call stack. @@ -1170,6 +1290,10 @@ fn render_error_response(response: Response, wire_format: WireFormat) -> Respons error.into_response(wire_format) } +fn anthropic_error_response(response: Response) -> Response { + render_error_response(response, WireFormat::AnthropicMessages) +} + fn anthropic_error_type(status: StatusCode) -> &'static str { match status { StatusCode::BAD_REQUEST => "invalid_request_error", @@ -1487,6 +1611,9 @@ fn endpoint_listing(has_routing_log: bool) -> String { " POST /v1/chat/completions OpenAI Chat Completions", " POST /v1/messages Anthropic Messages", " POST /v1/responses OpenAI Responses", + " POST /v1/messages/count_tokens", + " POST /v1/responses/input_tokens", + " POST /v1/responses/compact", " ANY unmatched paths optional fallback client", " GET /v1/models configured routes", " GET /v1/stats routing stats", diff --git a/crates/switchyard-server/tests/server.rs b/crates/switchyard-server/tests/server.rs index 0d13a8407..1da34ccc7 100644 --- a/crates/switchyard-server/tests/server.rs +++ b/crates/switchyard-server/tests/server.rs @@ -11,7 +11,7 @@ use std::sync::Arc; use axum::body::{Body, Bytes}; use axum::extract::{DefaultBodyLimit, State}; -use axum::http::{HeaderMap, HeaderValue, Request as HttpRequest, StatusCode}; +use axum::http::{HeaderMap, HeaderValue, Request as HttpRequest, StatusCode, Uri}; use axum::response::sse::{Event, Sse}; use axum::response::{IntoResponse, Response as HttpResponse}; use axum::routing::post; @@ -57,6 +57,12 @@ impl MockUpstream { post(upstream_responses_requires_forwarded_auth), ) .route("/capture", post(upstream_redirect_capture)) + .route("/v1/messages/count_tokens", post(upstream_count_tokens)) + .route( + "/v1/responses/input_tokens", + post(upstream_responses_auxiliary), + ) + .route("/v1/responses/compact", post(upstream_responses_auxiliary)) .route("/future/provider/endpoint", post(upstream_fallback)) .layer(DefaultBodyLimit::disable()) .with_state(Arc::clone(&calls)); @@ -422,6 +428,32 @@ async fn upstream_redirect_capture( StatusCode::OK.into_response() } +async fn upstream_count_tokens( + State(calls): State>>>, + Json(body): Json, +) -> HttpResponse { + calls.lock().await.push(body.clone()); + Json(json!({"input_tokens": 7})).into_response() +} + +async fn upstream_responses_auxiliary( + State(calls): State>>>, + uri: Uri, + headers: HeaderMap, + Json(body): Json, +) -> HttpResponse { + calls.lock().await.push(json!({ + "path": uri.path(), + "body": body, + "configured_header": headers.get("x-configured-client").and_then(|value| value.to_str().ok()) + })); + if uri.path().ends_with("/input_tokens") { + Json(json!({"input_tokens": 11})).into_response() + } else { + Json(json!({"id": "resp_compacted", "object": "response", "output": []})).into_response() + } +} + async fn upstream_fallback( State(calls): State>>>, headers: HeaderMap, @@ -1731,6 +1763,119 @@ response_format_type = "json_object" Ok(()) } +#[tokio::test] +async fn model_bearing_auxiliary_endpoints_use_configured_route_targets() -> TestResult { + let upstream = MockUpstream::start().await?; + let state = load_test_config(&format!( + r#" +schema_version = 1 + +[llm_clients.claude] +format = "anthropic_messages" +base_url = "{base_url}" + +[llm_clients.responses] +format = "openai_responses" +base_url = "{base_url}" +extra_headers = {{ "x-configured-client" = "responses" }} + +[targets.responses] +id = "real/responses-model" +llm_client = "responses" + +[targets.strong] +id = "real/opus" +llm_client = "claude" + +[targets.other] +id = "real/sonnet" +llm_client = "claude" + +[routes.random] +id = "switchyard/random" +type = "random" +targets = ["responses", "other", "strong"] +"#, + base_url = upstream.base_url + ))?; + let app = build_switchyard_router(state); + + let count_tokens = send( + &app, + "POST", + "/v1/messages/count_tokens", + Some(json!({ + "model": "switchyard/random", + "messages": [{"role": "user", "content": "hi"}] + })), + ) + .await?; + assert_eq!(count_tokens.status, StatusCode::OK); + assert_eq!(count_tokens.json()?["input_tokens"], 7); + + let input_tokens = send( + &app, + "POST", + "/v1/responses/input_tokens", + Some(json!({"model": "switchyard/random", "input": "count me"})), + ) + .await?; + assert_eq!(input_tokens.status, StatusCode::OK); + assert_eq!(input_tokens.json()?["input_tokens"], 11); + + let compact = send( + &app, + "POST", + "/v1/responses/compact", + Some(json!({"model": "switchyard/random", "input": "compact me"})), + ) + .await?; + assert_eq!(compact.status, StatusCode::OK); + assert_eq!(compact.json()?["id"], "resp_compacted"); + + let calls = upstream.calls.lock().await; + assert_eq!(calls.len(), 3); + assert_eq!(calls[0]["model"], "real/opus"); + assert_eq!( + calls[1], + json!({ + "path": "/v1/responses/input_tokens", + "body": {"model": "real/responses-model", "input": "count me"}, + "configured_header": "responses" + }) + ); + assert_eq!( + calls[2], + json!({ + "path": "/v1/responses/compact", + "body": {"model": "real/responses-model", "input": "compact me"}, + "configured_header": "responses" + }) + ); + drop(calls); + + let unsupported = random_state(&upstream.base_url, &[(ROUTE_MODEL, &["model/weak"])])?; + let unsupported = build_switchyard_router(unsupported); + for (path, body) in [ + ( + "/v1/messages/count_tokens", + json!({"model": ROUTE_MODEL, "messages": [{"role": "user", "content": "hi"}]}), + ), + ( + "/v1/responses/input_tokens", + json!({"model": ROUTE_MODEL, "input": "count me"}), + ), + ( + "/v1/responses/compact", + json!({"model": ROUTE_MODEL, "input": "compact me"}), + ), + ] { + let response = send(&unsupported, "POST", path, Some(body)).await?; + assert_eq!(response.status, StatusCode::BAD_REQUEST, "{path}"); + } + Ok(()) +} + #[tokio::test] async fn fallback_client_forwards_unmatched_requests_and_is_optional() -> TestResult { let upstream = MockUpstream::start().await?; From b21da5f5cb2b48afe61a53b6256bf1655337c676 Mon Sep 17 00:00:00 2001 From: nachiketb Date: Fri, 28 Aug 2026 15:33:13 -0700 Subject: [PATCH 4/5] refactor(server): consolidate auxiliary operations Signed-off-by: nachiketb --- crates/libsy-llm-client/src/client.rs | 121 +++++++++--------------- crates/libsy-llm-client/src/lib.rs | 2 +- crates/switchyard-runner/src/config.rs | 70 +++++++------- crates/switchyard-runner/src/failure.rs | 5 +- crates/switchyard-runner/src/lib.rs | 3 +- crates/switchyard-runner/src/route.rs | 79 +++++----------- crates/switchyard-server/src/lib.rs | 53 +++++------ 7 files changed, 136 insertions(+), 197 deletions(-) diff --git a/crates/libsy-llm-client/src/client.rs b/crates/libsy-llm-client/src/client.rs index a19715c38..9f47828aa 100644 --- a/crates/libsy-llm-client/src/client.rs +++ b/crates/libsy-llm-client/src/client.rs @@ -82,6 +82,34 @@ impl ModelConfig { } } +/// A model-bearing provider operation outside the normal completion endpoint. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum AuxiliaryOperation { + /// Anthropic Messages input-token counting. + AnthropicCountTokens, + /// OpenAI Responses input-token counting. + ResponsesInputTokens, + /// OpenAI Responses compaction. + ResponsesCompact, +} + +impl AuxiliaryOperation { + const fn wire_format(self) -> WireFormat { + match self { + Self::AnthropicCountTokens => WireFormat::AnthropicMessages, + Self::ResponsesInputTokens | Self::ResponsesCompact => WireFormat::OpenAiResponses, + } + } + + fn url(self, backend: &Backend) -> String { + match self { + Self::AnthropicCountTokens => backend.count_tokens_url(), + Self::ResponsesInputTokens => format!("{}/input_tokens", backend.url()), + Self::ResponsesCompact => format!("{}/compact", backend.url()), + } + } +} + /// A client that dispatches neutral-IR requests to per-model HTTP backends. /// /// Construct it with a list of [`ModelConfig`]s — one per model, each naming a @@ -146,80 +174,27 @@ impl TranslatingLlmClient { }) } - /// Whether `model` has an Anthropic backend that supports token counting. - pub fn supports_count_tokens(&self, model: &ModelId) -> bool { - self.backend_for(model, WireFormat::AnthropicMessages) - .is_some() + /// Whether `model` has a backend for `operation`. + pub fn supports_auxiliary(&self, model: &ModelId, operation: AuxiliaryOperation) -> bool { + self.backend_for(model, operation.wire_format()).is_some() } - /// Whether `model` has an OpenAI Responses backend for auxiliary operations. - pub fn supports_responses_auxiliary(&self, model: &ModelId) -> bool { - self.backend_for(model, WireFormat::OpenAiResponses) - .is_some() - } - - /// Counts input tokens with `model`'s Anthropic backend. + /// Calls a model-bearing auxiliary provider operation. /// - /// Returns an error when the model has no Anthropic backend or the upstream + /// Returns an error when the model has no compatible backend or the upstream /// request fails or returns invalid JSON. - pub async fn count_tokens(&self, model: &ModelId, request: Request) -> Result { - let backend = self - .backend_for(model, WireFormat::AnthropicMessages) - .ok_or_else(|| LlmClientError::Configuration { - message: format!("model {model} has no Anthropic backend for count_tokens"), - })?; - let Request { - mut llm_request, - metadata, - .. - } = request; - llm_request.model = Some(model.to_string()); - let http_response = self - .send_encoded( - backend, - WireFormat::AnthropicMessages, - llm_request, - metadata.as_ref(), - model, - UpstreamEndpoint::CountTokens, - ) - .await?; - let body = match http_response { - EncodedResponse::Buffered { body, .. } => body, - EncodedResponse::Streaming(_) => { - return Err(LlmClientError::InvalidRequest { - message: "count_tokens does not support streaming requests".to_string(), - }); - } - }; - serde_json::from_slice(&body).map_err(|error| LlmClientError::InvalidResponse { - source: Box::new(error), - }) - } - - /// Counts input tokens with `model`'s OpenAI Responses backend. - pub async fn responses_input_tokens(&self, model: &ModelId, request: Request) -> Result { - self.responses_auxiliary(model, request, UpstreamEndpoint::ResponsesInputTokens) - .await - } - - /// Compacts a request with `model`'s OpenAI Responses backend. - pub async fn responses_compact(&self, model: &ModelId, request: Request) -> Result { - self.responses_auxiliary(model, request, UpstreamEndpoint::ResponsesCompact) - .await - } - - async fn responses_auxiliary( + pub async fn call_auxiliary( &self, model: &ModelId, request: Request, - endpoint: UpstreamEndpoint, + operation: AuxiliaryOperation, ) -> Result { - let backend = self - .backend_for(model, WireFormat::OpenAiResponses) - .ok_or_else(|| LlmClientError::Configuration { - message: format!("model {model} has no OpenAI Responses backend"), - })?; + let wire_format = operation.wire_format(); + let backend = + self.backend_for(model, wire_format) + .ok_or_else(|| LlmClientError::Configuration { + message: format!("model {model} has no backend for {operation:?}"), + })?; let Request { mut llm_request, metadata, @@ -229,16 +204,16 @@ impl TranslatingLlmClient { let http_response = self .send_encoded( backend, - WireFormat::OpenAiResponses, + wire_format, llm_request, metadata.as_ref(), model, - endpoint, + UpstreamEndpoint::Auxiliary(operation), ) .await?; let EncodedResponse::Buffered { body, .. } = http_response else { return Err(LlmClientError::InvalidRequest { - message: "Responses auxiliary endpoints do not support streaming".to_string(), + message: "auxiliary endpoints do not support streaming".to_string(), }); }; serde_json::from_slice(&body).map_err(|error| LlmClientError::InvalidResponse { @@ -616,18 +591,14 @@ impl RoutedLlmClient for TranslatingLlmClient { #[derive(Clone, Copy)] enum UpstreamEndpoint { Completion, - CountTokens, - ResponsesInputTokens, - ResponsesCompact, + Auxiliary(AuxiliaryOperation), } impl UpstreamEndpoint { fn url(self, backend: &Backend) -> String { match self { UpstreamEndpoint::Completion => backend.url(), - UpstreamEndpoint::CountTokens => backend.count_tokens_url(), - UpstreamEndpoint::ResponsesInputTokens => format!("{}/input_tokens", backend.url()), - UpstreamEndpoint::ResponsesCompact => format!("{}/compact", backend.url()), + UpstreamEndpoint::Auxiliary(operation) => operation.url(backend), } } diff --git a/crates/libsy-llm-client/src/lib.rs b/crates/libsy-llm-client/src/lib.rs index 059675bbe..36175dd53 100644 --- a/crates/libsy-llm-client/src/lib.rs +++ b/crates/libsy-llm-client/src/lib.rs @@ -26,7 +26,7 @@ pub mod raw; pub mod run; pub use backend::{Backend, DEFAULT_MAX_RETRIES, HttpBackendConfig}; -pub use client::{ModelConfig, TranslatingLlmClient}; +pub use client::{AuxiliaryOperation, ModelConfig, TranslatingLlmClient}; pub use error::{LlmClientError, Result}; pub use observation::{LlmCallObservation, RunObservation, RunObserver}; pub use raw::RawResponse; diff --git a/crates/switchyard-runner/src/config.rs b/crates/switchyard-runner/src/config.rs index ddf8a9bf3..4b211ae81 100644 --- a/crates/switchyard-runner/src/config.rs +++ b/crates/switchyard-runner/src/config.rs @@ -12,14 +12,14 @@ use serde::de::DeserializeOwned; use serde::{Deserialize, Deserializer}; use serde_json::Value; use switchyard_llm_client::{ - Backend, ClientRouter, DEFAULT_MAX_RETRIES, HttpBackendConfig, ModelConfig, + AuxiliaryOperation, Backend, ClientRouter, DEFAULT_MAX_RETRIES, HttpBackendConfig, ModelConfig, TranslatingLlmClient, }; use switchyard_protocol::{ModelId, RoutedLlmClient, WireFormat}; use crate::{ - AlgorithmSpec, CallerAuthKind, CountTokensTarget, DecisionTarget, ModelCapabilities, - ResponsesTarget, Route, Runner, RunnerError, + AlgorithmSpec, AuxiliaryTarget, CallerAuthKind, DecisionTarget, ModelCapabilities, Route, + Runner, RunnerError, }; const SUPPORTED_SCHEMA_VERSION: u32 = 1; @@ -190,8 +190,10 @@ impl DeploymentConfig { .map_err(|error| RunnerError::configuration_source(error.to_string(), error))?; let (route_clients, caller_auth) = self.build_route_clients(route_name, config, &clients)?; - let count_tokens_target = self.build_count_tokens_target(config, &clients); - let responses_target = self.build_responses_target(config, &clients); + let anthropic_auxiliary_target = + self.build_anthropic_auxiliary_target(config, &clients); + let responses_auxiliary_target = + self.build_responses_auxiliary_target(config, &clients); let decision_targets = config .routing_target_names() .into_iter() @@ -202,8 +204,8 @@ impl DeploymentConfig { route_clients, caller_auth, capabilities, - count_tokens_target, - responses_target, + anthropic_auxiliary_target, + responses_auxiliary_target, decision_targets, ); routes.push((config.id.clone(), route)); @@ -308,48 +310,52 @@ impl DeploymentConfig { Ok(Some(config.base_url.as_str().to_string())) } - fn build_count_tokens_target( + fn build_anthropic_auxiliary_target( &self, route: &RouteConfig, clients: &BTreeMap>, - ) -> Option { + ) -> Option { route .routing_target_names() .into_iter() .enumerate() .filter_map(|(index, name)| { - let target = self.targets.get(name)?; - let client = clients.get(&target.llm_client)?; - client.supports_count_tokens(&target.id).then_some(( - count_tokens_priority(name, &target.id), - index, - target, - client, - )) - }) - .min_by_key(|(priority, index, _, _)| (*priority, *index)) - .map(|(_, _, target, client)| CountTokensTarget { - model: target.id.clone(), - client: client.clone(), + let target = self.build_auxiliary_target( + name, + clients, + AuxiliaryOperation::AnthropicCountTokens, + )?; + Some((count_tokens_priority(name, &target.model), index, target)) }) + .min_by_key(|(priority, index, _)| (*priority, *index)) + .map(|(_, _, target)| target) } - fn build_responses_target( + fn build_responses_auxiliary_target( &self, route: &RouteConfig, clients: &BTreeMap>, - ) -> Option { + ) -> Option { route.routing_target_names().into_iter().find_map(|name| { - let target = self.targets.get(name)?; - let client = clients.get(&target.llm_client)?; - client - .supports_responses_auxiliary(&target.id) - .then(|| ResponsesTarget { - model: target.id.clone(), - client: client.clone(), - }) + self.build_auxiliary_target(name, clients, AuxiliaryOperation::ResponsesInputTokens) }) } + + fn build_auxiliary_target( + &self, + name: &str, + clients: &BTreeMap>, + operation: AuxiliaryOperation, + ) -> Option { + let target = self.targets.get(name)?; + let client = clients.get(&target.llm_client)?; + client + .supports_auxiliary(&target.id, operation) + .then(|| AuxiliaryTarget { + model: target.id.clone(), + client: client.clone(), + }) + } } fn count_tokens_priority(target_name: &str, model_id: &ModelId) -> usize { diff --git a/crates/switchyard-runner/src/failure.rs b/crates/switchyard-runner/src/failure.rs index 598d8c31b..8113ef25b 100644 --- a/crates/switchyard-runner/src/failure.rs +++ b/crates/switchyard-runner/src/failure.rs @@ -100,8 +100,7 @@ impl RunnerError { ), Self::UnknownRouteModel(_) | Self::IncompatibleCallerFormat(_) - | Self::CountTokensUnsupported - | Self::ResponsesAuxiliaryUnsupported => summary( + | Self::AuxiliaryUnsupported => summary( RouteErrorKind::InvalidRequest, RouteErrorPhase::BeforeResponse, None, @@ -330,7 +329,7 @@ mod tests { RouteErrorKind::Configuration )); - let unsupported = RunnerError::CountTokensUnsupported; + let unsupported = RunnerError::AuxiliaryUnsupported; assert!(matches!( unsupported.execution_error_summary().kind, RouteErrorKind::InvalidRequest diff --git a/crates/switchyard-runner/src/lib.rs b/crates/switchyard-runner/src/lib.rs index 5e0cbb620..efb0441e4 100644 --- a/crates/switchyard-runner/src/lib.rs +++ b/crates/switchyard-runner/src/lib.rs @@ -15,7 +15,6 @@ pub use algorithm::{ }; pub use failure::{RouteErrorKind, RouteErrorPhase, RouteErrorSummary, stream_error_summary}; pub use route::{ - CallerAuthKind, CountTokensTarget, ModelCapabilities, ResponsesTarget, Route, RunOutput, - RunnerError, + AuxiliaryTarget, CallerAuthKind, ModelCapabilities, Route, RunOutput, RunnerError, }; pub use runner::{DecisionDescription, DecisionTarget, ModelInfo, Runner}; diff --git a/crates/switchyard-runner/src/route.rs b/crates/switchyard-runner/src/route.rs index 4e5b683a1..43022a9b5 100644 --- a/crates/switchyard-runner/src/route.rs +++ b/crates/switchyard-runner/src/route.rs @@ -8,7 +8,7 @@ use std::sync::Arc; use libsy::{Algorithm, CallModel, LibsyError, RoutingOutcome, drive}; use serde_json::Value; -use switchyard_llm_client::{ClientRouter, RunObserver, TranslatingLlmClient}; +use switchyard_llm_client::{AuxiliaryOperation, ClientRouter, RunObserver, TranslatingLlmClient}; use switchyard_protocol::{LlmClientError, ModelId, Request, Response, WireFormat}; use thiserror::Error; @@ -56,16 +56,9 @@ impl CallerAuthKind { } } -/// Exact upstream model used for Anthropic token counting. +/// Exact upstream model and client used for an auxiliary provider operation. #[derive(Clone)] -pub struct CountTokensTarget { - pub model: ModelId, - pub client: Arc, -} - -/// Exact upstream model used for OpenAI Responses auxiliary operations. -#[derive(Clone)] -pub struct ResponsesTarget { +pub struct AuxiliaryTarget { pub model: ModelId, pub client: Arc, } @@ -83,10 +76,8 @@ pub enum RunnerError { UnknownRouteModel(String), #[error("caller format is incompatible with {} credentials", .0.as_str())] IncompatibleCallerFormat(CallerAuthKind), - #[error("route has no Anthropic target for token counting")] - CountTokensUnsupported, - #[error("route has no OpenAI Responses target")] - ResponsesAuxiliaryUnsupported, + #[error("route has no compatible target for the auxiliary operation")] + AuxiliaryUnsupported, #[error(transparent)] Algorithm(#[from] LibsyError), #[error(transparent)] @@ -121,8 +112,8 @@ pub struct Route { clients: ClientRouter, caller_auth: Option, capabilities: ModelCapabilities, - count_tokens_target: Option, - responses_target: Option, + anthropic_auxiliary_target: Option, + responses_auxiliary_target: Option, decision_targets: Vec, } @@ -139,8 +130,8 @@ impl Route { clients: ClientRouter, caller_auth: Option, capabilities: ModelCapabilities, - count_tokens_target: Option, - responses_target: Option, + anthropic_auxiliary_target: Option, + responses_auxiliary_target: Option, decision_targets: Vec, ) -> Self { Self { @@ -148,8 +139,8 @@ impl Route { clients, caller_auth, capabilities, - count_tokens_target, - responses_target, + anthropic_auxiliary_target, + responses_auxiliary_target, decision_targets, } } @@ -215,41 +206,23 @@ impl Route { .map_err(Into::into) } - /// Counts tokens using the configured Anthropic-capable target. - pub async fn count_tokens(&self, request: Request) -> Result { - let target = self - .count_tokens_target - .as_ref() - .ok_or(RunnerError::CountTokensUnsupported)?; - target - .client - .count_tokens(&target.model, request) - .await - .map_err(Into::into) - } - - /// Counts tokens using the configured OpenAI Responses-capable target. - pub async fn responses_input_tokens(&self, request: Request) -> Result { - let target = self - .responses_target - .as_ref() - .ok_or(RunnerError::ResponsesAuxiliaryUnsupported)?; - target - .client - .responses_input_tokens(&target.model, request) - .await - .map_err(Into::into) - } - - /// Compacts a request using the configured OpenAI Responses-capable target. - pub async fn responses_compact(&self, request: Request) -> Result { - let target = self - .responses_target - .as_ref() - .ok_or(RunnerError::ResponsesAuxiliaryUnsupported)?; + /// Executes a model-bearing provider operation through a compatible target. + pub async fn call_auxiliary( + &self, + request: Request, + operation: AuxiliaryOperation, + ) -> Result { + let target = match operation { + AuxiliaryOperation::AnthropicCountTokens => &self.anthropic_auxiliary_target, + AuxiliaryOperation::ResponsesInputTokens | AuxiliaryOperation::ResponsesCompact => { + &self.responses_auxiliary_target + } + } + .as_ref() + .ok_or(RunnerError::AuxiliaryUnsupported)?; target .client - .responses_compact(&target.model, request) + .call_auxiliary(&target.model, request, operation) .await .map_err(Into::into) } diff --git a/crates/switchyard-server/src/lib.rs b/crates/switchyard-server/src/lib.rs index 120d1ee58..ae7671b61 100644 --- a/crates/switchyard-server/src/lib.rs +++ b/crates/switchyard-server/src/lib.rs @@ -37,7 +37,7 @@ use libsy::{Algorithm, LibsyError, RoutingOutcome}; use parking_lot::Mutex; use serde::{Deserialize, Serialize}; use serde_json::{Value, json}; -use switchyard_llm_client::{ClientRouter, RunObservation, RunObserver}; +use switchyard_llm_client::{AuxiliaryOperation, ClientRouter, RunObservation, RunObserver}; use switchyard_protocol::{LlmClientError, Metadata, ModelId, Request, Usage}; use switchyard_runner::{ CallerAuthKind, DecisionTarget, ModelCapabilities, Route, RunOutput, Runner, RunnerError, @@ -723,28 +723,22 @@ async fn anthropic_count_tokens( Ok(resolved) => resolved, Err(response) => return anthropic_error_response(response), }; - anthropic_error_response(match route.count_tokens(request).await { - Ok(payload) => (StatusCode::OK, Json(payload)).into_response(), - Err(RunnerError::CountTokensUnsupported) => error_response( - StatusCode::BAD_REQUEST, - "route has no Anthropic target for token counting", - "invalid_request_error", - "count_tokens_unsupported", - ), - Err(RunnerError::Client(error)) => count_tokens_error(error), - Err(error) => server_error(error.to_string()), - }) -} - -/// Maps a token-count client failure with the same policy as a routed client call. -fn count_tokens_error(error: LlmClientError) -> Response { - client_error(&error) -} - -#[derive(Clone, Copy)] -enum ResponsesAuxiliaryEndpoint { - InputTokens, - Compact, + anthropic_error_response( + match route + .call_auxiliary(request, AuxiliaryOperation::AnthropicCountTokens) + .await + { + Ok(payload) => (StatusCode::OK, Json(payload)).into_response(), + Err(RunnerError::AuxiliaryUnsupported) => error_response( + StatusCode::BAD_REQUEST, + "route has no Anthropic target for token counting", + "invalid_request_error", + "count_tokens_unsupported", + ), + Err(RunnerError::Client(error)) => client_error(&error), + Err(error) => server_error(error.to_string()), + }, + ) } async fn openai_responses_input_tokens( @@ -756,7 +750,7 @@ async fn openai_responses_input_tokens( state, headers, body, - ResponsesAuxiliaryEndpoint::InputTokens, + AuxiliaryOperation::ResponsesInputTokens, ) .await } @@ -766,7 +760,7 @@ async fn openai_responses_compact( headers: HeaderMap, body: std::result::Result, JsonRejection>, ) -> Response { - openai_responses_auxiliary(state, headers, body, ResponsesAuxiliaryEndpoint::Compact).await + openai_responses_auxiliary(state, headers, body, AuxiliaryOperation::ResponsesCompact).await } // Resolves route aliases and configured credentials before calling a model-bearing Responses API. @@ -774,7 +768,7 @@ async fn openai_responses_auxiliary( state: ServerState, headers: HeaderMap, body: std::result::Result, JsonRejection>, - endpoint: ResponsesAuxiliaryEndpoint, + operation: AuxiliaryOperation, ) -> Response { let body = match llm_json_body(body) { Ok(body) => body, @@ -794,14 +788,11 @@ async fn openai_responses_auxiliary( Ok(resolved) => resolved, Err(response) => return render_error_response(response, WireFormat::OpenAiResponses), }; - let result = match endpoint { - ResponsesAuxiliaryEndpoint::InputTokens => route.responses_input_tokens(request).await, - ResponsesAuxiliaryEndpoint::Compact => route.responses_compact(request).await, - }; + let result = route.call_auxiliary(request, operation).await; render_error_response( match result { Ok(payload) => (StatusCode::OK, Json(payload)).into_response(), - Err(RunnerError::ResponsesAuxiliaryUnsupported) => error_response( + Err(RunnerError::AuxiliaryUnsupported) => error_response( StatusCode::BAD_REQUEST, "route has no OpenAI Responses target", "invalid_request_error", From 347d6d1d95345a65975ca8712c4e4733f08e361c Mon Sep 17 00:00:00 2001 From: nachiketb Date: Mon, 31 Aug 2026 12:31:20 -0700 Subject: [PATCH 5/5] fix(server): build fallback proxy URLs with parser Signed-off-by: nachiketb --- crates/switchyard-server/src/lib.rs | 51 +++++++++++++++++++++++++---- 1 file changed, 44 insertions(+), 7 deletions(-) diff --git a/crates/switchyard-server/src/lib.rs b/crates/switchyard-server/src/lib.rs index ae7671b61..eea1c7a2e 100644 --- a/crates/switchyard-server/src/lib.rs +++ b/crates/switchyard-server/src/lib.rs @@ -564,10 +564,21 @@ async fn proxy_unmatched(State(state): State, request: HttpRequest) .map_or_else(|| parts.uri.path().to_string(), ToString::to_string); strip_hop_by_hop_headers(&mut parts.headers); parts.headers.remove(axum::http::header::HOST); + let url = match fallback_url(base_url, &path_and_query) { + Ok(url) => url, + Err(error) => { + return error_response( + StatusCode::BAD_GATEWAY, + error, + "upstream_error", + "upstream_error", + ); + } + }; let body = reqwest::Body::wrap_stream(body.into_data_stream()); match state .fallback_http - .request(parts.method, fallback_url(base_url, &path_and_query)) + .request(parts.method, url) .headers(parts.headers) .body(body) .send() @@ -588,18 +599,44 @@ async fn proxy_unmatched(State(state): State, request: HttpRequest) } } -fn fallback_url(base_url: &str, path_and_query: &str) -> String { - let base_url = base_url.trim_end_matches('/'); - let root = [ +fn fallback_url(base_url: &str, path_and_query: &str) -> Result { + let mut url = reqwest::Url::parse(base_url.trim()) + .map_err(|error| format!("invalid fallback base URL: {error}"))?; + let base_query = url.query().map(str::to_owned); + url.set_query(None); + url.set_fragment(None); + + let base_path = url.path().trim_end_matches('/'); + let root_path = [ "/v1/chat/completions", "/v1/responses", "/v1/messages", "/v1", ] .iter() - .find_map(|suffix| base_url.strip_suffix(suffix)) - .unwrap_or(base_url); - format!("{root}{path_and_query}") + .find_map(|suffix| base_path.strip_suffix(suffix)) + .unwrap_or(base_path); + let (request_path, request_query) = path_and_query + .split_once('?') + .map_or((path_and_query, None), |(path, query)| (path, Some(query))); + let request_path = request_path.strip_prefix('/').unwrap_or(request_path); + let target_path = if root_path.is_empty() { + format!("/{request_path}") + } else if request_path.is_empty() { + format!("{root_path}/") + } else { + format!("{root_path}/{request_path}") + }; + url.set_path(&target_path); + + let query = match (base_query, request_query) { + (Some(base_query), Some(request_query)) => Some(format!("{base_query}&{request_query}")), + (Some(base_query), None) => Some(base_query), + (None, Some(request_query)) => Some(request_query.to_owned()), + (None, None) => None, + }; + url.set_query(query.as_deref()); + Ok(url) } fn strip_hop_by_hop_headers(headers: &mut HeaderMap) {