fix: address greptile review feedback (greploop iteration 2)

- Cache invalidation: add _invalidate_byok_cred_cache() and call it after
  store_user_credential() in both token endpoint and management endpoint
- Unbounded cache: add _BYOK_CRED_CACHE_MAX_SIZE=4096 with clear-on-overflow
- Unbounded auth codes: add _AUTH_CODES_MAX_SIZE=1000 with 503 on overflow
- Double DB query: merge _check_byok_credential + _get_byok_credential into
  single _get_byok_credential call; raise 401 inline if None returned
- Sidebar: remove byok-demo entry (page was deleted in prior commit)
- JWT comment: document why byok_session HS256 token can't be used as proxy auth

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,

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

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)

ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx

@@ -228,12 +228,6 @@ const menuItems: MenuItemCfg[] = [
         icon: <DatabaseOutlined style={{ fontSize: 18 }} />,
         roles: all_admin_roles,
       },
-      {
-        key: "29",
-        page: "byok-demo",
-        label: "BYOK Demo",
-        icon: <KeyOutlined style={{ fontSize: 18 }} />,
-      },
     ],
   },
   {