diff --git a/.circleci/config.yml b/.circleci/config.yml index 133a7184f9..e03d108628 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1153,7 +1153,7 @@ jobs: pip install "pytest-asyncio==0.21.1" pip install "respx==0.22.0" pip install "pydantic==2.10.2" - pip install "mcp==1.10.1" + pip install "mcp==1.21.2" # Run pytest and generate JUnit XML report - run: name: Run tests diff --git a/tests/mcp_tests/mcp_server.py b/tests/mcp_tests/mcp_server.py index 99a67edd02..bc6accbb72 100644 --- a/tests/mcp_tests/mcp_server.py +++ b/tests/mcp_tests/mcp_server.py @@ -1,9 +1,33 @@ # math_server.py +import argparse +import os + from mcp.server.fastmcp import FastMCP mcp = FastMCP("Math") +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="MCP math test server") + parser.add_argument( + "--transport", + default=os.getenv("MCP_TRANSPORT", "stdio"), + help="Transport to use (stdio or http)", + ) + parser.add_argument( + "--host", + default=os.getenv("MCP_HOST", "127.0.0.1"), + help="Host to bind when serving over HTTP", + ) + parser.add_argument( + "--port", + type=int, + default=int(os.getenv("MCP_PORT", "0")), + help="Port to bind when serving over HTTP", + ) + return parser.parse_args() + + @mcp.tool() def add(a: int, b: int) -> int: """Add two numbers""" @@ -16,5 +40,24 @@ def multiply(a: int, b: int) -> int: return a * b +def main() -> None: + args = _parse_args() + transport = (args.transport or "stdio").lower() + + if transport == "stdio": + mcp.run(transport="stdio") + return + + if transport in {"http", "streamable_http", "streamable-http"}: + if args.port <= 0: + raise ValueError("HTTP transport requires a valid --port value") + mcp.settings.host = args.host + mcp.settings.port = args.port + mcp.run(transport="streamable-http") + return + + raise ValueError(f"Unsupported transport: {transport}") + + if __name__ == "__main__": - mcp.run(transport="stdio") + main() diff --git a/tests/mcp_tests/test_configs/test_config_mcp_e2e.yaml b/tests/mcp_tests/test_configs/test_config_mcp_e2e.yaml new file mode 100644 index 0000000000..37d5359e3c --- /dev/null +++ b/tests/mcp_tests/test_configs/test_config_mcp_e2e.yaml @@ -0,0 +1,23 @@ +general_settings: + master_key: sk-1234 + +litellm_settings: + drop_params: true + +model_list: + - model_name: openai-gpt-4o-mini + litellm_params: + model: gpt-4o-mini + - model_name: anthropic-claude-haiku-4-5 + litellm_params: + model: anthropic/claude-haiku-4-5 + +mcp_servers: + math_stdio: + transport: stdio + command: python3 + args: + - tests/mcp_tests/mcp_server.py + math_streamable_http: + transport: http + url: http://127.0.0.1:0/mcp diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 66dd2bdc6b..e8a1231c6f 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -1001,34 +1001,39 @@ async def test_mcp_server_manager_access_groups_from_config(): MCPRequestHandler, ) - # Patch global_mcp_server_manager for this test + # Patch global_mcp_server_manager for this test and restore afterwards to + # avoid leaking state into other tests (e.g. the proxy MCP e2e suite). import litellm.proxy._experimental.mcp_server.mcp_server_manager as mcp_server_manager_mod + original_manager = mcp_server_manager_mod.global_mcp_server_manager mcp_server_manager_mod.global_mcp_server_manager = test_manager - # Should find config_server for group-a, both for group-b, other_server for group-c - import asyncio + try: + # Should find config_server for group-a, both for group-b, other_server for group-c + import asyncio - server_ids_a = await MCPRequestHandler._get_mcp_servers_from_access_groups([ - "group-a" - ]) - server_ids_b = await MCPRequestHandler._get_mcp_servers_from_access_groups([ - "group-b" - ]) - server_ids_c = await MCPRequestHandler._get_mcp_servers_from_access_groups([ - "group-c" - ]) - assert any(config_server.server_id == sid for sid in server_ids_a) - assert set(server_ids_b) == set( - [ - s.server_id + server_ids_a = await MCPRequestHandler._get_mcp_servers_from_access_groups([ + "group-a" + ]) + server_ids_b = await MCPRequestHandler._get_mcp_servers_from_access_groups([ + "group-b" + ]) + server_ids_c = await MCPRequestHandler._get_mcp_servers_from_access_groups([ + "group-c" + ]) + assert any(config_server.server_id == sid for sid in server_ids_a) + assert set(server_ids_b) == set( + [ + s.server_id + for s in test_manager.config_mcp_servers.values() + if "group-b" in s.access_groups + ] + ) + assert any( + s.name == "other_server" and s.server_id in server_ids_c for s in test_manager.config_mcp_servers.values() - if "group-b" in s.access_groups - ] - ) - assert any( - s.name == "other_server" and s.server_id in server_ids_c - for s in test_manager.config_mcp_servers.values() - ) + ) + finally: + mcp_server_manager_mod.global_mcp_server_manager = original_manager async def test_mcp_server_manager_config_integration_with_database(): diff --git a/tests/mcp_tests/test_proxy_mcp_e2e.py b/tests/mcp_tests/test_proxy_mcp_e2e.py new file mode 100644 index 0000000000..2b8cde5471 --- /dev/null +++ b/tests/mcp_tests/test_proxy_mcp_e2e.py @@ -0,0 +1,236 @@ +import asyncio +import os +import socket +import subprocess +import sys +import threading +import time +import typing +from pathlib import Path + +import pytest +import uvicorn +import yaml +from mcp import ClientSession +from mcp.client.streamable_http import streamablehttp_client + +from litellm.proxy.proxy_server import ( + app as proxy_app, + cleanup_router_config_variables, + initialize, +) + + +CONFIG_TEMPLATE_PATH = Path("tests/mcp_tests/test_configs/test_config_mcp_e2e.yaml") +MCP_SERVER_SCRIPT = Path("tests/mcp_tests/mcp_server.py") +PROJECT_ROOT = Path(__file__).resolve().parents[2] +PROXY_START_TIMEOUT = 30 + + +PROXY_AUTHORIZATION_HEADER = "Bearer sk-1234" + + +@pytest.fixture(scope="session", autouse=True) +def _clear_proxy_database_env() -> typing.Iterator[None]: + """Ensure local proxy DB settings don't leak into tests.""" + mp = pytest.MonkeyPatch() + mp.delenv("DATABASE_URL", raising=False) + try: + yield + finally: + mp.undo() + + +def _initialize_proxy(config_path: str) -> None: + cleanup_router_config_variables() + asyncio.run(initialize(config=config_path, debug=True)) + + +def _start_proxy_server(config_path: str) -> tuple[str, uvicorn.Server, threading.Thread, socket.socket]: + _initialize_proxy(config_path) + + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + sock.bind(("127.0.0.1", 0)) + host, port = sock.getsockname() + + config = uvicorn.Config(proxy_app, host=host, port=port, log_level="warning") + server = uvicorn.Server(config) + + def _run() -> None: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + loop.run_until_complete(server.serve(sockets=[sock])) + + thread = threading.Thread(target=_run, daemon=True) + thread.start() + + start_time = time.time() + while not server.started: + if not thread.is_alive(): + raise RuntimeError("Proxy server failed to start") + if time.time() - start_time > PROXY_START_TIMEOUT: + raise TimeoutError("Proxy server did not start in time") + time.sleep(0.05) + + return f"http://{host}:{port}", server, thread, sock + + +@pytest.fixture(scope="session") +def math_streamable_http_server() -> str: + host = "127.0.0.1" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind((host, 0)) + _, port = sock.getsockname() + + cmd = [ + sys.executable, + str(MCP_SERVER_SCRIPT), + "--transport", + "http", + "--host", + host, + "--port", + str(port), + ] + + env = os.environ.copy() + server_process = subprocess.Popen( + cmd, + cwd=str(PROJECT_ROOT), + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + + start_time = time.time() + while True: + if server_process.poll() is not None: + stdout, stderr = server_process.communicate() + raise RuntimeError( + f"Streamable HTTP MCP server exited early.\nSTDOUT: {stdout.decode()}\nSTDERR: {stderr.decode()}" + ) + try: + with socket.create_connection((host, port), timeout=0.1): + break + except OSError: + if time.time() - start_time > PROXY_START_TIMEOUT: + server_process.terminate() + raise TimeoutError("Streamable HTTP MCP server did not start in time") + time.sleep(0.05) + + yield f"http://{host}:{port}" + + server_process.terminate() + try: + server_process.wait(timeout=5) + except subprocess.TimeoutExpired: + server_process.kill() + + +@pytest.fixture(scope="session") +def proxy_server_url( + tmp_path_factory: pytest.TempPathFactory, math_streamable_http_server: str +): + config_dir = tmp_path_factory.mktemp("mcp_e2e") + config_path = config_dir / "config.yaml" + config = yaml.safe_load(CONFIG_TEMPLATE_PATH.read_text()) + config["mcp_servers"]["math_streamable_http"][ + "url" + ] = f"{math_streamable_http_server}/mcp" + config_path.write_text(yaml.safe_dump(config)) + + server_url, server, thread, sock = _start_proxy_server(str(config_path)) + + yield server_url + + server.should_exit = True + thread.join(timeout=10) + sock.close() + + +class TestProxyMcpSimpleConnections: + @pytest.mark.asyncio + async def test_proxy_mcp_stdio_roundtrip(self, proxy_server_url: str) -> None: + async with asyncio.timeout(20): + async with streamablehttp_client( + url=f"{proxy_server_url}/mcp", + headers={ + "Authorization": PROXY_AUTHORIZATION_HEADER, + "x-mcp-servers": "math_stdio", + }, + ) as (read, write, _get_session_id): + async with ClientSession(read, write) as session: + await session.initialize() + tools_result = await session.list_tools() + assert any(tool.name.endswith("add") for tool in tools_result.tools) + + result = await session.call_tool( + "add", arguments={"a": 3, "b": 4} + ) + assert result.content + first_content = result.content[0] + text = getattr(first_content, "text", None) + assert text == "7" + + @pytest.mark.asyncio + async def test_proxy_mcp_streamable_http_roundtrip( + self, proxy_server_url: str + ) -> None: + async with asyncio.timeout(20): + async with streamablehttp_client( + url=f"{proxy_server_url}/mcp", + headers={ + "Authorization": PROXY_AUTHORIZATION_HEADER, + "x-mcp-servers": "math_streamable_http", + }, + ) as (read, write, _get_session_id): + async with ClientSession(read, write) as session: + await session.initialize() + tools_result = await session.list_tools() + assert any(tool.name.endswith("add") for tool in tools_result.tools) + + result = await session.call_tool( + "add", arguments={"a": 5, "b": 6} + ) + assert result.content + first_content = result.content[0] + text = getattr(first_content, "text", None) + assert text == "11" + + @pytest.mark.asyncio + async def test_proxy_mcp_lists_all_servers_without_header( + self, proxy_server_url: str + ) -> None: + async with asyncio.timeout(20): + async with streamablehttp_client( + url=f"{proxy_server_url}/mcp", + headers={"Authorization": PROXY_AUTHORIZATION_HEADER}, + ) as (read, write, _get_session_id): + async with ClientSession(read, write) as session: + await session.initialize() + tools_result = await session.list_tools() + tool_names = {tool.name for tool in tools_result.tools} + expected_tool_names = { + "math_stdio-add", + "math_stdio-multiply", + "math_streamable_http-add", + "math_streamable_http-multiply", + } + assert expected_tool_names <= tool_names + + async def _call_and_get_text( + tool_name: str, *, a: int, b: int + ) -> str | None: + result = await session.call_tool(tool_name, arguments={"a": a, "b": b}) + assert result.content + first_content = result.content[0] + return getattr(first_content, "text", None) + + stdio_result = await _call_and_get_text( + "math_stdio-add", a=2, b=3 + ) + streamable_result = await _call_and_get_text( + "math_streamable_http-add", a=4, b=5 + ) + assert stdio_result == "5" + assert streamable_result == "9"