From 4626862580ad6413489da5a4f8558fa8eb6cff90 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 10 Mar 2026 20:28:23 -0700 Subject: [PATCH] address greptile review feedback --- .../mcp_management_endpoints.py | 59 +++--- .../test_mcp_management_endpoints.py | 180 ++++++++++++++++++ .../src/components/chat/OAuthConnectModal.tsx | 5 +- .../src/components/networking.tsx | 24 ++- .../src/hooks/useMcpOAuthFlow.tsx | 12 +- .../src/hooks/useUserMcpOAuthFlow.tsx | 17 +- ui/litellm-dashboard/src/utils/errorUtils.ts | 35 ++++ 7 files changed, 275 insertions(+), 57 deletions(-) create mode 100644 ui/litellm-dashboard/src/utils/errorUtils.ts diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index e16e0ff481b..bcc69b34775 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -36,6 +36,11 @@ from fastapi import ( ) from fastapi.responses import JSONResponse +try: + from prisma.errors import RecordNotFoundError +except ImportError: + RecordNotFoundError = Exception # type: ignore + import litellm from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._uuid import uuid @@ -84,6 +89,7 @@ if MCP_AVAILABLE: delete_user_credential, get_all_mcp_servers_for_user, get_mcp_server, + get_mcp_servers, get_mcp_submissions, get_user_oauth_credential, list_user_oauth_credentials, @@ -1427,14 +1433,11 @@ if MCP_AVAILABLE: expires_in=payload.expires_in, scopes=payload.scopes, ) - from datetime import timedelta - from datetime import timezone as _tz - - expires_at: Optional[str] = None - if payload.expires_in is not None: - expires_at = ( - datetime.now(_tz.utc) + timedelta(seconds=payload.expires_in) - ).isoformat() + # Read back the persisted record so the response reflects the stored + # expires_at rather than recomputing it here (which could diverge by + # milliseconds or if the storage logic ever adds a grace period). + stored = await get_user_oauth_credential(prisma_client, user_id, server_id) + expires_at: Optional[str] = stored.get("expires_at") if stored else None return MCPOAuthUserCredentialStatus( server_id=server_id, has_credential=True, @@ -1463,10 +1466,15 @@ if MCP_AVAILABLE: status_code=status.HTTP_400_BAD_REQUEST, detail={"error": "User ID not found in token"}, ) - try: - await delete_user_credential(prisma_client, user_id, server_id) - except Exception: - pass # Already gone + # Only delete if the stored credential is actually an OAuth2 token. + # This prevents accidentally deleting a BYOK credential if one exists + # for the same (user_id, server_id) pair. + cred_to_delete = await get_user_oauth_credential(prisma_client, user_id, server_id) + if cred_to_delete is not None: + try: + await delete_user_credential(prisma_client, user_id, server_id) + except RecordNotFoundError: + pass # Already gone — treat as a successful delete return MCPOAuthUserCredentialStatus( server_id=server_id, has_credential=False, @@ -1485,8 +1493,6 @@ if MCP_AVAILABLE: user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """Return credential status (has_credential, expiry) without exposing the token.""" - from datetime import timezone as _tz - prisma_client = get_prisma_client_or_throw( "Database not connected. Connect a database to your proxy" ) @@ -1505,10 +1511,8 @@ if MCP_AVAILABLE: is_expired = False if expires_at: try: - from datetime import datetime as _dt - - exp = _dt.fromisoformat(expires_at) - is_expired = exp < _dt.now(_tz.utc) + exp = datetime.fromisoformat(expires_at) + is_expired = exp < datetime.now(timezone.utc) except Exception: pass return MCPOAuthUserCredentialStatus( @@ -1530,8 +1534,6 @@ if MCP_AVAILABLE: user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """Return all servers the calling user has connected via OAuth2.""" - from datetime import timezone as _tz - prisma_client = get_prisma_client_or_throw( "Database not connected. Connect a database to your proxy" ) @@ -1544,13 +1546,12 @@ if MCP_AVAILABLE: oauth_creds = await list_user_oauth_credentials(prisma_client, user_id) if not oauth_creds: return [] - # Fetch server metadata for display names + # Fetch server metadata for display names — single batch query instead of N+1. server_ids = [c["server_id"] for c in oauth_creds] - servers = {} - for sid in server_ids: - srv = await get_mcp_server(prisma_client, sid) - if srv is not None: - servers[sid] = srv + servers = { + srv.server_id: srv + for srv in await get_mcp_servers(prisma_client, server_ids) + } items: List[MCPUserCredentialListItem] = [] for cred in oauth_creds: sid = cred["server_id"] @@ -1559,9 +1560,7 @@ if MCP_AVAILABLE: is_expired = False if expires_at: try: - from datetime import datetime as _dt - - is_expired = _dt.fromisoformat(expires_at) < _dt.now(_tz.utc) + is_expired = datetime.fromisoformat(expires_at) < datetime.now(timezone.utc) except Exception: pass items.append( @@ -1571,7 +1570,7 @@ if MCP_AVAILABLE: alias=getattr(srv, "alias", None) if srv else None, credential_type="oauth2", has_credential=True, - expires_at=None if is_expired else expires_at, + expires_at=expires_at, # always pass the raw timestamp; client computes expiry state connected_at=cred.get("connected_at"), ) ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index 30b3be4a3eb..ebb666cfb0e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -1854,3 +1854,183 @@ class TestValidateMCPRequiredFields: _validate_mcp_required_fields(payload) assert exc_info.value.status_code == 500 assert "source_Url" in str(exc_info.value.detail) + + +# ── OAuth user credential endpoint unit tests ────────────────────────────────── + + +def _make_user_auth(user_id: str = "user-abc") -> "UserAPIKeyAuth": + return UserAPIKeyAuth( + api_key="sk-test", + user_id=user_id, + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + +def _make_prisma_client(): + """Return a minimal mock PrismaClient accepted by get_prisma_client_or_throw.""" + client = MagicMock() + client.db = MagicMock() + return client + + +@pytest.mark.asyncio +async def test_store_mcp_oauth_user_credential_returns_status(): + """store_mcp_oauth_user_credential persists the token and echoes back status.""" + from litellm.proxy._types import ( + MCPOAuthUserCredentialRequest, + MCPOAuthUserCredentialStatus, + ) + + if not mgmt_endpoints.MCP_AVAILABLE: + pytest.skip("MCP module not installed") + + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + store_mcp_oauth_user_credential, + ) + + server_id = "srv-1" + user_id = "user-123" + stored_payload = { + "type": "oauth2", + "access_token": "tok", + "expires_at": "2099-01-01T00:00:00+00:00", + "connected_at": "2026-01-01T00:00:00+00:00", + "server_id": server_id, + } + + mock_prisma = _make_prisma_client() + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=mock_prisma, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + new=AsyncMock(return_value=generate_mock_mcp_server_db_record(server_id=server_id)), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.store_user_oauth_credential", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_user_oauth_credential", + new=AsyncMock(return_value=stored_payload), + ), + ): + result = await store_mcp_oauth_user_credential( + server_id=server_id, + payload=MCPOAuthUserCredentialRequest( + access_token="tok", + expires_in=3600, + ), + user_api_key_dict=_make_user_auth(user_id), + ) + + assert isinstance(result, MCPOAuthUserCredentialStatus) + assert result.has_credential is True + assert result.server_id == server_id + # expires_at should come from the stored record, not be recomputed + assert result.expires_at == "2099-01-01T00:00:00+00:00" + + +@pytest.mark.asyncio +async def test_delete_mcp_oauth_user_credential_only_deletes_oauth(): + """delete_mcp_oauth_user_credential only deletes OAuth2 credentials, not BYOK.""" + from litellm.proxy._types import MCPOAuthUserCredentialStatus + + if not mgmt_endpoints.MCP_AVAILABLE: + pytest.skip("MCP module not installed") + + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + delete_mcp_oauth_user_credential, + ) + + server_id = "srv-2" + user_id = "user-456" + delete_mock = AsyncMock(return_value=None) + + # When get_user_oauth_credential returns None (no OAuth cred), delete should NOT be called. + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=_make_prisma_client(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_user_oauth_credential", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.delete_user_credential", + new=delete_mock, + ), + ): + result = await delete_mcp_oauth_user_credential( + server_id=server_id, + user_api_key_dict=_make_user_auth(user_id), + ) + + delete_mock.assert_not_called() + assert isinstance(result, MCPOAuthUserCredentialStatus) + assert result.has_credential is False + + +@pytest.mark.asyncio +async def test_list_mcp_user_credentials_batch_server_fetch(): + """list_mcp_user_credentials uses a single batch DB call, not N+1 queries.""" + from litellm.proxy._types import MCPUserCredentialListItem + + if not mgmt_endpoints.MCP_AVAILABLE: + pytest.skip("MCP module not installed") + + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + list_mcp_user_credentials, + ) + + user_id = "user-789" + server_id = "srv-3" + stored_creds = [ + { + "type": "oauth2", + "access_token": "tok", + "expires_at": "2099-01-01T00:00:00+00:00", + "connected_at": "2026-01-01T00:00:00+00:00", + "server_id": server_id, + } + ] + mock_server = generate_mock_mcp_server_db_record(server_id=server_id, alias="My Server") + # get_mcp_servers (batch) should be called once; get_mcp_server (single) must not be called. + batch_mock = AsyncMock(return_value=[mock_server]) + single_mock = AsyncMock(return_value=mock_server) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=_make_prisma_client(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.list_user_oauth_credentials", + new=AsyncMock(return_value=stored_creds), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_servers", + new=batch_mock, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + new=single_mock, + ), + ): + result = await list_mcp_user_credentials( + user_api_key_dict=_make_user_auth(user_id), + ) + + batch_mock.assert_called_once() + single_mock.assert_not_called() + assert len(result) == 1 + assert isinstance(result[0], MCPUserCredentialListItem) + assert result[0].server_id == server_id + assert result[0].alias == "My Server" + # expires_at should always be the raw timestamp (not set to None when expired) + assert result[0].expires_at == "2099-01-01T00:00:00+00:00" diff --git a/ui/litellm-dashboard/src/components/chat/OAuthConnectModal.tsx b/ui/litellm-dashboard/src/components/chat/OAuthConnectModal.tsx index 3efbef9ad9e..fb9568463bd 100644 --- a/ui/litellm-dashboard/src/components/chat/OAuthConnectModal.tsx +++ b/ui/litellm-dashboard/src/components/chat/OAuthConnectModal.tsx @@ -66,11 +66,12 @@ export const OAuthConnectModal: React.FC = ({ scopes, onSuccess: () => { onSuccess(server.server_id); - handleClose(); + // handleClose is invoked by the useEffect below to avoid calling it twice. }, }); - // If we return from OAuth callback, close the modal so the user isn't stuck. + // Close the modal whenever the flow reaches "success" — covers both the + // in-page path (status set directly) and the post-redirect resume path. useEffect(() => { if (status === "success") { handleClose(); diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 2176ef584a4..4535f51e5a2 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -9600,7 +9600,17 @@ export const storeMCPOAuthUserCredential = async ( }); if (!response.ok) { const err = await response.json().catch(() => ({})); - throw new Error((err as { detail?: { error?: string } })?.detail?.error || "Failed to store OAuth credential"); + const errObj = err as { detail?: unknown }; + const detail = errObj?.detail; + const detailMsg = + Array.isArray(detail) + ? detail.map((d: unknown) => (d && typeof d === "object" ? (d as Record).msg ?? JSON.stringify(d) : String(d))).join("; ") + : typeof detail === "string" + ? detail + : detail && typeof (detail as Record).error === "string" + ? (detail as Record).error as string + : undefined; + throw new Error(detailMsg || "Failed to store OAuth credential"); } return response.json(); }; @@ -9618,7 +9628,17 @@ export const deleteMCPOAuthUserCredential = async ( }); if (!response.ok) { const err = await response.json().catch(() => ({})); - throw new Error((err as { detail?: { error?: string } })?.detail?.error || "Failed to revoke OAuth credential"); + const errObj = err as { detail?: unknown }; + const detail = errObj?.detail; + const detailMsg = + Array.isArray(detail) + ? detail.map((d: unknown) => (d && typeof d === "object" ? (d as Record).msg ?? JSON.stringify(d) : String(d))).join("; ") + : typeof detail === "string" + ? detail + : detail && typeof (detail as Record).error === "string" + ? (detail as Record).error as string + : undefined; + throw new Error(detailMsg || "Failed to revoke OAuth credential"); } return response.json(); }; diff --git a/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx b/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx index 97d48510c2e..b39466320d7 100644 --- a/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx +++ b/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx @@ -10,20 +10,10 @@ import { registerMcpOAuthClient, serverRootPath, } from "@/components/networking"; +import { extractErrorMessage } from "@/utils/errorUtils"; export type McpOAuthStatus = "idle" | "authorizing" | "exchanging" | "success" | "error"; -function extractErrorMessage(err: unknown): string { - if (err instanceof Error) return err.message; - if (err && typeof err === "object") { - const e = err as Record; - if (typeof e.detail === "string") return e.detail; - if (typeof e.message === "string") return e.message; - return JSON.stringify(err); - } - return String(err); -} - interface UseMcpOAuthFlowOptions { accessToken: string | null; getCredentials: () => { diff --git a/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx b/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx index 53465d5554c..ef9b6dedbed 100644 --- a/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx +++ b/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx @@ -22,20 +22,10 @@ import { storeMCPOAuthUserCredential, } from "@/components/networking"; import NotificationsManager from "@/components/molecules/notifications_manager"; +import { extractErrorMessage } from "@/utils/errorUtils"; export type UserMcpOAuthStatus = "idle" | "authorizing" | "exchanging" | "success" | "error"; -function extractErrorMessage(err: unknown): string { - if (err instanceof Error) return err.message; - if (err && typeof err === "object") { - const e = err as Record; - if (typeof e.detail === "string") return e.detail; - if (typeof e.message === "string") return e.message; - return JSON.stringify(err); - } - return String(err); -} - interface UseUserMcpOAuthFlowOptions { accessToken: string; serverId: string; @@ -89,8 +79,11 @@ const genChallenge = async (verifier: string) => { const setStorage = (key: string, value: string) => { try { + // Use sessionStorage only — do not write to localStorage. + // The flow state may contain the LiteLLM access token; writing it to + // localStorage would persist it across browser sessions and make it + // readable by any injected script (XSS). window.sessionStorage.setItem(key, value); - window.localStorage.setItem(key, value); } catch (_) {} }; diff --git a/ui/litellm-dashboard/src/utils/errorUtils.ts b/ui/litellm-dashboard/src/utils/errorUtils.ts new file mode 100644 index 00000000000..665e4e9cb72 --- /dev/null +++ b/ui/litellm-dashboard/src/utils/errorUtils.ts @@ -0,0 +1,35 @@ +/** + * Shared error-message extraction utility. + * + * Handles the common shapes returned by LiteLLM / FastAPI: + * - Error instances (err.message) + * - { detail: "string" } + * - { detail: [{ msg, loc, type }] } (FastAPI 422) + * - { detail: { error: "string" } } + * - { message: "string" } + * - anything else → JSON.stringify / String() + */ +export function extractErrorMessage(err: unknown): string { + if (err instanceof Error) return err.message; + if (err && typeof err === "object") { + const e = err as Record; + const detail = e.detail; + if (typeof detail === "string") return detail; + if (Array.isArray(detail)) { + return detail.map((d: unknown) => { + if (d && typeof d === "object") { + const item = d as Record; + return typeof item.msg === "string" ? item.msg : JSON.stringify(d); + } + return String(d); + }).join("; "); + } + if (detail && typeof detail === "object") { + const detailObj = detail as Record; + if (typeof detailObj.error === "string") return detailObj.error; + } + if (typeof e.message === "string") return e.message; + return JSON.stringify(err); + } + return String(err); +}