fix(mcp): keep the BYOK pending-code cap with the shared store and bound the submitted key size

With redis_usage_cache the pending codes no longer live in the in-memory dict, so the 1000-code cap checked against that dict no longer applied. Count issued codes per TTL window in the shared store with an atomic increment (fail closed when it cannot count) and reject API keys longer than 8192 characters before a code is issued.

litellm/proxy/mcp_server/byok_oauth_endpoints.py

@@ -50,11 +50,16 @@ from litellm.proxy.middleware.per_request_root_path_middleware import get_server
 _byok_auth_codes: Final[dict[str, dict]] = {}
 _AUTH_CODE_CACHE_PREFIX: Final = "mcp_byok_auth_code:"
 _AUTH_CODE_CLAIM_PREFIX: Final = "mcp_byok_auth_code_claim:"
+# Codes issued in the current TTL window, counted in the shared store; bounds pending codes
+# there the way len(_byok_auth_codes) bounds them in memory.
+_AUTH_CODE_WINDOW_KEY: Final = "mcp_byok_auth_code_window"
 
 # Authorization codes expire after 5 minutes.
 _AUTH_CODE_TTL_SECONDS: Final = 300
 # Hard cap to prevent memory exhaustion from incomplete OAuth flows.
 _AUTH_CODES_MAX_SIZE: Final = 1000
+# Upper bound for a submitted API key, so a pending code cannot carry an arbitrarily large value.
+_MAX_API_KEY_LENGTH: Final = 8192
 
 router: Final = APIRouter(tags=["mcp"])
 
@@ -133,6 +138,21 @@ async def _discard_auth_code(code: str) -> None:
         await shared.async_delete_cache(_AUTH_CODE_CACHE_PREFIX + code)
 
 
+async def _reserve_auth_code_slot() -> bool:
+    """Whether another pending code may be issued. In memory this is the dict size; with the
+    shared store it is an atomic per-window counter, since the dict stays empty there. A shared
+    store that cannot count refuses (fail closed)."""
+    shared: Final = _shared_code_store()
+    if shared is None:
+        return len(_byok_auth_codes) < _AUTH_CODES_MAX_SIZE
+    try:
+        issued: Final = await shared.async_increment(_AUTH_CODE_WINDOW_KEY, 1, ttl=_AUTH_CODE_TTL_SECONDS)
+    except Exception as e:  # noqa: BLE001  # any shared-store fault refuses the new code
+        verbose_proxy_logger.warning("byok authorize: pending-code counter unavailable: %s", e)
+        return False
+    return issued <= _AUTH_CODES_MAX_SIZE
+
+
 async def _claim_auth_code(code: str) -> bool:
     """Atomic one-time use across workers (RFC 6749 section 4.1.2): of two concurrent
     redemptions of the same code only the first claim succeeds, and a claim that cannot be
@@ -904,7 +924,9 @@ async def byok_authorize_post(
 
     # 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:
+    if len(api_key) > _MAX_API_KEY_LENGTH:
+        raise HTTPException(status_code=400, detail="api_key is too long")
+    if not await _reserve_auth_code_slot():
         raise HTTPException(status_code=503, detail="Too many pending authorization flows")
 
     if code_challenge_method != "S256":

tests/unit/proxy/mcp_server/test_byok_oauth_endpoints.py

@@ -2019,7 +2019,7 @@ class _SharedStore:
         self.values.pop(key, None)
 
     async def async_increment(self, key: str, value: int, **_: Any) -> int:
-        if self.fail_increment:
+        if self.fail_increment and key.startswith("mcp_byok_auth_code_claim:"):
             raise ConnectionError("redis down")
         self.counts[key] = self.counts.get(key, 0) + value
         return self.counts[key]
@@ -2114,3 +2114,46 @@ async def test_a_shared_code_is_not_redeemed_when_the_single_use_claim_cannot_be
         result = await _redeem(code, verifier, store_credential)
     assert result.status_code == 400
     store_credential.assert_not_awaited()
+
+
+def test_the_pending_code_cap_holds_with_the_shared_store(client, monkeypatch):
+    from litellm.proxy.mcp_server.byok_oauth_endpoints import _AUTH_CODE_WINDOW_KEY, _AUTH_CODES_MAX_SIZE
+
+    monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-key")
+    shared = _SharedStore()
+    shared.counts[_AUTH_CODE_WINDOW_KEY] = _AUTH_CODES_MAX_SIZE
+    with patch("litellm.proxy.proxy_server.redis_usage_cache", shared):
+        resp = client.post(
+            "/v1/mcp/oauth/authorize",
+            data={
+                "client_id": "user-123",
+                "redirect_uri": "http://127.0.0.1:3000/callback",
+                "code_challenge": _make_challenge("a_code_verifier_that_is_long_enough_for_pkce_43"),
+                "code_challenge_method": "S256",
+                "state": "st",
+                "server_id": "server-xyz",
+                "api_key": "sk-supersecretkey",
+            },
+            follow_redirects=False,
+        )
+    assert resp.status_code == 503
+    assert shared.values == {}
+
+
+def test_an_oversized_api_key_is_rejected_before_a_code_is_issued(client):
+    from litellm.proxy.mcp_server.byok_oauth_endpoints import _MAX_API_KEY_LENGTH
+
+    resp = client.post(
+        "/v1/mcp/oauth/authorize",
+        data={
+            "client_id": "user-123",
+            "redirect_uri": "http://127.0.0.1:3000/callback",
+            "code_challenge": _make_challenge("a_code_verifier_that_is_long_enough_for_pkce_43"),
+            "code_challenge_method": "S256",
+            "state": "st",
+            "server_id": "server-xyz",
+            "api_key": "k" * (_MAX_API_KEY_LENGTH + 1),
+        },
+        follow_redirects=False,
+    )
+    assert resp.status_code == 400