diff --git a/backend/main.py b/backend/main.py index 4092cd63f69..83b42b0db06 100644 --- a/backend/main.py +++ b/backend/main.py @@ -20,14 +20,25 @@ DatabaseURLSettings.from_env().apply_to_env() from litellm.proxy.proxy_server import app +from backend.routers import ADMIN_PREFIX from backend.routes.allowlist import BACKEND_EXACT_PATHS, BACKEND_PATH_PREFIXES +_ADMIN_PREFIX = ADMIN_PREFIX + "/" + def _is_backend_route(route) -> bool: - """Keep the route on the backend if its path is in the management surface.""" + """Keep the route on the backend if it is an explicit admin route or its path + is in the (legacy, allowlisted) management surface. + + The explicit ``backend.routers`` admin routes are the source of truth and are + always kept; the allowlist covers the remaining proxy management surface that + has not been migrated onto explicit routers yet. + """ path = getattr(route, "path", None) if path is None: return False + if path == ADMIN_PREFIX or path.startswith(_ADMIN_PREFIX): + return True if isinstance(route, Mount): # Static UI mounts are served by the dedicated UI container, not here. return False diff --git a/backend/routers/__init__.py b/backend/routers/__init__.py new file mode 100644 index 00000000000..b997c78f687 --- /dev/null +++ b/backend/routers/__init__.py @@ -0,0 +1,15 @@ +"""Admin routes owned by the backend (control plane). + +These routers are the source of truth for the admin surface: they are defined +here and imported by ``litellm.proxy.proxy_server`` (which mounts them when the +backend package is importable) and served directly by ``backend.main``. Each +route authenticates through the ``auth_v2`` ``AuthSecurity`` stored on +``app.state.auth_v2`` rather than the legacy ``user_api_key_auth`` dependency. +""" + +from .teams import ADMIN_PREFIX +from .teams import router as admin_teams_router + +admin_routers = (admin_teams_router,) + +__all__ = ["ADMIN_PREFIX", "admin_routers", "admin_teams_router"] diff --git a/backend/routers/dependencies.py b/backend/routers/dependencies.py new file mode 100644 index 00000000000..52bcd84f437 --- /dev/null +++ b/backend/routers/dependencies.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +from typing import Callable, cast + +from fastapi import Request +from fastapi.security import SecurityScopes + +from litellm.proxy.auth_v2 import Principal, Role +from litellm.proxy.auth_v2.errors import forbidden_role +from litellm.proxy.auth_v2.resolvers import ProvisioningStore +from litellm.proxy.auth_v2.security import AuthSecurity + + +def get_auth(request: Request) -> AuthSecurity: + return cast(AuthSecurity, request.app.state.auth_v2) + + +def team_store(request: Request) -> ProvisioningStore: + return cast(ProvisioningStore, get_auth(request).resolver) + + +def require_roles(*allowed: Role) -> Callable[[Request], object]: + """Request-scoped role gate built on the ``auth_v2`` Security layer. + + Mirrors ``AuthSecurity.require_roles`` but reaches the per-app ``AuthSecurity`` + via ``request.app.state`` at request time, since these routers are wired into + the app after import rather than closing over an instance. + """ + + async def dependency(request: Request) -> Principal: + auth = get_auth(request) + principal = await auth.principal(SecurityScopes(scopes=[]), request) + if not auth.authorizer.has_any_role(principal, allowed): + raise forbidden_role() + return principal + + return dependency diff --git a/backend/routers/teams.py b/backend/routers/teams.py new file mode 100644 index 00000000000..f7465927443 --- /dev/null +++ b/backend/routers/teams.py @@ -0,0 +1,89 @@ +from __future__ import annotations + +from typing import List, Optional + +from fastapi import APIRouter, Depends, HTTPException, Request, Response, status +from pydantic import BaseModel, Field +from scim2_models import Group as ScimGroup +from scim2_models import GroupMember + +from litellm.proxy.auth_v2 import Role + +from .dependencies import require_roles, team_store + +ADMIN_PREFIX = "/admin" + +router = APIRouter(prefix=f"{ADMIN_PREFIX}/teams", tags=["admin"]) + +_team_admin = require_roles(Role.ORG_ADMIN, Role.PLATFORM_ADMIN) +_protected = [Depends(_team_admin)] + + +class TeamUpsert(BaseModel): + name: str + members: List[str] = Field(default_factory=list) + + +class TeamView(BaseModel): + id: str + name: str + members: List[str] + + +def _to_scim(team_id: Optional[str], body: TeamUpsert) -> ScimGroup: + group = ScimGroup( + display_name=body.name, + members=[GroupMember(value=user_id) for user_id in body.members] or None, + ) + if team_id is not None: + group.id = team_id + return group + + +def _to_view(group: ScimGroup) -> TeamView: + return TeamView( + id=group.id, + name=group.display_name, + members=[m.value for m in (group.members or []) if m.value], + ) + + +@router.post("", status_code=status.HTTP_201_CREATED, dependencies=_protected) +async def create_team(body: TeamUpsert, request: Request) -> TeamView: + stored = await team_store(request).upsert_group(_to_scim(None, body)) + return _to_view(stored) + + +@router.get("", dependencies=_protected) +async def list_teams(request: Request) -> List[TeamView]: + groups = await team_store(request).list_groups(None) + return [_to_view(group) for group in groups] + + +@router.get("/{team_id}", dependencies=_protected) +async def get_team(team_id: str, request: Request) -> TeamView: + group = await team_store(request).get_group(team_id) + if group is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, detail=f"team {team_id} not found" + ) + return _to_view(group) + + +@router.put("/{team_id}", dependencies=_protected) +async def upsert_team(team_id: str, body: TeamUpsert, request: Request) -> TeamView: + stored = await team_store(request).upsert_group(_to_scim(team_id, body)) + return _to_view(stored) + + +@router.delete( + "/{team_id}", status_code=status.HTTP_204_NO_CONTENT, dependencies=_protected +) +async def delete_team(team_id: str, request: Request) -> Response: + store = team_store(request) + if await store.get_group(team_id) is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, detail=f"team {team_id} not found" + ) + await store.delete_group(team_id) + return Response(status_code=status.HTTP_204_NO_CONTENT) diff --git a/litellm/proxy/auth_v2/resolvers.py b/litellm/proxy/auth_v2/resolvers.py index fe043b74a24..0f097375911 100644 --- a/litellm/proxy/auth_v2/resolvers.py +++ b/litellm/proxy/auth_v2/resolvers.py @@ -1,7 +1,7 @@ from __future__ import annotations import uuid -from typing import TYPE_CHECKING, List, Optional, Protocol, runtime_checkable +from typing import TYPE_CHECKING, List, Optional, Protocol, cast, runtime_checkable from scim2_models import Group as ScimGroup from scim2_models import User as ScimUser @@ -38,7 +38,9 @@ from litellm.repositories.user_repository import UserRepository if TYPE_CHECKING: from litellm.caching.caching import DualCache + from litellm.models.team import LiteLLM_TeamTable from litellm.models.user import LiteLLM_UserTable + from litellm.proxy.auth_v2.authorization import Role from litellm.proxy.utils import PrismaClient @@ -111,7 +113,29 @@ class DbIdentityStore(IdentityStore): raise errors.invalid_token() from exc if key.blocked: raise errors.account_disabled() - return self._principal_from_key(credential, key) + return self._principal_from_key(credential, key, await self._key_role(key)) + + async def _key_role(self, key: UserAPIKeyAuth) -> Optional[Role]: + """Platform role for an API key. + + ``get_key_object`` does not join the owning user's role onto the token, so + a key with a ``user_id`` but no ``user_role`` is resolved against the user + table (cache-backed) rather than coming back role-less. + """ + if key.user_role is not None: + return map_role(key.user_role) + if key.user_id is None: + return None + try: + user = await get_user_object( + user_id=key.user_id, + prisma_client=self._prisma, + user_api_key_cache=self._cache, + user_id_upsert=False, + ) + except Exception: + return None + return map_role(user.user_role) if user is not None else None async def _resolve_subject(self, credential: Credential) -> Principal: email = credential.claims.get("email") @@ -142,14 +166,16 @@ class DbIdentityStore(IdentityStore): ) def _principal_from_key( - self, credential: Credential, key: UserAPIKeyAuth + self, credential: Credential, key: UserAPIKeyAuth, role: Optional[Role] ) -> Principal: teams: List[TeamIdentity] = [] if key.team_id is not None: - role = ( + membership = ( team_role(key.team_member.role) if key.team_member else TeamRole.MEMBER ) - teams.append(TeamIdentity(id=key.team_id, name=key.team_alias, role=role)) + teams.append( + TeamIdentity(id=key.team_id, name=key.team_alias, role=membership) + ) organization = ( OrganizationIdentity(id=key.org_id, name=key.organization_alias) if key.org_id is not None @@ -160,7 +186,6 @@ class DbIdentityStore(IdentityStore): if key.user_id is not None else None ) - mapped = map_role(key.user_role) return Principal( principal_type=( PrincipalType.HUMAN if key.user_id else PrincipalType.SERVICE_ACCOUNT @@ -170,7 +195,7 @@ class DbIdentityStore(IdentityStore): user=user, organization=organization, teams=teams, - roles=[mapped] if mapped else [], + roles=[role] if role else [], scopes=list(credential.scopes), auth_method=credential.method, credential_ref=credential.credential_ref, @@ -266,29 +291,27 @@ class DbIdentityStore(IdentityStore): return [db_user_to_scim(row) for row in rows] async def upsert_group(self, group: ScimGroup) -> ScimGroup: + from prisma import Json + repo = TeamRepository(self._prisma) data = scim_group_to_db(group) - existing = ( - await repo.table.find_unique(where={"team_id": group.id}) - if group.id - else None - ) + data["members_with_roles"] = Json(cast(list, data["members_with_roles"])) + existing = await repo.find_by_id(group.id, "team_id") if group.id else None if existing is None: data["team_id"] = group.id or str(uuid.uuid4()) - stored = await repo.table.create(data=data) + stored: Optional[LiteLLM_TeamTable] = await repo.create(data) else: - stored = await repo.table.update(where={"team_id": group.id}, data=data) + stored = await repo.update(group.id, data, id_field="team_id") + assert stored is not None return db_team_to_scim(stored) async def get_group(self, resource_id: str) -> Optional[ScimGroup]: - stored = await TeamRepository(self._prisma).table.find_unique( - where={"team_id": resource_id} - ) + stored = await TeamRepository(self._prisma).find_by_id(resource_id, "team_id") return db_team_to_scim(stored) if stored is not None else None async def delete_group(self, resource_id: str) -> None: await TeamRepository(self._prisma).table.delete(where={"team_id": resource_id}) async def list_groups(self, filter_expr: Optional[str]) -> List[ScimGroup]: - rows = await TeamRepository(self._prisma).table.find_many() + rows = await TeamRepository(self._prisma).find_many() return [db_team_to_scim(row) for row in rows] diff --git a/litellm/proxy/auth_v2/utils.py b/litellm/proxy/auth_v2/utils.py index 9eda35cc89c..128af63c196 100644 --- a/litellm/proxy/auth_v2/utils.py +++ b/litellm/proxy/auth_v2/utils.py @@ -3,7 +3,7 @@ from __future__ import annotations import hashlib from typing import TYPE_CHECKING, Dict, List, Optional -from scim2_models import Email, Name +from scim2_models import Email, GroupMember, Name from scim2_models import Group as ScimGroup from scim2_models import User as ScimUser @@ -93,4 +93,11 @@ def scim_group_to_db(group: ScimGroup) -> Dict[str, object]: def db_team_to_scim(team: "LiteLLM_TeamTable") -> ScimGroup: result = ScimGroup(display_name=team.team_alias or team.team_id) result.id = team.team_id + members = [ + GroupMember(value=member.user_id) + for member in (team.members_with_roles or []) + if member.user_id + ] + if members: + result.members = members return result diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 96c9cd1e8fb..02ebecdf8af 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2,6 +2,7 @@ import asyncio import copy import enum import importlib +import importlib.util import inspect import io import os @@ -887,6 +888,17 @@ async def proxy_startup_event(app: FastAPI): # noqa: PLR0915 asyncio.create_task(_run_pw_migration()) + try: + from litellm.proxy.auth_v2 import AuthConfig, AuthSecurity + from litellm.proxy.auth_v2.resolvers import DbIdentityStore + + app.state.auth_v2 = AuthSecurity( + AuthConfig(), + DbIdentityStore(prisma_client, user_api_key_cache), + ) + except Exception as e: + verbose_proxy_logger.warning(f"auth_v2 wiring skipped: {e}") + ProxyStartupEvent._initialize_startup_logging( llm_router=llm_router, proxy_logging_obj=proxy_logging_obj, @@ -15818,6 +15830,15 @@ app.include_router(ui_discovery_endpoints_router) # Eager: /models/{name}:method overlaps with the OpenAI /models endpoint. app.include_router(google_router) +# Admin routes are defined in (and owned by) the `backend.routers` package; the +# proxy mounts them here as the source of truth when that package is importable +# (it ships with the source tree and Docker image, but not the pip wheel). +if importlib.util.find_spec("backend") is not None: + from backend.routers import admin_routers + + for _admin_router in admin_routers: + app.include_router(_admin_router) + attach_lazy_features(app) app.add_middleware( RequestSizeLimitMiddleware, diff --git a/tests/test_litellm/proxy/auth_v2/test_resolver.py b/tests/test_litellm/proxy/auth_v2/test_resolver.py index f2d40258081..84dec45ba53 100644 --- a/tests/test_litellm/proxy/auth_v2/test_resolver.py +++ b/tests/test_litellm/proxy/auth_v2/test_resolver.py @@ -77,6 +77,19 @@ async def test_api_key_resolves_to_principal_with_db_role(): assert principal.roles == [Role.ORG_ADMIN] +async def test_api_key_role_falls_back_to_owning_user(): + # get_key_object does not join the user's role onto the token, so a key with a + # user_id but no user_role must resolve the role from the user table + raw = "sk-live-noroll" + key = UserAPIKeyAuth(token=hash_token(raw), user_id="u-7") + user = LiteLLM_UserTable(user_id="u-7", user_role="proxy_admin") + store = _store({hash_token(raw): key, "u-7": user}) + + principal = await store.resolve(_api_key_credential(raw)) + + assert principal.roles == [Role.PLATFORM_ADMIN] + + async def test_api_key_lookup_is_keyed_on_hashed_token(): raw = "sk-live-abc" key = UserAPIKeyAuth(token=hash_token(raw), user_id="u-1") diff --git a/tests/test_litellm/proxy/auth_v2/test_utils.py b/tests/test_litellm/proxy/auth_v2/test_utils.py new file mode 100644 index 00000000000..ff82f4a245d --- /dev/null +++ b/tests/test_litellm/proxy/auth_v2/test_utils.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +from litellm.proxy._types import LiteLLM_TeamTable, Member +from litellm.proxy.auth_v2.utils import db_team_to_scim + + +def test_db_team_to_scim_carries_members(): + team = LiteLLM_TeamTable( + team_id="team-eng", + team_alias="eng", + members_with_roles=[ + Member(user_id="u-1", role="user"), + Member(user_id="u-2", role="admin"), + ], + ) + + group = db_team_to_scim(team) + + assert group.id == "team-eng" + assert group.display_name == "eng" + assert [member.value for member in group.members] == ["u-1", "u-2"] + + +def test_db_team_to_scim_without_members_omits_member_list(): + team = LiteLLM_TeamTable(team_id="team-solo", team_alias="solo") + + group = db_team_to_scim(team) + + assert group.id == "team-solo" + assert group.members is None diff --git a/tests/test_litellm/proxy/test_backend_admin_teams.py b/tests/test_litellm/proxy/test_backend_admin_teams.py new file mode 100644 index 00000000000..c594d52ccd1 --- /dev/null +++ b/tests/test_litellm/proxy/test_backend_admin_teams.py @@ -0,0 +1,188 @@ +"""Regression tests for the backend admin teams router (``backend/routers/teams.py``). + +These pin the two properties the router is responsible for: every route is gated +through the ``auth_v2`` ``AuthSecurity`` Security layer (role-checked, not the +legacy ``user_api_key_auth``), and team CRUD plus membership round-trips through +the injected identity store. +""" + +from __future__ import annotations + +import os +import sys +from typing import Dict, List, Optional + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient +from scim2_models import Group as ScimGroup + +# backend/ lives at the repo root, not inside litellm/. +_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")) +if _REPO_ROOT not in sys.path: + sys.path.insert(0, _REPO_ROOT) + +from backend.routers.teams import router as teams_router # noqa: E402 + +from litellm.proxy.auth_v2 import errors # noqa: E402 +from litellm.proxy.auth_v2.authenticators import APIKeyAuthenticator # noqa: E402 +from litellm.proxy.auth_v2.config import ApiKeySchemeConfig, AuthConfig # noqa: E402 +from litellm.proxy.auth_v2.authorization import Role # noqa: E402 +from litellm.proxy.auth_v2.models import ( # noqa: E402 + AuthMethod, + Credential, + Principal, + PrincipalType, +) +from litellm.proxy.auth_v2.security import AuthSecurity # noqa: E402 +from litellm.proxy.auth_v2.utils import hash_api_key # noqa: E402 + +ADMIN_KEY = "sk-admin" +READER_KEY = "sk-reader" + + +def _principal(subject: str, roles: List[Role]) -> Principal: + return Principal( + principal_type=PrincipalType.HUMAN, + subject=subject, + auth_method=AuthMethod.API_KEY, + roles=roles, + ) + + +class _FakeStore: + """IdentityResolver + team-group ProvisioningStore backed by a dict. + + Resolution maps API keys to fully-formed Principals (so role gating can be + driven directly), and the group methods store the exact ScimGroup handed in + so membership round-trips without a database. + """ + + def __init__(self, principals: Dict[str, Principal]) -> None: + self._by_key = principals + self._groups: Dict[str, ScimGroup] = {} + self._seq = 0 + + async def resolve(self, credential: Credential) -> Principal: + raw = credential.claims.get("_raw_api_key") + principal = ( + self._by_key.get(hash_api_key(raw)) if isinstance(raw, str) else None + ) + if principal is None: + raise errors.invalid_token() + return principal.model_copy() + + async def upsert_group(self, group: ScimGroup) -> ScimGroup: + if not group.id: + self._seq += 1 + group.id = f"team-{self._seq}" + self._groups[group.id] = group + return group + + async def get_group(self, resource_id: str) -> Optional[ScimGroup]: + return self._groups.get(resource_id) + + async def delete_group(self, resource_id: str) -> None: + self._groups.pop(resource_id, None) + + async def list_groups(self, filter_expr: Optional[str]) -> List[ScimGroup]: + return list(self._groups.values()) + + +@pytest.fixture +def client() -> TestClient: + resolver = _FakeStore( + { + hash_api_key(ADMIN_KEY): _principal("admin", [Role.ORG_ADMIN]), + hash_api_key(READER_KEY): _principal("reader", []), + } + ) + auth = AuthSecurity( + AuthConfig(), + resolver, + authenticators=[APIKeyAuthenticator(ApiKeySchemeConfig())], + ) + app = FastAPI() + app.state.auth_v2 = auth + app.include_router(teams_router) + return TestClient(app) + + +def _admin(headers: Optional[dict] = None) -> dict: + return {"x-litellm-api-key": ADMIN_KEY, **(headers or {})} + + +def test_create_requires_authentication(client): + response = client.post("/admin/teams", json={"name": "eng"}) + assert response.status_code == 401 + assert response.json()["detail"] == "Not authenticated" + + +def test_create_rejects_unknown_key(client): + response = client.post( + "/admin/teams", + json={"name": "eng"}, + headers={"x-litellm-api-key": "sk-bogus"}, + ) + assert response.status_code == 401 + + +def test_create_denied_without_admin_role(client): + response = client.post( + "/admin/teams", + json={"name": "eng"}, + headers={"x-litellm-api-key": READER_KEY}, + ) + assert response.status_code == 403 + assert response.json()["detail"] == "Insufficient role" + + +def test_create_then_get_round_trips_membership(client): + created = client.post( + "/admin/teams", + json={"name": "eng", "members": ["u-1", "u-2"]}, + headers=_admin(), + ) + assert created.status_code == 201 + body = created.json() + assert body["name"] == "eng" + assert body["members"] == ["u-1", "u-2"] + team_id = body["id"] + + fetched = client.get(f"/admin/teams/{team_id}", headers=_admin()) + assert fetched.status_code == 200 + assert fetched.json() == {"id": team_id, "name": "eng", "members": ["u-1", "u-2"]} + + +def test_list_returns_created_team(client): + client.post("/admin/teams", json={"name": "eng"}, headers=_admin()) + listed = client.get("/admin/teams", headers=_admin()) + assert listed.status_code == 200 + assert [team["name"] for team in listed.json()] == ["eng"] + + +def test_update_replaces_membership(client): + team_id = client.post( + "/admin/teams", json={"name": "eng", "members": ["u-1"]}, headers=_admin() + ).json()["id"] + + updated = client.put( + f"/admin/teams/{team_id}", + json={"name": "eng", "members": ["u-2", "u-3"]}, + headers=_admin(), + ) + assert updated.status_code == 200 + assert updated.json()["members"] == ["u-2", "u-3"] + + +def test_delete_removes_team(client): + team_id = client.post( + "/admin/teams", json={"name": "eng"}, headers=_admin() + ).json()["id"] + + assert client.delete(f"/admin/teams/{team_id}", headers=_admin()).status_code == 204 + assert client.get(f"/admin/teams/{team_id}", headers=_admin()).status_code == 404 + + +def test_delete_missing_team_is_404(client): + assert client.delete("/admin/teams/nope", headers=_admin()).status_code == 404 diff --git a/tests/test_litellm/proxy/test_component_allowlists.py b/tests/test_litellm/proxy/test_component_allowlists.py index ad25856b972..ee6153bb2b9 100644 --- a/tests/test_litellm/proxy/test_component_allowlists.py +++ b/tests/test_litellm/proxy/test_component_allowlists.py @@ -42,6 +42,7 @@ _REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", if _REPO_ROOT not in sys.path: sys.path.insert(0, _REPO_ROOT) +from backend.routers import ADMIN_PREFIX from backend.routes.allowlist import BACKEND_EXACT_PATHS, BACKEND_PATH_PREFIXES from gateway.routes.allowlist import GATEWAY_EXACT_PATHS, GATEWAY_PATH_PREFIXES from litellm.proxy.proxy_server import app @@ -53,7 +54,7 @@ for _key, _previous in _PRE_EXISTING_ENV.items(): os.environ[_key] = _previous -def _component_paths(routes, exact_paths, path_prefixes) -> set[str]: +def _component_paths(routes, exact_paths, path_prefixes, keep_prefixes=()) -> set[str]: """Reproduce ``gateway.main._is_gateway_route`` / ``backend.main._is_backend_route``.""" out: set[str] = set() for route in routes: @@ -62,7 +63,9 @@ def _component_paths(routes, exact_paths, path_prefixes) -> set[str]: path = getattr(route, "path", None) if path is None: continue - if path in exact_paths or any(path.startswith(p) for p in path_prefixes): + if any(path == p or path.startswith(p + "/") for p in keep_prefixes): + out.add(path) + elif path in exact_paths or any(path.startswith(p) for p in path_prefixes): out.add(path) return out @@ -78,7 +81,10 @@ def test_gateway_plus_backend_covers_full_app(): app.router.routes, GATEWAY_EXACT_PATHS, GATEWAY_PATH_PREFIXES ) backend_paths = _component_paths( - app.router.routes, BACKEND_EXACT_PATHS, BACKEND_PATH_PREFIXES + app.router.routes, + BACKEND_EXACT_PATHS, + BACKEND_PATH_PREFIXES, + keep_prefixes=(ADMIN_PREFIX,), ) uncovered = all_paths - (gateway_paths | backend_paths)