Clear stale auth in reused workers; bound elicitation; version explicit tasks

Three review fixes. A Docket worker may reuse an asyncio context across tasks,
so snapshot restore now always resets auth and headers to the current task's
state — an anonymous task following an authenticated one no longer inherits the
prior caller's identity. A stalled in-task elicitation handler is now bounded by
the call's remaining timeout, like polling and sleeps. And call_tool_task takes
a version= to task a specific component version rather than the highest.

fastmcp_tasks/fastmcp_tasks/client.py

@@ -168,6 +168,22 @@ async def _answer_input_requests(
             "ask for input."
         )
 
+    # Bound the whole answer phase — elicitation callbacks included — by the
+    # call's remaining budget: a stalled handler must not outlast `timeout=N`
+    # any more than a stalled poll does, matching the synchronous path.
+    loop = asyncio.get_event_loop()
+    deadline = (
+        None if read_timeout_seconds is None else loop.time() + read_timeout_seconds
+    )
+
+    def _remaining() -> float | None:
+        if deadline is None:
+            return None
+        left = deadline - loop.time()
+        if left <= 0:
+            raise TimeoutError(f"Task {task_id} timed out awaiting input")
+        return left
+
     responses: dict[str, Any] = {}
     for surfaced_key, payload in input_requests.items():
         method = payload.get("method") if isinstance(payload, dict) else None
