mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-25 05:54:53 +00:00
feat(key_management)!: apply default_key_generate_params.allowed_routes on key generation (#41236)
* feat(key_management): apply default_key_generate_params.allowed_routes on key generation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(key_management): avoid mutating default routes payload Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(proxy): apply default allowed_routes to key_type presets it covers * fix(proxy): keep management and read_only presets, skip service account keys, and show the default routes to custom_key_policy * fix(proxy): keep Full Access keys out of the default allowed_routes, let owners re-pick a key type, and show custom_key_policy the member key's routes * fix(proxy): keep owners from changing routes a key got from default_key_generate_params * docs(proxy): describe key_type when default_key_generate_params.allowed_routes is set, and test custom_key_policy input with the default off * test(proxy): type the default_key_generate_params allowed_routes tests * fix(proxy): let key owners send unchanged allowed_routes back on /key/update * fix(proxy): ignore a route default with a non-string entry and match resent routes exactly A default_key_generate_params.allowed_routes list with a non-string entry made /key/generate fail at the database, so it is now ignored like a non-list. An owner's resent allowed_routes on /key/update skip the route check only when they equal the stored list in order, so reordered or duplicated lists are checked as a change. Adds integration cells for the default and the owner update paths. * test(proxy): cover a non-string allowed_routes default in the key integration tests * test(proxy): read key rows from the owned proxy's database in the default routes integration tests * fix(proxy): leave allowed_routes out of a key update that resends them unchanged --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
c6c303a618
commit
1660c8ba1f
4 changed files with 709 additions and 12 deletions
|
|
@ -471,7 +471,13 @@ def _effective_key_for_generate(data: GenerateKeyRequest, now: datetime) -> Lite
|
|||
)
|
||||
metadata: Final = data.metadata or MappingProxyType({})
|
||||
folded_metadata: Final = {**metadata, **metadata_fields}
|
||||
columns: Final = handle_key_type(data, {**column_fields})
|
||||
preset_columns: Final = handle_key_type(data, {**column_fields})
|
||||
default_allowed_routes: Final = _default_allowed_routes_for(data, preset_columns.get("allowed_routes"))
|
||||
columns: Final = (
|
||||
preset_columns
|
||||
if default_allowed_routes is None
|
||||
else {**preset_columns, "allowed_routes": list(default_allowed_routes)}
|
||||
)
|
||||
expires: Final = (
|
||||
now + timedelta(seconds=duration_in_seconds(duration=data.duration)) if data.duration is not None else None
|
||||
)
|
||||
|
|
@ -936,7 +942,19 @@ def common_key_access_checks(
|
|||
router: Final = APIRouter()
|
||||
|
||||
|
||||
def handle_key_type(data: GenerateKeyRequest, data_json: dict) -> dict:
|
||||
_KEY_TYPE_PRESET_ROUTES: Final = MappingProxyType(
|
||||
{
|
||||
LiteLLMKeyType.LLM_API: ("llm_api_routes",),
|
||||
LiteLLMKeyType.MANAGEMENT: ("management_routes",),
|
||||
LiteLLMKeyType.READ_ONLY: ("info_routes",),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
_KeyJsonT: Final = TypeVar("_KeyJsonT", bound=dict[str, object])
|
||||
|
||||
|
||||
def handle_key_type(data: GenerateKeyRequest, data_json: _KeyJsonT) -> _KeyJsonT:
|
||||
"""
|
||||
Handle the key type.
|
||||
"""
|
||||
|
|
@ -945,15 +963,44 @@ 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"]
|
||||
preset_routes: Final = _KEY_TYPE_PRESET_ROUTES.get(key_type)
|
||||
if preset_routes is not None:
|
||||
data_json["allowed_routes"] = list(preset_routes)
|
||||
return data_json
|
||||
|
||||
|
||||
def _creates_team_service_account(request: GenerateKeyRequest) -> bool:
|
||||
return (
|
||||
request.user_id is None
|
||||
and request.team_id is not None
|
||||
and request.metadata is not None
|
||||
and request.metadata.get("service_account_id") is not None
|
||||
)
|
||||
|
||||
|
||||
def _configured_default_allowed_routes() -> tuple[str, ...] | None:
|
||||
default_routes: Final[object] = (litellm.default_key_generate_params or {}).get("allowed_routes")
|
||||
if not isinstance(default_routes, list) or not default_routes:
|
||||
return None
|
||||
routes: Final[tuple[object, ...]] = tuple(default_routes)
|
||||
string_routes: Final = tuple(route for route in routes if isinstance(route, str))
|
||||
return string_routes if len(string_routes) == len(routes) else None
|
||||
|
||||
|
||||
def _default_allowed_routes_for(request: GenerateKeyRequest, key_routes: object) -> tuple[str, ...] | None:
|
||||
configured_routes: Final = _configured_default_allowed_routes()
|
||||
if configured_routes is None or _creates_team_service_account(request):
|
||||
return None
|
||||
if "key_type" not in request.model_fields_set or request.key_type is None:
|
||||
return None if key_routes else configured_routes
|
||||
match request.key_type:
|
||||
case LiteLLMKeyType.LLM_API:
|
||||
llm_api_routes: Final = _KEY_TYPE_PRESET_ROUTES[LiteLLMKeyType.LLM_API]
|
||||
return configured_routes if all(route in configured_routes for route in llm_api_routes) else None
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
_NON_ADMIN_SAFE_ALLOWED_ROUTES_PRESETS: Final = frozenset({"llm_api_routes", "info_routes"})
|
||||
|
||||
|
||||
|
|
@ -1068,11 +1115,24 @@ def _is_safe_preset_route_transition(
|
|||
return existing != _READ_ONLY_ALLOWED_ROUTES_PRESET or incoming == existing
|
||||
|
||||
|
||||
def _resends_stored_routes(
|
||||
data: UpdateKeyRequest | RegenerateKeyRequest,
|
||||
existing_key_row: LiteLLM_VerificationToken,
|
||||
) -> bool:
|
||||
return (
|
||||
"allowed_routes" in data.model_fields_set
|
||||
and data.allowed_routes is not None
|
||||
and list(data.allowed_routes) == list(existing_key_row.allowed_routes or ())
|
||||
)
|
||||
|
||||
|
||||
def _enforce_allowed_routes_update_permission(
|
||||
data: UpdateKeyRequest,
|
||||
existing_key_row: LiteLLM_VerificationToken,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> None:
|
||||
if _resends_stored_routes(data, existing_key_row):
|
||||
return
|
||||
if _is_safe_preset_route_transition(
|
||||
incoming_allowed_routes=data.allowed_routes,
|
||||
existing_allowed_routes=existing_key_row.allowed_routes,
|
||||
|
|
@ -1499,7 +1559,9 @@ async def _common_key_generation_helper(
|
|||
build_policy_request=lambda: CustomKeyPolicyRequest(
|
||||
operation="generate",
|
||||
existing_key=None,
|
||||
effective_key=_effective_key_for_generate(data=data, now=datetime.now(timezone.utc)),
|
||||
effective_key=_effective_key_for_generate(
|
||||
data=_member_key_request(data, acting_as_team_admin), now=datetime.now(timezone.utc)
|
||||
),
|
||||
request=data,
|
||||
),
|
||||
)
|
||||
|
|
@ -1549,7 +1611,8 @@ async def _common_key_generation_helper(
|
|||
|
||||
data_json = data.model_dump(exclude_unset=True, exclude_none=True)
|
||||
|
||||
data_json = handle_key_type(_member_key_request(data, acting_as_team_admin), data_json)
|
||||
member_request: Final = _member_key_request(data, acting_as_team_admin)
|
||||
data_json = handle_key_type(member_request, data_json)
|
||||
|
||||
# Re-check allowed_routes after handle_key_type, since key_type can derive
|
||||
# an elevated bucket (e.g. ["management_routes"]) that wasn't present in
|
||||
|
|
@ -1561,6 +1624,9 @@ async def _common_key_generation_helper(
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
allow_safe_presets=True,
|
||||
)
|
||||
default_allowed_routes: Final = _default_allowed_routes_for(member_request, data_json.get("allowed_routes"))
|
||||
if default_allowed_routes is not None:
|
||||
data_json["allowed_routes"] = list(default_allowed_routes)
|
||||
|
||||
# if we get max_budget passed to /key/generate, then use it as key_max_budget. Since generate_key_helper_fn is used to make new users
|
||||
if "max_budget" in data_json:
|
||||
|
|
@ -2268,7 +2334,7 @@ async def generate_key_fn(
|
|||
- allowed_passthrough_routes: Optional[list] - List of allowed pass through endpoints for the key. Store the actual endpoint or store a wildcard pattern for a set of endpoints. Example - ["/my-custom-endpoint"]. Use this instead of allowed_routes, if you just want to specify which pass through endpoints the key can access, without specifying the routes. If allowed_routes is specified, allowed_pass_through_endpoints is ignored.
|
||||
- denied_passthrough_routes: Optional[list] - List of pass through routes the key may not call, even if allowed by `allowed_passthrough_routes` or `allowed_routes`. Matches exact paths, path prefixes, and trailing `*` wildcards. Applies together with the team's `denied_passthrough_routes`. Example - ["/my-custom-endpoint/admin"].
|
||||
- object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission.
|
||||
- key_type: Optional[str] - Type of key that determines default allowed routes. Options: "llm_api" (can call LLM API routes), "management" (can call management routes), "read_only" (can only call info/read routes), "default" (uses default allowed routes). Defaults to "default".
|
||||
- key_type: Optional[str] - Type of key that determines default allowed routes. Options: "llm_api" (can call LLM API routes), "management" (can call management routes), "read_only" (can only call info/read routes), "default" (Full Access in the Admin UI, skips `default_key_generate_params.allowed_routes`). When omitted, the key gets `default_key_generate_params.allowed_routes` if that is set and the request has no `allowed_routes`, else it works like "default".
|
||||
- prompts: Optional[List[str]] - List of allowed prompts for the key. If specified, the key will only be able to use these specific prompts.
|
||||
- auto_rotate: Optional[bool] - Whether this key should be automatically rotated (regenerated)
|
||||
- rotation_interval: Optional[str] - How often to auto-rotate this key (e.g., '30s', '30m', '30h', '30d'). Required if auto_rotate=True.
|
||||
|
|
@ -2758,6 +2824,8 @@ async def prepare_key_update_data(
|
|||
data_json.pop("key", None)
|
||||
data_json.pop("new_key", None)
|
||||
data_json.pop("grace_period", None) # Request-only param, not a DB column
|
||||
if _resends_stored_routes(data, existing_key_row):
|
||||
data_json.pop("allowed_routes", None)
|
||||
if (
|
||||
data.metadata is not None
|
||||
and data.metadata.get("service_account_id") is not None
|
||||
|
|
|
|||
327
tests/integration/management/test_key_default_allowed_routes.py
Normal file
327
tests/integration/management/test_key_default_allowed_routes.py
Normal file
|
|
@ -0,0 +1,327 @@
|
|||
import uuid
|
||||
from collections.abc import Iterator, Mapping
|
||||
from datetime import date
|
||||
from hashlib import sha256
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
from tests.integration._support.client import (
|
||||
JSON_OBJECT,
|
||||
Gateway,
|
||||
Scenario,
|
||||
eventually,
|
||||
gateway_from_environment,
|
||||
list_value,
|
||||
object_value,
|
||||
string_value,
|
||||
team_admin_permissions,
|
||||
)
|
||||
from tests.integration._support.process import graceful_stop_seconds, owned_proxy
|
||||
|
||||
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
|
||||
|
||||
DEFAULT_ROUTES: Final = ("llm_api_routes", "/key/spend/report")
|
||||
_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
|
||||
_PROBE_PREFIX: Final = "route-probe-"
|
||||
_RACE_PREFIX: Final = "route-race-"
|
||||
_ROUTES_AFTER_ADMIN_REMOVAL: Final = ["llm_api_routes"]
|
||||
_POLICY_SOURCE: Final = f'''import json
|
||||
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
|
||||
async def route_policy(request):
|
||||
alias = request.request.key_alias or ""
|
||||
if request.operation == "update" and alias.startswith("{_RACE_PREFIX}"):
|
||||
await proxy_server.prisma_client.db.litellm_verificationtoken.update(
|
||||
where={{"token": request.existing_key.token}},
|
||||
data={{"allowed_routes": {_ROUTES_AFTER_ADMIN_REMOVAL!r}}},
|
||||
)
|
||||
return {{"decision": True}}
|
||||
if not alias.startswith("{_PROBE_PREFIX}"):
|
||||
return {{"decision": True}}
|
||||
return {{"decision": False, "message": "routes=" + json.dumps(request.effective_key.allowed_routes or [])}}
|
||||
'''
|
||||
|
||||
|
||||
def _config_with_default_routes(
|
||||
directory: Path, default_routes: list[JsonValue], general_settings: Mapping[str, JsonValue]
|
||||
) -> Path:
|
||||
raw: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
|
||||
configured: Final[dict[str, JsonValue]] = {
|
||||
**raw,
|
||||
"litellm_settings": {
|
||||
**object_value(raw.get("litellm_settings") or {}),
|
||||
"default_key_generate_params": {"allowed_routes": default_routes},
|
||||
},
|
||||
"general_settings": {**object_value(raw.get("general_settings") or {}), **general_settings},
|
||||
}
|
||||
config: Final = directory / "default-allowed-routes.yaml"
|
||||
config.write_text(yaml.safe_dump(configured))
|
||||
return config
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def defaults(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]:
|
||||
directory: Final = tmp_path_factory.mktemp("default-allowed-routes")
|
||||
(directory / "route_policy.py").write_text(_POLICY_SOURCE)
|
||||
config: Final = _config_with_default_routes(
|
||||
directory, list(DEFAULT_ROUTES), {"custom_key_policy": "route_policy.route_policy"}
|
||||
)
|
||||
with gateway_from_environment() as gateway, owned_proxy(gateway, directory, {}, config=config) as candidate:
|
||||
yield candidate
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def non_string_default(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]:
|
||||
directory: Final = tmp_path_factory.mktemp("non-string-default-allowed-routes")
|
||||
config: Final = _config_with_default_routes(directory, ["llm_api_routes", 123], {})
|
||||
with gateway_from_environment() as gateway, owned_proxy(gateway, directory, {}, config=config) as candidate:
|
||||
yield candidate
|
||||
|
||||
|
||||
def _hash(key: str) -> str:
|
||||
return sha256(key.encode()).hexdigest()
|
||||
|
||||
|
||||
def _stored_routes(gateway: Gateway, key: str) -> JsonValue:
|
||||
rows: Final = gateway.rows('SELECT allowed_routes FROM "LiteLLM_VerificationToken" WHERE token = %s', (_hash(key),))
|
||||
assert len(rows) == 1, rows
|
||||
return rows[0]["allowed_routes"]
|
||||
|
||||
|
||||
def _created_key(scenario: Scenario, created: dict[str, JsonValue]) -> str:
|
||||
token: Final = string_value(created["key"])
|
||||
scenario.cleanups.callback(scenario.delete_key, token)
|
||||
return token
|
||||
|
||||
|
||||
def _full_access_key(scenario: Scenario, **fields: JsonValue) -> str:
|
||||
return scenario.key(key_type="default", **fields)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("body", "expected"),
|
||||
(
|
||||
({}, list(DEFAULT_ROUTES)),
|
||||
({"key_type": None}, list(DEFAULT_ROUTES)),
|
||||
({"allowed_routes": []}, list(DEFAULT_ROUTES)),
|
||||
({"key_type": "llm_api"}, list(DEFAULT_ROUTES)),
|
||||
({"allowed_routes": ["info_routes"]}, ["info_routes"]),
|
||||
({"key_type": "management"}, ["management_routes"]),
|
||||
({"key_type": "read_only"}, ["info_routes"]),
|
||||
({"key_type": "default"}, []),
|
||||
),
|
||||
ids=("unset", "key_type_null", "empty_list", "llm_api", "explicit_list", "management", "read_only", "full_access"),
|
||||
)
|
||||
def test_admin_created_key_routes(defaults: Gateway, body: dict[str, JsonValue], expected: list[str]) -> None:
|
||||
with defaults.scenario() as scenario:
|
||||
user: Final = scenario.user(user_role="internal_user")
|
||||
created: Final = defaults.post("/key/generate", {"user_id": user, **body})
|
||||
token: Final = _created_key(scenario, created)
|
||||
|
||||
assert (created["allowed_routes"], _stored_routes(defaults, token)) == (expected, expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("body", "expected"),
|
||||
(({}, []), ({"key_type": "llm_api"}, ["llm_api_routes"])),
|
||||
ids=("unset", "llm_api"),
|
||||
)
|
||||
def test_default_with_a_non_string_route_is_ignored(
|
||||
non_string_default: Gateway, body: dict[str, JsonValue], expected: list[str]
|
||||
) -> None:
|
||||
with non_string_default.scenario() as scenario:
|
||||
user: Final = scenario.user(user_role="internal_user")
|
||||
created: Final = non_string_default.post("/key/generate", {"user_id": user, **body})
|
||||
token: Final = _created_key(scenario, created)
|
||||
|
||||
assert (created["allowed_routes"], _stored_routes(non_string_default, token)) == (expected, expected)
|
||||
|
||||
|
||||
def test_default_routes_open_own_spend_report_and_nothing_wider(defaults: Gateway) -> None:
|
||||
with defaults.scenario() as scenario:
|
||||
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
|
||||
user: Final = scenario.user(user_role="internal_user")
|
||||
token: Final = _created_key(
|
||||
scenario, defaults.post("/key/generate", {"user_id": user, "key_type": "llm_api", "models": [model]})
|
||||
)
|
||||
other: Final = _full_access_key(scenario, user_id=scenario.user(user_role="internal_user"))
|
||||
defaults.chat(model, key=token, text=f"default routes {uuid.uuid4().hex}")
|
||||
logged: Final = eventually(
|
||||
lambda: defaults.rows(
|
||||
"""SELECT to_char("startTime", 'YYYY-MM-DD') AS day FROM "LiteLLM_SpendLogs" WHERE api_key = %s""",
|
||||
(_hash(token),),
|
||||
),
|
||||
lambda rows: len(rows) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
day: Final = date.fromisoformat(string_value(logged[0]["day"])).isoformat()
|
||||
window: Final = {"start_date": day, "end_date": day}
|
||||
|
||||
own: Final = defaults.request("GET", "/key/spend/report", key=token, params=window)
|
||||
assert own.status_code == 200, own.text
|
||||
assert [object_value(row)["api_key"] for row in list_value(_JSON.validate_json(own.content))] == [_hash(token)]
|
||||
|
||||
foreign: Final = defaults.request(
|
||||
"GET", "/key/spend/report", key=token, params={**window, "api_key": _hash(other)}
|
||||
)
|
||||
assert foreign.status_code == 403, foreign.text
|
||||
global_report: Final = defaults.request("GET", "/global/spend/report", key=token, params=window)
|
||||
assert global_report.status_code == 403, global_report.text
|
||||
minted: Final = defaults.request("POST", "/key/generate", {}, key=token)
|
||||
assert minted.status_code == 403, minted.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize("body", ({}, {"key_type": "llm_api"}), ids=("unset", "llm_api"))
|
||||
def test_internal_user_own_key_routes(defaults: Gateway, body: dict[str, JsonValue]) -> None:
|
||||
with defaults.scenario() as scenario:
|
||||
user: Final = scenario.user(user_role="internal_user")
|
||||
own_key: Final = _full_access_key(scenario, user_id=user)
|
||||
created: Final = defaults.request("POST", "/key/generate", body, key=own_key)
|
||||
assert created.status_code == 200, created.text
|
||||
token: Final = _created_key(scenario, JSON_OBJECT.validate_json(created.content))
|
||||
|
||||
assert defaults.rows(
|
||||
'SELECT user_id, allowed_routes FROM "LiteLLM_VerificationToken" WHERE token = %s', (_hash(token),)
|
||||
) == [{"user_id": user, "allowed_routes": list(DEFAULT_ROUTES)}]
|
||||
|
||||
|
||||
def test_service_account_skips_default_and_its_team_key_gets_it(defaults: Gateway) -> None:
|
||||
with defaults.scenario() as scenario:
|
||||
team: Final = scenario.team()
|
||||
account: Final = _created_key(scenario, defaults.post("/key/service-account/generate", {"team_id": team}))
|
||||
assert _stored_routes(defaults, account) == []
|
||||
|
||||
defaults.post("/team/permissions_update", {"team_id": team, "team_member_permissions": ["/key/generate"]})
|
||||
team_key: Final = _created_key(scenario, defaults.post("/key/generate", {"team_id": team}, key=account))
|
||||
|
||||
assert defaults.rows(
|
||||
'SELECT user_id, team_id, allowed_routes FROM "LiteLLM_VerificationToken" WHERE token = %s',
|
||||
(_hash(team_key),),
|
||||
) == [{"user_id": None, "team_id": team, "allowed_routes": list(DEFAULT_ROUTES)}]
|
||||
|
||||
|
||||
def _policy_seen_routes(response_text: str) -> JsonValue:
|
||||
message: Final = string_value(object_value(JSON_OBJECT.validate_json(response_text)["error"])["message"])
|
||||
assert message.startswith("routes="), message
|
||||
return _JSON.validate_json(message.removeprefix("routes="))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("team_admin", (False, True), ids=("unset", "team_admin_member_key"))
|
||||
def test_custom_key_policy_sees_stored_routes(defaults: Gateway, team_admin: bool) -> None:
|
||||
with defaults.scenario() as scenario:
|
||||
team: Final = scenario.team()
|
||||
admin: Final = scenario.member(team, role="admin")
|
||||
member: Final = scenario.member(team)
|
||||
caller: Final = _full_access_key(scenario, user_id=admin, team_id=team) if team_admin else defaults.key
|
||||
body: Final[dict[str, JsonValue]] = {"team_id": team, "user_id": member}
|
||||
scenario.cleanups.enter_context(team_admin_permissions(defaults, ("create_member_keys",)))
|
||||
|
||||
probe: Final = defaults.request(
|
||||
"POST", "/key/generate", {**body, "key_alias": f"{_PROBE_PREFIX}{uuid.uuid4().hex}"}, key=caller
|
||||
)
|
||||
assert probe.status_code == 403, probe.text
|
||||
created: Final = defaults.request(
|
||||
"POST", "/key/generate", {**body, "key_alias": f"route-twin-{uuid.uuid4().hex}"}, key=caller
|
||||
)
|
||||
assert created.status_code == 200, created.text
|
||||
token: Final = _created_key(scenario, JSON_OBJECT.validate_json(created.content))
|
||||
|
||||
assert _policy_seen_routes(probe.text) == _stored_routes(defaults, token) == list(DEFAULT_ROUTES)
|
||||
|
||||
|
||||
_ADMIN_SET_ROUTES: Final = ("llm_api_routes", "/key/spend/report")
|
||||
|
||||
|
||||
def _owner_key_with_admin_set_routes(gateway: Gateway, scenario: Scenario) -> tuple[str, str]:
|
||||
owner_key: Final = scenario.key(user_id=scenario.user(user_role="internal_user"))
|
||||
target: Final = _created_key(scenario, gateway.post("/key/generate", {}, key=owner_key))
|
||||
gateway.post("/key/update", {"key": target, "allowed_routes": list(_ADMIN_SET_ROUTES)})
|
||||
return owner_key, target
|
||||
|
||||
|
||||
def _stored_alias_and_routes(gateway: Gateway, key: str) -> tuple[JsonValue, list[str]]:
|
||||
rows: Final = gateway.rows(
|
||||
'SELECT key_alias, allowed_routes FROM "LiteLLM_VerificationToken" WHERE token = %s', (_hash(key),)
|
||||
)
|
||||
assert len(rows) == 1, rows
|
||||
return rows[0]["key_alias"], [str(route) for route in list_value(rows[0]["allowed_routes"])]
|
||||
|
||||
|
||||
def test_owner_resends_unchanged_routes(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
owner_key, target = _owner_key_with_admin_set_routes(gateway, scenario)
|
||||
alias: Final = f"resave-{uuid.uuid4().hex}"
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/key/update",
|
||||
{"key": target, "key_alias": alias, "allowed_routes": list(_ADMIN_SET_ROUTES)},
|
||||
key=owner_key,
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert _stored_alias_and_routes(gateway, target) == (alias, list(_ADMIN_SET_ROUTES))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"routes",
|
||||
(
|
||||
[],
|
||||
["llm_api_routes"],
|
||||
[*_ADMIN_SET_ROUTES, "/global/spend/report"],
|
||||
list(reversed(_ADMIN_SET_ROUTES)),
|
||||
[*_ADMIN_SET_ROUTES, _ADMIN_SET_ROUTES[-1]],
|
||||
),
|
||||
ids=("clear", "ai_apis", "widen", "reordered", "duplicated"),
|
||||
)
|
||||
def test_owner_cannot_change_admin_set_routes(gateway: Gateway, routes: list[JsonValue]) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
owner_key, target = _owner_key_with_admin_set_routes(gateway, scenario)
|
||||
before: Final = _stored_alias_and_routes(gateway, target)
|
||||
response: Final = gateway.request(
|
||||
"POST", "/key/update", {"key": target, "allowed_routes": routes}, key=owner_key
|
||||
)
|
||||
|
||||
assert response.status_code == 403, response.text
|
||||
assert _stored_alias_and_routes(gateway, target) == before
|
||||
|
||||
|
||||
def test_non_owner_resending_unchanged_routes_is_refused(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
_owner_key, target = _owner_key_with_admin_set_routes(gateway, scenario)
|
||||
outsider_key: Final = scenario.key(user_id=scenario.user(user_role="internal_user"))
|
||||
before: Final = _stored_alias_and_routes(gateway, target)
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/key/update",
|
||||
{"key": target, "key_alias": f"hijack-{uuid.uuid4().hex}", "allowed_routes": list(_ADMIN_SET_ROUTES)},
|
||||
key=outsider_key,
|
||||
)
|
||||
|
||||
assert response.status_code == 403, response.text
|
||||
assert _stored_alias_and_routes(gateway, target) == before
|
||||
|
||||
|
||||
def test_owner_resend_keeps_an_admin_route_removal_made_mid_update(defaults: Gateway) -> None:
|
||||
with defaults.scenario() as scenario:
|
||||
owner_key: Final = _full_access_key(scenario, user_id=scenario.user(user_role="internal_user"))
|
||||
created: Final = defaults.request("POST", "/key/generate", {}, key=owner_key)
|
||||
assert created.status_code == 200, created.text
|
||||
target: Final = _created_key(scenario, JSON_OBJECT.validate_json(created.content))
|
||||
assert _stored_routes(defaults, target) == list(DEFAULT_ROUTES)
|
||||
alias: Final = f"{_RACE_PREFIX}{uuid.uuid4().hex}"
|
||||
|
||||
response: Final = defaults.request(
|
||||
"POST",
|
||||
"/key/update",
|
||||
{"key": target, "key_alias": alias, "allowed_routes": list(DEFAULT_ROUTES)},
|
||||
key=owner_key,
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert _stored_alias_and_routes(defaults, target) == (alias, _ROUTES_AFTER_ADMIN_REMOVAL)
|
||||
|
|
@ -19639,6 +19639,285 @@ async def _generate_key_and_get_persisted_row(data: GenerateKeyRequest, mock_ins
|
|||
return key_call.kwargs["data"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_key_generate_params_allowed_routes_fills_unset(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"default_key_generate_params",
|
||||
{"allowed_routes": ["llm_api_routes", "/key/spend/report"]},
|
||||
)
|
||||
mock_insert_data: Final = _wire_key_generation_prisma(monkeypatch)
|
||||
|
||||
response: Final = await generate_key_fn(
|
||||
data=GenerateKeyRequest(),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
api_key="sk-internal",
|
||||
user_id="internal-user",
|
||||
),
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
assert response is not None
|
||||
key_call: Final = next(c for c in mock_insert_data.call_args_list if c.kwargs["table_name"] == "key")
|
||||
assert key_call.kwargs["data"]["allowed_routes"] == ["llm_api_routes", "/key/spend/report"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_key_generate_params_allowed_routes_keeps_explicit_value(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"default_key_generate_params",
|
||||
{"allowed_routes": ["llm_api_routes", "/key/spend/report"]},
|
||||
)
|
||||
mock_insert_data: Final = _wire_key_generation_prisma(monkeypatch)
|
||||
|
||||
key_row: Final = await _generate_key_and_get_persisted_row(
|
||||
GenerateKeyRequest(allowed_routes=["info_routes"]),
|
||||
mock_insert_data,
|
||||
)
|
||||
|
||||
assert key_row["allowed_routes"] == ["info_routes"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_key_generate_params_allowed_routes_extends_key_type_preset_it_covers(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from litellm.proxy._types import LiteLLMKeyType
|
||||
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"default_key_generate_params",
|
||||
{"allowed_routes": ["llm_api_routes", "/key/spend/report"]},
|
||||
)
|
||||
mock_insert_data: Final = _wire_key_generation_prisma(monkeypatch)
|
||||
|
||||
await generate_key_fn(
|
||||
data=GenerateKeyRequest(key_type=LiteLLMKeyType.LLM_API),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
api_key="sk-internal",
|
||||
user_id="internal-user",
|
||||
),
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
key_call: Final = next(c for c in mock_insert_data.call_args_list if c.kwargs["table_name"] == "key")
|
||||
assert key_call.kwargs["data"]["allowed_routes"] == ["llm_api_routes", "/key/spend/report"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("key_type", "preset_routes"),
|
||||
[(LiteLLMKeyType.MANAGEMENT, ["management_routes"]), (LiteLLMKeyType.READ_ONLY, ["info_routes"])],
|
||||
)
|
||||
async def test_default_key_generate_params_allowed_routes_keeps_management_and_read_only_presets(
|
||||
monkeypatch: pytest.MonkeyPatch, key_type: LiteLLMKeyType, preset_routes: list[str]
|
||||
) -> None:
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"default_key_generate_params",
|
||||
{"allowed_routes": ["llm_api_routes", "management_routes", "info_routes", "/key/spend/report"]},
|
||||
)
|
||||
mock_insert_data: Final = _wire_key_generation_prisma(monkeypatch)
|
||||
|
||||
key_row: Final = await _generate_key_and_get_persisted_row(
|
||||
GenerateKeyRequest(key_type=key_type),
|
||||
mock_insert_data,
|
||||
)
|
||||
|
||||
assert key_row["allowed_routes"] == preset_routes
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"data",
|
||||
[GenerateKeyRequest(), GenerateKeyRequest(key_type=LiteLLMKeyType.LLM_API)],
|
||||
ids=["unset", "llm_api"],
|
||||
)
|
||||
def test_effective_key_for_generate_shows_default_allowed_routes_to_the_policy(
|
||||
monkeypatch: pytest.MonkeyPatch, data: GenerateKeyRequest
|
||||
) -> None:
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"default_key_generate_params",
|
||||
{"allowed_routes": ["llm_api_routes", "/key/spend/report"]},
|
||||
)
|
||||
|
||||
effective_key: Final = _effective_key_for_generate(data=data, now=datetime(2026, 1, 1, tzinfo=timezone.utc))
|
||||
|
||||
assert effective_key.allowed_routes == ["llm_api_routes", "/key/spend/report"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_key_generate_params_allowed_routes_skips_team_service_account_keys(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"default_key_generate_params",
|
||||
{"allowed_routes": ["llm_api_routes", "/key/spend/report"]},
|
||||
)
|
||||
mock_insert_data: Final = _wire_key_generation_prisma(monkeypatch)
|
||||
|
||||
key_row: Final = await _generate_key_and_get_persisted_row(
|
||||
GenerateKeyRequest(team_id="team-a", metadata={"service_account_id": "sa-1"}),
|
||||
mock_insert_data,
|
||||
)
|
||||
|
||||
assert not key_row.get("allowed_routes")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_key_generate_params_allowed_routes_skips_full_access_keys(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"default_key_generate_params",
|
||||
{"allowed_routes": ["llm_api_routes", "/key/spend/report"]},
|
||||
)
|
||||
mock_insert_data: Final = _wire_key_generation_prisma(monkeypatch)
|
||||
|
||||
key_row: Final = await _generate_key_and_get_persisted_row(
|
||||
GenerateKeyRequest(key_type=LiteLLMKeyType.DEFAULT),
|
||||
mock_insert_data,
|
||||
)
|
||||
|
||||
assert not key_row.get("allowed_routes")
|
||||
|
||||
|
||||
def _owner_route_update_status(existing: list[str], incoming: list[str]) -> int | None:
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_enforce_allowed_routes_update_permission,
|
||||
)
|
||||
|
||||
try:
|
||||
_enforce_allowed_routes_update_permission(
|
||||
data=UpdateKeyRequest(key="sk-owned", allowed_routes=incoming),
|
||||
existing_key_row=LiteLLM_VerificationToken(
|
||||
token="hashed-owned", user_id="internal-user", allowed_routes=existing
|
||||
),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="internal-user", user_role=LitellmUserRoles.INTERNAL_USER),
|
||||
)
|
||||
except HTTPException as exc:
|
||||
return exc.status_code
|
||||
return None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("existing", "incoming", "expected_status"),
|
||||
[
|
||||
(["llm_api_routes"], [], None),
|
||||
(["llm_api_routes", "/key/spend/report"], [], 403),
|
||||
(["llm_api_routes", "/key/spend/report"], ["llm_api_routes"], 403),
|
||||
(["llm_api_routes", "/key/spend/report"], ["llm_api_routes", "/key/spend/report"], None),
|
||||
(["llm_api_routes", "/key/spend/report"], ["/key/spend/report", "llm_api_routes"], 403),
|
||||
(
|
||||
["llm_api_routes", "/key/spend/report"],
|
||||
["llm_api_routes", "/key/spend/report", "/key/spend/report"],
|
||||
403,
|
||||
),
|
||||
(["/global/spend/report"], ["/global/spend/report"], None),
|
||||
],
|
||||
ids=[
|
||||
"preset_to_full_access",
|
||||
"default_list_to_full_access",
|
||||
"default_list_to_ai_apis",
|
||||
"default_list_sent_back_unchanged",
|
||||
"default_list_sent_back_reordered",
|
||||
"default_list_sent_back_with_a_duplicate",
|
||||
"admin_set_list_sent_back_unchanged",
|
||||
],
|
||||
)
|
||||
def test_owner_cannot_change_the_routes_a_key_got_from_default_key_generate_params(
|
||||
monkeypatch: pytest.MonkeyPatch, existing: list[str], incoming: list[str], expected_status: int | None
|
||||
) -> None:
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"default_key_generate_params",
|
||||
{"allowed_routes": ["llm_api_routes", "/key/spend/report"]},
|
||||
)
|
||||
|
||||
assert _owner_route_update_status(existing, incoming) == expected_status
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("incoming", "written"),
|
||||
[
|
||||
(["llm_api_routes", "/key/spend/report"], None),
|
||||
(["/key/spend/report", "llm_api_routes"], ["/key/spend/report", "llm_api_routes"]),
|
||||
(["llm_api_routes"], ["llm_api_routes"]),
|
||||
],
|
||||
ids=["unchanged_resend", "reordered", "changed"],
|
||||
)
|
||||
async def test_key_update_writes_allowed_routes_only_when_they_change(
|
||||
incoming: list[str], written: list[str] | None
|
||||
) -> None:
|
||||
existing_key: Final = LiteLLM_VerificationToken(
|
||||
token="hashed", allowed_routes=["llm_api_routes", "/key/spend/report"]
|
||||
)
|
||||
|
||||
update_values: Final = await prepare_key_update_data(
|
||||
data=UpdateKeyRequest(key="sk-1", key_alias="renamed", allowed_routes=incoming), existing_key_row=existing_key
|
||||
)
|
||||
|
||||
assert (update_values["key_alias"], update_values.get("allowed_routes")) == ("renamed", written)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("request_data", "expected_routes"),
|
||||
[(GenerateKeyRequest(), []), (GenerateKeyRequest(key_type=LiteLLMKeyType.LLM_API), ["llm_api_routes"])],
|
||||
ids=["no_key_type", "ai_apis"],
|
||||
)
|
||||
async def test_default_key_generate_params_allowed_routes_with_a_non_string_entry_is_ignored(
|
||||
monkeypatch: pytest.MonkeyPatch, request_data: GenerateKeyRequest, expected_routes: list[str]
|
||||
) -> None:
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"default_key_generate_params",
|
||||
{"allowed_routes": ["llm_api_routes", 123]},
|
||||
)
|
||||
mock_insert_data: Final = _wire_key_generation_prisma(monkeypatch)
|
||||
|
||||
key_row: Final = await _generate_key_and_get_persisted_row(request_data, mock_insert_data)
|
||||
|
||||
assert (key_row.get("allowed_routes") or []) == expected_routes
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_key_generate_params_allowed_routes_empty_list_treated_as_unset(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"default_key_generate_params",
|
||||
{"allowed_routes": ["llm_api_routes", "/key/spend/report"]},
|
||||
)
|
||||
mock_insert_data: Final = _wire_key_generation_prisma(monkeypatch)
|
||||
|
||||
key_row: Final = await _generate_key_and_get_persisted_row(
|
||||
GenerateKeyRequest(allowed_routes=[]),
|
||||
mock_insert_data,
|
||||
)
|
||||
|
||||
assert key_row["allowed_routes"] == ["llm_api_routes", "/key/spend/report"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_key_generate_params_allowed_routes_ignores_a_string_value(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "default_key_generate_params", {"allowed_routes": "llm_api_routes"})
|
||||
mock_insert_data: Final = _wire_key_generation_prisma(monkeypatch)
|
||||
|
||||
key_row: Final = await _generate_key_and_get_persisted_row(GenerateKeyRequest(), mock_insert_data)
|
||||
|
||||
assert not key_row.get("allowed_routes")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_generate_explicit_null_budget_duration_beats_default_key_generate_params(monkeypatch):
|
||||
"""An explicit `"budget_duration": null` asks for a budget that never resets.
|
||||
|
|
@ -22647,6 +22926,29 @@ async def test_team_admin_creates_an_llm_api_key_owned_by_an_internal_user_teamm
|
|||
assert (saved["key_type"], saved["allowed_routes"]) == ("llm_api", ["llm_api_routes"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"default_key_generate_params",
|
||||
[None, {"allowed_routes": ["/chat/completions"]}],
|
||||
ids=["default_off", "default_without_llm_api_routes"],
|
||||
)
|
||||
async def test_custom_key_policy_sees_the_routes_a_member_key_is_saved_with(
|
||||
monkeypatch: pytest.MonkeyPatch, default_key_generate_params: dict[str, list[str]] | None
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "default_key_generate_params", default_key_generate_params)
|
||||
policy: Final = AsyncMock(return_value={"decision": True})
|
||||
generate_key: Final = AsyncMock(return_value={"key": "sk-member", "user_id": "member-1", "team_id": "team-1"})
|
||||
|
||||
with patch("litellm.proxy.proxy_server.user_custom_key_policy", policy): # test-quality-ok: inject policy hook
|
||||
await _team_admin_generates_member_key(
|
||||
AsyncMock(return_value=MagicMock(user_role=LitellmUserRoles.INTERNAL_USER.value)), generate_key
|
||||
)
|
||||
|
||||
policy_request: Final = policy.await_args.args[0]
|
||||
assert policy_request.effective_key.allowed_routes == ["llm_api_routes"]
|
||||
assert policy_request.effective_key.allowed_routes == generate_key.await_args.kwargs["allowed_routes"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_admin_cannot_create_a_key_owned_by_a_proxy_admin_teammate():
|
||||
generate_key: Final = AsyncMock()
|
||||
|
|
|
|||
2
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
2
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -8441,7 +8441,7 @@ export interface paths {
|
|||
* - allowed_passthrough_routes: Optional[list] - List of allowed pass through endpoints for the key. Store the actual endpoint or store a wildcard pattern for a set of endpoints. Example - ["/my-custom-endpoint"]. Use this instead of allowed_routes, if you just want to specify which pass through endpoints the key can access, without specifying the routes. If allowed_routes is specified, allowed_pass_through_endpoints is ignored.
|
||||
* - denied_passthrough_routes: Optional[list] - List of pass through routes the key may not call, even if allowed by `allowed_passthrough_routes` or `allowed_routes`. Matches exact paths, path prefixes, and trailing `*` wildcards. Applies together with the team's `denied_passthrough_routes`. Example - ["/my-custom-endpoint/admin"].
|
||||
* - object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission.
|
||||
* - key_type: Optional[str] - Type of key that determines default allowed routes. Options: "llm_api" (can call LLM API routes), "management" (can call management routes), "read_only" (can only call info/read routes), "default" (uses default allowed routes). Defaults to "default".
|
||||
* - key_type: Optional[str] - Type of key that determines default allowed routes. Options: "llm_api" (can call LLM API routes), "management" (can call management routes), "read_only" (can only call info/read routes), "default" (Full Access in the Admin UI, skips `default_key_generate_params.allowed_routes`). When omitted, the key gets `default_key_generate_params.allowed_routes` if that is set and the request has no `allowed_routes`, else it works like "default".
|
||||
* - prompts: Optional[List[str]] - List of allowed prompts for the key. If specified, the key will only be able to use these specific prompts.
|
||||
* - auto_rotate: Optional[bool] - Whether this key should be automatically rotated (regenerated)
|
||||
* - rotation_interval: Optional[str] - How often to auto-rotate this key (e.g., '30s', '30m', '30h', '30d'). Required if auto_rotate=True.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue