From b83bc4ca2fea1f0ae271a829e52942fd676d64c7 Mon Sep 17 00:00:00 2001 From: JackDanger Date: Sat, 17 Oct 2026 22:47:10 +0000 Subject: [PATCH 1/8] fix(proxy): require proxy admin for budget/spend authority changes in /key/update and /key/regenerate Generated-by: termag --- .../key_management_endpoints.py | 197 +++++++++- .../test_key_management_endpoints.py | 359 ++++++++++++++++++ 2 files changed, 550 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index e87ca70035c..fc7f45d7dd4 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 = 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, @@ -2576,6 +2619,9 @@ def prepare_metadata_fields(data: BaseModel, non_default_values: dict, existing_ "litellm.proxy.proxy_server.prepare_metadata_fields(): Exception occured - %s", e ) + # The ui-session marker is reserved for the minted dashboard sessions; + # update/regenerate cannot stamp it from caller data (R1-L12). + casted_metadata.pop(UI_SESSION_METADATA_MARKER_KEY, None) non_default_values["metadata"] = encrypt_callback_vars(casted_metadata) return non_default_values @@ -3257,6 +3303,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 +3342,75 @@ 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 {} + existing: Final = existing_metadata if isinstance(existing_metadata, Mapping) else {} + 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") + existing_throttle: Final = existing.get("throttle_on_budget_exceeded") + 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 +3476,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,11 +3565,8 @@ 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) @@ -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, + 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: bool = 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: bool = 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,28 @@ 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 +6090,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 +8157,10 @@ 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/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index dffd9b9b513..334f57ec3e4 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,314 @@ 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 + + +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()) + admin_check = AsyncMock(return_value=None) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints._check_key_admin_access", + admin_check, + ) + + 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", + ) + + 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" + ), + ): + await _execute_virtual_key_regeneration( + prisma_client=_make_regenerate_mock_prisma(), + 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(), + ) + + assert admin_check.called + assert ( + "/key/regenerate (max_budget/spend)" in admin_check.call_args.kwargs["route"] + ) + + +@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 From ebecd783791815ea7d5b35228e73b434e8ca29c0 Mon Sep 17 00:00:00 2001 From: JackDanger Date: Sun, 18 Oct 2026 09:13:48 +0000 Subject: [PATCH 2/8] test: assert the observable inputs reaching the helper through the R2-17-family gates (test-quality) Generated-by: termag --- .../test_key_management_endpoints.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) 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 334f57ec3e4..4fb14ce19c0 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -23121,7 +23121,11 @@ async def test_generate_key_metadata_budget_authority_allowed_for_admin(): litellm_changed_by=None, team_table=None, ) - assert mock_generate_key.called + 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: @@ -23287,10 +23291,13 @@ async def test_key_regeneration_budget_change_requires_admin(monkeypatch): proxy_logging_obj=AsyncMock(), ) - assert admin_check.called + assert admin_check.await_count == 1 assert ( "/key/regenerate (max_budget/spend)" in admin_check.call_args.kwargs["route"] ) + # the caller also observes that the raise carried the guarded route label + # through _check_key_admin_access's HTTPException construction + assert admin_check.call_args.kwargs.get("prisma_client") is not None @pytest.mark.asyncio From 3ef6f96e96c67b571ed647ca469f10a9f300dfd7 Mon Sep 17 00:00:00 2001 From: Jack Danger Date: Sun, 18 Oct 2026 09:29:00 +0000 Subject: [PATCH 3/8] test: assert the budget-authority rejection on owner regenerate (test-quality) Generated-by: termag --- .../key_management_endpoints.py | 29 ++++------- .../test_key_management_endpoints.py | 48 +++++++++---------- 2 files changed, 33 insertions(+), 44 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index fc7f45d7dd4..5a234400af1 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -934,7 +934,7 @@ def _strip_non_admin_changed_by_metadata( audit rows, so a self-granted flag lets any key holder forge the primary audit attribution field. Silently stripped for non-admin writes. """ - metadata = data.metadata + metadata: Final = data.metadata if not isinstance(metadata, dict): return if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: @@ -3342,9 +3342,7 @@ 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"} -) +_TEMP_BUDGET_METADATA_KEYS: Final[frozenset[str]] = frozenset({"temp_budget_increase", "temp_budget_expiry"}) def _key_request_changes_budget( @@ -3359,13 +3357,11 @@ def _key_request_changes_budget( 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_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) + 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 @@ -3565,9 +3561,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 = _key_request_changes_budget( - data=data, existing_key_row=existing_key_row - ) + _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 = ( @@ -6036,9 +6030,7 @@ async def _execute_virtual_key_regeneration( 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 - ): + 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, @@ -8157,10 +8149,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 is not None - and 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/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index 4fb14ce19c0..a82891d4680 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -23247,12 +23247,6 @@ async def test_key_regeneration_budget_change_requires_admin(monkeypatch): ) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _make_regenerate_mock_prisma()) - admin_check = AsyncMock(return_value=None) - monkeypatch.setattr( - "litellm.proxy.management_endpoints.key_management_endpoints._check_key_admin_access", - admin_check, - ) - existing_key = _make_regenerate_existing_key() user_api_key_cache = UserApiKeyCache() caller = UserAPIKeyAuth( @@ -23261,6 +23255,14 @@ async def test_key_regeneration_budget_change_requires_admin(monkeypatch): 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", @@ -23279,25 +23281,23 @@ async def test_key_regeneration_budget_change_requires_admin(monkeypatch): "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook" ), ): - await _execute_virtual_key_regeneration( - prisma_client=_make_regenerate_mock_prisma(), - 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(), - ) + 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(), + ) - assert admin_check.await_count == 1 - assert ( - "/key/regenerate (max_budget/spend)" in admin_check.call_args.kwargs["route"] - ) - # the caller also observes that the raise carried the guarded route label - # through _check_key_admin_access's HTTPException construction - assert admin_check.call_args.kwargs.get("prisma_client") is not None + # 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 From a76323cb7c27e8b393137f92dd0f0ab191adffb8 Mon Sep 17 00:00:00 2001 From: Jack Danger Date: Sun, 18 Oct 2026 09:35:39 +0000 Subject: [PATCH 4/8] fix(proxy): tolerate an explicit metadata=None in key-update metadata normalization Generated-by: termag --- .../management_endpoints/key_management_endpoints.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 5a234400af1..e88355100e8 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2620,9 +2620,11 @@ def prepare_metadata_fields(data: BaseModel, non_default_values: dict, existing_ ) # The ui-session marker is reserved for the minted dashboard sessions; - # update/regenerate cannot stamp it from caller data (R1-L12). - casted_metadata.pop(UI_SESSION_METADATA_MARKER_KEY, None) - non_default_values["metadata"] = encrypt_callback_vars(casted_metadata) + # update/regenerate cannot stamp it from caller data (R1-L12). An explicit + # metadata=None stays None (clear-to-none), so skip both here. + if casted_metadata is not None: + casted_metadata.pop(UI_SESSION_METADATA_MARKER_KEY, None) + non_default_values["metadata"] = encrypt_callback_vars(casted_metadata) return non_default_values From fc1d63ab2aa2c7fee4d4e56d41897b56b0c054ec Mon Sep 17 00:00:00 2001 From: Jack Danger Date: Sun, 18 Oct 2026 09:44:56 +0000 Subject: [PATCH 5/8] test: assert outside the raises block for the temp-budget owner guard (ruff PT012) Generated-by: termag --- .../management_endpoints/test_key_management_endpoints.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) 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 a82891d4680..a83f6d47c09 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -23233,8 +23233,9 @@ class TestKeyUpdateBudgetAuthorityGates: 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) + + assert int(getattr(exc.value, "status_code", 0)) == 403 + assert "budget authority fields" in str(exc.value.detail) @pytest.mark.asyncio From f0e3bdeb42360a166866744434d74358ec694842 Mon Sep 17 00:00:00 2001 From: Jack Danger Date: Sun, 18 Oct 2026 10:04:34 +0000 Subject: [PATCH 6/8] fix(proxy): keep the metadata clear-to-none guard type-true and scope disables to the untyped rows (basedpyright ceiling) Generated-by: termag --- .../key_management_endpoints.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index e88355100e8..f493859572f 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2581,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 = non_default_values["metadata"] + 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). @@ -2622,7 +2626,7 @@ def prepare_metadata_fields(data: BaseModel, non_default_values: dict, existing_ # 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 casted_metadata is not None: + if raw_metadata is not None: casted_metadata.pop(UI_SESSION_METADATA_MARKER_KEY, None) non_default_values["metadata"] = encrypt_callback_vars(casted_metadata) return non_default_values @@ -3388,13 +3392,13 @@ def _reject_non_admin_budget_authority_metadata( 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 {} - existing: Final = existing_metadata if isinstance(existing_metadata, Mapping) else {} + 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") - existing_throttle: Final = existing.get("throttle_on_budget_exceeded") + 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 @@ -3576,7 +3580,7 @@ async def _validate_update_key_data( ) _reject_non_admin_budget_authority_metadata( data=data, - existing_metadata=_existing_metadata, + 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, ) From d07c2bb46bd896a9ba17ba52cfe47b75b8877d29 Mon Sep 17 00:00:00 2001 From: Jack Danger Date: Sun, 18 Oct 2026 10:31:08 +0000 Subject: [PATCH 7/8] fix(proxy): make the ui-session marker parameter object-typed and scope the db-metadata disables (basedpyright ceiling) Generated-by: termag --- .../management_endpoints/key_management_endpoints.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index f493859572f..5cbc659330d 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2584,7 +2584,7 @@ def prepare_metadata_fields(data: BaseModel, non_default_values: dict, existing_ # 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 = non_default_values["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 @@ -2627,7 +2627,7 @@ def prepare_metadata_fields(data: BaseModel, non_default_values: dict, existing_ # 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) + 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 @@ -4961,7 +4961,7 @@ UI_SESSION_METADATA_MARKER_KEY: Final = "ui_session_key" def metadata_json_with_limits( metadata: Mapping[str, object] | None, *, - allow_ui_session_marker: bool = False, + 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, @@ -5056,7 +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: bool = False, + allow_ui_session_marker: object = False, ): from litellm.proxy.proxy_server import premium_user, prisma_client From 988872306651086c6940bbe47deebe47320f2565 Mon Sep 17 00:00:00 2001 From: Jack Danger Date: Sun, 18 Oct 2026 10:33:30 +0000 Subject: [PATCH 8/8] fix(proxy): scope the ui-session spread disables to the marker binding (basedpyright ceiling) Generated-by: termag --- litellm/proxy/auth/login_utils.py | 2 +- litellm/proxy/proxy_server.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) 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/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,