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

"""Tests for Context background task support (SEP-1686).

Tests Context API surface (unit) and background task elicitation (integration).
Integration tests use Client(mcp) with the real memory:// Docket backend —
no mocking of Redis, Docket, or session internals.
"""

import asyncio
import json
from datetime import datetime, timezone
from typing import Any, cast
from unittest.mock import AsyncMock, patch

import pytest
from mcp import ServerSession
from mcp.server.auth.middleware.auth_context import auth_context_var
from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser
from mcp_types import (
    ClientCapabilities,
    CreateMessageResult,
    Implementation,
    InitializeRequestParams,
    TextContent,
)
from pydantic import BaseModel

from fastmcp import FastMCP
from fastmcp.client import Client
from fastmcp.client.elicitation import ElicitResult
from fastmcp.dependencies import CurrentDocket
from fastmcp.server.auth import AccessToken
from fastmcp.server.context import Context
from fastmcp.server.dependencies import get_access_token
from fastmcp.server.elicitation import (
    AcceptedElicitation,
    CancelledElicitation,
    DeclinedElicitation,
)
from fastmcp.server.tasks.context import (
    TaskContextInfo,
    TaskContextSnapshot,
    _remember_snapshot,
    get_task_scope,
)
from fastmcp.server.tasks.elicitation import handle_task_input
from fastmcp.server.tasks.keys import (
    task_redis_prefix,
)

# =============================================================================
# Unit tests: Context API surface (no Redis/Docket needed)
# =============================================================================


class TestContextBackgroundTaskSupport:
    """Tests for Context.is_background_task and related functionality."""

    def test_context_not_background_task_by_default(self):
        """Context should not be a background task by default."""
        mcp = FastMCP("test")
        ctx = Context(mcp)
        assert ctx.is_background_task is False
        assert ctx.task_id is None

    def test_context_is_background_task_when_task_id_provided(self):
        """Context should be a background task when task_id is provided."""
        mcp = FastMCP("test")
        ctx = Context(mcp, task_id="test-task-123")
        assert ctx.is_background_task is True
        assert ctx.task_id == "test-task-123"

    def test_context_task_id_is_readonly(self):
        """task_id should be a read-only property."""
        mcp = FastMCP("test")
        ctx = Context(mcp, task_id="test-task-123")
        with pytest.raises(AttributeError):
            setattr(ctx, "task_id", "new-id")


class TestContextSessionProperty:
    """Tests for Context.session property in different modes."""

    def test_session_raises_when_no_session_available(self):
        """session should raise RuntimeError when no session is available."""
        mcp = FastMCP("test")
        ctx = Context(mcp)  # No session, not a background task

        with pytest.raises(RuntimeError, match="session is not available"):
            _ = ctx.session

    def test_session_uses_stored_session_in_background_task(self):
        """session should use _session in background task mode."""
        mcp = FastMCP("test")

        class MockSession:
            _fastmcp_state_prefix = "test-session"

        mock_session = MockSession()
        ctx = Context(
            mcp, session=cast(ServerSession, mock_session), task_id="test-task-123"
        )

        assert ctx.session is mock_session

    def test_session_uses_stored_session_during_on_initialize(self):
        """session should use _session during on_initialize (no request context)."""
        mcp = FastMCP("test")

        class MockSession:
            _fastmcp_state_prefix = "test-session"

        mock_session = MockSession()
        ctx = Context(mcp, session=cast(ServerSession, mock_session))

        assert ctx.session is mock_session


class TestContextBackgroundTaskLogging:
    """Tests for per-session log gating in background task mode."""

    def _make_task_context(
        self, mcp: FastMCP, session_id: str
    ) -> tuple[Context, AsyncMock]:
        send_log_message = AsyncMock()

        class MockConnection:
            def __init__(self, session_id: str) -> None:
                self.session_id = session_id

        class MockSession:
            def __init__(self, session_id: str) -> None:
                self._connection = MockConnection(session_id)
                self._fastmcp_state_prefix = session_id
                self.send_log_message = send_log_message

        session = MockSession(session_id)
        ctx = Context(
            mcp, session=cast(ServerSession, session), task_id="test-task-123"
        )
        return ctx, send_log_message

    async def test_background_task_honors_session_level(self):
        """A background task has a session but no request context; the
        per-session minimum registered via logging/setLevel must still gate
        its logs, so sub-threshold messages are not sent to the client."""
        mcp = FastMCP("test")
        session_id = "session-abc"
        mcp._client_log_levels[session_id] = "error"

        ctx, send_log_message = self._make_task_context(mcp, session_id)
        assert ctx.is_background_task is True
        assert ctx.request_context is None

        await ctx.info("info msg")
        send_log_message.assert_not_called()

        await ctx.error("error msg")
        send_log_message.assert_called_once()

    async def test_background_task_without_session_level_sends_all(self):
        """When no per-session level is registered, background-task logs fall
        back to the server default (which allows everything by default)."""
        mcp = FastMCP("test")
        ctx, send_log_message = self._make_task_context(mcp, "session-xyz")

        await ctx.info("info msg")
        send_log_message.assert_called_once()


