Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 13 additions & 8 deletions litellm/proxy/auth/litellm_license.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@


AUTO_ROUTER_LICENSE_FEATURE: Final = "auto_router"
LICENSE_ALL_FEATURES: Final = "*"
AUTO_ROUTER_LICENSE_REMEDY: Final = "A LiteLLM license with the 'auto_router' feature lifts the limit."


Expand Down Expand Up @@ -153,17 +154,21 @@ def is_team_count_over_limit(self, team_count: int) -> bool:
return False
return team_count > _max_teams_in_license

def grants_feature(self, feature: str) -> bool:
if self.airgapped_license_data is None:
return False
allowed_features: Final = self.airgapped_license_data.get("allowed_features")
granted: Final = allowed_features if isinstance(allowed_features, list) else (allowed_features,)
return feature in granted or LICENSE_ALL_FEATURES in granted

def auto_router_capability_limit(self) -> int | None:
"""
How many auto-routers may claim each licensed capability (heuristic_v2, operator-defined
tier_definitions): unlimited (None) only when the signed license lists the auto_router
feature, otherwise one per capability. A license verified through the API carries no
feature list, so it does not lift the limit either.
How many auto-routers may claim each gated classifier or customization capability:
unlimited (None) only when the signed license lists the auto_router feature or the
"*" wildcard that grants every feature, otherwise one per capability. A license verified
through the API carries no feature list, so it does not lift the limit either.
"""
if self.airgapped_license_data is None:
return 1
allowed_features: Final = self.airgapped_license_data.get("allowed_features")
if isinstance(allowed_features, list) and AUTO_ROUTER_LICENSE_FEATURE in allowed_features:
if self.grants_feature(AUTO_ROUTER_LICENSE_FEATURE):
return None
return 1

Expand Down
35 changes: 31 additions & 4 deletions tests/test_litellm/proxy/auth/test_litellm_license.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,8 +35,8 @@ def test_is_over_limit():


def test_auto_router_capability_limit() -> None:
"""Only the signed license's auto_router feature lifts the one-router limit; an API-verified
license (no airgapped data) and an airgapped license without the feature keep it."""
"""The signed license's auto_router feature or its "*" wildcard lifts the one-router limit; an
API-verified license (no airgapped data) and an airgapped license without either keep it."""
license_check = LicenseCheck()
license_check.airgapped_license_data = {"expiration_date": "2999-01-01", "allowed_features": ["auto_router"]}
assert license_check.auto_router_capability_limit() is None
Expand All @@ -47,25 +47,36 @@ def test_auto_router_capability_limit() -> None:
}
assert license_check.auto_router_capability_limit() is None

license_check.airgapped_license_data = {"expiration_date": "2999-01-01", "allowed_features": ["*"]}
assert license_check.auto_router_capability_limit() is None

license_check.airgapped_license_data = {"expiration_date": "2999-01-01", "allowed_features": ["sso", "*"]}
assert license_check.auto_router_capability_limit() is None

license_check.airgapped_license_data = {"expiration_date": "2999-01-01", "allowed_features": ["sso"]}
assert license_check.auto_router_capability_limit() == 1

license_check.airgapped_license_data = {"expiration_date": "2999-01-01", "allowed_features": "*"}
assert license_check.auto_router_capability_limit() is None

license_check.airgapped_license_data = {"expiration_date": "2999-01-01"}
assert license_check.auto_router_capability_limit() == 1

license_check.airgapped_license_data = None
assert license_check.auto_router_capability_limit() == 1


def _signed_license(expiration_date: str) -> tuple[RSAPublicKey, str]:
def _signed_license(
expiration_date: str, allowed_features: tuple[str, ...] = ("auto_router",)
) -> tuple[RSAPublicKey, str]:
import base64

from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding, rsa

private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
message = json.dumps(
{"expiration_date": expiration_date, "user_id": "u", "allowed_features": ["auto_router"]}
{"expiration_date": expiration_date, "user_id": "u", "allowed_features": list(allowed_features)}
).encode()
signature = private_key.sign(
message,
Expand Down Expand Up @@ -99,3 +110,19 @@ def test_valid_signed_license_with_auto_router_lifts_the_limit() -> None:

assert license_check.verify_license_without_api_request(public_key=public_key, license_key=license_key) is True
assert license_check.auto_router_capability_limit() is None


def test_valid_signed_wildcard_license_lifts_the_limit() -> None:
"""The license generator defaults allowed_features to ["*"], meaning every feature, so a wildcard
license grants auto_router the same way a license that names it does."""
license_check = LicenseCheck()
public_key, license_key = _signed_license("2999-01-01", allowed_features=("*",))

assert license_check.verify_license_without_api_request(public_key=public_key, license_key=license_key) is True
assert license_check.grants_feature("auto_router") is True
assert license_check.auto_router_capability_limit() is None

named_public_key, named_key = _signed_license("2999-01-01", allowed_features=("sso", "audit_logs"))
assert license_check.verify_license_without_api_request(public_key=named_public_key, license_key=named_key) is True
assert license_check.grants_feature("auto_router") is False
assert license_check.auto_router_capability_limit() == 1
Loading