From 7bd342d7eadba4e9ebbc393c44d36719bb29c7f6 Mon Sep 17 00:00:00 2001 From: paveln83 <349044532+paveln83@users.noreply.github.com> Date: Tue, 3 Nov 2026 01:56:08 -0800 Subject: [PATCH] fix(proxy): let a config license unlock enterprise settings and spell "environment" right in license errors (#49745) * fix(proxy): spell "environment" right in the allowed_ips and AWS KMS V2 license errors * fix(proxy): read general_settings.litellm_license before the allowed_ips and enforced_params gates * fix(proxy): read the config license before callbacks and gates, keeping a valid env license * fix(proxy): keep an invalid config license reported and sync license data with the license in use A config license that fails no longer overwrites the license string when the env license is not valid either, so /health/license still reports has_license. A failed boot-time env license is not re-checked. Both license writes go through one helper that also refreshes premium_user_data, so /health/license and the seat limits follow the license in use. * fix(proxy): keep the env license's boot result when a config license fails, without a second license check --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 38 +++- litellm/secret_managers/aws_secret_manager.py | 2 +- .../proxy/proxy_server/test_proxy_config.py | 212 ++++++++++++++++++ ...test_proxy_server_endpoints_and_startup.py | 6 +- .../test_aws_secret_manager.py | 14 ++ 5 files changed, 261 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 424e68ae3bc..6fcf789dc1b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1185,6 +1185,27 @@ server_root_path: Final = get_server_root_path() _license_check = LicenseCheck() premium_user: bool = _license_check.is_premium() premium_user_data: Optional["EnterpriseLicenseData"] = _license_check.airgapped_license_data + + +def _set_premium_user(is_premium: bool) -> None: + global premium_user, premium_user_data + premium_user = is_premium # rebind-ok: config load can switch the license in use + premium_user_data = _license_check.airgapped_license_data # rebind-ok: tracks the license in use + + +def _use_config_license(config_license: str, env_license_is_premium: bool) -> bool: + env_license: Final = _license_check.license_str + env_license_data: Final = _license_check.airgapped_license_data + _license_check.license_str = config_license + if _license_check.is_premium(): + return True + if not env_license_is_premium: + return False + _license_check.license_str = env_license + _license_check.airgapped_license_data = env_license_data + return True + + global_max_parallel_request_retries_env: Final[str | None] = os.getenv("LITELLM_GLOBAL_MAX_PARALLEL_REQUEST_RETRIES") proxy_state: Final = ProxyState() SENSITIVE_DATA_MASKER: Final = SensitiveDataMasker() @@ -6686,7 +6707,6 @@ class ProxyConfig: def _load_environment_variables(self, config: dict): ## ENVIRONMENT VARIABLES - global premium_user environment_variables: Final = config.get("environment_variables", None) if environment_variables: for key, value in environment_variables.items(): @@ -6726,7 +6746,7 @@ class ProxyConfig: # check if litellm_license in general_settings if "LITELLM_LICENSE" in environment_variables: _license_check.license_str = os.getenv("LITELLM_LICENSE", None) - premium_user = _license_check.is_premium() + _set_premium_user(_license_check.is_premium()) def _set_config_environment_variable(self, key: str, value: str) -> None: if key in os.environ and os.environ[key] != value: @@ -6778,7 +6798,6 @@ class ProxyConfig: prompt_injection_detection_obj, \ redis_usage_cache, \ store_model_in_db, \ - premium_user, \ open_telemetry_logger, \ health_check_details, \ proxy_batch_polling_interval, \ @@ -6791,6 +6810,12 @@ class ProxyConfig: self._load_environment_variables(config=config) + config_license: Final = _GENERAL_SETTINGS_VIEW.validate_python( + _GENERAL_SETTINGS_VIEW.validate_python(config).get("general_settings") or {} + ).get("litellm_license") + if isinstance(config_license, str): + _set_premium_user(_use_config_license(config_license, env_license_is_premium=premium_user)) + ## Coordination Redis (before cache init, so the explicit block wins) coordination_redis_cache: Final = self._init_coordination_redis(config=config) if coordination_redis_cache is not None: @@ -7367,7 +7392,7 @@ class ProxyConfig: allowed_ips: Final = general_settings.get("allowed_ips", None) if allowed_ips is not None and premium_user is False: raise ValueError( - "allowed_ips is an Enterprise Feature. Please add a valid LITELLM_LICENSE to your envionment." + "allowed_ips is an Enterprise Feature. Please add a valid LITELLM_LICENSE to your environment." ) ## BUDGET RESCHEDULER ## proxy_budget_rescheduler_min_time = general_settings.get( @@ -7423,11 +7448,6 @@ class ProxyConfig: if general_settings.get("enforced_params") is not None and premium_user is not True: raise ValueError("Trying to use `enforced_params`" + CommonProxyErrors.not_premium_user.value) - # check if litellm_license in general_settings - if "litellm_license" in general_settings: - _license_check.license_str = general_settings["litellm_license"] - premium_user = _license_check.is_premium() - router_params: Final[dict] = { "cache_responses": litellm.cache is not None, # cache if user passed in cache values } diff --git a/litellm/secret_managers/aws_secret_manager.py b/litellm/secret_managers/aws_secret_manager.py index 88642ae141e..c87edc4b1dd 100644 --- a/litellm/secret_managers/aws_secret_manager.py +++ b/litellm/secret_managers/aws_secret_manager.py @@ -70,7 +70,7 @@ class AWSKeyManagementService_V2: is_litellm_license_in_env = True if is_litellm_license_in_env is False: raise ValueError( - "AWSKeyManagementService V2 is an Enterprise Feature. Please add a valid LITELLM_LICENSE to your envionment." + "AWSKeyManagementService V2 is an Enterprise Feature. Please add a valid LITELLM_LICENSE to your environment." ) def load_aws_kms(self, use_aws_kms: bool | None): diff --git a/tests/unit/proxy/proxy_server/test_proxy_config.py b/tests/unit/proxy/proxy_server/test_proxy_config.py index 55226fc629b..712b4134c05 100644 --- a/tests/unit/proxy/proxy_server/test_proxy_config.py +++ b/tests/unit/proxy/proxy_server/test_proxy_config.py @@ -9,6 +9,7 @@ Pins covered: from __future__ import annotations import asyncio +import base64 import json import logging import os @@ -25,11 +26,15 @@ from unittest.mock import AsyncMock, MagicMock import httpx import pytest +from cryptography.hazmat.primitives import hashes +from cryptography.hazmat.primitives.asymmetric import padding, rsa +from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey, RSAPublicKey from prisma.errors import DataError from pydantic import JsonValue, TypeAdapter, ValidationError import litellm from litellm._logging import turn_on_json +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.proxy._types import CommonProxyErrors, ConfigGeneralSettings from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper from litellm.proxy.config_resolvers.settings_store import SettingsStore @@ -2671,6 +2676,213 @@ async def test_load_config_logs_disabled_budget_reservation_once(tmp_path, monke assert [record.levelno for record in records] == ([logging.INFO] if setting == "true" else []) +@pytest.mark.asyncio +async def test_load_config_refuses_allowed_ips_without_a_license(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + config_file: Final = tmp_path / "allowed_ips.yaml" + config_file.write_text('model_list: []\nlitellm_settings: {}\ngeneral_settings:\n allowed_ips: ["127.0.0.1"]\n') + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + + expected_message: Final = ( + "allowed_ips is an Enterprise Feature. Please add a valid LITELLM_LICENSE to your environment." + ) + with pytest.raises(ValueError, match=f"^{re.escape(expected_message)}$"): + await ProxyConfig().load_config(router=None, config_file_path=str(config_file)) + + +def _signed_test_license(private_key: RSAPrivateKey, max_users: int) -> str: + message: Final = json.dumps({"expiration_date": "9999-12-31", "user_id": "test", "max_users": max_users}).encode() + signature: Final = private_key.sign( + message, + padding.PSS(mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.MAX_LENGTH), + hashes.SHA256(), + ) + return base64.b64encode(message + b"." + signature).decode() + + +def _write_config(path: Path, general_settings: dict[str, object], litellm_settings: dict[str, object]) -> str: + path.write_text( + json.dumps({"model_list": [], "litellm_settings": litellm_settings, "general_settings": general_settings}) + ) + return str(path) + + +def _boot_proxy( + monkeypatch: pytest.MonkeyPatch, + public_key: RSAPublicKey, + env_license: str | None, + license_server_urls: list[str], + server_accepted_license: str | None = None, +) -> None: + from litellm.proxy import proxy_server + + def _license_server(request: httpx.Request) -> httpx.Response: + license_server_urls.append(str(request.url)) + return httpx.Response( + 200, json={"verify": server_accepted_license is not None and server_accepted_license in str(request.url)} + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.setattr("litellm.proxy.proxy_server._license_check.public_key", public_key) + monkeypatch.setattr("litellm.proxy.proxy_server._license_check.airgapped_license_data", None) + monkeypatch.setattr( + "litellm.proxy.proxy_server._license_check.http_handler", + HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_license_server))), + ) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + if env_license is None: + monkeypatch.delenv("LITELLM_LICENSE", raising=False) + else: + monkeypatch.setenv("LITELLM_LICENSE", env_license) + monkeypatch.setattr("litellm.proxy.proxy_server._license_check.license_str", env_license) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", proxy_server._license_check.is_premium()) + monkeypatch.setattr( + "litellm.proxy.proxy_server.premium_user_data", proxy_server._license_check.airgapped_license_data + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "gated_config", + [ + {"general_settings": {"allowed_ips": ["127.0.0.1"]}}, + {"general_settings": {"enforced_params": ["user"]}}, + {"litellm_settings": {"callbacks": ["openai_moderations"]}}, + ], +) +async def test_load_config_lets_a_config_license_unlock_its_enterprise_gates( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, gated_config: dict[str, dict[str, list[str]]] +): + from litellm.proxy import proxy_server + + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + license_key: Final = _signed_test_license(private_key, max_users=10) + config_file: Final = _write_config( + tmp_path / "licensed.yaml", + {"litellm_license": license_key, **gated_config.get("general_settings", {})}, + dict(gated_config.get("litellm_settings", {})), + ) + _boot_proxy(monkeypatch, private_key.public_key(), env_license=None, license_server_urls=[]) + + await ProxyConfig().load_config(router=None, config_file_path=config_file) + + assert proxy_server.premium_user is True + + +@pytest.mark.asyncio +async def test_load_config_keeps_a_valid_env_license_when_the_config_license_is_invalid( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy import proxy_server + + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + config_file: Final = _write_config( + tmp_path / "stale_license.yaml", {"litellm_license": "not-a-valid-license", "allowed_ips": ["127.0.0.1"]}, {} + ) + _boot_proxy( + monkeypatch, + private_key.public_key(), + env_license=_signed_test_license(private_key, max_users=5), + license_server_urls=[], + ) + + await ProxyConfig().load_config(router=None, config_file_path=config_file) + + assert proxy_server.premium_user is True + assert proxy_server._license_check.user_seat_limit() == 5 + + +@pytest.mark.asyncio +async def test_load_config_prefers_the_config_license_when_both_licenses_are_valid( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_license_endpoint + + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + config_file: Final = _write_config( + tmp_path / "licensed.yaml", {"litellm_license": _signed_test_license(private_key, max_users=100)}, {} + ) + _boot_proxy( + monkeypatch, + private_key.public_key(), + env_license=_signed_test_license(private_key, max_users=5), + license_server_urls=[], + ) + + await ProxyConfig().load_config(router=None, config_file_path=config_file) + + assert proxy_server._license_check.user_seat_limit() == 100 + assert (await health_license_endpoint(user_api_key_dict=UserAPIKeyAuth()))["limits"]["max_users"] == 100 + + +@pytest.mark.asyncio +async def test_load_config_still_reports_an_invalid_config_license_as_configured( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_license_endpoint + + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + config_file: Final = _write_config(tmp_path / "stale_license.yaml", {"litellm_license": "not-a-valid-license"}, {}) + _boot_proxy(monkeypatch, private_key.public_key(), env_license=None, license_server_urls=[]) + + await ProxyConfig().load_config(router=None, config_file_path=config_file) + + license_status: Final = await health_license_endpoint(user_api_key_dict=UserAPIKeyAuth()) + assert (license_status["has_license"], license_status["license_type"]) == (True, "community") + + +@pytest.mark.asyncio +async def test_load_config_does_not_recheck_an_env_license_that_failed_at_boot( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy import proxy_server + + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + config_file: Final = _write_config(tmp_path / "stale_license.yaml", {"litellm_license": "not-a-valid-license"}, {}) + license_server_urls: Final[list[str]] = [] + _boot_proxy(monkeypatch, private_key.public_key(), "also-not-a-valid-license", license_server_urls) + asked_at_boot: Final = len(license_server_urls) + + await ProxyConfig().load_config(router=None, config_file_path=config_file) + + assert proxy_server.premium_user is False + assert [url for url in license_server_urls[asked_at_boot:] if "also-not-a-valid-license" in url] == [] + assert any("not-a-valid-license" in url for url in license_server_urls[asked_at_boot:]) + + +@pytest.mark.asyncio +async def test_load_config_keeps_the_boot_verdict_of_a_server_checked_env_license( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy import proxy_server + + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + config_file: Final = _write_config( + tmp_path / "stale_license.yaml", {"litellm_license": "not-a-valid-license", "allowed_ips": ["127.0.0.1"]}, {} + ) + license_server_urls: Final[list[str]] = [] + _boot_proxy( + monkeypatch, + private_key.public_key(), + "server-checked-license", + license_server_urls, + server_accepted_license="server-checked-license", + ) + asked_at_boot: Final = len(license_server_urls) + + await ProxyConfig().load_config(router=None, config_file_path=config_file) + + assert proxy_server.premium_user is True + assert proxy_server._license_check.license_str == "server-checked-license" + assert [url for url in license_server_urls[asked_at_boot:] if "server-checked-license" in url] == [] + + @pytest.mark.asyncio @pytest.mark.parametrize( "setting", ["vector_store_deny_by_default", "search_tool_deny_by_default", "claude_code_marketplace_auth_required"] diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 82db8820447..9914247bc47 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -4117,7 +4117,11 @@ async def test_load_environment_variables_litellm_license_and_edge_cases(): mock_license_check = MagicMock() mock_license_check.is_premium.return_value = True - with patch("litellm.proxy.proxy_server._license_check", mock_license_check): + with ( + patch("litellm.proxy.proxy_server._license_check", mock_license_check), + patch("litellm.proxy.proxy_server.premium_user", False), + patch("litellm.proxy.proxy_server.premium_user_data", None), + ): with patch.dict(os.environ, {}, clear=False): # Call the method under test proxy_config._load_environment_variables(test_config_with_license) diff --git a/tests/unit/secret_managers/test_aws_secret_manager.py b/tests/unit/secret_managers/test_aws_secret_manager.py index d9922b43ec8..e16da7a0844 100644 --- a/tests/unit/secret_managers/test_aws_secret_manager.py +++ b/tests/unit/secret_managers/test_aws_secret_manager.py @@ -1,4 +1,6 @@ import base64 +import re +from typing import Final import pytest from botocore.stub import Stubber @@ -34,3 +36,15 @@ def test_decrypt_value_turns_only_boolean_literals_into_bools(monkeypatch, plain assert decrypted == expected assert type(decrypted) is type(expected) + + +def test_kms_v2_without_a_license_tells_the_user_to_set_one(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("AWS_REGION_NAME", "us-west-2") + monkeypatch.delenv("LITELLM_LICENSE", raising=False) + monkeypatch.delenv("LITELLM_SECRET_AWS_KMS_LITELLM_LICENSE", raising=False) + + expected_message: Final = ( + "AWSKeyManagementService V2 is an Enterprise Feature. Please add a valid LITELLM_LICENSE to your environment." + ) + with pytest.raises(ValueError, match=f"^{re.escape(expected_message)}$"): + AWSKeyManagementService_V2()