@@ -181,14 +197,16 @@ async def _answer_input_requests(
         context = ClientRequestContext(
             session=session, request_id=f"task-{task_id}-{surfaced_key}"
         )
-        answer = await elicitation_callback(context, request.params)
+        budget = _remaining()
+        call = elicitation_callback(context, request.params)
+        answer = await (asyncio.wait_for(call, budget) if budget is not None else call)
         if isinstance(answer, mcp_types.ErrorData):
             raise ToolError(f"Elicitation for task {task_id} failed: {answer.message}")
         responses[surfaced_key] = answer.model_dump(
             by_alias=True, mode="json", exclude_none=True
         )
 
-    await _send_update(session, task_id, responses, read_timeout_seconds)
+    await _send_update(session, task_id, responses, _remaining())
 
 
 # ---------------------------------------------------------------------------
@@ -475,6 +493,7 @@ async def call_tool_task(
     *,
     timeout: float | int | None = None,
     raise_on_error: bool = True,
+    version: str | None = None,
     meta: dict[str, Any] | None = None,
 ) -> ToolTask:
     """Call a tool as a background task and return a `ToolTask` handle immediately.
@@ -484,9 +503,18 @@ async def call_tool_task(
     work and drive the task through the handle. Requires the server to run the
     call as a task (a `task=True` tool on a task-serving backend); a call the
     server runs synchronously raises `ToolError`.
+
+    `version` targets a specific component version, the same as
+    `client.call_tool(..., version=...)`: the server tasks that version rather
+    than the highest. It is carried in the request metadata FastMCP reads.
     """
     read_timeout_seconds = normalize_timeout_to_seconds(timeout)
-    request_meta = cast("mcp_types.RequestParamsMeta | None", meta)
+    combined_meta: dict[str, Any] = dict(meta) if meta else {}
+    if version is not None:
+        fastmcp_meta = dict(combined_meta.get("fastmcp") or {})
+        fastmcp_meta["version"] = version
+        combined_meta["fastmcp"] = fastmcp_meta
+    request_meta = cast("mcp_types.RequestParamsMeta | None", combined_meta or None)
     raw = await client._await_with_session_monitoring(
         client.session.call_tool(
             name=name,

fastmcp_tasks/fastmcp_tasks/context.py

@@ -458,27 +458,34 @@ def _apply_snapshot_to_context(snapshot: TaskContextSnapshot) -> None:
     still raise inside a task — there is no request. Runs inside
     ``restore_task_snapshot`` (a Docket dependency), whose context vars propagate
     to the tool the same way the snapshot var already does.
+
+    Both vars are set unconditionally to *this* snapshot's state (``None`` when
+    it carries no token/headers), never left as-is: a Docket worker may reuse an
+    asyncio context across tasks, so an anonymous task following an authenticated
+    one must not inherit the prior caller's identity or headers.
     """
+    import time
+
+    from mcp.server.auth.middleware.auth_context import auth_context_var
+    from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser
+
+    from fastmcp.server.auth import AccessToken
+    from fastmcp.server.dependencies import _background_task_headers
+
+    user: AuthenticatedUser | None = None
     if snapshot.access_token_json is not None:
-        import time
-
-        from mcp.server.auth.middleware.auth_context import auth_context_var
-        from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser
-
-        from fastmcp.server.auth import AccessToken
-
         token = AccessToken.model_validate_json(snapshot.access_token_json)
         # A task may sit queued past its submitter's token expiry. Install it
         # only if still valid — mirroring the SDK's bearer check — so a delayed
         # task never runs under credentials a live request would reject (401).
         # An expired token leaves the worker unauthenticated, the honest state.
         if token.expires_at is None or token.expires_at >= int(time.time()):
-            auth_context_var.set(AuthenticatedUser(token))
+            user = AuthenticatedUser(token)
+    auth_context_var.set(user)
 
-    if snapshot.http_headers:
-        from fastmcp.server.dependencies import _background_task_headers
-
-        _background_task_headers.set(dict(snapshot.http_headers))
+    _background_task_headers.set(
+        dict(snapshot.http_headers) if snapshot.http_headers else None
+    )
 
 
 async def make_task_context() -> Context | None:

tests/tasks/client/test_transparent_tasks.py

@@ -80,6 +80,26 @@ async def test_failed_task_raises_tool_error(task_server: FastMCP):
             await client.call_tool("boom", {})
 
 
+async def test_call_tool_task_forwards_requested_version():
+    """`call_tool_task(..., version=...)` tasks the requested version, not the highest."""
+    mcp = FastMCP("versioned-task-client")
+    mcp.add_extension(TasksExtension())
+
+    @mcp.tool(name="pick", version="1.0", task=True)
+    async def pick_v1() -> str:
+        return "v1"
+
+    @mcp.tool(name="pick", version="2.0", task=True)
+    async def pick_v2() -> str:
+        return "v2"
+
+    async with Client(mcp, mode="auto") as client:
+        task = await call_tool_task(client, "pick", version="1.0")
+        result = await task.result()
+
+    assert result.data == "v1"
+
+
 async def test_raw_create_task_result_is_exposed(task_server: FastMCP):
     """The raw claimed CreateTaskResult is reachable via the session/handle path."""
     async with Client(task_server, mode="auto") as client:
@@ -174,6 +194,25 @@ async def test_in_task_input_without_handler_errors(guard_server: FastMCP):
             await client.call_tool("plan_dinner", {})
 
 
+async def test_call_tool_timeout_bounds_a_stalled_elicitation(guard_server: FastMCP):
+    """A stalled elicitation handler cannot outlast the call's timeout.
+
+    The deadline covers the whole drive, elicitation callbacks included: a
+    handler that hangs must abort the tasked call once `timeout=N` elapses,
+    matching the synchronous path rather than blocking forever inside the
+    callback.
+    """
+
+    async def slow_elicitation(message, response_type, params, context):
+        await asyncio.sleep(5)
+        return DinnerPrefs(cuisine="Thai", vegetarian=True)
+
+    client = Client(guard_server, mode="auto", elicitation_handler=slow_elicitation)
+    async with client:
+        with pytest.raises((TimeoutError, ToolError, MCPError)):
+            await client.call_tool("plan_dinner", {}, timeout=0.3)
+
+
 async def test_in_task_input_answered_by_handler_set_after_construction(
     guard_server: FastMCP,
 ):

tests/tasks/server/test_snapshot_restore.py

@@ -169,6 +169,33 @@ def test_apply_snapshot_skips_expired_token():
     contextvars.copy_context().run(run_in_clean_worker_context)
 
 
+def test_apply_snapshot_clears_prior_auth_in_reused_context():
+    """An anonymous task must not inherit a prior task's identity or headers.
+
+    A Docket worker may reuse an asyncio context across executions. Applying a
+    tokenless snapshot after an authenticated one must clear the earlier
+    caller's `auth_context_var` and headers rather than leave them installed.
+    """
+    prior = AccessToken(token="jwt-prior", client_id="prior-client", scopes=["read"])
+    authed = TaskContextSnapshot(
+        access_token_json=prior.model_dump_json(),
+        http_headers={"x-trace-id": "prior"},
+    )
+    anonymous = TaskContextSnapshot()
+
+    def run_in_reused_worker_context() -> None:
+        _apply_snapshot_to_context(authed)
+        assert get_access_token() is not None
+        assert get_http_headers()["x-trace-id"] == "prior"
+
+        # Same context, next task carries no auth/headers.
+        _apply_snapshot_to_context(anonymous)
+        assert get_access_token() is None
+        assert get_http_headers() == {}
+
+    contextvars.copy_context().run(run_in_reused_worker_context)
+
+
 async def test_restore_failure_is_nonfatal():
     """If deserialization blows up, the task still runs to completion and
     the snapshot cache stays empty."""