scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

Fix active stateful MCP session cleanup

Browse source
This commit is contained in:
Cursor Agent 2026-05-06 00:35:38 +00:00
parent ea0a12f9ba
commit 6cd7bb401e
No known key found for this signature in database
2 changed files with 82 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

55
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -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

37
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
View file

@ -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():
"""