scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 227

Merge pull request #40904 from BerriAI/litellm_jwt_agent_id_claim

Browse source 
feat(proxy): bind JWT claims to registered agents via agent_id_jwt_field
This commit is contained in:
Yassin Kortam 2026-09-15 12:46:18 -07:00 • committed by GitHub
commit 60808520df
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 601 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
litellm/proxy/_types.py
View file

@ -4716,6 +4716,7 @@ class JWTAuthBuilderResult(TypedDict):
org_id: str | None
team_membership: LiteLLM_TeamMembership | None
jwt_claims: dict # Decoded JWT token claims (avoids re-decoding)
agent_id: ReadOnly[str | None]
class ClientSideFallbackModel(TypedDict, total=False):
@ -4954,6 +4955,14 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
user_allowed_roles: list[str] | None = None
user_id_upsert: bool = Field(default=False, description="If user doesn't exist, upsert them into the db.")
end_user_id_jwt_field: str | None = None
agent_id_jwt_field: str | None = Field(
default=None,
description=(
"The field in the JWT token that identifies the calling agent (e.g. 'azp' for a Microsoft Entra ID "
"app token). Supports dot notation. The value is matched against a registered agent's agent_id, "
"then agent_name, and the request is rejected when it matches neither."
),
)
public_key_ttl: float = 600
public_key_stale_ttl: float = Field(
default=DEFAULT_JWKS_STALE_TTL,

71
litellm/proxy/auth/handle_jwt.py
View file

@ -14,7 +14,7 @@ import hashlib
import os
import re
import time
from collections.abc import Awaitable, Callable, Sequence
from collections.abc import Awaitable, Callable, Mapping, Sequence
from typing import Any, Final, Literal, NoReturn, Protocol, TypeVar, cast
import httpx
@ -61,6 +61,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
)
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.repositories.user_repository import UserRepository
from litellm.types.agents import AgentResponse
from .auth_checks import (
_allowed_routes_check,
@ -127,6 +128,26 @@ class _UserInfoResponse(Protocol):
def json(self) -> dict[str, object]: ...
class AgentLookup(Protocol):
"""The registered-agent lookups a JWT agent claim is matched against."""
def get_agent_by_id(self, agent_id: str) -> AgentResponse | None:
"""The agent registered under ``agent_id``, if any."""
def get_agent_by_name(self, agent_name: str) -> AgentResponse | None:
"""The agent registered under ``agent_name``, if any."""
class _NoRegisteredAgents:
"""The lookup in force until the proxy binds its agent registry: no agent is registered, so no claim matches."""
def get_agent_by_id(self, agent_id: str) -> None:
return None
def get_agent_by_name(self, agent_name: str) -> None:
return None
def _discovery_document(response: _OIDCDiscoveryResponse) -> _OIDCDiscoveryBody:
"""Decode an OIDC discovery response body."""
return response.json()
@ -198,6 +219,10 @@ class JWTHandler:
self.leeway = 0
# Per-cache-key locks so a TTL lapse triggers one refresh instead of one per in-flight request.
self._refresh_locks: dict[str, asyncio.Lock] = {} # mutable-ok: lock registry, keyed by JWKS url
self.agent_lookup: AgentLookup = _NoRegisteredAgents()
def bind_agent_lookup(self, agent_lookup: AgentLookup) -> None:
self.agent_lookup = agent_lookup
def update_environment(
self,
@ -623,6 +648,12 @@ class JWTHandler:
object_id = default_value
return object_id
def get_agent_claim(self, token: Mapping[str, object]) -> str | None:
if self.litellm_jwtauth.agent_id_jwt_field is None:
return None
claim: Final[object] = get_nested_value(data=token, key_path=self.litellm_jwtauth.agent_id_jwt_field)
return claim if isinstance(claim, str) and claim else None
def get_org_id(self, token: dict, default_value: str | None) -> str | None:
if self._has_trusted_issuer_normalized_claim(token=token, claim=self.LITELLM_ORG_ID_CLAIM):
return token.get(self.LITELLM_ORG_ID_CLAIM)
@ -1380,6 +1411,7 @@ class JWTAuthManager:
api_key: str,
jwt_valid_token: dict | None = None,
user_email: str | None = None,
agent_id: str | None = None,
) -> JWTAuthBuilderResult | None:
"""Check admin status and route access permissions"""
if not jwt_handler.is_admin(scopes=scopes):
@ -1409,8 +1441,28 @@ class JWTAuthManager:
org_id=org_id,
team_membership=None,
jwt_claims=jwt_valid_token or {},
agent_id=agent_id,
)
@staticmethod
def resolve_agent_id(
jwt_handler: JWTHandler,
jwt_valid_token: Mapping[str, object],
agent_registry: AgentLookup,
) -> str | None:
agent_claim: Final = jwt_handler.get_agent_claim(token=jwt_valid_token)
if agent_claim is None:
return None
agent: Final = agent_registry.get_agent_by_id(agent_id=agent_claim) or agent_registry.get_agent_by_name(
agent_name=agent_claim
)
if agent is None:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail=f"No registered agent matches JWT claim {jwt_handler.litellm_jwtauth.agent_id_jwt_field}={agent_claim}",
)
return agent.agent_id
@staticmethod
async def find_and_validate_specific_team_id(
jwt_handler: JWTHandler,
@ -2268,9 +2320,23 @@ class JWTAuthManager:
elif rbac_role == LitellmUserRoles.INTERNAL_USER:
user_id = object_id
agent_id: Final = JWTAuthManager.resolve_agent_id(
jwt_handler=jwt_handler,
jwt_valid_token=jwt_valid_token,
agent_registry=jwt_handler.agent_lookup,
)
# Check admin access
admin_result: Final = await JWTAuthManager.check_admin_access(
jwt_handler, scopes, route, user_id, org_id, api_key, jwt_valid_token, user_email=user_email
jwt_handler,
scopes,
route,
user_id,
org_id,
api_key,
jwt_valid_token,
user_email=user_email,
agent_id=agent_id,
)
if admin_result:
await JWTAuthManager._attach_team_from_header_for_admin(
@ -2514,4 +2580,5 @@ class JWTAuthManager:
token=api_key,
team_membership=team_membership_object,
jwt_claims=jwt_valid_token,
agent_id=agent_id,
)

6
litellm/proxy/auth/user_api_key_auth.py
View file

@ -852,6 +852,7 @@ async def _auto_register_jwt_mapping(
user_id: str | None = None,
org_id: str | None = None,
end_user_id: str | None = None,
agent_id: str | None = None,
) -> UserAPIKeyAuth | None:
"""
Auto-register: create a new virtual key + mapping for an unrecognised JWT
@ -884,6 +885,7 @@ async def _auto_register_jwt_mapping(
team_id=team_id,
user_id=user_id,
organization_id=org_id,
agent_id=agent_id,
metadata={
"auto_registered": True,
"jwt_claim_field": virtual_key_claim_field,
@ -1567,6 +1569,7 @@ async def _user_api_key_auth_builder(
org_id: Final = result["org_id"]
team_membership: Final[LiteLLM_TeamMembership | None] = result.get("team_membership", None)
jwt_claims = result.get("jwt_claims", None)
agent_id: Final[str | None] = result.get("agent_id")
if is_proxy_admin:
# Proxy admins authenticate via auth_builder (full
@ -1592,6 +1595,7 @@ async def _user_api_key_auth_builder(
end_user_id=end_user_id,
parent_otel_span=parent_otel_span,
jwt_claims=jwt_claims,
agent_id=agent_id,
**team_grants(team_object=team_object, team_membership=team_membership, user_id=user_id),
)
@ -1612,6 +1616,7 @@ async def _user_api_key_auth_builder(
user_rpm_limit=(user_object.rpm_limit if user_object is not None else None),
user_model_max_budget=(user_object.model_max_budget if user_object is not None else None),
jwt_claims=jwt_claims,
agent_id=agent_id,
**team_grants(team_object=team_object, team_membership=team_membership, user_id=user_id),
)
@ -1635,6 +1640,7 @@ async def _user_api_key_auth_builder(
user_id=user_id,
org_id=org_id,
end_user_id=end_user_id,
agent_id=agent_id,
)
if auto_registered is not None:
auto_registered.jwt_claims = jwt_claims

3
litellm/proxy/proxy_server.py
View file

@ -9516,6 +9516,9 @@ class ProxyStartupEvent:
user_api_key_cache=user_api_key_cache,
litellm_jwtauth=litellm_jwtauth,
)
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
jwt_handler.bind_agent_lookup(global_agent_registry)
@classmethod
def _add_proxy_budget_to_db(cls):

179
tests/test_litellm/proxy/auth/test_handle_jwt.py
View file

@ -23,6 +23,7 @@ from litellm.proxy._types import (
ProxyException,
)
from litellm.caching.dual_cache import DualCache
from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry
from litellm.proxy.auth.handle_jwt import (
JWKS_FETCH_ATTEMPTS,
STALE_CACHE_KEY_PREFIX,
@ -32,6 +33,7 @@ from litellm.proxy.auth.handle_jwt import (
JWTHandler,
NoMatchingJWTPublicKeyError,
)
from litellm.types.agents import AgentResponse
@pytest.mark.asyncio
@ -6786,3 +6788,180 @@ async def test_sync_user_role_and_teams_singular_claim_only_recognized_under_fla
}
assert mock_patch.call_args.kwargs["teams_ids_to_add_user_to"] == []
assert user.teams == []
def _entra_agent_registry() -> AgentRegistry:
registry = AgentRegistry()
registry.register_agent(
AgentResponse(
agent_id="canonical-agent-id",
agent_name="2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21",
agent_card_params={"name": "research-agent", "url": "http://localhost:9999/a2a", "version": "1.0.0"},
litellm_params={"require_trace_id_on_calls_by_agent": True},
)
)
return registry
def _entra_agent_jwt_handler(agent_id_jwt_field: str | None) -> JWTHandler:
jwt_handler = JWTHandler()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=DualCache(),
litellm_jwtauth=LiteLLM_JWTAuth(user_id_jwt_field="sub", agent_id_jwt_field=agent_id_jwt_field),
)
return jwt_handler
@pytest.mark.parametrize(
"claim_value",
["canonical-agent-id", "2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21"],
ids=["matches_agent_id", "matches_agent_name"],
)
def test_resolve_agent_id_returns_canonical_agent_id(claim_value: str):
"""An Entra app token's azp claim binds to the registered agent by id or by name and yields its canonical id."""
jwt_handler = _entra_agent_jwt_handler(agent_id_jwt_field="azp")
resolved = JWTAuthManager.resolve_agent_id(
jwt_handler=jwt_handler,
jwt_valid_token={"sub": "sp-object-id-1234", "azp": claim_value},
agent_registry=_entra_agent_registry(),
)
assert resolved == "canonical-agent-id"
def test_resolve_agent_id_reads_nested_claim():
jwt_handler = _entra_agent_jwt_handler(agent_id_jwt_field="entra.client_id")
resolved = JWTAuthManager.resolve_agent_id(
jwt_handler=jwt_handler,
jwt_valid_token={"sub": "sp-object-id-1234", "entra": {"client_id": "canonical-agent-id"}},
agent_registry=_entra_agent_registry(),
)
assert resolved == "canonical-agent-id"
def test_resolve_agent_id_rejects_claim_for_unregistered_agent():
"""A configured agent claim naming no registered agent fails closed with 403 instead of falling back to an unbound identity."""
jwt_handler = _entra_agent_jwt_handler(agent_id_jwt_field="azp")
with pytest.raises(HTTPException) as exc_info:
JWTAuthManager.resolve_agent_id(
jwt_handler=jwt_handler,
jwt_valid_token={"sub": "sp-object-id-1234", "azp": "00000000-0000-0000-0000-000000000000"},
agent_registry=_entra_agent_registry(),
)
assert exc_info.value.status_code == 403
@pytest.mark.parametrize(
"token",
[
{"sub": "sp-object-id-1234"},
{"sub": "sp-object-id-1234", "azp": ""},
{"sub": "sp-object-id-1234", "azp": ["canonical-agent-id"]},
],
ids=["claim_absent", "claim_empty", "claim_not_a_string"],
)
def test_resolve_agent_id_returns_none_when_claim_unusable(token: dict):
jwt_handler = _entra_agent_jwt_handler(agent_id_jwt_field="azp")
assert (
JWTAuthManager.resolve_agent_id(
jwt_handler=jwt_handler, jwt_valid_token=token, agent_registry=_entra_agent_registry()
)
is None
)
def test_resolve_agent_id_ignores_claim_when_field_not_configured():
"""Without agent_id_jwt_field an azp claim (even an unknown one) leaves JWT auth behaviour unchanged."""
jwt_handler = _entra_agent_jwt_handler(agent_id_jwt_field=None)
resolved = JWTAuthManager.resolve_agent_id(
jwt_handler=jwt_handler,
jwt_valid_token={"sub": "sp-object-id-1234", "azp": "00000000-0000-0000-0000-000000000000"},
agent_registry=_entra_agent_registry(),
)
assert resolved is None
def _entra_signed_app_token(monkeypatch, azp: str, scope: str) -> tuple[JWTHandler, str]:
"""A JWTHandler that verifies RS256 tokens against a pre-cached JWKS, plus a signed Entra-style app token."""
jwks_url = "https://login.microsoftonline.test/discovery/v2.0/keys"
monkeypatch.setenv("JWT_PUBLIC_KEY_URL", jwks_url)
monkeypatch.delenv("JWT_AUDIENCE", raising=False)
private_key, jwk = _get_rsa_key_and_jwk(kid="entra-kid")
cache = DualCache()
cache.set_cache(key=f"litellm_jwt_auth_keys_{jwks_url}", value=[jwk])
jwt_handler = JWTHandler()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=cache,
litellm_jwtauth=LiteLLM_JWTAuth(agent_id_jwt_field="azp"),
)
token = _encode_rsa_jwt(
private_key,
issuer="https://login.microsoftonline.test/lit7664-tenant/v2.0",
audience="api://litellm",
kid="entra-kid",
extra_claims={"sub": "sp-object-id-1234", "azp": azp, "scope": scope},
)
return jwt_handler, token
@pytest.mark.asyncio
@pytest.mark.parametrize("is_admin_token", [False, True], ids=["standard_jwt", "proxy_admin_jwt"])
async def test_auth_builder_propagates_agent_id_from_jwt_claim(monkeypatch, is_admin_token: bool):
"""auth_builder carries the resolved agent id into JWTAuthBuilderResult on both the admin and standard paths."""
jwt_handler, token = _entra_signed_app_token(
monkeypatch,
azp="2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21",
scope=LiteLLM_JWTAuth().admin_jwt_scope if is_admin_token else "",
)
jwt_handler.bind_agent_lookup(_entra_agent_registry())
result = await JWTAuthManager.auth_builder(
api_key=token,
jwt_handler=jwt_handler,
request_data={"model": "gpt-5.6"},
general_settings={"enforce_rbac": False},
route="/key/info" if is_admin_token else "/chat/completions",
prisma_client=None,
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
)
assert result["is_proxy_admin"] is is_admin_token
assert result["agent_id"] == "canonical-agent-id"
@pytest.mark.asyncio
async def test_auth_builder_denies_jwt_naming_unregistered_agent_before_admin_check(monkeypatch):
"""An unknown agent claim is rejected even when the token would otherwise be a proxy admin."""
jwt_handler, token = _entra_signed_app_token(
monkeypatch,
azp="00000000-0000-0000-0000-000000000000",
scope=LiteLLM_JWTAuth().admin_jwt_scope,
)
jwt_handler.bind_agent_lookup(_entra_agent_registry())
with pytest.raises(HTTPException) as exc_info:
await JWTAuthManager.auth_builder(
api_key=token,
jwt_handler=jwt_handler,
request_data={"model": "gpt-5.6"},
general_settings={"enforce_rbac": False},
route="/key/info",
prisma_client=None,
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
)
assert exc_info.value.status_code == 403

