test_run_task_flow.py

"""
Tests for the simplified task API: enable_tasks() + run_task()

This tests the recommended user flow:
1. server.experimental.enable_tasks() - one-line setup
2. ctx.experimental.run_task(work) - spawns work, returns CreateTaskResult
3. work function uses ServerTaskContext for elicit/create_message

These are integration tests that verify the complete flow works end-to-end.
"""

from typing import Any
from unittest.mock import Mock

import anyio
import pytest
from anyio import Event

from mcp.client.session import ClientSession
from mcp.server import Server
from mcp.server.experimental.request_context import Experimental
from mcp.server.experimental.task_context import ServerTaskContext
from mcp.server.experimental.task_support import TaskSupport
from mcp.server.lowlevel import NotificationOptions
from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
from mcp.shared.experimental.tasks.message_queue import InMemoryTaskMessageQueue
from mcp.shared.message import SessionMessage
from mcp.types import (
    TASK_REQUIRED,
    CallToolResult,
    CancelTaskRequest,
    CancelTaskResult,
    CreateTaskResult,
    GetTaskPayloadRequest,
    GetTaskPayloadResult,
    GetTaskRequest,
    GetTaskResult,
    ListTasksRequest,
    ListTasksResult,
    TextContent,
    Tool,
    ToolExecution,
)


@pytest.mark.anyio
async def test_run_task_basic_flow() -> None:
    """
    Test the basic run_task flow without elicitation.

    1. enable_tasks() sets up handlers
    2. Client calls tool with task field
    3. run_task() spawns work, returns CreateTaskResult
    4. Work completes in background
    5. Client polls and sees completed status
    """
    server = Server("test-run-task")

    # One-line setup
    server.experimental.enable_tasks()

    # Track when work completes and capture received meta
    work_completed = Event()
    received_meta: list[str | None] = [None]

    @server.list_tools()
    async def list_tools() -> list[Tool]:
        return [
            Tool(
                name="simple_task",
                description="A simple task",
                inputSchema={"type": "object", "properties": {"input": {"type": "string"}}},
                execution=ToolExecution(taskSupport=TASK_REQUIRED),
            )
        ]

    @server.call_tool()
    async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult | CreateTaskResult:
        ctx = server.request_context
        ctx.experimental.validate_task_mode(TASK_REQUIRED)

        # Capture the meta from the request (if present)
        if ctx.meta is not None and ctx.meta.model_extra:  # pragma: no branch
            received_meta[0] = ctx.meta.model_extra.get("custom_field")

        async def work(task: ServerTaskContext) -> CallToolResult:
            await task.update_status("Working...")
            input_val = arguments.get("input", "default")
            result = CallToolResult(content=[TextContent(type="text", text=f"Processed: {input_val}")])
            work_completed.set()
            return result

        return await ctx.experimental.run_task(work)

    # Set up streams
    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)

    async def run_server() -> None:
        await server.run(
            client_to_server_receive,
            server_to_client_send,
            server.create_initialization_options(
                notification_options=NotificationOptions(),
                experimental_capabilities={},
            ),
        )

    async def run_client() -> None:
        async with ClientSession(server_to_client_receive, client_to_server_send) as client_session:
            # Initialize
            await client_session.initialize()

            # Call tool as task (with meta to test that code path)
            result = await client_session.experimental.call_tool_as_task(
                "simple_task",
                {"input": "hello"},
                meta={"custom_field": "test_value"},
            )

            # Should get CreateTaskResult
            task_id = result.task.taskId
            assert result.task.status == "working"

            # Wait for work to complete
            with anyio.fail_after(5):
                await work_completed.wait()

            # Poll until task status is completed
            with anyio.fail_after(5):
                while True:
                    task_status = await client_session.experimental.get_task(task_id)
                    if task_status.status == "completed":  # pragma: no branch
                        break

    async with anyio.create_task_group() as tg:
        tg.start_soon(run_server)
        tg.start_soon(run_client)

    # Verify the meta was passed through correctly
    assert received_meta[0] == "test_value"


