diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2f1dc270cf6..7e159ec90e7 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -26,6 +26,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypeV import fastapi import yaml from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, status +from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict import litellm @@ -99,6 +100,11 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( _add_model_to_db, ) from litellm.proxy.management_endpoints.router_weights import validate_router_settings_weights +from litellm.proxy.management_endpoints.team_admin_field_permissions import ( + team_admin_key_edit_verdict, + team_admin_key_request_or_raise, + team_admin_may_edit_member_key_budgets, +) from litellm.proxy.management_helpers.access_group_key_sync import ( sync_key_access_group_membership, sync_key_regeneration_access_group_membership, @@ -3008,6 +3014,55 @@ async def _validate_end_user_budget_id_change( raise HTTPException(status_code=400, detail=missing_detail) +_GENERAL_SETTINGS: Final = TypeAdapter(dict[str, object]) + + +def _general_settings() -> Mapping[str, object]: + from litellm.proxy.proxy_server import ( + general_settings, # pyright: ignore[reportUnknownVariableType] # untyped module-level dict in proxy_server + ) + + return _GENERAL_SETTINGS.validate_python(general_settings) + + +async def _acting_as_team_admin_for_key_update( + data: UpdateKeyRequest, + existing_key_row: LiteLLM_VerificationToken, + user_api_key_dict: UserAPIKeyAuth, + checked_prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + is_proxy_admin: bool, +) -> bool: + """Whether the caller acts as a team admin on another member's team key. + + Raises 403 when the caller administers the key's team but the request edits fields + outside the member_key_budgets permission (or that permission is disabled). + """ + if ( + is_proxy_admin + or existing_key_row.team_id is None + or existing_key_row.user_id is None + or existing_key_row.user_id == user_api_key_dict.user_id + ): + return False + team_for_grant: Final = await get_team_object( + team_id=existing_key_row.team_id, + prisma_client=checked_prisma_client, + user_api_key_cache=user_api_key_cache, + check_db_only=True, + ) + if not _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_for_grant): + return False + team_admin_key_request_or_raise( + team_admin_key_edit_verdict( + data=data, + existing=existing_key_row, + enabled=team_admin_may_edit_member_key_budgets(_general_settings()), + ) + ) + return True + + async def _validate_update_key_data( data: UpdateKeyRequest, existing_key_row: LiteLLM_VerificationToken, @@ -3058,10 +3113,19 @@ async def _validate_update_key_data( ) is_project_change: Final = "project_id" in data.model_fields_set and data.project_id != existing_key_row.project_id + acting_as_team_admin: Final = await _acting_as_team_admin_for_key_update( + data=data, + existing_key_row=existing_key_row, + user_api_key_dict=user_api_key_dict, + checked_prisma_client=checked_prisma_client, + user_api_key_cache=user_api_key_cache, + is_proxy_admin=_is_proxy_admin, + ) + common_key_access_checks( user_api_key_dict=user_api_key_dict, data=data, - user_id=existing_key_row.user_id, + user_id=user_api_key_dict.user_id if acting_as_team_admin else existing_key_row.user_id, llm_router=llm_router, premium_user=premium_user, ) diff --git a/litellm/proxy/management_endpoints/team_admin_field_permissions.py b/litellm/proxy/management_endpoints/team_admin_field_permissions.py index 6038775d96b..5146f5e0979 100644 --- a/litellm/proxy/management_endpoints/team_admin_field_permissions.py +++ b/litellm/proxy/management_endpoints/team_admin_field_permissions.py @@ -1,5 +1,6 @@ """Proxy-wide allow-list of what a team admin may do on the teams they administer: team-settings fields on -/team/update, plus the ``projects`` permission for /project/new and /project/update.""" +/team/update, the ``projects`` permission for /project/new and /project/update, and the +``member_key_budgets`` permission for budget fields on other members' keys via /key/update.""" from collections.abc import Mapping from dataclasses import dataclass @@ -12,9 +13,11 @@ from typing_extensions import assert_never from litellm._logging import verbose_proxy_logger from litellm.models.team import LiteLLM_TeamTable +from litellm.models.verification_token import LiteLLM_VerificationToken from litellm.proxy._types import ( LiteLLM_ManagementEndpoint_MetadataFields, LiteLLM_ManagementEndpoint_MetadataFields_Premium, + UpdateKeyRequest, UpdateTeamRequest, ) @@ -23,12 +26,20 @@ TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING: Final = "team_admin_editable_team_field # TODO(LIT-5722): add the remaining team settings one per PR, each with its value-diff tests and dashboard field SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS: Final[frozenset[str]] = frozenset({"tpm_limit", "rpm_limit", "max_budget"}) TEAM_ADMIN_PROJECTS_PERMISSION: Final = "projects" +TEAM_ADMIN_MEMBER_KEY_BUDGETS_PERMISSION: Final = "member_key_budgets" SUPPORTED_TEAM_ADMIN_PERMISSIONS: Final[frozenset[str]] = SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS | { - TEAM_ADMIN_PROJECTS_PERMISSION + TEAM_ADMIN_PROJECTS_PERMISSION, + TEAM_ADMIN_MEMBER_KEY_BUDGETS_PERMISSION, } +# spend is deliberately excluded: the stored row lags the live cross-pod counter, so a value-diff gate +# would let a team admin overwrite real usage. +KEY_BUDGET_FIELDS: Final[frozenset[str]] = frozenset({"max_budget", "budget_duration", "soft_budget", "budget_limits"}) +_KEY_REQUEST_IDENTITY: Final[frozenset[str]] = frozenset({"key", "token", "metadata"}) + _FIELD_LIST: Final = TypeAdapter(list[str]) _JSON_OBJECT: Final = TypeAdapter(dict[str, object]) +_WINDOW_LIST: Final = TypeAdapter(list[dict[str, object]]) _EMPTY: Final[Mapping[str, object]] = MappingProxyType({}) _METADATA_FOLDED_FIELDS: Final[frozenset[str]] = frozenset( (*LiteLLM_ManagementEndpoint_MetadataFields, *LiteLLM_ManagementEndpoint_MetadataFields_Premium) @@ -89,6 +100,12 @@ def team_admin_may_manage_projects(general_settings: Mapping[str, object]) -> bo ) +def team_admin_may_edit_member_key_budgets(general_settings: Mapping[str, object]) -> bool: + return TEAM_ADMIN_MEMBER_KEY_BUDGETS_PERMISSION in resolve_team_admin_editable_fields( + general_settings, frozenset({TEAM_ADMIN_MEMBER_KEY_BUDGETS_PERMISSION}) + ) + + def _as_object(value: object) -> Mapping[str, object]: try: return _JSON_OBJECT.validate_json(value) if isinstance(value, str) else _JSON_OBJECT.validate_python(value) @@ -101,7 +118,7 @@ def _stored_metadata(existing: Mapping[str, object]) -> Mapping[str, object]: def _submitted_metadata( - data: UpdateTeamRequest, submitted: Mapping[str, object], existing: Mapping[str, object] + data: UpdateTeamRequest | UpdateKeyRequest, submitted: Mapping[str, object], existing: Mapping[str, object] ) -> Mapping[str, object]: """Metadata as it would be stored: the caller's dict (or the stored one) with top-level folded fields laid over.""" base: Final = ( @@ -112,7 +129,7 @@ def _submitted_metadata( def _metadata_changes( - data: UpdateTeamRequest, submitted: Mapping[str, object], existing: Mapping[str, object] + data: UpdateTeamRequest | UpdateKeyRequest, submitted: Mapping[str, object], existing: Mapping[str, object] ) -> frozenset[str]: merged: Final = _submitted_metadata(data, submitted, existing) stored: Final = _stored_metadata(existing) @@ -200,3 +217,109 @@ def team_admin_request_or_raise(verdict: TeamAdminEditVerdict) -> UpdateTeamRequ ) case _: assert_never(verdict) + + +def _budget_windows(value: object) -> frozenset[tuple[object, object]] | None: + """(budget_duration, max_budget) pairs for a stored or submitted budget_limits value. + + Stored windows carry server-added keys like ``reset_at``; only the caller-owned pair matters. + ``None`` means the value is not a list of windows and needs a plain comparison. + """ + if value is None: + return frozenset() + if not isinstance(value, list): + return None + try: + windows_input: Final = _WINDOW_LIST.validate_python(value) + except ValidationError: + return None + windows: Final = frozenset((window.get("budget_duration"), window.get("max_budget")) for window in windows_input) + if len(windows) != len(windows_input): + return None + return windows + + +def _key_column_changed(field: str, submitted: Mapping[str, object], existing: Mapping[str, object]) -> bool: + if field == "budget_limits": + sent: Final = _budget_windows(submitted.get(field)) + stored: Final = _budget_windows(existing.get(field)) + if sent is not None and stored is not None: + return sent != stored + if field in LiteLLM_VerificationToken.model_fields: + return submitted.get(field) != existing.get(field) + return True + + +def changed_key_fields(data: UpdateKeyRequest, existing_row: LiteLLM_VerificationToken) -> frozenset[str]: + """Logical field names whose stored value the key-update request would change. + + Same JSON-value comparison as :func:`changed_team_fields`: columns compare against the stored row, + fields the key endpoint folds into ``metadata`` compare against ``existing_row.metadata``, other + ``metadata`` keys are attributed to ``metadata``, and fields with no stored counterpart count as + changed whenever they are sent. ``budget_limits`` compares (budget_duration, max_budget) pairs so + order and server-computed ``reset_at`` values do not read as edits. + """ + submitted: Final = _JSON_OBJECT.validate_json(data.model_dump_json(exclude_unset=True)) + existing: Final = _JSON_OBJECT.validate_json(existing_row.model_dump_json()) + column_fields: Final = frozenset(data.model_fields_set) - _KEY_REQUEST_IDENTITY - _METADATA_FOLDED_FIELDS + column_changes: Final = frozenset( + field for field in column_fields if _key_column_changed(field, submitted, existing) + ) + return column_changes | _metadata_changes(data, submitted, existing) + + +@dataclass(frozen=True, slots=True) +class TeamAdminKeyEditAllowed: + changed: frozenset[str] + kind: Literal["allowed"] = "allowed" + + +@dataclass(frozen=True, slots=True) +class TeamAdminMemberKeyEditingDisabled: + kind: Literal["disabled"] = "disabled" + + +TeamAdminKeyEditVerdict: TypeAlias = ( + TeamAdminKeyEditAllowed | TeamAdminMemberKeyEditingDisabled | TeamAdminFieldNotPermitted +) + + +def team_admin_key_edit_verdict( + data: UpdateKeyRequest, + existing: LiteLLM_VerificationToken, + enabled: bool, +) -> TeamAdminKeyEditVerdict: + if not enabled: + return TeamAdminMemberKeyEditingDisabled() + changed: Final = changed_key_fields(data, existing) + blocked: Final = sorted( + (changed | (frozenset({"spend"}) if "spend" in data.model_fields_set else frozenset())) - KEY_BUDGET_FIELDS + ) + if blocked: + return TeamAdminFieldNotPermitted(field=blocked[0]) + return TeamAdminKeyEditAllowed(changed=changed) + + +def team_admin_key_request_or_raise(verdict: TeamAdminKeyEditVerdict) -> None: + match verdict: + case TeamAdminKeyEditAllowed(): + return + case TeamAdminMemberKeyEditingDisabled(): + raise HTTPException( + status_code=403, + detail=( + "Team admins on this proxy cannot update budgets on other members' keys. " + f"Ask a proxy admin to enable '{TEAM_ADMIN_MEMBER_KEY_BUDGETS_PERMISSION}' " + f"under {_SETTINGS_LOCATION}." + ), + ) + case TeamAdminFieldNotPermitted(field=field): + raise HTTPException( + status_code=403, + detail=( + "Team admins on this proxy may only update budget fields on other members' keys, " + f"not '{field}'. Ask a proxy admin to add it under {_SETTINGS_LOCATION}." + ), + ) + case _: + assert_never(verdict) diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index d520177965c..c91b1afd64a 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -328,6 +328,7 @@ class UISettings(BaseModel): description=( "Team settings fields a team admin may change on the teams they administer. " "Include 'projects' to let team admins create and update projects for those teams. " + "Include 'member_key_budgets' to let team admins update budget fields on keys owned by other members of those teams. " "Empty means team admins cannot edit team settings or manage projects at all. " "Proxy admins and org admins are not affected." ), diff --git a/tests/integration/authorization/test_warmed_policy.py b/tests/integration/authorization/test_warmed_policy.py index a9bee196ddd..b03610c894b 100644 --- a/tests/integration/authorization/test_warmed_policy.py +++ b/tests/integration/authorization/test_warmed_policy.py @@ -1,14 +1,14 @@ +import os from collections.abc import Iterator from contextlib import ExitStack, contextmanager from hashlib import sha256 from typing import Final -import os import psycopg import pytest -from pydantic import JsonValue from hypothesis import strategies as st from hypothesis.stateful import RuleBasedStateMachine, invariant, rule, run_state_machine_as_test +from pydantic import JsonValue from tests.integration._support.client import Gateway, eventually, object_value from tests.integration._support.database import read_rows @@ -195,6 +195,71 @@ def test_warmed_team_role_demotion_prevents_later_management_writes(gateway: Gat assert_serving(gateway, model, caller, 200) +def _key_row(key: str) -> dict[str, JsonValue]: + rows: Final = read_rows( + 'SELECT max_budget, key_alias FROM "LiteLLM_VerificationToken" WHERE token = %s', + (sha256(key.encode()).hexdigest(),), + ) + assert len(rows) == 1 + return rows[0] + + +@pytest.mark.covers("mgmt.key.update.team_admin_member_key_budget_requires_opt_in") +def test_team_admin_changes_member_key_budget_only_when_opted_in(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model() + admin: Final = scenario.user(user_role="internal_user") + member: Final = scenario.user(user_role="internal_user") + team: Final = scenario.team( + models=[model], + members_with_roles=[{"user_id": admin, "role": "admin"}, {"user_id": member, "role": "user"}], + ) + other_team: Final = scenario.team(models=[model], members_with_roles=[{"user_id": member, "role": "user"}]) + member_key: Final = scenario.key( + user_id=member, team_id=team, models=[model], max_budget=10, key_alias="member" + ) + personal_key: Final = scenario.key(user_id=member, models=[model], max_budget=10) + foreign_key: Final = scenario.key(user_id=member, team_id=other_team, models=[model], max_budget=10) + admin_key: Final = scenario.key( + user_id=admin, team_id=team, models=[model], allowed_routes=["/key/update", "/v1/chat/completions"] + ) + member_caller: Final = scenario.key( + user_id=member, team_id=team, models=[model], allowed_routes=["/key/update", "/v1/chat/completions"] + ) + assert_serving(gateway, model, member_key, 200) + with _team_admins_may_edit(gateway, []): + denied: Final = gateway.request("POST", "/key/update", {"key": member_key, "max_budget": 0}, key=admin_key) + assert denied.status_code == 403, denied.text + assert _key_row(member_key) == {"max_budget": 10.0, "key_alias": "member"} + with _team_admins_may_edit(gateway, ["member_key_budgets"]): + for target in (personal_key, foreign_key): + out_of_scope: Final = gateway.request( + "POST", "/key/update", {"key": target, "max_budget": 0}, key=admin_key + ) + assert out_of_scope.status_code == 403, out_of_scope.text + assert _key_row(target)["max_budget"] == 10.0 + by_member: Final = gateway.request( + "POST", "/key/update", {"key": admin_key, "max_budget": 0}, key=member_caller + ) + assert by_member.status_code == 403, by_member.text + not_budget: Final = gateway.request( + "POST", "/key/update", {"key": member_key, "key_alias": "renamed"}, key=admin_key + ) + assert not_budget.status_code == 403, not_budget.text + assert _key_row(member_key) == {"max_budget": 10.0, "key_alias": "member"} + changed: Final = gateway.request( + "POST", "/key/update", {"key": member_key, "max_budget": 0, "budget_duration": "30d"}, key=admin_key + ) + assert changed.status_code == 200, changed.text + assert _key_row(member_key) == {"max_budget": 0.0, "key_alias": "member"} + assert_serving(gateway, model, member_key, 422, "budget_exceeded") + restored: Final = gateway.request( + "POST", "/key/update", {"key": member_key, "max_budget": 10}, key=admin_key + ) + assert restored.status_code == 200, restored.text + assert_serving(gateway, model, member_key, 200) + + @pytest.mark.covers("mgmt.key.update.expiry_changes_reach_warmed_workers") def test_expiry_and_explicit_clear_reach_both_warmed_workers(gateway: Gateway, peer: Gateway) -> None: with gateway.scenario() as scenario: diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 3b86f1f6d20..aa6be328f4a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -20931,3 +20931,165 @@ async def test_key_update_evicts_object_permission_before_key_object(monkeypatch assert deleted.index(object_permission_cache_key(permission_id)) < deleted.index( _hash_token_if_needed("sk-lit5479") ), deleted + + +class TestTeamAdminMemberKeyBudgetUpdate: + """LIT-5647: a team admin may update budget fields on another member's team key + only when the proxy enables the 'member_key_budgets' permission.""" + + def _member_key_row(self): + return LiteLLM_VerificationToken( + token="hashed_member_key", + user_id="member-1", + team_id="team-1", + key_alias="member", + models=["m"], + max_budget=10.0, + metadata={}, + ) + + def _caller(self, user_id="team-admin-1"): + return UserAPIKeyAuth( + user_id=user_id, + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + def _team(self, members): + return LiteLLM_TeamTableCachedObj(team_id="team-1", members_with_roles=members) + + def _setup(self, monkeypatch, team_obj, editable_fields): + mock_get_team = AsyncMock(return_value=team_obj) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", + mock_get_team, + ) + monkeypatch.setattr( + "litellm.proxy.management_helpers.team_member_permission_checks.get_team_object", + mock_get_team, + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"team_admin_editable_team_fields": editable_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), + ), + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints._check_team_key_limits", + AsyncMock(return_value=None), + ) + + @pytest.mark.asyncio + async def test_team_admin_updates_member_key_budget_when_enabled(self, monkeypatch): + self._setup( + monkeypatch, + self._team([Member(user_id="team-admin-1", role="admin"), Member(user_id="member-1", role="user")]), + ["member_key_budgets"], + ) + 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="sk-member", max_budget=0, budget_duration="30d"), + 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(), + ) + admin_check.assert_called_once() + + @pytest.mark.asyncio + async def test_team_admin_denied_when_permission_disabled(self, monkeypatch): + self._setup( + monkeypatch, + self._team([Member(user_id="team-admin-1", role="admin"), Member(user_id="member-1", role="user")]), + ["tpm_limit"], + ) + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + data=UpdateKeyRequest(key="sk-member", max_budget=0), + 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 exc.value.status_code == 403 + assert "member_key_budgets" in str(exc.value.detail) + assert "only create keys for themselves" not in str(exc.value.detail) + + @pytest.mark.asyncio + async def test_enabled_but_non_budget_field_is_denied(self, monkeypatch): + self._setup( + monkeypatch, + self._team([Member(user_id="team-admin-1", role="admin"), Member(user_id="member-1", role="user")]), + ["member_key_budgets"], + ) + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + data=UpdateKeyRequest(key="sk-member", 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 exc.value.status_code == 403 + assert "'key_alias'" in str(exc.value.detail) + + @pytest.mark.asyncio + async def test_ordinary_member_still_denied_on_another_members_key(self, monkeypatch): + self._setup( + monkeypatch, + self._team([Member(user_id="member-2", role="user"), Member(user_id="member-1", role="user")]), + ["member_key_budgets"], + ) + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + data=UpdateKeyRequest(key="sk-member", max_budget=0), + existing_key_row=self._member_key_row(), + user_api_key_dict=self._caller(user_id="member-2"), + llm_router=None, + premium_user=True, + prisma_client=AsyncMock(), + user_api_key_cache=MagicMock(), + ) + assert exc.value.status_code == 403 + assert "member_key_budgets" not in str(exc.value.detail) + + @pytest.mark.asyncio + async def test_personal_key_owned_by_someone_else_still_denied(self, monkeypatch): + self._setup( + monkeypatch, + self._team([Member(user_id="team-admin-1", role="admin")]), + ["member_key_budgets"], + ) + personal_row = LiteLLM_VerificationToken( + token="hashed_personal", + user_id="member-1", + team_id=None, + max_budget=10.0, + metadata={}, + ) + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + data=UpdateKeyRequest(key="sk-personal", max_budget=0), + existing_key_row=personal_row, + user_api_key_dict=self._caller(), + llm_router=None, + premium_user=True, + prisma_client=AsyncMock(), + user_api_key_cache=MagicMock(), + ) + assert exc.value.status_code == 403 + assert "member_key_budgets" not in str(exc.value.detail) diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_admin_field_permissions.py b/tests/test_litellm/proxy/management_endpoints/test_team_admin_field_permissions.py index 1a72d1de393..02cda355621 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_admin_field_permissions.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_admin_field_permissions.py @@ -1,14 +1,22 @@ import pytest from fastapi import HTTPException -from litellm.proxy._types import LiteLLM_ModelTable, LiteLLM_TeamTable, UpdateTeamRequest +from litellm.models.team import BudgetLimitEntry +from litellm.models.verification_token import LiteLLM_VerificationToken +from litellm.proxy._types import LiteLLM_ModelTable, LiteLLM_TeamTable, UpdateKeyRequest, UpdateTeamRequest from litellm.proxy.management_endpoints.team_admin_field_permissions import ( TeamAdminEditAllowed, TeamAdminEditingDisabled, TeamAdminFieldNotPermitted, + TeamAdminKeyEditAllowed, + TeamAdminMemberKeyEditingDisabled, + changed_key_fields, changed_team_fields, resolve_team_admin_editable_fields, team_admin_edit_verdict, + team_admin_key_edit_verdict, + team_admin_key_request_or_raise, + team_admin_may_edit_member_key_budgets, team_admin_may_manage_projects, team_admin_request_or_raise, ) @@ -156,3 +164,125 @@ class TestTeamAdminRequestOrRaise: team_admin_request_or_raise(TeamAdminFieldNotPermitted(field="blocked")) assert exc.value.status_code == 403 assert "'blocked'" in exc.value.detail + + +def _key(**overrides): + return LiteLLM_VerificationToken(token="hashed", **overrides) + + +class TestTeamAdminMayEditMemberKeyBudgets: + def test_missing_setting_denies(self): + assert team_admin_may_edit_member_key_budgets({}) is False + + def test_team_fields_alone_do_not_grant(self): + configured = {"team_admin_editable_team_fields": ["tpm_limit", "max_budget", "projects"]} + assert team_admin_may_edit_member_key_budgets(configured) is False + + def test_member_key_budgets_entry_grants(self): + configured = {"team_admin_editable_team_fields": ["member_key_budgets"]} + assert team_admin_may_edit_member_key_budgets(configured) is True + + @pytest.mark.parametrize("raw", ["member_key_budgets", 7, [1, 2]]) + def test_malformed_setting_denies(self, raw): + assert team_admin_may_edit_member_key_budgets({"team_admin_editable_team_fields": raw}) is False + + +class TestChangedKeyFields: + def test_key_alone_changes_nothing(self): + assert changed_key_fields(UpdateKeyRequest(key="sk-1"), _key()) == frozenset() + + def test_columns_echoing_stored_values_are_not_a_change(self): + data = UpdateKeyRequest(key="sk-1", max_budget=10.0, models=["m"], tpm_limit=5) + existing = _key(max_budget=10.0, models=["m"], tpm_limit=5) + assert changed_key_fields(data, existing) == frozenset() + + def test_column_with_different_value_is_a_change(self): + data = UpdateKeyRequest(key="sk-1", max_budget=0) + assert changed_key_fields(data, _key(max_budget=10.0)) == frozenset({"max_budget"}) + + def test_metadata_folded_field_echo_is_not_a_change(self): + data = UpdateKeyRequest(key="sk-1", tag_rpm_limit={"fast": 3}) + existing = _key(metadata={"tag_rpm_limit": {"fast": 3}}) + assert changed_key_fields(data, existing) == frozenset() + + def test_metadata_folded_field_difference_is_named_not_metadata(self): + data = UpdateKeyRequest(key="sk-1", tag_rpm_limit={"fast": 4}) + existing = _key(metadata={"tag_rpm_limit": {"fast": 3}}) + assert changed_key_fields(data, existing) == frozenset({"tag_rpm_limit"}) + + def test_budget_limits_echo_ignores_order_and_reset_at(self): + windows = [ + {"budget_duration": "1d", "max_budget": 5.0, "reset_at": "2030-01-01T00:00:00"}, + {"budget_duration": "7d", "max_budget": 50.0, "reset_at": "2030-01-07T00:00:00"}, + ] + data = UpdateKeyRequest( + key="sk-1", + budget_limits=[ + BudgetLimitEntry(budget_duration="7d", max_budget=50.0), + BudgetLimitEntry(budget_duration="1d", max_budget=5.0), + ], + ) + assert changed_key_fields(data, _key(budget_limits=windows)) == frozenset() + + def test_budget_limits_difference_is_a_change(self): + data = UpdateKeyRequest(key="sk-1", budget_limits=[BudgetLimitEntry(budget_duration="1d", max_budget=9.0)]) + existing = _key(budget_limits=[{"budget_duration": "1d", "max_budget": 5.0, "reset_at": "2030-01-01"}]) + assert changed_key_fields(data, existing) == frozenset({"budget_limits"}) + + def test_explicit_null_clearing_a_stored_column_is_a_change(self): + data = UpdateKeyRequest(key="sk-1", budget_duration=None) + assert changed_key_fields(data, _key(budget_duration="30d")) == frozenset({"budget_duration"}) + + def test_field_without_a_stored_counterpart_counts_as_changed_when_sent(self): + data = UpdateKeyRequest(key="sk-1", duration="1h") + assert changed_key_fields(data, _key()) == frozenset({"duration"}) + + +class TestTeamAdminKeyEditVerdict: + def test_disabled_even_for_a_no_op(self): + verdict = team_admin_key_edit_verdict(UpdateKeyRequest(key="sk-1"), _key(), enabled=False) + assert verdict == TeamAdminMemberKeyEditingDisabled() + + def test_budget_only_change_is_allowed(self): + data = UpdateKeyRequest(key="sk-1", max_budget=0, budget_duration="30d") + verdict = team_admin_key_edit_verdict(data, _key(max_budget=10.0), enabled=True) + assert verdict == TeamAdminKeyEditAllowed(changed=frozenset({"max_budget", "budget_duration"})) + + def test_key_alias_change_is_blocked_and_named(self): + data = UpdateKeyRequest(key="sk-1", key_alias="renamed") + verdict = team_admin_key_edit_verdict(data, _key(key_alias="member"), enabled=True) + assert verdict == TeamAdminFieldNotPermitted(field="key_alias") + + def test_spend_is_blocked(self): + data = UpdateKeyRequest(key="sk-1", spend=0) + verdict = team_admin_key_edit_verdict(data, _key(spend=3.5), enabled=True) + assert verdict == TeamAdminFieldNotPermitted(field="spend") + + def test_spend_echo_is_blocked_even_when_unchanged(self): + data = UpdateKeyRequest(key="sk-1", spend=4.5, max_budget=0) + verdict = team_admin_key_edit_verdict(data, _key(spend=4.5, max_budget=10.0), enabled=True) + assert verdict == TeamAdminFieldNotPermitted(field="spend") + + def test_budget_plus_non_budget_names_the_non_budget_field(self): + data = UpdateKeyRequest(key="sk-1", max_budget=0, key_alias="renamed") + verdict = team_admin_key_edit_verdict(data, _key(max_budget=10.0, key_alias="member"), enabled=True) + assert verdict == TeamAdminFieldNotPermitted(field="key_alias") + + +class TestTeamAdminKeyRequestOrRaise: + def test_allowed_returns_none(self): + verdict = TeamAdminKeyEditAllowed(changed=frozenset({"max_budget"})) + assert team_admin_key_request_or_raise(verdict) is None + + def test_disabled_is_a_403_pointing_at_member_key_budgets(self): + with pytest.raises(HTTPException) as exc: + team_admin_key_request_or_raise(TeamAdminMemberKeyEditingDisabled()) + assert exc.value.status_code == 403 + assert "member_key_budgets" in exc.value.detail + assert "Settings > UI > Team admin editable fields" in exc.value.detail + + def test_field_not_permitted_is_a_403_naming_the_field(self): + with pytest.raises(HTTPException) as exc: + team_admin_key_request_or_raise(TeamAdminFieldNotPermitted(field="key_alias")) + assert exc.value.status_code == 403 + assert "'key_alias'" in exc.value.detail diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index 0140fcaba21..08d542df16c 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -3854,6 +3854,28 @@ class TestTeamAdminEditableTeamFieldsSetting: assert stored["team_admin_editable_team_fields"] == ["projects"] assert team_admin_may_manage_projects(general_settings) is True + def test_patch_accepts_the_member_key_budgets_permission(self, monkeypatch): + from litellm.proxy.management_endpoints.team_admin_field_permissions import ( + team_admin_may_edit_member_key_budgets, + ) + + mock_prisma = self._as_proxy_admin(monkeypatch) + general_settings: dict = {} + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings) + assert team_admin_may_edit_member_key_budgets(general_settings) is False + + try: + response = client.patch( + "/update/ui_settings", json={"team_admin_editable_team_fields": ["member_key_budgets"]} + ) + finally: + app.dependency_overrides.clear() + + assert response.status_code == 200 + stored = json.loads(mock_prisma.db.litellm_uisettings.upsert.call_args.kwargs["data"]["create"]["ui_settings"]) + assert stored["team_admin_editable_team_fields"] == ["member_key_budgets"] + assert team_admin_may_edit_member_key_budgets(general_settings) is True + def test_patch_with_an_empty_list_turns_team_admin_editing_off_again(self, monkeypatch): mock_prisma = self._as_proxy_admin(monkeypatch) general_settings: dict = {"team_admin_editable_team_fields": ["tpm_limit"]} diff --git a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts index 5706eeafe82..1d8cd90c517 100644 --- a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts +++ b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts @@ -48,6 +48,7 @@ const TEAM_ADMIN_FIELD_LABELS: ReadonlyMap = new Map([ ["rpm_limit", "Requests per minute Limit (RPM)"], ["max_budget", "Max Budget (USD)"], ["projects", "Create and update projects"], + ["member_key_budgets", "Update budgets on team members' keys"], ]); export const teamAdminFieldLabel = (field: string): string => TEAM_ADMIN_FIELD_LABELS.get(field) ?? field; diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx index 4d0afde6f25..9e3e19f1e1b 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx @@ -303,6 +303,7 @@ describe("KeyEditView", () => { fallbacks: [{ "gpt-4": ["gpt-4o", "gpt-4o-mini"] }], }), }), + expect.any(Array), ); }); }); @@ -323,6 +324,7 @@ describe("KeyEditView", () => { fallbacks: null, }), }), + expect.any(Array), ); }); }); @@ -632,7 +634,10 @@ describe("KeyEditView", () => { await userEvent.click(screen.getByRole("button", { name: /save changes/i })); await waitFor(() => { - expect(onSubmitMock).toHaveBeenCalledWith(expect.objectContaining({ throttle_on_budget_exceeded: true })); + expect(onSubmitMock).toHaveBeenCalledWith( + expect.objectContaining({ throttle_on_budget_exceeded: true }), + expect.any(Array), + ); }); }); @@ -662,7 +667,10 @@ describe("KeyEditView", () => { await userEvent.click(screen.getByRole("button", { name: /save changes/i })); await waitFor(() => { - expect(onSubmitMock).toHaveBeenCalledWith(expect.objectContaining({ enable_prompt_caching: true })); + expect(onSubmitMock).toHaveBeenCalledWith( + expect.objectContaining({ enable_prompt_caching: true }), + expect.any(Array), + ); }); }); @@ -1526,7 +1534,10 @@ describe("KeyEditView", () => { await userEvent.click(screen.getByRole("button", { name: /save changes/i })); await waitFor(() => { - expect(onSubmit).toHaveBeenCalledWith(expect.objectContaining({ organization_id: null, team_id: null })); + expect(onSubmit).toHaveBeenCalledWith( + expect.objectContaining({ organization_id: null, team_id: null }), + expect.any(Array), + ); }); expect(JSON.parse(JSON.stringify(onSubmit.mock.calls[0][0]))).toMatchObject({ organization_id: null, @@ -1568,7 +1579,9 @@ describe("KeyEditView", () => { await userEvent.click(await screen.findByRole("button", { name: "Detach from project" })); await userEvent.click(screen.getByRole("button", { name: /save changes/i })); const expectedDetach = { project_id: null, organization_id: "org-1", team_id: "group-maple", models: key.models }; - await waitFor(() => expect(onSubmit).toHaveBeenCalledWith(expect.objectContaining(expectedDetach))); + await waitFor(() => + expect(onSubmit).toHaveBeenCalledWith(expect.objectContaining(expectedDetach), expect.any(Array)), + ); expect(screen.getByRole("combobox", { name: "Team ID" })).toBeDisabled(); view.rerender(renderEditor({ ...key, project_id: null })); expect(screen.getByRole("combobox", { name: "Team ID" })).toBeEnabled(); @@ -1866,7 +1879,10 @@ describe("KeyEditView", () => { await save(); await waitFor(() => { - expect(onSubmit).toHaveBeenCalledWith(expect.objectContaining({ end_user_budget_id: "svc-b-budget" })); + expect(onSubmit).toHaveBeenCalledWith( + expect.objectContaining({ end_user_budget_id: "svc-b-budget" }), + expect.any(Array), + ); }); }); @@ -1878,7 +1894,7 @@ describe("KeyEditView", () => { await save(); await waitFor(() => { - expect(onSubmit).toHaveBeenCalledWith(expect.objectContaining({ end_user_budget_id: "" })); + expect(onSubmit).toHaveBeenCalledWith(expect.objectContaining({ end_user_budget_id: "" }), expect.any(Array)); }); }); diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx index 9cd97f4ef98..1342b97d1a0 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -80,7 +80,7 @@ import VectorStoreSelector from "../vector_store_management/VectorStoreSelector" interface KeyEditViewProps { keyData: KeyResponse; onCancel: () => void; - onSubmit: (values: any) => Promise; + onSubmit: (values: any, dirtyFields: readonly string[]) => Promise; teams?: any[] | null; accessToken: string | null; userID: string | null; @@ -317,6 +317,7 @@ export function KeyEditView({ ...values, ...(detachProject && enableProjectsUI && canDetachProject ? { project_id: null } : {}), }), + [...Object.keys(form.formState.dirtyFields), ...(budgetLimitsUnchanged ? [] : ["budget_limits"])], ); } finally { setIsKeySaving(false); diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index 63693fd1af5..7eb09926caf 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -49,6 +49,7 @@ import { RegenerateKeyModal } from "../organisms/RegenerateKeyModal"; import { parseErrorMessage } from "../shared/errorUtils"; import { InheritedBudgetHint, inheritedBudgetGates, keyOwnerBudgetSource } from "../shared/InheritedBudgetHint"; import { KeyEditView } from "./key_edit_view"; +import { isTeamAdminEditingMemberKey, teamAdminMemberKeyPayload } from "./teamAdminMemberKeyPayload"; export function needsLifetimeSpendBackfill(spend: number, totalSpend: number | null | undefined): boolean { return (totalSpend ?? 0) < spend; @@ -187,7 +188,7 @@ export default function KeyInfoView({ ); } - const handleKeyUpdate = async (formValues: Record) => { + const handleKeyUpdate = async (formValues: Record, dirtyFields: readonly string[] = []) => { try { if (!accessToken) return; @@ -359,6 +360,25 @@ export default function KeyInfoView({ formValues.budget_duration = wordToCanonical[formValues.budget_duration] ?? formValues.budget_duration; } + const memberKeyEditContext = { + userRole: userRole || "", + userId: userID || "", + keyUserId: currentKeyData.user_id, + keyTeamId: currentKeyData.team_id, + teamMembers: teamsData?.find((team) => team.team_id === currentKeyData.team_id)?.members_with_roles, + }; + const editingMemberKeyAsTeamAdmin = isTeamAdminEditingMemberKey(memberKeyEditContext); + if (editingMemberKeyAsTeamAdmin) { + const trimmed = teamAdminMemberKeyPayload(formValues, dirtyFields); + if (trimmed.kind === "blocked") { + toast.error( + `Team admins can only change budget fields on other members' keys, not ${trimmed.fields.join(", ")}`, + ); + return; + } + formValues = trimmed.payload; + } + const newKeyValues = await keyUpdateCall(accessToken, formValues); // Update local state diff --git a/ui/litellm-dashboard/src/components/templates/teamAdminMemberKeyPayload.test.ts b/ui/litellm-dashboard/src/components/templates/teamAdminMemberKeyPayload.test.ts new file mode 100644 index 00000000000..a43b41abeb0 --- /dev/null +++ b/ui/litellm-dashboard/src/components/templates/teamAdminMemberKeyPayload.test.ts @@ -0,0 +1,96 @@ +import { describe, expect, it } from "vitest"; +import { Member } from "@/components/networking"; +import { isTeamAdminEditingMemberKey, KEY_BUDGET_FIELDS, teamAdminMemberKeyPayload } from "./teamAdminMemberKeyPayload"; + +const members = (role: string): Member[] => [{ user_id: "admin-user", role, user_email: null } as unknown as Member]; + +const baseArgs = { + userRole: "Internal User", + userId: "admin-user", + keyUserId: "member-user", + keyTeamId: "team-1", +}; + +describe("isTeamAdminEditingMemberKey", () => { + it("is false for a proxy admin", () => { + expect(isTeamAdminEditingMemberKey({ ...baseArgs, userRole: "Admin", teamMembers: members("admin") })).toBe(false); + }); + + it("is false when the caller owns the key", () => { + expect(isTeamAdminEditingMemberKey({ ...baseArgs, keyUserId: "admin-user", teamMembers: members("admin") })).toBe( + false, + ); + }); + + it("is false for a personal key with no team", () => { + expect(isTeamAdminEditingMemberKey({ ...baseArgs, keyTeamId: null, teamMembers: members("admin") })).toBe(false); + }); + + it("is false when the caller is not a team admin", () => { + expect(isTeamAdminEditingMemberKey({ ...baseArgs, teamMembers: members("user") })).toBe(false); + expect(isTeamAdminEditingMemberKey({ ...baseArgs, teamMembers: null })).toBe(false); + expect(isTeamAdminEditingMemberKey({ ...baseArgs, teamMembers: undefined })).toBe(false); + }); + + it("is true for a team admin editing another member's team key", () => { + expect(isTeamAdminEditingMemberKey({ ...baseArgs, teamMembers: members("admin") })).toBe(true); + }); +}); + +describe("teamAdminMemberKeyPayload", () => { + it("keeps only dirty budget fields from the form values plus the key", () => { + const formValues = { + key: "sk-1", + max_budget: 25, + soft_budget: 10, + key_alias: "renamed", + metadata: { tags: ["a"] }, + tpm_limit: null, + }; + const result = teamAdminMemberKeyPayload(formValues, ["max_budget", "soft_budget"]); + expect(result).toEqual({ + kind: "ok", + payload: { key: "sk-1", max_budget: 25, soft_budget: 10 }, + }); + }); + + it("drops budget fields present in the form but not dirty", () => { + const formValues = { + key: "sk-1", + max_budget: 25, + budget_duration: "30d", + budget_limits: [{ budget_duration: "1d", max_budget: 5 }], + }; + const result = teamAdminMemberKeyPayload(formValues, ["max_budget"]); + expect(result).toEqual({ kind: "ok", payload: { key: "sk-1", max_budget: 25 } }); + }); + + it("drops budget_duration when it is an empty string but keeps null", () => { + const cleared = teamAdminMemberKeyPayload({ key: "sk-1", budget_duration: "" }, ["budget_duration"]); + expect(cleared).toEqual({ kind: "ok", payload: { key: "sk-1" } }); + const kept = teamAdminMemberKeyPayload({ key: "sk-1", budget_duration: null }, ["budget_duration"]); + expect(kept).toEqual({ kind: "ok", payload: { key: "sk-1", budget_duration: null } }); + }); + + it("keeps budget_limits when present", () => { + const windows = [{ budget_duration: "1d", max_budget: 5 }]; + const result = teamAdminMemberKeyPayload({ key: "sk-1", budget_limits: windows }, ["budget_limits"]); + expect(result).toEqual({ kind: "ok", payload: { key: "sk-1", budget_limits: windows } }); + }); + + it("is blocked when a dirty field is not a budget field, naming it", () => { + const result = teamAdminMemberKeyPayload({ key: "sk-1", key_alias: "renamed" }, ["key_alias", "max_budget"]); + expect(result).toEqual({ kind: "blocked", fields: ["key_alias"] }); + }); + + it("is ok when every dirty field is a budget field and ignores token/key", () => { + const result = teamAdminMemberKeyPayload({ key: "sk-1", max_budget: 5 }, ["token", "key", "max_budget"]); + expect(result).toEqual({ kind: "ok", payload: { key: "sk-1", max_budget: 5 } }); + }); + + it("covers exactly the backend budget field set", () => { + expect([...KEY_BUDGET_FIELDS].sort()).toEqual( + ["budget_duration", "budget_limits", "max_budget", "soft_budget"].sort(), + ); + }); +}); diff --git a/ui/litellm-dashboard/src/components/templates/teamAdminMemberKeyPayload.ts b/ui/litellm-dashboard/src/components/templates/teamAdminMemberKeyPayload.ts new file mode 100644 index 00000000000..abd8653dde5 --- /dev/null +++ b/ui/litellm-dashboard/src/components/templates/teamAdminMemberKeyPayload.ts @@ -0,0 +1,41 @@ +import { Member } from "@/components/networking"; +import { isProxyAdminRole, isUserTeamAdminForSingleTeam } from "@/utils/roles"; + +export const KEY_BUDGET_FIELDS = ["max_budget", "soft_budget", "budget_duration", "budget_limits"] as const; + +export const isTeamAdminEditingMemberKey = (args: { + userRole: string; + userId: string; + keyUserId: string | null | undefined; + keyTeamId: string | null | undefined; + teamMembers: Member[] | null | undefined; +}): boolean => { + if (isProxyAdminRole(args.userRole)) return false; + if (!args.keyTeamId) return false; + if (args.keyUserId === args.userId) return false; + return isUserTeamAdminForSingleTeam(args.teamMembers ?? null, args.userId); +}; + +export type TeamAdminMemberKeyPayload = + | { kind: "ok"; payload: Record } + | { kind: "blocked"; fields: readonly string[] }; + +export const teamAdminMemberKeyPayload = ( + formValues: Record, + dirtyFields: readonly string[], +): TeamAdminMemberKeyPayload => { + const disallowed = dirtyFields.filter( + (field) => field !== "token" && field !== "key" && !(KEY_BUDGET_FIELDS as readonly string[]).includes(field), + ); + if (disallowed.length > 0) { + return { kind: "blocked", fields: disallowed }; + } + const payload: Record = { key: formValues.key }; + for (const field of KEY_BUDGET_FIELDS) { + if (!dirtyFields.includes(field)) continue; + if (formValues[field] === undefined) continue; + if (field === "budget_duration" && formValues[field] === "") continue; + payload[field] = formValues[field]; + } + return { kind: "ok", payload }; +};