From 03f1c11e9b30d0762eb16f5b5d2052e38caedfc5 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Wed, 17 Jun 2026 13:57:04 -0700 Subject: [PATCH] feat(auth_v2): add project and end_user identity to Principal Principal carried user, organization and teams but omitted the project and end-user attribution axes that the existing key path tracks. Add ProjectIdentity and EndUserIdentity sub-models and the matching optional fields, and project them off the key object in _principal_from_key: project_id/project_alias map to ProjectIdentity, end_user_id to EndUserIdentity. Both stay None when absent. project_id is a column on the verification token, so it resolves from the combined-view key object directly. end_user_id is request-scoped and will be stamped at the seam, the same way network context is; the resolver maps it whenever the key carries it. --- litellm/proxy/auth_v2/README.md | 5 +-- litellm/proxy/auth_v2/models.py | 11 +++++++ litellm/proxy/auth_v2/resolvers.py | 12 +++++++ .../test_litellm/proxy/auth_v2/test_models.py | 14 +++++++++ .../proxy/auth_v2/test_resolver.py | 31 +++++++++++++++++++ 5 files changed, 71 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/auth_v2/README.md b/litellm/proxy/auth_v2/README.md index b316a05e4c8..2bdd8143f97 100644 --- a/litellm/proxy/auth_v2/README.md +++ b/litellm/proxy/auth_v2/README.md @@ -16,8 +16,9 @@ the scheme, the subject, scopes/claims from the token, and, for exchangeable bea tokens, the raw token for downstream RFC 8693 token exchange. `Principal` is what the route handler receives: a normalized identity with the user, -organization, teams, roles and scopes filled in. It holds identity only, no budget or -policy state. Scope checking lives here as `Principal.has_required_scopes`. +organization, teams, project, end user, roles and scopes filled in. It holds identity +only, no budget or policy state. Scope checking lives here as +`Principal.has_required_scopes`. The split matters: authentication proves the credential, resolution turns the proven credential into a known identity, and only then can authorization run. diff --git a/litellm/proxy/auth_v2/models.py b/litellm/proxy/auth_v2/models.py index c94417489a5..d5afb46fdbf 100644 --- a/litellm/proxy/auth_v2/models.py +++ b/litellm/proxy/auth_v2/models.py @@ -68,6 +68,15 @@ class TeamIdentity(BaseModel): role: TeamRole = TeamRole.MEMBER +class ProjectIdentity(BaseModel): + id: str + name: Optional[str] = None + + +class EndUserIdentity(BaseModel): + id: str + + class CredentialRef(BaseModel): key_id: Optional[str] = None token_id: Optional[str] = None @@ -118,6 +127,8 @@ class Principal(BaseModel): user: Optional[UserIdentity] = None organization: Optional[OrganizationIdentity] = None teams: List[TeamIdentity] = Field(default_factory=list) + project: Optional[ProjectIdentity] = None + end_user: Optional[EndUserIdentity] = None roles: List[Role] = Field(default_factory=list) scopes: List[str] = Field(default_factory=list) diff --git a/litellm/proxy/auth_v2/resolvers.py b/litellm/proxy/auth_v2/resolvers.py index fe043b74a24..156c435ec3b 100644 --- a/litellm/proxy/auth_v2/resolvers.py +++ b/litellm/proxy/auth_v2/resolvers.py @@ -17,9 +17,11 @@ from litellm.proxy.auth_v2 import errors from litellm.proxy.auth_v2.models import ( AuthMethod, Credential, + EndUserIdentity, OrganizationIdentity, Principal, PrincipalType, + ProjectIdentity, TeamIdentity, TeamRole, UserIdentity, @@ -160,6 +162,14 @@ class DbIdentityStore(IdentityStore): if key.user_id is not None else None ) + project = ( + ProjectIdentity(id=key.project_id, name=key.project_alias) + if key.project_id is not None + else None + ) + end_user = ( + EndUserIdentity(id=key.end_user_id) if key.end_user_id is not None else None + ) mapped = map_role(key.user_role) return Principal( principal_type=( @@ -170,6 +180,8 @@ class DbIdentityStore(IdentityStore): user=user, organization=organization, teams=teams, + project=project, + end_user=end_user, roles=[mapped] if mapped else [], scopes=list(credential.scopes), auth_method=credential.method, diff --git a/tests/test_litellm/proxy/auth_v2/test_models.py b/tests/test_litellm/proxy/auth_v2/test_models.py index aa330e14b20..e36c83b8484 100644 --- a/tests/test_litellm/proxy/auth_v2/test_models.py +++ b/tests/test_litellm/proxy/auth_v2/test_models.py @@ -6,8 +6,10 @@ from pydantic import ValidationError from litellm.proxy.auth_v2.models import ( AuthMethod, Credential, + EndUserIdentity, Principal, PrincipalType, + ProjectIdentity, SecuritySchemeType, TeamIdentity, TeamRole, @@ -70,6 +72,8 @@ def test_principal_default_network_and_collections(): ) assert principal.teams == [] assert principal.scopes == [] + assert principal.project is None + assert principal.end_user is None assert principal.network.client_ip is None assert principal.network.via_trusted_proxy is False @@ -91,3 +95,13 @@ def test_user_identity_optional_fields_default_none(): user = UserIdentity(id="u1") assert user.email is None assert user.external_id is None + + +def test_project_identity_name_is_optional(): + assert ProjectIdentity(id="p1").name is None + assert ProjectIdentity(id="p1", name="Acme").name == "Acme" + + +def test_end_user_identity_requires_id(): + with pytest.raises(ValidationError): + EndUserIdentity() diff --git a/tests/test_litellm/proxy/auth_v2/test_resolver.py b/tests/test_litellm/proxy/auth_v2/test_resolver.py index f2d40258081..7ff048466ea 100644 --- a/tests/test_litellm/proxy/auth_v2/test_resolver.py +++ b/tests/test_litellm/proxy/auth_v2/test_resolver.py @@ -77,6 +77,37 @@ async def test_api_key_resolves_to_principal_with_db_role(): assert principal.roles == [Role.ORG_ADMIN] +async def test_api_key_principal_carries_project_and_end_user(): + raw = "sk-live-proj" + key = UserAPIKeyAuth( + token=hash_token(raw), + user_id="u-1", + project_id="proj-1", + project_alias="Acme Prod", + end_user_id="cust-7", + ) + store = _store({hash_token(raw): key}) + + principal = await store.resolve(_api_key_credential(raw)) + + assert principal.project is not None + assert principal.project.id == "proj-1" + assert principal.project.name == "Acme Prod" + assert principal.end_user is not None + assert principal.end_user.id == "cust-7" + + +async def test_api_key_principal_omits_project_and_end_user_when_absent(): + raw = "sk-live-bare" + key = UserAPIKeyAuth(token=hash_token(raw), user_id="u-1") + store = _store({hash_token(raw): key}) + + principal = await store.resolve(_api_key_credential(raw)) + + assert principal.project is None + assert principal.end_user is None + + 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")