From 2e999ee1ad00e8817fb9090c42ff8c6e9b2ec35c Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 12 Jun 2026 21:17:09 -0700 Subject: [PATCH] refactor(auth_v2): flatten resolvers package into resolvers.py + utils.py There was only ever one resolver, so the resolvers/ package (base, db, utils) collapses into a single resolvers.py holding the protocols and DbIdentityStore, plus a utils.py for the pure SCIM/role-mapping helpers. Drops the unused roles_from_claims/public_claims helpers. --- litellm/proxy/auth_v2/README.md | 8 +- .../auth_v2/{resolvers/db.py => resolvers.py} | 96 ++++++++++++++++--- litellm/proxy/auth_v2/resolvers/__init__.py | 15 --- litellm/proxy/auth_v2/resolvers/base.py | 45 --------- .../proxy/auth_v2/{resolvers => }/utils.py | 18 +--- .../proxy/auth_v2/test_resolver.py | 2 +- .../proxy/auth_v2/test_security.py | 2 +- 7 files changed, 91 insertions(+), 95 deletions(-) rename litellm/proxy/auth_v2/{resolvers/db.py => resolvers.py} (72%) delete mode 100644 litellm/proxy/auth_v2/resolvers/__init__.py delete mode 100644 litellm/proxy/auth_v2/resolvers/base.py rename litellm/proxy/auth_v2/{resolvers => }/utils.py (84%) diff --git a/litellm/proxy/auth_v2/README.md b/litellm/proxy/auth_v2/README.md index 2c3c4818d14..b316a05e4c8 100644 --- a/litellm/proxy/auth_v2/README.md +++ b/litellm/proxy/auth_v2/README.md @@ -86,10 +86,10 @@ built once at the composition root. dispatch to it by where its credential lives). `build_authenticators` constructs and orders them from `AuthConfig`. JWT verification for OIDC/OAuth2 is shared via `JWTVerifier`. -`resolvers/` holds the `IdentityResolver` / `IdentityStore` protocols and their -implementations (`DbIdentityStore` against Prisma, an in-memory store for tests). The -store also handles SCIM user/group provisioning so a provisioned user is immediately -resolvable. +`resolvers.py` holds the `IdentityResolver` / `IdentityStore` protocols and the single +`DbIdentityStore` implementation against Prisma. The store also handles SCIM user/group +provisioning so a provisioned user is immediately resolvable. `utils.py` holds the pure +SCIM/role-mapping helpers the store uses. `authorization/` holds the `Authorizer` protocol and its implementations: `RBACEngine` (Casbin role hierarchy and policy) and `ABACEngine`. `Role` and the JWT/OIDC/SAML role diff --git a/litellm/proxy/auth_v2/resolvers/db.py b/litellm/proxy/auth_v2/resolvers.py similarity index 72% rename from litellm/proxy/auth_v2/resolvers/db.py rename to litellm/proxy/auth_v2/resolvers.py index 7f7957d6023..fe043b74a24 100644 --- a/litellm/proxy/auth_v2/resolvers/db.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 +from typing import TYPE_CHECKING, List, Optional, Protocol, runtime_checkable from scim2_models import Group as ScimGroup from scim2_models import User as ScimUser @@ -24,8 +24,7 @@ from litellm.proxy.auth_v2.models import ( TeamRole, UserIdentity, ) -from litellm.proxy.auth_v2.resolvers.base import IdentityStore -from litellm.proxy.auth_v2.resolvers.utils import ( +from litellm.proxy.auth_v2.utils import ( db_team_to_scim, db_user_to_scim, map_role, @@ -43,6 +42,43 @@ if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient +@runtime_checkable +class IdentityResolver(Protocol): + async def resolve(self, credential: Credential) -> Principal: + """Resolve a verified credential to a Principal. + + Must return a freshly constructed Principal, never a cached or shared + instance. The caller stamps request-scoped state (the network context) + onto the returned object, so handing back a shared one would leak that + state across concurrent requests for the same identity. Cache the + underlying identity lookups (as the DB resolver does), not the assembled + Principal. + """ + ... + + +@runtime_checkable +class ProvisioningStore(Protocol): + async def upsert_user(self, user: ScimUser) -> ScimUser: ... + async def get_user(self, resource_id: str) -> Optional[ScimUser]: ... + async def deactivate_user(self, resource_id: str) -> None: ... + async def list_users(self, filter_expr: Optional[str]) -> List[ScimUser]: ... + async def upsert_group(self, group: ScimGroup) -> ScimGroup: ... + async def get_group(self, resource_id: str) -> Optional[ScimGroup]: ... + async def delete_group(self, resource_id: str) -> None: ... + async def list_groups(self, filter_expr: Optional[str]) -> List[ScimGroup]: ... + + +@runtime_checkable +class IdentityStore(IdentityResolver, ProvisioningStore, Protocol): + """An identity backend: resolves credentials and provisions SCIM users/groups. + + This is the single interface every implementation satisfies (in-memory, + database, ...). Resolution and provisioning live behind one store so a + provisioned user is immediately resolvable. + """ + + class DbIdentityStore(IdentityStore): """Resolves credentials against the proxy's Prisma tables and provisions SCIM users/groups into ``LiteLLM_UserTable`` / ``LiteLLM_TeamTable``. @@ -105,18 +141,30 @@ class DbIdentityStore(IdentityStore): credential_ref=credential.credential_ref, ) - def _principal_from_key(self, credential: Credential, key: UserAPIKeyAuth) -> Principal: + def _principal_from_key( + self, credential: Credential, key: UserAPIKeyAuth + ) -> Principal: teams: List[TeamIdentity] = [] if key.team_id is not None: - role = team_role(key.team_member.role) if key.team_member else TeamRole.MEMBER + role = ( + 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)) organization = ( - OrganizationIdentity(id=key.org_id, name=key.organization_alias) if key.org_id is not None else None + OrganizationIdentity(id=key.org_id, name=key.organization_alias) + if key.org_id is not None + else None + ) + user = ( + UserIdentity(id=key.user_id, email=key.user_email) + if key.user_id is not None + else None ) - user = UserIdentity(id=key.user_id, email=key.user_email) 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), + principal_type=( + PrincipalType.HUMAN if key.user_id else PrincipalType.SERVICE_ACCOUNT + ), subject=key.user_id or key.key_alias or credential.subject, issuer=credential.issuer, user=user, @@ -128,7 +176,9 @@ class DbIdentityStore(IdentityStore): credential_ref=credential.credential_ref, ) - async def _principal_from_user(self, credential: Credential, user: "LiteLLM_UserTable") -> Principal: + async def _principal_from_user( + self, credential: Credential, user: "LiteLLM_UserTable" + ) -> Principal: teams: List[TeamIdentity] = [] for team_id in user.teams or []: try: @@ -164,7 +214,9 @@ class DbIdentityStore(IdentityStore): credential_ref=credential.credential_ref, ) - async def _organization(self, user: "LiteLLM_UserTable") -> Optional[OrganizationIdentity]: + async def _organization( + self, user: "LiteLLM_UserTable" + ) -> Optional[OrganizationIdentity]: if user.organization_id is None: return None try: @@ -180,7 +232,11 @@ class DbIdentityStore(IdentityStore): async def upsert_user(self, user: ScimUser) -> ScimUser: repo = UserRepository(self._prisma) data = scim_user_to_db(user) - existing = await repo.table.find_unique(where={"user_id": user.id}) if user.id else None + existing = ( + await repo.table.find_unique(where={"user_id": user.id}) + if user.id + else None + ) if existing is None: data["user_id"] = user.id or str(uuid.uuid4()) stored = await repo.table.create(data=data) @@ -189,7 +245,9 @@ class DbIdentityStore(IdentityStore): return db_user_to_scim(stored) async def get_user(self, resource_id: str) -> Optional[ScimUser]: - stored = await UserRepository(self._prisma).table.find_unique(where={"user_id": resource_id}) + stored = await UserRepository(self._prisma).table.find_unique( + where={"user_id": resource_id} + ) return db_user_to_scim(stored) if stored is not None else None async def deactivate_user(self, resource_id: str) -> None: @@ -199,7 +257,9 @@ class DbIdentityStore(IdentityStore): return metadata = dict(getattr(stored, "metadata", None) or {}) metadata["scim_active"] = False - await repo.table.update(where={"user_id": resource_id}, data={"metadata": metadata}) + await repo.table.update( + where={"user_id": resource_id}, data={"metadata": metadata} + ) async def list_users(self, filter_expr: Optional[str]) -> List[ScimUser]: rows = await UserRepository(self._prisma).table.find_many() @@ -208,7 +268,11 @@ class DbIdentityStore(IdentityStore): async def upsert_group(self, group: ScimGroup) -> ScimGroup: 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 + existing = ( + await repo.table.find_unique(where={"team_id": group.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) @@ -217,7 +281,9 @@ class DbIdentityStore(IdentityStore): 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).table.find_unique( + where={"team_id": resource_id} + ) return db_team_to_scim(stored) if stored is not None else None async def delete_group(self, resource_id: str) -> None: diff --git a/litellm/proxy/auth_v2/resolvers/__init__.py b/litellm/proxy/auth_v2/resolvers/__init__.py deleted file mode 100644 index 546366a1b94..00000000000 --- a/litellm/proxy/auth_v2/resolvers/__init__.py +++ /dev/null @@ -1,15 +0,0 @@ -from litellm.proxy.auth_v2.resolvers.base import ( - IdentityResolver, - IdentityStore, - ProvisioningStore, -) - -# DbIdentityStore is intentionally not re-exported here: it pulls in the v1 -# proxy DB machinery (auth_checks, repositories). Import it directly from -# litellm.proxy.auth_v2.resolvers.db when wiring a database-backed store. - -__all__ = [ - "IdentityResolver", - "ProvisioningStore", - "IdentityStore", -] diff --git a/litellm/proxy/auth_v2/resolvers/base.py b/litellm/proxy/auth_v2/resolvers/base.py deleted file mode 100644 index 046f7f75304..00000000000 --- a/litellm/proxy/auth_v2/resolvers/base.py +++ /dev/null @@ -1,45 +0,0 @@ -from __future__ import annotations - -from typing import List, Optional, Protocol, runtime_checkable - -from scim2_models import Group as ScimGroup -from scim2_models import User as ScimUser - -from litellm.proxy.auth_v2.models import Credential, Principal - - -@runtime_checkable -class IdentityResolver(Protocol): - async def resolve(self, credential: Credential) -> Principal: - """Resolve a verified credential to a Principal. - - Must return a freshly constructed Principal, never a cached or shared - instance. The caller stamps request-scoped state (the network context) - onto the returned object, so handing back a shared one would leak that - state across concurrent requests for the same identity. Cache the - underlying identity lookups (as the DB resolver does), not the assembled - Principal. - """ - ... - - -@runtime_checkable -class ProvisioningStore(Protocol): - async def upsert_user(self, user: ScimUser) -> ScimUser: ... - async def get_user(self, resource_id: str) -> Optional[ScimUser]: ... - async def deactivate_user(self, resource_id: str) -> None: ... - async def list_users(self, filter_expr: Optional[str]) -> List[ScimUser]: ... - async def upsert_group(self, group: ScimGroup) -> ScimGroup: ... - async def get_group(self, resource_id: str) -> Optional[ScimGroup]: ... - async def delete_group(self, resource_id: str) -> None: ... - async def list_groups(self, filter_expr: Optional[str]) -> List[ScimGroup]: ... - - -@runtime_checkable -class IdentityStore(IdentityResolver, ProvisioningStore, Protocol): - """An identity backend: resolves credentials and provisions SCIM users/groups. - - This is the single interface every implementation satisfies (in-memory, - database, ...). Resolution and provisioning live behind one store so a - provisioned user is immediately resolvable. - """ diff --git a/litellm/proxy/auth_v2/resolvers/utils.py b/litellm/proxy/auth_v2/utils.py similarity index 84% rename from litellm/proxy/auth_v2/resolvers/utils.py rename to litellm/proxy/auth_v2/utils.py index 9a85ced7a10..9eda35cc89c 100644 --- a/litellm/proxy/auth_v2/resolvers/utils.py +++ b/litellm/proxy/auth_v2/utils.py @@ -1,7 +1,7 @@ from __future__ import annotations import hashlib -from typing import TYPE_CHECKING, Any, Dict, List, Optional +from typing import TYPE_CHECKING, Dict, List, Optional from scim2_models import Email, Name from scim2_models import Group as ScimGroup @@ -19,18 +19,6 @@ def hash_api_key(raw: str) -> str: return hashlib.sha256(raw.encode("utf-8")).hexdigest() -def roles_from_claims(claims: Dict[str, Any]) -> List[Role]: - raw = claims.get("roles", []) - if not isinstance(raw, list): - return [] - valid = {role.value for role in Role} - return [Role(value) for value in raw if value in valid] - - -def public_claims(claims: Dict[str, Any]) -> Dict[str, Any]: - return {key: value for key, value in claims.items() if not key.startswith("_")} - - _ROLE_MAP: Dict[str, Role] = { "proxy_admin": Role.PLATFORM_ADMIN, "proxy_admin_viewer": Role.PLATFORM_VIEWER, @@ -96,7 +84,9 @@ def db_user_to_scim(user: "LiteLLM_UserTable") -> ScimUser: def scim_group_to_db(group: ScimGroup) -> Dict[str, object]: - members = [{"user_id": member.value, "role": "user"} for member in (group.members or [])] + members = [ + {"user_id": member.value, "role": "user"} for member in (group.members or []) + ] return {"team_alias": group.display_name, "members_with_roles": members} diff --git a/tests/test_litellm/proxy/auth_v2/test_resolver.py b/tests/test_litellm/proxy/auth_v2/test_resolver.py index a53b04a84e1..f2d40258081 100644 --- a/tests/test_litellm/proxy/auth_v2/test_resolver.py +++ b/tests/test_litellm/proxy/auth_v2/test_resolver.py @@ -14,7 +14,7 @@ from litellm.proxy.auth_v2.models import ( PrincipalType, SecuritySchemeType, ) -from litellm.proxy.auth_v2.resolvers.db import DbIdentityStore +from litellm.proxy.auth_v2.resolvers import DbIdentityStore class _FakeCache: diff --git a/tests/test_litellm/proxy/auth_v2/test_security.py b/tests/test_litellm/proxy/auth_v2/test_security.py index 9f01dfc29d1..1f41d898ce3 100644 --- a/tests/test_litellm/proxy/auth_v2/test_security.py +++ b/tests/test_litellm/proxy/auth_v2/test_security.py @@ -29,7 +29,7 @@ from litellm.proxy.auth_v2.models import ( SecuritySchemeType, ) from litellm.proxy.auth_v2.authorization import RBACEngine, Role -from litellm.proxy.auth_v2.resolvers.utils import hash_api_key +from litellm.proxy.auth_v2.utils import hash_api_key from litellm.proxy.auth_v2.security import AuthSecurity from auth_v2_helpers import TEST_AUDIENCE, TEST_ISSUER, FakeJwksClient