resolve

litellm/proxy/_experimental/mcp_server/server.py

@@ -112,6 +112,15 @@ try:
     )
     _session_id_auth_storage: Dict[uuid.UUID, "MCPAuthenticatedUser"] = {}
 
+    _captured_session_id_var: contextvars.ContextVar[Optional[uuid.UUID]] = (
+        contextvars.ContextVar("captured_session_id", default=None)
+    )
+
+    class _SessionIdCapturingDict(dict):
+        def __setitem__(self, key, value):
+            _captured_session_id_var.set(key)
+            super().__setitem__(key, value)
+
     active_mcp_session_var: contextvars.ContextVar[Optional[_McpServerSession]] = (
         contextvars.ContextVar("active_mcp_session", default=None)
     )
@@ -245,6 +254,8 @@ if MCP_AVAILABLE:
     from mcp.server.sse import SseServerTransport as _McpSseServerTransport
 
     sse = _McpSseServerTransport("/mcp/messages")
+    # Wrap the dict to capture session IDs race-free during connect_sse
+    sse._read_stream_writers = _SessionIdCapturingDict(sse._read_stream_writers)  # type: ignore
     # Create session managers (StreamableHTTP — stateless by default)
     session_manager = StreamableHTTPSessionManager(
         app=server,
@@ -2678,17 +2689,13 @@ if MCP_AVAILABLE:
                 verbose_logger.info("Initializing SSE session...")
                 options = server.create_initialization_options()
                 # Capture existing session IDs to find the newly created one
-                old_session_ids = set(getattr(sse, "_read_stream_writers", {}).keys())
+                _captured_session_id_var.set(None)
                 async with sse.connect_sse(scope, receive, send) as streams:
                     verbose_logger.info(
                         "SSE connection established, running server loop..."
                     )
 
-                    new_session_ids = set(
-                        getattr(sse, "_read_stream_writers", {}).keys()
-                    )
-                    diff = new_session_ids - old_session_ids
-                    session_id = diff.pop() if diff else None
+                    session_id = _captured_session_id_var.get()
 
                     # ContextVars are lost when the MCP SDK spawns internal tasks
                     # (e.g. _receive_loop), so tool handlers can't read auth_context_var reliably.
@@ -2759,10 +2766,6 @@ if MCP_AVAILABLE:
             ) = await extract_mcp_auth_context(scope, path)
             _sse_client_ip = IPAddressUtils.get_mcp_client_ip(request)
             # set_auth_context here is a no-op for actual tool execution since the SDK
-            # processes messages in background tasks that don't inherit this ContextVar.
-            # Authentication must be recovered from the session-auth-storage during execution.
-
-            # P1 Security: Bind POST auth to the existing SSE session
             session_id_param = request.query_params.get(
                 "sessionId"
             ) or request.query_params.get("session_id")
@@ -2771,31 +2774,64 @@ if MCP_AVAILABLE:
 
                 try:
                     session_id = uuid.UUID(hex=session_id_param)
-                    writer = getattr(sse, "_read_stream_writers", {}).get(session_id)
-                    if writer:
-                        session_auth = _session_id_auth_storage.get(session_id)
+                    # Look up auth associated with this session ID
+                    session_auth = _session_id_auth_storage.get(session_id)
 