@pytest.mark.anyio
async def test_run_task_auto_fails_on_exception() -> None:
    """
    Test that run_task automatically fails the task when work raises.
    """
    server = Server("test-run-task-fail")
    server.experimental.enable_tasks()

    work_failed = Event()

    @server.list_tools()
    async def list_tools() -> list[Tool]:
        return [
            Tool(
                name="failing_task",
                description="A task that fails",
                inputSchema={"type": "object"},
                execution=ToolExecution(taskSupport=TASK_REQUIRED),
            )
        ]

    @server.call_tool()
    async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult | CreateTaskResult:
        ctx = server.request_context
        ctx.experimental.validate_task_mode(TASK_REQUIRED)

        async def work(task: ServerTaskContext) -> CallToolResult:
            work_failed.set()
            raise RuntimeError("Something went wrong!")

        return await ctx.experimental.run_task(work)

    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)

    async def run_server() -> None:
        await server.run(
            client_to_server_receive,
            server_to_client_send,
            server.create_initialization_options(),
        )

    async def run_client() -> None:
        async with ClientSession(server_to_client_receive, client_to_server_send) as client_session:
            await client_session.initialize()

            result = await client_session.experimental.call_tool_as_task("failing_task", {})
            task_id = result.task.taskId

            # Wait for work to fail
            with anyio.fail_after(5):
                await work_failed.wait()

            # Poll until task status is failed
            with anyio.fail_after(5):
                while True:
                    task_status = await client_session.experimental.get_task(task_id)
                    if task_status.status == "failed":  # pragma: no branch
                        break

            assert "Something went wrong" in (task_status.statusMessage or "")

    async with anyio.create_task_group() as tg:
        tg.start_soon(run_server)
        tg.start_soon(run_client)


@pytest.mark.anyio
async def test_enable_tasks_auto_registers_handlers() -> None:
    """
    Test that enable_tasks() auto-registers get_task, list_tasks, cancel_task handlers.
    """
    server = Server("test-enable-tasks")

    # Before enable_tasks, no task capabilities
    caps_before = server.get_capabilities(NotificationOptions(), {})
    assert caps_before.tasks is None

    # Enable tasks
    server.experimental.enable_tasks()

    # After enable_tasks, should have task capabilities
    caps_after = server.get_capabilities(NotificationOptions(), {})
    assert caps_after.tasks is not None
    assert caps_after.tasks.list is not None
    assert caps_after.tasks.cancel is not None
    # Verify nested call capability is present
    assert caps_after.tasks.requests is not None
    assert caps_after.tasks.requests.tools is not None
    assert caps_after.tasks.requests.tools.call is not None


@pytest.mark.anyio
async def test_enable_tasks_with_custom_store_and_queue() -> None:
    """Test that enable_tasks() uses provided store and queue instead of defaults."""
    server = Server("test-custom-store-queue")

    # Create custom store and queue
    custom_store = InMemoryTaskStore()
    custom_queue = InMemoryTaskMessageQueue()

    # Enable tasks with custom implementations
    task_support = server.experimental.enable_tasks(store=custom_store, queue=custom_queue)

    # Verify our custom implementations are used
    assert task_support.store is custom_store
    assert task_support.queue is custom_queue


@pytest.mark.anyio
async def test_enable_tasks_skips_default_handlers_when_custom_registered() -> None:
    """Test that enable_tasks() doesn't override already-registered handlers."""
    server = Server("test-custom-handlers")

    # Register custom handlers BEFORE enable_tasks (never called, just for registration)
    @server.experimental.get_task()
    async def custom_get_task(req: GetTaskRequest) -> GetTaskResult:
        raise NotImplementedError

    @server.experimental.get_task_result()
    async def custom_get_task_result(req: GetTaskPayloadRequest) -> GetTaskPayloadResult:
        raise NotImplementedError

    @server.experimental.list_tasks()
    async def custom_list_tasks(req: ListTasksRequest) -> ListTasksResult:
        raise NotImplementedError

    @server.experimental.cancel_task()
    async def custom_cancel_task(req: CancelTaskRequest) -> CancelTaskResult:
        raise NotImplementedError

    # Now enable tasks - should NOT override our custom handlers
    server.experimental.enable_tasks()

    # Verify our custom handlers are still registered (not replaced by defaults)
    # The handlers dict should contain our custom handlers
    assert GetTaskRequest in server.request_handlers
    assert GetTaskPayloadRequest in server.request_handlers
    assert ListTasksRequest in server.request_handlers
    assert CancelTaskRequest in server.request_handlers


