From bc43b959629db2a9f1efee5e507a5c44fd43458b Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 9 Nov 2026 00:38:23 +0000 Subject: [PATCH] test(auth): freeze independent Phase 1 budget references --- .../auth/phase1_budget_auth_reference.py | 915 ++++++++++++++++++ .../auth/phase1_budget_context_reference.py | 352 +++++++ .../auth/phase1_budget_route_reference.py | 151 +++ .../auth/phase1_budget_throttle_reference.py | 45 + .../proxy/auth/test_budget_context_adapter.py | 10 +- .../proxy/auth/test_budget_pricing_adapter.py | 2 +- .../proxy/auth/test_common_budget_adapter.py | 5 +- .../auth/test_end_user_context_adapter.py | 35 +- .../auth/test_key_budget_context_adapter.py | 4 +- 9 files changed, 1502 insertions(+), 17 deletions(-) create mode 100644 tests/unit/proxy/auth/phase1_budget_auth_reference.py create mode 100644 tests/unit/proxy/auth/phase1_budget_context_reference.py create mode 100644 tests/unit/proxy/auth/phase1_budget_route_reference.py create mode 100644 tests/unit/proxy/auth/phase1_budget_throttle_reference.py diff --git a/tests/unit/proxy/auth/phase1_budget_auth_reference.py b/tests/unit/proxy/auth/phase1_budget_auth_reference.py new file mode 100644 index 00000000000..42c4673d874 --- /dev/null +++ b/tests/unit/proxy/auth/phase1_budget_auth_reference.py @@ -0,0 +1,915 @@ +"""Policy frozen from 8bd7f5089a7500d9c915a5670c1b314ee051f9c9, with independently supplied raw IO.""" + +from __future__ import annotations + +import asyncio +import json +import math +import re +from collections.abc import Mapping, Sequence +from datetime import datetime, timedelta, timezone +from types import MappingProxyType +from typing import Final, Optional, Protocol + +from opentelemetry.trace import Span +from pydantic import ConfigDict, TypeAdapter + +import litellm +from litellm._internal_context import with_service_target +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.decisions_pricing import decisions_pricing_fields, decisions_request_price +from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.proxy._types import ( + MODEL_LIST_ROUTES, + CallInfo, + LiteLLM_BudgetTable, + LiteLLM_EndUserTable, + Litellm_EntityType, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LiteLLM_UserTable, + UserAPIKeyAuth, +) +from litellm.proxy.auth.auth_checks import ( + _dictable_table, # pyright: ignore[reportPrivateUsage] # data-only Prisma row projection +) +from litellm.proxy.common_utils.user_api_key_cache import ( + AUTH_OBJECTS_TARGET, + UserApiKeyCache, + get_management_object_ttl, +) +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.repositories.budget_repository import BudgetRepository +from litellm.router import Router +from litellm.router_api import RegistryReader +from litellm.types.router import DeploymentTypedDict + +from .phase1_budget_route_reference import RouteChecks +from .phase1_budget_throttle_reference import budget_throttle_percentage, should_throttle_budget_exceeded + + +class _RawCacheRead(Protocol): + async def async_get_cache(self, *, key: str, local_only: bool = ...) -> object: ... + + +def _raw_cache(cache: _RawCacheRead) -> _RawCacheRead: + return cache + + +def _get_router_zero_cost_cache(llm_router: RegistryReader) -> dict[str, bool] | None: # mutable-ok: Phase 1 public collection and identity contract + + cache: Final = llm_router.registry().zero_cost() + return cache if isinstance(cache, dict) else None + + +def is_model_cost_zero(model: str | list[str] | None, llm_router: RegistryReader | None) -> bool: # mutable-ok: Phase 1 public collection and identity contract + + if model is None or llm_router is None: + return False + + model_list: Final = [model] if isinstance(model, str) else model + + zero_cost_cache: Final = _get_router_zero_cost_cache(llm_router) + + for model_name in model_list: + if zero_cost_cache is not None: + cached = zero_cost_cache.get(model_name) + if cached is not None: + if cached is False: + return False + continue + try: + model_group_info = llm_router.registry().group(model_name, include_hidden=True) + + if model_group_info is None: + verbose_proxy_logger.debug("No model group info found for %s, assuming it has cost", model_name) + if zero_cost_cache is not None: + zero_cost_cache[model_name] = False + return False + + input_cost = model_group_info.input_cost_per_token + output_cost = model_group_info.output_cost_per_token + + if input_cost is None or output_cost is None: + verbose_proxy_logger.debug( + "Model %s has undefined cost (input: %s, output: %s), assuming it has cost", + model_name, + input_cost, + output_cost, + ) + if zero_cost_cache is not None: + zero_cost_cache[model_name] = False + return False + + if input_cost > 0 or output_cost > 0: + verbose_proxy_logger.debug( + "Model %s has non-zero cost (input: %s, output: %s)", model_name, input_cost, output_cost + ) + if zero_cost_cache is not None: + zero_cost_cache[model_name] = False + return False + + safe_name = str(model_name).replace("\n", "").replace("\r", "") + 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 " + "cost (enforce budget)", + safe_name, + ) + if zero_cost_cache is not None: + zero_cost_cache[model_name] = False + return False + + 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)", + safe_name, + ) + if zero_cost_cache is not None: + zero_cost_cache[model_name] = False + return False + + verbose_proxy_logger.debug( + "Model %s has zero cost explicitly configured (input: %s, output: %s)", + safe_name, + input_cost, + output_cost, + ) + if zero_cost_cache is not None: + zero_cost_cache[model_name] = True + + except Exception as e: + verbose_proxy_logger.debug("Error checking cost for model %s: %s, assuming it has cost", model_name, e) + return False + + return True + + +_NO_MODEL_INFO: Final[Mapping[str, object]] = MappingProxyType({}) + + +def _deployments_served_for(model: str, llm_router: RegistryReader) -> Sequence[DeploymentTypedDict]: + + registry: Final = llm_router.registry() + routed_group: Final = registry.resolve(model) + if registry.resolve(routed_group) == routed_group: + return registry.deployments(routed_group) or () + named: Final = registry.deployments(model, source="alias") + if named: + return named + return tuple( + DeploymentTypedDict(**deployment) # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + for deployment in registry.patterns().get_deployments_by_pattern(model=routed_group) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + ) + + +def _has_ptu_flat_cost(deployments: Sequence[DeploymentTypedDict]) -> bool: + + for deployment in deployments: + model_info = deployment.get("model_info") or _NO_MODEL_INFO # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + if model_info.get("ptu_count") is not None and model_info.get("cost_per_ptu_per_hour") is not None: # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + return True + return False + + +def _is_cost_explicitly_configured(deployments: Sequence[DeploymentTypedDict]) -> bool: + + for deployment in deployments: + model_id = (deployment.get("model_info") or _EMPTY_COST_ENTRY).get("id") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + if model_id is None: + continue + raw_entry = litellm.model_cost.get(model_id, _EMPTY_COST_ENTRY) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + if "input_cost_per_token" in raw_entry or "output_cost_per_token" in raw_entry: + return True + return False + + +def _has_explicit_positive_cost(deployments: Sequence[DeploymentTypedDict]) -> bool: + + for deployment in deployments: + model_id = (deployment.get("model_info") or _EMPTY_COST_ENTRY).get("id") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + if model_id is None: + continue + raw_entry = as_str_mapping(litellm.model_cost.get(model_id)) or _EMPTY_COST_ENTRY # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + if ( + _is_positive_cost(raw_entry.get("input_cost_per_token")) + or _is_positive_cost(raw_entry.get("output_cost_per_token")) + or _is_positive_cost(_deployment_query_price(deployment["litellm_params"], raw_entry)) + ): + return True + return False + + +def _deployment_query_price(litellm_params: Mapping[str, object], cost_entry: Mapping[str, object]) -> float | None: + provider: Final = litellm_params.get("custom_llm_provider") or str(litellm_params.get("model")).partition("/")[0] + return decisions_request_price({**cost_entry, "litellm_provider": provider}, decisions_pricing_fields(cost_entry)) + + +_EMPTY_COST_ENTRY: Final[Mapping[str, object]] = MappingProxyType({}) + + +def _is_positive_cost(value: object) -> bool: + return isinstance(value, (int, float)) and not isinstance(value, bool) and value > 0 + + +MODEL_DISCOVERY_ROUTES: Final = frozenset( + { + *MODEL_LIST_ROUTES, + "/model/info", + "/v1/model/info", + "/v2/model/info", + "/model_group/info", + "/utils/model_info", + } +) + + +BUDGET_ENFORCED_SIDE_EFFECT_ROUTES: Final = frozenset( + { + "/health", + "/health/services", + "/health/test_connection", + "/auto_router/test_routing", + } +) + + +MCP_DISCOVERY_ROUTES: Final = frozenset({"/mcp-rest/tools/list", "/v1/mcp/tools"}) + + +MCP_ZERO_SPEND_JSONRPC_METHODS: Final = frozenset({"initialize", "notifications/initialized", "ping", "tools/list"}) + + +MCP_TOOL_CALL_ROUTES: Final = frozenset({"/mcp/tools/call", "/mcp-rest/tools/call"}) + + +def is_mcp_discovery_request(route: str, request_body: Mapping[str, object]) -> bool: + if route in MCP_DISCOVERY_ROUTES: + return True + if route in MCP_TOOL_CALL_ROUTES or not (route == "/mcp" or route.startswith("/mcp/")): + return False + method: Final = request_body.get("method") + return isinstance(method, str) and method in MCP_ZERO_SPEND_JSONRPC_METHODS + + +def route_skips_budget_checks(route: str) -> bool: + if route in BUDGET_ENFORCED_SIDE_EFFECT_ROUTES: + return False + if route in MODEL_DISCOVERY_ROUTES: + return True + if RouteChecks.is_llm_api_route(route=route): + return False + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( # noqa: PLC0415 # pass-through endpoints import the auth stack + InitPassThroughEndpointHelpers, + ) + + return not InitPassThroughEndpointHelpers.is_registered_pass_through_route(route=route) + + +def request_skips_budget_checks( + route: str, + model: str | list[str] | None, # mutable-ok: Phase 1 public collection and identity contract + llm_router: Router | None, + *, + request_body: Mapping[str, object] | None = None, +) -> bool: + if route_skips_budget_checks(route=route): + return True + if request_body is not None and is_mcp_discovery_request(route=route, request_body=request_body): + return True + return is_model_cost_zero(model=model, llm_router=llm_router) + + +KEY_END_USER_BUDGET_ID_METADATA_FIELD: Final = "end_user_budget_id" + + +def get_key_end_user_budget_id(key_metadata: Mapping[str, object] | None) -> str | None: + + if key_metadata is None: + return None + budget_id: Final = key_metadata.get(KEY_END_USER_BUDGET_ID_METADATA_FIELD) + return budget_id if isinstance(budget_id, str) and budget_id != "" else None + + +@with_service_target(AUTH_OBJECTS_TARGET) +async def get_default_end_user_budget( + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None = None, + budget_id: str | None = None, +) -> LiteLLM_BudgetTable | None: + + default_budget_id: Final = budget_id if budget_id is not None else litellm.max_end_user_budget_id + if prisma_client is None or default_budget_id is None: + return None + + cache_key: Final = f"default_end_user_budget:{default_budget_id}" + + cached_budget: Final = await user_api_key_cache.async_get_cache( + key=cache_key, + model_type=LiteLLM_BudgetTable, + ) + if cached_budget is not None: + return cached_budget + + try: + budget_record: Final = await _dictable_table(BudgetRepository(prisma_client), "budget").find_unique( + where={"budget_id": default_budget_id} + ) + + if budget_record is None: + verbose_proxy_logger.warning( + "Default end user budget not found in database: %s", + default_budget_id.replace("\r", "").replace("\n", ""), + ) + return None + + _budget_obj: Final = LiteLLM_BudgetTable.model_validate(budget_record.dict()) + + await user_api_key_cache.async_set_cache( + key=cache_key, + value=_budget_obj, + model_type=LiteLLM_BudgetTable, + ttl=get_management_object_ttl(user_api_key_cache), + ) + + return _budget_obj + + except Exception as e: + verbose_proxy_logger.error("Error fetching default end user budget: %s", e) + return None + + +async def get_team_member_default_budget( + budget_id: str, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + raise_on_lookup_error: bool = False, +) -> LiteLLM_BudgetTable | None: + raise AssertionError("Frozen reference raw IO must be supplied by the test: get_team_member_default_budget") + + +def resolve_effective_team_member_budget( + member_budget: LiteLLM_BudgetTable | None, + team_default_budget: LiteLLM_BudgetTable | None, + now: datetime, +) -> float | None: + if member_budget is not None and member_budget.max_budget is not None: + return _effective_member_max(member_budget, now) + if team_default_budget is None or team_default_budget.max_budget is None or team_default_budget.max_budget <= 0: + return None + return team_default_budget.max_budget + ( + _active_member_increase(member_budget, now) if member_budget is not None else 0.0 + ) + + +async def resolve_default_end_user_budget( + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + key_end_user_budget_id: str | None, + parent_otel_span: Span | None = None, +) -> LiteLLM_BudgetTable | None: + + if key_end_user_budget_id is not None: + key_budget: Final = await get_default_end_user_budget( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + budget_id=key_end_user_budget_id, + ) + if key_budget is not None: + return key_budget + + if litellm.max_end_user_budget_id is None: + return None + + return await get_default_end_user_budget( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + ) + + +async def _apply_default_budget_to_end_user( + end_user_obj: LiteLLM_EndUserTable, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None = None, + key_end_user_budget_id: str | None = None, +) -> LiteLLM_EndUserTable: + + if end_user_obj.budget_id is not None and end_user_obj.litellm_budget_table is not None: + return end_user_obj + + if key_end_user_budget_id is None and litellm.max_end_user_budget_id is None: + return end_user_obj + + default_budget: Final = await resolve_default_end_user_budget( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + key_end_user_budget_id=key_end_user_budget_id, + parent_otel_span=parent_otel_span, + ) + + if default_budget is None: + return end_user_obj + + verbose_proxy_logger.debug( + "Applied default budget %s to end user %s", default_budget.budget_id, end_user_obj.user_id + ) + return end_user_obj.model_copy(update=MappingProxyType({"litellm_budget_table": default_budget})) + + +async def check_end_user_budget( + end_user_obj: LiteLLM_EndUserTable, + route: str, +) -> None: + + if RouteChecks.is_info_route(route): + return + + if end_user_obj.litellm_budget_table is None: + return + + end_user_budget: Final = end_user_obj.litellm_budget_table.max_budget + if end_user_budget is None: + return + + from litellm.proxy.proxy_server import get_current_spend + + end_user_spend: Final = await get_current_spend( + counter_key=f"spend:end_user:{end_user_obj.user_id}", + fallback_spend=end_user_obj.spend or 0.0, + max_budget=end_user_budget, + fallback_authoritative=True, + ) + if end_user_spend > end_user_budget: + raise litellm.BudgetExceededError( + current_cost=end_user_spend, + max_budget=end_user_budget, + message=f"ExceededBudget: End User={end_user_obj.user_id} over budget. Spend={end_user_spend}, Budget={end_user_budget}", + entity_type=Litellm_EntityType.END_USER.value, + entity_id=end_user_obj.user_id, + ) + + +async def _load_end_user_restricted_registry( + prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache +) -> frozenset[str] | None: + raise AssertionError("Frozen reference raw IO must be supplied by the test: _load_end_user_restricted_registry") + + +async def _end_user_is_known_unrestricted( + end_user_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + token_end_user_max_budget: float | None, + key_end_user_budget_id: str | None = None, +) -> bool: + + if ( + litellm.max_end_user_budget_id is not None + or key_end_user_budget_id is not None + or litellm.validate_end_user_id_in_db + or token_end_user_max_budget is not None + ): + return False + + registry: Final = await _load_end_user_restricted_registry( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + return registry is not None and end_user_id not in registry + + +async def get_end_user_object( + end_user_id: str | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + route: str | None = "", + parent_otel_span: Span | None = None, + proxy_logging_obj: ProxyLogging | None = None, + token_end_user_max_budget: float | None = None, + key_end_user_budget_id: str | None = None, +) -> LiteLLM_EndUserTable | None: + raise AssertionError("Frozen reference raw IO must be supplied by the test: get_end_user_object") + + +_END_USER_VALIDATION_NEGATIVE_TTL: Final = 60 + + +_END_USER_VALIDATION_POSITIVE_TTL: Final = 300 + + +@with_service_target(AUTH_OBJECTS_TARGET) +async def resolve_and_validate_end_user_id( + raw_end_user_id: str | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None = None, + proxy_logging_obj: ProxyLogging | None = None, + route: str = "", + key_end_user_budget_id: str | None = None, +) -> str | None: + + if raw_end_user_id is None: + return None + if not litellm.validate_end_user_id_in_db: + return raw_end_user_id + if prisma_client is None: + return raw_end_user_id + + has_default_budget: Final = bool(litellm.max_end_user_budget_id) or key_end_user_budget_id is not None + cache_key: Final = f"end_user_validation:{raw_end_user_id}" + cached: Final = await _raw_cache(user_api_key_cache).async_get_cache(key=cache_key) + if cached == "valid": + return raw_end_user_id + if cached == "invalid": + return raw_end_user_id if has_default_budget else None + + is_valid: Final = await _end_user_id_exists_in_db( + end_user_id=raw_end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + route=route, + ) + + await user_api_key_cache.async_set_cache( + key=cache_key, + value="valid" if is_valid else "invalid", + ttl=(_END_USER_VALIDATION_POSITIVE_TTL if is_valid else _END_USER_VALIDATION_NEGATIVE_TTL), + ) + + return raw_end_user_id if is_valid or has_default_budget else None + + +async def _end_user_id_exists_in_db( + end_user_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None = None, + proxy_logging_obj: ProxyLogging | None = None, + route: str = "", +) -> bool: + + try: + end_user_obj: Final = await get_end_user_object( + end_user_id=end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + route=route, + ) + if end_user_obj is not None: + return True + except Exception as e: + verbose_proxy_logger.debug("end_user validation: get_end_user_object lookup failed: %s", e) + + try: + user_obj: Final = await get_user_object( + user_id=end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + check_db_only=False, + user_email=end_user_id if "@" in end_user_id else None, + ) + if user_obj is not None: + return True + except Exception as e: + verbose_proxy_logger.debug("end_user validation: get_user_object lookup failed: %s", e) + + return False + + +async def get_team_membership( + user_id: str, + team_id: str, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None = None, + proxy_logging_obj: ProxyLogging | None = None, + *, + check_db_only: bool = False, +) -> Optional["LiteLLM_TeamMembership"]: + raise AssertionError("Frozen reference raw IO must be supplied by the test: get_team_membership") + + +async def get_user_object( + user_id: str | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + user_id_upsert: bool, + parent_otel_span: Span | None = None, + proxy_logging_obj: ProxyLogging | None = None, + sso_user_id: str | None = None, + user_email: str | None = None, + check_db_only: bool | None = None, +) -> LiteLLM_UserTable | None: + raise AssertionError("Frozen reference raw IO must be supplied by the test: get_user_object") + + +_MASKED_KEY_NAME_RE: Final = re.compile(r"sk-\.\.\.(?:[^\s\x00-\x1f\x7f-\x9f]{4})?") + + +def _apply_budget_exceeded_throttle(valid_token: UserAPIKeyAuth) -> bool: + + pct: Final = budget_throttle_percentage() + if pct is None or not should_throttle_budget_exceeded(valid_token): + return False + valid_token.budget_throttle_pct = pct # rebind-ok: preserve Phase 1 mutation order and identity + return True + + +def send_token_budget_alert( + valid_token: UserAPIKeyAuth, + spend: float, + proxy_logging_obj: ProxyLogging, + user_obj: LiteLLM_UserTable | None, +) -> None: + call_info: Final = CallInfo( + token=valid_token.token, + spend=spend, + max_budget=valid_token.max_budget, + soft_budget=valid_token.soft_budget, + user_id=valid_token.user_id, + team_id=valid_token.team_id, + organization_id=valid_token.org_id, + user_email=user_obj.user_email if user_obj is not None else None, + key_alias=valid_token.key_alias, + event_group=Litellm_EntityType.KEY, + ) + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="token_budget", + user_info=call_info, + ) + ) + + +def key_budget_exceeded_error( + valid_token: UserAPIKeyAuth, spend: float, max_budget: float +) -> litellm.BudgetExceededError: + + key_label: Final = valid_token.key_alias or "key" + key_descriptor: Final = ( + f"{key_label} ({valid_token.key_name})" + if valid_token.key_name and _MASKED_KEY_NAME_RE.fullmatch(valid_token.key_name) + else key_label + ) + return litellm.BudgetExceededError( + current_cost=spend, + max_budget=max_budget, + message=f"Budget has been exceeded! Key={key_descriptor} Current cost: {spend}, Max budget: {max_budget}", + entity_type=Litellm_EntityType.KEY.value, + entity_id=valid_token.token, + ) + + +async def virtual_key_max_budget_check( + valid_token: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, + user_obj: LiteLLM_UserTable | None = None, +) -> None: + + if valid_token.max_budget is not None: + from litellm.proxy.proxy_server import get_current_spend + + spend: Final = await get_current_spend( + counter_key=f"spend:key:{valid_token.token}", + fallback_spend=valid_token.spend or 0.0, + max_budget=valid_token.max_budget, + ) + send_token_budget_alert(valid_token, spend, proxy_logging_obj, user_obj) + + if math.isfinite(valid_token.max_budget) and spend >= valid_token.max_budget: + if _apply_budget_exceeded_throttle(valid_token): + return + raise key_budget_exceeded_error(valid_token, spend, valid_token.max_budget) + + +async def _virtual_key_multi_budget_check( + valid_token: UserAPIKeyAuth, +): + + if not valid_token.budget_limits: # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + return + + from litellm.proxy.proxy_server import get_current_spend + + for window in valid_token.budget_limits: # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + w: dict = window if isinstance(window, dict) else window.model_dump() # pyright: ignore[reportMissingTypeArgument, reportUnknownVariableType, reportUnnecessaryIsInstance] # frozen Phase 1 accepts unvalidated legacy values # mutable-ok: Phase 1 public collection and identity contract + counter_key = f"spend:key:{valid_token.token}:window:{w['budget_duration']}" + window_spend = await get_current_spend( + counter_key=counter_key, + fallback_spend=0.0, + max_budget=w["max_budget"], # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + window_entity_type="Key", + window_entity_id=valid_token.token, + window_duration=str(w["budget_duration"]), # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + window_start=get_budget_window_start(w), # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + ) + if math.isfinite(w["max_budget"]) and window_spend >= w["max_budget"]: # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + raise litellm.BudgetExceededError( + current_cost=window_spend, + max_budget=w["max_budget"], # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + message=( + f"ExceededBudget: Key over {w['budget_duration']} budget. " + f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}" + ), + entity_type=Litellm_EntityType.KEY.value, + entity_id=valid_token.token, + ) + + +async def team_max_budget_check( + team_object: LiteLLM_TeamTable | None, + valid_token: UserAPIKeyAuth | None, + proxy_logging_obj: ProxyLogging, +) -> None: + + if team_object is not None and team_object.max_budget is not None: + from litellm.proxy.proxy_server import get_current_spend + + spend: Final = await get_current_spend( + counter_key=f"spend:team:{team_object.team_id}", + fallback_spend=team_object.spend or 0.0, + max_budget=team_object.max_budget, + ) + + if math.isfinite(team_object.max_budget) and spend > team_object.max_budget: + if valid_token: + call_info: Final = CallInfo( + token=valid_token.token, + spend=spend, + max_budget=team_object.max_budget, + user_id=valid_token.user_id, + team_id=valid_token.team_id, + team_alias=valid_token.team_alias, + organization_id=valid_token.org_id, + event_group=Litellm_EntityType.TEAM, + ) + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="team_budget", + user_info=call_info, + ) + ) + + raise litellm.BudgetExceededError( + current_cost=spend, + max_budget=team_object.max_budget, + message=f"Budget has been exceeded! Team={team_object.team_id} Current cost: {spend}, Max budget: {team_object.max_budget}", + entity_type=Litellm_EntityType.TEAM.value, + entity_id=team_object.team_id, + ) + + +async def _team_multi_budget_check( + team_object: LiteLLM_TeamTable | None, +): + + if team_object is None or not team_object.budget_limits: + return + + from litellm.proxy.proxy_server import get_current_spend + + for window in team_object.budget_limits: + w: dict = window if isinstance(window, dict) else window.model_dump() # pyright: ignore[reportMissingTypeArgument] # frozen Phase 1 accepts unvalidated legacy values # mutable-ok: Phase 1 public collection and identity contract + counter_key = f"spend:team:{team_object.team_id}:window:{w['budget_duration']}" + window_spend = await get_current_spend( + counter_key=counter_key, + fallback_spend=0.0, + max_budget=w["max_budget"], # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + window_entity_type="Team", + window_entity_id=team_object.team_id, + window_duration=str(w["budget_duration"]), # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + window_start=get_budget_window_start(w), # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + ) + if math.isfinite(w["max_budget"]) and window_spend >= w["max_budget"]: # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + raise litellm.BudgetExceededError( + current_cost=window_spend, + max_budget=w["max_budget"], # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + message=( + f"ExceededBudget: Team={team_object.team_id} over {w['budget_duration']} budget. " + f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}" + ), + entity_type=Litellm_EntityType.TEAM.value, + entity_id=team_object.team_id, + ) + + +class ModelMaxBudgetHolder(Protocol): + @property + def model_max_budget(self) -> Mapping[str, object] | None: ... + + +def model_max_budget_of(holder: ModelMaxBudgetHolder) -> Mapping[str, object] | None: + return holder.model_max_budget + + +def _active_member_increase(self: LiteLLM_BudgetTable, now: datetime) -> float: + if self.temp_budget_increase is None or self.temp_budget_expiry is None: + return 0.0 + expiry: Final = ( + self.temp_budget_expiry.replace(tzinfo=timezone.utc) + if self.temp_budget_expiry.tzinfo is None + else self.temp_budget_expiry + ) + return 0.0 if expiry <= now else self.temp_budget_increase + + +def _effective_member_max(self: LiteLLM_BudgetTable, now: datetime) -> float | None: + if self.max_budget is None: + return None + return self.max_budget + _active_member_increase(self, now) + + +_WINDOW_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + + +def _coerce_window(window: object) -> Mapping[str, object]: + if isinstance(window, Mapping): + return window # pyright: ignore[reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + if isinstance(window, str): + try: + return _WINDOW_OBJECT.validate_python(json.loads(window)) + except Exception: + return {} + model_dump: Final = getattr(window, "model_dump", None) + if not callable(model_dump): + return {} + dumped: Final[object] = model_dump() + return dumped if isinstance(dumped, Mapping) else {} # pyright: ignore[reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + + +def get_budget_window_start(window: object) -> datetime | None: + window_dict: Final = _coerce_window(window) + budget_duration: Final = window_dict.get("budget_duration") + if budget_duration is None: + return None + try: + duration_seconds: Final = duration_in_seconds(str(budget_duration)) + except Exception: + return None + + reset_at = _coerce_datetime(window_dict.get("reset_at")) # rebind-ok: preserve Phase 1 mutation order and identity + if reset_at is None: + return datetime.now(timezone.utc) - timedelta(seconds=duration_seconds) + if reset_at.tzinfo is None: + reset_at = reset_at.replace(tzinfo=timezone.utc) # rebind-ok: preserve Phase 1 mutation order and identity + return reset_at - timedelta(seconds=duration_seconds) + + +def _coerce_datetime(value: object) -> datetime | None: + if value is None: + return None + if isinstance(value, datetime): + return value + if isinstance(value, str): + try: + return datetime.fromisoformat(value.replace("Z", "+00:00")) + except ValueError: + return None + return None + + +__all__ = ( + "_raw_cache", + "_get_router_zero_cost_cache", + "_deployments_served_for", + "_has_ptu_flat_cost", + "_is_cost_explicitly_configured", + "_has_explicit_positive_cost", + "_deployment_query_price", + "_is_positive_cost", + "_apply_default_budget_to_end_user", + "_load_end_user_restricted_registry", + "_end_user_is_known_unrestricted", + "_end_user_id_exists_in_db", + "_apply_budget_exceeded_throttle", + "_virtual_key_multi_budget_check", + "_team_multi_budget_check", + "_active_member_increase", + "_effective_member_max", + "_coerce_window", + "_coerce_datetime", +) diff --git a/tests/unit/proxy/auth/phase1_budget_context_reference.py b/tests/unit/proxy/auth/phase1_budget_context_reference.py new file mode 100644 index 00000000000..46900a23d09 --- /dev/null +++ b/tests/unit/proxy/auth/phase1_budget_context_reference.py @@ -0,0 +1,352 @@ +"""Policy frozen from 8bd7f5089a7500d9c915a5670c1b314ee051f9c9, with independently supplied raw IO.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from datetime import datetime, timezone +from typing import Final, Protocol, cast + +from opentelemetry.trace import Span + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.models.end_user import LiteLLM_EndUserTable +from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_TeamTable, UserAPIKeyAuth +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging + +from .phase1_budget_auth_reference import ( + get_end_user_object, + get_key_end_user_budget_id, + get_team_member_default_budget, + get_team_membership, + model_max_budget_of, + resolve_default_end_user_budget, +) + + +class _TeamModelBudgetLimiter(Protocol): + async def is_team_within_model_budget( + self, + team_id: str, + team_model_max_budget: Mapping[str, object], + key_model_max_budget: Mapping[str, object] | None, + model: str, + ) -> bool: ... + + +class _EndUserModelBudgetLimiter(Protocol): + async def is_end_user_within_model_budget( + self, + end_user_id: str, + end_user_model_max_budget: Mapping[str, object], + model: str, + team_id: str | None = None, + ) -> bool: ... + + +async def _check_team_model_budget( + valid_token: UserAPIKeyAuth, + model_max_budget_limiter: _TeamModelBudgetLimiter, + models: list[str], # mutable-ok: Phase 1 public collection and identity contract +) -> None: + + team_model_max_budget: Final = valid_token.team_model_max_budget + if valid_token.team_id is None or not team_model_max_budget: + return + key_model_max_budget: Final[Mapping[str, object] | None] = valid_token.model_max_budget # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + for model_name in models: + await model_max_budget_limiter.is_team_within_model_budget( + team_id=valid_token.team_id, + team_model_max_budget=team_model_max_budget, + key_model_max_budget=key_model_max_budget, # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values + model=model_name, + ) + + +async def _end_user_model_max_budget( + end_user_id: str | None, + end_user_object: LiteLLM_EndUserTable | None, + key_end_user_budget_id: str | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, +) -> Mapping[str, object] | None: + budget_table: Final = ( + end_user_object.litellm_budget_table + if end_user_object is not None + else ( + await resolve_default_end_user_budget( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + key_end_user_budget_id=key_end_user_budget_id, + parent_otel_span=parent_otel_span, + ) + if end_user_id and prisma_client is not None + else None + ) + ) + return model_max_budget_of(budget_table) if budget_table is not None else None + + +async def _check_end_user_model_budget( + end_user_id: str | None, + end_user_model_max_budget: Mapping[str, object] | None, + model_max_budget_limiter: _EndUserModelBudgetLimiter, + models: Sequence[str], + team_id: str | None, +) -> None: + if not end_user_id or not end_user_model_max_budget: + return + for model_name in models: + await model_max_budget_limiter.is_end_user_within_model_budget( + end_user_id=end_user_id, + end_user_model_max_budget=end_user_model_max_budget, + model=model_name, + team_id=team_id, + ) + + +def _apply_budget_limits_to_end_user_params( + end_user_params: dict, # pyright: ignore[reportMissingTypeArgument, reportUnknownParameterType] # frozen Phase 1 accepts unvalidated legacy values # mutable-ok: Phase 1 public collection and identity contract + budget_info: LiteLLM_BudgetTable, + end_user_id: str | None, +) -> None: + + if budget_info.tpm_limit is not None: + end_user_params["end_user_tpm_limit"] = budget_info.tpm_limit # rebind-ok: preserve Phase 1 mutation order and identity + + if budget_info.rpm_limit is not None: + end_user_params["end_user_rpm_limit"] = budget_info.rpm_limit # rebind-ok: preserve Phase 1 mutation order and identity + + if budget_info.tpd_limit is not None: + end_user_params["end_user_tpd_limit"] = budget_info.tpd_limit # rebind-ok: preserve Phase 1 mutation order and identity + + if budget_info.max_budget is not None: + end_user_params["end_user_max_budget"] = budget_info.max_budget # rebind-ok: preserve Phase 1 mutation order and identity + + if budget_info.model_max_budget is not None: # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + end_user_params["end_user_model_max_budget"] = budget_info.model_max_budget # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values # rebind-ok: preserve Phase 1 mutation order and identity + + verbose_proxy_logger.debug("Applied budget limits to end user %s", end_user_id) + + +def update_valid_token_with_end_user_params(valid_token: UserAPIKeyAuth, end_user_params: dict) -> UserAPIKeyAuth: # pyright: ignore[reportMissingTypeArgument, reportUnknownParameterType] # frozen Phase 1 accepts unvalidated legacy values # mutable-ok: Phase 1 public collection and identity contract + valid_token.end_user_id = end_user_params.get("end_user_id") # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values # rebind-ok: preserve Phase 1 mutation order and identity + + if end_user_params.get("end_user_tpm_limit") is not None: # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + valid_token.end_user_tpm_limit = end_user_params["end_user_tpm_limit"] # rebind-ok: preserve Phase 1 mutation order and identity + if end_user_params.get("end_user_rpm_limit") is not None: # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + valid_token.end_user_rpm_limit = end_user_params["end_user_rpm_limit"] # rebind-ok: preserve Phase 1 mutation order and identity + if end_user_params.get("end_user_tpd_limit") is not None: # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + valid_token.end_user_tpd_limit = end_user_params["end_user_tpd_limit"] # rebind-ok: preserve Phase 1 mutation order and identity + if end_user_params.get("allowed_model_region") is not None: # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + valid_token.allowed_model_region = end_user_params["allowed_model_region"] # rebind-ok: preserve Phase 1 mutation order and identity + if end_user_params.get("end_user_model_max_budget") is not None: # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + valid_token.end_user_model_max_budget = end_user_params["end_user_model_max_budget"] # rebind-ok: preserve Phase 1 mutation order and identity + return valid_token + + +async def _load_team_member_default_model_budget( + user_api_key_auth_obj: UserAPIKeyAuth, + team_object: LiteLLM_TeamTable | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + skip_budget_checks: bool, + proxy_logging_obj: ProxyLogging | None, +) -> None: + team_metadata_value: Final[object] = getattr(team_object, "metadata", None) if team_object is not None else None + team_metadata: Final[Mapping[str, object] | None] = ( + cast( # cast-ok: team metadata is a string-keyed JSON object + Mapping[str, object], team_metadata_value + ) + if isinstance(team_metadata_value, Mapping) + else None + ) + team_member_budget_id: Final[object | None] = ( + team_metadata.get("team_member_budget_id") if isinstance(team_metadata, Mapping) else None + ) + if ( + not isinstance(user_api_key_auth_obj.user_id, str) + or not isinstance(user_api_key_auth_obj.team_id, str) + or prisma_client is None + ): + return + + try: + membership: Final = await get_team_membership( + user_id=user_api_key_auth_obj.user_id, + team_id=user_api_key_auth_obj.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + member_budget_row: Final = membership.litellm_budget_table if membership is not None else None + member_budget_value: Final[object] = ( + getattr(member_budget_row, "model_max_budget", None) if member_budget_row is not None else None + ) + if ( + membership is not None + and membership.budget_id is not None + and membership.budget_id != team_member_budget_id + and isinstance(member_budget_value, Mapping) + and member_budget_value + ): + user_api_key_auth_obj.team_member_model_max_budget = ( # rebind-ok: auth context carries member budget downstream + member_budget_value + ) + return + if not isinstance(team_member_budget_id, str): + return + team_member_default_budget: Final = await get_team_member_default_budget( + budget_id=team_member_budget_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + raise_on_lookup_error=True, + ) + except Exception: + if not skip_budget_checks: + raise + verbose_proxy_logger.warning("Skipping team-member model budget lookup after a lookup error", exc_info=True) + return + + default_budget_value: Final[object] = ( + getattr(team_member_default_budget, "model_max_budget", None) + if team_member_default_budget is not None + else None + ) + user_api_key_auth_obj.team_member_model_max_budget = ( # rebind-ok: auth context carries default budget downstream + default_budget_value if isinstance(default_budget_value, Mapping) else None + ) + + +async def _apply_key_end_user_default_budget_to_token( + valid_token: UserAPIKeyAuth, + end_user_object: LiteLLM_EndUserTable | None, + key_end_user_budget_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + keep_token_limits: bool, +) -> None: + + default_budget: Final = ( + end_user_object.litellm_budget_table + if end_user_object is not None + else await resolve_default_end_user_budget( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + key_end_user_budget_id=key_end_user_budget_id, + parent_otel_span=parent_otel_span, + ) + ) + if default_budget is None: + return + + if not keep_token_limits or valid_token.end_user_max_budget is None: + valid_token.end_user_max_budget = default_budget.max_budget # rebind-ok: preserve Phase 1 mutation order and identity + if not keep_token_limits or valid_token.end_user_tpm_limit is None: + valid_token.end_user_tpm_limit = default_budget.tpm_limit # rebind-ok: preserve Phase 1 mutation order and identity + if not keep_token_limits or valid_token.end_user_rpm_limit is None: + valid_token.end_user_rpm_limit = default_budget.rpm_limit # rebind-ok: preserve Phase 1 mutation order and identity + if not keep_token_limits or valid_token.end_user_tpd_limit is None: + valid_token.end_user_tpd_limit = default_budget.tpd_limit # rebind-ok: preserve Phase 1 mutation order and identity + if not keep_token_limits or valid_token.end_user_model_max_budget is None: + valid_token.end_user_model_max_budget = default_budget.model_max_budget # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values # rebind-ok: preserve Phase 1 mutation order and identity + + +def _get_temp_budget_increase(valid_token: UserAPIKeyAuth): # pyright: ignore[reportUnknownParameterType] # frozen Phase 1 accepts unvalidated legacy values + valid_token_metadata: Final = valid_token.metadata # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + if "temp_budget_increase" in valid_token_metadata and "temp_budget_expiry" in valid_token_metadata: + expiry = datetime.fromisoformat(valid_token_metadata["temp_budget_expiry"]) # pyright: ignore[reportUnknownArgumentType] # frozen Phase 1 accepts unvalidated legacy values # rebind-ok: preserve Phase 1 mutation order and identity + if expiry.tzinfo is None: + expiry = expiry.replace(tzinfo=timezone.utc) # rebind-ok: preserve Phase 1 mutation order and identity + if expiry > datetime.now(timezone.utc): + return valid_token_metadata["temp_budget_increase"] # pyright: ignore[reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + return None + + +def _update_key_budget_with_temp_budget_increase( + valid_token: UserAPIKeyAuth, +) -> UserAPIKeyAuth: + if valid_token.max_budget is None: + return valid_token + temp_budget_increase: Final = _get_temp_budget_increase(valid_token) # pyright: ignore[reportUnknownVariableType] # frozen Phase 1 accepts unvalidated legacy values + if not temp_budget_increase: + return valid_token + return valid_token.model_copy(update={"max_budget": valid_token.max_budget + temp_budget_increase}) + + +async def _lookup_end_user_and_apply_budget( + valid_token: UserAPIKeyAuth, + route: str, + parent_otel_span: Span | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging | None, +): + + end_user_object = None # rebind-ok: preserve Phase 1 mutation order and identity + key_end_user_budget_id: Final = get_key_end_user_budget_id(valid_token.metadata) # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + try: + end_user_object = await get_end_user_object( # rebind-ok: preserve Phase 1 mutation order and identity + end_user_id=valid_token.end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + route=route, + token_end_user_max_budget=valid_token.end_user_max_budget, + key_end_user_budget_id=key_end_user_budget_id, + ) + if end_user_object is not None: + end_user_params = { # rebind-ok: preserve Phase 1 mutation order and identity + "end_user_id": valid_token.end_user_id, + "allowed_model_region": end_user_object.allowed_model_region, + } + if end_user_object.litellm_budget_table is not None: + _apply_budget_limits_to_end_user_params( + end_user_params=end_user_params, + budget_info=end_user_object.litellm_budget_table, + end_user_id=valid_token.end_user_id or "", + ) + valid_token = update_valid_token_with_end_user_params( # rebind-ok: preserve Phase 1 mutation order and identity + valid_token=valid_token, end_user_params=end_user_params + ) + elif key_end_user_budget_id is not None or litellm.max_end_user_budget_id is not None: + default_budget: Final = await resolve_default_end_user_budget( + prisma_client=prisma_client, # pyright: ignore[reportArgumentType] # frozen Phase 1 forwards the optional raw client + user_api_key_cache=user_api_key_cache, + key_end_user_budget_id=key_end_user_budget_id, + parent_otel_span=parent_otel_span, + ) + if default_budget is not None: + end_user_params = {"end_user_id": valid_token.end_user_id} # rebind-ok: preserve Phase 1 mutation order and identity + _apply_budget_limits_to_end_user_params( + end_user_params=end_user_params, + budget_info=default_budget, + end_user_id=valid_token.end_user_id or "", + ) + valid_token = update_valid_token_with_end_user_params( # rebind-ok: preserve Phase 1 mutation order and identity + valid_token=valid_token, end_user_params=end_user_params + ) + if valid_token.end_user_max_budget is None: + valid_token.end_user_max_budget = default_budget.max_budget # rebind-ok: preserve Phase 1 mutation order and identity + except Exception as e: + if isinstance(e, litellm.BudgetExceededError): + raise e + verbose_proxy_logger.debug("Unable to find user in db. Error - %s", e) + return valid_token, end_user_object + + +__all__ = ( + "_check_team_model_budget", + "_end_user_model_max_budget", + "_check_end_user_model_budget", + "_apply_budget_limits_to_end_user_params", + "_load_team_member_default_model_budget", + "_apply_key_end_user_default_budget_to_token", + "_get_temp_budget_increase", + "_update_key_budget_with_temp_budget_increase", + "_lookup_end_user_and_apply_budget", +) diff --git a/tests/unit/proxy/auth/phase1_budget_route_reference.py b/tests/unit/proxy/auth/phase1_budget_route_reference.py new file mode 100644 index 00000000000..f409731dc7e --- /dev/null +++ b/tests/unit/proxy/auth/phase1_budget_route_reference.py @@ -0,0 +1,151 @@ +"""Route classification frozen from 8bd7f5089a7500d9c915a5670c1b314ee051f9c9 for independent budget references.""" + +from __future__ import annotations + +import re +from collections.abc import Collection +from typing import Final + +from litellm.proxy._types import LiteLLMRoutes + +_EU_JEV_DECISION_ROUTE: Final = "/eu_jev/v1/systemone" + + +class RouteChecks: + @staticmethod + def is_llm_api_route(route: str) -> bool: + + if not isinstance(route, str): # pyright: ignore[reportUnnecessaryIsInstance] # frozen Phase 1 accepts unvalidated legacy values + return False + + if route in LiteLLMRoutes.openai_routes.value: + return True + + if route in LiteLLMRoutes.anthropic_routes.value: + return True + + if route in LiteLLMRoutes.google_routes.value: + return True + + if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.mcp_inference_routes.value): + return True + + if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.agent_inference_routes.value): + return True + + if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.litellm_native_routes.value): + return True + + for openai_route in LiteLLMRoutes.openai_routes.value: + if "{" in openai_route: + if RouteChecks._route_matches_pattern(route=route, pattern=openai_route): + return True + + if RouteChecks.is_wildcard_pattern(pattern=openai_route): + if RouteChecks.route_matches_wildcard_pattern(route=route, pattern=openai_route): + return True + + for google_route in LiteLLMRoutes.google_routes.value: + if "{" in google_route: + if RouteChecks._route_matches_pattern(route=route, pattern=google_route): + return True + + for anthropic_route in LiteLLMRoutes.anthropic_routes.value: + if "{" in anthropic_route: + if RouteChecks._route_matches_pattern(route=route, pattern=anthropic_route): + return True + + if RouteChecks._is_azure_openai_route(route=route): + return True + + for _llm_passthrough_route in LiteLLMRoutes.mapped_pass_through_routes.value: + if route == _llm_passthrough_route or route.startswith(_llm_passthrough_route + "/"): + return True + return RouteChecks.is_builtin_eu_jev_decision_route(route=route) + + @staticmethod + def is_builtin_eu_jev_decision_route(route: str) -> bool: + if route.rstrip("/") != _EU_JEV_DECISION_ROUTE: + return False + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( # noqa: PLC0415 # pass_through_endpoints imports the auth stack + InitPassThroughEndpointHelpers, + ) + + return not InitPassThroughEndpointHelpers.has_configured_route_ignoring_trailing_slash( + route=_EU_JEV_DECISION_ROUTE + ) + + @staticmethod + def is_info_route(route: str) -> bool: + + return route in LiteLLMRoutes.info_routes.value + + @staticmethod + def _is_azure_openai_route(route: str) -> bool: + + if not isinstance(route, str): # pyright: ignore[reportUnnecessaryIsInstance] # frozen Phase 1 accepts unvalidated legacy values + return False + + deployment_pattern: Final = r"^/openai/deployments/[^/]+/[^/]+/chat/completions$" + engine_pattern: Final = r"^/engines/[^/]+/chat/completions$" + + if re.match(deployment_pattern, route) or re.match(engine_pattern, route): + return True + return False + + @staticmethod + def _route_matches_pattern(route: str, pattern: str) -> bool: + + if not isinstance(route, str): # pyright: ignore[reportUnnecessaryIsInstance] # frozen Phase 1 accepts unvalidated legacy values + return False + + def _placeholder_to_regex(match: re.Match[str]) -> str: + placeholder: Final = match.group(0).strip("{}") + if not placeholder.endswith(":path"): + return r"[^/]+" + + return r"[^:]+" if ":" in match.string[match.end() :] else r"[\s\S]+" + + pattern = re.sub(r"\{[^}]+\}", _placeholder_to_regex, pattern) # rebind-ok: preserve Phase 1 mutation order and identity + + pattern = f"^{pattern}$" # rebind-ok: preserve Phase 1 mutation order and identity + if re.match(pattern, route): + return True + return False + + @staticmethod + def is_wildcard_pattern(pattern: str) -> bool: + + return pattern.endswith("*") + + @staticmethod + def route_matches_wildcard_pattern(route: str, pattern: str) -> bool: + + if pattern.endswith("*"): + prefix: Final = pattern[:-1] + return route.startswith(prefix) + else: + return route == pattern + + @staticmethod + def check_route_access(route: str, allowed_routes: Collection[str]) -> bool: + + if route in allowed_routes: + return True + + if any( + RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route) + for allowed_route in allowed_routes + if RouteChecks.is_wildcard_pattern(pattern=allowed_route) + ): + return True + + if any( + RouteChecks._route_matches_pattern(route=route, pattern=allowed_route) for allowed_route in allowed_routes + ): + return True + + return False + + +__all__ = () diff --git a/tests/unit/proxy/auth/phase1_budget_throttle_reference.py b/tests/unit/proxy/auth/phase1_budget_throttle_reference.py new file mode 100644 index 00000000000..5b3d8469d5a --- /dev/null +++ b/tests/unit/proxy/auth/phase1_budget_throttle_reference.py @@ -0,0 +1,45 @@ +""" +Throttle a key after it exceeds its own ``max_budget`` instead of blocking it. + +When a key opts in via ``throttle_on_budget_exceeded`` and a global +``budget_exceeded_throttle_percentage`` is configured, an over-budget key keeps +serving requests but at a reduced TPM/RPM (the configured percentage of its +configured limits). The decision (over budget + opted in) is made once during +auth; the scaling is recomputed from the key's original limits on every request +so it never compounds across requests. +""" + +import math +from typing import Final + +import litellm +from litellm.proxy._types import UserAPIKeyAuth + + +def budget_throttle_percentage() -> float | None: + + pct: Final = litellm.budget_exceeded_throttle_percentage + if not isinstance(pct, (int, float)) or isinstance(pct, bool): + return None + if not 0 < pct <= 1: + return None + return float(pct) + + +def should_throttle_budget_exceeded(valid_token: UserAPIKeyAuth) -> bool: + + if (valid_token.metadata or {}).get("throttle_on_budget_exceeded") is not True: # pyright: ignore[reportUnknownMemberType] # frozen Phase 1 accepts unvalidated legacy values + return False + if valid_token.tpm_limit is None and valid_token.rpm_limit is None: + return False + return budget_throttle_percentage() is not None + + +def throttled_limit(limit: int | None, pct: float | None) -> int | None: + + if limit is None or pct is None: + return limit + return max(1, math.floor(limit * pct)) + + +__all__ = () diff --git a/tests/unit/proxy/auth/test_budget_context_adapter.py b/tests/unit/proxy/auth/test_budget_context_adapter.py index 707b918352c..c8a1cb31fee 100644 --- a/tests/unit/proxy/auth/test_budget_context_adapter.py +++ b/tests/unit/proxy/auth/test_budget_context_adapter.py @@ -11,8 +11,9 @@ from litellm.models.budget import LiteLLM_BudgetTable from litellm.proxy._types import LiteLLM_EndUserTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, UserAPIKeyAuth from litellm.proxy.auth import auth_checks from litellm.proxy.auth import budget_context_adapter as adapter -from litellm.proxy.auth import user_api_key_auth as legacy from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from tests.unit.proxy.auth import phase1_budget_auth_reference as auth_reference +from tests.unit.proxy.auth import phase1_budget_context_reference as legacy @pytest.fixture(params=(False, True), ids=("reference", "native")) @@ -185,6 +186,7 @@ async def test_end_user_source_read_order_and_identity( ) read: Final = AsyncMock(side_effect=(budget,) if key_found else (None, budget)) monkeypatch.setattr(auth_checks, "get_default_end_user_budget", read) + monkeypatch.setattr(auth_reference, "get_default_end_user_budget", read) monkeypatch.setattr(litellm, "max_end_user_budget_id", "proxy") fn: Final = adapter.end_user_model_max_budget if native else legacy._end_user_model_max_budget # pyright: ignore[reportPrivateUsage] # Phase 1 leaf is the unchanged reference implementation value: Final = await fn(identity, end_user, "key", Mock(), UserApiKeyCache(), None) @@ -246,8 +248,10 @@ async def test_member_default_read_and_mutation_order( member_read: Final = AsyncMock(return_value=membership) default_read: Final = AsyncMock(return_value=default_budget) monkeypatch.setattr(auth_checks, "get_team_membership", member_read) + monkeypatch.setattr(auth_reference, "get_team_membership", member_read) monkeypatch.setattr(legacy, "get_team_membership", member_read) monkeypatch.setattr(auth_checks, "get_team_member_default_budget", default_read) + monkeypatch.setattr(auth_reference, "get_team_member_default_budget", default_read) monkeypatch.setattr(legacy, "get_team_member_default_budget", default_read) token: Final = UserAPIKeyAuth(user_id="", team_id="", team_member_model_max_budget={"old": {}}) prior: Final = token.team_member_model_max_budget @@ -279,8 +283,10 @@ async def test_member_lookup_error_policy( member_read: Final = AsyncMock(return_value=None, side_effect=None if during_default else error) default_read: Final = AsyncMock(side_effect=error) monkeypatch.setattr(auth_checks, "get_team_membership", member_read) + monkeypatch.setattr(auth_reference, "get_team_membership", member_read) monkeypatch.setattr(legacy, "get_team_membership", member_read) monkeypatch.setattr(auth_checks, "get_team_member_default_budget", default_read) + monkeypatch.setattr(auth_reference, "get_team_member_default_budget", default_read) monkeypatch.setattr(legacy, "get_team_member_default_budget", default_read) token: Final = UserAPIKeyAuth(user_id="u", team_id="t", team_member_model_max_budget={"old": {}}) prior: Final = token.team_member_model_max_budget @@ -303,6 +309,7 @@ async def test_member_native_failure_never_becomes_skippable_lookup_error(monkey error: Final = NativeAuthRuntimeError("native planner failed") read: Final = AsyncMock(side_effect=error) monkeypatch.setattr(auth_checks, "get_team_membership", read) + monkeypatch.setattr(auth_reference, "get_team_membership", read) with pytest.raises(NativeAuthRuntimeError, match="native planner failed") as raised: await adapter.load_team_member_default_model_budget( UserAPIKeyAuth(user_id="u", team_id="t"), None, Mock(), UserApiKeyCache(), True, None @@ -319,6 +326,7 @@ async def test_member_planner_failure_stops_before_io(monkeypatch: pytest.Monkey read: Final = AsyncMock() monkeypatch.setattr(_native, "request_auth_context_member", native) monkeypatch.setattr(auth_checks, "get_team_membership", read) + monkeypatch.setattr(auth_reference, "get_team_membership", read) with pytest.raises(NativeAuthRuntimeError, match="decode failed"): await adapter.load_team_member_default_model_budget( UserAPIKeyAuth(user_id="u", team_id="t"), None, Mock(), UserApiKeyCache(), True, None diff --git a/tests/unit/proxy/auth/test_budget_pricing_adapter.py b/tests/unit/proxy/auth/test_budget_pricing_adapter.py index b0dacfdba2a..ff5fdd62a7d 100644 --- a/tests/unit/proxy/auth/test_budget_pricing_adapter.py +++ b/tests/unit/proxy/auth/test_budget_pricing_adapter.py @@ -6,8 +6,8 @@ import pytest import litellm from litellm._auth.budget_pricing import ModelPriceFacts -from litellm.proxy.auth import auth_checks as reference from litellm.proxy.auth import budget_pricing_adapter as native +from tests.unit.proxy.auth import phase1_budget_auth_reference as reference def router(input_price: float | None, output_price: float | None, deployments: tuple[object, ...] = ()) -> Mock: diff --git a/tests/unit/proxy/auth/test_common_budget_adapter.py b/tests/unit/proxy/auth/test_common_budget_adapter.py index 2043a8426e6..ace55803f6b 100644 --- a/tests/unit/proxy/auth/test_common_budget_adapter.py +++ b/tests/unit/proxy/auth/test_common_budget_adapter.py @@ -20,12 +20,11 @@ from litellm.proxy._types import ( LiteLLM_TeamTable, UserAPIKeyAuth, ) -from litellm.proxy.auth import auth_checks as reference -from litellm.proxy.auth import budget_windows_adapter +from litellm.proxy.auth import budget_windows_adapter, common_budget_lookups from litellm.proxy.auth import common_budget_adapter as native -from litellm.proxy.auth import common_budget_lookups from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.utils import ProxyLogging +from tests.unit.proxy.auth import phase1_budget_auth_reference as reference @dataclass(frozen=True, slots=True) diff --git a/tests/unit/proxy/auth/test_end_user_context_adapter.py b/tests/unit/proxy/auth/test_end_user_context_adapter.py index 66d66804629..8eaee4db5a9 100644 --- a/tests/unit/proxy/auth/test_end_user_context_adapter.py +++ b/tests/unit/proxy/auth/test_end_user_context_adapter.py @@ -8,9 +8,11 @@ import litellm from litellm._auth.failure import NativeAuthRuntimeError from litellm.models.budget import LiteLLM_BudgetTable from litellm.proxy._types import LiteLLM_EndUserTable, LiteLLM_UserTable, UserAPIKeyAuth -from litellm.proxy.auth import auth_checks, user_api_key_auth as legacy +from litellm.proxy.auth import auth_checks from litellm.proxy.auth import end_user_context_adapter as adapter from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from tests.unit.proxy.auth import phase1_budget_auth_reference as auth_reference +from tests.unit.proxy.auth import phase1_budget_context_reference as legacy @pytest.fixture(params=(False, True), ids=("reference", "native")) @@ -22,7 +24,7 @@ def native(request: pytest.FixtureRequest) -> bool: "value,expected", ((None, None), ("", None), (" ", " "), ("budget", "budget"), (0, None), (False, None), ({}, None)) ) def test_metadata_id_keeps_only_nonempty_strings(native: bool, value: object, expected: str | None) -> None: - function: Final = adapter.get_key_end_user_budget_id if native else auth_checks.get_key_end_user_budget_id + function: Final = adapter.get_key_end_user_budget_id if native else auth_reference.get_key_end_user_budget_id assert function({"end_user_budget_id": value}) == expected assert function(None) is None @@ -39,7 +41,7 @@ async def test_default_budget_id_and_cache_selection( selected: Final = explicit if explicit is not None else global_id budget: Final = LiteLLM_BudgetTable(budget_id=selected, max_budget=0) await cache.async_set_cache(key=f"default_end_user_budget:{selected}", value=budget, model_type=LiteLLM_BudgetTable) - function: Final = adapter.get_default_end_user_budget if native else auth_checks.get_default_end_user_budget + function: Final = adapter.get_default_end_user_budget if native else auth_reference.get_default_end_user_budget result: Final = await function(Mock() if db_present else None, cache, budget_id=explicit) assert result == (budget if db_present and selected is not None else None) @@ -67,8 +69,9 @@ async def test_default_application_preserves_own_budget_and_copies_only_resolved ) load: Final = AsyncMock(return_value=default) monkeypatch.setattr(auth_checks, "get_default_end_user_budget", load) + monkeypatch.setattr(auth_reference, "get_default_end_user_budget", load) function: Final = ( - adapter.apply_default_budget_to_end_user if native else auth_checks._apply_default_budget_to_end_user # pyright: ignore[reportPrivateUsage] # Phase 1 parity reference + adapter.apply_default_budget_to_end_user if native else auth_reference._apply_default_budget_to_end_user # pyright: ignore[reportPrivateUsage] # Phase 1 parity reference ) # pyright: ignore[reportPrivateUsage] # Phase 1 parity reference result: Final = await function(row, Mock(), UserApiKeyCache(), key_end_user_budget_id=key_id) unchanged: Final = (own_id is not None and own_present) or (key_id is None and global_id is None) or not found @@ -105,7 +108,8 @@ async def test_unrestricted_registry_requires_absence_of_all_budget_and_validati monkeypatch.setattr(litellm, "validate_end_user_id_in_db", validate) load: Final = AsyncMock(return_value=ids) monkeypatch.setattr(auth_checks, "_load_end_user_restricted_registry", load) - function: Final = adapter.end_user_is_known_unrestricted if native else auth_checks._end_user_is_known_unrestricted # pyright: ignore[reportPrivateUsage] # Phase 1 parity reference + monkeypatch.setattr(auth_reference, "_load_end_user_restricted_registry", load) + function: Final = adapter.end_user_is_known_unrestricted if native else auth_reference._end_user_is_known_unrestricted # pyright: ignore[reportPrivateUsage] # Phase 1 parity reference result: Final = await function("eu", Mock(), UserApiKeyCache(), max_budget, key_id) blocked: Final = global_id is not None or key_id is not None or validate or max_budget is not None assert result is (not blocked and ids is not None and "eu" not in ids) @@ -208,9 +212,11 @@ async def test_validation_uses_raw_row_presence_and_preserves_unknown_ids_with_d end: Final = AsyncMock(return_value=LiteLLM_EndUserTable(user_id=identity, blocked=False) if end_found else None) user: Final = AsyncMock(return_value=LiteLLM_UserTable(user_id=identity) if user_found else None) monkeypatch.setattr(auth_checks, "get_end_user_object", end) + monkeypatch.setattr(auth_reference, "get_end_user_object", end) monkeypatch.setattr(auth_checks, "get_user_object", user) + monkeypatch.setattr(auth_reference, "get_user_object", user) function: Final = ( - adapter.resolve_and_validate_end_user_id if native else auth_checks.resolve_and_validate_end_user_id + adapter.resolve_and_validate_end_user_id if native else auth_reference.resolve_and_validate_end_user_id ) result: Final = await function(identity, Mock(), cache, key_end_user_budget_id=key_id) valid: Final = cached == "valid" or (cached != "invalid" and (end_found or user_found)) @@ -238,7 +244,7 @@ async def test_validation_disabled_or_without_db_skips_all_lookups( read: Final = AsyncMock(side_effect=AssertionError("cache must not be read")) cache.async_get_cache = read function: Final = ( - adapter.resolve_and_validate_end_user_id if native else auth_checks.resolve_and_validate_end_user_id + adapter.resolve_and_validate_end_user_id if native else auth_reference.resolve_and_validate_end_user_id ) assert await function(identity, Mock() if db else None, cache) == identity read.assert_not_called() @@ -270,8 +276,10 @@ async def test_lookup_uses_row_before_default_and_only_fills_missing_max_budget( end: Final = AsyncMock(return_value=row) default: Final = AsyncMock(return_value=budget if found else None) monkeypatch.setattr(auth_checks, "get_end_user_object", end) + monkeypatch.setattr(auth_reference, "get_end_user_object", end) monkeypatch.setattr(legacy, "get_end_user_object", end) monkeypatch.setattr(auth_checks, "get_default_end_user_budget", default) + monkeypatch.setattr(auth_reference, "get_default_end_user_budget", default) token: Final = UserAPIKeyAuth( end_user_id="eu", end_user_max_budget=initial_max, end_user_tpm_limit=9, metadata={"end_user_budget_id": key_id} ) @@ -308,7 +316,9 @@ async def test_lookup_io_failure_does_not_select_defaults_and_native_failure_is_ load: Final = AsyncMock(side_effect=failure) default: Final = AsyncMock(side_effect=AssertionError("no defaults after lookup error")) monkeypatch.setattr(auth_checks, "get_end_user_object", load) + monkeypatch.setattr(auth_reference, "get_end_user_object", load) monkeypatch.setattr(auth_checks, "get_default_end_user_budget", default) + monkeypatch.setattr(auth_reference, "get_default_end_user_budget", default) token: Final = UserAPIKeyAuth(end_user_id="eu", metadata={"end_user_budget_id": "key"}) if native_failure: with pytest.raises(NativeAuthRuntimeError) as raised: @@ -327,8 +337,10 @@ async def test_native_validation_failure_stops_before_user_fallback(monkeypatch: monkeypatch.setattr(litellm, "validate_end_user_id_in_db", True) failure: Final = NativeAuthRuntimeError("native") monkeypatch.setattr(auth_checks, "get_end_user_object", AsyncMock(side_effect=failure)) + monkeypatch.setattr(auth_reference, "get_end_user_object", AsyncMock(side_effect=failure)) user: Final = AsyncMock() monkeypatch.setattr(auth_checks, "get_user_object", user) + monkeypatch.setattr(auth_reference, "get_user_object", user) with pytest.raises(NativeAuthRuntimeError) as raised: await adapter.resolve_and_validate_end_user_id("eu", Mock(), UserApiKeyCache()) assert raised.value is failure @@ -345,7 +357,7 @@ async def test_default_budget_cache_miss_reads_database_and_caches_only_found_ro row: Final = LiteLLM_BudgetTable(budget_id="budget", max_budget=0) if found else None read: Final = AsyncMock(return_value=row) db: Final = Mock(db=Mock(litellm_budgettable=Mock(find_unique=read))) - function: Final = adapter.get_default_end_user_budget if native else auth_checks.get_default_end_user_budget + function: Final = adapter.get_default_end_user_budget if native else auth_reference.get_default_end_user_budget assert await function(db, cache, budget_id="budget") == row read.assert_awaited_once_with(where={"budget_id": "budget"}) assert await cache.async_get_cache(key="default_end_user_budget:budget", model_type=LiteLLM_BudgetTable) == row @@ -375,7 +387,7 @@ async def test_default_budget_cache_failure_is_not_swallowed(native: bool) -> No cache: Final = Mock(spec=UserApiKeyCache) read: Final = AsyncMock(side_effect=failure) cache.async_get_cache = read - function: Final = adapter.get_default_end_user_budget if native else auth_checks.get_default_end_user_budget + function: Final = adapter.get_default_end_user_budget if native else auth_reference.get_default_end_user_budget with pytest.raises(RuntimeError) as raised: await function(Mock(), cache, budget_id="budget") assert raised.value is failure @@ -398,8 +410,10 @@ async def test_validation_lookup_errors_select_secondary_identity_read( else AsyncMock(side_effect=RuntimeError("user read")) ) monkeypatch.setattr(auth_checks, "get_end_user_object", own) + monkeypatch.setattr(auth_reference, "get_end_user_object", own) monkeypatch.setattr(auth_checks, "get_user_object", user) - function: Final = adapter.end_user_id_exists_in_db if native else auth_checks._end_user_id_exists_in_db # pyright: ignore[reportPrivateUsage] # Phase 1 parity reference + monkeypatch.setattr(auth_reference, "get_user_object", user) + function: Final = adapter.end_user_id_exists_in_db if native else auth_reference._end_user_id_exists_in_db # pyright: ignore[reportPrivateUsage] # Phase 1 parity reference assert await function("eu", Mock(), UserApiKeyCache()) is user_found own.assert_awaited_once() user.assert_awaited_once() @@ -410,6 +424,7 @@ async def test_lookup_budget_denial_preserves_exception_identity(native: bool, m failure: Final = litellm.BudgetExceededError(current_cost=2, max_budget=1) read: Final = AsyncMock(side_effect=failure) monkeypatch.setattr(auth_checks, "get_end_user_object", read) + monkeypatch.setattr(auth_reference, "get_end_user_object", read) monkeypatch.setattr(legacy, "get_end_user_object", read) invocation: Final = ( adapter.lookup_end_user_and_apply_budget(UserAPIKeyAuth(), "", None, Mock(), UserApiKeyCache(), None) diff --git a/tests/unit/proxy/auth/test_key_budget_context_adapter.py b/tests/unit/proxy/auth/test_key_budget_context_adapter.py index cb83565419e..843f74a30f5 100644 --- a/tests/unit/proxy/auth/test_key_budget_context_adapter.py +++ b/tests/unit/proxy/auth/test_key_budget_context_adapter.py @@ -8,9 +8,9 @@ import pytest import litellm from litellm._auth import budget_context as core from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.auth import budget_throttle from litellm.proxy.auth import key_budget_context_adapter as adapter -from litellm.proxy.auth import user_api_key_auth as legacy +from tests.unit.proxy.auth import phase1_budget_context_reference as legacy +from tests.unit.proxy.auth import phase1_budget_throttle_reference as budget_throttle _NOW: Final = datetime(2025, 1, 1, tzinfo=timezone.utc)