prefecthq--fastmcp/tests/server/transforms/test_catalog.py

"""Tests for CatalogTransform base class."""

from __future__ import annotations

import ast
from collections.abc import Sequence

from mcp_types import TextContent

from fastmcp import FastMCP
from fastmcp.server.context import Context
from fastmcp.server.transforms import GetToolNext
from fastmcp.server.transforms.catalog import CatalogTransform
from fastmcp.server.transforms.version_filter import VersionFilter
from fastmcp.tools.base import Tool
from fastmcp.utilities.versions import VersionSpec

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_versioned_tool(name: str, version: str) -> Tool:
    """Create a tool with a specific version for testing."""

    async def _fn() -> str:
        return f"{name}@{version}"

    return Tool.from_function(fn=_fn, name=name, version=version)


class CatalogReader(CatalogTransform):
    """Minimal CatalogTransform that exposes get_tool_catalog via a tool.

    After calling the ``read_catalog`` tool, the full catalog is available
    on ``self.last_catalog`` for assertions beyond tool names.
    """

    def __init__(self) -> None:
        super().__init__()
        self.last_catalog: Sequence[Tool] = []

    async def transform_tools(self, tools: Sequence[Tool]) -> Sequence[Tool]:
        return [self._make_reader_tool()]

    async def get_tool(
        self, name: str, call_next: GetToolNext, *, version: VersionSpec | None = None
    ) -> Tool | None:
        if name == "read_catalog":
            return self._make_reader_tool()
        return await call_next(name, version=version)

    def _make_reader_tool(self) -> Tool:
        transform = self

        async def read_catalog(ctx: Context = None) -> list[str]:  # type: ignore[assignment]  # ty:ignore[invalid-parameter-default]
            """Return names of tools visible in the catalog."""
            transform.last_catalog = await transform.get_tool_catalog(ctx)
            return [t.name for t in transform.last_catalog]

        return Tool.from_function(fn=read_catalog, name="read_catalog")


class ReplacingTransform(CatalogTransform):
    """Minimal subclass that replaces tools with a synthetic tool.

    Uses ``get_tool_catalog()`` to read the real catalog inside the
    synthetic tool's handler, verifying that the bypass mechanism works.
    """

    async def transform_tools(self, tools: Sequence[Tool]) -> Sequence[Tool]:
        return [self._make_synthetic_tool()]

    async def get_tool(
        self, name: str, call_next: GetToolNext, *, version: VersionSpec | None = None
    ) -> Tool | None:
        if name == "count_tools":
            return self._make_synthetic_tool()
        return await call_next(name, version=version)

    def _make_synthetic_tool(self) -> Tool:
        transform = self

        async def count_tools(ctx: Context = None) -> int:  # type: ignore[assignment]  # ty:ignore[invalid-parameter-default]
            """Return the number of real tools in the catalog."""
            catalog = await transform.get_tool_catalog(ctx)
            return len(catalog)

        return Tool.from_function(fn=count_tools, name="count_tools")


class TestCatalogTransformBypass:
    async def test_list_tools_replaced_by_subclass(self):
        mcp = FastMCP("test")

        @mcp.tool
        def add(a: int, b: int) -> int:
            return a + b

        @mcp.tool
        def multiply(x: float, y: float) -> float:
            return x * y

        mcp.add_transform(ReplacingTransform())
        tools = await mcp.list_tools()
        names = {t.name for t in tools}
        assert names == {"count_tools"}

    async def test_get_tool_catalog_returns_real_tools(self):
        mcp = FastMCP("test")

        @mcp.tool
        def add(a: int, b: int) -> int:
            return a + b

        @mcp.tool
        def multiply(x: float, y: float) -> float:
            return x * y

        mcp.add_transform(ReplacingTransform())
        result = await mcp.call_tool("count_tools", {})
        assert any("2" in c.text for c in result.content if isinstance(c, TextContent))

    async def test_multiple_instances_have_independent_bypass(self):
        """Each CatalogTransform instance has its own bypass ContextVar."""
        t1 = ReplacingTransform()
        t2 = ReplacingTransform()
        assert t1._instance_id != t2._instance_id
        assert t1._bypass is not t2._bypass