class TestContextClientExtensionBackgroundTask:
    """Tests for Context.client_supports_extension() in background task mode.

    A background task has a live snapshot session but no request context. The
    client's advertised capabilities are preserved on the snapshot session's
    ``client_params``, so extension detection must read from the session rather
    than gating on ``request_context``.
    """

    def _make_task_context(
        self, mcp: FastMCP, extensions: dict[str, dict[str, Any]] | None
    ) -> Context:
        capabilities = ClientCapabilities(extensions=extensions)
        client_params = InitializeRequestParams(
            protocol_version="2025-06-18",
            capabilities=capabilities,
            client_info=Implementation(name="test-client", version="1.0"),
        )

        class MockSession:
            _fastmcp_state_prefix = "session-ext"

            def __init__(self) -> None:
                self.client_params = client_params

        session = MockSession()
        return Context(
            mcp, session=cast(ServerSession, session), task_id="test-task-ext"
        )

    def test_background_task_detects_advertised_extension(self):
        """The snapshot session preserves the client's initialize params, so an
        advertised extension is detected even with no request context."""
        mcp = FastMCP("test")
        ctx = self._make_task_context(mcp, {"ext-abc": {}})

        assert ctx.is_background_task is True
        assert ctx.request_context is None
        assert ctx.client_supports_extension("ext-abc") is True
        assert ctx.client_supports_extension("ext-missing") is False

    def test_background_task_no_extensions_returns_false(self):
        """When the client advertised no extensions, detection returns False."""
        mcp = FastMCP("test")
        ctx = self._make_task_context(mcp, None)

        assert ctx.client_supports_extension("ext-abc") is False

    def test_no_session_returns_false(self):
        """With no session available at all (e.g. distributed worker), the
        method degrades to False rather than raising."""
        mcp = FastMCP("test")
        ctx = Context(mcp, task_id="test-task-ext")

        assert ctx.client_supports_extension("ext-abc") is False


class TestContextElicitBackgroundTask:
    """Tests for Context.elicit() in background task mode."""

    async def test_elicit_raises_when_background_task_but_no_docket(self):
        """elicit() should raise when in background task mode but Docket unavailable."""
        mcp = FastMCP("test")
        ctx = Context(mcp, task_id="test-task-123")

        class MockSession:
            _fastmcp_state_prefix = "test-session"

        ctx._session = cast(ServerSession, MockSession())

        with pytest.raises(RuntimeError, match="Docket"):
            await ctx.elicit("Need input", str)


class TestElicitFailFast:
    """Tests for elicit_for_task fail-fast on notification push failure."""

    async def test_elicit_returns_cancel_when_notification_push_fails(self):
        """elicit_for_task should return cancel immediately when push_notification fails.

        If the client can't receive the input_required notification, waiting
        for a response that will never come would block for up to 1 hour.
        Instead, we return cancel immediately (fail-fast).

        This test patches ONLY push_notification — all other components
        (Docket, Redis, session) are real via the memory:// backend.
        """
        mcp = FastMCP("failfast-test")
        elicit_started = asyncio.Event()
        captured: dict[str, object] = {}

        @mcp.tool(task=True)
        async def failfast_tool(ctx: Context) -> str:
            elicit_started.set()
            result = await ctx.elicit("This notification will fail", str)
            captured["result_type"] = type(result).__name__
            captured["is_cancelled"] = isinstance(result, CancelledElicitation)
            return "done"

        # Patch push_notification BEFORE starting client so it's active
        # when the tool runs in the Docket worker
        with patch(
            "fastmcp.server.tasks.notifications.push_notification",
            side_effect=ConnectionError("Redis queue unavailable"),
        ):
            async with Client(mcp) as client:
                task = await client.call_tool("failfast_tool", {}, task=True)
                await asyncio.wait_for(elicit_started.wait(), timeout=5.0)
                await task.wait(timeout=10.0)
                result = await task.result()
                assert result.data == "done"

        # The tool should have received CancelledElicitation (fail-fast)
        assert captured["is_cancelled"] is True
        assert captured["result_type"] == "CancelledElicitation"


