modelcontextprotocol--python-sdk/tests/experimental/tasks/server/test_server_task_context.py

"""Tests for ServerTaskContext."""

import asyncio
from unittest.mock import AsyncMock, Mock

import anyio
import pytest

from mcp.server.experimental.task_context import ServerTaskContext
from mcp.server.experimental.task_result_handler import TaskResultHandler
from mcp.shared.exceptions import MCPError
from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
from mcp.shared.experimental.tasks.message_queue import InMemoryTaskMessageQueue
from mcp.types import (
    CallToolResult,
    ClientCapabilities,
    ClientTasksCapability,
    ClientTasksRequestsCapability,
    Implementation,
    InitializeRequestParams,
    JSONRPCRequest,
    SamplingMessage,
    TaskMetadata,
    TasksCreateElicitationCapability,
    TasksCreateMessageCapability,
    TasksElicitationCapability,
    TasksSamplingCapability,
    TextContent,
)


@pytest.mark.anyio
async def test_server_task_context_properties() -> None:
    """Test ServerTaskContext property accessors."""
    store = InMemoryTaskStore()
    mock_session = Mock()
    queue = InMemoryTaskMessageQueue()
    task = await store.create_task(TaskMetadata(ttl=60000), task_id="test-123")

    ctx = ServerTaskContext(
        task=task,
        store=store,
        session=mock_session,
        queue=queue,
    )

    assert ctx.task_id == "test-123"
    assert ctx.task.task_id == "test-123"
    assert ctx.is_cancelled is False

    store.cleanup()


@pytest.mark.anyio
async def test_server_task_context_request_cancellation() -> None:
    """Test ServerTaskContext.request_cancellation()."""
    store = InMemoryTaskStore()
    mock_session = Mock()
    queue = InMemoryTaskMessageQueue()
    task = await store.create_task(TaskMetadata(ttl=60000))

    ctx = ServerTaskContext(
        task=task,
        store=store,
        session=mock_session,
        queue=queue,
    )

    assert ctx.is_cancelled is False
    ctx.request_cancellation()
    assert ctx.is_cancelled is True

    store.cleanup()


@pytest.mark.anyio
async def test_server_task_context_update_status_with_notify() -> None:
    """Test update_status sends notification when notify=True."""
    store = InMemoryTaskStore()
    mock_session = Mock()
    mock_session.send_notification = AsyncMock()
    queue = InMemoryTaskMessageQueue()
    task = await store.create_task(TaskMetadata(ttl=60000))

    ctx = ServerTaskContext(
        task=task,
        store=store,
        session=mock_session,
        queue=queue,
    )

    await ctx.update_status("Working...", notify=True)

    mock_session.send_notification.assert_called_once()
    store.cleanup()


@pytest.mark.anyio
async def test_server_task_context_update_status_without_notify() -> None:
    """Test update_status skips notification when notify=False."""
    store = InMemoryTaskStore()
    mock_session = Mock()
    mock_session.send_notification = AsyncMock()
    queue = InMemoryTaskMessageQueue()
    task = await store.create_task(TaskMetadata(ttl=60000))

    ctx = ServerTaskContext(
        task=task,
        store=store,
        session=mock_session,
        queue=queue,
    )

    await ctx.update_status("Working...", notify=False)

    mock_session.send_notification.assert_not_called()
    store.cleanup()


@pytest.mark.anyio
async def test_server_task_context_complete_with_notify() -> None:
    """Test complete sends notification when notify=True."""
    store = InMemoryTaskStore()
    mock_session = Mock()
    mock_session.send_notification = AsyncMock()
    queue = InMemoryTaskMessageQueue()
    task = await store.create_task(TaskMetadata(ttl=60000))

    ctx = ServerTaskContext(
        task=task,
        store=store,
        session=mock_session,
        queue=queue,
    )

    result = CallToolResult(content=[TextContent(type="text", text="Done")])
    await ctx.complete(result, notify=True)

    mock_session.send_notification.assert_called_once()
    store.cleanup()


