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.

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"

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.