From 806b8be4512e41dc29690a6b54c3535e3e168be5 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 20 May 2026 15:08:54 +0000 Subject: [PATCH] fix(mcp,jwt): address greptile review concerns - Cache _get_agent_object_permission via user_api_key_cache (sentinel for no-permission rows) so MCP requests from agent keys don't hit the DB on every tool-list / tool-call. - Re-raise HTTPException in handle_sse_mcp so 401 + WWW-Authenticate challenges (and other HTTP errors) propagate to SSE clients instead of being swallowed as 500. - Normalise booleans in _validate_token_response so admin rules written as JSON-style "true" / "false" match upstream responses that return Python True / False. - Treat configured JWT issuer claim mappings as advisory: when a mapped field is absent or empty, leave the normalised claim unset instead of raising, matching the global litellm_jwtauth path. Co-authored-by: Claude --- .../mcp_server/auth/user_api_key_auth_mcp.py | 40 ++++++++++++++++--- .../mcp_server/discoverable_endpoints.py | 19 ++++++++- .../proxy/_experimental/mcp_server/server.py | 6 ++- litellm/proxy/auth/handle_jwt.py | 19 +++++---- tests/mcp_tests/test_per_user_oauth_cache.py | 25 ++++++++++++ tests/proxy_unit_tests/test_jwt.py | 15 ++++--- 6 files changed, 103 insertions(+), 21 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 c1bb6edd772..95f36989dd7 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 @@ -1128,16 +1128,21 @@ class MCPRequestHandler: ) return [] + # Sentinel stored in cache when an agent has no object_permission, so we + # don't re-query the DB on every MCP request for that agent. + _AGENT_NO_PERMISSION_SENTINEL = "__agent_no_mcp_permission__" + @staticmethod async def _get_agent_object_permission( user_api_key_auth: Optional[UserAPIKeyAuth] = None, ): """ - Fetch the agent's object_permission from the DB (single query). + Get agent object_permission, using user_api_key_cache to avoid DB hits on every request. - Returns the object_permission object or None. + 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. """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache if not user_api_key_auth or not user_api_key_auth.agent_id: return None @@ -1146,15 +1151,40 @@ class MCPRequestHandler: verbose_logger.debug("prisma_client is None") 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 + 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": user_api_key_auth.agent_id}, + where={"agent_id": agent_id}, include={"object_permission": True}, ) 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, + ) return None - return agent_row.object_permission + obj_perm = LiteLLM_ObjectPermissionTable( + **agent_row.object_permission.dict() + ) + await user_api_key_cache.async_set_cache( + key=cache_key, value=obj_perm.dict() + ) + 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/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index fc8930fd81f..ffc25d5374e 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -141,6 +141,17 @@ def _resolve_oauth2_server_for_root_endpoints( return None +def _normalize_for_token_comparison(value: Any) -> str: + """Stringify ``value`` for token-rule comparison. + + Booleans are lower-cased so Python's ``True`` / ``False`` line up with + JSON-style ``"true"`` / ``"false"`` rules from admin config. + """ + if isinstance(value, bool): + return "true" if value else "false" + return str(value) + + def _validate_token_response( token_response: Dict[str, Any], validation_rules: Dict[str, Any], @@ -152,7 +163,9 @@ def _validate_token_response( ``token_response["team"]["enterprise_id"]``). Top-level keys are tried first, then dot-split traversal. All comparisons are string-coerced so that numeric values in the response (e.g. ``"org_id": 12345``) match string rules - (``"org_id": "12345"``). + (``"org_id": "12345"``). Booleans are normalised to JSON-style ``"true"`` / + ``"false"`` so admin rules written as ``{"verified": "true"}`` match upstream + responses of ``{"verified": true}``. """ for key, expected in validation_rules.items(): actual: Any = token_response.get(key) @@ -179,7 +192,9 @@ def _validate_token_response( ), }, ) - if str(actual) != str(expected): + if _normalize_for_token_comparison(actual) != _normalize_for_token_comparison( + expected + ): raise HTTPException( status_code=403, detail={ diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 93354ad2153..a8b98b3e2e8 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -3045,9 +3045,13 @@ if MCP_AVAILABLE: _sse_client_ip, ): await sse_session_manager.handle_request(scope, receive, send) + except HTTPException: + # Re-raise HTTP exceptions to preserve status codes and details + # (e.g. 401 + WWW-Authenticate challenges from OAuth pass-through). + raise except Exception as e: verbose_logger.exception(f"Error handling MCP request: {e}") - # Instead of re-raising, try to send a graceful error response + # Try to send a graceful error response for non-HTTP exceptions try: # Send a proper HTTP error response instead of letting the exception bubble up from starlette.responses import JSONResponse diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index f47a307cf72..800df5a8861 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -867,9 +867,13 @@ class JWTHandler: # its jwks_uri, matching JWTIssuerConfig.jwks_url's documented fallback. return f"{issuer_config.issuer.rstrip('/')}/.well-known/openid-configuration" - def _get_claim_value_for_issuer_mapping( - self, token: dict, claim_field: str, issuer: str - ) -> Any: + def _get_claim_value_for_issuer_mapping(self, token: dict, claim_field: str) -> Any: + """Resolve a mapped claim from ``token``. + + Returns ``None`` when the field is absent or empty so that mapped claims + behave like the global ``litellm_jwtauth`` path — present claims override + the normalised value, missing ones simply leave it ``None``. + """ sentinel = object() claim_value = get_nested_value( data=token, @@ -877,9 +881,7 @@ class JWTHandler: default=sentinel, ) if claim_value is sentinel or claim_value is None or claim_value == "": - raise Exception( - f"JWT issuer {issuer} missing required mapped claim: {claim_field}" - ) + return None return claim_value def _apply_issuer_claim_mappings( @@ -902,11 +904,12 @@ class JWTHandler: for source_claim, normalized_claim in claim_mappings: if source_claim is None: continue - token[normalized_claim] = self._get_claim_value_for_issuer_mapping( + claim_value = self._get_claim_value_for_issuer_mapping( token=source_token, claim_field=source_claim, - issuer=issuer_config.issuer, ) + if claim_value is not None: + token[normalized_claim] = claim_value return token diff --git a/tests/mcp_tests/test_per_user_oauth_cache.py b/tests/mcp_tests/test_per_user_oauth_cache.py index 43e514b32ae..141b906fce9 100644 --- a/tests/mcp_tests/test_per_user_oauth_cache.py +++ b/tests/mcp_tests/test_per_user_oauth_cache.py @@ -183,6 +183,31 @@ class TestValidateTokenResponse: server_id="atlassian", ) + def test_boolean_value_matches_lowercase_string_rule(self): + """Boolean ``True`` in token response must match the JSON-style rule ``"true"``. + + Admin config is typically written as ``{"verified": "true"}`` (lower-case + from JSON / YAML), but the OAuth response returns ``{"verified": true}`` + (Python ``True``). The normaliser must align them. + """ + _validate_token_response = _import_validate() + token_response = {"access_token": "tok", "verified": True} + # Should not raise + _validate_token_response( + token_response=token_response, + validation_rules={"verified": "true"}, + server_id="test", + ) + + def test_boolean_false_matches_lowercase_string_rule(self): + _validate_token_response = _import_validate() + token_response = {"access_token": "tok", "is_admin": False} + _validate_token_response( + token_response=token_response, + validation_rules={"is_admin": "false"}, + server_id="test", + ) + # ── _compute_per_user_token_ttl ────────────────────────────────────────────── diff --git a/tests/proxy_unit_tests/test_jwt.py b/tests/proxy_unit_tests/test_jwt.py index e25a5d86513..c825bd0aab7 100644 --- a/tests/proxy_unit_tests/test_jwt.py +++ b/tests/proxy_unit_tests/test_jwt.py @@ -1861,7 +1861,13 @@ async def test_multi_issuer_jwt_same_kid_does_not_cross_issuer_keys(monkeypatch) @pytest.mark.asyncio -async def test_multi_issuer_jwt_missing_mapped_claim_fails_closed(monkeypatch): +async def test_multi_issuer_jwt_missing_mapped_claim_is_optional(monkeypatch): + """Configured issuer claim mappings are advisory, not mandatory. + + When the token simply omits a mapped field (e.g. a service-to-service token + with no ``email`` claim), JWT auth still succeeds and the normalized claim + is just absent — matching the global ``litellm_jwtauth`` behaviour. + """ monkeypatch.delenv("JWT_AUDIENCE", raising=False) monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False) @@ -1886,11 +1892,10 @@ async def test_multi_issuer_jwt_missing_mapped_claim_fails_closed(monkeypatch): kid="issuer-key", ) - with pytest.raises(Exception) as exc: - await jwt_handler.auth_jwt(token=token) + claims = await jwt_handler.auth_jwt(token=token) - assert "missing required mapped claim: email" in str(exc.value) - assert "Validation fails" not in str(exc.value) + assert claims[JWTHandler.LITELLM_JWT_ISSUER_CLAIM] == issuer + assert JWTHandler.LITELLM_USER_ID_CLAIM not in claims @pytest.mark.asyncio