ceos/fastmcp
1
0
Fork 0
mirror of https://github.com/PrefectHQ/fastmcp.git synced 2026-08-19 20:14:17 +02:00
Code Issues Projects Releases Packages Wiki Activity

Address review: set status key on cancel fallback, prevent relay task GC

Browse source 
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Chris Guidry 2026-02-10 15:35:50 -05:00
commit aa4db3d00e
1 changed files with 16 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

18
src/fastmcp/server/tasks/notifications.py
View file

@ -199,11 +199,18 @@ async def _send_mcp_notification(
related_task = meta.get("modelcontextprotocol.io/related-task", {})
elicitation = related_task.get("elicitation")
if elicitation:
task_id = params["taskId"]
asyncio.create_task( # noqa: RUF006
task_id = params.get("taskId")
if not task_id:
logger.warning(
"input_required notification missing taskId, skipping relay"
)
return
task = asyncio.create_task(
_relay_elicitation(session, session_id, task_id, elicitation, docket),
name=f"elicitation-relay-{task_id[:8]}",
)
_background_tasks.add(task)
task.add_done_callback(_background_tasks.discard)
async def _relay_elicitation(
@ -268,12 +275,16 @@ async def _relay_elicitation(
response_key = docket.key(
ELICIT_RESPONSE_KEY.format(session_id=session_id, task_id=task_id)
)
status_key = docket.key(
ELICIT_STATUS_KEY.format(session_id=session_id, task_id=task_id)
)
cancel = {"action": "cancel", "content": None}
async with docket.redis() as redis:
await redis.lpush( # type: ignore[invalid-await]
response_key, json.dumps(cancel)
)
await redis.expire(response_key, ELICIT_TTL_SECONDS)
await redis.set(status_key, "responded", ex=ELICIT_TTL_SECONDS)
except Exception as cancel_error:
logger.warning(
"Failed to push cancel response for task %s "
@ -287,6 +298,9 @@ async def _relay_elicitation(
# Subscriber Management
# =============================================================================
# Strong references to fire-and-forget relay tasks (prevent GC mid-flight)
_background_tasks: set[asyncio.Task[None]] = set()
# Registry of active subscribers per session (prevents duplicates)
# Uses weakref to session to detect disconnects
_active_subscribers: dict[