From 2d3ad9ca0de9b26b6f55dff1c120c5b1aa845af3 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Sun, 26 Jul 2026 16:40:33 -0400 Subject: [PATCH] Make a contended tasks/update wait instead of dropping its answer Partial fulfillment means two in-flight updates can carry different answers, so acknowledging the one that loses the update lock stranded the task on a key the client had already sent. --- fastmcp_tasks/fastmcp_tasks/handlers.py | 22 +++++-- tests/server/middleware/test_caching.py | 53 ++++++---------- tests/tasks/server/test_guard_reentrant.py | 74 ++++++++++++++++++++++ 3 files changed, 111 insertions(+), 38 deletions(-) diff --git a/fastmcp_tasks/fastmcp_tasks/handlers.py b/fastmcp_tasks/fastmcp_tasks/handlers.py index 836441c98..5193d5b84 100644 --- a/fastmcp_tasks/fastmcp_tasks/handlers.py +++ b/fastmcp_tasks/fastmcp_tasks/handlers.py @@ -39,7 +39,6 @@ from fastmcp_tasks.creation import ( registered_component_for_key, ) from fastmcp_tasks.input_store import ( - acquire_update_lock, acquire_update_lock_blocking, clear_outstanding, discard_outstanding, @@ -351,9 +350,24 @@ async def tasks_update( ) # Serialize concurrent updates for this task so two racing answers cannot - # each enqueue a next leg (double execution). A loser is an idempotent no-op. - if not await acquire_update_lock(docket, task_scope, task_id): - return UpdateTaskResult() + # each enqueue a next leg (double execution). Waiting rather than dropping + # the loser matters for partial fulfillment: SEP-2663 invites a client to + # answer a multi-request ask one key at a time, so two in-flight updates may + # carry *different* answers. Acknowledging the loser without storing its + # answer would strand the task waiting on a key the client believes it has + # already sent. Once the winner finishes, the loser re-reads the leg's + # outstanding state: a genuinely duplicate answer finds nothing left to + # match and is the idempotent no-op SEP-2663 asks for. + if not await acquire_update_lock_blocking(docket, task_scope, task_id): + # The holder is wedged. Report a retryable failure rather than a false + # acknowledgement, which would silently lose this answer. + raise MCPError( + code=mcp_types.INTERNAL_ERROR, + message=( + f"Task {task_id} has an update in progress that did not complete " + "in time; retry this update." + ), + ) try: # A cancelled task never re-enters: clearing outstanding on cancel makes # translate return None already, but check explicitly so a cancel that diff --git a/tests/server/middleware/test_caching.py b/tests/server/middleware/test_caching.py index dd1782a7e..a88efbd3f 100644 --- a/tests/server/middleware/test_caching.py +++ b/tests/server/middleware/test_caching.py @@ -26,6 +26,7 @@ from pydantic import AnyUrl, BaseModel from fastmcp import Context, FastMCP from fastmcp.client.client import CallToolResult, Client +from fastmcp.client.elicitation import ElicitResult from fastmcp.client.transports import FastMCPTransport from fastmcp.prompts.base import Message, Prompt from fastmcp.prompts.function_prompt import FunctionPrompt @@ -922,22 +923,18 @@ class TestCachingWithInputRequiredResults: @staticmethod def _ask() -> mcp_types.InputRequiredResult: - return mcp_types.InputRequiredResult( - result_type="input_required", - input_requests={ - "q": mcp_types.ElicitRequest( - method="elicitation/create", - params=mcp_types.ElicitRequestFormParams( - message="Which quarter?", - requested_schema={ - "type": "object", - "properties": {"q": {"type": "string"}}, - "required": ["q"], - }, - ), - ) + params = mcp_types.ElicitRequestFormParams( + message="Which quarter?", + requested_schema={ + "type": "object", + "properties": {"q": {"type": "string"}}, + "required": ["q"], }, ) + request = mcp_types.ElicitRequest(method="elicitation/create", params=params) + return mcp_types.InputRequiredResult( + result_type="input_required", input_requests={"q": request} + ) @classmethod def _server(cls) -> FastMCP: @@ -966,32 +963,24 @@ class TestCachingWithInputRequiredResults: @staticmethod async def _handler(message, response_type, params, ctx): - from fastmcp.client.elicitation import ElicitResult - return ElicitResult(action="accept", content=response_type(q="Q3")) - async def test_tool_guard_completes_under_caching(self): - async with Client( - self._server(), mode="auto", elicitation_handler=self._handler - ) as client: - result = await client.call_tool("summarize_tool", {}) + def _client(self) -> Client: + return Client(self._server(), mode="auto", elicitation_handler=self._handler) + async def test_tool_guard_completes_under_caching(self): + async with self._client() as client: + result = await client.call_tool("summarize_tool", {}) assert result.data == "Summary for Q3" async def test_prompt_guard_completes_under_caching(self): - async with Client( - self._server(), mode="auto", elicitation_handler=self._handler - ) as client: + async with self._client() as client: result = await client.get_prompt("summarize") - assert result.messages[0].content.text == "Summary for Q3" async def test_resource_guard_completes_under_caching(self): - async with Client( - self._server(), mode="auto", elicitation_handler=self._handler - ) as client: + async with self._client() as client: result = await client.read_resource("report://x") - assert result[0].text == "Report for Q3" async def test_guard_still_asks_on_a_second_fresh_call(self): @@ -1000,12 +989,8 @@ class TestCachingWithInputRequiredResults: A second fresh flow has to be asked the same question; serving it a cached final answer would skip the component's own per-round logic. """ - mcp = self._server() - async with Client( - mcp, mode="auto", elicitation_handler=self._handler - ) as client: + async with self._client() as client: first = await client.get_prompt("summarize") second = await client.get_prompt("summarize") - assert first.messages[0].content.text == "Summary for Q3" assert second.messages[0].content.text == "Summary for Q3" diff --git a/tests/tasks/server/test_guard_reentrant.py b/tests/tasks/server/test_guard_reentrant.py index 93123734e..fb4cc52a7 100644 --- a/tests/tasks/server/test_guard_reentrant.py +++ b/tests/tasks/server/test_guard_reentrant.py @@ -11,9 +11,12 @@ the real interceptor and handlers via `task_helpers`. from __future__ import annotations +import asyncio from typing import Any import mcp_types +from fastmcp_tasks.context import get_task_scope +from fastmcp_tasks.input_store import acquire_update_lock, release_update_lock from mcp.shared.exceptions import MCPError from mcp_types import INTERNAL_ERROR @@ -311,6 +314,77 @@ async def test_partial_update_keeps_task_parked_on_remaining_request(): assert final.result["content"][0]["text"] == "one+two" +async def test_partial_update_waits_for_a_held_update_lock(): + """An update that arrives while another holds the lock must still land. + + SEP-2663 invites a client to answer a multi-request ask one key at a time, + so two updates can be in flight carrying *different* answers. Acknowledging + the one that loses the lock without storing its answer would leave the task + waiting forever on a key the client believes it already sent. + + The lock is taken out of band here so the contention is deterministic rather + than dependent on scheduling. + """ + mcp = FastMCP("lock-contention") + mcp.add_extension(TasksExtension()) + + @mcp.tool(task=True) + async def two_questions(ctx: Context) -> str | mcp_types.InputRequiredResult: + responses = ctx.input_responses + if responses is None: + return _input_required( + { + "first": _elicit_request("First?"), + "second": _elicit_request("Second?"), + } + ) + return f"{_answer(responses, 'first')}+{_answer(responses, 'second')}" + + async with running_task_server(mcp): + created = await submit_task(mcp, "two_questions", {}) + parked = await wait_for_task( + mcp, created.task_id, target_states=frozenset({"input_required"}) + ) + assert parked.input_requests is not None + first = _key_asking(parked.input_requests, "First?") + second = _key_asking(parked.input_requests, "Second?") + + docket = mcp._docket + assert docket is not None + scope = get_task_scope() + + # Simulate a concurrent update in progress. + assert await acquire_update_lock(docket, scope, created.task_id) + pending = asyncio.create_task( + update_task( + mcp, + created.task_id, + {first: {"action": "accept", "content": {"value": "one"}}}, + ) + ) + await asyncio.sleep(0.05) + assert not pending.done(), "update returned while the lock was held" + await release_update_lock(docket, scope, created.task_id) + await pending + + # The blocked answer landed, so only the other key remains outstanding. + still_parked = await get_task(mcp, created.task_id) + assert still_parked.status == "input_required" + assert still_parked.input_requests is not None + assert list(still_parked.input_requests) == [second] + + await update_task( + mcp, + created.task_id, + {second: {"action": "accept", "content": {"value": "two"}}}, + ) + final = await wait_for_task(mcp, created.task_id) + + assert final.status == "completed" + assert final.result is not None + assert final.result["content"][0]["text"] == "one+two" + + async def test_final_answer_keeps_task_parked_until_next_leg_is_durable(): """The last answer must not retire its outstanding marker early.