286
tests/test_litellm/proxy/auth/test_user_api_key_auth.py
View file

@ -1937,6 +1937,76 @@ async def test_standard_jwt_auth_propagates_user_email():
assert result.api_key is None
@pytest.mark.asyncio
@pytest.mark.parametrize("is_proxy_admin", [False, True], ids=["standard_jwt", "proxy_admin_jwt"])
async def test_jwt_auth_propagates_agent_id_to_user_api_key_auth(is_proxy_admin: bool):
"""The agent id resolved by auth_builder must land on UserAPIKeyAuth.agent_id so
agent-scoped checks (trace id requirement, MCP server/tool restrictions, spend
attribution) apply to JWT callers the same way they apply to agent-bound keys."""
jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature"
general_settings = {"enable_jwt_auth": True}
user_api_key_cache = DualCache()
jwt_handler = MagicMock()
jwt_handler.is_jwt.return_value = True
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(agent_id_jwt_field="azp")
user_object = LiteLLM_UserTable(user_id="sp-object-id-1234", user_role="internal_user")
mock_jwt_result = {
"is_proxy_admin": is_proxy_admin,
"team_object": None,
"user_object": user_object,
"end_user_object": None,
"org_object": None,
"token": jwt_token,
"team_id": None,
"user_id": "sp-object-id-1234",
"user_email": None,
"end_user_id": None,
"org_id": None,
"team_membership": None,
"jwt_claims": {"sub": "sp-object-id-1234", "azp": "2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21"},
"agent_id": "canonical-agent-id",
}
mock_request = MagicMock()
mock_request.url.path = "/v1/chat/completions"
mock_request.method = "POST"
mock_request.headers = {"authorization": f"Bearer {jwt_token}"}
mock_request.query_params = {}
mock_request.state = SimpleNamespace()
with (
patch.multiple( # test-quality-ok: production auth reads these module globals; no dependency injection seam exists
"litellm.proxy.proxy_server",
general_settings=general_settings,
premium_user=True,
master_key="sk-master",
prisma_client=None,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=MagicMock(),
jwt_handler=jwt_handler,
),
patch( # test-quality-ok: the builder calls this static method directly; no dependency injection seam exists
"litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder",
new_callable=AsyncMock,
return_value=mock_jwt_result,
),
):
result = await _user_api_key_auth_builder(
request=mock_request,
api_key=jwt_token,
azure_api_key_header="",
anthropic_api_key_header=None,
google_ai_studio_api_key_header=None,
azure_apim_header=None,
request_data={"model": "gpt-5.6"},
)
assert result.agent_id == "canonical-agent-id"
assert result.user_id == "sp-object-id-1234"
assert result.api_key is None
@pytest.mark.asyncio
async def test_auto_register_binds_api_key_to_token_hash():
"""
@ -2106,6 +2176,222 @@ async def test_auto_register_first_request_propagates_user_email():
assert result.api_key == "hashed-auto-key"
@pytest.mark.asyncio
async def test_auto_register_stamps_new_key_with_jwt_agent_id():
"""The virtual key AUTO_REGISTER creates must carry the agent id auth_builder bound
from the JWT claim, and the first request's principal must carry it too, or the
mapped-key path would drop the agent policies on that request and every later one."""
from litellm.proxy.auth.auth_method import AuthMethod
from litellm.proxy.auth.resolvers.models import CredentialRef
from litellm.proxy.auth.resolvers.store import IdentityStore
from litellm.proxy.auth.user_api_key_auth import _auto_register_jwt_mapping
from litellm.proxy.proxy_server import hash_token
plaintext = "sk-auto-registered-agent"
token_hash = hash_token(plaintext)
persisted_principal = IdentityStore._principal_from_key(
UserAPIKeyAuth(token=token_hash, user_id="validated-user", team_id="validated-team", agent_id="canonical-agent-id"),
auth_method=AuthMethod.API_KEY,
credential_ref=CredentialRef(token_id=token_hash),
)
prisma_client = MagicMock()
prisma_client.db.litellm_jwtkeymapping.create = AsyncMock()
user_api_key_cache = MagicMock()
user_api_key_cache.async_set_cache = AsyncMock()
jwt_handler = MagicMock()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(virtual_key_mapping_cache_ttl=300)
generate_key = AsyncMock(return_value={"token": plaintext})
with (
patch( # test-quality-ok: key creation is an inline import inside the helper; no dependency injection seam exists
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
generate_key,
),
patch( # test-quality-ok: the helper constructs IdentityStore itself; no dependency injection seam exists
"litellm.proxy.auth.resolvers.store.IdentityStore.resolve",
new_callable=AsyncMock,
return_value=persisted_principal,
),
):
result = await _auto_register_jwt_mapping(
virtual_key_claim_field="appid",
claim_value="2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21",
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
cache_key="jwt_key_mapping:appid:2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21",
team_id="validated-team",
user_id="validated-user",
agent_id="canonical-agent-id",
)
assert generate_key.await_args is not None
assert generate_key.await_args.kwargs["agent_id"] == "canonical-agent-id"
assert result is not None
assert result.agent_id == "canonical-agent-id"
@pytest.mark.asyncio
@pytest.mark.parametrize("losing_agent_id", ["other-agent", None], ids=["different_agent", "no_agent_claim"])
async def test_auto_register_race_loser_keeps_winners_agent_id(losing_agent_id: str | None):
"""When two requests race to AUTO_REGISTER the same mapping claim, the loser must run as
the persisted key, agent binding included. Every later request on that mapping uses the
winner's key, so stamping the loser's own (or missing) agent id on it would give one request
different agent policies and spend attribution than all the others."""
from litellm.proxy.auth.auth_method import AuthMethod
from litellm.proxy.auth.resolvers.models import CredentialRef
from litellm.proxy.auth.resolvers.store import IdentityStore
from litellm.proxy.auth.user_api_key_auth import _auto_register_jwt_mapping
winner_hash = "winner-key-hash"
winner_principal = IdentityStore._principal_from_key(
UserAPIKeyAuth(token=winner_hash, user_id="validated-user", team_id="validated-team", agent_id="winner-agent"),
auth_method=AuthMethod.API_KEY,
credential_ref=CredentialRef(token_id=winner_hash),
)
prisma_client = MagicMock()
prisma_client.db.litellm_jwtkeymapping.create = AsyncMock(
side_effect=Exception("Unique constraint failed on the fields: (`jwt_claim_name`,`jwt_claim_value`)")
)
prisma_client.db.litellm_verificationtoken.delete = AsyncMock()
user_api_key_cache = MagicMock()
user_api_key_cache.async_set_cache = AsyncMock()
jwt_handler = MagicMock()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(virtual_key_mapping_cache_ttl=300)
with (
patch( # test-quality-ok: key creation is an inline import inside the helper; no dependency injection seam exists
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
new_callable=AsyncMock,
return_value={"token": "sk-orphaned-loser-key"},
),
patch( # test-quality-ok: module-level helper called by the builder; no dependency injection seam exists
"litellm.proxy.auth.user_api_key_auth.get_jwt_key_mapping_object",
new_callable=AsyncMock,
return_value=winner_hash,
),
patch( # test-quality-ok: the helper constructs IdentityStore itself; no dependency injection seam exists
"litellm.proxy.auth.resolvers.store.IdentityStore.resolve",
new_callable=AsyncMock,
return_value=winner_principal,
),
):
result = await _auto_register_jwt_mapping(
virtual_key_claim_field="tid",
claim_value="shared-tenant",
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
cache_key="jwt_key_mapping:tid:shared-tenant",
team_id="validated-team",
user_id="validated-user",
agent_id=losing_agent_id,
)
assert result is not None
assert result.token == winner_hash
assert result.agent_id == "winner-agent"
@pytest.mark.asyncio
async def test_jwt_auto_register_forwards_bound_agent_id():
"""When a JWT under AUTO_REGISTER also carries the configured agent claim, the agent
id auth_builder resolved must reach the key creation, not be dropped when
valid_token is swapped for the freshly registered key."""
jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature"
user_api_key_cache = DualCache()
jwt_handler = MagicMock()
jwt_handler.is_jwt.return_value = True
jwt_handler.auth_jwt = AsyncMock(return_value={"sub": "user1", "appid": "2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21"})
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
virtual_key_claim_field="sub",
virtual_key_mapping_cache_ttl=300,
agent_id_jwt_field="appid",
)
user_object = LiteLLM_UserTable(user_id="validated-user", user_role="internal_user")
mock_jwt_result = {
"is_proxy_admin": False,
"team_object": None,
"user_object": user_object,
"end_user_object": None,
"org_object": None,
"token": jwt_token,
"team_id": "validated-team",
"user_id": "validated-user",
"user_email": None,
"end_user_id": None,
"org_id": None,
"team_membership": None,
"jwt_claims": {"sub": "user1", "appid": "2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21"},
"agent_id": "canonical-agent-id",
}
auto_register = AsyncMock(
return_value=UserAPIKeyAuth(
token="hashed-auto-key",
api_key="hashed-auto-key",
team_id="validated-team",
user_id="validated-user",
agent_id="canonical-agent-id",
)
)
mock_request = MagicMock()
mock_request.url.path = "/v1/chat/completions"
mock_request.method = "POST"
mock_request.headers = {"authorization": f"Bearer {jwt_token}"}
mock_request.query_params = {}
mock_request.state = SimpleNamespace()
with (
patch.multiple( # test-quality-ok: production auth reads these module globals; no dependency injection seam exists
"litellm.proxy.proxy_server",
general_settings={"enable_jwt_auth": True},
premium_user=True,
master_key="sk-master",
prisma_client=MagicMock(),
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=MagicMock(),
jwt_handler=jwt_handler,
),
patch( # test-quality-ok: module-level helper called by the builder; no dependency injection seam exists
"litellm.proxy.auth.user_api_key_auth._resolve_jwt_to_virtual_key",
new_callable=AsyncMock,
return_value=_PendingAutoRegister(
claim_field="sub",
claim_value="user1",
cache_key="jwt_key_mapping:sub:user1",
),
),
patch( # test-quality-ok: the builder calls this static method directly; no dependency injection seam exists
"litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder",
new_callable=AsyncMock,
return_value=mock_jwt_result,
),
patch( # test-quality-ok: module-level helper called by the builder; no dependency injection seam exists
"litellm.proxy.auth.user_api_key_auth._auto_register_jwt_mapping",
auto_register,
),
):
result = await _user_api_key_auth_builder(
request=mock_request,
api_key=jwt_token,
azure_api_key_header="",
anthropic_api_key_header=None,
google_ai_studio_api_key_header=None,
azure_apim_header=None,
request_data={"model": "gpt-5.6"},
)
assert auto_register.await_args is not None
assert auto_register.await_args.kwargs["agent_id"] == "canonical-agent-id"
assert result.agent_id == "canonical-agent-id"
assert result.api_key == "hashed-auto-key"
class TestJWTOAuth2Coexistence:
"""
Test that JWT and OAuth2 auth can coexist on the same instance.

