From 443711debd9f5c56fdbcd7ff5817d170a0e276c3 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Mon, 9 Feb 2026 17:53:54 -0800 Subject: [PATCH] fix: backend endpoints --- .../mcp_server/rest_endpoints.py | 63 ++++++-- .../mcp_server/test_rest_endpoints.py | 152 ++++++++++++++++++ 2 files changed, 204 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 65d3a2caa16..e618e88e454 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1,6 +1,6 @@ import importlib from datetime import datetime -from typing import Dict, List, Optional, Union +from typing import Any, Awaitable, Callable, Dict, List, Optional, Union from fastapi import APIRouter, Depends, HTTPException, Query, Request @@ -501,24 +501,50 @@ if MCP_AVAILABLE: NewMCPServerRequest, ) + def _extract_credentials( + request: NewMCPServerRequest, + ) -> tuple: + """ + Extract OAuth credentials from the nested ``request.credentials`` dict. + + Returns: + (client_id, client_secret, scopes) — any value may be ``None``. + """ + creds = request.credentials if isinstance(request.credentials, dict) else {} + client_id: Optional[str] = creds.get("client_id") + client_secret: Optional[str] = creds.get("client_secret") + scopes_raw = creds.get("scopes") + scopes: Optional[List[str]] = scopes_raw if isinstance(scopes_raw, list) else None + return client_id, client_secret, scopes + async def _execute_with_mcp_client( request: NewMCPServerRequest, - operation, + operation: Callable[..., Awaitable[Any]], mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None, oauth2_headers: Optional[Dict[str, str]] = None, raw_headers: Optional[Dict[str, str]] = None, - ): + ) -> dict: """ - Common helper to create MCP client, execute operation, and ensure proper cleanup. + Create a temporary MCP client from *request*, run *operation*, and return the result. + + For M2M OAuth servers (those with ``client_id``, ``client_secret``, and + ``token_url``), the incoming ``oauth2_headers`` are dropped so that + ``resolve_mcp_auth`` can auto-fetch a token via ``client_credentials``. Args: - request: MCP server configuration - operation: Async function that takes a client and returns the operation result + request: MCP server configuration submitted by the UI. + operation: Async callable that receives the created client and returns a result dict. + mcp_auth_header: Pre-resolved credential header (API-key / bearer token). + oauth2_headers: Headers extracted from the incoming request (may contain the + litellm API key — must NOT be forwarded for M2M servers). + raw_headers: Raw request headers forwarded for stdio env construction. Returns: - Operation result or error response + The dict returned by *operation*, or an error dict on failure. """ try: + client_id, client_secret, scopes = _extract_credentials(request) + server_model = MCPServer( server_id=request.server_id or "", name=request.alias or request.server_name or "", @@ -530,14 +556,26 @@ if MCP_AVAILABLE: args=request.args, env=request.env, static_headers=request.static_headers, + client_id=client_id, + client_secret=client_secret, + token_url=request.token_url, + scopes=scopes, + authorization_url=request.authorization_url, + registration_url=request.registration_url, ) stdio_env = global_mcp_server_manager._build_stdio_env( server_model, raw_headers ) + # For M2M OAuth servers, drop the incoming Authorization header so that + # resolve_mcp_auth can auto-fetch a token via client_credentials. + effective_oauth2_headers = ( + None if server_model.has_client_credentials else oauth2_headers + ) + merged_headers = merge_mcp_headers( - extra_headers=oauth2_headers, + extra_headers=effective_oauth2_headers, static_headers=request.static_headers, ) @@ -550,11 +588,14 @@ if MCP_AVAILABLE: return await operation(client) - except Exception as e: - verbose_logger.error(f"Error in MCP operation: {e}", exc_info=True) + except (KeyboardInterrupt, SystemExit): + raise + except BaseException as e: + verbose_logger.error("Error in MCP operation: %s", e, exc_info=True) return { "status": "error", - "message": "An internal error has occurred while testing the MCP server.", + "error": True, + "message": f"Failed to connect to MCP server: {e}", } @router.post("/test/connection", dependencies=[Depends(user_api_key_auth)]) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index f6c8115f7bc..f86f1e448cb 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -155,6 +155,158 @@ class TestExecuteWithMcpClient: } + @pytest.mark.asyncio + async def test_m2m_credentials_forwarded_to_server_model(self, monkeypatch): + """M2M OAuth credentials (client_id, client_secret) from the nested + ``credentials`` dict must be forwarded to the MCPServer model so that + ``has_client_credentials`` returns True and the proxy auto-fetches tokens.""" + captured: dict = {} + + def fake_build_stdio_env(server, raw_headers): + return None + + async def fake_create_client(*args, **kwargs): + captured["server"] = kwargs.get("server") + return object() + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_build_stdio_env", + fake_build_stdio_env, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_create_mcp_client", + fake_create_client, + raising=False, + ) + + async def ok_operation(client): + return {"status": "ok"} + + payload = NewMCPServerRequest( + server_name="m2m-server", + url="https://example.com", + auth_type=MCPAuth.oauth2, + token_url="https://auth.example.com/token", + credentials={ + "client_id": "my-id", + "client_secret": "my-secret", + "scopes": ["read", "write"], + }, + ) + + result = await rest_endpoints._execute_with_mcp_client( + payload, ok_operation + ) + + assert result["status"] == "ok" + server = captured["server"] + assert server.client_id == "my-id" + assert server.client_secret == "my-secret" + assert server.token_url == "https://auth.example.com/token" + assert server.scopes == ["read", "write"] + assert server.has_client_credentials is True + + @pytest.mark.asyncio + async def test_m2m_drops_incoming_oauth2_headers(self, monkeypatch): + """For M2M OAuth servers the incoming Authorization header (which carries + the litellm API key) must NOT be forwarded as extra_headers — otherwise + it overwrites the auto-fetched M2M token.""" + captured: dict = {} + + def fake_build_stdio_env(server, raw_headers): + return None + + async def fake_create_client(*args, **kwargs): + captured["extra_headers"] = kwargs.get("extra_headers") + return object() + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_build_stdio_env", + fake_build_stdio_env, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_create_mcp_client", + fake_create_client, + raising=False, + ) + + async def ok_operation(client): + return {"status": "ok"} + + payload = NewMCPServerRequest( + server_name="m2m-server", + url="https://example.com", + auth_type=MCPAuth.oauth2, + token_url="https://auth.example.com/token", + credentials={ + "client_id": "my-id", + "client_secret": "my-secret", + }, + ) + + incoming_oauth2 = {"Authorization": "Bearer sk-litellm-api-key"} + result = await rest_endpoints._execute_with_mcp_client( + payload, + ok_operation, + oauth2_headers=incoming_oauth2, + ) + + assert result["status"] == "ok" + # The incoming Authorization must be dropped — extra_headers should + # contain no oauth2 headers (only static_headers, which are None here). + assert captured["extra_headers"] is None or "Authorization" not in captured["extra_headers"] + + @pytest.mark.asyncio + async def test_catches_exception_group(self, monkeypatch): + """MCP SDK's anyio TaskGroup raises BaseExceptionGroup which does not + inherit from Exception. The handler must catch it and return an error + dict instead of letting a raw 500 propagate.""" + + def fake_build_stdio_env(server, raw_headers): + return None + + async def fake_create_client(*args, **kwargs): + raise BaseExceptionGroup( + "test group", [RuntimeError("Cancelled via cancel scope")] + ) + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_build_stdio_env", + fake_build_stdio_env, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_create_mcp_client", + fake_create_client, + raising=False, + ) + + async def ok_operation(client): + return {"status": "ok"} + + payload = NewMCPServerRequest( + server_name="bad-server", + url="https://example.com", + auth_type=MCPAuth.none, + ) + + result = await rest_endpoints._execute_with_mcp_client( + payload, ok_operation + ) + + assert result["status"] == "error" + assert result["error"] is True + assert "Failed to connect to MCP server" in result["message"] + + class TestTestConnection: def test_requires_auth_dependency(self): route = _get_route("/mcp-rest/test/connection", "POST")