From 6cd7bb401e079246d62e1b62a7f99eedb21a037a Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 6 May 2026 00:35:38 +0000 Subject: [PATCH] Fix active stateful MCP session cleanup --- .../proxy/_experimental/mcp_server/server.py | 55 +++++++++++++++---- .../mcp_server/test_mcp_server.py | 37 +++++++++++++ 2 files changed, 82 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 5a063456dc3..57eb63752f5 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -266,6 +266,7 @@ if MCP_AVAILABLE: # auth headers / mcp_servers / oauth state while in-flight callbacks are # still reading the shared object. _stateful_session_locks: Dict[str, asyncio.Lock] = {} + _stateful_session_active_request_counts: Dict[str, int] = {} # Keep this alias so existing references to session_manager still work session_manager = session_manager_stateless @@ -290,18 +291,22 @@ if MCP_AVAILABLE: """Terminate expired stateful sessions and drop their auth contexts.""" now = now or time.monotonic() server_instances = getattr(session_manager_stateful, "_server_instances", {}) - expired_session_ids = [ - session_id - for session_id, last_seen in _stateful_session_auth_context_last_seen.items() - if now - last_seen >= _STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS - or session_id not in server_instances - ] + expired_session_ids = [] + for session_id, last_seen in _stateful_session_auth_context_last_seen.items(): + if _stateful_session_active_request_counts.get(session_id, 0) > 0: + continue + if ( + now - last_seen >= _STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS + or session_id not in server_instances + ): + expired_session_ids.append(session_id) for session_id in expired_session_ids: _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) + _stateful_session_active_request_counts.pop(session_id, None) transport = server_instances.pop(session_id, None) if transport is not None: await transport.terminate() @@ -311,6 +316,7 @@ if MCP_AVAILABLE: _stateful_session_auth_context_last_seen.pop(session_id, None) _stateful_session_owners.pop(session_id, None) _stateful_session_locks.pop(session_id, None) + _stateful_session_active_request_counts.pop(session_id, None) async def _cleanup_expired_stateful_session_auth_contexts() -> None: while True: @@ -3068,6 +3074,12 @@ if MCP_AVAILABLE: session_id, asyncio.Lock() ) + track_active_stateful_request = bool(use_stateful and session_id) + if track_active_stateful_request and session_id: + _stateful_session_active_request_counts[session_id] = ( + _stateful_session_active_request_counts.get(session_id, 0) + 1 + ) + async def _dispatch() -> None: auth_user = _set_or_update_auth_context( user_api_key_auth=user_api_key_auth, @@ -3108,12 +3120,35 @@ if MCP_AVAILABLE: ) _stateful_session_owners.pop(session_id, None) _stateful_session_locks.pop(session_id, None) + _stateful_session_active_request_counts.pop( + session_id, None + ) - if session_lock is not None: - async with session_lock: + try: + if session_lock is not None: + async with session_lock: + await _dispatch() + else: await _dispatch() - else: - await _dispatch() + finally: + if track_active_stateful_request and session_id: + active_request_count = ( + _stateful_session_active_request_counts.get(session_id, 0) - 1 + ) + if active_request_count > 0: + _stateful_session_active_request_counts[session_id] = ( + active_request_count + ) + else: + _stateful_session_active_request_counts.pop(session_id, None) + + if ( + scope.get("method") != "DELETE" + and session_id in _stateful_session_auth_contexts + ): + _stateful_session_auth_context_last_seen[session_id] = ( + time.monotonic() + ) 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 10ea6559af5..26ace369905 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 @@ -1426,6 +1426,43 @@ async def test_stateful_mcp_auth_contexts_expire_with_idle_sessions(): transport.terminate.assert_awaited_once() +@pytest.mark.asyncio +async def test_stateful_mcp_auth_contexts_do_not_expire_active_sessions(): + """Active stateful sessions should not be terminated by idle cleanup.""" + try: + from litellm.proxy._experimental.mcp_server import server as mcp_server + except ImportError: + pytest.skip("MCP server not available") + + session_id = "active-stateful-session" + auth_user = UserAPIKeyAuth(api_key="active-key", user_id="active-user") + transport = MagicMock() + transport.terminate = AsyncMock() + now = 1000.0 + + mcp_server._stateful_session_auth_contexts[session_id] = auth_user + mcp_server._stateful_session_auth_context_last_seen[session_id] = ( + now - mcp_server._STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS + ) + mcp_server._stateful_session_active_request_counts[session_id] = 1 + + try: + with patch.object( + mcp_server.session_manager_stateful, + "_server_instances", + {session_id: transport}, + ): + await mcp_server._purge_expired_stateful_session_auth_contexts(now=now) + + assert session_id in mcp_server._stateful_session_auth_contexts + assert session_id in mcp_server._stateful_session_auth_context_last_seen + transport.terminate.assert_not_awaited() + finally: + mcp_server._stateful_session_auth_contexts.pop(session_id, None) + mcp_server._stateful_session_auth_context_last_seen.pop(session_id, None) + mcp_server._stateful_session_active_request_counts.pop(session_id, None) + + @pytest.mark.asyncio async def test_owner_fingerprint_distinguishes_oauth_callers(): """