diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/internal_user_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/internal_user_endpoints.py index 2f53f9e92819..1d3268da9a0d 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/internal_user_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/internal_user_endpoints.py @@ -28,6 +28,8 @@ async def available_enterprise_users( premium_user_data, prisma_client, ) + from litellm.repositories.team_repository import TeamRepository + from litellm.repositories.user_repository import UserRepository if prisma_client is None: raise HTTPException( @@ -44,9 +46,8 @@ async def available_enterprise_users( max_users=5, ) - # Count number of rows in LiteLLM_UserTable - user_count = await prisma_client.db.litellm_usertable.count() - team_count = await prisma_client.db.litellm_teamtable.count() + user_count = await UserRepository(prisma_client).count_billable_users() + team_count = await TeamRepository(prisma_client).count() if ( not premium_user_data diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index b517cb0c38d7..fcec551f25aa 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -497,6 +497,12 @@ def __init__( labelnames=[], ) + self.litellm_active_users_metric = self._gauge_factory( + "litellm_active_users", + "Number of billable users in LiteLLM (excludes SCIM-deactivated users)", + labelnames=[], + ) + self.litellm_teams_count_metric = self._gauge_factory( "litellm_teams_count", "Total number of teams in LiteLLM", @@ -3033,6 +3039,7 @@ async def _initialize_user_and_team_count_metrics(self): Updates: - litellm_total_users: Total count of users in the database + - litellm_active_users: Count of billable users (excludes SCIM-deactivated) - litellm_teams_count: Total count of teams in the database """ from litellm.proxy.proxy_server import prisma_client @@ -3047,6 +3054,10 @@ async def _initialize_user_and_team_count_metrics(self): self.litellm_total_users_metric.set(total_users) verbose_logger.debug(f"Prometheus: set litellm_total_users to {total_users}") + billable_users = await UserRepository(prisma_client).count_billable_users() + self.litellm_active_users_metric.set(billable_users) + verbose_logger.debug(f"Prometheus: set litellm_active_users to {billable_users}") + # Get total team count total_teams = await TeamRepository(prisma_client).table.count() self.litellm_teams_count_metric.set(total_teams) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 77e2f354bd40..9374aa3180b3 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -420,8 +420,8 @@ async def new_user( await _check_duplicate_user_email(data.user_email, prisma_client) # Check if license is over limit - total_users = await UserRepository(prisma_client).table.count() - if total_users and _license_check.is_over_limit(total_users=total_users): + billable_users = await UserRepository(prisma_client).count_billable_users() + if billable_users and _license_check.is_over_limit(total_users=billable_users): raise HTTPException( status_code=403, detail="License is over limit. Please contact support@berri.ai to upgrade your license.", diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 73ec56a82e62..89c1a925eeb1 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -858,8 +858,8 @@ async def google_login( if premium_user is not True: # Check if under 'free SSO user' limit if prisma_client is not None: - total_users = await UserRepository(prisma_client).table.count() - if total_users and total_users > 5: + billable_users = await UserRepository(prisma_client).count_billable_users() + if billable_users and billable_users > 5: raise ProxyException( message="You must be a LiteLLM Enterprise user to use SSO for more than 5 users. If you have a license please set `LITELLM_LICENSE` in your env. If you want to obtain a license meet with us here: https://enterprise.litellm.ai/demo You are seeing this error message because You set one of `MICROSOFT_CLIENT_ID`, `GOOGLE_CLIENT_ID`, or `GENERIC_CLIENT_ID` in your env. Please unset this", type=ProxyErrorTypes.auth_error, diff --git a/litellm/repositories/user_repository.py b/litellm/repositories/user_repository.py index 2697f15a6c06..f0b5dfd8bc50 100644 --- a/litellm/repositories/user_repository.py +++ b/litellm/repositories/user_repository.py @@ -59,6 +59,20 @@ async def find_by_team_id(self, team_id: str) -> List[LiteLLM_UserTable]: records = await self.table.find_many(where={"teams": {"has": team_id}}) return self._to_model_list(records) + async def count_billable_users(self) -> int: + """Number of users that count toward the license seat limit. + + Every user is billable except those SCIM-deactivated + (metadata.scim_active == false). Rows where scim_active is absent, + null, or true all count, so seats are counted as total users minus + the deactivated ones. + """ + from prisma import Json # pyright: ignore[reportUnknownVariableType] + + total = await self.count() + deactivated = await self.count(where={"metadata": {"path": ["scim_active"], "equals": Json(False)}}) + return max(0, total - deactivated) + async def create_user( self, user_id: str, diff --git a/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_internal_user_endpoints.py index 32685c5cbd34..fbfb6a99726b 100644 --- a/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_internal_user_endpoints.py @@ -25,6 +25,16 @@ def mock_user_api_key_auth(): yield mock_auth +def _user_count(total, deactivated=0): + """Where-aware count() fake: the filtered query (deactivated users) is + subtracted from the total to yield the billable count.""" + + async def _count(*args, where=None, **kwargs): + return deactivated if where is not None else total + + return _count + + class TestAvailableEnterpriseUsers: @pytest.mark.asyncio async def test_available_users_with_max_users_set( @@ -43,7 +53,7 @@ async def test_available_users_with_max_users_set( ), ): # Mock database count - mock_prisma.db.litellm_usertable.count = AsyncMock(return_value=5) + mock_prisma.db.litellm_usertable.count = _user_count(5) mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=2) # Override the dependency @@ -65,6 +75,36 @@ async def test_available_users_with_max_users_set( # Ensure no negative values assert data["total_users_remaining"] >= 0 + @pytest.mark.asyncio + async def test_available_users_excludes_scim_deactivated( + self, client, mock_user_api_key_auth + ): + """SCIM-deactivated users must not consume a seat: with 5 rows of which + 2 are deactivated, the displayed usage is 3 and a seat is freed.""" + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.premium_user", True), + patch( + "litellm.proxy.proxy_server.premium_user_data", + {"max_users": 10}, + ), + ): + mock_prisma.db.litellm_usertable.count = _user_count(5, deactivated=2) + mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=2) + + client.app.dependency_overrides[mock_user_api_key_auth] = lambda: { + "user_id": "test_user" + } + + response = client.get("/user/available_users") + + assert response.status_code == 200 + data = response.json() + + assert data["total_users"] == 10 + assert data["total_users_used"] == 3 + assert data["total_users_remaining"] == 7 + @pytest.mark.asyncio async def test_available_users_without_max_users_set( self, client, mock_user_api_key_auth @@ -82,7 +122,7 @@ async def test_available_users_without_max_users_set( ), ): # Mock database count - mock_prisma.db.litellm_usertable.count = AsyncMock(return_value=3) + mock_prisma.db.litellm_usertable.count = _user_count(3) mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=1) # Override the dependency @@ -119,7 +159,7 @@ async def test_available_users_negative_remaining_bug( ), ): # Mock database count higher than max_users to trigger the bug - mock_prisma.db.litellm_usertable.count = AsyncMock(return_value=8) + mock_prisma.db.litellm_usertable.count = _user_count(8) mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=3) # Override the dependency diff --git a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py b/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py index 90a9d1fcceba..22a8e8221d46 100644 --- a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py +++ b/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py @@ -191,6 +191,39 @@ async def test_initialize_remaining_budget_metrics_includes_user_team_counts( prometheus_logger._initialize_api_key_budget_metrics.assert_called_once() prometheus_logger._initialize_user_and_team_count_metrics.assert_called_once() + def test_active_users_metric_initialized(self, prometheus_logger): + """litellm_active_users gauge must exist alongside litellm_total_users.""" + assert hasattr(prometheus_logger, "litellm_active_users_metric") + assert prometheus_logger.litellm_active_users_metric is not None + + @pytest.mark.asyncio + async def test_initialize_counts_total_and_active_users(self, prometheus_logger): + """litellm_total_users counts every row; litellm_active_users counts only + billable (non SCIM-deactivated) users.""" + import sys + + prometheus_logger.litellm_total_users_metric = MagicMock() + prometheus_logger.litellm_active_users_metric = MagicMock() + prometheus_logger.litellm_teams_count_metric = MagicMock() + + async def _user_count(*args, where=None, **kwargs): + # 10 rows, 2 of them SCIM-deactivated -> 8 billable + return 2 if where is not None else 10 + + mock_prisma = MagicMock() + mock_prisma.db.litellm_usertable.count = _user_count + mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=4) + + mock_proxy_server = MagicMock() + mock_proxy_server.prisma_client = mock_prisma + + with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}): + await prometheus_logger._initialize_user_and_team_count_metrics() + + prometheus_logger.litellm_total_users_metric.set.assert_called_once_with(10) + prometheus_logger.litellm_active_users_metric.set.assert_called_once_with(8) + prometheus_logger.litellm_teams_count_metric.set.assert_called_once_with(4) + def test_metrics_have_correct_type(self, prometheus_logger): """Test that metrics are Gauge type (not Counter or Histogram)""" from prometheus_client import Gauge 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 27e82df90c10..56eeea82223f 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 @@ -801,9 +801,10 @@ async def test_new_user_license_over_limit(mocker): # Mock the prisma client mock_prisma_client = mocker.MagicMock() - # Setup the mock count response to return a high number of users - async def mock_count(*args, **kwargs): - return 1000 # High user count + # 1000 billable users (no SCIM-deactivated rows): the filtered count used + # for "deactivated" returns 0, so billable == total == 1000 + async def mock_count(*args, where=None, **kwargs): + return 0 if where is not None else 1000 mock_prisma_client.db.litellm_usertable.count = mock_count @@ -852,6 +853,69 @@ async def mock_check_duplicate_user_id(*args, **kwargs): mock_license_check.is_over_limit.assert_called_once_with(total_users=1000) +@pytest.mark.asyncio +async def test_new_user_license_gate_counts_only_billable_users(mocker): + """ + The /user/new license gate must count billable users only (excluding + SCIM-deactivated rows). Deactivated users that push the raw total over + max_users must not block creation, while active users over the limit must. + """ + from litellm.proxy.auth.litellm_license import LicenseCheck + + async def _noop(*args, **kwargs): + return None + + mocker.patch( + "litellm.proxy.management_endpoints.internal_user_endpoints._check_duplicate_user_email", + _noop, + ) + mocker.patch( + "litellm.proxy.management_endpoints.internal_user_endpoints._check_duplicate_user_id", + _noop, + ) + + license_check = LicenseCheck() + license_check.airgapped_license_data = {"max_users": 2} # type: ignore + mocker.patch("litellm.proxy.proxy_server._license_check", license_check) + + key_gen = mocker.patch( + "litellm.proxy.management_endpoints.internal_user_endpoints.generate_key_helper_fn", + new=mocker.AsyncMock(side_effect=RuntimeError("reached key generation")), + ) + + def _prisma(total, deactivated): + client = mocker.MagicMock() + + async def _count(*args, where=None, **kwargs): + return deactivated if where is not None else total + + client.db.litellm_usertable.count = _count + return client + + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + request = NewUserRequest(user_role="internal_user") + + # 2 active + 3 deactivated -> billable 2, not over max_users 2: gate passes + mocker.patch( + "litellm.proxy.proxy_server.prisma_client", _prisma(total=5, deactivated=3) + ) + with pytest.raises(ProxyException) as passed: + await new_user(data=request, user_api_key_dict=admin) + assert key_gen.call_count == 1 + assert "License is over limit" not in str(passed.value.message) + + # 3 active, 0 deactivated -> billable 3, over max_users 2: gate blocks + key_gen.reset_mock() + mocker.patch( + "litellm.proxy.proxy_server.prisma_client", _prisma(total=3, deactivated=0) + ) + with pytest.raises(ProxyException) as blocked: + await new_user(data=request, user_api_key_dict=admin) + assert blocked.value.code == 403 or blocked.value.code == "403" + assert "License is over limit" in str(blocked.value.message) + assert key_gen.call_count == 0 + + @pytest.mark.asyncio async def test_new_user_non_admin_cannot_create_admin(mocker): """ diff --git a/tests/test_litellm/repositories/test_repositories.py b/tests/test_litellm/repositories/test_repositories.py index f22debbae349..af2eea823f4c 100644 --- a/tests/test_litellm/repositories/test_repositories.py +++ b/tests/test_litellm/repositories/test_repositories.py @@ -2182,3 +2182,82 @@ def test_each_repository_binds_its_own_table_name(self): assert name not in seen, f"duplicate table_name {name}" seen.add(name) assert repo_cls(prisma_client).table is getattr(prisma_client.db, name) + + +def _json_path_equals( + metadata: Optional[Dict[str, Any]], path: List[str], expected: Any +) -> bool: + """Reproduce Postgres jsonb path-equals semantics: a missing path yields + SQL NULL, which never matches `equals`.""" + value: Any = metadata + for key in path: + if not isinstance(value, dict) or key not in value: + return False + value = value[key] + return value == expected + + +class _ScimAwareUserTable: + """Fake LiteLLM_UserTable whose count() applies the JSON `where` filter the + way Postgres would, so count_billable_users is checked against an + independent model of the filter rather than echoing its own where dict.""" + + def __init__(self, metadatas: List[Optional[Dict[str, Any]]]): + self._metadatas = metadatas + + async def count(self, where: Optional[Dict[str, Any]] = None) -> int: + if where is None: + return len(self._metadatas) + json_filter = where["metadata"] + path = json_filter["path"] + expected = getattr(json_filter["equals"], "data", json_filter["equals"]) + return sum( + 1 + for metadata in self._metadatas + if _json_path_equals(metadata, path, expected) + ) + + +class TestCountBillableUsers: + def _repo(self, metadatas: List[Optional[Dict[str, Any]]]) -> UserRepository: + client = MockPrismaClient() + client.db.litellm_usertable = _ScimAwareUserTable(metadatas) + return UserRepository(client) + + @pytest.mark.asyncio + async def test_excludes_only_scim_deactivated_users(self): + repo = self._repo( + [ + {}, + {"scim_active": True}, + {"scim_active": True}, + {"scim_active": None}, + {"other": "x"}, + {"scim_active": False}, + ] + ) + assert await repo.count_billable_users() == 5 + + @pytest.mark.asyncio + async def test_absent_null_and_true_all_count_as_billable(self): + repo = self._repo([{}, {"scim_active": None}, {"scim_active": True}]) + assert await repo.count_billable_users() == 3 + + @pytest.mark.asyncio + async def test_all_deactivated_returns_zero(self): + repo = self._repo([{"scim_active": False}, {"scim_active": False}]) + assert await repo.count_billable_users() == 0 + + @pytest.mark.asyncio + async def test_floors_at_zero_when_deactivated_exceeds_total(self): + """The total and deactivated counts are separate queries; a burst of + deactivations between them must never yield a negative seat count.""" + + class _RacyTable: + async def count(self, where=None): + return 5 if where is not None else 2 + + client = MockPrismaClient() + client.db.litellm_usertable = _RacyTable() + repo = UserRepository(client) + assert await repo.count_billable_users() == 0