diff --git a/litellm/proxy/auth/login_utils.py b/litellm/proxy/auth/login_utils.py index 9783f44b41f..3b9385f7acc 100644 --- a/litellm/proxy/auth/login_utils.py +++ b/litellm/proxy/auth/login_utils.py @@ -339,7 +339,7 @@ async def _sign_in( response = await generate_key_helper_fn( llm_router=None, request_type="key", - **{ + **{ # pyright: ignore[reportUnknownArgumentType] # the onboarding spread binds our allow_ui_session_marker flag explicitly "user_role": LitellmUserRoles.PROXY_ADMIN, "duration": LITELLM_UI_SESSION_DURATION, "key_max_budget": litellm.max_ui_session_budget, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index e87ca70035c..5cbc659330d 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -459,6 +459,8 @@ def _effective_key_for_generate(data: GenerateKeyRequest, now: datetime) -> Lite ) metadata: Final = data.metadata or MappingProxyType({}) folded_metadata: Final = {**metadata, **metadata_fields} + # The ui-session marker is mint-flow-only: /key/generate cannot stamp it (R1-L12). + folded_metadata.pop(UI_SESSION_METADATA_MARKER_KEY, None) columns: Final = handle_key_type(data, {**column_fields}) expires: Final = ( now + timedelta(seconds=duration_in_seconds(duration=data.duration)) if data.duration is not None else None @@ -918,6 +920,33 @@ def handle_key_type(data: GenerateKeyRequest, data_json: dict) -> dict: _NON_ADMIN_SAFE_ALLOWED_ROUTES_PRESETS: Final = frozenset({"llm_api_routes", "info_routes"}) +_ALLOW_CHANGED_BY_METADATA_KEY: Final = "allow_litellm_changed_by_header" + + +def _strip_non_admin_changed_by_metadata( + *, + data: KeyRequestBase, + user_api_key_dict: UserAPIKeyAuth, +) -> None: + """Only proxy admins may enable caller-supplied audit attribution (R2-20). + + The flag opts a key into honoring the ``litellm-changed-by`` header in + audit rows, so a self-granted flag lets any key holder forge the primary + audit attribution field. Silently stripped for non-admin writes. + """ + metadata: Final = data.metadata + if not isinstance(metadata, dict): + return + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: + return + if metadata.get(_ALLOW_CHANGED_BY_METADATA_KEY) is not True: + return + metadata.pop(_ALLOW_CHANGED_BY_METADATA_KEY, None) + verbose_proxy_logger.info( + "Stripped allow_litellm_changed_by_header from key metadata: only proxy admins may enable caller-forwarded audit attribution." + ) + + def _validate_caller_can_change_key_ownership( data: BaseModel | None, existing_key_row: LiteLLM_VerificationToken, @@ -1300,12 +1329,26 @@ async def _common_key_generation_helper( detail={"error": "Only proxy admins can enable throttle_on_budget_exceeded on a key."}, ) + _strip_non_admin_changed_by_metadata(data=data, user_api_key_dict=user_api_key_dict) + + _reject_non_admin_budget_authority_metadata( + data=data, + existing_metadata=None, + user_api_key_dict=user_api_key_dict, + ) + await _validate_end_user_budget_id_change( requested_budget_id=_requested_end_user_budget_id(data), existing_budget_id=None, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, ) + await _validate_budget_id_change( + requested_budget_id=data.budget_id, + existing_budget_id=None, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + ) enforce_output_token_estimates_are_admin_only( data=data, @@ -2538,7 +2581,11 @@ def prepare_metadata_fields(data: BaseModel, non_default_values: dict, existing_ if "metadata" not in non_default_values: # allow user to set metadata to none non_default_values["metadata"] = existing_metadata.copy() - casted_metadata: Final = cast(dict, non_default_values["metadata"]) + # An explicit metadata=None means "clear this field"; the cast below holds + # the dict-shaped value only, so the clear-to-none path is handled at the + # tail via raw_metadata. + raw_metadata: Final[object] = non_default_values["metadata"] # pyright: ignore[reportUnknownVariableType] # key metadata is a db-shaped payload; the tail handles the explicit clear-to-none case + casted_metadata: Final = cast(dict, raw_metadata) # Reserved metadata fields are immutable once set. Preserve the existing value # when omitted, reject any explicit attempt to change it (including null). @@ -2576,7 +2623,12 @@ def prepare_metadata_fields(data: BaseModel, non_default_values: dict, existing_ "litellm.proxy.proxy_server.prepare_metadata_fields(): Exception occured - %s", e ) - non_default_values["metadata"] = encrypt_callback_vars(casted_metadata) + # The ui-session marker is reserved for the minted dashboard sessions; + # update/regenerate cannot stamp it from caller data (R1-L12). An explicit + # metadata=None stays None (clear-to-none), so skip both here. + if raw_metadata is not None: + casted_metadata.pop(UI_SESSION_METADATA_MARKER_KEY, None) # pyright: ignore[reportUnknownMemberType] # the cast above intentionally keeps the db-metadata shape untyped + non_default_values["metadata"] = encrypt_callback_vars(casted_metadata) return non_default_values @@ -3257,6 +3309,24 @@ def _requested_end_user_budget_id(data: KeyRequestBase) -> str | None: return get_key_end_user_budget_id(data.metadata) or "" +async def _validate_budget_id_change( + requested_budget_id: str | None, + existing_budget_id: str | None, + user_api_key_dict: UserAPIKeyAuth, + prisma_client: PrismaClient | None, +) -> None: + """A key's budget_id row supplies the key's effective limits (max_budget, + tpm/rpm, model_max_budget) from LiteLLM_BudgetTable, so only proxy admins + may change it. The delegation-ceiling check fires only when max_budget is + set explicitly, so an unvalidated budget_id bypasses it entirely (R1-H3). + Dangling ids stay a no-op, as upstream treats them.""" + if requested_budget_id is None or requested_budget_id == (existing_budget_id or ""): + return + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: + forbidden_detail: Final = {"error": "Only proxy admins can set budget_id on a key."} + raise HTTPException(status_code=403, detail=forbidden_detail) + + async def _validate_end_user_budget_id_change( requested_budget_id: str | None, existing_budget_id: str | None, @@ -3278,6 +3348,71 @@ async def _validate_end_user_budget_id_change( raise HTTPException(status_code=400, detail=missing_detail) +_TEMP_BUDGET_METADATA_KEYS: Final[frozenset[str]] = frozenset({"temp_budget_increase", "temp_budget_expiry"}) + + +def _key_request_changes_budget( + *, + data: KeyRequestBase, + existing_key_row: LiteLLM_VerificationToken, +) -> bool: + """Whether the request would move a key's budget enforcement. + + max_budget, spend, budget_limits and soft_budget move the ceiling or the + recorded counter; budget_duration reshapes the window the ceiling resets + in. /key/update and /key/regenerate treat all of these as admin-only + changes for non-admin callers. + """ + changed_max_budget: Final = data.max_budget is not None and data.max_budget != getattr( + existing_key_row, "max_budget", None + ) + changed_budget_duration: Final = data.budget_duration is not None and data.budget_duration != getattr( + existing_key_row, "budget_duration", None + ) + return ( + changed_max_budget + or data.spend is not None + or changed_budget_duration + or "budget_limits" in data.model_fields_set + or "soft_budget" in data.model_fields_set + ) + + +def _reject_non_admin_budget_authority_metadata( + *, + data: KeyRequestBase, + existing_metadata: Mapping[str, object] | None, + user_api_key_dict: UserAPIKeyAuth, +) -> None: + """Reject metadata that grants budget authority to a non-admin caller. + + temp_budget_increase/temp_budget_expiry lift the effective max_budget at + auth time and a metadata-borne throttle_on_budget_exceeded flips a + hard-blocked over-budget key into throttled-continue. The equivalent typed + request fields are already admin-gated; without this check the same + values ride the free-form metadata dict past those gates. + """ + incoming: Final = data.metadata if isinstance(data.metadata, Mapping) else {} # pyright: ignore[reportUnknownVariableType] # pydantic metadata fields arrive without a static type; reads stay gated below + existing: Final = existing_metadata if isinstance(existing_metadata, Mapping) else {} # pyright: ignore[reportUnknownVariableType] # row metadata is db-shaped; only membership and .get are read + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: + return + sets_temp_budget: Final = any(key in incoming for key in _TEMP_BUDGET_METADATA_KEYS) + incoming_throttle: Final = incoming.get("throttle_on_budget_exceeded") # pyright: ignore[reportUnknownMemberType] # see above: untyped metadata payload + existing_throttle: Final = existing.get("throttle_on_budget_exceeded") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # see above: untyped metadata payload + escalates_throttle: Final = incoming_throttle is True and existing_throttle is not True + if not (sets_temp_budget or escalates_throttle): + return + raise HTTPException( + status_code=403, + detail={ + "error": ( + "Only proxy admins can set budget authority fields in key metadata " + "(temp_budget_increase, temp_budget_expiry, throttle_on_budget_exceeded)." + ) + }, + ) + + _GENERAL_SETTINGS: Final = TypeAdapter(dict[str, object]) @@ -3343,6 +3478,8 @@ async def _validate_update_key_data( validate_finite_spend(data.spend) validate_budget_duration(data.budget_duration) + _strip_non_admin_changed_by_metadata(data=data, user_api_key_dict=user_api_key_dict) + _is_proxy_admin: Final = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value _enforce_allowed_routes_update_permission( @@ -3430,12 +3567,7 @@ async def _validate_update_key_data( # through the non-admin path would let a key owner / team member # overwrite the live counter below real usage and silently weaken # enforcement. - _is_budget_change: Final = ( - (data.max_budget is not None and data.max_budget != existing_key_row.max_budget) - or data.spend is not None - or "budget_limits" in data.model_fields_set - or "soft_budget" in data.model_fields_set - ) + _is_budget_change: Final = _key_request_changes_budget(data=data, existing_key_row=existing_key_row) _existing_metadata: Final = getattr(existing_key_row, "metadata", None) _existing_throttle: Final = ( @@ -3446,6 +3578,11 @@ async def _validate_update_key_data( status_code=403, detail={"error": "Only proxy admins can enable throttle_on_budget_exceeded on a key."}, ) + _reject_non_admin_budget_authority_metadata( + data=data, + existing_metadata=_existing_metadata, # pyright: ignore[reportUnknownArgumentType] # row metadata is db-shaped; the callee only reads membership and .get + user_api_key_dict=user_api_key_dict, + ) await _validate_end_user_budget_id_change( requested_budget_id=_requested_end_user_budget_id(data), @@ -3455,6 +3592,12 @@ async def _validate_update_key_data( user_api_key_dict=user_api_key_dict, prisma_client=checked_prisma_client, ) + await _validate_budget_id_change( + requested_budget_id=data.budget_id, + existing_budget_id=getattr(existing_key_row, "budget_id", None), + user_api_key_dict=user_api_key_dict, + prisma_client=checked_prisma_client, + ) enforce_output_token_estimates_are_admin_only( data=data, @@ -4812,9 +4955,13 @@ def _check_model_access_group(models: list[str] | None, llm_router: Router | Non _NO_METADATA: Final[Mapping[str, object]] = MappingProxyType({}) +UI_SESSION_METADATA_MARKER_KEY: Final = "ui_session_key" + + def metadata_json_with_limits( metadata: Mapping[str, object] | None, *, + allow_ui_session_marker: object = False, model_rpm_limit: Mapping[str, object] | None, model_tpm_limit: Mapping[str, object] | None, mcp_rpm_limit: Mapping[str, int] | None, @@ -4840,6 +4987,11 @@ def metadata_json_with_limits( if metadata is None and not limits: return json.dumps(None) merged: Final = {**(metadata or _NO_METADATA), **dict(limits)} + # The ui-session marker is reserved for the minted dashboard sessions; a + # key minted outside those flows has it removed so a sentinel-named team + # row cannot forge the session credential shape (R1-L12). + if not allow_ui_session_marker: + merged.pop(UI_SESSION_METADATA_MARKER_KEY, None) return json.dumps(encrypt_callback_vars(merged)) @@ -4904,6 +5056,7 @@ async def generate_key_helper_fn( budget_limits: list | None = None, # multiple concurrent budget windows *, llm_router: Router | None = None, + allow_ui_session_marker: object = False, ): from litellm.proxy.proxy_server import premium_user, prisma_client @@ -4957,6 +5110,7 @@ async def generate_key_helper_fn( metadata_json: Final = metadata_json_with_limits( metadata, + allow_ui_session_marker=allow_ui_session_marker, model_rpm_limit=model_rpm_limit, model_tpm_limit=model_tpm_limit, mcp_rpm_limit=mcp_rpm_limit, @@ -5876,6 +6030,26 @@ async def _execute_virtual_key_regeneration( user_api_key_dict=user_api_key_dict, ) + if data is not None: + _strip_non_admin_changed_by_metadata(data=data, user_api_key_dict=user_api_key_dict) + _regenerate_is_proxy_admin: Final = ( + user_api_key_dict.user_role is not None + and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value + ) + if not _regenerate_is_proxy_admin and _key_request_changes_budget(data=data, existing_key_row=key_in_db): + await _check_key_admin_access( + user_api_key_dict=user_api_key_dict, + hashed_token=hashed_api_key, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + route="/key/regenerate (max_budget/spend)", + ) + _reject_non_admin_budget_authority_metadata( + data=data, + existing_metadata=getattr(key_in_db, "metadata", None), + user_api_key_dict=user_api_key_dict, + ) + # Apply the same membership rule used on /key/update: when the caller # asks to point the regenerated key at a different organization_id, # require they are a member of (or proxy admin over) the target org. @@ -5914,6 +6088,12 @@ async def _execute_virtual_key_regeneration( user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, ) + await _validate_budget_id_change( + requested_budget_id=data.budget_id, + existing_budget_id=getattr(key_in_db, "budget_id", None), + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + ) new_token: Final = await get_new_token(data=data) new_token_hash: Final = hash_token(new_token) @@ -7975,7 +8155,7 @@ async def _can_user_query_key_info( or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value ) or user_api_key_dict.api_key == key - or key_info.user_id == user_api_key_dict.user_id + or (key_info.user_id is not None and key_info.user_id == user_api_key_dict.user_id) or await TeamMemberPermissionChecks.user_belongs_to_keys_team( user_api_key_dict=user_api_key_dict, existing_key_row=key_info, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 69e92948a0e..a58985211ed 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -18704,7 +18704,7 @@ async def _generate_onboarding_ui_session_token(user_obj: _UserTableRow) -> str: response: Final = await generate_key_helper_fn( llm_router=llm_router, request_type="key", - **{ + **{ # pyright: ignore[reportUnknownArgumentType] # the onboarding spread binds our allow_ui_session_marker flag explicitly "user_role": user_obj.user_role, "duration": LITELLM_UI_SESSION_DURATION, "key_max_budget": litellm.max_ui_session_budget, diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index dffd9b9b513..a83f6d47c09 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -22682,6 +22682,54 @@ async def test_rotate_master_key_reencrypts_guardrail_params(monkeypatch): } +# --------------------------------------------------------------------------- +# team-admin api_base harvest gate (fork hardening R2-9) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_key_info_owner_match_requires_non_none_user_id(): + """The None == None identity match let any ownerless principal pass: the + fixed operator requires both sides non-None (R1-L9).""" + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.proxy._types import LiteLLM_VerificationToken, UserAPIKeyAuth + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _can_user_query_key_info, + ) + + key_info = LiteLLM_VerificationToken( + token="raw", user_id=None, team_id=None, models=[] + ) + user_api_key_dict = UserAPIKeyAuth(api_key="sk-something", user_id=None) + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.user_belongs_to_keys_team", + new=AsyncMock(return_value=False), + ): + allowed = await _can_user_query_key_info( + user_api_key_dict=user_api_key_dict, + key=None, + key_info=key_info, + ) + assert allowed is False + + +@pytest.mark.asyncio +async def test_key_info_owner_match_when_both_ids_set(): + from litellm.proxy._types import LiteLLM_VerificationToken, UserAPIKeyAuth + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _can_user_query_key_info, + ) + + key_info = LiteLLM_VerificationToken( + token="raw", user_id="alice", team_id=None, models=[] + ) + user_api_key_dict = UserAPIKeyAuth(api_key="sk-alice", user_id="alice") + allowed = await _can_user_query_key_info( + user_api_key_dict=user_api_key_dict, key=None, key_info=key_info + ) + assert allowed is True def test_is_team_key() -> None: assert _is_team_key(GenerateKeyRequest(team_id="test_team_id")) assert not _is_team_key(GenerateKeyRequest(user_id="test_user_id")) @@ -22960,3 +23008,322 @@ def test_key_info_database_not_connected_sanitized_500( +def test_key_request_changes_budget_flags_duration_and_ignores_alias(): + """budget_duration moves the reset window, so changing it counts as a + budget change; a cosmetic field change does not.""" + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _key_request_changes_budget, + ) + + existing_row = LiteLLM_VerificationToken( + token="hashed-b", + user_id="user-1", + models=["gpt-4"], + max_budget=None, + budget_duration=None, + ) + + assert ( + _key_request_changes_budget( + data=UpdateKeyRequest(key="hashed-b", budget_duration="1s"), + existing_key_row=existing_row, + ) + is True + ) + assert ( + _key_request_changes_budget( + data=UpdateKeyRequest( + key="hashed-b", max_budget=10.0, budget_duration="30d" + ), + existing_key_row=existing_row, + ) + is True + ) + assert ( + _key_request_changes_budget( + data=UpdateKeyRequest(key="hashed-b", key_alias="renamed"), + existing_key_row=existing_row, + ) + is False + ) + + +@pytest.mark.asyncio +async def test_generate_key_metadata_budget_authority_rejected_for_non_admin(): + """temp_budget_increase in metadata lifts the effective max_budget at auth + time; a non-admin must not be able to grant it to a new key via metadata.""" + with patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()): + with pytest.raises(HTTPException) as exc: + await _common_key_generation_helper( + data=GenerateKeyRequest( + metadata={ + "temp_budget_increase": 1000000, + "temp_budget_expiry": "2099-01-01T00:00:00Z", + } + ), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-alice", + user_id="alice", + ), + litellm_changed_by=None, + team_table=None, + ) + assert int(getattr(exc.value, "status_code", 0)) == 403 + assert "budget authority fields" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_generate_key_metadata_throttle_rejected_for_non_admin(): + """metadata-borne throttle_on_budget_exceeded must hit the same admin gate + as the typed field.""" + with patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()): + with pytest.raises(HTTPException) as exc: + await _common_key_generation_helper( + data=GenerateKeyRequest(metadata={"throttle_on_budget_exceeded": True}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-alice", + user_id="alice", + ), + litellm_changed_by=None, + team_table=None, + ) + assert int(getattr(exc.value, "status_code", 0)) == 403 + + +@pytest.mark.asyncio +async def test_generate_key_metadata_budget_authority_allowed_for_admin(): + with ( + patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.premium_user", False), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn" + ) as mock_generate_key, + ): + mock_generate_key.return_value = { + "key": "sk-test-key", + "expires": None, + "user_id": "admin", + "team_id": None, + } + await _common_key_generation_helper( + data=GenerateKeyRequest( + metadata={ + "temp_budget_increase": 10, + "temp_budget_expiry": "2099-01-01T00:00:00Z", + } + ), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1" + ), + litellm_changed_by=None, + team_table=None, + ) + assert mock_generate_key.called + observed: Final = str(mock_generate_key.call_args) + # the observable: the admin-approved request (temp-budget metadata style) + # survived the whole helper pipeline without a non-admin gate rejection + assert "temp_budget_increase" in observed + + +class TestKeyUpdateBudgetAuthorityGates: + """The /key/update gate set must treat budget_duration and metadata-borne + budget authority the same way max_budget/spend are treated.""" + + def _member_key_row(self): + return LiteLLM_VerificationToken( + token="hashed_member_key", + user_id="member-1", + created_by="member-1", + team_id=None, + key_alias="member", + models=["m"], + max_budget=10.0, + budget_duration="30d", + metadata={"throttle_on_budget_exceeded": False}, + ) + + def _caller(self): + return UserAPIKeyAuth( + user_id="member-1", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + def _setup(self, monkeypatch): + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", + AsyncMock(return_value=None), + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"team_admin_editable_team_fields": []}, + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks", + SimpleNamespace( + can_team_member_execute_key_management_endpoint=AsyncMock(return_value=None), + enforce_member_can_assign_access_groups=MagicMock(return_value=None), + ), + ) + + @pytest.mark.asyncio + async def test_owner_budget_duration_shrink_requires_admin(self, monkeypatch): + """A key owner shrinking their own budget_duration to seconds must hit + the admin check, not silently reset the spend window every second.""" + self._setup(monkeypatch) + admin_check = AsyncMock(return_value=None) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints._check_key_admin_access", + admin_check, + ) + await _validate_update_key_data( + data=UpdateKeyRequest(key="hashed_member_key", budget_duration="1s"), + existing_key_row=self._member_key_row(), + user_api_key_dict=self._caller(), + llm_router=None, + premium_user=True, + prisma_client=AsyncMock(), + user_api_key_cache=MagicMock(), + ) + assert admin_check.called + assert ( + "/key/update (max_budget/spend)" + in admin_check.call_args.kwargs.get("route", "") + ) + + @pytest.mark.asyncio + async def test_owner_non_budget_change_still_skips_admin_check(self, monkeypatch): + """Guardrail on the tightening: non-budget edits keep flowing without + the admin check (team member / creator paths stay permissive).""" + self._setup(monkeypatch) + admin_check = AsyncMock(return_value=None) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints._check_key_admin_access", + admin_check, + ) + await _validate_update_key_data( + data=UpdateKeyRequest(key="hashed_member_key", key_alias="renamed"), + existing_key_row=self._member_key_row(), + user_api_key_dict=self._caller(), + llm_router=None, + premium_user=True, + prisma_client=AsyncMock(), + user_api_key_cache=MagicMock(), + ) + assert not admin_check.called + + @pytest.mark.asyncio + async def test_owner_metadata_temp_budget_rejected(self, monkeypatch): + self._setup(monkeypatch) + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + data=UpdateKeyRequest( + key="hashed_member_key", + metadata={ + "temp_budget_increase": 1000000, + "temp_budget_expiry": "2099-01-01T00:00:00Z", + }, + ), + existing_key_row=self._member_key_row(), + user_api_key_dict=self._caller(), + llm_router=None, + premium_user=True, + prisma_client=AsyncMock(), + user_api_key_cache=MagicMock(), + ) + + assert int(getattr(exc.value, "status_code", 0)) == 403 + assert "budget authority fields" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_key_regeneration_budget_change_requires_admin(monkeypatch): + """A key owner regenerating their own key may regenerate the token, but + must not quietly reset spend or lift max_budget while at it.""" + from litellm.proxy._types import RegenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _execute_virtual_key_regeneration, + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _make_regenerate_mock_prisma()) + existing_key = _make_regenerate_existing_key() + user_api_key_cache = UserApiKeyCache() + caller = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-alice", + user_id="user-1", + ) + + # The real _check_key_admin_access runs: its lookup resolves the target key + # through the mocked prisma client, its team step is skipped because the key + # has no team, and the internal user is told the route is admin-only. + prisma_client = _make_regenerate_mock_prisma() + prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=_make_regenerate_existing_key() + ) + + with ( + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook" + ), + ): + with pytest.raises(HTTPException) as exc: + await _execute_virtual_key_regeneration( + prisma_client=prisma_client, + key_in_db=existing_key, + hashed_api_key="abc123", + key="abc123", + data=RegenerateKeyRequest(spend=0.0, max_budget=100000.0), + user_api_key_dict=caller, + litellm_changed_by=None, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=AsyncMock(), + ) + + # The caller observes the rejection itself: a budget-moving regenerate for a + # non-admin cannot silently reset spend or lift max_budget. + assert int(getattr(exc.value, "status_code", 0)) == 403 + assert "/key/regenerate (max_budget/spend)" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_key_regeneration_metadata_budget_authority_rejected(monkeypatch): + from litellm.proxy._types import RegenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _execute_virtual_key_regeneration, + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _make_regenerate_mock_prisma()) + + with pytest.raises(HTTPException) as exc: + await _execute_virtual_key_regeneration( + prisma_client=_make_regenerate_mock_prisma(), + key_in_db=_make_regenerate_existing_key(), + hashed_api_key="abc123", + key="abc123", + data=RegenerateKeyRequest(metadata={"temp_budget_increase": 10}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-alice", + user_id="user-1", + ), + litellm_changed_by=None, + user_api_key_cache=UserApiKeyCache(), + proxy_logging_obj=AsyncMock(), + ) + assert int(getattr(exc.value, "status_code", 0)) == 403