@pytest.mark.anyio
async def test_run_task_without_enable_tasks_raises() -> None:
    """Test that run_task raises when enable_tasks() wasn't called."""
    experimental = Experimental(
        task_metadata=None,
        _client_capabilities=None,
        _session=None,
        _task_support=None,  # Not enabled
    )

    async def work(task: ServerTaskContext) -> CallToolResult:
        raise NotImplementedError

    with pytest.raises(RuntimeError, match="Task support not enabled"):
        await experimental.run_task(work)


@pytest.mark.anyio
async def test_task_support_task_group_before_run_raises() -> None:
    """Test that accessing task_group before run() raises RuntimeError."""
    task_support = TaskSupport.in_memory()

    with pytest.raises(RuntimeError, match="TaskSupport not running"):
        _ = task_support.task_group


@pytest.mark.anyio
async def test_run_task_without_session_raises() -> None:
    """Test that run_task raises when session is not available."""
    task_support = TaskSupport.in_memory()

    experimental = Experimental(
        task_metadata=None,
        _client_capabilities=None,
        _session=None,  # No session
        _task_support=task_support,
    )

    async def work(task: ServerTaskContext) -> CallToolResult:
        raise NotImplementedError

    with pytest.raises(RuntimeError, match="Session not available"):
        await experimental.run_task(work)


@pytest.mark.anyio
async def test_run_task_without_task_metadata_raises() -> None:
    """Test that run_task raises when request is not task-augmented."""
    task_support = TaskSupport.in_memory()
    mock_session = Mock()

    experimental = Experimental(
        task_metadata=None,  # Not a task-augmented request
        _client_capabilities=None,
        _session=mock_session,
        _task_support=task_support,
    )

    async def work(task: ServerTaskContext) -> CallToolResult:
        raise NotImplementedError

    with pytest.raises(RuntimeError, match="Request is not task-augmented"):
        await experimental.run_task(work)


@pytest.mark.anyio
async def test_run_task_with_model_immediate_response() -> None:
    """Test that run_task includes model_immediate_response in CreateTaskResult._meta."""
    server = Server("test-run-task-immediate")
    server.experimental.enable_tasks()

    work_completed = Event()
    immediate_response_text = "Processing your request..."

    @server.list_tools()
    async def list_tools() -> list[Tool]:
        return [
            Tool(
                name="task_with_immediate",
                description="A task with immediate response",
                inputSchema={"type": "object"},
                execution=ToolExecution(taskSupport=TASK_REQUIRED),
            )
        ]

    @server.call_tool()
    async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult | CreateTaskResult:
        ctx = server.request_context
        ctx.experimental.validate_task_mode(TASK_REQUIRED)

        async def work(task: ServerTaskContext) -> CallToolResult:
            work_completed.set()
            return CallToolResult(content=[TextContent(type="text", text="Done")])

        return await ctx.experimental.run_task(work, model_immediate_response=immediate_response_text)

    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)

    async def run_server() -> None:
        await server.run(
            client_to_server_receive,
            server_to_client_send,
            server.create_initialization_options(),
        )

    async def run_client() -> None:
        async with ClientSession(server_to_client_receive, client_to_server_send) as client_session:
            await client_session.initialize()

            result = await client_session.experimental.call_tool_as_task("task_with_immediate", {})

            # Verify the immediate response is in _meta
            assert result.meta is not None
            assert "io.modelcontextprotocol/model-immediate-response" in result.meta
            assert result.meta["io.modelcontextprotocol/model-immediate-response"] == immediate_response_text

            with anyio.fail_after(5):
                await work_completed.wait()

    async with anyio.create_task_group() as tg:
        tg.start_soon(run_server)
        tg.start_soon(run_client)


