From 4d4d3fb18a28bb071089b163835551f90cbfa360 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 14 Sep 2026 21:09:14 +0000 Subject: [PATCH] fix(proxy): bind agent registry into JWTHandler and keep persisted agent id on AUTO_REGISTER race Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/handle_jwt.py | 23 ++++-- litellm/proxy/auth/user_api_key_auth.py | 1 - litellm/proxy/proxy_server.py | 2 + .../proxy/auth/test_handle_jwt.py | 4 +- .../proxy/auth/test_user_api_key_auth.py | 70 ++++++++++++++++++- .../proxy/proxy_server/test_proxy_config.py | 41 +++++++++++ 6 files changed, 128 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index fcbcf35dba9..94ca3047f45 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -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 diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 4d428fc6eb8..1ef0c7abd80 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -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 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d9f8e04ebda..919357498af 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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) diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index 2fe8729b78e..814e31535e0 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -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 diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 92c87df5060..866ea0b20e4 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -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 diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index e109b650da7..0a93e607313 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -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",