diff --git a/litellm/proxy/_experimental/mcp_server/openapi_oauth2_endpoints.py b/litellm/proxy/_experimental/mcp_server/openapi_oauth2_endpoints.py index b39093d465e..4de3db81087 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_oauth2_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_oauth2_endpoints.py @@ -16,7 +16,7 @@ import json import secrets import time from typing import Dict, Optional -from urllib.parse import parse_qs, urlencode +from urllib.parse import parse_qs, urlencode, urlparse, urlunparse import httpx from fastapi import APIRouter, Depends, HTTPException, Query, Request @@ -248,9 +248,12 @@ async def openapi_oauth2_connect( if server.scopes: params["scope"] = " ".join(server.scopes) - # Use "&" if the base URL already contains query parameters, otherwise "?" - sep = "&" if "?" in server.authorization_url else "?" - authorization_url = f"{server.authorization_url}{sep}{urlencode(params)}" + # Merge params into the base URL safely using urllib.parse to handle edge cases + # (existing query string, fragments, trailing separators, etc.). + parsed = urlparse(server.authorization_url) + existing_qs = parse_qs(parsed.query, keep_blank_values=True) + merged_qs = {**existing_qs, **{k: [v] for k, v in params.items()}} + authorization_url = urlunparse(parsed._replace(query=urlencode(merged_qs, doseq=True))) server_name = server.server_name or server.name or server_id verbose_proxy_logger.debug( @@ -614,6 +617,21 @@ async def openapi_oauth2_status( ) connected = False + # Seed the cache for the not-connected case so subsequent 2-second polls skip + # the DB during the pre-connection waiting window. When not-connected, we write + # None (the "no credential" sentinel) which _check_byok_credential also uses. + # For connected=True we do NOT write here: the callback's _invalidate_byok_cred_cache + # already cleared the entry, and polling stops immediately after we return True. + if not connected: + try: + from litellm.proxy._experimental.mcp_server.server import ( + _write_byok_cred_cache, + ) + + _write_byok_cred_cache(user_id, server_id, None) + except Exception: + pass # Best-effort; never block the response + return JSONResponse( {"connected": connected, "server_id": server_id, "server_name": server_name} ) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index cdabc653f83..af39c4241fa 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1585,6 +1585,7 @@ if MCP_AVAILABLE: if cached is not None: credential, ts = cached if time.monotonic() - ts < _BYOK_CRED_CACHE_TTL: + _byok_cred_cache.move_to_end(cache_key) # promote to MRU return credential from litellm.proxy._experimental.mcp_server.db import get_user_credential @@ -1657,6 +1658,7 @@ if MCP_AVAILABLE: if cached is not None: cached_cred, ts = cached if time.monotonic() - ts < _BYOK_CRED_CACHE_TTL: + _byok_cred_cache.move_to_end(cache_key) # promote to MRU if cached_cred is None: raise HTTPException( status_code=401,