From 974da4c2c97ebeb14947be5b5e3bd99731c1b050 Mon Sep 17 00:00:00 2001 From: Ayush Date: Mon, 14 Sep 2026 17:31:05 +0530 Subject: [PATCH 01/18] feat(router): continue chat streams on fallback after a mid-stream break A chat-completions stream that breaks after content has been delivered was re-raised and the fallback deployment never ran (#40404), while the Responses-API path already continues via a prefilled assistant turn. Add the same for chat completions behind `enable_mid_stream_fallback_continuation` (opt-in, async-only), so a post-content break re-enters the fallback chain with the partial text as an assistant prefill instead of surfacing the error. Continuation is only attempted when it is safe: the stream emitted plain assistant text (no tool/function calls, thinking blocks, reasoning items, audio or images), the request is not constrained output (response_format or a forced tool_choice) and not merge-reasoning mode, and the fallback target's model supports assistant prefill. The target check runs as a deployment pre-call filter, so a chain with no prefill-capable deployment empties and the original error is surfaced rather than a duplicated or rejected request being sent. Plain reasoning_content is treated as out-of-band and does not block a continuation, matching the Responses-API path; Anthropic thinking is excluded because its signed thinking blocks cannot ride a text-only prefill. CustomStreamWrapper now latches whether a disqualifying delta was streamed and carries it on MidStreamFallbackError so the router can decide without re-scanning chunks. Default behavior is unchanged. --- litellm/exceptions.py | 2 + .../litellm_core_utils/streaming_handler.py | 49 ++++- litellm/router.py | 84 ++++++++- .../continuation_prefill_check.py | 51 +++++ litellm/utils.py | 14 ++ .../test_streaming_handler.py | 44 +++++ .../test_continuation_prefill_check.py | 69 +++++++ tests/test_litellm/test_router.py | 176 ++++++++++++++++++ 8 files changed, 476 insertions(+), 13 deletions(-) create mode 100644 litellm/router_utils/pre_call_checks/continuation_prefill_check.py create mode 100644 tests/test_litellm/router_utils/pre_call_checks/test_continuation_prefill_check.py diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 3f22a4b2dcd..dfda31b11c0 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -1124,6 +1124,7 @@ class MidStreamFallbackError(ServiceUnavailableError): num_retries: int | None = None, generated_content: str = "", is_pre_first_chunk: bool = False, + emitted_disqualifying_content: bool = False, ): original_status: Final = getattr(original_exception, "status_code", None) self.status_code = int(original_status) if original_status is not None else 503 @@ -1136,6 +1137,7 @@ class MidStreamFallbackError(ServiceUnavailableError): self.num_retries = num_retries self.generated_content = generated_content self.is_pre_first_chunk = is_pre_first_chunk + self.emitted_disqualifying_content = emitted_disqualifying_content # Create a response if one wasn't provided if response is None: diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 6a5a8832cc6..f62a23d2c42 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -272,6 +272,7 @@ class CustomStreamWrapper: self.holding_chunk = "" self.complete_response = "" self.response_uptil_now = "" + self._emitted_disqualifying_content = False _model_info: Final[dict] = litellm_params.model_info or {} _api_base: Final = get_api_base( @@ -1950,9 +1951,7 @@ class CustomStreamWrapper: if response.choices: choice = response.choices[0] if isinstance(choice, StreamingChoices): - self.response_uptil_now += choice.delta.get("content", "") or "" - else: - self.response_uptil_now += "" + self._accumulate_streamed_delta(choice.delta) self.rules.post_call_rules(input=self.response_uptil_now, model=self.model) # HANDLE STREAM OPTIONS self.chunks.append(response) @@ -2150,9 +2149,7 @@ class CustomStreamWrapper: if processed_chunk.choices: choice = processed_chunk.choices[0] if isinstance(choice, StreamingChoices): - self.response_uptil_now += choice.delta.get("content", "") or "" - else: - self.response_uptil_now += "" + self._accumulate_streamed_delta(choice.delta) self.rules.post_call_rules(input=self.response_uptil_now, model=self.model) # Add mcp_list_tools to first chunk if present if not self.sent_first_chunk and processed_chunk.choices: @@ -2216,9 +2213,7 @@ class CustomStreamWrapper: choice = processed_chunk.choices[0] if isinstance(choice, StreamingChoices): - self.response_uptil_now += choice.delta.get("content", "") or "" - else: - self.response_uptil_now += "" + self._accumulate_streamed_delta(choice.delta) self.rules.post_call_rules(input=self.response_uptil_now, model=self.model) # RETURN RESULT self.chunks.append(processed_chunk) @@ -2395,6 +2390,41 @@ class CustomStreamWrapper: recover_error, ) + _CONTINUATION_DISQUALIFYING_DELTA_FIELDS: Final = ( + "tool_calls", + "function_call", + "thinking_blocks", + "reasoning_items", + "audio", + "images", + "annotations", + ) + + @classmethod + def _delta_disqualifies_continuation(cls, delta: object) -> bool: + """ + True when a streamed delta carries output a text-only prefill + continuation cannot represent: tool/function calls, signed Anthropic + thinking blocks, structured reasoning items, audio or image parts, or + annotations. Plain ``reasoning_content`` is deliberately not here - it + is out-of-band, never reaches the caller as answer text, and so does + not block a continuation (parity with the Responses-API path). + """ + get: Final = getattr(delta, "get", None) + if not callable(get): + return False + return any(get(field) for field in cls._CONTINUATION_DISQUALIFYING_DELTA_FIELDS) + + def _accumulate_streamed_delta(self, delta: object) -> None: + """Grow the running answer text and latch whether anything a + continuation cannot carry has been streamed. One home for both so the + three iteration sites (sync, async, non-aiohttp) stay in step.""" + get: Final = getattr(delta, "get", None) + content: Final = get("content", "") if callable(get) else "" + self.response_uptil_now += content or "" + if not self._emitted_disqualifying_content and self._delta_disqualifies_continuation(delta): + self._emitted_disqualifying_content = True + def _handle_stream_fallback_error(self, e: Exception) -> "NoReturn": """ Common error handling for both __next__ and __anext__. @@ -2466,6 +2496,7 @@ class CustomStreamWrapper: original_exception=mapped_exception, generated_content=self.response_uptil_now, is_pre_first_chunk=not self.sent_first_chunk, + emitted_disqualifying_content=self._emitted_disqualifying_content, ) @staticmethod diff --git a/litellm/router.py b/litellm/router.py index 8865543badd..a10acc63fe8 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -189,6 +189,10 @@ from litellm.router_utils.handle_error import ( send_llm_exception_alert, ) from litellm.router_utils.health_state_cache import DeploymentHealthCache +from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( + MID_STREAM_CONTINUATION_KWARG, + ContinuationPrefillDeploymentCheck, +) from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, warn_on_unknown_model_group_affinity_flags, @@ -764,6 +768,7 @@ class Router: health_check_ignore_transient_errors: bool = False, background_health_check_model_groups: Sequence[str] | None = None, enable_weighted_failover: bool = False, + enable_mid_stream_fallback_continuation: bool = False, fallback_access_check: FallbackAccessCheck | None = None, auto_router_capability_limit: AutoRouterCapabilityLimit | None = None, ) -> None: @@ -802,6 +807,7 @@ class Router: deployment_affinity_ttl_seconds (int): TTL for user-key -> deployment affinity mapping. Defaults to 3600. ignore_invalid_deployments (bool): Ignores invalid deployments, and continues with other deployments. Default is to raise an error. enable_weighted_failover (bool): When True and the routing strategy is "simple-shuffle", a retryable failure on one deployment causes the request to re-pick (weighted) across the other deployments in the same model group before any cross-group fallback runs. Bounded by `max_fallbacks`. Async-only: currently honored by `router.acompletion()` and other async entrypoints. The sync `router.completion()` path falls back to the regular fallback flow. Defaults to False. + enable_mid_stream_fallback_continuation (bool): When True, a chat-completions stream that breaks after plain assistant text has been delivered continues on a fallback deployment via assistant prefill instead of surfacing the error. Only deployments whose model supports assistant prefill are eligible, so the partial text is continued, not regenerated; if none is, the original error is surfaced. Streams that emitted tool calls, thinking blocks, audio/images, or constrained (JSON / forced tool_choice) output are never continued. Async-only. Defaults to False. fallback_access_check (Optional[FallbackAccessCheck]): Awaited before each cross-model-group fallback attempt on the async path; a fallback target it rejects is skipped. Defaults to None (every configured fallback is attempted). Returns: Router: An instance of the litellm.Router class. @@ -984,6 +990,7 @@ class Router: self.disable_cooldowns = disable_cooldowns self.enable_health_check_routing = enable_health_check_routing self.enable_weighted_failover = enable_weighted_failover + self.enable_mid_stream_fallback_continuation = enable_mid_stream_fallback_continuation self.health_check_ignore_transient_errors = health_check_ignore_transient_errors self.background_health_check_model_groups: frozenset[str] | None = ( frozenset(background_health_check_model_groups) @@ -1178,6 +1185,14 @@ class Router: default_pre_call_checks: Final[OptionalPreCallChecks] = [] self.add_optional_pre_call_checks(default_pre_call_checks) + if self.enable_mid_stream_fallback_continuation: + if self.optional_callbacks is None: + self.optional_callbacks = [] + if not any(isinstance(cb, ContinuationPrefillDeploymentCheck) for cb in self.optional_callbacks): + continuation_check: Final = ContinuationPrefillDeploymentCheck() + self.optional_callbacks.append(continuation_check) + litellm.logging_callback_manager.add_litellm_callback(continuation_check) + def discard(self): """ Pseudo-destructor to be invoked to clean up global data structures when router is no longer used. @@ -2802,9 +2817,17 @@ class Router: with anyio.CancelScope(shield=True): await close_model_response() await held_slot.aclose() - if not e.is_pre_first_chunk and ( - e.generated_content or _stream_chunks_have_generated_content(model_response.chunks) - ): + committed: Final = bool( + not e.is_pre_first_chunk + and (e.generated_content or _stream_chunks_have_generated_content(model_response.chunks)) + ) + continue_after_content: Final = committed and self._mid_stream_continuation_eligible( + e=e, request_kwargs=initial_kwargs + ) + # Content already reached the caller and we cannot safely + # continue it (feature off, or tool/thinking/constrained output): + # surface the real error rather than restart into the same stream. + if committed and not continue_after_content: if e.original_exception is not None: raise e.original_exception from e raise @@ -2827,7 +2850,13 @@ class Router: "content_policy_fallbacks", self.content_policy_fallbacks ) initial_kwargs["original_function"] = self._acompletion - initial_kwargs["messages"] = messages + if continue_after_content: + initial_kwargs["messages"] = self._build_completion_continuation_input( + messages, e.generated_content + ) + initial_kwargs[MID_STREAM_CONTINUATION_KWARG] = True + else: + initial_kwargs["messages"] = messages self._update_kwargs_before_fallbacks(model=model_group, kwargs=initial_kwargs) fallback_response = await self.async_function_with_fallbacks_common_utils( e=e, @@ -3017,6 +3046,51 @@ class Router: total_tokens=(partial_usage.total_tokens or 0) + (fb.total_tokens or 0), ) + def _mid_stream_continuation_eligible( + self, + e: "MidStreamFallbackError", + request_kwargs: Mapping[str, object], + ) -> bool: + """ + Whether a chat-completions stream that broke after content may be + continued on a fallback deployment via assistant prefill, instead of + re-raising. Only plain assistant text is safe: a continuation built from + ``generated_content`` (text-only) cannot carry tool calls, signed + thinking blocks, audio or images, and a constrained (JSON / forced + tool_choice) or merged-reasoning response cannot be resumed from an + arbitrary cut point. The fallback target's prefill support is enforced + separately at deployment selection. + """ + if not self.enable_mid_stream_fallback_continuation: + return False + if not e.generated_content or e.emitted_disqualifying_content: + return False + # Any structured-output request (response_format, or a forced tool call) + # produces a partial that cannot be resumed from an arbitrary cut point. + if request_kwargs.get("response_format") is not None: + return False + tool_choice: Final = request_kwargs.get("tool_choice") + if tool_choice == "required" or isinstance(tool_choice, Mapping): + return False + if request_kwargs.get("merge_reasoning_content_in_choices") is True: + return False + return True + + @staticmethod + def _build_completion_continuation_input( + messages: list[dict[str, str]], + generated_content: str, + ) -> Sequence[Mapping[str, object]]: + """ + Append the partial assistant output as a prefill so a prefill-capable + fallback continues where the broken stream stopped instead of + regenerating text already delivered to the caller. The deployment filter + guarantees the target supports ``prefix: True`` (parity with + ``_build_responses_continuation_input`` for the Responses-API path). + """ + prefill: dict[str, object] = {"role": "assistant", "content": generated_content, "prefix": True} + return [*messages, prefill] + @staticmethod def _build_responses_continuation_input( input_val: Union[str, "ResponseInputParam"] | None, @@ -3566,6 +3640,7 @@ class Router: } input_kwargs.pop("silent_model", None) input_kwargs.pop("include_fallback_errors", None) + input_kwargs.pop(MID_STREAM_CONTINUATION_KWARG, None) _response: Final = litellm.acompletion(**input_kwargs) @@ -11959,6 +12034,7 @@ class Router: "retry_policy", "model_group_alias", "enable_weighted_failover", + "enable_mid_stream_fallback_continuation", "enable_tag_filtering", "tag_routing_prefix", ] diff --git a/litellm/router_utils/pre_call_checks/continuation_prefill_check.py b/litellm/router_utils/pre_call_checks/continuation_prefill_check.py new file mode 100644 index 00000000000..75d1ecb225c --- /dev/null +++ b/litellm/router_utils/pre_call_checks/continuation_prefill_check.py @@ -0,0 +1,51 @@ +""" +Mid-stream fallback continuation: keep the fallback on a deployment that can +actually continue a prefilled assistant message. + +When a chat-completions stream breaks after content and the Router re-enters the +fallback chain to continue it (the request is marked with +``MID_STREAM_CONTINUATION_KWARG``), only a deployment whose model supports +assistant prefill can pick up the partial text without regenerating it. +Deployments that cannot are dropped, so selection lands on a +continuation-capable one. If a group has none it empties and the fallback chain +moves on, surfacing the original error rather than sending a request the target +would reject or duplicate. A request without the marker is passed through +untouched. +""" + +from typing import Final + +from pydantic import TypeAdapter, ValidationError + +from litellm.integrations.custom_logger import CustomLogger, Span +from litellm.types.llms.openai import AllMessageValues +from litellm.utils import supports_assistant_prefill + +MID_STREAM_CONTINUATION_KWARG: Final = "_mid_stream_continuation" + +_STR_KEYED_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) + + +def _deployment_supports_prefill(deployment: object) -> bool: + try: + params: Final = _STR_KEYED_DICT_ADAPTER.validate_python( + _STR_KEYED_DICT_ADAPTER.validate_python(deployment).get("litellm_params") + ) + except ValidationError: + return False + model: Final = params.get("model") + return isinstance(model, str) and bool(model) and supports_assistant_prefill(model=model) + + +class ContinuationPrefillDeploymentCheck(CustomLogger): + async def async_filter_deployments( + self, + model: str, + healthy_deployments: list[dict[str, object]], + messages: list[AllMessageValues] | None, + request_kwargs: dict[str, object] | None = None, + parent_otel_span: Span | None = None, + ) -> list[dict[str, object]]: + if not (request_kwargs or {}).get(MID_STREAM_CONTINUATION_KWARG): + return healthy_deployments + return [deployment for deployment in healthy_deployments if _deployment_supports_prefill(deployment)] diff --git a/litellm/utils.py b/litellm/utils.py index 7732cd88cb5..72c4c673547 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2796,6 +2796,20 @@ def supports_prompt_cache_breakpoint(model: str, custom_llm_provider: str | None ) +def supports_assistant_prefill(model: str, custom_llm_provider: str | None = None) -> bool: + """ + Whether the model can continue a prefilled assistant message (Anthropic's + ``prefix: True`` trick and equivalents). Missing metadata reads as False, + so a mid-stream fallback continuation is only routed to a model known to + support it. + """ + return _supports_factory( + model=model, + custom_llm_provider=custom_llm_provider, + key="supports_assistant_prefill", + ) + + def supports_computer_use(model: str, custom_llm_provider: str | None = None) -> bool: """ Check if the given model supports computer use and return a boolean value. diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 37e2031fdf4..f233d6e9bb2 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -4917,3 +4917,47 @@ class TestStableStreamingResponseId: ) wrapper.response_id = "chatcmpl-from-provider" assert wrapper.model_response_creator().id == "chatcmpl-from-provider" + + +class TestContinuationDisqualifiers: + """The mid-stream continuation eligibility hinges on classifying which deltas + carry output a text-only prefill cannot represent.""" + + @pytest.mark.parametrize( + "field", + ["tool_calls", "function_call", "thinking_blocks", "reasoning_items", "audio", "images", "annotations"], + ) + def test_disqualifying_fields_flagged(self, field): + assert CustomStreamWrapper._delta_disqualifies_continuation({field: [{"x": 1}]}) is True + + @pytest.mark.parametrize("delta", [{"content": "hi"}, {"reasoning_content": "thinking"}, {}, {"role": "assistant"}]) + def test_plain_text_and_reasoning_content_not_flagged(self, delta): + # plain reasoning_content is out-of-band and must NOT block a continuation + assert CustomStreamWrapper._delta_disqualifies_continuation(delta) is False + + def test_non_mapping_delta_is_safe(self): + assert CustomStreamWrapper._delta_disqualifies_continuation(object()) is False + + def test_accumulate_grows_text_and_leaves_flag_clear_for_plain_text(self): + wrapper = object.__new__(CustomStreamWrapper) + wrapper.response_uptil_now = "" + wrapper._emitted_disqualifying_content = False + + wrapper._accumulate_streamed_delta({"content": "Hel"}) + wrapper._accumulate_streamed_delta({"content": "lo"}) + + assert wrapper.response_uptil_now == "Hello" + assert wrapper._emitted_disqualifying_content is False + + def test_accumulate_latches_flag_on_disqualifying_delta(self): + wrapper = object.__new__(CustomStreamWrapper) + wrapper.response_uptil_now = "" + wrapper._emitted_disqualifying_content = False + + wrapper._accumulate_streamed_delta({"content": "Hi"}) + wrapper._accumulate_streamed_delta({"tool_calls": [{"index": 0}]}) + # a later plain-text delta must not clear the latch + wrapper._accumulate_streamed_delta({"content": "there"}) + + assert wrapper.response_uptil_now == "Hithere" + assert wrapper._emitted_disqualifying_content is True diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_continuation_prefill_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_continuation_prefill_check.py new file mode 100644 index 00000000000..28f220100fa --- /dev/null +++ b/tests/test_litellm/router_utils/pre_call_checks/test_continuation_prefill_check.py @@ -0,0 +1,69 @@ +import pytest + +from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( + MID_STREAM_CONTINUATION_KWARG, + ContinuationPrefillDeploymentCheck, + _deployment_supports_prefill, +) + +PREFILL_MODEL = "anthropic/claude-3-opus-20240229" # supports_assistant_prefill: True in the cost map +NON_PREFILL_MODEL = "openai/gpt-4o" # capability absent -> treated as unsupported + + +def _deployment(model: str, dep_id: str) -> dict: + return {"litellm_params": {"model": model}, "model_info": {"id": dep_id}} + + +def test_deployment_supports_prefill_reads_capability(): + assert _deployment_supports_prefill(_deployment(PREFILL_MODEL, "a")) is True + assert _deployment_supports_prefill(_deployment(NON_PREFILL_MODEL, "b")) is False + + +def test_deployment_supports_prefill_rejects_malformed_deployments(): + assert _deployment_supports_prefill({}) is False + assert _deployment_supports_prefill({"litellm_params": {}}) is False + assert _deployment_supports_prefill("not-a-dict") is False + + +@pytest.mark.asyncio +async def test_filter_is_noop_without_continuation_marker(): + """A normal (non-continuation) request must be passed through untouched, even + if some deployments cannot prefill.""" + check = ContinuationPrefillDeploymentCheck() + deployments = [_deployment(PREFILL_MODEL, "a"), _deployment(NON_PREFILL_MODEL, "b")] + + for request_kwargs in ({}, None, {MID_STREAM_CONTINUATION_KWARG: False}): + result = await check.async_filter_deployments( + model="group", healthy_deployments=deployments, messages=None, request_kwargs=request_kwargs + ) + assert result == deployments + + +@pytest.mark.asyncio +async def test_filter_keeps_only_prefill_capable_on_continuation(): + check = ContinuationPrefillDeploymentCheck() + deployments = [_deployment(PREFILL_MODEL, "a"), _deployment(NON_PREFILL_MODEL, "b")] + + result = await check.async_filter_deployments( + model="group", + healthy_deployments=deployments, + messages=None, + request_kwargs={MID_STREAM_CONTINUATION_KWARG: True}, + ) + assert [d["model_info"]["id"] for d in result] == ["a"] + + +@pytest.mark.asyncio +async def test_filter_empties_group_when_no_prefill_capable_deployment(): + """No prefill-capable deployment -> empty result, so the router advances the + fallback chain and ultimately surfaces the original error.""" + check = ContinuationPrefillDeploymentCheck() + deployments = [_deployment(NON_PREFILL_MODEL, "b"), _deployment("openai/gpt-4.1", "c")] + + result = await check.async_filter_deployments( + model="group", + healthy_deployments=deployments, + messages=None, + request_kwargs={MID_STREAM_CONTINUATION_KWARG: True}, + ) + assert result == [] diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index b8a0d70f5bc..1a1ca08b9b2 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -2176,6 +2176,182 @@ async def test_acompletion_streaming_iterator(): print("\n=== All tests passed! ===") +def _make_midstream_source(error, chunks=None): + """A minimal stand-in for the CustomStreamWrapper: yields one content chunk, + then raises ``error`` on the next pull. Carries the attributes + _acompletion_streaming_iterator reads.""" + from unittest.mock import MagicMock + + first_chunk = MagicMock(choices=[MagicMock(delta=MagicMock(content="Hello"))]) + + class _Source: + def __init__(self): + self.index = 0 + self.chunks = chunks if chunks is not None else [] + self.model = "gpt-4" + self.custom_llm_provider = "openai" + self.logging_obj = MagicMock() + + def __aiter__(self): + return self + + async def __anext__(self): + if self.index == 0: + self.index += 1 + return first_chunk + raise error + + return _Source() + + +class _FakeFallbackStream: + def __init__(self, item): + self._item = item + self._done = False + + def __aiter__(self): + return self + + async def __anext__(self): + if self._done: + raise StopAsyncIteration + self._done = True + return self._item + + +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator_continues_after_content_when_eligible(): + """Flag on + plain-text break: the router re-enters the fallback chain with an + assistant-prefill continuation and the mid-stream marker, then streams the + fallback's output instead of re-raising.""" + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.exceptions import MidStreamFallbackError + from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( + MID_STREAM_CONTINUATION_KWARG, + ) + + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, + {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, + ], + fallbacks=[{"gpt-4": ["backup"]}], + enable_mid_stream_fallback_continuation=True, + ) + + error = MidStreamFallbackError( + message="Connection lost", + model="gpt-4", + llm_provider="openai", + generated_content="Hello", + is_pre_first_chunk=False, + emitted_disqualifying_content=False, + ) + source = _make_midstream_source(error) + fallback_chunk = litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "world"}}]) + initial_kwargs = {"model": "gpt-4", "stream": True} + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=_FakeFallbackStream(fallback_chunk)), + ) as mock_fallback: + result = await router._acompletion_streaming_iterator( + model_response=source, messages=[{"role": "user", "content": "Hi"}], initial_kwargs=initial_kwargs + ) + collected = [chunk async for chunk in result] + + mock_fallback.assert_awaited_once() + passed_kwargs = mock_fallback.await_args.kwargs["kwargs"] + assert passed_kwargs[MID_STREAM_CONTINUATION_KWARG] is True + assert passed_kwargs["messages"][-1] == {"role": "assistant", "content": "Hello", "prefix": True} + assert fallback_chunk in collected + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "error_kwargs,request_kwargs", + [ + ({"emitted_disqualifying_content": True}, {}), + ({"emitted_disqualifying_content": False}, {"response_format": {"type": "json_object"}}), + ({"emitted_disqualifying_content": False}, {"tool_choice": "required"}), + ], + ids=["tool_or_thinking_emitted", "json_mode", "forced_tool_choice"], +) +async def test_acompletion_streaming_iterator_declines_ineligible_after_content(error_kwargs, request_kwargs): + """Flag on but the break is not continuation-safe: the router re-raises and + never enters the fallback chain, so no duplicated/rejected request is sent.""" + from unittest.mock import AsyncMock, patch + + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, + {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, + ], + fallbacks=[{"gpt-4": ["backup"]}], + enable_mid_stream_fallback_continuation=True, + ) + error = MidStreamFallbackError( + message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", + is_pre_first_chunk=False, **error_kwargs, + ) + source = _make_midstream_source(error) + initial_kwargs = {"model": "gpt-4", "stream": True, **request_kwargs} + + with patch.object(router, "async_function_with_fallbacks_common_utils", new=AsyncMock()) as mock_fallback: + result = await router._acompletion_streaming_iterator( + model_response=source, messages=[{"role": "user", "content": "Hi"}], initial_kwargs=initial_kwargs + ) + with pytest.raises(MidStreamFallbackError): + async for _ in result: + pass + + mock_fallback.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator_flag_off_declines_after_content(): + """Regression guard for the opt-in default: with the flag off, a post-content + break re-raises and never falls back, exactly as before this feature.""" + from unittest.mock import AsyncMock, patch + + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, + {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, + ], + fallbacks=[{"gpt-4": ["backup"]}], + ) + error = MidStreamFallbackError( + message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", + is_pre_first_chunk=False, emitted_disqualifying_content=False, + ) + source = _make_midstream_source(error) + + with patch.object(router, "async_function_with_fallbacks_common_utils", new=AsyncMock()) as mock_fallback: + result = await router._acompletion_streaming_iterator( + model_response=source, messages=[{"role": "user", "content": "Hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + with pytest.raises(MidStreamFallbackError): + async for _ in result: + pass + + mock_fallback.assert_not_awaited() + + +def test_build_completion_continuation_input_appends_assistant_prefill(): + messages = [{"role": "user", "content": "hi"}] + built = litellm.Router._build_completion_continuation_input(messages, "partial answer") + assert built[:-1] == messages + assert built[-1] == {"role": "assistant", "content": "partial answer", "prefix": True} + + @pytest.mark.asyncio async def test_acompletion_streaming_iterator_reraises_original_exception_when_available(): """Async: when the mid-stream MidStreamFallbackError wraps a real provider From c75e3a69e4782ae2c3614e0082d5892ae75347ea Mon Sep 17 00:00:00 2001 From: Ayush Date: Mon, 14 Sep 2026 18:12:27 +0530 Subject: [PATCH 02/18] fix(router): harden mid-stream continuation eligibility and prefill folding Address review of the mid-stream continuation: - keep response_format={"type": "text"} eligible; only json_object / json_schema and other structured formats decline, since "text" is the unconstrained default - fold a new partial into an existing trailing assistant prefill instead of appending a second one, so a nested mid-stream break stays a single prefill turn even on providers that do not merge consecutive assistant messages - split the deployment prefill-capability lookup into two readable steps - cover the merge_reasoning decline branch, the text response_format pass, and the prefill-folding path with tests --- litellm/router.py | 13 ++++++++- .../continuation_prefill_check.py | 7 ++--- tests/test_litellm/test_router.py | 29 ++++++++++++++++++- 3 files changed, 43 insertions(+), 6 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index a10acc63fe8..3de4a70d5c4 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -3067,7 +3067,9 @@ class Router: return False # Any structured-output request (response_format, or a forced tool call) # produces a partial that cannot be resumed from an arbitrary cut point. - if request_kwargs.get("response_format") is not None: + # `{"type": "text"}` is the unconstrained default and stays eligible. + response_format: Final = request_kwargs.get("response_format") + if response_format is not None and response_format != {"type": "text"}: return False tool_choice: Final = request_kwargs.get("tool_choice") if tool_choice == "required" or isinstance(tool_choice, Mapping): @@ -3087,7 +3089,16 @@ class Router: regenerating text already delivered to the caller. The deployment filter guarantees the target supports ``prefix: True`` (parity with ``_build_responses_continuation_input`` for the Responses-API path). + + A nested mid-stream break can re-enter here with a prefill already + appended; the new partial is folded into that trailing assistant turn so + the request keeps a single prefill rather than two consecutive assistant + messages a non-merging provider would reject. """ + if messages and messages[-1].get("role") == "assistant" and messages[-1].get("prefix"): + last: Final = messages[-1] + merged: dict[str, object] = {**last, "content": str(last.get("content") or "") + generated_content} + return [*messages[:-1], merged] prefill: dict[str, object] = {"role": "assistant", "content": generated_content, "prefix": True} return [*messages, prefill] 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 75d1ecb225c..14e637280d8 100644 --- a/litellm/router_utils/pre_call_checks/continuation_prefill_check.py +++ b/litellm/router_utils/pre_call_checks/continuation_prefill_check.py @@ -28,12 +28,11 @@ _STR_KEYED_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) def _deployment_supports_prefill(deployment: object) -> bool: try: - params: Final = _STR_KEYED_DICT_ADAPTER.validate_python( - _STR_KEYED_DICT_ADAPTER.validate_python(deployment).get("litellm_params") - ) + deployment_map: Final = _STR_KEYED_DICT_ADAPTER.validate_python(deployment) + litellm_params: Final = _STR_KEYED_DICT_ADAPTER.validate_python(deployment_map.get("litellm_params")) except ValidationError: return False - model: Final = params.get("model") + model: Final = litellm_params.get("model") return isinstance(model, str) and bool(model) and supports_assistant_prefill(model=model) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 1a1ca08b9b2..0c3bf43182d 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -2276,8 +2276,9 @@ async def test_acompletion_streaming_iterator_continues_after_content_when_eligi ({"emitted_disqualifying_content": True}, {}), ({"emitted_disqualifying_content": False}, {"response_format": {"type": "json_object"}}), ({"emitted_disqualifying_content": False}, {"tool_choice": "required"}), + ({"emitted_disqualifying_content": False}, {"merge_reasoning_content_in_choices": True}), ], - ids=["tool_or_thinking_emitted", "json_mode", "forced_tool_choice"], + ids=["tool_or_thinking_emitted", "json_mode", "forced_tool_choice", "merged_reasoning"], ) async def test_acompletion_streaming_iterator_declines_ineligible_after_content(error_kwargs, request_kwargs): """Flag on but the break is not continuation-safe: the router re-raises and @@ -2352,6 +2353,32 @@ def test_build_completion_continuation_input_appends_assistant_prefill(): assert built[-1] == {"role": "assistant", "content": "partial answer", "prefix": True} +def test_build_completion_continuation_input_folds_into_existing_prefill(): + """A nested break must not leave two trailing assistant turns: the new partial + folds into the prior prefill so a non-merging provider still gets one.""" + once = litellm.Router._build_completion_continuation_input([{"role": "user", "content": "hi"}], "part one ") + twice = litellm.Router._build_completion_continuation_input(list(once), "part two") + assert [m["role"] for m in twice] == ["user", "assistant"] + assert twice[-1] == {"role": "assistant", "content": "part one part two", "prefix": True} + + +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.""" + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[{"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k"}}], + enable_mid_stream_fallback_continuation=True, + ) + e = MidStreamFallbackError( + message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", + is_pre_first_chunk=False, emitted_disqualifying_content=False, + ) + assert router._mid_stream_continuation_eligible(e=e, request_kwargs={"response_format": {"type": "text"}}) is True + assert router._mid_stream_continuation_eligible(e=e, request_kwargs={"response_format": {"type": "json_object"}}) is False + + @pytest.mark.asyncio async def test_acompletion_streaming_iterator_reraises_original_exception_when_available(): """Async: when the mid-stream MidStreamFallbackError wraps a real provider From ba3374732e3e00b2e78afb86a8b8886788a175eb Mon Sep 17 00:00:00 2001 From: Ayush Date: Tue, 15 Sep 2026 02:59:51 +0530 Subject: [PATCH 03/18] fix(router): address review of mid-stream continuation - register the continuation deployment filter as a process-global singleton instead of tracking it per router, so discarding one router no longer removes the type-deduplicated filter that other live routers still depend on - disqualify reasoning_content: it reaches the caller as visible reasoning that a text-only prefill cannot carry, so a reasoning-capable fallback would re-derive it and produce an incoherent response - move MID_STREAM_CONTINUATION_KWARG to constants and lazy-import the filter class in the router, breaking the module-level import cycle CodeQL flagged - trim the added docstrings to the repository comment policy --- litellm/constants.py | 4 ++ .../litellm_core_utils/streaming_handler.py | 16 ++---- litellm/router.py | 50 ++++++------------- .../continuation_prefill_check.py | 18 ++----- .../test_streaming_handler.py | 18 +++++-- 5 files changed, 44 insertions(+), 62 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 5751e6e46af..487c1a1e020 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -2039,3 +2039,7 @@ BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY: Final = "batch_enqueued_token_limit" # Shared read-only empty mapping, for defaulting optional Mapping parameters without # constructing a fresh mutable dict at each call site. EMPTY_MAPPING: Final = MappingProxyType({}) + +# Marks a fallback re-entry as a mid-stream continuation, read by the deployment +# pre-call filter. Lives here so router and the filter share it without an import cycle. +MID_STREAM_CONTINUATION_KWARG: Final = "_mid_stream_continuation" diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index f62a23d2c42..b6301929ce3 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2390,9 +2390,12 @@ class CustomStreamWrapper: recover_error, ) + # Delta fields a text-only prefill continuation cannot carry, so a stream + # that emitted any of them is not eligible for mid-stream continuation. _CONTINUATION_DISQUALIFYING_DELTA_FIELDS: Final = ( "tool_calls", "function_call", + "reasoning_content", "thinking_blocks", "reasoning_items", "audio", @@ -2402,23 +2405,14 @@ class CustomStreamWrapper: @classmethod def _delta_disqualifies_continuation(cls, delta: object) -> bool: - """ - True when a streamed delta carries output a text-only prefill - continuation cannot represent: tool/function calls, signed Anthropic - thinking blocks, structured reasoning items, audio or image parts, or - annotations. Plain ``reasoning_content`` is deliberately not here - it - is out-of-band, never reaches the caller as answer text, and so does - not block a continuation (parity with the Responses-API path). - """ get: Final = getattr(delta, "get", None) if not callable(get): return False return any(get(field) for field in cls._CONTINUATION_DISQUALIFYING_DELTA_FIELDS) def _accumulate_streamed_delta(self, delta: object) -> None: - """Grow the running answer text and latch whether anything a - continuation cannot carry has been streamed. One home for both so the - three iteration sites (sync, async, non-aiohttp) stay in step.""" + # Shared by the sync, async, and non-aiohttp iteration sites so answer + # text and the disqualifying-content latch stay in step across all three. get: Final = getattr(delta, "get", None) content: Final = get("content", "") if callable(get) else "" self.response_uptil_now += content or "" diff --git a/litellm/router.py b/litellm/router.py index 3de4a70d5c4..9f770b23608 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -62,6 +62,7 @@ from litellm.constants import ( DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER, DEFAULT_MAX_LRU_CACHE_SIZE, INTERNAL_CALL_ORIGIN_METADATA_KEY, + MID_STREAM_CONTINUATION_KWARG, OUTPUT_TOKEN_CEILING_PARAMS, ROUTING_REQUEST_TAGS_METADATA_KEY, RUNTIME_UPDATABLE_ROUTER_SETTINGS, @@ -189,10 +190,6 @@ from litellm.router_utils.handle_error import ( send_llm_exception_alert, ) from litellm.router_utils.health_state_cache import DeploymentHealthCache -from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( - MID_STREAM_CONTINUATION_KWARG, - ContinuationPrefillDeploymentCheck, -) from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, warn_on_unknown_model_group_affinity_flags, @@ -1186,12 +1183,14 @@ class Router: self.add_optional_pre_call_checks(default_pre_call_checks) if self.enable_mid_stream_fallback_continuation: - if self.optional_callbacks is None: - self.optional_callbacks = [] - if not any(isinstance(cb, ContinuationPrefillDeploymentCheck) for cb in self.optional_callbacks): - continuation_check: Final = ContinuationPrefillDeploymentCheck() - self.optional_callbacks.append(continuation_check) - litellm.logging_callback_manager.add_litellm_callback(continuation_check) + from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( + ContinuationPrefillDeploymentCheck, + ) + + # Registered on the process-global callback list, never tracked per + # router, so discarding one router cannot drop the filter another + # still needs. It is inert unless a request carries the marker. + litellm.logging_callback_manager.add_litellm_callback(ContinuationPrefillDeploymentCheck()) def discard(self): """ @@ -3051,22 +3050,14 @@ class Router: e: "MidStreamFallbackError", request_kwargs: Mapping[str, object], ) -> bool: - """ - Whether a chat-completions stream that broke after content may be - continued on a fallback deployment via assistant prefill, instead of - re-raising. Only plain assistant text is safe: a continuation built from - ``generated_content`` (text-only) cannot carry tool calls, signed - thinking blocks, audio or images, and a constrained (JSON / forced - tool_choice) or merged-reasoning response cannot be resumed from an - arbitrary cut point. The fallback target's prefill support is enforced - separately at deployment selection. - """ + """Whether a stream that broke after plain assistant text may be + continued via prefill. The fallback target's prefill support is checked + separately at deployment selection.""" if not self.enable_mid_stream_fallback_continuation: return False if not e.generated_content or e.emitted_disqualifying_content: return False - # Any structured-output request (response_format, or a forced tool call) - # produces a partial that cannot be resumed from an arbitrary cut point. + # Structured output cannot be resumed from an arbitrary cut point; # `{"type": "text"}` is the unconstrained default and stays eligible. response_format: Final = request_kwargs.get("response_format") if response_format is not None and response_format != {"type": "text"}: @@ -3083,18 +3074,9 @@ class Router: messages: list[dict[str, str]], generated_content: str, ) -> Sequence[Mapping[str, object]]: - """ - Append the partial assistant output as a prefill so a prefill-capable - fallback continues where the broken stream stopped instead of - regenerating text already delivered to the caller. The deployment filter - guarantees the target supports ``prefix: True`` (parity with - ``_build_responses_continuation_input`` for the Responses-API path). - - A nested mid-stream break can re-enter here with a prefill already - appended; the new partial is folded into that trailing assistant turn so - the request keeps a single prefill rather than two consecutive assistant - messages a non-merging provider would reject. - """ + """Append the partial output as an assistant prefill for a prefill-capable + fallback to continue. A nested break folds the new partial into an + existing trailing prefill rather than appending a second assistant turn.""" if messages and messages[-1].get("role") == "assistant" and messages[-1].get("prefix"): last: Final = messages[-1] merged: dict[str, object] = {**last, "content": str(last.get("content") or "") + generated_content} 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 14e637280d8..2a527c8b249 100644 --- a/litellm/router_utils/pre_call_checks/continuation_prefill_check.py +++ b/litellm/router_utils/pre_call_checks/continuation_prefill_check.py @@ -1,28 +1,20 @@ """ Mid-stream fallback continuation: keep the fallback on a deployment that can -actually continue a prefilled assistant message. - -When a chat-completions stream breaks after content and the Router re-enters the -fallback chain to continue it (the request is marked with -``MID_STREAM_CONTINUATION_KWARG``), only a deployment whose model supports -assistant prefill can pick up the partial text without regenerating it. -Deployments that cannot are dropped, so selection lands on a -continuation-capable one. If a group has none it empties and the fallback chain -moves on, surfacing the original error rather than sending a request the target -would reject or duplicate. A request without the marker is passed through -untouched. +continue a prefilled assistant message. When a request carries +``MID_STREAM_CONTINUATION_KWARG``, deployments whose model does not support +assistant prefill are dropped, so the partial text is continued rather than +regenerated or rejected. Requests without the marker pass through untouched. """ from typing import Final from pydantic import TypeAdapter, ValidationError +from litellm.constants import MID_STREAM_CONTINUATION_KWARG from litellm.integrations.custom_logger import CustomLogger, Span from litellm.types.llms.openai import AllMessageValues from litellm.utils import supports_assistant_prefill -MID_STREAM_CONTINUATION_KWARG: Final = "_mid_stream_continuation" - _STR_KEYED_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index f233d6e9bb2..7aa244a8433 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -4925,14 +4925,24 @@ class TestContinuationDisqualifiers: @pytest.mark.parametrize( "field", - ["tool_calls", "function_call", "thinking_blocks", "reasoning_items", "audio", "images", "annotations"], + [ + "tool_calls", + "function_call", + "reasoning_content", + "thinking_blocks", + "reasoning_items", + "audio", + "images", + "annotations", + ], ) def test_disqualifying_fields_flagged(self, field): + # reasoning_content included: it reaches the caller as visible reasoning a + # text-only prefill cannot carry, so a reasoning fallback would re-derive it assert CustomStreamWrapper._delta_disqualifies_continuation({field: [{"x": 1}]}) is True - @pytest.mark.parametrize("delta", [{"content": "hi"}, {"reasoning_content": "thinking"}, {}, {"role": "assistant"}]) - def test_plain_text_and_reasoning_content_not_flagged(self, delta): - # plain reasoning_content is out-of-band and must NOT block a continuation + @pytest.mark.parametrize("delta", [{"content": "hi"}, {}, {"role": "assistant"}]) + def test_plain_text_not_flagged(self, delta): assert CustomStreamWrapper._delta_disqualifies_continuation(delta) is False def test_non_mapping_delta_is_safe(self): From d5132231bbaab3e1b1d5b7dea901bc09494dc7b9 Mon Sep 17 00:00:00 2001 From: Ayush Date: Tue, 15 Sep 2026 03:34:46 +0530 Subject: [PATCH 04/18] fix(router): keep mid-stream continuation within the type-discipline budget The LIT* type-discipline gate counts mutable-collection use the ruff pass does not see. Annotate the reads with read-only Sequence/Mapping views, compare response_format against a frozen MappingProxyType sentinel, and mark the genuinely mutable litellm message and deployment-list boundaries with mutable-ok reasons, so LIT001 and LIT002 stay under their limits. --- litellm/router.py | 15 +++++++++------ .../pre_call_checks/continuation_prefill_check.py | 12 +++++++----- .../litellm_core_utils/test_streaming_handler.py | 1 - 3 files changed, 16 insertions(+), 12 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 9f770b23608..64db136a791 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -422,6 +422,9 @@ def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream]) _NO_SESSION_KWARGS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType({}) _SESSION_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +# The only response_format that leaves a stream continuation-eligible: the rest +# ask for structured output that cannot resume from an arbitrary cut point. +_UNCONSTRAINED_RESPONSE_FORMAT: Final[Mapping[str, str]] = MappingProxyType({"type": "text"}) def _as_retry_skipped_deployment_ids(value: object) -> tuple[str, ...]: @@ -3060,7 +3063,7 @@ class Router: # Structured output cannot be resumed from an arbitrary cut point; # `{"type": "text"}` is the unconstrained default and stays eligible. response_format: Final = request_kwargs.get("response_format") - if response_format is not None and response_format != {"type": "text"}: + if response_format is not None and response_format != _UNCONSTRAINED_RESPONSE_FORMAT: return False tool_choice: Final = request_kwargs.get("tool_choice") if tool_choice == "required" or isinstance(tool_choice, Mapping): @@ -3071,17 +3074,17 @@ class Router: @staticmethod def _build_completion_continuation_input( - messages: list[dict[str, str]], + messages: Sequence[Mapping[str, str]], generated_content: str, ) -> Sequence[Mapping[str, object]]: """Append the partial output as an assistant prefill for a prefill-capable fallback to continue. A nested break folds the new partial into an existing trailing prefill rather than appending a second assistant turn.""" - if messages and messages[-1].get("role") == "assistant" and messages[-1].get("prefix"): - last: Final = messages[-1] - merged: dict[str, object] = {**last, "content": str(last.get("content") or "") + generated_content} + last: Final = messages[-1] if messages else None + if last is not None and last.get("role") == "assistant" and last.get("prefix"): + merged: Final = {**last, "content": str(last.get("content") or "") + generated_content} return [*messages[:-1], merged] - prefill: dict[str, object] = {"role": "assistant", "content": generated_content, "prefix": True} + prefill: Final = {"role": "assistant", "content": generated_content, "prefix": True} return [*messages, prefill] @staticmethod 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 2a527c8b249..688470b2b77 100644 --- a/litellm/router_utils/pre_call_checks/continuation_prefill_check.py +++ b/litellm/router_utils/pre_call_checks/continuation_prefill_check.py @@ -6,6 +6,7 @@ assistant prefill are dropped, so the partial text is continued rather than regenerated or rejected. Requests without the marker pass through untouched. """ +from collections.abc import Mapping, Sequence from typing import Final from pydantic import TypeAdapter, ValidationError @@ -32,11 +33,12 @@ class ContinuationPrefillDeploymentCheck(CustomLogger): async def async_filter_deployments( self, model: str, - healthy_deployments: list[dict[str, object]], - messages: list[AllMessageValues] | None, - request_kwargs: dict[str, object] | None = None, + healthy_deployments: list[dict[str, object]], # mutable-ok: CustomLogger deployment-list contract + messages: Sequence[AllMessageValues] | None, + request_kwargs: Mapping[str, object] | None = None, parent_otel_span: Span | None = None, - ) -> list[dict[str, object]]: + ) -> list[dict[str, object]]: # mutable-ok: returns a mutable deployment list if not (request_kwargs or {}).get(MID_STREAM_CONTINUATION_KWARG): return healthy_deployments - return [deployment for deployment in healthy_deployments if _deployment_supports_prefill(deployment)] + eligible: Final = (deployment for deployment in healthy_deployments if _deployment_supports_prefill(deployment)) + return list(eligible) # mutable-ok: downstream deployment selection consumes a mutable list diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 7aa244a8433..f5cc8fb6b03 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -4966,7 +4966,6 @@ class TestContinuationDisqualifiers: wrapper._accumulate_streamed_delta({"content": "Hi"}) wrapper._accumulate_streamed_delta({"tool_calls": [{"index": 0}]}) - # a later plain-text delta must not clear the latch wrapper._accumulate_streamed_delta({"content": "there"}) assert wrapper.response_uptil_now == "Hithere" From 8ae8dbe5570a8a8aa6cccd5a76afb4f06362b55c Mon Sep 17 00:00:00 2001 From: Ayush Date: Tue, 15 Sep 2026 12:03:03 +0530 Subject: [PATCH 05/18] fix(router): clear basedpyright operator error and the import-cycle alert - narrow the accumulated stream delta content to str before appending, so the shared helper no longer trips reportOperatorIssue (str += object) - keep the continuation marker in the pre-call-check module and lazy-import it at the two router call sites instead of routing it through constants, so the router has no module-level dependency that closes an import cycle --- litellm/constants.py | 4 ---- litellm/litellm_core_utils/streaming_handler.py | 4 ++-- litellm/router.py | 9 ++++++++- .../pre_call_checks/continuation_prefill_check.py | 5 ++++- 4 files changed, 14 insertions(+), 8 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 487c1a1e020..5751e6e46af 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -2039,7 +2039,3 @@ BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY: Final = "batch_enqueued_token_limit" # Shared read-only empty mapping, for defaulting optional Mapping parameters without # constructing a fresh mutable dict at each call site. EMPTY_MAPPING: Final = MappingProxyType({}) - -# Marks a fallback re-entry as a mid-stream continuation, read by the deployment -# pre-call filter. Lives here so router and the filter share it without an import cycle. -MID_STREAM_CONTINUATION_KWARG: Final = "_mid_stream_continuation" diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index b6301929ce3..cd560f4c0d1 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2414,8 +2414,8 @@ class CustomStreamWrapper: # Shared by the sync, async, and non-aiohttp iteration sites so answer # text and the disqualifying-content latch stay in step across all three. get: Final = getattr(delta, "get", None) - content: Final = get("content", "") if callable(get) else "" - self.response_uptil_now += content or "" + content: Final = get("content") if callable(get) else None + self.response_uptil_now += content if isinstance(content, str) else "" if not self._emitted_disqualifying_content and self._delta_disqualifies_continuation(delta): self._emitted_disqualifying_content = True diff --git a/litellm/router.py b/litellm/router.py index 64db136a791..226b5da7a83 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -62,7 +62,6 @@ from litellm.constants import ( DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER, DEFAULT_MAX_LRU_CACHE_SIZE, INTERNAL_CALL_ORIGIN_METADATA_KEY, - MID_STREAM_CONTINUATION_KWARG, OUTPUT_TOKEN_CEILING_PARAMS, ROUTING_REQUEST_TAGS_METADATA_KEY, RUNTIME_UPDATABLE_ROUTER_SETTINGS, @@ -2853,6 +2852,10 @@ class Router: ) initial_kwargs["original_function"] = self._acompletion if continue_after_content: + from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( + MID_STREAM_CONTINUATION_KWARG, + ) + initial_kwargs["messages"] = self._build_completion_continuation_input( messages, e.generated_content ) @@ -3634,6 +3637,10 @@ class Router: "client": model_client, **kwargs, } + from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( + MID_STREAM_CONTINUATION_KWARG, + ) + input_kwargs.pop("silent_model", None) input_kwargs.pop("include_fallback_errors", None) input_kwargs.pop(MID_STREAM_CONTINUATION_KWARG, None) 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 688470b2b77..a4e05672879 100644 --- a/litellm/router_utils/pre_call_checks/continuation_prefill_check.py +++ b/litellm/router_utils/pre_call_checks/continuation_prefill_check.py @@ -11,11 +11,14 @@ from typing import Final from pydantic import TypeAdapter, ValidationError -from litellm.constants import MID_STREAM_CONTINUATION_KWARG from litellm.integrations.custom_logger import CustomLogger, Span from litellm.types.llms.openai import AllMessageValues from litellm.utils import supports_assistant_prefill +# Marks a fallback re-entry as a mid-stream continuation. Router sets it (via a +# lazy import) and this filter reads it; kept here to avoid a module-level cycle. +MID_STREAM_CONTINUATION_KWARG: Final = "_mid_stream_continuation" + _STR_KEYED_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) From 8543b3cafd8dfca15c4d2e1882294b03a129bb1a Mon Sep 17 00:00:00 2001 From: Ayush Date: Tue, 15 Sep 2026 12:16:55 +0530 Subject: [PATCH 06/18] fix(router): make the mid-stream continuation marker unforgeable The deployment filter is process-global and read the continuation marker as a truthy value from request_kwargs. Since the proxy can forward arbitrary request-body fields into the router, a client could send `"_mid_stream_continuation": true` on a normal request to filter a mixed model group down to prefill-capable deployments and steer their prompt onto that provider. Set a private sentinel object internally and check it by type instead of truthiness: a JSON request body cannot construct one, so a forged flag is ignored. Adds a regression test that forged client values (true, "true", 1, a dict) leave the deployment list untouched. --- litellm/router.py | 3 ++- .../continuation_prefill_check.py | 25 +++++++++++++------ .../test_continuation_prefill_check.py | 23 +++++++++++++++-- tests/test_litellm/test_router.py | 3 ++- 4 files changed, 43 insertions(+), 11 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 226b5da7a83..6a43a6a05c0 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -2854,12 +2854,13 @@ class Router: if continue_after_content: from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( MID_STREAM_CONTINUATION_KWARG, + MID_STREAM_CONTINUATION_MARKER, ) initial_kwargs["messages"] = self._build_completion_continuation_input( messages, e.generated_content ) - initial_kwargs[MID_STREAM_CONTINUATION_KWARG] = True + initial_kwargs[MID_STREAM_CONTINUATION_KWARG] = MID_STREAM_CONTINUATION_MARKER else: initial_kwargs["messages"] = messages self._update_kwargs_before_fallbacks(model=model_group, kwargs=initial_kwargs) 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 a4e05672879..f9750b4933b 100644 --- a/litellm/router_utils/pre_call_checks/continuation_prefill_check.py +++ b/litellm/router_utils/pre_call_checks/continuation_prefill_check.py @@ -1,9 +1,9 @@ """ Mid-stream fallback continuation: keep the fallback on a deployment that can -continue a prefilled assistant message. When a request carries -``MID_STREAM_CONTINUATION_KWARG``, deployments whose model does not support -assistant prefill are dropped, so the partial text is continued rather than -regenerated or rejected. Requests without the marker pass through untouched. +continue a prefilled assistant message. When the router marks a fallback +re-entry as a continuation, deployments whose model does not support assistant +prefill are dropped, so the partial text is continued rather than regenerated or +rejected. Requests without the marker pass through untouched. """ from collections.abc import Mapping, Sequence @@ -15,10 +15,20 @@ from litellm.integrations.custom_logger import CustomLogger, Span from litellm.types.llms.openai import AllMessageValues from litellm.utils import supports_assistant_prefill -# Marks a fallback re-entry as a mid-stream continuation. Router sets it (via a -# lazy import) and this filter reads it; kept here to avoid a module-level cycle. +# The router marks a continuation re-entry by placing MID_STREAM_CONTINUATION_MARKER +# under this key. Since the proxy can forward arbitrary request-body fields into the +# router, the marker is a private object checked by type rather than a truthy value: +# a JSON request body cannot construct one, so a client cannot forge the flag to steer +# deployment selection toward prefill-capable deployments. MID_STREAM_CONTINUATION_KWARG: Final = "_mid_stream_continuation" + +class _ContinuationMarker: + """Unforgeable sentinel; only the router can produce an instance.""" + + +MID_STREAM_CONTINUATION_MARKER: Final = _ContinuationMarker() + _STR_KEYED_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) @@ -41,7 +51,8 @@ class ContinuationPrefillDeploymentCheck(CustomLogger): request_kwargs: Mapping[str, object] | None = None, parent_otel_span: Span | None = None, ) -> list[dict[str, object]]: # mutable-ok: returns a mutable deployment list - if not (request_kwargs or {}).get(MID_STREAM_CONTINUATION_KWARG): + marker: Final = (request_kwargs or {}).get(MID_STREAM_CONTINUATION_KWARG) + if not isinstance(marker, _ContinuationMarker): return healthy_deployments eligible: Final = (deployment for deployment in healthy_deployments if _deployment_supports_prefill(deployment)) return list(eligible) # mutable-ok: downstream deployment selection consumes a mutable list diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_continuation_prefill_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_continuation_prefill_check.py index 28f220100fa..2e1c41bf6ab 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_continuation_prefill_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_continuation_prefill_check.py @@ -2,6 +2,7 @@ import pytest from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( MID_STREAM_CONTINUATION_KWARG, + MID_STREAM_CONTINUATION_MARKER, ContinuationPrefillDeploymentCheck, _deployment_supports_prefill, ) @@ -39,6 +40,24 @@ async def test_filter_is_noop_without_continuation_marker(): assert result == deployments +@pytest.mark.asyncio +async def test_filter_ignores_forged_client_flag(): + """A client cannot steer routing: a plain truthy value under the marker key + (which the proxy could forward from the request body) is not the internal + sentinel, so the filter leaves the deployment list untouched.""" + check = ContinuationPrefillDeploymentCheck() + deployments = [_deployment(PREFILL_MODEL, "a"), _deployment(NON_PREFILL_MODEL, "b")] + + for forged in (True, "true", 1, {"any": "json"}): + result = await check.async_filter_deployments( + model="group", + healthy_deployments=deployments, + messages=None, + request_kwargs={MID_STREAM_CONTINUATION_KWARG: forged}, + ) + assert result == deployments + + @pytest.mark.asyncio async def test_filter_keeps_only_prefill_capable_on_continuation(): check = ContinuationPrefillDeploymentCheck() @@ -48,7 +67,7 @@ async def test_filter_keeps_only_prefill_capable_on_continuation(): model="group", healthy_deployments=deployments, messages=None, - request_kwargs={MID_STREAM_CONTINUATION_KWARG: True}, + request_kwargs={MID_STREAM_CONTINUATION_KWARG: MID_STREAM_CONTINUATION_MARKER}, ) assert [d["model_info"]["id"] for d in result] == ["a"] @@ -64,6 +83,6 @@ async def test_filter_empties_group_when_no_prefill_capable_deployment(): model="group", healthy_deployments=deployments, messages=None, - request_kwargs={MID_STREAM_CONTINUATION_KWARG: True}, + request_kwargs={MID_STREAM_CONTINUATION_KWARG: MID_STREAM_CONTINUATION_MARKER}, ) assert result == [] diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 0c3bf43182d..c37a23df7d1 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -2229,6 +2229,7 @@ async def test_acompletion_streaming_iterator_continues_after_content_when_eligi from litellm.exceptions import MidStreamFallbackError from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( MID_STREAM_CONTINUATION_KWARG, + MID_STREAM_CONTINUATION_MARKER, ) router = litellm.Router( @@ -2264,7 +2265,7 @@ async def test_acompletion_streaming_iterator_continues_after_content_when_eligi mock_fallback.assert_awaited_once() passed_kwargs = mock_fallback.await_args.kwargs["kwargs"] - assert passed_kwargs[MID_STREAM_CONTINUATION_KWARG] is True + assert passed_kwargs[MID_STREAM_CONTINUATION_KWARG] is MID_STREAM_CONTINUATION_MARKER assert passed_kwargs["messages"][-1] == {"role": "assistant", "content": "Hello", "prefix": True} assert fallback_chunk in collected From 177f65a89de83808326417cb82b7a1a353bf19c3 Mon Sep 17 00:00:00 2001 From: Ayush Date: Wed, 16 Sep 2026 00:05:26 +0530 Subject: [PATCH 07/18] feat(router): honor per-deployment model_info for continuation prefill capability The deployment filter only read supports_assistant_prefill from the cost map, so a model configured with `model_info: {supports_assistant_prefill: true}` on the deployment (the usual way to declare a capability for a model that is not in the cost map, or to override it) was wrongly skipped. Check the deployment's model_info first (True or False both authoritative), then fall back to the cost map. Adds tests for the override in both directions. --- .../continuation_prefill_check.py | 13 +++++++++++++ .../test_continuation_prefill_check.py | 19 +++++++++++++++++++ 2 files changed, 32 insertions(+) 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 f9750b4933b..ccd454bc207 100644 --- a/litellm/router_utils/pre_call_checks/continuation_prefill_check.py +++ b/litellm/router_utils/pre_call_checks/continuation_prefill_check.py @@ -35,6 +35,19 @@ _STR_KEYED_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) def _deployment_supports_prefill(deployment: object) -> bool: try: deployment_map: Final = _STR_KEYED_DICT_ADAPTER.validate_python(deployment) + except ValidationError: + return False + # A per-deployment model_info override wins, so a model that is not in the cost + # map (or is registered generically) can still opt in or out explicitly with + # `model_info: {"supports_assistant_prefill": true|false}`. + try: + model_info: Final = _STR_KEYED_DICT_ADAPTER.validate_python(deployment_map.get("model_info")) + declared: Final = model_info.get("supports_assistant_prefill") + if isinstance(declared, bool): + return declared + except ValidationError: + pass + try: litellm_params: Final = _STR_KEYED_DICT_ADAPTER.validate_python(deployment_map.get("litellm_params")) except ValidationError: return False diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_continuation_prefill_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_continuation_prefill_check.py index 2e1c41bf6ab..1d6b667464e 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_continuation_prefill_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_continuation_prefill_check.py @@ -20,6 +20,25 @@ def test_deployment_supports_prefill_reads_capability(): assert _deployment_supports_prefill(_deployment(NON_PREFILL_MODEL, "b")) is False +def test_deployment_model_info_override_wins_over_cost_map(): + # model_info True opts in a model that is not in the cost map + assert ( + _deployment_supports_prefill( + {"litellm_params": {"model": "vendor/custom-model"}, "model_info": {"supports_assistant_prefill": True}} + ) + is True + ) + # model_info False opts out a model the cost map would otherwise allow + assert ( + _deployment_supports_prefill( + {"litellm_params": {"model": PREFILL_MODEL}, "model_info": {"supports_assistant_prefill": False}} + ) + is False + ) + # model_info without the key falls through to the cost map + assert _deployment_supports_prefill(_deployment(PREFILL_MODEL, "z")) is True + + def test_deployment_supports_prefill_rejects_malformed_deployments(): assert _deployment_supports_prefill({}) is False assert _deployment_supports_prefill({"litellm_params": {}}) is False From 6f3bac769f57a7761786bd871ebb033e94b3ea82 Mon Sep 17 00:00:00 2001 From: Ayush Date: Wed, 16 Sep 2026 00:26:40 +0530 Subject: [PATCH 08/18] 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.""" From ce36fc6851b1b19e14184274a0d99042d72e199a Mon Sep 17 00:00:00 2001 From: Ayush Date: Wed, 16 Sep 2026 00:34:43 +0530 Subject: [PATCH 09/18] refactor(router): extract mid-stream decline to stay within the complexity budget The output-ceiling branch pushed stream_with_fallbacks past the C901 limit. Both decline paths (ineligible after content, and exhausted allowance) re-raised the original provider error the same way, so pull that into _raise_original_mid_stream_error. It removes the duplicated branching, drops the function back under the complexity ceiling, and reads clearer. --- litellm/router.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 474c284c389..9268674121d 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -33,7 +33,7 @@ from collections.abc import ( ) from functools import lru_cache, partial from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypeVar, Union, cast +from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias, TypeVar, Union, cast import anyio import httpx @@ -2829,9 +2829,7 @@ class Router: # continue it (feature off, or tool/thinking/constrained output): # surface the real error rather than restart into the same stream. if committed and not continue_after_content: - if e.original_exception is not None: - raise e.original_exception from e - raise + self._raise_original_mid_stream_error(e) from litellm.main import stream_chunk_builder @@ -2864,9 +2862,7 @@ class Router: 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 + self._raise_original_mid_stream_error(e) initial_kwargs.update(reduced_ceilings) initial_kwargs["messages"] = self._build_completion_continuation_input( messages, e.generated_content @@ -3063,6 +3059,14 @@ class Router: total_tokens=(partial_usage.total_tokens or 0) + (fb.total_tokens or 0), ) + @staticmethod + def _raise_original_mid_stream_error(e: "MidStreamFallbackError") -> NoReturn: + """Decline a continuation by surfacing the real provider error the stream + wrapper carried, rather than leaking the internal MidStreamFallbackError.""" + if e.original_exception is not None: + raise e.original_exception from e + raise e + def _mid_stream_continuation_eligible( self, e: "MidStreamFallbackError", From ba472dfa71235a1baab0dd1093542d16a775f27d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 1 Oct 2026 16:27:52 -0700 Subject: [PATCH 10/18] fix(cost-map): mark Claude Sonnet 4.6 as not supporting assistant prefill The Anthropic API and Bedrock both reject an assistant-last message on claude-sonnet-4-6 with "This model does not support assistant message prefill", so the continuation filter must not route a mid-stream continuation to it. Flip the flag on every claude-sonnet-4-6 entry and on databricks' opus-4-6 entry, whose anthropic sibling already says false. The prefill check tests now register synthetic cost-map entries instead of pinning real models, which broke when claude-3-opus left the map. --- ...odel_prices_and_context_window_backup.json | 28 +++++++++---------- model_prices_and_context_window.json | 28 +++++++++---------- .../test_continuation_prefill_check.py | 20 +++++++++++-- 3 files changed, 45 insertions(+), 31 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 40fdf083cf8..5d5dab6f746 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -2749,7 +2749,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -2784,7 +2784,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -2819,7 +2819,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -2854,7 +2854,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -2889,7 +2889,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -2924,7 +2924,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -3712,7 +3712,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 1.5e-05, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -15056,7 +15056,7 @@ }, "supports_adaptive_thinking": true, "supports_legacy_thinking": true, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -19314,7 +19314,7 @@ "output_cost_per_token": 2.5000010000000002e-05, "output_dbu_cost_per_token": 0.000357143, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_function_calling": true, "supports_legacy_thinking": true, "supports_prompt_caching": true, @@ -19529,7 +19529,7 @@ "output_cost_per_token": 1.5000020000000002e-05, "output_dbu_cost_per_token": 0.000214286, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_function_calling": true, "supports_legacy_thinking": true, "supports_prompt_caching": true, @@ -41833,7 +41833,7 @@ "output_cost_per_token": 1.5e-05, "output_cost_per_token_above_200k_tokens": 2.25e-05, "source": "https://openrouter.ai/api/v1/models", - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_prompt_caching": true, @@ -50030,7 +50030,7 @@ "mode": "chat", "output_cost_per_token": 1.5e-05, "output_cost_per_token_batches": 7.5e-06, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -57738,7 +57738,7 @@ "mode": "chat", "output_cost_per_token": 1.5e-05, "output_cost_per_token_batches": 7.5e-06, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -78427,7 +78427,7 @@ "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 40fdf083cf8..5d5dab6f746 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -2749,7 +2749,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -2784,7 +2784,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -2819,7 +2819,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -2854,7 +2854,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -2889,7 +2889,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -2924,7 +2924,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -3712,7 +3712,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 1.5e-05, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -15056,7 +15056,7 @@ }, "supports_adaptive_thinking": true, "supports_legacy_thinking": true, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -19314,7 +19314,7 @@ "output_cost_per_token": 2.5000010000000002e-05, "output_dbu_cost_per_token": 0.000357143, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_function_calling": true, "supports_legacy_thinking": true, "supports_prompt_caching": true, @@ -19529,7 +19529,7 @@ "output_cost_per_token": 1.5000020000000002e-05, "output_dbu_cost_per_token": 0.000214286, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_function_calling": true, "supports_legacy_thinking": true, "supports_prompt_caching": true, @@ -41833,7 +41833,7 @@ "output_cost_per_token": 1.5e-05, "output_cost_per_token_above_200k_tokens": 2.25e-05, "source": "https://openrouter.ai/api/v1/models", - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_prompt_caching": true, @@ -50030,7 +50030,7 @@ "mode": "chat", "output_cost_per_token": 1.5e-05, "output_cost_per_token_batches": 7.5e-06, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -57738,7 +57738,7 @@ "mode": "chat", "output_cost_per_token": 1.5e-05, "output_cost_per_token_batches": 7.5e-06, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -78427,7 +78427,7 @@ "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, diff --git a/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py b/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py index 1d6b667464e..bf3ff3c6b4f 100644 --- a/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py @@ -1,5 +1,6 @@ import pytest +import litellm from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( MID_STREAM_CONTINUATION_KWARG, MID_STREAM_CONTINUATION_MARKER, @@ -7,8 +8,21 @@ from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( _deployment_supports_prefill, ) -PREFILL_MODEL = "anthropic/claude-3-opus-20240229" # supports_assistant_prefill: True in the cost map -NON_PREFILL_MODEL = "openai/gpt-4o" # capability absent -> treated as unsupported +PREFILL_MODEL = "anthropic/prefill-capable-test-model" +NON_PREFILL_MODEL = "openai/prefill-unknown-test-model" +UNMAPPED_MODEL = "openai/unmapped-test-model" + + +@pytest.fixture(autouse=True) +def _synthetic_cost_map_entries() -> None: + litellm.register_model( + { + PREFILL_MODEL: {"litellm_provider": "anthropic", "mode": "chat", "supports_assistant_prefill": True}, + NON_PREFILL_MODEL: {"litellm_provider": "openai", "mode": "chat"}, + }, + persist_across_reloads=False, + ) + litellm.get_model_info.cache_clear() def _deployment(model: str, dep_id: str) -> dict: @@ -96,7 +110,7 @@ async def test_filter_empties_group_when_no_prefill_capable_deployment(): """No prefill-capable deployment -> empty result, so the router advances the fallback chain and ultimately surfaces the original error.""" check = ContinuationPrefillDeploymentCheck() - deployments = [_deployment(NON_PREFILL_MODEL, "b"), _deployment("openai/gpt-4.1", "c")] + deployments = [_deployment(NON_PREFILL_MODEL, "b"), _deployment(UNMAPPED_MODEL, "c")] result = await check.async_filter_deployments( model="group", From eab6d76e9a751730d475c4473da61e5b4ddc916e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 1 Oct 2026 17:04:07 -0700 Subject: [PATCH 11/18] refactor(router): type the streaming fallback kwargs and drop narration comments Removes an unused mutable-ok marker (LIT013) and the explanatory comments Greptile flagged, and types initial_kwargs on both streaming fallback iterators as dict[str, Any] like their Responses siblings so the new continuation helpers take a typed mapping and the basedpyright budget stays at or below main. --- litellm/router.py | 39 +++++-------------- .../continuation_prefill_check.py | 21 ++-------- 2 files changed, 12 insertions(+), 48 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 0fd97ebccad..d479519c18e 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -473,8 +473,6 @@ def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream]) _NO_SESSION_KWARGS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType({}) _SESSION_ADAPTER: Final = TypeAdapter(Mapping[str, object]) -# The only response_format that leaves a stream continuation-eligible: the rest -# ask for structured output that cannot resume from an arbitrary cut point. _UNCONSTRAINED_RESPONSE_FORMAT: Final[Mapping[str, str]] = MappingProxyType({"type": "text"}) _SILENT_MODEL_ADAPTER: Final = TypeAdapter(str | list[str]) @@ -1266,9 +1264,7 @@ class Router: ContinuationPrefillDeploymentCheck, ) - # Registered on the process-global callback list, never tracked per - # router, so discarding one router cannot drop the filter another - # still needs. It is inert unless a request carries the marker. + # Process-global on purpose: Router.discard() must not drop a filter another router still needs. litellm.logging_callback_manager.add_litellm_callback(ContinuationPrefillDeploymentCheck()) def discard(self): @@ -2559,7 +2555,7 @@ class Router: try: # Capture kwargs before deployment selection so the streaming # fallback iterator can re-dispatch with the original model group. - input_kwargs_for_streaming_fallback: Final = kwargs.copy() + input_kwargs_for_streaming_fallback: Final[dict[str, Any]] = kwargs.copy() input_kwargs_for_streaming_fallback["model"] = model # pick the one that is available (lowest TPM/RPM) @@ -2900,7 +2896,7 @@ class Router: self, model_response: CustomStreamWrapper, messages: list[dict[str, str]], - initial_kwargs: dict, + initial_kwargs: dict[str, Any], deployment_slot: contextlib.AsyncExitStack | None = None, ) -> CustomStreamWrapper: """ @@ -2966,9 +2962,6 @@ class Router: continue_after_content: Final = committed and self._mid_stream_continuation_eligible( e=e, request_kwargs=initial_kwargs ) - # Content already reached the caller and we cannot safely - # continue it (feature off, or tool/thinking/constrained output): - # surface the real error rather than restart into the same stream. if committed and not continue_after_content: self._raise_original_mid_stream_error(e) @@ -3001,8 +2994,6 @@ class Router: ) 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. self._raise_original_mid_stream_error(e) initial_kwargs.update(reduced_ceilings) initial_kwargs["messages"] = self._build_completion_continuation_input( @@ -3202,8 +3193,7 @@ class Router: @staticmethod def _raise_original_mid_stream_error(e: "MidStreamFallbackError") -> NoReturn: - """Decline a continuation by surfacing the real provider error the stream - wrapper carried, rather than leaking the internal MidStreamFallbackError.""" + """Surface the provider error the stream wrapper carried instead of the internal MidStreamFallbackError.""" if e.original_exception is not None: raise e.original_exception from e raise e @@ -3213,15 +3203,11 @@ class Router: e: "MidStreamFallbackError", request_kwargs: Mapping[str, object], ) -> bool: - """Whether a stream that broke after plain assistant text may be - continued via prefill. The fallback target's prefill support is checked - separately at deployment selection.""" + """Whether a stream that broke after plain assistant text may continue via prefill.""" if not self.enable_mid_stream_fallback_continuation: return False if not e.generated_content or e.emitted_disqualifying_content: return False - # Structured output cannot be resumed from an arbitrary cut point; - # `{"type": "text"}` is the unconstrained default and stays eligible. response_format: Final = request_kwargs.get("response_format") if response_format is not None and response_format != _UNCONSTRAINED_RESPONSE_FORMAT: return False @@ -3237,12 +3223,7 @@ class Router: 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.""" + """The caller's output ceilings minus the tokens already emitted, or None once nothing is left.""" ceilings: Final = MappingProxyType( { key: value - emitted_tokens @@ -3259,9 +3240,7 @@ class Router: messages: Sequence[Mapping[str, str]], generated_content: str, ) -> Sequence[Mapping[str, object]]: - """Append the partial output as an assistant prefill for a prefill-capable - fallback to continue. A nested break folds the new partial into an - existing trailing prefill rather than appending a second assistant turn.""" + """Append the partial output as an assistant prefill, extending an existing trailing prefill in place.""" last: Final = messages[-1] if messages else None if last is not None and last.get("role") == "assistant" and last.get("prefix"): merged: Final = {**last, "content": str(last.get("content") or "") + generated_content} @@ -3620,7 +3599,7 @@ class Router: self, model_response: CustomStreamWrapper, messages: list[dict[str, str]], - initial_kwargs: dict, + initial_kwargs: dict[str, Any], ) -> CustomStreamWrapper: """ Sync equivalent of _acompletion_streaming_iterator. @@ -3783,7 +3762,7 @@ class Router: deployment = None _timeout_debug_deployment_dict = {} # this is a temporary dict to debug timeout issues try: - input_kwargs_for_streaming_fallback: Final = kwargs.copy() + input_kwargs_for_streaming_fallback: Final[dict[str, Any]] = kwargs.copy() input_kwargs_for_streaming_fallback["model"] = model parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs) 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 c640374858a..b8d6c68a353 100644 --- a/litellm/router_utils/pre_call_checks/continuation_prefill_check.py +++ b/litellm/router_utils/pre_call_checks/continuation_prefill_check.py @@ -1,10 +1,4 @@ -""" -Mid-stream fallback continuation: keep the fallback on a deployment that can -continue a prefilled assistant message. When the router marks a fallback -re-entry as a continuation, deployments whose model does not support assistant -prefill are dropped, so the partial text is continued rather than regenerated or -rejected. Requests without the marker pass through untouched. -""" +"""Keeps a mid-stream continuation on deployments whose model supports assistant prefill.""" from collections.abc import Mapping, Sequence from typing import Final @@ -15,16 +9,11 @@ from litellm.integrations.custom_logger import CustomLogger, Span from litellm.types.llms.openai import AllMessageValues from litellm.utils import supports_assistant_prefill -# The router marks a continuation re-entry by placing MID_STREAM_CONTINUATION_MARKER -# under this key. Since the proxy can forward arbitrary request-body fields into the -# router, the marker is a private object checked by type rather than a truthy value: -# a JSON request body cannot construct one, so a client cannot forge the flag to steer -# deployment selection toward prefill-capable deployments. MID_STREAM_CONTINUATION_KWARG: Final = "_mid_stream_continuation" class _ContinuationMarker: - """Unforgeable sentinel; only the router can produce an instance.""" + """Only the router constructs one, so a JSON request body cannot forge the continuation flag.""" MID_STREAM_CONTINUATION_MARKER: Final = _ContinuationMarker() @@ -37,16 +26,12 @@ def _deployment_supports_prefill(deployment: object) -> bool: deployment_map: Final = _STR_KEYED_DICT_ADAPTER.validate_python(deployment) except ValidationError: return False - # A per-deployment model_info override wins, so a model that is not in the cost - # map (or is registered generically) can still opt in or out explicitly with - # `model_info: {"supports_assistant_prefill": true|false}`. try: model_info: Final = _STR_KEYED_DICT_ADAPTER.validate_python(deployment_map.get("model_info")) declared: Final = model_info.get("supports_assistant_prefill") 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")) @@ -69,4 +54,4 @@ class ContinuationPrefillDeploymentCheck(CustomLogger): if not isinstance(marker, _ContinuationMarker): return healthy_deployments eligible: Final = (deployment for deployment in healthy_deployments if _deployment_supports_prefill(deployment)) - return list(eligible) # mutable-ok: downstream deployment selection consumes a mutable list + return list(eligible) From 6e52b45dbff214226447a348b95dc7155f2908c2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 1 Oct 2026 17:14:48 -0700 Subject: [PATCH 12/18] fix(router): run post_call_rules over the joined text of a continued stream Each provider stream wrapper validates only the fragment it produced, so a rule that trips on the primary's prefix plus the continuation's text never fired. The router now feeds litellm.post_call_rules the emitted text plus every continuation chunk before yielding it. Also drops a narrating test-helper docstring. --- litellm/router.py | 16 ++++-- .../continuation_prefill_check.py | 26 ++++++++++ tests/unit/test_router/test_router.py | 49 +++++++++++++++++-- 3 files changed, 83 insertions(+), 8 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index d479519c18e..849edfc66a7 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -2983,12 +2983,16 @@ class Router: "content_policy_fallbacks", self.content_policy_fallbacks ) initial_kwargs["original_function"] = self._acompletion - if continue_after_content: - from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( - MID_STREAM_CONTINUATION_KWARG, - MID_STREAM_CONTINUATION_MARKER, - ) + from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( + MID_STREAM_CONTINUATION_KWARG, + MID_STREAM_CONTINUATION_MARKER, + ContinuationOutputRules, + ) + continuation_rules: Final = ( + ContinuationOutputRules(e.generated_content, model_group) if continue_after_content else None + ) + if continue_after_content: emitted_tokens: Final = int( getattr(complete_response_object_usage, "completion_tokens", 0) or 0 ) @@ -3035,6 +3039,8 @@ class Router: and hasattr(fallback_item, "usage") ): self._combine_fallback_usage(fallback_item, complete_response_object_usage) + if continuation_rules is not None and isinstance(fallback_item, ModelResponseStream): + continuation_rules.observe(fallback_item) yield fallback_item else: # If fallback returns a non-streaming response, yield None 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 b8d6c68a353..0d6d462d417 100644 --- a/litellm/router_utils/pre_call_checks/continuation_prefill_check.py +++ b/litellm/router_utils/pre_call_checks/continuation_prefill_check.py @@ -5,8 +5,11 @@ from typing import Final from pydantic import TypeAdapter, ValidationError +import litellm from litellm.integrations.custom_logger import CustomLogger, Span +from litellm.litellm_core_utils.rules import Rules from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import ModelResponseStream from litellm.utils import supports_assistant_prefill MID_STREAM_CONTINUATION_KWARG: Final = "_mid_stream_continuation" @@ -55,3 +58,26 @@ class ContinuationPrefillDeploymentCheck(CustomLogger): return healthy_deployments eligible: Final = (deployment for deployment in healthy_deployments if _deployment_supports_prefill(deployment)) return list(eligible) + + +def _delta_text(chunk: ModelResponseStream) -> str: + if not chunk.choices: + return "" + delta: Final = chunk.choices[0].delta + content: Final = delta.content if delta is not None else None + return content if isinstance(content, str) else "" + + +class ContinuationOutputRules: + """Runs litellm.post_call_rules over the primary's text plus the continuation's, since each stream wrapper only sees its own fragment.""" + + def __init__(self, emitted_text: str, model: str) -> None: + self._text = emitted_text + self._model: Final = model + self._rules: Final = Rules() + + def observe(self, chunk: ModelResponseStream) -> None: + if not litellm.post_call_rules: + return + self._text += _delta_text(chunk) + self._rules.post_call_rules(input=self._text, model=self._model) diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 753161a3bbc..5e1b2f3b8ef 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -2663,9 +2663,6 @@ async def test_acompletion_streaming_iterator(): def _make_midstream_source(error, chunks=None): - """A minimal stand-in for the CustomStreamWrapper: yields one content chunk, - then raises ``error`` on the next pull. Carries the attributes - _acompletion_streaming_iterator reads.""" from unittest.mock import MagicMock first_chunk = MagicMock(choices=[MagicMock(delta=MagicMock(content="Hello"))]) @@ -2756,6 +2753,52 @@ async def test_acompletion_streaming_iterator_continues_after_content_when_eligi assert fallback_chunk in collected +@pytest.mark.asyncio +@pytest.mark.parametrize( + "continuation_text,joined_output_trips_rule", + [("6789", True), ("6788", False)], + ids=["rule_trips_only_on_the_joined_text", "clean_continuation_streams"], +) +async def test_acompletion_streaming_iterator_continuation_runs_post_call_rules_on_the_joined_text( + monkeypatch, continuation_text, joined_output_trips_rule +): + from unittest.mock import AsyncMock, patch + + from litellm.exceptions import MidStreamFallbackError + + monkeypatch.setattr(litellm, "post_call_rules", [lambda output: "123-45-6789" not in output]) + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, + {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, + ], + fallbacks=[{"gpt-4": ["backup"]}], + enable_mid_stream_fallback_continuation=True, + ) + error = MidStreamFallbackError( + message="Connection lost", model="gpt-4", llm_provider="openai", generated_content="123-45-", + is_pre_first_chunk=False, emitted_disqualifying_content=False, + ) + fallback_chunk = litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": continuation_text}}]) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=_FakeFallbackStream(fallback_chunk)), + ): + result = await router._acompletion_streaming_iterator( + model_response=_make_midstream_source(error), + messages=[{"role": "user", "content": "Hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + if joined_output_trips_rule: + with pytest.raises(litellm.APIResponseValidationError): + async for _ in result: + pass + else: + assert fallback_chunk in [chunk async for chunk in result] + + @pytest.mark.asyncio @pytest.mark.parametrize( "error_kwargs,request_kwargs", From 430d931ff841fb39356d68c658095fa0cbf473c1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 1 Oct 2026 17:31:16 -0700 Subject: [PATCH 13/18] refactor(router): read a deployment's declared prefill support without an empty except --- .../continuation_prefill_check.py | 19 ++++++++++++------- .../test_continuation_prefill_check.py | 2 ++ 2 files changed, 14 insertions(+), 7 deletions(-) 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 0d6d462d417..2bde8ca920d 100644 --- a/litellm/router_utils/pre_call_checks/continuation_prefill_check.py +++ b/litellm/router_utils/pre_call_checks/continuation_prefill_check.py @@ -24,18 +24,23 @@ MID_STREAM_CONTINUATION_MARKER: Final = _ContinuationMarker() _STR_KEYED_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) +def _declared_prefill_support(deployment_map: Mapping[str, object]) -> bool | None: + try: + model_info: Final = _STR_KEYED_DICT_ADAPTER.validate_python(deployment_map.get("model_info")) + except ValidationError: + return None + declared: Final = model_info.get("supports_assistant_prefill") + return declared if isinstance(declared, bool) else None + + def _deployment_supports_prefill(deployment: object) -> bool: try: deployment_map: Final = _STR_KEYED_DICT_ADAPTER.validate_python(deployment) except ValidationError: return False - try: - model_info: Final = _STR_KEYED_DICT_ADAPTER.validate_python(deployment_map.get("model_info")) - declared: Final = model_info.get("supports_assistant_prefill") - if isinstance(declared, bool): - return declared - except ValidationError: - pass + declared: Final = _declared_prefill_support(deployment_map) + if declared is not None: + return declared try: litellm_params: Final = _STR_KEYED_DICT_ADAPTER.validate_python(deployment_map.get("litellm_params")) except ValidationError: diff --git a/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py b/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py index bf3ff3c6b4f..4cbe48eda54 100644 --- a/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py @@ -32,6 +32,8 @@ def _deployment(model: str, dep_id: str) -> dict: def test_deployment_supports_prefill_reads_capability(): assert _deployment_supports_prefill(_deployment(PREFILL_MODEL, "a")) is True assert _deployment_supports_prefill(_deployment(NON_PREFILL_MODEL, "b")) is False + assert _deployment_supports_prefill({"litellm_params": {"model": PREFILL_MODEL}}) is True + assert _deployment_supports_prefill({"litellm_params": {"model": PREFILL_MODEL}, "model_info": "bogus"}) is True def test_deployment_model_info_override_wins_over_cost_map(): From e6eefdf3d682ae5e903642ea9344d03791499376 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 1 Oct 2026 17:43:34 -0700 Subject: [PATCH 14/18] test(router): cover _raise_original_mid_stream_error directly for the router coverage gate --- tests/unit/test_router/test_router.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 5e1b2f3b8ef..8676e2079f4 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -2929,6 +2929,25 @@ def test_mid_stream_continuation_eligible_allows_text_response_format(): assert router._mid_stream_continuation_eligible(e=e, request_kwargs={"response_format": {"type": "json_object"}}) is False +def test_raise_original_mid_stream_error_surfaces_the_provider_exception(): + from litellm.exceptions import MidStreamFallbackError, RateLimitError + + provider_error = RateLimitError(message="rate limited", llm_provider="openai", model="gpt-4") + wrapped = MidStreamFallbackError( + message="rate limited", model="gpt-4", llm_provider="openai", + original_exception=provider_error, generated_content="Hello", + ) + with pytest.raises(RateLimitError) as raised: + litellm.Router._raise_original_mid_stream_error(wrapped) + assert raised.value is provider_error + assert raised.value.__cause__ is wrapped + + bare = MidStreamFallbackError(message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello") + with pytest.raises(MidStreamFallbackError) as bare_raised: + litellm.Router._raise_original_mid_stream_error(bare) + assert bare_raised.value is bare + + @pytest.mark.asyncio async def test_acompletion_streaming_iterator_reraises_original_exception_when_available(): """Async: when the mid-stream MidStreamFallbackError wraps a real provider From 197a58b6d222a8d31f974e2287a86dad464713e4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 1 Oct 2026 17:59:56 -0700 Subject: [PATCH 15/18] fix(router): fold the continuation prefill into a trailing assistant turn A request that already ends in an assistant message took the partial output as a second assistant turn, which prefill providers reject. The partial now extends that turn in place, with prefix set, so the continuation request never ends in two assistant messages. --- litellm/router.py | 7 ++++--- tests/unit/test_router/test_router.py | 9 +++++++++ 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 849edfc66a7..2d3d6971e3b 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -3246,10 +3246,11 @@ class Router: messages: Sequence[Mapping[str, str]], generated_content: str, ) -> Sequence[Mapping[str, object]]: - """Append the partial output as an assistant prefill, extending an existing trailing prefill in place.""" + """Append the partial output as an assistant prefill, or extend a trailing assistant turn in place so the + request never ends in two assistant messages.""" last: Final = messages[-1] if messages else None - if last is not None and last.get("role") == "assistant" and last.get("prefix"): - merged: Final = {**last, "content": str(last.get("content") or "") + generated_content} + if last is not None and last.get("role") == "assistant": + merged: Final = {**last, "content": str(last.get("content") or "") + generated_content, "prefix": True} return [*messages[:-1], merged] prefill: Final = {"role": "assistant", "content": generated_content, "prefix": True} return [*messages, prefill] diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 8676e2079f4..d59674f1c57 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -2892,6 +2892,15 @@ def test_build_completion_continuation_input_folds_into_existing_prefill(): assert twice[-1] == {"role": "assistant", "content": "part one part two", "prefix": True} +def test_build_completion_continuation_input_folds_into_trailing_plain_assistant_turn(): + """A request that already ends in an assistant message takes the partial as its prefill + instead of gaining a second assistant turn, which prefill providers reject.""" + messages = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "Sure, "}] + built = litellm.Router._build_completion_continuation_input(messages, "here it is") + assert [m["role"] for m in built] == ["user", "assistant"] + assert built[-1] == {"role": "assistant", "content": "Sure, here it is", "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.""" From d2ab01a88bba980da53fbda20d084845a5172d9a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 1 Oct 2026 18:11:49 -0700 Subject: [PATCH 16/18] fix(router): keep content blocks on a trailing assistant turn when extending it with the partial output A trailing assistant message whose content is a list of blocks was stringified before the partial output was appended, so the fallback received the Python repr of the blocks as its prefill. The partial now lands as one more text block, and a content shape that cannot be extended surfaces the original stream error instead of a garbled continuation. --- litellm/router.py | 32 +++++++++++++++++---------- tests/unit/test_router/test_router.py | 15 +++++++++++++ 2 files changed, 35 insertions(+), 12 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 2d3d6971e3b..bb8485e0ef5 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -475,6 +475,7 @@ _NO_SESSION_KWARGS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType _SESSION_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _UNCONSTRAINED_RESPONSE_FORMAT: Final[Mapping[str, str]] = MappingProxyType({"type": "text"}) _SILENT_MODEL_ADAPTER: Final = TypeAdapter(str | list[str]) +_CONTENT_BLOCKS_ADAPTER: Final = TypeAdapter(list[dict[str, object]]) def _as_retry_skipped_deployment_ids(value: object) -> tuple[str, ...]: @@ -2997,12 +2998,13 @@ class Router: 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: - self._raise_original_mid_stream_error(e) - initial_kwargs.update(reduced_ceilings) - initial_kwargs["messages"] = self._build_completion_continuation_input( + continuation_messages: Final = self._build_completion_continuation_input( messages, e.generated_content ) + if reduced_ceilings is None or continuation_messages is None: + self._raise_original_mid_stream_error(e) + initial_kwargs.update(reduced_ceilings) + initial_kwargs["messages"] = continuation_messages initial_kwargs[MID_STREAM_CONTINUATION_KWARG] = MID_STREAM_CONTINUATION_MARKER else: initial_kwargs["messages"] = messages @@ -3243,17 +3245,23 @@ class Router: @staticmethod def _build_completion_continuation_input( - messages: Sequence[Mapping[str, str]], + messages: Sequence[Mapping[str, object]], generated_content: str, - ) -> Sequence[Mapping[str, object]]: + ) -> Sequence[Mapping[str, object]] | None: """Append the partial output as an assistant prefill, or extend a trailing assistant turn in place so the - request never ends in two assistant messages.""" + request never ends in two assistant messages. None when that turn's content has a shape this cannot extend.""" last: Final = messages[-1] if messages else None - if last is not None and last.get("role") == "assistant": - merged: Final = {**last, "content": str(last.get("content") or "") + generated_content, "prefix": True} - return [*messages[:-1], merged] - prefill: Final = {"role": "assistant", "content": generated_content, "prefix": True} - return [*messages, prefill] + if last is None or last.get("role") != "assistant": + return [*messages, {"role": "assistant", "content": generated_content, "prefix": True}] + content: Final = last.get("content") + if content is None or isinstance(content, str): + return [*messages[:-1], {**last, "content": f"{content or ''}{generated_content}", "prefix": True}] + try: + blocks: Final = _CONTENT_BLOCKS_ADAPTER.validate_python(content) + except ValidationError: + return None + extended: Final = [*blocks, {"type": "text", "text": generated_content}] + return [*messages[:-1], {**last, "content": extended, "prefix": True}] @staticmethod def _build_responses_continuation_input( diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index d59674f1c57..3d592737e1e 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -2901,6 +2901,21 @@ def test_build_completion_continuation_input_folds_into_trailing_plain_assistant assert built[-1] == {"role": "assistant", "content": "Sure, here it is", "prefix": True} +def test_build_completion_continuation_input_keeps_structured_assistant_content(): + """Content blocks on a trailing assistant turn stay blocks: the partial lands as one more text + block instead of a stringified list, and a content shape that cannot be extended declines.""" + messages = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": [{"type": "text", "text": "Sure, "}]}] + built = litellm.Router._build_completion_continuation_input(messages, "here it is") + assert built is not None + assert [m["role"] for m in built] == ["user", "assistant"] + assert built[-1] == { + "role": "assistant", + "content": [{"type": "text", "text": "Sure, "}, {"type": "text", "text": "here it is"}], + "prefix": True, + } + assert litellm.Router._build_completion_continuation_input([{"role": "assistant", "content": 42}], "x") is None + + 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.""" From 1fc3520f7d6ce7d59ff006c3ffa6488a2657e88e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 1 Oct 2026 18:40:26 -0700 Subject: [PATCH 17/18] fix(router): continue the trailing text block in place when extending a prefill with content blocks Anthropic strips trailing whitespace per content block, so appending the partial output as a new text block after an assistant turn that ends in "Sure, " would send the prefill as "Sure,here". The continuation now extends the last text block itself and only opens a new text block when the trailing block is not text (an image, say), so the prefill reads exactly as the text the caller already received --- litellm/router.py | 10 +++++++++- tests/unit/test_router/test_router.py | 16 +++++++++------- 2 files changed, 18 insertions(+), 8 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index bb8485e0ef5..0b75a422a45 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -478,6 +478,14 @@ _SILENT_MODEL_ADAPTER: Final = TypeAdapter(str | list[str]) _CONTENT_BLOCKS_ADAPTER: Final = TypeAdapter(list[dict[str, object]]) +def _extend_trailing_text_block(blocks: Sequence[Mapping[str, object]], text: str) -> Sequence[Mapping[str, object]]: + """Continue the last text block in place so the prefill reads exactly as the text the caller already received.""" + last: Final = blocks[-1] if blocks else None + if last is not None and last.get("type") == "text" and isinstance(last_text := last.get("text"), str): + return [*blocks[:-1], {**last, "text": f"{last_text}{text}"}] + return [*blocks, {"type": "text", "text": text}] + + def _as_retry_skipped_deployment_ids(value: object) -> tuple[str, ...]: return tuple(item for item in value if isinstance(item, str)) if isinstance(value, tuple) else () @@ -3260,7 +3268,7 @@ class Router: blocks: Final = _CONTENT_BLOCKS_ADAPTER.validate_python(content) except ValidationError: return None - extended: Final = [*blocks, {"type": "text", "text": generated_content}] + extended: Final = _extend_trailing_text_block(blocks, generated_content) return [*messages[:-1], {**last, "content": extended, "prefix": True}] @staticmethod diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 3d592737e1e..86778b54158 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -2902,17 +2902,19 @@ def test_build_completion_continuation_input_folds_into_trailing_plain_assistant def test_build_completion_continuation_input_keeps_structured_assistant_content(): - """Content blocks on a trailing assistant turn stay blocks: the partial lands as one more text - block instead of a stringified list, and a content shape that cannot be extended declines.""" + """Content blocks on a trailing assistant turn stay blocks: the partial continues the last text + block in place (a block boundary would let a provider drop the space between "Sure, " and "here"), + lands as a new text block only after a non-text block, and a content shape that cannot be extended declines.""" + from litellm.router import _extend_trailing_text_block + messages = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": [{"type": "text", "text": "Sure, "}]}] built = litellm.Router._build_completion_continuation_input(messages, "here it is") assert built is not None assert [m["role"] for m in built] == ["user", "assistant"] - assert built[-1] == { - "role": "assistant", - "content": [{"type": "text", "text": "Sure, "}, {"type": "text", "text": "here it is"}], - "prefix": True, - } + assert built[-1] == {"role": "assistant", "content": [{"type": "text", "text": "Sure, here it is"}], "prefix": True} + image_block = {"type": "image_url", "image_url": {"url": "https://example.test/a.png"}} + assert _extend_trailing_text_block([image_block], "here it is") == [image_block, {"type": "text", "text": "here it is"}] + assert _extend_trailing_text_block([], "here it is") == [{"type": "text", "text": "here it is"}] assert litellm.Router._build_completion_continuation_input([{"role": "assistant", "content": 42}], "x") is None From 4fa01f97022107889ccdcd9f6a442af9be02cae9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 1 Oct 2026 18:47:34 -0700 Subject: [PATCH 18/18] test(router): type the prefill check's deployment helper --- .../pre_call_checks/test_continuation_prefill_check.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py b/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py index 4cbe48eda54..77a3db58336 100644 --- a/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py @@ -25,7 +25,7 @@ def _synthetic_cost_map_entries() -> None: litellm.get_model_info.cache_clear() -def _deployment(model: str, dep_id: str) -> dict: +def _deployment(model: str, dep_id: str) -> dict[str, object]: return {"litellm_params": {"model": model}, "model_info": {"id": dep_id}}