@pytest.mark.anyio
async def test_run_task_doesnt_complete_if_already_terminal() -> None:
    """Test that run_task doesn't auto-complete if work manually completed the task."""
    server = Server("test-already-complete")
    server.experimental.enable_tasks()

    work_completed = Event()

    @server.list_tools()
    async def list_tools() -> list[Tool]:
        return [
            Tool(
                name="manual_complete_task",
                description="A task that manually completes",
                inputSchema={"type": "object"},
                execution=ToolExecution(taskSupport=TASK_REQUIRED),
            )
        ]

    @server.call_tool()
    async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult | CreateTaskResult:
        ctx = server.request_context
        ctx.experimental.validate_task_mode(TASK_REQUIRED)

        async def work(task: ServerTaskContext) -> CallToolResult:
            # Manually complete the task before returning
            manual_result = CallToolResult(content=[TextContent(type="text", text="Manually completed")])
            await task.complete(manual_result, notify=False)
            work_completed.set()
            # Return a different result - but it should be ignored since task is already terminal
            return CallToolResult(content=[TextContent(type="text", text="This should be ignored")])

        return await ctx.experimental.run_task(work)

    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)

    async def run_server() -> None:
        await server.run(
            client_to_server_receive,
            server_to_client_send,
            server.create_initialization_options(),
        )

    async def run_client() -> None:
        async with ClientSession(server_to_client_receive, client_to_server_send) as client_session:
            await client_session.initialize()

            result = await client_session.experimental.call_tool_as_task("manual_complete_task", {})
            task_id = result.task.taskId

            with anyio.fail_after(5):
                await work_completed.wait()

            # Poll until task status is completed
            with anyio.fail_after(5):
                while True:
                    status = await client_session.experimental.get_task(task_id)
                    if status.status == "completed":  # pragma: no branch
                        break

    async with anyio.create_task_group() as tg:
        tg.start_soon(run_server)
        tg.start_soon(run_client)


@pytest.mark.anyio
async def test_run_task_doesnt_fail_if_already_terminal() -> None:
    """Test that run_task doesn't auto-fail if work manually failed/cancelled the task."""
    server = Server("test-already-failed")
    server.experimental.enable_tasks()

    work_completed = Event()

    @server.list_tools()
    async def list_tools() -> list[Tool]:
        return [
            Tool(
                name="manual_cancel_task",
                description="A task that manually cancels then raises",
                inputSchema={"type": "object"},
                execution=ToolExecution(taskSupport=TASK_REQUIRED),
            )
        ]

    @server.call_tool()
    async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult | CreateTaskResult:
        ctx = server.request_context
        ctx.experimental.validate_task_mode(TASK_REQUIRED)

        async def work(task: ServerTaskContext) -> CallToolResult:
            # Manually fail the task first
            await task.fail("Manually failed", notify=False)
            work_completed.set()
            # Then raise - but the auto-fail should be skipped since task is already terminal
            raise RuntimeError("This error should not change status")

        return await ctx.experimental.run_task(work)

    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)

    async def run_server() -> None:
        await server.run(
            client_to_server_receive,
            server_to_client_send,
            server.create_initialization_options(),
        )

    async def run_client() -> None:
        async with ClientSession(server_to_client_receive, client_to_server_send) as client_session:
            await client_session.initialize()

            result = await client_session.experimental.call_tool_as_task("manual_cancel_task", {})
            task_id = result.task.taskId

            with anyio.fail_after(5):
                await work_completed.wait()

            # Poll until task status is failed
            with anyio.fail_after(5):
                while True:
                    status = await client_session.experimental.get_task(task_id)
                    if status.status == "failed":  # pragma: no branch
                        break

            # Task should still be failed (from manual fail, not auto-fail from exception)
            assert status.statusMessage == "Manually failed"  # Not "This error should not change status"

    async with anyio.create_task_group() as tg:
        tg.start_soon(run_server)
        tg.start_soon(run_client)