@pytest.mark.anyio
async def test_server_task_context_fail_with_notify() -> None:
    """Test fail sends notification when notify=True."""
    store = InMemoryTaskStore()
    mock_session = Mock()
    mock_session.send_notification = AsyncMock()
    queue = InMemoryTaskMessageQueue()
    task = await store.create_task(TaskMetadata(ttl=60000))

    ctx = ServerTaskContext(
        task=task,
        store=store,
        session=mock_session,
        queue=queue,
    )

    await ctx.fail("Something went wrong", notify=True)

    mock_session.send_notification.assert_called_once()
    store.cleanup()


@pytest.mark.anyio
async def test_elicit_raises_when_client_lacks_capability() -> None:
    """Test that elicit() raises MCPError when client doesn't support elicitation."""
    store = InMemoryTaskStore()
    mock_session = Mock()
    mock_session.check_client_capability = Mock(return_value=False)
    queue = InMemoryTaskMessageQueue()
    handler = TaskResultHandler(store, queue)
    task = await store.create_task(TaskMetadata(ttl=60000))

    ctx = ServerTaskContext(
        task=task,
        store=store,
        session=mock_session,
        queue=queue,
        handler=handler,
    )

    with pytest.raises(MCPError) as exc_info:
        await ctx.elicit(message="Test?", requested_schema={"type": "object"})

    assert "elicitation capability" in exc_info.value.error.message
    mock_session.check_client_capability.assert_called_once()
    store.cleanup()


@pytest.mark.anyio
async def test_create_message_raises_when_client_lacks_capability() -> None:
    """Test that create_message() raises MCPError when client doesn't support sampling."""
    store = InMemoryTaskStore()
    mock_session = Mock()
    mock_session.check_client_capability = Mock(return_value=False)
    queue = InMemoryTaskMessageQueue()
    handler = TaskResultHandler(store, queue)
    task = await store.create_task(TaskMetadata(ttl=60000))

    ctx = ServerTaskContext(
        task=task,
        store=store,
        session=mock_session,
        queue=queue,
        handler=handler,
    )

    with pytest.raises(MCPError) as exc_info:
        await ctx.create_message(messages=[], max_tokens=100)

    assert "sampling capability" in exc_info.value.error.message
    mock_session.check_client_capability.assert_called_once()
    store.cleanup()


@pytest.mark.anyio
async def test_elicit_raises_without_handler() -> None:
    """Test that elicit() raises when handler is not provided."""
    store = InMemoryTaskStore()
    mock_session = Mock()
    mock_session.check_client_capability = Mock(return_value=True)
    queue = InMemoryTaskMessageQueue()
    task = await store.create_task(TaskMetadata(ttl=60000))

    ctx = ServerTaskContext(
        task=task,
        store=store,
        session=mock_session,
        queue=queue,
        handler=None,
    )

    with pytest.raises(RuntimeError, match="handler is required"):
        await ctx.elicit(message="Test?", requested_schema={"type": "object"})

    store.cleanup()


@pytest.mark.anyio
async def test_elicit_url_raises_without_handler() -> None:
    """Test that elicit_url() raises when handler is not provided."""
    store = InMemoryTaskStore()
    mock_session = Mock()
    mock_session.check_client_capability = Mock(return_value=True)
    queue = InMemoryTaskMessageQueue()
    task = await store.create_task(TaskMetadata(ttl=60000))

    ctx = ServerTaskContext(
        task=task,
        store=store,
        session=mock_session,
        queue=queue,
        handler=None,
    )

    with pytest.raises(RuntimeError, match="handler is required for elicit_url"):
        await ctx.elicit_url(
            message="Please authorize",
            url="https://example.com/oauth",
            elicitation_id="oauth-123",
        )

    store.cleanup()


@pytest.mark.anyio
async def test_create_message_raises_without_handler() -> None:
    """Test that create_message() raises when handler is not provided."""
    store = InMemoryTaskStore()
    mock_session = Mock()
    mock_session.check_client_capability = Mock(return_value=True)
    queue = InMemoryTaskMessageQueue()
    task = await store.create_task(TaskMetadata(ttl=60000))

    ctx = ServerTaskContext(
        task=task,
        store=store,
        session=mock_session,
        queue=queue,
        handler=None,
    )

    with pytest.raises(RuntimeError, match="handler is required"):
        await ctx.create_message(messages=[], max_tokens=100)

    store.cleanup()


