scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(mcp-oauth2): fix 3 bugs from greptile 2/5 review

Browse source 
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
This commit is contained in:
Ishaan Jaffer 2026-03-06 18:40:02 -08:00
parent d52db24931
commit c3387c476f
3 changed files with 34 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
litellm/proxy/_experimental/mcp_server/openapi_oauth2_endpoints.py
View file

@ -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,

9
tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_oauth2_endpoints.py
View file

@ -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,
):

13
ui/litellm-dashboard/src/components/mcp_tools/OAuth2ConnectButton.tsx
View file

@ -59,9 +59,20 @@ export const OAuth2ConnectButton: React.FC<OAuth2ConnectButtonProps> = ({
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;
}