scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-26 01:12:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 99

fix(mcp,jwt): address greptile concurrency and code-quality concerns

Browse source 
- _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 <claude@anthropic.com>
This commit is contained in:
Claude 2026-05-20 16:18:55 +00:00
parent 087a4da116
commit 22e9064673
No known key found for this signature in database
3 changed files with 66 additions and 61 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

91
litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
View file

@ -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(

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

@ -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 = {}

21
litellm/proxy/auth/ip_address_utils.py
View file

@ -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)