@pytest.mark.anyio
async def test_elicit_queues_request_and_waits_for_response() -> None:
    """Test that elicit() queues request and waits for response."""
    store = InMemoryTaskStore()
    queue = InMemoryTaskMessageQueue()
    handler = TaskResultHandler(store, queue)
    task = await store.create_task(TaskMetadata(ttl=60000))

    mock_session = Mock()
    mock_session.check_client_capability = Mock(return_value=True)
    mock_session._build_elicit_form_request = Mock(
        return_value=JSONRPCRequest(
            jsonrpc="2.0",
            id="test-req-1",
            method="elicitation/create",
            params={"message": "Test?", "_meta": {}},
        )
    )

    ctx = ServerTaskContext(
        task=task,
        store=store,
        session=mock_session,
        queue=queue,
        handler=handler,
    )

    elicit_result = None

    async def run_elicit() -> None:
        nonlocal elicit_result
        elicit_result = await ctx.elicit(
            message="Test?",
            requested_schema={"type": "object"},
        )

    async with anyio.create_task_group() as tg:
        tg.start_soon(run_elicit)

        # Wait for request to be queued
        await queue.wait_for_message(task.task_id)

        # Verify task is in input_required status
        updated_task = await store.get_task(task.task_id)
        assert updated_task is not None
        assert updated_task.status == "input_required"

        # Dequeue and simulate response
        msg = await queue.dequeue(task.task_id)
        assert msg is not None
        assert msg.resolver is not None

        # Resolve with mock elicitation response
        msg.resolver.set_result({"action": "accept", "content": {"name": "Alice"}})

    # Verify result
    assert elicit_result is not None
    assert elicit_result.action == "accept"
    assert elicit_result.content == {"name": "Alice"}

    # Verify task is back to working
    final_task = await store.get_task(task.task_id)
    assert final_task is not None
    assert final_task.status == "working"

    store.cleanup()


@pytest.mark.anyio
async def test_elicit_url_queues_request_and_waits_for_response() -> None:
    """Test that elicit_url() queues request and waits for response."""
    store = InMemoryTaskStore()
    queue = InMemoryTaskMessageQueue()
    handler = TaskResultHandler(store, queue)
    task = await store.create_task(TaskMetadata(ttl=60000))

    mock_session = Mock()
    mock_session.check_client_capability = Mock(return_value=True)
    mock_session._build_elicit_url_request = Mock(
        return_value=JSONRPCRequest(
            jsonrpc="2.0",
            id="test-url-req-1",
            method="elicitation/create",
            params={"message": "Authorize", "url": "https://example.com", "elicitationId": "123", "mode": "url"},
        )
    )

    ctx = ServerTaskContext(
        task=task,
        store=store,
        session=mock_session,
        queue=queue,
        handler=handler,
    )

    elicit_result = None

    async def run_elicit_url() -> None:
        nonlocal elicit_result
        elicit_result = await ctx.elicit_url(
            message="Authorize",
            url="https://example.com/oauth",
            elicitation_id="oauth-123",
        )

    async with anyio.create_task_group() as tg:
        tg.start_soon(run_elicit_url)

        # Wait for request to be queued
        await queue.wait_for_message(task.task_id)

        # Verify task is in input_required status
        updated_task = await store.get_task(task.task_id)
        assert updated_task is not None
        assert updated_task.status == "input_required"

        # Dequeue and simulate response
        msg = await queue.dequeue(task.task_id)
        assert msg is not None
        assert msg.resolver is not None

        # Resolve with mock elicitation response (URL mode just returns action)
        msg.resolver.set_result({"action": "accept"})

    # Verify result
    assert elicit_result is not None
    assert elicit_result.action == "accept"

    # Verify task is back to working
    final_task = await store.get_task(task.task_id)
    assert final_task is not None
    assert final_task.status == "working"

    store.cleanup()


