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

@@ -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]],

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"))

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."""