From e2932d22c0b6d6e0a01c4346fe304525eed9bf8e Mon Sep 17 00:00:00 2001 From: Nitish Agarwal <1592163+nitishagar@users.noreply.github.com> Date: Thu, 18 Jun 2026 14:06:36 +0530 Subject: [PATCH 001/166] fix(spend-tracking): hash raw api keys before persisting to spend logs (#30670) The Request Logs "Key Hash" column binds to the SpendLogs metadata user_api_key field, which _get_spend_logs_metadata copied verbatim from request metadata without hashing. The top-level api_key column only hashed values starting with sk-. So a non-sk- passthrough provider key, or a Bearer-prefixed key on paths that do not strip it, could be persisted and shown in plaintext. A single helper now redacts both fields at the SpendLogs builder: it strips a Bearer prefix, passes through values already a sha256 hash or hashed-jwt- identifier, and otherwise hashes with hash_token. sk- keys still map to the same canonical hash, so log and spend correlation is unchanged --- .../spend_tracking/spend_tracking_utils.py | 42 ++- .../test_spend_tracking_utils.py | 270 ++++++++++++++++-- 2 files changed, 290 insertions(+), 22 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index aef06a3c668..e8aa2325a18 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -24,6 +24,7 @@ from litellm.litellm_core_utils.core_helpers import ( get_litellm_metadata_from_kwargs, reconstruct_model_name, ) +from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash from litellm.litellm_core_utils.safe_json_dumps import safe_dumps, strip_null_bytes from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error @@ -65,6 +66,24 @@ def _is_master_key(api_key: Optional[str], _master_key: Optional[str]) -> bool: return secrets.compare_digest(api_key, _master_key) +_HASHED_JWT_RE = re.compile(r"^hashed-jwt-[a-fA-F0-9]{64}$") + + +def _redact_logged_api_key( + value: str | None, *, already_hashed: bool = False +) -> str | None: + if not isinstance(value, str) or not value: + return None + stripped = re.sub(r"(?i)^bearer ", "", value) + if not stripped: + return None + if already_hashed and is_valid_sha256_hash(stripped): + return stripped + if _HASHED_JWT_RE.match(stripped): + return stripped + return hash_token(stripped) + + def _get_spend_logs_metadata( metadata: Optional[dict], applied_guardrails: Optional[List[str]] = None, @@ -121,6 +140,16 @@ def _get_spend_logs_metadata( key: metadata.get(key) for key in SpendLogsMetadata.__annotations__.keys() } ) + _raw_key = clean_metadata.get("user_api_key") + _trusted_hash = metadata.get("user_api_key_hash") if metadata else None + _already_hashed = ( + isinstance(_trusted_hash, str) + and is_valid_sha256_hash(_trusted_hash) + and _trusted_hash == _raw_key + ) + clean_metadata["user_api_key"] = _redact_logged_api_key( + _raw_key, already_hashed=_already_hashed + ) clean_metadata["applied_guardrails"] = applied_guardrails clean_metadata["batch_models"] = batch_models clean_metadata["mcp_tool_call_metadata"] = mcp_tool_call_metadata @@ -283,10 +312,13 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs "completion_tokens", 0 ) standard_logging_total_tokens = standard_logging_payload.get("total_tokens", 0) - if api_key is not None and isinstance(api_key, str): - if api_key.startswith("sk-"): - # hash the api_key - api_key = hash_token(api_key) + _trusted_hash = metadata.get("user_api_key_hash") if metadata else None + _key_already_hashed = ( + isinstance(_trusted_hash, str) + and is_valid_sha256_hash(_trusted_hash) + and _trusted_hash == api_key + ) + api_key = _redact_logged_api_key(api_key, already_hashed=_key_already_hashed) or "" if ( standard_logging_payload is not None @@ -299,8 +331,6 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs end_user_id = end_user_id or standard_logging_payload["metadata"].get( "user_api_key_end_user_id" ) - # BUG FIX: Don't overwrite api_key when standard_logging_payload is None - # The api_key was already extracted from metadata (line 243) and hashed (lines 256-259) request_tags = ( safe_dumps(metadata.get("tags", [])) if isinstance(metadata.get("tags", []), list) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 0c7511589de..67d57d2d6bb 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -7,7 +7,6 @@ from datetime import timezone from typing import Any, cast import pytest -from fastapi.testclient import TestClient sys.path.insert( 0, os.path.abspath("../../../..") @@ -15,7 +14,6 @@ sys.path.insert( from unittest.mock import AsyncMock, MagicMock, patch -import litellm from litellm.constants import ( LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, @@ -30,12 +28,14 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( _get_spend_logs_metadata, _get_vector_store_request_for_spend_logs_payload, _is_master_key, + _redact_logged_api_key, _redact_prompt_leaks_in_error_string, _sanitize_error_information_for_spend_logs, _sanitize_request_body_for_spend_logs_payload, _should_store_prompts_and_responses_in_spend_logs, get_logging_payload, ) +from litellm.proxy.utils import hash_token from litellm.types.utils import ( StandardLoggingHiddenParams, StandardLoggingMetadata, @@ -661,8 +661,6 @@ def test_get_logging_payload_api_key_preserved_when_standard_logging_payload_is_ assert payload["model"] == "openai/gpt-4.1" assert payload["user"] == "test_user" - print(f"✅ Test passed! api_key preserved: {payload['api_key']}") - @pytest.mark.asyncio @patch("litellm.proxy.proxy_server.master_key", "sk-master-key") @@ -810,18 +808,6 @@ async def test_api_key_preserved_through_failure_hook_to_database(): assert payload.get("model") == "gpt-3.5-turbo" assert payload.get("user") == "test_user" - print("\n" + "=" * 80) - print("✅ CRITICAL E2E TEST PASSED") - print("=" * 80) - print(f"Token: {data['token']}") - print(f"Payload api_key: {payload_api_key}") - print(f"Match: {data['token'] == payload_api_key}") - print("=" * 80) - print("Production incident bug is FIXED and protected:") - print("- Failed requests preserve api_key through entire flow") - print("- Both SpendLogs AND DailyUserSpend will have correct api_key") - print("=" * 80 + "\n") - @patch("litellm.proxy.proxy_server.master_key", None) @patch("litellm.proxy.proxy_server.general_settings", {}) @@ -2073,3 +2059,255 @@ def test_sanitize_error_information_redacts_pydantic_assignment_form( assert sanitized is not None assert "leaked-via-pydantic-msg" not in sanitized["error_message"] assert REDACTED_BY_LITELM_STRING in sanitized["error_message"] + + +# ── _redact_logged_api_key unit tests ────────────────────────────────────── + + +def test_redact_logged_api_key_none_returns_none(): + assert _redact_logged_api_key(None) is None + + +def test_redact_logged_api_key_empty_string_returns_none(): + assert _redact_logged_api_key("") is None + + +def test_redact_logged_api_key_sk_key_is_hashed(): + raw = "sk-1234secret" + result = _redact_logged_api_key(raw) + assert result == hash_token(raw) + assert result is not None + assert not result.startswith("sk-") + assert len(result) == 64 + + +def test_redact_logged_api_key_bearer_sk_equals_sk_hash(): + raw = "sk-1234secret" + result_plain = _redact_logged_api_key(raw) + result_bearer = _redact_logged_api_key(f"Bearer {raw}") + assert result_bearer == result_plain + + +def test_redact_logged_api_key_bearer_case_insensitive(): + raw = "sk-1234secret" + result_lower = _redact_logged_api_key(f"bearer {raw}") + result_upper = _redact_logged_api_key(f"BEARER {raw}") + expected = hash_token(raw) + assert result_lower == expected + assert result_upper == expected + + +def test_redact_logged_api_key_non_sk_raw_key_is_hashed(): + raw = "anthropic-raw-key-xyz" + result = _redact_logged_api_key(raw) + assert result is not None + assert result != raw + assert len(result) == 64 + assert result == hash_token(raw) + + +def test_redact_logged_api_key_already_valid_sha256_passes_through_with_flag(): + already_hashed = hash_token("sk-some-key") + assert len(already_hashed) == 64 + result = _redact_logged_api_key(already_hashed, already_hashed=True) + assert result == already_hashed + assert hash_token(already_hashed) != result # no double-hash + + +def test_redact_logged_api_key_sha256_without_flag_is_hashed(): + already_hashed = hash_token("sk-some-key") + assert len(already_hashed) == 64 + result = _redact_logged_api_key(already_hashed) + assert result is not None + assert result != already_hashed + assert len(result) == 64 + assert result == hash_token(already_hashed) + + +def test_redact_logged_api_key_hashed_jwt_passes_through(): + jwt_hash = "hashed-jwt-" + "a" * 64 + result = _redact_logged_api_key(jwt_hash) + assert result == jwt_hash + + +def test_redact_logged_api_key_hashed_jwt_short_suffix_is_hashed(): + short_jwt = "hashed-jwt-tooshort" + result = _redact_logged_api_key(short_jwt) + assert result is not None + assert result != short_jwt + assert len(result) == 64 + assert result == hash_token(short_jwt) + + +def test_redact_logged_api_key_bearer_only_returns_none(): + # "bearer " with nothing after stripping is equivalent to no key + assert _redact_logged_api_key("bearer ") is None + assert _redact_logged_api_key("Bearer ") is None + assert _redact_logged_api_key("BEARER ") is None + + +# ── _get_spend_logs_metadata key-hash invariant tests ───────────────────── + + +def test_get_spend_logs_metadata_sk_key_hashed(): + raw = "sk-1234secret" + meta = _get_spend_logs_metadata({"user_api_key": raw}) + assert meta["user_api_key"] == hash_token(raw) + assert meta["user_api_key"] is not None + result = meta["user_api_key"] + assert result is not None + assert not result.startswith("sk-") + assert len(result) == 64 + + +def test_get_spend_logs_metadata_bearer_sk_key_hashed_same_as_plain(): + raw = "sk-1234secret" + meta_plain = _get_spend_logs_metadata({"user_api_key": raw}) + meta_bearer = _get_spend_logs_metadata({"user_api_key": f"Bearer {raw}"}) + assert meta_bearer["user_api_key"] == meta_plain["user_api_key"] + + +def test_get_spend_logs_metadata_non_sk_raw_key_hashed(): + raw = "anthropic-raw-key-xyz" + meta = _get_spend_logs_metadata({"user_api_key": raw}) + result = meta["user_api_key"] + assert result is not None + assert result != raw + assert len(result) == 64 + + +def test_get_spend_logs_metadata_already_hashed_unchanged_with_provenance(): + already_hashed = hash_token("sk-some-key") + meta = _get_spend_logs_metadata( + {"user_api_key": already_hashed, "user_api_key_hash": already_hashed} + ) + assert meta["user_api_key"] == already_hashed + assert hash_token(already_hashed) != meta["user_api_key"] # no double-hash + + +def test_get_spend_logs_metadata_already_hashed_no_provenance_is_rehashed(): + already_hashed = hash_token("sk-some-key") + meta = _get_spend_logs_metadata({"user_api_key": already_hashed}) + assert meta["user_api_key"] != already_hashed + assert meta["user_api_key"] == hash_token(already_hashed) + + +def test_get_spend_logs_metadata_provenance_bypass_requires_hash_match(): + already_hashed = hash_token("sk-some-key") + different_hash = hash_token("sk-other-key") + meta = _get_spend_logs_metadata( + {"user_api_key": already_hashed, "user_api_key_hash": different_hash} + ) + assert meta["user_api_key"] == hash_token(already_hashed) + + +def test_get_spend_logs_metadata_hashed_jwt_unchanged(): + jwt_hash = "hashed-jwt-" + "b" * 64 + meta = _get_spend_logs_metadata({"user_api_key": jwt_hash}) + assert meta["user_api_key"] == jwt_hash + + +def test_get_spend_logs_metadata_none_key_is_none(): + meta = _get_spend_logs_metadata({"user_api_key": None}) + assert meta["user_api_key"] is None + + +# ── get_logging_payload key-hash invariant tests ─────────────────────────── + + +@patch("litellm.proxy.proxy_server.master_key", None) +@patch("litellm.proxy.proxy_server.general_settings", {}) +def test_get_logging_payload_non_sk_raw_key_both_fields_hashed(): + raw = "anthropic-raw-key-xyz" + kwargs = { + "model": "openai/gpt-4.1", + "messages": [{"role": "user", "content": "Hello"}], + "call_type": "acompletion", + "litellm_params": { + "metadata": { + "user_api_key": raw, + "user_api_key_user_id": "test_user", + "user_api_key_team_id": "test_team", + } + }, + } + payload = get_logging_payload( + kwargs=kwargs, + response_obj=Exception("error"), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + + assert payload["api_key"] != raw + assert len(payload["api_key"]) == 64 + + parsed_meta = json.loads(payload["metadata"]) + assert parsed_meta["user_api_key"] != raw + assert parsed_meta["user_api_key"] is not None + assert len(parsed_meta["user_api_key"]) == 64 + + +@patch("litellm.proxy.proxy_server.master_key", None) +@patch("litellm.proxy.proxy_server.general_settings", {}) +def test_get_logging_payload_empty_key_slp_none_is_empty_string_not_none_literal(): + kwargs = { + "model": "openai/gpt-4.1", + "messages": [{"role": "user", "content": "Hello"}], + "call_type": "acompletion", + "litellm_params": { + "metadata": { + "user_api_key_user_id": "test_user", + } + }, + } + payload = get_logging_payload( + kwargs=kwargs, + response_obj=Exception("error"), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + + assert payload["api_key"] == "", ( + f"Expected empty string but got {payload['api_key']!r}; " + "dropping _redact_logged_api_key's 'or \"\"' guard would yield 'None' here" + ) + + +def test_get_spend_logs_metadata_sibling_fields_preserved(): + raw = "anthropic-raw-key-xyz" + meta = _get_spend_logs_metadata( + { + "user_api_key": raw, + "user_api_key_alias": "my-alias", + "user_api_key_team_id": "team-123", + } + ) + assert meta["user_api_key"] == hash_token(raw) + assert meta["user_api_key_alias"] == "my-alias" + assert meta["user_api_key_team_id"] == "team-123" + + +def test_redact_logged_api_key_partial_sha256_is_hashed(): + partial_hex = "a" * 63 + result = _redact_logged_api_key(partial_hex) + assert result is not None + assert result != partial_hex + assert len(result) == 64 + assert result == hash_token(partial_hex) + + +def test_redact_logged_api_key_bearer_already_hashed_passes_through_with_flag(): + already_hashed = hash_token("sk-some-key") + assert len(already_hashed) == 64 + result = _redact_logged_api_key(f"Bearer {already_hashed}", already_hashed=True) + assert result == already_hashed + assert hash_token(already_hashed) != result + + +def test_redact_logged_api_key_bearer_sha256_without_flag_is_hashed(): + already_hashed = hash_token("sk-some-key") + assert len(already_hashed) == 64 + result = _redact_logged_api_key(f"Bearer {already_hashed}") + assert result is not None + assert result != already_hashed + assert result == hash_token(already_hashed) From b5c59b67874c2212d708f3f2e3d2b0ee6fead017 Mon Sep 17 00:00:00 2001 From: bhuvan2134686 Date: Mon, 27 Jul 2026 16:51:18 +1000 Subject: [PATCH 002/166] feat(providers): add SCX.ai as OpenAI-compatible provider --- litellm/constants.py | 2 + .../get_llm_provider_logic.py | 3 + litellm/llms/openai_like/providers.json | 12 ++ litellm/types/utils.py | 1 + .../llms/openai_like/test_scx_ai_provider.py | 130 ++++++++++++++++++ 5 files changed, 148 insertions(+) create mode 100644 tests/test_litellm/llms/openai_like/test_scx_ai_provider.py diff --git a/litellm/constants.py b/litellm/constants.py index a9edf135731..cc6d9e2b1e0 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -728,6 +728,7 @@ openai_compatible_endpoints: List = [ "https://api.libertai.io/v1", "https://pinstripes.io/v1", "https://api.meta.ai/v1", + "https://api.scx.ai/v1", ] @@ -795,6 +796,7 @@ openai_compatible_providers: List = [ "pinstripes", # Pinstripes - JSON-configured provider "darkbloom", "meta", # Meta Model API (Muse Spark) - JSON-configured provider + "scx-ai", ] openai_text_completion_compatible_providers: List = [ # providers that support `/v1/completions` "together_ai", diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 487a7b7e25f..52a4f866fbd 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -349,6 +349,9 @@ def get_llm_provider( elif endpoint == "https://api.meta.ai/v1": custom_llm_provider = "meta" dynamic_api_key = get_secret_str("META_API_KEY") + elif endpoint == "https://api.scx.ai/v1": + custom_llm_provider = "scx-ai" + dynamic_api_key = get_secret_str("SCX_API_KEY") if api_base is not None and not isinstance(api_base, str): raise Exception("api base needs to be a string. api_base={}".format(api_base)) diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index 164100d4194..d796d140878 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -183,5 +183,17 @@ "max_completion_tokens": "max_tokens" }, "supported_endpoints": ["/v1/chat/completions", "/v1/responses", "/v1/embeddings"] + }, + "scx-ai": { + "base_url": "https://api.scx.ai/v1", + "api_key_env": "SCX_API_KEY", + "api_base_env": "SCX_API_BASE", + "param_mappings": { + "max_completion_tokens": "max_tokens" + }, + "constraints": { + "temperature_max": 1.0 + }, + "supported_endpoints": ["/v1/chat/completions"] } } diff --git a/litellm/types/utils.py b/litellm/types/utils.py index e4dfac48141..feafce9ace1 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3508,6 +3508,7 @@ class LlmProviders(str, Enum): TENSORMESH = "tensormesh" LIBERTAI = "libertai" PINSTRIPES = "pinstripes" + SCX_AI = "scx-ai" DARKBLOOM = "darkbloom" META = "meta" LITELLM_AGENT = "litellm_agent" diff --git a/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py b/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py new file mode 100644 index 00000000000..bebd079646e --- /dev/null +++ b/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py @@ -0,0 +1,130 @@ +""" +Tests for SCX.ai provider configuration and integration. +""" + +import litellm + + +class TestSCXAIProviderConfig: + def test_scx_ai_in_provider_list(self): + from litellm import LlmProviders + + assert hasattr(LlmProviders, "SCX_AI") + assert LlmProviders.SCX_AI.value == "scx-ai" + assert "scx-ai" in litellm.provider_list + + def test_scx_ai_json_config_exists(self): + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + assert JSONProviderRegistry.exists("scx-ai") + + scx = JSONProviderRegistry.get("scx-ai") + assert scx is not None + assert scx.base_url == "https://api.scx.ai/v1" + assert scx.api_key_env == "SCX_API_KEY" + assert scx.param_mappings.get("max_completion_tokens") == "max_tokens" + assert scx.constraints.get("temperature_max") == 1.0 + + def test_scx_ai_in_openai_compatible_providers(self): + from litellm.constants import openai_compatible_providers + + assert "scx-ai" in openai_compatible_providers + + def test_scx_ai_provider_resolution(self): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + model, provider, api_key, api_base = get_llm_provider( + model="scx-ai/gpt-oss-120b", + custom_llm_provider=None, + api_base=None, + api_key=None, + ) + + assert model == "gpt-oss-120b" + assert provider == "scx-ai" + assert api_base == "https://api.scx.ai/v1" + + def test_scx_ai_api_base_override(self): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + model, provider, api_key, api_base = get_llm_provider( + model="scx-ai/gpt-oss-120b", + custom_llm_provider=None, + api_base="https://custom.scx.ai/v1", + api_key="sk-test", + ) + + assert provider == "scx-ai" + assert api_base == "https://custom.scx.ai/v1" + assert api_key == "sk-test" + + def test_scx_ai_url_autodetection(self): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + model, provider, api_key, api_base = get_llm_provider( + model="gpt-oss-120b", + custom_llm_provider=None, + api_base="https://api.scx.ai/v1", + api_key=None, + ) + assert provider == "scx-ai" + assert api_base == "https://api.scx.ai/v1" + + def test_scx_ai_temperature_clamped_to_max(self): + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + provider = JSONProviderRegistry.get("scx-ai") + assert provider is not None + config = create_config_class(provider)() + + optional_params = config.map_openai_params( + non_default_params={"temperature": 1.7}, + optional_params={}, + model="gpt-oss-120b", + drop_params=False, + ) + assert optional_params["temperature"] == 1.0 + + optional_params = config.map_openai_params( + non_default_params={"temperature": 0.4}, + optional_params={}, + model="gpt-oss-120b", + drop_params=False, + ) + assert optional_params["temperature"] == 0.4 + + def test_scx_ai_max_completion_tokens_mapped(self): + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + provider = JSONProviderRegistry.get("scx-ai") + assert provider is not None + config = create_config_class(provider)() + + optional_params = config.map_openai_params( + non_default_params={"max_completion_tokens": 256}, + optional_params={}, + model="gpt-oss-120b", + drop_params=False, + ) + assert optional_params["max_tokens"] == 256 + assert "max_completion_tokens" not in optional_params + + def test_scx_ai_router_config(self): + from litellm import Router + + router = Router( + model_list=[ + { + "model_name": "scx-chat", + "litellm_params": { + "model": "scx-ai/gpt-oss-120b", + "api_key": "test-key", + }, + } + ] + ) + + assert len(router.model_list) == 1 + assert router.model_list[0]["model_name"] == "scx-chat" From 912c7c7f42edecd099076cd8f279cc5621ad32e9 Mon Sep 17 00:00:00 2001 From: bhuvan2134686 Date: Mon, 27 Jul 2026 16:51:18 +1000 Subject: [PATCH 003/166] feat(providers): register scx-ai in the endpoint support matrix --- litellm/provider_endpoints_support_backup.json | 17 +++++++++++++++++ provider_endpoints_support.json | 17 +++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index dd7712aabca..a7ab8187bd3 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -2010,6 +2010,23 @@ "interactions": true } }, + "scx-ai": { + "display_name": "SCX.ai (`scx-ai`)", + "url": "https://docs.litellm.ai/docs/providers/scx-ai", + "endpoints": { + "chat_completions": true, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false + } + }, "snowflake": { "display_name": "Snowflake (`snowflake`)", "url": "https://docs.litellm.ai/docs/providers/snowflake", diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 65db63dc045..53b6f6e9fa0 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -2244,6 +2244,23 @@ "interactions": true } }, + "scx-ai": { + "display_name": "SCX.ai (`scx-ai`)", + "url": "https://docs.litellm.ai/docs/providers/scx-ai", + "endpoints": { + "chat_completions": true, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false + } + }, "snowflake": { "display_name": "Snowflake (`snowflake`)", "url": "https://docs.litellm.ai/docs/providers/snowflake", From 8f61073af80ded2e637ed0fdb8fd39bf5e73c607 Mon Sep 17 00:00:00 2001 From: bhuvan2134686 Date: Mon, 27 Jul 2026 16:51:18 +1000 Subject: [PATCH 004/166] feat(ui): add SCX.ai to the dashboard provider list with logo --- ui/litellm-dashboard/public/assets/logos/scx_ai.svg | 1 + .../src/components/provider_info_helpers.test.tsx | 11 +++++++++++ .../src/components/provider_info_helpers.tsx | 4 ++++ 3 files changed, 16 insertions(+) create mode 100644 ui/litellm-dashboard/public/assets/logos/scx_ai.svg diff --git a/ui/litellm-dashboard/public/assets/logos/scx_ai.svg b/ui/litellm-dashboard/public/assets/logos/scx_ai.svg new file mode 100644 index 00000000000..545176a945b --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/scx_ai.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx index 777cdc62987..55983b20bf4 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx @@ -62,6 +62,17 @@ describe("provider_info_helpers", () => { expect(result.logo).toBe(providerLogoMap[Providers.Groq]); }); + it("should map scx-ai slug and SCX_AI enum key to the SCX.ai display name and logo", () => { + const fromSlug = getProviderLogoAndName("scx-ai"); + expect(fromSlug.displayName).toBe(Providers.SCX_AI); + expect(fromSlug.logo).toBe(providerLogoMap[Providers.SCX_AI]); + expect(fromSlug.logo).toBeTruthy(); + + const fromEnumKey = getProviderLogoAndName("SCX_AI"); + expect(fromEnumKey.displayName).toBe(Providers.SCX_AI); + expect(fromEnumKey.logo).toBe(providerLogoMap[Providers.SCX_AI]); + }); + it("should map bedrock_mantle slug to Bedrock Mantle display name and logo", () => { const result = getProviderLogoAndName("bedrock_mantle"); expect(result.displayName).toBe(Providers.BedrockMantle); diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx index fa6b3c79230..4f5c0b6000b 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -50,6 +50,7 @@ import replicateLogo from "../../public/assets/logos/replicate.svg"; import runwayLogo from "../../public/assets/logos/runway.png"; import sambanovaLogo from "../../public/assets/logos/sambanova.svg"; import sapLogo from "../../public/assets/logos/sap.png"; +import scxAiLogo from "../../public/assets/logos/scx_ai.svg"; import snowflakeLogo from "../../public/assets/logos/snowflake.svg"; import sonioxLogo from "../../public/assets/logos/soniox.svg"; import togetheraiLogo from "../../public/assets/logos/togetherai.svg"; @@ -151,6 +152,7 @@ export enum Providers { SAGEMAKER_LEGACY = "Sagemaker", Sambanova = "Sambanova", SAP = "SAP Generative AI Hub", + SCX_AI = "SCX.ai", Snowflake = "Snowflake", Soniox = "Soniox", TEXT_COMPLETION_CODESTRAL = "Text-Completion-Codestral", @@ -260,6 +262,7 @@ export const provider_map: Record = { SageMaker: "sagemaker_chat", Sambanova: "sambanova", SAP: "sap", + SCX_AI: "scx-ai", Snowflake: "snowflake", Soniox: "soniox", TEXT_COMPLETION_CODESTRAL: "text-completion-codestral", @@ -351,6 +354,7 @@ export const providerLogoMap: Partial> = { [Providers.SAGEMAKER_LEGACY]: bedrockLogo.src, [Providers.Sambanova]: sambanovaLogo.src, [Providers.SAP]: sapLogo.src, + [Providers.SCX_AI]: scxAiLogo.src, [Providers.Snowflake]: snowflakeLogo.src, [Providers.Soniox]: sonioxLogo.src, [Providers.TEXT_COMPLETION_CODESTRAL]: mistralLogo.src, From 7f48431e22b4664d5fa0035779393e4242c3dc8e Mon Sep 17 00:00:00 2001 From: bhuvan2134686 Date: Mon, 27 Jul 2026 17:21:58 +1000 Subject: [PATCH 005/166] feat(models): add pricing and metadata for 5 scx-ai models --- ...odel_prices_and_context_window_backup.json | 66 +++++++++++++++++++ model_prices_and_context_window.json | 66 +++++++++++++++++++ .../llms/openai_like/test_scx_ai_provider.py | 40 +++++++++++ 3 files changed, 172 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 2cef600ea32..2f47643da50 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -33546,6 +33546,72 @@ "supports_vision": true, "source": "https://cloud.sambanova.ai/plans/pricing" }, + "scx-ai/gemma-4-31B-it": { + "input_cost_per_token": 3e-07, + "litellm_provider": "scx-ai", + "max_input_tokens": 131000, + "max_tokens": 131000, + "mode": "chat", + "output_cost_per_token": 9.1e-07, + "source": "https://scx.ai/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "scx-ai/gpt-oss-120b": { + "input_cost_per_token": 1.7e-07, + "litellm_provider": "scx-ai", + "max_input_tokens": 131000, + "max_tokens": 131000, + "mode": "chat", + "output_cost_per_token": 5.5e-07, + "source": "https://scx.ai/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "scx-ai/Llama-4-Maverick-17B-128E-Instruct": { + "input_cost_per_token": 5.3e-07, + "litellm_provider": "scx-ai", + "max_input_tokens": 131000, + "max_tokens": 131000, + "mode": "chat", + "output_cost_per_token": 1.62e-06, + "source": "https://scx.ai/pricing", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "scx-ai/MiniMax-M2.7": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 4.8e-07, + "litellm_provider": "scx-ai", + "max_input_tokens": 192000, + "max_tokens": 192000, + "mode": "chat", + "output_cost_per_token": 1.79e-06, + "source": "https://scx.ai/pricing", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "scx-ai/Qwen3-32B": { + "input_cost_per_token": 3.6e-07, + "litellm_provider": "scx-ai", + "max_input_tokens": 33000, + "max_tokens": 33000, + "mode": "chat", + "output_cost_per_token": 8.7e-07, + "source": "https://scx.ai/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "snowflake/claude-3-5-sonnet": { "litellm_provider": "snowflake", "max_input_tokens": 200000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index db28118d52b..8464aec769f 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -33637,6 +33637,72 @@ "supports_vision": true, "source": "https://cloud.sambanova.ai/plans/pricing" }, + "scx-ai/gemma-4-31B-it": { + "input_cost_per_token": 3e-07, + "litellm_provider": "scx-ai", + "max_input_tokens": 131000, + "max_tokens": 131000, + "mode": "chat", + "output_cost_per_token": 9.1e-07, + "source": "https://scx.ai/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "scx-ai/gpt-oss-120b": { + "input_cost_per_token": 1.7e-07, + "litellm_provider": "scx-ai", + "max_input_tokens": 131000, + "max_tokens": 131000, + "mode": "chat", + "output_cost_per_token": 5.5e-07, + "source": "https://scx.ai/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "scx-ai/Llama-4-Maverick-17B-128E-Instruct": { + "input_cost_per_token": 5.3e-07, + "litellm_provider": "scx-ai", + "max_input_tokens": 131000, + "max_tokens": 131000, + "mode": "chat", + "output_cost_per_token": 1.62e-06, + "source": "https://scx.ai/pricing", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "scx-ai/MiniMax-M2.7": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 4.8e-07, + "litellm_provider": "scx-ai", + "max_input_tokens": 192000, + "max_tokens": 192000, + "mode": "chat", + "output_cost_per_token": 1.79e-06, + "source": "https://scx.ai/pricing", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "scx-ai/Qwen3-32B": { + "input_cost_per_token": 3.6e-07, + "litellm_provider": "scx-ai", + "max_input_tokens": 33000, + "max_tokens": 33000, + "mode": "chat", + "output_cost_per_token": 8.7e-07, + "source": "https://scx.ai/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "snowflake/claude-3-5-sonnet": { "litellm_provider": "snowflake", "max_input_tokens": 200000, diff --git a/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py b/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py index bebd079646e..7994a133a83 100644 --- a/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py +++ b/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py @@ -128,3 +128,43 @@ class TestSCXAIProviderConfig: assert len(router.model_list) == 1 assert router.model_list[0]["model_name"] == "scx-chat" + + +class TestSCXAIModelMetadata: + SCX_MODELS = ( + "scx-ai/Llama-4-Maverick-17B-128E-Instruct", + "scx-ai/gemma-4-31B-it", + "scx-ai/Qwen3-32B", + "scx-ai/MiniMax-M2.7", + "scx-ai/gpt-oss-120b", + ) + VISION_MODELS = ("scx-ai/Llama-4-Maverick-17B-128E-Instruct", "scx-ai/gemma-4-31B-it") + + @staticmethod + def _load(path_parts): + import json + from pathlib import Path + + json_path = Path(__file__).parents[4].joinpath(*path_parts) + with open(json_path) as f: + return json.load(f) + + def test_scx_ai_models_registered_with_correct_metadata(self): + model_cost = self._load(("model_prices_and_context_window.json",)) + for model in self.SCX_MODELS: + info = model_cost.get(model) + assert info is not None, f"{model} missing from model_prices_and_context_window.json" + assert info["litellm_provider"] == "scx-ai" + assert info["mode"] == "chat" + assert info["input_cost_per_token"] > 0 + assert info["output_cost_per_token"] > 0 + assert info["supports_function_calling"] is True + assert info["supports_tool_choice"] is True + assert info.get("supports_vision", False) is (model in self.VISION_MODELS) + + def test_scx_ai_models_synced_to_backup(self): + model_cost = self._load(("model_prices_and_context_window.json",)) + backup = self._load(("litellm", "model_prices_and_context_window_backup.json")) + for model in self.SCX_MODELS: + assert model in backup, f"{model} missing from backup json" + assert backup[model] == model_cost[model], f"{model} differs between root and backup json" From b8e2848fcd7e3be397de758de9a0d2061e6a10d5 Mon Sep 17 00:00:00 2001 From: bhuvan2134686 Date: Mon, 27 Jul 2026 17:21:58 +1000 Subject: [PATCH 006/166] fix(ui): make SCX.ai selectable in the Add Model provider dropdown --- .../provider_create_fields.json | 28 +++++++++++++++++++ .../llms/openai_like/test_scx_ai_provider.py | 27 ++++++++++++++++++ 2 files changed, 55 insertions(+) diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index fcc6aac1c14..36db02d814e 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -2660,6 +2660,34 @@ ], "default_model_placeholder": "sap/gpt-4" }, + { + "provider": "SCX_AI", + "provider_display_name": "SCX.ai", + "litellm_provider": "scx-ai", + "credential_fields": [ + { + "key": "api_base", + "label": "API Base", + "placeholder": "https://api.scx.ai/v1", + "tooltip": null, + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "api_key", + "label": "API Key", + "placeholder": null, + "tooltip": null, + "required": true, + "field_type": "password", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "scx-ai/gpt-oss-120b" + }, { "provider": "Snowflake", "provider_display_name": "Snowflake", diff --git a/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py b/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py index 7994a133a83..92921067728 100644 --- a/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py +++ b/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py @@ -168,3 +168,30 @@ class TestSCXAIModelMetadata: for model in self.SCX_MODELS: assert model in backup, f"{model} missing from backup json" assert backup[model] == model_cost[model], f"{model} differs between root and backup json" + + +class TestSCXAIDashboardRegistration: + @staticmethod + def _provider_create_fields(): + import json + from pathlib import Path + + import litellm + + path = Path(litellm.__file__).parent / "proxy" / "public_endpoints" / "provider_create_fields.json" + with open(path) as f: + return json.load(f) + + def test_scx_ai_is_selectable_in_the_add_model_form(self): + entries = [e for e in self._provider_create_fields() if e["litellm_provider"] == "scx-ai"] + assert len(entries) == 1, "scx-ai must appear exactly once in provider_create_fields.json" + + entry = entries[0] + assert entry["provider"] == "SCX_AI" + assert entry["provider_display_name"] == "SCX.ai" + assert entry["default_model_placeholder"].startswith("scx-ai/") + + fields = {f["key"]: f for f in entry["credential_fields"]} + assert fields["api_key"]["required"] is True + assert fields["api_key"]["field_type"] == "password" + assert fields["api_base"]["required"] is False From cabbc7ebfb9e3aa84649ab184bbab5c692f40baf Mon Sep 17 00:00:00 2001 From: bhuvan2134686 Date: Tue, 28 Jul 2026 09:07:01 +1000 Subject: [PATCH 007/166] feat(ui): default the SCX.ai Add Model placeholder to MiniMax-M2.7 --- litellm/proxy/public_endpoints/provider_create_fields.json | 2 +- .../src/components/provider_info_helpers.test.tsx | 4 ++++ ui/litellm-dashboard/src/components/provider_info_helpers.tsx | 2 ++ 3 files changed, 7 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index 36db02d814e..4cfe9e4ef5e 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -2686,7 +2686,7 @@ "default_value": null } ], - "default_model_placeholder": "scx-ai/gpt-oss-120b" + "default_model_placeholder": "scx-ai/MiniMax-M2.7" }, { "provider": "Snowflake", diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx index 55983b20bf4..8a4c9dfd24d 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx @@ -172,6 +172,10 @@ describe("provider_info_helpers", () => { expect(getPlaceholder(Providers.Vertex_AI)).toBe("gemini-pro"); }); + it("should return an scx-ai model placeholder for SCX_AI provider", () => { + expect(getPlaceholder(Providers.SCX_AI)).toBe("scx-ai/MiniMax-M2.7"); + }); + it("should return claude-3-opus placeholder for Anthropic provider", () => { expect(getPlaceholder(Providers.Anthropic)).toBe("claude-3-opus"); }); diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx index 4f5c0b6000b..fa1f9c11eb3 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -443,6 +443,8 @@ export const getPlaceholder = (selectedProvider: string): string => { return "cursor/claude-4-sonnet"; } else if (selectedProvider === Providers.ZAI) { return "zai/glm-4.5"; + } else if (selectedProvider === Providers.SCX_AI) { + return "scx-ai/MiniMax-M2.7"; } else { return "gpt-3.5-turbo"; } From 8aa9d3dfe58bffa4efee69db579bfd7be5e9fc02 Mon Sep 17 00:00:00 2001 From: bhuvan2134686 Date: Fri, 7 Aug 2026 11:22:01 +1000 Subject: [PATCH 008/166] feat(models): swap SCX.ai catalog to GLM-5.2 and Qwen3.8 Max Replaces the five launch models with the two that SCX.ai now leads on. Both are live on api.scx.ai and both were verified against it for tool calling, json_object and json_schema output, reasoning, prompt caching, and, for Qwen3.8 Max, image input Pricing follows SCX's published USD rates. GLM-5.2 lands at $0.55/M input and $1.9255/M output, tracking the recent GLM-5.2 market repricing; Qwen3.8 Max at $1.815/M and $5.4461/M sits under the only other seller of that model, and is the first Qwen3.8 Max entry in the catalog Also corrects a metadata bug the removed entries carried: they set max_tokens equal to max_input_tokens, conflating the context window with the output cap. Both new entries declare a max_output_tokens of 131072, which is what the endpoint's own validator enforces The Add Model placeholder moves to scx-ai/GLM-5.2 now that MiniMax-M2.7 is no longer in the catalog --- ...odel_prices_and_context_window_backup.json | 84 ++++++------------- .../provider_create_fields.json | 2 +- model_prices_and_context_window.json | 84 ++++++------------- .../llms/openai_like/test_scx_ai_provider.py | 34 ++++---- .../components/provider_info_helpers.test.tsx | 2 +- .../src/components/provider_info_helpers.tsx | 2 +- 6 files changed, 75 insertions(+), 133 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 2f47643da50..dee3cbdd054 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -33546,72 +33546,40 @@ "supports_vision": true, "source": "https://cloud.sambanova.ai/plans/pricing" }, - "scx-ai/gemma-4-31B-it": { - "input_cost_per_token": 3e-07, + "scx-ai/GLM-5.2": { + "cache_read_input_token_cost": 1.375e-07, + "input_cost_per_token": 5.5e-07, "litellm_provider": "scx-ai", - "max_input_tokens": 131000, - "max_tokens": 131000, + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 9.1e-07, - "source": "https://scx.ai/pricing", + "output_cost_per_token": 1.9255e-06, + "source": "https://llmgateway.io/models/glm-5.2/scx-ai-gp", "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "scx-ai/Qwen3.8-Max": { + "cache_read_input_token_cost": 2.1e-07, + "input_cost_per_token": 1.815e-06, + "litellm_provider": "scx-ai", + "max_input_tokens": 1000000, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 5.4461e-06, + "source": "https://llmgateway.io/models/qwen3.8-max/scx-ai-gp", + "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true }, - "scx-ai/gpt-oss-120b": { - "input_cost_per_token": 1.7e-07, - "litellm_provider": "scx-ai", - "max_input_tokens": 131000, - "max_tokens": 131000, - "mode": "chat", - "output_cost_per_token": 5.5e-07, - "source": "https://scx.ai/pricing", - "supports_function_calling": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_tool_choice": true - }, - "scx-ai/Llama-4-Maverick-17B-128E-Instruct": { - "input_cost_per_token": 5.3e-07, - "litellm_provider": "scx-ai", - "max_input_tokens": 131000, - "max_tokens": 131000, - "mode": "chat", - "output_cost_per_token": 1.62e-06, - "source": "https://scx.ai/pricing", - "supports_function_calling": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "scx-ai/MiniMax-M2.7": { - "cache_read_input_token_cost": 5e-08, - "input_cost_per_token": 4.8e-07, - "litellm_provider": "scx-ai", - "max_input_tokens": 192000, - "max_tokens": 192000, - "mode": "chat", - "output_cost_per_token": 1.79e-06, - "source": "https://scx.ai/pricing", - "supports_function_calling": true, - "supports_response_schema": true, - "supports_tool_choice": true - }, - "scx-ai/Qwen3-32B": { - "input_cost_per_token": 3.6e-07, - "litellm_provider": "scx-ai", - "max_input_tokens": 33000, - "max_tokens": 33000, - "mode": "chat", - "output_cost_per_token": 8.7e-07, - "source": "https://scx.ai/pricing", - "supports_function_calling": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_tool_choice": true - }, "snowflake/claude-3-5-sonnet": { "litellm_provider": "snowflake", "max_input_tokens": 200000, diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index 4cfe9e4ef5e..99d0de262a6 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -2686,7 +2686,7 @@ "default_value": null } ], - "default_model_placeholder": "scx-ai/MiniMax-M2.7" + "default_model_placeholder": "scx-ai/GLM-5.2" }, { "provider": "Snowflake", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 8464aec769f..30624d375f7 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -33637,72 +33637,40 @@ "supports_vision": true, "source": "https://cloud.sambanova.ai/plans/pricing" }, - "scx-ai/gemma-4-31B-it": { - "input_cost_per_token": 3e-07, + "scx-ai/GLM-5.2": { + "cache_read_input_token_cost": 1.375e-07, + "input_cost_per_token": 5.5e-07, "litellm_provider": "scx-ai", - "max_input_tokens": 131000, - "max_tokens": 131000, + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 9.1e-07, - "source": "https://scx.ai/pricing", + "output_cost_per_token": 1.9255e-06, + "source": "https://llmgateway.io/models/glm-5.2/scx-ai-gp", "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "scx-ai/Qwen3.8-Max": { + "cache_read_input_token_cost": 2.1e-07, + "input_cost_per_token": 1.815e-06, + "litellm_provider": "scx-ai", + "max_input_tokens": 1000000, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 5.4461e-06, + "source": "https://llmgateway.io/models/qwen3.8-max/scx-ai-gp", + "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true }, - "scx-ai/gpt-oss-120b": { - "input_cost_per_token": 1.7e-07, - "litellm_provider": "scx-ai", - "max_input_tokens": 131000, - "max_tokens": 131000, - "mode": "chat", - "output_cost_per_token": 5.5e-07, - "source": "https://scx.ai/pricing", - "supports_function_calling": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_tool_choice": true - }, - "scx-ai/Llama-4-Maverick-17B-128E-Instruct": { - "input_cost_per_token": 5.3e-07, - "litellm_provider": "scx-ai", - "max_input_tokens": 131000, - "max_tokens": 131000, - "mode": "chat", - "output_cost_per_token": 1.62e-06, - "source": "https://scx.ai/pricing", - "supports_function_calling": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "scx-ai/MiniMax-M2.7": { - "cache_read_input_token_cost": 5e-08, - "input_cost_per_token": 4.8e-07, - "litellm_provider": "scx-ai", - "max_input_tokens": 192000, - "max_tokens": 192000, - "mode": "chat", - "output_cost_per_token": 1.79e-06, - "source": "https://scx.ai/pricing", - "supports_function_calling": true, - "supports_response_schema": true, - "supports_tool_choice": true - }, - "scx-ai/Qwen3-32B": { - "input_cost_per_token": 3.6e-07, - "litellm_provider": "scx-ai", - "max_input_tokens": 33000, - "max_tokens": 33000, - "mode": "chat", - "output_cost_per_token": 8.7e-07, - "source": "https://scx.ai/pricing", - "supports_function_calling": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_tool_choice": true - }, "snowflake/claude-3-5-sonnet": { "litellm_provider": "snowflake", "max_input_tokens": 200000, diff --git a/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py b/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py index 92921067728..6b293cae303 100644 --- a/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py +++ b/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py @@ -34,13 +34,13 @@ class TestSCXAIProviderConfig: from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider model, provider, api_key, api_base = get_llm_provider( - model="scx-ai/gpt-oss-120b", + model="scx-ai/GLM-5.2", custom_llm_provider=None, api_base=None, api_key=None, ) - assert model == "gpt-oss-120b" + assert model == "GLM-5.2" assert provider == "scx-ai" assert api_base == "https://api.scx.ai/v1" @@ -48,7 +48,7 @@ class TestSCXAIProviderConfig: from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider model, provider, api_key, api_base = get_llm_provider( - model="scx-ai/gpt-oss-120b", + model="scx-ai/GLM-5.2", custom_llm_provider=None, api_base="https://custom.scx.ai/v1", api_key="sk-test", @@ -62,7 +62,7 @@ class TestSCXAIProviderConfig: from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider model, provider, api_key, api_base = get_llm_provider( - model="gpt-oss-120b", + model="GLM-5.2", custom_llm_provider=None, api_base="https://api.scx.ai/v1", api_key=None, @@ -81,7 +81,7 @@ class TestSCXAIProviderConfig: optional_params = config.map_openai_params( non_default_params={"temperature": 1.7}, optional_params={}, - model="gpt-oss-120b", + model="GLM-5.2", drop_params=False, ) assert optional_params["temperature"] == 1.0 @@ -89,7 +89,7 @@ class TestSCXAIProviderConfig: optional_params = config.map_openai_params( non_default_params={"temperature": 0.4}, optional_params={}, - model="gpt-oss-120b", + model="GLM-5.2", drop_params=False, ) assert optional_params["temperature"] == 0.4 @@ -105,7 +105,7 @@ class TestSCXAIProviderConfig: optional_params = config.map_openai_params( non_default_params={"max_completion_tokens": 256}, optional_params={}, - model="gpt-oss-120b", + model="GLM-5.2", drop_params=False, ) assert optional_params["max_tokens"] == 256 @@ -119,7 +119,7 @@ class TestSCXAIProviderConfig: { "model_name": "scx-chat", "litellm_params": { - "model": "scx-ai/gpt-oss-120b", + "model": "scx-ai/GLM-5.2", "api_key": "test-key", }, } @@ -132,13 +132,10 @@ class TestSCXAIProviderConfig: class TestSCXAIModelMetadata: SCX_MODELS = ( - "scx-ai/Llama-4-Maverick-17B-128E-Instruct", - "scx-ai/gemma-4-31B-it", - "scx-ai/Qwen3-32B", - "scx-ai/MiniMax-M2.7", - "scx-ai/gpt-oss-120b", + "scx-ai/GLM-5.2", + "scx-ai/Qwen3.8-Max", ) - VISION_MODELS = ("scx-ai/Llama-4-Maverick-17B-128E-Instruct", "scx-ai/gemma-4-31B-it") + VISION_MODELS = ("scx-ai/Qwen3.8-Max",) @staticmethod def _load(path_parts): @@ -160,8 +157,17 @@ class TestSCXAIModelMetadata: assert info["output_cost_per_token"] > 0 assert info["supports_function_calling"] is True assert info["supports_tool_choice"] is True + assert info["supports_reasoning"] is True + assert info["supports_response_schema"] is True assert info.get("supports_vision", False) is (model in self.VISION_MODELS) + assert info["supports_prompt_caching"] is True + assert 0 < info["cache_read_input_token_cost"] < info["input_cost_per_token"] + + assert info["max_output_tokens"] == 131072 + assert info["max_tokens"] == info["max_output_tokens"] + assert info["max_input_tokens"] >= 1_000_000 + def test_scx_ai_models_synced_to_backup(self): model_cost = self._load(("model_prices_and_context_window.json",)) backup = self._load(("litellm", "model_prices_and_context_window_backup.json")) diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx index 8a4c9dfd24d..75af4cffa06 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx @@ -173,7 +173,7 @@ describe("provider_info_helpers", () => { }); it("should return an scx-ai model placeholder for SCX_AI provider", () => { - expect(getPlaceholder(Providers.SCX_AI)).toBe("scx-ai/MiniMax-M2.7"); + expect(getPlaceholder(Providers.SCX_AI)).toBe("scx-ai/GLM-5.2"); }); it("should return claude-3-opus placeholder for Anthropic provider", () => { diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx index fa1f9c11eb3..41898ce27ed 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -444,7 +444,7 @@ export const getPlaceholder = (selectedProvider: string): string => { } else if (selectedProvider === Providers.ZAI) { return "zai/glm-4.5"; } else if (selectedProvider === Providers.SCX_AI) { - return "scx-ai/MiniMax-M2.7"; + return "scx-ai/GLM-5.2"; } else { return "gpt-3.5-turbo"; } From a028c8857e2d9ff23308ff256c2723220868b200 Mon Sep 17 00:00:00 2001 From: bhuvan2134686 Date: Fri, 7 Aug 2026 11:37:38 +1000 Subject: [PATCH 009/166] fix(scx-ai): correct the temperature ceiling to match the endpoint The constraint was 1.0, so anything above that was silently clamped down. SCX accepts [0.0, 2.0), verified live against both GLM-5.2 and Qwen3.8 Max: 1.5, 1.99 and 1.999 all return 200, while 2.0 returns 400 with "Temperature should be in [0.0, 2.0)" Since the clamp is an inclusive min(), 2.0 cannot be the ceiling or it would pass through a value the endpoint rejects. 1.99 is the practical maximum The clamp test now pins both ends: 2.5 comes back as 1.99, and 1.7 rides through untouched where it used to be flattened to 1.0 --- litellm/llms/openai_like/providers.json | 2 +- .../llms/openai_like/test_scx_ai_provider.py | 12 ++++++++++-- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index d796d140878..b43a44c2d3e 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -192,7 +192,7 @@ "max_completion_tokens": "max_tokens" }, "constraints": { - "temperature_max": 1.0 + "temperature_max": 1.99 }, "supported_endpoints": ["/v1/chat/completions"] } diff --git a/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py b/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py index 6b293cae303..1ce2da65fef 100644 --- a/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py +++ b/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py @@ -23,7 +23,7 @@ class TestSCXAIProviderConfig: assert scx.base_url == "https://api.scx.ai/v1" assert scx.api_key_env == "SCX_API_KEY" assert scx.param_mappings.get("max_completion_tokens") == "max_tokens" - assert scx.constraints.get("temperature_max") == 1.0 + assert scx.constraints.get("temperature_max") == 1.99 def test_scx_ai_in_openai_compatible_providers(self): from litellm.constants import openai_compatible_providers @@ -78,13 +78,21 @@ class TestSCXAIProviderConfig: assert provider is not None config = create_config_class(provider)() + optional_params = config.map_openai_params( + non_default_params={"temperature": 2.5}, + optional_params={}, + model="GLM-5.2", + drop_params=False, + ) + assert optional_params["temperature"] == 1.99 + optional_params = config.map_openai_params( non_default_params={"temperature": 1.7}, optional_params={}, model="GLM-5.2", drop_params=False, ) - assert optional_params["temperature"] == 1.0 + assert optional_params["temperature"] == 1.7 optional_params = config.map_openai_params( non_default_params={"temperature": 0.4}, From 0e3f52a4c0b071e8f297decce690a6ba2b6615ff Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Thu, 13 Aug 2026 19:18:22 -0500 Subject: [PATCH 010/166] feat(model_prices): add gemini-3.1-flash-lite-image Register Nano Banana 2 Lite on the unprefixed, gemini/, and vertex_ai/ keys so completion_cost and pass-through spend tracking no longer treat the model as unmapped --- ...odel_prices_and_context_window_backup.json | 95 +++++++ model_prices_and_context_window.json | 95 +++++++ .../llm_cost_calc/test_llm_cost_calc_utils.py | 1 + ...ini_3_1_flash_lite_image_model_metadata.py | 242 ++++++++++++++++++ tests/test_litellm/test_utils.py | 2 + 5 files changed, 435 insertions(+) create mode 100644 tests/test_litellm/test_gemini_3_1_flash_lite_image_model_metadata.py diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 1eb72c887b5..ca67d5d6844 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -18141,6 +18141,44 @@ }, "web_search_billing_unit": "per_query" }, + "gemini-3.1-flash-lite-image": { + "cache_read_input_token_cost": 2.5e-08, + "input_cost_per_image": 0.00028, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "image_generation", + "output_cost_per_image": 0.0336, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_video_input": true, + "supports_vision": true + }, "gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, "input_cost_per_audio_token": 5e-07, @@ -19949,6 +19987,42 @@ }, "web_search_billing_unit": "per_query" }, + "gemini/gemini-3.1-flash-lite-image": { + "input_cost_per_image": 0.00028, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "gemini", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "image_generation", + "output_cost_per_image": 0.0336, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "rpm": 1000, + "tpm": 4000000, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-lite-image", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_vision": true + }, "gemini/deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -38760,6 +38834,27 @@ "supports_reasoning": false, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models" }, + "vertex_ai/gemini-3.1-flash-lite-image": { + "cache_read_input_token_cost": 2.5e-08, + "input_cost_per_image": 0.00028, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "image_generation", + "output_cost_per_image": 0.0336, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_vision": true, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, "vertex_ai/gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, "input_cost_per_audio_token": 5e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 1eb72c887b5..ca67d5d6844 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -18141,6 +18141,44 @@ }, "web_search_billing_unit": "per_query" }, + "gemini-3.1-flash-lite-image": { + "cache_read_input_token_cost": 2.5e-08, + "input_cost_per_image": 0.00028, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "image_generation", + "output_cost_per_image": 0.0336, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_video_input": true, + "supports_vision": true + }, "gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, "input_cost_per_audio_token": 5e-07, @@ -19949,6 +19987,42 @@ }, "web_search_billing_unit": "per_query" }, + "gemini/gemini-3.1-flash-lite-image": { + "input_cost_per_image": 0.00028, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "gemini", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "image_generation", + "output_cost_per_image": 0.0336, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "rpm": 1000, + "tpm": 4000000, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-lite-image", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_vision": true + }, "gemini/deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -38760,6 +38834,27 @@ "supports_reasoning": false, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models" }, + "vertex_ai/gemini-3.1-flash-lite-image": { + "cache_read_input_token_cost": 2.5e-08, + "input_cost_per_image": 0.00028, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "image_generation", + "output_cost_per_image": 0.0336, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_vision": true, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, "vertex_ai/gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, "input_cost_per_audio_token": 5e-07, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 3aa41e18f1e..36fc98a1f09 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -1546,6 +1546,7 @@ def test_service_tier_fallback_pricing(): [ "gemini-3-pro-image-preview", "gemini-3.1-flash-image-preview", + "gemini-3.1-flash-lite-image", ], ) def test_gemini_image_generation_cost_with_zero_text_tokens(model: str): diff --git a/tests/test_litellm/test_gemini_3_1_flash_lite_image_model_metadata.py b/tests/test_litellm/test_gemini_3_1_flash_lite_image_model_metadata.py new file mode 100644 index 00000000000..aa6f03a47ff --- /dev/null +++ b/tests/test_litellm/test_gemini_3_1_flash_lite_image_model_metadata.py @@ -0,0 +1,242 @@ +import json +from pathlib import Path + +import pytest + +import litellm +from litellm import completion_cost +from litellm.cost_calculator import cost_per_token +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token +from litellm.llms.gemini.image_generation.cost_calculator import ( + cost_calculator as gemini_image_generation_cost_calculator, +) +from litellm.llms.vertex_ai.image_generation.cost_calculator import ( + cost_calculator as vertex_image_generation_cost_calculator, +) +from litellm.types.utils import ( + CompletionTokensDetailsWrapper, + ImageObject, + ImageResponse, + ImageUsage, + ImageUsageInputTokensDetails, + ModelResponse, + PromptTokensDetailsWrapper, + Usage, +) + +REPO_ROOT = Path(__file__).parents[2] +MAIN_PATH = REPO_ROOT / "model_prices_and_context_window.json" +BACKUP_PATH = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" + +UNPREFIXED = "gemini-3.1-flash-lite-image" +GEMINI = "gemini/gemini-3.1-flash-lite-image" +VERTEX = "vertex_ai/gemini-3.1-flash-lite-image" +ALL_KEYS = (UNPREFIXED, GEMINI, VERTEX) + +INPUT_COST = 2.5e-07 +INPUT_COST_BATCHES = 1.25e-07 +OUTPUT_TEXT_COST = 1.5e-06 +OUTPUT_TEXT_COST_BATCHES = 7.5e-07 +OUTPUT_IMAGE_TOKEN_COST = 3e-05 +OUTPUT_COST_PER_1K_IMAGE = 0.0336 +INPUT_COST_PER_IMAGE = 0.00028 +CACHE_READ_COST = 2.5e-08 +MAX_INPUT_TOKENS = 65536 +MAX_OUTPUT_TOKENS = 4096 +TOKENS_PER_1K_IMAGE = 1120 + + +def _load(path: Path) -> dict: + with open(path) as f: + return json.load(f) + + +@pytest.fixture +def local_model_cost_map(monkeypatch): + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + +@pytest.mark.parametrize("model", ALL_KEYS) +def test_gemini_3_1_flash_lite_image_is_registered(model: str): + info = _load(MAIN_PATH).get(model) + assert info is not None, f"{model} not found in model_prices_and_context_window.json" + + assert info["mode"] == "image_generation" + assert info["input_cost_per_token"] == INPUT_COST + assert info["input_cost_per_token_batches"] == INPUT_COST_BATCHES + assert info["output_cost_per_token"] == OUTPUT_TEXT_COST + assert info["output_cost_per_token_batches"] == OUTPUT_TEXT_COST_BATCHES + assert info["output_cost_per_image"] == OUTPUT_COST_PER_1K_IMAGE + assert info["output_cost_per_image_token"] == OUTPUT_IMAGE_TOKEN_COST + assert info["max_input_tokens"] == MAX_INPUT_TOKENS + assert info["max_output_tokens"] == MAX_OUTPUT_TOKENS + assert info["max_tokens"] == MAX_OUTPUT_TOKENS + assert info["supports_reasoning"] is False + assert info["supports_response_schema"] is False + assert info["supports_vision"] is True + for field in ("supports_web_search", "search_context_cost_per_query", "web_search_billing_unit"): + assert field not in info + + +def test_gemini_3_1_flash_lite_image_provider_specific_fields(): + cost_map = _load(MAIN_PATH) + + unprefixed = cost_map[UNPREFIXED] + assert unprefixed["litellm_provider"] == "vertex_ai-language-models" + assert unprefixed["cache_read_input_token_cost"] == CACHE_READ_COST + assert unprefixed["input_cost_per_image"] == INPUT_COST_PER_IMAGE + assert unprefixed["supports_function_calling"] is False + assert unprefixed["supports_prompt_caching"] is True + assert unprefixed["supports_pdf_input"] is True + assert unprefixed["supports_video_input"] is True + assert unprefixed["supported_modalities"] == ["text", "image", "video"] + + gemini = cost_map[GEMINI] + assert gemini["litellm_provider"] == "gemini" + assert gemini["supports_function_calling"] is True + assert gemini["supports_prompt_caching"] is False + assert "cache_read_input_token_cost" not in gemini + assert gemini["supported_modalities"] == ["text", "image"] + assert gemini["supported_output_modalities"] == ["text", "image"] + assert gemini["rpm"] == 1000 + assert gemini["tpm"] == 4000000 + assert gemini["input_cost_per_image"] == INPUT_COST_PER_IMAGE + + vertex = cost_map[VERTEX] + assert vertex["litellm_provider"] == "vertex_ai-language-models" + assert vertex["cache_read_input_token_cost"] == CACHE_READ_COST + assert vertex["input_cost_per_image"] == INPUT_COST_PER_IMAGE + assert vertex["supports_function_calling"] is False + assert vertex["supports_prompt_caching"] is True + + +def test_one_k_image_price_matches_official_token_math(): + assert TOKENS_PER_1K_IMAGE * OUTPUT_IMAGE_TOKEN_COST == OUTPUT_COST_PER_1K_IMAGE + assert TOKENS_PER_1K_IMAGE * INPUT_COST == INPUT_COST_PER_IMAGE + + +@pytest.mark.parametrize("model", ALL_KEYS) +def test_backup_matches_main(model: str): + main_cost = _load(MAIN_PATH) + backup_cost = _load(BACKUP_PATH) + assert backup_cost.get(model) == main_cost.get(model), f"{model} differs between main and backup model cost maps" + + +def test_gemini_prefix_routes_to_gemini(): + routed_model, provider, _, _ = get_llm_provider(model=GEMINI) + assert routed_model == UNPREFIXED + assert provider == "gemini" + + +def test_vertex_prefix_routes_to_vertex(): + routed_model, provider, _, _ = get_llm_provider(model=VERTEX) + assert routed_model == UNPREFIXED + assert provider == "vertex_ai" + + +def test_text_token_cost(local_model_cost_map): + prompt_cost, text_completion_cost = cost_per_token(model=GEMINI, prompt_tokens=1000, completion_tokens=500) + assert prompt_cost == pytest.approx(1000 * INPUT_COST) + assert text_completion_cost == pytest.approx(500 * OUTPUT_TEXT_COST) + + +def test_completion_cost_bills_one_k_image(local_model_cost_map): + response = ModelResponse() + response.model = UNPREFIXED + response.usage = Usage( + prompt_tokens=7, + completion_tokens=TOKENS_PER_1K_IMAGE, + total_tokens=7 + TOKENS_PER_1K_IMAGE, + completion_tokens_details=CompletionTokensDetailsWrapper(image_tokens=TOKENS_PER_1K_IMAGE, text_tokens=0), + ) + billed = completion_cost( + completion_response=response, + model=UNPREFIXED, + custom_llm_provider="vertex_ai", + ) + expected = TOKENS_PER_1K_IMAGE * OUTPUT_IMAGE_TOKEN_COST + 7 * INPUT_COST + assert billed == pytest.approx(expected) + + +def test_image_tokens_are_not_billed_as_text(local_model_cost_map): + usage = Usage( + completion_tokens=1345, + prompt_tokens=10, + total_tokens=1355, + completion_tokens_details=CompletionTokensDetailsWrapper( + accepted_prediction_tokens=None, + audio_tokens=None, + reasoning_tokens=225, + rejected_prediction_tokens=None, + text_tokens=0, + image_tokens=TOKENS_PER_1K_IMAGE, + ), + prompt_tokens_details=PromptTokensDetailsWrapper( + audio_tokens=None, cached_tokens=None, text_tokens=10, image_tokens=None + ), + ) + + _, image_completion_cost = generic_cost_per_token( + model=UNPREFIXED, + usage=usage, + custom_llm_provider="vertex_ai", + ) + + expected_completion_cost = TOKENS_PER_1K_IMAGE * OUTPUT_IMAGE_TOKEN_COST + 225 * OUTPUT_TEXT_COST + bugged_text_only_cost = 1345 * OUTPUT_TEXT_COST + assert image_completion_cost > bugged_text_only_cost * 2 + assert image_completion_cost == pytest.approx(expected_completion_cost) + + +def test_gemini_image_generation_uses_token_pricing(local_model_cost_map): + image_response = ImageResponse( + data=[ImageObject(b64_json="img1")], + usage=ImageUsage( + input_tokens=50 + TOKENS_PER_1K_IMAGE, + input_tokens_details=ImageUsageInputTokensDetails( + text_tokens=50, + image_tokens=TOKENS_PER_1K_IMAGE, + ), + output_tokens=TOKENS_PER_1K_IMAGE, + total_tokens=50 + TOKENS_PER_1K_IMAGE + TOKENS_PER_1K_IMAGE, + ), + ) + + cost = gemini_image_generation_cost_calculator(model=GEMINI, image_response=image_response) + expected = (50 + TOKENS_PER_1K_IMAGE) * INPUT_COST + TOKENS_PER_1K_IMAGE * OUTPUT_IMAGE_TOKEN_COST + assert cost == pytest.approx(expected) + assert cost != OUTPUT_COST_PER_1K_IMAGE + + +def test_vertex_image_generation_uses_token_pricing(local_model_cost_map): + image_response = ImageResponse( + data=[ImageObject(b64_json="img1")], + usage=ImageUsage( + input_tokens=50 + TOKENS_PER_1K_IMAGE, + input_tokens_details=ImageUsageInputTokensDetails( + text_tokens=50, + image_tokens=TOKENS_PER_1K_IMAGE, + ), + output_tokens=TOKENS_PER_1K_IMAGE, + total_tokens=50 + TOKENS_PER_1K_IMAGE + TOKENS_PER_1K_IMAGE, + ), + ) + + cost = vertex_image_generation_cost_calculator(model=UNPREFIXED, image_response=image_response) + expected = (50 + TOKENS_PER_1K_IMAGE) * INPUT_COST + TOKENS_PER_1K_IMAGE * OUTPUT_IMAGE_TOKEN_COST + assert cost == pytest.approx(expected) + + +def test_vertex_image_generation_falls_back_to_flat_image_price(local_model_cost_map): + image_response = ImageResponse(data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")]) + cost = vertex_image_generation_cost_calculator(model=UNPREFIXED, image_response=image_response) + assert cost == pytest.approx(2 * OUTPUT_COST_PER_1K_IMAGE) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 8e9e6167fb9..ada053ee38c 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4312,11 +4312,13 @@ class TestVertexEmbeddingEncodingFormat: "vertex_ai/gemini-3-pro-image-preview", "vertex_ai/gemini-3.1-flash-image", "vertex_ai/gemini-3.1-flash-image-preview", + "vertex_ai/gemini-3.1-flash-lite-image", "gemini/gemini-2.5-flash-image", "gemini/gemini-3-pro-image", "gemini/gemini-3-pro-image-preview", "gemini/gemini-3.1-flash-image", "gemini/gemini-3.1-flash-image-preview", + "gemini/gemini-3.1-flash-lite-image", ], ) def test_gemini_image_models_do_not_support_reasoning( From 9c922f4aa4f6ab3a261326332d01b74e3717e63c Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 15 Aug 2026 02:49:10 +0000 Subject: [PATCH 011/166] fix(bedrock): forward provider response headers on chat completions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/streaming_handler.py | 2 +- litellm/llms/bedrock/chat/converse_handler.py | 18 ++++-- litellm/llms/bedrock/chat/invoke_handler.py | 8 +-- .../base_invoke_transformation.py | 64 +++++++++---------- litellm/llms/custom_httpx/llm_http_handler.py | 2 + litellm/types/utils.py | 7 +- .../llm_translation/test_bedrock_moonshot.py | 19 ++---- .../llms/bedrock/chat/test_invoke_handler.py | 25 ++++++++ .../llms/chat/test_converse_handler.py | 55 ++++++++++++++++ 9 files changed, 143 insertions(+), 57 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 99b1c1a2ab7..fe107d349c6 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -164,7 +164,7 @@ class CustomStreamWrapper: custom_llm_provider: str | None = None, stream_options=None, make_call: Callable | None = None, - _response_headers: dict | None = None, + _response_headers: dict | httpx.Headers | None = None, ): self.model = model self.make_call = make_call diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 25e544f4521..c9ea06d37dc 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -33,7 +33,7 @@ def make_sync_call( json_mode: bool | None = False, fake_stream: bool = False, stream_chunk_size: int | None = None, -): +) -> tuple[Any, httpx.Headers]: if client is None: client = _get_httpx_client() # Create a new client if none provided @@ -74,7 +74,7 @@ def make_sync_call( additional_args={"complete_input_dict": data}, ) - return completion_stream + return completion_stream, response.headers class BedrockConverseLLM(BaseAWSLLM): @@ -132,7 +132,7 @@ class BedrockConverseLLM(BaseAWSLLM): }, ) - completion_stream: Final = await make_call( + completion_stream, response_headers = await make_call( client=client, api_base=api_base, headers=dict(prepped.headers), @@ -149,6 +149,7 @@ class BedrockConverseLLM(BaseAWSLLM): model=model, custom_llm_provider="bedrock", logging_obj=logging_obj, + _response_headers=response_headers, ) return streaming_response @@ -225,7 +226,7 @@ class BedrockConverseLLM(BaseAWSLLM): except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") - return litellm.AmazonConverseConfig()._transform_response( + transformed_response: Final = litellm.AmazonConverseConfig()._transform_response( model=model, response=response, model_response=model_response, @@ -237,6 +238,8 @@ class BedrockConverseLLM(BaseAWSLLM): optional_params=optional_params, encoding=encoding, ) + transformed_response.set_provider_response_headers(response.headers) + return transformed_response def completion( self, @@ -440,7 +443,7 @@ class BedrockConverseLLM(BaseAWSLLM): client = client if stream is not None and stream is True: - completion_stream: Final = make_sync_call( + completion_stream, response_headers = make_sync_call( client=(client if client is not None and isinstance(client, HTTPHandler) else None), api_base=proxy_endpoint_url, headers=prepped.headers, @@ -457,6 +460,7 @@ class BedrockConverseLLM(BaseAWSLLM): model=model, custom_llm_provider="bedrock", logging_obj=logging_obj, + _response_headers=response_headers, ) return streaming_response @@ -477,7 +481,7 @@ class BedrockConverseLLM(BaseAWSLLM): except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") - return litellm.AmazonConverseConfig()._transform_response( + sync_transformed_response: Final = litellm.AmazonConverseConfig()._transform_response( model=model, response=response, model_response=model_response, @@ -489,3 +493,5 @@ class BedrockConverseLLM(BaseAWSLLM): optional_params=optional_params, encoding=encoding, ) + sync_transformed_response.set_provider_response_headers(response.headers) + return sync_transformed_response diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 8d2b3dae71b..21892af3aae 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -163,7 +163,7 @@ async def make_call( json_mode: bool | None = False, bedrock_invoke_provider: litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL | None = None, stream_chunk_size: int | None = None, -): +) -> tuple[Any, httpx.Headers]: try: if client is None: client = get_async_httpx_client( @@ -225,7 +225,7 @@ async def make_call( additional_args={"complete_input_dict": data}, ) - return completion_stream + return completion_stream, response.headers except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code raise BedrockError(status_code=error_code, message=err.response.text) @@ -248,7 +248,7 @@ def make_sync_call( json_mode: bool | None = False, bedrock_invoke_provider: litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL | None = None, stream_chunk_size: int | None = None, -): +) -> tuple[Any, httpx.Headers]: try: if client is None: client = _get_httpx_client( @@ -309,7 +309,7 @@ def make_sync_call( additional_args={"complete_input_dict": data}, ) - return completion_stream + return completion_stream, response.headers except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code raise BedrockError(status_code=error_code, message=err.response.text) diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index 430d0a92b51..0b5855ccbd4 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -1,7 +1,6 @@ import copy import json import time -from functools import partial from typing import TYPE_CHECKING, Any, Final, cast, get_args import httpx @@ -444,24 +443,24 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): json_mode: bool | None = None, signed_json_body: bytes | None = None, ) -> CustomStreamWrapper: + completion_stream, response_headers = await make_call( + client=client, + api_base=api_base, + headers=headers, + data=json.dumps(data), + model=model, + messages=messages, + logging_obj=logging_obj, + fake_stream=True if "ai21" in api_base else False, + bedrock_invoke_provider=self.get_bedrock_invoke_provider(model), + json_mode=json_mode, + ) streaming_response: Final = CustomStreamWrapper( - completion_stream=None, - make_call=partial( - make_call, - client=client, - api_base=api_base, - headers=headers, - data=json.dumps(data), - model=model, - messages=messages, - logging_obj=logging_obj, - fake_stream=True if "ai21" in api_base else False, - bedrock_invoke_provider=self.get_bedrock_invoke_provider(model), - json_mode=json_mode, - ), + completion_stream=completion_stream, model=model, custom_llm_provider="bedrock", logging_obj=logging_obj, + _response_headers=response_headers, ) return streaming_response @@ -479,27 +478,28 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): json_mode: bool | None = None, signed_json_body: bytes | None = None, ) -> CustomStreamWrapper: - if client is None or isinstance(client, AsyncHTTPHandler): - client = _get_httpx_client(params={}) + sync_client: Final = ( + _get_httpx_client(params={}) if client is None or isinstance(client, AsyncHTTPHandler) else client + ) + completion_stream, response_headers = make_sync_call( + client=sync_client, + api_base=api_base, + headers=headers, + data=json.dumps(data), + signed_json_body=signed_json_body, + model=model, + messages=messages, + logging_obj=logging_obj, + fake_stream=True if "ai21" in api_base else False, + bedrock_invoke_provider=self.get_bedrock_invoke_provider(model), + json_mode=json_mode, + ) streaming_response: Final = CustomStreamWrapper( - completion_stream=None, - make_call=partial( - make_sync_call, - client=client, - api_base=api_base, - headers=headers, - data=json.dumps(data), - signed_json_body=signed_json_body, - model=model, - messages=messages, - logging_obj=logging_obj, - fake_stream=True if "ai21" in api_base else False, - bedrock_invoke_provider=self.get_bedrock_invoke_provider(model), - json_mode=json_mode, - ), + completion_stream=completion_stream, model=model, custom_llm_provider="bedrock", logging_obj=logging_obj, + _response_headers=response_headers, ) return streaming_response diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 721b9545ac1..36bcf58a243 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -608,6 +608,7 @@ class BaseLLMHTTPHandler: model=model, custom_llm_provider=custom_llm_provider, logging_obj=logging_obj, + _response_headers=headers, ) if client is None or not isinstance(client, HTTPHandler): @@ -771,6 +772,7 @@ class BaseLLMHTTPHandler: model=model, custom_llm_provider=custom_llm_provider, logging_obj=logging_obj, + _response_headers=_response_headers, ) return streamwrapper diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 220826ccbca..f49b5735f04 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -11,6 +11,7 @@ from typing import ( get_args, ) +import httpx from openai._models import BaseModel as OpenAIObject from openai.types.audio.transcription_create_params import ( FileTypes as FileTypes, @@ -49,7 +50,7 @@ from litellm.types.llms.base import ( ) from litellm.types.mcp import MCPServerCostInfo -from ..litellm_core_utils.core_helpers import map_finish_reason +from ..litellm_core_utils.core_helpers import map_finish_reason, process_response_headers from .agents import LiteLLMSendMessageResponse from .guardrails import GuardrailEventHooks from .llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse @@ -1896,6 +1897,10 @@ class ModelResponseBase(OpenAIObject): _response_headers: dict | None = None + def set_provider_response_headers(self, headers: httpx.Headers) -> None: + """Surface a provider's raw response headers to the caller as `llm_provider-*` headers.""" + self._hidden_params["additional_headers"] = process_response_headers(headers) + def model_dump(self, **kwargs): """Default to exclude_unset to avoid Pydantic serializer warnings for OpenAIObject-derived types.""" if "exclude_unset" not in kwargs and "exclude_none" not in kwargs: diff --git a/tests/llm_translation/test_bedrock_moonshot.py b/tests/llm_translation/test_bedrock_moonshot.py index a9f4a86b3b6..c777305d562 100644 --- a/tests/llm_translation/test_bedrock_moonshot.py +++ b/tests/llm_translation/test_bedrock_moonshot.py @@ -209,14 +209,12 @@ class TestBedrockMoonshotInvoke(BaseLLMChatTest): not exercised here — moonshot streaming delegates to the OpenAI parser and is covered by the OpenAI test suite. - Note: bedrock invoke streaming cannot be intercepted by patching - the caller-supplied client, because ``CustomStreamWrapper.fetch_sync_stream`` - at streaming_handler.py invokes the stored ``make_call`` partial with - ``client=litellm.module_level_client``, which overrides any client the - caller passed. Patch ``make_sync_call`` at its import site in - ``base_invoke_transformation`` so we observe the exact kwargs the - partial was built with at stream-wrapper construction time. + Patch ``make_sync_call`` at its import site in + ``base_invoke_transformation`` so we observe the exact kwargs it is + called with at stream-wrapper construction time. """ + import httpx + from litellm.utils import CustomStreamWrapper captured: dict = {} @@ -225,7 +223,7 @@ class TestBedrockMoonshotInvoke(BaseLLMChatTest): captured.update(kwargs) # Return an empty iterator so the stream wrapper's iteration # doesn't try to parse real bytes. - return iter([]) + return iter([]), httpx.Headers() with patch( "litellm.llms.bedrock.chat.invoke_transformations." @@ -246,11 +244,6 @@ class TestBedrockMoonshotInvoke(BaseLLMChatTest): aws_region_name="us-west-2", ) assert isinstance(response, CustomStreamWrapper) - # Trigger fetch_sync_stream → make_call(...) → fake_make_sync_call. - try: - next(iter(response)) - except StopIteration: - pass assert captured, "make_sync_call was never invoked" assert captured["api_base"].endswith("/invoke-with-response-stream") diff --git a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py index ee50b9db015..f211cdac475 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py @@ -2,17 +2,20 @@ import os import sys from unittest.mock import AsyncMock, MagicMock +import httpx import pytest sys.path.insert( 0, os.path.abspath("../../../../..") ) # Adds the parent directory to the system path +import litellm from litellm.llms.bedrock.chat.invoke_handler import ( AWSEventStreamDecoder, make_call, make_sync_call, ) +from litellm.llms.custom_httpx.http_handler import HTTPHandler def test_transform_thinking_blocks_with_redacted_content(): @@ -293,3 +296,25 @@ def test_make_sync_call_honors_explicit_stream_chunk_size(): response.iter_bytes.assert_called_once_with(chunk_size=2048) + +def test_invoke_streaming_forwards_bedrock_response_headers(): + """Streaming callers need `x-amzn-requestid` to correlate a LiteLLM request with AWS support.""" + response = MagicMock() + response.status_code = 200 + response.iter_bytes = MagicMock(return_value=iter([])) + response.headers = httpx.Headers({"x-amzn-requestid": "req-789"}) + client = HTTPHandler() + client.post = MagicMock(return_value=response) + + stream = litellm.completion( + model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert stream._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-789" + diff --git a/tests/test_litellm/llms/chat/test_converse_handler.py b/tests/test_litellm/llms/chat/test_converse_handler.py index 2a3db5982ef..89ab29122d7 100644 --- a/tests/test_litellm/llms/chat/test_converse_handler.py +++ b/tests/test_litellm/llms/chat/test_converse_handler.py @@ -1,7 +1,9 @@ +import json import os import sys from unittest.mock import MagicMock +import httpx import pytest import litellm @@ -202,6 +204,59 @@ def test_make_sync_call_honors_explicit_stream_chunk_size(): response.iter_bytes.assert_called_once_with(chunk_size=2048) +def _converse_response_body() -> dict: + return { + "output": {"message": {"role": "assistant", "content": [{"text": "hi"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2}, + } + + +def test_converse_completion_forwards_bedrock_response_headers(): + """Bedrock returns x-amzn-requestid on every converse call, which customers need to + correlate proxy requests with AWS support cases, so it must reach the caller as + llm_provider-x-amzn-requestid.""" + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json = MagicMock(return_value=_converse_response_body()) + mock_response.text = json.dumps(_converse_response_body()) + mock_response.headers = httpx.Headers({"x-amzn-requestid": "req-123"}) + client = HTTPHandler() + client.post = MagicMock(return_value=mock_response) + + response = litellm.completion( + model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert response._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-123" + + +def test_converse_streaming_forwards_bedrock_response_headers(): + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.iter_bytes = MagicMock(return_value=iter([])) + mock_response.headers = httpx.Headers({"x-amzn-requestid": "req-456"}) + client = HTTPHandler() + client.post = MagicMock(return_value=mock_response) + + response = litellm.completion( + model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert response._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-456" + + def test_completion_plumbs_stream_chunk_size_through_converse(): iter_bytes_spy = _stream_completion_with_spied_iter_bytes( model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0" From 726db1a4c118a2ab92214ba741efacdce7fe25b5 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 15 Aug 2026 03:07:14 +0000 Subject: [PATCH 012/166] test(bedrock): cover async header forwarding for converse and invoke Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llm_translation/test_bedrock_moonshot.py | 3 +- .../llms/bedrock/chat/test_invoke_handler.py | 29 +++++++++- .../llms/chat/test_converse_handler.py | 55 +++++++++++++++++-- 3 files changed, 78 insertions(+), 9 deletions(-) diff --git a/tests/llm_translation/test_bedrock_moonshot.py b/tests/llm_translation/test_bedrock_moonshot.py index c777305d562..61364cf2caa 100644 --- a/tests/llm_translation/test_bedrock_moonshot.py +++ b/tests/llm_translation/test_bedrock_moonshot.py @@ -12,6 +12,7 @@ This test suite verifies: """ from base_llm_unit_tests import BaseLLMChatTest +import httpx import pytest import sys import os @@ -213,8 +214,6 @@ class TestBedrockMoonshotInvoke(BaseLLMChatTest): ``base_invoke_transformation`` so we observe the exact kwargs it is called with at stream-wrapper construction time. """ - import httpx - from litellm.utils import CustomStreamWrapper captured: dict = {} diff --git a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py index f211cdac475..e8964910c69 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py @@ -15,7 +15,7 @@ from litellm.llms.bedrock.chat.invoke_handler import ( make_call, make_sync_call, ) -from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler def test_transform_thinking_blocks_with_redacted_content(): @@ -298,7 +298,6 @@ def test_make_sync_call_honors_explicit_stream_chunk_size(): def test_invoke_streaming_forwards_bedrock_response_headers(): - """Streaming callers need `x-amzn-requestid` to correlate a LiteLLM request with AWS support.""" response = MagicMock() response.status_code = 200 response.iter_bytes = MagicMock(return_value=iter([])) @@ -318,3 +317,29 @@ def test_invoke_streaming_forwards_bedrock_response_headers(): assert stream._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-789" + +@pytest.mark.asyncio +async def test_async_invoke_streaming_forwards_bedrock_response_headers(): + async def _no_bytes(chunk_size=None): + return + yield b"" + + response = MagicMock() + response.status_code = 200 + response.aiter_bytes = _no_bytes + response.headers = httpx.Headers({"x-amzn-requestid": "req-987"}) + client = AsyncHTTPHandler() + client.post = AsyncMock(return_value=response) + + stream = await litellm.acompletion( + model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert stream._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-987" + diff --git a/tests/test_litellm/llms/chat/test_converse_handler.py b/tests/test_litellm/llms/chat/test_converse_handler.py index 89ab29122d7..6f8a2788c38 100644 --- a/tests/test_litellm/llms/chat/test_converse_handler.py +++ b/tests/test_litellm/llms/chat/test_converse_handler.py @@ -1,7 +1,7 @@ import json import os import sys -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import httpx import pytest @@ -10,7 +10,7 @@ import litellm from litellm.llms.bedrock.chat import BedrockConverseLLM from litellm.llms.bedrock.chat.converse_handler import make_sync_call from litellm.llms.bedrock.common_utils import _get_all_bedrock_regions -from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler sys.path.insert( 0, os.path.abspath("../../../../..") @@ -213,9 +213,6 @@ def _converse_response_body() -> dict: def test_converse_completion_forwards_bedrock_response_headers(): - """Bedrock returns x-amzn-requestid on every converse call, which customers need to - correlate proxy requests with AWS support cases, so it must reach the caller as - llm_provider-x-amzn-requestid.""" mock_response = MagicMock() mock_response.status_code = 200 mock_response.json = MagicMock(return_value=_converse_response_body()) @@ -257,6 +254,54 @@ def test_converse_streaming_forwards_bedrock_response_headers(): assert response._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-456" +@pytest.mark.asyncio +async def test_async_converse_completion_forwards_bedrock_response_headers(): + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json = MagicMock(return_value=_converse_response_body()) + mock_response.text = json.dumps(_converse_response_body()) + mock_response.headers = httpx.Headers({"x-amzn-requestid": "req-abc"}) + client = AsyncHTTPHandler() + client.post = AsyncMock(return_value=mock_response) + + response = await litellm.acompletion( + model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert response._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-abc" + + +@pytest.mark.asyncio +async def test_async_converse_streaming_forwards_bedrock_response_headers(): + async def _no_bytes(chunk_size=None): + return + yield b"" + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.aiter_bytes = _no_bytes + mock_response.headers = httpx.Headers({"x-amzn-requestid": "req-def"}) + client = AsyncHTTPHandler() + client.post = AsyncMock(return_value=mock_response) + + response = await litellm.acompletion( + model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert response._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-def" + + def test_completion_plumbs_stream_chunk_size_through_converse(): iter_bytes_spy = _stream_completion_with_spied_iter_bytes( model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0" From 5d5dc4523fb950e131a235bea9a4f767ba7e0e17 Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 19 Aug 2026 03:56:32 +0000 Subject: [PATCH 013/166] fix(cost): price streamed Messages usage via calculate_usage and the logging obj Streamed `/v1/messages` `usage.cost` disagreed with the cost the logging callback recorded in three ways: `input_tokens` was read as the whole prompt total, but Anthropic reports it excluding cache tokens, so the non-cached input went unbilled on cache hits; the `cache_creation` 5m/1h split was dropped, billing 1h writes at the 5m rate; and costing by model name alone ignored the deployment's custom pricing, so a negotiated discount still streamed sticker price. Anthropic usage now goes through `AnthropicConfig.calculate_usage`, the same transformation the non-streaming path uses, and the chunk is priced through the call's logging object when there is one so it inherits `custom_pricing`, `custom_llm_provider`, `base_model` and `router_model_id`, falling back to `completion_cost` by model name. `calculate_usage` only reads its `usage_object`, so it now takes a `Mapping`. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/anthropic/chat/transformation.py | 2 +- litellm/proxy/common_request_processing.py | 130 ++++++++++------ .../streaming_handler.py | 2 +- .../proxy/test_common_request_processing.py | 140 ++++++++++++++++++ 4 files changed, 226 insertions(+), 48 deletions(-) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index ef4ad7011c5..b4f040b0a8c 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -2152,7 +2152,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): def calculate_usage( self, - usage_object: dict, + usage_object: Mapping[str, Any], reasoning_content: str | None, completion_response: dict | None = None, speed: str | None = None, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index a0b69ecb0bf..f5299b273b4 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -158,7 +158,7 @@ ProxyRouteType: TypeAlias = Literal[ "acancel_run", "adelete_run", ] -from litellm.types.utils import ServerToolUse +from litellm.llms.anthropic.chat.transformation import AnthropicConfig # Type alias for streaming chunk serializer (chunk after hooks + cost injection -> wire format) StreamChunkSerializer = Callable[[Any], str] @@ -3321,7 +3321,9 @@ class ProxyBaseLLMRequestProcessing: str_so_far += str(chunk.get("content", "")) model_name = request_data.get("model", "") - chunk = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection(chunk, model_name) + chunk = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( + chunk, model_name, request_data.get("litellm_logging_obj") + ) # Set before the yield: an async generator suspends at the yield, # so a GeneratorExit on client disconnect is raised there and any @@ -3418,20 +3420,28 @@ class ProxyBaseLLMRequestProcessing: @overload @staticmethod - def _process_chunk_with_cost_injection(chunk: bytes, model_name: str) -> bytes: ... + def _process_chunk_with_cost_injection( + chunk: bytes, model_name: str, litellm_logging_obj: LiteLLMLoggingObj | None = None + ) -> bytes: ... @overload @staticmethod - def _process_chunk_with_cost_injection(chunk: object, model_name: str) -> object: ... + def _process_chunk_with_cost_injection( + chunk: object, model_name: str, litellm_logging_obj: LiteLLMLoggingObj | None = None + ) -> object: ... @staticmethod - def _process_chunk_with_cost_injection(chunk: object, model_name: str) -> object: + def _process_chunk_with_cost_injection( + chunk: object, model_name: str, litellm_logging_obj: LiteLLMLoggingObj | None = None + ) -> object: """ Process a streaming chunk and inject cost information if enabled. Args: chunk: The streaming chunk (dict, str, bytes, or bytearray) model_name: Model name for cost calculation + litellm_logging_obj: The call's logging object, used to price the chunk with + the same custom/deployment pricing as the logging callback Returns: The processed chunk with cost information injected if applicable @@ -3441,21 +3451,27 @@ class ProxyBaseLLMRequestProcessing: try: if isinstance(chunk, dict): - maybe_modified: Final = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(chunk, model_name) + maybe_modified: Final = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict( + chunk, model_name, litellm_logging_obj + ) if maybe_modified is not None: return maybe_modified elif isinstance(chunk, (bytes, bytearray)): try: s: Final = chunk.decode("utf-8") if s.endswith(("\n\n", "\r\n\r\n")): - maybe_mod = ProxyBaseLLMRequestProcessing._inject_cost_into_sse_frame_str(s, model_name) + maybe_mod = ProxyBaseLLMRequestProcessing._inject_cost_into_sse_frame_str( + s, model_name, litellm_logging_obj + ) if maybe_mod is not None: return maybe_mod.encode("utf-8") except Exception: pass elif isinstance(chunk, str): # Try to parse SSE frame and inject cost into the data line - maybe_mod = ProxyBaseLLMRequestProcessing._inject_cost_into_sse_frame_str(chunk, model_name) + maybe_mod = ProxyBaseLLMRequestProcessing._inject_cost_into_sse_frame_str( + chunk, model_name, litellm_logging_obj + ) if maybe_mod is not None: # Ensure trailing frame separator return maybe_mod if maybe_mod.endswith("\n\n") else (maybe_mod + "\n\n") @@ -3466,13 +3482,16 @@ class ProxyBaseLLMRequestProcessing: return chunk @staticmethod - def _inject_cost_into_sse_frame_str(frame_str: str, model_name: str) -> str | None: + def _inject_cost_into_sse_frame_str( + frame_str: str, model_name: str, litellm_logging_obj: LiteLLMLoggingObj | None = None + ) -> str | None: """ Inject cost information into an SSE frame string by modifying the JSON in the 'data:' line. Args: frame_str: SSE frame string that may contain multiple lines model_name: Model name for cost calculation + litellm_logging_obj: The call's logging object, forwarded for pricing Returns: Modified SSE frame string with cost injected, or None if no modification needed @@ -3486,7 +3505,9 @@ class ProxyBaseLLMRequestProcessing: json_part = stripped_ln.split("data:", 1)[1].strip() if json_part and json_part != "[DONE]": obj = json.loads(json_part) - maybe_modified = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(obj, model_name) + maybe_modified = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict( + obj, model_name, litellm_logging_obj + ) if maybe_modified is not None: lines[idx] = "data: " + safe_dumps(maybe_modified) + ("\r" if ln.endswith("\r") else "") return "\n".join(lines) @@ -3494,34 +3515,6 @@ class ProxyBaseLLMRequestProcessing: except Exception: return None - @staticmethod - def _anthropic_stream_usage_kwargs(usage: Mapping[str, Any]) -> Mapping[str, Any]: - prompt_tokens: Final = int(usage.get("input_tokens", 0) or 0) - completion_tokens: Final = int(usage.get("output_tokens", 0) or 0) - total_tokens: Final = int( - usage.get("total_tokens", prompt_tokens + completion_tokens) or (prompt_tokens + completion_tokens) - ) - web_search_requests: Final = usage.get("web_search_requests") - server_tool_use: Final = ( - ServerToolUse(web_search_requests=web_search_requests) if web_search_requests is not None else None - ) - return MappingProxyType( - { - key: value - for key, value in ( - ("prompt_tokens", prompt_tokens), - ("completion_tokens", completion_tokens), - ("total_tokens", total_tokens), - ("completion_tokens_details", usage.get("completion_tokens_details")), - ("prompt_tokens_details", usage.get("prompt_tokens_details")), - ("cache_creation_input_tokens", usage.get("cache_creation_input_tokens")), - ("cache_read_input_tokens", usage.get("cache_read_input_tokens")), - ("server_tool_use", server_tool_use), - ) - if value is not None - } - ) - @staticmethod def _openai_stream_usage_kwargs(usage: Mapping[str, Any]) -> Mapping[str, Any]: prompt_tokens: Final = int(usage.get("prompt_tokens", 0) or 0) @@ -3544,11 +3537,19 @@ class ProxyBaseLLMRequestProcessing: ) @staticmethod - def _stream_usage_kwargs_for_event(obj: Mapping[str, object], usage: Mapping[str, Any]) -> Mapping[str, Any] | None: + def _stream_usage_for_event(obj: Mapping[str, object], usage: Mapping[str, Any]) -> Usage | None: + """ + Build the ``Usage`` to price a streamed usage event. + + Anthropic goes through ``AnthropicConfig.calculate_usage``, the same transformation + the non-streaming path uses, so ``prompt_tokens`` is the full input total and the + cache read plus 5m/1h cache creation split land in ``prompt_tokens_details`` where + the pricer looks for them. + """ if obj.get("type") == "message_delta": - return ProxyBaseLLMRequestProcessing._anthropic_stream_usage_kwargs(usage) + return AnthropicConfig().calculate_usage(usage_object=usage, reasoning_content=None) if obj.get("object") == "chat.completion.chunk": - return ProxyBaseLLMRequestProcessing._openai_stream_usage_kwargs(usage) + return Usage(**ProxyBaseLLMRequestProcessing._openai_stream_usage_kwargs(usage)) return None @staticmethod @@ -3563,7 +3564,41 @@ class ProxyBaseLLMRequestProcessing: return None @staticmethod - def _inject_cost_into_usage_dict(obj: dict, model_name: str) -> dict | None: + def _logging_obj_cost_or_none( + model_response: ModelResponse, litellm_logging_obj: LiteLLMLoggingObj + ) -> float | None: + try: + cost: Final = litellm_logging_obj._response_cost_calculator(result=model_response) # pyright: ignore[reportPrivateUsage] # reuse the call's own cost calc for pricing parity with the logging callback + except Exception: + return None + return float(cost) if isinstance(cost, (int, float)) and not isinstance(cost, bool) else None + + @staticmethod + def _streamed_usage_cost( + model_response: ModelResponse, + model_name: str, + service_tier: str | None, + litellm_logging_obj: LiteLLMLoggingObj | None, + ) -> float | None: + """ + Price a streamed usage response through the call's logging object when there is one, + so the streamed cost picks up the same custom/deployment pricing (``custom_pricing``, + ``custom_llm_provider``, ``base_model``, ``router_model_id``) as the logging callback + rather than the model's sticker price; fall back to ``completion_cost`` by model name. + """ + cost_from_logging_obj: Final = ( + ProxyBaseLLMRequestProcessing._logging_obj_cost_or_none(model_response, litellm_logging_obj) + if litellm_logging_obj is not None + else None + ) + if cost_from_logging_obj is not None: + return cost_from_logging_obj + return ProxyBaseLLMRequestProcessing._completion_cost_or_none(model_response, model_name, service_tier) + + @staticmethod + def _inject_cost_into_usage_dict( + obj: dict, model_name: str, litellm_logging_obj: LiteLLMLoggingObj | None = None + ) -> dict | None: """ Inject cost information into the usage object of a streamed usage event (Anthropic ``message_delta`` or OpenAI ``chat.completion.chunk``). @@ -3571,6 +3606,8 @@ class ProxyBaseLLMRequestProcessing: Args: obj: Dictionary containing the SSE event data model_name: Model name for cost calculation + litellm_logging_obj: The call's logging object, used so the injected cost + matches the cost the logging callback records Returns: Modified dictionary with cost injected, or None if no modification needed @@ -3578,14 +3615,15 @@ class ProxyBaseLLMRequestProcessing: usage: Final = obj.get("usage") if not isinstance(usage, dict): return None - usage_kwargs: Final = ProxyBaseLLMRequestProcessing._stream_usage_kwargs_for_event(obj, usage) - if usage_kwargs is None: + stream_usage: Final = ProxyBaseLLMRequestProcessing._stream_usage_for_event(obj, usage) + if stream_usage is None: return None service_tier: Final = obj.get("service_tier") - cost_val: Final = ProxyBaseLLMRequestProcessing._completion_cost_or_none( - ModelResponse(usage=Usage(**usage_kwargs)), + cost_val: Final = ProxyBaseLLMRequestProcessing._streamed_usage_cost( + ModelResponse(usage=stream_usage), model_name, service_tier if isinstance(service_tier, str) else None, + litellm_logging_obj, ) if cost_val is None: return None diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 697eb7b96eb..b71622fc33d 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -106,7 +106,7 @@ class PassThroughStreamingHandler: ) # rebind-ok: SSE frame reassembly buffer across transport chunks if complete_frames: yield ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( - complete_frames, resolved_model_name + complete_frames, resolved_model_name, litellm_logging_obj ) if pending: yield pending diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 355c6d27eb2..083982778b9 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -6082,6 +6082,121 @@ class TestInjectCostIntoUsageDict: injected = json.loads(result.split("\n")[0].split("data:", 1)[1].strip()) assert injected["usage"]["cost"] == pytest.approx(self._expected_cost("gpt-4o-mini", 11, 4)) + def test_message_delta_cost_charges_the_non_cached_input_tokens(self): + """Anthropic reports ``input_tokens`` excluding cache tokens, so reading it as the whole + prompt total drops the non-cached input from the bill on every cache hit.""" + model = "claude-haiku-4-5" + pricing = litellm.model_cost[model] + event = { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": { + "input_tokens": 14, + "output_tokens": 8, + "cache_read_input_tokens": 3202, + "cache_creation_input_tokens": 0, + }, + } + + result = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(event, model) + + assert result is not None + expected = ( + 14 * pricing["input_cost_per_token"] + + 3202 * pricing["cache_read_input_token_cost"] + + 8 * pricing["output_cost_per_token"] + ) + dropped_input = expected - 14 * pricing["input_cost_per_token"] + assert result["usage"]["cost"] == pytest.approx(expected) + assert result["usage"]["cost"] > dropped_input + + def test_message_delta_prices_1h_cache_creation_above_the_5m_rate(self): + """The ``cache_creation`` 5m/1h split has to survive into ``prompt_tokens_details``, + otherwise a 1h write is billed at the cheaper 5m rate.""" + model = "claude-haiku-4-5" + pricing = litellm.model_cost[model] + event = { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": { + "input_tokens": 14, + "output_tokens": 8, + "cache_read_input_tokens": 0, + "cache_creation_input_tokens": 2000, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 2000}, + }, + } + + result = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(event, model) + + assert result is not None + base = 14 * pricing["input_cost_per_token"] + 8 * pricing["output_cost_per_token"] + expected_1h = base + 2000 * pricing["cache_creation_input_token_cost_above_1hr"] + flat_5m = base + 2000 * pricing["cache_creation_input_token_cost"] + assert expected_1h != pytest.approx(flat_5m) + assert result["usage"]["cost"] == pytest.approx(expected_1h) + + def test_message_delta_prices_through_the_logging_obj_so_custom_pricing_applies(self): + """Costing by model name alone yields sticker price, so a deployment with a negotiated + discount streamed a ``usage.cost`` that disagreed with the callback's ``response_cost``.""" + + class _StubLoggingObj: + def __init__(self, cost): + self._cost = cost + self.captured_result = None + + def _response_cost_calculator(self, result): + self.captured_result = result + return self._cost + + model = "claude-haiku-4-5" + discounted_cost = 0.00099 + stub = _StubLoggingObj(discounted_cost) + event = { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": { + "input_tokens": 14, + "output_tokens": 8, + "cache_read_input_tokens": 3202, + "cache_creation_input_tokens": 500, + "cache_creation": {"ephemeral_5m_input_tokens": 100, "ephemeral_1h_input_tokens": 400}, + }, + } + + result = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(event, model, stub) + + assert result is not None + assert result["usage"]["cost"] == discounted_cost + assert result["usage"]["cost"] != pytest.approx(self._expected_cost(model, 14 + 500 + 3202, 8)) + usage = stub.captured_result.usage + assert usage.prompt_tokens == 14 + 500 + 3202 + details = usage.prompt_tokens_details.cache_creation_token_details + assert details.ephemeral_5m_input_tokens == 100 + assert details.ephemeral_1h_input_tokens == 400 + + def test_message_delta_falls_back_to_model_pricing_when_the_logging_obj_returns_no_cost(self): + class _StubLoggingObj: + def _response_cost_calculator(self, result): + return None + + model = "claude-haiku-4-5" + pricing = litellm.model_cost[model] + event = { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": {"input_tokens": 14, "output_tokens": 8, "cache_read_input_tokens": 3202}, + } + + result = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(event, model, _StubLoggingObj()) + + assert result is not None + assert result["usage"]["cost"] == pytest.approx( + 14 * pricing["input_cost_per_token"] + + 3202 * pricing["cache_read_input_token_cost"] + + 8 * pricing["output_cost_per_token"] + ) + class TestProcessChunkWithCostInjection: def test_complete_usage_frame_chunk_is_injected(self, monkeypatch): @@ -6116,6 +6231,31 @@ class TestProcessChunkWithCostInjection: assert ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection(chunk, "gpt-4o-mini") == chunk + def test_message_delta_frame_is_priced_with_the_logging_obj(self, monkeypatch): + """Pins that the logging object reaches the pricer through the byte-frame entry point, + which is how the proxy actually calls this on a streamed Messages API request.""" + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) + + class _StubLoggingObj: + def _response_cost_calculator(self, result): + return 0.00042 + + chunk = ( + b"event: message_delta\n" + b'data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},' + b'"usage":{"input_tokens":14,"output_tokens":8,"cache_read_input_tokens":3202}}\n\n' + ) + + result = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( + chunk, "claude-haiku-4-5", _StubLoggingObj() + ) + + assert result != chunk + data_line = next(ln for ln in result.decode("utf-8").splitlines() if ln.startswith("data:")) + payload = json.loads(data_line.split("data:", 1)[1].strip()) + assert payload["usage"]["cost"] == 0.00042 + assert payload["usage"]["cache_read_input_tokens"] == 3202 + # --------------------------------------------------------------------------- # SSE keepalive during the time-to-first-token (issue #34819) From 6f4844bfd9dd50901e3e4542f296602c9c8c9e56 Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 19 Aug 2026 03:58:24 +0000 Subject: [PATCH 014/166] refactor(cost): trim streamed cost helper docstrings to the non-obvious bits Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 24 ++++++---------------- 1 file changed, 6 insertions(+), 18 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index f5299b273b4..d533df8616d 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3440,8 +3440,7 @@ class ProxyBaseLLMRequestProcessing: Args: chunk: The streaming chunk (dict, str, bytes, or bytearray) model_name: Model name for cost calculation - litellm_logging_obj: The call's logging object, used to price the chunk with - the same custom/deployment pricing as the logging callback + litellm_logging_obj: The call's logging object, used for pricing Returns: The processed chunk with cost information injected if applicable @@ -3538,14 +3537,8 @@ class ProxyBaseLLMRequestProcessing: @staticmethod def _stream_usage_for_event(obj: Mapping[str, object], usage: Mapping[str, Any]) -> Usage | None: - """ - Build the ``Usage`` to price a streamed usage event. - - Anthropic goes through ``AnthropicConfig.calculate_usage``, the same transformation - the non-streaming path uses, so ``prompt_tokens`` is the full input total and the - cache read plus 5m/1h cache creation split land in ``prompt_tokens_details`` where - the pricer looks for them. - """ + # Anthropic reports input_tokens excluding cache tokens, so reuse the non-streaming + # transformation to total the prompt and keep the 5m/1h cache creation split if obj.get("type") == "message_delta": return AnthropicConfig().calculate_usage(usage_object=usage, reasoning_content=None) if obj.get("object") == "chat.completion.chunk": @@ -3580,12 +3573,8 @@ class ProxyBaseLLMRequestProcessing: service_tier: str | None, litellm_logging_obj: LiteLLMLoggingObj | None, ) -> float | None: - """ - Price a streamed usage response through the call's logging object when there is one, - so the streamed cost picks up the same custom/deployment pricing (``custom_pricing``, - ``custom_llm_provider``, ``base_model``, ``router_model_id``) as the logging callback - rather than the model's sticker price; fall back to ``completion_cost`` by model name. - """ + # Pricing via the logging object inherits the deployment's custom pricing, so the + # streamed cost matches what the logging callback records instead of sticker price cost_from_logging_obj: Final = ( ProxyBaseLLMRequestProcessing._logging_obj_cost_or_none(model_response, litellm_logging_obj) if litellm_logging_obj is not None @@ -3606,8 +3595,7 @@ class ProxyBaseLLMRequestProcessing: Args: obj: Dictionary containing the SSE event data model_name: Model name for cost calculation - litellm_logging_obj: The call's logging object, used so the injected cost - matches the cost the logging callback records + litellm_logging_obj: The call's logging object, used for pricing Returns: Modified dictionary with cost injected, or None if no modification needed From b6ee13803d1c21bc0b07dac6e6b9ad2ede7ad7aa Mon Sep 17 00:00:00 2001 From: HarryZhou <2373256746@qq.com> Date: Sun, 9 Aug 2026 22:59:25 +0800 Subject: [PATCH 015/166] fix(responses-bridge): preserve reasoning input items as reasoning_content --- .../transformation.py | 164 +++++++++++++++++- .../test_reasoning_input_item_preservation.py | 147 ++++++++++++++++ 2 files changed, 310 insertions(+), 1 deletion(-) create mode 100644 tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 4892e3b348c..5d3ed0477e3 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -557,7 +557,108 @@ class LiteLLMCompletionResponsesConfig: continue messages.extend(chat_completion_messages) - return messages + return LiteLLMCompletionResponsesConfig._merge_reasoning_only_assistant_messages(messages) + + @staticmethod + def _merge_reasoning_only_assistant_messages( + messages: list[ + AllMessageValues + | GenericChatCompletionMessage + | ChatCompletionMessageToolCall + | ChatCompletionResponseMessage + ], + ) -> list[ + AllMessageValues | GenericChatCompletionMessage | ChatCompletionMessageToolCall | ChatCompletionResponseMessage + ]: + """ + Responses API emits prior-turn reasoning as its own ``reasoning`` input + item, which becomes a standalone assistant message with + ``content=None`` + ``reasoning_content``. Chat-completions providers + (e.g. DeepSeek V4, Kimi K2.6) expect the chain-of-thought on the + assistant message that carries the answer or tool calls. This pass + merges standalone reasoning-only assistant messages into the + immediately following assistant message. + + If the reasoning item is not followed by an assistant message (e.g. a + stateless chain replays ``reasoning`` + ``user``), the standalone + reasoning message is preserved so the reasoning is still passed back. + """ + + def _role(msg: Any) -> str: + if isinstance(msg, dict): + return str(msg.get("role") or "") + return str(getattr(msg, "role", "") or "") + + def _reasoning_text(msg: Any) -> str | None: + if isinstance(msg, dict): + value = msg.get("reasoning_content") + else: + value = getattr(msg, "reasoning_content", None) + return value if isinstance(value, str) and value else None + + def _content(msg: Any) -> Any: + if isinstance(msg, dict): + return msg.get("content") + return getattr(msg, "content", None) + + def _tool_calls(msg: Any) -> Any: + if isinstance(msg, dict): + return msg.get("tool_calls") + return getattr(msg, "tool_calls", None) + + merged: list[ + AllMessageValues + | GenericChatCompletionMessage + | ChatCompletionMessageToolCall + | ChatCompletionResponseMessage + ] = [] + pending_reasoning: list[str] = [] + + for msg in messages: + if ( + _role(msg) == "assistant" + and _content(msg) is None + and not _tool_calls(msg) + and _reasoning_text(msg) is not None + ): + pending_reasoning.append(_reasoning_text(msg) or "") + continue + + if pending_reasoning and _role(msg) == "assistant": + combined = "\n".join(pending_reasoning) + existing = _reasoning_text(msg) + if existing: + combined = existing + "\n" + combined + if isinstance(msg, dict): + msg["reasoning_content"] = combined + else: + setattr(msg, "reasoning_content", combined) + pending_reasoning = [] + elif pending_reasoning: + # Not followed by an assistant message — keep the reasoning + # standalone instead of dropping it. + for text in pending_reasoning: + merged.append( + ChatCompletionResponseMessage( + role="assistant", + content=None, + reasoning_content=text, + ) + ) + pending_reasoning = [] + + merged.append(msg) + + for text in pending_reasoning: + merged.append( + ChatCompletionResponseMessage( + role="assistant", + content=None, + reasoning_content=text, + ) + ) + + return merged @staticmethod def _merged_trailing_assistant_message( @@ -1026,6 +1127,25 @@ class LiteLLMCompletionResponsesConfig: return LiteLLMCompletionResponsesConfig._transform_responses_api_function_call_to_chat_completion_message( function_call=input_item ) + elif input_item.get("type") == "reasoning": + # A ResponseReasoningItemParam carries the prior-turn chain-of-thought. + # Chat-completions providers (DeepSeek V4, Kimi K2.6, ...) expect this + # to be replayed as `reasoning_content` on an assistant message, not as + # visible `content` (prompt pollution) and not dropped (DeepSeek V4 + # rejects multi-turn requests with a missing `reasoning_content`). + reasoning_text = LiteLLMCompletionResponsesConfig._extract_reasoning_text_from_input_item(input_item) + if not reasoning_text: + # No plaintext reasoning is available (e.g. encrypted_content only). + # Chat-completions providers cannot consume opaque encrypted blobs, + # so skip the item instead of polluting the prompt. + return [] + return [ + ChatCompletionResponseMessage( + role="assistant", + content=None, + reasoning_content=reasoning_text, + ) + ] else: content: Final[object] = input_item.get("content") # Handle None content: Responses API allows None content, but GenericChatCompletionMessage requires content @@ -1041,6 +1161,48 @@ class LiteLLMCompletionResponsesConfig: ) ] + @staticmethod + def _extract_reasoning_text_from_input_item(input_item: Mapping[str, object]) -> str | None: + """ + Extract plaintext reasoning from a ResponseReasoningItemParam. + + Handles: + - content as a string + - content as a list of blocks (output_text / summary_text / text) + - summary as a list of summary_text blocks (fallback) + + Returns None when only opaque forms (e.g. encrypted_content) are present. + """ + content: Final[object] = input_item.get("content") + if isinstance(content, str) and content.strip(): + return content + if isinstance(content, list): + text_parts: list[str] = [] + for block in content: + if not isinstance(block, Mapping): + continue + block_type = block.get("type") + if block_type in ("encrypted_content", "redacted_thinking"): + continue + text = block.get("text") + if isinstance(text, str) and text.strip(): + text_parts.append(text.strip()) + if text_parts: + return "\n".join(text_parts) + + summary: Final[object] = input_item.get("summary") + if isinstance(summary, list): + text_parts = [] + for block in summary: + if not isinstance(block, Mapping): + continue + text = block.get("text") + if isinstance(text, str) and text.strip(): + text_parts.append(text.strip()) + if text_parts: + return "\n".join(text_parts) + return None + @staticmethod def _is_input_item_tool_call_output(input_item: Mapping[str, object]) -> bool: """ diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py new file mode 100644 index 00000000000..5fcd4df3ff8 --- /dev/null +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py @@ -0,0 +1,147 @@ +""" +Unit tests for preserving prior-turn ``reasoning`` input items when the +Responses API is bridged to chat completions. + +Without this handling, a ``ResponseReasoningItemParam`` falls through to the +generic message branch, polluting the prompt as visible assistant ``content`` +or being silently dropped. Chat-completions providers such as DeepSeek V4 and +Kimi K2.6 require the chain-of-thought to be replayed as ``reasoning_content`` +on an assistant message. +""" + +from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, +) + + +def _transform_item(item): + return LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message( + input_item=item + ) + + +def _transform_input(input_items): + return LiteLLMCompletionResponsesConfig._transform_response_input_param_to_chat_completion_message( + input=input_items + ) + + +class TestReasoningInputItemHandler: + """Reasoning input items map to assistant ``reasoning_content``.""" + + def test_reasoning_item_with_output_text_content(self): + """Standard Responses-API reasoning item with output_text blocks.""" + item = { + "type": "reasoning", + "id": "rs_abc", + "summary": [], + "content": [{"type": "output_text", "text": "step 1: think about X"}], + } + messages = _transform_item(item) + assert len(messages) == 1 + assert messages[0]["role"] == "assistant" + assert messages[0]["content"] is None + assert messages[0]["reasoning_content"] == "step 1: think about X" + + def test_reasoning_item_with_string_content(self): + """Variant: reasoning content as a plain string.""" + item = {"type": "reasoning", "id": "rs_1", "content": "step 1: ..."} + messages = _transform_item(item) + assert messages[0]["reasoning_content"] == "step 1: ..." + + def test_reasoning_item_with_summary_only(self): + """SDK form: reasoning carried in summary list, no content.""" + item = { + "type": "reasoning", + "id": "rs_2", + "summary": [{"type": "summary_text", "text": "..."}], + } + messages = _transform_item(item) + assert messages[0]["reasoning_content"] == "..." + + def test_reasoning_item_with_encrypted_content_only_dropped(self): + """Opaque encrypted reasoning cannot be forwarded to chat completions.""" + item = {"type": "reasoning", "id": "rs_3", "encrypted_content": "opaque-blob"} + assert _transform_item(item) == [] + + def test_reasoning_item_empty_dropped(self): + """Reasoning item with neither content nor summary drops cleanly.""" + assert _transform_item({"type": "reasoning", "id": "rs_4"}) == [] + + +class TestReasoningInputItemMerging: + """Standalone reasoning messages merge into the following assistant turn.""" + + def test_reasoning_merged_into_following_assistant_message(self): + """Reasoning + assistant answer become one assistant message.""" + messages = _transform_input( + [ + { + "type": "reasoning", + "id": "rs_1", + "content": [{"type": "output_text", "text": "secret reasoning"}], + }, + {"type": "message", "role": "assistant", "content": "The answer."}, + ] + ) + assert len(messages) == 1 + assert messages[0]["role"] == "assistant" + assert messages[0]["content"] == "The answer." + assert messages[0]["reasoning_content"] == "secret reasoning" + + def test_reasoning_preserved_when_followed_by_user_message(self): + """Stateless chain: reasoning + user prompt keeps the reasoning turn.""" + messages = _transform_input( + [ + { + "type": "reasoning", + "id": "rs_1", + "content": [{"type": "output_text", "text": "secret BLUEBERRY"}], + }, + {"role": "user", "content": "What is the secret word?"}, + ] + ) + assert len(messages) == 2 + assert messages[0]["role"] == "assistant" + assert messages[0]["content"] is None + assert messages[0]["reasoning_content"] == "secret BLUEBERRY" + assert messages[1]["role"] == "user" + + def test_reasoning_merged_into_function_call_assistant(self): + """Reasoning + function_call becomes one assistant tool-call message.""" + messages = _transform_input( + [ + { + "type": "reasoning", + "id": "rs_1", + "content": [{"type": "output_text", "text": "I should look this up"}], + }, + { + "type": "function_call", + "call_id": "call_1", + "name": "lookup", + "arguments": '{"cwe": "79"}', + }, + ] + ) + assert len(messages) == 1 + assert messages[0]["role"] == "assistant" + assert messages[0]["reasoning_content"] == "I should look this up" + assert len(messages[0]["tool_calls"]) == 1 + + +class TestNonReasoningInputItemUnchanged: + """Non-reasoning items still flow through the existing branches.""" + + def test_user_message_unchanged(self): + item = {"role": "user", "content": "hello"} + out = _transform_item(item) + assert len(out) == 1 + assert out[0]["role"] == "user" + + def test_assistant_message_unchanged(self): + item = {"role": "assistant", "content": "hi"} + out = _transform_item(item) + assert len(out) == 1 + assert out[0]["role"] == "assistant" + assert out[0]["content"] == "hi" From 3a77556dc14660e88a7d20f54e9762c39f24b749 Mon Sep 17 00:00:00 2001 From: HarryZhou <2373256746@qq.com> Date: Mon, 10 Aug 2026 00:45:25 +0800 Subject: [PATCH 016/166] fix(responses-bridge): preserve reasoning merge order when assistant already has reasoning_content --- .../transformation.py | 2 +- .../test_reasoning_input_item_preservation.py | 12 ++++++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 5d3ed0477e3..0604c3636ff 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -628,7 +628,7 @@ class LiteLLMCompletionResponsesConfig: combined = "\n".join(pending_reasoning) existing = _reasoning_text(msg) if existing: - combined = existing + "\n" + combined + combined = combined + "\n" + existing if isinstance(msg, dict): msg["reasoning_content"] = combined else: diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py index 5fcd4df3ff8..ecc024b7d04 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py @@ -129,6 +129,18 @@ class TestReasoningInputItemMerging: assert messages[0]["reasoning_content"] == "I should look this up" assert len(messages[0]["tool_calls"]) == 1 + def test_reasoning_merged_into_assistant_with_existing_reasoning_content(self): + """Old reasoning precedes existing reasoning on the target assistant turn.""" + messages = LiteLLMCompletionResponsesConfig._merge_reasoning_only_assistant_messages( + [ + {"role": "assistant", "content": None, "reasoning_content": "old reasoning"}, + {"role": "assistant", "content": "The answer.", "reasoning_content": "new reasoning"}, + ] + ) + assert len(messages) == 1 + assert messages[0]["content"] == "The answer." + assert messages[0]["reasoning_content"] == "old reasoning\nnew reasoning" + class TestNonReasoningInputItemUnchanged: """Non-reasoning items still flow through the existing branches.""" From 5911124f1dbba1e9c58f3b53619c3f875752a20f Mon Sep 17 00:00:00 2001 From: HarryZhou <2373256746@qq.com> Date: Mon, 10 Aug 2026 00:58:47 +0800 Subject: [PATCH 017/166] fix(responses-bridge): satisfy ruff strict-rule budget in reasoning merge --- .../transformation.py | 26 +++++++++++-------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 0604c3636ff..2c506d4a4c7 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -584,24 +584,24 @@ class LiteLLMCompletionResponsesConfig: reasoning message is preserved so the reasoning is still passed back. """ - def _role(msg: Any) -> str: + def _role(msg: object) -> str: if isinstance(msg, dict): return str(msg.get("role") or "") return str(getattr(msg, "role", "") or "") - def _reasoning_text(msg: Any) -> str | None: + def _reasoning_text(msg: object) -> str | None: if isinstance(msg, dict): value = msg.get("reasoning_content") else: value = getattr(msg, "reasoning_content", None) return value if isinstance(value, str) and value else None - def _content(msg: Any) -> Any: + def _content(msg: object) -> object | None: if isinstance(msg, dict): return msg.get("content") return getattr(msg, "content", None) - def _tool_calls(msg: Any) -> Any: + def _tool_calls(msg: object) -> object | None: if isinstance(msg, dict): return msg.get("tool_calls") return getattr(msg, "tool_calls", None) @@ -632,31 +632,35 @@ class LiteLLMCompletionResponsesConfig: if isinstance(msg, dict): msg["reasoning_content"] = combined else: - setattr(msg, "reasoning_content", combined) + setattr(msg, "reasoning_content", combined) # noqa: B010 pending_reasoning = [] elif pending_reasoning: # Not followed by an assistant message — keep the reasoning # standalone instead of dropping it. - for text in pending_reasoning: - merged.append( + merged.extend( + [ ChatCompletionResponseMessage( role="assistant", content=None, reasoning_content=text, ) - ) + for text in pending_reasoning + ] + ) pending_reasoning = [] merged.append(msg) - for text in pending_reasoning: - merged.append( + merged.extend( + [ ChatCompletionResponseMessage( role="assistant", content=None, reasoning_content=text, ) - ) + for text in pending_reasoning + ] + ) return merged From 438c1850fe1223feec1e2e6e5b48f0a6c15a1328 Mon Sep 17 00:00:00 2001 From: HarryZhou <2373256746@qq.com> Date: Mon, 10 Aug 2026 01:09:26 +0800 Subject: [PATCH 018/166] fix(responses-bridge): satisfy type-discipline budget in reasoning merge --- .../transformation.py | 34 ++++++++++--------- 1 file changed, 18 insertions(+), 16 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 2c506d4a4c7..e3e62ab3c55 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -561,13 +561,13 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _merge_reasoning_only_assistant_messages( - messages: list[ + messages: list[ # mutable-ok: input sequence AllMessageValues | GenericChatCompletionMessage | ChatCompletionMessageToolCall | ChatCompletionResponseMessage ], - ) -> list[ + ) -> list[ # mutable-ok: fresh merged list AllMessageValues | GenericChatCompletionMessage | ChatCompletionMessageToolCall | ChatCompletionResponseMessage ]: """ @@ -591,9 +591,9 @@ class LiteLLMCompletionResponsesConfig: def _reasoning_text(msg: object) -> str | None: if isinstance(msg, dict): - value = msg.get("reasoning_content") + value = msg.get("reasoning_content") # rebind-ok: branch lookup else: - value = getattr(msg, "reasoning_content", None) + value = getattr(msg, "reasoning_content", None) # rebind-ok: branch lookup return value if isinstance(value, str) and value else None def _content(msg: object) -> object | None: @@ -606,13 +606,13 @@ class LiteLLMCompletionResponsesConfig: return msg.get("tool_calls") return getattr(msg, "tool_calls", None) - merged: list[ + merged: list[ # mutable-ok: accumulator # rebind-ok: accumulator AllMessageValues | GenericChatCompletionMessage | ChatCompletionMessageToolCall | ChatCompletionResponseMessage - ] = [] - pending_reasoning: list[str] = [] + ] = [] # mutable-ok: accumulator + pending_reasoning: list[str] = [] # mutable-ok: accumulator # rebind-ok: accumulator for msg in messages: if ( @@ -633,11 +633,11 @@ class LiteLLMCompletionResponsesConfig: msg["reasoning_content"] = combined else: setattr(msg, "reasoning_content", combined) # noqa: B010 - pending_reasoning = [] + pending_reasoning = [] # mutable-ok: reset accumulator elif pending_reasoning: # Not followed by an assistant message — keep the reasoning # standalone instead of dropping it. - merged.extend( + merged.extend( # mutable-ok: append reasoning messages [ ChatCompletionResponseMessage( role="assistant", @@ -647,11 +647,11 @@ class LiteLLMCompletionResponsesConfig: for text in pending_reasoning ] ) - pending_reasoning = [] + pending_reasoning = [] # mutable-ok: reset accumulator merged.append(msg) - merged.extend( + merged.extend( # mutable-ok: append trailing reasoning [ ChatCompletionResponseMessage( role="assistant", @@ -1137,13 +1137,15 @@ class LiteLLMCompletionResponsesConfig: # to be replayed as `reasoning_content` on an assistant message, not as # visible `content` (prompt pollution) and not dropped (DeepSeek V4 # rejects multi-turn requests with a missing `reasoning_content`). - reasoning_text = LiteLLMCompletionResponsesConfig._extract_reasoning_text_from_input_item(input_item) + reasoning_text = LiteLLMCompletionResponsesConfig._extract_reasoning_text_from_input_item( # rebind-ok: extraction result + input_item + ) if not reasoning_text: # No plaintext reasoning is available (e.g. encrypted_content only). # Chat-completions providers cannot consume opaque encrypted blobs, # so skip the item instead of polluting the prompt. - return [] - return [ + return [] # mutable-ok: empty drop result + return [ # mutable-ok: single message result ChatCompletionResponseMessage( role="assistant", content=None, @@ -1181,7 +1183,7 @@ class LiteLLMCompletionResponsesConfig: if isinstance(content, str) and content.strip(): return content if isinstance(content, list): - text_parts: list[str] = [] + text_parts: list[str] = [] # mutable-ok: text accumulator # rebind-ok: text accumulator for block in content: if not isinstance(block, Mapping): continue @@ -1196,7 +1198,7 @@ class LiteLLMCompletionResponsesConfig: summary: Final[object] = input_item.get("summary") if isinstance(summary, list): - text_parts = [] + text_parts = [] # mutable-ok: text accumulator # rebind-ok: text accumulator for block in summary: if not isinstance(block, Mapping): continue From de95372dfbd7bbba8c478815340dd49c1b21da11 Mon Sep 17 00:00:00 2001 From: HarryZhou <2373256746@qq.com> Date: Mon, 10 Aug 2026 01:24:24 +0800 Subject: [PATCH 019/166] fix(responses-bridge): type-safe reasoning_content assignment in merge pass --- .../litellm_completion_transformation/transformation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index e3e62ab3c55..d7b6b8c7b8f 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -630,7 +630,7 @@ class LiteLLMCompletionResponsesConfig: if existing: combined = combined + "\n" + existing if isinstance(msg, dict): - msg["reasoning_content"] = combined + cast(dict[str, Any], msg)["reasoning_content"] = combined # cast-ok: mutable reasoning carrier else: setattr(msg, "reasoning_content", combined) # noqa: B010 pending_reasoning = [] # mutable-ok: reset accumulator From 2d4e6afe1c7d6d3233a18668d53fafe4cffa50b3 Mon Sep 17 00:00:00 2001 From: HarryZhou <2373256746@qq.com> Date: Tue, 18 Aug 2026 21:11:29 +0800 Subject: [PATCH 020/166] fix(guardrails): inspect responses reasoning content and summary text --- litellm/proxy/guardrails/_content_utils.py | 59 +++++++++++++------ .../transformation.py | 6 +- .../proxy/guardrails/test_content_utils.py | 54 +++++++++++++++++ 3 files changed, 100 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index 6ed6f0013df..ae92adcb1ee 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -8,7 +8,7 @@ skip the other shapes — these helpers normalise that so every hook sees every text fragment. """ -from collections.abc import Callable, Iterator +from collections.abc import Callable, Iterator, Mapping from typing import Any, Final # Call types whose body carries free-form chat / prompt text that @@ -33,7 +33,9 @@ def is_text_content_call_type(call_type: str) -> bool: return call_type in TEXT_CONTENT_CALL_TYPES -TEXT_PART_TYPES: Final[frozenset[str]] = frozenset({"text", "input_text", "output_text"}) +TEXT_PART_TYPES: Final[frozenset[str]] = frozenset( + {"text", "input_text", "output_text", "summary_text", "reasoning_text"} +) # Responses-API item types whose ``output`` field carries user/tool text # that guardrails should inspect. ``function_call_output`` is the @@ -42,6 +44,16 @@ TEXT_PART_TYPES: Final[frozenset[str]] = frozenset({"text", "input_text", "outpu _OUTPUT_ITEM_TYPES: Final[frozenset[str]] = frozenset({"function_call_output", "custom_tool_call_output"}) +def _part_text(part: Mapping[str, object]) -> str | None: + """Return non-empty plaintext from any content part that carries ``text``.""" + if not isinstance(part, dict): + return None + text = part.get("text") + if isinstance(text, str) and text: + return text + return None + + def _iter_text_parts_in_content(content: Any) -> Iterator[str]: """Yield text fragments from a ``message.content`` value (string or multimodal list). Non-text parts (images, audio, …) are skipped.""" @@ -58,10 +70,9 @@ def _iter_text_parts_in_content(content: Any) -> Iterator[str]: continue if not isinstance(part, dict): continue - if part.get("type") in TEXT_PART_TYPES: - text = part.get("text") - if isinstance(text, str) and text: - yield text + text = _part_text(part) + if text is not None: + yield text def _coerce_input_to_messages(input_value: Any) -> list[dict[str, Any]]: @@ -75,8 +86,23 @@ def _coerce_input_to_messages(input_value: Any) -> list[dict[str, Any]]: if isinstance(item, str): messages.append({"role": "user", "content": item}) elif isinstance(item, dict): - if item.get("type") in TEXT_PART_TYPES: + if _part_text(item) is not None: messages.append({"role": item.get("role") or "user", "content": [item]}) + elif item.get("type") == "reasoning": + if "content" in item: + messages.append( + { # mutable-ok: append reasoning content + "role": item.get("role") or "assistant", + "content": item["content"], + } + ) + if isinstance(item.get("summary"), list): + messages.append( + { # mutable-ok: append reasoning summary + "role": item.get("role") or "assistant", + "content": item["summary"], + } + ) elif "content" in item: messages.append({"role": item.get("role") or "user", "content": item["content"]}) elif item.get("type") in _OUTPUT_ITEM_TYPES and "output" in item: @@ -126,12 +152,7 @@ def walk_user_text(data: dict[str, Any], visit: Callable[[str], str]) -> int: if isinstance(part, str) and part: visited += 1 new_parts.append(visit(part)) - elif ( - isinstance(part, dict) - and part.get("type") in TEXT_PART_TYPES - and isinstance(part.get("text"), str) - and part["text"] - ): + elif isinstance(part, dict) and _part_text(part) is not None: visited += 1 new_parts.append({**part, "text": visit(part["text"])}) else: @@ -158,10 +179,14 @@ def walk_user_text(data: dict[str, Any], visit: Callable[[str], str]) -> int: visited += 1 input_value[idx] = visit(item) elif isinstance(item, dict): - if item.get("type") in TEXT_PART_TYPES: - if isinstance(item.get("text"), str) and item["text"]: - visited += 1 - input_value[idx] = {**item, "text": visit(item["text"])} + if _part_text(item) is not None: + visited += 1 + input_value[idx] = {**item, "text": visit(item["text"])} # mutable-ok: rewrite text part in place + elif item.get("type") == "reasoning": + if "content" in item: + item["content"] = _rewrite_content(item["content"]) + if isinstance(item.get("summary"), list): + item["summary"] = _rewrite_content(item["summary"]) elif "content" in item: item["content"] = _rewrite_content(item["content"]) elif item.get("type") in _OUTPUT_ITEM_TYPES and "output" in item: diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index d7b6b8c7b8f..c5f40242bfd 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -638,7 +638,7 @@ class LiteLLMCompletionResponsesConfig: # Not followed by an assistant message — keep the reasoning # standalone instead of dropping it. merged.extend( # mutable-ok: append reasoning messages - [ + [ # mutable-ok: append reasoning messages ChatCompletionResponseMessage( role="assistant", content=None, @@ -652,7 +652,7 @@ class LiteLLMCompletionResponsesConfig: merged.append(msg) merged.extend( # mutable-ok: append trailing reasoning - [ + [ # mutable-ok: append trailing reasoning ChatCompletionResponseMessage( role="assistant", content=None, @@ -1196,6 +1196,8 @@ class LiteLLMCompletionResponsesConfig: if text_parts: return "\n".join(text_parts) + # Guardrail traversal in litellm/proxy/guardrails/_content_utils.py + # inspects and rewrites these summary blocks before they are forwarded. summary: Final[object] = input_item.get("summary") if isinstance(summary, list): text_parts = [] # mutable-ok: text accumulator # rebind-ok: text accumulator diff --git a/tests/test_litellm/proxy/guardrails/test_content_utils.py b/tests/test_litellm/proxy/guardrails/test_content_utils.py index 3dfb98c12ea..d9e079c6d92 100644 --- a/tests/test_litellm/proxy/guardrails/test_content_utils.py +++ b/tests/test_litellm/proxy/guardrails/test_content_utils.py @@ -149,6 +149,22 @@ def test_iter_message_text_responses_api_tool_call_taxonomy(): assert list(iter_message_text(data)) == ["hello", "sunny"] +def test_iter_message_text_inspects_reasoning_content_and_summary(): + """VERIA: reasoning items forwarded as ``reasoning_content`` must be + inspected, including ``summary`` blocks the bridge reads as a fallback.""" + data = { + "input": [ + { + "type": "reasoning", + "id": "rs_1", + "content": [{"type": "summary_text", "text": "content secret"}], + "summary": [{"type": "summary_text", "text": "summary secret"}], + } + ] + } + assert list(iter_message_text(data)) == ["content secret", "summary secret"] + + # ── walk_user_text ──────────────────────────────────────────────────────────── @@ -308,6 +324,27 @@ def test_walk_user_text_redacts_mixed_list_input(): assert data["input"][2] == {"type": "image_url", "image_url": {"url": "..."}} +def test_walk_user_text_redacts_reasoning_content_and_summary(): + """VERIA: in-place redaction must cover both plaintext shapes the bridge + forwards from a reasoning item.""" + data = { + "input": [ + { + "type": "reasoning", + "id": "rs_1", + "content": [{"type": "summary_text", "text": "AKIAEXAMPLE content"}], + "summary": [{"type": "summary_text", "text": "AKIAEXAMPLE summary"}], + } + ] + } + visited = walk_user_text(data, lambda s: s.replace("AKIAEXAMPLE", "[REDACTED]")) + assert visited == 2 + item = data["input"][0] + assert item["content"][0]["text"] == "[REDACTED] content" + assert item["summary"][0]["text"] == "[REDACTED] summary" + assert item["id"] == "rs_1" + + # ── build_inspection_messages ───────────────────────────────────────────────── @@ -462,6 +499,23 @@ def test_build_inspection_messages_empty_data(): assert build_inspection_messages({"input": ""}) == [] +def test_build_inspection_messages_includes_reasoning_summary(): + """VERIA: remote guardrail APIs must see reasoning summaries even when + the reasoning item has no ``content`` field.""" + data = { + "input": [ + { + "type": "reasoning", + "id": "rs_1", + "summary": [{"type": "summary_text", "text": "secret summary"}], + } + ] + } + assert build_inspection_messages(data) == [ + {"role": "assistant", "content": "secret summary"} + ] + + # ── has_non_string_content ──────────────────────────────────────────────────── From 34e692c903c9d75b52065b4093c7d80b7eb2e00b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 20:16:22 -0700 Subject: [PATCH 021/166] fix(batches): decode model-encoded output file id so completed batches book spend Adds e2e coverage for batches terminal state and cost write-back, failure paths, per-backend file content downloads, and two-gateway routing (LIT-5730). --- litellm/batches/batch_utils.py | 40 +- tests/e2e/batches/COVERAGE.md | 87 +++- tests/e2e/batches/batch_client.py | 29 +- tests/e2e/batches/capabilities.py | 10 + tests/e2e/batches/test_batches_e2e.py | 471 +++++++++++++++++- .../llm_nonconversational.yaml | 10 + .../test_litellm/batches/test_batch_utils.py | 37 ++ 7 files changed, 646 insertions(+), 38 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 0cf22d82ca6..6eb13d2cba7 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -296,6 +296,32 @@ def calculate_vertex_ai_batch_cost_and_usage( ) +def _provider_output_file_id(output_file_id: str) -> str: + """ + Resolve the file id the provider actually knows: unified ids yield their embedded + llm_output_file_id, model-encoded ids decode to the raw provider id, raw ids pass through. + """ + from litellm.proxy.openai_files_endpoints.common_utils import ( + _is_base64_encoded_unified_file_id, + get_original_file_id, + ) + + unified_file_id: Final = _is_base64_encoded_unified_file_id(output_file_id) + if not unified_file_id: + return get_original_file_id(output_file_id) + try: + extracted: Final = unified_file_id.split("llm_output_file_id,")[1].split(";")[0] + except (IndexError, AttributeError) as e: + verbose_logger.error( + "Failed to extract LLM output file ID from unified file ID: %s, error: %s", + output_file_id, + e, + ) + return output_file_id + verbose_logger.debug("Extracted LLM output file ID from unified file ID: %s", extracted) + return extracted + + async def _fetch_batch_output_file_content( batch: Batch, custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai", @@ -311,23 +337,11 @@ async def _fetch_batch_output_file_content( Required for Azure and other providers that need authentication """ from litellm.files.main import afile_content - from litellm.proxy.openai_files_endpoints.common_utils import ( - _is_base64_encoded_unified_file_id, - ) if batch.output_file_id is None: raise ValueError("Output file id is None cannot retrieve file content") - file_id = batch.output_file_id - is_base64_unified_file_id: Final = _is_base64_encoded_unified_file_id(file_id) - if is_base64_unified_file_id: - try: - file_id = is_base64_unified_file_id.split("llm_output_file_id,")[1].split(";")[0] - verbose_logger.debug("Extracted LLM output file ID from unified file ID: %s", file_id) - except (IndexError, AttributeError) as e: - verbose_logger.error( - "Failed to extract LLM output file ID from unified file ID: %s, error: %s", batch.output_file_id, e - ) + file_id: Final = _provider_output_file_id(batch.output_file_id) # Build kwargs for afile_content with credentials from litellm_params file_content_kwargs: Final = { diff --git a/tests/e2e/batches/COVERAGE.md b/tests/e2e/batches/COVERAGE.md index ca48204962a..f02d4eb4fe4 100644 --- a/tests/e2e/batches/COVERAGE.md +++ b/tests/e2e/batches/COVERAGE.md @@ -1,9 +1,11 @@ # Batches Test Coverage Matrix Live e2e coverage of the Batches API over a real proxy, real provider keys, and -real cost. Synchronous tier only: a batch's completion window is 24h, so these -tests never wait for `completed`. They assert the proxy accepts, routes, retrieves, -cancels, and lists a batch; everything created is deleted on teardown. +real cost. Mostly synchronous tier: a batch's completion window is 24h, so the +lifecycle matrix never waits for `completed`. It asserts the proxy accepts, routes, +retrieves, cancels, and lists a batch; everything created is deleted on teardown. +The exception is `TestBatchTerminalState`, which covers the completed state and +cost write-back via a cross-run marker baton (design below). ## Provider x operation @@ -12,19 +14,26 @@ row per supported (provider, scenario) pair, so there are no skipped cells in th parametrized run. The batches suite never skips: missing provider creds or upstream failures are hard test failures (see `tests/e2e/CLAUDE.md`). -| Provider | create | retrieve | cancel | list | file backing | -|-----------|--------|----------|--------|------|--------------| -| OpenAI | yes | yes | yes | yes | OpenAI Files | -| Azure | yes | yes | yes | yes | Azure Files | -| Vertex AI | yes | yes | yes | yes | GCS (`gcs_bucket_name` / `GCS_BUCKET_NAME` on model) | -| Bedrock | yes (unified only) | yes | no (limited upstream) | no | S3 (`s3_bucket_name` + `aws_*` + `AWS_BATCH_ROLE_ARN` on model) | +| Provider | create | retrieve | cancel | list | content download | file backing | +|-----------|--------|----------|--------|------|------------------|--------------| +| OpenAI | yes | yes | yes | yes | yes (lifecycle + terminal output) | OpenAI Files | +| Azure | yes | yes | yes | yes | yes (byte-verbatim) | Azure Files | +| Vertex AI | yes | yes | yes | yes | yes (provider-transformed) | GCS (`gcs_bucket_name` / `GCS_BUCKET_NAME` on model) | +| Bedrock | yes (unified only) | yes | no (limited upstream) | no | yes (provider-transformed) | S3 (`s3_bucket_name` + `aws_*` + `AWS_BATCH_ROLE_ARN` on model) | Bedrock cancel is unreliable upstream and list is unsupported, so both are gated off -(`can_cancel=False`, `can_list=False`) when that provider is enabled in the matrix. +(`can_cancel=False`, `can_list=False`) when that provider is enabled in the matrix; +flipping those gates is tracked in LIT-4774 and deliberately not part of this suite. Bedrock file upload requires a model on the request (`encoded` / `unified` scenarios only); `model_param` and `provider_fallback` are omitted because `POST /bedrock/v1/files` has no model-less passthrough path. +`GET /v1/files/{id}/content` is exercised for the unified upload path per backend in +`test_unified_file_content_downloads`. Azure stores the JSONL verbatim, so its download +is asserted byte-equal to the upload. Vertex (GCS) and Bedrock (S3) transform lines at +upload time, so those assert a 200 with non-empty parseable JSON lines instead. Gemini +(non-Vertex) raises `NotImplementedError` for file content and has no cell here. + ## Routing scenarios (per `litellm/proxy/batches_endpoints/endpoints.py`) Each create-capable provider runs all four. The test asserts the returned file id @@ -71,11 +80,59 @@ File delete asserts `object=="file"` and `deleted==True`. | `batch_client.py` | typed file upload/download + batch create/retrieve/cancel/list/delete over the shared ProxyClient; runtime batch model registration via /model/new; denial helpers | | `capabilities.py` | the provider x scenario matrix + per-provider /model/new params + id-shape classifiers + per-provider raw-id assertion | | `conftest.py` | session-scoped batch deployment registration and teardown | -| `test_batches_e2e.py` | parametrized lifecycle with per-endpoint output assertions, file upload/delete outputs, key-model-access denial | +| `test_batches_e2e.py` | parametrized lifecycle with per-endpoint output assertions, file upload/delete outputs, key-model-access denial, per-backend content download, failure paths, second-hop routing, terminal state + cost | + +## Failure paths + +`TestBatchFailurePaths` pins the customer-facing error contracts. A malformed input +file is a 400 at upload naming the bad content. A JSONL line whose url contradicts +the batch endpoint passes create (providers validate asynchronously) and drives the +batch to `failed` with structured `errors.data` (code/line/message), a null +`output_file_id`, and a $0 spend row keyed `{batch_id}_batch_cost` (LIT-4852: a +failed batch books $0 instead of crashing cost tracking). Cancelling that failed +batch is a 409 naming the terminal status. A file id encoded for one deployment wins +over a conflicting `model` param on create: the batch routes and re-encodes by the +file's embedded model (foreign-id precedence). + +## Second hop (two chained gateways) + +`TestBatchSecondHop` registers a `litellm_proxy/` deployment pointing at +the proxy's own base URL with a freshly minted virtual key, so unified upload and +create traverse gateway -> gateway -> OpenAI (LIT-5347, PR #36240). The pin: +`target_model_names` is rewritten to the inner deployment on the second hop and the +nested managed ids round-trip retrieve. This self-chaining only needs the proxy to +reach its own `PROXY_BASE_URL`, which holds both locally and on the e2e stage. + +## Terminal state + cost write-back (cross-run marker baton) + +The 24h completion window rules out submit-and-wait inside one run, so +`TestBatchTerminalState` amortizes across runs. Each run submits a 1-line marker +batch (stable metadata key/value plus a per-run field) and deliberately never +cancels or deletes it or its input file: the marker is the baton the next run picks +up (OpenAI files expire on their own after ~30 days). Polling is list-only, up to 5 +minutes, because retrieving a non-terminal batch books a $0 spend row whose +request_id then blocks the later real-cost row (`skip_duplicates`); the single +retrieve happens only once a completed marker exists. The assertion target is the +newest completed marker from ANY run: run-scoped deployment names mean the list +re-encodes prior-run batches under new encoded ids, so their spend keys are fresh +and a prior-run marker is billable by this run. On the 6h stage cadence the full +assertions are therefore deterministic from run 2 onward. On a cold start (no +completed marker within the poll budget) the test passes on the submission +assertions alone: a documented vacuous pass, not a skip. Markers aged past the 24h +window (25h-73h band, within the newest 100-item list page) must be terminal. + +The cost assertion is the LIT-5730 headline: retrieving a completed model-encoded +batch must write a positive spend row with call_type `aretrieve_batch` and token +usage. Before the fix in `litellm/batches/batch_utils.py`, the retrieve endpoint +re-encoded the response's `output_file_id` in place before the queued logging +worker ran, the worker sent that encoded id to OpenAI, got a 404, and the spend row +never landed. ## Out of scope (intentionally) -Driving a batch to `completed`, cost tracking on completion, and the DB write-back -are not covered here; the 24h window makes them unfit for a synchronous gate. That -logic belongs in a DI-stubbed proxy integration test under `tests/test_litellm/proxy/` -where the provider client is injected to return `completed` deterministically. +Unified (managed) batch cost is owned by the hourly `CheckBatchCost` poller, and a +terminal DB status short-circuits retrieve for those ids, so the terminal-state cell +uses the encoded path; poller timing does not fit an e2e gate and belongs in a +DI-stubbed proxy integration test under `tests/test_litellm/proxy/`. Bedrock +cancel/list stay gated pending LIT-4774. Gemini (non-Vertex) file content raises +`NotImplementedError` upstream and is not a coverage cell. diff --git a/tests/e2e/batches/batch_client.py b/tests/e2e/batches/batch_client.py index 5cc5d1dae3b..21a56f3398f 100644 --- a/tests/e2e/batches/batch_client.py +++ b/tests/e2e/batches/batch_client.py @@ -44,6 +44,17 @@ class FileList(BaseModel): data: list[FileObject] = [] +class BatchErrorItem(BaseModel): + code: str | None = None + line: int | None = None + message: str | None = None + + +class BatchErrorList(BaseModel): + object: str | None = None + data: list[BatchErrorItem] = [] + + class BatchObject(BaseModel): id: str object: str | None = None @@ -51,6 +62,9 @@ class BatchObject(BaseModel): endpoint: str | None = None input_file_id: str | None = None output_file_id: str | None = None + error_file_id: str | None = None + errors: BatchErrorList | None = None + metadata: dict[str, str] | None = None completion_window: str | None = None created_at: int | None = None model: str | None = None @@ -72,12 +86,18 @@ class BatchCreateBody(BaseModel): endpoint: str = "/v1/chat/completions" completion_window: str = "24h" model: str | None = None + metadata: dict[str, str] | None = None class ModelQuery(BaseModel): model: str | None = None +class BatchListQuery(BaseModel): + model: str | None = None + limit: int | None = None + + def is_model_access_denied(resp: StreamingResponse) -> bool: """True if the proxy rejected the call because the key may not access the model.""" return resp.status_code == 403 and "key_model_access_denied" in resp.body @@ -168,12 +188,17 @@ class BatchClient: ) def list_batches( - self, *, key: str, provider: str | None = None + self, + *, + key: str, + provider: str | None = None, + model: str | None = None, + limit: int | None = None, ) -> Result[BatchList]: return self.proxy.transport.get( _batches_path(provider), headers=self.proxy.transport.bearer(key), - params=NoBody(), + params=BatchListQuery(model=model, limit=limit), response_type=BatchList, ) diff --git a/tests/e2e/batches/capabilities.py b/tests/e2e/batches/capabilities.py index 3988fb5e7e1..ce1f68184a7 100644 --- a/tests/e2e/batches/capabilities.py +++ b/tests/e2e/batches/capabilities.py @@ -210,6 +210,16 @@ def is_model_encoded_id(id_str: str) -> bool: return False +def decoded_model_from_id(id_str: str) -> str | None: + """Deployment name embedded in a model-encoded file/batch id, or None.""" + for prefix in ("file-", "batch_"): + if id_str.startswith(prefix): + decoded = _b64_decode(id_str[len(prefix) :]) + if decoded.startswith("litellm:") and ";model," in decoded: + return decoded.split(";model,", 1)[1].split(";")[0] + return None + + def matches_id_shape(shape: IdShape, id_str: str) -> bool: if shape == "managed": return is_managed_id(id_str) diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index 1376bdbed38..75b3a6cd758 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -1,11 +1,12 @@ """Live e2e for the Batches API across every provider LiteLLM supports. -Synchronous tier only: a batch's completion window is 24h, so these never wait for -"completed". Each case uploads a tiny JSONL, creates the batch through one of the -four routing scenarios, asserts it was accepted (non-terminal status) and routed to -the right provider, then retrieves / cancels / lists where the provider supports it. -Everything created is deleted on teardown. Completion + cost tracking are out of -scope here (see COVERAGE.md). +Mostly synchronous tier: a batch's completion window is 24h, so the lifecycle +matrix never waits for "completed". Each case uploads a tiny JSONL, creates the +batch through one of the four routing scenarios, asserts it was accepted +(non-terminal status) and routed to the right provider, then retrieves / cancels / +lists where the provider supports it. Everything created is deleted on teardown. +The exception is TestBatchTerminalState, which carries completed-state + cost +write-back coverage via a cross-run marker baton (design in COVERAGE.md). Routing signal: for provider_fallback the raw batch id discriminates the provider; for the encoded/unified/model_param scenarios the proxy re-encodes the id, so the @@ -22,8 +23,9 @@ from datetime import datetime, timedelta, timezone from typing import Callable import pytest +from pydantic import BaseModel -from e2e_config import unique_marker +from e2e_config import PROXY_BASE_URL, unique_marker from batch_client import ( UPLOAD_FILENAME, @@ -40,9 +42,12 @@ from capabilities import ( CAPABILITIES, FILE_ID_SHAPE, OPENAI_BATCH_MODEL, + PROVIDERS, Capability, + Provider, batch_model_name, coverage_cells_for_lifecycle, + decoded_model_from_id, is_managed_id, matches_id_shape, raw_id_matches_provider, @@ -475,9 +480,22 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row( OPENAI_FILE_CONTENT_BACKEND = "gpt-4o-mini" +FILE_CONTENT_CELLS = { + "azure": "llm.files.azure_openai.content.nonstream.works", + "vertex_ai": "llm.files.vertex.content.nonstream.works", + "bedrock": "llm.files.bedrock.content.nonstream.works", +} +BYTE_FIDELITY_CONTENT_PROVIDERS = frozenset({"azure"}) + class TestBatchFileContent: - """GET /v1/files/{id}/content returns the uploaded batch JSONL bytes.""" + """GET /v1/files/{id}/content returns the uploaded batch JSONL bytes. + + Azure stores the upload verbatim, so its download is asserted byte-equal. + Vertex (GCS) and Bedrock (S3) transform each JSONL line into the provider's + request format at upload time, so their downloads assert 200 plus non-empty + parseable JSON lines instead of byte equality. + """ @pytest.mark.covers( "llm.files.openai.content.nonstream.works", @@ -521,6 +539,62 @@ class TestBatchFileContent: "downloaded file content must match the uploaded JSONL bytes" ) + @pytest.mark.parametrize( + "provider", + [ + pytest.param( + p, + id=p.name, + marks=pytest.mark.covers( + FILE_CONTENT_CELLS[p.name], exercised_on=["files"] + ), + ) + for p in PROVIDERS + if p.name in FILE_CONTENT_CELLS + ], + ) + def test_unified_file_content_downloads( + self, + provider: Provider, + client: BatchClient, + resources: ResourceManager, + batch_deployments: None, + ) -> None: + key = resources.key() + payload = render_jsonl(provider.raw_model) + file = unwrap( + client.upload_file( + content=payload, + form=FileUploadForm(purpose="batch", target_model_names=provider.model), + key=key, + ) + ) + resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + assert_file_object(file, provider=provider.name) + assert is_managed_id(file.id), ( + f"{provider.name}: unified upload must return a managed file id, got {file.id!r}" + ) + + downloaded = client.proxy.transport.download( + f"/v1/files/{file.id}/content", + headers=client.proxy.transport.bearer(key), + ) + assert downloaded.status_code == 200, ( + f"{provider.name}: file content must be 200, " + f"got {downloaded.status_code}: {downloaded.body[:300]}" + ) + body = downloaded.body.strip() + assert body, f"{provider.name}: file content download returned an empty body" + if provider.name in BYTE_FIDELITY_CONTENT_PROVIDERS: + assert body == payload.decode().strip(), ( + f"{provider.name}: downloaded content must match the uploaded JSONL bytes" + ) + else: + for line in body.splitlines(): + assert json.loads(line), ( + f"{provider.name}: content line is not JSON: {line[:200]}" + ) + class TestOpenAIFiles: """GET /v1/files (list) and GET /v1/files/{id} (retrieve) over the OpenAI route. @@ -866,3 +940,384 @@ class TestHostedVllmBatch: f"hosted_vllm batch has non-transitional status {batch.status!r}" ) assert_batch_object(batch) + + +BATCH_TERMINAL_STATUSES = frozenset({"completed", "failed", "expired", "cancelled"}) +FAILED_BATCH_POLL_SECONDS = 120.0 +FAILED_BATCH_POLL_INTERVAL_SECONDS = 5.0 + +AZURE_BATCH_RAW_MODEL = next(p.raw_model for p in PROVIDERS if p.name == "azure") + + +def _mismatched_endpoint_jsonl(model: str) -> bytes: + line = { + "custom_id": "req-1", + "method": "POST", + "url": "/v1/embeddings", + "body": {"model": model, "input": "ping"}, + } + return (json.dumps(line) + "\n").encode() + + +def _poll_until_terminal(client: BatchClient, batch_id: str, key: str) -> BatchObject: + deadline = time.monotonic() + FAILED_BATCH_POLL_SECONDS + fetched = retrieve_batch(client, batch_id, key=key, provider=None) + while fetched.status not in BATCH_TERMINAL_STATUSES and time.monotonic() < deadline: + time.sleep(FAILED_BATCH_POLL_INTERVAL_SECONDS) + fetched = retrieve_batch(client, batch_id, key=key, provider=None) + return fetched + + +class TestBatchFailurePaths: + """Customer-facing failure contracts for /v1/batches. + + A malformed input file is rejected at upload with a 400 naming the bad + content. A JSONL line whose url contradicts the batch endpoint is accepted + at create (providers validate asynchronously) and drives the batch to + "failed" with structured per-line errors, a null output_file_id, and a + zero-cost spend row (LIT-4852: a failed batch must book $0, not crash cost + tracking). Cancelling that already-failed batch returns a 409 naming the + terminal status. A file id encoded for one deployment wins over a + conflicting model param on create: the batch routes (and re-encodes) by the + file's embedded model, pinning that precedence. + """ + + @pytest.mark.covers( + "llm.batches.openai.malformed_jsonl.nonstream.works", + exercised_on=["files"], + ) + def test_malformed_jsonl_upload_rejected( + self, client: BatchClient, resources: ResourceManager, batch_deployments: None + ) -> None: + result = client.upload_file( + content=b"this is not json\n", + form=FileUploadForm(purpose="batch"), + model=OPENAI_BATCH_MODEL, + key=resources.key(), + ) + match result: + case UnknownApiError(status_code=400, body=body): + assert "json" in body.lower(), ( + f"400 must name the malformed JSONL so users can fix the file, got: {body[:300]}" + ) + case _: + pytest.fail(f"malformed JSONL upload must be rejected with a 400, got: {result}") + + @pytest.mark.covers( + "llm.batches.openai.jsonl_endpoint_mismatch.nonstream.works", + "llm.batches.openai.cancel_terminal.nonstream.works", + exercised_on=["batches", "files"], + ) + def test_endpoint_mismatch_fails_batch_and_cancel_conflicts( + self, client: BatchClient, resources: ResourceManager, batch_deployments: None + ) -> None: + key = resources.key() + file = unwrap( + client.upload_file( + content=_mismatched_endpoint_jsonl("gpt-4o-mini"), + form=FileUploadForm(purpose="batch"), + model=OPENAI_BATCH_MODEL, + key=key, + ) + ) + resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + + created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) + require_successful_call(created) + batch = BatchObject.model_validate_json(created.body) + + fetched = _poll_until_terminal(client, batch.id, key) + assert fetched.status == "failed", ( + f"endpoint-mismatched batch must fail, got {fetched.status!r}" + ) + assert fetched.output_file_id is None, ( + f"failed batch must have no output file, got {fetched.output_file_id!r}" + ) + assert fetched.errors is not None and fetched.errors.data, ( + "failed batch must surface structured errors so users can fix the JSONL" + ) + first_error = fetched.errors.data[0] + assert first_error.message, "batch error item has no message" + assert first_error.code, "batch error item has no code" + + rows = client.proxy.poll_logs_for_request_id(f"{fetched.id}_batch_cost") + assert rows, ( + f"failed batch {fetched.id} wrote no spend row; retrieve must book $0 (LIT-4852)" + ) + assert all((row.spend or 0) == 0 for row in rows), ( + f"failed batch must cost $0, got {[(r.request_id, r.spend) for r in rows]}" + ) + assert rows[0].call_type == "aretrieve_batch", ( + f"batch cost row call_type={rows[0].call_type!r}" + ) + + conflict = client.cancel_batch(batch.id, key=key) + match conflict: + case UnknownApiError(status_code=409, body=body): + assert "failed" in body.lower(), ( + f"409 must name the terminal status blocking the cancel, got: {body[:300]}" + ) + case _: + pytest.fail(f"cancel of a failed batch must return a 409 conflict, got: {conflict}") + + @pytest.mark.covers( + "llm.batches.openai.foreign_file_id.nonstream.works", + exercised_on=["batches", "files"], + ) + def test_foreign_encoded_file_id_routes_by_file_model( + self, client: BatchClient, resources: ResourceManager, batch_deployments: None + ) -> None: + key = resources.key() + file = unwrap( + client.upload_file( + content=render_jsonl(AZURE_BATCH_RAW_MODEL), + form=FileUploadForm(purpose="batch"), + model=AZURE_BATCH_MODEL, + key=key, + ) + ) + resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + assert decoded_model_from_id(file.id) == AZURE_BATCH_MODEL, ( + f"upload did not encode the azure deployment into the file id: {file.id!r}" + ) + + created = client.create_batch( + body=BatchCreateBody(input_file_id=file.id, model=OPENAI_BATCH_MODEL), key=key + ) + require_successful_call(created) + batch = BatchObject.model_validate_json(created.body) + resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) + + assert decoded_model_from_id(batch.id) == AZURE_BATCH_MODEL, ( + "create with a foreign encoded file id must route by the file's embedded model, " + f"but the batch id encodes {decoded_model_from_id(batch.id)!r} " + f"(model param was {OPENAI_BATCH_MODEL!r})" + ) + fetched = retrieve_batch(client, batch.id, key=key, provider=None) + assert fetched.id == batch.id + assert fetched.status, "retrieved foreign-file batch has no status" + + +class TestBatchSecondHop: + """Two-proxy batch routing: a litellm_proxy deployment chained to the gateway + itself (LIT-5347, PR #36240). + + The hop deployment's litellm_params point litellm_proxy/ at this + gateway's own base URL with a freshly minted virtual key, so the unified + upload and batch create traverse gateway -> gateway -> OpenAI. The regression + this pins: target_model_names must be rewritten to the inner deployment on + the second hop and the nested managed ids must round-trip retrieve. + """ + + @pytest.mark.covers( + "llm.batches.openai.second_hop.nonstream.works", + exercised_on=["batches", "files"], + ) + def test_unified_create_and_retrieve_via_chained_gateway( + self, client: BatchClient, resources: ResourceManager, batch_deployments: None + ) -> None: + key = resources.key() + hop_name = batch_model_name("openai-batch-hop") + model_id = client.create_model( + hop_name, + LiteLLMParamsBody( + model=f"litellm_proxy/{OPENAI_BATCH_MODEL}", + api_base=PROXY_BASE_URL, + api_key=key, + ), + ) + resources.defer(lambda: client.delete_model(model_id)) + + file = unwrap( + client.upload_file( + content=render_jsonl("gpt-4o-mini"), + form=FileUploadForm(purpose="batch", target_model_names=hop_name), + key=key, + ) + ) + resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + assert is_managed_id(file.id), ( + f"second-hop unified upload must return a managed file id, got {file.id!r}" + ) + + created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) + require_successful_call(created) + batch = BatchObject.model_validate_json(created.body) + resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) + + assert is_managed_id(batch.id), ( + f"second-hop create must return a managed batch id, got {batch.id!r}" + ) + assert batch.status in CREATED_BATCH_STATUSES, ( + f"second-hop batch has non-transitional status {batch.status!r}" + ) + assert_batch_object(batch) + + fetched = retrieve_batch(client, batch.id, key=key, provider=None) + assert fetched.id == batch.id + assert fetched.status, "second-hop retrieve returned no status" + + +class BatchOutputBody(BaseModel): + choices: list[object] = [] + + +class BatchOutputResponse(BaseModel): + status_code: int | None = None + body: BatchOutputBody | None = None + + +class BatchOutputLine(BaseModel): + response: BatchOutputResponse + + +TERMINAL_MARKER_KEY = "litellm_e2e_suite" +TERMINAL_MARKER_VALUE = "batches-terminal-baton" +TERMINAL_POLL_SECONDS = 300.0 +TERMINAL_POLL_INTERVAL_SECONDS = 10.0 +TERMINAL_LIST_LIMIT = 100 +TERMINAL_BAND_MIN_AGE_SECONDS = 25 * 3600 +TERMINAL_BAND_MAX_AGE_SECONDS = 73 * 3600 + + +def _marker_batches(client: BatchClient, key: str) -> list[BatchObject]: + listed = unwrap( + client.list_batches(key=key, model=OPENAI_BATCH_MODEL, limit=TERMINAL_LIST_LIMIT) + ) + return [ + b + for b in listed.data + if (b.metadata or {}).get(TERMINAL_MARKER_KEY) == TERMINAL_MARKER_VALUE + ] + + +def _await_completed_marker( + client: BatchClient, key: str +) -> tuple[BatchObject | None, list[BatchObject]]: + deadline = time.monotonic() + TERMINAL_POLL_SECONDS + while True: + markers = _marker_batches(client, key) + completed = max( + (b for b in markers if b.status == "completed"), + key=lambda b: b.created_at or 0, + default=None, + ) + if completed is not None or time.monotonic() >= deadline: + return completed, markers + time.sleep(TERMINAL_POLL_INTERVAL_SECONDS) + + +def _assert_aged_markers_terminal(markers: list[BatchObject]) -> None: + now = time.time() + stuck = [ + b + for b in markers + if b.created_at is not None + and TERMINAL_BAND_MIN_AGE_SECONDS <= now - b.created_at <= TERMINAL_BAND_MAX_AGE_SECONDS + and b.status not in BATCH_TERMINAL_STATUSES + ] + assert not stuck, ( + "marker batches past their 24h completion window must be terminal; stuck: " + f"{[(b.id, b.status, b.created_at) for b in stuck]}" + ) + + +class TestBatchTerminalState: + """Terminal state + cost write-back via a cross-run marker baton. + + Each run submits a 1-line marker batch (stable metadata key/value plus a + per-run field) and never cancels or deletes it: the marker is the baton the + next run picks up. Polling is list-only for up to 5 minutes because a + retrieve of a non-terminal batch books a $0 spend row whose request_id then + blocks the real-cost row (skip_duplicates); the single retrieve happens only + once a completed marker exists. The assertion target is the newest completed + marker from ANY run, so on the 6h stage cadence the full assertions are + deterministic from run 2 onward. On a cold start (no marker has ever + completed within the poll budget) the test passes on the submission + assertions alone: that is a documented vacuous pass, not a skip, and this + run's marker becomes the next run's target. Markers aged past OpenAI's 24h + completion window (25h-73h band, within the newest list page) must be + terminal. The cost assertion is the LIT-5730 headline: retrieving a + completed model-encoded batch must write a positive spend row keyed + {batch_id}_batch_cost; before the fix the logging worker fetched the + re-encoded output_file_id, 404d, and the row never landed. + """ + + @pytest.mark.covers( + "llm.batches.openai.terminal_state.nonstream.works", + "llm.batches.openai.terminal_state.nonstream.cost_logged", + exercised_on=["batches", "files"], + ) + def test_completed_batch_downloads_output_and_books_cost( + self, client: BatchClient, resources: ResourceManager, batch_deployments: None + ) -> None: + key = resources.key() + file = unwrap( + client.upload_file( + content=render_jsonl("gpt-4o-mini"), + form=FileUploadForm(purpose="batch"), + model=OPENAI_BATCH_MODEL, + key=key, + ) + ) + created = client.create_batch( + body=BatchCreateBody( + input_file_id=file.id, + metadata={ + TERMINAL_MARKER_KEY: TERMINAL_MARKER_VALUE, + "run": unique_marker(), + }, + ), + key=key, + ) + require_successful_call(created) + submitted = BatchObject.model_validate_json(created.body) + assert submitted.status in CREATED_BATCH_STATUSES, ( + f"marker batch has non-transitional status {submitted.status!r}" + ) + assert (submitted.metadata or {}).get(TERMINAL_MARKER_KEY) == TERMINAL_MARKER_VALUE, ( + f"create dropped the marker metadata: {submitted.metadata!r}" + ) + + completed, markers = _await_completed_marker(client, key) + _assert_aged_markers_terminal(markers) + if completed is None: + return + + fetched = retrieve_batch(client, completed.id, key=key, provider=None) + assert fetched.status == "completed", ( + f"listed-completed marker retrieved as {fetched.status!r}" + ) + assert fetched.output_file_id, "completed batch has no output_file_id" + + downloaded = client.proxy.transport.download( + f"/v1/files/{fetched.output_file_id}/content", + headers=client.proxy.transport.bearer(key), + ) + assert downloaded.status_code == 200, ( + f"output content must be 200, got {downloaded.status_code}: {downloaded.body[:300]}" + ) + first_line = BatchOutputLine.model_validate_json(downloaded.body.strip().splitlines()[0]) + assert first_line.response.status_code == 200, ( + f"batch output line reports failure: {downloaded.body[:400]}" + ) + assert first_line.response.body is not None and first_line.response.body.choices, ( + "batch output line has no choices" + ) + + rows = client.proxy.poll_logs_for_request_id( + f"{fetched.id}_batch_cost", + predicate=lambda found: any((row.spend or 0) > 0 for row in found), + ) + priced = [row for row in rows if (row.spend or 0) > 0] + assert priced, ( + f"completed batch {fetched.id} wrote no positive-cost spend row under " + f"request_id {fetched.id}_batch_cost; cost write-back is broken (LIT-5730)" + ) + cost_row = priced[0] + assert cost_row.call_type == "aretrieve_batch", ( + f"batch cost row call_type={cost_row.call_type!r}" + ) + assert (cost_row.total_tokens or 0) > 0, ( + f"batch cost row has no token usage: {cost_row.total_tokens!r}" + ) diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index bb7169509eb..3de462683d2 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -24,6 +24,13 @@ - {id: llm.batches.hosted_vllm.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: hosted_vllm, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "hosted_vllm OpenAI-compatible batch create"} - {id: llm.batches.openai.key_model_access_denied.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Key model restriction 403 on upload/create"} - {id: llm.batches.openai.input_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "vendor strategy §9.18 / LIT-4778", rationale: "Missing input_file_id and invalid batch id rejected"} +- {id: llm.batches.openai.terminal_state.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "A batch actually reaches completed and its output file downloads through GET /v1/files/{id}/content with per-line provider responses"} +- {id: llm.batches.openai.terminal_state.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "test_batches_e2e.py / LIT-5730", fail_before_fix: proven, rationale: "Retrieving a completed model-encoded batch writes a positive spend row keyed {batch_id}_batch_cost (pins LIT-4852/LIT-5666; before the fix the logging worker 404d fetching the re-encoded output_file_id and the row was never written)"} +- {id: llm.batches.openai.malformed_jsonl.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "Uploading a non-JSON batch file is rejected with a 400 naming the bad line"} +- {id: llm.batches.openai.jsonl_endpoint_mismatch.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "JSONL line url that contradicts the batch endpoint drives the batch to failed with structured errors, retrieve stays clean, and the terminal retrieve books a zero-cost spend row (LIT-4852)"} +- {id: llm.batches.openai.cancel_terminal.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "Cancelling an already-terminal batch returns a 409 conflict naming the terminal status"} +- {id: llm.batches.openai.foreign_file_id.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "Create with one deployment's encoded file id and a conflicting model param routes by the file's embedded model; the returned batch id pins that precedence"} +- {id: llm.batches.openai.second_hop.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5347", rationale: "A litellm_proxy deployment chained to the gateway itself preserves target_model_names through nested unified ids; upload, create, and retrieve work over the two-hop chain (PR #36240)"} - {id: llm.files.openai.upload.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "openai_files_endpoints/files_endpoints.py:46", rationale: "File upload returns OpenAIFileObject"} - {id: llm.files.openai.input_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "vendor strategy §9.16 / LIT-4778", rationale: "File upload without purpose rejected"} - {id: llm.files.openai.retrieve.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "files_endpoints.py", rationale: "File retrieve by id"} @@ -36,6 +43,9 @@ - {id: llm.files.hosted_vllm.upload.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: hosted_vllm, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "hosted_vllm OpenAI-compatible file upload"} - {id: llm.rerank.cohere.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: rerank, route: cohere, capability: basic, streaming: nonstream, assertions: [works], source: "test_rerank_e2e.py:29", rationale: "Cohere rerank, top_n + relevance_score"} - {id: llm.files.openai.content.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "GET /v1/files/{id}/content returns uploaded batch JSONL bytes"} +- {id: llm.files.azure_openai.content.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: azure_openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "GET /v1/files/{id}/content on an Azure unified file returns the uploaded JSONL bytes verbatim"} +- {id: llm.files.vertex.content.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "GET /v1/files/{id}/content on a Vertex unified file streams the GCS object back (provider-transformed JSONL, so asserts non-empty JSON lines rather than byte equality)"} +- {id: llm.files.bedrock.content.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "GET /v1/files/{id}/content on a Bedrock unified file streams the S3 object back (provider-transformed JSONL, so asserts non-empty JSON lines rather than byte equality)"} - {id: llm.realtime.bedrock_converse.basic.stream.works, module: llm, tier: P0, subject_endpoint: realtime, route: bedrock_converse, capability: basic, streaming: stream, assertions: [works], source: "test_realtime_bedrock_e2e.py", rationale: "Nova Sonic realtime session emits response.done (LIT-2239)"} - {id: llm.google_native.gemini.basic.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: google_native, route: gemini, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "LIT-4076 / proxy/google_endpoints/endpoints.py", fail_before_fix: proven, rationale: "google-native generateContent must stamp x-litellm-response-cost so SDK traffic reconciles against spend"} - {id: llm.google_native.gemini.basic.stream.works, module: llm, tier: P0, subject_endpoint: google_native, route: gemini, capability: basic, streaming: stream, assertions: [works], source: "PR #28213 / proxy/proxy_server.py async_data_generator", fail_before_fix: proven, rationale: "streamGenerateContent must relay single-prefixed SSE frames with no [DONE] sentinel; doubled data: prefixes and the OpenAI terminator both break the Vertex Java SDK"} diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index ebe093c591c..08cdf945b80 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -800,6 +800,43 @@ async def test_output_file_content_vertex_unified_file_id_extracts_gcs_uri(monke assert captured["custom_llm_provider"] == "vertex_ai" +@pytest.mark.asyncio +async def test_output_file_content_model_encoded_file_id_decoded_to_provider_id(monkeypatch): + import litellm.files.main as files_main + from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model + + captured: dict = {} + + async def fake_afile_content(**kw): + captured.update(kw) + return type("R", (), {"content": b'{"a": 1}'})() + + monkeypatch.setattr(files_main, "afile_content", fake_afile_content) + encoded_id = encode_file_id_with_model("file-Y3FHrMpi7uCkDpY6fgWGeR", "my-batch-model") + + await bu._fetch_batch_output_file_content(_batch(encoded_id), custom_llm_provider="openai") + + assert captured["file_id"] == "file-Y3FHrMpi7uCkDpY6fgWGeR" + assert captured["custom_llm_provider"] == "openai" + + +@pytest.mark.asyncio +async def test_output_file_content_raw_openai_file_id_passes_through(monkeypatch): + import litellm.files.main as files_main + + captured: dict = {} + + async def fake_afile_content(**kw): + captured.update(kw) + return type("R", (), {"content": b'{"a": 1}'})() + + monkeypatch.setattr(files_main, "afile_content", fake_afile_content) + + await bu._fetch_batch_output_file_content(_batch("file-abc123"), custom_llm_provider="openai") + + assert captured["file_id"] == "file-abc123" + + def _vertex_predictions_row(custom_id, prompt_tokens, completion_tokens): return { "request": { From 42cffe93c821ce60a8a6f24d95cd01d995e52f32 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Wed, 19 Aug 2026 16:45:52 -0700 Subject: [PATCH 022/166] Add moonshot/kimi-k3 to model prices and context window map Pricing per https://platform.kimi.ai/docs/pricing/chat-k3: - $3.00/M input (cache miss), $0.30/M cache read, $15.00/M output - 1,048,576 context window; max_completion_tokens settable up to 1,048,576 - Supports reasoning (reasoning_effort low/high/max), tool calling, structured output, vision and video input Co-Authored-By: Claude Fable 5 --- .../model_prices_and_context_window_backup.json | 17 +++++++++++++++++ model_prices_and_context_window.json | 17 +++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d0eca17272d..a7c9825ee7a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -30213,6 +30213,23 @@ "supports_video_input": true, "supports_vision": true }, + "moonshot/kimi-k3": { + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "litellm_provider": "moonshot", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://platform.kimi.ai/docs/pricing/chat-k3", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true + }, "moonshot/kimi-latest": { "cache_read_input_token_cost": 1.5e-07, "deprecation_date": "2026-01-28", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d0eca17272d..a7c9825ee7a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -30213,6 +30213,23 @@ "supports_video_input": true, "supports_vision": true }, + "moonshot/kimi-k3": { + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "litellm_provider": "moonshot", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://platform.kimi.ai/docs/pricing/chat-k3", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true + }, "moonshot/kimi-latest": { "cache_read_input_token_cost": 1.5e-07, "deprecation_date": "2026-01-28", From f8b31f493a62a7b43a2effced84c8a9557929ffd Mon Sep 17 00:00:00 2001 From: 72004 Date: Fri, 21 Aug 2026 02:05:05 +0500 Subject: [PATCH 023/166] fix: don't retire a completed batch from cost recovery while output_file_id is still lagging --- .../openai_files_endpoints/common_utils.py | 21 +++++++++- .../test_files_common_utils.py | 42 +++++++++++++++++++ 2 files changed, 62 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index b9af01e9aea..f8896771077 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -1288,6 +1288,25 @@ def batch_cost_poller_is_active() -> bool: return False +def _completed_batch_safe_to_retire(response) -> bool: + """Whether a "completed" batch may be retired from cost recovery. + + ``batch_processed=True`` is the sole re-pickup gate for CheckBatchCost's + cost-recovery poller, so setting it retires the batch permanently. A batch can + reach ``status="completed"`` while ``output_file_id`` is still ``None`` (the + provider response briefly lags before the output id populates). Retiring in that + window loses the spend record forever. Retire only once we can prove there is + nothing left to recover: the output file has actually arrived, or the provider + reports no successful request lines. When counts are unknown, stay eligible so + the next poller pass revisits it. (#37713) + """ + if getattr(response, "output_file_id", None) is not None: + return True + request_counts = getattr(response, "request_counts", None) + completed = getattr(request_counts, "completed", None) + return completed == 0 + + async def update_batch_in_database( batch_id: str, unified_batch_id: str | Literal[False], @@ -1369,7 +1388,7 @@ async def update_batch_in_database( } poller_owns: Final = batch_cost_poller_is_active() if poller_owns_accounting is None else poller_owns_accounting - if db_status == "complete" and not poller_owns: + if db_status == "complete" and not poller_owns and _completed_batch_safe_to_retire(response): update_data["batch_processed"] = True try: diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py index 6ffb7daaa2d..eb6596e274c 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py @@ -431,3 +431,45 @@ def test_add_internal_model_credentials_survives_a_failing_deployment_lookup(): add_internal_model_credentials(data=data, llm_router=router, model_id="deployment-gone") assert data == {"batch_id": "unified-batch-id"} + + +from litellm.proxy.openai_files_endpoints.common_utils import ( + _completed_batch_safe_to_retire, +) + + +def _completed_batch(output_file_id, completed=None) -> LiteLLMBatch: + kwargs = dict( + id="batch-1", + completion_window="24h", + created_at=1234567890, + endpoint="/v1/chat/completions", + input_file_id="file-in", + object="batch", + status="completed", + output_file_id=output_file_id, + error_file_id=None, + ) + if completed is not None: + kwargs["request_counts"] = {"total": completed, "completed": completed, "failed": 0} + return LiteLLMBatch(**kwargs) + + +class TestCompletedBatchSafeToRetire: + """A completed batch is only safe to retire from cost recovery once its output + file has arrived or the provider proves no successful lines (#37713).""" + + def test_output_file_present_is_safe(self): + assert _completed_batch_safe_to_retire(_completed_batch("file-out")) is True + + def test_no_output_and_no_successful_lines_is_safe(self): + # Every request line errored -> nothing left to recover. + assert _completed_batch_safe_to_retire(_completed_batch(None, completed=0)) is True + + def test_no_output_but_successful_lines_is_not_safe(self): + # The bug: output_file_id is lagging; retiring here loses the spend record. + assert _completed_batch_safe_to_retire(_completed_batch(None, completed=5)) is False + + def test_no_output_and_unknown_counts_is_not_safe(self): + # Counts unknown -> stay eligible so the next poller pass revisits it. + assert _completed_batch_safe_to_retire(_completed_batch(None)) is False From 67d16a499dc208f69ddab20e17648b503acc07ec Mon Sep 17 00:00:00 2001 From: 72004 Date: Fri, 21 Aug 2026 02:13:11 +0500 Subject: [PATCH 024/166] Type the batch-retire helpers and rename test helper to avoid shadowing existing _completed_batch --- litellm/proxy/openai_files_endpoints/common_utils.py | 2 +- .../openai_files_endpoint/test_files_common_utils.py | 12 +++++++----- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index f8896771077..2e8ae6af7a9 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -1288,7 +1288,7 @@ def batch_cost_poller_is_active() -> bool: return False -def _completed_batch_safe_to_retire(response) -> bool: +def _completed_batch_safe_to_retire(response: "LiteLLMBatch") -> bool: """Whether a "completed" batch may be retired from cost recovery. ``batch_processed=True`` is the sole re-pickup gate for CheckBatchCost's diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py index eb6596e274c..3de9e61463f 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py @@ -438,7 +438,9 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( ) -def _completed_batch(output_file_id, completed=None) -> LiteLLMBatch: +def _completed_batch_for_retire( + output_file_id: str | None, completed: int | None = None +) -> LiteLLMBatch: kwargs = dict( id="batch-1", completion_window="24h", @@ -460,16 +462,16 @@ class TestCompletedBatchSafeToRetire: file has arrived or the provider proves no successful lines (#37713).""" def test_output_file_present_is_safe(self): - assert _completed_batch_safe_to_retire(_completed_batch("file-out")) is True + assert _completed_batch_safe_to_retire(_completed_batch_for_retire("file-out")) is True def test_no_output_and_no_successful_lines_is_safe(self): # Every request line errored -> nothing left to recover. - assert _completed_batch_safe_to_retire(_completed_batch(None, completed=0)) is True + assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None, completed=0)) is True def test_no_output_but_successful_lines_is_not_safe(self): # The bug: output_file_id is lagging; retiring here loses the spend record. - assert _completed_batch_safe_to_retire(_completed_batch(None, completed=5)) is False + assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None, completed=5)) is False def test_no_output_and_unknown_counts_is_not_safe(self): # Counts unknown -> stay eligible so the next poller pass revisits it. - assert _completed_batch_safe_to_retire(_completed_batch(None)) is False + assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None)) is False From a3b676278869f171863b1fb96e742249ff76f841 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 15:39:08 -0700 Subject: [PATCH 025/166] fix(streaming): price partial-stream spend rows at the real model and keep prompt and cache fields A streaming chat completion that ends early (client disconnect, or the proxy cutting the stream at LITELLM_MAX_STREAMING_DURATION_SECONDS) wrote a spend log row with spend 0.0, prompt_tokens 0 on the proxy-cut path, and no cache fields in usage_object. The proxy restamps chunk.model in place to the client-facing alias, so the partial response rebuilt from those chunks priced the unmapped alias and came out at 0. The failure path also rebuilt usage without the request messages, so prompt tokens counted to 0, and a cut stream never sees the final usage event that normally zero-fills the cache fields. Restamp the rebuilt partial response with the wrapper's real model before cost calculation on both the disconnect and the failure paths, pass the request messages when rebuilding usage on the failure path, and zero-fill missing cache usage fields the way completed streams already do. --- .../litellm_core_utils/streaming_handler.py | 19 +++- litellm/proxy/common_request_processing.py | 9 ++ .../test_streaming_handler.py | 104 ++++++++++++++++++ .../proxy/test_common_request_processing.py | 85 ++++++++++++++ 4 files changed, 216 insertions(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 485091bccd0..3565872f09e 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2315,10 +2315,18 @@ class CustomStreamWrapper: if self.logging_obj is None or not self.chunks: return try: - partial_response: Final = litellm.stream_chunk_builder(chunks=self.chunks) + partial_response: Final = litellm.stream_chunk_builder( + chunks=self.chunks, + messages=self.messages if isinstance(self.messages, list) else None, + ) + if partial_response is None: + return usage: Final = cast(Usage | None, getattr(partial_response, "usage", None)) if usage is None: return + if self.model: + partial_response.model = self.model + zero_fill_missing_cache_usage_fields(usage) self.logging_obj.model_call_details["combined_usage_object"] = usage self.logging_obj.model_call_details["response_cost"] = ( self.logging_obj._response_cost_calculator(result=partial_response) or 0.0 @@ -2439,6 +2447,15 @@ class CustomStreamWrapper: return chunk +def zero_fill_missing_cache_usage_fields(usage: Usage) -> None: + if getattr(usage, "cache_creation_input_tokens", None) is None: + usage.cache_creation_input_tokens = 0 # rebind-ok: in-place zero-fill is the contract + if getattr(usage, "cache_read_input_tokens", None) is None: + usage.cache_read_input_tokens = 0 # rebind-ok: in-place zero-fill is the contract + if usage.prompt_tokens_details is None: + usage.prompt_tokens_details = PromptTokensDetailsWrapper(cached_tokens=0) # rebind-ok: zero-fill in place + + _TokenDetails = TypeVar("_TokenDetails", PromptTokensDetailsWrapper, CompletionTokensDetailsWrapper) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index a0b69ecb0bf..d430d776350 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -43,6 +43,9 @@ from litellm.litellm_core_utils.llm_response_utils.get_headers import ( get_response_headers, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.litellm_core_utils.streaming_handler import ( + zero_fill_missing_cache_usage_fields, +) from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.auth.auth_checks import can_key_call_resolved_model from litellm.proxy.auth.auth_utils import check_response_size_is_safe @@ -324,6 +327,12 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons return False if partial_response is None: return False + wrapper_model: Final = getattr(response, "model", None) + if isinstance(wrapper_model, str) and wrapper_model: + partial_response.model = wrapper_model + partial_usage: Final = getattr(partial_response, "usage", None) + if isinstance(partial_usage, Usage): + zero_fill_missing_cache_usage_fields(partial_usage) try: await logging_obj.dispatch_success_handlers( partial_response, diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 05b44fffbc5..a550160cd73 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -3388,6 +3388,110 @@ def test_record_partial_usage_for_failure_noop_without_chunks(): assert "combined_usage_object" not in logging_obj.model_call_details +def _wrapper_with_partial_chunks( + chunk_model: str, + usage: Optional[Usage] = None, + model: str = "gpt-4o-mini", + custom_llm_provider: str = "openai", +) -> tuple: + logging_obj = Logging( + model=model, + messages=[{"role": "user", "content": "Tell me a long story"}], + stream=True, + call_type="completion", + start_time=time.time(), + litellm_call_id="partial-usage-alias", + function_id="1245", + ) + logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider + logging_obj.optional_params = {} + wrapper = CustomStreamWrapper( + completion_stream=None, + model=model, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + ) + wrapper.chunks = [ + ModelResponseStream( + id="chatcmpl-partial-alias-1", + created=1742056047, + model=chunk_model, + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="The Roman Empire began when", role="assistant" + ), + ) + ], + usage=usage, + ) + ] + return wrapper, logging_obj + + +def test_record_partial_usage_for_failure_prices_alias_restamped_chunks_at_real_model(): + wrapper, logging_obj = _wrapper_with_partial_chunks( + chunk_model="bedrock-claude-opus-5", + usage=Usage(prompt_tokens=40, completion_tokens=5, total_tokens=45), + model="us.anthropic.claude-opus-5", + custom_llm_provider="bedrock", + ) + assert "bedrock/bedrock-claude-opus-5" not in litellm.model_cost + + wrapper._record_partial_usage_for_failure() + + stashed = logging_obj.model_call_details["combined_usage_object"] + assert stashed.completion_tokens == 5 + rates = litellm.model_cost["us.anthropic.claude-opus-5"] + expected = 40 * rates["input_cost_per_token"] + 5 * rates["output_cost_per_token"] + assert logging_obj.model_call_details["response_cost"] == pytest.approx(expected) + + +def test_record_partial_usage_for_failure_counts_prompt_tokens_from_request_messages(): + wrapper, logging_obj = _wrapper_with_partial_chunks(chunk_model="my-public-alias") + + wrapper._record_partial_usage_for_failure() + + stashed = logging_obj.model_call_details["combined_usage_object"] + assert stashed.prompt_tokens > 0 + + +def test_record_partial_usage_for_failure_zero_fills_missing_cache_fields(): + wrapper, logging_obj = _wrapper_with_partial_chunks(chunk_model="gpt-4o-mini") + + wrapper._record_partial_usage_for_failure() + + stashed = logging_obj.model_call_details["combined_usage_object"] + assert stashed.cache_creation_input_tokens == 0 + assert stashed.cache_read_input_tokens == 0 + assert stashed.prompt_tokens_details is not None + assert stashed.prompt_tokens_details.cached_tokens == 0 + + +def test_record_partial_usage_for_failure_keeps_cache_values_recovered_from_chunks(): + recovered = Usage( + prompt_tokens=40, + completion_tokens=5, + total_tokens=45, + cache_read_input_tokens=7, + cache_creation_input_tokens=3, + ) + wrapper, logging_obj = _wrapper_with_partial_chunks( + chunk_model="gpt-4o-mini", usage=recovered + ) + + wrapper._record_partial_usage_for_failure() + + stashed = logging_obj.model_call_details["combined_usage_object"] + assert stashed.cache_read_input_tokens == 7 + assert stashed.cache_creation_input_tokens == 3 + assert stashed.prompt_tokens_details is not None + assert stashed.prompt_tokens_details.cached_tokens == 7 + + @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio async def test_stream_chunk_builder_raise_at_end_of_stream_still_recovers_usage( diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 510fb977a61..609d0dfe1b6 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -5519,6 +5519,91 @@ class TestStreamingClientDisconnectBilling: proxy_logging_obj._arelease_max_parallel_requests_on_disconnect.assert_awaited_once() + async def _bill_and_collect_success_event(self, prepare=None): + recorder = _RecordingSuccessLogger() + original_callbacks = litellm.callbacks + litellm.callbacks = [recorder] + try: + response = await self._start_partial_stream() + if prepare is not None: + prepare(response) + billed = await _bill_partial_streamed_spend_on_disconnect( + {"litellm_logging_obj": response.logging_obj}, response + ) + assert billed is True + for _ in range(50): + if recorder.success_events: + break + await asyncio.sleep(0.1) + finally: + litellm.callbacks = original_callbacks + assert len(recorder.success_events) == 1 + return recorder.success_events[0] + + @pytest.mark.asyncio + async def test_disconnect_billing_prices_alias_restamped_chunks_at_real_model(self): + assert "openai/my-public-alias" not in litellm.model_cost + + def restamp_chunks_to_alias(response): + for chunk in response.chunks: + chunk.model = "my-public-alias" + + event = await self._bill_and_collect_success_event(restamp_chunks_to_alias) + + assert event["response_obj"].model == "gpt-4o-mini" + standard_logging_object = event["kwargs"]["standard_logging_object"] + assert standard_logging_object["response_cost"] > 0.0 + + @pytest.mark.asyncio + async def test_disconnect_billing_zero_fills_missing_cache_fields(self): + event = await self._bill_and_collect_success_event() + + usage = event["response_obj"].usage + assert getattr(usage, "cache_creation_input_tokens", None) == 0 + assert getattr(usage, "cache_read_input_tokens", None) == 0 + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.cached_tokens == 0 + + @pytest.mark.asyncio + async def test_disconnect_billing_keeps_cache_values_recovered_from_chunks(self): + from litellm.types.utils import ( + Delta, + ModelResponseStream, + StreamingChoices, + Usage, + ) + + def append_usage_chunk(response): + response.chunks.append( + ModelResponseStream( + id=response.chunks[0].id, + model="gpt-4o-mini", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content=" and more", role="assistant"), + ) + ], + usage=Usage( + prompt_tokens=40, + completion_tokens=5, + total_tokens=45, + cache_read_input_tokens=7, + cache_creation_input_tokens=3, + ), + ) + ) + + event = await self._bill_and_collect_success_event(append_usage_chunk) + + usage = event["response_obj"].usage + assert getattr(usage, "cache_read_input_tokens", None) == 7 + assert getattr(usage, "cache_creation_input_tokens", None) == 3 + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.cached_tokens == 7 + def _apply_stream_usage_tracking( data: dict, From de7dcbbc677b3d52461c74f0595b27a3a38be996 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 17:41:03 -0700 Subject: [PATCH 026/166] Carry real cache counts up instead of zeroing them on partial rows cache_read_input_tokens and cache_creation_input_tokens are pydantic extras on Usage, not declared fields, so filling them in created keys that were not there before rather than replacing a None. Readers that test for presence then took the new zero as authoritative: the spend log writer skipped its own copy from prompt_tokens_details, turning a real cache read of 500 into 0, and the prometheus provider cache counters stopped incrementing. Carry the prompt_tokens_details counts up before defaulting to zero, so a partial row reports the same cache numbers a complete one does. Renamed the helper to say what it now does. --- .../litellm_core_utils/streaming_handler.py | 32 +++++++++++--- litellm/proxy/common_request_processing.py | 4 +- .../test_streaming_handler.py | 20 ++++++++- .../proxy/test_common_request_processing.py | 44 ++++++++++++++++++- 4 files changed, 90 insertions(+), 10 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 3565872f09e..651b169a9b2 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2326,7 +2326,7 @@ class CustomStreamWrapper: return if self.model: partial_response.model = self.model - zero_fill_missing_cache_usage_fields(usage) + backfill_missing_cache_usage_fields(usage) self.logging_obj.model_call_details["combined_usage_object"] = usage self.logging_obj.model_call_details["response_cost"] = ( self.logging_obj._response_cost_calculator(result=partial_response) or 0.0 @@ -2447,13 +2447,33 @@ class CustomStreamWrapper: return chunk -def zero_fill_missing_cache_usage_fields(usage: Usage) -> None: - if getattr(usage, "cache_creation_input_tokens", None) is None: - usage.cache_creation_input_tokens = 0 # rebind-ok: in-place zero-fill is the contract +def _cache_token_count(details: PromptTokensDetailsWrapper | None, keys: tuple[str, ...]) -> int: + for key in keys: + value = getattr(details, key, None) + if isinstance(value, int) and not isinstance(value, bool) and value: + return value + return 0 + + +def backfill_missing_cache_usage_fields(usage: Usage) -> None: + """Give partial-stream usage the same cache fields a complete stream reports. + + Carries OpenAI-style ``prompt_tokens_details`` counts up to the Anthropic-style + top-level keys, defaulting to zero. It must carry the real count rather than a + flat zero: downstream readers treat these keys as authoritative once present and + skip their own normalization, so a zero here would overwrite a real cache read. + """ + details: Final = usage.prompt_tokens_details if getattr(usage, "cache_read_input_tokens", None) is None: - usage.cache_read_input_tokens = 0 # rebind-ok: in-place zero-fill is the contract + usage.cache_read_input_tokens = _cache_token_count( # rebind-ok: in-place backfill is the contract + details, ("cached_tokens",) + ) + if getattr(usage, "cache_creation_input_tokens", None) is None: + usage.cache_creation_input_tokens = _cache_token_count( # rebind-ok: in-place backfill is the contract + details, ("cache_write_tokens", "cache_creation_tokens") + ) if usage.prompt_tokens_details is None: - usage.prompt_tokens_details = PromptTokensDetailsWrapper(cached_tokens=0) # rebind-ok: zero-fill in place + usage.prompt_tokens_details = PromptTokensDetailsWrapper(cached_tokens=0) # rebind-ok: backfill in place _TokenDetails = TypeVar("_TokenDetails", PromptTokensDetailsWrapper, CompletionTokensDetailsWrapper) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index d430d776350..72cc298d37d 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -44,7 +44,7 @@ from litellm.litellm_core_utils.llm_response_utils.get_headers import ( ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.streaming_handler import ( - zero_fill_missing_cache_usage_fields, + backfill_missing_cache_usage_fields, ) from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.auth.auth_checks import can_key_call_resolved_model @@ -332,7 +332,7 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons partial_response.model = wrapper_model partial_usage: Final = getattr(partial_response, "usage", None) if isinstance(partial_usage, Usage): - zero_fill_missing_cache_usage_fields(partial_usage) + backfill_missing_cache_usage_fields(partial_usage) try: await logging_obj.dispatch_success_handlers( partial_response, diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index a550160cd73..9f04d63b6ad 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -3459,7 +3459,7 @@ def test_record_partial_usage_for_failure_counts_prompt_tokens_from_request_mess assert stashed.prompt_tokens > 0 -def test_record_partial_usage_for_failure_zero_fills_missing_cache_fields(): +def test_record_partial_usage_for_failure_backfills_missing_cache_fields(): wrapper, logging_obj = _wrapper_with_partial_chunks(chunk_model="gpt-4o-mini") wrapper._record_partial_usage_for_failure() @@ -3471,6 +3471,24 @@ def test_record_partial_usage_for_failure_zero_fills_missing_cache_fields(): assert stashed.prompt_tokens_details.cached_tokens == 0 +def test_record_partial_usage_for_failure_carries_up_openai_style_cached_tokens(): + recovered = Usage( + prompt_tokens=1000, + completion_tokens=10, + total_tokens=1010, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=500), + ) + wrapper, logging_obj = _wrapper_with_partial_chunks( + chunk_model="gpt-4o-mini", usage=recovered + ) + + wrapper._record_partial_usage_for_failure() + + stashed = logging_obj.model_call_details["combined_usage_object"] + assert stashed.cache_read_input_tokens == 500 + assert stashed.cache_creation_input_tokens == 0 + + def test_record_partial_usage_for_failure_keeps_cache_values_recovered_from_chunks(): recovered = Usage( prompt_tokens=40, diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 609d0dfe1b6..4d78bf164b1 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -5555,7 +5555,7 @@ class TestStreamingClientDisconnectBilling: assert standard_logging_object["response_cost"] > 0.0 @pytest.mark.asyncio - async def test_disconnect_billing_zero_fills_missing_cache_fields(self): + async def test_disconnect_billing_backfills_missing_cache_fields(self): event = await self._bill_and_collect_success_event() usage = event["response_obj"].usage @@ -5564,6 +5564,48 @@ class TestStreamingClientDisconnectBilling: assert usage.prompt_tokens_details is not None assert usage.prompt_tokens_details.cached_tokens == 0 + @pytest.mark.asyncio + async def test_disconnect_billing_carries_up_openai_style_cached_tokens(self): + from litellm.types.utils import ( + Delta, + ModelResponseStream, + PromptTokensDetailsWrapper, + StreamingChoices, + Usage, + ) + + def append_openai_style_cached_usage_chunk(response): + response.chunks.append( + ModelResponseStream( + id=response.chunks[0].id, + model="gpt-4o-mini", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content=" and more", role="assistant"), + ) + ], + usage=Usage( + prompt_tokens=1000, + completion_tokens=10, + total_tokens=1010, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=500 + ), + ), + ) + ) + + event = await self._bill_and_collect_success_event( + append_openai_style_cached_usage_chunk + ) + + usage = event["response_obj"].usage + assert getattr(usage, "cache_read_input_tokens", None) == 500 + assert getattr(usage, "cache_creation_input_tokens", None) == 0 + @pytest.mark.asyncio async def test_disconnect_billing_keeps_cache_values_recovered_from_chunks(self): from litellm.types.utils import ( From 27c3f87aef8b6ddee9a5f894983988f7f8bffb6d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 18:00:27 -0700 Subject: [PATCH 027/166] fix(fal_ai): price gpt-image-2 per size and quality from request params --- .../litellm_core_utils/llm_cost_calc/utils.py | 1 + litellm/llms/fal_ai/cost_calculator.py | 66 ++- ...odel_prices_and_context_window_backup.json | 548 +++++++++++++++++- model_prices_and_context_window.json | 548 +++++++++++++++++- .../test_fal_ai_gpt_image_2_transformation.py | 16 +- .../llms/fal_ai/test_cost_calculator.py | 128 ++++ 6 files changed, 1284 insertions(+), 23 deletions(-) create mode 100644 tests/test_litellm/llms/fal_ai/test_cost_calculator.py diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 0793fe20b21..0a52e1d283e 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -1371,6 +1371,7 @@ class CostCalculatorUtils: return fal_ai_image_cost_calculator( model=model, image_response=completion_response, + optional_params=optional_params, ) elif custom_llm_provider == litellm.LlmProviders.RUNWAYML.value: from litellm.llms.runwayml.cost_calculator import ( diff --git a/litellm/llms/fal_ai/cost_calculator.py b/litellm/llms/fal_ai/cost_calculator.py index 8c5ad5a8c64..b2320de247e 100644 --- a/litellm/llms/fal_ai/cost_calculator.py +++ b/litellm/llms/fal_ai/cost_calculator.py @@ -1,25 +1,73 @@ -from typing import Any, Final +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final import litellm from litellm.types.utils import ImageResponse +FAL_KEYED_PRICING_DEFAULT_QUALITY: Final[str] = "high" +FAL_TEXT_TO_IMAGE_DEFAULT_SIZE: Final[str] = "1024-x-768" +FAL_NAMED_IMAGE_SIZES: Final[Mapping[str, str]] = MappingProxyType( + { + "square_hd": "1024-x-1024", + "square": "512-x-512", + "portrait_4_3": "768-x-1024", + "portrait_16_9": "576-x-1024", + "landscape_4_3": "1024-x-768", + "landscape_16_9": "1024-x-576", + } +) + + +def _keyed_size(model: str, optional_params: Mapping[str, object]) -> str | None: + image_size: Final = optional_params.get("image_size") + if image_size is None: + return None if model.endswith("/edit") else FAL_TEXT_TO_IMAGE_DEFAULT_SIZE + if isinstance(image_size, Mapping): + width: Final = image_size.get("width") + height: Final = image_size.get("height") + if isinstance(width, int) and isinstance(height, int): + return f"{width}-x-{height}" + return None + if isinstance(image_size, str): + return FAL_NAMED_IMAGE_SIZES.get(image_size) + return None + + +def _keyed_cost_per_image(model: str, optional_params: Mapping[str, object] | None) -> float | None: + if optional_params is None: + return None + size: Final = _keyed_size(model=model, optional_params=optional_params) + if size is None: + return None + raw_quality: Final = optional_params.get("quality") + quality: Final = ( + raw_quality if isinstance(raw_quality, str) and raw_quality != "auto" else FAL_KEYED_PRICING_DEFAULT_QUALITY + ) + keyed_entry: Final = litellm.model_cost.get(f"fal_ai/{quality}/{size}/{model}") + if keyed_entry is None: + return None + keyed_cost: Final = keyed_entry.get("output_cost_per_image") + return float(keyed_cost) if isinstance(keyed_cost, (int, float)) else None + def cost_calculator( model: str, - image_response: Any, + image_response: object, + optional_params: Mapping[str, object] | None = None, ) -> float: """ fal.ai image generation cost calculator """ + if not isinstance(image_response, ImageResponse): + raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}") + num_images: Final[int] = len(image_response.data) if image_response.data else 0 + keyed_cost_per_image: Final = _keyed_cost_per_image(model=model, optional_params=optional_params) + if keyed_cost_per_image is not None: + return keyed_cost_per_image * num_images _model_info: Final = litellm.get_model_info( model=model, custom_llm_provider=litellm.LlmProviders.FAL_AI.value, ) output_cost_per_image: Final[float] = _model_info.get("output_cost_per_image") or 0.0 - num_images: int = 0 - if isinstance(image_response, ImageResponse): - if image_response.data: - num_images = len(image_response.data) - return output_cost_per_image * num_images - else: - raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}") + return output_cost_per_image * num_images diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 5a1c988c21a..d700e75f6aa 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -17400,7 +17400,7 @@ "fal_ai/openai/gpt-image-2": { "litellm_provider": "fal_ai", "metadata": { - "notes": "OpenAI gpt-image-2 served through fal.ai. fal bills by token, so the flat output_cost_per_image here is fal's published per-image rate for a default request (quality=high, image_size=landscape_4_3 at 1024x768). Other canonical sizes at high quality: 1024x1024 $0.211, 1024x1536 $0.165, 1920x1080 $0.158, 2560x1440 $0.222, 3840x2160 $0.401" + "notes": "OpenAI gpt-image-2 served through fal.ai. fal bills by token but publishes deterministic per-image prices per size and quality, mirrored here as keyed entries fal_ai/{quality}/{width}-x-{height}/openai/gpt-image-2 that litellm's fal_ai cost calculator picks from the request params. This flat entry is the fallback when no keyed entry matches and carries the default request rate (quality=high, image_size=landscape_4_3 at 1024x768). quality=auto is priced as high" }, "mode": "image_generation", "output_cost_per_image": 0.145, @@ -17410,10 +17410,190 @@ ], "supports_vision": true }, + "fal_ai/low/1024-x-768/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.005, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1024-x-1024/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.006, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1024-x-1536/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.005, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1920-x-1080/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.005, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/2560-x-1440/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.007, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/3840-x-2160/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.012, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-768/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.037, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-1024/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.053, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-1536/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.042, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1920-x-1080/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.04, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/2560-x-1440/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.056, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/3840-x-2160/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.101, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-768/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.145, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-1024/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.211, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-1536/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.165, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1920-x-1080/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.158, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/2560-x-1440/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.222, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/3840-x-2160/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.401, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, "fal_ai/gpt-image-2": { "litellm_provider": "fal_ai", "metadata": { - "notes": "Alias of fal_ai/openai/gpt-image-2, which litellm also accepts without the openai/ prefix. Same rate, see that entry for the size and quality caveat" + "notes": "Alias of fal_ai/openai/gpt-image-2, which litellm also accepts without the openai/ prefix. Same rates, including the keyed fal_ai/{quality}/{width}-x-{height}/gpt-image-2 entries; see that entry for details" }, "mode": "image_generation", "output_cost_per_image": 0.145, @@ -17423,13 +17603,373 @@ ], "supports_vision": true }, + "fal_ai/low/1024-x-768/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.005, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1024-x-1024/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.006, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1024-x-1536/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.005, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1920-x-1080/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.005, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/2560-x-1440/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.007, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/3840-x-2160/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.012, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-768/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.037, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-1024/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.053, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-1536/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.042, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1920-x-1080/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.04, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/2560-x-1440/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.056, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/3840-x-2160/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.101, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-768/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.145, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-1024/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.211, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-1536/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.165, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1920-x-1080/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.158, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/2560-x-1440/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.222, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/3840-x-2160/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.401, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, "fal_ai/openai/gpt-image-2/edit": { "litellm_provider": "fal_ai", "metadata": { - "notes": "Editing endpoint of gpt-image-2 on fal.ai, reached through the image generation path with fal's image_urls param since /v1/images/edits is not wired for fal_ai. Same rate as fal_ai/openai/gpt-image-2, see that entry for the size and quality caveat" + "notes": "Editing endpoint of gpt-image-2 on fal.ai, reached through the image generation path with fal's image_urls param since /v1/images/edits is not wired for fal_ai. Prices include one input image and live in keyed entries fal_ai/{quality}/{width}-x-{height}/openai/gpt-image-2/edit. This flat entry is the fallback for the default edit request (quality=high, image_size=auto, inferred from the input image, priced as 1024x768 high)" }, "mode": "image_generation", - "output_cost_per_image": 0.145, + "output_cost_per_image": 0.151, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1024-x-768/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.011, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1024-x-1024/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.015, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1024-x-1536/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.018, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1920-x-1080/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.017, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/2560-x-1440/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.019, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/3840-x-2160/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.024, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-768/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.043, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-1024/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.061, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-1536/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.054, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1920-x-1080/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.053, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/2560-x-1440/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.068, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/3840-x-2160/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.113, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-768/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.151, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-1024/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.219, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-1536/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.178, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1920-x-1080/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.158, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/2560-x-1440/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.234, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/3840-x-2160/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.413, "source": "https://fal.ai/models/openai/gpt-image-2/edit", "supported_endpoints": [ "/v1/images/generations" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 5a1c988c21a..d700e75f6aa 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -17400,7 +17400,7 @@ "fal_ai/openai/gpt-image-2": { "litellm_provider": "fal_ai", "metadata": { - "notes": "OpenAI gpt-image-2 served through fal.ai. fal bills by token, so the flat output_cost_per_image here is fal's published per-image rate for a default request (quality=high, image_size=landscape_4_3 at 1024x768). Other canonical sizes at high quality: 1024x1024 $0.211, 1024x1536 $0.165, 1920x1080 $0.158, 2560x1440 $0.222, 3840x2160 $0.401" + "notes": "OpenAI gpt-image-2 served through fal.ai. fal bills by token but publishes deterministic per-image prices per size and quality, mirrored here as keyed entries fal_ai/{quality}/{width}-x-{height}/openai/gpt-image-2 that litellm's fal_ai cost calculator picks from the request params. This flat entry is the fallback when no keyed entry matches and carries the default request rate (quality=high, image_size=landscape_4_3 at 1024x768). quality=auto is priced as high" }, "mode": "image_generation", "output_cost_per_image": 0.145, @@ -17410,10 +17410,190 @@ ], "supports_vision": true }, + "fal_ai/low/1024-x-768/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.005, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1024-x-1024/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.006, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1024-x-1536/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.005, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1920-x-1080/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.005, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/2560-x-1440/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.007, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/3840-x-2160/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.012, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-768/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.037, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-1024/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.053, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-1536/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.042, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1920-x-1080/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.04, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/2560-x-1440/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.056, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/3840-x-2160/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.101, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-768/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.145, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-1024/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.211, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-1536/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.165, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1920-x-1080/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.158, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/2560-x-1440/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.222, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/3840-x-2160/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.401, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, "fal_ai/gpt-image-2": { "litellm_provider": "fal_ai", "metadata": { - "notes": "Alias of fal_ai/openai/gpt-image-2, which litellm also accepts without the openai/ prefix. Same rate, see that entry for the size and quality caveat" + "notes": "Alias of fal_ai/openai/gpt-image-2, which litellm also accepts without the openai/ prefix. Same rates, including the keyed fal_ai/{quality}/{width}-x-{height}/gpt-image-2 entries; see that entry for details" }, "mode": "image_generation", "output_cost_per_image": 0.145, @@ -17423,13 +17603,373 @@ ], "supports_vision": true }, + "fal_ai/low/1024-x-768/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.005, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1024-x-1024/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.006, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1024-x-1536/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.005, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1920-x-1080/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.005, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/2560-x-1440/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.007, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/3840-x-2160/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.012, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-768/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.037, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-1024/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.053, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-1536/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.042, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1920-x-1080/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.04, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/2560-x-1440/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.056, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/3840-x-2160/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.101, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-768/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.145, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-1024/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.211, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-1536/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.165, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1920-x-1080/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.158, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/2560-x-1440/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.222, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/3840-x-2160/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.401, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, "fal_ai/openai/gpt-image-2/edit": { "litellm_provider": "fal_ai", "metadata": { - "notes": "Editing endpoint of gpt-image-2 on fal.ai, reached through the image generation path with fal's image_urls param since /v1/images/edits is not wired for fal_ai. Same rate as fal_ai/openai/gpt-image-2, see that entry for the size and quality caveat" + "notes": "Editing endpoint of gpt-image-2 on fal.ai, reached through the image generation path with fal's image_urls param since /v1/images/edits is not wired for fal_ai. Prices include one input image and live in keyed entries fal_ai/{quality}/{width}-x-{height}/openai/gpt-image-2/edit. This flat entry is the fallback for the default edit request (quality=high, image_size=auto, inferred from the input image, priced as 1024x768 high)" }, "mode": "image_generation", - "output_cost_per_image": 0.145, + "output_cost_per_image": 0.151, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1024-x-768/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.011, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1024-x-1024/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.015, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1024-x-1536/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.018, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1920-x-1080/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.017, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/2560-x-1440/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.019, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/3840-x-2160/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.024, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-768/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.043, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-1024/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.061, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-1536/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.054, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1920-x-1080/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.053, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/2560-x-1440/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.068, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/3840-x-2160/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.113, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-768/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.151, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-1024/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.219, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-1536/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.178, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1920-x-1080/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.158, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/2560-x-1440/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.234, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/3840-x-2160/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.413, "source": "https://fal.ai/models/openai/gpt-image-2/edit", "supported_endpoints": [ "/v1/images/generations" diff --git a/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_gpt_image_2_transformation.py b/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_gpt_image_2_transformation.py index 3c8cf9f9e0a..1a527230f1b 100644 --- a/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_gpt_image_2_transformation.py +++ b/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_gpt_image_2_transformation.py @@ -128,19 +128,23 @@ def test_transform_image_generation_request(): @pytest.mark.parametrize( - "model", + ("model", "expected_cost_for_two_images"), [ - "openai/gpt-image-2", - "gpt-image-2", - "openai/gpt-image-2/edit", + ("openai/gpt-image-2", 0.29), + ("gpt-image-2", 0.29), + ("openai/gpt-image-2/edit", 0.302), ], ) -def test_cost_calculator_uses_registry_price(model, monkeypatch: pytest.MonkeyPatch): +def test_cost_calculator_uses_registry_price( + model, expected_cost_for_two_images, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() response = ImageResponse( data=[ ImageObject(url="https://v3b.fal.media/files/b/one.png"), ImageObject(url="https://v3b.fal.media/files/b/two.png"), ] ) - assert cost_calculator(model=model, image_response=response) == pytest.approx(0.29) + assert cost_calculator(model=model, image_response=response) == pytest.approx(expected_cost_for_two_images) diff --git a/tests/test_litellm/llms/fal_ai/test_cost_calculator.py b/tests/test_litellm/llms/fal_ai/test_cost_calculator.py new file mode 100644 index 00000000000..689b620a90b --- /dev/null +++ b/tests/test_litellm/llms/fal_ai/test_cost_calculator.py @@ -0,0 +1,128 @@ +import pytest + +import litellm +from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils +from litellm.llms.fal_ai.cost_calculator import cost_calculator +from litellm.types.utils import ImageObject, ImageResponse + + +@pytest.fixture(autouse=True) +def _use_local_model_cost_map(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() + + +def _image_response(num_images: int = 1) -> ImageResponse: + return ImageResponse(data=[ImageObject(url="https://example.com/img.png") for _ in range(num_images)]) + + +def test_high_quality_1024x1024_uses_keyed_price(): + cost = cost_calculator( + model="openai/gpt-image-2", + image_response=_image_response(), + optional_params={"quality": "high", "image_size": {"width": 1024, "height": 1024}}, + ) + assert cost == pytest.approx(0.211) + + +def test_alias_model_uses_keyed_price(): + cost = cost_calculator( + model="gpt-image-2", + image_response=_image_response(), + optional_params={"quality": "high", "image_size": {"width": 1024, "height": 1024}}, + ) + assert cost == pytest.approx(0.211) + + +def test_default_request_priced_at_default_size_and_quality(): + cost = cost_calculator( + model="openai/gpt-image-2", + image_response=_image_response(), + optional_params={}, + ) + assert cost == pytest.approx(0.145) + + +def test_auto_quality_priced_as_high(): + cost = cost_calculator( + model="openai/gpt-image-2", + image_response=_image_response(), + optional_params={"quality": "auto", "image_size": {"width": 1024, "height": 1024}}, + ) + assert cost == pytest.approx(0.211) + + +def test_low_quality_4k_uses_keyed_price(): + cost = cost_calculator( + model="openai/gpt-image-2", + image_response=_image_response(), + optional_params={"quality": "low", "image_size": {"width": 3840, "height": 2160}}, + ) + assert cost == pytest.approx(0.012) + + +def test_named_fal_size_uses_keyed_price(): + cost = cost_calculator( + model="openai/gpt-image-2", + image_response=_image_response(), + optional_params={"quality": "high", "image_size": "square_hd"}, + ) + assert cost == pytest.approx(0.211) + + +def test_edit_model_uses_keyed_edit_price(): + cost = cost_calculator( + model="openai/gpt-image-2/edit", + image_response=_image_response(), + optional_params={"quality": "high", "image_size": {"width": 1024, "height": 1024}}, + ) + assert cost == pytest.approx(0.219) + + +def test_edit_model_without_size_falls_back_to_flat_price(): + cost = cost_calculator( + model="openai/gpt-image-2/edit", + image_response=_image_response(), + optional_params={"quality": "high"}, + ) + assert cost == pytest.approx(0.151) + + +def test_missing_optional_params_falls_back_to_flat_price(): + cost = cost_calculator( + model="openai/gpt-image-2", + image_response=_image_response(), + optional_params=None, + ) + assert cost == pytest.approx(0.145) + + +def test_unlisted_size_falls_back_to_flat_price(): + cost = cost_calculator( + model="openai/gpt-image-2", + image_response=_image_response(), + optional_params={"quality": "high", "image_size": {"width": 999, "height": 999}}, + ) + assert cost == pytest.approx(0.145) + + +def test_keyed_price_multiplies_per_image(): + cost = cost_calculator( + model="openai/gpt-image-2", + image_response=_image_response(num_images=2), + optional_params={"quality": "high", "image_size": {"width": 1024, "height": 1024}}, + ) + assert cost == pytest.approx(0.422) + + +def test_route_image_generation_passes_optional_params_to_fal(): + cost = CostCalculatorUtils.route_image_generation_cost_calculator( + model="openai/gpt-image-2", + completion_response=_image_response(), + custom_llm_provider="fal_ai", + optional_params={"quality": "high", "image_size": {"width": 1024, "height": 1024}}, + ) + assert cost == pytest.approx(0.211) From 2471e85f544e7767a95f6fa1319866373a15812c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 18:28:59 -0700 Subject: [PATCH 028/166] chore(lint): note why the streamed cost fallback swallows pricing errors --- litellm/proxy/common_request_processing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index d533df8616d..7975db5cb85 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3562,7 +3562,7 @@ class ProxyBaseLLMRequestProcessing: ) -> float | None: try: cost: Final = litellm_logging_obj._response_cost_calculator(result=model_response) # pyright: ignore[reportPrivateUsage] # reuse the call's own cost calc for pricing parity with the logging callback - except Exception: + except Exception: # noqa: BLE001 # a pricing failure falls back to model-name pricing instead of breaking the stream return None return float(cost) if isinstance(cost, (int, float)) and not isinstance(cost, bool) else None From d5e6a0c9b8b262e097033f56e2aa4a8d5275bd22 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Thu, 20 Aug 2026 18:36:48 -0700 Subject: [PATCH 029/166] fix(fal_ai): strip provider prefix before keyed cost lookup --- litellm/llms/fal_ai/cost_calculator.py | 2 ++ .../llms/fal_ai/test_cost_calculator.py | 28 +++++++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/litellm/llms/fal_ai/cost_calculator.py b/litellm/llms/fal_ai/cost_calculator.py index b2320de247e..74848784c5b 100644 --- a/litellm/llms/fal_ai/cost_calculator.py +++ b/litellm/llms/fal_ai/cost_calculator.py @@ -61,6 +61,8 @@ def cost_calculator( """ if not isinstance(image_response, ImageResponse): raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}") + # the proxy cost path passes the provider-prefixed model name + model = model.removeprefix(f"{litellm.LlmProviders.FAL_AI.value}/") num_images: Final[int] = len(image_response.data) if image_response.data else 0 keyed_cost_per_image: Final = _keyed_cost_per_image(model=model, optional_params=optional_params) if keyed_cost_per_image is not None: diff --git a/tests/test_litellm/llms/fal_ai/test_cost_calculator.py b/tests/test_litellm/llms/fal_ai/test_cost_calculator.py index 689b620a90b..f167aceaa95 100644 --- a/tests/test_litellm/llms/fal_ai/test_cost_calculator.py +++ b/tests/test_litellm/llms/fal_ai/test_cost_calculator.py @@ -37,6 +37,24 @@ def test_alias_model_uses_keyed_price(): assert cost == pytest.approx(0.211) +def test_provider_prefixed_model_uses_keyed_price(): + cost = cost_calculator( + model="fal_ai/openai/gpt-image-2", + image_response=_image_response(), + optional_params={"quality": "high", "image_size": {"width": 1024, "height": 1024}}, + ) + assert cost == pytest.approx(0.211) + + +def test_provider_prefixed_edit_model_uses_keyed_edit_price(): + cost = cost_calculator( + model="fal_ai/openai/gpt-image-2/edit", + image_response=_image_response(), + optional_params={"quality": "high", "image_size": {"width": 1024, "height": 1024}}, + ) + assert cost == pytest.approx(0.219) + + def test_default_request_priced_at_default_size_and_quality(): cost = cost_calculator( model="openai/gpt-image-2", @@ -126,3 +144,13 @@ def test_route_image_generation_passes_optional_params_to_fal(): optional_params={"quality": "high", "image_size": {"width": 1024, "height": 1024}}, ) assert cost == pytest.approx(0.211) + + +def test_route_image_generation_with_provider_prefixed_model_uses_keyed_price(): + cost = CostCalculatorUtils.route_image_generation_cost_calculator( + model="fal_ai/openai/gpt-image-2", + completion_response=_image_response(), + custom_llm_provider="fal_ai", + optional_params={"quality": "high", "image_size": {"width": 1024, "height": 1024}}, + ) + assert cost == pytest.approx(0.211) From 03a253a1f9878d9a0d3a990d683ebb40acad0892 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 18:45:51 -0700 Subject: [PATCH 030/166] Keep the model Azure Model Router recovered from later chunks The disconnect billing path was stamping the wrapper's model over whatever stream_chunk_builder assembled. For Azure Model Router that throws away the routed model: the proxy deliberately leaves those chunks unrestamped so the builder can pick the real model off a later chunk, and overwriting it prices the row at the router alias instead. Only apply the wrapper's model when the builder did not find a model beyond the first chunk's, which is every case except Model Router. --- litellm/proxy/common_request_processing.py | 26 ++++++++++++++++++- .../proxy/test_common_request_processing.py | 13 ++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 72cc298d37d..7333111b6b4 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -277,6 +277,26 @@ def _deferred_stream_logging_is_armed(request_data: dict) -> bool: ) +def _assembled_model_came_from_a_later_chunk(chunks: list, assembled_model: object) -> bool: + """Report whether stream_chunk_builder picked a model the first chunk did not carry. + + Azure Model Router puts the routed model on the chunks after the first one, and the + proxy deliberately leaves those chunks unrestamped so the builder can recover it. The + assembled model is then more specific than the wrapper's, so the caller has to leave + it alone rather than stamping the wrapper's model over it. + """ + first_chunk: Final = chunks[0] + first_chunk_model: Final = ( + first_chunk.get("model") if isinstance(first_chunk, dict) else getattr(first_chunk, "model", None) + ) + return ( + isinstance(first_chunk_model, str) + and isinstance(assembled_model, str) + and bool(assembled_model) + and assembled_model != first_chunk_model + ) + + async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, response: object) -> bool: """ A client disconnect throws GeneratorExit/CancelledError into the streaming @@ -328,7 +348,11 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons if partial_response is None: return False wrapper_model: Final = getattr(response, "model", None) - if isinstance(wrapper_model, str) and wrapper_model: + if ( + isinstance(wrapper_model, str) + and wrapper_model + and not _assembled_model_came_from_a_later_chunk(chunks, partial_response.model) + ): partial_response.model = wrapper_model partial_usage: Final = getattr(partial_response, "usage", None) if isinstance(partial_usage, Usage): diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 4d78bf164b1..fd48d62a651 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -5554,6 +5554,19 @@ class TestStreamingClientDisconnectBilling: standard_logging_object = event["kwargs"]["standard_logging_object"] assert standard_logging_object["response_cost"] > 0.0 + @pytest.mark.asyncio + async def test_disconnect_billing_keeps_the_model_azure_model_router_picked(self): + def restamp_like_azure_model_router(response): + response.chunks[0].model = "azure-model-router" + for chunk in response.chunks[1:]: + chunk.model = "gpt-4.1-nano-2025-04-14" + + event = await self._bill_and_collect_success_event(restamp_like_azure_model_router) + + assert event["response_obj"].model == "gpt-4.1-nano-2025-04-14" + standard_logging_object = event["kwargs"]["standard_logging_object"] + assert standard_logging_object["response_cost"] > 0.0 + @pytest.mark.asyncio async def test_disconnect_billing_backfills_missing_cache_fields(self): event = await self._bill_and_collect_success_event() From 101ef7e16717495bccdb756f56b15f71e6ed2d8b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 18:46:38 -0700 Subject: [PATCH 031/166] test(cost): cover the chat.completion.chunk and raising-pricer branches The logging-object pricing applies to streamed /v1/chat/completions too, not just Anthropic message_delta, so a deployment with negotiated per-token prices now gets that price in the streamed usage.cost there as well. Nothing asserted that half. Adds the discounted and the sticker-fallback case for the OpenAI chunk shape, plus the branch where the pricer raises and the frame falls back to model-name pricing instead of breaking the stream. --- .../proxy/test_common_request_processing.py | 73 +++++++++++++++++++ 1 file changed, 73 insertions(+) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index b314a25ed09..58ea736ba12 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -6197,6 +6197,79 @@ class TestInjectCostIntoUsageDict: + 8 * pricing["output_cost_per_token"] ) + def test_message_delta_falls_back_to_model_pricing_when_the_logging_obj_raises(self): + """A pricing failure mid-stream must not break the frame, so the raise falls back to + model-name pricing rather than propagating into the response body.""" + + class _StubLoggingObj: + def _response_cost_calculator(self, result): + raise ValueError("no pricing for this deployment") + + model = "claude-haiku-4-5" + pricing = litellm.model_cost[model] + event = { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": {"input_tokens": 14, "output_tokens": 8, "cache_read_input_tokens": 3202}, + } + + result = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(event, model, _StubLoggingObj()) + + assert result is not None + assert result["usage"]["cost"] == pytest.approx( + 14 * pricing["input_cost_per_token"] + + 3202 * pricing["cache_read_input_token_cost"] + + 8 * pricing["output_cost_per_token"] + ) + + def test_openai_chunk_prices_through_the_logging_obj_so_custom_pricing_applies(self): + """The chat.completion.chunk path rides the same pricer, so a discounted deployment + streaming /v1/chat/completions gets its negotiated price instead of sticker.""" + + class _StubLoggingObj: + def __init__(self, cost): + self._cost = cost + self.captured_result = None + + def _response_cost_calculator(self, result): + self.captured_result = result + return self._cost + + discounted_cost = 0.00031 + stub = _StubLoggingObj(discounted_cost) + event = { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "choices": [], + "usage": {"prompt_tokens": 1000, "completion_tokens": 100, "total_tokens": 1100}, + } + + result = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(event, "gpt-4o-mini", stub) + + assert result is not None + assert result["usage"]["cost"] == discounted_cost + assert result["usage"]["cost"] != pytest.approx(self._expected_cost("gpt-4o-mini", 1000, 100)) + usage = stub.captured_result.usage + assert usage.prompt_tokens == 1000 + assert usage.completion_tokens == 100 + + def test_openai_chunk_falls_back_to_model_pricing_when_the_logging_obj_returns_no_cost(self): + class _StubLoggingObj: + def _response_cost_calculator(self, result): + return None + + event = { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "choices": [], + "usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}, + } + + result = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(event, "gpt-4o-mini", _StubLoggingObj()) + + assert result is not None + assert result["usage"]["cost"] == pytest.approx(self._expected_cost("gpt-4o-mini", 11, 4)) + class TestProcessChunkWithCostInjection: def test_complete_usage_frame_chunk_is_injected(self, monkeypatch): From 655d10775ca88dca9fc0bed63a2042ab3da5fe54 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 19:01:16 -0700 Subject: [PATCH 032/166] fix(cost): keep mid-stream pricing from leaking a breakdown into the spend log Pricing a frame through the request's own logging object is what makes custom deployment pricing work, but _response_cost_calculator does not only return a number. It also stamps cost_breakdown onto the live logging object, and on a pricing failure it writes response_cost_failure_debug_information into model_call_details. On an ordinary proxy stream that is harmless, because the success handler recomputes cost_breakdown at end of stream and overwrites whatever the frames left behind. The pass-through handlers are the problem: they compute their final cost with a bare completion_cost call and never touch cost_breakdown again, so a breakdown derived from one mid-stream frame would survive to the end and land in the spend log's metadata. response_cost itself is unaffected either way, so this was a reporting surface bug rather than a billing one, but the spend row would have gone from null to a populated breakdown for a partial frame. Snapshot both writes and put them back once the cost is read, so pricing a frame stays a read as far as the rest of the request is concerned. The returned cost is unchanged, so nothing about the injected usage.cost moves. --- litellm/proxy/common_request_processing.py | 17 ++++++ .../proxy/test_common_request_processing.py | 60 +++++++++++++++++++ 2 files changed, 77 insertions(+) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 7975db5cb85..d0dad1a4204 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3560,10 +3560,27 @@ class ProxyBaseLLMRequestProcessing: def _logging_obj_cost_or_none( model_response: ModelResponse, litellm_logging_obj: LiteLLMLoggingObj ) -> float | None: + # Pricing a frame stamps cost_breakdown and, on failure, the cost-failure debug key onto + # the live logging object. The pass-through handlers never recompute either one, so a + # frame-derived breakdown would outlive the stream and land in the spend log. Snapshot + # both and put them back, so pricing here stays a read as far as the request is concerned + breakdown_before: Final = getattr(litellm_logging_obj, "cost_breakdown", None) + call_details: Final = getattr(litellm_logging_obj, "model_call_details", None) + debug_key: Final = "response_cost_failure_debug_information" + debug_missing: Final = object() + debug_before: Final = call_details.get(debug_key, debug_missing) if isinstance(call_details, dict) else None try: cost: Final = litellm_logging_obj._response_cost_calculator(result=model_response) # pyright: ignore[reportPrivateUsage] # reuse the call's own cost calc for pricing parity with the logging callback except Exception: # noqa: BLE001 # a pricing failure falls back to model-name pricing instead of breaking the stream return None + finally: + if hasattr(litellm_logging_obj, "cost_breakdown"): + litellm_logging_obj.cost_breakdown = breakdown_before + if isinstance(call_details, dict): + if debug_before is debug_missing: + call_details.pop(debug_key, None) + else: + call_details[debug_key] = debug_before return float(cost) if isinstance(cost, (int, float)) and not isinstance(cost, bool) else None @staticmethod diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 58ea736ba12..6b762d7bd74 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -6222,6 +6222,66 @@ class TestInjectCostIntoUsageDict: + 8 * pricing["output_cost_per_token"] ) + def test_pricing_a_frame_leaves_the_real_logging_obj_unchanged(self): + """Pricing runs against the live logging object, and the pass-through handlers never + recompute cost_breakdown, so a frame-derived breakdown would reach the spend log.""" + from litellm.litellm_core_utils.litellm_logging import ( + Logging as LiteLLMLoggingObj, + ) + from litellm.types.utils import ModelResponse, Usage + + logging_obj = LiteLLMLoggingObj( + model="claude-haiku-4-5", + messages=[{"role": "user", "content": "test"}], + stream=True, + call_type="completion", + start_time=None, + litellm_call_id="lit4902-breakdown-test", + function_id="lit4902-breakdown-test", + ) + logging_obj.update_environment_variables(litellm_params={}, optional_params={}) + logging_obj.model_call_details["custom_llm_provider"] = "anthropic" + assert logging_obj.cost_breakdown is None + + model_response = ModelResponse( + usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224) + ) + cost = ProxyBaseLLMRequestProcessing._logging_obj_cost_or_none(model_response, logging_obj) + + assert cost is not None and cost > 0 + assert logging_obj.cost_breakdown is None + assert "response_cost_failure_debug_information" not in logging_obj.model_call_details + + def test_pricing_a_frame_restores_a_breakdown_the_request_already_had(self): + from litellm.litellm_core_utils.litellm_logging import ( + Logging as LiteLLMLoggingObj, + ) + from litellm.types.utils import ModelResponse, Usage + + logging_obj = LiteLLMLoggingObj( + model="claude-haiku-4-5", + messages=[{"role": "user", "content": "test"}], + stream=True, + call_type="completion", + start_time=None, + litellm_call_id="lit4902-breakdown-restore", + function_id="lit4902-breakdown-restore", + ) + logging_obj.update_environment_variables(litellm_params={}, optional_params={}) + logging_obj.model_call_details["custom_llm_provider"] = "anthropic" + logging_obj.set_cost_breakdown( + input_cost=0.5, output_cost=0.25, total_cost=0.75, cost_for_built_in_tools_cost_usd_dollar=0.0 + ) + existing = logging_obj.cost_breakdown + + model_response = ModelResponse( + usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224) + ) + ProxyBaseLLMRequestProcessing._logging_obj_cost_or_none(model_response, logging_obj) + + assert logging_obj.cost_breakdown is existing + assert logging_obj.cost_breakdown["total_cost"] == 0.75 + def test_openai_chunk_prices_through_the_logging_obj_so_custom_pricing_applies(self): """The chat.completion.chunk path rides the same pricer, so a discounted deployment streaming /v1/chat/completions gets its negotiated price instead of sticker.""" From c010bd6a7ccdc658a734e49f3a38e979c9c4275f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 19:08:11 -0700 Subject: [PATCH 033/166] Only keep the builder's model when the client did not ask for it A chunk carrying usage is stored as a pre-restamp copy, so an alias-restamped stream reaches disconnect billing with its first chunk still on the deployment model and every later chunk on the client's name. That is the same shape Azure Model Router produces, and the previous guard read it as a routed model and left the alias on the row, which is the unpriced name this PR set out to stop. Compare the assembled model against the name the proxy stamps chunks with, so the alias goes back to the deployment's model and the routed model stays. --- litellm/proxy/common_request_processing.py | 29 ++++++++++++----- .../proxy/test_common_request_processing.py | 31 +++++++++++++++++-- 2 files changed, 49 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 7333111b6b4..1c6237f1c56 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -281,9 +281,11 @@ def _assembled_model_came_from_a_later_chunk(chunks: list, assembled_model: obje """Report whether stream_chunk_builder picked a model the first chunk did not carry. Azure Model Router puts the routed model on the chunks after the first one, and the - proxy deliberately leaves those chunks unrestamped so the builder can recover it. The - assembled model is then more specific than the wrapper's, so the caller has to leave - it alone rather than stamping the wrapper's model over it. + proxy deliberately leaves those chunks unrestamped so the builder can recover it. + + A stored chunk that carries usage is a pre-restamp copy of the one the proxy saw, so + an alias-restamped stream reaches the builder with the same shape: a first chunk that + disagrees with the rest. Those two are only told apart by what the client asked for. """ first_chunk: Final = chunks[0] first_chunk_model: Final = ( @@ -297,6 +299,18 @@ def _assembled_model_came_from_a_later_chunk(chunks: list, assembled_model: obje ) +def _assembled_model_is_the_name_the_client_asked_for(request_data: dict, assembled_model: object) -> bool: + """Report whether the assembled model is the public name the proxy stamps onto chunks. + + That stamp is what leaves an unpriced alias on the partial response, so the deployment's + own model has to go back on before the row is costed. + """ + return assembled_model in ( + request_data.get("_litellm_client_requested_model"), + request_data.get("model"), + ) + + async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, response: object) -> bool: """ A client disconnect throws GeneratorExit/CancelledError into the streaming @@ -348,11 +362,10 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons if partial_response is None: return False wrapper_model: Final = getattr(response, "model", None) - if ( - isinstance(wrapper_model, str) - and wrapper_model - and not _assembled_model_came_from_a_later_chunk(chunks, partial_response.model) - ): + builder_recovered_the_routed_model: Final = _assembled_model_came_from_a_later_chunk( + chunks, partial_response.model + ) and not _assembled_model_is_the_name_the_client_asked_for(request_data, partial_response.model) + if isinstance(wrapper_model, str) and wrapper_model and not builder_recovered_the_routed_model: partial_response.model = wrapper_model partial_usage: Final = getattr(partial_response, "usage", None) if isinstance(partial_usage, Usage): diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index fd48d62a651..9e1fe5f8dac 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -5519,7 +5519,7 @@ class TestStreamingClientDisconnectBilling: proxy_logging_obj._arelease_max_parallel_requests_on_disconnect.assert_awaited_once() - async def _bill_and_collect_success_event(self, prepare=None): + async def _bill_and_collect_success_event(self, prepare=None, request_data=None): recorder = _RecordingSuccessLogger() original_callbacks = litellm.callbacks litellm.callbacks = [recorder] @@ -5528,7 +5528,7 @@ class TestStreamingClientDisconnectBilling: if prepare is not None: prepare(response) billed = await _bill_partial_streamed_spend_on_disconnect( - {"litellm_logging_obj": response.logging_obj}, response + {"litellm_logging_obj": response.logging_obj, **(request_data or {})}, response ) assert billed is True for _ in range(50): @@ -5554,6 +5554,28 @@ class TestStreamingClientDisconnectBilling: standard_logging_object = event["kwargs"]["standard_logging_object"] assert standard_logging_object["response_cost"] > 0.0 + @pytest.mark.asyncio + async def test_disconnect_billing_prices_a_partly_restamped_chunk_list_at_real_model(self): + """ + A chunk that carries usage is stored as a copy before the proxy restamps the + one it forwards, so an aliased stream can reach billing with its first chunk + still on the deployment model and the rest on the client's name. + """ + assert "openai/my-public-alias" not in litellm.model_cost + + def restamp_only_the_chunks_the_proxy_forwarded(response): + for chunk in response.chunks[1:]: + chunk.model = "my-public-alias" + + event = await self._bill_and_collect_success_event( + restamp_only_the_chunks_the_proxy_forwarded, + request_data={"model": "my-public-alias"}, + ) + + assert event["response_obj"].model == "gpt-4o-mini" + standard_logging_object = event["kwargs"]["standard_logging_object"] + assert standard_logging_object["response_cost"] > 0.0 + @pytest.mark.asyncio async def test_disconnect_billing_keeps_the_model_azure_model_router_picked(self): def restamp_like_azure_model_router(response): @@ -5561,7 +5583,10 @@ class TestStreamingClientDisconnectBilling: for chunk in response.chunks[1:]: chunk.model = "gpt-4.1-nano-2025-04-14" - event = await self._bill_and_collect_success_event(restamp_like_azure_model_router) + event = await self._bill_and_collect_success_event( + restamp_like_azure_model_router, + request_data={"model": "azure-model-router"}, + ) assert event["response_obj"].model == "gpt-4.1-nano-2025-04-14" standard_logging_object = event["kwargs"]["standard_logging_object"] From 08014933477131f40357808bc12f93c90f7e4ad3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 19:25:46 -0700 Subject: [PATCH 034/166] Match the client-name check to the name the proxy actually stamps Pre-call processing rewrites request_data["model"] for aliasing and routing, so matching either key let a routed model count as the client's own name and put the wrapper model back on an Azure Model Router row. --- litellm/proxy/common_request_processing.py | 11 +++++--- .../proxy/test_common_request_processing.py | 25 +++++++++++++++++++ 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 1c6237f1c56..e7983ba3c9b 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -303,12 +303,15 @@ def _assembled_model_is_the_name_the_client_asked_for(request_data: dict, assemb """Report whether the assembled model is the public name the proxy stamps onto chunks. That stamp is what leaves an unpriced alias on the partial response, so the deployment's - own model has to go back on before the row is costed. + own model has to go back on before the row is costed. Pre-call processing rewrites + `request_data["model"]` for aliasing and routing, so the client's own name wins when it + is there, in the same order the proxy picks the name it stamps. """ - return assembled_model in ( - request_data.get("_litellm_client_requested_model"), - request_data.get("model"), + client_requested_model: Final = request_data.get("_litellm_client_requested_model") + stamped_model: Final = ( + client_requested_model if isinstance(client_requested_model, str) else request_data.get("model") ) + return isinstance(stamped_model, str) and assembled_model == stamped_model async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, response: object) -> bool: diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 9e1fe5f8dac..9dd76b4eb88 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -5592,6 +5592,31 @@ class TestStreamingClientDisconnectBilling: standard_logging_object = event["kwargs"]["standard_logging_object"] assert standard_logging_object["response_cost"] > 0.0 + @pytest.mark.asyncio + async def test_disconnect_billing_keeps_the_routed_model_when_request_data_model_was_rewritten(self): + """ + Pre-call processing rewrites request_data["model"] for aliasing and routing, so the + routed model on the later chunks can end up matching it. Only the name the client + sent says whether the proxy restamped this stream. + """ + + def restamp_like_azure_model_router(response): + response.chunks[0].model = "azure-model-router" + for chunk in response.chunks[1:]: + chunk.model = "gpt-4.1-nano-2025-04-14" + + event = await self._bill_and_collect_success_event( + restamp_like_azure_model_router, + request_data={ + "model": "gpt-4.1-nano-2025-04-14", + "_litellm_client_requested_model": "azure-model-router", + }, + ) + + assert event["response_obj"].model == "gpt-4.1-nano-2025-04-14" + standard_logging_object = event["kwargs"]["standard_logging_object"] + assert standard_logging_object["response_cost"] > 0.0 + @pytest.mark.asyncio async def test_disconnect_billing_backfills_missing_cache_fields(self): event = await self._bill_and_collect_success_event() From f1c4145f86b7501fe6f693f4106c8a68af702c98 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 20 Aug 2026 19:32:57 -0700 Subject: [PATCH 035/166] fix(scim): resolve group members by SSO identity or email before creating a placeholder (#37686) SCIM group members were matched against litellm user ids only. An identity provider that lists people by email or by the OIDC subject therefore matched nothing, and the member fell through to placeholder creation. Since #37688 made a failed member creation fail the group sync rather than drop the member, that fallthrough is no longer quiet: the placeholder is created with user_email set to the member value, the duplicate-email check rejects it, and the whole group push answers 500. So on current staging a group listing anyone by their email fails outright, every other member in the payload included. An unmatched member id is now looked up across sso_user_id and user_email in one query. Searching either field first would hide a value that names one account by its SSO identity and another by its email, and hand the group to whichever was searched first. The two are not compared alike: an email is matched the way new_user matches one before accepting a new account, case-insensitively, because matching more strictly than the layer that would reject the placeholder is what turned an id whose casing differed from the stored email into that same 500. An SSO identity is matched exactly, since OIDC defines sub as case-sensitive and nothing folds its case on the way in. An exact user_id hit is checked the same way rather than trusted outright, since a value can be one account's id and another's SSO identity or email. That is not a corner case: the placeholders this bug provisioned are keyed by the very id the provider keeps pushing, so on a tenant that already has them the placeholder wins the id lookup and the real account can never be matched. Refusing names the problem instead of silently landing on the placeholder again. Those rows still have to be deleted before the real account resolves; making the sync heal itself needs a trustworthy way to tell a placeholder from an account someone created, and created_via lives in caller-writable metadata, so it is left to a follow-up. A value that names more than one account is refused with a 400 naming the id rather than attributed to one of them. Removals resolve too, since the roster holds canonical user ids and a directory removes people by the id it added them with. A removal counts the members one value names: the id as written when the roster holds it verbatim, which is how an earlier release recorded a member it could not match, together with the members it resolves to. Counting only the accounts on the roster keeps someone removable after a second account takes their email, which resolving table-wide would not, and counting both ways of naming a member together stops one value revoking two people when it is one member's canonical id and another's email. A value naming two of the group's own members is undecidable and fails rather than guessing or reporting a removal it did not perform. Resolves LIT-5383 Co-authored-by: Yassin Kortam --- .../management_endpoints/scim/scim_v2.py | 202 ++++- .../scim/test_scim_v2_endpoints.py | 706 +++++++++++++++++- 2 files changed, 888 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index 8d255859571..7183e6cb402 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -487,13 +487,18 @@ class _UnknownMember(NamedTuple): value: str -_ClassifiedGroupMember = Union[_ResolvedUserMember, _SkippedGroupMember, _UnknownMember] +class _AmbiguousMember(NamedTuple): + value: str + + +_ClassifiedGroupMember = Union[_ResolvedUserMember, _SkippedGroupMember, _UnknownMember, _AmbiguousMember] class _PartitionedMembers(NamedTuple): resolved_ids: tuple[str, ...] skipped: tuple[_SkippedGroupMember, ...] unknown_ids: tuple[str, ...] + ambiguous_values: tuple[str, ...] def _member_value(member: SCIMMember) -> str: @@ -536,6 +541,44 @@ def _team_metadata_has_scim_provenance(team_metadata: object) -> bool: return bool(fields.get(SCIM_MANAGED_TEAM_METADATA_KEY)) or fields.get(SCIM_TEAM_DATA_METADATA_KEY) is not None +class _CaseInsensitiveMatch(TypedDict): + equals: ReadOnly[str] + mode: ReadOnly[str] + + +async def _users_named_by_member_value( + value: str, prisma_client: PrismaClient, *, take: int | None = 2 +) -> tuple[str, ...]: + """Every user id this member value names, by SSO identity or by email. + + Both fields are searched in one pass, because searching either first would hide a + value that names one account by its SSO identity and another by its email, and + hand the group to whichever field was searched first. + + They are not compared alike. An email is matched the way ``new_user`` matches one + before it accepts a new account, case-insensitively: matching more strictly than + the layer that would reject the placeholder is what turned a member id whose + casing differed from the stored email into a 500 on the whole push. An SSO + identity is matched exactly, because OIDC defines ``sub`` as case-sensitive and + nothing folds its case on the way in, so treating two subjects that differ in case + as one would hand the group to an account the provider never named. + + ``take`` bounds the read for a caller that only needs to know whether the value + names one account or several; ``user_email`` carries no index, so letting the scan + stop early is worth the two rows. A caller that has to know *which* accounts, as a + removal does, passes None. That set is the accounts sharing one identity, which is + a handful at worst. + """ + subject: Final = value.strip() + email: Final[_CaseInsensitiveMatch] = {"equals": subject, "mode": "insensitive"} + rows: Final = await _table(UserRepository(prisma_client)).find_many( + # mutable-ok: the Prisma serializer requires concrete dicts and a concrete list + where={"OR": [{"sso_user_id": subject}, {"user_email": email}]}, + take=take, + ) + return tuple(dict.fromkeys(row.user_id for row in rows)) + + async def _classify_group_member(member: SCIMMember, prisma_client: PrismaClient) -> _ClassifiedGroupMember: """ Decide what a single SCIM group member refers to. @@ -557,6 +600,20 @@ async def _classify_group_member(member: SCIMMember, prisma_client: PrismaClient one the identity provider writes. An id the IdP called a User is a user even if some team happens to share the id, and a team created here rather than through SCIM is not evidence of anything about the member. + + When those checks miss on an otherwise user-shaped member, its value is looked + up as an SSO identity or an email, and a match resolves to that user's + ``user_id``. A value that names more than one account is ambiguous rather than + unknown: it names a real person we cannot identify, so it is neither guessed at + nor provisioned. + + An exact ``user_id`` hit is checked the same way rather than trusted outright. A + value can be one account's id and another's SSO identity or email, and taking the + id on sight would hand the group to whichever account happened to be keyed by it. + The placeholders this bug provisioned are that shape exactly, since they are keyed + by the very id the provider keeps pushing, so on a tenant that already has them + the membership is refused and named rather than silently landing on the + placeholder again. """ value: Final = _member_value(member) member_type: Final = _normalized_member_type(member) @@ -566,6 +623,18 @@ async def _classify_group_member(member: SCIMMember, prisma_client: PrismaClient user: Final = await _table(UserRepository(prisma_client)).find_unique(where={"user_id": value}) if user is not None: + shared_with: Final = tuple( + other for other in await _users_named_by_member_value(value, prisma_client) if other != value + ) + if shared_with: + verbose_proxy_logger.warning( + "SCIM: group member '%s' is one account's user id and is also account '%s' by SSO identity or email, " + "so the membership cannot be attributed. A placeholder an earlier release provisioned under this id " + "looks exactly like this and should be deleted so the real account can be matched", + value, + shared_with[0], + ) + return _AmbiguousMember(value=value) return _ResolvedUserMember(user_id=value) if member_type is not None and member_type != "user": @@ -576,6 +645,22 @@ async def _classify_group_member(member: SCIMMember, prisma_client: PrismaClient if team is not None and _team_metadata_has_scim_provenance(team.metadata): return _SkippedGroupMember(value=value, reason="existing_team") + named: Final = await _users_named_by_member_value(value, prisma_client) + if len(named) == 1: + verbose_proxy_logger.info( + "SCIM: group member '%s' matched user_id '%s' by SSO identity or email", + value, + named[0], + ) + return _ResolvedUserMember(user_id=named[0]) + if len(named) > 1: + verbose_proxy_logger.warning( + "SCIM: group member '%s' names more than one account by SSO identity or email and cannot be resolved " + "unambiguously", + value, + ) + return _AmbiguousMember(value=value) + return _UnknownMember(value=value) @@ -583,11 +668,13 @@ def _bucketed_member(entry: _ClassifiedGroupMember) -> _PartitionedMembers: """The single-member partition one classified entry contributes.""" match entry: case _ResolvedUserMember(user_id=user_id): - return _PartitionedMembers(resolved_ids=(user_id,), skipped=(), unknown_ids=()) + return _PartitionedMembers(resolved_ids=(user_id,), skipped=(), unknown_ids=(), ambiguous_values=()) case _SkippedGroupMember(): - return _PartitionedMembers(resolved_ids=(), skipped=(entry,), unknown_ids=()) + return _PartitionedMembers(resolved_ids=(), skipped=(entry,), unknown_ids=(), ambiguous_values=()) case _UnknownMember(value=value): - return _PartitionedMembers(resolved_ids=(), skipped=(), unknown_ids=(value,)) + return _PartitionedMembers(resolved_ids=(), skipped=(), unknown_ids=(value,), ambiguous_values=()) + case _AmbiguousMember(value=value): + return _PartitionedMembers(resolved_ids=(), skipped=(), unknown_ids=(), ambiguous_values=(value,)) case _: assert_never(entry) @@ -599,6 +686,7 @@ def _partition_classified_members(classified: Iterable[_ClassifiedGroupMember]) resolved_ids=tuple(chain.from_iterable(bucket.resolved_ids for bucket in bucketed)), skipped=tuple(chain.from_iterable(bucket.skipped for bucket in bucketed)), unknown_ids=tuple(chain.from_iterable(bucket.unknown_ids for bucket in bucketed)), + ambiguous_values=tuple(chain.from_iterable(bucket.ambiguous_values for bucket in bucketed)), ) @@ -608,7 +696,7 @@ def _admitted_member_id(entry: _ClassifiedGroupMember, created_ids: frozenset[st return user_id case _UnknownMember(value=value): return value if value in created_ids else None - case _SkippedGroupMember(): + case _SkippedGroupMember() | _AmbiguousMember(): return None case _: assert_never(entry) @@ -662,6 +750,70 @@ async def _ensure_group_member_user( raise HTTPException(status_code=500, detail=detail) +def _roster_entries_named_by(value: str, roster: frozenset[str], resolved: tuple[str, ...]) -> tuple[str, ...]: + """The members of this group a removal value names. + + Both ways of naming one count together. The id as written counts when the roster + holds it verbatim, which is how an earlier release recorded a member it could not + match, and the accounts it resolves to count when they are on the roster. Counting + only the resolved ones would let a value that is one member's canonical id and + another member's email revoke both, since each looks singular on its own. + """ + return tuple( + dict.fromkeys( + chain( + (value,) if value in roster else (), + (user_id for user_id in resolved if user_id in roster), + ) + ) + ) + + +async def _member_ids_to_drop( + members: Sequence[SCIMMember], roster: frozenset[str], prisma_client: PrismaClient +) -> frozenset[str]: + """The members a ``remove`` clears, one per id the request names. + + The roster holds canonical user ids, so a directory that added someone by their + email or SSO identity has to be able to remove them by that same value, and a + member an earlier release recorded under the raw id has to stay removable by it. + + Ambiguity is a property of the table as it stands, not of the value, so a value + that named one person when they were admitted can name two later. Resolving a + removal against the whole table would then drop nobody while answering 200, and + the person the directory just took out of the group would keep the team. So a + removal keeps only the accounts already on the roster: one is unambiguous however + many strangers share the address, none means there is nothing to revoke, and only + a value naming two of this group's own members is genuinely undecidable. That last + case fails rather than reporting a removal it did not perform, or revoking both. + + Raises: + HTTPException: 400 when a member id names more than one current member. + """ + written: Final = frozenset(_member_value(member) for member in members) + matched: Final = tuple( + [ + ( + value, + _roster_entries_named_by( + value, roster, await _users_named_by_member_value(value, prisma_client, take=None) + ), + ) + for value in sorted(written) + ] + ) + undecidable: Final = tuple(value for value, entries in matched if len(entries) > 1) + if undecidable: + raise HTTPException( + status_code=400, + detail={ + "error": f"Member ID '{undecidable[0]}' names more than one member of this group, so the removal " + "cannot be attributed. Send the LiteLLM user ID as the member value, or resolve the duplicate." + }, + ) + return frozenset(chain.from_iterable(entries for _, entries in matched)) + + async def _resolve_group_member_ids( members: Sequence[SCIMMember], created_via: str, @@ -670,17 +822,18 @@ async def _resolve_group_member_ids( """ Resolve SCIM group members to LiteLLM user ids, dropping members that are not users. - Only the operations that put ids onto a roster resolve their members: an id - that resolves to nothing is created when litellm_settings.scim_upsert_user is - True (default) and rejected per SCIM 2.0 otherwise. Removals do not come - through here; dropping an id is idempotent, so it needs neither a lookup nor a - user to drop. + Member ids are matched by ``user_id`` first, then by SSO identity or email. An + id that resolves to nothing is created when litellm_settings.scim_upsert_user is + True (default) and rejected per SCIM 2.0 otherwise. Removals do not come through + here: they resolve through ``_member_ids_to_drop`` instead, which neither creates + a user nor fails on an id it cannot place. Raises: - HTTPException: 400 when a member id is empty, or when scim_upsert_user is - False and a member id is neither an existing user, an existing team, nor a - member declared to be something other than a user. 500 when a member's - user row can neither be created nor found. + HTTPException: 400 when a member id is empty, when a member id names more + than one user, or when scim_upsert_user is False and a member id is neither + an existing user, an existing team, nor a member declared to be something + other than a user. 500 when a member's user row can neither be created nor + found. """ classified: Final = tuple([await _classify_group_member(member, prisma_client) for member in members]) partition: Final = _partition_classified_members(classified) @@ -692,6 +845,16 @@ async def _resolve_group_member_ids( skipped.reason, ) + if partition.ambiguous_values: + raise HTTPException( + status_code=400, + detail={ + "error": f"Member ID '{partition.ambiguous_values[0]}' names more than one LiteLLM user, so the " + "group membership cannot be attributed. Resolve the duplicate, which for an id that also matches a " + "SCIM-provisioned placeholder means deleting that placeholder." + }, + ) + if partition.unknown_ids and not await _get_scim_upsert_user_setting(): raise HTTPException( status_code=400, @@ -702,6 +865,13 @@ async def _resolve_group_member_ids( ) unique_unknown_ids: Final = tuple(dict.fromkeys(partition.unknown_ids)) + for user_id in unique_unknown_ids: + verbose_proxy_logger.warning( + "SCIM: creating placeholder user for group member '%s'; matched no user by user_id, sso_user_id or " + "user_email. An SSO-provisioned user's real account stays teamless if this is a mismatch", + user_id, + ) + creations: Final = tuple( [ ( @@ -2428,7 +2598,9 @@ async def _process_group_patch_operations( ) if op_type == "remove": - final_members = final_members - {_member_value(member) for member in patched_members} + final_members = final_members - await _member_ids_to_drop( + patched_members, frozenset(final_members), prisma_client + ) else: member_result = await _resolve_group_member_ids( members=patched_members, diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index f933cf6655e..0a9efd40b48 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py @@ -1,8 +1,13 @@ +import logging import time -from unittest.mock import AsyncMock +from collections.abc import Mapping +from itertools import chain +from typing import Final +from unittest.mock import AsyncMock, MagicMock, call import pytest from fastapi import HTTPException +from pytest_mock import MockerFixture from litellm.proxy._types import ( LiteLLM_TeamTable, @@ -72,6 +77,7 @@ async def test_create_user_existing_user_conflict(mocker): mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value={"user_id": "existing-user"}) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) # Mock the _get_prisma_client_or_raise_exception to return our mock mocker.patch( @@ -108,6 +114,7 @@ async def test_create_user_defaults_to_viewer(mocker, monkeypatch): mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) monkeypatch.setattr("litellm.default_internal_user_params", None, raising=False) @@ -158,6 +165,7 @@ async def test_create_user_ingests_enterprise_extension(mocker, monkeypatch): mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) monkeypatch.setattr("litellm.default_internal_user_params", None, raising=False) @@ -214,6 +222,7 @@ async def test_create_user_ingests_entitlements_and_roles(mocker, monkeypatch): mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) monkeypatch.setattr("litellm.default_internal_user_params", None, raising=False) @@ -263,6 +272,7 @@ async def test_create_user_uses_default_internal_user_params_role(mocker, monkey mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) # Set default_internal_user_params with a specific role @@ -362,6 +372,7 @@ async def test_scim_create_user_respects_default_role_set_via_ui(mocker, monkeyp mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) mocker.patch( @@ -1282,6 +1293,7 @@ async def test_update_group_metadata_serialization_issue(mocker): mock_user.user_email = "user1@example.com" # Add proper string value for user_email mock_user.teams = [group_id] mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=mock_user) # Mock the _get_prisma_client_or_raise_exception to return our mock @@ -1484,6 +1496,7 @@ async def test_update_group_e2e(mocker): mock_user = mocker.MagicMock() mock_user.user_id = "test-user" mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) # Mock dependencies mocker.patch( @@ -1618,6 +1631,8 @@ async def test_create_group_with_nonexistent_users_rejects(mocker, monkeypatch): return None # new-user-1 and new-user-2 don't exist mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[]) # Mock dependencies mocker.patch( @@ -1702,6 +1717,8 @@ async def test_update_group_with_nonexistent_users_rejects(mocker, monkeypatch): return None # new-user-3 and new-user-4 don't exist mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[]) # Mock dependencies mocker.patch( @@ -1771,6 +1788,8 @@ async def test_create_group_with_nonexistent_users_creates_when_flag_true(mocker return None # new-user-1 and new-user-2 don't exist mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[]) # Mock user creation created_user_1 = NewUserResponse(user_id="new-user-1", key="test-key-1") @@ -1859,6 +1878,8 @@ async def test_extract_group_member_ids_with_flag_true_creates_users(mocker, mon return None # new-user-1 doesn't exist mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[]) # Mock user creation created_user = NewUserResponse(user_id="new-user-1", key="test-key-1") @@ -1927,6 +1948,8 @@ async def test_extract_group_member_ids_with_flag_false_rejects(mocker, monkeypa return None # new-user-1 doesn't exist mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[]) # Mock dependencies mocker.patch( @@ -1976,6 +1999,8 @@ async def test_process_group_patch_operations_with_flag_true_creates_users(mocke # Mock user lookup - new-user-1 doesn't exist mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[]) mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) # Mock user creation @@ -2031,6 +2056,8 @@ async def test_process_group_patch_operations_with_flag_false_rejects(mocker, mo # Mock user lookup - new-user-1 doesn't exist mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[]) mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) # Execute the function - should raise HTTPException @@ -2070,6 +2097,7 @@ async def test_create_user_grants_admin_when_in_scim_admin_group(mocker, monkeyp mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) mocker.patch( @@ -2114,6 +2142,7 @@ async def test_create_user_keeps_default_when_not_in_scim_admin_group(mocker, mo mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) mocker.patch( @@ -2463,6 +2492,7 @@ def _scim_admin_prisma(mocker, *, user_teams): prisma.db = mocker.MagicMock() prisma.db.litellm_usertable = mocker.MagicMock() prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=user) + prisma.db.litellm_usertable.find_many = AsyncMock(return_value=()) prisma.db.litellm_usertable.update = AsyncMock(return_value=user) prisma.db.litellm_teamtable = mocker.MagicMock() prisma.db.litellm_teamtable.find_unique = AsyncMock(side_effect=_team_find_unique) @@ -2561,6 +2591,7 @@ async def test_update_group_recomputes_roles_for_changed_members(mocker): mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=existing_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -2618,6 +2649,7 @@ async def test_patch_group_recomputes_roles_for_changed_members(mocker): mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=existing_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -2671,6 +2703,7 @@ async def test_delete_group_recomputes_roles_for_members(mocker): mock_prisma_client.db.litellm_teamtable.delete = AsyncMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=member) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mock_prisma_client.db.litellm_usertable.update = AsyncMock() mocker.patch( @@ -2794,6 +2827,7 @@ async def test_create_user_existing_email_upsert_demotes_when_admin_group_set(mo mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=existing_user) mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value={"user_id": "returning-user"}) @@ -2843,6 +2877,7 @@ async def test_create_group_recomputes_roles_for_members(mocker): mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -2900,6 +2935,7 @@ async def test_update_group_rename_recomputes_retained_members(mocker): mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=existing_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -2954,6 +2990,7 @@ async def test_patch_group_rename_recomputes_retained_members(mocker): mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=existing_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -3019,6 +3056,7 @@ async def test_process_group_patch_operations_add_retains_existing_members(mocke mock_prisma_client.db.litellm_usertable = mocker.MagicMock() # new-user already exists in the DB mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock(user_id="new-user")) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) _, final_members, _ = await _process_group_patch_operations( patch_ops=patch_ops, @@ -3059,6 +3097,7 @@ async def test_process_group_patch_operations_remove_uses_members_with_roles(moc mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock(user_id="drop-user")) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) _, final_members, _ = await _process_group_patch_operations( patch_ops=patch_ops, @@ -3090,6 +3129,7 @@ async def test_get_groups_reports_members_from_members_with_roles(mocker): mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( return_value=mocker.MagicMock(user_id="member-1", user_email="member-1@example.com") ) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -3310,6 +3350,7 @@ async def test_patch_group_add_applies_delta_and_keeps_concurrent_add(mocker): mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=final_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -3403,6 +3444,7 @@ async def test_patch_group_replace_stays_absolute_against_concurrent_roster(mock mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=final_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -3503,6 +3545,7 @@ async def test_process_group_patch_remove_filtered_path_without_value(mocker): prisma_client.db = mocker.MagicMock() prisma_client.db.litellm_usertable = mocker.MagicMock() prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=LiteLLM_UserTable(user_id="user-1")) + prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) _, final_members, _ = await _process_group_patch_operations( patch_ops=patch_ops, @@ -3532,6 +3575,7 @@ async def test_process_group_patch_add_filtered_path_without_value(mocker): prisma_client.db = mocker.MagicMock() prisma_client.db.litellm_usertable = mocker.MagicMock() prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=LiteLLM_UserTable(user_id="user-3")) + prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) _, final_members, _ = await _process_group_patch_operations( patch_ops=patch_ops, @@ -3565,6 +3609,7 @@ async def test_process_group_patch_replace_empty_value_does_not_use_path_filter( prisma_client.db = mocker.MagicMock() prisma_client.db.litellm_usertable = mocker.MagicMock() prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=LiteLLM_UserTable(user_id="user-1")) + prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) _, final_members, _ = await _process_group_patch_operations( patch_ops=patch_ops, @@ -3575,7 +3620,16 @@ async def test_process_group_patch_replace_empty_value_does_not_use_path_filter( assert final_members == set() -def _member_resolution_prisma(mocker, *, users: set, teams: set, unmanaged_teams: frozenset = frozenset()): +def _member_resolution_prisma( + mocker: MockerFixture, + *, + users: set[str], + teams: set[str], + unmanaged_teams: frozenset[str] = frozenset(), + email_to_user_id: Mapping[str, str] | None = None, + email_to_user_ids: Mapping[str, tuple[str, ...]] | None = None, + sso_user_id_to_user_id: Mapping[str, str] | None = None, +) -> MagicMock: """Prisma mock where only the given ids resolve to a user row / team row. ``teams`` are teams a SCIM group write created, so they carry provenance; @@ -3589,14 +3643,78 @@ def _member_resolution_prisma(mocker, *, users: set, teams: set, unmanaged_teams return LiteLLM_TeamTable(team_id=team_id, metadata={}) return None + def user_row(where: Mapping[str, str]) -> LiteLLM_UserTable | None: + user_id: Final = where["user_id"] + if user_id in users: + return LiteLLM_UserTable(user_id=user_id) + return None + prisma_client = mocker.MagicMock() prisma_client.db = mocker.MagicMock() prisma_client.db.litellm_usertable = mocker.MagicMock() - prisma_client.db.litellm_usertable.find_unique = AsyncMock( - side_effect=lambda where: LiteLLM_UserTable(user_id=where["user_id"]) if where["user_id"] in users else None + prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=user_row) + + emails_to_ids: Final[Mapping[str, tuple[str, ...]]] = ( + dict(email_to_user_ids) + if email_to_user_ids is not None + else ({email: (user_id,) for email, user_id in email_to_user_id.items()} if email_to_user_id else {}) ) + ssos_to_ids: Final[Mapping[str, str]] = dict(sso_user_id_to_user_id) if sso_user_id_to_user_id else {} + + def identity_rows(where: Mapping[str, object], take: int | None = None) -> tuple[LiteLLM_UserTable, ...]: + """Stand-in for the cross-field lookup, honouring the comparison mode + production actually asks for per field, so a field that stops folding case, or + starts folding it, fails here instead of passing. + + A caller that must know which accounts match rather than merely how many + passes take=None, so an unbounded read returns every match. + """ + clauses: Final = where["OR"] + assert isinstance(clauses, list) + fields: Final = tuple(next(iter(clause)) for clause in clauses) + assert fields == ("sso_user_id", "user_email"), fields + + def comparison(clause: Mapping[str, object]) -> tuple[str, bool]: + """The needle and whether production asked for a case-insensitive compare, + read per field so a field that stops folding case fails here.""" + criterion = next(iter(clause.values())) + if isinstance(criterion, str): + return criterion, False + assert isinstance(criterion, dict), criterion + return criterion["equals"], criterion.get("mode") == "insensitive" + + sso_needle, sso_insensitive = comparison(clauses[0]) + email_needle, email_insensitive = comparison(clauses[1]) + + def same(stored: str, needle: str, insensitive: bool) -> bool: + return stored.casefold() == needle.casefold() if insensitive else stored == needle + + matched: Final = tuple( + chain( + ( + user_id + for sso_user_id, user_id in ssos_to_ids.items() + if same(sso_user_id, sso_needle, sso_insensitive) + ), + ( + user_id + for email, user_ids in emails_to_ids.items() + if same(email, email_needle, email_insensitive) + for user_id in user_ids + ), + ) + ) + found: Final = tuple(dict.fromkeys(matched)) + return tuple(LiteLLM_UserTable(user_id=user_id) for user_id in (found[:take] if take else found)) + + def team_lookup(where: Mapping[str, str]) -> LiteLLM_TeamTable | None: + team_id: Final = where["team_id"] + return team_row(team_id) + + prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + prisma_client.db.litellm_usertable.find_many = AsyncMock(side_effect=identity_rows) prisma_client.db.litellm_teamtable = mocker.MagicMock() - prisma_client.db.litellm_teamtable.find_unique = AsyncMock(side_effect=lambda where: team_row(where["team_id"])) + prisma_client.db.litellm_teamtable.find_unique = AsyncMock(side_effect=team_lookup) return prisma_client @@ -4363,6 +4481,581 @@ async def test_resolve_group_member_ids_dedupes_repeated_member(mocker, scim_ups assert result.all_member_ids == ["dup-user"] +def _identity_lookup(value: str) -> object: + """The single cross-field lookup the classifier is expected to issue.""" + return call( + where={"OR": [{"sso_user_id": value}, {"user_email": {"equals": value, "mode": "insensitive"}}]}, + take=2, + ) + + +@pytest.mark.asyncio +async def test_resolve_group_member_ids_matches_sso_user_id(mocker, scim_upsert_user_enabled): + """An OIDC subject in a group payload must resolve to the existing user's + internal id instead of provisioning a placeholder.""" + prisma_client = _member_resolution_prisma( + mocker, + users=set(), + teams=set(), + sso_user_id_to_user_id={"member-sub": "sso-user"}, + ) + create_user_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists", + AsyncMock(), + ) + + result = await _resolve_group_member_ids( + members=[SCIMMember(value="member-sub")], + created_via="scim_group_membership", + prisma_client=prisma_client, + ) + + create_user_mock.assert_not_called() + assert result.existing_member_ids == ["sso-user"] + assert result.created_users == [] + assert result.all_member_ids == ["sso-user"] + assert prisma_client.db.litellm_usertable.find_many.await_args_list == [_identity_lookup("member-sub")] + + +@pytest.mark.asyncio +async def test_resolve_group_member_ids_matches_user_email(mocker, scim_upsert_user_enabled): + """A group member email must resolve to the existing user's internal id + when the identity provider sends email rather than the user id.""" + prisma_client = _member_resolution_prisma( + mocker, + users=set(), + teams=set(), + email_to_user_id={"member@example.com": "email-user"}, + ) + create_user_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists", + AsyncMock(), + ) + + result = await _resolve_group_member_ids( + members=[SCIMMember(value="member@example.com")], + created_via="scim_group_membership", + prisma_client=prisma_client, + ) + + create_user_mock.assert_not_called() + assert result.existing_member_ids == ["email-user"] + assert result.created_users == [] + assert result.all_member_ids == ["email-user"] + assert prisma_client.db.litellm_usertable.find_many.await_args_list == [_identity_lookup("member@example.com")] + + +@pytest.mark.parametrize( + "pushed", + ["MEMBER@EXAMPLE.COM", "Member@Example.com", " member@example.com "], + ids=["upper", "mixed", "padded"], +) +@pytest.mark.asyncio +async def test_resolve_group_member_ids_matches_user_email_as_the_write_path_would( + mocker, scim_upsert_user_enabled, pushed +): + """The member value must be compared the way the layer that would reject a + placeholder compares it. + + ``new_user`` refuses a duplicate email case-insensitively and after stripping, so + a lookup that is stricter than that resolves nothing, creates a placeholder, and + is refused by that same layer, which surfaces as a 500 on the whole group push. + """ + prisma_client = _member_resolution_prisma( + mocker, + users=set(), + teams=set(), + email_to_user_id={"member@example.com": "email-user"}, + ) + create_user_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists", + AsyncMock(return_value=None), + ) + + result = await _resolve_group_member_ids( + members=[SCIMMember(value=pushed)], + created_via="scim_group_membership", + prisma_client=prisma_client, + ) + + create_user_mock.assert_not_called() + assert result.all_member_ids == ["email-user"] + + +@pytest.mark.parametrize( + "population", + [ + {"email_to_user_ids": {"duplicate@example.com": ("email-user-a", "email-user-b")}}, + {"email_to_user_ids": {"duplicate@example.com": ("email-user-a",), "DUPLICATE@EXAMPLE.COM": ("email-user-b",)}}, + { + "sso_user_id_to_user_id": {"duplicate@example.com": "sso-user"}, + "email_to_user_id": {"duplicate@example.com": "email-user"}, + }, + ], + ids=["same-email-twice", "emails-differing-only-in-case", "one-account-by-sso-another-by-email"], +) +@pytest.mark.asyncio +async def test_resolve_group_member_ids_rejects_a_value_naming_two_accounts( + mocker, scim_upsert_user_enabled, caplog, population +): + """A value that names two accounts names a real person we cannot identify, so + the write is refused rather than attributed to one of them. + + Every shape of collision is refused, not just two rows holding the same email + verbatim: rows whose emails differ only in case are one row to the layer that + rejects duplicates, and a value that is one account's SSO identity and another's + email would otherwise be handed to whichever field happened to be searched first. + + It must not fall through to placeholder creation. That path can only fail: the + placeholder carries ``user_email`` set to the member value, which the duplicate + email check rejects, and the recovery lookup that follows searches by ``user_id`` + and so misses the very rows that caused the collision. The operator's data problem + then surfaces as an HTTP 500 the identity provider retries forever. + """ + prisma_client = _member_resolution_prisma(mocker, users=set(), teams=set(), **population) + create_user_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists", + AsyncMock(return_value=None), + ) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + with pytest.raises(HTTPException) as exc_info: + await _resolve_group_member_ids( + members=[SCIMMember(value="duplicate@example.com")], + created_via="scim_group_membership", + prisma_client=prisma_client, + ) + + assert exc_info.value.status_code == 400 + assert "duplicate@example.com" in str(exc_info.value.detail) + assert "more than one" in str(exc_info.value.detail) + create_user_mock.assert_not_called() + assert any( + record.levelno >= logging.WARNING + and "duplicate@example.com" in record.getMessage() + and "more than one account" in record.getMessage() + for record in caplog.records + ) + + +@pytest.mark.asyncio +async def test_resolve_group_member_ids_does_not_fold_case_on_the_sso_identity(mocker, scim_upsert_user_enabled): + """An email and an SSO identity are not comparable the same way. + + OIDC defines ``sub`` as case-sensitive and nothing folds its case on the way in, + so two subjects differing only in case are two people. Folding it would hand the + group to an account the provider never named, which is the mis-grant the email + comparison is deliberately loose enough to avoid and this one is not. + """ + prisma_client = _member_resolution_prisma( + mocker, + users=set(), + teams=set(), + sso_user_id_to_user_id={"AbC-subject": "other-user"}, + ) + create_user_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists", + AsyncMock(return_value=NewUserResponse(user_id="abc-subject", key="placeholder-key")), + ) + + result = await _resolve_group_member_ids( + members=[SCIMMember(value="abc-subject")], + created_via="scim_group_membership", + prisma_client=prisma_client, + ) + + assert result.existing_member_ids == [] + assert result.all_member_ids == ["abc-subject"] + create_user_mock.assert_awaited_once_with(user_id="abc-subject", created_via="scim_group_membership") + + +@pytest.mark.asyncio +async def test_resolve_group_member_ids_ambiguous_email_outranks_upsert_rejection(mocker, scim_upsert_user_disabled): + """Ambiguity does not depend on scim_upsert_user, so the operator gets the + actionable message on either setting rather than being told to create a user that + already exists twice.""" + prisma_client = _member_resolution_prisma( + mocker, + users=set(), + teams=set(), + email_to_user_ids={"duplicate@example.com": ("email-user-a", "email-user-b")}, + ) + + with pytest.raises(HTTPException) as exc_info: + await _resolve_group_member_ids( + members=[SCIMMember(value="duplicate@example.com")], + created_via="scim_group_membership", + prisma_client=prisma_client, + ) + + assert exc_info.value.status_code == 400 + assert "more than one" in str(exc_info.value.detail) + assert "does not exist" not in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_create_group_rejects_ambiguous_member_email(mocker, scim_upsert_user_enabled): + """The refusal reaches the endpoint, so the identity provider sees a 400 on the + group write rather than a 500 it will retry.""" + scim_group = SCIMGroup( + schemas=["urn:ietf:params:scim:schemas:core:2.0:Group"], + id="ambiguous-group", + displayName="Ambiguous Group", + members=[SCIMMember(value="duplicate@example.com")], + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock( + return_value=_member_resolution_prisma( + mocker, + users=set(), + teams=set(), + email_to_user_ids={"duplicate@example.com": ("email-user-a", "email-user-b")}, + ) + ), + ) + create_user_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists", + AsyncMock(return_value=None), + ) + + with pytest.raises(ProxyException) as exc_info: + await create_group(group=scim_group) + + assert int(exc_info.value.code) == 400 + assert "duplicate@example.com" in str(exc_info.value.message) + create_user_mock.assert_not_called() + + +@pytest.mark.parametrize( + "removed_by", + ["member@example.com", "member-sub"], + ids=["by-email", "by-sso-subject"], +) +@pytest.mark.asyncio +async def test_process_group_patch_remove_by_the_id_the_directory_added_with( + mocker, scim_upsert_user_enabled, removed_by +): + """A directory removes people by the same id it added them with. + + Resolving on add and not on remove would let someone keep a team after the + directory took them out of the group: the roster holds the canonical user id, so + subtracting the email or the subject the request names would match nothing. + """ + patch_ops = SCIMPatchOp( + schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], + Operations=[SCIMPatchOperation(op="remove", path="members", value=[{"value": removed_by}])], + ) + existing_team = LiteLLM_TeamTable( + team_id="parent-group", + team_alias="Parent Group", + members=[], + members_with_roles=[Member(user_id="real-user", role="user"), Member(user_id="keep-user", role="user")], + ) + create_user_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists", + AsyncMock(return_value=NewUserResponse(user_id="phantom-user", key="phantom-key")), + ) + + _, final_members, _ = await _process_group_patch_operations( + patch_ops=patch_ops, + existing_team=existing_team, + prisma_client=_member_resolution_prisma( + mocker, + users={"real-user", "keep-user"}, + teams=set(), + email_to_user_id={"member@example.com": "real-user"}, + sso_user_id_to_user_id={"member-sub": "real-user"}, + ), + ) + + create_user_mock.assert_not_called() + assert final_members == {"keep-user"} + + +@pytest.mark.asyncio +async def test_process_group_patch_remove_still_drops_a_placeholder_by_its_literal_id( + mocker, scim_upsert_user_enabled +): + """An earlier release put unmatched ids on the roster verbatim, so a remove has to + keep clearing the id as written even once it also resolves.""" + patch_ops = SCIMPatchOp( + schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], + Operations=[SCIMPatchOperation(op="remove", path="members", value=[{"value": "legacy@example.com"}])], + ) + existing_team = LiteLLM_TeamTable( + team_id="parent-group", + team_alias="Parent Group", + members=[], + members_with_roles=[Member(user_id="legacy@example.com", role="user"), Member(user_id="keep-user", role="user")], + ) + + _, final_members, _ = await _process_group_patch_operations( + patch_ops=patch_ops, + existing_team=existing_team, + prisma_client=_member_resolution_prisma(mocker, users={"keep-user"}, teams=set()), + ) + + assert final_members == {"keep-user"} + + +@pytest.mark.asyncio +async def test_process_group_patch_remove_when_the_id_turned_ambiguous_after_admission( + mocker, scim_upsert_user_enabled +): + """Ambiguity is a property of the table as it stands, not of the value. + + Someone admitted while their email was theirs alone must stay removable after a + second account takes that email. Resolving the removal against the whole table + would find two accounts, decline to pick, drop nobody, and still answer 200, + leaving the person the directory just removed holding the team. + """ + patch_ops = SCIMPatchOp( + schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], + Operations=[SCIMPatchOperation(op="remove", path="members", value=[{"value": "shared@example.com"}])], + ) + existing_team = LiteLLM_TeamTable( + team_id="parent-group", + team_alias="Parent Group", + members=[], + members_with_roles=[Member(user_id="admitted-user", role="user"), Member(user_id="keep-user", role="user")], + ) + + _, final_members, _ = await _process_group_patch_operations( + patch_ops=patch_ops, + existing_team=existing_team, + # the newcomer took the address but never joined the group + prisma_client=_member_resolution_prisma( + mocker, + users={"admitted-user", "keep-user"}, + teams=set(), + email_to_user_ids={"shared@example.com": ("admitted-user", "newcomer")}, + ), + ) + + assert final_members == {"keep-user"} + + +@pytest.mark.asyncio +async def test_process_group_patch_remove_refuses_a_value_naming_one_member_by_id_and_another_by_email( + mocker, scim_upsert_user_enabled +): + """One value must never revoke two people. + + A SCIM-provisioned account is keyed by its userName, so a canonical user id that + looks like an email is ordinary rather than exotic, and a second account can hold + that address as its email. Counting the id as written and the resolved accounts + separately makes each look singular, and the removal then takes both. + """ + patch_ops = SCIMPatchOp( + schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], + Operations=[SCIMPatchOperation(op="remove", path="members", value=[{"value": "shared@example.com"}])], + ) + existing_team = LiteLLM_TeamTable( + team_id="parent-group", + team_alias="Parent Group", + members=[], + members_with_roles=[ + Member(user_id="shared@example.com", role="user"), + Member(user_id="other-account", role="user"), + ], + ) + + with pytest.raises(HTTPException) as exc_info: + await _process_group_patch_operations( + patch_ops=patch_ops, + existing_team=existing_team, + prisma_client=_member_resolution_prisma( + mocker, + users={"shared@example.com", "other-account"}, + teams=set(), + email_to_user_id={"shared@example.com": "other-account"}, + ), + ) + + assert exc_info.value.status_code == 400 + assert "shared@example.com" in str(exc_info.value.detail) + assert "more than one member of this group" in str(exc_info.value.detail) + + +@pytest.mark.parametrize("position", [0, 1, 2], ids=["first", "middle", "last"]) +@pytest.mark.asyncio +async def test_process_group_patch_remove_finds_the_member_past_the_bounded_read( + mocker, scim_upsert_user_enabled, position +): + """A removal has to know *which* accounts a value names, not merely whether it + names several, so it reads them all. + + An add stops after two matches, which is all it needs to decide the value is + ambiguous. Reusing that bounded read here would silently drop the member whenever + the one on the roster sorted past the cap, which no fixture smaller than the cap + can show. The member is placed at each position so the test cannot pass by luck + of ordering. + """ + strangers = ["stranger-one", "stranger-two"] + sharers = tuple(strangers[:position] + ["admitted-user"] + strangers[position:]) + patch_ops = SCIMPatchOp( + schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], + Operations=[SCIMPatchOperation(op="remove", path="members", value=[{"value": "shared@example.com"}])], + ) + existing_team = LiteLLM_TeamTable( + team_id="parent-group", + team_alias="Parent Group", + members=[], + members_with_roles=[Member(user_id="admitted-user", role="user"), Member(user_id="keep-user", role="user")], + ) + + _, final_members, _ = await _process_group_patch_operations( + patch_ops=patch_ops, + existing_team=existing_team, + prisma_client=_member_resolution_prisma( + mocker, + users={"admitted-user", "keep-user"}, + teams=set(), + email_to_user_ids={"shared@example.com": sharers}, + ), + ) + + assert final_members == {"keep-user"} + + +@pytest.mark.asyncio +async def test_process_group_patch_remove_refuses_when_two_members_share_the_id(mocker, scim_upsert_user_enabled): + """When both accounts a value names are on the roster the removal is genuinely + undecidable, so it fails rather than reporting a removal it did not perform or + revoking a membership the directory did not name.""" + patch_ops = SCIMPatchOp( + schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], + Operations=[SCIMPatchOperation(op="remove", path="members", value=[{"value": "shared@example.com"}])], + ) + existing_team = LiteLLM_TeamTable( + team_id="parent-group", + team_alias="Parent Group", + members=[], + members_with_roles=[Member(user_id="member-a", role="user"), Member(user_id="member-b", role="user")], + ) + + with pytest.raises(HTTPException) as exc_info: + await _process_group_patch_operations( + patch_ops=patch_ops, + existing_team=existing_team, + prisma_client=_member_resolution_prisma( + mocker, + users={"member-a", "member-b"}, + teams=set(), + email_to_user_ids={"shared@example.com": ("member-a", "member-b")}, + ), + ) + + assert exc_info.value.status_code == 400 + assert "shared@example.com" in str(exc_info.value.detail) + assert "more than one member of this group" in str(exc_info.value.detail) + + + +@pytest.mark.asyncio +async def test_resolve_group_member_ids_exact_user_id_wins_when_it_names_nobody_else( + mocker, scim_upsert_user_enabled +): + """The canonical user id stays authoritative, including when the same account also + holds that value as its email, which is how a SCIM-provisioned account is keyed.""" + prisma_client = _member_resolution_prisma( + mocker, + users={"member-id"}, + teams=set(), + email_to_user_id={"member-id": "member-id"}, + ) + create_user_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists", + AsyncMock(), + ) + + result = await _resolve_group_member_ids( + members=[SCIMMember(value="member-id")], + created_via="scim_group_membership", + prisma_client=prisma_client, + ) + + create_user_mock.assert_not_called() + assert result.existing_member_ids == ["member-id"] + assert result.all_member_ids == ["member-id"] + + +@pytest.mark.parametrize( + "population", + [ + {"sso_user_id_to_user_id": {"member-id": "someone-else"}}, + {"email_to_user_id": {"member-id": "someone-else"}}, + ], + ids=["another-account-by-sso", "another-account-by-email"], +) +@pytest.mark.asyncio +async def test_resolve_group_member_ids_refuses_a_user_id_that_names_another_account( + mocker, scim_upsert_user_enabled, caplog, population +): + """An exact user id is checked for collisions like every other match. + + Taking it on sight would hand the group to whichever account happened to be keyed + by the value. The placeholders this bug provisioned are exactly that shape, since + they are keyed by the very id the provider keeps pushing, so on a tenant that + already has them the real account can never win. Refusing names the problem + instead of silently landing on the placeholder again. + """ + prisma_client = _member_resolution_prisma(mocker, users={"member-id"}, teams=set(), **population) + create_user_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists", + AsyncMock(return_value=None), + ) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + with pytest.raises(HTTPException) as exc_info: + await _resolve_group_member_ids( + members=[SCIMMember(value="member-id")], + created_via="scim_group_membership", + prisma_client=prisma_client, + ) + + assert exc_info.value.status_code == 400 + assert "member-id" in str(exc_info.value.detail) + create_user_mock.assert_not_called() + assert any( + record.levelno >= logging.WARNING and "someone-else" in record.getMessage() for record in caplog.records + ) + + + +@pytest.mark.asyncio +async def test_resolve_group_member_ids_warns_before_creating_unmatched_placeholder( + mocker, scim_upsert_user_enabled, caplog +): + """An unmatched member still follows upsert behavior, but operators receive + a warning before the placeholder can leave an SSO user teamless.""" + prisma_client = _member_resolution_prisma(mocker, users=set(), teams=set()) + create_user_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists", + AsyncMock(return_value=NewUserResponse(user_id="placeholder", key="placeholder-key")), + ) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _resolve_group_member_ids( + members=[SCIMMember(value="unmatched-id")], + created_via="scim_group_membership", + prisma_client=prisma_client, + ) + + create_user_mock.assert_awaited_once_with(user_id="unmatched-id", created_via="scim_group_membership") + assert result.existing_member_ids == [] + assert result.created_users == [NewUserResponse(user_id="placeholder", key="placeholder-key")] + assert result.all_member_ids == ["unmatched-id"] + assert any( + record.levelno >= logging.WARNING + and "unmatched-id" in record.getMessage() + and "matched no user by user_id, sso_user_id or user_email" in record.getMessage() + and "real account stays teamless" in record.getMessage() + for record in caplog.records + ) + + @pytest.mark.parametrize( "operation", [ @@ -4429,6 +5122,7 @@ async def test_get_groups_members_are_typed_as_users(mocker): mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( return_value=mocker.MagicMock(user_id="member-1", user_email="member-1@example.com") ) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -4663,6 +5357,7 @@ async def test_update_group_roster_failure_propagates(mocker): mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=existing_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -4715,6 +5410,7 @@ async def test_resolve_group_member_ids_admits_member_created_concurrently(mocke prisma_client.db.litellm_usertable.find_unique = AsyncMock( side_effect=[None, LiteLLM_UserTable(user_id="raced-user")] ) + prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists", AsyncMock(return_value=None), From a112ba5f63dd9db389862b5ea46e12866a681ced Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 20 Aug 2026 19:36:26 -0700 Subject: [PATCH 036/166] test: enforce PT012 so a pytest.raises block cannot hide dead assertions (#37748) * test: enforce PT012 so a pytest.raises block cannot hide dead assertions `with pytest.raises(...)` stops at the first statement that raises. Anything sequenced after it inside the block never runs, so an assertion written there is never checked and the test still reports green. Two sites were doing exactly that, and both assertions turned out to be wrong once they started running. tests/llm_translation/test_prompt_factory.py asserted the bedrock rejection names "requires at least one non-system message", which holds. tests/proxy_unit_tests/test_proxy_server.py asserted the prisma startup failure mentions "httpx.ConnectError", which never appears: the failure is an httpx.ConnectError whose message is "All connection attempts failed", so that test now asserts the type. Its DATABASE_URL override moves to monkeypatch, since the old restore sat below the assertion and leaked the invalid URL into every later DB test the moment the assertion started being able to fail. The remaining 72 sites are rewritten without changing what they exercise: setup that cannot raise moves above the block, a nested `patch` moves outside it, and bodies with real control flow (a stream drain, an if/else on sync_mode, a retry loop) move into a local closure the block calls. Fixing PT012 unmasked two B017s, since ruff only reports a blind pytest.raises(Exception) once the block holds a single statement. tests/proxy_unit_tests/test_auth_checks.py narrows to the ProxyException can_key_call_model actually raises. tests/local_testing/test_completion_cost.py was asserting vertex_ai/medlm-medium has no cost entry, which stopped being true at some point; that dead first half is gone and the rest of the test, which checks medlm pricing resolves above zero, now runs instead of being skipped. * chore(ci): ratchet TQ004 to 768 after the prisma test moved to monkeypatch --- ruff-tests.toml | 5 ++- test-quality-budget.json | 2 +- .../test_bedrock_guardrails.py | 11 +++-- .../test_openai_responses_api.py | 5 ++- tests/llm_translation/test_prompt_factory.py | 3 +- tests/local_testing/test_aim_guardrails.py | 41 ++++++++++--------- tests/local_testing/test_completion_cost.py | 8 ---- tests/local_testing/test_exceptions.py | 5 ++- tests/local_testing/test_function_calling.py | 5 +-- tests/local_testing/test_mock_request.py | 3 +- .../test_router_budget_limiter.py | 12 ++---- tests/local_testing/test_router_fallbacks.py | 5 ++- .../test_router_max_parallel_requests.py | 5 ++- tests/local_testing/test_streaming.py | 7 +++- .../test_access_group_team_sync.py | 5 ++- tests/proxy_unit_tests/test_auth_checks.py | 6 +-- tests/proxy_unit_tests/test_jwt.py | 3 +- tests/proxy_unit_tests/test_proxy_server.py | 18 +++----- .../test_router_helper_utils.py | 14 ++++--- .../test_a2a_exception_mapping_utils.py | 7 +++- .../caching/test_redis_semantic_cache.py | 9 ++-- .../test_mcp_client.py | 7 ++-- .../bitbucket/test_bitbucket_integration.py | 31 ++++++++------ .../test_streaming_handler.py | 10 ++++- .../test_anthropic_chat_transformation.py | 6 ++- .../chat/test_bytez_chat_transformation.py | 7 ++-- .../custom_httpx/test_aiohttp_transport.py | 25 ++++++++--- .../test_credential_leak_prevention.py | 14 ++++--- .../oci/chat/test_oci_chat_transformation.py | 6 +-- .../llms/openai/test_openai_common_utils.py | 10 ++++- ...test_vertex_and_google_ai_studio_gemini.py | 5 ++- .../volcengine/test_volcengine_embedding.py | 7 ++-- .../test_async_streaming_error_propagation.py | 5 ++- .../passthrough/test_passthrough_main.py | 5 ++- ...test_streaming_interrupt_spend_tracking.py | 5 ++- .../proxy/auth/test_auth_exception_handler.py | 11 ++--- .../openai/test_moderations.py | 9 ++-- .../guardrail_hooks/test_cato_networks.py | 41 ++++++++++--------- .../guardrail_hooks/test_microsoft_purview.py | 15 +++++-- .../guardrail_hooks/test_panw_prisma_airs.py | 5 ++- .../test_prompt_security_guardrails.py | 6 +-- .../test_key_management_endpoints.py | 3 +- .../test_team_metadata_validation.py | 5 ++- .../proxy/test_budget_reservation.py | 15 +++++-- .../test_proxy_logging_hook_detection.py | 10 ++++- .../proxy/test_route_llm_request.py | 5 ++- .../repositories/test_unit_of_work.py | 10 ++++- .../test_streaming_iterator_error_events.py | 5 ++- .../test_custom_secret_manager.py | 8 ++-- tests/test_litellm/test_router.py | 15 +++++-- tests/test_ratelimit.py | 19 ++++++--- 51 files changed, 311 insertions(+), 193 deletions(-) diff --git a/ruff-tests.toml b/ruff-tests.toml index 8e90d6432df..6e77f4792a7 100644 --- a/ruff-tests.toml +++ b/ruff-tests.toml @@ -17,6 +17,9 @@ # B017 `pytest.raises(Exception)` accepts the TypeError a refactor introduced just as # readily as the rejection under test, so a crash reads as a pass. Narrow to the # real type, or add `match=` where the code genuinely raises a bare Exception +# PT012 a `pytest.raises` block that runs on past the raising call. Everything after +# that call is dead, so an `assert` sitting there is never checked. Keep the +# block to the call itself and put the assertions below it # # No target-version here on purpose: it resolves from requires-python (>=3.10), so # 3.11-only builtins like BaseExceptionGroup are correctly flagged in a tree that @@ -24,4 +27,4 @@ line-length = 120 -lint.select = ["F821", "B011", "B015", "B017", "B018", "PT015", "PLR0133", "PLW0127"] +lint.select = ["F821", "B011", "B015", "B017", "B018", "PT012", "PT015", "PLR0133", "PLW0127"] diff --git a/test-quality-budget.json b/test-quality-budget.json index fcb29c3191d..1613c8c75cb 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -9,7 +9,7 @@ "limit": 1078 }, "TQ004": { - "limit": 770 + "limit": 768 }, "TQ005": { "limit": 2832 diff --git a/tests/guardrails_tests/test_bedrock_guardrails.py b/tests/guardrails_tests/test_bedrock_guardrails.py index c7a5b79bbce..8b22cc0eb73 100644 --- a/tests/guardrails_tests/test_bedrock_guardrails.py +++ b/tests/guardrails_tests/test_bedrock_guardrails.py @@ -205,7 +205,7 @@ async def test_bedrock_guardrails_with_streaming(): mock_user_api_key_cache = MagicMock(spec=DualCache) mock_user_api_key_dict = UserAPIKeyAuth() - with pytest.raises(HTTPException): + async def _stream_through_guardrail(): proxy_logging_obj = ProxyLogging( user_api_key_cache=mock_user_api_key_cache, premium_user=True, @@ -240,6 +240,9 @@ async def test_bedrock_guardrails_with_streaming(): async for chunk in response: print(chunk) + with pytest.raises(HTTPException): + await _stream_through_guardrail() + @pytest.mark.asyncio async def test_bedrock_guardrails_with_streaming_no_violation(): @@ -1502,7 +1505,7 @@ async def test_bedrock_guardrail_disable_exception_on_block_streaming(): mock_post.return_value = mock_bedrock_response # Should raise exception during streaming processing - with pytest.raises(HTTPException): + async def _drain(): result_generator = ( guardrail_default.async_post_call_streaming_iterator_hook( user_api_key_dict=mock_user_api_key_dict, @@ -1511,10 +1514,12 @@ async def test_bedrock_guardrail_disable_exception_on_block_streaming(): ) ) - # Try to consume the generator - should raise exception async for chunk in result_generator: pass + with pytest.raises(HTTPException): + await _drain() + # Test 2: disable_exception_on_block=True. Streaming can't raise up to the # endpoint handler (SSE headers already flushed), so the block is delivered # as a synthetic stream with finish_reason=content_filter and the block diff --git a/tests/llm_responses_api_testing/test_openai_responses_api.py b/tests/llm_responses_api_testing/test_openai_responses_api.py index bd1517dbffb..d19fa09451c 100644 --- a/tests/llm_responses_api_testing/test_openai_responses_api.py +++ b/tests/llm_responses_api_testing/test_openai_responses_api.py @@ -1643,10 +1643,13 @@ async def test_openai_responses_api_token_limit_error(): model="gpt-5-mini", input=oversized_text, stream=True ) - with pytest.raises(litellm.APIError) as exc_info: + async def _drain(): async for event in response: print(event) + with pytest.raises(litellm.APIError) as exc_info: + await _drain() + assert exc_info.value.status_code == 400 assert "exceeds the context window" in str(exc_info.value) diff --git a/tests/llm_translation/test_prompt_factory.py b/tests/llm_translation/test_prompt_factory.py index ae215602e31..c3519fcb40f 100644 --- a/tests/llm_translation/test_prompt_factory.py +++ b/tests/llm_translation/test_prompt_factory.py @@ -1288,7 +1288,8 @@ def test_just_system_message(): model="anthropic.claude-3-sonnet-20240229-v1:0", llm_provider="bedrock", ) - assert "bedrock requires at least one non-system message" in str(e.value) + + assert "bedrock requires at least one non-system message" in str(e.value) def test_convert_generic_image_chunk_to_openai_image_obj(): diff --git a/tests/local_testing/test_aim_guardrails.py b/tests/local_testing/test_aim_guardrails.py index 2cb7f9cd357..5e5fb0d5459 100644 --- a/tests/local_testing/test_aim_guardrails.py +++ b/tests/local_testing/test_aim_guardrails.py @@ -101,26 +101,26 @@ async def test_block_callback(mode: str): ], } - with pytest.raises(ProxyException, match="Jailbreak detected") as exc_info: - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - return_value=Response( - json={ - "analysis_result": { - "analysis_time_ms": 212, - "policy_drill_down": {}, - "session_entities": [], - }, - "required_action": { - "action_type": "block_action", - "detection_message": "Jailbreak detected", - "policy_name": "blocking policy", - }, + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=Response( + json={ + "analysis_result": { + "analysis_time_ms": 212, + "policy_drill_down": {}, + "session_entities": [], }, - status_code=200, - request=Request(method="POST", url="http://aim"), - ), - ): + "required_action": { + "action_type": "block_action", + "detection_message": "Jailbreak detected", + "policy_name": "blocking policy", + }, + }, + status_code=200, + request=Request(method="POST", url="http://aim"), + ), + ): + async def _call_guardrail(): if mode == "pre_call": await aim_guardrail.async_pre_call_hook( data=data, @@ -135,6 +135,9 @@ async def test_block_callback(mode: str): call_type="completion", ) + with pytest.raises(ProxyException, match="Jailbreak detected") as exc_info: + await _call_guardrail() + exc = exc_info.value assert exc.code == "400" assert exc.type == "invalid_request_error" diff --git a/tests/local_testing/test_completion_cost.py b/tests/local_testing/test_completion_cost.py index e34d5c349c5..7dfcb55e29a 100644 --- a/tests/local_testing/test_completion_cost.py +++ b/tests/local_testing/test_completion_cost.py @@ -625,17 +625,9 @@ def test_vertex_ai_completion_cost(): print("calculated_input_cost: {}".format(calculated_input_cost)) -@pytest.mark.skip(reason="new test - WIP, working on fixing this") def test_vertex_ai_medlm_completion_cost(): """Test for medlm completion cost .""" - with pytest.raises(Exception) as e: - model = "vertex_ai/medlm-medium" - messages = [{"role": "user", "content": "Test MedLM completion cost."}] - predictive_cost = completion_cost( - model=model, messages=messages, custom_llm_provider="vertex_ai" - ) - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") diff --git a/tests/local_testing/test_exceptions.py b/tests/local_testing/test_exceptions.py index 8c1df52e28e..edf847f4cef 100644 --- a/tests/local_testing/test_exceptions.py +++ b/tests/local_testing/test_exceptions.py @@ -1417,7 +1417,7 @@ async def test_exception_bubbling_up(sync_mode, stream_mode, model): import litellm litellm.set_verbose = True - with pytest.raises(Exception) as exc_info: + async def _call_with_bad_role(): if sync_mode: litellm.completion( model=model, @@ -1433,6 +1433,9 @@ async def test_exception_bubbling_up(sync_mode, stream_mode, model): sync_stream=sync_mode, ) + with pytest.raises(Exception) as exc_info: + await _call_with_bad_role() + assert exc_info.value.code == "invalid_value" assert exc_info.value.param is not None assert exc_info.value.type == "invalid_request_error" diff --git a/tests/local_testing/test_function_calling.py b/tests/local_testing/test_function_calling.py index 4095962f91d..d6adde84400 100644 --- a/tests/local_testing/test_function_calling.py +++ b/tests/local_testing/test_function_calling.py @@ -357,14 +357,13 @@ def test_parallel_function_call_anthropic_error_msg( if expect_unsupported_params_error: with pytest.raises(litellm.UnsupportedParamsError) as e: - second_response = litellm.completion( + litellm.completion( model=model, messages=messages, temperature=0.2, seed=22, drop_params=True, - ) # get a new response from the model where it can see the function response - print("second response\n", second_response) + ) else: second_response = litellm.completion( model=model, diff --git a/tests/local_testing/test_mock_request.py b/tests/local_testing/test_mock_request.py index 710024b61b1..c9cd14633ba 100644 --- a/tests/local_testing/test_mock_request.py +++ b/tests/local_testing/test_mock_request.py @@ -128,13 +128,12 @@ def test_router_mock_request_with_mock_timeout(): ], ) with pytest.raises(litellm.Timeout): - response = router.completion( + router.completion( model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hey, I'm a mock request"}], timeout=3, mock_timeout=True, ) - print(response) end_time = time.time() assert end_time - start_time >= 3, f"Time taken: {end_time - start_time}" diff --git a/tests/local_testing/test_router_budget_limiter.py b/tests/local_testing/test_router_budget_limiter.py index 1a36e9de8f2..48915137138 100644 --- a/tests/local_testing/test_router_budget_limiter.py +++ b/tests/local_testing/test_router_budget_limiter.py @@ -161,12 +161,10 @@ async def test_provider_budgets_e2e_test_expect_to_fail(): for _ in range(3): with pytest.raises(Exception) as exc_info: - response = await router.acompletion( + await router.acompletion( messages=[{"role": "user", "content": "Hello, how are you?"}], model="anthropic/claude-sonnet-4-5-20250929", ) - print(response) - print("response.hidden_params", response._hidden_params) await asyncio.sleep(0.5) # Verify the error is related to budget exceeded @@ -597,12 +595,10 @@ async def test_deployment_budgets_e2e_test_expect_to_fail(): for _ in range(3): with pytest.raises(Exception) as exc_info: - response = await router.acompletion( + await router.acompletion( messages=[{"role": "user", "content": "Hello, how are you?"}], model="openai/gpt-4o-mini", ) - print(response) - print("response.hidden_params", response._hidden_params) await asyncio.sleep(0.5) # Verify the error is related to budget exceeded @@ -651,13 +647,11 @@ async def test_tag_budgets_e2e_test_expect_to_fail(): for _ in range(3): with pytest.raises(Exception) as exc_info: - response = await router.acompletion( + await router.acompletion( messages=[{"role": "user", "content": "Hello, how are you?"}], model="openai/gpt-4o-mini", metadata={"tags": [TAG_NAME]}, ) - print(response) - print("response.hidden_params", response._hidden_params) await asyncio.sleep(0.5) # Verify the error is related to budget exceeded diff --git a/tests/local_testing/test_router_fallbacks.py b/tests/local_testing/test_router_fallbacks.py index 7c09c978029..15c6c5fec59 100644 --- a/tests/local_testing/test_router_fallbacks.py +++ b/tests/local_testing/test_router_fallbacks.py @@ -1416,7 +1416,7 @@ async def test_router_fallbacks_default_and_model_specific_fallbacks(sync_mode): default_fallbacks=["bad-model"], ) - with pytest.raises(Exception) as exc_info: + async def _call_bad_model(): if sync_mode: resp = router.completion( model="bad-model", @@ -1429,6 +1429,9 @@ async def test_router_fallbacks_default_and_model_specific_fallbacks(sync_mode): model="bad-model", messages=[{"role": "user", "content": "Hey, how's it going?"}], ) + + with pytest.raises(Exception) as exc_info: + await _call_bad_model() assert isinstance( exc_info.value, litellm.AuthenticationError ), f"Expected AuthenticationError, but got {type(exc_info.value).__name__}" diff --git a/tests/local_testing/test_router_max_parallel_requests.py b/tests/local_testing/test_router_max_parallel_requests.py index 1b81b9eb999..7bb40dd7a2f 100644 --- a/tests/local_testing/test_router_max_parallel_requests.py +++ b/tests/local_testing/test_router_max_parallel_requests.py @@ -205,9 +205,12 @@ async def test_max_parallel_requests_tpm_rate_limiting_base_case(): num_retries=0, ) - with pytest.raises(litellm.RateLimitError): + async def _exceed_limit(): for _ in range(2): await router.acompletion( model="gpt-4o-2024-08-06", messages=_messages, ) + + with pytest.raises(litellm.RateLimitError): + await _exceed_limit() diff --git a/tests/local_testing/test_streaming.py b/tests/local_testing/test_streaming.py index a4f564b227f..1fe9a1ab297 100644 --- a/tests/local_testing/test_streaming.py +++ b/tests/local_testing/test_streaming.py @@ -2926,11 +2926,14 @@ def test_unit_test_custom_stream_wrapper_repeating_chunk( print(f"expected_chunk_fail: {expected_chunk_fail}") if (loop_amount > litellm.REPEATED_STREAMING_CHUNK_LIMIT) and expected_chunk_fail: + def _drain(): + for chunk in response: + continue + with pytest.raises( (litellm.InternalServerError, litellm.exceptions.MidStreamFallbackError) ): - for chunk in response: - continue + _drain() else: for chunk in response: continue diff --git a/tests/proxy_admin_ui_tests/test_access_group_team_sync.py b/tests/proxy_admin_ui_tests/test_access_group_team_sync.py index 629d77f20fc..f7092d3ec00 100644 --- a/tests/proxy_admin_ui_tests/test_access_group_team_sync.py +++ b/tests/proxy_admin_ui_tests/test_access_group_team_sync.py @@ -170,12 +170,15 @@ async def test_a_failed_mirror_takes_the_new_team_row_with_it(): async with _clean_db() as db: await _seed(db, {GROUPS[0]: [], GROUPS[1]: [OTHER_TEAM]}) - with pytest.raises(RuntimeError): + async def _blow_up_after_reconcile(): async with db.tx() as tx: await tx.litellm_teamtable.create(data={"team_id": TEAM, "access_group_ids": [GROUPS[0]]}) await reconcile_team_access_group_membership(tx, TEAM) raise RuntimeError("the cache handoff blew up") + with pytest.raises(RuntimeError): + await _blow_up_after_reconcile() + assert await _read(db) == {GROUPS[0]: [], GROUPS[1]: [OTHER_TEAM]} assert await db.litellm_teamtable.find_unique(where={"team_id": TEAM}) is None diff --git a/tests/proxy_unit_tests/test_auth_checks.py b/tests/proxy_unit_tests/test_auth_checks.py index a7fd68def2a..ffcbe472be7 100644 --- a/tests/proxy_unit_tests/test_auth_checks.py +++ b/tests/proxy_unit_tests/test_auth_checks.py @@ -242,8 +242,8 @@ async def test_can_team_call_model(model, expect_to_work): ) @pytest.mark.asyncio async def test_can_key_call_model_wildcard_access(key_models, model, expect_to_work): + from litellm.proxy._types import ProxyException from litellm.proxy.auth.auth_checks import can_key_call_model - from fastapi import HTTPException llm_model_list = [ { @@ -294,7 +294,7 @@ async def test_can_key_call_model_wildcard_access(key_models, model, expect_to_w llm_router=router, ) else: - with pytest.raises(Exception) as e: + with pytest.raises(ProxyException): await can_key_call_model( model=model, llm_model_list=llm_model_list, @@ -302,8 +302,6 @@ async def test_can_key_call_model_wildcard_access(key_models, model, expect_to_w llm_router=router, ) - print(e) - @pytest.mark.parametrize( "key_models, model, expect_to_work", diff --git a/tests/proxy_unit_tests/test_jwt.py b/tests/proxy_unit_tests/test_jwt.py index a5836c59694..4db47a1cde4 100644 --- a/tests/proxy_unit_tests/test_jwt.py +++ b/tests/proxy_unit_tests/test_jwt.py @@ -1047,8 +1047,7 @@ async def test_allow_access_by_email( else: # Expect the call to fail with pytest.raises(ProxyException): - resp = await user_api_key_auth(request=request, api_key=bearer_token) - print(resp) + await user_api_key_auth(request=request, api_key=bearer_token) def test_get_public_key_from_jwk_url(): diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index bfbc92adc74..04bc80bf0d6 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -2412,12 +2412,14 @@ async def test_proxy_server_prisma_setup(): @pytest.mark.asyncio -async def test_proxy_server_prisma_setup_invalid_db(): +async def test_proxy_server_prisma_setup_invalid_db(monkeypatch): """ PROD TEST: Test that proxy server startup fails when it's unable to connect to the database Think 2-3 times before editing / deleting this test, it's important for PROD """ + import httpx + from litellm.proxy.proxy_server import ProxyStartupEvent from litellm.proxy.utils import ProxyLogging from litellm.caching import DualCache @@ -2425,24 +2427,14 @@ async def test_proxy_server_prisma_setup_invalid_db(): user_api_key_cache = DualCache() invalid_db_url = "postgresql://invalid:invalid@localhost:5432/nonexistent" - _old_db_url = os.getenv("DATABASE_URL") - os.environ["DATABASE_URL"] = invalid_db_url + monkeypatch.setenv("DATABASE_URL", invalid_db_url) - with pytest.raises(Exception) as exc_info: + with pytest.raises(httpx.ConnectError): await ProxyStartupEvent._setup_prisma_client( database_url=invalid_db_url, proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache), user_api_key_cache=user_api_key_cache, ) - print("GOT EXCEPTION=", exc_info) - - assert "httpx.ConnectError" in str(exc_info.value) - - # # Verify the error message indicates a database connection issue - # assert any(x in str(exc_info.value).lower() for x in ["database", "connection", "authentication"]) - - if _old_db_url: - os.environ["DATABASE_URL"] = _old_db_url @pytest.mark.asyncio diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index c3db9e67f9c..1fef0f01df8 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -423,10 +423,11 @@ def test_get_timeout(model_list): def test_handle_mock_testing_fallbacks(model_list, fallback_kwarg, expected_error): """Test if the '_handle_mock_testing_fallbacks' function is working correctly""" router = Router(model_list=model_list) + data = { + fallback_kwarg: True, + } + with pytest.raises(expected_error): - data = { - fallback_kwarg: True, - } router._handle_mock_testing_fallbacks( kwargs=data, ) @@ -435,10 +436,11 @@ def test_handle_mock_testing_fallbacks(model_list, fallback_kwarg, expected_erro def test_handle_mock_testing_rate_limit_error(model_list): """Test if the '_handle_mock_testing_rate_limit_error' function is working correctly""" router = Router(model_list=model_list) + data = { + "mock_testing_rate_limit_error": True, + } + with pytest.raises(litellm.RateLimitError): - data = { - "mock_testing_rate_limit_error": True, - } router._handle_mock_testing_rate_limit_error( kwargs=data, ) diff --git a/tests/test_litellm/a2a_protocol/test_a2a_exception_mapping_utils.py b/tests/test_litellm/a2a_protocol/test_a2a_exception_mapping_utils.py index 06191d1a370..c31d50960b1 100644 --- a/tests/test_litellm/a2a_protocol/test_a2a_exception_mapping_utils.py +++ b/tests/test_litellm/a2a_protocol/test_a2a_exception_mapping_utils.py @@ -171,9 +171,12 @@ async def test_stream_with_retry_raises_after_localhost_retries_exhausted(): api_base="https://agent.example", agent_name="test-agent", ) + async def _drain(): + async for _chunk in stream: + pytest.fail("expected retry exhaustion to raise before yielding") + with pytest.raises( RuntimeError, match="no response received after retry attempts", ): - async for _chunk in stream: - pytest.fail("expected retry exhaustion to raise before yielding") + await _drain() diff --git a/tests/test_litellm/caching/test_redis_semantic_cache.py b/tests/test_litellm/caching/test_redis_semantic_cache.py index 54f1fa721a2..66271579d31 100644 --- a/tests/test_litellm/caching/test_redis_semantic_cache.py +++ b/tests/test_litellm/caching/test_redis_semantic_cache.py @@ -310,11 +310,12 @@ def test_redis_semantic_cache_reraises_unexpected_isolated_index_error(monkeypat monkeypatch.setenv("REDIS_PORT", "6379") monkeypatch.setenv("REDIS_PASSWORD", "test_password") + cache = RedisSemanticCache( + similarity_threshold=0.8, + index_name="existing_index", + ) + with pytest.raises(ValueError, match="connection failed"): - cache = RedisSemanticCache( - similarity_threshold=0.8, - index_name="existing_index", - ) _ = cache.llmcache diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index 6c3c852395b..1ddb2cc1c8d 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -75,11 +75,10 @@ class TestMCPClient: # Test missing stdio_config client = MCPClient(transport_type=MCPTransport.stdio) + async def _noop(session): + return None + with pytest.raises(ValueError, match="stdio_config is required for stdio transport"): - - async def _noop(session): - return None - await client.run_with_session(_noop) @pytest.mark.asyncio diff --git a/tests/test_litellm/integrations/bitbucket/test_bitbucket_integration.py b/tests/test_litellm/integrations/bitbucket/test_bitbucket_integration.py index 46cd1d6e765..142be536f6b 100644 --- a/tests/test_litellm/integrations/bitbucket/test_bitbucket_integration.py +++ b/tests/test_litellm/integrations/bitbucket/test_bitbucket_integration.py @@ -88,39 +88,44 @@ def test_bitbucket_prompt_manager_error_handling(mock_client_class): "access_token": "test-token", } + manager = BitBucketPromptManager(config, prompt_id="test_prompt") + with pytest.raises( Exception, match="Failed to load prompt 'test_prompt' from BitBucket" ): - manager = BitBucketPromptManager(config, prompt_id="test_prompt") - _ = manager.prompt_manager # This triggers the error + _ = manager.prompt_manager def test_bitbucket_prompt_manager_config_validation(): """Test BitBucketPromptManager configuration validation.""" # Test missing required fields - validation happens when prompt_manager is accessed - with pytest.raises( - ValueError, match="workspace, repository, and access_token are required" - ): - manager = BitBucketPromptManager({}) - _ = manager.prompt_manager # This triggers validation + manager = BitBucketPromptManager({}) with pytest.raises( ValueError, match="workspace, repository, and access_token are required" ): - manager = BitBucketPromptManager({"workspace": "test"}) - _ = manager.prompt_manager # This triggers validation + _ = manager.prompt_manager + + manager = BitBucketPromptManager({"workspace": "test"}) with pytest.raises( ValueError, match="workspace, repository, and access_token are required" ): - manager = BitBucketPromptManager({"repository": "test"}) - _ = manager.prompt_manager # This triggers validation + _ = manager.prompt_manager + + manager = BitBucketPromptManager({"repository": "test"}) with pytest.raises( ValueError, match="workspace, repository, and access_token are required" ): - manager = BitBucketPromptManager({"access_token": "test"}) - _ = manager.prompt_manager # This triggers validation + _ = manager.prompt_manager + + manager = BitBucketPromptManager({"access_token": "test"}) + + with pytest.raises( + ValueError, match="workspace, repository, and access_token are required" + ): + _ = manager.prompt_manager @patch("litellm.integrations.bitbucket.bitbucket_prompt_manager.BitBucketClient") diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index fc7c81a9bab..bc457578e1e 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -2143,10 +2143,13 @@ def test_raise_on_model_repetition( chunks = _build_chunks(chunks_pattern, len(chunks_pattern)) if should_raise: - with pytest.raises(litellm.InternalServerError) as exc_info: + def _feed(): for chunk in chunks: wrapper.chunks.append(chunk) wrapper.raise_on_model_repetition() + + with pytest.raises(litellm.InternalServerError) as exc_info: + _feed() assert "repeating the same chunk" in str(exc_info.value) else: for chunk in chunks: @@ -3616,10 +3619,13 @@ async def test_transport_read_error_before_finish_reason_raises(logging_obj: Log ) received = [] - with pytest.raises(MidStreamFallbackError): + async def _drain(): async for chunk in response: received.append(chunk) + with pytest.raises(MidStreamFallbackError): + await _drain() + fabricated_finish_reasons = [ chunk.choices[0].finish_reason for chunk in received diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 391bd8566a2..d6aa384e03d 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -1986,10 +1986,11 @@ def test_effort_validation(): ) assert result["output_config"]["effort"] == effort + optional_params = {"output_config": {"effort": "invalid"}} + with pytest.raises( litellm.exceptions.BadRequestError, match="Invalid effort value" ): - optional_params = {"output_config": {"effort": "invalid"}} config.transform_request( model="claude-opus-4-5-20251101", messages=messages, @@ -2043,11 +2044,12 @@ def test_max_effort_rejected_for_opus_45(): messages = [{"role": "user", "content": "Test"}] + optional_params = {"output_config": {"effort": "max"}} + with pytest.raises( litellm.exceptions.BadRequestError, match="effort='max' is not supported by this model", ): - optional_params = {"output_config": {"effort": "max"}} config.transform_request( model="claude-opus-4-5-20251101", messages=messages, diff --git a/tests/test_litellm/llms/bytez/chat/test_bytez_chat_transformation.py b/tests/test_litellm/llms/bytez/chat/test_bytez_chat_transformation.py index 2f8cc5484ba..e2421437720 100644 --- a/tests/test_litellm/llms/bytez/chat/test_bytez_chat_transformation.py +++ b/tests/test_litellm/llms/bytez/chat/test_bytez_chat_transformation.py @@ -35,11 +35,10 @@ class TestBytezChatConfig: assert result["user-agent"] == f"litellm/{version}" def test_missing_api_key(self): + config = BytezChatConfig() + headers = {} + with pytest.raises(Exception) as excinfo: - config = BytezChatConfig() - - headers = {} - config.validate_environment( headers=headers, model=TEST_MODEL, diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py b/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py index 6a8cd29692f..2dc7fbfd62a 100644 --- a/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py +++ b/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py @@ -127,10 +127,13 @@ async def test_client_payload_error_mid_stream_raises_read_error(): stream = AiohttpResponseStream(mock_response) # type: ignore received_chunks = [] - with pytest.raises(httpx.ReadError): + async def _drain(): async for chunk in stream: received_chunks.append(chunk) + with pytest.raises(httpx.ReadError): + await _drain() + assert received_chunks == [b"chunk1"] assert mock_response.closed is True @@ -151,10 +154,13 @@ async def test_client_payload_error_before_first_chunk_raises_read_error(): stream = AiohttpResponseStream(mock_response) # type: ignore received_chunks = [] - with pytest.raises(httpx.ReadError): + async def _drain(): async for chunk in stream: received_chunks.append(chunk) + with pytest.raises(httpx.ReadError): + await _drain() + assert received_chunks == [] assert mock_response.closed is True @@ -171,10 +177,13 @@ async def test_connection_closed_runtime_error_raises_read_error(): stream = AiohttpResponseStream(mock_response) # type: ignore received_chunks = [] - with pytest.raises(httpx.ReadError): + async def _drain(): async for chunk in stream: received_chunks.append(chunk) + with pytest.raises(httpx.ReadError): + await _drain() + assert received_chunks == [b"data1"] assert mock_response.closed is True @@ -209,10 +218,13 @@ async def test_transfer_encoding_error_raises_read_error(): stream = AiohttpResponseStream(mock_response) # type: ignore received_chunks = [] - with pytest.raises(httpx.ReadError): + async def _drain(): async for chunk in stream: received_chunks.append(chunk) + with pytest.raises(httpx.ReadError): + await _drain() + assert received_chunks == [b"data1"] assert mock_response.closed is True @@ -254,10 +266,13 @@ async def test_timeout_exception_gets_mapped(): received_chunks = [] # This should raise httpx.TimeoutException (mapped from aiohttp.ServerTimeoutError) - with pytest.raises(httpx.TimeoutException): + async def _drain(): async for chunk in stream: received_chunks.append(chunk) + with pytest.raises(httpx.TimeoutException): + await _drain() + # Should have received the first chunk before the error assert received_chunks == [b"chunk1"] diff --git a/tests/test_litellm/llms/custom_httpx/test_credential_leak_prevention.py b/tests/test_litellm/llms/custom_httpx/test_credential_leak_prevention.py index 0a3bf403bf8..bd9db87a765 100644 --- a/tests/test_litellm/llms/custom_httpx/test_credential_leak_prevention.py +++ b/tests/test_litellm/llms/custom_httpx/test_credential_leak_prevention.py @@ -287,10 +287,11 @@ class TestHTTPHandlerErrorPaths: "send", side_effect=_make_httpx_status_error(url="https://api.test.com?key=SECRET"), ): + kwargs = {"url": "https://api.test.com?key=SECRET"} + if method != "delete": + kwargs["data"] = {"test": 1} + with pytest.raises(MaskedHTTPStatusError) as exc_info: - kwargs = {"url": "https://api.test.com?key=SECRET"} - if method != "delete": - kwargs["data"] = {"test": 1} getattr(sync_handler, method)(**kwargs) assert "SECRET" not in str(exc_info.value.request.url) @@ -304,10 +305,11 @@ class TestHTTPHandlerErrorPaths: new_callable=AsyncMock, side_effect=_make_httpx_status_error(url="https://api.test.com?key=SECRET"), ): + kwargs = {"url": "https://api.test.com?key=SECRET"} + if method != "delete": + kwargs["data"] = {"test": 1} + with pytest.raises(MaskedHTTPStatusError) as exc_info: - kwargs = {"url": "https://api.test.com?key=SECRET"} - if method != "delete": - kwargs["data"] = {"test": 1} await getattr(async_handler, method)(**kwargs) assert "SECRET" not in str(exc_info.value.request.url) diff --git a/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py b/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py index 0f0033cae36..8be0780d86f 100644 --- a/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py +++ b/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py @@ -95,10 +95,10 @@ class TestOCIChatConfig: modified_params = params.copy() del modified_params[key] - with pytest.raises(Exception) as excinfo: - config = OCIChatConfig() - headers = {} + config = OCIChatConfig() + headers = {} + with pytest.raises(Exception) as excinfo: config.validate_environment( headers=headers, model=TEST_MODEL, diff --git a/tests/test_litellm/llms/openai/test_openai_common_utils.py b/tests/test_litellm/llms/openai/test_openai_common_utils.py index a28e133700e..bfd681cc06e 100644 --- a/tests/test_litellm/llms/openai/test_openai_common_utils.py +++ b/tests/test_litellm/llms/openai/test_openai_common_utils.py @@ -375,20 +375,26 @@ async def test_async_streaming_output_limit_400_maps_to_length_truncated_stream( @pytest.mark.parametrize("provider", ["openai", "azure"]) @pytest.mark.parametrize("stream", [False, True]) def test_sync_genuine_bad_request_still_raises(provider, stream): - with pytest.raises(litellm.BadRequestError): + def _call_and_drain(): result = litellm.completion( **_completion_kwargs(provider, _sync_client_raising(provider, GENUINE_400_MESSAGE), stream=stream) ) list(result) + with pytest.raises(litellm.BadRequestError): + _call_and_drain() + @pytest.mark.parametrize("provider", ["openai", "azure"]) @pytest.mark.parametrize("stream", [False, True]) @pytest.mark.asyncio async def test_async_genuine_bad_request_still_raises(provider, stream): - with pytest.raises(litellm.BadRequestError): + async def _call_and_drain(): result = await litellm.acompletion( **_completion_kwargs(provider, _async_client_raising(provider, GENUINE_400_MESSAGE), stream=stream) ) async for _ in result: pass + + with pytest.raises(litellm.BadRequestError): + await _call_and_drain() diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 51cc2857252..b7265ed62e9 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -5246,11 +5246,14 @@ def test_mid_stream_429_error_raises_during_iteration(): # Iterate the stream: first chunks should succeed, then 429 error should be raised results = [] - with pytest.raises(VertexAIError) as exc_info: + def _drain(): for chunk in streaming_obj: if chunk is not None: results.append(chunk) + with pytest.raises(VertexAIError) as exc_info: + _drain() + # Verify: received normal chunks before the error assert ( len(results) >= 1 diff --git a/tests/test_litellm/llms/volcengine/test_volcengine_embedding.py b/tests/test_litellm/llms/volcengine/test_volcengine_embedding.py index 6a035bcd7f0..07298f03f86 100644 --- a/tests/test_litellm/llms/volcengine/test_volcengine_embedding.py +++ b/tests/test_litellm/llms/volcengine/test_volcengine_embedding.py @@ -198,10 +198,11 @@ def test_volcengine_embedding_error_scenarios(): mock_embedding.side_effect = ValueError("Unsupported encoding_format") # Test that errors are properly raised + test_params = { + k: v for k, v in scenario.items() if k != "expected_error_pattern" + } + with pytest.raises(Exception) as exc_info: - test_params = { - k: v for k, v in scenario.items() if k != "expected_error_pattern" - } litellm.embedding(input=["test"], **test_params) # Verify error message contains expected pattern diff --git a/tests/test_litellm/passthrough/test_async_streaming_error_propagation.py b/tests/test_litellm/passthrough/test_async_streaming_error_propagation.py index d262063584b..faf4ea46c43 100644 --- a/tests/test_litellm/passthrough/test_async_streaming_error_propagation.py +++ b/tests/test_litellm/passthrough/test_async_streaming_error_propagation.py @@ -65,7 +65,7 @@ async def test_async_streaming_429_raises(): return mock_response chunks = [] - with pytest.raises(httpx.HTTPStatusError) as exc_info: + async def _drain(): async for chunk in _async_streaming( response=response_coro(), litellm_logging_obj=_make_mock_logging_obj(), @@ -73,6 +73,9 @@ async def test_async_streaming_429_raises(): ): chunks.append(chunk) + with pytest.raises(httpx.HTTPStatusError) as exc_info: + await _drain() + assert exc_info.value.response.status_code == 429 assert len(chunks) == 0 diff --git a/tests/test_litellm/passthrough/test_passthrough_main.py b/tests/test_litellm/passthrough/test_passthrough_main.py index 58a0185ea8c..965f9fd8f7d 100644 --- a/tests/test_litellm/passthrough/test_passthrough_main.py +++ b/tests/test_litellm/passthrough/test_passthrough_main.py @@ -721,10 +721,13 @@ async def test_allm_passthrough_route_429_streaming_raises(): # result is an async generator — consuming it must raise, not silently yield error bytes chunks = [] - with pytest.raises(httpx.HTTPStatusError) as exc_info: + async def _drain(): async for chunk in result: # type: ignore[union-attr] chunks.append(chunk) + with pytest.raises(httpx.HTTPStatusError) as exc_info: + await _drain() + assert exc_info.value.response.status_code == 429 assert len(chunks) == 0, "No chunks should be yielded before the 429 raises" diff --git a/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py b/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py index f3fe3ae5c38..3783e218e4e 100644 --- a/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py +++ b/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py @@ -164,7 +164,7 @@ async def test_async_streaming_flushes_on_upstream_exception_with_partial_data() provider_config = MagicMock() received = [] - with pytest.raises(httpx.ReadError): + async def _drain(): async for chunk in _async_streaming( response=response_coro(), litellm_logging_obj=mock_logging_obj, @@ -172,6 +172,9 @@ async def test_async_streaming_flushes_on_upstream_exception_with_partial_data() ): received.append(chunk) + with pytest.raises(httpx.ReadError): + await _drain() + assert received == partial_chunks await asyncio.sleep(0) diff --git a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py index b4725a81823..721857e5411 100644 --- a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py +++ b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py @@ -338,12 +338,13 @@ async def test_handle_authentication_error_budget_exceeded(): mock_api_key = "test-key" # Test with budget exceeded error - with pytest.raises(ProxyException) as exc_info: - from litellm.exceptions import BudgetExceededError + from litellm.exceptions import BudgetExceededError - budget_error = BudgetExceededError( - message="Budget exceeded", current_cost=100, max_budget=100 - ) + budget_error = BudgetExceededError( + message="Budget exceeded", current_cost=100, max_budget=100 + ) + + with pytest.raises(ProxyException) as exc_info: await handler._handle_authentication_error( budget_error, mock_request, diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py index 9002d1f81a3..729dcb54309 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py @@ -487,17 +487,18 @@ async def test_openai_moderation_guardrail_streaming_harmful_content(): # Should raise HTTPException when processing streaming harmful content from fastapi import HTTPException - with pytest.raises(HTTPException) as exc_info: + async def _drain(): result_chunks = [] - async for ( - chunk - ) in unified_guardrail.async_post_call_streaming_iterator_hook( + async for chunk in unified_guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=user_api_key_dict, response=mock_stream(), request_data=request_data, ): result_chunks.append(chunk) + with pytest.raises(HTTPException) as exc_info: + await _drain() + assert exc_info.value.status_code == 400 assert "Violated OpenAI moderation policy" in str(exc_info.value.detail) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cato_networks.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cato_networks.py index 428f2faf041..c23fbc0234e 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cato_networks.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cato_networks.py @@ -93,26 +93,26 @@ async def test_block_callback(mode: str): ], } - with pytest.raises(HTTPException, match="Jailbreak detected"): - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - return_value=Response( - json={ - "analysis_result": { - "analysis_time_ms": 212, - "policy_drill_down": {}, - "session_entities": [], - }, - "required_action": { - "action_type": "block_action", - "detection_message": "Jailbreak detected", - "policy_name": "blocking policy", - }, + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=Response( + json={ + "analysis_result": { + "analysis_time_ms": 212, + "policy_drill_down": {}, + "session_entities": [], }, - status_code=200, - request=Request(method="POST", url="http://cato"), - ), - ): + "required_action": { + "action_type": "block_action", + "detection_message": "Jailbreak detected", + "policy_name": "blocking policy", + }, + }, + status_code=200, + request=Request(method="POST", url="http://cato"), + ), + ): + async def _call_guardrail(): if mode == "pre_call": await cato_guardrail.async_pre_call_hook( data=data, @@ -127,6 +127,9 @@ async def test_block_callback(mode: str): call_type="completion", ) + with pytest.raises(HTTPException, match="Jailbreak detected"): + await _call_guardrail() + @pytest.mark.asyncio @pytest.mark.parametrize("mode", ["pre_call", "during_call"]) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py index cc89cea58d2..4a7a14fceaa 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py @@ -2441,7 +2441,7 @@ class TestStreamingIteratorHook: ), ): chunks = [] - with pytest.raises(HTTPException) as exc_info: + async def _drain(): async for chunk in guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=UserAPIKeyAuth( api_key="test", user_id="user-123" @@ -2451,6 +2451,9 @@ class TestStreamingIteratorHook: ): chunks.append(chunk) + with pytest.raises(HTTPException) as exc_info: + await _drain() + assert exc_info.value.status_code == 400 assert len(chunks) == 0 # No chunks yielded before the block @@ -2477,7 +2480,7 @@ class TestStreamingIteratorHook: "litellm.main.stream_chunk_builder", return_value=assembled_response ): chunks = [] - with pytest.raises(HTTPException) as exc_info: + async def _drain(): async for chunk in guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=UserAPIKeyAuth(api_key="test"), # no user_id response=fake_response_stream(), @@ -2485,6 +2488,9 @@ class TestStreamingIteratorHook: ): chunks.append(chunk) + with pytest.raises(HTTPException) as exc_info: + await _drain() + assert exc_info.value.status_code == 400 assert len(chunks) == 0 @@ -2625,7 +2631,7 @@ class TestStreamingIteratorHook: ), ): chunks = [] - with pytest.raises(HTTPException) as exc_info: + async def _drain(): async for chunk in guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=UserAPIKeyAuth( api_key="test", user_id="user-123" @@ -2635,6 +2641,9 @@ class TestStreamingIteratorHook: ): chunks.append(chunk) + with pytest.raises(HTTPException) as exc_info: + await _drain() + assert exc_info.value.status_code == 400 assert len(chunks) == 0 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py index 86a7ac1dabe..2284f2b678a 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py @@ -5902,7 +5902,7 @@ class TestPanwAirsBlockedErrorDetailPassthrough: with patch.object( base_handler, "_call_panw_api", return_value=copy.deepcopy(self._FULL_BLOCK_RESPONSE) ): - with pytest.raises(HTTPException) as exc_info: + async def _call_hook(): if is_response: await base_handler.async_post_call_success_hook( data=safe_prompt_data, @@ -5917,6 +5917,9 @@ class TestPanwAirsBlockedErrorDetailPassthrough: call_type="completion", ) + with pytest.raises(HTTPException) as exc_info: + await _call_hook() + error = exc_info.value.detail["error"] for field, value in self._FULL_BLOCK_RESPONSE.items(): if field == "category": diff --git a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py index f35d64b89e3..c8f22e6c15e 100644 --- a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py @@ -472,9 +472,9 @@ async def test_file_sanitization_block(): async def mock_get(*args, **kwargs): return mock_poll_response - with pytest.raises(HTTPException) as excinfo: - with patch.object(guardrail.async_handler, "post", side_effect=mock_post): - with patch.object(guardrail.async_handler, "get", side_effect=mock_get): + with patch.object(guardrail.async_handler, "post", side_effect=mock_post): + with patch.object(guardrail.async_handler, "get", side_effect=mock_get): + with pytest.raises(HTTPException) as excinfo: await guardrail.apply_guardrail( inputs=inputs, request_data=request_data, diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 8e661af8daa..d4e9ccdca5e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -11037,8 +11037,7 @@ class TestKeyAliasSkipValidationOnUnchanged: assert new_alias != existing_alias with pytest.raises(ProxyException): - if new_alias != existing_alias: - _validate_key_alias_format(new_alias) + _validate_key_alias_format(new_alias) @pytest.mark.asyncio async def test_update_key_changed_to_valid_alias_passes( diff --git a/tests/test_litellm/proxy/management_helpers/test_team_metadata_validation.py b/tests/test_litellm/proxy/management_helpers/test_team_metadata_validation.py index 4c66f4aadf1..e4b031ade57 100644 --- a/tests/test_litellm/proxy/management_helpers/test_team_metadata_validation.py +++ b/tests/test_litellm/proxy/management_helpers/test_team_metadata_validation.py @@ -573,12 +573,15 @@ async def test_http_validator_service_outage_fails_closed(monkeypatch, kind, exi monkeypatch.setenv("TEAM_METADATA_VALIDATION_SERVICE_URL", _closed_port_url()) with _configured(impls.validate_via_http): - with pytest.raises(ProxyException) as exc_info: + async def _drive(): if kind == "create": await _drive_create(metadata=request_payload) else: await _drive_update(kind, existing_metadata, request_payload) + with pytest.raises(ProxyException) as exc_info: + await _drive() + assert str(exc_info.value.code) == "503" assert DEFAULT_TEAM_METADATA_VALIDATION_UNAVAILABLE_MESSAGE in str(exc_info.value.message) diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 133b53bb18d..2388654bf4b 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -2412,10 +2412,13 @@ async def test_streaming_cancel_before_any_chunk_reconciles_to_input_cost( generator, streaming_logging_obj = _drive_streaming_cancel(valid_token, cancel_before_chunk) received = [] - with pytest.raises(asyncio.CancelledError): + async def _drain(): async for chunk in generator: received.append(chunk) + with pytest.raises(asyncio.CancelledError): + await _drain() + assert received == [] # no chunk delivered, but the provider already received the input, so the # reservation is reconciled to the input cost (0.5), not refunded to zero @@ -2444,10 +2447,13 @@ async def test_streaming_cancel_after_chunk_keeps_reservation( generator, streaming_logging_obj = _drive_streaming_cancel(valid_token, cancel_after_chunk) received = [] - with pytest.raises(asyncio.CancelledError): + async def _drain(): async for chunk in generator: received.append(chunk) + with pytest.raises(asyncio.CancelledError): + await _drain() + assert received == ["data: chunk\n\n"] # a consumed stream must NOT be refunded assert counter_cache.in_memory_cache.get_cache( @@ -2508,10 +2514,13 @@ async def test_streaming_cancel_in_slow_path_before_yield_refunds(spend_counter_ received = [] # include_cost_in_streaming_usage forces fast_path off, so the hook above runs with patch.object(litellm, "include_cost_in_streaming_usage", True, create=True): - with pytest.raises(asyncio.CancelledError): + async def _drain(): async for chunk in generator: received.append(chunk) + with pytest.raises(asyncio.CancelledError): + await _drain() + assert received == [] # cancellation happened before any chunk reached the client, but the # provider already received the input -> reconcile to the input cost (0.5) diff --git a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py index 133156f9321..542572e1e56 100644 --- a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py +++ b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py @@ -291,7 +291,7 @@ async def test_post_call_stream_guardrail_blocks_anthropic_messages_stream(monke yield chunk delivered = [] - with pytest.raises(HTTPException) as exc_info: + async def _drain(): async for chunk in proxy_logging.async_post_call_streaming_iterator_hook( response=fake_stream(), user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages"), @@ -299,6 +299,9 @@ async def test_post_call_stream_guardrail_blocks_anthropic_messages_stream(monke ): delivered.append(chunk) + with pytest.raises(HTTPException) as exc_info: + await _drain() + detail = exc_info.value.detail assert detail["guardrail_name"] == "output-filter" assert detail["keyword"] == "zebra" @@ -411,7 +414,7 @@ async def test_post_call_stream_guardrail_reroutes_inherited_apply_guardrail(mon yield chunk delivered = [] - with pytest.raises(HTTPException) as exc_info: + async def _drain(): async for chunk in proxy_logging.async_post_call_streaming_iterator_hook( response=fake_stream(), user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages"), @@ -419,6 +422,9 @@ async def test_post_call_stream_guardrail_reroutes_inherited_apply_guardrail(mon ): delivered.append(chunk) + with pytest.raises(HTTPException) as exc_info: + await _drain() + assert exc_info.value.detail["keyword"] == "zebra" assert delivered == [] diff --git a/tests/test_litellm/proxy/test_route_llm_request.py b/tests/test_litellm/proxy/test_route_llm_request.py index 1e716f7c148..23e0bbfb3ee 100644 --- a/tests/test_litellm/proxy/test_route_llm_request.py +++ b/tests/test_litellm/proxy/test_route_llm_request.py @@ -169,7 +169,7 @@ async def test_route_request_proxy_admin_can_call_all_team_scoped_deployments_wi ) ) - with pytest.raises(litellm.BadRequestError, match="multiple teams"): + async def _route_and_await(): ambiguous_call = await route_request( data=data, llm_router=router, @@ -179,6 +179,9 @@ async def test_route_request_proxy_admin_can_call_all_team_scoped_deployments_wi ) await ambiguous_call + with pytest.raises(litellm.BadRequestError, match="multiple teams"): + await _route_and_await() + router.add_deployment( Deployment( model_name="team-azure", diff --git a/tests/test_litellm/repositories/test_unit_of_work.py b/tests/test_litellm/repositories/test_unit_of_work.py index c270a570ad9..1ebfd917e36 100644 --- a/tests/test_litellm/repositories/test_unit_of_work.py +++ b/tests/test_litellm/repositories/test_unit_of_work.py @@ -59,11 +59,14 @@ async def test_updates_across_tables_share_one_batch_and_commit_once(): async def test_raising_inside_block_skips_commit(): batch = FakeBatch() - with pytest.raises(RuntimeError, match="boom"): + async def _blow_up_mid_transaction(): async with spend_reset_unit_of_work(lambda: batch) as uow: uow.keys.queue_spend_reset(token="tok-1", budget_reset_at=None) raise RuntimeError("boom") + with pytest.raises(RuntimeError, match="boom"): + await _blow_up_mid_transaction() + assert batch.commit_count == 0 @@ -119,9 +122,12 @@ async def test_budget_cascade_raising_inside_block_skips_commit(): the tier is still due on the next tick.""" batch = FakeBatch() - with pytest.raises(RuntimeError, match="boom"): + async def _blow_up_mid_transaction(): async with budget_cascade_unit_of_work(lambda: batch) as uow: uow.team_memberships.queue_spend_zero(where={"budget_id": {"in": ["budget-1"]}}) raise RuntimeError("boom") + with pytest.raises(RuntimeError, match="boom"): + await _blow_up_mid_transaction() + assert batch.commit_count == 0 diff --git a/tests/test_litellm/responses/test_streaming_iterator_error_events.py b/tests/test_litellm/responses/test_streaming_iterator_error_events.py index 3b87246ebdb..321abe4cc6d 100644 --- a/tests/test_litellm/responses/test_streaming_iterator_error_events.py +++ b/tests/test_litellm/responses/test_streaming_iterator_error_events.py @@ -208,9 +208,12 @@ async def test_async_iterator_error_after_first_chunk_carries_generated_content( ) chunks = [] - with pytest.raises(MidStreamFallbackError) as exc_info: + async def _drain(): async for chunk in iterator: chunks.append(chunk) + + with pytest.raises(MidStreamFallbackError) as exc_info: + await _drain() assert len(chunks) == 2 assert exc_info.value.status_code == 500 assert exc_info.value.is_pre_first_chunk is False diff --git a/tests/test_litellm/secret_managers/test_custom_secret_manager.py b/tests/test_litellm/secret_managers/test_custom_secret_manager.py index 1f4f9a47671..0426c5973cc 100644 --- a/tests/test_litellm/secret_managers/test_custom_secret_manager.py +++ b/tests/test_litellm/secret_managers/test_custom_secret_manager.py @@ -243,17 +243,17 @@ def test_minimal_custom_secret_manager(): assert value == "sync-TEST_KEY-value" # Write should raise NotImplementedError - with pytest.raises(NotImplementedError) as exc_info: - import asyncio + import asyncio + with pytest.raises(NotImplementedError) as exc_info: asyncio.run(secret_manager.async_write_secret("KEY", "value")) assert "Write operations are not implemented" in str(exc_info.value) # Delete should raise NotImplementedError - with pytest.raises(NotImplementedError) as exc_info: - import asyncio + import asyncio + with pytest.raises(NotImplementedError) as exc_info: asyncio.run(secret_manager.async_delete_secret("KEY")) assert "Delete operations are not implemented" in str(exc_info.value) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 9894fcef163..b50dc92c220 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1999,10 +1999,13 @@ async def test_acompletion_streaming_iterator(): # Collect streamed chunks — the first chunk succeeds, then the error re-raises collected_chunks = [] - with pytest.raises(MidStreamFallbackError): + async def _drain(): async for chunk in result: collected_chunks.append(chunk) + with pytest.raises(MidStreamFallbackError): + await _drain() + assert len(collected_chunks) == 1, "one chunk yielded before the error" print("✓ MidStreamFallbackError re-raised correctly when content was already generated") @@ -5557,10 +5560,13 @@ async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_f initial_kwargs=dict(initial_kwargs), ) collected = [] - with pytest.raises(MidStreamFallbackError): + async def _drain(): async for chunk in result: collected.append(chunk) + with pytest.raises(MidStreamFallbackError): + await _drain() + assert len(collected) == 1 logging_obj.dispatch_success_handlers.assert_not_called() @@ -5580,10 +5586,13 @@ async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_f initial_kwargs=dict(initial_kwargs), ) collected = [] - with pytest.raises(MidStreamFallbackError): + async def _drain(): async for chunk in result: collected.append(chunk) + with pytest.raises(MidStreamFallbackError): + await _drain() + assert len(collected) == 1, "only the partial chunk before the error" mock_fallback.assert_not_called() logging_obj.dispatch_success_handlers.assert_not_called() diff --git a/tests/test_ratelimit.py b/tests/test_ratelimit.py index 0469ded3f42..121dfbd99b7 100644 --- a/tests/test_ratelimit.py +++ b/tests/test_ratelimit.py @@ -149,19 +149,26 @@ def test_async_rate_limit( router: Router = router_factory(rpm, tpm, routing_strategy) print(f"router: {router.model_list}") - with pytest.raises(expected_exception) as excinfo: # asserts correct type raised - if sync_mode: - results = sync_call(router, list_of_messages) - else: - results = asyncio.run(async_call(router, list_of_messages)) + received = [] + + def _send_and_check(): + results = ( + sync_call(router, list_of_messages) + if sync_mode + else asyncio.run(async_call(router, list_of_messages)) + ) + received.extend(results) print(results) if len([i for i in results if i is not None]) != num_try_send: # since not all results got returned, raise rate limit error raise ValueError("No deployments available for selected model") raise ExpectNoException + with pytest.raises(expected_exception) as excinfo: # asserts correct type raised + _send_and_check() + print(expected_exception, excinfo) if expected_exception is ValueError: assert "No deployments available for selected model" in str(excinfo.value) else: - assert len([i for i in results if i is not None]) == num_try_send + assert len([i for i in received if i is not None]) == num_try_send From 7d999a15864a2821dfd85f023c5064fe497eb897 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 19:41:12 -0700 Subject: [PATCH 037/166] fix(cognition): price swe-1.7 at the standard tier, add swe-1.7-lightning The cost map shipped cognition/swe-1.7 at $2.50 in / $12.50 out per million with $1.00 cache reads. Those are the Lightning numbers. Cognition's own model list at https://docs.devin.ai/desktop/models has uid swe-1-7 at $0.50 / $2.50 with $0.20 cache reads, and uid swe-1-7-lightning at $2.50 / $12.50 with $1.00 cache reads, so every swe-1.7 call has been costed at 5x since the entry landed. swe-1.7 now carries the standard rates and the Lightning tier gets its own entry, in both cost map copies. The source field on both moves to the desktop models page, which is the one that lists both tiers. --- ...odel_prices_and_context_window_backup.json | 12 +++- model_prices_and_context_window.json | 12 +++- .../openai_like/test_cognition_provider.py | 58 ++++++++++++++++--- 3 files changed, 72 insertions(+), 10 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index efbe0d2ebb7..c5579f03a3e 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -48949,6 +48949,16 @@ "source": "https://docs.devin.ai/windsurf/plugins/cascade/models" }, "cognition/swe-1.7": { + "input_cost_per_token": 5e-07, + "output_cost_per_token": 2.5e-06, + "cache_read_input_token_cost": 2e-07, + "litellm_provider": "cognition", + "mode": "chat", + "supports_function_calling": true, + "supports_prompt_caching": true, + "source": "https://docs.devin.ai/desktop/models" + }, + "cognition/swe-1.7-lightning": { "input_cost_per_token": 2.5e-06, "output_cost_per_token": 1.25e-05, "cache_read_input_token_cost": 1e-06, @@ -48956,7 +48966,7 @@ "mode": "chat", "supports_function_calling": true, "supports_prompt_caching": true, - "source": "https://docs.devin.ai/windsurf/plugins/cascade/models" + "source": "https://docs.devin.ai/desktop/models" }, "pinstripes/ps/glm-4.5-air": { "max_tokens": 128000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index efbe0d2ebb7..c5579f03a3e 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -48949,6 +48949,16 @@ "source": "https://docs.devin.ai/windsurf/plugins/cascade/models" }, "cognition/swe-1.7": { + "input_cost_per_token": 5e-07, + "output_cost_per_token": 2.5e-06, + "cache_read_input_token_cost": 2e-07, + "litellm_provider": "cognition", + "mode": "chat", + "supports_function_calling": true, + "supports_prompt_caching": true, + "source": "https://docs.devin.ai/desktop/models" + }, + "cognition/swe-1.7-lightning": { "input_cost_per_token": 2.5e-06, "output_cost_per_token": 1.25e-05, "cache_read_input_token_cost": 1e-06, @@ -48956,7 +48966,7 @@ "mode": "chat", "supports_function_calling": true, "supports_prompt_caching": true, - "source": "https://docs.devin.ai/windsurf/plugins/cascade/models" + "source": "https://docs.devin.ai/desktop/models" }, "pinstripes/ps/glm-4.5-air": { "max_tokens": 128000, diff --git a/tests/test_litellm/llms/openai_like/test_cognition_provider.py b/tests/test_litellm/llms/openai_like/test_cognition_provider.py index c358d178f60..5c71b60e08a 100644 --- a/tests/test_litellm/llms/openai_like/test_cognition_provider.py +++ b/tests/test_litellm/llms/openai_like/test_cognition_provider.py @@ -111,33 +111,51 @@ class TestCognitionProviderIdentity: class TestCognitionCostTracking: @pytest.mark.parametrize( - "model, input_cost, output_cost", + "model, input_cost, output_cost, cache_read_cost", [ - ("cognition/swe-1.6", 5e-07, 2.5e-06), - ("cognition/swe-1.7", 2.5e-06, 1.25e-05), + ("cognition/swe-1.6", 5e-07, 2.5e-06, 2e-07), + ("cognition/swe-1.7", 5e-07, 2.5e-06, 2e-07), + ("cognition/swe-1.7-lightning", 2.5e-06, 1.25e-05, 1e-06), ], ) - def test_cost_map_entries(self, model: str, input_cost: float, output_cost: float): + def test_cost_map_entries(self, model: str, input_cost: float, output_cost: float, cache_read_cost: float): info = litellm.get_model_info(model=model) assert info["litellm_provider"] == "cognition" assert info["mode"] == "chat" assert info["input_cost_per_token"] == input_cost assert info["output_cost_per_token"] == output_cost + assert info["cache_read_input_token_cost"] == cache_read_cost - def test_cost_differs_from_openai_pricing(self): + @pytest.mark.parametrize( + "model, expected_prompt_cost, expected_completion_cost", + [ + ("cognition/swe-1.7", 0.5, 2.5), + ("cognition/swe-1.7-lightning", 2.5, 12.5), + ], + ) + def test_cost_differs_from_openai_pricing( + self, model: str, expected_prompt_cost: float, expected_completion_cost: float + ): """A cognition-prefixed model must never be priced off an OpenAI cost entry.""" from litellm.cost_calculator import cost_per_token prompt_cost, completion_cost = cost_per_token( - model="cognition/swe-1.7", + model=model, prompt_tokens=1_000_000, completion_tokens=1_000_000, custom_llm_provider="cognition", ) - assert prompt_cost == pytest.approx(2.5) - assert completion_cost == pytest.approx(12.5) + assert prompt_cost == pytest.approx(expected_prompt_cost) + assert completion_cost == pytest.approx(expected_completion_cost) + + def test_lightning_is_five_times_the_standard_tier(self): + standard = litellm.get_model_info(model="cognition/swe-1.7") + lightning = litellm.get_model_info(model="cognition/swe-1.7-lightning") + + assert lightning["input_cost_per_token"] == pytest.approx(standard["input_cost_per_token"] * 5) + assert lightning["output_cost_per_token"] == pytest.approx(standard["output_cost_per_token"] * 5) def test_supported_endpoints_matrix(self): matrix = json.loads((Path(litellm.__file__).parent / "provider_endpoints_support_backup.json").read_text()) @@ -170,6 +188,30 @@ class TestCognitionRouting: mock_response="hello from swe", ) + usage = response.usage + expected = usage.prompt_tokens * 5e-07 + usage.completion_tokens * 2.5e-06 + assert response._hidden_params["response_cost"] == pytest.approx(expected) + + @pytest.mark.asyncio + async def test_router_spend_uses_the_lightning_entry_for_lightning(self): + """The Lightning tier is its own model, costed off its own entry.""" + from litellm import Router + + router = Router( + model_list=[ + { + "model_name": "swe-lightning", + "litellm_params": {"model": "cognition/swe-1.7-lightning", "api_key": "sk-test"}, + } + ] + ) + + response = await router.acompletion( + model="swe-lightning", + messages=[{"role": "user", "content": "hi"}], + mock_response="hello from swe lightning", + ) + usage = response.usage expected = usage.prompt_tokens * 2.5e-06 + usage.completion_tokens * 1.25e-05 assert response._hidden_params["response_cost"] == pytest.approx(expected) From 16cd08054fddd9effd266f07b251db82c909ad9f Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 20 Aug 2026 19:20:47 -0700 Subject: [PATCH 038/166] fix: populate team member emails missing from the roster snapshot `members_with_roles` is a denormalized JSON snapshot written at add-time. `_update_team_members_list` backfilled `user_id` from `user_email` but never the reverse, so a member added by `user_id` alone was stored with `user_email=None` permanently - and `/team/info` returns that blob verbatim with no join to `LiteLLM_UserTable`, so the Admin UI's member table renders "-" for a user that plainly has an email. Fix both ends: - write path: `_resolve_member_identity` resolves identity both ways off the user rows the add just touched, so new roster entries stop being born blank. - read path: `/team/info` fills blank emails from `LiteLLM_UserTable` in one indexed `user_id IN (...)` query, repairing rows already in the database. Members that already carry an email are passed through untouched and cost no query, so this only ever turns a null into the right value. --- .../management_endpoints/team_endpoints.py | 136 ++++++++---- .../test_team_endpoints.py | 194 ++++++++++++++++++ 2 files changed, 291 insertions(+), 39 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 82e22bb5bbf..a8e545a8551 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -2640,51 +2640,63 @@ async def _process_team_members( return updated_users, updated_team_memberships +def _resolve_member_identity(member: Member, updated_users: Sequence[LiteLLM_UserTable]) -> Member: + """Return ``member`` with whichever of ``user_id`` / ``user_email`` the caller left out filled in. + + The roster entry is a snapshot, so whatever is missing here is missing for good. + Resolution runs both ways off the user rows the add just touched: added by email + -> stamp the user_id, added by user_id -> stamp the email. A value the caller + supplied is never overwritten. + """ + resolved_user_id: Final = member.user_id or next( + ( + user.user_id + for user in updated_users + if member.user_email is not None and user.user_email == member.user_email + ), + None, + ) + resolved_user_email: Final = member.user_email or next( + ( + user.user_email + for user in updated_users + if resolved_user_id is not None and user.user_id == resolved_user_id and user.user_email is not None + ), + None, + ) + return member.model_copy( + update={ # mutable-ok: pydantic update payload + "user_id": resolved_user_id, + "user_email": resolved_user_email, + } + ) + + +def _member_already_in_team(member: Member, complete_team_data: LiteLLM_TeamTable) -> bool: + return any( + (member.user_id is not None and existing_member.user_id == member.user_id) + or (member.user_email is not None and existing_member.user_email == member.user_email) + for existing_member in complete_team_data.members_with_roles + ) + + async def _update_team_members_list( data: TeamMemberAddRequest, complete_team_data: LiteLLM_TeamTable, updated_users: list[LiteLLM_UserTable], ) -> None: """Update the team's members_with_roles list.""" - if isinstance(data.member, Member): - new_member: Final = data.member.model_copy() + requested_members: Final[Sequence[Member]] = ( + (data.member,) if isinstance(data.member, Member) else tuple(data.member) + ) + resolved_members: Final = tuple(_resolve_member_identity(m, updated_users) for m in requested_members) - # get user id - if new_member.user_id is None and new_member.user_email is not None: - for user in updated_users: - if user.user_email is not None and user.user_email == new_member.user_email: - new_member.user_id = user.user_id - - # Check if member already exists in team before adding - member_already_exists = False - for existing_member in complete_team_data.members_with_roles: - if (new_member.user_id is not None and existing_member.user_id == new_member.user_id) or ( - new_member.user_email is not None and existing_member.user_email == new_member.user_email - ): - member_already_exists = True - break - - if not member_already_exists: - complete_team_data.members_with_roles.append(new_member) - - elif isinstance(data.member, list): - for nm in data.member: - if nm.user_id is None and nm.user_email is not None: - for user in updated_users: - if user.user_email is not None and user.user_email == nm.user_email: - nm.user_id = user.user_id - - # Check if member already exists in team before adding - member_already_exists = False - for existing_member in complete_team_data.members_with_roles: - if (nm.user_id is not None and existing_member.user_id == nm.user_id) or ( - nm.user_email is not None and existing_member.user_email == nm.user_email - ): - member_already_exists = True - break - - if not member_already_exists: - complete_team_data.members_with_roles.append(nm) + # extend() consumes the generator as it appends, so a member already added by this + # same call is seen by the next _member_already_in_team check - the batch dedupes + # against itself exactly as the append-one-at-a-time loop this replaced did. + complete_team_data.members_with_roles.extend( # rebind-ok: this helper's contract is to grow the caller's roster in place + m for m in resolved_members if not _member_already_in_team(m, complete_team_data) + ) async def _add_team_members_to_team( @@ -4086,6 +4098,39 @@ async def _add_team_member_budget_table( return team_info_response_object +async def _hydrate_member_emails( + prisma_client: PrismaClient, + members: Sequence[Member], +) -> tuple[Member, ...]: + """Fill in ``user_email`` for roster entries that were stored without one. + + ``members_with_roles`` is a denormalized snapshot written at add-time, so an entry + stored with ``user_email=None`` keeps that null even once the user row has an email. + Look the missing ones up in ``LiteLLM_UserTable`` (one indexed query) and fill them + in. A stored email is never overwritten - the snapshot stays the source of truth + wherever it has a value. + """ + missing_user_ids: Final = frozenset(m.user_id for m in members if not m.user_email and m.user_id is not None) + if not missing_user_ids: + return tuple(members) + + user_rows: Final[Sequence[LiteLLM_UserTable]] = await _user_db(prisma_client).find_many( + where={ # mutable-ok: Prisma query filters are dict-shaped + "user_id": { # mutable-ok: Prisma query filters are dict-shaped + "in": sorted(missing_user_ids) + } + } + ) + email_by_user_id: Final = MappingProxyType({u.user_id: u.user_email for u in user_rows if u.user_email}) + + return tuple( + m.model_copy(update={"user_email": email_by_user_id[m.user_id]}) # mutable-ok: pydantic update payload + if not m.user_email and m.user_id in email_by_user_id + else m + for m in members + ) + + async def _resolve_team_access_group_resources( _team_info: TeamInfoResponseObjectTeamTable, ) -> TeamInfoResponseObjectTeamTable: @@ -4221,9 +4266,22 @@ async def team_info( # Resolve resources inherited from access groups resolved_team_info: Final = await _resolve_team_access_group_resources(_team_info) + # Fill in emails the add-time roster snapshot never captured + hydrated_members: Final = await _hydrate_member_emails( + prisma_client=prisma_client, + members=resolved_team_info.members_with_roles, + ) + hydrated_team_info: Final = resolved_team_info.model_copy( + update={ # mutable-ok: pydantic update payload + # list(), not the tuple: model_copy skips validation, so the field has + # to be handed the list[Member] the response model declares. + "members_with_roles": list(hydrated_members) # mutable-ok: declared list[Member] + } + ) + response_object: Final = TeamInfoResponseObject( team_id=team_id, - team_info=resolved_team_info, + team_info=hydrated_team_info, keys=keys, team_memberships=returned_tm, ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 0767288d0bc..e39b09ae073 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -10253,6 +10253,65 @@ async def test_team_info_returns_model_aliases(): assert litellm_model_table.model_aliases == {"gpt-4o": "gpt-4o-team-1"} +@pytest.mark.asyncio +async def test_team_info_hydrates_member_emails_from_the_user_table(): + """/team/info must fill in emails missing from the members_with_roles snapshot. + + members_with_roles is written at add-time, so a member added by user_id alone + carries user_email=None forever. Without this join the Admin UI's member table + shows "-" for a user that has an email on their user row. A stored email is left + exactly as-is. + """ + from fastapi import Request + + from litellm.proxy.management_endpoints import team_endpoints + + team_row = LiteLLM_TeamTable( + team_id="team-1", + members_with_roles=[ + Member(user_id="no-email-on-roster", role="admin"), + Member(user_id="already-stored", user_email="stored@example.com", role="user"), + ], + ) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + mock_prisma.get_data = AsyncMock(return_value=[]) + + find_many = AsyncMock( + return_value=[ + LiteLLM_UserTable( + user_id="no-email-on-roster", + user_email="real@example.com", + max_budget=None, + spend=0.0, + models=[], + ) + ] + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch.object(team_endpoints, "get_all_team_memberships", AsyncMock(return_value=[])), + patch.object(team_endpoints, "UserRepository") as repo, + ): + repo.return_value.table.find_many = find_many + + response = await team_endpoints.team_info( + http_request=MagicMock(spec=Request), + team_id="team-1", + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + members = response["team_info"].members_with_roles + assert [(m.user_id, m.user_email) for m in members] == [ + ("no-email-on-roster", "real@example.com"), + ("already-stored", "stored@example.com"), + ] + # only the member actually missing an email is looked up + assert find_many.await_args.kwargs["where"] == {"user_id": {"in": ["no-email-on-roster"]}} + + @pytest.mark.asyncio async def test_update_model_table_clears_aliases_with_empty_map(): """``model_aliases={}`` on /team/update must persist an empty map (json.dumps({})) @@ -11369,6 +11428,141 @@ async def test_resolve_existing_member_user_ids_skips_the_query_when_no_user_ids repo.return_value.table.find_many.assert_not_awaited() +def _user_row(user_id: str, user_email: str | None) -> LiteLLM_UserTable: + return LiteLLM_UserTable( + user_id=user_id, user_email=user_email, max_budget=None, spend=0.0, models=[] + ) + + +@pytest.mark.asyncio +async def test_hydrate_member_emails_fills_in_emails_the_roster_snapshot_never_captured(): + """A member added by user_id alone has user_email=None on the stored roster entry. + + /team/info has to fill it in from the user row, or the UI renders "-" for a user + that plainly has an email. + """ + from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_emails + + find_many = AsyncMock(return_value=[_user_row("by-id", "found@example.com")]) + + with patch("litellm.proxy.management_endpoints.team_endpoints.UserRepository") as repo: + repo.return_value.table.find_many = find_many + + hydrated = await _hydrate_member_emails( + prisma_client=MagicMock(), + members=[Member(user_id="by-id", role="admin")], + ) + + assert [(m.user_id, m.user_email, m.role) for m in hydrated] == [("by-id", "found@example.com", "admin")] + find_many.assert_awaited_once() + assert find_many.await_args.kwargs["where"] == {"user_id": {"in": ["by-id"]}} + + +@pytest.mark.asyncio +async def test_hydrate_member_emails_never_overwrites_a_stored_email(): + """The snapshot wins wherever it has a value - hydration only fills blanks. + + Overwriting would be a real behavior change to /team/info; filling a null is not. + """ + from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_emails + + find_many = AsyncMock(return_value=[_user_row("has-email", "current@example.com")]) + + with patch("litellm.proxy.management_endpoints.team_endpoints.UserRepository") as repo: + repo.return_value.table.find_many = find_many + + hydrated = await _hydrate_member_emails( + prisma_client=MagicMock(), + members=[Member(user_id="has-email", user_email="stored@example.com", role="user")], + ) + + assert hydrated[0].user_email == "stored@example.com" + # nothing was missing, so no round-trip either + find_many.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_hydrate_member_emails_leaves_members_alone_when_the_user_row_has_no_email(): + """A user row with no email leaves the member as-is rather than inventing one.""" + from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_emails + + with patch("litellm.proxy.management_endpoints.team_endpoints.UserRepository") as repo: + repo.return_value.table.find_many = AsyncMock(return_value=[_user_row("no-email", None)]) + + hydrated = await _hydrate_member_emails( + prisma_client=MagicMock(), + members=[Member(user_id="no-email", role="user"), Member(user_email="e@example.com", role="user")], + ) + + assert [m.user_email for m in hydrated] == [None, "e@example.com"] + + +@pytest.mark.asyncio +async def test_hydrate_member_emails_skips_the_query_when_every_member_has_one(): + """No blanks means /team/info pays for no extra query.""" + from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_emails + + with patch("litellm.proxy.management_endpoints.team_endpoints.UserRepository") as repo: + repo.return_value.table.find_many = AsyncMock() + + hydrated = await _hydrate_member_emails( + prisma_client=MagicMock(), + members=[Member(user_id="a", user_email="a@example.com", role="user")], + ) + + assert hydrated[0].user_email == "a@example.com" + repo.return_value.table.find_many.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_update_team_members_list_stamps_email_for_a_member_added_by_user_id(): + """Identity resolution runs both ways, so new roster entries stop being born blank. + + Previously only user_id was backfilled (from email); a member added by user_id + was written with user_email=None forever. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + _update_team_members_list, + ) + + mock_team = MagicMock(spec=LiteLLM_TeamTable) + mock_team.members_with_roles = [] + + await _update_team_members_list( + data=TeamMemberAddRequest(team_id="test-team-123", member=Member(user_id="new-user-123", role="user")), + complete_team_data=mock_team, + updated_users=[_user_row("new-user-123", "new@example.com")], + ) + + assert len(mock_team.members_with_roles) == 1 + assert mock_team.members_with_roles[0].user_email == "new@example.com" + + +@pytest.mark.asyncio +async def test_update_team_members_list_stamps_email_for_each_member_in_a_bulk_add(): + """Same both-ways resolution for the list branch.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _update_team_members_list, + ) + + mock_team = MagicMock(spec=LiteLLM_TeamTable) + mock_team.members_with_roles = [] + + await _update_team_members_list( + data=TeamMemberAddRequest( + team_id="test-team-123", + member=[Member(user_id="u1", role="user"), Member(user_email="u2@example.com", role="admin")], + ), + complete_team_data=mock_team, + updated_users=[_user_row("u1", "u1@example.com"), _user_row("u2", "u2@example.com")], + ) + + assert [(m.user_id, m.user_email) for m in mock_team.members_with_roles] == [ + ("u1", "u1@example.com"), + ("u2", "u2@example.com"), + ] + + def test_pre_existing_user_ids_counts_ids_filled_in_by_member_resolution(): """An id the member-resolution step filled in came from a matched row, so it pre-existed. From 722c650bfdf5375fc9e2d444485156b97c085762 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 19:47:05 -0700 Subject: [PATCH 039/166] test: cover generic HTTP streaming provider header forwarding Add sync and async regression tests for the BaseLLMHTTPHandler streaming path, which forwards provider response headers for the ~30 providers that ride the generic handler and had no coverage. Also drop redundant setup prose from the moonshot invoke test docstring. --- .../llm_translation/test_bedrock_moonshot.py | 4 - .../custom_httpx/test_llm_http_handler.py | 76 +++++++++++++++++++ 2 files changed, 76 insertions(+), 4 deletions(-) diff --git a/tests/llm_translation/test_bedrock_moonshot.py b/tests/llm_translation/test_bedrock_moonshot.py index 61364cf2caa..a82d1c6f029 100644 --- a/tests/llm_translation/test_bedrock_moonshot.py +++ b/tests/llm_translation/test_bedrock_moonshot.py @@ -209,10 +209,6 @@ class TestBedrockMoonshotInvoke(BaseLLMChatTest): endpoint with the messages body. Iteration of the stream itself is not exercised here — moonshot streaming delegates to the OpenAI parser and is covered by the OpenAI test suite. - - Patch ``make_sync_call`` at its import site in - ``base_invoke_transformation`` so we observe the exact kwargs it is - called with at stream-wrapper construction time. """ from litellm.utils import CustomStreamWrapper diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index fddd8d09dfc..b78b313e05c 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -2071,3 +2071,79 @@ async def test_anthropic_invalid_thinking_signature_retry_resigns_bedrock_reques retry_authorization = posts[1]["headers"]["Authorization"] assert retry_authorization.startswith("AWS4-HMAC-SHA256") assert retry_authorization != first_attempt_headers["Authorization"] + + +_GENERIC_STREAM_SSE = ( + b'data: {"id":"chatcmpl-1","object":"chat.completion.chunk","created":1,' + b'"model":"test-model","choices":[{"index":0,"delta":{"content":"hi"},' + b'"finish_reason":null}]}\n\n' + b"data: [DONE]\n\n" +) + + +def _generic_stream_upstream_response() -> httpx.Response: + return httpx.Response( + 200, + headers={ + "x-request-id": "generic-req-123", + "x-ratelimit-remaining-requests": "42", + }, + content=_GENERIC_STREAM_SSE, + request=httpx.Request("POST", "https://fake-vllm.test/v1/chat/completions"), + ) + + +def test_generic_http_handler_sync_streaming_forwards_provider_response_headers(): + """ + Regression test for the generic BaseLLMHTTPHandler streaming path used by + ~30 providers (deepseek, groq, hosted_vllm, databricks, openrouter, ...). + + The sync `completion()` streaming branch builds the CustomStreamWrapper from + `make_sync_call`, which returns the upstream response headers alongside the + stream. Those headers must reach the caller as `llm_provider-*` entries in + `_hidden_params["additional_headers"]`, which is what the proxy merges into + the client-facing response headers. + """ + mock_client = Mock(spec=HTTPHandler) + mock_client.post = Mock(return_value=_generic_stream_upstream_response()) + + response = litellm.completion( + model="hosted_vllm/test-model", + messages=[{"role": "user", "content": "Hello"}], + api_base="https://fake-vllm.test/v1", + api_key="sk-test", + stream=True, + client=mock_client, + ) + + additional_headers = response._hidden_params["additional_headers"] + assert additional_headers["llm_provider-x-request-id"] == "generic-req-123" + assert additional_headers["llm_provider-x-ratelimit-remaining-requests"] == "42" + + assert "".join([chunk.choices[0].delta.content or "" for chunk in response]) == "hi" + + +@pytest.mark.asyncio +async def test_generic_http_handler_async_streaming_forwards_provider_response_headers(): + """ + Companion to the sync test above for `acompletion_stream_function`, which + builds its CustomStreamWrapper from `make_async_call_stream_helper`. + """ + mock_client = AsyncMock(spec=AsyncHTTPHandler) + mock_client.post = AsyncMock(return_value=_generic_stream_upstream_response()) + + response = await litellm.acompletion( + model="hosted_vllm/test-model", + messages=[{"role": "user", "content": "Hello"}], + api_base="https://fake-vllm.test/v1", + api_key="sk-test", + stream=True, + client=mock_client, + ) + + additional_headers = response._hidden_params["additional_headers"] + assert additional_headers["llm_provider-x-request-id"] == "generic-req-123" + assert additional_headers["llm_provider-x-ratelimit-remaining-requests"] == "42" + + collected = [chunk async for chunk in response] + assert "".join([chunk.choices[0].delta.content or "" for chunk in collected]) == "hi" From 2ea633d223e57777edf39edb2c3ca6b18c0266d2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 19:57:46 -0700 Subject: [PATCH 040/166] fix(sagemaker_chat): send the inference component header and honor hf_model_name sagemaker_chat never put X-Amzn-SageMaker-Inference-Component on the request, so any endpoint backed by inference components answered 400 INFERENCE_COMPONENT_NAME_MISSING and the call never reached the container. The legacy sagemaker provider has built that header from model_id since #8889, and this brings the chat provider in line. It goes on in validate_environment, which runs before the request is SigV4-signed, so the signature covers it The request body also always named the endpoint rather than the served model, which containers that validate the body's model answer with a 404. hf_model_name now becomes the body's model when it is set, and endpoints that do not set it keep sending exactly what they send today --- litellm/llms/sagemaker/chat/transformation.py | 25 +++++- .../test_sagemaker_chat_transformation.py | 84 +++++++++++++++++++ 2 files changed, 108 insertions(+), 1 deletion(-) diff --git a/litellm/llms/sagemaker/chat/transformation.py b/litellm/llms/sagemaker/chat/transformation.py index 99543e7add1..37ddd813d6f 100644 --- a/litellm/llms/sagemaker/chat/transformation.py +++ b/litellm/llms/sagemaker/chat/transformation.py @@ -54,7 +54,30 @@ class SagemakerChatConfig(OpenAIGPTConfig, BaseAWSLLM): api_key: str | None = None, api_base: str | None = None, ) -> dict: - return headers + inference_component_name: Final = optional_params.get("model_id") + if not isinstance(inference_component_name, str): + return headers + return {**headers, "X-Amzn-SageMaker-Inference-Component": inference_component_name} + + def transform_request( + self, + model: str, + messages: list[AllMessageValues], # mutable-ok: matches the base chat transform signature + optional_params: dict, # mutable-ok: matches the base chat transform signature + litellm_params: dict, # mutable-ok: matches the base chat transform signature + headers: dict, # mutable-ok: matches the base chat transform signature + ) -> dict: # mutable-ok: the handler sends this body straight to httpx + request: Final = super().transform_request( + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + headers=headers, + ) + served_model_name: Final = litellm_params.get("hf_model_name") + if not isinstance(served_model_name, str): + return request + return {**request, "model": served_model_name} def get_complete_url( self, diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py b/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py index 54e6f95c795..da6caca4f05 100644 --- a/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py +++ b/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py @@ -19,6 +19,8 @@ from unittest.mock import MagicMock import httpx import pytest +import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.sagemaker.chat.transformation import SagemakerChatConfig @@ -233,3 +235,85 @@ def test_decoder_reassembles_frames_across_arbitrary_byte_boundaries(split_size) ] assert texts == [f"token{i} " for i in range(len(frames))] + + +_INFERENCE_COMPONENT_HEADER = "X-Amzn-SageMaker-Inference-Component" + +_STUB_COMPLETION_RESPONSE = { + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1700000000, + "model": "served-model", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} + + +class _RequestCapturingHTTPHandler(HTTPHandler): + """Injected transport that records exactly what sagemaker_chat put on the wire.""" + + def __init__(self) -> None: + super().__init__() + self.request_headers: dict[str, str] = {} + self.request_body: dict = {} + + def post(self, url: str, headers=None, data=None, **kwargs) -> httpx.Response: + self.request_headers = dict(headers or {}) + self.request_body = json.loads(data) + return httpx.Response(200, json=_STUB_COMPLETION_RESPONSE, request=httpx.Request("POST", url)) + + +def _invoke_sagemaker_chat(monkeypatch, **extra_params) -> _RequestCapturingHTTPHandler: + """Drive one sagemaker_chat completion against an injected transport. + + A Bedrock API key short-circuits SigV4 inside `BaseAWSLLM._sign_request`, which would hide + whether the inference-component header is really covered by the signature, so it is cleared. + """ + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + client = _RequestCapturingHTTPHandler() + litellm.completion( + model="sagemaker_chat/my-endpoint", + messages=[{"role": "user", "content": "hi"}], + aws_access_key_id="AKIATESTTESTTESTTEST", + aws_secret_access_key="test-secret-key", + aws_region_name="us-east-1", + client=client, + **extra_params, + ) + return client + + +def test_model_id_is_sent_as_a_signed_inference_component_header(monkeypatch): + """`model_id` names an inference component and must reach SageMaker as a signed header. + + Endpoints backed by inference components reject any request without + `X-Amzn-SageMaker-Inference-Component` with HTTP 400 INFERENCE_COMPONENT_NAME_MISSING, so the + header has to be built before `sign_request` runs and end up inside SignedHeaders. + """ + client = _invoke_sagemaker_chat(monkeypatch, model_id="my-inference-component") + + assert client.request_headers[_INFERENCE_COMPONENT_HEADER] == "my-inference-component" + assert "x-amzn-sagemaker-inference-component" in client.request_headers["Authorization"] + + +def test_no_inference_component_header_when_model_id_is_unset(monkeypatch): + """Plain endpoints must not receive the header at all, not even an empty one.""" + client = _invoke_sagemaker_chat(monkeypatch) + + assert not any(name.lower() == _INFERENCE_COMPONENT_HEADER.lower() for name in client.request_headers) + + +def test_hf_model_name_becomes_the_body_model(monkeypatch): + """`hf_model_name` names the served model, and containers that validate the body's `model` + 404 on the endpoint name, so it has to replace it rather than ride along as an extra field.""" + client = _invoke_sagemaker_chat(monkeypatch, hf_model_name="org/served-model") + + assert client.request_body["model"] == "org/served-model" + assert "hf_model_name" not in client.request_body + + +def test_body_model_stays_the_endpoint_name_when_hf_model_name_is_unset(monkeypatch): + """Without `hf_model_name` the body must keep the model it has today.""" + client = _invoke_sagemaker_chat(monkeypatch) + + assert client.request_body["model"] == "my-endpoint" From b76def0e5df1ce874c131fdbd4b4e10def90f3df Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 20 Aug 2026 20:24:49 -0700 Subject: [PATCH 041/166] test: require a `match=` on broad pytest.raises, and drop duplicate parametrize cases (#37769) `pytest.raises(Exception)` with no `match=` passes on any error that broad. A TypeError from a refactor, a botched fixture, an import that moved: all of them read as the rejection the test claims to police, so the test goes green for the wrong reason and stays green after the behaviour it guards is gone. PT011 closes that gap for the 317 sites B017 could not reach, because B017 only fires on a single-statement body with no `as e` binding. Each pattern here is the message the code actually raised, recorded by running the sites under a plugin that logged the concrete type and text per call site, so the assertions describe observed behaviour rather than a guess. Where a site raises more than one message across its parametrize cases, the pattern is an alternation of what was seen; where the exception carries an empty `str()` and puts the text on `.message`, the site keeps a narrow `noqa` with the reason. PT014 removes four parametrize cases that were listed twice. The duplicate re-runs an assertion that already passed, and it usually marks a case someone meant to vary and forgot to edit. --- ruff-tests.toml | 8 +- .../integrations/test_prometheus.py | 8 +- .../test_bedrock_apply_guardrail.py | 2 +- .../proxy/hooks/test_managed_files.py | 8 +- .../test_project_endpoints_prisma.py | 10 +- .../test_dynamoai_guardrails.py | 2 +- .../test_eu_ai_act_article5.py | 2 +- .../test_eu_ai_act_french_3_scenarios.py | 8 +- .../test_sg_mas_ai_guardrails.py | 2 +- .../test_sg_pdpa_guardrails.py | 2 +- tests/litellm_utils_tests/test_hashicorp.py | 2 +- tests/litellm_utils_tests/test_utils.py | 2 +- .../test_validate_tool_choice.py | 10 +- .../test_responses_hooks.py | 2 +- .../test_bedrock_completion.py | 2 +- .../test_convert_dict_to_chat_completion.py | 6 +- tests/llm_translation/test_prompt_factory.py | 6 +- tests/llm_translation/test_triton.py | 2 +- .../test_unit_test_bedrock_invoke.py | 2 +- tests/local_testing/test_auth_utils.py | 4 - tests/local_testing/test_exceptions.py | 2 +- tests/local_testing/test_file_types.py | 4 +- tests/local_testing/test_get_model_info.py | 1 - .../test_router_budget_limiter.py | 6 +- tests/local_testing/test_router_fallbacks.py | 2 +- .../test_standard_logging_payload.py | 2 +- .../test_update_team_e2e.py | 6 +- .../test_ocr_azure_document_intelligence.py | 2 +- tests/otel_tests/test_e2e_model_access.py | 9 +- .../test_key_management.py | 2 +- .../test_role_based_access.py | 4 +- tests/proxy_unit_tests/test_auth_checks.py | 4 +- tests/proxy_unit_tests/test_jwt.py | 10 +- tests/proxy_unit_tests/test_proxy_utils.py | 4 +- tests/proxy_unit_tests/test_update_spend.py | 2 +- .../test_router_helper_utils.py | 4 +- .../test_mcp_servers.py | 2 +- .../send_emails/test_sendgrid_email.py | 2 +- .../proxy/test_managed_files_hook.py | 2 +- .../cloudzero/test_cloudzero_database.py | 2 +- .../cloudzero/test_cz_stream_api.py | 2 +- .../integrations/focus/test_focus_database.py | 2 +- .../integrations/focus/test_s3_destination.py | 2 +- .../integrations/gitlab/test_gitlab_client.py | 10 +- .../integrations/levo/test_levo.py | 2 +- .../integrations/otel/test_otel_v2_metrics.py | 2 +- .../integrations/test_langfuse.py | 2 +- ...llm_core_utils_prompt_templates_factory.py | 2 +- ...test_initialize_dynamic_callback_params.py | 4 +- .../litellm_core_utils/test_llm_judge.py | 2 +- .../test_streaming_handler.py | 4 +- .../litellm_core_utils/test_token_counter.py | 20 +- .../litellm_core_utils/test_url_utils.py | 4 +- ...st_aiml_image_generation_transformation.py | 2 +- .../messages/test_mcp_handler.py | 4 +- .../test_azure_ai_rerank_transformation.py | 4 +- .../llms/bedrock/test_base_aws_llm.py | 4 +- ...bedrock_mantle_responses_transformation.py | 6 +- .../test_bedrock_mantle_transformation.py | 4 +- .../chat/test_bytez_chat_transformation.py | 2 +- .../test_deepinfra_rerank_transformation.py | 8 +- .../test_fal_ai_nano_banana_transformation.py | 2 +- .../test_featherless_chat_transformation.py | 6 +- ...test_fireworks_ai_rerank_transformation.py | 2 +- .../test_gemini_image_edit_transformation.py | 2 +- .../llms/gemini/test_gemini_client_setup.py | 4 +- .../test_hosted_vllm_rerank_transformation.py | 2 +- .../chat/test_langflow_chat_transformation.py | 2 +- ...est_modelscope_image_gen_transformation.py | 6 +- .../chat/test_novita_chat_transformation.py | 2 +- .../oci/chat/test_oci_chat_transformation.py | 16 +- .../test_pg_vector_transformation.py | 4 +- .../test_recraft_image_edit_transformation.py | 2 +- .../test_recraft_image_gen_transformation.py | 6 +- .../test_stability_image_generation.py | 6 +- .../llms/tinyfish/test_tinyfish_search.py | 12 +- .../files/test_vertex_ai_files_integration.py | 2 +- .../vertex_ai/test_vertex_ai_common_utils.py | 2 +- ...est_volcengine_responses_transformation.py | 2 +- .../volcengine/test_volcengine_embedding.py | 2 +- .../test_voyage_rerank_transformation.py | 4 +- .../test_voyage_multimodal_embedding.py | 4 +- .../llms/xai/test_xai_key_fallback.py | 2 +- .../auth/test_user_api_key_auth_mcp.py | 2 +- .../mcp_server/test_discoverable_endpoints.py | 7 - .../mcp_server/test_short_mcp_tool_prefix.py | 2 +- .../proxy/auth/test_auth_checks.py | 4 +- .../test_auth_hot_path_network_requests.py | 1042 ++++++++--------- .../proxy/auth/test_auth_utils.py | 18 +- .../proxy/auth/test_handle_jwt.py | 18 +- .../proxy/auth/test_oauth2_proxy_hook.py | 2 +- .../proxy/auth/test_route_checks.py | 14 +- .../proxy/auth/test_user_api_key_auth.py | 4 +- .../proxy/client/cli/test_auth_commands.py | 10 +- .../proxy/client/cli/test_pkce_login.py | 2 +- .../test_litellm/proxy/client/test_models.py | 4 +- .../proxy/common_utils/test_callback_utils.py | 6 +- .../proxy/common_utils/test_path_utils.py | 2 +- .../proxy/common_utils/test_timezone_utils.py | 6 +- .../test_spend_logs_partition_manager.py | 4 +- .../db/test_prisma_planned_engine_restart.py | 2 +- .../proxy/db/test_spend_log_tool_index.py | 2 +- .../test_bedrock_invoke_guardrail_checks.py | 4 +- .../guardrail_hooks/test_enkryptai.py | 2 +- .../test_generic_guardrail_api.py | 6 +- .../guardrail_hooks/test_model_armor.py | 6 +- .../guardrail_hooks/test_panw_prisma_airs.py | 2 +- .../guardrail_hooks/test_straiker.py | 6 +- .../guardrail_hooks/test_tool_permission.py | 4 +- .../proxy/guardrails/test_llm_as_a_judge.py | 2 +- .../hooks/test_sensitive_data_routing.py | 2 +- .../proxy/hooks/test_tpm_concurrent.py | 28 +- .../test_key_management_endpoints.py | 2 +- .../test_mcp_management_endpoints.py | 2 +- .../test_model_management_endpoints.py | 12 +- .../test_ptu_model_settings.py | 18 +- .../usage_endpoints/test_ai_usage_chat.py | 2 +- .../test_team_metadata_validation.py | 2 +- .../test_batch_guardrails.py | 2 +- .../test_llm_pass_through_endpoints.py | 8 +- .../policy_engine/test_policy_versioning.py | 4 +- .../proxy/proxy_server/test_proxy_config.py | 4 +- .../test_spend_management_endpoints.py | 4 +- .../proxy/test_common_request_processing.py | 2 +- .../proxy/test_enforce_user_param.py | 6 +- .../proxy/test_litellm_pre_call_utils.py | 6 +- tests/test_litellm/proxy/test_proxy_server.py | 2 +- .../proxy/test_spend_log_cleanup.py | 4 +- .../test_litellm/proxy/test_team_org_move.py | 2 +- .../proxy/utils/helpers/test_team_configs.py | 2 +- .../test_proxy_update_spend.py | 2 +- .../test_callback_capabilities_class.py | 2 +- .../test_responses_websocket_all_providers.py | 2 +- .../adaptive_router/test_bandit.py | 2 +- .../test_router_tag_routing.py | 26 +- .../test_litellm/sandbox/test_e2b_sandbox.py | 2 +- .../test_base_secret_manager.py | 2 +- .../test_github_close_low_quality_prs.py | 2 +- .../test_github_triage_with_llm.py | 4 +- .../test_redact_string_in_error_paths.py | 2 +- tests/test_litellm/test_redis.py | 4 +- tests/test_litellm/test_router.py | 10 +- .../test_router_model_cost_isolation.py | 4 +- tests/test_litellm/test_utils.py | 2 +- tests/test_litellm/types/test_router.py | 2 +- 145 files changed, 844 insertions(+), 867 deletions(-) diff --git a/ruff-tests.toml b/ruff-tests.toml index 6e77f4792a7..60438d355f0 100644 --- a/ruff-tests.toml +++ b/ruff-tests.toml @@ -20,6 +20,12 @@ # PT012 a `pytest.raises` block that runs on past the raising call. Everything after # that call is dead, so an `assert` sitting there is never checked. Keep the # block to the call itself and put the assertions below it +# PT011 `pytest.raises(Exception)` / `(ValueError)` / `(OSError)` with no `match=`. The +# block passes on any error that broad, so the TypeError a refactor introduced +# reads as the rejection under test. Pin the message the code actually raises +# PT014 the same `parametrize` case listed twice. The copy re-runs an assertion that +# already passed and adds no coverage, and it usually marks a case someone meant +# to vary and forgot to edit # # No target-version here on purpose: it resolves from requires-python (>=3.10), so # 3.11-only builtins like BaseExceptionGroup are correctly flagged in a tree that @@ -27,4 +33,4 @@ line-length = 120 -lint.select = ["F821", "B011", "B015", "B017", "B018", "PT012", "PT015", "PLR0133", "PLW0127"] +lint.select = ["F821", "B011", "B015", "B017", "B018", "PT011", "PT012", "PT014", "PT015", "PLR0133", "PLW0127"] diff --git a/tests/enterprise/litellm_enterprise/integrations/test_prometheus.py b/tests/enterprise/litellm_enterprise/integrations/test_prometheus.py index 0cd6055e09d..bdf73b6ab03 100644 --- a/tests/enterprise/litellm_enterprise/integrations/test_prometheus.py +++ b/tests/enterprise/litellm_enterprise/integrations/test_prometheus.py @@ -400,7 +400,7 @@ def test_invalid_metric_name_validation(): litellm.prometheus_metrics_config = test_config # Creating PrometheusLogger should raise ValueError due to invalid metric - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Configuration validation failed') as exc_info: PrometheusLogger() # Verify error message contains information about invalid metric @@ -429,7 +429,7 @@ def test_invalid_labels_validation(): litellm.prometheus_metrics_config = test_config # Creating PrometheusLogger should raise ValueError due to invalid labels - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Configuration validation failed') as exc_info: PrometheusLogger() # Verify error message contains information about invalid labels @@ -598,7 +598,7 @@ def test_invalid_exclude_metric_name_raises(reset_prometheus_exclude_settings): litellm.prometheus_exclude_labels = None litellm.prometheus_exclude_metrics = ["not_a_real_metric"] - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Prometheus exclude configuration validation failed') as exc_info: PrometheusLogger() assert "not_a_real_metric" in str(exc_info.value) @@ -612,7 +612,7 @@ def test_invalid_exclude_label_name_raises(reset_prometheus_exclude_settings): litellm.prometheus_exclude_metrics = None litellm.prometheus_exclude_labels = ["not_a_real_label"] - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Prometheus exclude configuration validation failed') as exc_info: PrometheusLogger() assert "not_a_real_label" in str(exc_info.value) diff --git a/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py b/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py index f257b47404e..6b6b5d768dd 100644 --- a/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py +++ b/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py @@ -141,7 +141,7 @@ async def test_bedrock_apply_guardrail_api_failure(): mock_api_request.side_effect = Exception("API connection failed") # Test the apply_guardrail method should raise an exception - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Bedrock guardrail failed: API connection failed') as exc_info: await guardrail.apply_guardrail( inputs={"texts": ["This is a test message"]}, request_data={}, diff --git a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py index 714f3be6df9..2d845a445b5 100644 --- a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py +++ b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py @@ -1653,7 +1653,7 @@ async def test_afile_retrieve_raises_error_when_no_router_and_file_object_none() unified_file_id = "test-unified-file-id" - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='LiteLLM Managed File object with id=test-unified-file-id') as exc_info: await proxy_managed_files.afile_retrieve( file_id=unified_file_id, litellm_parent_otel_span=None, @@ -1719,7 +1719,7 @@ async def test_afile_retrieve_raises_error_for_non_managed_file(): # Mock get_unified_file_id to return None (file not found) proxy_managed_files.get_unified_file_id = AsyncMock(return_value=None) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='LiteLLM Managed File object with id=non-existent-file-id') as exc_info: await proxy_managed_files.afile_retrieve( file_id="non-existent-file-id", litellm_parent_otel_span=None, @@ -2027,7 +2027,7 @@ async def test_list_batches_from_managed_objects_table_provider_filter_raises_ex ) # Filtering by provider should raise Exception - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="Filtering by 'provider' is not supported when using managed") as exc_info: await proxy_managed_files.list_user_batches( user_api_key_dict=UserAPIKeyAuth(user_id="test-user"), limit=10, @@ -2053,7 +2053,7 @@ async def test_list_batches_from_managed_objects_table_target_model_name_filter_ ) # Filtering by provider should raise Exception - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="Filtering by 'target_model_names' is not supported when") as exc_info: await proxy_managed_files.list_user_batches( user_api_key_dict=UserAPIKeyAuth(user_id="test-user"), limit=10, diff --git a/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py b/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py index bd6637ffcac..ed6735a7126 100644 --- a/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py +++ b/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py @@ -448,7 +448,7 @@ def test_check_team_project_limits_models_not_in_team(): models=["gpt-5.5", "claude-3"], # claude-3 not in team ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="not in team's allowed models\\. Team allowed models") as exc_info: _check_team_project_limits(team_object=team, data=data) assert "claude-3" in str(exc_info.value.detail) @@ -476,7 +476,7 @@ def test_check_team_project_limits_budget_exceeds_team(): max_budget=150.0, # exceeds team's 100.0 ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Project max_budget') as exc_info: _check_team_project_limits(team_object=team, data=data) assert "exceeds team's max_budget" in str(exc_info.value.detail) @@ -551,7 +551,7 @@ def test_check_team_project_limits_tpm_exceeds_team(): tpm_limit=20000, # exceeds team's 10000 ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Project tpm_limit') as exc_info: _check_team_project_limits(team_object=team, data=data) assert "exceeds team's tpm_limit" in str(exc_info.value.detail) @@ -577,7 +577,7 @@ def test_check_team_project_limits_negative_budget(): max_budget=-10.0, ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='max_budget cannot be negative\\. Received') as exc_info: _check_team_project_limits(team_object=team, data=data) assert "cannot be negative" in str(exc_info.value.detail) @@ -604,7 +604,7 @@ def test_check_team_project_limits_soft_budget_gte_max(): soft_budget=100.0, # equal to max, should fail ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='must be strictly lower than max_budget') as exc_info: _check_team_project_limits(team_object=team, data=data) assert "must be strictly lower" in str(exc_info.value.detail) diff --git a/tests/guardrails_tests/test_dynamoai_guardrails.py b/tests/guardrails_tests/test_dynamoai_guardrails.py index 98f676a71d5..6f0ea00165b 100644 --- a/tests/guardrails_tests/test_dynamoai_guardrails.py +++ b/tests/guardrails_tests/test_dynamoai_guardrails.py @@ -61,7 +61,7 @@ async def test_dynamoai_blocks_content_with_block_action(): guardrail.should_run_guardrail = MagicMock(return_value=True) # Test that the guardrail raises ValueError for blocked content - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='violation\\(s\\) detected') as exc_info: await guardrail.async_pre_call_hook( data=request_data, user_api_key_dict=UserAPIKeyAuth(), diff --git a/tests/guardrails_tests/test_eu_ai_act_article5.py b/tests/guardrails_tests/test_eu_ai_act_article5.py index 0903e6c5416..f7384667481 100644 --- a/tests/guardrails_tests/test_eu_ai_act_article5.py +++ b/tests/guardrails_tests/test_eu_ai_act_article5.py @@ -211,7 +211,7 @@ class TestEUAIActArticle5ConditionalMatching: # Apply guardrail if expected == "BLOCK": # Should raise an exception or return modified response indicating block - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Content blocked: eu_ai_act_article') as exc_info: await content_filter_guardrail.apply_guardrail( inputs={"texts": [sentence]}, request_data=request_data, diff --git a/tests/guardrails_tests/test_eu_ai_act_french_3_scenarios.py b/tests/guardrails_tests/test_eu_ai_act_french_3_scenarios.py index 6b9774d9cde..221ca5aa6e6 100644 --- a/tests/guardrails_tests/test_eu_ai_act_french_3_scenarios.py +++ b/tests/guardrails_tests/test_eu_ai_act_french_3_scenarios.py @@ -83,7 +83,7 @@ class TestEUAIActFrench3Scenarios: print(f"{'='*70}\n") # Should raise an exception (blocked) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="prohibited_practices_fr conditional match 'concevoir \\+") as exc_info: await content_filter_guardrail.apply_guardrail( inputs={"texts": [sentence]}, request_data=request_data, @@ -123,7 +123,7 @@ class TestEUAIActFrench3Scenarios: print(f"{'='*70}\n") # Should raise an exception (blocked) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="prohibited_practices_fr conditional match 'créer \\+") as exc_info: await content_filter_guardrail.apply_guardrail( inputs={"texts": [sentence]}, request_data=request_data, @@ -194,7 +194,7 @@ class TestEUAIActFrench3Scenarios: print(f"{'='*70}\n") # Should raise an exception (blocked by conditional matching) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="prohibited_practices_fr conditional match 'développer \\+") as exc_info: await content_filter_guardrail.apply_guardrail( inputs={"texts": [sentence]}, request_data=request_data, @@ -278,7 +278,7 @@ class TestFrenchEdgeCases: request_data = {"messages": [{"role": "user", "content": sentence}]} # Should still block (no exception bypass) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="prohibited_practices_fr conditional match 'créer \\+ crédit") as exc_info: await content_filter_guardrail.apply_guardrail( inputs={"texts": [sentence]}, request_data=request_data, diff --git a/tests/guardrails_tests/test_sg_mas_ai_guardrails.py b/tests/guardrails_tests/test_sg_mas_ai_guardrails.py index 668ee704692..e587d666a79 100644 --- a/tests/guardrails_tests/test_sg_mas_ai_guardrails.py +++ b/tests/guardrails_tests/test_sg_mas_ai_guardrails.py @@ -55,7 +55,7 @@ def _make_guardrail(yaml_filename: str, category_name: str) -> ContentFilterGuar async def _expect_block(guardrail: ContentFilterGuardrail, sentence: str, reason: str): request_data = {"messages": [{"role": "user", "content": sentence}]} - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Content blocked: sg_mas_') as exc_info: await guardrail.apply_guardrail( inputs={"texts": [sentence]}, request_data=request_data, diff --git a/tests/guardrails_tests/test_sg_pdpa_guardrails.py b/tests/guardrails_tests/test_sg_pdpa_guardrails.py index fd7133bc745..42c3a15f9f6 100644 --- a/tests/guardrails_tests/test_sg_pdpa_guardrails.py +++ b/tests/guardrails_tests/test_sg_pdpa_guardrails.py @@ -62,7 +62,7 @@ def _make_guardrail(yaml_filename: str, category_name: str) -> ContentFilterGuar async def _expect_block(guardrail: ContentFilterGuardrail, sentence: str, reason: str): """Assert that the guardrail BLOCKS the sentence.""" request_data = {"messages": [{"role": "user", "content": sentence}]} - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Content blocked: sg_pdpa_') as exc_info: await guardrail.apply_guardrail( inputs={"texts": [sentence]}, request_data=request_data, diff --git a/tests/litellm_utils_tests/test_hashicorp.py b/tests/litellm_utils_tests/test_hashicorp.py index 9aff7ddc10e..fa39a045227 100644 --- a/tests/litellm_utils_tests/test_hashicorp.py +++ b/tests/litellm_utils_tests/test_hashicorp.py @@ -432,7 +432,7 @@ def test_hashicorp_get_url_rejects_path_traversal(monkeypatch, malicious_secret_ monkeypatch.setenv("HCP_VAULT_TOKEN", "test-token-for-get-url-only") manager = HashicorpSecretManager() - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Invalid secret_name'): manager.get_url(malicious_secret_name) diff --git a/tests/litellm_utils_tests/test_utils.py b/tests/litellm_utils_tests/test_utils.py index 2b539b97c9b..3c73224d7a1 100644 --- a/tests/litellm_utils_tests/test_utils.py +++ b/tests/litellm_utils_tests/test_utils.py @@ -2147,7 +2147,7 @@ def test_validate_user_messages_invalid_content_type(): messages = [{"content": [{"type": "invalid_type", "text": "Hello"}]}] - with pytest.raises(Exception) as e: + with pytest.raises(Exception, match='Please ensure all messages are valid OpenAI chat completion') as e: validate_chat_completion_user_messages(messages) assert "Invalid message" in str(e) diff --git a/tests/litellm_utils_tests/test_validate_tool_choice.py b/tests/litellm_utils_tests/test_validate_tool_choice.py index 0e6294a7cd4..07f8c9ed8f4 100644 --- a/tests/litellm_utils_tests/test_validate_tool_choice.py +++ b/tests/litellm_utils_tests/test_validate_tool_choice.py @@ -37,27 +37,27 @@ def test_validate_tool_choice_cursor_format(): def test_validate_tool_choice_invalid_dict(): """Test that invalid dict formats raise exceptions.""" # Missing both type and function - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Invalid tool choice, tool_choice=\\{\\}\\. Please ensure') as exc_info: validate_chat_completion_tool_choice({}) assert "Invalid tool choice" in str(exc_info.value) # Invalid type value - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="Invalid tool choice, tool_choice=\\{'type': 'invalid'\\}\\.") as exc_info: validate_chat_completion_tool_choice({"type": "invalid"}) assert "Invalid tool choice" in str(exc_info.value) # Has type but missing function when type is "function" - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="Invalid tool choice, tool_choice=\\{'type': 'function'\\}\\.") as exc_info: validate_chat_completion_tool_choice({"type": "function"}) assert "Invalid tool choice" in str(exc_info.value) def test_validate_tool_choice_invalid_type(): """Test that invalid types raise exceptions.""" - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="\\. Expecting str, or dict\\. Please ensure") as exc_info: validate_chat_completion_tool_choice(123) assert "Got=" in str(exc_info.value) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="Invalid tool choice, tool_choice=\\[\\]\\. Got=\\.") as exc_info: validate_chat_completion_tool_choice([]) assert "Got=" in str(exc_info.value) diff --git a/tests/llm_responses_api_testing/test_responses_hooks.py b/tests/llm_responses_api_testing/test_responses_hooks.py index 2344a62de4d..66dbb29dba5 100644 --- a/tests/llm_responses_api_testing/test_responses_hooks.py +++ b/tests/llm_responses_api_testing/test_responses_hooks.py @@ -295,7 +295,7 @@ async def test_responses_streaming_failure_triggers_failure_handlers(): call_type=CallTypes.responses.value, ) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="boom"): iterator._process_chunk('{"delta": "chunk"}') # allow failure callbacks to run diff --git a/tests/llm_translation/test_bedrock_completion.py b/tests/llm_translation/test_bedrock_completion.py index 3303fafafb0..9534bc8de3c 100644 --- a/tests/llm_translation/test_bedrock_completion.py +++ b/tests/llm_translation/test_bedrock_completion.py @@ -1890,7 +1890,7 @@ def test_bedrock_completion_test_4(modify_params): ] assert transformed_messages == expected_messages else: - with pytest.raises(Exception) as e: + with pytest.raises(Exception, match=r"litellm\.modify_params") as e: litellm.completion(**data) assert "litellm.modify_params" in str(e.value) diff --git a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py index cc493cc5a28..8c7390d3d04 100644 --- a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py +++ b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py @@ -982,7 +982,7 @@ def test_convert_to_model_response_object_with_real_error(): }, } - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception) as exc_info: # noqa: PT011 # message rides on .message, str() is empty convert_to_model_response_object( model_response_object=ModelResponse(), response_object=response_object, @@ -1243,7 +1243,7 @@ def test_convert_to_model_response_object_with_error_code_only(): }, } - with pytest.raises(Exception) as exc_info: # noqa: B017 # bare Exception raised, so status_code is the assertion + with pytest.raises(Exception) as exc_info: # noqa: B017, PT011 # bare Exception, empty message, so status_code is the assertion convert_to_model_response_object( model_response_object=ModelResponse(), response_object=response_object, @@ -1423,7 +1423,7 @@ def test_error_message_includes_function_args(): "choices": [{"index": 0}], } - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='in convert_to_model_response_object') as exc_info: convert_to_model_response_object( model_response_object=ModelResponse(), response_object=response_object, diff --git a/tests/llm_translation/test_prompt_factory.py b/tests/llm_translation/test_prompt_factory.py index c3519fcb40f..1b4c8a82cf4 100644 --- a/tests/llm_translation/test_prompt_factory.py +++ b/tests/llm_translation/test_prompt_factory.py @@ -1845,7 +1845,7 @@ def test_parse_tool_call_arguments_malformed_json(): parse_tool_call_arguments, ) - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match="Failed to parse tool call arguments for tool 'load_skill") as exc_info: parse_tool_call_arguments( '{"skill_name": "pptx', tool_name="load_skill", @@ -1877,7 +1877,7 @@ def test_convert_to_anthropic_tool_invoke_malformed_json(): } ] - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match="Failed to parse tool call arguments for tool 'bad_tool") as exc_info: convert_to_anthropic_tool_invoke(tool_calls) error_msg = str(exc_info.value) @@ -2023,7 +2023,7 @@ def test_parse_tool_call_arguments_still_raises_for_unrepairable(): parse_tool_call_arguments, ) - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match="Failed to parse tool call arguments for tool 'test_tool") as exc_info: parse_tool_call_arguments( '{"key": "unterminated', tool_name="test_tool", diff --git a/tests/llm_translation/test_triton.py b/tests/llm_translation/test_triton.py index 21887e8d848..f4a26360a6c 100644 --- a/tests/llm_translation/test_triton.py +++ b/tests/llm_translation/test_triton.py @@ -45,7 +45,7 @@ def test_split_embedding_by_shape_fails_with_shape_value_error(): "data": [1, 2, 3, 4, 5, 6], } ] - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Shape must be of length'): TritonEmbeddingConfig.split_embedding_by_shape( data[0]["data"], data[0]["shape"] ) diff --git a/tests/llm_translation/test_unit_test_bedrock_invoke.py b/tests/llm_translation/test_unit_test_bedrock_invoke.py index 14f08c759c5..39f02263f03 100644 --- a/tests/llm_translation/test_unit_test_bedrock_invoke.py +++ b/tests/llm_translation/test_unit_test_bedrock_invoke.py @@ -59,7 +59,7 @@ def test_transform_request_invalid_provider(bedrock_transformer): """Test request transformation with invalid provider""" messages = [{"role": "user", "content": "Hello"}] - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Bedrock Invoke HTTPX: Unknown provider=None') as exc_info: bedrock_transformer.transform_request( model="invalid.model", messages=messages, diff --git a/tests/local_testing/test_auth_utils.py b/tests/local_testing/test_auth_utils.py index 9aecb7e10e4..88e8c02a606 100644 --- a/tests/local_testing/test_auth_utils.py +++ b/tests/local_testing/test_auth_utils.py @@ -264,10 +264,6 @@ def test_get_end_user_id_from_request_body_backwards_compatibility(): ["gpt-3.5-turbo", "gpt-4o-mini-general-deployment"], ), ({"model": "gpt-3.5-turbo"}, "gpt-3.5-turbo"), - ( - {"model": "gpt-3.5-turbo, gpt-4o-mini-general-deployment"}, - ["gpt-3.5-turbo", "gpt-4o-mini-general-deployment"], - ), ], ) def test_get_model_from_request(request_data, expected_model): diff --git a/tests/local_testing/test_exceptions.py b/tests/local_testing/test_exceptions.py index edf847f4cef..8dd90cbfb37 100644 --- a/tests/local_testing/test_exceptions.py +++ b/tests/local_testing/test_exceptions.py @@ -1433,7 +1433,7 @@ async def test_exception_bubbling_up(sync_mode, stream_mode, model): sync_stream=sync_mode, ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='litellm\\.BadRequestError: OpenAIException - Invalid value') as exc_info: await _call_with_bad_role() assert exc_info.value.code == "invalid_value" diff --git a/tests/local_testing/test_file_types.py b/tests/local_testing/test_file_types.py index db83ba0e74b..7fda81ebd45 100644 --- a/tests/local_testing/test_file_types.py +++ b/tests/local_testing/test_file_types.py @@ -23,13 +23,13 @@ class TestFileConsts: def test_get_file_extension_from_mime_type(self): assert get_file_extension_from_mime_type("audio/aac") == "aac" assert get_file_extension_from_mime_type("application/pdf") == "pdf" - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Unknown extension for mime type: application'): get_file_extension_from_mime_type("application/unknown") def test_get_file_type_from_extension(self): assert get_file_type_from_extension("aac") == FileType.AAC assert get_file_type_from_extension("pdf") == FileType.PDF - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Unknown file type for extension: unknown'): get_file_type_from_extension("unknown") def test_get_file_extension_for_file_type(self): diff --git a/tests/local_testing/test_get_model_info.py b/tests/local_testing/test_get_model_info.py index 385be25fb07..cef05050ac9 100644 --- a/tests/local_testing/test_get_model_info.py +++ b/tests/local_testing/test_get_model_info.py @@ -134,7 +134,6 @@ def test_get_model_info_bedrock_region(): "ft:gpt-3.5-turbo:my-org:custom_suffix:id", "ft:gpt-4-0613:my-org:custom_suffix:id", "ft:davinci-002:my-org:custom_suffix:id", - "ft:gpt-4-0613:my-org:custom_suffix:id", "ft:babbage-002:my-org:custom_suffix:id", "gpt-35-turbo", "ada", diff --git a/tests/local_testing/test_router_budget_limiter.py b/tests/local_testing/test_router_budget_limiter.py index 48915137138..4ef99ec8c12 100644 --- a/tests/local_testing/test_router_budget_limiter.py +++ b/tests/local_testing/test_router_budget_limiter.py @@ -160,7 +160,7 @@ async def test_provider_budgets_e2e_test_expect_to_fail(): await asyncio.sleep(2.5) for _ in range(3): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="Exceeded budget for provider") as exc_info: await router.acompletion( messages=[{"role": "user", "content": "Hello, how are you?"}], model="anthropic/claude-sonnet-4-5-20250929", @@ -594,7 +594,7 @@ async def test_deployment_budgets_e2e_test_expect_to_fail(): await asyncio.sleep(2.5) for _ in range(3): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="Exceeded budget for deployment") as exc_info: await router.acompletion( messages=[{"role": "user", "content": "Hello, how are you?"}], model="openai/gpt-4o-mini", @@ -646,7 +646,7 @@ async def test_tag_budgets_e2e_test_expect_to_fail(): await asyncio.sleep(2.5) for _ in range(3): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match=f"Exceeded budget for tag='{TAG_NAME}'") as exc_info: await router.acompletion( messages=[{"role": "user", "content": "Hello, how are you?"}], model="openai/gpt-4o-mini", diff --git a/tests/local_testing/test_router_fallbacks.py b/tests/local_testing/test_router_fallbacks.py index 15c6c5fec59..86dec406332 100644 --- a/tests/local_testing/test_router_fallbacks.py +++ b/tests/local_testing/test_router_fallbacks.py @@ -1430,7 +1430,7 @@ async def test_router_fallbacks_default_and_model_specific_fallbacks(sync_mode): messages=[{"role": "user", "content": "Hey, how's it going?"}], ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='litellm\\.AuthenticationError: AuthenticationError') as exc_info: await _call_bad_model() assert isinstance( exc_info.value, litellm.AuthenticationError diff --git a/tests/logging_callback_tests/test_standard_logging_payload.py b/tests/logging_callback_tests/test_standard_logging_payload.py index d13cdf1337a..6a632c32fc2 100644 --- a/tests/logging_callback_tests/test_standard_logging_payload.py +++ b/tests/logging_callback_tests/test_standard_logging_payload.py @@ -293,7 +293,7 @@ def test_cleanup_timestamps(): assert all(isinstance(x, float) for x in result) # Test invalid input - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="start_time is required, got=invalid of type "): StandardLoggingPayloadSetup.cleanup_timestamps( "invalid", end_float, completion_float ) diff --git a/tests/multi_instance_e2e_tests/test_update_team_e2e.py b/tests/multi_instance_e2e_tests/test_update_team_e2e.py index dfbfbd310ee..13091fd3df6 100644 --- a/tests/multi_instance_e2e_tests/test_update_team_e2e.py +++ b/tests/multi_instance_e2e_tests/test_update_team_e2e.py @@ -143,7 +143,7 @@ async def test_team_blocking_behavior_multi_instance(): assert team_info_4001["blocked"] is True, "Team should be blocked after update" # 8. Make a chat completion request on port 4000 with a new prompt; expect it to be blocked. - with pytest.raises(Exception) as excinfo: + with pytest.raises(Exception, match="(?i)blocked") as excinfo: await chat_completion_on_port( session, key=key, @@ -157,7 +157,7 @@ async def test_team_blocking_behavior_multi_instance(): ), f"Expected error indicating team blocked, got: {error_msg}" # 9. Make a chat completion request on port 4000 with a new prompt; expect it to be blocked. - with pytest.raises(Exception) as excinfo: + with pytest.raises(Exception, match="(?i)blocked") as excinfo: await chat_completion_on_port( session, key=key, @@ -171,7 +171,7 @@ async def test_team_blocking_behavior_multi_instance(): ), f"Expected error indicating team blocked, got: {error_msg}" # 9. Repeat the chat completion request with another new prompt; expect it to be blocked. - with pytest.raises(Exception) as excinfo_second: + with pytest.raises(Exception, match="(?i)blocked") as excinfo_second: await chat_completion_on_port( session, key=key, diff --git a/tests/ocr_tests/test_ocr_azure_document_intelligence.py b/tests/ocr_tests/test_ocr_azure_document_intelligence.py index 5736bd797e3..e6a2e5e5735 100644 --- a/tests/ocr_tests/test_ocr_azure_document_intelligence.py +++ b/tests/ocr_tests/test_ocr_azure_document_intelligence.py @@ -101,7 +101,7 @@ class TestAzureDocumentIntelligencePagesParam: cfg.map_ocr_params({"pages": [True, False]}, {}, "prebuilt-layout") def test_map_ocr_params_unsupported_type_raises(self, cfg): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='based, Mistral-style\\) or a string like'): cfg.map_ocr_params({"pages": 5}, {}, "prebuilt-layout") def test_get_complete_url_appends_pages_query(self, cfg): diff --git a/tests/otel_tests/test_e2e_model_access.py b/tests/otel_tests/test_e2e_model_access.py index 5b5f2a89c8d..e5e93c0b179 100644 --- a/tests/otel_tests/test_e2e_model_access.py +++ b/tests/otel_tests/test_e2e_model_access.py @@ -3,6 +3,7 @@ import asyncio import aiohttp import json from httpx import AsyncClient +from openai import PermissionDeniedError from typing import Any, Optional, List, Literal @@ -134,7 +135,7 @@ async def test_model_access_update(): await mock_chat_completion(session=session, key=key, model="openai/gpt-5.5") # Should fail with gpt-5-mini - with pytest.raises(Exception) as exc_info: + with pytest.raises(PermissionDeniedError) as exc_info: await mock_chat_completion( session=session, key=key, model="openai/gpt-5-mini" ) @@ -157,7 +158,7 @@ async def test_model_access_update(): ) # Non-OpenAI model should still fail - with pytest.raises(Exception) as exc_info: + with pytest.raises(PermissionDeniedError) as exc_info: await mock_chat_completion( session=session, key=key, model="anthropic/claude-2" ) @@ -254,7 +255,7 @@ async def test_team_model_access_update(): await mock_chat_completion(session=session, key=key, model="openai/gpt-5.5") # Should fail with gpt-5-mini - with pytest.raises(Exception) as exc_info: + with pytest.raises(PermissionDeniedError) as exc_info: await mock_chat_completion( session=session, key=key, model="openai/gpt-5-mini" ) @@ -279,7 +280,7 @@ async def test_team_model_access_update(): ) # Non-OpenAI model should still fail - with pytest.raises(Exception) as exc_info: + with pytest.raises(PermissionDeniedError) as exc_info: await mock_chat_completion( session=session, key=key, model="anthropic/claude-2" ) diff --git a/tests/proxy_admin_ui_tests/test_key_management.py b/tests/proxy_admin_ui_tests/test_key_management.py index 4c5a045509a..7e8494b77fc 100644 --- a/tests/proxy_admin_ui_tests/test_key_management.py +++ b/tests/proxy_admin_ui_tests/test_key_management.py @@ -1340,6 +1340,6 @@ async def test_team_model_alias(prisma_client, requested_model, should_pass): }, "Expected model aliases to be present" else: # Verify the key fails with non-aliased models - with pytest.raises(Exception) as exc_info: + with pytest.raises(ProxyException) as exc_info: await user_api_key_auth(request=request, api_key=f"Bearer {generated_key}") assert exc_info.value.type == ProxyErrorTypes.key_model_access_denied diff --git a/tests/proxy_admin_ui_tests/test_role_based_access.py b/tests/proxy_admin_ui_tests/test_role_based_access.py index 9398428bd67..f9506fb694b 100644 --- a/tests/proxy_admin_ui_tests/test_role_based_access.py +++ b/tests/proxy_admin_ui_tests/test_role_based_access.py @@ -9,7 +9,7 @@ from litellm._uuid import uuid from datetime import datetime from dotenv import load_dotenv -from fastapi import Request +from fastapi import HTTPException, Request from fastapi.routing import APIRoute load_dotenv() @@ -530,7 +530,7 @@ async def test_user_role_permissions(prisma_client, route, user_role, expected_r print(f"Auth passed as expected for {route} with role {user_role}") else: # Should raise an error - with pytest.raises(Exception) as exc_info: + with pytest.raises((ProxyException, HTTPException)) as exc_info: await user_api_key_auth(request=request, api_key=bearer_token) print(f"Auth failed as expected for {route} with role {user_role}") print(f"Error message: {str(exc_info.value)}") diff --git a/tests/proxy_unit_tests/test_auth_checks.py b/tests/proxy_unit_tests/test_auth_checks.py index ffcbe472be7..ef3cbd0ae95 100644 --- a/tests/proxy_unit_tests/test_auth_checks.py +++ b/tests/proxy_unit_tests/test_auth_checks.py @@ -173,7 +173,7 @@ async def test_can_key_call_model(model, expect_to_work): if expect_to_work: await can_key_call_model(**args) else: - with pytest.raises(Exception) as e: + with pytest.raises(Exception, match='key not allowed to access model\\. This key can only access') as e: await can_key_call_model(**args) print(e) @@ -958,7 +958,7 @@ async def test_can_key_call_model_with_aliases(model, alias_map, expect_to_work) llm_router=router, ) else: - with pytest.raises(Exception) as e: + with pytest.raises(Exception, match='key not allowed to access model\\. This key can only access') as e: await can_key_call_model( model=model, llm_model_list=llm_model_list, diff --git a/tests/proxy_unit_tests/test_jwt.py b/tests/proxy_unit_tests/test_jwt.py index 4db47a1cde4..686d7021257 100644 --- a/tests/proxy_unit_tests/test_jwt.py +++ b/tests/proxy_unit_tests/test_jwt.py @@ -1583,7 +1583,7 @@ async def test_auth_jwt_mismatched_key_fails(monkeypatch): h = JWTHandler() with patch.object(h, "get_public_key", new=AsyncMock(return_value=rsa_jwk)): - with pytest.raises(Exception) as exc: + with pytest.raises(Exception, match='Validation fails: Expecting a PEM-formatted key\\.') as exc: await h.auth_jwt(token) assert "Validation fails" in str(exc.value) @@ -1826,7 +1826,7 @@ async def test_multi_issuer_jwt_unknown_issuer_without_global_jwks_rejected( kid="issuer-key", ) - with pytest.raises(Exception) as exc: + with pytest.raises(Exception, match='Missing JWT Public Key URL from environment\\.') as exc: await jwt_handler.auth_jwt(token=token) assert "Missing JWT Public Key URL" in str(exc.value) @@ -1857,7 +1857,7 @@ async def test_multi_issuer_jwt_rejects_wrong_audience(monkeypatch): kid="issuer-key", ) - with pytest.raises(Exception) as exc: + with pytest.raises(Exception, match="Validation fails: Audience doesn't match") as exc: await jwt_handler.auth_jwt(token=token) assert "Validation fails" in str(exc.value) @@ -1900,7 +1900,7 @@ async def test_multi_issuer_jwt_same_kid_does_not_cross_issuer_keys(monkeypatch) kid=shared_kid, ) - with pytest.raises(Exception) as exc: + with pytest.raises(Exception, match='Validation fails: Signature verification failed') as exc: await jwt_handler.auth_jwt(token=token) assert "Validation fails" in str(exc.value) @@ -1953,7 +1953,7 @@ def test_multi_issuer_jwt_requires_audience_unless_explicitly_disabled( issuer = "https://issuer.example.com" jwks_url = f"{issuer}/keys" - with pytest.raises(Exception) as exc: + with pytest.raises(Exception, match='must configure audience or set') as exc: LiteLLM_JWTAuth( issuers=[ { diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index 00ed6d13f63..de2a9282300 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -1026,7 +1026,7 @@ def test_enforced_params_check( from litellm.proxy.litellm_pre_call_utils import _enforced_params_check if expected_error: - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='in request body\\. This is a required param'): _enforced_params_check( request_body=request_body, general_settings=general_settings, @@ -2626,7 +2626,7 @@ async def test_during_call_hook_parallel_execution_with_error(): try: litellm.callbacks = [FailingGuardrail()] - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Guardrail violation detected!') as exc_info: await proxy_logging.during_call_hook( data={ "model": "gpt-4", diff --git a/tests/proxy_unit_tests/test_update_spend.py b/tests/proxy_unit_tests/test_update_spend.py index 0d1d6dcf3c6..6b8973fbad2 100644 --- a/tests/proxy_unit_tests/test_update_spend.py +++ b/tests/proxy_unit_tests/test_update_spend.py @@ -166,7 +166,7 @@ async def test_update_spend_logs_non_connection_error(): prisma_client.db.litellm_spendlogs.create_many = create_many_mock # Execute and verify it raises immediately without retrying - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Unexpected database error') as exc_info: await update_spend(prisma_client, None, proxy_logging_obj) # Verify error message diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index 1fef0f01df8..f81578dbd99 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -90,7 +90,7 @@ def test_routing_strategy_init_invalid_strategy(model_list): router = Router(model_list=model_list) # Test common mistake: "simple" instead of "simple-shuffle" - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match="usage-based-routing', 'provider-budget-routing'\\]\\. Check") as exc_info: router.routing_strategy_init( routing_strategy="simple", routing_strategy_args={} ) @@ -106,7 +106,7 @@ def test_routing_strategy_init_invalid_strategy(model_list): assert "Router SDK" in error_msg # Test completely invalid strategy - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match="usage-based-routing', 'provider-budget-routing'\\]\\. Check") as exc_info: router.routing_strategy_init( routing_strategy="not-a-real-strategy", routing_strategy_args={} ) diff --git a/tests/store_model_in_db_tests/test_mcp_servers.py b/tests/store_model_in_db_tests/test_mcp_servers.py index e9c26221580..735d5d71ad3 100644 --- a/tests/store_model_in_db_tests/test_mcp_servers.py +++ b/tests/store_model_in_db_tests/test_mcp_servers.py @@ -471,7 +471,7 @@ def test_validate_mcp_server_name_direct(): validate_mcp_server_name("valid name") # Test that invalid names with hyphens raise exceptions - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="Server name cannot contain '-'\\. Use an alternative") as exc_info: validate_mcp_server_name("invalid-name") assert "cannot contain" in str(exc_info.value) diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py index 40439a78a49..5fe4b217e4f 100644 --- a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py +++ b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py @@ -104,7 +104,7 @@ async def test_send_email_missing_api_key(): try: logger = SendGridEmailLogger() - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='SENDGRID_API_KEY is not set'): await logger.send_email( from_email="test@example.com", to_email=["recipient@example.com"], diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py index 6cc31f991a3..fcd03e77aa2 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -471,7 +471,7 @@ async def test_afile_content_error_reports_unified_id_not_provider_uri(): mock_router.get_deployment_credentials_with_provider = MagicMock(return_value=None) mock_router.afile_content = AsyncMock(side_effect=Exception("deployment failed")) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='LiteLLM Managed File object with') as exc_info: await managed_files.afile_content( file_id=unified_file_id, litellm_parent_otel_span=None, diff --git a/tests/test_litellm/integrations/cloudzero/test_cloudzero_database.py b/tests/test_litellm/integrations/cloudzero/test_cloudzero_database.py index 89a5028011c..7f930f90247 100644 --- a/tests/test_litellm/integrations/cloudzero/test_cloudzero_database.py +++ b/tests/test_litellm/integrations/cloudzero/test_cloudzero_database.py @@ -51,7 +51,7 @@ async def test_get_usage_data_rejects_invalid_limit(monkeypatch: pytest.MonkeyPa """limit must coerce to int or raise ValueError before hitting the DB.""" db, query_mock = _setup_db(monkeypatch, []) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='limit must be an integer'): await db.get_usage_data(limit="invalid") assert query_mock.await_count == 0 diff --git a/tests/test_litellm/integrations/cloudzero/test_cz_stream_api.py b/tests/test_litellm/integrations/cloudzero/test_cz_stream_api.py index 440ce39e021..a715116e5ee 100644 --- a/tests/test_litellm/integrations/cloudzero/test_cz_stream_api.py +++ b/tests/test_litellm/integrations/cloudzero/test_cz_stream_api.py @@ -108,7 +108,7 @@ class TestCloudZeroStreamer: """Test _parse_and_convert_timestamp method with invalid timestamp.""" streamer = CloudZeroStreamer("test-key", "test-connection") - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="Could not parse timestamp 'invalid-timestamp': Invalid"): streamer._parse_and_convert_timestamp("invalid-timestamp") def test_prepare_batch_payload(self): diff --git a/tests/test_litellm/integrations/focus/test_focus_database.py b/tests/test_litellm/integrations/focus/test_focus_database.py index d77af2dd170..5c13665f1f1 100644 --- a/tests/test_litellm/integrations/focus/test_focus_database.py +++ b/tests/test_litellm/integrations/focus/test_focus_database.py @@ -68,7 +68,7 @@ async def test_should_accept_string_timestamps(monkeypatch: pytest.MonkeyPatch): async def test_should_reject_invalid_limit(monkeypatch: pytest.MonkeyPatch): db, query_mock = _setup_db(monkeypatch, []) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='limit must be an integer'): await db.get_usage_data(limit="invalid") assert query_mock.await_count == 0 diff --git a/tests/test_litellm/integrations/focus/test_s3_destination.py b/tests/test_litellm/integrations/focus/test_s3_destination.py index f915b2c56a3..8e54b561f82 100644 --- a/tests/test_litellm/integrations/focus/test_s3_destination.py +++ b/tests/test_litellm/integrations/focus/test_s3_destination.py @@ -20,7 +20,7 @@ def _window(freq: str = "hourly", hour: int = 5) -> FocusTimeWindow: def test_should_require_bucket_name(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='bucket_name must be provided for S'): FocusS3Destination(prefix="focus", config={}) diff --git a/tests/test_litellm/integrations/gitlab/test_gitlab_client.py b/tests/test_litellm/integrations/gitlab/test_gitlab_client.py index 4556950cd3e..529868ca06a 100644 --- a/tests/test_litellm/integrations/gitlab/test_gitlab_client.py +++ b/tests/test_litellm/integrations/gitlab/test_gitlab_client.py @@ -95,9 +95,9 @@ def enc_project(p): # how client encodes project in urls # Constructor / config tests # ----------------------------- def test_init_requires_project_and_token(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='project and access_token are required'): GitLabClient({"project": "p"}) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='project and access_token are required'): GitLabClient({"access_token": "t"}) @@ -127,7 +127,7 @@ def test_set_ref_updates_effective_ref(): c = make_client(branch="main") c.set_ref("feature/x") assert c.ref == "feature/x" - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='ref must be a non-empty string'): c.set_ref("") @@ -193,12 +193,12 @@ def test_get_file_content_permission_errors_are_mapped(): raw_url = f"https://gitlab.example.com/api/v4/projects/{enc_project('group/sub/repo')}/repository/files/secure%2Ffile.prompt/raw?ref=main" # raise_for_status will be called, so return 403 response (not an exception from transport) c.http_handler.routes[raw_url] = FakeResponse(status_code=403) - with pytest.raises(Exception) as ei: + with pytest.raises(Exception, match="Check your GitLab permissions for project 'group") as ei: c.get_file_content("secure/file.prompt") assert "Access denied" in str(ei.value) c.http_handler.routes[raw_url] = FakeResponse(status_code=401) - with pytest.raises(Exception) as ei2: + with pytest.raises(Exception, match='Authentication failed\\. Check your GitLab token and') as ei2: c.get_file_content("secure/file.prompt") assert "Authentication failed" in str(ei2.value) diff --git a/tests/test_litellm/integrations/levo/test_levo.py b/tests/test_litellm/integrations/levo/test_levo.py index 98b0327dbf2..903be644671 100644 --- a/tests/test_litellm/integrations/levo/test_levo.py +++ b/tests/test_litellm/integrations/levo/test_levo.py @@ -198,7 +198,7 @@ class TestLevoIntegration(unittest.TestCase): """Test health check returns unhealthy status when required vars are missing.""" # Try to create logger without required env vars # This should fail during config, but we can test health check logic - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='LEVOAI_API_KEY environment variable is required for Levo'): LevoLogger.get_levo_config() @patch.dict( diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_metrics.py b/tests/test_litellm/integrations/otel/test_otel_v2_metrics.py index e1b8e4b5721..b810ffdc6be 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_metrics.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_metrics.py @@ -554,7 +554,7 @@ def test_token_type_rejected_from_either_list(attributes, monkeypatch): recorder rather than silently ignored, so the misconfig is caught at all.""" recorder = _recorder(monkeypatch, attributes) kwargs, response_obj, start, end = _build_call() - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='otel\\.attributes: gen_ai\\.token\\.type is a structural') as exc_info: recorder.record(kwargs, response_obj, start, end) # The dedicated discriminator guard, not the generic unknown-name path: assert # the specific reason so dropping that guard (and falling through to "unknown diff --git a/tests/test_litellm/integrations/test_langfuse.py b/tests/test_litellm/integrations/test_langfuse.py index 6e57a36c5b6..3c7dd51bff8 100644 --- a/tests/test_litellm/integrations/test_langfuse.py +++ b/tests/test_litellm/integrations/test_langfuse.py @@ -1163,7 +1163,7 @@ def test_max_langfuse_clients_limit(): assert litellm.initialized_langfuse_clients == 2 # Third client should fail with exception - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Max langfuse clients reached') as exc_info: logger3 = LangFuseLogger( langfuse_public_key="test_key_3", langfuse_secret="test_secret_3", diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index fffbc884782..08d8c17cc2e 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -1169,7 +1169,7 @@ def test_bedrock_image_processor_content_type_fallback_failure(): # Test with URL without recognizable extension image_url = "https://example.com/unknown-file" - with pytest.raises(ValueError) as excinfo: + with pytest.raises(ValueError, match='Unable to determine content type from URL: https') as excinfo: BedrockImageProcessor._post_call_image_processing(mock_response, image_url) assert "Unable to determine content type" in str(excinfo.value) diff --git a/tests/test_litellm/litellm_core_utils/test_initialize_dynamic_callback_params.py b/tests/test_litellm/litellm_core_utils/test_initialize_dynamic_callback_params.py index 0dca4f3a1b1..956f86a9292 100644 --- a/tests/test_litellm/litellm_core_utils/test_initialize_dynamic_callback_params.py +++ b/tests/test_litellm/litellm_core_utils/test_initialize_dynamic_callback_params.py @@ -110,7 +110,7 @@ def test_top_level_kwargs_overrides_metadata_slots(): def test_env_reference_at_top_level_raises_with_guidance(): kwargs = {"langfuse_public_key": "os.environ/LANGFUSE_PUBLIC_KEY"} - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match="Callback param 'langfuse_public_key' \\(from request body\\)") as exc_info: initialize_standard_callback_dynamic_params(kwargs) message = str(exc_info.value) @@ -127,7 +127,7 @@ def test_env_reference_in_metadata_raises_with_guidance(): } } - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match="Callback param 'langsmith_api_key' \\(from metadata\\) contains") as exc_info: initialize_standard_callback_dynamic_params(kwargs) message = str(exc_info.value) diff --git a/tests/test_litellm/litellm_core_utils/test_llm_judge.py b/tests/test_litellm/litellm_core_utils/test_llm_judge.py index 5c092caa7c3..a0a2311914b 100644 --- a/tests/test_litellm/litellm_core_utils/test_llm_judge.py +++ b/tests/test_litellm/litellm_core_utils/test_llm_judge.py @@ -27,7 +27,7 @@ def test_parse_json_verdict_tolerates_fences_and_prose(raw, expected): def test_parse_json_verdict_rejects_non_object(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='judge response is not a JSON object'): parse_json_verdict('["not", "an", "object"]') with pytest.raises((json.JSONDecodeError, ValueError)): parse_json_verdict("no json here at all") diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 456d4db4afe..fbdfcac1adc 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -982,7 +982,7 @@ async def test_bedrock_validation_error_raises_directly(logging_obj: Logging): make_call=_raise_400, ) - with pytest.raises(Exception) as excinfo: + with pytest.raises(Exception, match='litellm\\.BadRequestError: BedrockException') as excinfo: await response.__anext__() assert not isinstance(excinfo.value, MidStreamFallbackError) assert getattr(excinfo.value, "status_code", None) == 400 @@ -2722,7 +2722,7 @@ def test_dispatch_text_completion_codestral_requires_string( is a programming error and must surface loudly.""" initialized_custom_stream_wrapper.custom_llm_provider = "text-completion-codestral" - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="chunk is not a string: \\{'not': 'a string'\\}"): _run_dispatch(initialized_custom_stream_wrapper, {"not": "a string"}) diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index 3c33ee13c3f..eec4b307c87 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -763,24 +763,6 @@ class TestTokenizerSelection(unittest.TestCase): ], } ], - [ - { - "role": "user", - "content": [ - { - "type": "text", - "text": "These are some sample images from a movie. Based on these images, what do you think the tone of the movie is?", - }, - { - "type": "text", - "image_url": { - "url": "https://gratisography.com/wp-content/uploads/2024/11/gratisography-augmented-reality-800x525.jpg", - "detail": "high", - }, - }, - ], - } - ], ], ) def test_bad_input_token_counter(model, messages): @@ -1174,7 +1156,7 @@ def test_count_content_list_rejects_unknown_type(): """ from litellm.litellm_core_utils.token_counter import _count_content_list - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Error getting number of tokens from content list: Invalid') as exc_info: _count_content_list( count_function=len, content_list=[{"type": "totally_unknown_block"}], diff --git a/tests/test_litellm/litellm_core_utils/test_url_utils.py b/tests/test_litellm/litellm_core_utils/test_url_utils.py index cef09f3f2b0..751b548adcd 100644 --- a/tests/test_litellm/litellm_core_utils/test_url_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_url_utils.py @@ -100,12 +100,12 @@ class TestEncodeUrlPathSegment: @pytest.mark.parametrize("value", ["", ".", "..", None]) def test_rejects_empty_and_dot_segments(self, value): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="resource_id (is required|cannot be a dot path segment)"): encode_url_path_segment(value, field_name="resource_id") @pytest.mark.parametrize("value", ["../model", "model/../other", "/model"]) def test_rejects_dot_segments_in_multi_segment_paths(self, value): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="model (is required|cannot be a dot path segment)"): encode_url_path_segments(value, field_name="model") diff --git a/tests/test_litellm/llms/aiml/image_generation/test_aiml_image_generation_transformation.py b/tests/test_litellm/llms/aiml/image_generation/test_aiml_image_generation_transformation.py index 7485f2121df..dd74379a883 100644 --- a/tests/test_litellm/llms/aiml/image_generation/test_aiml_image_generation_transformation.py +++ b/tests/test_litellm/llms/aiml/image_generation/test_aiml_image_generation_transformation.py @@ -113,7 +113,7 @@ def test_flux_style_request_still_remaps_to_legacy_fields(): def test_openai_style_unsupported_param_raises_without_drop_params(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Supported parameters are'): AimlImageGenerationConfig().map_openai_params( non_default_params={"image_size": {"width": 1024, "height": 1024}}, optional_params={}, diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py index 654b0097546..f3cb2956aeb 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py @@ -61,7 +61,7 @@ def test_anthropic_messages_handler_skips_the_gateway_on_recursion(): "litellm.llms.anthropic.experimental_pass_through.messages.mcp_handler.anthropic_messages_with_mcp", new=AsyncMock(return_value={"routed": True}), ) as routed: - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='anthropic_messages_handler is not implemented for sync calls'): anthropic_messages_handler( max_tokens=100, messages=[{"role": "user", "content": "hi"}], @@ -80,7 +80,7 @@ def test_anthropic_messages_handler_leaves_native_tools_alone(): "litellm.llms.anthropic.experimental_pass_through.messages.mcp_handler.anthropic_messages_with_mcp", new=AsyncMock(return_value={"routed": True}), ) as routed: - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='anthropic_messages_handler is not implemented for sync calls'): anthropic_messages_handler( max_tokens=100, messages=[{"role": "user", "content": "hi"}], diff --git a/tests/test_litellm/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py b/tests/test_litellm/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py index ffabce6e00c..602cbf68f3f 100644 --- a/tests/test_litellm/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py +++ b/tests/test_litellm/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py @@ -16,7 +16,7 @@ class TestAzureAIRerankConfigGetCompleteUrl: self.model = "azure_ai/cohere-rerank-v3-english" def test_api_base_required(self): - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Azure AI API Base is required\\. api_base=None\\. Set in') as exc_info: self.config.get_complete_url(api_base=None, model=self.model) assert "api_base=None" in str(exc_info.value) @@ -31,7 +31,7 @@ class TestAzureAIRerankConfigGetCompleteUrl: ], ) def test_api_base_requires_scheme(self, api_base): - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Azure AI API Base must be an absolute URL including scheme') as exc_info: self.config.get_complete_url(api_base=api_base, model=self.model) error_message = str(exc_info.value).lower() diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index cfe9930e76e..b9f8283b78e 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -1944,7 +1944,7 @@ def test_role_assumption_access_denied_raises_when_different_role(): with patch.object( base_aws_llm, "_is_already_running_as_role", return_value=False ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='An error occurred \\(AccessDenied\\) when calling the') as exc_info: base_aws_llm._auth_with_aws_role( aws_access_key_id=None, aws_secret_access_key=None, @@ -1969,7 +1969,7 @@ def test_role_assumption_non_access_denied_error_propagated(): ) with patch("boto3.client", return_value=mock_sts_client): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='An error occurred \\(MalformedPolicyDocument\\) when calling') as exc_info: base_aws_llm._auth_with_aws_role( aws_access_key_id=None, aws_secret_access_key=None, diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py index 8281f3387d9..28c8e5c7ed6 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py @@ -109,7 +109,7 @@ class TestBedrockMantleResponsesURL: monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) monkeypatch.delenv("AWS_REGION", raising=False) cfg = BedrockMantleResponsesAPIConfig() - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="api\\.aws\\.attacker\\.example/'\\. Region names must contain only"): cfg.get_complete_url( api_base=None, litellm_params={ @@ -1418,7 +1418,7 @@ class TestBedrockMantleResponsesSigV4: signer.get_credentials = MagicMock(side_effect=NoCredentialsError()) cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) - with pytest.raises(ValueError) as exc: + with pytest.raises(ValueError, match='Bedrock Mantle auth failed: no Bearer token and no usable') as exc: cfg.sign_request( headers={}, optional_params={"aws_region_name": "us-east-2"}, @@ -1448,7 +1448,7 @@ class TestBedrockMantleResponsesSigV4: signer.get_credentials = MagicMock(side_effect=cred_error) cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) - with pytest.raises(ValueError) as exc: + with pytest.raises(ValueError, match='Bedrock Mantle auth failed: no Bearer token and no usable') as exc: cfg.sign_request( headers={}, optional_params={"aws_region_name": "us-east-2"}, diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index 275fb460b9f..07910b0b56f 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -107,7 +107,7 @@ class TestBedrockMantleConfig: monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) monkeypatch.delenv("AWS_REGION", raising=False) cfg = BedrockMantleChatConfig() - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="api\\.aws\\.attacker\\.example/'\\. Region names must contain only"): cfg._get_openai_compatible_provider_info( None, None, @@ -416,7 +416,7 @@ class TestBedrockMantleChatAuth: signer.get_credentials = MagicMock(side_effect=NoCredentialsError()) cfg = BedrockMantleChatConfig(aws_signer=signer) - with pytest.raises(ValueError) as exc: + with pytest.raises(ValueError, match='Bedrock Mantle auth failed: no Bearer token and no usable') as exc: cfg.sign_request( headers={}, optional_params={"aws_region_name": "us-east-2"}, diff --git a/tests/test_litellm/llms/bytez/chat/test_bytez_chat_transformation.py b/tests/test_litellm/llms/bytez/chat/test_bytez_chat_transformation.py index e2421437720..94b8c51dd52 100644 --- a/tests/test_litellm/llms/bytez/chat/test_bytez_chat_transformation.py +++ b/tests/test_litellm/llms/bytez/chat/test_bytez_chat_transformation.py @@ -38,7 +38,7 @@ class TestBytezChatConfig: config = BytezChatConfig() headers = {} - with pytest.raises(Exception) as excinfo: + with pytest.raises(Exception, match='Missing api_key, make sure you pass in your api key') as excinfo: config.validate_environment( headers=headers, model=TEST_MODEL, diff --git a/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_transformation.py b/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_transformation.py index a5411078cf7..ae3c166e7aa 100644 --- a/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_transformation.py +++ b/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_transformation.py @@ -258,7 +258,7 @@ class TestDeepinfraRerankTransform: status_code = 401 headers = {"content-type": "application/json"} - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Authentication failed') as exc_info: self.config.get_error_class(error_message, status_code, headers) # The method should raise a BaseLLMException @@ -271,7 +271,7 @@ class TestDeepinfraRerankTransform: status_code = 404 headers = {"content-type": "application/json"} - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Model not found') as exc_info: self.config.get_error_class(error_message, status_code, headers) # Should extract the nested error message @@ -284,7 +284,7 @@ class TestDeepinfraRerankTransform: status_code = 503 headers = {"content-type": "application/json"} - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Service unavailable') as exc_info: self.config.get_error_class(error_message, status_code, headers) # Should extract the string detail @@ -296,7 +296,7 @@ class TestDeepinfraRerankTransform: status_code = 500 headers = {"content-type": "application/json"} - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Invalid JSON error message') as exc_info: self.config.get_error_class(error_message, status_code, headers) # Should use the original error message when JSON parsing fails diff --git a/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_nano_banana_transformation.py b/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_nano_banana_transformation.py index 593593bfa73..c0f74eff51b 100644 --- a/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_nano_banana_transformation.py +++ b/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_nano_banana_transformation.py @@ -113,7 +113,7 @@ def test_response_format_is_ignored(): def test_unsupported_param_raises_without_drop_params(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="Supported parameters are \\['n', 'response_format', 'size'\\]\\."): FalAINanoBananaConfig().map_openai_params( non_default_params={"style": "vivid"}, optional_params={}, diff --git a/tests/test_litellm/llms/featherless_ai/chat/test_featherless_chat_transformation.py b/tests/test_litellm/llms/featherless_ai/chat/test_featherless_chat_transformation.py index 4dc467575a0..bf40abd7016 100644 --- a/tests/test_litellm/llms/featherless_ai/chat/test_featherless_chat_transformation.py +++ b/tests/test_litellm/llms/featherless_ai/chat/test_featherless_chat_transformation.py @@ -44,7 +44,7 @@ class TestFeatherlessAIConfig: """Test error handling when API key is missing""" config = FeatherlessAIConfig() - with pytest.raises(ValueError) as excinfo: + with pytest.raises(ValueError, match='Missing Featherless AI API Key') as excinfo: config.validate_environment( headers={}, model="featherless-ai/Qwerky-72B", @@ -112,7 +112,7 @@ class TestFeatherlessAIConfig: "tool_choice": {"type": "function", "function": {"name": "get_weather"}} } optional_params = {} - with pytest.raises(Exception) as excinfo: + with pytest.raises(Exception, match="litellm\\.UnsupportedParamsError: Featherless AI doesn't") as excinfo: config.map_openai_params( non_default_params=non_default_params, optional_params=optional_params, @@ -138,7 +138,7 @@ class TestFeatherlessAIConfig: assert "tools" not in result # Test with tools and drop_params=False - with pytest.raises(Exception) as excinfo: + with pytest.raises(Exception, match="litellm\\.UnsupportedParamsError: Featherless AI doesn't") as excinfo: config.map_openai_params( non_default_params=non_default_params, optional_params=optional_params, diff --git a/tests/test_litellm/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py b/tests/test_litellm/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py index 30bf5860dee..521ea4f8263 100644 --- a/tests/test_litellm/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py @@ -301,7 +301,7 @@ class TestFireworksAIRerankTransform: mock_logging = MagicMock() model_response = RerankResponse() - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Failed to parse response: Invalid JSON: line') as exc_info: self.config.transform_rerank_response( model=self.model, raw_response=mock_response, diff --git a/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py b/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py index 9b57e1991de..bd9b7006e58 100644 --- a/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py +++ b/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py @@ -244,7 +244,7 @@ class TestGeminiImageEditTransformation: def test_transform_image_edit_request_without_image_raises(self) -> None: optional_params = {} - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Gemini image edit requires at least one image\\.'): self.config.transform_image_edit_request( model=self.model, prompt=self.prompt, diff --git a/tests/test_litellm/llms/gemini/test_gemini_client_setup.py b/tests/test_litellm/llms/gemini/test_gemini_client_setup.py index 51c6fedf5b8..48b010aca48 100644 --- a/tests/test_litellm/llms/gemini/test_gemini_client_setup.py +++ b/tests/test_litellm/llms/gemini/test_gemini_client_setup.py @@ -28,7 +28,7 @@ def test_gemini_completion_no_api_key(): del os.environ[key] # Test without mock_response to ensure actual API key validation - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='in _complete_vertex_ai_beta') as exc_info: completion( model="gemini/gemini-1.5-flash", messages=[{"role": "user", "content": "Test message"}], @@ -60,7 +60,7 @@ def test_gemini_completion_no_api_key_with_mock(): with patch("litellm.get_secret") as mock_get_secret: mock_get_secret.return_value = None - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='in _complete_vertex_ai_beta') as exc_info: completion( model="gemini/gemini-1.5-flash", messages=[{"role": "user", "content": "Test message"}], diff --git a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py index 6425e815db0..e6e6aa946d5 100644 --- a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py +++ b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py @@ -109,7 +109,7 @@ class TestHostedVLLMRerankTransform: ) assert url2 == "https://api.example.com/rerank" # Raises if api_base is None - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='api_base must be provided for Hosted VLLM rerank'): self.config.get_complete_url(None, self.model) def test_transform_response(self): diff --git a/tests/test_litellm/llms/langflow/chat/test_langflow_chat_transformation.py b/tests/test_litellm/llms/langflow/chat/test_langflow_chat_transformation.py index b3c4e0f1858..0c241add77b 100644 --- a/tests/test_litellm/llms/langflow/chat/test_langflow_chat_transformation.py +++ b/tests/test_litellm/llms/langflow/chat/test_langflow_chat_transformation.py @@ -46,7 +46,7 @@ def test_langflow_config_get_complete_url(): def test_langflow_config_get_complete_url_requires_api_base(): config = LangFlowConfig() - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='api_base is required for LangFlow\\. Set it via'): config.get_complete_url( api_base=None, api_key=None, diff --git a/tests/test_litellm/llms/modelscope/image_generation/test_modelscope_image_gen_transformation.py b/tests/test_litellm/llms/modelscope/image_generation/test_modelscope_image_gen_transformation.py index 7f00f53c451..fbcec3d4d2e 100644 --- a/tests/test_litellm/llms/modelscope/image_generation/test_modelscope_image_gen_transformation.py +++ b/tests/test_litellm/llms/modelscope/image_generation/test_modelscope_image_gen_transformation.py @@ -154,7 +154,7 @@ class TestModelScopeImageGenerationTransformation: mock_get_secret.return_value = None headers = {} - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='MODELSCOPE_API_KEY is not set\\. Please set it via') as exc_info: self.config.validate_environment( headers=headers, model=self.model, @@ -367,7 +367,7 @@ class TestModelScopeImageGenerationTransformation: model_response = ImageResponse(data=[]) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='litellm\\.BadRequestError: ModelScope error: Invalid prompt') as exc_info: self.config.transform_image_generation_response( model=self.model, raw_response=mock_response, @@ -393,7 +393,7 @@ class TestModelScopeImageGenerationTransformation: model_response = ImageResponse(data=[]) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='litellm\\.InternalServerError: Error parsing ModelScope') as exc_info: self.config.transform_image_generation_response( model=self.model, raw_response=mock_response, diff --git a/tests/test_litellm/llms/novita/chat/test_novita_chat_transformation.py b/tests/test_litellm/llms/novita/chat/test_novita_chat_transformation.py index 7a00b361252..ade5e4176e8 100644 --- a/tests/test_litellm/llms/novita/chat/test_novita_chat_transformation.py +++ b/tests/test_litellm/llms/novita/chat/test_novita_chat_transformation.py @@ -47,7 +47,7 @@ class TestNovitaConfig: """Test error handling when API key is missing""" config = NovitaConfig() - with pytest.raises(ValueError) as excinfo: + with pytest.raises(ValueError, match='Missing Novita AI API Key - A call is being made to novita') as excinfo: config.validate_environment( headers={}, model="novita/meta-llama/llama-3.3-70b-instruct", diff --git a/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py b/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py index 8be0780d86f..5aa96a66d2d 100644 --- a/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py +++ b/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py @@ -98,7 +98,7 @@ class TestOCIChatConfig: config = OCIChatConfig() headers = {} - with pytest.raises(Exception) as excinfo: + with pytest.raises(Exception, match='Missing required parameters: oci_user, oci_fingerprint') as excinfo: config.validate_environment( headers=headers, model=TEST_MODEL, @@ -272,7 +272,7 @@ class TestOCIChatConfig: "oci_serving_mode": "INVALID_MODE", } - with pytest.raises(Exception) as excinfo: + with pytest.raises(Exception, match="kwarg `oci_serving_mode` must be either 'ON_DEMAND' or") as excinfo: config.transform_request( model=TEST_MODEL_NAME, messages=TEST_MESSAGES, # type: ignore @@ -892,7 +892,7 @@ class TestOCISignerSupport: optional_params = {"oci_signer": MockSigner(), "method": "INVALID"} - with pytest.raises(ValueError) as excinfo: + with pytest.raises(ValueError, match='Unsupported HTTP method: INVALID') as excinfo: config.sign_request( headers={}, optional_params=optional_params, @@ -1604,7 +1604,7 @@ class TestOCIKeyNormalization: # We can't fully test signing without a real key, but we can verify # the error message indicates the key was processed (not a type error) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='why-can-t-i-import-my-pem-file for more details\\.') as exc_info: sign_with_manual_credentials( headers={}, optional_params=optional_params, @@ -1630,7 +1630,7 @@ class TestOCIKeyNormalization: "oci_key": crlf_pem, } - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='why-can-t-i-import-my-pem-file for more details\\.') as exc_info: sign_with_manual_credentials( headers={}, optional_params=optional_params, @@ -1692,7 +1692,7 @@ class TestOCIValidateEnvironment: def test_missing_required_credentials_raises_error(self, config): """Test that missing required credentials raise an error.""" - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Missing required parameters: oci_user, oci_fingerprint') as exc_info: config.validate_environment( headers={}, model="oci/xai.grok-3", @@ -1875,7 +1875,7 @@ class TestOCIImageUrlTransformation: } ] - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Prop `image_url` must be a string or an object with a `url`') as exc_info: adapt_messages_to_generic_oci_standard(messages) assert "image_url" in str(exc_info.value) @@ -1899,7 +1899,7 @@ class TestOCIImageUrlTransformation: } ] - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Prop `image_url` must be a string or an object with a `url`') as exc_info: adapt_messages_to_generic_oci_standard(messages) assert "image_url" in str(exc_info.value) diff --git a/tests/test_litellm/llms/pg_vector/vector_stores/test_pg_vector_transformation.py b/tests/test_litellm/llms/pg_vector/vector_stores/test_pg_vector_transformation.py index 56953a574d6..1d44b2bc278 100644 --- a/tests/test_litellm/llms/pg_vector/vector_stores/test_pg_vector_transformation.py +++ b/tests/test_litellm/llms/pg_vector/vector_stores/test_pg_vector_transformation.py @@ -42,7 +42,7 @@ class TestPGVectorStoreConfig: litellm_params = GenericLiteLLMParams() headers = {} - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='PG Vector API key is required\\. Set PG_VECTOR_API_KEY') as exc_info: config.validate_environment(headers, litellm_params) assert "PG Vector API key is required" in str(exc_info.value) @@ -84,7 +84,7 @@ class TestPGVectorStoreConfig: config = PGVectorStoreConfig() litellm_params = {} - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='PG Vector API base URL is required\\. Set') as exc_info: config.get_complete_url(None, litellm_params) assert "PG Vector API base URL is required" in str(exc_info.value) diff --git a/tests/test_litellm/llms/recraft/image_edit/test_recraft_image_edit_transformation.py b/tests/test_litellm/llms/recraft/image_edit/test_recraft_image_edit_transformation.py index 0acabd05805..47811321133 100644 --- a/tests/test_litellm/llms/recraft/image_edit/test_recraft_image_edit_transformation.py +++ b/tests/test_litellm/llms/recraft/image_edit/test_recraft_image_edit_transformation.py @@ -167,7 +167,7 @@ class TestRecraftImageEditTransformation: mock_response.status_code = 500 mock_response.headers = {} - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Error transforming image edit response: Invalid JSON: line') as exc_info: self.config.transform_image_edit_response( model=self.model, raw_response=mock_response, diff --git a/tests/test_litellm/llms/recraft/image_generation/test_recraft_image_gen_transformation.py b/tests/test_litellm/llms/recraft/image_generation/test_recraft_image_gen_transformation.py index 70311201969..ccc72dde7b8 100644 --- a/tests/test_litellm/llms/recraft/image_generation/test_recraft_image_gen_transformation.py +++ b/tests/test_litellm/llms/recraft/image_generation/test_recraft_image_gen_transformation.py @@ -64,7 +64,7 @@ class TestRecraftImageGenerationTransformation: non_default_params = {"n": 2, "unsupported_param": "value"} optional_params = {} - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Supported parameters are') as exc_info: self.config.map_openai_params( non_default_params=non_default_params, optional_params=optional_params, @@ -171,7 +171,7 @@ class TestRecraftImageGenerationTransformation: mock_get_secret.return_value = None headers = {} - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='RECRAFT_API_KEY is not set') as exc_info: self.config.validate_environment( headers=headers, model=self.model, @@ -248,7 +248,7 @@ class TestRecraftImageGenerationTransformation: model_response = ImageResponse(data=[]) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Error transforming image generation response: Invalid JSON') as exc_info: self.config.transform_image_generation_response( model=self.model, raw_response=mock_response, diff --git a/tests/test_litellm/llms/stability/image_generation/test_stability_image_generation.py b/tests/test_litellm/llms/stability/image_generation/test_stability_image_generation.py index 6f1a04e78d3..c5b3c8fbdc5 100644 --- a/tests/test_litellm/llms/stability/image_generation/test_stability_image_generation.py +++ b/tests/test_litellm/llms/stability/image_generation/test_stability_image_generation.py @@ -83,7 +83,7 @@ class TestStabilityImageGenerationConfig: non_default_params = {"unsupported_param": "value"} optional_params = {} - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match="Supported parameters are \\['n', 'size',") as exc_info: self.config.map_openai_params( non_default_params=non_default_params, optional_params=optional_params, @@ -168,7 +168,7 @@ class TestStabilityImageGenerationConfig: def test_validate_environment_raises_without_api_key(self): """Test that validate_environment raises error without API key""" - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='STABILITY_API_KEY is not set\\. Please set it via') as exc_info: self.config.validate_environment( headers={}, model="stability/sd3", @@ -251,7 +251,7 @@ class TestStabilityImageGenerationConfig: model_response = ImageResponse(data=[]) mock_logging = MagicMock() - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Content was filtered by Stability AI safety systems') as exc_info: self.config.transform_image_generation_response( model="stability/sd3", raw_response=mock_response, diff --git a/tests/test_litellm/llms/tinyfish/test_tinyfish_search.py b/tests/test_litellm/llms/tinyfish/test_tinyfish_search.py index 2dcccb8ea7e..69afbb416aa 100644 --- a/tests/test_litellm/llms/tinyfish/test_tinyfish_search.py +++ b/tests/test_litellm/llms/tinyfish/test_tinyfish_search.py @@ -697,7 +697,7 @@ class TestErrorHandling: } } mock_response = _make_mock_response(body, status_code=400) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='TinyFish Search: query is required\\. See https') as exc_info: config.transform_search_response( raw_response=mock_response, logging_obj=None ) @@ -713,7 +713,7 @@ class TestErrorHandling: mock_response = _make_mock_response( body, status_code=429, headers={"Retry-After": "60"} ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='TinyFish Search: rate limit exceeded\\. See https') as exc_info: config.transform_search_response( raw_response=mock_response, logging_obj=None ) @@ -728,7 +728,7 @@ class TestErrorHandling: config = TinyfishSearchConfig() body = {"errors": [{"code": "10000", "message": "Internal"}]} mock_response = _make_mock_response(body, status_code=502) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='TinyFish Search') as exc_info: config.transform_search_response( raw_response=mock_response, logging_obj=None ) @@ -742,7 +742,7 @@ class TestErrorHandling: mock_response = _make_mock_response( json_data=None, status_code=502, text="Bad Gateway" ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='TinyFish Search: Bad Gateway<') as exc_info: config.transform_search_response( raw_response=mock_response, logging_obj=None ) @@ -756,7 +756,7 @@ class TestErrorHandling: mock_response = _make_mock_response( json_data=None, status_code=200, text="not json" ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='TinyFish Search: Expected JSON response, got: not json\\.') as exc_info: config.transform_search_response( raw_response=mock_response, logging_obj=None ) @@ -785,7 +785,7 @@ class TestErrorHandling: # check TinyFish's schema, not their own input. config = TinyfishSearchConfig() mock_response = _make_mock_response({"query": "x"}) # no `results` key - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='validation error for SearchResponse') as exc_info: config.transform_search_response( raw_response=mock_response, logging_obj=None ) diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_integration.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_integration.py index 272565990bd..8f9acafa49d 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_integration.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_integration.py @@ -159,7 +159,7 @@ class TestVertexAIFilesIntegration: # This test ensures the type annotations and error messages include vertex_ai # Test that calling with unsupported provider raises appropriate error - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="unsupported_provider' is not a valid LlmProviders") as exc_info: litellm.file_content( file_id="test-file-id", custom_llm_provider="unsupported_provider", # This should fail diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py index b83d4742b64..c189cdd0ea7 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -33,7 +33,7 @@ def test_validate_vertex_location_accepts_valid(location): ["attacker.example/", "evil.com#", "us.attacker.example", "us/../..", "US", "us_central1", "-us", "", None], ) def test_validate_vertex_location_rejects_invalid(location): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="vertex_location is required|Invalid vertex_location format"): validate_vertex_location(location) diff --git a/tests/test_litellm/llms/volcengine/responses/test_volcengine_responses_transformation.py b/tests/test_litellm/llms/volcengine/responses/test_volcengine_responses_transformation.py index 891d1c15c61..7922331d19f 100644 --- a/tests/test_litellm/llms/volcengine/responses/test_volcengine_responses_transformation.py +++ b/tests/test_litellm/llms/volcengine/responses/test_volcengine_responses_transformation.py @@ -137,7 +137,7 @@ class TestVolcengineResponsesAPITransformation: monkeypatch.delenv("ARK_API_KEY", raising=False) monkeypatch.delenv("VOLCENGINE_API_KEY", raising=False) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Volcengine API key is required\\. Set ARK_API_KEY /'): config.validate_environment(headers={}, model="volcengine/demo", litellm_params={}) def test_unsupported_params_are_dropped_with_extra_body(self): diff --git a/tests/test_litellm/llms/volcengine/test_volcengine_embedding.py b/tests/test_litellm/llms/volcengine/test_volcengine_embedding.py index 07298f03f86..1670dac0e9d 100644 --- a/tests/test_litellm/llms/volcengine/test_volcengine_embedding.py +++ b/tests/test_litellm/llms/volcengine/test_volcengine_embedding.py @@ -202,7 +202,7 @@ def test_volcengine_embedding_error_scenarios(): k: v for k, v in scenario.items() if k != "expected_error_pattern" } - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match=f"(?i){scenario['expected_error_pattern']}") as exc_info: litellm.embedding(input=["test"], **test_params) # Verify error message contains expected pattern diff --git a/tests/test_litellm/llms/voyage/rerank/test_voyage_rerank_transformation.py b/tests/test_litellm/llms/voyage/rerank/test_voyage_rerank_transformation.py index 8f99609e3f5..f466b7e19b5 100644 --- a/tests/test_litellm/llms/voyage/rerank/test_voyage_rerank_transformation.py +++ b/tests/test_litellm/llms/voyage/rerank/test_voyage_rerank_transformation.py @@ -227,7 +227,7 @@ class TestVoyageRerankTransform: mock_logging = MagicMock() model_response = RerankResponse() - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Unauthorized') as exc_info: self.config.transform_rerank_response( model=self.model, raw_response=mock_response, @@ -248,7 +248,7 @@ class TestVoyageRerankTransform: mock_logging = MagicMock() model_response = RerankResponse() - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Failed to parse response: Invalid JSON response') as exc_info: self.config.transform_rerank_response( model=self.model, raw_response=mock_response, diff --git a/tests/test_litellm/llms/voyage/test_voyage_multimodal_embedding.py b/tests/test_litellm/llms/voyage/test_voyage_multimodal_embedding.py index f283e7fe0df..f3e6885cbe6 100644 --- a/tests/test_litellm/llms/voyage/test_voyage_multimodal_embedding.py +++ b/tests/test_litellm/llms/voyage/test_voyage_multimodal_embedding.py @@ -195,7 +195,7 @@ class TestVoyageMultimodalEmbeddings: monkeypatch.setattr(module, "get_secret_str", lambda name: None) config = VoyageMultimodalEmbeddingConfig() - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Voyage API key is required for multimodal embeddings\\. Set') as exc_info: config.validate_environment( {}, "voyage-multimodal-3.5", [], {}, {}, api_key=None ) @@ -207,7 +207,7 @@ class TestVoyageMultimodalEmbeddings: ) config = VoyageMultimodalEmbeddingConfig() - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Voyage multimodal embeddings require a non-empty') as exc_info: config._normalize_content_item({"type": "image_url", "image_url": {}}) assert "image_url" in str(exc_info.value) diff --git a/tests/test_litellm/llms/xai/test_xai_key_fallback.py b/tests/test_litellm/llms/xai/test_xai_key_fallback.py index 4c769c572ac..ec3eb83309c 100644 --- a/tests/test_litellm/llms/xai/test_xai_key_fallback.py +++ b/tests/test_litellm/llms/xai/test_xai_key_fallback.py @@ -168,7 +168,7 @@ def test_responses_config_raises_when_no_key_is_available(monkeypatch): monkeypatch.setattr(litellm, "api_key", None) monkeypatch.delenv("XAI_API_KEY", raising=False) - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='XAI API key is required\\. Set api_key, litellm\\.xai_key') as exc_info: XAIResponsesAPIConfig().validate_environment({}, "xai/grok-3-mini", None) error_message = str(exc_info.value) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 0209abee510..d5936b2ae86 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -8185,7 +8185,7 @@ class TestGetUserObjectPermission: return_value=None, ), ): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="user 'human-dangling' names object_permission_id"): await MCPRequestHandler._get_user_object_permission(auth) async def test_no_user_id_places_no_ceiling(self): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 34852850de6..b4d3782ba43 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -2695,13 +2695,6 @@ async def test_token_endpoint_respects_x_forwarded_host(): "443", "https://internal.local", ), - ( - "http://localhost:4000/", - "https", - "proxy.example.com", - "8443", - "https://proxy.example.com:8443", - ), ( "http://localhost:4000/", "https", diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py index 6e3ac014840..941e5deee93 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py @@ -79,7 +79,7 @@ class TestShortPrefixHelpers: assert compute_short_server_prefix("abc") != compute_short_server_prefix("abd") def test_short_prefix_requires_server_id(self): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='compute_short_server_prefix requires a non-empty server_id'): compute_short_server_prefix("") def test_flag_defaults_to_false(self): diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 6b40fa1b324..762d2cbf3c7 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -268,7 +268,7 @@ def test_get_experimental_ui_login_jwt_auth_token_invalid( invalid_sso_user_defined_values, ): """Test generating JWT token with missing user role""" - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='User role is required for experimental UI login') as exc_info: ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token( invalid_sso_user_defined_values ) @@ -883,7 +883,7 @@ async def test_get_user_object_wraps_db_outage_as_valueerror_preserving_context( mock_cache.async_set_cache = AsyncMock() with patch("litellm.proxy.auth.auth_checks._should_check_db", return_value=True): - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match="User doesn't exist in db\\.") as exc_info: await get_user_object( user_id="outage-contract-probe-user", prisma_client=mock_prisma_client, diff --git a/tests/test_litellm/proxy/auth/test_auth_hot_path_network_requests.py b/tests/test_litellm/proxy/auth/test_auth_hot_path_network_requests.py index 22752f767ce..6b2d2babedc 100644 --- a/tests/test_litellm/proxy/auth/test_auth_hot_path_network_requests.py +++ b/tests/test_litellm/proxy/auth/test_auth_hot_path_network_requests.py @@ -1,521 +1,521 @@ -""" -Test to count and track the number of network requests (DB queries, cache lookups) -made on the hot path for keys that have team_id and user_id attached. - -This test ensures we don't regress on the number of network requests made during -request authentication, which directly impacts proxy latency. - -The hot path covers auth functions called on every LLM API request: -- get_key_object: lookup the API key -- get_team_object: lookup the team (for keys with team_id) -- get_user_object: lookup the user (for keys with user_id) -- get_team_membership: lookup team member budget (when team_member_spend set) - -Each function does: cache read -> (on miss) DB query -> cache write. -We count these to catch regressions in the number of network requests. - -NOTE: This test does NOT require proxy extras (apscheduler, etc.) because -it tests at the auth_checks level, not the full proxy_server level. -""" - -import os -import sys -import time -from typing import Any, Dict, List, Optional -from unittest.mock import AsyncMock, MagicMock - -import pytest - -sys.path.insert(0, os.path.abspath("../../..")) - -from litellm.caching.dual_cache import DualCache -from litellm.caching.in_memory_cache import InMemoryCache -from litellm.proxy._types import ( - LiteLLM_TeamTableCachedObj, - LiteLLM_UserTable, - LitellmUserRoles, - UserAPIKeyAuth, - LiteLLM_TeamMembership, - hash_token, -) -from litellm.proxy.auth.auth_checks import ( - get_key_object, - get_team_membership, - get_team_object, - get_user_object, -) - - -class CacheCallTracker: - """ - Tracks cache read/write operations by wrapping DualCache methods. - This is used to count network-level operations on the hot path. - """ - - def __init__(self): - self.cache_reads: List[Dict[str, Any]] = [] - self.cache_writes: List[Dict[str, Any]] = [] - self.db_queries: List[Dict[str, Any]] = [] - - def get_summary(self) -> Dict[str, Any]: - return { - "total_cache_reads": len(self.cache_reads), - "total_cache_writes": len(self.cache_writes), - "total_db_queries": len(self.db_queries), - "total_network_requests": len(self.cache_reads) - + len(self.cache_writes) - + len(self.db_queries), - "cache_read_keys": [r["key"] for r in self.cache_reads], - "cache_write_keys": [w["key"] for w in self.cache_writes], - "db_query_details": self.db_queries, - } - - -def _wrap_cache_with_tracker(cache: DualCache, tracker: CacheCallTracker) -> DualCache: - """Wrap a DualCache to track all reads and writes.""" - original_async_get = cache.async_get_cache - original_async_set = cache.async_set_cache - - async def tracked_async_get(key, *args, **kwargs): - result = await original_async_get(key, *args, **kwargs) - tracker.cache_reads.append( - {"key": key, "hit": result is not None, "method": "async_get_cache"} - ) - return result - - async def tracked_async_set(key, value, *args, **kwargs): - tracker.cache_writes.append({"key": key, "method": "async_set_cache"}) - return await original_async_set(key, value, *args, **kwargs) - - cache.async_get_cache = tracked_async_get - cache.async_set_cache = tracked_async_set - return cache - - -def _create_valid_token( - api_key: str, - team_id: str, - user_id: str, - has_team_member_spend: bool = False, - org_id: Optional[str] = None, -) -> UserAPIKeyAuth: - """Create a UserAPIKeyAuth with team_id and user_id set.""" - hashed = hash_token(api_key) - return UserAPIKeyAuth( - token=hashed, - api_key=api_key, - team_id=team_id, - user_id=user_id, - org_id=org_id, - models=["gpt-4", "gpt-3.5-turbo"], - max_budget=100.0, - spend=10.0, - team_spend=50.0, - team_max_budget=1000.0, - team_models=["gpt-4", "gpt-3.5-turbo"], - team_member_spend=5.0 if has_team_member_spend else None, - last_refreshed_at=time.time(), - user_role=LitellmUserRoles.INTERNAL_USER, - ) - - -def _create_team_object(team_id: str) -> LiteLLM_TeamTableCachedObj: - """Create a team table object for caching.""" - return LiteLLM_TeamTableCachedObj( - team_id=team_id, - models=["gpt-4", "gpt-3.5-turbo"], - max_budget=1000.0, - spend=50.0, - tpm_limit=10000, - rpm_limit=100, - last_refreshed_at=time.time(), - ) - - -def _create_user_object(user_id: str) -> LiteLLM_UserTable: - """Create a user table object for caching.""" - return LiteLLM_UserTable( - user_id=user_id, - max_budget=500.0, - spend=25.0, - models=["gpt-4"], - tpm_limit=5000, - rpm_limit=50, - user_role=LitellmUserRoles.INTERNAL_USER, - user_email="test@example.com", - ) - - -# ============================================================================ -# TEST: get_key_object cache behavior -# ============================================================================ - - -@pytest.mark.asyncio -async def test_get_key_object_warm_cache(): - """ - Test get_key_object with a warm cache - should hit cache, no DB query. - """ - api_key = "sk-test-key-warm" - team_id = "team-123" - user_id = "user-456" - hashed_token = hash_token(api_key) - - valid_token = _create_valid_token(api_key, team_id, user_id) - - # Create cache with pre-populated data - cache = DualCache(in_memory_cache=InMemoryCache()) - await cache.async_set_cache(key=hashed_token, value=valid_token) - - # Track cache operations - tracker = CacheCallTracker() - tracked_cache = _wrap_cache_with_tracker(cache, tracker) - - # Mock prisma client (should NOT be called for warm cache) - mock_prisma = MagicMock() - mock_prisma.get_data = AsyncMock() - - result = await get_key_object( - hashed_token=hashed_token, - prisma_client=mock_prisma, - user_api_key_cache=tracked_cache, - parent_otel_span=None, - proxy_logging_obj=None, - ) - - summary = tracker.get_summary() - - # Should have exactly 1 cache read - assert summary["total_cache_reads"] == 1 - assert hashed_token in summary["cache_read_keys"] - - # Prisma should NOT have been called - mock_prisma.get_data.assert_not_called() - - # Result should be the cached token - assert result.token == hashed_token - - -@pytest.mark.asyncio -async def test_get_key_object_cold_cache(): - """ - Test get_key_object with a cold cache - should miss cache, query DB. - """ - api_key = "sk-test-key-cold" - team_id = "team-123" - user_id = "user-456" - hashed_token = hash_token(api_key) - - valid_token = _create_valid_token(api_key, team_id, user_id) - - # Create empty cache - cache = DualCache(in_memory_cache=InMemoryCache()) - - tracker = CacheCallTracker() - tracked_cache = _wrap_cache_with_tracker(cache, tracker) - - # Mock prisma client to return token on DB query - mock_prisma = MagicMock() - mock_prisma.get_data = AsyncMock(return_value=valid_token) - - await get_key_object( - hashed_token=hashed_token, - prisma_client=mock_prisma, - user_api_key_cache=tracked_cache, - parent_otel_span=None, - proxy_logging_obj=None, - ) - - summary = tracker.get_summary() - - # Should have 1 cache read (miss) and at least 1 cache write (populate cache) - assert summary["total_cache_reads"] >= 1 - - # Prisma SHOULD have been called - mock_prisma.get_data.assert_called_once() - - -# ============================================================================ -# TEST: get_team_object cache behavior -# ============================================================================ - - -@pytest.mark.asyncio -async def test_get_team_object_warm_cache(): - """ - Test get_team_object with a warm cache - should hit cache, no DB query. - """ - team_id = "team-warm-123" - team_obj = _create_team_object(team_id) - - cache = DualCache(in_memory_cache=InMemoryCache()) - cache_key = f"team_id:{team_id}" - await cache.async_set_cache(key=cache_key, value=team_obj) - - tracker = CacheCallTracker() - tracked_cache = _wrap_cache_with_tracker(cache, tracker) - - mock_prisma = MagicMock() - mock_prisma.db = MagicMock() - mock_prisma.db.litellm_teamtable = MagicMock() - mock_prisma.db.litellm_teamtable.find_unique = AsyncMock() - - await get_team_object( - team_id=team_id, - prisma_client=mock_prisma, - user_api_key_cache=tracked_cache, - parent_otel_span=None, - proxy_logging_obj=None, - ) - - summary = tracker.get_summary() - - assert summary["total_cache_reads"] >= 1 - assert cache_key in summary["cache_read_keys"] - - # DB should NOT have been called - mock_prisma.db.litellm_teamtable.find_unique.assert_not_called() - - -# ============================================================================ -# TEST: get_user_object cache behavior -# ============================================================================ - - -@pytest.mark.asyncio -async def test_get_user_object_warm_cache(): - """ - Test get_user_object with a warm cache - should hit cache, no DB query. - """ - user_id = "user-warm-456" - user_obj = _create_user_object(user_id) - - cache = DualCache(in_memory_cache=InMemoryCache()) - await cache.async_set_cache(key=user_id, value=user_obj) - - tracker = CacheCallTracker() - tracked_cache = _wrap_cache_with_tracker(cache, tracker) - - mock_prisma = MagicMock() - mock_prisma.db = MagicMock() - mock_prisma.db.litellm_usertable = MagicMock() - mock_prisma.db.litellm_usertable.find_unique = AsyncMock() - - await get_user_object( - user_id=user_id, - prisma_client=mock_prisma, - user_api_key_cache=tracked_cache, - parent_otel_span=None, - proxy_logging_obj=None, - user_id_upsert=False, - ) - - summary = tracker.get_summary() - - assert summary["total_cache_reads"] >= 1 - assert user_id in summary["cache_read_keys"] - - # DB should NOT have been called - mock_prisma.db.litellm_usertable.find_unique.assert_not_called() - - -# ============================================================================ -# TEST: get_team_membership cache behavior -# ============================================================================ - - -@pytest.mark.asyncio -async def test_get_team_membership_warm_cache(): - """ - Test get_team_membership with a warm cache - should hit cache, no DB query. - """ - user_id = "user-tm-456" - team_id = "team-tm-123" - - membership_dict = { - "user_id": user_id, - "team_id": team_id, - "spend": 3.0, - "budget_id": None, - "litellm_budget_table": None, - } - - cache = DualCache(in_memory_cache=InMemoryCache()) - # Cache key format used by get_team_membership - cache_key = f"team_membership:{user_id}:{team_id}" - await cache.async_set_cache(key=cache_key, value=membership_dict) - - tracker = CacheCallTracker() - tracked_cache = _wrap_cache_with_tracker(cache, tracker) - - mock_prisma = MagicMock() - mock_prisma.db = MagicMock() - mock_prisma.db.litellm_teammembership = MagicMock() - mock_prisma.db.litellm_teammembership.find_unique = AsyncMock() - - await get_team_membership( - user_id=user_id, - team_id=team_id, - prisma_client=mock_prisma, - user_api_key_cache=tracked_cache, - parent_otel_span=None, - proxy_logging_obj=None, - ) - - summary = tracker.get_summary() - - assert summary["total_cache_reads"] >= 1 - assert cache_key in summary["cache_read_keys"] - - # DB should NOT have been called - mock_prisma.db.litellm_teammembership.find_unique.assert_not_called() - - -# ============================================================================ -# TEST: Document duplicate team membership cache key issue -# ============================================================================ - - -@pytest.mark.asyncio -async def test_team_membership_cache_key_duplication(): - """ - Document the team membership duplicate cache key issue: - - Team membership is queried via TWO different cache keys: - 1. "{team_id}_{user_id}" - used in user_api_key_auth.py:1048 - 2. "team_membership:{user_id}:{team_id}" - used in auth_checks.py:960 (get_team_membership) - - This test documents that both keys refer to the same data but use different - cache key formats, potentially leading to duplicate lookups. - """ - user_id = "user-dup-456" - team_id = "team-dup-123" - - # The two different cache keys used for the same data - key_format_1 = f"{team_id}_{user_id}" # user_api_key_auth format - key_format_2 = f"team_membership:{user_id}:{team_id}" # auth_checks format - - _ = { - "user_id": user_id, - "team_id": team_id, - "spend": 3.0, - } - - # Document that these are different keys - assert ( - key_format_1 != key_format_2 - ), "Cache keys should be different (this is the bug)" - - # Document that these are different keys - assert ( - key_format_1 != key_format_2 - ), "Cache keys should be different (this is the bug)" - - -# ============================================================================ -# TEST: Full hot path network count summary -# ============================================================================ - - -@pytest.mark.asyncio -async def test_full_hot_path_network_count(): - """ - Summary test that counts all network operations when processing - a request with a key that has team_id and user_id attached. - - This test verifies the baseline number of cache operations expected - on a fully warm cache path. - """ - api_key = "sk-test-full-path" - team_id = "team-full-123" - user_id = "user-full-456" - hashed_token = hash_token(api_key) - - # Create all objects - valid_token = _create_valid_token( - api_key, team_id, user_id, has_team_member_spend=True - ) - team_obj = _create_team_object(team_id) - user_obj = _create_user_object(user_id) - membership_data = LiteLLM_TeamMembership( - user_id=user_id, - team_id=team_id, - spend=3.0, - budget_id=None, - litellm_budget_table=None, - ) - - # Pre-populate cache with all data - cache = DualCache(in_memory_cache=InMemoryCache()) - await cache.async_set_cache(key=hashed_token, value=valid_token) - await cache.async_set_cache(key=f"team_id:{team_id}", value=team_obj) - await cache.async_set_cache(key=user_id, value=user_obj) - await cache.async_set_cache( - key=f"team_membership:{user_id}:{team_id}", value=membership_data.model_dump() - ) - await cache.async_set_cache( - key=f"{team_id}_{user_id}", value=membership_data.model_dump() - ) - - # Create tracker AFTER populating cache - tracker = CacheCallTracker() - tracked_cache = _wrap_cache_with_tracker(cache, tracker) - - # Mock prisma (should not be called on warm cache) - mock_prisma = MagicMock() - - # Call each function to simulate the hot path - await get_key_object( - hashed_token=hashed_token, - prisma_client=mock_prisma, - user_api_key_cache=tracked_cache, - parent_otel_span=None, - proxy_logging_obj=None, - ) - - await get_team_object( - team_id=team_id, - prisma_client=mock_prisma, - user_api_key_cache=tracked_cache, - parent_otel_span=None, - proxy_logging_obj=None, - ) - - await get_user_object( - user_id=user_id, - prisma_client=mock_prisma, - user_api_key_cache=tracked_cache, - parent_otel_span=None, - proxy_logging_obj=None, - user_id_upsert=False, - ) - - await get_team_membership( - user_id=user_id, - team_id=team_id, - prisma_client=mock_prisma, - user_api_key_cache=tracked_cache, - parent_otel_span=None, - proxy_logging_obj=None, - ) - - summary = tracker.get_summary() - - # Assertions for expected baseline - # On warm cache: 4 reads (key, team, user, team_membership) - assert ( - summary["total_cache_reads"] == 4 - ), f"Expected 4 cache reads on warm path, got {summary['total_cache_reads']}" - - # No DB queries on warm cache - assert ( - summary["total_db_queries"] == 0 - ), f"Expected 0 DB queries on warm path, got {summary['total_db_queries']}" - - # Total network requests should be exactly 4 on warm cache - assert ( - summary["total_network_requests"] == 4 - ), f"Expected 4 total network requests on warm path, got {summary['total_network_requests']}" +""" +Test to count and track the number of network requests (DB queries, cache lookups) +made on the hot path for keys that have team_id and user_id attached. + +This test ensures we don't regress on the number of network requests made during +request authentication, which directly impacts proxy latency. + +The hot path covers auth functions called on every LLM API request: +- get_key_object: lookup the API key +- get_team_object: lookup the team (for keys with team_id) +- get_user_object: lookup the user (for keys with user_id) +- get_team_membership: lookup team member budget (when team_member_spend set) + +Each function does: cache read -> (on miss) DB query -> cache write. +We count these to catch regressions in the number of network requests. + +NOTE: This test does NOT require proxy extras (apscheduler, etc.) because +it tests at the auth_checks level, not the full proxy_server level. +""" + +import os +import sys +import time +from typing import Any, Dict, List, Optional +from unittest.mock import AsyncMock, MagicMock + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.proxy._types import ( + LiteLLM_TeamTableCachedObj, + LiteLLM_UserTable, + LitellmUserRoles, + UserAPIKeyAuth, + LiteLLM_TeamMembership, + hash_token, +) +from litellm.proxy.auth.auth_checks import ( + get_key_object, + get_team_membership, + get_team_object, + get_user_object, +) + + +class CacheCallTracker: + """ + Tracks cache read/write operations by wrapping DualCache methods. + This is used to count network-level operations on the hot path. + """ + + def __init__(self): + self.cache_reads: List[Dict[str, Any]] = [] + self.cache_writes: List[Dict[str, Any]] = [] + self.db_queries: List[Dict[str, Any]] = [] + + def get_summary(self) -> Dict[str, Any]: + return { + "total_cache_reads": len(self.cache_reads), + "total_cache_writes": len(self.cache_writes), + "total_db_queries": len(self.db_queries), + "total_network_requests": len(self.cache_reads) + + len(self.cache_writes) + + len(self.db_queries), + "cache_read_keys": [r["key"] for r in self.cache_reads], + "cache_write_keys": [w["key"] for w in self.cache_writes], + "db_query_details": self.db_queries, + } + + +def _wrap_cache_with_tracker(cache: DualCache, tracker: CacheCallTracker) -> DualCache: + """Wrap a DualCache to track all reads and writes.""" + original_async_get = cache.async_get_cache + original_async_set = cache.async_set_cache + + async def tracked_async_get(key, *args, **kwargs): + result = await original_async_get(key, *args, **kwargs) + tracker.cache_reads.append( + {"key": key, "hit": result is not None, "method": "async_get_cache"} + ) + return result + + async def tracked_async_set(key, value, *args, **kwargs): + tracker.cache_writes.append({"key": key, "method": "async_set_cache"}) + return await original_async_set(key, value, *args, **kwargs) + + cache.async_get_cache = tracked_async_get + cache.async_set_cache = tracked_async_set + return cache + + +def _create_valid_token( + api_key: str, + team_id: str, + user_id: str, + has_team_member_spend: bool = False, + org_id: Optional[str] = None, +) -> UserAPIKeyAuth: + """Create a UserAPIKeyAuth with team_id and user_id set.""" + hashed = hash_token(api_key) + return UserAPIKeyAuth( + token=hashed, + api_key=api_key, + team_id=team_id, + user_id=user_id, + org_id=org_id, + models=["gpt-4", "gpt-3.5-turbo"], + max_budget=100.0, + spend=10.0, + team_spend=50.0, + team_max_budget=1000.0, + team_models=["gpt-4", "gpt-3.5-turbo"], + team_member_spend=5.0 if has_team_member_spend else None, + last_refreshed_at=time.time(), + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + +def _create_team_object(team_id: str) -> LiteLLM_TeamTableCachedObj: + """Create a team table object for caching.""" + return LiteLLM_TeamTableCachedObj( + team_id=team_id, + models=["gpt-4", "gpt-3.5-turbo"], + max_budget=1000.0, + spend=50.0, + tpm_limit=10000, + rpm_limit=100, + last_refreshed_at=time.time(), + ) + + +def _create_user_object(user_id: str) -> LiteLLM_UserTable: + """Create a user table object for caching.""" + return LiteLLM_UserTable( + user_id=user_id, + max_budget=500.0, + spend=25.0, + models=["gpt-4"], + tpm_limit=5000, + rpm_limit=50, + user_role=LitellmUserRoles.INTERNAL_USER, + user_email="test@example.com", + ) + + +# ============================================================================ +# TEST: get_key_object cache behavior +# ============================================================================ + + +@pytest.mark.asyncio +async def test_get_key_object_warm_cache(): + """ + Test get_key_object with a warm cache - should hit cache, no DB query. + """ + api_key = "sk-test-key-warm" + team_id = "team-123" + user_id = "user-456" + hashed_token = hash_token(api_key) + + valid_token = _create_valid_token(api_key, team_id, user_id) + + # Create cache with pre-populated data + cache = DualCache(in_memory_cache=InMemoryCache()) + await cache.async_set_cache(key=hashed_token, value=valid_token) + + # Track cache operations + tracker = CacheCallTracker() + tracked_cache = _wrap_cache_with_tracker(cache, tracker) + + # Mock prisma client (should NOT be called for warm cache) + mock_prisma = MagicMock() + mock_prisma.get_data = AsyncMock() + + result = await get_key_object( + hashed_token=hashed_token, + prisma_client=mock_prisma, + user_api_key_cache=tracked_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + summary = tracker.get_summary() + + # Should have exactly 1 cache read + assert summary["total_cache_reads"] == 1 + assert hashed_token in summary["cache_read_keys"] + + # Prisma should NOT have been called + mock_prisma.get_data.assert_not_called() + + # Result should be the cached token + assert result.token == hashed_token + + +@pytest.mark.asyncio +async def test_get_key_object_cold_cache(): + """ + Test get_key_object with a cold cache - should miss cache, query DB. + """ + api_key = "sk-test-key-cold" + team_id = "team-123" + user_id = "user-456" + hashed_token = hash_token(api_key) + + valid_token = _create_valid_token(api_key, team_id, user_id) + + # Create empty cache + cache = DualCache(in_memory_cache=InMemoryCache()) + + tracker = CacheCallTracker() + tracked_cache = _wrap_cache_with_tracker(cache, tracker) + + # Mock prisma client to return token on DB query + mock_prisma = MagicMock() + mock_prisma.get_data = AsyncMock(return_value=valid_token) + + await get_key_object( + hashed_token=hashed_token, + prisma_client=mock_prisma, + user_api_key_cache=tracked_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + summary = tracker.get_summary() + + # Should have 1 cache read (miss) and at least 1 cache write (populate cache) + assert summary["total_cache_reads"] >= 1 + + # Prisma SHOULD have been called + mock_prisma.get_data.assert_called_once() + + +# ============================================================================ +# TEST: get_team_object cache behavior +# ============================================================================ + + +@pytest.mark.asyncio +async def test_get_team_object_warm_cache(): + """ + Test get_team_object with a warm cache - should hit cache, no DB query. + """ + team_id = "team-warm-123" + team_obj = _create_team_object(team_id) + + cache = DualCache(in_memory_cache=InMemoryCache()) + cache_key = f"team_id:{team_id}" + await cache.async_set_cache(key=cache_key, value=team_obj) + + tracker = CacheCallTracker() + tracked_cache = _wrap_cache_with_tracker(cache, tracker) + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.litellm_teamtable = MagicMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock() + + await get_team_object( + team_id=team_id, + prisma_client=mock_prisma, + user_api_key_cache=tracked_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + summary = tracker.get_summary() + + assert summary["total_cache_reads"] >= 1 + assert cache_key in summary["cache_read_keys"] + + # DB should NOT have been called + mock_prisma.db.litellm_teamtable.find_unique.assert_not_called() + + +# ============================================================================ +# TEST: get_user_object cache behavior +# ============================================================================ + + +@pytest.mark.asyncio +async def test_get_user_object_warm_cache(): + """ + Test get_user_object with a warm cache - should hit cache, no DB query. + """ + user_id = "user-warm-456" + user_obj = _create_user_object(user_id) + + cache = DualCache(in_memory_cache=InMemoryCache()) + await cache.async_set_cache(key=user_id, value=user_obj) + + tracker = CacheCallTracker() + tracked_cache = _wrap_cache_with_tracker(cache, tracker) + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.litellm_usertable = MagicMock() + mock_prisma.db.litellm_usertable.find_unique = AsyncMock() + + await get_user_object( + user_id=user_id, + prisma_client=mock_prisma, + user_api_key_cache=tracked_cache, + parent_otel_span=None, + proxy_logging_obj=None, + user_id_upsert=False, + ) + + summary = tracker.get_summary() + + assert summary["total_cache_reads"] >= 1 + assert user_id in summary["cache_read_keys"] + + # DB should NOT have been called + mock_prisma.db.litellm_usertable.find_unique.assert_not_called() + + +# ============================================================================ +# TEST: get_team_membership cache behavior +# ============================================================================ + + +@pytest.mark.asyncio +async def test_get_team_membership_warm_cache(): + """ + Test get_team_membership with a warm cache - should hit cache, no DB query. + """ + user_id = "user-tm-456" + team_id = "team-tm-123" + + membership_dict = { + "user_id": user_id, + "team_id": team_id, + "spend": 3.0, + "budget_id": None, + "litellm_budget_table": None, + } + + cache = DualCache(in_memory_cache=InMemoryCache()) + # Cache key format used by get_team_membership + cache_key = f"team_membership:{user_id}:{team_id}" + await cache.async_set_cache(key=cache_key, value=membership_dict) + + tracker = CacheCallTracker() + tracked_cache = _wrap_cache_with_tracker(cache, tracker) + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.litellm_teammembership = MagicMock() + mock_prisma.db.litellm_teammembership.find_unique = AsyncMock() + + await get_team_membership( + user_id=user_id, + team_id=team_id, + prisma_client=mock_prisma, + user_api_key_cache=tracked_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + summary = tracker.get_summary() + + assert summary["total_cache_reads"] >= 1 + assert cache_key in summary["cache_read_keys"] + + # DB should NOT have been called + mock_prisma.db.litellm_teammembership.find_unique.assert_not_called() + + +# ============================================================================ +# TEST: Document duplicate team membership cache key issue +# ============================================================================ + + +@pytest.mark.asyncio +async def test_team_membership_cache_key_duplication(): + """ + Document the team membership duplicate cache key issue: + + Team membership is queried via TWO different cache keys: + 1. "{team_id}_{user_id}" - used in user_api_key_auth.py:1048 + 2. "team_membership:{user_id}:{team_id}" - used in auth_checks.py:960 (get_team_membership) + + This test documents that both keys refer to the same data but use different + cache key formats, potentially leading to duplicate lookups. + """ + user_id = "user-dup-456" + team_id = "team-dup-123" + + # The two different cache keys used for the same data + key_format_1 = f"{team_id}_{user_id}" # user_api_key_auth format + key_format_2 = f"team_membership:{user_id}:{team_id}" # auth_checks format + + _ = { + "user_id": user_id, + "team_id": team_id, + "spend": 3.0, + } + + # Document that these are different keys + assert ( + key_format_1 != key_format_2 + ), "Cache keys should be different (this is the bug)" + + # Document that these are different keys + assert ( + key_format_1 != key_format_2 + ), "Cache keys should be different (this is the bug)" + + +# ============================================================================ +# TEST: Full hot path network count summary +# ============================================================================ + + +@pytest.mark.asyncio +async def test_full_hot_path_network_count(): + """ + Summary test that counts all network operations when processing + a request with a key that has team_id and user_id attached. + + This test verifies the baseline number of cache operations expected + on a fully warm cache path. + """ + api_key = "sk-test-full-path" + team_id = "team-full-123" + user_id = "user-full-456" + hashed_token = hash_token(api_key) + + # Create all objects + valid_token = _create_valid_token( + api_key, team_id, user_id, has_team_member_spend=True + ) + team_obj = _create_team_object(team_id) + user_obj = _create_user_object(user_id) + membership_data = LiteLLM_TeamMembership( + user_id=user_id, + team_id=team_id, + spend=3.0, + budget_id=None, + litellm_budget_table=None, + ) + + # Pre-populate cache with all data + cache = DualCache(in_memory_cache=InMemoryCache()) + await cache.async_set_cache(key=hashed_token, value=valid_token) + await cache.async_set_cache(key=f"team_id:{team_id}", value=team_obj) + await cache.async_set_cache(key=user_id, value=user_obj) + await cache.async_set_cache( + key=f"team_membership:{user_id}:{team_id}", value=membership_data.model_dump() + ) + await cache.async_set_cache( + key=f"{team_id}_{user_id}", value=membership_data.model_dump() + ) + + # Create tracker AFTER populating cache + tracker = CacheCallTracker() + tracked_cache = _wrap_cache_with_tracker(cache, tracker) + + # Mock prisma (should not be called on warm cache) + mock_prisma = MagicMock() + + # Call each function to simulate the hot path + await get_key_object( + hashed_token=hashed_token, + prisma_client=mock_prisma, + user_api_key_cache=tracked_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + await get_team_object( + team_id=team_id, + prisma_client=mock_prisma, + user_api_key_cache=tracked_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + await get_user_object( + user_id=user_id, + prisma_client=mock_prisma, + user_api_key_cache=tracked_cache, + parent_otel_span=None, + proxy_logging_obj=None, + user_id_upsert=False, + ) + + await get_team_membership( + user_id=user_id, + team_id=team_id, + prisma_client=mock_prisma, + user_api_key_cache=tracked_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + summary = tracker.get_summary() + + # Assertions for expected baseline + # On warm cache: 4 reads (key, team, user, team_membership) + assert ( + summary["total_cache_reads"] == 4 + ), f"Expected 4 cache reads on warm path, got {summary['total_cache_reads']}" + + # No DB queries on warm cache + assert ( + summary["total_db_queries"] == 0 + ), f"Expected 0 DB queries on warm path, got {summary['total_db_queries']}" + + # Total network requests should be exactly 4 on warm cache + assert ( + summary["total_network_requests"] == 4 + ), f"Expected 4 total network requests on warm path, got {summary['total_network_requests']}" # ============================================================================ @@ -540,7 +540,7 @@ async def test_get_user_object_missing_user_negative_cache(): mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) for _ in range(3): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="User doesn't exist in db\\."): await get_user_object( user_id=user_id, prisma_client=mock_prisma, @@ -570,7 +570,7 @@ async def test_get_user_object_missing_user_rechecks_after_expiry(): mock_prisma.db.litellm_usertable = MagicMock() mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="User doesn't exist in db\\."): await get_user_object( user_id=user_id, prisma_client=mock_prisma, @@ -586,7 +586,7 @@ async def test_get_user_object_missing_user_rechecks_after_expiry(): time.time() - (db_cache_expiry + 1), ) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="User doesn't exist in db\\."): await get_user_object( user_id=user_id, prisma_client=mock_prisma, diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index ecf7f89d487..9301176f3ed 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -1588,7 +1588,7 @@ class TestCheckCompleteCredentialsBlocksSSRF: "litellm.proxy.auth.auth_utils.validate_url", side_effect=SSRFError(f"blocked: {blocked_url}"), ): - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='is rejected by the SSRF guard') as exc_info: check_complete_credentials( { "model": "gpt-4", @@ -2144,7 +2144,7 @@ class TestIsRequestBodySafeBlocksEndpointTargetingFields: ], ) def test_endpoint_targeting_field_in_request_body_is_rejected(self, field): - with pytest.raises(ValueError) as exc: + with pytest.raises(ValueError, match='Rejected Request') as exc: is_request_body_safe( request_body={"model": "gpt-4", field: "https://attacker.example"}, general_settings={}, @@ -2165,7 +2165,7 @@ class TestIsRequestBodySafeBlocksEndpointTargetingFields: # on the blocklist into an SSRF / credential-exfil hole. Verify # that supplying an api_key (alongside the banned param) does NOT # bypass the gate — it can only be opened by an admin opt-in. - with pytest.raises(ValueError) as exc: + with pytest.raises(ValueError, match='Rejected Request') as exc: is_request_body_safe( request_body={ "model": "gpt-4", @@ -2722,7 +2722,7 @@ class TestObservabilityCallbackBans: ], ) def test_observability_field_in_request_body_root_is_rejected(self, field): - with pytest.raises(ValueError) as exc: + with pytest.raises(ValueError, match='Rejected Request') as exc: is_request_body_safe( request_body={"model": "gpt-4", field: "attacker-value"}, general_settings={}, @@ -2752,7 +2752,7 @@ class TestObservabilityCallbackBans: # Verifies the metadata walk: a value smuggled inside ``metadata`` # or ``litellm_metadata`` is just as dangerous as the same field # at the body root, and must hit the same gate. - with pytest.raises(ValueError) as exc: + with pytest.raises(ValueError, match='Rejected Request') as exc: is_request_body_safe( request_body={ "model": "gpt-4", @@ -2787,7 +2787,7 @@ class TestObservabilityCallbackBans: ) def test_observability_field_in_litellm_params_metadata_is_rejected(self): - with pytest.raises(ValueError) as exc: + with pytest.raises(ValueError, match='Rejected Request: turn_off_message_logging is not allowed') as exc: is_request_body_safe( request_body={ "model": "gpt-4", @@ -2814,7 +2814,7 @@ class TestObservabilityCallbackBans: # the ``isinstance(dict)`` guard. import json - with pytest.raises(ValueError) as exc: + with pytest.raises(ValueError, match='Rejected Request: langfuse_host is not allowed in request') as exc: is_request_body_safe( request_body={ "model": "gpt-4", @@ -2887,7 +2887,7 @@ def test_model_level_allow_does_not_skip_subsequent_banned_params(monkeypatch): lambda model, param, request_body_value, llm_router: param == "api_base", ) - with pytest.raises(ValueError) as exc: + with pytest.raises(ValueError, match='Rejected Request: langfuse_host is not allowed in request') as exc: is_request_body_safe( request_body={ "model": "gpt-4", @@ -2958,7 +2958,7 @@ class TestPricingInjectionBlocked: ], ) def test_pricing_field_rejected_by_default(self, field, value): - with pytest.raises(ValueError) as exc: + with pytest.raises(ValueError, match='Rejected Request') as exc: is_request_body_safe( request_body={"model": "gpt-4", field: value}, general_settings={}, diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index a9e12beb54b..99a0a4c0a8b 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -2589,7 +2589,7 @@ async def test_find_and_validate_raises_when_required_team_not_found(): # Token without team info jwt_token = {"sub": "user-1"} - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="No team found in token\\. Checked team_id field 'None' and") as exc_info: await JWTAuthManager.find_and_validate_specific_team_id( jwt_handler=jwt_handler, jwt_valid_token=jwt_token, @@ -2916,7 +2916,7 @@ async def test_find_and_validate_specific_team_id_hints_bracket_notation(): # token has roles as a list — dot-notation won't find anything token = {"roles": ["team1"]} - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="is not supported\\. Use 'roles' instead — LiteLLM") as exc_info: await JWTAuthManager.find_and_validate_specific_team_id( jwt_handler=handler, jwt_valid_token=token, @@ -2947,7 +2947,7 @@ async def test_find_and_validate_specific_team_id_hints_bracket_index_notation() handler = _make_jwt_handler("roles[0]") token = {"roles": ["team1"]} - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="is not supported in team_id_jwt_field\\. Use 'roles' instead") as exc_info: await JWTAuthManager.find_and_validate_specific_team_id( jwt_handler=handler, jwt_valid_token=token, @@ -2977,7 +2977,7 @@ async def test_find_and_validate_specific_team_id_no_hint_for_valid_field(): handler = _make_jwt_handler("appid") token = {} # no appid — triggers the "no team found" path - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="No team found in token\\. Checked team_id field 'appid' and") as exc_info: await JWTAuthManager.find_and_validate_specific_team_id( jwt_handler=handler, jwt_valid_token=token, @@ -4807,7 +4807,7 @@ async def test_multi_issuer_jwt_unknown_issuer_falls_back_to_global_jwks(monkeyp kid="issuer-key", ) - with pytest.raises(Exception) as exc: + with pytest.raises(Exception, match='Missing JWT Public Key URL from environment\\.') as exc: await jwt_handler.auth_jwt(token=token) assert "Missing JWT Public Key URL from environment." in str(exc.value) @@ -4838,7 +4838,7 @@ async def test_multi_issuer_jwt_rejects_wrong_audience(monkeypatch): kid="issuer-key", ) - with pytest.raises(Exception) as exc: + with pytest.raises(Exception, match="Validation fails: Audience doesn't match") as exc: await jwt_handler.auth_jwt(token=token) assert "Validation fails" in str(exc.value) @@ -4881,7 +4881,7 @@ async def test_multi_issuer_jwt_same_kid_does_not_cross_issuer_keys(monkeypatch) kid=shared_kid, ) - with pytest.raises(Exception) as exc: + with pytest.raises(Exception, match='Validation fails: Signature verification failed') as exc: await jwt_handler.auth_jwt(token=token) assert "Validation fails" in str(exc.value) @@ -4936,7 +4936,7 @@ def test_multi_issuer_jwt_requires_audience_unless_explicitly_disabled( issuer = "https://issuer.example.com" jwks_url = f"{issuer}/keys" - with pytest.raises(Exception) as exc: + with pytest.raises(Exception, match='must configure audience or set') as exc: LiteLLM_JWTAuth( issuers=[ { @@ -4953,7 +4953,7 @@ def test_multi_issuer_jwt_rejects_audience_with_disable_audience_validation(): issuer = "https://issuer.example.com" jwks_url = f"{issuer}/keys" - with pytest.raises(Exception) as exc: + with pytest.raises(Exception, match='cannot set audience and disable_audience_validation=True') as exc: LiteLLM_JWTAuth( issuers=[ { diff --git a/tests/test_litellm/proxy/auth/test_oauth2_proxy_hook.py b/tests/test_litellm/proxy/auth/test_oauth2_proxy_hook.py index dcbfd281e01..2d81d48de1e 100644 --- a/tests/test_litellm/proxy/auth/test_oauth2_proxy_hook.py +++ b/tests/test_litellm/proxy/auth/test_oauth2_proxy_hook.py @@ -141,7 +141,7 @@ async def test_refuses_to_map_non_identity_fields(configure_proxy, privileged_fi configure_proxy(mappings={privileged_field: f"x-{privileged_field}"}) request = _request_with_headers({f"x-{privileged_field}": "proxy_admin"}) - with pytest.raises(ValueError) as exc: + with pytest.raises(ValueError, match='proxy auth refuses to map non-identity UserAPIKeyAuth') as exc: await handle_oauth2_proxy_request(request) assert privileged_field in str(exc.value) diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 6d6e20e9c36..636c5480d67 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -42,7 +42,7 @@ def test_non_admin_config_update_route_rejected(): request.query_params = {} # Test that calling /config/update route raises HTTPException with 403 status - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info: RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER.value, @@ -134,7 +134,7 @@ def test_user_banner_update_rejected_for_non_admin(): request = MagicMock(spec=Request) request.query_params = {} - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info: RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER.value, @@ -1814,7 +1814,7 @@ def test_internal_user_blocked_from_global_spend_routes(route): request = MagicMock(spec=Request) request.query_params = {} - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info: RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER.value, @@ -1843,7 +1843,7 @@ def test_internal_user_view_only_blocked_from_global_spend_routes(route): request = MagicMock(spec=Request) request.query_params = {} - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info: RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, @@ -2046,7 +2046,7 @@ def test_internal_user_blocked_from_admin_viewer_logs_routes(route): if route not in INTERNAL_USER_BLOCKED_SUBSET: return - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info: RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER.value, @@ -2530,7 +2530,7 @@ def test_non_admin_non_team_admin_cannot_access_config_update_but_can_attempt_re ) # /config/update is still blocked - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info: RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER.value, @@ -3188,7 +3188,7 @@ def test_internal_user_blocked_from_search_tool_writes(route): request = MagicMock(spec=Request) request.query_params = {} - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info: RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER.value, diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index ab7e3d9701c..043bbb5b76a 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -5337,7 +5337,7 @@ async def test_random_non_sk_token_is_rejected(monkeypatch): patch("litellm.proxy.proxy_server.master_key", "sk-master"), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='LiteLLM Virtual Key expected\\.') as exc_info: await user_api_key_auth( request=mock_request, api_key="Bearer not-a-real-token", @@ -5539,7 +5539,7 @@ async def test_real_jwt_still_requires_license_when_jwt_auth_enabled(monkeypatch patch("litellm.proxy.proxy_server.master_key", "sk-master"), patch("litellm.proxy.proxy_server.prisma_client", None), ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='JWT Auth is an enterprise only feature\\. You must be a') as exc_info: await user_api_key_auth( request=mock_request, api_key=f"Bearer {jwt_token}", diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index eb40f54a1f3..be29269fe25 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -84,7 +84,7 @@ class TestPollingErrorSurfacing: } with patch("requests.get", return_value=mock_response) as mock_get, patch("time.sleep"): - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Your litellm CLI is out of date and uses a login flow') as exc_info: _poll_for_ready_data("http://test/sso/cli/poll/sk-legacy") assert mock_get.call_count == 1 @@ -151,7 +151,7 @@ class TestStartCliSsoFlowErrors: mock_response.status_code = 404 with patch("requests.post", return_value=mock_response): - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Either --base-url is wrong, or the proxy is older than') as exc_info: _start_cli_sso_flow("https://old-proxy.example.com") message = str(exc_info.value) @@ -167,7 +167,7 @@ class TestStartCliSsoFlowErrors: mock_response.json.return_value = {"detail": "Too many CLI login attempts. Try again later."} with patch("requests.post", return_value=mock_response): - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Too many CLI login attempts\\. Try again later\\.') as exc_info: _start_cli_sso_flow("https://test.example.com") assert "HTTP 429" in str(exc_info.value) @@ -183,7 +183,7 @@ class TestStartCliSsoFlowErrors: mock_response.text = "Sign in to corporate VPN" with patch("requests.post", return_value=mock_response): - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='A proxy, load balancer, or auth gateway in front of') as exc_info: _start_cli_sso_flow("https://test.example.com") message = str(exc_info.value) @@ -197,7 +197,7 @@ class TestStartCliSsoFlowErrors: from litellm.proxy.client.cli.commands.auth import _start_cli_sso_flow with patch("requests.post", side_effect=requests.ConnectionError("Connection refused")): - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Connection refused\\. Check that the proxy is running') as exc_info: _start_cli_sso_flow("https://unreachable.example.com") message = str(exc_info.value) diff --git a/tests/test_litellm/proxy/client/cli/test_pkce_login.py b/tests/test_litellm/proxy/client/cli/test_pkce_login.py index 70f481d5cfa..f58bd0ff412 100644 --- a/tests/test_litellm/proxy/client/cli/test_pkce_login.py +++ b/tests/test_litellm/proxy/client/cli/test_pkce_login.py @@ -622,7 +622,7 @@ def test_fresh_api_key_never_hands_out_a_rotated_key_it_could_not_save(): def save(_record): raise OSError("disk full") - with pytest.raises(OSError): + with pytest.raises(OSError, match="disk full"): _fresh(STORED, save, http, now=lambda: 999_950.0) diff --git a/tests/test_litellm/proxy/client/test_models.py b/tests/test_litellm/proxy/client/test_models.py index b2485032a37..33f963b74af 100644 --- a/tests/test_litellm/proxy/client/test_models.py +++ b/tests/test_litellm/proxy/client/test_models.py @@ -472,14 +472,14 @@ def test_get_invalid_params(): client = ModelsManagementClient(base_url="http://localhost:8000") # Test with no parameters - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Exactly one of model_id or model_name must be provided') as exc_info: client.get() assert "Exactly one of model_id or model_name must be provided" in str( exc_info.value ) # Test with both parameters - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Exactly one of model_id or model_name must be provided') as exc_info: client.get(model_id="123", model_name="gpt-4") assert "Exactly one of model_id or model_name must be provided" in str( exc_info.value diff --git a/tests/test_litellm/proxy/common_utils/test_callback_utils.py b/tests/test_litellm/proxy/common_utils/test_callback_utils.py index 515a7b27c7b..77ada4c11a9 100644 --- a/tests/test_litellm/proxy/common_utils/test_callback_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_callback_utils.py @@ -586,7 +586,7 @@ def test_initialize_callbacks_on_proxy_rejects_class_valued_entry(probe_config_p silently never run the hook. Config load must fail instead.""" entry = f"{_PROBE_MODULE_NAME}.FloorMaxTokens" - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='litellm_settings\\.callbacks entry') as exc_info: _load_callbacks([entry], probe_config_path) message = str(exc_info.value) @@ -609,7 +609,7 @@ def test_initialize_callbacks_on_proxy_rejects_non_dispatchable_values( ): entry = f"{_PROBE_MODULE_NAME}.{attribute}" - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='litellm_settings\\.callbacks entry') as exc_info: _load_callbacks([entry], probe_config_path) message = str(exc_info.value) @@ -621,7 +621,7 @@ def test_initialize_callbacks_on_proxy_rejects_non_dispatchable_values( def test_initialize_callbacks_on_proxy_rejects_class_valued_non_list_value(probe_config_path): entry = f"{_PROBE_MODULE_NAME}.FloorMaxTokens" - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='litellm_settings\\.callbacks entry') as exc_info: _load_callbacks(entry, probe_config_path) assert entry in str(exc_info.value) diff --git a/tests/test_litellm/proxy/common_utils/test_path_utils.py b/tests/test_litellm/proxy/common_utils/test_path_utils.py index c8d58fa8259..8936d910777 100644 --- a/tests/test_litellm/proxy/common_utils/test_path_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_path_utils.py @@ -42,5 +42,5 @@ class TestSafeFilename: safe_filename("..") def test_empty_rejected(self): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Empty or unsafe filename'): safe_filename("") diff --git a/tests/test_litellm/proxy/common_utils/test_timezone_utils.py b/tests/test_litellm/proxy/common_utils/test_timezone_utils.py index 7f686c53c95..dc3917cb48e 100644 --- a/tests/test_litellm/proxy/common_utils/test_timezone_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_timezone_utils.py @@ -130,16 +130,16 @@ def test_parse_budget_reset_time_unset_defaults_to_midnight(): def test_parse_budget_reset_time_invalid_string_raises(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="hour 'HH:MM' or 'HH:MM:SS' string, e\\.g\\."): parse_budget_reset_time("25:00") - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="Invalid budget_reset_time 'noon'; expected a"): parse_budget_reset_time("noon") def test_parse_budget_reset_time_non_string_raises(): # Unquoted "12:00" in YAML parses to the int 720; it must fail loudly, # not silently fall back to midnight. - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="hour 'HH:MM' string, e\\.g\\."): parse_budget_reset_time(720) diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py index e949afce57b..4ea655b8871 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py @@ -308,9 +308,9 @@ async def test_partition_maintenance_issues_nothing_when_the_budget_is_already_s def test_unsupported_interval_raises(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Unsupported partition interval: year'): period_start(date(2026, 6, 1), "year") - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Unsupported partition interval: year'): next_period_start(date(2026, 6, 1), "year") diff --git a/tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py b/tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py index c8e0338eeaa..95e794012ec 100644 --- a/tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py +++ b/tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py @@ -913,7 +913,7 @@ async def test_health_check_alerts_for_non_connection_errors_during_a_replacemen await _yield_to_loop() assert wrapper._reconnection_lock.locked() is True - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='malformed SELECT'): await client.health_check() gate.set() diff --git a/tests/test_litellm/proxy/db/test_spend_log_tool_index.py b/tests/test_litellm/proxy/db/test_spend_log_tool_index.py index 71073fd216e..9c4fbbf41aa 100644 --- a/tests/test_litellm/proxy/db/test_spend_log_tool_index.py +++ b/tests/test_litellm/proxy/db/test_spend_log_tool_index.py @@ -327,7 +327,7 @@ class TestFlushToolUsageTransactions: async def test_non_connection_errors_do_not_retry(self): prisma = MagicMock() prisma.db.batch_ = MagicMock(side_effect=ValueError("bad data")) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="bad data"): await flush_tool_usage_transactions(prisma_client=prisma, transactions=[_transaction("r1")]) prisma.db.batch_.assert_called_once() diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py index 3adf8b8407d..ceb59571389 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py @@ -56,7 +56,7 @@ def _patched(guardrail: BedrockGuardrail, http_response): def test_init_rejects_both_identifier_and_checks(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Bedrock guardrail accepts either'): BedrockGuardrail(guardrailIdentifier="gid", checks=CONTENT_FILTER_CHECKS) @@ -304,7 +304,7 @@ async def test_truncated_pii_ignored_when_pii_check_not_configured(): @pytest.mark.asyncio async def test_checks_with_guardrail_version_rejected(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Bedrock guardrail accepts either'): BedrockGuardrail(checks=CONTENT_FILTER_CHECKS, guardrailVersion="DRAFT") diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_enkryptai.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_enkryptai.py index e6c94a4c3cd..d31f462a185 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_enkryptai.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_enkryptai.py @@ -178,7 +178,7 @@ class TestEnkryptAIGuardrailHooks: with patch.object( enkryptai_guardrail.async_handler, "post", return_value=mock_response ): - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='violation\\(s\\) detected') as exc_info: await enkryptai_guardrail.async_pre_call_hook( user_api_key_dict=mock_user_api_key_dict, cache=MagicMock(), diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index 5be0d43c250..523ec1a37b4 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -767,7 +767,7 @@ class TestErrorHandling: "API Error", request=MagicMock(), response=MagicMock(status_code=500) ), ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Generic Guardrail API failed: API Error') as exc_info: await generic_guardrail.apply_guardrail( inputs={"texts": ["test"]}, request_data=mock_request_data_input, @@ -786,7 +786,7 @@ class TestErrorHandling: "post", side_effect=httpx.RequestError("Connection failed", request=MagicMock()), ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Generic Guardrail API failed: Connection failed') as exc_info: await generic_guardrail.apply_guardrail( inputs={"texts": ["test"]}, request_data=mock_request_data_input, @@ -810,7 +810,7 @@ class TestErrorHandling: "post", side_effect=httpx.RequestError("Connection failed", request=MagicMock()), ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Generic Guardrail API failed: Connection failed') as exc_info: await guardrail.apply_guardrail( inputs={"texts": ["test"]}, request_data=mock_request_data_input, diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py index 89b6af27719..14c0d2f9435 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -2334,7 +2334,7 @@ async def test_async_moderation_hook_api_error_fail_on_error_true(): } # Should raise the exception since fail_on_error is True - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="API Error") as exc_info: await guardrail.async_moderation_hook( data=request_data, user_api_key_dict=mock_user_api_key_dict, @@ -2374,7 +2374,7 @@ async def test_async_moderation_hook_api_error_fail_on_error_false(): # Even with fail_on_error=False, the decorator may still raise the exception # This test verifies that the exception is properly logged and handled - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="API Error") as exc_info: await guardrail.async_moderation_hook( data=request_data, user_api_key_dict=mock_user_api_key_dict, @@ -2865,7 +2865,7 @@ async def test_skip_unscannable_still_fails_closed_on_api_error(): "post", AsyncMock(side_effect=Exception("model armor upstream 500")), ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='model armor upstream') as exc_info: await guardrail.async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth(), cache=MagicMock(spec=DualCache), diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py index 2284f2b678a..8f29ba66814 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py @@ -5652,7 +5652,7 @@ class TestPanwAirsTimeoutCoercion: assert isinstance(params.timeout, float) def test_litellm_params_rejects_garbage_timeout(self): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='validation error for LitellmParams'): LitellmParams( guardrail="panw_prisma_airs", mode="pre_call", diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py index a3d86034f70..81604e22c87 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py @@ -90,12 +90,12 @@ def test_config_model_wiring(): def test_init_rejects_empty_api_key(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='api_key must be non-empty'): StraikerGuardrail(api_key="") def test_init_rejects_invalid_fallback(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="unreachable_fallback must be 'fail_open' or 'fail_closed';"): StraikerGuardrail(api_key="k", unreachable_fallback="nope") @@ -109,7 +109,7 @@ def test_supported_hooks_limited_to_pre_and_post(): def test_during_call_mode_rejected_at_init(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Event hook GuardrailEventHooks\\.during_call is not in the'): StraikerGuardrail(api_key="k", event_hook="during_call") diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py index 4b381b67f0e..0c5addbc143 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -124,7 +124,7 @@ class TestToolPermissionGuardrail: assert rule_id is None def test_rule_requires_name_or_type(self): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='validation error for ToolPermissionRule'): ToolPermissionGuardrail( guardrail_name="invalid-rule", rules=[{"id": "no_target", "decision": "allow"}], @@ -1042,7 +1042,7 @@ class TestToolPermissionGuardrailInMemoryUpdate: assert guardrail._check_tool_permission("Secret")[0] is False assert guardrail._check_tool_permission("Other")[0] is True - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="Invalid regex for tool_name in rule 'bad': unterminated"): guardrail.update_in_memory_litellm_params( LitellmParams( guardrail="tool_permission", diff --git a/tests/test_litellm/proxy/guardrails/test_llm_as_a_judge.py b/tests/test_litellm/proxy/guardrails/test_llm_as_a_judge.py index 45dec4ddb2d..bd2553b3280 100644 --- a/tests/test_litellm/proxy/guardrails/test_llm_as_a_judge.py +++ b/tests/test_litellm/proxy/guardrails/test_llm_as_a_judge.py @@ -230,7 +230,7 @@ def test_parse_judge_verdict_reraises_when_no_json(): def test_parse_judge_verdict_rejects_json_non_object(): """Valid JSON that is not an object (e.g. a bare list) raises ValueError.""" - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='judge response is not a JSON object'): _parse_judge_verdict("[1, 2, 3]") diff --git a/tests/test_litellm/proxy/hooks/test_sensitive_data_routing.py b/tests/test_litellm/proxy/hooks/test_sensitive_data_routing.py index 78d2c3af0f3..35c0f8deaf1 100644 --- a/tests/test_litellm/proxy/hooks/test_sensitive_data_routing.py +++ b/tests/test_litellm/proxy/hooks/test_sensitive_data_routing.py @@ -227,7 +227,7 @@ class TestCustomGuardrailSensitiveDataRouting: request_data = {"model": "gpt-4"} - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Cannot route sensitive data without a session_id\\. Ensure') as exc_info: guardrail.raise_sensitive_data_route_exception( route_to_model="on-premise-model", request_data=request_data, diff --git a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py index 8d03857c917..2839acab6b0 100644 --- a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py +++ b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py @@ -150,7 +150,7 @@ async def test_no_leak_on_over_limit_rejection(rate_limiter): f"estimated={estimated}, limit={user_api_key_dict.tpm_limit}" ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Limit type: tokens\\. Current limit') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, @@ -685,7 +685,7 @@ async def test_contentless_request_reserves_minimum(rate_limiter): f"counter should be 2, got {counter_after_two}" ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Limit type: tokens\\. Current limit') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, @@ -1319,7 +1319,7 @@ async def test_project_otpm_rejects_multiple_completion_candidates(rate_limiter) "n": 10, } - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_otpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, @@ -1347,7 +1347,7 @@ async def test_project_otpm_reserves_largest_conflicting_output_cap(rate_limiter "max_completion_tokens": 100, } - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_otpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, @@ -1377,7 +1377,7 @@ async def test_project_otpm_rejects_google_genai_native_output_cap( project_metadata={"model_otpm_limit": {model: 50}}, ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_otpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, @@ -1411,7 +1411,7 @@ async def test_project_otpm_rejects_google_genai_native_candidate_count( project_metadata={"model_otpm_limit": {model: 150}}, ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_otpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, @@ -1500,7 +1500,7 @@ async def test_project_otpm_over_limit_rolls_back_itpm_reservation(rate_limiter) "max_tokens": 500, # blows past the 10-token OTPM limit } - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_otpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, @@ -2003,7 +2003,7 @@ async def test_otpm_rejection_does_not_double_refund_combined_tpm(rate_limiter): rate_limit_type="tokens", ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_otpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, @@ -2060,7 +2060,7 @@ async def test_project_itpm_rejects_pretokenized_embedding_input( "input": embedding_input, } - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_itpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, @@ -2250,7 +2250,7 @@ async def test_itpm_reservation_accounts_for_audio_content_not_just_text(rate_li ], } - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_itpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, @@ -2361,7 +2361,7 @@ async def test_itpm_rejects_large_audio_payload_that_would_pass_flat_estimate( ], } - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_itpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, @@ -2622,7 +2622,7 @@ async def test_explicit_zero_output_responses_call_reserves_effective_provider_m }, ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_otpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, @@ -2850,7 +2850,7 @@ async def test_otpm_rejection_releases_stashed_parallel_slot(rate_limiter): "rate_limit": {"tokens_per_unit": 5, "window_size": 60}, } - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_otpm') as exc_info: await handler._reserve_project_io_tokens_or_raise( descriptors=[otpm_descriptor], data=data, @@ -3296,7 +3296,7 @@ async def test_rerank_query_and_documents_enforce_project_itpm( project_metadata={"model_itpm_limit": {"rerank-model": 100}}, ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_itpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index d4e9ccdca5e..069cfa01178 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -10179,7 +10179,7 @@ async def test_update_key_creator_reassigned_key_blocked(monkeypatch): mock_request = MagicMock() mock_request.query_params = {} - with pytest.raises(Exception) as exc: + with pytest.raises(Exception, match='User can only create keys for themselves\\. Got') as exc: await update_key_fn( request=mock_request, data=UpdateKeyRequest(key=test_hashed_token, key_alias="hijacked"), diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index 84dee5b05c5..01c0760bd27 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -2325,7 +2325,7 @@ class TestTemporaryMCPSessionEndpoints: "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload", MagicMock(), ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='User does not have permission to create temporary mcp') as exc_info: await add_session_mcp_server( payload=payload, user_api_key_dict=non_admin, diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 7e4596d154b..42e96ad8659 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -140,7 +140,7 @@ class TestModelManagementAuthChecks: @pytest.mark.asyncio async def test_can_user_make_team_model_call_non_premium_fails(self): """Test that non-premium users cannot make team model calls""" - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='You must be a LiteLLM Enterprise user to use this feature\\.') as exc_info: ModelManagementAuthChecks.can_user_make_team_model_call( team_id="test_team", user_api_key_dict=self.admin_user, @@ -195,7 +195,7 @@ class TestModelManagementAuthChecks: ) prisma_client = MockPrismaClient(team_exists=True) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='You must be a LiteLLM Enterprise user to use this feature\\.') as exc_info: await ModelManagementAuthChecks.allow_team_model_action( model_params=model_params, user_api_key_dict=self.admin_user, @@ -216,7 +216,7 @@ class TestModelManagementAuthChecks: ) prisma_client = MockPrismaClient(team_exists=False) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="Team id=nonexistent_team does not exist in db'\\}") as exc_info: await ModelManagementAuthChecks.allow_team_model_action( model_params=model_params, user_api_key_dict=self.admin_user, @@ -257,7 +257,7 @@ class TestModelManagementAuthChecks: ) prisma_client = MockPrismaClient(team_exists=True, user_admin=False) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="Team ID=test_team does not match the API key's team") as exc_info: await ModelManagementAuthChecks.can_user_make_model_call( model_params=model_params, user_api_key_dict=self.normal_user, @@ -1483,7 +1483,7 @@ class TestTeamModelUpdate: "litellm.proxy.proxy_server.premium_user", True, ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="does not match the API key's team ID=None, OR you are") as exc_info: await _update_team_model_in_db( db_model=db_model, patch_data=patch_data, @@ -3256,7 +3256,7 @@ class TestPatchModelBlockedAuthGate: new=AsyncMock(return_value=None), ), ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="Only proxy admins can change a model's blocked flag\\.") as exc_info: await patch_model( model_id="m1", patch_data=updateDeployment(blocked=True), diff --git a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py index 92a34b5ee7c..a1c38d26b9d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py @@ -58,7 +58,7 @@ def test_model_info_accepts_valid_ptu_fields(): def test_model_info_rejects_non_positive_count(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='value_error, input_value'): ModelInfo( id="x", team_id="t", @@ -69,7 +69,7 @@ def test_model_info_rejects_non_positive_count(): def test_model_info_rejects_negative_rate(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='value_error, input_value'): ModelInfo( id="x", team_id="t", @@ -82,7 +82,7 @@ def test_model_info_rejects_negative_rate(): def test_model_info_rejects_a_count_beyond_the_cap(): """flat cost multiplies the count by a float, and an unbounded int overflows that conversion, which aborted the rollup for every team rather than skipping one model.""" - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='validation error for ModelInfo'): ModelInfo(id="x", team_id="t", ptu_count=10**400, cost_per_ptu_per_hour=2.0) @@ -95,12 +95,12 @@ def test_model_info_accepts_a_count_at_the_cap(): def test_model_info_rejects_a_non_finite_rate(rate): """NaN compares False against every bound, so a bare `< 0` check let it through and the deployment then accrued a flat cost of nan.""" - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='value_error, input_value'): ModelInfo(id="x", team_id="t", ptu_count=5, cost_per_ptu_per_hour=rate) def test_model_info_rejects_a_rate_beyond_the_cap(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='validation error for ModelInfo'): ModelInfo(id="x", team_id="t", ptu_count=5, cost_per_ptu_per_hour=ModelInfo.MAX_COST_PER_PTU_PER_HOUR * 2) @@ -148,7 +148,7 @@ def test_validate_helper_passes_full_config(): def test_model_info_rejects_effective_to_before_from(): import datetime - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='validation error for ModelInfo'): ModelInfo( id="x", team_id="t", @@ -186,7 +186,7 @@ def test_model_info_compares_mixed_naive_and_aware_timestamps(): ) assert info.ptu_effective_to is not None - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='validation error for ModelInfo'): ModelInfo( id="x", team_id="t", @@ -698,7 +698,7 @@ class TestAddNewModelPtuGate: with ExitStack() as stack: for active_patch in patches: stack.enter_context(active_patch) - with pytest.raises(Exception) as exc: + with pytest.raises(Exception, match='PTU cost attribution is disabled, so ptu_count') as exc: await add_new_model(model_params=self._ptu_deployment("ptu-gate-model"), user_api_key_dict=admin) assert PTU_COST_ATTRIBUTION_ENV_VAR in str(exc.value) @@ -1273,7 +1273,7 @@ class TestPtuDeploymentsAreNotBilledPerToken: with ExitStack() as stack: for active_patch in patches: stack.enter_context(active_patch) - with pytest.raises(Exception) as exc: + with pytest.raises(Exception, match='A PTU deployment bills by reserved capacity, so') as exc: await add_new_model(model_params=deployment, user_api_key_dict=admin) assert "input_cost_per_token" in str(exc.value) diff --git a/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py b/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py index e8a74e41dae..3a32b3cc128 100644 --- a/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py +++ b/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py @@ -458,7 +458,7 @@ class TestUsageAiChatServiceAccountGuard: _resolve_fetch_kwargs, ) - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Non-admin caller has user_id=None; refusing to issue an') as exc_info: _resolve_fetch_kwargs( fn_name="get_usage_data", fn_args={"start_date": "2025-01-01", "end_date": "2025-01-31"}, diff --git a/tests/test_litellm/proxy/management_helpers/test_team_metadata_validation.py b/tests/test_litellm/proxy/management_helpers/test_team_metadata_validation.py index e4b031ade57..1acb8e7e016 100644 --- a/tests/test_litellm/proxy/management_helpers/test_team_metadata_validation.py +++ b/tests/test_litellm/proxy/management_helpers/test_team_metadata_validation.py @@ -671,7 +671,7 @@ async def test_non_callable_validator_is_rejected_with_clean_500(): def test_parse_schema_duplicate_error_lists_offending_keys(): - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='team_metadata_schema contains duplicate keys: app_name') as exc_info: parse_team_metadata_schema( [{"key": "cost_center"}, {"key": "app_name"}, {"key": "cost_center"}, {"key": "app_name"}] ) diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py b/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py index a05b8ae530c..6ce7af1e2ee 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py @@ -781,7 +781,7 @@ async def test_the_rewrite_closes_its_own_output_when_it_cannot_finish(): original_read = bg._read_spooled bg._read_spooled = _boom try: - with pytest.raises(OSError): + with pytest.raises(OSError, match='no space left on device'): rewrite_batch_input_file(source, result) finally: bg.tempfile.SpooledTemporaryFile = real diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index f994fba371b..d3237f5f49d 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -2714,7 +2714,7 @@ class TestMilvusProxyRoute: None ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Vector store not found for missing-store') as exc_info: await milvus_proxy_route( endpoint="vectors/search", request=mock_request, @@ -2779,7 +2779,7 @@ class TestMilvusProxyRoute: mock_vector_store ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='api_base not found in vector store configuration for') as exc_info: await milvus_proxy_route( endpoint="vectors/search", request=mock_request, @@ -2988,7 +2988,7 @@ class TestOpenAIPassthroughRoute: "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials", return_value=None, ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="Required 'OPENAI_API_KEY' in environment to make") as exc_info: await openai_proxy_route( endpoint="v1/chat/completions", request=mock_request, @@ -3177,7 +3177,7 @@ class TestCursorProxyRoute: [], ), ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Cursor API key not found\\. Add Cursor credentials via') as exc_info: await cursor_proxy_route( endpoint="v0/agents", request=mock_request, diff --git a/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py b/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py index ebebfde5cd3..b6633779326 100644 --- a/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py +++ b/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py @@ -157,7 +157,7 @@ class TestUpdatePolicyDraftOnly: prod_row = _make_row(policy_id="pid-1", version_status="production") prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=prod_row) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Error updating policy in DB: Only draft versions can be') as exc_info: await registry.update_policy_in_db( policy_id="pid-1", policy_request=PolicyUpdateRequest(description="new"), @@ -341,7 +341,7 @@ class TestUpdateVersionStatus: draft = _make_row(policy_id="d-1", version_status="draft") prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=draft) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Error updating version status: Cannot promote draft') as exc_info: await registry.update_version_status( policy_id="d-1", new_status="production", diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 54ae279d005..47f01fe096d 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -1504,7 +1504,7 @@ async def test_ProxyConfig__init_non_llm_configs_premium_invalid_worker_registry async def test_ProxyConfig__init_non_llm_configs_worker_registry_requires_premium(monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False) pc = ProxyConfig() - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Trying to use `worker_registry`You must be a LiteLLM') as exc_info: await pc._init_non_llm_configs( config={ "worker_registry": [ @@ -1769,7 +1769,7 @@ def test_ProxyConfig_initialize_secret_manager_none_noop(): def test_ProxyConfig_initialize_secret_manager_invalid_kms_raises(): pc = ProxyConfig() - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Invalid Key Management System selected'): pc.initialize_secret_manager(key_management_system="not-a-real-kms") diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 15a3e6609f0..b2ec500d045 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -3263,7 +3263,7 @@ async def test_provider_budget_over(disable_budget_sync): model_list=MODEL_LIST, ) - with pytest.raises(Exception) as e: + with pytest.raises(Exception, match='No deployments available - crossed budget: Exceeded budget') as e: await router.acompletion( model="azure-gpt-4o", messages=[{"role": "user", "content": "Hello, world!"}], @@ -5096,7 +5096,7 @@ def test_resolve_spend_report_scope_missing_caller_value_400(): @pytest.mark.parametrize("bad_column", ["metadata", "end_user", "evil; DROP TABLE", ""]) def test_scoped_spend_report_sql_rejects_unknown_column(bad_column): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Unsupported spend report scope column'): spend_management_endpoints._scoped_spend_report_sql(scope_column=bad_column) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index c693f5ab2cb..716fba370df 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -435,7 +435,7 @@ class TestProxyBaseLLMRequestProcessing: # Test with invalid header value (should raise ValueError when converting to float) headers_with_invalid = {"x-litellm-stream-timeout": "invalid"} - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="could not convert string to float: 'invalid"): LiteLLMProxyRequestSetup._get_stream_timeout_from_request(headers_with_invalid) @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/test_enforce_user_param.py b/tests/test_litellm/proxy/test_enforce_user_param.py index 6891123e70e..1001372aeb5 100644 --- a/tests/test_litellm/proxy/test_enforce_user_param.py +++ b/tests/test_litellm/proxy/test_enforce_user_param.py @@ -56,7 +56,7 @@ class TestEnforceUserParamPostGetFiltering: new_callable=AsyncMock, return_value=True, ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="user' param not passed in\\. 'enforce_user_param'=True") as exc_info: await common_checks( request_body=request_body, team_object=None, @@ -175,7 +175,7 @@ class TestEnforceUserParamPostGetFiltering: new_callable=AsyncMock, return_value=True, ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="user' param not passed in\\. 'enforce_user_param'=True") as exc_info: await common_checks( request_body=request_body, team_object=None, @@ -405,7 +405,7 @@ class TestEnforceUserParamEdgeCases: new_callable=AsyncMock, return_value=True, ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="user' param not passed in\\. 'enforce_user_param'=True") as exc_info: await common_checks( request_body=request_body, team_object=None, diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index b1071150f3b..636974d5deb 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -2279,12 +2279,12 @@ def test_get_num_retries_from_request(): # Test case 7: Header present with invalid value (should raise ValueError when int() is called) headers_with_invalid = {"x-litellm-num-retries": "invalid"} - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='invalid literal for int\\(\\) with base'): LiteLLMProxyRequestSetup._get_num_retries_from_request(headers_with_invalid) # Test case 8: Header present with float string (should raise ValueError when int() is called) headers_with_float = {"x-litellm-num-retries": "3.5"} - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='invalid literal for int\\(\\) with base'): LiteLLMProxyRequestSetup._get_num_retries_from_request(headers_with_float) # Test case 9: Header present with negative number @@ -2324,7 +2324,7 @@ def test_get_keepalive_seconds_from_request(): # Header present with invalid value raises ValueError, matching the other # x-litellm-* numeric header helpers (_get_timeout_from_request, etc.) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="could not convert string to float: 'not-a-number"): LiteLLMProxyRequestSetup._get_keepalive_seconds_from_request( {"x-litellm-keepalive-seconds": "not-a-number"} ) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 75aa716bb85..83e9095c8ec 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1507,7 +1507,7 @@ def test_team_info_masking(): "langfuse_public_key": "public-test-key", } - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="secr\\*\\*\\*\\*\\*\\*\\*-key', 'langfuse_public_key':") as exc_info: proxy_config._get_team_config( team_id="test_dev", all_teams_config=[team1_info], diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index ce0b6b755cc..bf1538183ab 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -82,10 +82,10 @@ def test_spend_log_cleanup_cron_scheduling(): assert trigger_weekly is not None # Invalid cron expression should raise ValueError - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Wrong number of fields; got'): CronTrigger.from_crontab("invalid cron") - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='is higher than the maximum value'): CronTrigger.from_crontab("60 25 * * *") # Invalid minute and hour diff --git a/tests/test_litellm/proxy/test_team_org_move.py b/tests/test_litellm/proxy/test_team_org_move.py index 2dc961bec85..064e9de550e 100644 --- a/tests/test_litellm/proxy/test_team_org_move.py +++ b/tests/test_litellm/proxy/test_team_org_move.py @@ -97,7 +97,7 @@ class TestValidateTeamOrgChange: team = _make_team(member_ids=["sso-user-001"]) org = _make_org(members=[]) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Cannot move team to organization\\. Team has user_id') as exc_info: validate_team_org_change( team=team, organization=org, llm_router=router, is_proxy_admin=False ) diff --git a/tests/test_litellm/proxy/utils/helpers/test_team_configs.py b/tests/test_litellm/proxy/utils/helpers/test_team_configs.py index 0e0906892b0..185d4d26ff4 100644 --- a/tests/test_litellm/proxy/utils/helpers/test_team_configs.py +++ b/tests/test_litellm/proxy/utils/helpers/test_team_configs.py @@ -66,7 +66,7 @@ def test_is_valid_team_configs_short_circuits_when_team_id_none(): def test_is_valid_team_configs_raises_on_model_not_in_team_models(): team_config = {"models": ["gpt-4o"]} request_data = {"model": "claude-haiku"} - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='claude-haiku\\. Valid models for team are') as exc_info: _is_valid_team_configs( team_id="team-1", team_config=team_config, diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py index 7057a112c83..93c99c7fd04 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py @@ -561,7 +561,7 @@ async def test_update_spend_logs_does_not_requeue_non_transport_failures( proxy_logging.failure_handler = AsyncMock() mock_prisma_client.spend_log_transactions = [] - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="bad payload"): await ProxyUpdateSpend.update_spend_logs( n_retry_times=1, prisma_client=mock_prisma_client, diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_callback_capabilities_class.py b/tests/test_litellm/proxy/utils/proxy_logging/test_callback_capabilities_class.py index 9452e8042bd..75a91177f00 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_callback_capabilities_class.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_callback_capabilities_class.py @@ -181,7 +181,7 @@ def test_has_streaming_callbacks_error_when_resolution_fails(monkeypatch): "get_custom_logger_compatible_class", lambda *a, **kw: (_ for _ in ()).throw(ValueError("nope")), ) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="nope"): ProxyLogging.has_streaming_callbacks() diff --git a/tests/test_litellm/responses/test_responses_websocket_all_providers.py b/tests/test_litellm/responses/test_responses_websocket_all_providers.py index 2d523bfdeb3..e8333214ea8 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -84,7 +84,7 @@ class TestResponsesAPIWebSocketSupport: def test_azure_websocket_url_requires_api_base(self): config = AzureOpenAIResponsesAPIConfig() - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='api_base is required for Azure WebSocket'): config.get_websocket_url(api_base=None, litellm_params={}) def test_azure_model_not_in_websocket_url(self): diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_bandit.py b/tests/test_litellm/router_strategy/adaptive_router/test_bandit.py index ab322f0fb37..78390cc1193 100644 --- a/tests/test_litellm/router_strategy/adaptive_router/test_bandit.py +++ b/tests/test_litellm/router_strategy/adaptive_router/test_bandit.py @@ -100,7 +100,7 @@ def test_score_combines_quality_and_cost(): def test_pick_best_empty_dict_raises(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='pick_best called with no models'): pick_best({}, {}) diff --git a/tests/test_litellm/router_strategy/test_router_tag_routing.py b/tests/test_litellm/router_strategy/test_router_tag_routing.py index 73491490b14..60b1166de73 100644 --- a/tests/test_litellm/router_strategy/test_router_tag_routing.py +++ b/tests/test_litellm/router_strategy/test_router_tag_routing.py @@ -656,7 +656,7 @@ async def test_negation_all_excluded_raises(): enable_tag_filtering=True, ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: await router.acompletion( model="gpt-4", messages=[{"role": "user", "content": "hi"}], @@ -699,7 +699,7 @@ async def test_negation_ban_only_cannot_escape_default_pool(): # Sending only "!default" must NOT route to the paid deployment. # The base pool for ban-only is the default pool; banning the only # default deployment should raise rather than falling through to paid. - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: await router.acompletion( model="gpt-4", messages=[{"role": "user", "content": "hi"}], @@ -969,7 +969,7 @@ async def test_negation_exhausts_entire_fallback_chain(): enable_tag_filtering=True, ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: await router.acompletion( model="primary", messages=[{"role": "user", "content": "hi"}], @@ -1719,7 +1719,7 @@ async def test_required_and_unmatched_raises_by_default(): enable_tag_filtering=True, ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: await router.acompletion( model="gpt-4", messages=[{"role": "user", "content": "hi"}], @@ -1751,7 +1751,7 @@ async def test_required_and_combined_with_positive_unmatched_raises_by_default() enable_tag_filtering=True, ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: await router.acompletion( model="gpt-4", messages=[{"role": "user", "content": "hi"}], @@ -1973,7 +1973,7 @@ async def test_negation_combined_with_positive_unmatched_raises_by_default(): enable_tag_filtering=True, ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: await router.acompletion( model="gpt-4", messages=[{"role": "user", "content": "hi"}], @@ -2131,7 +2131,7 @@ async def test_mixed_constraint_survivor_unmatched_by_positive_tag_raises_by_def enable_tag_filtering=True, ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: await router.acompletion( model="gpt-4", messages=[{"role": "user", "content": "hi"}], @@ -2224,7 +2224,7 @@ async def test_allow_fail_open_denied_when_request_includes_unknown_tag(): enable_tag_filtering=True, ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: await router.acompletion( model="gpt-4", messages=[{"role": "user", "content": "hi"}], @@ -2538,7 +2538,7 @@ async def test_plain_tag_exhaustion_with_universal_default_tag_raises_by_default "litellm.router._async_get_cooldown_deployments", new=AsyncMock(return_value=["quality-high-1", "quality-high-2"]), ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: await router.acompletion( model="gpt-4", messages=[{"role": "user", "content": "hi"}], @@ -2767,7 +2767,7 @@ async def test_allow_fail_open_raises_when_inherited_constraint_alone_is_unsatis # allow_fail_open unset. router = _eu_region_router() - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: await router.acompletion( model="chat", messages=[{"role": "user", "content": "hi"}], @@ -2941,7 +2941,7 @@ async def test_tagged_request_direct_to_plain_group_still_rejected(): # tag filtering must reject exactly as before. router = _tagged_marker_router() - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: await router.acompletion( model="gemini-flash", messages=[{"role": "user", "content": "hi"}], @@ -2962,7 +2962,7 @@ async def test_caller_forged_consumption_stamp_is_neutralized_by_the_hook(): # tag filtering runs. router = _tagged_marker_router() - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: await router.acompletion( model="gemini-flash", messages=[{"role": "user", "content": "hi"}], @@ -2984,7 +2984,7 @@ async def test_inherited_constraint_still_applies_to_the_routed_tier(): # ®ion:eu comes from key/team policy (present in inherited_tags): # consuming the router-selecting "route" tag must not also discard the # inherited requirement, so a tier without the tag still raises... - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: await _tagged_marker_router().acompletion( model="gpt4o", messages=[{"role": "user", "content": "hi"}], diff --git a/tests/test_litellm/sandbox/test_e2b_sandbox.py b/tests/test_litellm/sandbox/test_e2b_sandbox.py index cc5b12156a1..e01b9120416 100644 --- a/tests/test_litellm/sandbox/test_e2b_sandbox.py +++ b/tests/test_litellm/sandbox/test_e2b_sandbox.py @@ -293,7 +293,7 @@ async def test_public_lifecycle_create_run_delete(): @pytest.mark.asyncio async def test_unsupported_provider_raises(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="not-a-provider' is not a valid SandboxProviders"): await litellm.acreate_sandbox(provider="not-a-provider") diff --git a/tests/test_litellm/secret_managers/test_base_secret_manager.py b/tests/test_litellm/secret_managers/test_base_secret_manager.py index cba6a99ab7f..e1ccb91c381 100644 --- a/tests/test_litellm/secret_managers/test_base_secret_manager.py +++ b/tests/test_litellm/secret_managers/test_base_secret_manager.py @@ -32,7 +32,7 @@ from litellm.secret_managers.base_secret_manager import raise_if_unsafe_secret_n ], ) def test_raise_if_unsafe_secret_name_rejects_traversal_and_line_breaks(secret_name): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Invalid secret_name'): raise_if_unsafe_secret_name(secret_name) diff --git a/tests/test_litellm/test_github_close_low_quality_prs.py b/tests/test_litellm/test_github_close_low_quality_prs.py index e3b653dde64..2a891ca72f5 100644 --- a/tests/test_litellm/test_github_close_low_quality_prs.py +++ b/tests/test_litellm/test_github_close_low_quality_prs.py @@ -697,7 +697,7 @@ class TestListOpenItemsNoCap: def test_list_open_items_rejects_unknown_kind(self, closer_module): shared = self._shared(closer_module) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="kind must be 'pr' or 'issue', got 'both"): shared.list_open_items("both", repo="o/r", fields="number") def test_fetch_open_prs_delegates_with_no_cap(self, closer_module, monkeypatch): diff --git a/tests/test_litellm/test_github_triage_with_llm.py b/tests/test_litellm/test_github_triage_with_llm.py index 96b77e80457..ddffb978b48 100644 --- a/tests/test_litellm/test_github_triage_with_llm.py +++ b/tests/test_litellm/test_github_triage_with_llm.py @@ -665,11 +665,11 @@ class TestParseVerdict: assert triage_module.parse_verdict(raw)["verdict"] == "pass" def test_should_raise_for_unparseable_text(self, triage_module): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='could not extract JSON from LLM response: not even close to'): triage_module.parse_verdict("not even close to json") def test_should_raise_for_empty(self, triage_module): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='empty LLM response'): triage_module.parse_verdict("") diff --git a/tests/test_litellm/test_redact_string_in_error_paths.py b/tests/test_litellm/test_redact_string_in_error_paths.py index 8eddc2b1a5a..1c4d91397d1 100644 --- a/tests/test_litellm/test_redact_string_in_error_paths.py +++ b/tests/test_litellm/test_redact_string_in_error_paths.py @@ -234,7 +234,7 @@ class TestRouterFallbackFailureTracebackRedaction: raise ValueError(f"primary deployment failed api_key={secret}") except ValueError as original_exception: with caplog.at_level(logging.DEBUG, logger="LiteLLM Router"): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='primary deployment failed api_key=sk-testsecretvalu'): await router.async_function_with_fallbacks_common_utils( e=original_exception, disable_fallbacks=False, diff --git a/tests/test_litellm/test_redis.py b/tests/test_litellm/test_redis.py index 3aa4bc58f13..c645a67ef84 100644 --- a/tests/test_litellm/test_redis.py +++ b/tests/test_litellm/test_redis.py @@ -131,7 +131,7 @@ def test_get_redis_url_from_environment_missing_host_port(monkeypatch): monkeypatch.delenv("REDIS_PORT", raising=False) # Call the function and expect a ValueError - with pytest.raises(ValueError) as excinfo: + with pytest.raises(ValueError, match="Either 'REDIS_URL' or both 'REDIS_HOST' and 'REDIS_PORT") as excinfo: get_redis_url_from_environment() # Check the error message @@ -149,7 +149,7 @@ def test_get_redis_url_from_environment_missing_port(monkeypatch): monkeypatch.setenv("REDIS_HOST", "redis-server") # Call the function and expect a ValueError - with pytest.raises(ValueError) as excinfo: + with pytest.raises(ValueError, match="Either 'REDIS_URL' or both 'REDIS_HOST' and 'REDIS_PORT") as excinfo: get_redis_url_from_environment() # Check the error message diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index b50dc92c220..a47525749f6 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -953,7 +953,7 @@ async def test_arouter_filter_team_based_models(): assert result is not None # FAILS - with pytest.raises(Exception) as e: + with pytest.raises(Exception, match='No deployments available for selected model, Try again in') as e: result = await router.acompletion( model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hello, world!"}], @@ -1225,7 +1225,7 @@ def test_add_invalid_provider_to_router(): ], ) - with pytest.raises(Exception) as e: + with pytest.raises(Exception, match='Unsupported provider - vertex_ai_eu') as e: router.add_deployment( Deployment( model_name="vertex_ai/*", @@ -1320,7 +1320,7 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): with patch.object(router, "async_get_available_deployment") as mock_get_deployment: mock_get_deployment.side_effect = Exception("No deployment available") - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='No deployment available') as exc_info: await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_generic_function, @@ -1394,7 +1394,7 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): with patch.object( router, "async_routing_strategy_pre_call_checks" ) as mock_pre_call_checks: - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Mock failure') as exc_info: await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_failing_function, @@ -3737,7 +3737,7 @@ def test_count_pre_call_check_tokens_across_api_surfaces(): assert string_input_tokens > 0 assert list_input_tokens > 0 - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Either messages or input must be provided to count tokens'): router._count_pre_call_check_tokens(messages=None, input=None) diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index 5bb854c12e0..dc210f900bf 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -1706,7 +1706,7 @@ def test_an_incomplete_reservation_is_refused_rather_than_served(dropped): state the operator was trying to leave.""" incomplete = {k: v for k, v in _PTU_MODEL_INFO.items() if k != dropped} - with pytest.raises(ValueError) as raised: + with pytest.raises(ValueError, match="PTU configuration on model 'gpt") as raised: _ptu_router(model_info=incomplete, litellm_params={"input_cost_per_token": 5e-06}) assert "gpt-4o-ptu" in str(raised.value) @@ -1726,7 +1726,7 @@ def test_the_refusal_reason_is_the_one_the_model_endpoint_answers_with(dropped, incomplete = {k: v for k, v in _PTU_MODEL_INFO.items() if k != dropped} assert ptu_config_error(incomplete) == expected - with pytest.raises(ValueError) as raised: + with pytest.raises(ValueError, match="PTU configuration on model 'gpt") as raised: _ptu_router(model_info=incomplete) assert expected in str(raised.value) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 041c60e0ba6..075b455e4b5 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4370,7 +4370,7 @@ class TestVertexEmbeddingEncodingFormat: assert "encoding_format" not in optional_params def test_encoding_format_base64_still_rejected_without_drop_params(self): - with pytest.raises(Exception) as excinfo: + with pytest.raises(Exception, match='To drop these, set `litellm\\.drop_params=True` or for proxy') as excinfo: litellm.utils.get_optional_params_embeddings( model="gemini-embedding-001", encoding_format="base64", diff --git a/tests/test_litellm/types/test_router.py b/tests/test_litellm/types/test_router.py index 5ce5eca4954..accd3b32a0d 100644 --- a/tests/test_litellm/types/test_router.py +++ b/tests/test_litellm/types/test_router.py @@ -87,5 +87,5 @@ def test_pricing_strings_are_coerced_to_float(): def test_invalid_pricing_is_rejected(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='validation error for ModelInfo'): ModelInfo(id="x", input_cost_per_token="free") From 9697748f92d1a02c4b2cfc2768c20c39242503d4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 22:28:43 -0700 Subject: [PATCH 042/166] test: gate the already-hashed pass-through on the provenance flag The spend-log helper no longer treats a 64-hex shape as proof a value was already hashed, so this case has to say where the hash came from. Reconciles the test that came in with #31799 against that change. --- .../proxy/spend_tracking/test_spend_tracking_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 706466d5033..e34d4f389e1 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -2927,7 +2927,7 @@ class TestSpendLogKeyRedaction: def test_already_hashed_key_unchanged(self): hashed = "bcfe8173f5447f10be0e7fb37aaa8b97829d5c9e0498232152f9d123456789ab" - assert _redact_logged_api_key(hashed) == hashed + assert _redact_logged_api_key(hashed, already_hashed=True) == hashed def test_bearer_prefixed_non_sk_key_is_hashed(self): raw = "Bearer some-other-token-format" From fb417a556300fd6a983af66e28e0ee759d41a041 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 22:38:04 -0700 Subject: [PATCH 043/166] fix(spend-tracking): tie the already-hashed pass-through to provenance The hashed-jwt branch trusted the value's shape alone, so a caller-supplied key in that shape was stored unhashed. Both pass-throughs now require the value to match the auth-time user_api_key_hash, and the shape check is a full match. --- .../spend_tracking/spend_tracking_utils.py | 14 ++++++----- .../test_spend_tracking_utils.py | 25 +++++++++++++++++-- 2 files changed, 31 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 4592add1032..692200b856c 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -64,7 +64,11 @@ def _is_master_key(api_key: str | None, _master_key: str | None) -> bool: return secrets.compare_digest(api_key, _master_key) -_HASHED_JWT_RE = re.compile(r"^hashed-jwt-[a-fA-F0-9]{64}$") +_HASHED_JWT_RE = re.compile(r"hashed-jwt-[a-fA-F0-9]{64}") + + +def _is_prehashed_key_shape(value: str) -> bool: + return is_valid_sha256_hash(value) or _HASHED_JWT_RE.fullmatch(value) is not None def _redact_logged_api_key(value: str | None, *, already_hashed: bool = False) -> str | None: @@ -73,9 +77,7 @@ def _redact_logged_api_key(value: str | None, *, already_hashed: bool = False) - stripped: Final = re.sub(r"(?i)^bearer ", "", value) if not stripped: return None - if already_hashed and is_valid_sha256_hash(stripped): - return stripped - if _HASHED_JWT_RE.match(stripped): + if already_hashed and _is_prehashed_key_shape(stripped): return stripped return hash_token(stripped) @@ -136,7 +138,7 @@ def _get_spend_logs_metadata( _raw_key: Final = clean_metadata.get("user_api_key") _trusted_hash: Final = metadata.get("user_api_key_hash") _already_hashed: Final = ( - isinstance(_trusted_hash, str) and is_valid_sha256_hash(_trusted_hash) and _trusted_hash == _raw_key + isinstance(_trusted_hash, str) and _is_prehashed_key_shape(_trusted_hash) and _trusted_hash == _raw_key ) clean_metadata["user_api_key"] = _redact_logged_api_key(_raw_key, already_hashed=_already_hashed) clean_metadata["applied_guardrails"] = applied_guardrails @@ -296,7 +298,7 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs standard_logging_total_tokens = standard_logging_payload.get("total_tokens", 0) _trusted_hash = metadata.get("user_api_key_hash") _key_already_hashed = ( - isinstance(_trusted_hash, str) and is_valid_sha256_hash(_trusted_hash) and _trusted_hash == api_key + isinstance(_trusted_hash, str) and _is_prehashed_key_shape(_trusted_hash) and _trusted_hash == api_key ) api_key = _redact_logged_api_key(api_key, already_hashed=_key_already_hashed) or "" diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index e34d4f389e1..92e78e04512 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -2653,10 +2653,24 @@ def test_redact_logged_api_key_long_opaque_token_is_hashed(): def test_redact_logged_api_key_hashed_jwt_passes_through(): jwt_hash = "hashed-jwt-" + "a" * 64 - result = _redact_logged_api_key(jwt_hash) + result = _redact_logged_api_key(jwt_hash, already_hashed=True) assert result == jwt_hash +def test_redact_logged_api_key_hashed_jwt_shape_without_provenance_is_hashed(): + lookalike = "hashed-jwt-" + "a" * 64 + result = _redact_logged_api_key(lookalike) + assert result == hash_token(lookalike) + assert result != lookalike + + +def test_redact_logged_api_key_hashed_jwt_trailing_newline_is_hashed(): + trailing = "hashed-jwt-" + "a" * 64 + "\n" + result = _redact_logged_api_key(trailing, already_hashed=True) + assert result == hash_token(trailing) + assert result != trailing + + def test_redact_logged_api_key_hashed_jwt_short_suffix_is_hashed(): short_jwt = "hashed-jwt-tooshort" result = _redact_logged_api_key(short_jwt) @@ -2730,10 +2744,17 @@ def test_get_spend_logs_metadata_provenance_bypass_requires_hash_match(): def test_get_spend_logs_metadata_hashed_jwt_unchanged(): jwt_hash = "hashed-jwt-" + "b" * 64 - meta = _get_spend_logs_metadata({"user_api_key": jwt_hash}) + meta = _get_spend_logs_metadata({"user_api_key": jwt_hash, "user_api_key_hash": jwt_hash}) assert meta["user_api_key"] == jwt_hash +def test_get_spend_logs_metadata_hashed_jwt_shape_without_provenance_is_hashed(): + lookalike = "hashed-jwt-" + "b" * 64 + meta = _get_spend_logs_metadata({"user_api_key": lookalike}) + assert meta["user_api_key"] == hash_token(lookalike) + assert meta["user_api_key"] != lookalike + + def test_get_spend_logs_metadata_none_key_is_none(): meta = _get_spend_logs_metadata({"user_api_key": None}) assert meta["user_api_key"] is None From c7b34da079e962d006e3b109739203703fa152f6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 23:19:17 -0700 Subject: [PATCH 044/166] test(proxy): keep a leaked llm_router out of the next test in the worker The proxy conftest already snapshots master_key and prisma_client around every test, because a value left behind on litellm.proxy.proxy_server poisons the rest of the xdist worker. llm_router has the same problem. The PTU rollup reads the running router out of sys.modules, so a router a sibling test left behind lands in its deployment scan and three test_ptu_flat_cost_rollup tests fail or pass depending on how xdist happens to split the shard. --- tests/test_litellm/proxy/conftest.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/conftest.py b/tests/test_litellm/proxy/conftest.py index 61752997f0f..65e12b7d777 100644 --- a/tests/test_litellm/proxy/conftest.py +++ b/tests/test_litellm/proxy/conftest.py @@ -18,6 +18,7 @@ from prisma.errors import ClientNotConnectedError _PROXY_MODULE_GLOBALS_TO_ISOLATE = ( "master_key", "prisma_client", + "llm_router", ) @@ -56,7 +57,10 @@ def pytest_runtest_setup(item): Without this, a leaked value (e.g. master_key set by a sibling test) flips the auth short-circuit in user_api_key_auth and causes unrelated - tests in the same xdist worker to return 401 instead of 200. + tests in the same xdist worker to return 401 instead of 200. A leaked + llm_router does the same to anything that reads the running router out + of sys.modules, such as the PTU rollup's deployment scan, which then + counts a sibling test's deployments as if the proxy owned them. This must be a hook pair, not an autouse fixture: an autouse fixture in the root conftest requests monkeypatch, so monkeypatch's undo stack From a50590f32454ee845775a04087d6cd42d249f37d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 23:41:55 -0700 Subject: [PATCH 045/166] fix(spend-tracking): keep the master key alias readable in spend logs Master-key auth stamps the stable alias litellm_proxy_master_key instead of the raw key, so spend logs carry a readable, non-secret identifier for those rows. The new redaction path only recognized sha256 and hashed-jwt shapes, so it hashed that alias and broke continuity with every master-key row written before this change. The alias joins the recognized non-secret values, still behind the same provenance gate, so a caller who sends the alias string as their own bearer token still gets it hashed. --- .../spend_tracking/spend_tracking_utils.py | 27 ++++---- .../test_spend_tracking_utils.py | 64 +++++++++++++++++-- 2 files changed, 75 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 692200b856c..cba1f9069d3 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -10,6 +10,7 @@ from pydantic import BaseModel import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import ( + LITELLM_PROXY_MASTER_KEY_ALIAS, LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, REDACTED_BY_LITELM_STRING, @@ -67,17 +68,21 @@ def _is_master_key(api_key: str | None, _master_key: str | None) -> bool: _HASHED_JWT_RE = re.compile(r"hashed-jwt-[a-fA-F0-9]{64}") -def _is_prehashed_key_shape(value: str) -> bool: - return is_valid_sha256_hash(value) or _HASHED_JWT_RE.fullmatch(value) is not None +def _is_non_secret_key_value(value: str) -> bool: + return ( + value == LITELLM_PROXY_MASTER_KEY_ALIAS + or is_valid_sha256_hash(value) + or _HASHED_JWT_RE.fullmatch(value) is not None + ) -def _redact_logged_api_key(value: str | None, *, already_hashed: bool = False) -> str | None: +def _redact_logged_api_key(value: str | None, *, already_redacted: bool = False) -> str | None: if not isinstance(value, str) or not value: return None stripped: Final = re.sub(r"(?i)^bearer ", "", value) if not stripped: return None - if already_hashed and _is_prehashed_key_shape(stripped): + if already_redacted and _is_non_secret_key_value(stripped): return stripped return hash_token(stripped) @@ -137,10 +142,10 @@ def _get_spend_logs_metadata( clean_metadata: Final = SpendLogsMetadata(**{key: metadata.get(key) for key in SpendLogsMetadata.__annotations__}) _raw_key: Final = clean_metadata.get("user_api_key") _trusted_hash: Final = metadata.get("user_api_key_hash") - _already_hashed: Final = ( - isinstance(_trusted_hash, str) and _is_prehashed_key_shape(_trusted_hash) and _trusted_hash == _raw_key + _already_redacted: Final = ( + isinstance(_trusted_hash, str) and _is_non_secret_key_value(_trusted_hash) and _trusted_hash == _raw_key ) - clean_metadata["user_api_key"] = _redact_logged_api_key(_raw_key, already_hashed=_already_hashed) + clean_metadata["user_api_key"] = _redact_logged_api_key(_raw_key, already_redacted=_already_redacted) clean_metadata["applied_guardrails"] = applied_guardrails clean_metadata["batch_models"] = batch_models clean_metadata["mcp_tool_call_metadata"] = mcp_tool_call_metadata @@ -297,10 +302,10 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs standard_logging_completion_tokens = standard_logging_payload.get("completion_tokens", 0) standard_logging_total_tokens = standard_logging_payload.get("total_tokens", 0) _trusted_hash = metadata.get("user_api_key_hash") - _key_already_hashed = ( - isinstance(_trusted_hash, str) and _is_prehashed_key_shape(_trusted_hash) and _trusted_hash == api_key + _key_already_redacted = ( + isinstance(_trusted_hash, str) and _is_non_secret_key_value(_trusted_hash) and _trusted_hash == api_key ) - api_key = _redact_logged_api_key(api_key, already_hashed=_key_already_hashed) or "" + api_key = _redact_logged_api_key(api_key, already_redacted=_key_already_redacted) or "" if ( standard_logging_payload is not None @@ -308,7 +313,7 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs api_key = ( api_key or _redact_logged_api_key( - standard_logging_payload["metadata"].get("user_api_key_hash"), already_hashed=True + standard_logging_payload["metadata"].get("user_api_key_hash"), already_redacted=True ) or "" ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 92e78e04512..f2dd66ee677 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -2625,7 +2625,7 @@ def test_redact_logged_api_key_non_sk_raw_key_is_hashed(): def test_redact_logged_api_key_already_valid_sha256_passes_through_with_flag(): already_hashed = hash_token("sk-some-key") assert len(already_hashed) == 64 - result = _redact_logged_api_key(already_hashed, already_hashed=True) + result = _redact_logged_api_key(already_hashed, already_redacted=True) assert result == already_hashed assert hash_token(already_hashed) != result # no double-hash @@ -2653,7 +2653,7 @@ def test_redact_logged_api_key_long_opaque_token_is_hashed(): def test_redact_logged_api_key_hashed_jwt_passes_through(): jwt_hash = "hashed-jwt-" + "a" * 64 - result = _redact_logged_api_key(jwt_hash, already_hashed=True) + result = _redact_logged_api_key(jwt_hash, already_redacted=True) assert result == jwt_hash @@ -2666,7 +2666,7 @@ def test_redact_logged_api_key_hashed_jwt_shape_without_provenance_is_hashed(): def test_redact_logged_api_key_hashed_jwt_trailing_newline_is_hashed(): trailing = "hashed-jwt-" + "a" * 64 + "\n" - result = _redact_logged_api_key(trailing, already_hashed=True) + result = _redact_logged_api_key(trailing, already_redacted=True) assert result == hash_token(trailing) assert result != trailing @@ -2680,6 +2680,33 @@ def test_redact_logged_api_key_hashed_jwt_short_suffix_is_hashed(): assert result == hash_token(short_jwt) +def test_redact_logged_api_key_master_key_alias_passes_through(): + from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS + + result = _redact_logged_api_key(LITELLM_PROXY_MASTER_KEY_ALIAS, already_redacted=True) + assert result == LITELLM_PROXY_MASTER_KEY_ALIAS + + +def test_redact_logged_api_key_master_key_alias_without_provenance_is_hashed(): + from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS + + result = _redact_logged_api_key(LITELLM_PROXY_MASTER_KEY_ALIAS) + assert result == hash_token(LITELLM_PROXY_MASTER_KEY_ALIAS) + assert result != LITELLM_PROXY_MASTER_KEY_ALIAS + + +def test_get_spend_logs_metadata_keeps_master_key_alias_readable(): + from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS + + meta = _get_spend_logs_metadata( + { + "user_api_key": LITELLM_PROXY_MASTER_KEY_ALIAS, + "user_api_key_hash": LITELLM_PROXY_MASTER_KEY_ALIAS, + } + ) + assert meta["user_api_key"] == LITELLM_PROXY_MASTER_KEY_ALIAS + + def test_redact_logged_api_key_bearer_only_returns_none(): # "bearer " with nothing after stripping is equivalent to no key assert _redact_logged_api_key("bearer ") is None @@ -2948,7 +2975,7 @@ class TestSpendLogKeyRedaction: def test_already_hashed_key_unchanged(self): hashed = "bcfe8173f5447f10be0e7fb37aaa8b97829d5c9e0498232152f9d123456789ab" - assert _redact_logged_api_key(hashed, already_hashed=True) == hashed + assert _redact_logged_api_key(hashed, already_redacted=True) == hashed def test_bearer_prefixed_non_sk_key_is_hashed(self): raw = "Bearer some-other-token-format" @@ -2995,6 +3022,33 @@ def test_get_logging_payload_non_sk_raw_key_both_fields_hashed(): assert len(parsed_meta["user_api_key"]) == 64 +def test_get_logging_payload_keeps_master_key_alias_readable(): + from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS + + kwargs = { + "model": "openai/gpt-4.1", + "messages": [{"role": "user", "content": "Hello"}], + "call_type": "acompletion", + "litellm_params": { + "metadata": { + "user_api_key": LITELLM_PROXY_MASTER_KEY_ALIAS, + "user_api_key_hash": LITELLM_PROXY_MASTER_KEY_ALIAS, + "user_api_key_user_id": "test_user", + } + }, + } + payload = get_logging_payload( + kwargs=kwargs, + response_obj=Exception("error"), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + + assert payload["api_key"] == LITELLM_PROXY_MASTER_KEY_ALIAS + parsed_meta = json.loads(payload["metadata"]) + assert parsed_meta["user_api_key"] == LITELLM_PROXY_MASTER_KEY_ALIAS + + @patch("litellm.proxy.proxy_server.master_key", None) @patch("litellm.proxy.proxy_server.general_settings", {}) def test_get_logging_payload_hashes_bearer_prefixed_api_key(): @@ -3503,7 +3557,7 @@ def test_redact_logged_api_key_partial_sha256_is_hashed(): def test_redact_logged_api_key_bearer_already_hashed_passes_through_with_flag(): already_hashed = hash_token("sk-some-key") assert len(already_hashed) == 64 - result = _redact_logged_api_key(f"Bearer {already_hashed}", already_hashed=True) + result = _redact_logged_api_key(f"Bearer {already_hashed}", already_redacted=True) assert result == already_hashed assert hash_token(already_hashed) != result From 7d9e3756980135699a43f5c3c3d892a87b7ec842 Mon Sep 17 00:00:00 2001 From: bhuvan2134686 Date: Fri, 21 Aug 2026 17:28:12 +1000 Subject: [PATCH 046/166] fix(scx-ai): use the published scx.ai rates and the scx_ai docs url Applies the review suggestions. The cost map now carries the rates published on https://scx.ai/pricing, GLM-5.2 at 0.61 in, 0.22 cached, 1.98 out and Qwen3.8-Max at 1.65 in, 0.21 cached, 4.99 out per million tokens, and cites that page as the source rather than a third party gateway. The provider link is corrected to https://docs.litellm.ai/docs/providers/scx_ai to match the page that shipped as scx_ai.md. Both the primary files and their backup mirrors are updated. --- .../model_prices_and_context_window_backup.json | 14 +++++++------- litellm/provider_endpoints_support_backup.json | 2 +- model_prices_and_context_window.json | 14 +++++++------- provider_endpoints_support.json | 2 +- 4 files changed, 16 insertions(+), 16 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 222f4dd4db6..6c86c56d52a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -36725,15 +36725,15 @@ "source": "https://cloud.sambanova.ai/plans/pricing" }, "scx-ai/GLM-5.2": { - "cache_read_input_token_cost": 1.375e-07, - "input_cost_per_token": 5.5e-07, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 6.1e-07, "litellm_provider": "scx-ai", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 1.9255e-06, - "source": "https://llmgateway.io/models/glm-5.2/scx-ai-gp", + "output_cost_per_token": 1.98e-06, + "source": "https://scx.ai/pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -36743,14 +36743,14 @@ }, "scx-ai/Qwen3.8-Max": { "cache_read_input_token_cost": 2.1e-07, - "input_cost_per_token": 1.815e-06, + "input_cost_per_token": 1.65e-06, "litellm_provider": "scx-ai", "max_input_tokens": 1000000, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 5.4461e-06, - "source": "https://llmgateway.io/models/qwen3.8-max/scx-ai-gp", + "output_cost_per_token": 4.99e-06, + "source": "https://scx.ai/pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index c1928c34349..86c14fb4cd8 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -2029,7 +2029,7 @@ }, "scx-ai": { "display_name": "SCX.ai (`scx-ai`)", - "url": "https://docs.litellm.ai/docs/providers/scx-ai", + "url": "https://docs.litellm.ai/docs/providers/scx_ai", "endpoints": { "chat_completions": true, "messages": false, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 222f4dd4db6..6c86c56d52a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -36725,15 +36725,15 @@ "source": "https://cloud.sambanova.ai/plans/pricing" }, "scx-ai/GLM-5.2": { - "cache_read_input_token_cost": 1.375e-07, - "input_cost_per_token": 5.5e-07, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 6.1e-07, "litellm_provider": "scx-ai", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 1.9255e-06, - "source": "https://llmgateway.io/models/glm-5.2/scx-ai-gp", + "output_cost_per_token": 1.98e-06, + "source": "https://scx.ai/pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -36743,14 +36743,14 @@ }, "scx-ai/Qwen3.8-Max": { "cache_read_input_token_cost": 2.1e-07, - "input_cost_per_token": 1.815e-06, + "input_cost_per_token": 1.65e-06, "litellm_provider": "scx-ai", "max_input_tokens": 1000000, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 5.4461e-06, - "source": "https://llmgateway.io/models/qwen3.8-max/scx-ai-gp", + "output_cost_per_token": 4.99e-06, + "source": "https://scx.ai/pricing", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 3da0ec7d6b4..1d8d374c2c4 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -2263,7 +2263,7 @@ }, "scx-ai": { "display_name": "SCX.ai (`scx-ai`)", - "url": "https://docs.litellm.ai/docs/providers/scx-ai", + "url": "https://docs.litellm.ai/docs/providers/scx_ai", "endpoints": { "chat_completions": true, "messages": false, From 86efa2bcfde555e519140f6d9721e9570aada2d2 Mon Sep 17 00:00:00 2001 From: longwind48 Date: Fri, 21 Aug 2026 17:27:28 +0800 Subject: [PATCH 047/166] feat(bedrock): serve gpt-5.6 cross-region inference profiles on bedrock runtime GPT-5.6 Sol, Terra and Luna reached the bedrock-runtime data plane on 2026-08-17, separately from the existing bedrock-mantle path. On runtime they are served only through cross-region inference profiles, so bedrock/us.openai.gpt-5.6-* had no cost map entry and fell through to the Invoke route, which rewrites the token cap to max_tokens and is rejected as unsupported_parameter on both /v1/chat/completions and /v1/responses. Register the Geo and Global profiles as bedrock_converse so routing reaches Converse, which AWS documents and serves for these models, and price each profile from its own published rate table. No bare key: the control plane reports inferenceTypesSupported INFERENCE_PROFILE with no on-demand throughput, so a bare id is not invocable. Declare the published cache-read and cache-write rates. Bedrock rejects an explicit cachePoint block for these models, so supports_prompt_caching stays off, but it caches long prefixes implicitly and reports the cache tokens in usage either way. Without the cost fields a cache-read turn bills only its uncached tokens: measured against live Bedrock, a 15609-token cached prefix came to $0.000176 instead of $0.00876095. Clients that resend a long prefix every turn are the worst affected. Reasoning stays unadvertised. Converse rejects the Anthropic-shaped thinking block LiteLLM sends for reasoning_effort; the shape these models accept is additionalModelRequestFields {"reasoning": {"effort": ...}}, which needs a transform change tracked by #34105. Advertising it without that change is what made the earlier attempt in #37307 fail. --- ...odel_prices_and_context_window_backup.json | 150 ++++++++++ model_prices_and_context_window.json | 150 ++++++++++ ..._cross_region_inference_profile_mapping.py | 258 +++++++++++++++++- 3 files changed, 557 insertions(+), 1 deletion(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 91c10d13e8e..be8e6da5f59 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -48545,6 +48545,156 @@ "supports_tool_choice": true, "supports_vision": true }, + "us.openai.gpt-5.6-sol": { + "input_cost_per_token": 5.5e-06, + "input_cost_per_token_above_272k_tokens": 1.1e-05, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1.375e-05, + "cache_read_input_token_cost": 5.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, + "output_cost_per_token": 3.3e-05, + "output_cost_per_token_above_272k_tokens": 4.95e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "global.openai.gpt-5.6-sol": { + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1.25e-05, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "us.openai.gpt-5.6-terra": { + "input_cost_per_token": 2.2e-06, + "input_cost_per_token_above_272k_tokens": 4.4e-06, + "cache_creation_input_token_cost": 2.75e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5.5e-06, + "cache_read_input_token_cost": 2.2e-07, + "cache_read_input_token_cost_above_272k_tokens": 4.4e-07, + "output_cost_per_token": 1.32e-05, + "output_cost_per_token_above_272k_tokens": 1.98e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "global.openai.gpt-5.6-terra": { + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_272k_tokens": 4e-06, + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_272k_tokens": 4e-07, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_272k_tokens": 1.8e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "us.openai.gpt-5.6-luna": { + "input_cost_per_token": 2.2e-07, + "input_cost_per_token_above_272k_tokens": 4.4e-07, + "cache_creation_input_token_cost": 2.75e-07, + "cache_creation_input_token_cost_above_272k_tokens": 5.5e-07, + "cache_read_input_token_cost": 2.2e-08, + "cache_read_input_token_cost_above_272k_tokens": 4.4e-08, + "output_cost_per_token": 1.32e-06, + "output_cost_per_token_above_272k_tokens": 1.98e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "global.openai.gpt-5.6-luna": { + "input_cost_per_token": 2e-07, + "input_cost_per_token_above_272k_tokens": 4e-07, + "cache_creation_input_token_cost": 2.5e-07, + "cache_creation_input_token_cost_above_272k_tokens": 5e-07, + "cache_read_input_token_cost": 2e-08, + "cache_read_input_token_cost_above_272k_tokens": 4e-08, + "output_cost_per_token": 1.2e-06, + "output_cost_per_token_above_272k_tokens": 1.8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, "bedrock_mantle/openai.gpt-5.5": { "input_cost_per_token": 5.5e-06, "cache_read_input_token_cost": 5.5e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 91c10d13e8e..be8e6da5f59 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -48545,6 +48545,156 @@ "supports_tool_choice": true, "supports_vision": true }, + "us.openai.gpt-5.6-sol": { + "input_cost_per_token": 5.5e-06, + "input_cost_per_token_above_272k_tokens": 1.1e-05, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1.375e-05, + "cache_read_input_token_cost": 5.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, + "output_cost_per_token": 3.3e-05, + "output_cost_per_token_above_272k_tokens": 4.95e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "global.openai.gpt-5.6-sol": { + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1.25e-05, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "us.openai.gpt-5.6-terra": { + "input_cost_per_token": 2.2e-06, + "input_cost_per_token_above_272k_tokens": 4.4e-06, + "cache_creation_input_token_cost": 2.75e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5.5e-06, + "cache_read_input_token_cost": 2.2e-07, + "cache_read_input_token_cost_above_272k_tokens": 4.4e-07, + "output_cost_per_token": 1.32e-05, + "output_cost_per_token_above_272k_tokens": 1.98e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "global.openai.gpt-5.6-terra": { + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_272k_tokens": 4e-06, + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_272k_tokens": 4e-07, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_272k_tokens": 1.8e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "us.openai.gpt-5.6-luna": { + "input_cost_per_token": 2.2e-07, + "input_cost_per_token_above_272k_tokens": 4.4e-07, + "cache_creation_input_token_cost": 2.75e-07, + "cache_creation_input_token_cost_above_272k_tokens": 5.5e-07, + "cache_read_input_token_cost": 2.2e-08, + "cache_read_input_token_cost_above_272k_tokens": 4.4e-08, + "output_cost_per_token": 1.32e-06, + "output_cost_per_token_above_272k_tokens": 1.98e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "global.openai.gpt-5.6-luna": { + "input_cost_per_token": 2e-07, + "input_cost_per_token_above_272k_tokens": 4e-07, + "cache_creation_input_token_cost": 2.5e-07, + "cache_creation_input_token_cost_above_272k_tokens": 5e-07, + "cache_read_input_token_cost": 2e-08, + "cache_read_input_token_cost_above_272k_tokens": 4e-08, + "output_cost_per_token": 1.2e-06, + "output_cost_per_token_above_272k_tokens": 1.8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, "bedrock_mantle/openai.gpt-5.5": { "input_cost_per_token": 5.5e-06, "cache_read_input_token_cost": 5.5e-07, diff --git a/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py b/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py index 3a27f3ed002..22aba59fb5d 100644 --- a/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py +++ b/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py @@ -1,13 +1,132 @@ """Test Bedrock cross-region inference profile model mapping""" +import json import os import sys +from functools import lru_cache +from pathlib import Path +from typing import NamedTuple + +import pytest sys.path.insert(0, os.path.abspath("../../../..")) +import litellm +from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig +from litellm.llms.bedrock.common_utils import BedrockModelInfo from litellm.utils import _get_model_info_helper from litellm.cost_calculator import completion_cost -from litellm.types.utils import ModelResponse, Usage, Choices, Message +from litellm.types.utils import ( + Choices, + Message, + ModelResponse, + PromptTokensDetailsWrapper, + Usage, +) + + +@pytest.fixture +def local_model_cost_map(monkeypatch): + """Resolve models against this checkout's cost map instead of the network-fetched + ``main`` copy, which lags this branch until merge.""" + original_converse_models = set(litellm.bedrock_converse_models) + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + try: + litellm.bedrock_converse_models.update( + key + for key, value in litellm.model_cost.items() + if isinstance(value, dict) + and value.get("litellm_provider") == "bedrock_converse" + ) + yield + finally: + litellm.bedrock_converse_models.clear() + litellm.bedrock_converse_models.update(original_converse_models) + litellm.get_model_info.cache_clear() + + +class GptProfile(NamedTuple): + model_id: str + input_cost: float + input_cost_above_272k: float + cache_write: float + cache_write_above_272k: float + cache_read: float + cache_read_above_272k: float + output_cost: float + output_cost_above_272k: float + + +GPT_5_6_PROFILES = [ + GptProfile( + model_id="us.openai.gpt-5.6-sol", + input_cost=5.5e-06, input_cost_above_272k=1.1e-05, + cache_write=6.875e-06, cache_write_above_272k=1.375e-05, + cache_read=5.5e-07, cache_read_above_272k=1.1e-06, + output_cost=3.3e-05, output_cost_above_272k=4.95e-05, + ), + GptProfile( + model_id="global.openai.gpt-5.6-sol", + input_cost=5e-06, input_cost_above_272k=1e-05, + cache_write=6.25e-06, cache_write_above_272k=1.25e-05, + cache_read=5e-07, cache_read_above_272k=1e-06, + output_cost=3e-05, output_cost_above_272k=4.5e-05, + ), + GptProfile( + model_id="us.openai.gpt-5.6-terra", + input_cost=2.2e-06, input_cost_above_272k=4.4e-06, + cache_write=2.75e-06, cache_write_above_272k=5.5e-06, + cache_read=2.2e-07, cache_read_above_272k=4.4e-07, + output_cost=1.32e-05, output_cost_above_272k=1.98e-05, + ), + GptProfile( + model_id="global.openai.gpt-5.6-terra", + input_cost=2e-06, input_cost_above_272k=4e-06, + cache_write=2.5e-06, cache_write_above_272k=5e-06, + cache_read=2e-07, cache_read_above_272k=4e-07, + output_cost=1.2e-05, output_cost_above_272k=1.8e-05, + ), + GptProfile( + model_id="us.openai.gpt-5.6-luna", + input_cost=2.2e-07, input_cost_above_272k=4.4e-07, + cache_write=2.75e-07, cache_write_above_272k=5.5e-07, + cache_read=2.2e-08, cache_read_above_272k=4.4e-08, + output_cost=1.32e-06, output_cost_above_272k=1.98e-06, + ), + GptProfile( + model_id="global.openai.gpt-5.6-luna", + input_cost=2e-07, input_cost_above_272k=4e-07, + cache_write=2.5e-07, cache_write_above_272k=5e-07, + cache_read=2e-08, cache_read_above_272k=4e-08, + output_cost=1.2e-06, output_cost_above_272k=1.8e-06, + ), +] + + +@lru_cache(maxsize=1) +def _packaged_cost_map(): + """The map litellm actually resolves against, for fields ModelInfoBase drops.""" + path = Path(litellm.__file__).parent / "model_prices_and_context_window_backup.json" + return json.loads(path.read_text()) + + +def _bedrock_response(model, usage): + return ModelResponse( + id="test", + created=1234567890, + model=model, + object="chat.completion", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message(content="OK", role="assistant"), + ) + ], + usage=usage, + ) def test_bedrock_cross_region_inference_profile_mapping(): @@ -52,3 +171,140 @@ def test_proxy_cost_calculation_scenario(): ) expected_cost = (100 * 8e-07) + (50 * 4e-06) assert cost == expected_cost + + +@pytest.mark.parametrize("profile", GPT_5_6_PROFILES, ids=lambda p: p.model_id) +def test_bedrock_gpt_5_6_profiles_route_to_converse(profile, local_model_cost_map): + """GPT-5.6 is served by Converse on bedrock-runtime, never by Invoke.""" + assert BedrockModelInfo.get_bedrock_route(f"bedrock/{profile.model_id}") == "converse" + + +@pytest.mark.parametrize("profile", GPT_5_6_PROFILES, ids=lambda p: p.model_id) +def test_bedrock_gpt_5_6_published_rates(profile, local_model_cost_map): + """Geo and Global profiles carry their own published rates, per context tier.""" + model_info = _get_model_info_helper( + model=f"bedrock/{profile.model_id}", custom_llm_provider="bedrock" + ) + + assert model_info["litellm_provider"] == "bedrock_converse" + assert model_info["mode"] == "chat" + assert model_info["max_input_tokens"] == 1000000 + assert model_info["input_cost_per_token"] == profile.input_cost + assert ( + model_info["input_cost_per_token_above_272k_tokens"] + == profile.input_cost_above_272k + ) + assert model_info["output_cost_per_token"] == profile.output_cost + assert ( + model_info["output_cost_per_token_above_272k_tokens"] + == profile.output_cost_above_272k + ) + assert model_info["cache_creation_input_token_cost"] == profile.cache_write + assert ( + model_info["cache_creation_input_token_cost_above_272k_tokens"] + == profile.cache_write_above_272k + ) + assert model_info["cache_read_input_token_cost"] == profile.cache_read + assert ( + model_info["cache_read_input_token_cost_above_272k_tokens"] + == profile.cache_read_above_272k + ) + + +def test_bedrock_gpt_5_6_above_272k_tier_applies_to_cost(local_model_cost_map): + """A prompt over 272K tokens is billed at the long-context rate, not the base rate.""" + response = _bedrock_response( + "bedrock/us.openai.gpt-5.6-sol", + Usage(prompt_tokens=300000, completion_tokens=1000, total_tokens=301000), + ) + + cost = completion_cost( + completion_response=response, + model="bedrock/us.openai.gpt-5.6-sol", + custom_llm_provider="bedrock", + ) + + assert cost == pytest.approx((300000 * 1.1e-05) + (1000 * 4.95e-05), rel=1e-9) + + +def test_bedrock_gpt_5_6_bills_cache_read_tokens(local_model_cost_map): + """Bedrock caches long prefixes implicitly and reports them, so a cache-read turn + must be billed at the cache rate rather than dropped to zero.""" + usage = Usage( + prompt_tokens=15611, + completion_tokens=5, + total_tokens=15616, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=15609), + ) + response = _bedrock_response("bedrock/us.openai.gpt-5.6-sol", usage) + + cost = completion_cost( + completion_response=response, + model="bedrock/us.openai.gpt-5.6-sol", + custom_llm_provider="bedrock", + ) + + expected = (2 * 5.5e-06) + (15609 * 5.5e-07) + (5 * 3.3e-05) + assert cost == pytest.approx(expected, rel=1e-9) + # Without cache_read_input_token_cost the cached prefix bills at zero. + assert cost > (15611 * 5.5e-06) * 0.1 + + +def test_bedrock_gpt_5_6_bills_cache_write_tokens(local_model_cost_map): + """The write side of the same cache cycle is billed at the 30m cache-write rate.""" + usage = Usage( + prompt_tokens=15611, + completion_tokens=5, + total_tokens=15616, + cache_creation_input_tokens=15609, + ) + response = _bedrock_response("bedrock/us.openai.gpt-5.6-sol", usage) + + cost = completion_cost( + completion_response=response, + model="bedrock/us.openai.gpt-5.6-sol", + custom_llm_provider="bedrock", + ) + + expected = (2 * 5.5e-06) + (15609 * 6.875e-06) + (5 * 3.3e-05) + assert cost == pytest.approx(expected, rel=1e-9) + + +@pytest.mark.parametrize("profile", GPT_5_6_PROFILES, ids=lambda p: p.model_id) +def test_bedrock_gpt_5_6_advertises_only_converse_supported_features( + profile, local_model_cost_map +): + model_info = _get_model_info_helper( + model=f"bedrock/{profile.model_id}", custom_llm_provider="bedrock" + ) + + assert model_info["supports_function_calling"] is True + assert model_info["supports_tool_choice"] is True + assert model_info["supports_vision"] is True + + # Bedrock rejects an explicit cachePoint block for these models, so the flag that + # offers caller-driven caching stays off even though the cache rates are declared. + assert not model_info.get("supports_prompt_caching") + + # ModelInfoBase drops these two, so they are read from the map litellm resolves. + raw = _packaged_cost_map()[profile.model_id] + assert raw["supported_modalities"] == ["text", "image"] + assert raw["supported_output_modalities"] == ["text"] + # No bedrock_converse entry declares supported_endpoints; these models are reachable + # on chat completions and on the Responses API without it. + assert "supported_endpoints" not in raw + + +@pytest.mark.parametrize("profile", GPT_5_6_PROFILES, ids=lambda p: p.model_id) +def test_bedrock_gpt_5_6_offers_tools_but_not_reasoning(profile, local_model_cost_map): + """Converse rejects the Anthropic-shaped thinking block LiteLLM emits for + reasoning_effort, so neither reasoning param may be offered yet, while the tool + params these models do accept must be.""" + supported = AmazonConverseConfig().get_supported_openai_params( + model=f"bedrock/{profile.model_id}" + ) + + assert "tools" in supported + assert "tool_choice" in supported + assert "reasoning_effort" not in supported + assert "thinking" not in supported From 13d4074492aa03b4d35a62fc8ffb8de2ef40e8dc Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 21 Aug 2026 04:41:21 -0700 Subject: [PATCH 048/166] test(mcp): retire the last file of the dead tests/litellm mirror tests/litellm/ was a second mirror beside tests/test_litellm/ that no workflow, Makefile target, or CircleCI job ever named. Its other 33 files were reconciled during August 2026; this one stayed behind under a ci-coverage-allowlist entry asking a later pass to decide which of its five orphan behaviours still hold. They no longer hold as written: 25 of its 32 cases fail against today's code, because the file froze on the day it stopped being collected and the endpoints kept moving. Three of the five are already covered by the live twin, and better. test_get_request_base_url_xff_trust_gate parametrizes the trust gate in both directions, including the exact untrusted-caller case the orphan asserted, and the standard and legacy protected-resource shapes are both exercised through use_standard_pattern. The other two were the only tests anywhere for validate_trusted_redirect_uri under that same gate, so they are ported rather than dropped, rebuilt on the live file's request-mock conventions. Both directions are load-bearing: forcing is_request_from_trusted_proxy to True fails the untrusted case, forcing it to False fails the trusted one. 313 tests pass in the live file, up from 311. Dropping the dead file clears one zero-assert TQ001 violation, so its ceiling ratchets down with it. --- .github/ci-coverage-allowlist.yml | 10 - test-quality-budget.json | 2 +- .../mcp_server/test_discoverable_endpoints.py | 1268 ----------------- .../mcp_server/test_discoverable_endpoints.py | 49 + 4 files changed, 50 insertions(+), 1279 deletions(-) delete mode 100644 tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py diff --git a/.github/ci-coverage-allowlist.yml b/.github/ci-coverage-allowlist.yml index ff8fa864d4a..918589f84d1 100644 --- a/.github/ci-coverage-allowlist.yml +++ b/.github/ci-coverage-allowlist.yml @@ -48,16 +48,6 @@ test_paths: choice it informed is settled paths: - tests/code_coverage_tests/test_aio_http_image_conversion.py - - reason: >- - The last file of a second mirror that sat beside tests/test_litellm and ran nowhere. Its - other 33 files landed in the real mirror during August 2026, 30 as moves and 3 by merging - their bodies into the live file of the same name. This one cannot follow either route yet: - its live twin was rewritten from 1268 lines to 9434, and of the 19 tests here 5 have no - counterpart while 25 assertions fail against today's code, so what survives that rewrite - is a judgement about the endpoints, not a merge. Revisit by deciding which of the five - behaviours still hold - paths: - - tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py - reason: >- No job invokes this suite and its files mix pure transformation tests with ones driving live vendor vector stores, so assigning them needs a per-file decision diff --git a/test-quality-budget.json b/test-quality-budget.json index 1613c8c75cb..91ae881c83a 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -1,6 +1,6 @@ { "TQ001": { - "limit": 750 + "limit": 746 }, "TQ002": { "limit": 742 diff --git a/tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py deleted file mode 100644 index 2a8768df722..00000000000 --- a/tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ /dev/null @@ -1,1268 +0,0 @@ -"""Tests for MCP OAuth discoverable endpoints""" - -import pytest -from fastapi import HTTPException -from unittest.mock import AsyncMock, MagicMock, patch - -TRUSTED_PROXY_IP = "10.0.0.5" -TRUSTED_PROXY_RANGES = ["10.0.0.0/8"] - - -def set_request_from_trusted_proxy(mock_request): - mock_request.client = MagicMock() - mock_request.client.host = TRUSTED_PROXY_IP - - -@pytest.fixture -def trusted_proxy_origin_headers(): - with ( - patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.IPAddressUtils.is_request_from_trusted_proxy", - return_value=True, - ), - patch( - "litellm.proxy._experimental.mcp_server.oauth_utils.IPAddressUtils.is_request_from_trusted_proxy", - return_value=True, - ), - ): - yield - - -@pytest.mark.asyncio -async def test_authorize_endpoint_includes_response_type(): - """Test that authorize endpoint includes response_type=code parameter (fixes #15684)""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - authorize, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server - oauth2_server = MCPServer( - server_id="test_oauth_server", - name="test_oauth", - server_name="test_oauth", - alias="test_oauth", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="test_client_id", - client_secret="test_client_secret", - authorization_url="https://provider.com/oauth/authorize", - token_url="https://provider.com/oauth/token", - scopes=["read", "write"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://litellm.example.com/" - mock_request.headers = {} - - # Mock the encryption functions to avoid needing a signing key - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" - ) as mock_encrypt: - mock_encrypt.return_value = "mocked_encrypted_state" - - # Call authorize endpoint - response = await authorize( - request=mock_request, - client_id="test_client_id", - mcp_server_name="test_oauth", - redirect_uri="http://127.0.0.1:60108/callback", - state="test_state", - ) - - # Verify response is a redirect - assert response.status_code == 307 # FastAPI RedirectResponse default - - # Verify response_type is in the redirect URL - assert "response_type=code" in response.headers["location"] - assert "https://provider.com/oauth/authorize" in response.headers["location"] - assert "client_id=test_client_id" in response.headers["location"] - assert "scope=read+write" in response.headers["location"] - - -@pytest.mark.asyncio -async def test_authorize_endpoint_forwards_pkce_parameters(): - """Test that authorize endpoint forwards PKCE parameters (code_challenge and code_challenge_method)""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - authorize, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server (simulating Google OAuth) - oauth2_server = MCPServer( - server_id="google_mcp", - name="google_mcp", - server_name="google_mcp", - alias="google_mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="669428968603-test.apps.googleusercontent.com", - client_secret="GOCSPX-test_secret", - authorization_url="https://accounts.google.com/o/oauth2/v2/auth", - token_url="https://oauth2.googleapis.com/token", - scopes=["https://www.googleapis.com/auth/drive", "openid", "email"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://litellm-proxy.example.com/" - mock_request.headers = {} - - # Mock the encryption function - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" - ) as mock_encrypt: - mock_encrypt.return_value = "mocked_encrypted_state_with_pkce" - - # Call authorize endpoint with PKCE parameters - response = await authorize( - request=mock_request, - client_id="669428968603-test.apps.googleusercontent.com", - mcp_server_name="google_mcp", - redirect_uri="http://localhost:60108/callback", - state="test_client_state", - code_challenge="x6YH_qgwbvOzbsHDuL1sW9gYkR9-gObUiIB5RkPwxDk", - code_challenge_method="S256", - ) - - # Verify response is a redirect - assert response.status_code == 307 - - # Verify PKCE parameters are included in the redirect URL - location = response.headers["location"] - assert "https://accounts.google.com/o/oauth2/v2/auth" in location - assert "code_challenge=x6YH_qgwbvOzbsHDuL1sW9gYkR9-gObUiIB5RkPwxDk" in location - assert "code_challenge_method=S256" in location - assert "client_id=669428968603-test.apps.googleusercontent.com" in location - assert "response_type=code" in location - - -@pytest.mark.asyncio -async def test_token_endpoint_forwards_code_verifier(): - """Test that token endpoint forwards code_verifier for PKCE flow""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - token_endpoint, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server - oauth2_server = MCPServer( - server_id="google_mcp", - name="google_mcp", - server_name="google_mcp", - alias="google_mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="669428968603-test.apps.googleusercontent.com", - client_secret="GOCSPX-test_secret", - authorization_url="https://accounts.google.com/o/oauth2/v2/auth", - token_url="https://oauth2.googleapis.com/token", - scopes=["https://www.googleapis.com/auth/drive", "openid", "email"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://litellm-proxy.example.com/" - mock_request.headers = {} - - # Mock httpx client response - mock_response = MagicMock() - mock_response.json.return_value = { - "access_token": "ya29.test_access_token", - "token_type": "Bearer", - "expires_in": 3599, - "scope": "openid email https://www.googleapis.com/auth/drive", - } - mock_response.raise_for_status = MagicMock() - - # Mock the async httpx client with AsyncMock for async methods - from unittest.mock import AsyncMock - - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client" - ) as mock_get_client: - mock_async_client = MagicMock() - # Use AsyncMock for the async post method - mock_async_client.post = AsyncMock(return_value=mock_response) - mock_get_client.return_value = mock_async_client - - # Call token endpoint with code_verifier - response = await token_endpoint( - request=mock_request, - grant_type="authorization_code", - code="4/test_authorization_code", - redirect_uri="http://localhost:60108/callback", - client_id="669428968603-test.apps.googleusercontent.com", - mcp_server_name="google_mcp", - client_secret="GOCSPX-test_secret", - code_verifier="test_code_verifier_from_client", - ) - - # Verify that the token endpoint was called with code_verifier - mock_async_client.post.assert_called_once() - call_args = mock_async_client.post.call_args - - # Check the data parameter includes code_verifier - assert call_args[1]["data"]["code_verifier"] == "test_code_verifier_from_client" - assert call_args[1]["data"]["code"] == "4/test_authorization_code" - assert ( - call_args[1]["data"]["client_id"] - == "669428968603-test.apps.googleusercontent.com" - ) - assert call_args[1]["data"]["client_secret"] == "GOCSPX-test_secret" - assert call_args[1]["data"]["grant_type"] == "authorization_code" - - # Verify response - response_data = response.body - import json - - token_data = json.loads(response_data) - assert token_data["access_token"] == "ya29.test_access_token" - assert token_data["token_type"] == "Bearer" - - -@pytest.mark.asyncio -async def test_register_client_without_mcp_server_name_returns_dummy(): - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - register_client, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - global_mcp_server_manager.registry.clear() - - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://proxy.litellm.example/" - mock_request.headers = {} - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", - new=AsyncMock(return_value={}), - ): - result = await register_client(request=mock_request) - - assert result == { - "client_id": "dummy_client", - "client_secret": "dummy", - "redirect_uris": ["https://proxy.litellm.example/callback"], - } - - -@pytest.mark.asyncio -async def test_register_client_returns_existing_server_credentials(): - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - register_client, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - global_mcp_server_manager.registry.clear() - oauth2_server = MCPServer( - server_id="stored_server", - name="stored_server", - server_name="stored_server", - alias="stored_server", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="existing-client", - client_secret="existing-secret", - authorization_url="https://provider.example/oauth/authorize", - token_url="https://provider.example/oauth/token", - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://proxy.litellm.example/" - mock_request.headers = {} - - try: - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", - new=AsyncMock(return_value={}), - ): - result = await register_client( - request=mock_request, mcp_server_name=oauth2_server.server_name - ) - finally: - global_mcp_server_manager.registry.clear() - - assert result == { - "client_id": "stored_server", - "client_secret": "dummy", - "redirect_uris": ["https://proxy.litellm.example/callback"], - } - - -@pytest.mark.asyncio -async def test_register_client_remote_registration_success(): - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - register_client, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - global_mcp_server_manager.registry.clear() - oauth2_server = MCPServer( - server_id="remote_server", - name="remote_server", - server_name="remote_server", - alias="remote_server", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id=None, - client_secret=None, - authorization_url="https://provider.example/oauth/authorize", - token_url="https://provider.example/oauth/token", - registration_url="https://provider.example/oauth/register", - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://proxy.litellm.example/" - mock_request.headers = {} - - request_payload = { - "client_name": "Litellm Proxy", - "grant_types": ["authorization_code", "refresh_token"], - "response_types": ["code"], - "token_endpoint_auth_method": "client_secret_post", - } - - mock_response = MagicMock() - mock_response.json.return_value = { - "client_id": "generated-client", - "client_secret": "generated-secret", - } - mock_response.raise_for_status = MagicMock() - mock_async_client = MagicMock() - mock_async_client.post = AsyncMock(return_value=mock_response) - - try: - with ( - patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", - new=AsyncMock(return_value=request_payload), - ), - patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", - return_value=mock_async_client, - ), - ): - response = await register_client( - request=mock_request, mcp_server_name=oauth2_server.server_name - ) - finally: - global_mcp_server_manager.registry.clear() - - import json - - assert response.status_code == 200 - payload = json.loads(response.body.decode("utf-8")) - assert payload == mock_response.json.return_value - - mock_async_client.post.assert_called_once() - call_args = mock_async_client.post.call_args - assert call_args.args[0] == oauth2_server.registration_url - assert call_args.kwargs["headers"] == { - "Content-Type": "application/json", - "Accept": "application/json", - } - assert call_args.kwargs["json"]["redirect_uris"] == [ - "https://proxy.litellm.example/callback" - ] - assert call_args.kwargs["json"]["grant_types"] == request_payload["grant_types"] - assert ( - call_args.kwargs["json"]["token_endpoint_auth_method"] - == request_payload["token_endpoint_auth_method"] - ) - - -@pytest.mark.asyncio -async def test_authorize_endpoint_respects_x_forwarded_proto( - trusted_proxy_origin_headers, -): - """Test that authorize endpoint uses X-Forwarded-Proto header to construct correct redirect_uri""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - authorize, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server - oauth2_server = MCPServer( - server_id="test_oauth_server", - name="test_oauth", - server_name="test_oauth", - alias="test_oauth", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="test_client_id", - client_secret="test_client_secret", - authorization_url="https://provider.com/oauth/authorize", - token_url="https://provider.com/oauth/token", - scopes=["read", "write"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request with http base_url but X-Forwarded-Proto: https - mock_request = MagicMock(spec=Request) - mock_request.base_url = "http://litellm.example.com/" # HTTP - mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy - set_request_from_trusted_proxy(mock_request) - - # Mock the encryption functions - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" - ) as mock_encrypt: - mock_encrypt.return_value = "mocked_encrypted_state" - - # Call authorize endpoint - response = await authorize( - request=mock_request, - client_id="test_client_id", - mcp_server_name="test_oauth", - redirect_uri="http://127.0.0.1:60108/callback", - state="test_state", - ) - - # Verify redirect URL uses HTTPS in the redirect_uri parameter - location = response.headers["location"] - - # The redirect_uri parameter sent to the OAuth provider should use HTTPS - assert ( - "redirect_uri=https%3A%2F%2Flitellm.example.com%2Fcallback" in location - or "redirect_uri=https://litellm.example.com/callback" in location - ) - - -@pytest.mark.asyncio -async def test_token_endpoint_respects_x_forwarded_proto( - trusted_proxy_origin_headers, -): - """Test that token endpoint uses X-Forwarded-Proto header for redirect_uri""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - token_endpoint, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server - oauth2_server = MCPServer( - server_id="google_mcp", - name="google_mcp", - server_name="google_mcp", - alias="google_mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="test_client_id", - client_secret="test_secret", - authorization_url="https://accounts.google.com/o/oauth2/v2/auth", - token_url="https://oauth2.googleapis.com/token", - scopes=["openid", "email"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request with http base_url but X-Forwarded-Proto: https - mock_request = MagicMock(spec=Request) - mock_request.base_url = "http://litellm-proxy.example.com/" # HTTP - mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy - set_request_from_trusted_proxy(mock_request) - - # Mock httpx client response - mock_response = MagicMock() - mock_response.json.return_value = { - "access_token": "test_token", - "token_type": "Bearer", - "expires_in": 3599, - } - mock_response.raise_for_status = MagicMock() - - # Mock the async httpx client - mock_async_client = MagicMock() - mock_async_client.post = AsyncMock(return_value=mock_response) - - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client" - ) as mock_get_client: - mock_get_client.return_value = mock_async_client - - # Call token endpoint - await token_endpoint( - request=mock_request, - grant_type="authorization_code", - code="test_code", - redirect_uri="http://localhost:60108/callback", - client_id="test_client_id", - mcp_server_name="google_mcp", - client_secret="test_secret", - ) - - # Verify that the redirect_uri sent to the provider uses HTTPS - call_args = mock_async_client.post.call_args - assert ( - call_args[1]["data"]["redirect_uri"] - == "https://litellm-proxy.example.com/callback" - ) - - -@pytest.mark.asyncio -async def test_oauth_protected_resource_standard_pattern(): - """Test that oauth_protected_resource_mcp_standard returns standard MCP URL pattern (/mcp/{server_name})""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - oauth_protected_resource_mcp_standard, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server - oauth2_server = MCPServer( - server_id="test_server", - name="test_server", - server_name="test_server", - alias="test_server", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="test_client_id", - client_secret="test_client_secret", - authorization_url="https://provider.com/oauth/authorize", - token_url="https://provider.com/oauth/token", - scopes=["read", "write"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://litellm.example.com/" - mock_request.headers = {} - - # Call the standard pattern endpoint - response = await oauth_protected_resource_mcp_standard( - request=mock_request, - mcp_server_name="test_server", - ) - - # Verify response uses standard MCP pattern: /mcp/{server_name} - assert response["resource"] == "https://litellm.example.com/mcp/test_server" - assert ( - response["authorization_servers"][0] - == "https://litellm.example.com/test_server" - ) - assert response["scopes_supported"] == oauth2_server.scopes - - -@pytest.mark.asyncio -async def test_oauth_protected_resource_legacy_pattern(): - """Test that oauth_protected_resource_mcp returns legacy URL pattern (/{server_name}/mcp)""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - oauth_protected_resource_mcp, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server - oauth2_server = MCPServer( - server_id="test_server", - name="test_server", - server_name="test_server", - alias="test_server", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="test_client_id", - client_secret="test_client_secret", - authorization_url="https://provider.com/oauth/authorize", - token_url="https://provider.com/oauth/token", - scopes=["read", "write"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://litellm.example.com/" - mock_request.headers = {} - - # Call the legacy pattern endpoint - response = await oauth_protected_resource_mcp( - request=mock_request, - mcp_server_name="test_server", - ) - - # Verify response uses legacy pattern: /{server_name}/mcp - assert response["resource"] == "https://litellm.example.com/test_server/mcp" - assert ( - response["authorization_servers"][0] - == "https://litellm.example.com/test_server" - ) - assert response["scopes_supported"] == oauth2_server.scopes - - -@pytest.mark.asyncio -async def test_oauth_protected_resource_respects_x_forwarded_proto( - trusted_proxy_origin_headers, -): - """Test that oauth_protected_resource_mcp uses X-Forwarded-Proto for URLs""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - oauth_protected_resource_mcp, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server - oauth2_server = MCPServer( - server_id="test_oauth_server", - name="test_oauth", - server_name="test_oauth", - alias="test_oauth", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="test_client_id", - client_secret="test_client_secret", - authorization_url="https://provider.com/oauth/authorize", - token_url="https://provider.com/oauth/token", - scopes=["read", "write"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request with http base_url but X-Forwarded-Proto: https - mock_request = MagicMock(spec=Request) - mock_request.base_url = "http://litellm.example.com/" # HTTP - mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy - set_request_from_trusted_proxy(mock_request) - - # Call the endpoint - response = await oauth_protected_resource_mcp( - request=mock_request, - mcp_server_name="test_oauth", - ) - - # Verify response uses HTTPS URLs - assert response["authorization_servers"][0].startswith( - "https://litellm.example.com/" - ) - assert response["scopes_supported"] == oauth2_server.scopes - - -@pytest.mark.asyncio -async def test_oauth_authorization_server_respects_x_forwarded_proto( - trusted_proxy_origin_headers, -): - """Test that oauth_authorization_server_mcp uses X-Forwarded-Proto for URLs""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - oauth_authorization_server_mcp, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server - oauth2_server = MCPServer( - server_id="test_oauth_server", - name="test_oauth", - server_name="test_oauth", - alias="test_oauth", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="test_client_id", - client_secret="test_client_secret", - authorization_url="https://provider.com/oauth/authorize", - token_url="https://provider.com/oauth/token", - scopes=["read", "write"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request with http base_url but X-Forwarded-Proto: https - mock_request = MagicMock(spec=Request) - mock_request.base_url = "http://litellm.example.com/" # HTTP - mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy - set_request_from_trusted_proxy(mock_request) - - # Call the endpoint - response = await oauth_authorization_server_mcp( - request=mock_request, - mcp_server_name="test_oauth", - ) - - # Verify response uses HTTPS URLs - assert response["authorization_endpoint"].startswith("https://litellm.example.com/") - assert response["token_endpoint"].startswith("https://litellm.example.com/") - assert response["registration_endpoint"].startswith("https://litellm.example.com/") - assert response["grant_types_supported"] == ["authorization_code", "refresh_token"] - assert response["scopes_supported"] == oauth2_server.scopes - - -@pytest.mark.asyncio -async def test_register_client_respects_x_forwarded_proto( - trusted_proxy_origin_headers, -): - """Test that register_client uses X-Forwarded-Proto for redirect_uris""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - register_client, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - global_mcp_server_manager.registry.clear() - - # Mock request with http base_url but X-Forwarded-Proto: https - mock_request = MagicMock(spec=Request) - mock_request.base_url = "http://proxy.litellm.example/" # HTTP - mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy - set_request_from_trusted_proxy(mock_request) - - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", - new=AsyncMock(return_value={}), - ): - result = await register_client(request=mock_request) - - # Verify the redirect_uris use HTTPS - assert result == { - "client_id": "dummy_client", - "client_secret": "dummy", - "redirect_uris": ["https://proxy.litellm.example/callback"], - } - - -@pytest.mark.asyncio -async def test_authorize_endpoint_respects_x_forwarded_host( - trusted_proxy_origin_headers, -): - """Test that authorize endpoint uses X-Forwarded-Host and X-Forwarded-Proto to construct correct redirect_uri""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - authorize, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server - oauth2_server = MCPServer( - server_id="test_oauth_server", - name="test_oauth", - server_name="test_oauth", - alias="test_oauth", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="test_client_id", - client_secret="test_client_secret", - authorization_url="https://provider.com/oauth/authorize", - token_url="https://provider.com/oauth/token", - scopes=["read", "write"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request simulating nginx proxy: - # Internal: http://localhost:8888/github/mcp - # External: https://proxy.example.com/github/mcp - mock_request = MagicMock(spec=Request) - mock_request.base_url = "http://localhost:8888/github/mcp" - mock_request.headers = { - "X-Forwarded-Proto": "https", - "X-Forwarded-Host": "proxy.example.com", - } - set_request_from_trusted_proxy(mock_request) - - # Mock the encryption functions - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" - ) as mock_encrypt: - mock_encrypt.return_value = "mocked_encrypted_state" - - # Call authorize endpoint - response = await authorize( - request=mock_request, - client_id="test_client_id", - mcp_server_name="test_oauth", - redirect_uri="http://127.0.0.1:60108/callback", - state="test_state", - ) - - # Verify redirect URL uses the forwarded host and scheme - location = response.headers["location"] - - # The redirect_uri parameter should use the external URL - assert ( - "redirect_uri=https%3A%2F%2Fproxy.example.com%2Fgithub%2Fmcp%2Fcallback" - in location - or "redirect_uri=https://proxy.example.com/github/mcp/callback" in location - ) - - -@pytest.mark.asyncio -async def test_token_endpoint_respects_x_forwarded_host( - trusted_proxy_origin_headers, -): - """Test that token endpoint uses X-Forwarded-Host and X-Forwarded-Proto for redirect_uri""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - token_endpoint, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server - oauth2_server = MCPServer( - server_id="google_mcp", - name="google_mcp", - server_name="google_mcp", - alias="google_mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="test_client_id", - client_secret="test_secret", - authorization_url="https://accounts.google.com/o/oauth2/v2/auth", - token_url="https://oauth2.googleapis.com/token", - scopes=["openid", "email"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request simulating nginx proxy without port in host - mock_request = MagicMock(spec=Request) - mock_request.base_url = "http://localhost:8888/github/mcp" - mock_request.headers = { - "X-Forwarded-Proto": "https", - "X-Forwarded-Host": "proxy.example.com", - } - set_request_from_trusted_proxy(mock_request) - - # Mock httpx client response - mock_response = MagicMock() - mock_response.json.return_value = { - "access_token": "test_token", - "token_type": "Bearer", - "expires_in": 3599, - } - mock_response.raise_for_status = MagicMock() - - # Mock the async httpx client - mock_async_client = MagicMock() - mock_async_client.post = AsyncMock(return_value=mock_response) - - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client" - ) as mock_get_client: - mock_get_client.return_value = mock_async_client - - # Call token endpoint - await token_endpoint( - request=mock_request, - grant_type="authorization_code", - code="test_code", - redirect_uri="http://localhost:60108/callback", - client_id="test_client_id", - mcp_server_name="google_mcp", - client_secret="test_secret", - ) - - # Verify that the redirect_uri sent to the provider uses the external URL - call_args = mock_async_client.post.call_args - assert ( - call_args[1]["data"]["redirect_uri"] - == "https://proxy.example.com/github/mcp/callback" - ) - - -@pytest.mark.parametrize( - "base_url,x_forwarded_proto,x_forwarded_host,x_forwarded_port,expected_url", - [ - # Case 1: No forwarded headers - use original URL as-is (no trailing slash) - ( - "http://localhost:4000/", - None, - None, - None, - "http://localhost:4000", - ), - # Case 2: Only X-Forwarded-Proto - change scheme only - ( - "http://localhost:4000/", - "https", - None, - None, - "https://localhost:4000", - ), - # Case 3: X-Forwarded-Proto + X-Forwarded-Host - change scheme and host - ( - "http://localhost:4000/", - "https", - "proxy.example.com", - None, - "https://proxy.example.com", - ), - # Case 4: X-Forwarded-Host with port included in host header - ( - "http://localhost:4000/", - "https", - "proxy.example.com:8080", - None, - "https://proxy.example.com:8080", - ), - # Case 5: X-Forwarded-Host + X-Forwarded-Port as separate headers - ( - "http://localhost:4000/", - "https", - "proxy.example.com", - "8443", - "https://proxy.example.com:8443", - ), - # Case 6: Only X-Forwarded-Host without proto - use original scheme - ( - "http://localhost:4000/", - None, - "proxy.example.com", - None, - "http://proxy.example.com", - ), - # Case 7: Only X-Forwarded-Port without host - preserves original port if present - # (This is safer behavior - X-Forwarded-Port alone is unusual) - ( - "http://localhost:4000/", - None, - None, - "8443", - "http://localhost:4000", # Original port preserved when already present - ), - # Case 8: Complex internal URL with path (path is preserved) - ( - "http://localhost:8888/github/mcp", - "https", - "proxy.example.com", - None, - "https://proxy.example.com/github/mcp", - ), - # Case 9: IPv6 address in X-Forwarded-Host (should not treat :: as port separator) - ( - "http://localhost:4000/", - "https", - "[2001:db8::1]", - None, - "https://[2001:db8::1]", - ), - # Case 10: IPv6 address with port - ( - "http://localhost:4000/", - "https", - "[2001:db8::1]:8080", - None, - "https://[2001:db8::1]:8080", - ), - # Case 11: X-Forwarded-Host already has port, X-Forwarded-Port also provided (host wins) - ( - "http://localhost:4000/", - "https", - "proxy.example.com:9000", - "8443", - "https://proxy.example.com:9000", - ), - # Case 12: Standard proxy setup (most common case) - ( - "http://127.0.0.1:8888/", - "https", - "chatproxy.company.com", - None, - "https://chatproxy.company.com", - ), - # Case 13: Internal URL already has port, X-Forwarded-Port does NOT override - # (safer behavior - preserves original port when X-Forwarded-Host not provided) - ( - "http://localhost:4000/", - None, - None, - "443", - "http://localhost:4000", # Original port preserved - ), - # Case 14: Original URL with existing port in netloc, X-Forwarded-Host replaces it - ( - "http://internal.local:8888/", - "https", - "external.com", - None, - "https://external.com", - ), - ], -) -def test_get_request_base_url_comprehensive( - base_url, - x_forwarded_proto, - x_forwarded_host, - x_forwarded_port, - expected_url, - trusted_proxy_origin_headers, -): - """Comprehensive test for get_request_base_url with various header combinations""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - get_request_base_url, - ) - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - # Create mock request - mock_request = MagicMock(spec=Request) - mock_request.base_url = base_url - set_request_from_trusted_proxy(mock_request) - - # Build headers dict - headers = {} - if x_forwarded_proto: - headers["X-Forwarded-Proto"] = x_forwarded_proto - if x_forwarded_host: - headers["X-Forwarded-Host"] = x_forwarded_host - if x_forwarded_port: - headers["X-Forwarded-Port"] = x_forwarded_port - - # Mock headers.get() to return our test values - def mock_get(header_name, default=None): - return headers.get(header_name, default) - - mock_request.headers.get = mock_get - - # Test the function - result = get_request_base_url(mock_request) - - # Verify result - assert result == expected_url, ( - f"Expected '{expected_url}' but got '{result}'\n" - f"Input: base_url={base_url}, " - f"X-Forwarded-Proto={x_forwarded_proto}, " - f"X-Forwarded-Host={x_forwarded_host}, " - f"X-Forwarded-Port={x_forwarded_port}" - ) - - -def test_get_request_base_url_ignores_forwarded_headers_from_untrusted_client(): - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - get_request_base_url, - ) - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://gateway.example.com/mcp" - mock_request.headers = { - "X-Forwarded-Proto": "https", - "X-Forwarded-Host": "attacker.example.com", - "X-Forwarded-Port": "443", - } - mock_request.client = MagicMock() - mock_request.client.host = "203.0.113.10" - - with patch( - "litellm.proxy.proxy_server.general_settings", - { - "use_x_forwarded_for": True, - "mcp_trusted_proxy_ranges": TRUSTED_PROXY_RANGES, - }, - create=True, - ): - assert get_request_base_url(mock_request) == "https://gateway.example.com/mcp" - - -def test_validate_trusted_redirect_uri_rejects_spoofed_forwarded_host(): - try: - from litellm.proxy._experimental.mcp_server.oauth_utils import ( - validate_trusted_redirect_uri, - ) - from fastapi import Request - except ImportError: - pytest.skip("MCP OAuth utilities not available") - - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://gateway.example.com/" - mock_request.headers = { - "X-Forwarded-Proto": "https", - "X-Forwarded-Host": "attacker.example.com", - } - mock_request.client = MagicMock() - mock_request.client.host = "203.0.113.10" - - with ( - patch( - "litellm.proxy.proxy_server.general_settings", - { - "use_x_forwarded_for": True, - "mcp_trusted_proxy_ranges": TRUSTED_PROXY_RANGES, - }, - create=True, - ), - pytest.raises(HTTPException), - ): - validate_trusted_redirect_uri( - mock_request, - "https://attacker.example.com/callback", - ) - - -def test_validate_trusted_redirect_uri_allows_forwarded_origin_from_trusted_proxy( - trusted_proxy_origin_headers, -): - try: - from litellm.proxy._experimental.mcp_server.oauth_utils import ( - validate_trusted_redirect_uri, - ) - from fastapi import Request - except ImportError: - pytest.skip("MCP OAuth utilities not available") - - mock_request = MagicMock(spec=Request) - mock_request.base_url = "http://localhost:4000/" - mock_request.headers = { - "X-Forwarded-Proto": "https", - "X-Forwarded-Host": "proxy.example.com", - } - set_request_from_trusted_proxy(mock_request) - - validate_trusted_redirect_uri( - mock_request, - "https://proxy.example.com/callback", - ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index b4d3782ba43..442bfe8a090 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -2957,6 +2957,55 @@ def test_validate_trusted_redirect_uri_logs_diagnostic_on_rejection(caplog, monk assert "X-Forwarded-Host" in msg +@pytest.mark.parametrize( + "direct_ip,expect_accepted", + [ + ("10.0.0.7", True), + ("203.0.113.5", False), + ], +) +def test_validate_trusted_redirect_uri_follows_the_xff_trust_gate(direct_ip, expect_accepted, monkeypatch): + try: + from fastapi import HTTPException, Request + + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + validate_trusted_redirect_uri, + ) + except ImportError: + pytest.skip("MCP oauth_utils not available") + + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + monkeypatch.delenv("MCP_TRUSTED_REDIRECT_ORIGINS", raising=False) + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://localhost:4000/" + mock_request.client = MagicMock() + mock_request.client.host = direct_ip + + headers = { + "X-Forwarded-Proto": "https", + "X-Forwarded-Host": "proxy.example.com", + } + mock_request.headers.get = lambda name, default=None: headers.get(name, default) + mock_request.headers.__contains__ = lambda self_, name: name in headers + + redirect_uri = "https://proxy.example.com/callback" + general_settings = { + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": ["10.0.0.0/8"], + } + + with patch("litellm.proxy.proxy_server.general_settings", general_settings, create=True): + if expect_accepted: + validate_trusted_redirect_uri(mock_request, redirect_uri) + return + with pytest.raises(HTTPException) as exc_info: + validate_trusted_redirect_uri(mock_request, redirect_uri) + + assert exc_info.value.status_code == 400 + assert "proxy.example.com" in str(exc_info.value.detail) + + @pytest.mark.parametrize( "bad_value", [ From f9f8320972f6589dd5aac0877a1bc89c4f200028 Mon Sep 17 00:00:00 2001 From: Rithvik Mysore Suresh Date: Fri, 21 Aug 2026 11:47:04 -0400 Subject: [PATCH 049/166] fix(files): list unscoped managed files Read owner-scoped managed rows directly when no provider or model is supplied, avoiding an unauthenticated OpenAI fallback. Refs #35362 --- .../proxy/hooks/managed_files.py | 19 ++++-- litellm/llms/base_llm/files/transformation.py | 6 +- .../openai_files_endpoints/files_endpoints.py | 36 ++++++----- .../proxy/test_managed_files_hook.py | 33 +++++++++++ .../test_files_endpoint.py | 59 +++++++++++++++++++ 5 files changed, 130 insertions(+), 23 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index c986e835e4f..9b62284072d 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -1365,12 +1365,23 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): async def afile_list( self, - purpose: Optional[OpenAIFilesPurpose], + purpose: str | None, litellm_parent_otel_span: Optional[Span], + user_api_key_dict: UserAPIKeyAuth, **data: Dict, - ) -> List[OpenAIFileObject]: - """Handled in files_endpoints.py""" - return [] + ) -> Dict[str, object]: + owner_filter: Final = build_owner_filter(user_api_key_dict) + if owner_filter is None: + return build_list_page([]) + + rows: Final = await _managed_file_table(self.prisma_client).find_many(where=owner_filter) + files: Final = [ + parsed_file_object.model_copy(update={"id": row.unified_file_id}) + for row in rows + if (parsed_file_object := _parse_managed_file_object(row.file_object, row.unified_file_id)) is not None + and (purpose is None or parsed_file_object.purpose == purpose) + ] + return build_list_page(files) def _is_batch_polling_enabled(self) -> bool: """ diff --git a/litellm/llms/base_llm/files/transformation.py b/litellm/llms/base_llm/files/transformation.py index 174be93448b..7c19326b627 100644 --- a/litellm/llms/base_llm/files/transformation.py +++ b/litellm/llms/base_llm/files/transformation.py @@ -13,7 +13,6 @@ from litellm.types.llms.openai import ( FileContentRequest, OpenAICreateFileRequestOptionalParams, OpenAIFileObject, - OpenAIFilesPurpose, ) from litellm.types.utils import LlmProviders, ModelResponse @@ -240,10 +239,11 @@ class BaseFileEndpoints(ABC): @abstractmethod async def afile_list( self, - purpose: OpenAIFilesPurpose | None, + purpose: str | None, litellm_parent_otel_span: Span | None, + user_api_key_dict: UserAPIKeyAuth, **data: dict, - ) -> list[OpenAIFileObject]: + ) -> dict[str, object]: pass @abstractmethod diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 37cfd9d073d..a482cc54748 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -1488,24 +1488,28 @@ async def list_files( or get_custom_llm_provider_from_request_headers(request=request) or get_custom_llm_provider_from_request_query(request=request) or await get_custom_llm_provider_from_request_body(request=request) - or "openai" ) + managed_files_obj: Final = proxy_logging_obj.get_proxy_hook("managed_files") + if custom_llm_provider is None and isinstance(managed_files_obj, BaseFileEndpoints): + response = await managed_files_obj.afile_list( + purpose=purpose, + litellm_parent_otel_span=user_api_key_dict.parent_otel_span, + user_api_key_dict=user_api_key_dict, + ) + else: + resolved_custom_llm_provider: Final = custom_llm_provider or "openai" + apply_team_provider_credentials( + data=data, + llm_router=llm_router, + user_api_key_dict=user_api_key_dict, + custom_llm_provider=resolved_custom_llm_provider, + ) - # No model/target_model_names pinned: resolve upstream credentials from - # the team's deployment for this provider so the call is authenticated - # against the team's own account (e.g. the team's openai deployment). - apply_team_provider_credentials( - data=data, - llm_router=llm_router, - user_api_key_dict=user_api_key_dict, - custom_llm_provider=custom_llm_provider, - ) - - response = await litellm.afile_list( - custom_llm_provider=custom_llm_provider, - purpose=purpose, - **data, - ) + response = await litellm.afile_list( + custom_llm_provider=resolved_custom_llm_provider, + purpose=purpose, + **data, + ) if response is None: raise HTTPException( diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py index fcd03e77aa2..b39d2ef8559 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -190,6 +190,39 @@ async def test_get_user_created_file_ids_remaps_stored_raw_provider_id_to_unifie assert files[0].purpose == raw_provider_object.purpose +@pytest.mark.asyncio +async def test_afile_list_returns_owner_scoped_managed_files(): + managed_files = _make_managed_files_instance() + managed_files.prisma_client.db.litellm_managedfiletable.find_many = AsyncMock( + return_value=[ + MagicMock( + file_object=_make_file_object("file-provider-id").model_dump(), + unified_file_id="unified-file-id", + ), + MagicMock( + file_object=_make_file_object("file-other-purpose").model_copy( + update={"purpose": "batch"} + ).model_dump(), + unified_file_id="unified-other-purpose", + ), + ] + ) + + response = await managed_files.afile_list( + purpose="batch_output", + litellm_parent_otel_span=None, + user_api_key_dict=_make_user_api_key_dict(), + ) + + managed_files.prisma_client.db.litellm_managedfiletable.find_many.assert_awaited_once_with( + where={"created_by": "test-user"} + ) + assert [file.id for file in response["data"]] == ["unified-file-id"] + assert response["first_id"] == "unified-file-id" + assert response["last_id"] == "unified-file-id" + assert response["has_more"] is False + + @pytest.mark.asyncio async def test_parse_managed_file_object_warning_omits_rejected_values(caplog): from litellm_enterprise.proxy.hooks.managed_files import ( diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index bf9323cdc6a..e6101d3edd8 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -2468,6 +2468,65 @@ def test_list_files_without_target_model_names_uses_team_openai_deployment( proxy_logging_obj.post_call_failure_hook.assert_not_called() +def test_unscoped_list_files_uses_managed_file_store( + mocker: MockerFixture, monkeypatch, llm_router: Router +): + import litellm.proxy.proxy_server as ps + from litellm.llms.base_llm.files.transformation import BaseFileEndpoints + from litellm.proxy._types import LitellmUserRoles + + managed_file = OpenAIFileObject( + id="unified-file-id", + object="file", + bytes=100, + created_at=1700000000, + filename="output.jsonl", + purpose="batch_output", + status="processed", + ) + + proxy_logging_obj = setup_proxy_logging_object(monkeypatch, llm_router) + managed_files = mocker.MagicMock(spec=BaseFileEndpoints) + managed_files.afile_list = mocker.AsyncMock( + return_value={ + "object": "list", + "data": [managed_file], + "first_id": managed_file.id, + "last_id": managed_file.id, + "has_more": False, + } + ) + proxy_logging_obj.proxy_hook_mapping["managed_files"] = managed_files + proxy_logging_obj.update_request_status = mocker.AsyncMock() + proxy_logging_obj.post_call_success_hook = mocker.AsyncMock(return_value=None) + proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + provider_list = mocker.patch.object(litellm, "afile_list", new=mocker.AsyncMock()) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="test-key", + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="test-user", + ) + + try: + response = client.get( + "/v1/files", + headers={"Authorization": "Bearer test-key"}, + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 200, response.text + assert response.json()["data"][0]["id"] == "unified-file-id" + managed_files.afile_list.assert_awaited_once() + assert managed_files.afile_list.await_args.kwargs["user_api_key_dict"].user_id == "test-user" + provider_list.assert_not_awaited() + proxy_logging_obj.post_call_failure_hook.assert_not_called() + + def test_list_files_restricted_team_does_not_leak_global_openai_credentials( mocker: MockerFixture, monkeypatch ): From 7da34e8aed341b3368b14810d91052ebccb97117 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 21 Aug 2026 09:47:52 -0700 Subject: [PATCH 050/166] fix(proxy): make per-model budgets track spend, enforce, and report the same counter (#37736) Per-model budgets were three separate things pretending to be one. The enforcement check, the post-call increment and the info endpoints each derived their own cache key, so a budget could refuse traffic at 429 while /key/info reported zero usage, and a Bedrock model id never matched a budget keyed on the bare family name. /user/new echoed a model_max_budget back and stored an empty dict, and nothing enforced a user-scoped per-model budget at all. One owner now builds the counter key from the configured budget model, and enforcement, the increment and the info endpoints all read it. Bedrock ids resolve through the model-cost map. Auth carries the user's budget onto the token on every branch that reaches the spend hook, including JWT and auto-registration. Native passthrough attaches the three budget metadata keys its StandardLoggingUserAPIKeyMetadata does not carry, so /anthropic/... and /bedrock/... traffic is counted and capped like /v1/chat/completions. The dashboard gains the per-model budget editor it never had, on the key create, key edit and internal-user edit forms. It is read-only without an enterprise license, matching the write gate the proxy already enforces, and an untouched budget is left out of an update so an unrelated edit cannot trip that gate. The editor hydrates from either BudgetConfig spelling, since model_max_budget is a plain dict that the proxy stores exactly as the client sent it, and it carries through the fields it does not model. Without both, editing one model would drop another model row entirely and silently discard its tpm_limit and rpm_limit. /user/info refreshes its local copy of the user field by field after a save, so model_max_budget joins that list. Left out, a saved cap read back as the old one when the form was reopened, and clearing the row to recover would then wipe the value that had actually persisted. A zero-dollar cap is the strictest limit expressible, not the absence of one, so it is enforced rather than skipped on falsiness, spend exactly at the cap is refused the way every sibling budget check already refuses it, and a counter that was never written reads as zero spend rather than as unknown. The usage endpoints read every counter in one batched lookup, so a large model_max_budget cannot fan out into one concurrent cache call per configured model. Every auth path honours the same zero-cost skip flag, so none of them can refuse a free request that another serves. The custom-auth helper gains the flag it never had, which also changes its pre-existing key and end-user checks. The compaction summary gate checks the user scope alongside the key and end-user ones. This file propagates all three budgets into the summary subrequest, so enforcing only two let compaction increment a counter it could not be refused by. Custom auth attaches the user's budget to the token unconditionally, since the post-call spend hook reads it there: gating the attach on the same condition as enforcement left the counter uncharged whenever the request was not itself enforceable. An entry that will not validate is skipped rather than raised on, so one malformed scope cannot abort every other scope's increment or turn a config typo into a 500. The edit forms re-seed the budget editor when a different key or user is loaded. Its rows are seeded once and cannot re-read their own value prop, so without this a save wrote the previously loaded record's budgets onto the current one. Only the built-in provider pass-through routes carry the budget metadata. get_model_from_request deliberately resolves no model for a user-defined pass-through, since its body is forwarded verbatim and names an upstream model, so attaching there would charge a counter nothing on that route can refuse. --- .../context_management/editors/compact.py | 32 +- litellm/proxy/_types.py | 6 + litellm/proxy/auth/auth_utils.py | 4 +- litellm/proxy/auth/user_api_key_auth.py | 140 ++- .../proxy/hooks/model_max_budget_limiter.py | 584 +++++---- litellm/proxy/litellm_pre_call_utils.py | 2 + .../internal_user_endpoints.py | 21 +- .../key_management_endpoints.py | 84 +- .../pass_through_endpoints.py | 17 + ...test_unit_test_max_model_budget_limiter.py | 1054 ++++++++++++++--- .../test_user_api_key_auth.py | 670 ++++++++++- .../context_management/test_compact.py | 76 ++ .../test_internal_user_endpoints.py | 75 ++ .../test_key_management_endpoints.py | 104 +- .../test_pass_through_endpoints.py | 840 +++++-------- .../users/_components/BulkEditUsers.tsx | 3 + .../users/_components/user_edit_view.test.tsx | 121 +- .../users/_components/user_edit_view.tsx | 25 + .../user_info_view.integration.test.tsx | 44 +- .../_components/view_users/user_info_view.tsx | 8 + .../ModelMaxBudgetEditor.integration.test.tsx | 69 ++ .../ModelMaxBudgetEditor.test.ts | 140 +++ .../key_team_helpers/ModelMaxBudgetEditor.tsx | 233 ++++ .../components/key_team_helpers/key_list.tsx | 4 +- .../modelMaxBudgetPayload.test.ts | 71 ++ .../key_team_helpers/modelMaxBudgetPayload.ts | 44 + .../useModelMaxBudgetField.ts | 34 + .../key_team_helpers/useSeededState.ts | 26 + .../src/components/networking.tsx | 3 + .../organisms/createKeyPayload.test.ts | 25 + .../components/organisms/createKeyPayload.ts | 3 + .../organisms/create_key_button.tsx | 19 + .../templates/key_edit_view.test.tsx | 89 ++ .../components/templates/key_edit_view.tsx | 15 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 12 + 35 files changed, 3656 insertions(+), 1041 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.integration.test.tsx create mode 100644 ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.test.ts create mode 100644 ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx create mode 100644 ui/litellm-dashboard/src/components/key_team_helpers/modelMaxBudgetPayload.test.ts create mode 100644 ui/litellm-dashboard/src/components/key_team_helpers/modelMaxBudgetPayload.ts create mode 100644 ui/litellm-dashboard/src/components/key_team_helpers/useModelMaxBudgetField.ts create mode 100644 ui/litellm-dashboard/src/components/key_team_helpers/useSeededState.ts diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py index a7c462a8fb0..2a87afb5990 100644 --- a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py @@ -56,7 +56,7 @@ from ..result import PolyfillResult # so the summary's spend is attributed to the same scopes. The list mirrors the # fields populated by # ``LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata``. -# ``user_api_key_model_max_budget`` / ``user_api_key_end_user_model_max_budget`` +# The three ``*_model_max_budget`` fields # are what ``_PROXY_VirtualKeyModelMaxBudgetLimiter`` reads post-call to update # the per-model spend caches, so without them the summary spend would never # count against the caller's model budget. ``user_api_key_end_user_id`` / @@ -76,6 +76,7 @@ _PROPAGATED_METADATA_KEYS: Final = ( "user_api_key_end_user_id", "user_api_end_user_max_budget", "user_api_key_model_max_budget", + "user_api_key_user_model_max_budget", "user_api_key_end_user_model_max_budget", "litellm_call_id", "litellm_parent_otel_span", @@ -317,10 +318,14 @@ async def _check_summary_model_budget( The summary subrequest never passes back through ``user_api_key_auth``, so without this gate a caller whose ``model_max_budget`` for ``context_management_summary_model`` is exhausted could keep consuming that - model via compaction. Mirrors the ``model_max_budget`` / - ``end_user_model_max_budget`` enforcement that ``user_api_key_auth`` runs for - the client-requested model. Returns True outside the proxy or when no + model via compaction. Mirrors the per-model budget enforcement that + ``user_api_key_auth`` runs for the client-requested model. Returns True outside the proxy or when no per-model budget is configured. + + All three scopes are checked because the summary's spend is charged to all + three: this file propagates the key, user and end-user budgets into the + subrequest's metadata, so enforcing only two of them would let compaction + increment a counter it can never be refused by. """ if user_api_key_auth is None: return True @@ -347,6 +352,25 @@ async def _check_summary_model_budget( ) return False + user_model_max_budget: Final = getattr(user_api_key_auth, "user_model_max_budget", None) + user_id: Final = getattr(user_api_key_auth, "user_id", None) + if isinstance(user_model_max_budget, dict) and user_model_max_budget and user_id is not None: + try: + await model_max_budget_limiter.is_user_within_model_budget( + user_id=user_id, + user_model_max_budget=user_model_max_budget, + model=summary_model, + ) + except litellm.BudgetExceededError: + return False + except Exception as e: # noqa: BLE001 # a budget gate denies on any failure, as the key and end-user scopes do + verbose_logger.warning( + "compact_20260112: unexpected error during user model-budget check for summary_model=%s; denying: %s", + summary_model, + e, + ) + return False + end_user_model_max_budget: Final = getattr(user_api_key_auth, "end_user_model_max_budget", None) end_user_id: Final = getattr(user_api_key_auth, "end_user_id", None) if isinstance(end_user_model_max_budget, dict) and end_user_model_max_budget and end_user_id is not None: diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 00cbd13cfdc..e51a3138d2d 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2805,6 +2805,10 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob user_email: str | None = None user_spend: float | None = None user_max_budget: float | None = None + # Values stay `object` rather than BudgetConfig: this is the raw JSON column, + # and validating it here would make one malformed row fail auth outright. + # resolve_model_budget validates the single entry a request actually needs. + user_model_max_budget: dict[str, object] | None = None request_route: str | None = None is_session_token: bool = False # Server-only marker set exclusively by the MCP gateway admission path @@ -2982,6 +2986,8 @@ class UserInfoV2Response(LiteLLMPydanticObjectBase): sso_user_id: str | None = None teams: list[str] = [] # Just team IDs, not full team objects object_permission: LiteLLM_ObjectPermissionTable | None = None + model_max_budget: dict | None = None + model_max_budget_usage: dict | None = None from litellm.models.config import LiteLLM_Config as LiteLLM_Config # noqa: E402 diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index ce662ee0374..d04a71535ef 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -1801,7 +1801,7 @@ def _format_model_candidates( return candidates -def _request_dispatched_to_pass_through_endpoint(request: Request | None) -> bool: +def request_dispatched_to_pass_through_endpoint(request: Request | None) -> bool: """Whether FastAPI resolved this request to a user-defined pass-through handler. Reads the marker set by ``create_pass_through_route`` off the dispatched endpoint @@ -1842,7 +1842,7 @@ def get_model_from_request( and does not carry the marker. Built-in provider passthrough routes (``/vertex_ai``, ``/gemini``, ...) are separate handlers and keep model enforcement. """ - if _request_dispatched_to_pass_through_endpoint(request): + if request_dispatched_to_pass_through_endpoint(request): return None candidates: Final = _extract_model_candidates_from_request( diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 99592d44f9b..fe4f1ee4ae5 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -11,6 +11,7 @@ import asyncio import fnmatch import re import secrets +from collections.abc import Mapping from datetime import datetime, timezone from typing import Any, Final, NamedTuple, Protocol, Union, cast @@ -186,6 +187,62 @@ class _KeyModelBudgetLimiter(Protocol): async def get_fallback_model_within_budget(self, user_api_key_dict: UserAPIKeyAuth, model: str) -> str | None: ... +class _UserModelBudgetLimiter(Protocol): + async def is_user_within_model_budget( + self, user_id: str, user_model_max_budget: Mapping[str, object], model: str + ) -> bool: ... + + +async def _read_user_model_max_budget( + user_id: str | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: object, + proxy_logging_obj: ProxyLogging, +) -> dict | None: + """The user row's `model_max_budget`, or None when the row cannot be read. + + A user whose row is missing must not be refused: this is a budget lookup, + and the main auth path likewise treats an unreadable user as no user. + """ + if user_id is None or prisma_client is None: + return None + try: + user_obj: Final = await get_user_object( + user_id=user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + parent_otel_span=parent_otel_span, # pyright: ignore[reportArgumentType] # Span is a runtime union, not usable in an annotation here + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: # noqa: BLE001 # mirrors the main path's tolerance + verbose_logger.debug("Unable to read user for the per-model budget check: %s", e) + return None + return getattr(user_obj, "model_max_budget", None) + + +async def _check_user_model_budget( + valid_token: UserAPIKeyAuth, + model_max_budget_limiter: _UserModelBudgetLimiter, + models: list[str], +) -> None: + """Enforce the internal user's own `model_max_budget` across the request's models. + + Separate from the key check: a user's per-model budget caps every key they + own, so a caller cannot escape it by minting another key. + """ + user_model_max_budget: Final = valid_token.user_model_max_budget + if valid_token.user_id is None or not isinstance(user_model_max_budget, Mapping) or not user_model_max_budget: + return + for model_name in models: + await model_max_budget_limiter.is_user_within_model_budget( + user_id=valid_token.user_id, + user_model_max_budget=user_model_max_budget, + model=model_name, + ) + + async def _check_key_model_budget_with_fallback( valid_token: UserAPIKeyAuth, model_max_budget_limiter: _KeyModelBudgetLimiter, @@ -1390,6 +1447,7 @@ async def _user_api_key_auth_builder( end_user_id=end_user_id, user_tpm_limit=(user_object.tpm_limit if user_object is not None else None), user_rpm_limit=(user_object.rpm_limit if user_object is not None else None), + user_model_max_budget=(user_object.model_max_budget if user_object is not None else None), team_member_rpm_limit=( team_membership.safe_get_team_member_rpm_limit() if team_membership is not None else None ), @@ -1427,6 +1485,13 @@ async def _user_api_key_auth_builder( if auto_registered is not None: auto_registered.jwt_claims = jwt_claims auto_registered.user_email = user_email + # The auto-registered token is built from the new key's + # columns, which carry no user budget. Carry over the + # already-loaded user row rather than re-reading it, or + # the budget check below has nothing to enforce. + auto_registered.user_model_max_budget = ( + user_object.model_max_budget if user_object is not None else None + ) valid_token = auto_registered api_key = valid_token.token or "" @@ -1458,6 +1523,28 @@ async def _user_api_key_auth_builder( valid_token.project_metadata = _jwt_project_obj.metadata valid_token.project_alias = _jwt_project_obj.project_alias + # JWT auth returns here rather than falling through to the + # virtual-key checks below, so the user's per-model budget + # has to be enforced on this path too. Without it the + # post-call increment still charges the counter and nothing + # ever reads it, which is worse than not tracking at all. + # Guarded by the same flag the virtual-key path uses, or a + # zero-cost model would be refused here and allowed there, + # while the log above claims all budget checks were skipped. + if not skip_budget_checks: + await _check_user_model_budget( + valid_token=cast(UserAPIKeyAuth, valid_token), + model_max_budget_limiter=model_max_budget_limiter, + models=_get_model_names_for_budget_checks( + model=_get_model_from_request_context( + request_data=request_data, + route=route, + request=request, + llm_router=llm_router, + ) + ), + ) + return cast(UserAPIKeyAuth, valid_token) #### ELSE #### @@ -1811,6 +1898,12 @@ async def _user_api_key_auth_builder( ) user_obj = None + if user_obj is not None: + # The joint verification-token view carries the key's columns only, so the + # user's own per-model budget reaches enforcement and the post-call + # increment through the row fetched here. + valid_token.user_model_max_budget = user_obj.model_max_budget + if ( user_obj is not None and isinstance(user_obj.metadata, dict) @@ -1974,6 +2067,14 @@ async def _user_api_key_auth_builder( ) current_models = _get_model_names_for_budget_checks(model=current_model) + # Check 5a. Internal user model_max_budget + if current_models: + await _check_user_model_budget( + valid_token=valid_token, + model_max_budget_limiter=model_max_budget_limiter, + models=current_models, + ) + # Check 5b. End-user model max budget end_user_mmb: Final = valid_token.end_user_model_max_budget if ( @@ -2757,6 +2858,7 @@ async def _return_user_api_key_auth_obj( user_email=user_obj.user_email, user_spend=getattr(user_obj, "spend", None), user_max_budget=getattr(user_obj, "max_budget", None), + user_model_max_budget=getattr(user_obj, "model_max_budget", None), ) if user_obj is not None and _is_user_proxy_admin(user_obj=user_obj): user_api_key_kwargs.update( @@ -3020,10 +3122,21 @@ async def _run_post_custom_auth_checks( ) current_models = _get_model_names_for_budget_checks(model=current_model) + # A zero-cost model cannot move any counter, so refusing it means refusing on + # spend some other model accrued. The JWT and virtual-key paths already skip + # every budget check for these; this path did not, so the same request could + # be refused under custom auth and served under the other two. + skip_budget_checks: Final = ( + _is_model_cost_zero(model=current_model, llm_router=llm_router) + if current_model is not None and llm_router is not None + else False + ) + # 3. Check key-level model_max_budget max_budget_per_model: Final = valid_token.model_max_budget if ( - max_budget_per_model is not None + not skip_budget_checks + and max_budget_per_model is not None and isinstance(max_budget_per_model, dict) and len(max_budget_per_model) > 0 and current_models @@ -3050,10 +3163,33 @@ async def _run_post_custom_auth_checks( ) current_models = _get_model_names_for_budget_checks(model=current_model) + # 3b. Attach and check the internal user's model_max_budget. + # Custom auth builds its own token, so unlike the main path nothing has + # loaded the user row yet. The attach is unconditional because the post-call + # spend hook reads this field off the token: gating it on the same condition + # as enforcement would leave the user's counter uncharged whenever this + # request was not itself enforceable, which is the untracked-spend bug this + # PR exists to fix. + user_budget: Final = await _read_user_model_max_budget( + user_id=valid_token.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, + ) + valid_token.user_model_max_budget = user_budget # rebind-ok: the spend hook reads it off this token + if not skip_budget_checks and current_models: + await _check_user_model_budget( + valid_token=valid_token, + model_max_budget_limiter=model_max_budget_limiter, + models=current_models, + ) + # 4. Check end-user model_max_budget end_user_mmb: Final = valid_token.end_user_model_max_budget if ( - end_user_mmb is not None + not skip_budget_checks + and end_user_mmb is not None and isinstance(end_user_mmb, dict) and len(end_user_mmb) > 0 and current_models diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 215969ef899..c5d10b2749b 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -1,21 +1,253 @@ import json +import time +from collections.abc import Iterable, Mapping +from dataclasses import dataclass +from types import MappingProxyType from typing import Final import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import Span +from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.llms.bedrock.common_utils import get_bedrock_base_model from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import ( - BudgetConfig, - GenericBudgetConfigType, - StandardLoggingPayload, -) +from litellm.types.utils import BudgetConfig, StandardLoggingPayload VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX: Final = "virtual_key_spend" END_USER_SPEND_CACHE_KEY_PREFIX: Final = "end_user_model_spend" +USER_SPEND_CACHE_KEY_PREFIX: Final = "user_model_spend" + +_SPEND_CACHE_KEY_PREFIXES: Final = MappingProxyType( + { + Litellm_EntityType.KEY: VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX, + Litellm_EntityType.USER: USER_SPEND_CACHE_KEY_PREFIX, + Litellm_EntityType.END_USER: END_USER_SPEND_CACHE_KEY_PREFIX, + } +) + +_LEGACY_REQUEST_MODEL_SCOPES: Final = frozenset({Litellm_EntityType.KEY, Litellm_EntityType.END_USER}) + +_PROCESS_STARTED_AT: Final = time.monotonic() + +_BUDGET_START_TIME_KEY_PREFIXES: Final = MappingProxyType( + { + Litellm_EntityType.KEY: "virtual_key_budget_start_time", + Litellm_EntityType.USER: "user_model_budget_start_time", + Litellm_EntityType.END_USER: "end_user_budget_start_time", + } +) + + +@dataclass(frozen=True, slots=True) +class ResolvedModelBudget: + """The `model_max_budget` entry a request resolved to. + + ``budget_model`` is the key as the operator configured it, not the model + name on the request. Every counter is keyed on it so enforcement, the + post-call increment and the `/key/info` + `/user/info` usage reads cannot + disagree about which counter a request belongs to. + """ + + budget_model: str + budget_config: BudgetConfig + + +def model_budget_spend_cache_key( + entity_type: Litellm_EntityType, + entity_id: str | None, + budget_model: str, + budget_duration: str | None, +) -> str: + """Sole owner of the per-model spend counter key, shared by its writer and all of its readers.""" + return f"{_SPEND_CACHE_KEY_PREFIXES[entity_type]}:{entity_id}:{budget_model}:{budget_duration}" + + +def _legacy_request_model_spend_cache_key( + entity_type: Litellm_EntityType, + entity_id: str | None, + model: str, + resolved: ResolvedModelBudget, +) -> str | None: + """The counter this request was billed to before the budget model owned the key, or None. + + Upgrading proxies carry live counters keyed on the model as REQUESTED + (`openai/gpt-4`) rather than as configured (`gpt-4`), and those were the + counters the previous version enforced on. Nothing writes that spelling once + this version is running, so the pre-upgrade and post-upgrade counters hold + disjoint halves of one window and adding them is the window's real spend. + + Only the key and end-user scopes ever had one. The user scope is introduced + by this change, so it has no counter to carry. + + The carry stops one budget window after start-up, because a legacy counter + belongs to a window that was already open when this process replaced the one + writing it. Past that point the lookup could only ever miss. + """ + budget_duration: Final = resolved.budget_config.budget_duration + if entity_type not in _LEGACY_REQUEST_MODEL_SCOPES or budget_duration is None: + return None + if time.monotonic() - _PROCESS_STARTED_AT >= duration_in_seconds(budget_duration): + return None + return model_budget_spend_cache_key( + entity_type=entity_type, + entity_id=entity_id, + budget_model=model, + budget_duration=budget_duration, + ) + + +def model_budget_start_time_cache_key( + entity_type: Litellm_EntityType, + entity_id: str | None, + budget_model: str, + budget_duration: str | None, +) -> str: + """Window start for one (entity, budget model) pair. + + Scoped per budget model because an entity may budget two models over + different periods, and a shared start time lets the shorter period restart + the longer one's window. + """ + return f"{_BUDGET_START_TIME_KEY_PREFIXES[entity_type]}:{entity_id}:{budget_model}:{budget_duration}" + + +def resolve_model_budget(model: str, model_max_budget: Mapping[str, object]) -> ResolvedModelBudget | None: + """Find the `model_max_budget` entry that governs `model`, or None.""" + for candidate in _budget_model_candidates(model): + raw_budget_config = model_max_budget.get(candidate) + if raw_budget_config is None: + continue + if (budget_config := _usable_budget_config(raw_budget_config)) is None: + # An entry that will not validate cannot be keyed, so it cannot be + # enforced or incremented. Skip to the next candidate rather than + # raising: raising would abort every other scope's increment and turn + # a config typo into a 500, and stopping here would let one malformed + # specific entry disable a perfectly good bare-family budget beside + # it. The candidate chain already falls through an ABSENT entry, and + # an unparseable one is indistinguishable from absent to enforcement. + # `validate_model_max_budget` rejects these on the write path, so + # reaching here means config.yaml or a direct DB edit. + verbose_proxy_logger.warning( + "Ignoring unusable model_max_budget entry for %s; it cannot be enforced or tracked", + candidate, + ) + continue + return ResolvedModelBudget(budget_model=candidate, budget_config=budget_config) + return None + + +def _budget_model_candidates(model: str) -> tuple[str, ...]: + """Names a budget may be configured under for a request on `model`, most specific first. + + Beyond the model as sent, a budget may be keyed on the model without its + ``{custom_llm_provider}/`` prefix (``gpt-4o`` governs ``openai/gpt-4o``), on + the Bedrock base model (``anthropic.claude-opus-4-8`` governs the + cross-region ``us.anthropic.claude-opus-4-8``), or on the bare family name + that Bedrock id shares with its direct-provider twin (``claude-opus-4-8``). + """ + return tuple(dict.fromkeys((model, model.split("/")[-1], *_bedrock_candidates(model)))) + + +def _bedrock_candidates(model: str) -> tuple[str, ...]: + """Bedrock-only candidates, empty unless litellm prices `model` as a Bedrock model. + + Gating on the cost map rather than on a vendor allowlist is what makes + splitting the leading dotted segment safe: most dotted model ids are not + Bedrock ids at all (``azure/gpt-4.1``, ``gpt-image-1.5``), and splitting one + of those would produce a garbage candidate. + """ + base_model: Final = get_bedrock_base_model(model) + cost_entry: Final = litellm.model_cost.get(base_model) + if not isinstance(cost_entry, dict) or not str(cost_entry.get("litellm_provider", "")).startswith("bedrock"): + return () + _, _, without_vendor = base_model.partition(".") + return (base_model, without_vendor) if without_vendor else (base_model,) + + +async def build_model_max_budget_usage( + entity_type: Litellm_EntityType, + entity_id: str | None, + model_max_budget: Mapping[str, object] | None, + cache: DualCache | None, +) -> dict[str, dict[str, object]]: + """Current-window spend per configured budget model, as `/key/info` and `/user/info` report it. + + `cache` must be the DualCache the limiter writes the counters to; callers + read it off the limiter rather than re-deriving it, so a scope that is being + blocked can never report zero usage. + """ + if cache is None or entity_id is None or not model_max_budget: + return {} + + budgets: Final = tuple( + (budget_model, budget_config) + for budget_model, raw_budget_config in model_max_budget.items() + for budget_config in (_usable_budget_config(raw_budget_config),) + if budget_config is not None + ) + if not budgets: + return {} + spend_keys: Final = tuple( + model_budget_spend_cache_key( + entity_type=entity_type, + entity_id=entity_id, + budget_model=budget_model, + budget_duration=budget_config.budget_duration, + ) + for budget_model, budget_config in budgets + ) + batched: Final = await cache.async_batch_get_cache( + keys=list(spend_keys) # mutable-ok: async_batch_get_cache annotates keys as list, so one must exist here + ) + # async_batch_get_cache returns None if it fails internally, and its result is + # index-aligned with `keys` otherwise. An unusable result reads as a miss, + # which is what a never-written counter already reads as. + current_spends: Final = ( + tuple(batched) if isinstance(batched, list) and len(batched) == len(budgets) else (None,) * len(budgets) + ) + return { + budget_model: { + "current_spend": round(_as_spend(current_spend), 4), + "budget_limit": budget_config.max_budget, + "time_period": budget_config.budget_duration, + } + for (budget_model, budget_config), current_spend in zip(budgets, current_spends, strict=True) + } + + +def _usable_budget_config(raw_budget_config: object) -> BudgetConfig | None: + try: + budget_config: Final = BudgetConfig.model_validate(raw_budget_config) + if budget_config.budget_duration is None: + return None + duration_in_seconds(budget_config.budget_duration) + except Exception: # noqa: BLE001 # a malformed entry must not fail the whole report + return None + return budget_config + + +def _as_spend(current_spend: object) -> float: + try: + return float(current_spend or 0.0) # pyright: ignore[reportArgumentType] # non-numeric falls to the except + except (TypeError, ValueError): + return 0.0 + + +def _resolve_entity_model_budgets( + model: str, + entity_budgets: Iterable[tuple[Litellm_EntityType, str | None, object]], +) -> tuple[tuple[Litellm_EntityType, str, ResolvedModelBudget], ...]: + """Drop the scopes that do not budget `model`, keeping only what can be incremented.""" + return tuple( + (entity_type, entity_id, resolved) + for entity_type, entity_id, model_max_budget in entity_budgets + if entity_id is not None and isinstance(model_max_budget, Mapping) and model_max_budget + for resolved in (resolve_model_budget(model=model, model_max_budget=model_max_budget),) + if resolved is not None and resolved.budget_config.budget_duration is not None + ) class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): @@ -41,47 +273,17 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): Raises: BudgetExceededError: If the user_api_key_dict has exceeded the model budget """ - _model_max_budget: Final = user_api_key_dict.model_max_budget - internal_model_max_budget: Final[GenericBudgetConfigType] = {} - - for _model, _budget_info in _model_max_budget.items(): - internal_model_max_budget[_model] = BudgetConfig(**_budget_info) - - verbose_proxy_logger.debug( - "internal_model_max_budget %s", - json.dumps(internal_model_max_budget, indent=4, default=str), + return await self._is_entity_within_model_budget( + entity_type=Litellm_EntityType.KEY, + entity_id=user_api_key_dict.token, + model_max_budget=user_api_key_dict.model_max_budget, + model=model, + exceeded_message=( + f"LiteLLM Virtual Key: {user_api_key_dict.token}, key_alias: {user_api_key_dict.key_alias}, " + f"exceeded budget for model={model}" + ), ) - # check if current model is in internal_model_max_budget - _current_model_budget_info: Final = self._get_request_model_budget_config( - model=model, internal_model_max_budget=internal_model_max_budget - ) - if _current_model_budget_info is None: - verbose_proxy_logger.debug("Model %s not found in internal_model_max_budget", model) - return True - - # check if current model is within budget - if _current_model_budget_info.max_budget and _current_model_budget_info.max_budget > 0: - _current_spend: Final = await self._get_virtual_key_spend_for_model( - user_api_key_hash=user_api_key_dict.token, - model=model, - key_budget_config=_current_model_budget_info, - ) - if ( - _current_spend is not None - and _current_model_budget_info.max_budget is not None - and _current_spend > _current_model_budget_info.max_budget - ): - raise litellm.BudgetExceededError( - message=f"LiteLLM Virtual Key: {user_api_key_dict.token}, key_alias: {user_api_key_dict.key_alias}, exceeded budget for model={model}", - current_cost=_current_spend, - max_budget=_current_model_budget_info.max_budget, - entity_type=Litellm_EntityType.KEY.value, - entity_id=user_api_key_dict.token, - ) - - return True - async def get_fallback_model_within_budget( self, user_api_key_dict: UserAPIKeyAuth, @@ -96,10 +298,30 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): continue return None + async def is_user_within_model_budget( + self, + user_id: str, + user_model_max_budget: Mapping[str, object], + model: str, + ) -> bool: + """ + Check if the internal user is within the model budget + + Raises: + BudgetExceededError: If the user has exceeded the model budget + """ + return await self._is_entity_within_model_budget( + entity_type=Litellm_EntityType.USER, + entity_id=user_id, + model_max_budget=user_model_max_budget, + model=model, + exceeded_message=f"LiteLLM User: {user_id}, exceeded budget for model={model}", + ) + async def is_end_user_within_model_budget( self, end_user_id: str, - end_user_model_max_budget: dict, + end_user_model_max_budget: Mapping[str, object], model: str, ) -> bool: """ @@ -108,116 +330,81 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): Raises: BudgetExceededError: If the end_user has exceeded the model budget """ - internal_model_max_budget: Final[GenericBudgetConfigType] = {} - - for _model, _budget_info in end_user_model_max_budget.items(): - internal_model_max_budget[_model] = BudgetConfig(**_budget_info) - - verbose_proxy_logger.debug( - "end_user internal_model_max_budget %s", - json.dumps(internal_model_max_budget, indent=4, default=str), + return await self._is_entity_within_model_budget( + entity_type=Litellm_EntityType.END_USER, + entity_id=end_user_id, + model_max_budget=end_user_model_max_budget, + model=model, + exceeded_message=f"LiteLLM End User: {end_user_id}, exceeded budget for model={model}", ) - # check if current model is in internal_model_max_budget - _current_model_budget_info: Final = self._get_request_model_budget_config( - model=model, internal_model_max_budget=internal_model_max_budget - ) - if _current_model_budget_info is None: - verbose_proxy_logger.debug("Model %s not found in end_user_model_max_budget", model) + async def _is_entity_within_model_budget( + self, + entity_type: Litellm_EntityType, + entity_id: str | None, + model_max_budget: Mapping[str, object] | None, + model: str, + exceeded_message: str, + ) -> bool: + if not model_max_budget: + return True + resolved: Final = resolve_model_budget(model=model, model_max_budget=model_max_budget) + if resolved is None: + verbose_proxy_logger.debug("Model %s not found in %s model_max_budget", model, entity_type.value) return True - # check if current model is within budget - if _current_model_budget_info.max_budget and _current_model_budget_info.max_budget > 0: - _current_spend: Final = await self._get_end_user_spend_for_model( - end_user_id=end_user_id, - model=model, - key_budget_config=_current_model_budget_info, - ) - if ( - _current_spend is not None - and _current_model_budget_info.max_budget is not None - and _current_spend > _current_model_budget_info.max_budget - ): - raise litellm.BudgetExceededError( - message=f"LiteLLM End User: {end_user_id}, exceeded budget for model={model}", - current_cost=_current_spend, - max_budget=_current_model_budget_info.max_budget, - entity_type=Litellm_EntityType.END_USER.value, - entity_id=end_user_id, - ) + max_budget: Final = resolved.budget_config.max_budget + if max_budget is None or max_budget < 0: + return True + current_spend: Final = await self._get_spend_for_model_budget( + entity_type=entity_type, + entity_id=entity_id, + model=model, + resolved=resolved, + ) + if current_spend >= max_budget: + raise litellm.BudgetExceededError( + message=exceeded_message, + current_cost=current_spend, + max_budget=max_budget, + entity_type=entity_type.value, + entity_id=entity_id, + ) return True - async def _get_end_user_spend_for_model( + async def _get_spend_for_model_budget( self, - end_user_id: str, + entity_type: Litellm_EntityType, + entity_id: str | None, model: str, - key_budget_config: BudgetConfig, - ) -> float | None: - # 1. model: directly look up `model` - end_user_model_spend_cache_key = ( - f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{model}:{key_budget_config.budget_duration}" - ) - _current_spend = await self.dual_cache.async_get_cache( - key=end_user_model_spend_cache_key, - ) + resolved: ResolvedModelBudget, + ) -> float: + """Spend charged to this budget in the current window, legacy counter included. - if _current_spend is None: - # 2. If 1, does not exist, check if passed as {custom_llm_provider}/model - end_user_model_spend_cache_key = f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{self._get_model_without_custom_llm_provider(model)}:{key_budget_config.budget_duration}" - _current_spend = await self.dual_cache.async_get_cache( - key=end_user_model_spend_cache_key, - ) - return _current_spend - - async def _get_virtual_key_spend_for_model( - self, - user_api_key_hash: str | None, - model: str, - key_budget_config: BudgetConfig, - ) -> float | None: + A counter that was never written is zero spend, not unknown spend. The + distinction only shows up at a zero-dollar cap, where skipping the + comparison would let the strictest possible limit admit every request. """ - Get the current spend for a virtual key for a model - - Lookup model in this order: - 1. model: directly look up `model` - 2. If 1, does not exist, check if passed as {custom_llm_provider}/model - """ - - # 1. model: directly look up `model` - virtual_key_model_spend_cache_key = ( - f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{user_api_key_hash}:{model}:{key_budget_config.budget_duration}" + spend_key: Final = model_budget_spend_cache_key( + entity_type=entity_type, + entity_id=entity_id, + budget_model=resolved.budget_model, + budget_duration=resolved.budget_config.budget_duration, ) - _current_spend = await self.dual_cache.async_get_cache( - key=virtual_key_model_spend_cache_key, + legacy_spend_key: Final = _legacy_request_model_spend_cache_key( + entity_type=entity_type, + entity_id=entity_id, + model=model, + resolved=resolved, ) + current_spend: Final = _as_spend(await self._cached_spend(spend_key)) + if legacy_spend_key is None or legacy_spend_key == spend_key: + return current_spend + return current_spend + _as_spend(await self._cached_spend(legacy_spend_key)) - if _current_spend is None: - # 2. If 1, does not exist, check if passed as {custom_llm_provider}/model - # if "/" in model, remove first part before "/" - eg. openai/o1-preview -> o1-preview - virtual_key_model_spend_cache_key = f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{user_api_key_hash}:{self._get_model_without_custom_llm_provider(model)}:{key_budget_config.budget_duration}" - _current_spend = await self.dual_cache.async_get_cache( - key=virtual_key_model_spend_cache_key, - ) - return _current_spend - - def _get_request_model_budget_config( - self, model: str, internal_model_max_budget: GenericBudgetConfigType - ) -> BudgetConfig | None: - """ - Get the budget config for the request model - - 1. Check if `model` is in `internal_model_max_budget` - 2. If not, check if `model` without custom llm provider is in `internal_model_max_budget` - """ - return internal_model_max_budget.get(model, None) or internal_model_max_budget.get( - self._get_model_without_custom_llm_provider(model), None - ) - - def _get_model_without_custom_llm_provider(self, model: str) -> str: - if "/" in model: - return model.split("/")[-1] - return model + async def _cached_spend(self, spend_key: str) -> float | None: + return await self.dual_cache.async_get_cache(key=spend_key) async def async_filter_deployments( self, @@ -245,80 +432,63 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): _litellm_params: Final[dict] = kwargs.get("litellm_params", {}) or {} _metadata: Final[dict] = _litellm_params.get("metadata", {}) or {} - user_api_key_model_max_budget: Final[dict | None] = _metadata.get("user_api_key_model_max_budget", None) - user_api_key_end_user_model_max_budget: Final[dict | None] = _metadata.get( - "user_api_key_end_user_model_max_budget", None - ) - if (user_api_key_model_max_budget is None or len(user_api_key_model_max_budget) == 0) and ( - user_api_key_end_user_model_max_budget is None or len(user_api_key_end_user_model_max_budget) == 0 - ): - verbose_proxy_logger.debug( - "Not running _PROXY_VirtualKeyModelMaxBudgetLimiter.async_log_success_event because user_api_key_model_max_budget and user_api_key_end_user_model_max_budget are None or empty." - ) - return + payload_metadata: Final = standard_logging_payload.get("metadata") or {} - response_cost: Final[float] = standard_logging_payload.get("response_cost", 0) # Use model_group (the user-facing model alias, e.g. "gpt-4o") when - # available. The enforcement path (is_key_within_model_budget) receives - # the model name from request_data["model"] which is the model group - # alias, so the spend tracking cache key must use the same name. - # Falling back to the deployment-level "model" field preserves - # behaviour for non-proxy or non-router deployments where model_group - # is None. + # available. The enforcement path receives the model name from + # request_data["model"] which is the model group alias, so the spend + # tracking cache key must resolve from the same name. Falling back to + # the deployment-level "model" field preserves behaviour for non-proxy + # or non-router deployments where model_group is None. model: Final = standard_logging_payload.get("model_group") or standard_logging_payload.get("model") - virtual_key: Final = standard_logging_payload.get("metadata", {}).get("user_api_key_hash") - end_user_id = standard_logging_payload.get("end_user") or standard_logging_payload.get("metadata", {}).get( - "user_api_key_end_user_id" - ) - if model is None: return - if ( - virtual_key is not None - and user_api_key_model_max_budget is not None - and len(user_api_key_model_max_budget) > 0 - ): - internal_model_max_budget: GenericBudgetConfigType = {} - for _model, _budget_info in user_api_key_model_max_budget.items(): - internal_model_max_budget[_model] = BudgetConfig(**_budget_info) - key_budget_config = self._get_request_model_budget_config( - model=model, internal_model_max_budget=internal_model_max_budget - ) - if key_budget_config is not None and key_budget_config.budget_duration: - virtual_spend_key: Final = ( - f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model}:{key_budget_config.budget_duration}" - ) - virtual_start_time_key: Final = f"virtual_key_budget_start_time:{virtual_key}" - await self._increment_spend_for_key( - budget_config=key_budget_config, - spend_key=virtual_spend_key, - start_time_key=virtual_start_time_key, - response_cost=response_cost, - ) + response_cost: Final[float] = standard_logging_payload.get("response_cost", 0) + entity_budgets: Final = ( + ( + Litellm_EntityType.KEY, + payload_metadata.get("user_api_key_hash"), + _metadata.get("user_api_key_model_max_budget"), + ), + ( + Litellm_EntityType.USER, + payload_metadata.get("user_api_key_user_id"), + _metadata.get("user_api_key_user_model_max_budget"), + ), + ( + Litellm_EntityType.END_USER, + standard_logging_payload.get("end_user") or payload_metadata.get("user_api_key_end_user_id"), + _metadata.get("user_api_key_end_user_model_max_budget"), + ), + ) - if ( - end_user_id is not None - and user_api_key_end_user_model_max_budget is not None - and len(user_api_key_end_user_model_max_budget) > 0 - ): - internal_model_max_budget: GenericBudgetConfigType = {} - for _model, _budget_info in user_api_key_end_user_model_max_budget.items(): - internal_model_max_budget[_model] = BudgetConfig(**_budget_info) - key_budget_config = self._get_request_model_budget_config( - model=model, internal_model_max_budget=internal_model_max_budget + resolved_budgets: Final = _resolve_entity_model_budgets(model=model, entity_budgets=entity_budgets) + if not resolved_budgets: + verbose_proxy_logger.debug( + "Not running _PROXY_VirtualKeyModelMaxBudgetLimiter.async_log_success_event: " + "no key, user or end-user model_max_budget covers model=%s", + model, + ) + return + + for entity_type, entity_id, resolved in resolved_budgets: + await self._increment_spend_for_key( + budget_config=resolved.budget_config, + spend_key=model_budget_spend_cache_key( + entity_type=entity_type, + entity_id=entity_id, + budget_model=resolved.budget_model, + budget_duration=resolved.budget_config.budget_duration, + ), + start_time_key=model_budget_start_time_cache_key( + entity_type=entity_type, + entity_id=entity_id, + budget_model=resolved.budget_model, + budget_duration=resolved.budget_config.budget_duration, + ), + response_cost=response_cost, ) - if key_budget_config is not None and key_budget_config.budget_duration: - end_user_spend_key: Final = ( - f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{model}:{key_budget_config.budget_duration}" - ) - end_user_start_time_key: Final = f"end_user_budget_start_time:{end_user_id}" - await self._increment_spend_for_key( - budget_config=key_budget_config, - spend_key=end_user_spend_key, - start_time_key=end_user_start_time_key, - response_cost=response_cost, - ) if self.dual_cache.redis_cache is not None: await self._push_in_memory_increments_to_redis() diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 2ec5c34958c..1541b8acfdc 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1943,6 +1943,8 @@ async def add_litellm_data_to_request( # Follow same pattern as team and API key budgets data[_metadata_variable_name]["user_api_key_user_spend"] = user_api_key_dict.user_spend data[_metadata_variable_name]["user_api_key_user_max_budget"] = user_api_key_dict.user_max_budget + user_model_budget: Final = user_api_key_dict.user_model_max_budget + data[_metadata_variable_name]["user_api_key_user_model_max_budget"] = user_model_budget # rebind-ok: out-param data[_metadata_variable_name]["user_api_key_metadata"] = strip_callback_config(user_api_key_dict.metadata) data[_metadata_variable_name]["user_api_key_team_metadata"] = strip_callback_config(user_api_key_dict.team_metadata) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 9c725c54d08..c2f5b8eeb8b 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -32,6 +32,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( object_permission_cache_key, user_object_permission_id_cache_key, ) +from litellm.proxy.hooks.model_max_budget_limiter import build_model_max_budget_usage from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks from litellm.proxy.management_endpoints.common_daily_activity import ( DailySpendRecord, @@ -817,6 +818,7 @@ def _build_user_info_response( keys: list[LiteLLM_VerificationToken] | None, team_list: list[TeamListResponseObject], teams_1: list[TeamListResponseObject] | None, + model_max_budget_usage: dict[str, dict[str, object]] | None = None, ) -> UserInfoResponse: """Create UserInfoResponse while filtering sensitive fields.""" if user_info is None and keys is not None: @@ -830,6 +832,8 @@ def _build_user_info_response( if isinstance(_user_info, dict): _user_info.pop("password", None) _user_info["metadata"] = _redact_scim_enterprise_metadata(_user_info.get("metadata")) + if model_max_budget_usage is not None: + _user_info["model_max_budget_usage"] = model_max_budget_usage return UserInfoResponse( user_id=user_id, @@ -864,7 +868,7 @@ async def user_info( --header 'Authorization: Bearer sk-1234' ``` """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import model_max_budget_limiter, prisma_client try: user_id = _normalize_user_info_user_id(request=request, user_id=user_id) @@ -910,6 +914,12 @@ async def user_info( keys=keys, team_list=team_list, teams_1=teams_1, + model_max_budget_usage=await build_model_max_budget_usage( + entity_type=Litellm_EntityType.USER, + entity_id=user_id, + model_max_budget=getattr(user_info, "model_max_budget", None), + cache=model_max_budget_limiter.dual_cache, + ), ) return response_data @@ -1007,7 +1017,7 @@ async def user_info_v2( --header 'Authorization: Bearer sk-1234' ``` """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import model_max_budget_limiter, prisma_client try: if prisma_client is None: @@ -1062,6 +1072,13 @@ async def user_info_v2( sso_user_id=user_data.get("sso_user_id"), teams=user_data.get("teams") or [], object_permission=user_data.get("object_permission"), + model_max_budget=user_data.get("model_max_budget"), + model_max_budget_usage=await build_model_max_budget_usage( + entity_type=Litellm_EntityType.USER, + entity_id=user_data.get("user_id", user_id), + model_max_budget=user_data.get("model_max_budget"), + cache=model_max_budget_limiter.dual_cache, + ), ) except Exception as e: verbose_proxy_logger.exception("litellm.proxy.proxy_server.user_info_v2(): Exception occured - %s", e) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 71218d6114b..bf42aeeec05 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -29,6 +29,7 @@ from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, s import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid +from litellm.caching.dual_cache import DualCache from litellm.constants import ( LENGTH_OF_LITELLM_GENERATED_KEY, LITELLM_PROXY_ADMIN_NAME, @@ -47,7 +48,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_s rotate_sso_identity_assertions_master_key, ) from litellm.proxy._types import * -from litellm.proxy._types import LiteLLM_VerificationToken, hash_token +from litellm.proxy._types import Litellm_EntityType, LiteLLM_VerificationToken, hash_token from litellm.proxy.auth.auth_checks import ( _delete_cache_key_object, can_team_access_model, @@ -73,9 +74,7 @@ from litellm.proxy.common_utils.rbac_utils import check_org_admin_can_generate_k from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks -from litellm.proxy.hooks.model_max_budget_limiter import ( - VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX, -) +from litellm.proxy.hooks.model_max_budget_limiter import build_model_max_budget_usage from litellm.proxy.management_endpoints.common_utils import ( _check_passthrough_routes_caller_permission, _is_user_org_admin_for_team, @@ -3511,62 +3510,17 @@ async def delete_key_fn( raise handle_exception_on_proxy(e) -async def _get_model_max_budget_current_spend( - api_key_hash: str, - model: str, - budget_config: BudgetConfig, - user_api_key_cache: UserApiKeyCache, -) -> float: - virtual_key_model_spend_cache_key = ( - f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{api_key_hash}:{model}:{budget_config.budget_duration}" - ) - current_spend: float | None = await user_api_key_cache.async_get_cache( - key=virtual_key_model_spend_cache_key, - ) - if current_spend is None: - model_without_prefix: Final = model.split("/")[-1] if "/" in model else model - virtual_key_model_spend_cache_key = ( - f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:" - f"{api_key_hash}:{model_without_prefix}:{budget_config.budget_duration}" - ) - current_spend = await user_api_key_cache.async_get_cache( - key=virtual_key_model_spend_cache_key, - ) - try: - return float(current_spend or 0.0) - except (TypeError, ValueError): - return 0.0 - - async def _build_model_max_budget_usage( api_key_hash: str, model_max_budget: Mapping[str, Mapping[str, object]], - user_api_key_cache: UserApiKeyCache | None, + user_api_key_cache: DualCache | None, ) -> dict[str, dict[str, object]]: - if user_api_key_cache is None or not model_max_budget: - return {} - - result: Final[dict[str, dict[str, object]]] = {} - for model, budget_info in model_max_budget.items(): - try: - budget_config = BudgetConfig.model_validate(budget_info) - if budget_config.budget_duration is None: - continue - duration_in_seconds(budget_config.budget_duration) - except Exception: # noqa: BLE001 - continue - spend = await _get_model_max_budget_current_spend( - api_key_hash=api_key_hash, - model=model, - budget_config=budget_config, - user_api_key_cache=user_api_key_cache, - ) - result[model] = { - "current_spend": round(spend, 4), - "budget_limit": budget_config.max_budget, - "time_period": budget_config.budget_duration, - } - return result + return await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id=api_key_hash, + model_max_budget=model_max_budget, + cache=user_api_key_cache, + ) @router.post( @@ -3596,7 +3550,10 @@ async def info_key_fn_v2( -d {"keys": ["sk-1", "sk-2", "sk-3"]} ``` """ - from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + from litellm.proxy.proxy_server import ( + model_max_budget_limiter, + prisma_client, + ) try: if prisma_client is None: @@ -3648,7 +3605,7 @@ async def info_key_fn_v2( k_dict["model_max_budget_usage"] = await _build_model_max_budget_usage( api_key_hash=k_token_hash, model_max_budget=model_max_budget, - user_api_key_cache=user_api_key_cache, + user_api_key_cache=model_max_budget_limiter.dual_cache, ) filtered_key_info.append(k_dict) @@ -3707,7 +3664,10 @@ async def info_key_fn( -H "Authorization: Bearer sk-test-example-key-123" ``` """ - from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + from litellm.proxy.proxy_server import ( + model_max_budget_limiter, + prisma_client, + ) try: if prisma_client is None: @@ -3760,7 +3720,7 @@ async def info_key_fn( key_info["model_max_budget_usage"] = await _build_model_max_budget_usage( api_key_hash=key_token_hash, model_max_budget=model_max_budget, - user_api_key_cache=user_api_key_cache, + user_api_key_cache=model_max_budget_limiter.dual_cache, ) # Attach object_permission if object_permission_id is set @@ -3953,6 +3913,10 @@ async def generate_key_helper_fn( } if teams is not None: user_data["teams"] = teams + if model_max_budget: + # Only when supplied: the SSO and default-key callers reach this with the + # empty default, and writing that would clear an existing user's budgets. + user_data["model_max_budget"] = model_max_budget_json key_data: Final = { "token": token, "key_alias": key_alias, diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index d45421489e7..1915a853983 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -64,6 +64,7 @@ from litellm.proxy._types import ( ProxyException, UserAPIKeyAuth, ) +from litellm.proxy.auth.auth_utils import request_dispatched_to_pass_through_endpoint from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, @@ -568,6 +569,22 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): _metadata["user_api_key"] = user_api_key_dict.api_key _metadata["litellm_parent_otel_span"] = user_api_key_dict.parent_otel_span _metadata["user_api_key_budget_reservation"] = user_api_key_dict.budget_reservation + # The per-model budget counters are keyed off these. get_sanitized_user_information_from_key + # returns StandardLoggingUserAPIKeyMetadata, which carries no budget field, so without this + # the post-call increment finds nothing and every passthrough request goes untracked and + # unenforced. Set after the client merge so a request body cannot supply its own budget. + # + # Only for the built-in provider routes. `get_model_from_request` returns + # None for a user-defined pass-through, deliberately: its body is forwarded + # verbatim, so `model` there names an UPSTREAM model rather than a + # LiteLLM-managed one. Enforcement is therefore skipped on those routes, and + # charging a counter anyway would track spend that nothing can refuse, and + # would attribute it to a budget the operator scoped to a LiteLLM model that + # merely shares the name. + if not request_dispatched_to_pass_through_endpoint(request): + _metadata["user_api_key_model_max_budget"] = user_api_key_dict.model_max_budget + _metadata["user_api_key_user_model_max_budget"] = user_api_key_dict.user_model_max_budget + _metadata["user_api_key_end_user_model_max_budget"] = user_api_key_dict.end_user_model_max_budget _metadata.update( LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict) ) diff --git a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py index 55459721906..fcdb3c9246c 100644 --- a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py +++ b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py @@ -2,16 +2,21 @@ import os import sys from unittest.mock import AsyncMock, patch -sys.path.insert( - 0, os.path.abspath("../..") -) # Adds the parent directory to the system-path +sys.path.insert(0, os.path.abspath("../..")) # Adds the parent directory to the system-path import pytest import litellm from litellm.caching.caching import DualCache +from datetime import datetime, timezone + +from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.proxy._types import Litellm_EntityType from litellm.proxy.hooks.model_max_budget_limiter import ( + _budget_model_candidates, _PROXY_VirtualKeyModelMaxBudgetLimiter, + build_model_max_budget_usage, + resolve_model_budget, ) from litellm.proxy._types import UserAPIKeyAuth from litellm.types.utils import BudgetConfig as GenericBudgetInfo @@ -24,41 +29,95 @@ def budget_limiter(): return _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) -# Test _get_model_without_custom_llm_provider -def test_get_model_without_custom_llm_provider(budget_limiter): +# Test _budget_model_candidates +def test_budget_model_candidates(): # Test with custom provider - assert ( - budget_limiter._get_model_without_custom_llm_provider("openai/gpt-4") == "gpt-4" - ) + assert _budget_model_candidates("openai/gpt-4") == ("openai/gpt-4", "gpt-4") - # Test without custom provider - assert budget_limiter._get_model_without_custom_llm_provider("gpt-4") == "gpt-4" + # Test without custom provider: no duplicate candidate + assert _budget_model_candidates("gpt-4") == ("gpt-4",) -# Test _get_request_model_budget_config -def test_get_request_model_budget_config(budget_limiter): - internal_budget = { - "gpt-4": GenericBudgetInfo(budget_limit=100.0, time_period="1d"), - "claude-3": GenericBudgetInfo(budget_limit=50.0, time_period="1d"), +@pytest.mark.parametrize( + "model,expected", + [ + ( + "bedrock/anthropic.claude-opus-4-8", + ( + "bedrock/anthropic.claude-opus-4-8", + "anthropic.claude-opus-4-8", + "claude-opus-4-8", + ), + ), + ( + "us.anthropic.claude-opus-4-8", + ( + "us.anthropic.claude-opus-4-8", + "anthropic.claude-opus-4-8", + "claude-opus-4-8", + ), + ), + ( + "bedrock/converse/us.amazon.nova-pro-v1:0", + ( + "bedrock/converse/us.amazon.nova-pro-v1:0", + "us.amazon.nova-pro-v1:0", + "amazon.nova-pro-v1:0", + "nova-pro-v1:0", + ), + ), + ], +) +def test_budget_model_candidates_reach_the_bedrock_family_name(model, expected): + """ + Bedrock ids carry a dotted vendor segment ("anthropic.", "amazon.") on top of + the optional cross-region prefix, so a budget configured under the bare + family name would otherwise never match Bedrock traffic: no enforcement and + no spend tracking at all. + """ + assert _budget_model_candidates(model) == expected + + +@pytest.mark.parametrize( + "model", + [ + "azure/gpt-4.1", + "gpt-image-1.5", + "not-a-real-model.with.dots", + "ft:gpt-4o:acme::abc", + ], +) +def test_budget_model_candidates_never_split_a_non_bedrock_dotted_name(model): + """ + Most dotted model ids are versions, not Bedrock vendor prefixes. Splitting one + would offer a garbage candidate ("gpt-4.1" -> "1") that could collide with an + unrelated budget entry, so the split is gated on litellm pricing the model as + a Bedrock model. + """ + for candidate in _budget_model_candidates(model): + assert candidate in (model, model.split("/")[-1]) + + +# Test resolve_model_budget +def test_resolve_model_budget(): + model_max_budget = { + "gpt-4": {"budget_limit": 100.0, "time_period": "1d"}, + "claude-3": {"budget_limit": 50.0, "time_period": "1d"}, } # Test direct model match - config = budget_limiter._get_request_model_budget_config( - model="gpt-4", internal_model_max_budget=internal_budget - ) - assert config.max_budget == 100.0 + resolved = resolve_model_budget(model="gpt-4", model_max_budget=model_max_budget) + assert resolved.budget_model == "gpt-4" + assert resolved.budget_config.max_budget == 100.0 - # Test model with provider - config = budget_limiter._get_request_model_budget_config( - model="openai/gpt-4", internal_model_max_budget=internal_budget - ) - assert config.max_budget == 100.0 + # Test model with provider: the counter is keyed on the CONFIGURED name, + # not the request name, so every reader looks it up the same way. + resolved = resolve_model_budget(model="openai/gpt-4", model_max_budget=model_max_budget) + assert resolved.budget_model == "gpt-4" + assert resolved.budget_config.max_budget == 100.0 # Test non-existent model - config = budget_limiter._get_request_model_budget_config( - model="non-existent", internal_model_max_budget=internal_budget - ) - assert config is None + assert resolve_model_budget(model="non-existent", model_max_budget=model_max_budget) is None # Test is_key_within_model_budget @@ -72,47 +131,47 @@ async def test_is_key_within_model_budget(budget_limiter): ) # Test when model is within budget - with patch.object( - budget_limiter, "_get_virtual_key_spend_for_model", return_value=50.0 - ): - assert ( - await budget_limiter.is_key_within_model_budget(user_api_key, "gpt-4") - is True - ) + with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=50.0): + assert await budget_limiter.is_key_within_model_budget(user_api_key, "gpt-4") is True # Test when model exceeds budget - with patch.object( - budget_limiter, "_get_virtual_key_spend_for_model", return_value=150.0 - ): + with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=150.0): with pytest.raises(litellm.BudgetExceededError): await budget_limiter.is_key_within_model_budget(user_api_key, "gpt-4") # Test model not in budget config - assert ( - await budget_limiter.is_key_within_model_budget(user_api_key, "non-existent") - is True + assert await budget_limiter.is_key_within_model_budget(user_api_key, "non-existent") is True + + +# Test _get_spend_for_model_budget +@pytest.mark.asyncio +async def test_get_spend_for_model_budget_reads_the_configured_model_key( + budget_limiter, +): + from litellm.proxy.hooks.model_max_budget_limiter import ( + VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX, ) + model_max_budget = {"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}} + # openai/gpt-4 resolves to the configured "gpt-4" entry, so the lookup must + # hit the same key async_log_success_event writes. + resolved = resolve_model_budget(model="openai/gpt-4", model_max_budget=model_max_budget) -# Test _get_virtual_key_spend_for_model -@pytest.mark.asyncio -async def test_get_virtual_key_spend_for_model(budget_limiter): - budget_config = GenericBudgetInfo(budget_limit=100.0, time_period="1d") + async def _spend(key): + return 50.0 if key == f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:test-key:gpt-4:1d" else None - # Mock cache get - with patch.object(budget_limiter.dual_cache, "async_get_cache", return_value=50.0): - spend = await budget_limiter._get_virtual_key_spend_for_model( - user_api_key_hash="test-key", model="gpt-4", key_budget_config=budget_config - ) - assert spend == 50.0 - - # Test with provider prefix - spend = await budget_limiter._get_virtual_key_spend_for_model( - user_api_key_hash="test-key", + with patch.object(budget_limiter.dual_cache, "async_get_cache", side_effect=_spend) as mock_get: + spend = await budget_limiter._get_spend_for_model_budget( + entity_type=Litellm_EntityType.KEY, + entity_id="test-key", model="openai/gpt-4", - key_budget_config=budget_config, + resolved=resolved, ) assert spend == 50.0 + assert [call.kwargs["key"] for call in mock_get.call_args_list] == [ + f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:test-key:gpt-4:1d", + f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:test-key:openai/gpt-4:1d", + ] @pytest.mark.asyncio @@ -138,9 +197,7 @@ async def test_async_log_success_event_uses_per_model_budget_duration(budget_lim "metadata": {"user_api_key_hash": virtual_key}, }, "litellm_params": { - "metadata": { - "user_api_key_model_max_budget": user_api_key_model_max_budget - }, + "metadata": {"user_api_key_model_max_budget": user_api_key_model_max_budget}, }, } with patch.object( @@ -148,15 +205,11 @@ async def test_async_log_success_event_uses_per_model_budget_duration(budget_lim "_increment_spend_for_key", new_callable=AsyncMock, ) as mock_increment: - await budget_limiter.async_log_success_event( - kwargs, response_obj=None, start_time=None, end_time=None - ) + await budget_limiter.async_log_success_event(kwargs, response_obj=None, start_time=None, end_time=None) mock_increment.assert_awaited_once() call_kwargs = mock_increment.call_args.kwargs spend_key = call_kwargs["spend_key"] - assert spend_key == ( - f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model}:{budget_duration}" - ) + assert spend_key == (f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model}:{budget_duration}") assert call_kwargs["response_cost"] == 0.05 @@ -164,9 +217,7 @@ async def test_async_log_success_event_uses_per_model_budget_duration(budget_lim @pytest.mark.asyncio async def test_is_end_user_within_model_budget(budget_limiter): # Test when model is within budget - with patch.object( - budget_limiter, "_get_end_user_spend_for_model", return_value=50.0 - ): + with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=50.0): assert ( await budget_limiter.is_end_user_within_model_budget( "test-user", @@ -177,9 +228,7 @@ async def test_is_end_user_within_model_budget(budget_limiter): ) # Test when model exceeds budget - with patch.object( - budget_limiter, "_get_end_user_spend_for_model", return_value=150.0 - ): + with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=150.0): with pytest.raises(litellm.BudgetExceededError): await budget_limiter.is_end_user_within_model_budget( "test-user", @@ -198,25 +247,31 @@ async def test_is_end_user_within_model_budget(budget_limiter): ) -# Test _get_end_user_spend_for_model +# Test _get_spend_for_model_budget for the end-user scope @pytest.mark.asyncio -async def test_get_end_user_spend_for_model(budget_limiter): - budget_config = GenericBudgetInfo(budget_limit=100.0, time_period="1d") +async def test_get_spend_for_end_user_model_budget(budget_limiter): + from litellm.proxy.hooks.model_max_budget_limiter import ( + END_USER_SPEND_CACHE_KEY_PREFIX, + ) - # Mock cache get - with patch.object(budget_limiter.dual_cache, "async_get_cache", return_value=50.0): - spend = await budget_limiter._get_end_user_spend_for_model( - end_user_id="test-user", model="gpt-4", key_budget_config=budget_config - ) - assert spend == 50.0 + model_max_budget = {"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}} + resolved = resolve_model_budget(model="openai/gpt-4", model_max_budget=model_max_budget) - # Test with provider prefix - spend = await budget_limiter._get_end_user_spend_for_model( - end_user_id="test-user", + async def _spend(key): + return 50.0 if key == f"{END_USER_SPEND_CACHE_KEY_PREFIX}:test-user:gpt-4:1d" else None + + with patch.object(budget_limiter.dual_cache, "async_get_cache", side_effect=_spend) as mock_get: + spend = await budget_limiter._get_spend_for_model_budget( + entity_type=Litellm_EntityType.END_USER, + entity_id="test-user", model="openai/gpt-4", - key_budget_config=budget_config, + resolved=resolved, ) assert spend == 50.0 + assert [call.kwargs["key"] for call in mock_get.call_args_list] == [ + f"{END_USER_SPEND_CACHE_KEY_PREFIX}:test-user:gpt-4:1d", + f"{END_USER_SPEND_CACHE_KEY_PREFIX}:test-user:openai/gpt-4:1d", + ] @pytest.mark.asyncio @@ -261,16 +316,12 @@ async def test_async_log_success_event_uses_model_group_for_cache_key(budget_lim "_increment_spend_for_key", new_callable=AsyncMock, ) as mock_increment: - await budget_limiter.async_log_success_event( - kwargs, response_obj=None, start_time=None, end_time=None - ) + await budget_limiter.async_log_success_event(kwargs, response_obj=None, start_time=None, end_time=None) mock_increment.assert_awaited_once() call_kwargs = mock_increment.call_args.kwargs spend_key = call_kwargs["spend_key"] # The cache key must use the model_group name, NOT the deployment name - assert spend_key == ( - f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model_group}:{budget_duration}" - ) + assert spend_key == (f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model_group}:{budget_duration}") assert call_kwargs["response_cost"] == 0.10 @@ -310,15 +361,11 @@ async def test_async_log_success_event_falls_back_to_model_when_no_model_group( "_increment_spend_for_key", new_callable=AsyncMock, ) as mock_increment: - await budget_limiter.async_log_success_event( - kwargs, response_obj=None, start_time=None, end_time=None - ) + await budget_limiter.async_log_success_event(kwargs, response_obj=None, start_time=None, end_time=None) mock_increment.assert_awaited_once() call_kwargs = mock_increment.call_args.kwargs spend_key = call_kwargs["spend_key"] - assert spend_key == ( - f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model}:{budget_duration}" - ) + assert spend_key == (f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model}:{budget_duration}") @pytest.mark.asyncio @@ -357,15 +404,11 @@ async def test_async_log_success_event_end_user_uses_model_group(budget_limiter) "_increment_spend_for_key", new_callable=AsyncMock, ) as mock_increment: - await budget_limiter.async_log_success_event( - kwargs, response_obj=None, start_time=None, end_time=None - ) + await budget_limiter.async_log_success_event(kwargs, response_obj=None, start_time=None, end_time=None) mock_increment.assert_awaited_once() call_kwargs = mock_increment.call_args.kwargs spend_key = call_kwargs["spend_key"] - assert spend_key == ( - f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{model_group}:{budget_duration}" - ) + assert spend_key == (f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{model_group}:{budget_duration}") @pytest.mark.asyncio @@ -393,9 +436,7 @@ async def test_async_log_success_event_uses_end_user_model_budget_duration( "metadata": {"user_api_key_end_user_id": end_user_id}, }, "litellm_params": { - "metadata": { - "user_api_key_end_user_model_max_budget": user_api_key_end_user_model_max_budget - }, + "metadata": {"user_api_key_end_user_model_max_budget": user_api_key_end_user_model_max_budget}, }, } with patch.object( @@ -403,15 +444,11 @@ async def test_async_log_success_event_uses_end_user_model_budget_duration( "_increment_spend_for_key", new_callable=AsyncMock, ) as mock_increment: - await budget_limiter.async_log_success_event( - kwargs, response_obj=None, start_time=None, end_time=None - ) + await budget_limiter.async_log_success_event(kwargs, response_obj=None, start_time=None, end_time=None) mock_increment.assert_awaited_once() call_kwargs = mock_increment.call_args.kwargs spend_key = call_kwargs["spend_key"] - assert spend_key == ( - f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{model}:{budget_duration}" - ) + assert spend_key == (f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{model}:{budget_duration}") assert call_kwargs["response_cost"] == 0.05 @@ -446,9 +483,7 @@ async def test_async_log_success_event_pushes_redis_increments_when_redis_config "_push_in_memory_increments_to_redis", new_callable=AsyncMock, ) as mock_push: - await limiter.async_log_success_event( - kwargs, response_obj=None, start_time=None, end_time=None - ) + await limiter.async_log_success_event(kwargs, response_obj=None, start_time=None, end_time=None) mock_push.assert_awaited_once() @@ -457,10 +492,7 @@ async def test_get_fallback_model_within_budget_returns_none_without_fallbacks( budget_limiter, ): user_api_key = UserAPIKeyAuth(token="test-key", budget_fallbacks={}) - assert ( - await budget_limiter.get_fallback_model_within_budget(user_api_key, "gpt-4") - is None - ) + assert await budget_limiter.get_fallback_model_within_budget(user_api_key, "gpt-4") is None @pytest.mark.asyncio @@ -472,12 +504,8 @@ async def test_get_fallback_model_within_budget_returns_first_within_budget( model_max_budget={"gpt-4o-mini": {"budget_limit": 100.0, "time_period": "1d"}}, budget_fallbacks={"gpt-4": ["gpt-4o-mini", "claude-haiku"]}, ) - with patch.object( - budget_limiter, "_get_virtual_key_spend_for_model", return_value=1.0 - ): - result = await budget_limiter.get_fallback_model_within_budget( - user_api_key, "gpt-4" - ) + with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=1.0): + result = await budget_limiter.get_fallback_model_within_budget(user_api_key, "gpt-4") assert result == "gpt-4o-mini" @@ -494,17 +522,15 @@ async def test_get_fallback_model_within_budget_skips_exhausted_fallback( budget_fallbacks={"gpt-4": ["gpt-4o-mini", "claude-haiku"]}, ) - async def _spend_for_model(user_api_key_hash, model, key_budget_config): - return 150.0 if model == "gpt-4o-mini" else 1.0 + async def _spend_for_model(entity_type, entity_id, model, resolved): + return 150.0 if resolved.budget_model == "gpt-4o-mini" else 1.0 with patch.object( budget_limiter, - "_get_virtual_key_spend_for_model", + "_get_spend_for_model_budget", side_effect=_spend_for_model, ): - result = await budget_limiter.get_fallback_model_within_budget( - user_api_key, "gpt-4" - ) + result = await budget_limiter.get_fallback_model_within_budget(user_api_key, "gpt-4") assert result == "claude-haiku" @@ -520,12 +546,8 @@ async def test_get_fallback_model_within_budget_returns_none_when_chain_exhauste }, budget_fallbacks={"gpt-4": ["gpt-4o-mini", "claude-haiku"]}, ) - with patch.object( - budget_limiter, "_get_virtual_key_spend_for_model", return_value=150.0 - ): - result = await budget_limiter.get_fallback_model_within_budget( - user_api_key, "gpt-4" - ) + with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=150.0): + result = await budget_limiter.get_fallback_model_within_budget(user_api_key, "gpt-4") assert result is None @@ -554,7 +576,761 @@ async def test_async_log_success_event_skips_redis_push_without_redis(budget_lim "_push_in_memory_increments_to_redis", new_callable=AsyncMock, ) as mock_push: - await budget_limiter.async_log_success_event( - kwargs, response_obj=None, start_time=None, end_time=None - ) + await budget_limiter.async_log_success_event(kwargs, response_obj=None, start_time=None, end_time=None) mock_push.assert_not_awaited() + + +def _success_kwargs( + *, + model_group, + deployment_model=None, + response_cost=0.5, + key_hash=None, + key_model_max_budget=None, + user_id=None, + user_model_max_budget=None, + end_user_id=None, + end_user_model_max_budget=None, +): + return { + "standard_logging_object": { + "response_cost": response_cost, + "model": deployment_model or model_group, + "model_group": model_group, + "end_user": end_user_id, + "metadata": { + "user_api_key_hash": key_hash, + "user_api_key_user_id": user_id, + "user_api_key_end_user_id": end_user_id, + }, + }, + "litellm_params": { + "metadata": { + "user_api_key_model_max_budget": key_model_max_budget, + "user_api_key_user_model_max_budget": user_model_max_budget, + "user_api_key_end_user_model_max_budget": end_user_model_max_budget, + }, + }, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request_model", + ["gpt-4", "openai/gpt-4"], + ids=["request_model_matches_budget_key", "request_model_carries_provider_prefix"], +) +async def test_logged_spend_is_visible_to_key_info_usage_and_enforcement(request_model): + """ + The counter written post-call, the counter enforcement reads and the counter + /key/info reports must be one and the same, including when the request model + is not byte-identical to the configured budget key. + + Regression: the increment used to be keyed on the REQUEST model while + /key/info only ever looked up the CONFIGURED model, so a key could be + actively blocked at 429 while reporting current_spend 0. + """ + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + key_hash = "vk-hash" + model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1d"}} + + await limiter.async_log_success_event( + _success_kwargs( + model_group=request_model, + response_cost=0.75, + key_hash=key_hash, + key_model_max_budget=model_max_budget, + ), + response_obj=None, + start_time=None, + end_time=None, + ) + + usage = await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id=key_hash, + model_max_budget=model_max_budget, + cache=dual_cache, + ) + assert usage == { + "gpt-4": { + "current_spend": 0.75, + "budget_limit": 1.0, + "time_period": "1d", + } + } + + user_api_key = UserAPIKeyAuth(token=key_hash, model_max_budget=model_max_budget) + # Still under the 1.0 limit. + assert await limiter.is_key_within_model_budget(user_api_key, request_model) is True + + await limiter.async_log_success_event( + _success_kwargs( + model_group=request_model, + response_cost=0.75, + key_hash=key_hash, + key_model_max_budget=model_max_budget, + ), + response_obj=None, + start_time=None, + end_time=None, + ) + with pytest.raises(litellm.BudgetExceededError): + await limiter.is_key_within_model_budget(user_api_key, request_model) + + usage_after = await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id=key_hash, + model_max_budget=model_max_budget, + cache=dual_cache, + ) + assert usage_after["gpt-4"]["current_spend"] == 1.5 + + +@pytest.mark.asyncio +async def test_user_model_budget_is_tracked_and_enforced(): + """ + An internal user's own model_max_budget must be incremented post-call and + enforced, independently of any key-level budget. + """ + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + user_id = "user-1" + user_model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1mo"}} + + assert ( + await limiter.is_user_within_model_budget( + user_id=user_id, + user_model_max_budget=user_model_max_budget, + model="openai/gpt-4", + ) + is True + ) + + await limiter.async_log_success_event( + _success_kwargs( + model_group="openai/gpt-4", + response_cost=1.5, + user_id=user_id, + user_model_max_budget=user_model_max_budget, + ), + response_obj=None, + start_time=None, + end_time=None, + ) + + assert await build_model_max_budget_usage( + entity_type=Litellm_EntityType.USER, + entity_id=user_id, + model_max_budget=user_model_max_budget, + cache=dual_cache, + ) == {"gpt-4": {"current_spend": 1.5, "budget_limit": 1.0, "time_period": "1mo"}} + + with pytest.raises(litellm.BudgetExceededError) as exc: + await limiter.is_user_within_model_budget( + user_id=user_id, + user_model_max_budget=user_model_max_budget, + model="openai/gpt-4", + ) + assert exc.value.entity_type == Litellm_EntityType.USER.value + + +@pytest.mark.asyncio +async def test_user_model_budget_counter_is_separate_from_the_key_counter(): + """ + A key budget and a user budget over the same model are two independent + counters, so one request must charge each exactly once. + """ + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + model_max_budget = {"gpt-4": {"budget_limit": 10.0, "time_period": "1d"}} + + await limiter.async_log_success_event( + _success_kwargs( + model_group="gpt-4", + response_cost=2.0, + key_hash="vk-hash", + key_model_max_budget=model_max_budget, + user_id="user-1", + user_model_max_budget=model_max_budget, + ), + response_obj=None, + start_time=None, + end_time=None, + ) + + assert await dual_cache.async_get_cache(key="virtual_key_spend:vk-hash:gpt-4:1d") == 2.0 + assert await dual_cache.async_get_cache(key="user_model_spend:user-1:gpt-4:1d") == 2.0 + + +@pytest.mark.asyncio +async def test_two_models_on_one_key_do_not_share_a_budget_window(): + """ + A key budgeting two models over different periods must own one window start + per model: a shared start lets the shorter period restart the longer one. + """ + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + model_max_budget = { + "gpt-4": {"budget_limit": 10.0, "time_period": "1d"}, + "claude-3": {"budget_limit": 10.0, "time_period": "30d"}, + } + + start_time_keys = [] + with patch.object(limiter, "_increment_spend_for_key", new_callable=AsyncMock) as mock_increment: + for model in ("gpt-4", "claude-3"): + await limiter.async_log_success_event( + _success_kwargs( + model_group=model, + key_hash="vk-hash", + key_model_max_budget=model_max_budget, + ), + response_obj=None, + start_time=None, + end_time=None, + ) + start_time_keys = [call.kwargs["start_time_key"] for call in mock_increment.call_args_list] + + assert start_time_keys == [ + "virtual_key_budget_start_time:vk-hash:gpt-4:1d", + "virtual_key_budget_start_time:vk-hash:claude-3:30d", + ] + assert len(set(start_time_keys)) == 2 + + +@pytest.mark.asyncio +async def test_no_increment_when_no_scope_budgets_the_model(): + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + with patch.object(limiter, "_increment_spend_for_key", new_callable=AsyncMock) as mock_increment: + await limiter.async_log_success_event( + _success_kwargs( + model_group="gpt-4", + key_hash="vk-hash", + key_model_max_budget={"claude-3": {"budget_limit": 1.0, "time_period": "1d"}}, + user_id="user-1", + user_model_max_budget={}, + ), + response_obj=None, + start_time=None, + end_time=None, + ) + mock_increment.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_build_model_max_budget_usage_skips_unusable_entries(): + """A malformed or period-less entry must be omitted, not crash the report.""" + dual_cache = DualCache() + await dual_cache.async_set_cache(key="virtual_key_spend:vk:gpt-4:1d", value=3.0) + + usage = await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id="vk", + model_max_budget={ + "gpt-4": {"budget_limit": 10.0, "time_period": "1d"}, + "no-period": {"budget_limit": 10.0}, + "bad-period": {"budget_limit": 10.0, "time_period": "not-a-duration"}, + }, + cache=dual_cache, + ) + assert usage == {"gpt-4": {"current_spend": 3.0, "budget_limit": 10.0, "time_period": "1d"}} + + +@pytest.mark.asyncio +async def test_bedrock_traffic_charges_the_bare_family_name_budget(): + """ + The reported case: a budget configured as "claude-opus-4-8" with traffic on + "bedrock/anthropic.claude-opus-4-8". Before the fix nothing matched, so spend + was never tracked and the budget was never enforced no matter how far over it + the key went. + """ + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + key_hash = "vk-hash" + model_max_budget = {"claude-opus-4-8": {"budget_limit": 1.0, "time_period": "18h"}} + user_api_key = UserAPIKeyAuth(token=key_hash, model_max_budget=model_max_budget) + + await limiter.async_log_success_event( + _success_kwargs( + model_group="bedrock/anthropic.claude-opus-4-8", + response_cost=1.5, + key_hash=key_hash, + key_model_max_budget=model_max_budget, + ), + response_obj=None, + start_time=None, + end_time=None, + ) + + assert await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id=key_hash, + model_max_budget=model_max_budget, + cache=dual_cache, + ) == { + "claude-opus-4-8": { + "current_spend": 1.5, + "budget_limit": 1.0, + "time_period": "18h", + } + } + + with pytest.raises(litellm.BudgetExceededError): + await limiter.is_key_within_model_budget(user_api_key, "bedrock/anthropic.claude-opus-4-8") + + +@pytest.mark.asyncio +async def test_user_model_budget_window_resets_when_the_period_elapses(): + """ + A monthly user budget must start a fresh window once the period elapses, + and the window start must be scoped to that one budget model so a second + model on a shorter period cannot drag it forward. + """ + from litellm.proxy.hooks.model_max_budget_limiter import ( + model_budget_spend_cache_key, + model_budget_start_time_cache_key, + ) + + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + user_id = "user-1" + user_model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1mo"}} + spend_key = model_budget_spend_cache_key( + entity_type=Litellm_EntityType.USER, + entity_id=user_id, + budget_model="gpt-4", + budget_duration="1mo", + ) + start_time_key = model_budget_start_time_cache_key( + entity_type=Litellm_EntityType.USER, + entity_id=user_id, + budget_model="gpt-4", + budget_duration="1mo", + ) + + kwargs = _success_kwargs( + model_group="gpt-4", + response_cost=1.5, + user_id=user_id, + user_model_max_budget=user_model_max_budget, + ) + await limiter.async_log_success_event(kwargs, response_obj=None, start_time=None, end_time=None) + assert await dual_cache.async_get_cache(key=spend_key) == 1.5 + with pytest.raises(litellm.BudgetExceededError): + await limiter.is_user_within_model_budget( + user_id=user_id, + user_model_max_budget=user_model_max_budget, + model="gpt-4", + ) + + # Age the window past its period. The next charge opens a new window rather + # than adding to the exhausted one. + elapsed = duration_in_seconds("1mo") + 60 + await dual_cache.async_set_cache( + key=start_time_key, + value=datetime.now(timezone.utc).timestamp() - elapsed, + ttl=elapsed, + ) + + await limiter.async_log_success_event(kwargs, response_obj=None, start_time=None, end_time=None) + assert await dual_cache.async_get_cache(key=spend_key) == 1.5 + assert await build_model_max_budget_usage( + entity_type=Litellm_EntityType.USER, + entity_id=user_id, + model_max_budget=user_model_max_budget, + cache=dual_cache, + ) == {"gpt-4": {"current_spend": 1.5, "budget_limit": 1.0, "time_period": "1mo"}} + + +@pytest.mark.asyncio +async def test_a_zero_dollar_cap_blocks_the_model(): + """ + 0 is the operator saying "nobody may spend anything on this model", which is + the strictest cap expressible, not the absence of one. Skipping it on + falsiness turned the strictest setting into no setting at all, so the model + stayed wide open. The dashboard editor can produce this value, so it has to + mean something. + """ + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + key = UserAPIKeyAuth( + token="hash-zero", + model_max_budget={"gpt-4": {"budget_limit": 0, "time_period": "1d"}}, + ) + + with pytest.raises(litellm.BudgetExceededError) as exc: + await limiter.is_key_within_model_budget(user_api_key_dict=key, model="gpt-4") + assert exc.value.max_budget == 0 + + +@pytest.mark.asyncio +async def test_a_zero_dollar_cap_is_reported_as_a_cap_not_as_absent(): + """The usage endpoints must show the 0 too, or an operator cannot see the block they configured.""" + assert await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id="hash-zero", + model_max_budget={"gpt-4": {"budget_limit": 0, "time_period": "1d"}}, + cache=DualCache(), + ) == {"gpt-4": {"current_spend": 0.0, "budget_limit": 0.0, "time_period": "1d"}} + + +@pytest.mark.asyncio +async def test_spend_exactly_at_the_cap_is_refused(): + """ + Spending the whole budget exhausts it. `>` let a caller sit exactly on the + limit and keep going, and every sibling budget check in the codebase + (RouterBudgetLimiting, the key and team budget checks) uses `>=`. + """ + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + budget = {"gpt-4": {"budget_limit": 2.0, "time_period": "1d"}} + key = UserAPIKeyAuth(token="hash-exact", model_max_budget=budget) + + await limiter.async_log_success_event( + _success_kwargs(model_group="gpt-4", response_cost=2.0, key_hash="hash-exact", key_model_max_budget=budget), + response_obj=None, + start_time=None, + end_time=None, + ) + + with pytest.raises(litellm.BudgetExceededError): + await limiter.is_key_within_model_budget(user_api_key_dict=key, model="gpt-4") + + +@pytest.mark.asyncio +async def test_usage_report_reads_every_counter_in_one_batched_lookup(): + """ + model_max_budget is caller-supplied and unbounded in size, so one cache + coroutine per configured model let a large map fan out into an unbounded + number of concurrent lookups on an endpoint anyone holding the key can call. + One batched read keeps it to a single round trip whatever the map's size. + """ + dual_cache = DualCache() + budget = {f"model-{i}": {"budget_limit": 1.0, "time_period": "1d"} for i in range(50)} + + with ( + patch.object(dual_cache, "async_batch_get_cache", new=AsyncMock(return_value=[None] * 50)) as batched, + patch.object(dual_cache, "async_get_cache", new=AsyncMock()) as single, + ): + usage = await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id="hash-many", + model_max_budget=budget, + cache=dual_cache, + ) + + assert batched.await_count == 1 + assert len(batched.await_args.kwargs["keys"]) == 50 + assert single.await_count == 0 + assert len(usage) == 50 + + +@pytest.mark.asyncio +async def test_usage_report_survives_a_batch_lookup_that_returns_nothing(): + """ + async_batch_get_cache swallows its own failures and returns None. Zipping + that against the budgets would raise and take the whole /key/info response + with it, so an unusable result has to read as a miss instead. + """ + dual_cache = DualCache() + with patch.object(dual_cache, "async_batch_get_cache", new=AsyncMock(return_value=None)): + assert await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id="hash-none", + model_max_budget={"gpt-4": {"budget_limit": 1.0, "time_period": "1d"}}, + cache=dual_cache, + ) == {"gpt-4": {"current_spend": 0.0, "budget_limit": 1.0, "time_period": "1d"}} + + +@pytest.mark.asyncio +async def test_one_malformed_scope_does_not_abort_the_other_scopes(): + """ + Every scope is resolved before any of them is incremented, so a single + unusable entry used to raise out of resolution and leave the key counter + unwritten too. The key's budget is well formed here and must still be + charged despite the user's entry being garbage. + """ + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + key_budget = {"gpt-4": {"budget_limit": 10.0, "time_period": "1d"}} + + await limiter.async_log_success_event( + _success_kwargs( + model_group="gpt-4", + response_cost=0.25, + key_hash="hash-mixed", + key_model_max_budget=key_budget, + user_id="user-mixed", + user_model_max_budget={"gpt-4": {"budget_limit": "not-a-number", "time_period": "1d"}}, + ), + response_obj=None, + start_time=None, + end_time=None, + ) + + assert await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id="hash-mixed", + model_max_budget=key_budget, + cache=dual_cache, + ) == {"gpt-4": {"current_spend": 0.25, "budget_limit": 10.0, "time_period": "1d"}} + + +@pytest.mark.asyncio +async def test_an_unusable_budget_entry_is_not_enforced_instead_of_raising(): + """ + A config typo must not turn every request for that model into a 500. It + cannot be keyed, so it cannot be enforced; the write path rejects these, so + reaching here means config.yaml or a direct DB edit. + """ + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + key = UserAPIKeyAuth( + token="hash-malformed", + model_max_budget={"gpt-4": {"budget_limit": "not-a-number", "time_period": "1d"}}, + ) + + assert await limiter.is_key_within_model_budget(user_api_key_dict=key, model="gpt-4") is True + + +def test_resolve_model_budget_returns_none_for_an_unusable_entry(): + assert ( + resolve_model_budget( + model="gpt-4", + model_max_budget={"gpt-4": {"budget_limit": "not-a-number", "time_period": "1d"}}, + ) + is None + ) + + +def test_a_malformed_specific_entry_does_not_hide_a_usable_family_budget(): + """ + The candidate chain is most-specific-first and already falls through an entry + that is ABSENT. An entry that will not parse is indistinguishable from absent + as far as enforcement goes, so it has to fall through too: otherwise one bad + provider-prefixed entry silently disables the valid bare-family budget sitting + next to it, and the model goes uncapped. + """ + resolved = resolve_model_budget( + model="openai/gpt-4", + model_max_budget={ + "openai/gpt-4": {"budget_limit": "not-a-number", "time_period": "1d"}, + "gpt-4": {"budget_limit": 7.0, "time_period": "1d"}, + }, + ) + + assert resolved is not None + assert resolved.budget_model == "gpt-4" + assert resolved.budget_config.max_budget == 7.0 + + +@pytest.mark.asyncio +async def test_a_malformed_specific_entry_still_enforces_the_family_budget(): + """The fall-through has to reach enforcement, not just resolution.""" + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + budget = { + "openai/gpt-4": {"budget_limit": "not-a-number", "time_period": "1d"}, + "gpt-4": {"budget_limit": 1.0, "time_period": "1d"}, + } + key = UserAPIKeyAuth(token="hash-fallthrough", model_max_budget=budget) + + await limiter.async_log_success_event( + _success_kwargs( + model_group="openai/gpt-4", + response_cost=2.0, + key_hash="hash-fallthrough", + key_model_max_budget=budget, + ), + response_obj=None, + start_time=None, + end_time=None, + ) + + with pytest.raises(litellm.BudgetExceededError): + await limiter.is_key_within_model_budget(user_api_key_dict=key, model="openai/gpt-4") + + +def test_documented_budget_spelling_survives_model_validate(): + """ + `budget_limit` / `time_period` are the spelling the docs, the CRUD endpoints + and the dashboard editor all use, and BudgetConfig maps them onto + `max_budget` / `budget_duration` inside its `__init__`. + + Pydantic v2 normally bypasses a custom `__init__` in `model_validate`, and + this code path validates rather than constructing. It works today, but that + is a property of the installed Pydantic rather than of anything in this + repository, so an upgrade could silently stop applying the mapping and + quietly disable every budget written in the documented spelling. Pinned here + so that becomes a red test instead of an outage. + """ + from litellm.types.utils import BudgetConfig + + validated = BudgetConfig.model_validate({"budget_limit": 5, "time_period": "1d"}) + assert validated.max_budget == 5.0 + assert validated.budget_duration == "1d" + + # Control: an unrecognised key must NOT populate max_budget, or the assertion + # above would also pass against a model that accepted anything at all. + ignored = BudgetConfig.model_validate({"bogus_limit": 5, "time_period": "1d"}) + assert ignored.max_budget is None + + +def test_resolution_accepts_both_documented_spellings(): + """The resolver is what enforcement, tracking and reporting all go through.""" + for budget in ( + {"gpt-4": {"budget_limit": 5, "time_period": "1d"}}, + {"gpt-4": {"max_budget": 5, "budget_duration": "1d"}}, + ): + resolved = resolve_model_budget(model="gpt-4", model_max_budget=budget) + assert resolved is not None, f"{budget} resolved to nothing" + assert resolved.budget_config.max_budget == 5.0 + assert resolved.budget_config.budget_duration == "1d" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "entity_type, prefix", + [ + (Litellm_EntityType.KEY, "virtual_key_spend"), + (Litellm_EntityType.END_USER, "end_user_model_spend"), + ], +) +async def test_a_pre_upgrade_counter_keyed_on_the_request_model_still_enforces(entity_type, prefix): + """An upgrading proxy must not hand out a second allowance for the window it is already in. + + Before the counter key moved to the configured budget model, spend for a + request on `openai/gpt-4` against a budget configured as `gpt-4` was both + written to and enforced on `{prefix}:{id}:openai/gpt-4:1d`. Reading only the + configured-model key finds that counter empty and admits another full budget + until the window expires. + """ + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + model_max_budget = {"gpt-4": {"budget_limit": 10.0, "time_period": "1d"}} + await limiter.dual_cache.async_set_cache(key=f"{prefix}:entity-1:openai/gpt-4:1d", value=25.0, ttl=86400) + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + if entity_type == Litellm_EntityType.KEY: + await limiter.is_key_within_model_budget( + user_api_key_dict=UserAPIKeyAuth(token="entity-1", model_max_budget=model_max_budget), + model="openai/gpt-4", + ) + else: + await limiter.is_end_user_within_model_budget( + end_user_id="entity-1", + end_user_model_max_budget=model_max_budget, + model="openai/gpt-4", + ) + assert exc_info.value.current_cost == 25.0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "legacy_spend, current_spend, expect_blocked", + [(6.0, 5.0, True), (2.0, 3.0, False)], +) +async def test_the_pre_upgrade_and_post_upgrade_counters_add_up_over_one_window( + legacy_spend, current_spend, expect_blocked +): + """The two counters hold disjoint halves of one window, so the window's spend is their sum. + + Nothing writes the request-model spelling once this version is running, so + the legacy counter is frozen at whatever the previous version charged and + the configured-model counter carries everything since. Either one alone + under-reports the window: 6 + 5 is over a cap of 10 that neither half + reaches on its own. + """ + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + model_max_budget = {"gpt-4": {"budget_limit": 10.0, "time_period": "1d"}} + await limiter.dual_cache.async_set_cache( + key="virtual_key_spend:entity-1:openai/gpt-4:1d", value=legacy_spend, ttl=86400 + ) + await limiter.dual_cache.async_set_cache(key="virtual_key_spend:entity-1:gpt-4:1d", value=current_spend, ttl=86400) + + async def enforce(): + return await limiter.is_key_within_model_budget( + user_api_key_dict=UserAPIKeyAuth(token="entity-1", model_max_budget=model_max_budget), + model="openai/gpt-4", + ) + + if expect_blocked: + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await enforce() + assert exc_info.value.current_cost == legacy_spend + current_spend + else: + assert await enforce() is True + + +@pytest.mark.asyncio +async def test_the_configured_model_counter_is_never_counted_twice(): + """When the request names the budget exactly there is no legacy counter, only the one key. + + Both keys are `virtual_key_spend:entity-1:gpt-4:1d` here, so a lookup that + added them without noticing would charge 12 against a cap of 10 and refuse a + key that has spent 6. + """ + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + await limiter.dual_cache.async_set_cache(key="virtual_key_spend:entity-1:gpt-4:1d", value=6.0, ttl=86400) + + assert ( + await limiter.is_key_within_model_budget( + user_api_key_dict=UserAPIKeyAuth( + token="entity-1", + model_max_budget={"gpt-4": {"budget_limit": 10.0, "time_period": "1d"}}, + ), + model="gpt-4", + ) + is True + ) + + +@pytest.mark.asyncio +async def test_the_pre_upgrade_counter_is_no_longer_read_a_window_after_start_up(monkeypatch): + """The carry is bounded, so it cannot become a permanent second lookup on every request. + + A counter written by the previous version belongs to a window that was + already open when this process replaced it, so once a full window has passed + since start-up there is nothing left for the lookup to find. + """ + import litellm.proxy.hooks.model_max_budget_limiter as limiter_module + + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + model_max_budget = {"gpt-4": {"budget_limit": 10.0, "time_period": "1d"}} + await limiter.dual_cache.async_set_cache(key="virtual_key_spend:entity-1:openai/gpt-4:1d", value=25.0, ttl=86400) + user_api_key = UserAPIKeyAuth(token="entity-1", model_max_budget=model_max_budget) + + # Control: within the first window since start-up the same counter blocks, + # so the assertion below cannot pass against a lookup that never worked. + with pytest.raises(litellm.BudgetExceededError): + await limiter.is_key_within_model_budget(user_api_key_dict=user_api_key, model="openai/gpt-4") + + monkeypatch.setattr(limiter_module, "_PROCESS_STARTED_AT", limiter_module.time.monotonic() - 86401) + assert await limiter.is_key_within_model_budget(user_api_key_dict=user_api_key, model="openai/gpt-4") is True + + +@pytest.mark.asyncio +async def test_the_user_scope_has_no_pre_upgrade_counter_to_carry(): + """The user scope is introduced by this change, so a request-model key under it is not one of ours. + + Reading one would invent a counter no previous version ever wrote, which is + the opposite of preserving one. + """ + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + model_max_budget = {"gpt-4": {"budget_limit": 10.0, "time_period": "1d"}} + await limiter.dual_cache.async_set_cache(key="user_model_spend:u1:openai/gpt-4:1d", value=25.0, ttl=86400) + + assert ( + await limiter.is_user_within_model_budget( + user_id="u1", user_model_max_budget=model_max_budget, model="openai/gpt-4" + ) + is True + ) + + # Control: the same overspend under the key this scope does own must block, + # or the assertion above would pass against a scope that enforces nothing. + await limiter.dual_cache.async_set_cache(key="user_model_spend:u1:gpt-4:1d", value=25.0, ttl=86400) + with pytest.raises(litellm.BudgetExceededError): + await limiter.is_user_within_model_budget( + user_id="u1", user_model_max_budget=model_max_budget, model="openai/gpt-4" + ) diff --git a/tests/proxy_unit_tests/test_user_api_key_auth.py b/tests/proxy_unit_tests/test_user_api_key_auth.py index 58dbe3ad370..e9566254dbc 100644 --- a/tests/proxy_unit_tests/test_user_api_key_auth.py +++ b/tests/proxy_unit_tests/test_user_api_key_auth.py @@ -7,9 +7,7 @@ import sys import litellm.proxy import litellm.proxy.proxy_server -sys.path.insert( - 0, os.path.abspath("../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../..")) # Adds the parent directory to the system path from typing import Dict, List, Optional from unittest.mock import MagicMock, patch, AsyncMock @@ -50,9 +48,7 @@ class Request: ), # Request with no client IP should not be allowed ], ) -def test_check_valid_ip( - allowed_ips: Optional[List[str]], client_ip: Optional[str], expected_result: bool -): +def test_check_valid_ip(allowed_ips: Optional[List[str]], client_ip: Optional[str], expected_result: bool): from litellm.proxy.auth.auth_utils import _check_valid_ip request = Request(client_ip) @@ -121,9 +117,7 @@ async def test_check_blocked_team(): last_refreshed_at=time.time(), ) await asyncio.sleep(1) - team_obj = LiteLLM_TeamTableCachedObj( - team_id=_team_id, blocked=False, last_refreshed_at=time.time() - ) + team_obj = LiteLLM_TeamTableCachedObj(team_id=_team_id, blocked=False, last_refreshed_at=time.time()) hashed_token = hash_token(user_key) print(f"STORING TOKEN UNDER KEY={hashed_token}") user_api_key_cache.set_cache(key=hashed_token, value=valid_token) @@ -173,9 +167,7 @@ async def test_team_object_has_object_permission_id(): request = Request(scope={"type": "http"}) request._url = URL(url="/chat/completions") - with patch( - "litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock - ) as mock_common_checks: + with patch("litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock) as mock_common_checks: mock_common_checks.return_value = True await user_api_key_auth(request=request, api_key="Bearer " + user_key) @@ -200,9 +192,7 @@ async def test_returned_user_api_key_auth(user_role, expected_role): from datetime import datetime new_obj = await _return_user_api_key_auth_obj( - user_obj=LiteLLM_UserTable( - user_role=user_role, user_id="", max_budget=None, user_email="" - ), + user_obj=LiteLLM_UserTable(user_role=user_role, user_id="", max_budget=None, user_email=""), api_key="hello-world", parent_otel_span=None, valid_token_dict={}, @@ -258,9 +248,7 @@ async def test_aaauser_personal_budgets(key_ownership): spend=20, ) - user_obj = LiteLLM_UserTable( - user_id=_user_id, spend=11, max_budget=10, user_email="" - ) + user_obj = LiteLLM_UserTable(user_id=_user_id, spend=11, max_budget=10, user_email="") user_api_key_cache.set_cache(key=hash_token(user_key), value=valid_token) user_api_key_cache.set_cache(key="{}".format(_user_id), value=user_obj) @@ -273,10 +261,7 @@ async def test_aaauser_personal_budgets(key_ownership): test_user_cache = getattr(litellm.proxy.proxy_server, "user_api_key_cache") - assert ( - test_user_cache.get_cache(key=hash_token(user_key), model_type=UserAPIKeyAuth) - == valid_token - ) + assert test_user_cache.get_cache(key=hash_token(user_key), model_type=UserAPIKeyAuth) == valid_token if key_ownership == "user_key": with pytest.raises(ProxyException) as exc_info: @@ -311,9 +296,7 @@ async def test_user_api_key_auth_fails_with_prohibited_params(prohibited_param): request.body = return_body try: - response = await user_api_key_auth( - request=request, api_key="Bearer " + user_key - ) + response = await user_api_key_auth(request=request, api_key="Bearer " + user_key) except Exception as e: print("error str=", str(e)) error_message = str(e.message) @@ -519,9 +502,7 @@ def _assert_api_key_from_custom_header(headers, custom_header_name, expected_api verbose_proxy_logger.setLevel(logging.DEBUG) request = MagicMock(spec=Request) request.headers = headers - api_key = get_api_key_from_custom_header( - request=request, custom_litellm_key_header_name=custom_header_name - ) + api_key = get_api_key_from_custom_header(request=request, custom_litellm_key_header_name=custom_header_name) assert api_key == expected_api_key @@ -572,9 +553,7 @@ from litellm.proxy._types import LitellmUserRoles (LitellmUserRoles.TEAM, "1234", "1234", True), ], ) -def test_allowed_route_inside_route( - user_role, auth_user_id, requested_user_id, expected_result -): +def test_allowed_route_inside_route(user_role, auth_user_id, requested_user_id, expected_result): from litellm.proxy.auth.auth_checks import allowed_route_check_inside_route from litellm.proxy._types import UserAPIKeyAuth, LitellmUserRoles @@ -715,9 +694,7 @@ async def test_soft_budget_alert(): try: # Call user_api_key_auth - response = await user_api_key_auth( - request=request, api_key="Bearer " + user_key - ) + response = await user_api_key_auth(request=request, api_key="Bearer " + user_key) # Assert the request was allowed (no exception raised) assert response is not None @@ -883,9 +860,7 @@ async def test_user_api_key_auth_websocket(): mock_websocket.url = URL(url="/ws") # Mock the return value of `user_api_key_auth` when it's called within the `user_api_key_auth_websocket` function - with patch( - "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", autospec=True - ) as mock_user_api_key_auth: + with patch("litellm.proxy.auth.user_api_key_auth.user_api_key_auth", autospec=True) as mock_user_api_key_auth: # Make the call to the WebSocket function await user_api_key_auth_websocket(mock_websocket) @@ -896,17 +871,11 @@ async def test_user_api_key_auth_websocket(): request_arg = mock_user_api_key_auth.call_args.kwargs["request"] # Verify that the request has headers set - assert hasattr( - request_arg, "headers" - ), "Request object should have headers attribute" - assert ( - "authorization" in request_arg.headers - ), "Request headers should contain authorization" + assert hasattr(request_arg, "headers"), "Request object should have headers attribute" + assert "authorization" in request_arg.headers, "Request headers should contain authorization" assert request_arg.headers["authorization"] == "Bearer some_api_key" - assert ( - mock_user_api_key_auth.call_args.kwargs["api_key"] == "Bearer some_api_key" - ) + assert mock_user_api_key_auth.call_args.kwargs["api_key"] == "Bearer some_api_key" @pytest.mark.asyncio @@ -929,9 +898,7 @@ async def test_user_api_key_auth_websocket_carries_asgi_path(): } mock_websocket.url = URL(url="/v1/realtime") - with patch( - "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", autospec=True - ) as mock_user_api_key_auth: + with patch("litellm.proxy.auth.user_api_key_auth.user_api_key_auth", autospec=True) as mock_user_api_key_auth: await user_api_key_auth_websocket(mock_websocket) request_arg = mock_user_api_key_auth.call_args.kwargs["request"] @@ -1127,9 +1094,7 @@ async def test_jwt_non_admin_team_route_access(monkeypatch): ) request._url = URL(url="/team/new") - monkeypatch.setattr( - litellm.proxy.proxy_server, "general_settings", {"enable_jwt_auth": True} - ) + monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", {"enable_jwt_auth": True}) # Initialize jwt_handler with a default LiteLLM_JWTAuth so that the # virtual_key_claim_field check in user_api_key_auth doesn't fail with @@ -1158,9 +1123,7 @@ async def test_jwt_non_admin_team_route_access(monkeypatch): ): try: await user_api_key_auth(request=request, api_key="Bearer fake.jwt.token") - pytest.fail( - "Expected this call to fail. Non-admin user should not access team routes." - ) + pytest.fail("Expected this call to fail. Non-admin user should not access team routes.") except ProxyException as e: print("e", e) assert "Only proxy admin can be used to generate" in str(e.message) @@ -1220,9 +1183,7 @@ async def test_user_api_key_from_query_param(): from litellm.proxy.proxy_server import hash_token, user_api_key_cache user_key = "sk-query-1234" - user_api_key_cache.set_cache( - key=hash_token(user_key), value=UserAPIKeyAuth(token=hash_token(user_key)) - ) + user_api_key_cache.set_cache(key=hash_token(user_key), value=UserAPIKeyAuth(token=hash_token(user_key))) setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") @@ -1235,9 +1196,7 @@ async def test_user_api_key_from_query_param(): "query_string": f"alt=sse&key={user_key}".encode(), } ) - request._url = URL( - url=f"/v1beta/models/gemini:streamGenerateContent?alt=sse&key={user_key}" - ) + request._url = URL(url=f"/v1beta/models/gemini:streamGenerateContent?alt=sse&key={user_key}") async def return_body(): return b"{}" @@ -1246,3 +1205,592 @@ async def test_user_api_key_from_query_param(): valid_token = await user_api_key_auth(request=request, api_key="") assert valid_token.token == hash_token(user_key) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "user_id,user_model_max_budget,expected_calls", + [ + ("u-1", {"gpt-4": {"budget_limit": 1.0, "time_period": "1mo"}}, 1), + ("u-1", {}, 0), + ("u-1", None, 0), + (None, {"gpt-4": {"budget_limit": 1.0, "time_period": "1mo"}}, 0), + ], + ids=["enforced", "empty_budget", "no_budget", "no_user_id"], +) +async def test_check_user_model_budget(user_id, user_model_max_budget, expected_calls): + """ + An internal user's model_max_budget must reach the limiter. Before this it was + stored on LiteLLM_UserTable, accepted by /user/new and /user/update, and read + by nothing, so a user-level per-model budget never blocked anything. + """ + from litellm.proxy.auth.user_api_key_auth import _check_user_model_budget + + calls = [] + + class _Limiter: + async def is_user_within_model_budget(self, user_id, user_model_max_budget, model): + calls.append((user_id, user_model_max_budget, model)) + return True + + valid_token = UserAPIKeyAuth( + token="hash", + user_id=user_id, + user_model_max_budget=user_model_max_budget, + ) + await _check_user_model_budget( + valid_token=valid_token, + model_max_budget_limiter=_Limiter(), + models=["gpt-4"], + ) + assert len(calls) == expected_calls + if expected_calls: + assert calls[0] == ("u-1", user_model_max_budget, "gpt-4") + + +@pytest.mark.asyncio +async def test_user_model_max_budget_is_threaded_onto_the_auth_object(): + """ + The limiter can only enforce what auth carries. Regression for the user row's + model_max_budget being dropped on the way into UserAPIKeyAuth. + """ + from datetime import datetime + + from litellm.proxy.auth.user_api_key_auth import _return_user_api_key_auth_obj + + budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1mo"}} + user_obj = LiteLLM_UserTable( + user_id="u-1", + max_budget=None, + spend=0.0, + user_email=None, + models=[], + model_max_budget=budget, + ) + + auth_obj = await _return_user_api_key_auth_obj( + user_obj=user_obj, + api_key="sk-1234", + parent_otel_span=None, + valid_token_dict={"token": "hash"}, + route="/chat/completions", + start_time=datetime.now(), + user_role=LitellmUserRoles.INTERNAL_USER, + ) + assert auth_obj.user_model_max_budget == budget + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "over_budget,expect_refusal", + [(True, True), (False, False)], + ids=["over_budget_is_refused", "under_budget_is_served"], +) +async def test_user_model_budget_is_enforced_through_user_api_key_auth(over_budget, expect_refusal): + """ + Drive the real auth entry point, not the helper. + + The user's model_max_budget lives on the user row, and the joint + verification-token view auth builds its token from does not carry it. A test + that only exercises the helper passes while the whole path is inert, so this + one goes through user_api_key_auth with a key that has no per-model budget of + its own and asserts the USER's budget decides the outcome. + """ + from fastapi import Request + from starlette.datastructures import URL + + from litellm.proxy._types import LiteLLM_UserTable, Litellm_EntityType, UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.hooks.model_max_budget_limiter import model_budget_spend_cache_key + from litellm.proxy.proxy_server import ( + hash_token, + model_max_budget_limiter, + user_api_key_cache, + ) + + user_id = "user-model-budget" + model = "gpt-4o" + key = "sk-user-model-budget" + hashed = hash_token(key) + user_model_max_budget = {model: {"budget_limit": 1.0, "time_period": "1mo"}} + + setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache) + setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + setattr(litellm.proxy.proxy_server, "prisma_client", "present") + + await user_api_key_cache.async_set_cache( + key=hashed, + value=UserAPIKeyAuth(token=hashed, user_id=user_id, models=[], model_max_budget={}), + model_type=UserAPIKeyAuth, + ) + await model_max_budget_limiter.dual_cache.async_set_cache( + key=model_budget_spend_cache_key( + entity_type=Litellm_EntityType.USER, + entity_id=user_id, + budget_model=model, + budget_duration="1mo", + ), + value=5.0 if over_budget else 0.25, + ttl=600, + ) + + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + async def return_body(): + return f'{{"model": "{model}"}}'.encode() + + request.body = return_body + + async def fake_get_user_object(**kwargs): + return LiteLLM_UserTable( + user_id=user_id, + max_budget=None, + spend=0.0, + user_email=None, + models=[], + model_max_budget=user_model_max_budget, + ) + + with patch( + "litellm.proxy.auth.user_api_key_auth.get_user_object", + new=fake_get_user_object, + ): + if expect_refusal: + with pytest.raises(Exception) as exc: + await user_api_key_auth(request=request, api_key="Bearer " + key) + assert "budget" in str(exc.value).lower() + assert user_id in str(exc.value) + else: + result = await user_api_key_auth(request=request, api_key="Bearer " + key) + # The budget must also reach the token, or the post-call increment + # has nothing to charge and the counter never grows. + assert result.user_model_max_budget == user_model_max_budget + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "over_budget,expect_refusal", + [(True, True), (False, False)], + ids=["over_budget_is_refused", "under_budget_is_served"], +) +async def test_jwt_user_model_budget_is_enforced_before_the_jwt_path_returns(over_budget, expect_refusal): + """ + JWT auth returns its own token instead of falling through to the + virtual-key budget checks, so the user's per-model budget has to be enforced + on that path explicitly. + + The dangerous shape is not "no tracking": the post-call increment charges the + JWT user's counter either way, so without this check the counter grows and + nothing ever reads it, which looks enforced and is not. + """ + from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import _check_user_model_budget + from litellm.proxy.hooks.model_max_budget_limiter import model_budget_spend_cache_key + from litellm.proxy.proxy_server import model_max_budget_limiter + + user_id = "jwt-user-model-budget" + model = "gpt-4o" + user_model_max_budget = {model: {"budget_limit": 1.0, "time_period": "1mo"}} + + await model_max_budget_limiter.dual_cache.async_set_cache( + key=model_budget_spend_cache_key( + entity_type=Litellm_EntityType.USER, + entity_id=user_id, + budget_model=model, + budget_duration="1mo", + ), + value=5.0 if over_budget else 0.25, + ttl=600, + ) + + # The token the JWT branch builds and returns. + valid_token = UserAPIKeyAuth( + api_key=None, + user_id=user_id, + user_model_max_budget=user_model_max_budget, + ) + + if expect_refusal: + with pytest.raises(litellm.BudgetExceededError) as exc: + await _check_user_model_budget( + valid_token=valid_token, + model_max_budget_limiter=model_max_budget_limiter, + models=[model], + ) + assert exc.value.entity_type == Litellm_EntityType.USER.value + else: + await _check_user_model_budget( + valid_token=valid_token, + model_max_budget_limiter=model_max_budget_limiter, + models=[model], + ) + + +def test_jwt_path_enforces_the_user_model_budget_before_returning(): + """ + The JWT branch returns early, so the enforcement call has to sit before that + return rather than in the virtual-key block. Assert on the call graph, since + a helper-level test passes whether or not the JWT path ever calls it. + """ + import ast + import inspect + import textwrap + + from litellm.proxy.auth import user_api_key_auth as auth_module + + tree = ast.parse(textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder))) + + def calls_before_each_return(node): + seen_check = [] + for child in ast.walk(node): + if isinstance(child, ast.Call): + fn = child.func + name = getattr(fn, "id", None) or getattr(fn, "attr", None) + if name == "_check_user_model_budget": + seen_check.append(child.lineno) + return seen_check + + check_lines = calls_before_each_return(tree) + assert check_lines, "_user_api_key_auth_builder never enforces the user model budget" + + jwt_returns = [ + n.lineno + for n in ast.walk(tree) + if isinstance(n, ast.Return) and isinstance(n.value, ast.Call) and getattr(n.value.func, "id", None) == "cast" + ] + assert jwt_returns, "expected the JWT branch's `return cast(UserAPIKeyAuth, valid_token)`" + assert any(check < jwt_return for check in check_lines for jwt_return in jwt_returns), ( + "the user model-budget check must run before the JWT branch returns" + ) + + +def test_every_jwt_branch_carries_the_user_model_budget(): + """ + Each JWT branch that builds or replaces `valid_token` has to put the user's + model budget on it, or the enforcement call a few lines later has nothing to + read and silently admits the request. + + The auto-register branch is the one that regressed: it REPLACES the token + built above it with a key-scoped one whose columns carry no user budget. + """ + import ast + import inspect + import textwrap + + from litellm.proxy.auth import user_api_key_auth as auth_module + + tree = ast.parse(textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder))) + + assignments = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Assign) + and any(isinstance(t, ast.Attribute) and t.attr == "user_model_max_budget" for t in node.targets) + ] + targets = { + t.value.id + for node in assignments + for t in node.targets + if isinstance(t, ast.Attribute) and isinstance(t.value, ast.Name) + } + assert "auto_registered" in targets, ( + f"the auto-registered JWT token must carry the user's model budget; only these are populated: {sorted(targets)}" + ) + assert "valid_token" in targets, "the virtual-key path must carry the user's model budget" + + +@pytest.mark.asyncio +async def test_user_budget_lookup_tolerates_an_unreadable_user(): + """ + `get_user_object(user_id_upsert=False)` raises a bare Exception when the row + is simply ABSENT, which is the ordinary state for a custom-auth deployment + that never writes users to the proxy DB. Refusing on that exception would + turn "no user row" into a 4xx for every such request, and a transient DB + blip into a full outage. + + The virtual-key path makes the same call and swallows the same exception + ("Unable to get user from db/cache. Setting user_obj to None"), so this is + the established contract, not a shortcut. There is also nothing to enforce: + the budget being looked up lives on the row that could not be read. + """ + from litellm.caching.dual_cache import DualCache + from litellm.proxy.auth.user_api_key_auth import _read_user_model_max_budget + + prisma_client = MagicMock() + + with patch( + "litellm.proxy.auth.user_api_key_auth.get_user_object", + new=AsyncMock(side_effect=Exception("No user table row")), + ): + budget = await _read_user_model_max_budget( + user_id="user-with-no-row", + prisma_client=prisma_client, + user_api_key_cache=DualCache(), + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + + assert budget is None + + +@pytest.mark.asyncio +async def test_user_budget_lookup_is_also_unenforced_when_the_database_is_down(): + """ + KNOWN LIMITATION, pinned deliberately rather than discovered later. + + `get_user_object` cannot tell "row absent" from "database unreachable": the + absent case raises inside its own try (auth_checks.py:2177) and the handler + at :2213 rewrites every exception into the same + `ValueError("User doesn't exist in db...")`. A connection error, a query + timeout and a malformed row all reach us as that one type and message. + + So tolerating the absent case, which the test above requires, unavoidably + tolerates an outage too, and a user who DOES have a per-model budget goes + unenforced while the DB is unreachable. This is pre-existing behaviour of + `get_user_object` that the virtual-key path inherits identically; it is not + introduced here. Distinguishing them needs a dedicated exception type for + the absent case and a change to both auth paths. + """ + from litellm.caching.dual_cache import DualCache + from litellm.proxy.auth.user_api_key_auth import _read_user_model_max_budget + + db_down = ValueError("User doesn't exist in db. 'user_id'=u-1. Got error - Connection refused") + + with patch( + "litellm.proxy.auth.user_api_key_auth.get_user_object", + new=AsyncMock(side_effect=db_down), + ): + budget = await _read_user_model_max_budget( + user_id="u-1", + prisma_client=MagicMock(), + user_api_key_cache=DualCache(), + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + + assert budget is None + + +@pytest.mark.asyncio +async def test_user_budget_lookup_returns_the_budget_when_the_row_reads(): + """Positive control: the tolerance above must not be swallowing every result.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.auth.user_api_key_auth import _read_user_model_max_budget + + stored = {"claude-opus-4-8": {"budget_limit": 1.0, "time_period": "18h"}} + user_obj = MagicMock() + user_obj.model_max_budget = stored + + with patch( + "litellm.proxy.auth.user_api_key_auth.get_user_object", + new=AsyncMock(return_value=user_obj), + ): + budget = await _read_user_model_max_budget( + user_id="user-1", + prisma_client=MagicMock(), + user_api_key_cache=DualCache(), + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + + assert budget == stored + + +def test_zero_cost_models_skip_the_user_budget_check_on_every_path(): + """ + `skip_budget_checks` is computed per request for zero-cost models, and the + JWT branch logs "Skipping all budget checks" when it is set. Any enforcement + call that ignores it makes the same request behave differently depending on + whether the caller used a JWT or a virtual key, and makes that log a lie. + + Structural rather than behavioural on purpose: the defect is a call site + sitting outside a guard, and driving both auth paths to a zero-cost model + would prove it for the two requests exercised rather than for every site. + """ + import ast + import inspect + import textwrap + + from litellm.proxy.auth import user_api_key_auth as auth_module + + tree = ast.parse(textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder))) + + def guarded_by_skip(node: ast.AST, target: ast.AST) -> bool: + for parent in ast.walk(node): + if not isinstance(parent, ast.If): + continue + test = parent.test + is_skip_guard = ( + isinstance(test, ast.UnaryOp) + and isinstance(test.op, ast.Not) + and isinstance(test.operand, ast.Name) + and test.operand.id == "skip_budget_checks" + ) + if is_skip_guard and any(sub is target for sub in ast.walk(parent)): + return True + return False + + calls = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "_check_user_model_budget" + ] + assert len(calls) == 2, f"expected the JWT and virtual-key call sites, found {len(calls)}" + + unguarded = [c for c in calls if not guarded_by_skip(tree, c)] + assert not unguarded, ( + f"{len(unguarded)} _check_user_model_budget call(s) run even when " + "skip_budget_checks is set, so a zero-cost model is enforced on one auth path and not the other" + ) + + +def test_custom_auth_also_skips_budget_checks_for_zero_cost_models(): + """ + The custom-auth helper runs its own key, user and end-user per-model budget + checks. If it does not honour the zero-cost skip that the JWT and + virtual-key paths honour, the same free request is refused under one auth + method and served under the others. + + Asserted structurally, on the same reasoning as the sibling test: the defect + is a check sitting outside a guard, and it must hold for checks added later + rather than only for whichever request a behavioural test happened to drive. + """ + import ast + import inspect + import textwrap + + from litellm.proxy.auth import user_api_key_auth as auth_module + + src = textwrap.dedent(inspect.getsource(auth_module._run_post_custom_auth_checks)) + tree = ast.parse(src) + + assert "skip_budget_checks" in src, "the custom-auth path never computes the zero-cost skip flag" + + budget_calls = ( + "_check_key_model_budget_with_fallback", + "_check_user_model_budget", + "is_end_user_within_model_budget", + ) + + def guarding_ifs(target: ast.AST) -> list[ast.If]: + return [ + node for node in ast.walk(tree) if isinstance(node, ast.If) and any(sub is target for sub in ast.walk(node)) + ] + + def mentions_skip(node: ast.If) -> bool: + return any(isinstance(sub, ast.Name) and sub.id == "skip_budget_checks" for sub in ast.walk(node.test)) + + for call_name in budget_calls: + calls = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and ( + (isinstance(node.func, ast.Name) and node.func.id == call_name) + or (isinstance(node.func, ast.Attribute) and node.func.attr == call_name) + ) + ] + assert calls, f"{call_name} is no longer called here; update this invariant" + for call in calls: + assert any(mentions_skip(node) for node in guarding_ifs(call)), ( + f"{call_name} runs even for a zero-cost model, so custom auth refuses " + "requests the JWT and virtual-key paths serve" + ) + + +def test_custom_auth_attaches_the_user_budget_even_when_it_does_not_enforce(): + """ + The post-call spend hook reads `user_model_max_budget` off the token, so the + attach has to happen whether or not THIS request was enforceable. Gating it + on the same condition as the check leaves the user's counter uncharged for + every request with no resolvable model or a zero-cost one, which is exactly + the untracked-spend defect this PR fixes. + + Structural, because the failure is an assignment sitting inside a guard. + """ + import ast + import inspect + import textwrap + + from litellm.proxy.auth import user_api_key_auth as auth_module + + tree = ast.parse(textwrap.dedent(inspect.getsource(auth_module._run_post_custom_auth_checks))) + + attaches = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Assign) + and any(isinstance(t, ast.Attribute) and t.attr == "user_model_max_budget" for t in node.targets) + ] + assert attaches, "custom auth no longer attaches the user budget at all" + + for attach in attaches: + enclosing_ifs = [ + node for node in ast.walk(tree) if isinstance(node, ast.If) and any(sub is attach for sub in ast.walk(node)) + ] + assert not enclosing_ifs, ( + "the user budget is attached inside a conditional, so the spend hook " + "cannot charge the user counter whenever that condition is false" + ) + + +def test_mapped_key_jwt_falls_through_to_the_shared_user_budget_attach(): + """ + A JWT that maps to an existing virtual key resolves through the resolver + store, which builds the token from the KEY row alone and therefore carries + no user-level per-model budget. That branch sets `do_standard_jwt_auth = + False` precisely so it falls through to the shared virtual-key checks, where + the user row is loaded and its budget copied onto the token. + + Reviewed as a bypass three times, so the two halves it depends on are pinned + here: the branch must not return before the shared block, and the shared + block must copy the user row's budget onto the token. Structural on purpose, + because the claim is about control flow reaching a statement, and it has to + hold for branches added later rather than for one mocked request. + """ + import ast + import inspect + import textwrap + + from litellm.proxy.auth import user_api_key_auth as auth_module + + tree = ast.parse(textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder))) + + # Half one: the shared block copies the user row's budget onto the token. + copies_user_row = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Assign) + and any(isinstance(t, ast.Attribute) and t.attr == "user_model_max_budget" for t in node.targets) + and any(isinstance(v, ast.Attribute) and v.attr == "model_max_budget" for v in ast.walk(node.value)) + ] + assert copies_user_row, ( + "nothing copies the user row's model_max_budget onto the token, so a mapped-key " + "JWT reaches enforcement carrying the key's columns only" + ) + + # Half two: the mapped-key branch does not return before reaching it. + disables_standard_auth = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Assign) + and any(isinstance(t, ast.Name) and t.id == "do_standard_jwt_auth" for t in node.targets) + and isinstance(node.value, ast.Constant) + and node.value.value is False + ] + assert len(disables_standard_auth) == 1, "expected exactly one mapped-key branch" + marker = disables_standard_auth[0] + + enclosing = [ + node for node in ast.walk(tree) if isinstance(node, ast.If) and any(sub is marker for sub in node.body) + ] + assert enclosing, "could not locate the mapped-key branch body" + + returns_after = [ + node for node in ast.walk(enclosing[0]) if isinstance(node, ast.Return) and node.lineno > marker.lineno + ] + assert not returns_after, ( + "the mapped-key branch returns before the shared virtual-key checks, so the " + "user's per-model budget is never attached and never enforced" + ) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py index 6cc1d9e5add..28c82fdf528 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py @@ -1202,6 +1202,8 @@ def _fake_user_api_key_auth( model_max_budget=None, end_user_model_max_budget=None, end_user_id=None, + user_model_max_budget=None, + user_id=None, token=None, ): """Build a minimal stand-in for ``UserAPIKeyAuth`` with just the fields @@ -1220,6 +1222,8 @@ def _fake_user_api_key_auth( auth.model_max_budget = model_max_budget auth.end_user_model_max_budget = end_user_model_max_budget auth.end_user_id = end_user_id + auth.user_model_max_budget = user_model_max_budget + auth.user_id = user_id auth.token = token return auth @@ -1548,6 +1552,78 @@ async def test_summary_model_denied_when_key_over_model_budget(): assert result.applied_edits[0].get("error") == "summary_model_budget_exceeded" +async def test_summary_model_denied_when_user_over_model_budget(): + """Internal-user per-model budget is enforced for the summary subrequest too. + + This file propagates `user_api_key_user_model_max_budget` into the summary + subrequest's metadata, so its spend charges the user's counter. Enforcing + only the key and end-user scopes would let compaction increment a counter it + can never be refused by, which is the asymmetry this PR exists to remove. + """ + import litellm + + messages = _simple_messages() + mock_call = AsyncMock(return_value=_make_mock_response("x")) + + auth = _fake_user_api_key_auth( + key_models=["all-proxy-models"], + user_model_max_budget={"claude-haiku-4-5": {"budget_limit": 5}}, + user_id="user-over-budget", + token="hashed-token", + ) + + limiter = MagicMock() + limiter.is_user_within_model_budget = AsyncMock( + side_effect=litellm.BudgetExceededError( + message="over budget", current_cost=10, max_budget=5 + ) + ) + + with ( + patch( + "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting", + return_value="claude-haiku-4-5", + ), + patch("litellm.token_counter", return_value=200_000), + patch( + "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model", + mock_call, + ), + patch("litellm.proxy.proxy_server.model_max_budget_limiter", limiter), + ): + result = await apply_compact_20260112( + model=MODEL, + messages=messages, + tools=None, + system=None, + edit_spec=_EDIT_SPEC_DEFAULT, + user_api_key_auth=auth, + ) + + mock_call.assert_not_awaited() + assert result.applied_edits[0].get("error") == "summary_model_budget_exceeded" + + # The limiter is a mock, so it would accept any kwargs. Pin the call shape and + # check it against the real method, or a rename there would keep this test + # green while breaking compaction in production. + limiter.is_user_within_model_budget.assert_awaited_once_with( + user_id="user-over-budget", + user_model_max_budget={"claude-haiku-4-5": {"budget_limit": 5}}, + model="claude-haiku-4-5", + ) + import inspect + + from litellm.proxy.hooks.model_max_budget_limiter import ( + _PROXY_VirtualKeyModelMaxBudgetLimiter, + ) + + real_params = inspect.signature( + _PROXY_VirtualKeyModelMaxBudgetLimiter.is_user_within_model_budget + ).parameters + for kwarg in ("user_id", "user_model_max_budget", "model"): + assert kwarg in real_params, f"compact.py passes {kwarg}=, which the limiter no longer accepts" + + async def test_summary_model_denied_when_end_user_over_model_budget(): """End-user per-model budget is enforced for the summary subrequest too.""" import litellm diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 8a036d7e62f..da51d513b39 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -2975,6 +2975,7 @@ async def test_user_info_v2_response_shape(mocker): "updated_at": datetime(2024, 6, 1, tzinfo=timezone.utc), "sso_user_id": None, "teams": ["team-a", "team-b"], + "model_max_budget": {"gpt-3.5-turbo": {"budget_limit": 5.0, "time_period": "30d"}}, } async def mock_find_unique(*args, **kwargs): @@ -3018,9 +3019,20 @@ async def test_user_info_v2_response_shape(mocker): "sso_user_id", "teams", "object_permission", + "model_max_budget", + "model_max_budget_usage", } assert set(response_dict.keys()) == expected_fields + # The dashboard's user edit form hydrates its per-model budget rows from + # these two, so dropping them makes a save replace the user's budgets. + assert response_dict["model_max_budget"] == { + "gpt-3.5-turbo": {"budget_limit": 5.0, "time_period": "30d"} + } + assert response_dict["model_max_budget_usage"] == { + "gpt-3.5-turbo": {"current_spend": 0.0, "budget_limit": 5.0, "time_period": "30d"} + } + # Verify teams is a list of strings (team IDs), not team objects assert isinstance(response.teams, list) assert all(isinstance(t, str) for t in response.teams) @@ -4150,3 +4162,66 @@ async def test_user_info_v2_returns_the_mcp_entitlement(mocker): assert response.object_permission.mcp_tool_permissions == { "github": ["list_issues"] } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model_max_budget,expected_written", + [ + ( + {"claude-opus-4-8": {"budget_limit": 200.0, "time_period": "1mo"}}, + '{"claude-opus-4-8": {"budget_limit": 200.0, "time_period": "1mo"}}', + ), + (None, None), + ({}, None), + ], + ids=["supplied", "omitted", "empty"], +) +async def test_user_new_persists_model_max_budget( + monkeypatch, model_max_budget, expected_written +): + """ + /user/new used to echo model_max_budget back while writing {} to the user row, + so a per-model budget looked configured and was read by nothing. + + The omitted/empty cases are the other half: SSO and default-key callers reach + generate_key_helper_fn with no budget, and writing "{}" for them would clear + an existing user's budgets. + """ + from litellm.proxy.management_endpoints import key_management_endpoints + + captured = {} + + class _FakeUserRow: + models = [] + + class _FakePrisma: + async def insert_data(self, data, table_name): + if table_name == "user": + captured["user_data"] = dict(data) + return _FakeUserRow() + captured["key_data"] = dict(data) + return SimpleNamespace( + token=data.get("token"), + litellm_budget_table=None, + created_at=None, + updated_at=None, + ) + + async def get_data(self, *args, **kwargs): + return None + + import litellm.proxy.proxy_server as proxy_server + + monkeypatch.setattr(proxy_server, "prisma_client", _FakePrisma(), raising=False) + # model_max_budget is an enterprise feature; without this the call is rejected + # before it ever reaches the write this test is about. + monkeypatch.setattr(proxy_server, "premium_user", True, raising=False) + + await key_management_endpoints.generate_key_helper_fn( + request_type="user", + user_id="u-1", + model_max_budget=model_max_budget, + ) + + assert captured["user_data"].get("model_max_budget") == expected_written diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 069cfa01178..fff6368cfc6 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -13507,10 +13507,15 @@ async def test_info_key_fn_includes_model_max_budget_usage(monkeypatch): mock_prisma_client = AsyncMock() monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_user_api_key_cache = AsyncMock() - mock_user_api_key_cache.async_get_cache = AsyncMock(return_value=0.23) monkeypatch.setattr( "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache ) + # A real cache seeded at the real counter key: the spend only comes back if + # the endpoint computed virtual_key_spend:hashed_token_budget_test:gpt-4o:1d. + monkeypatch.setattr( + "litellm.proxy.proxy_server.model_max_budget_limiter.dual_cache", + await _budget_cache({"virtual_key_spend:hashed_token_budget_test:gpt-4o:1d": 0.23}), + ) mock_key_info = MagicMock(spec=LiteLLM_VerificationToken) mock_key_info.token = test_key_token @@ -13568,6 +13573,10 @@ async def test_info_key_fn_no_model_max_budget_skips_usage(monkeypatch): monkeypatch.setattr( "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.model_max_budget_limiter.dual_cache", + mock_user_api_key_cache, + ) mock_key_info = MagicMock(spec=LiteLLM_VerificationToken) mock_key_info.token = test_key_token @@ -13621,10 +13630,15 @@ async def test_info_key_fn_v2_includes_model_max_budget_usage(monkeypatch): mock_prisma_client = AsyncMock() monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_user_api_key_cache = AsyncMock() - mock_user_api_key_cache.async_get_cache = AsyncMock(return_value=0.55) monkeypatch.setattr( "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache ) + # A real cache seeded at the real counter key: the spend only comes back if + # the endpoint computed virtual_key_spend:hashed_token_v2_test:gpt-4o:7d. + monkeypatch.setattr( + "litellm.proxy.proxy_server.model_max_budget_limiter.dual_cache", + await _budget_cache({"virtual_key_spend:hashed_token_v2_test:gpt-4o:7d": 0.55}), + ) mock_key = MagicMock(spec=LiteLLM_VerificationToken) mock_key.token = test_key_token @@ -13680,10 +13694,15 @@ async def test_info_key_fn_budget_table_fallback(monkeypatch): mock_prisma_client = AsyncMock() monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_user_api_key_cache = AsyncMock() - mock_user_api_key_cache.async_get_cache = AsyncMock(return_value=1.20) monkeypatch.setattr( "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache ) + # A real cache seeded at the real counter key: the spend only comes back if + # the endpoint computed virtual_key_spend:hashed_token_budget_table_test:bedrock/anthropic.claude-opus-4:30d. + monkeypatch.setattr( + "litellm.proxy.proxy_server.model_max_budget_limiter.dual_cache", + await _budget_cache({"virtual_key_spend:hashed_token_budget_table_test:bedrock/anthropic.claude-opus-4:30d": 1.20}), + ) mock_key_info = MagicMock(spec=LiteLLM_VerificationToken) mock_key_info.token = test_key_token @@ -13748,10 +13767,15 @@ async def test_info_key_fn_v2_budget_table_fallback(monkeypatch): mock_prisma_client = AsyncMock() monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_user_api_key_cache = AsyncMock() - mock_user_api_key_cache.async_get_cache = AsyncMock(return_value=2.50) monkeypatch.setattr( "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache ) + # A real cache seeded at the real counter key: the spend only comes back if + # the endpoint computed virtual_key_spend:hashed_token_v2_bt_test:bedrock/anthropic.claude-opus-4:30d. + monkeypatch.setattr( + "litellm.proxy.proxy_server.model_max_budget_limiter.dual_cache", + await _budget_cache({"virtual_key_spend:hashed_token_v2_bt_test:bedrock/anthropic.claude-opus-4:30d": 2.50}), + ) mock_key = MagicMock(spec=LiteLLM_VerificationToken) mock_key.token = test_key_token @@ -13793,8 +13817,13 @@ async def test_info_key_fn_v2_budget_table_fallback(monkeypatch): @pytest.mark.asyncio -async def test_info_key_fn_provider_prefix_spend_fallback(monkeypatch): - """Cached spend for 'gpt-4o' matches budget key 'openai/gpt-4o' via suffix match.""" +async def test_info_key_fn_reads_the_configured_budget_model_key(monkeypatch): + """/key/info reads the one counter enforcement reads: the configured budget model. + + It used to probe a second, provider-stripped key because the counter was + written under the request model instead, which is what let a key report zero + usage while being blocked at 429. + """ from unittest.mock import AsyncMock, MagicMock from litellm.proxy._types import LiteLLM_VerificationToken @@ -13808,10 +13837,15 @@ async def test_info_key_fn_provider_prefix_spend_fallback(monkeypatch): mock_prisma_client = AsyncMock() monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_user_api_key_cache = AsyncMock() - mock_user_api_key_cache.async_get_cache = AsyncMock(side_effect=[None, 0.75]) monkeypatch.setattr( "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache ) + # A real cache seeded at the real counter key: the spend only comes back if + # the endpoint computed virtual_key_spend:hashed_token_prefix_test:openai/gpt-4o:7d. + monkeypatch.setattr( + "litellm.proxy.proxy_server.model_max_budget_limiter.dual_cache", + await _budget_cache({"virtual_key_spend:hashed_token_prefix_test:openai/gpt-4o:7d": 0.75}), + ) mock_key_info = MagicMock(spec=LiteLLM_VerificationToken) mock_key_info.token = test_key_token @@ -13847,7 +13881,22 @@ async def test_info_key_fn_provider_prefix_spend_fallback(monkeypatch): assert "model_max_budget_usage" in result["info"] usage = result["info"]["model_max_budget_usage"] assert usage["openai/gpt-4o"]["current_spend"] == 0.75 - assert mock_user_api_key_cache.async_get_cache.await_count == 2 + + +async def _budget_cache(seeded): + """A real DualCache holding spend at the given LITERAL counter keys. + + The keys are spelled out in full on purpose. Seeding via + model_budget_spend_cache_key would move the seed and the read together, so + any change to the key format would still match itself and these tests could + never fail, which is the exact bug they exist to catch. + """ + from litellm.caching.caching import DualCache + + cache = DualCache() + for key, spend in seeded.items(): + await cache.async_set_cache(key, spend) + return cache @pytest.mark.asyncio @@ -13872,19 +13921,16 @@ async def test_build_model_max_budget_usage_reads_current_cache_window(): _build_model_max_budget_usage, ) - mock_user_api_key_cache = AsyncMock() - mock_user_api_key_cache.async_get_cache = AsyncMock(return_value=0.30) + cache = await _budget_cache({"virtual_key_spend:some-hash:gpt-4o:30d": 0.30}) result = await _build_model_max_budget_usage( api_key_hash="some-hash", model_max_budget={"gpt-4o": {"budget_limit": 1.0, "time_period": "30d"}}, - user_api_key_cache=mock_user_api_key_cache, + user_api_key_cache=cache, ) + # 0.30 comes back only if the key matched virtual_key_spend:some-hash:gpt-4o:30d. assert result["gpt-4o"]["current_spend"] == 0.30 - mock_user_api_key_cache.async_get_cache.assert_awaited_once_with( - key="virtual_key_spend:some-hash:gpt-4o:30d" - ) @pytest.mark.asyncio @@ -13915,8 +13961,7 @@ async def test_build_model_max_budget_usage_skips_model_without_duration(): _build_model_max_budget_usage, ) - mock_user_api_key_cache = AsyncMock() - mock_user_api_key_cache.async_get_cache = AsyncMock(return_value=0.10) + cache = await _budget_cache({"virtual_key_spend:some-hash:gpt-4o:1d": 0.10}) result = await _build_model_max_budget_usage( api_key_hash="some-hash", @@ -13924,11 +13969,10 @@ async def test_build_model_max_budget_usage_skips_model_without_duration(): "gpt-4o": {"budget_limit": 1.0, "time_period": "1d"}, "gpt-3.5-turbo": {"budget_limit": 0.5}, }, - user_api_key_cache=mock_user_api_key_cache, + user_api_key_cache=cache, ) - assert "gpt-4o" in result + assert result["gpt-4o"]["current_spend"] == 0.10 assert "gpt-3.5-turbo" not in result - assert mock_user_api_key_cache.async_get_cache.await_count == 1 @pytest.mark.asyncio @@ -13961,8 +14005,7 @@ async def test_build_model_max_budget_usage_invalid_budget_config_skipped(): _build_model_max_budget_usage, ) - mock_user_api_key_cache = AsyncMock() - mock_user_api_key_cache.async_get_cache = AsyncMock(return_value=0.20) + cache = await _budget_cache({"virtual_key_spend:some-hash:gpt-3.5-turbo:7d": 0.20}) result = await _build_model_max_budget_usage( api_key_hash="some-hash", @@ -13970,32 +14013,35 @@ async def test_build_model_max_budget_usage_invalid_budget_config_skipped(): "gpt-4o": {"max_budget": "not-a-number", "budget_duration": "1d"}, "gpt-3.5-turbo": {"budget_limit": 0.5, "time_period": "7d"}, }, - user_api_key_cache=mock_user_api_key_cache, + user_api_key_cache=cache, ) assert "gpt-4o" not in result - assert "gpt-3.5-turbo" in result - assert mock_user_api_key_cache.async_get_cache.await_count == 1 + assert result["gpt-3.5-turbo"]["current_spend"] == 0.20 @pytest.mark.asyncio -async def test_build_model_max_budget_usage_provider_prefix_cache_fallback(): +async def test_build_model_max_budget_usage_reads_only_the_configured_model_key(): + """One lookup, at the configured budget model. + + The counter is written under the name the operator configured, so probing a + provider-stripped variant would read a key nothing writes. + """ from unittest.mock import AsyncMock from litellm.proxy.management_endpoints.key_management_endpoints import ( _build_model_max_budget_usage, ) - mock_user_api_key_cache = AsyncMock() - mock_user_api_key_cache.async_get_cache = AsyncMock(side_effect=[None, 0.55]) + cache = await _budget_cache({"virtual_key_spend:test-hash:openai/gpt-4o:7d": 0.55}) result = await _build_model_max_budget_usage( api_key_hash="test-hash", model_max_budget={"openai/gpt-4o": {"budget_limit": 2.0, "time_period": "7d"}}, - user_api_key_cache=mock_user_api_key_cache, + user_api_key_cache=cache, ) + # Seeded only under the configured name, so a provider-stripped probe reads 0.0. assert result["openai/gpt-4o"]["current_spend"] == 0.55 - assert mock_user_api_key_cache.async_get_cache.await_count == 2 def test_list_keys_substring_matching_param_defaults_to_false(): diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 090acf2dbb0..4c6ba23c88c 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -15,9 +15,7 @@ from fastapi import Request, UploadFile from starlette.datastructures import FormData, Headers, QueryParams from starlette.datastructures import UploadFile as StarletteUploadFile -sys.path.insert( - 0, os.path.abspath("../../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS, @@ -73,9 +71,7 @@ async def test_build_request_files_from_upload_file(): upload_file = UploadFile(file=file, filename="test.txt", headers=headers) upload_file.read = AsyncMock(return_value=file_content) - result = await HttpPassThroughEndpointHelpers._build_request_files_from_upload_file( - upload_file - ) + result = await HttpPassThroughEndpointHelpers._build_request_files_from_upload_file(upload_file) assert result == ("test.txt", file_content, "text/plain") # Test with Starlette UploadFile @@ -87,9 +83,7 @@ async def test_build_request_files_from_upload_file(): ) starlette_file.read = AsyncMock(return_value=file_content) - result = await HttpPassThroughEndpointHelpers._build_request_files_from_upload_file( - starlette_file - ) + result = await HttpPassThroughEndpointHelpers._build_request_files_from_upload_file(starlette_file) assert result == ("test2.txt", file_content, "text/plain") @@ -275,9 +269,7 @@ async def test_non_streaming_http_request_handler_multipart_with_non_empty_parse """ request = MagicMock(spec=Request) request.method = "POST" - request.headers = Headers( - {"content-type": "multipart/form-data; boundary=------------------------test"} - ) + request.headers = Headers({"content-type": "multipart/form-data; boundary=------------------------test"}) file_content = b"test file content" file = BytesIO(file_content) @@ -316,9 +308,7 @@ async def test_pass_through_request_failure_handler(): Critical Test: When a users pass through endpoint request fails, we must log the failure code, exception in litellm spend logs. """ with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: - with patch( - "litellm.llms.custom_httpx.http_handler.get_async_httpx_client" - ) as mock_get_client: + with patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client") as mock_get_client: with patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints.ProxyBaseLLMRequestProcessing" ) as mock_processing: @@ -329,9 +319,7 @@ async def test_pass_through_request_failure_handler(): # Setup mock for httpx client mock_client = MagicMock() mock_client.client = MagicMock() - mock_client.client.request = AsyncMock( - side_effect=httpx.HTTPError("Request failed") - ) + mock_client.client.request = AsyncMock(side_effect=httpx.HTTPError("Request failed")) mock_get_client.return_value = mock_client # Mock headers for custom headers @@ -364,9 +352,7 @@ async def test_pass_through_request_failure_handler(): # Verify the arguments to post_call_failure_hook call_args = mock_proxy_logging.post_call_failure_hook.call_args[1] assert call_args["user_api_key_dict"] == mock_user_api_key_dict - assert isinstance( - call_args["original_exception"], TypeError - ) # Now expecting TypeError + assert isinstance(call_args["original_exception"], TypeError) # Now expecting TypeError assert "traceback_str" in call_args @@ -410,27 +396,14 @@ def test_is_langfuse_route(): handler = PassThroughEndpointLogging() # Test positive cases - assert ( - handler.is_langfuse_route("http://localhost:4000/langfuse/api/public/traces") - is True - ) - assert ( - handler.is_langfuse_route( - "https://proxy.example.com/langfuse/api/public/sessions" - ) - is True - ) + assert handler.is_langfuse_route("http://localhost:4000/langfuse/api/public/traces") is True + assert handler.is_langfuse_route("https://proxy.example.com/langfuse/api/public/sessions") is True assert handler.is_langfuse_route("/langfuse/api/public/ingestion") is True assert handler.is_langfuse_route("http://localhost:4000/langfuse/") is True # Test negative cases - assert ( - handler.is_langfuse_route("https://api.openai.com/v1/chat/completions") is False - ) - assert ( - handler.is_langfuse_route("http://localhost:4000/anthropic/v1/messages") - is False - ) + assert handler.is_langfuse_route("https://api.openai.com/v1/chat/completions") is False + assert handler.is_langfuse_route("http://localhost:4000/anthropic/v1/messages") is False assert handler.is_langfuse_route("https://example.com/other") is False assert handler.is_langfuse_route("") is False @@ -447,17 +420,9 @@ def test_is_vertex_route_ignores_plain_predict_path_segment(): """ handler = PassThroughEndpointLogging() - assert ( - handler.is_vertex_route( - "https://upstream.example.com/ml/api/v1/time-series-forecast/predict" - ) - is False - ) + assert handler.is_vertex_route("https://upstream.example.com/ml/api/v1/time-series-forecast/predict") is False assert handler.is_vertex_route("https://upstream.example.com/api/v1/search") is False - assert ( - handler.is_vertex_route("https://upstream.example.com/predict/generateContent") - is False - ) + assert handler.is_vertex_route("https://upstream.example.com/predict/generateContent") is False assert ( handler.is_vertex_route( @@ -483,10 +448,7 @@ def test_is_vertex_route_ignores_plain_predict_path_segment(): ) is True ) - assert ( - handler.is_vertex_route("https://discoveryengine.googleapis.com/v1/x:search") - is True - ) + assert handler.is_vertex_route("https://discoveryengine.googleapis.com/v1/x:search") is True assert ( handler.is_vertex_route( @@ -545,9 +507,7 @@ async def test_custom_passthrough_predict_path_logs_via_generic_handler(): mock_vertex_handler.assert_not_called() handler._handle_logging.assert_awaited_once() - logged_object = handler._handle_logging.call_args.kwargs[ - "standard_logging_response_object" - ] + logged_object = handler._handle_logging.call_args.kwargs["standard_logging_response_object"] assert logged_object == {"response": '{"forecast": [1, 2, 3]}'} @@ -600,10 +560,7 @@ async def test_langfuse_passthrough_no_logging(): assert result is None # Verify that the passthrough_logging_payload was still set (this happens before the langfuse check) - assert ( - mock_logging_obj.model_call_details["passthrough_logging_payload"] - == passthrough_logging_payload - ) + assert mock_logging_obj.model_call_details["passthrough_logging_payload"] == passthrough_logging_payload def test_construct_target_url_with_subpath(): @@ -1051,9 +1008,7 @@ async def test_create_pass_through_route_with_cost_per_request(): # Mock the pass_through_request function to capture its call with ( - patch( - "litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_request" - ) as mock_pass_through, + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_request") as mock_pass_through, patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints.InitPassThroughEndpointHelpers.is_registered_pass_through_route" ) as mock_is_registered, @@ -1100,10 +1055,7 @@ def test_resolve_pass_through_request_timeout_precedence(): assert resolve_pass_through_request_timeout(endpoint_timeout=800) == 800.0 with patch("litellm.proxy.proxy_server.general_settings", {}): - assert ( - resolve_pass_through_request_timeout() - == DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS - ) + assert resolve_pass_through_request_timeout() == DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS def test_resolve_llm_passthrough_timeout_precedence(): @@ -1135,15 +1087,11 @@ async def test_pass_through_request_uses_resolved_timeout(): with patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client" ) as mock_get_client: - mock_proxy_logging.pre_call_hook = AsyncMock( - side_effect=lambda **kwargs: kwargs["data"] - ) + mock_proxy_logging.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) mock_client = MagicMock() mock_client.client = MagicMock() - mock_client.client.request = AsyncMock( - side_effect=httpx.HTTPError("Request failed") - ) + mock_client.client.request = AsyncMock(side_effect=httpx.HTTPError("Request failed")) mock_get_client.return_value = mock_client mock_request = MagicMock(spec=Request) @@ -1181,9 +1129,7 @@ async def test_create_pass_through_route_forwards_timeout(): ) with ( - patch( - "litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_request" - ) as mock_pass_through, + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_request") as mock_pass_through, patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints.InitPassThroughEndpointHelpers.is_registered_pass_through_route" ) as mock_is_registered, @@ -1296,9 +1242,7 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs(): "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_response_body" ) as mock_get_response_body: # Setup mock for pre_call_hook and post_call_failure_hook - mock_proxy_logging.pre_call_hook = AsyncMock( - return_value={"test": "data"} - ) + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={"test": "data"}) mock_proxy_logging.post_call_failure_hook = AsyncMock() mock_proxy_logging.post_call_response_headers_hook = AsyncMock( return_value={"x-callback-test": "value"} @@ -1308,9 +1252,7 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs(): mock_response = MagicMock() mock_response.status_code = 200 mock_response.headers = {} - mock_response.aread = AsyncMock( - return_value=b'{"success": true}' - ) + mock_response.aread = AsyncMock(return_value=b'{"success": true}') mock_response.text = '{"success": true}' mock_response.raise_for_status = MagicMock() @@ -1330,9 +1272,7 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs(): mock_request = MagicMock(spec=Request) mock_request.method = "POST" mock_request.url = "http://test-proxy.com/api/endpoint" - mock_request.body = AsyncMock( - return_value=b'{"message": "test request"}' - ) + mock_request.body = AsyncMock(return_value=b'{"message": "test request"}') mock_request.headers = Headers({}) mock_request.query_params = QueryParams({}) @@ -1411,9 +1351,7 @@ async def test_pass_through_request_streaming_marks_logging_obj_as_stream(): with patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints.PassThroughStreamingHandler.chunk_processor" ) as mock_chunk_processor: - mock_proxy_logging.pre_call_hook = AsyncMock( - return_value={"model": "claude-3", "stream": True} - ) + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={"model": "claude-3", "stream": True}) mock_proxy_logging.post_call_failure_hook = AsyncMock() mock_proxy_logging.post_call_response_headers_hook = AsyncMock( return_value={"x-callback-test": "value"} @@ -1438,9 +1376,7 @@ async def test_pass_through_request_streaming_marks_logging_obj_as_stream(): mock_request = MagicMock(spec=Request) mock_request.method = "POST" mock_request.url = "http://test-proxy.com/v1/messages" - mock_request.body = AsyncMock( - return_value=b'{"model": "claude-3", "stream": true}' - ) + mock_request.body = AsyncMock(return_value=b'{"model": "claude-3", "stream": true}') mock_request.headers = Headers({}) mock_request.query_params = QueryParams({}) @@ -1456,9 +1392,7 @@ async def test_pass_through_request_streaming_marks_logging_obj_as_stream(): assert async_client.send.call_args.kwargs["stream"] is True mock_chunk_processor.assert_called_once() - logging_obj = mock_chunk_processor.call_args.kwargs[ - "litellm_logging_obj" - ] + logging_obj = mock_chunk_processor.call_args.kwargs["litellm_logging_obj"] assert logging_obj.stream is True assert logging_obj.model_call_details["stream"] is True @@ -1479,9 +1413,7 @@ async def test_pass_through_request_sse_response_marks_logging_obj_as_stream(): with patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints.PassThroughStreamingHandler.chunk_processor" ) as mock_chunk_processor: - mock_proxy_logging.pre_call_hook = AsyncMock( - return_value={"model": "claude-3"} - ) + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={"model": "claude-3"}) mock_proxy_logging.post_call_failure_hook = AsyncMock() mock_proxy_logging.post_call_response_headers_hook = AsyncMock( return_value={"x-callback-test": "value"} @@ -1521,9 +1453,7 @@ async def test_pass_through_request_sse_response_marks_logging_obj_as_stream(): async_client.send.assert_awaited_once() mock_chunk_processor.assert_called_once() - logging_obj = mock_chunk_processor.call_args.kwargs[ - "litellm_logging_obj" - ] + logging_obj = mock_chunk_processor.call_args.kwargs["litellm_logging_obj"] assert logging_obj.stream is True assert logging_obj.model_call_details["stream"] is True @@ -1550,16 +1480,10 @@ async def test_create_pass_through_endpoint(): ) # Mock the database functions - with patch( - "litellm.proxy.proxy_server.get_config_general_settings" - ) as mock_get_config: - with patch( - "litellm.proxy.proxy_server.update_config_general_settings" - ) as mock_update_config: + with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config: + with patch("litellm.proxy.proxy_server.update_config_general_settings") as mock_update_config: # Mock existing config (empty list) - mock_get_config.return_value = ConfigFieldInfo( - field_name="pass_through_endpoints", field_value=[] - ) + mock_get_config.return_value = ConfigFieldInfo(field_name="pass_through_endpoints", field_value=[]) # Create test endpoint data test_endpoint = PassThroughGenericEndpoint( @@ -1629,12 +1553,8 @@ async def test_update_pass_through_endpoint(): ) # Mock the database functions - with patch( - "litellm.proxy.proxy_server.get_config_general_settings" - ) as mock_get_config: - with patch( - "litellm.proxy.proxy_server.update_config_general_settings" - ) as mock_update_config: + with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config: + with patch("litellm.proxy.proxy_server.update_config_general_settings") as mock_update_config: # Create existing endpoint data existing_endpoint_id = "test-endpoint-123" existing_endpoints = [ @@ -1731,18 +1651,14 @@ async def test_create_pass_through_endpoint_auth_true_enforces_allowlist(): registry: dict = {} with ( - patch( - "litellm.proxy.proxy_server.get_config_general_settings" - ) as mock_get_config, + patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config, patch("litellm.proxy.proxy_server.update_config_general_settings"), patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes", registry, ), ): - mock_get_config.return_value = ConfigFieldInfo( - field_name="pass_through_endpoints", field_value=[] - ) + mock_get_config.return_value = ConfigFieldInfo(field_name="pass_through_endpoints", field_value=[]) # auth is not passed -> defaults to True on PassThroughGenericEndpoint endpoint = PassThroughGenericEndpoint( @@ -1757,19 +1673,12 @@ async def test_create_pass_through_endpoint_auth_true_enforces_allowlist(): ) assert any(value.get("auth") is True for value in registry.values()) - assert ( - RouteChecks.is_auth_enforced_pass_through_route( - route="/secure-passthrough", method="POST" - ) - is True - ) + assert RouteChecks.is_auth_enforced_pass_through_route(route="/secure-passthrough", method="POST") is True post_request = MagicMock(spec=Request) post_request.method = "POST" - without_allowlist = UserAPIKeyAuth( - user_id="u", allowed_routes=["llm_api_routes"] - ) + without_allowlist = UserAPIKeyAuth(user_id="u", allowed_routes=["llm_api_routes"]) with pytest.raises(HTTPException) as exc_info: RouteChecks.is_virtual_key_allowed_to_call_route( route="/secure-passthrough", @@ -1826,9 +1735,7 @@ async def test_update_pass_through_endpoint_auth_true_enforces_allowlist(): ] with ( - patch( - "litellm.proxy.proxy_server.get_config_general_settings" - ) as mock_get_config, + patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config, patch("litellm.proxy.proxy_server.update_config_general_settings"), patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes", @@ -1851,19 +1758,12 @@ async def test_update_pass_through_endpoint_auth_true_enforces_allowlist(): user_api_key_dict=MagicMock(spec=UserAPIKeyAuth), ) - assert ( - RouteChecks.is_auth_enforced_pass_through_route( - route="/edited-passthrough", method="POST" - ) - is True - ) + assert RouteChecks.is_auth_enforced_pass_through_route(route="/edited-passthrough", method="POST") is True post_request = MagicMock(spec=Request) post_request.method = "POST" - without_allowlist = UserAPIKeyAuth( - user_id="u", allowed_routes=["llm_api_routes"] - ) + without_allowlist = UserAPIKeyAuth(user_id="u", allowed_routes=["llm_api_routes"]) with pytest.raises(HTTPException) as exc_info: RouteChecks.is_virtual_key_allowed_to_call_route( route="/edited-passthrough", @@ -1905,12 +1805,8 @@ async def test_update_pass_through_endpoint_preserves_auth_false(): ] with ( - patch( - "litellm.proxy.proxy_server.get_config_general_settings" - ) as mock_get_config, - patch( - "litellm.proxy.proxy_server.update_config_general_settings" - ) as mock_update_config, + patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config, + patch("litellm.proxy.proxy_server.update_config_general_settings") as mock_update_config, patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes", registry, @@ -1937,12 +1833,7 @@ async def test_update_pass_through_endpoint_preserves_auth_false(): persisted = mock_update_config.call_args[1]["data"].field_value[0] assert persisted["auth"] is False - assert ( - RouteChecks.is_auth_enforced_pass_through_route( - route="/public-passthrough", method="POST" - ) - is False - ) + assert RouteChecks.is_auth_enforced_pass_through_route(route="/public-passthrough", method="POST") is False @pytest.mark.asyncio @@ -1962,9 +1853,7 @@ async def test_update_pass_through_endpoint_not_found(): ) # Mock the database functions - with patch( - "litellm.proxy.proxy_server.get_config_general_settings" - ) as mock_get_config: + with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config: # Mock existing config with different endpoint existing_endpoints = [ { @@ -1982,9 +1871,7 @@ async def test_update_pass_through_endpoint_not_found(): ) # Create update data - update_data = PassThroughGenericEndpoint( - path="/test/endpoint", target="http://newapi.com/v2" - ) + update_data = PassThroughGenericEndpoint(path="/test/endpoint", target="http://newapi.com/v2") # Mock user API key dict mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) @@ -2023,12 +1910,8 @@ async def test_delete_pass_through_endpoint(): ) # Mock the database functions - with patch( - "litellm.proxy.proxy_server.get_config_general_settings" - ) as mock_get_config: - with patch( - "litellm.proxy.proxy_server.update_config_general_settings" - ) as mock_update_config: + with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config: + with patch("litellm.proxy.proxy_server.update_config_general_settings") as mock_update_config: # Create existing endpoint data endpoint_to_delete_id = "test-endpoint-123" other_endpoint_id = "other-endpoint-456" @@ -2106,9 +1989,7 @@ async def test_delete_pass_through_endpoint_not_found(): ) # Mock the database functions - with patch( - "litellm.proxy.proxy_server.get_config_general_settings" - ) as mock_get_config: + with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config: # Mock existing config with different endpoint existing_endpoints = [ { @@ -2199,14 +2080,8 @@ async def test_get_pass_through_endpoints_includes_config_and_db(): with patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints._get_pass_through_endpoints_from_config" ) as mock_get_config: - db_objects = [ - PassThroughGenericEndpoint(**ep, is_from_config=False) - for ep in db_endpoints - ] - config_objects = [ - PassThroughGenericEndpoint(**ep, is_from_config=True) - for ep in config_endpoints - ] + db_objects = [PassThroughGenericEndpoint(**ep, is_from_config=False) for ep in db_endpoints] + config_objects = [PassThroughGenericEndpoint(**ep, is_from_config=True) for ep in config_endpoints] mock_get_db.return_value = db_objects mock_get_config.return_value = config_objects @@ -2280,13 +2155,9 @@ async def test_delete_pass_through_endpoint_empty_list(): ) # Mock the database functions - with patch( - "litellm.proxy.proxy_server.get_config_general_settings" - ) as mock_get_config: + with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config: # Mock empty config - mock_get_config.return_value = ConfigFieldInfo( - field_name="pass_through_endpoints", field_value=None - ) + mock_get_config.return_value = ConfigFieldInfo(field_name="pass_through_endpoints", field_value=None) # Mock user API key dict mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) @@ -2325,9 +2196,7 @@ async def test_pass_through_request_query_params_forwarding(): ) as mock_get_response_body: # Setup mock for pre_call_hook test_body = {"name": "Azure Assistant", "model": "gpt-4o"} - mock_proxy_logging.pre_call_hook = AsyncMock( - return_value=test_body - ) + mock_proxy_logging.pre_call_hook = AsyncMock(return_value=test_body) mock_proxy_logging.post_call_response_headers_hook = AsyncMock( return_value={"x-callback-test": "value"} ) @@ -2336,9 +2205,7 @@ async def test_pass_through_request_query_params_forwarding(): mock_response = MagicMock() mock_response.status_code = 200 mock_response.headers = {"content-type": "application/json"} - mock_response.aread = AsyncMock( - return_value=b'{"id": "asst_123", "object": "assistant"}' - ) + mock_response.aread = AsyncMock(return_value=b'{"id": "asst_123", "object": "assistant"}') mock_response.text = '{"id": "asst_123", "object": "assistant"}' mock_response.raise_for_status = MagicMock() @@ -2360,20 +2227,12 @@ async def test_pass_through_request_query_params_forwarding(): # Create mock request with query parameters (Azure API version) mock_request = MagicMock(spec=Request) mock_request.method = "POST" - mock_request.url = ( - "http://localhost:4000/azure-assistant/openai/assistants" - ) - mock_request.body = AsyncMock( - return_value=json.dumps(test_body).encode() - ) - mock_request.headers = Headers( - {"Content-Type": "application/json"} - ) + mock_request.url = "http://localhost:4000/azure-assistant/openai/assistants" + mock_request.body = AsyncMock(return_value=json.dumps(test_body).encode()) + mock_request.headers = Headers({"Content-Type": "application/json"}) # Create QueryParams with api-version parameter - mock_request.query_params = QueryParams( - [("api-version", "2025-01-01-preview")] - ) + mock_request.query_params = QueryParams([("api-version", "2025-01-01-preview")]) # Create mock user API key dict mock_user_api_key_dict = MagicMock() @@ -2395,9 +2254,7 @@ async def test_pass_through_request_query_params_forwarding(): # The key assertion: query parameters should be preserved and passed to the HTTP handler assert "requested_query_params" in call_kwargs - assert call_kwargs["requested_query_params"] == { - "api-version": "2025-01-01-preview" - } + assert call_kwargs["requested_query_params"] == {"api-version": "2025-01-01-preview"} assert call_kwargs.get("forward_multipart") is False # Verify the target URL is correct @@ -2447,9 +2304,7 @@ async def _run_pass_through_and_capture_wire_url( "PassThroughEndpoint client not found in in_memory_llm_clients_cache; " "get_async_httpx_client may not be caching this provider." ) - cache_dict[cache_key] = SimpleNamespace( - client=httpx.AsyncClient(transport=httpx.MockTransport(transport_handler)) - ) + cache_dict[cache_key] = SimpleNamespace(client=httpx.AsyncClient(transport=httpx.MockTransport(transport_handler))) mock_request = MagicMock(spec=Request) mock_request.method = "GET" @@ -2458,18 +2313,14 @@ async def _run_pass_through_and_capture_wire_url( mock_request.body = AsyncMock(return_value=b"") mock_proxy_logging = MagicMock() - mock_proxy_logging.pre_call_hook = AsyncMock( - side_effect=lambda user_api_key_dict, data, call_type: data - ) + mock_proxy_logging.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data) mock_proxy_logging.post_call_failure_hook = AsyncMock() mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={}) mock_proxy_logging.get_proxy_hook = MagicMock(return_value=managed_files_hook) try: with ExitStack() as stack: - stack.enter_context( - patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging) - ) + stack.enter_context(patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging)) if managed_files_hook is not None: stack.enter_context( patch( @@ -2637,26 +2488,16 @@ async def test_filter_endpoints_by_team_allowed_routes_with_filter(): # Create test endpoints endpoints = [ - PassThroughGenericEndpoint( - id="endpoint-1", path="/api/allowed1", target="http://example.com/api1" - ), - PassThroughGenericEndpoint( - id="endpoint-2", path="/api/allowed2", target="http://example.com/api2" - ), - PassThroughGenericEndpoint( - id="endpoint-3", path="/api/notallowed", target="http://example.com/api3" - ), + PassThroughGenericEndpoint(id="endpoint-1", path="/api/allowed1", target="http://example.com/api1"), + PassThroughGenericEndpoint(id="endpoint-2", path="/api/allowed2", target="http://example.com/api2"), + PassThroughGenericEndpoint(id="endpoint-3", path="/api/notallowed", target="http://example.com/api3"), ] # Mock prisma client mock_prisma_client = MagicMock() mock_team = MagicMock() - mock_team.metadata = { - "allowed_passthrough_routes": ["/api/allowed1", "/api/allowed2"] - } - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_team - ) + mock_team.metadata = {"allowed_passthrough_routes": ["/api/allowed1", "/api/allowed2"]} + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_team) # Call the function result = await _filter_endpoints_by_team_allowed_routes( @@ -2671,9 +2512,7 @@ async def test_filter_endpoints_by_team_allowed_routes_with_filter(): assert result[1].path == "/api/allowed2" # Verify database call - mock_prisma_client.db.litellm_teamtable.find_unique.assert_called_once_with( - where={"team_id": "test-team-123"} - ) + mock_prisma_client.db.litellm_teamtable.find_unique.assert_called_once_with(where={"team_id": "test-team-123"}) @pytest.mark.asyncio @@ -2691,9 +2530,7 @@ async def test_filter_endpoints_by_team_allowed_routes_team_not_found(): # Create test endpoints endpoints = [ - PassThroughGenericEndpoint( - id="endpoint-1", path="/api/test", target="http://example.com/api" - ), + PassThroughGenericEndpoint(id="endpoint-1", path="/api/test", target="http://example.com/api"), ] # Mock prisma client to return None (team not found) @@ -2726,21 +2563,15 @@ async def test_filter_endpoints_by_team_allowed_routes_no_metadata(): # Create test endpoints endpoints = [ - PassThroughGenericEndpoint( - id="endpoint-1", path="/api/test1", target="http://example.com/api1" - ), - PassThroughGenericEndpoint( - id="endpoint-2", path="/api/test2", target="http://example.com/api2" - ), + PassThroughGenericEndpoint(id="endpoint-1", path="/api/test1", target="http://example.com/api1"), + PassThroughGenericEndpoint(id="endpoint-2", path="/api/test2", target="http://example.com/api2"), ] # Mock prisma client with team that has None metadata mock_prisma_client = MagicMock() mock_team = MagicMock() mock_team.metadata = None - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_team) # Call the function result = await _filter_endpoints_by_team_allowed_routes( @@ -2768,21 +2599,15 @@ async def test_filter_endpoints_by_team_allowed_routes_no_allowed_routes_key(): # Create test endpoints endpoints = [ - PassThroughGenericEndpoint( - id="endpoint-1", path="/api/test1", target="http://example.com/api1" - ), - PassThroughGenericEndpoint( - id="endpoint-2", path="/api/test2", target="http://example.com/api2" - ), + PassThroughGenericEndpoint(id="endpoint-1", path="/api/test1", target="http://example.com/api1"), + PassThroughGenericEndpoint(id="endpoint-2", path="/api/test2", target="http://example.com/api2"), ] # Mock prisma client with team that has metadata but no allowed_passthrough_routes mock_prisma_client = MagicMock() mock_team = MagicMock() mock_team.metadata = {"some_other_key": "some_value"} - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_team) # Call the function result = await _filter_endpoints_by_team_allowed_routes( @@ -2810,21 +2635,15 @@ async def test_filter_endpoints_by_team_allowed_routes_empty_allowed_list(): # Create test endpoints endpoints = [ - PassThroughGenericEndpoint( - id="endpoint-1", path="/api/test1", target="http://example.com/api1" - ), - PassThroughGenericEndpoint( - id="endpoint-2", path="/api/test2", target="http://example.com/api2" - ), + PassThroughGenericEndpoint(id="endpoint-1", path="/api/test1", target="http://example.com/api1"), + PassThroughGenericEndpoint(id="endpoint-2", path="/api/test2", target="http://example.com/api2"), ] # Mock prisma client with team that has empty allowed_passthrough_routes mock_prisma_client = MagicMock() mock_team = MagicMock() mock_team.metadata = {"allowed_passthrough_routes": []} - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_team) # Call the function result = await _filter_endpoints_by_team_allowed_routes( @@ -2850,29 +2669,21 @@ async def test_filter_endpoints_by_team_allowed_routes_partial_match(): # Create test endpoints endpoints = [ - PassThroughGenericEndpoint( - id="endpoint-1", path="/api/openai", target="http://example.com/openai" - ), + PassThroughGenericEndpoint(id="endpoint-1", path="/api/openai", target="http://example.com/openai"), PassThroughGenericEndpoint( id="endpoint-2", path="/api/anthropic", target="http://example.com/anthropic", ), - PassThroughGenericEndpoint( - id="endpoint-3", path="/api/azure", target="http://example.com/azure" - ), - PassThroughGenericEndpoint( - id="endpoint-4", path="/api/cohere", target="http://example.com/cohere" - ), + PassThroughGenericEndpoint(id="endpoint-3", path="/api/azure", target="http://example.com/azure"), + PassThroughGenericEndpoint(id="endpoint-4", path="/api/cohere", target="http://example.com/cohere"), ] # Mock prisma client with team that allows only 2 routes mock_prisma_client = MagicMock() mock_team = MagicMock() mock_team.metadata = {"allowed_passthrough_routes": ["/api/openai", "/api/azure"]} - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_team - ) + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_team) # Call the function result = await _filter_endpoints_by_team_allowed_routes( @@ -2904,9 +2715,7 @@ async def test_bedrock_router_passthrough_metadata_initialization(): ) # Mock ProxyBaseLLMRequestProcessing to verify it's used - with patch( - "litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing" - ) as mock_processing_class: + with patch("litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing") as mock_processing_class: # Setup mock instance mock_processor = MagicMock() mock_processing_class.return_value = mock_processor @@ -2914,12 +2723,8 @@ async def test_bedrock_router_passthrough_metadata_initialization(): # Mock successful response mock_response = MagicMock() mock_response.status_code = 200 - mock_response.aread = AsyncMock( - return_value=b'{"content": [{"text": "Hello"}]}' - ) - mock_processor.base_passthrough_process_llm_request = AsyncMock( - return_value=mock_response - ) + mock_response.aread = AsyncMock(return_value=b'{"content": [{"text": "Hello"}]}') + mock_processor.base_passthrough_process_llm_request = AsyncMock(return_value=mock_response) # Create mock request with headers mock_request = MagicMock(spec=Request) @@ -2986,18 +2791,10 @@ async def test_bedrock_router_passthrough_metadata_initialization(): call_kwargs = mock_processor.base_passthrough_process_llm_request.call_args[1] # These are the critical parameters that ensure metadata is properly initialized: - assert ( - call_kwargs["request"] == mock_request - ), "Request must be passed for header extraction" - assert ( - call_kwargs["user_api_key_dict"] == mock_user_api_key_dict - ), "User API key dict needed for metadata" - assert ( - call_kwargs["proxy_logging_obj"] == mock_proxy_logging - ), "Logging obj needed for hooks" - assert ( - call_kwargs["llm_router"] == mock_router - ), "Router needed for model routing" + assert call_kwargs["request"] == mock_request, "Request must be passed for header extraction" + assert call_kwargs["user_api_key_dict"] == mock_user_api_key_dict, "User API key dict needed for metadata" + assert call_kwargs["proxy_logging_obj"] == mock_proxy_logging, "Logging obj needed for hooks" + assert call_kwargs["llm_router"] == mock_router, "Router needed for model routing" assert call_kwargs["model"] == "my-bedrock-model", "Model name must be passed" # Verify response was returned @@ -3060,18 +2857,12 @@ async def test_add_litellm_data_to_request_adds_headers_to_metadata(): # Bedrock passthrough uses litellm_metadata to prevent key-level # tags from leaking into the provider payload (GH#30629). assert "litellm_metadata" in result, "litellm_metadata should be present in result" - assert ( - "headers" in result["litellm_metadata"] - ), "headers should be present in litellm_metadata" - assert isinstance( - result["litellm_metadata"]["headers"], dict - ), "headers should be a dictionary" + assert "headers" in result["litellm_metadata"], "headers should be present in litellm_metadata" + assert isinstance(result["litellm_metadata"]["headers"], dict), "headers should be a dictionary" # Verify specific headers are accessible (important for guardrails) headers = result["litellm_metadata"]["headers"] - assert ( - "user-agent" in headers or "User-Agent" in headers - ), "User-Agent header should be accessible in metadata" + assert "user-agent" in headers or "User-Agent" in headers, "User-Agent header should be accessible in metadata" # Also verify proxy_server_request has headers (original location) assert "proxy_server_request" in result @@ -3106,9 +2897,7 @@ async def test_create_pass_through_route_custom_body_url_target(): ) with ( - patch( - "litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_request" - ) as mock_pass_through, + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_request") as mock_pass_through, patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints.InitPassThroughEndpointHelpers.is_registered_pass_through_route" ) as mock_is_registered, @@ -3145,9 +2934,7 @@ async def test_create_pass_through_route_custom_body_url_target(): "retrievalQuery": {"text": "What is in the knowledge base?"}, } - setattr( - mock_request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, bedrock_body - ) + setattr(mock_request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, bedrock_body) await endpoint_func( request=mock_request, @@ -3185,9 +2972,7 @@ async def test_pass_through_request_non_streaming_uses_content_for_state_raw_bod mock_request.headers = Headers({"Content-Type": "application/json"}) mock_request.state = SimpleNamespace() setattr(mock_request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, raw_signed) - mock_request.body = AsyncMock( - return_value=json.dumps(parsed_from_wire).encode("utf-8") - ) + mock_request.body = AsyncMock(return_value=json.dumps(parsed_from_wire).encode("utf-8")) mock_user = MagicMock() mock_user.api_key = "sk-test" @@ -3260,9 +3045,7 @@ async def test_pass_through_request_streaming_uses_content_for_state_raw_body(): mock_request.headers = Headers({"Content-Type": "application/json"}) mock_request.state = SimpleNamespace() setattr(mock_request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, raw_signed) - mock_request.body = AsyncMock( - return_value=json.dumps(parsed_from_wire).encode("utf-8") - ) + mock_request.body = AsyncMock(return_value=json.dumps(parsed_from_wire).encode("utf-8")) mock_user = MagicMock() mock_user.api_key = "sk-test" @@ -3336,9 +3119,7 @@ async def test_create_pass_through_route_no_custom_body_falls_back(): ) with ( - patch( - "litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_request" - ) as mock_pass_through, + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_request") as mock_pass_through, patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints.InitPassThroughEndpointHelpers.is_registered_pass_through_route" ) as mock_is_registered, @@ -3407,32 +3188,12 @@ def test_is_registered_pass_through_route_with_custom_root(): } with patch("litellm.proxy.utils.get_server_root_path", return_value="/proxy"): - assert ( - InitPassThroughEndpointHelpers.is_registered_pass_through_route( - "/proxy/api/endpoint" - ) - is True - ) - assert ( - InitPassThroughEndpointHelpers.is_registered_pass_through_route( - "/api/endpoint" - ) - is True - ) + assert InitPassThroughEndpointHelpers.is_registered_pass_through_route("/proxy/api/endpoint") is True + assert InitPassThroughEndpointHelpers.is_registered_pass_through_route("/api/endpoint") is True with patch("litellm.proxy.utils.get_server_root_path", return_value="/"): - assert ( - InitPassThroughEndpointHelpers.is_registered_pass_through_route( - "/api/endpoint" - ) - is True - ) - assert ( - InitPassThroughEndpointHelpers.is_registered_pass_through_route( - "/proxy/api/endpoint" - ) - is False - ) + assert InitPassThroughEndpointHelpers.is_registered_pass_through_route("/api/endpoint") is True + assert InitPassThroughEndpointHelpers.is_registered_pass_through_route("/proxy/api/endpoint") is False # Clean up _registered_pass_through_routes.clear() @@ -3464,24 +3225,18 @@ def test_get_registered_pass_through_route_with_custom_root(): with patch("litellm.proxy.utils.get_server_root_path", return_value="/litellm"): # Prefixed incoming route - result = InitPassThroughEndpointHelpers.get_registered_pass_through_route( - "/litellm/chat/completions" - ) + result = InitPassThroughEndpointHelpers.get_registered_pass_through_route("/litellm/chat/completions") assert result is not None assert result["target"] == "http://api.example.com/v1/chat/completions" assert result["headers"]["Authorization"] == "Bearer token123" # Bare incoming route (get_request_route convention) - result = InitPassThroughEndpointHelpers.get_registered_pass_through_route( - "/chat/completions" - ) + result = InitPassThroughEndpointHelpers.get_registered_pass_through_route("/chat/completions") assert result is not None assert result["target"] == "http://api.example.com/v1/chat/completions" with patch("litellm.proxy.utils.get_server_root_path", return_value="/"): - result = InitPassThroughEndpointHelpers.get_registered_pass_through_route( - "/chat/completions" - ) + result = InitPassThroughEndpointHelpers.get_registered_pass_through_route("/chat/completions") assert result is not None assert result["target"] == "http://api.example.com/v1/chat/completions" @@ -3535,12 +3290,7 @@ def test_db_registered_pass_through_route_bare_path_convention( "litellm.proxy.utils.get_server_root_path", return_value=server_root_path, ): - assert ( - InitPassThroughEndpointHelpers.is_registered_pass_through_route( - incoming_route - ) - is should_match - ) + assert InitPassThroughEndpointHelpers.is_registered_pass_through_route(incoming_route) is should_match _registered_pass_through_routes.clear() @@ -3559,25 +3309,13 @@ def test_mapped_pass_through_routes_with_server_root_path(): with patch("litellm.proxy.utils.get_server_root_path", return_value="/litellm"): # prefixed route should match mapped routes like /vertex_ai assert ( - InitPassThroughEndpointHelpers.is_registered_pass_through_route( - "/litellm/vertex_ai/v1/projects/foo" - ) - is True - ) - assert ( - InitPassThroughEndpointHelpers.is_registered_pass_through_route( - "/litellm/bedrock/model/invoke" - ) + InitPassThroughEndpointHelpers.is_registered_pass_through_route("/litellm/vertex_ai/v1/projects/foo") is True ) + assert InitPassThroughEndpointHelpers.is_registered_pass_through_route("/litellm/bedrock/model/invoke") is True # bare route without prefix should not match when root is set - assert ( - InitPassThroughEndpointHelpers.is_registered_pass_through_route( - "/vertex_ai/v1/projects/foo" - ) - is False - ) + assert InitPassThroughEndpointHelpers.is_registered_pass_through_route("/vertex_ai/v1/projects/foo") is False @pytest.mark.asyncio @@ -3594,24 +3332,18 @@ async def test_multipart_passthrough_preserves_boundary(): mock_response = MagicMock() mock_response.status_code = 200 mock_response.headers = httpx.Headers({"content-type": "application/json"}) - mock_response.aread = AsyncMock( - return_value=b'{"filename": "test.txt", "size": 17}' - ) + mock_response.aread = AsyncMock(return_value=b'{"filename": "test.txt", "size": 17}') mock_response.text = '{"filename": "test.txt", "size": 17}' async def mock_httpx_request(method, url, **kwargs): # Verify that files parameter is passed (not json) assert "files" in kwargs, "Files should be passed for multipart requests" - file_parts = [ - value for name, value in kwargs["files"] if name == "file" - ] + file_parts = [value for name, value in kwargs["files"] if name == "file"] assert len(file_parts) == 1, "File field should be in files" # Verify content-type is NOT in headers (httpx will set it with correct boundary) headers = kwargs.get("headers", {}) - assert ( - "content-type" not in headers - ), "content-type should be removed for multipart" + assert "content-type" not in headers, "content-type should be removed for multipart" filename, content, content_type = file_parts[0] assert filename == "test.txt" @@ -3684,9 +3416,7 @@ def test_get_response_headers_strips_server_and_date(): "connection", "keep-alive", ): - assert ( - stripped not in lowered_keys - ), f"{stripped!r} must not be forwarded by passthrough" + assert stripped not in lowered_keys, f"{stripped!r} must not be forwarded by passthrough" # Application/business headers must still pass through. lowered = {k.lower(): v for k, v in result.items()} @@ -3724,9 +3454,7 @@ class TestStaleRouteCleanupOnReload: ) stack.enter_context(patch("litellm.proxy.proxy_server.premium_user", True)) mock_set_env = stack.enter_context( - patch( - "litellm.proxy.pass_through_endpoints.pass_through_endpoints.set_env_variables_in_header" - ) + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.set_env_variables_in_header") ) mock_set_env.return_value = {} return stack @@ -3771,14 +3499,10 @@ class TestStaleRouteCleanupOnReload: so the registry would hold both paths instead of only ``/b``. """ with self._patches(): - await initialize_pass_through_endpoints( - [{"path": "/a", "target": "http://example.com"}] - ) + await initialize_pass_through_endpoints([{"path": "/a", "target": "http://example.com"}]) assert self._paths_in_registry() == ["/a"] - await initialize_pass_through_endpoints( - [{"path": "/b", "target": "http://example.com"}] - ) + await initialize_pass_through_endpoints([{"path": "/b", "target": "http://example.com"}]) assert self._paths_in_registry() == ["/b"] @@ -3801,12 +3525,8 @@ class TestStaleRouteCleanupOnReload: ] ) - assert InitPassThroughEndpointHelpers.is_registered_pass_through_route( - "/live-passthrough" - ) - assert InitPassThroughEndpointHelpers.is_registered_pass_through_route( - "/live-passthrough/some/subpath" - ) + assert InitPassThroughEndpointHelpers.is_registered_pass_through_route("/live-passthrough") + assert InitPassThroughEndpointHelpers.is_registered_pass_through_route("/live-passthrough/some/subpath") # Regression (LIT-3538): a pre-call guardrail block on a passthrough endpoint @@ -3907,40 +3627,26 @@ async def _drive_pass_through_block(raised_exception): 400, ), ( - _FastAPIHTTPException( - status_code=400, detail={"error": "Violated moderation policy"} - ), + _FastAPIHTTPException(status_code=400, detail={"error": "Violated moderation policy"}), 400, ), ], ) -async def test_pre_call_guardrail_block_logs_warning_not_exception( - guardrail_exception, expected_code -): +async def test_pre_call_guardrail_block_logs_warning_not_exception(guardrail_exception, expected_code): status_code, logger = await _drive_pass_through_block(guardrail_exception) assert int(status_code) == expected_code - assert ( - logger.exception.call_count == 0 - ), "guardrail block must not be logged as an ERROR with a traceback" - assert ( - logger.warning.call_count == 1 - ), "guardrail block must be logged once at WARNING" + assert logger.exception.call_count == 0, "guardrail block must not be logged as an ERROR with a traceback" + assert logger.warning.call_count == 1, "guardrail block must be logged once at WARNING" @pytest.mark.asyncio async def test_non_guardrail_exception_still_logs_with_traceback(): - status_code, logger = await _drive_pass_through_block( - RuntimeError("upstream connection reset") - ) + status_code, logger = await _drive_pass_through_block(RuntimeError("upstream connection reset")) assert int(status_code) == 500 - assert ( - logger.exception.call_count == 1 - ), "a genuine failure must still be logged via verbose_proxy_logger.exception" - assert ( - logger.warning.call_count == 0 - ), "a genuine failure must not be downgraded to WARNING" + assert logger.exception.call_count == 1, "a genuine failure must still be logged via verbose_proxy_logger.exception" + assert logger.warning.call_count == 0, "a genuine failure must not be downgraded to WARNING" # Regression: generic config-based passthrough (`pass_through_request`) used to @@ -3979,9 +3685,7 @@ async def test_pass_through_request_non_streaming_upstream_error_returned_unchan ) as mock_success_handler: mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) mock_proxy_logging.post_call_failure_hook = AsyncMock() - mock_proxy_logging.post_call_response_headers_hook = AsyncMock( - return_value=None - ) + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) mock_processing.get_custom_headers.return_value = {} mock_success_handler.return_value = None @@ -4067,9 +3771,7 @@ async def test_pass_through_request_upstream_error_failure_hook_exception_is_swa mock_proxy_logging.post_call_failure_hook = AsyncMock( side_effect=RuntimeError("alerting integration misconfigured") ) - mock_proxy_logging.post_call_response_headers_hook = AsyncMock( - return_value=None - ) + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) mock_processing.get_custom_headers.return_value = {} mock_success_handler.return_value = None @@ -4119,9 +3821,7 @@ async def test_pass_through_request_streaming_upstream_error_returned_unchanged( ) as mock_success_handler: mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) mock_proxy_logging.post_call_failure_hook = AsyncMock() - mock_proxy_logging.post_call_response_headers_hook = AsyncMock( - return_value=None - ) + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) mock_success_handler.return_value = None async_client = MagicMock() @@ -4149,10 +3849,7 @@ async def test_pass_through_request_streaming_upstream_error_returned_unchanged( streamed_chunks = [chunk async for chunk in response.body_iterator] await asyncio.sleep(0) - streamed_bytes = b"".join( - chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") - for chunk in streamed_chunks - ) + streamed_bytes = b"".join(chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") for chunk in streamed_chunks) assert streamed_bytes == upstream_content assert json.loads(streamed_bytes) == _UPSTREAM_ERROR_BODY @@ -4198,9 +3895,7 @@ async def test_pass_through_request_non_streaming_success_unchanged(): ) as mock_success_handler: mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) mock_proxy_logging.post_call_failure_hook = AsyncMock() - mock_proxy_logging.post_call_response_headers_hook = AsyncMock( - return_value=None - ) + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) mock_processing.get_custom_headers.return_value = {} mock_success_handler.return_value = None @@ -4244,9 +3939,7 @@ async def test_pass_through_request_internal_failure_still_raises_proxy_exceptio from litellm.proxy._types import ProxyException with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: - mock_proxy_logging.pre_call_hook = AsyncMock( - side_effect=RuntimeError("auth backend unavailable") - ) + mock_proxy_logging.pre_call_hook = AsyncMock(side_effect=RuntimeError("auth backend unavailable")) mock_proxy_logging.post_call_failure_hook = AsyncMock() mock_request = MagicMock(spec=Request) @@ -4335,9 +4028,7 @@ def _inject_fake_passthrough_client(transport, timeout): def _enter_relay_logging_mocks(stack, parsed_body): from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER - mock_proxy_logging = stack.enter_context( - patch("litellm.proxy.proxy_server.proxy_logging_obj") - ) + mock_proxy_logging = stack.enter_context(patch("litellm.proxy.proxy_server.proxy_logging_obj")) mock_proxy_logging.pre_call_hook = AsyncMock(return_value=parsed_body) mock_proxy_logging.post_call_failure_hook = AsyncMock() mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) @@ -4347,11 +4038,7 @@ def _enter_relay_logging_mocks(stack, parsed_body): ) ) mock_success_handler.return_value = None - stack.enter_context( - patch.object( - GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue", new=MagicMock() - ) - ) + stack.enter_context(patch.object(GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue", new=MagicMock())) return mock_proxy_logging, mock_success_handler @@ -4437,10 +4124,7 @@ async def test_pass_through_request_relays_non_json_body_without_buffering(): mock_success_handler.assert_called_once() success_kwargs = mock_success_handler.call_args.kwargs assert success_kwargs["response_body"] is None - assert ( - success_kwargs["url_route"] - == "http://upstream.test/v1/messages/batches/b1/results" - ) + assert success_kwargs["url_route"] == "http://upstream.test/v1/messages/batches/b1/results" finally: cleanup() await fake_client.aclose() @@ -4518,9 +4202,7 @@ async def test_pass_through_request_upstream_error_body_stays_buffered(): ) try: with ExitStack() as stack: - mock_proxy_logging, mock_success_handler = _enter_relay_logging_mocks( - stack, {} - ) + mock_proxy_logging, mock_success_handler = _enter_relay_logging_mocks(stack, {}) response = await pass_through_request( request=_relay_client_request(), @@ -4590,18 +4272,11 @@ async def test_pass_through_relay_client_disconnect_logs_partial_relay_warning(c partial_relay_warnings = [ record.getMessage() for record in caplog.records - if record.levelno == logging.WARNING - and _PARTIAL_RELAY_WARNING_MARKER in record.getMessage() + if record.levelno == logging.WARNING and _PARTIAL_RELAY_WARNING_MARKER in record.getMessage() ] assert len(partial_relay_warnings) == 1 - assert ( - "http://upstream.test/v1/messages/batches/b1/results" - in partial_relay_warnings[0] - ) - assert ( - f"{len(first_chunk)} bytes were sent to the client" - in partial_relay_warnings[0] - ) + assert "http://upstream.test/v1/messages/batches/b1/results" in partial_relay_warnings[0] + assert f"{len(first_chunk)} bytes were sent to the client" in partial_relay_warnings[0] assert upstream_stream.closed is True mock_success_handler.assert_called_once() @@ -4648,10 +4323,7 @@ async def test_pass_through_relay_full_consumption_logs_no_partial_relay_warning relayed = [chunk async for chunk in response.body_iterator] assert b"".join(relayed) == b"".join(upstream_chunks) - assert not any( - _PARTIAL_RELAY_WARNING_MARKER in record.getMessage() - for record in caplog.records - ) + assert not any(_PARTIAL_RELAY_WARNING_MARKER in record.getMessage() for record in caplog.records) mock_success_handler.assert_called_once() finally: cleanup() @@ -4689,9 +4361,7 @@ def _enter_upstream_usage_mocks(stack, parsed_body): logging worker would have run so the test can await them.""" from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER - mock_proxy_logging = stack.enter_context( - patch("litellm.proxy.proxy_server.proxy_logging_obj") - ) + mock_proxy_logging = stack.enter_context(patch("litellm.proxy.proxy_server.proxy_logging_obj")) mock_proxy_logging.pre_call_hook = AsyncMock(return_value=parsed_body) mock_proxy_logging.post_call_failure_hook = AsyncMock() mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) @@ -4708,9 +4378,7 @@ def _enter_upstream_usage_mocks(stack, parsed_body): return mock_proxy_logging, enqueued -async def _run_upstream_reporting_passthrough( - upstream_headers, status_code=200, cost_per_request=None -): +async def _run_upstream_reporting_passthrough(upstream_headers, status_code=200, cost_per_request=None): """Drive a generic pass-through against an upstream that reports its own cost/usage. Returns (recorded standard logging payloads, proxy logging mock).""" from litellm.proxy._types import UserAPIKeyAuth @@ -4731,9 +4399,7 @@ async def _run_upstream_reporting_passthrough( request=_relay_client_request(method="POST"), target="http://internal-api.test/v1/summarize", custom_headers={}, - user_api_key_dict=UserAPIKeyAuth( - api_key="sk-upstream-usage", team_id="team-fil" - ), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-upstream-usage", team_id="team-fil"), cost_per_request=cost_per_request, ) for coroutine in enqueued: @@ -4802,23 +4468,17 @@ async def test_passthrough_records_upstream_reported_cost_on_error_response(): ) mock_proxy_logging.post_call_failure_hook.assert_awaited_once() - request_data = mock_proxy_logging.post_call_failure_hook.await_args.kwargs[ - "request_data" - ] + request_data = mock_proxy_logging.post_call_failure_hook.await_args.kwargs["request_data"] assert request_data["response_cost"] == 0.00021 assert request_data["combined_usage_object"] == litellm.Usage(total_tokens=930) @pytest.mark.asyncio async def test_passthrough_error_response_without_usage_headers_records_no_spend(): - _, mock_proxy_logging = await _run_upstream_reporting_passthrough( - {}, status_code=500 - ) + _, mock_proxy_logging = await _run_upstream_reporting_passthrough({}, status_code=500) mock_proxy_logging.post_call_failure_hook.assert_awaited_once() - request_data = mock_proxy_logging.post_call_failure_hook.await_args.kwargs[ - "request_data" - ] + request_data = mock_proxy_logging.post_call_failure_hook.await_args.kwargs["request_data"] assert "combined_usage_object" not in request_data @@ -4853,14 +4513,10 @@ async def test_streaming_passthrough_records_cost_and_tokens_reported_by_upstrea request=_relay_client_request(method="POST"), target="http://internal-api.test/v1/summarize", custom_headers={}, - user_api_key_dict=UserAPIKeyAuth( - api_key="sk-upstream-usage", team_id="team-fil" - ), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-upstream-usage", team_id="team-fil"), ) assert isinstance(response, StreamingResponse) - assert [chunk async for chunk in response.body_iterator] == [ - b'data: {"delta": "hi"}\n\n' - ] + assert [chunk async for chunk in response.body_iterator] == [b'data: {"delta": "hi"}\n\n'] for coroutine in enqueued: await coroutine finally: @@ -4968,9 +4624,7 @@ async def test_websocket_passthrough_forwards_non_ascii_first_frame(): "litellm.proxy.pass_through_endpoints.pass_through_endpoints.connect", return_value=FakeUpstreamConnect(upstream_ws), ), - patch( - "litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER" - ) as mock_worker, + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER") as mock_worker, ): mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) mock_proxy_logging.post_call_success_hook = AsyncMock() @@ -5052,9 +4706,7 @@ def _patched_websocket_passthrough_environment(upstream_ws): "litellm.proxy.pass_through_endpoints.pass_through_endpoints.connect", return_value=FakeUpstreamConnect(upstream_ws), ), - patch( - "litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER" - ) as mock_worker, + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER") as mock_worker, ): mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) mock_proxy_logging.post_call_success_hook = AsyncMock() @@ -5199,9 +4851,7 @@ async def test_websocket_passthrough_does_not_relay_unsendable_upstream_close(rc "abnormal": Close(1006, "connection died"), "no_status": Close(1005, ""), }[rcvd_close] - upstream_ws = ClosingUpstreamWebSocket( - ConnectionClosedError(rcvd=rcvd, sent=None, rcvd_then_sent=None) - ) + upstream_ws = ClosingUpstreamWebSocket(ConnectionClosedError(rcvd=rcvd, sent=None, rcvd_then_sent=None)) websocket = _client_websocket(_pending_receive) with _patched_websocket_passthrough_environment(upstream_ws): @@ -5276,14 +4926,15 @@ async def test_websocket_passthrough_does_not_close_twice_when_success_logging_f def _passthrough_kwargs_for_reservation( - user_api_key_dict: UserAPIKeyAuth, parsed_body: Optional[dict] = None + user_api_key_dict: UserAPIKeyAuth, + parsed_body: Optional[dict] = None, + user_defined_route: bool = False, ) -> dict: mock_request = MagicMock(spec=Request) mock_request.method = "POST" - mock_request.url = ( - "http://0.0.0.0:4000/gemini/v1beta/models/gemini-2.5-flash:generateContent" - ) + mock_request.url = "http://0.0.0.0:4000/gemini/v1beta/models/gemini-2.5-flash:generateContent" mock_request.headers = Headers({}) + mock_request.scope = {"endpoint": _marked_pass_through_endpoint()} if user_defined_route else {} return HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint( request=mock_request, @@ -5352,10 +5003,7 @@ async def test_passthrough_success_reconciles_budget_reservation(): reservation = user_api_key_dict.budget_reservation kwargs = _passthrough_kwargs_for_reservation(user_api_key_dict) - assert ( - kwargs["litellm_params"]["metadata"]["user_api_key_budget_reservation"] - is reservation - ) + assert kwargs["litellm_params"]["metadata"]["user_api_key_budget_reservation"] is reservation increment_spend_counters = await _track_cost_for_passthrough_kwargs(kwargs) @@ -5381,9 +5029,7 @@ async def test_passthrough_body_cannot_forge_budget_reservation(): user_api_key_dict, parsed_body={"litellm_metadata": {"user_api_key_budget_reservation": forged}}, ) - assert ( - kwargs["litellm_params"]["metadata"]["user_api_key_budget_reservation"] is None - ) + assert kwargs["litellm_params"]["metadata"]["user_api_key_budget_reservation"] is None increment_spend_counters = await _track_cost_for_passthrough_kwargs(kwargs) @@ -5391,9 +5037,7 @@ async def test_passthrough_body_cannot_forge_budget_reservation(): assert increment_spend_counters.await_args.kwargs["budget_reservation"] is None -async def _drive_streaming_pass_through( - upstream_content_type, chunk_delay_seconds, client_asked_for_stream=True -): +async def _drive_streaming_pass_through(upstream_content_type, chunk_delay_seconds, client_asked_for_stream=True): """Drive pass_through_request against an upstream that stalls before its first byte. ``client_asked_for_stream`` picks which of pass_through_request's two streaming @@ -5405,22 +5049,14 @@ async def _drive_streaming_pass_through( ) with ExitStack() as stack: - mock_proxy_logging = stack.enter_context( - patch("litellm.proxy.proxy_server.proxy_logging_obj") - ) + mock_proxy_logging = stack.enter_context(patch("litellm.proxy.proxy_server.proxy_logging_obj")) mock_get_client = stack.enter_context( - patch( - "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client" - ) - ) - mock_chunk_processor = stack.enter_context( - patch.object(PassThroughStreamingHandler, "chunk_processor") + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client") ) + mock_chunk_processor = stack.enter_context(patch.object(PassThroughStreamingHandler, "chunk_processor")) mock_proxy_logging.pre_call_hook = AsyncMock( - return_value={"model": "claude-3", "stream": True} - if client_asked_for_stream - else {"model": "claude-3"} + return_value={"model": "claude-3", "stream": True} if client_asked_for_stream else {"model": "claude-3"} ) mock_proxy_logging.post_call_failure_hook = AsyncMock() mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={}) @@ -5510,9 +5146,7 @@ async def test_pass_through_binary_event_stream_is_never_given_an_sse_comment(): @pytest.mark.asyncio @pytest.mark.parametrize("configured_interval, expect_ping", [(0.05, True), (None, False)]) -async def test_pass_through_route_pings_while_the_upstream_call_is_still_running( - configured_interval, expect_ping -): +async def test_pass_through_route_pings_while_the_upstream_call_is_still_running(configured_interval, expect_ping): """The upstream withholds its response headers until its first token, so the whole time-to-first-token is spent inside pass_through_request with nothing on the wire (issue #34819).""" @@ -5542,9 +5176,7 @@ async def test_pass_through_route_pings_while_the_upstream_call_is_still_running ) ) stack.enter_context(patch(f"{module}.pass_through_request", slow_pass_through)) - stack.enter_context( - patch.object(litellm, "sse_keepalive_ping_interval_seconds", configured_interval) - ) + stack.enter_context(patch.object(litellm, "sse_keepalive_ping_interval_seconds", configured_interval)) endpoint_func = create_pass_through_route( endpoint="/v1/messages", @@ -5571,3 +5203,147 @@ async def test_pass_through_route_pings_while_the_upstream_call_is_still_running assert (collected[0] == b": ping\n\n") is expect_ping assert collected[-1] in (MESSAGE_START_SSE_FRAME, MESSAGE_START_SSE_FRAME.decode()) + + +def test_passthrough_carries_the_per_model_budgets(): + """ + Native passthrough builds its logging metadata from + StandardLoggingUserAPIKeyMetadata, which has no budget field, and never calls + add_litellm_data_to_request. Without these three keys the post-call increment + exits early, so a /bedrock/... request is costed but its per-model counter is + never written: the budget reports zero forever and enforces nothing. + """ + key_budget = {"claude-opus-4-8": {"budget_limit": 1.0, "time_period": "18h"}} + user_budget = {"claude-opus-4-8": {"budget_limit": 2.0, "time_period": "1mo"}} + end_user_budget = {"claude-opus-4-8": {"budget_limit": 3.0, "time_period": "1d"}} + + kwargs = _passthrough_kwargs_for_reservation( + UserAPIKeyAuth( + token="hash", + user_id="u-1", + model_max_budget=key_budget, + user_model_max_budget=user_budget, + end_user_model_max_budget=end_user_budget, + ) + ) + + metadata = kwargs["litellm_params"]["metadata"] + assert metadata["user_api_key_model_max_budget"] == key_budget + assert metadata["user_api_key_user_model_max_budget"] == user_budget + assert metadata["user_api_key_end_user_model_max_budget"] == end_user_budget + + +def test_passthrough_budget_metadata_cannot_be_forged_by_the_request_body(): + """ + These keys decide budget enforcement, so a caller-supplied body must not be + able to raise its own cap. They are set after the client metadata merge for + the same reason user_api_key and the parent span are. + """ + key_budget = {"claude-opus-4-8": {"budget_limit": 1.0, "time_period": "18h"}} + + kwargs = _passthrough_kwargs_for_reservation( + UserAPIKeyAuth(token="hash", user_id="u-1", model_max_budget=key_budget), + parsed_body={ + "litellm_metadata": { + "user_api_key_model_max_budget": {"claude-opus-4-8": {"budget_limit": 999999.0, "time_period": "18h"}} + } + }, + ) + + metadata = kwargs["litellm_params"]["metadata"] + assert metadata["user_api_key_model_max_budget"] == key_budget + + +def _marked_pass_through_endpoint(): + """An endpoint carrying the marker ``create_pass_through_route`` sets.""" + from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + LITELLM_PASS_THROUGH_ENDPOINT_MARKER, + ) + + def _endpoint(): # pragma: no cover - identity only + return None + + setattr(_endpoint, LITELLM_PASS_THROUGH_ENDPOINT_MARKER, True) # noqa: B010 # name is a module constant + return _endpoint + + +def test_user_defined_passthrough_is_neither_tracked_nor_enforced(): + """ + `get_model_from_request` returns None for a user-defined pass-through on + purpose: the body is forwarded verbatim, so its `model` names an UPSTREAM + model rather than a LiteLLM-managed one, and enforcing key/team allowlists + against it would reject valid requests. Enforcement is therefore skipped + on those routes. + + Attaching the budget metadata anyway would charge a counter that nothing on + that route can refuse, and would attribute the spend to a budget the operator + scoped to a LiteLLM model that merely shares the name. Tracking and + enforcement have to agree: both on for the built-in provider routes, both off + here. + """ + kwargs = _passthrough_kwargs_for_reservation( + UserAPIKeyAuth( + token="hash", + user_id="u-1", + model_max_budget={"claude-opus-4-8": {"budget_limit": 1.0, "time_period": "18h"}}, + ), + user_defined_route=True, + ) + + metadata = kwargs["litellm_params"]["metadata"] + for field in ( + "user_api_key_model_max_budget", + "user_api_key_user_model_max_budget", + "user_api_key_end_user_model_max_budget", + ): + assert field not in metadata, f"{field} was attached on a route that never enforces it" + + +@pytest.mark.parametrize( + "handler_name", + [ + "anthropic_proxy_route", + "bedrock_proxy_route", + "gemini_proxy_route", + "cohere_proxy_route", + "vllm_proxy_route", + "mistral_proxy_route", + ], +) +def test_builtin_provider_routes_do_not_carry_the_user_defined_marker(handler_name): + """ + The budget metadata is attached only when the dispatched endpoint is NOT a + user-defined pass-through, so the built-in provider handlers must not carry + that marker or native provider spend would stop being tracked and enforced. + + These handlers DO call `create_pass_through_route` internally, and that + factory sets the marker on what it returns. But the result is awaited + immediately rather than registered, so FastAPI puts the decorated handler in + `request.scope["endpoint"]`, and that is what the marker check reads. This + test pins the distinction between calling the factory and being dispatched as + its product, which is easy to misread from a grep alone. + """ + from litellm.proxy.pass_through_endpoints import llm_passthrough_endpoints + from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + LITELLM_PASS_THROUGH_ENDPOINT_MARKER, + ) + + handler = getattr(llm_passthrough_endpoints, handler_name) + assert getattr(handler, LITELLM_PASS_THROUGH_ENDPOINT_MARKER, False) is False, ( + f"{handler_name} is marked as a user-defined pass-through, so per-model budget " + "metadata would be skipped and native provider spend would go untracked" + ) + + +def test_the_marker_check_distinguishes_the_two_route_kinds(): + """Positive control: the factory's product IS marked, so the check can discriminate.""" + from litellm.proxy.auth.auth_utils import request_dispatched_to_pass_through_endpoint + from litellm.proxy.pass_through_endpoints import llm_passthrough_endpoints + + marked = MagicMock(spec=Request) + marked.scope = {"endpoint": _marked_pass_through_endpoint()} + assert request_dispatched_to_pass_through_endpoint(marked) is True + + builtin = MagicMock(spec=Request) + builtin.scope = {"endpoint": llm_passthrough_endpoints.anthropic_proxy_route} + assert request_dispatched_to_pass_through_endpoint(builtin) is False diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/BulkEditUsers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/BulkEditUsers.tsx index b3326df1ff8..459e3fd8c92 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/BulkEditUsers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/BulkEditUsers.tsx @@ -10,6 +10,7 @@ import { Checkbox } from "@/components/ui/checkbox"; import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; import { Separator } from "@/components/ui/separator"; import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; interface BulkEditUserModalProps { open: boolean; @@ -36,6 +37,7 @@ const BulkEditUserModal: React.FC = ({ userModels, allowAllUsers = false, }) => { + const { premiumUser } = useAuthorized(); const [loading, setLoading] = useState(false); const [selectedTeams, setSelectedTeams] = useState([]); const [teamBudget, setTeamBudget] = useState(null); @@ -362,6 +364,7 @@ const BulkEditUserModal: React.FC = ({ userModels={userModels} possibleUIRoles={possibleUIRoles} isBulkEdit={true} + premiumUser={premiumUser === true} /> {loading && ( diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.test.tsx index 0ca68cc665e..8fb94ce477e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.test.tsx @@ -1,4 +1,4 @@ -import { cleanup, screen, waitFor } from "@testing-library/react"; +import { cleanup, fireEvent, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders } from "../../../../../tests/test-utils"; @@ -612,6 +612,125 @@ describe("UserEditView", () => { expect(onSubmit).not.toHaveBeenCalled(); }); + // /user/new validates model_max_budget behind an enterprise license, so a + // form that re-sends what is already stored turns an unrelated edit into a + // 400 on a proxy without one. + describe("per-model budgets", () => { + const withStoredBudgets = { + ...MOCK_USER_DATA, + user_info: { + ...MOCK_USER_DATA.user_info, + model_max_budget: { "gpt-4": { budget_limit: 5, time_period: "30d" } }, + }, + }; + + it("should leave model_max_budget out of an edit that did not touch it", async () => { + const payload = await submittedPayload({ userData: withStoredBudgets, premiumUser: true }); + + expect(payload).not.toHaveProperty("model_max_budget"); + }); + + // The proxy stores model_max_budget as a plain dict, exactly as the client + // sent it, and BudgetConfig documents the max_budget/budget_duration + // spelling. A row hydrated from the spelling the editor does not read mounts + // with an empty cap, and every edit re-emits ALL rows, so touching one + // model's budget silently deletes another's. + it("should keep a row stored under the BudgetConfig aliases when a sibling row is edited", async () => { + const onSubmit = vi.fn(); + renderWithProviders( + , + ); + + const [aliasRow, canonicalRow] = await screen.findAllByPlaceholderText("Max spend ($)"); + expect(aliasRow).toHaveValue(5); + + fireEvent.change(canonicalRow, { target: { value: "3" } }); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmit).toHaveBeenCalled(); + }); + expect(onSubmit.mock.calls[0][0].model_max_budget).toEqual({ + "gpt-4": { budget_limit: 5, time_period: "30d" }, + "gpt-3.5-turbo": { budget_limit: 3, time_period: "1h" }, + }); + }); + + // The effect already re-seeds the form on a userData change, so that change + // does happen while this component stays mounted. The editor holds its rows + // in state seeded once, so without a matching re-seed the rows on screen + // keep describing the previously loaded user and a save overwrites theirs. + it("re-seeds the editor when a different user is loaded", async () => { + const withBudget = (limit: number, id: string) => ({ + ...MOCK_USER_DATA, + user_id: id, + user_info: { + ...MOCK_USER_DATA.user_info, + model_max_budget: { "gpt-4": { budget_limit: limit, time_period: "1h" } }, + }, + }); + + const { rerender } = renderWithProviders( + , + ); + expect(await screen.findByPlaceholderText("Max spend ($)")).toHaveValue(5); + + rerender(); + + expect(await screen.findByPlaceholderText("Max spend ($)")).toHaveValue(99); + }); + + // BulkEditUsers copies a fixed field list into its payload and never reads + // model_max_budget, so an editor rendered here would take input and throw + // it away. It also has no single stored budget to diff against, since its + // userData stands in for every selected user. + it("does not offer the editor in bulk edit, where the value would be discarded", async () => { + renderWithProviders( + , + ); + + await screen.findByRole("button", { name: /save changes/i }); + expect(screen.queryByPlaceholderText("Max spend ($)")).not.toBeInTheDocument(); + }); + + it("should lock the editor when the proxy has no enterprise license", async () => { + renderWithProviders(); + + expect(await screen.findByPlaceholderText("Max spend ($)")).toBeDisabled(); + }); + + it("should leave the editor usable when the proxy has one", async () => { + renderWithProviders(); + + expect(await screen.findByPlaceholderText("Max spend ($)")).toBeEnabled(); + }); + }); + it("should send an empty-string metadata through untouched rather than as an object", async () => { const onSubmit = vi.fn(); renderWithProviders(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx index 5612f5cd5f6..ed0c08adf38 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx @@ -2,6 +2,9 @@ import React, { useMemo, useState } from "react"; import { z } from "zod/v4"; import { all_admin_roles } from "@/utils/roles"; import BudgetDurationDropdown from "@/components/common_components/budget_duration_dropdown"; +import { ModelMaxBudget, ModelMaxBudgetField } from "@/components/key_team_helpers/ModelMaxBudgetEditor"; +import { modelMaxBudgetUpdate } from "@/components/key_team_helpers/modelMaxBudgetPayload"; +import { useSeededState } from "@/components/key_team_helpers/useSeededState"; import { getModelDisplayName } from "@/components/key_team_helpers/fetch_available_models_team_key"; import MCPServerSelector from "@/components/mcp_server_management/MCPServerSelector"; import MCPToolPermissions from "@/components/mcp_server_management/MCPToolPermissions"; @@ -30,6 +33,7 @@ interface UserEditViewProps { possibleUIRoles: Record> | null; isBulkEdit?: boolean; objectPermission?: ObjectPermission | null; + premiumUser?: boolean; } const MCP_SELECTION_SHAPE = z.object({ @@ -135,9 +139,14 @@ export function UserEditView({ possibleUIRoles, isBulkEdit = false, objectPermission, + premiumUser = false, }: UserEditViewProps) { const canEditMcpPermissions = !isBulkEdit && all_admin_roles.includes(userRole || ""); const [unlimitedBudget, setUnlimitedBudget] = useState(false); + const [modelMaxBudget, setModelMaxBudget] = useSeededState( + userData.user_id, + () => userData.user_info?.model_max_budget ?? {}, + ); const schema = useMemo(() => budgetSchema(unlimitedBudget), [unlimitedBudget]); const form = useZodForm(schema, { defaultValues: toFormValues(userData, objectPermission, isBulkEdit, canEditMcpPermissions), @@ -162,9 +171,11 @@ export function UserEditView({ return; } + const modelBudgets = modelMaxBudgetUpdate(modelMaxBudget, userData.user_info?.model_max_budget); onSubmit({ ...values, ...("metadata" in values ? { metadata: metadata.value } : {}), + ...(modelBudgets !== undefined && { model_max_budget: modelBudgets }), max_budget: unlimitedBudget || values.max_budget === "" || values.max_budget === undefined ? null : values.max_budget, }); @@ -282,6 +293,20 @@ export function UserEditView({ {({ id, value, onChange }) => } + {/* Bulk edit forwards a fixed field list and has no single stored budget to + diff against, so the editor would silently discard whatever was typed. */} + {!isBulkEdit && ( + + )} + {({ ref, value, ...control }) => (