49
tests/test_litellm/proxy/proxy_server/test_proxy_config.py
View file

@ -3764,6 +3764,55 @@ async def test_ProxyConfig__init_agents_in_db_keeps_config_defined_agents(clean_
]
@pytest.mark.asyncio
@pytest.mark.parametrize("agents_source", ["config", "db", "api"])
async def test_ProxyStartupEvent_jwt_auth_resolves_agent_claims_against_live_registry(
clean_agent_registry, agents_source
):
"""A JWT agent claim must resolve against every agent the proxy knows, including ones created after startup."""
from litellm.proxy import proxy_server
from litellm.proxy._types import LiteLLM_JWTAuth
from litellm.proxy.auth.handle_jwt import JWTAuthManager
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.types.agents import AgentResponse
original_lookup = proxy_server.jwt_handler.agent_lookup
try:
proxy_server.ProxyStartupEvent._initialize_jwt_auth(
general_settings={"litellm_jwtauth": {"agent_id_jwt_field": "appid"}},
prisma_client=None,
user_api_key_cache=UserApiKeyCache(),
)
if agents_source == "config":
await ProxyConfig()._init_non_llm_configs(
config={"agents": [_config_agent("loaded-agent")]},
config_file_path=None,
)
elif agents_source == "db":
prisma_client = MagicMock()
prisma_client.db.litellm_agentstable.find_many = AsyncMock(
return_value=[_FakeAgentRow("db-id", "loaded-agent")]
)
await ProxyConfig()._init_agents_in_db(prisma_client=prisma_client)
else:
clean_agent_registry.register_agent(
agent_config=AgentResponse(agent_id="api-id", **_config_agent("loaded-agent"))
)
resolved = JWTAuthManager.resolve_agent_id(
jwt_handler=proxy_server.jwt_handler,
jwt_valid_token={"appid": "loaded-agent"},
agent_registry=proxy_server.jwt_handler.agent_lookup,
)
finally:
proxy_server.jwt_handler.bind_agent_lookup(original_lookup)
proxy_server.jwt_handler.update_environment(
prisma_client=None, user_api_key_cache=UserApiKeyCache(), litellm_jwtauth=LiteLLM_JWTAuth()
)
assert resolved == clean_agent_registry.get_agent_by_name(agent_name="loaded-agent").agent_id
@pytest.mark.asyncio
@pytest.mark.parametrize(
"config, expected_agent_names",