fix(proxy): price the model every alias hop dispatches to in the zero-cost budget check

litellm/proxy/auth/auth_checks.py

@@ -89,7 +89,7 @@ from litellm.proxy.common_utils.http_parsing_utils import (
     _safe_get_request_headers,
     _safe_get_request_query_params,
 )
-from litellm.proxy.common_utils.model_listing_utils import alias_map
+from litellm.proxy.common_utils.model_listing_utils import alias_map, alias_target, caller_alias_maps
 from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
 from litellm.proxy.common_utils.user_api_key_cache import (
     END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL,
@@ -564,15 +564,33 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None
     return True
 
 
-def _dispatched_model_name(model_name: str, valid_token: UserAPIKeyAuth) -> str:
-    after_team_alias: Final = alias_map(valid_token.team_model_aliases).get(model_name, model_name)
-    return alias_map(valid_token.aliases).get(after_team_alias, after_team_alias)
+def _dispatched_model_name(
+    model_name: str, valid_token: UserAPIKeyAuth, router_settings_aliases: object, router_settings_rewrite_pending: bool
+) -> str:
+    requested: Final = (
+        resolve_model_group_alias(router_settings_aliases, model_name) or model_name
+        if router_settings_rewrite_pending
+        else model_name
+    )
+    caller_aliases: Final = caller_alias_maps(
+        valid_token.aliases, valid_token.team_model_aliases, valid_token.team_id, None
+    )
+    aliased: Final = alias_target(requested, caller_aliases) or requested
+    return resolve_model_group_alias(router_settings_aliases, aliased) or aliased
 
 
 def is_dispatched_model_cost_zero(
-    model: str | list[str] | None, llm_router: Router | None, valid_token: UserAPIKeyAuth
+    model: str | list[str] | None,
+    llm_router: Router | None,
+    valid_token: UserAPIKeyAuth,
+    router_settings_aliases: object,
+    router_settings_rewrite_pending: bool,
 ) -> bool:
-    dispatched_model: Final = _dispatched_model_name(model, valid_token) if isinstance(model, str) else model
+    dispatched_model: Final = (
+        _dispatched_model_name(model, valid_token, router_settings_aliases, router_settings_rewrite_pending)
+        if isinstance(model, str)
+        else model
+    )
     return _is_model_cost_zero(model=dispatched_model, llm_router=llm_router)
 
 

litellm/proxy/auth/user_api_key_auth.py

@@ -1795,8 +1795,8 @@ async def _user_api_key_auth_builder(
                     )
                     skip_budget_checks = False
                     if model is not None and llm_router is not None:
-                        skip_budget_checks = is_dispatched_model_cost_zero(
-                            model=model, llm_router=llm_router, valid_token=valid_token
+                        skip_budget_checks = await _is_dispatched_model_cost_zero(
+                            model=model, llm_router=llm_router, valid_token=valid_token, request=request
                         )
                         if skip_budget_checks:
                             verbose_proxy_logger.info("Skipping all budget checks for zero-cost model: %s", model)
@@ -2237,8 +2237,8 @@ async def _user_api_key_auth_builder(
             )
             skip_budget_checks = False
             if model is not None and llm_router is not None:
-                skip_budget_checks = is_dispatched_model_cost_zero(
-                    model=model, llm_router=llm_router, valid_token=valid_token
+                skip_budget_checks = await _is_dispatched_model_cost_zero(
+                    model=model, llm_router=llm_router, valid_token=valid_token, request=request
                 )
                 if skip_budget_checks:
                     verbose_proxy_logger.info("Skipping all budget checks for zero-cost model: %s", model)
@@ -2960,7 +2960,7 @@ async def _run_centralized_common_checks(
             keep_token_limits=user_custom_auth is not None,
         )
 
-    skip_budget_checks: Final = _should_skip_budget_checks(
+    skip_budget_checks: Final = await _should_skip_budget_checks(
         request_data=request_data,
         route=route,
         request=request,
@@ -3131,7 +3131,7 @@ async def _reserve_budget_after_common_checks(
         request.state.budget_reservation = reservation  # rebind-ok: read by the release middleware
 
 
-def _should_skip_budget_checks(
+async def _should_skip_budget_checks(
     request_data: dict,
     route: str,
     request: Request | None,
@@ -3146,10 +3146,30 @@ def _should_skip_budget_checks(
         team_id=valid_token.team_id,
     )
     if model is not None and llm_router is not None:
-        return is_dispatched_model_cost_zero(model=model, llm_router=llm_router, valid_token=valid_token)
+        return await _is_dispatched_model_cost_zero(
+            model=model, llm_router=llm_router, valid_token=valid_token, request=request
+        )
     return False
 
 
+async def _is_dispatched_model_cost_zero(
+    model: str | list[str], llm_router: litellm.Router, valid_token: UserAPIKeyAuth, request: Request | None
+) -> bool:
+    from litellm.proxy.proxy_server import prisma_client, proxy_config, proxy_logging_obj
+
+    settings: Final = await proxy_config.get_hierarchical_router_settings(
+        user_api_key_dict=valid_token, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj
+    )
+    return is_dispatched_model_cost_zero(
+        model=model,
+        llm_router=llm_router,
+        valid_token=valid_token,
+        router_settings_aliases=settings.get("model_group_alias") if isinstance(settings, Mapping) else None,
+        router_settings_rewrite_pending=request is not None
+        and request.scope.get(MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY) is not True,
+    )
+
+
 def _resolve_request_principal(request: Request, valid_token: UserAPIKeyAuth) -> Principal:
     """Project the resolved identity into one per-request Principal, off the key
     object the builder already fetched, and stamp the request network context
@@ -3715,7 +3735,9 @@ async def _run_post_custom_auth_checks(
     # every budget check for these; this path did not, so the same request could
     # be refused under custom auth and served under the other two.
     skip_budget_checks: Final = (
-        is_dispatched_model_cost_zero(model=current_model, llm_router=llm_router, valid_token=valid_token)
+        await _is_dispatched_model_cost_zero(
+            model=current_model, llm_router=llm_router, valid_token=valid_token, request=request
+        )
         if current_model is not None and llm_router is not None
         else False
     )

tests/unit/proxy/auth/test_auth_checks.py

@@ -1598,4 +1598,53 @@ def test_zero_cost_check_prices_the_model_the_key_and_team_aliases_dispatch_to(
 ) -> None:
     router: Final = _alias_router(_aliases_to("free-model"))
     token: Final = UserAPIKeyAuth(aliases=key_aliases or {}, team_model_aliases=team_model_aliases)
-    assert is_dispatched_model_cost_zero(model=model, llm_router=router, valid_token=token) is expected
+    assert (
+        is_dispatched_model_cost_zero(
+            model=model,
+            llm_router=router,
+            valid_token=token,
+            router_settings_aliases=None,
+            router_settings_rewrite_pending=False,
+        )
+        is expected
+    )
+
+
+@pytest.mark.parametrize(
+    "key_aliases, router_settings_aliases, model_alias_map, rewrite_pending, model, expected",
+    [
+        ({"my-free": "free-model"}, {"free-model": "paid-model"}, {}, True, "my-free", False),
+        ({"my-free": "free-model"}, {}, {"free-model": "paid-model"}, True, "my-free", False),
+        ({}, {"free-model": "paid-model"}, {}, True, "free-model", False),
+        ({}, {"router-free": "free-model"}, {}, True, "router-free", True),
+        ({}, {"router-free": {"model": "visible", "hidden": True}}, {}, True, "router-free", True),
+        ({"global-free": "free-model"}, {}, {"my-free": "global-free"}, True, "my-free", True),
+        ({"global-free": "paid-model"}, {}, {"my-free": "global-free"}, True, "my-free", False),
+        ({}, {}, {"my-free": "visible"}, True, "my-free", True),
+        ({"free-model": "paid-model"}, {"router-free": "free-model"}, {}, True, "router-free", False),
+        ({}, {"router-free": "free-model", "free-model": "paid-model"}, {}, True, "router-free", False),
+        ({}, {"router-free": "free-model", "free-model": "paid-model"}, {}, False, "router-free", True),
+    ],
+)
+def test_zero_cost_check_prices_the_model_after_every_alias_hop_dispatch_applies(
+    monkeypatch: pytest.MonkeyPatch,
+    key_aliases: dict[str, str],
+    router_settings_aliases: dict[str, str | dict[str, str | bool]],
+    model_alias_map: dict[str, str],
+    rewrite_pending: bool,
+    model: str,
+    expected: bool,
+) -> None:
+    monkeypatch.setattr(litellm, "model_alias_map", model_alias_map)
+    router: Final = _alias_router(_aliases_to("free-model"))
+    token: Final = UserAPIKeyAuth(aliases=key_aliases)
+    assert (
+        is_dispatched_model_cost_zero(
+            model=model,
+            llm_router=router,
+            valid_token=token,
+            router_settings_aliases=router_settings_aliases,
+            router_settings_rewrite_pending=rewrite_pending,
+        )
+        is expected
+    )

tests/unit/proxy/auth/test_user_api_key_auth.py

@@ -1796,14 +1796,28 @@ def test_mapped_key_jwt_falls_through_to_the_shared_user_budget_attach():
     )
 
 
+@pytest.mark.asyncio
 @pytest.mark.parametrize(
-    "key_aliases, expected",
+    "model, key_aliases, key_router_settings, router_settings_rewritten, expected",
     [
-        ({}, True),
-        ({"free-alias": "paid-model"}, False),
+        ("free-alias", {}, None, None, True),
+        ("free-alias", {"free-alias": "paid-model"}, None, None, False),
+        ("my-alias", {"my-alias": "free-alias"}, None, None, True),
+        ("my-alias", {"my-alias": "free-alias"}, {"model_group_alias": {"free-alias": "paid-model"}}, None, False),
+        ("rs-alias", {"free-model": "paid-model"}, {"model_group_alias": {"rs-alias": "free-model"}}, False, False),
+        ("rs-alias", {"free-model": "paid-model"}, {"model_group_alias": {"rs-alias": "free-model"}}, True, True),
     ],
 )
-def test_budget_skip_judges_the_model_a_key_alias_dispatches_to(key_aliases: dict[str, str], expected: bool) -> None:
+async def test_budget_skip_judges_the_model_the_key_aliases_dispatch_to(
+    model: str,
+    key_aliases: dict[str, str],
+    key_router_settings: dict[str, dict[str, str]] | None,
+    router_settings_rewritten: bool | None,
+    expected: bool,
+) -> None:
+    from starlette.requests import Request
+
+    from litellm.constants import MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY
     from litellm.proxy.auth.user_api_key_auth import _should_skip_budget_checks
     from litellm.router import Router
 
@@ -1830,11 +1844,24 @@ def test_budget_skip_judges_the_model_a_key_alias_dispatches_to(key_aliases: dic
         ],
         model_group_alias={"free-alias": "free-model"},
     )
-    skipped = _should_skip_budget_checks(
-        request_data={"model": "free-alias"},
+    skipped = await _should_skip_budget_checks(
+        request_data={"model": model},
         route="/chat/completions",
-        request=None,
+        request=(
+            None
+            if router_settings_rewritten is None
+            else Request(
+                {
+                    "type": "http",
+                    "method": "POST",
+                    "path": "/chat/completions",
+                    "headers": [],
+                    "query_string": b"",
+                    MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY: router_settings_rewritten,
+                }
+            )
+        ),
         llm_router=router,
-        valid_token=UserAPIKeyAuth(aliases=key_aliases),
+        valid_token=UserAPIKeyAuth(aliases=key_aliases, router_settings=key_router_settings),
     )
     assert skipped is expected