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

"""
Tests for task protocol methods.

Tests the tasks/get, tasks/result, and tasks/list JSON-RPC protocol methods.
"""

import asyncio

import pytest
from mcp.shared.exceptions import MCPError

from fastmcp import FastMCP
from fastmcp.client import Client


@pytest.fixture
async def endpoint_server():
    """Create a server with background tasks and HTTP transport."""
    mcp = FastMCP("endpoint-test-server")

    @mcp.tool(task=True)  # Enable background execution
    async def quick_tool(value: int) -> int:
        """Returns the value immediately."""
        return value * 2

    @mcp.tool(task=True)  # Enable background execution
    async def error_tool() -> str:
        """Always raises an error."""
        raise RuntimeError("Task failed!")

    @mcp.tool(task=True)  # Enable background execution
    async def slow_tool() -> str:
        """A slow tool for testing cancellation."""
        await asyncio.sleep(10)
        return "done"

    return mcp


async def test_tasks_get_endpoint_returns_status(endpoint_server):
    """POST /tasks/get returns task status."""
    async with Client(endpoint_server) as client:
        # Submit a task
        task = await client.call_tool("quick_tool", {"value": 21}, task=True)

        # Check status immediately - should be submitted or working
        status = await task.status()
        assert status.task_id == task.task_id
        assert status.status in ["working", "completed"]

        # Wait for completion
        await task.wait(timeout=2.0)

        # Check again - should be completed
        status = await task.status()
        assert status.status == "completed"


async def test_tasks_get_endpoint_includes_poll_interval(endpoint_server):
    """Task status includes pollFrequency hint."""
    async with Client(endpoint_server) as client:
        task = await client.call_tool("quick_tool", {"value": 42}, task=True)

        status = await task.status()
        assert status.poll_interval is not None
        assert isinstance(status.poll_interval, int)


async def test_tasks_result_endpoint_returns_result_when_completed(endpoint_server):
    """POST /tasks/result returns the tool result when completed."""
    async with Client(endpoint_server) as client:
        task = await client.call_tool("quick_tool", {"value": 21}, task=True)

        # Wait for completion and get result
        result = await task.result()
        assert result.data == 42  # 21 * 2


async def test_tasks_result_endpoint_errors_if_not_completed(endpoint_server):
    """POST /tasks/result returns error if task not completed yet."""
    # Create a task that won't complete until signaled
    completion_signal = asyncio.Event()

    @endpoint_server.tool(task=True)  # Enable background execution
    async def blocked_tool() -> str:
        await completion_signal.wait()
        return "done"

    async with Client(endpoint_server) as client:
        task = await client.call_tool("blocked_tool", task=True)

        # Try to get result immediately (task still running)
        with pytest.raises(Exception):  # Should raise or return error
            await client.get_task_result(task.task_id)

        # Cleanup - signal completion
        completion_signal.set()


async def test_tasks_result_endpoint_errors_if_task_not_found(endpoint_server):
    """POST /tasks/result returns error for non-existent task."""
    async with Client(endpoint_server) as client:
        # Try to get result for non-existent task
        with pytest.raises(Exception):
            await client.get_task_result("non-existent-task-id")


async def test_tasks_result_endpoint_returns_error_for_failed_task(endpoint_server):
    """POST /tasks/result returns error information for failed tasks."""
    async with Client(endpoint_server) as client:
        task = await client.call_tool("error_tool", task=True)

        # Wait for task to fail
        await task.wait(state="failed", timeout=2.0)

        # Getting result should raise or return error info
        with pytest.raises(Exception) as exc_info:
            await task.result()

        assert (
            "failed" in str(exc_info.value).lower()
            or "error" in str(exc_info.value).lower()
        )


async def test_tasks_list_endpoint_session_isolation(endpoint_server):
    """list_tasks returns only tasks submitted by this client."""
    # Since client tracks tasks locally, this tests client-side tracking
    async with Client(endpoint_server) as client:
        # Submit multiple tasks (server generates IDs)
        tasks = []
        for i in range(3):
            task = await client.call_tool("quick_tool", {"value": i}, task=True)
            tasks.append(task)

        # Wait for all to complete
        for task in tasks:
            await task.wait(timeout=2.0)

        # List tasks - should see all 3
        response = await client.list_tasks()
        returned_ids = [t["taskId"] for t in response["tasks"]]
        task_ids = [t.task_id for t in tasks]
        assert len(returned_ids) == 3
        assert all(tid in task_ids for tid in returned_ids)


async def test_get_status_nonexistent_task_raises_error(endpoint_server):
    """Getting status for nonexistent task raises MCP error (per SEP-1686 SDK behavior)."""
    async with Client(endpoint_server) as client:
        # Try to get status for task that was never created
        # Per SDK implementation: raises ValueError which becomes JSON-RPC error
        with pytest.raises(MCPError, match="Task nonexistent-task-id not found"):
            await client.get_task_status("nonexistent-task-id")


async def test_task_cancellation_workflow(endpoint_server):
    """Task can be cancelled, transitioning to cancelled state."""
    async with Client(endpoint_server) as client:
        # Submit slow task
        task = await client.call_tool("slow_tool", {}, task=True)

        # Give it a moment to start
        await asyncio.sleep(0.1)

        # Cancel the task
        await task.cancel()

        # Give cancellation a moment to process
        await asyncio.sleep(0.1)

        # Task should be in cancelled state
        status = await task.status()
        assert status.status == "cancelled"


@pytest.mark.timeout(10)
async def test_task_cancellation_interrupts_running_coroutine(endpoint_server):
    """Task cancellation actually interrupts the running coroutine.

    This verifies that when a task is cancelled, the underlying asyncio
    coroutine receives CancelledError rather than continuing to completion.
    Requires pydocket >= 0.16.2.

    See: https://github.com/PrefectHQ/fastmcp/issues/2679
    """
    started = asyncio.Event()
    was_interrupted = asyncio.Event()
    completed_normally = asyncio.Event()

    @endpoint_server.tool(task=True)
    async def interruptible_tool() -> str:
        started.set()
        try:
            await asyncio.sleep(60)
            completed_normally.set()
            return "completed"
        except asyncio.CancelledError:
            was_interrupted.set()
            raise

    async with Client(endpoint_server) as client:
        task = await client.call_tool("interruptible_tool", {}, task=True)

        # Wait for the tool to actually start executing
        await asyncio.wait_for(started.wait(), timeout=5.0)

        # Cancel the task
        await task.cancel()

        # Wait for cancellation to propagate
        await asyncio.wait_for(was_interrupted.wait(), timeout=5.0)

        # The coroutine should have been interrupted, not completed normally
        assert was_interrupted.is_set(), "Task was not interrupted by cancellation"
        assert not completed_normally.is_set(), (
            "Task completed instead of being cancelled"
        )