fix(router): keep the context window retry inside the deployments the router would pick

Count the deployments left to try with the same pick filters the router
uses (order, pinning, tags, pre-call checks), resolve the refused model
from wildcard copies, only let a matching context_window_fallbacks chain
skip the retry, and hand the original context window error to the caller
and fallbacks when the retries fail for another reason.

litellm-rust/crates/router-core/src/call.rs

@@ -289,16 +289,54 @@ fn serving_the_refused_model<'a>(
     refused: &'a str,
     healthy: &'a [Arc<Deployment>],
 ) -> impl Iterator<Item = String> + 'a {
-    let model = snapshot
-        .registry
-        .deployment(refused)
-        .map(|deployment| deployment_model(&deployment));
+    let model = healthy
+        .iter()
+        .find(|d| d.id == refused)
+        .map(|d| deployment_model(d))
+        .or_else(|| {
+            snapshot
+                .registry
+                .deployment(refused)
+                .map(|deployment| deployment_model(&deployment))
+        });
     healthy
         .iter()
         .filter(move |d| d.id == refused || model.as_ref() == Some(&deployment_model(d)))
         .map(|d| d.id.clone())
 }
 
+fn context_window_refusal(failed: &Failed) -> Option<Failed> {
+    match &*failed.error {
+        Error::Attempt(failure) if failure.is_a("ContextWindowExceededError") => Some(Failed {
+            error: Box::new(Error::Attempt(failure.clone())),
+            deployment_id: failed.deployment_id.clone(),
+            routed: failed.routed.clone(),
+            served_as: failed.served_as.clone(),
+        }),
+        _ => None,
+    }
+}
+
+/// `fallback_lookup_groups`: the group a pre-routing hook selected, the group being served, then
+/// the group the request was for.
+fn fallback_lookup<'a>(
+    routed: Option<&'a str>,
+    served_as: Option<&'a str>,
+    target: &'a Target,
+    request: &'a RouterRequest,
+) -> Vec<&'a str> {
+    let mut lookup: Vec<&str> = routed.into_iter().collect();
+    for group in served_as
+        .into_iter()
+        .chain([target.group.as_str(), request.model.as_str()])
+    {
+        if !lookup.contains(&group) {
+            lookup.push(group);
+        }
+    }
+    lookup
+}
+
 impl Core {
     pub(crate) async fn call(
         &self,
@@ -775,17 +813,7 @@ impl Core {
         let failure = retried_failure(error);
         let is_a = |name: &str| failure.as_ref().is_some_and(|f| f.is_a(name));
         let lists = Lists::of(snapshot, request);
-        // `fallback_lookup_groups`: the group a pre-routing hook selected, the group being served,
-        // then the group the request was for.
-        let mut lookup: Vec<&str> = routed.into_iter().collect();
-        for group in served_as
-            .into_iter()
-            .chain([target.group.as_str(), request.model.as_str()])
-        {
-            if !lookup.contains(&group) {
-                lookup.push(group);
-            }
-        }
+        let lookup = fallback_lookup(routed, served_as, target, request);
         // `order_model_group`: the pre-routing selection's order levels, else the level's own.
         let order_group = routed.unwrap_or(target.group.as_str());
         let chain = |list: &Fallbacks| {
@@ -987,6 +1015,11 @@ impl Core {
             .flatten();
         let policy_applies = policy_retries.is_some();
         num_retries = policy_retries.unwrap_or(num_retries);
+        let refusal = if policy_applies {
+            None
+        } else {
+            context_window_refusal(&failed)
+        };
         drop(first_failure);
         // `_async_get_healthy_deployments` finds no group for a deployment id, so a call by id
         // counts no deployments when deciding whether to retry.
@@ -1021,32 +1054,41 @@ impl Core {
             // keeps going, else `should_retry_this_error` decides (it raises when nothing is
             // healthy).
             let Some(failure) = retried_failure(&failed.error) else {
-                return Err(failed);
+                return Err(refusal.unwrap_or(failed));
             };
             let refused_by = failure
                 .failed_deployment_id
                 .as_ref()
                 .or(failed.deployment_id.as_ref());
-            if let Some(refused) = refused_by
-                && !policy_applies
-                && failure.is_a("ContextWindowExceededError")
-            {
-                let serving: Vec<String> =
-                    serving_the_refused_model(snapshot, refused, &healthy_deployments)
-                        .filter(|id| !inputs.skipped_ids.contains(id))
-                        .collect();
-                inputs.skipped_ids.extend(serving);
-            }
-            let untried = refused_by.map_or(0, |refused| {
-                healthy_deployments
-                    .iter()
-                    .filter(|d| d.id != *refused && !inputs.skipped_ids.contains(&d.id))
-                    .count()
-            });
+            let untried = match refused_by {
+                Some(refused) if !policy_applies && failure.is_a("ContextWindowExceededError") => {
+                    let serving: Vec<String> =
+                        serving_the_refused_model(snapshot, refused, &healthy_deployments)
+                            .filter(|id| !inputs.skipped_ids.contains(id))
+                            .collect();
+                    inputs.skipped_ids.extend(serving);
+                    self.untried(snapshot, op, request, &target.group, &inputs)
+                        .await
+                }
+                _ => 0,
+            };
+            let lookup = fallback_lookup(
+                failed.routed.as_deref(),
+                failed.served_as.as_deref(),
+                target,
+                request,
+            );
             if !policy_applies
-                && !self.should_retry(snapshot, request, &failure, (healthy, untried), all)
+                && !self.should_retry(
+                    snapshot,
+                    request,
+                    &failure,
+                    &lookup,
+                    (healthy, untried),
+                    all,
+                )
             {
-                return Err(failed);
+                return Err(refusal.unwrap_or(failed));
             }
             let uncounted = request.uncounted_retries.filter(|_| !policy_applies);
             if retry == 0 && uncounted.is_none() {
@@ -1061,15 +1103,15 @@ impl Core {
                     model: target.group.clone(),
                     cause: RetryCause::of(&failed.error),
                 });
-                return Err(Failed {
+                return Err(refusal.unwrap_or(Failed {
                     error: Box::new(Error::Attempt(AttemptError::host_raised(&["Exception"]))),
                     deployment_id: None,
                     routed: None,
                     served_as: None,
-                });
+                }));
             }
             if num_retries == 0 {
-                return Err(failed);
+                return Err(refusal.unwrap_or(failed));
             }
             if retry == 0 {
                 self.log_retry(target, at, 0, &failed.error);
@@ -1087,7 +1129,7 @@ impl Core {
                     .iter()
                     .all(|d| inputs.skipped_ids.contains(&d.id))
             {
-                return Err(failed);
+                return Err(refusal.unwrap_or(failed));
             }
             let delay = self.time_to_sleep(
                 snapshot,
@@ -1130,7 +1172,7 @@ impl Core {
                     return Ok(response);
                 }
                 Err(next) if next.error.attempt().is_some_and(|f| f.agentic_followup) => {
-                    return Err(next);
+                    return Err(refusal.unwrap_or(next));
                 }
                 Err(next) => {
                     self.pick_failed(at, target, &next.error);
@@ -1138,7 +1180,25 @@ impl Core {
                 }
             }
         }
-        Err(failed)
+        Err(refusal.unwrap_or(failed))
+    }
+
+    async fn untried(
+        &self,
+        snapshot: &RouterStateSnapshot,
+        op: &Op,
+        request: &RouterRequest,
+        group: &str,
+        inputs: &SelectInputs,
+    ) -> usize {
+        self.candidates(snapshot, op, request, group, inputs)
+            .await
+            .map_or(0, |c| {
+                c.healthy
+                    .iter()
+                    .filter(|d| !inputs.skipped_ids.contains(&d.id))
+                    .count()
+            })
     }
 
     /// `error` is a pick's: reported as it happens, whatever the level then does with it.
@@ -1203,6 +1263,7 @@ impl Core {
         snapshot: &RouterStateSnapshot,
         request: &RouterRequest,
         failure: &AttemptError,
+        lookup: &[&str],
         (healthy, untried): (usize, usize),
         all: usize,
     ) -> bool {
@@ -1213,7 +1274,9 @@ impl Core {
                 healthy,
                 untried,
                 all,
-                has_context_window_fallbacks: lists.context_window.is_some(),
+                has_context_window_fallbacks: lists
+                    .context_window
+                    .is_some_and(|list| list.lookup_groups(lookup, lists.keys).is_some()),
                 has_content_policy_fallbacks: lists.content_policy.is_some(),
                 has_regular_fallbacks: lists.has_regular(),
             },

litellm-rust/crates/router-core/src/call/stream.rs

@@ -10,7 +10,7 @@ use std::sync::Arc;
 use serde::Deserialize;
 
 use super::{
-    Failed, FallbackChain, FallbackHop, Lists, Target, Tried, group_members,
+    Failed, FallbackChain, FallbackHop, Lists, Target, Tried, fallback_lookup, group_members,
     serving_the_refused_model,
 };
 use crate::{
@@ -247,34 +247,39 @@ impl Router {
             .or(origin.deployment_id.as_deref());
         let (budget, policy_applies) =
             stream_retry_budget(&snapshot, &request, &origin, failure.kind);
-        let skipped_ids: Vec<String> = match refused_by {
-            Some(refused) if !policy_applies && failure.is_a("ContextWindowExceededError") => {
-                serving_the_refused_model(&snapshot, refused, &healthy_deployments).collect()
-            }
-            _ => Vec::new(),
-        };
-        let untried = refused_by.map_or(0, |refused| {
-            healthy_deployments
-                .iter()
-                .filter(|d| d.id != refused && !skipped_ids.contains(&d.id))
-                .count()
-        });
-        let inputs = SelectInputs {
+        let refusal = (!policy_applies && failure.is_a("ContextWindowExceededError"))
+            .then(|| failure.clone());
+        let mut inputs = SelectInputs {
             target_order: target.order,
             requested: Some(target.group.clone()),
             fallback_hop: at.level > 0,
-            skipped_ids,
             ..SelectInputs::default()
         };
+        let untried = match refused_by {
+            Some(refused) if refusal.is_some() => {
+                inputs.skipped_ids =
+                    serving_the_refused_model(&snapshot, refused, &healthy_deployments).collect();
+                self.core
+                    .untried(&snapshot, &op, &request, &target.group, &inputs)
+                    .await
+            }
+            _ => 0,
+        };
+        let lookup = fallback_lookup(origin.routed.as_deref(), None, &target, &request);
         let mut last = failure;
         for attempt in origin.attempted_retries..budget {
             if last.agentic_followup {
                 break;
             }
             if !policy_applies
-                && !self
-                    .core
-                    .should_retry(&snapshot, &request, &last, (healthy, untried), all)
+                && !self.core.should_retry(
+                    &snapshot,
+                    &request,
+                    &last,
+                    &lookup,
+                    (healthy, untried),
+                    all,
+                )
             {
                 break;
             }
@@ -320,11 +325,11 @@ impl Router {
                             break;
                         }
                     }
-                    other => return Err(other),
+                    other => return Err(refusal.map_or(other, Error::Attempt)),
                 },
             }
         }
-        Err(Error::Attempt(last))
+        Err(Error::Attempt(refusal.unwrap_or(last)))
     }
 }
 

litellm-rust/crates/router-core/src/retry.rs

@@ -65,7 +65,8 @@ pub struct RetryContext {
     pub healthy: usize,
     pub untried: usize,
     pub all: usize,
-    /// `context_window_fallbacks is not None`.
+    /// A `context_window_fallbacks` chain matches the failed group
+    /// (`context_window_fallback_model_group`).
     pub has_context_window_fallbacks: bool,
     /// `content_policy_fallbacks is not None`.
     pub has_content_policy_fallbacks: bool,

litellm-rust/crates/router-core/tests/call_loop.rs

@@ -973,6 +973,139 @@ async fn a_context_window_error_no_deployment_raised_is_not_retried(config: Rout
     assert_eq!((host.retries(), host.selected().len()), (0, 0));
 }
 
+fn refusing_once() -> Arc<ScriptedHost> {
+    Arc::new(ScriptedHost {
+        first_failure: Mutex::new(Some(context_window())),
+        ..ScriptedHost::default()
+    })
+}
+
+#[rstest]
+#[tokio::test]
+async fn a_context_window_error_takes_the_next_order_tier_instead_of_the_refusing_one() {
+    let host = refusing_once();
+    let router = router(tiers(true), json!({"num_retries": 2}), host.clone());
+    let response = router
+        .call(Op::from("acompletion"), request())
+        .await
+        .unwrap();
+    assert_eq!(
+        (
+            response.deployment_id.as_str(),
+            host.selected(),
+            host.retries()
+        ),
+        ("p2", vec!["p1".to_owned(), "p2".to_owned()], 0)
+    );
+}
+
+#[rstest]
+#[tokio::test]
+async fn a_context_window_error_on_a_pinned_deployment_is_not_retried(config: RouterConfig) {
+    let host = refusing_once();
+    let router = router(config, json!({"num_retries": 2}), host.clone());
+    let pinned = RouterRequest {
+        pinned_deployment_id: Some("a1".to_owned()),
+        ..request()
+    };
+    let refused = router.call(Op::from("acompletion"), pinned).await;
+    assert!(
+        matches!(refused, Err(Error::Attempt(failed)) if failed.kind == FailureKind::ContextWindowExceeded)
+    );
+    assert_eq!(host.selected(), ["a1"]);
+}
+
+#[rstest]
+#[tokio::test]
+async fn a_context_window_error_skips_the_wildcard_deployments_serving_that_model() {
+    let config: RouterConfig = serde_json::from_value(json!({
+        "model_list": [
+            {"model_name": "openai/*", "litellm_params": {"model": "openai/*"}, "model_info": {"id": "w1"}},
+            {"model_name": "openai/*", "litellm_params": {"model": "openai/*"}, "model_info": {"id": "w2"}},
+        ]
+    }))
+    .unwrap();
+    let host = refusing_once();
+    let router = router(config, json!({"num_retries": 2}), host.clone());
+    let wildcard = RouterRequest {
+        model: "openai/gpt-4o".to_owned(),
+        ..RouterRequest::default()
+    };
+    let refused = router.call(Op::from("acompletion"), wildcard).await;
+    assert!(
+        matches!(refused, Err(Error::Attempt(failed)) if failed.kind == FailureKind::ContextWindowExceeded)
+    );
+    assert_eq!(host.selected().len(), 1);
+}
+
+#[rstest]
+#[tokio::test]
+async fn a_context_window_list_for_another_group_leaves_the_retry_in_place(config: RouterConfig) {
+    let host = Arc::new(ScriptedHost::failing("primary", vec![context_window()]));
+    let router = router(
+        config,
+        json!({"num_retries": 2, "context_window_fallbacks": [{"backup": ["primary"]}]}),
+        host.clone(),
+    );
+    let response = router
+        .call(Op::from("acompletion"), request())
+        .await
+        .unwrap();
+    assert_eq!(
+        (response.model_group.as_str(), response.attempted_retries),
+        ("primary", 1)
+    );
+}
+
+fn provider_bad_request() -> AttemptError {
+    let mut failed = AttemptError::new(FailureKind::BadRequest, Some(400));
+    failed.names_provider = true;
+    failed
+}
+
+#[rstest]
+#[tokio::test]
+async fn a_retry_that_fails_after_a_context_window_error_falls_back_on_that_error(
+    config: RouterConfig,
+) {
+    let config = RouterConfig {
+        fallback_on_unsupported_params: Some(false),
+        ..config
+    };
+    let host = Arc::new(ScriptedHost::failing(
+        "primary",
+        vec![context_window(), provider_bad_request()],
+    ));
+    let router = router(
+        config,
+        json!({"num_retries": 2, "fallbacks": [{"primary": ["backup"]}]}),
+        host.clone(),
+    );
+    let response = router
+        .call(Op::from("acompletion"), request())
+        .await
+        .unwrap();
+    assert_eq!(
+        (response.deployment_id.as_str(), host.selected().len()),
+        ("b1", 3)
+    );
+}
+
+#[rstest]
+#[tokio::test]
+async fn a_retry_that_fails_after_a_context_window_error_raises_that_error(config: RouterConfig) {
+    let host = Arc::new(ScriptedHost::failing(
+        "primary",
+        vec![context_window(), server_error(), server_error()],
+    ));
+    let router = router(config, json!({"num_retries": 2}), host.clone());
+    let refused = router.call(Op::from("acompletion"), request()).await;
+    assert!(
+        matches!(refused, Err(Error::Attempt(failed)) if failed.kind == FailureKind::ContextWindowExceeded)
+    );
+    assert_eq!(host.selected().len(), 3);
+}
+
 #[rstest]
 #[tokio::test]
 async fn a_failed_agentic_followup_is_not_retried_but_still_falls_back(config: RouterConfig) {
@@ -1113,6 +1246,57 @@ async fn a_stream_context_window_error_retries_on_another_deployment(config: Rou
     );
 }
 
+#[rstest]
+#[tokio::test]
+async fn a_stream_context_window_error_in_an_order_tier_is_not_retried() {
+    let host = Arc::new(ScriptedHost::default());
+    let router = router(tiers(true), json!({}), host.clone());
+    let origin = StreamOrigin {
+        deployment_id: Some("p1".to_owned()),
+        ..stream_origin()
+    };
+    let refused = router
+        .retry_stream(
+            Op::from("acompletion"),
+            stream_request(),
+            origin,
+            context_window(),
+        )
+        .await;
+    assert!(
+        matches!(refused, Err(Error::Attempt(failed)) if failed.kind == FailureKind::ContextWindowExceeded)
+    );
+    assert!(host.selected().is_empty());
+}
+
+#[rstest]
+#[tokio::test]
+async fn a_stream_retry_that_fails_after_a_context_window_error_raises_that_error(
+    config: RouterConfig,
+) {
+    let host = Arc::new(ScriptedHost::failing(
+        "primary",
+        vec![server_error(), server_error(), server_error()],
+    ));
+    let router = router(config, json!({}), host.clone());
+    let origin = StreamOrigin {
+        deployment_id: Some("a1".to_owned()),
+        ..stream_origin()
+    };
+    let refused = router
+        .retry_stream(
+            Op::from("acompletion"),
+            stream_request(),
+            origin,
+            context_window(),
+        )
+        .await;
+    assert!(
+        matches!(refused, Err(Error::Attempt(failed)) if failed.kind == FailureKind::ContextWindowExceeded)
+    );
+    assert_eq!(host.selected().len(), 3);
+}
+
 fn fallback_list(value: serde_json::Value) -> Fallbacks {
     serde_json::from_value(value).unwrap()
 }

tests/unit/test_router/test_router_retries.py

@@ -439,7 +439,7 @@ def test_raise_context_window_exceeded_error():
 
 def test_raise_context_window_exceeded_error_no_retry():
     """
-    Do not Retry Context Window Exceeded Error, when context_window_fallbacks is None
+    Do not retry a Context Window Exceeded Error when no other healthy deployment is left to try
     """
     context_window_error = litellm.ContextWindowExceededError(
         message="Context window exceeded",