fastmcp/tests/tasks/server/test_task_dependencies.py

"""Tests for dependency injection in background tasks.

These tests verify that Docket's dependency system works correctly when tool
functions are queued as background tasks. Dependencies like CurrentDocket(),
CurrentFastMCP(), and Depends() should be resolved in the worker context.

SEP-2663 is tools-only, so only tools carry a task-capable config; the removed
prompt/resource task cases are gone.
"""

from __future__ import annotations

from contextlib import asynccontextmanager
from typing import Any, cast

import pytest
from fastmcp_tasks.dependencies import CurrentDocket
from uncalled_for import Depends

from fastmcp import FastMCP
from fastmcp.server.dependencies import CurrentFastMCP
from fastmcp_tasks import TasksExtension
from tests.tasks.task_helpers import (
    call_tool_without_optin,
    run_task,
    running_task_server,
)


@pytest.fixture
def dependency_server() -> FastMCP:
    """A FastMCP server with dependency-using background tools."""
    mcp = FastMCP("dependency-test-server")
    mcp.add_extension(TasksExtension())

    injected_values: list[tuple[str, Any]] = []

    @mcp.tool(task=True)
    async def tool_with_docket_dependency(docket=CurrentDocket()) -> str:
        injected_values.append(("docket", docket))
        return f"Docket: {docket is not None}"

    @mcp.tool(task=True)
    async def tool_with_server_dependency(server=CurrentFastMCP()) -> str:
        injected_values.append(("server", server))
        return f"Server: {server.name}"

    @mcp.tool(task=True)
    async def tool_with_custom_dependency(
        value: int, multiplier: int = Depends(lambda: 10)
    ) -> int:
        injected_values.append(("multiplier", multiplier))
        return value * multiplier

    @mcp.tool(task=True)
    async def tool_with_multiple_dependencies(
        name: str,
        docket=CurrentDocket(),
        server=CurrentFastMCP(),
    ) -> str:
        injected_values.append(("multi_docket", docket))
        injected_values.append(("multi_server", server))
        return f"{name} on {server.name}"

    mcp._injected_values = injected_values  # type: ignore[attr-defined]  # ty:ignore[unresolved-attribute]

    return mcp


async def test_background_tool_receives_docket_dependency(dependency_server):
    """Background tools can use CurrentDocket() and it resolves in the worker."""
    async with running_task_server(dependency_server):
        final = await run_task(dependency_server, "tool_with_docket_dependency", {})

    assert final.status == "completed"
    assert final.result is not None
    assert final.result["structuredContent"] == {"result": "Docket: True"}
    assert len(dependency_server._injected_values) == 1
    dep_type, dep_value = dependency_server._injected_values[0]
    assert dep_type == "docket"
    assert dep_value is not None


async def test_background_tool_receives_server_dependency(dependency_server):
    """Background tools can use CurrentFastMCP() and get the actual server."""
    dependency_server._injected_values.clear()

    async with running_task_server(dependency_server):
        final = await run_task(dependency_server, "tool_with_server_dependency", {})

    assert final.status == "completed"
    assert final.result is not None
    assert final.result["structuredContent"] == {
        "result": f"Server: {dependency_server.name}"
    }
    assert len(dependency_server._injected_values) == 1
    dep_type, dep_value = dependency_server._injected_values[0]
    assert dep_type == "server"
    assert dep_value is dependency_server  # Same instance!


async def test_background_tool_receives_custom_depends(dependency_server):
    """Background tools can use Depends() with custom functions."""
    dependency_server._injected_values.clear()

    async with running_task_server(dependency_server):
        final = await run_task(
            dependency_server, "tool_with_custom_dependency", {"value": 5}
        )

    assert final.status == "completed"
    assert final.result is not None
    assert final.result["structuredContent"] == {"result": 50}  # 5 * 10
    assert len(dependency_server._injected_values) == 1
    dep_type, dep_value = dependency_server._injected_values[0]
    assert dep_type == "multiplier"
    assert dep_value == 10


async def test_background_tool_with_multiple_dependencies(dependency_server):
    """Background tools can have multiple dependencies injected at once."""
    dependency_server._injected_values.clear()

    async with running_task_server(dependency_server):
        final = await run_task(
            dependency_server, "tool_with_multiple_dependencies", {"name": "test"}
        )

    assert final.status == "completed"
    assert final.result is not None
    assert final.result["structuredContent"] == {
        "result": f"test on {dependency_server.name}"
    }

    dep_types = {item[0] for item in dependency_server._injected_values}
    assert "multi_docket" in dep_types
    assert "multi_server" in dep_types

    server_dep = next(
        v for t, v in dependency_server._injected_values if t == "multi_server"
    )
    assert server_dep is dependency_server


async def test_foreground_tool_dependencies_unaffected(dependency_server):
    """Synchronous tools still get their dependencies as before."""
    dependency_server._injected_values.clear()

    @dependency_server.tool
    async def sync_tool(server=CurrentFastMCP()) -> str:
        dependency_server._injected_values.append(("sync_server", server))
        return f"Sync: {server.name}"

    async with running_task_server(dependency_server):
        await call_tool_without_optin(dependency_server, "sync_tool", {})

    assert len(dependency_server._injected_values) == 1
    assert dependency_server._injected_values[0][1] is dependency_server


async def test_dependency_context_managers_cleaned_up_in_background():
    """Context-manager dependencies are cleaned up after a background task."""
    cleanup_called: list[str] = []

    mcp = FastMCP("cleanup-test")
    mcp.add_extension(TasksExtension())

    @asynccontextmanager
    async def tracked_connection():
        try:
            cleanup_called.append("enter")
            yield "connection"
        finally:
            cleanup_called.append("exit")

    @mcp.tool(task=True)
    async def use_connection(name: str, conn: str = Depends(tracked_connection)) -> str:
        assert conn == "connection"
        assert "enter" in cleanup_called
        assert "exit" not in cleanup_called  # Still open during execution
        return f"Used: {conn}"

    async with running_task_server(mcp):
        final = await run_task(mcp, "use_connection", {"name": "test"})

    assert final.status == "completed"
    assert final.result is not None
    assert final.result["structuredContent"] == {"result": "Used: connection"}
    assert cleanup_called == ["enter", "exit"]


async def test_dependency_errors_propagate_to_task_failure():
    """If dependency resolution fails, the background task should fail."""
    mcp = FastMCP("error-test")
    mcp.add_extension(TasksExtension())

    async def failing_dependency():
        raise ValueError("Dependency failed!")

    @mcp.tool(task=True)
    async def tool_with_failing_dep(
        value: str, dep: str = cast(Any, Depends(failing_dependency))
    ) -> str:
        return f"Got: {dep}"

    async with running_task_server(mcp):
        final = await run_task(mcp, "tool_with_failing_dep", {"value": "test"})

    assert final.status == "failed"
    assert final.error is not None