@pytest.mark.anyio
async def test_create_message_queues_request_and_waits_for_response() -> None:
    """Test that create_message() queues request and waits for response."""
    store = InMemoryTaskStore()
    queue = InMemoryTaskMessageQueue()
    handler = TaskResultHandler(store, queue)
    task = await store.create_task(TaskMetadata(ttl=60000))

    mock_session = Mock()
    mock_session.check_client_capability = Mock(return_value=True)
    mock_session._build_create_message_request = Mock(
        return_value=JSONRPCRequest(
            jsonrpc="2.0",
            id="test-req-2",
            method="sampling/createMessage",
            params={"messages": [], "maxTokens": 100, "_meta": {}},
        )
    )

    ctx = ServerTaskContext(
        task=task,
        store=store,
        session=mock_session,
        queue=queue,
        handler=handler,
    )

    sampling_result = None

    async def run_sampling() -> None:
        nonlocal sampling_result
        sampling_result = await ctx.create_message(
            messages=[SamplingMessage(role="user", content=TextContent(type="text", text="Hello"))],
            max_tokens=100,
        )

    async with anyio.create_task_group() as tg:
        tg.start_soon(run_sampling)

        # Wait for request to be queued
        await queue.wait_for_message(task.task_id)

        # Verify task is in input_required status
        updated_task = await store.get_task(task.task_id)
        assert updated_task is not None
        assert updated_task.status == "input_required"

        # Dequeue and simulate response
        msg = await queue.dequeue(task.task_id)
        assert msg is not None
        assert msg.resolver is not None

        # Resolve with mock sampling response
        msg.resolver.set_result(
            {
                "role": "assistant",
                "content": {"type": "text", "text": "Hello back!"},
                "model": "test-model",
                "stopReason": "endTurn",
            }
        )

    # Verify result
    assert sampling_result is not None
    assert sampling_result.role == "assistant"
    assert sampling_result.model == "test-model"

    # Verify task is back to working
    final_task = await store.get_task(task.task_id)
    assert final_task is not None
    assert final_task.status == "working"

    store.cleanup()


@pytest.mark.anyio
async def test_elicit_restores_status_on_cancellation() -> None:
    """Test that elicit() restores task status to working when cancelled."""
    store = InMemoryTaskStore()
    queue = InMemoryTaskMessageQueue()
    handler = TaskResultHandler(store, queue)
    task = await store.create_task(TaskMetadata(ttl=60000))

    mock_session = Mock()
    mock_session.check_client_capability = Mock(return_value=True)
    mock_session._build_elicit_form_request = Mock(
        return_value=JSONRPCRequest(
            jsonrpc="2.0",
            id="test-req-cancel",
            method="elicitation/create",
            params={"message": "Test?", "_meta": {}},
        )
    )

    ctx = ServerTaskContext(
        task=task,
        store=store,
        session=mock_session,
        queue=queue,
        handler=handler,
    )

    cancelled_error_raised = False

    async with anyio.create_task_group() as tg:

        async def do_elicit() -> None:
            nonlocal cancelled_error_raised
            try:
                await ctx.elicit(
                    message="Test?",
                    requested_schema={"type": "object"},
                )
            except anyio.get_cancelled_exc_class():
                cancelled_error_raised = True
                # Don't re-raise - let the test continue

        tg.start_soon(do_elicit)

        # Wait for request to be queued
        await queue.wait_for_message(task.task_id)

        # Verify task is in input_required status
        updated_task = await store.get_task(task.task_id)
        assert updated_task is not None
        assert updated_task.status == "input_required"

        # Get the queued message and set cancellation exception on its resolver
        msg = await queue.dequeue(task.task_id)
        assert msg is not None
        assert msg.resolver is not None

        # Trigger cancellation by setting exception (use asyncio.CancelledError directly)
        msg.resolver.set_exception(asyncio.CancelledError())

    # Verify task is back to working after cancellation
    final_task = await store.get_task(task.task_id)
    assert final_task is not None
    assert final_task.status == "working"
    assert cancelled_error_raised

    store.cleanup()


