diff --git a/fastmcp_slim/fastmcp/server/tasks/context.py b/fastmcp_slim/fastmcp/server/tasks/context.py index 7ade43dfc..b89572001 100644 --- a/fastmcp_slim/fastmcp/server/tasks/context.py +++ b/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: diff --git a/tests/server/tasks/test_context_background_task.py b/tests/server/tasks/test_context_background_task.py index e0b8b3dc3..45060ea7a 100644 --- a/tests/server/tasks/test_context_background_task.py +++ b/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."""