fix(proxy): judge the free-model budget waiver by the group an alias routes to (#44638)
Some checks are pending
CI Coverage / assert-ci-coverage (push) Waiting to run
CodeQL / Analyze (actions) (push) Waiting to run
CodeQL / Analyze (javascript-typescript) (push) Waiting to run
CodeQL / Analyze (python) (push) Waiting to run
CodSpeed Benchmarks / benchmarks (push) Waiting to run
Helm unit test / unit-test (push) Waiting to run
Lens Worker Image / lens-worker-image (push) Waiting to run
Publish basedpyright base counts / publish (push) Waiting to run
Scorecard supply-chain security / Scorecard analysis (push) Waiting to run
Code Quality Checks / code-quality (push) Waiting to run
Code Quality Checks / python-310-import-smoke (push) Waiting to run
UI Unit Tests / ui-unit-tests (push) Waiting to run
LiteLLM Rust / rust-lint (push) Waiting to run
LiteLLM Rust / rust-test (push) Waiting to run
LiteLLM Rust / rust-wheel (push) Waiting to run
Unit Tests: Documentation Validation / documentation (push) Waiting to run
Unit Tests: Proxy DB Operations / assert-shard-coverage (push) Waiting to run
Unit Tests: Proxy DB Operations / auth-checks (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / budgets (push) Blocked by required conditions
Unit Tests / Build the Rust bridge (push) Waiting to run
Unit Tests / caching-local (push) Blocked by required conditions
VS Code Extension / vscode-extension (push) Waiting to run
Unit Tests: Proxy DB Operations / custom-logging (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / db-and-spend (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / endpoints-and-responses (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / guardrails-hooks (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / jwt-and-keys (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / key-generation (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / logging-misc (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / proxy-runtime (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / proxy-server-core (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / proxy-utils (push) Blocked by required conditions
Unit Tests / core-utils (push) Blocked by required conditions
Unit Tests / enterprise-managed-files (push) Blocked by required conditions
Unit Tests / enterprise-package (push) Blocked by required conditions
Unit Tests / enterprise-routing (push) Blocked by required conditions
Unit Tests / integrations (push) Blocked by required conditions
Unit Tests / OpenAI and Meta Providers (push) Blocked by required conditions
Unit Tests / All Other Providers (push) Blocked by required conditions
Unit Tests / Vertex AI (push) Blocked by required conditions
Unit Tests / misc (push) Blocked by required conditions
Unit Tests / misc-dirs (push) Blocked by required conditions
Unit Tests / proxy-auth (push) Blocked by required conditions
Unit Tests / proxy-endpoints (push) Blocked by required conditions
Unit Tests / proxy-extras (push) Blocked by required conditions
Unit Tests / proxy-feature-endpoints (push) Blocked by required conditions
Unit Tests / proxy-hooks-client (push) Blocked by required conditions
Unit Tests / proxy-infra (push) Blocked by required conditions
Unit Tests / proxy-infra-root (push) Blocked by required conditions
Unit Tests / proxy-server (push) Blocked by required conditions
Unit Tests / responses-caching-types (push) Blocked by required conditions
Unit Tests / unit (push) Blocked by required conditions
GitHub Actions Security Analysis / zizmor (push) Waiting to run

* fix(proxy): judge the free-model budget waiver by the group an alias routes to

A hidden model_group_alias can reuse the name of a real model_name. The
router serves that name from the alias target, but two of the three checks
behind the free-model budget waiver read the deployments of both the alias
target and the real model that shares the name.

So an over-budget key was served through an alias whose name belongs to a
model with an explicit $0 price when the target is $0 only by the cost map,
which the same key is refused on by its own name. The mirror case refused a
free target because the shadowed name belongs to a PTU-priced deployment.

Resolve the alias once and have the explicit-cost and PTU checks read the
routed group, the same group the price check already reads.

* test(proxy): cover the plain alias form of a shadowed free model name

The same wrong verdict exists for a plain string alias, so the unpriced-target case now runs for both alias shapes.

* fix(proxy): judge an alias chain's budget waiver by the deployments it is served from

The explicit-price and PTU checks read the alias target through
Router.get_model_list(), which follows a second alias hop when the target is
itself an alias key. The router never takes that hop, so an alias chain was
judged by a deployment the request never reaches. The checks now take the
deployments named after the routed group, or the wildcard deployment serving
it when none carries its name.

* fix(proxy): refuse the budget waiver when an alias chain is served by a priced wildcard route

* test(integration): cover the shadowing alias budget gate end to end

Adds the audit cells for a hidden alias whose name shadows an explicitly
free group: streamed SDK refusals on chat, responses and messages,
embeddings, the free wildcard and mixed-group paths, per-model budgets,
JWT and custom auth callers, cache hits, alias removal under traffic,
provider failure and fallback, a concurrent outage burst, and a worker
kill on an owned two-worker proxy

* test(integration): cite the cost-map rows the shadowing alias cells rely on

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-10-06 02:44:35 +00:00 • committed by GitHub
parent ff5084687c
commit 4a963a6c0b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 1189 additions and 25 deletions

View file

@ -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": "<hash>"} 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, ...]:

View file

@ -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 ``<prefix>/...`` 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}

View file

@ -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