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: ,
- },
],
},
{