scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-26 06:03:19 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 35

Merge 9888723066 into 581b7d1a02

Browse source
This commit is contained in:
Jack Danger 2026-10-19 07:36:58 +00:00 • committed by GitHub
commit f3589de9e4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 558 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/proxy/auth/login_utils.py
View file

@ -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,

198
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -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,

2
litellm/proxy/proxy_server.py
View file

@ -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,

367
tests/unit/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -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