diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 1978c0a1b0b..c586d67873b 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1260,6 +1260,7 @@ class UpdateKeyRequest(KeyRequestBase): auto_rotate: bool | None = None rotation_interval: str | None = None organization_id: str | None = None + key_type: LiteLLMKeyType | None = None @model_validator(mode="after") def validate_temp_budget(self) -> "UpdateKeyRequest": diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 97999cbb6d7..384a5413eb3 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -571,7 +571,19 @@ def common_key_access_checks( router: Final = APIRouter() -def handle_key_type(data: GenerateKeyRequest, data_json: dict) -> dict: +def _key_type_allowed_routes( + key_type: LiteLLMKeyType | None, +) -> tuple[str, ...] | None: + if key_type == LiteLLMKeyType.LLM_API: + return ("llm_api_routes",) + if key_type == LiteLLMKeyType.MANAGEMENT: + return ("management_routes",) + if key_type == LiteLLMKeyType.READ_ONLY: + return ("info_routes",) + return None + + +def handle_key_type(data: GenerateKeyRequest | UpdateKeyRequest, data_json: dict) -> dict: """ Handle the key type. """ @@ -580,12 +592,9 @@ def handle_key_type(data: GenerateKeyRequest, data_json: dict) -> dict: data_json.pop("key_type", None) return data_json data_json["key_type"] = key_type.value - if key_type == LiteLLMKeyType.LLM_API: - data_json["allowed_routes"] = ["llm_api_routes"] - elif key_type == LiteLLMKeyType.MANAGEMENT: - data_json["allowed_routes"] = ["management_routes"] - elif key_type == LiteLLMKeyType.READ_ONLY: - data_json["allowed_routes"] = ["info_routes"] + allowed_routes: Final = _key_type_allowed_routes(key_type) + if allowed_routes is not None: + data_json["allowed_routes"] = list(allowed_routes) return data_json @@ -639,7 +648,7 @@ def _validate_caller_can_change_key_ownership( def _check_allowed_routes_caller_permission( - allowed_routes: list | None, + allowed_routes: Sequence[str] | None, user_api_key_dict: UserAPIKeyAuth, *, allowed_routes_was_provided: bool = False, @@ -681,6 +690,37 @@ def _check_allowed_routes_caller_permission( ) +def _check_key_type_transition_against_existing_routes( + data: BaseModel | None, + existing_key_row: LiteLLM_VerificationToken, + user_api_key_dict: UserAPIKeyAuth, +) -> None: + """Keep non-admin key-type presets from replacing custom route restrictions.""" + if ( + data is None + or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value + or "key_type" not in data.model_fields_set + or data.key_type + not in ( + LiteLLMKeyType.DEFAULT, + LiteLLMKeyType.LLM_API, + LiteLLMKeyType.READ_ONLY, + ) + ): + return + existing_allowed_routes: Final = tuple(getattr(existing_key_row, "allowed_routes", None) or ()) + if existing_allowed_routes not in ( + (), + ("llm_api_routes",), + ("info_routes",), + ): + _check_allowed_routes_caller_permission( + allowed_routes=existing_allowed_routes, + user_api_key_dict=user_api_key_dict, + allowed_routes_was_provided=True, + ) + + def _check_permissions_caller_permission( data: GenerateRequestBase, user_api_key_dict: UserAPIKeyAuth, @@ -2025,6 +2065,10 @@ async def prepare_key_update_data( existing_key_row: LiteLLM_VerificationToken, ): data_json: Final[dict] = data.model_dump(exclude_unset=True) + if "key_type" in data.model_fields_set: + handle_key_type(data, data_json) + if data.key_type == LiteLLMKeyType.DEFAULT and "allowed_routes" not in data.model_fields_set: + data_json["allowed_routes"] = [] data_json.pop("key", None) data_json.pop("new_key", None) data_json.pop("grace_period", None) # Request-only param, not a DB column @@ -2478,11 +2522,52 @@ async def _validate_update_key_data( validate_budget_duration(data.budget_duration) _is_proxy_admin: Final = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value + existing_allowed_routes: Final = tuple(getattr(existing_key_row, "allowed_routes", None) or ()) + default_key_type_clears_routes: Final = ( + "key_type" in data.model_fields_set + and data.key_type == LiteLLMKeyType.DEFAULT + and "allowed_routes" not in data.model_fields_set + ) + derived_allowed_routes: Final = _key_type_allowed_routes(data.key_type) + target_allowed_routes: Final = ( + tuple(data.allowed_routes or ()) + if "allowed_routes" in data.model_fields_set + else (() if default_key_type_clears_routes else derived_allowed_routes) + ) + safe_key_type_transition: Final = ( + target_allowed_routes + in ( + (), + ("llm_api_routes",), + ("info_routes",), + ) + and "key_type" in data.model_fields_set + and data.key_type + in ( + LiteLLMKeyType.DEFAULT, + LiteLLMKeyType.LLM_API, + LiteLLMKeyType.READ_ONLY, + ) + and existing_allowed_routes + in ( + (), + ("llm_api_routes",), + ("info_routes",), + ) + ) _check_allowed_routes_caller_permission( allowed_routes=data.allowed_routes, user_api_key_dict=user_api_key_dict, - allowed_routes_was_provided="allowed_routes" in data.model_fields_set, + allowed_routes_was_provided=( + ("allowed_routes" in data.model_fields_set or default_key_type_clears_routes) + and not safe_key_type_transition + ), + ) + _check_allowed_routes_caller_permission( + allowed_routes=derived_allowed_routes, + user_api_key_dict=user_api_key_dict, + allow_safe_presets=safe_key_type_transition, ) _check_passthrough_routes_caller_permission( data=data, @@ -2734,6 +2819,7 @@ async def update_key_fn( Parameters: - key: Optional[str] - The key to update. Either key or key_alias must be provided. - key_alias: Optional[str] - User-friendly key alias. If key is omitted, also identifies the key to update (must match exactly one key, same as /key/delete's key_aliases) + - key_type: Optional[LiteLLMKeyType] - Access preset to apply to the key. - user_id: Optional[str] - User ID associated with key - team_id: Optional[str] - Team ID associated with key - agent_id: Optional[str] - The agent id associated with the key. @@ -4731,6 +4817,11 @@ async def _execute_virtual_key_regeneration( existing_key_row=key_in_db, user_api_key_dict=user_api_key_dict, ) + _check_key_type_transition_against_existing_routes( + data=data, + existing_key_row=key_in_db, + 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, 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 6045b64023d..bf14f82d114 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 @@ -13,8 +13,8 @@ import inspect from litellm.proxy._types import ( GenerateKeyRequest, - NewUserRequest, LiteLLM_BudgetTable, + LiteLLMKeyType, LiteLLM_OrganizationTable, LiteLLM_ProjectTableCachedObj, LiteLLM_TeamTableCachedObj, @@ -22,6 +22,7 @@ from litellm.proxy._types import ( LiteLLM_VerificationToken, LitellmUserRoles, Member, + NewUserRequest, ProxyException, ResetSpendRequest, UpdateKeyRequest, @@ -11349,6 +11350,49 @@ def _make_regenerate_existing_key(): ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key_type", + [ + LiteLLMKeyType.DEFAULT, + LiteLLMKeyType.LLM_API, + LiteLLMKeyType.READ_ONLY, + ], +) +async def test_execute_virtual_key_regeneration_rejects_preset_over_custom_routes( + key_type, +): + """A non-admin regeneration preset must preserve a custom route allowlist.""" + from litellm.proxy._types import RegenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _execute_virtual_key_regeneration, + ) + + existing_key = _make_regenerate_existing_key() + existing_key.allowed_routes = ["/custom/admin-defined-route"] + mock_prisma_client = _make_regenerate_mock_prisma() + + with pytest.raises(HTTPException) as exc_info: + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=existing_key, + hashed_api_key="abc123", + key="abc123", + data=RegenerateKeyRequest(key_type=key_type), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="user-1", + ), + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + + assert exc_info.value.status_code == 403 + assert "allowed_routes" in str(exc_info.value.detail) + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0 + + @pytest.mark.asyncio async def test_execute_virtual_key_regeneration_rejects_over_limit_duration(monkeypatch): """Regenerate must reject durations exceeding upperbound_key_generate_params.duration.""" @@ -11837,6 +11881,278 @@ class TestAllowedRoutesCallerPermission: assert str(exc_info.value.code) == "403" assert "allowed_routes" in str(exc_info.value.message) + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("key_type", "expected_routes"), + [ + (LiteLLMKeyType.DEFAULT, []), + (LiteLLMKeyType.LLM_API, ["llm_api_routes"]), + (LiteLLMKeyType.READ_ONLY, ["info_routes"]), + ], + ) + async def test_non_admin_update_key_accepts_empty_routes_with_safe_key_type( + self, + key_type: LiteLLMKeyType, + expected_routes: list[str], + ): + from types import SimpleNamespace + + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_update_key_data, + prepare_key_update_data, + ) + + data = UpdateKeyRequest( + key="sk-test", + allowed_routes=[], + key_type=key_type, + ) + existing_key = SimpleNamespace( + created_by="internal-user-123", + key_alias=None, + max_budget=None, + metadata={}, + organization_id=None, + project_id=None, + team_id=None, + token="hashed-key", + user_id="internal-user-123", + ) + user_api_key_dict = UserAPIKeyAuth( + user_id="internal-user-123", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + await _validate_update_key_data( + data=data, + existing_key_row=existing_key, + user_api_key_dict=user_api_key_dict, + llm_router=None, + premium_user=True, + prisma_client=None, + user_api_key_cache=MagicMock(), + ) + + update_data = await prepare_key_update_data( + data=data, + existing_key_row=existing_key, + ) + assert update_data["key_type"] == key_type.value + assert update_data["allowed_routes"] == expected_routes + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("key_type", "existing_routes"), + [ + (LiteLLMKeyType.DEFAULT, ["llm_api_routes"]), + (LiteLLMKeyType.LLM_API, ["info_routes"]), + (LiteLLMKeyType.READ_ONLY, ["llm_api_routes"]), + ], + ) + async def test_non_admin_update_key_accepts_key_type_only_over_safe_preset( + self, + key_type: LiteLLMKeyType, + existing_routes: list[str], + ): + from types import SimpleNamespace + + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_update_key_data, + ) + + res = await _validate_update_key_data( + data=UpdateKeyRequest(key="sk-test", key_type=key_type), + existing_key_row=SimpleNamespace( + allowed_routes=existing_routes, + max_budget=None, + team_id=None, + user_id="internal-user-123", + ), + user_api_key_dict=UserAPIKeyAuth( + user_id="internal-user-123", + user_role=LitellmUserRoles.INTERNAL_USER, + ), + llm_router=None, + premium_user=True, + prisma_client=None, + user_api_key_cache=MagicMock(), + ) + assert res is None + + @pytest.mark.asyncio + async def test_non_admin_update_key_rejects_empty_routes_with_management_key_type( + self, + ): + from types import SimpleNamespace + + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_update_key_data, + ) + + data = UpdateKeyRequest( + key="sk-test", + allowed_routes=[], + key_type=LiteLLMKeyType.MANAGEMENT, + ) + + with pytest.raises(HTTPException) as exc_info: + await _validate_update_key_data( + data=data, + existing_key_row=SimpleNamespace( + max_budget=None, + user_id="internal-user-123", + ), + user_api_key_dict=UserAPIKeyAuth( + user_id="internal-user-123", + user_role=LitellmUserRoles.INTERNAL_USER, + ), + llm_router=None, + premium_user=True, + prisma_client=None, + user_api_key_cache=MagicMock(), + ) + assert exc_info.value.status_code == 403 + assert "allowed_routes" in str(exc_info.value.detail) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "key_type", + [ + LiteLLMKeyType.DEFAULT, + LiteLLMKeyType.LLM_API, + LiteLLMKeyType.READ_ONLY, + ], + ) + async def test_non_admin_update_key_rejects_preset_over_custom_routes( + self, key_type + ): + """A preset must not replace an admin-defined route allowlist.""" + from types import SimpleNamespace + + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_update_key_data, + ) + + data = UpdateKeyRequest( + key="sk-test", + allowed_routes=[], + key_type=key_type, + ) + + with pytest.raises(HTTPException) as exc_info: + await _validate_update_key_data( + data=data, + existing_key_row=SimpleNamespace( + allowed_routes=["/custom/admin-defined-route"], + max_budget=None, + user_id="internal-user-123", + ), + user_api_key_dict=UserAPIKeyAuth( + user_id="internal-user-123", + user_role=LitellmUserRoles.INTERNAL_USER, + ), + llm_router=None, + premium_user=True, + prisma_client=None, + user_api_key_cache=MagicMock(), + ) + + assert exc_info.value.status_code == 403 + assert "allowed_routes" in str(exc_info.value.detail) + + @pytest.mark.asyncio + async def test_regenerate_key_omitting_key_type_preserves_classification(self): + """Regeneration must not materialize the request model's default key type.""" + from types import SimpleNamespace + + from litellm.proxy._types import RegenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + prepare_key_update_data, + ) + + data = RegenerateKeyRequest(key="sk-test") + assert "key_type" not in data.model_fields_set + + update_data = await prepare_key_update_data( + data=data, + existing_key_row=SimpleNamespace( + metadata={}, + team_id=None, + ), + ) + + assert "key_type" not in update_data + assert "allowed_routes" not in update_data + + @pytest.mark.asyncio + async def test_update_key_default_type_clears_existing_preset_routes(self): + """An explicit default preset must remove the stored route restriction.""" + from types import SimpleNamespace + + from litellm.proxy.management_endpoints.key_management_endpoints import ( + prepare_key_update_data, + ) + + update_data = await prepare_key_update_data( + data=UpdateKeyRequest( + key="sk-test", + key_type=LiteLLMKeyType.DEFAULT, + ), + existing_key_row=SimpleNamespace( + metadata={}, + team_id=None, + ), + ) + + assert update_data["key_type"] == "default" + assert update_data["allowed_routes"] == [] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "key_type", + [ + LiteLLMKeyType.DEFAULT, + LiteLLMKeyType.LLM_API, + LiteLLMKeyType.READ_ONLY, + ], + ) + async def test_non_admin_update_key_rejects_preset_over_custom_routes_without_allowed_routes( + self, + key_type, + ): + """The key_type-only payload cannot clear an admin-defined allowlist.""" + from types import SimpleNamespace + + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_update_key_data, + ) + + data = UpdateKeyRequest( + key="sk-test", + key_type=key_type, + ) + + with pytest.raises(HTTPException) as exc_info: + await _validate_update_key_data( + data=data, + existing_key_row=SimpleNamespace( + allowed_routes=["/custom/admin-defined-route"], + max_budget=None, + user_id="internal-user-123", + ), + user_api_key_dict=UserAPIKeyAuth( + user_id="internal-user-123", + user_role=LitellmUserRoles.INTERNAL_USER, + ), + llm_router=None, + premium_user=True, + prisma_client=None, + user_api_key_cache=MagicMock(), + ) + + assert exc_info.value.status_code == 403 + assert "allowed_routes" in str(exc_info.value.detail) + @pytest.mark.asyncio async def test_non_admin_regenerate_key_explicit_empty_allowed_routes_rejected(self): """`regenerate_key_fn` rejects a non-admin when `allowed_routes` is diff --git a/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx b/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx index 1ab7b9be52c..f0f3d110b74 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx @@ -25,7 +25,7 @@ export const KeyTypeSelect = ({ onChange, }: { id: string; - value: string; + value?: string; onChange: (value: string) => void; }) => (