From c3387c476f120cc4ad5d204f76f7e712ea767934 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Fri, 6 Mar 2026 18:40:02 -0800 Subject: [PATCH] fix(mcp-oauth2): fix 3 bugs from greptile 2/5 review 1. Separate try-blocks for store_user_credential and _invalidate_byok_cred_cache: a cache-flush failure no longer returns a 'Storage error' page when the write succeeded 2. Popup-close race fix in OAuth2ConnectButton: do one final status check when popup is detected closed so fast OAuth flows (popup auto-closes before the first 2s poll fires) are not silently missed 3. Fix test patches: `master_key` and `prisma_client` are imported inline in the function body via `from proxy_server import X`; patch `litellm.proxy.proxy_server.*` not the endpoint module --- .../mcp_server/openapi_oauth2_endpoints.py | 20 ++++++++++++++----- .../test_openapi_oauth2_endpoints.py | 9 +++++++-- .../mcp_tools/OAuth2ConnectButton.tsx | 13 +++++++++++- 3 files changed, 34 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/openapi_oauth2_endpoints.py b/litellm/proxy/_experimental/mcp_server/openapi_oauth2_endpoints.py index 226ca7520e3..74fa23c8cf7 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_oauth2_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_oauth2_endpoints.py @@ -451,11 +451,6 @@ async def openapi_oauth2_callback( server_id=server_id, credential=credential_to_store, ) - 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( "openapi_oauth2_callback: failed to store credential user=%s server=%s: %s", @@ -471,6 +466,21 @@ async def openapi_oauth2_callback( status_code=500, ) + # Best-effort cache flush; a failure here must NOT mask the successful write above. + try: + 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.warning( + "openapi_oauth2_callback: cache invalidation failed (credential was stored) user=%s server=%s: %s", + user_id, + server_id, + exc, + ) + verbose_proxy_logger.info( "openapi_oauth2_callback: connected user=%s server=%s", user_id, 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 cbe08a9fc1c..5e7145db66b 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 @@ -105,10 +105,12 @@ async def test_connect_missing_client_secret_raises_400(): mock_server.client_id = "my-client-id" mock_server.client_secret = None # missing + # master_key is imported inline via `from litellm.proxy.proxy_server import master_key`; + # patch at the source so the function sees the mock value. 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.master_key", + "litellm.proxy.proxy_server.master_key", "sk-test", create=True, ): @@ -207,10 +209,13 @@ async def test_status_no_prisma_returns_not_connected(): mock_server.server_name = "GitHub" mock_server.name = "github" + # prisma_client is imported inside the function via + # `from litellm.proxy.proxy_server import prisma_client`, so we must patch + # it at the source module rather than on the endpoint module. 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.prisma_client", + "litellm.proxy.proxy_server.prisma_client", None, create=True, ): diff --git a/ui/litellm-dashboard/src/components/mcp_tools/OAuth2ConnectButton.tsx b/ui/litellm-dashboard/src/components/mcp_tools/OAuth2ConnectButton.tsx index d1e6be825b2..9d62a09b9ed 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/OAuth2ConnectButton.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/OAuth2ConnectButton.tsx @@ -59,9 +59,20 @@ export const OAuth2ConnectButton: React.FC = ({ return; } - // Stop if popup was closed by user + // Stop if popup was closed by user; do one final status check first so + // that fast OAuth flows (popup auto-closes before the first poll fires) + // are not silently missed. if (popupRef.current && popupRef.current.closed) { stopPolling(); + try { + const finalStatus = await getMcpOAuth2Status(server.server_id, accessToken); + if (finalStatus.connected) { + handleConnected(); + return; + } + } catch { + // ignore — popup was closed by user without completing auth + } setLoading(false); return; }