From 9f595912dbdd8e1bcb4a19eab21438ff3e87ad3f Mon Sep 17 00:00:00 2001 From: Yug Date: Thu, 30 Apr 2026 21:54:47 +0530 Subject: [PATCH] resolve --- .../proxy/_experimental/mcp_server/server.py | 104 +++++--- .../test_configs/test_config_mcp_e2e.yaml | 2 +- tests/mcp_tests/test_mcp_session_binding.py | 235 ++++++++++++++++++ tests/mcp_tests/test_proxy_mcp_e2e.py | 6 +- 4 files changed, 314 insertions(+), 33 deletions(-) create mode 100644 tests/mcp_tests/test_mcp_session_binding.py diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 213446f0d70..0fe510891d9 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/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 diff --git a/tests/mcp_tests/test_configs/test_config_mcp_e2e.yaml b/tests/mcp_tests/test_configs/test_config_mcp_e2e.yaml index ad68a03781d..d097e6b0f4e 100644 --- a/tests/mcp_tests/test_configs/test_config_mcp_e2e.yaml +++ b/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 diff --git a/tests/mcp_tests/test_mcp_session_binding.py b/tests/mcp_tests/test_mcp_session_binding.py new file mode 100644 index 00000000000..c3d2313b046 --- /dev/null +++ b/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 diff --git a/tests/mcp_tests/test_proxy_mcp_e2e.py b/tests/mcp_tests/test_proxy_mcp_e2e.py index 2dd57e13d3b..67b73f4c951 100644 --- a/tests/mcp_tests/test_proxy_mcp_e2e.py +++ b/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={