From 1ef0fe9790afc007c54b32c8167faeac54ec5498 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 3 Oct 2026 16:35:11 -0700 Subject: [PATCH] fix(auth): resolve hidden model_group_alias entries in the zero-cost budget check (#43741) * fix(auth): resolve hidden model_group_alias entries in the zero-cost budget check Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(auth): key the zero-cost cache by the resolved model group Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(auth): keep the zero-cost verdict per requested name and include hidden aliases * test(integration): audit the zero-cost bypass through hidden model_group_alias names * test(router): cover the extracted routing strategy switch * test(integration): record a pre-flip burst before the alias flip --------- Co-authored-by: jesus Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 2 +- litellm/router.py | 35 +- .../authorization/_hidden_alias_budget.py | 423 +++++++++++++ .../test_hidden_alias_budget_bypass.py | 577 ++++++++++++++++++ .../test_hidden_alias_budget_bypass_chaos.py | 208 +++++++ .../test_unmapped_model_budget_enforcement.py | 150 ++++- tests/unit/test_router/test_router.py | 67 ++ 7 files changed, 1430 insertions(+), 32 deletions(-) create mode 100644 tests/integration/authorization/_hidden_alias_budget.py create mode 100644 tests/integration/authorization/test_hidden_alias_budget_bypass.py create mode 100644 tests/integration/authorization/test_hidden_alias_budget_bypass_chaos.py diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 918428bcf0e..dfdf7dc4e66 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -484,7 +484,7 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None continue try: # Use router's get_model_group_info method directly for better reliability - model_group_info = llm_router.get_model_group_info(model_group=model_name) + model_group_info = llm_router.get_model_group_info(model_group=model_name, include_hidden=True) if model_group_info is None: # Model not found or no pricing info available diff --git a/litellm/router.py b/litellm/router.py index 49e9b9b5a78..54fffae8dbd 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -11246,13 +11246,13 @@ class Router: return model_group_info - def get_model_group_info(self, model_group: str) -> ModelGroupInfo | None: + def get_model_group_info(self, model_group: str, *, include_hidden: bool = False) -> ModelGroupInfo | None: """ For a given model group name, return the combined model info Returns: - ModelGroupInfo if able to construct a model group - - None if error constructing model group info or hidden model group + - None if error constructing model group info or hidden model group (unless include_hidden) """ ## Check if model group alias if model_group in self.model_group_alias: @@ -11260,7 +11260,7 @@ class Router: if isinstance(item, str): _router_model_group = item elif isinstance(item, dict): - if item["hidden"] is True: + if item["hidden"] is True and not include_hidden: return None else: _router_model_group = item["model"] @@ -12302,6 +12302,16 @@ class Router: ] return _settings_to_return + def _switch_routing_strategy(self, routing_strategy: str | None, kwargs: Mapping[str, object]) -> None: + if routing_strategy == "lar1": + from litellm.router_strategy.lar1_routing import apply_lar1_routing_strategy + + apply_lar1_routing_strategy(self, kwargs.get("routing_strategy_args")) + return + self.routing_strategy_init( + routing_strategy=routing_strategy, routing_strategy_args=kwargs.get("routing_strategy_args", {}) + ) + def update_settings(self, **kwargs): """ Update the router settings. @@ -12315,6 +12325,7 @@ class Router: ] _existing_router_settings: Final = self.get_settings() + model_group_alias_before: Final = self.model_group_alias rebuild_routing_groups = False routing_args_updated = False for var in kwargs: @@ -12338,20 +12349,7 @@ class Router: if var == "routing_strategy": value = self._normalize_strategy(value) if _existing_router_settings["routing_strategy"] != value: - if value == "lar1": - from litellm.router_strategy.lar1_routing import ( - apply_lar1_routing_strategy, - ) - - apply_lar1_routing_strategy( - self, - kwargs.get("routing_strategy_args"), - ) - else: - self.routing_strategy_init( - routing_strategy=value, - routing_strategy_args=kwargs.get("routing_strategy_args", {}), - ) + self._switch_routing_strategy(value, kwargs) rebuild_routing_groups = True elif var == "routing_strategy_args": routing_args_updated = value != self.routing_strategy_args @@ -12362,6 +12360,9 @@ class Router: if routing_args_updated: self._apply_updated_routing_strategy_args() + if self.model_group_alias != model_group_alias_before: + self._invalidate_model_group_info_cache() + if rebuild_routing_groups: routing_groups_input: Final = kwargs.get("routing_groups", self._routing_groups_input) self._init_routing_groups(routing_groups_input) diff --git a/tests/integration/authorization/_hidden_alias_budget.py b/tests/integration/authorization/_hidden_alias_budget.py new file mode 100644 index 00000000000..eb099e56026 --- /dev/null +++ b/tests/integration/authorization/_hidden_alias_budget.py @@ -0,0 +1,423 @@ +import json +import os +import uuid +from collections import Counter +from collections.abc import Callable, Generator, Mapping +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from dataclasses import dataclass +from hashlib import sha256 +from types import MappingProxyType +from typing import Final + +import httpx +from anthropic import Anthropic, AsyncAnthropic +from integration._support.anthropic_thinking import JSON_LIST, JSON_OBJECT +from integration._support.client import ( + GATEWAY_LIMITS, + Gateway, + Scenario, + eventually, + gateway_from_environment, + object_value, +) +from integration._support.database import read_rows +from integration._support.upstream import delete_scenario, register_scenario +from integration.cost_calculation.cost_tracking_case import JsonResponse, SseResponse +from openai import AsyncOpenAI, OpenAI +from pydantic import JsonValue + +BUDGET: Final = 0.05 +CHAT_REPLY: Final = "Hello! This is a mock response from the fake OpenAI endpoint." +RESPONSES_REPLY: Final = "free reply" +BUDGET_EXCEEDED: Final = 422 +GATEWAY_BURST: Final = 24 +PEER_BURST: Final = 4 +SPEND_MARKER_HEADER: Final = "x-litellm-spend-logs-metadata" +PROXY_BUDGET_USER: Final = "litellm-proxy-budget" + +_RESPONSE: Final[JsonValue] = { + "id": "resp_$UNIQUE_ID", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "id": "msg_$UNIQUE_ID", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": RESPONSES_REPLY, "annotations": []}], + } + ], + "usage": {"input_tokens": 20, "output_tokens": 20, "total_tokens": 40}, +} +_RESPONSE_EVENTS: Final[tuple[Mapping[str, JsonValue], ...]] = ( + { + "type": "response.created", + "sequence_number": 0, + "response": {**_RESPONSE, "status": "in_progress", "output": []}, + }, + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": "msg_$UNIQUE_ID", + "output_index": 0, + "content_index": 0, + "delta": RESPONSES_REPLY, + }, + {"type": "response.completed", "sequence_number": 2, "response": _RESPONSE}, +) + + +@dataclass(frozen=True, slots=True) +class AliasRig: + gateway: Gateway + peer: Gateway + free: str + paid: str + failing_free: str + failing_provider_model: str + hidden_free: str + visible_free: str + hidden_paid: str + hidden_responses: str + hidden_responses_stream: str + hidden_failing: str + shown_free: str + null_hidden_free: str + hidden_unpriced: str + hidden_missing: str + + +def base_url(candidate: Gateway) -> str: + return str(candidate.client.base_url).rstrip("/") + + +def spend_marker(marker: str) -> Mapping[str, str]: + return {SPEND_MARKER_HEADER: json.dumps({"marker": marker})} + + +def fresh_post(candidate: Gateway, path: str, body: Mapping[str, JsonValue], key: str, marker: str) -> httpx.Response: + return httpx.post( + f"{base_url(candidate)}{path}", + json=dict(body), + headers={"Authorization": f"Bearer {key}", **spend_marker(marker)}, + timeout=60, + trust_env=False, + ) + + +NO_EXTRA: Final[Mapping[str, JsonValue]] = MappingProxyType({}) + + +def chat_body(model: JsonValue, marker: str, extra: Mapping[str, JsonValue] = NO_EXTRA) -> Mapping[str, JsonValue]: + return {"model": model, "messages": [{"role": "user", "content": marker}], **extra} + + +def fresh_chat( + candidate: Gateway, model: str, key: str, marker: str, extra: Mapping[str, JsonValue] = NO_EXTRA +) -> httpx.Response: + return fresh_post(candidate, "/v1/chat/completions", chat_body(model, marker, extra), key, marker) + + +def fresh_response( + candidate: Gateway, model: str, key: str, marker: str, extra: Mapping[str, JsonValue] = NO_EXTRA +) -> httpx.Response: + return fresh_post(candidate, "/v1/responses", {"model": model, "input": marker, **extra}, key, marker) + + +def fresh_message( + candidate: Gateway, model: str, key: str, marker: str, extra: Mapping[str, JsonValue] = NO_EXTRA +) -> httpx.Response: + return fresh_post( + candidate, + "/v1/messages", + {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": marker}], **extra}, + key, + marker, + ) + + +def statuses(send: Callable[[str], httpx.Response], count: int) -> frozenset[int]: + markers: Final = tuple(uuid.uuid4().hex for _ in range(count)) + with ThreadPoolExecutor(max_workers=count) as pool: + return frozenset(response.status_code for response in pool.map(send, markers)) + + +def error_type(response: httpx.Response) -> str: + return str(object_value(JSON_OBJECT.validate_json(response.content)["error"])["type"]) + + +def settle( + send: Callable[[str], httpx.Response], status: int, *, burst: int = GATEWAY_BURST, seconds: float = 60 +) -> None: + eventually( + lambda: statuses(send, burst) | statuses(send, burst), + lambda seen: seen == frozenset({status}), + seconds=seconds, + ) + + +def chat_statuses( + candidate: Gateway, model: str, key: str, count: int, extra: Mapping[str, JsonValue] = NO_EXTRA +) -> frozenset[int]: + return statuses(lambda marker: fresh_chat(candidate, model, key, marker, extra), count) + + +def settle_candidate( + candidate: Gateway, + model: str, + key: str, + status: int, + *, + burst: int = GATEWAY_BURST, + seconds: float = 60, + extra: Mapping[str, JsonValue] = NO_EXTRA, +) -> None: + settle(lambda marker: fresh_chat(candidate, model, key, marker, extra), status, burst=burst, seconds=seconds) + + +def settle_chat( + rig: AliasRig, + model: str, + key: str, + status: int, + *, + seconds: float = 60, + extra: Mapping[str, JsonValue] = NO_EXTRA, +) -> None: + settle_candidate(rig.gateway, model, key, status, seconds=seconds, extra=extra) + settle_candidate(rig.peer, model, key, status, burst=PEER_BURST, seconds=seconds, extra=extra) + + +def exhausted_key(rig: AliasRig, scenario: Scenario) -> str: + key: Final = scenario.key(max_budget=BUDGET) + first: Final = fresh_chat(rig.gateway, rig.paid, key, "exhaust-" + uuid.uuid4().hex) + assert first.status_code == 200, first.text + settle_chat(rig, rig.paid, key, BUDGET_EXCEEDED) + return key + + +def upstream_requests(upstream_url: str) -> tuple[str, ...]: + drained: Final = JSON_OBJECT.validate_json( + httpx.get(f"{upstream_url}/__observations", timeout=15, trust_env=False).content + ) + return tuple(json.dumps(entry) for entry in JSON_LIST.validate_python(drained["requests"])) + + +def upstream_hits(observed: tuple[str, ...], marker: str) -> int: + return sum(1 for entry in observed if marker in entry) + + +def script_provider(rig: AliasRig, failures: int) -> None: + scripted: Final = httpx.post( + f"{rig.gateway.upstream_url}/__scripts/{rig.failing_provider_model}", + json={"statuses": [500] * failures}, + timeout=15, + trust_env=False, + ) + assert scripted.status_code == 200, scripted.text + + +def clear_provider_script(rig: AliasRig) -> None: + cleared: Final = httpx.delete( + f"{rig.gateway.upstream_url}/__scripts/{rig.failing_provider_model}", timeout=15, trust_env=False + ) + assert cleared.status_code in (200, 404), cleared.text + + +def spend_rows(key: str) -> tuple[Mapping[str, JsonValue], ...]: + return tuple( + read_rows( + "SELECT request_id, spend, model_group, status, call_type, " + "metadata->'spend_logs_metadata'->>'marker' AS marker " + 'FROM "LiteLLM_SpendLogs" WHERE api_key = %s', + (sha256(key.encode()).hexdigest(),), + ) + ) + + +def landed(key: str, marker: str) -> tuple[Mapping[str, JsonValue], ...]: + return tuple(row for row in spend_rows(key) if row["marker"] == marker) + + +def landed_once(key: str, marker: str) -> Mapping[str, JsonValue]: + rows: Final = eventually(lambda: landed(key, marker), lambda found: len(found) >= 1, seconds=70) + assert len(rows) == 1, rows + return rows[0] + + +def marker_counts(key: str, markers: frozenset[str]) -> Mapping[str, int]: + return dict(Counter(str(row["marker"]) for row in spend_rows(key) if row["marker"] in markers)) + + +def landed_all_once(key: str, markers: frozenset[str]) -> tuple[Mapping[str, JsonValue], ...]: + counts: Final = eventually( + lambda: marker_counts(key, markers), lambda found: frozenset(found) == markers, seconds=90 + ) + assert counts == dict.fromkeys(markers, 1), counts + return tuple(row for row in spend_rows(key) if row["marker"] in markers) + + +def assert_free_row(row: Mapping[str, JsonValue], model_group: str) -> None: + assert float(str(row["spend"])) == 0.0, row + assert row["model_group"] == model_group, row + assert row["status"] == "success", row + + +def alias_map(gateway: Gateway) -> Mapping[str, JsonValue]: + current: Final = object_value(gateway.get("/router/settings")["current_values"]).get("model_group_alias") + return object_value(current) if current is not None else {} + + +def write_aliases(gateway: Gateway, aliases: Mapping[str, JsonValue]) -> None: + gateway.post("/config/update", {"router_settings": {"model_group_alias": dict(aliases)}}) + + +def install_aliases(gateway: Gateway, aliases: Mapping[str, JsonValue]) -> None: + write_aliases(gateway, {**alias_map(gateway), **aliases}) + + +def remove_aliases(gateway: Gateway, names: frozenset[str]) -> None: + write_aliases(gateway, {name: target for name, target in alias_map(gateway).items() if name not in names}) + + +def hidden(group: str) -> JsonValue: + return {"model": group, "hidden": True} + + +def openai_client(candidate: Gateway, key: str) -> OpenAI: + return OpenAI( + api_key=key, + base_url=base_url(candidate) + "/v1", + max_retries=0, + http_client=httpx.Client(timeout=60, trust_env=False), + ) + + +def async_openai_client(candidate: Gateway, key: str) -> AsyncOpenAI: + return AsyncOpenAI( + api_key=key, + base_url=base_url(candidate) + "/v1", + max_retries=0, + http_client=httpx.AsyncClient(timeout=60, trust_env=False), + ) + + +def anthropic_client(candidate: Gateway, key: str) -> Anthropic: + return Anthropic( + api_key=key, + base_url=base_url(candidate), + max_retries=0, + http_client=httpx.Client(timeout=60, trust_env=False), + ) + + +def async_anthropic_client(candidate: Gateway, key: str) -> AsyncAnthropic: + return AsyncAnthropic( + api_key=key, + base_url=base_url(candidate), + max_retries=0, + http_client=httpx.AsyncClient(timeout=60, trust_env=False), + ) + + +def _zero_cost_responses_group( + scenario: Scenario, gateway: Gateway, name: str, response: JsonResponse | SseResponse +) -> str: + handle: Final = register_scenario(name, response, control_url=gateway.upstream_url) + scenario.cleanups.callback(delete_scenario, handle) + return scenario.model(api_base=f"{handle.api_base()}/v1", input_cost_per_token=0, output_cost_per_token=0) + + +def settle_responses( + rig: AliasRig, model: str, key: str, status: int, extra: Mapping[str, JsonValue] = NO_EXTRA +) -> None: + settle(lambda marker: fresh_response(rig.gateway, model, key, marker, extra), status) + settle(lambda marker: fresh_response(rig.peer, model, key, marker, extra), status, burst=PEER_BURST) + + +def _await_rig(rig: AliasRig) -> None: + admin: Final = rig.gateway.key + for model in ( + rig.hidden_free, + rig.visible_free, + rig.hidden_paid, + rig.hidden_failing, + rig.shown_free, + rig.null_hidden_free, + rig.hidden_unpriced, + ): + settle_chat(rig, model, admin, 200) + settle_responses(rig, rig.hidden_responses, admin, 200) + settle_responses(rig, rig.hidden_responses_stream, admin, 200, extra={"stream": True}) + + +@contextmanager +def alias_rig() -> Generator[AliasRig]: + suffix: Final = uuid.uuid4().hex[:12] + with ( + gateway_from_environment() as gateway, + httpx.Client( + base_url=os.environ["INTEGRATION_PEER_URL"], timeout=15, trust_env=False, limits=GATEWAY_LIMITS + ) as peer_client, + gateway.scenario() as scenario, + ): + free: Final = scenario.model(input_cost_per_token=0, output_cost_per_token=0) + paid: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + unpriced: Final = scenario.model() + failing_provider_model: Final = f"hidden-alias-failing-{suffix}" + failing_free: Final = scenario.model( + model=f"openai/{failing_provider_model}", input_cost_per_token=0, output_cost_per_token=0 + ) + responses_free: Final = _zero_cost_responses_group( + scenario, + gateway, + f"hidden-alias-json-{suffix}", + JsonResponse(content_type="application/json", body=_RESPONSE), + ) + responses_stream_free: Final = _zero_cost_responses_group( + scenario, + gateway, + f"hidden-alias-sse-{suffix}", + SseResponse( + content_type="text/event-stream", + frames=tuple(f"event: {event['type']}\ndata: {json.dumps(event)}" for event in _RESPONSE_EVENTS), + ), + ) + rig: Final = AliasRig( + gateway=gateway, + peer=Gateway(peer_client, gateway.key, gateway.upstream_url), + free=free, + paid=paid, + failing_free=failing_free, + failing_provider_model=failing_provider_model, + hidden_free=f"hidden-free-{suffix}", + visible_free=f"visible-free-{suffix}", + hidden_paid=f"hidden-paid-{suffix}", + hidden_responses=f"hidden-responses-{suffix}", + hidden_responses_stream=f"hidden-responses-stream-{suffix}", + hidden_failing=f"hidden-failing-{suffix}", + shown_free=f"shown-free-{suffix}", + null_hidden_free=f"null-hidden-free-{suffix}", + hidden_unpriced=f"hidden-unpriced-{suffix}", + hidden_missing=f"hidden-missing-{suffix}", + ) + aliases: Final[Mapping[str, JsonValue]] = { + rig.hidden_free: hidden(free), + rig.visible_free: free, + rig.hidden_paid: hidden(paid), + rig.hidden_responses: hidden(responses_free), + rig.hidden_responses_stream: hidden(responses_stream_free), + rig.hidden_failing: hidden(failing_free), + rig.shown_free: {"model": free, "hidden": False}, + rig.null_hidden_free: {"model": free, "hidden": None}, + rig.hidden_unpriced: hidden(unpriced), + rig.hidden_missing: hidden(f"missing-group-{suffix}"), + } + install_aliases(gateway, aliases) + scenario.cleanups.callback(remove_aliases, gateway, frozenset(aliases)) + _await_rig(rig) + yield rig diff --git a/tests/integration/authorization/test_hidden_alias_budget_bypass.py b/tests/integration/authorization/test_hidden_alias_budget_bypass.py new file mode 100644 index 00000000000..46d0852ca8a --- /dev/null +++ b/tests/integration/authorization/test_hidden_alias_budget_bypass.py @@ -0,0 +1,577 @@ +import json +import math +import uuid +from collections.abc import Iterator, Mapping +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +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.authorization._hidden_alias_budget import ( + BUDGET, + BUDGET_EXCEEDED, + CHAT_REPLY, + PROXY_BUDGET_USER, + RESPONSES_REPLY, + AliasRig, + alias_rig, + anthropic_client, + assert_free_row, + async_anthropic_client, + async_openai_client, + base_url, + chat_statuses, + clear_provider_script, + error_type, + exhausted_key, + fresh_chat, + hidden, + install_aliases, + landed, + landed_once, + openai_client, + remove_aliases, + script_provider, + settle_candidate, + settle_chat, + spend_marker, + upstream_hits, + upstream_requests, +) +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(240) + + +@pytest.fixture(scope="module") +def rig() -> Iterator[AliasRig]: + with alias_rig() as built: + yield built + + +def test_exhausted_key_reaches_hidden_free_alias_through_openai_chat(rig: AliasRig) -> None: + marker: Final = "chat-sync-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + key: Final = exhausted_key(rig, scenario) + with openai_client(rig.gateway, key) as client: + completion: Final = client.chat.completions.create( + model=rig.hidden_free, + messages=[{"role": "user", "content": marker}], + extra_headers=spend_marker(marker), + ) + assert completion.choices[0].message.content == CHAT_REPLY, completion + assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 1 + row: Final = landed_once(key, marker) + assert row["request_id"] == completion.id, row + assert_free_row(row, rig.hidden_free) + + +async def test_exhausted_key_reaches_hidden_free_alias_through_streamed_openai_chat(rig: AliasRig) -> None: + marker: Final = "chat-stream-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + key: Final = exhausted_key(rig, scenario) + async with async_openai_client(rig.gateway, key) as client: + stream: Final = await client.chat.completions.create( + model=rig.hidden_free, + messages=[{"role": "user", "content": marker}], + stream=True, + stream_options={"include_usage": True}, + extra_headers=spend_marker(marker), + ) + chunks: Final = tuple([chunk async for chunk in stream]) + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == CHAT_REPLY + assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 1 + row: Final = landed_once(key, marker) + assert row["request_id"] == chunks[0].id, row + assert_free_row(row, rig.hidden_free) + + +def test_exhausted_key_reaches_hidden_free_alias_through_anthropic_messages(rig: AliasRig) -> None: + marker: Final = "messages-sync-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + key: Final = exhausted_key(rig, scenario) + with anthropic_client(rig.gateway, key) as client: + message: Final = client.messages.create( + model=rig.hidden_responses, + max_tokens=16, + messages=[{"role": "user", "content": marker}], + extra_headers=spend_marker(marker), + ) + assert [block.text for block in message.content if block.type == "text"] == [RESPONSES_REPLY], message + assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 1 + row: Final = landed_once(key, marker) + assert row["request_id"] == message.id, row + assert_free_row(row, rig.hidden_responses) + + +async def test_exhausted_key_reaches_hidden_free_alias_through_streamed_anthropic_messages(rig: AliasRig) -> None: + marker: Final = "messages-stream-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + key: Final = exhausted_key(rig, scenario) + async with ( + async_anthropic_client(rig.gateway, key) as client, + client.messages.stream( + model=rig.hidden_responses_stream, + max_tokens=16, + messages=[{"role": "user", "content": marker}], + extra_headers=spend_marker(marker), + ) as stream, + ): + text: Final = "".join([piece async for piece in stream.text_stream]) + final: Final = await stream.get_final_message() + assert text == RESPONSES_REPLY, final + assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 1 + row: Final = landed_once(key, marker) + assert row["request_id"] == final.id, row + assert_free_row(row, rig.hidden_responses_stream) + + +def test_exhausted_key_reaches_hidden_free_alias_through_openai_responses(rig: AliasRig) -> None: + marker: Final = "responses-sync-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + key: Final = exhausted_key(rig, scenario) + with openai_client(rig.gateway, key) as client: + response: Final = client.responses.create( + model=rig.hidden_responses, input=marker, extra_headers=spend_marker(marker) + ) + assert response.output_text == RESPONSES_REPLY, response + assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 1 + row: Final = landed_once(key, marker) + assert row["request_id"] == response.id, row + assert_free_row(row, rig.hidden_responses) + + +async def test_exhausted_key_reaches_hidden_free_alias_through_streamed_openai_responses(rig: AliasRig) -> None: + marker: Final = "responses-stream-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + key: Final = exhausted_key(rig, scenario) + async with async_openai_client(rig.gateway, key) as client: + stream: Final = await client.responses.create( + model=rig.hidden_responses_stream, input=marker, stream=True, extra_headers=spend_marker(marker) + ) + events: Final = tuple([event async for event in stream]) + assert events[-1].type == "response.completed", events + assert "".join(event.delta for event in events if event.type == "response.output_text.delta") == RESPONSES_REPLY + assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 1 + assert_free_row(landed_once(key, marker), rig.hidden_responses_stream) + + +def _assert_raw_chat_served(rig: AliasRig, candidate: Gateway, key: str) -> None: + marker: Final = "raw-" + uuid.uuid4().hex + response: Final = fresh_chat(candidate, rig.hidden_free, key, marker) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == CHAT_REPLY, response.text + assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 1 + row: Final = landed_once(key, marker) + assert row["request_id"] == response.json()["id"], row + assert_free_row(row, rig.hidden_free) + + +def test_exhausted_key_reaches_hidden_free_alias_over_raw_http_on_both_replicas(rig: AliasRig) -> None: + with rig.gateway.scenario() as scenario: + key: Final = exhausted_key(rig, scenario) + _assert_raw_chat_served(rig, rig.gateway, key) + _assert_raw_chat_served(rig, rig.peer, key) + + +def _duplicate_model_post(rig: AliasRig, key: str, first: str, last: str, marker: str) -> httpx.Response: + messages: Final = json.dumps([{"role": "user", "content": marker}]) + return httpx.post( + f"{base_url(rig.gateway)}/v1/chat/completions", + content=f'{{"model": {json.dumps(first)}, "model": {json.dumps(last)}, "messages": {messages}}}'.encode(), + headers={"Authorization": f"Bearer {key}", "Content-Type": "application/json", **spend_marker(marker)}, + timeout=60, + trust_env=False, + ) + + +def test_duplicate_model_field_is_judged_by_its_last_value(rig: AliasRig) -> None: + free_marker: Final = "duplicate-free-" + uuid.uuid4().hex + paid_marker: Final = "duplicate-paid-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + key: Final = exhausted_key(rig, scenario) + served: Final = _duplicate_model_post(rig, key, rig.hidden_paid, rig.hidden_free, free_marker) + assert served.status_code == 200, served.text + refused: Final = _duplicate_model_post(rig, key, rig.hidden_free, rig.hidden_paid, paid_marker) + assert refused.status_code == BUDGET_EXCEEDED, refused.text + observed: Final = upstream_requests(rig.gateway.upstream_url) + assert upstream_hits(observed, free_marker) == 1 + assert upstream_hits(observed, paid_marker) == 0 + assert_free_row(landed_once(key, free_marker), rig.hidden_free) + + +def test_provider_failure_behind_hidden_free_alias_reaches_the_caller(rig: AliasRig) -> None: + failed_marker: Final = "provider-failure-" + uuid.uuid4().hex + recovered_marker: Final = "provider-recovered-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + key: Final = exhausted_key(rig, scenario) + scenario.cleanups.callback(clear_provider_script, rig) + script_provider(rig, 1) + failed: Final = fresh_chat(rig.gateway, rig.hidden_failing, key, failed_marker) + assert failed.status_code == 500, failed.text + assert "Controlled provider failure" in failed.text, failed.text + assert "budget" not in failed.text.lower(), failed.text + clear_provider_script(rig) + recovered: Final = fresh_chat(rig.gateway, rig.hidden_failing, key, recovered_marker) + assert recovered.status_code == 200, recovered.text + observed: Final = upstream_requests(rig.gateway.upstream_url) + assert upstream_hits(observed, failed_marker) == 1 + assert upstream_hits(observed, recovered_marker) == 1 + assert_free_row(landed_once(key, recovered_marker), rig.hidden_failing) + + +def test_exhausted_user_budget_still_reaches_hidden_free_alias(rig: AliasRig) -> None: + with rig.gateway.scenario() as scenario: + user: Final = scenario.user(max_budget=BUDGET) + key: Final = scenario.key(user_id=user, max_budget=5.0) + first: Final = fresh_chat(rig.gateway, rig.paid, key, "user-exhaust-" + uuid.uuid4().hex) + assert first.status_code == 200, first.text + settle_chat(rig, rig.paid, key, BUDGET_EXCEEDED, seconds=90) + assert chat_statuses(rig.gateway, rig.hidden_free, key, 8) == {200} + assert chat_statuses(rig.peer, rig.hidden_free, key, 8) == {200} + settle_chat(rig, rig.hidden_paid, key, BUDGET_EXCEEDED, seconds=90) + refused: Final = fresh_chat(rig.gateway, rig.hidden_paid, key, "user-refused-" + uuid.uuid4().hex) + assert f"User={user}" in refused.text, refused.text + + +def test_exhausted_team_budget_still_reaches_hidden_free_alias(rig: AliasRig) -> None: + with rig.gateway.scenario() as scenario: + team: Final = scenario.team(max_budget=BUDGET) + key: Final = scenario.key(team_id=team, max_budget=5.0) + first: Final = fresh_chat(rig.gateway, rig.paid, key, "team-exhaust-" + uuid.uuid4().hex) + assert first.status_code == 200, first.text + settle_chat(rig, rig.paid, key, BUDGET_EXCEEDED, seconds=90) + assert chat_statuses(rig.gateway, rig.hidden_free, key, 8) == {200} + assert chat_statuses(rig.peer, rig.hidden_free, key, 8) == {200} + settle_chat(rig, rig.hidden_paid, key, BUDGET_EXCEEDED, seconds=90) + refused: Final = fresh_chat(rig.gateway, rig.hidden_paid, key, "team-refused-" + uuid.uuid4().hex) + assert f"Team={team}" in refused.text, refused.text + + +def test_exhausted_tag_budget_still_reaches_hidden_free_alias(rig: AliasRig) -> None: + tag: Final = "hidden-alias-tag-" + uuid.uuid4().hex + tagged: Final[Mapping[str, JsonValue]] = {"metadata": {"tags": [tag]}} + with rig.gateway.scenario() as scenario: + rig.gateway.post("/tag/new", {"name": tag, "max_budget": BUDGET}) + scenario.cleanups.callback(rig.gateway.post, "/tag/delete", {"name": tag}) + key: Final = scenario.key(max_budget=5.0) + first: Final = fresh_chat(rig.gateway, rig.paid, key, "tag-exhaust-" + uuid.uuid4().hex, tagged) + assert first.status_code == 200, first.text + settle_chat(rig, rig.paid, key, BUDGET_EXCEEDED, seconds=90, extra=tagged) + assert chat_statuses(rig.gateway, rig.hidden_free, key, 8, tagged) == {200} + assert chat_statuses(rig.peer, rig.hidden_free, key, 8, tagged) == {200} + settle_chat(rig, rig.hidden_paid, key, BUDGET_EXCEEDED, seconds=90, extra=tagged) + refused: Final = fresh_chat(rig.gateway, rig.hidden_paid, key, "tag-refused-" + uuid.uuid4().hex, tagged) + assert f"Tag={tag}" in refused.text, refused.text + untagged: Final = fresh_chat(rig.gateway, rig.hidden_paid, key, "tag-untagged-" + uuid.uuid4().hex) + assert untagged.status_code == 200, untagged.text + + +def test_hidden_alias_repointed_between_paid_and_free_groups_follows_the_target(rig: AliasRig) -> None: + alias: Final = "hidden-repoint-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + install_aliases(rig.gateway, {alias: hidden(rig.paid)}) + scenario.cleanups.callback(remove_aliases, rig.gateway, frozenset({alias})) + settle_chat(rig, alias, rig.gateway.key, 200) + key: Final = exhausted_key(rig, scenario) + settle_chat(rig, alias, key, BUDGET_EXCEEDED) + install_aliases(rig.gateway, {alias: hidden(rig.free)}) + settle_chat(rig, alias, key, 200) + install_aliases(rig.gateway, {alias: hidden(rig.paid)}) + settle_chat(rig, alias, key, BUDGET_EXCEEDED) + + +def test_visible_alias_repointed_to_a_paid_group_loses_the_bypass(rig: AliasRig) -> None: + alias: Final = "visible-repoint-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + install_aliases(rig.gateway, {alias: rig.free}) + scenario.cleanups.callback(remove_aliases, rig.gateway, frozenset({alias})) + key: Final = exhausted_key(rig, scenario) + settle_chat(rig, alias, key, 200) + install_aliases(rig.gateway, {alias: rig.paid}) + settle_chat(rig, alias, key, BUDGET_EXCEEDED) + install_aliases(rig.gateway, {alias: rig.free}) + settle_chat(rig, alias, key, 200) + + +def test_failed_free_primary_falls_back_to_hidden_free_alias_for_exhausted_key(rig: AliasRig) -> None: + free_marker: Final = "fallback-free-" + uuid.uuid4().hex + paid_marker: Final = "fallback-paid-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + key: Final = exhausted_key(rig, scenario) + scenario.cleanups.callback(clear_provider_script, rig) + script_provider(rig, 2) + served: Final = fresh_chat(rig.gateway, rig.failing_free, key, free_marker, {"fallbacks": [rig.hidden_free]}) + assert served.status_code == 200, served.text + assert served.headers["x-litellm-model-group"] == rig.hidden_free, dict(served.headers) + refused: Final = fresh_chat(rig.gateway, rig.failing_free, key, paid_marker, {"fallbacks": [rig.hidden_paid]}) + assert refused.status_code == 500, refused.text + assert "Controlled provider failure" in refused.text, refused.text + observed: Final = upstream_requests(rig.gateway.upstream_url) + assert upstream_hits(observed, free_marker) == 2 + assert upstream_hits(observed, paid_marker) == 1 + assert_free_row(landed_once(key, free_marker), rig.hidden_free) + + +def test_exhausted_key_is_served_a_cached_reply_through_hidden_free_alias(rig: AliasRig) -> None: + marker: Final = "cache-twin-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + key: Final = exhausted_key(rig, scenario) + first: Final = fresh_chat(rig.gateway, rig.hidden_free, key, marker) + assert first.status_code == 200, first.text + second: Final = fresh_chat(rig.gateway, rig.hidden_free, key, marker) + assert second.status_code == 200, second.text + assert second.json()["id"] == first.json()["id"], second.text + assert "x-litellm-cache-key" in second.headers, dict(second.headers) + assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 1 + rows: Final = landed(key, marker) + assert all(float(str(row["spend"])) == 0.0 for row in rows), rows + + +def _base_config() -> Mapping[str, JsonValue]: + return JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + + +def _own_config(directory: Path, name: str, section: str, value: JsonValue) -> Path: + path: Final = directory / name + path.write_text(yaml.safe_dump({**_base_config(), section: value})) + return path + + +def _delete_proxy_budget_row(rig: AliasRig) -> None: + deleted: Final = rig.gateway.request("POST", "/user/delete", {"user_ids": [PROXY_BUDGET_USER]}) + assert deleted.status_code == 200, deleted.text + + +@pytest.mark.timeout(480) +def test_exhausted_proxy_budget_still_reaches_hidden_free_alias(rig: AliasRig, tmp_path: Path) -> None: + settings: Final = object_value(_base_config()["litellm_settings"]) + config: Final = _own_config( + tmp_path, + "proxy-budget.yaml", + "litellm_settings", + {**settings, "max_budget": BUDGET, "budget_duration": "30d"}, + ) + with rig.gateway.scenario() as scenario: + key: Final = scenario.key() + scenario.cleanups.callback(_delete_proxy_budget_row, rig) + with owned_proxy(rig.gateway, tmp_path, {}, config=config, workers=2) as candidate: + settle_candidate(candidate, rig.hidden_free, key, 200, seconds=120) + first: Final = fresh_chat(candidate, rig.paid, key, "proxy-exhaust-" + uuid.uuid4().hex) + assert first.status_code == 200, first.text + settle_candidate(candidate, rig.paid, key, BUDGET_EXCEEDED, seconds=120) + refused: Final = fresh_chat(candidate, rig.paid, key, "proxy-refused-" + uuid.uuid4().hex) + assert error_type(refused) == "budget_exceeded", refused.text + assert "Key=" not in refused.text, refused.text + assert chat_statuses(candidate, rig.hidden_free, key, 8) == {200} + settle_candidate(candidate, rig.hidden_paid, key, BUDGET_EXCEEDED, seconds=120) + unbudgeted: Final = fresh_chat(rig.gateway, rig.paid, key, "proxy-unbudgeted-" + uuid.uuid4().hex) + assert unbudgeted.status_code == 200, unbudgeted.text + + +_TAG_ADDER: Final = """from litellm.integrations.custom_guardrail import CustomGuardrail + + +class TagAdder(CustomGuardrail): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + metadata = data.setdefault("metadata", {}) + metadata["tags"] = [*(metadata.get("tags") or []), "__TAG__"] + return data +""" + + +@pytest.mark.timeout(480) +def test_guardrail_added_tag_over_budget_still_reaches_hidden_free_alias(rig: AliasRig, tmp_path: Path) -> None: + tag: Final = "hidden-alias-guardrail-tag-" + uuid.uuid4().hex + module: Final = "tag_adder_" + uuid.uuid4().hex + served_marker: Final = "guardrail-served-" + uuid.uuid4().hex + (tmp_path / f"{module}.py").write_text(_TAG_ADDER.replace("__TAG__", tag)) + config: Final = _own_config( + tmp_path, + "guardrail-tag.yaml", + "guardrails", + [ + { + "guardrail_name": "tag-adder-" + uuid.uuid4().hex, + "litellm_params": {"guardrail": f"{module}.TagAdder", "mode": "pre_call", "default_on": True}, + } + ], + ) + tagged: Final[Mapping[str, JsonValue]] = {"metadata": {"tags": [tag]}} + with rig.gateway.scenario() as scenario: + rig.gateway.post("/tag/new", {"name": tag, "max_budget": BUDGET}) + scenario.cleanups.callback(rig.gateway.post, "/tag/delete", {"name": tag}) + key: Final = scenario.key(max_budget=5.0) + first: Final = fresh_chat(rig.gateway, rig.paid, key, "guardrail-exhaust-" + uuid.uuid4().hex, tagged) + assert first.status_code == 200, first.text + settle_chat(rig, rig.paid, key, BUDGET_EXCEEDED, seconds=90, extra=tagged) + untagged: Final = fresh_chat(rig.gateway, rig.hidden_paid, key, "guardrail-untagged-" + uuid.uuid4().hex) + assert untagged.status_code == 200, untagged.text + with owned_proxy(rig.gateway, tmp_path, {}, config=config, workers=2) as candidate: + settle_candidate(candidate, rig.hidden_paid, key, BUDGET_EXCEEDED, seconds=120) + refused: Final = fresh_chat(candidate, rig.hidden_paid, key, "guardrail-refused-" + uuid.uuid4().hex) + assert f"Tag={tag}" in refused.text, refused.text + assert chat_statuses(candidate, rig.hidden_free, key, 8) == {200} + served: Final = fresh_chat(candidate, rig.hidden_free, key, served_marker) + assert served.status_code == 200, served.text + row: Final = landed_once(key, served_marker) + assert row["request_id"] == served.json()["id"], row + assert_free_row(row, rig.hidden_free) + + +def _assert_free_alias_served(rig: AliasRig, candidate: Gateway, alias: str, key: str, prefix: str) -> None: + marker: Final = f"{prefix}-" + uuid.uuid4().hex + response: Final = fresh_chat(candidate, alias, key, marker) + assert response.status_code == 200, response.text + assert_free_row(landed_once(key, marker), alias) + + +def test_exhausted_key_reaches_visible_free_alias(rig: AliasRig) -> None: + with rig.gateway.scenario() as scenario: + key: Final = exhausted_key(rig, scenario) + _assert_free_alias_served(rig, rig.gateway, rig.visible_free, key, "visible") + _assert_free_alias_served(rig, rig.peer, rig.visible_free, key, "visible") + + +def _assert_hidden_paid_refused(rig: AliasRig, candidate: Gateway, key: str) -> None: + marker: Final = "hidden-paid-" + uuid.uuid4().hex + response: Final = fresh_chat(candidate, rig.hidden_paid, key, marker) + assert response.status_code == BUDGET_EXCEEDED, response.text + assert error_type(response) == "budget_exceeded", response.text + assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 0 + + +def test_exhausted_key_is_refused_on_hidden_paid_alias(rig: AliasRig) -> None: + with rig.gateway.scenario() as scenario: + key: Final = exhausted_key(rig, scenario) + _assert_hidden_paid_refused(rig, rig.gateway, key) + _assert_hidden_paid_refused(rig, rig.peer, key) + + +def test_exhausted_key_reaches_free_group_by_its_own_name(rig: AliasRig) -> None: + marker: Final = "plain-free-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + key: Final = exhausted_key(rig, scenario) + response: Final = fresh_chat(rig.gateway, rig.free, key, marker) + assert response.status_code == 200, response.text + assert_free_row(landed_once(key, marker), rig.free) + + +def test_key_with_headroom_is_billed_through_hidden_paid_alias(rig: AliasRig) -> None: + paid_marker: Final = "headroom-paid-" + uuid.uuid4().hex + free_marker: Final = "headroom-free-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(max_budget=5.0) + paid: Final = fresh_chat(rig.gateway, rig.hidden_paid, key, paid_marker) + assert paid.status_code == 200, paid.text + free: Final = fresh_chat(rig.gateway, rig.hidden_free, key, free_marker) + assert free.status_code == 200, free.text + billed: Final = landed_once(key, paid_marker) + assert math.isclose(float(str(billed["spend"])), 20 * 0.001 + 20 * 0.002), billed + assert billed["model_group"] == rig.hidden_paid, billed + assert_free_row(landed_once(key, free_marker), rig.hidden_free) + + +def test_key_restricted_to_the_free_group_reaches_its_hidden_alias(rig: AliasRig) -> None: + marker: Final = "restricted-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(models=[rig.free], max_budget=5.0) + response: Final = fresh_chat(rig.gateway, rig.hidden_free, key, marker) + assert response.status_code == 200, response.text + refused: Final = fresh_chat(rig.gateway, rig.hidden_paid, key, "restricted-paid-" + uuid.uuid4().hex) + assert refused.status_code == 403, refused.text + assert error_type(refused) == "key_model_access_denied", refused.text + assert_free_row(landed_once(key, marker), rig.hidden_free) + + +@pytest.mark.parametrize("flag", ["false", "null"]) +def test_alias_with_a_non_hidden_flag_keeps_the_bypass(rig: AliasRig, flag: str) -> None: + alias: Final = {"false": rig.shown_free, "null": rig.null_hidden_free}[flag] + with rig.gateway.scenario() as scenario: + key: Final = exhausted_key(rig, scenario) + _assert_free_alias_served(rig, rig.gateway, alias, key, f"flag-{flag}") + _assert_free_alias_served(rig, rig.peer, alias, key, f"flag-{flag}") + + +def test_hidden_alias_to_a_group_priced_by_the_cost_map_stays_budgeted(rig: AliasRig) -> None: + marker: Final = "unpriced-" + uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + key: Final = exhausted_key(rig, scenario) + response: Final = fresh_chat(rig.gateway, rig.hidden_unpriced, key, marker) + assert response.status_code == BUDGET_EXCEEDED, response.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) + refused: Final = fresh_chat(rig.gateway, rig.hidden_missing, key, "missing-" + uuid.uuid4().hex) + assert refused.status_code == BUDGET_EXCEEDED, refused.text + unroutable: Final = fresh_chat(rig.gateway, rig.hidden_missing, rig.gateway.key, "missing-" + uuid.uuid4().hex) + assert unroutable.status_code == 400, unroutable.text + assert "no healthy deployments" in unroutable.text, unroutable.text + for candidate in (rig.gateway, rig.peer): + assert candidate.request("GET", "/health/liveliness").status_code == 200 + assert candidate.request("GET", "/model/info").status_code == 200 + assert candidate.request("GET", "/v1/models").status_code == 200 + served: Final = fresh_chat(rig.gateway, rig.hidden_free, key, "missing-control-" + uuid.uuid4().hex) + assert served.status_code == 200, served.text + + +@pytest.mark.parametrize( + ("shape", "status"), + [("int", BUDGET_EXCEEDED), ("list", 400), ("empty", BUDGET_EXCEEDED), ("oversized", BUDGET_EXCEEDED)], +) +def test_malformed_model_value_never_takes_the_bypass(rig: AliasRig, shape: str, status: int) -> None: + marker: Final = f"malformed-{shape}-" + uuid.uuid4().hex + models: Final[Mapping[str, JsonValue]] = { + "int": 5, + "list": [rig.hidden_free], + "empty": "", + "oversized": rig.hidden_free + "x" * 5120, + } + with rig.gateway.scenario() as scenario: + key: Final = exhausted_key(rig, scenario) + response: Final = rig.gateway.request( + "POST", + "/v1/chat/completions", + {"model": models[shape], "messages": [{"role": "user", "content": marker}]}, + key=key, + ) + assert response.status_code == status, response.text + assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 0 + served: Final = fresh_chat(rig.gateway, rig.hidden_free, key, "malformed-control-" + uuid.uuid4().hex) + assert served.status_code == 200, served.text + + +def test_unauthenticated_request_to_hidden_alias_is_rejected(rig: AliasRig) -> None: + marker: Final = "unauthenticated-" + uuid.uuid4().hex + response: Final = httpx.post( + f"{base_url(rig.gateway)}/v1/chat/completions", + json={"model": rig.hidden_free, "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 + + +def _assert_hidden_aliases_unlisted(rig: AliasRig, candidate: Gateway) -> None: + models: Final = candidate.get("/v1/models")["data"] + groups: Final = candidate.get("/model_group/info")["data"] + assert isinstance(models, list) and isinstance(groups, list) + listed: Final = frozenset(str(object_value(entry)["id"]) for entry in models) + described: Final = frozenset(str(object_value(entry)["model_group"]) for entry in groups) + assert rig.visible_free in listed and rig.visible_free in described + assert rig.shown_free in listed and rig.shown_free in described + for name in (rig.hidden_free, rig.hidden_paid, rig.hidden_responses, rig.hidden_missing): + assert name not in listed and name not in described, name + + +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) diff --git a/tests/integration/authorization/test_hidden_alias_budget_bypass_chaos.py b/tests/integration/authorization/test_hidden_alias_budget_bypass_chaos.py new file mode 100644 index 00000000000..a7da04f1364 --- /dev/null +++ b/tests/integration/authorization/test_hidden_alias_budget_bypass_chaos.py @@ -0,0 +1,208 @@ +import os +import re +import signal +import threading +import uuid +from collections.abc import Callable, Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from queue import SimpleQueue +from typing import Final + +import httpx +import pytest +from pydantic import JsonValue +from integration._support.client import Gateway, eventually +from integration._support.process import owned_proxy_process, owned_upstream +from integration.authorization._hidden_alias_budget import ( + BUDGET_EXCEEDED, + AliasRig, + alias_rig, + assert_free_row, + chat_statuses, + exhausted_key, + fresh_chat, + fresh_message, + fresh_response, + hidden, + install_aliases, + landed_all_once, + remove_aliases, + settle_candidate, + settle_chat, + upstream_hits, + upstream_requests, +) + +pytestmark: Final = pytest.mark.timeout(240) + +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") +_WAVE: Final = 8 +_STREAM: Final[Mapping[str, JsonValue]] = {"stream": True} + + +@pytest.fixture(scope="module") +def rig() -> Iterator[AliasRig]: + with alias_rig() as built: + yield built + + +def _burst_markers(prefix: str, count: int) -> tuple[str, ...]: + return tuple(f"{prefix}-{index}-{uuid.uuid4().hex}" for index in range(count)) + + +def _send_all(send: Callable[[str], httpx.Response], markers: tuple[str, ...]) -> tuple[httpx.Response, ...]: + with ThreadPoolExecutor(max_workers=len(markers)) as pool: + return tuple(pool.map(send, markers)) + + +def _mixed_call(rig: AliasRig, key: str, marker: str) -> httpx.Response: + senders: Final[Mapping[str, Callable[[], httpx.Response]]] = { + "chat": lambda: fresh_chat(rig.gateway, rig.hidden_free, key, marker), + "chatstream": lambda: fresh_chat(rig.gateway, rig.hidden_free, key, marker, _STREAM), + "messages": lambda: fresh_message(rig.gateway, rig.hidden_responses, key, marker), + "messagesstream": lambda: fresh_message(rig.gateway, rig.hidden_responses_stream, key, marker, _STREAM), + "responses": lambda: fresh_response(rig.gateway, rig.hidden_responses, key, marker), + "responsesstream": lambda: fresh_response(rig.gateway, rig.hidden_responses_stream, key, marker, _STREAM), + } + return senders[marker.split("-")[1]]() + + +_KINDS: Final = ("chat", "chatstream", "messages", "messagesstream", "responses", "responsesstream") + + +def test_mixed_concurrent_burst_through_hidden_free_aliases_lands_each_call_once(rig: AliasRig) -> None: + markers: Final = tuple(f"burst-{_KINDS[index % len(_KINDS)]}-{index}-{uuid.uuid4().hex}" for index in range(30)) + with rig.gateway.scenario() as scenario: + key: Final = exhausted_key(rig, scenario) + upstream_requests(rig.gateway.upstream_url) + responses: Final = _send_all(lambda marker: _mixed_call(rig, key, marker), markers) + assert [response.status_code for response in responses] == [200] * len(markers), [ + response.text for response in responses if response.status_code != 200 + ] + observed: Final = upstream_requests(rig.gateway.upstream_url) + assert {marker: upstream_hits(observed, marker) for marker in markers} == dict.fromkeys(markers, 1) + rows: Final = landed_all_once(key, frozenset(markers)) + assert len(rows) == len(markers), rows + for row in rows: + assert float(str(row["spend"])) == 0.0, row + assert row["status"] == "success", row + + +def test_alias_flipped_to_a_free_group_during_a_burst_only_ever_serves_or_refuses(rig: AliasRig) -> None: + alias: Final = "flip-" + uuid.uuid4().hex + seen: Final[SimpleQueue[tuple[str, int]]] = SimpleQueue() + with rig.gateway.scenario() as scenario: + install_aliases(rig.gateway, {alias: hidden(rig.paid)}) + scenario.cleanups.callback(remove_aliases, rig.gateway, frozenset({alias})) + settle_chat(rig, alias, rig.gateway.key, 200) + key: Final = exhausted_key(rig, scenario) + settle_chat(rig, alias, key, BUDGET_EXCEEDED) + upstream_requests(rig.gateway.upstream_url) + + def wave(candidate: Gateway) -> frozenset[int]: + markers: Final = _burst_markers("flip", _WAVE) + responses: Final = _send_all(lambda marker: fresh_chat(candidate, alias, key, marker), markers) + for marker, response in zip(markers, responses, strict=True): + seen.put((marker, response.status_code)) + return frozenset(response.status_code for response in responses) + + assert wave(rig.gateway) == frozenset({BUDGET_EXCEEDED}) + flip: Final = threading.Thread(target=install_aliases, args=(rig.gateway, {alias: hidden(rig.free)})) + flip.start() + eventually(lambda: wave(rig.gateway) | wave(rig.gateway), lambda found: found == frozenset({200}), seconds=60) + flip.join(timeout=30) + assert not flip.is_alive() + settle_chat(rig, alias, key, 200) + collected: Final = tuple(seen.get() for _ in range(seen.qsize())) + assert {status for _, status in collected} <= {200, BUDGET_EXCEEDED}, collected + served: Final = frozenset(marker for marker, status in collected if status == 200) + refused: Final = frozenset(marker for marker, status in collected if status == BUDGET_EXCEEDED) + assert served and refused, collected + observed: Final = upstream_requests(rig.gateway.upstream_url) + assert all(upstream_hits(observed, marker) == 1 for marker in served), collected + assert all(upstream_hits(observed, marker) == 0 for marker in refused), collected + for row in landed_all_once(key, served): + assert_free_row(row, alias) + + +def _tolerant_status(candidate: Gateway, model: str, key: str, marker: str) -> int | None: + try: + return fresh_chat(candidate, model, key, marker).status_code + except httpx.TransportError: + return None + + +@pytest.mark.timeout(480) +def test_upstream_outage_behind_hidden_free_alias_is_a_provider_error_and_recovers( + rig: AliasRig, tmp_path: Path +) -> None: + alias: Final = "hidden-outage-" + uuid.uuid4().hex + with owned_upstream(tmp_path) as slot, rig.gateway.scenario() as scenario: + group: Final = scenario.model(api_base=f"{slot.url}/v1", input_cost_per_token=0, output_cost_per_token=0) + install_aliases(rig.gateway, {alias: hidden(group)}) + scenario.cleanups.callback(remove_aliases, rig.gateway, frozenset({alias})) + settle_chat(rig, alias, rig.gateway.key, 200) + key: Final = exhausted_key(rig, scenario) + before: Final = _burst_markers("outage-before", 10) + served_before: Final = _send_all(lambda marker: fresh_chat(rig.gateway, alias, key, marker), before) + assert [response.status_code for response in served_before] == [200] * 10 + slot.stop() + during: Final = _burst_markers("outage-during", 10) + failed: Final = _send_all(lambda marker: fresh_chat(rig.gateway, alias, key, marker), during) + for response in failed: + assert response.status_code >= 500, response.text + assert "budget" not in response.text.lower(), response.text + assert rig.gateway.request("GET", "/health/liveliness").status_code == 200 + unrelated: Final = fresh_chat(rig.gateway, rig.hidden_free, key, "outage-unrelated-" + uuid.uuid4().hex) + assert unrelated.status_code == 200, unrelated.text + slot.start() + settle_candidate(rig.gateway, alias, key, 200, seconds=90) + after: Final = _burst_markers("outage-after", 10) + served_after: Final = _send_all(lambda marker: fresh_chat(rig.gateway, alias, key, marker), after) + assert [response.status_code for response in served_after] == [200] * 10 + observed: Final = upstream_requests(slot.url) + assert {marker: upstream_hits(observed, marker) for marker in after} == dict.fromkeys(after, 1) + for row in landed_all_once(key, frozenset(before + after)): + assert_free_row(row, alias) + + +def _worker_startups(log: Path) -> tuple[tuple[int, ...], int]: + text: Final = log.read_text() + started: Final = tuple(int(found[1]) for found in _STARTED_WORKER.finditer(text)) + return started, text.count("Application startup complete.") + + +@pytest.mark.timeout(480) +def test_killed_worker_leaves_the_sibling_serving_hidden_free_aliases(rig: AliasRig, tmp_path: Path) -> None: + with rig.gateway.scenario() as scenario: + key: Final = exhausted_key(rig, scenario) + with owned_proxy_process(rig.gateway, tmp_path, {}, workers=2) as owned: + candidate: Final = owned.gateway + workers, _ = eventually( + lambda: _worker_startups(owned.log), + lambda found: len(found[0]) == 2 and found[1] == 2, + seconds=120, + ) + settle_candidate(candidate, rig.hidden_free, key, 200, seconds=120) + settle_candidate(candidate, rig.hidden_paid, key, BUDGET_EXCEEDED, seconds=120) + os.kill(workers[0], signal.SIGKILL) + eventually( + lambda: _tolerant_status(candidate, rig.hidden_free, key, "kill-probe-" + uuid.uuid4().hex), + lambda found: found == 200, + seconds=60, + ) + assert chat_statuses(candidate, rig.hidden_free, key, 8) == {200} + assert chat_statuses(candidate, rig.hidden_paid, key, 8) == {BUDGET_EXCEEDED} + eventually( + lambda: _worker_startups(owned.log), + lambda found: len(found[0]) == 3 and found[1] == 3, + seconds=180, + ) + settle_candidate(candidate, rig.hidden_free, key, 200, seconds=120) + settle_candidate(candidate, rig.hidden_paid, key, BUDGET_EXCEEDED, seconds=120) + markers: Final = _burst_markers("kill-after", 10) + served: Final = _send_all(lambda marker: fresh_chat(candidate, rig.hidden_free, key, marker), markers) + assert [response.status_code for response in served] == [200] * 10 + for row in landed_all_once(key, frozenset(markers)): + assert_free_row(row, rig.hidden_free) 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 7665008a6a6..515b6e45dc4 100644 --- a/tests/unit/proxy/auth/test_unmapped_model_budget_enforcement.py +++ b/tests/unit/proxy/auth/test_unmapped_model_budget_enforcement.py @@ -9,6 +9,8 @@ See: https://github.com/BerriAI/litellm/issues/24770 import copy +import pytest + import litellm from litellm.proxy.auth.auth_checks import _is_model_cost_zero from litellm.router import Router @@ -39,10 +41,7 @@ class TestUnmappedModelBudgetEnforcement: ] ) result = _is_model_cost_zero(model="custom-model", llm_router=router) - assert result is False, ( - "Unmapped model should enforce budget (return False), " - "not bypass it (return True)" - ) + assert result is False, "Unmapped model should enforce budget (return False), not bypass it (return True)" def test_explicitly_free_model_bypasses_budget(self): """A model with explicit cost=0 in model_info should bypass budget.""" @@ -65,9 +64,7 @@ class TestUnmappedModelBudgetEnforcement: ] ) result = _is_model_cost_zero(model="free-model", llm_router=router) - assert ( - result is True - ), "Explicitly free model should bypass budget (return True)" + assert result is True, "Explicitly free model should bypass budget (return True)" def test_known_paid_model_enforces_budget(self): """A model in the cost map with non-zero costs should enforce budget.""" @@ -101,9 +98,7 @@ class TestUnmappedModelBudgetEnforcement: ] ) result = _is_model_cost_zero(model="free-via-params", llm_router=router) - assert ( - result is True - ), "Model with explicit cost=0 in litellm_params should bypass budget" + assert result is True, "Model with explicit cost=0 in litellm_params should bypass budget" def test_cache_invalidates_on_in_place_pricing_update(self): """ @@ -285,9 +280,12 @@ class TestUnmappedModelBudgetEnforcement: "An aliased PTU group must not be read as free" ) - def test_hidden_model_group_alias_enforces_budget(self): - """A hidden alias keeps budget enforced: get_model_group_info() returns None for it, - so the cost is unknown before the configuration gate is reached.""" + def test_hidden_model_group_alias_to_free_model_bypasses_budget(self): + """A hidden alias to an explicitly free group bypasses budget, like the group itself. + + ``get_model_group_info`` returns None for hidden aliases, so the alias must be + resolved to its target group before the cost lookup. + """ router = Router( model_list=[ { @@ -304,7 +302,22 @@ class TestUnmappedModelBudgetEnforcement: model_group_alias={"hidden-alias": {"model": "free-model", "hidden": True}}, ) - assert _is_model_cost_zero(model="hidden-alias", llm_router=router) is False + assert _is_model_cost_zero(model="hidden-alias", llm_router=router) is True + + def test_hidden_model_group_alias_to_paid_model_enforces_budget(self): + """A hidden alias to a priced group keeps budget enforced.""" + router = Router( + model_list=[ + { + "model_name": "paid-model", + "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "sk-fake"}, + "model_info": {"id": "paid-model-id"}, + }, + ], + model_group_alias={"hidden-paid-alias": {"model": "paid-model", "hidden": True}}, + ) + + assert _is_model_cost_zero(model="hidden-paid-alias", llm_router=router) is False def test_dangling_model_group_alias_enforces_budget(self): """An alias pointing at a group that does not exist keeps budget enforced.""" @@ -326,6 +339,115 @@ class TestUnmappedModelBudgetEnforcement: assert _is_model_cost_zero(model="dangling-alias", llm_router=router) is False + def test_repointed_hidden_alias_does_not_reuse_cached_free_result(self): + """Repointing a hidden alias from a free group to a paid group re-evaluates the cost. + + ``Router.update_settings`` is the one runtime path that rewrites the alias map (the + proxy's config update applies ``router_settings`` through it), so the cached verdict + has to drop there. + """ + router = Router( + model_list=[ + { + "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_name": "paid-model", + "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "sk-fake"}, + "model_info": {"id": "paid-model-id"}, + }, + ], + model_group_alias={"hidden-alias": {"model": "free-model", "hidden": True}}, + ) + + assert _is_model_cost_zero(model="hidden-alias", llm_router=router) is True + router.update_settings(model_group_alias={"hidden-alias": {"model": "paid-model", "hidden": True}}) + 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. + + 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. + """ + router = Router( + model_list=[ + { + "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_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_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} + + assert [_is_model_cost_zero(model=name, llm_router=router) for name in order] == [ + expected[name] for name in order + ] + assert [_is_model_cost_zero(model=name, llm_router=router) for name in order] == [ + expected[name] for name in order + ], "the cached verdicts must match the first evaluation" + + 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. + + The router serves ``chain-smart`` with the real ``chain-legacy`` deployment, which is priced, + so following the second hop to the free group would waive the budget for a paid call. + """ + router = Router( + model_list=[ + { + "model_name": "chain-legacy", + "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": "chain-legacy-id"}, + }, + { + "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={"chain-smart": "chain-legacy", "chain-legacy": "free-model"}, + ) + + assert _is_model_cost_zero(model="chain-smart", llm_router=router) is False + 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 diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 823429d7ed6..5d9b37e24fc 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -2278,6 +2278,73 @@ def test_model_group_info_cost_none_for_unpriced_deployment_but_zero_when_declar assert priced.output_cost_per_token is not None and priced.output_cost_per_token > 0 +def _alias_cost_router() -> Router: + return Router( + model_list=[ + { + "model_name": "vllm-free", + "litellm_params": { + "model": "openai/my-vllm-free", + "api_key": "fake", + "api_base": "http://localhost:8000/v1", + "input_cost_per_token": 0, + "output_cost_per_token": 0, + }, + }, + { + "model_name": "gpt-priced", + "litellm_params": {"model": "gpt-4o", "api_key": "fake"}, + }, + ], + model_group_alias={"hidden-free": {"model": "vllm-free", "hidden": True}, "visible": "vllm-free"}, + ) + + +def test_get_model_group_info_include_hidden_resolves_a_hidden_alias(): + router = _alias_cost_router() + + assert router.get_model_group_info(model_group="hidden-free") is None + + hidden: Final = router.get_model_group_info(model_group="hidden-free", include_hidden=True) + assert hidden is not None + assert hidden.model_group == "hidden-free" + assert hidden.input_cost_per_token == 0 + assert hidden.output_cost_per_token == 0 + + +def test_update_settings_model_group_alias_drops_cached_group_info(): + router = _alias_cost_router() + before: Final = router.cached_model_group_info("visible") + assert before is not None and before.input_cost_per_token == 0 + + router.update_settings(model_group_alias={"visible": "gpt-priced"}) + + after: Final = router.cached_model_group_info("visible") + assert after is not None + assert after.input_cost_per_token is not None and after.input_cost_per_token > 0 + + +def test_switch_routing_strategy_installs_lar1_then_restores_the_default_selector(): + router = _alias_cost_router() + + router._switch_routing_strategy( + "lar1", + { + "routing_strategy_args": { + "confidence_threshold_low": 0.1, + "confidence_threshold_medium": 0.3, + "confidence_threshold_high": 0.9, + } + }, + ) + assert router.routing_strategy == "lar1" + assert "async_get_available_deployment" in router.__dict__ + + router._switch_routing_strategy("usage-based-routing-v2", {}) + assert router.lowesttpm_logger_v2 is not None + assert "async_get_available_deployment" not in router.__dict__ + + @pytest.mark.parametrize( "value,expected", [