diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 2944587c3e1..460b25ddac7 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -148,6 +148,7 @@ from litellm.repositories.user_repository import UserRepository from litellm.router import Router from litellm.types.proxy.auth.auth_checks import UserNotFoundError from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget +from litellm.types.router import DeploymentTypedDict from litellm.utils import get_utc_datetime from .auth_checks_organization import ( @@ -527,7 +528,8 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None # not from defaulted sparse auto-registration entries. # See: https://github.com/BerriAI/litellm/issues/24770 safe_name = str(model_name).replace("\n", "").replace("\r", "") - if not _is_cost_explicitly_configured(model_name, llm_router): + served_deployments = _deployments_served_for(model_name, llm_router) + if not _is_cost_explicitly_configured(served_deployments): verbose_proxy_logger.debug( "Model %s has zero cost but no explicit cost " "configuration in model_cost entry — treating as unknown " @@ -538,7 +540,16 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None zero_cost_cache[model_name] = False return False - if _has_ptu_flat_cost(model_name, llm_router): + if _has_explicit_positive_cost(served_deployments): + verbose_proxy_logger.debug( + "Model %s routes to a deployment with an explicit positive per-token price (enforce budget)", + safe_name, + ) + if zero_cost_cache is not None: + zero_cost_cache[model_name] = False + return False + + if _has_ptu_flat_cost(served_deployments): verbose_proxy_logger.debug( "Model %s prices reserved PTU capacity as a flat cost, so its zero per-token " "rate is not a free model (enforce budget)", @@ -570,40 +581,56 @@ _NO_MODEL_INFO: Final[Mapping[str, object]] = MappingProxyType({}) _TEAM_GRANT_RELATIONS: Final[Mapping[str, object]] = MappingProxyType({"litellm_model_table": True}) -def _has_ptu_flat_cost(model: str, llm_router: "Router") -> bool: - """Whether any deployment in the model group bills reserved PTU capacity as a flat cost. +def _deployments_served_for(model: str, llm_router: "Router") -> Sequence[DeploymentTypedDict]: + """ + The deployments a request to ``model`` routes to, for the gates that waive budget checks. + + The router resolves a ``model_group_alias`` name to its target exactly once and serves that + group, never a real deployment that shares the alias's name. ``Router.get_model_list()`` on the + routed group reads those deployments unless the group is itself an alias key, where it would + follow the alias a second hop the router never takes; then the deployments named after the + group count, or the wildcard route serving the group when none carries its name, the same + fallback ``get_model_list()`` takes. + """ + routed_group: Final = llm_router.routable_model_group(model) + if llm_router.routable_model_group(routed_group) == routed_group: + return llm_router.get_model_list(model_name=routed_group) or () + named: Final = llm_router.get_model_list_from_model_alias(model_name=model) + if named: + return named + return tuple( + DeploymentTypedDict(**deployment) + for deployment in llm_router.pattern_router.get_deployments_by_pattern(model=routed_group) + ) + + +def _has_ptu_flat_cost(deployments: Sequence[DeploymentTypedDict]) -> bool: + """Whether any of the deployments bills reserved PTU capacity as a flat cost. Such a deployment carries an explicit zero per-token price so the flat cost is not charged twice, which otherwise reads here as a free model and waives every budget check for it. - - Resolved through ``Router.get_model_list()``, which includes ``model_group_alias``, because - this runs after the explicit-cost gate: resolving that gate alone would let an aliased PTU - group through as free. """ - for deployment in llm_router.get_model_list(model_name=model) or (): + for deployment in deployments: model_info = deployment.get("model_info") or _NO_MODEL_INFO if model_info.get("ptu_count") is not None and model_info.get("cost_per_ptu_per_hour") is not None: return True return False -def _is_cost_explicitly_configured(model: str, llm_router: "Router") -> bool: +def _is_cost_explicitly_configured(deployments: Sequence[DeploymentTypedDict]) -> bool: """ - Check if any deployment in the model group has cost fields explicitly - set in its litellm.model_cost entry. + Check if any of the deployments has cost fields explicitly set in its + litellm.model_cost entry. When Router._create_deployment() registers a model not in the global cost map, it creates a sparse entry like {"id": ""} with no cost fields. _get_model_info_helper() then defaults missing costs to 0. This function detects that scenario by checking the raw model_cost entry. - The group is resolved through ``Router.get_model_list()``, the same resolution - ``get_model_group_info()`` applies when the caller reads the cost a few lines earlier, so the - two lookups cannot disagree, including for names defined in ``Router.model_group_alias``. It also reaches a deployment that prices itself through its ``model_info`` block, whose entry lands in the cost map under the deployment id. """ - for deployment in llm_router.get_model_list(model_name=model) or (): + for deployment in deployments: model_id = (deployment.get("model_info") or _EMPTY_COST_ENTRY).get("id") if model_id is None: continue @@ -613,6 +640,26 @@ def _is_cost_explicitly_configured(model: str, llm_router: "Router") -> bool: return False +def _has_explicit_positive_cost(deployments: Sequence[DeploymentTypedDict]) -> bool: + """Whether any of the deployments carries an explicit positive per-token price. + + The group's price is read from the deployments named after the routed group and from + that group's own alias target, while an alias chain is served from the deployments named + after its routed group or from a wildcard route; a priced route among those serves the + request at its own rate whatever the group's price reads. + """ + for deployment in deployments: + model_id = (deployment.get("model_info") or _EMPTY_COST_ENTRY).get("id") + if model_id is None: + continue + raw_entry = litellm.model_cost.get(model_id, _EMPTY_COST_ENTRY) + if _is_positive_cost(raw_entry.get("input_cost_per_token")) or _is_positive_cost( + raw_entry.get("output_cost_per_token") + ): + return True + return False + + _EMPTY_COST_ENTRY: Final[Mapping[str, object]] = MappingProxyType({}) @@ -672,7 +719,7 @@ def model_has_no_cost_mapping(model: str | None, llm_router: Router | None) -> b if _model_group_has_pricing(model=model, llm_router=llm_router): return False - return not _is_cost_explicitly_configured(model=model, llm_router=llm_router) + return not _is_cost_explicitly_configured(llm_router.get_model_list(model_name=model) or ()) def _unpriced_models_in_request(model: str | list[str] | None, llm_router: Router | None) -> tuple[str, ...]: diff --git a/tests/integration/authorization/test_hidden_alias_budget_bypass.py b/tests/integration/authorization/test_hidden_alias_budget_bypass.py index 46d0852ca8a..eabdb01602e 100644 --- a/tests/integration/authorization/test_hidden_alias_budget_bypass.py +++ b/tests/integration/authorization/test_hidden_alias_budget_bypass.py @@ -1,20 +1,35 @@ +import base64 import json import math +import signal +import time import uuid -from collections.abc import Iterator, Mapping +from collections.abc import Callable, Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from hashlib import sha256 from pathlib import Path from typing import Final +import anthropic import httpx +import jwt +import openai +import psutil import pytest import yaml +from cryptography.hazmat.primitives.asymmetric import rsa from integration._support.anthropic_thinking import JSON_OBJECT -from integration._support.client import Gateway, object_value -from integration._support.process import owned_proxy +from integration._support.client import Gateway, Scenario, eventually, object_value +from integration._support.database import write_rows +from integration._support.process import group_members, owned_proxy, owned_proxy_process +from integration._support.wire import Reply, Request, wire_server from integration.authorization._hidden_alias_budget import ( BUDGET, BUDGET_EXCEEDED, CHAT_REPLY, + GATEWAY_BURST, + PEER_BURST, PROXY_BUDGET_USER, RESPONSES_REPLY, AliasRig, @@ -29,9 +44,13 @@ from integration.authorization._hidden_alias_budget import ( error_type, exhausted_key, fresh_chat, + fresh_message, + fresh_post, + fresh_response, hidden, install_aliases, landed, + landed_all_once, landed_once, openai_client, remove_aliases, @@ -46,6 +65,11 @@ from pydantic import JsonValue pytestmark: Final = pytest.mark.timeout(240) +# $0 comes from litellm's own cost-map row (model_prices_and_context_window.json, added 2025-10-14 in 900a57f5b1f): +# Lemonade is a local inference server, so there is no vendor price for the row to drift from. A row that gains a +# price is stale, not a regression: pick another chat row the map prices at $0 and serves over the OpenAI wire +UNPRICED_FREE_PROVIDER_MODEL: Final = "lemonade/Gemma-3-4b-it-GGUF" + @pytest.fixture(scope="module") def rig() -> Iterator[AliasRig]: @@ -506,6 +530,222 @@ def test_hidden_alias_to_a_group_priced_by_the_cost_map_stays_budgeted(rig: Alia assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 0 +def _reported_per_token_price(rig: AliasRig, group: str) -> tuple[float, float]: + entries: Final = rig.gateway.get("/model/info")["data"] + assert isinstance(entries, list), entries + infos: Final = tuple( + object_value(object_value(entry)["model_info"]) + for entry in entries + if object_value(entry)["model_name"] == group + ) + assert len(infos) == 1, infos + return float(str(infos[0]["input_cost_per_token"])), float(str(infos[0]["output_cost_per_token"])) + + +def _serving_deployment(candidate: Gateway, model: str) -> str | None: + response: Final = fresh_chat(candidate, model, candidate.key, "route-" + uuid.uuid4().hex) + return response.headers.get("x-litellm-model-id") if response.status_code == 200 else None + + +def _routes_alike(candidate: Gateway, name: str, target: str) -> bool: + served: Final = _serving_deployment(candidate, name) + return served is not None and served == _serving_deployment(candidate, target) + + +def _await_alias(candidate: Gateway, name: str, target: str) -> None: + """Wait until ``candidate`` serves ``name`` from ``target``'s deployment. + + ``/router/settings`` answers from the stored config before a replica's next reload applies it, and a + shadowing alias's name answers 200 either way, so the serving deployment is the only signal on the wire. + """ + eventually(lambda: _routes_alike(candidate, name, target), bool, seconds=90) + + +def _shadowing_alias( + rig: AliasRig, scenario: Scenario, target: str, candidates: tuple[Gateway, ...], **shadow: JsonValue +) -> str: + """A hidden alias whose name is also a real group, live on every candidate so the name routes to ``target``.""" + name: Final = scenario.model(**shadow) + install_aliases(rig.gateway, {name: hidden(target)}) + scenario.cleanups.callback(remove_aliases, rig.gateway, frozenset({name})) + for candidate in candidates: + _await_alias(candidate, name, target) + return name + + +def _unpriced_free_target(rig: AliasRig, scenario: Scenario, provider_model: str = UNPRICED_FREE_PROVIDER_MODEL) -> str: + target: Final = scenario.model(model=provider_model) + assert _reported_per_token_price(rig, target) == (0.0, 0.0), ( + f"{provider_model} must price at $0 in the cost map with no explicit price on the deployment, " + "or the shadowed alias is refused by the price gate and never reaches the one under test" + ) + return target + + +def _shadow_over_unpriced_target(rig: AliasRig, scenario: Scenario) -> tuple[str, str]: + target: Final = _unpriced_free_target(rig, scenario) + shadow: Final = _shadowing_alias( + rig, scenario, target, (rig.gateway, rig.peer), input_cost_per_token=0, output_cost_per_token=0 + ) + return target, shadow + + +def _refusal(response: httpx.Response) -> tuple[int, str]: + return response.status_code, error_type(response) + + +def _assert_refused_like_target(rig: AliasRig, candidate: Gateway, key: str, shadow: str, target: str) -> None: + for route, send in (("chat", fresh_chat), ("responses", fresh_response), ("messages", fresh_message)): + shadow_marker: Final = f"shadow-{route}-" + uuid.uuid4().hex + target_marker: Final = f"target-{route}-" + uuid.uuid4().hex + via_shadow: Final = send(candidate, shadow, key, shadow_marker) + by_name: Final = send(candidate, target, key, target_marker) + assert _refusal(by_name) == (BUDGET_EXCEEDED, "budget_exceeded"), by_name.text + assert via_shadow.status_code == by_name.status_code, via_shadow.text + assert error_type(via_shadow) == error_type(by_name), via_shadow.text + observed: Final = upstream_requests(rig.gateway.upstream_url) + assert upstream_hits(observed, shadow_marker) == 0 + assert upstream_hits(observed, target_marker) == 0 + + +def test_hidden_alias_shadowing_an_explicitly_free_group_is_refused_like_its_unpriced_target(rig: AliasRig) -> None: + with rig.gateway.scenario() as scenario: + target, shadow = _shadow_over_unpriced_target(rig, scenario) + key: Final = exhausted_key(rig, scenario) + settle_chat(rig, target, key, BUDGET_EXCEEDED) + _assert_refused_like_target(rig, rig.gateway, key, shadow, target) + _assert_refused_like_target(rig, rig.peer, key, shadow, target) + + +def _assert_refused_on_every_route(rig: AliasRig, candidate: Gateway, key: str, model: str) -> None: + for route, send in (("chat", fresh_chat), ("responses", fresh_response), ("messages", fresh_message)): + marker: Final = f"chain-{route}-" + uuid.uuid4().hex + refused: Final = send(candidate, model, key, marker) + assert refused.status_code == BUDGET_EXCEEDED, refused.text + assert error_type(refused) == "budget_exceeded", refused.text + assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 0 + + +def test_hidden_alias_shadow_chain_is_judged_by_the_deployment_it_is_served_from(rig: AliasRig) -> None: + """A shadowing alias routes one hop to its target's own deployment, even when that target is an alias too.""" + with rig.gateway.scenario() as scenario: + middle, shadow = _shadow_over_unpriced_target(rig, scenario) + served_from: Final = _serving_deployment(rig.gateway, middle) + install_aliases(rig.gateway, {middle: hidden(rig.free)}) + scenario.cleanups.callback(remove_aliases, rig.gateway, frozenset({middle})) + for candidate in (rig.gateway, rig.peer): + _await_alias(candidate, middle, rig.free) + assert _serving_deployment(candidate, shadow) == served_from + key: Final = exhausted_key(rig, scenario) + for candidate in (rig.gateway, rig.peer): + marker: Final = "chain-target-" + uuid.uuid4().hex + by_name: Final = fresh_chat(candidate, middle, key, marker) + assert by_name.status_code == 200, by_name.text + assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 1 + _assert_refused_on_every_route(rig, candidate, key, shadow) + + +def _await_route(candidate: Gateway, name: str, route: str) -> None: + eventually(lambda: _serving_deployment(candidate, name), lambda served: served == route, seconds=90) + + +def _deployment_named(rig: AliasRig, scenario: Scenario, name: str, **parameters: JsonValue) -> str: + created: Final = rig.gateway.post( + "/model/new", + { + "model_name": name, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "integration-provider-key", + "api_base": f"{rig.gateway.upstream_url}/v1", + **parameters, + }, + "model_info": {}, + }, + ) + identity: Final = str(object_value(created["model_info"])["id"]) + scenario.cleanups.callback(scenario.delete_model, identity) + return identity + + +def _priced_wildcard_route(rig: AliasRig, scenario: Scenario, prefix: str) -> str: + """A wildcard deployment serving every ``/...`` name at an explicit positive per-token price.""" + return _deployment_named( + rig, scenario, f"{prefix}/*", model="openai/*", input_cost_per_token=0.00001, output_cost_per_token=0.00002 + ) + + +@pytest.mark.parametrize("entry_on_pattern", [False, True], ids=["entry_off_pattern", "entry_on_pattern"]) +def test_hidden_alias_chain_served_by_a_priced_wildcard_route_is_refused(rig: AliasRig, entry_on_pattern: bool) -> None: + """A chain the router serves from a priced wildcard route is budgeted at that route's price. + + The chain's middle name is only an alias key to the free group, so the group's price reads $0 + through it while the request is served by the wildcard route matching the middle name. + """ + prefix: Final = "wild" + uuid.uuid4().hex[:8] + middle: Final = f"{prefix}/gpt-4o-mini" + entry: Final = (f"{prefix}/" if entry_on_pattern else "chain-entry-") + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + route: Final = _priced_wildcard_route(rig, scenario, prefix) + install_aliases(rig.gateway, {entry: hidden(middle), middle: hidden(rig.free)}) + scenario.cleanups.callback(remove_aliases, rig.gateway, frozenset({entry, middle})) + for candidate in (rig.gateway, rig.peer): + _await_route(candidate, entry, route) + _await_alias(candidate, middle, rig.free) + key: Final = exhausted_key(rig, scenario) + settle_chat(rig, rig.paid, key, BUDGET_EXCEEDED) + for candidate in (rig.gateway, rig.peer): + _assert_refused_on_every_route(rig, candidate, key, entry) + + +def _ptu_shadow_config(rig: AliasRig, directory: Path, name: str) -> Path: + base: Final = _base_config()["model_list"] + assert isinstance(base, list), base + deployment: Final[JsonValue] = { + "model_name": name, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "integration-provider-key", + "api_base": f"{rig.gateway.upstream_url}/v1", + "input_cost_per_token": 0, + "output_cost_per_token": 0, + }, + "model_info": {"ptu_count": 100, "cost_per_ptu_per_hour": 2.0}, + } + return _own_config(directory, "ptu-shadow.yaml", "model_list", [*base, deployment]) + + +@pytest.mark.timeout(480) +def test_hidden_alias_shadowing_a_ptu_group_is_served_through_its_free_target(rig: AliasRig, tmp_path: Path) -> None: + shadow: Final = "ptu-shadow-" + uuid.uuid4().hex + marker: Final = "ptu-shadow-served-" + uuid.uuid4().hex + config: Final = _ptu_shadow_config(rig, tmp_path, shadow) + with rig.gateway.scenario() as scenario: + install_aliases(rig.gateway, {shadow: hidden(rig.free)}) + scenario.cleanups.callback(remove_aliases, rig.gateway, frozenset({shadow})) + key: Final = exhausted_key(rig, scenario) + with owned_proxy(rig.gateway, tmp_path, {}, config=config) as candidate: + _await_alias(candidate, shadow, rig.free) + settle_candidate(candidate, rig.paid, key, BUDGET_EXCEEDED, seconds=120) + served: Final = fresh_chat(candidate, shadow, key, marker) + assert served.status_code == 200, served.text + assert served.json()["choices"][0]["message"]["content"] == CHAT_REPLY, served.text + assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 1 + assert_free_row(landed_once(key, marker), shadow) + + +def test_hidden_alias_shadowing_an_explicitly_free_group_to_a_paid_target_stays_budgeted(rig: AliasRig) -> None: + marker: Final = "paid-shadow-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + shadow: Final = _shadowing_alias( + rig, scenario, rig.paid, (rig.gateway, rig.peer), input_cost_per_token=0, output_cost_per_token=0 + ) + key: Final = exhausted_key(rig, scenario) + refused: Final = fresh_chat(rig.gateway, shadow, key, marker) + assert _refusal(refused) == (BUDGET_EXCEEDED, "budget_exceeded"), refused.text + assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 0 + + def test_hidden_alias_to_a_missing_group_is_refused_and_the_proxy_stays_healthy(rig: AliasRig) -> None: with rig.gateway.scenario() as scenario: key: Final = exhausted_key(rig, scenario) @@ -575,3 +815,578 @@ def _assert_hidden_aliases_unlisted(rig: AliasRig, candidate: Gateway) -> None: def test_hidden_alias_stays_out_of_model_listings(rig: AliasRig) -> None: _assert_hidden_aliases_unlisted(rig, rig.gateway) _assert_hidden_aliases_unlisted(rig, rig.peer) + + +def _deployment_id(rig: AliasRig, name: str) -> str: + entries: Final = rig.gateway.get("/model/info")["data"] + assert isinstance(entries, list), entries + identities: Final = tuple( + str(object_value(object_value(entry)["model_info"])["id"]) + for entry in entries + if object_value(entry)["model_name"] == name + ) + assert len(identities) == 1, identities + return identities[0] + + +def _assert_budget_refusal(status_code: int, body: str) -> None: + assert status_code == BUDGET_EXCEEDED, body + assert "budget_exceeded" in body, body + + +def _assert_no_served_row(key: str, marker: str) -> None: + assert all(row["status"] != "success" for row in landed(key, marker)), landed(key, marker) + + +def test_exhausted_key_is_refused_on_shadowing_alias_through_streamed_openai_chat(rig: AliasRig) -> None: + marker: Final = "shadow-chat-stream-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + _, shadow = _shadow_over_unpriced_target(rig, scenario) + key: Final = exhausted_key(rig, scenario) + with openai_client(rig.gateway, key) as client, pytest.raises(openai.APIStatusError) as refused: + client.chat.completions.create( + model=shadow, + messages=[{"role": "user", "content": marker}], + stream=True, + extra_headers=spend_marker(marker), + ) + _assert_budget_refusal(refused.value.status_code, str(refused.value.body)) + assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 0 + + +async def test_exhausted_key_is_refused_on_shadowing_alias_through_streamed_openai_responses(rig: AliasRig) -> None: + marker: Final = "shadow-responses-stream-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + _, shadow = _shadow_over_unpriced_target(rig, scenario) + key: Final = exhausted_key(rig, scenario) + async with async_openai_client(rig.gateway, key) as client: + with pytest.raises(openai.APIStatusError) as refused: + await client.responses.create( + model=shadow, input=marker, stream=True, extra_headers=spend_marker(marker) + ) + _assert_budget_refusal(refused.value.status_code, str(refused.value.body)) + assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 0 + + +async def test_exhausted_key_is_refused_on_shadowing_alias_through_streamed_anthropic_messages(rig: AliasRig) -> None: + marker: Final = "shadow-messages-stream-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + _, shadow = _shadow_over_unpriced_target(rig, scenario) + key: Final = exhausted_key(rig, scenario) + async with async_anthropic_client(rig.gateway, key) as client: + with pytest.raises(anthropic.APIStatusError) as refused: + async with client.messages.stream( + model=shadow, + max_tokens=16, + messages=[{"role": "user", "content": marker}], + extra_headers=spend_marker(marker), + ) as stream: + await stream.get_final_message() + _assert_budget_refusal(refused.value.status_code, str(refused.value.body)) + assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 0 + + +def _await_served_within(candidate: Gateway, name: str, deployments: frozenset[str]) -> None: + eventually(lambda: _serving_deployment(candidate, name), lambda served: served in deployments, seconds=90) + + +def _assert_served_by( + rig: AliasRig, candidate: Gateway, name: str, key: str, deployments: frozenset[str], prefix: str +) -> str: + marker: Final = f"{prefix}-" + uuid.uuid4().hex + served: Final = fresh_chat(candidate, name, key, marker) + assert served.status_code == 200, served.text + assert served.headers["x-litellm-model-id"] in deployments, dict(served.headers) + assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 1 + response_id: Final = str(JSON_OBJECT.validate_json(served.content)["id"]) + row: Final = landed_once(key, marker) + assert row["request_id"] == response_id, row + assert_free_row(row, name) + return response_id + + +def test_key_with_headroom_is_served_by_the_target_through_a_shadowing_alias(rig: AliasRig) -> None: + with rig.gateway.scenario() as scenario: + target, shadow = _shadow_over_unpriced_target(rig, scenario) + key: Final = scenario.key(max_budget=BUDGET) + for candidate in (rig.gateway, rig.peer): + _assert_served_by(rig, candidate, shadow, key, frozenset({_deployment_id(rig, target)}), "headroom-shadow") + + +def test_exhausted_key_reaches_hidden_alias_to_a_name_served_by_a_free_wildcard_route(rig: AliasRig) -> None: + prefix: Final = "freewild" + uuid.uuid4().hex[:8] + entry: Final = "free-wildcard-entry-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + route: Final = _deployment_named( + rig, scenario, f"{prefix}/*", model="openai/*", input_cost_per_token=0, output_cost_per_token=0 + ) + install_aliases(rig.gateway, {entry: hidden(f"{prefix}/gpt-4o-mini")}) + scenario.cleanups.callback(remove_aliases, rig.gateway, frozenset({entry})) + for candidate in (rig.gateway, rig.peer): + _await_route(candidate, entry, route) + key: Final = exhausted_key(rig, scenario) + for candidate in (rig.gateway, rig.peer): + _assert_served_by(rig, candidate, entry, key, frozenset({route}), "free-wildcard") + + +def test_exhausted_key_reaches_a_group_with_one_explicitly_free_deployment(rig: AliasRig) -> None: + group: Final = "mixed-free-" + uuid.uuid4().hex + entry: Final = "mixed-free-entry-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + members: Final = frozenset( + { + _deployment_named(rig, scenario, group, input_cost_per_token=0, output_cost_per_token=0), + _deployment_named(rig, scenario, group, model=UNPRICED_FREE_PROVIDER_MODEL), + } + ) + install_aliases(rig.gateway, {entry: hidden(group)}) + scenario.cleanups.callback(remove_aliases, rig.gateway, frozenset({entry})) + for candidate in (rig.gateway, rig.peer): + _await_served_within(candidate, entry, members) + key: Final = exhausted_key(rig, scenario) + for candidate in (rig.gateway, rig.peer): + for name in (group, entry): + _assert_served_by(rig, candidate, name, key, members, "mixed") + + +def _script_target(rig: AliasRig, statuses: tuple[int, ...]) -> None: + provider_model: Final = UNPRICED_FREE_PROVIDER_MODEL.split("/", 1)[1] + scripted: Final = httpx.post( + f"{rig.gateway.upstream_url}/__scripts/{provider_model}", + json={"statuses": list(statuses)}, + timeout=15, + trust_env=False, + ) + assert scripted.status_code == 200, scripted.text + + +def _clear_target_script(rig: AliasRig) -> None: + provider_model: Final = UNPRICED_FREE_PROVIDER_MODEL.split("/", 1)[1] + cleared: Final = httpx.delete(f"{rig.gateway.upstream_url}/__scripts/{provider_model}", timeout=15, trust_env=False) + assert cleared.status_code in (200, 404), cleared.text + + +def _assert_provider_failure(response: httpx.Response) -> None: + assert response.status_code == 500, response.text + assert "Controlled provider failure" in response.text, response.text + assert "budget" not in response.text.lower(), response.text + + +def test_provider_failure_behind_a_shadowing_alias_reaches_the_caller(rig: AliasRig) -> None: + failed_marker: Final = "shadow-provider-failure-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + target, shadow = _shadow_over_unpriced_target(rig, scenario) + healthy: Final = scenario.key(max_budget=BUDGET) + exhausted: Final = exhausted_key(rig, scenario) + settle_chat(rig, target, exhausted, BUDGET_EXCEEDED) + _assert_refused_like_target(rig, rig.gateway, exhausted, shadow, target) + scenario.cleanups.callback(_clear_target_script, rig) + _script_target(rig, (500,)) + _assert_provider_failure(fresh_chat(rig.gateway, shadow, healthy, failed_marker)) + assert upstream_hits(upstream_requests(rig.gateway.upstream_url), failed_marker) == 1 + _clear_target_script(rig) + _assert_served_by( + rig, rig.gateway, shadow, healthy, frozenset({_deployment_id(rig, target)}), "shadow-provider-recovered" + ) + _assert_no_served_row(healthy, failed_marker) + _assert_refused_like_target(rig, rig.gateway, exhausted, shadow, target) + + +def test_unauthenticated_request_to_a_shadowing_alias_is_rejected(rig: AliasRig) -> None: + marker: Final = "shadow-unauthenticated-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + _, shadow = _shadow_over_unpriced_target(rig, scenario) + response: Final = httpx.post( + f"{base_url(rig.gateway)}/v1/chat/completions", + json={"model": shadow, "messages": [{"role": "user", "content": marker}]}, + timeout=60, + trust_env=False, + ) + assert response.status_code == 401, response.text + assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 0 + + +# $0 comes from litellm's own cost-map row (model_prices_and_context_window.json, added 2025-12-08 in #17692). The cell +# needs an embedding row priced $0 only by the map and served over the OpenAI wire; a row that gains a price is stale, +# not a regression: pick another such row +UNPRICED_FREE_EMBEDDING_MODEL: Final = "fireworks_ai/accounts/fireworks/models/qwen3-embedding-0p6b" + + +def _embed(candidate: Gateway, model: str, key: str, marker: str) -> httpx.Response: + return fresh_post(candidate, "/v1/embeddings", {"model": model, "input": marker}, key, marker) + + +def _assert_embeddings_judged_like_target( + candidate: Gateway, shadow: str, target: str, healthy: str, exhausted: str +) -> None: + marker: Final = "embed-shadow-" + uuid.uuid4().hex + served: Final = _embed(candidate, shadow, healthy, marker) + assert served.status_code == 200, served.text + vectors: Final = JSON_OBJECT.validate_json(served.content)["data"] + assert isinstance(vectors, list) and len(vectors) == 1, served.text + assert_free_row(landed_once(healthy, marker), shadow) + via_shadow: Final = _embed(candidate, shadow, exhausted, "embed-refused-shadow-" + uuid.uuid4().hex) + by_name: Final = _embed(candidate, target, exhausted, "embed-refused-target-" + uuid.uuid4().hex) + assert _refusal(by_name) == (BUDGET_EXCEEDED, "budget_exceeded"), by_name.text + assert _refusal(via_shadow) == _refusal(by_name), via_shadow.text + + +def test_exhausted_key_is_refused_on_a_shadowing_alias_through_embeddings(rig: AliasRig) -> None: + with rig.gateway.scenario() as scenario: + target: Final = _unpriced_free_target(rig, scenario, UNPRICED_FREE_EMBEDDING_MODEL) + shadow: Final = _shadowing_alias( + rig, scenario, target, (rig.gateway, rig.peer), input_cost_per_token=0, output_cost_per_token=0 + ) + healthy: Final = scenario.key(max_budget=BUDGET) + exhausted: Final = exhausted_key(rig, scenario) + for candidate in (rig.gateway, rig.peer): + _assert_embeddings_judged_like_target(candidate, shadow, target, healthy, exhausted) + + +def test_failed_primary_does_not_fall_back_to_a_shadowing_alias_for_exhausted_key(rig: AliasRig) -> None: + marker: Final = "fallback-shadow-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + _, shadow = _shadow_over_unpriced_target(rig, scenario) + key: Final = exhausted_key(rig, scenario) + scenario.cleanups.callback(clear_provider_script, rig) + script_provider(rig, 1) + _assert_provider_failure(fresh_chat(rig.gateway, rig.failing_free, key, marker, {"fallbacks": [shadow]})) + assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 1 + _assert_no_served_row(key, marker) + + +def test_removing_a_shadowing_alias_restores_the_explicitly_free_deployment(rig: AliasRig) -> None: + with rig.gateway.scenario() as scenario: + target, shadow = _shadow_over_unpriced_target(rig, scenario) + real: Final = _deployment_id(rig, shadow) + key: Final = exhausted_key(rig, scenario) + settle_chat(rig, target, key, BUDGET_EXCEEDED) + for candidate in (rig.gateway, rig.peer): + _assert_refused_like_target(rig, candidate, key, shadow, target) + remove_aliases(rig.gateway, frozenset({shadow})) + for candidate in (rig.gateway, rig.peer): + _await_route(candidate, shadow, real) + _assert_served_by(rig, candidate, shadow, key, frozenset({real}), "restored") + + +ZERO_MODEL_BUDGET: Final[JsonValue] = {"budget_limit": 0, "time_period": "1d"} +MODEL_BUDGET_SQL: Final = 'UPDATE "LiteLLM_VerificationToken" SET model_max_budget = %s::jsonb WHERE token = %s' + + +def _key_with_zero_model_budgets(scenario: Scenario, names: tuple[str, ...]) -> str: + key: Final = scenario.key(max_budget=BUDGET) + budgets: Final = json.dumps({name: ZERO_MODEL_BUDGET for name in names}) + write_rows(MODEL_BUDGET_SQL, (budgets, sha256(key.encode()).hexdigest())) + return key + + +def _assert_model_budget_refused_alike(rig: AliasRig, candidate: Gateway, key: str, shadow: str, target: str) -> None: + shadow_marker: Final = "model-budget-shadow-" + uuid.uuid4().hex + target_marker: Final = "model-budget-target-" + uuid.uuid4().hex + via_shadow: Final = fresh_chat(candidate, shadow, key, shadow_marker) + by_name: Final = fresh_chat(candidate, target, key, target_marker) + assert by_name.status_code == BUDGET_EXCEEDED, by_name.text + assert _refusal(via_shadow) == _refusal(by_name), via_shadow.text + observed: Final = upstream_requests(rig.gateway.upstream_url) + assert upstream_hits(observed, shadow_marker) == 0 + assert upstream_hits(observed, target_marker) == 0 + + +def test_per_model_budget_on_a_shadowing_alias_is_judged_like_its_target(rig: AliasRig) -> None: + with rig.gateway.scenario() as scenario: + target, shadow = _shadow_over_unpriced_target(rig, scenario) + key: Final = _key_with_zero_model_budgets(scenario, (shadow, target)) + for candidate in (rig.gateway, rig.peer): + _assert_model_budget_refused_alike(rig, candidate, key, shadow, target) + + +def test_exhausted_key_is_not_served_a_cached_reply_through_a_shadowing_alias(rig: AliasRig) -> None: + marker: Final = "shadow-cache-twin-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + _, shadow = _shadow_over_unpriced_target(rig, scenario) + healthy: Final = scenario.key(max_budget=BUDGET) + exhausted: Final = exhausted_key(rig, scenario) + first: Final = fresh_chat(rig.gateway, shadow, healthy, marker) + assert first.status_code == 200, first.text + warmed: Final = fresh_chat(rig.gateway, shadow, healthy, marker) + assert warmed.status_code == 200, warmed.text + assert warmed.json()["id"] == first.json()["id"], warmed.text + assert "x-litellm-cache-key" in warmed.headers, dict(warmed.headers) + refused: Final = fresh_chat(rig.gateway, shadow, exhausted, marker) + _assert_budget_refusal(refused.status_code, refused.text) + assert "x-litellm-cache-key" not in refused.headers, dict(refused.headers) + assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 1 + _assert_no_served_row(exhausted, marker) + + +BUDGET_VERDICTS: Final = frozenset({200, BUDGET_EXCEEDED}) + + +def _only_budget_verdicts(seen: frozenset[int]) -> bool: + assert seen <= BUDGET_VERDICTS, seen + return seen == frozenset({200}) + + +def test_shadowing_alias_removed_under_traffic_keeps_every_answer_a_budget_verdict(rig: AliasRig) -> None: + with rig.gateway.scenario() as scenario: + target, shadow = _shadow_over_unpriced_target(rig, scenario) + real: Final = _deployment_id(rig, shadow) + key: Final = exhausted_key(rig, scenario) + settle_chat(rig, shadow, key, BUDGET_EXCEEDED) + remove_aliases(rig.gateway, frozenset({shadow})) + eventually(lambda: chat_statuses(rig.gateway, shadow, key, GATEWAY_BURST), _only_budget_verdicts, seconds=90) + eventually(lambda: chat_statuses(rig.peer, shadow, key, PEER_BURST), _only_budget_verdicts, seconds=90) + for candidate in (rig.gateway, rig.peer): + _assert_served_by(rig, candidate, shadow, key, frozenset({real}), "after-removal") + assert _serving_deployment(rig.gateway, target) == _deployment_id(rig, target) + + +JWT_KEY_ID: Final = "integration-shadow-jwt-key" + + +def _jwt_config(directory: Path) -> Path: + general: Final = object_value(_base_config()["general_settings"]) + settings: Final[JsonValue] = { + **general, + "enable_jwt_auth": True, + "litellm_jwtauth": {"team_id_jwt_field": "team_id"}, + } + return _own_config(directory, "shadow-jwt.yaml", "general_settings", settings) + + +def _jwk_integer(value: int) -> str: + return base64.urlsafe_b64encode(value.to_bytes((value.bit_length() + 7) // 8, "big")).rstrip(b"=").decode() + + +def _jwks(private_key: rsa.RSAPrivateKey) -> bytes: + numbers: Final = private_key.public_key().public_numbers() + key: Final = {"kty": "RSA", "use": "sig", "alg": "RS256", "kid": JWT_KEY_ID} + return json.dumps({"keys": [{**key, "n": _jwk_integer(numbers.n), "e": _jwk_integer(numbers.e)}]}).encode() + + +@pytest.mark.timeout(480) +def test_jwt_caller_over_team_budget_is_refused_on_a_shadowing_alias(rig: AliasRig, tmp_path: Path) -> None: + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + jwks_body: Final = _jwks(private_key) + + def respond(request: Request) -> Reply: + assert request.method == "GET", request + return Reply(body=jwks_body) + + with wire_server(respond) as jwks, rig.gateway.scenario() as scenario: + target, shadow = _shadow_over_unpriced_target(rig, scenario) + team: Final = scenario.team(max_budget=BUDGET) + issued: Final = int(time.time()) + token: Final = jwt.encode( + {"sub": "integration-shadow-jwt-caller", "team_id": team, "iat": issued, "exp": issued + 900}, + private_key, + algorithm="RS256", + headers={"kid": JWT_KEY_ID}, + ) + with owned_proxy( + rig.gateway, tmp_path, {"JWT_PUBLIC_KEY_URL": jwks.url}, config=_jwt_config(tmp_path) + ) as candidate: + _await_alias(candidate, shadow, target) + first: Final = fresh_chat(candidate, rig.paid, token, "jwt-exhaust-" + uuid.uuid4().hex) + assert first.status_code == 200, first.text + settle_candidate(candidate, rig.paid, token, BUDGET_EXCEEDED, seconds=120) + settle_candidate(candidate, target, token, BUDGET_EXCEEDED, seconds=120) + _assert_refused_like_target(rig, candidate, token, shadow, target) + free_marker: Final = "jwt-free-" + uuid.uuid4().hex + served: Final = fresh_chat(candidate, rig.hidden_free, token, free_marker) + assert served.status_code == 200, served.text + assert upstream_hits(upstream_requests(rig.gateway.upstream_url), free_marker) == 1 + + +CUSTOM_AUTH_DOUBLE: Final = """ +import os + +from fastapi import Request + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + +async def user_api_key_auth(request: Request, api_key: str) -> UserAPIKeyAuth: + if api_key == os.environ["LITELLM_MASTER_KEY"]: + return UserAPIKeyAuth(api_key=api_key, user_role=LitellmUserRoles.PROXY_ADMIN) + if api_key != os.environ["INTEGRATION_CUSTOM_AUTH_KEY"]: + raise Exception("unknown caller") + budgets = { + name: {"budget_limit": 0, "time_period": "1d"} + for name in os.environ["INTEGRATION_ZERO_BUDGET_MODELS"].split(",") + } + return UserAPIKeyAuth(api_key=api_key, user_id="integration-custom-caller", model_max_budget=budgets) +""" + + +def _custom_auth_config(directory: Path) -> Path: + (directory / "custom_auth_double.py").write_text(CUSTOM_AUTH_DOUBLE) + base: Final = _base_config() + general: Final = object_value(base["general_settings"]) + litellm_settings: Final = object_value(base["litellm_settings"]) + path: Final = directory / "shadow-custom-auth.yaml" + path.write_text( + yaml.safe_dump( + { + **base, + "general_settings": {**general, "custom_auth": "custom_auth_double.user_api_key_auth"}, + "litellm_settings": {**litellm_settings, "enable_post_custom_auth_checks": True}, + } + ) + ) + return path + + +@pytest.mark.timeout(480) +def test_custom_auth_caller_with_zero_model_budgets_is_refused_on_a_shadowing_alias( + rig: AliasRig, tmp_path: Path +) -> None: + caller: Final = "sk-custom-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + target, shadow = _shadow_over_unpriced_target(rig, scenario) + overrides: Final = { + "INTEGRATION_CUSTOM_AUTH_KEY": caller, + "INTEGRATION_ZERO_BUDGET_MODELS": ",".join((shadow, target, rig.free)), + } + with owned_proxy(rig.gateway, tmp_path, overrides, config=_custom_auth_config(tmp_path)) as candidate: + _await_alias(candidate, shadow, target) + settle_candidate(candidate, target, caller, BUDGET_EXCEEDED, seconds=120) + _assert_refused_like_target(rig, candidate, caller, shadow, target) + free_marker: Final = "custom-free-" + uuid.uuid4().hex + served: Final = fresh_chat(candidate, rig.free, caller, free_marker) + assert served.status_code == 200, served.text + assert upstream_hits(upstream_requests(rig.gateway.upstream_url), free_marker) == 1 + unknown: Final = fresh_chat(candidate, shadow, "sk-unknown-" + uuid.uuid4().hex, "custom-unknown") + assert unknown.status_code == 401, unknown.text + + +BURST_SERVED: Final = 40 +BURST_FAILURES: Final = 8 +BURST_PROBES: Final = 8 +BURST_KINDS: Final = ("chat", "chat-stream", "responses", "responses-stream", "messages", "messages-stream") + + +@dataclass(frozen=True, slots=True) +class _BurstCall: + kind: str + marker: str + + +def _burst_calls() -> tuple[_BurstCall, ...]: + served: Final = tuple( + _BurstCall(BURST_KINDS[index % len(BURST_KINDS)], "burst-" + uuid.uuid4().hex) for index in range(BURST_SERVED) + ) + probes: Final = tuple(_BurstCall("probe", "probe-" + uuid.uuid4().hex) for _ in range(BURST_PROBES)) + liveliness: Final = tuple(_BurstCall("liveliness", "live-" + uuid.uuid4().hex) for _ in range(BURST_PROBES)) + return served + probes + liveliness + + +def _send_burst_call(rig: AliasRig, shadow: str, healthy: str, exhausted: str, call: _BurstCall) -> httpx.Response: + stream: Final[Mapping[str, JsonValue]] = {"stream": True} + match call.kind: + case "chat": + return fresh_chat(rig.gateway, shadow, healthy, call.marker) + case "chat-stream": + return fresh_chat(rig.gateway, shadow, healthy, call.marker, stream) + case "responses": + return fresh_response(rig.gateway, shadow, healthy, call.marker) + case "responses-stream": + return fresh_response(rig.gateway, shadow, healthy, call.marker, stream) + case "messages": + return fresh_message(rig.gateway, shadow, healthy, call.marker) + case "messages-stream": + return fresh_message(rig.gateway, shadow, healthy, call.marker, stream) + case "probe": + return fresh_chat(rig.gateway, shadow, exhausted, call.marker) + case "liveliness": + return httpx.get(f"{base_url(rig.gateway)}/health/liveliness", timeout=60, trust_env=False) + case _: + raise AssertionError(call.kind) + + +def _served_calls(responses: Mapping[_BurstCall, httpx.Response]) -> Mapping[_BurstCall, httpx.Response]: + return {call: response for call, response in responses.items() if call.kind in BURST_KINDS} + + +def _calls_of_kind(responses: Mapping[_BurstCall, httpx.Response], kind: str) -> tuple[_BurstCall, ...]: + return tuple(call for call in responses if call.kind == kind) + + +def test_shadowing_alias_burst_with_provider_outage_lands_every_served_id_once(rig: AliasRig) -> None: + with rig.gateway.scenario() as scenario: + _, shadow = _shadow_over_unpriced_target(rig, scenario) + healthy: Final = scenario.key(max_budget=BUDGET) + exhausted: Final = exhausted_key(rig, scenario) + scenario.cleanups.callback(_clear_target_script, rig) + _script_target(rig, (500,) * BURST_FAILURES + (200,) * BURST_SERVED) + calls: Final = _burst_calls() + + def send(call: _BurstCall) -> httpx.Response: + return _send_burst_call(rig, shadow, healthy, exhausted, call) + + with ThreadPoolExecutor(max_workers=len(calls)) as pool: + responses: Final = dict(zip(calls, pool.map(send, calls), strict=True)) + served: Final = _served_calls(responses) + failed: Final = {call: response for call, response in served.items() if response.status_code != 200} + assert len(failed) == BURST_FAILURES, {call.kind: response.status_code for call, response in failed.items()} + for response in failed.values(): + _assert_provider_failure(response) + for call in _calls_of_kind(responses, "probe"): + _assert_budget_refusal(responses[call].status_code, responses[call].text) + for call in _calls_of_kind(responses, "liveliness"): + assert responses[call].status_code == 200, responses[call].text + served_markers: Final = frozenset(call.marker for call in served if call not in failed) + assert len(served_markers) == BURST_SERVED - BURST_FAILURES + observed: Final = upstream_requests(rig.gateway.upstream_url) + assert all(upstream_hits(observed, call.marker) == 1 for call in served) + assert all(upstream_hits(observed, call.marker) == 0 for call in _calls_of_kind(responses, "probe")) + for row in landed_all_once(healthy, served_markers): + assert_free_row(row, shadow) + for call in failed: + _assert_no_served_row(healthy, call.marker) + for call in _calls_of_kind(responses, "probe"): + _assert_no_served_row(exhausted, call.marker) + + +def _is_worker(process: psutil.Process, root: int) -> bool: + try: + return ( + process.ppid() == root + and process.status() != psutil.STATUS_ZOMBIE + and "spawn_main" in " ".join(process.cmdline()) + ) + except psutil.Error: + return False + + +def _workers(root: int) -> tuple[psutil.Process, ...]: + return tuple(process for process in group_members(root) if _is_worker(process, root)) + + +def _replaced(victim: psutil.Process) -> Callable[[tuple[psutil.Process, ...]], bool]: + return lambda workers: len(workers) == 2 and victim.pid not in {worker.pid for worker in workers} + + +@pytest.mark.timeout(600) +def test_killing_one_worker_keeps_the_shadowing_alias_refused(rig: AliasRig, tmp_path: Path) -> None: + config: Final = _own_config( + tmp_path, "shadow-two-workers.yaml", "general_settings", _base_config()["general_settings"] + ) + with rig.gateway.scenario() as scenario: + target, shadow = _shadow_over_unpriced_target(rig, scenario) + healthy: Final = scenario.key(max_budget=BUDGET) + exhausted: Final = exhausted_key(rig, scenario) + with owned_proxy_process(rig.gateway, tmp_path, {}, config=config, workers=2) as owned: + candidate: Final = owned.gateway + root: Final = owned.process.pid + before: Final = eventually(lambda: _workers(root), lambda workers: len(workers) == 2, seconds=60) + _await_alias(candidate, shadow, target) + settle_candidate(candidate, target, exhausted, BUDGET_EXCEEDED, seconds=120) + settle_candidate(candidate, shadow, exhausted, BUDGET_EXCEEDED, seconds=120) + victim: Final = before[0] + victim.send_signal(signal.SIGKILL) + eventually(lambda: _workers(root), _replaced(victim), seconds=120) + settle_candidate(candidate, shadow, exhausted, BUDGET_EXCEEDED, seconds=120) + assert chat_statuses(candidate, shadow, exhausted, GATEWAY_BURST) == {BUDGET_EXCEEDED} + assert chat_statuses(candidate, shadow, healthy, GATEWAY_BURST) == {200} + assert chat_statuses(candidate, target, exhausted, GATEWAY_BURST) == {BUDGET_EXCEEDED} diff --git a/tests/unit/proxy/auth/test_unmapped_model_budget_enforcement.py b/tests/unit/proxy/auth/test_unmapped_model_budget_enforcement.py index 515b6e45dc4..be7b438a442 100644 --- a/tests/unit/proxy/auth/test_unmapped_model_budget_enforcement.py +++ b/tests/unit/proxy/auth/test_unmapped_model_budget_enforcement.py @@ -16,6 +16,54 @@ from litellm.proxy.auth.auth_checks import _is_model_cost_zero from litellm.router import Router +UNPRICED_ZERO_COST_MODEL = "ollama/unpriced-zero-cost-target" + + +def _explicitly_free(name: str, model: str = "gpt-3.5-turbo", **model_info: float) -> dict: + return { + "model_name": name, + "litellm_params": { + "model": model, + "api_key": "sk-fake", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, + "model_info": {"id": f"{name}-id", **model_info}, + } + + +def _free_by_cost_map_only(name: str) -> dict: + litellm.model_cost[UNPRICED_ZERO_COST_MODEL] = { + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "litellm_provider": "ollama", + "mode": "chat", + } + return { + "model_name": name, + "litellm_params": {"model": UNPRICED_ZERO_COST_MODEL, "api_base": "http://localhost:11434"}, + "model_info": {"id": f"{name}-id"}, + } + + +def _explicitly_free_ollama_wildcard() -> dict: + return { + "model_name": "ollama/*", + "litellm_params": { + "model": "ollama/*", + "api_base": "http://localhost:11434", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, + "model_info": {"id": "free-wildcard-id"}, + } + + +def _served_model(router: Router, name: str) -> str: + deployment = router.get_available_deployment(model=name, messages=[{"role": "user", "content": "hi"}]) + return deployment["litellm_params"]["model"] + + class TestUnmappedModelBudgetEnforcement: """Unmapped models must NOT bypass budget checks.""" @@ -372,11 +420,12 @@ class TestUnmappedModelBudgetEnforcement: assert _is_model_cost_zero(model="hidden-alias", llm_router=router) is False @pytest.mark.parametrize("alias_name_first", [True, False]) - def test_alias_shadowing_a_real_group_gives_each_name_its_own_verdict(self, alias_name_first: bool): - """An alias whose name is also a real PTU-priced group never shares a verdict with its target. + def test_alias_shadowing_a_real_group_answers_for_its_target_in_either_order(self, alias_name_first: bool): + """An alias whose name is also a real PTU-priced group is judged by the free target it routes to. - The verdict is cached per requested name, so whichever name is asked first, the free target - stays free and the shadowed PTU name stays enforced. + The router resolves the alias before it looks at deployments, so the PTU deployment sharing + the alias's name is never served under it. The verdict is cached per requested name, so + whichever name is asked first, both names read as free. """ router = Router( model_list=[ @@ -405,7 +454,7 @@ class TestUnmappedModelBudgetEnforcement: model_group_alias={"ptu-model": "free-model"}, ) order = ("ptu-model", "free-model") if alias_name_first else ("free-model", "ptu-model") - expected = {"ptu-model": False, "free-model": True} + expected = {"ptu-model": True, "free-model": True} assert [_is_model_cost_zero(model=name, llm_router=router) for name in order] == [ expected[name] for name in order @@ -414,6 +463,124 @@ class TestUnmappedModelBudgetEnforcement: expected[name] for name in order ], "the cached verdicts must match the first evaluation" + @pytest.mark.parametrize( + "alias_entry", + ["unpriced-target", {"model": "unpriced-target", "hidden": True}], + ids=["plain_alias", "hidden_alias"], + ) + def test_alias_shadowing_an_explicitly_free_group_answers_for_its_unpriced_target( + self, alias_entry: str | dict[str, str | bool] + ): + """An alias keyed like an explicitly free real group is judged by its target alone, hidden or not. + + The router serves the alias name from its target, a group whose cost-map price is zero with + no explicit price on the deployment, so the budget stays enforced exactly as it is for the + target by name. The shadowed explicitly free deployment is never served under that name and + must not lend it the bypass. + """ + litellm.model_cost["ollama/unpriced-zero-cost-target"] = { + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "litellm_provider": "ollama", + "mode": "chat", + } + router = Router( + model_list=[ + { + "model_name": "free-model", + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "sk-fake", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, + "model_info": {"id": "free-model-id"}, + }, + { + "model_name": "unpriced-target", + "litellm_params": { + "model": "ollama/unpriced-zero-cost-target", + "api_base": "http://localhost:11434", + }, + "model_info": {"id": "unpriced-target-id"}, + }, + ], + model_group_alias={"free-model": alias_entry}, + ) + + assert _is_model_cost_zero(model="unpriced-target", llm_router=router) is False + assert _is_model_cost_zero(model="free-model", llm_router=router) is False, ( + "the alias routes to the unpriced target, so it must be refused like the target by name" + ) + + def test_hidden_alias_shadowing_a_ptu_group_answers_for_its_free_target(self): + """A hidden alias keyed like a PTU-priced real group reads as free when its target is free. + + The PTU deployment sharing the alias's name is never served under it, so its flat cost must + not keep the alias enforced while the router serves every call from the free target. + """ + router = Router( + model_list=[ + { + "model_name": "ptu-model", + "litellm_params": { + "model": "azure/ptu-deployment", + "api_base": "https://fake.openai.azure.com", + "api_key": "sk-fake", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, + "model_info": {"id": "ptu-model-id", "ptu_count": 100, "cost_per_ptu_per_hour": 2.0}, + }, + { + "model_name": "free-model", + "litellm_params": { + "model": "ollama/llama2", + "api_base": "http://localhost:11434", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, + "model_info": {"id": "free-model-id"}, + }, + ], + model_group_alias={"ptu-model": {"model": "free-model", "hidden": True}}, + ) + + assert _is_model_cost_zero(model="free-model", llm_router=router) is True + assert _is_model_cost_zero(model="ptu-model", llm_router=router) is True, ( + "the alias routes to the free target, so it must bypass budget like the target by name" + ) + + def test_hidden_alias_shadowing_an_explicitly_free_group_to_a_priced_target_enforces_budget(self): + """A hidden alias keyed like an explicitly free real group stays enforced when its target is priced.""" + router = Router( + model_list=[ + { + "model_name": "free-model", + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "sk-fake", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, + "model_info": {"id": "free-model-id"}, + }, + { + "model_name": "paid-model", + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "sk-fake", + "input_cost_per_token": 0.0000002, + "output_cost_per_token": 0.0000012, + }, + "model_info": {"id": "paid-model-id"}, + }, + ], + model_group_alias={"free-model": {"model": "paid-model", "hidden": True}}, + ) + + assert _is_model_cost_zero(model="free-model", llm_router=router) is False + def test_alias_chain_through_a_priced_group_enforces_budget(self): """An alias to a group that is itself an alias key resolves one hop, like the router does. @@ -448,6 +615,141 @@ class TestUnmappedModelBudgetEnforcement: assert _is_model_cost_zero(model="chain-smart", llm_router=router) is False + @pytest.mark.parametrize("hidden", [False, True], ids=["plain_alias", "hidden_alias"]) + def test_alias_to_an_unpriced_group_that_is_also_an_alias_enforces_budget(self, hidden: bool): + """An alias is judged by the group it routes to, never by where that group's own alias points. + + The router serves ``chain-entry`` from the real ``chain-middle`` deployment, whose zero price + comes from the cost map alone. ``chain-middle`` is also an alias key to an explicitly free + group, a second hop the router never takes for ``chain-entry``, so that group must not lend + it the bypass. + """ + + def alias(target: str) -> str | dict[str, str | bool]: + return {"model": target, "hidden": True} if hidden else target + + router = Router( + model_list=[_free_by_cost_map_only("chain-middle"), _explicitly_free("free-model")], + model_group_alias={"chain-entry": alias("chain-middle"), "chain-middle": alias("free-model")}, + ) + + assert _served_model(router, "chain-entry") == UNPRICED_ZERO_COST_MODEL + assert _is_model_cost_zero(model="chain-entry", llm_router=router) is False + assert _is_model_cost_zero(model="chain-middle", llm_router=router) is True, ( + "asked by its own name, chain-middle routes to the explicitly free group" + ) + + def test_alias_to_a_free_group_that_is_also_an_alias_to_a_ptu_group_bypasses_budget(self): + """A PTU group one alias hop past the group that serves the request does not enforce the budget. + + The router serves ``chain-entry`` from the real, explicitly free ``chain-middle`` deployment. + ``chain-middle`` is also an alias key to a PTU-priced group, which only a request for + ``chain-middle`` itself routes to. + """ + router = Router( + model_list=[ + _explicitly_free("chain-middle"), + _explicitly_free("ptu-model", model="azure/ptu-deployment", ptu_count=100, cost_per_ptu_per_hour=2.0), + ], + model_group_alias={"chain-entry": "chain-middle", "chain-middle": "ptu-model"}, + ) + + assert _served_model(router, "chain-entry") == "gpt-3.5-turbo" + assert _is_model_cost_zero(model="chain-entry", llm_router=router) is True + assert _is_model_cost_zero(model="chain-middle", llm_router=router) is False, ( + "asked by its own name, chain-middle routes to the PTU group" + ) + + def test_alias_chain_served_by_a_paid_wildcard_route_enforces_budget(self): + """An alias whose target is only an alias key is never judged by that second alias's free group. + + ``chain-entry`` resolves one hop to ``gpt-4o-mini``, which is no deployment's name, so the + router serves it from the paid wildcard route. The free group ``gpt-4o-mini`` is aliased + to is only reached by a request for ``gpt-4o-mini`` itself. + """ + router = Router( + model_list=[ + _explicitly_free("free-model", model="ollama/llama2"), + {"model_name": "*", "litellm_params": {"model": "openai/*", "api_key": "sk-fake"}}, + ], + model_group_alias={"chain-entry": "gpt-4o-mini", "gpt-4o-mini": "free-model"}, + ) + + assert _served_model(router, "chain-entry") == "openai/gpt-4o-mini" + assert _is_model_cost_zero(model="chain-entry", llm_router=router) is False + + @pytest.mark.parametrize("alias_name", ["openai/smart", "smart"], ids=["alias_on_pattern", "alias_off_pattern"]) + def test_alias_chain_served_by_an_explicitly_priced_wildcard_route_enforces_budget(self, alias_name: str): + """A chain the router serves from a priced wildcard route is budgeted at that route's price. + + The group's price reads $0 through the second alias to the free group, and the wildcard + route's cost-map entry carries explicit prices, so only their sign tells that the served + deployment is not free. + """ + router = Router( + model_list=[ + _explicitly_free("free-model", model="ollama/llama2"), + { + "model_name": "openai/*", + "litellm_params": { + "model": "openai/*", + "api_key": "sk-fake", + "input_cost_per_token": 0.00001, + "output_cost_per_token": 0.00002, + }, + "model_info": {"id": "priced-wildcard-id"}, + }, + ], + model_group_alias={alias_name: "openai/gpt-4o-mini", "openai/gpt-4o-mini": "free-model"}, + ) + + assert _served_model(router, alias_name) == "openai/gpt-4o-mini" + assert _is_model_cost_zero(model=alias_name, llm_router=router) is False + + def test_alias_shadowing_a_free_group_is_judged_by_its_unpriced_target_through_an_alias_chain(self): + """A shadowing alias stays enforced when its unpriced target is itself an alias key to a free group.""" + router = Router( + model_list=[ + _explicitly_free("shadowed-free"), + _free_by_cost_map_only("chain-middle"), + _explicitly_free("free-model"), + ], + model_group_alias={"shadowed-free": "chain-middle", "chain-middle": "free-model"}, + ) + + assert _served_model(router, "shadowed-free") == UNPRICED_ZERO_COST_MODEL + assert _is_model_cost_zero(model="shadowed-free", llm_router=router) is False + + @pytest.mark.parametrize("alias_name", ["ollama/fast", "fast"], ids=["alias_on_pattern", "alias_off_pattern"]) + def test_alias_to_a_name_served_by_an_explicitly_free_wildcard_route_bypasses_budget(self, alias_name: str): + """An alias to a name only a wildcard route serves reads that route's deployment, like the name itself.""" + router = Router( + model_list=[_explicitly_free_ollama_wildcard()], + model_group_alias={alias_name: "ollama/llama3"}, + ) + + assert _served_model(router, alias_name) == "ollama/llama3" + assert _is_model_cost_zero(model="ollama/llama3", llm_router=router) is True + assert _is_model_cost_zero(model=alias_name, llm_router=router) is True + + def test_alias_chain_to_a_name_served_by_an_explicitly_free_wildcard_route_bypasses_budget(self): + """An alias to an alias key no deployment is named after reads the wildcard route serving it. + + ``ollama/fast`` resolves one hop to ``ollama/llama3``, which is only an alias key, so the + router serves it from the explicitly free wildcard route. The group ``ollama/llama3`` is + aliased to is only reached by a request for ``ollama/llama3`` itself. + """ + router = Router( + model_list=[_explicitly_free_ollama_wildcard(), _free_by_cost_map_only("unpriced-model")], + model_group_alias={"ollama/fast": "ollama/llama3", "ollama/llama3": "unpriced-model"}, + ) + + assert _served_model(router, "ollama/fast") == "ollama/llama3" + assert _is_model_cost_zero(model="ollama/fast", llm_router=router) is True + assert _is_model_cost_zero(model="ollama/llama3", llm_router=router) is False, ( + "asked by its own name, ollama/llama3 routes to the unpriced group" + ) + def test_handles_router_without_zero_cost_cache_attribute(self): """Tolerate router-like objects (e.g. ``MagicMock`` stand-ins) that do not expose ``_zero_cost_cache`` — the auth check must still