fix: allow safe key type updates for non-admins

Rebase onto current staging, preserve LF line endings, and extract preset normalization to satisfy the UI max-lines rule without suppressions.
This commit is contained in:
Patrick Ribbsaeter 2026-08-11 15:47:01 +02:00
parent ca0b951a43
commit ff6996a682
8 changed files with 502 additions and 19 deletions

View file

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

View file

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

View file

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

View file

@ -25,7 +25,7 @@ export const KeyTypeSelect = ({
onChange,
}: {
id: string;
value: string;
value?: string;
onChange: (value: string) => void;
}) => (
<Select

View file

@ -10,14 +10,22 @@ export const canonicalBudgetDuration = (duration: string | null | undefined): st
duration ? WORD_FORM_BUDGET_DURATIONS[duration] ?? duration : null;
// Determine the key_type display value from allowed_routes
export const keyTypeFromRoutes = (allowedRoutes: string[] | null | undefined): string => {
if (!allowedRoutes || allowedRoutes.length === 0) return "default";
if (allowedRoutes.includes("llm_api_routes")) return "llm_api";
if (allowedRoutes.includes("management_routes")) return "management";
if (allowedRoutes.includes("info_routes")) return "read_only";
return "default";
const KEY_TYPE_BY_PRESET_ROUTE: Record<string, string> = {
llm_api_routes: "llm_api",
management_routes: "management",
info_routes: "read_only",
};
export const exactKeyTypePresetFromRoutes = (allowedRoutes: string[] | null | undefined): string | undefined => {
if (!allowedRoutes || allowedRoutes.length === 0) return "default";
if (allowedRoutes.length !== 1) return undefined;
return KEY_TYPE_BY_PRESET_ROUTE[allowedRoutes[0]];
};
// Determine the key_type display value from allowed_routes
export const keyTypeFromRoutes = (allowedRoutes: string[] | null | undefined): string =>
exactKeyTypePresetFromRoutes(allowedRoutes) ?? "default";
export const parseAllowedRoutes = (value: unknown): string[] =>
typeof value === "string" && value.trim() !== ""
? value
@ -43,3 +51,4 @@ export const currentValuePlaceholder = (
if (!premiumUser) return premiumHint;
return Array.isArray(current) && current.length > 0 ? `Current: ${current.join(", ")}` : emptyHint;
};

View file

@ -716,7 +716,7 @@ describe("KeyEditView", () => {
});
});
it("should handle empty allowed routes string on submit", async () => {
it("should submit the default preset without derived allowed routes", async () => {
const onSubmitMock = vi.fn().mockResolvedValue(undefined);
const keyDataWithRoutes = {
...MOCK_KEY_DATA,
@ -747,7 +747,59 @@ describe("KeyEditView", () => {
await waitFor(() => {
expect(onSubmitMock).toHaveBeenCalled();
const callArgs = onSubmitMock.mock.calls[0][0];
expect(callArgs.allowed_routes).toEqual([]);
expect(callArgs.allowed_routes).toBeUndefined();
expect(callArgs.key_type).toBe("default");
});
});
it("should submit the AI APIs preset without derived allowed routes", async () => {
const onSubmitMock = vi.fn().mockResolvedValue(undefined);
renderWithProviders(
<KeyEditView
keyData={MOCK_KEY_DATA}
onCancel={() => {}}
onSubmit={onSubmitMock}
accessToken={"test-token"}
userID={"test-user"}
userRole={"admin"}
premiumUser={false}
/>,
);
const allowedRoutesInput = await screen.findByLabelText(/allowed routes/i);
await userEvent.type(allowedRoutesInput, "llm_api_routes");
await userEvent.click(screen.getByRole("button", { name: /save changes/i }));
await waitFor(() => {
const callArgs = onSubmitMock.mock.calls[0][0];
expect(callArgs.allowed_routes).toBeUndefined();
expect(callArgs.key_type).toBe("llm_api");
});
});
it("should preserve mixed route allowlists without applying a preset", async () => {
const onSubmitMock = vi.fn().mockResolvedValue(undefined);
renderWithProviders(
<KeyEditView
keyData={MOCK_KEY_DATA}
onCancel={() => {}}
onSubmit={onSubmitMock}
accessToken={"test-token"}
userID={"test-user"}
userRole={"admin"}
premiumUser={false}
/>,
);
const allowedRoutesInput = await screen.findByLabelText(/allowed routes/i);
await userEvent.clear(allowedRoutesInput);
await userEvent.type(allowedRoutesInput, "llm_api_routes, /custom");
await userEvent.click(screen.getByRole("button", { name: /save changes/i }));
await waitFor(() => {
const callArgs = onSubmitMock.mock.calls[0][0];
expect(callArgs.allowed_routes).toEqual(["llm_api_routes", "/custom"]);
expect(callArgs.key_type).toBeUndefined();
});
});

View file

@ -27,7 +27,9 @@ import RouterSettingsAccordion, { RouterSettingsAccordionRef } from "../common_c
import { routerSettingsEditorValue, routerSettingsUpdate } from "../common_components/routerSettingsPayload";
import { estimateTooltips, withNormalizedEstimates } from "./estimatedOutputTokens";
import {
canonicalBudgetDuration,
currentValuePlaceholder,
exactKeyTypePresetFromRoutes,
keyTypeFromRoutes,
modelSentinelOptions,
parseAllowedRoutes,
@ -236,6 +238,16 @@ export function KeyEditView({
[...submittedRoutesSet].every((r) => originalRoutesSet.has(r));
if (allowedRoutesUnchanged) {
delete values.allowed_routes;
} else {
const exactKeyTypePreset = exactKeyTypePresetFromRoutes(values.allowed_routes);
if (exactKeyTypePreset) {
values.key_type = exactKeyTypePreset;
// The backend derives allowed_routes from key_type. Omitting the
// derived routes also lets authorized non-admins use safe presets.
delete values.allowed_routes;
} else {
delete values.key_type;
}
}
if (neverExpire) {

View file

@ -7779,6 +7779,7 @@ export interface paths {
* 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.
@ -36311,6 +36312,7 @@ export interface components {
key?: string | null;
/** Key Alias */
key_alias?: string | null;
key_type?: components["schemas"]["LiteLLMKeyType"] | null;
/** Max Budget */
max_budget?: number | null;
/** Max Parallel Requests */