From a804d2cc66eb41a509bc9b980881d5fa0f7ee6f2 Mon Sep 17 00:00:00 2001 From: Artur Sidorenko Date: Mon, 2 Nov 2026 11:58:26 +0100 Subject: [PATCH] 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()