diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 7fff640c497..39857984247 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -690,6 +690,13 @@ class LiteLLMRoutes(enum.Enum): "/organization/delete", "/organization/member_add", "/organization/member_update", + # member_delete is equally destructive as member_add / member_update + # and must be scoped the same way — otherwise it falls through to + # the management_routes / self_managed_routes path and lets any + # non-PROXY_ADMIN caller that reaches the route delete arbitrary + # org memberships without the organization_role_based_access_check + # that member_add / member_update trigger. + "/organization/member_delete", ] # Routes accessible by Admin Viewer (read-only admin access) diff --git a/litellm/proxy/auth/auth_checks_organization.py b/litellm/proxy/auth/auth_checks_organization.py index 50efe137209..d89afcffa9a 100644 --- a/litellm/proxy/auth/auth_checks_organization.py +++ b/litellm/proxy/auth/auth_checks_organization.py @@ -144,7 +144,7 @@ def _user_is_org_admin( user_object: Optional[LiteLLM_UserTable] = None, ) -> bool: """ - Helper function to check if user is an org admin for any of the passed organizations. + Helper function to check if user is an org admin for all of the passed organizations. Checks both: - `organization_id` (singular string) — legacy callers @@ -168,9 +168,13 @@ def _user_is_org_admin( if not candidate_org_ids: return False - for _membership in user_object.organization_memberships: - if _membership.organization_id in candidate_org_ids: - if _membership.user_role == LitellmUserRoles.ORG_ADMIN.value: - return True + # Build set of orgs where user is admin + admin_org_ids = { + _membership.organization_id + for _membership in user_object.organization_memberships + if _membership.user_role == LitellmUserRoles.ORG_ADMIN.value + and _membership.organization_id is not None + } - return False + # User must be admin of ALL requested orgs, not just any one + return all(org_id in admin_org_ids for org_id in candidate_org_ids) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 1772da3d15e..e343429ff12 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -1180,6 +1180,23 @@ async def _update_single_user_helper( "error": "User does not have permission to update this user. Only PROXY_ADMIN can update other users." }, ) + else: + # Silent-create guard: if the target user doesn't exist, the update + # path falls through to an upsert that creates a new user with + # caller-supplied fields (models, metadata, budgets, …). Only + # PROXY_ADMIN is allowed to create users this way; otherwise an org + # admin could spawn arbitrary users attached to nothing by supplying + # a fresh email, bypassing the /user/new org/team-scoping checks. + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: + raise HTTPException( + status_code=404, + detail={ + "error": ( + "User not found. Only PROXY_ADMIN can create users " + "via /user/update; use /user/new instead." + ) + }, + ) existing_metadata = ( cast(Dict, getattr(existing_user_row, "metadata", {}) or {}) @@ -2057,6 +2074,54 @@ async def delete_user( if data.user_ids is None: raise HTTPException(status_code=400, detail={"error": "No user id passed in"}) + # Per-target authorization: the route-level gate accepts this call when + # the caller is PROXY_ADMIN or an ORG_ADMIN of *any* org named in + # request_data["organization_id"]/["organizations"]. That gate does NOT + # cross-check data.user_ids against the caller's scope, so without this + # loop an org-admin of org-A could delete users in org-B by supplying + # {"user_ids": [victim_in_org_B], "organization_id": "org-A"}. + caller_is_proxy_admin = ( + user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value + ) + caller_admin_org_ids: set = set() + if not caller_is_proxy_admin: + caller_memberships = ( + await prisma_client.db.litellm_organizationmembership.find_many( + where={ + "user_id": user_api_key_dict.user_id, + "user_role": LitellmUserRoles.ORG_ADMIN.value, + } + ) + if user_api_key_dict.user_id + else [] + ) + caller_admin_org_ids = { + m.organization_id for m in caller_memberships if m.organization_id + } + if not caller_admin_org_ids: + raise HTTPException( + status_code=403, + detail={ + "error": "Only PROXY_ADMIN or ORG_ADMIN users may delete users." + }, + ) + + # Batch-fetch target memberships once before the per-user loop. Avoids + # an N+1 DB call when delete_user is called with a large user_ids list. + target_org_ids_by_user: Dict[str, set] = {} + if not caller_is_proxy_admin: + all_target_memberships = ( + await prisma_client.db.litellm_organizationmembership.find_many( + where={"user_id": {"in": data.user_ids}} + ) + ) + for m in all_target_memberships: + if not m.organization_id: + continue + target_org_ids_by_user.setdefault(m.user_id, set()).add( + m.organization_id + ) + # check that all teams passed exist for user_id in data.user_ids: user_row = await prisma_client.db.litellm_usertable.find_unique( @@ -2068,30 +2133,49 @@ async def delete_user( status_code=404, detail={"error": f"User not found, passed user_id={user_id}"}, ) - else: - # Enterprise Feature - Audit Logging. Enable with litellm.store_audit_logs = True - # we do this after the first for loop, since first for loop is for validation. we only want this inserted after validation passes - if litellm.store_audit_logs is True: - # make an audit log for each team deleted - _user_row = user_row.json(exclude_none=True) - asyncio.create_task( - create_audit_log_for_update( - request_data=LiteLLM_AuditLogs( - id=str(uuid.uuid4()), - updated_at=datetime.now(timezone.utc), - changed_by=litellm_changed_by - or user_api_key_dict.user_id - or litellm_proxy_admin_name, - changed_by_api_key=user_api_key_dict.api_key, - table_name=LitellmTableNames.USER_TABLE_NAME, - object_id=user_id, - action="deleted", - updated_values="{}", - before_value=_user_row, + if not caller_is_proxy_admin: + target_org_ids = target_org_ids_by_user.get(user_id, set()) + # Org-admin may only delete users whose entire org membership is + # within their admin scope. A target with ANY org outside the + # caller's scope (or no org at all) requires PROXY_ADMIN. + if not target_org_ids or not target_org_ids.issubset( + caller_admin_org_ids + ): + raise HTTPException( + status_code=403, + detail={ + "error": ( + f"User {user_id} is not within your admin scope. " + "Only PROXY_ADMIN may delete users outside your " + "administered organizations." ) + }, + ) + + # Enterprise Feature - Audit Logging. Enable with litellm.store_audit_logs = True + # we do this after the first for loop, since first for loop is for validation. we only want this inserted after validation passes + if litellm.store_audit_logs is True: + # make an audit log for each team deleted + _user_row = user_row.json(exclude_none=True) + + asyncio.create_task( + create_audit_log_for_update( + request_data=LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=litellm_changed_by + or user_api_key_dict.user_id + or litellm_proxy_admin_name, + changed_by_api_key=user_api_key_dict.api_key, + table_name=LitellmTableNames.USER_TABLE_NAME, + object_id=user_id, + action="deleted", + updated_values="{}", + before_value=_user_row, ) ) + ) ## CLEANUP MEMBERS_WITH_ROLES fetch_all_teams = await prisma_client.db.litellm_teamtable.find_many( diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index f69d9d2f8d4..32227b3dd6d 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1961,22 +1961,48 @@ async def _validate_update_key_data( user_api_key_cache=user_api_key_cache, ) - # Admin-only: only proxy admins, team admins, or org admins can modify max_budget or spend - if ( - data.max_budget is not None and data.max_budget != existing_key_row.max_budget + # Cross-key authorization. Previously only gated on max_budget/spend + # changes, which let a non-admin blanket-rewrite any OTHER field on + # any key (models, alias, metadata, tpm_limit, rpm_limit, + # allowed_routes, guardrails, blocked, duration, permissions, …) as + # long as they avoided budget/spend. + # + # Policy: + # - Key owner (same user_id): may update non-budget fields on their + # own key without the admin check. + # - Anyone else (non-PROXY_ADMIN, not the owner): must pass + # _check_key_admin_access (PROXY_ADMIN / key-owner / team-admin / + # org-admin of the key). + # - max_budget / spend: always require the admin check, even for the + # key owner (matches the existing admin-only budget semantics). + is_key_owner = ( + user_api_key_dict.user_id is not None + and existing_key_row.user_id == user_api_key_dict.user_id + ) + _is_budget_change = ( + data.max_budget is not None + and data.max_budget != existing_key_row.max_budget ) or ( data.spend is not None and data.spend != getattr(existing_key_row, "spend", None) + ) + if ( + (not _is_proxy_admin) + and prisma_client is not None + and (not is_key_owner or _is_budget_change) ): - if prisma_client is not None: - hashed_key = existing_key_row.token - await _check_key_admin_access( - user_api_key_dict=user_api_key_dict, - hashed_token=hashed_key, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - route="/key/update (max_budget/spend)", - ) + hashed_key = existing_key_row.token + await _check_key_admin_access( + user_api_key_dict=user_api_key_dict, + hashed_token=hashed_key, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + route=( + "/key/update (max_budget/spend)" + if _is_budget_change + else "/key/update" + ), + ) # Check team limits if key has a team_id (from request or existing key) team_obj: Optional[LiteLLM_TeamTableCachedObj] = None diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 25df9f0b0f7..670946ac0d8 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -1064,6 +1064,33 @@ async def organization_member_update( }, ) + # Reject attempts to change the role of a global PROXY_ADMIN via + # org-scoped operations. An org-admin of any org could otherwise + # alter a PROXY_ADMIN user's per-org role, which has downstream + # effects on admin UI filtering and scope derivation. + target_user_row = await prisma_client.db.litellm_usertable.find_unique( + where={"user_id": data.user_id} + ) + if target_user_row is not None and getattr( + target_user_row, "user_role", None + ) in ( + LitellmUserRoles.PROXY_ADMIN.value, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + ): + if ( + user_api_key_dict.user_role + != LitellmUserRoles.PROXY_ADMIN.value + ): + raise HTTPException( + status_code=403, + detail={ + "error": ( + "Only PROXY_ADMIN may modify the organization " + "role of a user who is a global PROXY_ADMIN." + ) + }, + ) + # Update member role if data.role is not None: await prisma_client.db.litellm_organizationmembership.update( diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 138469312e1..0bb6c3d90f3 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -120,6 +120,7 @@ def _sanitize_for_log(value: Any) -> str: text = repr(value) return text.replace("\r", "").replace("\n", "") + async def _verify_team_access( team_obj: LiteLLM_TeamTable, user_api_key_dict: UserAPIKeyAuth, @@ -314,10 +315,8 @@ class TeamMemberBudgetHandler: return # Batch-fetch existing memberships for this team (avoids N+1 queries) - existing_memberships = ( - await prisma_client.db.litellm_teammembership.find_many( - where={"team_id": team_id} - ) + existing_memberships = await prisma_client.db.litellm_teammembership.find_many( + where={"team_id": team_id} ) existing_user_ids = {m.user_id for m in existing_memberships} @@ -1545,6 +1544,42 @@ async def update_team( # noqa: PLR0915 if ( data.organization_id is not None and len(data.organization_id) > 0 ): # allow unsetting the organization_id + # If the caller is relocating the team to a different org, they + # must also be PROXY_ADMIN or an org-admin of the DESTINATION org. + # _verify_team_access above only checked the team's CURRENT org, + # so without this gate an org-admin could hand their team to any + # other org (or capture a team from another org they once + # administered into a new destination). + current_org_id = getattr(existing_team_row, "organization_id", None) + if ( + data.organization_id != current_org_id + and user_api_key_dict.user_role + != LitellmUserRoles.PROXY_ADMIN.value + ): + # Is the caller org_admin of the destination org? + caller_memberships = ( + await prisma_client.db.litellm_organizationmembership.find_many( + where={ + "user_id": user_api_key_dict.user_id, + "organization_id": data.organization_id, + "user_role": LitellmUserRoles.ORG_ADMIN.value, + } + ) + if user_api_key_dict.user_id + else [] + ) + if not caller_memberships: + raise HTTPException( + status_code=403, + detail={ + "error": ( + "Relocating a team to a different organization " + "requires PROXY_ADMIN or org-admin of the " + "destination org." + ) + }, + ) + await fetch_and_validate_organization( organization_id=data.organization_id, existing_team_row=existing_team_row, @@ -1659,12 +1694,12 @@ async def update_team( # noqa: PLR0915 updated_kv["router_settings"] = safe_dumps(updated_kv["router_settings"]) 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 + 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: @@ -2411,13 +2446,13 @@ async def team_member_delete( ) # Fetch keys before deletion to persist them - keys_to_delete: List[ - LiteLLM_VerificationToken - ] = await prisma_client.db.litellm_verificationtoken.find_many( - where={ - "user_id": {"in": list(user_ids_to_delete)}, - "team_id": data.team_id, - } + keys_to_delete: List[LiteLLM_VerificationToken] = ( + await prisma_client.db.litellm_verificationtoken.find_many( + where={ + "user_id": {"in": list(user_ids_to_delete)}, + "team_id": data.team_id, + } + ) ) if keys_to_delete: @@ -2683,6 +2718,20 @@ async def bulk_team_member_add( ) if data.all_users: + # `all_users=True` pulls every user in the database into this team, + # regardless of org. Any team admin could use it to capture every + # user across every org into a team they control. Restrict to + # PROXY_ADMIN. + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: + raise HTTPException( + status_code=403, + detail={ + "error": ( + "`all_users=true` is restricted to PROXY_ADMIN. " + "Org/team admins must specify explicit member lists." + ) + }, + ) # get all users from the database all_users_in_db = await prisma_client.db.litellm_usertable.find_many( order={"created_at": "desc"} @@ -2801,10 +2850,10 @@ async def delete_team( team_rows: List[LiteLLM_TeamTable] = [] for team_id in data.team_ids: try: - team_row_base: Optional[ - BaseModel - ] = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id} + team_row_base: Optional[BaseModel] = ( + await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id} + ) ) if team_row_base is None: raise Exception @@ -2870,10 +2919,10 @@ async def delete_team( _persist_deleted_verification_tokens, ) - keys_to_delete: List[ - LiteLLM_VerificationToken - ] = await prisma_client.db.litellm_verificationtoken.find_many( - where={"team_id": {"in": data.team_ids}} + keys_to_delete: List[LiteLLM_VerificationToken] = ( + await prisma_client.db.litellm_verificationtoken.find_many( + where={"team_id": {"in": data.team_ids}} + ) ) if keys_to_delete: @@ -3110,11 +3159,11 @@ async def team_info( ) try: - team_info: Optional[ - BaseModel - ] = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id}, - include={"object_permission": True}, + team_info: Optional[BaseModel] = ( + await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id}, + include={"object_permission": True}, + ) ) if team_info is None: raise Exception @@ -3405,6 +3454,7 @@ async def _get_org_admin_org_ids( m.organization_id for m in (caller_user.organization_memberships or []) if m.user_role == LitellmUserRoles.ORG_ADMIN.value + and m.organization_id is not None ] return org_ids if org_ids else None @@ -3439,13 +3489,8 @@ async def _build_team_list_where_conditions( if organization_id: where_conditions["organization_id"] = organization_id - elif org_admin_org_ids is not None and not user_id: - # Org admin without explicit org or user filter: scope to their orgs. - # NOTE: when user_id is provided, no org filter is applied — the - # query returns all teams the target user belongs to across all - # organisations. This matches the legacy /team/list behaviour in - # _authorize_and_filter_teams which fetches direct-membership teams - # without an org constraint. + elif org_admin_org_ids is not None: + # Org admin: always scope to their orgs, even when filtering by user_id. where_conditions["organization_id"] = {"in": org_admin_org_ids} if user_id: @@ -3815,7 +3860,7 @@ async def _authorize_and_filter_teams( Authorize the /team/list request and return filtered teams. - Proxy admins: all teams (or filtered by user_id if provided). - - Org admins: teams from their orgs + teams they are direct members of. + - Org admins: teams from their orgs (scoped to user_id if provided). - Own query (user_id matches caller): teams the user is a member of. - Others: 401. """ @@ -3843,6 +3888,7 @@ async def _authorize_and_filter_teams( m.organization_id for m in (caller_user.organization_memberships or []) if m.user_role == LitellmUserRoles.ORG_ADMIN.value + and m.organization_id is not None ] if not allowed_org_ids: allowed_org_ids = None @@ -3865,20 +3911,13 @@ async def _authorize_and_filter_teams( ) if not user_id: return list(org_teams) - # Also include teams the user is a direct member of (outside their orgs) - seen_team_ids = {team.team_id for team in org_teams} - all_teams = list(org_teams) - # Prisma doesn't support filtering JSON array fields, so we fetch by membership separately - member_teams = await prisma_client.db.litellm_teamtable.find_many( - where={"team_id": {"not_in": list(seen_team_ids)}} if seen_team_ids else {}, - include={"litellm_model_table": True}, - ) - for team in member_teams: - if team.members_with_roles and any( - m.get("user_id") == user_id for m in team.members_with_roles - ): - all_teams.append(team) - return all_teams + # Filter org teams to only those where the target user is a member + return [ + team + for team in org_teams + if team.members_with_roles + and any(m.get("user_id") == user_id for m in team.members_with_roles) + ] elif user_id: # Regular user: fetch all and filter by membership (Prisma can't filter JSON arrays) response = await prisma_client.db.litellm_teamtable.find_many( diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index f1344a302d7..bca6b9e78d9 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -1058,8 +1058,15 @@ class TestModelsRouteExemptFromDisableLLMEndpoints: local_file = os.path.join( os.path.dirname(__file__), - "..", "..", "..", "..", "enterprise", - "litellm_enterprise", "proxy", "auth", "route_checks.py", + "..", + "..", + "..", + "..", + "enterprise", + "litellm_enterprise", + "proxy", + "auth", + "route_checks.py", ) local_file = os.path.abspath(local_file) @@ -1075,10 +1082,15 @@ class TestModelsRouteExemptFromDisableLLMEndpoints: """Test that /models is allowed even when LLM API routes are disabled""" EnterpriseRouteChecks = self._get_enterprise_route_checks() - with patch.object( - EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True - ), patch.object( - EnterpriseRouteChecks, "is_management_routes_disabled", return_value=False + with ( + patch.object( + EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True + ), + patch.object( + EnterpriseRouteChecks, + "is_management_routes_disabled", + return_value=False, + ), ): # /models should NOT raise - it's exempt EnterpriseRouteChecks.should_call_route("/models") @@ -1088,10 +1100,15 @@ class TestModelsRouteExemptFromDisableLLMEndpoints: """Test that /v1/models is allowed even when LLM API routes are disabled""" EnterpriseRouteChecks = self._get_enterprise_route_checks() - with patch.object( - EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True - ), patch.object( - EnterpriseRouteChecks, "is_management_routes_disabled", return_value=False + with ( + patch.object( + EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True + ), + patch.object( + EnterpriseRouteChecks, + "is_management_routes_disabled", + return_value=False, + ), ): # /v1/models should NOT raise - it's exempt EnterpriseRouteChecks.should_call_route("/v1/models") @@ -1101,10 +1118,15 @@ class TestModelsRouteExemptFromDisableLLMEndpoints: """Test that non-exempt LLM routes like /v1/chat/completions are still blocked""" EnterpriseRouteChecks = self._get_enterprise_route_checks() - with patch.object( - EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True - ), patch.object( - EnterpriseRouteChecks, "is_management_routes_disabled", return_value=False + with ( + patch.object( + EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True + ), + patch.object( + EnterpriseRouteChecks, + "is_management_routes_disabled", + return_value=False, + ), ): with pytest.raises(HTTPException) as exc_info: EnterpriseRouteChecks.should_call_route("/v1/chat/completions") @@ -1119,10 +1141,15 @@ class TestModelsRouteExemptFromDisableLLMEndpoints: """Test that /v1/embeddings is still blocked when LLM API routes are disabled""" EnterpriseRouteChecks = self._get_enterprise_route_checks() - with patch.object( - EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True - ), patch.object( - EnterpriseRouteChecks, "is_management_routes_disabled", return_value=False + with ( + patch.object( + EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True + ), + patch.object( + EnterpriseRouteChecks, + "is_management_routes_disabled", + return_value=False, + ), ): with pytest.raises(HTTPException) as exc_info: EnterpriseRouteChecks.should_call_route("/v1/embeddings") @@ -1134,10 +1161,15 @@ class TestModelsRouteExemptFromDisableLLMEndpoints: """Test that /models works normally when LLM API routes are not disabled""" EnterpriseRouteChecks = self._get_enterprise_route_checks() - with patch.object( - EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=False - ), patch.object( - EnterpriseRouteChecks, "is_management_routes_disabled", return_value=False + with ( + patch.object( + EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=False + ), + patch.object( + EnterpriseRouteChecks, + "is_management_routes_disabled", + return_value=False, + ), ): # Should not raise EnterpriseRouteChecks.should_call_route("/models") @@ -1359,6 +1391,38 @@ def test_non_org_admin_with_organizations_list(): assert _user_is_org_admin({"organizations": ["org-1"]}, user_obj) is False +def test_org_admin_cannot_escalate_to_other_org(): + """Regression: admin of org-A requesting [org-A, org-B] must be rejected.""" + user_obj = _make_org_admin_user("org-A") + assert _user_is_org_admin({"organizations": ["org-A", "org-B"]}, user_obj) is False + + +def test_org_admin_of_multiple_orgs_can_operate_on_both(): + """Admin of both org-A and org-B can operate on both.""" + memberships = [ + LiteLLM_OrganizationMembershipTable( + user_id="multi-admin", + organization_id="org-A", + user_role=LitellmUserRoles.ORG_ADMIN.value, + created_at=datetime(2024, 1, 1), + updated_at=datetime(2024, 1, 1), + ), + LiteLLM_OrganizationMembershipTable( + user_id="multi-admin", + organization_id="org-B", + user_role=LitellmUserRoles.ORG_ADMIN.value, + created_at=datetime(2024, 1, 1), + updated_at=datetime(2024, 1, 1), + ), + ] + user_obj = LiteLLM_UserTable( + user_id="multi-admin", + user_role=LitellmUserRoles.INTERNAL_USER.value, + organization_memberships=memberships, + ) + assert _user_is_org_admin({"organizations": ["org-A", "org-B"]}, user_obj) is True + + @pytest.mark.asyncio async def test_initialize_pass_through_registers_wildcard_for_auth_subpath(): """ @@ -1389,15 +1453,19 @@ async def test_initialize_pass_through_registers_wildcard_for_auth_subpath(): original_routes = LiteLLMRoutes.openai_routes.value[:] try: - with patch( - "litellm.proxy.proxy_server.app", - MagicMock(), - ), patch( - "litellm.proxy.proxy_server.premium_user", - True, - ), patch( - "litellm.proxy.proxy_server.config_passthrough_endpoints", - None, + with ( + patch( + "litellm.proxy.proxy_server.app", + MagicMock(), + ), + patch( + "litellm.proxy.proxy_server.premium_user", + True, + ), + patch( + "litellm.proxy.proxy_server.config_passthrough_endpoints", + None, + ), ): await initialize_pass_through_endpoints([endpoint_config]) @@ -1417,7 +1485,9 @@ async def test_initialize_pass_through_registers_wildcard_for_auth_subpath(): # Removing the endpoint should clean up openai_routes # remove_endpoint_routes takes endpoint_id (UUID portion of # the route key "{id}:exact:{path}:{methods}") - registered = InitPassThroughEndpointHelpers.get_all_registered_pass_through_routes() + registered = ( + InitPassThroughEndpointHelpers.get_all_registered_pass_through_routes() + ) endpoint_ids = {k.split(":")[0] for k in registered} for eid in endpoint_ids: InitPassThroughEndpointHelpers.remove_endpoint_routes(eid) @@ -1427,8 +1497,8 @@ async def test_initialize_pass_through_registers_wildcard_for_auth_subpath(): LiteLLMRoutes.openai_routes.value[:] = original_routes # Clean up any routes registered during this test to avoid # polluting the module-level _registered_pass_through_routes - registered = InitPassThroughEndpointHelpers.get_all_registered_pass_through_routes() + registered = ( + InitPassThroughEndpointHelpers.get_all_registered_pass_through_routes() + ) for k in registered: - InitPassThroughEndpointHelpers.remove_endpoint_routes( - k.split(":")[0] - ) + InitPassThroughEndpointHelpers.remove_endpoint_routes(k.split(":")[0]) diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 0f90d236aed..a1ba7ecd677 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -1878,6 +1878,121 @@ async def test_delete_user_cleans_up_created_by_invitation_links(mocker): assert condition[field] == {"in": ["admin-creator"]} +@pytest.mark.asyncio +async def test_delete_user_rejects_org_admin_deleting_outside_scope(mocker): + """Regression: an org admin of org-A must not be able to delete a user + whose org memberships include org-B. + + Route-level gate accepts the request when the caller supplies an + `organization_id` they administer; without per-user org authorization + the handler would cascade-delete the victim's keys, memberships, and + user row regardless of where the victim actually belongs. + """ + from fastapi import HTTPException + + from litellm.proxy._types import DeleteUserRequest, UserAPIKeyAuth + from litellm.proxy.management_endpoints.internal_user_endpoints import delete_user + + mock_prisma_client = mocker.MagicMock() + + # Target user exists and is a member of org-B only. + mock_target_user = mocker.MagicMock() + mock_target_user.user_id = "victim" + mock_target_user.user_email = "victim@example.com" + mock_target_user.teams = [] + mock_target_user.json.return_value = "{}" + + async def mock_find_unique(*args, **kwargs): + return mock_target_user + + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( + side_effect=mock_find_unique + ) + + # Caller (org_admin_user) administers org-A. + caller_membership = mocker.MagicMock() + caller_membership.organization_id = "org-A" + + # Target user is a member of org-B (outside caller's scope). + target_membership = mocker.MagicMock() + target_membership.organization_id = "org-B" + + async def mock_find_memberships(*args, **kwargs): + where = kwargs.get("where") or (args[0] if args else {}) + user_id_filter = where.get("user_id") + # Batched lookup: {"user_id": {"in": [...]}} returns target memberships. + # Caller role lookup: {"user_id": "", "user_role": ...}. + if isinstance(user_id_filter, dict) and "in" in user_id_filter: + if "victim" in user_id_filter["in"]: + # Attach user_id on the mock so the caller can build its + # per-user dict from the batch result. + target_membership.user_id = "victim" + return [target_membership] + return [] + if user_id_filter == "org_admin_user": + return [caller_membership] + return [] + + mock_prisma_client.db.litellm_organizationmembership.find_many = mocker.AsyncMock( + side_effect=mock_find_memberships + ) + + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + data = DeleteUserRequest(user_ids=["victim"]) + user_api_key_dict = UserAPIKeyAuth( + user_id="org_admin_user", user_role=LitellmUserRoles.ORG_ADMIN + ) + + with pytest.raises(HTTPException) as exc: + await delete_user(data=data, user_api_key_dict=user_api_key_dict) + assert exc.value.status_code == 403 + + # Critical: no delete_many calls should have executed. + assert not hasattr( + mock_prisma_client.db.litellm_verificationtoken.delete_many, "mock_calls" + ) or len( + mock_prisma_client.db.litellm_verificationtoken.delete_many.mock_calls + ) == 0 + + +@pytest.mark.asyncio +async def test_user_update_rejects_silent_create_for_non_proxy_admin(mocker): + """Regression: `/user/update` with an unknown user_email used to fall + through to an INSERT, silently creating a new user with caller-supplied + budget, models, and metadata. An org admin could use this to spawn + arbitrary users outside the /user/new authorization flow.""" + from fastapi import HTTPException + + from litellm.proxy._types import UpdateUserRequest, UserAPIKeyAuth + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _update_single_user_helper, + ) + + mock_prisma_client = mocker.MagicMock() + # user_email lookup yields None → would silently create pre-fix. + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( + return_value=None + ) + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + user_request = UpdateUserRequest( + user_email="newcomer@example.com", + max_budget=1_000_000, + models=["gpt-4"], + ) + org_admin = UserAPIKeyAuth( + user_id="org-admin", + user_role=LitellmUserRoles.ORG_ADMIN, + ) + + with pytest.raises(HTTPException) as exc: + await _update_single_user_helper( + user_request=user_request, user_api_key_dict=org_admin + ) + assert exc.value.status_code == 404 + + # ===================================================================== # /v2/user/info endpoint tests # ===================================================================== diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 479defbff5c..41ab60fcf67 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -8030,6 +8030,74 @@ async def test_update_key_non_budget_fields_allowed_for_internal_user(monkeypatc assert result is not None +@pytest.mark.asyncio +async def test_update_key_non_budget_rejects_cross_user_modification(monkeypatch): + """Regression: previously _check_key_admin_access was gated on + max_budget/spend changes only, so an internal user could rewrite any + OTHER field (alias, models, tpm_limit, blocked, metadata, …) on any + key they weren't admin of as long as they avoided budget/spend. This + confirms that a non-admin user updating a key that belongs to another + user fails with 403 even for non-budget fields.""" + from litellm.proxy.management_endpoints.key_management_endpoints import ( + update_key_fn, + ) + + mock_prisma_client = AsyncMock() + test_hashed_token = ( + "cafebabe" * 8 + ) + + mock_existing_key = MagicMock() + mock_existing_key.token = test_hashed_token + mock_existing_key.user_id = "victim_user" # owned by someone else + mock_existing_key.team_id = None + mock_existing_key.project_id = None + mock_existing_key.max_budget = 10.0 + mock_existing_key.key_alias = "original" + mock_existing_key.models = [] + mock_existing_key.model_dump.return_value = { + "token": test_hashed_token, + "user_id": "victim_user", + "team_id": None, + "max_budget": 10.0, + } + + mock_prisma_client.get_data = AsyncMock(return_value=mock_existing_key) + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=mock_existing_key + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", AsyncMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr( + "litellm.proxy.proxy_server.hash_token", lambda t: test_hashed_token + ) + + mock_request = MagicMock() + mock_request.query_params = {} + attacker = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-attacker", + user_id="attacker_user", # NOT the owner + ) + + # Trying to blanket-rewrite a non-budget field on someone else's key + # must now fail. + with pytest.raises(ProxyException) as exc: + await update_key_fn( + request=mock_request, + data=UpdateKeyRequest( + key=test_hashed_token, key_alias="pwned", blocked=True + ), + user_api_key_dict=attacker, + litellm_changed_by=None, + ) + assert str(exc.value.code) == "403" + + # ============================================================================ # LIT-1884: Internal users cannot create invalid keys # ============================================================================ 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 bee6642dec7..896e1efe0e4 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -1010,12 +1010,15 @@ async def test_validate_team_member_add_permissions_non_admin(): team.organization_id = None # Mock the helper functions to return False - with patch( - "litellm.proxy.management_endpoints.team_endpoints._is_user_team_admin", - return_value=False, - ), patch( - "litellm.proxy.management_endpoints.team_endpoints._is_available_team", - return_value=False, + with ( + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_user_team_admin", + return_value=False, + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_available_team", + return_value=False, + ), ): # Should raise HTTPException for non-admin with pytest.raises(HTTPException) as exc_info: @@ -1257,19 +1260,17 @@ async def test_update_team_team_member_budget_not_passed_to_db(): user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id" ) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, patch( - "litellm.proxy.proxy_server.llm_router" - ) as mock_llm_router, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ) as mock_cache, patch( - "litellm.proxy.proxy_server.proxy_logging_obj" - ) as mock_logging, patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" - ), patch( - "litellm.proxy.auth.auth_checks._cache_team_object" - ) as mock_cache_team, patch( - "litellm.proxy.management_endpoints.team_endpoints.TeamMemberBudgetHandler.upsert_team_member_budget_table" - ) as mock_upsert_budget: + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, + patch("litellm.proxy.proxy_server.llm_router") as mock_llm_router, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_logging, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch("litellm.proxy.auth.auth_checks._cache_team_object") as mock_cache_team, + patch( + "litellm.proxy.management_endpoints.team_endpoints.TeamMemberBudgetHandler.upsert_team_member_budget_table" + ) as mock_upsert_budget, + ): # Setup mock prisma client mock_existing_team = MagicMock() mock_existing_team.model_dump.return_value = { @@ -1690,19 +1691,17 @@ async def test_update_team_with_team_member_budget_duration(): user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id" ) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, patch( - "litellm.proxy.proxy_server.llm_router" - ) as mock_llm_router, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ) as mock_cache, patch( - "litellm.proxy.proxy_server.proxy_logging_obj" - ) as mock_logging, patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" - ), patch( - "litellm.proxy.auth.auth_checks._cache_team_object" - ) as mock_cache_team, patch( - "litellm.proxy.management_endpoints.team_endpoints.TeamMemberBudgetHandler.upsert_team_member_budget_table" - ) as mock_upsert_budget: + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, + patch("litellm.proxy.proxy_server.llm_router") as mock_llm_router, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_logging, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch("litellm.proxy.auth.auth_checks._cache_team_object") as mock_cache_team, + patch( + "litellm.proxy.management_endpoints.team_endpoints.TeamMemberBudgetHandler.upsert_team_member_budget_table" + ) as mock_upsert_budget, + ): mock_existing_team = MagicMock() mock_existing_team.model_dump.return_value = { "team_id": "test_team_id", @@ -1777,7 +1776,9 @@ async def test_backfill_team_member_budget_entries_creates_missing_memberships() from unittest.mock import AsyncMock, MagicMock from litellm.proxy._types import Member - from litellm.proxy.management_endpoints.team_endpoints import TeamMemberBudgetHandler + from litellm.proxy.management_endpoints.team_endpoints import ( + TeamMemberBudgetHandler, + ) team_id = "team-abc" budget_id = "budget-xyz" @@ -1847,7 +1848,9 @@ async def test_backfill_team_member_budget_entries_no_op_when_all_exist(): from unittest.mock import AsyncMock, MagicMock from litellm.proxy._types import Member - from litellm.proxy.management_endpoints.team_endpoints import TeamMemberBudgetHandler + from litellm.proxy.management_endpoints.team_endpoints import ( + TeamMemberBudgetHandler, + ) team_id = "team-abc" budget_id = "budget-xyz" @@ -1886,7 +1889,9 @@ async def test_backfill_team_member_budget_entries_empty_members(): """ from unittest.mock import AsyncMock, MagicMock - from litellm.proxy.management_endpoints.team_endpoints import TeamMemberBudgetHandler + from litellm.proxy.management_endpoints.team_endpoints import ( + TeamMemberBudgetHandler, + ) mock_prisma = MagicMock() mock_prisma.db.litellm_teammembership.find_many = AsyncMock(return_value=[]) @@ -2092,11 +2097,14 @@ async def test_bulk_team_member_add_all_users_flag(): updated_team_memberships=[], ) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.management_endpoints.team_endpoints.team_member_add", - new_callable=AsyncMock, - return_value=mock_team_response, - ) as mock_team_member_add: + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch( + "litellm.proxy.management_endpoints.team_endpoints.team_member_add", + new_callable=AsyncMock, + return_value=mock_team_response, + ) as mock_team_member_add, + ): # Mock the database find_many call mock_prisma.db.litellm_usertable.find_many = AsyncMock( return_value=mock_db_users @@ -2213,12 +2221,15 @@ async def test_list_team_v2_security_check_non_admin_user(): user_id="non_admin_user_123", ) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ), patch("litellm.proxy.proxy_server.proxy_logging_obj"), patch( - "litellm.proxy.management_endpoints.team_endpoints.get_user_object", - new_callable=AsyncMock, - return_value=None, + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, + patch("litellm.proxy.proxy_server.user_api_key_cache"), + patch("litellm.proxy.proxy_server.proxy_logging_obj"), + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + new_callable=AsyncMock, + return_value=None, + ), ): mock_prisma_client.return_value = MagicMock() # Mock non-None prisma client @@ -2260,12 +2271,15 @@ async def test_list_team_v2_security_check_non_admin_user_other_user(): user_id="non_admin_user_123", ) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ), patch("litellm.proxy.proxy_server.proxy_logging_obj"), patch( - "litellm.proxy.management_endpoints.team_endpoints.get_user_object", - new_callable=AsyncMock, - return_value=None, + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, + patch("litellm.proxy.proxy_server.user_api_key_cache"), + patch("litellm.proxy.proxy_server.proxy_logging_obj"), + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + new_callable=AsyncMock, + return_value=None, + ), ): mock_prisma_client.return_value = MagicMock() # Mock non-None prisma client @@ -2305,9 +2319,11 @@ async def test_list_team_v2_security_check_non_admin_user_own_teams(): user_id="non_admin_user_123", ) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ), patch("litellm.proxy.proxy_server.proxy_logging_obj"): + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, + patch("litellm.proxy.proxy_server.user_api_key_cache"), + patch("litellm.proxy.proxy_server.proxy_logging_obj"), + ): # Mock prisma client and database operations mock_db = Mock() mock_prisma_client.db = mock_db @@ -2509,12 +2525,15 @@ async def test_list_team_v2_org_admin_sees_org_teams(): ], ) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ), patch("litellm.proxy.proxy_server.proxy_logging_obj"), patch( - "litellm.proxy.management_endpoints.team_endpoints.get_user_object", - new_callable=AsyncMock, - return_value=mock_user, + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache"), + patch("litellm.proxy.proxy_server.proxy_logging_obj"), + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + new_callable=AsyncMock, + return_value=mock_user, + ), ): mock_db = Mock() mock_prisma.db = mock_db @@ -2592,12 +2611,15 @@ async def test_list_team_v2_org_admin_cannot_view_other_orgs(): ], ) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ), patch("litellm.proxy.proxy_server.proxy_logging_obj"), patch( - "litellm.proxy.management_endpoints.team_endpoints.get_user_object", - new_callable=AsyncMock, - return_value=mock_user, + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache"), + patch("litellm.proxy.proxy_server.proxy_logging_obj"), + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + new_callable=AsyncMock, + return_value=mock_user, + ), ): mock_prisma.db = Mock() @@ -2680,11 +2702,14 @@ async def test_list_team_v2_org_admin_with_user_id_returns_user_teams(): return mock_org_admin return mock_target_user - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ), patch("litellm.proxy.proxy_server.proxy_logging_obj"), patch( - "litellm.proxy.management_endpoints.team_endpoints.get_user_object", - side_effect=mock_get_user_object, + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache"), + patch("litellm.proxy.proxy_server.proxy_logging_obj"), + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + side_effect=mock_get_user_object, + ), ): mock_db = Mock() mock_prisma.db = mock_db @@ -2714,10 +2739,10 @@ async def test_list_team_v2_org_admin_with_user_id_returns_user_teams(): assert result["total"] == 1 - # Verify the where clause filters by user's teams, not org scope + # Verify the where clause filters by user's teams AND org scope where = mock_db.litellm_teamtable.find_many.call_args.kwargs["where"] assert where["team_id"] == {"in": ["team_X", "team_Y"]} - assert "organization_id" not in where + assert where["organization_id"] == {"in": ["org_A"]} @pytest.mark.asyncio @@ -2913,15 +2938,15 @@ async def test_new_team_max_budget_exceeds_user_max_budget(): 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: + 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 @@ -2982,15 +3007,15 @@ async def test_new_team_max_budget_within_user_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: + 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, + ): # Setup mocks mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_license.is_team_count_over_limit.return_value = False @@ -3111,17 +3136,18 @@ async def test_new_team_org_scoped_budget_bypasses_user_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: + 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 @@ -3253,17 +3279,18 @@ async def test_new_team_org_scoped_models_bypasses_user_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: + 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 @@ -3393,13 +3420,14 @@ async def test_new_team_standalone_validates_against_user_models(monkeypatch): 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: + 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 @@ -3460,15 +3488,15 @@ async def test_new_team_standalone_validates_against_user_budget(): 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: + 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 @@ -3534,17 +3562,18 @@ async def test_new_team_org_scoped_budget_exceeds_org_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: + 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 @@ -3613,17 +3642,18 @@ async def test_new_team_org_scoped_models_not_in_org_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._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: + 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 @@ -3688,13 +3718,14 @@ async def test_update_team_standalone_budget_exceeds_user_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: + 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" @@ -3778,16 +3809,18 @@ async def test_update_team_org_scoped_budget_exceeds_org_limit(): 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: + 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" @@ -3852,13 +3885,14 @@ async def test_update_team_standalone_models_exceeds_user_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: + 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" @@ -3936,16 +3970,18 @@ async def test_update_team_org_scoped_budget_bypasses_user_limit(): 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: + 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" @@ -4044,16 +4080,18 @@ async def test_update_team_org_scoped_models_bypasses_user_limit(): 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: + 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" @@ -4145,16 +4183,18 @@ async def test_update_team_org_scoped_models_not_in_org_models(): 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: + 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" @@ -4231,16 +4271,18 @@ async def test_update_team_org_scoped_models_with_all_proxy_models(): mock_org.models = [SpecialModelNames.all_proxy_models.value] # Allows all models 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: + 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-all-proxy-models-123" @@ -4333,10 +4375,10 @@ async def test_update_team_tpm_limit_exceeds_user_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" + 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() @@ -4397,10 +4439,10 @@ async def test_update_team_rpm_limit_exceeds_user_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" + 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() @@ -4479,15 +4521,15 @@ async def test_new_team_org_scoped_tpm_exceeds_org_limit(): 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), + 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) @@ -4555,15 +4597,15 @@ async def test_new_team_org_scoped_rpm_exceeds_org_limit(): 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), + 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) @@ -4634,20 +4676,22 @@ async def test_new_team_org_scoped_tpm_rpm_bypasses_user_limit(): 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(), + 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) @@ -4736,13 +4780,14 @@ async def test_update_team_org_scoped_tpm_exceeds_org_limit(): 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), + 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() @@ -4822,13 +4867,14 @@ async def test_update_team_org_scoped_rpm_exceeds_org_limit(): 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), + 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() @@ -4911,15 +4957,15 @@ async def test_update_team_org_scoped_tpm_rpm_bypasses_user_limit(): 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), + 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() @@ -5036,17 +5082,18 @@ async def test_update_team_guardrails_with_org_id(): "teams": [], } - 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() - ), patch( - "litellm.proxy.proxy_server.premium_user", - True, # Required for guardrails feature - ), patch( - "litellm.proxy.proxy_server.llm_router", MagicMock() + 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() + ), + patch( + "litellm.proxy.proxy_server.premium_user", + True, # Required for guardrails feature + ), + patch("litellm.proxy.proxy_server.llm_router", MagicMock()), ): # Mock existing team - must have compatible models with organization mock_existing_team = MagicMock() @@ -5081,6 +5128,16 @@ async def test_update_team_guardrails_with_org_id(): return_value=mock_org ) + # Destination-org guard in update_team queries for the caller's + # ORG_ADMIN membership on the destination org. Return a match so + # the guardrails-update path (the subject under test) proceeds. + mock_org_admin_membership = MagicMock() + mock_org_admin_membership.user_id = "org-admin-guardrails-test" + mock_org_admin_membership.organization_id = "test-org-guardrails" + mock_prisma.db.litellm_organizationmembership.find_many = AsyncMock( + return_value=[mock_org_admin_membership] + ) + # Mock team update mock_updated_team = MagicMock(spec=LiteLLM_TeamTable) mock_updated_team.team_id = "team-guardrails-123" @@ -5601,15 +5658,15 @@ async def test_new_team_soft_budget_validation( 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: + 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, + ): # Setup mocks mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_license.is_team_count_over_limit.return_value = False @@ -5799,13 +5856,14 @@ async def test_update_team_soft_budget_validation( 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: + 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 team with existing budgets mock_existing_team = MagicMock() mock_existing_team.team_id = "test-team-123" @@ -6794,14 +6852,18 @@ async def test_list_team_v1_batches_key_queries(): key3 = MagicMock() key3.team_id = "team-2" - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, patch( - "litellm.proxy.management_endpoints.team_endpoints._authorize_and_filter_teams", - new_callable=AsyncMock, - return_value=[team1, team2], - ), patch( - "litellm.proxy.management_endpoints.team_endpoints.get_all_team_memberships", - new_callable=AsyncMock, - return_value=[], + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, + patch( + "litellm.proxy.management_endpoints.team_endpoints._authorize_and_filter_teams", + new_callable=AsyncMock, + return_value=[team1, team2], + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_all_team_memberships", + new_callable=AsyncMock, + return_value=[], + ), ): async def filtered_find_many(**kwargs): @@ -7070,16 +7132,17 @@ async def test_update_team_rejects_unauthorized_caller(): from litellm.proxy._types import UpdateTeamRequest - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, patch( - "litellm.proxy.proxy_server.llm_router" - ), patch("litellm.proxy.proxy_server.user_api_key_cache"), patch( - "litellm.proxy.proxy_server.proxy_logging_obj" - ), patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" - ), patch( - "litellm.proxy.management_endpoints.team_endpoints._is_user_org_admin_for_team", - new_callable=AsyncMock, - return_value=False, + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, + patch("litellm.proxy.proxy_server.llm_router"), + patch("litellm.proxy.proxy_server.user_api_key_cache"), + patch("litellm.proxy.proxy_server.proxy_logging_obj"), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_user_org_admin_for_team", + new_callable=AsyncMock, + return_value=False, + ), ): mock_existing_team = MagicMock() mock_existing_team.model_dump.return_value = {