From 1b0b75ccd7eec00d5b62b363741e8bc4243d80c4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 21 May 2026 19:14:18 +0000 Subject: [PATCH] fix(jwt/mcp): warn on unscoped JWT fallback; route agent permission lookup through shared helper - _build_decode_kwargs no longer suppresses the unscoped-fallback warning when LiteLLM_JWTAuth.issuers is set: tokens whose iss does not match any configured issuer still fall through to the global path, and that fallback is itself unscoped when JWT_AUDIENCE/JWT_ISSUER are absent. - _get_agent_object_permission now caches the agent_id -> object_permission_id mapping and delegates the permission lookup to the shared get_object_permission helper, so the agent path reuses the same cache entries as the org / team / key paths. --- .../mcp_server/auth/user_api_key_auth_mcp.py | 74 +++++++++-------- litellm/proxy/auth/handle_jwt.py | 17 ++-- .../auth/test_user_api_key_auth_mcp.py | 83 +++++++++++++++++++ .../proxy/auth/test_handle_jwt.py | 14 ++-- 4 files changed, 135 insertions(+), 53 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 16f8fe3fb74..aa11ae59328 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -1289,12 +1289,17 @@ class MCPRequestHandler: user_api_key_auth: Optional[UserAPIKeyAuth] = None, ): """ - Get agent object_permission, using user_api_key_cache to avoid DB hits on every request. - - Caches both positive results and the absence of an object_permission so that agents - with no MCP permissions configured do not trigger a DB query on every request. + Get agent object_permission via the established ``get_object_permission`` + helper. Caches the ``agent_id -> object_permission_id`` mapping so we + avoid re-reading the agent row on every request, and reuses the shared + ``object_permission_id`` cache populated by the org / team / key paths. """ - from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + from litellm.proxy.auth.auth_checks import get_object_permission + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) if not user_api_key_auth or not user_api_key_auth.agent_id: return None @@ -1304,42 +1309,41 @@ class MCPRequestHandler: return None agent_id = user_api_key_auth.agent_id - cache_key = f"agent_object_permission:{agent_id}" - - from litellm.proxy._types import LiteLLM_ObjectPermissionTable + cache_key = f"agent_object_permission_id:{agent_id}" try: - cached = await user_api_key_cache.async_get_cache(key=cache_key) - if cached is not None: - if cached == MCPRequestHandler._AGENT_NO_PERMISSION_SENTINEL: - return None - # Redis deserialises to a plain dict; reconstruct the Pydantic model - # so callers can access .mcp_servers / .mcp_tool_permissions as attrs. - if isinstance(cached, dict): - return LiteLLM_ObjectPermissionTable(**cached) - return cached - - agent_row = await prisma_client.db.litellm_agentstable.find_unique( - where={"agent_id": agent_id}, - include={"object_permission": True}, + object_permission_id: Optional[str] = ( + await user_api_key_cache.async_get_cache(key=cache_key) ) - if agent_row is None or agent_row.object_permission is None: - await user_api_key_cache.async_set_cache( - key=cache_key, - value=MCPRequestHandler._AGENT_NO_PERMISSION_SENTINEL, - ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, - ) + + if object_permission_id == MCPRequestHandler._AGENT_NO_PERMISSION_SENTINEL: return None - obj_perm = LiteLLM_ObjectPermissionTable( - **agent_row.object_permission.dict() + if object_permission_id is None: + agent_row = await prisma_client.db.litellm_agentstable.find_unique( + where={"agent_id": agent_id}, + ) + object_permission_id = ( + getattr(agent_row, "object_permission_id", None) + if agent_row is not None + else None + ) + await user_api_key_cache.async_set_cache( + key=cache_key, + value=object_permission_id + or MCPRequestHandler._AGENT_NO_PERMISSION_SENTINEL, + ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, + ) + if not object_permission_id: + return None + + return await get_object_permission( + object_permission_id=object_permission_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_auth.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, ) - await user_api_key_cache.async_set_cache( - key=cache_key, - value=obj_perm.dict(), - ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, - ) - return obj_perm except Exception as e: verbose_logger.warning(f"Failed to get agent object permission: {str(e)}") return None diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 314903825ba..582acab68d8 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -843,7 +843,7 @@ class JWTHandler: _unscoped_jwt_warning_emitted = False @classmethod - def _build_decode_kwargs(cls, has_issuer_config: bool = False) -> dict: + def _build_decode_kwargs(cls) -> dict: """Build the audience/issuer/options kwargs for ``jwt.decode``. Setting ``JWT_AUDIENCE`` (and optionally ``JWT_ISSUER``) turns on the @@ -852,10 +852,10 @@ class JWTHandler: When both are unset PyJWT only checks the signature and expiry, which is preserved for backward compatibility but logged once as a warning. - ``has_issuer_config`` suppresses the warning when the caller has - configured per-issuer scoping via ``LiteLLM_JWTAuth.issuers``: this - global path is then just the fallback for tokens whose ``iss`` did not - match any configured issuer, not the only scoping mechanism. + The warning fires even in mixed deployments that also configure + ``LiteLLM_JWTAuth.issuers``: tokens whose ``iss`` does not match any + configured issuer fall through to this global path, and if env-var + scoping is absent that fallback is itself unscoped. """ audience = os.getenv("JWT_AUDIENCE") issuer = os.getenv("JWT_ISSUER") @@ -863,7 +863,6 @@ class JWTHandler: if ( audience is None and issuer is None - and not has_issuer_config and not cls._unscoped_jwt_warning_emitted ): verbose_proxy_logger.warning( @@ -1057,11 +1056,7 @@ class JWTHandler: kid=kid, ) - decode_kwargs = self._build_decode_kwargs( - has_issuer_config=bool( - getattr(getattr(self, "litellm_jwtauth", None), "issuers", None) - ) - ) + decode_kwargs = self._build_decode_kwargs() public_key = await self.get_public_key(kid=kid) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 984052e9b16..5ad915a670e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -3202,6 +3202,89 @@ class TestAgentMCPPermissions: ) assert sorted(result) == ["tool_a", "tool_b"] + async def test_get_agent_object_permission_uses_shared_helper(self): + """``_get_agent_object_permission`` must resolve the agent's + ``object_permission_id`` and then defer to the shared + ``get_object_permission`` helper so cache entries are shared with the + org / team / key paths.""" + from litellm.caching.dual_cache import DualCache + + cache = DualCache() + agent_row = MagicMock() + agent_row.object_permission_id = "perm-xyz" + prisma_client = MagicMock() + prisma_client.db.litellm_agentstable.find_unique = AsyncMock( + return_value=agent_row + ) + user_api_key_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + agent_id="agent-shared", + ) + expected_perm = MagicMock() + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma_client), + patch("litellm.proxy.proxy_server.user_api_key_cache", cache), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + patch( + "litellm.proxy.auth.auth_checks.get_object_permission", + new_callable=AsyncMock, + return_value=expected_perm, + ) as mock_get_perm, + ): + result = await MCPRequestHandler._get_agent_object_permission( + user_api_key_auth + ) + assert result is expected_perm + mock_get_perm.assert_awaited_once() + assert mock_get_perm.await_args.kwargs["object_permission_id"] == "perm-xyz" + + # Second call: the agent_id -> object_permission_id mapping is + # cached, so the agent row is not re-fetched. + prisma_client.db.litellm_agentstable.find_unique.reset_mock() + await MCPRequestHandler._get_agent_object_permission(user_api_key_auth) + prisma_client.db.litellm_agentstable.find_unique.assert_not_called() + + async def test_get_agent_object_permission_caches_missing_permission(self): + """When the agent has no ``object_permission_id`` the sentinel must be + cached so subsequent requests do not hit the DB again.""" + from litellm.caching.dual_cache import DualCache + + cache = DualCache() + agent_row = MagicMock() + agent_row.object_permission_id = None + prisma_client = MagicMock() + prisma_client.db.litellm_agentstable.find_unique = AsyncMock( + return_value=agent_row + ) + user_api_key_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + agent_id="agent-no-perm", + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma_client), + patch("litellm.proxy.proxy_server.user_api_key_cache", cache), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + patch( + "litellm.proxy.auth.auth_checks.get_object_permission", + new_callable=AsyncMock, + ) as mock_get_perm, + ): + assert ( + await MCPRequestHandler._get_agent_object_permission(user_api_key_auth) + is None + ) + assert ( + await MCPRequestHandler._get_agent_object_permission(user_api_key_auth) + is None + ) + + mock_get_perm.assert_not_awaited() + prisma_client.db.litellm_agentstable.find_unique.assert_awaited_once() + @pytest.mark.asyncio async def test_tool_permission_servers_included_in_allowed_servers(): diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index 4becac13ccf..d4c9074d92b 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -3341,24 +3341,24 @@ async def test_multi_issuer_jwt_does_not_emit_unscoped_global_warning( assert JWTHandler._unscoped_jwt_warning_emitted is False -def test_build_decode_kwargs_no_warning_when_issuer_config_scoped( +def test_build_decode_kwargs_warns_for_unscoped_global_fallback_in_mixed_deployment( monkeypatch, _reset_unscoped_warning_flag, caplog ): - """When per-issuer config (``LiteLLM_JWTAuth.issuers``) scopes the proxy, - the global path's unscoped-fallback warning must not fire — admins who - intentionally use issuer config without env-var scoping otherwise see a - misleading warning implying their JWT auth is unscoped.""" + """The unscoped-fallback warning must fire even when per-issuer configs + are set. In mixed deployments, tokens whose ``iss`` does not match any + configured issuer fall through to the global path; if env-var scoping is + absent that fallback IS unscoped, and the operator needs to be told.""" import logging monkeypatch.delenv("JWT_AUDIENCE", raising=False) monkeypatch.delenv("JWT_ISSUER", raising=False) caplog.set_level(logging.WARNING) - JWTHandler._build_decode_kwargs(has_issuer_config=True) + JWTHandler._build_decode_kwargs() matching = [ r for r in caplog.records if "neither JWT_AUDIENCE nor JWT_ISSUER" in r.getMessage() ] - assert matching == [] + assert len(matching) == 1