prefecthq--fastmcp/tests/client/tasks/test_client_task_notifications.py

"""
Tests for client-side handling of notifications/tasks/status (SEP-1686 lines 436-444).

Verifies that Task objects receive notifications, update their cache, wake up wait() calls,
and invoke user callbacks.
"""

import asyncio
import time
from datetime import datetime, timezone

import pytest
from mcp_types import GetTaskResult

from fastmcp import FastMCP
from fastmcp.client import Client


@pytest.fixture
async def task_notification_server():
    """Server that sends task status notifications."""
    mcp = FastMCP("task-notification-test")

    @mcp.tool(task=True)
    async def quick_task(value: int) -> int:
        """Quick background task."""
        await asyncio.sleep(0.05)
        return value * 2

    @mcp.tool(task=True)
    async def slow_task(duration: float = 0.2) -> str:
        """Slow background task."""
        await asyncio.sleep(duration)
        return "done"

    @mcp.tool(task=True)
    async def failing_task() -> str:
        """Task that fails."""
        raise ValueError("Intentional failure")

    return mcp


async def test_task_receives_status_notification(task_notification_server):
    """Task object receives and processes status notifications."""
    async with Client(task_notification_server) as client:
        task = await client.call_tool("quick_task", {"value": 5}, task=True)

        # Wait for task to complete (notification should arrive)
        status = await task.wait(timeout=2.0)

        # Verify task completed
        assert status.status == "completed"


async def test_status_cache_updated_by_notification(task_notification_server):
    """Cached status is updated when notification arrives."""
    async with Client(task_notification_server) as client:
        task = await client.call_tool("quick_task", {"value": 10}, task=True)

        # Wait for completion (notification should update cache)
        await task.wait(timeout=2.0)

        # Status should be cached (no server call needed)
        # Call status() twice - should return same cached object
        status1 = await task.status()
        status2 = await task.status()

        # Should be the exact same object (from cache)
        assert status1 is status2
        assert status1.status == "completed"


async def test_callback_invoked_on_notification(task_notification_server):
    """User callback is invoked when notification arrives."""
    callback_invocations = []

    def status_callback(status: GetTaskResult):
        """Sync callback."""
        callback_invocations.append(status)

    async with Client(task_notification_server) as client:
        task = await client.call_tool("quick_task", {"value": 7}, task=True)

        # Register callback
        task.on_status_change(status_callback)

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

        # Give callbacks a moment to fire
        await asyncio.sleep(0.1)

    # Callback should have been invoked at least once
    assert len(callback_invocations) > 0

    # Should have received completed status
    completed_statuses = [s for s in callback_invocations if s.status == "completed"]
    assert len(completed_statuses) > 0


async def test_async_callback_invoked(task_notification_server):
    """Async callback is invoked when notification arrives."""
    callback_invocations = []

    async def async_status_callback(status: GetTaskResult):
        """Async callback."""
        await asyncio.sleep(0.01)  # Simulate async work
        callback_invocations.append(status)

    async with Client(task_notification_server) as client:
        task = await client.call_tool("quick_task", {"value": 3}, task=True)

        # Register async callback
        task.on_status_change(async_status_callback)

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

        # Give async callbacks time to complete
        await asyncio.sleep(0.2)

    # Async callback should have been invoked
    assert len(callback_invocations) > 0


async def test_multiple_callbacks_all_invoked(task_notification_server):
    """Multiple callbacks are all invoked."""
    callback1_calls = []
    callback2_calls = []

    def callback1(status: GetTaskResult):
        callback1_calls.append(status.status)

    def callback2(status: GetTaskResult):
        callback2_calls.append(status.status)

    async with Client(task_notification_server) as client:
        task = await client.call_tool("quick_task", {"value": 8}, task=True)

        task.on_status_change(callback1)
        task.on_status_change(callback2)

        await task.wait(timeout=2.0)
        await asyncio.sleep(0.1)

    # Both callbacks should have been invoked
    assert len(callback1_calls) > 0
    assert len(callback2_calls) > 0


async def test_callback_error_doesnt_break_notification(task_notification_server):
    """Callback errors don't prevent other callbacks from running."""
    callback1_calls = []
    callback2_calls = []

    def failing_callback(status: GetTaskResult):
        callback1_calls.append("called")
        raise ValueError("Callback intentionally fails")

    def working_callback(status: GetTaskResult):
        callback2_calls.append(status.status)

    async with Client(task_notification_server) as client:
        task = await client.call_tool("quick_task", {"value": 12}, task=True)

        task.on_status_change(failing_callback)
        task.on_status_change(working_callback)

        await task.wait(timeout=2.0)
        await asyncio.sleep(0.1)

    # Failing callback was called (and errored)
    assert len(callback1_calls) > 0

    # Working callback should still have been invoked
    assert len(callback2_calls) > 0


async def test_wait_wakes_early_on_notification(task_notification_server):
    """wait() wakes up immediately when notification arrives, not after poll interval."""
    async with Client(task_notification_server) as client:
        task = await client.call_tool("quick_task", {"value": 15}, task=True)

        # Record timing
        start = time.time()
        status = await task.wait(timeout=5.0)
        elapsed = time.time() - start

        # Should complete much faster than the fallback poll interval (500ms)
        # With notifications, should be < 200ms for quick task
        # Without notifications, would take 500ms+ due to polling
        assert elapsed < 1.0  # Very generous bound
        assert status.status == "completed"


async def test_notification_with_failed_task(task_notification_server):
    """Notifications work for failed tasks too."""
    async with Client(task_notification_server) as client:
        task = await client.call_tool("failing_task", {}, task=True)

        with pytest.raises(Exception):
            await task

        # Should have cached the failed status from notification
        status = await task.status()
        assert status.status == "failed"
        assert (
            status.status_message is not None
        )  # Error details in statusMessage per spec


async def test_wait_returns_on_input_required(task_notification_server):
    """wait() should return immediately when task enters input_required, not hang."""
    async with Client(task_notification_server) as client:
        task = await client.call_tool("quick_task", {"value": 1}, task=True)

        # Directly inject an input_required status into the cache and signal the event.
        # SDK v2 types the Task timestamps as ISO 8601 strings.
        now = datetime.now(timezone.utc).isoformat()
        input_required_status = GetTaskResult(
            task_id=task._task_id,
            status="input_required",
            status_message="Waiting for user input",
            created_at=now,
            last_updated_at=now,
            ttl=None,
        )
        task._status_cache = input_required_status
        if task._status_event is None:
            task._status_event = asyncio.Event()
        task._status_event.set()

        # Should return immediately with input_required, not hang for 300s
        status = await task.wait(timeout=2.0)
        assert status.status == "input_required"