diff --git a/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py b/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py index d4b09770750..b4d3157b17e 100644 --- a/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py @@ -39,6 +39,8 @@ _byok_auth_codes: Dict[str, dict] = {} # Authorization codes expire after 5 minutes. _AUTH_CODE_TTL_SECONDS = 300 +# Hard cap to prevent memory exhaustion from incomplete OAuth flows. +_AUTH_CODES_MAX_SIZE = 1000 router = APIRouter(tags=["mcp"]) @@ -648,6 +650,11 @@ async def byok_authorize_post( if parsed_uri.scheme not in ("http", "https"): raise HTTPException(status_code=400, detail="Invalid redirect_uri scheme") + # Reject new codes if the store is at capacity (prevents memory exhaustion + # from a burst of abandoned OAuth flows). + if len(_byok_auth_codes) >= _AUTH_CODES_MAX_SIZE: + raise HTTPException(status_code=503, detail="Too many pending authorization flows") + if code_challenge_method != "S256": raise HTTPException( status_code=400, detail="Only S256 code_challenge_method is supported" @@ -733,6 +740,12 @@ async def byok_token( server_id=server_id, credential=api_key_value, ) + # Invalidate any cached negative result so the user isn't blocked + # for up to the TTL period after completing the OAuth flow. + from litellm.proxy._experimental.mcp_server.server import ( + _invalidate_byok_cred_cache, + ) + _invalidate_byok_cred_cache(user_id, server_id) except Exception as exc: verbose_proxy_logger.error( "byok_token: failed to store user credential for user=%s server=%s: %s", @@ -755,6 +768,9 @@ async def byok_token( payload = { "user_id": user_id, "server_id": server_id, + # "type" distinguishes this from regular proxy auth tokens. + # The proxy's SSO JWT path uses asymmetric keys (RS256/ES256), so an + # HS256 token signed with master_key cannot be accepted there. "type": "byok_session", "iat": now, "exp": now + 3600, diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 3bccbdbc879..6e9e24b3d46 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -59,6 +59,16 @@ from litellm.utils import Rules, client, function_setup # Keyed by (user_id, server_id); value is (credential_exists, monotonic_timestamp). _byok_cred_cache: Dict[Tuple[str, str], Tuple[bool, float]] = {} _BYOK_CRED_CACHE_TTL = 60 # seconds +_BYOK_CRED_CACHE_MAX_SIZE = 4096 # cap to prevent unbounded growth + + +def _invalidate_byok_cred_cache(user_id: str, server_id: str) -> None: + """Remove a (user_id, server_id) entry from the BYOK credential cache. + + Call this after storing a new credential so the user isn't blocked by a + previously cached negative result. + """ + _byok_cred_cache.pop((user_id, server_id), None) # Check if MCP is available # "mcp" requires python 3.10 or higher, but several litellm users use python 3.8 @@ -1591,6 +1601,8 @@ if MCP_AVAILABLE: user_id=user_id, server_id=mcp_server.server_id, ) + if len(_byok_cred_cache) >= _BYOK_CRED_CACHE_MAX_SIZE: + _byok_cred_cache.clear() _byok_cred_cache[cache_key] = (credential_exists, time.monotonic()) if not credential_exists: raise HTTPException( @@ -1698,17 +1710,30 @@ if MCP_AVAILABLE: "mcp_tool_call_metadata" ] = standard_logging_mcp_tool_call - # BYOK check: if this server requires a per-user key and the - # user has not stored one yet, issue a 401 OAuth challenge so - # that an MCP client can trigger the authorization flow. - await _check_byok_credential(mcp_server, user_api_key_auth) - - # For BYOK servers, inject the user's stored credential as the - # auth header if no explicit override was provided by the caller. + # BYOK: retrieve the stored per-user credential. A single DB call + # both checks existence and fetches the value, avoiding a double query. if mcp_server.is_byok and not mcp_auth_header: - mcp_auth_header = await _get_byok_credential( - mcp_server, user_api_key_auth - ) + byok_cred = await _get_byok_credential(mcp_server, user_api_key_auth) + if byok_cred is None: + raise HTTPException( + status_code=401, + detail={ + "error": "byok_auth_required", + "server_id": mcp_server.server_id, + "server_name": mcp_server.server_name or mcp_server.name, + "message": ( + "No stored credential found for this BYOK server. " + "Complete the OAuth authorization flow to provide your API key." + ), + }, + headers={ + "WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"' + }, + ) + mcp_auth_header = byok_cred + elif mcp_server.is_byok: + # External auth header supplied; still enforce user-identity check. + await _check_byok_credential(mcp_server, user_api_key_auth) # Check if tool exists in local registry first (for OpenAPI-based tools) # These tools are registered with their prefixed names diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index b08f6943358..912c161418d 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -1095,6 +1095,10 @@ if MCP_AVAILABLE: ) if payload.save: await store_user_credential(prisma_client, user_id, server_id, payload.credential) + from litellm.proxy._experimental.mcp_server.server import ( + _invalidate_byok_cred_cache, + ) + _invalidate_byok_cred_cache(user_id, server_id) return MCPUserCredentialResponse(server_id=server_id, has_credential=True) # save=False: credential not persisted return MCPUserCredentialResponse(server_id=server_id, has_credential=False) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx b/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx index 397e55aa4ec..dbc1c4d10e2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx @@ -228,12 +228,6 @@ const menuItems: MenuItemCfg[] = [ icon: , roles: all_admin_roles, }, - { - key: "29", - page: "byok-demo", - label: "BYOK Demo", - icon: , - }, ], }, {