diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 6d4faae5fd8c..b697e01a6ef8 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) + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) - existing_team_row = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": data.team_id} - ) + 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) - if existing_team_row is None: - raise HTTPException( - status_code=404, - detail={"error": f"Team not found, passed team_id={data.team_id}"}, + existing_team_row = await prisma_client.db.litellm_teamtable.find_unique( + where={"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 + if existing_team_row is None: + raise HTTPException( + status_code=404, + detail={"error": f"Team not found, passed team_id={data.team_id}"}, + ) - # 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 ( + 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_table is not None: - await _check_org_team_limits( - org_table=org_table, - data=data, + 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, + ) - updated_kv = data.json(exclude_unset=True) + # 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, + ) - # 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 + updated_kv = data.json(exclude_unset=True) - reset_at = get_budget_reset_time(budget_duration=data.budget_duration) + # 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 - # set the budget_reset_at in DB - updated_kv["budget_reset_at"] = reset_at + reset_at = get_budget_reset_time(budget_duration=data.budget_duration) - 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, + # 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) - 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, + # 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, ) - 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 + # 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, + ) - 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 - ) - ) + 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 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 + ) ) - 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, - ) + 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)}, + ) - # 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, + 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, ) - return {"team_id": team_row.team_id, "data": team_row} + # 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} + 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 86b23c98ba55..06ec71a84f84 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