From 50931e4b37f15f11df8bcb4cb4f270ef4382eb92 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 5 May 2026 20:32:58 +0000 Subject: [PATCH] fix(mcp): serialize concurrent requests on same stateful session Bugbot's 'Concurrent requests share context' finding: _update_auth_context mutates the single MCPAuthenticatedUser stored per session in place on every request, so two requests sharing one mcp-session-id can overwrite each other's mcp_servers / auth headers / oauth state / client_ip while in-flight callbacks are still reading the same object. Owner-binding alone narrows this to same-principal racing, but the in-place mutation race remains. Add a per-session asyncio.Lock around handle_request so concurrent same-session requests run sequentially. The lock is allocated on demand and torn down with the rest of the session state on DELETE / idle expiry. Co-authored-by: Mateo Wang --- .../proxy/_experimental/mcp_server/server.py | 86 ++++++++++++------ .../mcp_server/test_mcp_server.py | 87 +++++++++++++++++++ 2 files changed, 146 insertions(+), 27 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 5cf94aa48c7..628856e46b4 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -258,6 +258,13 @@ if MCP_AVAILABLE: # Without this, a leaked mcp-session-id could be driven (or terminated) # by any other authenticated proxy user. _stateful_session_owners: Dict[str, str] = {} + # Per-session lock that serializes ``handle_request`` for the same + # mcp-session-id. The stored ``MCPAuthenticatedUser`` is mutated in place + # by ``_update_auth_context`` each request; without this lock, two + # concurrent requests on the same session would clobber each other's + # auth headers / mcp_servers / oauth state while in-flight callbacks are + # still reading the shared object. + _stateful_session_locks: Dict[str, asyncio.Lock] = {} # Keep this alias so existing references to session_manager still work session_manager = session_manager_stateless @@ -293,6 +300,7 @@ if MCP_AVAILABLE: _stateful_session_auth_contexts.pop(session_id, None) _stateful_session_auth_context_last_seen.pop(session_id, None) _stateful_session_owners.pop(session_id, None) + _stateful_session_locks.pop(session_id, None) transport = server_instances.pop(session_id, None) if transport is not None: await transport.terminate() @@ -301,6 +309,7 @@ if MCP_AVAILABLE: if session_id not in _stateful_session_auth_contexts: _stateful_session_auth_context_last_seen.pop(session_id, None) _stateful_session_owners.pop(session_id, None) + _stateful_session_locks.pop(session_id, None) async def _cleanup_expired_stateful_session_auth_contexts() -> None: while True: @@ -3020,35 +3029,58 @@ if MCP_AVAILABLE: receive = wrapped_receive - auth_user = _set_or_update_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=_client_ip, - session_id=session_id if use_stateful else None, - ) - if use_stateful and is_initialize: - send = _wrap_send_with_stateful_session_auth_context( - send, - auth_user, - _owner_fingerprint_for(user_api_key_auth), + # Serialize requests on the same stateful session so concurrent + # callers don't clobber each other's auth context mid-flight. + session_lock: Optional[asyncio.Lock] = None + if use_stateful and session_id: + session_lock = _stateful_session_locks.setdefault( + session_id, asyncio.Lock() ) - async with _gateway_initialize_instructions_request_scope( - user_api_key_auth, - mcp_servers, - _client_ip, - ): - try: - await target_manager.handle_request(scope, receive, send) - finally: - if use_stateful and session_id and scope.get("method") == "DELETE": - _stateful_session_auth_contexts.pop(session_id, None) - _stateful_session_auth_context_last_seen.pop(session_id, None) - _stateful_session_owners.pop(session_id, None) + async def _dispatch() -> None: + auth_user = _set_or_update_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=_client_ip, + session_id=session_id if use_stateful else None, + ) + local_send = send + if use_stateful and is_initialize: + local_send = _wrap_send_with_stateful_session_auth_context( + local_send, + auth_user, + _owner_fingerprint_for(user_api_key_auth), + ) + + async with _gateway_initialize_instructions_request_scope( + user_api_key_auth, + mcp_servers, + _client_ip, + ): + try: + await target_manager.handle_request(scope, receive, local_send) + finally: + if ( + use_stateful + and session_id + and scope.get("method") == "DELETE" + ): + _stateful_session_auth_contexts.pop(session_id, None) + _stateful_session_auth_context_last_seen.pop( + session_id, None + ) + _stateful_session_owners.pop(session_id, None) + _stateful_session_locks.pop(session_id, None) + + if session_lock is not None: + async with session_lock: + await _dispatch() + else: + await _dispatch() except HTTPException: # Re-raise HTTP exceptions to preserve status codes and details raise diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index ab6dc1752b9..0ecc896fb47 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -1508,6 +1508,93 @@ async def test_stateful_mcp_session_owner_mismatch_returns_403(): mcp_server._stateful_session_owners.pop(session_id, None) +@pytest.mark.asyncio +async def test_stateful_mcp_session_serializes_concurrent_requests(): + """ + Concurrent requests on the same stateful mcp-session-id must be + serialized so they cannot observe each other's mutation of the shared + MCPAuthenticatedUser while in-flight callbacks are still running. + """ + try: + from litellm.proxy._experimental.mcp_server import server as mcp_server + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + session_manager_stateful, + ) + except ImportError: + pytest.skip("MCP server not available") + + session_id = "serialized-session-1" + owner_auth = UserAPIKeyAuth(api_key="owner-key", user_id="owner") + mcp_server._stateful_session_auth_contexts[session_id] = ( + mcp_server.MCPAuthenticatedUser(user_api_key_auth=owner_auth) + ) + mcp_server._stateful_session_owners[session_id] = mcp_server._owner_fingerprint_for( + owner_auth + ) + + inside = 0 + max_inside = 0 + gate = asyncio.Event() + + async def slow_handle(s, r, se): + nonlocal inside, max_inside + inside += 1 + max_inside = max(max_inside, inside) + await gate.wait() + inside -= 1 + + async def make_request(): + scope = { + "type": "http", + "method": "POST", + "path": "/mcp", + "headers": [(b"mcp-session-id", session_id.encode())], + } + receive = AsyncMock( + return_value={ + "type": "http.request", + "body": b'{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}', + "more_body": False, + } + ) + send = AsyncMock() + await handle_streamable_http_mcp(scope, receive, send) + + try: + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(owner_auth, None, None, None, None, None), + ), + patch( + "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", + True, + ), + patch.object( + session_manager_stateful, "handle_request", side_effect=slow_handle + ), + patch.object( + session_manager_stateful, + "_server_instances", + {session_id: MagicMock()}, + ), + ): + tasks = [asyncio.create_task(make_request()) for _ in range(3)] + await asyncio.sleep(0.05) + gate.set() + await asyncio.gather(*tasks) + finally: + mcp_server._stateful_session_auth_contexts.pop(session_id, None) + mcp_server._stateful_session_owners.pop(session_id, None) + mcp_server._stateful_session_locks.pop(session_id, None) + + assert ( + max_inside == 1 + ), "concurrent requests on same stateful session must be serialized" + + @pytest.mark.asyncio @pytest.mark.no_parallel async def test_mcp_routing_with_conflicting_alias_and_group_name():