From 6daddbcdab4f27907a5d54f03c636b075b58f679 Mon Sep 17 00:00:00 2001 From: Andrew Bernat <31580217+bernata@users.noreply.github.com> Date: Wed, 24 Dec 2025 21:41:17 -0800 Subject: [PATCH] add license endpoint (#18311) * Add license health endpoint * Address review feedback for license endpoint * Refactor license health endpoint helpers * Remove unused timezone import from health_endpoints.py --------- Co-authored-by: Plan42.ai --- .../health_endpoints/_health_endpoints.py | 88 ++++++++++++++++++- .../health_endpoints/test_health_endpoints.py | 70 ++++++++++++++- 2 files changed, 155 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 65de1bd7393..d27e0036235 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -4,7 +4,7 @@ import os import time import traceback from datetime import datetime, timedelta -from typing import Dict, Literal, Optional, Union +from typing import Any, Dict, Literal, Optional, Union, cast import fastapi from fastapi import APIRouter, Depends, HTTPException, Request, Response, status @@ -16,6 +16,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import ( AlertType, CallInfo, + EnterpriseLicenseData, Litellm_EntityType, ProxyErrorTypes, ProxyException, @@ -960,6 +961,91 @@ async def shared_health_check_status_endpoint( ) +def _read_license_data() -> Optional[Dict[str, Any]]: + from litellm.proxy.proxy_server import ( + _license_check, + premium_user_data, + ) + + license_data: Optional[EnterpriseLicenseData] = ( + premium_user_data or _license_check.airgapped_license_data + ) + + if ( + license_data is None + and getattr(_license_check, "license_str", None) + and getattr(_license_check, "public_key", None) + ): + try: + verification_result = _license_check.verify_license_without_api_request( + public_key=_license_check.public_key, + license_key=_license_check.license_str, + ) + if verification_result is True: + license_data = _license_check.airgapped_license_data + except Exception: + pass + + if license_data is None: + return None + return cast(Dict[str, Any], license_data) + + +def _read_allowed_features(license_data: Dict[str, Any]) -> list: + raw_allowed_features = license_data.get("allowed_features") + if isinstance(raw_allowed_features, list): + return list(raw_allowed_features) + if raw_allowed_features is None: + return [] + return [raw_allowed_features] + + +@router.get( + "/health/license", + tags=["health"], + dependencies=[Depends(user_api_key_auth)], +) +async def health_license_endpoint( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """Return metadata about the configured LiteLLM license without exposing the key.""" + from litellm.proxy.proxy_server import ( + _license_check, + premium_user, + ) + + license_data = _read_license_data() + has_license = bool(getattr(_license_check, "license_str", None)) + license_type = "enterprise" if premium_user else "community" + + if license_data is None: + return { + "has_license": has_license, + "license_type": license_type, + "expiration_date": None, + "allowed_features": [], + "limits": { + "max_users": None, + "max_teams": None, + }, + } + + expiration_date = license_data.get("expiration_date") + max_users = license_data.get("max_users") + max_teams = license_data.get("max_teams") + + return { + "has_license": has_license, + "license_type": license_type, + "expiration_date": expiration_date, + "allowed_features": _read_allowed_features(license_data), + "limits": { + "max_users": max_users, + "max_teams": max_teams, + }, + } + + db_health_cache = {"status": "unknown", "last_updated": datetime.now()} diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index 23b3b0287ee..edfdd9e4065 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -2,7 +2,8 @@ import os import sys import time from datetime import datetime, timedelta -from unittest.mock import MagicMock, patch, AsyncMock +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch sys.path.insert( 0, os.path.abspath("../../..") @@ -10,10 +11,14 @@ sys.path.insert( import pytest from prisma.errors import ClientNotConnectedError, HTTPClientClosedError, PrismaError + from litellm.proxy.health_endpoints._health_endpoints import ( _db_health_readiness_check, db_health_cache, + health_license_endpoint, health_services_endpoint, +) +from litellm.proxy.health_endpoints._health_endpoints import ( test_model_connection as health_test_model_connection, ) @@ -128,6 +133,68 @@ async def test_health_services_endpoint_sqs(status, error_message): mock_instance.async_health_check.assert_awaited_once() +@pytest.mark.asyncio +async def test_health_license_endpoint_with_active_license(): + license_data = { + "expiration_date": "2099-01-01", + "allowed_features": ["feature-a"], + "max_users": 100, + "max_teams": 5, + } + mock_license_check = SimpleNamespace( + license_str="test-license", + public_key=None, + airgapped_license_data=license_data, + verify_license_without_api_request=MagicMock(return_value=True), + ) + + with patch( + "litellm.proxy.proxy_server._license_check", + mock_license_check, + ), patch( + "litellm.proxy.proxy_server.premium_user", + True, + ), patch( + "litellm.proxy.proxy_server.premium_user_data", + license_data, + ): + response = await health_license_endpoint(user_api_key_dict=MagicMock()) + + assert response["has_license"] is True + assert response["license_type"] == "enterprise" + assert response["expiration_date"] == "2099-01-01" + assert response["allowed_features"] == ["feature-a"] + assert response["limits"] == {"max_users": 100, "max_teams": 5} + + +@pytest.mark.asyncio +async def test_health_license_endpoint_without_valid_license(): + mock_license_check = SimpleNamespace( + license_str="invalid-key", + public_key=None, + airgapped_license_data=None, + verify_license_without_api_request=MagicMock(return_value=False), + ) + + 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, + ): + response = await health_license_endpoint(user_api_key_dict=MagicMock()) + + assert response["has_license"] is True + assert response["license_type"] == "community" + assert response["expiration_date"] is None + assert response["allowed_features"] == [] + assert response["limits"] == {"max_users": None, "max_teams": None} + + @pytest.mark.asyncio async def test_test_model_connection_loads_config_from_router(): """ @@ -374,4 +441,3 @@ def test_health_readiness(proxy_client): f"Unexpected db status: {db_status}" print("="*60 + "\n") -