Repository: BerriAI/litellm
Pull request: [Feat] MCP expose streamable https endpoint for LiteLLM Proxy
## [Feat] MCP expose streamable https endpoint for LiteLLM Proxy

Closes https://github.com/BerriAI/litellm/issues/11603

This PR exposes a streamable HTTPS endpoint for the LiteLLM Proxy MCP serve

h/t @wagnerjt for the initial scoping work and help on this

<!-- e.g. "Implement user authentication feature" -->

## Relevant issues

<!-- e.g. "Fixes #000" -->

## Pre-Submission checklist

**Please complete all items before asking a LiteLLM maintainer to review your PR**

- [x] I have Added testing in the [`tests/litellm/`](https://github.com/BerriAI/litellm/tree/main/tests/litellm) directory, **Adding at least 1 test is a hard requirement** - [see details](https://docs.litellm.ai/docs/extras/contributing_code)
- [x] I have added a screenshot of my new test passing locally 
- [x] My PR passes all unit tests on [`make test-unit`](https://docs.litellm.ai/docs/extras/contributing_code)
- [x] My PR's scope is as isolated as possible, it only solves 1 specific problem


## Type

<!-- Select the type of Pull Request -->
<!-- Keep only the necessary ones -->

🆕 New Feature
🚄 Infrastructure
✅ Test

## Changes




Problem statement:
[Feature]: Upgrade to the MCP streamable http transport
### The Feature

Upgrade allowing SSE and the newer streamable http support

### Motivation, pitch

No lag behind

### LiteLLM is hiring a founding backend engineer, are you interested in joining us and shipping to all our users?

No

### Twitter / LinkedIn details

_No response_

Hints from issue discussion:
What is the ETA on this? @wagnerjt @ishaan-jaff

Commit message: Update litellm/proxy/_experimental/mcp_server/server.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Patch to review:
diff --git a/litellm/constants.py b/litellm/constants.py
index c4adb09dc6ab..c96e1dfaba82 100644
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -714,6 +714,7 @@
     "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 11a52c6bda8f..fbc780dadd49 100644
--- a/litellm/proxy/_experimental/mcp_server/server.py
+++ b/litellm/proxy/_experimental/mcp_server/server.py
@@ -3,12 +3,13 @@
 """
 
 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 fastapi.responses import JSONResponse
+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 +25,10 @@
     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 +53,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 +89,80 @@ class ListMCPToolsRestAPIResponseObject(MCPTool):
     ########################################################
     ############ 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 +241,15 @@ async def call_mcp_tool(
             "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 +302,40 @@ async def _handle_local_mcp_tool(
         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()
+    @router.get("/health")
+    async def health_check(request: Request) -> JSONResponse:
+        """Health check endpoint."""
+        return JSONResponse({"status": "healthy", "message": "MCP Server is running"})
+
+    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)
+
+            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)
 
-    @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 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 +412,19 @@ async def call_tool_rest_api(
         )
         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 42fc6a242284..87cfda32a3ec 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -26,6 +26,7 @@
 )
 
 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 @@ def generate_feedback_box():
 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 @@ async def get_routes():
 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 @@ async def get_routes():
 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)


Code review comments:
diff_hunk: @@ -76,9 +89,80 @@ class ListMCPToolsRestAPIResponseObject(MCPTool):
     ########################################################
     ############ 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
path: litellm/proxy/_experimental/mcp_server/server.py
text: @user1:
Global variable naming inconsistency: the shutdown function uses '_session_managers_initialized' and '_session_manager_task', but the corresponding globals are declared as '_SESSION_MANAGERS_INITIALIZED' and '_SESSION_MANAGER_TASK'. Use consistent naming to ensure proper state management.
```suggestion
        _SESSION_MANAGERS_INITIALIZED = False
        _SESSION_MANAGER_TASK = None
```