From 8cf090b3683c4a5fe31ebc9ef6f742befd4398a9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 01:46:33 -0700 Subject: [PATCH] fix(router): arm the provider-scoped fallback pin only on resource-operating handlers --- .../router_utils/fallback_event_handlers.py | 27 ++++++- .../test_fallback_event_handlers.py | 80 ++++++++++++++++++- 2 files changed, 101 insertions(+), 6 deletions(-) diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 23eeb495d61..924574537f3 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -253,6 +253,17 @@ def get_fallback_model_group(fallbacks: list[Any], model_group: str) -> tuple[li PROVIDER_SCOPED_RESOURCE_KEYS: Final = ("input_file_id", "training_file", "batch_id", "file_id", "fine_tuning_job_id") +PROVIDER_SCOPED_RESOURCE_FUNCTION_NAMES: Final = frozenset( + { + "_acreate_batch", + "_acancel_batch", + "acreate_fine_tuning_job", + "acancel_fine_tuning_job", + "aretrieve_fine_tuning_job", + "afile_content", + "afile_delete", + } +) PROVIDER_SCOPED_CREATION_FUNCTION_NAMES: Final = frozenset({"_acreate_file"}) @@ -284,13 +295,23 @@ async def _is_fallback_target_authorized( def references_provider_scoped_resource(kwargs: Mapping[str, object]) -> bool: """ - True when the request names a file, batch, or fine-tuning job that only exists under - one provider's credentials. + True when a file, batch, or fine-tuning job operation names an id that only exists + under one provider's credentials. Each of those ids lives in the account of the deployment that issued it. Handing it to a different model group asks a provider about an id it never issued, which costs an - extra round trip that can only answer not-found. + extra round trip that can only answer not-found. Generic calls dispatched through + `Router._ageneric_api_call_with_fallbacks` carry the real handler in + `original_generic_function`, so both slots are checked. Gating on the handler name + keeps completion-style requests eligible for cross-group fallback even when a caller + passes a stray extra body field that happens to share one of these key names. """ + handler_names: Final = tuple( + getattr(kwargs.get(function_key), "__name__", None) + for function_key in ("original_function", "original_generic_function") + ) + if all(name not in PROVIDER_SCOPED_RESOURCE_FUNCTION_NAMES for name in handler_names): + return False return any(kwargs.get(key) for key in PROVIDER_SCOPED_RESOURCE_KEYS) diff --git a/tests/test_litellm/router_utils/test_fallback_event_handlers.py b/tests/test_litellm/router_utils/test_fallback_event_handlers.py index c993739e364..8336926c050 100644 --- a/tests/test_litellm/router_utils/test_fallback_event_handlers.py +++ b/tests/test_litellm/router_utils/test_fallback_event_handlers.py @@ -180,6 +180,30 @@ async def _acreate_file(*args: object, **kwargs: object) -> NoReturn: raise AssertionError("only used for its __name__") +async def _acancel_batch(*args: object, **kwargs: object) -> NoReturn: + raise AssertionError("only used for its __name__") + + +async def _acompletion(*args: object, **kwargs: object) -> NoReturn: + raise AssertionError("only used for its __name__") + + +async def _ageneric_api_call_with_fallbacks_helper(*args: object, **kwargs: object) -> NoReturn: + raise AssertionError("only used for its __name__") + + +async def acreate_fine_tuning_job(*args: object, **kwargs: object) -> NoReturn: + raise AssertionError("only used for its __name__") + + +async def aretrieve_fine_tuning_job(*args: object, **kwargs: object) -> NoReturn: + raise AssertionError("only used for its __name__") + + +async def afile_content(*args: object, **kwargs: object) -> NoReturn: + raise AssertionError("only used for its __name__") + + @pytest.mark.asyncio async def test_run_async_fallback_keeps_uploaded_file_requests_in_their_model_group(): """An input_file_id only exists under the credentials of the group it was uploaded @@ -217,6 +241,8 @@ async def test_run_async_fallback_keeps_fine_tuning_requests_in_their_model_grou fallback_depth=0, model="openai-group", training_file="file-owned-by-openai", + original_function=_ageneric_api_call_with_fallbacks_helper, + original_generic_function=acreate_fine_tuning_job, ) assert router.attempted_model_groups == [] @@ -300,10 +326,33 @@ async def test_run_async_fallback_still_crosses_model_groups_without_an_uploaded @pytest.mark.asyncio -@pytest.mark.parametrize("resource_key", ["batch_id", "file_id", "fine_tuning_job_id"]) -async def test_run_async_fallback_keeps_provider_scoped_ids_in_their_model_group(resource_key: str): +@pytest.mark.parametrize( + ("resource_key", "handler_kwargs"), + [ + ("batch_id", {"original_function": _acancel_batch}), + ( + "file_id", + { + "original_function": _ageneric_api_call_with_fallbacks_helper, + "original_generic_function": afile_content, + }, + ), + ( + "fine_tuning_job_id", + { + "original_function": _ageneric_api_call_with_fallbacks_helper, + "original_generic_function": aretrieve_fine_tuning_job, + }, + ), + ], +) +async def test_run_async_fallback_keeps_provider_scoped_ids_in_their_model_group( + resource_key: str, handler_kwargs: dict +): """A batch, file, or fine-tuning job id only exists under the credentials of the group - that issued it, so a cross-group fallback asks a provider about an id it never saw.""" + that issued it, so a cross-group fallback asks a provider about an id it never saw. + Generic API calls carry the real handler in original_generic_function, so the pin + must recognize it there too.""" router = AttemptRecordingRouter() with pytest.raises(RuntimeError, match="openai connection error"): @@ -316,11 +365,35 @@ async def test_run_async_fallback_keeps_provider_scoped_ids_in_their_model_group fallback_depth=0, model="openai-group", **{resource_key: "owned-by-openai"}, + **handler_kwargs, ) assert router.attempted_model_groups == [] +@pytest.mark.asyncio +@pytest.mark.parametrize("resource_key", ["batch_id", "file_id", "fine_tuning_job_id"]) +async def test_run_async_fallback_ignores_stray_resource_ids_on_completion_calls(resource_key: str): + """A caller-supplied top-level field like file_id on a chat completion is application + data, never a provider resource reference, so it must not cost the request its + cross-group fallbacks.""" + router = AttemptRecordingRouter() + + await run_async_fallback( + litellm_router=router, + fallback_model_group=["azure-group"], + original_model_group="openai-group", + original_exception=RuntimeError("openai connection error"), + max_fallbacks=3, + fallback_depth=0, + model="openai-group", + original_function=_acompletion, + **{resource_key: "caller-app-data"}, + ) + + assert router.attempted_model_groups == ["azure-group"] + + @pytest.mark.asyncio async def test_run_async_fallback_allows_same_model_group_retry_for_batch_cancel(): router = AttemptRecordingRouter() @@ -334,6 +407,7 @@ async def test_run_async_fallback_allows_same_model_group_retry_for_batch_cancel fallback_depth=0, model="openai-group", batch_id="owned-by-openai", + original_function=_acancel_batch, ) assert router.attempted_model_groups == ["openai-group"]