Clean up disconnected task sessions (#4519)

fastmcp_slim/fastmcp/server/tasks/context.py

@@ -290,7 +290,12 @@ def register_task_session(session_id: str, session: ServerSession) -> None:
     stored as a weakref so it doesn't prevent garbage collection when the
     client disconnects.
     """
-    _task_sessions[session_id] = weakref.ref(session)
+
+    def remove_session(ref: weakref.ref[ServerSession]) -> None:
+        if _task_sessions.get(session_id) is ref:
+            _task_sessions.pop(session_id)
+
+    _task_sessions[session_id] = weakref.ref(session, remove_session)
 
 
 def get_task_session(session_id: str) -> ServerSession | None:

tests/server/tasks/test_context_background_task.py

@@ -6,6 +6,7 @@ no mocking of Redis, Docket, or session internals.
 """
 
 import asyncio
+import gc
 import json
 from datetime import datetime, timezone
 from typing import Any, cast
@@ -40,7 +41,10 @@ from fastmcp.server.tasks.context import (
     TaskContextInfo,
     TaskContextSnapshot,
     _remember_snapshot,
+    _task_sessions,
     get_task_scope,
+    get_task_session,
+    register_task_session,
 )
 from fastmcp.server.tasks.elicitation import handle_task_input
 from fastmcp.server.tasks.keys import (
@@ -77,6 +81,41 @@ class TestContextBackgroundTaskSupport:
             setattr(ctx, "task_id", "new-id")
 
 
+async def test_task_session_is_released_after_client_disconnect():
+    _task_sessions.clear()
+    mcp = FastMCP("test")
+
+    @mcp.tool(task=True)
+    async def work() -> str:
+        return "done"
+
+    async with Client(mcp) as client:
+        task = await client.call_tool("work", task=True)
+        await task.result()
+        assert len(_task_sessions) == 1
+
+    assert _task_sessions == {}
+
+
+def test_replaced_task_session_is_not_removed_by_old_weakref():
+    _task_sessions.clear()
+
+    class MockSession:
+        pass
+
+    old_session = MockSession()
+    new_session = MockSession()
+    register_task_session("shared", cast(ServerSession, old_session))
+    old_ref = _task_sessions["shared"]
+    register_task_session("shared", cast(ServerSession, new_session))
+
+    del old_session
+    gc.collect()
+
+    assert old_ref() is None
+    assert get_task_session("shared") is new_session
+
+
 class TestContextSessionProperty:
     """Tests for Context.session property in different modes."""