fix(mcp): drop orphan per-session lock when auth context absent

Defensive cleanup for _stateful_session_locks entries created on
sessions that never enter _stateful_session_auth_contexts. The
periodic cleanup loop only iterates auth_context_last_seen, so such
locks would otherwise live forever. Add a test that reproduces the
leak and verifies the request finalizer pops the lock.

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>
This commit is contained in:
Cursor Agent 2026-05-11 22:58:19 +00:00
parent ac6986bd94
commit 8125c5066b
No known key found for this signature in database
2 changed files with 84 additions and 0 deletions

View file

@ -3193,6 +3193,19 @@ if MCP_AVAILABLE:
_stateful_session_auth_context_last_seen[session_id] = (
time.monotonic()
)
# Defensive cleanup: if the session has no auth-context
# tracking and no in-flight requests, the per-session lock
# would otherwise be orphaned (the periodic cleanup loop
# iterates ``_stateful_session_auth_context_last_seen`` and
# would never see it). This guards against narrow paths
# where a lock is created for a session_id that never
# entered ``_stateful_session_auth_contexts``.
if (
active_request_count <= 0
and session_id not in _stateful_session_auth_contexts
):
_stateful_session_locks.pop(session_id, None)
except HTTPException:
# Re-raise HTTP exceptions to preserve status codes and details
raise

View file

@ -1817,6 +1817,77 @@ async def test_stateful_mcp_session_serializes_concurrent_requests():
), "concurrent requests on same stateful session must be serialized"
@pytest.mark.asyncio
async def test_stateful_mcp_lock_does_not_leak_when_auth_context_missing():
"""
If a per-session lock is created for a session_id that is not tracked in
``_stateful_session_auth_contexts`` (e.g., a defensive path), the request
finalizer must drop the lock so it isn't orphaned. The periodic cleanup
loop only iterates ``_stateful_session_auth_context_last_seen``, so a
leaked lock would otherwise live forever.
"""
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 = "untracked-session-1"
owner_auth = UserAPIKeyAuth(api_key="owner-key", user_id="owner")
async def handle(s, r, se):
return None
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,
}
)
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=handle
),
patch.object(
session_manager_stateful,
"_server_instances",
{session_id: MagicMock()},
),
):
assert session_id not in mcp_server._stateful_session_auth_contexts
await handle_streamable_http_mcp(scope, receive, AsyncMock())
assert (
session_id not in mcp_server._stateful_session_locks
), "lock entry must be cleaned up for untracked stateful session"
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)
mcp_server._stateful_session_active_request_counts.pop(session_id, None)
@pytest.mark.asyncio
async def test_stateful_mcp_get_stream_does_not_block_post():
"""