-                        if (
-                            session_auth
-                            and session_auth.user_api_key_auth
-                            and user_api_key_auth
-                        ):
-                            session_api_key = getattr(
-                                session_auth.user_api_key_auth, "api_key", None
+                    if session_auth is not None:
+                        # P1 Security: Bind POST auth to the existing SSE session
+                        session_api_key = getattr(
+                            session_auth.user_api_key_auth, "api_key", None
+                        )
+                        session_has_auth = (
+                            session_api_key is not None or session_auth.oauth2_headers
+                        )
+
+                        if session_has_auth:
+                            post_api_key = (
+                                getattr(user_api_key_auth, "api_key", None)
+                                if user_api_key_auth
+                                else None
                             )
-                            post_api_key = getattr(user_api_key_auth, "api_key", None)
-                            if (
-                                session_api_key
-                                and post_api_key
-                                and session_api_key != post_api_key
-                            ):
+
+                            # Match by API Key
+                            if session_api_key is not None:
+                                if session_api_key != post_api_key:
+                                    raise HTTPException(
+                                        status_code=403,
+                                        detail="Authentication mismatch: POST auth does not match the session owner.",
+                                    )
+                            # Match by OAuth2 Token if no API key
+                            elif session_auth.oauth2_headers:
+                                if (
+                                    not oauth2_headers
+                                    or session_auth.oauth2_headers.get("Authorization")
+                                    != oauth2_headers.get("Authorization")
+                                ):
+                                    raise HTTPException(
+                                        status_code=403,
+                                        detail="Authentication mismatch: POST OAuth2 token does not match the session owner.",
+                                    )
+                            # If session was authenticated but POST has no auth, deny.
+                            elif user_api_key_auth is None and not oauth2_headers:
                                 raise HTTPException(
                                     status_code=403,
-                                    detail="Authentication mismatch: POST auth does not match the session owner.",
+                                    detail="Authentication required: Target session is authenticated but POST lacks credentials.",
                                 )
                 except ValueError:
+                    # Invalid UUID format in session_id_param
                     pass
 
+            # Set auth context for this POST request (important for SDK background tasks)
+            set_auth_context(
+                user_api_key_auth=user_api_key_auth,
+                mcp_auth_header=mcp_auth_header,
+                mcp_servers=mcp_servers,
+                mcp_server_auth_headers=mcp_server_auth_headers,
+                oauth2_headers=oauth2_headers,
+                raw_headers=raw_headers,
+                client_ip=_sse_client_ip,
+            )
+            await sse.handle_post_message(scope, receive, send)
         except HTTPException:
             raise
         except Exception as e:
@@ -2822,7 +2858,17 @@ if MCP_AVAILABLE:
 
     def get_active_mcp_session() -> Optional[_McpServerSession]:
         """Get the active downstream MCP session from the current context."""
-        return active_mcp_session_var.get()
+        # Try context var first
+        session = active_mcp_session_var.get()
+        if session:
+            return session
+        # Fallback to request_ctx (available in SDK request handlers)
+        try:
+            from mcp.server.lowlevel.server import request_ctx
+
+            return request_ctx.get().session
+        except Exception:
+            return None
 
     def get_active_auth_context() -> Optional[MCPAuthenticatedUser]:
         """Get the active auth context from the server object or context var."""
@@ -2832,7 +2878,7 @@ if MCP_AVAILABLE:
             return auth
         elif auth:
             return cast(MCPAuthenticatedUser, auth)
-        # Fallback to session read_stream
+        # Fallback to session read_stream (available in SDK request handlers)
         try:
             from mcp.server.lowlevel.server import request_ctx
 

tests/mcp_tests/test_configs/test_config_mcp_e2e.yaml

@@ -15,7 +15,7 @@ model_list:
 mcp_servers:
   math_stdio:
     transport: stdio
-    command: python3
+    command: python
     args:
       - tests/mcp_tests/mcp_server.py
     allow_all_keys: true  # Needed for master_key auth when DB/object_permission not used

tests/mcp_tests/test_mcp_session_binding.py

@@ -0,0 +1,235 @@
+import pytest
+import uuid
+import asyncio
+from unittest.mock import AsyncMock, MagicMock, patch
+from fastapi import HTTPException
+from contextlib import asynccontextmanager
+from litellm.proxy._experimental.mcp_server.server import (
+    handle_sse_mcp_endpoint,
+    handle_sse_post_messages,
+    _captured_session_id_var,
+    _session_id_auth_storage,
+)
+from litellm.proxy._types import UserAPIKeyAuth
+
+
+@pytest.mark.asyncio
+async def test_session_id_capture_and_binding():
+    """
+    Test that session_id is correctly captured during SSE connection
+    and that subsequent POST messages are correctly bound to that session.
+    """
+    session_id = uuid.uuid4()
+
+    # Mock SSE transport
+    mock_sse = MagicMock()
+    mock_sse.handle_post_message = AsyncMock()
+
+    # Mock connect_sse to simulate the SDK behavior
+    @asynccontextmanager
+    async def mock_connect_sse(scope, receive, send):
+        # Directly set the ContextVar to simulate our capturing dict's behavior
+        _captured_session_id_var.set(session_id)
+        yield (AsyncMock(), AsyncMock())
+
+    mock_sse.connect_sse = MagicMock(side_effect=mock_connect_sse)
+    mock_sse._read_stream_writers = {session_id: MagicMock()}
+
+    mock_auth_context = (
+        UserAPIKeyAuth(api_key="session-key", user_id="user1"),
+        "auth-header",
+        ["server1"],
+        {},
+        {},
+        {},
+    )
+
+    mock_scope = {"type": "http", "method": "GET", "path": "/mcp/sse", "headers": []}
+    mock_receive = AsyncMock()
+    mock_send = AsyncMock()
+
+    # We'll use an event to coordinate with the mocked server.run
+    run_started = asyncio.Event()
+    finish_run = asyncio.Event()
+
+    async def mock_run(*args, **kwargs):
+        run_started.set()
+        await finish_run.wait()
+
+    # 1. Test Session ID Capture and POST binding during active session
+    with (
+        patch("litellm.proxy._experimental.mcp_server.server.sse", mock_sse),
+        patch(
+            "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
+            AsyncMock(return_value=mock_auth_context),
+        ),
+        patch(
+            "litellm.proxy._experimental.mcp_server.server._gateway_initialize_instructions_request_scope",
+            MagicMock(
+                return_value=MagicMock(__aenter__=AsyncMock(), __aexit__=AsyncMock())
+            ),
+        ),
+        patch(
+            "litellm.proxy._experimental.mcp_server.server.server",
+            MagicMock(run=mock_run, create_initialization_options=MagicMock()),
+        ),
+        patch(
+            "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED",
+            True,
+        ),
+    ):
+        # Start the endpoint in the background
+        task = asyncio.create_task(
+            handle_sse_mcp_endpoint(mock_scope, mock_receive, mock_send)
+        )
+
+        # Wait for it to start and set the auth
+        await asyncio.wait_for(run_started.wait(), timeout=2.0)
+
+        # Verify it was stored in the global storage
+        assert session_id in _session_id_auth_storage
+        stored_auth = _session_id_auth_storage[session_id]
+        assert stored_auth.user_api_key_auth.api_key == "session-key"
+
+        # 2. Test POST Message Binding (Success)
+        post_auth_context = (
+            UserAPIKeyAuth(api_key="session-key", user_id="user1"),
+            "auth-header",
+            ["server1"],
+            {},
+            {},
+            {},
+        )
+
+        post_scope = {
+            "type": "http",
+            "method": "POST",
+            "path": "/mcp/messages",
+            "query_string": f"session_id={session_id.hex}".encode(),
+            "headers": [],
+        }
+
+        with patch(
+            "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
+            AsyncMock(return_value=post_auth_context),
+        ):
+            # Should NOT raise HTTPException
+            await handle_sse_post_messages(post_scope, mock_receive, mock_send)
+
+        # 3. Test POST Message Binding (Mismatch - Security Fix)
+        wrong_post_auth_context = (
+            UserAPIKeyAuth(api_key="wrong-key", user_id="user2"),
+            "auth-header",
+            ["server1"],
+            {},
+            {},
+            {},
+        )
+
+        with patch(
+            "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
+            AsyncMock(return_value=wrong_post_auth_context),
+        ):
+            with pytest.raises(HTTPException) as exc:
+                await handle_sse_post_messages(post_scope, mock_receive, mock_send)
+            assert exc.value.status_code == 403
+            assert "Authentication mismatch" in exc.value.detail
+
+        # 3b. Test POST Message Binding (No Auth - Security Fix)
+        no_post_auth_context = (
+            None,  # user_api_key_auth is None
+            None,
+            ["server1"],
+            {},
+            {},
+            {},
+        )
+
+        with patch(
+            "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
+            AsyncMock(return_value=no_post_auth_context),
+        ):
+            with pytest.raises(HTTPException) as exc:
+                await handle_sse_post_messages(post_scope, mock_receive, mock_send)
+            assert exc.value.status_code == 403
+            assert "Authentication mismatch" in exc.value.detail
+
+        # Finish the run
+        finish_run.set()
+        await task
+
+    # 4. Verify cleanup after session ends
+    assert session_id not in _session_id_auth_storage
+
+
+@pytest.mark.asyncio
+async def test_anonymous_session_still_works():
+    """
+    Test that if a session was started without auth (anonymous),
+    subsequent POST messages without auth are still allowed.
+    """
+    session_id = uuid.uuid4()
+    mock_sse = MagicMock()
+    mock_sse.handle_post_message = AsyncMock()
+
+    @asynccontextmanager
+    async def mock_connect_sse(scope, receive, send):
+        _captured_session_id_var.set(session_id)
+        yield (AsyncMock(), AsyncMock())
+
+    mock_sse.connect_sse = MagicMock(side_effect=mock_connect_sse)
+    mock_sse._read_stream_writers = {session_id: MagicMock()}
+
+    # SSE session started WITHOUT auth
+    anon_auth_context = (None, None, ["server1"], {}, {}, {})
+    mock_scope = {"type": "http", "method": "GET", "path": "/mcp/sse", "headers": []}
+
+    finish_run = asyncio.Event()
+
+    async def mock_run(*args, **kwargs):
+        await finish_run.wait()
+
+    with (
+        patch("litellm.proxy._experimental.mcp_server.server.sse", mock_sse),
+        patch(
+            "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
+            AsyncMock(return_value=anon_auth_context),
+        ),
+        patch(
+            "litellm.proxy._experimental.mcp_server.server.server",
+            MagicMock(run=mock_run, create_initialization_options=MagicMock()),
+        ),
+        patch(
+            "litellm.proxy._experimental.mcp_server.server._gateway_initialize_instructions_request_scope",
+            MagicMock(
+                return_value=MagicMock(__aenter__=AsyncMock(), __aexit__=AsyncMock())
+            ),
+        ),
+        patch(
+            "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED",
+            True,
+        ),
+    ):
+        task = asyncio.create_task(
+            handle_sse_mcp_endpoint(mock_scope, AsyncMock(), AsyncMock())
+        )
+        await asyncio.sleep(0.1)
+
+        # POST without auth
+        post_scope = {
+            "type": "http",
+            "method": "POST",
+            "path": "/mcp/messages",
+            "query_string": f"session_id={session_id.hex}".encode(),
+            "headers": [],
+        }
+
+        with patch(
+            "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
+            AsyncMock(return_value=anon_auth_context),
+        ):
+            # Should NOT raise HTTPException
+            await handle_sse_post_messages(post_scope, AsyncMock(), AsyncMock())
+
+        finish_run.set()
+        await task

tests/mcp_tests/test_proxy_mcp_e2e.py

@@ -24,7 +24,7 @@ from litellm.proxy.proxy_server import (
 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_START_TIMEOUT = 60
 
 
 PROXY_AUTHORIZATION_HEADER = "Bearer sk-1234"
@@ -157,7 +157,7 @@ def proxy_server_url(
 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 asyncio.timeout(120):
             async with streamablehttp_client(
                 url=f"{proxy_server_url}/mcp",
                 headers={
@@ -180,7 +180,7 @@ class TestProxyMcpSimpleConnections:
     async def test_proxy_mcp_streamable_http_roundtrip(
         self, proxy_server_url: str
     ) -> None:
-        async with asyncio.timeout(20):
+        async with asyncio.timeout(120):
             async with streamablehttp_client(
                 url=f"{proxy_server_url}/mcp",
                 headers={