From 0cffe411156bf5f9e08612595ca853c9d4a4d2e4 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Sat, 27 Jun 2026 10:03:59 -0400 Subject: [PATCH] fix: validate task tool arguments against declared types (#4373) --- fastmcp_slim/fastmcp/server/tasks/handlers.py | 7 ++ fastmcp_slim/fastmcp/tools/function_tool.py | 60 +++++++++++++ fastmcp_slim/fastmcp/utilities/components.py | 13 +++ tests/server/tasks/test_task_tools.py | 86 +++++++++++++++++++ 4 files changed, 166 insertions(+) diff --git a/fastmcp_slim/fastmcp/server/tasks/handlers.py b/fastmcp_slim/fastmcp/server/tasks/handlers.py index aa9d73bdc..b427a0dfe 100644 --- a/fastmcp_slim/fastmcp/server/tasks/handlers.py +++ b/fastmcp_slim/fastmcp/server/tasks/handlers.py @@ -66,6 +66,13 @@ async def submit_to_docket( Returns: CreateTaskResult: Task stub with proper Task object """ + # Validate and coerce arguments before creating any task state. A failure + # here must surface before the Redis metadata and initial "working" + # notification below are written, otherwise an invalid input would orphan a + # task the client has already observed (#4349). + if arguments is not None: + arguments = component.coerce_task_arguments(arguments) + # Generate server-side task ID per SEP-1686 final spec (line 375-377) # Server MUST generate task IDs, clients no longer provide them server_task_id = str(uuid.uuid4()) diff --git a/fastmcp_slim/fastmcp/tools/function_tool.py b/fastmcp_slim/fastmcp/tools/function_tool.py index 18466a0f7..a882edc4e 100644 --- a/fastmcp_slim/fastmcp/tools/function_tool.py +++ b/fastmcp_slim/fastmcp/tools/function_tool.py @@ -2,6 +2,7 @@ from __future__ import annotations +import functools import inspect import warnings from collections.abc import Callable @@ -15,6 +16,7 @@ from typing import ( Protocol, TypeVar, cast, + get_type_hints, overload, runtime_checkable, ) @@ -89,6 +91,31 @@ class ToolMeta: run_in_thread: bool = True +def _resolve_param_hints(fn: Callable[..., Any]) -> dict[str, Any]: + """Resolve a callable's parameter type hints, tolerating partials. + + ``get_type_hints`` rejects ``functools.partial`` objects (and other + non-function callables), which the synchronous TypeAdapter path handles + natively. For those, resolve hints against the underlying function and keep + only the parameters that remain in the partially-bound signature. + """ + try: + return get_type_hints(fn, include_extras=True) + except TypeError: + target = fn + while isinstance(target, functools.partial): + target = target.func + try: + resolved = get_type_hints(target, include_extras=True) + except TypeError: + return {} + return { + name: resolved[name] + for name in inspect.signature(fn).parameters + if name in resolved + } + + class FunctionTool(Tool): fn: SkipJsonSchema[Callable[..., Any]] return_type: Annotated[SkipJsonSchema[Any], Field(exclude=True)] = None @@ -393,6 +420,39 @@ class FunctionTool(Tool): kwargs["key"] = task_key return await docket.add(lookup_key, **kwargs)(**arguments) + def coerce_task_arguments(self, arguments: dict[str, Any]) -> dict[str, Any]: + """Validate client arguments against their declared parameter types. + + The synchronous ``run()`` path validates arguments through the + function's Pydantic TypeAdapter, so a parameter typed as a model + arrives as a model instance. The task path hands the raw arguments to + Docket, which binds them to the function signature without coercion — + so without this a model-typed parameter would reach the function as a + raw dict (#4349). ``submit_to_docket`` calls this up front so coerced + values are what get queued, and validation errors surface before any + task state is created. Coerced values survive the trip to the worker + because Docket serializes task arguments with cloudpickle. + + Injected dependency parameters (Context, Depends()) are excluded via + the same wrapper used by the synchronous path, so only client-supplied + arguments are coerced and Docket's dependency resolution is untouched. + """ + from fastmcp.server.dependencies import without_injected_parameters + + wrapper_fn = without_injected_parameters( + self.fn, run_in_thread=self.run_in_thread + ) + hints = _resolve_param_hints(wrapper_fn) + + coerced = dict(arguments) + for name, value in arguments.items(): + annotation = hints.get(name) + if annotation is None: + continue + adapter = get_cached_typeadapter(annotation) + coerced[name] = adapter.validate_python(value) + return coerced + @overload def tool(fn: F) -> F: ... diff --git a/fastmcp_slim/fastmcp/utilities/components.py b/fastmcp_slim/fastmcp/utilities/components.py index e481aac9e..d371ca495 100644 --- a/fastmcp_slim/fastmcp/utilities/components.py +++ b/fastmcp_slim/fastmcp/utilities/components.py @@ -232,6 +232,19 @@ class FastMCPComponent(FastMCPBaseModel): """ # Base implementation: no-op (subclasses override) + def coerce_task_arguments(self, arguments: dict[str, Any]) -> dict[str, Any]: + """Validate and coerce task arguments before any task state is created. + + Called by ``submit_to_docket`` up front, so invalid inputs raise before + the task's Redis metadata and initial status notification exist — + otherwise a coercion failure during queueing would orphan a task the + client has already observed. The base implementation is a no-op; + components that splat arguments into a typed Python callable (e.g. + ``FunctionTool``) override this to mirror the synchronous validation + path. + """ + return arguments + async def add_to_docket( self, docket: Docket, *args: Any, **kwargs: Any ) -> Execution: diff --git a/tests/server/tasks/test_task_tools.py b/tests/server/tasks/test_task_tools.py index 282c2b46c..26fa1483c 100644 --- a/tests/server/tasks/test_task_tools.py +++ b/tests/server/tasks/test_task_tools.py @@ -6,12 +6,18 @@ and test_task_resources.py. """ import asyncio +import functools +import mcp.types import pytest +from pydantic import BaseModel from fastmcp import FastMCP from fastmcp.client import Client +from fastmcp.client.messages import MessageHandler from fastmcp.client.tasks import ToolTask +from fastmcp.exceptions import ToolError +from fastmcp.tools.function_tool import _resolve_param_hints @pytest.fixture @@ -32,6 +38,86 @@ async def tool_server(): return mcp +class _Item(BaseModel): + value: str + + +async def test_task_tool_validates_model_arguments(): + """Model-typed args are coerced to model instances for task calls (#4349). + + The synchronous path validates arguments through the function's + TypeAdapter, so a parameter typed as a Pydantic model arrives as a model + instance. The task path must coerce the same way rather than passing the + raw dict through to the function. + """ + mcp = FastMCP("tool-task-validation-server") + + @mcp.tool(task=True) + async def inspect_items(item: _Item, items: list[_Item]) -> dict[str, str]: + return {"item": type(item).__name__, "element": type(items[0]).__name__} + + arguments = {"item": {"value": "a"}, "items": [{"value": "b"}]} + expected = {"item": "_Item", "element": "_Item"} + + async with Client(mcp) as client: + sync_result = await client.call_tool("inspect_items", arguments) + task = await client.call_tool("inspect_items", arguments, task=True) + task_result = await task.result() + + assert sync_result.data == expected + assert task_result.data == expected + + +async def test_task_tool_invalid_arguments_fail_before_task_state(): + """Invalid task arguments are rejected before any task state is created. + + Coercion runs up front in submit_to_docket, so a validation failure surfaces + before the task's Redis metadata and initial "working" status notification + are written. Otherwise an invalid input would orphan a task the client had + already observed via that notification. + """ + + class _Recorder(MessageHandler): + def __init__(self): + super().__init__() + self.methods: list[str] = [] + + async def on_notification(self, message: mcp.types.ServerNotification) -> None: + self.methods.append(message.root.method) + + server = FastMCP("tool-task-invalid-args-server") + + @server.tool(task=True) + async def needs_item(item: _Item) -> str: + return item.value + + recorder = _Recorder() + async with Client(server, message_handler=recorder) as client: + # `item` is missing its required `value` field. + task = await client.call_tool("needs_item", {"item": {}}, task=True) + assert task.returned_immediately + with pytest.raises(ToolError): + await task.result() + + assert "notifications/tasks/status" not in recorder.methods + + +def test_resolve_param_hints_handles_partials(): + """Partials aren't introspectable by get_type_hints; resolve via the func. + + Argument coercion must not raise for partial-wrapped callables — it should + resolve hints for the still-unbound parameters. + """ + + async def base(prefix: str, items: list[_Item]) -> str: + return prefix + + partial_fn = functools.partial(base, "bound") + hints = _resolve_param_hints(partial_fn) + + assert hints["items"] == list[_Item] + + async def test_synchronous_tool_call_unchanged(tool_server): """Tools without task metadata execute synchronously as before.""" async with Client(tool_server) as client: