From b26c0128fa4f3a99d80adfc92cf81023978135e5 Mon Sep 17 00:00:00 2001 From: Elyas Mehtabuddin Date: Thu, 6 Aug 2026 11:08:43 -0700 Subject: [PATCH 1/2] fix(routing): scope context-overflow history by session and agent Signed-off-by: Elyas Mehtabuddin --- crates/libsy/src/algorithms/fall_through.rs | 22 ++- crates/libsy/src/algorithms/util/affinity.rs | 38 ++--- crates/libsy/src/core/algorithm.rs | 102 +++++++++---- crates/switchyard-server/tests/server.rs | 146 +++++++++++++++++++ 4 files changed, 240 insertions(+), 68 deletions(-) diff --git a/crates/libsy/src/algorithms/fall_through.rs b/crates/libsy/src/algorithms/fall_through.rs index 8d1314a56..2ea2e23a2 100644 --- a/crates/libsy/src/algorithms/fall_through.rs +++ b/crates/libsy/src/algorithms/fall_through.rs @@ -26,7 +26,9 @@ use async_trait::async_trait; use parking_lot::Mutex; use tokio::sync::Mutex as AsyncMutex; -use crate::core::algorithm::{self, Algorithm, Driver, LlmTarget, LlmTargetSet, SessionEvictions}; +use crate::core::algorithm::{ + self, Algorithm, Driver, LlmTarget, LlmTargetSet, RoutingIdentity, SessionEvictions, +}; use crate::core::classifier::{Classification, Classifier, Score}; use crate::core::processor::{Event, Processor}; use crate::{LibsyError, Result}; @@ -198,9 +200,7 @@ where .as_ref() .and_then(|metadata| metadata.session_final) == Some(true); - let result = self - .execute_session(ctx, driver, request, session.as_deref()) - .await; + let result = self.execute_session(ctx, driver, request).await; if session_final && let Some(session) = session.as_deref() { self.remove_session(session); } @@ -223,13 +223,19 @@ where ctx: Context, driver: Driver, request: Request, - session: Option<&str>, ) -> Result { // The request is threaded mutably through the whole fold: any component may rewrite // it, later components see the rewrite, and the final value reaches the model. let mut request = request; let mut ctx = ctx; - algorithm::exclude_evicted(&mut ctx, &self.targets, &self.session_evictions, session); + // Processors may rewrite the request; overflow history stays with its inbound identity. + let identity = RoutingIdentity::from_request(&request); + algorithm::exclude_evicted( + &mut ctx, + &self.targets, + &self.session_evictions, + identity.as_ref(), + ); let session_state = self.session_state(&request); let (target, decision, served) = match session_state { Some(state) => { @@ -257,7 +263,7 @@ where target, decision, request, - session, + identity.as_ref(), &self.session_evictions, |from, to| self.fallback_decision(from, to), ) @@ -271,7 +277,7 @@ where if let Some(states) = &self.session_states { states.lock().remove(session); } - self.session_evictions.remove(session); + self.session_evictions.remove_session(session); } /// The decision published when an overflow sends the request to a different target. diff --git a/crates/libsy/src/algorithms/util/affinity.rs b/crates/libsy/src/algorithms/util/affinity.rs index 04aa59221..28a701d08 100644 --- a/crates/libsy/src/algorithms/util/affinity.rs +++ b/crates/libsy/src/algorithms/util/affinity.rs @@ -24,7 +24,7 @@ use async_trait::async_trait; use parking_lot::Mutex; use switchyard_protocol::{Request, Role}; -use crate::core::algorithm::Driver; +use crate::core::algorithm::{Driver, RoutingIdentity}; use crate::core::classifier::{Classification, Classifier, Score}; use crate::core::processor::{Event, Processor}; @@ -32,18 +32,6 @@ use crate::core::processor::{Event, Processor}; /// without limit; the oldest entry is evicted once the bound is reached. const MAX_ASSIGNMENTS: usize = 4096; -/// The stable identity a model assignment is retained against. -/// -/// A sub-agent request is keyed by `session + agent` and a root request by session alone, -/// so a sub-agent's assignment is scoped within — but distinct from — its session's. -#[derive(Clone, Hash, PartialEq, Eq)] -enum AffinityKey { - /// One model per session, for root-agent traffic. - Session(String), - /// One model per identified child agent within a session. - Subagent { session: String, agent: String }, -} - /// Retains a model per request identity and forces it on later matching requests. /// /// Register the same instance as both a processor and a classifier; the two roles share @@ -66,7 +54,7 @@ pub struct AffinityRouter { /// /// Held on the instance so the two roles share one process-local map through a /// single registered [`Arc`](std::sync::Arc); bounded by [`MAX_ASSIGNMENTS`]. - assignments: Mutex>, + assignments: Mutex>, } impl AffinityRouter { @@ -106,19 +94,11 @@ impl AffinityRouter { } /// Derives the stable identity this router should retain for `request`. - fn affinity_key(&self, request: &Request) -> Option { - if let Some(metadata) = request.metadata.as_ref() - && let Some(session) = metadata.session_id.clone() - { - return if metadata.is_subagent { - metadata - .agent_id - .clone() - .map(|agent| AffinityKey::Subagent { session, agent }) - } else if self.subagents_only { - None - } else { - Some(AffinityKey::Session(session)) + fn affinity_key(&self, request: &Request) -> Option { + if let Some(identity) = RoutingIdentity::from_request(request) { + return match identity { + RoutingIdentity::Session(_) if self.subagents_only => None, + identity => Some(identity), }; } @@ -131,7 +111,7 @@ impl AffinityRouter { .then(|| { first_user_message_hash(request).map(|hash| { tracing::debug!(affinity_key = %hash, "affinity using message hash fallback"); - AffinityKey::Session(hash) + RoutingIdentity::Session(hash) }) }) .flatten() @@ -201,7 +181,7 @@ where } /// Evicts one arbitrary assignment when the map has reached [`MAX_ASSIGNMENTS`]. -fn evict_if_full(assignments: &mut HashMap) { +fn evict_if_full(assignments: &mut HashMap) { if assignments.len() >= MAX_ASSIGNMENTS && let Some(evicted) = assignments.keys().next().cloned() { diff --git a/crates/libsy/src/core/algorithm.rs b/crates/libsy/src/core/algorithm.rs index d63294ee7..51db44c55 100644 --- a/crates/libsy/src/core/algorithm.rs +++ b/crates/libsy/src/core/algorithm.rs @@ -407,51 +407,91 @@ impl LlmTargetSet { } } -/// Bounds process-local overflow history. Dropping a live session's entry costs one -/// rediscovered overflow, so the victim choice does not need to be exact. -const MAX_EVICTION_SESSIONS: usize = 1_024; +/// Key for overflow history: a root request by its session, a child request by its session +/// and agent. Keying a child finer than its session keeps one child's overflow from evicting +/// a target for the parent or a sibling sharing the session. +#[derive(Clone, Hash, PartialEq, Eq)] +pub(crate) enum RoutingIdentity { + /// Root request, keyed by session ID. + Session(String), + /// Child request, keyed by session and agent IDs. + Subagent { session: String, agent: String }, +} + +impl RoutingIdentity { + /// Builds a root or child identity from non-empty request metadata. + /// + /// A child request missing either ID returns `None`, so it keeps no routing history + /// rather than sharing the parent's. + pub(crate) fn from_request(request: &Request) -> Option { + let metadata = request.metadata.as_ref()?; + let session = metadata.session_id.as_deref().filter(|id| !id.is_empty())?; + if metadata.is_subagent { + let agent = metadata.agent_id.as_deref().filter(|id| !id.is_empty())?; + Some(Self::Subagent { + session: session.to_string(), + agent: agent.to_string(), + }) + } else { + Some(Self::Session(session.to_string())) + } + } -/// Per-session record of the targets that overflowed their context window. + /// The session this identity belongs to; shared by a session's root and its children. + fn session(&self) -> &str { + match self { + Self::Session(session) | Self::Subagent { session, .. } => session, + } + } +} + +/// Bounds process-local overflow history. Dropping a live entry costs one rediscovered +/// overflow, so the victim choice does not need to be exact. +const MAX_EVICTION_IDENTITIES: usize = 1_024; + +/// Per-identity record of the targets that overflowed their context window. /// /// A conversation only grows, so a target that could not fit one turn will not fit a /// later one; remembering it lets the next turn skip a call certain to fail. Requests -/// without a session id are not tracked — there is nothing to remember them by. +/// without a routing identity are not tracked — there is nothing to remember them by. #[derive(Default)] pub(crate) struct SessionEvictions { - sessions: Mutex>>, + by_identity: Mutex>>, } impl SessionEvictions { - /// Forgets overflow history for a completed session. - pub(crate) fn remove(&self, session: &str) { - self.sessions.lock().remove(session); + /// Forgets overflow history for a completed session, including every child of it. + pub(crate) fn remove_session(&self, session: &str) { + self.by_identity + .lock() + .retain(|identity, _| identity.session() != session); } - /// The targets `session` has already overflowed; empty for an untracked request. - fn evicted_in(&self, session: Option<&str>) -> Vec { - let Some(session) = session else { + /// The targets `identity` has already overflowed; empty for an untracked request. + fn evicted_for(&self, identity: Option<&RoutingIdentity>) -> Vec { + let Some(identity) = identity else { return Vec::new(); }; - self.sessions + self.by_identity .lock() - .get(session) + .get(identity) .map(|targets| targets.iter().cloned().collect()) .unwrap_or_default() } - /// Remembers that `target` overflowed in `session`, tracking at most - /// [`MAX_EVICTION_SESSIONS`] sessions. - fn record(&self, session: Option<&str>, target: &str) { - let Some(session) = session else { return }; - let mut sessions = self.sessions.lock(); - if sessions.len() >= MAX_EVICTION_SESSIONS - && !sessions.contains_key(session) - && let Some(oldest) = sessions.keys().next().cloned() + /// Remembers that `target` overflowed for `identity`, tracking at most + /// [`MAX_EVICTION_IDENTITIES`] identities. + fn record(&self, identity: Option<&RoutingIdentity>, target: &str) { + let Some(identity) = identity else { return }; + let mut histories = self.by_identity.lock(); + if histories.len() >= MAX_EVICTION_IDENTITIES + && !histories.contains_key(identity) + && let Some(oldest) = histories.keys().next().cloned() { - sessions.remove(&oldest); + histories.remove(&oldest); } - sessions - .entry(session.to_string()) + histories + .entry(identity.clone()) .or_default() .insert(target.to_string()); } @@ -466,15 +506,15 @@ fn eligible_targets(targets: &LlmTargetSet, ctx: &Context) -> usize { .count() } -/// Bars the targets `session` has already overflowed from this request, so routing does +/// Bars the targets `identity` has already overflowed from this request, so routing does /// not select one that is certain to fail again. pub(crate) fn exclude_evicted( ctx: &mut Context, targets: &LlmTargetSet, evictions: &SessionEvictions, - session: Option<&str>, + identity: Option<&RoutingIdentity>, ) { - for target in evictions.evicted_in(session) { + for target in evictions.evicted_for(identity) { // Never seed the pool empty: a later turn may be small enough to serve, and the // caller should get the upstream's answer rather than a routing error. if eligible_targets(targets, ctx) <= 1 { @@ -490,7 +530,7 @@ pub(crate) fn exclude_evicted( /// Routing is deliberately not re-run: the fallback replaces the target in place, so the /// caller's request-side work and retained state still see exactly one turn. /// `fallback_decision` builds the [`Decision`] published for a `from -> to` hop, and each -/// overflow is recorded against `session` so later turns skip that target outright. +/// overflow is recorded for `identity` so later turns skip that target outright. #[allow(clippy::too_many_arguments)] pub(crate) async fn call_llm_with_overflow_fallback( mut ctx: Context, @@ -499,7 +539,7 @@ pub(crate) async fn call_llm_with_overflow_fallback( mut target: LlmTarget, mut decision: Arc, request: Request, - session: Option<&str>, + identity: Option<&RoutingIdentity>, evictions: &SessionEvictions, fallback_decision: impl Fn(&LlmTarget, &LlmTarget) -> Arc, ) -> Result { @@ -520,7 +560,7 @@ pub(crate) async fn call_llm_with_overflow_fallback( if !ctx.exclude_target(failed) { return Err(error); } - evictions.record(session, failed); + evictions.record(identity, failed); let Ok(next) = targets.resolve_target(&target.semantic_name, &ctx) else { return Err(error); }; diff --git a/crates/switchyard-server/tests/server.rs b/crates/switchyard-server/tests/server.rs index 045f118b0..0dde3915d 100644 --- a/crates/switchyard-server/tests/server.rs +++ b/crates/switchyard-server/tests/server.rs @@ -90,6 +90,18 @@ async fn upstream_chat( } let model = body["model"].as_str().unwrap_or("unknown").to_string(); + if model == "model/weak" && body["messages"][0]["content"] == "overflow" { + return ( + StatusCode::BAD_REQUEST, + Json(json!({ + "error": { + "code": "context_length_exceeded", + "message": "request exceeds this model's context window" + } + })), + ) + .into_response(); + } if body["stream"].as_bool() == Some(true) { if body["messages"][0]["content"] == "stream-error" { let events = [ @@ -513,8 +525,52 @@ fn load_test_config(toml: &str) -> TestResult { Ok(load_server_state(config.path())?) } +/// A `random` route that always selects `first` (weight 1 vs 0), so a test can drive the +/// overflow fallback from `first` to `second` deterministically. +fn overflow_fallback_state(base_url: &str) -> TestResult { + load_test_config(&format!( + r#" +schema_version = 1 + +[llm_clients.mock] +format = "openai_chat" +base_url = "{base_url}" +max_retries = 0 + +[targets.first] +id = "{first}" +llm_client = "mock" + +[targets.second] +id = "{second}" +llm_client = "mock" + +[routes.random] +id = "{ROUTE_MODEL}" +type = "random" +targets = ["first", "second"] +weights = [1, 0] +"#, + first = "model/weak", + second = "model/strong", + )) +} + async fn send(app: &Router, method: &str, path: &str, body: Option) -> TestResult { + send_with_headers(app, method, path, body, &[]).await +} + +async fn send_with_headers( + app: &Router, + method: &str, + path: &str, + body: Option, + headers: &[(&str, &str)], +) -> TestResult { let mut builder = HttpRequest::builder().method(method).uri(path); + for (name, value) in headers { + builder = builder.header(*name, *value); + } let request_body = if let Some(body) = body { builder = builder.header("content-type", "application/json"); Body::from(serde_json::to_vec(&body)?) @@ -1377,6 +1433,96 @@ async fn routing_log_exposes_session_stats() -> TestResult { Ok(()) } +// Overflow history is isolated per child, cleared with the session, and not retained when a +// child lacks an agent ID. +#[tokio::test] +async fn overflow_history_is_scoped_to_agent_and_session_lifetime() -> TestResult { + let upstream = MockUpstream::start().await?; + let state = overflow_fallback_state(&upstream.base_url)?; + let app = build_switchyard_router(state); + let child_a = [ + ("x-switchyard-session-id", "shared-session"), + ("x-switchyard-agent-id", "child-a"), + ("x-switchyard-is-subagent", "true"), + ]; + let root = [ + ("x-switchyard-session-id", "shared-session"), + ("x-switchyard-agent-id", "root"), + ("x-switchyard-is-subagent", "false"), + ]; + let child_b = [ + ("x-switchyard-session-id", "shared-session"), + ("x-switchyard-agent-id", "child-b"), + ("x-switchyard-is-subagent", "true"), + ]; + let child_without_agent_id = [ + ("x-switchyard-session-id", "shared-session"), + ("x-switchyard-is-subagent", "true"), + ]; + let final_root = [ + ("x-switchyard-session-id", "shared-session"), + ("x-switchyard-agent-id", "root"), + ("x-switchyard-is-subagent", "false"), + ("x-switchyard-session-final", "true"), + ]; + type Case<'a> = (&'a str, &'a [(&'a str, &'a str)], &'a [&'a str]); + let cases: [Case<'_>; 8] = [ + ( + "overflow", + child_a.as_slice(), + &["model/weak", "model/strong"], + ), + ("fits", child_a.as_slice(), &["model/strong"]), + ("fits", root.as_slice(), &["model/weak"]), + ("fits", child_b.as_slice(), &["model/weak"]), + ("fits", final_root.as_slice(), &["model/weak"]), + ("fits", child_a.as_slice(), &["model/weak"]), + ( + "overflow", + child_without_agent_id.as_slice(), + &["model/weak", "model/strong"], + ), + ( + "overflow", + child_without_agent_id.as_slice(), + &["model/weak", "model/strong"], + ), + ]; + + for (content, headers, expected_calls) in cases { + let previous_call_count = upstream.calls.lock().await.len(); + let response = send_with_headers( + &app, + "POST", + "/v1/chat/completions", + Some(json!({ + "model": ROUTE_MODEL, + "messages": [{"role": "user", "content": content}] + })), + headers, + ) + .await?; + assert_eq!(response.status, StatusCode::OK); + let expected_model = expected_calls.last().copied(); + assert_eq!( + response + .headers + .get("x-model-router-selected-model") + .and_then(|value| value.to_str().ok()), + expected_model + ); + let calls = upstream.calls.lock().await; + assert_eq!( + calls[previous_call_count..] + .iter() + .map(|call| call["model"].as_str().unwrap_or("")) + .collect::>(), + expected_calls + ); + } + Ok(()) +} + #[tokio::test] async fn streaming_response_is_framed_for_the_inbound_api() -> TestResult { let (_upstream, app) = test_app(&[(ROUTE_MODEL, &["model/a"])]).await?; From 366623a7bd095aade82076bbc0c910532ddc3271 Mon Sep 17 00:00:00 2001 From: Elyas Mehtabuddin Date: Thu, 6 Aug 2026 17:43:25 -0700 Subject: [PATCH 2/2] fix(routing): fail over unavailable targets Signed-off-by: Elyas Mehtabuddin --- crates/libsy/src/algorithms/fall_through.rs | 131 ++++++++++++++---- crates/libsy/src/algorithms/llm_class.rs | 2 +- crates/libsy/src/algorithms/util/affinity.rs | 13 ++ crates/libsy/src/core/algorithm.rs | 106 ++++++++++++-- crates/libsy/src/core/classifier.rs | 5 + crates/protocol/src/client.rs | 23 +++ crates/switchyard-server/src/lib.rs | 8 ++ crates/switchyard-server/src/routing_log.rs | 14 +- .../src/stats/accumulator.rs | 24 ++++ crates/switchyard-server/tests/server.rs | 116 +++++++++++++++- 10 files changed, 396 insertions(+), 46 deletions(-) diff --git a/crates/libsy/src/algorithms/fall_through.rs b/crates/libsy/src/algorithms/fall_through.rs index 2ea2e23a2..96a7e5093 100644 --- a/crates/libsy/src/algorithms/fall_through.rs +++ b/crates/libsy/src/algorithms/fall_through.rs @@ -15,6 +15,7 @@ //! //! Every composition retains one thing regardless: a target that overflows its context window is //! remembered for the rest of its session and skipped on later turns. +//! An unavailable target is skipped only for the current request. use std::{ collections::HashMap, @@ -32,7 +33,7 @@ use crate::core::algorithm::{ use crate::core::classifier::{Classification, Classifier, Score}; use crate::core::processor::{Event, Processor}; use crate::{LibsyError, Result}; -use switchyard_protocol::{Context, Decision, Request, Response}; +use switchyard_protocol::{Context, Decision, Request, Response, RoutingFallbackReason}; struct SessionState { state: Arc>, @@ -57,6 +58,7 @@ pub struct FallThroughDecision { /// Human-readable explanation of the selection. pub reasoning: String, tier: Option<&'static str>, + fallback_reason: Option, } impl Decision for FallThroughDecision { @@ -68,6 +70,10 @@ impl Decision for FallThroughDecision { self.tier } + fn fallback_reason(&self) -> Option { + self.fallback_reason + } + fn reasoning(&self) -> Option<&str> { Some(&self.reasoning) } @@ -237,7 +243,7 @@ where identity.as_ref(), ); let session_state = self.session_state(&request); - let (target, decision, served) = match session_state { + let (target, decision, served, deciding) = match session_state { Some(state) => { let mut state = state.lock().await; self.route(&mut state, &ctx, &driver, &mut request).await? @@ -256,7 +262,7 @@ where match served { Some(response) => Ok(response), None => { - algorithm::call_llm_with_overflow_fallback( + algorithm::call_llm_with_fallback( ctx, &driver, &self.targets, @@ -265,7 +271,12 @@ where request, identity.as_ref(), &self.session_evictions, - |from, to| self.fallback_decision(from, to), + |request, target| { + for classifier in &self.classifiers { + classifier.target_unavailable(request, target); + } + }, + |from, to, reason| self.fallback_decision(deciding.as_ref(), from, to, reason), ) .await } @@ -280,18 +291,26 @@ where self.session_evictions.remove_session(session); } - /// The decision published when an overflow sends the request to a different target. - fn fallback_decision(&self, from: &LlmTarget, to: &LlmTarget) -> Arc { + /// The decision published when a route-level failure selects a different target. + fn fallback_decision( + &self, + deciding: &dyn Classifier, + from: &LlmTarget, + to: &LlmTarget, + reason: RoutingFallbackReason, + ) -> Arc { + let failure = match reason { + RoutingFallbackReason::ContextWindow => "exceeded its context window", + RoutingFallbackReason::Unavailable => "was unavailable", + }; Arc::new(FallThroughDecision { selected_model: to.semantic_name.clone(), reasoning: format!( - "{} exceeded its context window; fell back to {}", - from.semantic_name, to.semantic_name + "{} {failure}; fell back to {}", + from.semantic_name, to.semantic_name, ), - tier: self - .classifiers - .iter() - .find_map(|c| c.routing_tier(&to.semantic_name)), + tier: deciding.routing_tier(&to.semantic_name), + fallback_reason: Some(reason), }) } @@ -315,7 +334,12 @@ where ctx: &Context, driver: &Driver, request: &mut Request, - ) -> Result<(LlmTarget, Arc, Option)> { + ) -> Result<( + LlmTarget, + Arc, + Option, + Arc>, + )> { // 1. Processor chain accumulates request-side facts into the composition's state. for processor in &self.processors { processor.process(state, Event::Request(request)).await?; @@ -323,21 +347,17 @@ where // 2. Fall through the cascade: the first classifier to score decides (argmax). The // per-request driver is offered to each — driver-backed classifiers use it. - let mut maybe_score: Option = None; - let mut deciding: Option<&Arc>> = None; - let mut served: Option = None; + let mut routed = None; for classifier in &self.classifiers { let (scores, response) = classifier.score(state, request, Some(driver)).await?; - maybe_score = scores.argmax(false)?; - if maybe_score.is_some() { - deciding = Some(classifier); + if let Some(score) = scores.argmax(false)? { // Only the deciding classifier's response answers the turn; an abstaining // classifier selected nothing for it to be the answer to. - served = response; + routed = Some((score, Arc::clone(classifier), response)); break; } } - let Some(score) = maybe_score else { + let Some((score, deciding, served)) = routed else { return Err(LibsyError::AlgorithmError { message: "every classifier abstained".to_string(), }); @@ -357,7 +377,8 @@ where let decision: Arc = Arc::new(FallThroughDecision { selected_model: target.semantic_name.clone(), reasoning, - tier: deciding.and_then(|c| c.routing_tier(&target.semantic_name)), + tier: deciding.routing_tier(&target.semantic_name), + fallback_reason: None, }); driver.info(ctx.clone(), decision.clone()).await?; @@ -371,7 +392,7 @@ where processor.process(state, event).await?; } - Ok((target, decision, served)) + Ok((target, decision, served, deciding)) } } @@ -440,7 +461,7 @@ mod tests { use super::*; use crate::algorithms::util::prompts; use crate::core::classifier::Classification; - use crate::{SystemPromptProcessor, TargetPrompts}; + use crate::{AffinityRouter, SystemPromptProcessor, TargetPrompts}; use switchyard_protocol::{ LlmClientError, LlmRequest, LlmResponse, Message, Metadata, Role, RoutedLlmClient, @@ -797,6 +818,45 @@ mod tests { calls: Option>>>, } + /// Returns 503 for `weak` and records every routed call. + struct UnavailableClient(Arc>>); + + #[async_trait] + impl RoutedLlmClient for UnavailableClient { + async fn call( + &self, + _ctx: Context, + _request: Request, + decision: Arc, + ) -> std::result::Result { + let model = decision.selected_model().to_string(); + self.0.lock().push(model.clone()); + if model == "weak" { + return Err(LlmClientError::UpstreamHttp { + status: 503, + body: "unavailable".to_string(), + }); + } + Ok(Response { + llm_response: LlmResponse::Agg(text_response(None, model)), + metadata: None, + }) + } + } + + fn unavailable_targets(calls: Arc>>) -> LlmTargetSet { + let client: Arc = Arc::new(UnavailableClient(calls)); + LlmTargetSet::new( + ["weak", "strong"] + .into_iter() + .map(|name| LlmTarget { + semantic_name: name.to_string(), + llm_client: Some(Arc::clone(&client)), + }) + .collect(), + ) + } + #[async_trait] impl RoutedLlmClient for OverflowClient { async fn call( @@ -930,6 +990,29 @@ mod tests { Ok(()) } + #[tokio::test] + async fn unavailable_target_clears_matching_affinity_before_the_next_turn() -> Result<()> { + let calls = Arc::new(Mutex::new(Vec::new())); + let affinity = Arc::new(AffinityRouter::new()); + let router = Arc::new( + FallThrough::<()>::new(unavailable_targets(Arc::clone(&calls))) + .with_processor(affinity.clone()) + .with_classifier(affinity) + .with_classifier(Arc::new(DefaultTarget::new("weak"))), + ); + + for _ in 0..2 { + let (model, trace) = run_turn(&router).await?; + assert_eq!(model, "strong"); + assert_eq!( + trace.last().and_then(|decision| decision.fallback_reason()), + Some(RoutingFallbackReason::Unavailable) + ); + } + assert_eq!(&*calls.lock(), &["weak", "strong", "weak", "strong"]); + Ok(()) + } + #[tokio::test] async fn overflowing_targets_are_retried_until_one_fits() -> Result<()> { let router = FallThrough::<()>::new(target_set_with_overflow( diff --git a/crates/libsy/src/algorithms/llm_class.rs b/crates/libsy/src/algorithms/llm_class.rs index c2cce86e5..95f50d312 100644 --- a/crates/libsy/src/algorithms/llm_class.rs +++ b/crates/libsy/src/algorithms/llm_class.rs @@ -517,7 +517,7 @@ impl Classifier for EscalationClassifier { // // If the efficient model exceeds its context window, fall through to capable: returning // `(decisive(capable), None)` tells FallThrough::execute to call - // call_llm_with_overflow_fallback with the capable target instead of surfacing the error. + // call_llm_with_fallback with the capable target instead of surfacing the error. let efficient_response = match driver .call_llm_target( Context::default(), diff --git a/crates/libsy/src/algorithms/util/affinity.rs b/crates/libsy/src/algorithms/util/affinity.rs index 28a701d08..da07eb7b4 100644 --- a/crates/libsy/src/algorithms/util/affinity.rs +++ b/crates/libsy/src/algorithms/util/affinity.rs @@ -157,6 +157,19 @@ impl Classifier for AffinityRouter where S: Send + 'static, { + fn target_unavailable(&self, request: &Request, target: &str) { + let Some(key) = self.affinity_key(request) else { + return; + }; + let mut assignments = self.assignments.lock(); + if assignments + .get(&key) + .is_some_and(|assigned| assigned == target) + { + assignments.remove(&key); + } + } + async fn score( &self, _state: &mut S, diff --git a/crates/libsy/src/core/algorithm.rs b/crates/libsy/src/core/algorithm.rs index 51db44c55..0d4e76f9f 100644 --- a/crates/libsy/src/core/algorithm.rs +++ b/crates/libsy/src/core/algorithm.rs @@ -25,7 +25,8 @@ use tracing::Instrument; /// [`switchyard_protocol::LlmResponse`] carries either a live /// [`switchyard_protocol::LlmResponseStream`] or the terminal aggregate. use switchyard_protocol::{ - Context, Decision, LlmClientError, Request, Response, RoutedLlmClient, Signals, Usage, + Context, Decision, LlmClientError, Request, Response, RoutedLlmClient, RoutingFallbackReason, + Signals, Usage, }; use super::driver::{DriverRequest, DriverStep, TypeErasedDriver}; @@ -524,15 +525,35 @@ pub(crate) fn exclude_evicted( } } -/// Calls `target`, falling back to the next eligible target in `targets` whenever one -/// overflows its context window, until a call succeeds or every target has been tried. +/// Returns the failed target and routing fallback policy for a terminal client error. +fn classify_fallback(error: &LibsyError) -> Option<(&str, RoutingFallbackReason)> { + let LibsyError::ClientCall { target, source } = error else { + return None; + }; + let reason = match source { + LlmClientError::ContextWindowExceeded { .. } => RoutingFallbackReason::ContextWindow, + LlmClientError::Transport { .. } | LlmClientError::Timeout { .. } => { + RoutingFallbackReason::Unavailable + } + LlmClientError::UpstreamHttp { status, .. } + if matches!(*status, 403 | 408 | 429) || (500..=599).contains(status) => + { + RoutingFallbackReason::Unavailable + } + _ => return None, + }; + Some((target, reason)) +} + +/// Calls `target`, falling back to the next eligible target after a route-level failure, +/// until a call succeeds or every target has been tried. /// /// Routing is deliberately not re-run: the fallback replaces the target in place, so the /// caller's request-side work and retained state still see exactly one turn. -/// `fallback_decision` builds the [`Decision`] published for a `from -> to` hop, and each -/// overflow is recorded for `identity` so later turns skip that target outright. +/// `fallback_decision` builds the [`Decision`] published for a `from -> to` hop. Context +/// overflows are recorded for `identity`; unavailable targets remain request-local. #[allow(clippy::too_many_arguments)] -pub(crate) async fn call_llm_with_overflow_fallback( +pub(crate) async fn call_llm_with_fallback( mut ctx: Context, driver: &Driver, targets: &LlmTargetSet, @@ -541,30 +562,30 @@ pub(crate) async fn call_llm_with_overflow_fallback( request: Request, identity: Option<&RoutingIdentity>, evictions: &SessionEvictions, - fallback_decision: impl Fn(&LlmTarget, &LlmTarget) -> Arc, + target_unavailable: impl Fn(&Request, &str), + fallback_decision: impl Fn(&LlmTarget, &LlmTarget, RoutingFallbackReason) -> Arc, ) -> Result { loop { let result = driver .call_llm_target(ctx.clone(), &target, request.clone(), decision.clone()) .await; let Err(error) = result else { return result }; - let LibsyError::ClientCall { - target: failed, - source: LlmClientError::ContextWindowExceeded { .. }, - } = &error - else { + let Some((failed, reason)) = classify_fallback(&error) else { return Err(error); }; // A target already excluded means the pool is spent; surface the client error - // so the caller still sees a context overflow rather than an internal failure. + // so the caller still sees the concrete upstream failure. if !ctx.exclude_target(failed) { return Err(error); } - evictions.record(identity, failed); + match reason { + RoutingFallbackReason::ContextWindow => evictions.record(identity, failed), + RoutingFallbackReason::Unavailable => target_unavailable(&request, failed), + } let Ok(next) = targets.resolve_target(&target.semantic_name, &ctx) else { return Err(error); }; - decision = fallback_decision(&target, &next); + decision = fallback_decision(&target, &next, reason); target = next; driver.info(ctx.clone(), decision.clone()).await?; } @@ -825,6 +846,61 @@ mod tests { LibsyError::external("test", TestError(message)) } + fn classified_client_error(source: LlmClientError) -> Option { + classify_fallback(&LibsyError::client_call("target", source)).map(|(_, reason)| reason) + } + + #[test] + fn route_fallback_only_accepts_context_and_unavailable_failures() { + assert_eq!( + classified_client_error(LlmClientError::ContextWindowExceeded { + model: "target".to_string(), + message: "too long".to_string(), + }), + Some(RoutingFallbackReason::ContextWindow) + ); + for source in [ + LlmClientError::Transport { + source: Box::new(std::io::Error::other("connection failed")), + }, + LlmClientError::Timeout { + source: Box::new(std::io::Error::other("request timed out")), + }, + ] { + assert_eq!( + classified_client_error(source), + Some(RoutingFallbackReason::Unavailable) + ); + } + for (status, expected) in [ + (400, None), + (401, None), + (403, Some(RoutingFallbackReason::Unavailable)), + (404, None), + (408, Some(RoutingFallbackReason::Unavailable)), + (409, None), + (429, Some(RoutingFallbackReason::Unavailable)), + (499, None), + (500, Some(RoutingFallbackReason::Unavailable)), + (599, Some(RoutingFallbackReason::Unavailable)), + (600, None), + ] { + assert_eq!( + classified_client_error(LlmClientError::UpstreamHttp { + status, + body: "failed".to_string(), + }), + expected + ); + } + assert_eq!( + classified_client_error(LlmClientError::InvalidResponse { + source: Box::new(std::io::Error::other("invalid response")), + }), + None + ); + } + /// Mock client that echoes back the target name it was called with. struct EchoClient; diff --git a/crates/libsy/src/core/classifier.rs b/crates/libsy/src/core/classifier.rs index ff2dc50f4..7a7451f23 100644 --- a/crates/libsy/src/core/classifier.rs +++ b/crates/libsy/src/core/classifier.rs @@ -76,6 +76,11 @@ pub trait Classifier: Send + Sync { None } + /// Drops retained routing state when `target` was unavailable for `request`. + /// + /// Stateless classifiers do not need to implement this hook. + fn target_unavailable(&self, _request: &Request, _target: &str) {} + /// Score the classifier's targets given the current state and request. /// /// When present, `driver` lets a classifier offload model calls. It is `None` diff --git a/crates/protocol/src/client.rs b/crates/protocol/src/client.rs index 05124509b..131ce00dc 100644 --- a/crates/protocol/src/client.rs +++ b/crates/protocol/src/client.rs @@ -110,6 +110,25 @@ pub enum LlmClientError { General(String), } +/// Why routing replaced a selected target with another eligible target. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum RoutingFallbackReason { + /// The selected target rejected the request because its context window was too small. + ContextWindow, + /// The selected target was unavailable after its client retries finished. + Unavailable, +} + +impl RoutingFallbackReason { + /// Stable value used by logs and statistics. + pub const fn as_str(self) -> &'static str { + match self { + Self::ContextWindow => "context_window", + Self::Unavailable => "unavailable", + } + } +} + /// A decision/trace object produced by an algorithm. /// /// Carried as a trait object (not a generic parameter) so a stream consumer can @@ -127,6 +146,10 @@ pub trait Decision: Send + Sync { fn is_routed_call(&self) -> bool { true } + /// Why this decision replaced an earlier selected target, when it did. + fn fallback_reason(&self) -> Option { + None + } /// A human-readable explanation of the decision, for logs and traces. fn reasoning(&self) -> Option<&str>; /// Downcast handle: a consumer that knows the algorithm can recover the diff --git a/crates/switchyard-server/src/lib.rs b/crates/switchyard-server/src/lib.rs index 820a800b4..8f5c03461 100644 --- a/crates/switchyard-server/src/lib.rs +++ b/crates/switchyard-server/src/lib.rs @@ -689,12 +689,20 @@ async fn handle_llm_request( Ok(result) => result, Err(error) => return algorithm_error(error), }; + for reason in trace + .iter() + .filter_map(|decision| decision.fallback_reason()) + { + state.stats.record_routing_fallback(reason); + } // Metrics, response body, and routing header all read the same decision, so // the model they name can never disagree. An empty trace leaves the body with // the id the upstream reported. let decision = trace.last(); let response = if let Some(decision) = decision { + let routing_log_context = routing_log_context + .map(|context| context.with_fallback_reason(decision.fallback_reason())); let cache_eligible = cache_probe .as_ref() .map(|probe| { diff --git a/crates/switchyard-server/src/routing_log.rs b/crates/switchyard-server/src/routing_log.rs index 88156dc92..c60f47c5e 100644 --- a/crates/switchyard-server/src/routing_log.rs +++ b/crates/switchyard-server/src/routing_log.rs @@ -12,7 +12,7 @@ use std::time::SystemTime; use humantime::format_rfc3339_millis; use serde::{Deserialize, Serialize}; -use switchyard_protocol::Usage; +use switchyard_protocol::{RoutingFallbackReason, Usage}; use crate::usage_metrics::token_usage; use crate::{ServerError, ServerResult}; @@ -56,6 +56,7 @@ impl RoutingLog { session_id: context.session_id.map(Cow::Owned), model: model.into(), tier: tier.unwrap_or("").into(), + fallback_reason: context.fallback_reason.map(Cow::Borrowed), prompt_tokens: usage.prompt_tokens, cached_tokens: usage.cached_tokens, cache_creation_tokens: usage.cache_creation_tokens, @@ -95,12 +96,13 @@ pub(crate) fn snapshot( Ok((snapshot.total_calls > 0).then_some(snapshot)) } -/// Request headers retained until terminal usage is available. +/// Request fields retained until terminal usage and routing are available. #[derive(Clone)] pub(crate) struct RoutingLogContext { task: Option, trial_id: Option, session_id: Option, + fallback_reason: Option<&'static str>, } impl RoutingLogContext { @@ -109,8 +111,14 @@ impl RoutingLogContext { task: nonempty_header(headers, TASK_HEADER).map(|s| s.to_string()), trial_id: nonempty_header(headers, TRIAL_ID_HEADER).map(|s| s.to_string()), session_id: nonempty_header(headers, SESSION_ID_HEADER).map(|s| s.to_string()), + fallback_reason: None, } } + + pub(crate) fn with_fallback_reason(mut self, reason: Option) -> Self { + self.fallback_reason = reason.map(RoutingFallbackReason::as_str); + self + } } /// One appended routing record, and the read schema [`snapshot`] parses back, @@ -128,6 +136,8 @@ struct RoutingRecord<'a> { session_id: Option>, model: Cow<'a, str>, tier: Cow<'a, str>, + #[serde(borrow, skip_serializing_if = "Option::is_none")] + fallback_reason: Option>, prompt_tokens: u64, cached_tokens: u64, cache_creation_tokens: u64, diff --git a/crates/switchyard-server/src/stats/accumulator.rs b/crates/switchyard-server/src/stats/accumulator.rs index 51dc96a3d..c7a287991 100644 --- a/crates/switchyard-server/src/stats/accumulator.rs +++ b/crates/switchyard-server/src/stats/accumulator.rs @@ -8,6 +8,7 @@ use std::sync::Arc; use parking_lot::{Mutex, MutexGuard}; use serde::Serialize; +use switchyard_protocol::RoutingFallbackReason; use super::cache_eligibility::PrefixProbe; @@ -103,6 +104,19 @@ impl StatsAccumulator { self.lock().routing_overhead.record(routing_overhead_ms); } + /// Records one target replacement by its route-level cause. + pub(crate) fn record_routing_fallback(&self, reason: RoutingFallbackReason) { + let fallbacks = &mut self.lock().routing_fallbacks; + match reason { + RoutingFallbackReason::ContextWindow => { + fallbacks.context_window = fallbacks.context_window.saturating_add(1); + } + RoutingFallbackReason::Unavailable => { + fallbacks.unavailable = fallbacks.unavailable.saturating_add(1); + } + } + } + /// Records one successful classifier or judge call. pub(crate) fn record_classifier_success( &self, @@ -168,6 +182,7 @@ struct StatsAccumulatorInner { total_requests: u64, total_errors: u64, routing_overhead: LatencyHistogram, + routing_fallbacks: RoutingFallbackStats, by_classifier: BTreeMap, classifier_requests: u64, classifier_errors: u64, @@ -202,6 +217,7 @@ impl StatsAccumulatorInner { models, tiers: tier_snapshots(&self.by_tier, total_tokens.total, self.total_requests), routing_overhead: self.routing_overhead.snapshot(), + routing_fallbacks: self.routing_fallbacks, classifier, } } @@ -330,9 +346,17 @@ pub(crate) struct StatsSnapshot { pub models: BTreeMap, pub tiers: BTreeMap, pub routing_overhead: LatencyHistogramSnapshot, + pub routing_fallbacks: RoutingFallbackStats, pub classifier: ClassifierStatsSnapshot, } +/// Route-level target replacements grouped by their fixed, low-cardinality cause. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Serialize)] +pub(crate) struct RoutingFallbackStats { + pub context_window: u64, + pub unavailable: u64, +} + #[derive(Clone, Debug, Default, PartialEq, Serialize)] pub(crate) struct ClassifierStatsSnapshot { pub total_requests: u64, diff --git a/crates/switchyard-server/tests/server.rs b/crates/switchyard-server/tests/server.rs index 0dde3915d..3f3907558 100644 --- a/crates/switchyard-server/tests/server.rs +++ b/crates/switchyard-server/tests/server.rs @@ -90,6 +90,14 @@ async fn upstream_chat( } let model = body["model"].as_str().unwrap_or("unknown").to_string(); + let prompt = body["messages"][0]["content"].as_str().unwrap_or(""); + if (model == "model/weak" && prompt == "unavailable") || prompt == "all-unavailable" { + return ( + StatusCode::SERVICE_UNAVAILABLE, + Json(json!({"error": {"message": "upstream is unavailable"}})), + ) + .into_response(); + } if model == "model/weak" && body["messages"][0]["content"] == "overflow" { return ( StatusCode::BAD_REQUEST, @@ -262,6 +270,10 @@ async fn stats_exposes_the_exact_empty_schema_and_no_legacy_alias() -> TestResul "p50_ms": 0.0, "p99_ms": 0.0 }, + "routing_fallbacks": { + "context_window": 0, + "unavailable": 0 + }, "classifier": { "total_requests": 0, "total_errors": 0, @@ -525,9 +537,8 @@ fn load_test_config(toml: &str) -> TestResult { Ok(load_server_state(config.path())?) } -/// A `random` route that always selects `first` (weight 1 vs 0), so a test can drive the -/// overflow fallback from `first` to `second` deterministically. -fn overflow_fallback_state(base_url: &str) -> TestResult { +/// A `random` route that selects `first` before any request-local fallback. +fn fallback_state(base_url: &str) -> TestResult { load_test_config(&format!( r#" schema_version = 1 @@ -1438,7 +1449,7 @@ async fn routing_log_exposes_session_stats() -> TestResult { #[tokio::test] async fn overflow_history_is_scoped_to_agent_and_session_lifetime() -> TestResult { let upstream = MockUpstream::start().await?; - let state = overflow_fallback_state(&upstream.base_url)?; + let state = fallback_state(&upstream.base_url)?; let app = build_switchyard_router(state); let child_a = [ ("x-switchyard-session-id", "shared-session"), @@ -1523,6 +1534,103 @@ async fn overflow_history_is_scoped_to_agent_and_session_lifetime() -> TestResul Ok(()) } +#[tokio::test] +async fn unavailable_target_fails_over_across_endpoints_and_stops_when_exhausted() -> TestResult { + let upstream = MockUpstream::start().await?; + let temp_dir = tempfile::tempdir()?; + let log_path = temp_dir.path().join("routing.jsonl"); + let state = fallback_state(&upstream.base_url)?.with_routing_log(&log_path)?; + let app = build_switchyard_router(state); + let cases = [ + ( + "/v1/chat/completions", + json!({ + "model": ROUTE_MODEL, + "messages": [{"role": "user", "content": "unavailable"}] + }), + ), + ( + "/v1/messages", + json!({ + "model": ROUTE_MODEL, + "max_tokens": 16, + "messages": [{"role": "user", "content": "unavailable"}] + }), + ), + ( + "/v1/responses", + json!({"model": ROUTE_MODEL, "input": "unavailable"}), + ), + ]; + + for (path, body) in cases { + let previous_call_count = upstream.calls.lock().await.len(); + let response = send(&app, "POST", path, Some(body)).await?; + assert_eq!(response.status, StatusCode::OK); + assert_eq!( + response + .headers + .get("x-model-router-selected-model") + .and_then(|value| value.to_str().ok()), + Some("model/strong") + ); + assert_eq!( + response + .headers + .get("x-model-router-rationale") + .and_then(|value| value.to_str().ok()), + Some("model/weak was unavailable; fell back to model/strong") + ); + let calls = upstream.calls.lock().await; + assert_eq!( + calls[previous_call_count..] + .iter() + .map(|call| call["model"].as_str().unwrap_or("")) + .collect::>(), + ["model/weak", "model/strong"] + ); + } + + let stats = send(&app, "GET", "/v1/stats", None).await?.json()?; + assert_eq!(stats["routing_fallbacks"]["unavailable"], 3); + assert_eq!(stats["routing_fallbacks"]["context_window"], 0); + + let records = std::fs::read_to_string(&log_path)?; + let records = records + .lines() + .map(serde_json::from_str::) + .collect::, _>>()?; + assert_eq!(records.len(), 3); + assert!(records.iter().all(|record| { + record["model"] == "model/strong" && record["fallback_reason"] == "unavailable" + })); + + let previous_call_count = upstream.calls.lock().await.len(); + let response = send( + &app, + "POST", + "/v1/chat/completions", + Some(json!({ + "model": ROUTE_MODEL, + "messages": [{"role": "user", "content": "all-unavailable"}] + })), + ) + .await?; + assert_eq!(response.status, StatusCode::SERVICE_UNAVAILABLE); + let error = response.json()?; + assert_eq!(error["error"]["type"], "upstream_error"); + assert_eq!(error["error"]["code"], "upstream_error"); + let calls = upstream.calls.lock().await; + assert_eq!( + calls[previous_call_count..] + .iter() + .map(|call| call["model"].as_str().unwrap_or("")) + .collect::>(), + ["model/weak", "model/strong"] + ); + Ok(()) +} + #[tokio::test] async fn streaming_response_is_framed_for_the_inbound_api() -> TestResult { let (_upstream, app) = test_app(&[(ROUTE_MODEL, &["model/a"])]).await?;