From 70126d91302233bdb4e4b6aecf5d81b462a94527 Mon Sep 17 00:00:00 2001 From: rioiart Date: Tue, 2 Dec 2025 04:51:42 +0100 Subject: [PATCH] Fix/new org team validate against org (#17333) * fix: skip user budget/model validation for org-scoped teams When creating a team with organization_id, budget and model constraints should be validated against the organization's limits, not the user's personal limits. This allows org admins with restrictive personal budgets to create teams within their organization's more generous limits. Adds 4 unit tests to verify: - Org-scoped teams bypass user budget validation - Org-scoped teams bypass user model validation - Standalone teams still validate against user limits * fix: enforce user budget/model limits for standalone teams in update_team - Add user-level budget and model validation to update_team endpoint for standalone teams, matching the existing pattern in new_team - Org-scoped teams correctly bypass user validation and use organization limits instead - Add 5 new comprehensive tests covering standalone/org team budget/model validation * fix: Add direct TPM/RPM org limit validation and consolidate user team limit checks - Add direct TPM/RPM comparison against org limits in _check_org_team_limits() - Consolidate budget/models/TPM/RPM user validation into _check_user_team_limits() helper - Ensure user limits only apply to standalone teams (organization_id=None) - Org-scoped teams now validate TPM/RPM against org limits (not user limits) - Add 8 tests for TPM/RPM validation scenarios (org and user limits) - Reduce code duplication between new_team() and update_team() --- .../management_endpoints/team_endpoints.py | 496 +++-- .../test_team_endpoints.py | 1602 +++++++++++++++++ 2 files changed, 1901 insertions(+), 197 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 6d4faae5fd8..b697e01a6ef 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -461,11 +461,65 @@ async def _check_org_team_limits( prisma_client: PrismaClient, ) -> None: """ - Check if the organization team is allocating guaranteed throughput limits. If so, raise an error if we're overallocating. - - Only runs check if tpm_limit_type or rpm_limit_type is "guaranteed_throughput" + Check organization team limits including: + - Team budget vs organization's max_budget + - Team models vs organization's allowed models + - Guaranteed throughput limits (tpm/rpm) if applicable """ + # Validate team budget against organization's max_budget + if ( + data.max_budget is not None + and org_table.litellm_budget_table is not None + and org_table.litellm_budget_table.max_budget is not None + and data.max_budget > org_table.litellm_budget_table.max_budget + ): + raise HTTPException( + status_code=400, + detail={ + "error": f"Team max_budget ({data.max_budget}) exceeds organization's max_budget ({org_table.litellm_budget_table.max_budget}). Organization: {org_table.organization_id}" + }, + ) + + # Validate team models against organization's allowed models + if data.models is not None and len(org_table.models) > 0: + for m in data.models: + if m not in org_table.models: + raise HTTPException( + status_code=400, + detail={ + "error": f"Model '{m}' not in organization's allowed models. Organization allowed models={org_table.models}. Organization: {org_table.organization_id}" + }, + ) + + # Validate team TPM/RPM against organization's TPM/RPM limits (direct comparison) + if ( + data.tpm_limit is not None + and org_table.litellm_budget_table is not None + and org_table.litellm_budget_table.tpm_limit is not None + and data.tpm_limit > org_table.litellm_budget_table.tpm_limit + ): + raise HTTPException( + status_code=400, + detail={ + "error": f"Team tpm_limit ({data.tpm_limit}) exceeds organization's tpm_limit ({org_table.litellm_budget_table.tpm_limit}). Organization: {org_table.organization_id}" + }, + ) + + if ( + data.rpm_limit is not None + and org_table.litellm_budget_table is not None + and org_table.litellm_budget_table.rpm_limit is not None + and data.rpm_limit > org_table.litellm_budget_table.rpm_limit + ): + raise HTTPException( + status_code=400, + detail={ + "error": f"Team rpm_limit ({data.rpm_limit}) exceeds organization's rpm_limit ({org_table.litellm_budget_table.rpm_limit}). Organization: {org_table.organization_id}" + }, + ) + + # Check guaranteed throughput limits (only if applicable) rpm_limit_type = getattr(data, "rpm_limit_type", None) or ( data.metadata.get("rpm_limit_type", None) if data.metadata else None ) @@ -503,6 +557,80 @@ async def _check_org_team_limits( ) +async def _check_user_team_limits( + data: Union[NewTeamRequest, UpdateTeamRequest], + user_api_key_dict: UserAPIKeyAuth, + prisma_client: PrismaClient, + user_api_key_cache: Any, +) -> None: + """ + Check user team limits for standalone teams (not org-scoped). + + This validates: + - Team budget vs user's max_budget + - Team models vs user's allowed models + + Should only be called for standalone teams (when organization_id is None). + For org-scoped teams, use _check_org_team_limits() instead. + """ + # Validate team budget against user's max_budget + if data.max_budget is not None and user_api_key_dict.user_id is not None: + user_obj = await get_user_object( + user_id=user_api_key_dict.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + ) + + if ( + user_obj is not None + and user_obj.max_budget is not None + and data.max_budget > user_obj.max_budget + ): + raise HTTPException( + status_code=400, + detail={ + "error": f"max budget higher than user max. User max budget={user_obj.max_budget}. User role={user_api_key_dict.user_role}" + }, + ) + + # Validate team models against user's allowed models + if data.models is not None and len(user_api_key_dict.models) > 0: + for m in data.models: + if m not in user_api_key_dict.models: + raise HTTPException( + status_code=400, + detail={ + "error": f"Model not in allowed user models. User allowed models={user_api_key_dict.models}. User id={user_api_key_dict.user_id}" + }, + ) + + # Validate team TPM/RPM against user's TPM/RPM limits + if ( + data.tpm_limit is not None + and user_api_key_dict.tpm_limit is not None + and data.tpm_limit > user_api_key_dict.tpm_limit + ): + raise HTTPException( + status_code=400, + detail={ + "error": f"tpm limit higher than user max. User tpm limit={user_api_key_dict.tpm_limit}. User role={user_api_key_dict.user_role}" + }, + ) + + if ( + data.rpm_limit is not None + and user_api_key_dict.rpm_limit is not None + and data.rpm_limit > user_api_key_dict.rpm_limit + ): + raise HTTPException( + status_code=400, + detail={ + "error": f"rpm limit higher than user max. User rpm limit={user_api_key_dict.rpm_limit}. User role={user_api_key_dict.user_role}" + }, + ) + + #### TEAM MANAGEMENT #### @router.post( "/team/new", @@ -665,61 +793,16 @@ async def new_team( # noqa: PLR0915 user_api_key_dict.user_role is None or user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN ): # don't restrict proxy admin - if ( - data.tpm_limit is not None - and user_api_key_dict.tpm_limit is not None - and data.tpm_limit > user_api_key_dict.tpm_limit - ): - raise HTTPException( - status_code=400, - detail={ - "error": f"tpm limit higher than user max. User tpm limit={user_api_key_dict.tpm_limit}. User role={user_api_key_dict.user_role}" - }, - ) - - if ( - data.rpm_limit is not None - and user_api_key_dict.rpm_limit is not None - and data.rpm_limit > user_api_key_dict.rpm_limit - ): - raise HTTPException( - status_code=400, - detail={ - "error": f"rpm limit higher than user max. User rpm limit={user_api_key_dict.rpm_limit}. User role={user_api_key_dict.user_role}" - }, - ) - - if data.max_budget is not None and user_api_key_dict.user_id is not None: - # Fetch user object to get max_budget - user_obj = await get_user_object( - user_id=user_api_key_dict.user_id, + # Only validate user budget/models/tpm/rpm for standalone teams (not org-scoped) + # For org-scoped teams, validation is done by _check_org_team_limits() + if data.organization_id is None: + await _check_user_team_limits( + data=data, + user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, - user_id_upsert=False, ) - if ( - user_obj is not None - and user_obj.max_budget is not None - and data.max_budget > user_obj.max_budget - ): - raise HTTPException( - status_code=400, - detail={ - "error": f"max budget higher than user max. User max budget={user_obj.max_budget}. User role={user_api_key_dict.user_role}" - }, - ) - - if data.models is not None and len(user_api_key_dict.models) > 0: - for m in data.models: - if m not in user_api_key_dict.models: - raise HTTPException( - status_code=400, - detail={ - "error": f"Model not in allowed user models. User allowed models={user_api_key_dict.models}. User id={user_api_key_dict.user_id}" - }, - ) - if user_api_key_dict.user_id is not None: creating_user_in_list = False for member in data.members_with_roles: @@ -1151,168 +1234,187 @@ async def update_team( }' ``` """ - from litellm.proxy.auth.auth_checks import _cache_team_object - from litellm.proxy.proxy_server import ( - litellm_proxy_admin_name, - llm_router, - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) - - if prisma_client is None: - raise HTTPException( - status_code=500, - detail={"error": CommonProxyErrors.db_not_connected_error.value}, + try: + from litellm.proxy.auth.auth_checks import _cache_team_object + from litellm.proxy.proxy_server import ( + litellm_proxy_admin_name, + llm_router, + prisma_client, + proxy_logging_obj, + user_api_key_cache, ) - if data.team_id is None: - raise HTTPException(status_code=400, detail={"error": "No team id passed in"}) - verbose_proxy_logger.debug("/team/update - %s", data) - - existing_team_row = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": data.team_id} - ) - - if existing_team_row is None: - raise HTTPException( - status_code=404, - detail={"error": f"Team not found, passed team_id={data.team_id}"}, - ) - - if ( - data.organization_id is not None and len(data.organization_id) > 0 - ): # allow unsetting the organization_id - await fetch_and_validate_organization( - organization_id=data.organization_id, - existing_team_row=existing_team_row, - llm_router=llm_router, - prisma_client=prisma_client, - ) - elif data.organization_id is not None and len(data.organization_id) == 0: - # unsetting the organization_id - data.organization_id = None - - # check org team limits - if updating team that belongs to an org - org_id_to_check = ( - data.organization_id - if data.organization_id is not None - else existing_team_row.organization_id - ) - if ( - org_id_to_check is not None - and isinstance(org_id_to_check, str) - and prisma_client is not None - ): - org_table = await get_org_object( - org_id=org_id_to_check, - user_api_key_cache=user_api_key_cache, - prisma_client=prisma_client, - ) - if org_table is not None: - await _check_org_team_limits( - org_table=org_table, - data=data, - prisma_client=prisma_client, + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - updated_kv = data.json(exclude_unset=True) + if data.team_id is None: + raise HTTPException(status_code=400, detail={"error": "No team id passed in"}) + verbose_proxy_logger.debug("/team/update - %s", data) - # Check budget_duration and budget_reset_at - if data.budget_duration is not None: - from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time + existing_team_row = await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": data.team_id} + ) - reset_at = get_budget_reset_time(budget_duration=data.budget_duration) + if existing_team_row is None: + raise HTTPException( + status_code=404, + detail={"error": f"Team not found, passed team_id={data.team_id}"}, + ) - # set the budget_reset_at in DB - updated_kv["budget_reset_at"] = reset_at + if ( + data.organization_id is not None and len(data.organization_id) > 0 + ): # allow unsetting the organization_id + await fetch_and_validate_organization( + organization_id=data.organization_id, + existing_team_row=existing_team_row, + llm_router=llm_router, + prisma_client=prisma_client, + ) + elif data.organization_id is not None and len(data.organization_id) == 0: + # unsetting the organization_id + data.organization_id = None - if TeamMemberBudgetHandler.should_create_budget( - team_member_budget=data.team_member_budget, - team_member_rpm_limit=data.team_member_rpm_limit, - team_member_tpm_limit=data.team_member_tpm_limit, - ): - updated_kv = await TeamMemberBudgetHandler.upsert_team_member_budget_table( - team_table=existing_team_row, - user_api_key_dict=user_api_key_dict, - updated_kv=updated_kv, + # check org team limits - if updating team that belongs to an org + org_id_to_check = ( + data.organization_id + if data.organization_id is not None + else existing_team_row.organization_id + ) + if ( + org_id_to_check is not None + and isinstance(org_id_to_check, str) + and prisma_client is not None + ): + org_table = await get_org_object( + org_id=org_id_to_check, + user_api_key_cache=user_api_key_cache, + prisma_client=prisma_client, + ) + if org_table is not None: + await _check_org_team_limits( + org_table=org_table, + data=data, + prisma_client=prisma_client, + ) + + # Check user limits for standalone teams (not org-scoped) + # Skip for PROXY_ADMIN users + if ( + user_api_key_dict.user_role is None + or user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN + ): + # Only validate user budget/models for standalone teams + # For org-scoped teams, validation is done by _check_org_team_limits() above + if org_id_to_check is None: + await _check_user_team_limits( + data=data, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + + updated_kv = data.json(exclude_unset=True) + + # Check budget_duration and budget_reset_at + if data.budget_duration is not None: + from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time + + reset_at = get_budget_reset_time(budget_duration=data.budget_duration) + + # set the budget_reset_at in DB + updated_kv["budget_reset_at"] = reset_at + + if TeamMemberBudgetHandler.should_create_budget( team_member_budget=data.team_member_budget, team_member_rpm_limit=data.team_member_rpm_limit, team_member_tpm_limit=data.team_member_tpm_limit, - ) - else: - TeamMemberBudgetHandler._clean_team_member_fields(updated_kv) - - # Check object permission - if data.object_permission is not None: - updated_kv = await handle_update_object_permission( - data_json=updated_kv, - existing_team_row=existing_team_row, - ) - - # update team metadata fields - _team_metadata_fields = LiteLLM_ManagementEndpoint_MetadataFields_Premium - for field in _team_metadata_fields: - if field in updated_kv and updated_kv[field] is not None: - _update_metadata_field( + ): + updated_kv = await TeamMemberBudgetHandler.upsert_team_member_budget_table( + team_table=existing_team_row, + user_api_key_dict=user_api_key_dict, updated_kv=updated_kv, - field_name=field, + team_member_budget=data.team_member_budget, + team_member_rpm_limit=data.team_member_rpm_limit, + team_member_tpm_limit=data.team_member_tpm_limit, + ) + else: + TeamMemberBudgetHandler._clean_team_member_fields(updated_kv) + + # Check object permission + if data.object_permission is not None: + updated_kv = await handle_update_object_permission( + data_json=updated_kv, + existing_team_row=existing_team_row, ) - for field in LiteLLM_ManagementEndpoint_MetadataFields: - if field in updated_kv and updated_kv[field] is not None: - _update_metadata_field( - updated_kv=updated_kv, - field_name=field, + # update team metadata fields + _team_metadata_fields = LiteLLM_ManagementEndpoint_MetadataFields_Premium + for field in _team_metadata_fields: + if field in updated_kv and updated_kv[field] is not None: + _update_metadata_field( + updated_kv=updated_kv, + field_name=field, + ) + + for field in LiteLLM_ManagementEndpoint_MetadataFields: + if field in updated_kv and updated_kv[field] is not None: + _update_metadata_field( + updated_kv=updated_kv, + field_name=field, + ) + + if "model_aliases" in updated_kv: + updated_kv.pop("model_aliases") + _model_id = await _update_model_table( + data=data, + model_id=existing_team_row.model_id, + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ) + if _model_id is not None: + updated_kv["model_id"] = _model_id + + updated_kv = prisma_client.jsonify_team_object(db_data=updated_kv) + team_row: Optional[LiteLLM_TeamTable] = ( + await prisma_client.db.litellm_teamtable.update( + where={"team_id": data.team_id}, + data=updated_kv, + include={"litellm_model_table": True}, # type: ignore + ) + ) + + if team_row is None or team_row.team_id is None: + raise HTTPException( + status_code=400, + detail={"error": "Team doesn't exist. Got={}".format(team_row)}, ) - if "model_aliases" in updated_kv: - updated_kv.pop("model_aliases") - _model_id = await _update_model_table( - data=data, - model_id=existing_team_row.model_id, - prisma_client=prisma_client, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, - ) - if _model_id is not None: - updated_kv["model_id"] = _model_id - - updated_kv = prisma_client.jsonify_team_object(db_data=updated_kv) - team_row: Optional[LiteLLM_TeamTable] = ( - await prisma_client.db.litellm_teamtable.update( - where={"team_id": data.team_id}, - data=updated_kv, - include={"litellm_model_table": True}, # type: ignore - ) - ) - - if team_row is None or team_row.team_id is None: - raise HTTPException( - status_code=400, - detail={"error": "Team doesn't exist. Got={}".format(team_row)}, + verbose_proxy_logger.info("Successfully updated team - %s, info", team_row.team_id) + await _cache_team_object( + team_id=team_row.team_id, + team_table=LiteLLM_TeamTableCachedObj(**team_row.model_dump()), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, ) - verbose_proxy_logger.info("Successfully updated team - %s, info", team_row.team_id) - await _cache_team_object( - team_id=team_row.team_id, - team_table=LiteLLM_TeamTableCachedObj(**team_row.model_dump()), - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) + # Enterprise Feature - Audit Logging. Enable with litellm.store_audit_logs = True + if litellm.store_audit_logs is True: + await _create_team_update_audit_log( + existing_team_row=existing_team_row, + updated_kv=updated_kv, + team_id=data.team_id, + litellm_changed_by=litellm_changed_by, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ) - # Enterprise Feature - Audit Logging. Enable with litellm.store_audit_logs = True - if litellm.store_audit_logs is True: - await _create_team_update_audit_log( - existing_team_row=existing_team_row, - updated_kv=updated_kv, - team_id=data.team_id, - litellm_changed_by=litellm_changed_by, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, - ) - - return {"team_id": team_row.team_id, "data": team_row} + return {"team_id": team_row.team_id, "data": team_row} + except Exception as e: + raise handle_exception_on_proxy(e) async def handle_update_object_permission( 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 86b23c98ba5..06ec71a84f8 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -2015,3 +2015,1605 @@ async def test_new_team_max_budget_within_user_limit(): assert result is not None assert result["team_id"] == "team-within-budget-789" assert result["max_budget"] == 50.0 + + +@pytest.mark.asyncio +async def test_new_team_org_scoped_budget_bypasses_user_limit(): + """ + Test that /team/new with organization_id does NOT validate budget against user's personal max_budget. + + This is the bug fix for: When an org admin creates an org-scoped team, the team's budget should + be validated against the organization's limits, not the user's personal limits. + + Scenario: + - Organization has max_budget=$100 + - User (org admin) has personal max_budget=$3 + - Team is created with organization_id and max_budget=$50 + - Expected: Should succeed (within org's $100 limit) + - Bug behavior: Would fail with "max budget higher than user max. User max budget=3.0" + """ + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest, UserAPIKeyAuth, LiteLLM_UserTable, LiteLLM_OrganizationTable + from litellm.proxy.management_endpoints.team_endpoints import new_team + + # Create non-admin user with very restrictive personal budget ($3) + org_admin_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="org-admin-user-123", + user_max_budget=3.0, # Restrictive personal budget + models=[], # Empty models list to bypass model validation + ) + + # Create team request with budget ($50) that's within org's limit but exceeds user's personal limit + team_request = NewTeamRequest( + team_alias="org-scoped-team", + max_budget=50.0, # Within org's $100 limit, but exceeds user's $3 limit + organization_id="test-org-123", # This makes it an org-scoped team + ) + + dummy_request = MagicMock(spec=Request) + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( + "litellm.proxy.proxy_server.user_api_key_cache" + ) as mock_cache, patch( + "litellm.proxy.proxy_server._license_check" + ) as mock_license, patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + ), patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit, patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object" + ) as mock_get_org: + + # Setup mocks + mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) + mock_license.is_team_count_over_limit.return_value = False + mock_prisma.jsonify_team_object = lambda db_data: db_data + mock_prisma.get_data = AsyncMock(return_value=None) + mock_prisma.update_data = AsyncMock() + + # Mock organization with $100 budget + mock_org = MagicMock(spec=LiteLLM_OrganizationTable) + mock_org.organization_id = "test-org-123" + mock_org.max_budget = 100.0 + mock_org.models = ["gpt-4", "gpt-3.5-turbo", "claude-3-opus"] + mock_org.litellm_budget_table = None # No budget table for this test + mock_get_org.return_value = mock_org + + # Mock user cache to return user with restrictive personal budget + mock_user_obj = LiteLLM_UserTable( + user_id="org-admin-user-123", + max_budget=3.0, # Restrictive personal budget + ) + mock_cache.async_get_cache = AsyncMock(return_value=mock_user_obj) + + # Mock team creation + mock_created_team = MagicMock() + mock_created_team.team_id = "team-org-scoped-789" + mock_created_team.team_alias = "org-scoped-team" + mock_created_team.max_budget = 50.0 + mock_created_team.organization_id = "test-org-123" + mock_created_team.members_with_roles = [] + mock_created_team.metadata = None + mock_created_team.model_dump.return_value = { + "team_id": "team-org-scoped-789", + "team_alias": "org-scoped-team", + "max_budget": 50.0, + "organization_id": "test-org-123", + "members_with_roles": [], + } + mock_prisma.db.litellm_teamtable.create = AsyncMock(return_value=mock_created_team) + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_created_team) + + # Mock model table + mock_prisma.db.litellm_modeltable = MagicMock() + mock_prisma.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model123")) + + # Mock user table operations + mock_user = MagicMock() + mock_user.user_id = "org-admin-user-123" + mock_user.model_dump.return_value = {"user_id": "org-admin-user-123", "teams": ["team-org-scoped-789"]} + mock_prisma.db.litellm_usertable = MagicMock() + mock_prisma.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user) + mock_prisma.db.litellm_usertable.update = AsyncMock(return_value=mock_user) + + # Mock team membership table + mock_membership = MagicMock() + mock_membership.model_dump.return_value = { + "team_id": "team-org-scoped-789", + "user_id": "org-admin-user-123", + "budget_id": None, + } + mock_prisma.db.litellm_teammembership = MagicMock() + mock_prisma.db.litellm_teammembership.create = AsyncMock(return_value=mock_membership) + + # Should NOT raise an exception - the fix should bypass user budget validation for org-scoped teams + result = await new_team( + data=team_request, + http_request=dummy_request, + user_api_key_dict=org_admin_user, + ) + + # Verify the team was created successfully with the higher budget + assert result is not None + assert result["team_id"] == "team-org-scoped-789" + assert result["max_budget"] == 50.0 + assert result["organization_id"] == "test-org-123" + + +@pytest.mark.asyncio +async def test_new_team_org_scoped_models_bypasses_user_limit(): + """ + Test that /team/new with organization_id does NOT validate models against user's personal models. + + This is the bug fix for: When an org admin creates an org-scoped team, the team's models should + be validated against the organization's models, not the user's personal models. + + Scenario: + - Organization has models=['gpt-4', 'gpt-3.5-turbo', 'claude-3-opus'] + - User (org admin) has personal models=['no-default-models'] + - Team is created with organization_id and models=['gpt-4'] + - Expected: Should succeed (within org's allowed models) + - Bug behavior: Would fail with "Model not in allowed user models" + """ + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest, UserAPIKeyAuth, LiteLLM_UserTable, LiteLLM_OrganizationTable + from litellm.proxy.management_endpoints.team_endpoints import new_team + + # Create non-admin user with restrictive personal models + org_admin_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="org-admin-user-456", + user_max_budget=None, # No budget restriction for this test + models=["no-default-models"], # Restrictive personal models + ) + + # Create team request with models that are within org's allowed models but not user's + team_request = NewTeamRequest( + team_alias="org-scoped-models-team", + models=["gpt-4"], # Within org's allowed models, but not in user's personal models + organization_id="test-org-456", # This makes it an org-scoped team + ) + + dummy_request = MagicMock(spec=Request) + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( + "litellm.proxy.proxy_server.user_api_key_cache" + ) as mock_cache, patch( + "litellm.proxy.proxy_server._license_check" + ) as mock_license, patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + ), patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit, patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object" + ) as mock_get_org: + + # Setup mocks + mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) + mock_license.is_team_count_over_limit.return_value = False + mock_prisma.jsonify_team_object = lambda db_data: db_data + mock_prisma.get_data = AsyncMock(return_value=None) + mock_prisma.update_data = AsyncMock() + + # Mock organization with allowed models + mock_org = MagicMock(spec=LiteLLM_OrganizationTable) + mock_org.organization_id = "test-org-456" + mock_org.max_budget = 100.0 + mock_org.models = ["gpt-4", "gpt-3.5-turbo", "claude-3-opus"] + mock_org.litellm_budget_table = None + mock_get_org.return_value = mock_org + + # Mock user cache + mock_user_obj = LiteLLM_UserTable( + user_id="org-admin-user-456", + max_budget=None, + ) + mock_cache.async_get_cache = AsyncMock(return_value=mock_user_obj) + + # Mock team creation + mock_created_team = MagicMock() + mock_created_team.team_id = "team-org-scoped-models-789" + mock_created_team.team_alias = "org-scoped-models-team" + mock_created_team.max_budget = None + mock_created_team.organization_id = "test-org-456" + mock_created_team.models = ["gpt-4"] + mock_created_team.members_with_roles = [] + mock_created_team.metadata = None + mock_created_team.model_dump.return_value = { + "team_id": "team-org-scoped-models-789", + "team_alias": "org-scoped-models-team", + "max_budget": None, + "organization_id": "test-org-456", + "models": ["gpt-4"], + "members_with_roles": [], + } + mock_prisma.db.litellm_teamtable.create = AsyncMock(return_value=mock_created_team) + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_created_team) + + # Mock model table + mock_prisma.db.litellm_modeltable = MagicMock() + mock_prisma.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model123")) + + # Mock user table operations + mock_user = MagicMock() + mock_user.user_id = "org-admin-user-456" + mock_user.model_dump.return_value = {"user_id": "org-admin-user-456", "teams": ["team-org-scoped-models-789"]} + mock_prisma.db.litellm_usertable = MagicMock() + mock_prisma.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user) + mock_prisma.db.litellm_usertable.update = AsyncMock(return_value=mock_user) + + # Mock team membership table + mock_membership = MagicMock() + mock_membership.model_dump.return_value = { + "team_id": "team-org-scoped-models-789", + "user_id": "org-admin-user-456", + "budget_id": None, + } + mock_prisma.db.litellm_teammembership = MagicMock() + mock_prisma.db.litellm_teammembership.create = AsyncMock(return_value=mock_membership) + + # Should NOT raise an exception - the fix should bypass user model validation for org-scoped teams + result = await new_team( + data=team_request, + http_request=dummy_request, + user_api_key_dict=org_admin_user, + ) + + # Verify the team was created successfully with the org's models + assert result is not None + assert result["team_id"] == "team-org-scoped-models-789" + assert result["models"] == ["gpt-4"] + assert result["organization_id"] == "test-org-456" + + +@pytest.mark.asyncio +async def test_new_team_standalone_validates_against_user_models(): + """ + Test that /team/new WITHOUT organization_id still validates models against user's personal models. + + This ensures that standalone teams (not org-scoped) still use user-level validation. + + Scenario: + - User has personal models=['no-default-models'] + - Team is created WITHOUT organization_id and models=['gpt-4'] + - Expected: Should fail with "Model not in allowed user models" + """ + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest, ProxyException, UserAPIKeyAuth + from litellm.proxy.management_endpoints.team_endpoints import new_team + + # Create non-admin user with restrictive personal models + non_admin_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="non-admin-user-789", + user_max_budget=None, + models=["no-default-models"], # Restrictive personal models + ) + + # Create standalone team request (no organization_id) with models not in user's list + team_request = NewTeamRequest( + team_alias="standalone-team", + models=["gpt-4"], # Not in user's allowed models + # Note: No organization_id - this is a standalone team + ) + + dummy_request = MagicMock(spec=Request) + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( + "litellm.proxy.proxy_server._license_check" + ) as mock_license, patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + ), patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit: + # Setup basic mocks + mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) + mock_license.is_team_count_over_limit.return_value = False + mock_prisma.get_data = AsyncMock(return_value=None) + + # Should raise ProxyException because gpt-4 is not in user's allowed models + with pytest.raises(ProxyException) as exc_info: + await new_team( + data=team_request, + http_request=dummy_request, + user_api_key_dict=non_admin_user, + ) + + # Verify exception details + assert exc_info.value.code == '400' + assert "Model not in allowed user models" in str(exc_info.value.message) + assert "no-default-models" in str(exc_info.value.message) + + +@pytest.mark.asyncio +async def test_new_team_standalone_validates_against_user_budget(): + """ + Test that /team/new WITHOUT organization_id still validates budget against user's personal max_budget. + + This ensures that standalone teams (not org-scoped) still use user-level validation. + This is essentially the same as test_new_team_max_budget_exceeds_user_max_budget but + explicitly showing the contrast with org-scoped teams. + + Scenario: + - User has personal max_budget=$3 + - Team is created WITHOUT organization_id and max_budget=$50 + - Expected: Should fail with "max budget higher than user max" + """ + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest, ProxyException, UserAPIKeyAuth, LiteLLM_UserTable + from litellm.proxy.management_endpoints.team_endpoints import new_team + + # Create non-admin user with restrictive personal budget + non_admin_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="non-admin-user-budget-789", + user_max_budget=100.0, # This is for key auth, actual budget is from user object + models=[], # Empty models list to bypass model validation + ) + + # Create standalone team request (no organization_id) with budget exceeding user's limit + team_request = NewTeamRequest( + team_alias="standalone-budget-team", + max_budget=50.0, # Exceeds user's personal budget + # Note: No organization_id - this is a standalone team + ) + + dummy_request = MagicMock(spec=Request) + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( + "litellm.proxy.proxy_server._license_check" + ) as mock_license, patch( + "litellm.proxy.proxy_server.user_api_key_cache" + ) as mock_cache, patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + ), patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit: + # Setup basic mocks + mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) + mock_license.is_team_count_over_limit.return_value = False + mock_prisma.get_data = AsyncMock(return_value=None) + + # Mock user cache to return user with restrictive personal budget ($3) + mock_user_obj = LiteLLM_UserTable( + user_id="non-admin-user-budget-789", + max_budget=3.0, # Restrictive personal budget + ) + mock_cache.async_get_cache = AsyncMock(return_value=mock_user_obj) + + # Should raise ProxyException because budget exceeds user's max_budget + with pytest.raises(ProxyException) as exc_info: + await new_team( + data=team_request, + http_request=dummy_request, + user_api_key_dict=non_admin_user, + ) + + # Verify exception details + assert exc_info.value.code == '400' + assert "max budget higher than user max" in str(exc_info.value.message) + assert "3.0" in str(exc_info.value.message) # User's max_budget should be mentioned + + +@pytest.mark.asyncio +async def test_new_team_org_scoped_budget_exceeds_org_limit(): + """ + Test that /team/new with organization_id fails when team budget exceeds organization's max_budget. + + Scenario: + - Organization has max_budget=$100 + - Team is created with organization_id and max_budget=$150 + - Expected: Should fail with error about exceeding org budget + """ + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest, ProxyException, UserAPIKeyAuth, LiteLLM_OrganizationTable, LiteLLM_BudgetTable + from litellm.proxy.management_endpoints.team_endpoints import new_team + + # Create user (org admin) + org_admin_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="org-admin-user-budget-test", + models=[], + ) + + # Create team request with budget ($150) that exceeds org's limit ($100) + team_request = NewTeamRequest( + team_alias="org-team-exceeds-budget", + max_budget=150.0, # Exceeds org's $100 limit + organization_id="test-org-budget-limit", + ) + + dummy_request = MagicMock(spec=Request) + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( + "litellm.proxy.proxy_server.user_api_key_cache" + ) as mock_cache, patch( + "litellm.proxy.proxy_server._license_check" + ) as mock_license, patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + ), patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit, patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object" + ) as mock_get_org: + + # Setup mocks + mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) + mock_license.is_team_count_over_limit.return_value = False + mock_prisma.get_data = AsyncMock(return_value=None) + + # Mock organization with $100 budget limit + mock_budget_table = MagicMock(spec=LiteLLM_BudgetTable) + mock_budget_table.max_budget = 100.0 + + mock_org = MagicMock(spec=LiteLLM_OrganizationTable) + mock_org.organization_id = "test-org-budget-limit" + mock_org.models = ["gpt-4", "gpt-3.5-turbo"] + mock_org.litellm_budget_table = mock_budget_table + mock_get_org.return_value = mock_org + + # Should raise ProxyException because team budget exceeds org budget + with pytest.raises(ProxyException) as exc_info: + await new_team( + data=team_request, + http_request=dummy_request, + user_api_key_dict=org_admin_user, + ) + + # Verify exception details + assert exc_info.value.code == '400' + assert "exceeds organization" in str(exc_info.value.message).lower() or "organization" in str(exc_info.value.message).lower() + + +@pytest.mark.asyncio +async def test_new_team_org_scoped_models_not_in_org_models(): + """ + Test that /team/new with organization_id fails when team models are not in organization's allowed models. + + Scenario: + - Organization has models=['gpt-4', 'gpt-3.5-turbo'] + - Team is created with organization_id and models=['claude-3-opus'] + - Expected: Should fail with error about model not in org's allowed models + """ + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest, ProxyException, UserAPIKeyAuth, LiteLLM_OrganizationTable, LiteLLM_BudgetTable + from litellm.proxy.management_endpoints.team_endpoints import new_team + + # Create user (org admin) + org_admin_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="org-admin-user-models-test", + models=[], + ) + + # Create team request with model not in org's allowed list + team_request = NewTeamRequest( + team_alias="org-team-invalid-model", + models=["claude-3-opus"], # Not in org's allowed models + organization_id="test-org-models-limit", + ) + + dummy_request = MagicMock(spec=Request) + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( + "litellm.proxy.proxy_server.user_api_key_cache" + ) as mock_cache, patch( + "litellm.proxy.proxy_server._license_check" + ) as mock_license, patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + ), patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit, patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object" + ) as mock_get_org: + + # Setup mocks + mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) + mock_license.is_team_count_over_limit.return_value = False + mock_prisma.get_data = AsyncMock(return_value=None) + + # Mock organization with specific allowed models (not including claude-3-opus) + mock_org = MagicMock(spec=LiteLLM_OrganizationTable) + mock_org.organization_id = "test-org-models-limit" + mock_org.models = ["gpt-4", "gpt-3.5-turbo"] # claude-3-opus is NOT allowed + mock_org.litellm_budget_table = None + mock_get_org.return_value = mock_org + + # Should raise ProxyException because claude-3-opus is not in org's allowed models + with pytest.raises(ProxyException) as exc_info: + await new_team( + data=team_request, + http_request=dummy_request, + user_api_key_dict=org_admin_user, + ) + + # Verify exception details + assert exc_info.value.code == '400' + assert "claude-3-opus" in str(exc_info.value.message) or "organization" in str(exc_info.value.message).lower() + + +@pytest.mark.asyncio +async def test_update_team_standalone_budget_exceeds_user_limit(): + """ + Test that /team/update for a standalone team fails when new budget exceeds user's max_budget. + + Scenario: + - User has personal max_budget=$50 + - Standalone team exists (no organization_id) + - User tries to update team budget to $100 + - Expected: Should fail with error about exceeding user budget + """ + from fastapi import Request + + from litellm.proxy._types import UpdateTeamRequest, ProxyException, UserAPIKeyAuth, LiteLLM_UserTable + from litellm.proxy.management_endpoints.team_endpoints import update_team + + # Create non-admin user with restrictive personal budget + non_admin_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="non-admin-update-test", + models=[], + ) + + # Create update request with budget exceeding user's limit + update_request = UpdateTeamRequest( + team_id="standalone-team-123", + max_budget=100.0, # Exceeds user's $50 limit + ) + + dummy_request = MagicMock(spec=Request) + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( + "litellm.proxy.proxy_server.user_api_key_cache" + ) as mock_cache, patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + ), patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit: + + # Mock existing standalone team (no organization_id) + mock_existing_team = MagicMock() + mock_existing_team.team_id = "standalone-team-123" + mock_existing_team.organization_id = None # Standalone team + mock_existing_team.max_budget = 30.0 + mock_existing_team.model_dump.return_value = { + "team_id": "standalone-team-123", + "organization_id": None, + "max_budget": 30.0, + } + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) + + # Mock user cache to return user with restrictive budget + mock_user_obj = LiteLLM_UserTable( + user_id="non-admin-update-test", + max_budget=50.0, # User's budget limit + ) + mock_cache.async_get_cache = AsyncMock(return_value=mock_user_obj) + + # Should raise ProxyException because new budget exceeds user's max_budget + with pytest.raises(ProxyException) as exc_info: + await update_team( + data=update_request, + http_request=dummy_request, + user_api_key_dict=non_admin_user, + ) + + # Verify exception details + assert exc_info.value.code == '400' + assert "budget" in str(exc_info.value.message).lower() + + +@pytest.mark.asyncio +async def test_update_team_org_scoped_budget_exceeds_org_limit(): + """ + Test that /team/update for an org-scoped team fails when new budget exceeds organization's max_budget. + + Scenario: + - Organization has max_budget=$100 + - Org-scoped team exists + - User tries to update team budget to $150 + - Expected: Should fail with error about exceeding org budget + """ + from fastapi import Request + + from litellm.proxy._types import UpdateTeamRequest, ProxyException, UserAPIKeyAuth, LiteLLM_OrganizationTable, LiteLLM_BudgetTable + from litellm.proxy.management_endpoints.team_endpoints import update_team + + # Create user (org admin) + org_admin_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="org-admin-update-test", + models=[], + ) + + # Create update request with budget exceeding org's limit + update_request = UpdateTeamRequest( + team_id="org-team-456", + max_budget=150.0, # Exceeds org's $100 limit + ) + + dummy_request = MagicMock(spec=Request) + + # Mock organization with $100 budget limit + mock_budget_table = MagicMock(spec=LiteLLM_BudgetTable) + mock_budget_table.max_budget = 100.0 + + mock_org = MagicMock(spec=LiteLLM_OrganizationTable) + mock_org.organization_id = "test-org-update" + mock_org.models = ["gpt-4"] + mock_org.litellm_budget_table = mock_budget_table + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( + "litellm.proxy.proxy_server.user_api_key_cache" + ) as mock_cache, patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + ), patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit, patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object", + new=AsyncMock(return_value=mock_org) + ) as mock_get_org: + + # Mock existing org-scoped team + mock_existing_team = MagicMock() + mock_existing_team.team_id = "org-team-456" + mock_existing_team.organization_id = "test-org-update" + mock_existing_team.max_budget = 80.0 + mock_existing_team.model_dump.return_value = { + "team_id": "org-team-456", + "organization_id": "test-org-update", + "max_budget": 80.0, + } + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) + + # Should raise ProxyException because new budget exceeds org's max_budget + with pytest.raises(ProxyException) as exc_info: + await update_team( + data=update_request, + http_request=dummy_request, + user_api_key_dict=org_admin_user, + ) + + # Verify exception details + assert exc_info.value.code == '400' + assert "organization" in str(exc_info.value.message).lower() or "budget" in str(exc_info.value.message).lower() + + +@pytest.mark.asyncio +async def test_update_team_standalone_models_exceeds_user_limit(): + """ + Test that /team/update for a standalone team fails when models are not in user's allowed models. + + Scenario: + - User has personal models=['gpt-3.5-turbo'] + - Standalone team exists (no organization_id) + - User tries to update team models to ['gpt-4'] (not in user's allowed models) + - Expected: Should fail with error about model not in user's allowed models + """ + from fastapi import Request + + from litellm.proxy._types import UpdateTeamRequest, ProxyException, UserAPIKeyAuth + from litellm.proxy.management_endpoints.team_endpoints import update_team + + # Create non-admin user with restrictive personal models + non_admin_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="non-admin-update-models-test", + models=["gpt-3.5-turbo"], # Restrictive model list + ) + + # Create update request with model not in user's allowed list + update_request = UpdateTeamRequest( + team_id="standalone-team-models-123", + models=["gpt-4"], # Not in user's allowed models + ) + + dummy_request = MagicMock(spec=Request) + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( + "litellm.proxy.proxy_server.user_api_key_cache" + ) as mock_cache, patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + ), patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit: + + # Mock existing standalone team (no organization_id) + mock_existing_team = MagicMock() + mock_existing_team.team_id = "standalone-team-models-123" + mock_existing_team.organization_id = None # Standalone team + mock_existing_team.models = ["gpt-3.5-turbo"] + mock_existing_team.model_dump.return_value = { + "team_id": "standalone-team-models-123", + "organization_id": None, + "models": ["gpt-3.5-turbo"], + } + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) + + # Should raise ProxyException because model not in user's allowed models + with pytest.raises(ProxyException) as exc_info: + await update_team( + data=update_request, + http_request=dummy_request, + user_api_key_dict=non_admin_user, + ) + + # Verify exception details + assert exc_info.value.code == '400' + assert "model" in str(exc_info.value.message).lower() + + +@pytest.mark.asyncio +async def test_update_team_org_scoped_budget_bypasses_user_limit(): + """ + Test that /team/update for an org-scoped team does NOT validate budget against user's personal max_budget. + + Scenario: + - Organization has max_budget=$100 + - User (org admin) has personal max_budget=$3 + - Org-scoped team exists with current budget=$30 + - User tries to update team budget to $50 (within org limit, exceeds user limit) + - Expected: Should succeed (validated against org, not user) + """ + from fastapi import Request + + from litellm.proxy._types import UpdateTeamRequest, UserAPIKeyAuth, LiteLLM_UserTable, LiteLLM_OrganizationTable, LiteLLM_BudgetTable + from litellm.proxy.management_endpoints.team_endpoints import update_team + + # Create user with very restrictive personal budget ($3) + org_admin_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="org-admin-update-budget-test", + models=[], + ) + + # Create update request with budget within org limit but exceeding user limit + update_request = UpdateTeamRequest( + team_id="org-team-update-budget-123", + max_budget=50.0, # Within org's $100 limit, exceeds user's $3 limit + ) + + dummy_request = MagicMock(spec=Request) + + # Mock organization with $100 budget limit + mock_budget_table = MagicMock(spec=LiteLLM_BudgetTable) + mock_budget_table.max_budget = 100.0 + + mock_org = MagicMock(spec=LiteLLM_OrganizationTable) + mock_org.organization_id = "test-org-update-budget" + mock_org.models = ["gpt-4", "gpt-3.5-turbo"] + mock_org.litellm_budget_table = mock_budget_table + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( + "litellm.proxy.proxy_server.user_api_key_cache" + ) as mock_cache, patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + ), patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit, patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object", + new=AsyncMock(return_value=mock_org) + ) as mock_get_org: + + # Mock existing org-scoped team + mock_existing_team = MagicMock() + mock_existing_team.team_id = "org-team-update-budget-123" + mock_existing_team.organization_id = "test-org-update-budget" + mock_existing_team.max_budget = 30.0 + mock_existing_team.model_id = None + mock_existing_team.model_dump.return_value = { + "team_id": "org-team-update-budget-123", + "organization_id": "test-org-update-budget", + "max_budget": 30.0, + } + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) + mock_prisma.jsonify_team_object = lambda db_data: db_data + + # Mock user cache to return user with restrictive budget + mock_user_obj = LiteLLM_UserTable( + user_id="org-admin-update-budget-test", + max_budget=3.0, # Restrictive personal budget + ) + mock_cache.async_get_cache = AsyncMock(return_value=mock_user_obj) + mock_cache.async_set_cache = AsyncMock() # Mock cache set for _cache_team_object + + # Mock team update + mock_updated_team = MagicMock() + mock_updated_team.team_id = "org-team-update-budget-123" + mock_updated_team.organization_id = "test-org-update-budget" + mock_updated_team.max_budget = 50.0 + mock_updated_team.litellm_model_table = None + mock_updated_team.model_dump.return_value = { + "team_id": "org-team-update-budget-123", + "organization_id": "test-org-update-budget", + "max_budget": 50.0, + } + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_updated_team) + + # Should NOT raise an exception - bypass user budget validation for org-scoped teams + result = await update_team( + data=update_request, + http_request=dummy_request, + user_api_key_dict=org_admin_user, + ) + + # Verify the team was updated successfully with the higher budget + assert result is not None + assert result["data"].max_budget == 50.0 + + +@pytest.mark.asyncio +async def test_update_team_org_scoped_models_bypasses_user_limit(): + """ + Test that /team/update for an org-scoped team does NOT validate models against user's personal models. + + Scenario: + - Organization has models=['gpt-4', 'gpt-3.5-turbo', 'claude-3-opus'] + - User (org admin) has personal models=['no-default-models'] + - Org-scoped team exists + - User tries to update team models to ['gpt-4'] (in org's allowed, not in user's) + - Expected: Should succeed (validated against org, not user) + """ + from fastapi import Request + + from litellm.proxy._types import UpdateTeamRequest, UserAPIKeyAuth, LiteLLM_OrganizationTable + from litellm.proxy.management_endpoints.team_endpoints import update_team + + # Create user with very restrictive personal models + org_admin_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="org-admin-update-models-test", + models=["no-default-models"], # Restrictive model list + ) + + # Create update request with models in org's allowed but not in user's + update_request = UpdateTeamRequest( + team_id="org-team-update-models-123", + models=["gpt-4"], # In org's allowed, not in user's + ) + + dummy_request = MagicMock(spec=Request) + + # Mock organization with generous model list + mock_org = MagicMock(spec=LiteLLM_OrganizationTable) + mock_org.organization_id = "test-org-update-models" + mock_org.models = ["gpt-4", "gpt-3.5-turbo", "claude-3-opus"] + mock_org.litellm_budget_table = None + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( + "litellm.proxy.proxy_server.user_api_key_cache" + ) as mock_cache, patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + ), patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit, patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object", + new=AsyncMock(return_value=mock_org) + ) as mock_get_org: + + # Mock existing org-scoped team + mock_existing_team = MagicMock() + mock_existing_team.team_id = "org-team-update-models-123" + mock_existing_team.organization_id = "test-org-update-models" + mock_existing_team.models = ["gpt-3.5-turbo"] + mock_existing_team.model_id = None + mock_existing_team.model_dump.return_value = { + "team_id": "org-team-update-models-123", + "organization_id": "test-org-update-models", + "models": ["gpt-3.5-turbo"], + } + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) + mock_prisma.jsonify_team_object = lambda db_data: db_data + mock_cache.async_set_cache = AsyncMock() # Mock cache set for _cache_team_object + + # Mock team update + mock_updated_team = MagicMock() + mock_updated_team.team_id = "org-team-update-models-123" + mock_updated_team.organization_id = "test-org-update-models" + mock_updated_team.models = ["gpt-4"] + mock_updated_team.litellm_model_table = None + mock_updated_team.model_dump.return_value = { + "team_id": "org-team-update-models-123", + "organization_id": "test-org-update-models", + "models": ["gpt-4"], + } + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_updated_team) + + # Should NOT raise an exception - bypass user models validation for org-scoped teams + result = await update_team( + data=update_request, + http_request=dummy_request, + user_api_key_dict=org_admin_user, + ) + + # Verify the team was updated successfully with the new models + assert result is not None + assert result["data"].models == ["gpt-4"] + + +@pytest.mark.asyncio +async def test_update_team_org_scoped_models_not_in_org_models(): + """ + Test that /team/update for an org-scoped team fails when models are not in organization's allowed models. + + Scenario: + - Organization has models=['gpt-4', 'gpt-3.5-turbo'] + - Org-scoped team exists + - User tries to update team models to ['claude-3-opus'] (not in org's allowed models) + - Expected: Should fail with error about model not in org's allowed models + """ + from fastapi import Request + + from litellm.proxy._types import UpdateTeamRequest, ProxyException, UserAPIKeyAuth, LiteLLM_OrganizationTable + from litellm.proxy.management_endpoints.team_endpoints import update_team + + # Create user (org admin) + org_admin_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="org-admin-update-models-fail-test", + models=[], + ) + + # Create update request with model not in org's allowed list + update_request = UpdateTeamRequest( + team_id="org-team-update-models-fail-123", + models=["claude-3-opus"], # Not in org's allowed models + ) + + dummy_request = MagicMock(spec=Request) + + # Mock organization with restricted model list (no claude-3-opus) + mock_org = MagicMock(spec=LiteLLM_OrganizationTable) + mock_org.organization_id = "test-org-update-models-fail" + mock_org.models = ["gpt-4", "gpt-3.5-turbo"] # claude-3-opus is NOT allowed + mock_org.litellm_budget_table = None + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( + "litellm.proxy.proxy_server.user_api_key_cache" + ) as mock_cache, patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + ), patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit, patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object", + new=AsyncMock(return_value=mock_org) + ) as mock_get_org: + + # Mock existing org-scoped team + mock_existing_team = MagicMock() + mock_existing_team.team_id = "org-team-update-models-fail-123" + mock_existing_team.organization_id = "test-org-update-models-fail" + mock_existing_team.models = ["gpt-4"] + mock_existing_team.model_dump.return_value = { + "team_id": "org-team-update-models-fail-123", + "organization_id": "test-org-update-models-fail", + "models": ["gpt-4"], + } + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) + + # Should raise ProxyException because claude-3-opus is not in org's allowed models + with pytest.raises(ProxyException) as exc_info: + await update_team( + data=update_request, + http_request=dummy_request, + user_api_key_dict=org_admin_user, + ) + + # Verify exception details + assert exc_info.value.code == '400' + assert "claude-3-opus" in str(exc_info.value.message) or "organization" in str(exc_info.value.message).lower() + + +@pytest.mark.asyncio +async def test_update_team_tpm_limit_exceeds_user_limit(): + """ + Test that /team/update fails when TPM limit exceeds user's TPM limit. + + Scenario: + - User has tpm_limit=1000 + - User tries to update team with tpm_limit=5000 + - Expected: Should fail with error about exceeding user TPM limit + """ + from fastapi import Request + + from litellm.proxy._types import UpdateTeamRequest, ProxyException, UserAPIKeyAuth + from litellm.proxy.management_endpoints.team_endpoints import update_team + + # Create non-admin user with TPM limit + non_admin_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="tpm-limit-user", + models=[], + tpm_limit=1000, # User's TPM limit + ) + + # Create update request with TPM exceeding user's limit + update_request = UpdateTeamRequest( + team_id="team-tpm-test-123", + tpm_limit=5000, # Exceeds user's 1000 limit + ) + + dummy_request = MagicMock(spec=Request) + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( + "litellm.proxy.proxy_server.user_api_key_cache" + ) as mock_cache, patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + ): + + # Mock existing standalone team + mock_existing_team = MagicMock() + mock_existing_team.team_id = "team-tpm-test-123" + mock_existing_team.organization_id = None + mock_existing_team.tpm_limit = 500 + mock_existing_team.model_dump.return_value = { + "team_id": "team-tpm-test-123", + "organization_id": None, + "tpm_limit": 500, + } + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) + + # Should raise ProxyException because new TPM exceeds user's limit + with pytest.raises(ProxyException) as exc_info: + await update_team( + data=update_request, + http_request=dummy_request, + user_api_key_dict=non_admin_user, + ) + + # Verify exception details + assert exc_info.value.code == '400' + assert "tpm" in str(exc_info.value.message).lower() + + +@pytest.mark.asyncio +async def test_update_team_rpm_limit_exceeds_user_limit(): + """ + Test that /team/update fails when RPM limit exceeds user's RPM limit. + + Scenario: + - User has rpm_limit=100 + - User tries to update team with rpm_limit=500 + - Expected: Should fail with error about exceeding user RPM limit + """ + from fastapi import Request + + from litellm.proxy._types import UpdateTeamRequest, ProxyException, UserAPIKeyAuth + from litellm.proxy.management_endpoints.team_endpoints import update_team + + # Create non-admin user with RPM limit + non_admin_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="rpm-limit-user", + models=[], + rpm_limit=100, # User's RPM limit + ) + + # Create update request with RPM exceeding user's limit + update_request = UpdateTeamRequest( + team_id="team-rpm-test-123", + rpm_limit=500, # Exceeds user's 100 limit + ) + + dummy_request = MagicMock(spec=Request) + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( + "litellm.proxy.proxy_server.user_api_key_cache" + ) as mock_cache, patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + ): + + # Mock existing standalone team + mock_existing_team = MagicMock() + mock_existing_team.team_id = "team-rpm-test-123" + mock_existing_team.organization_id = None + mock_existing_team.rpm_limit = 50 + mock_existing_team.model_dump.return_value = { + "team_id": "team-rpm-test-123", + "organization_id": None, + "rpm_limit": 50, + } + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) + + # Should raise ProxyException because new RPM exceeds user's limit + with pytest.raises(ProxyException) as exc_info: + await update_team( + data=update_request, + http_request=dummy_request, + user_api_key_dict=non_admin_user, + ) + + # Verify exception details + assert exc_info.value.code == '400' + assert "rpm" in str(exc_info.value.message).lower() + + +@pytest.mark.asyncio +async def test_new_team_org_scoped_tpm_exceeds_org_limit(): + """ + Test that /team/new for an org-scoped team fails when TPM exceeds organization's TPM limit. + + Scenario: + - Organization has tpm_limit=10000 + - User tries to create org-scoped team with tpm_limit=20000 + - Expected: Should fail with error about exceeding org TPM limit + """ + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest, ProxyException, UserAPIKeyAuth, LiteLLM_OrganizationTable, LiteLLM_BudgetTable + from litellm.proxy.management_endpoints.team_endpoints import new_team + + # Create user (with restrictive personal TPM limit that should be bypassed) + org_admin_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="org-admin-tpm-test", + models=[], + tpm_limit=1000, # User's personal limit (should be bypassed for org teams) + ) + + # Create team request with TPM exceeding org's limit + team_request = NewTeamRequest( + team_alias="org-tpm-test-team", + organization_id="test-org-tpm", + tpm_limit=20000, # Exceeds org's 10000 limit + ) + + dummy_request = MagicMock(spec=Request) + + # Mock organization with TPM limit + mock_budget_table = MagicMock(spec=LiteLLM_BudgetTable) + mock_budget_table.tpm_limit = 10000 # Org's TPM limit + mock_budget_table.rpm_limit = None + mock_budget_table.max_budget = None + + mock_org = MagicMock(spec=LiteLLM_OrganizationTable) + mock_org.organization_id = "test-org-tpm" + mock_org.models = ["gpt-4"] + mock_org.litellm_budget_table = mock_budget_table + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( + "litellm.proxy.proxy_server.user_api_key_cache" + ) as mock_cache, patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + ), patch( + "litellm.proxy.proxy_server._license_check" + ) as mock_license, patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object", + new=AsyncMock(return_value=mock_org) + ): + mock_license.is_team_count_over_limit.return_value = False + mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) + mock_prisma.get_data = AsyncMock(return_value=None) + + # Should raise ProxyException because TPM exceeds org limit + with pytest.raises(ProxyException) as exc_info: + await new_team( + data=team_request, + http_request=dummy_request, + user_api_key_dict=org_admin_user, + ) + + # Verify exception details + assert exc_info.value.code == '400' + assert "tpm" in str(exc_info.value.message).lower() + + +@pytest.mark.asyncio +async def test_new_team_org_scoped_rpm_exceeds_org_limit(): + """ + Test that /team/new for an org-scoped team fails when RPM exceeds organization's RPM limit. + + Scenario: + - Organization has rpm_limit=1000 + - User tries to create org-scoped team with rpm_limit=2000 + - Expected: Should fail with error about exceeding org RPM limit + """ + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest, ProxyException, UserAPIKeyAuth, LiteLLM_OrganizationTable, LiteLLM_BudgetTable + from litellm.proxy.management_endpoints.team_endpoints import new_team + + # Create user (with restrictive personal RPM limit that should be bypassed) + org_admin_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="org-admin-rpm-test", + models=[], + rpm_limit=100, # User's personal limit (should be bypassed for org teams) + ) + + # Create team request with RPM exceeding org's limit + team_request = NewTeamRequest( + team_alias="org-rpm-test-team", + organization_id="test-org-rpm", + rpm_limit=2000, # Exceeds org's 1000 limit + ) + + dummy_request = MagicMock(spec=Request) + + # Mock organization with RPM limit + mock_budget_table = MagicMock(spec=LiteLLM_BudgetTable) + mock_budget_table.tpm_limit = None + mock_budget_table.rpm_limit = 1000 # Org's RPM limit + mock_budget_table.max_budget = None + + mock_org = MagicMock(spec=LiteLLM_OrganizationTable) + mock_org.organization_id = "test-org-rpm" + mock_org.models = ["gpt-4"] + mock_org.litellm_budget_table = mock_budget_table + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( + "litellm.proxy.proxy_server.user_api_key_cache" + ) as mock_cache, patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + ), patch( + "litellm.proxy.proxy_server._license_check" + ) as mock_license, patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object", + new=AsyncMock(return_value=mock_org) + ): + mock_license.is_team_count_over_limit.return_value = False + mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) + mock_prisma.get_data = AsyncMock(return_value=None) + + # Should raise ProxyException because RPM exceeds org limit + with pytest.raises(ProxyException) as exc_info: + await new_team( + data=team_request, + http_request=dummy_request, + user_api_key_dict=org_admin_user, + ) + + # Verify exception details + assert exc_info.value.code == '400' + assert "rpm" in str(exc_info.value.message).lower() + + +@pytest.mark.asyncio +async def test_new_team_org_scoped_tpm_rpm_bypasses_user_limit(): + """ + Test that /team/new for an org-scoped team bypasses user's TPM/RPM limits. + + Scenario: + - User has tpm_limit=1000, rpm_limit=100 + - Organization has tpm_limit=50000, rpm_limit=5000 + - User creates org-scoped team with tpm_limit=10000, rpm_limit=1000 + - Expected: Should succeed (bypasses user limits, within org limits) + """ + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest, UserAPIKeyAuth, LiteLLM_OrganizationTable, LiteLLM_BudgetTable, LiteLLM_TeamTable + from litellm.proxy.management_endpoints.team_endpoints import new_team + + # Create user with restrictive personal limits + org_admin_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="org-admin-bypass-test", + models=[], + tpm_limit=1000, # Restrictive user TPM limit + rpm_limit=100, # Restrictive user RPM limit + ) + + # Create team request exceeding user limits but within org limits + team_request = NewTeamRequest( + team_alias="org-bypass-test-team", + organization_id="test-org-bypass", + tpm_limit=10000, # Exceeds user's 1000 but within org's 50000 + rpm_limit=1000, # Exceeds user's 100 but within org's 5000 + ) + + dummy_request = MagicMock(spec=Request) + + # Mock organization with generous limits + mock_budget_table = MagicMock(spec=LiteLLM_BudgetTable) + mock_budget_table.tpm_limit = 50000 # Generous org TPM limit + mock_budget_table.rpm_limit = 5000 # Generous org RPM limit + mock_budget_table.max_budget = None + + mock_org = MagicMock(spec=LiteLLM_OrganizationTable) + mock_org.organization_id = "test-org-bypass" + mock_org.models = ["gpt-4"] + mock_org.litellm_budget_table = mock_budget_table + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( + "litellm.proxy.proxy_server.user_api_key_cache" + ) as mock_cache, patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + ), patch( + "litellm.proxy.proxy_server._license_check" + ) as mock_license, patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ), patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object", + new=AsyncMock(return_value=mock_org) + ), patch( + "litellm.proxy.management_endpoints.team_endpoints._add_team_members_to_team", + new=AsyncMock() + ): + mock_license.is_team_count_over_limit.return_value = False + mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) + mock_prisma.get_data = AsyncMock(return_value=None) + + # Mock team creation + mock_created_team = MagicMock(spec=LiteLLM_TeamTable) + mock_created_team.team_id = "new-bypass-team-id" + mock_created_team.team_alias = "org-bypass-test-team" + mock_created_team.tpm_limit = 10000 + mock_created_team.rpm_limit = 1000 + mock_created_team.metadata = None + mock_created_team.members_with_roles = [] + mock_created_team.model_dump.return_value = { + "team_id": "new-bypass-team-id", + "team_alias": "org-bypass-test-team", + "tpm_limit": 10000, + "rpm_limit": 1000, + "metadata": None, + "members_with_roles": [], + } + mock_prisma.db.litellm_teamtable.create = AsyncMock(return_value=mock_created_team) + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_created_team) + mock_prisma.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) + + # Should succeed - bypasses user limits since org-scoped + result = await new_team( + data=team_request, + http_request=dummy_request, + user_api_key_dict=org_admin_user, + ) + + # Verify team was created + assert result["team_id"] == "new-bypass-team-id" + + +@pytest.mark.asyncio +async def test_update_team_org_scoped_tpm_exceeds_org_limit(): + """ + Test that /team/update for an org-scoped team fails when TPM exceeds organization's TPM limit. + + Scenario: + - Organization has tpm_limit=10000 + - User tries to update org-scoped team with tpm_limit=20000 + - Expected: Should fail with error about exceeding org TPM limit + """ + from fastapi import Request + + from litellm.proxy._types import UpdateTeamRequest, ProxyException, UserAPIKeyAuth, LiteLLM_OrganizationTable, LiteLLM_BudgetTable + from litellm.proxy.management_endpoints.team_endpoints import update_team + + # Create user (with restrictive personal TPM limit that should be bypassed) + org_admin_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="org-admin-update-tpm-test", + models=[], + tpm_limit=1000, # User's personal limit (should be bypassed for org teams) + ) + + # Create update request with TPM exceeding org's limit + update_request = UpdateTeamRequest( + team_id="org-team-update-tpm-123", + tpm_limit=20000, # Exceeds org's 10000 limit + ) + + dummy_request = MagicMock(spec=Request) + + # Mock organization with TPM limit + mock_budget_table = MagicMock(spec=LiteLLM_BudgetTable) + mock_budget_table.tpm_limit = 10000 # Org's TPM limit + mock_budget_table.rpm_limit = None + mock_budget_table.max_budget = None + + mock_org = MagicMock(spec=LiteLLM_OrganizationTable) + mock_org.organization_id = "test-org-update-tpm" + mock_org.models = ["gpt-4"] + mock_org.litellm_budget_table = mock_budget_table + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( + "litellm.proxy.proxy_server.user_api_key_cache" + ) as mock_cache, patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + ), patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object", + new=AsyncMock(return_value=mock_org) + ): + + # Mock existing org-scoped team + mock_existing_team = MagicMock() + mock_existing_team.team_id = "org-team-update-tpm-123" + mock_existing_team.organization_id = "test-org-update-tpm" + mock_existing_team.tpm_limit = 5000 + mock_existing_team.model_dump.return_value = { + "team_id": "org-team-update-tpm-123", + "organization_id": "test-org-update-tpm", + "tpm_limit": 5000, + } + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) + + # Should raise ProxyException because TPM exceeds org limit + with pytest.raises(ProxyException) as exc_info: + await update_team( + data=update_request, + http_request=dummy_request, + user_api_key_dict=org_admin_user, + ) + + # Verify exception details + assert exc_info.value.code == '400' + assert "tpm" in str(exc_info.value.message).lower() + + +@pytest.mark.asyncio +async def test_update_team_org_scoped_rpm_exceeds_org_limit(): + """ + Test that /team/update for an org-scoped team fails when RPM exceeds organization's RPM limit. + + Scenario: + - Organization has rpm_limit=1000 + - User tries to update org-scoped team with rpm_limit=2000 + - Expected: Should fail with error about exceeding org RPM limit + """ + from fastapi import Request + + from litellm.proxy._types import UpdateTeamRequest, ProxyException, UserAPIKeyAuth, LiteLLM_OrganizationTable, LiteLLM_BudgetTable + from litellm.proxy.management_endpoints.team_endpoints import update_team + + # Create user (with restrictive personal RPM limit that should be bypassed) + org_admin_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="org-admin-update-rpm-test", + models=[], + rpm_limit=100, # User's personal limit (should be bypassed for org teams) + ) + + # Create update request with RPM exceeding org's limit + update_request = UpdateTeamRequest( + team_id="org-team-update-rpm-123", + rpm_limit=2000, # Exceeds org's 1000 limit + ) + + dummy_request = MagicMock(spec=Request) + + # Mock organization with RPM limit + mock_budget_table = MagicMock(spec=LiteLLM_BudgetTable) + mock_budget_table.tpm_limit = None + mock_budget_table.rpm_limit = 1000 # Org's RPM limit + mock_budget_table.max_budget = None + + mock_org = MagicMock(spec=LiteLLM_OrganizationTable) + mock_org.organization_id = "test-org-update-rpm" + mock_org.models = ["gpt-4"] + mock_org.litellm_budget_table = mock_budget_table + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( + "litellm.proxy.proxy_server.user_api_key_cache" + ) as mock_cache, patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + ), patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object", + new=AsyncMock(return_value=mock_org) + ): + + # Mock existing org-scoped team + mock_existing_team = MagicMock() + mock_existing_team.team_id = "org-team-update-rpm-123" + mock_existing_team.organization_id = "test-org-update-rpm" + mock_existing_team.rpm_limit = 500 + mock_existing_team.model_dump.return_value = { + "team_id": "org-team-update-rpm-123", + "organization_id": "test-org-update-rpm", + "rpm_limit": 500, + } + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) + + # Should raise ProxyException because RPM exceeds org limit + with pytest.raises(ProxyException) as exc_info: + await update_team( + data=update_request, + http_request=dummy_request, + user_api_key_dict=org_admin_user, + ) + + # Verify exception details + assert exc_info.value.code == '400' + assert "rpm" in str(exc_info.value.message).lower() + + +@pytest.mark.asyncio +async def test_update_team_org_scoped_tpm_rpm_bypasses_user_limit(): + """ + Test that /team/update for an org-scoped team bypasses user's TPM/RPM limits. + + Scenario: + - User has tpm_limit=1000, rpm_limit=100 + - Organization has tpm_limit=50000, rpm_limit=5000 + - User updates org-scoped team with tpm_limit=10000, rpm_limit=1000 + - Expected: Should succeed (bypasses user limits, within org limits) + """ + from fastapi import Request + + from litellm.proxy._types import UpdateTeamRequest, UserAPIKeyAuth, LiteLLM_OrganizationTable, LiteLLM_BudgetTable, LiteLLM_TeamTable + from litellm.proxy.management_endpoints.team_endpoints import update_team + + # Create user with restrictive personal limits + org_admin_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="org-admin-update-bypass-test", + models=[], + tpm_limit=1000, # Restrictive user TPM limit + rpm_limit=100, # Restrictive user RPM limit + ) + + # Create update request exceeding user limits but within org limits + update_request = UpdateTeamRequest( + team_id="org-team-update-bypass-123", + tpm_limit=10000, # Exceeds user's 1000 but within org's 50000 + rpm_limit=1000, # Exceeds user's 100 but within org's 5000 + ) + + dummy_request = MagicMock(spec=Request) + + # Mock organization with generous limits + mock_budget_table = MagicMock(spec=LiteLLM_BudgetTable) + mock_budget_table.tpm_limit = 50000 # Generous org TPM limit + mock_budget_table.rpm_limit = 5000 # Generous org RPM limit + mock_budget_table.max_budget = None + + mock_org = MagicMock(spec=LiteLLM_OrganizationTable) + mock_org.organization_id = "test-org-update-bypass" + mock_org.models = ["gpt-4"] + mock_org.litellm_budget_table = mock_budget_table + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( + "litellm.proxy.proxy_server.user_api_key_cache" + ) as mock_cache, patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + ), patch( + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_logging, patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object", + new=AsyncMock(return_value=mock_org) + ): + + # Mock existing org-scoped team + mock_existing_team = MagicMock() + mock_existing_team.team_id = "org-team-update-bypass-123" + mock_existing_team.organization_id = "test-org-update-bypass" + mock_existing_team.tpm_limit = 5000 + mock_existing_team.rpm_limit = 500 + mock_existing_team.model_id = None + mock_existing_team.model_dump.return_value = { + "team_id": "org-team-update-bypass-123", + "organization_id": "test-org-update-bypass", + "tpm_limit": 5000, + "rpm_limit": 500, + } + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) + mock_cache.async_set_cache = AsyncMock() + + # Mock team update + mock_updated_team = MagicMock(spec=LiteLLM_TeamTable) + mock_updated_team.team_id = "org-team-update-bypass-123" + mock_updated_team.tpm_limit = 10000 + mock_updated_team.rpm_limit = 1000 + mock_updated_team.model_dump.return_value = { + "team_id": "org-team-update-bypass-123", + "tpm_limit": 10000, + "rpm_limit": 1000, + } + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_updated_team) + mock_prisma.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) + + # Should succeed - bypasses user limits since org-scoped + result = await update_team( + data=update_request, + http_request=dummy_request, + user_api_key_dict=org_admin_user, + ) + + # Verify team was updated + assert result["team_id"] == "org-team-update-bypass-123" \ No newline at end of file