From 74e01d5e08cc3c4d63642a82fb54e8061b31e0cb Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Wed, 22 Jul 2026 15:21:10 -0400 Subject: [PATCH] Add SEP-2663 client half: transparent call_tool, ResultClaim, Task handle A FastMCP client now transparently completes tasked tools/call: the tasks ClientExtension advertises the capability and claims the CreateTaskResult, and the resolver drives the tasks/get poll loop to completion, answering in-task input through the client's elicitation handler and returning the tool's real result. call_tool is transparent, call_tool_mcp exposes the raw result, and call_tool_task yields a Task handle. The client half moves to fastmcp-tasks; the [tasks] client extension auto-wires into Client (ProxyClient opts out). Co-Authored-By: Claude --- examples/task_elicitation.py | 84 +- examples/tasks/client.py | 160 ++- examples/tasks/server.py | 6 +- fastmcp_slim/fastmcp/client/client.py | 55 +- .../fastmcp/client/extension_hooks.py | 65 ++ .../fastmcp/server/providers/proxy.py | 6 + fastmcp_slim/fastmcp/settings.py | 20 - fastmcp_tasks/fastmcp_tasks/__init__.py | 10 +- .../fastmcp_tasks/_client_task_management.py | 232 ----- fastmcp_tasks/fastmcp_tasks/client.py | 951 +++++++----------- fastmcp_tasks/fastmcp_tasks/client_models.py | 130 +++ fastmcp_tasks/fastmcp_tasks/settings.py | 35 + pyproject.toml | 3 - tests/client/client/test_client.py | 40 +- .../telemetry/test_client_task_tracing.py | 98 -- tests/client/test_client_extensions.py | 185 ++-- .../client/test_client_task_notifications.py | 283 ------ .../tasks/client/test_client_task_protocol.py | 89 -- tests/tasks/client/test_client_tool_tasks.py | 202 ++-- tests/tasks/client/test_poll_interval.py | 99 +- .../client/test_task_context_validation.py | 224 ----- .../tasks/client/test_task_result_caching.py | 341 ------- tests/tasks/client/test_transparent_tasks.py | 158 +++ 23 files changed, 1174 insertions(+), 2302 deletions(-) create mode 100644 fastmcp_slim/fastmcp/client/extension_hooks.py delete mode 100644 fastmcp_tasks/fastmcp_tasks/_client_task_management.py create mode 100644 fastmcp_tasks/fastmcp_tasks/client_models.py delete mode 100644 tests/client/telemetry/test_client_task_tracing.py delete mode 100644 tests/tasks/client/test_client_task_notifications.py delete mode 100644 tests/tasks/client/test_client_task_protocol.py delete mode 100644 tests/tasks/client/test_task_context_validation.py delete mode 100644 tests/tasks/client/test_task_result_caching.py create mode 100644 tests/tasks/client/test_transparent_tasks.py diff --git a/examples/task_elicitation.py b/examples/task_elicitation.py index 510e21fdf..18d56f2e7 100644 --- a/examples/task_elicitation.py +++ b/examples/task_elicitation.py @@ -1,8 +1,16 @@ """ -Background task elicitation demo. +Background task input demo (SEP-2663 guard pattern). -A background task (Docket) that pauses mid-execution to ask the user a -question, waits for the answer, then resumes and finishes. +A background task that pauses to ask the user a question, waits for the answer, +then resumes and finishes. Under SEP-2663 a task gathers input by the *guard +pattern*: instead of awaiting `ctx.elicit()` (which would block a worker), the +tool *returns* an `InputRequiredResult`. That ends the leg; the client answers +via the tasks protocol; the framework re-runs the tool with the answer on +`ctx.input_responses`. No worker is ever blocked. + +The client side is transparent: `client.call_tool(...)` drives the whole +round-trip β€” poll, answer via the `elicitation_handler`, poll again β€” and returns +the finished result. Works with both in-memory and Redis backends: @@ -22,13 +30,15 @@ Requires the `docket` extra (included in dev dependencies). import asyncio from dataclasses import dataclass +import mcp_types from mcp_types import TextContent from fastmcp import Context, FastMCP from fastmcp.client import Client -from fastmcp.server.elicitation import AcceptedElicitation +from fastmcp_tasks import TasksExtension mcp = FastMCP("Task Elicitation Demo") +mcp.add_extension(TasksExtension()) @dataclass @@ -37,44 +47,60 @@ class DinnerPrefs: vegetarian: bool -@mcp.tool(task=True) -async def plan_dinner(ctx: Context) -> str: - """Plan a dinner menu, asking the user what they're in the mood for.""" - - await ctx.report_progress(0, 2, "Asking what you'd like...") - - result = await ctx.elicit( - "What kind of dinner are you in the mood for?", - response_type=DinnerPrefs, +def _ask_dinner_prefs() -> mcp_types.InputRequiredResult: + """Return the input request that pauses the task until the client answers.""" + request = mcp_types.ElicitRequest( + params=mcp_types.ElicitRequestFormParams( + message="What kind of dinner are you in the mood for?", + requested_schema={ + "type": "object", + "properties": { + "cuisine": {"type": "string"}, + "vegetarian": {"type": "boolean"}, + }, + "required": ["cuisine", "vegetarian"], + }, + ) + ) + return mcp_types.InputRequiredResult( + result_type="input_required", + input_requests={"prefs": request}, ) - if not isinstance(result, AcceptedElicitation): + +@mcp.tool(task=True) +async def plan_dinner(ctx: Context) -> str | mcp_types.InputRequiredResult: + """Plan a dinner menu, asking the user what they're in the mood for.""" + responses = ctx.input_responses + if responses is None: + # First leg: ask for preferences and end the leg. + return _ask_dinner_prefs() + + # Re-entered leg: the client's answer is on ctx.input_responses. + answer = responses["prefs"] + assert isinstance(answer, mcp_types.ElicitResult) + if answer.action != "accept" or answer.content is None: return "Dinner cancelled!" - prefs = result.data - assert isinstance(prefs, DinnerPrefs) - await ctx.report_progress(1, 2, "Planning your menu...") - await asyncio.sleep(1) - await ctx.report_progress(2, 2, "Done!") - - veg = "vegetarian " if prefs.vegetarian else "" - return f"Tonight's menu: a lovely {veg}{prefs.cuisine} dinner!" + await asyncio.sleep(1) # "planning the menu" + veg = "vegetarian " if answer.content["vegetarian"] else "" + return f"Tonight's menu: a lovely {veg}{answer.content['cuisine']} dinner!" async def handle_elicitation(message, response_type, params, context): - """Handle elicitation requests from background tasks.""" + """Answer elicitation requests raised by the background task.""" print(f" Server asks: {message}") print(" Responding with: cuisine=Thai, vegetarian=True") return DinnerPrefs(cuisine="Thai", vegetarian=True) async def main(): - async with Client(mcp, elicitation_handler=handle_elicitation) as client: - print("Starting background task...") - task = await client.call_tool("plan_dinner", {}, task=True) - print(f" task_id = {task.task_id}\n") - - result = await task.result() + client = Client(mcp, mode="auto", elicitation_handler=handle_elicitation) + async with client: + print("Calling plan_dinner (runs as a background task)...") + # call_tool drives the whole round-trip transparently: it polls, answers + # the task's input request via handle_elicitation, and returns the result. + result = await client.call_tool("plan_dinner", {}) assert isinstance(result.content[0], TextContent) print(f"\nResult: {result.content[0].text}") diff --git a/examples/tasks/client.py b/examples/tasks/client.py index c8581667f..1c039d6f7 100644 --- a/examples/tasks/client.py +++ b/examples/tasks/client.py @@ -1,18 +1,23 @@ """ -FastMCP Tasks Example Client +FastMCP Tasks Example Client (SEP-2663) -Demonstrates calling tools both immediately and as background tasks, -with real-time progress updates via status callbacks. +Demonstrates the two client task surfaces: + +- Transparent: `client.call_tool(...)` drives the background task to completion + under the hood and returns the tool's real result. The caller writes ordinary + tool-call code and never sees that the server ran the call as a task. +- Explicit handle: `call_tool_task(...)` returns a `ToolTask` immediately, so the + client can do other work and poll the task itself before collecting the result. Usage: # Make sure environment is configured (source .envrc or use direnv) source .envrc - # Background task execution with progress callbacks (default) + # Transparent background task (default) python client.py --duration 10 - # Immediate execution (blocks until complete) - python client.py immediate --duration 5 + # Return-quickly handle, driven by the client + python client.py handle --duration 5 """ import asyncio @@ -21,10 +26,11 @@ from pathlib import Path from typing import Annotated import cyclopts -from mcp_types import GetTaskResult, TextContent +from mcp_types import TextContent from rich.console import Console from fastmcp.client import Client +from fastmcp_tasks import call_tool_task console = Console() app = cyclopts.App(name="tasks-client", help="FastMCP Tasks Example Client") @@ -41,52 +47,14 @@ def load_server(): return server_module.mcp -# Track last message to deduplicate consecutive identical notifications -# Note: Docket fires separate events for progress.increment() and progress.set_message(), -# but MCP's status_message field only carries the text message (no numerical progress). -# This means we often get duplicate notifications with identical messages. -_last_notification_message = None - - -def print_notification(status: GetTaskResult) -> None: - """Callback function for push notifications from server. - - This is called automatically when the server sends notifications/tasks/status. - Deduplicates identical consecutive messages to keep output clean. - """ - global _last_notification_message - - # Skip if this is the same message we just printed - if status.status_message == _last_notification_message: - return - - _last_notification_message = status.status_message - - color = { - "working": "yellow", - "completed": "green", - "failed": "red", - }.get(status.status, "yellow") - - icon = { - "working": "πŸš€", - "completed": "βœ…", - "failed": "❌", - }.get(status.status, "⚠️") - - console.print( - f"[{color}]πŸ“’ Notification: {status.status} {icon} - {status.status_message}[/{color}]" - ) - - @app.default -async def task( +async def transparent( duration: Annotated[ int, cyclopts.Parameter(help="Duration of computation in seconds (1-60)"), ] = 10, ): - """Execute as background task with real-time progress callbacks.""" + """Call the tool transparently: the client drives the task to completion.""" if duration < 1 or duration > 60: console.print("[red]Error: Duration must be between 1 and 60 seconds[/red]") sys.exit(1) @@ -94,58 +62,11 @@ async def task( server = load_server() console.print(f"\n[bold]Calling slow_computation(duration={duration})[/bold]") - console.print("Mode: [cyan]Background task[/cyan]\n") + console.print("Mode: [cyan]Transparent (server may run it as a task)[/cyan]\n") - async with Client(server) as client: - task_obj = await client.call_tool( - "slow_computation", - arguments={"duration": duration}, - task=True, - ) - - console.print(f"Task started: [cyan]{task_obj.task_id}[/cyan]\n") - - # Register callback for real-time push notifications - task_obj.on_status_change(print_notification) - - console.print( - "[dim]Notifications will appear as the server sends them...[/dim]\n" - ) - - # Do other work while task runs in background - for i in range(3): - await asyncio.sleep(0.5) - console.print(f"[dim]Client doing other work... ({i + 1}/3)[/dim]") - - console.print() - - # Wait for task to complete - console.print("[dim]Waiting for final result...[/dim]") - result = await task_obj.result() - - console.print("\n[bold]Result:[/bold]") - assert isinstance(result.content[0], TextContent) - console.print(f" {result.content[0].text}") - - -@app.command -async def immediate( - duration: Annotated[ - int, - cyclopts.Parameter(help="Duration of computation in seconds (1-60)"), - ] = 5, -): - """Execute the tool immediately (blocks until complete).""" - if duration < 1 or duration > 60: - console.print("[red]Error: Duration must be between 1 and 60 seconds[/red]") - sys.exit(1) - - server = load_server() - - console.print(f"\n[bold]Calling slow_computation(duration={duration})[/bold]") - console.print("Mode: [cyan]Immediate execution[/cyan]\n") - - async with Client(server) as client: + # mode="auto" negotiates the modern protocol, so the server may run the call + # as a background task; the client resolves it transparently. + async with Client(server, mode="auto") as client: result = await client.call_tool( "slow_computation", arguments={"duration": duration}, @@ -156,5 +77,48 @@ async def immediate( console.print(f" {result.content[0].text}") +@app.command +async def handle( + duration: Annotated[ + int, + cyclopts.Parameter(help="Duration of computation in seconds (1-60)"), + ] = 5, +): + """Use the explicit handle: return immediately, then drive the task.""" + if duration < 1 or duration > 60: + console.print("[red]Error: Duration must be between 1 and 60 seconds[/red]") + sys.exit(1) + + server = load_server() + + console.print(f"\n[bold]Calling slow_computation(duration={duration})[/bold]") + console.print("Mode: [cyan]Explicit ToolTask handle[/cyan]\n") + + async with Client(server, mode="auto") as client: + task = await call_tool_task( + client, + "slow_computation", + arguments={"duration": duration}, + ) + + console.print(f"Task started: [cyan]{task.task_id}[/cyan]\n") + + # Do other work while the task runs in the background. + for i in range(3): + await asyncio.sleep(0.5) + status = await task.status() + console.print( + f"[dim]Client doing other work... ({i + 1}/3) " + f"β€” task is {status.status}[/dim]" + ) + + console.print("\n[dim]Waiting for the final result...[/dim]") + result = await task.result() + + console.print("\n[bold]Result:[/bold]") + assert isinstance(result.content[0], TextContent) + console.print(f" {result.content[0].text}") + + if __name__ == "__main__": app() diff --git a/examples/tasks/server.py b/examples/tasks/server.py index 77b3cde82..8405ab717 100644 --- a/examples/tasks/server.py +++ b/examples/tasks/server.py @@ -20,13 +20,17 @@ from docket import Logged from fastmcp import FastMCP from fastmcp.dependencies import Progress +from fastmcp_tasks import TasksExtension # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) -# Create server +# Create server and enable background tasks (SEP-2663). The extension reads the +# FASTMCP_DOCKET_* environment for its backend (memory:// by default, Redis for +# distributed execution). mcp = FastMCP("Tasks Example") +mcp.add_extension(TasksExtension()) @mcp.tool(task=True) diff --git a/fastmcp_slim/fastmcp/client/client.py b/fastmcp_slim/fastmcp/client/client.py index 74cb76055..28c02f1e1 100644 --- a/fastmcp_slim/fastmcp/client/client.py +++ b/fastmcp_slim/fastmcp/client/client.py @@ -41,7 +41,11 @@ from mcp.client.extension import ( NotificationBinding, ResultClaim, ) -from mcp.client.session import ClientRequestContext, MessageHandlerFnT +from mcp.client.session import ( + ClientRequestContext, + ElicitationFnT, + MessageHandlerFnT, +) from mcp_types.methods import validate_server_result from mcp_types.version import HANDSHAKE_PROTOCOL_VERSIONS, MODERN_PROTOCOL_VERSIONS from pydantic import AnyUrl, ValidationError @@ -52,6 +56,7 @@ from fastmcp.client.elicitation import ( ElicitationHandler, create_elicitation_callback, ) +from fastmcp.client.extension_hooks import build_internal_client_extensions from fastmcp.client.logging import ( LogHandler, create_log_callback, @@ -334,6 +339,13 @@ class Client( ``` """ + #: Whether FastMCP-internal client extensions (e.g. the tasks extension) are + #: folded in automatically at construction. `ProxyClient` overrides this to + #: `False`: a proxy forwards calls and must not advertise task support to its + #: backend, since proxied tools run synchronously (forbidden mode) and the + #: proxy has no path to drive a backend task on the front connection's behalf. + _auto_internal_extensions: bool = True + @overload def __init__(self: Client[T], transport: T, *args: Any, **kwargs: Any) -> None: ... @@ -504,6 +516,16 @@ class Client( # `_build_extension_kwargs`. self._claim_by_model: dict[type[mcp_types.Result], ResultClaim[Any]] = {} + # Build the elicitation callback up front: it is threaded both into the + # session (to answer server-initiated elicitation) and into the internal + # client extensions (so a task resolver can answer in-task input), and + # `_build_extension_kwargs` β€” called below β€” needs it. + self._elicitation_callback: ElicitationFnT | None = ( + create_elicitation_callback(elicitation_handler) + if elicitation_handler is not None + else None + ) + self._session_kwargs: SessionKwargs = { "sampling_callback": None, "list_roots_callback": None, @@ -527,10 +549,8 @@ class Client( else mcp_types.SamplingCapability() ) - if elicitation_handler is not None: - self._session_kwargs["elicitation_callback"] = create_elicitation_callback( - elicitation_handler - ) + if self._elicitation_callback is not None: + self._session_kwargs["elicitation_callback"] = self._elicitation_callback # Maximum time to wait for a clean disconnect before giving up. # Normally disconnects complete in <100ms; this is a safety net for @@ -1191,8 +1211,31 @@ class Client( Also rebuilds `self._claim_by_model`, the modelβ†’claim index the resolution path uses to finish a claimed `tools/call` result, covering both the folded extension claims and the explicit `result_claims` extras. + + FastMCP-internal extensions (e.g. the tasks extension from `fastmcp-tasks`, + registered via `register_internal_client_extension_factory`) are folded in + automatically so an ordinary `Client` transparently drives a server's + background tasks. They lead the fold order; a user extension declaring the + same identifier wins, so the internal one is dropped rather than colliding. """ - folded = _fold_extensions(self._extensions_arg) + user_extensions = list(self._extensions_arg or ()) + user_identifiers = { + identifier + for extension in user_extensions + if (identifier := getattr(extension, "identifier", None)) is not None + } + internal_extensions = ( + [ + extension + for extension in build_internal_client_extensions( + self._elicitation_callback + ) + if extension.identifier not in user_identifiers + ] + if self._auto_internal_extensions + else [] + ) + folded = _fold_extensions([*internal_extensions, *user_extensions]) claims: dict[str, tuple[ResultClaim[Any], ...]] = dict(folded.claims or {}) by_model: dict[type[mcp_types.Result], ResultClaim[Any]] = dict(folded.by_model) diff --git a/fastmcp_slim/fastmcp/client/extension_hooks.py b/fastmcp_slim/fastmcp/client/extension_hooks.py new file mode 100644 index 000000000..02f367be1 --- /dev/null +++ b/fastmcp_slim/fastmcp/client/extension_hooks.py @@ -0,0 +1,65 @@ +"""Registry for FastMCP-internal client extensions (SEP-2133). + +Core ships the client wiring for opt-in extensions but no extension of its own. +A companion package (``fastmcp-tasks``) provides an extension the ``Client`` +should register *automatically* β€” so an ordinary ``Client(url)`` transparently +drives a server's background tasks without the caller passing anything. The +package cannot reach into core's ``Client`` constructor, so core exposes this +hook instead: the package registers a factory on import, and ``Client`` folds +the factory's extension in alongside the user's own. + +This mirrors the server-side ``set_background_context_factory`` hook: core +declares the extension point, the tasks package fills it. With no package +imported, the registry is empty and ``Client`` behaves exactly as before. + +A factory receives the client's elicitation callback (so a task resolver can +answer in-task input prompts) and returns a ``ClientExtension`` to register, or +``None`` to contribute nothing for this client. +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from mcp.client.extension import ClientExtension + from mcp.client.session import ElicitationFnT + +#: A factory that builds a FastMCP-internal client extension for one ``Client``, +#: given that client's elicitation callback (``None`` when the client has no +#: elicitation handler). +InternalClientExtensionFactory = Callable[ + ["ElicitationFnT | None"], "ClientExtension | None" +] + +_internal_client_extension_factories: list[InternalClientExtensionFactory] = [] + + +def register_internal_client_extension_factory( + factory: InternalClientExtensionFactory, +) -> None: + """Register a factory whose extension every ``Client`` folds in automatically. + + Idempotent: registering the same factory object twice is a no-op, so a + package importing more than once does not double-register. + """ + if factory not in _internal_client_extension_factories: + _internal_client_extension_factories.append(factory) + + +def build_internal_client_extensions( + elicitation_callback: ElicitationFnT | None, +) -> list[ClientExtension]: + """Build the internal extensions to fold into a ``Client`` under construction. + + Each registered factory is invoked with the client's elicitation callback; + factories that return ``None`` contribute nothing. Empty when no package has + registered a factory (plain core). + """ + extensions: list[ClientExtension] = [] + for factory in _internal_client_extension_factories: + extension = factory(elicitation_callback) + if extension is not None: + extensions.append(extension) + return extensions diff --git a/fastmcp_slim/fastmcp/server/providers/proxy.py b/fastmcp_slim/fastmcp/server/providers/proxy.py index 866322177..2011e160a 100644 --- a/fastmcp_slim/fastmcp/server/providers/proxy.py +++ b/fastmcp_slim/fastmcp/server/providers/proxy.py @@ -1353,6 +1353,12 @@ class ProxyClient(Client[ClientTransportT]): _proxy_rc_ref: list[Any] _proxy_restoring_handler_keys: set[str] + # A proxy forwards calls; it must not advertise task support to its backend. + # Proxied tools run synchronously (forbidden mode), and the proxy has no path + # to drive a backend task on the front connection's behalf, so the internal + # tasks client extension is not folded into a proxy's backend client. + _auto_internal_extensions: bool = False + def __init__( self, transport: ClientTransportT diff --git a/fastmcp_slim/fastmcp/settings.py b/fastmcp_slim/fastmcp/settings.py index 17cd884c8..352b383e2 100644 --- a/fastmcp_slim/fastmcp/settings.py +++ b/fastmcp_slim/fastmcp/settings.py @@ -181,26 +181,6 @@ class Settings(BaseSettings): ), ] = 5 - # May move to the fastmcp-tasks package alongside the client task senders - # when client task support is rebuilt on the SEP-2663 extension. - client_task_poll_interval: Annotated[ - float, - Field( - description=inspect.cleandoc( - """ - Ceiling, in seconds, for the fallback poll backoff while waiting on a - background task (SEP-1686). Applies only when the server does not - advertise its own pollInterval: in that case Task.wait() starts polling - fast (~20ms) and doubles up to this ceiling, so quick tasks resolve - promptly while long-running tasks don't hammer the server. When the - server does advertise a pollInterval, that interval is honored exactly - and this setting is ignored. Must be positive. - """ - ), - gt=0, - ), - ] = 0.5 - # Transport settings transport: Literal["stdio", "http", "sse", "streamable-http"] = "stdio" diff --git a/fastmcp_tasks/fastmcp_tasks/__init__.py b/fastmcp_tasks/fastmcp_tasks/__init__.py index 32c500f0e..7d6420426 100644 --- a/fastmcp_tasks/fastmcp_tasks/__init__.py +++ b/fastmcp_tasks/fastmcp_tasks/__init__.py @@ -2,6 +2,8 @@ from importlib.metadata import PackageNotFoundError, version +from fastmcp.client.extension_hooks import register_internal_client_extension_factory +from fastmcp_tasks.client import ToolTask, _build_tasks_client_extension, call_tool_task from fastmcp_tasks.extension import TasksExtension try: @@ -9,4 +11,10 @@ try: except PackageNotFoundError: __version__ = "0.0.0" -__all__ = ["TasksExtension", "__version__"] +# Register the client half so every FastMCP `Client` transparently drives a +# task-serving backend's background tasks (see `fastmcp_tasks.client`). Importing +# this package β€” which any task deployment does, server or client side β€” is what +# turns on client task support. +register_internal_client_extension_factory(_build_tasks_client_extension) + +__all__ = ["TasksExtension", "ToolTask", "call_tool_task", "__version__"] diff --git a/fastmcp_tasks/fastmcp_tasks/_client_task_management.py b/fastmcp_tasks/fastmcp_tasks/_client_task_management.py deleted file mode 100644 index 8b3617d3f..000000000 --- a/fastmcp_tasks/fastmcp_tasks/_client_task_management.py +++ /dev/null @@ -1,232 +0,0 @@ -"""Task management methods for FastMCP Client.""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, Any, cast - -import mcp_types -from mcp import MCPError -from mcp_types import Result -from pydantic import ConfigDict - -if TYPE_CHECKING: - from fastmcp.client.client import Client -from mcp_types import ( - CancelTaskRequest, - CancelTaskRequestParams, - GetTaskPayloadRequest, - GetTaskPayloadRequestParams, - GetTaskRequest, - GetTaskRequestParams, - GetTaskResult, - ListTasksRequest, - PaginatedRequestParams, -) - -from fastmcp.client.telemetry import client_span -from fastmcp.telemetry import inject_trace_context -from fastmcp.utilities.logging import get_logger - -logger = get_logger(__name__) - - -class _RawTaskPayloadResult(Result): - """Permissive result type for `tasks/result` responses. - - Per the v2 spec, a `tasks/result` payload arrives as extra wire fields whose - shape matches the original request's result type (CallToolResult, - GetPromptResult, ReadResourceResult, ...). `GetTaskPayloadResult` is a bare - `Result` that drops those fields on validation, so this subclass retains them - with `extra="allow"`; callers re-parse the resulting dict into the concrete - result type. - """ - - model_config = ConfigDict( - alias_generator=Result.model_config.get("alias_generator"), - populate_by_name=True, - extra="allow", - ) - - -class ClientTaskManagementMixin: - """Mixin providing task management methods for Client.""" - - async def get_task_status(self: Client, task_id: str) -> GetTaskResult: - """Query the status of a background task. - - Sends a 'tasks/get' MCP protocol request over the existing transport. - - Args: - task_id: The task ID returned from call_tool_as_task - - Returns: - GetTaskResult: Status information including taskId, status, pollInterval, etc. - - Raises: - RuntimeError: If client not connected - MCPError: If the request results in a TimeoutError | JSONRPCError - """ - with client_span( - "tasks/get", - "tasks/get", - task_id, - session_id=self.transport.get_session_id(), - ): - request_meta = cast( - "mcp_types.RequestParamsMeta | None", inject_trace_context() - ) - request = GetTaskRequest( - params=GetTaskRequestParams( - task_id=task_id, - _meta=request_meta, # type: ignore[unknown-argument] - ) - ) - return await self._await_with_session_monitoring( - self.session.send_request( - request=request, # type: ignore[arg-type] - result_type=GetTaskResult, - ) - ) - - async def get_task_result(self: Client, task_id: str) -> Any: - """Retrieve the raw result of a completed background task. - - Sends a 'tasks/result' MCP protocol request over the existing transport. - Returns the raw result - callers should parse it appropriately. - - Args: - task_id: The task ID returned from call_tool_as_task - - Returns: - Any: The raw result (could be tool, prompt, or resource result) - - Raises: - RuntimeError: If client not connected, task not found, or task failed - MCPError: If the request results in a TimeoutError | JSONRPCError - """ - with client_span( - "tasks/result", - "tasks/result", - task_id, - session_id=self.transport.get_session_id(), - ): - request_meta = cast( - "mcp_types.RequestParamsMeta | None", inject_trace_context() - ) - request = GetTaskPayloadRequest( - params=GetTaskPayloadRequestParams( - task_id=task_id, - _meta=request_meta, # type: ignore[unknown-argument] - ) - ) - # Return raw result - Task classes handle type-specific parsing - result = await self._await_with_session_monitoring( - self.session.send_request( - request=request, # type: ignore[arg-type] - result_type=_RawTaskPayloadResult, - ) - ) - # Return as dict for compatibility with Task class parsing. The payload - # fields (content, structuredContent, messages, contents, ...) survive - # via the permissive result type's extra="allow". - return result.model_dump(exclude_none=True, by_alias=True) - - async def list_tasks( - self: Client, - cursor: str | None = None, - limit: int = 50, - ) -> dict[str, Any]: - """List background tasks. - - Sends a 'tasks/list' MCP protocol request to the server. If the server - returns an empty list (indicating client-side tracking), falls back to - querying status for locally tracked task IDs. - - Args: - cursor: Optional pagination cursor - limit: Maximum number of tasks to return (default 50) - - Returns: - dict: Response with structure: - - tasks: List of task status dicts with taskId, status, etc. - - nextCursor: Optional cursor for next page - - Raises: - RuntimeError: If client not connected - MCPError: If the request results in a TimeoutError | JSONRPCError - """ - with client_span( - "tasks/list", - "tasks/list", - "", - session_id=self.transport.get_session_id(), - ): - request_meta = cast( - "mcp_types.RequestParamsMeta | None", inject_trace_context() - ) - - # Send protocol request - params = PaginatedRequestParams.model_validate( - {"cursor": cursor, "limit": limit, "_meta": request_meta} - ) - request = ListTasksRequest(params=params) - server_response = await self._await_with_session_monitoring( - self.session.send_request( - request=request, # type: ignore[invalid-argument-type] - result_type=mcp_types.ListTasksResult, - ) - ) - - # If server returned tasks, use those - if server_response.tasks: - return server_response.model_dump(by_alias=True) - - # Server returned empty - fall back to client-side tracking - tasks = [] - for task_id in list(self._submitted_task_ids)[:limit]: # ty: ignore[unresolved-attribute] - try: - status = await self.get_task_status(task_id) # ty: ignore[unresolved-attribute] - tasks.append(status.model_dump(by_alias=True)) - except MCPError: - # Task may have expired or been deleted, skip it - continue - - return {"tasks": tasks, "nextCursor": None} - - async def cancel_task(self: Client, task_id: str) -> mcp_types.CancelTaskResult: - """Cancel a task, transitioning it to cancelled state. - - Sends a 'tasks/cancel' MCP protocol request. Task will halt execution - and transition to cancelled state. - - Args: - task_id: The task ID to cancel - - Returns: - CancelTaskResult: The task status showing cancelled state - - Raises: - RuntimeError: If task doesn't exist - MCPError: If the request results in a TimeoutError | JSONRPCError - """ - with client_span( - "tasks/cancel", - "tasks/cancel", - task_id, - session_id=self.transport.get_session_id(), - ): - request_meta = cast( - "mcp_types.RequestParamsMeta | None", inject_trace_context() - ) - request = CancelTaskRequest( - params=CancelTaskRequestParams( - task_id=task_id, - _meta=request_meta, # type: ignore[unknown-argument] - ) - ) - return await self._await_with_session_monitoring( - self.session.send_request( - request=request, # type: ignore[invalid-argument-type] - result_type=mcp_types.CancelTaskResult, - ) - ) diff --git a/fastmcp_tasks/fastmcp_tasks/client.py b/fastmcp_tasks/fastmcp_tasks/client.py index 5e4c0502d..3412ce268 100644 --- a/fastmcp_tasks/fastmcp_tasks/client.py +++ b/fastmcp_tasks/fastmcp_tasks/client.py @@ -1,626 +1,443 @@ -"""SEP-1686 client Task classes.""" +"""SEP-2663 client task support: the tasks extension, resolver, and handle. + +FastMCP drives a server's background tasks transparently. When a `task=True` +tool runs a call as a task, the server answers `tools/call` with a claimed +`CreateTaskResult` (SEP-2133) instead of the tool's result. This module supplies +the client half: + +- `TasksClientExtension` advertises the tasks capability (so the server *may* + task the call) and declares a `ResultClaim` for `resultType: "task"`. It is + registered on every FastMCP `Client` automatically, so the caller opts in to + nothing. +- The claim's resolver polls `tasks/get` to completion under the hood and returns + the tool's real result as a `CallToolResult` β€” the caller of `call_tool` never + learns the call was tasked. A task that pauses for input is answered through the + client's `elicitation_handler` via `tasks/update`, then polling resumes. +- `ToolTask` is the explicit handle for callers who want to return immediately and + drive the task themselves (`status`/`wait`/`result`/`cancel`), built via + `call_tool_task`. + +Tasks are modern-protocol only: on a legacy connection the SDK strips the +capability ad, the server never tasks, and this extension is inert. +""" from __future__ import annotations -import abc import asyncio -import inspect -import time -import weakref -from collections.abc import Awaitable, Callable -from datetime import datetime, timezone -from typing import TYPE_CHECKING, Generic, TypeVar +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any, cast import mcp_types -from mcp_types import GetTaskResult, TaskStatusNotification +from mcp.client.extension import ClaimContext, ClientExtension, ResultClaim +from mcp.client.session import ClientRequestContext, ClientSession, ElicitationFnT +from mcp_types import CallToolResult +from mcp_types.version import MODERN_PROTOCOL_VERSIONS -import fastmcp -from fastmcp.client.messages import Message, MessageHandler from fastmcp.exceptions import ToolError from fastmcp.utilities.logging import get_logger +from fastmcp.utilities.tasks import TASKS_EXTENSION_ID +from fastmcp.utilities.timeout import normalize_timeout_to_seconds +from fastmcp_tasks.client_models import ( + CancelTaskRequest, + CancelTaskRequestParams, + ClientCreateTaskResult, + ClientGetTaskResult, + GetTaskRequest, + GetTaskRequestParams, + UpdateTaskRequest, + UpdateTaskRequestParams, +) +from fastmcp_tasks.settings import client_settings + +if TYPE_CHECKING: + from fastmcp.client.client import CallToolResult as FastMCPCallToolResult + from fastmcp.client.client import Client logger = get_logger(__name__) -# Floor for the fallback poll interval in Task.wait() (seconds). When the server -# does not advertise a pollInterval, each wait() call starts its backoff ramp -# here so fast tasks resolve quickly even if a status notification is missed. -# When the server does advertise one, this is only a safety floor that keeps a -# server sending `pollInterval: 0` from spinning the client in a tight loop. +#: Floor for the fallback poll interval (seconds). When the server does not +#: advertise a `pollIntervalMs`, each drive starts its backoff ramp here so quick +#: tasks resolve fast; when it does advertise one, this floors it so a server +#: sending `0` cannot spin the client in a tight loop. MIN_POLL_INTERVAL = 0.02 -if TYPE_CHECKING: - from fastmcp.client.client import CallToolResult, Client +_TERMINAL_STATES = frozenset({"completed", "failed", "cancelled"}) -class TaskNotificationHandler(MessageHandler): - """MessageHandler that routes task status notifications to Task objects.""" - - def __init__(self, client: Client): - super().__init__() - self._client_ref: weakref.ref[Client] = weakref.ref(client) - - async def dispatch(self, message: Message) -> None: - """Dispatch messages, including task status notifications.""" - # SDK v2 delivers notifications unwrapped (no `.root` wrapper). - if isinstance(message, TaskStatusNotification): - client = self._client_ref() - if client: - client._handle_task_status_notification(message) # ty: ignore[unresolved-attribute] - - await super().dispatch(message) +# --------------------------------------------------------------------------- +# Wire senders (tasks/get, tasks/update, tasks/cancel) over a ClientSession +# --------------------------------------------------------------------------- -TaskResultT = TypeVar("TaskResultT") +async def _send_get(session: ClientSession, task_id: str) -> ClientGetTaskResult: + """Send `tasks/get` and parse the detailed task response.""" + request = GetTaskRequest(params=GetTaskRequestParams(task_id=task_id)) + return await session.send_request(request, ClientGetTaskResult) -class Task(abc.ABC, Generic[TaskResultT]): +async def _send_update( + session: ClientSession, task_id: str, input_responses: dict[str, Any] +) -> None: + """Send `tasks/update` delivering the caller's answers to a parked task.""" + request = UpdateTaskRequest( + params=UpdateTaskRequestParams(task_id=task_id, input_responses=input_responses) + ) + await session.send_request(request, mcp_types.Result) + + +async def _send_cancel(session: ClientSession, task_id: str) -> None: + """Send `tasks/cancel` to cooperatively cancel a task.""" + request = CancelTaskRequest(params=CancelTaskRequestParams(task_id=task_id)) + await session.send_request(request, mcp_types.Result) + + +# --------------------------------------------------------------------------- +# Poll cadence +# --------------------------------------------------------------------------- + + +def _poll_ceiling(poll_interval_ms: float | None) -> float: + """The upper bound for the poll backoff, in seconds. + + A server-advertised `pollIntervalMs` is a deliberate statement about how much + load the server wants to take, so it caps the backoff. A zero, negative, or + absent value falls back to the `poll_interval` client setting; the ceiling is + never below `MIN_POLL_INTERVAL` so a hostile `0` cannot spin the client. """ - Abstract base class for MCP background tasks (SEP-1686). + if poll_interval_ms is not None and poll_interval_ms > 0: + return max(poll_interval_ms / 1000, MIN_POLL_INTERVAL) + return client_settings.poll_interval - Provides a uniform API whether the server accepts background execution - or executes synchronously (graceful degradation per SEP-1686). - Subclasses: - - ToolTask: For tool calls (result type: CallToolResult) - - PromptTask: For prompts (future, result type: GetPromptResult) - - ResourceTask: For resources (future, result type: ReadResourceResult) +def _next_poll_delay( + poll_interval_ms: float | None, backoff: float +) -> tuple[float, float]: + """Delay before the next poll, plus the backoff for the round after. + + With no status notifications on the modern protocol, polling is the only + signal, so a fixed cadence at the server's advertised interval would make a + quick task take that full interval to observe as done. Instead the backoff + ramps from `MIN_POLL_INTERVAL`, doubling each round up to the ceiling + (`_poll_ceiling`): a quick task resolves in ~20ms while a long one settles to + the server's advertised cadence, hammering neither. + """ + ceiling = _poll_ceiling(poll_interval_ms) + return min(backoff, ceiling), min(backoff * 2, ceiling) + + +# --------------------------------------------------------------------------- +# In-task input: answer a parked task's requests via the elicitation handler +# --------------------------------------------------------------------------- + + +async def _answer_input_requests( + session: ClientSession, + task_id: str, + input_requests: dict[str, Any], + elicitation_callback: ElicitationFnT | None, +) -> None: + """Answer a task's outstanding input requests, then deliver via `tasks/update`. + + Each request is surfaced by a server-minted key and carries a serialized + `ElicitRequest`. The client's elicitation handler produces each answer; the + keyed answers are sent back with `tasks/update`, which re-enters the task. + Sampling and roots requests are not supported on the modern protocol. + """ + if elicitation_callback is None: + raise ToolError( + f"Task {task_id} requires input but the client has no elicitation " + "handler; pass elicitation_handler= to Client() to drive tasks that " + "ask for input." + ) + + responses: dict[str, Any] = {} + for surfaced_key, payload in input_requests.items(): + method = payload.get("method") if isinstance(payload, dict) else None + if method != "elicitation/create": + raise ToolError( + f"Task {task_id} requested in-task input via {method!r}, which the " + "client cannot answer; only elicitation is supported on the modern " + "protocol (sampling and roots are deprecated)." + ) + request = mcp_types.ElicitRequest.model_validate(payload) + context = ClientRequestContext( + session=session, request_id=f"task-{task_id}-{surfaced_key}" + ) + answer = await elicitation_callback(context, request.params) + 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) + + +# --------------------------------------------------------------------------- +# The shared poll loop +# --------------------------------------------------------------------------- + + +async def _drive_to_terminal( + session: ClientSession, + task_id: str, + elicitation_callback: ElicitationFnT | None, +) -> ClientGetTaskResult: + """Poll `tasks/get` until the task reaches a terminal state. + + `working` sleeps and polls again; `input_required` answers the outstanding + requests through the elicitation handler and re-enters; a terminal state + (completed / failed / cancelled) is returned. Shared by the transparent + resolver and `ToolTask.result()`. + """ + backoff = MIN_POLL_INTERVAL + while True: + current = await _send_get(session, task_id) + if current.status in _TERMINAL_STATES: + return current + if current.status == "input_required": + await _answer_input_requests( + session, task_id, current.input_requests or {}, elicitation_callback + ) + backoff = MIN_POLL_INTERVAL + continue + # working + delay, backoff = _next_poll_delay(current.poll_interval_ms, backoff) + await asyncio.sleep(delay) + + +def _inlined_call_tool_result(result: dict[str, Any] | None) -> CallToolResult: + """Parse a completed task's inlined result dict into a `CallToolResult`.""" + return CallToolResult.model_validate(result or {}) + + +def _terminal_error_message(final: ClientGetTaskResult) -> str: + """The best available error message for a failed task.""" + if isinstance(final.error, dict): + message = final.error.get("message") + if isinstance(message, str) and message: + return message + if final.status_message: + return final.status_message + return f"Task {final.task_id} failed" + + +# --------------------------------------------------------------------------- +# The tasks client extension and its claim resolver +# --------------------------------------------------------------------------- + + +class TasksClientExtension(ClientExtension): + """The client half of the `io.modelcontextprotocol/tasks` extension (SEP-2663). + + Advertising this extension tells the server the client can drive tasks, so a + `task=True` tool may run as a task; the declared `ResultClaim` then resolves + the `CreateTaskResult` the server returns by polling `tasks/get` to the real + result. Registered automatically on every FastMCP `Client`. + """ + + identifier = TASKS_EXTENSION_ID + + def __init__(self, elicitation_callback: ElicitationFnT | None = None) -> None: + self._elicitation_callback = elicitation_callback + + def settings(self) -> dict[str, Any]: + """The tasks extension advertises no per-extension settings.""" + return {} + + def claims(self) -> Sequence[ResultClaim[Any]]: + return ( + ResultClaim( + result_type="task", + model=ClientCreateTaskResult, + resolve=self._resolve_task, + protocol_versions=frozenset(MODERN_PROTOCOL_VERSIONS), + ), + ) + + async def _resolve_task( + self, create_result: ClientCreateTaskResult, ctx: ClaimContext + ) -> CallToolResult: + """Finish a tasked `tools/call` by polling `tasks/get` to completion. + + Returns the tool's real result on completion; a failed or cancelled task + becomes an error `CallToolResult` so the ordinary `call_tool` error path + (raise `ToolError`) applies uniformly, and the completed inlined result is + schema-valid so the SDK's output-schema revalidation passes. + """ + final = await _drive_to_terminal( + ctx.session, create_result.task_id, self._elicitation_callback + ) + if final.status == "completed": + return _inlined_call_tool_result(final.result) + if final.status == "failed": + message = _terminal_error_message(final) + else: + message = f"Task {final.task_id} was cancelled" + return CallToolResult( + content=[mcp_types.TextContent(type="text", text=message)], + is_error=True, + ) + + +def _build_tasks_client_extension( + elicitation_callback: ElicitationFnT | None, +) -> ClientExtension: + """Factory registered with core so every `Client` folds in task support.""" + return TasksClientExtension(elicitation_callback) + + +# --------------------------------------------------------------------------- +# The explicit task handle (return-quickly surface) +# --------------------------------------------------------------------------- + + +class ToolTask: + """A handle to a tool call the server is running as a background task. + + Returned by `call_tool_task`. Lets a caller return immediately and then drive + the task: check `status`, `wait` for a state, get the finished `result` + (answering any input prompts through the client's elicitation handler), or + `cancel`. Awaiting the handle is shorthand for `result()`. """ def __init__( self, client: Client, - task_id: str, - immediate_result: TaskResultT | None = None, - ): - """ - Create a Task wrapper. - - Args: - client: The FastMCP client - task_id: The task identifier - immediate_result: If server executed synchronously, the immediate result - """ + tool_name: str, + create_result: ClientCreateTaskResult, + *, + raise_on_error: bool = True, + ) -> None: self._client = client - self._task_id = task_id - self._immediate_result = immediate_result - self._is_immediate = immediate_result is not None - - # Notification-based optimization (SEP-1686 notifications/tasks/status) - self._status_cache: GetTaskResult | None = None - self._status_event: asyncio.Event | None = None # Lazy init - self._status_callbacks: list[ - Callable[[GetTaskResult], None | Awaitable[None]] - ] = [] - self._cached_result: TaskResultT | None = None - - def _check_client_connected(self) -> None: - """Validate that client context is still active. - - Raises: - RuntimeError: If accessed outside client context (unless immediate) - """ - if self._is_immediate: - return # Already resolved, no client needed - - try: - _ = self._client.session - except RuntimeError as e: - raise RuntimeError( - "Cannot access task results outside client context. " - "Task futures must be used within 'async with client:' block." - ) from e + self._tool_name = tool_name + self._create_result = create_result + self._raise_on_error = raise_on_error + self._cached_result: FastMCPCallToolResult | None = None @property def task_id(self) -> str: - """Get the task ID.""" - return self._task_id + """The server-generated task id.""" + return self._create_result.task_id @property - def returned_immediately(self) -> bool: - """Check if server executed the task immediately. + def create_result(self) -> ClientCreateTaskResult: + """The raw `CreateTaskResult` the server returned for the tasked call.""" + return self._create_result - Returns: - True if server executed synchronously (graceful degradation or no task support) - False if server accepted background execution - """ - return self._is_immediate + @property + def _session(self) -> ClientSession: + return self._client.session - def _handle_status_notification(self, status: GetTaskResult) -> None: - """Process incoming notifications/tasks/status (internal). + @property + def _elicitation_callback(self) -> ElicitationFnT | None: + return self._client._elicitation_callback - Called by Client when a notification is received for this task. - Updates cache, triggers events, and invokes user callbacks. - - Args: - status: Task status from notification - """ - # Update cache for next status() call - self._status_cache = status - - # Wake up any wait() calls - if self._status_event is not None: - self._status_event.set() - - # Invoke user callbacks - for callback in self._status_callbacks: - try: - result = callback(status) - if inspect.isawaitable(result): - # Fire and forget async callbacks - asyncio.create_task(result) # type: ignore[arg-type] # noqa: RUF006 # ty:ignore[invalid-argument-type] - except Exception as e: - logger.warning(f"Task callback error: {e}", exc_info=True) - - def on_status_change( - self, - callback: Callable[[GetTaskResult], None | Awaitable[None]], - ) -> None: - """Register callback for status change notifications. - - The callback will be invoked when a notifications/tasks/status is received - for this task (optional server feature per SEP-1686 lines 436-444). - - Supports both sync and async callbacks (auto-detected). - - Args: - callback: Function to call with GetTaskResult when status changes. - Can return None (sync) or Awaitable[None] (async). - - Example: - >>> task = await client.call_tool("slow_operation", {}, task=True) - >>> - >>> def on_update(status: GetTaskResult): - ... print(f"Task {status.task_id} is now {status.status}") - >>> - >>> task.on_status_change(on_update) - >>> result = await task # Callback fires when status changes - """ - self._status_callbacks.append(callback) - - async def status(self) -> GetTaskResult: - """Get current task status. - - If server executed immediately, returns synthetic completed status. - Otherwise queries the server for current status. - """ - self._check_client_connected() - - if self._is_immediate: - # Return synthetic completed status. SDK v2 types the task - # timestamps as ISO 8601 strings. - now = datetime.now(timezone.utc).isoformat() - return GetTaskResult( - task_id=self._task_id, - status="completed", - created_at=now, - last_updated_at=now, - ttl=None, - poll_interval=1000, - ) - - # Return cached status if available (from notification) - if self._status_cache is not None: - cached = self._status_cache - # Don't clear cache - keep it for next call - return cached - - # Query server and cache the result - self._status_cache = await self._client.get_task_status(self._task_id) # ty: ignore[unresolved-attribute] - return self._status_cache - - @abc.abstractmethod - async def result(self) -> TaskResultT: - """Wait for and return the task result. - - Must be implemented by subclasses to return the appropriate result type. - """ - ... + async def status(self) -> ClientGetTaskResult: + """Fetch the task's current status via `tasks/get`.""" + return await _send_get(self._session, self.task_id) async def wait( self, *, state: str | None = None, timeout: float = 300.0 - ) -> GetTaskResult: - """Wait for task to reach a specific state or complete. + ) -> ClientGetTaskResult: + """Poll until the task reaches `state` (or any terminal state if `None`). - Uses event-based waiting when notifications are available (fast), - with fallback to polling (reliable). Optimally wakes up immediately - on status changes when server sends notifications/tasks/status. - - The fallback poll cadence has two modes. If the server advertises a - `pollInterval`, that interval is honored exactly (subject only to a - 20ms safety floor), because it is a deliberate statement about how - much load the server wants to take. If it does not, the poll starts at - 20ms and doubles up to the `client_task_poll_interval` setting. - - Args: - state: Desired state ('working', 'input_required', 'completed', 'failed', 'cancelled'). - If None, waits until the task exits the 'working' state (completed, failed, cancelled, input_required, etc.) - timeout: Maximum time to wait in seconds - - Returns: - GetTaskResult: Final task status - - Raises: - TimeoutError: If desired state not reached within timeout + Does not answer input prompts: a caller that wants automatic answering + should use `result()`. `wait(state="input_required")` lets a caller + observe the parked state and answer it manually. """ - self._check_client_connected() - - if self._is_immediate: - # Already done - return await self.status() - - # Initialize event for notification wake-ups - if self._status_event is None: - self._status_event = asyncio.Event() - - start = time.time() - in_progress_states = {"working"} - # Backoff state for the unadvertised-interval mode; resets per wait() - # call. Notifications still short-circuit the wait via the status event, - # so this only governs the fallback poll. + loop = asyncio.get_event_loop() + deadline = loop.time() + timeout backoff = MIN_POLL_INTERVAL - while True: - # Check cached status first (updated by notifications) - if self._status_cache: - current = self._status_cache.status - if state is None: - if current not in in_progress_states: - return self._status_cache - elif current == state: - return self._status_cache - - # Check timeout - elapsed = time.time() - start - if elapsed >= timeout: + current = await self.status() + if state is not None: + if current.status == state: + return current + elif current.status in _TERMINAL_STATES: + return current + if loop.time() >= deadline: raise TimeoutError( - f"Task {self._task_id} did not reach {state or 'terminal state'} within {timeout}s" + f"Task {self.task_id} did not reach " + f"{state or 'a terminal state'} within {timeout}s" ) + delay, backoff = _next_poll_delay(current.poll_interval_ms, backoff) + await asyncio.sleep(delay) - remaining = timeout - elapsed - interval, backoff = self._next_poll_delay(backoff) + async def result(self) -> FastMCPCallToolResult: + """Drive the task to completion and return its parsed result. - # Wait for notification event OR poll timeout - try: - await asyncio.wait_for( - self._status_event.wait(), timeout=min(interval, remaining) - ) - self._status_event.clear() - except asyncio.TimeoutError: - # Fallback: poll server (notification didn't arrive in time) - self._status_cache = await self._client.get_task_status(self._task_id) # ty: ignore[unresolved-attribute] - - def _next_poll_delay(self, backoff: float) -> tuple[float, float]: - """Delay before the next fallback poll, plus the backoff for the round after. - - Advertised interval -> honor it; no advertised interval -> ramp. - - A server that advertises `pollInterval` (milliseconds) is making a - deliberate statement about how much load it wants to take, so that - interval is used verbatim as the delay with no backoff ramp. The only - adjustment is `MIN_POLL_INTERVAL` as a safety floor, so a server sending - a zero or negative interval cannot spin this client in a tight request - loop. - - When the server advertises nothing, there is no guidance to honor, so - the poll starts at `MIN_POLL_INTERVAL` and doubles each round up to the - `client_task_poll_interval` setting. + Answers any input prompts through the client's elicitation handler. + Raises `ToolError` on a failed or cancelled task when `raise_on_error` + (the default); otherwise returns an error result. The result is cached, so + repeated calls return the same object. """ - cache = self._status_cache - if cache is not None and cache.poll_interval is not None: - return max(cache.poll_interval / 1000, MIN_POLL_INTERVAL), backoff + if self._cached_result is not None: + return self._cached_result - ceiling = fastmcp.settings.client_task_poll_interval - return min(backoff, ceiling), min(backoff * 2, ceiling) + final = await _drive_to_terminal( + self._session, self.task_id, self._elicitation_callback + ) + if final.status == "completed": + mcp_result = _inlined_call_tool_result(final.result) + else: + if final.status == "failed": + message = _terminal_error_message(final) + else: + message = f"Task {self.task_id} was cancelled" + if self._raise_on_error: + raise ToolError(message) + mcp_result = CallToolResult( + content=[mcp_types.TextContent(type="text", text=message)], + is_error=True, + ) - async def _wait_terminal(self, timeout: float = 300.0) -> GetTaskResult: - """Wait until task reaches a terminal state (completed, failed, cancelled). - - Unlike wait(), this will not return on input_required β€” it continues - waiting until the task fully resolves. Used internally by result(). - """ - terminal_states = {"completed", "failed", "cancelled"} - status = await self.wait(timeout=timeout) - while status.status not in terminal_states: - # Task is in a non-terminal state (e.g. input_required) β€” reset - # cache so the next wait() call blocks instead of returning immediately. - self._status_cache = None - status = await self.wait(timeout=timeout) - return status + parsed = await self._client._parse_call_tool_result( + self._tool_name, mcp_result, raise_on_error=self._raise_on_error + ) + self._cached_result = parsed + return parsed async def cancel(self) -> None: - """Cancel this task, transitioning it to cancelled state. - - Sends a tasks/cancel protocol request. The server will attempt to halt - execution and move the task to cancelled state. - - Note: If server executed immediately (graceful degradation), this is a no-op - as there's no server-side task to cancel. - """ - if self._is_immediate: - # No server-side task to cancel - return - self._check_client_connected() - await self._client.cancel_task(self._task_id) # ty: ignore[unresolved-attribute] - # Invalidate cache to force fresh status fetch - self._status_cache = None + """Request cooperative cancellation of the task via `tasks/cancel`.""" + await _send_cancel(self._session, self.task_id) def __await__(self): - """Allow 'await task' to get result.""" return self.result().__await__() -class ToolTask(Task["CallToolResult"]): +async def call_tool_task( + client: Client, + name: str, + arguments: dict[str, Any] | None = None, + *, + timeout: float | int | None = None, + raise_on_error: bool = True, + meta: dict[str, Any] | None = None, +) -> ToolTask: + """Call a tool as a background task and return a `ToolTask` handle immediately. + + Unlike `client.call_tool` (which polls to completion transparently), this + returns as soon as the server accepts the task, so the caller can do other + 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`. """ - Represents a tool call that may execute in background or immediately. - - Provides a uniform API whether the server accepts background execution - or executes synchronously (graceful degradation per SEP-1686). - - Usage: - task = await client.call_tool_as_task("analyze", args) - - # Check status - status = await task.status() - - # Wait for completion - await task.wait() - - # Get result (waits if needed) - result = await task.result() # Returns CallToolResult - - # Or just await the task directly - result = await task - """ - - def __init__( - self, - client: Client, - task_id: str, - tool_name: str, - immediate_result: CallToolResult | None = None, - raise_on_error: bool = True, - ): - """ - Create a ToolTask wrapper. - - Args: - client: The FastMCP client - task_id: The task identifier - tool_name: Name of the tool being executed - immediate_result: If server executed synchronously, the immediate result - raise_on_error: Whether task.result() should raise ToolError on errors - """ - super().__init__(client, task_id, immediate_result) - self._tool_name = tool_name - self._raise_on_error = raise_on_error - - async def result(self) -> CallToolResult: - """Wait for and return the tool result. - - If server executed immediately, returns the immediate result. - Otherwise waits for background task to complete and retrieves result. - - Returns: - CallToolResult: The parsed tool result (same as call_tool returns) - """ - # Check cache first - if self._cached_result is not None: - return self._cached_result - - if self._is_immediate: - assert self._immediate_result is not None # Type narrowing - result = self._immediate_result - if result.is_error and self._raise_on_error: - if result.content and isinstance( - result.content[0], mcp_types.TextContent - ): - msg = result.content[0].text - else: - msg = f"Tool '{self._tool_name}' returned an error" - raise ToolError(msg) - else: - # Check client connected - self._check_client_connected() - - # Wait for completion using event-based wait (respects notifications) - await self._wait_terminal() - - # Get the raw result (dict or CallToolResult) - raw_result = await self._client.get_task_result(self._task_id) # ty: ignore[unresolved-attribute] - - # Convert to CallToolResult if needed and parse - if isinstance(raw_result, dict): - # Raw dict from get_task_result - parse as CallToolResult - mcp_result = mcp_types.CallToolResult.model_validate(raw_result) - result = await self._client._parse_call_tool_result( - self._tool_name, - mcp_result, - raise_on_error=self._raise_on_error, - ) - elif isinstance(raw_result, mcp_types.CallToolResult): - # Already a CallToolResult from MCP protocol - parse it - result = await self._client._parse_call_tool_result( - self._tool_name, - raw_result, - raise_on_error=self._raise_on_error, - ) - else: - # Legacy ToolResult format - convert to MCP type - if hasattr(raw_result, "content") and hasattr( - raw_result, "structured_content" - ): - mcp_result = mcp_types.CallToolResult( - content=raw_result.content, - structured_content=raw_result.structured_content, - _meta=raw_result.meta, # type: ignore[call-arg] # _meta is Pydantic alias for meta field - ) - result = await self._client._parse_call_tool_result( - self._tool_name, - mcp_result, - raise_on_error=self._raise_on_error, - ) - else: - # Unknown type - just return it - result = raw_result - - # Cache before returning - self._cached_result = result - return result - - -class PromptTask(Task[mcp_types.GetPromptResult]): - """ - Represents a prompt call that may execute in background or immediately. - - Provides a uniform API whether the server accepts background execution - or executes synchronously (graceful degradation per SEP-1686). - - Usage: - task = await client.get_prompt_as_task("analyze", args) - result = await task # Returns GetPromptResult - """ - - def __init__( - self, - client: Client, - task_id: str, - prompt_name: str, - immediate_result: mcp_types.GetPromptResult | None = None, - ): - """ - Create a PromptTask wrapper. - - Args: - client: The FastMCP client - task_id: The task identifier - prompt_name: Name of the prompt being executed - immediate_result: If server executed synchronously, the immediate result - """ - super().__init__(client, task_id, immediate_result) - self._prompt_name = prompt_name - - async def result(self) -> mcp_types.GetPromptResult: - """Wait for and return the prompt result. - - If server executed immediately, returns the immediate result. - Otherwise waits for background task to complete and retrieves result. - - Returns: - GetPromptResult: The prompt result with messages and description - """ - # Check cache first - if self._cached_result is not None: - return self._cached_result - - if self._is_immediate: - assert self._immediate_result is not None - result = self._immediate_result - else: - # Check client connected - self._check_client_connected() - - # Wait for completion using event-based wait (respects notifications) - await self._wait_terminal() - - # Get the raw MCP result - mcp_result = await self._client.get_task_result(self._task_id) # ty: ignore[unresolved-attribute] - - # Parse as GetPromptResult - result = mcp_types.GetPromptResult.model_validate(mcp_result) - - # Cache before returning - self._cached_result = result - return result - - -class ResourceTask( - Task[list[mcp_types.TextResourceContents | mcp_types.BlobResourceContents]] -): - """ - Represents a resource read that may execute in background or immediately. - - Provides a uniform API whether the server accepts background execution - or executes synchronously (graceful degradation per SEP-1686). - - Usage: - task = await client.read_resource_as_task("file://data.txt") - contents = await task # Returns list[ReadResourceContents] - """ - - def __init__( - self, - client: Client, - task_id: str, - uri: str, - immediate_result: list[ - mcp_types.TextResourceContents | mcp_types.BlobResourceContents - ] - | None = None, - ): - """ - Create a ResourceTask wrapper. - - Args: - client: The FastMCP client - task_id: The task identifier - uri: URI of the resource being read - immediate_result: If server executed synchronously, the immediate result - """ - super().__init__(client, task_id, immediate_result) - self._uri = uri - - async def result( - self, - ) -> list[mcp_types.TextResourceContents | mcp_types.BlobResourceContents]: - """Wait for and return the resource contents. - - If server executed immediately, returns the immediate result. - Otherwise waits for background task to complete and retrieves result. - - Returns: - list[ReadResourceContents]: The resource contents - """ - # Check cache first - if self._cached_result is not None: - return self._cached_result - - if self._is_immediate: - assert self._immediate_result is not None - result = self._immediate_result - else: - # Check client connected - self._check_client_connected() - - # Wait for completion using event-based wait (respects notifications) - await self._wait_terminal() - - # Get the raw MCP result - mcp_result = await self._client.get_task_result(self._task_id) # ty: ignore[unresolved-attribute] - - # Parse as ReadResourceResult or extract contents - if isinstance(mcp_result, mcp_types.ReadResourceResult): - # Already parsed by TasksResponse - extract contents - result = list(mcp_result.contents) - elif isinstance(mcp_result, dict) and "contents" in mcp_result: - # Dict format - parse each content item - parsed_contents = [] - for item in mcp_result["contents"]: - if isinstance(item, dict): - if "blob" in item: - parsed_contents.append( - mcp_types.BlobResourceContents.model_validate(item) - ) - else: - parsed_contents.append( - mcp_types.TextResourceContents.model_validate(item) - ) - else: - parsed_contents.append(item) - result = parsed_contents - else: - # Fallback - might be the list directly - result = mcp_result if isinstance(mcp_result, list) else [mcp_result] - - # Cache before returning - self._cached_result = result - return result + read_timeout_seconds = normalize_timeout_to_seconds(timeout) + request_meta = cast("mcp_types.RequestParamsMeta | None", meta) + raw = await client._await_with_session_monitoring( + client.session.call_tool( + name=name, + arguments=arguments or {}, + read_timeout_seconds=read_timeout_seconds, + meta=request_meta, + allow_claimed=True, + ) + ) + if isinstance(raw, ClientCreateTaskResult): + return ToolTask(client, name, raw, raise_on_error=raise_on_error) + raise ToolError( + f"Tool {name!r} did not run as a task: the server returned a " + f"{type(raw).__name__} instead of a task. Ensure the tool is declared " + "task=True and the connection is modern (mode='auto')." + ) diff --git a/fastmcp_tasks/fastmcp_tasks/client_models.py b/fastmcp_tasks/fastmcp_tasks/client_models.py new file mode 100644 index 000000000..25b9cc5e2 --- /dev/null +++ b/fastmcp_tasks/fastmcp_tasks/client_models.py @@ -0,0 +1,130 @@ +"""Client-side wire models for the SEP-2663 tasks extension. + +These mirror the server models in ``models.py`` but flip the alias direction: +the server *produces* the wire (``serialization_alias`` -> camelCase dump), while +the client *consumes* it. The SDK validates both a claimed ``tools/call`` result +and a ``tasks/get`` response with ``model_validate(raw, by_name=False)``, so these +models declare **validation** aliases (``Field(alias="taskId")``) to read the +camelCase wire keys. + +``ClientCreateTaskResult`` is the claim shape the tasks ``ResultClaim`` resolves. +It must subclass ``mcp_types.Result`` (not ``CallToolResult`` / +``InputRequiredResult``) and pin ``result_type`` to ``Literal["task"]`` β€” the +SDK's ``ResultClaim.__post_init__`` enforces exactly this. ``ClientGetTaskResult`` +is the typed ``tasks/get`` response: the flat task fields plus exactly one of +``result`` (completed), ``error`` (failed), or ``inputRequests`` (input_required). +""" + +from __future__ import annotations + +from typing import Any, Literal + +import mcp_types +from mcp_types import RequestParams, Result +from pydantic import ConfigDict, Field + +__all__ = [ + "TaskStatus", + "ClientCreateTaskResult", + "ClientGetTaskResult", + "GetTaskRequest", + "GetTaskRequestParams", + "UpdateTaskRequest", + "UpdateTaskRequestParams", + "CancelTaskRequest", + "CancelTaskRequestParams", +] + +TaskStatus = Literal["working", "input_required", "completed", "failed", "cancelled"] + + +class _ClientTaskFields(Result): + """The flat task fields shared by every SEP-2663 task result, read from the wire. + + Validation aliases (camelCase) because the SDK validates the server's + ``model_dump(by_alias=True)`` output with ``by_name=False``. + """ + + model_config = ConfigDict(populate_by_name=True) + + task_id: str = Field(alias="taskId") + status: TaskStatus + created_at: str = Field(alias="createdAt") + last_updated_at: str = Field(alias="lastUpdatedAt") + ttl_ms: float | None = Field(default=None, alias="ttlMs") + status_message: str | None = Field(default=None, alias="statusMessage") + poll_interval_ms: float | None = Field(default=None, alias="pollIntervalMs") + + +class ClientCreateTaskResult(_ClientTaskFields): + """The claimed ``tools/call`` result the server returns when it runs a call as a task. + + Pinned to ``resultType: "task"`` so the tasks ``ResultClaim`` can key on it. + The resolver polls ``tasks/get`` from here to the finished result. + """ + + result_type: Literal["task"] = Field(alias="resultType") + + +class ClientGetTaskResult(_ClientTaskFields): + """The typed ``tasks/get`` response: task fields plus the inlined outcome. + + Exactly one of ``result`` / ``error`` / ``input_requests`` is set, matching + the task's status. ``result_type`` is ``"complete"`` because ``tasks/get`` + itself always completes normally, whatever the task's own status. + """ + + result_type: Literal["complete"] = Field(alias="resultType") + result: dict[str, Any] | None = None + error: dict[str, Any] | None = None + input_requests: dict[str, Any] | None = Field(default=None, alias="inputRequests") + + +class GetTaskRequestParams(RequestParams): + """Params for ``tasks/get`` / ``tasks/cancel``: the target task id. + + These are outbound (client -> server), so they carry *serialization* aliases: + the client constructs them by field name and `send_request` dumps them to the + camelCase wire shape with `by_alias=True`. + """ + + model_config = ConfigDict(populate_by_name=True) + + task_id: str = Field(serialization_alias="taskId") + + +CancelTaskRequestParams = GetTaskRequestParams + + +class UpdateTaskRequestParams(RequestParams): + """Params for ``tasks/update``: task id plus the caller's input responses.""" + + model_config = ConfigDict(populate_by_name=True) + + task_id: str = Field(serialization_alias="taskId") + input_responses: dict[str, Any] = Field(serialization_alias="inputResponses") + + +class GetTaskRequest(mcp_types.Request[GetTaskRequestParams, Literal["tasks/get"]]): + """``tasks/get`` request envelope for ``ClientSession.send_request``.""" + + method: Literal["tasks/get"] = "tasks/get" + params: GetTaskRequestParams + + +class UpdateTaskRequest( + mcp_types.Request[UpdateTaskRequestParams, Literal["tasks/update"]] +): + """``tasks/update`` request envelope for ``ClientSession.send_request``.""" + + method: Literal["tasks/update"] = "tasks/update" + params: UpdateTaskRequestParams + + +class CancelTaskRequest( + mcp_types.Request[CancelTaskRequestParams, Literal["tasks/cancel"]] +): + """``tasks/cancel`` request envelope for ``ClientSession.send_request``.""" + + method: Literal["tasks/cancel"] = "tasks/cancel" + params: CancelTaskRequestParams diff --git a/fastmcp_tasks/fastmcp_tasks/settings.py b/fastmcp_tasks/fastmcp_tasks/settings.py index 4f9fec622..3b22d10cc 100644 --- a/fastmcp_tasks/fastmcp_tasks/settings.py +++ b/fastmcp_tasks/fastmcp_tasks/settings.py @@ -120,3 +120,38 @@ class DocketSettings(BaseSettings): docket_settings = DocketSettings() + + +class TasksClientSettings(BaseSettings): + """Client-side settings for driving background tasks. + + Moved here from core ``fastmcp.settings`` during the SEP-1686 -> SEP-2663 + migration: the entire client task-driving path now lives in + ``fastmcp-tasks``, so its one tunable does too. + """ + + model_config = SettingsConfigDict( + env_prefix="FASTMCP_TASKS_CLIENT_", + extra="ignore", + ) + + poll_interval: Annotated[ + float, + Field( + description=inspect.cleandoc( + """ + Ceiling, in seconds, for the fallback poll backoff while the client + waits on a background task. Applies only when the server does not + advertise its own pollIntervalMs: in that case the client starts + polling fast (~20ms) and doubles up to this ceiling, so quick tasks + resolve promptly while long-running tasks don't hammer the server. + When the server advertises a pollIntervalMs, that interval is honored + exactly and this setting is ignored. Must be positive. + """ + ), + gt=0, + ), + ] = 0.5 + + +client_settings = TasksClientSettings() diff --git a/pyproject.toml b/pyproject.toml index d34a8fc68..71bad7539 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -155,9 +155,6 @@ exclude = [ "examples/providers/sqlite", # needs aiosqlite "examples/memory.py", # needs asyncpg, numpy, pydantic_ai, pgvector "examples/get_file.py", # needs aiohttp - # Skipped pending client task support; rewritten in the client-task follow-up. - "tests/tasks/client/test_task_context_validation.py", - "tests/tasks/client/test_task_result_caching.py", ] [tool.ty.environment] diff --git a/tests/client/client/test_client.py b/tests/client/client/test_client.py index 38f733f63..d0edc194a 100644 --- a/tests/client/client/test_client.py +++ b/tests/client/client/test_client.py @@ -7,7 +7,6 @@ from typing import Any, cast import anyio import pytest -from fastmcp_tasks.client import TaskNotificationHandler from mcp import ClientSession, MCPError from mcp_types import TextContent from pydantic import AnyUrl @@ -886,34 +885,19 @@ async def test_client_list_dict_return_type(): assert result.data == [{"city": "NYC", "temp": 72}, {"city": "LA", "temp": 85}] -@pytest.mark.skip(reason="Phase 4: requires client task support (SEP-2663)") -def test_client_new_resets_mutable_task_state(fastmcp_server): - """Client.new() should not share mutable task tracking structures.""" - client = Client(transport=FastMCPTransport(fastmcp_server)) +def test_client_new_preserves_internal_task_extension(fastmcp_server): + """Client.new() rebuilds the clone with the auto-registered tasks claim. - client._task_registry["task-1"] = lambda: None # type: ignore[assignment] # ty: ignore - client._submitted_task_ids.add("task-1") # ty: ignore + The tasks client extension (from fastmcp-tasks, imported above) is folded into + every Client automatically; a clone must carry it too so tasked calls still + resolve transparently on the clone. + """ + from fastmcp_tasks.client_models import ClientCreateTaskResult + + client = Client(transport=FastMCPTransport(fastmcp_server)) + assert ClientCreateTaskResult in client._claim_by_model clone = client.new() - assert clone is not client - assert clone._task_registry == {} # ty: ignore - assert clone._submitted_task_ids == set() # ty: ignore - assert clone._task_registry is not client._task_registry # ty: ignore - assert clone._submitted_task_ids is not client._submitted_task_ids # ty: ignore - - -@pytest.mark.skip(reason="Phase 4: requires client task support (SEP-2663)") -def test_client_new_rebinds_default_task_notification_handler(fastmcp_server): - """Client.new() should bind the default task handler to the cloned client.""" - client = Client(transport=FastMCPTransport(fastmcp_server)) - - handler = client._session_kwargs.get("message_handler") - assert isinstance(handler, TaskNotificationHandler) - - clone = client.new() - - clone_handler = clone._session_kwargs.get("message_handler") - assert isinstance(clone_handler, TaskNotificationHandler) - assert clone_handler is not handler - assert clone_handler._client_ref() is clone + assert ClientCreateTaskResult in clone._claim_by_model + assert clone._claim_by_model is not client._claim_by_model diff --git a/tests/client/telemetry/test_client_task_tracing.py b/tests/client/telemetry/test_client_task_tracing.py deleted file mode 100644 index 4c93a2fed..000000000 --- a/tests/client/telemetry/test_client_task_tracing.py +++ /dev/null @@ -1,98 +0,0 @@ -"""Tests for client OpenTelemetry tracing on task operations.""" - -import asyncio - -import pytest -from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( - InMemorySpanExporter, -) -from opentelemetry.trace import SpanKind - -from fastmcp import Client, FastMCP - -pytestmark = pytest.mark.skip(reason="Phase 4: requires client task support (SEP-2663)") - - -def assert_propagating_client_span( - trace_exporter: InMemorySpanExporter, - method: str, - component_key: str, -) -> None: - all_spans = trace_exporter.get_finished_spans() - spans = [span for span in all_spans if span.name == method] - client_span = next( - span - for span in spans - if span.attributes is not None and "fastmcp.server.name" not in span.attributes - ) - server_span = next( - span - for span in spans - if span.attributes is not None and "fastmcp.server.name" in span.attributes - ) - - assert client_span.kind == SpanKind.CLIENT - assert client_span.attributes is not None - assert client_span.attributes["mcp.method.name"] == method - assert client_span.attributes["fastmcp.component.key"] == component_key - assert server_span.parent is not None - assert server_span.context.trace_id == client_span.context.trace_id - - spans_by_id = {span.context.span_id: span for span in all_spans} - current = server_span - while current.parent is not None: - parent = spans_by_id.get(current.parent.span_id) - assert parent is not None - if parent.context.span_id == client_span.context.span_id: - break - current = parent - else: - raise AssertionError("Server span should descend from the client span") - - -async def test_list_tasks_creates_propagating_client_span( - trace_exporter: InMemorySpanExporter, -): - server = FastMCP("test-server") - - async with Client(server, mode="legacy") as client: - await client.list_tasks() - - assert_propagating_client_span(trace_exporter, "tasks/list", "") - - -async def test_task_id_operations_create_propagating_client_spans( - trace_exporter: InMemorySpanExporter, -): - started = asyncio.Event() - server = FastMCP("test-server") - - @server.tool(task=True) - async def quick_tool() -> str: - return "done" - - @server.tool(task=True) - async def slow_tool() -> str: - started.set() - # Never completes on its own - the test cancels this task well - # before any real-time completion would matter. - await asyncio.Event().wait() - return "done" - - async with Client(server, mode="legacy") as client: - completed_task = await client.call_tool("quick_tool", task=True) - await completed_task.wait(timeout=2) - trace_exporter.clear() - - await client.get_task_status(completed_task.task_id) - await client.get_task_result(completed_task.task_id) - - running_task = await client.call_tool("slow_tool", task=True) - await asyncio.wait_for(started.wait(), timeout=2) - await client.cancel_task(running_task.task_id) - - assert_propagating_client_span(trace_exporter, "tasks/get", completed_task.task_id) - assert_propagating_client_span( - trace_exporter, "tasks/result", completed_task.task_id - ) - assert_propagating_client_span(trace_exporter, "tasks/cancel", running_task.task_id) diff --git a/tests/client/test_client_extensions.py b/tests/client/test_client_extensions.py index 03af1e89c..1a054ae83 100644 --- a/tests/client/test_client_extensions.py +++ b/tests/client/test_client_extensions.py @@ -1,17 +1,21 @@ """Tests for surfacing SEP-2133 client extensions on ``fastmcp.Client``. Covers that ``extensions=`` / ``result_claims=`` are folded into the underlying -``ClientSession`` kwargs on construction, that user-supplied notification -bindings *compose* with FastMCP's internal task-status binding rather than -clobbering it, that both bindings actually fire against a live server, and that -a claimed ``tools/call`` result is resolved end-to-end through the owning -extension's resolver. +``ClientSession`` kwargs on construction, that a claimed ``tools/call`` result is +resolved end-to-end through the owning extension's resolver, and that FastMCP's +internal tasks extension (from ``fastmcp-tasks``, imported below) is folded in +automatically and *composes* with a user's own extensions rather than being +clobbered by them. + +Importing ``fastmcp_tasks`` registers the internal client extension factory +process-wide, so every ``Client`` built here carries the tasks capability ad and +its ``resultType: "task"`` claim. These tests assert that composition explicitly. """ -import asyncio from typing import Any, Literal import pytest +from fastmcp_tasks.client_models import ClientCreateTaskResult from mcp.client.extension import ( ClaimContext, ClientExtension, @@ -26,12 +30,15 @@ from mcp_types import CallToolRequestParams, CallToolResult, Result, TextContent from mcp_types.version import LATEST_MODERN_VERSION from pydantic import BaseModel +# Importing the package registers the internal tasks client extension factory, so +# every Client below folds the tasks extension in. Kept as an explicit import so +# the composition assertions are deterministic regardless of test import order. +import fastmcp_tasks # noqa: F401 from fastmcp import FastMCP from fastmcp.client import Client -from fastmcp.server.dependencies import get_context +from fastmcp.utilities.tasks import TASKS_EXTENSION_ID CUSTOM_METHOD = "notifications/x-test/ping" -TASK_STATUS_METHOD = "notifications/tasks/status" EXTENSION_ID = "test.example.com/demo" CLAIMED_TYPE = "x-test/claimed" @@ -120,19 +127,19 @@ def _claiming_server() -> SDKServer: return server -def _binding_methods(client: Client) -> list[str]: - bindings = client._session_kwargs.get("notification_bindings") or [] - return [b.method for b in bindings] - - def test_extension_folds_into_session_kwargs(): - """A ClientExtension's ad, claim, and binding reach the session kwargs.""" + """A ClientExtension's ad and claim reach the session kwargs, alongside tasks.""" client = Client(FastMCP("srv"), extensions=[_DemoExtension()]) - assert client._session_kwargs.get("extensions") == {EXTENSION_ID: {"enabled": True}} + # The tasks extension is auto-folded in beside the user's own. + assert client._session_kwargs.get("extensions") == { + TASKS_EXTENSION_ID: {}, + EXTENSION_ID: {"enabled": True}, + } result_claims = client._session_kwargs.get("result_claims") assert result_claims is not None assert [c.result_type for c in result_claims[EXTENSION_ID]] == [CLAIMED_TYPE] + assert [c.result_type for c in result_claims[TASKS_EXTENSION_ID]] == ["task"] def test_extension_populates_claim_by_model_index(): @@ -140,42 +147,62 @@ def test_extension_populates_claim_by_model_index(): client = Client(FastMCP("srv"), extensions=[_DemoExtension()]) assert client._claim_by_model[ClaimedResult].result_type == CLAIMED_TYPE + # The auto-folded tasks claim is indexed too. + assert client._claim_by_model[ClientCreateTaskResult].result_type == "task" -@pytest.mark.skip(reason="Phase 4: requires client task support (SEP-2663)") -def test_binding_composes_with_internal_task_binding(): - """User binding is appended to (not replacing) the task-status binding.""" - client = Client(FastMCP("srv"), extensions=[_DemoExtension()]) - - methods = _binding_methods(client) - assert TASK_STATUS_METHOD in methods - assert CUSTOM_METHOD in methods - # The internal task binding must lead so user bindings extend it. - assert methods[0] == TASK_STATUS_METHOD - - -@pytest.mark.skip(reason="Phase 4: requires client task support (SEP-2663)") -def test_no_extensions_leaves_only_task_binding(): - """Without extensions, only the internal task-status binding is registered.""" +def test_internal_tasks_extension_present_without_user_extensions(): + """Even with no user extensions, the tasks claim is auto-registered.""" client = Client(FastMCP("srv")) - assert _binding_methods(client) == [TASK_STATUS_METHOD] - assert "extensions" not in client._session_kwargs - assert "result_claims" not in client._session_kwargs + assert client._session_kwargs.get("extensions") == {TASKS_EXTENSION_ID: {}} + assert client._claim_by_model[ClientCreateTaskResult].result_type == "task" + + +def test_user_extension_composes_with_internal_tasks_extension(): + """A user extension is folded in beside the internal tasks extension.""" + client = Client(FastMCP("srv"), extensions=[_DemoExtension()]) + + ad = client._session_kwargs.get("extensions") or {} + assert TASKS_EXTENSION_ID in ad + assert EXTENSION_ID in ad + # Both claims are resolvable. + assert set(client._claim_by_model) == {ClaimedResult, ClientCreateTaskResult} + + +def test_user_extension_may_override_internal_tasks_extension(): + """A user extension declaring the tasks identifier wins; the internal one drops. + + Composition prefers the user's extension: rather than colliding on the shared + identifier (which the fold rejects), the internal tasks extension is dropped so + a power user can supply their own task-handling extension. + """ + + class CustomTasks(ClientExtension): + identifier = TASKS_EXTENSION_ID + + def settings(self) -> dict[str, Any]: + return {"custom": True} + + client = Client(FastMCP("srv"), extensions=[CustomTasks()]) + + assert client._session_kwargs.get("extensions") == { + TASKS_EXTENSION_ID: {"custom": True} + } + # The user extension declares no claim, so no task claim is registered. assert client._claim_by_model == {} -@pytest.mark.skip(reason="Phase 4: requires client task support (SEP-2663)") def test_new_preserves_extension_composition(): - """new() rebuilds the clone with both the task binding and user bindings.""" + """new() rebuilds the clone with both the tasks extension and user extensions.""" client = Client(FastMCP("srv"), extensions=[_DemoExtension()]) clone = client.new() - methods = _binding_methods(clone) - assert methods[0] == TASK_STATUS_METHOD - assert CUSTOM_METHOD in methods - assert clone._session_kwargs.get("extensions") == {EXTENSION_ID: {"enabled": True}} + ad = clone._session_kwargs.get("extensions") or {} + assert TASKS_EXTENSION_ID in ad + assert EXTENSION_ID in ad assert clone._claim_by_model[ClaimedResult].result_type == CLAIMED_TYPE + assert clone._claim_by_model[ClientCreateTaskResult].result_type == "task" def test_result_claims_merge_with_extension_claims(): @@ -203,80 +230,12 @@ def test_result_claims_merge_with_extension_claims(): assert result_claims is not None tags = {c.result_type for c in result_claims[EXTENSION_ID]} assert tags == {CLAIMED_TYPE, "x-test/extra"} - # Both the extension claim and the explicit extra claim are resolvable. - assert set(client._claim_by_model) == {ClaimedResult, ExtraClaimed} - - -@pytest.mark.skip(reason="Phase 4: requires client task support (SEP-2663)") -async def test_user_binding_clobbering_task_method_is_rejected(): - """A user extension binding the task-status method cannot silently replace it. - - Composition means the internal task binding always leads; a user extension - that binds the same method collides with it, and the SDK session rejects the - duplicate at connect time rather than letting one silently win. - """ - - class TaskClobberExtension(ClientExtension): - identifier = "test.example.com/clobber" - - def notifications(self): - async def _handler(params: PingParams) -> None: - return None - - return ( - NotificationBinding( - method=TASK_STATUS_METHOD, - params_type=PingParams, - handler=_handler, - ), - ) - - client = Client(FastMCP("srv"), extensions=[TaskClobberExtension()]) - with pytest.raises(RuntimeError, match="duplicate notification binding"): - async with client: - pass - - -@pytest.mark.skip(reason="Phase 4: requires client task support (SEP-2663)") -async def test_both_bindings_fire_against_live_server(): - """The internal task binding and a user extension binding both fire. - - A ``task=True`` tool drives ``notifications/tasks/status`` (the internal - binding) while a second tool emits a custom notification the user extension - observes, proving the two coexist on one live connection. Pinned to - ``mode="legacy"`` because FastMCP task submission is a legacy-era feature. - """ - received: list[PingParams] = [] - mcp = FastMCP("compose-server") - - @mcp.tool - async def emit(value: int) -> int: - ctx = get_context() - # Emit a custom (non-core) notification straight onto the outbound - # channel; unknown methods route to the client's notification bindings. - await ctx.session._connection.notify(CUSTOM_METHOD, {"value": value}) - return value - - @mcp.tool(task=True) - async def background(value: int) -> int: - await asyncio.sleep(0.02) - return value * 2 - - client = Client(mcp, extensions=[_DemoExtension(received)], mode="legacy") - - async with client: - # The user extension binding fires on the custom notification. - await client.call_tool("emit", {"value": 21}) - # The internal task binding fires on the task-status notification. - task = await client.call_tool("background", {"value": 5}, task=True) # ty: ignore - status = await task.wait(timeout=2.0) # ty: ignore - # Give the custom-notification queue a moment to drain. - await asyncio.sleep(0.1) - - # Internal task binding fired: the task completed via a status notification. - assert status.status == "completed" - # User extension binding fired: it observed the custom notification. - assert [p.value for p in received] == [21] + # The extension claim, the explicit extra claim, and the tasks claim resolve. + assert set(client._claim_by_model) == { + ClaimedResult, + ExtraClaimed, + ClientCreateTaskResult, + } class TestClaimedResultResolution: diff --git a/tests/tasks/client/test_client_task_notifications.py b/tests/tasks/client/test_client_task_notifications.py deleted file mode 100644 index 74cb5b207..000000000 --- a/tests/tasks/client/test_client_task_notifications.py +++ /dev/null @@ -1,283 +0,0 @@ -""" -Tests for client-side handling of notifications/tasks/status (SEP-1686 lines 436-444). - -Verifies that Task objects receive notifications, update their cache, wake up wait() calls, -and invoke user callbacks. -""" - -import asyncio -import time -from collections.abc import Callable -from datetime import datetime, timezone - -import pytest -from mcp_types import GetTaskResult - -from fastmcp import FastMCP -from fastmcp.client import Client - -pytestmark = pytest.mark.skip(reason="Phase 4: requires client task support (SEP-2663)") - - -async def _wait_until(condition: Callable[[], bool], timeout: float = 5.0) -> None: - """Poll until condition() is true or timeout elapses. - - Used in place of a fixed sleep when waiting for an async callback or - notification to be delivered/dispatched after the awaited call returns. - """ - deadline = time.monotonic() + timeout - while not condition() and time.monotonic() < deadline: - await asyncio.sleep(0.005) - - -@pytest.fixture -async def task_notification_server(): - """Server that sends task status notifications.""" - mcp = FastMCP("task-notification-test") - - @mcp.tool(task=True) - async def quick_task(value: int) -> int: - """Quick background task with a brief, measurable delay (contrast with instant_task).""" - await asyncio.sleep(0.01) - return value * 2 - - @mcp.tool(task=True) - async def instant_task(value: int) -> int: - """Background task that completes with no delay.""" - return value * 2 - - @mcp.tool(task=True) - async def failing_task() -> str: - """Task that fails.""" - raise ValueError("Intentional failure") - - return mcp - - -async def test_task_receives_status_notification(task_notification_server): - """Task object receives and processes status notifications.""" - async with Client(task_notification_server, mode="legacy") as client: - task = await client.call_tool("quick_task", {"value": 5}, task=True) - - # Wait for task to complete (notification should arrive) - status = await task.wait(timeout=2.0) - - # Verify task completed - assert status.status == "completed" - - -async def test_status_cache_updated_by_notification(task_notification_server): - """Cached status is updated when notification arrives.""" - async with Client(task_notification_server, mode="legacy") as client: - task = await client.call_tool("quick_task", {"value": 10}, task=True) - - # Wait for completion (notification should update cache) - await task.wait(timeout=2.0) - - # Status should be cached (no server call needed) - # Call status() twice - should return same cached object - status1 = await task.status() - status2 = await task.status() - - # Should be the exact same object (from cache) - assert status1 is status2 - assert status1.status == "completed" - - -async def test_callback_invoked_on_notification(task_notification_server): - """User callback is invoked when notification arrives.""" - callback_invocations = [] - - def status_callback(status: GetTaskResult): - """Sync callback.""" - callback_invocations.append(status) - - async with Client(task_notification_server, mode="legacy") as client: - task = await client.call_tool("quick_task", {"value": 7}, task=True) - - # Register callback - task.on_status_change(status_callback) - - # Wait for completion - await task.wait(timeout=2.0) - - # Wait for the status this test actually asserts on. Waiting merely for - # "some callback fired" would be satisfied by the earlier `working` - # notification and race the `completed` one. - await _wait_until( - lambda: any(s.status == "completed" for s in callback_invocations) - ) - - # Callback should have been invoked at least once - assert len(callback_invocations) > 0 - - # Should have received completed status - completed_statuses = [s for s in callback_invocations if s.status == "completed"] - assert len(completed_statuses) > 0 - - -async def test_async_callback_invoked(task_notification_server): - """Async callback is invoked when notification arrives.""" - callback_invocations = [] - - async def async_status_callback(status: GetTaskResult): - """Async callback.""" - await asyncio.sleep(0.01) # Simulate async work - callback_invocations.append(status) - - async with Client(task_notification_server, mode="legacy") as client: - task = await client.call_tool("quick_task", {"value": 3}, task=True) - - # Register async callback - task.on_status_change(async_status_callback) - - # Wait for completion - await task.wait(timeout=2.0) - - # Give async callbacks time to complete - await _wait_until(lambda: len(callback_invocations) > 0) - - # Async callback should have been invoked - assert len(callback_invocations) > 0 - - -async def test_multiple_callbacks_all_invoked(task_notification_server): - """Multiple callbacks are all invoked.""" - callback1_calls = [] - callback2_calls = [] - - def callback1(status: GetTaskResult): - callback1_calls.append(status.status) - - def callback2(status: GetTaskResult): - callback2_calls.append(status.status) - - async with Client(task_notification_server, mode="legacy") as client: - task = await client.call_tool("quick_task", {"value": 8}, task=True) - - task.on_status_change(callback1) - task.on_status_change(callback2) - - await task.wait(timeout=2.0) - await _wait_until(lambda: bool(callback1_calls) and bool(callback2_calls)) - - # Both callbacks should have been invoked - assert len(callback1_calls) > 0 - assert len(callback2_calls) > 0 - - -async def test_callback_error_doesnt_break_notification(task_notification_server): - """Callback errors don't prevent other callbacks from running.""" - callback1_calls = [] - callback2_calls = [] - - def failing_callback(status: GetTaskResult): - callback1_calls.append("called") - raise ValueError("Callback intentionally fails") - - def working_callback(status: GetTaskResult): - callback2_calls.append(status.status) - - async with Client(task_notification_server, mode="legacy") as client: - task = await client.call_tool("quick_task", {"value": 12}, task=True) - - task.on_status_change(failing_callback) - task.on_status_change(working_callback) - - await task.wait(timeout=2.0) - await _wait_until(lambda: bool(callback1_calls) and bool(callback2_calls)) - - # Failing callback was called (and errored) - assert len(callback1_calls) > 0 - - # Working callback should still have been invoked - assert len(callback2_calls) > 0 - - -async def test_wait_wakes_early_on_notification(task_notification_server): - """wait() wakes up immediately when notification arrives, not after poll interval.""" - async with Client(task_notification_server, mode="legacy") as client: - task = await client.call_tool("quick_task", {"value": 15}, task=True) - - # Record timing - start = time.time() - status = await task.wait(timeout=5.0) - elapsed = time.time() - start - - # Should complete much faster than the fallback poll interval (500ms) - # With notifications, should be < 200ms for quick task - # Without notifications, would take 500ms+ due to polling - assert elapsed < 1.0 # Very generous bound - assert status.status == "completed" - - -async def test_notification_with_failed_task(task_notification_server): - """Notifications work for failed tasks too.""" - async with Client(task_notification_server, mode="legacy") as client: - task = await client.call_tool("failing_task", {}, task=True) - - with pytest.raises(Exception): - await task - - # Should have cached the failed status from notification - status = await task.status() - assert status.status == "failed" - assert ( - status.status_message is not None - ) # Error details in statusMessage per spec - - -async def test_fast_task_completion_delivered_via_notification( - task_notification_server, -): - """A near-instant task still delivers its completion via a status notification. - - Regression test for the Docket subscribe() setup-window race: a task that - finishes before the pub/sub subscription goes live had its terminal state - publish lost, so no completion notification ever reached the client and - wait() fell back to a full poll interval. The server now reconciles the - execution against Redis to close that gap. - - Callbacks fire only for received notifications β€” client-side polling updates - the status cache directly without invoking them β€” so a "completed" callback - proves the notification path (not the poll fallback) was exercised. - """ - received: list[str] = [] - - async with Client(task_notification_server, mode="legacy") as client: - task = await client.call_tool("instant_task", {"value": 21}, task=True) - task.on_status_change(lambda status: received.append(status.status)) - - result = await task - assert result.data == 42 - - # Allow the completion notification to arrive and dispatch. - await _wait_until(lambda: "completed" in received) - - assert "completed" in received - - -async def test_wait_returns_on_input_required(task_notification_server): - """wait() should return immediately when task enters input_required, not hang.""" - async with Client(task_notification_server, mode="legacy") as client: - task = await client.call_tool("quick_task", {"value": 1}, task=True) - - # Directly inject an input_required status into the cache and signal the event. - # SDK v2 types the Task timestamps as ISO 8601 strings. - now = datetime.now(timezone.utc).isoformat() - input_required_status = GetTaskResult( - task_id=task._task_id, - status="input_required", - status_message="Waiting for user input", - created_at=now, - last_updated_at=now, - ttl=None, - ) - task._status_cache = input_required_status - if task._status_event is None: - task._status_event = asyncio.Event() - task._status_event.set() - - # Should return immediately with input_required, not hang for 300s - status = await task.wait(timeout=2.0) - assert status.status == "input_required" diff --git a/tests/tasks/client/test_client_task_protocol.py b/tests/tasks/client/test_client_task_protocol.py deleted file mode 100644 index 7b9698bb2..000000000 --- a/tests/tasks/client/test_client_task_protocol.py +++ /dev/null @@ -1,89 +0,0 @@ -""" -Tests for client-side task protocol. - -Generic protocol tests that use tools as test fixtures. -""" - -import asyncio - -import pytest - -from fastmcp import FastMCP -from fastmcp.client import Client - -pytestmark = pytest.mark.skip(reason="Phase 4: requires client task support (SEP-2663)") - - -async def test_end_to_end_task_flow(): - """Complete end-to-end flow: submit, poll, retrieve.""" - start_signal = asyncio.Event() - complete_signal = asyncio.Event() - - mcp = FastMCP("protocol-test") - - @mcp.tool(task=True) - async def controlled_tool(message: str) -> str: - """Tool with controlled execution.""" - start_signal.set() - await complete_signal.wait() - return f"Processed: {message}" - - async with Client(mcp, mode="legacy") as client: - # Submit task - task = await client.call_tool( - "controlled_tool", {"message": "integration test"}, task=True - ) - - # Wait for execution to start - await asyncio.wait_for(start_signal.wait(), timeout=2.0) - - # Check status while running - status = await task.status() - assert status.status in ["working"] - - # Signal completion - complete_signal.set() - - # Wait for task to finish and retrieve result - result = await task.result() - assert result.data == "Processed: integration test" - - -async def test_multiple_concurrent_tasks(): - """Multiple tasks can run concurrently.""" - mcp = FastMCP("concurrent-test") - - @mcp.tool(task=True) - async def multiply(a: int, b: int) -> int: - return a * b - - async with Client(mcp, mode="legacy") as client: - # Submit multiple tasks - tasks = [] - for i in range(5): - task = await client.call_tool("multiply", {"a": i, "b": 2}, task=True) - tasks.append((task, i * 2)) - - # Wait for all to complete and verify results - for task, expected in tasks: - result = await task.result() - assert result.data == expected - - -async def test_task_id_auto_generation(): - """Task IDs are auto-generated if not provided.""" - mcp = FastMCP("id-test") - - @mcp.tool(task=True) - async def echo(message: str) -> str: - return f"Echo: {message}" - - async with Client(mcp, mode="legacy") as client: - # Submit without custom task ID - task_1 = await client.call_tool("echo", {"message": "first"}, task=True) - task_2 = await client.call_tool("echo", {"message": "second"}, task=True) - - # Should generate different IDs - assert task_1.task_id != task_2.task_id - assert len(task_1.task_id) > 0 - assert len(task_2.task_id) > 0 diff --git a/tests/tasks/client/test_client_tool_tasks.py b/tests/tasks/client/test_client_tool_tasks.py index 0a8220140..815adaeba 100644 --- a/tests/tasks/client/test_client_tool_tasks.py +++ b/tests/tasks/client/test_client_tool_tasks.py @@ -1,158 +1,152 @@ -""" -Tests for client-side tool task methods. +"""The explicit `ToolTask` handle (the return-quickly surface, SEP-2663). -Tests the client's tool-specific task functionality, parallel to -test_client_prompt_tasks.py and test_client_resource_tasks.py. +`call_tool_task` returns a `ToolTask` as soon as the server accepts the task, so +the caller can do other work and drive it: `status`, `wait`, `result`, `cancel`, +or `await`. This contrasts with `client.call_tool`, which polls to completion +transparently. All tests use a real `Client(mode="auto")` over the in-memory +transport, since tasks are modern-only. """ +from __future__ import annotations + +import asyncio + import pytest -from fastmcp_tasks.client import ToolTask +from fastmcp_tasks.models import MISSING_REQUIRED_CLIENT_CAPABILITY +from mcp.shared.exceptions import MCPError -from fastmcp import FastMCP +from fastmcp import Context, FastMCP from fastmcp.client import Client from fastmcp.exceptions import ToolError - -pytestmark = pytest.mark.skip(reason="Phase 4: requires client task support (SEP-2663)") +from fastmcp.utilities.tasks import TaskConfig +from fastmcp_tasks import TasksExtension, ToolTask, call_tool_task @pytest.fixture -async def tool_task_server(): - """Create a test server with task-enabled tools.""" +def tool_task_server() -> FastMCP: mcp = FastMCP("tool-task-test") + mcp.add_extension(TasksExtension()) @mcp.tool(task=True) async def echo(message: str) -> str: - """Echo back the message.""" return f"Echo: {message}" @mcp.tool(task=True) async def multiply(a: int, b: int) -> int: - """Multiply two numbers.""" return a * b + @mcp.tool(task=True) + async def boom() -> str: + raise ValueError("background task failure") + return mcp -async def test_call_tool_as_task_returns_tool_task(tool_task_server): - """call_tool with task=True returns a ToolTask object.""" - async with Client(tool_task_server, mode="legacy") as client: - task = await client.call_tool("echo", {"message": "hello"}, task=True) +async def test_call_tool_task_returns_tool_task(tool_task_server: FastMCP): + async with Client(tool_task_server, mode="auto") as client: + task = await call_tool_task(client, "echo", {"message": "hello"}) assert isinstance(task, ToolTask) assert isinstance(task.task_id, str) - assert len(task.task_id) > 0 + assert task.task_id -async def test_tool_task_server_generated_id(tool_task_server): - """call_tool with task=True gets server-generated task ID.""" - async with Client(tool_task_server, mode="legacy") as client: - task = await client.call_tool("echo", {"message": "test"}, task=True) - - # Server should generate a UUID task ID - assert task.task_id is not None - assert isinstance(task.task_id, str) - # UUIDs have hyphens - assert "-" in task.task_id - - -async def test_tool_task_result_returns_call_tool_result(tool_task_server): - """ToolTask.result() returns CallToolResult with tool data.""" - async with Client(tool_task_server, mode="legacy") as client: - task = await client.call_tool("multiply", {"a": 6, "b": 7}, task=True) - assert not task.returned_immediately - +async def test_tool_task_result_returns_parsed_result(tool_task_server: FastMCP): + async with Client(tool_task_server, mode="auto") as client: + task = await call_tool_task(client, "multiply", {"a": 6, "b": 7}) result = await task.result() assert result.data == 42 -async def test_tool_task_await_syntax(tool_task_server): - """Tool tasks can be awaited directly to get result.""" - async with Client(tool_task_server, mode="legacy") as client: - task = await client.call_tool("multiply", {"a": 7, "b": 6}, task=True) - - # Can await task directly (syntactic sugar for task.result()) +async def test_tool_task_await_syntax(tool_task_server: FastMCP): + async with Client(tool_task_server, mode="auto") as client: + task = await call_tool_task(client, "multiply", {"a": 7, "b": 6}) result = await task assert result.data == 42 -async def test_tool_task_status_and_wait(tool_task_server): - """ToolTask.status() returns GetTaskResult.""" - async with Client(tool_task_server, mode="legacy") as client: - task = await client.call_tool("echo", {"message": "test"}, task=True) +async def test_tool_task_status_and_wait(tool_task_server: FastMCP): + async with Client(tool_task_server, mode="auto") as client: + task = await call_tool_task(client, "echo", {"message": "test"}) status = await task.status() assert status.task_id == task.task_id - assert status.status in ["working", "completed"] + assert status.status in {"working", "completed"} - # Wait for completion - await task.wait(timeout=2.0) - final_status = await task.status() - assert final_status.status == "completed" + final = await task.wait(timeout=2.0) + assert final.status == "completed" -async def test_immediate_tool_task_respects_raise_on_error_true(): - """Immediate task fallback should still raise ToolError when requested.""" - mcp = FastMCP("immediate-tool-task-error") +async def test_tool_task_result_is_cached(tool_task_server: FastMCP): + """Repeated result() calls return the same cached object without re-polling.""" + async with Client(tool_task_server, mode="auto") as client: + task = await call_tool_task(client, "multiply", {"a": 2, "b": 5}) - @mcp.tool - def failing_tool() -> str: - raise ValueError("immediate task failure") - - async with Client(mcp, mode="legacy") as client: - task = await client.call_tool("failing_tool", task=True, raise_on_error=True) - - assert task.returned_immediately - with pytest.raises( - ToolError, match="does not support task-augmented execution" - ): - await task.result() + result1 = await task.result() + result2 = await task.result() + result3 = await task + assert result1 is result2 is result3 + assert result1.data == 10 -async def test_immediate_tool_task_respects_raise_on_error_false(): - """Immediate task fallback should return error results when requested.""" - mcp = FastMCP("immediate-tool-task-no-raise") - - @mcp.tool - def failing_tool() -> str: - raise ValueError("immediate task failure") - - async with Client(mcp, mode="legacy") as client: - task = await client.call_tool("failing_tool", task=True, raise_on_error=False) - - assert task.returned_immediately - result = await task.result() - assert result.is_error is True - assert "does not support task-augmented execution" in str(result) - - -async def test_background_tool_task_respects_raise_on_error_true(): - """Background tasks should still raise ToolError by default on errors.""" - mcp = FastMCP("background-tool-task-error") - - @mcp.tool(task=True) - async def failing_tool() -> str: - raise ValueError("background task failure") - - async with Client(mcp, mode="legacy") as client: - task = await client.call_tool("failing_tool", task=True, raise_on_error=True) - - assert not task.returned_immediately +async def test_background_task_raises_on_error_by_default(tool_task_server: FastMCP): + async with Client(tool_task_server, mode="auto") as client: + task = await call_tool_task(client, "boom", {}) with pytest.raises(ToolError, match="background task failure"): await task.result() -async def test_background_tool_task_respects_raise_on_error_false(): - """Background tasks should return error results when raise_on_error is disabled.""" - mcp = FastMCP("background-tool-task-no-raise") +async def test_background_task_returns_error_when_not_raising( + tool_task_server: FastMCP, +): + async with Client(tool_task_server, mode="auto") as client: + task = await call_tool_task(client, "boom", {}, raise_on_error=False) + result = await task.result() + assert result.is_error + assert "background task failure" in str(result) + + +async def test_multiple_concurrent_tool_tasks(tool_task_server: FastMCP): + async with Client(tool_task_server, mode="auto") as client: + tasks = [ + (await call_tool_task(client, "multiply", {"a": i, "b": 2}), i * 2) + for i in range(5) + ] + for task, expected in tasks: + result = await task.result() + assert result.data == expected + + +async def test_tool_task_cancel(): + """A long-running task can be cancelled through the handle.""" + mcp = FastMCP("cancel-test") + mcp.add_extension(TasksExtension()) @mcp.tool(task=True) - async def failing_tool() -> str: - raise ValueError("background task failure") + async def forever(ctx: Context) -> str: + await asyncio.Event().wait() + return "never" + + async with Client(mcp, mode="auto") as client: + task = await call_tool_task(client, "forever", {}) + await task.wait(state="working", timeout=2.0) + await task.cancel() + final = await task.wait(timeout=2.0) + assert final.status == "cancelled" + + +async def test_required_mode_without_optin_raises_32003(): + """A legacy client never negotiates the tasks capability, so a required-mode + tool call is rejected with the -32003 missing-capability error.""" + mcp = FastMCP("required-test") + mcp.add_extension(TasksExtension()) + + @mcp.tool(task=TaskConfig(mode="required")) + async def must_task(x: int) -> int: + return x async with Client(mcp, mode="legacy") as client: - task = await client.call_tool("failing_tool", task=True, raise_on_error=False) + with pytest.raises(MCPError) as excinfo: + await client.call_tool("must_task", {"x": 1}) - assert not task.returned_immediately - result = await task.result() - assert result.is_error is True - assert "background task failure" in str(result) + assert excinfo.value.error.code == MISSING_REQUIRED_CLIENT_CAPABILITY diff --git a/tests/tasks/client/test_poll_interval.py b/tests/tasks/client/test_poll_interval.py index fa4a25a4b..5d0afab3f 100644 --- a/tests/tasks/client/test_poll_interval.py +++ b/tests/tasks/client/test_poll_interval.py @@ -1,92 +1,61 @@ -"""Fallback poll cadence for client-side task waiting. +"""Fallback poll cadence for client-side task waiting (SEP-2663). -Two modes: a server-advertised pollInterval is honored exactly, while an -unadvertised one falls back to an exponential ramp up to the client setting. +The modern protocol has no task status notifications, so the client polls. The +backoff ramps from a fast floor, doubling up to a ceiling: the server-advertised +``pollIntervalMs`` when present (a statement about server load), else the client +``poll_interval`` setting. A quick task resolves in ~20ms; a long one settles to +the advertised cadence. """ +from __future__ import annotations + import pytest -from fastmcp_tasks.client import MIN_POLL_INTERVAL, ToolTask -from mcp_types import GetTaskResult +from fastmcp_tasks.client import MIN_POLL_INTERVAL, _next_poll_delay, _poll_ceiling +from fastmcp_tasks.settings import TasksClientSettings, client_settings from pydantic import ValidationError -from fastmcp import Client, FastMCP -from fastmcp.settings import Settings -from fastmcp.utilities.tests import temporary_settings - -pytestmark = pytest.mark.skip(reason="Phase 4: requires client task support (SEP-2663)") - @pytest.mark.parametrize("value", [0, -0.5, -1]) def test_non_positive_poll_interval_setting_is_rejected(value: float): with pytest.raises(ValidationError): - Settings(client_task_poll_interval=value) + TasksClientSettings(poll_interval=value) def test_positive_poll_interval_setting_is_accepted(): - settings = Settings(client_task_poll_interval=0.25) - assert settings.client_task_poll_interval == 0.25 + settings = TasksClientSettings(poll_interval=0.25) + assert settings.poll_interval == 0.25 -@pytest.fixture -def task() -> ToolTask: - client = Client(FastMCP()) - return ToolTask(client=client, task_id="t1", tool_name="echo") +@pytest.mark.parametrize("poll_interval_ms", [2000, 30_000]) +def test_advertised_interval_caps_the_ramp(poll_interval_ms: int): + """An advertised interval is the ceiling the ramp tops out at.""" + assert _poll_ceiling(poll_interval_ms) == poll_interval_ms / 1000 -def _status(poll_interval: int | None) -> GetTaskResult: - return GetTaskResult( - task_id="t1", - status="working", - created_at="2026-01-01T00:00:00+00:00", - last_updated_at="2026-01-01T00:00:00+00:00", - ttl=None, - poll_interval=poll_interval, - ) +def test_large_advertised_interval_is_honored(): + day_ms = 24 * 60 * 60 * 1000 + assert _poll_ceiling(day_ms) == 24 * 60 * 60 -@pytest.mark.parametrize("poll_interval", [2000, 30_000]) -def test_advertised_interval_is_used_verbatim_without_backoff( - task: ToolTask, poll_interval: int -): - """An advertised interval is the delay itself, not a ceiling to ramp toward.""" - task._status_cache = _status(poll_interval) - expected = poll_interval / 1000 +@pytest.mark.parametrize("poll_interval_ms", [None, 0, -1, -5000]) +def test_absent_or_hostile_interval_falls_back_to_setting(poll_interval_ms): + """An absent, zero, or negative server value cannot spin the client: use the setting.""" + assert _poll_ceiling(poll_interval_ms) == client_settings.poll_interval + +def test_ramp_doubles_from_floor_up_to_advertised_ceiling(): + """Even with an advertised interval, the poll ramps fast then caps at it.""" + ceiling_ms = 500 # 0.5s ceiling + delays = [] backoff = MIN_POLL_INTERVAL - for _ in range(5): - delay, backoff = task._next_poll_delay(backoff) - assert delay == expected - - -def test_large_advertised_interval_is_honored(task: ToolTask): - task._status_cache = _status(24 * 60 * 60 * 1000) - delay, _ = task._next_poll_delay(MIN_POLL_INTERVAL) - assert delay == 24 * 60 * 60 - - -@pytest.mark.parametrize("poll_interval", [0, -1, -5000]) -def test_non_positive_advertised_interval_is_floored( - task: ToolTask, poll_interval: int -): - """A buggy or hostile server must not be able to spin the client.""" - task._status_cache = _status(poll_interval) - delay, _ = task._next_poll_delay(MIN_POLL_INTERVAL) - assert delay == MIN_POLL_INTERVAL - - -def test_unadvertised_interval_ramps_up_to_setting(task: ToolTask): - task._status_cache = _status(None) - with temporary_settings(client_task_poll_interval=0.5): - delays = [] - backoff = MIN_POLL_INTERVAL - for _ in range(7): - delay, backoff = task._next_poll_delay(backoff) - delays.append(delay) + for _ in range(7): + delay, backoff = _next_poll_delay(ceiling_ms, backoff) + delays.append(delay) assert delays == [0.02, 0.04, 0.08, 0.16, 0.32, 0.5, 0.5] -def test_missing_status_cache_ramps_from_floor(task: ToolTask): - delay, backoff = task._next_poll_delay(MIN_POLL_INTERVAL) +def test_first_delay_is_the_floor(): + delay, backoff = _next_poll_delay(30_000, MIN_POLL_INTERVAL) assert delay == MIN_POLL_INTERVAL assert backoff == MIN_POLL_INTERVAL * 2 diff --git a/tests/tasks/client/test_task_context_validation.py b/tests/tasks/client/test_task_context_validation.py deleted file mode 100644 index 9d5f44e1e..000000000 --- a/tests/tasks/client/test_task_context_validation.py +++ /dev/null @@ -1,224 +0,0 @@ -""" -Tests for Task client context validation. - -Verifies that Task methods properly validate client context and that -cached results remain accessible outside context. -""" - -import pytest - -from fastmcp import FastMCP -from fastmcp.client import Client - -pytestmark = pytest.mark.skip(reason="Phase 4: requires client task support (SEP-2663)") - - -@pytest.fixture -async def task_server(): - """Create a test server with background tasks.""" - mcp = FastMCP("context-test-server") - - @mcp.tool(task=True) - async def background_tool(value: str) -> str: - """Tool that runs in background.""" - return f"Result: {value}" - - @mcp.prompt(task=True) - async def background_prompt(topic: str) -> str: - """Prompt that runs in background.""" - return f"Prompt about {topic}" - - @mcp.resource("file://background.txt", task=True) - async def background_resource() -> str: - """Resource that runs in background.""" - return "Background resource content" - - return mcp - - -async def test_task_status_outside_context_raises(task_server): - """Calling task.status() outside client context raises error.""" - task = None - async with Client(task_server, mode="legacy") as client: - task = await client.call_tool("background_tool", {"value": "test"}, task=True) - assert not task.returned_immediately - # Now outside context - - with pytest.raises(RuntimeError, match="outside client context"): - await task.status() - - -async def test_task_result_outside_context_raises(task_server): - """Calling task.result() outside context raises error.""" - task = None - async with Client(task_server, mode="legacy") as client: - task = await client.call_tool("background_tool", {"value": "test"}, task=True) - assert not task.returned_immediately - # Now outside context - - with pytest.raises(RuntimeError, match="outside client context"): - await task.result() - - -async def test_task_wait_outside_context_raises(task_server): - """Calling task.wait() outside context raises error.""" - task = None - async with Client(task_server, mode="legacy") as client: - task = await client.call_tool("background_tool", {"value": "test"}, task=True) - assert not task.returned_immediately - # Now outside context - - with pytest.raises(RuntimeError, match="outside client context"): - await task.wait() - - -async def test_task_cancel_outside_context_raises(task_server): - """Calling task.cancel() outside context raises error.""" - task = None - async with Client(task_server, mode="legacy") as client: - task = await client.call_tool("background_tool", {"value": "test"}, task=True) - assert not task.returned_immediately - # Now outside context - - with pytest.raises(RuntimeError, match="outside client context"): - await task.cancel() - - -async def test_cached_tool_task_accessible_outside_context(task_server): - """Tool tasks with cached results work outside context.""" - task = None - async with Client(task_server, mode="legacy") as client: - task = await client.call_tool("background_tool", {"value": "test"}, task=True) - assert not task.returned_immediately - - # Get result once to cache it - result1 = await task.result() - assert result1.data == "Result: test" - # Now outside context - - # Should work because result is cached - result2 = await task.result() - assert result2 is result1 # Same object - assert result2.data == "Result: test" - - -@pytest.mark.xfail( - reason="SDK v2 has no `task` field on GetPromptRequestParams / " - "ReadResourceRequestParams; prompt/resource task submission is not " - "wire-expressible and always graceful-degrades (sdk-feedback #3).", - strict=True, -) -async def test_cached_prompt_task_accessible_outside_context(task_server): - """Prompt tasks with cached results work outside context.""" - task = None - async with Client(task_server, mode="legacy") as client: - task = await client.get_prompt( - "background_prompt", {"topic": "test"}, task=True - ) - assert not task.returned_immediately - - # Get result once to cache it - result1 = await task.result() - assert result1.description == "Prompt that runs in background." - # Now outside context - - # Should work because result is cached - result2 = await task.result() - assert result2 is result1 # Same object - assert result2.description == "Prompt that runs in background." - - -@pytest.mark.xfail( - reason="SDK v2 has no `task` field on GetPromptRequestParams / " - "ReadResourceRequestParams; prompt/resource task submission is not " - "wire-expressible and always graceful-degrades (sdk-feedback #3).", - strict=True, -) -async def test_cached_resource_task_accessible_outside_context(task_server): - """Resource tasks with cached results work outside context.""" - task = None - async with Client(task_server, mode="legacy") as client: - task = await client.read_resource("file://background.txt", task=True) - assert not task.returned_immediately - - # Get result once to cache it - result1 = await task.result() - assert len(result1) > 0 - # Now outside context - - # Should work because result is cached - result2 = await task.result() - assert result2 is result1 # Same object - - -async def test_uncached_status_outside_context_raises(task_server): - """Even after caching result, status() still requires client context.""" - task = None - async with Client(task_server, mode="legacy") as client: - task = await client.call_tool("background_tool", {"value": "test"}, task=True) - assert not task.returned_immediately - - # Cache the result - await task.result() - # Now outside context - - # result() works (cached) - result = await task.result() - assert result.data == "Result: test" - - # But status() still needs client connection - with pytest.raises(RuntimeError, match="outside client context"): - await task.status() - - -async def test_task_await_syntax_outside_context_raises(task_server): - """Using await task syntax outside context raises error for background tasks.""" - task = None - async with Client(task_server, mode="legacy") as client: - task = await client.call_tool("background_tool", {"value": "test"}, task=True) - assert not task.returned_immediately - # Now outside context - - with pytest.raises(RuntimeError, match="outside client context"): - await task # Same as await task.result() - - -async def test_task_await_syntax_works_for_cached_results(task_server): - """Using await task syntax works outside context when result is cached.""" - task = None - async with Client(task_server, mode="legacy") as client: - task = await client.call_tool("background_tool", {"value": "test"}, task=True) - result1 = await task # Cache it - # Now outside context - - result2 = await task # Should work (cached) - assert result2 is result1 - assert result2.data == "Result: test" - - -async def test_multiple_result_calls_return_same_cached_object(task_server): - """Multiple result() calls return the same cached object.""" - async with Client(task_server, mode="legacy") as client: - task = await client.call_tool("background_tool", {"value": "test"}, task=True) - - result1 = await task.result() - result2 = await task.result() - result3 = await task.result() - - # Should all be the same object (cached) - assert result1 is result2 - assert result2 is result3 - - -async def test_background_task_properties_accessible_outside_context(task_server): - """Background task properties like task_id accessible outside context.""" - task = None - async with Client(task_server, mode="legacy") as client: - task = await client.call_tool("background_tool", {"value": "test"}, task=True) - task_id_inside = task.task_id - assert not task.returned_immediately - # Now outside context - - # Properties should still be accessible (they don't need client connection) - assert task.task_id == task_id_inside - assert task.returned_immediately is False diff --git a/tests/tasks/client/test_task_result_caching.py b/tests/tasks/client/test_task_result_caching.py deleted file mode 100644 index f7670cc6e..000000000 --- a/tests/tasks/client/test_task_result_caching.py +++ /dev/null @@ -1,341 +0,0 @@ -""" -Tests for Task result caching behavior. - -Verifies that Task.result() and await task cache results properly to avoid -redundant server calls and ensure consistent object identity. -""" - -import pytest - -from fastmcp import FastMCP -from fastmcp.client import Client - -pytestmark = pytest.mark.skip(reason="Phase 4: requires client task support (SEP-2663)") - - -async def test_tool_task_result_cached_on_first_call(): - """First call caches result, subsequent calls return cached value.""" - call_count = 0 - mcp = FastMCP("test") - - @mcp.tool(task=True) - async def counting_tool() -> int: - nonlocal call_count - call_count += 1 - return call_count - - async with Client(mcp, mode="legacy") as client: - task = await client.call_tool("counting_tool", task=True) - - result1 = await task.result() - result2 = await task.result() - result3 = await task.result() - - # All should return 1 (first execution value) - assert result1.data == 1 - assert result2.data == 1 - assert result3.data == 1 - - # Verify they're the same object (cached) - assert result1 is result2 is result3 - - -async def test_prompt_task_result_cached(): - """PromptTask caches results on first call.""" - call_count = 0 - mcp = FastMCP("test") - - @mcp.prompt(task=True) - async def counting_prompt() -> str: - nonlocal call_count - call_count += 1 - return f"Call number: {call_count}" - - async with Client(mcp, mode="legacy") as client: - task = await client.get_prompt("counting_prompt", task=True) - - result1 = await task.result() - result2 = await task.result() - result3 = await task.result() - - # All should return same content - assert result1.messages[0].content.text == "Call number: 1" - assert result2.messages[0].content.text == "Call number: 1" - assert result3.messages[0].content.text == "Call number: 1" - - # Verify they're the same object (cached) - assert result1 is result2 is result3 - - -async def test_resource_task_result_cached(): - """ResourceTask caches results on first call.""" - call_count = 0 - mcp = FastMCP("test") - - @mcp.resource("file://counter.txt", task=True) - async def counting_resource() -> str: - nonlocal call_count - call_count += 1 - return f"Count: {call_count}" - - async with Client(mcp, mode="legacy") as client: - task = await client.read_resource("file://counter.txt", task=True) - - result1 = await task.result() - result2 = await task.result() - result3 = await task.result() - - # All should return same content - assert result1[0].text == "Count: 1" - assert result2[0].text == "Count: 1" - assert result3[0].text == "Count: 1" - - # Verify they're the same object (cached) - assert result1 is result2 is result3 - - -async def test_multiple_await_returns_same_object(): - """Multiple await task calls return identical object.""" - mcp = FastMCP("test") - - @mcp.tool(task=True) - async def sample_tool() -> str: - return "result" - - async with Client(mcp, mode="legacy") as client: - task = await client.call_tool("sample_tool", task=True) - - result1 = await task - result2 = await task - result3 = await task - - # Should be exact same object in memory - assert result1 is result2 is result3 - assert id(result1) == id(result2) == id(result3) - - -async def test_result_and_await_share_cache(): - """task.result() and await task share the same cache.""" - mcp = FastMCP("test") - - @mcp.tool(task=True) - async def sample_tool() -> str: - return "cached" - - async with Client(mcp, mode="legacy") as client: - task = await client.call_tool("sample_tool", task=True) - - # Call result() first - result_via_method = await task.result() - - # Then await directly - result_via_await = await task - - # Should be the same cached object - assert result_via_method is result_via_await - assert id(result_via_method) == id(result_via_await) - - -async def test_forbidden_mode_tool_caches_error_result(): - """Tools with task=False (mode=forbidden) cache error results.""" - mcp = FastMCP("test") - - @mcp.tool(task=False) - async def non_task_tool() -> int: - return 1 - - async with Client(mcp, mode="legacy") as client: - # Request as task, but mode="forbidden" will reject with error - task = await client.call_tool("non_task_tool", task=True, raise_on_error=False) - - # Should be immediate (error returned immediately) - assert task.returned_immediately - - result1 = await task.result() - result2 = await task.result() - result3 = await task.result() - - # All should return cached error - assert result1.is_error - assert "does not support task-augmented execution" in str(result1) - - # Verify they're the same object (cached) - assert result1 is result2 is result3 - - -@pytest.mark.xfail( - reason="SDK v2 has no `task` field on GetPromptRequestParams / " - "ReadResourceRequestParams; prompt/resource task submission is not " - "wire-expressible and always graceful-degrades (sdk-feedback #3).", - strict=True, -) -async def test_forbidden_mode_prompt_raises_error(): - """Prompts with task=False (mode=forbidden) raise error.""" - import pytest - from mcp.shared.exceptions import MCPError - - mcp = FastMCP("test") - - @mcp.prompt(task=False) - async def non_task_prompt() -> str: - return "Immediate" - - async with Client(mcp, mode="legacy") as client: - # Prompts with mode="forbidden" raise MCPError when called with task=True - with pytest.raises(MCPError): - await client.get_prompt("non_task_prompt", task=True) - - -@pytest.mark.xfail( - reason="SDK v2 has no `task` field on GetPromptRequestParams / " - "ReadResourceRequestParams; prompt/resource task submission is not " - "wire-expressible and always graceful-degrades (sdk-feedback #3).", - strict=True, -) -async def test_forbidden_mode_resource_raises_error(): - """Resources with task=False (mode=forbidden) raise error.""" - import pytest - from mcp.shared.exceptions import MCPError - - mcp = FastMCP("test") - - @mcp.resource("file://immediate.txt", task=False) - async def non_task_resource() -> str: - return "Immediate" - - async with Client(mcp, mode="legacy") as client: - # Resources with mode="forbidden" raise MCPError when called with task=True - with pytest.raises(MCPError): - await client.read_resource("file://immediate.txt", task=True) - - -async def test_immediate_task_caches_result(): - """Immediate tasks (optional mode called without background) cache results.""" - call_count = 0 - mcp = FastMCP("test", tasks=True) - - # Tool with task=True (optional mode) - but without docket will execute immediately - @mcp.tool(task=True) - async def task_tool() -> int: - nonlocal call_count - call_count += 1 - return call_count - - async with Client(mcp, mode="legacy") as client: - # Call with task=True - task = await client.call_tool("task_tool", task=True) - - # Get result multiple times - result1 = await task.result() - result2 = await task.result() - result3 = await task.result() - - # All should return cached value - assert result1.data == 1 - assert result2.data == 1 - assert result3.data == 1 - - # Verify they're the same object (cached) - assert result1 is result2 is result3 - - -async def test_cache_persists_across_mixed_access_patterns(): - """Cache works correctly when mixing result() and await.""" - mcp = FastMCP("test") - - @mcp.tool(task=True) - async def mixed_tool() -> str: - return "mixed" - - async with Client(mcp, mode="legacy") as client: - task = await client.call_tool("mixed_tool", task=True) - - # Access in various orders - result1 = await task - result2 = await task.result() - result3 = await task - result4 = await task.result() - - # All should be the same cached object - assert result1 is result2 is result3 is result4 - - -async def test_different_tasks_have_separate_caches(): - """Different task instances maintain separate caches.""" - mcp = FastMCP("test") - - @mcp.tool(task=True) - async def separate_tool(value: str) -> str: - return f"Result: {value}" - - async with Client(mcp, mode="legacy") as client: - task1 = await client.call_tool("separate_tool", {"value": "A"}, task=True) - task2 = await client.call_tool("separate_tool", {"value": "B"}, task=True) - - result1 = await task1.result() - result2 = await task2.result() - - # Different results - assert result1.data == "Result: A" - assert result2.data == "Result: B" - - # Not the same object - assert result1 is not result2 - - # But each task's cache works independently - result1_again = await task1.result() - result2_again = await task2.result() - - assert result1 is result1_again - assert result2 is result2_again - - -async def test_cache_survives_status_checks(): - """Calling status() doesn't affect result caching.""" - mcp = FastMCP("test") - - @mcp.tool(task=True) - async def status_check_tool() -> str: - return "status" - - async with Client(mcp, mode="legacy") as client: - task = await client.call_tool("status_check_tool", task=True) - - # Check status multiple times - await task.status() - await task.status() - - result1 = await task.result() - - # Check status again - await task.status() - - result2 = await task.result() - - # Cache should still work - assert result1 is result2 - - -async def test_cache_survives_wait_calls(): - """Calling wait() doesn't affect result caching.""" - mcp = FastMCP("test") - - @mcp.tool(task=True) - async def wait_test_tool() -> str: - return "waited" - - async with Client(mcp, mode="legacy") as client: - task = await client.call_tool("wait_test_tool", task=True) - - # Wait for completion - await task.wait() - - result1 = await task.result() - - # Wait again (no-op since completed) - await task.wait() - - result2 = await task.result() - - # Cache should still work - assert result1 is result2 diff --git a/tests/tasks/client/test_transparent_tasks.py b/tests/tasks/client/test_transparent_tasks.py new file mode 100644 index 000000000..a0b95301e --- /dev/null +++ b/tests/tasks/client/test_transparent_tasks.py @@ -0,0 +1,158 @@ +"""The transparent client task flow over a real in-memory connection. + +A real `Client(server, mode="auto")` calls a `task=True` tool; the server runs it +as a task and answers `tools/call` with a `CreateTaskResult`; the client's +auto-registered tasks extension resolves it by polling `tasks/get` to completion. +The caller of `call_tool` sees only the tool's real result β€” never that the call +was tasked. This is the whole point of the client half. +""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass + +import mcp_types +import pytest + +from fastmcp import Context, FastMCP +from fastmcp.client import Client +from fastmcp.exceptions import ToolError +from fastmcp_tasks import TasksExtension, call_tool_task + + +@pytest.fixture +def task_server() -> FastMCP: + mcp = FastMCP("transparent-tasks") + mcp.add_extension(TasksExtension()) + + @mcp.tool(task=True) + async def multiply(a: int, b: int) -> int: + await asyncio.sleep(0.01) + return a * b + + @mcp.tool(task=True) + async def boom() -> str: + raise ValueError("kaboom") + + return mcp + + +async def test_call_tool_transparently_completes_a_task(task_server: FastMCP): + """call_tool returns the tool's real result; the caller never sees a task.""" + async with Client(task_server, mode="auto") as client: + result = await client.call_tool("multiply", {"a": 6, "b": 7}) + + assert result.data == 42 + + +async def test_call_tool_mcp_returns_completed_result(task_server: FastMCP): + """call_tool_mcp resolves the tasked call into an ordinary CallToolResult.""" + async with Client(task_server, mode="auto") as client: + result = await client.call_tool_mcp("multiply", {"a": 3, "b": 4}) + + assert result.structured_content == {"result": 12} + assert not result.is_error + + +async def test_failed_task_raises_tool_error(task_server: FastMCP): + """A task whose tool raises surfaces as a ToolError through call_tool.""" + async with Client(task_server, mode="auto") as client: + with pytest.raises(ToolError, match="kaboom"): + await client.call_tool("boom", {}) + + +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: + task = await call_tool_task(client, "multiply", {"a": 2, "b": 5}) + # The raw claimed shape is exposed on the handle. + assert task.create_result.result_type == "task" + assert task.create_result.status == "working" + assert isinstance(task.task_id, str) and task.task_id + + result = await task.result() + assert result.data == 10 + + +async def test_legacy_client_never_tasks(task_server: FastMCP): + """A legacy-era client never negotiates the capability, so nothing is tasked. + + The optional-mode tool simply runs synchronously and returns its result + directly (no CreateTaskResult on the wire). + """ + async with Client(task_server, mode="legacy") as client: + result = await client.call_tool("multiply", {"a": 8, "b": 9}) + + assert result.data == 72 + + +# --- In-task input over the wire ------------------------------------------- + + +@dataclass +class DinnerPrefs: + cuisine: str + vegetarian: bool + + +def _elicit_request(message: str) -> mcp_types.ElicitRequest: + return mcp_types.ElicitRequest( + params=mcp_types.ElicitRequestFormParams( + message=message, + requested_schema={ + "type": "object", + "properties": { + "cuisine": {"type": "string"}, + "vegetarian": {"type": "boolean"}, + }, + "required": ["cuisine", "vegetarian"], + }, + ) + ) + + +@pytest.fixture +def guard_server() -> FastMCP: + mcp = FastMCP("guard-tasks") + mcp.add_extension(TasksExtension()) + + @mcp.tool(task=True) + async def plan_dinner( + ctx: Context, + ) -> str | mcp_types.InputRequiredResult: + responses = ctx.input_responses + if responses is None: + return mcp_types.InputRequiredResult( + result_type="input_required", + input_requests={"prefs": _elicit_request("What's for dinner?")}, + ) + answer = responses["prefs"] + assert isinstance(answer, mcp_types.ElicitResult) + assert answer.content is not None + veg = "vegetarian " if answer.content["vegetarian"] else "" + return f"Tonight: a {veg}{answer.content['cuisine']} dinner!" + + return mcp + + +async def test_in_task_input_answered_transparently(guard_server: FastMCP): + """A guard task that asks for input is answered via the elicitation handler.""" + + async def handle_elicitation(message, response_type, params, context): + return DinnerPrefs(cuisine="Thai", vegetarian=True) + + client = Client( + guard_server, mode="auto", elicitation_handler=handle_elicitation + ) + async with client: + result = await client.call_tool("plan_dinner", {}) + + assert result.data == "Tonight: a vegetarian Thai dinner!" + + +async def test_in_task_input_without_handler_errors(guard_server: FastMCP): + """A guard task with no elicitation handler surfaces a clear error.""" + async with Client(guard_server, mode="auto") as client: + with pytest.raises(ToolError, match="no elicitation handler"): + await client.call_tool("plan_dinner", {})