diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 891b52db7a..4c947b99ba 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1,5 +1,4 @@ import importlib -import traceback from typing import Dict, List, Optional, Union from fastapi import APIRouter, Depends, Query, Request @@ -347,17 +346,16 @@ if MCP_AVAILABLE: except Exception as e: verbose_logger.error(f"Error in MCP operation: {e}", exc_info=True) - stack_trace = traceback.format_exc() return { "status": "error", - "message": f"An internal error has occurred: {str(e)}", - "stack_trace": stack_trace, + "message": "An internal error has occurred while testing the MCP server.", } - @router.post("/test/connection") + @router.post("/test/connection", dependencies=[Depends(user_api_key_auth)]) async def test_connection( request: Request, new_mcp_server_request: NewMCPServerRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ Test if we can connect to the provided MCP server before adding it 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 ce38ee59e4..d9aba73de8 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 @@ -7,6 +7,7 @@ from litellm.proxy._experimental.mcp_server import rest_endpoints from litellm.proxy._experimental.mcp_server.auth import ( user_api_key_auth_mcp as auth_mcp, ) +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy._types import NewMCPServerRequest, UserAPIKeyAuth from litellm.types.mcp import MCPAuth @@ -31,6 +32,60 @@ def _build_request(headers: Optional[Dict[str, str]] = None) -> Request: return Request(scope, receive=receive) +def _get_route(path: str, method: str): + for route in rest_endpoints.router.routes: + if getattr(route, "path", None) == path and method in getattr( + route, "methods", set() + ): + return route + raise AssertionError(f"Route {method} {path} not found") + + +def _route_has_dependency(route, dependency) -> bool: + if any( + getattr(dep, "dependency", None) == dependency + for dep in getattr(route, "dependencies", []) + ): + return True + dependant = getattr(route, "dependant", None) + if dependant is None: + return False + return any(getattr(dep, "call", None) == dependency for dep in dependant.dependencies) + + +@pytest.mark.asyncio +async def test_execute_with_mcp_client_redacts_stack_trace(monkeypatch): + def fake_create_client(*args, **kwargs): + return object() + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_create_mcp_client", + fake_create_client, + ) + + async def failing_operation(client): + raise RuntimeError("boom") + + payload = NewMCPServerRequest( + server_name="example", + url="https://example.com", + auth_type=MCPAuth.none, + ) + + result = await rest_endpoints._execute_with_mcp_client( + payload, failing_operation + ) + + assert result["status"] == "error" + assert "stack_trace" not in result + + +def test_test_connection_requires_auth_dependency(): + route = _get_route("/test/connection", "POST") + assert _route_has_dependency(route, user_api_key_auth) + + @pytest.mark.asyncio async def test_test_tools_list_forwards_mcp_auth_header(monkeypatch): """Ensure credential-based auth forwards the auth_value to the MCP client."""