fix(routing): scope context-overflow history by session and agent

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,
@@ -26,11 +27,13 @@ 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};
-use switchyard_protocol::{Context, Decision, Request, Response};
+use switchyard_protocol::{Context, Decision, Request, Response, RoutingFallbackReason};
 
 struct SessionState<S> {
     state: Arc<AsyncMutex<S>>,
@@ -55,6 +58,7 @@ pub struct FallThroughDecision {
     /// Human-readable explanation of the selection.
     pub reasoning: String,
     tier: Option<&'static str>,
+    fallback_reason: Option<RoutingFallbackReason>,
 }
 
 impl Decision for FallThroughDecision {
@@ -66,6 +70,10 @@ impl Decision for FallThroughDecision {
         self.tier
     }
 
+    fn fallback_reason(&self) -> Option<RoutingFallbackReason> {
+        self.fallback_reason
+    }
+
     fn reasoning(&self) -> Option<&str> {
         Some(&self.reasoning)
     }
@@ -198,9 +206,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,15 +229,21 @@ where
         ctx: Context,
         driver: Driver,
         request: Request,
-        session: Option<&str>,
     ) -> Result<Response> {
         // 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 {
+        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?
@@ -250,16 +262,21 @@ where
         match served {
             Some(response) => Ok(response),
             None => {
-                algorithm::call_llm_with_overflow_fallback(
+                algorithm::call_llm_with_fallback(
                     ctx,
                     &driver,
                     &self.targets,
                     target,
                     decision,
                     request,
-                    session,
+                    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
             }
@@ -271,21 +288,29 @@ 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.
-    fn fallback_decision(&self, from: &LlmTarget, to: &LlmTarget) -> Arc<dyn Decision> {
+    /// The decision published when a route-level failure selects a different target.
+    fn fallback_decision(
+        &self,
+        deciding: &dyn Classifier<S>,
+        from: &LlmTarget,
+        to: &LlmTarget,
+        reason: RoutingFallbackReason,
+    ) -> Arc<dyn Decision> {
+        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),
         })
     }
 
@@ -309,29 +334,30 @@ where
         ctx: &Context,
         driver: &Driver,
         request: &mut Request,
-    ) -> Result<(LlmTarget, Arc<dyn Decision>, Option<Response>)> {
+    ) -> Result<(
+        LlmTarget,
+        Arc<dyn Decision>,
+        Option<Response>,
+        Arc<dyn Classifier<S>>,
+    )> {
         // 1. Processor chain accumulates request-side facts into the composition's state.
         for processor in &self.processors {
             processor.process(state, Event::Request(request)).await?;
         }
 
         // 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<Score> = None;
-        let mut deciding: Option<&Arc<dyn Classifier<S>>> = None;
-        let mut served: Option<Response> = 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(),
             });
@@ -351,7 +377,8 @@ where
         let decision: Arc<dyn Decision> = 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?;
 
@@ -365,7 +392,7 @@ where
             processor.process(state, event).await?;
         }
 
-        Ok((target, decision, served))
+        Ok((target, decision, served, deciding))
     }
 }
 
@@ -434,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,
@@ -791,6 +818,45 @@ mod tests {
         calls: Option<Arc<Mutex<Vec<String>>>>,
     }
 
+    /// Returns 503 for `weak` and records every routed call.
+    struct UnavailableClient(Arc<Mutex<Vec<String>>>);
+
+    #[async_trait]
+    impl RoutedLlmClient for UnavailableClient {
+        async fn call(
+            &self,
+            _ctx: Context,
+            _request: Request,
+            decision: Arc<dyn Decision>,
+        ) -> std::result::Result<Response, LlmClientError> {
+            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<Mutex<Vec<String>>>) -> LlmTargetSet {
+        let client: Arc<dyn RoutedLlmClient> = 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(
@@ -924,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(

crates/libsy/src/algorithms/llm_class.rs

@@ -517,7 +517,7 @@ impl Classifier<State> 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(),

crates/libsy/src/algorithms/util/affinity.rs

@@ -24,26 +24,14 @@ 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};
 
 /// Upper bound on retained assignments, keeping the process-local map from growing
 /// 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<HashMap<AffinityKey, String>>,
+    assignments: Mutex<HashMap<RoutingIdentity, String>>,
 }
 
 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<AffinityKey> {
-        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<RoutingIdentity> {
+        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()
@@ -177,6 +157,19 @@ impl<S> Classifier<S> 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,
@@ -201,7 +194,7 @@ where
 }
 
 /// Evicts one arbitrary assignment when the map has reached [`MAX_ASSIGNMENTS`].
-fn evict_if_full(assignments: &mut HashMap<AffinityKey, String>) {
+fn evict_if_full(assignments: &mut HashMap<RoutingIdentity, String>) {
     if assignments.len() >= MAX_ASSIGNMENTS
         && let Some(evicted) = assignments.keys().next().cloned()
     {