From e3291d79ae4f6b0e306dd945ed54638a133597dc Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 26 Sep 2026 02:46:31 +0000 Subject: [PATCH 1/5] fix(router): treat a max_budget of zero as a hard stop A provider, deployment, or tag cap of 0 was skipped because the spend check used the dollar amount as a boolean. Compare against None instead, and keep checking deployment and tag caps when a provider entry has no dollar limit. Co-authored-by: abeed09 --- litellm/router_strategy/budget_limiter.py | 43 +- .../router_strategy/test_budget_limiter.py | 434 ++++++++++++++++++ 2 files changed, 459 insertions(+), 18 deletions(-) diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 64252cbbfb3..8ecc7db0af9 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -237,20 +237,20 @@ class RouterBudgetLimiting(CustomLogger): provider = self._get_llm_provider_for_deployment(deployment) if provider in provider_configs: config = provider_configs[provider] - if config.max_budget is None: - continue - current_spend = spend_map.get(f"provider_spend:{provider}:{config.budget_duration}", 0.0) - self._track_provider_remaining_budget_prometheus( - provider=provider, - spend=current_spend, - budget_limit=config.max_budget, - ) - - if config.max_budget and current_spend >= config.max_budget: - debug_msg = f"Exceeded budget for provider {provider}: {current_spend} >= {config.max_budget}" - deployment_above_budget_info += f"{debug_msg}\n" - is_within_budget = False - continue + if config.max_budget is not None: + current_spend = spend_map.get(f"provider_spend:{provider}:{config.budget_duration}", 0.0) + self._track_provider_remaining_budget_prometheus( + provider=provider, + spend=current_spend, + budget_limit=config.max_budget, + ) + if current_spend >= config.max_budget: + debug_msg = ( + f"Exceeded budget for provider {provider}: {current_spend} >= {config.max_budget}" + ) + deployment_above_budget_info += f"{debug_msg}\n" + is_within_budget = False + continue # Check deployment budget if self.deployment_budget_config and is_within_budget: @@ -261,8 +261,12 @@ class RouterBudgetLimiting(CustomLogger): if model_id in deployment_configs: config = deployment_configs[model_id] current_spend = spend_map.get(f"deployment_spend:{model_id}:{config.budget_duration}", 0.0) - if config.max_budget and current_spend >= config.max_budget: - debug_msg = f"Exceeded budget for deployment model_name: {_model_name}, litellm_params.model: {_litellm_model_name}, model_id: {model_id}: {current_spend} >= {config.budget_duration}" + if config.max_budget is not None and current_spend >= config.max_budget: + debug_msg = ( + f"Exceeded budget for deployment model_name: {_model_name}, " + f"litellm_params.model: {_litellm_model_name}, model_id: {model_id}: " + f"{current_spend} >= {config.max_budget}" + ) verbose_router_logger.debug(debug_msg) deployment_above_budget_info += f"{debug_msg}\n" is_within_budget = False @@ -276,8 +280,11 @@ class RouterBudgetLimiting(CustomLogger): f"tag_spend:{_tag}:{_tag_budget_config.budget_duration}", 0.0, ) - if _tag_budget_config.max_budget and _tag_spend >= _tag_budget_config.max_budget: - debug_msg = f"Exceeded budget for tag='{_tag}', tag_spend={_tag_spend}, tag_budget_limit={_tag_budget_config.max_budget}" + if _tag_budget_config.max_budget is not None and _tag_spend >= _tag_budget_config.max_budget: + debug_msg = ( + f"Exceeded budget for tag='{_tag}', tag_spend={_tag_spend}, " + f"tag_budget_limit={_tag_budget_config.max_budget}" + ) verbose_router_logger.debug(debug_msg) deployment_above_budget_info += f"{debug_msg}\n" is_within_budget = False diff --git a/tests/unit/router_strategy/test_budget_limiter.py b/tests/unit/router_strategy/test_budget_limiter.py index 62de1586fdd..68d9ba022c8 100644 --- a/tests/unit/router_strategy/test_budget_limiter.py +++ b/tests/unit/router_strategy/test_budget_limiter.py @@ -6,12 +6,17 @@ anthropic_messages, embedding and rerank surfaces leave it unset, which used to the callback raise before any spend was recorded, so those budgets never moved. """ +import asyncio from typing import Final import pytest +import litellm from litellm.caching.caching import DualCache +from litellm.router import Router from litellm.router_strategy.budget_limiter import RouterBudgetLimiting +from litellm.types.router import RouterErrors +from litellm.types.utils import BudgetConfig @pytest.fixture @@ -135,3 +140,432 @@ async def test_deployment_budget_tracked_when_provider_is_unresolvable(disable_b ) assert await limiter.dual_cache.async_get_cache("deployment_spend:deployment-1:1d") == 0.25 + + +_DURATION: Final = "1d" +_MODEL_ID: Final = "dep-1" +_TAG: Final = "prod" +_BUDGET_CASES: Final = ( + (None, 50.0, True), + (0.0, 0.0, False), + (0.0, 50.0, False), + (0.01, 0.0, True), + (0.01, 0.01, False), + (0.01, 50.0, False), +) +_BUDGET_CASE_IDS: Final = ( + "unset-cap-stays-routable", + "zero-cap-blocks-with-no-spend", + "zero-cap-blocks-recorded-spend", + "under-cap-stays-routable", + "spend-equal-to-cap-is-blocked", + "spend-above-cap-is-blocked", +) + + +@pytest.fixture +def isolated_callbacks(disable_budget_sync): + before: Final = tuple(litellm.callbacks) if isinstance(litellm.callbacks, list) else () + if isinstance(litellm.callbacks, list): + litellm.callbacks[:] = [ + callback for callback in litellm.callbacks if not isinstance(callback, RouterBudgetLimiting) + ] + yield + if isinstance(litellm.callbacks, list): + litellm.callbacks[:] = list(before) + + +def _budget(max_budget: float | None) -> BudgetConfig: + return BudgetConfig(max_budget=max_budget, budget_duration=_DURATION) + + +def _deployment( + *, + model: str = "openai/gpt-4", + model_id: str = _MODEL_ID, + model_name: str = "gpt-4", +) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": {"model": model}, + "model_info": {"id": model_id}, + } + + +def _limiter() -> RouterBudgetLimiting: + return RouterBudgetLimiting(dual_cache=DualCache(), provider_budget_config=None) + + +def _filter( + limiter: RouterBudgetLimiting, + deployment: dict[str, object], + *, + provider_configs: dict[str, BudgetConfig], + deployment_configs: dict[str, BudgetConfig], + deployment_providers: list[str | None], + spend_map: dict[str, float], + request_tags: list[str], +) -> tuple[list[dict[str, object]], str]: + return limiter._filter_out_deployments_above_budget( + potential_deployments=[], + healthy_deployments=[deployment], + provider_configs=provider_configs, + deployment_configs=deployment_configs, + deployment_providers=deployment_providers, + spend_map=spend_map, + request_tags=request_tags, + ) + + +def _assert_routed( + kept: list[dict[str, object]], + debug_info: str, + deployment: dict[str, object], + *, + stays: bool, + blocked_debug: str, +) -> None: + if stays: + assert kept == [deployment] + assert debug_info == "" + return + assert kept == [] + assert debug_info == blocked_debug + + +async def _drain_background_tasks() -> None: + current: Final = asyncio.current_task() + pending: Final = tuple(task for task in asyncio.all_tasks() if task is not current and not task.done()) + if pending: + await asyncio.gather(*pending) + + +def _router_model( + *, + model_name: str, + model: str, + mock_response: str, + model_id: str, + max_budget: float | None = None, +) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": { + "model": model, + "mock_response": mock_response, + "api_key": "sk-fake", + **({} if max_budget is None else {"max_budget": max_budget, "budget_duration": _DURATION}), + }, + "model_info": {"id": model_id}, + } + + +async def _complete(router: Router, model: str, **extra: object) -> str: + response: Final = await router.acompletion( + model=model, + messages=[{"role": "user", "content": "hi"}], + num_retries=0, + **extra, + ) + content: Final = response.choices[0].message.content + assert isinstance(content, str) + return content + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("max_budget", "spend", "stays"), _BUDGET_CASES, ids=_BUDGET_CASE_IDS) +async def test_provider_max_budget_treats_zero_as_a_cap_and_none_as_uncapped( + isolated_callbacks, + max_budget: float | None, + spend: float, + stays: bool, +) -> None: + config: Final = _budget(max_budget) + deployment: Final = _deployment() + limiter: Final = _limiter() + limiter.provider_budget_config = {"openai": config} + kept, debug_info = _filter( + limiter, + deployment, + provider_configs={"openai": config}, + deployment_configs={}, + deployment_providers=["openai"], + spend_map={f"provider_spend:openai:{_DURATION}": spend}, + request_tags=[], + ) + _assert_routed( + kept, + debug_info, + deployment, + stays=stays, + blocked_debug=f"Exceeded budget for provider openai: {spend} >= {config.max_budget}\n", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("max_budget", "spend", "stays"), _BUDGET_CASES, ids=_BUDGET_CASE_IDS) +async def test_deployment_max_budget_treats_zero_as_a_cap_and_none_as_uncapped( + isolated_callbacks, + max_budget: float | None, + spend: float, + stays: bool, +) -> None: + config: Final = _budget(max_budget) + deployment: Final = _deployment() + limiter: Final = _limiter() + limiter.deployment_budget_config = {_MODEL_ID: config} + kept, debug_info = _filter( + limiter, + deployment, + provider_configs={}, + deployment_configs={_MODEL_ID: config}, + deployment_providers=[], + spend_map={f"deployment_spend:{_MODEL_ID}:{_DURATION}": spend}, + request_tags=[], + ) + _assert_routed( + kept, + debug_info, + deployment, + stays=stays, + blocked_debug=( + "Exceeded budget for deployment model_name: gpt-4, litellm_params.model: openai/gpt-4, " + f"model_id: {_MODEL_ID}: {spend} >= {config.max_budget}\n" + ), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("max_budget", "spend", "stays"), _BUDGET_CASES, ids=_BUDGET_CASE_IDS) +async def test_tag_max_budget_treats_zero_as_a_cap_and_none_as_uncapped( + isolated_callbacks, + max_budget: float | None, + spend: float, + stays: bool, +) -> None: + config: Final = _budget(max_budget) + deployment: Final = _deployment() + limiter: Final = _limiter() + limiter.tag_budget_config = {_TAG: config} + kept, debug_info = _filter( + limiter, + deployment, + provider_configs={}, + deployment_configs={}, + deployment_providers=[], + spend_map={f"tag_spend:{_TAG}:{_DURATION}": spend}, + request_tags=[_TAG], + ) + _assert_routed( + kept, + debug_info, + deployment, + stays=stays, + blocked_debug=f"Exceeded budget for tag='{_TAG}', tag_spend={spend}, tag_budget_limit={config.max_budget}\n", + ) + + +@pytest.mark.asyncio +async def test_provider_zero_cap_blocks_when_no_spend_has_been_recorded(isolated_callbacks) -> None: + config: Final = _budget(0) + deployment: Final = _deployment() + limiter: Final = _limiter() + limiter.provider_budget_config = {"openai": config} + kept, debug_info = _filter( + limiter, + deployment, + provider_configs={"openai": config}, + deployment_configs={}, + deployment_providers=["openai"], + spend_map={}, + request_tags=[], + ) + _assert_routed( + kept, + debug_info, + deployment, + stays=False, + blocked_debug=f"Exceeded budget for provider openai: 0.0 >= {config.max_budget}\n", + ) + + +@pytest.mark.asyncio +async def test_unset_provider_cap_still_enforces_a_zero_deployment_cap(isolated_callbacks) -> None: + provider: Final = _budget(None) + deployment_cap: Final = _budget(0) + deployment: Final = _deployment() + limiter: Final = _limiter() + limiter.provider_budget_config = {"openai": provider} + limiter.deployment_budget_config = {_MODEL_ID: deployment_cap} + kept, debug_info = _filter( + limiter, + deployment, + provider_configs={"openai": provider}, + deployment_configs={_MODEL_ID: deployment_cap}, + deployment_providers=["openai"], + spend_map={}, + request_tags=[], + ) + _assert_routed( + kept, + debug_info, + deployment, + stays=False, + blocked_debug=( + "Exceeded budget for deployment model_name: gpt-4, litellm_params.model: openai/gpt-4, " + f"model_id: {_MODEL_ID}: 0.0 >= {deployment_cap.max_budget}\n" + ), + ) + + +@pytest.mark.asyncio +async def test_router_refuses_a_provider_whose_max_budget_is_zero(isolated_callbacks) -> None: + router: Final = Router( + model_list=[ + _router_model(model_name="gpt-4", model="openai/gpt-4", mock_response="served", model_id=_MODEL_ID) + ], + provider_budget_config={"openai": _budget(0)}, + ) + await _drain_background_tasks() + + with pytest.raises(ValueError) as exc_info: + await _complete(router, "gpt-4") + + assert str(exc_info.value) == ( + f"{RouterErrors.no_deployments_with_provider_budget_routing.value}: " + "Exceeded budget for provider openai: 0.0 >= 0.0\n" + ) + + +@pytest.mark.asyncio +async def test_router_refuses_a_provider_that_already_spent_against_a_zero_cap(isolated_callbacks) -> None: + router: Final = Router( + model_list=[ + _router_model(model_name="gpt-4", model="openai/gpt-4", mock_response="served", model_id=_MODEL_ID) + ], + provider_budget_config={"openai": _budget(0)}, + ) + await _drain_background_tasks() + logger: Final = router.router_budget_logger + assert logger is not None + await logger.dual_cache.async_set_cache(key=f"provider_spend:openai:{_DURATION}", value=50.0) + + with pytest.raises(ValueError) as exc_info: + await _complete(router, "gpt-4") + + assert str(exc_info.value) == ( + f"{RouterErrors.no_deployments_with_provider_budget_routing.value}: " + "Exceeded budget for provider openai: 50.0 >= 0.0\n" + ) + + +@pytest.mark.asyncio +async def test_router_serves_a_provider_under_its_cap(isolated_callbacks) -> None: + router: Final = Router( + model_list=[ + _router_model(model_name="gpt-4", model="openai/gpt-4", mock_response="served", model_id=_MODEL_ID) + ], + provider_budget_config={"openai": _budget(0.01)}, + ) + await _drain_background_tasks() + + assert await _complete(router, "gpt-4") == "served" + + +@pytest.mark.asyncio +async def test_router_serves_when_the_provider_cap_is_unset(isolated_callbacks) -> None: + router: Final = Router( + model_list=[ + _router_model(model_name="gpt-4", model="openai/gpt-4", mock_response="served", model_id=_MODEL_ID) + ], + provider_budget_config={"openai": _budget(None)}, + ) + await _drain_background_tasks() + + assert await _complete(router, "gpt-4") == "served" + + +@pytest.mark.asyncio +async def test_router_keeps_an_uncapped_provider_when_another_is_at_zero(isolated_callbacks) -> None: + router: Final = Router( + model_list=[ + _router_model(model_name="chat", model="openai/gpt-4", mock_response="from openai", model_id="dep-openai"), + _router_model( + model_name="chat", + model="anthropic/claude-3-5-sonnet-latest", + mock_response="from anthropic", + model_id="dep-anthropic", + ), + ], + provider_budget_config={"openai": _budget(0)}, + ) + await _drain_background_tasks() + + served: Final = tuple([await _complete(router, "chat") for _ in range(20)]) + assert served == ("from anthropic",) * 20 + + +@pytest.mark.asyncio +async def test_router_refuses_a_deployment_whose_max_budget_is_zero(isolated_callbacks) -> None: + router: Final = Router( + model_list=[ + _router_model( + model_name="gpt-4", + model="openai/gpt-4", + mock_response="served", + model_id=_MODEL_ID, + max_budget=0, + ) + ], + ) + await _drain_background_tasks() + + with pytest.raises(ValueError) as exc_info: + await _complete(router, "gpt-4") + + assert str(exc_info.value) == ( + f"{RouterErrors.no_deployments_with_provider_budget_routing.value}: " + "Exceeded budget for deployment model_name: gpt-4, litellm_params.model: openai/gpt-4, " + f"model_id: {_MODEL_ID}: 0.0 >= 0.0\n" + ) + + +@pytest.mark.asyncio +async def test_router_serves_a_deployment_under_its_cap(isolated_callbacks) -> None: + router: Final = Router( + model_list=[ + _router_model( + model_name="gpt-4", + model="openai/gpt-4", + mock_response="served", + model_id=_MODEL_ID, + max_budget=0.01, + ) + ], + ) + await _drain_background_tasks() + + assert await _complete(router, "gpt-4") == "served" + + +@pytest.mark.asyncio +async def test_router_tag_cap_of_zero_blocks_only_tagged_requests(isolated_callbacks) -> None: + router: Final = Router( + model_list=[ + _router_model(model_name="gpt-4", model="openai/gpt-4", mock_response="served", model_id=_MODEL_ID) + ], + provider_budget_config={"openai": _budget(1.0)}, + ) + await _drain_background_tasks() + logger: Final = router.router_budget_logger + assert logger is not None + logger.tag_budget_config = {_TAG: _budget(0)} + + assert await _complete(router, "gpt-4") == "served" + with pytest.raises(ValueError) as exc_info: + await _complete(router, "gpt-4", metadata={"tags": [_TAG]}) + + assert str(exc_info.value) == ( + f"{RouterErrors.no_deployments_with_provider_budget_routing.value}: " + f"Exceeded budget for tag='{_TAG}', tag_spend=0.0, tag_budget_limit=0.0\n" + ) From 6214d5d082838d7c5026187b471595dcee86bb1b Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 26 Sep 2026 03:04:45 +0000 Subject: [PATCH 2/5] test(router): keep the zero-budget cases in one matrix The provider, deployment, and tag checks share one parametrized test, plus the case where an empty provider cap must still honor a deployment cap. Co-authored-by: abeed09 --- litellm/router_strategy/budget_limiter.py | 11 +- .../router_strategy/test_budget_limiter.py | 476 +++--------------- 2 files changed, 79 insertions(+), 408 deletions(-) diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 8ecc7db0af9..c3209b2e080 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -262,11 +262,7 @@ class RouterBudgetLimiting(CustomLogger): config = deployment_configs[model_id] current_spend = spend_map.get(f"deployment_spend:{model_id}:{config.budget_duration}", 0.0) if config.max_budget is not None and current_spend >= config.max_budget: - debug_msg = ( - f"Exceeded budget for deployment model_name: {_model_name}, " - f"litellm_params.model: {_litellm_model_name}, model_id: {model_id}: " - f"{current_spend} >= {config.max_budget}" - ) + debug_msg = f"Exceeded budget for deployment model_name: {_model_name}, litellm_params.model: {_litellm_model_name}, model_id: {model_id}: {current_spend} >= {config.max_budget}" verbose_router_logger.debug(debug_msg) deployment_above_budget_info += f"{debug_msg}\n" is_within_budget = False @@ -281,10 +277,7 @@ class RouterBudgetLimiting(CustomLogger): 0.0, ) if _tag_budget_config.max_budget is not None and _tag_spend >= _tag_budget_config.max_budget: - debug_msg = ( - f"Exceeded budget for tag='{_tag}', tag_spend={_tag_spend}, " - f"tag_budget_limit={_tag_budget_config.max_budget}" - ) + debug_msg = f"Exceeded budget for tag='{_tag}', tag_spend={_tag_spend}, tag_budget_limit={_tag_budget_config.max_budget}" verbose_router_logger.debug(debug_msg) deployment_above_budget_info += f"{debug_msg}\n" is_within_budget = False diff --git a/tests/unit/router_strategy/test_budget_limiter.py b/tests/unit/router_strategy/test_budget_limiter.py index 68d9ba022c8..a6d4842a27b 100644 --- a/tests/unit/router_strategy/test_budget_limiter.py +++ b/tests/unit/router_strategy/test_budget_limiter.py @@ -6,16 +6,12 @@ anthropic_messages, embedding and rerank surfaces leave it unset, which used to the callback raise before any spend was recorded, so those budgets never moved. """ -import asyncio -from typing import Final +from typing import Final, Literal import pytest -import litellm from litellm.caching.caching import DualCache -from litellm.router import Router from litellm.router_strategy.budget_limiter import RouterBudgetLimiting -from litellm.types.router import RouterErrors from litellm.types.utils import BudgetConfig @@ -142,430 +138,112 @@ async def test_deployment_budget_tracked_when_provider_is_unresolvable(disable_b assert await limiter.dual_cache.async_get_cache("deployment_spend:deployment-1:1d") == 0.25 +_Scope = Literal["provider", "deployment", "tag"] _DURATION: Final = "1d" _MODEL_ID: Final = "dep-1" -_TAG: Final = "prod" -_BUDGET_CASES: Final = ( - (None, 50.0, True), - (0.0, 0.0, False), - (0.0, 50.0, False), - (0.01, 0.0, True), - (0.01, 0.01, False), - (0.01, 50.0, False), -) -_BUDGET_CASE_IDS: Final = ( - "unset-cap-stays-routable", - "zero-cap-blocks-with-no-spend", - "zero-cap-blocks-recorded-spend", - "under-cap-stays-routable", - "spend-equal-to-cap-is-blocked", - "spend-above-cap-is-blocked", -) +_DEPLOYMENT: Final[dict[str, object]] = { + "model_name": "gpt-4", + "litellm_params": {"model": "openai/gpt-4"}, + "model_info": {"id": _MODEL_ID}, +} -@pytest.fixture -def isolated_callbacks(disable_budget_sync): - before: Final = tuple(litellm.callbacks) if isinstance(litellm.callbacks, list) else () - if isinstance(litellm.callbacks, list): - litellm.callbacks[:] = [ - callback for callback in litellm.callbacks if not isinstance(callback, RouterBudgetLimiting) - ] - yield - if isinstance(litellm.callbacks, list): - litellm.callbacks[:] = list(before) +def _blocked_debug(scope: _Scope, spend: float, cap: float | None) -> str: + if scope == "provider": + return f"Exceeded budget for provider openai: {spend} >= {cap}\n" + if scope == "deployment": + return ( + "Exceeded budget for deployment model_name: gpt-4, litellm_params.model: openai/gpt-4, " + f"model_id: {_MODEL_ID}: {spend} >= {cap}\n" + ) + return f"Exceeded budget for tag='prod', tag_spend={spend}, tag_budget_limit={cap}\n" -def _budget(max_budget: float | None) -> BudgetConfig: - return BudgetConfig(max_budget=max_budget, budget_duration=_DURATION) - - -def _deployment( - *, - model: str = "openai/gpt-4", - model_id: str = _MODEL_ID, - model_name: str = "gpt-4", -) -> dict[str, object]: - return { - "model_name": model_name, - "litellm_params": {"model": model}, - "model_info": {"id": model_id}, - } - - -def _limiter() -> RouterBudgetLimiting: - return RouterBudgetLimiting(dual_cache=DualCache(), provider_budget_config=None) - - -def _filter( - limiter: RouterBudgetLimiting, - deployment: dict[str, object], - *, - provider_configs: dict[str, BudgetConfig], - deployment_configs: dict[str, BudgetConfig], - deployment_providers: list[str | None], - spend_map: dict[str, float], - request_tags: list[str], +def _filter_one( + scope: _Scope, + config: BudgetConfig, + spend: float, ) -> tuple[list[dict[str, object]], str]: + limiter: Final = RouterBudgetLimiting(dual_cache=DualCache(), provider_budget_config=None) + provider_configs: dict[str, BudgetConfig] = {} + deployment_configs: dict[str, BudgetConfig] = {} + deployment_providers: list[str | None] = [] + request_tags: list[str] = [] + spend_key: str = f"provider_spend:openai:{_DURATION}" + if scope == "provider": + limiter.provider_budget_config = {"openai": config} + provider_configs = {"openai": config} + deployment_providers = ["openai"] + elif scope == "deployment": + limiter.deployment_budget_config = {_MODEL_ID: config} + deployment_configs = {_MODEL_ID: config} + spend_key = f"deployment_spend:{_MODEL_ID}:{_DURATION}" + else: + limiter.tag_budget_config = {"prod": config} + request_tags = ["prod"] + spend_key = f"tag_spend:prod:{_DURATION}" return limiter._filter_out_deployments_above_budget( potential_deployments=[], - healthy_deployments=[deployment], + healthy_deployments=[_DEPLOYMENT], provider_configs=provider_configs, deployment_configs=deployment_configs, deployment_providers=deployment_providers, - spend_map=spend_map, + spend_map={spend_key: spend}, request_tags=request_tags, ) -def _assert_routed( - kept: list[dict[str, object]], - debug_info: str, - deployment: dict[str, object], - *, +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("scope", "max_budget", "spend", "stays"), + [ + ("provider", None, 50.0, True), + ("provider", 0.0, 0.0, False), + ("provider", 0.01, 0.0, True), + ("provider", 0.01, 0.01, False), + ("deployment", None, 50.0, True), + ("deployment", 0.0, 0.0, False), + ("deployment", 0.01, 0.0, True), + ("deployment", 0.01, 0.01, False), + ("tag", None, 50.0, True), + ("tag", 0.0, 0.0, False), + ("tag", 0.01, 0.0, True), + ("tag", 0.01, 0.01, False), + ], +) +async def test_max_budget_of_zero_blocks_and_none_does_not( + disable_budget_sync, + scope: _Scope, + max_budget: float | None, + spend: float, stays: bool, - blocked_debug: str, ) -> None: + config: Final = BudgetConfig(max_budget=max_budget, budget_duration=_DURATION) + kept, debug_info = _filter_one(scope, config, spend) if stays: - assert kept == [deployment] + assert kept == [_DEPLOYMENT] assert debug_info == "" return assert kept == [] - assert debug_info == blocked_debug - - -async def _drain_background_tasks() -> None: - current: Final = asyncio.current_task() - pending: Final = tuple(task for task in asyncio.all_tasks() if task is not current and not task.done()) - if pending: - await asyncio.gather(*pending) - - -def _router_model( - *, - model_name: str, - model: str, - mock_response: str, - model_id: str, - max_budget: float | None = None, -) -> dict[str, object]: - return { - "model_name": model_name, - "litellm_params": { - "model": model, - "mock_response": mock_response, - "api_key": "sk-fake", - **({} if max_budget is None else {"max_budget": max_budget, "budget_duration": _DURATION}), - }, - "model_info": {"id": model_id}, - } - - -async def _complete(router: Router, model: str, **extra: object) -> str: - response: Final = await router.acompletion( - model=model, - messages=[{"role": "user", "content": "hi"}], - num_retries=0, - **extra, - ) - content: Final = response.choices[0].message.content - assert isinstance(content, str) - return content + assert debug_info == _blocked_debug(scope, spend, config.max_budget) @pytest.mark.asyncio -@pytest.mark.parametrize(("max_budget", "spend", "stays"), _BUDGET_CASES, ids=_BUDGET_CASE_IDS) -async def test_provider_max_budget_treats_zero_as_a_cap_and_none_as_uncapped( - isolated_callbacks, - max_budget: float | None, - spend: float, - stays: bool, -) -> None: - config: Final = _budget(max_budget) - deployment: Final = _deployment() - limiter: Final = _limiter() - limiter.provider_budget_config = {"openai": config} - kept, debug_info = _filter( - limiter, - deployment, - provider_configs={"openai": config}, - deployment_configs={}, - deployment_providers=["openai"], - spend_map={f"provider_spend:openai:{_DURATION}": spend}, - request_tags=[], - ) - _assert_routed( - kept, - debug_info, - deployment, - stays=stays, - blocked_debug=f"Exceeded budget for provider openai: {spend} >= {config.max_budget}\n", - ) - - -@pytest.mark.asyncio -@pytest.mark.parametrize(("max_budget", "spend", "stays"), _BUDGET_CASES, ids=_BUDGET_CASE_IDS) -async def test_deployment_max_budget_treats_zero_as_a_cap_and_none_as_uncapped( - isolated_callbacks, - max_budget: float | None, - spend: float, - stays: bool, -) -> None: - config: Final = _budget(max_budget) - deployment: Final = _deployment() - limiter: Final = _limiter() - limiter.deployment_budget_config = {_MODEL_ID: config} - kept, debug_info = _filter( - limiter, - deployment, - provider_configs={}, - deployment_configs={_MODEL_ID: config}, - deployment_providers=[], - spend_map={f"deployment_spend:{_MODEL_ID}:{_DURATION}": spend}, - request_tags=[], - ) - _assert_routed( - kept, - debug_info, - deployment, - stays=stays, - blocked_debug=( - "Exceeded budget for deployment model_name: gpt-4, litellm_params.model: openai/gpt-4, " - f"model_id: {_MODEL_ID}: {spend} >= {config.max_budget}\n" - ), - ) - - -@pytest.mark.asyncio -@pytest.mark.parametrize(("max_budget", "spend", "stays"), _BUDGET_CASES, ids=_BUDGET_CASE_IDS) -async def test_tag_max_budget_treats_zero_as_a_cap_and_none_as_uncapped( - isolated_callbacks, - max_budget: float | None, - spend: float, - stays: bool, -) -> None: - config: Final = _budget(max_budget) - deployment: Final = _deployment() - limiter: Final = _limiter() - limiter.tag_budget_config = {_TAG: config} - kept, debug_info = _filter( - limiter, - deployment, - provider_configs={}, - deployment_configs={}, - deployment_providers=[], - spend_map={f"tag_spend:{_TAG}:{_DURATION}": spend}, - request_tags=[_TAG], - ) - _assert_routed( - kept, - debug_info, - deployment, - stays=stays, - blocked_debug=f"Exceeded budget for tag='{_TAG}', tag_spend={spend}, tag_budget_limit={config.max_budget}\n", - ) - - -@pytest.mark.asyncio -async def test_provider_zero_cap_blocks_when_no_spend_has_been_recorded(isolated_callbacks) -> None: - config: Final = _budget(0) - deployment: Final = _deployment() - limiter: Final = _limiter() - limiter.provider_budget_config = {"openai": config} - kept, debug_info = _filter( - limiter, - deployment, - provider_configs={"openai": config}, - deployment_configs={}, - deployment_providers=["openai"], - spend_map={}, - request_tags=[], - ) - _assert_routed( - kept, - debug_info, - deployment, - stays=False, - blocked_debug=f"Exceeded budget for provider openai: 0.0 >= {config.max_budget}\n", - ) - - -@pytest.mark.asyncio -async def test_unset_provider_cap_still_enforces_a_zero_deployment_cap(isolated_callbacks) -> None: - provider: Final = _budget(None) - deployment_cap: Final = _budget(0) - deployment: Final = _deployment() - limiter: Final = _limiter() - limiter.provider_budget_config = {"openai": provider} +async def test_unset_provider_cap_still_enforces_a_zero_deployment_cap(disable_budget_sync) -> None: + provider: Final = BudgetConfig(max_budget=None, budget_duration=_DURATION) + deployment_cap: Final = BudgetConfig(max_budget=0, budget_duration=_DURATION) + limiter: Final = RouterBudgetLimiting(dual_cache=DualCache(), provider_budget_config={"openai": provider}) limiter.deployment_budget_config = {_MODEL_ID: deployment_cap} - kept, debug_info = _filter( - limiter, - deployment, + + kept, debug_info = limiter._filter_out_deployments_above_budget( + potential_deployments=[], + healthy_deployments=[_DEPLOYMENT], provider_configs={"openai": provider}, deployment_configs={_MODEL_ID: deployment_cap}, deployment_providers=["openai"], spend_map={}, request_tags=[], ) - _assert_routed( - kept, - debug_info, - deployment, - stays=False, - blocked_debug=( - "Exceeded budget for deployment model_name: gpt-4, litellm_params.model: openai/gpt-4, " - f"model_id: {_MODEL_ID}: 0.0 >= {deployment_cap.max_budget}\n" - ), - ) - -@pytest.mark.asyncio -async def test_router_refuses_a_provider_whose_max_budget_is_zero(isolated_callbacks) -> None: - router: Final = Router( - model_list=[ - _router_model(model_name="gpt-4", model="openai/gpt-4", mock_response="served", model_id=_MODEL_ID) - ], - provider_budget_config={"openai": _budget(0)}, - ) - await _drain_background_tasks() - - with pytest.raises(ValueError) as exc_info: - await _complete(router, "gpt-4") - - assert str(exc_info.value) == ( - f"{RouterErrors.no_deployments_with_provider_budget_routing.value}: " - "Exceeded budget for provider openai: 0.0 >= 0.0\n" - ) - - -@pytest.mark.asyncio -async def test_router_refuses_a_provider_that_already_spent_against_a_zero_cap(isolated_callbacks) -> None: - router: Final = Router( - model_list=[ - _router_model(model_name="gpt-4", model="openai/gpt-4", mock_response="served", model_id=_MODEL_ID) - ], - provider_budget_config={"openai": _budget(0)}, - ) - await _drain_background_tasks() - logger: Final = router.router_budget_logger - assert logger is not None - await logger.dual_cache.async_set_cache(key=f"provider_spend:openai:{_DURATION}", value=50.0) - - with pytest.raises(ValueError) as exc_info: - await _complete(router, "gpt-4") - - assert str(exc_info.value) == ( - f"{RouterErrors.no_deployments_with_provider_budget_routing.value}: " - "Exceeded budget for provider openai: 50.0 >= 0.0\n" - ) - - -@pytest.mark.asyncio -async def test_router_serves_a_provider_under_its_cap(isolated_callbacks) -> None: - router: Final = Router( - model_list=[ - _router_model(model_name="gpt-4", model="openai/gpt-4", mock_response="served", model_id=_MODEL_ID) - ], - provider_budget_config={"openai": _budget(0.01)}, - ) - await _drain_background_tasks() - - assert await _complete(router, "gpt-4") == "served" - - -@pytest.mark.asyncio -async def test_router_serves_when_the_provider_cap_is_unset(isolated_callbacks) -> None: - router: Final = Router( - model_list=[ - _router_model(model_name="gpt-4", model="openai/gpt-4", mock_response="served", model_id=_MODEL_ID) - ], - provider_budget_config={"openai": _budget(None)}, - ) - await _drain_background_tasks() - - assert await _complete(router, "gpt-4") == "served" - - -@pytest.mark.asyncio -async def test_router_keeps_an_uncapped_provider_when_another_is_at_zero(isolated_callbacks) -> None: - router: Final = Router( - model_list=[ - _router_model(model_name="chat", model="openai/gpt-4", mock_response="from openai", model_id="dep-openai"), - _router_model( - model_name="chat", - model="anthropic/claude-3-5-sonnet-latest", - mock_response="from anthropic", - model_id="dep-anthropic", - ), - ], - provider_budget_config={"openai": _budget(0)}, - ) - await _drain_background_tasks() - - served: Final = tuple([await _complete(router, "chat") for _ in range(20)]) - assert served == ("from anthropic",) * 20 - - -@pytest.mark.asyncio -async def test_router_refuses_a_deployment_whose_max_budget_is_zero(isolated_callbacks) -> None: - router: Final = Router( - model_list=[ - _router_model( - model_name="gpt-4", - model="openai/gpt-4", - mock_response="served", - model_id=_MODEL_ID, - max_budget=0, - ) - ], - ) - await _drain_background_tasks() - - with pytest.raises(ValueError) as exc_info: - await _complete(router, "gpt-4") - - assert str(exc_info.value) == ( - f"{RouterErrors.no_deployments_with_provider_budget_routing.value}: " - "Exceeded budget for deployment model_name: gpt-4, litellm_params.model: openai/gpt-4, " - f"model_id: {_MODEL_ID}: 0.0 >= 0.0\n" - ) - - -@pytest.mark.asyncio -async def test_router_serves_a_deployment_under_its_cap(isolated_callbacks) -> None: - router: Final = Router( - model_list=[ - _router_model( - model_name="gpt-4", - model="openai/gpt-4", - mock_response="served", - model_id=_MODEL_ID, - max_budget=0.01, - ) - ], - ) - await _drain_background_tasks() - - assert await _complete(router, "gpt-4") == "served" - - -@pytest.mark.asyncio -async def test_router_tag_cap_of_zero_blocks_only_tagged_requests(isolated_callbacks) -> None: - router: Final = Router( - model_list=[ - _router_model(model_name="gpt-4", model="openai/gpt-4", mock_response="served", model_id=_MODEL_ID) - ], - provider_budget_config={"openai": _budget(1.0)}, - ) - await _drain_background_tasks() - logger: Final = router.router_budget_logger - assert logger is not None - logger.tag_budget_config = {_TAG: _budget(0)} - - assert await _complete(router, "gpt-4") == "served" - with pytest.raises(ValueError) as exc_info: - await _complete(router, "gpt-4", metadata={"tags": [_TAG]}) - - assert str(exc_info.value) == ( - f"{RouterErrors.no_deployments_with_provider_budget_routing.value}: " - f"Exceeded budget for tag='{_TAG}', tag_spend=0.0, tag_budget_limit=0.0\n" - ) + assert kept == [] + assert debug_info == _blocked_debug("deployment", 0.0, deployment_cap.max_budget) From 25fc7894cad8c868993abc6a283046e18da16719 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 26 Sep 2026 04:03:08 +0000 Subject: [PATCH 3/5] test(router): cover a zero provider budget on the routing path A hand-built spend map never asks the router which deployment it would send. Route a mocked completion through Router and require the uncapped sibling, with the zero-cap provider weighted so it would win if the filter left it eligible. Co-authored-by: abeed09 --- .../router_strategy/test_budget_limiter.py | 31 +++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/tests/unit/router_strategy/test_budget_limiter.py b/tests/unit/router_strategy/test_budget_limiter.py index a6d4842a27b..cbadebf8096 100644 --- a/tests/unit/router_strategy/test_budget_limiter.py +++ b/tests/unit/router_strategy/test_budget_limiter.py @@ -10,6 +10,7 @@ from typing import Final, Literal import pytest +from litellm import Router from litellm.caching.caching import DualCache from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.types.utils import BudgetConfig @@ -247,3 +248,33 @@ async def test_unset_provider_cap_still_enforces_a_zero_deployment_cap(disable_b assert kept == [] assert debug_info == _blocked_debug("deployment", 0.0, deployment_cap.max_budget) + + +def _router_deployment(model_name: str, model: str, model_id: str, *, weight: int = 0) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": {"model": model, "api_key": "sk-fake", "weight": weight}, + "model_info": {"id": model_id}, + } + + +@pytest.mark.asyncio +async def test_router_serves_the_uncapped_sibling_when_provider_max_budget_is_zero(disable_budget_sync) -> None: + router: Final = Router( + model_list=[ + _router_deployment("chat", "openai/gpt-4o-mini", "openai-capped", weight=100), + _router_deployment("chat", "anthropic/claude-haiku-4-5", "anthropic-open"), + _router_deployment("openai-only", "openai/gpt-4o-mini", "openai-only", weight=100), + ], + provider_budget_config={"openai": BudgetConfig(budget_limit=0, time_period="1d")}, + num_retries=0, + ) + messages: Final = [{"role": "user", "content": "hi"}] + + served: Final = await router.acompletion(model="chat", messages=messages, mock_response="served") + + assert served._hidden_params["model_id"] == "anthropic-open" + assert served._hidden_params["custom_llm_provider"] == "anthropic" + + with pytest.raises(ValueError, match=r"Exceeded budget for provider openai: 0\.0 >= 0\.0"): + await router.acompletion(model="openai-only", messages=messages, mock_response="served") From f370e7f342cf52714bd133f9f7b61b439946d577 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 26 Sep 2026 04:18:17 +0000 Subject: [PATCH 4/5] test(router): stop reassigning budget filter inputs Each zero-budget case is built once and passed through. The router deployment uses DeploymentTypedDict, the model list entry type. Co-authored-by: abeed09 --- .../router_strategy/test_budget_limiter.py | 93 +++++++++++++------ 1 file changed, 66 insertions(+), 27 deletions(-) diff --git a/tests/unit/router_strategy/test_budget_limiter.py b/tests/unit/router_strategy/test_budget_limiter.py index cbadebf8096..336e1cd96d9 100644 --- a/tests/unit/router_strategy/test_budget_limiter.py +++ b/tests/unit/router_strategy/test_budget_limiter.py @@ -6,13 +6,16 @@ anthropic_messages, embedding and rerank surfaces leave it unset, which used to the callback raise before any spend was recorded, so those budgets never moved. """ -from typing import Final, Literal +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Final, Literal, assert_never import pytest from litellm import Router from litellm.caching.caching import DualCache from litellm.router_strategy.budget_limiter import RouterBudgetLimiting +from litellm.types.router import DeploymentTypedDict from litellm.types.utils import BudgetConfig @@ -160,40 +163,75 @@ def _blocked_debug(scope: _Scope, spend: float, cap: float | None) -> str: return f"Exceeded budget for tag='prod', tag_spend={spend}, tag_budget_limit={cap}\n" -def _filter_one( - scope: _Scope, - config: BudgetConfig, +def _apply_budget_filter( + *, + provider_budget: Mapping[str, BudgetConfig] | None, + deployment_budget: Mapping[str, BudgetConfig] | None, + tag_budget: Mapping[str, BudgetConfig] | None, + provider_configs: Mapping[str, BudgetConfig], + deployment_configs: Mapping[str, BudgetConfig], + deployment_providers: Sequence[str | None], + request_tags: Sequence[str], + spend_key: str, spend: float, -) -> tuple[list[dict[str, object]], str]: +): limiter: Final = RouterBudgetLimiting(dual_cache=DualCache(), provider_budget_config=None) - provider_configs: dict[str, BudgetConfig] = {} - deployment_configs: dict[str, BudgetConfig] = {} - deployment_providers: list[str | None] = [] - request_tags: list[str] = [] - spend_key: str = f"provider_spend:openai:{_DURATION}" - if scope == "provider": - limiter.provider_budget_config = {"openai": config} - provider_configs = {"openai": config} - deployment_providers = ["openai"] - elif scope == "deployment": - limiter.deployment_budget_config = {_MODEL_ID: config} - deployment_configs = {_MODEL_ID: config} - spend_key = f"deployment_spend:{_MODEL_ID}:{_DURATION}" - else: - limiter.tag_budget_config = {"prod": config} - request_tags = ["prod"] - spend_key = f"tag_spend:prod:{_DURATION}" + limiter.provider_budget_config = provider_budget + limiter.deployment_budget_config = deployment_budget + limiter.tag_budget_config = tag_budget return limiter._filter_out_deployments_above_budget( - potential_deployments=[], - healthy_deployments=[_DEPLOYMENT], + potential_deployments=[], # mutable-ok: the filter appends each deployment it keeps + healthy_deployments=(_DEPLOYMENT,), provider_configs=provider_configs, deployment_configs=deployment_configs, deployment_providers=deployment_providers, - spend_map={spend_key: spend}, + spend_map=MappingProxyType({spend_key: spend}), request_tags=request_tags, ) +def _filter_one(scope: _Scope, config: BudgetConfig, spend: float): + match scope: + case "provider": + return _apply_budget_filter( + provider_budget=MappingProxyType({"openai": config}), + deployment_budget=None, + tag_budget=None, + provider_configs=MappingProxyType({"openai": config}), + deployment_configs=MappingProxyType({}), + deployment_providers=("openai",), + request_tags=(), + spend_key=f"provider_spend:openai:{_DURATION}", + spend=spend, + ) + case "deployment": + return _apply_budget_filter( + provider_budget=None, + deployment_budget=MappingProxyType({_MODEL_ID: config}), + tag_budget=None, + provider_configs=MappingProxyType({}), + deployment_configs=MappingProxyType({_MODEL_ID: config}), + deployment_providers=(), + request_tags=(), + spend_key=f"deployment_spend:{_MODEL_ID}:{_DURATION}", + spend=spend, + ) + case "tag": + return _apply_budget_filter( + provider_budget=None, + deployment_budget=None, + tag_budget=MappingProxyType({"prod": config}), + provider_configs=MappingProxyType({}), + deployment_configs=MappingProxyType({}), + deployment_providers=(), + request_tags=("prod",), + spend_key=f"tag_spend:prod:{_DURATION}", + spend=spend, + ) + case _: + assert_never(scope) + + @pytest.mark.asyncio @pytest.mark.parametrize( ("scope", "max_budget", "spend", "stays"), @@ -250,12 +288,13 @@ async def test_unset_provider_cap_still_enforces_a_zero_deployment_cap(disable_b assert debug_info == _blocked_debug("deployment", 0.0, deployment_cap.max_budget) -def _router_deployment(model_name: str, model: str, model_id: str, *, weight: int = 0) -> dict[str, object]: - return { +def _router_deployment(model_name: str, model: str, model_id: str, *, weight: int = 0) -> DeploymentTypedDict: + deployment: Final[DeploymentTypedDict] = { "model_name": model_name, "litellm_params": {"model": model, "api_key": "sk-fake", "weight": weight}, "model_info": {"id": model_id}, } + return deployment @pytest.mark.asyncio From 72e1c69e6d0349ad69355ad82d161e4e2481ed3a Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 26 Sep 2026 04:25:44 +0000 Subject: [PATCH 5/5] test(router): annotate the budget filter return values The two helpers return the deployments the filter kept and the debug string. Name that as a sequence of model-list entries. Co-authored-by: abeed09 --- tests/unit/router_strategy/test_budget_limiter.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/router_strategy/test_budget_limiter.py b/tests/unit/router_strategy/test_budget_limiter.py index 336e1cd96d9..169d4b58ca3 100644 --- a/tests/unit/router_strategy/test_budget_limiter.py +++ b/tests/unit/router_strategy/test_budget_limiter.py @@ -174,7 +174,7 @@ def _apply_budget_filter( request_tags: Sequence[str], spend_key: str, spend: float, -): +) -> tuple[Sequence[DeploymentTypedDict], str]: limiter: Final = RouterBudgetLimiting(dual_cache=DualCache(), provider_budget_config=None) limiter.provider_budget_config = provider_budget limiter.deployment_budget_config = deployment_budget @@ -190,7 +190,7 @@ def _apply_budget_filter( ) -def _filter_one(scope: _Scope, config: BudgetConfig, spend: float): +def _filter_one(scope: _Scope, config: BudgetConfig, spend: float) -> tuple[Sequence[DeploymentTypedDict], str]: match scope: case "provider": return _apply_budget_filter(