class TestCatalogDeduplication:
    """get_tool_catalog() deduplicates versioned tools, keeping only the highest."""

    async def test_returns_highest_version_only(self):
        mcp = FastMCP("test")
        mcp.add_tool(_make_versioned_tool("greet", "1"))
        mcp.add_tool(_make_versioned_tool("greet", "2"))
        mcp.add_tool(_make_versioned_tool("greet", "3"))

        reader = CatalogReader()
        mcp.add_transform(reader)

        result = await mcp.call_tool("read_catalog", {})
        names = _extract_result(result)
        assert names == ["greet"]

    async def test_version_metadata_injected(self):
        """Highest-version tool has meta.fastmcp.versions listing all available."""
        mcp = FastMCP("test")
        mcp.add_tool(_make_versioned_tool("greet", "1"))
        mcp.add_tool(_make_versioned_tool("greet", "3"))

        reader = CatalogReader()
        mcp.add_transform(reader)

        await mcp.call_tool("read_catalog", {})
        assert len(reader.last_catalog) == 1
        tool = reader.last_catalog[0]
        assert tool.name == "greet"
        assert tool.version == "3"
        assert tool.meta is not None
        versions = tool.meta["fastmcp"]["versions"]
        assert versions == ["3", "1"]

    async def test_mixed_versioned_and_unversioned(self):
        mcp = FastMCP("test")

        @mcp.tool
        def standalone() -> str:
            return "hi"

        mcp.add_tool(_make_versioned_tool("greet", "1"))
        mcp.add_tool(_make_versioned_tool("greet", "2"))

        reader = CatalogReader()
        mcp.add_transform(reader)

        result = await mcp.call_tool("read_catalog", {})
        names = sorted(_extract_result(result))
        assert names == ["greet", "standalone"]

    async def test_version_filter_applied_before_catalog(self):
        """A VersionFilter added before the CatalogTransform restricts what the catalog sees."""
        mcp = FastMCP("test")
        mcp.add_tool(_make_versioned_tool("greet", "1"))
        mcp.add_tool(_make_versioned_tool("greet", "2"))
        mcp.add_tool(_make_versioned_tool("greet", "3"))

        # Filter keeps only versions < 3 — so v1 and v2 survive, v3 is excluded.
        mcp.add_transform(VersionFilter(version_lt="3"))

        reader = CatalogReader()
        mcp.add_transform(reader)

        await mcp.call_tool("read_catalog", {})
        assert len(reader.last_catalog) == 1
        tool = reader.last_catalog[0]
        assert tool.name == "greet"
        assert tool.version == "2"


class TestCatalogVisibility:
    """get_tool_catalog() respects visibility (disabled tools are excluded)."""

    async def test_disabled_tool_excluded(self):
        mcp = FastMCP("test")

        @mcp.tool
        def public() -> str:
            return "visible"

        @mcp.tool
        def secret() -> str:
            return "hidden"

        mcp.disable(names={"secret"}, components={"tool"})
        reader = CatalogReader()
        mcp.add_transform(reader)

        result = await mcp.call_tool("read_catalog", {})
        names = _extract_result(result)
        assert names == ["public"]


class TestCatalogAuth:
    """get_tool_catalog() respects tool-level auth filtering."""

    async def test_auth_rejected_tool_excluded(self):
        mcp = FastMCP("test")

        @mcp.tool
        def public() -> str:
            return "visible"

        @mcp.tool(auth=lambda _ctx: False)
        def protected() -> str:
            return "hidden"

        reader = CatalogReader()
        mcp.add_transform(reader)

        result = await mcp.call_tool("read_catalog", {})
        names = _extract_result(result)
        assert names == ["public"]


# ---------------------------------------------------------------------------
# Test helpers
# ---------------------------------------------------------------------------


def _extract_result(result: object) -> list[str]:
    """Extract the list of names from a call_tool ToolResult."""
    for c in result.content:  # type: ignore[union-attr]  # ty:ignore[unresolved-attribute]
        if isinstance(c, TextContent):
            return ast.literal_eval(c.text)
    raise AssertionError("No text content found")