From 6f3bac769f57a7761786bd871ebb033e94b3ea82 Mon Sep 17 00:00:00 2001
From: Ayush <shekharayush56@gmail.com>
Date: Wed, 16 Sep 2026 00:26:40 +0530
Subject: [PATCH] fix(router): bound mid-stream continuation to the caller's
 output allowance

A continuation handed the fallback the full max_tokens / max_completion_tokens
even though part of the answer had already streamed, so the total could exceed
the caller's limit and repeated interruptions could re-grant the allowance on
every hop, amplifying spend. Reduce each configured ceiling by the tokens already
emitted before re-entry, and decline the continuation (surface the original
error) when the allowance is exhausted. The reductions compound across hops
because the trimmed ceiling is what the next hop sees.

Also add the missing explanatory comment on the model_info fall-through except.
---
 litellm/router.py                             | 33 +++++++++++++++++++
 .../continuation_prefill_check.py             |  1 +
 tests/test_litellm/test_router.py             | 20 +++++++++++
 3 files changed, 54 insertions(+)

diff --git a/litellm/router.py b/litellm/router.py
index 6a43a6a05c0..474c284c389 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -2857,6 +2857,17 @@ class Router:
                             MID_STREAM_CONTINUATION_MARKER,
                         )
 
+                        emitted_tokens: Final = int(
+                            getattr(complete_response_object_usage, "completion_tokens", 0) or 0
+                        )
+                        reduced_ceilings: Final = self._continuation_output_ceilings(initial_kwargs, emitted_tokens)
+                        if reduced_ceilings is None:
+                            # The caller's output allowance is already spent; surface the
+                            # error rather than grant a fresh allowance on this fallback hop.
+                            if e.original_exception is not None:
+                                raise e.original_exception from e
+                            raise
+                        initial_kwargs.update(reduced_ceilings)
                         initial_kwargs["messages"] = self._build_completion_continuation_input(
                             messages, e.generated_content
                         )
@@ -3076,6 +3087,28 @@ class Router:
             return False
         return True
 
+    @staticmethod
+    def _continuation_output_ceilings(
+        request_kwargs: Mapping[str, object],
+        emitted_tokens: int,
+    ) -> Mapping[str, int] | None:
+        """The max_tokens / max_completion_tokens a continuation should carry, each
+        reduced by the tokens already emitted so the whole answer stays within the
+        caller's original allowance instead of getting a fresh one on every fallback
+        hop. Reductions compound across hops because the trimmed ceiling is what the
+        next hop sees. Returns None when the allowance is already exhausted, so the
+        stream must not be continued."""
+        ceilings: Final = MappingProxyType(
+            {
+                key: value - emitted_tokens
+                for key in ("max_tokens", "max_completion_tokens")
+                if isinstance(value := request_kwargs.get(key), int)
+            }
+        )
+        if ceilings and min(ceilings.values()) <= 0:
+            return None
+        return ceilings
+
     @staticmethod
     def _build_completion_continuation_input(
         messages: Sequence[Mapping[str, str]],
diff --git a/litellm/router_utils/pre_call_checks/continuation_prefill_check.py b/litellm/router_utils/pre_call_checks/continuation_prefill_check.py
index ccd454bc207..c640374858a 100644
--- a/litellm/router_utils/pre_call_checks/continuation_prefill_check.py
+++ b/litellm/router_utils/pre_call_checks/continuation_prefill_check.py
@@ -46,6 +46,7 @@ def _deployment_supports_prefill(deployment: object) -> bool:
         if isinstance(declared, bool):
             return declared
     except ValidationError:
+        # No usable model_info override; fall through to the cost-map lookup.
         pass
     try:
         litellm_params: Final = _STR_KEYED_DICT_ADAPTER.validate_python(deployment_map.get("litellm_params"))
diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py
index c37a23df7d1..987b96b3228 100644
--- a/tests/test_litellm/test_router.py
+++ b/tests/test_litellm/test_router.py
@@ -2363,6 +2363,26 @@ def test_build_completion_continuation_input_folds_into_existing_prefill():
     assert twice[-1] == {"role": "assistant", "content": "part one part two", "prefix": True}
 
 
+def test_continuation_output_ceilings_reduces_by_emitted_tokens():
+    """A continuation must complete within the caller's original allowance, so each
+    output ceiling is reduced by the tokens already emitted."""
+    assert litellm.Router._continuation_output_ceilings({"max_tokens": 100}, 30) == {"max_tokens": 70}
+    assert litellm.Router._continuation_output_ceilings({"max_tokens": 100, "max_completion_tokens": 40}, 25) == {
+        "max_tokens": 75,
+        "max_completion_tokens": 15,
+    }
+    # no ceiling configured -> nothing to reduce, continuation proceeds as before
+    assert litellm.Router._continuation_output_ceilings({}, 50) == {}
+
+
+def test_continuation_output_ceilings_none_when_allowance_exhausted():
+    """When the emitted tokens already meet or exceed a ceiling, there is no budget
+    left to continue, so the helper signals a decline rather than a fresh allowance."""
+    assert litellm.Router._continuation_output_ceilings({"max_tokens": 20}, 20) is None
+    assert litellm.Router._continuation_output_ceilings({"max_tokens": 20}, 25) is None
+    assert litellm.Router._continuation_output_ceilings({"max_tokens": 100, "max_completion_tokens": 10}, 10) is None
+
+
 def test_mid_stream_continuation_eligible_allows_text_response_format():
     """response_format={"type": "text"} is the unconstrained default and must stay
     eligible, unlike json_object / json_schema."""