prefecthq--fastmcp/tests/server/tasks/test_task_tools.py

"""
Tests for server-side tool task behavior.

Tests tool-specific task handling, parallel to test_task_prompts.py
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
async def tool_server():
    """Create a FastMCP server with task-enabled tools."""
    mcp = FastMCP("tool-task-server")

    @mcp.tool(task=True)
    async def simple_tool(message: str) -> str:
        """A simple tool for testing."""
        return f"Processed: {message}"

    @mcp.tool(task=False)
    async def sync_only_tool(message: str) -> str:
        """Tool with task=False."""
        return f"Sync: {message}"

    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.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


async def test_task_submission_honors_strict_input_validation():
    """Strict input validation applies to task submissions, not just sync calls.

    With ``strict_input_validation=True``, a lax coercion like ``{"n": "1"}``
    for an ``int`` parameter is rejected on the synchronous path. The task
    submission path must reject it identically rather than silently coercing
    and queueing it — otherwise ``task=True`` would bypass the strict flag.
    """

    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.method)

    server = FastMCP("strict-task-server", strict_input_validation=True)

    @server.tool(task=True)
    async def square(n: int) -> int:
        return n * n

    recorder = _Recorder()
    async with Client(server, message_handler=recorder) as client:
        # Sync path rejects the string-for-int coercion under strict validation.
        with pytest.raises(ToolError):
            await client.call_tool("square", {"n": "1"})

        # Task path must reject it too, before any task state is created — so no
        # status notification is emitted for the orphaned submission.
        task = await client.call_tool("square", {"n": "1"}, task=True)
        assert task.returned_immediately
        with pytest.raises(ToolError):
            await task.result()

    assert "notifications/tasks/status" not in recorder.methods


async def test_task_submission_valid_argument_under_strict_validation():
    """A well-typed argument still submits fine when strict validation is on."""
    server = FastMCP("strict-task-valid-server", strict_input_validation=True)

    @server.tool(task=True)
    async def square(n: int) -> int:
        return n * n

    async with Client(server) as client:
        task = await client.call_tool("square", {"n": 4}, task=True)
        assert not task.returned_immediately
        result = await task.result()
        assert result.data == 16


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:
        # Regular call without task metadata
        result = await client.call_tool("simple_tool", {"message": "hello"})

        # Should execute immediately and return result
        assert "Processed: hello" in str(result)


async def test_tool_with_task_metadata_returns_immediately(tool_server):
    """Tools with task metadata return immediately with ToolTask object."""
    async with Client(tool_server) as client:
        # Call with task metadata
        task = await client.call_tool("simple_tool", {"message": "test"}, task=True)
        assert task
        assert not task.returned_immediately

        assert isinstance(task, ToolTask)
        assert isinstance(task.task_id, str)
        assert len(task.task_id) > 0


async def test_tool_task_executes_in_background(tool_server):
    """Tool task is submitted to Docket and executes in background."""
    execution_started = asyncio.Event()
    execution_completed = asyncio.Event()

    @tool_server.tool(task=True)
    async def coordinated_tool() -> str:
        """Tool with coordination points."""
        execution_started.set()
        await execution_completed.wait()
        return "completed"

    async with Client(tool_server) as client:
        task = await client.call_tool("coordinated_tool", task=True)
        assert task
        assert not task.returned_immediately

        # Wait for execution to start
        await asyncio.wait_for(execution_started.wait(), timeout=2.0)

        # Task should still be working
        status = await task.status()
        assert status.status in ["working"]

        # Signal completion
        execution_completed.set()
        await task.wait(timeout=2.0)

        result = await task.result()
        assert result.data == "completed"


async def test_forbidden_mode_tool_rejects_task_calls(tool_server):
    """Tools with task=False (mode=forbidden) reject task-augmented calls."""
    async with Client(tool_server) as client:
        # Calling with task=True when task=False should return error
        task = await client.call_tool(
            "sync_only_tool", {"message": "test"}, task=True, raise_on_error=False
        )
        assert task
        assert task.returned_immediately

        result = await task.result()
        # New behavior: mode="forbidden" returns an error
        assert result.is_error
        assert "does not support task-augmented execution" in str(result)