diff --git a/litellm/proxy/management_helpers/bulk_user_creation.py b/litellm/proxy/management_helpers/bulk_user_creation.py index cca37ebd7ed..62228f6faa8 100644 --- a/litellm/proxy/management_helpers/bulk_user_creation.py +++ b/litellm/proxy/management_helpers/bulk_user_creation.py @@ -31,6 +31,7 @@ from litellm.proxy._types import ( from litellm.proxy.auth.auth_checks import invalidate_team_member_spend_state from litellm.proxy.auth.auth_utils import enforce_batch_limits_are_admin_only from litellm.proxy.auth.litellm_license import LicenseCheck +from litellm.proxy.common_utils.callback_config_validation import credential_reference_rejection from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks @@ -221,7 +222,14 @@ def _row_error(item: BulkNewUserItem, user_api_key_dict: UserAPIKeyAuth) -> str enforce_batch_limits_are_admin_only(item, None, user_api_key_dict, "key") except Exception as exc: # noqa: BLE001 # any validation failure is reported on this row only return _error_message(exc) - return None + rejection: Final = ( + credential_reference_rejection( + item.metadata, None, is_proxy_admin=user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN + ) + if item.auto_create_key + else None + ) + return None if rejection is None else rejection.message def _normalized_email(email: str | None) -> str | None: diff --git a/tests/integration/observability/test_arize_key_credential_reference.py b/tests/integration/observability/test_arize_key_credential_reference.py index a445269837e..57e74b370e3 100644 --- a/tests/integration/observability/test_arize_key_credential_reference.py +++ b/tests/integration/observability/test_arize_key_credential_reference.py @@ -220,3 +220,20 @@ def test_only_proxy_admins_create_a_user_whose_key_references_a_credential(gatew ) assert single.status_code == 403, single.text assert "litellm_credential_name" in single.text, single.text + + delegate_key: Final = scenario.key( + user_id=scenario.user(user_role="internal_user"), allowed_routes=["/management/v1/users/bulk"] + ) + bulk: Final = rig.proxy.request( + "POST", + "/management/v1/users/bulk", + { + "users": [ + {"user_id": "user-" + uuid.uuid4().hex, "auto_create_key": True, "metadata": referencing_metadata} + ] + }, + key=delegate_key, + ) + assert bulk.status_code == 200, bulk.text + assert [row["success"] for row in bulk.json()["data"]] == [False], bulk.text + assert "litellm_credential_name" in bulk.json()["data"][0]["error"], bulk.text diff --git a/tests/unit/proxy/management_helpers/test_bulk_user_creation.py b/tests/unit/proxy/management_helpers/test_bulk_user_creation.py index 2b966a8ebeb..96c84578111 100644 --- a/tests/unit/proxy/management_helpers/test_bulk_user_creation.py +++ b/tests/unit/proxy/management_helpers/test_bulk_user_creation.py @@ -7,6 +7,7 @@ import pytest from prisma.errors import UniqueViolationError from pydantic import BaseModel, ConfigDict, ValidationError +import litellm from litellm.caching.caching import DualCache from litellm.proxy._types import LiteLLM_TeamTable, LitellmUserRoles, Member, UserAPIKeyAuth from litellm.proxy.list_api.common import ManagementProblem @@ -15,6 +16,7 @@ from litellm.types.proxy.management_endpoints.internal_user_endpoints import ( BulkNewUserItem, BulkNewUserRequest, ) +from litellm.types.utils import CredentialItem ADMIN: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) INTERNAL: Final = UserAPIKeyAuth(user_id="someone", user_role=LitellmUserRoles.INTERNAL_USER) @@ -434,6 +436,39 @@ async def test_only_proxy_admin_sets_batch_limits_on_a_created_key(): assert calls[-1]["metadata"] == limits +@pytest.mark.asyncio +async def test_only_proxy_admin_creates_a_key_whose_logging_references_a_credential( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + litellm, + "credential_list", + [CredentialItem(credential_name="space-a", credential_values={"arize_api_key": "k"}, credential_info={})], + ) + reference: Final = {"logging": [{"callback_name": "arize", "litellm_credential_name": "space-a", "callback_vars": {}}]} + calls: list[dict[str, object]] = [] + + async def generate_key(**kwargs: object) -> dict[str, object]: + calls.append(kwargs) + return {"token": f"sk-{kwargs['user_id']}"} + + rows: Final = [ + {"user_id": "u1", "auto_create_key": True, "metadata": reference}, + {"user_id": "u2", "auto_create_key": True}, + ] + + response = await _run(_FakePrisma(), rows, caller=INTERNAL, generate_key=generate_key) + + assert [r.success for r in response.data] == [False, True] + assert "litellm_credential_name" in (response.data[0].error or "") + assert [call["user_id"] for call in calls] == ["u2"] + + admin_response = await _run(_FakePrisma(), rows[:1], caller=ADMIN, generate_key=generate_key) + + assert [r.key for r in admin_response.data] == ["sk-u1"] + assert calls[-1]["metadata"] == reference + + @pytest.mark.asyncio async def test_license_is_checked_once_against_the_whole_batch(): prisma = _FakePrisma()