From 4fb589ab2ebb0e0d32e3269cedf2fedb16f69054 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 7 Mar 2026 16:28:58 -0800 Subject: [PATCH] refactor: add get_cached_byok_credential helper, move prisma_client import, add form-encoded tests - Add get_cached_byok_credential() public helper to server.py to expose cache reads without tight coupling to private internals (_byok_cred_cache, TTL const) - Update status endpoint to use get_cached_byok_credential() instead of importing private cache internals directly - Move prisma_client import in callback to after all early-exit paths so it is only imported when actually needed for token exchange - Add test_callback_url_encoded_success_response: verifies form-encoded token response (GitHub default) is stored as plain string - Add test_callback_url_encoded_error_response: verifies form-encoded error body returns 502 with error key in HTML response --- .../mcp_server/openapi_oauth2_endpoints.py | 41 +++--- .../proxy/_experimental/mcp_server/server.py | 18 +++ .../test_openapi_oauth2_endpoints.py | 135 ++++++++++++++++++ 3 files changed, 171 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/openapi_oauth2_endpoints.py b/litellm/proxy/_experimental/mcp_server/openapi_oauth2_endpoints.py index 4d1e79f21a8..9a4116abb36 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_oauth2_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_oauth2_endpoints.py @@ -279,8 +279,6 @@ async def openapi_oauth2_callback( # noqa: PLR0915 error: Optional[str] = Query(default=None), error_description: Optional[str] = Query(default=None), ) -> Response: - from litellm.proxy.proxy_server import prisma_client - if error: # Consume the state on denial/error so orphaned entries don't fill the store. # RFC 6749 §4.1.2.1 requires the provider to echo back the state in error @@ -335,6 +333,10 @@ async def openapi_oauth2_callback( # noqa: PLR0915 server_id: str = state_data["server_id"] user_id: str = state_data["user_id"] + # Import prisma_client here, after all early-exit paths, so it is only + # imported when actually needed for the token exchange path. + from litellm.proxy.proxy_server import prisma_client + server = global_mcp_server_manager.get_mcp_server_by_id(server_id) if server is None: return HTMLResponse( @@ -574,31 +576,24 @@ async def openapi_oauth2_status( ) # Check the shared credential cache before issuing a DB query. - # The cache is written by _get_byok_credential / _check_byok_credential when - # a tool call fetches the real token. If the token is already cached (i.e. - # the user has made at least one tool call), we can skip the DB entirely. - # We intentionally do NOT write to the cache here: writing a sentinel value - # would create dual interpretations of the same cache key across status and - # auth functions, making the cache contract subtle and fragile. + # The cache is populated by _get_byok_credential / _check_byok_credential + # when a tool call fetches the real token. We access it via the public + # get_cached_byok_credential() helper to avoid tight coupling to internals. try: - import time as _time - from litellm.proxy._experimental.mcp_server.server import ( - _BYOK_CRED_CACHE_TTL, - _byok_cred_cache, + get_cached_byok_credential, ) - cached = _byok_cred_cache.get((user_id, server_id)) - if cached is not None: - cached_cred, ts = cached - if _time.monotonic() - ts < _BYOK_CRED_CACHE_TTL: - return JSONResponse( - { - "connected": bool(cached_cred), - "server_id": server_id, - "server_name": server_name, - } - ) + result = get_cached_byok_credential(user_id, server_id) + if result is not None: + cached_cred, _ = result + return JSONResponse( + { + "connected": bool(cached_cred), + "server_id": server_id, + "server_name": server_name, + } + ) except Exception: pass # If cache import fails, fall through to DB diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index af39c4241fa..53d904367a5 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -101,6 +101,24 @@ def _invalidate_byok_cred_cache(user_id: str, server_id: str) -> None: _byok_cred_cache.pop((user_id, server_id), None) +def get_cached_byok_credential( + user_id: str, server_id: str +) -> Optional[Tuple[Optional[str], bool]]: + """Return (credential, True) if a valid cache entry exists, else None. + + Promotes the entry to MRU on hit. Returns None for expired or missing + entries so callers fall through to the DB without accessing cache internals. + """ + cached = _byok_cred_cache.get((user_id, server_id)) + if cached is None: + return None + cred, ts = cached + if time.monotonic() - ts >= _BYOK_CRED_CACHE_TTL: + return None + _byok_cred_cache.move_to_end((user_id, server_id)) + return cred, True + + def _write_byok_cred_cache( user_id: str, server_id: str, credential: Optional[str] ) -> None: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_oauth2_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_oauth2_endpoints.py index 4eaf9b562b7..a367b62f1d0 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_oauth2_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_oauth2_endpoints.py @@ -462,6 +462,141 @@ async def test_callback_uses_basic_auth_when_token_endpoint_auth_method_is_basic assert "client_secret" not in (call["data"] or {}) +# --------------------------------------------------------------------------- +# URL-encoded form response path (provider compatibility) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_callback_url_encoded_success_response(): + """URL-encoded form response (e.g. GitHub default) — token extracted as plain string.""" + from litellm.proxy._experimental.mcp_server.openapi_oauth2_endpoints import ( + _pending_oauth2_states, + openapi_oauth2_callback, + ) + + state = "test-state-url-encoded-ok" + now = time.time() + _pending_oauth2_states[state] = { + "server_id": "server1", + "user_id": "user1", + "timestamp": now, + "expires_at": now + 600, + "callback_url": "http://localhost:4000/v1/mcp/oauth2/callback", + } + + mock_server = MagicMock() + mock_server.token_url = "https://github.com/login/oauth/access_token" + mock_server.client_id = "cid" + mock_server.client_secret = "csecret" + mock_server.server_name = "GitHub" + mock_server.name = "github" + + mock_response = MagicMock() + # GitHub returns URL-encoded form by default when Accept header is not set + mock_response.headers = {"content-type": "application/x-www-form-urlencoded"} + mock_response.text = "access_token=ghu_form_token&scope=repo&token_type=bearer" + mock_response.raise_for_status = MagicMock() + + stored_credentials: list = [] + + async def fake_store(prisma_client, user_id, server_id, credential): + stored_credentials.append(credential) + + with patch( + "litellm.proxy._experimental.mcp_server.openapi_oauth2_endpoints.global_mcp_server_manager" + ) as mock_mgr, patch( + "litellm.proxy._experimental.mcp_server.openapi_oauth2_endpoints.store_user_credential", + side_effect=fake_store, + ), patch( + "litellm.proxy.proxy_server.prisma_client", + MagicMock(), + create=True, + ), patch( + "litellm.proxy._experimental.mcp_server.server._invalidate_byok_cred_cache", + MagicMock(), + ), patch( + "httpx.AsyncClient" + ) as mock_client_cls: + mock_mgr.get_mcp_server_by_id.return_value = mock_server + mock_async_client = AsyncMock() + mock_async_client.post = AsyncMock(return_value=mock_response) + mock_client_cls.return_value.__aenter__ = AsyncMock(return_value=mock_async_client) + mock_client_cls.return_value.__aexit__ = AsyncMock(return_value=None) + + await openapi_oauth2_callback( + request=MagicMock(), + code="form-auth-code", + state=state, + error=None, + error_description=None, + ) + + assert len(stored_credentials) == 1 + assert stored_credentials[0] == "ghu_form_token" + + +@pytest.mark.asyncio +async def test_callback_url_encoded_error_response(): + """URL-encoded form error body from provider returns 502 error page.""" + from fastapi.responses import HTMLResponse + + from litellm.proxy._experimental.mcp_server.openapi_oauth2_endpoints import ( + _pending_oauth2_states, + openapi_oauth2_callback, + ) + + state = "test-state-url-encoded-err" + now = time.time() + _pending_oauth2_states[state] = { + "server_id": "server1", + "user_id": "user1", + "timestamp": now, + "expires_at": now + 600, + "callback_url": "http://localhost:4000/v1/mcp/oauth2/callback", + } + + mock_server = MagicMock() + mock_server.token_url = "https://github.com/login/oauth/access_token" + mock_server.client_id = "cid" + mock_server.client_secret = "csecret" + mock_server.server_name = "GitHub" + mock_server.name = "github" + + mock_response = MagicMock() + mock_response.headers = {"content-type": "application/x-www-form-urlencoded"} + mock_response.text = "error=bad_verification_code&error_description=The+code+is+expired" + mock_response.raise_for_status = MagicMock() + + with patch( + "litellm.proxy._experimental.mcp_server.openapi_oauth2_endpoints.global_mcp_server_manager" + ) as mock_mgr, patch( + "litellm.proxy.proxy_server.prisma_client", + MagicMock(), + create=True, + ), patch( + "httpx.AsyncClient" + ) as mock_client_cls: + mock_mgr.get_mcp_server_by_id.return_value = mock_server + mock_async_client = AsyncMock() + mock_async_client.post = AsyncMock(return_value=mock_response) + mock_client_cls.return_value.__aenter__ = AsyncMock(return_value=mock_async_client) + mock_client_cls.return_value.__aexit__ = AsyncMock(return_value=None) + + result = await openapi_oauth2_callback( + request=MagicMock(), + code="expired-auth-code", + state=state, + error=None, + error_description=None, + ) + + assert isinstance(result, HTMLResponse) + assert result.status_code == 502 + body_bytes = bytes(result.body) if not isinstance(result.body, bytes) else result.body + assert b"bad_verification_code" in body_bytes + + # --------------------------------------------------------------------------- # _extract_access_token (server.py helper) # ---------------------------------------------------------------------------