class TestContextDocumentation:
    """Tests to verify Context documentation and API surface."""

    def test_is_background_task_has_docstring(self):
        """is_background_task property should have documentation."""
        assert Context.is_background_task.__doc__ is not None
        assert "background task" in Context.is_background_task.__doc__.lower()

    def test_task_id_has_docstring(self):
        """task_id property should have documentation."""
        assert Context.task_id.fget.__doc__ is not None
        assert "task ID" in Context.task_id.fget.__doc__

    def test_session_has_docstring(self):
        """session property should document background task support."""
        assert Context.session.fget.__doc__ is not None
        assert "background task" in Context.session.fget.__doc__.lower()


# =============================================================================
# Integration tests: Client(mcp) + memory:// Docket backend
# =============================================================================


class TestBackgroundTaskIntegration:
    """Integration tests for background task context using real Docket memory backend.

    These tests use Client(mcp) with the memory:// broker — no mocking.
    The memory:// backend provides a fully functional in-memory Redis store
    that Docket uses automatically when running tests.
    """

    async def test_report_progress_in_background_task(self):
        """report_progress() should complete without error in a background task."""
        mcp = FastMCP("progress-test")
        progress_reported = asyncio.Event()

        @mcp.tool(task=True)
        async def progress_tool(ctx: Context) -> str:
            await ctx.report_progress(0, 100, "Starting...")
            await ctx.report_progress(50, 100, "Half done")
            await ctx.report_progress(100, 100, "Complete")
            progress_reported.set()
            return "done"

        async with Client(mcp) as client:
            task = await client.call_tool("progress_tool", {}, task=True)
            await asyncio.wait_for(progress_reported.wait(), timeout=5.0)
            await task.wait(timeout=5.0)
            result = await task.result()
            assert result.data == "done"

    async def test_context_wiring_in_background_task(self):
        """Context should be properly wired with task_id and session_id."""
        mcp = FastMCP("wiring-test")
        task_completed = asyncio.Event()
        captured: dict[str, object] = {}

        @mcp.tool(task=True)
        async def verify_wiring(ctx: Context) -> str:
            captured["task_id"] = ctx.task_id
            captured["session_id"] = ctx.session_id
            captured["is_background"] = ctx.is_background_task
            task_completed.set()
            return "ok"

        async with Client(mcp) as client:
            task = await client.call_tool("verify_wiring", {}, task=True)
            await asyncio.wait_for(task_completed.wait(), timeout=5.0)
            await task.wait(timeout=5.0)
            result = await task.result()
            assert result.data == "ok"

        assert captured["task_id"] is not None
        assert captured["session_id"] is not None
        assert captured["is_background"] is True

    async def test_origin_request_id_round_trips_through_background_task(self):
        """E2E: origin_request_id captured at submit time is restored in worker.

        We validate this by comparing ctx.origin_request_id with the value
        stored in Docket's Redis for this task.
        """

        mcp = FastMCP("origin-request-id-roundtrip")

        @mcp.tool(task=True)
        async def check_origin_request_id(ctx: Context, docket=CurrentDocket()) -> str:
            assert ctx.is_background_task is True
            assert ctx.request_context is None
            assert ctx.task_id is not None

            origin = ctx.origin_request_id
            assert origin is not None
            assert isinstance(origin, str)
            assert origin != ""

            # Verify the snapshot in Redis contains the same value
            task_scope = get_task_scope()
            key = docket.key(f"{task_redis_prefix(task_scope)}:{ctx.task_id}:snapshot")
            async with docket.redis() as redis:
                raw = await redis.get(key)

            assert raw is not None
            if isinstance(raw, bytes):
                raw = raw.decode()
            snapshot = json.loads(raw)
            assert snapshot["origin_request_id"] == origin
            return "ok"

        async with Client(mcp) as client:
            task = await client.call_tool("check_origin_request_id", {}, task=True)
            result = await task.result()
            assert result.data == "ok"

    @pytest.mark.xfail(
        reason="Background-task sampling has no back-channel under SDK v2: the "
        "per-request ServerSession that would carry sampling/createMessage is "
        "gone once the submitting request completes, so ctx.sample() from a "
        "worker raises NoBackChannelError. Needs a relay like elicit() "
        "(context.py TODO); tracked in sdk-feedback.",
        strict=True,
    )
    async def test_sample_uses_origin_request_id_in_background_task(self):
        """E2E: ctx.sample() works in a task without an active request context."""
        mcp = FastMCP("sample-background-test")
        captured: dict[str, object] = {}

        @mcp.tool(task=True)
        async def ask_client(ctx: Context) -> str:
            assert ctx.is_background_task is True
            assert ctx.request_context is None
            assert ctx.origin_request_id is not None
            result = await ctx.sample("Say hello")
            return result.text or ""

        def sampling_handler(messages, params, ctx):
            captured["called"] = True
            return CreateMessageResult(
                role="assistant",
                content=TextContent(type="text", text="hello from background"),
                model="test-model",
                stop_reason="endTurn",
            )

        async with Client(mcp, sampling_handler=sampling_handler) as client:
            task = await client.call_tool("ask_client", {}, task=True)
            result = await task.result()

        assert result.data == "hello from background"
        assert captured["called"] is True

    async def test_elicit_accept_flow(self):
        """E2E: tool elicits input, client accepts via elicitation_handler."""
        mcp = FastMCP("elicit-accept-test")

        @mcp.tool(task=True)
        async def ask_name(ctx: Context) -> str:
            result = await ctx.elicit("What is your name?", str)
            if isinstance(result, AcceptedElicitation):
                return f"Hello, {result.data}!"
            return "No name provided"

        async def handler(message, response_type, params, ctx):
            return ElicitResult(action="accept", content={"value": "Bob"})

        async with Client(mcp, elicitation_handler=handler) as client:
            task = await client.call_tool("ask_name", {}, task=True)
            await task.wait(timeout=10.0)
            result = await task.result()
            assert result.data == "Hello, Bob!"

    async def test_elicit_decline_flow(self):
        """E2E: tool elicits input, client declines via elicitation_handler."""
        mcp = FastMCP("elicit-decline-test")

        @mcp.tool(task=True)
        async def optional_input(ctx: Context) -> str:
            result = await ctx.elicit("Want to provide a name?", str)
            if isinstance(result, DeclinedElicitation):
                return "User declined"
            if isinstance(result, AcceptedElicitation):
                return f"Got: {result.data}"
            return "Cancelled"

        async def handler(message, response_type, params, ctx):
            return ElicitResult(action="decline")

        async with Client(mcp, elicitation_handler=handler) as client:
            task = await client.call_tool("optional_input", {}, task=True)
            await task.wait(timeout=10.0)
            result = await task.result()
            assert result.data == "User declined"

    async def test_elicit_with_pydantic_model(self):
        """E2E: tool elicits structured Pydantic input via elicitation_handler."""

        class UserInfo(BaseModel):
            name: str
            age: int

        mcp = FastMCP("elicit-pydantic-test")

        @mcp.tool(task=True)
        async def get_user_info(ctx: Context) -> str:
            result = await ctx.elicit("Provide user info", UserInfo)
            if isinstance(result, AcceptedElicitation):
                assert isinstance(result.data, UserInfo)
                return f"{result.data.name} is {result.data.age}"
            return "No info"

        async def handler(message, response_type, params, ctx):
            return ElicitResult(action="accept", content={"name": "Alice", "age": 30})

        async with Client(mcp, elicitation_handler=handler) as client:
            task = await client.call_tool("get_user_info", {}, task=True)
            await task.wait(timeout=10.0)
            result = await task.result()
            assert result.data == "Alice is 30"

    async def test_handle_task_input_rejects_when_not_waiting(self):
        """handle_task_input returns False when no task is waiting for input."""
        mcp = FastMCP("reject-test")

        @mcp.tool(task=True)
        async def simple_tool() -> str:
            return "done"

        async with Client(mcp) as client:
            task = await client.call_tool("simple_tool", {}, task=True)
            await task.wait(timeout=5.0)

            # Task already completed — no elicitation waiting
            success = await handle_task_input(
                task_id=task.task_id,
                task_scope="nonexistent-scope",
                action="accept",
                content={"value": "too late"},
                fastmcp=mcp,
            )
            assert success is False


