From 95d87d0d0dee9e10baba5dbb0d81a5b763e53c60 Mon Sep 17 00:00:00 2001 From: Jaeyeon Kim Date: Mon, 2 Nov 2026 13:03:07 +0000 Subject: [PATCH] fix(router): schedule generic mirrors once per request chain Signed-off-by: Jaeyeon Kim --- litellm/router_python_core.py | 22 ++++++- tests/unit/test_router_silent_experiment.py | 68 +++++++++++++++++++++ 2 files changed, 89 insertions(+), 1 deletion(-) diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index a3ad51fa6e5..a52c4eead9a 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -1067,6 +1067,13 @@ class _StreamRetryOpened: max_retries: int +@dataclass(slots=True) +class _GenericMirrorState: + """Shared by the attempts of one request, including mid-stream fallback hops.""" + + scheduled: bool = False + + @dataclass(frozen=True, slots=True) class _StreamRetriesExhausted: error: "MidStreamFallbackError" @@ -6827,6 +6834,7 @@ class PythonCore(DeprecatedRouterMembers): kwargs["model"] = model kwargs["original_generic_function"] = original_function kwargs["original_function"] = attempt_function or self._ageneric_api_call_with_fallbacks_helper + kwargs["_silent_experiment_state"] = _GenericMirrorState() # rebind-ok: request-owned attempt state if attempt_function is not None: controls: Final = per_request_fallback_controls(kwargs) kwargs[MID_STREAM_FALLBACK_CONTROLS_KEY] = controls # rebind-ok: forwarded to every hop @@ -6874,12 +6882,21 @@ class PythonCore(DeprecatedRouterMembers): kwargs["endpoint"] = replace_path_segment(kwargs["endpoint"], model, replacement_model_name) return kwargs - async def _ageneric_api_call_with_fallbacks_helper(self, model: str, original_generic_function: Callable, **kwargs): + async def _ageneric_api_call_with_fallbacks_helper( + self, + model: str, + original_generic_function: Callable, + _silent_experiment_state: object = None, + **kwargs, + ): """ Helper function to make a generic LLM API call through the router, this allows you to use retries/fallbacks with litellm router """ passthrough_on_no_deployment: Final = kwargs.pop("passthrough_on_no_deployment", False) + mirror_state: Final = ( + _silent_experiment_state if isinstance(_silent_experiment_state, _GenericMirrorState) else None + ) function_name: Final = "_ageneric_api_call_with_fallbacks" handler_name: Final[str | None] = getattr(original_generic_function, "__name__", None) deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked @@ -6913,12 +6930,15 @@ class PythonCore(DeprecatedRouterMembers): silent_model is not None and handler_name in SILENT_MODEL_MIRROR_ALLOWED_CALL_TYPES and not deployment_has_mcp + and (mirror_state is None or not mirror_state.scheduled) ): # Snapshot now: `_update_kwargs_with_deployment` below mutates kwargs in place, # and the task only starts running once this coroutine yields. for silent_target in _silent_experiment_targets(silent_model): silent_kwargs = self._silent_experiment_generic_kwargs(**kwargs) if silent_kwargs is not None: + if mirror_state is not None: + mirror_state.scheduled = True asyncio.create_task( self._silent_experiment_ageneric( silent_model=silent_target, diff --git a/tests/unit/test_router_silent_experiment.py b/tests/unit/test_router_silent_experiment.py index ff7f8148310..0869d585cc8 100644 --- a/tests/unit/test_router_silent_experiment.py +++ b/tests/unit/test_router_silent_experiment.py @@ -1381,3 +1381,71 @@ async def test_calls_a_failure_function_callback_starts_for_a_mirrored_call_are_ await _wait_for_judges(sync_and_async_recording_logger, {"shadow-a"}) assert sync_and_async_recording_logger.judge_flags() == {"shadow-a": None} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", ["aresponses", "anthropic_messages"]) +@pytest.mark.parametrize("use_fallback", [False, True], ids=["retry", "fallback"]) +async def test_generic_mirrors_only_once_across_attempts(call_type: str, use_fallback: bool) -> None: + calls: Final[list[tuple[str, bool]]] = [] + + async def provider_call(**kwargs): + model = kwargs["model"] + mirror = in_mirrored_call() + calls.append((model, mirror)) + if not mirror and len([call for call in calls if not call[1]]) == 1: + raise litellm.InternalServerError(message="retry me", llm_provider="openai", model=model) + return "provider-result" + + provider: Final = AsyncMock(side_effect=provider_call) + provider.__name__ = call_type + models = _generic_silent_model_list() + models.append( + { + "model_name": "fallback-model", + "litellm_params": {"model": "openai/gpt-4.1-mini", "api_key": "fake-key", "silent_model": "silent-model"}, + } + ) + router: Final = Router( + model_list=models, + num_retries=0 if use_fallback else 1, + retry_after=0, + disable_cooldowns=True, + fallbacks=[{"primary-model": ["fallback-model"]}] if use_fallback else [], + ) + endpoint: Final = router.factory_function(provider, call_type=call_type) + result = await endpoint(model="primary-model", input="hi") + mirrors = tuple( + task for task in asyncio.all_tasks() if task.get_coro().__qualname__.endswith("._silent_experiment_ageneric") + ) + await asyncio.gather(*mirrors) + + assert result == "provider-result" + assert len([call for call in calls if not call[1]]) == 2 + assert len([call for call in calls if call[1]]) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", ["aresponses", "anthropic_messages"]) +async def test_generic_mirror_attempt_state_is_per_request(call_type: str) -> None: + provider: Final = AsyncMock(return_value="provider-result") + provider.__name__ = call_type + router: Final = Router(model_list=_generic_silent_model_list(), num_retries=0) + endpoint: Final = router.factory_function(provider, call_type=call_type) + + assert await asyncio.gather( + endpoint(model="primary-model", input="request-a"), + endpoint(model="primary-model", input="request-b"), + ) == ["provider-result", "provider-result"] + mirrors: Final = tuple( + task for task in asyncio.all_tasks() if task.get_coro().__qualname__.endswith("._silent_experiment_ageneric") + ) + await asyncio.gather(*mirrors) + + assert provider.await_count == 4 + for request in ("request-a", "request-b"): + assert {call.kwargs["model"] for call in provider.await_args_list if call.kwargs["input"] == request} == { + "openai/gpt-4o-mini", + "openai/gpt-4o", + } + assert all("_silent_experiment_state" not in call.kwargs for call in provider.await_args_list)