From c3ad5e9ecb93d24662be8d865a51910aee32a9c3 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Thu, 23 Jul 2026 18:25:48 -0400 Subject: [PATCH] Clear stale auth in reused workers; bound elicitation; version explicit tasks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 | 34 +++++++++++++++-- fastmcp_tasks/fastmcp_tasks/context.py | 31 ++++++++++------ tests/tasks/client/test_transparent_tasks.py | 39 ++++++++++++++++++++ tests/tasks/server/test_snapshot_restore.py | 27 ++++++++++++++ 4 files changed, 116 insertions(+), 15 deletions(-) diff --git a/fastmcp_tasks/fastmcp_tasks/client.py b/fastmcp_tasks/fastmcp_tasks/client.py index 006104e02..19def1a46 100644 --- a/fastmcp_tasks/fastmcp_tasks/client.py +++ b/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, diff --git a/fastmcp_tasks/fastmcp_tasks/context.py b/fastmcp_tasks/fastmcp_tasks/context.py index 9975c088f..22bdf0259 100644 --- a/fastmcp_tasks/fastmcp_tasks/context.py +++ b/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: diff --git a/tests/tasks/client/test_transparent_tasks.py b/tests/tasks/client/test_transparent_tasks.py index 7af25f10f..0c7dd046b 100644 --- a/tests/tasks/client/test_transparent_tasks.py +++ b/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, ): diff --git a/tests/tasks/server/test_snapshot_restore.py b/tests/tasks/server/test_snapshot_restore.py index 422b90114..aeb3e97d9 100644 --- a/tests/tasks/server/test_snapshot_restore.py +++ b/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."""