From 26cb228b5dd3b855e390553dea5f476e324e47cd Mon Sep 17 00:00:00 2001 From: Yug Date: Thu, 30 Apr 2026 22:07:28 +0530 Subject: [PATCH] resolve --- .../proxy/_experimental/mcp_server/server.py | 29 ++++++++++++------- tests/mcp_tests/test_proxy_mcp_e2e.py | 2 -- 2 files changed, 19 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index e241b0f9c05..6ec286abaf5 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -112,13 +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) - ) + _captured_session_id_container_var: contextvars.ContextVar[ + Optional[Dict[str, uuid.UUID]] + ] = contextvars.ContextVar("captured_session_id_container", default=None) class _SessionIdCapturingDict(dict): def __setitem__(self, key, value): - _captured_session_id_var.set(key) + container = _captured_session_id_container_var.get() + if container is not None: + container["session_id"] = key super().__setitem__(key, value) active_mcp_session_var: contextvars.ContextVar[Optional[_McpServerSession]] = ( @@ -2689,13 +2691,18 @@ if MCP_AVAILABLE: verbose_logger.info("Initializing SSE session...") options = server.create_initialization_options() # Capture existing session IDs to find the newly created one - _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..." - ) + # We use a mutable container because ContextVar writes are task-local + _capture_container: Dict[str, uuid.UUID] = {} + _capture_token = _captured_session_id_container_var.set( + _capture_container + ) + try: + async with sse.connect_sse(scope, receive, send) as streams: + verbose_logger.info( + "SSE connection established, running server loop..." + ) - session_id = _captured_session_id_var.get() + session_id = _capture_container.get("session_id") # ContextVars are lost when the MCP SDK spawns internal tasks # (e.g. _receive_loop), so tool handlers can't read auth_context_var reliably. @@ -2725,6 +2732,8 @@ if MCP_AVAILABLE: finally: if session_id: _session_id_auth_storage.pop(session_id, None) + finally: + _captured_session_id_container_var.reset(_capture_token) except HTTPException: raise diff --git a/tests/mcp_tests/test_proxy_mcp_e2e.py b/tests/mcp_tests/test_proxy_mcp_e2e.py index 67b73f4c951..569b6f82cda 100644 --- a/tests/mcp_tests/test_proxy_mcp_e2e.py +++ b/tests/mcp_tests/test_proxy_mcp_e2e.py @@ -1,5 +1,4 @@ import asyncio -import os import socket import subprocess import sys @@ -100,7 +99,6 @@ def math_streamable_http_server() -> str: str(port), ] - env = os.environ.copy() server_process = subprocess.Popen( cmd, cwd=str(PROJECT_ROOT),