diff --git a/.circleci/config.yml b/.circleci/config.yml index 6e51e9ac71..a1908d476b 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1238,6 +1238,7 @@ jobs: pip install "pytest-asyncio==0.21.1" pip install "pytest-cov==5.0.0" pip install "tomli==2.2.1" + pip install "mcp==1.9.3" - run: name: Run tests command: | diff --git a/litellm/constants.py b/litellm/constants.py index c4adb09dc6..c96e1dfaba 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -714,6 +714,7 @@ BEDROCK_AGENT_RUNTIME_PASS_THROUGH_ROUTES = [ "generateQuery/", "optimize-prompt/", ] +BASE_MCP_ROUTE = "/mcp" BATCH_STATUS_POLL_INTERVAL_SECONDS = int( os.getenv("BATCH_STATUS_POLL_INTERVAL_SECONDS", 3600) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 11a52c6bda..d61ae75716 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -3,12 +3,12 @@ LiteLLM MCP Server Routes """ import asyncio -from typing import Any, Dict, List, Optional, Union +import contextlib +from typing import Any, AsyncIterator, Dict, List, Optional, Union -from anyio import BrokenResourceError -from fastapi import APIRouter, Depends, HTTPException, Query, Request -from fastapi.responses import StreamingResponse -from pydantic import ConfigDict, ValidationError +from fastapi import APIRouter, Depends, FastAPI, HTTPException, Query, Request +from pydantic import ConfigDict +from starlette.types import Receive, Scope, Send from litellm._logging import verbose_logger from litellm.constants import MCP_TOOL_NAME_PREFIX @@ -24,6 +24,10 @@ router = APIRouter( tags=["mcp"], ) +LITELLM_MCP_SERVER_NAME = "litellm-mcp-server" +LITELLM_MCP_SERVER_VERSION = "1.0.0" +LITELLM_MCP_SERVER_DESCRIPTION = "MCP Server for LiteLLM" + # Check if MCP is available # "mcp" requires python 3.10 or higher, but several litellm users use python 3.8 # We're making this conditional import to avoid breaking users who use python 3.8. @@ -48,17 +52,25 @@ def get_mcp_server_enabled() -> Dict[str, bool]: return {"enabled": MCP_AVAILABLE} +# Global variables to track initialization +_SESSION_MANAGERS_INITIALIZED = False +_SESSION_MANAGER_TASK = None + if MCP_AVAILABLE: - from mcp.server import NotificationOptions, Server - from mcp.server.models import InitializationOptions + from mcp.server import Server + from mcp.server.streamable_http_manager import StreamableHTTPSessionManager from mcp.types import EmbeddedResource as MCPEmbeddedResource from mcp.types import ImageContent as MCPImageContent from mcp.types import TextContent as MCPTextContent from mcp.types import Tool as MCPTool - from .mcp_server_manager import global_mcp_server_manager - from .sse_transport import SseServerTransport - from .tool_registry import global_mcp_tool_registry + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._experimental.mcp_server.sse_transport import SseServerTransport + from litellm.proxy._experimental.mcp_server.tool_registry import ( + global_mcp_tool_registry, + ) ###################################################### ############ MCP Tools List REST API Response Object # @@ -76,9 +88,80 @@ if MCP_AVAILABLE: ######################################################## ############ Initialize the MCP Server ################# ######################################################## - server: Server = Server("litellm-mcp-server") + server: Server = Server( + name=LITELLM_MCP_SERVER_NAME, + version=LITELLM_MCP_SERVER_VERSION, + ) sse: SseServerTransport = SseServerTransport("/mcp/sse/messages") + # Create session managers + session_manager = StreamableHTTPSessionManager( + app=server, + event_store=None, + json_response=True, # Use JSON responses instead of SSE by default + stateless=True, + ) + + # Create SSE session manager + sse_session_manager = StreamableHTTPSessionManager( + app=server, + event_store=None, + json_response=False, # Use SSE responses for this endpoint + stateless=True, + ) + + async def initialize_session_managers(): + """Initialize the session managers. Can be called from main app lifespan.""" + global _SESSION_MANAGERS_INITIALIZED, _SESSION_MANAGER_TASK + + if _SESSION_MANAGERS_INITIALIZED: + return + + verbose_logger.info("Initializing MCP session managers...") + + # Create a task to run the session managers + async def run_session_managers(): + async with session_manager.run(): + async with sse_session_manager.run(): + verbose_logger.info( + "MCP Server started with StreamableHTTP and SSE session managers!" + ) + try: + # Keep running until cancelled + while True: + await asyncio.sleep(1) + except asyncio.CancelledError: + verbose_logger.info("MCP session managers shutting down...") + raise + + _SESSION_MANAGER_TASK = asyncio.create_task(run_session_managers()) + _SESSION_MANAGERS_INITIALIZED = True + verbose_logger.info("MCP session managers initialization completed!") + + async def shutdown_session_managers(): + """Shutdown the session managers.""" + global _SESSION_MANAGERS_INITIALIZED, _SESSION_MANAGER_TASK + + if _SESSION_MANAGER_TASK and not _SESSION_MANAGER_TASK.done(): + verbose_logger.info("Shutting down MCP session managers...") + _SESSION_MANAGER_TASK.cancel() + try: + await _SESSION_MANAGER_TASK + except asyncio.CancelledError: + pass + + _SESSION_MANAGERS_INITIALIZED = False + _SESSION_MANAGER_TASK = None + + @contextlib.asynccontextmanager + async def lifespan(app) -> AsyncIterator[None]: + """Application lifespan context manager.""" + await initialize_session_managers() + try: + yield + finally: + await shutdown_session_managers() + ######################################################## ############### MCP Server Routes ####################### ######################################################## @@ -157,15 +240,15 @@ if MCP_AVAILABLE: "litellm_logging_obj", None ) if litellm_logging_obj: - litellm_logging_obj.model_call_details[ - "mcp_tool_call_metadata" - ] = standard_logging_mcp_tool_call - litellm_logging_obj.model_call_details[ - "model" - ] = f"{MCP_TOOL_NAME_PREFIX}: {standard_logging_mcp_tool_call.get('name') or ''}" - litellm_logging_obj.model_call_details[ - "custom_llm_provider" - ] = standard_logging_mcp_tool_call.get("mcp_server_name") + litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = ( + standard_logging_mcp_tool_call + ) + litellm_logging_obj.model_call_details["model"] = ( + f"{MCP_TOOL_NAME_PREFIX}: {standard_logging_mcp_tool_call.get('name') or ''}" + ) + litellm_logging_obj.model_call_details["custom_llm_provider"] = ( + standard_logging_mcp_tool_call.get("mcp_server_name") + ) # Try managed server tool first if name in global_mcp_server_manager.tool_name_to_mcp_server_name_mapping: @@ -218,27 +301,35 @@ if MCP_AVAILABLE: except Exception as e: return [MCPTextContent(text=f"Error: {str(e)}", type="text")] - @router.get("/", response_class=StreamingResponse) - async def handle_sse(request: Request): - verbose_logger.info("new incoming SSE connection established") - async with sse.connect_sse(request) as streams: - try: - await server.run(streams[0], streams[1], options) - except BrokenResourceError: - pass - except asyncio.CancelledError: - pass - except ValidationError: - pass - except Exception: - raise - await request.close() + async def handle_streamable_http_mcp( + scope: Scope, receive: Receive, send: Send + ) -> None: + """Handle MCP requests through StreamableHTTP.""" + try: + # Ensure session managers are initialized + if not _SESSION_MANAGERS_INITIALIZED: + await initialize_session_managers() + # Give it a moment to start up + await asyncio.sleep(0.1) - @router.post("/sse/messages") - async def handle_messages(request: Request): - verbose_logger.info("incoming SSE message received") - await sse.handle_post_message(request.scope, request.receive, request._send) - await request.close() + await session_manager.handle_request(scope, receive, send) + except Exception as e: + verbose_logger.exception(f"Error handling MCP request: {e}") + raise e + + async def handle_sse_mcp(scope: Scope, receive: Receive, send: Send) -> None: + """Handle MCP requests through SSE.""" + try: + # Ensure session managers are initialized + if not _SESSION_MANAGERS_INITIALIZED: + await initialize_session_managers() + # Give it a moment to start up + await asyncio.sleep(0.1) + + await sse_session_manager.handle_request(scope, receive, send) + except Exception as e: + verbose_logger.exception(f"Error handling MCP request: {e}") + raise e ######################################################## ############ MCP Server REST API Routes ################# @@ -315,11 +406,19 @@ if MCP_AVAILABLE: ) return await call_mcp_tool(**data) - options = InitializationOptions( - server_name="litellm-mcp-server", - server_version="0.1.0", - capabilities=server.get_capabilities( - notification_options=NotificationOptions(), - experimental_capabilities={}, - ), + app = FastAPI( + title=LITELLM_MCP_SERVER_NAME, + description=LITELLM_MCP_SERVER_DESCRIPTION, + version=LITELLM_MCP_SERVER_VERSION, + lifespan=lifespan, ) + + # Include the MCP router + app.include_router(router) + + # Mount the MCP handlers + app.mount("/", handle_streamable_http_mcp) + app.mount("/sse", handle_sse_mcp) + +else: + app = FastAPI() diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 42fc6a2422..87cfda32a3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -26,6 +26,7 @@ from typing import ( ) from litellm.constants import ( + BASE_MCP_ROUTE, DEFAULT_MAX_RECURSE_DEPTH, DEFAULT_SLACK_ALERTING_THRESHOLD, LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS, @@ -143,7 +144,7 @@ from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler -from litellm.proxy._experimental.mcp_server.server import router as mcp_router +from litellm.proxy._experimental.mcp_server.server import app as mcp_app from litellm.proxy._experimental.mcp_server.tool_registry import ( global_mcp_tool_registry, ) @@ -8271,7 +8272,6 @@ app.include_router(image_router) app.include_router(fine_tuning_router) app.include_router(credential_router) app.include_router(llm_passthrough_router) -app.include_router(mcp_router) app.include_router(mcp_management_router) app.include_router(anthropic_router) app.include_router(langfuse_router) @@ -8297,3 +8297,7 @@ app.include_router(model_management_router) app.include_router(tag_management_router) app.include_router(enterprise_router) app.include_router(ui_discovery_endpoints_router) +######################################################## +# MCP Server +######################################################## +app.mount(path=BASE_MCP_ROUTE, app=mcp_app) diff --git a/tests/local_testing/test_basic_python_version.py b/tests/local_testing/test_basic_python_version.py index c629ef3df8..7d06f9c8be 100644 --- a/tests/local_testing/test_basic_python_version.py +++ b/tests/local_testing/test_basic_python_version.py @@ -24,8 +24,8 @@ def test_using_litellm(): def test_litellm_proxy_server(): - # Install the litellm[proxy] package - subprocess.run(["pip", "install", "litellm[proxy]"]) + # Install the local litellm[proxy] package in development mode + subprocess.run(["pip", "install", "-e", ".[proxy]"]) # Import the proxy_server module try: @@ -91,11 +91,11 @@ import requests def test_litellm_proxy_server_config_no_general_settings(): - # Install the litellm[proxy] package - # Start the server + # Install the local litellm packages in development mode + server_process = None try: - subprocess.run(["pip", "install", "litellm[proxy]"]) - subprocess.run(["pip", "install", "litellm[extra_proxy]"]) + subprocess.run(["pip", "install", "-e", ".[proxy]"]) + subprocess.run(["pip", "install", "-e", ".[extra_proxy]"]) filepath = os.path.dirname(os.path.abspath(__file__)) config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml" server_process = subprocess.Popen( @@ -136,8 +136,9 @@ def test_litellm_proxy_server_config_no_general_settings(): pytest.fail("Failed to connect to the server") finally: # Shut down the server - server_process.terminate() - server_process.wait() + if server_process: + server_process.terminate() + server_process.wait() # Additional assertions can be added here assert True diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index a78805ca15..b4af1dca28 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -321,3 +321,57 @@ async def test_mcp_http_transport_tool_not_found(): ) +@pytest.mark.asyncio +async def test_streamable_http_mcp_handler_mock(): + """Test the streamable HTTP MCP handler functionality""" + + # Mock the session manager and its methods + mock_session_manager = AsyncMock() + mock_session_manager.handle_request = AsyncMock() + + # Mock scope, receive, send + mock_scope = {"type": "http", "method": "POST", "path": "/mcp"} + mock_receive = AsyncMock() + mock_send = AsyncMock() + + with patch('litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED', True), \ + patch('litellm.proxy._experimental.mcp_server.server.session_manager', mock_session_manager): + + from litellm.proxy._experimental.mcp_server.server import handle_streamable_http_mcp + + # Call the handler + await handle_streamable_http_mcp(mock_scope, mock_receive, mock_send) + + # Verify session manager handle_request was called + mock_session_manager.handle_request.assert_called_once_with( + mock_scope, mock_receive, mock_send + ) + + +@pytest.mark.asyncio +async def test_sse_mcp_handler_mock(): + """Test the SSE MCP handler functionality""" + + # Mock the SSE session manager and its methods + mock_sse_session_manager = AsyncMock() + mock_sse_session_manager.handle_request = AsyncMock() + + # Mock scope, receive, send + mock_scope = {"type": "http", "method": "GET", "path": "/mcp/sse"} + mock_receive = AsyncMock() + mock_send = AsyncMock() + + with patch('litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED', True), \ + patch('litellm.proxy._experimental.mcp_server.server.sse_session_manager', mock_sse_session_manager): + + from litellm.proxy._experimental.mcp_server.server import handle_sse_mcp + + # Call the handler + await handle_sse_mcp(mock_scope, mock_receive, mock_send) + + # Verify SSE session manager handle_request was called + mock_sse_session_manager.handle_request.assert_called_once_with( + mock_scope, mock_receive, mock_send + ) + +