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 102

fix(proxy): bind agent registry into JWTHandler and keep persisted agent id on AUTO_REGISTER race

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-09-14 21:09:14 +00:00
parent 4435aa601d
commit 4d4d3fb18a
6 changed files with 128 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -131,15 +131,21 @@ class _UserInfoResponse(Protocol):
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: ...
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: ...
def get_agent_by_name(self, agent_name: str) -> AgentResponse | None:
"""The agent registered under ``agent_name``, if any."""
def _global_agent_lookup() -> AgentLookup:
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
class _NoRegisteredAgents:
"""The lookup in force until the proxy binds its agent registry: no agent is registered, so no claim matches."""
return global_agent_registry
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:
@ -213,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,
@ -2251,7 +2261,6 @@ class JWTAuthManager:
proxy_logging_obj: ProxyLogging,
request_headers: dict | None = None,
request_method: str | None = None,
agent_registry: AgentLookup | None = None,
) -> JWTAuthBuilderResult:
"""Main authentication and authorization builder"""
# Check if OIDC UserInfo endpoint is enabled, but fall back to standard
@ -2314,7 +2323,7 @@ class JWTAuthManager:
agent_id: Final = JWTAuthManager.resolve_agent_id(
jwt_handler=jwt_handler,
jwt_valid_token=jwt_valid_token,
agent_registry=agent_registry if agent_registry is not None else _global_agent_lookup(),
agent_registry=jwt_handler.agent_lookup,
)
# Check admin access

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

@ -971,7 +971,6 @@ async def _auto_register_jwt_mapping(
if auto_registered_key is not None:
auto_registered_key.org_id = org_id
auto_registered_key.end_user_id = end_user_id
auto_registered_key.agent_id = agent_id
auto_registered_key.api_key = auto_registered_key.token
return auto_registered_key

2
litellm/proxy/proxy_server.py
View file

@ -6217,6 +6217,7 @@ class ProxyConfig:
)
global_agent_registry.load_agents_from_config(agent_config)
jwt_handler.bind_agent_lookup(global_agent_registry)
mcp_servers_config: Final = config.get("mcp_servers", None)
if mcp_servers_config:
@ -8182,6 +8183,7 @@ class ProxyConfig:
global_agent_registry as AGENT_REGISTRY,
)
jwt_handler.bind_agent_lookup(AGENT_REGISTRY)
try:
async with AGENT_RECONCILE_LOCK:
db_agents: Final = await AGENT_REGISTRY.get_all_agents_from_db(prisma_client=prisma_client)

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

@ -6923,6 +6923,7 @@ async def test_auth_builder_propagates_agent_id_from_jwt_claim(monkeypatch, is_a
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,
@ -6934,7 +6935,6 @@ async def test_auth_builder_propagates_agent_id_from_jwt_claim(monkeypatch, is_a
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
agent_registry=_entra_agent_registry(),
)
assert result["is_proxy_admin"] is is_admin_token
@ -6949,6 +6949,7 @@ async def test_auth_builder_denies_jwt_naming_unregistered_agent_before_admin_ch
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(
@ -6961,7 +6962,6 @@ async def test_auth_builder_denies_jwt_naming_unregistered_agent_before_admin_ch
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
agent_registry=_entra_agent_registry(),
)
assert exc_info.value.status_code == 403

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

@ -2189,8 +2189,8 @@ async def test_auto_register_stamps_new_key_with_jwt_agent_id():
plaintext = "sk-auto-registered-agent"
token_hash = hash_token(plaintext)
principal = IdentityStore._principal_from_key(
UserAPIKeyAuth(token=token_hash, user_id="validated-user", team_id="validated-team"),
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),
)
@ -2210,7 +2210,7 @@ async def test_auto_register_stamps_new_key_with_jwt_agent_id():
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=principal,
return_value=persisted_principal,
),
):
result = await _auto_register_jwt_mapping(
@ -2233,6 +2233,70 @@ async def test_auto_register_stamps_new_key_with_jwt_agent_id():
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

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

@ -3740,6 +3740,47 @@ async def test_ProxyConfig__init_agents_in_db_keeps_config_defined_agents(clean_
]
@pytest.mark.asyncio
@pytest.mark.parametrize("agents_source", ["config", "db"])
async def test_ProxyConfig_agent_loading_binds_registry_to_jwt_agent_claims(clean_agent_registry, agents_source):
"""A JWT agent claim must resolve against the agents the proxy loaded, whichever source registered them."""
from litellm.proxy import proxy_server
from litellm.proxy._types import LiteLLM_JWTAuth
from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
jwt_handler = JWTHandler()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=UserApiKeyCache(),
litellm_jwtauth=LiteLLM_JWTAuth(agent_id_jwt_field="appid"),
)
original_lookup = proxy_server.jwt_handler.agent_lookup
proxy_server.jwt_handler.bind_agent_lookup(jwt_handler.agent_lookup)
try:
if agents_source == "config":
await ProxyConfig()._init_non_llm_configs(
config={"agents": [_config_agent("loaded-agent")]},
config_file_path=None,
)
else:
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)
resolved = JWTAuthManager.resolve_agent_id(
jwt_handler=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)
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",