From 22e9064673539dd612e929af018a969707cb453b Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 20 May 2026 16:18:55 +0000 Subject: [PATCH] fix(mcp,jwt): address greptile concurrency and code-quality concerns - _apply_issuer_claim_mappings now builds a new dict and reads from the original token, rather than mutating its input. The change is behaviour-preserving (caller passes a fresh jwt.decode result), but avoids the surprise-mutation pattern flagged by greptile. - is_network_error uses isinstance(exc, httpx.TransportError) instead of matching type(exc).__name__ against a hand-maintained string set, so ReadError / WriteError / ProxyError / etc. are also treated as transport-level failures and surfaced as HTTP 502. - fetch_upstream_oauth_protected_resource now coalesces concurrent discovery requests per (server_id, resource_url) through an asyncio.Lock so concurrent .well-known calls share a single upstream fetch + cache write. - Drop the redundant 'if trusted_ranges:' branch in get_mcp_client_ip; it is always true on the path that reaches it (the prior 'if not trusted_ranges:' early-returns). Co-authored-by: Claude --- .../mcp_server/discoverable_endpoints.py | 91 ++++++++++--------- litellm/proxy/auth/handle_jwt.py | 15 ++- litellm/proxy/auth/ip_address_utils.py | 21 ++--- 3 files changed, 66 insertions(+), 61 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index b648f5a9dc7..b478fecaa84 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -1,8 +1,10 @@ +import asyncio import json import time from typing import Any, Dict, Optional, Tuple from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse +import httpx from fastapi import APIRouter, Form, HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse @@ -32,6 +34,9 @@ from litellm.types.mcp_server.mcp_server_manager import MCPServer _OAUTH_METADATA_CACHE: Dict[Tuple[str, str], Tuple[float, dict]] = {} _OAUTH_METADATA_CACHE_TTL_SECONDS = 300 _OAUTH_METADATA_CACHE_MAX_SIZE = 128 +# Per-(server_id, resource_url) async locks so concurrent discovery requests +# coalesce onto a single upstream fetch instead of issuing N parallel calls. +_OAUTH_METADATA_FETCH_LOCKS: Dict[Tuple[str, str], asyncio.Lock] = {} router = APIRouter( tags=["mcp"], @@ -746,59 +751,61 @@ async def fetch_upstream_oauth_protected_resource( if cached is not None and cached[0] > now: return cached[1] - host_base = f"{upstream.scheme}://{upstream.netloc}" - candidates = [f"{host_base}/.well-known/oauth-protected-resource"] - # RFC 9728 §3.1 path fallback - if upstream.path and upstream.path not in ("", "/"): - candidates.append( - f"{host_base}/.well-known/oauth-protected-resource" - f"{upstream.path.rstrip('/')}" + lock = _OAUTH_METADATA_FETCH_LOCKS.setdefault(cache_key, asyncio.Lock()) + async with lock: + now = time.time() + cached = _OAUTH_METADATA_CACHE.get(cache_key) + if cached is not None and cached[0] > now: + return cached[1] + + host_base = f"{upstream.scheme}://{upstream.netloc}" + candidates = [f"{host_base}/.well-known/oauth-protected-resource"] + # RFC 9728 §3.1 path fallback + if upstream.path and upstream.path not in ("", "/"): + candidates.append( + f"{host_base}/.well-known/oauth-protected-resource" + f"{upstream.path.rstrip('/')}" + ) + + async_client = get_async_httpx_client( + llm_provider=httpxSpecialProvider.Oauth2Check ) - async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) - - network_errors: list[Exception] = [] - for candidate in candidates: - try: - response = await async_client.get( - candidate, - headers={"Accept": "application/json"}, - ) - except Exception as exc: # network / connect errors - if is_network_error(exc): - network_errors.append(exc) - continue - if response.status_code == 200: + network_errors: list[Exception] = [] + for candidate in candidates: try: - payload = response.json() - except Exception: - continue - if isinstance(payload, dict): - now = time.time() - _OAUTH_METADATA_CACHE[cache_key] = ( - now + _OAUTH_METADATA_CACHE_TTL_SECONDS, - payload, + response = await async_client.get( + candidate, + headers={"Accept": "application/json"}, ) - _prune_oauth_metadata_cache(now) - return payload + except Exception as exc: # network / connect errors + if is_network_error(exc): + network_errors.append(exc) + continue + if response.status_code == 200: + try: + payload = response.json() + except Exception: + continue + if isinstance(payload, dict): + now = time.time() + _OAUTH_METADATA_CACHE[cache_key] = ( + now + _OAUTH_METADATA_CACHE_TTL_SECONDS, + payload, + ) + _prune_oauth_metadata_cache(now) + return payload - if len(network_errors) == len(candidates): - raise network_errors[-1] + if len(network_errors) == len(candidates): + raise network_errors[-1] - return None + return None def is_network_error(exc: Exception) -> bool: """True for transport-layer failures (connection refused, DNS, TLS, timeout) as opposed to HTTP protocol errors (4xx/5xx with a valid response).""" - name = type(exc).__name__ - return name in { - "ConnectError", - "ConnectTimeout", - "ReadTimeout", - "PoolTimeout", - "RemoteProtocolError", - } + return isinstance(exc, httpx.TransportError) async def _build_oauth_protected_resource_response( diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 685c00cbf16..b7c18eb2921 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -922,11 +922,10 @@ class JWTHandler: def _apply_issuer_claim_mappings( self, token: dict, issuer_config: JWTIssuerConfig ) -> dict: - source_token = {**token} - for claim in self.LITELLM_INTERNAL_CLAIMS: - token.pop(claim, None) - - token[self.LITELLM_JWT_ISSUER_CLAIM] = issuer_config.issuer + normalized: dict = { + k: v for k, v in token.items() if k not in self.LITELLM_INTERNAL_CLAIMS + } + normalized[self.LITELLM_JWT_ISSUER_CLAIM] = issuer_config.issuer claim_mappings = [ (issuer_config.user_id_jwt_field, self.LITELLM_USER_ID_CLAIM), (issuer_config.user_email_jwt_field, self.LITELLM_USER_EMAIL_CLAIM), @@ -940,13 +939,13 @@ class JWTHandler: if source_claim is None: continue claim_value = self._get_claim_value_for_issuer_mapping( - token=source_token, + token=token, claim_field=source_claim, ) if claim_value is not None: - token[normalized_claim] = claim_value + normalized[normalized_claim] = claim_value - return token + return normalized def _get_jwk_from_public_key(self, public_key: dict) -> dict: jwk = {} diff --git a/litellm/proxy/auth/ip_address_utils.py b/litellm/proxy/auth/ip_address_utils.py index 16a5bc61d31..0218bc2211a 100644 --- a/litellm/proxy/auth/ip_address_utils.py +++ b/litellm/proxy/auth/ip_address_utils.py @@ -206,16 +206,15 @@ class IPAddressUtils: request, general_settings=general_settings ) return _get_request_ip_address(request, use_x_forwarded_for=False) - if trusted_ranges: - # Validate direct connection is from trusted proxy - direct_ip = request.client.host if request.client else None - trusted_networks = IPAddressUtils.parse_trusted_proxy_networks( - trusted_ranges + # Validate direct connection is from trusted proxy + direct_ip = request.client.host if request.client else None + trusted_networks = IPAddressUtils.parse_trusted_proxy_networks( + trusted_ranges + ) + if not IPAddressUtils.is_trusted_proxy(direct_ip, trusted_networks): + # Untrusted source trying to set XFF - ignore XFF, use direct IP + verbose_proxy_logger.warning( + "XFF header from untrusted IP %s, ignoring", direct_ip ) - if not IPAddressUtils.is_trusted_proxy(direct_ip, trusted_networks): - # Untrusted source trying to set XFF - ignore XFF, use direct IP - verbose_proxy_logger.warning( - "XFF header from untrusted IP %s, ignoring", direct_ip - ) - return direct_ip + return direct_ip return _get_request_ip_address(request, use_x_forwarded_for=use_xff)