class TestAccessTokenInBackgroundTasks:
    """Tests for access token availability in background tasks (#3095).

    Integration tests use Client(mcp) with the real memory:// Docket backend.
    The token snapshot/restore round-trip flows through actual Redis (fakeredis).

    Note: async tests run in isolated asyncio tasks, so ContextVar changes
    are automatically scoped — no cleanup required.
    """

    async def test_token_round_trips_through_background_task(self):
        """E2E: token set at submit time is available inside the worker."""
        mcp = FastMCP("token-roundtrip")

        @mcp.tool(task=True)
        async def check_token(ctx: Context) -> str:
            token = get_access_token()
            if token is None:
                return "no-token"
            return f"{token.token}|{token.client_id}"

        test_token = AccessToken(
            token="roundtrip-jwt",
            client_id="test-client",
            scopes=["read"],
            claims={"sub": "user-1"},
        )
        auth_context_var.set(AuthenticatedUser(test_token))

        async with Client(mcp) as client:
            task = await client.call_tool("check_token", {}, task=True)
            result = await task.result()
            assert result.data == "roundtrip-jwt|test-client"

    async def test_no_token_when_unauthenticated(self):
        """E2E: background task gets no token when nothing was set."""
        mcp = FastMCP("no-auth")

        @mcp.tool(task=True)
        async def check_token(ctx: Context) -> str:
            token = get_access_token()
            return "no-token" if token is None else token.token

        async with Client(mcp) as client:
            task = await client.call_tool("check_token", {}, task=True)
            result = await task.result()
            assert result.data == "no-token"

    async def test_expired_token_returns_none(self):
        """get_access_token() returns None when task token has expired."""
        expired = AccessToken(
            token="expired-jwt",
            client_id="test-client",
            scopes=["read"],
            expires_at=int(datetime.now(timezone.utc).timestamp()) - 3600,
        )
        _remember_snapshot(
            "test-task",
            TaskContextSnapshot(access_token_json=expired.model_dump_json()),
        )
        fake_ctx = TaskContextInfo(task_id="test-task", task_scope="s")
        with patch(
            "fastmcp.server.dependencies.get_task_context", return_value=fake_ctx
        ):
            assert get_access_token() is None

    async def test_valid_token_with_future_expiry(self):
        """get_access_token() returns token when expiry is in the future."""
        valid = AccessToken(
            token="valid-jwt",
            client_id="test-client",
            scopes=["read"],
            expires_at=int(datetime.now(timezone.utc).timestamp()) + 3600,
        )
        _remember_snapshot(
            "test-task",
            TaskContextSnapshot(access_token_json=valid.model_dump_json()),
        )
        fake_ctx = TaskContextInfo(task_id="test-task", task_scope="s")
        with patch(
            "fastmcp.server.dependencies.get_task_context", return_value=fake_ctx
        ):
            result = get_access_token()
            assert result is not None
            assert result.token == "valid-jwt"

    async def test_token_without_expiry_always_valid(self):
        """get_access_token() returns token when no expires_at is set."""
        no_expiry = AccessToken(
            token="eternal-jwt",
            client_id="test-client",
            scopes=["read"],
        )
        _remember_snapshot(
            "test-task",
            TaskContextSnapshot(access_token_json=no_expiry.model_dump_json()),
        )
        fake_ctx = TaskContextInfo(task_id="test-task", task_scope="s")
        with patch(
            "fastmcp.server.dependencies.get_task_context", return_value=fake_ctx
        ):
            result = get_access_token()
            assert result is not None
            assert result.token == "eternal-jwt"


class TestLifespanContextInBackgroundTasks:
    """Tests for lifespan_context availability in background tasks (#3095)."""

    def test_lifespan_context_falls_back_to_server_result(self):
        """lifespan_context reads from server when request_context is None."""
        mcp = FastMCP("test")
        mcp._lifespan_result = {"db": "mock-db-connection", "cache": "mock-cache"}

        ctx = Context(mcp, task_id="test-task")
        assert ctx.request_context is None
        assert ctx.lifespan_context == {
            "db": "mock-db-connection",
            "cache": "mock-cache",
        }

    def test_lifespan_context_returns_empty_dict_when_no_lifespan(self):
        """lifespan_context returns {} when no lifespan is configured."""
        mcp = FastMCP("test")
        ctx = Context(mcp, task_id="test-task")
        assert ctx.request_context is None
        assert ctx.lifespan_context == {}