@pytest.mark.anyio
async def test_create_message_restores_status_on_cancellation() -> None:
    """Test that create_message() restores task status to working when cancelled."""
    store = InMemoryTaskStore()
    queue = InMemoryTaskMessageQueue()
    handler = TaskResultHandler(store, queue)
    task = await store.create_task(TaskMetadata(ttl=60000))

    mock_session = Mock()
    mock_session.check_client_capability = Mock(return_value=True)
    mock_session._build_create_message_request = Mock(
        return_value=JSONRPCRequest(
            jsonrpc="2.0",
            id="test-req-cancel-2",
            method="sampling/createMessage",
            params={"messages": [], "maxTokens": 100, "_meta": {}},
        )
    )

    ctx = ServerTaskContext(
        task=task,
        store=store,
        session=mock_session,
        queue=queue,
        handler=handler,
    )

    cancelled_error_raised = False

    async with anyio.create_task_group() as tg:

        async def do_sampling() -> None:
            nonlocal cancelled_error_raised
            try:
                await ctx.create_message(
                    messages=[SamplingMessage(role="user", content=TextContent(type="text", text="Hello"))],
                    max_tokens=100,
                )
            except anyio.get_cancelled_exc_class():
                cancelled_error_raised = True
                # Don't re-raise

        tg.start_soon(do_sampling)

        # Wait for request to be queued
        await queue.wait_for_message(task.task_id)

        # Verify task is in input_required status
        updated_task = await store.get_task(task.task_id)
        assert updated_task is not None
        assert updated_task.status == "input_required"

        # Get the queued message and set cancellation exception on its resolver
        msg = await queue.dequeue(task.task_id)
        assert msg is not None
        assert msg.resolver is not None

        # Trigger cancellation by setting exception (use asyncio.CancelledError directly)
        msg.resolver.set_exception(asyncio.CancelledError())

    # Verify task is back to working after cancellation
    final_task = await store.get_task(task.task_id)
    assert final_task is not None
    assert final_task.status == "working"
    assert cancelled_error_raised

    store.cleanup()


@pytest.mark.anyio
async def test_elicit_as_task_raises_without_handler() -> None:
    """Test that elicit_as_task() raises when handler is not provided."""
    store = InMemoryTaskStore()
    queue = InMemoryTaskMessageQueue()
    task = await store.create_task(TaskMetadata(ttl=60000))

    # Create mock session with proper client capabilities
    mock_session = Mock()
    mock_session.client_params = InitializeRequestParams(
        protocol_version="2025-01-01",
        capabilities=ClientCapabilities(
            tasks=ClientTasksCapability(
                requests=ClientTasksRequestsCapability(
                    elicitation=TasksElicitationCapability(create=TasksCreateElicitationCapability())
                )
            )
        ),
        client_info=Implementation(name="test", version="1.0"),
    )

    ctx = ServerTaskContext(
        task=task,
        store=store,
        session=mock_session,
        queue=queue,
        handler=None,
    )

    with pytest.raises(RuntimeError, match="handler is required for elicit_as_task"):
        await ctx.elicit_as_task(message="Test?", requested_schema={"type": "object"})

    store.cleanup()


@pytest.mark.anyio
async def test_create_message_as_task_raises_without_handler() -> None:
    """Test that create_message_as_task() raises when handler is not provided."""
    store = InMemoryTaskStore()
    queue = InMemoryTaskMessageQueue()
    task = await store.create_task(TaskMetadata(ttl=60000))

    # Create mock session with proper client capabilities
    mock_session = Mock()
    mock_session.client_params = InitializeRequestParams(
        protocol_version="2025-01-01",
        capabilities=ClientCapabilities(
            tasks=ClientTasksCapability(
                requests=ClientTasksRequestsCapability(
                    sampling=TasksSamplingCapability(create_message=TasksCreateMessageCapability())
                )
            )
        ),
        client_info=Implementation(name="test", version="1.0"),
    )

    ctx = ServerTaskContext(
        task=task,
        store=store,
        session=mock_session,
        queue=queue,
        handler=None,
    )

    with pytest.raises(RuntimeError, match="handler is required for create_message_as_task"):
        await ctx.create_message_as_task(
            messages=[SamplingMessage(role="user", content=TextContent(type="text", text="Hello"))],
            max_tokens=100,
        )

    store.cleanup()