Fix/non standard mcp url pattern (#19738)

* fix(mcp): Add standard MCP URL pattern support for OAuth discovery (#17272) OAuth discovery endpoints now support both URL patterns: - Standard MCP pattern: /mcp/{server_name} (new) - Legacy LiteLLM pattern: /{server_name}/mcp (backward compatible) The standard pattern is required by MCP-compliant clients like mcp-inspector and VSCode Copilot, which expect resource URLs following the /mcp/{server_name} convention per RFC 9728. Changes: - Add _build_oauth_protected_resource_response() helper - Add oauth_protected_resource_mcp_standard() endpoint - Add oauth_authorization_server_mcp_standard() endpoint - Keep legacy endpoints for backward compatibility - Add tests for both URL patterns Fixes #17272 * fix(mcp): Add standard MCP URL pattern support for OAuth discovery (#17272) OAuth discovery endpoints now support both URL patterns: - Standard MCP pattern: /mcp/{server_name} (new) - Legacy LiteLLM pattern: /{server_name}/mcp (backward compatible) The standard pattern is required by MCP-compliant clients like mcp-inspector and VSCode Copilot, which expect resource URLs following the /mcp/{server_name} convention per RFC 9728. Changes: - Add _build_oauth_protected_resource_response() helper - Add oauth_protected_resource_mcp_standard() endpoint - Add oauth_authorization_server_mcp_standard() endpoint - Keep legacy endpoints for backward compatibility - Add tests for both URL patterns Fixes #17272 * Test was relocated * refactor(mcp): Extract helper methods from run_with_session to fix PLR0915 Split the large run_with_session method (55 statements) into smaller helper methods to satisfy ruff's PLR0915 rule (max 50 statements): - _create_transport_context(): Creates transport based on type - _execute_session_operation(): Handles session lifecycle Also changed cleanup exception handling from Exception to BaseException to properly catch asyncio.CancelledError (which is a BaseException subclass in Python 3.8+). Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> * test(mcp): Fix flaky test by mocking health_check_server The test_mcp_server_manager_config_integration_with_database test was making real network calls to fake URLs which caused timeouts and CancelledError exceptions. Fixed by mocking health_check_server to return a proper LiteLLM_MCPServerTable object instead of making network calls. * test(mcp): Fix skip condition to properly detect claude model names The skip condition for missing API keys was checking for "anthropic" in the model name, but the test uses "claude-haiku-4-5" which doesn't match. Updated to check for both "anthropic" and "claude" model patterns. Also added skip condition for OpenAI models when OPENAI_API_KEY is not set. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> * test(mcp): Fix skip condition to properly detect claude model names The skip condition for missing API keys was checking for "anthropic" in the model name, but the test uses "claude-haiku-4-5" which doesn't match. Updated to check for both "anthropic" and "claude" model patterns. Also added skip condition for OpenAI models when OPENAI_API_KEY is not set. --------- Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
2026-08-14 02:22:54 +00:00 · 2026-01-25 23:00:53 -08:00
parent 582d324a76
commit 5666c725ce
5 changed files with 1361 additions and 105 deletions
@@ -4,7 +4,7 @@ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers.

 import asyncio
 import base64
-from typing import Awaitable, Callable, Dict, List, Optional, TypeVar, Union
+from typing import Any, Awaitable, Callable, Dict, List, Optional, Tuple, TypeVar, Union

 import httpx
 from mcp import ClientSession, ReadResourceResult, Resource, StdioServerParameters
@@ -74,97 +74,102 @@ class MCPClient:
        if auth_value:
            self.update_auth_value(auth_value)

+    def _create_transport_context(
+        self,
+    ) -> Tuple[Any, Optional[httpx.AsyncClient]]:
+        """
+        Create the appropriate transport context based on transport type.
+
+        Returns:
+            Tuple of (transport_context, http_client).
+            http_client is only set for HTTP transport and needs cleanup.
+        """
+        http_client: Optional[httpx.AsyncClient] = None
+
+        if self.transport_type == MCPTransport.stdio:
+            if not self.stdio_config:
+                raise ValueError("stdio_config is required for stdio transport")
+            server_params = StdioServerParameters(
+                command=self.stdio_config.get("command", ""),
+                args=self.stdio_config.get("args", []),
+                env=self.stdio_config.get("env", {}),
+            )
+            return stdio_client(server_params), None
+
+        if self.transport_type == MCPTransport.sse:
+            headers = self._get_auth_headers()
+            httpx_client_factory = self._create_httpx_client_factory()
+            return sse_client(
+                url=self.server_url,
+                timeout=self.timeout,
+                headers=headers,
+                httpx_client_factory=httpx_client_factory,
+            ), None
+
+        # HTTP transport (default)
+        headers = self._get_auth_headers()
+        httpx_client_factory = self._create_httpx_client_factory()
+        verbose_logger.debug(
+            "litellm headers for streamable_http_client: %s", headers
+        )
+        http_client = httpx_client_factory(
+            headers=headers,
+            timeout=httpx.Timeout(self.timeout),
+        )
+        transport_ctx = streamable_http_client(
+            url=self.server_url,
+            http_client=http_client,
+        )
+        return transport_ctx, http_client
+
+    async def _execute_session_operation(
+        self,
+        transport_ctx: Any,
+        operation: Callable[[ClientSession], Awaitable[TSessionResult]],
+    ) -> TSessionResult:
+        """
+        Execute an operation within a transport and session context.
+
+        Handles entering/exiting contexts and running the operation.
+        """
+        transport = await transport_ctx.__aenter__()
+        try:
+            read_stream, write_stream = transport[0], transport[1]
+            session_ctx = ClientSession(read_stream, write_stream)
+            session = await session_ctx.__aenter__()
+            try:
+                await session.initialize()
+                return await operation(session)
+            finally:
+                try:
+                    await session_ctx.__aexit__(None, None, None)
+                except BaseException as e:
+                    verbose_logger.debug(f"Error during session context exit: {e}")
+        finally:
+            try:
+                await transport_ctx.__aexit__(None, None, None)
+            except BaseException as e:
+                verbose_logger.debug(f"Error during transport context exit: {e}")
+
    async def run_with_session(
        self, operation: Callable[[ClientSession], Awaitable[TSessionResult]]
    ) -> TSessionResult:
        """Open a session, run the provided coroutine, and clean up."""
-        transport_ctx = None
        http_client: Optional[httpx.AsyncClient] = None
-        transport = None
-        session_ctx = None
-
        try:
-            if self.transport_type == MCPTransport.stdio:
-                if not self.stdio_config:
-                    raise ValueError("stdio_config is required for stdio transport")
-
-                server_params = StdioServerParameters(
-                    command=self.stdio_config.get("command", ""),
-                    args=self.stdio_config.get("args", []),
-                    env=self.stdio_config.get("env", {}),
-                )
-                transport_ctx = stdio_client(server_params)
-            elif self.transport_type == MCPTransport.sse:
-                headers = self._get_auth_headers()
-                httpx_client_factory = self._create_httpx_client_factory()
-                transport_ctx = sse_client(
-                    url=self.server_url,
-                    timeout=self.timeout,
-                    headers=headers,
-                    httpx_client_factory=httpx_client_factory,
-                )
-            else:
-                headers = self._get_auth_headers()
-                httpx_client_factory = self._create_httpx_client_factory()
-                verbose_logger.debug(
-                    "litellm headers for streamable_http_client: %s", headers
-                )
-                http_client = httpx_client_factory(
-                    headers=headers,
-                    timeout=httpx.Timeout(self.timeout),
-                )
-                transport_ctx = streamable_http_client(
-                    url=self.server_url,
-                    http_client=http_client,
-                )
-
-            if transport_ctx is None:
-                raise RuntimeError("Failed to create transport context")
-
-            # Enter transport context
-            transport = await transport_ctx.__aenter__()
-            try:
-                read_stream, write_stream = transport[0], transport[1]
-                session_ctx = ClientSession(read_stream, write_stream)
-                
-                # Enter session context
-                session = await session_ctx.__aenter__()
-                try:
-                    await session.initialize()
-                    result = await operation(session)
-                    return result
-                finally:
-                    # Ensure session context is properly exited
-                    if session_ctx is not None:
-                        try:
-                            await session_ctx.__aexit__(None, None, None)
-                        except Exception as e:
-                            verbose_logger.debug(
-                                f"Error during session context exit: {e}"
-                            )
-            finally:
-                # Ensure transport context is properly exited
-                if transport_ctx is not None:
-                    try:
-                        await transport_ctx.__aexit__(None, None, None)
-                    except Exception as e:
-                        verbose_logger.debug(
-                            f"Error during transport context exit: {e}"
-                        )
+            transport_ctx, http_client = self._create_transport_context()
+            return await self._execute_session_operation(transport_ctx, operation)
        except Exception:
            verbose_logger.warning(
                "MCP client run_with_session failed for %s", self.server_url or "stdio"
            )
            raise
        finally:
-            # Always clean up http_client if it was created
            if http_client is not None:
                try:
                    await http_client.aclose()
-                except Exception as e:
-                    verbose_logger.debug(
-                        f"Error during http_client cleanup: {e}"
-                    )
+                except BaseException as e:
+                    verbose_logger.debug(f"Error during http_client cleanup: {e}")

    def update_auth_value(self, mcp_auth_value: Union[str, Dict[str, str]]):
        """
@@ -387,25 +387,57 @@ async def callback(code: str, state: str):
    1. Try resource_metadata from WWW-Authenticate header (if present)
    2. Fall back to path-based well-known URI: /.well-known/oauth-protected-resource/{path}
    (
-    If the resource identifier value contains a path or query component, any terminating slash (/) 
-    following the host component MUST be removed before inserting /.well-known/ and the well-known 
-    URI path suffix between the host component and the path(include root path) and/or query components. 
+    If the resource identifier value contains a path or query component, any terminating slash (/)
+    following the host component MUST be removed before inserting /.well-known/ and the well-known
+    URI path suffix between the host component and the path(include root path) and/or query components.
    https://datatracker.ietf.org/doc/html/rfc9728#section-3.1)
    3. Fall back to root-based well-known URI: /.well-known/oauth-protected-resource
+
+    Dual Pattern Support:
+    - Standard MCP pattern: /mcp/{server_name} (recommended, used by mcp-inspector, VSCode Copilot)
+    - LiteLLM legacy pattern: /{server_name}/mcp (backward compatibility)
+
+    The resource URL returned matches the pattern used in the discovery request.
 """
-@router.get(f"/.well-known/oauth-protected-resource{'' if get_server_root_path() == '/' else get_server_root_path()}/{{mcp_server_name}}/mcp")
-@router.get("/.well-known/oauth-protected-resource")
-async def oauth_protected_resource_mcp(
-    request: Request, mcp_server_name: Optional[str] = None
-):
+
+
+def _build_oauth_protected_resource_response(
+    request: Request,
+    mcp_server_name: Optional[str],
+    use_standard_pattern: bool,
+) -> dict:
+    """
+    Build OAuth protected resource response with the appropriate URL pattern.
+
+    Args:
+        request: FastAPI Request object
+        mcp_server_name: Name of the MCP server
+        use_standard_pattern: If True, use /mcp/{server_name} pattern;
+                             if False, use /{server_name}/mcp pattern
+
+    Returns:
+        OAuth protected resource metadata dict
+    """
    from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
        global_mcp_server_manager,
    )
-    # Get the correct base URL considering X-Forwarded-* headers
+
    request_base_url = get_request_base_url(request)
    mcp_server: Optional[MCPServer] = None
    if mcp_server_name:
        mcp_server = global_mcp_server_manager.get_mcp_server_by_name(mcp_server_name)
+
+    # Build resource URL based on the pattern
+    if mcp_server_name:
+        if use_standard_pattern:
+            # Standard MCP pattern: /mcp/{server_name}
+            resource_url = f"{request_base_url}/mcp/{mcp_server_name}"
+        else:
+            # LiteLLM legacy pattern: /{server_name}/mcp
+            resource_url = f"{request_base_url}/{mcp_server_name}/mcp"
+    else:
+        resource_url = f"{request_base_url}/mcp"
+
    return {
        "authorization_servers": [
            (
@@ -414,14 +446,55 @@ async def oauth_protected_resource_mcp(
                else f"{request_base_url}"
            )
        ],
-        "resource": (
-            f"{request_base_url}/{mcp_server_name}/mcp"
-            if mcp_server_name
-            else f"{request_base_url}/mcp"
-        ),  # this is what Claude will call
+        "resource": resource_url,
        "scopes_supported": mcp_server.scopes if mcp_server else [],
    }

+
+# Standard MCP pattern: /.well-known/oauth-protected-resource/mcp/{server_name}
+# This is the pattern expected by standard MCP clients (mcp-inspector, VSCode Copilot)
+@router.get(f"/.well-known/oauth-protected-resource{'' if get_server_root_path() == '/' else get_server_root_path()}/mcp/{{mcp_server_name}}")
+async def oauth_protected_resource_mcp_standard(
+    request: Request, mcp_server_name: str
+):
+    """
+    OAuth protected resource discovery endpoint using standard MCP URL pattern.
+
+    Standard pattern: /mcp/{server_name}
+    Discovery path: /.well-known/oauth-protected-resource/mcp/{server_name}
+
+    This endpoint is compliant with MCP specification and works with standard
+    MCP clients like mcp-inspector and VSCode Copilot.
+    """
+    return _build_oauth_protected_resource_response(
+        request=request,
+        mcp_server_name=mcp_server_name,
+        use_standard_pattern=True,
+    )
+
+
+# LiteLLM legacy pattern: /.well-known/oauth-protected-resource/{server_name}/mcp
+# Kept for backward compatibility with existing deployments
+@router.get(f"/.well-known/oauth-protected-resource{'' if get_server_root_path() == '/' else get_server_root_path()}/{{mcp_server_name}}/mcp")
+@router.get("/.well-known/oauth-protected-resource")
+async def oauth_protected_resource_mcp(
+    request: Request, mcp_server_name: Optional[str] = None
+):
+    """
+    OAuth protected resource discovery endpoint using LiteLLM legacy URL pattern.
+
+    Legacy pattern: /{server_name}/mcp
+    Discovery path: /.well-known/oauth-protected-resource/{server_name}/mcp
+
+    This endpoint is kept for backward compatibility. New integrations should
+    use the standard MCP pattern (/mcp/{server_name}) instead.
+    """
+    return _build_oauth_protected_resource_response(
+        request=request,
+        mcp_server_name=mcp_server_name,
+        use_standard_pattern=False,
+    )
+
 """
    https://datatracker.ietf.org/doc/html/rfc8414#section-3.1
    RFC 8414: Path-aware OAuth discovery
@@ -430,15 +503,26 @@ async def oauth_protected_resource_mcp(
    the well-known URI suffix between the host component and the path(include root path)
    component.
 """
-@router.get(f"/.well-known/oauth-authorization-server{'' if get_server_root_path() == '/' else get_server_root_path()}/{{mcp_server_name}}")
-@router.get("/.well-known/oauth-authorization-server")
-async def oauth_authorization_server_mcp(
-    request: Request, mcp_server_name: Optional[str] = None
-):
+
+
+def _build_oauth_authorization_server_response(
+    request: Request,
+    mcp_server_name: Optional[str],
+) -> dict:
+    """
+    Build OAuth authorization server metadata response.
+
+    Args:
+        request: FastAPI Request object
+        mcp_server_name: Name of the MCP server
+
+    Returns:
+        OAuth authorization server metadata dict
+    """
    from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
        global_mcp_server_manager,
    )
-    # Get the correct base URL considering X-Forwarded-* headers
+
    request_base_url = get_request_base_url(request)

    authorization_endpoint = (
@@ -470,18 +554,58 @@ async def oauth_authorization_server_mcp(
    }


+# Standard MCP pattern: /.well-known/oauth-authorization-server/mcp/{server_name}
+@router.get(f"/.well-known/oauth-authorization-server{'' if get_server_root_path() == '/' else get_server_root_path()}/mcp/{{mcp_server_name}}")
+async def oauth_authorization_server_mcp_standard(
+    request: Request, mcp_server_name: str
+):
+    """
+    OAuth authorization server discovery endpoint using standard MCP URL pattern.
+
+    Standard pattern: /mcp/{server_name}
+    Discovery path: /.well-known/oauth-authorization-server/mcp/{server_name}
+    """
+    return _build_oauth_authorization_server_response(
+        request=request,
+        mcp_server_name=mcp_server_name,
+    )
+
+
+# LiteLLM legacy pattern and root endpoint
+@router.get(f"/.well-known/oauth-authorization-server{'' if get_server_root_path() == '/' else get_server_root_path()}/{{mcp_server_name}}")
+@router.get("/.well-known/oauth-authorization-server")
+async def oauth_authorization_server_mcp(
+    request: Request, mcp_server_name: Optional[str] = None
+):
+    """
+    OAuth authorization server discovery endpoint.
+
+    Supports both legacy pattern (/{server_name}) and root endpoint.
+    """
+    return _build_oauth_authorization_server_response(
+        request=request,
+        mcp_server_name=mcp_server_name,
+    )
+
+
 # Alias for standard OpenID discovery
@router.get("/.well-known/openid-configuration")
 async def openid_configuration(request: Request):
    return await oauth_authorization_server_mcp(request)


+# Additional legacy pattern support
@router.get("/.well-known/oauth-authorization-server/{mcp_server_name}/mcp")
-@router.get("/.well-known/oauth-authorization-server")
-async def oauth_authorization_server_root(
-    request: Request, mcp_server_name: Optional[str] = None
+async def oauth_authorization_server_legacy(
+    request: Request, mcp_server_name: str
 ):
-    return await oauth_authorization_server_mcp(request, mcp_server_name)
+    """
+    OAuth authorization server discovery for legacy /{server_name}/mcp pattern.
+    """
+    return _build_oauth_authorization_server_response(
+        request=request,
+        mcp_server_name=mcp_server_name,
+    )


@router.post("/{mcp_server_name}/register")
@@ -683,9 +683,11 @@ async def test_streaming_responses_api_with_mcp_tools(

    Return the user the result of request 2
    """
-    # Skip test if ANTHROPIC_API_KEY is not set for anthropic models
-    if "anthropic" in model.lower() and not os.getenv("ANTHROPIC_API_KEY"):
+    # Skip test if required API keys are not set
+    if ("anthropic" in model.lower() or "claude" in model.lower()) and not os.getenv("ANTHROPIC_API_KEY"):
        pytest.skip("ANTHROPIC_API_KEY not set, skipping anthropic model test")
+    if ("gpt" in model.lower() or "openai" in model.lower()) and not os.getenv("OPENAI_API_KEY"):
+        pytest.skip("OPENAI_API_KEY not set, skipping openai model test")
    
    from unittest.mock import AsyncMock, patch
    
@@ -1105,9 +1105,26 @@ async def test_mcp_server_manager_config_integration_with_database():

    test_manager.get_allowed_mcp_servers = mock_get_allowed_servers

-    # Test the method (this tests our second fix)
-    import asyncio
+    # Mock health_check_server to avoid real network calls that timeout
+    async def mock_health_check(server_id: str, mcp_auth_header=None):
+        server = test_manager.get_mcp_server_by_id(server_id)
+        if not server:
+            return None
+        return LiteLLM_MCPServerTable(
+            server_id=server_id,
+            server_name=server.name,
+            url=server.url,
+            transport=server.transport,
+            description=server.mcp_info.get("description") if server.mcp_info else None,
+            mcp_access_groups=server.access_groups,
+            status="healthy",
+            last_health_check=datetime.datetime.now(),
+            mcp_info=server.mcp_info,
+        )

+    test_manager.health_check_server = mock_health_check
+
+    # Test the method (this tests our second fix)
    servers_list = await test_manager.get_all_mcp_servers_with_health_and_teams(
        user_api_key_auth=mock_user_auth
    )