From a804d2cc66eb41a509bc9b980881d5fa0f7ee6f2 Mon Sep 17 00:00:00 2001 From: Artur Sidorenko Date: Mon, 2 Nov 2026 11:58:26 +0100 Subject: [PATCH 1/2] fix(mcp): share BYOK authorization codes across workers so the token exchange works on multi-worker proxies byok_authorize_post kept pending authorization codes in a module-level dict. The authorize POST and the token request are separate HTTP requests; on a proxy with several workers or replicas the token request often lands on a process that never saw the code and fails with invalid_grant. With redis_usage_cache configured, store the record there, encrypted with the proxy salt key because it carries the user's API key, and redeem it through the same atomic single-use claim the gateway OAuth flow uses (_SingleUseGuard), so two concurrent redemptions cannot both succeed and a Redis fault fails closed. Without Redis the dict stays the store, as before. --- .../proxy/mcp_server/byok_oauth_endpoints.py | 110 +++++++++++++--- .../mcp_server/test_byok_oauth_endpoints.py | 120 ++++++++++++++++++ 2 files changed, 213 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/mcp_server/byok_oauth_endpoints.py b/litellm/proxy/mcp_server/byok_oauth_endpoints.py index 868b73e39f1..9b4ecaf5fac 100644 --- a/litellm/proxy/mcp_server/byok_oauth_endpoints.py +++ b/litellm/proxy/mcp_server/byok_oauth_endpoints.py @@ -15,6 +15,7 @@ Endpoints implemented here: import base64 import hashlib import html as _html_module +import json import time import uuid from typing import Final, cast @@ -38,10 +39,17 @@ from litellm.proxy.mcp_server.oauth_utils import ( from litellm.proxy.middleware.per_request_root_path_middleware import get_server_root_paths # --------------------------------------------------------------------------- -# In-memory store for pending authorization codes. +# Pending authorization codes. # Each entry: {code: {api_key, server_id, code_challenge, redirect_uri, user_id, expires_at}} +# +# The authorize POST and the token request are two separate HTTP requests and can +# land on different workers or replicas, so with a shared Redis (redis_usage_cache) +# the record lives there, encrypted with the proxy salt key because it carries the +# user's API key. Without Redis (a single process) it stays in this dict. # --------------------------------------------------------------------------- _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:" # Authorization codes expire after 5 minutes. _AUTH_CODE_TTL_SECONDS: Final = 300 @@ -75,6 +83,69 @@ def _purge_expired_codes() -> None: del _byok_auth_codes[k] +def _shared_code_store(): + """The cross-worker store for authorization codes, or None on a single-process proxy.""" + from litellm.proxy.proxy_server import redis_usage_cache + + return redis_usage_cache + + +async def _store_auth_code(code: str, record: dict) -> None: + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper + + shared: Final = _shared_code_store() + if shared is None: + _byok_auth_codes[code] = record + return + await shared.async_set_cache( + _AUTH_CODE_CACHE_PREFIX + code, + encrypt_value_helper(json.dumps(record)), + ttl=_AUTH_CODE_TTL_SECONDS, + ) + + +async def _load_auth_code(code: str) -> dict | None: + from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper + + local: Final = _byok_auth_codes.get(code) + if local is not None: + return local + shared: Final = _shared_code_store() + if shared is None: + return None + sealed: Final = await shared.async_get_cache(_AUTH_CODE_CACHE_PREFIX + code) + if not isinstance(sealed, str): + return None + plain: Final = decrypt_value_helper(sealed, "mcp_byok_auth_code", exception_type="debug") + if not isinstance(plain, str): + return None + try: + record: Final = json.loads(plain) + except ValueError: + return None + return record if isinstance(record, dict) else None + + +async def _discard_auth_code(code: str) -> None: + _byok_auth_codes.pop(code, None) + shared: Final = _shared_code_store() + if shared is not None: + await shared.async_delete_cache(_AUTH_CODE_CACHE_PREFIX + code) + + +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 + recorded in the shared backend fails closed.""" + from litellm.proxy.mcp_server.gateway_dcr_flow import _SingleUseGuard + from litellm.proxy.proxy_server import user_api_key_cache + + outcome: Final = await _SingleUseGuard(user_api_key_cache).claim( + _AUTH_CODE_CLAIM_PREFIX + code, _AUTH_CODE_TTL_SECONDS + ) + return outcome == "first" + + def _oauth_token_error(code: str, status: int = 400) -> JSONResponse: """RFC 6749 §5.2 token-endpoint error body: ``{"error": ""}``. FastAPI's default ``HTTPException`` renders ``{"detail": ...}`` which @@ -848,18 +919,21 @@ async def byok_authorize_post( raise HTTPException(status_code=401, detail="login_required") auth_code: Final = str(uuid.uuid4()) - _byok_auth_codes[auth_code] = { - "api_key": api_key, - "server_id": server_id, - "code_challenge": code_challenge, - "redirect_uri": redirect_uri, - # RFC 6749 §4.1.3 defense-in-depth: if the authorization request - # declared a client_id, the token request must submit the same - # value. Stored even though we don't pre-register clients. - "client_id": client_id, - "user_id": user_id, - "expires_at": time.time() + _AUTH_CODE_TTL_SECONDS, - } + await _store_auth_code( + auth_code, + { + "api_key": api_key, + "server_id": server_id, + "code_challenge": code_challenge, + "redirect_uri": redirect_uri, + # RFC 6749 §4.1.3 defense-in-depth: if the authorization request + # declared a client_id, the token request must submit the same + # value. Stored even though we don't pre-register clients. + "client_id": client_id, + "user_id": user_id, + "expires_at": time.time() + _AUTH_CODE_TTL_SECONDS, + }, + ) params: Final = urlencode({"code": auth_code, "state": state}) separator: Final = "&" if "?" in redirect_uri else "?" @@ -895,12 +969,12 @@ async def byok_token( if grant_type != "authorization_code": return _oauth_token_error("unsupported_grant_type") - record: Final = _byok_auth_codes.get(code) + record: Final = await _load_auth_code(code) if record is None: return _oauth_token_error("invalid_grant") if time.time() > record["expires_at"]: - del _byok_auth_codes[code] + await _discard_auth_code(code) return _oauth_token_error("invalid_grant") # PKCE verification @@ -941,8 +1015,10 @@ async def byok_token( if master_key is None: return _oauth_token_error("server_error", status=500) - # Consume the code (one-time use) - del _byok_auth_codes[code] + # Consume the code (one-time use, atomic across workers) + if not await _claim_auth_code(code): + return _oauth_token_error("invalid_grant") + await _discard_auth_code(code) # Persist the BYOK credential if prisma_client is not None: diff --git a/tests/unit/proxy/mcp_server/test_byok_oauth_endpoints.py b/tests/unit/proxy/mcp_server/test_byok_oauth_endpoints.py index 37da00ab2ba..eed5d9fbfa1 100644 --- a/tests/unit/proxy/mcp_server/test_byok_oauth_endpoints.py +++ b/tests/unit/proxy/mcp_server/test_byok_oauth_endpoints.py @@ -1994,3 +1994,123 @@ async def test_authenticated_browser_identity_checks_embedded_session( ) request: Final = Request({"type": "http", "headers": [(b"cookie", ("token=" + cookie).encode())]}) assert await get_authenticated_browser_user_id(request) == expected + + +# --------------------------------------------------------------------------- +# Authorization codes shared across workers +# --------------------------------------------------------------------------- + + +class _SharedStore: + """Stands in for redis_usage_cache: one store every worker sees.""" + + def __init__(self, fail_increment: bool = False) -> None: + self.values: dict[str, Any] = {} + self.counts: dict[str, int] = {} + self.fail_increment = fail_increment + + async def async_set_cache(self, key: str, value: Any, **_: Any) -> None: + self.values[key] = value + + async def async_get_cache(self, key: str, **_: Any) -> Any: + return self.values.get(key) + + async def async_delete_cache(self, key: str) -> None: + self.values.pop(key, None) + + async def async_increment(self, key: str, value: int, **_: Any) -> int: + if self.fail_increment: + raise ConnectionError("redis down") + self.counts[key] = self.counts.get(key, 0) + value + return self.counts[key] + + +def _authorize(client, challenge: str) -> str: + from urllib.parse import parse_qs, urlparse + + resp = client.post( + "/v1/mcp/oauth/authorize", + data={ + "client_id": "user-123", + "redirect_uri": "http://127.0.0.1:3000/callback", + "code_challenge": challenge, + "code_challenge_method": "S256", + "state": "st", + "server_id": "server-xyz", + "api_key": "sk-supersecretkey", + }, + follow_redirects=False, + ) + assert resp.status_code == 302 + return parse_qs(urlparse(resp.headers["location"]).query)["code"][0] + + +async def _redeem(code: str, verifier: str, store_credential: AsyncMock): + from litellm.proxy.mcp_server.byok_oauth_endpoints import byok_token + + with ( + patch("litellm.proxy.mcp_server.byok_oauth_endpoints.store_user_credential", store_credential), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.master_key", "test_master_key_value"), + patch("litellm.proxy.mcp_server.server._invalidate_byok_cred_cache", AsyncMock(), create=True), + ): + return await byok_token( + request=MagicMock(), + grant_type="authorization_code", + code=code, + redirect_uri="http://127.0.0.1:3000/callback", + code_verifier=verifier, + client_id="user-123", + ) + + +@pytest.mark.asyncio +async def test_a_code_issued_on_one_worker_is_redeemed_on_another_through_the_shared_store(client, monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-key") + shared = _SharedStore() + verifier = "a_code_verifier_that_is_long_enough_for_pkce_43" + with patch("litellm.proxy.proxy_server.redis_usage_cache", shared): + code = _authorize(client, _make_challenge(verifier)) + assert code not in _byok_auth_codes + (sealed,) = shared.values.values() + assert "sk-supersecretkey" not in json.dumps(sealed) + + store_credential = AsyncMock() + result = await _redeem(code, verifier, store_credential) + assert result.status_code == 200 + assert json.loads(result.body)["access_token"] + store_credential.assert_awaited_once() + assert store_credential.await_args.kwargs["credential"] == "sk-supersecretkey" + assert shared.values == {} + + +@pytest.mark.asyncio +async def test_a_shared_code_is_redeemed_only_once(client, monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-key") + shared = _SharedStore() + verifier = "a_code_verifier_that_is_long_enough_for_pkce_43" + with patch("litellm.proxy.proxy_server.redis_usage_cache", shared): + code = _authorize(client, _make_challenge(verifier)) + sealed = dict(shared.values) + first = await _redeem(code, verifier, AsyncMock()) + # a concurrent redemption that read the record before the first one deleted it + shared.values.update(sealed) + store_credential = AsyncMock() + second = await _redeem(code, verifier, store_credential) + assert first.status_code == 200 + assert second.status_code == 400 + assert json.loads(second.body)["error"] == "invalid_grant" + store_credential.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_a_shared_code_is_not_redeemed_when_the_single_use_claim_cannot_be_recorded(client, monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-key") + shared = _SharedStore(fail_increment=True) + verifier = "a_code_verifier_that_is_long_enough_for_pkce_43" + with patch("litellm.proxy.proxy_server.redis_usage_cache", shared): + code = _authorize(client, _make_challenge(verifier)) + store_credential = AsyncMock() + result = await _redeem(code, verifier, store_credential) + assert result.status_code == 400 + store_credential.assert_not_awaited() From 36097f2c21742c698d736a2bfe60a7fe567cb26e Mon Sep 17 00:00:00 2001 From: Artur Sidorenko Date: Mon, 2 Nov 2026 13:03:51 +0100 Subject: [PATCH 2/2] 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. --- .../proxy/mcp_server/byok_oauth_endpoints.py | 24 +++++++++- .../mcp_server/test_byok_oauth_endpoints.py | 45 ++++++++++++++++++- 2 files changed, 67 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/mcp_server/byok_oauth_endpoints.py b/litellm/proxy/mcp_server/byok_oauth_endpoints.py index 9b4ecaf5fac..07bdb81a0f2 100644 --- a/litellm/proxy/mcp_server/byok_oauth_endpoints.py +++ b/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": diff --git a/tests/unit/proxy/mcp_server/test_byok_oauth_endpoints.py b/tests/unit/proxy/mcp_server/test_byok_oauth_endpoints.py index eed5d9fbfa1..f1073c5ed6b 100644 --- a/tests/unit/proxy/mcp_server/test_byok_oauth_endpoints.py +++ b/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