Clean up task sessions on connection exit (#4535)

fastmcp_slim/fastmcp/server/tasks/context.py

@@ -281,6 +281,13 @@ async def restore_task_snapshot(key: str = TaskKey()) -> None:
 # process boundaries (see notifications.py and elicitation.py).
 
 _task_sessions: dict[str, weakref.ref[ServerSession]] = {}
+_TASK_SESSION_CONNECTION_REF = "_fastmcp_task_session_ref"
+_TASK_SESSION_CLEANUP_REGISTERED = "_fastmcp_task_session_cleanup_registered"
+
+
+def _remove_task_session(session_id: str, ref: weakref.ref[ServerSession]) -> None:
+    if _task_sessions.get(session_id) is ref:
+        _task_sessions.pop(session_id)
 
 
 def register_task_session(session_id: str, session: ServerSession) -> None:
@@ -291,11 +298,28 @@ def register_task_session(session_id: str, session: ServerSession) -> None:
     client disconnects.
     """
 
-    def remove_session(ref: weakref.ref[ServerSession]) -> None:
-        if _task_sessions.get(session_id) is ref:
-            _task_sessions.pop(session_id)
+    session_ref = weakref.ref(
+        session, lambda ref: _remove_task_session(session_id, ref)
+    )
+    _task_sessions[session_id] = session_ref
 
-    _task_sessions[session_id] = weakref.ref(session, remove_session)
+    connection = getattr(session, "_connection", None)
+    if connection is None:
+        return
+
+    state = connection.state
+    state[_TASK_SESSION_CONNECTION_REF] = (session_id, session_ref)
+    if state.get(_TASK_SESSION_CLEANUP_REGISTERED):
+        return
+
+    def remove_connection_session() -> None:
+        registered = state.pop(_TASK_SESSION_CONNECTION_REF, None)
+        if registered is not None:
+            registered_session_id, registered_ref = registered
+            _remove_task_session(registered_session_id, registered_ref)
+
+    connection.exit_stack.callback(remove_connection_session)
+    state[_TASK_SESSION_CLEANUP_REGISTERED] = True
 
 
 def get_task_session(session_id: str) -> ServerSession | None:

tests/server/tasks/test_context_background_task.py

@@ -8,6 +8,7 @@ no mocking of Redis, Docket, or session internals.
 import asyncio
 import gc
 import json
+from contextlib import AsyncExitStack
 from datetime import datetime, timezone
 from typing import Any, cast
 from unittest.mock import AsyncMock, patch
@@ -97,6 +98,51 @@ async def test_task_session_is_released_after_client_disconnect():
     assert _task_sessions == {}
 
 
+async def test_live_task_session_is_released_on_connection_disconnect():
+    _task_sessions.clear()
+
+    class MockConnection:
+        def __init__(self) -> None:
+            self.state: dict[str, object] = {}
+            self.exit_stack = AsyncExitStack()
+
+    class MockSession:
+        def __init__(self, connection: MockConnection) -> None:
+            self._connection = connection
+
+    connection = MockConnection()
+    session = MockSession(connection)
+    async with connection.exit_stack:
+        register_task_session("session", cast(ServerSession, session))
+        session_ref = _task_sessions["session"]
+
+    assert session_ref() is session
+    assert _task_sessions == {}
+
+
+async def test_connection_cleanup_does_not_remove_replacement_session():
+    _task_sessions.clear()
+
+    class MockConnection:
+        def __init__(self) -> None:
+            self.state: dict[str, object] = {}
+            self.exit_stack = AsyncExitStack()
+
+    class MockSession:
+        def __init__(self, connection: MockConnection | None = None) -> None:
+            self._connection = connection
+
+    connection = MockConnection()
+    old_session = MockSession(connection)
+    new_session = MockSession()
+    async with connection.exit_stack:
+        register_task_session("shared", cast(ServerSession, old_session))
+        register_task_session("shared", cast(ServerSession, new_session))
+
+    assert get_task_session("shared") is new_session
+    _task_sessions.clear()
+
+
 def test_replaced_task_session_is_not_removed_by_old_weakref():
     _task_sessions.clear()